vllm 分布式批处理推理示例无法工作

j8ag8udp  于 6个月前  发布在  其他
关注(0)|答案(3)|浏览(42)

有一个从多个GPU进行批量推理的例子。
4abf633
然而,运行这段代码会抛出一个错误,因为llm示例不可序列化。
为了让它运行,我不得不修改代码如下:

class Predictor:
    def __init__(self):
        pass
 
    def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, list]:
        # prepare samples...
 
        llm = LLM(model="model name")
 
        # Generate translations from the prompts.
        outputs = llm.generate(prompts, sampling_params)
         
        del llm
        llm = None
        destroy_model_parallel()
         
        gc.collect()
 
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        if torch.backends.mps.is_available():
            torch.mps.empty_cache()
        torch.cuda.synchronize()
 
        # process...
 
        return result
hfyxw5xn

hfyxw5xn1#

cc @c21@zhe-thoughts

8wtpewkr

8wtpewkr2#

你好,seungduk-yanolja,你能分享一下完整的代码(当它不起作用时)和错误堆栈跟踪吗?谢谢。

不可序列化通常发生在对象在类之外创建的情况下。在这种情况下,LLM对象是在类__init__()内部创建的,因此不应该发生跨机器的序列化。

cidc1ykv

cidc1ykv3#

在这里需要注意的另一件事是——当将 LLMPredictor 传递给 Dataset.map_batches() 时,你应该传递类本身(代码)。
你是否顺便传递了对象 LLMPredictor()?这将触发头节点中对象的 __init__,并导致序列化到工作节点。

相关问题