refactor(model): build GLM-4.5 from a builder-backed model config

Move GLM-4.5 onto the ModelConfig/ModelBuilder path, following Muse Glimmer:

- GLM45Bridge maps Hugging Face configs in hf_config_to_model_config_kwargs().
  The builder's default MoE layer spec is the GPT decoder block spec, which
  honors moe_layer_freq; the deprecated provider path keeps its explicit
  decoder block spec. USE_MODEL_CONFIG_FOR_CONVERSION = True.
- All six GLM-4.5 H100 recipes use get_model_config().

Signed-off-by: Kamran Jafari <kjafarisadeg@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: kamran-nvidia

src/megatron/bridge/models/glm/glm45_bridge.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 
 from functools import partial
+from typing import Any
 
 import torch
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
 from megatron.core.models.gpt.gpt_model import GPTModel
-from transformers import Glm4MoeForCausalLM
+from transformers import Glm4MoeForCausalLM, PretrainedConfig
 
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
@@ -28,8 +29,6 @@
     GatedMLPMapping,
     QKVMapping,
 )
-from megatron.bridge.models.gpt_provider import GPTModelProvider
-from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
 
 
 try:
@@ -51,47 +50,58 @@ class GLM45Bridge(MegatronModelBridge):
     Example:
         >>> from megatron.bridge import AutoBridge
         >>> bridge = AutoBridge.from_hf_pretrained("zai-org/GLM-4.5")
-        >>> provider = bridge.to_megatron_provider()
+        >>> model_config = bridge.get_model_config()
     """
 
-    def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider:
-        """Convert HuggingFace config to GPTModelProvider."""
-        provider = super().provider_bridge(hf_pretrained)
-        hf_config = hf_pretrained.config
+    USE_MODEL_CONFIG_FOR_CONVERSION = True
 
-        # Use decoder block spec to properly handle moe_layer_freq (mixed dense/MoE layers)
-        provider.transformer_layer_spec = partial(get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE)
-        provider.normalization = "RMSNorm"
-        provider.gated_linear_unit = True
-        provider.add_bias_linear = False
-        provider.share_embeddings_and_output_weights = False
+    def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Convert a Hugging Face GLM-4.5 config to builder config kwargs.
 
-        provider.moe_shared_expert_overlap = True
-        provider.moe_token_dispatcher_type = "alltoall"
-        provider.moe_router_load_balancing_type = "seq_aux_loss"
-        provider.moe_router_pre_softmax = False
-        provider.moe_grouped_gemm = True
-        provider.moe_router_score_function = "sigmoid"
-        provider.moe_permute_fusion = True
-        provider.moe_router_enable_expert_bias = True
-        provider.moe_router_dtype = "fp32"
-        provider.moe_router_bias_update_rate = 0
-        provider.moe_aux_loss_coeff = 0.001
-
-        provider.persist_layer_norm = True
-        provider.bias_activation_fusion = True
-        provider.bias_dropout_fusion = True
-        provider.hidden_dropout = 0.0
-        provider.autocast_dtype = torch.bfloat16
-        provider.mtp_num_layers = getattr(hf_config, "num_nextn_predict_layers", None)
-        provider.mtp_loss_scaling_factor = 0.3
-        provider.moe_shared_expert_intermediate_size = hf_config.moe_intermediate_size
-
-        provider.moe_layer_freq = [0] * hf_config.first_k_dense_replace + [1] * (
-            hf_config.num_hidden_layers - hf_config.first_k_dense_replace
+        The builder's default MoE layer spec already uses the GPT decoder block
+        spec, which honors the mixed dense/MoE ``moe_layer_freq``.
+        """
+        config_kwargs = super().hf_config_to_model_config_kwargs(hf_config)
+        config_kwargs.update(
+            normalization="RMSNorm",
+            gated_linear_unit=True,
+            add_bias_linear=False,
+            share_embeddings_and_output_weights=False,
+            moe_shared_expert_overlap=True,
+            moe_token_dispatcher_type="alltoall",
+            moe_router_load_balancing_type="seq_aux_loss",
+            moe_router_pre_softmax=False,
+            moe_grouped_gemm=True,
+            moe_router_score_function="sigmoid",
+            moe_permute_fusion=True,
+            moe_router_enable_expert_bias=True,
+            moe_router_dtype="fp32",
+            moe_router_bias_update_rate=0,
+            moe_aux_loss_coeff=0.001,
+            persist_layer_norm=True,
+            bias_activation_fusion=True,
+            bias_dropout_fusion=True,
+            hidden_dropout=0.0,
+            autocast_dtype=torch.bfloat16,
+            masked_softmax_fusion=True,
+            rope_scaling=False,
+            rope_scaling_factor=1.0,
+            mtp_num_layers=getattr(hf_config, "num_nextn_predict_layers", None),
+            mtp_loss_scaling_factor=0.3,
+            moe_shared_expert_intermediate_size=hf_config.moe_intermediate_size,
+            moe_layer_freq=[0] * hf_config.first_k_dense_replace
+            + [1] * (hf_config.num_hidden_layers - hf_config.first_k_dense_replace),
         )
+        config_kwargs.setdefault("position_embedding_type", "rope")
+        return config_kwargs
 
-        return provider
+    def hf_config_to_provider_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
+        """Adapt the canonical builder mapping to the deprecated provider path."""
+        provider_kwargs = self.hf_config_to_model_config_kwargs(hf_config)
+        # The provider's default layer spec builds homogeneous layers; the decoder
+        # block spec is needed for mixed dense/MoE layers (moe_layer_freq).
+        provider_kwargs["transformer_layer_spec"] = partial(get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE)
+        return provider_kwargs
 
     def mapping_registry(self) -> MegatronMappingRegistry:
         mapping_list = []

src/megatron/bridge/recipes/glm/h100/glm45.py

@@ -32,7 +32,7 @@ def glm45_355b_pretrain_128gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _pretrain_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("zai-org/GLM-4.5").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("zai-org/GLM-4.5").get_model_config()
 
     # Tokenizer - uses NullTokenizer by default
     cfg.tokenizer.tokenizer_type = "NullTokenizer"
@@ -165,7 +165,7 @@ def glm45_air_106b_pretrain_32gpu_h100_bf16_config() -> ConfigContainer:
     cfg = _pretrain_common()
 
     # Model config
-    cfg.model = AutoBridge.from_hf_pretrained("zai-org/GLM-4.5-Air").to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained("zai-org/GLM-4.5-Air").get_model_config()
 
     # Tokenizer - uses NullTokenizer by default
     cfg.tokenizer.tokenizer_type = "NullTokenizer"
@@ -304,7 +304,7 @@ def glm45_355b_sft_128gpu_h100_bf16_config() -> ConfigContainer:
 
     # Model config from HuggingFace
     hf_path = "zai-org/GLM-4.5"
-    cfg.model = AutoBridge.from_hf_pretrained(hf_path).to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained(hf_path).get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = hf_path
@@ -441,7 +441,7 @@ def glm45_air_106b_sft_32gpu_h100_bf16_config() -> ConfigContainer:
 
     # Model config from HuggingFace
     hf_path = "zai-org/GLM-4.5-Air"
-    cfg.model = AutoBridge.from_hf_pretrained(hf_path).to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained(hf_path).get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = hf_path
@@ -588,7 +588,7 @@ def glm45_355b_peft_16gpu_h100_bf16_config(
 
     # Model config from HuggingFace
     hf_path = "zai-org/GLM-4.5"
-    cfg.model = AutoBridge.from_hf_pretrained(hf_path).to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained(hf_path).get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = hf_path
@@ -732,7 +732,7 @@ def glm45_air_106b_peft_8gpu_h100_bf16_config(
 
     # Model config from HuggingFace
     hf_path = "zai-org/GLM-4.5-Air"
-    cfg.model = AutoBridge.from_hf_pretrained(hf_path).to_megatron_provider(load_weights=False)
+    cfg.model = AutoBridge.from_hf_pretrained(hf_path).get_model_config()
 
     # Tokenizer
     cfg.tokenizer.tokenizer_model = hf_path

tests/unit_tests/models/glm/test_glm45_bridge.py

@@ -330,3 +330,84 @@ def test_dtype_mapping_default(self, glm45_355b_config):
         assert provider.fp16 is False
         assert provider.bf16 is False
         assert provider.params_dtype == torch.float32
+
+
+class TestGLM45BuilderConfig:
+    """Builder-backed ModelConfig path for GLM-4.5."""
+
+    @staticmethod
+    def _config(**overrides):
+        from transformers import Glm4MoeConfig
+
+        kwargs = dict(
+            hidden_size=64,
+            intermediate_size=128,
+            moe_intermediate_size=32,
+            num_hidden_layers=4,
+            num_attention_heads=8,
+            num_key_value_heads=4,
+            n_routed_experts=8,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            first_k_dense_replace=1,
+            max_position_embeddings=4096,
+            vocab_size=128,
+            torch_dtype="bfloat16",
+        )
+        kwargs.update(overrides)
+        config = Glm4MoeConfig(**kwargs)
+        config.architectures = ["Glm4MoeForCausalLM"]
+        config.num_nextn_predict_layers = 1
+        return config
+
+    def test_provider_bridge_is_inherited_compatibility_only(self):
+        assert "provider_bridge" not in GLM45Bridge.__dict__
+
+    def test_conversion_uses_builder_config(self):
+        assert GLM45Bridge.USE_MODEL_CONFIG_FOR_CONVERSION is True
+
+    def test_hf_config_to_model_config_uses_direct_mapping(self):
+        from unittest.mock import patch
+
+        from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
+
+        config = self._config()
+        bridge = GLM45Bridge()
+        with (
+            patch.object(bridge, "provider_bridge", side_effect=AssertionError("provider path used")),
+            patch.object(bridge, "hf_config_to_provider_kwargs", side_effect=AssertionError("provider kwargs used")),
+        ):
+            result = bridge.hf_config_to_model_config(config)
+
+        assert isinstance(result, BridgeGPTModelConfig)
+        # The builder's default MoE spec is the decoder block spec, so none is configured here.
+        assert result.transformer_layer_spec is None
+        assert result.moe_layer_freq == [0, 1, 1, 1]
+        assert result.moe_shared_expert_intermediate_size == config.moe_intermediate_size
+        assert result.moe_router_score_function == "sigmoid"
+        assert result.moe_router_enable_expert_bias is True
+        assert result.mtp_num_layers == 1
+        assert result.mtp_loss_scaling_factor == 0.3
+        assert result.position_embedding_type == "rope"
+
+    def test_model_config_matches_provider_runtime_config(self):
+        from dataclasses import fields
+
+        from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
+
+        config = self._config()
+        mock_pretrained = Mock(spec=PreTrainedCausalLM)
+        mock_pretrained.config = config
+        bridge = GLM45Bridge()
+
+        model_config = bridge.hf_config_to_model_config(config)
+        with pytest.warns(FutureWarning, match=r"deprecated.*get_model_config.*get_model"):
+            provider = bridge.provider_bridge(mock_pretrained)
+
+        provider_fields = {field.name for field in fields(provider)}
+        model_config_fields = {field.name for field in fields(model_config)}
+        model_config_fields.update(field.name for field in fields(model_config.transformer))
+        for field_name in sorted((provider_fields & model_config_fields) - {"transformer_layer_spec"}):
+            assert getattr(model_config, field_name) == getattr(provider, field_name), field_name
+        # The provider keeps its explicit decoder block spec for mixed dense/MoE layers.
+        assert provider.transformer_layer_spec.func is get_gpt_decoder_block_spec

tests/unit_tests/recipes/test_glm45_recipes.py

@@ -98,9 +98,12 @@ class _FakeBridge:
     def __init__(self):
         pass
 
-    def to_megatron_provider(self, load_weights: bool = False):
+    def get_model_config(self):
         return _FakeModelCfg()
 
+    def to_megatron_provider(self, load_weights: bool = False):
+        raise AssertionError("GLM-4.5 recipes must use get_model_config(), not the legacy provider API")
+
     @staticmethod
     def from_hf_pretrained(hf_path: str, **kwargs):
         return _FakeBridge()
@@ -112,7 +115,7 @@ class _RealDepthFakeBridge(_FakeBridge):
     def __init__(self, num_layers: int):
         self.num_layers = num_layers
 
-    def to_megatron_provider(self, load_weights: bool = False):
+    def get_model_config(self):
         model = _FakeModelCfg()
         model.num_layers = self.num_layers
         return model
@@ -546,3 +549,66 @@ def test_glm45_recompute_configuration(monkeypatch: pytest.MonkeyPatch):
     assert hasattr(cfg.model, "recompute_granularity")
     assert hasattr(cfg.model, "recompute_method")
     assert hasattr(cfg.model, "recompute_num_layers")
+
+
+class _BuilderOnlyBridge:
+    """Return a real strict ModelConfig at the recipe's GLM-4.5 depth while rejecting the provider API."""
+
+    def __init__(self, num_layers: int):
+        self.num_layers = num_layers
+
+    @staticmethod
+    def from_hf_pretrained(hf_path: str, **kwargs) -> "_BuilderOnlyBridge":
+        return _BuilderOnlyBridge(46 if hf_path.endswith("Air") else 92)
+
+    def get_model_config(self):
+        from transformers import Glm4MoeConfig
+
+        from megatron.bridge import AutoBridge
+
+        config = Glm4MoeConfig(
+            hidden_size=64,
+            intermediate_size=128,
+            moe_intermediate_size=32,
+            num_hidden_layers=self.num_layers,
+            num_attention_heads=8,
+            num_key_value_heads=4,
+            n_routed_experts=16,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            first_k_dense_replace=1,
+            max_position_embeddings=131072,
+            vocab_size=151552,
+        )
+        config.architectures = ["Glm4MoeForCausalLM"]
+        config.num_nextn_predict_layers = 1
+        return AutoBridge.from_hf_config(config).get_model_config()
+
+    def to_megatron_provider(self, load_weights: bool = False):
+        raise AssertionError("GLM-4.5 recipes must not call the legacy provider API")
+
+
+@pytest.mark.parametrize("recipe_func", _GLM45_RECIPE_FUNCS, ids=lambda func: func.__name__)
+def test_each_glm45_recipe_uses_strict_builder_config(recipe_func: Callable, monkeypatch: pytest.MonkeyPatch):
+    """Every GLM-4.5 recipe should configure a ModelConfig without provider fallback."""
+    from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
+
+    patch_recipe_module_global(monkeypatch, recipe_func, "AutoBridge", _BuilderOnlyBridge)
+    # Avoid HF I/O for SFT/PEFT tokenizers.
+    import transformers
+
+    monkeypatch.setattr(
+        transformers,
+        "AutoTokenizer",
+        type("FakeAutoTokenizer", (), {"from_pretrained": staticmethod(lambda *args, **kwargs: _FakeTokenizer())}),
+    )
+
+    if "peft" in recipe_func.__name__:
+        cfg = recipe_func(peft_scheme="lora")
+    else:
+        cfg = recipe_func()
+
+    assert isinstance(cfg.model, BridgeGPTModelConfig)
+    assert cfg.model.normalization == "RMSNorm"
+    assert cfg.model.moe_router_score_function == "sigmoid"
+    assert cfg.model.mtp_num_layers == 1