[model, bugfix] qwen3-omni: fix Megatron parity for external RL integration (#4988)

Signed-off-by: hbhflw2000 <417911774@qq.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Co-authored-by: yaoyu-33 <yaoyu.094@gmail.com>
Co-authored-by: Yuekai Zhang <zhangyuekai@foxmail.com>

Author: 3 people

examples/models/qwen/qwen3_omni/local_train_thinker_full.sh

@@ -255,6 +255,12 @@ PY
 
 mkdir -p "${THINKER_ONLY_MIRROR_DIR}"
 EFFECTIVE_HF_MODEL_PATH=$(prepare_thinker_only_hf_path "${HF_MODEL_PATH}" "${THINKER_ONLY_MIRROR_DIR}" "${PYTHON_BIN}")
+# Recipe construction calls AutoBridge before CLI overrides are applied.  Point
+# it at the local thinker-only mirror so all ranks avoid Hub access and its
+# shared config lock.
+export QWEN3_OMNI_HF_PATH="${EFFECTIVE_HF_MODEL_PATH}"
+export HF_HUB_OFFLINE=1
+export TRANSFORMERS_OFFLINE=1
 
 CMD=(
     "${PYTHON_BIN}" -m torch.distributed.run

src/megatron/bridge/models/qwen_omni/qwen3_omni_bridge.py

@@ -104,6 +104,9 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> Qwen3OmniModelPr
             spatial_merge_size=getattr(vision_config, "spatial_merge_size", 2),
             position_embedding_type="mrope",
             mrope_section=rope_scaling.get("mrope_section", [24, 20, 20]),
+            # HF/vLLM Qwen3-Omni expose interleaved keys in rope_scaling, but the
+            # thinker text path matches Megatron's non-interleaved rotary layout.
+            rotary_interleaved=False,
         )
         return provider
 

src/megatron/bridge/models/qwen_omni/qwen3_omni_provider.py

@@ -42,6 +42,12 @@
 from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention
 
 
+def _bool_config_enabled(value) -> bool:
+    if isinstance(value, str):
+        return value.lower() in {"1", "true", "yes", "on"}
+    return bool(value)
+
+
 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:
@@ -98,9 +104,12 @@ class Qwen3OmniModelProvider(GPTModelProvider):
     moe_router_load_balancing_type: str = "aux_loss"
     moe_aux_loss_coeff: float = 1e-3
     moe_router_pre_softmax: bool = False
+    moe_router_dtype: str = "fp32"
+    moe_router_score_function: str = "softmax"
     moe_token_dispatcher_type: str = "alltoall"
     moe_permute_fusion: bool = True
     masked_softmax_fusion: bool = False
+    moe_enable_routing_replay: bool = False
 
     image_token_id: int = 151655
     video_token_id: int = 151656
@@ -146,6 +155,11 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None):
             )
 
         use_local_attention = self.attention_backend in {AttnBackend.local, "local"}
+        if use_local_attention and _bool_config_enabled(self.sequence_parallel):
+            raise ValueError(
+                "Qwen3-Omni sequence_parallel requires a Transformer Engine attention backend; "
+                "attention_backend=local selects the torch local attention path."
+            )
         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,

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/attention.py

@@ -15,6 +15,7 @@
 
 import torch
 from einops import rearrange
+from megatron.core.extensions.transformer_engine import TEDotProductAttention
 from megatron.core.transformer.attention import (
     HAVE_FA3,
     BaseInferenceContext,
@@ -31,6 +32,23 @@
 from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.rope import apply_rotary_pos_emb_absolute
 
 
+def _qwen_attention_mask_for_core_attention(
+    core_attention: object,
+    attention_mask: Tensor | None,
+    packed_seq_params: PackedSeqParams | None,
+) -> Tensor | None:
+    """Translate Qwen's 2D valid-token mask to TE's masked-token convention."""
+    if (
+        not isinstance(core_attention, TEDotProductAttention)
+        or packed_seq_params is not None
+        or not isinstance(attention_mask, Tensor)
+        or attention_mask.ndim != 2
+        or attention_mask.dtype != torch.bool
+    ):
+        return attention_mask
+    return (~attention_mask).unsqueeze(1).unsqueeze(1)
+
+
 class Qwen3VLSelfAttention(SelfAttention):
     """
     Overrides the SelfAttention class, the difference is that qwen3vl uses apply_rotary_pos_emb_absolute
@@ -239,6 +257,11 @@ def forward(
         # ==================================
 
         nvtx_range_push(suffix="core_attention")
+        core_attention_mask = _qwen_attention_mask_for_core_attention(
+            self.core_attention,
+            attention_mask,
+            packed_seq_params,
+        )
         if (
             isinstance(self.core_attention, DotProductAttention)
             and packed_seq_params is None
@@ -252,7 +275,7 @@ def forward(
                 query,
                 key,
                 value,
-                attention_mask,
+                core_attention_mask,
                 attn_mask_type=attn_mask_type,
                 attention_bias=attention_bias,
                 packed_seq_params=packed_seq_params,
@@ -264,7 +287,7 @@ def forward(
                     query,
                     key,
                     value,
-                    attention_mask,
+                    core_attention_mask,
                     attn_mask_type=attn_mask_type,
                     attention_bias=attention_bias,
                     packed_seq_params=packed_seq_params,

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/rope.py

@@ -96,8 +96,7 @@ def get_packed_seq_attention_mask(input_ids: torch.Tensor, packed_seq_params: Pa
 
 
 class Qwen3VLMultimodalRotaryEmbedding(nn.Module):
-    """Multimodal Rotary Embedding for language model.
-    only support for qwen3vl
+    """Multimodal Rotary Embedding for Qwen3-VL/Qwen3-Omni language model.
 
     Args:
         kv_channels (int): Projection weights dimension in multi-head attention. Obtained
@@ -127,7 +126,6 @@ def __init__(
         if rotary_percent < 1.0:
             dim = int(dim * rotary_percent)
         self.rotary_interleaved = rotary_interleaved
-        assert not self.rotary_interleaved, "only support qwen3vl"
 
         self.seq_len_interpolation_factor = seq_len_interpolation_factor
         self.inv_freq = 1.0 / (

src/megatron/bridge/recipes/qwen_omni/h100/qwen3_omni.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 
 import torch
@@ -34,6 +35,11 @@
 _QWEN3_OMNI_HF_PATH = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
 
 
+def _qwen3_omni_hf_path() -> str:
+    """Resolve an explicit local checkpoint before falling back to the public id."""
+    return os.environ.get("QWEN3_OMNI_HF_PATH", _QWEN3_OMNI_HF_PATH)
+
+
 def qwen3_omni_30b_a3b_sft_8gpu_h100_bf16_config() -> "ConfigContainer":
     """Return an 8-GPU thinker-only SFT config for Qwen3-Omni 30B-A3B.
 
@@ -42,7 +48,8 @@ def qwen3_omni_30b_a3b_sft_8gpu_h100_bf16_config() -> "ConfigContainer":
 
     cfg = _sft_common_vlm()
 
-    cfg.model = AutoBridge.from_hf_pretrained(_QWEN3_OMNI_HF_PATH).to_megatron_provider(load_weights=False)
+    hf_model_path = _qwen3_omni_hf_path()
+    cfg.model = AutoBridge.from_hf_pretrained(hf_model_path).to_megatron_provider(load_weights=False)
     cfg.model.seq_length = 4096
     cfg.model.tensor_model_parallel_size = 1
     cfg.model.pipeline_model_parallel_size = 1
@@ -78,7 +85,7 @@ def qwen3_omni_30b_a3b_sft_8gpu_h100_bf16_config() -> "ConfigContainer":
     cfg.scheduler = scheduler_cfg
 
     cfg.dataset.seq_length = 4096
-    cfg.dataset.hf_processor_path = _QWEN3_OMNI_HF_PATH
+    cfg.dataset.hf_processor_path = hf_model_path
     cfg.dataset.enable_in_batch_packing = False
     cfg.dataset.skip_getting_attention_mask_from_dataset = False
 
@@ -105,7 +112,7 @@ def qwen3_omni_30b_a3b_sft_8gpu_h100_bf16_hf_json_config() -> "ConfigContainer":
     cfg.dataset = DirectHFSFTDatasetConfig(
         seq_length=cfg.model.seq_length,
         preprocessing=ChatSFTPreprocessingConfig(),
-        hf_processor_path=_QWEN3_OMNI_HF_PATH,
+        hf_processor_path=_qwen3_omni_hf_path(),
         source=HFDatasetSourceConfig(
             path_or_dataset="json",
             split="train",

tests/unit_tests/models/qwen_omni/test_qwen3_omni_bridge.py

@@ -41,7 +41,12 @@ def mock_text_config():
     text_config.max_position_embeddings = 32768
     text_config.rope_theta = 1000000.0
     text_config.attention_bias = False
-    text_config.rope_scaling = {"mrope_section": [24, 20, 20]}
+    text_config.rope_scaling = {
+        "rope_type": "default",
+        "interleaved": True,
+        "mrope_interleaved": True,
+        "mrope_section": [24, 20, 20],
+    }
     return text_config
 
 
@@ -104,6 +109,7 @@ def test_provider_bridge_basic_config(self, mock_hf_pretrained):
         assert provider.moe_router_topk == 8
         assert provider.share_embeddings_and_output_weights is False
         assert provider.mrope_section == [24, 20, 20]
+        assert provider.rotary_interleaved is False
         assert provider.language_max_sequence_length == 32768
         assert provider.patch_size == 32
         assert provider.temporal_patch_size == 4
@@ -145,6 +151,56 @@ def test_mapping_registry(self):
         )
         assert any("thinker.language_model.decoder.layers.*.mlp.router.weight" in name for name in mapping_names)
 
+    def test_mapping_registry_resolves_parallel_parameter_names(self):
+        bridge = Qwen3OmniBridge()
+        registry = bridge.mapping_registry()
+
+        expected = {
+            "thinker.language_model.decoder.layers.0.self_attention.linear_qkv.weight": {
+                "q": "thinker.model.layers.0.self_attn.q_proj.weight",
+                "k": "thinker.model.layers.0.self_attn.k_proj.weight",
+                "v": "thinker.model.layers.0.self_attn.v_proj.weight",
+            },
+            "thinker.language_model.decoder.layers.0.self_attention.linear_qkv.layer_norm_weight": (
+                "thinker.model.layers.0.input_layernorm.weight"
+            ),
+            "thinker.language_model.decoder.layers.0.input_layernorm.weight": (
+                "thinker.model.layers.0.input_layernorm.weight"
+            ),
+            "thinker.language_model.decoder.layers.0.pre_mlp_layernorm.weight": (
+                "thinker.model.layers.0.post_attention_layernorm.weight"
+            ),
+            "thinker.language_model.decoder.layers.0.self_attention.q_layernorm.weight": (
+                "thinker.model.layers.0.self_attn.q_norm.weight"
+            ),
+            "thinker.language_model.decoder.layers.0.self_attention.k_layernorm.weight": (
+                "thinker.model.layers.0.self_attn.k_norm.weight"
+            ),
+            "thinker.language_model.decoder.layers.0.self_attention.linear_proj.weight": (
+                "thinker.model.layers.0.self_attn.o_proj.weight"
+            ),
+            "thinker.language_model.decoder.layers.0.mlp.router.weight": "thinker.model.layers.0.mlp.gate.weight",
+            "thinker.language_model.decoder.layers.0.mlp.experts.linear_fc1.weight7": {
+                "gate": "thinker.model.layers.0.mlp.experts.7.gate_proj.weight",
+                "up": "thinker.model.layers.0.mlp.experts.7.up_proj.weight",
+            },
+            "thinker.language_model.decoder.layers.0.mlp.experts.linear_fc2.weight7": (
+                "thinker.model.layers.0.mlp.experts.7.down_proj.weight"
+            ),
+            "thinker.language_model.decoder.layers.0.mlp.experts.local_experts.7.linear_fc1.weight": {
+                "gate": "thinker.model.layers.0.mlp.experts.7.gate_proj.weight",
+                "up": "thinker.model.layers.0.mlp.experts.7.up_proj.weight",
+            },
+            "thinker.language_model.decoder.layers.0.mlp.experts.local_experts.7.linear_fc2.weight": (
+                "thinker.model.layers.0.mlp.experts.7.down_proj.weight"
+            ),
+        }
+
+        for megatron_param, hf_param in expected.items():
+            mapping = registry.megatron_to_hf_lookup(megatron_param)
+            assert mapping is not None, megatron_param
+            assert mapping.hf_param == hf_param
+
     def test_provider_bridge_warns_for_audio_output_stack(self, mock_hf_pretrained, caplog):
         mock_hf_pretrained.config.enable_audio_output = True
         caplog.set_level("WARNING")

tests/unit_tests/models/qwen_omni/test_qwen3_omni_provider.py

@@ -30,6 +30,13 @@
 
 
 class TestQwen3OmniModelProvider:
+    @pytest.mark.parametrize(
+        ("value", "expected"),
+        [("true", True), ("ON", True), ("false", False), (True, True), (0, False)],
+    )
+    def test_bool_config_enabled(self, value, expected):
+        assert qwen3_omni_provider._bool_config_enabled(value) is expected
+
     def test_qwen3_omni_model_provider_initialization(self):
         provider = Qwen3OmniModelProvider(
             num_layers=48,
@@ -43,7 +50,10 @@ def test_qwen3_omni_model_provider_initialization(self):
         assert provider.position_embedding_type == "mrope"
         assert provider.scatter_embedding_sequence_parallel is False
         assert provider.qk_layernorm is True
+        assert provider.moe_router_dtype == "fp32"
+        assert provider.moe_router_score_function == "softmax"
         assert provider.masked_softmax_fusion is False
+        assert provider.moe_enable_routing_replay is False
         assert provider.apply_rotary_pos_emb_in_fp32 is False
         assert provider.image_token_id == 151655
         assert provider.video_token_id == 151656
@@ -75,6 +85,16 @@ def test_qwen3_omni_custom_thinker_config(self):
 
         assert provider.thinker_config.text_config.hidden_size == 128
 
+    def test_qwen3_omni_uses_mcore_routing_replay_field(self):
+        provider = Qwen3OmniModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+            moe_enable_routing_replay=True,
+        )
+
+        assert provider.moe_enable_routing_replay is True
+
     def test_qwen3_omni_freeze_flags(self):
         provider = Qwen3OmniModelProvider(
             num_layers=48,
@@ -181,6 +201,51 @@ def test_default_backend_keeps_te_core_attention_when_available(self):
         assert self_attention_spec.module is Qwen3VLSelfAttention
         assert self_attention_spec.submodules.core_attention.__name__ == "TEDotProductAttention"
 
+    @pytest.mark.parametrize("attention_backend", [AttnBackend.unfused, "unfused"])
+    def test_unfused_backend_keeps_te_core_attention_with_sequence_parallel(self, attention_backend):
+        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=attention_backend,
+            sequence_parallel=True,
+        )
+
+        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.__name__ == "TEDotProductAttention"
+
+    def test_local_backend_with_sequence_parallel_hard_fails(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",
+            sequence_parallel=True,
+        )
+
+        with pytest.raises(ValueError, match="sequence_parallel requires a Transformer Engine attention backend"):
+            provider.provide()
+
     def test_default_backend_falls_back_to_local_spec_without_te(self):
         provider = Qwen3OmniModelProvider(
             num_layers=2,