Repository navigation
Fix bias fusion for reused Linear modules - #2538
Open
Linxiushen wants to merge 1 commit into
Open
Linxiushen wants to merge 1 commit into
Linxiushen wants to merge 1 commit into
Conversation
3 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
FuseBiasInLinearfails when the sametorch.nn.Linearis 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: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-depsusing 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.