From acdf4bf9cbdef5345535bf1b53359983c226a9cd Mon Sep 17 00:00:00 2001
From: JingyaHuang <huang_jingya@outlook.com>
Date: Thu, 20 Aug 2026 15:10:26 +0000
Subject: [PATCH 1/5] [core] Shard tensor-parallel checkpoints on load and save

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.
---
 .../en/training/distributed_inference.md      |  60 ++-
 src/diffusers/hooks/tensor_parallel.py        | 261 +++++++--
 src/diffusers/hooks/tensor_parallel_neuron.py | 177 ++----
 src/diffusers/models/_modeling_parallel.py    |   2 +
 src/diffusers/models/model_loading_utils.py   | 115 +++-
 src/diffusers/models/modeling_utils.py        | 502 +++++++++++++++---
 src/diffusers/utils/__init__.py               |   1 +
 src/diffusers/utils/constants.py              |   1 +
 tests/models/testing_utils/parallelism.py     | 202 +++++++
 9 files changed, 1033 insertions(+), 288 deletions(-)

diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md
index 856572c2ff08..25d4dba4e339 100644
--- a/docs/source/en/training/distributed_inference.md
+++ b/docs/source/en/training/distributed_inference.md
@@ -436,43 +436,42 @@ pipeline = DiffusionPipeline.from_pretrained(
 
 [Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) shards the weight matrices of a model across devices. Each device holds a column-wise (`"colwise"`) or row-wise (`"rowwise"`) slice of each layer, computes a partial result, and an `AllReduce`/`AllGather` at the layer boundary reconstructs the full output. Unlike context parallelism, it reduces the per-device *weight* memory, which is useful for models that do not fit on a single device.
 
-Pass a [`TensorParallelConfig`] to [`~ModelMixin.enable_parallelism`]. `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model must define a `_tp_plan` (a flat mapping of module-name globs to a `"colwise"`/`"rowwise"` style).
+Pass a [`TensorParallelConfig`] to the `parallel_config` argument of the model's [`~ModelMixin.from_pretrained`]. `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model must define a `_tp_plan` (a flat mapping of module-name globs to a `"colwise"`/`"rowwise"` style).
+
+Loading this way shards the checkpoint *while reading it*: each rank reads only its own slice of each sharded weight and places it straight onto its own device. Nothing full-size is ever materialized, so per-rank memory falls as `tp_degree` rises.
 
 ```py
 import torch
 from torch import distributed as dist
-from diffusers import DiffusionPipeline, TensorParallelConfig
+from diffusers import DiffusionPipeline, Flux2Transformer2DModel, TensorParallelConfig
 
-def setup_distributed():
-    if not dist.is_initialized():
-        dist.init_process_group(backend="nccl")
-    rank = dist.get_rank()
+def main():
+    dist.init_process_group(backend="nccl")
+    rank, world_size = dist.get_rank(), dist.get_world_size()
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    return device
 
-def main():
-    device = setup_distributed()
-    world_size = dist.get_world_size()
+    # Each rank reads only its own shard of every planned weight, straight onto `cuda:rank`.
+    transformer = Flux2Transformer2DModel.from_pretrained(
+        "black-forest-labs/FLUX.2-dev",
+        subfolder="transformer",
+        torch_dtype=torch.bfloat16,
+        parallel_config=TensorParallelConfig(tp_degree=world_size),
+    )
 
     pipeline = DiffusionPipeline.from_pretrained(
-        "black-forest-labs/FLUX.2-dev", torch_dtype=torch.bfloat16
-    )  # weights stay on CPU
-
-    # Shard the transformer first, then move only each rank's slice onto the accelerator.
-    pipeline.transformer.enable_parallelism(config=TensorParallelConfig(tp_degree=world_size))
-    pipeline.transformer.to(device)
-
-    # Move the remaining, non-sharded components onto the accelerator individually.
+        "black-forest-labs/FLUX.2-dev", transformer=transformer, torch_dtype=torch.bfloat16
+    )
+    # The transformer is already on its device; move the remaining components individually. Do not call
+    # `pipeline.to(device)` — that would move every rank's shards onto the same device.
     pipeline.text_encoder.to(device)
     pipeline.vae.to(device)
 
     generator = torch.Generator().manual_seed(42)
     image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
-    if dist.get_rank() == 0:
+    if rank == 0:
         image.save("output.png")
-    if dist.is_initialized():
-        dist.destroy_process_group()
+    dist.destroy_process_group()
 
 if __name__ == "__main__":
     main()
@@ -484,6 +483,25 @@ torchrun --nproc-per-node 4 tensor_parallel_flux.py
 
 `tp_degree` is taken from `world_size` above, so `--nproc-per-node 4` shards the transformer across 4 devices.
 
+A tensor-parallel `parallel_config` cannot be combined with `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True`, DDUF checkpoints, or non-safetensors weights; each raises rather than quietly falling back to loading the full checkpoint. To shard a model that is already in memory, call [`~ModelMixin.enable_parallelism`] with the same config instead — that loads everything first and reshards it, so it costs full checkpoint memory on every rank.
+
+### Saving a tensor-parallel model
+
+[`~ModelMixin.save_pretrained`] gathers the shards back into ordinary full tensors, so the result is a normal checkpoint that loads with or without tensor parallelism. Gathering is a collective, so call it on **every** rank; only rank 0 writes.
+
+```py
+# on all ranks
+pipeline.transformer.save_pretrained("flux2-transformer")
+```
+
+For a model too large to gather onto a single rank, pass `dcp=True` to write a [distributed checkpoint](https://pytorch.org/docs/stable/distributed.checkpoint.html) instead. Every rank writes its own shards, so no full tensor is ever formed.
+
+```py
+pipeline.transformer.save_pretrained("flux2-transformer-dcp", dcp=True)
+```
+
+`from_pretrained` detects such a directory automatically, and reads it back with the same `parallel_config` you saved it under. Because a packed projection's shards are stored interleaved by the writing degree, the checkpoint only loads at that same `tp_degree`, and only with tensor parallelism — anything else raises rather than silently returning wrong weights. It is also local-only: a distributed checkpoint is recognized by the `.metadata` file in its directory, so it cannot be pushed to or loaded from the Hub. To lift any of these restrictions, re-save with the default (gathered) path, which produces an ordinary checkpoint.
+
 ### Writing a tensor parallelism plan
 
 Tensor parallelism only works on models that define a `_tp_plan`, a flat class attribute mapping module-name globs to a sharding style. Writing one is mostly a matter of pairing each projection that *expands* the hidden dimension with the projection that *contracts* it back.
diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py
index b90a5761d043..d3cbd82d7980 100644
--- a/src/diffusers/hooks/tensor_parallel.py
+++ b/src/diffusers/hooks/tensor_parallel.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import NamedTuple
+
 import torch
 
 from ..models._modeling_parallel import TensorParallelConfig
@@ -65,6 +67,150 @@ def _blocks_to_block_sizes(total_size: int, blocks: "list[int]") -> "list[int]":
     return [b * unit for b in blocks]
 
 
+class TPShardSpec(NamedTuple):
+    """How one parameter is laid out across the tensor-parallel ranks.
+
+    `dim` is the dimension sharded across ranks, or `None` when the parameter is replicated on every rank (a rowwise
+    bias, which is added after the all-reduce). `block_sizes` partitions `dim` into independently sharded blocks; a
+    plain `"colwise"` / `"rowwise"` style has a single block covering the whole dimension, and packed styles have one
+    per fused projection.
+    """
+
+    dim: "int | None"
+    block_sizes: "list[int] | None"
+
+
+def _local_shard(tensor, dim: int, block_sizes: "list[int]", tp_mesh) -> torch.Tensor:
+    """Extract this rank's slice of `tensor` along `dim`.
+
+    `tensor` may be a `torch.Tensor` or a safetensors `PySafeSlice`, so the same arithmetic serves both resharding a
+    weight already in memory and reading only one rank's slice off disk. Note that `PySafeSlice` exposes `get_shape()`
+    rather than `.ndim`, and that a `dim`-1 slice comes back strided, hence the final `contiguous()` —
+    `DTensor.from_local` needs a contiguous local tensor.
+    """
+    rank = tp_mesh.get_local_rank()
+    tp_size = tp_mesh.size()
+    ndim = tensor.dim() if isinstance(tensor, torch.Tensor) else len(tensor.get_shape())
+
+    parts, offset = [], 0
+    for block_size in block_sizes:
+        # An uneven split is rejected rather than silently handed to `Shard`, which pads the tail
+        # and would break both the paired colwise/rowwise matmul and `_unshard_gathered`.
+        if block_size % tp_size != 0:
+            raise ValueError(
+                f"Cannot shard a block of size {block_size} across {tp_size} tensor-parallel ranks: "
+                f"{block_size} is not divisible by {tp_size}."
+            )
+        chunk = block_size // tp_size
+        index = [slice(None)] * ndim
+        index[dim] = slice(offset + rank * chunk, offset + (rank + 1) * chunk)
+        parts.append(tensor[tuple(index)])
+        offset += block_size
+
+    local = parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim)
+    return local.contiguous()
+
+
+def _unshard_gathered(gathered: torch.Tensor, dim: int, block_sizes: "list[int]", tp_size: int) -> torch.Tensor:
+    """Undo `_local_shard`'s block interleaving on an all-gathered tensor.
+
+    `DTensor.full_tensor()` concatenates the local shards rank-major, so a packed weight comes back as `[block0_rank0,
+    block1_rank0, block0_rank1, block1_rank1, ...]` and has to be regrouped by block. A single block is already in the
+    original order and passes through unchanged.
+    """
+    if len(block_sizes) == 1:
+        return gathered
+
+    local_sizes = [block_size // tp_size for block_size in block_sizes]
+    stride = sum(local_sizes)
+    parts = []
+    for i, local_size in enumerate(local_sizes):
+        offset = sum(local_sizes[:i])
+        parts.extend(gathered.narrow(dim, rank * stride + offset, local_size) for rank in range(tp_size))
+    return torch.cat(parts, dim=dim)
+
+
+def gather_tp_state_dict(state_dict: dict, specs: "dict[str, TPShardSpec]", config: TensorParallelConfig) -> dict:
+    """Reassemble a tensor-parallel `state_dict` into ordinary full tensors.
+
+    Every `DTensor` is all-gathered back to its full shape and, for the packed styles, reordered by `_unshard_gathered`
+    — `full_tensor()` alone would leave the fused blocks interleaved by rank. Replicated and unplanned parameters pass
+    through untouched.
+
+    `full_tensor()` is a collective, so this must run on **every** rank even though usually only rank 0 goes on to
+    write the result.
+    """
+    from torch.distributed.tensor import DTensor
+
+    tp_size = config._tp_degree
+    gathered = {}
+    for key, value in state_dict.items():
+        if not isinstance(value, DTensor):
+            gathered[key] = value
+            continue
+        # `full_tensor()` is the collective; the reorder after it is plain tensor arithmetic, so keep it off
+        # the accelerator — CPU is where this state dict is headed anyway, since it is about to be written.
+        full = value.full_tensor().cpu()
+        spec = specs[key]
+        if spec.dim is not None:
+            full = _unshard_gathered(full, spec.dim, spec.block_sizes, tp_size)
+        gathered[key] = full.contiguous()
+    return gathered
+
+
+def resolve_tp_shard_specs(model: torch.nn.Module, tp_plan: dict) -> "dict[str, TPShardSpec]":
+    """Map every `_tp_plan`-covered parameter name to its `TPShardSpec`.
+
+    Parameters absent from the result are untouched by tensor parallelism. Both `weight` and `bias` of each planned
+    module are covered.
+
+    The plan is expanded by `_resolve_tp_plan` so there is a single implementation of the glob rules; the `id(module)
+    -> name` map recovers qualified names from the submodules it returns. Going back through `_resolve_tp_plan` also
+    handles a model that reuses one block instance in two places, which re-expanding the globs here would get wrong.
+
+    Safe to call on a meta model: only shapes, `bias is not None`, and the packed-block attributes set in the module's
+    `__init__` are read.
+    """
+    names = {id(module): name for name, module in model.named_modules()}
+    specs: dict[str, TPShardSpec] = {}
+
+    for block, relative_plan in _resolve_tp_plan(model, tp_plan):
+        prefix = names[id(block)]
+        for relative_path, style in relative_plan.items():
+            submodule = block
+            for atom in relative_path.split("."):
+                submodule = getattr(submodule, atom)
+            path = f"{prefix}.{relative_path}" if prefix else relative_path
+
+            # `_tp_packed_*_blocks` hold absolute sizes rather than proportions; that works because
+            # they sum to the full dimension, so `_blocks_to_block_sizes` computes `unit == 1`.
+            if style == "colwise":
+                weight_spec = TPShardSpec(0, [submodule.weight.shape[0]])
+                bias_spec = weight_spec
+            elif style == "rowwise":
+                weight_spec = TPShardSpec(1, [submodule.weight.shape[1]])
+                bias_spec = TPShardSpec(None, None)
+            elif isinstance(style, PackedColwiseParallel):
+                blocks = style.blocks if style.blocks is not None else submodule._tp_packed_col_blocks
+                weight_spec = TPShardSpec(0, _blocks_to_block_sizes(submodule.weight.shape[0], blocks))
+                bias_spec = weight_spec
+            elif isinstance(style, PackedRowwiseParallel):
+                blocks = style.blocks if style.blocks is not None else submodule._tp_packed_row_blocks
+                weight_spec = TPShardSpec(1, _blocks_to_block_sizes(submodule.weight.shape[1], blocks))
+                bias_spec = TPShardSpec(None, None)
+            else:
+                raise ValueError(
+                    f"Unsupported tensor-parallel style '{style}' for '{path}'. "
+                    f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+                )
+
+            specs[f"{path}.weight"] = weight_spec
+            if submodule.bias is not None:
+                specs[f"{path}.bias"] = bias_spec
+
+    return specs
+
+
 def _resolve_tp_plan(model: torch.nn.Module, tp_plan: dict) -> list:
     """Group a flat `_tp_plan` into per-block `(submodule, {relative_path: style})` plans.
 
@@ -123,32 +269,24 @@ def _make_packed_col(marker: PackedColwiseParallel) -> ColwiseParallel:
 
         class _PackedColwiseImpl(ColwiseParallel):
             def _partition_linear_fn(self, name, module, device_mesh):
-                blocks = _blocks if _blocks is not None else getattr(module, "_tp_packed_col_blocks")
-                rank = device_mesh.get_local_rank()
-                tp_size = device_mesh.size()
+                blocks = _blocks if _blocks is not None else module._tp_packed_col_blocks
                 # Both weight (`[out, in]`) and bias (`[out]`) are sharded row-wise (dim 0) with the same per-block
                 # slicing so each rank's bias rows line up with its weight rows for the packed layout.
                 for param_name, param in module.named_parameters():
+                    # Replicate before slicing: the broadcast from `src_data_rank` is what makes one rank's
+                    # weights authoritative when the model was randomly initialized rather than loaded from a
+                    # checkpoint, in which case every rank starts with different values.
                     full = distribute_tensor(
                         param, device_mesh, [Replicate()], src_data_rank=self.src_data_rank
                     ).to_local()
-                    block_sizes = _blocks_to_block_sizes(full.shape[0], blocks)
-                    parts, offset = [], 0
-                    for bs in block_sizes:
-                        if bs % tp_size != 0:
-                            raise ValueError(
-                                f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: "
-                                f"{bs} is not divisible by {tp_size}."
-                            )
-                        chunk = bs // tp_size
-                        parts.append(full[offset + rank * chunk : offset + (rank + 1) * chunk].contiguous())
-                        offset += bs
-                    local = torch.cat(parts, dim=0)
-                    dist_param = nn.Parameter(
-                        DTensor.from_local(local, device_mesh, [Shard(0)], run_check=False),
-                        requires_grad=param.requires_grad,
+                    local = _local_shard(full, 0, _blocks_to_block_sizes(full.shape[0], blocks), device_mesh)
+                    module.register_parameter(
+                        param_name,
+                        nn.Parameter(
+                            DTensor.from_local(local, device_mesh, [Shard(0)], run_check=False),
+                            requires_grad=param.requires_grad,
+                        ),
                     )
-                    module.register_parameter(param_name, dist_param)
 
         return _PackedColwiseImpl()
 
@@ -157,26 +295,14 @@ def _make_packed_row(marker: PackedRowwiseParallel) -> RowwiseParallel:
 
         class _PackedRowwiseImpl(RowwiseParallel):
             def _partition_linear_fn(self, name, module, device_mesh):
-                blocks = _blocks if _blocks is not None else getattr(module, "_tp_packed_row_blocks")
-                rank = device_mesh.get_local_rank()
-                tp_size = device_mesh.size()
+                blocks = _blocks if _blocks is not None else module._tp_packed_row_blocks
                 for param_name, param in module.named_parameters():
                     if param_name == "weight":
+                        # See `_make_packed_col`: replicate first so one rank's weights win.
                         full = distribute_tensor(
                             param, device_mesh, [Replicate()], src_data_rank=self.src_data_rank
                         ).to_local()
-                        block_sizes = _blocks_to_block_sizes(full.shape[1], blocks)
-                        parts, offset = [], 0
-                        for bs in block_sizes:
-                            if bs % tp_size != 0:
-                                raise ValueError(
-                                    f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: "
-                                    f"{bs} is not divisible by {tp_size}."
-                                )
-                            chunk = bs // tp_size
-                            parts.append(full[:, offset + rank * chunk : offset + (rank + 1) * chunk].contiguous())
-                            offset += bs
-                        local = torch.cat(parts, dim=1)
+                        local = _local_shard(full, 1, _blocks_to_block_sizes(full.shape[1], blocks), device_mesh)
                         dist_param = nn.Parameter(
                             DTensor.from_local(local, device_mesh, [Shard(1)], run_check=False),
                             requires_grad=param.requires_grad,
@@ -193,7 +319,7 @@ def _partition_linear_fn(self, name, module, device_mesh):
     # `distribute_tensor` accepts an indivisible shard dim and just gives the trailing ranks a smaller (or empty)
     # slice, so an uneven split does not raise here — it surfaces much later as a shape or numerics error, because
     # the attention head split and the paired colwise/rowwise Linear both assume equal shards. Reject it up front,
-    # matching what the packed styles above and the Neuron pre-shard path already do.
+    # matching what `_local_shard` already does for the packed styles.
     def _make_checked_col(path: str) -> ColwiseParallel:
         class _CheckedColwiseImpl(ColwiseParallel):
             def _partition_linear_fn(self, name, module, device_mesh):
@@ -240,16 +366,68 @@ def _partition_linear_fn(self, name, module, device_mesh):
     return resolved
 
 
+def _hooks_only_styles(relative_plan: dict) -> dict:
+    """Map a `{relative_path: style}` plan to styles that partition nothing.
+
+    Used when the caller has already placed every planned parameter as a sharded `DTensor`. `parallelize_module` then
+    runs only to register the forward input/output hooks; `_partition_linear_fn` must not re-partition. Packed and
+    plain styles share hook behaviour, so both collapse onto the two styles here.
+
+    Note this is not purely additive: `distribute_module` still replicates any *remaining* plain parameter of the
+    targeted module into a `Replicate()` DTensor via a broadcast. Callers should therefore place every planned
+    parameter themselves, and must ensure none is left on `meta` — the broadcast would be issued on a meta tensor.
+    """
+    from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
+
+    class _NoPartitionColwise(ColwiseParallel):
+        def _partition_linear_fn(self, name, module, device_mesh):
+            pass  # weight already Shard(0)
+
+    class _NoPartitionRowwise(RowwiseParallel):
+        def _partition_linear_fn(self, name, module, device_mesh):
+            pass  # weight already Shard(1)
+
+    resolved = {}
+    for path, style in relative_plan.items():
+        if style == "colwise" or isinstance(style, PackedColwiseParallel):
+            resolved[path] = _NoPartitionColwise()
+        elif style == "rowwise" or isinstance(style, PackedRowwiseParallel):
+            resolved[path] = _NoPartitionRowwise()
+        else:
+            raise ValueError(
+                f"Unsupported tensor-parallel style '{style}' for '{path}'. "
+                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+            )
+    return resolved
+
+
 def apply_tensor_parallel(
     model: torch.nn.Module,
     config: TensorParallelConfig,
     tp_plan: dict,
+    weights_already_sharded: bool = False,
 ) -> None:
-    """Apply tensor parallel on a model from its flat `_tp_plan`."""
+    """Apply tensor parallel on a model from its flat `_tp_plan`.
+
+    Set `weights_already_sharded` when the planned parameters are already `DTensor` shards, as they are after a
+    streaming `from_pretrained` load; only the forward hooks are then registered. This is passed explicitly rather than
+    detected, because a planned parameter missing from the checkpoint would still be a meta tensor and would make
+    detection say "not sharded" for a model that is in fact half-sharded.
+    """
+    if tp_plan is None:
+        raise ValueError(
+            "`_tp_plan` must be set on the model class to use tensor parallelism. "
+            f"'{model.__class__.__name__}' does not define one."
+        )
+
     tp_mesh = config._mesh
     if tp_mesh is None:
         raise ValueError("`config._mesh` is None. Call `config.setup(rank, world_size, device)` before applying TP.")
 
+    num_heads = getattr(model.config, "num_attention_heads", None)
+    if num_heads is not None and num_heads % config._tp_degree != 0:
+        raise ValueError(f"`tp_degree` ({config._tp_degree}) must divide the number of attention heads ({num_heads}).")
+
     if tp_mesh.device_type not in _SUPPORTED_TP_DEVICES:
         raise ValueError(
             f"Tensor parallelism is not supported on device type '{tp_mesh.device_type}'. Supported device types are "
@@ -261,13 +439,18 @@ def apply_tensor_parallel(
     groups = _resolve_tp_plan(model, tp_plan)
     logger.debug(f"Applying tensor parallel (backend={backend}) over {len(groups)} module group(s) on mesh {tp_mesh}.")
 
+    from torch.distributed.tensor.parallel import parallelize_module
+
+    if weights_already_sharded:
+        for submodule, relative_plan in groups:
+            parallelize_module(submodule, tp_mesh, _hooks_only_styles(relative_plan))
+        return
+
     if backend == "neuron":
         from .tensor_parallel_neuron import _apply_tp_neuron
 
-        _apply_tp_neuron(model, tp_mesh, groups)
+        _apply_tp_neuron(model, tp_mesh, groups, resolve_tp_shard_specs(model, tp_plan))
         return
 
-    from torch.distributed.tensor.parallel import parallelize_module
-
     for submodule, relative_plan in groups:
         parallelize_module(submodule, tp_mesh, _styles(relative_plan))
diff --git a/src/diffusers/hooks/tensor_parallel_neuron.py b/src/diffusers/hooks/tensor_parallel_neuron.py
index 6b8f219a17ff..ffcba0973d81 100644
--- a/src/diffusers/hooks/tensor_parallel_neuron.py
+++ b/src/diffusers/hooks/tensor_parallel_neuron.py
@@ -17,162 +17,59 @@
 The difference from the generic path is a workaround for a Neuron NRT bug: consecutive `reduce_scatter` collectives for
 large weight tensors (≥ 5120×5120) can fail when all layers are distributed in a single `parallelize_module` call. The
 fix is to pre-shard each weight locally on CPU via `DTensor.from_local` *before* calling `parallelize_module`; the
-latter then sees already-placed DTensors, skips the collective for weights, but still registers the required
+latter then sees already-placed DTensors and skips the collective for weights, while still registering the required
 input/output hooks for the forward pass.
+
+Only needed for a model that is already in memory. `from_pretrained` with a tensor-parallel `parallel_config` streams
+each rank's slice straight off disk into its DTensor, which issues no weight collectives at all and so cannot hit the
+bug in the first place.
 """
 
 import torch
-import torch.distributed as dist
 import torch.nn as nn
 
-
-def _neuron_styles(relative_plan: dict) -> dict:
-    """Map a `{relative_path: style}` plan to no-op-partition styles for Neuron.
-
-    Weights (and biases) are pre-sharded in `_pre_shard_and_tp`, so `parallelize_module` runs only to register the
-    forward hooks; `_partition_linear_fn` must not re-partition. Packed and plain styles share hook behavior, so both
-    collapse onto the two no-op styles.
-    """
-    from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
-
-    from .tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel
-
-    class _NeuronColwise(ColwiseParallel):
-        def _partition_linear_fn(self, name, module, device_mesh):
-            pass  # weight already Shard(0) via DTensor.from_local; parallelize_module runs only for the hooks
-
-    class _NeuronRowwise(RowwiseParallel):
-        def _partition_linear_fn(self, name, module, device_mesh):
-            pass  # weight already Shard(1) via DTensor.from_local; parallelize_module runs only for the hooks
-
-    resolved = {}
-    for path, style in relative_plan.items():
-        if style == "colwise" or isinstance(style, PackedColwiseParallel):
-            resolved[path] = _NeuronColwise()
-        elif style == "rowwise" or isinstance(style, PackedRowwiseParallel):
-            resolved[path] = _NeuronRowwise()
-        else:
-            raise ValueError(
-                f"Unsupported tensor-parallel style '{style}' for '{path}'. "
-                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
-            )
-    return resolved
-
-
-def _pre_shard_and_tp(
-    module: nn.Module,
-    tp_mesh: "torch.distributed.device_mesh.DeviceMesh",
-    original_plan: dict,
-    rank: int,
-    tp_size: int,
-) -> None:
-    """Pre-shard Linear weights via `DTensor.from_local`, then call `parallelize_module`.
-
-    Workaround for a Neuron NRT bug where consecutive `reduce_scatter` calls for large weight tensors (≥ 5120×5120)
-    fail when all layers are distributed in a single `parallelize_module` call. Pre-sharding each weight on CPU means
-    it is already an on-device DTensor when `parallelize_module` runs (via `_neuron_styles`), so the collective is
-    skipped while the forward hooks are still registered.
-    """
-    from torch.distributed.tensor import DTensor, Replicate, Shard
-    from torch.distributed.tensor.parallel import parallelize_module
-
-    from .tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel, _blocks_to_block_sizes
-
-    device = torch.neuron.current_device()
-
-    for path, orig_style in original_plan.items():
-        # Resolve nested attribute path (e.g. "attn.to_q" or "attn.to_out.0")
-        submod = module
-        for part in path.split("."):
-            submod = getattr(submod, part)
-
-        if not hasattr(submod, "weight"):
-            raise ValueError(f"`_tp_plan` entry '{path}' does not resolve to a module with a `weight` parameter.")
-
-        w = submod.weight.data  # CPU at this point
-        b = submod.bias.data if submod.bias is not None else None
-        if isinstance(orig_style, PackedColwiseParallel):
-            blocks = orig_style.blocks if orig_style.blocks is not None else getattr(submod, "_tp_packed_col_blocks")
-            block_sizes = _blocks_to_block_sizes(w.shape[0], blocks)
-            parts, bias_parts, offset = [], [], 0
-            for bs in block_sizes:
-                if bs % tp_size != 0:
-                    raise ValueError(
-                        f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: "
-                        f"{bs} is not divisible by {tp_size}."
-                    )
-                chunk = bs // tp_size
-                sl = slice(offset + rank * chunk, offset + (rank + 1) * chunk)
-                parts.append(w[sl, :].contiguous())
-                if b is not None:
-                    bias_parts.append(b[sl].contiguous())
-                offset += bs
-            shard = torch.cat(parts, dim=0).to(device)
-            submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(0)]))
-            if b is not None:
-                bias_shard = torch.cat(bias_parts, dim=0).to(device)
-                submod.bias = nn.Parameter(DTensor.from_local(bias_shard, tp_mesh, [Shard(0)]))
-        elif isinstance(orig_style, PackedRowwiseParallel):
-            blocks = orig_style.blocks if orig_style.blocks is not None else getattr(submod, "_tp_packed_row_blocks")
-            block_sizes = _blocks_to_block_sizes(w.shape[1], blocks)
-            parts, offset = [], 0
-            for bs in block_sizes:
-                if bs % tp_size != 0:
-                    raise ValueError(
-                        f"Cannot shard packed block of size {bs} across {tp_size} tensor-parallel ranks: "
-                        f"{bs} is not divisible by {tp_size}."
-                    )
-                chunk = bs // tp_size
-                parts.append(w[:, offset + rank * chunk : offset + (rank + 1) * chunk].contiguous())
-                offset += bs
-            shard = torch.cat(parts, dim=1).to(device)
-            submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(1)]))
-            if b is not None:  # rowwise bias is added post-reduction → keep it replicated
-                submod.bias = nn.Parameter(DTensor.from_local(b.to(device), tp_mesh, [Replicate()]))
-        elif orig_style == "colwise":
-            if w.shape[0] % tp_size != 0:
-                raise ValueError(
-                    f"Cannot colwise-shard '{path}' weight rows ({w.shape[0]}) across {tp_size} "
-                    f"tensor-parallel ranks: not divisible by {tp_size}."
-                )
-            rows = w.shape[0] // tp_size
-            sl = slice(rank * rows, (rank + 1) * rows)
-            submod.weight = nn.Parameter(DTensor.from_local(w[sl, :].contiguous().to(device), tp_mesh, [Shard(0)]))
-            if b is not None:
-                submod.bias = nn.Parameter(DTensor.from_local(b[sl].contiguous().to(device), tp_mesh, [Shard(0)]))
-        elif orig_style == "rowwise":
-            if w.shape[1] % tp_size != 0:
-                raise ValueError(
-                    f"Cannot rowwise-shard '{path}' weight columns ({w.shape[1]}) across {tp_size} "
-                    f"tensor-parallel ranks: not divisible by {tp_size}."
-                )
-            cols = w.shape[1] // tp_size
-            shard = w[:, rank * cols : (rank + 1) * cols].contiguous().to(device)
-            submod.weight = nn.Parameter(DTensor.from_local(shard, tp_mesh, [Shard(1)]))
-            if b is not None:  # rowwise bias is added post-reduction → keep it replicated
-                submod.bias = nn.Parameter(DTensor.from_local(b.to(device), tp_mesh, [Replicate()]))
-
-    # parallelize_module is now a no-op for weight distribution (already DTensors)
-    # but still registers the input/output hooks required for the forward pass.
-    parallelize_module(module, tp_mesh, _neuron_styles(original_plan))
+from .tensor_parallel import TPShardSpec, _hooks_only_styles, _local_shard
 
 
 def _apply_tp_neuron(
     model: nn.Module,
     tp_mesh: "torch.distributed.device_mesh.DeviceMesh",
     groups: list,
+    specs: "dict[str, TPShardSpec]",
 ) -> None:
-    """Apply tensor parallelism on Neuron from resolved `_tp_plan` groups.
+    """Pre-shard the planned parameters via `DTensor.from_local`, then register the forward hooks.
 
-    `groups` is produced by `diffusers.hooks.tensor_parallel._resolve_tp_plan` — the same source of truth used by the
-    generic path, so the two backends shard identical layers. For each `(block, relative_plan)` group this pre-shards
-    the weights via `DTensor.from_local` (Neuron NRT consecutive-reduce-scatter workaround), then calls
-    `parallelize_module` to register the forward hooks.
+    `groups` and `specs` both come from the model's `_tp_plan` via `diffusers.hooks.tensor_parallel._resolve_tp_plan` /
+    `resolve_tp_shard_specs`, the same source of truth the generic path uses, so the two backends shard identical
+    layers.
 
     Model weights must be on CPU when this is called.
     """
-    rank = dist.get_rank()
-    tp_size = tp_mesh.size()
+    from torch.distributed.tensor import DTensor, Replicate, Shard
+    from torch.distributed.tensor.parallel import parallelize_module
 
+    device = torch.neuron.current_device()
+
+    for name, spec in specs.items():
+        path, _, param_name = name.rpartition(".")
+        module = model.get_submodule(path)
+        param = getattr(module, param_name)
+
+        if spec.dim is None:
+            # A rowwise bias is added after the all-reduce, so every rank needs the whole vector.
+            local, placement = param.data, Replicate()
+        else:
+            local, placement = _local_shard(param.data, spec.dim, spec.block_sizes, tp_mesh), Shard(spec.dim)
+
+        module.register_parameter(
+            param_name,
+            nn.Parameter(
+                DTensor.from_local(local.to(device), tp_mesh, [placement]),
+                requires_grad=param.requires_grad,
+            ),
+        )
+
+    # `parallelize_module` is now a no-op for weight distribution (they are already DTensors) but still registers the
+    # input/output hooks required for the forward pass.
     for block, relative_plan in groups:
-        _pre_shard_and_tp(block, tp_mesh, relative_plan, rank, tp_size)
+        parallelize_module(block, tp_mesh, _hooks_only_styles(relative_plan))
diff --git a/src/diffusers/models/_modeling_parallel.py b/src/diffusers/models/_modeling_parallel.py
index 86627284e078..b54e86d6b4f2 100644
--- a/src/diffusers/models/_modeling_parallel.py
+++ b/src/diffusers/models/_modeling_parallel.py
@@ -186,6 +186,8 @@ def __post_init__(self):
             raise ValueError("`tp_degree` must be >= 1.")
 
     def setup(self, rank: int, world_size: int, device: torch.device, mesh: torch.distributed.device_mesh.DeviceMesh):
+        if mesh.size() > world_size:
+            raise ValueError(f"Tensor parallel degree ({mesh.size()}) cannot exceed the world size ({world_size}).")
         self._rank = rank
         self._world_size = world_size
         self._device = device
diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py
index abbde8082bb5..d0ba37514b9e 100644
--- a/src/diffusers/models/model_loading_utils.py
+++ b/src/diffusers/models/model_loading_utils.py
@@ -388,6 +388,99 @@ def _load_shard_file(
     return offload_index, state_dict_index, mismatched_keys, error_msgs
 
 
+def _load_shard_file_tp(
+    shard_file,
+    model,
+    model_state_dict,
+    tp_shard_specs,
+    tp_config,
+    dtype=None,
+    keep_in_fp32_modules=None,
+    unexpected_keys=None,
+    ignore_mismatched_sizes=False,
+):
+    """Load one safetensors shard, reading only this rank's slice of each tensor-parallel parameter.
+
+    The counterpart of `_load_shard_file` for a model being sharded by `_tp_plan`, with the same return contract so it
+    can be swapped in as `load_fn`. Parameters covered by `tp_shard_specs` are sliced while still on disk and placed as
+    `DTensor`s; everything else is read whole and replicated on every rank, exactly as tensor parallelism requires.
+
+    Slicing before the dtype cast is the point of the whole exercise: `load_model_dict_into_meta` casts the full tensor
+    first, which would materialize it in full on every rank.
+    """
+    from safetensors import safe_open
+    from torch.distributed.tensor import DTensor, Replicate, Shard
+
+    from ..hooks.tensor_parallel import _local_shard
+
+    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
+
+    mismatched_keys = []
+
+    # The slices are lazy views over the file, so every read has to happen inside this block.
+    with safe_open(shard_file, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            if key not in model_state_dict:
+                unexpected_keys.append(key)
+                continue
+
+            checkpoint_slice = f.get_slice(key)
+            expected_shape = model_state_dict[key].shape
+            if tuple(checkpoint_slice.get_shape()) != tuple(expected_shape):
+                # Checkpoints always hold full tensors, so the comparison is against the unsharded shape.
+                if not ignore_mismatched_sizes:
+                    raise ValueError(
+                        f"Cannot load {key} because it has shape {tuple(checkpoint_slice.get_shape())} in the "
+                        f"checkpoint but shape {tuple(expected_shape)} in {model.__class__.__name__}. Pass "
+                        "`ignore_mismatched_sizes=True` to skip it and keep the randomly initialized weight."
+                    )
+                mismatched_keys.append((key, tuple(checkpoint_slice.get_shape()), tuple(expected_shape)))
+                continue
+
+            spec = tp_shard_specs.get(key)
+            if spec is None or spec.dim is None:
+                param = checkpoint_slice[...]
+            else:
+                param = _local_shard(checkpoint_slice, spec.dim, spec.block_sizes, tp_mesh)
+
+            # Mirror `load_model_dict_into_meta`: only floating point weights are cast, and modules held
+            # in fp32 override the requested dtype.
+            if dtype is not None and torch.is_floating_point(param):
+                if keep_in_fp32_modules is not None and any(
+                    module_to_keep_in_fp32 in key.split(".") for module_to_keep_in_fp32 in keep_in_fp32_modules
+                ):
+                    param = param.to(torch.float32)
+                else:
+                    param = param.to(dtype)
+
+            if spec is None:
+                set_module_tensor_to_device(model, key, device, value=param)
+                continue
+
+            path, _, param_name = key.rpartition(".")
+            module = model.get_submodule(path)
+            # A rowwise bias is added after the all-reduce, so it stays replicated. It still has to be a
+            # DTensor: a plain tensor next to a sharded weight fails the `addmm` dispatch.
+            placement = Replicate() if spec.dim is None else Shard(spec.dim)
+            module.register_parameter(
+                param_name,
+                torch.nn.Parameter(
+                    DTensor.from_local(param.to(device), tp_mesh, [placement], run_check=False),
+                    requires_grad=getattr(module, param_name).requires_grad,
+                ),
+            )
+
+    # `offload_index` / `state_dict_index` are always None here: offloading and tensor parallelism are
+    # rejected as a combination by `from_pretrained`.
+    return None, None, mismatched_keys, []
+
+
 def _load_shard_files_with_threadpool(
     shard_files,
     model,
@@ -452,28 +545,6 @@ def _load_shard_files_with_threadpool(
     return offload_index, state_dict_index, mismatched_keys, error_msgs
 
 
-def _find_mismatched_keys(
-    state_dict,
-    model_state_dict,
-    loaded_keys,
-    ignore_mismatched_sizes,
-):
-    mismatched_keys = []
-    if ignore_mismatched_sizes:
-        for checkpoint_key in loaded_keys:
-            model_key = checkpoint_key
-            # If the checkpoint is sharded, we may not have the key here.
-            if checkpoint_key not in state_dict:
-                continue
-
-            if model_key in model_state_dict and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape:
-                mismatched_keys.append(
-                    (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape)
-                )
-                del state_dict[checkpoint_key]
-    return mismatched_keys
-
-
 def _load_state_dict_into_model(
     model_to_load, state_dict: OrderedDict, assign_to_params_buffers: bool = False
 ) -> list[str]:
diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py
index 5af0ca0e6278..e690fff0b058 100644
--- a/src/diffusers/models/modeling_utils.py
+++ b/src/diffusers/models/modeling_utils.py
@@ -42,6 +42,7 @@
 from ..quantizers.quantization_config import QuantizationMethod
 from ..utils import (
     CONFIG_NAME,
+    DCP_CONFIG_NAME,
     FLASHPACK_WEIGHTS_NAME,
     HF_ENABLE_PARALLEL_LOADING,
     SAFE_WEIGHTS_INDEX_NAME,
@@ -76,6 +77,7 @@
     _fetch_index_file,
     _fetch_index_file_legacy,
     _load_shard_file,
+    _load_shard_file_tp,
     _load_shard_files_with_threadpool,
     load_state_dict,
 )
@@ -686,6 +688,7 @@ def save_pretrained(
         max_shard_size: int | str = "10GB",
         push_to_hub: bool = False,
         use_flashpack: bool = False,
+        dcp: bool = False,
         **kwargs,
     ):
         """
@@ -718,8 +721,18 @@ def save_pretrained(
                 Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
                 repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
                 namespace).
+            dcp (`bool`, *optional*, defaults to `False`):
+                Write a [`torch.distributed.checkpoint`](https://pytorch.org/docs/stable/distributed.checkpoint.html)
+                directory instead of safetensors files. Only valid for a tensor-parallel model: every rank writes its
+                own shards, so no full tensor is ever materialized, which matters for models too large to gather onto
+                one rank. Read it back with `from_pretrained`, which detects the directory automatically and can
+                reshard it to a different `tp_degree`.
             kwargs (`dict[str, Any]`, *optional*):
                 Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
+
+        A tensor-parallel model is gathered back into ordinary full tensors before saving, so the result is a normal
+        checkpoint that loads without tensor parallelism. Gathering is a collective: call `save_pretrained` on every
+        rank, not just the main process. Only rank 0 writes.
         """
         if os.path.isfile(save_directory):
             logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
@@ -742,6 +755,76 @@ def save_pretrained(
                     " the logger on the traceback to understand the reason why the quantized model is not serializable."
                 )
 
+        tp_config = None
+        if self._parallel_config is not None:
+            tp_config = self._parallel_config.tensor_parallel_config
+
+        if dcp:
+            if tp_config is None:
+                raise ValueError(
+                    "`dcp=True` is only meaningful for a tensor-parallel model, whose parameters are sharded "
+                    "across ranks. Save an unsharded model with the default safetensors path."
+                )
+            unsupported = [
+                name
+                for name, value in (
+                    ("use_flashpack", use_flashpack),
+                    ("variant", variant),
+                    ("safe_serialization=False", not safe_serialization),
+                    ("save_function", save_function),
+                )
+                if value
+            ]
+            if unsupported:
+                raise ValueError(
+                    f"{unsupported} cannot be combined with `dcp=True`: a distributed checkpoint is a directory "
+                    "of `.distcp` shards, not a single named weights file."
+                )
+            if push_to_hub:
+                # `from_pretrained` only recognizes a distributed checkpoint by looking for `.metadata` in a
+                # local directory, so one cannot be loaded back from the Hub.
+                raise ValueError(
+                    "`push_to_hub=True` cannot be combined with `dcp=True`: a distributed checkpoint can only "
+                    "be loaded from a local directory. Save it with the default safetensors path to push it."
+                )
+
+            import torch.distributed.checkpoint as dcp_api
+
+            os.makedirs(save_directory, exist_ok=True)
+            if tp_config._mesh.get_local_rank() == 0:
+                self.save_config(save_directory)
+                # A packed weight's local shard is `cat(block_0_shard, block_1_shard, ...)`, which DTensor —
+                # and therefore DCP — records as plain chunk `rank` of the global tensor. The stored layout is
+                # thus interleaved by the saving `tp_degree`, so the checkpoint can only be read back at that
+                # same degree. Record it so a mismatch fails clearly instead of silently loading garbage.
+                with open(os.path.join(save_directory, DCP_CONFIG_NAME), "w", encoding="utf-8") as f:
+                    json.dump({"tp_degree": tp_config._tp_degree}, f, indent=2)
+            # Written from the sharded state dict, so no rank ever holds a full tensor. Collective, so every
+            # rank takes part.
+            dcp_api.save(self.state_dict(), checkpoint_id=save_directory)
+            logger.info(f"Distributed checkpoint saved in {save_directory}")
+            return
+
+        # Under tensor parallelism the parameters are DTensor shards, so they have to be gathered before
+        # anything can be written. `state_dict()` is read here rather than further down because the gather is
+        # a collective: every rank must reach it, while only rank 0 may go on to touch the filesystem or the
+        # Hub. Non-TP saves keep the original ordering.
+        state_dict = None
+        if tp_config is not None:
+            if use_flashpack:
+                raise ValueError(
+                    "`use_flashpack=True` is not supported for a tensor-parallel model. Save it with "
+                    "`safe_serialization=True`, or use `dcp=True` to write a sharded checkpoint."
+                )
+            from ..hooks.tensor_parallel import gather_tp_state_dict, resolve_tp_shard_specs
+
+            state_dict = gather_tp_state_dict(
+                self.state_dict(), resolve_tp_shard_specs(self, self._tp_plan), tp_config
+            )
+            if tp_config._mesh.get_local_rank() != 0:
+                # `is_main_process` defaults to True on every rank, so it cannot be used for this.
+                return
+
         weights_name = WEIGHTS_NAME
         if use_flashpack:
             weights_name = FLASHPACK_WEIGHTS_NAME
@@ -772,7 +855,8 @@ def save_pretrained(
             model_to_save.save_config(save_directory)
 
         # Save the model
-        state_dict = model_to_save.state_dict()
+        if state_dict is None:
+            state_dict = model_to_save.state_dict()
         quantization_metadata = {}
         if hf_quantizer is not None:
             state_dict, quantization_metadata = hf_quantizer.get_state_dict_and_metadata(
@@ -1037,7 +1121,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
         quantization_config = kwargs.pop("quantization_config", None)
         dduf_entries: dict[str, DDUFEntry] | None = kwargs.pop("dduf_entries", None)
         disable_mmap = kwargs.pop("disable_mmap", False)
-        parallel_config: ParallelConfig | ContextParallelConfig | None = kwargs.pop("parallel_config", None)
+        parallel_config: ParallelConfig | ContextParallelConfig | TensorParallelConfig | None = kwargs.pop(
+            "parallel_config", None
+        )
         use_flashpack = kwargs.pop("use_flashpack", False)
         flashpack_kwargs = kwargs.pop("flashpack_kwargs", {})
 
@@ -1150,6 +1236,35 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
         # no in-place modification of the original config.
         config = copy.deepcopy(config)
 
+        # A `torch.distributed.checkpoint` directory written by `save_pretrained(..., dcp=True)` holds
+        # `.distcp` shards rather than safetensors, so it bypasses the checkpoint-file resolution below.
+        if os.path.isdir(pretrained_model_name_or_path):
+            dcp_dir = os.path.join(pretrained_model_name_or_path, subfolder or "")
+            if os.path.isfile(os.path.join(dcp_dir, ".metadata")):
+                # Checked here rather than in `_load_dcp_checkpoint` because this branch returns before the
+                # quantizer is built and before `_check_tp_streaming_supported` runs, so nothing else would
+                # look at these.
+                unsupported = [
+                    name
+                    for name, value in (
+                        ("device_map", device_map),
+                        ("quantization_config", quantization_config),
+                        ("use_flashpack", use_flashpack),
+                        ("variant", variant),
+                        ("dduf_entries", dduf_entries),
+                        ("low_cpu_mem_usage=False", not low_cpu_mem_usage),
+                    )
+                    if value
+                ]
+                if unsupported:
+                    raise ValueError(
+                        f"{unsupported} cannot be combined with the distributed checkpoint at {dcp_dir}: its "
+                        "shards are read in place onto each rank's device."
+                    )
+                return cls._load_dcp_checkpoint(
+                    dcp_dir, config, unused_kwargs, torch_dtype=torch_dtype, parallel_config=parallel_config
+                )
+
         # determine initial quantization config.
         #######################################
         pre_quantized = "quantization_config" in config and config["quantization_config"] is not None
@@ -1204,6 +1319,27 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
         else:
             keep_in_fp32_modules = []
 
+        # A tensor-parallel `parallel_config` makes `from_pretrained` shard while it reads, so each rank only
+        # ever materializes its own slice. Validate the combination before any file is fetched.
+        tp_config = None
+        if parallel_config is not None:
+            tp_config = (
+                parallel_config
+                if isinstance(parallel_config, TensorParallelConfig)
+                else parallel_config.tensor_parallel_config
+            )
+            if tp_config is not None and tp_config.tp_degree == 1 and tp_config.mesh is None:
+                # Nothing to shard, so take the ordinary loader rather than building 1-rank DTensors.
+                tp_config = None
+        if tp_config is not None:
+            cls._check_tp_streaming_supported(
+                device_map=device_map,
+                low_cpu_mem_usage=low_cpu_mem_usage,
+                use_flashpack=use_flashpack,
+                hf_quantizer=hf_quantizer,
+                dduf_entries=dduf_entries,
+            )
+
         is_sharded = False
         resolved_model_file = None
 
@@ -1320,6 +1456,27 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
         with ContextManagers(init_contexts):
             model = cls.from_config(config, **unused_kwargs)
 
+        # Resolve the tensor-parallel mesh before any weights are read, so each rank can stream only its own
+        # slice of every planned parameter straight into a DTensor instead of materializing the full
+        # checkpoint and resharding it afterwards.
+        tp_shard_specs = None
+        if tp_config is not None:
+            from ..hooks.tensor_parallel import resolve_tp_shard_specs
+
+            non_safetensors = [f for f in resolved_model_file if not str(f).endswith(".safetensors")]
+            if non_safetensors:
+                raise ValueError(
+                    f"A tensor-parallel `parallel_config` requires safetensors weights, so that each rank can "
+                    f"read only its own slice of each tensor. Got {non_safetensors}."
+                )
+
+            parallel_config = model._resolve_parallel_config(parallel_config)
+            tp_config = parallel_config.tensor_parallel_config
+            tp_shard_specs = resolve_tp_shard_specs(model, cls._tp_plan)
+            # Each rank opens every shard file but only reads its own slices, so threading the files buys
+            # nothing and would have several threads calling `register_parameter` on the same modules.
+            is_parallel_loading_enabled = False
+
         if use_flashpack:
             if is_flashpack_available():
                 import flashpack
@@ -1362,7 +1519,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
             torch.set_default_dtype(dtype_orig)
 
         state_dict = None
-        if not is_sharded:
+        if not is_sharded and tp_shard_specs is None:
             # Time to load the checkpoint
             state_dict = load_state_dict(resolved_model_file[0], disable_mmap=disable_mmap, dduf_entries=dduf_entries)
             # We only fix it for non sharded checkpoints as we don't need it yet for sharded one.
@@ -1370,6 +1527,13 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
 
         if is_sharded:
             loaded_keys = sharded_metadata["all_checkpoint_keys"]
+        elif tp_shard_specs is not None:
+            # Read the key names out of the safetensors header without materializing any tensor, and leave
+            # `state_dict` as None so `_load_pretrained_model` keeps reading from the file itself.
+            from safetensors import safe_open
+
+            with safe_open(resolved_model_file[0], framework="pt") as f:
+                loaded_keys = list(f.keys())
         else:
             loaded_keys = list(state_dict.keys())
 
@@ -1418,6 +1582,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
             dduf_entries=dduf_entries,
             is_parallel_loading_enabled=is_parallel_loading_enabled,
             disable_mmap=disable_mmap,
+            tp_shard_specs=tp_shard_specs,
+            tp_config=tp_config,
         )
         loading_info = {
             "missing_keys": missing_keys,
@@ -1457,7 +1623,13 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
         # Set model in evaluation mode to deactivate DropOut modules by default
         model.eval()
 
-        if parallel_config is not None:
+        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
+
+            apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True)
+        elif parallel_config is not None:
             model.enable_parallelism(config=parallel_config)
 
         if output_loading_info:
@@ -1604,25 +1776,182 @@ def compile_repeated_blocks(self, *args, **kwargs):
                 f"Regional compilation failed because {repeated_blocks} classes are not found in the model. "
             )
 
-    def enable_parallelism(
-        self,
+    @classmethod
+    def _load_dcp_checkpoint(
+        cls,
+        checkpoint_dir: str,
+        config: dict,
+        unused_kwargs: dict,
         *,
-        config: ParallelConfig | ContextParallelConfig | TensorParallelConfig,
-        cp_plan: dict[str, ContextParallelModelPlan] | None = None,
+        torch_dtype: torch.dtype | None,
+        parallel_config: ParallelConfig | ContextParallelConfig | TensorParallelConfig | None,
     ):
-        logger.warning(
-            "`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning."
-        )
+        """Load a `torch.distributed.checkpoint` directory written by `save_pretrained(..., dcp=True)`.
+
+        The shards are those of a tensor-parallel model, so a tensor-parallel `parallel_config` is required, at the
+        `tp_degree` the checkpoint was written with — see the note where it is written. Use the ordinary safetensors
+        path to move a model between degrees; it streams each rank's slice, so it costs no more memory than this
+        does.
+
+        DCP loads **in place**, so every parameter has to be allocated first with its local shape and on the device it
+        will end up on.
+        """
+        import torch.distributed.checkpoint as dcp
+        from torch.distributed.tensor import DTensor, Replicate, Shard
+
+        from ..hooks.tensor_parallel import apply_tensor_parallel, resolve_tp_shard_specs
+
+        with open(os.path.join(checkpoint_dir, DCP_CONFIG_NAME), encoding="utf-8") as f:
+            saved_tp_degree = json.load(f)["tp_degree"]
+
+        with ContextManagers([no_init_weights(), accelerate.init_empty_weights()]):
+            model = cls.from_config(config, **unused_kwargs)
+
+        tp_config = None
+        if parallel_config is not None:
+            tp_config = (
+                parallel_config
+                if isinstance(parallel_config, TensorParallelConfig)
+                else parallel_config.tensor_parallel_config
+            )
+        if tp_config is None:
+            raise ValueError(
+                f"The distributed checkpoint at {checkpoint_dir} holds the shards of a tensor-parallel model, so "
+                f"it can only be read back with a tensor-parallel `parallel_config` of `tp_degree="
+                f"{saved_tp_degree}`. To load it without tensor parallelism, re-save the model with "
+                f"`save_pretrained(...)`, which gathers the shards into ordinary safetensors."
+            )
+        # An explicit `mesh` overrides `tp_degree` (see `TensorParallelConfig`), and `_tp_degree` is only set
+        # by `setup()`, which has not run yet — so the effective degree has to be resolved by hand here.
+        requested_tp_degree = tp_config.mesh.size() if tp_config.mesh is not None else tp_config.tp_degree
+        if requested_tp_degree != saved_tp_degree:
+            raise ValueError(
+                f"The distributed checkpoint at {checkpoint_dir} was written with `tp_degree={saved_tp_degree}` "
+                f"and can only be loaded with the same degree, but {requested_tp_degree} was requested. Packed "
+                f"projections are stored interleaved by the writing degree, so reading at another degree would "
+                f"silently produce wrong weights. To change degree, re-save the model with "
+                f"`save_pretrained(...)` (which gathers to ordinary safetensors) and load that with "
+                f"`from_pretrained(..., parallel_config=...)`."
+            )
+        parallel_config = model._resolve_parallel_config(parallel_config)
+        tp_config = parallel_config.tensor_parallel_config
+        tp_shard_specs = resolve_tp_shard_specs(model, cls._tp_plan)
+        tp_mesh = tp_config._mesh
+        device = torch.neuron.current_device() if tp_mesh.device_type == "neuron" else tp_config._device
+
+        for name, meta_param in model.state_dict().items():
+            dtype = torch_dtype if torch_dtype is not None and meta_param.is_floating_point() else meta_param.dtype
+            spec = tp_shard_specs.get(name)
+            if spec is None or spec.dim is None:
+                local = torch.empty(meta_param.shape, dtype=dtype, device=device)
+            else:
+                shape = list(meta_param.shape)
+                shape[spec.dim] //= tp_config._tp_degree
+                local = torch.empty(shape, dtype=dtype, device=device)
+
+            module_path, _, param_name = name.rpartition(".")
+            module = model.get_submodule(module_path) if module_path else model
+            if spec is None:
+                value = local
+            else:
+                placement = Replicate() if spec.dim is None else Shard(spec.dim)
+                value = DTensor.from_local(local, tp_mesh, [placement], run_check=False)
+            if param_name in module._buffers:
+                module._buffers[param_name] = value
+            else:
+                module.register_parameter(param_name, torch.nn.Parameter(value, requires_grad=False))
 
-        if not torch.distributed.is_available() and not torch.distributed.is_initialized():
+        state_dict = model.state_dict()
+        dcp.load(state_dict, checkpoint_id=checkpoint_dir)
+
+        # `dcp.load` silently does nothing for a parameter left on `meta`, so a mistake above would
+        # otherwise produce a model of uninitialized weights with no diagnostic at all.
+        still_meta = sorted(name for name, value in state_dict.items() if value.device.type == "meta")
+        if still_meta:
             raise RuntimeError(
-                "torch.distributed must be available and initialized before calling `enable_parallelism`."
+                f"Loading the distributed checkpoint at {checkpoint_dir} left these parameters on the meta "
+                f"device: {still_meta}."
             )
 
-        from ..hooks.context_parallel import apply_context_parallel
-        from .attention import AttentionModuleMixin
-        from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
-        from .attention_processor import Attention, MochiAttention
+        # Non-persistent buffers are absent from both the state dict and the checkpoint, and
+        # `init_empty_weights` leaves them as real CPU tensors, so move them across explicitly.
+        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)
+
+        model.register_to_config(_name_or_path=checkpoint_dir)
+        model.eval()
+
+        apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True)
+
+        return model
+
+    @classmethod
+    def _check_tp_streaming_supported(
+        cls,
+        *,
+        device_map,
+        low_cpu_mem_usage: bool,
+        use_flashpack: bool,
+        hf_quantizer,
+        dduf_entries,
+    ) -> None:
+        """Reject the `from_pretrained` options that cannot be combined with a tensor-parallel load.
+
+        Sharding on load needs a meta-initialized model and lazily sliceable safetensors files. Rather than silently
+        falling back to loading the full checkpoint and resharding it — which would quietly give up the memory saving
+        that is the whole point — each unsupported combination raises.
+
+        Called before the checkpoint files are resolved, so that e.g. `use_flashpack` fails with the real reason
+        instead of a missing-file error. The weights-format check lives at the point where the resolved file list is
+        known.
+        """
+        if cls._tp_plan is None:
+            raise ValueError(
+                f"`_tp_plan` must be set on the model class to use tensor parallelism. "
+                f"'{cls.__name__}' does not define one."
+            )
+        if device_map is not None:
+            raise ValueError(
+                "`device_map` cannot be combined with a tensor-parallel `parallel_config`: tensor parallelism "
+                "already places each rank's shard on that rank's device. Drop `device_map`."
+            )
+        if hf_quantizer is not None:
+            raise ValueError(
+                "`quantization_config` cannot be combined with a tensor-parallel `parallel_config`. Load the "
+                "model unquantized, or shard it after loading with `enable_parallelism`."
+            )
+        if not low_cpu_mem_usage:
+            raise ValueError(
+                "`low_cpu_mem_usage=False` cannot be combined with a tensor-parallel `parallel_config`: "
+                "streaming each rank's shard requires the model to be initialized on the meta device."
+            )
+        if use_flashpack:
+            raise ValueError(
+                "`use_flashpack=True` cannot be combined with a tensor-parallel `parallel_config`; FlashPack "
+                "checkpoints cannot be sliced per rank."
+            )
+        if dduf_entries:
+            raise ValueError(
+                "DDUF checkpoints cannot be combined with a tensor-parallel `parallel_config`; their tensors "
+                "cannot be sliced per rank."
+            )
+
+    def _resolve_parallel_config(
+        self, config: ParallelConfig | ContextParallelConfig | TensorParallelConfig
+    ) -> ParallelConfig:
+        """Normalize `config`, build its device mesh, and record it on the model.
+
+        Split out of `enable_parallelism` because `from_pretrained` needs the mesh *before* it reads any weights, in
+        order to stream each rank's shard straight into place. Whichever of the two runs first builds the mesh exactly
+        once.
+        """
+        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+            raise RuntimeError(
+                "torch.distributed must be available and initialized before applying a `parallel_config`."
+            )
 
         if isinstance(config, ContextParallelConfig):
             config = ParallelConfig(context_parallel_config=config)
@@ -1635,6 +1964,51 @@ def enable_parallelism(
         device_module = torch.get_device_module(device_type)
         device = torch.device(device_type, rank % device_module.device_count())
 
+        mesh = None
+        if config.context_parallel_config is not None:
+            cp_config = config.context_parallel_config
+            mesh = cp_config.mesh or torch.distributed.device_mesh.init_device_mesh(
+                device_type=device_type,
+                mesh_shape=cp_config.mesh_shape,
+                mesh_dim_names=cp_config.mesh_dim_names,
+            )
+        elif config.tensor_parallel_config is not None:
+            tp_config = config.tensor_parallel_config
+            mesh = tp_config.mesh or torch.distributed.device_mesh.init_device_mesh(
+                device_type=device_type,
+                mesh_shape=(tp_config.tp_degree,),
+                mesh_dim_names=("tp",),
+            )
+
+        # `config.setup()` records the mesh resolved above onto the config; see `ParallelConfig.setup`.
+        config.setup(rank, world_size, device, mesh=mesh)
+        self._parallel_config = config
+        return config
+
+    def enable_parallelism(
+        self,
+        *,
+        config: ParallelConfig | ContextParallelConfig | TensorParallelConfig,
+        cp_plan: dict[str, ContextParallelModelPlan] | None = None,
+    ):
+        logger.warning(
+            "`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning."
+        )
+
+        from ..hooks.context_parallel import apply_context_parallel
+        from .attention import AttentionModuleMixin
+        from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
+        from .attention_processor import Attention, MochiAttention
+
+        if self._parallel_config is not None:
+            raise RuntimeError(
+                f"Parallelism is already applied to this {self.__class__.__name__}. `enable_parallelism` cannot be "
+                "called twice, and it must not be called on a model loaded with `from_pretrained(..., "
+                "parallel_config=...)` — that already sharded the weights while reading the checkpoint."
+            )
+
+        config = self._resolve_parallel_config(config)
+
         attention_classes = (Attention, MochiAttention, AttentionModuleMixin)
 
         if config.context_parallel_config is not None:
@@ -1665,26 +2039,6 @@ def enable_parallelism(
                 # iterate over all modules after checking the first processor
                 break
 
-        mesh = None
-        if config.context_parallel_config is not None:
-            cp_config = config.context_parallel_config
-            mesh = cp_config.mesh or torch.distributed.device_mesh.init_device_mesh(
-                device_type=device_type,
-                mesh_shape=cp_config.mesh_shape,
-                mesh_dim_names=cp_config.mesh_dim_names,
-            )
-        elif config.tensor_parallel_config is not None:
-            tp_config = config.tensor_parallel_config
-            mesh = tp_config.mesh or torch.distributed.device_mesh.init_device_mesh(
-                device_type=device_type,
-                mesh_shape=(tp_config.tp_degree,),
-                mesh_dim_names=("tp",),
-            )
-
-        # `config.setup()` records the mesh resolved above onto the config; see `ParallelConfig.setup`.
-        config.setup(rank, world_size, device, mesh=mesh)
-        self._parallel_config = config
-
         # Only context parallelism needs the config inside attention: it replaces the attention computation itself
         # (Ulysses all-to-all / ring). Tensor parallelism only shards `Linear` weights, so each rank runs the ordinary
         # attention op over its own heads and the processors must stay unaware of it.
@@ -1705,16 +2059,6 @@ def enable_parallelism(
             apply_context_parallel(self, config.context_parallel_config, cp_plan)
 
         if config.tensor_parallel_config is not None:
-            if self._tp_plan is None:
-                raise ValueError(
-                    "`_tp_plan` must be set on the model class to use tensor parallelism. "
-                    f"'{self.__class__.__name__}' does not define one."
-                )
-            tp_degree = config.tensor_parallel_config._tp_degree
-            num_heads = getattr(self.config, "num_attention_heads", None)
-            if num_heads is not None and num_heads % tp_degree != 0:
-                raise ValueError(f"`tp_degree` ({tp_degree}) must divide the number of attention heads ({num_heads}).")
-
             from ..hooks.tensor_parallel import apply_tensor_parallel
 
             apply_tensor_parallel(self, config.tensor_parallel_config, self._tp_plan)
@@ -1739,6 +2083,8 @@ def _load_pretrained_model(
         dduf_entries: dict[str, DDUFEntry] | None = None,
         is_parallel_loading_enabled: bool | None = False,
         disable_mmap: bool = False,
+        tp_shard_specs: dict | None = None,
+        tp_config: TensorParallelConfig | None = None,
     ):
         model_state_dict = model.state_dict()
         expected_keys = list(model_state_dict.keys())
@@ -1755,6 +2101,17 @@ def _load_pretrained_model(
         mismatched_keys = []
         error_msgs = []
 
+        if tp_shard_specs is not None:
+            # `_hooks_only_styles` lets `parallelize_module` broadcast any planned parameter it finds still
+            # plain, and for a key the checkpoint does not carry that broadcast would be issued on a `meta`
+            # tensor.
+            missing_planned_keys = sorted(set(tp_shard_specs) & set(missing_keys))
+            if missing_planned_keys:
+                raise ValueError(
+                    f"Cannot shard {cls.__name__} across tensor-parallel ranks because its `_tp_plan` covers "
+                    f"parameters that the checkpoint does not contain: {missing_planned_keys}."
+                )
+
         # Deal with offload
         if device_map is not None and "disk" in device_map.values():
             if offload_folder is None:
@@ -1790,25 +2147,38 @@ def _load_pretrained_model(
             resolved_model_file = [state_dict]
 
         # Prepare the loading function sharing the attributes shared between them.
-        load_fn = functools.partial(
-            _load_shard_files_with_threadpool if is_parallel_loading_enabled else _load_shard_file,
-            model=model,
-            model_state_dict=model_state_dict,
-            device_map=device_map,
-            dtype=dtype,
-            hf_quantizer=hf_quantizer,
-            keep_in_fp32_modules=keep_in_fp32_modules,
-            dduf_entries=dduf_entries,
-            loaded_keys=loaded_keys,
-            unexpected_keys=unexpected_keys,
-            offload_index=offload_index,
-            offload_folder=offload_folder,
-            state_dict_index=state_dict_index,
-            state_dict_folder=state_dict_folder,
-            ignore_mismatched_sizes=ignore_mismatched_sizes,
-            low_cpu_mem_usage=low_cpu_mem_usage,
-            disable_mmap=disable_mmap,
-        )
+        if tp_shard_specs is not None:
+            load_fn = functools.partial(
+                _load_shard_file_tp,
+                model=model,
+                model_state_dict=model_state_dict,
+                tp_shard_specs=tp_shard_specs,
+                tp_config=tp_config,
+                dtype=dtype,
+                keep_in_fp32_modules=keep_in_fp32_modules,
+                unexpected_keys=unexpected_keys,
+                ignore_mismatched_sizes=ignore_mismatched_sizes,
+            )
+        else:
+            load_fn = functools.partial(
+                _load_shard_files_with_threadpool if is_parallel_loading_enabled else _load_shard_file,
+                model=model,
+                model_state_dict=model_state_dict,
+                device_map=device_map,
+                dtype=dtype,
+                hf_quantizer=hf_quantizer,
+                keep_in_fp32_modules=keep_in_fp32_modules,
+                dduf_entries=dduf_entries,
+                loaded_keys=loaded_keys,
+                unexpected_keys=unexpected_keys,
+                offload_index=offload_index,
+                offload_folder=offload_folder,
+                state_dict_index=state_dict_index,
+                state_dict_folder=state_dict_folder,
+                ignore_mismatched_sizes=ignore_mismatched_sizes,
+                low_cpu_mem_usage=low_cpu_mem_usage,
+                disable_mmap=disable_mmap,
+            )
 
         if is_parallel_loading_enabled:
             offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(resolved_model_file)
diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py
index 0554d341022a..97ddb8a2589c 100644
--- a/src/diffusers/utils/__init__.py
+++ b/src/diffusers/utils/__init__.py
@@ -20,6 +20,7 @@
 from .. import __version__
 from .constants import (
     CONFIG_NAME,
+    DCP_CONFIG_NAME,
     DEFAULT_HF_PARALLEL_LOADING_WORKERS,
     DEPRECATED_REVISION_ARGS,
     DIFFUSERS_DYNAMIC_MODULE_NAME,
diff --git a/src/diffusers/utils/constants.py b/src/diffusers/utils/constants.py
index fcf0e4518800..597ccd3eebd1 100644
--- a/src/diffusers/utils/constants.py
+++ b/src/diffusers/utils/constants.py
@@ -35,6 +35,7 @@
 SAFETENSORS_FILE_EXTENSION = "safetensors"
 FLASHPACK_WEIGHTS_NAME = "model.flashpack"
 FLASHPACK_FILE_EXTENSION = "flashpack"
+DCP_CONFIG_NAME = "dcp_config.json"
 GGUF_FILE_EXTENSION = "gguf"
 ONNX_EXTERNAL_WEIGHTS_NAME = "weights.pb"
 HUGGINGFACE_CO_RESOLVE_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://huggingface.co")
diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py
index 63575abf6b7b..c5174c396561 100644
--- a/tests/models/testing_utils/parallelism.py
+++ b/tests/models/testing_utils/parallelism.py
@@ -20,9 +20,11 @@
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
+from safetensors.torch import load_file
 
 from diffusers.models._modeling_parallel import ContextParallelConfig, TensorParallelConfig
 from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
+from diffusers.utils.constants import SAFETENSORS_WEIGHTS_NAME
 
 from ...testing_utils import (
     is_attention,
@@ -295,6 +297,108 @@ def _tensor_parallel_worker(
             dist.destroy_process_group()
 
 
+def _tensor_parallel_from_pretrained_worker(
+    rank, world_size, master_port, model_class, checkpoint_dir, resave_dir, inputs_dict, return_dict
+):
+    """Worker for `from_pretrained(..., parallel_config=...)`, i.e. sharding while reading the checkpoint.
+
+    Each rank loads only its own slice of every `_tp_plan` parameter straight into a `DTensor`, runs a forward
+    pass, and (if `resave_dir` is given) saves the model back out, which has to gather the shards first. Rank
+    0 reports its output and the local/global shapes of one sharded weight so the caller can check both the
+    numerics and that sharding actually happened.
+    """
+    try:
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(master_port)
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+
+        device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
+        dist.init_process_group(backend=device_config["backend"], rank=rank, world_size=world_size)
+        device_config["module"].set_device(rank)
+
+        from torch.distributed.tensor import DTensor
+
+        model = model_class.from_pretrained(
+            checkpoint_dir, parallel_config=TensorParallelConfig(tp_degree=world_size)
+        ).eval()
+
+        device = torch.device(f"{torch_device}:{rank}")
+        inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+        with torch.no_grad():
+            output = model(**inputs_on_device, return_dict=False)[0]
+        if isinstance(output, DTensor):
+            output = output.full_tensor()
+
+        # Gathering is a collective, so every rank has to reach this even though only rank 0 writes.
+        if resave_dir is not None:
+            model.save_pretrained(resave_dir)
+
+        if rank == 0:
+            sharded = {k: v for k, v in model.state_dict().items() if isinstance(v, DTensor)}
+            assert sharded, "No parameter was sharded into a DTensor by the streaming load."
+            name, param = next(iter(sharded.items()))
+            return_dict["status"] = "success"
+            return_dict["num_sharded"] = len(sharded)
+            return_dict["shard_example"] = (name, list(param.to_local().shape), list(param.shape))
+            return_dict["output"] = output.float().cpu().tolist()
+
+    except Exception as e:
+        if rank == 0:
+            return_dict["status"] = "error"
+            return_dict["error"] = f"{type(e).__name__}: {e}"
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+def _tensor_parallel_dcp_worker(
+    rank, world_size, master_port, model_class, checkpoint_dir, dcp_dir, inputs_dict, return_dict
+):
+    """Worker for the `save_pretrained(..., dcp=True)` round trip.
+
+    Streams the checkpoint into shards, writes them as a distributed checkpoint (no rank ever holding a full
+    tensor), then loads that back and runs a forward pass. Rank 0 reports the output.
+    """
+    try:
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(master_port)
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+
+        device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
+        dist.init_process_group(backend=device_config["backend"], rank=rank, world_size=world_size)
+        device_config["module"].set_device(rank)
+
+        from torch.distributed.tensor import DTensor
+
+        tp_config = TensorParallelConfig(tp_degree=world_size)
+        model_class.from_pretrained(checkpoint_dir, parallel_config=tp_config).save_pretrained(dcp_dir, dcp=True)
+
+        reloaded = model_class.from_pretrained(
+            dcp_dir, parallel_config=TensorParallelConfig(tp_degree=world_size)
+        ).eval()
+
+        device = torch.device(f"{torch_device}:{rank}")
+        inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+        with torch.no_grad():
+            output = reloaded(**inputs_on_device, return_dict=False)[0]
+        if isinstance(output, DTensor):
+            output = output.full_tensor()
+
+        if rank == 0:
+            return_dict["status"] = "success"
+            return_dict["output"] = output.float().cpu().tolist()
+
+    except Exception as e:
+        if rank == 0:
+            return_dict["status"] = "error"
+            return_dict["error"] = f"{type(e).__name__}: {e}"
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
 @is_tensor_parallel
 @require_torch_multi_accelerator
 class TensorParallelTesterMixin:
@@ -344,6 +448,104 @@ def test_tensor_parallel_inference(self, batch_size: int = 1):
     def test_tensor_parallel_batch_inputs(self):
         self.test_tensor_parallel_inference(batch_size=2)
 
+    def _tp_checkpoint_and_reference(self, tmp_path, world_size):
+        """Write a checkpoint for the sharded loaders to read, and record its single-device output.
+
+        Returns `(checkpoint_dir, cpu_inputs, reference_output)`, or skips when the model cannot be sharded
+        across `world_size` ranks.
+        """
+        if not torch.distributed.is_available():
+            pytest.skip("torch.distributed is not available.")
+        if getattr(self.model_class, "_tp_plan", None) is None:
+            pytest.skip("Model does not define a `_tp_plan` for tensor parallel inference.")
+
+        init_dict = self.get_init_dict()
+        num_heads = init_dict.get("num_attention_heads")
+        if num_heads is not None and num_heads % world_size != 0:
+            pytest.skip(f"`num_attention_heads` ({num_heads}) is not divisible by tp_degree ({world_size}).")
+
+        inputs_dict = self.get_dummy_inputs()
+        model = self.model_class(**init_dict).eval().to(torch_device)
+        with torch.no_grad():
+            reference = model(**inputs_dict, return_dict=False)[0].float().cpu()
+
+        checkpoint_dir = str(tmp_path / "checkpoint")
+        model.save_pretrained(checkpoint_dir)
+
+        inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}
+        return checkpoint_dir, inputs_dict, reference
+
+    def test_tensor_parallel_from_pretrained(self, tmp_path):
+        """`from_pretrained(..., parallel_config=...)` shards while reading, and `save_pretrained` gathers back.
+
+        Covers both directions in one spawn: the streaming load must match the single-device reference, and the
+        checkpoint it writes back out must be byte-identical to the one it read. The round trip is what catches
+        a wrong packed-projection reorder — a plain colwise/rowwise mistake would pass the forward check alone.
+        """
+        world_size = 2
+        checkpoint_dir, inputs_dict, reference = self._tp_checkpoint_and_reference(tmp_path, world_size)
+        resave_dir = str(tmp_path / "resaved")
+
+        manager = mp.Manager()
+        return_dict = manager.dict()
+        mp.spawn(
+            _tensor_parallel_from_pretrained_worker,
+            args=(
+                world_size,
+                _find_free_port(),
+                self.model_class,
+                checkpoint_dir,
+                resave_dir,
+                inputs_dict,
+                return_dict,
+            ),
+            nprocs=world_size,
+            join=True,
+        )
+        assert return_dict.get("status") == "success", (
+            f"Tensor parallel `from_pretrained` failed: {return_dict.get('error', 'Unknown error')}"
+        )
+
+        name, local_shape, global_shape = return_dict["shard_example"]
+        assert local_shape != global_shape, (
+            f"'{name}' has local shape {local_shape} equal to its global shape, so it was not sharded."
+        )
+
+        # Sharded matmuls + all-reduce reorder the summation, so allow a small tolerance over the reference.
+        torch.testing.assert_close(reference, torch.tensor(return_dict["output"]), atol=1e-3, rtol=1e-3)
+
+        original = load_file(os.path.join(checkpoint_dir, SAFETENSORS_WEIGHTS_NAME))
+        resaved = load_file(os.path.join(resave_dir, SAFETENSORS_WEIGHTS_NAME))
+        assert original.keys() == resaved.keys()
+        for key, value in original.items():
+            torch.testing.assert_close(resaved[key], value, atol=0, rtol=0, msg=lambda m, key=key: f"{key}: {m}")
+
+    def test_tensor_parallel_dcp_roundtrip(self, tmp_path):
+        """`save_pretrained(..., dcp=True)` writes sharded and `from_pretrained` reads it back at the same degree."""
+        world_size = 2
+        checkpoint_dir, inputs_dict, reference = self._tp_checkpoint_and_reference(tmp_path, world_size)
+
+        manager = mp.Manager()
+        return_dict = manager.dict()
+        mp.spawn(
+            _tensor_parallel_dcp_worker,
+            args=(
+                world_size,
+                _find_free_port(),
+                self.model_class,
+                checkpoint_dir,
+                str(tmp_path / "dcp"),
+                inputs_dict,
+                return_dict,
+            ),
+            nprocs=world_size,
+            join=True,
+        )
+        assert return_dict.get("status") == "success", (
+            f"Tensor parallel DCP round trip failed: {return_dict.get('error', 'Unknown error')}"
+        )
+        torch.testing.assert_close(reference, torch.tensor(return_dict["output"]), atol=1e-3, rtol=1e-3)
+
 
 @is_context_parallel
 @require_torch_multi_accelerator

From 40ddb53dc0a7faa6bd43ba3bcfee39c76facec5d Mon Sep 17 00:00:00 2001
From: JingyaHuang <huang_jingya@outlook.com>
Date: Thu, 20 Aug 2026 22:51:19 +0000
Subject: [PATCH 2/5] Raise when tensor parallelism is combined with
 quantization, offloading or LoRA
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Addresses the remaining two items of the review on #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.
---
 .../en/training/distributed_inference.md      |   2 +-
 src/diffusers/hooks/tensor_parallel.py        |  55 +++++-
 src/diffusers/loaders/peft.py                 |  10 ++
 src/diffusers/models/modeling_utils.py        |  43 ++++-
 src/diffusers/pipelines/pipeline_utils.py     |  25 +++
 tests/models/test_parallelism_guards.py       | 170 ++++++++++++++++++
 6 files changed, 295 insertions(+), 10 deletions(-)
 create mode 100644 tests/models/test_parallelism_guards.py

diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md
index 25d4dba4e339..20618280c841 100644
--- a/docs/source/en/training/distributed_inference.md
+++ b/docs/source/en/training/distributed_inference.md
@@ -483,7 +483,7 @@ torchrun --nproc-per-node 4 tensor_parallel_flux.py
 
 `tp_degree` is taken from `world_size` above, so `--nproc-per-node 4` shards the transformer across 4 devices.
 
-A tensor-parallel `parallel_config` cannot be combined with `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True`, DDUF checkpoints, or non-safetensors weights; each raises rather than quietly falling back to loading the full checkpoint. To shard a model that is already in memory, call [`~ModelMixin.enable_parallelism`] with the same config instead — that loads everything first and reshards it, so it costs full checkpoint memory on every rank.
+A tensor-parallel `parallel_config` cannot be combined with `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True`, DDUF checkpoints, or non-safetensors weights; each raises rather than quietly falling back to loading the full checkpoint. Tensor parallelism also cannot be combined with quantization, offloading, or LoRA adapters at all — the parameters it shards have to be plain parameters owned by the model — so those raise however the model is sharded. To shard a model that is already in memory, call [`~ModelMixin.enable_parallelism`] with the same config instead — that loads everything first and reshards it, so it costs full checkpoint memory on every rank.
 
 ### Saving a tensor-parallel model
 
diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py
index d3cbd82d7980..f3856b438100 100644
--- a/src/diffusers/hooks/tensor_parallel.py
+++ b/src/diffusers/hooks/tensor_parallel.py
@@ -17,7 +17,7 @@
 import torch
 
 from ..models._modeling_parallel import TensorParallelConfig
-from ..utils import get_logger
+from ..utils import get_logger, is_peft_available
 
 
 logger = get_logger(__name__)  # pylint: disable=invalid-name
@@ -401,6 +401,55 @@ def _partition_linear_fn(self, name, module, device_mesh):
     return resolved
 
 
+def _check_tp_model_state(model: torch.nn.Module) -> None:
+    """Reject a model whose parameters tensor parallelism cannot take over.
+
+    Tensor parallelism replaces every planned `weight` and `bias` with a `DTensor` shard. That only works on plain
+    parameters owned by the model itself, so a model whose parameters are quantized, held elsewhere by an offloading
+    hook, or wrapped by an adapter is rejected up front rather than failing deep inside `parallelize_module` — or,
+    worse, sharding successfully and producing wrong numbers.
+
+    `from_pretrained` rejects the same combinations earlier and with a message naming the offending argument; this is
+    the only guard on the `enable_parallelism` path, where the model already exists and only its state can be read.
+    """
+    if getattr(model, "hf_quantizer", None) is not None or getattr(model, "is_quantized", False):
+        raise ValueError(
+            f"'{model.__class__.__name__}' is quantized, which cannot be combined with tensor parallelism: its "
+            "parameters are packed into a quantizer-specific layout that cannot be sharded into `DTensor`s. Load "
+            "the model unquantized to shard it."
+        )
+
+    from .group_offloading import _is_group_offload_enabled
+
+    if _is_group_offload_enabled(model):
+        raise ValueError(
+            f"'{model.__class__.__name__}' has group offloading enabled, which cannot be combined with tensor "
+            "parallelism: both decide where a parameter lives. Tensor parallelism already keeps only one shard of "
+            "each weight per rank, so offloading is not needed on top of it."
+        )
+
+    # `device_map` dispatch and accelerate's CPU offloading both leave an `_hf_hook` on every module they placed, and
+    # the weights they offloaded are `meta` tensors that `DTensor.from_local` cannot shard.
+    if getattr(model, "hf_device_map", None) is not None or any(
+        hasattr(module, "_hf_hook") for module in model.modules()
+    ):
+        raise ValueError(
+            f"'{model.__class__.__name__}' is placed by accelerate — through `device_map` or CPU offloading — which "
+            "cannot be combined with tensor parallelism: tensor parallelism already places each rank's shard on that "
+            "rank's device. Load the model without `device_map` and without offloading to shard it."
+        )
+
+    if is_peft_available():
+        from peft.tuners.tuners_utils import BaseTunerLayer
+
+        if any(isinstance(module, BaseTunerLayer) for module in model.modules()):
+            raise ValueError(
+                f"'{model.__class__.__name__}' has adapter (LoRA) layers injected, which cannot be combined with "
+                "tensor parallelism: `_tp_plan` covers the base `Linear` layers only, so the adapter weights would "
+                "stay unsharded and the result would be wrong. Unload the adapter before sharding."
+            )
+
+
 def apply_tensor_parallel(
     model: torch.nn.Module,
     config: TensorParallelConfig,
@@ -428,6 +477,10 @@ def apply_tensor_parallel(
     if num_heads is not None and num_heads % config._tp_degree != 0:
         raise ValueError(f"`tp_degree` ({config._tp_degree}) must divide the number of attention heads ({num_heads}).")
 
+    # Before the device-type check below, so that a quantized or offloaded model reports what is actually wrong with
+    # it rather than being turned away for its device type.
+    _check_tp_model_state(model)
+
     if tp_mesh.device_type not in _SUPPORTED_TP_DEVICES:
         raise ValueError(
             f"Tensor parallelism is not supported on device type '{tp_mesh.device_type}'. Supported device types are "
diff --git a/src/diffusers/loaders/peft.py b/src/diffusers/loaders/peft.py
index b0494207f48e..0f933f5ba096 100644
--- a/src/diffusers/loaders/peft.py
+++ b/src/diffusers/loaders/peft.py
@@ -154,6 +154,16 @@ def load_lora_adapter(
 
         from ..hooks.group_offloading import _maybe_remove_and_reapply_group_offloading
 
+        parallel_config = getattr(self, "_parallel_config", None)
+        if parallel_config is not None and parallel_config.tensor_parallel_config is not None:
+            # `_tp_plan` covers the base `Linear` layers only, so the injected adapter weights would stay unsharded
+            # and the sharded base layer would be added to a full-sized adapter output.
+            raise ValueError(
+                f"Cannot load a LoRA adapter into '{self.__class__.__name__}': it is sharded with tensor "
+                f"parallelism, and the adapter layers are not covered by the model's `_tp_plan`. Load the adapter "
+                f"before sharding the model."
+            )
+
         cache_dir = kwargs.pop("cache_dir", None)
         force_download = kwargs.pop("force_download", False)
         proxies = kwargs.pop("proxies", None)
diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py
index e690fff0b058..bd4ec03727dd 100644
--- a/src/diffusers/models/modeling_utils.py
+++ b/src/diffusers/models/modeling_utils.py
@@ -575,6 +575,12 @@ def enable_group_offload(
                 "2. Or, run a forward pass with tiling disabled (can still use small dummy inputs)."
             )
             logger.warning(msg)
+        if self._parallel_config is not None and self._parallel_config.tensor_parallel_config is not None:
+            raise ValueError(
+                f"'{self.__class__.__name__}' is sharded with tensor parallelism, which cannot be combined with group "
+                "offloading: both decide where a parameter lives. Tensor parallelism already keeps only one shard of "
+                "each weight per rank, so offloading is not needed on top of it."
+            )
         if not self._supports_group_offloading:
             raise ValueError(
                 f"{self.__class__.__name__} does not support group offloading. Please make sure to set the boolean attribute "
@@ -739,6 +745,22 @@ def save_pretrained(
             return
 
         hf_quantizer = getattr(self, "hf_quantizer", None)
+
+        tp_config = None
+        if self._parallel_config is not None:
+            tp_config = self._parallel_config.tensor_parallel_config
+
+        if hf_quantizer is not None and tp_config is not None:
+            # Checked before the serializability check below, so that the reason reported is this one rather than a
+            # generic "not serializable". Neither save path can honour both: the `dcp=True` branch returns before
+            # `hf_quantizer.get_state_dict_and_metadata` runs, which would leave the shards without their
+            # quantization metadata, and the gathered path would hand the quantizer tensors that have been through a
+            # DTensor round trip. Tensor parallelism and quantization cannot be combined in the first place.
+            raise ValueError(
+                "A quantized tensor-parallel model cannot be saved: tensor parallelism and quantization cannot be "
+                "combined in the first place."
+            )
+
         if hf_quantizer is not None:
             quantization_serializable = (
                 hf_quantizer is not None
@@ -755,10 +777,6 @@ def save_pretrained(
                     " the logger on the traceback to understand the reason why the quantized model is not serializable."
                 )
 
-        tp_config = None
-        if self._parallel_config is not None:
-            tp_config = self._parallel_config.tensor_parallel_config
-
         if dcp:
             if tp_config is None:
                 raise ValueError(
@@ -1249,6 +1267,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
                     for name, value in (
                         ("device_map", device_map),
                         ("quantization_config", quantization_config),
+                        # The config's own entry, not just the kwarg: this branch returns before `pre_quantized` is
+                        # computed, so a pre-quantized checkpoint directory would otherwise load silently.
+                        ("a quantized checkpoint", config.get("quantization_config") is not None),
                         ("use_flashpack", use_flashpack),
                         ("variant", variant),
                         ("dduf_entries", dduf_entries),
@@ -1261,6 +1282,12 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
                         f"{unsupported} cannot be combined with the distributed checkpoint at {dcp_dir}: its "
                         "shards are read in place onto each rank's device."
                     )
+                if cls._tp_plan is None:
+                    raise ValueError(
+                        f"`_tp_plan` must be set on the model class to read the distributed checkpoint at "
+                        f"{dcp_dir}, whose shards are those of a tensor-parallel model. '{cls.__name__}' does not "
+                        f"define one."
+                    )
                 return cls._load_dcp_checkpoint(
                     dcp_dir, config, unused_kwargs, torch_dtype=torch_dtype, parallel_config=parallel_config
                 )
@@ -1790,8 +1817,7 @@ def _load_dcp_checkpoint(
 
         The shards are those of a tensor-parallel model, so a tensor-parallel `parallel_config` is required, at the
         `tp_degree` the checkpoint was written with — see the note where it is written. Use the ordinary safetensors
-        path to move a model between degrees; it streams each rank's slice, so it costs no more memory than this
-        does.
+        path to move a model between degrees; it streams each rank's slice, so it costs no more memory than this does.
 
         DCP loads **in place**, so every parameter has to be allocated first with its local shape and on the device it
         will end up on.
@@ -1920,8 +1946,9 @@ def _check_tp_streaming_supported(
             )
         if hf_quantizer is not None:
             raise ValueError(
-                "`quantization_config` cannot be combined with a tensor-parallel `parallel_config`. Load the "
-                "model unquantized, or shard it after loading with `enable_parallelism`."
+                "`quantization_config` cannot be combined with a tensor-parallel `parallel_config`: quantized "
+                "parameters are packed into a quantizer-specific layout that cannot be sharded into `DTensor`s. "
+                "Load the model unquantized to shard it."
             )
         if not low_cpu_mem_usage:
             raise ValueError(
diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py
index 24fe0eabfa6f..37b563f3f79e 100644
--- a/src/diffusers/pipelines/pipeline_utils.py
+++ b/src/diffusers/pipelines/pipeline_utils.py
@@ -1208,6 +1208,7 @@ def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.devi
                 automatically detect the available accelerator and use.
         """
         self._maybe_raise_error_if_group_offload_active(raise_error=True)
+        self._maybe_raise_error_if_tensor_parallel_active(raise_error=True)
 
         is_pipeline_device_mapped = self._is_pipeline_device_mapped()
         if is_pipeline_device_mapped:
@@ -1326,6 +1327,7 @@ def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch
                 automatically detect the available accelerator and use.
         """
         self._maybe_raise_error_if_group_offload_active(raise_error=True)
+        self._maybe_raise_error_if_tensor_parallel_active(raise_error=True)
 
         if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
             from accelerate import cpu_offload
@@ -2272,6 +2274,29 @@ def _maybe_raise_error_if_group_offload_active(
                 return True
         return False
 
+    def _maybe_raise_error_if_tensor_parallel_active(
+        self, raise_error: bool = False, module: torch.nn.Module | None = None
+    ) -> bool:
+        """Whether any component is sharded with tensor parallelism, which CPU offloading cannot be applied on top of.
+
+        A tensor-parallel component's parameters are `DTensor` shards tied to that rank's device and process group;
+        moving them to CPU and back, as the offload hooks do, is not supported.
+        """
+        components = self.components.values() if module is None else [module]
+        components = [component for component in components if isinstance(component, torch.nn.Module)]
+        for component in components:
+            parallel_config = getattr(component, "_parallel_config", None)
+            if parallel_config is not None and parallel_config.tensor_parallel_config is not None:
+                if raise_error:
+                    raise ValueError(
+                        f"You are trying to apply model/sequential CPU offloading to a pipeline whose "
+                        f"'{component.__class__.__name__}' is sharded with tensor parallelism. This is not supported: "
+                        f"tensor parallelism already keeps only one shard of each weight per rank, so offloading is "
+                        f"not needed on top of it."
+                    )
+                return True
+        return False
+
     def _is_pipeline_device_mapped(self):
         # We support passing `device_map="cuda"`, for example. This is helpful, in case
         # users want to pass `device_map="cpu"` when initializing a pipeline. This explicit declaration is desirable
diff --git a/tests/models/test_parallelism_guards.py b/tests/models/test_parallelism_guards.py
new file mode 100644
index 000000000000..0521925b050d
--- /dev/null
+++ b/tests/models/test_parallelism_guards.py
@@ -0,0 +1,170 @@
+# coding=utf-8
+# Copyright 2026 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Guards rejecting tensor parallelism combined with quantization, offloading, or LoRA adapters.
+
+Unlike the rest of the tensor-parallel suite in `testing_utils/parallelism.py`, these tests need neither an
+accelerator nor more than one rank: every case asserts that a call raises before any collective is issued. They run
+single-process on gloo, so they run in ordinary CI.
+"""
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models._modeling_parallel import ParallelConfig, TensorParallelConfig
+from diffusers.models.modeling_utils import ModelMixin
+
+
+class TinyTPModel(ModelMixin, ConfigMixin):
+    """Smallest model carrying a `_tp_plan`: one colwise Linear feeding one rowwise Linear."""
+
+    config_name = "config.json"
+    _tp_plan = {"linear_1": "colwise", "linear_2": "rowwise"}
+    _supports_group_offloading = True
+
+    @register_to_config
+    def __init__(self, hidden_size: int = 8, num_attention_heads: int = 2):
+        super().__init__()
+        self.linear_1 = nn.Linear(hidden_size, hidden_size)
+        self.linear_2 = nn.Linear(hidden_size, hidden_size)
+
+    def forward(self, hidden_states):
+        return self.linear_2(self.linear_1(hidden_states))
+
+
+class _SerializableQuantizer:
+    """Stand-in that passes `save_pretrained`'s serializability check, so the TP guard is what raises."""
+
+    is_serializable = True
+    supports_safetensors_serialization = True
+
+
+@pytest.fixture(scope="module")
+def gloo_process_group():
+    """A single-rank CPU process group, enough for `_resolve_parallel_config` to build a mesh."""
+    if not dist.is_available():
+        pytest.skip("torch.distributed is not available.")
+    already_initialized = dist.is_initialized()
+    if not already_initialized:
+        dist.init_process_group(backend="gloo", init_method="tcp://127.0.0.1:29591", world_size=1, rank=0)
+    yield
+    if not already_initialized:
+        dist.destroy_process_group()
+
+
+def _shard(model):
+    model.enable_parallelism(config=TensorParallelConfig(tp_degree=1))
+
+
+def _mark_as_tensor_parallel(model):
+    """Put the model in the state it would be in after sharding, without needing a real mesh."""
+    model._parallel_config = ParallelConfig(tensor_parallel_config=TensorParallelConfig(tp_degree=2))
+    return model
+
+
+class TestTensorParallelModelStateGuards:
+    """`_check_tp_model_state` — a model whose parameters TP cannot take over."""
+
+    def test_clean_model_reaches_the_device_check(self, gloo_process_group):
+        """Ordering guard: with none of the bad states, the device-type check is what rejects CPU.
+
+        This is what keeps the tests below meaningful. If `_check_tp_model_state` ran after the
+        `_SUPPORTED_TP_DEVICES` check, every case would raise the device error instead of its own.
+        """
+        with pytest.raises(ValueError, match="not supported on device type"):
+            _shard(TinyTPModel())
+
+    def test_quantized_via_hf_quantizer(self, gloo_process_group):
+        model = TinyTPModel()
+        model.hf_quantizer = object()
+        with pytest.raises(ValueError, match="is quantized"):
+            _shard(model)
+
+    def test_quantized_via_is_quantized(self, gloo_process_group):
+        model = TinyTPModel()
+        model.is_quantized = True
+        with pytest.raises(ValueError, match="is quantized"):
+            _shard(model)
+
+    def test_device_map_dispatched(self, gloo_process_group):
+        model = TinyTPModel()
+        model.hf_device_map = {"": 0}
+        with pytest.raises(ValueError, match="placed by accelerate"):
+            _shard(model)
+
+    def test_accelerate_hook_on_submodule(self, gloo_process_group):
+        model = TinyTPModel()
+        model.linear_1._hf_hook = object()
+        with pytest.raises(ValueError, match="placed by accelerate"):
+            _shard(model)
+
+    def test_group_offloaded(self, gloo_process_group, monkeypatch):
+        import diffusers.hooks.group_offloading as group_offloading
+
+        monkeypatch.setattr(group_offloading, "_is_group_offload_enabled", lambda module: True)
+        with pytest.raises(ValueError, match="group offloading enabled"):
+            _shard(TinyTPModel())
+
+    def test_peft_adapter_injected(self, gloo_process_group):
+        peft = pytest.importorskip("peft")
+
+        model = TinyTPModel()
+        peft.inject_adapter_in_model(peft.LoraConfig(r=2, target_modules=["linear_1"]), model)
+        with pytest.raises(ValueError, match=r"adapter \(LoRA\) layers injected"):
+            _shard(model)
+
+
+class TestTensorParallelReverseDirectionGuards:
+    """The other order: a model already sharded, then asked to offload or take an adapter."""
+
+    def test_enable_group_offload_on_tp_model(self):
+        model = _mark_as_tensor_parallel(TinyTPModel())
+        with pytest.raises(ValueError, match="sharded with tensor parallelism"):
+            model.enable_group_offload(onload_device=torch.device("cpu"))
+
+    def test_pipeline_offload_helper_detects_tp_component(self):
+        from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+        model = _mark_as_tensor_parallel(TinyTPModel())
+        # The helper takes an explicit module, so it runs without building a whole pipeline.
+        with pytest.raises(ValueError, match="sharded with tensor parallelism"):
+            DiffusionPipeline._maybe_raise_error_if_tensor_parallel_active(
+                DiffusionPipeline, raise_error=True, module=model
+            )
+
+    def test_pipeline_offload_helper_passes_for_plain_model(self):
+        from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+        assert not DiffusionPipeline._maybe_raise_error_if_tensor_parallel_active(
+            DiffusionPipeline, raise_error=True, module=TinyTPModel()
+        )
+
+
+class TestTensorParallelSaveGuards:
+    """`save_pretrained` must not write a checkpoint that silently drops quantization."""
+
+    def test_dcp_save_rejects_quantized_model(self, tmp_path):
+        model = _mark_as_tensor_parallel(TinyTPModel())
+        model.hf_quantizer = _SerializableQuantizer()
+        with pytest.raises(ValueError, match="quantized tensor-parallel model cannot be saved"):
+            model.save_pretrained(str(tmp_path / "dcp"), dcp=True)
+
+    def test_tp_save_rejects_quantized_model(self, tmp_path):
+        model = _mark_as_tensor_parallel(TinyTPModel())
+        model.hf_quantizer = _SerializableQuantizer()
+        with pytest.raises(ValueError, match="quantized tensor-parallel model cannot be saved"):
+            model.save_pretrained(str(tmp_path / "full"))

From a8956f4d31583a41eb28074b792a1da2cecc4077 Mon Sep 17 00:00:00 2001
From: JingyaHuang <huang_jingya@outlook.com>
Date: Fri, 21 Aug 2026 11:08:58 +0000
Subject: [PATCH 3/5] Add tensor-parallel support for MiniMax-H3
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 .../transformers/transformer_minimax_h3.py    | 37 ++++++++-
 .../test_models_transformer_minimax_h3.py     | 83 +++++++++++++++----
 2 files changed, 100 insertions(+), 20 deletions(-)

diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py
index f49cdaca2eb6..41269215cbf3 100644
--- a/src/diffusers/models/transformers/transformer_minimax_h3.py
+++ b/src/diffusers/models/transformers/transformer_minimax_h3.py
@@ -19,6 +19,7 @@
 import torch.nn as nn
 
 from ...configuration_utils import ConfigMixin, register_to_config
+from ...hooks.tensor_parallel import PackedColwiseParallel
 from ...loaders import PeftAdapterMixin
 from ...utils import BaseOutput, apply_lora_scale, logging
 from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
@@ -179,9 +180,11 @@ def __call__(
             key = attn.to_k(hidden_states)
             value = attn.to_v(hidden_states)
 
-        query = query.unflatten(-1, (attn.heads, -1))
-        key = key.unflatten(-1, (attn.heads, -1))
-        value = value.unflatten(-1, (attn.heads, -1))
+        # Reshape by a fixed `head_dim` and let `-1` absorb the head count. Under tensor parallelism each rank
+        # holds a column-sharded slice (`attn.heads // tp_degree` heads); this keeps the processor TP-agnostic.
+        query = query.unflatten(-1, (-1, attn.head_dim))
+        key = key.unflatten(-1, (-1, attn.head_dim))
+        value = value.unflatten(-1, (-1, attn.head_dim))
 
         query = attn.norm_q(query)
         key = attn.norm_k(key)
@@ -449,6 +452,34 @@ class MiniMaxH3Transformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftA
         "audio_proj_out",
         "rope",
     ]
+    # Tensor-parallel plan: how each block's Linears shard across the TP mesh. Q/K/V and the attention output
+    # are unfused, so they are plain "colwise"/"rowwise" (torch's ColwiseParallel / RowwiseParallel). The one
+    # packed projection is the SwiGLU input `ff.net.0.proj`, a single Linear producing `[value; gate]` in equal
+    # halves, hence PackedColwiseParallel([1, 1]) so each half is sharded independently.
+    #
+    # Intentionally absent, i.e. replicated on every rank: the RMSNorms (`norm1`, `norm2`,
+    # `token_refiner.final_norm`, `norm_out.norm`); the QK-norms, which apply over `head_dim` after the heads are
+    # already split; the AdaLN modulation (`adaln_proj.linear`, `norm_out.linear`), which indexes the full hidden
+    # dim; and the patch/text embedders and the two output heads.
+    #
+    # `attn.to_qkv` is deliberately not listed: it only exists after `fuse_projections()`, and the plan is
+    # resolved by attribute lookup, so an unconditional entry would break the ordinary unfused model.
+    _tp_plan = {
+        # denoiser block stack
+        "transformer_blocks.*.attn.to_q": "colwise",
+        "transformer_blocks.*.attn.to_k": "colwise",
+        "transformer_blocks.*.attn.to_v": "colwise",
+        "transformer_blocks.*.attn.to_out.0": "rowwise",
+        "transformer_blocks.*.ff.net.0.proj": PackedColwiseParallel([1, 1]),
+        "transformer_blocks.*.ff.net.2": "rowwise",
+        # the token-refiner blocks are the same attention + SwiGLU FFN, minus AdaLN and rotary
+        "token_refiner.refiner_blocks.*.attn.to_q": "colwise",
+        "token_refiner.refiner_blocks.*.attn.to_k": "colwise",
+        "token_refiner.refiner_blocks.*.attn.to_v": "colwise",
+        "token_refiner.refiner_blocks.*.attn.to_out.0": "rowwise",
+        "token_refiner.refiner_blocks.*.ff.net.0.proj": PackedColwiseParallel([1, 1]),
+        "token_refiner.refiner_blocks.*.ff.net.2": "rowwise",
+    }
     # Context parallelism shards the packed sequence, so the split cannot happen on the inputs of `forward`: the rows
     # of the three modalities are scattered into the packed buffer with sequence-wide indices, which only address the
     # full sequence. The split therefore happens once the buffer is built, at the first block, and everything that is
diff --git a/tests/models/transformers/test_models_transformer_minimax_h3.py b/tests/models/transformers/test_models_transformer_minimax_h3.py
index 00baa37c84a0..e27cbd68cf21 100644
--- a/tests/models/transformers/test_models_transformer_minimax_h3.py
+++ b/tests/models/transformers/test_models_transformer_minimax_h3.py
@@ -13,13 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import subprocess
+import sys
+
 import torch
 
 from diffusers import MiniMaxH3Transformer3DModel
 from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput
 from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import enable_full_determinism, torch_device
+from ...testing_utils import enable_full_determinism, is_tensor_parallel, require_torch_neuron, torch_device
 from ..testing_utils import (
     AttentionTesterMixin,
     BaseModelTesterConfig,
@@ -27,6 +31,7 @@
     LoraTesterMixin,
     MemoryTesterMixin,
     ModelTesterMixin,
+    TensorParallelTesterMixin,
     TorchCompileTesterMixin,
     TrainingTesterMixin,
 )
@@ -84,7 +89,7 @@ def get_init_dict(self) -> dict:
             "rope_freq_dim": 2,
         }
 
-    def get_packed_layout(self, num_video_tokens: int = NUM_VIDEO_TOKENS) -> dict:
+    def get_packed_layout(self, num_video_tokens: int = NUM_VIDEO_TOKENS, device: str | torch.device = None) -> dict:
         r"""
         Build the structural arguments of one packed sequence.
 
@@ -92,29 +97,33 @@ def get_packed_layout(self, num_video_tokens: int = NUM_VIDEO_TOKENS) -> dict:
         modality and its noise level, and hands over the `(t, h, w)` grid plus the three index tensors. The layout
         here mirrors what the pipelines pack, with two distinct timesteps so the `(timestep, modality)` AdaLN table
         is addressed on more than one row.
+
+        `device` defaults to the test device; the Neuron TP spec asks for CPU because its worker builds the model on
+        CPU and moves it only after sharding.
         """
+        device = torch_device if device is None else device
         sequence_length = NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS + num_video_tokens
-        text_indices = torch.arange(NUM_TEXT_TOKENS, device=torch_device)
-        audio_indices = torch.arange(NUM_TEXT_TOKENS, NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS, device=torch_device)
-        video_indices = torch.arange(NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS, sequence_length, device=torch_device)
+        text_indices = torch.arange(NUM_TEXT_TOKENS, device=device)
+        audio_indices = torch.arange(NUM_TEXT_TOKENS, NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS, device=device)
+        video_indices = torch.arange(NUM_TEXT_TOKENS + NUM_AUDIO_TOKENS, sequence_length, device=device)
 
         # 0 = video, 1 = text, 2 = audio.
-        token_tags = torch.empty(sequence_length, dtype=torch.long, device=torch_device)
+        token_tags = torch.empty(sequence_length, dtype=torch.long, device=device)
         token_tags[text_indices] = 1
         token_tags[audio_indices] = 2
         token_tags[video_indices] = 0
 
         # The conditioning-free rows share the video timestep; the audio rows step down their own schedule.
-        timestep_indices = torch.zeros(sequence_length, dtype=torch.long, device=torch_device)
+        timestep_indices = torch.zeros(sequence_length, dtype=torch.long, device=device)
         timestep_indices[audio_indices] = 1
 
-        position_ids = torch.zeros(sequence_length, 3, dtype=torch.float32, device=torch_device)
-        position_ids[:, 0] = torch.arange(sequence_length, dtype=torch.float32, device=torch_device)
-        position_ids[video_indices, 1] = torch.arange(num_video_tokens, dtype=torch.float32, device=torch_device) % 4
-        position_ids[video_indices, 2] = torch.arange(num_video_tokens, dtype=torch.float32, device=torch_device) % 2
+        position_ids = torch.zeros(sequence_length, 3, dtype=torch.float32, device=device)
+        position_ids[:, 0] = torch.arange(sequence_length, dtype=torch.float32, device=device)
+        position_ids[video_indices, 1] = torch.arange(num_video_tokens, dtype=torch.float32, device=device) % 4
+        position_ids[video_indices, 2] = torch.arange(num_video_tokens, dtype=torch.float32, device=device) % 2
 
         return {
-            "timestep": torch.tensor([0.7, 0.3], device=torch_device),
+            "timestep": torch.tensor([0.7, 0.3], device=device),
             "timestep_indices": timestep_indices,
             "token_tags": token_tags,
             "position_ids": position_ids,
@@ -123,7 +132,10 @@ def get_packed_layout(self, num_video_tokens: int = NUM_VIDEO_TOKENS) -> dict:
             "text_indices": text_indices,
         }
 
-    def get_dummy_inputs(self, num_video_tokens: int = NUM_VIDEO_TOKENS, batch_size: int = 2) -> dict:
+    def get_dummy_inputs(
+        self, num_video_tokens: int = NUM_VIDEO_TOKENS, batch_size: int = 2, device: str | torch.device = None
+    ) -> dict:
+        device = torch_device if device is None else device
         generator = self.generator
         init_dict = self.get_init_dict()
         patch_size = init_dict["patch_size"]
@@ -131,17 +143,17 @@ def get_dummy_inputs(self, num_video_tokens: int = NUM_VIDEO_TOKENS, batch_size:
 
         return {
             "hidden_states": randn_tensor(
-                (batch_size, num_video_tokens, video_patch_dim), generator=generator, device=torch_device
+                (batch_size, num_video_tokens, video_patch_dim), generator=generator, device=device
             ),
             "audio_hidden_states": randn_tensor(
                 (batch_size, NUM_AUDIO_TOKENS, init_dict["audio_in_channels"]),
                 generator=generator,
-                device=torch_device,
+                device=device,
             ),
             "encoder_hidden_states": randn_tensor(
-                (batch_size, NUM_TEXT_TOKENS, init_dict["text_dim"]), generator=generator, device=torch_device
+                (batch_size, NUM_TEXT_TOKENS, init_dict["text_dim"]), generator=generator, device=device
             ),
-            **self.get_packed_layout(num_video_tokens),
+            **self.get_packed_layout(num_video_tokens, device=device),
         }
 
 
@@ -189,3 +201,40 @@ class TestMiniMaxH3TransformerContextParallel(MiniMaxH3TransformerTesterConfig,
 
 class TestMiniMaxH3TransformerLoRA(MiniMaxH3TransformerTesterConfig, LoraTesterMixin):
     """LoRA tests for the MiniMax-H3 transformer."""
+
+
+class TestMiniMaxH3TransformerTensorParallel(MiniMaxH3TransformerTesterConfig, TensorParallelTesterMixin):
+    """Tensor Parallel inference tests for the MiniMax-H3 transformer (CUDA/XPU multi-accelerator)."""
+
+
+def make_neuron_tp_spec():
+    """Model spec consumed by the generic Neuron TP worker (`_neuron_tp_worker.py`).
+
+    Returns `(model_class, init_dict, cpu_inputs)`. Defined here so all MiniMax-H3-specific test data lives in this
+    file while the worker stays model-agnostic. Reuses the shared tester config so the spec never drifts from the
+    rest of the MiniMax-H3 tests.
+    """
+    config = MiniMaxH3TransformerTesterConfig()
+    return MiniMaxH3Transformer3DModel, config.get_init_dict(), config.get_dummy_inputs(device="cpu")
+
+
+@is_tensor_parallel
+@require_torch_neuron
+class TestMiniMaxH3TransformerTensorParallelNeuron:
+    """Tensor Parallel inference test for the MiniMax-H3 transformer on AWS Neuron.
+
+    Neuron TP runs through `torchrun` with the `"neuron"` distributed backend, so it cannot use the
+    `torch.multiprocessing`/NCCL spawn path of `TensorParallelTesterMixin`. This launches the generic worker with
+    the MiniMax-H3 model spec (`make_neuron_tp_spec`); the worker asserts the sharded output matches a
+    single-device reference, and the test checks its exit code.
+    """
+
+    def test_tensor_parallel_neuron_inference(self):
+        worker = os.path.join(os.path.dirname(__file__), "_neuron_tp_worker.py")
+        spec = "tests.models.transformers.test_models_transformer_minimax_h3:make_neuron_tp_spec"
+        cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", worker, spec]
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        assert result.returncode == 0, (
+            f"Neuron tensor-parallel worker failed (exit {result.returncode}).\n"
+            f"--- stdout ---\n{result.stdout}\n--- stderr ---\n{result.stderr}"
+        )

From 7e6e38facffe29551ca04f003e034b61e732718c Mon Sep 17 00:00:00 2001
From: JingyaHuang <huang_jingya@outlook.com>
Date: Fri, 21 Aug 2026 13:45:24 +0000
Subject: [PATCH 4/5] Shard MiniMax-H3's adaln_proj to fit tensor parallelism
 on one device
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`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>
---
 src/diffusers/hooks/tensor_parallel.py        | 45 +++++++++++++++----
 .../transformers/transformer_minimax_h3.py    | 17 +++++--
 2 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py
index f3856b438100..a13fa229788c 100644
--- a/src/diffusers/hooks/tensor_parallel.py
+++ b/src/diffusers/hooks/tensor_parallel.py
@@ -50,6 +50,20 @@ def __init__(self, blocks: "list[int] | None" = None):
         self.blocks = blocks
 
 
+class ReplicatedInputRowwiseParallel:
+    """Row-wise sharding for a Linear whose input arrives replicated instead of column-sharded.
+
+    Plain `"rowwise"` is the second half of a colwise/rowwise pair, so it expects its input to already be `Shard(-1)`
+    — which it is when the preceding Linear was colwise-sharded. A Linear that instead reads a replicated activation,
+    such as a modulation projection off the shared timestep embedding, needs its input sharded on the way in (a local
+    narrow, no collective) and its partial output all-reduced on the way out.
+
+    Weight and bias shard exactly as for plain `"rowwise"`: the weight over its input columns, the bias replicated and
+    added after the all-reduce. Use this to shard a large standalone projection whose output must keep the full
+    feature dimension, where colwise sharding would need an extra all-gather to rebuild it.
+    """
+
+
 def _blocks_to_block_sizes(total_size: int, blocks: "list[int]") -> "list[int]":
     """Convert proportional block counts to absolute sizes.
 
@@ -187,7 +201,8 @@ def resolve_tp_shard_specs(model: torch.nn.Module, tp_plan: dict) -> "dict[str,
             if style == "colwise":
                 weight_spec = TPShardSpec(0, [submodule.weight.shape[0]])
                 bias_spec = weight_spec
-            elif style == "rowwise":
+            elif style == "rowwise" or isinstance(style, ReplicatedInputRowwiseParallel):
+                # Both place the weight the same way; they differ only in the forward input/output hooks.
                 weight_spec = TPShardSpec(1, [submodule.weight.shape[1]])
                 bias_spec = TPShardSpec(None, None)
             elif isinstance(style, PackedColwiseParallel):
@@ -201,7 +216,8 @@ def resolve_tp_shard_specs(model: torch.nn.Module, tp_plan: dict) -> "dict[str,
             else:
                 raise ValueError(
                     f"Unsupported tensor-parallel style '{style}' for '{path}'. "
-                    f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+                    f"Expected 'colwise', 'rowwise', PackedColwiseParallel, PackedRowwiseParallel, or "
+                    f"ReplicatedInputRowwiseParallel."
                 )
 
             specs[f"{path}.weight"] = weight_spec
@@ -256,9 +272,10 @@ def _resolve_tp_plan(model: torch.nn.Module, tp_plan: dict) -> list:
 def _styles(relative_plan: dict) -> dict:
     """Map a `{relative_path: style}` plan to `parallelize_module` style instances.
 
-    Values may be plain strings (`"colwise"` / `"rowwise"`) or `PackedColwiseParallel` / `PackedRowwiseParallel` marker
-    instances. Returns `{relative_path: ColwiseParallel() | RowwiseParallel() | <packed impl>}`, each subclassed to
-    reject a sharded dim that is not divisible by the TP degree.
+    Values may be plain strings (`"colwise"` / `"rowwise"`) or `PackedColwiseParallel` / `PackedRowwiseParallel` /
+    `ReplicatedInputRowwiseParallel` marker instances. Returns `{relative_path: ColwiseParallel() |
+    RowwiseParallel() | <packed impl>}`, each subclassed to reject a sharded dim that is not divisible by the TP
+    degree.
     """
     import torch.nn as nn
     from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor
@@ -334,7 +351,7 @@ def _partition_linear_fn(self, name, module, device_mesh):
 
         return _CheckedColwiseImpl()
 
-    def _make_checked_row(path: str) -> RowwiseParallel:
+    def _make_checked_row(path: str, replicated_input: bool = False) -> RowwiseParallel:
         class _CheckedRowwiseImpl(RowwiseParallel):
             def _partition_linear_fn(self, name, module, device_mesh):
                 tp_size = device_mesh.size()
@@ -346,7 +363,10 @@ def _partition_linear_fn(self, name, module, device_mesh):
                     )
                 super()._partition_linear_fn(name, module, device_mesh)
 
-        return _CheckedRowwiseImpl()
+        # `input_layouts=Replicate()` makes `prepare_input` narrow the replicated activation down to this rank's
+        # columns rather than trusting it to already be `Shard(-1)`; the default would read a full-width tensor as
+        # if it were one rank's shard.
+        return _CheckedRowwiseImpl(input_layouts=Replicate()) if replicated_input else _CheckedRowwiseImpl()
 
     resolved = {}
     for path, style in relative_plan.items():
@@ -354,6 +374,8 @@ def _partition_linear_fn(self, name, module, device_mesh):
             resolved[path] = _make_checked_col(path)
         elif style == "rowwise":
             resolved[path] = _make_checked_row(path)
+        elif isinstance(style, ReplicatedInputRowwiseParallel):
+            resolved[path] = _make_checked_row(path, replicated_input=True)
         elif isinstance(style, PackedColwiseParallel):
             resolved[path] = _make_packed_col(style)
         elif isinstance(style, PackedRowwiseParallel):
@@ -361,7 +383,8 @@ def _partition_linear_fn(self, name, module, device_mesh):
         else:
             raise ValueError(
                 f"Unsupported tensor-parallel style '{style}' for '{path}'. "
-                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, PackedRowwiseParallel, or "
+                f"ReplicatedInputRowwiseParallel."
             )
     return resolved
 
@@ -377,6 +400,7 @@ def _hooks_only_styles(relative_plan: dict) -> dict:
     targeted module into a `Replicate()` DTensor via a broadcast. Callers should therefore place every planned
     parameter themselves, and must ensure none is left on `meta` — the broadcast would be issued on a meta tensor.
     """
+    from torch.distributed.tensor import Replicate
     from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
 
     class _NoPartitionColwise(ColwiseParallel):
@@ -393,10 +417,13 @@ def _partition_linear_fn(self, name, module, device_mesh):
             resolved[path] = _NoPartitionColwise()
         elif style == "rowwise" or isinstance(style, PackedRowwiseParallel):
             resolved[path] = _NoPartitionRowwise()
+        elif isinstance(style, ReplicatedInputRowwiseParallel):
+            resolved[path] = _NoPartitionRowwise(input_layouts=Replicate())
         else:
             raise ValueError(
                 f"Unsupported tensor-parallel style '{style}' for '{path}'. "
-                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, PackedRowwiseParallel, or "
+                f"ReplicatedInputRowwiseParallel."
             )
     return resolved
 
diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py
index 41269215cbf3..f6038c817848 100644
--- a/src/diffusers/models/transformers/transformer_minimax_h3.py
+++ b/src/diffusers/models/transformers/transformer_minimax_h3.py
@@ -19,7 +19,7 @@
 import torch.nn as nn
 
 from ...configuration_utils import ConfigMixin, register_to_config
-from ...hooks.tensor_parallel import PackedColwiseParallel
+from ...hooks.tensor_parallel import PackedColwiseParallel, ReplicatedInputRowwiseParallel
 from ...loaders import PeftAdapterMixin
 from ...utils import BaseOutput, apply_lora_scale, logging
 from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
@@ -457,10 +457,20 @@ class MiniMaxH3Transformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftA
     # packed projection is the SwiGLU input `ff.net.0.proj`, a single Linear producing `[value; gate]` in equal
     # halves, hence PackedColwiseParallel([1, 1]) so each half is sharded independently.
     #
+    # `adaln_proj.linear` is the one projection that has to keep its full output width: the six modulation
+    # parameters it produces scale and shift the full hidden dim of the packed sequence, which is already all-reduced
+    # by the time they are applied. Sharding it colwise would need an all-gather to rebuild that width, so it is
+    # sharded `ReplicatedInputRowwiseParallel` instead — over its `time_embed_dim` input, reading the replicated
+    # `temb` and all-reducing the result. It is worth the collective: at `6 * hidden_size * MINIMAX_H3_MODALITY_NUM`
+    # outputs per block it is ~40% of the denoiser's weights, and leaving it replicated puts the per-rank floor above
+    # a single device's memory at any TP degree. The all-reduce itself is over `num_timesteps * 6 * hidden_size *
+    # MINIMAX_H3_MODALITY_NUM` elements, i.e. hundreds of KB, once per block per step.
+    #
     # Intentionally absent, i.e. replicated on every rank: the RMSNorms (`norm1`, `norm2`,
     # `token_refiner.final_norm`, `norm_out.norm`); the QK-norms, which apply over `head_dim` after the heads are
-    # already split; the AdaLN modulation (`adaln_proj.linear`, `norm_out.linear`), which indexes the full hidden
-    # dim; and the patch/text embedders and the two output heads.
+    # already split; `norm_out.linear`, which indexes the full hidden dim as `adaln_proj` does but is a single
+    # `2 * hidden_size` projection rather than one per block, so sharding it would buy nothing; and the patch/text
+    # embedders and the two output heads.
     #
     # `attn.to_qkv` is deliberately not listed: it only exists after `fuse_projections()`, and the plan is
     # resolved by attribute lookup, so an unconditional entry would break the ordinary unfused model.
@@ -472,6 +482,7 @@ class MiniMaxH3Transformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftA
         "transformer_blocks.*.attn.to_out.0": "rowwise",
         "transformer_blocks.*.ff.net.0.proj": PackedColwiseParallel([1, 1]),
         "transformer_blocks.*.ff.net.2": "rowwise",
+        "transformer_blocks.*.adaln_proj.linear": ReplicatedInputRowwiseParallel(),
         # the token-refiner blocks are the same attention + SwiGLU FFN, minus AdaLN and rotary
         "token_refiner.refiner_blocks.*.attn.to_q": "colwise",
         "token_refiner.refiner_blocks.*.attn.to_k": "colwise",

From bb236dccfef5529910711c4340e9c9660cc3bffb Mon Sep 17 00:00:00 2001
From: JingyaHuang <huang_jingya@outlook.com>
Date: Fri, 21 Aug 2026 14:32:06 +0000
Subject: [PATCH 5/5] Build MiniMax-H3's row timestep plan on CPU, as its
 caller expects
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`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>
---
 .../modular_pipelines/minimax_h3/before_denoise.py         | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py
index c670467a9307..c119559ad408 100644
--- a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py
+++ b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py
@@ -1208,6 +1208,13 @@ def build_row_timesteps(
         Returns:
             `tuple[torch.Tensor, torch.Tensor]`: the distinct timesteps, sorted, and the index of every row into them.
         """
+        # 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.
+        video_indices = video_indices.cpu()
+        audio_indices = audio_indices.cpu()
+
         sequence_length = int(video_indices.numel() + audio_indices.numel() + num_text_tokens)
         row_timesteps = torch.full((sequence_length,), video_timestep, dtype=torch.float32)
         row_timesteps[video_indices[:num_condition_video_rows]] = condition_video_timestep