From ce178f8a4bb58fff8c022e8b5445d4a3807e2ac0 Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Sat, 22 Aug 2026 04:33:24 +0000 Subject: [PATCH 1/4] perf(minimax-h3): align Sol Engine inference optimizations --- fastvideo/attention/layer.py | 33 +++++ fastvideo/attention/minimax_h3_relayout.py | 135 ++++++++++++++++++ fastvideo/distributed/communication_op.py | 8 ++ fastvideo/distributed/parallel_state.py | 22 +++ fastvideo/envs.py | 9 ++ fastvideo/models/dits/minimax_h3.py | 91 ++++++++++++ .../minimax_h3/stages/minimax_h3_denoising.py | 8 ++ .../test_minimax_h3_sol_optimizations.py | 104 ++++++++++++++ 8 files changed, 410 insertions(+) create mode 100644 fastvideo/attention/minimax_h3_relayout.py create mode 100644 fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index ac53492b7d..10a32bb046 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -8,6 +8,7 @@ from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather, + sequence_model_parallel_direct_all_to_all, sequence_model_parallel_all_to_all_4D) from fastvideo.distributed.parallel_state import (get_sp_parallel_rank, get_sp_world_size) from fastvideo.forward_context import ForwardContext, get_forward_context @@ -70,6 +71,7 @@ def __init__(self, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = "", + packed_qkv_relayout: bool = False, **extra_impl_args) -> None: super().__init__() if softmax_scale is None: @@ -103,6 +105,9 @@ def __init__(self, # inference loader may enable this one instance after validating the # transformer's resolved backend; no process-global default changes. self._compile_forward_enabled = not _attention_compile_disabled() + self.packed_qkv_relayout = packed_qkv_relayout + if packed_qkv_relayout and self.backend != AttentionBackendEnum.FLASH_ATTN: + raise ValueError("MiniMax-H3 packed QKV relayout currently supports only dense FLASH_ATTN") def _set_compile_forward_enabled(self, enabled: bool) -> None: self._compile_forward_enabled = enabled @@ -144,6 +149,34 @@ def forward( forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata + if self.packed_qkv_relayout and world_size > 1: + if batch_size != 1: + raise ValueError("MiniMax-H3 packed QKV relayout currently requires batch size 1") + if any(t is not None for t in (replicated_q, replicated_k, replicated_v, freqs_cis)): + raise ValueError("MiniMax-H3 packed QKV relayout does not support replicated tokens or deferred RoPE") + from fastvideo.attention.minimax_h3_relayout import (merge_heads, + pack_qkv_destination_major) + + rows_local = q.shape[1] + packed = pack_qkv_destination_major(q[0], k[0], v[0], world_size) + packed = sequence_model_parallel_direct_all_to_all(packed) + heads_local = num_heads // world_size + packed = packed.reshape(world_size * rows_local, heads_local, 3 * self.head_size) + q_full, k_full, v_full = packed.split(self.head_size, dim=-1) + original_seq_len = original_seq_len or q_full.shape[0] + pad_seq_len = q_full.shape[0] - original_seq_len + output = self.attn_impl.forward( + q_full[:original_seq_len].unsqueeze(0), + k_full[:original_seq_len].unsqueeze(0), + v_full[:original_seq_len].unsqueeze(0), + ctx_attn_metadata, + ) + output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) + output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len)) + output = sequence_model_parallel_direct_all_to_all(output.squeeze(0).contiguous()) + output = merge_heads(output.reshape(world_size, rows_local, heads_local, self.head_size)) + return output.unsqueeze(0), None + # Stack QKV qkv = torch.cat([q, k, v], dim=0) # [3*batch, seq_len, num_heads, head_dim] diff --git a/fastvideo/attention/minimax_h3_relayout.py b/fastvideo/attention/minimax_h3_relayout.py new file mode 100644 index 0000000000..4c203f1d4f --- /dev/null +++ b/fastvideo/attention/minimax_h3_relayout.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Compile-safe MiniMax-H3 Ulysses relayout kernels adapted from Sol-Engine.""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: # pragma: no cover + triton = None + tl = None + HAVE_TRITON = False + + +if HAVE_TRITON: + + @triton.jit + def _pack_qkv_kernel( + out_ptr, + q_ptr, + k_ptr, + v_ptr, + total_elements, + rows, + heads_local, + head_dim, + stride_q_row, + stride_q_head, + stride_k_row, + stride_k_head, + stride_v_row, + stride_v_head, + BLOCK: tl.constexpr, + ): + offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < total_elements + dim = offsets % head_dim + head_slot = offsets // head_dim + local_head = head_slot % heads_local + row_slot = head_slot // heads_local + row = row_slot % rows + destination = row_slot // rows + global_head = destination * heads_local + local_head + + q = tl.load(q_ptr + row * stride_q_row + global_head * stride_q_head + dim, mask=mask) + k = tl.load(k_ptr + row * stride_k_row + global_head * stride_k_head + dim, mask=mask) + v = tl.load(v_ptr + row * stride_v_row + global_head * stride_v_head + dim, mask=mask) + base = head_slot * (3 * head_dim) + dim + tl.store(out_ptr + base, q, mask=mask) + tl.store(out_ptr + base + head_dim, k, mask=mask) + tl.store(out_ptr + base + 2 * head_dim, v, mask=mask) + + @triton.jit + def _merge_heads_kernel(out_ptr, x_ptr, total_elements, world, rows, inner, BLOCK: tl.constexpr): + offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < total_elements + tail = offsets % inner + slot = offsets // inner + source = slot % world + row = slot // world + src = (source * rows + row) * inner + tail + tl.store(out_ptr + offsets, tl.load(x_ptr + src, mask=mask), mask=mask) + + @torch.library.triton_op("fastvideo::minimax_h3_pack_qkv", mutates_args={}) + def _pack_qkv_op(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, world: int) -> torch.Tensor: + rows, heads, head_dim = q.shape + heads_local = heads // world + out = torch.empty((world, rows, heads_local, 3 * head_dim), dtype=q.dtype, device=q.device) + total = rows * heads * head_dim + torch.library.wrap_triton(_pack_qkv_kernel)[(triton.cdiv(total, 1024), )]( + out, + q, + k, + v, + total, + rows, + heads_local, + head_dim, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + BLOCK=1024, + num_warps=8, + ) + return out + + @torch.library.triton_op("fastvideo::minimax_h3_merge_heads", mutates_args={}) + def _merge_heads_op(x: torch.Tensor) -> torch.Tensor: + world, rows, heads_local, head_dim = x.shape + out = torch.empty((rows, world, heads_local, head_dim), dtype=x.dtype, device=x.device) + total = out.numel() + torch.library.wrap_triton(_merge_heads_kernel)[(triton.cdiv(total, 1024), )]( + out, + x, + total, + world, + rows, + heads_local * head_dim, + BLOCK=1024, + num_warps=8, + ) + return out + + +def pack_qkv_destination_major(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, world: int) -> torch.Tensor: + """Move three ``(rows, heads, dim)`` tensors into destination-major QKV in one pass.""" + if not HAVE_TRITON: + raise RuntimeError("MiniMax-H3 packed sequence parallelism requires Triton") + if q.ndim != 3 or q.shape != k.shape or q.shape != v.shape: + raise ValueError("q, k, and v must have matching (rows, heads, head_dim) shapes") + if q.shape[1] % world: + raise ValueError(f"heads ({q.shape[1]}) must divide sequence parallel size ({world})") + if any(t.stride(-1) != 1 for t in (q, k, v)): + raise ValueError("q, k, and v must be contiguous in head_dim") + return _pack_qkv_op(q, k, v, world) + + +def merge_heads(x: torch.Tensor) -> torch.Tensor: + """Move source-major all-to-all output back to row-major head order in one pass.""" + if not HAVE_TRITON: + raise RuntimeError("MiniMax-H3 packed sequence parallelism requires Triton") + if x.ndim != 4 or not x.is_contiguous(): + raise ValueError("packed all-to-all output must be a contiguous 4D tensor") + world, rows, heads_local, head_dim = x.shape + return _merge_heads_op(x).reshape(rows, world * heads_local, head_dim) + + +__all__ = ["HAVE_TRITON", "merge_heads", "pack_qkv_destination_major"] diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/distributed/communication_op.py index 502284f06b..24a227a886 100644 --- a/fastvideo/distributed/communication_op.py +++ b/fastvideo/distributed/communication_op.py @@ -32,6 +32,14 @@ def sequence_model_parallel_all_to_all_4D(input_: torch.Tensor, return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim) +def sequence_model_parallel_direct_all_to_all(input_: torch.Tensor) -> torch.Tensor: + """Synchronous equal-split all-to-all used by the packed H3 inference path.""" + group = get_sp_group() + if group.world_size == 1: + return input_ + return torch.ops.fastvideo.direct_all_to_all_single(input_, group.unique_name) + + def sequence_model_parallel_all_gather(input_: torch.Tensor, dim: int = -1) -> torch.Tensor: """All-gather the input tensor across model parallel group.""" return get_sp_group().all_gather(input_, dim) diff --git a/fastvideo/distributed/parallel_state.py b/fastvideo/distributed/parallel_state.py index 5cba91c373..49f2e7e7e3 100644 --- a/fastvideo/distributed/parallel_state.py +++ b/fastvideo/distributed/parallel_state.py @@ -114,6 +114,28 @@ def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor: return torch.empty_like(tensor) +@torch.library.custom_op( + "fastvideo::direct_all_to_all_single", + mutates_args=(), + device_types="cuda", +) +def direct_all_to_all_single(tensor: torch.Tensor, group_name: str) -> torch.Tensor: + """Issue a synchronous all-to-all without the functional-collective wrapper.""" + assert group_name in _groups, f"Group {group_name} is not found." + group = _groups[group_name]() + if group is None: + raise ValueError(f"Group {group_name} is destroyed.") + output = torch.empty_like(tensor) + torch.distributed.all_to_all_single(output, tensor, group=group.device_group) + return output + + +@torch.library.register_fake("fastvideo::direct_all_to_all_single") +def direct_all_to_all_single_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor: + del group_name + return torch.empty_like(tensor) + + class GroupCoordinator: """ PyTorch ProcessGroup wrapper for a group of processes. diff --git a/fastvideo/envs.py b/fastvideo/envs.py index 86abc1bcde..86fc37ab59 100644 --- a/fastvideo/envs.py +++ b/fastvideo/envs.py @@ -28,6 +28,8 @@ FASTVIDEO_VAE_PARALLEL_ENCODE: bool = False FASTVIDEO_VAE_PARALLEL_DECODE_STRATEGY: str | None = None FASTVIDEO_ULYSSES_A2A: str = "off" + FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE: bool = False + FASTVIDEO_MINIMAX_H3_PACKED_SP: bool = False FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "spawn" FASTVIDEO_TARGET_DEVICE: str = "cuda" MAX_JOBS: str | None = None @@ -270,6 +272,13 @@ def maybe_convert_int(value: str | None) -> int | None: "FASTVIDEO_ULYSSES_A2A": lambda: os.getenv("FASTVIDEO_ULYSSES_A2A", "off").strip().lower(), + # Exact, inference-only MiniMax-H3 trajectory optimizations adapted from + # Sol-Engine. Kept separate from block-local fusions for clean A/B tests. + "FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE": + lambda: os.getenv("FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE", "0") != "0", + "FASTVIDEO_MINIMAX_H3_PACKED_SP": + lambda: os.getenv("FASTVIDEO_MINIMAX_H3_PACKED_SP", "0") != "0", + # Use dedicated multiprocess context for workers. "FASTVIDEO_WORKER_MULTIPROC_METHOD": lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "spawn"), diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index ec168cec18..0c3c62b47e 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -145,6 +145,7 @@ def __init__( prefix: str, fuse_qknorm_rope: bool = False, fa4_packed_varlen: bool = False, + packed_qkv_relayout: bool = False, ) -> None: super().__init__() self.num_attention_heads = num_attention_heads @@ -198,6 +199,7 @@ def __init__( supported_attention_backends=supported_attention_backends, prefix=prefix, fa4_packed_varlen=fa4_packed_varlen, + packed_qkv_relayout=packed_qkv_relayout, ) self.to_gate_compress: ReplicatedLinear | None = None # None = unchecked; the first forward tests the loaded weight once and @@ -418,6 +420,31 @@ def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]: return temb.view(-1, 6 * self.hidden_size).chunk(6, dim=-1) +class _MiniMaxH3StepCursor: + + __slots__ = ("signature", "step") + + def __init__(self, device: torch.device, signature: tuple[tuple[float, ...], ...]) -> None: + self.step = torch.zeros((), dtype=torch.long, device=device) + self.signature = signature + + def set(self, index: int) -> None: + self.step.fill_(index) + + +class _MiniMaxH3PrecomputedModulation(nn.Module): + + def __init__(self, table: torch.Tensor, cursor: _MiniMaxH3StepCursor) -> None: + super().__init__() + self.register_buffer("table", table, persistent=False) + self.cursor = cursor + + def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]: + del temb + rows = self.table.index_select(0, self.cursor.step.reshape(1))[0] + return rows.chunk(6, dim=-1) + + class MiniMaxH3AdaLayerNormOut(nn.Module): """Final RMSNorm with per-timestep row modulation.""" @@ -475,6 +502,7 @@ def __init__( fuse_qknorm_rope: bool = False, fuse_swiglu: bool = False, fa4_packed_varlen: bool = False, + packed_qkv_relayout: bool = False, ) -> None: super().__init__() self.norm1 = nn.RMSNorm(hidden_size, eps=norm_eps) @@ -488,6 +516,7 @@ def __init__( prefix=f"{prefix}.attn", fuse_qknorm_rope=fuse_qknorm_rope, fa4_packed_varlen=fa4_packed_varlen, + packed_qkv_relayout=packed_qkv_relayout, ) self.norm2 = nn.RMSNorm(hidden_size, eps=norm_eps) self.ff = MiniMaxH3FeedForward( @@ -596,6 +625,8 @@ def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None: super().__init__(config, hf_config) arch = config.arch_config self.enabled_fusions = _enabled_minimax_h3_fusions() + self.adaln_precompute_enabled = envs.FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE + self.packed_sp_enabled = envs.FASTVIDEO_MINIMAX_H3_PACKED_SP if self.enabled_fusions: if HAVE_TRITON: logger.info( @@ -707,6 +738,7 @@ def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None: fuse_qknorm_rope="qknorm_rope" in self.enabled_fusions, fuse_swiglu="swiglu" in self.enabled_fusions, fa4_packed_varlen=envs.FASTVIDEO_MINIMAX_H3_FA4_PACKED_VARLEN, + packed_qkv_relayout=self.packed_sp_enabled, ) for index in range(arch.num_layers) ]) self.norm_out = MiniMaxH3AdaLayerNormOut( @@ -733,6 +765,65 @@ def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None: ) self.__post_init__() + @torch.no_grad() + def prepare_adaln_trajectory(self, row_timestep_plan: list[tuple[torch.Tensor, + torch.Tensor]]) -> dict[str, float | int]: + """Replace full-rank per-block AdaLN projections by exact per-step tables.""" + signature = tuple(tuple(float(value) for value in timestep.detach().cpu().flatten()) + for timestep, _ in row_timestep_plan) + cursor = getattr(self, "_h3_adaln_cursor", None) + if cursor is not None: + if cursor.signature != signature: + raise RuntimeError( + "MiniMax-H3 AdaLN weights were replaced by a cached trajectory; reuse this transformer only " + "with the same denoising schedule.") + return {"steps": len(row_timestep_plan), "blocks": len(self.transformer_blocks), "installed": 0} + if self.adaln_rank is not None: + raise RuntimeError("trajectory AdaLN precompute expects the stock full-rank MiniMax-H3 checkpoint") + if not row_timestep_plan: + raise ValueError("row_timestep_plan must not be empty") + + device = next(self.parameters()).device + embeddings = [] + for timestep, _ in row_timestep_plan: + temb = self.time_proj(timestep.to(device)) + embeddings.append(self.time_embedder(temb.to(self.time_embedder.fc_in.weight.dtype))) + max_rows = max(int(timestep.numel()) for timestep, _ in row_timestep_plan) * MINIMAX_H3_MODALITY_NUM + + def padded(rows: torch.Tensor) -> torch.Tensor: + if rows.shape[0] == max_rows: + return rows + return torch.cat((rows, rows.new_zeros(max_rows - rows.shape[0], rows.shape[1]))) + + cursor = _MiniMaxH3StepCursor(device, signature) + table_bytes = 0 + freed_bytes = 0 + for block in self.transformer_blocks: + projection = block.adaln_proj + table = torch.stack([padded(torch.cat(projection(temb), dim=-1)) for temb in embeddings]) + table_bytes += table.numel() * table.element_size() + freed_bytes += sum(parameter.numel() * parameter.element_size() + for parameter in projection.parameters()) + block.adaln_proj = _MiniMaxH3PrecomputedModulation(table, cursor) + self._h3_adaln_cursor = cursor + torch.cuda.empty_cache() + stats: dict[str, float | int] = { + "steps": len(row_timestep_plan), + "blocks": len(self.transformer_blocks), + "table_gb": table_bytes / 1024**3, + "freed_gb": freed_bytes / 1024**3, + "installed": 1, + } + logger.info( + "MiniMax H3 AdaLN trajectory cached: %d blocks x %d steps, table %.2f GB, freed %.2f GB.", + stats["blocks"], stats["steps"], stats["table_gb"], stats["freed_gb"]) + return stats + + def set_adaln_step(self, index: int) -> None: + cursor = getattr(self, "_h3_adaln_cursor", None) + if cursor is not None: + cursor.set(index) + def prepare_for_compile(self) -> None: """Pipeline hook, called once right before torch.compile wraps the blocks. diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index 72f3669cc2..22bb67251c 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -127,6 +127,13 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward row_timestep_plan.append((unique.to(device), inverse.to(device))) batch.timesteps = video_timesteps + if bool(getattr(self.transformer, "adaln_precompute_enabled", False)): + if fastvideo_args.use_fsdp_inference: + raise RuntimeError( + "MiniMax-H3 trajectory AdaLN precompute replaces projection modules after loading; " + "use replicated inference weights, matching Sol-Engine, rather than FSDP inference.") + self.transformer.prepare_adaln_trajectory(row_timestep_plan) + position_ids = layout.position_ids.to(device) token_tags = layout.token_tags.to(device) video_indices = layout.video_indices.to(device) @@ -158,6 +165,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward for index, (video_timestep, audio_timestep) in enumerate(zip(video_timesteps, audio_timesteps, strict=True)): unique_timesteps, timestep_indices = row_timestep_plan[index] + self.transformer.set_adaln_step(index) attn_metadata = None if vsa_metadata_builder is not None: # Optional schedule: run the first N steps dense (sparsity 0 diff --git a/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py b/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py new file mode 100644 index 0000000000..e9c6681aee --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for Sol-Engine-aligned MiniMax-H3 optimizations.""" + +from types import SimpleNamespace + +import pytest +import torch + + +def test_minimax_h3_relayout_is_bit_exact_and_compile_safe() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + pytest.importorskip("triton") + + from fastvideo.attention.minimax_h3_relayout import ( + merge_heads, + pack_qkv_destination_major, + ) + + rows, heads, head_dim, world = 17, 8, 16, 4 + fused = torch.randn(rows, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16) + q, k, v = fused[:, 0], fused[:, 1], fused[:, 2] + heads_local = heads // world + + expected_pack = torch.empty( + world, + rows, + heads_local, + 3 * head_dim, + device=q.device, + dtype=q.dtype, + ) + for index, tensor in enumerate((q, k, v)): + shard = tensor.reshape(rows, world, heads_local, head_dim).permute(1, 0, 2, 3) + expected_pack[..., index * head_dim:(index + 1) * head_dim].copy_(shard) + + compiled_pack = torch.compile(pack_qkv_destination_major, fullgraph=True, dynamic=False) + actual_pack = compiled_pack(q, k, v, world) + assert torch.equal(actual_pack, expected_pack) + + packed_output = torch.randn_like(expected_pack[..., :head_dim]) + expected_merge = packed_output.permute(1, 0, 2, 3).contiguous().reshape(rows, heads, head_dim) + compiled_merge = torch.compile(merge_heads, fullgraph=True, dynamic=False) + actual_merge = compiled_merge(packed_output) + assert torch.equal(actual_merge, expected_merge) + + +def test_adaln_precompute_rejects_a_different_trajectory() -> None: + from fastvideo.models.dits.minimax_h3 import ( + MiniMaxH3Transformer3DModel, + _MiniMaxH3StepCursor, + ) + + cursor = _MiniMaxH3StepCursor(torch.device("cpu"), ((1.0,),)) + transformer = SimpleNamespace( + _h3_adaln_cursor=cursor, + transformer_blocks=[object()], + ) + same_plan = [(torch.tensor([1.0]), torch.tensor([0]))] + stats = MiniMaxH3Transformer3DModel.prepare_adaln_trajectory(transformer, same_plan) + assert stats["installed"] == 0 + + different_plan = [(torch.tensor([0.5]), torch.tensor([0]))] + with pytest.raises(RuntimeError, match="same denoising schedule"): + MiniMaxH3Transformer3DModel.prepare_adaln_trajectory(transformer, different_plan) + + +def test_minimax_h3_fusions_capture_in_fullgraph() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + pytest.importorskip("triton") + + from fastvideo.models.dits.minimax_h3_fusions import ( + fused_qknorm_rope, + fused_residual_gate_rmsnorm_modulate, + fused_rmsnorm_modulate, + minimax_h3_swiglu, + ) + + device, dtype = torch.device("cuda"), torch.bfloat16 + batch, rows, hidden = 1, 12, 128 + x = torch.randn(batch, rows, hidden, device=device, dtype=dtype) + branch = torch.randn_like(x) + weight = torch.randn(hidden, device=device, dtype=dtype) + table = torch.randn(18, hidden, device=device, dtype=dtype) + index = torch.arange(rows, device=device).remainder(table.shape[0]) + q = torch.randn(batch, rows, 1, hidden, device=device, dtype=dtype) + cos = torch.randn(rows, 96, device=device, dtype=dtype) + sin = torch.randn_like(cos) + + cases = ( + (fused_rmsnorm_modulate, (x, weight, table, table, index, 1e-5)), + (fused_residual_gate_rmsnorm_modulate, (x, branch, table, weight, table, table, index, 1e-5)), + (fused_qknorm_rope, (q, weight, cos, sin, 1e-5)), + (minimax_h3_swiglu, (torch.randn(batch, rows, 512, device=device, dtype=dtype),)), + ) + with torch.inference_mode(): + for function, args in cases: + expected = function(*args) + actual = torch.compile(function, fullgraph=True, dynamic=False)(*args) + expected_items = expected if isinstance(expected, tuple) else (expected,) + actual_items = actual if isinstance(actual, tuple) else (actual,) + for expected_item, actual_item in zip(expected_items, actual_items, strict=True): + torch.testing.assert_close(actual_item, expected_item, atol=0, rtol=0) From 4dbccffd010340672357b131a0797301d38b4e52 Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Sat, 22 Aug 2026 07:26:44 +0000 Subject: [PATCH 2/4] refactor(minimax-h3): colocate relayout kernels --- fastvideo/attention/layer.py | 3 +-- .../dits/minimax_h3_fusions/relayout.py} | 0 .../basic/minimax_h3/stages/minimax_h3_denoising.py | 5 ++--- .../tests/transformers/test_minimax_h3_sol_optimizations.py | 2 +- 4 files changed, 4 insertions(+), 6 deletions(-) rename fastvideo/{attention/minimax_h3_relayout.py => models/dits/minimax_h3_fusions/relayout.py} (100%) diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index 10a32bb046..600abf8184 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -154,8 +154,7 @@ def forward( raise ValueError("MiniMax-H3 packed QKV relayout currently requires batch size 1") if any(t is not None for t in (replicated_q, replicated_k, replicated_v, freqs_cis)): raise ValueError("MiniMax-H3 packed QKV relayout does not support replicated tokens or deferred RoPE") - from fastvideo.attention.minimax_h3_relayout import (merge_heads, - pack_qkv_destination_major) + from fastvideo.models.dits.minimax_h3_fusions.relayout import (merge_heads, pack_qkv_destination_major) rows_local = q.shape[1] packed = pack_qkv_destination_major(q[0], k[0], v[0], world_size) diff --git a/fastvideo/attention/minimax_h3_relayout.py b/fastvideo/models/dits/minimax_h3_fusions/relayout.py similarity index 100% rename from fastvideo/attention/minimax_h3_relayout.py rename to fastvideo/models/dits/minimax_h3_fusions/relayout.py diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index 22bb67251c..7054cb8b75 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -129,9 +129,8 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward if bool(getattr(self.transformer, "adaln_precompute_enabled", False)): if fastvideo_args.use_fsdp_inference: - raise RuntimeError( - "MiniMax-H3 trajectory AdaLN precompute replaces projection modules after loading; " - "use replicated inference weights, matching Sol-Engine, rather than FSDP inference.") + raise RuntimeError("MiniMax-H3 trajectory AdaLN precompute replaces projection modules after loading; " + "use replicated inference weights, matching Sol-Engine, rather than FSDP inference.") self.transformer.prepare_adaln_trajectory(row_timestep_plan) position_ids = layout.position_ids.to(device) diff --git a/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py b/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py index e9c6681aee..514768dc7c 100644 --- a/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py +++ b/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py @@ -12,7 +12,7 @@ def test_minimax_h3_relayout_is_bit_exact_and_compile_safe() -> None: pytest.skip("CUDA is required") pytest.importorskip("triton") - from fastvideo.attention.minimax_h3_relayout import ( + from fastvideo.models.dits.minimax_h3_fusions.relayout import ( merge_heads, pack_qkv_destination_major, ) From 85c78a9b0805f315c03bd33c57c9961a79211ced Mon Sep 17 00:00:00 2001 From: William Lin <8941107+SolitaryThinker@users.noreply.github.com> Date: Wed, 26 Aug 2026 00:02:43 -0700 Subject: [PATCH 3/4] [bugfix]: make H3 Sol optimizations lifecycle-safe --- docs/inference/optimizations.md | 39 +++ fastvideo/attention/layer.py | 17 +- fastvideo/distributed/communication_op.py | 22 +- fastvideo/models/dits/minimax_h3.py | 30 ++- .../minimax_h3/stages/minimax_h3_denoising.py | 17 +- .../distributed/test_minimax_h3_packed_sp.py | 162 ++++++++++++ .../test_minimax_h3_sol_optimizations.py | 234 ++++++++++++++++++ 7 files changed, 509 insertions(+), 12 deletions(-) create mode 100644 fastvideo/tests/distributed/test_minimax_h3_packed_sp.py diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 753fef2835..87b5467d49 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -113,6 +113,45 @@ the Preview checkpoint's sparse VSA blocks. Packed-varlen changes floating-point reduction order relative to fixed-length FA4, so treat it as a speed/quality evaluation option rather than an exact-parity mode. +### MiniMax-H3 AdaLN trajectory and packed sequence parallelism (opt-in) + +Two additional released Sol-Engine-style DiT optimizations are independently +available for controlled MiniMax-H3 inference experiments: + +```bash +export FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE=1 +export FASTVIDEO_MINIMAX_H3_PACKED_SP=1 +``` + +Both flags default to off. They do not enable VSA, cross-step caches, +quantization, or VAE optimizations. + +- `FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE` evaluates each full-rank block AdaLN + projection once for every step in the fixed denoising schedule, installs + non-persistent lookup tables, and releases the projection modules. The table + values match the original projections, but the loaded transformer is then + tied to that schedule for the rest of its lifetime. It requires a stock + full-rank checkpoint with replicated, materialized weights: + `dit_layerwise_offload=False` and `use_fsdp_inference=False`. Setup briefly + holds the original projections and all replacement tables together. In the + reported 50-step workload the tables used 0.88 GB and replaced 24.23 GB of + projection weights; other schedules and checkpoints have different costs. +- `FASTVIDEO_MINIMAX_H3_PACKED_SP` fuses the Q/K/V communication layout around + two direct NCCL all-to-all collectives. It engages only for inference with + sequence-parallel world size greater than one, batch size one, dense + FlashAttention, no replicated tokens or deferred RoPE, and no VSA. A + grad-enabled forward automatically retains the autograd-aware generic + Ulysses route. The packed direct collective is separate from the optional + `FASTVIDEO_ULYSSES_A2A=auto` transport. + +At SP=1 the packed flag is deliberately inert, so it provides no acceleration +on a one-GPU GB10 and does not force FlashAttention over the configured GB10 +backend. The contribution's performance evidence is scoped to dense BF16 FA4 +on 4× GB200, SP=4, 1344×768×124, with one post-warmup sample; it is not evidence +for a sparse/VSA route or a one-GPU system. These exact table/relayout features +also do not make regional `torch.compile` eager-parity-safe: the MiniMax-H3 +accuracy caveat below still applies to the complete serving profile. + ### FP4 Flash Attention 4 (Blackwell only) **`FLASH_ATTN`** with **`--nvfp4_fa4`** diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index 600abf8184..f0e34fd0b3 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -8,13 +8,13 @@ from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather, - sequence_model_parallel_direct_all_to_all, - sequence_model_parallel_all_to_all_4D) + sequence_model_parallel_all_to_all_4D, + sequence_model_parallel_direct_all_to_all) from fastvideo.distributed.parallel_state import (get_sp_parallel_rank, get_sp_world_size) from fastvideo.forward_context import ForwardContext, get_forward_context +from fastvideo.layers.rotary_embedding import _apply_rotary_emb from fastvideo.platforms import AttentionBackendEnum from fastvideo.utils import get_compute_dtype -from fastvideo.layers.rotary_embedding import _apply_rotary_emb def _attention_compile_disabled() -> bool: @@ -149,9 +149,16 @@ def forward( forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata - if self.packed_qkv_relayout and world_size > 1: + # The direct packed collective is deliberately inference-only. A + # training process may inherit the opt-in environment variable, but a + # grad-enabled forward must retain the established autograd-aware + # Ulysses path. + if self.packed_qkv_relayout and world_size > 1 and not torch.is_grad_enabled(): if batch_size != 1: raise ValueError("MiniMax-H3 packed QKV relayout currently requires batch size 1") + if num_heads % world_size: + raise ValueError( + f"MiniMax-H3 packed QKV relayout requires {num_heads} heads to be divisible by SP={world_size}") if any(t is not None for t in (replicated_q, replicated_k, replicated_v, freqs_cis)): raise ValueError("MiniMax-H3 packed QKV relayout does not support replicated tokens or deferred RoPE") from fastvideo.models.dits.minimax_h3_fusions.relayout import (merge_heads, pack_qkv_destination_major) @@ -163,6 +170,8 @@ def forward( packed = packed.reshape(world_size * rows_local, heads_local, 3 * self.head_size) q_full, k_full, v_full = packed.split(self.head_size, dim=-1) original_seq_len = original_seq_len or q_full.shape[0] + if original_seq_len < 1 or original_seq_len > q_full.shape[0]: + raise ValueError(f"original_seq_len must be in [1, {q_full.shape[0]}], got {original_seq_len}") pad_seq_len = q_full.shape[0] - original_seq_len output = self.attn_impl.forward( q_full[:original_seq_len].unsqueeze(0), diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/distributed/communication_op.py index 24a227a886..9fc19a316f 100644 --- a/fastvideo/distributed/communication_op.py +++ b/fastvideo/distributed/communication_op.py @@ -33,10 +33,30 @@ def sequence_model_parallel_all_to_all_4D(input_: torch.Tensor, def sequence_model_parallel_direct_all_to_all(input_: torch.Tensor) -> torch.Tensor: - """Synchronous equal-split all-to-all used by the packed H3 inference path.""" + """Synchronous equal-split all-to-all used by the packed H3 inference path. + + This primitive intentionally has no autograd formula. The owning attention + route falls back to ``all_to_all_4D`` whenever gradients are enabled; this + guard keeps other callers from discovering the restriction at backward. + """ group = get_sp_group() if group.world_size == 1: return input_ + if torch.is_grad_enabled() and input_.requires_grad: + raise RuntimeError("sequence_model_parallel_direct_all_to_all is inference-only and has no autograd formula; " + "use sequence_model_parallel_all_to_all_4D for grad-enabled execution") + if input_.ndim < 1 or input_.shape[0] % group.world_size: + raise ValueError( + "direct all-to-all requires the leading dimension to be evenly divisible by the SP world size; " + f"got shape {tuple(input_.shape)} and SP={group.world_size}") + if not input_.is_contiguous(): + raise ValueError("direct all-to-all requires a contiguous input tensor") + # CPU/Gloo is useful for the real multi-rank contract test. The production + # packed H3 route is CUDA/Triton and takes the compiler-visible custom op. + if not input_.is_cuda: + output = torch.empty_like(input_) + torch.distributed.all_to_all_single(output, input_, group=group.device_group) + return output return torch.ops.fastvideo.direct_all_to_all_single(input_, group.unique_name) diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 0c3c62b47e..e430c486dd 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -70,6 +70,13 @@ def _can_run_minimax_h3_fusion(tensor: torch.Tensor) -> bool: return HAVE_TRITON and tensor.is_cuda and not torch.is_grad_enabled() +def _packed_sp_active(requested: bool, world_size: int) -> bool: + """Packed H3 relayout is meaningful only for multi-rank inference.""" + if world_size < 1: + raise ValueError(f"sequence parallel world size must be positive, got {world_size}") + return requested and world_size > 1 + + class MiniMaxH3RotaryPosEmbed(nn.Module): """Three-axis rotary frequencies over packed `(t, h, w)` coordinates.""" @@ -429,6 +436,8 @@ def __init__(self, device: torch.device, signature: tuple[tuple[float, ...], ... self.signature = signature def set(self, index: int) -> None: + if index < 0 or index >= len(self.signature): + raise IndexError(f"AdaLN trajectory step must be in [0, {len(self.signature)}), got {index}") self.step.fill_(index) @@ -626,7 +635,15 @@ def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None: arch = config.arch_config self.enabled_fusions = _enabled_minimax_h3_fusions() self.adaln_precompute_enabled = envs.FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE - self.packed_sp_enabled = envs.FASTVIDEO_MINIMAX_H3_PACKED_SP + sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1 + packed_sp_requested = envs.FASTVIDEO_MINIMAX_H3_PACKED_SP + self.packed_sp_enabled = _packed_sp_active(packed_sp_requested, sp_world_size) + if packed_sp_requested and not self.packed_sp_enabled: + logger.info("FASTVIDEO_MINIMAX_H3_PACKED_SP is inert at SP=1; using the configured attention backend") + elif self.packed_sp_enabled: + logger.info( + "MiniMax H3 packed sequence parallelism enabled for SP=%d (dense FlashAttention, batch-1, " + "inference-only; grad-enabled forwards retain the autograd-aware Ulysses path)", sp_world_size) if self.enabled_fusions: if HAVE_TRITON: logger.info( @@ -637,7 +654,6 @@ def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None: logger.warning( "FASTVIDEO_MINIMAX_H3_FUSIONS requested %s but Triton is unavailable; " "every forward stays on the eager path.", ",".join(sorted(self.enabled_fusions))) - sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1 if arch.num_attention_heads % sp_world_size: raise ValueError(f"MiniMax H3 attention heads ({arch.num_attention_heads}) must be divisible by " f"sequence parallel size ({sp_world_size}).") @@ -798,15 +814,21 @@ def padded(rows: torch.Tensor) -> torch.Tensor: cursor = _MiniMaxH3StepCursor(device, signature) table_bytes = 0 freed_bytes = 0 + replacements: list[_MiniMaxH3PrecomputedModulation] = [] for block in self.transformer_blocks: projection = block.adaln_proj table = torch.stack([padded(torch.cat(projection(temb), dim=-1)) for temb in embeddings]) table_bytes += table.numel() * table.element_size() freed_bytes += sum(parameter.numel() * parameter.element_size() for parameter in projection.parameters()) - block.adaln_proj = _MiniMaxH3PrecomputedModulation(table, cursor) + replacements.append(_MiniMaxH3PrecomputedModulation(table, cursor)) + # Build every table before mutating the transformer. A projection or + # allocation failure therefore leaves all original modules intact. + for block, replacement in zip(self.transformer_blocks, replacements, strict=True): + block.adaln_proj = replacement self._h3_adaln_cursor = cursor - torch.cuda.empty_cache() + if device.type == "cuda": + torch.cuda.empty_cache() stats: dict[str, float | int] = { "steps": len(row_timestep_plan), "blocks": len(self.transformer_blocks), diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index 7054cb8b75..8e687cd766 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -26,6 +26,19 @@ from fastvideo.utils import get_compute_dtype +def _validate_adaln_precompute_configuration(fastvideo_args: FastVideoArgs) -> None: + """Reject loader modes whose parameter lifecycle precompute bypasses.""" + if fastvideo_args.dit_layerwise_offload: + raise RuntimeError( + "MiniMax-H3 trajectory AdaLN precompute cannot run with dit_layerwise_offload=True because " + "the projection weights are materialized only by each block forward hook; disable layerwise offload " + "or disable FASTVIDEO_MINIMAX_H3_ADALN_PRECOMPUTE") + if fastvideo_args.use_fsdp_inference: + raise RuntimeError( + "MiniMax-H3 trajectory AdaLN precompute replaces projection modules after loading; use replicated " + "inference weights, matching Sol-Engine, rather than FSDP inference") + + def _h3_vsa_metadata_builder(transformer: Any, fastvideo_args: FastVideoArgs) -> Any: """Builder instance when the transformer resolved to VSA-H3, else None. @@ -128,9 +141,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward batch.timesteps = video_timesteps if bool(getattr(self.transformer, "adaln_precompute_enabled", False)): - if fastvideo_args.use_fsdp_inference: - raise RuntimeError("MiniMax-H3 trajectory AdaLN precompute replaces projection modules after loading; " - "use replicated inference weights, matching Sol-Engine, rather than FSDP inference.") + _validate_adaln_precompute_configuration(fastvideo_args) self.transformer.prepare_adaln_trajectory(row_timestep_plan) position_ids = layout.position_ids.to(device) diff --git a/fastvideo/tests/distributed/test_minimax_h3_packed_sp.py b/fastvideo/tests/distributed/test_minimax_h3_packed_sp.py new file mode 100644 index 0000000000..90503d9de9 --- /dev/null +++ b/fastvideo/tests/distributed/test_minimax_h3_packed_sp.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Real world-4 contract coverage for MiniMax-H3 packed sequence parallelism.""" + +from __future__ import annotations + +import contextlib +import os +import socket +import subprocess +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn + +SP_WORLD_SIZE = 4 + + +class _TestGroup: + + def __init__(self, world_size: int) -> None: + self.world_size = world_size + self.device_group = dist.group.WORLD + self.unique_name = "minimax_h3_packed_sp_test" + + +class _IdentityAttentionImpl(nn.Module): + + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, metadata: object) -> torch.Tensor: + del k, v, metadata + return q + + def postprocess_output(self, output: torch.Tensor, metadata: object) -> torch.Tensor: + del metadata + return output + + +def _reference_pack(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, world_size: int) -> torch.Tensor: + rows, heads, head_dim = q.shape + heads_local = heads // world_size + output = torch.empty(world_size, + rows, + heads_local, + 3 * head_dim, + device=q.device, + dtype=q.dtype) + for index, tensor in enumerate((q, k, v)): + shard = tensor.reshape(rows, world_size, heads_local, head_dim).permute(1, 0, 2, 3) + output[..., index * head_dim:(index + 1) * head_dim].copy_(shard) + return output + + +def _reference_merge(output: torch.Tensor) -> torch.Tensor: + world, rows, heads_local, head_dim = output.shape + return output.permute(1, 0, 2, 3).contiguous().reshape(rows, world * heads_local, head_dim) + + +def _worker() -> None: + from fastvideo.attention.layer import DistributedAttention + from fastvideo.distributed import communication_op, parallel_state + from fastvideo.forward_context import set_forward_context + from fastvideo.models.dits.minimax_h3_fusions import relayout + + world = int(os.environ["WORLD_SIZE"]) + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + exact_cuda_route = torch.cuda.device_count() >= world + backend = "nccl" if exact_cuda_route else "gloo" + dist.init_process_group(backend=backend) + if exact_cuda_route: + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + dtype = torch.bfloat16 + else: + device = torch.device("cpu") + dtype = torch.float32 + + group = _TestGroup(world) + parallel_state._register_group(group) + attention = DistributedAttention.__new__(DistributedAttention) + nn.Module.__init__(attention) + attention.attn_impl = _IdentityAttentionImpl() + attention.head_size = 16 + attention.packed_qkv_relayout = True + attention._compile_forward_enabled = True + + rows_local, heads, head_dim = 7, 8, 16 + values = torch.arange(rows_local * heads * head_dim, device=device, dtype=torch.float32) + q = (values.reshape(1, rows_local, heads, head_dim) + rank * 10_000).to(dtype) + k = q + 100 + v = q + 200 + semantic_rows = world * rows_local - 3 + expected = q.clone() + if rank == world - 1: + expected[:, -3:] = 0 + + try: + with contextlib.ExitStack() as stack: + stack.enter_context(patch("fastvideo.attention.layer.get_sp_world_size", return_value=world)) + stack.enter_context(patch("fastvideo.attention.layer.get_sp_parallel_rank", return_value=rank)) + stack.enter_context(patch.object(communication_op, "get_sp_group", return_value=group)) + if not exact_cuda_route: + # The production relayout kernels are independently compiled + # and bit-exact-tested on CUDA. Gloo lets every development + # machine exercise the same two real world-4 collectives and + # rank ordering without pretending Triton runs on CPU. + stack.enter_context(patch.object(relayout, "pack_qkv_destination_major", _reference_pack)) + stack.enter_context(patch.object(relayout, "merge_heads", _reference_merge)) + stack.enter_context(torch.inference_mode()) + stack.enter_context(set_forward_context(current_timestep=0, attn_metadata=None)) + output, replicated = attention(q, k, v, original_seq_len=semantic_rows) + + assert replicated is None + assert torch.equal(output, expected), f"rank {rank} packed scatter/gather round-trip differed" + + with patch.object(communication_op, "get_sp_group", return_value=group): + with torch.no_grad(), pytest.raises(ValueError, match="leading dimension"): + communication_op.sequence_model_parallel_direct_all_to_all(torch.empty(world + 1, 2, device=device)) + with pytest.raises(RuntimeError, match="inference-only"): + communication_op.sequence_model_parallel_direct_all_to_all( + torch.empty(world, 2, device=device, requires_grad=True)) + + dist.barrier() + if rank == 0: + mode = "cuda-production" if exact_cuda_route else "gloo-contract" + print(f"MINIMAX_H3_PACKED_SP_OK world={world} mode={mode}", flush=True) + finally: + dist.destroy_process_group() + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_minimax_h3_packed_sp_world4_collective_contract() -> None: + environment = dict(os.environ, OMP_NUM_THREADS="1") + process = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc_per_node={SP_WORLD_SIZE}", + f"--master_port={_free_port()}", + str(Path(__file__).resolve()), + "--worker", + ], + env=environment, + capture_output=True, + text=True, + timeout=300, + ) + assert process.returncode == 0 and "MINIMAX_H3_PACKED_SP_OK" in process.stdout, ( + f"stdout:\n{process.stdout[-6000:]}\nstderr:\n{process.stderr[-6000:]}") + + +if __name__ == "__main__" and "--worker" in sys.argv: + _worker() diff --git a/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py b/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py index 514768dc7c..6600a063cb 100644 --- a/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py +++ b/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py @@ -2,9 +2,77 @@ """Regression tests for Sol-Engine-aligned MiniMax-H3 optimizations.""" from types import SimpleNamespace +from unittest.mock import patch import pytest import torch +import torch.nn as nn + + +class _TimeProjection(nn.Module): + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + return torch.stack((timesteps, timesteps.square()), dim=-1) + + +class _TimeEmbedder(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.fc_in = nn.Linear(2, 2, bias=False) + with torch.no_grad(): + self.fc_in.weight.copy_(torch.eye(2)) + + def forward(self, embeddings: torch.Tensor) -> torch.Tensor: + return self.fc_in(embeddings) + + +class _Projection(nn.Module): + + def __init__(self, offset: float) -> None: + super().__init__() + self.linear = nn.Linear(2, 18) + with torch.no_grad(): + values = torch.arange(36, dtype=torch.float32).reshape(18, 2) + self.linear.weight.copy_(values / 37 + offset) + self.linear.bias.copy_(torch.arange(18, dtype=torch.float32) / 19 + offset) + + def forward(self, embeddings: torch.Tensor) -> tuple[torch.Tensor, ...]: + rows = self.linear(embeddings).view(-1, 6) + return rows.chunk(6, dim=-1) + + +class _ProjectionBlock(nn.Module): + + def __init__(self, offset: float) -> None: + super().__init__() + self.adaln_proj = _Projection(offset) + + +class _TrajectoryTransformer(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.time_proj = _TimeProjection() + self.time_embedder = _TimeEmbedder() + self.transformer_blocks = nn.ModuleList([_ProjectionBlock(0.0), _ProjectionBlock(0.25)]) + self.adaln_rank = None + + +class _IdentityAttentionImpl(nn.Module): + + def preprocess_qkv(self, qkv: torch.Tensor, metadata: object) -> torch.Tensor: + del metadata + return qkv + + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, metadata: object) -> torch.Tensor: + del k, v, metadata + return q + + def postprocess_output(self, output: torch.Tensor, metadata: object) -> torch.Tensor: + del metadata + return output def test_minimax_h3_relayout_is_bit_exact_and_compile_safe() -> None: @@ -65,6 +133,172 @@ def test_adaln_precompute_rejects_a_different_trajectory() -> None: MiniMaxH3Transformer3DModel.prepare_adaln_trajectory(transformer, different_plan) +def test_adaln_precompute_matches_projection_tables_and_reuses_schedule() -> None: + from fastvideo.models.dits.minimax_h3 import MiniMaxH3Transformer3DModel + + transformer = _TrajectoryTransformer() + plan = [ + (torch.tensor([1.0, 0.5]), torch.tensor([0, 1])), + (torch.tensor([0.25]), torch.tensor([0])), + ] + original_projections = [block.adaln_proj for block in transformer.transformer_blocks] + embeddings = [transformer.time_embedder(transformer.time_proj(timestep)) for timestep, _ in plan] + expected = [] + for projection in original_projections: + per_step = [torch.cat(projection(embedding), dim=-1) for embedding in embeddings] + per_step[1] = torch.cat((per_step[1], torch.zeros_like(per_step[1]))) + expected.append(per_step) + + stats = MiniMaxH3Transformer3DModel.prepare_adaln_trajectory(transformer, plan) + assert stats["installed"] == 1 + assert stats["steps"] == 2 + assert stats["blocks"] == 2 + + for step in range(2): + MiniMaxH3Transformer3DModel.set_adaln_step(transformer, step) + for block, expected_steps in zip(transformer.transformer_blocks, expected, strict=True): + actual = torch.cat(block.adaln_proj(torch.empty(0)), dim=-1) + assert torch.equal(actual, expected_steps[step]) + + reused = MiniMaxH3Transformer3DModel.prepare_adaln_trajectory(transformer, plan) + assert reused["installed"] == 0 + with pytest.raises(IndexError, match="trajectory step"): + MiniMaxH3Transformer3DModel.set_adaln_step(transformer, 2) + + +def test_adaln_precompute_failure_does_not_partially_replace_blocks() -> None: + from fastvideo.models.dits.minimax_h3 import MiniMaxH3Transformer3DModel + + transformer = _TrajectoryTransformer() + originals = [block.adaln_proj for block in transformer.transformer_blocks] + + def fail_projection(embeddings: torch.Tensor) -> tuple[torch.Tensor, ...]: + del embeddings + raise RuntimeError("injected projection failure") + + transformer.transformer_blocks[1].adaln_proj.forward = fail_projection + plan = [(torch.tensor([1.0]), torch.tensor([0]))] + with pytest.raises(RuntimeError, match="injected projection failure"): + MiniMaxH3Transformer3DModel.prepare_adaln_trajectory(transformer, plan) + assert [block.adaln_proj for block in transformer.transformer_blocks] == originals + assert not hasattr(transformer, "_h3_adaln_cursor") + + +@pytest.mark.parametrize( + ("layerwise", "fsdp", "message"), + [ + (True, False, "dit_layerwise_offload=True"), + (False, True, "FSDP inference"), + ], +) +def test_adaln_precompute_rejects_incompatible_loader_lifecycles(layerwise: bool, fsdp: bool, + message: str) -> None: + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import ( + _validate_adaln_precompute_configuration, + ) + + args = SimpleNamespace(dit_layerwise_offload=layerwise, use_fsdp_inference=fsdp) + with pytest.raises(RuntimeError, match=message): + _validate_adaln_precompute_configuration(args) + + +def test_adaln_precompute_accepts_replicated_materialized_weights() -> None: + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import ( + _validate_adaln_precompute_configuration, + ) + + _validate_adaln_precompute_configuration( + SimpleNamespace(dit_layerwise_offload=False, use_fsdp_inference=False)) + + +def test_packed_sp_is_inert_on_one_rank() -> None: + from fastvideo.models.dits.minimax_h3 import _packed_sp_active + + assert not _packed_sp_active(False, 1) + assert not _packed_sp_active(True, 1) + assert _packed_sp_active(True, 4) + with pytest.raises(ValueError, match="must be positive"): + _packed_sp_active(True, 0) + + +def test_packed_sp_falls_back_to_autograd_aware_collective_when_grad_enabled() -> None: + from fastvideo.attention.layer import DistributedAttention + from fastvideo.forward_context import set_forward_context + + attention = DistributedAttention.__new__(DistributedAttention) + nn.Module.__init__(attention) + attention.attn_impl = _IdentityAttentionImpl() + attention.head_size = 2 + attention.packed_qkv_relayout = True + attention._compile_forward_enabled = True + + q = torch.randn(1, 5, 4, 2, requires_grad=True) + k = torch.randn_like(q, requires_grad=True) + v = torch.randn_like(q, requires_grad=True) + with ( + patch("fastvideo.attention.layer.get_sp_world_size", return_value=2), + patch("fastvideo.attention.layer.get_sp_parallel_rank", return_value=0), + patch("fastvideo.attention.layer.sequence_model_parallel_all_to_all_4D", side_effect=lambda tensor, **_: tensor) + as generic, + patch("fastvideo.attention.layer.sequence_model_parallel_direct_all_to_all", + side_effect=AssertionError("packed collective must not run with gradients")) as direct, + set_forward_context(current_timestep=0, attn_metadata=None), + ): + output, replicated = attention(q, k, v) + output.sum().backward() + assert replicated is None + assert torch.equal(q.grad, torch.ones_like(q)) + assert generic.call_count == 2 + direct.assert_not_called() + + +def test_direct_packed_collective_rejects_unsupported_inputs_before_launch() -> None: + from fastvideo.distributed.communication_op import sequence_model_parallel_direct_all_to_all + + group = SimpleNamespace(world_size=4) + with patch("fastvideo.distributed.communication_op.get_sp_group", return_value=group): + with pytest.raises(RuntimeError, match="inference-only"): + sequence_model_parallel_direct_all_to_all(torch.randn(4, 2, requires_grad=True)) + with torch.no_grad(), pytest.raises(ValueError, match="leading dimension"): + sequence_model_parallel_direct_all_to_all(torch.randn(5, 2)) + with torch.no_grad(), pytest.raises(ValueError, match="contiguous"): + sequence_model_parallel_direct_all_to_all(torch.randn(2, 4).transpose(0, 1)) + + +def test_direct_packed_collective_captures_as_a_fullgraph_custom_op(tmp_path) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if torch.distributed.is_initialized(): + pytest.skip("test owns a temporary world-1 process group") + + from fastvideo.distributed import parallel_state + + class Group: + unique_name = "minimax_h3_direct_compile_test" + device_group = None + + group = Group() + torch.distributed.init_process_group( + backend="nccl", + init_method=f"file://{tmp_path / 'store'}", + rank=0, + world_size=1, + ) + group.device_group = torch.distributed.group.WORLD + parallel_state._register_group(group) + try: + def direct_collective(tensor: torch.Tensor) -> torch.Tensor: + return torch.ops.fastvideo.direct_all_to_all_single(tensor, group.unique_name) + + tensor = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + compiled = torch.compile(direct_collective, fullgraph=True, dynamic=False) + output = compiled(tensor) + assert torch.equal(output, tensor) + finally: + torch.distributed.destroy_process_group() + + def test_minimax_h3_fusions_capture_in_fullgraph() -> None: if not torch.cuda.is_available(): pytest.skip("CUDA is required") From 7299c643e466e007b0c46e33efb958c1dd14d72b Mon Sep 17 00:00:00 2001 From: William Lin <8941107+SolitaryThinker@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:01:45 -0700 Subject: [PATCH 4/4] [bugfix]: harden H3 Sol optimization contracts --- .buildkite/scripts/lanes/ssim.sh | 6 + .buildkite/scripts/lanes/transformer.sh | 2 +- .github/scripts/plan_merge_ci.py | 7 + docs/contributing/ci_architecture.md | 11 +- docs/contributing/testing.md | 9 + docs/inference/optimizations.md | 8 +- fastvideo/attention/layer.py | 3 +- fastvideo/distributed/communication_op.py | 36 +++ fastvideo/distributed/parallel_state.py | 35 ++- fastvideo/models/dits/minimax_h3.py | 18 +- .../dits/minimax_h3_fusions/relayout.py | 12 +- .../minimax_h3/stages/minimax_h3_denoising.py | 130 +++++----- .../tests/contract/test_ci_test_collection.py | 17 +- .../tests/contract/test_merge_ci_plan.py | 9 + .../tests/contract/test_modal_fa4_policy.py | 3 +- .../distributed/test_minimax_h3_packed_sp.py | 76 ++++-- fastvideo/tests/modal/test_pr_test.py | 3 +- .../test_minimax_h3_sol_optimizations.py | 232 ++++++++++++++++++ 18 files changed, 514 insertions(+), 103 deletions(-) diff --git a/.buildkite/scripts/lanes/ssim.sh b/.buildkite/scripts/lanes/ssim.sh index 54bc0d4068..a359e9ef11 100755 --- a/.buildkite/scripts/lanes/ssim.sh +++ b/.buildkite/scripts/lanes/ssim.sh @@ -25,6 +25,12 @@ if [ "$selected" != all ]; then done fi +# This lane owns a whole four-GPU tray. Exercise the exact NCCL + Triton +# MiniMax-H3 packed-SP route before output-quality jobs, while retaining the +# scheduler-provided rendezvous port in the nested torchrun invocation. +FASTVIDEO_MINIMAX_H3_PACKED_SP_STRICT_CUDA=1 \ + pytest ./fastvideo/tests/distributed/test_minimax_h3_packed_sp.py -vs + # MoGe's utils3d dependency builds glcontext from source on ARM64. The current # runner image predates the baked-in X11 headers below, so keep this guarded # bootstrap until every deployed image digest contains libx11-dev. diff --git a/.buildkite/scripts/lanes/transformer.sh b/.buildkite/scripts/lanes/transformer.sh index 9041f10df7..57902222ff 100755 --- a/.buildkite/scripts/lanes/transformer.sh +++ b/.buildkite/scripts/lanes/transformer.sh @@ -2,4 +2,4 @@ # Canonical Slurm CI selection for the transformer lane. set -euo pipefail -exec pytest ./fastvideo/tests/transformers -vs +exec pytest ./fastvideo/tests/transformers ./fastvideo/tests/distributed/test_minimax_h3_packed_sp.py -vs diff --git a/.github/scripts/plan_merge_ci.py b/.github/scripts/plan_merge_ci.py index 074a3e652e..959bfe6a72 100644 --- a/.github/scripts/plan_merge_ci.py +++ b/.github/scripts/plan_merge_ci.py @@ -373,6 +373,13 @@ def classify_paths(paths: list[str]) -> MergePlan: plan.add_lanes("ssim", reason=f"shared SSIM harness/reference: {path}") continue + if path == "fastvideo/tests/distributed/test_minimax_h3_packed_sp.py": + # Fastcheck owns the portable world-4 Gloo contract. Also select + # the existing four-GPU SSIM lane, whose preflight makes the same + # test require the production NCCL + Triton route. + plan.add_ssim(("test_minimax_h3_similarity.py", ), reason=f"MiniMax-H3 packed-SP CUDA preflight: {path}") + continue + if path.startswith("fastvideo/tests/performance/") or path.startswith(".buildkite/performance-benchmarks/"): plan.add_lanes("performance", reason=f"performance coverage: {path}") continue diff --git a/docs/contributing/ci_architecture.md b/docs/contributing/ci_architecture.md index 1a630dd9fa..3311217161 100644 --- a/docs/contributing/ci_architecture.md +++ b/docs/contributing/ci_architecture.md @@ -164,6 +164,14 @@ selection never deletes or dynamically invents a Buildkite step. | Modular train framework | `train_framework` | 1 | `fastvideo/train/` and its tests | | Eval metrics | `eval` | 1 | `fastvideo/eval/` and its tests | +The transformer Fastcheck lane also runs the MiniMax-H3 packed sequence- +parallel world-4 contract with Gloo. That portable test proves rank ordering, +padding/trim behavior, and Q/K/V-dependent scatter/gather semantics; it is not +CUDA performance or NCCL/Triton evidence. The existing four-GPU SSIM lane runs +the same contract first in strict mode, which fails unless all four CUDA +devices are visible and the production NCCL plus Triton relayout route runs. +No separate distributed lane is introduced. + Golden-gate and SSIM selections are basenames, not arbitrary pytest arguments. The private host checks the comma-separated allowlist before staging, and the container checks it again before invoking pytest. Shared quality-harness @@ -352,7 +360,8 @@ table binds each internal `*_ci` type to that script, its GPU count, wall-clock limit, dependency extras, kernel-build policy, secrets, and artifacts. The internal suffix is an implementation detail; there is only one active backend. -SSIM uses `fastvideo/tests/ssim/ci_runner.py` inside a single four-GPU lease. +SSIM first runs the strict MiniMax-H3 packed-SP production preflight, then uses +`fastvideo/tests/ssim/ci_runner.py` inside the same four-GPU lease. It discovers `REQUIRED_GPUS` and `*_MODEL_TO_PARAMS` with AST parsing, then packs independent pytest subprocesses across the visible GPUs with fail-fast termination. Performance writes reports to a host-mounted artifact directory; diff --git a/docs/contributing/testing.md b/docs/contributing/testing.md index 78ebe0e9bd..382ad6d6ea 100644 --- a/docs/contributing/testing.md +++ b/docs/contributing/testing.md @@ -16,6 +16,7 @@ slash-command mappings, and workflow ownership live in | Inference tests | `fastvideo/tests/inference` | Validate specialized inference paths such as LoRA inference and V-MoBA. | | Performance tests | `fastvideo/tests/performance` | Gate latency, throughput, peak memory, and stage timings. See [Performance Benchmarks](performance_benchmarks.md). | | Eval tests | `fastvideo/tests/eval` | Check eval metrics against pinned reference scores and assets. | +| Distributed contracts | `fastvideo/tests/distributed` | Exercise selected real multi-rank communication contracts; most remain manual unless an existing lane names them explicitly. | | DreamVerse app tests | `apps/dreamverse` | Validate the DreamVerse backend, frontend, and mock-backed browser flows. | ## Running Tests Locally @@ -151,6 +152,14 @@ by the changed model family. Shared SSIM harness changes still select the complete lane. Independently, `main` runs the full SSIM matrix every Sunday at 05:00 UTC so infrequently touched model families retain periodic coverage. +Before scheduling SSIM cases, the lane uses its existing four-GPU allocation +to run the MiniMax-H3 packed-SP world-4 contract in strict CUDA mode. Strict +mode rejects a Gloo fallback and therefore covers the production NCCL and +Triton relayout route. The one-GPU transformer Fastcheck lane runs the same +test with its portable Gloo fallback; that result covers collective ordering +and Q/K/V semantics only, not CUDA behavior or performance. Both nested +torchrun invocations preserve the Slurm runner's assigned `MASTER_PORT`. + For a focused developer run, invoke pytest directly and optionally select one model from a parameterized test through `FASTVIDEO_SSIM_MODEL_ID`: diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 87b5467d49..0f9642ec42 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -130,7 +130,13 @@ quantization, or VAE optimizations. projection once for every step in the fixed denoising schedule, installs non-persistent lookup tables, and releases the projection modules. The table values match the original projections, but the loaded transformer is then - tied to that schedule for the rest of its lifetime. It requires a stock + tied to that schedule for the rest of its lifetime. The runtime cursor is a + non-persistent buffer and follows transformer device moves, including full + CPU offload. The replacement tables and cursor are deliberately omitted + from `state_dict`, while the original projection parameters no longer + exist: a post-precompute state dict is therefore not reloadable. Save or + convert checkpoints before enabling precompute, and reload the stock model + to change schedules. It requires a stock full-rank checkpoint with replicated, materialized weights: `dit_layerwise_offload=False` and `use_fsdp_inference=False`. Setup briefly holds the original projections and all replacement tables together. In the diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index f0e34fd0b3..053063c856 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -169,7 +169,8 @@ def forward( heads_local = num_heads // world_size packed = packed.reshape(world_size * rows_local, heads_local, 3 * self.head_size) q_full, k_full, v_full = packed.split(self.head_size, dim=-1) - original_seq_len = original_seq_len or q_full.shape[0] + if original_seq_len is None: + original_seq_len = q_full.shape[0] if original_seq_len < 1 or original_seq_len > q_full.shape[0]: raise ValueError(f"original_seq_len must be in [1, {q_full.shape[0]}], got {original_seq_len}") pad_seq_len = q_full.shape[0] - original_seq_len diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/distributed/communication_op.py index 9fc19a316f..c842dc14b5 100644 --- a/fastvideo/distributed/communication_op.py +++ b/fastvideo/distributed/communication_op.py @@ -14,6 +14,36 @@ _sp_warmup_done = False +def _validate_direct_all_to_all_group(group: object, input_: torch.Tensor, expected_backend: str) -> None: + """Validate the live process group immediately before a direct collective.""" + process_group = getattr(group, "device_group", None) + if process_group is None or not torch.distributed.is_initialized(): + raise RuntimeError("direct all-to-all requires a live distributed process group") + try: + actual_world = torch.distributed.get_world_size(process_group) + torch.distributed.get_rank(process_group) + backend = str(torch.distributed.get_backend(process_group)).lower() + except (RuntimeError, ValueError) as error: + raise RuntimeError("direct all-to-all process group is not live") from error + configured_world = int(getattr(group, "world_size", 0)) + if actual_world != configured_world: + raise RuntimeError( + f"direct all-to-all group world size mismatch: coordinator={configured_world}, process_group={actual_world}" + ) + if backend != expected_backend: + raise RuntimeError( + f"direct all-to-all requires the {expected_backend} backend for {input_.device.type} tensors, got {backend}" + ) + configured_device = getattr(group, "device", None) + if configured_device is None: + raise RuntimeError("direct all-to-all coordinator does not declare its collective device") + expected_device = torch.device(configured_device) + if (expected_device.type != input_.device.type + or (expected_device.index is not None and expected_device.index != input_.device.index)): + raise RuntimeError( + f"direct all-to-all tensor device {input_.device} does not match coordinator device {expected_device}") + + def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor: """All-reduce the input tensor across model parallel group.""" return get_tp_group().all_reduce(input_) @@ -40,6 +70,8 @@ def sequence_model_parallel_direct_all_to_all(input_: torch.Tensor) -> torch.Ten guard keeps other callers from discovering the restriction at backward. """ group = get_sp_group() + # Packed SP is deliberately inert at SP=1 and must not require distributed + # initialization merely to preserve the identity path. if group.world_size == 1: return input_ if torch.is_grad_enabled() and input_.requires_grad: @@ -51,6 +83,10 @@ def sequence_model_parallel_direct_all_to_all(input_: torch.Tensor) -> torch.Ten f"got shape {tuple(input_.shape)} and SP={group.world_size}") if not input_.is_contiguous(): raise ValueError("direct all-to-all requires a contiguous input tensor") + # Dynamo reaches the CUDA custom op below, whose opaque runtime body + # repeats these checks. Eager and CPU/Gloo execution validate here. + if not torch.compiler.is_compiling(): + _validate_direct_all_to_all_group(group, input_, "nccl" if input_.is_cuda else "gloo") # CPU/Gloo is useful for the real multi-rank contract test. The production # packed H3 route is CUDA/Triton and takes the compiler-visible custom op. if not input_.is_cuda: diff --git a/fastvideo/distributed/parallel_state.py b/fastvideo/distributed/parallel_state.py index 49f2e7e7e3..770af32788 100644 --- a/fastvideo/distributed/parallel_state.py +++ b/fastvideo/distributed/parallel_state.py @@ -121,10 +121,41 @@ def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor: ) def direct_all_to_all_single(tensor: torch.Tensor, group_name: str) -> torch.Tensor: """Issue a synchronous all-to-all without the functional-collective wrapper.""" - assert group_name in _groups, f"Group {group_name} is not found." + if group_name not in _groups: + raise RuntimeError(f"Group {group_name} is not registered.") group = _groups[group_name]() if group is None: - raise ValueError(f"Group {group_name} is destroyed.") + raise RuntimeError(f"Group {group_name} is destroyed.") + process_group = getattr(group, "device_group", None) + if process_group is None or not torch.distributed.is_initialized(): + raise RuntimeError("direct all-to-all requires a live NCCL process group") + try: + actual_world = torch.distributed.get_world_size(process_group) + torch.distributed.get_rank(process_group) + backend = str(torch.distributed.get_backend(process_group)).lower() + except (RuntimeError, ValueError) as error: + raise RuntimeError("direct all-to-all process group is not live") from error + configured_world = int(getattr(group, "world_size", 0)) + if actual_world != configured_world: + raise RuntimeError( + f"direct all-to-all group world size mismatch: coordinator={configured_world}, process_group={actual_world}" + ) + if backend != "nccl": + raise RuntimeError(f"direct all-to-all CUDA tensors require the NCCL backend, got {backend}") + configured_device = getattr(group, "device", None) + if configured_device is None: + raise RuntimeError("direct all-to-all coordinator does not declare its CUDA device") + expected_device = torch.device(configured_device) + if (expected_device.type != "cuda" + or (expected_device.index is not None and expected_device.index != tensor.device.index)): + raise RuntimeError( + f"direct all-to-all tensor device {tensor.device} does not match coordinator device {expected_device}") + if tensor.ndim < 1 or tensor.shape[0] % actual_world: + raise ValueError( + "direct all-to-all requires the leading dimension to be evenly divisible by the live group world size; " + f"got shape {tuple(tensor.shape)} and world={actual_world}") + if not tensor.is_contiguous(): + raise ValueError("direct all-to-all requires a contiguous input tensor") output = torch.empty_like(tensor) torch.distributed.all_to_all_single(output, tensor, group=group.device_group) return output diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index e430c486dd..9b5b1b241a 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -427,12 +427,14 @@ def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]: return temb.view(-1, 6 * self.hidden_size).chunk(6, dim=-1) -class _MiniMaxH3StepCursor: - - __slots__ = ("signature", "step") +class _MiniMaxH3StepCursor(nn.Module): def __init__(self, device: torch.device, signature: tuple[tuple[float, ...], ...]) -> None: - self.step = torch.zeros((), dtype=torch.long, device=device) + super().__init__() + # The cursor is runtime-only state shared by every precomputed block. + # Registering it makes model.to(...) follow the transformer lifecycle, + # while persistent=False keeps it out of checkpoint state. + self.register_buffer("step", torch.zeros((), dtype=torch.long, device=device), persistent=False) self.signature = signature def set(self, index: int) -> None: @@ -784,7 +786,13 @@ def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None: @torch.no_grad() def prepare_adaln_trajectory(self, row_timestep_plan: list[tuple[torch.Tensor, torch.Tensor]]) -> dict[str, float | int]: - """Replace full-rank per-block AdaLN projections by exact per-step tables.""" + """Replace full-rank per-block AdaLN projections by exact per-step tables. + + This is an irreversible inference transformation. The tables and their + device-following cursor are intentionally absent from ``state_dict``; + after installation, the transformed module cannot be checkpointed and + reloaded as either the stock model or an equivalent cached model. + """ signature = tuple(tuple(float(value) for value in timestep.detach().cpu().flatten()) for timestep, _ in row_timestep_plan) cursor = getattr(self, "_h3_adaln_cursor", None) diff --git a/fastvideo/models/dits/minimax_h3_fusions/relayout.py b/fastvideo/models/dits/minimax_h3_fusions/relayout.py index 4c203f1d4f..936e327b39 100644 --- a/fastvideo/models/dits/minimax_h3_fusions/relayout.py +++ b/fastvideo/models/dits/minimax_h3_fusions/relayout.py @@ -113,10 +113,18 @@ def pack_qkv_destination_major(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor """Move three ``(rows, heads, dim)`` tensors into destination-major QKV in one pass.""" if not HAVE_TRITON: raise RuntimeError("MiniMax-H3 packed sequence parallelism requires Triton") + if world < 1: + raise ValueError(f"sequence parallel world size must be positive, got {world}") if q.ndim != 3 or q.shape != k.shape or q.shape != v.shape: raise ValueError("q, k, and v must have matching (rows, heads, head_dim) shapes") + if q.device != k.device or q.device != v.device: + raise ValueError(f"q, k, and v must be on one device, got {q.device}, {k.device}, and {v.device}") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError(f"q, k, and v must have one dtype, got {q.dtype}, {k.dtype}, and {v.dtype}") + if not q.is_cuda: + raise ValueError("MiniMax-H3 packed QKV relayout requires CUDA tensors") if q.shape[1] % world: - raise ValueError(f"heads ({q.shape[1]}) must divide sequence parallel size ({world})") + raise ValueError(f"heads ({q.shape[1]}) must be divisible by sequence parallel size ({world})") if any(t.stride(-1) != 1 for t in (q, k, v)): raise ValueError("q, k, and v must be contiguous in head_dim") return _pack_qkv_op(q, k, v, world) @@ -128,6 +136,8 @@ def merge_heads(x: torch.Tensor) -> torch.Tensor: raise RuntimeError("MiniMax-H3 packed sequence parallelism requires Triton") if x.ndim != 4 or not x.is_contiguous(): raise ValueError("packed all-to-all output must be a contiguous 4D tensor") + if not x.is_cuda: + raise ValueError("MiniMax-H3 packed head relayout requires a CUDA tensor") world, rows, heads_local, head_dim = x.shape return _merge_heads_op(x).reshape(rows, world * heads_local, head_dim) diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py index 8e687cd766..ab81691aba 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -100,6 +100,14 @@ def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> V result.add_check("step_index", batch.step_index, V.non_negative_int) return result + def _release_denoising_resources(self, fastvideo_args: FastVideoArgs, full_cpu_offload: bool) -> None: + if bool(getattr(fastvideo_args, "dit_layerwise_offload", False)): + manager = getattr(self.transformer, "_layerwise_offload_manager", None) + if manager is not None and getattr(manager, "enabled", False): + manager.release_all() + if full_cpu_offload: + self.transformer.to("cpu") + @torch.no_grad() def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: """Denoise the packed H3 video and audio streams over one shared schedule.""" @@ -112,61 +120,66 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward full_cpu_offload = (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload and not fastvideo_args.use_fsdp_inference) device = get_local_torch_device() - if full_cpu_offload: - self.transformer.to(device) - batch.latents = batch.latents.to(device) - batch.audio_latents = batch.audio_latents.to(device) - - self.scheduler.set_timesteps(batch.num_inference_steps, device=device) - self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device) - video_timesteps = self.scheduler.timesteps - audio_timesteps = self.audio_scheduler.timesteps - if video_timesteps is None or audio_timesteps is None: - raise ValueError("MiniMax-H3 schedulers did not produce timesteps.") - if len(video_timesteps) != len(audio_timesteps): - raise ValueError("MiniMax-H3 video and audio schedules must have the same number of intervals.") - - row_timestep_plan = [] - for video_timestep, audio_timestep in zip(video_timesteps, audio_timesteps, strict=True): - video_value = float(video_timestep.item()) - audio_value = float(audio_timestep.item()) - unique, inverse = build_row_timesteps( - layout, - video_timestep=video_value, - audio_timestep=audio_value, - condition_video_timestep=max(video_value, MINIMAX_H3_KEYFRAME_NOISE_AUG), - condition_audio_timestep=1.0, - ) - row_timestep_plan.append((unique.to(device), inverse.to(device))) - batch.timesteps = video_timesteps - - if bool(getattr(self.transformer, "adaln_precompute_enabled", False)): - _validate_adaln_precompute_configuration(fastvideo_args) - self.transformer.prepare_adaln_trajectory(row_timestep_plan) - - position_ids = layout.position_ids.to(device) - token_tags = layout.token_tags.to(device) - video_indices = layout.video_indices.to(device) - audio_indices = layout.audio_indices.to(device) - text_indices = layout.text_indices.to(device) - prompt_embeds = batch.prompt_embeds[0].to(device) - - vsa_metadata_builder = _h3_vsa_metadata_builder(self.transformer, fastvideo_args) - if vsa_metadata_builder is not None: - vsa_patch_size = fastvideo_args.pipeline_config.dit_config.patch_size - vsa_prefix_segments = _h3_vsa_prefix_segments(layout, vsa_patch_size) - # Per-request knobs (sweeps flip these between generate_video calls - # without respawning workers); mode None defers to the env default. - vsa_mode = batch.extra.get("vsa_mode", "exempt") - if vsa_mode not in ("exempt", "compete"): - raise ValueError(f"vsa_mode must be 'exempt' or 'compete', got {vsa_mode!r}.") - vsa_exempt = vsa_mode == "exempt" - vsa_dense_layers = tuple(batch.extra.get("vsa_dense_layers", ())) - vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0)) - # Run-level tile geometry (256 default, 64 = native Triton path), - # plumbed like the run-level sparsity; the builder validates the - # value against VSA_H3_TILE_SHAPES. - vsa_tile_size = int(fastvideo_args.VSA_tile_size) + preparation_complete = False + try: + if full_cpu_offload: + self.transformer.to(device) + batch.latents = batch.latents.to(device) + batch.audio_latents = batch.audio_latents.to(device) + self.scheduler.set_timesteps(batch.num_inference_steps, device=device) + self.audio_scheduler.set_timesteps(batch.num_inference_steps, device=device) + video_timesteps = self.scheduler.timesteps + audio_timesteps = self.audio_scheduler.timesteps + if video_timesteps is None or audio_timesteps is None: + raise ValueError("MiniMax-H3 schedulers did not produce timesteps.") + if len(video_timesteps) != len(audio_timesteps): + raise ValueError("MiniMax-H3 video and audio schedules must have the same number of intervals.") + + row_timestep_plan = [] + for video_timestep, audio_timestep in zip(video_timesteps, audio_timesteps, strict=True): + video_value = float(video_timestep.item()) + audio_value = float(audio_timestep.item()) + unique, inverse = build_row_timesteps( + layout, + video_timestep=video_value, + audio_timestep=audio_value, + condition_video_timestep=max(video_value, MINIMAX_H3_KEYFRAME_NOISE_AUG), + condition_audio_timestep=1.0, + ) + row_timestep_plan.append((unique.to(device), inverse.to(device))) + batch.timesteps = video_timesteps + + if bool(getattr(self.transformer, "adaln_precompute_enabled", False)): + _validate_adaln_precompute_configuration(fastvideo_args) + self.transformer.prepare_adaln_trajectory(row_timestep_plan) + + position_ids = layout.position_ids.to(device) + token_tags = layout.token_tags.to(device) + video_indices = layout.video_indices.to(device) + audio_indices = layout.audio_indices.to(device) + text_indices = layout.text_indices.to(device) + prompt_embeds = batch.prompt_embeds[0].to(device) + + vsa_metadata_builder = _h3_vsa_metadata_builder(self.transformer, fastvideo_args) + if vsa_metadata_builder is not None: + vsa_patch_size = fastvideo_args.pipeline_config.dit_config.patch_size + vsa_prefix_segments = _h3_vsa_prefix_segments(layout, vsa_patch_size) + # Per-request knobs (sweeps flip these between generate_video calls + # without respawning workers); mode None defers to the env default. + vsa_mode = batch.extra.get("vsa_mode", "exempt") + if vsa_mode not in ("exempt", "compete"): + raise ValueError(f"vsa_mode must be 'exempt' or 'compete', got {vsa_mode!r}.") + vsa_exempt = vsa_mode == "exempt" + vsa_dense_layers = tuple(batch.extra.get("vsa_dense_layers", ())) + vsa_dense_first_n = int(batch.extra.get("vsa_dense_first_n_steps", 0)) + # Run-level tile geometry (256 default, 64 = native Triton path), + # plumbed like the run-level sparsity; the builder validates the + # value against VSA_H3_TILE_SHAPES. + vsa_tile_size = int(fastvideo_args.VSA_tile_size) + preparation_complete = True + finally: + if not preparation_complete: + self._release_denoising_resources(fastvideo_args, full_cpu_offload) try: # The stage range groups the complete denoising loop while the @@ -235,12 +248,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward batch.step_index = index batch.timestep = video_timestep finally: - if bool(getattr(fastvideo_args, "dit_layerwise_offload", False)): - manager = getattr(self.transformer, "_layerwise_offload_manager", None) - if manager is not None and getattr(manager, "enabled", False): - manager.release_all() - if full_cpu_offload: - self.transformer.to("cpu") + self._release_denoising_resources(fastvideo_args, full_cpu_offload) return batch diff --git a/fastvideo/tests/contract/test_ci_test_collection.py b/fastvideo/tests/contract/test_ci_test_collection.py index 4a0e0b1798..6d00fb657c 100644 --- a/fastvideo/tests/contract/test_ci_test_collection.py +++ b/fastvideo/tests/contract/test_ci_test_collection.py @@ -62,7 +62,6 @@ ALLOWLIST = { "attention": "no lane yet — GPU attention-backend tests, run manually", "audio": "no lane yet — audio encoder tests, run manually", - "distributed": "no lane yet — multi-GPU torchrun tests, run manually", "hooks": "no lane yet — run manually", "layers": "no lane yet — torchrun FSDP dispatch tests, run manually", "nightly": "by design: nightly cadence, not per-PR", @@ -197,6 +196,20 @@ def test_gpu_tests_preserve_a_launcher_assigned_rendezvous_port(): assert not overwrites, f"Tests overwrite the CI runner's per-lease MASTER_PORT: {overwrites}" +def test_minimax_h3_world4_contract_has_portable_and_strict_cuda_owners(): + transformer_lane = (REPO_ROOT / ".buildkite/scripts/lanes/transformer.sh").read_text() + ssim_lane = (REPO_ROOT / ".buildkite/scripts/lanes/ssim.sh").read_text() + contract = (TESTS_ROOT / "distributed/test_minimax_h3_packed_sp.py").read_text() + + assert "fastvideo/tests/distributed/test_minimax_h3_packed_sp.py" in transformer_lane + assert "fastvideo/tests/distributed/test_minimax_h3_packed_sp.py" in ssim_lane + assert "FASTVIDEO_MINIMAX_H3_PACKED_SP_STRICT_CUDA=1" in ssim_lane + assert 'os.environ.get("MASTER_PORT")' in contract + assert 'launcher.append(f"--master_port={inherited_master_port}")' in contract + assert 'launcher.append("--standalone")' in contract + assert "socket" not in contract + + def test_dreamverse_lane_keeps_arm64_browser_coverage_explicit(): lane = (REPO_ROOT / ".buildkite/scripts/lanes/dreamverse.sh").read_text() @@ -256,6 +269,8 @@ def test_ssim_lane_uses_the_local_four_gpu_scheduler(): assert "FASTVIDEO_SSIM_TEST_FILES" in lane_script assert 'if [ "${TEST_SCOPE:-}" = merge ]; then' in lane_script assert "Missing FASTVIDEO_SSIM_TEST_FILES for merge scope" in lane_script + assert "FASTVIDEO_MINIMAX_H3_PACKED_SP_STRICT_CUDA=1" in lane_script + assert "test_minimax_h3_packed_sp.py" in lane_script def test_golden_lane_accepts_only_focused_test_basenames(): diff --git a/fastvideo/tests/contract/test_merge_ci_plan.py b/fastvideo/tests/contract/test_merge_ci_plan.py index e7f23f841d..e421a5a062 100644 --- a/fastvideo/tests/contract/test_merge_ci_plan.py +++ b/fastvideo/tests/contract/test_merge_ci_plan.py @@ -51,6 +51,15 @@ def test_changed_ssim_test_selects_only_that_file(): assert plan.encoded_ssim_tests() == "test_flux_t2i_similarity.py" +def test_h3_packed_sp_contract_selects_existing_strict_four_gpu_lane(): + plan = PLAN_MERGE_CI.classify_paths([ + "fastvideo/tests/distributed/test_minimax_h3_packed_sp.py" + ]) + + assert plan.encoded_lanes() == ",ssim," + assert plan.encoded_ssim_tests() == "test_minimax_h3_similarity.py" + + def test_shared_golden_harness_or_reference_requires_full_golden_lane(): plan = PLAN_MERGE_CI.classify_paths(["fastvideo/tests/golden_gate/_harness.py"]) diff --git a/fastvideo/tests/contract/test_modal_fa4_policy.py b/fastvideo/tests/contract/test_modal_fa4_policy.py index ff64b84a90..97a5ea31ad 100644 --- a/fastvideo/tests/contract/test_modal_fa4_policy.py +++ b/fastvideo/tests/contract/test_modal_fa4_policy.py @@ -62,7 +62,8 @@ def test_pr_model_load_and_training_lanes_disable_fa4(): # scripts used by the Slurm runner. lanes = { "run_transformer_tests": - ("pytest ./fastvideo/tests/transformers -vs", ".buildkite/scripts/lanes/transformer.sh"), + ("pytest ./fastvideo/tests/transformers ./fastvideo/tests/distributed/test_minimax_h3_packed_sp.py -vs", + ".buildkite/scripts/lanes/transformer.sh"), "run_training_tests": ("pytest ./fastvideo/tests/training/Vanilla -srP", None), "run_training_lora_tests": ("pytest ./fastvideo/tests/training/lora/test_lora_training.py -srP", None), "run_training_tests_VSA": ("pytest ./fastvideo/tests/training/VSA -srP", None), diff --git a/fastvideo/tests/distributed/test_minimax_h3_packed_sp.py b/fastvideo/tests/distributed/test_minimax_h3_packed_sp.py index 90503d9de9..dc19a27b2e 100644 --- a/fastvideo/tests/distributed/test_minimax_h3_packed_sp.py +++ b/fastvideo/tests/distributed/test_minimax_h3_packed_sp.py @@ -5,7 +5,6 @@ import contextlib import os -import socket import subprocess import sys from pathlib import Path @@ -21,17 +20,18 @@ class _TestGroup: - def __init__(self, world_size: int) -> None: + def __init__(self, world_size: int, device: torch.device) -> None: self.world_size = world_size + self.device = device self.device_group = dist.group.WORLD self.unique_name = "minimax_h3_packed_sp_test" -class _IdentityAttentionImpl(nn.Module): +class _QKVOracleAttentionImpl(nn.Module): def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, metadata: object) -> torch.Tensor: - del k, v, metadata - return q + del metadata + return q + 2 * k + 3 * v def postprocess_output(self, output: torch.Tensor, metadata: object) -> torch.Tensor: del metadata @@ -58,6 +58,16 @@ def _reference_merge(output: torch.Tensor) -> torch.Tensor: return output.permute(1, 0, 2, 3).contiguous().reshape(rows, world * heads_local, head_dim) +def _exact_cuda_route(world: int) -> bool: + strict_cuda = os.environ.get("FASTVIDEO_MINIMAX_H3_PACKED_SP_STRICT_CUDA", "0") == "1" + visible_cuda_devices = torch.cuda.device_count() + if strict_cuda and visible_cuda_devices < world: + raise RuntimeError( + "strict MiniMax-H3 packed-SP preflight requires one visible CUDA device per rank; " + f"world={world}, visible={visible_cuda_devices}") + return visible_cuda_devices >= world + + def _worker() -> None: from fastvideo.attention.layer import DistributedAttention from fastvideo.distributed import communication_op, parallel_state @@ -67,22 +77,22 @@ def _worker() -> None: world = int(os.environ["WORLD_SIZE"]) rank = int(os.environ["RANK"]) local_rank = int(os.environ["LOCAL_RANK"]) - exact_cuda_route = torch.cuda.device_count() >= world - backend = "nccl" if exact_cuda_route else "gloo" - dist.init_process_group(backend=backend) + exact_cuda_route = _exact_cuda_route(world) if exact_cuda_route: torch.cuda.set_device(local_rank) device = torch.device(f"cuda:{local_rank}") dtype = torch.bfloat16 + dist.init_process_group(backend="nccl", device_id=device) else: device = torch.device("cpu") dtype = torch.float32 + dist.init_process_group(backend="gloo") - group = _TestGroup(world) + group = _TestGroup(world, device) parallel_state._register_group(group) attention = DistributedAttention.__new__(DistributedAttention) nn.Module.__init__(attention) - attention.attn_impl = _IdentityAttentionImpl() + attention.attn_impl = _QKVOracleAttentionImpl() attention.head_size = 16 attention.packed_qkv_relayout = True attention._compile_forward_enabled = True @@ -93,7 +103,7 @@ def _worker() -> None: k = q + 100 v = q + 200 semantic_rows = world * rows_local - 3 - expected = q.clone() + expected = q + 2 * k + 3 * v if rank == world - 1: expected[:, -3:] = 0 @@ -131,32 +141,44 @@ def _worker() -> None: dist.destroy_process_group() -def _free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind(("127.0.0.1", 0)) - return int(sock.getsockname()[1]) - - def test_minimax_h3_packed_sp_world4_collective_contract() -> None: environment = dict(os.environ, OMP_NUM_THREADS="1") + launcher = [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc_per_node={SP_WORLD_SIZE}", + ] + inherited_master_port = os.environ.get("MASTER_PORT") + if inherited_master_port: + launcher.append(f"--master_port={inherited_master_port}") + else: + # Local invocations have no scheduler-assigned port. Let torchrun own + # discovery instead of racing a bind-and-close port probe. + launcher.append("--standalone") + launcher.extend([ + str(Path(__file__).resolve()), + "--worker", + ]) process = subprocess.run( - [ - sys.executable, - "-m", - "torch.distributed.run", - f"--nproc_per_node={SP_WORLD_SIZE}", - f"--master_port={_free_port()}", - str(Path(__file__).resolve()), - "--worker", - ], + launcher, env=environment, capture_output=True, text=True, timeout=300, ) - assert process.returncode == 0 and "MINIMAX_H3_PACKED_SP_OK" in process.stdout, ( + strict_cuda = os.environ.get("FASTVIDEO_MINIMAX_H3_PACKED_SP_STRICT_CUDA", "0") == "1" + expected_mode = "mode=cuda-production" if strict_cuda else "MINIMAX_H3_PACKED_SP_OK" + assert process.returncode == 0 and expected_mode in process.stdout, ( f"stdout:\n{process.stdout[-6000:]}\nstderr:\n{process.stderr[-6000:]}") +def test_minimax_h3_packed_sp_strict_mode_rejects_insufficient_cuda(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FASTVIDEO_MINIMAX_H3_PACKED_SP_STRICT_CUDA", "1") + monkeypatch.setattr(torch.cuda, "device_count", lambda: SP_WORLD_SIZE - 1) + with pytest.raises(RuntimeError, match="one visible CUDA device per rank"): + _exact_cuda_route(SP_WORLD_SIZE) + + if __name__ == "__main__" and "--worker" in sys.argv: _worker() diff --git a/fastvideo/tests/modal/test_pr_test.py b/fastvideo/tests/modal/test_pr_test.py index f2360bfcc3..8f1b6fe9df 100644 --- a/fastvideo/tests/modal/test_pr_test.py +++ b/fastvideo/tests/modal/test_pr_test.py @@ -370,7 +370,8 @@ def test_wave1_lane_functions_use_shared_scripts(monkeypatch): "golden_gate.sh": 'exec pytest "$golden_root" -vs', "encoder.sh": "pytest ./fastvideo/tests/encoders -vs", "vae.sh": "pytest ./fastvideo/tests/vaes -vs", - "transformer.sh": "pytest ./fastvideo/tests/transformers -vs", + "transformer.sh": + "pytest ./fastvideo/tests/transformers ./fastvideo/tests/distributed/test_minimax_h3_packed_sp.py -vs", "inference_lora.sh": "pytest ./fastvideo/tests/inference/lora/test_lora_inference_similarity.py -vs", "distillation_dmd.sh": "pytest ./fastvideo/tests/training/distill/test_distill_dmd.py -vs", "train_framework.sh": "pytest ./fastvideo/tests/train/models ./fastvideo/tests/train/methods -vs", diff --git a/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py b/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py index 6600a063cb..b198487034 100644 --- a/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py +++ b/fastvideo/tests/transformers/test_minimax_h3_sol_optimizations.py @@ -60,6 +60,45 @@ def __init__(self) -> None: self.adaln_rank = None +class _TrackingTrajectoryTransformer(_TrajectoryTransformer): + + adaln_precompute_enabled = True + + def __init__(self) -> None: + super().__init__() + self.device_moves: list[torch.device | str] = [] + + def to(self, device: torch.device | str, *args: object, **kwargs: object) -> "_TrackingTrajectoryTransformer": + self.device_moves.append(device) + return super().to(device, *args, **kwargs) + + def prepare_adaln_trajectory(self, plan: list[tuple[torch.Tensor, torch.Tensor]]) -> dict[str, float | int]: + from fastvideo.models.dits.minimax_h3 import MiniMaxH3Transformer3DModel + + return MiniMaxH3Transformer3DModel.prepare_adaln_trajectory(self, plan) + + +class _StageLayout: + + def __init__(self) -> None: + self.position_ids = torch.zeros(1, dtype=torch.long) + self.token_tags = torch.zeros(1, dtype=torch.long) + self.video_indices = torch.zeros(1, dtype=torch.long) + self.audio_indices = torch.zeros(1, dtype=torch.long) + self.text_indices = torch.zeros(1, dtype=torch.long) + + +class _StageScheduler: + + def __init__(self, values: tuple[float, ...] = (1.0, )) -> None: + self.timesteps: torch.Tensor | None = None + self.values = values + + def set_timesteps(self, steps: int, device: torch.device) -> None: + del steps + self.timesteps = torch.tensor(self.values, device=device) + + class _IdentityAttentionImpl(nn.Module): def preprocess_qkv(self, qkv: torch.Tensor, metadata: object) -> torch.Tensor: @@ -113,6 +152,23 @@ def test_minimax_h3_relayout_is_bit_exact_and_compile_safe() -> None: assert torch.equal(actual_merge, expected_merge) +def test_minimax_h3_relayout_validates_world_device_dtype_and_cuda() -> None: + from fastvideo.models.dits.minimax_h3_fusions import relayout + + q = torch.zeros(2, 4, 8) + with patch.object(relayout, "HAVE_TRITON", True): + with pytest.raises(ValueError, match="world size must be positive"): + relayout.pack_qkv_destination_major(q, q, q, 0) + with pytest.raises(ValueError, match="one device"): + relayout.pack_qkv_destination_major(q, torch.zeros_like(q, device="meta"), q, 2) + with pytest.raises(ValueError, match="one dtype"): + relayout.pack_qkv_destination_major(q, q.double(), q, 2) + with pytest.raises(ValueError, match="requires CUDA"): + relayout.pack_qkv_destination_major(q, q, q, 2) + with pytest.raises(ValueError, match="requires a CUDA tensor"): + relayout.merge_heads(torch.zeros(2, 2, 2, 8)) + + def test_adaln_precompute_rejects_a_different_trajectory() -> None: from fastvideo.models.dits.minimax_h3 import ( MiniMaxH3Transformer3DModel, @@ -166,6 +222,36 @@ def test_adaln_precompute_matches_projection_tables_and_reuses_schedule() -> Non MiniMaxH3Transformer3DModel.set_adaln_step(transformer, 2) +def test_adaln_precompute_cursor_follows_transformer_device_and_is_nonpersistent() -> None: + from fastvideo.models.dits.minimax_h3 import MiniMaxH3Transformer3DModel + + transformer = _TrajectoryTransformer() + MiniMaxH3Transformer3DModel.prepare_adaln_trajectory( + transformer, [(torch.tensor([1.0]), torch.tensor([0]))]) + + named_buffers = dict(transformer.named_buffers()) + assert any(name.endswith("cursor.step") for name in named_buffers) + assert not any(name.endswith("cursor.step") for name in transformer.state_dict()) + transformer.to("meta") + assert transformer._h3_adaln_cursor.step.device.type == "meta" + + +def test_adaln_precompute_state_dict_is_explicitly_not_reloadable() -> None: + from fastvideo.models.dits.minimax_h3 import MiniMaxH3Transformer3DModel + + transformer = _TrajectoryTransformer() + original_keys = set(transformer.state_dict()) + assert any("adaln_proj.linear" in key for key in original_keys) + MiniMaxH3Transformer3DModel.prepare_adaln_trajectory( + transformer, [(torch.tensor([1.0]), torch.tensor([0]))]) + transformed_state = transformer.state_dict() + + assert not any("adaln_proj.linear" in key for key in transformed_state) + assert not any("adaln_proj.table" in key or "cursor.step" in key for key in transformed_state) + with pytest.raises(RuntimeError, match="Missing key"): + _TrajectoryTransformer().load_state_dict(transformed_state) + + def test_adaln_precompute_failure_does_not_partially_replace_blocks() -> None: from fastvideo.models.dits.minimax_h3 import MiniMaxH3Transformer3DModel @@ -184,6 +270,73 @@ def fail_projection(embeddings: torch.Tensor) -> tuple[torch.Tensor, ...]: assert not hasattr(transformer, "_h3_adaln_cursor") +@pytest.mark.parametrize( + ("failure", "message"), + [ + ("latent_transfer", "injected latent transfer failure"), + ("projection", "injected projection failure"), + ("rank_reduced", "stock full-rank"), + ("schedule_reuse", "same denoising schedule"), + ("scheduler_mismatch", "same number of intervals"), + ], +) +def test_denoising_preparation_failures_restore_full_cpu_offload(failure: str, message: str) -> None: + from fastvideo.models.dits.minimax_h3 import _MiniMaxH3StepCursor + from fastvideo.pipelines.basic.minimax_h3.stages import minimax_h3_denoising as denoising + + transformer = _TrackingTrajectoryTransformer() + if failure == "projection": + def fail_projection(embeddings: torch.Tensor) -> tuple[torch.Tensor, ...]: + del embeddings + raise RuntimeError("injected projection failure") + + transformer.transformer_blocks[1].adaln_proj.forward = fail_projection + elif failure == "rank_reduced": + transformer.adaln_rank = 4 + elif failure == "schedule_reuse": + transformer._h3_adaln_cursor = _MiniMaxH3StepCursor(torch.device("cpu"), ((0.5, ), )) + + audio_scheduler = _StageScheduler((1.0, 0.5)) if failure == "scheduler_mismatch" else _StageScheduler() + stage = denoising.MiniMaxH3DenoisingStage(transformer, _StageScheduler(), audio_scheduler) + layout = _StageLayout() + batch = SimpleNamespace( + extra={denoising.MINIMAX_H3_LAYOUT_KEY: layout}, + prompt_embeds=[torch.zeros(1, 1, 2)], + latents=torch.zeros(1, 2), + audio_latents=torch.zeros(1, 2), + num_inference_steps=1, + ) + if failure == "latent_transfer": + + class _FailingLatents: + + def to(self, _device: torch.device) -> torch.Tensor: + raise RuntimeError("injected latent transfer failure") + + batch.latents = _FailingLatents() + args = SimpleNamespace( + dit_cpu_offload=True, + dit_layerwise_offload=False, + use_fsdp_inference=False, + ) + with ( + patch.object(denoising, "MiniMaxH3PackedLayout", _StageLayout), + patch.object(denoising, "get_local_torch_device", return_value=torch.device("cpu")), + patch.object( + denoising, + "build_row_timesteps", + side_effect=lambda *_args, **_kwargs: (torch.tensor([1.0]), torch.tensor([0])), + ), + patch.object(denoising, "_h3_vsa_metadata_builder", return_value=None), + pytest.raises((RuntimeError, ValueError), match=message), + ): + stage.forward(batch, args) + + assert transformer.device_moves[0] == torch.device("cpu") + assert transformer.device_moves[-1] == "cpu" + assert len(transformer.device_moves) == 2 + + @pytest.mark.parametrize( ("layerwise", "fsdp", "message"), [ @@ -252,6 +405,32 @@ def test_packed_sp_falls_back_to_autograd_aware_collective_when_grad_enabled() - direct.assert_not_called() +def test_packed_sp_rejects_zero_semantic_rows_instead_of_treating_them_as_unspecified() -> None: + from fastvideo.attention.layer import DistributedAttention + from fastvideo.forward_context import set_forward_context + + attention = DistributedAttention.__new__(DistributedAttention) + nn.Module.__init__(attention) + attention.attn_impl = _IdentityAttentionImpl() + attention.head_size = 2 + attention.packed_qkv_relayout = True + attention._compile_forward_enabled = True + q = torch.randn(1, 2, 4, 2) + + packed = torch.zeros(2, 2, 2, 6) + with ( + patch("fastvideo.attention.layer.get_sp_world_size", return_value=2), + patch("fastvideo.attention.layer.get_sp_parallel_rank", return_value=0), + patch("fastvideo.models.dits.minimax_h3_fusions.relayout.pack_qkv_destination_major", + return_value=packed), + patch("fastvideo.attention.layer.sequence_model_parallel_direct_all_to_all", side_effect=lambda tensor: tensor), + torch.inference_mode(), + set_forward_context(current_timestep=0, attn_metadata=None), + pytest.raises(ValueError, match=r"original_seq_len must be in \[1, 4\], got 0"), + ): + attention(q, q, q, original_seq_len=0) + + def test_direct_packed_collective_rejects_unsupported_inputs_before_launch() -> None: from fastvideo.distributed.communication_op import sequence_model_parallel_direct_all_to_all @@ -265,6 +444,57 @@ def test_direct_packed_collective_rejects_unsupported_inputs_before_launch() -> sequence_model_parallel_direct_all_to_all(torch.randn(2, 4).transpose(0, 1)) +def test_direct_packed_collective_is_an_identity_at_sp1_without_a_process_group() -> None: + from fastvideo.distributed.communication_op import sequence_model_parallel_direct_all_to_all + + tensor = torch.randn(3, 2) + group = SimpleNamespace(world_size=1, device_group=None) + with patch("fastvideo.distributed.communication_op.get_sp_group", return_value=group): + assert sequence_model_parallel_direct_all_to_all(tensor) is tensor + + +@pytest.mark.parametrize( + ("actual_world", "backend", "device", "message"), + [ + (2, "nccl", torch.device("cpu"), "requires the gloo backend"), + (3, "gloo", torch.device("cpu"), "world size mismatch"), + (2, "gloo", torch.device("cuda:0"), "does not match coordinator device"), + (2, "gloo", None, "does not declare its collective device"), + ], +) +def test_direct_packed_collective_validates_live_group_backend_world_and_device( + actual_world: int, backend: str, device: torch.device | None, message: str) -> None: + from fastvideo.distributed.communication_op import sequence_model_parallel_direct_all_to_all + + process_group = object() + group = SimpleNamespace(world_size=2, device_group=process_group, device=device) + tensor = torch.randn(2, 3) + with ( + patch("fastvideo.distributed.communication_op.get_sp_group", return_value=group), + patch("torch.distributed.is_initialized", return_value=True), + patch("torch.distributed.get_world_size", return_value=actual_world), + patch("torch.distributed.get_rank", return_value=0), + patch("torch.distributed.get_backend", return_value=backend), + torch.no_grad(), + pytest.raises(RuntimeError, match=message), + ): + sequence_model_parallel_direct_all_to_all(tensor) + + +def test_direct_packed_collective_rejects_missing_or_dead_process_group() -> None: + from fastvideo.distributed.communication_op import sequence_model_parallel_direct_all_to_all + + tensor = torch.randn(2, 3) + group = SimpleNamespace(world_size=2, device_group=None, device=torch.device("cpu")) + with ( + patch("fastvideo.distributed.communication_op.get_sp_group", return_value=group), + patch("torch.distributed.is_initialized", return_value=False), + torch.no_grad(), + pytest.raises(RuntimeError, match="live distributed process group"), + ): + sequence_model_parallel_direct_all_to_all(tensor) + + def test_direct_packed_collective_captures_as_a_fullgraph_custom_op(tmp_path) -> None: if not torch.cuda.is_available(): pytest.skip("CUDA is required") @@ -276,6 +506,8 @@ def test_direct_packed_collective_captures_as_a_fullgraph_custom_op(tmp_path) -> class Group: unique_name = "minimax_h3_direct_compile_test" device_group = None + world_size = 1 + device = torch.device("cuda:0") group = Group() torch.distributed.init_process_group(