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

[whn09]

Description

Add NxDI model adapter for Qwen2.5-Omni-7B with full multimodal support and end-to-end Neuron speech synthesis: text generation, image understanding, audio understanding, and text-to-speech — all on Neuron.

All components run on Neuron at TP=4: Thinker (28 heads → 7/rank), Vision (16 → 4/rank), Audio (20 → 5/rank), Talker (12 heads → 3/rank), Token2Wav DiT (22 blocks, traced). This enables running on trn2.xlarge (4 NeuronCores) as well as larger instances.

Full Neuron Speech Pipeline

The complete text-to-speech pipeline runs entirely on Neuron hardware:

Component	Time	Notes
Thinker (7B, Neuron TP=4)	0.3s	24 text tokens
Talker (690M, Neuron TP=4)	2.1s	454 codec tokens, per-step thinker injection
Token2Wav (Neuron DiT + CPU BigVGAN)	9.9s	22 DiT blocks × 10 ODE steps
Total	~15s	9.1s audio, RTF ~1.7x

Compared to 267.9s on CPU (~18x overall speedup).

Key Technical Contributions

1. Talker per-step thinker state injection: At each autoregressive step, the corresponding thinker_reply_part[step] embedding is added to the codec token embedding via the vision_embeddings mechanism, matching HF's generation behavior.

2. Fused embedding: Collapsed embed_tokens (8448→3584) + thinker_to_talker_proj (3584→896) into a single 8448→896 lookup, eliminating the projection layer during inference.

3. Vision embeddings auto-padding: Compiled Neuron models require fixed bucket shapes. Vision embeddings are auto-padded to max_context_length in set_vision_embeddings() for bucket compatibility.

4. Token2Wav CPU fallback: Moved mel_len overflow check before input_embed (which doubles batch for classifier-free guidance) to prevent batch dimension mismatch on CPU fallback.

5. model_base extension: Added apply_vision_during_token_gen flag to allow vision_embeddings during token generation (not just context encoding), enabling the Talker's per-step injection pattern.

Model Information

Model Name: Qwen2.5-Omni-7B

Model Architecture: Multimodal encoder-decoder (Thinker + Vision + Audio + Talker + Token2Wav)

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

Architecture

Component	Runtime	TP	Key dims
Thinker (text)	Neuron	4	hidden=3584, heads=28, kv_heads=4, layers=28
Vision encoder	Neuron	4	embed=1280, heads=16, depth=32, SwiGLU
Audio encoder	CPU+Neuron	4	d_model=1280, heads=20, layers=32, chunked attn
Talker	Neuron	4	hidden=896, heads=12, kv_heads=4, layers=24, vocab=8448, fused embed
Token2Wav	CPU+Neuron (fp32)	N/A	DiT: dim=1024, 22 blocks (Neuron); BigVGAN: 6 upsample stages (CPU)

Checklist

Required Components

• Accuracy Test — All 6 tests pass on trn2.48xlarge with real Qwen2.5-Omni-7B weights
• README.md — Full multimodal documentation with architecture details, usage examples, vLLM serving instructions, and performance benchmarks
• Source Code (src/)
  • modeling_qwen25_omni.py — Text-only + multimodal orchestration, config handling, state dict conversion
  • modeling_qwen25_omni_vision.py — Vision encoder (SwiGLU, RMSNorm, separate QKV, PatchMerger)
  • modeling_qwen25_omni_audio.py — Audio encoder (hybrid CPU+Neuron, chunked attention)
  • modeling_qwen25_omni_talker.py — Talker (Neuron, fused embed, per-step thinker injection, auto-padding)
  • modeling_qwen25_omni_token2wav.py — Token2Wav (Neuron DiT + CPU BigVGAN, CPU fallback)

Optional Components

• Unit Tests — Integration tests on Trn2 with real weights
• End-to-end speech pipeline — Verified producing real human speech audio

Folder Structure

src/neuronx_distributed_inference/models/qwen25_omni/
  __init__.py
  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, TP=4)
  modeling_qwen25_omni_talker.py   # Talker (Neuron, TP=4, per-step injection)
  modeling_qwen25_omni_token2wav.py # Token2Wav (Neuron DiT + CPU BigVGAN)

contrib/models/Qwen2.5-Omni-7B/
  README.md                        # Full multimodal documentation
  test/integration/test_model.py   # 6 integration tests

perf_test/
  3_bench_qwen25_omni_7b.sh       # vLLM benchmark (BS=1/4, TP=4)
  apply_vllm_neuron_patch_qwen25omni.py  # vllm-neuron patch script

Testing

How did you test this change?

All 6 tests pass on trn2.48xlarge with Qwen2.5-Omni-7B weights (Neuron SDK 2.23, PyTorch 2.9):

1. Imports: All 7 module groups import correctly
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 (1-30s audio)
5. Talker Neuron: Compile + load + generate codec tokens on Neuron, TPOT 4.0ms
6. Text generation: Compile + load + generate on Neuron, TPOT ~11-13ms, correct outputs
7. Full speech pipeline: Thinker → Talker → Token2Wav all on Neuron, producing real speech audio

Compatibility

Tested with:

• Neuron SDK Version(s): 2.23
• Instance Type(s): Trn2 (trn2.48xlarge), should work on trn2.xlarge (4 cores)
• PyTorch Version: 2.9
• Python Version: 3.12

Related Issues

N/A

vLLM Integration

• This model/feature is intended for use with vLLM (text-only)
• vLLM-neuron patch included in 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

[whn09]

Closing and superseding this PR with a clean zero-invasion rewrite.

The original submission modified files under src/neuronx_distributed_inference/ (model registrations in utils/constants.py and inference_demo.py, plus small hooks in models/model_base.py and utils/hf_adapter.py), which is not friendly for a contrib merge. The new PR refactors Qwen2.5-Omni-7B to a 100% zero-invasion contrib layout:

• All model code lives under contrib/models/Qwen2.5-Omni-7B/{src,examples,test,perf_test}/
• git diff upstream/main..HEAD -- src/ is empty
• The Talker's previous need for an apply_vision_during_token_gen gate in model_base.py is handled by a get_model_output override on the contrib Talker subclass
• One small upstream bug in hf_adapter.py (undefined tensor_capture_hook reference) is worked around by an in-contrib runtime shim, and will be fixed via a separate 1-line upstream PR

New PR: will link once opened.