From f82d8be4bf90254243ea0d1e2a6dca8e8bac4ebb Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Mon, 31 Aug 2026 04:20:21 -0700 Subject: [PATCH] [perf] Sequential MiniMax H3 start with GPU-direct DiT load (#1793) Co-authored-by: Aryan Kumar --- .../inference_schema_parity_inventory.yaml | 1 + .../installation/spark_performance.md | 14 ++ docs/inference/offloading.md | 32 +++ examples/inference/basic/README.md | 2 +- examples/inference/basic/basic_fasth3.py | 7 + fastvideo/fastvideo_args.py | 15 ++ fastvideo/models/loader/fsdp_load.py | 6 +- .../basic/minimax_h3/minimax_h3_pipeline.py | 165 ++++++++++++++- .../inference/test_basic_fasth3_profile.py | 11 + .../test_minimax_h3_sequential_start.py | 199 ++++++++++++++++++ 10 files changed, 442 insertions(+), 10 deletions(-) create mode 100644 fastvideo/tests/stages/test_minimax_h3_sequential_start.py diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 1f1f10352..196b305d5 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -80,6 +80,7 @@ surfaces: VSA_tile_size: "VSA-H3 tile geometry request; model-specific optimization not yet represented in the typed public schema." inference_torch_compile: "Regional inference compile opt-in currently carried through PipelineSelection.experimental rather than CompileConfig." vae_parallel_decode: "MiniMax-H3 sequence-parallel VAE decode opt-in; model-specific optimization not yet represented in the typed public schema." + h3_sequential_load: "MiniMax-H3 sequential text-encoder then DiT/VAE load; model-specific optimization not yet represented in the typed public schema." vae_parallel_encode: "MiniMax-H3 sequence-parallel reference VAE encode opt-in; model-specific optimization not yet represented in the typed public schema." vae_parallel_decode_strategy: "Chunk-transport collective for vae_parallel_decode; model-specific optimization not yet represented in the typed public schema." attention_backend: "Process-wide default attention-backend request applied per component at load time; kernel-selection knob not yet represented in the typed public schema." diff --git a/docs/getting_started/installation/spark_performance.md b/docs/getting_started/installation/spark_performance.md index a2a3a1452..da92c0084 100644 --- a/docs/getting_started/installation/spark_performance.md +++ b/docs/getting_started/installation/spark_performance.md @@ -161,6 +161,15 @@ 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. +- **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). `h3_sequential_load` defaults to auto and turns this split on for + unified-memory devices. Do not pass `--no-h3-sequential-load` here. Force + `--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). ## Gotchas specific to the GB10 @@ -177,6 +186,11 @@ A few things that surprise people on this box (beyond the memory notes above): - **Cosmos-2.5** uses a Qwen2.5-VL text encoder; make sure you're on a FastVideo build recent enough to include its `transformers`-compatibility handling before running it. +- **MiniMax H3 worker init can look healthy and still die on the first generate** + if sequential load is off (`--no-h3-sequential-load`, or auto-off on a + misclassified device) and encoder, VAE, and DiT load together. Confirm the log + contains `Released MiniMax-H3 text encoder after conditioning` before + `Loading MiniMax-H3 denoise modules`. ## Reproduce these numbers diff --git a/docs/inference/offloading.md b/docs/inference/offloading.md index 61773f710..08ad0f6aa 100644 --- a/docs/inference/offloading.md +++ b/docs/inference/offloading.md @@ -21,6 +21,19 @@ 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. +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` +defaults to auto (`None`): on for unified-memory devices such as GB10, off on +discrete GPUs. Pass `--h3-sequential-load` to force it, or +`--no-h3-sequential-load` to keep the encoder resident for later `generate()` +calls on the same worker. Sequential load currently cannot re-encode a new prompt +on that worker; start a new generator until prompt-cache reload exists. The MLX +FastH3 runtime always uses this 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. + ## Behavior Explanation !!! note @@ -65,6 +78,25 @@ This option introduces negligible performance overhead. We recommend enabling this option for single GPU usage. This option is not compatible with FSDP. +### `h3_sequential_load` + +MiniMax H3 only. When enabled, the pipeline loads Qwen3-VL, runs conditioning, +releases that encoder, then loads the DiT and VAEs. Default is auto: enabled on +unified-memory accelerators, disabled on discrete GPUs. + +#### Performance Impact + +On GB10 / Spark this is required so encoder, DiT, and VAE weights are not all +resident in one unified pool. On discrete GPUs (for example 4×GB200) sequential +load is unnecessary and it blocks a second `generate()` on the same worker +because the encoder has been released. + +#### Usage Recommendation + +Leave the default on Spark / DGX Spark. Force `--h3-sequential-load` only when +you need the split on a discrete GPU. Use `--no-h3-sequential-load` when you +need more than one prompt per worker and have enough memory to keep the encoder. + ### `text_encoder_cpu_offload` When enabled, the text encoder model weights are offloaded to CPU memory, and text encoding is computed on CPU. diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md index 0bad54cbe..52e1e0711 100644 --- a/examples/inference/basic/README.md +++ b/examples/inference/basic/README.md @@ -118,7 +118,7 @@ python examples/inference/basic/basic_fasth3.py \ --warmup-seed 999 ``` -`all` enables the inference-only H3 fusions and regional compile. Both can change floating-point operation order, so this is a report-only performance profile rather than an exact-parity route. Use `--profile strict` to disable the H3 fusions while preserving regional compile, or `--profile strict --no-inference-torch-compile` for the eager strict route. Individual `--no-*` switches are available for portability and attribution; in particular, use `--vsa-kernel triton --no-fa4` if the Blackwell kernels are unavailable. The script preserves the warmup and each measured video under distinct paths, then prints per-request wall time plus a warmup-excluded median. +`all` enables the inference-only H3 fusions and regional compile. Both can change floating-point operation order, so this is a report-only performance profile rather than an exact-parity route. Use `--profile strict` to disable the H3 fusions while preserving regional compile, or `--profile strict --no-inference-torch-compile` for the eager strict route. Individual `--no-*` switches are available for portability and attribution; in particular, use `--vsa-kernel triton --no-fa4` if the Blackwell kernels are unavailable. `--h3-sequential-load` / `--no-h3-sequential-load` override the auto split that releases Qwen3-VL before DiT/VAE load (on by default on GB10, off on discrete GPUs). The script preserves the warmup and each measured video under distinct paths, then prints per-request wall time plus a warmup-excluded median. One script covers each validated duration; regional compile is the fastest measured DiT route for all three: diff --git a/examples/inference/basic/basic_fasth3.py b/examples/inference/basic/basic_fasth3.py index 696cab526..2e240c6fa 100644 --- a/examples/inference/basic/basic_fasth3.py +++ b/examples/inference/basic/basic_fasth3.py @@ -99,6 +99,11 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser: action=argparse.BooleanOptionalAction, default=True, help="round-robin VAE temporal chunks across sequence-parallel ranks") + parser.add_argument("--h3-sequential-load", + action=argparse.BooleanOptionalAction, + default=None, + help="encode with Qwen3-VL, release it, then load DiT/VAEs. Default auto: on for " + "unified-memory devices (GB10), off on discrete GPUs") parser.add_argument("--replicated-dit", action=argparse.BooleanOptionalAction, default=True, @@ -229,6 +234,8 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig: "vae_parallel_decode": args.parallel_vae, "vae_parallel_decode_strategy": "gather", } + if args.h3_sequential_load is not None: + experimental["h3_sequential_load"] = args.h3_sequential_load if use_vsa: experimental.update({ "VSA_sparsity": args.vsa_sparsity, diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 610bef3c5..ff812437f 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -159,6 +159,13 @@ class FastVideoArgs: vae_cpu_offload: bool = True pin_cpu_memory: bool = True + # MiniMax-H3 inference load order. ``None`` (auto) defers DiT/VAE load until + # after the Qwen3-VL encoder is released, but only on unified-memory + # devices (GB10 / Spark). Discrete GPUs keep the encoder resident so a + # later ``generate()`` on the same worker can re-encode. Explicit True / + # False overrides the probe. Training never defers. + h3_sequential_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 # round-robined across the sequence-parallel ranks and reassembled @@ -714,6 +721,14 @@ class FastVideoArgs: "Pin memory for CPU offload. Only added as a temp workaround if it throws \"CUDA error: invalid argument\". " "Should be enabled in almost all cases", ) + parser.add_argument( + "--h3-sequential-load", + action=argparse.BooleanOptionalAction, + default=None, + help="MiniMax-H3: encode with Qwen3-VL, release that encoder, then load DiT and VAEs. " + "Omit for auto (on for unified-memory devices such as GB10; off on discrete GPUs). " + "Pass --no-h3-sequential-load to keep the encoder resident for later generate() calls.", + ) parser.add_argument( "--vae-parallel-decode", action=StoreBoolean, diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 05d8222ca..2028d1b26 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -239,7 +239,11 @@ def maybe_load_fsdp_model( fsdp_shard_conditions=model._fsdp_shard_conditions, pin_cpu_memory=pin_cpu_memory) - weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True) + # Host offload is already disabled on unified memory (GB10). Staging the + # 35B FastH3 DiT on CPU and then copying to CUDA doubled that working set + # and took minutes. Follow cpu_offload: read onto the accelerator. + weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=cpu_offload) + logger.info("Loading transformer weights with to_cpu=%s", cpu_offload) param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping) dense_lora_patch = DenseLoRAPatch.from_adapter( lora_path, diff --git a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py index a45c57fa1..e1dceab49 100644 --- a/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py +++ b/fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py @@ -3,8 +3,17 @@ 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 from fastvideo.pipelines.basic.minimax_h3.stages import ( MiniMaxH3AudioDecodingStage, MiniMaxH3ConditioningStage, @@ -15,6 +24,37 @@ from fastvideo.pipelines.basic.minimax_h3.stages import ( ) from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +logger = init_logger(__name__) + +# Same split as the MLX runtime: condition, release the ~66 GB Qwen3-VL stack, +# then load DiT + VAEs. Keeping them resident together OOMs unified-memory +# boxes (GB10 / Spark) even though host offload is correctly disabled there. +_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): @@ -53,6 +93,11 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): "audio_scheduler", ] + def __init__(self, *args: Any, **kwargs: Any) -> None: + self._ref2va = False + self._denoise_stages_ready = False + super().__init__(*args, **kwargs) + @classmethod 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)) @@ -67,18 +112,84 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): 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: - 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") + def _defer_denoise_modules(self, fastvideo_args: FastVideoArgs) -> bool: + if not fastvideo_args.inference_mode or bool(getattr(fastvideo_args, "training_mode", False)): + return False + requested = fastvideo_args.h3_sequential_load + if requested is True: + return True + if requested is False: + return False + from fastvideo.pipelines import composed_pipeline_base + from fastvideo.platforms import current_platform + device = composed_pipeline_base.get_local_torch_device() + device_id = 0 if device.index is None else int(device.index) + unified = bool(current_platform.has_unified_memory(device_id)) + logger.info("MiniMax-H3 sequential module load auto=%s (unified_memory=%s)", unified, unified) + return unified + + def _denoise_modules_loaded(self) -> bool: + return all(self.get_module(name) is not None for name in _DENOISE_MODULE_NAMES) + + def load_modules(self, + fastvideo_args: FastVideoArgs, + loaded_modules: dict[str, torch.nn.Module] | None = None) -> dict[str, Any]: + """Load the Qwen3-VL conditioner first; defer DiT and VAEs until after encode.""" + if not self._defer_denoise_modules(fastvideo_args): + return super().load_modules(fastvideo_args, loaded_modules) + if loaded_modules is not None and all(name in loaded_modules for name in _DENOISE_MODULE_NAMES): + return super().load_modules(fastvideo_args, loaded_modules) + + saved = list(self.required_config_modules) + self._required_config_modules = [name for name in saved if name not in _DENOISE_MODULE_NAMES] + try: + logger.info("Loading MiniMax-H3 condition modules first: %s", self._required_config_modules) + return super().load_modules(fastvideo_args, loaded_modules) + finally: + self._required_config_modules = saved + + def _load_denoise_modules(self, fastvideo_args: FastVideoArgs) -> None: + if self._denoise_modules_loaded(): + return + saved = list(self.required_config_modules) + self._required_config_modules = [name for name in saved if name != "text_encoder"] + try: + logger.info("Loading MiniMax-H3 denoise modules after releasing the text encoder: %s", + [name for name in self._required_config_modules if name in _DENOISE_MODULE_NAMES]) + loaded = super().load_modules(fastvideo_args, loaded_modules=self.modules) + for name, module in loaded.items(): + self.add_module(name, module) + finally: + self._required_config_modules = saved + + def _release_text_encoder(self) -> None: + stage = self._stage_name_mapping.get("conditioning_stage") + if stage is not None: + stage.conditioner = None + encoder = self.modules.pop("text_encoder", None) + if encoder is None: + return + logger.info("Released MiniMax-H3 text encoder after conditioning") + del encoder + gc.collect() + 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=vae, - audio_vae=audio_vae if ref2va else None, + vae=self._input_vae(), + audio_vae=self._input_audio_vae(ref2va=ref2va), ref2va=ref2va, ), ) @@ -91,6 +202,15 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): ref2va=ref2va, ), ) + + def _add_denoise_stages(self, *, 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") + if transformer is None or vae is None or audio_vae is None: + raise RuntimeError("MiniMax-H3 denoise stages require transformer, vae, and audio_vae to be loaded.") self.add_stage( "latent_preparation_stage", MiniMaxH3LatentPreparationStage( @@ -111,6 +231,35 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase): ) self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=vae, transformer=transformer)) self.add_stage("audio_decoding_stage", MiniMaxH3AudioDecodingStage(audio_vae=audio_vae)) + self._denoise_stages_ready = True + + def _add_stages(self, *, ref2va: bool) -> None: + self._ref2va = ref2va + self._add_condition_stages(ref2va=ref2va) + if self._denoise_modules_loaded(): + self._add_denoise_stages(ref2va=ref2va) + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if not self.post_init_called: + self.post_init() + + if self._denoise_stages_ready: + return super().forward(batch, fastvideo_args) + + logger.info("Running MiniMax-H3 condition stages before loading DiT/VAE weights") + for stage in self.stages: + batch = stage(batch, fastvideo_args) + self._release_text_encoder() + self._load_denoise_modules(fastvideo_args) + self._add_denoise_stages(ref2va=self._ref2va) + for name in ( + "latent_preparation_stage", + "denoising_stage", + "video_decoding_stage", + "audio_decoding_stage", + ): + batch = self._stage_name_mapping[name](batch, fastvideo_args) + return batch class MiniMaxH3Pipeline(MiniMaxH3BasePipeline): diff --git a/fastvideo/tests/inference/test_basic_fasth3_profile.py b/fastvideo/tests/inference/test_basic_fasth3_profile.py index e4f8bc250..5c2a2baea 100644 --- a/fastvideo/tests/inference/test_basic_fasth3_profile.py +++ b/fastvideo/tests/inference/test_basic_fasth3_profile.py @@ -129,6 +129,17 @@ def test_strict_profile_changes_only_non_parity_fusions(): assert fasth3.build_generator_config(all_args) == fasth3.build_generator_config(strict_args) +def test_h3_sequential_load_is_opt_in_experimental(): + default = fasth3.build_generator_config(_args()) + assert "h3_sequential_load" not in default.pipeline.experimental + + enabled = fasth3.build_generator_config(_args("--h3-sequential-load")) + assert enabled.pipeline.experimental["h3_sequential_load"] is True + + disabled = fasth3.build_generator_config(_args("--no-h3-sequential-load")) + assert disabled.pipeline.experimental["h3_sequential_load"] is False + + def test_opt_outs_override_inherited_environment(monkeypatch): args = _args( "--vsa-kernel", diff --git a/fastvideo/tests/stages/test_minimax_h3_sequential_start.py b/fastvideo/tests/stages/test_minimax_h3_sequential_start.py new file mode 100644 index 000000000..6ddfc2d91 --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_sequential_start.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU contracts for MiniMax-H3 Mac-style sequential module loading.""" +from __future__ import annotations + +from contextlib import nullcontext +from types import SimpleNamespace + +import torch + +import fastvideo.pipelines.composed_pipeline_base as composed_pipeline_base +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.utils import FlexibleArgumentParser +from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import ( + MiniMaxH3Pipeline, + _DENOISE_MODULE_NAMES, +) +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + +class _Profiler: + + def region(self, name): + del name + return nullcontext() + + +def _stub_module(name: str) -> SimpleNamespace: + if name in {"scheduler"}: + return SimpleNamespace(shift=12.0, name=name) + if name in {"audio_scheduler"}: + return SimpleNamespace(shift=3.0, name=name) + if name == "transformer": + # LoRAPipeline reads exclude_lora_layers off the DiT arch config. + return SimpleNamespace( + name=name, + config=SimpleNamespace(arch_config=SimpleNamespace(exclude_lora_layers=[])), + ) + return SimpleNamespace(name=name) + + +def _patch_pipeline_construction(monkeypatch, events: list) -> None: + monkeypatch.setattr( + composed_pipeline_base, + "maybe_init_distributed_environment_and_model_parallel", + lambda *args, **kwargs: events.append(("distributed", None)), + ) + monkeypatch.setattr(composed_pipeline_base, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(composed_pipeline_base, "get_world_group", lambda: SimpleNamespace(local_rank=0)) + monkeypatch.setattr(composed_pipeline_base, "get_or_create_profiler", lambda trace_dir: _Profiler()) + monkeypatch.setattr(composed_pipeline_base, "warmup_sequence_parallel_communication", lambda: None) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: False) + monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False) + + +def test_inference_defers_dit_and_vae_until_after_conditioning(monkeypatch) -> None: + events: list = [] + _patch_pipeline_construction(monkeypatch, events) + loads: list[list[str]] = [] + + def fake_load(self, fastvideo_args, loaded_modules=None): + del fastvideo_args + requested = list(self.required_config_modules) + loads.append(requested) + modules = dict(loaded_modules or {}) + for name in requested: + modules.setdefault(name, _stub_module(name)) + return modules + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_load) + + args = FastVideoArgs( + model_path="unused/for-this-test", + enable_stage_verification=False, + h3_sequential_load=True, + ) + pipeline = MiniMaxH3Pipeline("unused/for-this-test", args) + pipeline.post_init() + + assert loads, "condition modules should load during construction" + assert "text_encoder" in loads[0] + assert all(name not in loads[0] for name in _DENOISE_MODULE_NAMES) + assert pipeline.get_module("text_encoder") is not None + assert pipeline.get_module("transformer") is None + assert list(pipeline._stage_name_mapping) == ["input_preparation_stage", "conditioning_stage"] + + condition_stage = pipeline._stage_name_mapping["conditioning_stage"] + passthrough = lambda batch, _args: batch + monkeypatch.setattr(pipeline._stage_name_mapping["input_preparation_stage"], "forward", passthrough) + monkeypatch.setattr(condition_stage, "forward", passthrough) + + original_add_denoise = pipeline._add_denoise_stages + + def fake_add_denoise(*, ref2va: bool) -> None: + original_add_denoise(ref2va=ref2va) + for name in ( + "latent_preparation_stage", + "denoising_stage", + "video_decoding_stage", + "audio_decoding_stage", + ): + monkeypatch.setattr(pipeline._stage_name_mapping[name], "forward", passthrough) + + monkeypatch.setattr(pipeline, "_add_denoise_stages", fake_add_denoise) + + batch = ForwardBatch(data_type="video", prompt="alpine dancer") + out = pipeline.forward(batch, args) + + assert out is batch + assert len(loads) == 2 + assert "transformer" in loads[1] + assert "vae" in loads[1] + assert "text_encoder" not in loads[1] + assert pipeline.get_module("text_encoder") is None + assert condition_stage.conditioner is None + assert pipeline.get_module("transformer") is not None + assert pipeline._denoise_stages_ready is True + + +def test_injected_denoise_weights_skip_the_deferred_split(monkeypatch) -> None: + events: list = [] + _patch_pipeline_construction(monkeypatch, events) + loads: list[list[str]] = [] + + def fake_load(self, fastvideo_args, loaded_modules=None): + del fastvideo_args + loads.append(list(self.required_config_modules)) + return dict(loaded_modules or {}) + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_load) + injected = {name: _stub_module(name) for name in MiniMaxH3Pipeline._required_config_modules} + args = FastVideoArgs(model_path="unused/for-this-test", h3_sequential_load=True) + MiniMaxH3Pipeline("unused/for-this-test", args, loaded_modules=injected) + + assert loads == [list(MiniMaxH3Pipeline._required_config_modules)] + + +def test_explicit_false_loads_encoder_dit_and_vae_together(monkeypatch) -> None: + events: list = [] + _patch_pipeline_construction(monkeypatch, events) + loads: list[list[str]] = [] + + def fake_load(self, fastvideo_args, loaded_modules=None): + del fastvideo_args, loaded_modules + loads.append(list(self.required_config_modules)) + return {name: _stub_module(name) for name in self.required_config_modules} + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_load) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: True) + args = FastVideoArgs(model_path="unused/for-this-test", h3_sequential_load=False) + MiniMaxH3Pipeline("unused/for-this-test", args) + + assert loads == [list(MiniMaxH3Pipeline._required_config_modules)] + assert "text_encoder" in loads[0] + assert all(name in loads[0] for name in _DENOISE_MODULE_NAMES) + + +def test_auto_defers_on_unified_memory(monkeypatch) -> None: + events: list = [] + _patch_pipeline_construction(monkeypatch, events) + loads: list[list[str]] = [] + + def fake_load(self, fastvideo_args, loaded_modules=None): + del fastvideo_args, loaded_modules + loads.append(list(self.required_config_modules)) + return {name: _stub_module(name) for name in self.required_config_modules} + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_load) + monkeypatch.setattr("fastvideo.platforms.current_platform.has_unified_memory", lambda device_id: True) + args = FastVideoArgs(model_path="unused/for-this-test") + MiniMaxH3Pipeline("unused/for-this-test", args) + + assert loads + assert "text_encoder" in loads[0] + assert all(name not in loads[0] for name in _DENOISE_MODULE_NAMES) + + +def test_auto_loads_together_without_unified_memory(monkeypatch) -> None: + events: list = [] + _patch_pipeline_construction(monkeypatch, events) + loads: list[list[str]] = [] + + def fake_load(self, fastvideo_args, loaded_modules=None): + del fastvideo_args, loaded_modules + loads.append(list(self.required_config_modules)) + return {name: _stub_module(name) for name in self.required_config_modules} + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_load) + args = FastVideoArgs(model_path="unused/for-this-test") + MiniMaxH3Pipeline("unused/for-this-test", args) + + assert loads == [list(MiniMaxH3Pipeline._required_config_modules)] + + +def test_cli_tri_state_h3_sequential_load() -> None: + parser = FastVideoArgs.add_cli_args(FlexibleArgumentParser()) + assert parser.parse_args([]).h3_sequential_load is None + assert parser.parse_args(["--h3-sequential-load"]).h3_sequential_load is True + assert parser.parse_args(["--no-h3-sequential-load"]).h3_sequential_load is False