InternLM / InternLM/lmdeploy

AsyncEngine 的 stream_infer 函数增加手动传入session_id,实现多次调用 stream_infer 时的并行推理[Feature]

Open
#1,590 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
8.1k
Forks
748
Avg merge
6d 2h
Merged PRs (30d)
54

Description

### Motivation

我自己手动实例化pipe之后使用pipe的stream_infer实现模型推理,代码如下
```python
response = ""
for _response in self.pipe.stream_infer(
prompts = prompts,
gen_config = self.gen_config,
do_preprocess = True,
adapter_name = None
):
response += _response.text
yield response, history + [[query, response]]
```
之后自己搭建gradio界面进行推理,gradio实现了多线程访问,在使用 transformers 库进行推理时,实现了多个用户同时和模型对话,但是在使用 lmdeploy 部署时发现对话是串行的,同一时刻只能有一个用户进行对话。

排查后发现是 AsyncEngine 的 stream_infer 函数中在调用 generate 时传递的 session_id 为对话列表的 index,所以同一时刻使用这个函数时,多个对话的 session_id 都为0,因此变为了串行的方式推理。
这一行传递了固定的列表 index https://github.com/InternLM/lmdeploy/blob/main/lmdeploy/serve/async_engine.py#L495

我将传入的 i 修改为了 随机数,就可以实现并行推理了,不过这样修改似乎并不优雅,希望提供手动传入 session_id 的方式
```python
for i, prompt in enumerate(prompts):
generators.append(
self.generate(prompt,
random.randint(0, 1e9), # 这一行原本是 i,我改为了 random 模块的 randint
gen_config=gen_config[i],
stream_response=True,
sequence_start=True,
sequence_end=True,
do_preprocess=do_preprocess,
adapter_name=adapter_name,
**kwargs))
```

### Related resources

这是我实现推理的完整代码
https://github.com/NagatoYuki0943/xlab-huanhuan/blob/master/lmdeploy/turbomind_gradio_stream.py
https://github.com/NagatoYuki0943/xlab-huanhuan/blob/master/lmdeploy/infer_engine.py
### Additional context

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.