Skip to content

Fix bias fusion for reused Linear modules - #2538

Open
Linxiushen wants to merge 1 commit into
huggingface:mainfrom
Linxiushen:codex/fuse-bias-shared-linear
Open

Linxiushen wants to merge 1 commit into
huggingface:mainfrom
Linxiushen:codex/fuse-bias-shared-linear

Conversation

@Linxiushen

Copy link
Copy Markdown

FuseBiasInLinear fails when the same torch.nn.Linear is called more than once. The first call site clears the module's bias, so subsequent sites are left with an unextended input and fail with a matrix multiplication shape mismatch. This affects both parallel uses and sequential reuse of a projection.

Rewrite all affected FX call sites before changing parameters, then append the bias column once per module. Reverse each call-site rewrite before restoring the shared module once. The regression covers two and three shared calls, sequential reuse, two input ranks, single-call and biasless controls, forward values, input/parameter gradients, idempotence, and two fuse/reverse cycles with exact state-dict restoration.

Validation on unchanged 0df84f23884f7c5b6cc76c0bdc9597edd306034f: the new tests produce six shape-mismatch failures and four passing controls. With this patch, all ten cases and two existing adjacent tests pass:

pytest tests/fx/optimization/test_shared_linear.py tests/fx/optimization/test_transformations.py::test_get_parent tests/fx/optimization/test_transformations.py::TransformationTester::test_transformation_signature
black --check .
ruff check .
git diff --check

Tested on Windows CPU with Python 3.12.14, PyTorch 2.9.1+cpu, Transformers 4.57.3 and pytest 8.4.2, with Hub access offline. Optimum was installed editable with --no-deps using existing Transformers 4.x dependencies (huggingface_hub 0.36.2); this was a bounded native FX test, not a fully dependency-resolved installation of the current Hub requirement. The full suite with downloaded models, GPU execution and model training were not run. The separate dtype issue in #2532 / #2533 is outside this change.

AI assistance: developed, tested and reviewed with OpenAI Codex.

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