From 3abac1ae4c602644ff0dee0b14d773d59254c412 Mon Sep 17 00:00:00 2001 From: Kavyansh Tyagi <142140238+KAVYANSHTYAGI@users.noreply.github.com> Date: Fri, 9 May 2025 12:30:30 +0530 Subject: [PATCH 1/3] Update transformer_engine.py --- src/lightning/fabric/plugins/precision/transformer_engine.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/lightning/fabric/plugins/precision/transformer_engine.py b/src/lightning/fabric/plugins/precision/transformer_engine.py index c3ef84a453e73..bf1e51ea6b2b0 100644 --- a/src/lightning/fabric/plugins/precision/transformer_engine.py +++ b/src/lightning/fabric/plugins/precision/transformer_engine.py @@ -171,7 +171,9 @@ def _convert_layers(module: torch.nn.Module) -> None: elif isinstance(child, torch.nn.LayerNorm): replacement = te.LayerNorm(child.normalized_shape[0], eps=child.eps) replacement.weight.data = child.weight.data.clone() - replacement.bias.data = child.bias.data.clone() + # Check if bias exists before attempting to clone its data + if child.bias is not None and replacement.bias is not None: + replacement.bias.data = child.bias.data.clone() log.debug(f"Replacing layer {name!r} with Transformer Engine equivalent") module.__setattr__(name, replacement) else: From ab1dccd3ea44613c261ede951d6dbe79136f9dbd Mon Sep 17 00:00:00 2001 From: Kavyansh Tyagi <142140238+KAVYANSHTYAGI@users.noreply.github.com> Date: Fri, 9 May 2025 12:37:26 +0530 Subject: [PATCH 2/3] Update test_transformer_engine.py --- .../precision/test_transformer_engine.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/tests_fabric/plugins/precision/test_transformer_engine.py b/tests/tests_fabric/plugins/precision/test_transformer_engine.py index 033484aca9c90..e6fd51cad4e38 100644 --- a/tests/tests_fabric/plugins/precision/test_transformer_engine.py +++ b/tests/tests_fabric/plugins/precision/test_transformer_engine.py @@ -115,3 +115,37 @@ class TELayerNormMock(Mock): ... assert isinstance(model.l1, TELinearMock) assert isinstance(model.l2, TELayerNormMock) assert isinstance(model.l3.l, TELinearMock) + +def test_convert_module_handles_linear_without_bias(monkeypatch): + + module = lightning.fabric.plugins.precision.transformer_engine # Set up mock transformer_engine + monkeypatch.setattr(module, "_TRANSFORMER_ENGINE_AVAILABLE", lambda: True) + + transformer_engine_mock = Mock() + monkeypatch.setitem(sys.modules, "transformer_engine", transformer_engine_mock) + monkeypatch.setitem(sys.modules, "transformer_engine.pytorch", transformer_engine_mock.pytorch) + monkeypatch.setitem(sys.modules, "transformer_engine.common.recipe", transformer_engine_mock.recipe) + + + class TELinearMock(torch.nn.Linear): # Mock the Linear replacement class + def __init__(self, in_features, out_features, bias=True): + super().__init__(in_features, out_features, bias) + + transformer_engine_mock.pytorch.Linear = TELinearMock + transformer_engine_mock.pytorch.LayerNorm = torch.nn.LayerNorm + transformer_engine_mock.recipe.DelayedScaling.return_value = None + + class BiaslessModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(16, 32, bias=False) # This was causing the bug + + model = BiaslessModel() + precision = TransformerEnginePrecision(weights_dtype=torch.float16) + precision.replace_layers = True + + + precision.convert_module(model) # This should no longer raise AttributeError + + assert isinstance(model.linear, TELinearMock) + assert model.linear.bias is None From 9dcce14c86a969c03542e79f41cec1a38e8cd6f2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 May 2025 07:23:19 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../plugins/precision/test_transformer_engine.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/tests_fabric/plugins/precision/test_transformer_engine.py b/tests/tests_fabric/plugins/precision/test_transformer_engine.py index e6fd51cad4e38..ed7c984b1ae64 100644 --- a/tests/tests_fabric/plugins/precision/test_transformer_engine.py +++ b/tests/tests_fabric/plugins/precision/test_transformer_engine.py @@ -116,9 +116,9 @@ class TELayerNormMock(Mock): ... assert isinstance(model.l2, TELayerNormMock) assert isinstance(model.l3.l, TELinearMock) + def test_convert_module_handles_linear_without_bias(monkeypatch): - - module = lightning.fabric.plugins.precision.transformer_engine # Set up mock transformer_engine + module = lightning.fabric.plugins.precision.transformer_engine # Set up mock transformer_engine monkeypatch.setattr(module, "_TRANSFORMER_ENGINE_AVAILABLE", lambda: True) transformer_engine_mock = Mock() @@ -126,8 +126,7 @@ def test_convert_module_handles_linear_without_bias(monkeypatch): monkeypatch.setitem(sys.modules, "transformer_engine.pytorch", transformer_engine_mock.pytorch) monkeypatch.setitem(sys.modules, "transformer_engine.common.recipe", transformer_engine_mock.recipe) - - class TELinearMock(torch.nn.Linear): # Mock the Linear replacement class + class TELinearMock(torch.nn.Linear): # Mock the Linear replacement class def __init__(self, in_features, out_features, bias=True): super().__init__(in_features, out_features, bias) @@ -138,14 +137,13 @@ def __init__(self, in_features, out_features, bias=True): class BiaslessModel(torch.nn.Module): def __init__(self): super().__init__() - self.linear = torch.nn.Linear(16, 32, bias=False) # This was causing the bug + self.linear = torch.nn.Linear(16, 32, bias=False) # This was causing the bug model = BiaslessModel() precision = TransformerEnginePrecision(weights_dtype=torch.float16) precision.replace_layers = True - - precision.convert_module(model) # This should no longer raise AttributeError + precision.convert_module(model) # This should no longer raise AttributeError assert isinstance(model.linear, TELinearMock) assert model.linear.bias is None