fix(ckpt): handle missing newer MLM checkpoint args (#5715)

Signed-off-by: caohch-1 <caohch2023@shanghaitech.edu.cn>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>

Author: caohch-1

src/megatron/bridge/training/mlm_compat/arguments.py

@@ -69,11 +69,13 @@ def _transformer_config_from_args(
     args: argparse.Namespace, config_class: type[TransformerConfig] = TransformerConfig
 ) -> TransformerConfig:
     """Build a variant of TransformerConfig based on contents of the MLM argparse args object."""
-    if args.multi_latent_attention:
+    multi_latent_attention = getattr(args, "multi_latent_attention", False)
+    if multi_latent_attention:
         config_class = MLATransformerConfig
 
-    if args.heterogeneous_layers_config_path is not None:
-        assert not args.multi_latent_attention, "Multi latent attention with heterogeneous layers is not supported."
+    heterogeneous_layers_config_path = getattr(args, "heterogeneous_layers_config_path", None)
+    if heterogeneous_layers_config_path is not None:
+        assert not multi_latent_attention, "Multi latent attention with heterogeneous layers is not supported."
         config_class = HeterogeneousTransformerConfig
 
     # Translate args to core transformer configuration
@@ -90,10 +92,13 @@ def _transformer_config_from_args(
     kw_args["pipeline_dtype"] = args.params_dtype
     kw_args["batch_p2p_comm"] = not args.overlap_p2p_comm
     kw_args["num_moe_experts"] = args.num_experts
-    kw_args["rotary_interleaved"] = args.rotary_interleaved
-    kw_args["num_layers_in_first_pipeline_stage"] = args.decoder_first_pipeline_num_layers
-    kw_args["num_layers_in_last_pipeline_stage"] = args.decoder_last_pipeline_num_layers
-    kw_args["fp8_param"] = args.fp8_param_gather
+    if hasattr(args, "decoder_first_pipeline_num_layers"):
+        kw_args["num_layers_in_first_pipeline_stage"] = args.decoder_first_pipeline_num_layers
+    if hasattr(args, "decoder_last_pipeline_num_layers"):
+        kw_args["num_layers_in_last_pipeline_stage"] = args.decoder_last_pipeline_num_layers
+    if hasattr(args, "fp8_param_gather"):
+        kw_args["fp8_param"] = args.fp8_param_gather
+
     if args.swiglu:
         kw_args["activation_func"] = F.silu
         kw_args["gated_linear_unit"] = True
@@ -109,12 +114,10 @@ def _transformer_config_from_args(
         kw_args["num_query_groups"] = args.num_query_groups
     else:
         kw_args["num_query_groups"] = None
-    kw_args["config_logger_dir"] = args.config_logger_dir
 
-    if len(args.cp_comm_type) == 1:
-        kw_args["cp_comm_type"] = args.cp_comm_type[0]
-    if args.is_hybrid_model:
-        kw_args["is_hybrid_model"] = args.is_hybrid_model
+    cp_comm_type = getattr(args, "cp_comm_type", ["p2p"])
+    if len(cp_comm_type) == 1:
+        kw_args["cp_comm_type"] = cp_comm_type[0]
 
     # handle quantization config
     # NOTE: Kitchen arguments are only added to the namespace when

tests/unit_tests/training/mlm_compat/test_arguments.py

@@ -173,6 +173,51 @@ def basic_args(self):
 
         return args
 
+    @pytest.mark.parametrize(
+        ("missing_arg", "config_attr", "expected"),
+        [
+            ("multi_latent_attention", "multi_latent_attention", False),
+            ("rotary_interleaved", "rotary_interleaved", False),
+            ("fp8_param_gather", "fp8_param", False),
+            ("config_logger_dir", "config_logger_dir", ""),
+            ("cp_comm_type", "cp_comm_type", "p2p"),
+            ("is_hybrid_model", "is_hybrid_model", False),
+        ],
+    )
+    def test_missing_newer_arg_uses_legacy_default(self, basic_args, missing_arg, config_attr, expected):
+        delattr(basic_args, missing_arg)
+
+        cfg = _transformer_config_from_args(basic_args)
+
+        assert getattr(cfg, config_attr) == expected
+
+    def test_missing_heterogeneous_layers_config_path_uses_default(self, basic_args):
+        delattr(basic_args, "heterogeneous_layers_config_path")
+
+        cfg = _transformer_config_from_args(basic_args)
+
+        assert type(cfg) is TransformerConfig
+
+    def test_missing_uneven_pipeline_args_uses_legacy_defaults(self, basic_args):
+        delattr(basic_args, "decoder_first_pipeline_num_layers")
+        delattr(basic_args, "decoder_last_pipeline_num_layers")
+
+        cfg = _transformer_config_from_args(basic_args)
+
+        assert cfg.num_layers_in_first_pipeline_stage is None
+        assert cfg.num_layers_in_last_pipeline_stage is None
+
+    def test_preserves_explicit_newer_args(self, basic_args):
+        basic_args.rotary_interleaved = True
+        basic_args.config_logger_dir = "/tmp/config-logger"
+        basic_args.is_hybrid_model = True
+
+        cfg = _transformer_config_from_args(basic_args)
+
+        assert cfg.rotary_interleaved is True
+        assert cfg.config_logger_dir == "/tmp/config-logger"
+        assert cfg.is_hybrid_model is True
+
     def test_basic_transformer_config(self, basic_args):
         """Test basic transformer config creation."""
         cfg = _transformer_config_from_args(basic_args)