Skip to content

refactor(model): build OLMoE and GLM-4.5 from builder-backed model configs #6383

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cuichenx merged 3 commits into from
Oct 16, 2026
Merged

refactor(model): build OLMoE and GLM-4.5 from builder-backed model configs #6383

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c91143d
refactor(model): build OLMoE from a builder-backed model config
kamran-nvidia Oct 9, 2026
4f859a5
refactor(model): build GLM-4.5 from a builder-backed model config
kamran-nvidia Oct 9, 2026
584dc25
refactor(model): opt the migrated families into builder-backed model …
kamran-nvidia Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 50 additions & 38 deletions src/megatron/bridge/models/glm/glm45_bridge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -28,8 +29,7 @@
GatedMLPMapping,
QKVMapping,
)
from megatron.bridge.models.gpt_provider import GPTModelProvider
from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig


try:
Expand All @@ -51,47 +51,59 @@ 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
MODEL_CONFIG_CLASS = BridgeGPTModelConfig
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 = []
Expand Down
90 changes: 45 additions & 45 deletions src/megatron/bridge/models/olmoe/olmoe_bridge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,11 @@
# limitations under the License.

import logging
from typing import Any

import torch
from megatron.core.models.gpt.gpt_model import GPTModel
from transformers import OlmoeForCausalLM
from transformers import OlmoeForCausalLM, PretrainedConfig

from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
Expand All @@ -25,8 +26,7 @@
GatedMLPMapping,
QKVMapping,
)
from megatron.bridge.models.gpt_provider import GPTModelProvider
from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
from megatron.bridge.models.gpt.model_config import BridgeGPTModelConfig
from megatron.bridge.models.olmoe.olmoe_provider import olmoe_layer_spec


Expand All @@ -45,56 +45,56 @@ class OlMoEBridge(MegatronModelBridge):
Example:
>>> from megatron.bridge import AutoBridge
>>> bridge = AutoBridge.from_hf_pretrained("allenai/OLMoE-1B-7B-0125")
>>> provider = bridge.to_megatron_provider()
>>> model_config = bridge.get_model_config()
"""

def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> GPTModelProvider:
"""Convert HuggingFace OlMoE config to Megatron GPTModelProvider.
MODEL_CONFIG_CLASS = BridgeGPTModelConfig
USE_MODEL_CONFIG_FOR_CONVERSION = True

Uses base class implementation for common conversion, then sets
OlMoE-specific config. OlMoE uses QK layernorm and mixture of experts.
def hf_config_to_model_config_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
"""Convert a Hugging Face OlMoE config to builder config kwargs.

Args:
hf_pretrained: HuggingFace PreTrainedCausalLM containing the OlMoE config

Returns:
GPTModelProvider configured for OlMoE architecture
OlMoE uses QK layernorm and mixture of experts.
"""
provider = super().provider_bridge(hf_pretrained)
hf_config = hf_pretrained.config

# OlMoE uses custom layer spec with OLMoESelfAttention for QK layernorm
provider.transformer_layer_spec = olmoe_layer_spec

# Set kv_channels (head_dim) - OLMoE HF config doesn't have head_dim, so calculate it
provider.kv_channels = getattr(hf_config, "head_dim", None) or (
hf_config.hidden_size // hf_config.num_attention_heads
)

# OlMoE-specific architecture settings
provider.normalization = "RMSNorm"
provider.gated_linear_unit = True
provider.add_bias_linear = False
provider.hidden_dropout = 0.0
provider.share_embeddings_and_output_weights = False
provider.qk_layernorm = True
provider.persist_layer_norm = True
provider.autocast_dtype = torch.bfloat16

# MoE-specific settings
provider.moe_ffn_hidden_size = hf_config.intermediate_size
provider.moe_aux_loss_coeff = hf_config.router_aux_loss_coef
provider.moe_token_dispatcher_type = "alltoall"
provider.moe_router_load_balancing_type = "seq_aux_loss"
provider.moe_router_pre_softmax = True
provider.moe_grouped_gemm = True
if hasattr(hf_config, "scoring_func") and hf_config.scoring_func != "softmax":
raise ValueError(f"OlMoE only supports scoring_func='softmax', got {hf_config.scoring_func!r}")
provider.moe_router_score_function = "softmax"
provider.moe_permute_fusion = True
provider.moe_router_dtype = "fp32"

return provider
config_kwargs = super().hf_config_to_model_config_kwargs(hf_config)
config_kwargs.update(
# OlMoE uses custom layer spec with OLMoESelfAttention for QK layernorm
transformer_layer_spec=olmoe_layer_spec,
# OLMoE HF config doesn't have head_dim, so calculate it
kv_channels=getattr(hf_config, "head_dim", None)
or (hf_config.hidden_size // hf_config.num_attention_heads),
# OlMoE-specific architecture settings
normalization="RMSNorm",
gated_linear_unit=True,
add_bias_linear=False,
hidden_dropout=0.0,
share_embeddings_and_output_weights=False,
qk_layernorm=True,
persist_layer_norm=True,
autocast_dtype=torch.bfloat16,
masked_softmax_fusion=True,
rope_scaling=False,
rope_scaling_factor=1.0,
# MoE-specific settings
moe_ffn_hidden_size=hf_config.intermediate_size,
moe_aux_loss_coeff=hf_config.router_aux_loss_coef,
moe_token_dispatcher_type="alltoall",
moe_router_load_balancing_type="seq_aux_loss",
moe_router_pre_softmax=True,
moe_grouped_gemm=True,
moe_router_score_function="softmax",
moe_permute_fusion=True,
moe_router_dtype="fp32",
)
config_kwargs.setdefault("position_embedding_type", "rope")
return config_kwargs

def hf_config_to_provider_kwargs(self, hf_config: PretrainedConfig) -> dict[str, Any]:
"""Adapt the canonical builder mapping to the deprecated provider path."""
return self.hf_config_to_model_config_kwargs(hf_config)

def mapping_registry(self) -> MegatronMappingRegistry:
mapping_list = []
Expand Down
14 changes: 11 additions & 3 deletions src/megatron/bridge/models/olmoe/olmoe_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from megatron.core.transformer.spec_utils import build_module
from megatron.core.transformer.transformer_config import TransformerConfig

from megatron.bridge.models.gpt import gpt_builder
from megatron.bridge.models.gpt_provider import GPTModelProvider, default_layer_spec


Expand All @@ -38,9 +39,16 @@
SplitAlongDim = None


def olmoe_layer_spec(config: "GPTModelProvider") -> ModuleSpec:
"""Layer spec for OlMoE models."""
layer_spec = default_layer_spec(config)
def olmoe_layer_spec(config: "GPTModelProvider | gpt_builder.GPTModelConfig") -> ModuleSpec:
"""Layer spec for OlMoE models.

Accepts the legacy provider or a builder-backed GPT model config, whose
transformer fields live on ``config.transformer``.
"""
if isinstance(config, GPTModelProvider):
layer_spec = default_layer_spec(config)
else:
layer_spec = gpt_builder.transformer_engine_layer_spec(config)
layer_spec.submodules.self_attention.module = OLMoESelfAttention
return layer_spec

Expand Down
12 changes: 6 additions & 6 deletions src/megatron/bridge/recipes/glm/h100/glm45.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading