rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Merge branch 'united' of https://github.com/henk717/KoboldAI into accelerate-offloading

Browse Source
This commit is contained in:
somebody
2023-07-07 14:54:12 -05:00
parent cfe1f5b514 76d21bb142
commit 35f3687667
5 changed files with 486 additions and 202 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
modeling/lazy_loader.py
View File

@@ -172,14 +172,11 @@ class TorchLazyTensor(LazyTensor):
CheckpointChunkCache.file_name = filename
CheckpointChunkCache.key = self.key
try:
CheckpointChunkCache.handle = checkpoint.open(
f"archive/data/{self.key}", "r"
)
except KeyError:
CheckpointChunkCache.handle = checkpoint.open(
f"{filename}/data/{self.key}", "r"
)
ziproot = checkpoint.namelist()[0].split("/")[0]
CheckpointChunkCache.handle = checkpoint.open(f"{ziproot}/data/{self.key}", "r")
else:
# Cache hit. Hip hip hooray! :^)
# print(".", end="", flush=True)
@@ -523,8 +520,8 @@ def use_lazy_load(
torch.load = torch_load
if HAS_SAFETENSORS:
patch_safetensors(callback)
#if HAS_SAFETENSORS:
#patch_safetensors(callback)
if dematerialized_modules:
# Most devices can just use Accelerate's implementation, but the Transformers on