feat(examples): Qwen-Image 20B MMDiT text-to-image on Trainium

[ymwangg]

Summary

Adds a device-only, TP-mandatory Qwen-Image 20B MMDiT text-to-image pipeline to examples/models/qwen_image/. The full pipeline runs on trn2 (TP=4):

• Denoiser — 60-block dual-stream MMDiT (3D RoPE, QK-RMSNorm, per-stream modulation) with fused on-device CFG + FlowMatchEuler sampling; the packed latent stays resident across all steps and the host feeds only per-step scalars.
• Text encoder — prefill Qwen2.5-VL decoder LM (28 layers, GQA 28/4, SwiGLU), reusing the qwen3 block kernels, Megatron-sharded. Device rel_l2 6.7e-3 bf16.
• VAE decoder — fp32 one-shot; at T=1 the WAN-style 3D causal video VAE collapses exactly to a 2D conv decoder. Device rel_l2 4.4e-6 at 64×64 → 512×512.
• Host keeps only trivial glue (tokenizer/chat template, embedding lookup, latent pack/denorm, scalar flow-match schedule).

Produces correct 512px images. End-to-end demo.sh reuses the compile cache, runs CPU correctness tests, then generates via torchrun.

Sample output

512px, TP=4, 50 steps, guidance 4.0 — prompt: "a coffee shop entrance with a chalkboard sign".

Notable details

• RoPE perf fix: interleave-scatter rewrite drops the denoiser from 14.1s → 0.30s/step. Qwen uses the interleaved (2i, 2i+1) pair convention, not qwen3's (i, i+half) split.
• VAE denorm: the pipeline stores latents_std as its reciprocal (latents / std_recip); dividing by raw std instead is off by ~std².
• 1024px limit: the 60-block denoiser unrolls into one HLO graph; at grid 64×64 the serialized proto exceeds protobuf's 2 GB limit. Works today at ≤512px; durable fix is a device-side scan over one compiled block.

Startup & serving perf

A cold launch was dominated by host-side weight prep (~350s extract/shard), not compile or generation. Three additions cut that:

• Shard cache (--weight-cache-dir, default on): extract+shard once, then reload per-rank safetensors shards on later launches (bf16 = what uploads; fp32 for the VAE). ~350s → memory-mapped read; peak RAM ~200GB → ~45GB. --no-weight-cache disables.
• Resident mode (--prompts-file / --interactive): keep the process and all TP ranks alive, generating one image per prompt so the weight upload (~25s) + NEFF load (~17s) that no disk cache can remove are paid once — additional images then cost only generate() (~2.8s at 8 steps).
• Text-length bucketing (--text-bucket, default 64): round the encoder seq and denoiser text length up to a bucket so varying-length prompts reuse one compiled kernel per bucket instead of recompiling per exact length. Exact (causal encoder + masked pad text tokens); --text-bucket 0 restores per-exact-length, bitwise-stable behavior.

Test plan

• CPU correctness tests: uv run pytest examples/models/qwen_image/tests/
• On-device TP check (opt-in, needs hardware): QWEN_IMAGE_TP_DEVICE_TEST=1 uv run pytest tests/test_tp_device.py
• Validated end-to-end on trn2 at TP=4, 512px.

See examples/models/qwen_image/README.md for full run instructions and the durable lessons/limits.