|
|
|
|
|
|
|
import torch |
|
|
|
|
|
cpu = torch.device('cpu') |
|
gpu = torch.device(f'cuda:{torch.cuda.current_device()}') |
|
gpu_complete_modules = [] |
|
|
|
|
|
class DynamicSwapInstaller: |
|
@staticmethod |
|
def _install_module(module: torch.nn.Module, **kwargs): |
|
original_class = module.__class__ |
|
module.__dict__['forge_backup_original_class'] = original_class |
|
|
|
def hacked_get_attr(self, name: str): |
|
if '_parameters' in self.__dict__: |
|
_parameters = self.__dict__['_parameters'] |
|
if name in _parameters: |
|
p = _parameters[name] |
|
if p is None: |
|
return None |
|
if p.__class__ == torch.nn.Parameter: |
|
return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad) |
|
else: |
|
return p.to(**kwargs) |
|
if '_buffers' in self.__dict__: |
|
_buffers = self.__dict__['_buffers'] |
|
if name in _buffers: |
|
return _buffers[name].to(**kwargs) |
|
return super(original_class, self).__getattr__(name) |
|
|
|
module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), { |
|
'__getattr__': hacked_get_attr, |
|
}) |
|
|
|
return |
|
|
|
@staticmethod |
|
def _uninstall_module(module: torch.nn.Module): |
|
if 'forge_backup_original_class' in module.__dict__: |
|
module.__class__ = module.__dict__.pop('forge_backup_original_class') |
|
return |
|
|
|
@staticmethod |
|
def install_model(model: torch.nn.Module, **kwargs): |
|
for m in model.modules(): |
|
DynamicSwapInstaller._install_module(m, **kwargs) |
|
return |
|
|
|
@staticmethod |
|
def uninstall_model(model: torch.nn.Module): |
|
for m in model.modules(): |
|
DynamicSwapInstaller._uninstall_module(m) |
|
return |
|
|
|
|
|
def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device): |
|
if hasattr(model, 'scale_shift_table'): |
|
model.scale_shift_table.data = model.scale_shift_table.data.to(target_device) |
|
return |
|
|
|
for k, p in model.named_modules(): |
|
if hasattr(p, 'weight'): |
|
p.to(target_device) |
|
return |
|
|
|
|
|
def get_cuda_free_memory_gb(device=None): |
|
if device is None: |
|
device = gpu |
|
|
|
memory_stats = torch.cuda.memory_stats(device) |
|
bytes_active = memory_stats['active_bytes.all.current'] |
|
bytes_reserved = memory_stats['reserved_bytes.all.current'] |
|
bytes_free_cuda, _ = torch.cuda.mem_get_info(device) |
|
bytes_inactive_reserved = bytes_reserved - bytes_active |
|
bytes_total_available = bytes_free_cuda + bytes_inactive_reserved |
|
return bytes_total_available / (1024 ** 3) |
|
|
|
|
|
def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0): |
|
print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB') |
|
|
|
for m in model.modules(): |
|
if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb: |
|
torch.cuda.empty_cache() |
|
return |
|
|
|
if hasattr(m, 'weight'): |
|
m.to(device=target_device) |
|
|
|
model.to(device=target_device) |
|
torch.cuda.empty_cache() |
|
return |
|
|
|
|
|
def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0): |
|
print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB') |
|
|
|
for m in model.modules(): |
|
if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb: |
|
torch.cuda.empty_cache() |
|
return |
|
|
|
if hasattr(m, 'weight'): |
|
m.to(device=cpu) |
|
|
|
model.to(device=cpu) |
|
torch.cuda.empty_cache() |
|
return |
|
|
|
|
|
def unload_complete_models(*args): |
|
for m in gpu_complete_modules + list(args): |
|
m.to(device=cpu) |
|
print(f'Unloaded {m.__class__.__name__} as complete.') |
|
|
|
gpu_complete_modules.clear() |
|
torch.cuda.empty_cache() |
|
return |
|
|
|
|
|
def load_model_as_complete(model, target_device, unload=True): |
|
if unload: |
|
unload_complete_models() |
|
|
|
model.to(device=target_device) |
|
print(f'Loaded {model.__class__.__name__} to {target_device} as complete.') |
|
|
|
gpu_complete_modules.append(model) |
|
return |
|
|
