[feat] Tensor parallelism for MiniMax-H3

docs/source/en/training/distributed_inference.md

@@ -511,6 +511,7 @@ Each key may contain **at most one `*`**, and the prefix before it must resolve
 |---|---|---|---|
 | `"colwise"` | output features (`weight` dim 0) | computes a slice of the output | `to_q`, `to_k`, `to_v`, FFN in-projection |
 | `"rowwise"` | input features (`weight` dim 1) | computes a partial sum | `to_out.0`, FFN out-projection |
+| `"rowwise_split_input"` | input features (`weight` dim 1) | slices its replicated input, then computes a partial sum | a standalone projection that reads a replicated activation and must output the full width, e.g. an AdaLN modulation projection off the timestep embedding |
 
 Always pair them in that order. A `"colwise"` projection leaves its output sharded, the following `"rowwise"` projection consumes that shard directly, and a single `AllReduce` at the block boundary reconstructs the result. Sharding the pair any other way forces a gather in the middle and communicates far more.
 

src/diffusers/hooks/tensor_parallel.py

@@ -149,7 +149,8 @@ def resolve_tp_shard_specs(model: torch.nn.Module, tp_plan: dict, tp_degree: int
             if style == "colwise":
                 weight_spec = TPShardSpec(0, [submodule.weight.shape[0]])
                 bias_spec = weight_spec
-            elif style == "rowwise":
+            elif style in ("rowwise", "rowwise_split_input"):
+                # 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):
@@ -163,7 +164,8 @@ def resolve_tp_shard_specs(model: torch.nn.Module, tp_plan: dict, tp_degree: int
             else:
                 raise ValueError(
                     f"Unsupported tensor-parallel style '{style}' for '{path}'. "
-                    f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+                    f"Expected 'colwise', 'rowwise', 'rowwise_split_input', PackedColwiseParallel, or "
+                    f"PackedRowwiseParallel."
                 )
 
             specs[f"{path}.weight"] = weight_spec
@@ -242,9 +244,9 @@ def _shard_packed_param(param, dim: int, blocks: "list[int]", device_mesh, src_d
 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>}`. Divisibility by the TP
-    degree is checked earlier, by `resolve_tp_shard_specs`.
+    Values may be plain strings (`"colwise"` / `"rowwise"` / `"rowwise_split_input"`) or `PackedColwiseParallel` /
+    `PackedRowwiseParallel` marker instances. Returns `{relative_path: ColwiseParallel() | RowwiseParallel() | <packed
+    impl>}`. Divisibility by the TP degree is checked earlier, by `resolve_tp_shard_specs`.
     """
     from torch.distributed.tensor import Replicate, distribute_tensor
     from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
@@ -285,14 +287,21 @@ def _partition_linear_fn(self, name, module, device_mesh):
             resolved[path] = ColwiseParallel()
         elif style == "rowwise":
             resolved[path] = RowwiseParallel()
+        elif style == "rowwise_split_input":
+            # Same weight split as "rowwise", for a Linear whose input arrives whole instead of split by a "colwise"
+            # layer before it; each rank slices its own columns locally. E.g. an AdaLN modulation projection reading
+            # the timestep embedding (MiniMax-H3's `adaln_proj.linear`), or an output projection after a gathered
+            # QK-norm. Same as the transformers style of the same name.
+            resolved[path] = RowwiseParallel(input_layouts=Replicate())
         elif isinstance(style, PackedColwiseParallel):
             resolved[path] = _make_packed_col(style)
         elif isinstance(style, PackedRowwiseParallel):
             resolved[path] = _make_packed_row(style)
         else:
             raise ValueError(
                 f"Unsupported tensor-parallel style '{style}' for '{path}'. "
-                f"Expected 'colwise', 'rowwise', PackedColwiseParallel, or PackedRowwiseParallel."
+                f"Expected 'colwise', 'rowwise', 'rowwise_split_input', PackedColwiseParallel, or "
+                f"PackedRowwiseParallel."
             )
     return resolved
 
@@ -308,6 +317,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):
@@ -324,14 +334,58 @@ def _partition_linear_fn(self, name, module, device_mesh):
             resolved[path] = _NoPartitionColwise()
         elif style == "rowwise" or isinstance(style, PackedRowwiseParallel):
             resolved[path] = _NoPartitionRowwise()
+        elif style == "rowwise_split_input":
+            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', 'rowwise_split_input', PackedColwiseParallel, or "
+                f"PackedRowwiseParallel."
             )
     return resolved
 
 
+def _pre_shard_and_parallelize(
+    model: torch.nn.Module,
+    tp_mesh: "torch.distributed.device_mesh.DeviceMesh",
+    groups: list,
+    specs: "dict[str, TPShardSpec]",
+    device: torch.device,
+) -> None:
+    """Slice every planned parameter on CPU, place only this rank's shard on `device`, then register the hooks.
+
+    Unlike the default path this does not broadcast from a single rank, so every rank must already hold the same
+    weights. Model weights must be on CPU when this is called.
+    """
+    import torch.nn as nn
+    from torch.distributed.tensor import DTensor, Replicate, Shard
+    from torch.distributed.tensor.parallel import parallelize_module
+
+    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:
+        parallelize_module(block, tp_mesh, _hooks_only_styles(relative_plan))
+
+
 def _tp_degree(tp_config) -> int:
     """The TP degree, read the same way `_resolve_parallel_config` will, for checks run before the mesh is built."""
     return tp_config.mesh.size() if tp_config.mesh is not None else tp_config.tp_degree
@@ -498,9 +552,8 @@ def apply_tensor_parallel(
             f"or from the active accelerator when the mesh is built from `tp_degree`."
         )
 
-    backend = "neuron" if tp_mesh.device_type == "neuron" else "default"
     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}.")
+    logger.debug(f"Applying tensor parallel over {len(groups)} module group(s) on mesh {tp_mesh}.")
 
     from torch.distributed.tensor.parallel import parallelize_module
 
@@ -512,10 +565,11 @@ def apply_tensor_parallel(
     # Also validates every planned shard, so an uneven split raises before any module is sharded.
     specs = resolve_tp_shard_specs(model, tp_plan, tp_mesh.size())
 
-    if backend == "neuron":
-        from .tensor_parallel_neuron import _apply_tp_neuron
-
-        _apply_tp_neuron(model, tp_mesh, groups, specs)
+    if tp_mesh.device_type in ("neuron", "tpu"):
+        # Pre-shard on CPU instead of letting `parallelize_module` distribute the weights. On Neuron, consecutive
+        # `reduce_scatter`s of large weights (>= 5120x5120) in one `parallelize_module` call can fail in the runtime. On
+        # TPU, it would materialize every full weight on each chip first, which runs out of HBM on large models.
+        _pre_shard_and_parallelize(model, tp_mesh, groups, specs, config._device)
         return
 
     for submodule, relative_plan in groups:

src/diffusers/models/model_loading_utils.py

@@ -407,12 +407,7 @@ def _load_shard_file_tp(
     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
+    device = tp_config._device
 
     mismatched_keys = []
 

src/diffusers/models/modeling_utils.py

@@ -1544,8 +1544,19 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | 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 torch.distributed.tensor import DTensor
+
                 from ..hooks.tensor_parallel import apply_tensor_parallel
 
+                # 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.
+                tp_device = tp_config._device
+                for name, buffer in model.named_buffers():
+                    if buffer.device != tp_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(tp_device)
+
                 apply_tensor_parallel(model, tp_config, cls._tp_plan, weights_already_sharded=True)
             else:
                 model.enable_parallelism(config=parallel_config)
@@ -1716,7 +1727,7 @@ def _resolve_parallel_config(
         rank = torch.distributed.get_rank()
         world_size = torch.distributed.get_world_size()
         device_type = torch._C._get_accelerator().type
-        device = torch.device(device_type, rank % torch.accelerator.device_count())
+        device = torch.device(device_type, rank % max(torch.accelerator.device_count(), 1))
 
         mesh = None
         if config.context_parallel_config is not None:

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 FromOriginalModelMixin, 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)
@@ -451,6 +454,45 @@ class MiniMaxH3Transformer3DModel(
         "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.
+    #
+    # `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 `"rowwise_split_input"` 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; `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.
+    _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",
+        "transformer_blocks.*.adaln_proj.linear": "rowwise_split_input",
+        # 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