[perf]: rebuild the deferred-release schedule when a stage is added late

The schedule maps each deferred component to the last stage that holds it,
derived once after create_pipeline_stages. Every pipeline in the tree builds
its stages there, but nothing enforced it. A stage appended afterwards could
hold a component an earlier stage had already been told to free, and would
then be handed a released component mid-run with no error.

add_stage now rebuilds the schedule and says so, turning an invariant nothing
checked into a visible self-correction.

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

Author: KyleNeverGivesUp

fastvideo/pipelines/composed_pipeline_base.py

@@ -48,6 +48,10 @@ class ComposedPipelineBase(ABC):
     trainable_transformer_names: list[str] = ["transformer"]
     trainable_transformer_modules: dict[str, torch.nn.Module] = {}
     post_init_called: bool = False
+    # Set once the deferred-release schedule has been derived from the stage
+    # 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
@@ -600,11 +604,13 @@ def _install_lazy_release_hooks(self) -> None:
                 "lazy_module_load is on but no deferred module is held by a stage, so nothing will be "
                 "freed mid-run. Pipeline %s may load its modules eagerly or hold them outside its stages.",
                 type(self).__name__)
+            self._lazy_release_hooks_installed = True
             return
 
         for index, names in sorted(schedule.items()):
             logger.info("Deferred modules to free after stage %d (%s): %s", index,
                         getattr(self._stages[index], "_pipeline_stage_name", "?"), names)
+        self._lazy_release_hooks_installed = True
 
     def add_stage(self, stage_name: str, stage: PipelineStage):
         assert self.modules is not None, "No modules are registered"
@@ -616,6 +622,15 @@ def add_stage(self, stage_name: str, stage: PipelineStage):
         self._stage_name_mapping[stage_name] = stage
         setattr(self, stage_name, stage)
 
+        if self._lazy_release_hooks_installed:
+            # The schedule maps each deferred module to its last holder. A
+            # stage appended afterwards may hold a module an earlier stage has
+            # already been told to free, which would hand it a released
+            # component mid-run. Rebuild rather than trust the stale plan.
+            logger.warning("Stage %s was added after the deferred-release schedule was built; rebuilding the schedule",
+                           stage_name)
+            self._install_lazy_release_hooks()
+
     # TODO(will): don't hardcode no_grad
     @torch.no_grad()
     def forward(

fastvideo/tests/stages/test_lazy_module_load.py

@@ -383,3 +383,25 @@ def test_a_stage_that_rebinds_through_to_can_still_be_released():
     stage(object(), SimpleNamespace(enable_stage_verification=False))
 
     assert not vae.is_materialized
+
+
+def test_a_stage_added_after_the_schedule_rebuilds_it(caplog):
+    # The schedule is derived from the stage list. A stage appended afterwards
+    # could hold a module an earlier stage was already told to free, which
+    # would hand it a released component mid-run.
+    vae = _lazy("vae")
+    first = _EchoStage(vae=vae)
+    pipeline = _FakePipeline({"vae": vae}, [])
+    pipeline._stage_name_mapping = {}
+    pipeline.add_stage("first", first)
+    pipeline._install_lazy_release_hooks()
+
+    assert first._lazy_modules_to_release == (vae, )
+
+    later = _EchoStage(vae=vae)
+    with caplog.at_level("WARNING"):
+        pipeline.add_stage("later", later)
+
+    assert "rebuilding the schedule" in caplog.text
+    assert first._lazy_modules_to_release == ()
+    assert later._lazy_modules_to_release == (vae, )