rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-01-10 15:22:59 +01:00
Code Issues Projects Releases Wiki Activity

Only enable TPU transpose optimization if loading from HF model

Browse Source
This commit is contained in:
vfbd 2022-11-21 13:47:18 -05:00
parent f2077b8e58
commit 9a3f0eaab2
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tpu_mtj_backend.py
View File

@ -1149,7 +1149,8 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
params[param] = default_params[param] params[param] = default_params[param]
# Use an optimization that will allow us to avoid one extra transpose operation # Use an optimization that will allow us to avoid one extra transpose operation
params["transposed_linear"] = True if hf_checkpoint:
params["transposed_linear"] = True
# Load tokenizer # Load tokenizer
if vars.model == "TPUMeshTransformerGPTNeoX": if vars.model == "TPUMeshTransformerGPTNeoX":