fix(conversion): preserve legacy model construction paths

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

Author: yaoyu-33

examples/conversion/compare_hf_and_megatron/compare.py

@@ -716,7 +716,7 @@ def _load_megatron_model(args):
             ),
             **_hf_revision_kwargs(args.hf_revision),
         )
-        if getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None:
+        if getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False):
             model_config = bridge.get_model_config()
             transformer = model_config.transformer
             transformer.tensor_model_parallel_size = tp
@@ -756,7 +756,7 @@ def _load_megatron_model(args):
             ),
             **_hf_revision_kwargs(args.hf_revision),
         )
-        if getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None:
+        if getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False):
             model_config = bridge.get_model_config()
             transformer = model_config.transformer
             transformer.tensor_model_parallel_size = tp

scripts/conversion/gpu_backend.py

@@ -178,7 +178,7 @@ def _configure_model_provider(
 
 def _uses_model_builder(bridge: AutoBridge) -> bool:
     """Return whether the selected bridge supports native builder construction."""
-    return getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None
+    return getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False)
 
 
 def _configure_model_config(

scripts/inference/vlm_generation.py

@@ -151,7 +151,7 @@ def _hf_revision_kwargs(revision: str | None) -> dict[str, str]:
 
 def _uses_model_builder(bridge: AutoBridge) -> bool:
     """Return whether the selected bridge supports native builder construction."""
-    return getattr(bridge._model_bridge, "MODEL_CONFIG_CLASS", None) is not None
+    return getattr(bridge._model_bridge, "USE_MODEL_CONFIG_FOR_CONVERSION", False)
 
 
 def _enable_deterministic_execution() -> None:

src/megatron/bridge/models/conversion/auto_bridge.py

@@ -1484,7 +1484,7 @@ def import_ckpt(
         with model_context:
             # Prefer the native ModelConfig/ModelBuilder path for migrated model
             # families while preserving the provider path for legacy bridges.
-            if bridge._model_bridge.MODEL_CONFIG_CLASS is not None:
+            if bridge._model_bridge.USE_MODEL_CONFIG_FOR_CONVERSION:
                 model_config = bridge.get_model_config()
                 model_config.transformer.use_cpu_initialization = True
                 megatron_model = bridge.get_model(

src/megatron/bridge/models/conversion/model_bridge.py

@@ -441,6 +441,12 @@ def mapping_registry(self) -> MegatronMappingRegistry:
     # transformer config class.
     MODEL_CONFIG_CLASS: ClassVar[type[ModelConfig] | None] = BridgeGPTModelConfig
 
+    # Conversion still uses the provider-backed construction path unless a model family has
+    # explicitly migrated its conversion lifecycle to ModelBuilder. Keeping this separate
+    # from MODEL_CONFIG_CLASS preserves get_model_config() for legacy families without
+    # silently changing how their checkpoints are constructed.
+    USE_MODEL_CONFIG_FOR_CONVERSION: ClassVar[bool] = False
+
     # Leave unset unless HF export must copy nonstandard files in addition to the usual artifacts,
     # for example ``["*reasoning_parser.py"]``.
     ADDITIONAL_FILE_PATTERNS = None

src/megatron/bridge/models/conversion/utils.py

@@ -256,9 +256,10 @@ def remove_non_pickleables(obj, max_depth: int = 3, current_depth: int = 0):
         # Create a copy to avoid modifying the original
         cleaned_obj = copy.copy(obj)
 
-        for attr_name in list(vars(cleaned_obj).keys()):
-            attr_value = getattr(cleaned_obj, attr_name)
-
+        # Read stored attributes from ``__dict__`` directly. Configuration classes may
+        # deliberately reject dynamic attribute access for values whose meaning is
+        # layer-dependent, even though the raw value still needs to be copied for IPC.
+        for attr_name, attr_value in list(vars(cleaned_obj).items()):
             # Recursively clean attribute
             cleaned_value = remove_non_pickleables(attr_value, max_depth, current_depth + 1)
 

src/megatron/bridge/models/muse_glimmer/muse_glimmer_bridge.py

@@ -139,6 +139,7 @@ class MuseGlimmerBridge(MegatronModelBridge):
     """Builder-backed bridge for the complete Muse Glimmer multimodal model."""
 
     MODEL_CONFIG_CLASS = MuseGlimmerModelConfig
+    USE_MODEL_CONFIG_FOR_CONVERSION = True
 
     @staticmethod
     def _validate_architecture(hf_config: PretrainedConfig) -> tuple[Any, Any]:

src/megatron/bridge/training/checkpointing.py

@@ -177,8 +177,12 @@ class _CpuTorchDistSaveShardedStrategy(TorchDistSaveShardedStrategy):
     """Run MCore's synchronous torch-dist writer without a CUDA staging barrier."""
 
     @staticmethod
-    def _run_cpu_finalize(finalize_fn: Callable[[], None]) -> None:
-        """Make MCore's failure-status tensor use the CPU for a Gloo save."""
+    def _run_finalize(finalize_fn: Callable[[], None]) -> None:
+        """Use a backend-compatible device for MCore's failure-status tensor."""
+        if torch.distributed.is_initialized() and torch.distributed.get_backend() == "nccl":
+            finalize_fn()
+            return
+
         current_device = torch.cuda.current_device
         try:
             torch.cuda.current_device = lambda: torch.device("cpu")
@@ -204,7 +208,7 @@ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> No
                 preload_fn=partial(preload_fn.func, *bound.args, **bound.kwargs),
             )
         async_request = async_request._replace(
-            finalize_fns=[partial(self._run_cpu_finalize, finalize_fn) for finalize_fn in async_request.finalize_fns],
+            finalize_fns=[partial(self._run_finalize, finalize_fn) for finalize_fn in async_request.finalize_fns],
         )
         async_request.execute_sync()
 

tests/functional_tests/test_groups/models/nemotron_omni/test_nemotron_omni_conversion.py

@@ -40,6 +40,15 @@
 # 52-layer pattern, which keeps a representative mamba (M), MoE-MLP (E) and
 # attention (*) layer mix.
 _LLM_LAYER_TYPES = ["mamba", "moe", "mamba", "moe", "mamba", "attention"]
+# Transformers 5.15 serializes the same hybrid blocks under their canonical names.
+_SERIALIZED_LLM_LAYER_TYPES = [
+    "linear_attention",
+    "moe",
+    "linear_attention",
+    "moe",
+    "linear_attention",
+    "full_attention",
+]
 _LLM_OVERRIDES = {
     "hybrid_override_pattern": "MEMEM*",
     "layer_types": _LLM_LAYER_TYPES,
@@ -220,5 +229,5 @@ def test_nemotron_omni_conversion_roundtrip(self, nemotron_omni_toy_model_path,
         assert "vision_config" in saved_config
         assert "sound_config" in saved_config
         assert saved_config["llm_config"]["num_hidden_layers"] == 6
-        assert saved_config["llm_config"]["layer_types"] == _LLM_LAYER_TYPES
-        assert saved_config["llm_config"]["layers_block_type"] == _LLM_LAYER_TYPES
+        assert saved_config["llm_config"]["layer_types"] == _SERIALIZED_LLM_LAYER_TYPES
+        assert saved_config["llm_config"]["layers_block_type"] == _SERIALIZED_LLM_LAYER_TYPES

tests/unit_tests/conversion/launcher/test_gpu_backend.py

@@ -115,6 +115,7 @@ def provide_distributed_model(self, *args, **kwargs):
 
 class _FakeModelBridge:
     MODEL_CONFIG_CLASS = None
+    USE_MODEL_CONFIG_FOR_CONVERSION = False
 
     def get_hf_tokenizer_kwargs(self):
         return {"padding_side": "left"}
@@ -191,6 +192,7 @@ def test_import_uses_builder_for_migrated_model(self, cli, monkeypatch):
 
         class BuilderModelBridge(_FakeModelBridge):
             MODEL_CONFIG_CLASS = object
+            USE_MODEL_CONFIG_FOR_CONVERSION = True
 
         class FakeBridge:
             _model_bridge = BuilderModelBridge()