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..231282c552 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,17 @@ 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. 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 @@ -109,9 +112,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 +129,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)