[JingyaHuang]

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

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

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]

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.

[whn09]

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]

@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]

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.

[whn09]

Re-tested 6845f333 on trn2 (torch-neuronx 2.14.3.0.2008+2dda91f8d).

Fixed, thanks: default TP+CP mesh (_neuron_hybrid_worker.py, TP=2 × Ulysses=4, no mesh: passes, max_abs_diff=7.6e-06), and on the sharded from_pretrained path CP is applied and the non-persistent buffers land on device. MiniMax-H3 TP=8 × Ulysses=4 on 32 cores runs at 3.97 s/step (1344×768).

One regression: with ulysses_anything=True and a sequence not divisible by the Ulysses degree (e.g. 10211 rows over 4), _all_gather_anything fails on Neuron:

RuntimeError: copy_(): source shape [1, 2552, 96] is not broadcastable to destination shape [1, 2553, 96]

Neuron's all-gather needs equal sizes per rank. Restoring the padding branch removed in 986ea9a fixes it (verified; output matches TP-only).

[JingyaHuang]

Thanks for contributing and helping out @whn09 ! I think the PR is in good shape for review @sayakpaul @DN6 , tested on both AWS Neuron and Google TPU v6.

[sayakpaul]

Can I get some explanation on when this case arrives?

[JingyaHuang]

this just shard the weight like "rowwise" but declares the input as replicated, it's applied to adaln_proj.linear of H3. Without it, adaln_proj would stay copied on every chip, and H3 wouldn't fit at tp=4 a 32 GB tpu chip.

[sayakpaul]

Should _pre_shard_and_parallelize then not be shape and device aware? Like for example, if we don't have such big shapes then doing it on CPU might be watseful, no?

[JingyaHuang]

actually here _pre_shard_and_parallelize should be applied to all device including cuda instead of just neuron/ tpu. Just removed the if condition.

And for the concern of slicing on CPU might be wasteful when shape is small. indeed small shape slicing on device could be faster than on CPU, but it also mean copy all weights on each rank that takes time then throw away (n_rank -1)/n_rank that takes time and could fail if shape is too big for memory-limited device. So I feel that pre-shard on CPU should be set as default, since it's memory-safe and should not make huge difference on speed when shape is small. Besides _pre_shard_and_parallelize doesn't move anything to the CPU, _local_shard slices the weight wherever it already is, if user insist slice on device and free parts for other rank afterward, they can always send to device by .to("cuda").

[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]

What is full width?

[JingyaHuang]

it's just to say all output features. With rowwise_split_input, each rank multiplies its slice of the input features by its weight shard, producing a partial sum for every output feature. AllReduce adds these partial sums, giving every rank the complete output. Let me make it less confusing.

[sayakpaul]

Where are these utilities going?

[JingyaHuang]

In this PR, the refactor moves weights sharding before hook registration with _pre_shard_and_parallelize (mostly for Neuron/TPU’s memory constraints, I removed the if condition in previous commit so it applies on cuda as well now) instead of using torch.distributed.tensor.parallel.parallelize_module, so we have

• _shard_packed_param -> covered by _sync_tp_model_state, _local_shard, and _pre_shard_and_parallelize
• _styles -> _hooks_only_styles (no weights partition tho since it's already done)

[sayakpaul]

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

[sayakpaul]

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?

[JingyaHuang]

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]

What is happening here?

[JingyaHuang]

not necessary at all, removed

[sayakpaul]

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