feat: Add support for Qwen 2.5-Omni model bridge

[martinzhang03]

Summary

This PR adds Qwen2.5-Omni support to Megatron-Bridge: HF thinker / talker / token2wav alignment, conversion bridge, recipes, docs, and tests.

Stacking / rebase note: This branch is built on top of community PR #2831 — Qwen3-Omni (Omni3 multimodal layout, tests, and shared qwen_omni direction). Qwen2.5-Omni work extends that foundation rather than duplicating parallel scaffolding.

Head branch workflow: Development continued on feat/omni-sp (branched from the original PR head feat/omni-support). The fork branch feat/omni-support was updated to match feat/omni-sp so PR #2296 continues to use the same head branch name while including all latest commits.

---

What changed (high level)

Qwen2.5-Omni package and integration

• megatron.bridge.models.qwen_omni: dedicated package for Qwen2.5-Omni (provider, bridge, modeling for thinker/talker/token2wav), with registration and exports aligned with the Omni3 layout from [model] Add initial Qwen3-Omni support in Megatron-Bridge #2831.
• Provider / layers: dense stack defaults to Megatron-Core local GPT layers when TE is not used; optional TE path via env opt-in; persist_layer_norm aligned with torch norm when using local layers (replace(..., persist_layer_norm=False) on the non-TE path).
• mRoPE / distributed: fixes in Qwen3-VL rope path used by the thinker LM where needed for CP / process-group behavior.
• Conversion example robustness: torchrun_main writes TorchElastic-compatible error.json where possible (including non-Exception exits), prefers normal process exit over hard os._exit for better diagnostics; conversion script avoids sys.exit for the “not launched under torchrun” path in favor of RuntimeError so @record can capture failures.

Functional test and subprocess harness

• tests/functional_tests/models/qwen_omni/test_qwen25_omni_conversion.py: runs torch.distributed.run with sys.executable, --log_dir + -r 3, no coverage run --parallel-mode (avoids opaque error_file: <N/A>); prints error.json on failure; sets conservative NCCL_IB_DISABLE / NCCL_NET defaults for single-node functional runs on hosts with broken IB / gIB plugin expectations.

Conversion / distributed edge cases (single rank / PP=1)

• model_bridge / peft_bridge: skip torch.distributed.all_gather_object when pipeline-parallel group size is 1 (avoids observed crashes on some PyTorch/NCCL stacks).
• auto_bridge.save_hf_weights: dist.barrier() only when world_size > 1; for NCCL multi-rank collectives, use barrier(device_ids=[torch.cuda.current_device()]) where appropriate to avoid device-guessing warnings and segfaults seen on T4-class single-node runs.

Misc

• HF toy checkpoint for tests: thinker_config serialization via to_dict(), spk_dict.pt for audio-off reload, tokenizer sidecar files as needed.
• Test data download helper hardened (e.g. optional GH_TOKEN, clearer failure behavior).
• Docs / recipes / launch scripts updated for the full Qwen2.5-Omni PR surface (as in this branch).

---

Tests run (local, all passed)

Commands and results (GPU VM, uv run --no-sync pytest):

Scope	Command	Result
Qwen2.5-Omni HF↔Megatron roundtrip (functional)	pytest -v tests/functional_tests/models/qwen_omni/test_qwen25_omni_conversion.py	2 passed
Qwen2.5-Omni modeling (unit)	pytest -v tests/unit_tests/models/qwen_omni/modeling_qwen25_omni/test_omni_model.py	4 passed
Qwen2.5-Omni recipes (unit)	pytest -v tests/unit_tests/recipes/qwen_vl/test_qwen25_omni_recipes.py	2 passed

Note: This is not a substitute for full upstream CI; NVIDIA maintainers still need to trigger/approve CI per repo policy.

---

Relationship to PR #2831 (Omni3)

• [model] Add initial Qwen3-Omni support in Megatron-Bridge #2831 introduces Qwen3-Omni support (provider, bridge, tests, docs).
• This PR rebases / layers Qwen2.5-Omni on that work so both models share consistent Omni packaging and patterns where possible.

If #2831 is updated before merge, this branch may need a follow-up rebase to stay aligned.

---

Future work / follow-ups

• Broader CI matrix (multi-GPU TP > 1, different NCCL / IB environments).
• Optional coverage path behind a flag (e.g. MEGATRON_BRIDGE_ENABLE_COVERAGE=1) once the default debug-first subprocess harness is stable.
• Any upstream API migrations (e.g. ModelConfig + ModelBuilder) called out in deprecation warnings in the conversion example.
• Coordinate merge order with [model] Add initial Qwen3-Omni support in Megatron-Bridge #2831 to minimize churn for reviewers (stacked PR or sequential merge as maintainers prefer).

---

Maintenance

This PR description will be edited again if additional commits land on the head branch (e.g. review feedback, rebase onto updated main or updated #2831).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces the Qwen25OmniBridge class, a new Megatron bridge that converts HuggingFace's Qwen2.5OmniForConditionalGeneration model to Megatron-Core's Qwen25VLModel. The bridge extracts and validates nested configurations, derives model parameters, and provides parameter mapping registries for weight transformation. Comprehensive unit tests validate configuration handling and edge cases.

Changes

Cohort / File(s)	Summary
Bridge Implementation src/megatron/bridge/models/qwen_vl/qwen25_omni_bridge.py	New bridge class with provider_bridge method that constructs a Qwen25VLModelProvider from HF PreTrainedVLM, extracting and validating thinker and text configs, deriving mRoPE and rope_theta with fallbacks, and populating vision/token IDs and dtype settings. Includes mapping_registry method providing parameter mappings for embeddings, layers, multimodal components (visual, audio_tower, talker, token2wav), QKV, and GatedMLP transformations.
Module Exports src/megatron/bridge/models/__init__.py, src/megatron/bridge/models/qwen_vl/__init__.py	Import and export Qwen25OmniBridge through module hierarchy to expose new bridge publicly.
Test Suite tests/unit_tests/models/qwen_vl/test_qwen25_omni_bridge.py	Comprehensive test module validating bridge initialization, provider creation, configuration wiring (num_layers, hidden_size, token IDs, dtype handling), token ID overrides and fallbacks, rope parameter extraction, embedding tying behavior, and error conditions for missing configs.

Sequence Diagram

sequenceDiagram
    participant Client
    participant Qwen25OmniBridge
    participant ConfigExtractor
    participant Qwen25VLModelProvider
    participant MappingRegistry

    Client->>Qwen25OmniBridge: provider_bridge(hf_pretrained)
    Qwen25OmniBridge->>ConfigExtractor: extract thinker_config & text_config
    ConfigExtractor-->>Qwen25OmniBridge: configs (or error if missing)
    Qwen25OmniBridge->>Qwen25OmniBridge: validate configs present
    Qwen25OmniBridge->>Qwen25OmniBridge: derive mRoPE & rope_theta
    Qwen25OmniBridge->>Qwen25OmniBridge: populate vision/token IDs<br/>& dtype settings
    Qwen25OmniBridge->>Qwen25VLModelProvider: construct with all params
    Qwen25VLModelProvider-->>Qwen25OmniBridge: provider instance
    Qwen25OmniBridge-->>Client: configured Qwen25VLModelProvider

    Client->>Qwen25OmniBridge: mapping_registry()
    Qwen25OmniBridge->>MappingRegistry: construct with layer & multimodal mappings
    MappingRegistry-->>Client: MegatronMappingRegistry

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'feat: Add support for Qwen 2.5-Omni model bridge' clearly and accurately summarizes the main change: adding support for the Qwen 2.5-Omni model bridge to Megatron, which is reflected across all modified files including the bridge implementation, public exports, and comprehensive tests.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Test Results For Major Changes	✅ Passed	PR documentation explicitly states '36 tests passed' for new comprehensive test suite covering configuration conversion, parameter mapping, and weight initialization for the major new Qwen 2.5-Omni model bridge.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In `@src/megatron/bridge/models/qwen_vl/qwen25_omni_bridge.py`:
- Around line 33-39: Update the class docstring to correctly reference the
bridge target Qwen25VLModel instead of "Megatron-Core GPTModel formats"; locate
the docstring for the Qwen 2.5-Omni Bridge (the class decorated with the target
Qwen25VLModel) and replace the inaccurate phrase so it describes conversion
between HuggingFace Qwen2_5OmniForConditionalGeneration and Qwen25VLModel
formats, including weight mappings and configuration translation for multimodal
models.
- Around line 107-109: The current assignment uses "or" which treats 0 as falsy
and can drop a legitimate token id; change the logic for image_token_id and
video_token_id to prefer thinker_config.image_token_id / image_token_index only
when the first is None (e.g., read getattr(..., None) into a variable and if
that variable is None then use getattr(..._index, <default>)); update references
to image_token_id and video_token_id accordingly. Also address the long
commented-out audio_token_id line: move the explanatory note into a separate
shorter comment above the commented-out declaration (or shorten it) so the line
stays under 119 chars and clearly documents why audio_token_id is commented out.

🧹 Nitpick comments (4)

src/megatron/bridge/models/qwen_vl/qwen25_omni_bridge.py (1)

86-103: Consider extracting rope parameter resolution into a helper.

The mrope_section and rope_theta extraction blocks share the same pattern (try text_config attr → try rope_parameters dict → try rope_parameters object attr → default). A small helper like _get_rope_param(text_config, key, default) would reduce the nesting and duplication.

tests/unit_tests/models/qwen_vl/test_qwen25_omni_bridge.py (3)

358-514: Extract the repeated mapping-names collection into a helper.

The same ~10-line block that gathers mapping_names from registry.mappings is copy-pasted across 8 test methods (test_mapping_registry_contains_required_mappings, test_mapping_registry_thinker_model_mappings, test_mapping_registry_visual_params, etc.). Extract it into a helper function or fixture to reduce duplication.

Example helper extraction

+def _collect_mapping_names(registry):
+    """Collect all Megatron and HF param names from the registry mappings."""
+    names = []
+    for mapping in registry.mappings:
+        if hasattr(mapping, "megatron_param"):
+            names.append(str(mapping.megatron_param))
+        hf = getattr(mapping, "hf_param", None)
+        if isinstance(hf, dict):
+            names.extend(str(v) for v in hf.values())
+        elif isinstance(hf, str):
+            names.append(hf)
+    return names
+
+
 class TestQwen25OmniBridgeMappingRegistry:
     """Test mapping_registry method functionality."""

     ...

     def test_mapping_registry_contains_required_mappings(self, qwen25_omni_bridge):
         """Test mapping_registry contains all required parameter mappings."""
         registry = qwen25_omni_bridge.mapping_registry()
-
-        # Extract mappings - registry should contain mappings for common parameters
-        mappings = registry.mappings
-        assert len(mappings) > 0
-
-        # Check that we have mappings for embeddings, output layer, layernorms
-        mapping_names = []
-        for mapping in mappings:
-            # Collect Megatron param pattern
-            if hasattr(mapping, "megatron_param"):
-                mapping_names.append(str(getattr(mapping, "megatron_param")))
-            # Collect HF param pattern(s)
-            hf = getattr(mapping, "hf_param", None)
-            if isinstance(hf, dict):
-                mapping_names.extend([str(v) for v in hf.values()])
-            elif isinstance(hf, str):
-                mapping_names.append(hf)
+        assert len(registry.mappings) > 0
+        mapping_names = _collect_mapping_names(registry)

         # Should contain word embeddings mapping (thinker.model.embed_tokens)
         has_embeddings = any("embed_tokens" in name or "word_embeddings" in name for name in mapping_names)

Then apply the same pattern to the other 7 test methods.

---

28-45: Mock auto-attributes may leak into tests as unexpected values.

Mock() auto-creates attributes on access, so hasattr(text_config, "rope_parameters") returns True even though rope_parameters was never explicitly set. This means the bridge's mrope_section extraction path will execute against a nested Mock chain, resulting in a Mock object (not a list or None) being assigned to mrope_section in tests that don't explicitly override rope_parameters.

The targeted mRoPE tests handle this via delattr, but the basic config tests silently pass a Mock as mrope_section to the provider. Consider using spec= on the Mock or explicitly initializing text_config.rope_parameters = None in the fixture to make the default behavior deterministic.

Proposed fix

 `@pytest.fixture`
 def mock_text_config():
     """Create a mock text_config for Qwen2.5-Omni."""
     text_config = Mock()
     text_config.num_hidden_layers = 32
     text_config.hidden_size = 4096
     text_config.intermediate_size = 11008
     text_config.num_attention_heads = 32
     text_config.num_key_value_heads = 32
     text_config.initializer_range = 0.02
     text_config.rms_norm_eps = 1e-6
     text_config.vocab_size = 151936
     text_config.max_position_embeddings = 4096
     text_config.rope_theta = 1000000.0
     text_config.tie_word_embeddings = False
     text_config.bos_token_id = 151643
     text_config.eos_token_id = 151645
+    text_config.rope_parameters = None
     return text_config

---

85-86: Missing pytest.mark categorization on test classes.

Per coding guidelines, tests should be categorized with pytest.mark (e.g., @pytest.mark.unit). None of the test classes in this file have markers.

As per coding guidelines: "Use 'pytest.mark' to categorize tests (unit, integration, system)."

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Inaccurate docstring: references GPTModel instead of Qwen25VLModel.

The class docstring mentions "Megatron-Core GPTModel formats" but the bridge target (as declared in the decorator on line 31) is Qwen25VLModel.

Proposed fix

     """
     Megatron Bridge for Qwen 2.5-Omni Conditional Generation.

     This bridge handles the conversion between HuggingFace Qwen2_5OmniForConditionalGeneration
-    and Megatron-Core GPTModel formats, including weight mappings and
+    and Megatron-Core Qwen25VLModel formats, including weight mappings and
     configuration translation for multimodal models.
     """

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	"""
	Megatron Bridge for Qwen 2.5-Omni Conditional Generation.
	This bridge handles the conversion between HuggingFace Qwen2_5OmniForConditionalGeneration
	and Megatron-Core GPTModel formats, including weight mappings and
	configuration translation for multimodal models.
	"""
	"""
	Megatron Bridge for Qwen 2.5-Omni Conditional Generation.
	This bridge handles the conversion between HuggingFace Qwen2_5OmniForConditionalGeneration
	and Megatron-Core Qwen25VLModel formats, including weight mappings and
	configuration translation for multimodal models.
	"""

🤖 Prompt for AI Agents

In `@src/megatron/bridge/models/qwen_vl/qwen25_omni_bridge.py` around lines 33 -
39, Update the class docstring to correctly reference the bridge target
Qwen25VLModel instead of "Megatron-Core GPTModel formats"; locate the docstring
for the Qwen 2.5-Omni Bridge (the class decorated with the target Qwen25VLModel)
and replace the inaccurate phrase so it describes conversion between HuggingFace
Qwen2_5OmniForConditionalGeneration and Qwen25VLModel formats, including weight
mappings and configuration translation for multimodal models.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Falsy token ID (e.g. 0) would silently fall through to the fallback.

The or operator treats 0 as falsy, so if image_token_id or video_token_id is legitimately 0, it would be ignored in favor of the *_token_index fallback or the hardcoded default. While token ID 0 is unlikely for Qwen special tokens, using is None checks would be more robust.

Also, line 109 (commented-out audio_token_id) far exceeds the 119-character line limit. Per coding guidelines, commented-out code should include a comment describing why it's commented out — the inline note is there but the line is too long. Consider moving the explanation to a standalone comment line above.

Proposed fix

-        image_token_id = getattr(thinker_config, "image_token_id", None) or getattr(thinker_config, "image_token_index", 151655)
-        video_token_id = getattr(thinker_config, "video_token_id", None) or getattr(thinker_config, "video_token_index", 151656)
-        # audio_token_id = getattr(thinker_config, "audio_token_id", None) or getattr(thinker_config, "audio_token_index", 151646) # audio_token_id is not used by Qwen25VLModelProvider, which only supports image and video tokens
+        image_token_id = getattr(thinker_config, "image_token_id", None)
+        if image_token_id is None:
+            image_token_id = getattr(thinker_config, "image_token_index", 151655)
+        video_token_id = getattr(thinker_config, "video_token_id", None)
+        if video_token_id is None:
+            video_token_id = getattr(thinker_config, "video_token_index", 151656)
+        # audio_token_id is not used by Qwen25VLModelProvider,
+        # which only supports image and video tokens.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	image_token_id = getattr(thinker_config, "image_token_id", None) or getattr(thinker_config, "image_token_index", 151655)
	video_token_id = getattr(thinker_config, "video_token_id", None) or getattr(thinker_config, "video_token_index", 151656)
	# audio_token_id = getattr(thinker_config, "audio_token_id", None) or getattr(thinker_config, "audio_token_index", 151646) # audio_token_id is not used by Qwen25VLModelProvider, which only supports image and video tokens
	image_token_id = getattr(thinker_config, "image_token_id", None)
	if image_token_id is None:
	image_token_id = getattr(thinker_config, "image_token_index", 151655)
	video_token_id = getattr(thinker_config, "video_token_id", None)
	if video_token_id is None:
	video_token_id = getattr(thinker_config, "video_token_index", 151656)
	# audio_token_id is not used by Qwen25VLModelProvider,
	# which only supports image and video tokens.

🤖 Prompt for AI Agents

In `@src/megatron/bridge/models/qwen_vl/qwen25_omni_bridge.py` around lines 107 -
109, The current assignment uses "or" which treats 0 as falsy and can drop a
legitimate token id; change the logic for image_token_id and video_token_id to
prefer thinker_config.image_token_id / image_token_index only when the first is
None (e.g., read getattr(..., None) into a variable and if that variable is None
then use getattr(..._index, <default>)); update references to image_token_id and
video_token_id accordingly. Also address the long commented-out audio_token_id
line: move the explanatory note into a separate shorter comment above the
commented-out declaration (or shorten it) so the line stays under 119 chars and
clearly documents why audio_token_id is commented out.

[yaoyu-33]

@martinzhang03 : thanks for contribution! It all looks good, but we are having a refactor on-going right now. Do you think you can rebase on top of https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/2250/changes?

[yuekaizhang]

@martinzhang03 Do you have any plans to support the Qwen Omni audio modality? If not, I would be happy to further integrate the audio modality based on your PR.

[martinzhang03]

@martinzhang03 : thanks for contribution! It all looks good, but we are having a refactor on-going right now. Do you think you can rebase on top of https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/2250/changes?

Hi @yaoyu-33 Thanks for the feedback. I'm currently working on rebasing my branch on top of PR #2250 and will ping you once its complete.

[martinzhang03]

@martinzhang03 Do you have any plans to support the Qwen Omni audio modality? If not, I would be happy to further integrate the audio modality based on your PR.

Hi @yuekaizhang Thanks. I don't have immediate plans to implement the audio modality. I would definitely welcome your help in further integration once the PR is stabilized. Please feel free to build upon this work.

[martinzhang03]

@martinzhang03 : thanks for contribution! It all looks good, but we are having a refactor on-going right now. Do you think you can rebase on top of https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/2250/changes?

Hi @yaoyu-33 Thanks for the feedback. I'm currently working on rebasing my branch on top of PR #2250 and will ping you once its complete.

Done.

[martinzhang03]

Hi @yaoyu-33 I have updated the PR to include full Sequence Parallelism (SP) support for the Omni model. This addition complements the base model registration by providing the high-performance training infra needed for industrial scale-up. I've also updated the PR description with the technical details.

[martinzhang03]

Hi @yaoyu-33 this PR has been updated. I've rebased it on top of PR #2831. Tests are added and all test cases passed.

[yaoyu-33]

https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/qwen_omni we added qwen25 omni in another pr.