Fix layernorm arg parsing (bigscience-workshop#281)

Author: Quentin-Anthony

megatron/arguments.py

@@ -628,7 +628,7 @@ def _add_network_size_args(parser):
                        'around zero. This improves numerical stability.')
     group.add_argument('--disable-mem-efficient-ln', action='store_false', 
                        help='Disable the memory-efficient fused LayerNorm optimization '
-                       'introduced in https://github.com/NVIDIA/apex/pull/1715')
+                       'introduced in https://github.com/NVIDIA/apex/pull/1715', dest='mem_efficient_ln')
     group.add_argument('--apply-residual-connection-post-layernorm',
                        action='store_true',
                        help='If set, use original BERT residula connection '

megatron/model/transformer.py

@@ -920,7 +920,7 @@ def __init__(self, config,
                     no_persist_layer_norm=args.no_persist_layer_norm,
                     sequence_parallel=config.sequence_parallel,
                     apply_layernorm_1p=args.apply_layernorm_1p,
-                    mem_efficient_ln=not args.disable_mem_efficient_ln)
+                    mem_efficient_ln=args.mem_efficient_ln)
             else:
                 self.input_layernorm = LayerNorm(
                     config.hidden_size,
@@ -946,7 +946,7 @@ def __init__(self, config,
                     no_persist_layer_norm=not config.persist_layer_norm,
                     sequence_parallel=config.sequence_parallel,
                     apply_layernorm_1p=args.apply_layernorm_1p,
-                    mem_efficient_ln=not args.disable_mem_efficient_ln)
+                    mem_efficient_ln=args.mem_efficient_ln)
             else:
                 self.post_attention_layernorm = LayerNorm(
                     config.hidden_size,
@@ -970,7 +970,7 @@ def __init__(self, config,
                     no_persist_layer_norm=not config.persist_layer_norm,
                     sequence_parallel=config.sequence_parallel,
                     apply_layernorm_1p=args.apply_layernorm_1p,
-                    mem_efficient_ln=not args.disable_mem_efficient_ln)
+                    mem_efficient_ln=args.mem_efficient_ln)
             else:
                 self.post_inter_attention_layernorm = MixedFusedRMSNorm(config.hidden_size, config.layernorm_epsilon)
 
@@ -1730,7 +1730,7 @@ def build_layer(layer_number, n_e):
                         no_persist_layer_norm=args.no_persist_layer_norm,
                         sequence_parallel=config.sequence_parallel,
                         apply_layernorm_1p=args.apply_layernorm_1p,
-                        mem_efficient_ln=not args.disable_mem_efficient_ln)
+                        mem_efficient_ln=args.mem_efficient_ln)
                 else:
                     self.final_layernorm = LayerNorm(
                         config.hidden_size,