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

Fix loading bar

Browse Source
This commit is contained in:
somebody
2023-06-21 16:27:22 -05:00
parent aca2b532d7
commit c56214c275
2 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
modeling/lazy_loader.py
View File

@@ -159,7 +159,7 @@ class TorchLazyTensor(LazyTensor):
):
# Cache miss. Assuming weights are loaded in order of
# (key, seek_offset), this means we need to invalidate the cache.
print("!", end="", flush=True)
# print("!", end="", flush=True)
CheckpointChunkCache.hit_data["misses"] += 1
CheckpointChunkCache.clear()
@@ -176,7 +176,7 @@ class TorchLazyTensor(LazyTensor):
)
else:
# Cache hit. Hip hip hooray! :^)
print(".", end="", flush=True)
# print(".", end="", flush=True)
CheckpointChunkCache.hit_data["hits"] += 1
size = reduce(lambda x, y: x * y, self.shape, 1)
@@ -481,6 +481,11 @@ def post_load_cleanup() -> None:
logger.debug(f"[lazy_load] CheckpointChunkCache Hit Data: {CheckpointChunkCache.hit_data}")
CheckpointChunkCache.clear(unload_model=True)
# Bar is initialized in
# patches.patch_transformers_for_lazyload._load_state_dict_into_meta_model,
# as it has access to the state dict (for getting tensor count)
utils.bar = None
for v in torch_checkpoint_file_handles.values():
v.close()
torch_checkpoint_file_handles = {}