github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 04:24:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

[infer] fix test bug (#4838)

Browse Source 
* fix test bug

* delete useless code

* fix typo
This commit is contained in:
Xu Kai
2023-10-04 10:01:03 +08:00
committed by GitHub
parent 013a4bedf0
commit d1fcc0fa4d
6 changed files with 56 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
examples/inference/bench_llama.py
View File

@@ -43,7 +43,6 @@ def run_llama_test(args):
tokenizer.pad_token_id = tokenizer.unk_token_id
model = LlamaForCausalLM.from_pretrained(llama_model_path, pad_token_id=tokenizer.eos_token_id)
model = model.half()
model_config = model.config
shard_config = ShardConfig(enable_tensor_parallelism=True if args.tp_size > 1 else False, inference_only=True)