From bd5ee4c228c54bc983ffe5868ec5a21f7b5415e4 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Mon, 31 Aug 2026 12:09:10 -0700 Subject: [PATCH 1/2] [bugfix]: read H3 geometry from checkpoint JSON so lazy load can drop the DiT before VAE decode Input prep and unpatchify were holding live VAE/DiT proxies just for two integers, which loaded the video VAE before Qwen and kept the DiT resident through decode. Auto-enable --lazy-module-load on unified memory and on single-GPU FastH3 examples. --- .../installation/spark_performance.md | 9 ++- docs/inference/offloading.md | 34 ++++++----- examples/inference/basic/basic_fasth3.py | 11 ++-- .../inference/basic/basic_minimax_h3_t2v.py | 11 ++-- fastvideo/api/schema.py | 4 +- fastvideo/fastvideo_args.py | 34 ++++++++--- .../basic/minimax_h3/minimax_h3_pipeline.py | 50 ++++++++++++---- .../pipelines/basic/minimax_h3/packing.py | 15 ++++- .../minimax_h3/stages/minimax_h3_decoding.py | 6 +- .../stages/minimax_h3_latent_preparation.py | 21 ++++--- fastvideo/tests/api/test_parser.py | 2 +- .../inference/test_basic_fasth3_profile.py | 12 ++++ .../platforms/test_unified_memory_offload.py | 16 +++++ .../tests/stages/test_lazy_module_load.py | 60 ++++++++++++++++++- .../stages/test_minimax_h3_vae_streaming.py | 33 ++++++---- .../benchmark_minimax_h3_video_vae_memory.py | 14 +++-- 16 files changed, 250 insertions(+), 82 deletions(-) diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index 96284efe89..db8d4b21c5 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -164,9 +164,12 @@ is power-cycled. To avoid it: - **MiniMax H3 / FastH3** still needs sequential loading on one GB10. The Qwen3-VL conditioner is tens of gigabytes of BF16. If the DiT and VAEs load while that encoder is still resident, the process is a typical `earlyoom` kill (Python is - preferred). The CUDA pipeline now encodes first, releases the encoder, then - loads DiT and VAEs onto the accelerator (`to_cpu` follows `cpu_offload`, which - is off here). See [Offloading](../../inference/offloading.md). + preferred). Pass `--lazy-module-load` (or omit it: FastVideo auto-enables the + flag on unified memory, and `basic_fasth3.py` defaults it on when + `--num-gpus 1`). That loads Qwen, releases it, loads the DiT, releases the DiT, + then loads the VAE. Geometry scalars come from checkpoint `config.json`, not + from live weights. A later `generate()` on the same worker reloads from disk. + See [Offloading](../../inference/offloading.md). ## Gotchas specific to the GB10 diff --git a/docs/inference/offloading.md b/docs/inference/offloading.md index 5d65e468f3..3dd2146c09 100644 --- a/docs/inference/offloading.md +++ b/docs/inference/offloading.md @@ -12,7 +12,7 @@ text_encoder_cpu_offload: bool = True image_encoder_cpu_offload: bool = True vae_cpu_offload: bool = True pin_cpu_memory: bool = True -lazy_module_load: bool = False +lazy_module_load: bool | None = None ``` On unified-memory accelerators such as NVIDIA GB10 and Apple silicon, FastVideo @@ -23,14 +23,18 @@ memory. CUDA FSDP sharding remains enabled when requested; MPS continues to disable FSDP. `pin_cpu_memory` is not an offload mode and is left unchanged. MiniMax H3 CUDA inference uses a second lever that does not copy weights to a -host pool. The pipeline loads the Qwen3-VL text encoder, runs conditioning, then -releases that encoder before it loads the DiT and video/audio VAEs. The MLX FastH3 -runtime uses the same phase order. When host offload is off, DiT safetensors are -read onto the accelerator instead of CPU-then-copy. Input-preparation geometry -(spatial ratio, latent channels, audio sample rate) comes from the VAE arch -configs until those weights load. A later `generate()` on the same worker -currently re-enters conditioning after the encoder has been released; start a -new generator for a new prompt until prompt-cache reload exists. +host pool. With `lazy_module_load`, the pipeline loads the Qwen3-VL text encoder, +runs conditioning, and releases that encoder before it loads the DiT. After +denoise it releases the DiT, then loads the video VAE for decode. Input +preparation and unpatchify read geometry from checkpoint `config.json` (VAE +spatial ratio / latent channels, DiT patch size) so those stages do not +materialize weights just to read two integers. The MLX FastH3 runtime uses the +same phase order. When host offload is off, DiT safetensors are read onto the +accelerator instead of CPU-then-copy. A later `generate()` on the same worker +reloads a released component from disk; start a new generator for a new prompt +until prompt-cache reload exists. The flag is auto-enabled on unified-memory +devices. Pass `--no-lazy-module-load` (or `lazy_module_load=False`) to keep every +component resident. ## Behavior Explanation @@ -109,9 +113,9 @@ By default a pipeline loads every component before the first stage runs, so peak memory is the sum of all of them even though no two are needed at the same moment. With `lazy_module_load` enabled, each heavy component loads on first use and is freed once the last stage that needs it has returned, so peak memory -becomes the largest overlapping set instead of the sum. For a text-to-video -pipeline that is roughly `max(text encoder, DiT + VAE)` rather than -`text encoder + DiT + VAE`. +becomes the largest overlapping set instead of the sum. MiniMax-H3 T2VA is +`max(text encoder, DiT, VAE)` rather than `text encoder + DiT + VAE`, because +the DiT is not held through VAE decode. #### Performance Impact @@ -126,8 +130,10 @@ and kernel caches when the component structure and input shapes are unchanged. Enable this when a model does not fit at load time, which the CPU offload options above cannot help with because they act after loading. It is particularly relevant on unified-memory devices, where host and device draw on -the same pool and moving weights to the host frees nothing. Leave it off when -the model already fits. +the same pool and moving weights to the host frees nothing. FastVideo +auto-enables it there (`lazy_module_load=None`). Leave it off when the model +already fits, or pass `--no-lazy-module-load` to keep components resident for +later `generate()` calls. This option applies to inference only. Training keeps every component resident and logs a warning if the flag is set. diff --git a/examples/inference/basic/basic_fasth3.py b/examples/inference/basic/basic_fasth3.py index 38c407cb99..78b9d6dc30 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/examples/inference/basic/basic_fasth3.py @@ -49,11 +49,13 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser: parser.add_argument("--prompt", required=True) parser.add_argument("--output", default="outputs/fasth3") parser.add_argument("--lazy-module-load", - action="store_true", + action=argparse.BooleanOptionalAction, + default=None, help="load each heavy component on first use and free it after the last stage that " "needs it, so peak memory is the largest overlapping set instead of the sum of every " - "component. Enable when the model does not fit at load time; costs a reload per " - "generation, so leave it off when it does fit") + "component. Default: on when --num-gpus is 1; FastVideo also auto-enables on unified " + "memory. Costs a reload per generation; pass --no-lazy-module-load to keep every " + "component resident") parser.add_argument("--profile", choices=("all", "strict"), default="all", @@ -259,7 +261,8 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig: text_encoder=True, vae=True, pin_cpu_memory=args.pin_cpu_memory, - lazy_module_load=args.lazy_module_load, + lazy_module_load=(True if args.lazy_module_load is None and args.num_gpus == 1 else + args.lazy_module_load), ), compile=CompileConfig( enabled=args.torch_compile, diff --git a/examples/inference/basic/basic_minimax_h3_t2v.py b/examples/inference/basic/basic_minimax_h3_t2v.py index ea15bacde6..719e46780e 100644 --- a/examples/inference/basic/basic_minimax_h3_t2v.py +++ b/examples/inference/basic/basic_minimax_h3_t2v.py @@ -49,11 +49,13 @@ def parse_args() -> argparse.Namespace: "First generation pays the inductor JIT (~1-2 min); use --repeats >= 2 and time " "the last repeat. FASTVIDEO_INFERENCE_TORCH_COMPILE=1 is equivalent") parser.add_argument("--lazy-module-load", - action="store_true", + action=argparse.BooleanOptionalAction, + default=None, help="load each heavy component on first use and free it after the last stage that " "needs it, so peak memory is the largest overlapping set instead of the sum of every " - "component. Enable when the model does not fit at load time; costs a reload per " - "generation, so leave it off when it does fit") + "component. Default: on when --num-gpus is 1; FastVideo also auto-enables on unified " + "memory. Costs a reload per generation; pass --no-lazy-module-load to keep every " + "component resident") parser.add_argument("--repeats", type=int, default=1, @@ -87,7 +89,8 @@ def main() -> None: text_encoder=True, vae=True, pin_cpu_memory=False, - lazy_module_load=args.lazy_module_load, + lazy_module_load=(True if args.lazy_module_load is None and args.num_gpus == 1 else + args.lazy_module_load), ), compile=CompileConfig( enabled=args.torch_compile, diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index f664384176..77224d0368 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -34,8 +34,8 @@ class OffloadConfig: # after the last stage that needs it, so peak memory is the largest # overlapping set rather than the sum. Grouped here because it is the same # decision the offload knobs answer, which is how much of the model has to - # be resident at once. - lazy_module_load: bool = False + # be resident at once. ``None`` auto-enables on unified-memory devices. + lazy_module_load: bool | None = None @dataclass diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 74c9244de1..551df9c127 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -162,11 +162,13 @@ class FastVideoArgs: # Load each heavy component on first use and free it once the last stage # that holds it has run, instead of keeping every component resident from # load time to shutdown. Peak memory becomes the largest overlapping set - # rather than the sum of all components. Off by default: a released - # component is re-read from disk on the next generation, so this trades - # per-request latency for headroom and only pays off when the sum does not - # fit. Inference only; training keeps every component resident. - lazy_module_load: bool = False + # rather than the sum of all components. ``None`` (auto) turns this on for + # unified-memory devices (GB10 / Spark) after the worker binds its device, + # and leaves it off on discrete GPUs. Explicit True / False overrides the + # probe. A released component is re-read from disk on the next generation, + # so this trades per-request latency for headroom. Inference only; training + # keeps every component resident. + lazy_module_load: bool | None = None # Sequence-parallel MiniMax-H3 VAE (opt-in, default off). With SP > 1 the # video VAE's temporal chunks (decode) and clips (reference encode) are @@ -718,10 +720,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: ) parser.add_argument( "--lazy-module-load", - action=StoreBoolean, + action=argparse.BooleanOptionalAction, + default=None, help="Load each heavy component on first use and free it after the last stage that needs it, " - "so peak memory is the largest overlapping set of components instead of their sum. Enable when a " - "model does not fit at load time. Costs a reload per generation, so leave it off when it does fit.", + "so peak memory is the largest overlapping set of components instead of their sum. " + "Omit for auto (on for unified-memory devices such as GB10; off on discrete GPUs). " + "Pass --no-lazy-module-load to keep every component resident.", ) parser.add_argument( "--pin-cpu-memory", @@ -964,6 +968,20 @@ def _resolve_device_offload_conflicts(self) -> None: def finalize_device_offload_policy(self, device_id: int = 0) -> bool: """Apply device-local memory policy, then resolve incompatible modes.""" has_unified_memory = self.disable_offload_on_unified_memory(device_id) + if self.lazy_module_load is None: + self.lazy_module_load = bool(has_unified_memory) and not self.training_mode + if self.lazy_module_load: + from fastvideo.platforms import current_platform + + try: + device_name = current_platform.get_device_name(device_id) + except Exception: + device_name = current_platform.device_name + logger.info( + "Enabling lazy_module_load: %s has unified memory, so encoder, DiT, and VAEs cannot stay " + "resident together. Pass --no-lazy-module-load to keep every component loaded.", + device_name, + ) self._resolve_device_offload_conflicts() return has_unified_memory diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index 25dc999856..c39e7e8614 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -3,8 +3,12 @@ from __future__ import annotations +from pathlib import Path + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.hf_transformer_utils import get_diffusers_config from fastvideo.pipelines.basic.minimax_h3.stages import ( MiniMaxH3AudioDecodingStage, MiniMaxH3ConditioningStage, @@ -16,6 +20,26 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline +logger = init_logger(__name__) + + +def _apply_h3_checkpoint_arch_configs(model_path: str, fastvideo_args: FastVideoArgs, + extra_config_module_map: dict[str, str]) -> None: + """Overlay checkpoint config.json onto pipeline configs without loading weights.""" + root = Path(model_path) + vae_dir = root / "vae" + if (vae_dir / "config.json").is_file(): + fastvideo_args.pipeline_config.vae_config.update_model_arch(get_diffusers_config(str(vae_dir))) + transformer_dir = root / extra_config_module_map.get("transformer", "transformer") + if (transformer_dir / "config.json").is_file(): + fastvideo_args.pipeline_config.dit_config.update_model_arch(get_diffusers_config(str(transformer_dir))) + logger.info( + "MiniMax-H3 geometry from config: patch_size=%s spatial_compression_ratio=%s latent_channels=%s", + tuple(fastvideo_args.pipeline_config.dit_config.patch_size), + int(fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio), + int(fastvideo_args.pipeline_config.vae_config.arch_config.latent_channels), + ) + class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): """Shared loading and target-generation path for MiniMax H3. @@ -52,9 +76,10 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): "scheduler", "audio_scheduler", ] - # Deferral is safe here: no stage reads a component's attributes while it - # is being constructed, and `initialize_pipeline` only inspects the - # schedulers, which are never deferred. + # Deferral is safe here: geometry scalars come from checkpoint config.json + # (applied in initialize_pipeline without loading weights), no stage + # constructor reads a deferred component, and initialize_pipeline only + # inspects the schedulers, which are never deferred. _lazy_module_names = ("text_encoder", "transformer", "vae", "audio_vae") @classmethod @@ -62,7 +87,7 @@ def get_hf_download_component_dirs(cls) -> tuple[str, ...]: return tuple(sorted(cls._extra_config_module_map.get(name, name) for name in cls._required_config_modules)) def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: - del fastvideo_args + _apply_h3_checkpoint_arch_configs(self.model_path, fastvideo_args, self._extra_config_module_map) for module_name, modality, expected_shift in ( ("scheduler", "video", 12.0), ("audio_scheduler", "audio", 3.0), @@ -71,17 +96,21 @@ def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: if shift is None or float(shift) != expected_shift: raise ValueError(f"MiniMax-H3 {modality} scheduler must expose shift={expected_shift:g}, got {shift}.") - def _add_stages(self, *, ref2va: bool) -> None: + def _add_stages(self, fastvideo_args: FastVideoArgs, *, ref2va: bool) -> None: transformer = self.get_module("transformer") vae = self.get_module("vae") audio_vae = self.get_module("audio_vae") scheduler = self.get_module("scheduler") audio_scheduler = self.get_module("audio_scheduler") + # Geometry scalars live on the checkpoint-updated arch config. Holding + # the live VAE/DiT here would materialize them on the first attribute + # read. Encode still needs the live VAE for FL2VA/Ref2VA. + video_geometry = fastvideo_args.pipeline_config.vae_config.arch_config self.add_stage( "input_preparation_stage", MiniMaxH3InputPreparationStage( - vae=vae, + vae=video_geometry, audio_vae=audio_vae if ref2va else None, ref2va=ref2va, ), @@ -98,7 +127,6 @@ def _add_stages(self, *, ref2va: bool) -> None: self.add_stage( "latent_preparation_stage", MiniMaxH3LatentPreparationStage( - transformer=transformer, vae=vae, audio_vae=audio_vae, scheduler=scheduler, @@ -113,7 +141,7 @@ def _add_stages(self, *, ref2va: bool) -> None: audio_scheduler=audio_scheduler, ), ) - self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=vae, transformer=transformer)) + self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=vae)) self.add_stage("audio_decoding_stage", MiniMaxH3AudioDecodingStage(audio_vae=audio_vae)) @@ -121,8 +149,7 @@ class MiniMaxH3Pipeline(MiniMaxH3BasePipeline): """One-request joint video/stereo-audio pipeline for T2VA and FL2VA.""" def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - del fastvideo_args - self._add_stages(ref2va=False) + self._add_stages(fastvideo_args, ref2va=False) class MiniMaxH3RefPipeline(MiniMaxH3BasePipeline): @@ -131,8 +158,7 @@ class MiniMaxH3RefPipeline(MiniMaxH3BasePipeline): _extra_config_module_map = {"transformer": "transformer_ref"} def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - del fastvideo_args - self._add_stages(ref2va=True) + self._add_stages(fastvideo_args, ref2va=True) class MiniMaxH3ModularPipeline(MiniMaxH3Pipeline): diff --git a/fastvideo/pipelines/basic/minimax_h3/packing.py b/fastvideo/pipelines/basic/minimax_h3/packing.py index 7ef3468512..1c67c49a35 100644 --- a/fastvideo/pipelines/basic/minimax_h3/packing.py +++ b/fastvideo/pipelines/basic/minimax_h3/packing.py @@ -4,7 +4,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -38,6 +38,19 @@ MINIMAX_H3_KEYFRAME_NOISE_AUG = 0.999 MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42 + +def h3_dit_patch_size(fastvideo_args: Any) -> tuple[int, int, int]: + """Read DiT patch size from pipeline config, not live transformer weights.""" + dit_config = getattr(getattr(fastvideo_args, "pipeline_config", None), "dit_config", None) + patch_size = getattr(dit_config, "patch_size", None) + if patch_size is None: + raise ValueError("MiniMax-H3 requires pipeline_config.dit_config.patch_size.") + values = tuple(int(axis) for axis in patch_size) + if len(values) != 3 or min(values) <= 0: + raise ValueError(f"MiniMax-H3 patch_size must be three positive ints, got {patch_size!r}.") + return values + + MINIMAX_H3_ROPE_FRAME_RESCALE = 5.0 / 3.0 MINIMAX_H3_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) _ROPE_SPATIAL_SCALE = 32 diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py index 7cbc0d6e28..13f45c97c3 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -16,6 +16,7 @@ from fastvideo.profiler import nvtx_range from fastvideo.pipelines.basic.minimax_h3.packing import ( MiniMaxH3PackedLayout, + h3_dit_patch_size, unpack_audio_tokens, unpatchify_video_tokens, ) @@ -58,10 +59,9 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): performance_component_metric = "vae_decode_time_s" - def __init__(self, vae: AutoencoderKLMiniMaxH3, transformer: Any) -> None: + def __init__(self, vae: AutoencoderKLMiniMaxH3) -> None: super().__init__() self.vae = vae - self.transformer = transformer def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() @@ -97,7 +97,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward latent_height, latent_width, channels, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), ) device = get_local_torch_device() self.vae.to(device) diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py index b30e599f7e..701ec1982b 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py @@ -21,6 +21,7 @@ audio_latent_num_frames, build_packed_sequence, build_ref2va_packed_sequence, + h3_dit_patch_size, keyframe_condition_noise, patchify_video_latents, ) @@ -65,7 +66,6 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): def __init__( self, - transformer: Any, vae: Any, audio_vae: Any, scheduler: Any, @@ -73,7 +73,6 @@ def __init__( ref2va: bool = False, ) -> None: super().__init__() - self.transformer = transformer self.vae = vae self.audio_vae = audio_vae self.scheduler = scheduler @@ -111,7 +110,7 @@ def _encode_visual_rows( device: torch.device, fastvideo_args: FastVideoArgs, ) -> list[torch.Tensor]: - patch_size = self.transformer.patch_size + patch_size = h3_dit_patch_size(fastvideo_args) # Reference encode runs on every rank (all ranks hold identical # prepared references), so clip-parallel encode keeps participation # uniform by construction: each rank encodes a clip subset and the @@ -184,7 +183,7 @@ def _encode_fl2va_conditions( for image in keyframes: clean_rows.append( patchify_video_latents(self._encode_keyframe_latents(image, vae_device), - self.transformer.patch_size)) + h3_dit_patch_size(fastvideo_args))) finally: if fastvideo_args.vae_cpu_offload: self.vae.to("cpu") @@ -193,7 +192,7 @@ def _encode_fl2va_conditions( shapes = ((1, latent_height, latent_width), ) * len(keyframes) noise = keyframe_condition_noise( shapes, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), self.vae.latent_channels, generator=batch.generator, device=device, @@ -241,7 +240,7 @@ def _encode_ref2va_conditions( for reference in references if reference.media_type != "audio") noise = keyframe_condition_noise( shapes, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), self.vae.latent_channels, generator=batch.generator, device=device, @@ -259,7 +258,7 @@ def _encode_ref2va_conditions( reference.waveform = None return video_conditions, audio_conditions - def _build_layout(self, batch: ForwardBatch) -> MiniMaxH3PackedLayout: + def _build_layout(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> MiniMaxH3PackedLayout: text_token_tags = batch.extra.get(MINIMAX_H3_TEXT_TOKEN_TAGS_KEY) if not isinstance(text_token_tags, torch.Tensor): raise ValueError("MiniMax-H3 conditioning must produce text token tags.") @@ -276,7 +275,7 @@ def _build_layout(self, batch: ForwardBatch) -> MiniMaxH3PackedLayout: height, width, num_audio_latents, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), ) anchors = batch.extra.get(MINIMAX_H3_KEYFRAME_ANCHORS_KEY, ()) if not isinstance(anchors, tuple): @@ -287,7 +286,7 @@ def _build_layout(self, batch: ForwardBatch) -> MiniMaxH3PackedLayout: height, width, num_audio_latents, - self.transformer.patch_size, + h3_dit_patch_size(fastvideo_args), anchors, ) @@ -301,7 +300,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward else: condition_video, condition_audio = self._encode_fl2va_conditions(batch, fastvideo_args, device) - layout = self._build_layout(batch) + layout = self._build_layout(batch, fastvideo_args) video_channels, num_frames, height, width = _video_geometry(batch) expected_video_shape = (1, video_channels, num_frames, height, width) if video_noise is None: @@ -315,7 +314,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward raise ValueError(f"MiniMax-H3 injected video latents must have shape {expected_video_shape}, " f"got {tuple(video_noise.shape)}.") video_rows = patchify_video_latents(video_noise.to(device=device, dtype=torch.float32), - self.transformer.patch_size) + h3_dit_patch_size(fastvideo_args)) num_audio_latents = layout.num_audio_latents expected_audio_shape = (MINIMAX_H3_AUDIO_CHANNELS, self.audio_vae.latent_channels, num_audio_latents) diff --git a/fastvideo/tests/api/test_parser.py b/fastvideo/tests/api/test_parser.py index e5086bfd7b..9b0f2104ee 100644 --- a/fastvideo/tests/api/test_parser.py +++ b/fastvideo/tests/api/test_parser.py @@ -114,7 +114,7 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: "image_encoder": True, "vae": True, "pin_cpu_memory": True, - "lazy_module_load": False, + "lazy_module_load": None, }, "compile": { "enabled": False, diff --git a/fastvideo/tests/inference/test_basic_fasth3_profile.py b/fastvideo/tests/inference/test_basic_fasth3_profile.py index e4f8bc2508..f0d89549a3 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_profile.py +++ b/fastvideo/tests/inference/test_basic_fasth3_profile.py @@ -84,6 +84,18 @@ def test_default_all_profile_matches_fastest_contract(tmp_path): assert request.sampling.guidance_scale == 1.0 assert request.sampling.batch_cfg is False assert request.output.output_path == str(tmp_path / "result.mp4") + assert config.engine.offload.lazy_module_load is None + + +def test_lazy_module_load_defaults_on_for_single_gpu(): + config = fasth3.build_generator_config(_args("--num-gpus", "1")) + assert config.engine.offload.lazy_module_load is True + + enabled = fasth3.build_generator_config(_args("--lazy-module-load")) + assert enabled.engine.offload.lazy_module_load is True + + disabled = fasth3.build_generator_config(_args("--no-lazy-module-load")) + assert disabled.engine.offload.lazy_module_load is False @pytest.mark.parametrize("num_frames", (124, 243, 345)) diff --git a/fastvideo/tests/platforms/test_unified_memory_offload.py b/fastvideo/tests/platforms/test_unified_memory_offload.py index aec74b0960..cc35c30019 100644 --- a/fastvideo/tests/platforms/test_unified_memory_offload.py +++ b/fastvideo/tests/platforms/test_unified_memory_offload.py @@ -104,6 +104,7 @@ def test_discrete_device_finalization_retains_layerwise_precedence(monkeypatch) assert args.text_encoder_cpu_offload is True assert args.image_encoder_cpu_offload is True assert args.vae_cpu_offload is True + assert args.lazy_module_load is False def test_workers_classify_their_own_device(monkeypatch) -> None: @@ -189,3 +190,18 @@ def unsupported_name(device_id): assert args.disable_offload_on_unified_memory() is True assert args.text_encoder_cpu_offload is False + + +def test_unified_device_auto_enables_lazy_module_load(as_unified_cuda) -> None: + args = FastVideoArgs(model_path="unused/for-this-test") + + assert args.lazy_module_load is None + assert args.finalize_device_offload_policy(device_id=6) is True + assert args.lazy_module_load is True + + +def test_explicit_false_lazy_module_load_stays_off_on_unified(as_unified_cuda) -> None: + args = FastVideoArgs(model_path="unused/for-this-test", lazy_module_load=False) + + args.finalize_device_offload_policy(device_id=6) + assert args.lazy_module_load is False diff --git a/fastvideo/tests/stages/test_lazy_module_load.py b/fastvideo/tests/stages/test_lazy_module_load.py index caa3088007..5f1d25feba 100644 --- a/fastvideo/tests/stages/test_lazy_module_load.py +++ b/fastvideo/tests/stages/test_lazy_module_load.py @@ -394,6 +394,7 @@ def test_schedule_is_empty_without_lazy_modules(): (True, False, True), (True, True, False), (False, True, False), + (None, False, False), ]) def test_training_mode_never_defers(lazy, training, expected): args = SimpleNamespace(lazy_module_load=lazy, training_mode=training) @@ -401,11 +402,11 @@ def test_training_mode_never_defers(lazy, training, expected): assert ComposedPipelineBase._lazy_module_load_enabled(args) is expected -def test_flag_defaults_to_off(): +def test_flag_defaults_to_auto(): from fastvideo.fastvideo_args import FastVideoArgs fields = {f.name: f for f in dataclasses.fields(FastVideoArgs)} - assert fields["lazy_module_load"].default is False + assert fields["lazy_module_load"].default is None # ---------------------------------------------------------------------- @@ -694,7 +695,9 @@ def test_building_the_real_h3_stages_materializes_nothing(): # `DenoisingStage.__init__` in the shared stage set reads # `transformer.hidden_size` to pick an attention backend, which would pull # the DiT in during post_init. H3's stages must not acquire that habit. + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline + from fastvideo.pipelines.composed_pipeline_base import _iter_held_objects loaded: list[str] = [] @@ -714,11 +717,62 @@ def tracked(name): "scheduler": object(), "audio_scheduler": object(), } + args = SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig()) - pipeline._add_stages(ref2va=False) + pipeline._add_stages(args, ref2va=False) assert loaded == [], f"building stages materialized {loaded}" assert len(pipeline._stages) == 6 + input_held = {id(obj) for obj in _iter_held_objects(pipeline._stage_name_mapping["input_preparation_stage"])} + latent_held = {id(obj) for obj in _iter_held_objects(pipeline._stage_name_mapping["latent_preparation_stage"])} + decode_held = {id(obj) for obj in _iter_held_objects(pipeline._stage_name_mapping["video_decoding_stage"])} + denoise_held = {id(obj) for obj in _iter_held_objects(pipeline._stage_name_mapping["denoising_stage"])} + assert id(pipeline.modules["vae"]) not in input_held + assert id(pipeline.modules["transformer"]) not in input_held + assert id(pipeline.modules["transformer"]) not in latent_held + assert id(pipeline.modules["transformer"]) not in decode_held + assert id(pipeline.modules["transformer"]) in denoise_held + assert id(pipeline.modules["vae"]) in decode_held + + +def test_h3_lazy_release_drops_dit_before_vae_decode(): + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline + + pipeline = MiniMaxH3Pipeline.__new__(MiniMaxH3Pipeline) + pipeline._stages = [] + pipeline._stage_name_mapping = {} + pipeline.modules = { + "text_encoder": LazyModule("text_encoder", lambda: _Component("text_encoder")), + "transformer": LazyModule("transformer", lambda: _Component("transformer")), + "vae": LazyModule("vae", lambda: _Component("vae")), + "audio_vae": LazyModule("audio_vae", lambda: _Component("audio_vae")), + "tokenizer": object(), + "processor": object(), + "scheduler": object(), + "audio_scheduler": object(), + } + args = SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig()) + pipeline._add_stages(args, ref2va=False) + schedule = pipeline._build_lazy_release_schedule() + names = {pipeline._stages[index]._pipeline_stage_name: modules for index, modules in schedule.items()} + assert names["conditioning_stage"] == ["text_encoder"] + assert names["denoising_stage"] == ["transformer"] + assert "transformer" not in names.get("video_decoding_stage", []) + assert "vae" in names["video_decoding_stage"] + + +def test_h3_checkpoint_json_updates_dit_patch_size_without_weights(tmp_path): + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import _apply_h3_checkpoint_arch_configs + + transformer_dir = tmp_path / "transformer" + transformer_dir.mkdir() + (transformer_dir / "config.json").write_text('{"patch_size": [1, 1, 1]}') + args = SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig()) + assert tuple(args.pipeline_config.dit_config.patch_size) == (1, 2, 2) + _apply_h3_checkpoint_arch_configs(str(tmp_path), args, {}) + assert tuple(args.pipeline_config.dit_config.patch_size) == (1, 1, 1) class _LoRAConfigComponent(torch.nn.Module): diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py index 4e35bf66fe..6202e11d02 100644 --- a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -51,7 +51,6 @@ def normalize_latents(self, latents): return latents stage = MiniMaxH3LatentPreparationStage( - transformer=SimpleNamespace(patch_size=(1, 1, 1)), vae=VAE(), audio_vae=None, scheduler=None, @@ -61,7 +60,10 @@ def normalize_latents(self, latents): media_type="video", frames=np.zeros((22, 16, 16, 3), dtype=np.uint8), ) - args = SimpleNamespace(vae_parallel_encode=False) + args = SimpleNamespace( + vae_parallel_encode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))), + ) rows = stage._encode_visual_rows([reference], torch.device("cpu"), args) assert observed["pixels"].dtype == torch.uint8 @@ -95,9 +97,15 @@ def decode_to_pixels(self, decoded_latents, output): output.fill_(0.25) monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu")) - result = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward( + result = MiniMaxH3VideoDecodingStage(VAE()).forward( batch, - SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=False, vae_parallel_decode=False), + SimpleNamespace( + output_type="pil", + pin_cpu_memory=False, + vae_cpu_offload=False, + vae_parallel_decode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))), + ), ) torch.testing.assert_close(observed["latents"], latents) @@ -120,7 +128,7 @@ def to(self, device): monkeypatch.setattr(minimax_h3_decoding, "get_world_group", lambda: SimpleNamespace(is_first_rank=False)) args = SimpleNamespace(output_type="pil", pin_cpu_memory=False, vae_cpu_offload=True, vae_parallel_decode=False) - video = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace()).forward(ForwardBatch(data_type="video"), args) + video = MiniMaxH3VideoDecodingStage(VAE()).forward(ForwardBatch(data_type="video"), args) assert video.output.shape == (0, 3, 0, 0, 0) audio_batch = ForwardBatch(data_type="audio", latents=torch.zeros(1), audio_latents=torch.zeros(1)) @@ -161,11 +169,14 @@ def fake_parallel(vae, latents, output, group, strategy): monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu")) monkeypatch.setattr(minimax_h3_decoding, "model_parallel_is_initialized", lambda: True) monkeypatch.setattr(minimax_h3_decoding, "decode_to_pixels_parallel", fake_parallel) - args = SimpleNamespace(output_type="pil", - pin_cpu_memory=False, - vae_cpu_offload=False, - vae_parallel_decode=True, - vae_parallel_decode_strategy="gather") + args = SimpleNamespace( + output_type="pil", + pin_cpu_memory=False, + vae_cpu_offload=False, + vae_parallel_decode=True, + vae_parallel_decode_strategy="gather", + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 1, 1))), + ) for rank, is_first in ((0, True), (2, False)): monkeypatch.setattr( @@ -175,7 +186,7 @@ def fake_parallel(vae, latents, output, group, strategy): rank_in_group=rank)) batch = ForwardBatch(data_type="video", latents=rows.clone(), raw_latent_shape=latent_shape) batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape) - result = MiniMaxH3VideoDecodingStage(VAE(), SimpleNamespace(patch_size=(1, 1, 1))).forward(batch, args) + result = MiniMaxH3VideoDecodingStage(VAE()).forward(batch, args) if is_first: assert result.output.shape == (1, 3, 5, 16, 16) assert torch.all(result.output == 0.5) 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 index 57e68a7426..2bd39b9c0a 100644 --- a/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py +++ b/tests/local_tests/vaes/benchmark_minimax_h3_video_vae_memory.py @@ -94,8 +94,13 @@ def _build_operation(args, vae, device): 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) + runtime_args = SimpleNamespace( + output_type="pil", + pin_cpu_memory=False, + vae_cpu_offload=True, + vae_parallel_encode=False, + pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=patch_size)), + ) if args.operation == "encode": if int(os.environ.get("WORLD_SIZE", "1")) != 1: @@ -107,7 +112,6 @@ def _build_operation(args, vae, device): dtype=np.uint8, ) stage = MiniMaxH3LatentPreparationStage( - transformer=transformer, vae=vae, audio_vae=None, scheduler=None, @@ -118,7 +122,7 @@ def run_once(): reference = MiniMaxH3PreparedReference(media_type="video", frames=frames) vae.to(device) try: - return stage._encode_visual_rows([reference], device)[0] + return stage._encode_visual_rows([reference], device, runtime_args)[0] finally: vae.to("cpu") @@ -139,7 +143,7 @@ def run_once(): 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) + stage = MiniMaxH3VideoDecodingStage(vae) def run_once(): batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape) From 770d2dfa284018a64b7388c76981d82e20157abf Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Mon, 31 Aug 2026 12:28:53 -0700 Subject: [PATCH 2/2] [docs]: drop leftover sequential-start wording from lazy-load offloading docs The MLX phase-order sentence and "start a new generator" instruction belong to a different design. A later generate() reloads in-process. --- docs/inference/offloading.md | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/docs/inference/offloading.md b/docs/inference/offloading.md index 3dd2146c09..231282c552 100644 --- a/docs/inference/offloading.md +++ b/docs/inference/offloading.md @@ -28,13 +28,12 @@ runs conditioning, and releases that encoder before it loads the DiT. After denoise it releases the DiT, then loads the video VAE for decode. Input preparation and unpatchify read geometry from checkpoint `config.json` (VAE spatial ratio / latent channels, DiT patch size) so those stages do not -materialize weights just to read two integers. The MLX FastH3 runtime uses the -same phase order. When host offload is off, DiT safetensors are read onto the -accelerator instead of CPU-then-copy. A later `generate()` on the same worker -reloads a released component from disk; start a new generator for a new prompt -until prompt-cache reload exists. The flag is auto-enabled on unified-memory -devices. Pass `--no-lazy-module-load` (or `lazy_module_load=False`) to keep every -component resident. +materialize weights just to read two integers. When host offload is off, DiT +safetensors are read onto the accelerator instead of CPU-then-copy. A later +`generate()` on the same worker reloads a released component from disk in +process; it does not need a new generator. The flag is auto-enabled on +unified-memory devices. Pass `--no-lazy-module-load` (or `lazy_module_load=False`) +to keep every component resident. ## Behavior Explanation