Skip to content

[model] fix: register both MTP submodule spellings in qwen3_next_bridge #11

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

Open
Zhichenzzz wants to merge 1 commit into
base: bridge
Choose a base branch
from
Open

[model] fix: register both MTP submodule spellings in qwen3_next_bridge #11

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
100 changes: 64 additions & 36 deletions src/megatron/bridge/models/qwen/qwen3_next_bridge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,16 +124,34 @@ def mapping_registry(self) -> MegatronMappingRegistry:
"mtp.layers.0.enorm.weight": "mtp.pre_fc_norm_embedding.weight",
"mtp.layers.0.hnorm.weight": "mtp.pre_fc_norm_hidden.weight",
"mtp.layers.0.final_layernorm.weight": "mtp.norm.weight",
# MTP MoE
"mtp.layers.0.transformer_layer.mlp.router.weight": "mtp.layers.0.mlp.gate.weight",
"mtp.layers.0.transformer_layer.pre_mlp_layernorm.weight": "mtp.layers.0.post_attention_layernorm.weight",
# MTP standard attention
"mtp.layers.0.transformer_layer.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight",
"mtp.layers.0.transformer_layer.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight",
"mtp.layers.0.transformer_layer.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight",
"mtp.layers.0.transformer_layer.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight",
}

# Register both MTP submodule spellings (mtp_model_layer / transformer_layer),
# like glm45_bridge, so the bridge works with either Megatron-LM version.
for mtp_layer_attr in ("mtp_model_layer", "transformer_layer"):
param_mappings.update(
{
# MTP MoE
f"mtp.layers.0.{mtp_layer_attr}.mlp.router.weight": "mtp.layers.0.mlp.gate.weight",
f"mtp.layers.0.{mtp_layer_attr}.pre_mlp_layernorm.weight": (
"mtp.layers.0.post_attention_layernorm.weight"
),
# MTP standard attention
f"mtp.layers.0.{mtp_layer_attr}.self_attention.linear_qkv.layer_norm_weight": (
"mtp.layers.0.input_layernorm.weight"
),
f"mtp.layers.0.{mtp_layer_attr}.self_attention.q_layernorm.weight": (
"mtp.layers.0.self_attn.q_norm.weight"
),
f"mtp.layers.0.{mtp_layer_attr}.self_attention.k_layernorm.weight": (
"mtp.layers.0.self_attn.k_norm.weight"
),
f"mtp.layers.0.{mtp_layer_attr}.self_attention.linear_proj.weight": (
"mtp.layers.0.self_attn.o_proj.weight"
),
}
)

mapping_list = []
# Convert each dictionary entry to AutoMapping(megatron_param, hf_param)
for megatron_param, hf_param in param_mappings.items():
Expand All @@ -152,12 +170,6 @@ def mapping_registry(self) -> MegatronMappingRegistry:
k="model.layers.*.self_attn.k_proj.weight",
v="model.layers.*.self_attn.v_proj.weight",
),
QKVMapping(
megatron_param="mtp.layers.*.transformer_layer.self_attention.linear_qkv.weight",
q="mtp.layers.*.self_attn.q_proj.weight",
k="mtp.layers.*.self_attn.k_proj.weight",
v="mtp.layers.*.self_attn.v_proj.weight",
),
# GDNLinear: Combine separate QKVZ_proj and BA_proj into single in_proj for GDN
# Note: Qwen3-Next does NOT have bias in the input linear projections
GDNConv1dMapping(
Expand All @@ -179,15 +191,6 @@ def mapping_registry(self) -> MegatronMappingRegistry:
megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*",
hf_param="model.layers.*.mlp.experts.*.down_proj.weight",
),
GatedMLPMapping(
megatron_param="mtp.layers.*.transformer_layer.mlp.experts.linear_fc1.weight*",
gate="mtp.layers.*.mlp.experts.*.gate_proj.weight",
up="mtp.layers.*.mlp.experts.*.up_proj.weight",
),
AutoMapping(
megatron_param="mtp.layers.*.transformer_layer.mlp.experts.linear_fc2.weight*",
hf_param="mtp.layers.*.mlp.experts.*.down_proj.weight",
),
# Gated MLP of shared expert
GatedMLPMapping(
megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight",
Expand All @@ -198,24 +201,11 @@ def mapping_registry(self) -> MegatronMappingRegistry:
megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc2.weight",
hf_param="model.layers.*.mlp.shared_expert.down_proj.weight",
),
GatedMLPMapping(
megatron_param="mtp.layers.*.transformer_layer.mlp.shared_experts.linear_fc1.weight",
gate="mtp.layers.*.mlp.shared_expert.gate_proj.weight",
up="mtp.layers.*.mlp.shared_expert.up_proj.weight",
),
AutoMapping(
megatron_param="mtp.layers.*.transformer_layer.mlp.shared_experts.linear_fc2.weight",
hf_param="mtp.layers.*.mlp.shared_expert.down_proj.weight",
),
# Shared expert gate
ReplicatedMapping(
megatron_param="decoder.layers.*.mlp.shared_experts.gate_weight",
hf_param="model.layers.*.mlp.shared_expert_gate.weight",
),
ReplicatedMapping(
megatron_param="mtp.layers.0.transformer_layer.mlp.shared_experts.gate_weight",
hf_param="mtp.layers.0.mlp.shared_expert_gate.weight",
),
# Qwen3-Next implements the output norm as a standard RMSNorm while initializing weight to ones,
# while other norms are regular zero-centered RMSNorms.
# To correctly load the output norm weight, we need to subtract 1 from it.
Expand All @@ -226,4 +216,42 @@ def mapping_registry(self) -> MegatronMappingRegistry:
]
)

# MTP transformation mappings, registered for both submodule spellings (see above).
for mtp_layer_attr in ("mtp_model_layer", "transformer_layer"):
mapping_list.extend(
[
QKVMapping(
megatron_param=f"mtp.layers.*.{mtp_layer_attr}.self_attention.linear_qkv.weight",
q="mtp.layers.*.self_attn.q_proj.weight",
k="mtp.layers.*.self_attn.k_proj.weight",
v="mtp.layers.*.self_attn.v_proj.weight",
),
# Gated MLP of experts
GatedMLPMapping(
megatron_param=f"mtp.layers.*.{mtp_layer_attr}.mlp.experts.linear_fc1.weight*",
gate="mtp.layers.*.mlp.experts.*.gate_proj.weight",
up="mtp.layers.*.mlp.experts.*.up_proj.weight",
),
AutoMapping(
megatron_param=f"mtp.layers.*.{mtp_layer_attr}.mlp.experts.linear_fc2.weight*",
hf_param="mtp.layers.*.mlp.experts.*.down_proj.weight",
),
# Gated MLP of shared expert
GatedMLPMapping(
megatron_param=f"mtp.layers.*.{mtp_layer_attr}.mlp.shared_experts.linear_fc1.weight",
gate="mtp.layers.*.mlp.shared_expert.gate_proj.weight",
up="mtp.layers.*.mlp.shared_expert.up_proj.weight",
),
AutoMapping(
megatron_param=f"mtp.layers.*.{mtp_layer_attr}.mlp.shared_experts.linear_fc2.weight",
hf_param="mtp.layers.*.mlp.shared_expert.down_proj.weight",
),
# Shared expert gate
ReplicatedMapping(
megatron_param=f"mtp.layers.0.{mtp_layer_attr}.mlp.shared_experts.gate_weight",
hf_param="mtp.layers.0.mlp.shared_expert_gate.weight",
),
]
)

return MegatronMappingRegistry(*mapping_list)
Loading