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 pull request #370 from Zurnaz/tpu_bfloat16

Browse Source 
fix: tpu support models saved in bfloat16
This commit is contained in:
henk717
2023-06-03 14:06:21 +02:00
committed by GitHub
parent 8008a1c0b0 18dc706944
commit 0ac2800601
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tpu_mtj_backend.py
View File

@@ -1255,7 +1255,11 @@ def load_model(path: str, model_type: str, badwordsids=koboldai_settings.badword
params["cores_per_replica"],
network.state["params"][spec["module"]][spec["param"]].shape,
)
tensor = jnp.array(tensor.detach())
tensor = tensor.detach()
# numpy does not support bfloat16
if tensor.dtype is torch.bfloat16:
tensor = tensor.to(torch.float32)
tensor = jnp.array(tensor)
if tensor.dtype is torch.float16 or tensor.dtype is torch.float32:
tensor = tensor.bfloat16()
network.state["params"][spec["module"]][spec["param"]] = move_xmap(