switch to MixedFusedLayerNorm (#262)

Author: stas00

megatron/mpu/layers.py

@@ -36,6 +36,7 @@
 from .utils import divide
 from .utils import split_tensor_along_last_dim
 from .utils import VocabUtility
+from ..model.fused_layer_norm import MixedFusedLayerNorm as LayerNorm
 from megatron import get_args, mpu
 import deepspeed.runtime.activation_checkpointing.checkpointing as ds_checkpointing
 
@@ -191,7 +192,7 @@ def __init__(self, num_embeddings, embedding_dim,
         # only the first stage embedding runs this class' forward. The head's embedding does its own
         # thing, so don't waste memory allocating LN weights.
         if mpu.is_pipeline_first_stage() and (args.use_bnb_optimizer or args.embed_layernorm):
-            self.norm = torch.nn.LayerNorm(embedding_dim)
+            self.norm = LayerNorm(embedding_dim)
 
         if args.use_bnb_optimizer:
             # for BNB we ignore the passed init_method and use torch.nn.init.xavier_uniform_