[contrib] Add Qwen2.5-Omni-7B with full Neuron speech pipeline (Thinker+Talker+Token2Wav, TP=4)

[whn09]

Description

Adds the Qwen2.5-Omni-7B multimodal model to contrib/ with full Neuron support for the entire speech pipeline: text generation, image understanding, audio understanding, and text-to-speech — all running on Neuron at TP=4.

Supersedes #122, refactored to a 100% zero-invasion contrib layout (no changes under src/neuronx_distributed_inference/).

Full Neuron Speech Pipeline

Steady-state per-run latency on trn2.48xlarge with all three Neuron models resident in the same process (TP=4, NeuronCores 0–3), averaged over 5 runs of the default prompt:

Component	Runtime	TP	Params	Per-run latency
Thinker (text)	Neuron	4	7B	0.34s
Vision encoder	Neuron	4	670M	— (on-demand, image input)
Audio encoder	CPU + Neuron	4	620M	20–34ms CPU pre/post (on-demand)
HF CPU forward (Phase 2 hidden states)	CPU bf16	—	7B	1.16s
Talker prep (CPU projection)	CPU	—	—	0.18s
Talker (codec tokens)	Neuron	4	690M	2.28s
Token2Wav DiT	Neuron	—	85M	(part of 14.14s)
Token2Wav BigVGAN	CPU	—	~15M	(part of 14.14s)
Pipeline total				18.10s

Audio: ~12s of speech per request. RTF: 1.51×.

One-time load (inside the same process, amortized across every subsequent run):

	Time
Thinker load + warmup	13.1s
HF CPU bf16 load	0.3s
Talker load	2.0s
Token2Wav DiT load	19.4s
Total cold start	~37s

So wall time for a single-prompt invocation is ~55s; for 5 back-to-back prompts it is ~129s. Once the process is up every subsequent prompt is 18s regardless of how many requests came before.

Zero-Invasion Design

All code lives under contrib/models/Qwen2.5-Omni-7B/. git diff upstream/main..HEAD -- src/ is empty. Three things that previously required modifying the core tree have been handled locally:

1. Talker per-step thinker state injection — The Talker needs encode_vision_to_input to run during both context encoding and token generation. Previously this required adding an apply_vision_during_token_gen gate in model_base.py. Now it's a get_model_output override on the contrib Talker subclass (src/modeling_qwen25_omni_talker.py), which pre-injects into inputs_embeds before calling super().get_model_output().
2. Model registration — The model registers itself via flat imports (from modeling_qwen25_omni import ...) from sys.path.insert, same pattern as upstream's contrib/Qwen2-Audio-7B. No entries added to utils/constants.py or inference_demo.py.
3. hf_adapter.py NameError — Upstream HuggingFaceGenerationAdapter.prepare_inputs_for_generation references an undefined tensor_capture_hook local, which breaks adapter.generate() for the Talker. A runtime shim in src/_upstream_compat.py patches this at import time. The underlying bug is addressed by companion PR Fix NameError in HuggingFaceGenerationAdapter.prepare_inputs_for_generation #136 — once Fix NameError in HuggingFaceGenerationAdapter.prepare_inputs_for_generation #136 merges the shim can be removed.

Single-Process Pipeline + NeuronCore Pinning

generate_qwen25_omni_speech.py loads all three Neuron-compiled models into the same Python process (Thinker TP=4, Talker TP=4, Token2Wav DiT single-device) and reuses them across runs. The earlier subprocess-per-component layout paid ~60s of cold start per request because each child forked, loaded the model, inferred, and exited.

Two details matter for keeping the DiT fast once it co-exists with the Thinker and Talker:

• NeuronCore pinning: the script sets NEURON_RT_VISIBLE_CORES=0-3 before any Neuron module is imported. Without this, the runtime places the single-device DiT NEFF on a different core group than the TP=4 Thinker/Talker and every DiT forward pays a cross-group scheduling penalty. This is documented in the README and users embedding the pipeline in their own entrypoint must set the env var the same way.
• HF CPU model in bf16: Phase 2 (hidden-state extraction for the Talker) loads the HuggingFace model in bfloat16 rather than float32. Downstream consumers round back to bf16 anyway, so float32 here is pure overhead — switching cut Phase 2 from ~10s to ~1s.

Default model path

generate_qwen25_omni*.py and test_model.py all resolve the model path via huggingface_hub.snapshot_download("Qwen/Qwen2.5-Omni-7B"), so a user who has never touched the checkpoint can run the examples directly and the weights auto-download on first use. Users with a pre-downloaded checkpoint can point QWEN25_OMNI_MODEL_PATH at a directory containing config.json to skip the HF cache entirely.

Model Information

Model Name: Qwen2.5-Omni-7B

Model Architecture: Multimodal encoder–decoder — Thinker (Qwen2 text) + Vision (SwiGLU ViT) + Audio (CPU frontend + Neuron transformer) + Talker (Qwen2-style with fused embedding) + Token2Wav (DiT + BigVGAN)

Purpose: Text generation, image-to-text, audio-to-text, and text-to-speech (full omni-modal)

HuggingFace Checkpoint: Qwen/Qwen2.5-Omni-7B

Checklist

Required Components

• Accuracy Test — contrib/models/Qwen2.5-Omni-7B/test/integration/test_model.py: 8-test suite (imports, config, state dict, audio CPU components, Talker CPU, text-only Thinker compile+generate, image understanding, audio understanding).
• README.md with:
  • Usage Example — text-only and full multimodal code snippets using the flat-import bootstrap.
  • Compatibility Matrix — Neuron SDK 2.29, PyTorch 2.9, Python 3.12, Trn2 (trn2.48xlarge).
  • Example Checkpoints — HuggingFace link + snapshot_download auto-fetch behavior.
  • Testing Instructions — how to run test_model.py and generate_qwen25_omni_speech.py against compiled artifacts.
• Source Code under contrib/models/Qwen2.5-Omni-7B/src/:
  • modeling_qwen25_omni.py — Text-only and multimodal orchestration, config, state dict conversion.
  • modeling_qwen25_omni_vision.py — Vision encoder (SwiGLU, RMSNorm, PatchMerger) on Neuron.
  • modeling_qwen25_omni_audio.py — Audio encoder with CPU frontend + Neuron transformer (chunked attention).
  • modeling_qwen25_omni_talker.py — Talker (Neuron, fused 8448→896 embedding, per-step thinker injection via get_model_output override).
  • modeling_qwen25_omni_token2wav.py — DiT transformer on Neuron + BigVGAN on CPU.
  • _model_path.py — resolve_model_path() helper; honors QWEN25_OMNI_MODEL_PATH, falls back to snapshot_download.
  • _upstream_compat.py — Runtime shim for the upstream hf_adapter.py tensor_capture_hook NameError (see companion PR Fix NameError in HuggingFaceGenerationAdapter.prepare_inputs_for_generation #136).

Optional Components

• Integration tests on Neuron — under test/integration/: test_model.py (8-test suite), test_talker_neuron.py, test_e2e_qwen25_omni.py.
• End-to-end speech pipeline — verified producing real speech audio on trn2.48xlarge; steady-state 18s/request.
• Examples — examples/generate_qwen25_omni.py (text/image/audio → text) and examples/generate_qwen25_omni_speech.py (full speech synthesis, three Neuron models resident in one process).
• Performance benchmarks — perf_test/3_bench_qwen25_omni_7b.sh (vLLM BS=1/4, TP=4) and perf_test/apply_vllm_neuron_patch_qwen25omni.py.

Folder Structure

contrib/models/Qwen2.5-Omni-7B/
  README.md
  src/
    __init__.py
    _model_path.py                   # snapshot_download-backed default path resolver
    _upstream_compat.py              # Runtime shim for hf_adapter.py NameError
    modeling_qwen25_omni.py          # Text-only + multimodal orchestration
    modeling_qwen25_omni_vision.py   # Vision encoder (Neuron, TP=4)
    modeling_qwen25_omni_audio.py    # Audio encoder (CPU + Neuron hybrid)
    modeling_qwen25_omni_talker.py   # Talker with get_model_output override
    modeling_qwen25_omni_token2wav.py # DiT (Neuron) + BigVGAN (CPU)
  examples/
    generate_qwen25_omni.py          # text / image / audio / full
    generate_qwen25_omni_speech.py   # end-to-end speech (single-process, core 0-3 pinned)
  test/
    integration/
      test_model.py                  # 8-test suite
      test_talker_neuron.py          # Talker-specific Neuron tests
      test_e2e_qwen25_omni.py        # CPU reference end-to-end
  perf_test/
    3_bench_qwen25_omni_7b.sh
    apply_vllm_neuron_patch_qwen25omni.py

Testing

How did you test this change?

All 8 tests pass on trn2.48xlarge with real Qwen2.5-Omni-7B weights (Neuron SDK 2.29, PyTorch 2.9, Python 3.12):

1. Imports — all module groups import correctly through the contrib flat-import bootstrap.
2. Config — TP=4 head divisibility verified (Thinker 7Q/1KV, Audio 5, Vision 4 per rank).
3. State dict — all 2448 keys convert correctly (text=339, audio=489, vision=518, talker=293, token2wav=809).
4. Audio CPU — frontend + postprocessor latency 20–34ms for 1–30s audio.
5. Talker Neuron — compile + load + generate codec tokens on Neuron.
6. Text generation — compile + load + generate on Neuron, outputs match HF reference.
7. Image understanding — vision encoder + Thinker, correct captioning.
8. Full speech pipeline — Thinker → Talker → Token2Wav all on Neuron, producing real speech audio.

End-to-end speech pipeline benchmark (generate_qwen25_omni_speech.py --num-runs 5 on trn2.48xlarge):

Phase	Avg / run
Thinker	0.34s
HF CPU hidden states	1.16s
Talker input prep	0.18s
Talker	2.28s
Token2Wav	14.14s
Pipeline	18.10s

Wall time for 5 runs: 129s (load + 5×18s). Audio duration: 12s; RTF 1.51×.

Verified behaviors:

• The _upstream_compat shim successfully patches HuggingFaceGenerationAdapter.prepare_inputs_for_generation at import time.
• NeuronTalkerModel.get_model_output override replaces the old apply_vision_during_token_gen class attribute (which used to require a model_base.py edit).
• All three Neuron-compiled models co-reside on NeuronCores 0–3 when NEURON_RT_VISIBLE_CORES=0-3 is set (verified via neuron-top: NC0/NC1/NC3 at ~5.6 GB each for the TP=4 Thinker+Talker shards, NC2 at ~8.2 GB with the extra DiT NEFF).

Compatibility

Tested with:

• Neuron SDK Version(s): 2.29
• Instance Type(s): Trn2 (trn2.48xlarge — primary; trn2.3xlarge supports the same 4-core TP=4 configuration but was not benchmarked in this PR).
• PyTorch Version: 2.9
• Python Version: 3.12

Additional Information

Known limitations:

• Token2Wav BigVGAN remains on CPU (6 upsample stages); only the DiT transformer core is compiled to Neuron. This dominates Token2Wav's 14s per-run time and is the largest remaining optimization target.
• _upstream_compat.py is a stopgap for an hf_adapter.py bug addressed by companion PR Fix NameError in HuggingFaceGenerationAdapter.prepare_inputs_for_generation #136; once that merges the shim can be deleted.
• Examples and tests resolve paths via Path(__file__).resolve().parents[N], so the contrib directory layout must be preserved.

Key technical contributions:

1. NeuronTalkerModel.get_model_output override replaces the previous apply_vision_during_token_gen flag-on-model_base approach, keeping the change fully inside contrib.
2. Fused embedding (8448→3584→896 collapsed to 8448→896) eliminates the thinker-to-talker projection at inference time.
3. Vision embeddings auto-pad to max_context_length buckets in set_vision_embeddings() for bucket compatibility.
4. Token2Wav CPU fallback correctly orders the mel_len overflow check before input_embed doubles the batch for classifier-free guidance.
5. Three-model single-process pipeline with NeuronCore 0-3 pinning (NEURON_RT_VISIBLE_CORES=0-3) + Phase 2 loaded in bf16 — the combination keeps per-request latency at ~18s even after the cold start, vs. ~90s when each component was in its own subprocess.

Related Issues / PRs

• Supersedes Add Qwen2.5-Omni-7B full Neuron speech pipeline (Thinker+Talker+Token2Wav, TP=4) #122.
• Companion bug fix: Fix NameError in HuggingFaceGenerationAdapter.prepare_inputs_for_generation #136 (1-line fix for the HuggingFaceGenerationAdapter.prepare_inputs_for_generation NameError).

vLLM Integration

• This model is intended for use with vLLM (text-only).
• vLLM-neuron patch included at perf_test/apply_vllm_neuron_patch_qwen25omni.py.

---

By submitting this PR, I confirm that:

• I have read and followed the contributing guidelines.
• This is a community contribution and may have limited testing compared to officially-supported models.
• The code follows best practices and is well-documented.
• All required components listed above are included.