[bugfix]: preserve compile across lazy module reloads

Register pipeline-level compilation as a lazy materialization transform so every freshly loaded component receives the same compile setup. Keep whole-module compilation behind the proxy, preserve release scheduling, and add regression coverage for conditional and whole-module compile paths.

Author: Satyam-53

docs/inference/offloading.md

@@ -107,7 +107,9 @@ pipeline that is roughly `max(text encoder, DiT + VAE)` rather than
 
 A freed component is read from disk again on the next generation, so a
 multi-prompt run pays one reload per component per request. For a large text
-encoder that is tens of seconds.
+encoder that is tens of seconds. If pipeline-level `torch.compile` is enabled,
+the compile setup is reapplied after each reload; PyTorch can reuse its graph
+and kernel caches when the component structure and input shapes are unchanged.
 
 #### Usage Recommendation
 

fastvideo/pipelines/composed_pipeline_base.py

@@ -9,6 +9,7 @@
 import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
+from functools import partial
 from typing import Any, cast
 
 import torch
@@ -201,49 +202,71 @@ def _compile_with_conditions(
                 compiled_count += 1
         return compiled_count
 
-    def _maybe_compile_pipeline_module(
-        self,
+    @staticmethod
+    def _compile_pipeline_module_instance(
         module_name: str,
+        module: torch.nn.Module,
         fsdp_module_cls: type | None,
         compile_kwargs: dict[str, Any],
-    ) -> None:
-        if module_name not in self.modules:
-            return
-
-        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
+    ) -> Any:
+        """Apply pipeline-level compile setup to one loaded component."""
         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",
                 module_name.capitalize(),
             )
-            return
+            return module
 
         prepare_for_compile = getattr(module, "prepare_for_compile", None)
         if callable(prepare_for_compile):
             logger.info("Running prepare_for_compile for %s", module_name)
             prepare_for_compile()
 
-        compiled_count = self._compile_with_conditions(module, compile_kwargs)
+        compiled_count = ComposedPipelineBase._compile_with_conditions(module, compile_kwargs)
         if compiled_count > 0:
             logger.info(
                 "Enabled torch.compile for %d submodules in %s via _compile_conditions with kwargs=%s",
                 compiled_count,
                 module_name,
                 compile_kwargs,
             )
-            return
+            return 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)
+        return torch.compile(module, **compile_kwargs)
+
+    def _maybe_compile_pipeline_module(
+        self,
+        module_name: str,
+        fsdp_module_cls: type | None,
+        compile_kwargs: dict[str, Any],
+    ) -> None:
+        if module_name not in self.modules:
+            return
+
+        entry = self.modules[module_name]
         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)
+            # Compilation is part of materialization, not a one-time mutation
+            # of the first loaded instance. The proxy remains in the module
+            # map, so whole-module and conditional compile both survive every
+            # release/reload cycle without making initialization eager.
+            entry.set_materialize_transform(
+                partial(
+                    ComposedPipelineBase._compile_pipeline_module_instance,
+                    module_name,
+                    fsdp_module_cls=fsdp_module_cls,
+                    compile_kwargs=dict(compile_kwargs),
+                ))
+            logger.info("Configured torch.compile for every materialization of deferred %s", module_name)
+            return
+
+        self.modules[module_name] = self._compile_pipeline_module_instance(
+            module_name,
+            entry,
+            fsdp_module_cls,
+            compile_kwargs,
+        )
 
     def post_init(self) -> None:
         assert self.fastvideo_args is not None, "fastvideo_args must be set"

fastvideo/pipelines/lazy_module.py

@@ -43,11 +43,12 @@ class LazyModule:
     early is a latency cost, never a correctness one.
     """
 
-    __slots__ = ("_lazy_name", "_lazy_loader", "_lazy_module")
+    __slots__ = ("_lazy_name", "_lazy_loader", "_lazy_materialize_transform", "_lazy_module")
 
     def __init__(self, name: str, loader: Callable[[], Any]) -> None:
         object.__setattr__(self, "_lazy_name", name)
         object.__setattr__(self, "_lazy_loader", loader)
+        object.__setattr__(self, "_lazy_materialize_transform", None)
         object.__setattr__(self, "_lazy_module", None)
 
     @property
@@ -70,13 +71,40 @@ def materialize(self) -> Any:
         module = loader()
         if module is None:
             raise ValueError(f"Deferred loader for module {name} returned None")
+
+        transform = object.__getattribute__(self, "_lazy_materialize_transform")
+        if transform is not None:
+            module = transform(module)
+            if module is None:
+                raise ValueError(f"Materialize transform for module {name} returned None")
         object.__setattr__(self, "_lazy_module", module)
 
         allocated = _cuda_allocated_gib()
         if allocated is not None:
             logger.info("Loaded deferred module %s, cuda allocated now %.2f GiB", name, allocated)
         return module
 
+    def set_materialize_transform(self, transform: Callable[[Any], Any]) -> None:
+        """Apply ``transform`` to this and every future loaded instance.
+
+        Registering a transform does not itself load a deferred component. If
+        something has already materialized the component, transform that
+        instance immediately so current and future instances have the same
+        setup. A transform may return a wrapper, as ``torch.compile`` does.
+        """
+        current_transform = object.__getattribute__(self, "_lazy_materialize_transform")
+        if current_transform is not None:
+            raise RuntimeError(f"Materialize transform for module {self.lazy_name} is already set")
+
+        module = object.__getattribute__(self, "_lazy_module")
+        transformed = transform(module) if module is not None else None
+        if module is not None and transformed is None:
+            raise ValueError(f"Materialize transform for module {self.lazy_name} returned None")
+
+        object.__setattr__(self, "_lazy_materialize_transform", transform)
+        if module is not None:
+            object.__setattr__(self, "_lazy_module", transformed)
+
     def release(self) -> bool:
         """Drop the real component. Returns True if something was released."""
         module = object.__getattribute__(self, "_lazy_module")

fastvideo/tests/stages/test_lazy_module_load.py

@@ -154,6 +154,30 @@ def test_release_then_reload_is_correct_not_broken():
     assert second.tag == "c"
 
 
+def test_materialize_transform_applies_to_every_loaded_instance_without_loading_eagerly():
+    loader, calls = _counting_loader()
+    transformed = []
+    module = LazyModule("vae", loader)
+
+    def transform(component):
+        transformed.append(component)
+        component.tag = f"compiled-{component.tag}"
+        return component
+
+    module.set_materialize_transform(transform)
+    assert calls == []
+
+    first = module.materialize()
+    assert first.tag == "compiled-c"
+    assert module.release() is True
+
+    second = module.materialize()
+    assert second.tag == "compiled-c"
+    assert second is not first
+    assert calls == ["c", "c"]
+    assert transformed == [first, second]
+
+
 def test_release_without_materializing_is_a_noop():
     loader, calls = _counting_loader()
     module = LazyModule("text_encoder", loader)
@@ -204,6 +228,93 @@ def _lazy(name):
     return LazyModule(name, lambda: _Component(name))
 
 
+def test_pipeline_compile_is_reapplied_after_lazy_release(monkeypatch):
+    loads = []
+    compile_calls = []
+
+    class _CompileAwareComponent(torch.nn.Module):
+        _compile_conditions = (lambda name, module: name == "block", )
+
+        def __init__(self):
+            super().__init__()
+            self.block = torch.nn.Linear(2, 2)
+            self.prepare_calls = 0
+
+        def prepare_for_compile(self):
+            self.prepare_calls += 1
+
+    def load_component():
+        component = _CompileAwareComponent()
+        loads.append(component)
+        return component
+
+    def fake_compile(target, **kwargs):
+        compile_calls.append((target, kwargs))
+        return target
+
+    monkeypatch.setattr(torch, "compile", fake_compile)
+    lazy = LazyModule("vae", load_component)
+    pipeline = _FakePipeline({"vae": lazy}, [])
+
+    pipeline._maybe_compile_pipeline_module("vae", None, {"mode": "reduce-overhead"})
+    assert loads == []
+
+    first = lazy.materialize()
+    assert first.prepare_calls == 1
+    assert lazy.release() is True
+
+    second = lazy.materialize()
+    assert second is not first
+    assert second.prepare_calls == 1
+    assert loads == [first, second]
+    assert len(compile_calls) == 2
+    assert [kwargs for _, kwargs in compile_calls] == [
+        {"mode": "reduce-overhead"},
+        {"mode": "reduce-overhead"},
+    ]
+
+
+def test_whole_module_compile_keeps_lazy_proxy_and_recompiles_after_release(monkeypatch):
+    loads = []
+    compile_calls = []
+
+    class _WholeComponent(torch.nn.Module):
+
+        def forward(self, value):
+            return value
+
+    class _Compiled:
+
+        def __init__(self, original):
+            self.original = original
+
+    def load_component():
+        component = _WholeComponent()
+        loads.append(component)
+        return component
+
+    def fake_compile(target, **kwargs):
+        compile_calls.append((target, kwargs))
+        return _Compiled(target)
+
+    monkeypatch.setattr(torch, "compile", fake_compile)
+    lazy = LazyModule("text_encoder", load_component)
+    pipeline = _FakePipeline({"text_encoder": lazy}, [])
+
+    pipeline._maybe_compile_pipeline_module("text_encoder", None, {"dynamic": True})
+    assert pipeline.modules["text_encoder"] is lazy
+    assert loads == []
+
+    first = lazy.materialize()
+    assert first.original is loads[0]
+    assert lazy.release() is True
+
+    second = lazy.materialize()
+    assert second.original is loads[1]
+    assert second is not first
+    assert len(compile_calls) == 2
+
+
 def test_schedule_releases_after_the_last_stage_that_holds_a_module():
     text_encoder = _lazy("text_encoder")
     transformer = _lazy("transformer")