Introduce LayerNorm optimization from latest Apex (bigscience-workshop#277)

* Introduce LayerNorm optimization from NVIDIA/apex#1715

* Fix args call

* Ad-hoc apex version check

* Remove unnecessary TransformerConfig arg

Author: Quentin-Anthony

megatron/arguments.py

@@ -626,6 +626,9 @@ def _add_network_size_args(parser):
     group.add_argument('--apply-layernorm-1p', action='store_true',
                        help='Adjust LayerNorm weights such that they are centered '
                        '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')
     group.add_argument('--apply-residual-connection-post-layernorm',
                        action='store_true',
                        help='If set, use original BERT residula connection '

megatron/model/fused_layer_norm.py

@@ -10,6 +10,7 @@
 from torch.nn import init
 import importlib
 from torch.nn import functional as F
+import inspect
 
 from megatron.core.utils import make_viewless_tensor
 
@@ -31,10 +32,12 @@ class MixedFusedLayerNorm(torch.nn.Module):
   def __init__(self, normalized_shape, eps=1e-5,
                no_persist_layer_norm=True,
                sequence_parallel=False,
-               apply_layernorm_1p=False):
+               apply_layernorm_1p=False,
+               mem_efficient_ln=True):
         super(MixedFusedLayerNorm, self).__init__()
 
         self.apply_layernorm_1p = apply_layernorm_1p
+        self.mem_efficient_ln = mem_efficient_ln
 
         global fused_layer_norm_cuda
         fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
@@ -83,7 +86,12 @@ def forward(self, input):
         return F.layer_norm(input, self.normalized_shape, weight, self.bias, self.eps)
 
     if self.no_persist_layer_norm:
-        return FusedLayerNormAffineFunction.apply(input, weight, self.bias, self.normalized_shape, self.eps)
+        # Apex does not have versions yet (https://github.com/NVIDIA/apex/pull/1648), so we need to inspect 
+        # the function manually on whether the extra arg introduced in https://github.com/NVIDIA/apex/pull/1715 exists yet
+        if 'memory_efficient' in inspect.getfullargspec(FusedLayerNormAffineFunction.forward).args:
+            return FusedLayerNormAffineFunction.apply(input, weight, self.bias, self.normalized_shape, self.eps, self.mem_efficient_ln)
+        else:
+            return FusedLayerNormAffineFunction.apply(input, weight, self.bias, self.normalized_shape, self.eps)
     else:
         output = FastLayerNormFN.apply(input, weight, self.bias, self.eps)
 

megatron/model/transformer.py

@@ -919,7 +919,8 @@ def __init__(self, config,
                     eps=config.layernorm_epsilon,
                     no_persist_layer_norm=args.no_persist_layer_norm,
                     sequence_parallel=config.sequence_parallel,
-                    apply_layernorm_1p=args.apply_layernorm_1p)
+                    apply_layernorm_1p=args.apply_layernorm_1p,
+                    mem_efficient_ln=not args.disable_mem_efficient_ln)
             else:
                 self.input_layernorm = LayerNorm(
                     config.hidden_size,
@@ -944,7 +945,8 @@ def __init__(self, config,
                     eps=config.layernorm_epsilon,
                     no_persist_layer_norm=not config.persist_layer_norm,
                     sequence_parallel=config.sequence_parallel,
-                    apply_layernorm_1p=args.apply_layernorm_1p)
+                    apply_layernorm_1p=args.apply_layernorm_1p,
+                    mem_efficient_ln=not args.disable_mem_efficient_ln)
             else:
                 self.post_attention_layernorm = LayerNorm(
                     config.hidden_size,
@@ -967,7 +969,8 @@ def __init__(self, config,
                     eps=config.layernorm_epsilon,
                     no_persist_layer_norm=not config.persist_layer_norm,
                     sequence_parallel=config.sequence_parallel,
-                    apply_layernorm_1p=args.apply_layernorm_1p)
+                    apply_layernorm_1p=args.apply_layernorm_1p,
+                    mem_efficient_ln=not args.disable_mem_efficient_ln)
             else:
                 self.post_inter_attention_layernorm = MixedFusedRMSNorm(config.hidden_size, config.layernorm_epsilon)
 
@@ -1726,7 +1729,8 @@ def build_layer(layer_number, n_e):
                         eps=config.layernorm_epsilon,
                         no_persist_layer_norm=args.no_persist_layer_norm,
                         sequence_parallel=config.sequence_parallel,
-                        apply_layernorm_1p=args.apply_layernorm_1p)
+                        apply_layernorm_1p=args.apply_layernorm_1p,
+                        mem_efficient_ln=not args.disable_mem_efficient_ln)
                 else:
                     self.final_layernorm = LayerNorm(
                         config.hidden_size,