[Bug] The Transformer Engine plugin seems to be incompatible with LayerNorm that has no weights.

[IDKiro]

System Info

When the elementwise_affine attribute of PyTorch’s LayerNorm is set to False, the LayerNorm has no weights or biases.
At this point, the following error will occur:

  File "/workspace/.venv/lib/python3.10/site-packages/accelerate/accelerator.py", line 1342, in prepare
    args = self._prepare_te(*args)
  File "/workspace/.venv/lib/python3.10/site-packages/accelerate/accelerator.py", line 1642, in _prepare_te
    convert_model(model)
  File "/workspace/.venv/lib/python3.10/site-packages/accelerate/utils/transformer_engine.py", line 77, in convert_model
    convert_model(
  File "/workspace/.venv/lib/python3.10/site-packages/accelerate/utils/transformer_engine.py", line 77, in convert_model
    convert_model(
  File "/workspace/.venv/lib/python3.10/site-packages/accelerate/utils/transformer_engine.py", line 77, in convert_model
    convert_model(
  [Previous line repeated 2 more times]
  File "/workspace/.venv/lib/python3.10/site-packages/accelerate/utils/transformer_engine.py", line 55, in convert_model
    te_module = te.LayerNorm(module.normalized_shape[0], eps=module.eps, params_dtype=module.weight.dtype)
AttributeError: 'NoneType' object has no attribute 'dtype'

Information

• The official example scripts
• My own modified scripts

Tasks

• One of the scripts in the examples/ folder of Accelerate or an officially supported no_trainer script in the examples folder of the transformers repo (such as run_no_trainer_glue.py)
• My own task or dataset (give details below)

Reproduction

1. when finetuning the flux model using FP8

Expected behavior

1. Just add a judgment