[model] fix: map optional Step3.7 vision parameters (#5558)

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

Author: yaoyu-33

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

@@ -500,6 +500,8 @@ def mapping_registry(self) -> MegatronMappingRegistry:
         vision_norm_param_mappings = {
             "vision_model.ln_pre.weight": "vision_model.ln_pre.weight",
             "vision_model.ln_pre.bias": "vision_model.ln_pre.bias",
+            "vision_model.ln_post.weight": "vision_model.ln_post.weight",
+            "vision_model.ln_post.bias": "vision_model.ln_post.bias",
             "vision_model.transformer.resblocks.*.ln_1.weight": "vision_model.transformer.resblocks.*.ln_1.weight",
             "vision_model.transformer.resblocks.*.ln_1.bias": "vision_model.transformer.resblocks.*.ln_1.bias",
             "vision_model.transformer.resblocks.*.ln_2.weight": "vision_model.transformer.resblocks.*.ln_2.weight",
@@ -519,6 +521,7 @@ def mapping_registry(self) -> MegatronMappingRegistry:
             "vision_model.vit_downsampler2.bias": "vision_model.vit_downsampler2.bias",
             # Top-level nn.Parameter on Step37VisionModel — AutoMapping
             # would see the top-level vision class as the owning module.
+            "vision_model.class_embedding": "vision_model.class_embedding",
             "vision_model.positional_embedding": "vision_model.positional_embedding",
             # Attention — in_proj_* are bare nn.Parameter on
             # EncoderVisionAttention; out_proj.* are plain nn.Linear.

tests/unit_tests/models/stepfun/test_step37_bridge.py

@@ -21,10 +21,11 @@
 
 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.configuration_step37 import Step37Config, Step37VisionConfig
 from megatron.bridge.models.stepfun.modelling_step37 import transformer_block as _step37_block_mod
 from megatron.bridge.models.stepfun.modelling_step37.image_insert_embedding import ImageInsertEmbedding
 from megatron.bridge.models.stepfun.modelling_step37.transformer_block import get_step37_text_layer_spec
+from megatron.bridge.models.stepfun.modelling_step37.vision_model import Step37VisionModel
 from megatron.bridge.models.stepfun.step35_bridge import build_step35_layer_spec
 from megatron.bridge.models.stepfun.step37_bridge import Step37Bridge
 
@@ -121,6 +122,53 @@ def test_projector_bias_schema_roundtrips(self):
         ) == ("vit_large_projector.bias", True)
 
 
+class TestStep37VisionMappings:
+    def test_optional_vision_parameters_have_checkpoint_mappings(self):
+        vision_config = Step37VisionConfig(
+            width=8,
+            layers=0,
+            heads=1,
+            image_size=14,
+            patch_size=14,
+            mlp_ratio=1,
+            use_cls_token=True,
+            use_ln_pre=False,
+            use_ln_post=True,
+        )
+        vision_model = Step37VisionModel(vision_config)
+        parameter_names = set(dict(vision_model.named_parameters()))
+        assert {"class_embedding", "ln_post.weight", "ln_post.bias"} <= parameter_names
+
+        disabled_model = Step37VisionModel(
+            Step37VisionConfig(
+                width=8,
+                layers=0,
+                heads=1,
+                image_size=14,
+                patch_size=14,
+                mlp_ratio=1,
+                use_cls_token=False,
+                use_ln_pre=False,
+                use_ln_post=False,
+            )
+        )
+        disabled_parameter_names = set(dict(disabled_model.named_parameters()))
+        assert "class_embedding" not in disabled_parameter_names
+        assert "ln_post.weight" not in disabled_parameter_names
+        assert "ln_post.bias" not in disabled_parameter_names
+
+        bridge = Step37Bridge()
+        bridge.hf_config = SimpleNamespace(
+            text_config=SimpleNamespace(num_hidden_layers=0, num_nextn_predict_layers=0)
+        )
+        registry = bridge.mapping_registry()
+        for parameter_name in ("class_embedding", "ln_post.weight", "ln_post.bias"):
+            megatron_name = f"vision_model.{parameter_name}"
+            mapping = registry.megatron_to_hf_lookup(megatron_name)
+            assert mapping is not None
+            assert mapping.hf_param == megatron_name
+
+
 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."""