[bugfix]: fix lazy-load, LoRA, and Ray NIC issues from the #1803 review

Keep deferral, compile, and later generate() from fighting each other, stop LoRA bookkeeping from pinning a released DiT, and leave per-node NCCL/Gloo interface names alone.

Author: aryan5v

docs/getting_started/installation/spark_pair.md

@@ -191,7 +191,7 @@ sm_100a VSA kernel is not on this chip, so denoise is slower than a GB200
 | `RayDistributedExecutor` TypeError / abstract `set_log_queue` | Use a FastVideo build that implements those methods on the Ray executor (this page). |
 | Worker SIGTERM during DiT shard 11/14 | `RAY_memory_monitor_refresh_ms=0` **before** `ray start`. Do not leave Ray's default 30% object store. |
 | NCCL hangs or uses Wi-Fi | `source spark_pair_env.sh`. Confirm `NCCL_SOCKET_IFNAME` is the QSFP NIC. |
-| Gloo `connectFullMesh` / `remote=[127.0.0.1]` | Two 1-GPU nodes must not use loopback as the Gloo store. Source `spark_pair_env.sh` so `GLOO_SOCKET_IFNAME` is the QSFP NIC. Use a FastVideo build that keys loopback on unique worker IPs. |
+| Gloo `connectFullMesh` / `remote=[127.0.0.1]` | Two 1-GPU nodes must not use loopback as the Gloo store. Source `spark_pair_env.sh` so `GLOO_SOCKET_IFNAME` is the QSFP NIC on **each** box. FastVideo no longer copies that NIC name from the driver onto workers. |
 | Second `generate()` crashes `NoneType.parameters` | Sequential load used to drop the text encoder without reloading it. This branch reloads Qwen for later requests so `--warmup --repeats N` works. |
 | OOM / `earlyoom` prefers Python | Lazy module load must stay on (do not pass `--no-lazy-module-load`). Peak GPU during 345-frame denoise is ~90 GiB/node. |
 | `num_gpus=2` on one Spark | Each Spark has one GPU. Use Ray across two nodes, or `num_gpus=1` on one box. |

docs/inference/offloading.md

@@ -173,7 +173,10 @@ when it enables block-sparse attention, and reading a component's attributes
 while stages are built, as the shared denoising stage does to pick an attention
 backend. A pipeline therefore lists the components it has checked in
 `_lazy_module_names`, which is empty in the base class. MiniMax-H3 opts in. On
-a pipeline that has not, the flag logs a warning and changes nothing.
+a pipeline that has not, the flag is a no-op: hooks are not installed and no
+warning is logged. Sequential MiniMax-H3 (`h3_sequential_load`) reloads the
+text encoder for a later `generate()` on the same worker; you do not need to
+start a new generator.
 
 ## General Recommendations
 

examples/inference/basic/basic_fasth3.py

@@ -53,8 +53,8 @@ def build_parser(description: str | None = None) -> argparse.ArgumentParser:
                         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. 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. Omit for auto (on for unified-memory devices such as GB10; off on discrete "
+                        "GPUs). Costs a reload per generation; pass --no-lazy-module-load to keep every "
                         "component resident")
     parser.add_argument("--profile",
                         choices=("all", "strict"),
@@ -282,8 +282,7 @@ def build_generator_config(args: argparse.Namespace) -> GeneratorConfig:
                 text_encoder=True,
                 vae=True,
                 pin_cpu_memory=args.pin_cpu_memory,
-                lazy_module_load=(True if args.lazy_module_load is None and args.num_gpus == 1 else
-                                  args.lazy_module_load),
+                lazy_module_load=args.lazy_module_load,
             ),
             compile=CompileConfig(
                 enabled=args.torch_compile,

examples/inference/basic/basic_minimax_h3_t2v.py

@@ -66,8 +66,8 @@ def parse_args() -> argparse.Namespace:
                         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. 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. Omit for auto (on for unified-memory devices such as GB10; off on discrete "
+                        "GPUs). Costs a reload per generation; pass --no-lazy-module-load to keep every "
                         "component resident")
     parser.add_argument("--repeats",
                         type=int,
@@ -104,8 +104,7 @@ def main() -> None:
                     text_encoder=True,
                     vae=True,
                     pin_cpu_memory=False,
-                    lazy_module_load=(True if args.lazy_module_load is None and args.num_gpus == 1 else
-                                      args.lazy_module_load),
+                    lazy_module_load=args.lazy_module_load,
                 ),
                 compile=CompileConfig(
                     enabled=args.torch_compile,

fastvideo/attention/backends/video_sparse_attn_h3.py

@@ -452,7 +452,6 @@ def __init__(
         # request-time env/probe/fallback behavior; only Dynamo capture reads
         # the prepared, static route.
         self._regional_compile_sm100a_enabled: bool | None = None
-        self._regional_compile_layer_idx: torch.Tensor | None = None
 
     def prepare_for_compile(self, device: torch.device) -> None:
         """Tensorize per-layer state shared by every torch.compile route."""
@@ -468,7 +467,8 @@ def prepare_for_regional_compile(self, device: torch.device) -> str | None:
         the loaded model's device now, then let ``forward`` specialize on the
         resulting plain bool while Dynamo is compiling.
         """
-        self.prepare_for_compile(device)
+        if self._compile_layer_idx is None:
+            self.prepare_for_compile(device)
         requested = os.environ.get(VSA_SM100A_ENV, "0") == "1"
         enabled = False
         reason = None if requested else f"{VSA_SM100A_ENV}=1 is required for compile-safe VSA-H3 attention"
@@ -495,10 +495,6 @@ def prepare_for_regional_compile(self, device: torch.device) -> str | None:
                 enabled = reason is None
 
         self._regional_compile_sm100a_enabled = enabled
-        # Keep this marker unset when preparation fails. Generic/training
-        # torch.compile must retain the established Triton attention route.
-        self._regional_compile_layer_idx = (torch.tensor(self.layer_idx, device=device, dtype=torch.int64)
-                                            if enabled else None)
         if enabled:
             route = ("native fastvideo-kernel mask entry" if callable(
                 getattr(_sm100a, "block_sparse_attn_sm100a_from_mask", None)) else
@@ -524,7 +520,7 @@ def tile(self, x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Te
         n_tiles = attn_metadata.variable_block_sizes.numel()
         grad_mode = torch.is_grad_enabled() and x.requires_grad
         compiling = torch.compiler.is_compiling()
-        regional_compiling = compiling and self._regional_compile_layer_idx is not None
+        regional_compiling = compiling and self._regional_compile_sm100a_enabled is True
         if regional_compiling:
             sm100a_requested = bool(self._regional_compile_sm100a_enabled)
         elif compiling:
@@ -570,7 +566,7 @@ def forward(  # type: ignore[override]
         attn_metadata: MiniMaxH3VSAMetadata,
     ) -> torch.Tensor:
         compiling = torch.compiler.is_compiling()
-        regional_compiling = compiling and self._regional_compile_layer_idx is not None
+        regional_compiling = compiling and self._regional_compile_sm100a_enabled is True
 
         tile_elems = attn_metadata.tile_elems
         if regional_compiling and tile_elems != 64:

fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py

@@ -63,18 +63,26 @@ def _apply_h3_checkpoint_arch_configs(model_path: str, fastvideo_args: FastVideo
                                       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"
+    vae_dir = root / extra_config_module_map.get("vae", "vae")
     if (vae_dir / "config.json").is_file():
         fastvideo_args.pipeline_config.vae_config.update_model_arch(get_diffusers_config(str(vae_dir)))
+    audio_vae_dir = root / extra_config_module_map.get("audio_vae", "audio_vae")
+    audio_vae_config = getattr(fastvideo_args.pipeline_config, "audio_vae_config", None)
+    if audio_vae_config is not None and (audio_vae_dir / "config.json").is_file():
+        audio_vae_config.update_model_arch(get_diffusers_config(str(audio_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),
-    )
+    dit_config = fastvideo_args.pipeline_config.dit_config
+    vae_arch = getattr(fastvideo_args.pipeline_config.vae_config, "arch_config", None)
+    patch_size = getattr(dit_config, "patch_size", None)
+    if patch_size is not None and vae_arch is not None:
+        logger.info(
+            "MiniMax-H3 geometry from config: patch_size=%s spatial_compression_ratio=%s latent_channels=%s",
+            tuple(patch_size),
+            int(getattr(vae_arch, "spatial_compression_ratio", 0)),
+            int(getattr(vae_arch, "latent_channels", 0)),
+        )
 
 
 class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
@@ -195,6 +203,7 @@ def _load_denoise_modules(self, fastvideo_args: FastVideoArgs) -> None:
             loaded = super().load_modules(fastvideo_args, loaded_modules=self.modules)
             for name, module in loaded.items():
                 self.add_module(name, module)
+            self._apply_inference_compile(tuple(name for name in loaded if name in _DENOISE_MODULE_NAMES))
         finally:
             self._required_config_modules = saved
 
@@ -226,6 +235,7 @@ def _ensure_text_encoder(self, fastvideo_args: FastVideoArgs) -> None:
             loaded = super().load_modules(fastvideo_args, loaded_modules=self.modules)
             for name, module in loaded.items():
                 self.add_module(name, module)
+            self._apply_inference_compile(("text_encoder", ))
         finally:
             self._required_config_modules = saved
         if stage is not None:
@@ -254,17 +264,20 @@ def _input_video_geometry(self, fastvideo_args: FastVideoArgs) -> Any:
             return arch
         return _default_video_geometry()
 
-    def _input_audio_vae(self, *, ref2va: bool) -> Any | None:
+    def _input_audio_vae(self, fastvideo_args: FastVideoArgs, *, ref2va: bool) -> Any | None:
         if not ref2va:
             return None
+        arch = getattr(getattr(fastvideo_args.pipeline_config, "audio_vae_config", None), "arch_config", None)
+        if arch is not None:
+            return arch
         return _default_audio_geometry()
 
     def _add_condition_stages(self, fastvideo_args: FastVideoArgs, *, ref2va: bool) -> None:
         self.add_stage(
             "input_preparation_stage",
             MiniMaxH3InputPreparationStage(
                 vae=self._input_video_geometry(fastvideo_args),
-                audio_vae=self._input_audio_vae(ref2va=ref2va),
+                audio_vae=self._input_audio_vae(fastvideo_args, ref2va=ref2va),
                 ref2va=ref2va,
             ),
         )
@@ -317,12 +330,23 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
         if not self.post_init_called:
             self.post_init()
 
-        self._ensure_text_encoder(fastvideo_args)
-        if self._denoise_stages_ready:
-            logger.info("Running MiniMax-H3 condition stages before denoise (subsequent request)")
-        else:
-            logger.info("Running MiniMax-H3 condition stages before loading DiT/VAE weights")
-        return self._run_condition_then_denoise(batch, fastvideo_args)
+        # Sequential encode-then-release is the H3-only fallback. Lazy and the
+        # fully-resident discrete-GPU path both keep a complete stage list and
+        # must use the base forward so abort cleanup and text_encoder_cpu_offload
+        # still apply. Releasing Qwen on every request was re-reading it from disk
+        # when neither deferral flag was on.
+        if self._defer_denoise_modules(fastvideo_args):
+            try:
+                self._ensure_text_encoder(fastvideo_args)
+                if self._denoise_stages_ready:
+                    logger.info("Running MiniMax-H3 condition stages before denoise (subsequent request)")
+                else:
+                    logger.info("Running MiniMax-H3 condition stages before loading DiT/VAE weights")
+                return self._run_condition_then_denoise(batch, fastvideo_args)
+            except BaseException:
+                self._release_all_lazy_modules()
+                raise
+        return super().forward(batch, fastvideo_args)
 
 
 class MiniMaxH3Pipeline(MiniMaxH3BasePipeline):

fastvideo/pipelines/basic/minimax_h3/packing.py

@@ -38,19 +38,36 @@
 MINIMAX_H3_KEYFRAME_NOISE_AUG = 0.999
 MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42
 
+_PATCH_SIZE_CACHE: dict[int, tuple[int, int, int]] = {}
+
 
 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)
+    cached = _PATCH_SIZE_CACHE.get(id(dit_config)) if dit_config is not None else None
+    if cached is not None:
+        return cached
     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:
+    axes = tuple(int(axis) for axis in patch_size)
+    if len(axes) != 3 or min(axes) <= 0:
         raise ValueError(f"MiniMax-H3 patch_size must be three positive ints, got {patch_size!r}.")
+    values = (axes[0], axes[1], axes[2])
+    if dit_config is not None:
+        _PATCH_SIZE_CACHE[id(dit_config)] = values
     return values
 
 
+def h3_latent_channels(model_config: Any, name: str) -> int:
+    """Read VAE latent width from arch config, not a live VAE proxy."""
+    arch = getattr(model_config, "arch_config", None)
+    value = getattr(arch, "latent_channels", None)
+    if value is None:
+        raise ValueError(f"MiniMax-H3 requires {name}.arch_config.latent_channels")
+    return int(value)
+
+
 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_latent_preparation.py

@@ -22,6 +22,7 @@
     build_packed_sequence,
     build_ref2va_packed_sequence,
     h3_dit_patch_size,
+    h3_latent_channels,
     keyframe_condition_noise,
     patchify_video_latents,
 )
@@ -59,6 +60,17 @@ def _sample_visual_posterior(posterior: Any) -> torch.Tensor:
     return posterior.sample(generator=generator)
 
 
+def _video_latent_channels(fastvideo_args: FastVideoArgs) -> int:
+    return h3_latent_channels(fastvideo_args.pipeline_config.vae_config, "vae_config")
+
+
+def _audio_latent_channels(fastvideo_args: FastVideoArgs) -> int:
+    return h3_latent_channels(
+        getattr(fastvideo_args.pipeline_config, "audio_vae_config", None),
+        "audio_vae_config",
+    )
+
+
 class MiniMaxH3LatentPreparationStage(PipelineStage):
     """Encode fixed conditions, build the row layout, then draw target noise."""
 
@@ -193,7 +205,7 @@ def _encode_fl2va_conditions(
         noise = keyframe_condition_noise(
             shapes,
             h3_dit_patch_size(fastvideo_args),
-            self.vae.latent_channels,
+            _video_latent_channels(fastvideo_args),
             generator=batch.generator,
             device=device,
         )
@@ -241,7 +253,7 @@ def _encode_ref2va_conditions(
         noise = keyframe_condition_noise(
             shapes,
             h3_dit_patch_size(fastvideo_args),
-            self.vae.latent_channels,
+            _video_latent_channels(fastvideo_args),
             generator=batch.generator,
             device=device,
         )
@@ -317,10 +329,11 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
                                             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)
+        audio_channels = _audio_latent_channels(fastvideo_args)
+        expected_audio_shape = (MINIMAX_H3_AUDIO_CHANNELS, audio_channels, num_audio_latents)
         if audio_noise is None:
             audio_rows = randn_tensor(
-                (num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS, self.audio_vae.latent_channels),
+                (num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS, audio_channels),
                 generator=batch.generator,
                 device=device,
                 dtype=torch.float32,
@@ -329,9 +342,7 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
             if tuple(audio_noise.shape) != expected_audio_shape:
                 raise ValueError(f"MiniMax-H3 injected audio latents must have shape {expected_audio_shape}, "
                                  f"got {tuple(audio_noise.shape)}.")
-            audio_rows = audio_noise.to(device=device,
-                                        dtype=torch.float32).permute(0, 2,
-                                                                     1).reshape(-1, self.audio_vae.latent_channels)
+            audio_rows = audio_noise.to(device=device, dtype=torch.float32).permute(0, 2, 1).reshape(-1, audio_channels)
 
         if condition_video is not None:
             video_rows = torch.cat((condition_video.to(device), video_rows))