Davide Locatelli
e1e2559359
These speed-ups for the MiniMax-H3 pipelines keep the output bit-identical. Each one has a switch that is off by default: - FASTVIDEO_MINIMAX_H3_EXACT_KERNELS: Triton RoPE, AdaLN modulation, gated residual and SwiGLU kernels that round to bf16 at the same points as the eager ops. - FASTVIDEO_H3_VSA_HEADS_FIRST_TILE: VSA-H3 tiles are written directly in the heads-first layout of the kernel. - FASTVIDEO_H3_VAE_TILE_PARALLEL: VAE spatial tiles are split across the sequence-parallel ranks. - FASTVIDEO_H3_REF2VA_MEMO_ENTRIES: repeat reference encodes are reused. The example basic_fasth3_omniref_pdd.py --lossless-accel turns them all on. On GB200 with 4 GPUs, MiniMax-H3 text to video with audio (124 frames, 1344x768) gives bit-identical frames and audio with the first three switches on and off.
313 lines
15 KiB
Python
313 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Eight-forward Ref2VA video+audio generation with a FastH3 OmniRef PDD student.
|
|
|
|
A FastH3 OmniRef PDD export is a Parallel Decoding Distillation (PDD) student of
|
|
MiniMax-H3's reference-conditioned ``transformer_ref`` partition. Its directory
|
|
(local, or a Hugging Face repo) carries only what distillation changed:
|
|
``transformer_ref/`` (output heads widened to the fine grid, trained VSA
|
|
compression gates), the two scheduler configs, and ``fastvideo_inference.json``.
|
|
The text encoder, tokenizer, processor, and both VAEs are base MiniMax-H3's.
|
|
|
|
This script composes the two into one local model directory of symlinks and
|
|
runs the recipe the contract records: fused blocks of the fine grid (one
|
|
transformer forward each), the video/audio shifts, and VIDEO_SPARSE_ATTN_H3
|
|
with its sparsity and tile size, where every reference video is its own
|
|
sparse region. The contract's 128-token tiles run only on the sm_100a/sm_103a
|
|
CUDA kernel (B200/B300/GB200/GB300) of a fastvideo-kernel build with the
|
|
Blackwell VSA extension; the VSA-H3 backend raises at the first attention call
|
|
when either is missing.
|
|
With ``--num-gpus`` above 1 the DiT is sharded across the GPUs (FSDP) and the
|
|
sequence is split across them (sequence parallelism).
|
|
|
|
References are ordered. Pass them in order with --image / --video / --audio,
|
|
for example ``--video dance.mp4 --image outfit.png --audio voice.wav``.
|
|
|
|
``--lossless-accel`` turns on every bit-identical speed-up the pipeline has
|
|
(LOSSLESS_ACCEL_ENV below, plus resident encoders, pinned host memory and the
|
|
sequence-parallel VAE); the output is bitwise the same as without it. It keeps
|
|
the 63 GB Qwen3-VL text encoder resident on every GPU, so it needs about
|
|
125 GB per GPU at four GPUs (B200/B300 class).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from collections.abc import Sequence
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from fastvideo import VideoGenerator
|
|
from fastvideo.api import (
|
|
ComponentConfig,
|
|
EngineConfig,
|
|
GenerationRequest,
|
|
GeneratorConfig,
|
|
InputConfig,
|
|
OffloadConfig,
|
|
OutputConfig,
|
|
ParallelismConfig,
|
|
PipelineSelection,
|
|
SamplingConfig,
|
|
)
|
|
from fastvideo.configs.pipelines.minimax_h3 import parse_base_model_revision
|
|
from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference
|
|
|
|
CONTRACT = "fastvideo_inference.json"
|
|
# Components the distilled export replaces, and the ones it shares with base MiniMax-H3.
|
|
EXPORT_COMPONENTS = ("transformer_ref", "scheduler", "audio_scheduler")
|
|
BASE_COMPONENTS = ("text_encoder", "tokenizer", "processor", "vae", "audio_vae")
|
|
MANIFESTS = ("modular_model_index.json", "model_index.json")
|
|
# Components the composed directory's manifest must declare, besides its transformer_ref.
|
|
MANIFEST_COMPONENTS = ("scheduler", "audio_scheduler", *BASE_COMPONENTS)
|
|
# Bit-identical speed-ups that FastVideo reads from the environment. Workers
|
|
# inherit it, so these must be set before the generator starts.
|
|
LOSSLESS_ACCEL_ENV = {
|
|
# Fused NVLink Ulysses all-to-all instead of NCCL plus permute copies.
|
|
"FASTVIDEO_ULYSSES_A2A": "auto",
|
|
# RoPE, AdaLN modulation / gated residual and SwiGLU kernels that round to
|
|
# bf16 exactly where the eager ops do (not the FP32 round-once fusions).
|
|
"FASTVIDEO_MINIMAX_H3_EXACT_KERNELS": "all",
|
|
# VSA tiles scattered straight into the kernels' heads-first layout.
|
|
"FASTVIDEO_H3_VSA_HEADS_FIRST_TILE": "1",
|
|
# VAE spatial tiles split across the sequence-parallel ranks.
|
|
"FASTVIDEO_H3_VAE_TILE_PARALLEL": "1",
|
|
# Content-keyed reuse of repeat reference encodes.
|
|
"FASTVIDEO_H3_REF2VA_MEMO_ENTRIES": "16",
|
|
}
|
|
|
|
|
|
class _AppendReference(argparse.Action):
|
|
"""Collect --image/--video/--audio into one list that keeps command-line order."""
|
|
|
|
def __call__(self, parser, namespace, value, option_string=None):
|
|
references = list(getattr(namespace, self.dest, None) or [])
|
|
references.append((self.const, value))
|
|
setattr(namespace, self.dest, references)
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument("--model-path",
|
|
required=True,
|
|
help="FastH3 OmniRef PDD export: a local directory or a Hugging Face repo id")
|
|
parser.add_argument("--revision", default=None, help="revision of --model-path when it is a Hugging Face repo")
|
|
parser.add_argument("--base-model-path",
|
|
default=None,
|
|
help="base MiniMax-H3 snapshot (local directory or repo id). Default: the repo and revision "
|
|
"pinned by the export's base_model_revision")
|
|
parser.add_argument("--base-revision",
|
|
default=None,
|
|
help="revision of --base-model-path when it is a repo id (the export's pinned revision "
|
|
"applies only to its own base repo); without --base-model-path, overrides that pin")
|
|
parser.add_argument("--composed-dir",
|
|
default=None,
|
|
help="where to write the composed model directory of symlinks (default: "
|
|
"<output>/fasth3_omniref_model)")
|
|
for flag in ("image", "video", "audio"):
|
|
parser.add_argument(f"--{flag}",
|
|
dest="references",
|
|
action=_AppendReference,
|
|
const=flag,
|
|
metavar="PATH",
|
|
help=f"an ordered {flag} reference (repeatable)")
|
|
parser.add_argument("--prompt", required=True)
|
|
parser.add_argument("--output", default="outputs/fasth3_omniref_pdd")
|
|
parser.add_argument("--height", type=int, default=480)
|
|
parser.add_argument("--width", type=int, default=832)
|
|
parser.add_argument("--num-frames", type=int, default=124)
|
|
parser.add_argument("--seed", type=int, default=0)
|
|
parser.add_argument("--num-gpus", type=int, default=1, help="sequence-parallel GPUs (must divide 56 heads)")
|
|
parser.add_argument("--lossless-accel",
|
|
action="store_true",
|
|
help="enable every bit-identical speed-up (output unchanged; about 125 GB per GPU)")
|
|
return parser
|
|
|
|
|
|
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
if not args.references:
|
|
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")
|
|
return args
|
|
|
|
|
|
def _snapshot(path_or_repo: str, revision: str | None, allow_patterns: list[str]) -> Path:
|
|
local = Path(path_or_repo).expanduser()
|
|
if local.is_dir():
|
|
return local
|
|
from huggingface_hub import snapshot_download
|
|
|
|
return Path(snapshot_download(repo_id=path_or_repo, revision=revision, allow_patterns=allow_patterns))
|
|
|
|
|
|
def load_contract(export_dir: Path) -> dict[str, Any]:
|
|
path = export_dir / CONTRACT
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"{export_dir} has no {CONTRACT}; it is not a FastH3 distilled export.")
|
|
contract = json.loads(path.read_text(encoding="utf-8"))
|
|
if "pdd_steps" not in contract or contract.get("model_type") != "ref2va":
|
|
raise ValueError(f"{path} is not a Ref2VA PDD contract; use the example matching the checkpoint.")
|
|
return contract
|
|
|
|
|
|
def base_model_source(args: argparse.Namespace, contract: dict[str, Any]) -> tuple[str, str | None]:
|
|
"""The base snapshot and revision: --base-model-path as given, else the export's pin."""
|
|
base_repo, base_revision = base_model_from_contract(contract)
|
|
if args.base_model_path:
|
|
return args.base_model_path, args.base_revision
|
|
return base_repo, args.base_revision or base_revision
|
|
|
|
|
|
def base_model_from_contract(contract: dict[str, Any]) -> tuple[str, str]:
|
|
"""The base repo id and revision that the export pins as ``hf://<repo id>@<revision>``."""
|
|
return parse_base_model_revision(contract.get("base_model_revision"))
|
|
|
|
|
|
def _link(destination: Path, source: Path) -> None:
|
|
if not source.exists():
|
|
raise FileNotFoundError(f"Missing checkpoint component: {source}")
|
|
source = source.resolve()
|
|
if destination.is_symlink():
|
|
if destination.resolve() == source:
|
|
return
|
|
destination.unlink()
|
|
elif destination.exists():
|
|
raise FileExistsError(f"{destination} exists and is not a symlink; choose another --composed-dir.")
|
|
destination.symlink_to(source, target_is_directory=source.is_dir())
|
|
|
|
|
|
def _declares_components(manifest: Path) -> bool:
|
|
"""Whether a Diffusers manifest declares every component of the composed directory."""
|
|
declared = {
|
|
name
|
|
for name, spec in json.loads(manifest.read_text(encoding="utf-8")).items()
|
|
if isinstance(spec, list) and spec and spec[0] is not None
|
|
}
|
|
return set(MANIFEST_COMPONENTS) <= declared and bool({"transformer", "transformer_ref"} & declared)
|
|
|
|
|
|
def select_manifest(export_dir: Path, base_dir: Path) -> Path:
|
|
"""The export's manifest when it declares every composed component, else the base's."""
|
|
for root in (export_dir, base_dir):
|
|
for name in MANIFESTS:
|
|
if (root / name).is_file() and _declares_components(root / name):
|
|
return root / name
|
|
raise FileNotFoundError("Neither the export nor the base snapshot has a Diffusers model manifest that declares "
|
|
f"{', '.join(MANIFEST_COMPONENTS)} and transformer_ref.")
|
|
|
|
|
|
def compose_model_dir(export_dir: Path, base_dir: Path, composed_dir: Path) -> Path:
|
|
"""One MiniMax-H3 model directory: distilled components from the export, the rest from the base."""
|
|
composed_dir.mkdir(parents=True, exist_ok=True)
|
|
for name in (*EXPORT_COMPONENTS, CONTRACT):
|
|
_link(composed_dir / name, export_dir / name)
|
|
manifest_source = select_manifest(export_dir, base_dir)
|
|
for name in MANIFESTS:
|
|
stale = composed_dir / name
|
|
if name != manifest_source.name and stale.is_symlink():
|
|
stale.unlink()
|
|
_link(composed_dir / manifest_source.name, manifest_source)
|
|
for name in BASE_COMPONENTS:
|
|
_link(composed_dir / name, base_dir / name)
|
|
return composed_dir
|
|
|
|
|
|
def resolve_model(args: argparse.Namespace) -> tuple[Path, dict[str, Any]]:
|
|
export_dir = _snapshot(args.model_path, args.revision,
|
|
[CONTRACT, *MANIFESTS, *(f"{name}/**" for name in EXPORT_COMPONENTS)])
|
|
contract = load_contract(export_dir)
|
|
base_source, base_revision = base_model_source(args, contract)
|
|
base_dir = _snapshot(base_source, base_revision, [*MANIFESTS, *(f"{name}/**" for name in BASE_COMPONENTS)])
|
|
composed_dir = Path(args.composed_dir) if args.composed_dir else Path(args.output) / "fasth3_omniref_model"
|
|
return compose_model_dir(export_dir, base_dir, composed_dir), contract
|
|
|
|
|
|
def apply_lossless_accel_env() -> None:
|
|
"""Set LOSSLESS_ACCEL_ENV, keeping any value already in the environment."""
|
|
for name, value in LOSSLESS_ACCEL_ENV.items():
|
|
os.environ.setdefault(name, value)
|
|
|
|
|
|
def build_generator_config(model_dir: Path, num_gpus: int, lossless_accel: bool = False) -> GeneratorConfig:
|
|
if lossless_accel:
|
|
# Resident encoders skip streaming the 63 GB text encoder (and the
|
|
# VAEs) from host memory on every clip; pinned buffers copy the
|
|
# decoded clip at full PCIe speed; the parallel VAE paths reassemble
|
|
# the serial result exactly.
|
|
offload = OffloadConfig(dit=False, dit_layerwise=False, text_encoder=False, vae=False, pin_cpu_memory=True)
|
|
experimental = {"vae_parallel_decode": num_gpus > 1, "vae_parallel_encode": num_gpus > 1}
|
|
else:
|
|
offload = OffloadConfig(dit=False, dit_layerwise=False, text_encoder=True, vae=True, pin_cpu_memory=False)
|
|
experimental = {}
|
|
return GeneratorConfig(
|
|
model_path=str(model_dir),
|
|
engine=EngineConfig(
|
|
num_gpus=num_gpus,
|
|
# Shard the DiT across the GPUs rather than holding a full copy on each.
|
|
use_fsdp_inference=num_gpus > 1,
|
|
parallelism=ParallelismConfig(tp_size=1, sp_size=num_gpus),
|
|
offload=offload,
|
|
),
|
|
pipeline=PipelineSelection(
|
|
workload_type="i2v",
|
|
# FastVideo reads the trained fused-block partition and attention
|
|
# settings from the composed directory's fastvideo_inference.json.
|
|
components=ComponentConfig(override_pipeline_cls_name="MiniMaxH3Ref2VAModularPipeline"),
|
|
experimental=experimental,
|
|
),
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
output_dir = Path(args.output)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
model_dir, contract = resolve_model(args)
|
|
print(f"Composed model directory: {model_dir}")
|
|
print(f"Contract: {contract['num_inference_steps']} fused blocks of a {contract['pdd_steps']}-interval grid, "
|
|
f"{contract['attention_backend']} with sparsity {contract['vsa_sparsity']}, "
|
|
f"{contract['vsa_tile_size']}-token tiles, reference keep rate {contract['vsa_ref_keep_rate']}")
|
|
references = [MiniMaxH3Reference(source=path, media_type=kind) for kind, path in args.references]
|
|
|
|
if args.lossless_accel:
|
|
apply_lossless_accel_env()
|
|
generator = VideoGenerator.from_config(build_generator_config(model_dir, args.num_gpus, args.lossless_accel))
|
|
try:
|
|
result = generator.generate(
|
|
GenerationRequest(
|
|
prompt=args.prompt,
|
|
negative_prompt="",
|
|
inputs=InputConfig(references=references),
|
|
sampling=SamplingConfig(
|
|
height=args.height,
|
|
width=args.width,
|
|
num_frames=args.num_frames,
|
|
fps=24,
|
|
# PDD counts fused blocks: one transformer forward each.
|
|
num_inference_steps=contract["num_inference_steps"],
|
|
guidance_scale=1.0,
|
|
batch_cfg=False,
|
|
seed=args.seed,
|
|
),
|
|
output=OutputConfig(
|
|
output_path=str(output_dir / "fasth3_omniref_pdd.mp4"),
|
|
save_video=True,
|
|
return_frames=False,
|
|
),
|
|
))
|
|
print(f"Output written to: {result.video_path}")
|
|
if result.generation_time is not None:
|
|
print(f"Generation time: {result.generation_time:.1f} s")
|
|
if result.peak_memory_mb is not None:
|
|
print(f"Peak GPU memory allocated (rank 0): {result.peak_memory_mb / 1024:.1f} GiB")
|
|
finally:
|
|
generator.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|