diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index da92c0084..cb5f575fd 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -161,6 +161,7 @@ is power-cycled. To avoid it: on: "CPU" offload uses the same unified RAM. Multi-GPU FSDP sharding remains available because it partitions weights without parking them in a separate host pool. +<<<<<<< Updated upstream - **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 @@ -169,7 +170,25 @@ is power-cycled. To avoid it: `--h3-sequential-load` only if auto-detect misses the device. The CUDA pipeline 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). +======= +- **MiniMax H3 / FastH3** requires lazy component loading on one GB10; otherwise + the text encoder, DiT, and VAEs are loaded together and can trigger an + `earlyoom` kill. The FastH3 example turns this on automatically with + `--num-gpus 1`, but make the requirement explicit in reproducible Spark + commands: + + ```bash + nice -n 19 python examples/inference/basic/basic_fasth3.py \ + --prompt "your prompt" --num-gpus 1 --lazy-module-load \ + --vsa-kernel triton --no-fa4 + ``` + + Lazy loading reads released components from disk again on every generation, + and it only lowers component residency; resolution-dependent activations must + still fit. With host offload disabled, DiT checkpoint tensors load directly + onto the accelerator (`to_cpu` follows `cpu_offload`, which is off here). See +>>>>>>> Stashed changes +[Offloading](../../inference/offloading.md). ## Gotchas specific to the GB10 diff --git a/docs/inference/offloading.md b/docs/inference/offloading.md index bfc46295e..47b685e1b 100644 --- a/docs/inference/offloading.md +++ b/docs/inference/offloading.md @@ -22,6 +22,7 @@ pool there, so offload adds transfers and duplicate residency instead of freeing 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. +<<<<<<< Updated upstream MiniMax H3 CUDA inference can use a second lever that does not copy weights to a host pool: load the Qwen3-VL text encoder, run conditioning, then release that encoder before loading the DiT and video/audio VAEs. `h3_sequential_load` @@ -35,6 +36,8 @@ 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. +======= +>>>>>>> Stashed changes ## Behavior Explanation !!! note @@ -127,21 +130,25 @@ We recommend enabling these options when OOM happens. Every option above moves weights between host and device. This one changes whether they are in memory at all. -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`. +The flag defaults to `False`. By default a pipeline loads every component before +the first stage runs, so peak memory is the sum of all of them even when they are +used in different phases. With `lazy_module_load` enabled, each opted-in heavy +component loads on first use and is freed once the last stage holding it has +returned. Peak component memory then becomes the largest live set required by a +pipeline phase instead of the sum of every component. + +The FastH3 example enables lazy loading by default when `--num-gpus 1` is used +and leaves it off for multi-GPU runs. `--lazy-module-load` and +`--no-lazy-module-load` override that example-specific default. Other entrypoints +keep the API default of `False` unless the option is set explicitly. #### Performance Impact -A freed component is read from disk again on the next generation, so a -multi-prompt run pays one reload per component per request. For a large text -encoder that is tens of seconds. If pipeline-level `torch.compile` is enabled, -the compile setup is reapplied after each reload; PyTorch can reuse its graph -and kernel caches when the component structure and input shapes are unchanged. +A freed component is read from disk again on every later `generate()`, so a +multi-prompt run pays one reload per released component per request. If +pipeline-level `torch.compile` is enabled, the compile setup is reapplied after +each reload; PyTorch can reuse its graph and kernel caches when the component +structure and input shapes are unchanged. #### Usage Recommendation diff --git a/examples/inference/basic/basic_fasth3.py b/examples/inference/basic/basic_fasth3.py index cad9f0788..c9a30f59d 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/examples/inference/basic/basic_fasth3.py @@ -49,11 +49,12 @@ 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. Defaults on for one GPU and off for multi-GPU; each generation reloads " + "released components from disk") parser.add_argument("--profile", choices=("all", "strict"), default="all", @@ -143,6 +144,8 @@ def validate_args(parser: argparse.ArgumentParser, args: argparse.Namespace) -> parser.error("--repeats must be at least 1") if args.num_gpus < 1: parser.error("--num-gpus must be at least 1") + if args.lazy_module_load is None: + args.lazy_module_load = args.num_gpus == 1 if not 0.0 <= args.vsa_sparsity < 1.0: parser.error("--vsa-sparsity must be in [0, 1)") if args.compile_mode is not None and args.inference_torch_compile: diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index f8f5de8dc..87e2e5cd6 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -4,13 +4,10 @@ from __future__ import annotations import gc -from dataclasses import dataclass from typing import Any import torch -from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEArchConfig -from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEArchConfig from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger @@ -34,29 +31,6 @@ logger = init_logger(__name__) _DENOISE_MODULE_NAMES = ("vae", "audio_vae", "transformer") -@dataclass(frozen=True) -class _H3VideoGeometry: - spatial_compression_ratio: int - latent_channels: int - - -@dataclass(frozen=True) -class _H3AudioGeometry: - sampling_rate: int - - -def _default_video_geometry() -> _H3VideoGeometry: - arch = MiniMaxH3VideoVAEArchConfig() - return _H3VideoGeometry( - spatial_compression_ratio=int(arch.spatial_compression_ratio), - latent_channels=int(arch.latent_channels), - ) - - -def _default_audio_geometry() -> _H3AudioGeometry: - return _H3AudioGeometry(sampling_rate=int(MiniMaxH3AudioVAEArchConfig().sampling_rate)) - - class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): """Shared loading and target-generation path for MiniMax H3. @@ -117,6 +91,9 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): raise ValueError(f"MiniMax-H3 {modality} scheduler must expose shift={expected_shift:g}, got {shift}.") def _defer_denoise_modules(self, fastvideo_args: FastVideoArgs) -> bool: + if self._lazy_module_load_enabled(fastvideo_args): + logger.info("lazy_module_load supersedes MiniMax-H3 sequential module loading") + return False if not fastvideo_args.inference_mode or bool(getattr(fastvideo_args, "training_mode", False)): return False requested = fastvideo_args.h3_sequential_load @@ -180,22 +157,10 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): if torch.cuda.is_available(): torch.cuda.empty_cache() - def _input_vae(self) -> Any: - return self.get_module("vae") or _default_video_geometry() - - def _input_audio_vae(self, *, ref2va: bool) -> Any | None: - if not ref2va: - return None - return self.get_module("audio_vae") or _default_audio_geometry() - def _add_condition_stages(self, *, ref2va: bool) -> None: self.add_stage( "input_preparation_stage", - MiniMaxH3InputPreparationStage( - vae=self._input_vae(), - audio_vae=self._input_audio_vae(ref2va=ref2va), - ref2va=ref2va, - ), + MiniMaxH3InputPreparationStage(ref2va=ref2va), ) self.add_stage( "conditioning_stage", @@ -233,7 +198,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): 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)) self._denoise_stages_ready = True 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 7cbc0d6e2..8878def3c 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py @@ -36,6 +36,14 @@ def _layout(batch: ForwardBatch) -> MiniMaxH3PackedLayout: return layout +def _patch_size(batch: ForwardBatch) -> tuple[int, int, int]: + patch_size = batch.minimax_h3_patch_size + if (not isinstance(patch_size, tuple) or len(patch_size) != 3 + or not all(isinstance(value, int) and value > 0 for value in patch_size)): + raise ValueError("MiniMax-H3 patch geometry is missing at decode.") + return patch_size + + def _decode_participation(fastvideo_args: FastVideoArgs, want_parallel: bool) -> tuple[Any, bool, bool]: """Resolve (sp_group, is_output_rank, parallel) for the VAE decode stages. @@ -58,16 +66,21 @@ 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() result.add_check("layout", batch.extra.get(MINIMAX_H3_LAYOUT_KEY), V.not_none) result.add_check("latents", batch.latents, V.with_dims(2)) result.add_check("raw_latent_shape", batch.raw_latent_shape, V.not_none) + result.add_check( + "minimax_h3_patch_size", + batch.minimax_h3_patch_size, + lambda value: isinstance(value, tuple) and len(value) == 3 and all( + isinstance(item, int) and item > 0 for item in value), + ) return result def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: @@ -90,6 +103,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): layout = _layout(batch) if batch.latents is None or batch.raw_latent_shape is None or len(batch.raw_latent_shape) != 5: raise ValueError("MiniMax-H3 video latents or raw geometry are missing at decode.") + patch_size = _patch_size(batch) _, channels, num_frames, latent_height, latent_width = batch.raw_latent_shape latents = unpatchify_video_tokens( batch.latents[layout.num_condition_video_rows:], @@ -97,7 +111,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage): latent_height, latent_width, channels, - self.transformer.patch_size, + patch_size, ) device = get_local_torch_device() self.vae.to(device) @@ -207,6 +221,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage): batch.latents = None batch.audio_latents = None batch.references = None + batch.minimax_h3_patch_size = None for key in tuple(batch.extra): if key.startswith("minimax_h3_"): batch.extra.pop(key) diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py index b5712362b..29eb4f594 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py @@ -3,11 +3,14 @@ from __future__ import annotations -from typing import Any +from typing import cast from PIL import Image, ImageOps import torch +from fastvideo.configs.models.vaes.minimax_h3_audio import MiniMaxH3AudioVAEArchConfig +from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEArchConfig +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.models.vision_utils import load_image from fastvideo.pipelines.basic.minimax_h3.packing import ( @@ -69,7 +72,11 @@ def prepare_common_request(batch: ForwardBatch) -> None: batch.fps = MINIMAX_H3_FPS -def resolve_target_canvas(batch: ForwardBatch, vae: Any, default_aspect: tuple[int, int]) -> tuple[int, int, int]: +def resolve_target_canvas( + batch: ForwardBatch, + spatial_compression_ratio: int, + default_aspect: tuple[int, int], +) -> tuple[int, int, int]: """Resolve target geometry independently of the selected condition mode.""" if (batch.height is None) != (batch.width is None): raise ValueError("MiniMax-H3 `height` and `width` must be passed together, or neither.") @@ -83,7 +90,7 @@ def resolve_target_canvas(batch: ForwardBatch, vae: Any, default_aspect: tuple[i raise ValueError(f"MiniMax-H3 `height` and `width` must be positive multiples of " f"{MINIMAX_H3_CANVAS_MULTIPLE}, got {height}x{width}.") - ratio = int(vae.spatial_compression_ratio) + ratio = int(spatial_compression_ratio) if height % ratio or width % ratio: raise ValueError(f"MiniMax-H3 canvas {height}x{width} is not divisible by VAE ratio {ratio}.") return height, width, ratio @@ -104,12 +111,8 @@ def resolve_target_num_frames(num_frames: object) -> int: class MiniMaxH3InputPreparationStage(PipelineStage): """Prepare FL2VA/T2VA or Ref2VA inputs without a parallel family state object.""" - def __init__(self, vae: Any, audio_vae: Any | None = None, *, ref2va: bool = False) -> None: + def __init__(self, *, ref2va: bool = False) -> None: super().__init__() - if ref2va and audio_vae is None: - raise ValueError("MiniMax-H3 Ref2VA input preparation requires an audio VAE.") - self.vae = vae - self.audio_vae = audio_vae self.ref2va = ref2va def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: @@ -153,11 +156,18 @@ class MiniMaxH3InputPreparationStage(PipelineStage): ) return result - def _write_target_geometry(self, batch: ForwardBatch, height: int, width: int, ratio: int, num_frames: int) -> None: + def _write_target_geometry( + self, + batch: ForwardBatch, + height: int, + width: int, + ratio: int, + num_frames: int, + latent_channels: int, + ) -> None: latent_height = height // ratio latent_width = width // ratio num_latent_frames = video_latent_num_frames(num_frames) - latent_channels = int(self.vae.latent_channels) batch.height = height batch.width = width @@ -166,7 +176,7 @@ class MiniMaxH3InputPreparationStage(PipelineStage): batch.width_latents = latent_width batch.raw_latent_shape = (1, latent_channels, num_latent_frames, latent_height, latent_width) - def _prepare_fl2va(self, batch: ForwardBatch) -> None: + def _prepare_fl2va(self, batch: ForwardBatch, spatial_ratio: int, latent_channels: int) -> None: if batch.references: raise ValueError("MiniMax-H3 references belong to the Ref2VA pipeline.") @@ -180,9 +190,9 @@ class MiniMaxH3InputPreparationStage(PipelineStage): for anchor, image in (("first", batch.pil_image), ("last", batch.last_image)) if image is not None] default_aspect = raw_keyframes[0][1].size if raw_keyframes else (16, 9) - height, width, ratio = resolve_target_canvas(batch, self.vae, default_aspect) + height, width, ratio = resolve_target_canvas(batch, spatial_ratio, default_aspect) num_frames = resolve_target_num_frames(batch.num_frames) - self._write_target_geometry(batch, height, width, ratio, num_frames) + self._write_target_geometry(batch, height, width, ratio, num_frames, latent_channels) batch.extra[MINIMAX_H3_KEYFRAME_ANCHORS_KEY] = tuple(anchor for anchor, _ in raw_keyframes) batch.extra[MINIMAX_H3_KEYFRAMES_KEY] = [ @@ -190,29 +200,41 @@ class MiniMaxH3InputPreparationStage(PipelineStage): for index, (_, image) in enumerate(raw_keyframes) ] - def _prepare_ref2va(self, batch: ForwardBatch) -> None: + def _prepare_ref2va( + self, + batch: ForwardBatch, + spatial_ratio: int, + latent_channels: int, + audio_sampling_rate: int, + ) -> None: if batch.image_path is not None or batch.pil_image is not None or batch.last_image is not None: raise ValueError("MiniMax-H3 Ref2VA accepts media through `references`, not FL2VA keyframe fields.") references = validate_references(list(batch.references or [])) - if self.audio_vae is None: - raise RuntimeError("MiniMax-H3 Ref2VA input preparation has no audio VAE.") - height, width, ratio = resolve_target_canvas(batch, self.vae, (16, 9)) + height, width, ratio = resolve_target_canvas(batch, spatial_ratio, (16, 9)) num_frames = resolve_target_num_frames(batch.num_frames) - target_sample_rate = int(self.audio_vae.sampling_rate) - batch.references = [prepare_reference(reference, num_frames, target_sample_rate) for reference in references] - self._write_target_geometry(batch, height, width, ratio, num_frames) + batch.references = [prepare_reference(reference, num_frames, audio_sampling_rate) for reference in references] + self._write_target_geometry(batch, height, width, ratio, num_frames, latent_channels) batch.extra[MINIMAX_H3_KEYFRAMES_KEY] = [] batch.extra[MINIMAX_H3_KEYFRAME_ANCHORS_KEY] = () @torch.no_grad() def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: - del fastvideo_args + pipeline_config = cast(MiniMaxH3PipelineConfig, fastvideo_args.pipeline_config) + video_arch = cast(MiniMaxH3VideoVAEArchConfig, pipeline_config.vae_config.arch_config) + audio_arch = cast(MiniMaxH3AudioVAEArchConfig, pipeline_config.audio_vae_config.arch_config) + spatial_ratio = int(video_arch.spatial_compression_ratio) + latent_channels = int(video_arch.latent_channels) prepare_common_request(batch) if self.ref2va: - self._prepare_ref2va(batch) + self._prepare_ref2va( + batch, + spatial_ratio, + latent_channels, + int(audio_arch.sampling_rate), + ) else: - self._prepare_fl2va(batch) + self._prepare_fl2va(batch, spatial_ratio, latent_channels) return batch 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 b30e599f7..b6ed410fc 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 @@ -293,6 +293,8 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): @torch.no_grad() def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + patch_t, patch_h, patch_w = self.transformer.patch_size + batch.minimax_h3_patch_size = (int(patch_t), int(patch_h), int(patch_w)) video_noise = batch.latents audio_noise = batch.audio_latents device = get_local_torch_device() diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 3a9ca1467..93ef66443 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -163,6 +163,9 @@ class ForwardBatch: height_latents: list[int] | int | None = None width_latents: list[int] | int | None = None num_frames: list[int] | int = 1 # Default for image models + # MiniMax H3 carries this from latent preparation so video decode does not + # retain the DiT solely to read its patch geometry. + minimax_h3_patch_size: tuple[int, int, int] | None = None # Original dimensions (before VAE scaling) height: list[int] | int | None = None diff --git a/fastvideo/tests/inference/test_basic_fasth3_profile.py b/fastvideo/tests/inference/test_basic_fasth3_profile.py index 5c2a2baea..5dcd6813f 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_profile.py +++ b/fastvideo/tests/inference/test_basic_fasth3_profile.py @@ -44,6 +44,7 @@ def test_default_all_profile_matches_fastest_contract(tmp_path): assert args.warmup_seed == 999 assert args.warmup is True assert args.repeats == 3 + assert args.lazy_module_load is False assert args.inference_torch_compile is True assert args.ulysses_a2a == "off" @@ -55,6 +56,7 @@ def test_default_all_profile_matches_fastest_contract(tmp_path): assert config.engine.offload.text_encoder is True assert config.engine.offload.vae is True assert config.engine.offload.pin_cpu_memory is True + assert config.engine.offload.lazy_module_load is False assert config.engine.compile.enabled is False assert config.engine.compile.vae_enabled is True assert config.pipeline.experimental == { @@ -86,6 +88,18 @@ def test_default_all_profile_matches_fastest_contract(tmp_path): assert request.output.output_path == str(tmp_path / "result.mp4") +@pytest.mark.parametrize(("overrides", "expected"), [ + (("--num-gpus", "1"), True), + (("--num-gpus", "1", "--no-lazy-module-load"), False), + (("--lazy-module-load", ), True), +]) +def test_lazy_module_loading_defaults_on_only_for_one_gpu(overrides, expected): + args = _args(*overrides) + + assert args.lazy_module_load is expected + assert fasth3.build_generator_config(args).engine.offload.lazy_module_load is expected + + @pytest.mark.parametrize("num_frames", (124, 243, 345)) def test_fast_profile_supports_measured_durations_without_separate_scripts(tmp_path, num_frames): args = _args("--num-frames", str(num_frames)) diff --git a/fastvideo/tests/stages/test_lazy_module_load.py b/fastvideo/tests/stages/test_lazy_module_load.py index caa308800..026565518 100644 --- a/fastvideo/tests/stages/test_lazy_module_load.py +++ b/fastvideo/tests/stages/test_lazy_module_load.py @@ -721,6 +721,78 @@ def test_building_the_real_h3_stages_materializes_nothing(): assert len(pipeline._stages) == 6 +def test_real_h3_input_prep_uses_arch_config_without_materializing_vaes(): + from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + loaded: list[str] = [] + + def tracked(name): + return LazyModule(name, lambda: loaded.append(name) or _Component(name)) + + pipeline = MiniMaxH3Pipeline.__new__(MiniMaxH3Pipeline) + pipeline._stages = [] + pipeline._stage_name_mapping = {} + pipeline.modules = { + "text_encoder": tracked("text_encoder"), + "transformer": tracked("transformer"), + "vae": tracked("vae"), + "audio_vae": tracked("audio_vae"), + "tokenizer": object(), + "processor": object(), + "scheduler": object(), + "audio_scheduler": object(), + } + pipeline._add_stages(ref2va=False) + batch = ForwardBatch( + data_type="video", + prompt="A test prompt", + height=192, + width=320, + num_frames=124, + ) + args = SimpleNamespace(pipeline_config=MiniMaxH3PipelineConfig()) + + result = pipeline._stages[0].forward(batch, args) + + assert result is batch + assert loaded == [] + assert batch.height_latents == 12 + assert batch.width_latents == 20 + assert batch.raw_latent_shape is not None + assert batch.raw_latent_shape[1] == 24 + + +def test_real_h3_release_schedule_drops_transformer_before_video_decode(): + from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline + + def tracked(name): + return LazyModule(name, lambda: _Component(name)) + + pipeline = MiniMaxH3Pipeline.__new__(MiniMaxH3Pipeline) + pipeline._stages = [] + pipeline._stage_name_mapping = {} + pipeline.modules = { + "text_encoder": tracked("text_encoder"), + "transformer": tracked("transformer"), + "vae": tracked("vae"), + "audio_vae": tracked("audio_vae"), + "tokenizer": object(), + "processor": object(), + "scheduler": object(), + "audio_scheduler": object(), + } + pipeline._add_stages(ref2va=False) + + assert pipeline._build_lazy_release_schedule() == { + 1: ["text_encoder"], + 3: ["transformer"], + 4: ["vae"], + 5: ["audio_vae"], + } + + class _LoRAConfigComponent(torch.nn.Module): def __init__(self, excluded_layers): diff --git a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py index 4e35bf66f..c910c2ef2 100644 --- a/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py +++ b/fastvideo/tests/stages/test_minimax_h3_vae_streaming.py @@ -73,7 +73,12 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None: latent_shape = (1, 4, 2, 4, 4) latents = torch.randn(latent_shape) rows = patchify_video_latents(latents, (1, 1, 1)) - batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape) + batch = ForwardBatch( + data_type="video", + latents=rows, + raw_latent_shape=latent_shape, + minimax_h3_patch_size=(1, 1, 1), + ) batch.extra[MINIMAX_H3_LAYOUT_KEY] = _layout(rows.shape[0], latent_shape) observed = {} @@ -95,7 +100,7 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None: 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), ) @@ -120,7 +125,7 @@ def test_decode_stages_skip_vae_on_non_output_rank(monkeypatch) -> None: 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)) @@ -173,9 +178,14 @@ def test_parallel_decode_runs_on_every_rank(monkeypatch) -> None: lambda rank=rank, is_first=is_first: SimpleNamespace(is_first_rank=is_first, world_size=4, rank_in_group=rank)) - batch = ForwardBatch(data_type="video", latents=rows.clone(), raw_latent_shape=latent_shape) + batch = ForwardBatch( + data_type="video", + latents=rows.clone(), + raw_latent_shape=latent_shape, + minimax_h3_patch_size=(1, 1, 1), + ) 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 57e68a742..34e10d84c 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 @@ -139,10 +139,15 @@ def _build_operation(args, vae, device): 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) + batch = ForwardBatch( + data_type="video", + latents=rows, + raw_latent_shape=latent_shape, + minimax_h3_patch_size=patch_size, + ) batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout return stage.forward(batch, runtime_args).output