From 74b409d7cfff92c95cb6b77c836f18f482b5be34 Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Wed, 19 Aug 2026 02:45:23 -0700 Subject: [PATCH] [test]: add H3 VAE parity and memory benchmark --- tests/local_tests/minimax_h3/README.md | 25 ++ tests/local_tests/minimax_h3/_reference.py | 23 ++ .../benchmark_minimax_h3_video_vae_memory.py | 300 ++++++++++++++++++ .../vaes/test_minimax_h3_video_vae_parity.py | 32 +- 4 files changed, 374 insertions(+), 6 deletions(-) create mode 100644 tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py diff --git a/tests/local_tests/minimax_h3/README.md b/tests/local_tests/minimax_h3/README.md index 192b87c90..82f26c85f 100644 --- a/tests/local_tests/minimax_h3/README.md +++ b/tests/local_tests/minimax_h3/README.md @@ -59,5 +59,30 @@ pytest \ With a gate enabled, missing CUDA, source, or weights is a failure. Recorded component evidence is exact for both DiT partitions, the video VAE, and all Qwen3-VL hidden states; audio decode has maximum absolute drift `2.4e-7`. +The video VAE test verifies the reference checkout at commit +`abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc` and compares the production CPU `uint8` `encode_pixels()` path against +the official posterior element by element. + +## Video VAE memory benchmark + +The benchmark uses one warmup and three measured runs with `vae_cpu_offload=True`. It reports absolute and +stage-incremental allocated/reserved CUDA peaks for every rank. For SP runs, the reported aggregate is explicitly the +sum of rank-local maxima, not a simultaneous node peak. + +```bash +python tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py \ + --source-root "$PWD" --model-root "$MINIMAX_H3_MODEL_ROOT" \ + --revision-label candidate --operation encode + +python -m torch.distributed.run --nproc_per_node=4 \ + tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py \ + --source-root "$PWD" --model-root "$MINIMAX_H3_MODEL_ROOT" \ + --revision-label candidate-sp4 --operation decode +``` + +Run the same script with `--source-root` pointed at the base checkout for a comparable baseline. The default workload +is deterministic `124 x 768 x 1344` video geometry with seed `20260803`; the JSON record includes source/model +revisions, software/allocator metadata, exact measurement boundaries, per-repetition values, and output shapes. + FastVideo joint audio/video generation and SP=1/SP=4 latent consistency have been validated. T2VA, FL2VA, and Ref2VA video/audio latents match the pinned Diffusers pipeline exactly. diff --git a/tests/local_tests/minimax_h3/_reference.py b/tests/local_tests/minimax_h3/_reference.py index 2dbb77d21..ac80d7c84 100644 --- a/tests/local_tests/minimax_h3/_reference.py +++ b/tests/local_tests/minimax_h3/_reference.py @@ -3,6 +3,7 @@ import hashlib import inspect import os +import subprocess from pathlib import Path import pytest @@ -13,6 +14,28 @@ REFERENCE_SRC = REFERENCE_ROOT / "src" PINNED_COMMIT = "abc5e9bf71fd38f53cd471bc3acaa84bc5ecbfdc" +def _run_git(*args: str) -> str: + try: + result = subprocess.run( + ["git", "-C", str(REFERENCE_ROOT), *args], + check=True, + capture_output=True, + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError) as error: + raise RuntimeError(f"Could not verify the MiniMax-H3 reference checkout at {REFERENCE_ROOT}.") from error + return result.stdout.strip() + + +def assert_reference_revision() -> None: + actual_commit = _run_git("rev-parse", "HEAD") + if actual_commit != PINNED_COMMIT: + raise RuntimeError(f"MiniMax-H3 parity requires Diffusers commit {PINNED_COMMIT}, got {actual_commit}.") + dirty_source = _run_git("status", "--short", "--untracked-files=no", "--", "src/diffusers") + if dirty_source: + raise RuntimeError(f"MiniMax-H3 reference source has tracked changes:\n{dirty_source}") + + def assert_pinned_reference(relative_path: str, sha256: str) -> Path: path = REFERENCE_ROOT / relative_path if not path.is_file(): diff --git a/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py b/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py new file mode 100644 index 000000000..57e68a742 --- /dev/null +++ b/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Measure MiniMax-H3 production VAE stage memory with CPU offload enabled.""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import os +from pathlib import Path +import statistics +import subprocess +import sys +import time +from types import SimpleNamespace + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--source-root", type=Path, required=True) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--revision-label", required=True) + parser.add_argument("--operation", choices=("encode", "decode"), required=True) + parser.add_argument("--height", type=int, default=768) + parser.add_argument("--width", type=int, default=1344) + parser.add_argument("--num-frames", type=int, default=124) + parser.add_argument("--seed", type=int, default=20260803) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--repetitions", type=int, default=3) + return parser.parse_args() + + +def _git(source_root: Path, *args: str) -> str: + result = subprocess.run( + ["git", "-C", str(source_root), *args], + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for block in iter(lambda: file.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _model_snapshot(model_root: Path) -> str | None: + parts = model_root.resolve().parts + try: + return parts[parts.index("snapshots") + 1] + except (ValueError, IndexError): + return None + + +def _make_layout(rows, latent_shape): + import torch + + from fastvideo.pipelines.basic.minimax_h3.packing import MiniMaxH3PackedLayout + + empty = torch.empty(0, dtype=torch.long, device=rows.device) + return MiniMaxH3PackedLayout( + sequence_length=rows.shape[0], + position_ids=empty, + token_tags=empty, + video_indices=empty, + audio_indices=empty, + text_indices=empty, + num_condition_video_rows=0, + num_condition_audio_rows=0, + num_video_latent_frames=latent_shape[2], + latent_height=latent_shape[3], + latent_width=latent_shape[4], + num_audio_latents=0, + ) + + +def _build_operation(args, vae, device): + import numpy as np + import torch + + from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config + from fastvideo.pipelines.basic.minimax_h3.packing import patchify_video_latents, video_latent_num_frames + from fastvideo.pipelines.basic.minimax_h3.reference import MiniMaxH3PreparedReference + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_decoding import MiniMaxH3VideoDecodingStage + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import ( + MINIMAX_H3_LAYOUT_KEY, + MiniMaxH3LatentPreparationStage, + ) + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + patch_size = MiniMaxH3Config().arch_config.patch_size + transformer = SimpleNamespace(patch_size=patch_size) + runtime_args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True) + + if args.operation == "encode": + if int(os.environ.get("WORLD_SIZE", "1")) != 1: + raise ValueError("The encode benchmark is single-rank; use one process.") + frames = np.random.default_rng(args.seed).integers( + 0, + 256, + size=(args.num_frames, args.height, args.width, 3), + dtype=np.uint8, + ) + stage = MiniMaxH3LatentPreparationStage( + transformer=transformer, + vae=vae, + audio_vae=None, + scheduler=None, + ref2va=True, + ) + + def run_once(): + reference = MiniMaxH3PreparedReference(media_type="video", frames=frames) + vae.to(device) + try: + return stage._encode_visual_rows([reference], device)[0] + finally: + vae.to("cpu") + + return run_once, { + "input": "NumPy PCG64 CPU uint8 RGB pixels", + "input_shape": [1, 3, args.num_frames, args.height, args.width], + "boundary": "before VAE CPU-to-GPU transfer through post-encode VAE CPU offload", + } + + latent_shape = ( + 1, + vae.latent_channels, + video_latent_num_frames(args.num_frames), + args.height // vae.spatial_compression_ratio, + args.width // vae.spatial_compression_ratio, + ) + generator = torch.Generator(device=device).manual_seed(args.seed) + latents = torch.randn(latent_shape, generator=generator, device=device, dtype=torch.float32) + rows = patchify_video_latents(latents, patch_size) + layout = _make_layout(rows, latent_shape) + stage = MiniMaxH3VideoDecodingStage(vae, transformer) + + def run_once(): + batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape) + batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout + return stage.forward(batch, runtime_args).output + + return run_once, { + "input": "PyTorch Philox normalized FP32 latents", + "input_shape": list(latent_shape), + "boundary": "MiniMaxH3VideoDecodingStage.forward including VAE CPU-to-GPU transfer and CPU offload", + } + + +def _measure(run_once, device, warmups: int, repetitions: int) -> dict: + import torch + + for _ in range(warmups): + with torch.inference_mode(): + result = run_once() + torch.cuda.synchronize(device) + del result + gc.collect() + torch.cuda.empty_cache() + + records = [] + for repetition in range(repetitions): + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize(device) + start_allocated = torch.cuda.memory_allocated(device) + start_reserved = torch.cuda.memory_reserved(device) + torch.cuda.reset_peak_memory_stats(device) + started = time.perf_counter() + with torch.inference_mode(): + result = run_once() + torch.cuda.synchronize(device) + elapsed = time.perf_counter() - started + records.append({ + "repetition": repetition, + "elapsed_seconds": elapsed, + "start_allocated_bytes": start_allocated, + "peak_allocated_bytes": torch.cuda.max_memory_allocated(device), + "incremental_peak_allocated_bytes": torch.cuda.max_memory_allocated(device) - start_allocated, + "start_reserved_bytes": start_reserved, + "peak_reserved_bytes": torch.cuda.max_memory_reserved(device), + "incremental_peak_reserved_bytes": torch.cuda.max_memory_reserved(device) - start_reserved, + "output_shape": list(result.shape), + }) + del result + return { + "repetitions": records, + "median_elapsed_seconds": statistics.median(record["elapsed_seconds"] for record in records), + "max_peak_allocated_bytes": max(record["peak_allocated_bytes"] for record in records), + "max_incremental_peak_allocated_bytes": max( + record["incremental_peak_allocated_bytes"] for record in records), + "max_peak_reserved_bytes": max(record["peak_reserved_bytes"] for record in records), + "max_incremental_peak_reserved_bytes": max( + record["incremental_peak_reserved_bytes"] for record in records), + } + + +def main() -> None: + args = _parse_args() + if args.warmups < 0 or args.repetitions < 1: + raise ValueError("warmups must be non-negative and repetitions must be positive.") + source_root = args.source_root.resolve() + sys.path.insert(0, str(source_root)) + + import torch + import torch.distributed as dist + + import fastvideo + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + from fastvideo.models.loader.component_loader import VAELoader + + imported_root = Path(fastvideo.__file__).resolve().parents[1] + if imported_root != source_root: + raise RuntimeError(f"Imported FastVideo from {imported_root}, expected {source_root}.") + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29673") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + world_size = int(os.environ["WORLD_SIZE"]) + maybe_init_distributed_environment_and_model_parallel(1, world_size) + rank = dist.get_rank() + device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") + + component_dir = args.model_root.resolve() / "vae" + config_path = component_dir / "config.json" + index_path = component_dir / "diffusion_pytorch_model.safetensors.index.json" + for path in (config_path, index_path): + if not path.is_file(): + raise FileNotFoundError(path) + + loader_args = SimpleNamespace( + pipeline_config=MiniMaxH3PipelineConfig(), + model_paths={}, + vae_cpu_offload=True, + ) + vae = VAELoader().load(str(component_dir), loader_args) + parameter_bytes = sum(parameter.numel() * parameter.element_size() for parameter in vae.parameters()) + run_once, workload = _build_operation(args, vae, device) + rank_result = _measure(run_once, device, args.warmups, args.repetitions) + rank_result.update({ + "rank": rank, + "device_index": device.index, + "device_name": torch.cuda.get_device_name(device), + "device_total_memory_bytes": torch.cuda.get_device_properties(device).total_memory, + }) + rank_results = [None] * world_size + dist.all_gather_object(rank_results, rank_result) + + if rank == 0: + result = { + "schema_version": 1, + "exit_status": 0, + "revision_label": args.revision_label, + "source_root": str(source_root), + "source_git_head": _git(source_root, "rev-parse", "HEAD"), + "source_tracked_status": _git(source_root, "status", "--short", "--untracked-files=no"), + "benchmark_script_sha256": _sha256(Path(__file__).resolve()), + "model_root": str(args.model_root.resolve()), + "model_snapshot": _model_snapshot(args.model_root), + "model_config_sha256": _sha256(config_path), + "model_index_sha256": _sha256(index_path), + "vae_parameter_bytes": parameter_bytes, + "operation": args.operation, + "vae_cpu_offload": True, + "pin_cpu_memory": False, + "warmups": args.warmups, + "measurement_repetitions": args.repetitions, + "seed": args.seed, + "workload": workload, + "world_size": world_size, + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "allocator_config": os.environ.get("PYTORCH_ALLOC_CONF") + or os.environ.get("PYTORCH_CUDA_ALLOC_CONF"), + "rank_results": rank_results, + "sum_rank_local_max_peak_allocated_bytes": sum( + item["max_peak_allocated_bytes"] for item in rank_results), + "sum_rank_local_max_incremental_peak_allocated_bytes": sum( + item["max_incremental_peak_allocated_bytes"] for item in rank_results), + "sum_rank_local_max_peak_reserved_bytes": sum( + item["max_peak_reserved_bytes"] for item in rank_results), + "sum_rank_local_max_incremental_peak_reserved_bytes": sum( + item["max_incremental_peak_reserved_bytes"] for item in rank_results), + } + print("MINIMAX_H3_VAE_MEMORY=" + json.dumps(result, sort_keys=True), flush=True) + + cleanup_dist_env_and_memory() + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py b/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py index 1a2d6adc0..98308a441 100644 --- a/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py +++ b/tests/local_tests/vaes/test_minimax_h3_video_vae_parity.py @@ -62,8 +62,9 @@ def _require_assets() -> tuple[torch.device, Path]: def _load_official(component_dir: Path, device: torch.device) -> torch.nn.Module: from diffusers.models.autoencoders.autoencoder_kl_minimax_h3 import AutoencoderKLMiniMaxH3 - from tests.local_tests.minimax_h3._reference import assert_reference_source + from tests.local_tests.minimax_h3._reference import assert_reference_revision, assert_reference_source + assert_reference_revision() assert_reference_source( AutoencoderKLMiniMaxH3, "src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py", @@ -101,11 +102,27 @@ def _load_fastvideo(component_dir: Path) -> torch.nn.Module: return model -def _make_video() -> torch.Tensor: +def _make_pixels() -> torch.Tensor: generator = torch.Generator(device="cpu").manual_seed(20260803) # Two logical 17-frame clips after padding are required for H3's three-token # temporal overlap contract; 32 px remains safely above reflect-pad minima. - return torch.randn(1, 3, 22, 32, 32, generator=generator, dtype=torch.float32) + return torch.randint(0, 256, (1, 3, 22, 32, 32), generator=generator, dtype=torch.uint8) + + +def _normalize_pixels(pixels: torch.Tensor, device: torch.device) -> torch.Tensor: + video = pixels.to(device=device, dtype=torch.float32).div_(255.0) + pixel_mean = torch.tensor((0.485, 0.456, 0.406), device=device).view(1, -1, 1, 1, 1) + pixel_std = torch.tensor((0.229, 0.224, 0.225), device=device).view(1, -1, 1, 1, 1) + return (video - pixel_mean) / pixel_std + + +def _run_encode_pixels(model: torch.nn.Module, pixels: torch.Tensor) -> dict[str, torch.Tensor]: + with torch.inference_mode(): + posterior = model.encode_pixels(pixels, return_dict=False)[0] + return { + "mean": posterior.mean.detach().cpu(), + "logvar": posterior.logvar.detach().cpu(), + } def _run( @@ -169,15 +186,16 @@ def _assert_tensor_parity(name: str, actual: torch.Tensor, expected: torch.Tenso def test_minimax_h3_video_vae_parity() -> None: """Match posterior, normalization, geometry, and deterministic decode.""" device, component_dir = _require_assets() - video = _make_video() + pixels = _make_pixels() official = _load_official(component_dir, device) - expected = _run(official, video.to(device)) + expected = _run(official, _normalize_pixels(pixels, device)) del official _reclaim_vram() fastvideo = _load_fastvideo(component_dir) - actual = _run(fastvideo, video.to(device), stream_output=True) + actual = _run(fastvideo, _normalize_pixels(pixels, device), stream_output=True) + streaming_encode = _run_encode_pixels(fastvideo, pixels) assert fastvideo.temporal_compression_ratio == 4 assert fastvideo.spatial_compression_ratio == 16 del fastvideo @@ -185,6 +203,8 @@ def test_minimax_h3_video_vae_parity() -> None: _assert_tensor_parity("video_vae.mean", actual["mean"], expected["mean"], 0.0) _assert_tensor_parity("video_vae.logvar", actual["logvar"], expected["logvar"], 0.0) + _assert_tensor_parity("video_vae.streaming_encode.mean", streaming_encode["mean"], expected["mean"], 0.0) + _assert_tensor_parity("video_vae.streaming_encode.logvar", streaming_encode["logvar"], expected["logvar"], 0.0) _assert_tensor_parity("video_vae.normalized", actual["normalized"], expected["normalized"], 0.0) _assert_tensor_parity("video_vae.decode", actual["decoded"], expected["decoded"], 0.0) _assert_tensor_parity("video_vae.streaming_decode", actual["streamed"], expected["pixels"], 0.0)