KoboldAI-Client/torch_lazy_loader.py

98 lines
3.9 KiB
Python
Raw Normal View History

2022-03-01 21:40:44 +01:00
import contextlib
from functools import reduce
import zipfile
2022-03-01 21:40:44 +01:00
import pickle
import torch
from typing import Any, Callable, Dict, Optional, Tuple, Type
class LazyTensor:
def __init__(self, storage_type: Type[torch._StorageBase], key: str, location: str, storage_offset: Optional[int] = None, shape: Optional[Tuple[int, ...]] = None, stride: Optional[Tuple[int, ...]] = None, requires_grad=False, backward_hooks: Any = None):
2022-03-01 21:40:44 +01:00
self.storage_type = storage_type
self.key = key
self.location = location
self.storage_offset = storage_offset
self.shape = shape
self.stride = stride
self.requires_grad = requires_grad
self.backward_hooks = backward_hooks
def __view(self, f: Callable):
return f"{type(self).__name__}(storage_type={f(self.storage_type)}, key={f(self.key)}, location={f(self.location)}, storage_offset={f(self.storage_offset)}, shape={f(self.shape)}, stride={f(self.stride)}, requires_grad={f(self.requires_grad)}, backward_hooks={f(self.backward_hooks)})"
2022-03-01 21:40:44 +01:00
def __repr__(self):
return self.__view(repr)
def materialize(self, checkpoint: zipfile.ZipFile, map_location=None) -> torch.Tensor:
size = reduce(lambda x, y: x * y, self.shape, 1)
dtype = self.storage_type(0).dtype
nbytes = size if dtype is torch.bool else size * ((torch.finfo if dtype.is_floating_point else torch.iinfo)(dtype).bits >> 3)
with checkpoint.open(f"archive/data/{self.key}", "r") as f:
f.seek(self.storage_offset)
storage = self.storage_type.from_buffer(f.read(nbytes), "little")
2022-03-01 21:40:44 +01:00
storage = torch.serialization._get_restore_location(map_location)(storage, self.location)
tensor = torch.tensor([], dtype=storage.dtype, device=storage.device)
tensor.set_(storage, 0, self.shape, self.stride)
2022-03-01 21:40:44 +01:00
tensor.requires_grad = self.requires_grad
tensor._backward_hooks = self.backward_hooks
return tensor
class _LazyUnpickler(pickle.Unpickler):
lazy_loaded_storages: Dict[str, LazyTensor]
def __init__(self, *args, **kwargs):
self.lazy_loaded_storages = {}
return super().__init__(*args, **kwargs)
def forced_persistent_load(self, saved_id):
assert isinstance(saved_id, tuple)
typename = saved_id[0]
assert typename == "storage", f"Unknown typename for persistent_load, expected 'storage' but got '{typename}'"
storage_type, key, location, _ = saved_id[1:]
return LazyTensor(storage_type, key, location)
2022-03-01 21:40:44 +01:00
def load(self, *args, **kwargs):
self.persistent_load = self.forced_persistent_load
retval = super().load(*args, **kwargs)
self.lazy_loaded_storages = {}
return retval
def _rebuild_tensor(lazy_storage: LazyTensor, storage_offset, shape, stride):
lazy_storage.storage_offset = storage_offset
lazy_storage.shape = shape
lazy_storage.stride = stride
return lazy_storage
@contextlib.contextmanager
def use_lazy_torch_load(enable=True, callback: Optional[Callable] = None):
if not enable:
yield False
return
try:
old_unpickler = pickle.Unpickler
pickle.Unpickler = _LazyUnpickler
2022-03-01 21:40:44 +01:00
old_rebuild_tensor = torch._utils._rebuild_tensor
torch._utils._rebuild_tensor = _rebuild_tensor
2022-03-01 21:40:44 +01:00
old_torch_load = torch.load
2022-03-01 21:40:44 +01:00
def torch_load(f, map_location=None, pickle_module=pickle, **pickle_load_args):
retval = old_torch_load(f=f, map_location=map_location, pickle_module=pickle_module, **pickle_load_args)
if callback is not None:
callback(retval, f=f, map_location=map_location, pickle_module=pickle_module, **pickle_load_args)
return retval
2022-03-01 21:40:44 +01:00
torch.load = torch_load
2022-03-01 21:40:44 +01:00
yield True
2022-03-01 21:40:44 +01:00
finally:
pickle.Unpickler = old_unpickler
torch._utils._rebuild_tensor = old_rebuild_tensor
torch.load = old_torch_load