From 98fbc99d9a9d6edb38444b0baa900bd0005db138 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sun, 12 Jul 2026 23:20:26 +0000 Subject: [PATCH 01/10] feat(qwen): migrate Qwen3 dense and MoE models to HybridModel Signed-off-by: Philip Petrakian --- .../bridge/models/hybrid/hybrid_provider.py | 10 +- .../bridge/models/qwen/qwen3_bridge.py | 130 +++++++---- .../bridge/models/qwen/qwen3_moe_bridge.py | 139 +++++++---- .../bridge/models/qwen/qwen_hybrid.py | 220 ++++++++++++++++++ .../recipes/test_qwen_recipes_pretrain.py | 2 +- .../models/qwen/test_qwen3_bridge.py | 71 ++++-- .../models/qwen/test_qwen3_moe_bridge.py | 25 +- .../models/qwen/test_qwen_hybrid.py | 95 ++++++++ 8 files changed, 573 insertions(+), 119 deletions(-) create mode 100644 src/megatron/bridge/models/qwen/qwen_hybrid.py create mode 100644 tests/unit_tests/models/qwen/test_qwen_hybrid.py diff --git a/src/megatron/bridge/models/hybrid/hybrid_provider.py b/src/megatron/bridge/models/hybrid/hybrid_provider.py index 84acdae2af..9e7e036ed3 100644 --- a/src/megatron/bridge/models/hybrid/hybrid_provider.py +++ b/src/megatron/bridge/models/hybrid/hybrid_provider.py @@ -193,9 +193,13 @@ def finalize(self) -> None: self.hybrid_override_pattern = None used_hybrid_override_pattern = True - # Combine hybrid_layer_pattern (main decoder) with mtp_hybrid_override_pattern - # into a single unified pattern that MCore HybridModel can parse. - if self.hybrid_layer_pattern is not None and self.mtp_hybrid_override_pattern: + # Combine the main and MTP patterns only when MTP is enabled. Providers may + # predeclare an MTP override before a recipe sets mtp_num_layers. + if ( + self.hybrid_layer_pattern is not None + and self.mtp_hybrid_override_pattern + and (self.mtp_num_layers or self.mtp_use_repeated_layer) + ): sep = Symbols.MTP_SEPARATOR main_pattern = self.hybrid_layer_pattern.split(sep)[0] # When mtp_use_repeated_layer=True, the shared MTP layer always exists diff --git a/src/megatron/bridge/models/qwen/qwen3_bridge.py b/src/megatron/bridge/models/qwen/qwen3_bridge.py index e5b1442f1b..d3ca04dca5 100644 --- a/src/megatron/bridge/models/qwen/qwen3_bridge.py +++ b/src/megatron/bridge/models/qwen/qwen3_bridge.py @@ -13,7 +13,8 @@ # limitations under the License. import torch -from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.models.hybrid.hybrid_model import HybridModel from transformers import Qwen3ForCausalLM from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry @@ -23,15 +24,26 @@ GatedMLPMapping, QKVMapping, ) +from megatron.bridge.models.qwen.qwen_hybrid import ( + QwenHybridModelProvider, + configure_qwen_hybrid_layers, + qwen_logical_layer_count, + qwen_physical_layer_indices, +) -@MegatronModelBridge.register_bridge(source=Qwen3ForCausalLM, target=GPTModel, model_type="qwen3") +@MegatronModelBridge.register_bridge( + source=Qwen3ForCausalLM, + target=HybridModel, + provider=QwenHybridModelProvider, + model_type="qwen3", +) class Qwen3Bridge(MegatronModelBridge): """ Megatron Bridge for Qwen3 Causal LM. This bridge handles the conversion between HuggingFace Qwen3ForCausalLM - and Megatron-Core GPTModel formats. Qwen3 differs from Qwen2 by using + and Megatron-Core HybridModel formats. Qwen3 differs from Qwen2 by using QK layernorm and no QKV bias. Example: @@ -40,9 +52,19 @@ class Qwen3Bridge(MegatronModelBridge): >>> provider = bridge.to_megatron_provider() """ + @classmethod + def megatron_to_hf_config(cls, provider) -> dict: + """Convert a Hybrid Qwen3 provider to a Hugging Face config dictionary.""" + hf_config = super().megatron_to_hf_config(provider) + logical_layer_count = qwen_logical_layer_count(provider.hybrid_layer_pattern) + if logical_layer_count is not None: + hf_config["num_hidden_layers"] = logical_layer_count + return hf_config + def provider_bridge(self, hf_pretrained): - """Convert HuggingFace Qwen3 config to GPTModelProvider.""" + """Convert a Hugging Face Qwen3 config to HybridModelProvider.""" provider = super().provider_bridge(hf_pretrained) + hf_config = hf_pretrained.config provider.normalization = "RMSNorm" provider.gated_linear_unit = True @@ -51,6 +73,14 @@ def provider_bridge(self, hf_pretrained): provider.hidden_dropout = 0.0 provider.qk_layernorm = True # Qwen3 uses QK layernorm provider.autocast_dtype = torch.bfloat16 + provider.share_embeddings_and_output_weights = getattr(hf_config, "tie_word_embeddings", True) + + configure_qwen_hybrid_layers( + provider, + num_logical_layers=hf_config.num_hidden_layers, + mlp_symbols=Symbols.MLP, + mtp_mlp_symbol=Symbols.MLP, + ) return provider @@ -65,67 +95,85 @@ def mapping_registry(self) -> MegatronMappingRegistry: The fused QKV matrix is handled by :class:`QKVMapping` and the gated MLP gate+up projection by :class:`GatedMLPMapping`. """ - # Dictionary maps Megatron parameter names -> HF parameter names - # Supports wildcard (*) patterns for layer-specific parameters param_mappings = { # Embedding and output "embedding.word_embeddings.weight": "model.embed_tokens.weight", "output_layer.weight": "lm_head.weight", - "decoder.final_layernorm.weight": "model.norm.weight", - # Decoder layer attention norms and projections - "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight", - "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.layers.*.post_attention_layernorm.weight", - "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_norm.weight", # Qwen3 specific - "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight", # Qwen3 specific - "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight", - "decoder.layers.*.mlp.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight", + "decoder.final_norm.weight": "model.norm.weight", # MTP projection and norms (used when mtp_num_layers >= 1) "mtp.layers.0.eh_proj.weight": "mtp.fc.weight", "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 transformer layer attention - "mtp.layers.0.mtp_model_layer.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", # MTP transformer layer MLP - "mtp.layers.0.mtp_model_layer.mlp.linear_fc1.layer_norm_weight": "mtp.layers.0.post_attention_layernorm.weight", - "mtp.layers.0.mtp_model_layer.mlp.linear_fc2.weight": "mtp.layers.0.mlp.down_proj.weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc1.layer_norm_weight": "mtp.layers.0.post_attention_layernorm.weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc2.weight": "mtp.layers.0.mlp.down_proj.weight", } - mapping_list = [] - # Convert each dictionary entry to AutoMapping(megatron_param, hf_param) - for megatron_param, hf_param in param_mappings.items(): - mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) + mapping_list = [AutoMapping(megatron_param=k, hf_param=v) for k, v in param_mappings.items()] + + for logical_layer_idx in range(self.hf_config.num_hidden_layers): + attention_layer_idx, mlp_layer_idx = qwen_physical_layer_indices(logical_layer_idx) + hf_layer = f"model.layers.{logical_layer_idx}" + attention_layer = f"decoder.layers.{attention_layer_idx}.self_attention" + mlp_layer = f"decoder.layers.{mlp_layer_idx}.mlp" + mapping_list.extend( + [ + AutoMapping( + megatron_param=f"{attention_layer}.linear_qkv.layer_norm_weight", + hf_param=f"{hf_layer}.input_layernorm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.q_layernorm.weight", + hf_param=f"{hf_layer}.self_attn.q_norm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.k_layernorm.weight", + hf_param=f"{hf_layer}.self_attn.k_norm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.linear_proj.weight", + hf_param=f"{hf_layer}.self_attn.o_proj.weight", + ), + AutoMapping( + megatron_param=f"{mlp_layer}.linear_fc1.layer_norm_weight", + hf_param=f"{hf_layer}.post_attention_layernorm.weight", + ), + AutoMapping( + megatron_param=f"{mlp_layer}.linear_fc2.weight", + hf_param=f"{hf_layer}.mlp.down_proj.weight", + ), + QKVMapping( + megatron_param=f"{attention_layer}.linear_qkv.weight", + q=f"{hf_layer}.self_attn.q_proj.weight", + k=f"{hf_layer}.self_attn.k_proj.weight", + v=f"{hf_layer}.self_attn.v_proj.weight", + ), + GatedMLPMapping( + megatron_param=f"{mlp_layer}.linear_fc1.weight", + gate=f"{hf_layer}.mlp.gate_proj.weight", + up=f"{hf_layer}.mlp.up_proj.weight", + ), + ] + ) - # Add special mappings that require parameter concatenation/transformation mapping_list.extend( [ - # QKV: Combine separate Q, K, V matrices into single QKV matrix - # Note: Qwen3 does NOT have bias in QKV projections (unlike Qwen2) - QKVMapping( - megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", - q="model.layers.*.self_attn.q_proj.weight", - k="model.layers.*.self_attn.k_proj.weight", - v="model.layers.*.self_attn.v_proj.weight", - ), - # Gated MLP: Combine gate and up projection matrices into single FC1 matrix - GatedMLPMapping( - megatron_param="decoder.layers.*.mlp.linear_fc1.weight", - gate="model.layers.*.mlp.gate_proj.weight", - up="model.layers.*.mlp.up_proj.weight", - ), # MTP QKV: same split/merge as decoder layers QKVMapping( - megatron_param="mtp.layers.*.mtp_model_layer.self_attention.linear_qkv.weight", + megatron_param="mtp.layers.*.mtp_model_layer.layers.0.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", ), # MTP Gated MLP GatedMLPMapping( - megatron_param="mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight", + megatron_param="mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc1.weight", gate="mtp.layers.0.mlp.gate_proj.weight", up="mtp.layers.0.mlp.up_proj.weight", ), diff --git a/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py b/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py index a6476f4f49..1f9fa5a46c 100755 --- a/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py +++ b/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py @@ -13,7 +13,8 @@ # limitations under the License. import torch -from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.models.hybrid.hybrid_model import HybridModel from transformers import Qwen3MoeForCausalLM from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry @@ -23,15 +24,26 @@ GatedMLPMapping, QKVMapping, ) +from megatron.bridge.models.qwen.qwen_hybrid import ( + QwenHybridModelProvider, + configure_qwen_hybrid_layers, + qwen_logical_layer_count, + qwen_physical_layer_indices, +) -@MegatronModelBridge.register_bridge(source=Qwen3MoeForCausalLM, target=GPTModel, model_type="qwen3_moe") +@MegatronModelBridge.register_bridge( + source=Qwen3MoeForCausalLM, + target=HybridModel, + provider=QwenHybridModelProvider, + model_type="qwen3_moe", +) class Qwen3MoEBridge(MegatronModelBridge): """ Megatron Bridge for Qwen3 MoE Causal LM. This bridge handles the conversion between HuggingFace Qwen3MoeForCausalLM - and Megatron-Core GPTModel formats. Qwen3 MoE models use mixture of experts + and Megatron-Core HybridModel formats. Qwen3 MoE models use mixture of experts architecture with QK layernorm. Example: @@ -44,13 +56,17 @@ class Qwen3MoEBridge(MegatronModelBridge): def megatron_to_hf_config(cls, provider) -> dict: """Convert Megatron provider config to HuggingFace Qwen3MoeConfig dict.""" hf_config = super().megatron_to_hf_config(provider) + logical_layer_count = qwen_logical_layer_count(provider.hybrid_layer_pattern) + if logical_layer_count is not None: + hf_config["num_hidden_layers"] = logical_layer_count hf_config["decoder_sparse_step"] = 1 # All layers are MoE in Qwen3 MoE hf_config["norm_topk_prob"] = not provider.moe_router_pre_softmax return hf_config def provider_bridge(self, hf_pretrained): - """Convert HuggingFace Qwen3 MoE config to GPTModelProvider.""" + """Convert a Hugging Face Qwen3 MoE config to HybridModelProvider.""" provider = super().provider_bridge(hf_pretrained) + hf_config = hf_pretrained.config provider.normalization = "RMSNorm" provider.gated_linear_unit = True @@ -66,6 +82,14 @@ def provider_bridge(self, hf_pretrained): provider.moe_router_pre_softmax = not hf_pretrained.config.norm_topk_prob provider.moe_token_dispatcher_type = "alltoall" provider.moe_permute_fusion = True + provider.share_embeddings_and_output_weights = getattr(hf_config, "tie_word_embeddings", True) + + configure_qwen_hybrid_layers( + provider, + num_logical_layers=hf_config.num_hidden_layers, + mlp_symbols=Symbols.MOE, + mtp_mlp_symbol=Symbols.MOE, + ) return provider @@ -73,57 +97,70 @@ def mapping_registry(self) -> MegatronMappingRegistry: # Return MegatronMappingRegistry containing parameter mappings from Megatron to HF format # First create simple 1:1 parameter mappings using a dictionary for readability - # Dictionary maps Megatron parameter names -> HF parameter names - # Supports wildcard (*) patterns for layer-specific parameters param_mappings = { "embedding.word_embeddings.weight": "model.embed_tokens.weight", "output_layer.weight": "lm_head.weight", - "decoder.final_layernorm.weight": "model.norm.weight", - "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight", - "decoder.layers.*.mlp.router.weight": "model.layers.*.mlp.gate.weight", - "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.post_attention_layernorm.weight", - "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_norm.weight", - "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight", - "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight", + "decoder.final_norm.weight": "model.norm.weight", } - mapping_list = [] - # Convert each dictionary entry to AutoMapping(megatron_param, hf_param) - for megatron_param, hf_param in param_mappings.items(): - mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) - - # Add special mappings that require parameter concatenation/transformation - mapping_list.extend( - [ - # QKV: Combine separate Q, K, V matrices into single QKV matrix - # Note: Qwen3 MoE does NOT have bias in QKV projections - QKVMapping( - megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", - q="model.layers.*.self_attn.q_proj.weight", - k="model.layers.*.self_attn.k_proj.weight", - v="model.layers.*.self_attn.v_proj.weight", - ), - # Expert mappings for TEGroupedMLP - GatedMLPMapping( - megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", - gate="model.layers.*.mlp.experts.*.gate_proj.weight", - up="model.layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", - hf_param="model.layers.*.mlp.experts.*.down_proj.weight", - ), - # Expert mappings for SequentialMLP (used by quantization) - GatedMLPMapping( - megatron_param="decoder.layers.*.mlp.experts.local_experts.*.linear_fc1.weight", - gate="model.layers.*.mlp.experts.*.gate_proj.weight", - up="model.layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight", - hf_param="model.layers.*.mlp.experts.*.down_proj.weight", - ), - ] - ) + mapping_list = [AutoMapping(megatron_param=k, hf_param=v) for k, v in param_mappings.items()] + + for logical_layer_idx in range(self.hf_config.num_hidden_layers): + attention_layer_idx, moe_layer_idx = qwen_physical_layer_indices(logical_layer_idx) + hf_layer = f"model.layers.{logical_layer_idx}" + attention_layer = f"decoder.layers.{attention_layer_idx}.self_attention" + moe_layer = f"decoder.layers.{moe_layer_idx}" + mapping_list.extend( + [ + AutoMapping( + megatron_param=f"{attention_layer}.linear_qkv.layer_norm_weight", + hf_param=f"{hf_layer}.input_layernorm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.q_layernorm.weight", + hf_param=f"{hf_layer}.self_attn.q_norm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.k_layernorm.weight", + hf_param=f"{hf_layer}.self_attn.k_norm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.linear_proj.weight", + hf_param=f"{hf_layer}.self_attn.o_proj.weight", + ), + AutoMapping( + megatron_param=f"{moe_layer}.pre_mlp_layernorm.weight", + hf_param=f"{hf_layer}.post_attention_layernorm.weight", + ), + AutoMapping( + megatron_param=f"{moe_layer}.mlp.router.weight", + hf_param=f"{hf_layer}.mlp.gate.weight", + ), + QKVMapping( + megatron_param=f"{attention_layer}.linear_qkv.weight", + q=f"{hf_layer}.self_attn.q_proj.weight", + k=f"{hf_layer}.self_attn.k_proj.weight", + v=f"{hf_layer}.self_attn.v_proj.weight", + ), + GatedMLPMapping( + megatron_param=f"{moe_layer}.mlp.experts.linear_fc1.weight*", + gate=f"{hf_layer}.mlp.experts.*.gate_proj.weight", + up=f"{hf_layer}.mlp.experts.*.up_proj.weight", + ), + AutoMapping( + megatron_param=f"{moe_layer}.mlp.experts.linear_fc2.weight*", + hf_param=f"{hf_layer}.mlp.experts.*.down_proj.weight", + ), + GatedMLPMapping( + megatron_param=f"{moe_layer}.mlp.experts.local_experts.*.linear_fc1.weight", + gate=f"{hf_layer}.mlp.experts.*.gate_proj.weight", + up=f"{hf_layer}.mlp.experts.*.up_proj.weight", + ), + AutoMapping( + megatron_param=f"{moe_layer}.mlp.experts.local_experts.*.linear_fc2.weight", + hf_param=f"{hf_layer}.mlp.experts.*.down_proj.weight", + ), + ] + ) return MegatronMappingRegistry(*mapping_list) diff --git a/src/megatron/bridge/models/qwen/qwen_hybrid.py b/src/megatron/bridge/models/qwen/qwen_hybrid.py new file mode 100644 index 0000000000..b5c30fa683 --- /dev/null +++ b/src/megatron/bridge/models/qwen/qwen_hybrid.py @@ -0,0 +1,220 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared HybridModel layout helpers for Qwen3 and newer models.""" + +from collections.abc import Sequence +from dataclasses import dataclass + +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider + + +def qwen_pipeline_layer_pattern( + hybrid_layer_pattern: str, + pipeline_model_parallel_size: int, + *, + account_for_embedding: bool = False, + account_for_loss: bool = False, +) -> str: + """Split a Qwen Hybrid pattern across PP stages without separating logical blocks.""" + if pipeline_model_parallel_size <= 0: + raise ValueError("pipeline_model_parallel_size must be positive") + + main_pattern, separator, mtp_pattern = hybrid_layer_pattern.partition(Symbols.MTP_SEPARATOR) + if Symbols.PIPE in main_pattern: + segment_count = main_pattern.count(Symbols.PIPE) + 1 + if segment_count != pipeline_model_parallel_size: + raise ValueError( + f"Qwen hybrid_layer_pattern defines {segment_count} pipeline segments, " + f"but pipeline_model_parallel_size is {pipeline_model_parallel_size}." + ) + return hybrid_layer_pattern + if pipeline_model_parallel_size == 1: + return hybrid_layer_pattern + if len(main_pattern) % 2: + raise ValueError("Qwen Hybrid patterns must contain two physical layers per logical block.") + + logical_blocks = [main_pattern[index : index + 2] for index in range(0, len(main_pattern), 2)] + attention_symbols = {Symbols.ATTENTION, Symbols.GDN} + mlp_symbols = {Symbols.MLP, Symbols.MOE} + invalid_blocks = [ + block for block in logical_blocks if block[0] not in attention_symbols or block[1] not in mlp_symbols + ] + if invalid_blocks: + raise ValueError(f"Unsupported Qwen logical blocks in hybrid_layer_pattern: {invalid_blocks}") + + total_pipeline_units = len(logical_blocks) + int(account_for_embedding) + int(account_for_loss) + units_per_stage, extra_units = divmod(total_pipeline_units, pipeline_model_parallel_size) + logical_blocks_per_stage = [ + units_per_stage + int(stage_index < extra_units) for stage_index in range(pipeline_model_parallel_size) + ] + logical_blocks_per_stage[0] -= int(account_for_embedding) + logical_blocks_per_stage[-1] -= int(account_for_loss) + if any(count <= 0 for count in logical_blocks_per_stage): + raise ValueError( + "Every pipeline stage must receive at least one Qwen logical block after embedding/loss balancing." + ) + + segments = [] + block_offset = 0 + for block_count in logical_blocks_per_stage: + segments.append("".join(logical_blocks[block_offset : block_offset + block_count])) + block_offset += block_count + if block_offset != len(logical_blocks): + raise RuntimeError("Failed to assign every Qwen logical block to a pipeline stage.") + + segmented_pattern = Symbols.PIPE.join(segments) + return segmented_pattern + (separator + mtp_pattern if separator else "") + + +@dataclass +class QwenHybridModelProvider(HybridModelProvider): + """HybridModel provider that keeps each Qwen logical block on one PP stage.""" + + def finalize(self) -> None: + if ( + self.hybrid_layer_pattern is not None + and self.pipeline_model_parallel_size > 1 + and self.pipeline_model_parallel_layout is None + and self.num_layers_in_first_pipeline_stage is None + and self.num_layers_in_last_pipeline_stage is None + ): + segmented_pattern = qwen_pipeline_layer_pattern( + self.hybrid_layer_pattern, + self.pipeline_model_parallel_size, + account_for_embedding=bool(self.account_for_embedding_in_pipeline_split), + account_for_loss=bool(self.account_for_loss_in_pipeline_split), + ) + if segmented_pattern != self.hybrid_layer_pattern: + self.hybrid_layer_pattern = segmented_pattern + # The explicit segments already include embedding/loss-aware balancing. + self.account_for_embedding_in_pipeline_split = False + self.account_for_loss_in_pipeline_split = False + super().finalize() + + +def qwen_attention_symbols( + num_layers: int, + linear_attention_freq: int | Sequence[int] | None = None, +) -> list[str]: + """Translate a logical Qwen attention schedule to HybridModel symbols.""" + if linear_attention_freq is None: + return [Symbols.ATTENTION] * num_layers + + if isinstance(linear_attention_freq, int): + if linear_attention_freq <= 0: + raise ValueError("linear_attention_freq must be positive") + return [ + Symbols.ATTENTION if (layer_idx + 1) % linear_attention_freq == 0 else Symbols.GDN + for layer_idx in range(num_layers) + ] + + linear_attention_pattern = list(linear_attention_freq) + if len(linear_attention_pattern) != num_layers: + raise ValueError( + "linear_attention_freq has " + f"{len(linear_attention_pattern)} entries, but num_hidden_layers is {num_layers}." + ) + invalid_values = sorted(set(linear_attention_pattern) - {0, 1}) + if invalid_values: + raise ValueError(f"Unsupported linear attention pattern values: {invalid_values}. Expected only 0 or 1.") + return [Symbols.GDN if is_linear else Symbols.ATTENTION for is_linear in linear_attention_pattern] + + +def qwen_hybrid_layer_pattern( + num_layers: int, + *, + mlp_symbols: str | Sequence[str], + linear_attention_freq: int | Sequence[int] | None = None, +) -> str: + """Build a two-physical-layer HybridModel pattern for each logical Qwen block.""" + if isinstance(mlp_symbols, str): + mlp_pattern = [mlp_symbols] * num_layers + else: + mlp_pattern = list(mlp_symbols) + if len(mlp_pattern) != num_layers: + raise ValueError(f"MLP pattern has {len(mlp_pattern)} entries, but num_hidden_layers is {num_layers}.") + + invalid_symbols = sorted(set(mlp_pattern) - {Symbols.MLP, Symbols.MOE}) + if invalid_symbols: + raise ValueError( + f"Unsupported Qwen MLP symbols: {invalid_symbols}. Expected '{Symbols.MLP}' or '{Symbols.MOE}'." + ) + + attention_pattern = qwen_attention_symbols(num_layers, linear_attention_freq) + return "".join( + attention_symbol + mlp_symbol for attention_symbol, mlp_symbol in zip(attention_pattern, mlp_pattern) + ) + + +def qwen_moe_layer_symbols( + num_layers: int, + *, + decoder_sparse_step: int = 1, + mlp_only_layers: Sequence[int] = (), +) -> list[str]: + """Translate Hugging Face Qwen MoE placement fields to Hybrid symbols.""" + if decoder_sparse_step <= 0: + raise ValueError("decoder_sparse_step must be positive") + dense_layers = set(mlp_only_layers) + invalid_layers = sorted(layer_idx for layer_idx in dense_layers if not 0 <= layer_idx < num_layers) + if invalid_layers: + raise ValueError(f"mlp_only_layers contains out-of-range indices: {invalid_layers}") + return [ + Symbols.MOE + if logical_layer_idx not in dense_layers and (logical_layer_idx + 1) % decoder_sparse_step == 0 + else Symbols.MLP + for logical_layer_idx in range(num_layers) + ] + + +def configure_qwen_hybrid_layers( + provider: HybridModelProvider, + *, + num_logical_layers: int, + mlp_symbols: str | Sequence[str], + linear_attention_freq: int | Sequence[int] | None = None, + mtp_mlp_symbol: str | None = None, +) -> None: + """Configure main and optional MTP physical layer patterns on a Qwen provider.""" + provider.hybrid_layer_pattern = qwen_hybrid_layer_pattern( + num_logical_layers, + mlp_symbols=mlp_symbols, + linear_attention_freq=linear_attention_freq, + ) + provider.num_layers = len(provider.hybrid_layer_pattern) + + if mtp_mlp_symbol is not None: + if mtp_mlp_symbol not in {Symbols.MLP, Symbols.MOE}: + raise ValueError( + f"Unsupported Qwen MTP MLP symbol: {mtp_mlp_symbol}. Expected '{Symbols.MLP}' or '{Symbols.MOE}'." + ) + provider.mtp_hybrid_override_pattern = Symbols.ATTENTION + mtp_mlp_symbol + + +def qwen_logical_layer_count(hybrid_layer_pattern: str | None) -> int | None: + """Return the number of logical Qwen blocks encoded by a HybridModel pattern.""" + if not hybrid_layer_pattern: + return None + main_pattern = hybrid_layer_pattern.split(Symbols.MTP_SEPARATOR)[0].replace(Symbols.PIPE, "") + attention_symbols = {Symbols.ATTENTION, Symbols.GDN} + return sum(symbol in attention_symbols for symbol in main_pattern) + + +def qwen_physical_layer_indices(logical_layer_idx: int) -> tuple[int, int]: + """Return attention and MLP/MoE physical indices for one logical Qwen block.""" + attention_layer_idx = 2 * logical_layer_idx + return attention_layer_idx, attention_layer_idx + 1 diff --git a/tests/functional_tests/test_groups/recipes/test_qwen_recipes_pretrain.py b/tests/functional_tests/test_groups/recipes/test_qwen_recipes_pretrain.py index 6550d86b56..6a96db259a 100644 --- a/tests/functional_tests/test_groups/recipes/test_qwen_recipes_pretrain.py +++ b/tests/functional_tests/test_groups/recipes/test_qwen_recipes_pretrain.py @@ -28,7 +28,7 @@ QWEN_PRETRAIN_RECIPES = [ # (config_func, name, parallelism_overrides, model_overrides) (qwen25_500m_config, "qwen25_500m", {}, {"num_layers": 2}), - (qwen3_600m_config, "qwen3_600m", {}, {"num_layers": 2}), + (qwen3_600m_config, "qwen3_600m", {}, {"num_layers": 4, "hybrid_layer_pattern": "*-*-"}), ] diff --git a/tests/unit_tests/models/qwen/test_qwen3_bridge.py b/tests/unit_tests/models/qwen/test_qwen3_bridge.py index 8bf74831c2..deeacc472d 100644 --- a/tests/unit_tests/models/qwen/test_qwen3_bridge.py +++ b/tests/unit_tests/models/qwen/test_qwen3_bridge.py @@ -14,6 +14,7 @@ import tempfile from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock, patch import pytest @@ -23,11 +24,20 @@ from megatron.bridge.models import AutoBridge from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from megatron.bridge.models.conversion.transformers_compat import rope_theta_from_hf -from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.models.qwen.qwen3_bridge import Qwen3Bridge +@pytest.fixture(autouse=True) +def _set_bridge_hf_config(): + """Give explicit mapping tests a small logical Qwen layer count.""" + previous = Qwen3Bridge.hf_config + Qwen3Bridge.hf_config = SimpleNamespace(num_hidden_layers=2) + yield + Qwen3Bridge.hf_config = previous + + class TestMegatronQwen3Bridge: """Test cases for MegatronQwen3Bridge class.""" @@ -93,11 +103,11 @@ def test_provider_bridge_basic(self, mock_pretrained_qwen3, qwen3_config): # Call provider_bridge result = bridge.provider_bridge(mock_pretrained_qwen3) - # Check that it returns a GPTModelProvider instance (after refactoring) - assert isinstance(result, GPTModelProvider) + assert isinstance(result, HybridModelProvider) # Check basic configuration mapping - assert result.num_layers == qwen3_config.num_hidden_layers + assert result.num_layers == 2 * qwen3_config.num_hidden_layers + assert result.hybrid_layer_pattern == "*-" * qwen3_config.num_hidden_layers assert result.hidden_size == qwen3_config.hidden_size assert result.num_attention_heads == qwen3_config.num_attention_heads assert result.seq_length == qwen3_config.max_position_embeddings @@ -200,8 +210,7 @@ def test_provider_bridge_with_custom_kwargs(self, mock_pretrained_qwen3): # Pass model only result = bridge.provider_bridge(mock_pretrained_qwen3) - # Just verify that we got a valid GPTModelProvider - assert isinstance(result, GPTModelProvider) + assert isinstance(result, HybridModelProvider) def test_provider_bridge_without_tie_embeddings(self, qwen3_config): """Test provider_bridge when tie_word_embeddings is not present.""" @@ -393,7 +402,7 @@ def test_from_pretrained_multiple_models(self, mock_autoconfig, mock_pretrained, "megatron.bridge.models.conversion.auto_bridge.model_bridge.get_model_bridge" ) as mock_get_bridge: mock_bridge = Mock() - mock_provider = Mock(spec=GPTModelProvider) + mock_provider = Mock(spec=HybridModelProvider) mock_bridge.provider_bridge.return_value = mock_provider mock_get_bridge.return_value = mock_bridge @@ -448,6 +457,38 @@ def test_supports_qwen3_architectures(self, qwen3_configs): assert AutoBridge.supports(non_causal_config) == False +class TestQwen3BridgeParameterMapping: + """Test parameter mapping functionality in Qwen3Bridge.""" + + def test_mapping_registry_has_qwen3_specific_mappings(self): + """Test that mapping registry includes Qwen3-specific QK norm mappings.""" + bridge = Qwen3Bridge() + mapping_registry = bridge.mapping_registry() + + mapping = mapping_registry.megatron_to_hf_lookup("decoder.layers.0.self_attention.q_layernorm.weight") + assert mapping is not None + assert mapping.hf_param == "model.layers.0.self_attn.q_norm.weight" + + def test_qwen3_qk_norm_mapping_difference(self): + """Test that Qwen3 bridge includes QK norm mappings not present in Qwen2.""" + bridge = Qwen3Bridge() + mapping_registry = bridge.mapping_registry() + + attention_mapping = mapping_registry.megatron_to_hf_lookup( + "decoder.layers.0.self_attention.k_layernorm.weight" + ) + mlp_mapping = mapping_registry.megatron_to_hf_lookup("decoder.layers.1.mlp.linear_fc2.weight") + assert attention_mapping.hf_param == "model.layers.0.self_attn.k_norm.weight" + assert mlp_mapping.hf_param == "model.layers.0.mlp.down_proj.weight" + + def test_qwen3_no_qkv_bias_mapping(self): + """Test that Qwen3 bridge doesn't include QKV bias mappings.""" + bridge = Qwen3Bridge() + mapping_registry = bridge.mapping_registry() + + assert mapping_registry.megatron_to_hf_lookup("decoder.layers.0.self_attention.linear_qkv.bias") is None + + class TestQwen3BridgeMTPMapping: """Tests for MTP (Multi-Token Prediction) weight mappings in Qwen3Bridge. @@ -466,15 +507,15 @@ class TestQwen3BridgeMTPMapping: "mtp.layers.0.enorm.weight", "mtp.layers.0.hnorm.weight", "mtp.layers.0.final_layernorm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.linear_qkv.layer_norm_weight", - "mtp.layers.0.mtp_model_layer.self_attention.q_layernorm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.k_layernorm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.linear_proj.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_qkv.layer_norm_weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.q_layernorm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.k_layernorm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_proj.weight", # QKVMapping stores a wildcard pattern, not a concrete layer index - "mtp.layers.*.mtp_model_layer.self_attention.linear_qkv.weight", - "mtp.layers.0.mtp_model_layer.mlp.linear_fc1.layer_norm_weight", - "mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight", - "mtp.layers.0.mtp_model_layer.mlp.linear_fc2.weight", + "mtp.layers.*.mtp_model_layer.layers.0.self_attention.linear_qkv.weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc1.layer_norm_weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc1.weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc2.weight", ) def _get_all_megatron_params(self, mapping_registry): diff --git a/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py b/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py index e9f23d16e6..eb28997d87 100644 --- a/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py +++ b/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py @@ -16,17 +16,26 @@ Unit tests for Qwen3 MoE bridge functionality. """ +from types import SimpleNamespace from unittest.mock import Mock import pytest import torch from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge -from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.models.qwen.qwen3_moe_bridge import Qwen3MoEBridge +@pytest.fixture(autouse=True) +def _set_bridge_hf_config(): + previous = Qwen3MoEBridge.hf_config + Qwen3MoEBridge.hf_config = SimpleNamespace(num_hidden_layers=2) + yield + Qwen3MoEBridge.hf_config = previous + + class TestQwen3MoEBridge: """Test cases for Qwen3MoEBridge class.""" @@ -124,11 +133,11 @@ def test_provider_bridge_basic(self, mock_pretrained_qwen3_moe, mock_qwen3_moe_c # Call provider_bridge result = bridge.provider_bridge(mock_pretrained_qwen3_moe) - # Check that it returns a GPTModelProvider instance (after refactoring) - assert isinstance(result, GPTModelProvider) + assert isinstance(result, HybridModelProvider) # Check basic configuration mapping - assert result.num_layers == mock_qwen3_moe_config.num_hidden_layers + assert result.num_layers == 2 * mock_qwen3_moe_config.num_hidden_layers + assert result.hybrid_layer_pattern == "*E" * mock_qwen3_moe_config.num_hidden_layers assert result.hidden_size == mock_qwen3_moe_config.hidden_size assert result.num_attention_heads == mock_qwen3_moe_config.num_attention_heads assert result.seq_length == mock_qwen3_moe_config.max_position_embeddings @@ -278,7 +287,7 @@ def test_provider_bridge_missing_tie_word_embeddings(self, mock_qwen3_moe_config bridge = Qwen3MoEBridge() result = bridge.provider_bridge(mock_pretrained) - # GPTModelProvider defaults share_embeddings_and_output_weights to True + # Preserve the historical default when the HF field is absent. assert result.share_embeddings_and_output_weights is True def test_provider_bridge_235b_config(self, qwen3_moe_235b_config_dict): @@ -294,7 +303,7 @@ def test_provider_bridge_235b_config(self, qwen3_moe_235b_config_dict): result = bridge.provider_bridge(mock_pretrained) # Check 235B-specific configuration - assert result.num_layers == 94 + assert result.num_layers == 188 assert result.hidden_size == 4096 assert result.num_attention_heads == 64 assert result.ffn_hidden_size == 12288 @@ -339,7 +348,7 @@ def test_mapping_registry_parameter_mappings(self): # Should have layer norm mappings assert "model.norm.weight" in hf_params - assert "decoder.final_layernorm.weight" in megatron_params + assert "decoder.final_norm.weight" in megatron_params def test_mapping_registry_qkv_mapping(self): """Test that mapping_registry contains QKV mapping.""" @@ -374,7 +383,7 @@ def test_mapping_registry_moe_mappings(self): # Check for MoE router mapping hf_params = [mapping.hf_param for mapping in auto_mappings] - assert "model.layers.*.mlp.gate.weight" in hf_params + assert "model.layers.0.mlp.gate.weight" in hf_params # Check for expert mappings in GatedMLPMapping assert len(gated_mlp_mappings) > 0 diff --git a/tests/unit_tests/models/qwen/test_qwen_hybrid.py b/tests/unit_tests/models/qwen/test_qwen_hybrid.py new file mode 100644 index 0000000000..2365be855a --- /dev/null +++ b/tests/unit_tests/models/qwen/test_qwen_hybrid.py @@ -0,0 +1,95 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + +from megatron.bridge.models.qwen.qwen_hybrid import ( + QwenHybridModelProvider, + configure_qwen_hybrid_layers, + qwen_pipeline_layer_pattern, +) + + +def _provider() -> QwenHybridModelProvider: + return QwenHybridModelProvider( + num_layers=2, + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=4, + bf16=False, + ) + + +def test_mtp_pattern_is_deferred_when_mtp_is_disabled(): + provider = _provider() + configure_qwen_hybrid_layers( + provider, + num_logical_layers=2, + mlp_symbols=Symbols.MLP, + mtp_mlp_symbol=Symbols.MLP, + ) + + provider.finalize() + + assert provider.hybrid_layer_pattern == "*-*-" + assert provider.mtp_hybrid_override_pattern == "*-" + + +def test_mtp_pattern_honors_recipe_override_set_after_conversion(): + provider = _provider() + configure_qwen_hybrid_layers( + provider, + num_logical_layers=2, + mlp_symbols=Symbols.MLP, + mtp_mlp_symbol=Symbols.MLP, + ) + provider.mtp_num_layers = 1 + + provider.finalize() + + assert provider.hybrid_layer_pattern == "*-*-/*-" + assert provider.num_layers == 4 + + +def test_pipeline_segmentation_preserves_logical_blocks_and_embedding_loss_balance(): + provider = QwenHybridModelProvider( + num_layers=94, + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=4, + pipeline_model_parallel_size=16, + account_for_embedding_in_pipeline_split=True, + account_for_loss_in_pipeline_split=True, + bf16=False, + ) + configure_qwen_hybrid_layers( + provider, + num_logical_layers=94, + mlp_symbols=Symbols.MOE, + ) + + provider.finalize() + + segments = provider.hybrid_layer_pattern.split(Symbols.PIPE) + assert [len(segment) for segment in segments] == [10] + [12] * 14 + [10] + assert all(segment == "*E" * (len(segment) // 2) for segment in segments) + assert provider.num_layers == 188 + assert provider.account_for_embedding_in_pipeline_split is False + assert provider.account_for_loss_in_pipeline_split is False + + +def test_pipeline_segmentation_rejects_stale_explicit_segments(): + with pytest.raises(ValueError, match="defines 2 pipeline segments"): + qwen_pipeline_layer_pattern("*-*-|*-*-", pipeline_model_parallel_size=1) From c274260e96fed18033ba4cd362523e83cb4f6529 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 13 Jul 2026 05:41:24 +0000 Subject: [PATCH 02/10] fix(qwen): restore HybridModel CI compatibility Signed-off-by: Philip Petrakian --- .../bridge/models/hybrid/hybrid_provider.py | 10 ++++++++-- .../ckpts/qwen3_4b/test_qwen3_4b_ckpt.py | 12 ++++++++---- .../training/test_seqpacking_cp_example.py | 6 ++++-- .../models/hybrid/test_hybrid_provider.py | 17 +++++++++++++++++ 4 files changed, 37 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/models/hybrid/hybrid_provider.py b/src/megatron/bridge/models/hybrid/hybrid_provider.py index 9e7e036ed3..746f1e096c 100644 --- a/src/megatron/bridge/models/hybrid/hybrid_provider.py +++ b/src/megatron/bridge/models/hybrid/hybrid_provider.py @@ -60,8 +60,9 @@ def _configure_mamba_chunk_size(stack_spec: ModuleSpec, chunk_size: int) -> Modu def modelopt_hybrid_stack_spec(config: "HybridModelProvider | None" = None) -> ModuleSpec: """Hybrid stack specification for quantization with ModelOpt. - Uses Norm instead of TENorm and ColumnParallelLinear/RowParallelLinear - instead of TE layers to enable proper quantizer insertion by ModelOpt. + Grouped-GEMM MoE checkpoints retain the default Transformer Engine stack so + their shared expert quantizers have the same state-dict paths when restored. + Other Hybrid models use ModelOpt's local stack specification. Args: config: Optional Hybrid configuration object. @@ -69,6 +70,11 @@ def modelopt_hybrid_stack_spec(config: "HybridModelProvider | None" = None) -> M Returns: Module specification for quantization-ready Hybrid stack. """ + if config is not None and config.num_moe_experts is not None and config.moe_grouped_gemm: + return get_hybrid_stack_modelopt_spec( + use_default_te_spec=True, + moe_grouped_gemm=True, + ) return get_hybrid_stack_modelopt_spec( local_core_attention=False, remap_te_layernorm=True, diff --git a/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py b/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py index 2e21c8507f..85885b2118 100644 --- a/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py +++ b/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py @@ -31,6 +31,7 @@ MBRIDGE_CKPT = f"{BASE_DIR}/mbridge" MCORE_CKPT = f"{BASE_DIR}/mcore" TB_DIR = f"{BASE_DIR}/tb" +QWEN3_4B_HYBRID_PATTERN = "*-" * 36 class TestQwen3Ckpt: @@ -60,7 +61,7 @@ def test_qwen3_4b_ckpt_mbridge(self): @pytest.mark.run_only_on("GPU") def test_qwen3_4b_ckpt_mcore(self, monkeypatch): - """Functional test for Qwen MCore checkpoint.""" + """Functional test for Qwen MCore Hybrid checkpoint.""" load_dir = MBRIDGE_CKPT if os.path.exists(MBRIDGE_CKPT) else None train_iters = 10 if load_dir else 5 @@ -76,7 +77,7 @@ def test_qwen3_4b_ckpt_mcore(self, monkeypatch): [ "torchrun", "--nproc_per_node=2", - "/opt/Megatron-Bridge/3rdparty/Megatron-LM/pretrain_gpt.py", + "/opt/Megatron-Bridge/3rdparty/Megatron-LM/pretrain_hybrid.py", "--init-method-std", "0.014", "--disable-bias-linear", @@ -88,8 +89,11 @@ def test_qwen3_4b_ckpt_mcore(self, monkeypatch): "--rotary-base", "1000000", "--use-rotary-position-embeddings", - "--num-layers", - "36", + "--hybrid-layer-pattern", + QWEN3_4B_HYBRID_PATTERN, + "--spec", + "megatron.core.models.hybrid.hybrid_layer_specs", + "hybrid_stack_spec", "--hidden-size", "2560", "--num-attention-heads", diff --git a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py index 4eb6f6add5..2c62c9a1cc 100644 --- a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py +++ b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py @@ -45,9 +45,11 @@ def _set_existing_attr(target: object, name: str, value: object) -> None: def _make_functional_test_model_small(model: object) -> None: # Keep this checkpoint-loading functional test far below runner memory limits. # The path under test is CP + sequence packing + pretrained checkpoint loading, - # not the full Llama 3.2 1B model shape. + # not the full Qwen3 model shape. Two logical Qwen blocks contain four physical + # HybridModel layers. for name, value in { - "num_layers": 2, + "num_layers": 4, + "hybrid_layer_pattern": "*-*-", "hidden_size": 128, "ffn_hidden_size": 512, "num_attention_heads": 4, diff --git a/tests/unit_tests/models/hybrid/test_hybrid_provider.py b/tests/unit_tests/models/hybrid/test_hybrid_provider.py index 4742ef7d6c..07e5ade321 100644 --- a/tests/unit_tests/models/hybrid/test_hybrid_provider.py +++ b/tests/unit_tests/models/hybrid/test_hybrid_provider.py @@ -58,6 +58,23 @@ def test_modelopt_spec_remaps_te_layernorm_keys(self): mock_fn.assert_called_once_with(local_core_attention=False, remap_te_layernorm=True) assert result is mock_spec + def test_modelopt_spec_preserves_grouped_moe_topology(self): + provider = HybridModelProvider( + hidden_size=128, + num_attention_heads=1, + num_moe_experts=4, + moe_grouped_gemm=True, + ) + mock_spec = Mock(spec=ModuleSpec) + with patch( + "megatron.bridge.models.hybrid.hybrid_provider.get_hybrid_stack_modelopt_spec", + return_value=mock_spec, + ) as mock_fn: + result = hybrid_provider.modelopt_hybrid_stack_spec(provider) + + mock_fn.assert_called_once_with(use_default_te_spec=True, moe_grouped_gemm=True) + assert result is mock_spec + def test_rejects_mamba_stack_spec_argument(self): module_spec = ModuleSpec(module=object) From eddbe88e3983d5e8b60ac12fa5cdb13d18bb1f05 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Thu, 23 Jul 2026 22:02:58 +0000 Subject: [PATCH 03/10] fix(qwen): preserve Hybrid YaRN and VPP configs Signed-off-by: Philip Petrakian --- .../bridge/models/hybrid/hybrid_provider.py | 9 ++- .../bridge/models/qwen/qwen_hybrid.py | 19 ++++--- .../models/qwen/test_qwen_hybrid.py | 56 +++++++++++++++++++ 3 files changed, 76 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/models/hybrid/hybrid_provider.py b/src/megatron/bridge/models/hybrid/hybrid_provider.py index 746f1e096c..060560b13c 100644 --- a/src/megatron/bridge/models/hybrid/hybrid_provider.py +++ b/src/megatron/bridge/models/hybrid/hybrid_provider.py @@ -135,10 +135,17 @@ class HybridModelProvider(TransformerConfig, ModelProviderMixin[MCoreHybridModel hybrid_layer_pattern: str | None = None seq_length: int = 8192 # HybridModel with no attention has no need for position embeddings, so none is default. - position_embedding_type: Literal["learned_absolute", "rope", "none"] = "none" + position_embedding_type: Literal["learned_absolute", "rope", "yarn", "none"] = "none" rotary_percent: float = 1.0 rotary_base: int = 10000 seq_len_interpolation_factor: float | None = None + yarn_rotary_scaling_factor: float | None = None + yarn_original_max_position_embeddings: int | None = None + yarn_beta_fast: float | None = None + yarn_beta_slow: float | None = None + yarn_mscale: float | None = None + yarn_mscale_all_dim: float | None = None + yarn_correction_range_round_to_int: bool | None = None apply_rope_fusion: bool = True make_vocab_size_divisible_by: int = 128 gated_linear_unit: bool = False diff --git a/src/megatron/bridge/models/qwen/qwen_hybrid.py b/src/megatron/bridge/models/qwen/qwen_hybrid.py index b5c30fa683..76c1407161 100644 --- a/src/megatron/bridge/models/qwen/qwen_hybrid.py +++ b/src/megatron/bridge/models/qwen/qwen_hybrid.py @@ -26,23 +26,27 @@ def qwen_pipeline_layer_pattern( hybrid_layer_pattern: str, pipeline_model_parallel_size: int, *, + virtual_pipeline_model_parallel_size: int | None = None, account_for_embedding: bool = False, account_for_loss: bool = False, ) -> str: - """Split a Qwen Hybrid pattern across PP stages without separating logical blocks.""" + """Split a Qwen Hybrid pattern across PP and VPP stages without separating logical blocks.""" if pipeline_model_parallel_size <= 0: raise ValueError("pipeline_model_parallel_size must be positive") + if virtual_pipeline_model_parallel_size is not None and virtual_pipeline_model_parallel_size <= 0: + raise ValueError("virtual_pipeline_model_parallel_size must be positive") main_pattern, separator, mtp_pattern = hybrid_layer_pattern.partition(Symbols.MTP_SEPARATOR) + pipeline_segment_count = pipeline_model_parallel_size * (virtual_pipeline_model_parallel_size or 1) if Symbols.PIPE in main_pattern: segment_count = main_pattern.count(Symbols.PIPE) + 1 - if segment_count != pipeline_model_parallel_size: + if segment_count != pipeline_segment_count: raise ValueError( f"Qwen hybrid_layer_pattern defines {segment_count} pipeline segments, " - f"but pipeline_model_parallel_size is {pipeline_model_parallel_size}." + f"but the PP/VPP configuration requires {pipeline_segment_count}." ) return hybrid_layer_pattern - if pipeline_model_parallel_size == 1: + if pipeline_segment_count == 1: return hybrid_layer_pattern if len(main_pattern) % 2: raise ValueError("Qwen Hybrid patterns must contain two physical layers per logical block.") @@ -57,9 +61,9 @@ def qwen_pipeline_layer_pattern( raise ValueError(f"Unsupported Qwen logical blocks in hybrid_layer_pattern: {invalid_blocks}") total_pipeline_units = len(logical_blocks) + int(account_for_embedding) + int(account_for_loss) - units_per_stage, extra_units = divmod(total_pipeline_units, pipeline_model_parallel_size) + units_per_stage, extra_units = divmod(total_pipeline_units, pipeline_segment_count) logical_blocks_per_stage = [ - units_per_stage + int(stage_index < extra_units) for stage_index in range(pipeline_model_parallel_size) + units_per_stage + int(stage_index < extra_units) for stage_index in range(pipeline_segment_count) ] logical_blocks_per_stage[0] -= int(account_for_embedding) logical_blocks_per_stage[-1] -= int(account_for_loss) @@ -87,7 +91,7 @@ class QwenHybridModelProvider(HybridModelProvider): def finalize(self) -> None: if ( self.hybrid_layer_pattern is not None - and self.pipeline_model_parallel_size > 1 + and self.pipeline_model_parallel_size * (self.virtual_pipeline_model_parallel_size or 1) > 1 and self.pipeline_model_parallel_layout is None and self.num_layers_in_first_pipeline_stage is None and self.num_layers_in_last_pipeline_stage is None @@ -95,6 +99,7 @@ def finalize(self) -> None: segmented_pattern = qwen_pipeline_layer_pattern( self.hybrid_layer_pattern, self.pipeline_model_parallel_size, + virtual_pipeline_model_parallel_size=self.virtual_pipeline_model_parallel_size, account_for_embedding=bool(self.account_for_embedding_in_pipeline_split), account_for_loss=bool(self.account_for_loss_in_pipeline_split), ) diff --git a/tests/unit_tests/models/qwen/test_qwen_hybrid.py b/tests/unit_tests/models/qwen/test_qwen_hybrid.py index 2365be855a..fc263bc12f 100644 --- a/tests/unit_tests/models/qwen/test_qwen_hybrid.py +++ b/tests/unit_tests/models/qwen/test_qwen_hybrid.py @@ -20,6 +20,7 @@ configure_qwen_hybrid_layers, qwen_pipeline_layer_pattern, ) +from megatron.bridge.training.utils.config_utils import _ConfigContainerBase def _provider() -> QwenHybridModelProvider: @@ -47,6 +48,38 @@ def test_mtp_pattern_is_deferred_when_mtp_is_disabled(): assert provider.mtp_hybrid_override_pattern == "*-" +def test_yarn_settings_are_serialized(): + provider = QwenHybridModelProvider( + num_layers=2, + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=4, + bf16=False, + position_embedding_type="yarn", + yarn_rotary_scaling_factor=3.125, + yarn_original_max_position_embeddings=40960, + yarn_beta_fast=32.0, + yarn_beta_slow=1.0, + yarn_mscale=1.0, + yarn_mscale_all_dim=1.0, + yarn_correction_range_round_to_int=False, + ) + expected_yarn_settings = { + "position_embedding_type": "yarn", + "yarn_rotary_scaling_factor": 3.125, + "yarn_original_max_position_embeddings": 40960, + "yarn_beta_fast": 32.0, + "yarn_beta_slow": 1.0, + "yarn_mscale": 1.0, + "yarn_mscale_all_dim": 1.0, + "yarn_correction_range_round_to_int": False, + } + + serialized_provider = _ConfigContainerBase._convert_value_to_dict(provider) + + assert {key: serialized_provider[key] for key in expected_yarn_settings} == expected_yarn_settings + + def test_mtp_pattern_honors_recipe_override_set_after_conversion(): provider = _provider() configure_qwen_hybrid_layers( @@ -90,6 +123,29 @@ def test_pipeline_segmentation_preserves_logical_blocks_and_embedding_loss_balan assert provider.account_for_loss_in_pipeline_split is False +def test_pipeline_segmentation_accounts_for_virtual_pipeline_stages(): + provider = QwenHybridModelProvider( + num_layers=16, + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=4, + pipeline_model_parallel_size=2, + virtual_pipeline_model_parallel_size=4, + bf16=False, + ) + configure_qwen_hybrid_layers( + provider, + num_logical_layers=16, + mlp_symbols=Symbols.MOE, + ) + + provider.finalize() + + segments = provider.hybrid_layer_pattern.split(Symbols.PIPE) + assert len(segments) == 8 + assert all(segment == "*E*E" for segment in segments) + + def test_pipeline_segmentation_rejects_stale_explicit_segments(): with pytest.raises(ValueError, match="defines 2 pipeline segments"): qwen_pipeline_layer_pattern("*-*-|*-*-", pipeline_model_parallel_size=1) From f6d8cf7e20e19c35c572dd2dd318ef29d77d9e57 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sun, 12 Jul 2026 23:20:26 +0000 Subject: [PATCH 04/10] feat(qwen): migrate Qwen3 dense and MoE models to HybridModel Signed-off-by: Philip Petrakian --- .../bridge/models/hybrid/hybrid_provider.py | 10 +- .../bridge/models/qwen/qwen3_bridge.py | 130 +++++++---- .../bridge/models/qwen/qwen3_moe_bridge.py | 139 +++++++---- .../bridge/models/qwen/qwen_hybrid.py | 220 ++++++++++++++++++ .../recipes/test_qwen_recipes_pretrain.py | 2 +- .../models/qwen/test_qwen3_bridge.py | 71 ++++-- .../models/qwen/test_qwen3_moe_bridge.py | 25 +- .../models/qwen/test_qwen_hybrid.py | 95 ++++++++ 8 files changed, 573 insertions(+), 119 deletions(-) create mode 100644 src/megatron/bridge/models/qwen/qwen_hybrid.py create mode 100644 tests/unit_tests/models/qwen/test_qwen_hybrid.py diff --git a/src/megatron/bridge/models/hybrid/hybrid_provider.py b/src/megatron/bridge/models/hybrid/hybrid_provider.py index 36aa49c143..df1cb92324 100644 --- a/src/megatron/bridge/models/hybrid/hybrid_provider.py +++ b/src/megatron/bridge/models/hybrid/hybrid_provider.py @@ -193,9 +193,13 @@ def finalize(self) -> None: self.hybrid_override_pattern = None used_hybrid_override_pattern = True - # Combine hybrid_layer_pattern (main decoder) with mtp_hybrid_override_pattern - # into a single unified pattern that MCore HybridModel can parse. - if self.hybrid_layer_pattern is not None and self.mtp_hybrid_override_pattern: + # Combine the main and MTP patterns only when MTP is enabled. Providers may + # predeclare an MTP override before a recipe sets mtp_num_layers. + if ( + self.hybrid_layer_pattern is not None + and self.mtp_hybrid_override_pattern + and (self.mtp_num_layers or self.mtp_use_repeated_layer) + ): sep = Symbols.MTP_SEPARATOR main_pattern = self.hybrid_layer_pattern.split(sep)[0] # When mtp_use_repeated_layer=True, the shared MTP layer always exists diff --git a/src/megatron/bridge/models/qwen/qwen3_bridge.py b/src/megatron/bridge/models/qwen/qwen3_bridge.py index e5b1442f1b..d3ca04dca5 100644 --- a/src/megatron/bridge/models/qwen/qwen3_bridge.py +++ b/src/megatron/bridge/models/qwen/qwen3_bridge.py @@ -13,7 +13,8 @@ # limitations under the License. import torch -from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.models.hybrid.hybrid_model import HybridModel from transformers import Qwen3ForCausalLM from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry @@ -23,15 +24,26 @@ GatedMLPMapping, QKVMapping, ) +from megatron.bridge.models.qwen.qwen_hybrid import ( + QwenHybridModelProvider, + configure_qwen_hybrid_layers, + qwen_logical_layer_count, + qwen_physical_layer_indices, +) -@MegatronModelBridge.register_bridge(source=Qwen3ForCausalLM, target=GPTModel, model_type="qwen3") +@MegatronModelBridge.register_bridge( + source=Qwen3ForCausalLM, + target=HybridModel, + provider=QwenHybridModelProvider, + model_type="qwen3", +) class Qwen3Bridge(MegatronModelBridge): """ Megatron Bridge for Qwen3 Causal LM. This bridge handles the conversion between HuggingFace Qwen3ForCausalLM - and Megatron-Core GPTModel formats. Qwen3 differs from Qwen2 by using + and Megatron-Core HybridModel formats. Qwen3 differs from Qwen2 by using QK layernorm and no QKV bias. Example: @@ -40,9 +52,19 @@ class Qwen3Bridge(MegatronModelBridge): >>> provider = bridge.to_megatron_provider() """ + @classmethod + def megatron_to_hf_config(cls, provider) -> dict: + """Convert a Hybrid Qwen3 provider to a Hugging Face config dictionary.""" + hf_config = super().megatron_to_hf_config(provider) + logical_layer_count = qwen_logical_layer_count(provider.hybrid_layer_pattern) + if logical_layer_count is not None: + hf_config["num_hidden_layers"] = logical_layer_count + return hf_config + def provider_bridge(self, hf_pretrained): - """Convert HuggingFace Qwen3 config to GPTModelProvider.""" + """Convert a Hugging Face Qwen3 config to HybridModelProvider.""" provider = super().provider_bridge(hf_pretrained) + hf_config = hf_pretrained.config provider.normalization = "RMSNorm" provider.gated_linear_unit = True @@ -51,6 +73,14 @@ def provider_bridge(self, hf_pretrained): provider.hidden_dropout = 0.0 provider.qk_layernorm = True # Qwen3 uses QK layernorm provider.autocast_dtype = torch.bfloat16 + provider.share_embeddings_and_output_weights = getattr(hf_config, "tie_word_embeddings", True) + + configure_qwen_hybrid_layers( + provider, + num_logical_layers=hf_config.num_hidden_layers, + mlp_symbols=Symbols.MLP, + mtp_mlp_symbol=Symbols.MLP, + ) return provider @@ -65,67 +95,85 @@ def mapping_registry(self) -> MegatronMappingRegistry: The fused QKV matrix is handled by :class:`QKVMapping` and the gated MLP gate+up projection by :class:`GatedMLPMapping`. """ - # Dictionary maps Megatron parameter names -> HF parameter names - # Supports wildcard (*) patterns for layer-specific parameters param_mappings = { # Embedding and output "embedding.word_embeddings.weight": "model.embed_tokens.weight", "output_layer.weight": "lm_head.weight", - "decoder.final_layernorm.weight": "model.norm.weight", - # Decoder layer attention norms and projections - "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight", - "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.layers.*.post_attention_layernorm.weight", - "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_norm.weight", # Qwen3 specific - "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight", # Qwen3 specific - "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight", - "decoder.layers.*.mlp.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight", + "decoder.final_norm.weight": "model.norm.weight", # MTP projection and norms (used when mtp_num_layers >= 1) "mtp.layers.0.eh_proj.weight": "mtp.fc.weight", "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 transformer layer attention - "mtp.layers.0.mtp_model_layer.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_qkv.layer_norm_weight": "mtp.layers.0.input_layernorm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.q_layernorm.weight": "mtp.layers.0.self_attn.q_norm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.k_layernorm.weight": "mtp.layers.0.self_attn.k_norm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_proj.weight": "mtp.layers.0.self_attn.o_proj.weight", # MTP transformer layer MLP - "mtp.layers.0.mtp_model_layer.mlp.linear_fc1.layer_norm_weight": "mtp.layers.0.post_attention_layernorm.weight", - "mtp.layers.0.mtp_model_layer.mlp.linear_fc2.weight": "mtp.layers.0.mlp.down_proj.weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc1.layer_norm_weight": "mtp.layers.0.post_attention_layernorm.weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc2.weight": "mtp.layers.0.mlp.down_proj.weight", } - mapping_list = [] - # Convert each dictionary entry to AutoMapping(megatron_param, hf_param) - for megatron_param, hf_param in param_mappings.items(): - mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) + mapping_list = [AutoMapping(megatron_param=k, hf_param=v) for k, v in param_mappings.items()] + + for logical_layer_idx in range(self.hf_config.num_hidden_layers): + attention_layer_idx, mlp_layer_idx = qwen_physical_layer_indices(logical_layer_idx) + hf_layer = f"model.layers.{logical_layer_idx}" + attention_layer = f"decoder.layers.{attention_layer_idx}.self_attention" + mlp_layer = f"decoder.layers.{mlp_layer_idx}.mlp" + mapping_list.extend( + [ + AutoMapping( + megatron_param=f"{attention_layer}.linear_qkv.layer_norm_weight", + hf_param=f"{hf_layer}.input_layernorm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.q_layernorm.weight", + hf_param=f"{hf_layer}.self_attn.q_norm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.k_layernorm.weight", + hf_param=f"{hf_layer}.self_attn.k_norm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.linear_proj.weight", + hf_param=f"{hf_layer}.self_attn.o_proj.weight", + ), + AutoMapping( + megatron_param=f"{mlp_layer}.linear_fc1.layer_norm_weight", + hf_param=f"{hf_layer}.post_attention_layernorm.weight", + ), + AutoMapping( + megatron_param=f"{mlp_layer}.linear_fc2.weight", + hf_param=f"{hf_layer}.mlp.down_proj.weight", + ), + QKVMapping( + megatron_param=f"{attention_layer}.linear_qkv.weight", + q=f"{hf_layer}.self_attn.q_proj.weight", + k=f"{hf_layer}.self_attn.k_proj.weight", + v=f"{hf_layer}.self_attn.v_proj.weight", + ), + GatedMLPMapping( + megatron_param=f"{mlp_layer}.linear_fc1.weight", + gate=f"{hf_layer}.mlp.gate_proj.weight", + up=f"{hf_layer}.mlp.up_proj.weight", + ), + ] + ) - # Add special mappings that require parameter concatenation/transformation mapping_list.extend( [ - # QKV: Combine separate Q, K, V matrices into single QKV matrix - # Note: Qwen3 does NOT have bias in QKV projections (unlike Qwen2) - QKVMapping( - megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", - q="model.layers.*.self_attn.q_proj.weight", - k="model.layers.*.self_attn.k_proj.weight", - v="model.layers.*.self_attn.v_proj.weight", - ), - # Gated MLP: Combine gate and up projection matrices into single FC1 matrix - GatedMLPMapping( - megatron_param="decoder.layers.*.mlp.linear_fc1.weight", - gate="model.layers.*.mlp.gate_proj.weight", - up="model.layers.*.mlp.up_proj.weight", - ), # MTP QKV: same split/merge as decoder layers QKVMapping( - megatron_param="mtp.layers.*.mtp_model_layer.self_attention.linear_qkv.weight", + megatron_param="mtp.layers.*.mtp_model_layer.layers.0.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", ), # MTP Gated MLP GatedMLPMapping( - megatron_param="mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight", + megatron_param="mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc1.weight", gate="mtp.layers.0.mlp.gate_proj.weight", up="mtp.layers.0.mlp.up_proj.weight", ), diff --git a/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py b/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py index a6476f4f49..1f9fa5a46c 100755 --- a/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py +++ b/src/megatron/bridge/models/qwen/qwen3_moe_bridge.py @@ -13,7 +13,8 @@ # limitations under the License. import torch -from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.models.hybrid.hybrid_model import HybridModel from transformers import Qwen3MoeForCausalLM from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry @@ -23,15 +24,26 @@ GatedMLPMapping, QKVMapping, ) +from megatron.bridge.models.qwen.qwen_hybrid import ( + QwenHybridModelProvider, + configure_qwen_hybrid_layers, + qwen_logical_layer_count, + qwen_physical_layer_indices, +) -@MegatronModelBridge.register_bridge(source=Qwen3MoeForCausalLM, target=GPTModel, model_type="qwen3_moe") +@MegatronModelBridge.register_bridge( + source=Qwen3MoeForCausalLM, + target=HybridModel, + provider=QwenHybridModelProvider, + model_type="qwen3_moe", +) class Qwen3MoEBridge(MegatronModelBridge): """ Megatron Bridge for Qwen3 MoE Causal LM. This bridge handles the conversion between HuggingFace Qwen3MoeForCausalLM - and Megatron-Core GPTModel formats. Qwen3 MoE models use mixture of experts + and Megatron-Core HybridModel formats. Qwen3 MoE models use mixture of experts architecture with QK layernorm. Example: @@ -44,13 +56,17 @@ class Qwen3MoEBridge(MegatronModelBridge): def megatron_to_hf_config(cls, provider) -> dict: """Convert Megatron provider config to HuggingFace Qwen3MoeConfig dict.""" hf_config = super().megatron_to_hf_config(provider) + logical_layer_count = qwen_logical_layer_count(provider.hybrid_layer_pattern) + if logical_layer_count is not None: + hf_config["num_hidden_layers"] = logical_layer_count hf_config["decoder_sparse_step"] = 1 # All layers are MoE in Qwen3 MoE hf_config["norm_topk_prob"] = not provider.moe_router_pre_softmax return hf_config def provider_bridge(self, hf_pretrained): - """Convert HuggingFace Qwen3 MoE config to GPTModelProvider.""" + """Convert a Hugging Face Qwen3 MoE config to HybridModelProvider.""" provider = super().provider_bridge(hf_pretrained) + hf_config = hf_pretrained.config provider.normalization = "RMSNorm" provider.gated_linear_unit = True @@ -66,6 +82,14 @@ def provider_bridge(self, hf_pretrained): provider.moe_router_pre_softmax = not hf_pretrained.config.norm_topk_prob provider.moe_token_dispatcher_type = "alltoall" provider.moe_permute_fusion = True + provider.share_embeddings_and_output_weights = getattr(hf_config, "tie_word_embeddings", True) + + configure_qwen_hybrid_layers( + provider, + num_logical_layers=hf_config.num_hidden_layers, + mlp_symbols=Symbols.MOE, + mtp_mlp_symbol=Symbols.MOE, + ) return provider @@ -73,57 +97,70 @@ def mapping_registry(self) -> MegatronMappingRegistry: # Return MegatronMappingRegistry containing parameter mappings from Megatron to HF format # First create simple 1:1 parameter mappings using a dictionary for readability - # Dictionary maps Megatron parameter names -> HF parameter names - # Supports wildcard (*) patterns for layer-specific parameters param_mappings = { "embedding.word_embeddings.weight": "model.embed_tokens.weight", "output_layer.weight": "lm_head.weight", - "decoder.final_layernorm.weight": "model.norm.weight", - "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight", - "decoder.layers.*.mlp.router.weight": "model.layers.*.mlp.gate.weight", - "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.post_attention_layernorm.weight", - "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_norm.weight", - "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight", - "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight", + "decoder.final_norm.weight": "model.norm.weight", } - mapping_list = [] - # Convert each dictionary entry to AutoMapping(megatron_param, hf_param) - for megatron_param, hf_param in param_mappings.items(): - mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) - - # Add special mappings that require parameter concatenation/transformation - mapping_list.extend( - [ - # QKV: Combine separate Q, K, V matrices into single QKV matrix - # Note: Qwen3 MoE does NOT have bias in QKV projections - QKVMapping( - megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", - q="model.layers.*.self_attn.q_proj.weight", - k="model.layers.*.self_attn.k_proj.weight", - v="model.layers.*.self_attn.v_proj.weight", - ), - # Expert mappings for TEGroupedMLP - GatedMLPMapping( - megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", - gate="model.layers.*.mlp.experts.*.gate_proj.weight", - up="model.layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", - hf_param="model.layers.*.mlp.experts.*.down_proj.weight", - ), - # Expert mappings for SequentialMLP (used by quantization) - GatedMLPMapping( - megatron_param="decoder.layers.*.mlp.experts.local_experts.*.linear_fc1.weight", - gate="model.layers.*.mlp.experts.*.gate_proj.weight", - up="model.layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param="decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight", - hf_param="model.layers.*.mlp.experts.*.down_proj.weight", - ), - ] - ) + mapping_list = [AutoMapping(megatron_param=k, hf_param=v) for k, v in param_mappings.items()] + + for logical_layer_idx in range(self.hf_config.num_hidden_layers): + attention_layer_idx, moe_layer_idx = qwen_physical_layer_indices(logical_layer_idx) + hf_layer = f"model.layers.{logical_layer_idx}" + attention_layer = f"decoder.layers.{attention_layer_idx}.self_attention" + moe_layer = f"decoder.layers.{moe_layer_idx}" + mapping_list.extend( + [ + AutoMapping( + megatron_param=f"{attention_layer}.linear_qkv.layer_norm_weight", + hf_param=f"{hf_layer}.input_layernorm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.q_layernorm.weight", + hf_param=f"{hf_layer}.self_attn.q_norm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.k_layernorm.weight", + hf_param=f"{hf_layer}.self_attn.k_norm.weight", + ), + AutoMapping( + megatron_param=f"{attention_layer}.linear_proj.weight", + hf_param=f"{hf_layer}.self_attn.o_proj.weight", + ), + AutoMapping( + megatron_param=f"{moe_layer}.pre_mlp_layernorm.weight", + hf_param=f"{hf_layer}.post_attention_layernorm.weight", + ), + AutoMapping( + megatron_param=f"{moe_layer}.mlp.router.weight", + hf_param=f"{hf_layer}.mlp.gate.weight", + ), + QKVMapping( + megatron_param=f"{attention_layer}.linear_qkv.weight", + q=f"{hf_layer}.self_attn.q_proj.weight", + k=f"{hf_layer}.self_attn.k_proj.weight", + v=f"{hf_layer}.self_attn.v_proj.weight", + ), + GatedMLPMapping( + megatron_param=f"{moe_layer}.mlp.experts.linear_fc1.weight*", + gate=f"{hf_layer}.mlp.experts.*.gate_proj.weight", + up=f"{hf_layer}.mlp.experts.*.up_proj.weight", + ), + AutoMapping( + megatron_param=f"{moe_layer}.mlp.experts.linear_fc2.weight*", + hf_param=f"{hf_layer}.mlp.experts.*.down_proj.weight", + ), + GatedMLPMapping( + megatron_param=f"{moe_layer}.mlp.experts.local_experts.*.linear_fc1.weight", + gate=f"{hf_layer}.mlp.experts.*.gate_proj.weight", + up=f"{hf_layer}.mlp.experts.*.up_proj.weight", + ), + AutoMapping( + megatron_param=f"{moe_layer}.mlp.experts.local_experts.*.linear_fc2.weight", + hf_param=f"{hf_layer}.mlp.experts.*.down_proj.weight", + ), + ] + ) return MegatronMappingRegistry(*mapping_list) diff --git a/src/megatron/bridge/models/qwen/qwen_hybrid.py b/src/megatron/bridge/models/qwen/qwen_hybrid.py new file mode 100644 index 0000000000..b5c30fa683 --- /dev/null +++ b/src/megatron/bridge/models/qwen/qwen_hybrid.py @@ -0,0 +1,220 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared HybridModel layout helpers for Qwen3 and newer models.""" + +from collections.abc import Sequence +from dataclasses import dataclass + +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider + + +def qwen_pipeline_layer_pattern( + hybrid_layer_pattern: str, + pipeline_model_parallel_size: int, + *, + account_for_embedding: bool = False, + account_for_loss: bool = False, +) -> str: + """Split a Qwen Hybrid pattern across PP stages without separating logical blocks.""" + if pipeline_model_parallel_size <= 0: + raise ValueError("pipeline_model_parallel_size must be positive") + + main_pattern, separator, mtp_pattern = hybrid_layer_pattern.partition(Symbols.MTP_SEPARATOR) + if Symbols.PIPE in main_pattern: + segment_count = main_pattern.count(Symbols.PIPE) + 1 + if segment_count != pipeline_model_parallel_size: + raise ValueError( + f"Qwen hybrid_layer_pattern defines {segment_count} pipeline segments, " + f"but pipeline_model_parallel_size is {pipeline_model_parallel_size}." + ) + return hybrid_layer_pattern + if pipeline_model_parallel_size == 1: + return hybrid_layer_pattern + if len(main_pattern) % 2: + raise ValueError("Qwen Hybrid patterns must contain two physical layers per logical block.") + + logical_blocks = [main_pattern[index : index + 2] for index in range(0, len(main_pattern), 2)] + attention_symbols = {Symbols.ATTENTION, Symbols.GDN} + mlp_symbols = {Symbols.MLP, Symbols.MOE} + invalid_blocks = [ + block for block in logical_blocks if block[0] not in attention_symbols or block[1] not in mlp_symbols + ] + if invalid_blocks: + raise ValueError(f"Unsupported Qwen logical blocks in hybrid_layer_pattern: {invalid_blocks}") + + total_pipeline_units = len(logical_blocks) + int(account_for_embedding) + int(account_for_loss) + units_per_stage, extra_units = divmod(total_pipeline_units, pipeline_model_parallel_size) + logical_blocks_per_stage = [ + units_per_stage + int(stage_index < extra_units) for stage_index in range(pipeline_model_parallel_size) + ] + logical_blocks_per_stage[0] -= int(account_for_embedding) + logical_blocks_per_stage[-1] -= int(account_for_loss) + if any(count <= 0 for count in logical_blocks_per_stage): + raise ValueError( + "Every pipeline stage must receive at least one Qwen logical block after embedding/loss balancing." + ) + + segments = [] + block_offset = 0 + for block_count in logical_blocks_per_stage: + segments.append("".join(logical_blocks[block_offset : block_offset + block_count])) + block_offset += block_count + if block_offset != len(logical_blocks): + raise RuntimeError("Failed to assign every Qwen logical block to a pipeline stage.") + + segmented_pattern = Symbols.PIPE.join(segments) + return segmented_pattern + (separator + mtp_pattern if separator else "") + + +@dataclass +class QwenHybridModelProvider(HybridModelProvider): + """HybridModel provider that keeps each Qwen logical block on one PP stage.""" + + def finalize(self) -> None: + if ( + self.hybrid_layer_pattern is not None + and self.pipeline_model_parallel_size > 1 + and self.pipeline_model_parallel_layout is None + and self.num_layers_in_first_pipeline_stage is None + and self.num_layers_in_last_pipeline_stage is None + ): + segmented_pattern = qwen_pipeline_layer_pattern( + self.hybrid_layer_pattern, + self.pipeline_model_parallel_size, + account_for_embedding=bool(self.account_for_embedding_in_pipeline_split), + account_for_loss=bool(self.account_for_loss_in_pipeline_split), + ) + if segmented_pattern != self.hybrid_layer_pattern: + self.hybrid_layer_pattern = segmented_pattern + # The explicit segments already include embedding/loss-aware balancing. + self.account_for_embedding_in_pipeline_split = False + self.account_for_loss_in_pipeline_split = False + super().finalize() + + +def qwen_attention_symbols( + num_layers: int, + linear_attention_freq: int | Sequence[int] | None = None, +) -> list[str]: + """Translate a logical Qwen attention schedule to HybridModel symbols.""" + if linear_attention_freq is None: + return [Symbols.ATTENTION] * num_layers + + if isinstance(linear_attention_freq, int): + if linear_attention_freq <= 0: + raise ValueError("linear_attention_freq must be positive") + return [ + Symbols.ATTENTION if (layer_idx + 1) % linear_attention_freq == 0 else Symbols.GDN + for layer_idx in range(num_layers) + ] + + linear_attention_pattern = list(linear_attention_freq) + if len(linear_attention_pattern) != num_layers: + raise ValueError( + "linear_attention_freq has " + f"{len(linear_attention_pattern)} entries, but num_hidden_layers is {num_layers}." + ) + invalid_values = sorted(set(linear_attention_pattern) - {0, 1}) + if invalid_values: + raise ValueError(f"Unsupported linear attention pattern values: {invalid_values}. Expected only 0 or 1.") + return [Symbols.GDN if is_linear else Symbols.ATTENTION for is_linear in linear_attention_pattern] + + +def qwen_hybrid_layer_pattern( + num_layers: int, + *, + mlp_symbols: str | Sequence[str], + linear_attention_freq: int | Sequence[int] | None = None, +) -> str: + """Build a two-physical-layer HybridModel pattern for each logical Qwen block.""" + if isinstance(mlp_symbols, str): + mlp_pattern = [mlp_symbols] * num_layers + else: + mlp_pattern = list(mlp_symbols) + if len(mlp_pattern) != num_layers: + raise ValueError(f"MLP pattern has {len(mlp_pattern)} entries, but num_hidden_layers is {num_layers}.") + + invalid_symbols = sorted(set(mlp_pattern) - {Symbols.MLP, Symbols.MOE}) + if invalid_symbols: + raise ValueError( + f"Unsupported Qwen MLP symbols: {invalid_symbols}. Expected '{Symbols.MLP}' or '{Symbols.MOE}'." + ) + + attention_pattern = qwen_attention_symbols(num_layers, linear_attention_freq) + return "".join( + attention_symbol + mlp_symbol for attention_symbol, mlp_symbol in zip(attention_pattern, mlp_pattern) + ) + + +def qwen_moe_layer_symbols( + num_layers: int, + *, + decoder_sparse_step: int = 1, + mlp_only_layers: Sequence[int] = (), +) -> list[str]: + """Translate Hugging Face Qwen MoE placement fields to Hybrid symbols.""" + if decoder_sparse_step <= 0: + raise ValueError("decoder_sparse_step must be positive") + dense_layers = set(mlp_only_layers) + invalid_layers = sorted(layer_idx for layer_idx in dense_layers if not 0 <= layer_idx < num_layers) + if invalid_layers: + raise ValueError(f"mlp_only_layers contains out-of-range indices: {invalid_layers}") + return [ + Symbols.MOE + if logical_layer_idx not in dense_layers and (logical_layer_idx + 1) % decoder_sparse_step == 0 + else Symbols.MLP + for logical_layer_idx in range(num_layers) + ] + + +def configure_qwen_hybrid_layers( + provider: HybridModelProvider, + *, + num_logical_layers: int, + mlp_symbols: str | Sequence[str], + linear_attention_freq: int | Sequence[int] | None = None, + mtp_mlp_symbol: str | None = None, +) -> None: + """Configure main and optional MTP physical layer patterns on a Qwen provider.""" + provider.hybrid_layer_pattern = qwen_hybrid_layer_pattern( + num_logical_layers, + mlp_symbols=mlp_symbols, + linear_attention_freq=linear_attention_freq, + ) + provider.num_layers = len(provider.hybrid_layer_pattern) + + if mtp_mlp_symbol is not None: + if mtp_mlp_symbol not in {Symbols.MLP, Symbols.MOE}: + raise ValueError( + f"Unsupported Qwen MTP MLP symbol: {mtp_mlp_symbol}. Expected '{Symbols.MLP}' or '{Symbols.MOE}'." + ) + provider.mtp_hybrid_override_pattern = Symbols.ATTENTION + mtp_mlp_symbol + + +def qwen_logical_layer_count(hybrid_layer_pattern: str | None) -> int | None: + """Return the number of logical Qwen blocks encoded by a HybridModel pattern.""" + if not hybrid_layer_pattern: + return None + main_pattern = hybrid_layer_pattern.split(Symbols.MTP_SEPARATOR)[0].replace(Symbols.PIPE, "") + attention_symbols = {Symbols.ATTENTION, Symbols.GDN} + return sum(symbol in attention_symbols for symbol in main_pattern) + + +def qwen_physical_layer_indices(logical_layer_idx: int) -> tuple[int, int]: + """Return attention and MLP/MoE physical indices for one logical Qwen block.""" + attention_layer_idx = 2 * logical_layer_idx + return attention_layer_idx, attention_layer_idx + 1 diff --git a/tests/functional_tests/test_groups/recipes/test_qwen_recipes_pretrain.py b/tests/functional_tests/test_groups/recipes/test_qwen_recipes_pretrain.py index 6550d86b56..6a96db259a 100644 --- a/tests/functional_tests/test_groups/recipes/test_qwen_recipes_pretrain.py +++ b/tests/functional_tests/test_groups/recipes/test_qwen_recipes_pretrain.py @@ -28,7 +28,7 @@ QWEN_PRETRAIN_RECIPES = [ # (config_func, name, parallelism_overrides, model_overrides) (qwen25_500m_config, "qwen25_500m", {}, {"num_layers": 2}), - (qwen3_600m_config, "qwen3_600m", {}, {"num_layers": 2}), + (qwen3_600m_config, "qwen3_600m", {}, {"num_layers": 4, "hybrid_layer_pattern": "*-*-"}), ] diff --git a/tests/unit_tests/models/qwen/test_qwen3_bridge.py b/tests/unit_tests/models/qwen/test_qwen3_bridge.py index 8bf74831c2..deeacc472d 100644 --- a/tests/unit_tests/models/qwen/test_qwen3_bridge.py +++ b/tests/unit_tests/models/qwen/test_qwen3_bridge.py @@ -14,6 +14,7 @@ import tempfile from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock, patch import pytest @@ -23,11 +24,20 @@ from megatron.bridge.models import AutoBridge from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from megatron.bridge.models.conversion.transformers_compat import rope_theta_from_hf -from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.models.qwen.qwen3_bridge import Qwen3Bridge +@pytest.fixture(autouse=True) +def _set_bridge_hf_config(): + """Give explicit mapping tests a small logical Qwen layer count.""" + previous = Qwen3Bridge.hf_config + Qwen3Bridge.hf_config = SimpleNamespace(num_hidden_layers=2) + yield + Qwen3Bridge.hf_config = previous + + class TestMegatronQwen3Bridge: """Test cases for MegatronQwen3Bridge class.""" @@ -93,11 +103,11 @@ def test_provider_bridge_basic(self, mock_pretrained_qwen3, qwen3_config): # Call provider_bridge result = bridge.provider_bridge(mock_pretrained_qwen3) - # Check that it returns a GPTModelProvider instance (after refactoring) - assert isinstance(result, GPTModelProvider) + assert isinstance(result, HybridModelProvider) # Check basic configuration mapping - assert result.num_layers == qwen3_config.num_hidden_layers + assert result.num_layers == 2 * qwen3_config.num_hidden_layers + assert result.hybrid_layer_pattern == "*-" * qwen3_config.num_hidden_layers assert result.hidden_size == qwen3_config.hidden_size assert result.num_attention_heads == qwen3_config.num_attention_heads assert result.seq_length == qwen3_config.max_position_embeddings @@ -200,8 +210,7 @@ def test_provider_bridge_with_custom_kwargs(self, mock_pretrained_qwen3): # Pass model only result = bridge.provider_bridge(mock_pretrained_qwen3) - # Just verify that we got a valid GPTModelProvider - assert isinstance(result, GPTModelProvider) + assert isinstance(result, HybridModelProvider) def test_provider_bridge_without_tie_embeddings(self, qwen3_config): """Test provider_bridge when tie_word_embeddings is not present.""" @@ -393,7 +402,7 @@ def test_from_pretrained_multiple_models(self, mock_autoconfig, mock_pretrained, "megatron.bridge.models.conversion.auto_bridge.model_bridge.get_model_bridge" ) as mock_get_bridge: mock_bridge = Mock() - mock_provider = Mock(spec=GPTModelProvider) + mock_provider = Mock(spec=HybridModelProvider) mock_bridge.provider_bridge.return_value = mock_provider mock_get_bridge.return_value = mock_bridge @@ -448,6 +457,38 @@ def test_supports_qwen3_architectures(self, qwen3_configs): assert AutoBridge.supports(non_causal_config) == False +class TestQwen3BridgeParameterMapping: + """Test parameter mapping functionality in Qwen3Bridge.""" + + def test_mapping_registry_has_qwen3_specific_mappings(self): + """Test that mapping registry includes Qwen3-specific QK norm mappings.""" + bridge = Qwen3Bridge() + mapping_registry = bridge.mapping_registry() + + mapping = mapping_registry.megatron_to_hf_lookup("decoder.layers.0.self_attention.q_layernorm.weight") + assert mapping is not None + assert mapping.hf_param == "model.layers.0.self_attn.q_norm.weight" + + def test_qwen3_qk_norm_mapping_difference(self): + """Test that Qwen3 bridge includes QK norm mappings not present in Qwen2.""" + bridge = Qwen3Bridge() + mapping_registry = bridge.mapping_registry() + + attention_mapping = mapping_registry.megatron_to_hf_lookup( + "decoder.layers.0.self_attention.k_layernorm.weight" + ) + mlp_mapping = mapping_registry.megatron_to_hf_lookup("decoder.layers.1.mlp.linear_fc2.weight") + assert attention_mapping.hf_param == "model.layers.0.self_attn.k_norm.weight" + assert mlp_mapping.hf_param == "model.layers.0.mlp.down_proj.weight" + + def test_qwen3_no_qkv_bias_mapping(self): + """Test that Qwen3 bridge doesn't include QKV bias mappings.""" + bridge = Qwen3Bridge() + mapping_registry = bridge.mapping_registry() + + assert mapping_registry.megatron_to_hf_lookup("decoder.layers.0.self_attention.linear_qkv.bias") is None + + class TestQwen3BridgeMTPMapping: """Tests for MTP (Multi-Token Prediction) weight mappings in Qwen3Bridge. @@ -466,15 +507,15 @@ class TestQwen3BridgeMTPMapping: "mtp.layers.0.enorm.weight", "mtp.layers.0.hnorm.weight", "mtp.layers.0.final_layernorm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.linear_qkv.layer_norm_weight", - "mtp.layers.0.mtp_model_layer.self_attention.q_layernorm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.k_layernorm.weight", - "mtp.layers.0.mtp_model_layer.self_attention.linear_proj.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_qkv.layer_norm_weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.q_layernorm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.k_layernorm.weight", + "mtp.layers.0.mtp_model_layer.layers.0.self_attention.linear_proj.weight", # QKVMapping stores a wildcard pattern, not a concrete layer index - "mtp.layers.*.mtp_model_layer.self_attention.linear_qkv.weight", - "mtp.layers.0.mtp_model_layer.mlp.linear_fc1.layer_norm_weight", - "mtp.layers.0.mtp_model_layer.mlp.linear_fc1.weight", - "mtp.layers.0.mtp_model_layer.mlp.linear_fc2.weight", + "mtp.layers.*.mtp_model_layer.layers.0.self_attention.linear_qkv.weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc1.layer_norm_weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc1.weight", + "mtp.layers.0.mtp_model_layer.layers.1.mlp.linear_fc2.weight", ) def _get_all_megatron_params(self, mapping_registry): diff --git a/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py b/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py index e9f23d16e6..eb28997d87 100644 --- a/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py +++ b/tests/unit_tests/models/qwen/test_qwen3_moe_bridge.py @@ -16,17 +16,26 @@ Unit tests for Qwen3 MoE bridge functionality. """ +from types import SimpleNamespace from unittest.mock import Mock import pytest import torch from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge -from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.models.qwen.qwen3_moe_bridge import Qwen3MoEBridge +@pytest.fixture(autouse=True) +def _set_bridge_hf_config(): + previous = Qwen3MoEBridge.hf_config + Qwen3MoEBridge.hf_config = SimpleNamespace(num_hidden_layers=2) + yield + Qwen3MoEBridge.hf_config = previous + + class TestQwen3MoEBridge: """Test cases for Qwen3MoEBridge class.""" @@ -124,11 +133,11 @@ def test_provider_bridge_basic(self, mock_pretrained_qwen3_moe, mock_qwen3_moe_c # Call provider_bridge result = bridge.provider_bridge(mock_pretrained_qwen3_moe) - # Check that it returns a GPTModelProvider instance (after refactoring) - assert isinstance(result, GPTModelProvider) + assert isinstance(result, HybridModelProvider) # Check basic configuration mapping - assert result.num_layers == mock_qwen3_moe_config.num_hidden_layers + assert result.num_layers == 2 * mock_qwen3_moe_config.num_hidden_layers + assert result.hybrid_layer_pattern == "*E" * mock_qwen3_moe_config.num_hidden_layers assert result.hidden_size == mock_qwen3_moe_config.hidden_size assert result.num_attention_heads == mock_qwen3_moe_config.num_attention_heads assert result.seq_length == mock_qwen3_moe_config.max_position_embeddings @@ -278,7 +287,7 @@ def test_provider_bridge_missing_tie_word_embeddings(self, mock_qwen3_moe_config bridge = Qwen3MoEBridge() result = bridge.provider_bridge(mock_pretrained) - # GPTModelProvider defaults share_embeddings_and_output_weights to True + # Preserve the historical default when the HF field is absent. assert result.share_embeddings_and_output_weights is True def test_provider_bridge_235b_config(self, qwen3_moe_235b_config_dict): @@ -294,7 +303,7 @@ def test_provider_bridge_235b_config(self, qwen3_moe_235b_config_dict): result = bridge.provider_bridge(mock_pretrained) # Check 235B-specific configuration - assert result.num_layers == 94 + assert result.num_layers == 188 assert result.hidden_size == 4096 assert result.num_attention_heads == 64 assert result.ffn_hidden_size == 12288 @@ -339,7 +348,7 @@ def test_mapping_registry_parameter_mappings(self): # Should have layer norm mappings assert "model.norm.weight" in hf_params - assert "decoder.final_layernorm.weight" in megatron_params + assert "decoder.final_norm.weight" in megatron_params def test_mapping_registry_qkv_mapping(self): """Test that mapping_registry contains QKV mapping.""" @@ -374,7 +383,7 @@ def test_mapping_registry_moe_mappings(self): # Check for MoE router mapping hf_params = [mapping.hf_param for mapping in auto_mappings] - assert "model.layers.*.mlp.gate.weight" in hf_params + assert "model.layers.0.mlp.gate.weight" in hf_params # Check for expert mappings in GatedMLPMapping assert len(gated_mlp_mappings) > 0 diff --git a/tests/unit_tests/models/qwen/test_qwen_hybrid.py b/tests/unit_tests/models/qwen/test_qwen_hybrid.py new file mode 100644 index 0000000000..2365be855a --- /dev/null +++ b/tests/unit_tests/models/qwen/test_qwen_hybrid.py @@ -0,0 +1,95 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols + +from megatron.bridge.models.qwen.qwen_hybrid import ( + QwenHybridModelProvider, + configure_qwen_hybrid_layers, + qwen_pipeline_layer_pattern, +) + + +def _provider() -> QwenHybridModelProvider: + return QwenHybridModelProvider( + num_layers=2, + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=4, + bf16=False, + ) + + +def test_mtp_pattern_is_deferred_when_mtp_is_disabled(): + provider = _provider() + configure_qwen_hybrid_layers( + provider, + num_logical_layers=2, + mlp_symbols=Symbols.MLP, + mtp_mlp_symbol=Symbols.MLP, + ) + + provider.finalize() + + assert provider.hybrid_layer_pattern == "*-*-" + assert provider.mtp_hybrid_override_pattern == "*-" + + +def test_mtp_pattern_honors_recipe_override_set_after_conversion(): + provider = _provider() + configure_qwen_hybrid_layers( + provider, + num_logical_layers=2, + mlp_symbols=Symbols.MLP, + mtp_mlp_symbol=Symbols.MLP, + ) + provider.mtp_num_layers = 1 + + provider.finalize() + + assert provider.hybrid_layer_pattern == "*-*-/*-" + assert provider.num_layers == 4 + + +def test_pipeline_segmentation_preserves_logical_blocks_and_embedding_loss_balance(): + provider = QwenHybridModelProvider( + num_layers=94, + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=4, + pipeline_model_parallel_size=16, + account_for_embedding_in_pipeline_split=True, + account_for_loss_in_pipeline_split=True, + bf16=False, + ) + configure_qwen_hybrid_layers( + provider, + num_logical_layers=94, + mlp_symbols=Symbols.MOE, + ) + + provider.finalize() + + segments = provider.hybrid_layer_pattern.split(Symbols.PIPE) + assert [len(segment) for segment in segments] == [10] + [12] * 14 + [10] + assert all(segment == "*E" * (len(segment) // 2) for segment in segments) + assert provider.num_layers == 188 + assert provider.account_for_embedding_in_pipeline_split is False + assert provider.account_for_loss_in_pipeline_split is False + + +def test_pipeline_segmentation_rejects_stale_explicit_segments(): + with pytest.raises(ValueError, match="defines 2 pipeline segments"): + qwen_pipeline_layer_pattern("*-*-|*-*-", pipeline_model_parallel_size=1) From a581f7c0b2fae4cce5d4e21fb91a883d6a92d39f Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 13 Jul 2026 05:41:24 +0000 Subject: [PATCH 05/10] fix(qwen): restore HybridModel CI compatibility Signed-off-by: Philip Petrakian --- .../bridge/models/hybrid/hybrid_provider.py | 10 ++++++++-- .../ckpts/qwen3_4b/test_qwen3_4b_ckpt.py | 12 ++++++++---- .../training/test_seqpacking_cp_example.py | 6 ++++-- .../models/hybrid/test_hybrid_provider.py | 17 +++++++++++++++++ 4 files changed, 37 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/models/hybrid/hybrid_provider.py b/src/megatron/bridge/models/hybrid/hybrid_provider.py index df1cb92324..95e298cd6c 100644 --- a/src/megatron/bridge/models/hybrid/hybrid_provider.py +++ b/src/megatron/bridge/models/hybrid/hybrid_provider.py @@ -60,8 +60,9 @@ def _configure_mamba_chunk_size(stack_spec: ModuleSpec, chunk_size: int) -> Modu def modelopt_hybrid_stack_spec(config: "HybridModelProvider | None" = None) -> ModuleSpec: """Hybrid stack specification for quantization with ModelOpt. - Uses Norm instead of TENorm and ColumnParallelLinear/RowParallelLinear - instead of TE layers to enable proper quantizer insertion by ModelOpt. + Grouped-GEMM MoE checkpoints retain the default Transformer Engine stack so + their shared expert quantizers have the same state-dict paths when restored. + Other Hybrid models use ModelOpt's local stack specification. Args: config: Optional Hybrid configuration object. @@ -69,6 +70,11 @@ def modelopt_hybrid_stack_spec(config: "HybridModelProvider | None" = None) -> M Returns: Module specification for quantization-ready Hybrid stack. """ + if config is not None and config.num_moe_experts is not None and config.moe_grouped_gemm: + return get_hybrid_stack_modelopt_spec( + use_default_te_spec=True, + moe_grouped_gemm=True, + ) return get_hybrid_stack_modelopt_spec( local_core_attention=False, remap_te_layernorm=True, diff --git a/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py b/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py index 2e21c8507f..85885b2118 100644 --- a/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py +++ b/tests/functional_tests/test_groups/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py @@ -31,6 +31,7 @@ MBRIDGE_CKPT = f"{BASE_DIR}/mbridge" MCORE_CKPT = f"{BASE_DIR}/mcore" TB_DIR = f"{BASE_DIR}/tb" +QWEN3_4B_HYBRID_PATTERN = "*-" * 36 class TestQwen3Ckpt: @@ -60,7 +61,7 @@ def test_qwen3_4b_ckpt_mbridge(self): @pytest.mark.run_only_on("GPU") def test_qwen3_4b_ckpt_mcore(self, monkeypatch): - """Functional test for Qwen MCore checkpoint.""" + """Functional test for Qwen MCore Hybrid checkpoint.""" load_dir = MBRIDGE_CKPT if os.path.exists(MBRIDGE_CKPT) else None train_iters = 10 if load_dir else 5 @@ -76,7 +77,7 @@ def test_qwen3_4b_ckpt_mcore(self, monkeypatch): [ "torchrun", "--nproc_per_node=2", - "/opt/Megatron-Bridge/3rdparty/Megatron-LM/pretrain_gpt.py", + "/opt/Megatron-Bridge/3rdparty/Megatron-LM/pretrain_hybrid.py", "--init-method-std", "0.014", "--disable-bias-linear", @@ -88,8 +89,11 @@ def test_qwen3_4b_ckpt_mcore(self, monkeypatch): "--rotary-base", "1000000", "--use-rotary-position-embeddings", - "--num-layers", - "36", + "--hybrid-layer-pattern", + QWEN3_4B_HYBRID_PATTERN, + "--spec", + "megatron.core.models.hybrid.hybrid_layer_specs", + "hybrid_stack_spec", "--hidden-size", "2560", "--num-attention-heads", diff --git a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py index 4eb6f6add5..2c62c9a1cc 100644 --- a/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py +++ b/tests/functional_tests/test_groups/training/test_seqpacking_cp_example.py @@ -45,9 +45,11 @@ def _set_existing_attr(target: object, name: str, value: object) -> None: def _make_functional_test_model_small(model: object) -> None: # Keep this checkpoint-loading functional test far below runner memory limits. # The path under test is CP + sequence packing + pretrained checkpoint loading, - # not the full Llama 3.2 1B model shape. + # not the full Qwen3 model shape. Two logical Qwen blocks contain four physical + # HybridModel layers. for name, value in { - "num_layers": 2, + "num_layers": 4, + "hybrid_layer_pattern": "*-*-", "hidden_size": 128, "ffn_hidden_size": 512, "num_attention_heads": 4, diff --git a/tests/unit_tests/models/hybrid/test_hybrid_provider.py b/tests/unit_tests/models/hybrid/test_hybrid_provider.py index fb54b984d1..d762668a20 100644 --- a/tests/unit_tests/models/hybrid/test_hybrid_provider.py +++ b/tests/unit_tests/models/hybrid/test_hybrid_provider.py @@ -61,6 +61,23 @@ def test_modelopt_spec_remaps_te_layernorm_keys(self): mock_fn.assert_called_once_with(local_core_attention=False, remap_te_layernorm=True) assert result is mock_spec + def test_modelopt_spec_preserves_grouped_moe_topology(self): + provider = HybridModelProvider( + hidden_size=128, + num_attention_heads=1, + num_moe_experts=4, + moe_grouped_gemm=True, + ) + mock_spec = Mock(spec=ModuleSpec) + with patch( + "megatron.bridge.models.hybrid.hybrid_provider.get_hybrid_stack_modelopt_spec", + return_value=mock_spec, + ) as mock_fn: + result = hybrid_provider.modelopt_hybrid_stack_spec(provider) + + mock_fn.assert_called_once_with(use_default_te_spec=True, moe_grouped_gemm=True) + assert result is mock_spec + def test_rejects_mamba_stack_spec_argument(self): module_spec = ModuleSpec(module=object) From b9216b0b06192aa1bb45869c4bb695bdca807dad Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Thu, 23 Jul 2026 22:02:58 +0000 Subject: [PATCH 06/10] fix(qwen): preserve Hybrid YaRN and VPP configs Signed-off-by: Philip Petrakian --- .../bridge/models/hybrid/hybrid_provider.py | 9 ++- .../bridge/models/qwen/qwen_hybrid.py | 19 ++++--- .../models/qwen/test_qwen_hybrid.py | 56 +++++++++++++++++++ 3 files changed, 76 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/models/hybrid/hybrid_provider.py b/src/megatron/bridge/models/hybrid/hybrid_provider.py index 95e298cd6c..165956b7c6 100644 --- a/src/megatron/bridge/models/hybrid/hybrid_provider.py +++ b/src/megatron/bridge/models/hybrid/hybrid_provider.py @@ -135,10 +135,17 @@ class HybridModelProvider(TransformerConfig, ModelProviderMixin[MCoreHybridModel hybrid_layer_pattern: str | None = None seq_length: int = 8192 # HybridModel with no attention has no need for position embeddings, so none is default. - position_embedding_type: Literal["learned_absolute", "rope", "none"] = "none" + position_embedding_type: Literal["learned_absolute", "rope", "yarn", "none"] = "none" rotary_percent: float = 1.0 rotary_base: int = 10000 seq_len_interpolation_factor: float | None = None + yarn_rotary_scaling_factor: float | None = None + yarn_original_max_position_embeddings: int | None = None + yarn_beta_fast: float | None = None + yarn_beta_slow: float | None = None + yarn_mscale: float | None = None + yarn_mscale_all_dim: float | None = None + yarn_correction_range_round_to_int: bool | None = None apply_rope_fusion: bool = True make_vocab_size_divisible_by: int = 128 gated_linear_unit: bool = False diff --git a/src/megatron/bridge/models/qwen/qwen_hybrid.py b/src/megatron/bridge/models/qwen/qwen_hybrid.py index b5c30fa683..76c1407161 100644 --- a/src/megatron/bridge/models/qwen/qwen_hybrid.py +++ b/src/megatron/bridge/models/qwen/qwen_hybrid.py @@ -26,23 +26,27 @@ def qwen_pipeline_layer_pattern( hybrid_layer_pattern: str, pipeline_model_parallel_size: int, *, + virtual_pipeline_model_parallel_size: int | None = None, account_for_embedding: bool = False, account_for_loss: bool = False, ) -> str: - """Split a Qwen Hybrid pattern across PP stages without separating logical blocks.""" + """Split a Qwen Hybrid pattern across PP and VPP stages without separating logical blocks.""" if pipeline_model_parallel_size <= 0: raise ValueError("pipeline_model_parallel_size must be positive") + if virtual_pipeline_model_parallel_size is not None and virtual_pipeline_model_parallel_size <= 0: + raise ValueError("virtual_pipeline_model_parallel_size must be positive") main_pattern, separator, mtp_pattern = hybrid_layer_pattern.partition(Symbols.MTP_SEPARATOR) + pipeline_segment_count = pipeline_model_parallel_size * (virtual_pipeline_model_parallel_size or 1) if Symbols.PIPE in main_pattern: segment_count = main_pattern.count(Symbols.PIPE) + 1 - if segment_count != pipeline_model_parallel_size: + if segment_count != pipeline_segment_count: raise ValueError( f"Qwen hybrid_layer_pattern defines {segment_count} pipeline segments, " - f"but pipeline_model_parallel_size is {pipeline_model_parallel_size}." + f"but the PP/VPP configuration requires {pipeline_segment_count}." ) return hybrid_layer_pattern - if pipeline_model_parallel_size == 1: + if pipeline_segment_count == 1: return hybrid_layer_pattern if len(main_pattern) % 2: raise ValueError("Qwen Hybrid patterns must contain two physical layers per logical block.") @@ -57,9 +61,9 @@ def qwen_pipeline_layer_pattern( raise ValueError(f"Unsupported Qwen logical blocks in hybrid_layer_pattern: {invalid_blocks}") total_pipeline_units = len(logical_blocks) + int(account_for_embedding) + int(account_for_loss) - units_per_stage, extra_units = divmod(total_pipeline_units, pipeline_model_parallel_size) + units_per_stage, extra_units = divmod(total_pipeline_units, pipeline_segment_count) logical_blocks_per_stage = [ - units_per_stage + int(stage_index < extra_units) for stage_index in range(pipeline_model_parallel_size) + units_per_stage + int(stage_index < extra_units) for stage_index in range(pipeline_segment_count) ] logical_blocks_per_stage[0] -= int(account_for_embedding) logical_blocks_per_stage[-1] -= int(account_for_loss) @@ -87,7 +91,7 @@ class QwenHybridModelProvider(HybridModelProvider): def finalize(self) -> None: if ( self.hybrid_layer_pattern is not None - and self.pipeline_model_parallel_size > 1 + and self.pipeline_model_parallel_size * (self.virtual_pipeline_model_parallel_size or 1) > 1 and self.pipeline_model_parallel_layout is None and self.num_layers_in_first_pipeline_stage is None and self.num_layers_in_last_pipeline_stage is None @@ -95,6 +99,7 @@ def finalize(self) -> None: segmented_pattern = qwen_pipeline_layer_pattern( self.hybrid_layer_pattern, self.pipeline_model_parallel_size, + virtual_pipeline_model_parallel_size=self.virtual_pipeline_model_parallel_size, account_for_embedding=bool(self.account_for_embedding_in_pipeline_split), account_for_loss=bool(self.account_for_loss_in_pipeline_split), ) diff --git a/tests/unit_tests/models/qwen/test_qwen_hybrid.py b/tests/unit_tests/models/qwen/test_qwen_hybrid.py index 2365be855a..fc263bc12f 100644 --- a/tests/unit_tests/models/qwen/test_qwen_hybrid.py +++ b/tests/unit_tests/models/qwen/test_qwen_hybrid.py @@ -20,6 +20,7 @@ configure_qwen_hybrid_layers, qwen_pipeline_layer_pattern, ) +from megatron.bridge.training.utils.config_utils import _ConfigContainerBase def _provider() -> QwenHybridModelProvider: @@ -47,6 +48,38 @@ def test_mtp_pattern_is_deferred_when_mtp_is_disabled(): assert provider.mtp_hybrid_override_pattern == "*-" +def test_yarn_settings_are_serialized(): + provider = QwenHybridModelProvider( + num_layers=2, + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=4, + bf16=False, + position_embedding_type="yarn", + yarn_rotary_scaling_factor=3.125, + yarn_original_max_position_embeddings=40960, + yarn_beta_fast=32.0, + yarn_beta_slow=1.0, + yarn_mscale=1.0, + yarn_mscale_all_dim=1.0, + yarn_correction_range_round_to_int=False, + ) + expected_yarn_settings = { + "position_embedding_type": "yarn", + "yarn_rotary_scaling_factor": 3.125, + "yarn_original_max_position_embeddings": 40960, + "yarn_beta_fast": 32.0, + "yarn_beta_slow": 1.0, + "yarn_mscale": 1.0, + "yarn_mscale_all_dim": 1.0, + "yarn_correction_range_round_to_int": False, + } + + serialized_provider = _ConfigContainerBase._convert_value_to_dict(provider) + + assert {key: serialized_provider[key] for key in expected_yarn_settings} == expected_yarn_settings + + def test_mtp_pattern_honors_recipe_override_set_after_conversion(): provider = _provider() configure_qwen_hybrid_layers( @@ -90,6 +123,29 @@ def test_pipeline_segmentation_preserves_logical_blocks_and_embedding_loss_balan assert provider.account_for_loss_in_pipeline_split is False +def test_pipeline_segmentation_accounts_for_virtual_pipeline_stages(): + provider = QwenHybridModelProvider( + num_layers=16, + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=4, + pipeline_model_parallel_size=2, + virtual_pipeline_model_parallel_size=4, + bf16=False, + ) + configure_qwen_hybrid_layers( + provider, + num_logical_layers=16, + mlp_symbols=Symbols.MOE, + ) + + provider.finalize() + + segments = provider.hybrid_layer_pattern.split(Symbols.PIPE) + assert len(segments) == 8 + assert all(segment == "*E*E" for segment in segments) + + def test_pipeline_segmentation_rejects_stale_explicit_segments(): with pytest.raises(ValueError, match="defines 2 pipeline segments"): qwen_pipeline_layer_pattern("*-*-|*-*-", pipeline_model_parallel_size=1) From 18af16efdbe9235cfe0f6f32b1c53295385c6ea0 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sat, 25 Jul 2026 16:00:23 -0700 Subject: [PATCH 07/10] Updating card --- .../qwen3-8b/card.yaml | 265 ++++++++++-------- 1 file changed, 144 insertions(+), 121 deletions(-) diff --git a/examples/model_verification_cards/qwen3-8b/card.yaml b/examples/model_verification_cards/qwen3-8b/card.yaml index 52fb0c6749..4a76ca3c7d 100644 --- a/examples/model_verification_cards/qwen3-8b/card.yaml +++ b/examples/model_verification_cards/qwen3-8b/card.yaml @@ -10,8 +10,10 @@ title: qwen3_8b summary: > Performance disclaimer: this model has not been performance-tuned; reported timing and throughput metrics are sanity checks, not optimized performance - results. Qwen3-8B support verification covers conversion, inference, and - training. + results. Qwen3-8B HybridModel bidirectional CPU and TP4 conversion, manual + forward correlation, deterministic Megatron inference, bounded pretraining, + full SFT with exported-HF inference, PEFT, and direct checkpoint resume are + verified on H100; long-context SFT verification remains pending. verification_index: model_level: verified: @@ -23,7 +25,8 @@ verification_index: - inference training: H100: - verified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] + verified: [pretrain, sft, sft_export_inference, peft, checkpoint_resume] + unverified: [sft_long_context] GB200: unverified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] model: @@ -32,8 +35,8 @@ model: architecture: Qwen3ForCausalLM min_transformers_version: "5.8.1" verification_environment: - base_container: nvcr.io/nvidia/pytorch:26.04-py3 - bridge_commit: 5c56eab34c540fad08544c38b7cc39d662fb7475 # pragma: allowlist secret + base_container: nvcr.io/nvidia/pytorch:26.06-py3 + bridge_commit: b9216b0b06192aa1bb45869c4bb695bdca807dad # pragma: allowlist secret items: hf_to_megatron_cpu: @@ -42,13 +45,15 @@ items: command: > ./scripts/conversion/convert.sh import --executor slurm --device cpu --nodes 1 --hf-model Qwen/Qwen3-8B + --hf-revision b968826d9c46dd6066d109eabc6255188de91218 --megatron-path work/model-verification/qwen3-8b/cpu-megatron --torch-dtype bfloat16 - last_verified: 2026-07-16 + last_verified: 2026-07-25 expected_result: > - The command exits successfully, creates iter_0000000, and the checkpoint - round-trips through CPU export with all 399 HF tensors matching the - recorded HF revision exactly in keys, shapes, dtypes, and values. + The pinned CPU import exits successfully and creates a complete + iter_0000000 checkpoint. Its paired CPU export reloads that checkpoint + and exactly matches all 399 source BF16 tensors in keys, shapes, dtypes, + and values across 16381470720 tensor payload bytes. hf_to_megatron_gpu: status: verified @@ -57,12 +62,15 @@ items: ./scripts/conversion/convert.sh import --executor slurm --device gpu --nodes 1 --gpus-per-node 4 --hf-model Qwen/Qwen3-8B + --hf-revision b968826d9c46dd6066d109eabc6255188de91218 --megatron-path work/model-verification/qwen3-8b/imported-megatron --torch-dtype bfloat16 --tp 4 - last_verified: 2026-07-16 + last_verified: 2026-07-25 expected_result: > - The command exits successfully, creates iter_0000000, and the checkpoint - reloads at TP=4 with weights exactly matching the recorded HF revision. + The pinned four-rank GPU import exits successfully, creates a complete + four-shard iter_0000000 checkpoint, reloads at TP=4, and round-trips + through distributed GPU export with all 399 source BF16 tensors and + 16381470720 tensor payload bytes matching exactly. megatron_to_hf_cpu: status: verified @@ -70,13 +78,16 @@ items: command: > ./scripts/conversion/convert.sh export --executor slurm --device cpu --nodes 1 --hf-model Qwen/Qwen3-8B + --hf-revision b968826d9c46dd6066d109eabc6255188de91218 --megatron-path work/model-verification/qwen3-8b/cpu-megatron/iter_0000000 --hf-path work/model-verification/qwen3-8b/cpu-hf-export - last_verified: 2026-07-16 + --torch-dtype bfloat16 + last_verified: 2026-07-25 expected_result: > - The command exits successfully; all 399 exported tensors match the - recorded HF revision exactly, and the export reloads on CPU as - Qwen3ForCausalLM. + The pinned CPU export exits successfully and all 399 BF16 tensors and + 16381470720 payload bytes match the source checkpoint exactly. + Transformers reloads it strictly as Qwen3ForCausalLM with 399 tensors and + [151936, 4096] embedding and output-head shapes. megatron_to_hf_gpu: status: verified @@ -85,13 +96,16 @@ items: ./scripts/conversion/convert.sh export --executor slurm --device gpu --nodes 1 --gpus-per-node 4 --hf-model Qwen/Qwen3-8B + --hf-revision b968826d9c46dd6066d109eabc6255188de91218 --megatron-path work/model-verification/qwen3-8b/imported-megatron/iter_0000000 --hf-path work/model-verification/qwen3-8b/hf-export --torch-dtype bfloat16 --export-weight-dtype bfloat16 --distributed-save --tp 4 - last_verified: 2026-07-16 + last_verified: 2026-07-25 expected_result: > - Strict export exits successfully and the Hugging Face output reloads with - AutoModelForCausalLM as Qwen3ForCausalLM. + The pinned distributed export exits successfully and all 399 BF16 tensors + and 16381470720 payload bytes match the source checkpoint exactly. + Transformers reloads it strictly as Qwen3ForCausalLM with 399 tensors and + [151936, 4096] embedding and output-head shapes. manual_forward_pass: status: verified @@ -100,17 +114,15 @@ items: uv run python -m torch.distributed.run --standalone --nproc_per_node=4 examples/conversion/compare_hf_and_megatron/compare.py --hf_model_path Qwen/Qwen3-8B + --hf-revision b968826d9c46dd6066d109eabc6255188de91218 --megatron_model_path work/model-verification/qwen3-8b/imported-megatron/iter_0000000 --tp 4 --prompt "The capital of France is the city of" - last_verified: 2026-07-17 + last_verified: 2026-07-25 expected_result: > - The one-step comparison of the eight-token prompt exits successfully; the - Hugging Face and Megatron next-token predictions match at token ID 12095 - (" Paris"), and cosine similarity is 0.999969, above the 0.99 correlation - gate. This historical result predates explicit helper revision pinning - and is retained against the recorded immutable HF revision. The maximum - and mean absolute logit differences are 0.187500 and 0.030649, - respectively; both are report-only diagnostic observations. + The pinned one-step comparison exits successfully and the next token + matches: Hugging Face and Megatron both predict token 12095 (' Paris'). + Cosine similarity is 0.999969, maximum absolute logit difference is + 0.250000, and mean absolute logit difference is 0.031181. inference: status: verified @@ -119,20 +131,21 @@ items: uv run python -m torch.distributed.run --standalone --nproc_per_node=4 examples/conversion/hf_to_megatron_generate_text.py --hf_model_path Qwen/Qwen3-8B + --hf-revision b968826d9c46dd6066d109eabc6255188de91218 --megatron_model_path work/model-verification/qwen3-8b/imported-megatron/iter_0000000 --tp 4 --prompt "The capital of France is" --max_new_tokens 32 - last_verified: 2026-07-16 + last_verified: 2026-07-25 expected_result: > - Two independent executions exit successfully after exactly 32 generation - steps and print this byte-identical completion, including its leading space, - " Paris. The capital of Italy is Rome. The capital of Spain is Madrid. The - capital of Germany is Berlin. The capital of the Netherlands is Amsterdam. The". + One deterministic greedy generation reloads the imported checkpoint and + exits successfully after exactly 32 generated tokens. Literal completion + including its leading space: " Paris. The capital of France is Paris. The + capital of France is Paris. The capital of France is Paris. The capital + of France is Paris. The capital" pretrain: H100: status: verified precision: bf16 - bridge_commit: 619cc20bd3c7eca1dc84e8ea0f822307dfbd5cc3 # pragma: allowlist secret enabled_features: {} command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 @@ -151,26 +164,25 @@ items: rerun_state_machine.check_for_nan_in_loss=true checkpoint.load=null --save_dir work/model-verification/qwen3-8b/pretrain-reference-checkpoints --save_interval 50 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-21 + last_verified: 2026-07-25 metrics: - initial_loss: 12.73617 - final_loss: 6.190218 - last_10_steps_step_time_ms_avg: 25070.320 - last_10_steps_model_tflops_per_gpu_avg: 512.700 + initial_loss: 12.731420 + final_loss: 6.140030 + last_10_steps_step_time_ms_avg: 24361.810 + last_10_steps_model_tflops_per_gpu_avg: 527.620 expected_result: > - The uninterrupted 16-GPU bounded RP2 run completes exactly 100 steps at - TP1/PP1/CP1, DP16, GBS/MBS 1024/1, and 64-way gradient accumulation. It - reaches peak learning rate at step 40 and completes cosine decay at step - 100. LM loss is finite from 12.73617 to 6.190218 with no skipped or NaN - iterations, all four metrics are recorded, the post-setup configuration - persists, and complete 16-shard iter_0000050 and iter_0000100 checkpoints - are saved. + The uninterrupted bounded RP2 run completes exactly 100 steps on 16 + H100s at TP1/PP1/CP1, DP16, GBS/MBS 1024/1, and 64-way gradient + accumulation. It reaches peak learning rate at step 40 and completes + cosine decay at step 100. LM loss is finite from 12.731420 to 6.140030 + with no skipped or NaN iterations, all four metrics are recorded, the + post-setup configuration persists, and complete 16-shard iter_0000050 + and iter_0000100 checkpoints are saved. sft: H100: status: verified precision: bf16 - bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret enabled_features: sequence_packing: offline command: > @@ -178,67 +190,73 @@ items: --recipe qwen3_8b_sft_4gpu_h100_bf16_config --mode sft --dataset tulu3 --pretrained_checkpoint work/model-verification/qwen3-8b/imported-megatron/iter_0000000 - --max_steps 100 --seq_length 2048 + --max_steps 100 --seq_length 8192 --lr 5e-6 --min_lr 0 --warmup_iters 10 'dataset.hf_dataset.split="train[:10000]"' 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' '++tokenizer.hf_tokenizer_kwargs.revision="b968826d9c46dd6066d109eabc6255188de91218"' - dataset.hf_output_root=work/data/tulu3/qwen3-8b-sft-b14afda60f1b + dataset.hf_output_root=work/data/tulu3/qwen3-8b-sft-v2-8k-b14afda60f1b dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 dataset.do_validation=false dataset.hf_validation_proportion=null dataset.enable_offline_packing=true +dataset.offline_packing_specs.pad_seq_to_mult=1 scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null - --save_dir work/model-verification/qwen3-8b/sft-checkpoints --save_interval 100 - logger.log_interval=1 logger.log_throughput=true - last_verified: 2026-07-19 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + --save_dir work/model-verification/qwen3-8b/sft-convergence-v2-checkpoints + --save_interval 100 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: 2026-07-25 metrics: - initial_loss: 1.657044 - final_loss: 0.9733383 - last_10_steps_step_time_ms_avg: 7424.240 - last_10_steps_model_tflops_per_gpu_avg: 101.630 + initial_loss: 1.530258 + final_loss: 0.899704 + last_10_steps_step_time_ms_avg: 2026.780 + last_10_steps_model_tflops_per_gpu_avg: 373.640 expected_result: > - Immutable Tulu3 pad-1 offline packing is 99.30% efficient. Full SFT uses - DP=1 with 32 gradient-accumulation steps and reaches step 100 with the four - recorded metrics, finite loss, and no skipped or NaN iterations. Across - 6,553,600 token slots, the sampled assistant-only loss masks contain - 4,350,004 supervised tokens. The complete four-shard iter_0000100 - full-model checkpoint is saved. + The pinned Tulu3 pad-1 offline-packed run completes exactly 100 full-SFT + steps at sequence length 8192 with TP4/PP1/CP1, DP1, GBS/MBS 8/1, and + eight-way gradient accumulation. LM loss is finite from 1.530258 to + 0.899704 with no skipped or NaN iterations. Packing is 99.87% efficient + and the 6553600 token slots contain 4342763 actual supervised tokens. + All four metrics and the post-setup configuration persist, and the + complete four-shard iter_0000100 checkpoint reloads for export. sft_export_inference: H100: status: verified precision: bf16 - bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret depends_on: sft commands: - > ./scripts/conversion/convert.sh export --executor slurm --device gpu --nodes 1 --gpus-per-node 4 --hf-model Qwen/Qwen3-8B - --megatron-path work/model-verification/qwen3-8b/sft-checkpoints/iter_0000100 - --hf-path work/model-verification/qwen3-8b/sft-hf-export + --hf-revision b968826d9c46dd6066d109eabc6255188de91218 + --megatron-path work/model-verification/qwen3-8b/sft-convergence-v2-checkpoints/iter_0000100 + --hf-path work/model-verification/qwen3-8b/sft-convergence-v2-hf-export --torch-dtype bfloat16 --export-weight-dtype bfloat16 --distributed-save --tp 4 - > uv run python skills/create-model-verification-card/scripts/verify_hf_inference.py - --hf-model work/model-verification/qwen3-8b/sft-hf-export + --hf-model work/model-verification/qwen3-8b/sft-convergence-v2-hf-export --prompt "Name the capital of France and explain its role in one sentence." --max-new-tokens 45 --chat-template --disable-thinking - last_verified: 2026-07-19 + last_verified: 2026-07-25 expected_result: > - The current full-SFT checkpoint exports as five indexed BF16 shards with - 399 weights, config vocabulary size 151936, and embedding and output-head - shapes [151936, 4096]. Transformers natively reloads all 399 weights, and - two independent runs using greedy generation produce byte-identical token - IDs and exactly 45 new tokens with this literal completion, including the - escaped blank line: "The capital of France is Paris. It is the political, economic, and cultural center of France and serves as the seat of the French government. ✅\n\nParis is also a major tourist destination, known for its iconic landmarks". + The final v2 full-SFT checkpoint exports as five indexed BF16 shards + with 399 weights, vocabulary size 151936, and embedding and output-head + shapes [151936, 4096]. Transformers reloads all 399 weights without + missing, unexpected, or mismatched tensors. One deterministic greedy + run produces exactly 45 new tokens with this literal completion, where + \n\n denotes two newline bytes: "The capital of France is Paris. It is + the political, economic, and cultural center of France and a major + global city. ✅\n\nParis is the capital of France and the center of the + country's political, economic," sft_long_context: H100: - status: verified + status: unverified precision: bf16 - bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret enabled_features: sequence_packing: offline context_parallel_size: 2 @@ -261,26 +279,28 @@ items: model.cp_comm_type=a2a model.cross_entropy_loss_fusion=false scheduler.lr_decay_iters=20 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null checkpoint.save=null - logger.log_interval=1 logger.log_throughput=true - last_verified: 2026-07-20 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: null metrics: - initial_loss: 1.611622 - final_loss: 1.490618 - last_10_steps_step_time_ms_avg: 71668.780 - last_10_steps_model_tflops_per_gpu_avg: 34.080 + initial_loss: 1.611647 + final_loss: 1.490598 + last_10_steps_step_time_ms_avg: 55077.330 + last_10_steps_model_tflops_per_gpu_avg: 44.360 expected_result: > - The immutable-revision 8-GPU run completes exactly 20 Tulu3 SFT steps at - sequence length 32768 with recipe-owned TP4/PP1/CP2/SP-on, GBS/MBS 8/1, - and explicit pad-8 offline packing. LM loss is 1.611622 to 1.490618; - skipped/NaN totals are 0/0. The persisted post-setup runtime config - matches the command, packing is 99.97%, and the sampled training window - contains 3,444,917 actual supervised tokens. + The pinned Tulu3 run completed exactly 20 SFT steps at sequence length + 32768 with TP4/PP1/CP2/SP-on, DP1, GBS/MBS 8/1, and explicit pad-8 + offline packing. LM loss was finite from 1.611647 to 1.490598 with no + skipped or NaN iterations. Packing was 99.97% efficient and the 5242880 + token slots contained 3463090 actual supervised tokens. The workload + remains unverified because it did not persist the required post-setup + configuration; rerun with an explicit logger.save_config_filepath. peft: H100: status: verified precision: bf16 - bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret enabled_features: sequence_packing: offline command: > @@ -288,38 +308,44 @@ items: --recipe qwen3_8b_peft_1gpu_h100_bf16_config --mode lora --dataset tulu3 --pretrained_checkpoint work/model-verification/qwen3-8b/imported-megatron/iter_0000000 - --max_steps 100 --seq_length 2048 + --max_steps 100 --seq_length 8192 --lr 1e-4 --min_lr 0 --warmup_iters 10 'dataset.hf_dataset.split="train[:10000]"' 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' '++tokenizer.hf_tokenizer_kwargs.revision="b968826d9c46dd6066d109eabc6255188de91218"' - dataset.hf_output_root=work/data/tulu3/qwen3-8b-peft-b14afda60f1b + dataset.hf_output_root=work/data/tulu3/qwen3-8b-peft-v2-8k-b14afda60f1b dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 dataset.do_validation=false dataset.hf_validation_proportion=null dataset.enable_offline_packing=true - +dataset.offline_packing_specs.pad_seq_to_mult=4 scheduler.lr_decay_iters=100 + +dataset.offline_packing_specs.pad_seq_to_mult=1 scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null - --save_dir work/model-verification/qwen3-8b/peft-checkpoints --save_interval 100 - logger.log_interval=1 logger.log_throughput=true - last_verified: 2026-07-19 + ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true + rerun_state_machine.check_for_nan_in_loss=true + --save_dir work/model-verification/qwen3-8b/peft-convergence-v2-checkpoints + --save_interval 100 + logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + last_verified: 2026-07-25 metrics: - initial_loss: 1.527745 - final_loss: 1.170038 - last_10_steps_step_time_ms_avg: 7774.790 - last_10_steps_model_tflops_per_gpu_avg: 387.430 + initial_loss: 1.529487 + final_loss: 0.975589 + last_10_steps_step_time_ms_avg: 4326.870 + last_10_steps_model_tflops_per_gpu_avg: 697.720 expected_result: > - Immutable Tulu3 pad-4 offline packing is 99.39% efficient. The 100 LoRA - steps use DP=1 with 32 gradient-accumulation steps and complete with the - four recorded metrics, finite loss, and no skipped or NaN iterations. - Across 6,553,600 token slots, the sampled assistant-only loss masks - contain 4,332,480 supervised tokens. The complete single-shard - iter_0000100 adapter checkpoint is saved. + The pinned Tulu3 pad-1 offline-packed run completes exactly 100 LoRA + steps at sequence length 8192 with TP1/PP1/CP1, DP1, GBS/MBS 8/1, and + eight-way gradient accumulation. LM loss is finite from 1.529487 to + 0.975589 with no skipped or NaN iterations. Packing is 99.87% efficient + and the 6553600 token slots contain 4342763 actual supervised tokens. + Exactly 5308416 rank-8, alpha-16, dropout-0 linear_qkv and linear_proj + adapter parameters train while base weights remain frozen. All four + metrics and the post-setup configuration persist, and the complete + single-shard iter_0000100 adapter checkpoint contains all 144 expected + BF16 adapter tensors. checkpoint_resume: H100: status: verified precision: bf16 - bridge_commit: 619cc20bd3c7eca1dc84e8ea0f822307dfbd5cc3 # pragma: allowlist secret depends_on: pretrain command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 @@ -340,12 +366,12 @@ items: --save_dir work/model-verification/qwen3-8b/pretrain-resumed-from-reference-checkpoints --save_interval 50 checkpoint.ckpt_step=50 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-21 + last_verified: 2026-07-25 metrics: - initial_loss: 7.015641 - final_loss: 6.190785 - last_10_steps_step_time_ms_avg: 24995.610 - last_10_steps_model_tflops_per_gpu_avg: 514.240 + initial_loss: 6.889990 + final_loss: 6.137956 + last_10_steps_step_time_ms_avg: 24260.990 + last_10_steps_model_tflops_per_gpu_avg: 529.800 resume_comparison: reference_item: pretrain sentinel_steps: [51, 100] @@ -353,15 +379,12 @@ items: loss_absolute_tolerance: 1.0e-6 sentinels_match: true expected_result: > - The direct 16-GPU continuation restores optimizer, scheduler, - data-order, and RNG state from iter_0000050, begins at step 51, and - finishes exactly at step 100 in a distinct output root at TP1/PP1/CP1, - DP16, GBS/MBS 1024/1, and 64-way gradient accumulation. All 50 losses - are finite with no skipped or NaN iterations, all four metrics are - recorded, the post-setup configuration persists, and a complete - 16-shard iter_0000100 checkpoint is saved. Step-51 resumed/reference - loss is 7.015641/7.015641. Step-100 resumed/reference loss is - 6.190785/6.190218, an absolute difference of 0.000567. This difference - does not meet the card's previous 1e-6-relative threshold; this verified - result uses the model-verification skill's standard 1%-relative gate, - which was fixed for this 16-GPU execution layout before the run. + The direct 16-GPU continuation restores model, optimizer, scheduler, + data-order, and RNG state from iter_0000050, executes only steps 51 + through 100, and saves a complete 16-shard iter_0000100 checkpoint to a + distinct output root. All 50 losses are finite with no skipped or NaN + iterations, all four metrics and the post-setup configuration persist, + and the declared loss sentinels pass: resumed/reference loss is + 6.889990/6.889990 at step 51 and 6.137956/6.140030 at step 100, whose + absolute difference of 0.002074 is within the standard + absolute-plus-1%-relative tolerance. From de7b4ce37a260df76ef3487b093ab07514d351e7 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sat, 25 Jul 2026 18:11:56 -0700 Subject: [PATCH 08/10] docs(model): verify Qwen3-8B long-context SFT Signed-off-by: Philip Petrakian --- .../qwen3-8b/card.yaml | 32 ++++++++++--------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/examples/model_verification_cards/qwen3-8b/card.yaml b/examples/model_verification_cards/qwen3-8b/card.yaml index 4a76ca3c7d..08da919e7e 100644 --- a/examples/model_verification_cards/qwen3-8b/card.yaml +++ b/examples/model_verification_cards/qwen3-8b/card.yaml @@ -13,7 +13,8 @@ summary: > results. Qwen3-8B HybridModel bidirectional CPU and TP4 conversion, manual forward correlation, deterministic Megatron inference, bounded pretraining, full SFT with exported-HF inference, PEFT, and direct checkpoint resume are - verified on H100; long-context SFT verification remains pending. + verified on H100, including long-context SFT with sequence packing and + context parallelism. verification_index: model_level: verified: @@ -25,8 +26,7 @@ verification_index: - inference training: H100: - verified: [pretrain, sft, sft_export_inference, peft, checkpoint_resume] - unverified: [sft_long_context] + verified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] GB200: unverified: [pretrain, sft, sft_export_inference, sft_long_context, peft, checkpoint_resume] model: @@ -255,8 +255,9 @@ items: sft_long_context: H100: - status: unverified + status: verified precision: bf16 + bridge_commit: 6d417c08fde67c88827bb95aa7048c65f481d5d1 # pragma: allowlist secret enabled_features: sequence_packing: offline context_parallel_size: 2 @@ -281,21 +282,22 @@ items: checkpoint.load=null checkpoint.save=null ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true - logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: null + logger.log_interval=1 logger.log_throughput=true + logger.save_config_filepath=work/model-verification/qwen3-8b/sft-long-context-runtime-config.yaml + logger.tensorboard_dir=null + last_verified: 2026-07-25 metrics: initial_loss: 1.611647 - final_loss: 1.490598 - last_10_steps_step_time_ms_avg: 55077.330 - last_10_steps_model_tflops_per_gpu_avg: 44.360 + final_loss: 1.490631 + last_10_steps_step_time_ms_avg: 55104.950 + last_10_steps_model_tflops_per_gpu_avg: 44.300 expected_result: > - The pinned Tulu3 run completed exactly 20 SFT steps at sequence length + The pinned Tulu3 run completes exactly 20 SFT steps at sequence length 32768 with TP4/PP1/CP2/SP-on, DP1, GBS/MBS 8/1, and explicit pad-8 - offline packing. LM loss was finite from 1.611647 to 1.490598 with no - skipped or NaN iterations. Packing was 99.97% efficient and the 5242880 - token slots contained 3463090 actual supervised tokens. The workload - remains unverified because it did not persist the required post-setup - configuration; rerun with an explicit logger.save_config_filepath. + offline packing. LM loss is finite from 1.611647 to 1.490631 with no + skipped or NaN iterations. Packing is 99.97% efficient and the 5242880 + token slots contain 3463090 actual supervised tokens. All four metrics + and the post-setup ConfigContainer persist. peft: H100: From 83caabba49ac32f76d8e365a36926366c6940001 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sat, 25 Jul 2026 20:36:49 -0700 Subject: [PATCH 09/10] Updating card Signed-off-by: Philip Petrakian --- .../qwen3-8b/card.yaml | 90 ++++++++++--------- 1 file changed, 46 insertions(+), 44 deletions(-) diff --git a/examples/model_verification_cards/qwen3-8b/card.yaml b/examples/model_verification_cards/qwen3-8b/card.yaml index 08da919e7e..b6ea03919a 100644 --- a/examples/model_verification_cards/qwen3-8b/card.yaml +++ b/examples/model_verification_cards/qwen3-8b/card.yaml @@ -183,6 +183,7 @@ items: H100: status: verified precision: bf16 + bridge_commit: de7b4ce37a260df76ef3487b093ab07514d351e7 # pragma: allowlist secret enabled_features: sequence_packing: offline command: > @@ -190,68 +191,69 @@ items: --recipe qwen3_8b_sft_4gpu_h100_bf16_config --mode sft --dataset tulu3 --pretrained_checkpoint work/model-verification/qwen3-8b/imported-megatron/iter_0000000 - --max_steps 100 --seq_length 8192 + --max_steps 100 --seq_length 2048 --lr 5e-6 --min_lr 0 --warmup_iters 10 'dataset.hf_dataset.split="train[:10000]"' 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' '++tokenizer.hf_tokenizer_kwargs.revision="b968826d9c46dd6066d109eabc6255188de91218"' - dataset.hf_output_root=work/data/tulu3/qwen3-8b-sft-v2-8k-b14afda60f1b + dataset.hf_output_root=work/data/tulu3/qwen3-8b-sft-regression-v1-2k-b14afda60f1b dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 dataset.do_validation=false dataset.hf_validation_proportion=null dataset.enable_offline_packing=true +dataset.offline_packing_specs.pad_seq_to_mult=1 scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null - ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true - rerun_state_machine.check_for_nan_in_loss=true - --save_dir work/model-verification/qwen3-8b/sft-convergence-v2-checkpoints + --save_dir work/model-verification/qwen3-8b/sft-regression-v1-checkpoints --save_interval 100 - logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + logger.log_interval=1 logger.log_throughput=true + logger.save_config_filepath=work/model-verification/qwen3-8b/sft-regression-v1-runtime-config.yaml last_verified: 2026-07-25 metrics: - initial_loss: 1.530258 - final_loss: 0.899704 - last_10_steps_step_time_ms_avg: 2026.780 - last_10_steps_model_tflops_per_gpu_avg: 373.640 + initial_loss: 1.657584 + final_loss: 0.9735701 + last_10_steps_step_time_ms_avg: 4360.140 + last_10_steps_model_tflops_per_gpu_avg: 173.630 expected_result: > The pinned Tulu3 pad-1 offline-packed run completes exactly 100 full-SFT - steps at sequence length 8192 with TP4/PP1/CP1, DP1, GBS/MBS 8/1, and - eight-way gradient accumulation. LM loss is finite from 1.530258 to - 0.899704 with no skipped or NaN iterations. Packing is 99.87% efficient - and the 6553600 token slots contain 4342763 actual supervised tokens. - All four metrics and the post-setup configuration persist, and the - complete four-shard iter_0000100 checkpoint reloads for export. + steps at sequence length 2048 with TP4/PP1/CP1, DP1, GBS/MBS 32/1, and + 32-way gradient accumulation. LM loss is finite from 1.657584 to + 0.9735701 with no skipped or NaN iterations. Packing is 99.30% efficient + and the 6553600 token slots contain 4350004 supervised tokens under the + historical packed-mask counting convention. All four metrics and the + post-setup configuration persist, and the complete four-shard + iter_0000100 checkpoint reloads for export. sft_export_inference: H100: status: verified precision: bf16 + bridge_commit: de7b4ce37a260df76ef3487b093ab07514d351e7 # pragma: allowlist secret depends_on: sft commands: - > ./scripts/conversion/convert.sh export --executor slurm --device gpu --nodes 1 --gpus-per-node 4 --hf-model Qwen/Qwen3-8B --hf-revision b968826d9c46dd6066d109eabc6255188de91218 - --megatron-path work/model-verification/qwen3-8b/sft-convergence-v2-checkpoints/iter_0000100 - --hf-path work/model-verification/qwen3-8b/sft-convergence-v2-hf-export + --megatron-path work/model-verification/qwen3-8b/sft-regression-v1-checkpoints/iter_0000100 + --hf-path work/model-verification/qwen3-8b/sft-regression-v1-hf-export --torch-dtype bfloat16 --export-weight-dtype bfloat16 --distributed-save --tp 4 - > uv run python skills/create-model-verification-card/scripts/verify_hf_inference.py - --hf-model work/model-verification/qwen3-8b/sft-convergence-v2-hf-export + --hf-model work/model-verification/qwen3-8b/sft-regression-v1-hf-export --prompt "Name the capital of France and explain its role in one sentence." --max-new-tokens 45 --chat-template --disable-thinking last_verified: 2026-07-25 expected_result: > - The final v2 full-SFT checkpoint exports as five indexed BF16 shards - with 399 weights, vocabulary size 151936, and embedding and output-head - shapes [151936, 4096]. Transformers reloads all 399 weights without - missing, unexpected, or mismatched tensors. One deterministic greedy - run produces exactly 45 new tokens with this literal completion, where - \n\n denotes two newline bytes: "The capital of France is Paris. It is - the political, economic, and cultural center of France and a major - global city. ✅\n\nParis is the capital of France and the center of the - country's political, economic," + The final controlled full-SFT checkpoint reloads and exports as five + indexed BF16 shards with 399 weights, vocabulary size 151936, and + embedding and output-head shapes [151936, 4096]. Transformers reloads + all 399 weights without missing, unexpected, or mismatched tensors. One + deterministic greedy run produces exactly 45 new tokens with this + literal completion, where \n\n denotes two newline bytes: "The capital + of France is Paris. It is the political, economic, and cultural center + of France and serves as the seat of the French government. ✅\n\nI hope + this helps! Let me know if you have any other" sft_long_context: H100: @@ -303,6 +305,7 @@ items: H100: status: verified precision: bf16 + bridge_commit: de7b4ce37a260df76ef3487b093ab07514d351e7 # pragma: allowlist secret enabled_features: sequence_packing: offline command: > @@ -310,34 +313,33 @@ items: --recipe qwen3_8b_peft_1gpu_h100_bf16_config --mode lora --dataset tulu3 --pretrained_checkpoint work/model-verification/qwen3-8b/imported-megatron/iter_0000000 - --max_steps 100 --seq_length 8192 + --max_steps 100 --seq_length 2048 --lr 1e-4 --min_lr 0 --warmup_iters 10 'dataset.hf_dataset.split="train[:10000]"' 'dataset.hf_dataset.load_kwargs={revision:"b14afda60f1bbebe55d5d2fa1e4df5042f97f8be"}' '++tokenizer.hf_tokenizer_kwargs.revision="b968826d9c46dd6066d109eabc6255188de91218"' - dataset.hf_output_root=work/data/tulu3/qwen3-8b-peft-v2-8k-b14afda60f1b + dataset.hf_output_root=work/data/tulu3/qwen3-8b-peft-regression-v1-2k-b14afda60f1b dataset.hf_rewrite=true dataset.seed=1234 rng.seed=5678 dataset.do_validation=false dataset.hf_validation_proportion=null dataset.enable_offline_packing=true - +dataset.offline_packing_specs.pad_seq_to_mult=1 scheduler.lr_decay_iters=100 + +dataset.offline_packing_specs.pad_seq_to_mult=4 scheduler.lr_decay_iters=100 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null - ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true - rerun_state_machine.check_for_nan_in_loss=true - --save_dir work/model-verification/qwen3-8b/peft-convergence-v2-checkpoints + --save_dir work/model-verification/qwen3-8b/peft-regression-v1-checkpoints --save_interval 100 - logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null + logger.log_interval=1 logger.log_throughput=true + logger.save_config_filepath=work/model-verification/qwen3-8b/peft-regression-v1-runtime-config.yaml last_verified: 2026-07-25 metrics: - initial_loss: 1.529487 - final_loss: 0.975589 - last_10_steps_step_time_ms_avg: 4326.870 - last_10_steps_model_tflops_per_gpu_avg: 697.720 + initial_loss: 1.527877 + final_loss: 1.170195 + last_10_steps_step_time_ms_avg: 5367.620 + last_10_steps_model_tflops_per_gpu_avg: 561.330 expected_result: > - The pinned Tulu3 pad-1 offline-packed run completes exactly 100 LoRA - steps at sequence length 8192 with TP1/PP1/CP1, DP1, GBS/MBS 8/1, and - eight-way gradient accumulation. LM loss is finite from 1.529487 to - 0.975589 with no skipped or NaN iterations. Packing is 99.87% efficient - and the 6553600 token slots contain 4342763 actual supervised tokens. + The pinned Tulu3 pad-4 offline-packed run completes exactly 100 LoRA + steps at sequence length 2048 with TP1/PP1/CP1, DP1, GBS/MBS 32/1, and + 32-way gradient accumulation. LM loss is finite from 1.527877 to + 1.170195 with no skipped or NaN iterations. Packing is 99.39% efficient + and the 6553600 token slots contain 4332480 supervised tokens. Exactly 5308416 rank-8, alpha-16, dropout-0 linear_qkv and linear_proj adapter parameters train while base weights remain frozen. All four metrics and the post-setup configuration persist, and the complete From 5ed17285c90aa525b540bb2ad48b029c6a824032 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Sun, 26 Jul 2026 14:55:45 -0700 Subject: [PATCH 10/10] updating 30b card Signed-off-by: Philip Petrakian --- .../qwen3-30b-a3b/card.yaml | 73 ++++++++++--------- 1 file changed, 37 insertions(+), 36 deletions(-) diff --git a/examples/model_verification_cards/qwen3-30b-a3b/card.yaml b/examples/model_verification_cards/qwen3-30b-a3b/card.yaml index 612207bee3..98299c8444 100644 --- a/examples/model_verification_cards/qwen3-30b-a3b/card.yaml +++ b/examples/model_verification_cards/qwen3-30b-a3b/card.yaml @@ -98,7 +98,7 @@ items: manual_forward_pass: status: verified precision: bf16 - bridge_commit: 24f3dcde1fe27a1c3f1097b94d69b38e37a57415 # pragma: allowlist secret + bridge_commit: 83caabba49ac32f76d8e365a36926366c6940001 # pragma: allowlist secret command: > uv run python -m torch.distributed.run --standalone --nproc_per_node=8 examples/conversion/compare_hf_and_megatron/compare.py @@ -106,17 +106,18 @@ items: --hf-revision ad44e777bcd18fa416d9da3bd8f70d33ebb85d39 --megatron_model_path work/model-verification/qwen3-30b-a3b/imported-megatron/iter_0000000 --tp 1 --pp 8 --prompt "The capital of France is the city of" - last_verified: 2026-07-20 + last_verified: 2026-07-26 expected_result: > The pinned-revision one-step comparison exits successfully; the Hugging Face and Megatron next-token predictions match at token ID 12095 - (" Paris"), and cosine similarity is 0.999317, above the 0.99 correlation - gate. The maximum and mean absolute logit differences are 0.593750 and - 0.130064, respectively; both are report-only diagnostic observations. + (" Paris"), and cosine similarity is 0.999725, above the 0.99 correlation + gate. The maximum and mean absolute logit differences are 0.390625 and + 0.067908, respectively; both are report-only diagnostic observations. inference: status: verified precision: bf16 + bridge_commit: 83caabba49ac32f76d8e365a36926366c6940001 # pragma: allowlist secret command: > uv run python -m torch.distributed.run --standalone --nproc_per_node=8 examples/conversion/hf_to_megatron_generate_text.py @@ -124,7 +125,7 @@ items: --megatron_model_path work/model-verification/qwen3-30b-a3b/imported-megatron/iter_0000000 --tp 4 --pp 2 --ep 4 --prompt "The capital of France is" --max_new_tokens 32 - last_verified: 2026-07-17 + last_verified: 2026-07-26 expected_result: > Two independent executions exit successfully after exactly 32 new tokens and print this byte-identical completion, including its leading @@ -136,7 +137,7 @@ items: H100: status: verified precision: bf16 - bridge_commit: 5b9d9cf501193277e1ca47a99b97c640b8f39f90 # pragma: allowlist secret + bridge_commit: 83caabba49ac32f76d8e365a36926366c6940001 # pragma: allowlist secret enabled_features: cuda_graph: implementation: transformer_engine @@ -161,18 +162,18 @@ items: --save_dir work/model-verification/qwen3-30b-a3b/pretrain-convergence-v1-reference-checkpoints --save_interval 50 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-19 + last_verified: 2026-07-26 metrics: - initial_loss: 12.41145 - final_loss: 6.139116 - last_10_steps_step_time_ms_avg: 30289.550 - last_10_steps_model_tflops_per_gpu_avg: 199.120 + initial_loss: 12.40963 + final_loss: 6.063886 + last_10_steps_step_time_ms_avg: 31601.840 + last_10_steps_model_tflops_per_gpu_avg: 190.840 expected_result: > On 16x H100, the public alias resolves to the 16-GPU recipe and completes exactly 100 bounded RP2 optimizer steps with TP1/PP1/CP1/EP16/ETP1, DP16, SP off, GBS/MBS 1024/1, and 64-way gradient accumulation. Natural routing, HybridEP, and Transformer Engine CUDA graphs for moe_router and - moe_preprocess remain active. Loss is finite from 12.41145 to 6.139116 + moe_preprocess remain active. Loss is finite from 12.40963 to 6.063886 with no skipped or NaN iterations, all four metrics are recorded, and complete iter_0000050 and iter_0000100 checkpoints are saved. @@ -259,7 +260,7 @@ items: H100: status: verified precision: bf16 - bridge_commit: f3ae2767b5e18aeb67b726cd8d5f1db58216dcc9 # pragma: allowlist secret + bridge_commit: 83caabba49ac32f76d8e365a36926366c6940001 # pragma: allowlist secret enabled_features: sequence_packing: offline context_parallel_size: 2 @@ -285,26 +286,26 @@ items: scheduler.lr_decay_iters=20 validation.eval_iters=0 validation.eval_interval=0 checkpoint.load=null checkpoint.save=null logger.log_interval=1 logger.log_throughput=true - last_verified: 2026-07-20 + last_verified: 2026-07-26 metrics: - initial_loss: 1.645009 - final_loss: 1.468103 - last_10_steps_step_time_ms_avg: 142663.710 - last_10_steps_model_tflops_per_gpu_avg: 26.120 + initial_loss: 1.645200 + final_loss: 1.468164 + last_10_steps_step_time_ms_avg: 37263.530 + last_10_steps_model_tflops_per_gpu_avg: 100.040 expected_result: > The immutable-revision 16-GPU run completes exactly 20 Tulu3 SFT steps at sequence length 32768 with TP8/PP1/CP2/EP8/SP-on, DeepEP, and explicit - pad-16 offline packing. LM loss is 1.645009 to 1.468103; skipped/NaN + pad-16 offline packing. LM loss is 1.645200 to 1.468164; skipped/NaN totals are 0/0. The persisted post-setup runtime config matches the command, packing is 99.28%, and the sampled training window contains - 13,573,663 actual supervised tokens. PP=1 keeps tokens, labels, loss + 13,644,614 actual supervised tokens. PP=1 keeps tokens, labels, loss masks, and packed-sequence boundaries on one pipeline stage. peft: H100: status: verified precision: bf16 - bridge_commit: 619cc20bd3c7eca1dc84e8ea0f822307dfbd5cc3 # pragma: allowlist secret + bridge_commit: 83caabba49ac32f76d8e365a36926366c6940001 # pragma: allowlist secret enabled_features: sequence_packing: offline moe_dispatcher: deepep @@ -330,12 +331,12 @@ items: --save_dir work/model-verification/qwen3-30b-a3b/peft-tp4-ep4-pad4-checkpoints --save_interval 100 logger.log_interval=1 logger.log_throughput=true - last_verified: 2026-07-21 + last_verified: 2026-07-26 metrics: - initial_loss: 1.575987 - final_loss: 1.113119 - last_10_steps_step_time_ms_avg: 22347.640 - last_10_steps_model_tflops_per_gpu_avg: 13.840 + initial_loss: 1.575093 + final_loss: 1.114007 + last_10_steps_step_time_ms_avg: 16181.250 + last_10_steps_model_tflops_per_gpu_avg: 19.120 expected_result: > The immutable-revision 4-GPU run completes exactly 100 PEFT steps at TP4/PP1/CP1/EP4/ETP1, DP1, SP on, GBS/MBS 32/1, and 32-way gradient @@ -343,7 +344,7 @@ items: 99.39% efficient, and the sampled 6,553,600 token slots contain 4,332,480 supervised tokens after label masking. Only rank-8, alpha-16, zero-dropout LoRA on linear_qkv and linear_proj is trainable. LM loss is - finite from 1.575987 to 1.113119 with no skipped or NaN iterations, all + finite from 1.575093 to 1.114007 with no skipped or NaN iterations, all four metrics are recorded, and the complete four-shard iter_0000100 adapter checkpoint covers all 192 expected adapter entries. @@ -351,7 +352,7 @@ items: H100: status: verified precision: bf16 - bridge_commit: 5b9d9cf501193277e1ca47a99b97c640b8f39f90 # pragma: allowlist secret + bridge_commit: 83caabba49ac32f76d8e365a36926366c6940001 # pragma: allowlist secret depends_on: pretrain command: > ./scripts/training/train.sh --nodes 2 --gpus-per-node 8 @@ -375,12 +376,12 @@ items: train.empty_unused_memory_level=2 logger.log_interval=1 logger.log_throughput=true logger.tensorboard_dir=null - last_verified: 2026-07-19 + last_verified: 2026-07-26 metrics: - initial_loss: 6.989006 - final_loss: 6.145390 - last_10_steps_step_time_ms_avg: 31142.650 - last_10_steps_model_tflops_per_gpu_avg: 193.640 + initial_loss: 6.704744 + final_loss: 6.060962 + last_10_steps_step_time_ms_avg: 34594.660 + last_10_steps_model_tflops_per_gpu_avg: 174.340 resume_comparison: reference_item: pretrain sentinel_steps: [51, 100] @@ -392,8 +393,8 @@ items: iter_0000050, begins at step 51, and finishes at step 100 in the distinct resumed root with finite losses and no skipped or NaN iterations. Releasing unused cached memory after optimizer steps is execution-only. Step-51 loss - 6.989006 matches the uninterrupted reference exactly; step-100 loss - 6.145390 differs from reference 6.139116 by 0.102197%, within the declared + 6.704744 matches the uninterrupted reference exactly; step-100 loss + 6.060962 differs from reference 6.063886 by 0.048220%, within the declared one-percent gate, so both sentinels match and all four metrics are recorded. pretrain_performance: