Skip to content

Preserve single-output bias shape when reversing fusion - #2539

Open
kocean9-freedom wants to merge 1 commit into
huggingface:mainfrom
kocean9-freedom:fix/fused-bias-single-output
Open

kocean9-freedom wants to merge 1 commit into
huggingface:mainfrom
kocean9-freedom:fix/fused-bias-single-output

Conversation

@kocean9-freedom

@kocean9-freedom kocean9-freedom commented Oct 22, 2026 •

Copy link
Copy Markdown

What does this PR do?

Preserve a one-element bias vector when reversing FuseBiasInLinear. module.weight[:, -1] already has shape [out_features]; the extra squeeze() turns it into a scalar when out_features=1. The restored layer's state dict then fails to load into the original nn.Linear, with size mismatch for bias: torch.Size([]) ... torch.Size([1]). Single-output classification/regression heads, including text rerankers, are affected.

Remove only that redundant squeeze. Multi-output and biasless layers keep their behavior; no dependency or API changes.

Validation on Python 3.12 / PyTorch 2.10 / Transformers 5.19:

  • New download-free native FX tests: 6 failures and 8 passing controls on upstream; 14 passed after the fix. They cover one/multiple outputs, 1D/2D/3D inputs, float32/float64, biasless controls, two fuse/reverse cycles, exact parameter shapes/values, strict checkpoint loading, and input/parameter gradients.
  • The actual single-output head from cached pretrained cross-encoder-testing/reranker-bert-tiny-gooaq-bce-STv54 reproduces the checkpoint-loading failure on upstream. With the fix, CPU and MPS paired-text scores, strict state loading, model save/reload, gradients, and one SGD update match the untransformed reference. These checks also pass from an isolated built wheel.
  • Full-repository Black and Ruff checks, git diff --check, sdist and wheel builds pass. The complete downloaded-model FX/Optimum test suite and CUDA were not run.

The shared-call fix in #2538 is a separate issue and its current head retains this squeeze. In a separate checkout, combining its patch with this one-line fix passes all 24 shared-call and single-output cases. That patch is not included here. The dtype issue in #2533 is also outside this change.

Before submitting

  • Added regression tests and checked the real model use case.
  • Public arguments and documented behavior are unchanged; no documentation update is needed.
  • Ran code style checks and recorded verification limits above.

AI assistance: implementation, validation, review, and PR preparation used OpenAI Codex.

Additional legacy source compatibility

A separate checkout combining #2535, #2539 and #2540 passes 33 offline cases with Transformers 4.57.6 on both PyTorch 2.10.0 and 2.5.1: the two new 14-case suites, #2535's three layouts, and two existing offline helper tests. This is a source compatibility overlay using Hub 0.36.2, which does not satisfy current main's Hub >=1.31 requirement; repository dependency pins are unchanged. It is not a fully resolved legacy installation or a downloaded-model/full-suite claim.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant