Deferred loading and release of heavy pipeline modules.
A pipeline normally materializes every component before the first stage runs, so peak memory is the sum of all components even though no two of them are needed at the same moment. On a unified-memory device that sum is charged against the same pool the activations come from, and a model whose components individually fit can still fail to load.
LazyModule turns that sum into a maximum. It stands in for a component, loads it on first use, and drops it once the last stage holding it has run. The pipeline decides when to release; this module only owns the proxying and the load/free mechanics.
Classes
A stand-in for a pipeline component that loads on first use.
Every attribute access, call, and isinstance check forwards to the real component, materializing it if needed. release drops the reference and frees the allocator cache; a later access re-runs the loader, so releasing early is a latency cost, never a correctness one.
Source code in fastvideo/pipelines/lazy_module.py
| 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)
object.__setattr__(self, "_lazy_release_callbacks", [])
|
Methods:
add_release_callback(callback: Callable[[], None]) -> None
Run callback each time the real component is dropped.
Source code in fastvideo/pipelines/lazy_module.py
| def add_release_callback(self, callback: Callable[[], None]) -> None:
"""Run ``callback`` each time the real component is dropped."""
object.__getattribute__(self, "_lazy_release_callbacks").append(callback)
|
Return the real component, loading it if this is the first use.
Source code in fastvideo/pipelines/lazy_module.py
| def materialize(self) -> Any:
"""Return the real component, loading it if this is the first use."""
module = object.__getattribute__(self, "_lazy_module")
if module is not None:
return module
name = object.__getattribute__(self, "_lazy_name")
loader = object.__getattribute__(self, "_lazy_loader")
logger.info("Loading deferred module %s", name)
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
|
Drop the real component. Returns True if something was released.
Source code in fastvideo/pipelines/lazy_module.py
| def release(self) -> bool:
"""Drop the real component. Returns True if something was released."""
module = object.__getattribute__(self, "_lazy_module")
if module is None:
return False
name = object.__getattribute__(self, "_lazy_name")
before = _cuda_allocated_gib()
object.__setattr__(self, "_lazy_module", None)
del module
for callback in list(object.__getattribute__(self, "_lazy_release_callbacks")):
callback()
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
after = _cuda_allocated_gib()
if before is not None and after is not None:
logger.info("Released deferred module %s, cuda allocated %.2f -> %.2f GiB, freed %.2f GiB", name, before,
after, before - after)
else:
logger.info("Released deferred module %s", name)
return True
|
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.
Source code in fastvideo/pipelines/lazy_module.py
| 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")
module = object.__getattribute__(self, "_lazy_module")
if current_transform is not None:
inner = current_transform
def chained(loaded: Any) -> Any:
return transform(inner(loaded))
stored: Callable[[Any], Any] = chained
# The resident instance already ran ``inner``; only apply the new outer.
immediate = transform
else:
stored = transform
immediate = transform
transformed = immediate(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", stored)
if module is not None:
object.__setattr__(self, "_lazy_module", transformed)
|
Functions:
Type test that does not materialize, unlike isinstance.
Source code in fastvideo/pipelines/lazy_module.py
| def is_lazy_module(obj: Any) -> TypeGuard[LazyModule]:
"""Type test that does not materialize, unlike ``isinstance``."""
return type(obj) is LazyModule
|