[shiyang-opensource]

Summary

MiniMax-H3 (T2VA, 33B DiT + Qwen3-VL conditioner + video/audio VAEs) on the PyTorch-native neuron backend (TorchNeuron, torch-neuronx 2.14 / SDK 2.33). The AOT trainium backend cannot move to SDK 2.33: neuronx-cc 2.28 removed neuronxcc.nki, which NxD 0.19 / NxDI 0.10 import. With this PR H3 runs on the current SDK without NxD/NxDI.

With DIFFLET_BACKEND=neuron (or --backend neuron), difflet generate runs H3 in one torchrun launch of --tp-degree ranks:

1. Text: Qwen3-VL's language model to hidden_states[50], TP-sharded (NeuronMiniMaxH3TextEncoderApplication; first 50 layers only, 11.7 GiB/rank at TP4). Released before the DiT loads.
2. DiT: NeuronMiniMaxH3TransformerApplication, per-block compiled. The AdaLN modulation is precomputed on the host, so the 13B adaln_proj branch is never loaded (9.4 GiB/rank).
3. Video VAE: one decoder per rank (NeuronMiniMaxH3VideoVaeApplication, fp32 by default). The official 256 px tiles are spread over the ranks (decode_video_on_ranks).

The audio VAE and the AV mux stay on the host (audio_vae stage). --host-text-encoder and --host-vae bring back the host text and video_vae stages.

Commits

1. refactor(minimax-h3): the DiT TP rewrite moves to models/minimax_h3/modeling_minimax_h3_tp.py. The Trainium module re-exports every name, so its graph is built from the same objects. This is the only change to the Trainium path, plus --model-path now being forwarded to every stage on both backends.
2. feat(neuron): the H3 DiT as a TorchNeuron application.
   • Fused SwiGLU via CheckpointSlice.
   • fp32 patch projections, timestep MLP and output heads on fp32 storage.
   • int32 masks.
3. feat: host CPU text encoder and decoders (host_components.py), used by the --host-* paths.
4. feat(neuron): the staged CLI on the torchrun backend.
   • The registry lists neuron for minimax_h3.
   • The c10 refusal tests now restrict H3's entry to trainium themselves, because no registered model lacks neuron any more.
5. feat(neuron): the video VAE on a NeuronCore.
6. feat(neuron): the Qwen3-VL conditioner on the device.
7. feat(neuron): text, DiT and video decode in one generate launch.

Evidence (trn3.48xlarge, device 7 = cores 28-31, TP4, 256x448x124, 30 steps, seed 42)

	AOT trainium (SDK 2.26 stack)	neuron (this PR)
DiT step (median of 20, same device)	782.6 ms	714.7 ms
difflet generate end to end	~150 s (generate only)	109 s
text / DiT 30 steps / video decode / audio	25 / ~29 / 18 / 53 s	8.9 / 21.1 / 6.6 / 2.9 s

Trainium path unchanged: after the TP-module move, the AOT DiT was recompiled from scratch (fresh NEFF cache, 438 s) and generated from the same text.pt. The video and audio latents are bit-identical to the pre-refactor run (max_abs 0.0).

Numerics vs references:

• DiT: tp1/tp2/tp4, eager and compile, matches diffusers' MiniMaxH3Transformer3DModel (fp32, atol 2e-5), with the host AdaLN table and with the resident branch.
• Text encoder: tp1/2/4 matches transformers' Qwen3VLForConditionalGeneration (fp32, 2e-5). On trn3, vs the AOT/NxDI encode: rel_l2 9.8e-4 (bf16); mask and token count equal.
• Video VAE: device fp32 vs the host fp32 decode is 80 dB per clip; the rank-spread decode is equal to the one-process decode. bf16 (--vae-dtype bf16) is 57 dB and 5x faster.
• End to end vs AOT: same scene (frames compared). Audio mean -58.6 vs -58.0 dB. The latents are not expected to be bit-equal across compilers and kernels after 30 bf16 steps.

Not done here:

• The audio VAE stays on the host. Device fp32 hits a neuronx-cc 2.28 internal error on BigVGAN's depthwise conv1d (NCC_ITIN902); bf16 runs but gives 13 dB SNR and no speed-up over the host's 2.9 s.
• Top-level difflet generate has no --model-path yet; this would be a shared cli/main.py change.

Test plan

• tests/unit/backends/test_neuron_minimax_h3_{l1,text_encoder,vae}.py (CPU/gloo, run on the trn3 host; the Mac torch 2.11 lacks all_gather_single)
• tests/unit/cli/test_orchestrator_minimax_h3.py, tests/unit/backends/test_neuron_launch_c10.py (109 passed), registry tests
• tests/unit/models/minimax_h3 in the Trainium lock venv (32 passed, including the re-export identity test)
• On-device: difflet compile + difflet generate (neuron), AOT recompile bit-identity (trainium)

🤖 Generated with Claude Code