Merge pull request #180 from VE-FORBRYDERNE/patch

Only enable TPU transpose optimization if loading from HF model
This commit is contained in:
henk717 2022-11-21 20:02:14 +01:00 committed by GitHub
commit 04d9172fcd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions

View File

@ -1149,7 +1149,8 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
params[param] = default_params[param]
# 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
if vars.model == "TPUMeshTransformerGPTNeoX":