Fix: TransformerEnginePrecision conversion for layers with bias=False #20805
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.
What does this PR do?
This PR resolves a runtime AttributeError occurring in the TransformerEnginePrecision plugin during layer conversion.
When the convert_module() method replaces standard PyTorch layers (nn.Linear, nn.LayerNorm, etc.) with transformer engine counterparts, it attempts to clone both weight and bias tensors. However, some layers may be instantiated with bias=False, resulting in child.bias being None. The following line in the plugin causes a crash when this is the case:
replacement.bias.data = child.bias.data.clone()
This PR introduces a safe check to ensure both child.bias and replacement.bias are not None before attempting to access .data. This ensures compatibility with modules that omit bias, such as:
nn.Linear(in_features=16, out_features=32, bias=False)
Additionally, the PR includes a targeted unit test (test_convert_module_handles_linear_without_bias) that defines a model using such a bias-less layer and verifies that the conversion logic runs without error. This prevents regression and guards against similar issues in future updates.
The fix is surgical, backward-compatible, and maintains the expected behavior for all other layers with bias enabled.
Fixes #20803
cc @Lightning-AI/fabric-maintainers
Summary of changes
Guarded .bias.data.clone() with if child.bias is not None and replacement.bias is not None
Added test test_convert_module_handles_linear_without_bias in test_transformer_engine.py
-No breaking changes introduced.
Before submitting
Was this discussed/agreed via a GitHub issue? No
Did you read the contributor guideline, Pull Request section? Yes
Did you make sure your PR does only one thing, instead of bundling different changes together? Yes
Did you write any new necessary tests? Yes, for the edge case bias=False
Did you verify new and existing tests pass locally with your changes? Yes
Not applicable to update the CHANGELOG for this bugfix
🙃 Did you have fun coding?
Absolutely. Excited to contribute more improvements soon!
📚 Documentation preview 📚: https://pytorch-lightning--20805.org.readthedocs.build/en/20805/