diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 0723cc97b..a72960cc7 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.buildkite/scripts/unit_test.sh @@ -29,6 +29,7 @@ exec pytest \ ./fastvideo/tests/training/test_trackers.py \ ./fastvideo/tests/training/test_ltx2_rope_fps.py \ ./fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py \ + ./fastvideo/tests/inference/test_basic_fasth3_token_reduction.py \ ./fastvideo/tests/inference/test_inference_regional_compile.py \ ./fastvideo/tests/attention/ \ ./fastvideo/tests/layers/test_pdd_linear.py \ diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index d878de1e3..678fff1f8 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -83,6 +83,7 @@ surfaces: vae_parallel_decode: "MiniMax-H3 sequence-parallel VAE decode opt-in; model-specific optimization not yet represented in the typed public schema." h3_sequential_load: "MiniMax-H3 sequential text-encoder then DiT/VAE load; model-specific optimization not yet represented in the typed public schema." h3_vsa_fp4: "MiniMax-H3 native FP4 attention options carried through generator.pipeline.experimental.h3_vsa_fp4 rather than a dedicated typed component setting." + h3_token_reduction: "MiniMax-H3 training-free token reduction options carried through generator.pipeline.experimental.h3_token_reduction rather than a dedicated typed component setting." h3_encoder_split: "MiniMax-H3 Ray-only split that runs the Qwen3-VL text encoder on dedicated nodes and the DiT/VAE on the rest; model-specific optimization not yet represented in the typed public schema." h3_encoder_nodes: "Leading node count that h3_encoder_split reserves for the MiniMax-H3 text encoder; model-specific optimization not yet represented in the typed public schema." video_decode_backend: "MiniMax-H3 video decoder selection (full VAE vs TAEH3 preview); model-specific optimization not yet represented in the typed public schema." diff --git a/docs/inference/fasth3-distilled.md b/docs/inference/fasth3-distilled.md index 5c9b9d958..0b497a085 100644 --- a/docs/inference/fasth3-distilled.md +++ b/docs/inference/fasth3-distilled.md @@ -198,6 +198,57 @@ forwards. With the DiT sharded, peak allocated memory is about 66 GiB per GPU (about 95 GiB in use), against 83 GiB (103 GiB) with a full DiT copy on each GPU; the output is identical. Model loading and decoding come on top of this. +## Token reduction + +`--token-reduction` is an opt-in, training-free speed mode that follows the +`--token-reduction` mode of [antirez/h3.c](https://github.com/antirez/h3.c). +Inside DiT blocks 4 to 29, every two horizontally adjacent tokens of the +generated video are averaged into one, so those blocks see the video at half +width. During the first 4 of the 8 forwards, where the coarse layout forms, +the interval extends to block 39. Before the first full-width block, each +token gets back its own value plus the update its pair received. Text, audio, +reference and condition tokens are never paired. + +FastH3 8-Step V2 and the OmniRef PDD export both support it: + +```bash +python examples/inference/basic/basic_fasth3_8step.py \ + --prompt 'Two people walk along a beach at sunset, talking and laughing' \ + --num-gpus 4 --token-reduction + +python examples/inference/basic/basic_fasth3_omniref_pdd.py \ + --model-path --image reference.png \ + --prompt ' ...' --num-gpus 4 --token-reduction +``` + +In a recipe, set `pipeline.experimental.h3_token_reduction` to +`{enabled: true}`. The optional keys `blocks`, `early` and `update_scale` +match the `--token-reduction-blocks`, `--token-reduction-early` and +`--token-reduction-scale` flags. Malformed values, and intervals past the +checkpoint's DiT depth, fail before any weights load. + +Measured on one GB200 tray with 4 GPUs, 1344x768, 124 frames and BF16 +weights, after warm-up: + +| Checkpoint | Denoising, 8 forwards | Whole request | +|---|---|---| +| FastH3 8-Step V2, default profile with regional compile | 5.08 s to 3.53 s, 30.6% less | 9.13 s to 7.58 s, 17.0% less | +| FastH3 8-Step V2, regional compile off | 6.10 s to 4.26 s, 30.2% less | 10.18 s to 8.41 s, 17.4% less | +| OmniRef PDD with one image reference, eager | 12.07 s to 9.33 s, 22.7% less | 24.67 s to 21.62 s, 12.3% less | + +The saving follows how many rows the middle blocks drop. A V2 request has a +short text and audio prefix, so the reduced blocks see about half the rows. +An OmniRef request also carries the reference tokens, which stay unpaired, so +it saves less. Decoding takes the same time either way. + +The reduced video keeps the quality of the unreduced one, but it is not the +same video. Most scenes keep their overall layout, while details can change, +such as one of two people in a scene or the venue of a concert. The +checkpoints were not trained with token reduction, and fine-tuning with it +enabled is the expected way to keep the original shot. Only the GB200 BF16 +path has been validated so far; the NVFP4 and FP8 consumer profiles have not +been tested with it. + ## Apple Silicon `mlx_fasth3.py` stays on FastH3 V1 and its uniform AdaLN cache. diff --git a/examples/inference/basic/basic_fasth3.py b/examples/inference/basic/basic_fasth3.py index 91314f8ca..5f59b2691 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/examples/inference/basic/basic_fasth3.py @@ -37,6 +37,8 @@ from fastvideo.api import ( PipelineSelection, SamplingConfig, ) +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig +from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule DEFAULT_MODEL = "FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2" @@ -163,12 +165,55 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser: default=None, help='whole-DiT torch.compile mode, e.g. "reduce-overhead"; requires ' "--no-inference-torch-compile") + parser.add_argument("--token-reduction", + action="store_true", + help="training-free speed mode after antirez/h3.c: pair horizontally adjacent generated-video " + "tokens inside the middle DiT blocks; faster, but can change composition") + parser.add_argument("--token-reduction-blocks", + default=None, + metavar="BEGIN:END", + help="blocks that run reduced (default 4:30)") + parser.add_argument("--token-reduction-early", + default=None, + metavar="STEPS:END", + help="for the first STEPS forwards, reduce up to block END instead (default 4:40; 0:0 " + "disables)") + parser.add_argument("--token-reduction-scale", + type=float, + default=None, + help="scale of the update each paired token gets from its pair (default 1.0); unpaired " + "tokens, such as the last column of an odd-width video, take the full update as in h3.c") return parser +def token_reduction_options(args: argparse.Namespace) -> dict[str, object] | None: + """``h3_token_reduction`` options from the command line, or None when it is off.""" + if not getattr(args, "token_reduction", False): + return None + options: dict[str, object] = {"enabled": True} + if args.token_reduction_blocks is not None: + options["blocks"] = args.token_reduction_blocks + if args.token_reduction_early is not None: + options["early"] = args.token_reduction_early + if args.token_reduction_scale is not None: + options["update_scale"] = args.token_reduction_scale + return options + + def validate_args(parser: argparse.ArgumentParser, args: argparse.Namespace) -> argparse.Namespace: if args.repeats < 1: parser.error("--repeats must be at least 1") + tuning = (args.token_reduction_blocks, args.token_reduction_early, args.token_reduction_scale) + if not args.token_reduction and any(value is not None for value in tuning): + parser.error("--token-reduction-blocks/-early/-scale need --token-reduction") + if args.token_reduction: + # Reject malformed values and intervals past H3's 50 blocks before any download; the + # transformer loader checks a shallower checkpoint, such as FastH3 Trim, against its own depth. + try: + MiniMaxH3TokenReductionSchedule.from_options(token_reduction_options(args) or {}, + num_blocks=MiniMaxH3ArchConfig().num_layers) + except ValueError as error: + parser.error(str(error)) if args.num_gpus < 1: parser.error("--num-gpus must be at least 1") if not 0.0 <= args.vsa_sparsity < 1.0: @@ -303,6 +348,9 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig: "VSA_sparsity": args.vsa_sparsity, "VSA_tile_size": args.vsa_tile_size, }) + token_reduction = token_reduction_options(args) + if token_reduction is not None: + experimental["h3_token_reduction"] = token_reduction return GeneratorConfig( model_path=args.model_path, pipeline=PipelineSelection( diff --git a/examples/inference/basic/basic_fasth3_omniref_pdd.py b/examples/inference/basic/basic_fasth3_omniref_pdd.py index a3f7aed2c..37494df67 100644 --- a/examples/inference/basic/basic_fasth3_omniref_pdd.py +++ b/examples/inference/basic/basic_fasth3_omniref_pdd.py @@ -61,7 +61,9 @@ from fastvideo.api import ( QuantizationConfig, SamplingConfig, ) +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig from fastvideo.configs.pipelines.minimax_h3 import parse_base_model_revision +from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference CONTRACT = "fastvideo_inference.json" @@ -154,9 +156,41 @@ def build_parser() -> argparse.ArgumentParser: help="precomputed AdaLN modulation table, so the 26 GB of adaln_proj never loads " "(write one once with FASTVIDEO_H3_ADALN_DUMP). Strongly recommended: the 24/32 GB " "profiles need it to keep the DiT resident") + parser.add_argument("--token-reduction", + action="store_true", + help="training-free speed mode after antirez/h3.c: pair horizontally adjacent generated-video " + "tokens inside the middle DiT blocks; faster, but can change composition") + parser.add_argument("--token-reduction-blocks", + default=None, + metavar="BEGIN:END", + help="blocks that run reduced (default 4:30)") + parser.add_argument("--token-reduction-early", + default=None, + metavar="STEPS:END", + help="for the first STEPS forwards, reduce up to block END instead (default 4:40; 0:0 " + "disables)") + parser.add_argument("--token-reduction-scale", + type=float, + default=None, + help="scale of the update each paired token gets from its pair (default 1.0); unpaired " + "tokens, such as the last column of an odd-width video, take the full update as in h3.c") return parser +def token_reduction_options(args: argparse.Namespace) -> dict[str, Any] | None: + """``h3_token_reduction`` options from the command line, or None when it is off.""" + if not args.token_reduction: + return None + options: dict[str, Any] = {"enabled": True} + if args.token_reduction_blocks is not None: + options["blocks"] = args.token_reduction_blocks + if args.token_reduction_early is not None: + options["early"] = args.token_reduction_early + if args.token_reduction_scale is not None: + options["update_scale"] = args.token_reduction_scale + return options + + def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = build_parser() args = parser.parse_args(argv) @@ -164,6 +198,17 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.error("pass at least one ordered reference with --image, --video, or --audio") if not any(kind != "audio" for kind, _ in args.references): parser.error("Ref2VA needs at least one image or video reference") + tuning = (args.token_reduction_blocks, args.token_reduction_early, args.token_reduction_scale) + if not args.token_reduction and any(value is not None for value in tuning): + parser.error("--token-reduction-blocks/-early/-scale need --token-reduction") + if args.token_reduction: + # Reject malformed values and intervals past H3's 50 blocks before any download; the + # transformer loader checks the checkpoint's own depth. + try: + MiniMaxH3TokenReductionSchedule.from_options(token_reduction_options(args), + num_blocks=MiniMaxH3ArchConfig().num_layers) + except ValueError as error: + parser.error(str(error)) return args @@ -421,7 +466,8 @@ def apply_profile(profile: DeviceProfile, decoder: str | None, adaln_table: str def build_generator_config(model_dir: Path, num_gpus: int, lossless_accel: bool = False, - profile: DeviceProfile | None = None) -> GeneratorConfig: + profile: DeviceProfile | None = None, + token_reduction: dict[str, Any] | None = None) -> GeneratorConfig: quantization = None if profile is not None: # `h3_dit_vsa` matches the packed export the consumer weights ship; the @@ -444,6 +490,8 @@ def build_generator_config(model_dir: Path, else: offload = OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=True, pin_cpu_memory=False) experimental = {} + if token_reduction is not None: + experimental["h3_token_reduction"] = dict(token_reduction) return GeneratorConfig( model_path=str(model_dir), engine=EngineConfig( @@ -496,7 +544,7 @@ def main() -> None: if args.height is None: args.height = profile.height if profile is not None else DEFAULT_HEIGHT generator = VideoGenerator.from_config( - build_generator_config(model_dir, args.num_gpus, args.lossless_accel, profile)) + build_generator_config(model_dir, args.num_gpus, args.lossless_accel, profile, token_reduction_options(args))) try: result = generator.generate( GenerationRequest( diff --git a/fastvideo/configs/pipelines/minimax_h3_token_reduction.py b/fastvideo/configs/pipelines/minimax_h3_token_reduction.py new file mode 100644 index 000000000..85209dad3 --- /dev/null +++ b/fastvideo/configs/pipelines/minimax_h3_token_reduction.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Options for MiniMax-H3 training-free token reduction, parsed without importing model code. + +``FastVideoArgs`` validates ``h3_token_reduction`` here, so building args never +imports a model module. The row maps and the pooling arithmetic live in +``fastvideo/models/dits/minimax_h3_token_reduction.py``. +""" +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +TOKEN_REDUCTION_OPTION_KEYS = frozenset({"enabled", "blocks", "early", "update_scale"}) + + +@dataclass(frozen=True) +class MiniMaxH3TokenReductionSchedule: + """Which DiT blocks run on paired tokens at each denoising step. + + Blocks ``begin <= block < end`` run reduced. During the first + ``early_steps`` steps the interval extends to ``early_end``. + + h3.c reduces blocks 4..29, and 4..39 during its first 10 of about 19 + forwards, where coarse structure forms. The default keeps that share of an + eight-forward PDD schedule: blocks 4..39 for the first four forwards. + """ + + begin: int = 4 + end: int = 30 + early_steps: int = 4 + early_end: int = 40 + update_scale: float = 1.0 + + @classmethod + def from_options(cls, + options: Mapping[str, Any], + num_blocks: int | None = None) -> MiniMaxH3TokenReductionSchedule | None: + """Parse ``h3_token_reduction`` options; None unless ``enabled`` is true. + + ``num_blocks``, when known, bounds both block intervals by the DiT depth. + """ + if not isinstance(options, Mapping) or options.keys() - TOKEN_REDUCTION_OPTION_KEYS: + raise ValueError( + f"h3_token_reduction must be a mapping with keys from {sorted(TOKEN_REDUCTION_OPTION_KEYS)}") + enabled = options.get("enabled", False) + if not isinstance(enabled, bool): + raise ValueError("h3_token_reduction.enabled must be a boolean") + if not enabled: + return None + begin, end = _int_pair(options.get("blocks", (cls.begin, cls.end)), "blocks") + if not 0 <= begin < end: + raise ValueError(f"h3_token_reduction.blocks must satisfy 0 <= BEGIN < END, got {begin}:{end}") + # Same rule as h3.c: the deeper early interval applies only when the + # base interval restores before block 40. + default_early = (cls.early_steps, cls.early_end) if end < cls.early_end else (0, end) + early_steps, early_end = _int_pair(options.get("early", default_early), "early") + if early_steps < 0: + raise ValueError("h3_token_reduction.early STEPS must be non-negative") + if early_steps == 0: + early_end = end + elif early_end <= end: + raise ValueError(f"h3_token_reduction.early END must exceed blocks END {end}, got {early_end}") + if num_blocks is not None and max(end, early_end) > num_blocks: + raise ValueError(f"h3_token_reduction block intervals end at {max(end, early_end)}, past the " + f"{num_blocks} DiT blocks") + scale = options.get("update_scale", cls.update_scale) + if isinstance(scale, bool) or not isinstance(scale, int | float) or not 0.0 <= float(scale) <= 2.0: + raise ValueError("h3_token_reduction.update_scale must be a number in [0, 2]") + return cls(begin=begin, end=end, early_steps=early_steps, early_end=early_end, update_scale=float(scale)) + + def end_for_step(self, step: int) -> int: + return self.early_end if step < self.early_steps else self.end + + +def _int_pair(value: Any, name: str) -> tuple[int, int]: + if isinstance(value, str): + value = value.split(":") + if (not isinstance(value, list | tuple) or len(value) != 2 + or any(isinstance(item, bool) or not isinstance(item, int | str) for item in value)): + raise ValueError(f"h3_token_reduction.{name} must be two integers, e.g. [4, 30] or '4:30'") + try: + return int(value[0]), int(value[1]) + except ValueError as error: + raise ValueError(f"h3_token_reduction.{name} must be two integers, got {value!r}") from error + + +__all__ = ["MiniMaxH3TokenReductionSchedule", "TOKEN_REDUCTION_OPTION_KEYS"] diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 133470b06..607660f36 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -12,6 +12,7 @@ from typing import Any, TYPE_CHECKING from fastvideo.configs.configs import PreprocessConfig from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule from fastvideo.configs.utils import clean_cli_args from fastvideo.layers.quantization import QUANTIZATION_METHODS, QuantizationMethods from fastvideo.logger import init_logger @@ -296,6 +297,11 @@ class FastVideoArgs: inference_torch_compile: bool = False h3_vsa_fp4: dict[str, Any] = field(default_factory=dict) + # Training-free horizontal pairing of MiniMax-H3's generated-video tokens in + # the middle DiT blocks (antirez/h3.c --token-reduction). Keys: enabled, + # blocks [BEGIN, END], early [STEPS, END], update_scale. See + # fastvideo/configs/pipelines/minimax_h3_token_reduction.py. + h3_token_reduction: dict[str, Any] = field(default_factory=dict) disable_autocast: bool = False @@ -453,6 +459,8 @@ class FastVideoArgs: raise ValueError("h3_vsa_fp4.prefix_precision must be empty, fp8 or bf16") if self.h3_vsa_fp4.get("enabled") and (self.inference_torch_compile or self.enable_torch_compile): raise ValueError("h3_vsa_fp4 requires eager DiT inference; disable DiT torch.compile") + # The DiT depth is checked against the checkpoint's own config when the transformer loads. + MiniMaxH3TokenReductionSchedule.from_options(self.h3_token_reduction) self._fold_vae_parallel_env() self._fold_h3_encoder_split_env() # Runs after FASTVIDEO_ATTENTION_BACKEND is copied into attention_backend, diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 06ad52e68..df1a8e5d8 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -38,6 +38,8 @@ from fastvideo.logger import init_logger from fastvideo.models.dits.base import BaseDiT from fastvideo.models.dits.minimax_h3_vsa_fp4 import (STAGES, vsa_fp4_attention, vsa_fp4_attention_sp, vsa_fp4_requested, vsa_tile_first_attention) +from fastvideo.models.dits.minimax_h3_token_reduction import (MiniMaxH3TokenReduction, expand_token_pairs, + pool_token_pairs) from fastvideo.models.dits.minimax_h3_fusions import exact as exact_kernels from fastvideo.models.dits.minimax_h3_fusions import mk as mk_kernels from fastvideo.models.dits.minimax_h3_fusions import ( @@ -967,6 +969,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT): # per-generation caches for loop-invariant work (see _rotary_for / # _refined_text); plain attrs, never in state_dict self._rope_cache: tuple | None = None + self._reduced_rope_cache: tuple | None = None self._text_cache: tuple | None = None self.token_refiner = MiniMaxH3TokenRefiner( arch.hidden_size, @@ -1229,24 +1232,61 @@ class MiniMaxH3Transformer3DModel(BaseDiT): (2 * arch.rope_freq_dim))) self.rope._buffers["inv_freq"] = inv_freq - def _rotary_for(self, position_ids: torch.Tensor, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + def _rotary_for(self, + position_ids: torch.Tensor, + dtype: torch.dtype, + cache_name: str = "_rope_cache") -> tuple[torch.Tensor, torch.Tensor]: """cos/sin depend only on positions; the denoising stage holds one position_ids tensor for the whole loop, so cache per layout and pre-cast once instead of rebuilding every step and re-casting in every block. Holding the source tensor keeps the identity check safe - against id reuse.""" + against id reuse. Token reduction keeps its pair positions in a + second slot so the two layouts do not evict each other.""" if torch.compiler.is_compiling(): cos, sin = self.rope(position_ids) return cos.to(dtype), sin.to(dtype) - cached = self._rope_cache + cached = getattr(self, cache_name) if (cached is not None and cached[0] is position_ids and cached[1] == dtype and cached[2][0].device == position_ids.device): return cached[2] cos, sin = self.rope(position_ids) value = (cos.to(dtype), sin.to(dtype)) - self._rope_cache = (position_ids, dtype, value) + setattr(self, cache_name, (position_ids, dtype, value)) return value + def _enter_token_reduction(self, packed_hidden_states: torch.Tensor, plan: MiniMaxH3TokenReduction, + sequence_length: int, + sp_world_size: int) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Pool video pairs; returns the reduced block input and the full rows plus pooled baseline to restore from. + + Under sequence parallelism pooling changes how many rows each rank + holds, so the full sequence is gathered, pooled and sharded again. + """ + with nvtx_range("minimax_h3.token_reduction.pool"): + if sp_world_size > 1: + packed_hidden_states = sequence_model_parallel_all_gather_with_unpad(packed_hidden_states, + sequence_length, + dim=1) + pooled = pool_token_pairs(packed_hidden_states, plan) + reduced = pooled + if sp_world_size > 1: + reduced = self._block_input(sequence_model_parallel_shard(pooled, dim=1)[0]) + return reduced, (packed_hidden_states, pooled) + + def _leave_token_reduction(self, packed_hidden_states: torch.Tensor, saved: tuple[torch.Tensor, torch.Tensor], + plan: MiniMaxH3TokenReduction, sp_world_size: int) -> torch.Tensor: + """Restore the full sequence from the reduced stack's output.""" + original, pooled = saved + with nvtx_range("minimax_h3.token_reduction.expand"): + if sp_world_size > 1: + packed_hidden_states = sequence_model_parallel_all_gather_with_unpad(packed_hidden_states, + plan.reduced_length, + dim=1) + restored = expand_token_pairs(original, packed_hidden_states, pooled, plan) + if sp_world_size > 1: + restored = self._block_input(sequence_model_parallel_shard(restored, dim=1)[0]) + return restored + def attach_step_splice(self, late: nn.Module, from_step: int) -> None: """Hand denoising steps ``from_step`` onward to ``late`` (same architecture, other weights). @@ -1339,13 +1379,18 @@ class MiniMaxH3Transformer3DModel(BaseDiT): audio_indices: torch.Tensor, text_indices: torch.Tensor, text_attn_metadata: AttentionMetadata | None = None, + token_reduction: MiniMaxH3TokenReduction | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - """Predict video and audio velocities from one caller-defined packed layout.""" + """Predict video and audio velocities from one caller-defined packed layout. + + ``token_reduction`` (inference only) runs blocks ``[begin, end)`` on + horizontally paired generated-video rows; see minimax_h3_token_reduction. + """ late = self.__dict__.get("_splice_late") if late is not None and get_forward_context().current_timestep >= self._splice_from_step: return late(hidden_states, audio_hidden_states, encoder_hidden_states, timestep, timestep_indices, token_tags, position_ids, video_indices, audio_indices, text_indices, - text_attn_metadata=text_attn_metadata) + text_attn_metadata=text_attn_metadata, token_reduction=token_reduction) if position_ids.ndim != 2 or position_ids.shape[-1] != 3: raise ValueError(f"position_ids must have shape (seq_len, 3), got {tuple(position_ids.shape)}.") sequence_length = position_ids.shape[0] @@ -1357,6 +1402,15 @@ class MiniMaxH3Transformer3DModel(BaseDiT): raise ValueError("audio_hidden_states row count must match audio_indices.") if encoder_hidden_states.shape[1] != text_indices.numel(): raise ValueError("encoder_hidden_states row count must match text_indices.") + if token_reduction is not None: + if torch.is_grad_enabled(): + raise ValueError("MiniMax H3 token reduction is inference-only.") + if token_reduction.sequence_length != sequence_length: + raise ValueError(f"Token reduction was built for {token_reduction.sequence_length} rows, " + f"the packed sequence has {sequence_length}.") + if not 0 <= token_reduction.begin < token_reduction.end <= len(self.transformer_blocks): + raise ValueError(f"Token reduction blocks {token_reduction.begin}:{token_reduction.end} must lie in " + f"0:{len(self.transformer_blocks)}.") video_embeds, _ = self.proj_in(hidden_states.to(self.proj_in.weight.dtype)) audio_embeds, _ = self.audio_proj_in(audio_hidden_states.to(self.audio_proj_in.weight.dtype)) @@ -1371,6 +1425,8 @@ class MiniMaxH3Transformer3DModel(BaseDiT): pad_tokens = sequence_length - real_tokens if pad_tokens < 0: raise ValueError("Real H3 token counts exceed the packed sequence length") + if pad_tokens and token_reduction is not None: + raise ValueError("MiniMax H3 token reduction does not support capacity padding rows.") if pad_tokens: metadata = get_forward_context().attn_metadata packed_sequence = getattr(metadata, "packed_sequence", None) @@ -1391,6 +1447,12 @@ class MiniMaxH3Transformer3DModel(BaseDiT): adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags local_timestep_indices = timestep_indices original_seq_len = sequence_length + reduced_rotary_emb = reduced_adaln_indices = None + if token_reduction is not None: + # A pair takes its first source's modulation row and the mean of both positions. + reduced_rotary_emb = self._rotary_for(token_reduction.position_ids, text_embeds.dtype, + "_reduced_rope_cache") + reduced_adaln_indices = adaln_indices.index_select(0, token_reduction.first) if sp_world_size > 1: packed_hidden_states = self._block_input(sequence_model_parallel_shard(packed_hidden_states, dim=1)[0]) @@ -1399,6 +1461,10 @@ class MiniMaxH3Transformer3DModel(BaseDiT): adaln_indices, _ = sequence_model_parallel_shard(adaln_indices, dim=0) local_timestep_indices, _ = sequence_model_parallel_shard(local_timestep_indices, dim=0) rotary_emb = (rotary_cos, rotary_sin) + if reduced_rotary_emb is not None: + reduced_rotary_emb = (sequence_model_parallel_shard(reduced_rotary_emb[0], dim=0)[0], + sequence_model_parallel_shard(reduced_rotary_emb[1], dim=0)[0]) + reduced_adaln_indices, _ = sequence_model_parallel_shard(reduced_adaln_indices, dim=0) if getattr(self, "_adaln_host_cache", False): # One host read per forward keys every block's modulation cache by the timestep values. @@ -1407,20 +1473,53 @@ class MiniMaxH3Transformer3DModel(BaseDiT): block.adaln_proj._cache_key = key self._move_adaln_tables(temb.device) - # The eager driver owns profiling markers while each block's compiled - # forward owns the graph that the marker surrounds. - for block_index, block in enumerate(self.transformer_blocks): - if STAGES.enabled: - logger.info("H3_MEMORY_BLOCK %d allocated=%.3f GiB reserved=%.3f GiB", block_index, - torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30) - with nvtx_range(f"minimax_h3.transformer_block.{block_index}"), STAGES.span("block_total"): - packed_hidden_states = block( - packed_hidden_states, - temb, - adaln_indices, - rotary_emb, - original_seq_len, - ) + # Reduced blocks attend over the shorter sequence, so they read its own + # attention metadata from the forward context; the full one is put back + # on exit and on error. + forward_context = get_forward_context() if token_reduction is not None else None + full_attn_metadata = forward_context.attn_metadata if forward_context is not None else None + reduction_saved = None + try: + # The eager driver owns profiling markers while each block's compiled + # forward owns the graph that the marker surrounds. + for block_index, block in enumerate(self.transformer_blocks): + if token_reduction is not None and block_index == token_reduction.begin: + assert forward_context is not None + packed_hidden_states, reduction_saved = self._enter_token_reduction( + packed_hidden_states, token_reduction, original_seq_len, sp_world_size) + forward_context.attn_metadata = token_reduction.attn_metadata + if token_reduction is not None and block_index == token_reduction.end: + assert forward_context is not None and reduction_saved is not None + packed_hidden_states = self._leave_token_reduction(packed_hidden_states, reduction_saved, + token_reduction, sp_world_size) + forward_context.attn_metadata = full_attn_metadata + reduction_saved = None + if reduction_saved is None: + block_adaln_indices, block_rotary_emb, block_seq_len = adaln_indices, rotary_emb, original_seq_len + else: + assert token_reduction is not None and reduced_adaln_indices is not None + assert reduced_rotary_emb is not None + block_adaln_indices, block_rotary_emb = reduced_adaln_indices, reduced_rotary_emb + block_seq_len = token_reduction.reduced_length + if STAGES.enabled: + logger.info("H3_MEMORY_BLOCK %d allocated=%.3f GiB reserved=%.3f GiB", block_index, + torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30) + with nvtx_range(f"minimax_h3.transformer_block.{block_index}"), STAGES.span("block_total"): + packed_hidden_states = block( + packed_hidden_states, + temb, + block_adaln_indices, + block_rotary_emb, + block_seq_len, + ) + if reduction_saved is not None: + assert token_reduction is not None + packed_hidden_states = self._leave_token_reduction(packed_hidden_states, reduction_saved, + token_reduction, sp_world_size) + reduction_saved = None + finally: + if forward_context is not None: + forward_context.attn_metadata = full_attn_metadata packed_hidden_states = self.norm_out( packed_hidden_states, diff --git a/fastvideo/models/dits/minimax_h3_token_reduction.py b/fastvideo/models/dits/minimax_h3_token_reduction.py new file mode 100644 index 000000000..993a88435 --- /dev/null +++ b/fastvideo/models/dits/minimax_h3_token_reduction.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Training-free horizontal token pairing for MiniMax H3, after antirez/h3.c ``--token-reduction``. + +Inside an interval of transformer blocks, every pair of horizontally adjacent +generated-video tokens is averaged into one row, so the middle blocks see the +target video at half width. Text, audio, references and conditioning rows are +never paired and pass through exactly. Before the first full-width block each +original row is restored as its own saved value plus the update its pair +received, so within-pair detail survives:: + + row 1, 3, 5, 9 -> pairs 2, 7 -> middle blocks 2.5, 6 -> restored 1.5, 3.5, 4, 8 + +Pair positions are the mean of the two source positions; a pair takes the +modulation row of its first source (every target-video row shares one). +Unpaired rows, such as the last column of an odd-width video, take the +reduced output directly, as in h3.c, so ``update_scale`` applies to pairs only. + +The option parsing and the default schedule live in +``fastvideo/configs/pipelines/minimax_h3_token_reduction.py``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from fastvideo.attention.backends.abstract import AttentionMetadata +from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule + + +@dataclass(frozen=True) +class MiniMaxH3TokenReduction: + """Row maps for one packed sequence whose generated video is its trailing ``(t, h, w)`` token grid. + + ``first``/``second`` name the full rows averaged into each reduced row + (equal for unpaired rows); ``parents`` names the reduced row every full + row restores from. ``pair_reduced_rows`` and ``pair_rows`` list the + reduced and full rows that belong to a true pair. + """ + + begin: int + end: int + first: torch.Tensor + second: torch.Tensor + parents: torch.Tensor + pair_reduced_rows: torch.Tensor + pair_rows: torch.Tensor + position_ids: torch.Tensor + update_scale: float = 1.0 + # Attention metadata of the reduced sequence (VSA-H3 tiles of the half-width grid). + attn_metadata: AttentionMetadata | None = None + + @property + def sequence_length(self) -> int: + return int(self.parents.numel()) + + @property + def reduced_length(self) -> int: + return int(self.first.numel()) + + def to(self, device: torch.device) -> MiniMaxH3TokenReduction: + return MiniMaxH3TokenReduction( + begin=self.begin, + end=self.end, + first=self.first.to(device), + second=self.second.to(device), + parents=self.parents.to(device), + pair_reduced_rows=self.pair_reduced_rows.to(device), + pair_rows=self.pair_rows.to(device), + position_ids=self.position_ids.to(device), + update_scale=self.update_scale, + attn_metadata=self.attn_metadata, + ) + + +def reduced_token_width(width: int) -> int: + return (width + 1) // 2 + + +def build_token_reduction( + position_ids: torch.Tensor, + video_start: int, + grid: tuple[int, int, int], + schedule: MiniMaxH3TokenReductionSchedule, +) -> MiniMaxH3TokenReduction: + """Pair horizontally adjacent rows of the row-major ``(t, h, w)`` video grid ending the sequence.""" + t, h, w = (int(v) for v in grid) + sequence_length = int(position_ids.shape[0]) + if min(t, h, w) <= 0 or video_start < 0 or video_start + t * h * w != sequence_length: + raise ValueError(f"Token reduction needs the generated video grid {grid} to end the packed sequence " + f"({video_start} + {t * h * w} != {sequence_length}).") + reduced_w = reduced_token_width(w) + prefix = torch.arange(video_start, dtype=torch.long) + + reduced_local = torch.arange(t * h * reduced_w, dtype=torch.long) + pair_column = 2 * (reduced_local % reduced_w) + first_video = video_start + (reduced_local // reduced_w) * w + pair_column + has_pair = pair_column + 1 < w + second_video = first_video + has_pair.long() + + full_local = torch.arange(t * h * w, dtype=torch.long) + column = full_local % w + parent_video = video_start + (full_local // w) * reduced_w + column // 2 + in_pair = 2 * (column // 2) + 1 < w + + first = torch.cat([prefix, first_video]) + second = torch.cat([prefix, second_video]) + # (x + x) / 2 == x, so unpaired rows keep their exact positions. + reduced_positions = (position_ids.index_select(0, first) + position_ids.index_select(0, second)) * 0.5 + return MiniMaxH3TokenReduction( + begin=schedule.begin, + end=schedule.end, + first=first, + second=second, + parents=torch.cat([prefix, parent_video]), + pair_reduced_rows=video_start + torch.nonzero(has_pair).flatten(), + pair_rows=video_start + torch.nonzero(in_pair).flatten(), + position_ids=reduced_positions, + update_scale=schedule.update_scale, + ) + + +def pool_token_pairs(hidden_states: torch.Tensor, plan: MiniMaxH3TokenReduction) -> torch.Tensor: + """``[B, S, C]`` -> ``[B, R, C]``: unpaired rows copied, pairs averaged in float32.""" + pooled = hidden_states.index_select(1, plan.first) + rows = plan.pair_reduced_rows + if rows.numel(): + pair_mean = (hidden_states.index_select(1, plan.first.index_select(0, rows)).float() + + hidden_states.index_select(1, plan.second.index_select(0, rows)).float()) * 0.5 + pooled.index_copy_(1, rows, pair_mean.to(hidden_states.dtype)) + return pooled + + +def expand_token_pairs(original: torch.Tensor, reduced: torch.Tensor, pooled: torch.Tensor, + plan: MiniMaxH3TokenReduction) -> torch.Tensor: + """Restore ``[B, S, C]`` from the reduced stack's output ``[B, R, C]``. + + Unpaired rows take the reduced output directly. A paired row becomes its + value before pooling plus ``update_scale`` times the change of its pair, + ``reduced - pooled``. + """ + restored = reduced.index_select(1, plan.parents) + rows = plan.pair_rows + if rows.numel(): + parents = plan.parents.index_select(0, rows) + update = reduced.index_select(1, parents).float() - pooled.index_select(1, parents).float() + paired = original.index_select(1, rows).float() + plan.update_scale * update + restored.index_copy_(1, rows, paired.to(restored.dtype)) + return restored + + +__all__ = [ + "MiniMaxH3TokenReduction", + "build_token_reduction", + "expand_token_pairs", + "pool_token_pairs", + "reduced_token_width", +] diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 9c9ffb90f..06c6bc693 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -1241,6 +1241,14 @@ class TransformerLoader(ComponentLoader): if not isinstance(dit_config, MiniMaxH3Config): raise ValueError("h3_vsa_fp4 is only supported by MiniMax-H3") dit_config.vsa_fp4 = dict(fastvideo_args.h3_vsa_fp4) + token_reduction = getattr(fastvideo_args, "h3_token_reduction", None) + if token_reduction and token_reduction.get("enabled"): + from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config + from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule + if not isinstance(dit_config, MiniMaxH3Config): + raise ValueError("h3_token_reduction is only supported by MiniMax-H3") + # The checkpoint's own config fixes the DiT depth; reject intervals past it before any weight loads. + MiniMaxH3TokenReductionSchedule.from_options(token_reduction, num_blocks=dit_config.arch_config.num_layers) # Generator-only QAT for DMD distillation: the teacher (real_score) and # critic (fake_score) transformers load with this flag set and must stay 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 d8333997f..46b3d83d8 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_denoising.py @@ -4,17 +4,24 @@ from __future__ import annotations import contextlib +import dataclasses from typing import Any import torch from fastvideo.attention.selector import component_attention_backend, get_attn_backend +from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context from fastvideo.hooks.activation_trace import trace_step from fastvideo.layers.pdd import PDDModalitySchedule, PDDSamplingPlan, build_pdd_sampling_plan from fastvideo.logger import init_logger +from fastvideo.models.dits.minimax_h3_token_reduction import ( + MiniMaxH3TokenReduction, + build_token_reduction, + reduced_token_width, +) from fastvideo.profiler import nvtx_range, profiler_region from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_KEYFRAME_NOISE_AUG, @@ -122,6 +129,35 @@ def _h3_vsa_ref2va_segments( return tuple(segments) +def _h3_token_reduction_plan( + layout: MiniMaxH3PackedLayout, + patch_size: tuple[int, int, int], + schedule: MiniMaxH3TokenReductionSchedule, +) -> MiniMaxH3TokenReduction: + """Pair rows of the generated video, which every H3 packing places last in row-major ``(t, h, w)`` order.""" + if layout.pad_tokens: + raise ValueError("MiniMax-H3 token reduction does not support capacity padding rows.") + grid = (layout.num_video_latent_frames // patch_size[0], layout.latent_height // patch_size[1], + layout.latent_width // patch_size[2]) + num_target_rows = grid[0] * grid[1] * grid[2] + video_start = layout.sequence_length - num_target_rows + target_rows = layout.video_indices[layout.num_condition_video_rows:] + if target_rows.numel() != num_target_rows or not torch.equal(target_rows.cpu(), + torch.arange(video_start, layout.sequence_length)): + raise ValueError("MiniMax-H3 token reduction needs the generated video to end the packed sequence.") + return build_token_reduction(layout.position_ids, video_start, grid, schedule) + + +def _h3_token_reduced_segments( + segments: tuple[int | tuple[int, int, int], ...], + patch_size: tuple[int, int, int], +) -> tuple[int | tuple[int, int, int], ...]: + """VSA-H3 segments of the reduced sequence: the generated video (last region) at half token width.""" + frames, height, width = segments[-1] # type: ignore[misc] + reduced_width = reduced_token_width(width // patch_size[2]) * patch_size[2] + return (*segments[:-1], (frames, height, reduced_width)) + + class MiniMaxH3DenoisingStage(PipelineStage): """Build both schedules and denoise both modalities in one transformer call.""" @@ -309,6 +345,31 @@ class MiniMaxH3DenoisingStage(PipelineStage): logger.info("MiniMax-H3 VSA-H3: %d reference video region(s); %s; %d-token tiles.", num_reference_regions, effect, vsa_tile_size) + token_reduction_schedule = MiniMaxH3TokenReductionSchedule.from_options( + getattr(fastvideo_args, "h3_token_reduction", {})) + token_reduction = None + reduced_vsa_metadata_builder = None + reduced_vsa_packed_segments: tuple[int | tuple[int, int, int], ...] = () + if token_reduction_schedule is not None: + dit_patch_size = fastvideo_args.pipeline_config.dit_config.patch_size + token_reduction = _h3_token_reduction_plan(layout, dit_patch_size, token_reduction_schedule).to(device) + logger.info( + "MiniMax-H3 token reduction: %d -> %d rows in blocks %d:%d (%d:%d for the first %d step(s)), " + "update scale %g.", token_reduction.sequence_length, token_reduction.reduced_length, + token_reduction_schedule.begin, token_reduction_schedule.end, token_reduction_schedule.begin, + token_reduction_schedule.early_end, token_reduction_schedule.early_steps, + token_reduction_schedule.update_scale) + if (token_reduction_schedule.early_steps >= len(video_timesteps) + and token_reduction_schedule.early_end != token_reduction_schedule.end): + logger.warning( + "MiniMax-H3 token reduction: the early interval covers all %d forwards, so every forward " + "reduces blocks %d:%d.", len(video_timesteps), token_reduction_schedule.begin, + token_reduction_schedule.early_end) + if vsa_metadata_builder is not None: + # Its own builder, so each sequence length keeps its own tile buffers. + reduced_vsa_metadata_builder = _h3_vsa_metadata_builder(self.transformer, fastvideo_args) + reduced_vsa_packed_segments = _h3_token_reduced_segments(vsa_packed_segments, vsa_patch_size) + try: if full_cpu_offload: self.transformer.to(device) @@ -327,6 +388,11 @@ class MiniMaxH3DenoisingStage(PipelineStage): break unique_timesteps, timestep_indices = row_timestep_plan[index] attn_metadata = None + step_token_reduction = None + if token_reduction is not None: + assert token_reduction_schedule is not None + step_token_reduction = dataclasses.replace(token_reduction, + end=token_reduction_schedule.end_for_step(index)) if vsa_metadata_builder is not None: # Optional schedule: run the first N steps dense (sparsity 0 # selects every tile — parity-proven ≡ dense ≤2e-4); early @@ -343,6 +409,23 @@ class MiniMaxH3DenoisingStage(PipelineStage): tile_size=vsa_tile_size, ref_keep_rate=vsa_ref_keep_rate, ) + if step_token_reduction is not None: + assert reduced_vsa_metadata_builder is not None + reduced_attn_metadata = reduced_vsa_metadata_builder.build( + current_timestep=index, + patch_size=vsa_patch_size, + VSA_sparsity=vsa_sparsity, + packed_segments=reduced_vsa_packed_segments, + device=device, + exempt=vsa_exempt, + dense_layers=vsa_dense_layers, + tile_size=vsa_tile_size, + ref_keep_rate=vsa_ref_keep_rate, + ) + step_token_reduction = dataclasses.replace(step_token_reduction, + attn_metadata=reduced_attn_metadata) + # Only pass the argument when reducing, so transformers without it keep working. + reduction_kwargs = {} if step_token_reduction is None else {"token_reduction": step_token_reduction} # Under torch.compile(mode="reduce-overhead") each denoising # step must be marked, or cudagraph trees flag cross-step # reuse of pooled outputs as "accessing tensor output of @@ -367,6 +450,7 @@ class MiniMaxH3DenoisingStage(PipelineStage): video_indices=video_indices, audio_indices=audio_indices, text_indices=text_indices, + **reduction_kwargs, ) video_start = layout.num_condition_video_rows diff --git a/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py index 468e01de7..8ed3ea595 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py +++ b/fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py @@ -349,3 +349,30 @@ def test_row_split_is_an_explicit_spark_opt_in(): assert envs.environment_variables["FASTVIDEO_NVFP4_MM_ROW_SPLIT"].default == "off" assert example.DEVICE_PROFILES["spark"].env["FASTVIDEO_NVFP4_MM_ROW_SPLIT"] == "table" assert "FASTVIDEO_NVFP4_MM_ROW_SPLIT" not in example.DEVICE_PROFILES["rtx-pro-6000"].env + + +@pytest.mark.parametrize("flags", [ + ("--token-reduction", "--token-reduction-blocks", "4:60"), + ("--token-reduction", "--token-reduction-early", "4:30"), + ("--token-reduction", "--token-reduction-scale", "3"), + ("--token-reduction-blocks", "4:20"), +]) +def test_token_reduction_flags_are_checked_before_download(flags): + with pytest.raises(SystemExit): + _args(*flags) + + +def test_token_reduction_flags_reach_the_generator_config(tmp_path): + args = _args("--token-reduction", "--token-reduction-blocks", "4:20", "--token-reduction-early", "0:0") + config = example.build_generator_config(tmp_path, 4, token_reduction=example.token_reduction_options(args)) + assert config.pipeline.experimental == {"h3_token_reduction": {"enabled": True, "blocks": "4:20", "early": "0:0"}} + assert example.build_generator_config(tmp_path, 4).pipeline.experimental == {} + + +def test_invalid_token_reduction_options_fail_when_args_are_built(): + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + + with pytest.raises(ValueError, match="h3_token_reduction"): + FastVideoArgs(model_path="unused", pipeline_config=MiniMaxH3PipelineConfig(), + h3_token_reduction={"enabled": True, "blocks": "30:4"}) diff --git a/fastvideo/tests/inference/test_basic_fasth3_token_reduction.py b/fastvideo/tests/inference/test_basic_fasth3_token_reduction.py new file mode 100644 index 000000000..8738cd766 --- /dev/null +++ b/fastvideo/tests/inference/test_basic_fasth3_token_reduction.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks of the token reduction flags in examples/inference/basic/basic_fasth3.py (no weights, no GPU). + +basic_fasth3_8step.py, the FastH3 8-Step V2 example, reuses this parser and generator config. +""" +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[3] +EXAMPLES = REPO_ROOT / "examples" / "inference" / "basic" + + +def _load_example(name: str, path: Path): + # basic_fasth3_8step.py imports basic_fasth3 from its own directory. + sys.path.insert(0, str(EXAMPLES)) + try: + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + finally: + sys.path.remove(str(EXAMPLES)) + + +fasth3 = _load_example("basic_fasth3_token_reduction_contract", EXAMPLES / "basic_fasth3.py") +eight_step = _load_example("basic_fasth3_8step_token_reduction_contract", EXAMPLES / "basic_fasth3_8step.py") + + +def _args(*overrides: str): + return fasth3.parse_args(["--prompt", "a test prompt", *overrides]) + + +def test_token_reduction_is_opt_in_experimental(): + assert "h3_token_reduction" not in fasth3.build_generator_config(_args()).pipeline.experimental + config = fasth3.build_generator_config(_args("--token-reduction", "--token-reduction-blocks", "4:20")) + assert config.pipeline.experimental["h3_token_reduction"] == {"enabled": True, "blocks": "4:20"} + + +@pytest.mark.parametrize("flags", [ + ("--token-reduction", "--token-reduction-blocks", "4:60"), + ("--token-reduction", "--token-reduction-early", "4:30"), + ("--token-reduction", "--token-reduction-scale", "3"), + ("--token-reduction-scale", "0.5"), +]) +def test_token_reduction_flags_are_checked_before_download(flags): + with pytest.raises(SystemExit): + _args(*flags) + + +def test_eight_step_example_accepts_token_reduction(): + args = eight_step.parse_args(["--prompt", "p", "--token-reduction", "--token-reduction-early", "0:0"]) + config = fasth3.build_generator_config(args) + assert config.model_path == "FastVideo/FastVideo-FastH3-8-Step-V2" + assert config.pipeline.experimental["h3_token_reduction"] == {"enabled": True, "early": "0:0"} diff --git a/fastvideo/tests/transformers/test_minimax_h3_token_reduction.py b/fastvideo/tests/transformers/test_minimax_h3_token_reduction.py new file mode 100644 index 000000000..0ee60fdec --- /dev/null +++ b/fastvideo/tests/transformers/test_minimax_h3_token_reduction.py @@ -0,0 +1,344 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 training-free token reduction (antirez/h3.c ``--token-reduction``). + +CPU checks cover option parsing, the pair row maps, the pool/restore +arithmetic and the denoising-stage helpers. The GPU checks run a tiny real +transformer on the dense SDPA backend. A one-token-wide video has no pairs, so +the reduced path must be bit-identical to the plain forward; that covers the +reduced positions, modulation rows and the entry and exit plumbing. Separate +checks record which attention metadata every block reads, and that the +full-sequence metadata comes back after the forward and after an error. +""" + +from __future__ import annotations + +import dataclasses +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo import envs +from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3ArchConfig, MiniMaxH3Config +from fastvideo.configs.pipelines.minimax_h3_token_reduction import MiniMaxH3TokenReductionSchedule +from fastvideo.forward_context import get_forward_context, set_forward_context +from fastvideo.models.dits.minimax_h3_token_reduction import ( + build_token_reduction, + expand_token_pairs, + pool_token_pairs, +) +from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence, build_row_timesteps +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import ( + _h3_token_reduced_segments, + _h3_token_reduction_plan, +) + +_PATCH = (1, 2, 2) + + +def _schedule(**options) -> MiniMaxH3TokenReductionSchedule: + schedule = MiniMaxH3TokenReductionSchedule.from_options({"enabled": True, **options}) + assert schedule is not None + return schedule + + +# ---------------------------------------------------------------------------- +# Options +# ---------------------------------------------------------------------------- + + +def test_schedule_is_off_unless_enabled() -> None: + assert MiniMaxH3TokenReductionSchedule.from_options({}) is None + assert MiniMaxH3TokenReductionSchedule.from_options({"enabled": False, "blocks": [1, 2]}) is None + + +def test_schedule_defaults_follow_h3c_share_of_eight_forwards() -> None: + schedule = _schedule() + assert (schedule.begin, schedule.end, schedule.early_steps, schedule.early_end) == (4, 30, 4, 40) + assert [schedule.end_for_step(step) for step in range(8)] == [40, 40, 40, 40, 30, 30, 30, 30] + + +def test_schedule_skips_early_interval_when_base_reaches_block_40() -> None: + schedule = _schedule(blocks="4:45") + assert (schedule.early_steps, schedule.end_for_step(0)) == (0, 45) + + +def test_schedule_accepts_strings_and_disabling_early() -> None: + schedule = _schedule(blocks="2:20", early="0:0", update_scale=0.5) + assert (schedule.begin, schedule.end, schedule.end_for_step(0), schedule.update_scale) == (2, 20, 20, 0.5) + + +@pytest.mark.parametrize("options,match", [ + ({"enabled": "yes"}, "enabled must be a boolean"), + ({"enabled": True, "blocks": [30, 4]}, "BEGIN < END"), + ({"enabled": True, "blocks": "4"}, "two integers"), + ({"enabled": True, "early": [2, 20]}, "must exceed blocks END"), + ({"enabled": True, "update_scale": 3.0}, r"\[0, 2\]"), + ({"enabled": True, "update_scale": True}, r"\[0, 2\]"), + ({"enabled": True, "begin": 4}, "keys from"), +]) +def test_schedule_rejects_bad_options(options, match) -> None: + with pytest.raises(ValueError, match=match): + MiniMaxH3TokenReductionSchedule.from_options(options) + + +def test_schedule_bounds_both_intervals_by_the_dit_depth() -> None: + assert MiniMaxH3TokenReductionSchedule.from_options({"enabled": True, "blocks": "4:50"}, num_blocks=50) is not None + for options in ({"blocks": "4:51"}, {"early": "4:60"}): + with pytest.raises(ValueError, match="past the 50 DiT blocks"): + MiniMaxH3TokenReductionSchedule.from_options({"enabled": True, **options}, num_blocks=50) + + +# ---------------------------------------------------------------------------- +# Row maps and arithmetic +# ---------------------------------------------------------------------------- + + +def _positions(sequence_length: int) -> torch.Tensor: + return torch.arange(sequence_length * 3, dtype=torch.float64).reshape(sequence_length, 3) + + +def test_pairs_cover_rows_and_leave_odd_column_alone() -> None: + # Two prefix rows, then a 1x2x3 video: columns (0, 1) pair, column 2 is alone. + plan = build_token_reduction(_positions(8), 2, (1, 2, 3), _schedule()) + assert plan.first.tolist() == [0, 1, 2, 4, 5, 7] + assert plan.second.tolist() == [0, 1, 3, 4, 6, 7] + assert plan.parents.tolist() == [0, 1, 2, 2, 3, 4, 4, 5] + assert plan.pair_reduced_rows.tolist() == [2, 4] + assert plan.pair_rows.tolist() == [2, 3, 5, 6] + assert (plan.sequence_length, plan.reduced_length) == (8, 6) + + +def test_pair_positions_are_the_mean_and_unpaired_positions_are_exact() -> None: + positions = _positions(8) + plan = build_token_reduction(positions, 2, (1, 2, 3), _schedule()) + torch.testing.assert_close(plan.position_ids[2], (positions[2] + positions[3]) / 2, rtol=0, atol=0) + for reduced_row, full_row in ((0, 0), (1, 1), (3, 4), (5, 7)): + assert torch.equal(plan.position_ids[reduced_row], positions[full_row]) + + +def test_plan_requires_the_video_to_end_the_sequence() -> None: + with pytest.raises(ValueError, match="end the packed sequence"): + build_token_reduction(_positions(9), 2, (1, 2, 3), _schedule()) + + +def test_h3c_worked_example() -> None: + # One prefix row (10), then a 1x1x4 video row 1, 3, 5, 9. + plan = build_token_reduction(_positions(5), 1, (1, 1, 4), _schedule()) + hidden = torch.tensor([10.0, 1.0, 3.0, 5.0, 9.0]).view(1, 5, 1) + pooled = pool_token_pairs(hidden, plan) + assert pooled.flatten().tolist() == [10.0, 2.0, 7.0] + # The middle blocks move the prefix to 11 and the pairs to 2.5 and 6. + reduced = torch.tensor([11.0, 2.5, 6.0]).view(1, 3, 1) + restored = expand_token_pairs(hidden, reduced, pooled, plan) + assert restored.flatten().tolist() == [11.0, 1.5, 3.5, 4.0, 8.0] + + +def test_update_scale_scales_only_the_pair_update() -> None: + plan = build_token_reduction(_positions(5), 1, (1, 1, 4), _schedule(update_scale=0.5)) + hidden = torch.tensor([10.0, 1.0, 3.0, 5.0, 9.0]).view(1, 5, 1) + pooled = pool_token_pairs(hidden, plan) + reduced = torch.tensor([11.0, 2.5, 6.0]).view(1, 3, 1) + assert expand_token_pairs(hidden, reduced, pooled, plan).flatten().tolist() == [11.0, 1.25, 3.25, 4.5, 8.5] + + +def test_unchanged_reduced_stack_restores_bf16_input_exactly() -> None: + plan = build_token_reduction(_positions(2 + 2 * 3 * 5), 2, (2, 3, 5), _schedule()) + hidden = torch.randn(1, plan.sequence_length, 8, generator=torch.Generator().manual_seed(0)).bfloat16() + pooled = pool_token_pairs(hidden, plan) + assert torch.equal(expand_token_pairs(hidden, pooled, pooled, plan), hidden) + + +def test_one_token_wide_video_is_an_identity() -> None: + plan = build_token_reduction(_positions(2 + 3), 2, (3, 1, 1), _schedule()) + assert plan.reduced_length == plan.sequence_length + assert plan.pair_rows.numel() == 0 + hidden = torch.randn(1, plan.sequence_length, 4).bfloat16() + assert torch.equal(pool_token_pairs(hidden, plan), hidden) + + +# ---------------------------------------------------------------------------- +# Denoising-stage helpers on a real packed layout +# ---------------------------------------------------------------------------- + + +def _layout(latent_width: int, keyframes: tuple[str, ...] = ()): + return build_packed_sequence( + text_token_tags=torch.ones(4, dtype=torch.long), + num_latent_frames=2, + latent_height=4, + latent_width=latent_width, + num_audio_latents=3, + patch_size=_PATCH, + keyframe_anchors=keyframes, + ) + + +@pytest.mark.parametrize("keyframes", [(), ("first", )]) +def test_stage_plan_pairs_only_the_generated_video(keyframes) -> None: + layout = _layout(latent_width=6, keyframes=keyframes) + plan = _h3_token_reduction_plan(layout, _PATCH, _schedule()) + video_start = layout.sequence_length - 2 * 2 * 3 + assert plan.sequence_length == layout.sequence_length + assert plan.reduced_length == video_start + 2 * 2 * 2 + assert plan.first[:video_start].tolist() == list(range(video_start)) + assert int(plan.pair_rows.min()) >= video_start + + +def test_reduced_segments_halve_the_last_region_width() -> None: + assert _h3_token_reduced_segments((4, 0, 6, (2, 4, 6)), _PATCH) == (4, 0, 6, (2, 4, 4)) + assert _h3_token_reduced_segments((4, (3, 4, 4), 6, (2, 4, 8)), _PATCH) == (4, (3, 4, 4), 6, (2, 4, 4)) + + +# ---------------------------------------------------------------------------- +# Tiny real transformer (GPU) +# ---------------------------------------------------------------------------- + + +def _tiny_transformer(device: torch.device): + from fastvideo.models.dits.minimax_h3 import MiniMaxH3Transformer3DModel + + arch = MiniMaxH3ArchConfig( + num_attention_heads=2, + attention_head_dim=16, + hidden_size=32, + num_layers=3, + num_refiner_layers=1, + ffn_dim=64, + freq_dim=16, + time_embed_hidden_dim=32, + time_embed_dim=16, + rope_freq_dim=2, + ) + transformer = MiniMaxH3Transformer3DModel(config=MiniMaxH3Config(arch_config=arch), hf_config={}) + generator = torch.Generator(device="cpu").manual_seed(0) + with torch.no_grad(): + for name, parameter in transformer.named_parameters(): + if parameter.ndim < 2 and "norm" in name: + parameter.fill_(1.0) + else: + parameter.copy_(torch.randn(parameter.shape, generator=generator) * 0.05) + return transformer.to(device).eval() + + +def _call_kwargs(layout, inputs) -> dict[str, torch.Tensor]: + device = inputs["hidden_states"].device + unique, inverse = build_row_timesteps(layout, 0.5, 0.5, 0.999, 1.0) + return { + "timestep": unique.to(device), + "timestep_indices": inverse.to(device), + "token_tags": layout.token_tags.to(device), + "position_ids": layout.position_ids.to(device), + "video_indices": layout.video_indices.to(device), + "audio_indices": layout.audio_indices.to(device), + "text_indices": layout.text_indices.to(device), + **inputs, + } + + +def _forward(transformer, layout, inputs, token_reduction=None, attn_metadata=None): + with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=attn_metadata): + outputs = transformer(token_reduction=token_reduction, **_call_kwargs(layout, inputs)) + # The full-sequence metadata is back in place after the forward. + assert get_forward_context().attn_metadata is attn_metadata + return outputs + + +def _inputs(transformer, layout, device: torch.device) -> dict[str, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(1) + + def rows(count: int, width: int) -> torch.Tensor: + return torch.randn(1, count, width, generator=generator).to(device) + + return { + "hidden_states": rows(layout.video_indices.numel(), transformer.proj_in.weight.shape[1]), + "audio_hidden_states": rows(layout.audio_indices.numel(), transformer.audio_proj_in.weight.shape[1]), + "encoder_hidden_states": rows(layout.text_indices.numel(), transformer.context_embedder.weight.shape[1]), + } + + +@pytest.mark.skipif(not torch.cuda.is_available(), + reason="MiniMax H3 attention layers require a CUDA device for a real forward") +def test_one_token_wide_reduced_forward_is_bit_identical(env_overrides, distributed_setup) -> None: + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + device = torch.device("cuda") + transformer = _tiny_transformer(device) + layout = _layout(latent_width=2) + inputs = _inputs(transformer, layout, device) + plan = _h3_token_reduction_plan(layout, _PATCH, _schedule(blocks="1:2", early="0:0")).to(device) + plain = _forward(transformer, layout, inputs) + reduced = _forward(transformer, layout, inputs, plan) + for expected, actual in zip(plain, reduced, strict=True): + assert torch.equal(expected, actual) + + +@pytest.mark.skipif(not torch.cuda.is_available(), + reason="MiniMax H3 attention layers require a CUDA device for a real forward") +@pytest.mark.parametrize("blocks", ["1:2", "0:3"]) +def test_reduced_forward_keeps_shapes_and_changes_video(env_overrides, distributed_setup, blocks) -> None: + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + device = torch.device("cuda") + transformer = _tiny_transformer(device) + layout = _layout(latent_width=8) + inputs = _inputs(transformer, layout, device) + plan = _h3_token_reduction_plan(layout, _PATCH, _schedule(blocks=blocks, early="0:0")).to(device) + plain_video, plain_audio = _forward(transformer, layout, inputs) + video, audio = _forward(transformer, layout, inputs, plan) + assert video.shape == plain_video.shape and audio.shape == plain_audio.shape + assert torch.isfinite(video).all() and torch.isfinite(audio).all() + assert not torch.equal(video, plain_video) + + +@pytest.mark.skipif(not torch.cuda.is_available(), + reason="MiniMax H3 attention layers require a CUDA device for a real forward") +@pytest.mark.parametrize("blocks", ["1:2", "1:3"]) +def test_reduced_blocks_read_their_own_attention_metadata(env_overrides, distributed_setup, blocks) -> None: + # "1:3" ends at the last block, so the exit and restore happen after the block loop. + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + device = torch.device("cuda") + transformer = _tiny_transformer(device) + layout = _layout(latent_width=8) + inputs = _inputs(transformer, layout, device) + full_metadata, reduced_metadata = SimpleNamespace(name="full"), SimpleNamespace(name="reduced") + plan = _h3_token_reduction_plan(layout, _PATCH, _schedule(blocks=blocks, early="0:0")).to(device) + plan = dataclasses.replace(plan, attn_metadata=reduced_metadata) + seen = [] + hooks = [ + block.register_forward_pre_hook(lambda module, args: seen.append(get_forward_context().attn_metadata)) + for block in transformer.transformer_blocks + ] + try: + _forward(transformer, layout, inputs, plan, attn_metadata=full_metadata) + finally: + for hook in hooks: + hook.remove() + expected = [ + reduced_metadata if plan.begin <= index < plan.end else full_metadata + for index in range(len(transformer.transformer_blocks)) + ] + assert [metadata.name for metadata in seen] == [metadata.name for metadata in expected] + + +@pytest.mark.skipif(not torch.cuda.is_available(), + reason="MiniMax H3 attention layers require a CUDA device for a real forward") +def test_attention_metadata_is_restored_after_an_error_in_a_reduced_block(env_overrides, distributed_setup) -> None: + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + device = torch.device("cuda") + transformer = _tiny_transformer(device) + layout = _layout(latent_width=8) + inputs = _inputs(transformer, layout, device) + full_metadata = SimpleNamespace(name="full") + plan = _h3_token_reduction_plan(layout, _PATCH, _schedule(blocks="1:2", early="0:0")).to(device) + plan = dataclasses.replace(plan, attn_metadata=SimpleNamespace(name="reduced")) + + def fail(module, args): + raise RuntimeError("block failed") + + hook = transformer.transformer_blocks[1].register_forward_pre_hook(fail) + try: + with torch.no_grad(), set_forward_context(current_timestep=0, attn_metadata=full_metadata): + with pytest.raises(RuntimeError, match="block failed"): + transformer(token_reduction=plan, **_call_kwargs(layout, inputs)) + assert get_forward_context().attn_metadata is full_metadata + finally: + hook.remove()