Transformer Engine plugin fails to check weight exists for LayerNorm

Bug description

At https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/plugins/precision/transformer_engine.py#L173

There is no check that the weights of the LayerNorm layer are not None

This means that if a LayerNorm layer is created using elementwise_affine=False we get the following error

Traceback (most recent call last):
  File "/usr/lib/python3.14/pdb.py", line 3656, in main
    pdb._run(target)
    ~~~~~~~~^^^^^^^^
  File "/usr/lib/python3.14/pdb.py", line 2566, in _run
    self.run(target.code)
    ~~~~~~~~^^^^^^^^^^^^^
  File "/usr/lib/python3.14/bdb.py", line 913, in run
    exec(cmd, globals, locals)
    ~~~~^^^^^^^^^^^^^^^^^^^^^^
  File "<string>", line 1, in <module>
  File "/home/nvme/budget-flow-matching/src/main.py", line 212, in <module>
    trainer.fit(model, dataloader, val_dataloaders=val_dataloader, ckpt_path=args.continue_from)
    ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/pytorch/trainer/trainer.py", line 584, in fit
    call._call_and_handle_interrupt(
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        self,
        ^^^^^
    ...<6 lines>...
        weights_only,
        ^^^^^^^^^^^^^
    )
    ^
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/pytorch/trainer/call.py", line 49, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/pytorch/trainer/trainer.py", line 630, in _fit_impl
    self._run(model, ckpt_path=ckpt_path, weights_only=weights_only)
    ~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/pytorch/trainer/trainer.py", line 1053, in _run
    self.strategy.setup(self)
    ~~~~~~~~~~~~~~~~~~~^^^^^^
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/pytorch/strategies/strategy.py", line 154, in setup
    self.model = self.precision_plugin.convert_module(self.model)
                 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/fabric/plugins/precision/transformer_engine.py", line 105, in convert_module
    _convert_layers(module)
    ~~~~~~~~~~~~~~~^^^^^^^^
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/fabric/plugins/precision/transformer_engine.py", line 182, in _convert_layers
    _convert_layers(child)
    ~~~~~~~~~~~~~~~^^^^^^^
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/fabric/plugins/precision/transformer_engine.py", line 182, in _convert_layers
    _convert_layers(child)
    ~~~~~~~~~~~~~~~^^^^^^^
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/fabric/plugins/precision/transformer_engine.py", line 182, in _convert_layers
    _convert_layers(child)
    ~~~~~~~~~~~~~~~^^^^^^^
  File "/home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/fabric/plugins/precision/transformer_engine.py", line 173, in _convert_layers
    replacement.weight.data = child.weight.data.clone()
                              ^^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'data'
Uncaught exception. Entering post mortem debugging
Running 'cont' or 'step' will restart the program
> /home/henry/.local/share/virtualenvs/budget-flow-matching-u9VKU3BC/lib/python3.14/site-packages/lightning/fabric/plugins/precision/transformer_engine.py(173)_convert_layers()
-> replacement.weight.data = child.weight.data.clone()
(Pdb) child
LayerNorm((896,), eps=1e-06, elementwise_affine=False, bias=False)

What version are you seeing the problem on?

v2.6

Reproduced in studio

No response

How to reproduce the bug

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset

import lightning as L


class DummyDataset(Dataset):
    def __len__(self):
        return 8

    def __getitem__(self, idx):
        return torch.randn(16), torch.randn(2)

class BuggyModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        # elementwise_affine=False sets weight and bias to None, triggering the bug
        self.layer_norm = nn.LayerNorm(16, elementwise_affine=False)
        self.linear = nn.Linear(16, 2)

    def forward(self, x):
        return self.linear(self.layer_norm(x))

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = nn.functional.mse_loss(self(x), y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)

if __name__ == "__main__":
    model = BuggyModel()
    dataset = DummyDataset()
    dataloader = DataLoader(dataset, batch_size=2)

    # Initialize the trainer with the Transformer Engine precision Lugin
    trainer = L.Trainer(
        max_epochs=1,
        accelerator="cuda",
        devices=1,
        precision="transformer-engine"
    )

    print("Starting trainer.fit()...")
    # This will crash inside trainer.fit() during the strategy setup phase
    trainer.fit(model, dataloader)
    def __getitem__(self, idx):
        return torch.randn(16), torch.randn(2)

Error messages and logs

# Error messages and logs here please

Environment

More info

Should be fixable with a simple

if child.weight is not None and replacement.weight is not None:

cc @ethanwharris

Start in src/lightning/fabric/plugins/precision/transformer_engine.py at _convert_layers around line 173, then run the provided LayerNorm reproduction with elementwise_affine=False and transformer-engine precision. Confirm that conversion handles missing LayerNorm weights without the reported AttributeError and that the training setup completes.