diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index 818488b4cfce..7d0282ea8836 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -495,7 +495,7 @@ torchrun --nproc-per-node 4 tensor_parallel_flux.py > > Combining tensor parallelism with quantization, offloading, or LoRA adapters isn't supported yet either, 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. +> 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. Each rank keeps its own slice of its own copy, so every rank must hold the same weights. Saving a tensor-parallel model isn't supported yet, and [`~ModelMixin.save_pretrained`] raises on one. Save the model before sharding it. @@ -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. diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index 19a772704e51..9f4d4152022e 100644 --- a/src/diffusers/hooks/tensor_parallel.py +++ b/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 @@ -223,80 +225,6 @@ def _resolve_tp_plan(model: torch.nn.Module, tp_plan: dict) -> list: return [grouped[key] for key in order] -def _shard_packed_param(param, dim: int, blocks: "list[int]", device_mesh, src_data_rank) -> torch.nn.Parameter: - """Shard a packed `param` along `dim`, splitting each fused block across the ranks separately. - - The parameter is replicated 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. - """ - from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor - - full = distribute_tensor(param, device_mesh, [Replicate()], src_data_rank=src_data_rank).to_local() - local = _local_shard(full, dim, _blocks_to_block_sizes(full.shape[dim], blocks), device_mesh) - return torch.nn.Parameter( - DTensor.from_local(local, device_mesh, [Shard(dim)], run_check=False), requires_grad=param.requires_grad - ) - - -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() | }`. 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 - - def _make_packed_col(marker: PackedColwiseParallel) -> ColwiseParallel: - class _PackedColwiseImpl(ColwiseParallel): - def _partition_linear_fn(self, name, module, device_mesh): - blocks = _packed_blocks(marker, module, name) - # 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(): - sharded = _shard_packed_param(param, 0, blocks, device_mesh, self.src_data_rank) - module.register_parameter(param_name, sharded) - - return _PackedColwiseImpl() - - def _make_packed_row(marker: PackedRowwiseParallel) -> RowwiseParallel: - class _PackedRowwiseImpl(RowwiseParallel): - def _partition_linear_fn(self, name, module, device_mesh): - blocks = _packed_blocks(marker, module, name) - # Only the weight (`[out, in]`) is sharded, column-wise (dim 1); the bias is added after the - # all-reduce, so every rank keeps it whole. - for param_name, param in module.named_parameters(): - if param_name == "weight": - sharded = _shard_packed_param(param, 1, blocks, device_mesh, self.src_data_rank) - else: - sharded = torch.nn.Parameter( - distribute_tensor(param, device_mesh, [Replicate()], src_data_rank=self.src_data_rank), - requires_grad=param.requires_grad, - ) - module.register_parameter(param_name, sharded) - - return _PackedRowwiseImpl() - - resolved = {} - for path, style in relative_plan.items(): - if style == "colwise": - resolved[path] = ColwiseParallel() - elif style == "rowwise": - resolved[path] = RowwiseParallel() - 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." - ) - return resolved - - def _hooks_only_styles(relative_plan: dict) -> dict: """Map a `{relative_path: style}` plan to styles that partition nothing. @@ -308,6 +236,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 +253,63 @@ 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": + # 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] = _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 out this rank's shard of every planned parameter, place it on `device`, then register the hooks. + + Each rank keeps its own slice and only that slice is copied to the device, so no weight collective is issued and no + full-size weight is moved to the device. Nothing is broadcast either, so every rank must already hold the same + weights, as it does after loading a checkpoint or seeding before a random initialization. + """ + 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 +476,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,11 +489,4 @@ 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) - return - - for submodule, relative_plan in groups: - parallelize_module(submodule, tp_mesh, _styles(relative_plan)) + _pre_shard_and_parallelize(model, tp_mesh, groups, specs, config._device) diff --git a/src/diffusers/hooks/tensor_parallel_neuron.py b/src/diffusers/hooks/tensor_parallel_neuron.py deleted file mode 100644 index ffcba0973d81..000000000000 --- a/src/diffusers/hooks/tensor_parallel_neuron.py +++ /dev/null @@ -1,75 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# 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. - -"""Neuron backend for tensor parallelism, dispatched from `apply_tensor_parallel` when the TP mesh is on Neuron. - -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 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.nn as nn - -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: - """Pre-shard the planned parameters via `DTensor.from_local`, then 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. - """ - 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: - parallelize_module(block, tp_mesh, _hooks_only_styles(relative_plan)) diff --git a/src/diffusers/models/model_loading_utils.py b/src/diffusers/models/model_loading_utils.py index b6d9d7997d81..655f50b5378a 100644 --- a/src/diffusers/models/model_loading_utils.py +++ b/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 = [] diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 469906613a5e..35774a129b8d 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/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: diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index 95ee1f2ac511..5e167ccfc36e 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 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 diff --git a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py index 94395a7c3ebe..bb1c1cd54b4f 100644 --- a/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py +++ b/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py @@ -1228,6 +1228,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 diff --git a/tests/models/transformers/test_models_transformer_minimax_h3.py b/tests/models/transformers/test_models_transformer_minimax_h3.py index 4165f9b5392d..d9258f04127e 100644 --- a/tests/models/transformers/test_models_transformer_minimax_h3.py +++ b/tests/models/transformers/test_models_transformer_minimax_h3.py @@ -28,6 +28,8 @@ MemoryTesterMixin, ModelTesterMixin, SingleFileTesterMixin, + TensorParallelNeuronTesterMixin, + TensorParallelTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -85,7 +87,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. @@ -93,29 +95,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, @@ -124,7 +130,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"] @@ -132,17 +141,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), } @@ -206,3 +215,11 @@ def pretrained_model_name_or_path(self): @property def pretrained_model_kwargs(self): return {"subfolder": "transformer"} + + +class TestMiniMaxH3TransformerTensorParallel(MiniMaxH3TransformerTesterConfig, TensorParallelTesterMixin): + """Tensor Parallel inference tests for the MiniMax-H3 transformer (CUDA/XPU multi-accelerator).""" + + +class TestMiniMaxH3TransformerTensorParallelNeuron(MiniMaxH3TransformerTesterConfig, TensorParallelNeuronTesterMixin): + """Tensor Parallel inference test for the MiniMax-H3 transformer on AWS Neuron."""