有一个从多个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
3条答案
按热度按时间hfyxw5xn1#
cc @c21@zhe-thoughts
8wtpewkr2#
你好,seungduk-yanolja,你能分享一下完整的代码(当它不起作用时)和错误堆栈跟踪吗?谢谢。
不可序列化通常发生在对象在类之外创建的情况下。在这种情况下,
LLM
对象是在类__init__()
内部创建的,因此不应该发生跨机器的序列化。cidc1ykv3#
在这里需要注意的另一件事是——当将
LLMPredictor
传递给Dataset.map_batches()
时,你应该传递类本身(代码)。你是否顺便传递了对象
LLMPredictor()
?这将触发头节点中对象的__init__
,并导致序列化到工作节点。