github-actions[bot]
36 lines
1.4 KiB
Python
36 lines
1.4 KiB
Python
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from typing import TypeVar
|
|
|
|
import torch
|
|
|
|
from ltx_core.devices import synchronize_device
|
|
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
|
from ltx_pipelines.utils.helpers import cleanup_memory
|
|
|
|
_M = TypeVar("_M", bound=torch.nn.Module)
|
|
|
|
|
|
@contextmanager
|
|
def gpu_model(model: _M, alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM) -> Iterator[_M]:
|
|
"""Context manager that yields a model and releases its memory on exit.
|
|
On ``TRIM`` (default): synchronize, move parameters/buffers to the ``meta``
|
|
device (releasing GPU+CPU storage), then ``cleanup_memory()`` to return
|
|
cached blocks to the OS. ``DEFER`` skips this -- the model's storage is
|
|
reclaimed by normal GC and the CUDA caching allocator stays warm for the
|
|
next build (cheaper for back-to-back runs).
|
|
Usage::
|
|
with gpu_model(build_encoder()) as encoder:
|
|
... # use encoder -- typed as the concrete class
|
|
# GPU + CPU memory freed automatically
|
|
"""
|
|
try:
|
|
yield model
|
|
finally:
|
|
if alloc_trim_strategy == AllocatorTrimStrategy.TRIM:
|
|
synchronize_device()
|
|
# .to("meta") releases storage for all parameters/buffers regardless
|
|
# of their original device (CUDA or CPU).
|
|
model.to("meta")
|
|
cleanup_memory()
|