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

Giving proper TPU load progress in UI

Browse Source
This commit is contained in:
ebolam
2022-09-25 20:03:15 -04:00
parent 74faad94ac
commit 55771c1edf
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tpu_mtj_backend.py
View File

@@ -1247,6 +1247,7 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
else: else:
num_tensors = len(model_dict) num_tensors = len(model_dict)
utils.bar = tqdm(total=num_tensors, desc="Loading model tensors") utils.bar = tqdm(total=num_tensors, desc="Loading model tensors")
koboldai_vars.total_layers = num_tensors
if utils.num_shards is not None: if utils.num_shards is not None:
utils.current_shard += 1 utils.current_shard += 1
@@ -1261,6 +1262,7 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
if model_spec_key is None: if model_spec_key is None:
model_dict[key] = torch.empty(model_dict[key].shape, dtype=model_dict[key].dtype, device="meta") model_dict[key] = torch.empty(model_dict[key].shape, dtype=model_dict[key].dtype, device="meta")
utils.bar.update(1) utils.bar.update(1)
koboldai_vars.loaded_layers += 1
continue continue
storage_key = model_dict[key].key storage_key = model_dict[key].key
@@ -1313,6 +1315,7 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
np.empty(params["cores_per_replica"]), np.empty(params["cores_per_replica"]),
) )
koboldai_vars.loaded_layers += 1
utils.bar.update(1) utils.bar.update(1)
if utils.num_shards is not None and utils.current_shard < utils.num_shards: if utils.num_shards is not None and utils.current_shard < utils.num_shards: