github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix test llama (#4884)

pull/4889/head
Xu Kai 2023-10-11 17:43:01 +08:00 committed by GitHub
parent 08a9f76b2f
commit fdec650bb4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 0 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tests/test_infer/test_llama_infer.py
View File

@ -38,7 +38,6 @@ def run_llama_test(test_config):
enable_tensor_parallelism=True if test_config["tp_size"] > 1 else False, inference_only=True
)
infer_engine = TPInferEngine(model, shard_config, BATCH_SIZE, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
init_to_get_rotary(model.model, base=10000)
generate_kwargs = dict(max_new_tokens=MAX_OUTPUT_LEN, do_sample=False)
input_tokens = {