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

[Inference] Fix Inference Generation Config and Sampling (#5710)

Browse Source 
* refactor and add

* config default values

* fix gen config passing

* fix rpc generation config
This commit is contained in:
Yuanheng Zhao
2024-05-19 15:08:42 +08:00
committed by GitHub
parent 8bcfe360fd
commit 283c407a19
6 changed files with 124 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/inference/core/rpc_engine.py
View File

@@ -257,7 +257,12 @@ class RPCInferenceEngine(InferenceEngine):
assert len(self.workers) == self.tp_size, "init workers first"
init_tasks = [
self.async_parallel_wrapper(worker.execute_model_forward, input_token_ids, input_meta_data.to_rpc_param())
self.async_parallel_wrapper(
worker.execute_model_forward,
input_token_ids,
input_meta_data.to_rpc_param(),
self.generation_config_dict,
)
for worker in self.workers
]
ret = await asyncio.gather(*init_tasks)