[perf]: make deferred loading opt-in per pipeline and clean up on failure

Review of the first version found that the flag was unsafe as a global
default. Releasing a component and loading it again is only correct when
nothing outside the loader has changed it, and two habits in this tree break
that without raising:

  * mutating a component after load. LongCatPipeline.initialize_pipeline turns
    on block-sparse attention and writes parameters into every transformer
    block. That runs once, so a re-materialized component silently comes back
    with the feature off and the second generation quietly degrades.
  * reading a component's attributes while stages are built. The shared
    DenoisingStage.__init__ derives the attention backend from
    transformer.hidden_size, which materializes the DiT during post_init and
    defeats the deferral it was meant to gain.

_lazy_module_names is now empty in the base class, so an unchecked pipeline
gets no deferral and says so. MiniMax-H3 opts in to the four components this
PR measured, and nothing else changes behaviour.

Also from review:

Releasing on the failure path. A stage's hook frees only what that stage is
the last user of, and it ran only after a successful forward. Both the stage
and the whole run now release on the way out, so the retry a memory
constrained caller attempts does not start from a worse position than the
request that just failed. A failing release cannot replace the exception
being propagated.

Walking into nested stages. Cosmos25AutoDenoisingStage keeps the transformer
inside child stages, so a one-level scan called it unreferenced and never
freed it. The scan now recurses through stages and containers with cycle
protection.

Keeping the proxy when torch.compile is skipped. The FSDP check ran after the
proxy had already been replaced by the real module, so an FSDP-wrapped
component lost both the compile and its release hook.

Not attaching the activation trace to a deferred component, since the hook
manager pins every module it wraps.

test_parser.py asserted on a whole serialized config dict, so adding a field
to OffloadConfig broke it. Grepping the field name could not find that; only
running the suite could.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: KyleNeverGivesUp

docs/inference/offloading.md

@@ -113,6 +113,15 @@ the model already fits.
 This option applies to inference only. Training keeps every component resident
 and logs a warning if the flag is set.
 
+Deferral is opt-in per pipeline. Releasing a component and loading it again is
+only safe when nothing outside the loader has changed it, and two common habits
+break that without raising: mutating a component after load, as LongCat does
+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.
+
 ## General Recommendations
 
 ### Single GPU Inference

fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py

@@ -30,6 +30,10 @@ class MiniMaxH3BasePipeline(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.
+    _lazy_module_names = ("text_encoder", "transformer", "vae", "audio_vae")
 
     @classmethod
     def get_hf_download_component_dirs(cls) -> tuple[str, ...]:

fastvideo/pipelines/composed_pipeline_base.py

@@ -8,6 +8,7 @@
 import argparse
 import os
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
 from typing import Any, cast
 
 import torch
@@ -29,6 +30,40 @@
 logger = init_logger(__name__)
 
 
+def _iter_held_objects(stage: PipelineStage) -> Iterator[Any]:
+    """Yield everything a stage holds, walking into nested stages.
+
+    A stage can compose others rather than hold a component directly:
+    ``Cosmos25AutoDenoisingStage`` keeps the transformer inside its ``_t2w``
+    and ``_v2w`` children. A scan that stopped at the outer stage would call
+    that transformer unreferenced and never free it, so the flag would quietly
+    deliver less than it promises on those pipelines.
+    """
+    stack: list[Any] = [stage]
+    visited: set[int] = set()
+    while stack:
+        obj = stack.pop()
+        # is_lazy_module first: isinstance() on a proxy forwards __class__ and
+        # would load every deferred component just to work out where to free it.
+        if is_lazy_module(obj):
+            yield obj
+            continue
+        if isinstance(obj, PipelineStage | list | tuple | dict):
+            if id(obj) in visited:
+                continue
+            visited.add(id(obj))
+        if isinstance(obj, PipelineStage):
+            for key, value in vars(obj).items():
+                if key == "_lazy_modules_to_release":
+                    # Installed by this schedule, not a real use.
+                    continue
+                stack.append(value)
+        elif isinstance(obj, list | tuple):
+            stack.extend(obj)
+        elif isinstance(obj, dict):
+            stack.extend(obj.values())
+
+
 class ComposedPipelineBase(ABC):
     """
     Base class for pipelines composed of multiple stages.
@@ -52,23 +87,23 @@ class ComposedPipelineBase(ABC):
     # list, so a stage added afterwards can rebuild it instead of running
     # against a plan that predates it.
     _lazy_release_hooks_installed: bool = False
-    # Components eligible for deferred loading under ``lazy_module_load``.
-    # These are the weight-bearing ones; tokenizers, processors, and
-    # schedulers are cheap and stay resident so the pipeline can inspect them
-    # at construction time. Names follow the diffusers manifest convention, so
-    # the default list covers most pipelines; override per pipeline if a
-    # component is named differently or must never be released.
-    _lazy_module_names: tuple[str, ...] = (
-        "transformer",
-        "transformer_2",
-        "transformer_ref",
-        "transformer_refine",
-        "text_encoder",
-        "text_encoder_2",
-        "image_encoder",
-        "vae",
-        "audio_vae",
-    )
+    # Components this pipeline allows ``lazy_module_load`` to defer and free.
+    # Empty by default: deferral is opt-in per pipeline, because releasing a
+    # component and loading it again is only safe when nothing outside the
+    # loader has changed it. Two habits break that and neither raises:
+    #
+    #   * mutating a component after load. ``LongCatPipeline.initialize_pipeline``
+    #     turns on block-sparse attention and writes parameters into every
+    #     transformer block. That runs once, so a re-materialized component
+    #     silently comes back with the feature off.
+    #   * reading a component's attributes while building stages. The shared
+    #     ``DenoisingStage.__init__`` derives the attention backend from
+    #     ``transformer.hidden_size``, which materializes the DiT before the
+    #     first request and defeats the deferral it was meant to gain.
+    #
+    # A pipeline opts in by listing the components it has checked. Names match
+    # the diffusers manifest.
+    _lazy_module_names: tuple[str, ...] = ()
 
     @classmethod
     def get_hf_download_component_dirs(cls) -> tuple[str, ...] | None:
@@ -163,14 +198,12 @@ def _maybe_compile_pipeline_module(
         if module_name not in self.modules:
             return
 
-        module = self.modules[module_name]
-        if is_lazy_module(module):
-            # torch.compile replaces the entry in self.modules, which would
-            # drop the proxy and with it the ability to release. Load now and
-            # keep this component resident for the run.
-            logger.info("torch.compile requested for %s; loading it eagerly instead of deferring", module_name)
-            module = module.materialize()
-            self.modules[module_name] = module
+        entry = self.modules[module_name]
+        # Materialize into a local. The dict keeps the proxy until we know a
+        # compiled callable is actually going to replace it, so a component
+        # that turns out to be FSDP-wrapped is neither compiled nor stripped of
+        # its release hook.
+        module = entry.materialize() if is_lazy_module(entry) else entry
         if fsdp_module_cls is not None and isinstance(module, fsdp_module_cls):
             logger.info(
                 "%s is already FSDP-wrapped; skipping torch.compile in pipeline",
@@ -195,6 +228,9 @@ def _maybe_compile_pipeline_module(
 
         # Backward-compatible fallback: compile full module if no condition matched.
         logger.info("Enabling torch.compile for %s with kwargs=%s", module_name, compile_kwargs)
+        if is_lazy_module(entry):
+            logger.info("Whole-module torch.compile replaces the deferred %s, so it stays resident for the run",
+                        module_name)
         self.modules[module_name] = torch.compile(module, **compile_kwargs)
 
     def post_init(self) -> None:
@@ -287,7 +323,15 @@ def post_init(self) -> None:
                     )
                     logger.info("Torch Compile enabled for audio VAE")
 
-        self._trace_mgr = attach_activation_trace(self.modules.get("transformer"))
+        trace_target = self.modules.get("transformer")
+        if is_lazy_module(trace_target):
+            # The hook manager keeps a strong reference to every module it
+            # wraps, so attaching here would materialize the DiT before the
+            # first request and pin that instance past any release.
+            logger.warning("Activation trace is not attached to a deferred transformer; "
+                           "turn off lazy_module_load to trace it")
+            trace_target = None
+        self._trace_mgr = attach_activation_trace(trace_target)
 
         if not self.fastvideo_args.training_mode:
             logger.info("Creating pipeline stages...")
@@ -561,25 +605,10 @@ def _build_lazy_release_schedule(self) -> dict[int, list[str]]:
 
         last_use: dict[str, int] = {}
         for index, stage in enumerate(self._stages):
-            for key, value in vars(stage).items():
-                if key == "_lazy_modules_to_release":
-                    # Installed by this schedule, not a real use.
-                    continue
-                # is_lazy_module first: isinstance() on a proxy forwards
-                # __class__ and would load every deferred module just to work
-                # out where to release it.
-                if is_lazy_module(value):
-                    candidates: tuple[Any, ...] = (value, )
-                elif isinstance(value, list | tuple):
-                    candidates = tuple(value)
-                elif isinstance(value, dict):
-                    candidates = tuple(value.values())
-                else:
-                    continue
-                for candidate in candidates:
-                    name = lazy_names_by_id.get(id(candidate))
-                    if name is not None:
-                        last_use[name] = index
+            for held in _iter_held_objects(stage):
+                name = lazy_names_by_id.get(id(held))
+                if name is not None:
+                    last_use[name] = index
 
         schedule: dict[int, list[str]] = {}
         for name, index in sorted(last_use.items()):
@@ -612,6 +641,17 @@ def _install_lazy_release_hooks(self) -> None:
                         getattr(self._stages[index], "_pipeline_stage_name", "?"), names)
         self._lazy_release_hooks_installed = True
 
+    def _release_all_lazy_modules(self) -> None:
+        """Free every deferred component that is currently materialized."""
+        for module_name, module in self.modules.items():
+            if not is_lazy_module(module):
+                continue
+            try:
+                module.release()
+            except Exception:
+                # Never let cleanup replace the exception being propagated.
+                logger.exception("Failed to release deferred module %s", module_name)
+
     def add_stage(self, stage_name: str, stage: PipelineStage):
         assert self.modules is not None, "No modules are registered"
         # Preserve the pipeline-unique stage key for structured metrics.
@@ -653,8 +693,17 @@ def forward(
         # Execute each stage
         logger.info("Running pipeline stages: %s", self._stage_name_mapping.keys())
         # logger.info("Batch: %s", batch)
-        for stage in self.stages:
-            batch = stage(batch, fastvideo_args)
+        try:
+            for stage in self.stages:
+                batch = stage(batch, fastvideo_args)
+        except BaseException:
+            # A stage's own hook frees only what that stage was the last user
+            # of. When the run aborts earlier, everything already materialized
+            # stays for the life of the generator, and the retry a
+            # memory-constrained caller is most likely to attempt starts from a
+            # worse position than the request that just failed.
+            self._release_all_lazy_modules()
+            raise
 
         # Return the output
         return batch

fastvideo/pipelines/stages/base.py

@@ -151,6 +151,34 @@ def __call__(
                 raise
 
         # Execute the actual stage logic
+        try:
+            result = self._execute(batch, fastvideo_args, stage_key, stage_class_name, stage_name)
+        except BaseException:
+            self._release_deferred_modules(stage_name)
+            raise
+
+        if enable_verification:
+            # Post-execution output verification
+            try:
+                output_result = self.verify_output(result, fastvideo_args)
+                self._run_verification(output_result, stage_name, "output")
+            except Exception as e:
+                logger.error("Output verification failed for %s: %s", stage_name, str(e))
+                self._release_deferred_modules(stage_name)
+                raise
+
+        self._release_deferred_modules(stage_name)
+        return result
+
+    def _execute(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+        stage_key: str,
+        stage_class_name: str,
+        stage_name: str,
+    ) -> ForwardBatch:
+        """Run forward, with the optional timing and logging wrapper."""
         if envs.FASTVIDEO_STAGE_LOGGING:
             logger.info("[%s] Starting execution", stage_name)
             torch.cuda.synchronize()
@@ -176,19 +204,23 @@ def __call__(
             # Direct execution (current behavior)
             result = self.forward(batch, fastvideo_args)
 
-        if enable_verification:
-            # Post-execution output verification
-            try:
-                output_result = self.verify_output(result, fastvideo_args)
-                self._run_verification(output_result, stage_name, "output")
-            except Exception as e:
-                logger.error("Output verification failed for %s: %s", stage_name, str(e))
-                raise
+        return result
 
-        for lazy_module in self._lazy_modules_to_release:
-            lazy_module.release()
+    def _release_deferred_modules(self, stage_name: str) -> None:
+        """Free the deferred components this stage is the last user of.
 
-        return result
+        Called on the way out whether or not the stage succeeded. A stage that
+        raises after materializing a multi-gigabyte component would otherwise
+        keep it for the life of the generator, and the retry that a
+        memory-constrained caller is most likely to attempt would start from a
+        worse position than the request that just failed.
+        """
+        for lazy_module in self._lazy_modules_to_release:
+            try:
+                lazy_module.release()
+            except Exception:
+                # Never let cleanup replace the exception being propagated.
+                logger.exception("Failed to release deferred module after %s", stage_name)
 
     @abstractmethod
     def forward(

fastvideo/tests/api/test_parser.py

@@ -114,6 +114,7 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None:
                     "image_encoder": True,
                     "vae": True,
                     "pin_cpu_memory": True,
+                    "lazy_module_load": False,
                 },
                 "compile": {
                     "enabled": False,