[feat] Tensor parallelism for MiniMax-H3 by JingyaHuang · Pull Request #14609 · huggingface/diffusers

JingyaHuang · 2026-08-26T12:18:43Z

What does this PR do?

Shards MiniMaxH3Transformer3DModel across a TP mesh.

Adds a _tp_plan for the transformer blocks and the token refiner, and reshapes attention by head_dim so the processor works on a rank's local heads.
Shards adaln_proj with a new ReplicatedInputRowwiseParallel style. It is ~40% of the weights, and left replicated it puts the per-rank memory above a 24 GiB core at any TP degree.
Builds the row timestep plan on CPU, as its caller expects.
Shares the CPU pre-shard path between Neuron and TPU, and moves non-persistent buffers to the device on the sharded from_pretrained path.

Combining TP with context parallelism is split out to #14725.

Validation

Mode	Neuron	TPU	CUDA
Eager	✅	✅
Compile	✅	✅

On Neuron (trn2), TP=8 eager generates 768×1344, 124 frames: peak 15.31 GiB per core. This needs the packed sequence to be a multiple of 512 so the NKI flash attention kernel is used.

Before submitting

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.

HuggingFaceDocBuilderDev · 2026-08-26T12:24:57Z

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.

Pages changed in this PR:

en/training/distributed_inference

whn09 · 2026-09-06T14:57:38Z

Tried this branch (aa50ad3) on a trn2.48xlarge (Trainium2, 16 devices × 4 logical cores, 24 GiB per core) with MiniMax-H3, and got 768p (1344×768, 124 frames) running end-to-end on the native Neuron eager path. Thanks — the TP plan works as-is. One real bug, plus numbers and a possible next step.

Setup. tp_degree=8 (H3 has 56 attention heads, so the divisors are 1/2/4/7/8; tp_degree=2 OOMs at ~31 GiB/rank). The Qwen3-VL conditioner (33.4B / 62 GiB) runs on the host CPU as its own pipeline and the VAEs are decoded in a separate process, so only the 33.1B denoiser is on device: 4.21B params/rank = 7.84 GiB, memory_allocated 8.008 GiB.

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

_load_dcp_checkpoint handles this correctly already (modeling_utils.py:1902-1908), but from_pretrained(..., parallel_config=...) against a plain safetensors checkpoint goes through _load_shard_file_tp (model_loading_utils.py:391), which only ever iterates checkpoint keys. All 638 parameters land on the right core; H3's rotary inv_freq (persistent=False, so absent from both the state dict and the checkpoint) is left on CPU, and the first thing forward does is multiply position_ids by it:

RuntimeError: input tensor is on cpu device, expected neuron
  transformer_minimax_h3.py:97

The old enable_parallelism + .to(device) route was fine because .to moves buffers too.

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 _load_dcp_checkpoint, replace the inline loop and the inline device resolution with move_non_persistent_buffers(model, tp_config) / tp_device(tp_config).

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. tp_degree must divide the head count, so H3 caps at 8, and ParallelConfig.__post_init__ (_modeling_parallel.py:221) rejects TP together with CP — on this instance that leaves 8 of 64 cores usable. H3 already has a _cp_plan on main, and CP=4 is worth 1.62× for us at this resolution, so composing TP × CP looks like the highest-value follow-up. Happy to help with the Neuron side if you're interested.

FYI for anyone else on Neuron eager (not a diffusers issue). torch-neuronx only dispatches SDPA to its NKI flash kernel when L % 512 == 0 and S % 512 == 0; otherwise it silently falls back to a materialised score matrix. H3 is packed full self-attention with no mask, so it can't be zero-padded — at 1344×768×124 the video+audio rows come to 37710, and an arbitrary prompt length gives e.g. 37726, which emits a 1x7x37726x37726xf32 dot_general (39.9 TB) that never finishes compiling, with no hint as to why. Trimming the prompt to 178 tokens makes the total 37888 = 74×512 and it compiles fine. I'm raising this with the Neuron compiler team separately.

whn09 · 2026-09-07T03:31:58Z

@JingyaHuang two things stacked on this branch, both verified on a trn2.48xlarge:

Allow tensor parallelism and context parallelism in one ParallelConfig #14725 (against main, draft): let ParallelConfig carry a context_parallel_config and a tensor_parallel_config at the same time, sharing one mesh. Needed for models whose head count caps tp_degree well below the available accelerators — H3 has 56 heads, so tp_degree <= 8 and a 64-core host leaves 56 cores idle. It also fixes the Neuron TP backend's shard index, which reads the global rank where it means the rank's coordinate on the TP mesh (equal only while that mesh spans the whole world).
Apply the context-parallel hooks on the sharded-load path too JingyaHuang/diffusers#2 (against this branch, for your review): from_pretrained(..., parallel_config=...) applies tensor parallelism but never the context-parallel hooks, so a config carrying both is silently reduced to TP on your streaming path — no error, right numbers, every rank redundantly computing the whole sequence. Same for _load_dcp_checkpoint. The CP half of enable_parallelism moves into _apply_context_parallel and both paths call it.

The second one matters specifically for your use case: at 33B parameters the enable_parallelism route is unusable (every rank would first hold the full checkpoint), so this loader is the only way H3 can go past 8 cores. With both, H3 768p goes 9.285 s/step on 8 cores -> 3.941 s/step on 32 (TP=8 x ulysses=4), 2.36x, with visually equivalent output.

The PR includes a Neuron test following your _neuron_tp_worker.py convention, plus its negative control — output values alone cannot catch a missing CP hook, since the model still returns the right answer, so it asserts the hooks and the processors' config instead. Details and the exact commands are in the PR description.

This is independent of the non-persistent-buffer fix I left above, which is on the same loading path.

whn09 · 2026-10-08T05:32:24Z

Tested the current head (7e4a7eff) on trn2 with the hybrid worker that's in this branch (TP=2 × Ulysses=4, 8 cores, default mesh). It fails at collective init:

ENC:enc_init_comm  failed (2) to init a collective algorithm. reason: no_hier no_mesh replica-group: [0,2,4,6]

With ("ring", "ulysses", "tp"), the Ulysses groups are strided by tp_degree, and Neuron's all-to-all rejects strided replica groups (the comment above the mesh says the same thing). Changing only the mesh to mesh_shape=(tp_degree, ring_degree, ulysses_degree), mesh_dim_names=("tp", "ring", "ulysses") makes it pass (max abs diff 7.6e-06 vs the single-device reference), identical to #14725. TP's all-reduce accepts strided groups, so TP is the axis that can give up contiguity.

Env: torch 2.14.0, torch-neuronx 2.14.3.0.2008+2dda91f8d. Once this lands I'll close #14725 in favour of this PR.

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.

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>

sayakpaul

Thanks, left some comments. Would be nice to simplify the comments to explain the critical changes.

Update: Ran the tests on 2 L4 GPUs, TP=2, PyTorch 2.10.0+cu129 using HF Jobs and they pass: https://huggingface.co/jobs/sayakpaul/6ae033333fb22a235dd25b98

sayakpaul · 2026-10-27T04:45:39Z

    return resolved


+def _pre_shard_and_parallelize(


Where is this coming from? Could I get some explanation?

Hmm, I see what we're doing, but this is also a behavioral change, right? It assumes all ranks already have identical weights. Is that something we can always assume?

True the assumption won't be always valid for in-memory model with these edge cases:

Random initialization differs: each rank calls Model.from_config(config) with a different RNG state.

Ranks modify weights differently: for example, only rank 0 loads an updated state_dict, or reinitializes a layer.

Ranks load different checkpoints or revisions.

Buffers differ: random non-persistent buffers aren’t restored from a checkpoint and may differ even when weights match.

As we have two paths for the sharding shard-on-read(recommended) and in-memory which is slower and consume more memory as we added in the TP doc, and the changes here is a refactor to ensure that the memory of tpu/neuron won't explode when user goes with the in-memory path while doing TP for H3 (shard-on-read path is always fine tho)

For the edge cases, I can add a _sync_tp_model_state helper to broadcast model state from the first TP rank... Or if we don't want the refactor, I can just restore _styles / _shard_packed_param where previously used, only keep the required rowwise_split_input support in this file.

It's really If we want to support every possibility...

sayakpaul · 2026-10-27T04:53:31Z

+        # The row plan is built on CPU and the caller moves the finished pair to the denoiser's device: `unique`
+        # has a data-dependent output shape, which an accelerator would rather not trace. The layout step hands the
+        # index tensors over already on that device, so bring them back for the scatter below — indexing a CPU
+        # tensor with an accelerator one does not work.


Can we simplify this? I don't fully understand why this is even needed.

I simplified the docstring, it's just that video_indices and audio_indices were on device

diffusers/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py

Line 453 in b674fd8

block_state.video_indices = video_indices.to(device)

but

diffusers/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py

Line 1235 in b674fd8

row_timesteps = torch.full((sequence_length,), video_timestep, dtype=torch.float32)

happens on cpu.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…into add-h3-tp-support

sayakpaul

Thanks! Just one question regarding putting the indices on CPU unconditionally.

I am also running a snippet with real values to see if the videos look as expected. Will keep this comment updated once done.

Update: Was able to successfully run and create a realistic video: https://huggingface.co/jobs/sayakpaul/6aec4c52c78863f9c6333a02

tp2.mp4

Script: https://gist.github.com/sayakpaul/b52339e37b4513822de8a26ccdaa8039

sayakpaul · 2026-11-05T09:19:46Z

+        # The layout step puts the indices on the execution device, but this plan is built on CPU.
+        video_indices = video_indices.cpu()
+        audio_indices = audio_indices.cpu()


Why is this change needed?

it's just that video_indices and audio_indices were on device

diffusers/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py

Line 453 in b674fd8

block_state.video_indices = video_indices.to(device)

but

diffusers/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py

Line 1235 in b674fd8

row_timesteps = torch.full((sequence_length,), video_timestep, dtype=torch.float32)

happens on cpu, and it failed without moving indices back to cpu.

But why does it particularly show up for TP? Does regular inference fail without it?

~~I'm almost sure that's bc I changed the definition of pipeline's device in my tp PR, but let me confirm.~~

ok @sayakpaul , on cuda actually it works w/o. the change since PyTorch itself copies a foreign-device index onto the indexed tensor's device:
https://github.com/pytorch/pytorch/blob/735a2ec9b17c91c93d6714f88115f97394b836c9/aten/src/ATen/native/IndexingUtils.h#L60-L61
so there was a hidden GPU→CPU copy, and it's not the case for Neuron which check the device consistency first:
https://github.com/aws-neuron/torch-neuronx/blob/d9d8423807874b11c8d59242b09378e68bc9a504/torch_neuronx/csrc/core/dispatch/PreprocessArgs.cpp#L31-L42

So here the move of indices to cpu is equivalent to what had happened under the hood in PyTorch.

sayakpaul

Let's go!

github-actions Bot added size/L PR with diff > 200 LOC documentation Improvements or additions to documentation lora models tests modular-pipelines utils pipelines hooks and removed size/L PR with diff > 200 LOC labels Aug 26, 2026

github-actions Bot added the size/L PR with diff > 200 LOC label Aug 26, 2026

This was referenced Sep 7, 2026

Allow tensor parallelism and context parallelism in one ParallelConfig #14725

Open

Apply the context-parallel hooks on the sharded-load path too JingyaHuang/diffusers#2

Open

github-actions Bot added quantization CI single-file labels Oct 2, 2026

JingyaHuang force-pushed the add-h3-tp-support branch from 7e4a7ef to e4ae9de Compare October 8, 2026 14:17

github-actions Bot removed documentation Improvements or additions to documentation lora quantization utils pipelines CI single-file labels Oct 8, 2026

JingyaHuang marked this pull request as ready for review October 8, 2026 15:22

github-actions Bot added the documentation Improvements or additions to documentation label Oct 15, 2026

JingyaHuang and others added 4 commits October 15, 2026 14:24

Merge branch 'main' into add-h3-tp-support

3ea9cc5

feat:shard tensor-parallel weights locally on every device

660bb6d

Merge branch 'main' into add-h3-tp-support

82e45d4

JingyaHuang requested a review from sayakpaul October 27, 2026 00:25

sayakpaul reviewed Oct 27, 2026

View reviewed changes

Merge branch 'main' into add-h3-tp-support

dd54cfb

sayakpaul added this to the Release 0.43.0 milestone Oct 27, 2026

yzhautouskay mentioned this pull request Oct 27, 2026

[Cosmos3] Add native tensor parallelism #15137

Open

1 of 10 tasks

JingyaHuang and others added 10 commits October 28, 2026 14:23

doc: improve the explanation of rowwise_split_input

9ec4666

review: remove useless zero division prevention

d13196b

review: simplify the docstring

ff24da5

revert: Restore standard TP sharding and isolate the Neuron workaround

63732ec

revert: no refactoring of in-memory sharding

8e2fa1d

Merge branch 'main' into add-h3-tp-support

669f503

Merge branch 'main' into add-h3-tp-support

4d13666

revert: keep the in-memory TP path as on main

b680a25

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

style: pika

39bcfef

Merge branch 'add-h3-tp-support' of github.com:JingyaHuang/diffusers …

8b78dfe

…into add-h3-tp-support

github-actions Bot added size/M PR with diff < 200 LOC and removed size/L PR with diff > 200 LOC labels Nov 4, 2026

style: make fix-copies

8c74bb1

JingyaHuang requested a review from sayakpaul November 4, 2026 18:43

Merge branch 'main' into add-h3-tp-support

3dc0194

sayakpaul reviewed Nov 5, 2026

View reviewed changes

Merge branch 'main' into add-h3-tp-support

510487e

sayakpaul approved these changes Nov 5, 2026

View reviewed changes

sayakpaul merged commit 7e46c94 into huggingface:main Nov 5, 2026
33 checks passed

yiyixuxu mentioned this pull request Nov 9, 2026

_tp_plan / _cp_plan: allow entries for modules that are not always present #15218

Open

Conversation

JingyaHuang commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Validation

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

whn09 commented Sep 6, 2026

Bug: non-persistent buffers stay on CPU in the safetensors TP path

Uh oh!

whn09 commented Sep 7, 2026

Uh oh!

whn09 commented Oct 8, 2026

Uh oh!

sayakpaul left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JingyaHuang Nov 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JingyaHuang Nov 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

JingyaHuang commented Aug 26, 2026 •

edited

Loading

HuggingFaceDocBuilderDev commented Aug 26, 2026 •

edited

Loading

sayakpaul left a comment •

edited

Loading

sayakpaul left a comment •

edited

Loading

JingyaHuang Nov 5, 2026 •

edited

Loading

JingyaHuang Nov 5, 2026 •

edited

Loading