[model] fix: preserve Step3.7 MTP export config (#5528)

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/models/stepfun/step37_bridge.py

@@ -112,6 +112,15 @@ class Step37Bridge(MegatronModelBridge):
     # uses the identical HF schema.
     CONFIG_MAPPING = Step35Bridge.CONFIG_MAPPING
 
+    @classmethod
+    def megatron_to_hf_config(cls, provider: Step37ModelProvider) -> dict:
+        """Convert a Step3.7 provider into its nested Hugging Face config."""
+        hf_config = super().megatron_to_hf_config(provider)
+        mtp_num_layers = hf_config.pop("num_nextn_predict_layers", None)
+        if mtp_num_layers is not None:
+            hf_config["text_config"] = {"num_nextn_predict_layers": mtp_num_layers}
+        return hf_config
+
     def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> Step37ModelProvider:
         """Convert a HuggingFace Step3.7 config into a :class:`Step37ModelProvider`.
 

tests/unit_tests/models/stepfun/test_step37_bridge.py

@@ -17,7 +17,9 @@
 from types import SimpleNamespace
 from unittest.mock import patch
 
+from megatron.bridge.models.conversion.utils import conform_config_to_reference
 from megatron.bridge.models.stepfun import step37_bridge as _step37_bridge_mod
+from megatron.bridge.models.stepfun.configuration_step37 import Step37Config
 from megatron.bridge.models.stepfun.modelling_step37 import transformer_block as _step37_block_mod
 from megatron.bridge.models.stepfun.modelling_step37.transformer_block import get_step37_text_layer_spec
 from megatron.bridge.models.stepfun.step35_bridge import build_step35_layer_spec
@@ -83,6 +85,18 @@ def test_layer_spec_target_passes_checkpoint_load_validation(self):
         _validate_target_prefix(target=target, full_key="transformer_layer_spec")
 
 
+class TestStep37BridgeReverseConfig:
+    def test_checkpoint_mtp_override_is_nested_in_exported_config(self):
+        reference_config = Step37Config(text_config={"num_nextn_predict_layers": 3})
+        provider = SimpleNamespace(mtp_num_layers=0)
+
+        checkpoint_config = Step37Bridge.megatron_to_hf_config(provider)
+        exported_config = conform_config_to_reference(checkpoint_config, reference_config.to_dict())
+        synthesized_config = Step37Config(**exported_config)
+
+        assert synthesized_config.text_config.num_nextn_predict_layers == 0
+
+
 class TestGetStep37TextLayerSpec:
     def test_delegates_to_step35_builder(self):
         """Step3.7's text decoder spec is Step-3.5's hybrid layer spec, passed through verbatim."""