[model] Add initial Qwen3-Omni support in Megatron-Bridge

docs/models/vlm/index.md

@@ -10,6 +10,7 @@ glm-45v.md
 ministral3.md
 nemotron-nano-v2-vl.md
 qwen2.5-vl.md
+qwen3-omni.md
 qwen3-vl.md
 qwen35-vl.md
 ```

docs/models/vlm/qwen3-omni.md

@@ -0,0 +1,42 @@
+# Qwen3-Omni
+
+Qwen3-Omni is a multimodal Qwen family model with text, image, video, and audio inputs. Megatron Bridge support for Qwen3-Omni reuses the existing Qwen3-VL language and vision path, and adds Qwen3-Omni-specific audio handling and checkpoint mappings.
+
+The current implementation is focused on checkpoint conversion and training-oriented multimodal forward paths. It supports single-rank functional validation for text, vision, and audio inputs, and keeps distributed parallel validation as a follow-up step.
+
+## Current Support
+
+- Hugging Face to Megatron Bridge checkpoint conversion for `Qwen/Qwen3-Omni-30B-A3B-Instruct`
+- Megatron Bridge to Hugging Face export for the same model family
+- Text, image, video, and audio multimodal forward paths
+- Qwen3-Omni-specific multimodal RoPE handling for Megatron Bridge runtime
+- Single-GPU smoke validation with a vertically trimmed checkpoint
+
+## Implementation Notes
+
+- The language backbone uses Megatron-Core modules through the existing Qwen3-VL GPT path.
+- The vision and audio towers currently reuse Hugging Face Qwen3-Omni components.
+- Multimodal embeddings are fused before the language decoder.
+- The current validation path focuses on correctness and checkpoint interoperability rather than distributed performance tuning.
+
+## Validation
+
+The current test coverage includes:
+
+- Provider unit tests
+- Bridge unit tests
+- Model unit tests for text, vision, and audio paths
+- Functional checkpoint conversion tests
+- End-to-end smoke tests using user-provided multimodal assets
+
+## Hugging Face Model Cards
+
+- Qwen3-Omni-30B-A3B-Instruct: `https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct`
+
+## Related Docs
+
+- Related VLM: [Qwen3-VL](qwen3-vl.md)
+- Related VLM: [Qwen 3.5](qwen35-vl.md)
+- Recipe usage: [Recipe usage](../../recipe-usage.md)
+- Customizing the training recipe configuration: [Configuration overview](../../training/config-container-overview.md)
+- Training entry points: [Entry points](../../training/entry-points.md)

src/megatron/bridge/models/__init__.py

@@ -182,6 +182,9 @@
     Qwen25OmniBridge,
     Qwen25OmniModel,
     Qwen25OmniModelProvider,
+    Qwen3OmniBridge,
+    Qwen3OmniModel,
+    Qwen3OmniModelProvider,
 )
 from megatron.bridge.models.qwen_vl import (
     Qwen25VLBridge,
@@ -359,6 +362,9 @@
     "Qwen25OmniModel",
     "Qwen25OmniBridge",
     "Qwen25OmniModelProvider",
+    "Qwen3OmniModel",
+    "Qwen3OmniBridge",
+    "Qwen3OmniModelProvider",
     "SarvamMLABridge",
     "SarvamMoEBridge",
 ]

src/megatron/bridge/models/qwen/qwen_provider.py

@@ -18,13 +18,21 @@
 
 import torch
 import torch.nn.functional as F
-from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
-    get_transformer_block_with_experimental_attention_variant_spec,
-)
 from megatron.core.transformer.spec_utils import ModuleSpec
 
 from megatron.bridge.models.gpt_provider import GPTModelProvider
 
+try:
+    from megatron.core.models.gpt.experimental_attention_variant_module_specs import (
+        get_transformer_block_with_experimental_attention_variant_spec,
+    )
+except ModuleNotFoundError:
+    def get_transformer_block_with_experimental_attention_variant_spec(*args, **kwargs):
+        raise NotImplementedError(
+            "Qwen3-Next requires Megatron-Core support for "
+            "experimental_attention_variant_module_specs."
+        )
+
 
 try:
     import transformer_engine  # type: ignore  # noqa: F401

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

@@ -13,12 +13,18 @@
 # limitations under the License.
 
 from megatron.bridge.models.qwen_omni.modeling_qwen25_omni.model import Qwen25OmniModel
+from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.model import Qwen3OmniModel
 from megatron.bridge.models.qwen_omni.qwen25_omni_bridge import Qwen25OmniBridge
 from megatron.bridge.models.qwen_omni.qwen25_omni_provider import Qwen25OmniModelProvider
+from megatron.bridge.models.qwen_omni.qwen3_omni_bridge import Qwen3OmniBridge
+from megatron.bridge.models.qwen_omni.qwen3_omni_provider import Qwen3OmniModelProvider
 
 
 __all__ = [
     "Qwen25OmniBridge",
     "Qwen25OmniModel",
     "Qwen25OmniModelProvider",
+    "Qwen3OmniBridge",
+    "Qwen3OmniModel",
+    "Qwen3OmniModelProvider",
 ]

src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2025, 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.
+
+from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.model import Qwen3OmniModel
+from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.transformer_config import (
+    Qwen3OmniTransformerConfig,
+)
+
+
+__all__ = [
+    "Qwen3OmniModel",
+    "Qwen3OmniTransformerConfig",
+]

src/megatron/bridge/models/qwen_omni/modeling_qwen3_omni/model.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2025, 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 torch
+from megatron.core import InferenceParams
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+    Qwen3OmniMoeCode2WavConfig,
+    Qwen3OmniMoeTalkerConfig,
+    Qwen3OmniMoeThinkerConfig,
+)
+
+from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.thinker_model import Qwen3OmniThinkerModel
+from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.transformer_config import (
+    Qwen3OmniTransformerConfig,
+)
+
+
+class Qwen3OmniModel(MegatronModule):
+    """Qwen3-Omni model wrapper."""
+
+    def __init__(
+        self,
+        language_transformer_config: Qwen3OmniTransformerConfig,
+        language_transformer_layer_spec: ModuleSpec,
+        thinker_transformer_config: Qwen3OmniMoeThinkerConfig,
+        talker_transformer_config: Qwen3OmniMoeTalkerConfig | None = None,
+        code2wav_transformer_config: Qwen3OmniMoeCode2WavConfig | None = None,
+        parallel_output: bool = True,
+        pre_process: bool = True,
+        post_process: bool = True,
+        add_encoder: bool = True,
+        add_decoder: bool = True,
+        pg_collection: ProcessGroupCollection | None = None,
+    ) -> None:
+        super().__init__(config=language_transformer_config)
+
+        self.thinker = Qwen3OmniThinkerModel(
+            language_transformer_config,
+            language_transformer_layer_spec,
+            thinker_transformer_config,
+            parallel_output,
+            pre_process,
+            post_process,
+            add_encoder,
+            add_decoder,
+            pg_collection,
+        )
+
+        self.talker_transformer_config = talker_transformer_config
+        self.code2wav_transformer_config = code2wav_transformer_config
+
+    def shared_embedding_or_output_weight(self):
+        return self.thinker.shared_embedding_or_output_weight()
+
+    def set_input_tensor(self, input_tensor) -> None:
+        return self.thinker.set_input_tensor(input_tensor)
+
+    def freeze(
+        self,
+        freeze_language_model: bool = False,
+        freeze_vision_model: bool = False,
+        freeze_audio_model: bool = False,
+    ):
+        return self.thinker.freeze(
+            freeze_language_model=freeze_language_model,
+            freeze_vision_model=freeze_vision_model,
+            freeze_audio_model=freeze_audio_model,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        loss_mask: torch.Tensor | None = None,
+        inference_params: InferenceParams | None = None,
+        packed_seq_params: PackedSeqParams | None = None,
+        extra_block_kwargs: dict | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return self.thinker(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+            loss_mask=loss_mask,
+            inference_params=inference_params,
+            packed_seq_params=packed_seq_params,
+            extra_block_kwargs=extra_block_kwargs,
+            **kwargs,
+        )