Add Qwen2.5-Omni-7B multimodal support (TP=4, verified on Trn2)

[whn09]

Description

Add NxDI model adapter for Qwen2.5-Omni-7B with full multimodal support: text-only Thinker, vision encoder, audio encoder, Talker (speech codec generation), and Token2Wav (waveform synthesis).

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

The Thinker's text backbone is architecturally identical to Qwen2.5-7B, so the text model reuses NeuronQwen2ForCausalLM with state-dict prefix remapping. The vision encoder uses SwiGLU MLP, RMSNorm, and separate QKV projections, compiled and run on Neuron. The audio encoder uses a hybrid CPU+Neuron architecture: Conv1d frontend + chunking on CPU, 32 transformer layers on Neuron, AvgPool + projection on CPU. The Talker stays on CPU (non-standard head_dim, 3D mRoPE, per-step thinker injection). Token2Wav runs on CPU in float32.

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)

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 (HF wrapper, codec token generation, CPU)
  • modeling_qwen25_omni_token2wav.py — Token2Wav (DiT + BigVGAN, waveform synthesis, CPU)

Optional Components

• Unit Tests — Integration tests on Trn2 with real weights (see Testing section)

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 codec generation (CPU)
  modeling_qwen25_omni_token2wav.py # Token2Wav waveform synthesis (CPU)

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 CPU: 293 keys loaded, codec tokens verified (bos=8293, eos=8294, pad=8292)
6. Text generation: Compile + load + generate on Neuron, TPOT ~11-13ms, correct outputs

Generation Results (greedy, TP=4):

Prompt	Tokens	Time	TPOT
"What is 2+3?"	1	0.10s	—
"Write a haiku about the ocean."	20	0.26s	~12.8ms
"Explain quantum computing in one sentence."	30	0.34s	~11.5ms

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

Additional Information

Architecture details:

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	CPU	N/A	hidden=896, heads=12, kv_heads=4, layers=24, vocab=8448
Token2Wav	CPU (fp32)	N/A	DiT: dim=1024, 22 blocks; BigVGAN: 6 upsample stages

Design decisions:

• TP=4: All attention head counts (28, 16, 20) are divisible by 4, enabling single trn2.xlarge
• Audio hybrid CPU+Neuron: Conv1d frontend + chunking on CPU (variable length), transformer on Neuron (TP=4), postprocessor on CPU
• Talker on CPU: Non-standard head_dim (128 ≠ 896/12), 3D mRoPE with per-step thinker-state injection, only ~690M params
• Token2Wav on CPU: ODE sampling (Runge-Kutta 4) requires float32 precision

Related Issues

N/A

vLLM Integration

• This model/feature is intended for use with vLLM
• 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