fused layernorm (#1105)

* Add simple util for CUDA timings

* Add fused layernorm kernel from Megatron

Closes #952

* change default fused layernorm to false

* Update test_setup.yml

* Update test_train_base.yml

---------

Co-authored-by: Yang Zhang <[email protected]>
Co-authored-by: jahatef <[email protected]>
Co-authored-by: Jacob Hatef <[email protected]>

configs/1-3B.yml

@@ -21,6 +21,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
    # init methods
    "init_method": "small_init",

configs/125M-json.yml

@@ -16,6 +16,7 @@
   "scaled_upper_triang_masked_softmax_fusion": false,
   "bias_gelu_fusion": false,
   "rope_fusion": false,
+  "layernorm_fusion": false,
 
   "init_method": "small_init",
   "output_layer_init_method": "wang_init",

configs/125M.yml

@@ -21,6 +21,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
    # init methods
    "init_method": "small_init",

configs/13B.yml

@@ -21,6 +21,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
    # init methods
    "init_method": "small_init",

configs/175B.yml

@@ -21,6 +21,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
    # init methods
    "init_method": "small_init",

configs/19M.yml

@@ -16,6 +16,7 @@
   "scaled_upper_triang_masked_softmax_fusion": false,
   "bias_gelu_fusion": false,
   "rope_fusion": false,
+  "layernorm_fusion": false,
 
   # init methods
   "init_method": "small_init",

configs/2-7B.yml

@@ -21,6 +21,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
    # init methods
    "init_method": "small_init",

configs/20B.yml

@@ -31,6 +31,7 @@
   "scaled_upper_triang_masked_softmax_fusion": true,
   "bias_gelu_fusion": true,
   "rope_fusion": false,
+  "layernorm_fusion": false,
 
   # init methods
   "init_method": "small_init",

configs/350M.yml

@@ -21,6 +21,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
    # init methods
    "init_method": "small_init",

configs/49M.yml

@@ -19,6 +19,7 @@
   "scaled_upper_triang_masked_softmax_fusion": false,
   "bias_gelu_fusion": false,
   "rope_fusion": false,
+  "layernorm_fusion": false,
 
   # init methods
   "init_method": "small_init",

configs/6-7B.yml

@@ -21,6 +21,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
    # init methods
    "init_method": "small_init",

configs/760M.yml

@@ -21,6 +21,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
    # init methods
    "init_method": "small_init",

configs/800M.yml

@@ -16,6 +16,7 @@
   "scaled_upper_triang_masked_softmax_fusion": false,
   "bias_gelu_fusion": false,
   "rope_fusion": false,
+  "layernorm_fusion": false,
 
   # init methods
   "init_method": "small_init",

configs/bf16_125M.yml

@@ -19,6 +19,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
 
    # optimizer settings

configs/bnb_125M.yml

@@ -20,6 +20,7 @@
    "scaled_upper_triang_masked_softmax_fusion": false,
    "bias_gelu_fusion": false,
    "rope_fusion": false,
+   "layernorm_fusion": false,
 
 
    # optimizer settings

configs/slurm_125M.yml

@@ -12,6 +12,7 @@
    "scaled_upper_triang_masked_softmax_fusion": true,
    "bias_gelu_fusion": true,
    "rope_fusion": false,
+   "layernorm_fusion": false,
    "optimizer": {
      "type": "Adam",
      "params": {

megatron/devutil.py

@@ -0,0 +1,51 @@
+import torch.cuda
+
+
+class Metric:
+    """
+    Dumb utility to collect and report average wall-time metrics.
+    """
+
+    def __init__(self, label):
+        self.label = label
+        self.measurements = []
+
+    def collect(self, measurement):
+        self.measurements.append(measurement)
+
+    def get_measurements(self):
+        return self.measurements[:]
+
+    def report(self):
+        print(
+            self.label,
+            torch.quantile(torch.tensor(self.measurements), torch.arange(10) / 10.0),
+        )
+
+
+def monitor_method_cuda_wall_times(metric, obj, methodname):
+    """
+    Measure timings for a method on an object or class.
+
+    For instance:
+
+    >>> metric = Metric('!LNORM')
+    >>> monitor_method_wall_times(metric, LayerNorm, 'forward')
+    """
+    oldmeth = getattr(obj, methodname)
+
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+
+    def newmeth(*args, **kw):
+        start_event.record()
+        try:
+            return oldmeth(*args, **kw)
+        finally:
+            end_event.record()
+            torch.cuda.synchronize()
+            elapsed = start_event.elapsed_time(end_event)
+            metric.collect(elapsed)
+            metric.report()
+
+    setattr(obj, methodname, newmeth)

megatron/model/fused_layer_norm.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""This code is copied from NVIDIA apex:
+      https://github.com/NVIDIA/apex
+   with some changes. """
+
+import numbers
+import torch
+from torch.nn.parameter import Parameter
+from torch.nn import init
+import importlib
+from torch.nn import functional as F
+import inspect
+
+from megatron.utils import make_viewless_tensor
+
+try:
+    from apex.contrib.layer_norm.layer_norm import FastLayerNormFN
+
+    HAVE_PERSIST_LAYER_NORM = True
+except:
+    HAVE_PERSIST_LAYER_NORM = False
+
+from apex.normalization.fused_layer_norm import FusedLayerNormAffineFunction
+
+
+global fused_layer_norm_cuda
+fused_layer_norm_cuda = None
+
+
+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,
+        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")
+
+        # List of hiddens sizes supported in the persistent layer norm kernel
+        # If the hidden size is not supported, fall back to the non-persistent
+        # kernel.
+        persist_ln_hidden_sizes = [
+            1024,
+            1536,
+            2048,
+            2304,
+            3072,
+            3840,
+            4096,
+            5120,
+            6144,
+            8192,
+            10240,
+            12288,
+            12800,
+            15360,
+            16384,
+            18432,
+            20480,
+            24576,
+            25600,
+            30720,
+            32768,
+            40960,
+            49152,
+            65536,
+        ]
+        if (
+            normalized_shape not in persist_ln_hidden_sizes
+            or not HAVE_PERSIST_LAYER_NORM
+        ):
+            no_persist_layer_norm = True
+
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        self.normalized_shape = torch.Size(normalized_shape)
+        self.eps = eps
+        self.weight = Parameter(torch.Tensor(*normalized_shape))
+        self.bias = Parameter(torch.Tensor(*normalized_shape))
+        self.reset_parameters()
+        self.no_persist_layer_norm = no_persist_layer_norm
+        self.sequence_parallel = sequence_parallel
+
+        # set sequence parallelism flag on weight and bias parameters
+        setattr(self.weight, "sequence_parallel", self.sequence_parallel)
+        setattr(self.bias, "sequence_parallel", self.sequence_parallel)
+
+    def reset_parameters(self):
+
+        if self.apply_layernorm_1p:
+            init.zeros_(self.weight)
+            init.zeros_(self.bias)
+        else:
+            init.ones_(self.weight)
+            init.zeros_(self.bias)
+
+    def forward(self, input):
+
+        weight = self.weight + 1 if self.apply_layernorm_1p else self.weight
+        # CPU path is here for unittest sake.
+        if not input.is_cuda:
+            print(
+                "WARNING! The input of FusedLayerNorm should be on the GPU."
+                "This warning should only be triggered in the FusedLayerNorm unit tests."
+            )
+            return F.layer_norm(
+                input, self.normalized_shape, weight, self.bias, self.eps
+            )
+
+        if self.no_persist_layer_norm:
+            # 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)
+
+            # Apex's fast layer norm function outputs a 'view' tensor (i.e., has
+            # a populated '_base' field). This will result in schedule.py's
+            # deallocate_output_tensor() throwing an error, so a viewless tensor is
+            # created to prevent this.
+            output = make_viewless_tensor(
+                inp=output, requires_grad=input.requires_grad, keep_graph=True
+            )
+
+            return output

megatron/model/norms.py

@@ -14,6 +14,7 @@
 
 import torch
 from torch.nn import LayerNorm as LayerNorm
+from .fused_layer_norm import MixedFusedLayerNorm
 
 
 def get_norm(neox_args):
@@ -22,7 +23,7 @@ def get_norm(neox_args):
         eps = neox_args.rms_norm_epsilon
     elif neox_args.norm == "layernorm":
         eps = neox_args.layernorm_epsilon
-        norm = LayerNorm
+        norm = MixedFusedLayerNorm if neox_args.layernorm_fusion else LayerNorm
     elif neox_args.norm == "scalenorm":
         eps = neox_args.scalenorm_epsilon
         norm = ScaleNorm

megatron/neox_arguments/neox_args.py

@@ -125,6 +125,11 @@ class NeoXArgsModel(NeoXArgsTemplate):
     Normalization layer to use. Choose from "layernorm", "rmsnorm", "scalenorm".
     """
 
+    layernorm_fusion: bool = False
+    """
+    Use fused layer norm kernel (if `norm` is `layernorm`).
+    """
+
     use_qk_layernorm: bool = False
     """
     Use QK Normalization