ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

enable inference mode for deepspeed tp serving (#11742)

Browse source
This commit is contained in:
Shaojun Liu 2024-08-08 14:38:30 +08:00 committed by GitHub
parent 9e65cf00b3
commit 107f7aafd0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
python/llm/example/GPU/Deepspeed-AutoTP-FastAPI/serving.py
View file

@ -121,6 +121,8 @@ def load_model(model_path, low_bit):
current_accel = XPU_Accelerator() current_accel = XPU_Accelerator()
set_accelerator(current_accel) set_accelerator(current_accel)
model=model.eval()
# Move model back to xpu # Move model back to xpu
model = model.to(f"xpu:{local_rank}") model = model.to(f"xpu:{local_rank}")
model = BenchmarkWrapper(model) model = BenchmarkWrapper(model)