Merge branch 'fix_checkpointing_with_te' into 'main'

Fix TE checkpoint loading

See merge request ADLR/megatron-lm!1220

megatron/checkpointing.py

@@ -677,7 +677,7 @@ def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', stri
         print_rank_0('could not find arguments in the checkpoint ...')
 
     # Model.
-    strict = False if args.retro_add_retriever or args.transformer_impl == 'transformer_engine' else strict
+    strict = False if args.retro_add_retriever else strict
     if len(model) == 1:
         model[0].load_state_dict(state_dict['model'], strict=strict)
     else:

megatron/model/transformer.py

@@ -1802,6 +1802,10 @@ def load_state_dict(self, state_dict, strict=True):
         # Handle renaming layernorm -> norm in component names
         state_dict_ = {}
         for key in state_dict.keys():
+            # Bypass TransformerEngine module parameters.
+            if "layernorm_qkv" in key or "layernorm_mlp" in key:
+                state_dict_[key] = state_dict[key]
+                continue
             newkey = key.replace("layernorm", "norm")
             state_dict_[newkey] = state_dict[key]