[feat] Tensor parallelism for MiniMax-H3 - #14609
Draft
JingyaHuang wants to merge 8 commits into
Draft
Conversation
Stream each rank's slice of a tensor-parallel checkpoint straight off disk instead of materializing the full checkpoint on every rank and resharding it afterwards, and gather the shards back on save. - `from_pretrained(..., parallel_config=TensorParallelConfig(...))` resolves the shard specs on the still-meta model, then slices each safetensors tensor before the dtype cast, so host memory peaks at ~1/tp_degree of the checkpoint. - `save_pretrained` all-gathers the DTensors into an ordinary checkpoint, or writes a distributed checkpoint with `dcp=True` so no full tensor is ever formed. The writing `tp_degree` is recorded, since a packed weight's stored layout is interleaved by it. - Factor the plan interpretation out of the Neuron pre-shard path into shared `TPShardSpec` / `resolve_tp_shard_specs` / `_local_shard` / `_hooks_only_styles` helpers, so both backends and both the load and save paths shard identically.
…ng or LoRA Addresses the remaining two items of the review on huggingface#13718: tensor parallelism was rejected alongside quantization and `device_map` only on the `from_pretrained` streaming path, while `enable_parallelism` — which the quantization error message itself recommended — accepted a quantized, offloaded or adapter-injected model and sharded it anyway. - Add `_check_tp_model_state`, called from `apply_tensor_parallel`, the one chokepoint every TP entry point funnels through. It rejects a model that is quantized, group-offloaded, placed by accelerate (`device_map` or CPU offload), or has PEFT layers injected. Placed before the device-type check so the reported reason is the useful one. - Guard the reverse order too: `enable_group_offload`, the two pipeline CPU-offload methods, and `load_lora_adapter` now refuse a tensor-parallel model. - `save_pretrained` refuses a quantized tensor-parallel model. Previously the `dcp=True` branch returned before the quantizer's serialization step, writing shards with no quantization metadata and no error. - The DCP load guard checked the `quantization_config` kwarg only, so a pre-quantized checkpoint directory loaded silently; check the config's own entry too, and add the missing `_tp_plan` check that otherwise surfaced as a raw `AttributeError`. - Correct the `from_pretrained` message and the doc sentence that pointed at `enable_parallelism` as a way to shard a quantized model. The new tests are the first tensor-parallel tests that need neither an accelerator nor more than one rank: every case asserts a raise before any collective, so they run single-process on gloo.
…sers into add-shard-ckpt-loading
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.
`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_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>
github-actions
Bot
added
size/L
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Shards
MiniMaxH3Transformer3DModelacross a TP mesh.Before submitting
self-reviewskill on the diff?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.