diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index 818488b4cfce..25de837a8c25 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/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. diff --git a/src/diffusers/hooks/tensor_parallel.py b/src/diffusers/hooks/tensor_parallel.py index 19a772704e51..951450a8dffb 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 @@ -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() | }`. 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() | }`. 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,6 +287,12 @@ 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): @@ -292,7 +300,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', '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: 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."""