[feat] Tensor parallelism for MiniMax-H3 - #14609
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. |
|
Tried this branch ( Setup. Numbers. 30 model evaluations in 354.1s. The first evaluation is 71s (kernel compile/warmup) and steady state is 9.30 s/step. For scale: our hand-written NeuronCore plugin for H3 does the same canvas at 2.90 s/step on 32 cores with CP=4, i.e. 92.8 core-seconds/step, against 74.4 core-seconds/step here. So the native path is not costing per-core throughput at all — it comes out ahead, which is expected since CP=4 has scaling loss that TP=8 doesn't. Its only limit is that 8 of the 64 cores are usable (see the last section). Bug: non-persistent buffers stay on CPU in the safetensors TP path
The old Suggested fix — hoist your existing loop into a helper so both sharded-load paths share it. This also folds the three copies of the Neuron device resolution into one: # src/diffusers/hooks/tensor_parallel.py (next to resolve_tp_shard_specs)
def tp_device(tp_config: TensorParallelConfig) -> torch.device:
"""The device this rank's shards belong on.
`TensorParallelConfig._device` is derived from the default accelerator, which is not meaningful on
Neuron: the core a rank owns is only known after `init_process_group`.
"""
if tp_config._mesh.device_type == "neuron":
return torch.neuron.current_device()
return tp_config._device
def move_non_persistent_buffers(model: torch.nn.Module, tp_config: TensorParallelConfig) -> None:
"""Place the buffers that no loading path can have placed.
Non-persistent buffers are absent from both the state dict and the checkpoint, so neither the
per-rank safetensors reader nor `dcp.load` ever sees them, and `init_empty_weights` leaves them
as real CPU tensors.
"""
from torch.distributed.tensor import DTensor
device = tp_device(tp_config)
for name, buffer in model.named_buffers():
if buffer.device != device and not isinstance(buffer, DTensor):
module_path, _, buffer_name = name.rpartition(".")
module = model.get_submodule(module_path) if module_path else model
module._buffers[buffer_name] = buffer.to(device)--- a/src/diffusers/models/modeling_utils.py
+++ b/src/diffusers/models/modeling_utils.py
@@ -1653,8 +1653,9 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
if tp_shard_specs is not None:
# The weights are already sharded, so this only registers the forward hooks. `_parallel_config`
# was recorded by `_resolve_parallel_config` before loading.
- from ..hooks.tensor_parallel import apply_tensor_parallel
+ from ..hooks.tensor_parallel import apply_tensor_parallel, move_non_persistent_buffers
+ move_non_persistent_buffers(model, tp_config)
apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True)--- a/src/diffusers/models/model_loading_utils.py
+++ b/src/diffusers/models/model_loading_utils.py
@@ -411,15 +411,10 @@ def _load_shard_file_tp(
- from ..hooks.tensor_parallel import _local_shard
+ from ..hooks.tensor_parallel import _local_shard, tp_device
tp_mesh = tp_config._mesh
- # `TensorParallelConfig._device` is derived from the default accelerator, which is not meaningful on
- # Neuron; resolve it the way the Neuron pre-shard backend does.
- if tp_mesh.device_type == "neuron":
- device = torch.neuron.current_device()
- else:
- device = tp_config._device
+ device = tp_device(tp_config)and in I ran this patch on the box: the explicit buffer sweep I had been carrying in my own script no longer finds anything to move, and 768p completes unchanged. Happy to open it as a PR against your branch if that's easier. Where the ceiling is, and a possible next step. FYI for anyone else on Neuron eager (not a diffusers issue). |
|
@JingyaHuang two things stacked on this branch, both verified on a
The second one matters specifically for your use case: at 33B parameters the The PR includes a Neuron test following your This is independent of the non-persistent-buffer fix I left above, which is on the same loading path. |
3 participants
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.