Merge branch 'moe_distckpt_compatibility' into 'main'

Fix moe dist-ckpt compatibility for !2230 See merge request ADLR/megatron-lm!2449
NVIDIA · Dec 13, 2024 · 40db706 · 40db706
2 parents 99f23d2 + 17b92eb
commit 40db706
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
@@ -576,7 +576,10 @@ def sharded_state_dict(
         non_homogeneous_layers = metadata is not None and metadata.get(
             'non_homogeneous_layers', False
         )
-        if self.config.num_moe_experts is not None:
+        if isinstance(self.config.moe_layer_freq, int):
+            if self.config.moe_layer_freq > 1:
+                non_homogeneous_layers = True
+        elif isinstance(self.config.moe_layer_freq, list):
             non_homogeneous_layers = True
 
         sharded_state_dict = {}