Maybe this is better

Author: thomasw21

megatron/model/shared_t5_model.py

@@ -91,8 +91,8 @@ def _to_16bit(inputs):
         self.specs.append(
             LayerSpec(
                 LayerNorm,
+                args.hidden_size,
                 forward_fn=lambda module, input_tokens, target_tokens: (module(input_tokens), target_tokens),
-                hidden_size=args.hidden_size,
                 eps=args.layernorm_epsilon
             ))
 
@@ -121,7 +121,7 @@ def _to_16bit(inputs):
         self.specs.append(
             LayerSpec(
                 LayerNorm,
-                hidden_size=args.hidden_size,
+                args.hidden_size,
                 eps=args.layernorm_epsilon
             ))