有一个从多个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