[bugfix]: read H3 geometry from checkpoint JSON so lazy load can drop the DiT before VAE decode

Input prep and unpatchify were holding live VAE/DiT proxies just for two integers, which loaded the video VAE before Qwen and kept the DiT resident through decode. Auto-enable --lazy-module-load on unified memory and on single-GPU FastH3 examples.

Author: aryan5v

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
 

docs/inference/offloading.md

@@ -12,7 +12,7 @@ text_encoder_cpu_offload: bool = True
 image_encoder_cpu_offload: bool = True
 vae_cpu_offload: bool = True
 pin_cpu_memory: bool = True
-lazy_module_load: bool = False
+lazy_module_load: bool | None = None
 ```
 
 On unified-memory accelerators such as NVIDIA GB10 and Apple silicon, FastVideo
@@ -23,14 +23,18 @@ memory. CUDA FSDP sharding remains enabled when requested; MPS continues to
 disable FSDP. `pin_cpu_memory` is not an offload mode and is left unchanged.
 
 MiniMax H3 CUDA inference uses a second lever that does not copy weights to a
-host pool. The pipeline loads the Qwen3-VL text encoder, runs conditioning, then
-releases that encoder before it loads the DiT and video/audio VAEs. The MLX FastH3
-runtime uses the same phase order. When host offload is off, DiT safetensors are
-read onto the accelerator instead of CPU-then-copy. Input-preparation geometry
-(spatial ratio, latent channels, audio sample rate) comes from the VAE arch
-configs until those weights load. A later `generate()` on the same worker
-currently re-enters conditioning after the encoder has been released; start a
-new generator for a new prompt until prompt-cache reload exists.
+host pool. With `lazy_module_load`, the pipeline loads the Qwen3-VL text encoder,
+runs conditioning, and releases that encoder before it loads the DiT. After
+denoise it releases the DiT, then loads the video VAE for decode. Input
+preparation and unpatchify read geometry from checkpoint `config.json` (VAE
+spatial ratio / latent channels, DiT patch size) so those stages do not
+materialize weights just to read two integers. The MLX FastH3 runtime uses the
+same phase order. When host offload is off, DiT safetensors are read onto the
+accelerator instead of CPU-then-copy. A later `generate()` on the same worker
+reloads a released component from disk; start a new generator for a new prompt
+until prompt-cache reload exists. The flag is auto-enabled on unified-memory
+devices. Pass `--no-lazy-module-load` (or `lazy_module_load=False`) to keep every
+component resident.
 
 ## Behavior Explanation
 
@@ -109,9 +113,9 @@ By default a pipeline loads every component before the first stage runs, so
 peak memory is the sum of all of them even though no two are needed at the same
 moment. With `lazy_module_load` enabled, each heavy component loads on first use
 and is freed once the last stage that needs it has returned, so peak memory
-becomes the largest overlapping set instead of the sum. For a text-to-video
-pipeline that is roughly `max(text encoder, DiT + VAE)` rather than
-`text encoder + DiT + VAE`.
+becomes the largest overlapping set instead of the sum. MiniMax-H3 T2VA is
+`max(text encoder, DiT, VAE)` rather than `text encoder + DiT + VAE`, because
+the DiT is not held through VAE decode.
 
 #### Performance Impact
 
@@ -126,8 +130,10 @@ and kernel caches when the component structure and input shapes are unchanged.
 Enable this when a model does not fit at load time, which the CPU offload
 options above cannot help with because they act after loading. It is
 particularly relevant on unified-memory devices, where host and device draw on
-the same pool and moving weights to the host frees nothing. Leave it off when
-the model already fits.
+the same pool and moving weights to the host frees nothing. FastVideo
+auto-enables it there (`lazy_module_load=None`). Leave it off when the model
+already fits, or pass `--no-lazy-module-load` to keep components resident for
+later `generate()` calls.
 
 This option applies to inference only. Training keeps every component resident
 and logs a warning if the flag is set.

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,

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,

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

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
 

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,17 +76,18 @@ 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
     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,16 +141,15 @@ 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))
 
 
 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):

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

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)