diffusers

7e46c94c - [feat] Tensor parallelism for MiniMax-H3 (#14609)

[feat] Tensor parallelism for MiniMax-H3 (#14609) * Add tensor-parallel support for MiniMax-H3 Shards `MiniMaxH3Transformer3DModel` across devices, following the plan already established for Flux1/Flux2/Qwen-Image. Validated on Trainium at TP=2 and TP=8. - `_tp_plan` with twelve entries: the same six shapes for the 50 denoiser blocks and for the two token-refiner blocks, which are the same attention + SwiGLU FFN minus AdaLN and rotary. Q/K/V and the attention output are unfused, so they are plain colwise/rowwise; the SwiGLU input `ff.net.0.proj` is one Linear producing `[value; gate]` in equal halves and takes PackedColwiseParallel([1, 1]). - The attention processor reshaped by the config head count, `unflatten(-1, (attn.heads, -1))`, which mis-splits under sharding: each rank holds `inner_dim / tp_degree` columns, so this yields `head_dim / tp_degree` per head instead of `heads / tp_degree` heads. Reshape by the fixed `attn.head_dim` instead and let `-1` absorb the head count, as Flux does. Numerically identical unsharded, since `inner_dim == heads * head_dim`. - Norms, QK-norms (head_dim-shaped, applied after the head split), AdaLN modulation and the patch/text embedders and output heads stay replicated. `attn.to_qkv` is deliberately not in the plan: it exists only after `fuse_projections()`, and the plan is resolved by attribute lookup. No RoPE change was needed — unlike Qwen-Image, H3's rotary is already real sin/cos and already broadcasts over the head axis. Tests mirror the Flux2/Qwen-Image layout: the CUDA/XPU `TensorParallelTesterMixin` class, a `make_neuron_tp_spec()` factory, and a Neuron launcher that shells out to the model-agnostic `_neuron_tp_worker.py`. `get_dummy_inputs` and `get_packed_layout` take an optional `device` so the Neuron spec can ask for CPU tensors, since its worker shards on CPU and moves to device after. * Shard MiniMax-H3's adaln_proj to fit tensor parallelism on one device `transformer_blocks.*.adaln_proj.linear` was left replicated on every rank, and at `[96768, 2688]` bf16 per block it is 24.23 GiB of the denoiser's 61.73 GiB — about 40%. That made the per-rank floor 24.40 GiB of weights (plus 5.13 GiB for the two VAEs) regardless of TP degree, so MiniMax-H3 could not fit a 24 GiB NeuronCore at *any* valid TP: 34.20 GiB/rank at TP=8, and still 30.20 GiB/rank at TP=56. Raising TP only divided the 60% that already sharded. (TP=16 is not an option either — 56 attention heads.) Shard it rowwise, over the `time_embed_dim` input, rather than colwise: the six modulation parameters scale and shift the *full* hidden dim of a sequence that is already all-reduced by the time they are applied, so a colwise split would need an all-gather to rebuild that width. Rowwise keeps the output full-width, leaving the module's `view`/`chunk` untouched, and all-reduces a few hundred KB per block per step. Plain `"rowwise"` could not be reused. It is normally the second half of a colwise/rowwise pair, so it defaults to `input_layouts=Shard(-1)` and would read the full-width `temb` as if it were one rank's shard. Hence `ReplicatedInputRowwiseParallel`: input narrowed locally on the way in (no collective), partial output all-reduced on the way out, bias replicated and added after the reduce. It is wired into `_styles`, `_hooks_only_styles` — the path the Neuron backend takes, since `_apply_tp_neuron` pre-shards on CPU and then registers hooks only — and `resolve_tp_shard_specs`. Replicated weights drop from 24.40 GiB to 0.15 GiB, putting TP=8 at 7.84 GiB of transformer plus 5.13 GiB of VAEs, i.e. 12.97 GiB/rank against a 24 GiB budget. Verified on CPU/gloo that both the generic and the pre-sharded hooks-only path shard the weight on its input dim and match a replicated reference to 3.6e-7. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Build MiniMax-H3's row timestep plan on CPU, as its caller expects `build_row_timesteps` allocates `row_timesteps` with `torch.full` and no device, then scatters into it at `video_indices` / `audio_indices`. The layout step hands those index tensors over already on the execution device, and indexing a CPU tensor with an accelerator one is an error — on Neuron it surfaces as "Non-scalar tensor arg0 is on cpu device, expected neuron", and on CUDA it would raise "indices should be either on cpu or on the same device". CPU is the right place for this to run, not the accelerator: `torch.unique` has a data-dependent output shape, which is precisely what a tracing backend cannot handle, and the caller already moves the finished `(timestep, timestep_indices)` pair to the device itself. So bring the two index tensors back to CPU for the scatter rather than allocating `row_timesteps` on their device. Only reachable once the denoiser is actually on an accelerator while the pipeline's execution device resolves there too, which is why it went unnoticed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * feat:combine tp+cp (cherry picked from commit a489f03e4e9888c0abb996f45d38885f03772ae4) * Allow tensor parallelism and context parallelism in one ParallelConfig `ParallelConfig.__post_init__` rejects a config carrying both a `context_parallel_config` and a `tensor_parallel_config`, so a model can only use one form of parallelism at a time. That caps some models well below the hardware they run on: MiniMax-H3 has 56 attention heads, so `tp_degree` cannot exceed 8, and on a 64-accelerator host tensor parallelism alone leaves 56 of them idle even though the model already ships a `_cp_plan`. Both configs already accept a caller-supplied `mesh`, documented as being for "combining CP with other parallelism strategies that share the same mesh", so the intent was there; what was missing was building that shared mesh and splitting it between the two. * `__post_init__` now only rejects a config that carries neither. * `ParallelConfig.setup` hands each config its own dimensions: CP keeps the whole mesh, since its own `setup` selects "ring" and "ulysses" out of it, while TP gets the "tp" dimension alone because it reads its degree from `mesh.size()`. * `enable_parallelism` builds one ("tp", "ring", "ulysses") mesh when both are requested, with the context-parallel dimensions varying fastest. That order is a default rather than a universal optimum, and the comment says so: pass `mesh=` on either config to choose the layout yourself. * The Neuron tensor-parallel backend used the global rank as its shard index. That equals the rank's coordinate on the TP mesh only while that mesh spans the whole world, so sharing a mesh made it index past the end of the weight; the generic backend beside it already reads the coordinate. Measured on a trn2.48xlarge (MiniMax-H3, 1344x768x124f, 30 steps): TP=8 alone is 9.285 s/step on 8 cores, TP=4 x ulysses=4 is 5.66 s/step on 16, and TP=8 x ulysses=4 is 3.941 s/step on 32 -- 2.36x faster than the widest configuration reachable before this change, with output that is visually equivalent to the tensor-parallel-only run. Tests: a `HybridParallelTesterMixin` next to the existing TP and CP mixins, wired into the Flux transformer tests, plus a Neuron `torchrun` worker following the `_neuron_tp_worker.py` convention already in the tree. The Neuron test runs at tp_degree=2 x ulysses_degree=4 and passes on a trn2.48xlarge (max_abs_diff 4.2e-05 against the single-device reference). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * feat: validate the support on TPU * fix: put tp first in the combined TP+CP mesh * fix: style * fix: apply context parallelism on the sharded-load path * revert: attention head padding * revert: drop the Neuron all-gather padding and restore the TPU pre-shard path * refactor: share the CPU pre-shard path between Neuron and TPU tensor parallelism * fix: new device count API for TPU since deprecation * fix: pad uneven all-gathers on Neuron and TPU in _all_gather_anything * fix: fall back to one device when no accelerator is present Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * Move TP+CP composition to #14725 Keep this PR to tensor parallelism for MiniMax-H3. Combining TP and CP in one ParallelConfig, with its tests, now lives in #14725. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * review: use `config._device` for the Neuron TP device `config._device` now resolves to the local Neuron core on every rank, so the `torch.neuron.current_device()` special case is no longer needed when pre-sharding, loading sharded checkpoints and moving buffers. * review: replace ReplicatedInputRowwiseParallel with a "rowwise_split_input" style The marker class carried no state, so make it a plain string style next to "colwise" and "rowwise". The name matches the transformers style of the same name, which is the same RowwiseParallel(input_layouts=Replicate()). Add it to the TP styles table in the distributed inference docs, with the AdaLN modulation projection as the example. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * review: explain when to use "rowwise_split_input" Replace the comment on how torch narrows the input with when the style is needed, plus two examples: H3's AdaLN modulation projection and an output projection after a gathered QK-norm. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * feat:shard tensor-parallel weights locally on every device * doc: improve the explanation of rowwise_split_input * review: remove useless zero division prevention * review: simplify the docstring * revert: Restore standard TP sharding and isolate the Neuron workaround * revert: no refactoring of in-memory sharding * revert: keep the in-memory TP path as on main Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * style: pika * style: make fix-copies --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: whn09 <whn09@users.noreply.github.com> Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>