diff --git a/src/megatron/bridge/models/qwen3_asr/modeling_qwen3_asr/thinker_model.py b/src/megatron/bridge/models/qwen3_asr/modeling_qwen3_asr/thinker_model.py index 98bf167b62..7da95b4261 100644 --- a/src/megatron/bridge/models/qwen3_asr/modeling_qwen3_asr/thinker_model.py +++ b/src/megatron/bridge/models/qwen3_asr/modeling_qwen3_asr/thinker_model.py @@ -27,8 +27,7 @@ ) from megatron.bridge.models.qwen3_asr.modeling_qwen3_asr.rope import get_rope_index from megatron.bridge.models.qwen3_asr.modeling_qwen3_asr.transformer_config import Qwen3ASRTransformerConfig -from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention -from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import Qwen3VLGPTModel +from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import Qwen3VLHybridModel from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.utils import ( split_data_cp_rank, ) @@ -56,8 +55,6 @@ def __init__( ) -> None: super().__init__(config=language_transformer_config) - language_transformer_layer_spec.submodules.self_attention.module = Qwen3VLSelfAttention - self.pre_process = pre_process self.post_process = post_process self.add_encoder = add_encoder @@ -95,13 +92,13 @@ def __init__( self.audio_model = Qwen3ASRAudioEncoderHF._from_config(thinker_transformer_config.audio_config) hook_hf_module_setattr_for_tp_grad_sync(self.audio_model) - self.language_model = Qwen3VLGPTModel( + self.language_model = Qwen3VLHybridModel( config=language_transformer_config, - transformer_layer_spec=language_transformer_layer_spec, + hybrid_stack_spec=language_transformer_layer_spec, vocab_size=language_transformer_config.vocab_size, max_sequence_length=language_transformer_config.language_max_sequence_length, + hybrid_layer_pattern=language_transformer_config.hybrid_layer_pattern, parallel_output=parallel_output, - position_embedding_type="mrope", rotary_percent=language_transformer_config.rotary_percent, pre_process=self.pre_process, post_process=self.post_process, diff --git a/src/megatron/bridge/models/qwen3_asr/modeling_qwen3_asr/transformer_config.py b/src/megatron/bridge/models/qwen3_asr/modeling_qwen3_asr/transformer_config.py index 98c132ce3f..c1d6ddac25 100644 --- a/src/megatron/bridge/models/qwen3_asr/modeling_qwen3_asr/transformer_config.py +++ b/src/megatron/bridge/models/qwen3_asr/modeling_qwen3_asr/transformer_config.py @@ -29,6 +29,7 @@ class Qwen3ASRTransformerConfig(TransformerConfig): share_embeddings_and_output_weights: bool = False rotary_percent: float = 1.0 rotary_base: float = 5000000.0 + hybrid_layer_pattern: str | None = None # Multimodal rope section for 3 dimensions (same position IDs across all dims for ASR) mrope_section: list[int] = field(default_factory=lambda: [24, 20, 20]) diff --git a/src/megatron/bridge/models/qwen3_asr/qwen3_asr_bridge.py b/src/megatron/bridge/models/qwen3_asr/qwen3_asr_bridge.py index 5786297acc..a311be59fb 100644 --- a/src/megatron/bridge/models/qwen3_asr/qwen3_asr_bridge.py +++ b/src/megatron/bridge/models/qwen3_asr/qwen3_asr_bridge.py @@ -13,6 +13,7 @@ # limitations under the License. import torch +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge @@ -23,6 +24,11 @@ ReplicatedMapping, ) from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.qwen.qwen_hybrid import ( + configure_qwen_hybrid_layers, + qwen_logical_layer_count, + qwen_physical_layer_indices, +) from megatron.bridge.models.qwen3_asr.modeling_qwen3_asr.model import Qwen3ASRModel from megatron.bridge.models.qwen3_asr.qwen3_asr_provider import Qwen3ASRModelProvider @@ -30,7 +36,12 @@ # Use string-based registration because Qwen3ASRForConditionalGeneration is not in # the standard transformers library (it's a custom model in qwen_asr package). # auto_bridge.py resolves custom architectures via config.auto_map or string fallback. -@MegatronModelBridge.register_bridge(source="Qwen3ASRForConditionalGeneration", target=Qwen3ASRModel) +@MegatronModelBridge.register_bridge( + source="Qwen3ASRForConditionalGeneration", + target=Qwen3ASRModel, + provider=Qwen3ASRModelProvider, + model_type="qwen3_asr", +) class Qwen3ASRBridge(MegatronModelBridge): """ Megatron Bridge for Qwen3-ASR Conditional Generation. @@ -79,31 +90,57 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> Qwen3ASRModelPro audio_start_token_id=getattr(thinker_config, "audio_start_token_id", 151647), mrope_section=(getattr(text_config, "rope_scaling", None) or {}).get("mrope_section", [24, 20, 20]), ) + configure_qwen_hybrid_layers( + provider, + num_logical_layers=text_config.num_hidden_layers, + mlp_symbols=Symbols.MLP, + mtp_mlp_symbol=Symbols.MLP, + ) return provider def mapping_registry(self) -> MegatronMappingRegistry: """Return MegatronMappingRegistry containing parameter mappings for Qwen3-ASR models.""" - # LLM parameter mappings (Qwen3-style with QK layernorm, prefixed with thinker.) param_mappings = { - # Embeddings and output layers "thinker.language_model.embedding.word_embeddings.weight": "thinker.model.embed_tokens.weight", "thinker.language_model.output_layer.weight": "thinker.lm_head.weight", - "thinker.language_model.decoder.final_layernorm.weight": "thinker.model.norm.weight", - # Layer normalization - "thinker.language_model.decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "thinker.model.layers.*.input_layernorm.weight", - "thinker.language_model.decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "thinker.model.layers.*.post_attention_layernorm.weight", - # QK layernorm (Qwen3-specific) - "thinker.language_model.decoder.layers.*.self_attention.q_layernorm.weight": "thinker.model.layers.*.self_attn.q_norm.weight", - "thinker.language_model.decoder.layers.*.self_attention.k_layernorm.weight": "thinker.model.layers.*.self_attn.k_norm.weight", - # Attention output projection - "thinker.language_model.decoder.layers.*.self_attention.linear_proj.weight": "thinker.model.layers.*.self_attn.o_proj.weight", - # MLP down projection - "thinker.language_model.decoder.layers.*.mlp.linear_fc2.weight": "thinker.model.layers.*.mlp.down_proj.weight", + "thinker.language_model.decoder.final_norm.weight": "thinker.model.norm.weight", } - mapping_list = [] - for megatron_param, hf_param in param_mappings.items(): - mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) + mapping_list = [AutoMapping(k, v) for k, v in param_mappings.items()] + + num_layers = self.hf_config.thinker_config.text_config.num_hidden_layers + for logical_layer_idx in range(num_layers): + attention_layer_idx, mlp_layer_idx = qwen_physical_layer_indices(logical_layer_idx) + hf_layer = f"thinker.model.layers.{logical_layer_idx}" + attention_layer = f"thinker.language_model.decoder.layers.{attention_layer_idx}.self_attention" + mlp_layer = f"thinker.language_model.decoder.layers.{mlp_layer_idx}.mlp" + mapping_list.extend( + [ + AutoMapping( + f"{attention_layer}.linear_qkv.layer_norm_weight", + f"{hf_layer}.input_layernorm.weight", + ), + AutoMapping(f"{attention_layer}.q_layernorm.weight", f"{hf_layer}.self_attn.q_norm.weight"), + AutoMapping(f"{attention_layer}.k_layernorm.weight", f"{hf_layer}.self_attn.k_norm.weight"), + AutoMapping(f"{attention_layer}.linear_proj.weight", f"{hf_layer}.self_attn.o_proj.weight"), + AutoMapping( + f"{mlp_layer}.linear_fc1.layer_norm_weight", + f"{hf_layer}.post_attention_layernorm.weight", + ), + AutoMapping(f"{mlp_layer}.linear_fc2.weight", 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", + ), + ] + ) mapping_list.extend( [ @@ -113,20 +150,16 @@ def mapping_registry(self) -> MegatronMappingRegistry: megatron_param="thinker.audio_model.**", hf_param="thinker.audio_tower.**", ), - # QKV weight: Combine separate Q, K, V weights into single QKV matrix (no bias for Qwen3) - QKVMapping( - megatron_param="thinker.language_model.decoder.layers.*.self_attention.linear_qkv.weight", - q="thinker.model.layers.*.self_attn.q_proj.weight", - k="thinker.model.layers.*.self_attn.k_proj.weight", - v="thinker.model.layers.*.self_attn.v_proj.weight", - ), - # Gated MLP: Combine gate and up projection matrices into single FC1 matrix - GatedMLPMapping( - megatron_param="thinker.language_model.decoder.layers.*.mlp.linear_fc1.weight", - gate="thinker.model.layers.*.mlp.gate_proj.weight", - up="thinker.model.layers.*.mlp.up_proj.weight", - ), ] ) return MegatronMappingRegistry(*mapping_list) + + @classmethod + def megatron_to_hf_config(cls, provider) -> dict: + """Restore the logical Qwen layer count when exporting HybridModel config.""" + 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 diff --git a/src/megatron/bridge/models/qwen3_asr/qwen3_asr_provider.py b/src/megatron/bridge/models/qwen3_asr/qwen3_asr_provider.py index 78ac343f5f..431708af5e 100644 --- a/src/megatron/bridge/models/qwen3_asr/qwen3_asr_provider.py +++ b/src/megatron/bridge/models/qwen3_asr/qwen3_asr_provider.py @@ -23,21 +23,26 @@ from typing import Callable import torch.nn.functional as F -from megatron.core.models.gpt import GPTModel as MCoreGPTModel -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.transformer.spec_utils import ModuleSpec -from megatron.bridge.models.gpt_provider import GPTModelProvider +from megatron.bridge.models.qwen.qwen_hybrid import QwenHybridModelProvider, configure_qwen_hybrid_layers from megatron.bridge.models.qwen3_asr.hf_qwen3_asr.configuration_qwen3_asr import ( Qwen3ASRThinkerConfig, ) from megatron.bridge.models.qwen3_asr.modeling_qwen3_asr.model import Qwen3ASRModel +from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import ( + Qwen3VLHybridModel, + get_qwen3_vl_hybrid_stack_spec, +) +from megatron.bridge.models.qwen_vl.qwen3_vl_provider import _provide_qwen3_vl_language_model @dataclass -class Qwen3ASRModelProvider(GPTModelProvider): +class Qwen3ASRModelProvider(QwenHybridModelProvider): """ Base model provider for Qwen3-ASR Models. - Inherits language model configuration from GPTModelProvider with Qwen3-specific defaults. + Inherits language model configuration from HybridModelProvider with Qwen3-specific defaults. Key characteristics: - Audio-only (no vision, no video) @@ -85,19 +90,36 @@ class Qwen3ASRModelProvider(GPTModelProvider): distribute_saved_activations: bool = False cp_comm_type: str = "p2p" gradient_accumulation_fusion: bool = False + mtp_num_layers: int | None = None + hybrid_stack_spec: ModuleSpec | Callable = get_qwen3_vl_hybrid_stack_spec + + def finalize(self) -> None: + if self.hybrid_layer_pattern is None: + if self.num_layers is None: + raise ValueError("num_layers must be configured for Qwen3-ASR") + configure_qwen_hybrid_layers( + self, + num_logical_layers=self.num_layers, + mlp_symbols=Symbols.MLP, + mtp_mlp_symbol=Symbols.MLP, + ) + super().finalize() def provide(self, pre_process=None, post_process=None, vp_stage=None): """Provide a Qwen3-ASR model instance with audio and language components.""" + if self.hybrid_layer_pattern is None: + if self.num_layers is None: + raise ValueError("num_layers must be configured for Qwen3-ASR") + configure_qwen_hybrid_layers( + self, + num_logical_layers=self.num_layers, + mlp_symbols=Symbols.MLP, + mtp_mlp_symbol=Symbols.MLP, + ) language_transformer_config = self thinker_config = self.thinker_config - # Qwen3 GPT layer spec with QK layernorm - language_transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( - num_experts=None, - moe_grouped_gemm=False, - qk_layernorm=self.qk_layernorm, - fp8=False, - ) + language_transformer_layer_spec = self._resolve_hybrid_stack_spec() model = Qwen3ASRModel( language_transformer_config=language_transformer_config, @@ -116,6 +138,6 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None): return model - def provide_language_model(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreGPTModel: + def provide_language_model(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VLHybridModel: """Provide just the language model component without audio.""" - return GPTModelProvider.provide(self, pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) + return _provide_qwen3_vl_language_model(self, pre_process, post_process, vp_stage) diff --git a/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/thinker_model.py b/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/thinker_model.py index 606e5a2453..50af126518 100644 --- a/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/thinker_model.py +++ b/src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/thinker_model.py @@ -32,7 +32,7 @@ from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.transformer_config import ( Qwen3OmniTransformerConfig, ) -from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import Qwen3VLGPTModel +from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import Qwen3VLHybridModel from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.utils import ( split_data_cp_rank, split_deepstack_embs, @@ -215,13 +215,13 @@ def __init__( _enable_multimodal_gradient_checkpointing(self.visual) _enable_multimodal_gradient_checkpointing(self.audio_model) - self.language_model = Qwen3VLGPTModel( + self.language_model = Qwen3VLHybridModel( config=language_transformer_config, - transformer_layer_spec=language_transformer_layer_spec, + hybrid_stack_spec=language_transformer_layer_spec, vocab_size=language_transformer_config.vocab_size, max_sequence_length=language_transformer_config.language_max_sequence_length, + hybrid_layer_pattern=language_transformer_config.hybrid_layer_pattern, parallel_output=parallel_output, - position_embedding_type="mrope", rotary_percent=language_transformer_config.rotary_percent, pre_process=self.pre_process, post_process=self.post_process, diff --git a/src/megatron/bridge/models/qwen_omni/qwen3_omni_bridge.py b/src/megatron/bridge/models/qwen_omni/qwen3_omni_bridge.py index fc60135351..eb7ed64c47 100644 --- a/src/megatron/bridge/models/qwen_omni/qwen3_omni_bridge.py +++ b/src/megatron/bridge/models/qwen_omni/qwen3_omni_bridge.py @@ -15,6 +15,7 @@ import logging import torch +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols from transformers import Qwen3OmniMoeForConditionalGeneration from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry @@ -26,6 +27,11 @@ ReplicatedMapping, ) from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.qwen.qwen_hybrid import ( + configure_qwen_hybrid_layers, + qwen_logical_layer_count, + qwen_physical_layer_indices, +) from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.model import Qwen3OmniModel from megatron.bridge.models.qwen_omni.qwen3_omni_provider import Qwen3OmniModelProvider @@ -42,6 +48,15 @@ class Qwen3OmniBridge(MegatronModelBridge): """Bridge for Qwen3-Omni.""" + @classmethod + def megatron_to_hf_config(cls, provider) -> dict: + """Restore the logical Qwen layer count when exporting HybridModel config.""" + 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: PreTrainedCausalLM) -> Qwen3OmniModelProvider: hf_config = hf_pretrained.config if getattr(hf_config, "enable_audio_output", False): @@ -105,19 +120,19 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> Qwen3OmniModelPr position_embedding_type="mrope", mrope_section=rope_scaling.get("mrope_section", [24, 20, 20]), ) + configure_qwen_hybrid_layers( + provider, + num_logical_layers=text_config.num_hidden_layers, + mlp_symbols=Symbols.MOE, + mtp_mlp_symbol=Symbols.MOE, + ) return provider def mapping_registry(self) -> MegatronMappingRegistry: param_mappings = { "thinker.language_model.embedding.word_embeddings.weight": "thinker.model.embed_tokens.weight", "thinker.language_model.output_layer.weight": "thinker.lm_head.weight", - "thinker.language_model.decoder.final_layernorm.weight": "thinker.model.norm.weight", - "thinker.language_model.decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "thinker.model.layers.*.input_layernorm.weight", - "thinker.language_model.decoder.layers.*.pre_mlp_layernorm.weight": "thinker.model.layers.*.post_attention_layernorm.weight", - "thinker.language_model.decoder.layers.*.self_attention.q_layernorm.weight": "thinker.model.layers.*.self_attn.q_norm.weight", - "thinker.language_model.decoder.layers.*.self_attention.k_layernorm.weight": "thinker.model.layers.*.self_attn.k_norm.weight", - "thinker.language_model.decoder.layers.*.self_attention.linear_proj.weight": "thinker.model.layers.*.self_attn.o_proj.weight", - "thinker.language_model.decoder.layers.*.mlp.router.weight": "thinker.model.layers.*.mlp.gate.weight", + "thinker.language_model.decoder.final_norm.weight": "thinker.model.norm.weight", } mapping_list = [ @@ -125,6 +140,53 @@ def mapping_registry(self) -> MegatronMappingRegistry: for megatron_param, hf_param in param_mappings.items() ] + num_layers = self.hf_config.thinker_config.text_config.num_hidden_layers + for logical_layer_idx in range(num_layers): + attention_layer_idx, moe_layer_idx = qwen_physical_layer_indices(logical_layer_idx) + hf_layer = f"thinker.model.layers.{logical_layer_idx}" + attention_layer = f"thinker.language_model.decoder.layers.{attention_layer_idx}.self_attention" + moe_layer = f"thinker.language_model.decoder.layers.{moe_layer_idx}" + mapping_list.extend( + [ + AutoMapping( + f"{attention_layer}.linear_qkv.layer_norm_weight", + f"{hf_layer}.input_layernorm.weight", + ), + AutoMapping(f"{attention_layer}.q_layernorm.weight", f"{hf_layer}.self_attn.q_norm.weight"), + AutoMapping(f"{attention_layer}.k_layernorm.weight", f"{hf_layer}.self_attn.k_norm.weight"), + AutoMapping(f"{attention_layer}.linear_proj.weight", f"{hf_layer}.self_attn.o_proj.weight"), + AutoMapping( + f"{moe_layer}.pre_mlp_layernorm.weight", + f"{hf_layer}.post_attention_layernorm.weight", + ), + AutoMapping(f"{moe_layer}.mlp.router.weight", 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( + f"{moe_layer}.mlp.experts.linear_fc2.weight*", + 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( + f"{moe_layer}.mlp.experts.local_experts.*.linear_fc2.weight", + f"{hf_layer}.mlp.experts.*.down_proj.weight", + ), + ] + ) + mapping_list.extend( [ ReplicatedMapping( @@ -135,30 +197,6 @@ def mapping_registry(self) -> MegatronMappingRegistry: megatron_param="thinker.audio_model.**", hf_param="thinker.audio_tower.**", ), - QKVMapping( - megatron_param="thinker.language_model.decoder.layers.*.self_attention.linear_qkv.weight", - q="thinker.model.layers.*.self_attn.q_proj.weight", - k="thinker.model.layers.*.self_attn.k_proj.weight", - v="thinker.model.layers.*.self_attn.v_proj.weight", - ), - GatedMLPMapping( - megatron_param="thinker.language_model.decoder.layers.*.mlp.experts.linear_fc1.weight*", - gate="thinker.model.layers.*.mlp.experts.*.gate_proj.weight", - up="thinker.model.layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param="thinker.language_model.decoder.layers.*.mlp.experts.linear_fc2.weight*", - hf_param="thinker.model.layers.*.mlp.experts.*.down_proj.weight", - ), - GatedMLPMapping( - megatron_param="thinker.language_model.decoder.layers.*.mlp.experts.local_experts.*.linear_fc1.weight", - gate="thinker.model.layers.*.mlp.experts.*.gate_proj.weight", - up="thinker.model.layers.*.mlp.experts.*.up_proj.weight", - ), - AutoMapping( - megatron_param="thinker.language_model.decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight", - hf_param="thinker.model.layers.*.mlp.experts.*.down_proj.weight", - ), ] ) diff --git a/src/megatron/bridge/models/qwen_omni/qwen3_omni_provider.py b/src/megatron/bridge/models/qwen_omni/qwen3_omni_provider.py index 0dae0d1762..78d69241e4 100644 --- a/src/megatron/bridge/models/qwen_omni/qwen3_omni_provider.py +++ b/src/megatron/bridge/models/qwen_omni/qwen3_omni_provider.py @@ -16,20 +16,8 @@ from typing import Callable import torch.nn.functional as F -from megatron.core.extensions.transformer_engine import HAVE_TE -from megatron.core.models.gpt import GPTModel as MCoreGPTModel -from megatron.core.models.gpt.gpt_layer_specs import ( - get_gpt_layer_local_spec, - get_gpt_layer_with_transformer_engine_spec, -) -from megatron.core.pipeline_parallel.utils import ( - is_pp_first_stage, - is_pp_last_stage, - is_vp_first_stage, - is_vp_last_stage, -) -from megatron.core.transformer.attention import SelfAttention -from megatron.core.transformer.enums import AttnBackend +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage from megatron.core.transformer.spec_utils import ModuleSpec from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeCode2WavConfig, @@ -37,44 +25,17 @@ Qwen3OmniMoeThinkerConfig, ) -from megatron.bridge.models.gpt_provider import GPTModelProvider +from megatron.bridge.models.qwen.qwen_hybrid import QwenHybridModelProvider, configure_qwen_hybrid_layers from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.model import Qwen3OmniModel -from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention - - -def _use_qwen3_vl_self_attention(layer_spec) -> None: - """Install the mRoPE-aware Qwen3-VL attention module on standard attention specs.""" - if layer_spec is None: - return - - if hasattr(layer_spec, "layer_specs"): - for sub_spec in layer_spec.layer_specs: - _use_qwen3_vl_self_attention(sub_spec) - return - - if not isinstance(layer_spec, ModuleSpec): - return - - submodules = getattr(layer_spec, "submodules", None) - if submodules is None: - return - - if hasattr(submodules, "mtp_model_layer"): - _use_qwen3_vl_self_attention(submodules.mtp_model_layer) - - attention_spec = getattr(submodules, "self_attention", None) - if attention_spec is None: - return - - attention_module = getattr(attention_spec, "module", None) - if attention_module is SelfAttention or ( - isinstance(attention_module, type) and issubclass(attention_module, SelfAttention) - ): - attention_spec.module = Qwen3VLSelfAttention +from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import ( + Qwen3VLHybridModel, + get_qwen3_vl_hybrid_stack_spec, +) +from megatron.bridge.models.qwen_vl.qwen3_vl_provider import _provide_qwen3_vl_language_model @dataclass -class Qwen3OmniModelProvider(GPTModelProvider): +class Qwen3OmniModelProvider(QwenHybridModelProvider): """Provider for Qwen3-Omni. The current implementation focuses on thinker-side multimodal training and @@ -128,39 +89,40 @@ class Qwen3OmniModelProvider(GPTModelProvider): freeze_audio_model: bool = False vit_gradient_checkpointing: bool = False multimodal_attn_impl: str = "auto" + mtp_num_layers: int | None = None + hybrid_stack_spec: ModuleSpec | Callable = get_qwen3_vl_hybrid_stack_spec + + def finalize(self) -> None: + if self.hybrid_layer_pattern is None: + if self.num_layers is None: + raise ValueError("num_layers must be configured for Qwen3-Omni") + configure_qwen_hybrid_layers( + self, + num_logical_layers=self.num_layers, + mlp_symbols=Symbols.MOE, + mtp_mlp_symbol=Symbols.MOE, + ) + super().finalize() def provide(self, pre_process=None, post_process=None, vp_stage=None): + if vp_stage is not None or self.virtual_pipeline_model_parallel_size is not None: + raise ValueError("Virtual pipeline parallelism is not supported by HybridModel.") + if self.hybrid_layer_pattern is None: + if self.num_layers is None: + raise ValueError("num_layers must be configured for Qwen3-Omni") + configure_qwen_hybrid_layers( + self, + num_logical_layers=self.num_layers, + mlp_symbols=Symbols.MOE, + mtp_mlp_symbol=Symbols.MOE, + ) pp_group = self._pg_collection.pp if self._pg_collection is not None else None - vp_size = self.virtual_pipeline_model_parallel_size if pre_process is None: - pre_process = ( - is_vp_first_stage(vp_stage=vp_stage, vp_size=vp_size) and is_pp_first_stage(pp_group) - if pp_group is not None - else True - ) + pre_process = is_pp_first_stage(pp_group) if pp_group is not None else True if post_process is None: - post_process = ( - is_vp_last_stage(vp_stage=vp_stage, vp_size=vp_size) and is_pp_last_stage(pp_group) - if pp_group is not None - else True - ) + post_process = is_pp_last_stage(pp_group) if pp_group is not None else True - use_local_attention = self.attention_backend in {AttnBackend.local, "local"} - if HAVE_TE and not use_local_attention: - language_transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( - num_experts=self.num_moe_experts, - moe_grouped_gemm=self.moe_grouped_gemm, - qk_layernorm=self.qk_layernorm, - fp8=False, - ) - else: - language_transformer_layer_spec = get_gpt_layer_local_spec( - num_experts=self.num_moe_experts, - moe_grouped_gemm=self.moe_grouped_gemm, - qk_layernorm=self.qk_layernorm, - normalization=self.normalization, - ) - _use_qwen3_vl_self_attention(language_transformer_layer_spec) + language_transformer_layer_spec = self._resolve_hybrid_stack_spec() model = Qwen3OmniModel( language_transformer_config=self, @@ -182,5 +144,5 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None): return model - def provide_language_model(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreGPTModel: - return super().provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) + def provide_language_model(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VLHybridModel: + return _provide_qwen3_vl_language_model(self, pre_process, post_process, vp_stage) diff --git a/src/megatron/bridge/models/qwen_omni/qwen3_omni_step.py b/src/megatron/bridge/models/qwen_omni/qwen3_omni_step.py index a5c0fe0fae..850b3740b2 100644 --- a/src/megatron/bridge/models/qwen_omni/qwen3_omni_step.py +++ b/src/megatron/bridge/models/qwen_omni/qwen3_omni_step.py @@ -21,8 +21,8 @@ from typing import TYPE_CHECKING, Any, Iterable import torch -from megatron.core.models.gpt import GPTModel from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage +from megatron.core.transformer import MegatronModule from megatron.core.utils import get_batch_on_this_cp_rank, get_model_config from megatron.bridge.training.losses import ( @@ -214,7 +214,7 @@ def _get_dense_batch_on_this_cp_rank(batch: dict[str, Any], cp_group) -> dict[st def forward_step( state: "GlobalState", data_iterator: Iterable, - model: GPTModel, + model: MegatronModule, return_schedule_plan: bool = False, ) -> tuple[torch.Tensor, partial]: """Forward training step for Qwen3-Omni thinker.""" diff --git a/tests/unit_tests/models/qwen3_asr/modeling_qwen3_asr/test_qwen3_asr_model.py b/tests/unit_tests/models/qwen3_asr/modeling_qwen3_asr/test_qwen3_asr_model.py index c0c370231f..46656aaf7b 100644 --- a/tests/unit_tests/models/qwen3_asr/modeling_qwen3_asr/test_qwen3_asr_model.py +++ b/tests/unit_tests/models/qwen3_asr/modeling_qwen3_asr/test_qwen3_asr_model.py @@ -28,7 +28,6 @@ import torch.distributed as dist import torch.nn.functional as F from megatron.core import parallel_state -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed @@ -38,6 +37,7 @@ ) from megatron.bridge.models.qwen3_asr.modeling_qwen3_asr.model import Qwen3ASRModel from megatron.bridge.models.qwen3_asr.modeling_qwen3_asr.transformer_config import Qwen3ASRTransformerConfig +from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import get_qwen3_vl_hybrid_stack_spec HIDDEN_SIZE = 128 @@ -147,7 +147,7 @@ def _make_language_config(thinker_config): mrope_section = (rope_scaling or {}).get("mrope_section", [4, 6, 6]) return Qwen3ASRTransformerConfig( - num_layers=2, + num_layers=4, hidden_size=text_cfg.hidden_size, num_attention_heads=text_cfg.num_attention_heads, num_query_groups=text_cfg.num_key_value_heads, @@ -171,26 +171,23 @@ def _make_language_config(thinker_config): attention_dropout=text_cfg.attention_dropout, audio_token_id=thinker_config.audio_token_id, audio_start_token_id=thinker_config.audio_start_token_id, + hybrid_layer_pattern="*-*-", ) @staticmethod - def _make_layer_spec(): - """Create a GPT layer spec for the language model (Qwen3 uses QK layernorm).""" - return get_gpt_layer_with_transformer_engine_spec( - num_experts=None, - moe_grouped_gemm=False, - qk_layernorm=True, - fp8=False, - ) + def _make_layer_spec(language_config): + """Create a Qwen multimodal Hybrid stack spec.""" + return get_qwen3_vl_hybrid_stack_spec(language_config) def _build_model(self, thinker_config, pre_process=True, post_process=True, add_encoder=True, add_decoder=True): """Helper to build a Qwen3ASRModel with the given flags.""" self._setup_parallel_state(tp_size=1, pp_size=1) pg_collection = ProcessGroupCollection.use_mpu_process_groups() + language_config = self._make_language_config(thinker_config) return Qwen3ASRModel( - language_transformer_config=self._make_language_config(thinker_config), - language_transformer_layer_spec=self._make_layer_spec(), + language_transformer_config=language_config, + language_transformer_layer_spec=self._make_layer_spec(language_config), thinker_transformer_config=thinker_config, parallel_output=True, pre_process=pre_process, diff --git a/tests/unit_tests/models/qwen3_asr/test_qwen3_asr_config.py b/tests/unit_tests/models/qwen3_asr/test_qwen3_asr_config.py index baa6385590..075c3def09 100644 --- a/tests/unit_tests/models/qwen3_asr/test_qwen3_asr_config.py +++ b/tests/unit_tests/models/qwen3_asr/test_qwen3_asr_config.py @@ -15,16 +15,28 @@ import pytest from megatron.bridge.models.conversion.utils import conform_config_to_reference +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.models.qwen3_asr.hf_qwen3_asr.configuration_qwen3_asr import ( Qwen3ASRConfig, Qwen3ASRThinkerConfig, ) +from megatron.bridge.models.qwen3_asr.qwen3_asr_provider import Qwen3ASRModelProvider from megatron.bridge.training.config import ConfigContainer pytestmark = [pytest.mark.unit] +def test_qwen3_asr_uses_hybrid_provider(): + assert issubclass(Qwen3ASRModelProvider, HybridModelProvider) + + provider = Qwen3ASRModelProvider(num_layers=2, hidden_size=128, num_attention_heads=4) + provider.finalize() + + assert provider.num_layers == 4 + assert provider.hybrid_layer_pattern == "*-*-" + + def test_qwen3_asr_config_default_constructs_thinker_config(): config = Qwen3ASRConfig() diff --git a/tests/unit_tests/models/qwen_omni/modeling_qwen3_omni/test_omni_model.py b/tests/unit_tests/models/qwen_omni/modeling_qwen3_omni/test_omni_model.py index 0720187155..fc9646b048 100644 --- a/tests/unit_tests/models/qwen_omni/modeling_qwen3_omni/test_omni_model.py +++ b/tests/unit_tests/models/qwen_omni/modeling_qwen3_omni/test_omni_model.py @@ -23,10 +23,6 @@ import torch.distributed as dist import torch.nn.functional as F from megatron.core import parallel_state -from megatron.core.models.gpt.gpt_layer_specs import ( - get_gpt_layer_local_spec, - get_gpt_layer_with_transformer_engine_spec, -) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( @@ -39,6 +35,7 @@ from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.transformer_config import ( Qwen3OmniTransformerConfig, ) +from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import get_qwen3_vl_hybrid_stack_spec HIDDEN_SIZE = 128 @@ -165,7 +162,7 @@ def teardown_method(self): @staticmethod def _make_language_config(): return Qwen3OmniTransformerConfig( - num_layers=2, + num_layers=4, hidden_size=HIDDEN_SIZE, num_attention_heads=4, num_query_groups=2, @@ -195,30 +192,20 @@ def _make_language_config(): audio_start_token_id=AUDIO_START_TOKEN_ID, position_id_per_seconds=25, seconds_per_chunk=2, + hybrid_layer_pattern="*E*E", ) @staticmethod - def _make_layer_spec(): - if not torch.cuda.is_available(): - return get_gpt_layer_local_spec( - num_experts=8, - moe_grouped_gemm=True, - qk_layernorm=True, - normalization="RMSNorm", - ) - return get_gpt_layer_with_transformer_engine_spec( - num_experts=8, - moe_grouped_gemm=True, - qk_layernorm=True, - fp8=False, - ) + def _make_layer_spec(language_config): + return get_qwen3_vl_hybrid_stack_spec(language_config) def _build_model(self, thinker_config): self._setup_parallel_state(tp_size=1, pp_size=1) pg_collection = ProcessGroupCollection.use_mpu_process_groups() + language_config = self._make_language_config() return Qwen3OmniModel( - language_transformer_config=self._make_language_config(), - language_transformer_layer_spec=self._make_layer_spec(), + language_transformer_config=language_config, + language_transformer_layer_spec=self._make_layer_spec(language_config), thinker_transformer_config=thinker_config, parallel_output=True, pre_process=True, @@ -261,7 +248,7 @@ def enable_input_require_grads(self): model = Qwen3OmniModel( language_transformer_config=cfg, - language_transformer_layer_spec=self._make_layer_spec(), + language_transformer_layer_spec=self._make_layer_spec(cfg), thinker_transformer_config=thinker_config, parallel_output=True, pre_process=True, diff --git a/tests/unit_tests/models/qwen_omni/test_qwen3_omni_bridge.py b/tests/unit_tests/models/qwen_omni/test_qwen3_omni_bridge.py index 2914cb78f6..2ca007a74c 100644 --- a/tests/unit_tests/models/qwen_omni/test_qwen3_omni_bridge.py +++ b/tests/unit_tests/models/qwen_omni/test_qwen3_omni_bridge.py @@ -91,7 +91,8 @@ def test_provider_bridge_basic_config(self, mock_hf_pretrained): provider = bridge.provider_bridge(mock_hf_pretrained) assert isinstance(provider, Qwen3OmniModelProvider) - assert provider.num_layers == 48 + assert provider.num_layers == 96 + assert provider.hybrid_layer_pattern == "*E" * 48 assert provider.hidden_size == 2048 assert provider.ffn_hidden_size == 6144 assert provider.moe_ffn_hidden_size == 768 @@ -128,8 +129,9 @@ def test_provider_bridge_dtype(self, mock_dtype_from_hf, mock_hf_pretrained): assert provider.fp16 is False assert provider.params_dtype == torch.bfloat16 - def test_mapping_registry(self): + def test_mapping_registry(self, mock_hf_config): bridge = Qwen3OmniBridge() + bridge.hf_config = mock_hf_config registry = bridge.mapping_registry() assert isinstance(registry, MegatronMappingRegistry) @@ -140,10 +142,10 @@ def test_mapping_registry(self): assert any("thinker.language_model.embedding.word_embeddings.weight" in name for name in mapping_names) assert any( - "thinker.language_model.decoder.layers.*.self_attention.linear_qkv.weight" in name + "thinker.language_model.decoder.layers.0.self_attention.linear_qkv.weight" in name for name in mapping_names ) - assert any("thinker.language_model.decoder.layers.*.mlp.router.weight" in name for name in mapping_names) + assert any("thinker.language_model.decoder.layers.1.mlp.router.weight" in name for name in mapping_names) def test_provider_bridge_warns_for_audio_output_stack(self, mock_hf_pretrained, caplog): mock_hf_pretrained.config.enable_audio_output = True diff --git a/tests/unit_tests/models/qwen_omni/test_qwen3_omni_provider.py b/tests/unit_tests/models/qwen_omni/test_qwen3_omni_provider.py index fbb743e28d..4fac3412ec 100644 --- a/tests/unit_tests/models/qwen_omni/test_qwen3_omni_provider.py +++ b/tests/unit_tests/models/qwen_omni/test_qwen3_omni_provider.py @@ -15,13 +15,11 @@ from unittest.mock import patch import pytest -from megatron.core.transformer.dot_product_attention import DotProductAttention -from megatron.core.transformer.enums import AttnBackend from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeThinkerConfig, ) -import megatron.bridge.models.qwen_omni.qwen3_omni_provider as qwen3_omni_provider +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.models.qwen_omni import Qwen3OmniModelProvider from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention @@ -30,6 +28,9 @@ class TestQwen3OmniModelProvider: + def test_inherits_from_hybrid_provider(self): + assert issubclass(Qwen3OmniModelProvider, HybridModelProvider) + def test_qwen3_omni_model_provider_initialization(self): provider = Qwen3OmniModelProvider( num_layers=48, @@ -108,99 +109,7 @@ def test_provide_defaults_to_single_stage_pre_post_process(self): _, kwargs = mock_model_cls.call_args assert kwargs["pre_process"] is True assert kwargs["post_process"] is True - - def test_provide_uses_qwen3_vl_attention_with_local_backend(self): - provider = Qwen3OmniModelProvider( - num_layers=2, - hidden_size=128, - ffn_hidden_size=256, - num_attention_heads=4, - num_query_groups=2, - kv_channels=32, - vocab_size=1024, - use_cpu_initialization=True, - bf16=False, - attention_backend=AttnBackend.local, - ) - - with patch("megatron.bridge.models.qwen_omni.qwen3_omni_provider.Qwen3OmniModel") as mock_model_cls: - provider.provide() - - _, kwargs = mock_model_cls.call_args - language_spec = kwargs["language_transformer_layer_spec"] - self_attention_spec = language_spec.submodules.self_attention - assert self_attention_spec.module is Qwen3VLSelfAttention - assert self_attention_spec.submodules.core_attention is DotProductAttention - - def test_string_local_backend_uses_local_spec_when_te_is_available(self): - provider = Qwen3OmniModelProvider( - num_layers=2, - hidden_size=128, - ffn_hidden_size=256, - num_attention_heads=4, - num_query_groups=2, - kv_channels=32, - vocab_size=1024, - use_cpu_initialization=True, - bf16=False, - attention_backend="local", + assert ( + kwargs["language_transformer_layer_spec"].submodules.attention_layer.submodules.self_attention.module + is Qwen3VLSelfAttention ) - - with ( - patch("megatron.bridge.models.qwen_omni.qwen3_omni_provider.HAVE_TE", True), - patch("megatron.bridge.models.qwen_omni.qwen3_omni_provider.Qwen3OmniModel") as mock_model_cls, - ): - provider.provide() - - _, kwargs = mock_model_cls.call_args - self_attention_spec = kwargs["language_transformer_layer_spec"].submodules.self_attention - assert self_attention_spec.module is Qwen3VLSelfAttention - assert self_attention_spec.submodules.core_attention is DotProductAttention - - def test_default_backend_keeps_te_core_attention_when_available(self): - if not qwen3_omni_provider.HAVE_TE: - pytest.skip("Transformer Engine is not available") - - provider = Qwen3OmniModelProvider( - num_layers=2, - hidden_size=128, - ffn_hidden_size=256, - num_attention_heads=4, - num_query_groups=2, - kv_channels=32, - vocab_size=1024, - use_cpu_initialization=True, - bf16=False, - ) - - with patch("megatron.bridge.models.qwen_omni.qwen3_omni_provider.Qwen3OmniModel") as mock_model_cls: - provider.provide() - - _, kwargs = mock_model_cls.call_args - self_attention_spec = kwargs["language_transformer_layer_spec"].submodules.self_attention - assert self_attention_spec.module is Qwen3VLSelfAttention - assert self_attention_spec.submodules.core_attention.__name__ == "TEDotProductAttention" - - def test_default_backend_falls_back_to_local_spec_without_te(self): - provider = Qwen3OmniModelProvider( - num_layers=2, - hidden_size=128, - ffn_hidden_size=256, - num_attention_heads=4, - num_query_groups=2, - kv_channels=32, - vocab_size=1024, - use_cpu_initialization=True, - bf16=False, - ) - - with ( - patch("megatron.bridge.models.qwen_omni.qwen3_omni_provider.HAVE_TE", False), - patch("megatron.bridge.models.qwen_omni.qwen3_omni_provider.Qwen3OmniModel") as mock_model_cls, - ): - provider.provide() - - _, kwargs = mock_model_cls.call_args - self_attention_spec = kwargs["language_transformer_layer_spec"].submodules.self_attention - assert self_attention_spec.module is Qwen3VLSelfAttention - assert self_attention_spec.submodules.core_attention is DotProductAttention diff --git a/tests/unit_tests/models/qwen_omni/test_qwen3_omni_training_smoke.py b/tests/unit_tests/models/qwen_omni/test_qwen3_omni_training_smoke.py index 48b7118ab9..39b3ec3279 100644 --- a/tests/unit_tests/models/qwen_omni/test_qwen3_omni_training_smoke.py +++ b/tests/unit_tests/models/qwen_omni/test_qwen3_omni_training_smoke.py @@ -21,10 +21,6 @@ import torch.distributed as dist import torch.nn.functional as F from megatron.core import parallel_state -from megatron.core.models.gpt.gpt_layer_specs import ( - get_gpt_layer_local_spec, - get_gpt_layer_with_transformer_engine_spec, -) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( @@ -35,6 +31,7 @@ from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.transformer_config import ( Qwen3OmniTransformerConfig, ) +from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import get_qwen3_vl_hybrid_stack_spec HIDDEN_SIZE = 128 @@ -94,7 +91,7 @@ def _make_toy_thinker_config(): def _make_language_config(): return Qwen3OmniTransformerConfig( - num_layers=2, + num_layers=4, hidden_size=HIDDEN_SIZE, num_attention_heads=4, num_query_groups=2, @@ -124,23 +121,12 @@ def _make_language_config(): audio_start_token_id=AUDIO_START_TOKEN_ID, position_id_per_seconds=25, seconds_per_chunk=2, + hybrid_layer_pattern="*E*E", ) -def _make_layer_spec(): - if not torch.cuda.is_available(): - return get_gpt_layer_local_spec( - num_experts=8, - moe_grouped_gemm=True, - qk_layernorm=True, - normalization="RMSNorm", - ) - return get_gpt_layer_with_transformer_engine_spec( - num_experts=8, - moe_grouped_gemm=True, - qk_layernorm=True, - fp8=False, - ) +def _make_layer_spec(config): + return get_qwen3_vl_hybrid_stack_spec(config) @pytest.fixture(scope="module", autouse=True) @@ -202,9 +188,10 @@ def _setup_parallel_state(): def _build_model(): _setup_parallel_state() pg_collection = ProcessGroupCollection.use_mpu_process_groups() + language_config = _make_language_config() return Qwen3OmniModel( - language_transformer_config=_make_language_config(), - language_transformer_layer_spec=_make_layer_spec(), + language_transformer_config=language_config, + language_transformer_layer_spec=_make_layer_spec(language_config), thinker_transformer_config=_make_toy_thinker_config(), parallel_output=True, pre_process=True,