模型生成器仅在一个 GPU 上加载，导致 CUDA 内存不足错误

Question

我正在 8 个 Nvidia A100 GPU（1 个节点）上运行 LLM 代码。
当我尝试加载大模型 (70B) 时，出现

CUDA out of memory

错误。

from haystack.components.generators import HuggingFaceLocalGenerator
generator = HuggingFaceLocalGenerator(model="meta-llama/Meta-Llama-3-70B-Instruct",
                                      generation_kwargs={"max_new_tokens": 100,
                             # tried with and without the device_map set to auto !
                                                         "device_map": "auto"})

然后，当运行

query_pipeline.run()

时，模型正在加载，我可以在

nvidia-smi

输出中看到（每 3 秒自动运行一次），开始时所有 8 个 GPU 都是空闲的（使用了 1-2 MB），然后只有第一个开始填满，其他的一直空着，直到最后压碎。

您能否帮助解决如何确保模型在所有 GPU 上加载？

Answer 1

谢谢，@Stefano Fiorucci - anakin87，您的建议解决了问题。 运行发电机的正确方法是：

generator = HuggingFaceLocalGenerator(model="meta-llama/Meta-Llama-3-70B-Instruct",
                                      generation_kwargs={"max_new_tokens": 100},
                                      huggingface_pipeline_kwargs={"device_map": "auto" 
                                                                   })

这样，我可以在

nvidia-smi

输出中看到负载分布在所有 GPU 之间（所有 GPU 的负载均达到 ±50%），并且模型运行时不会出现崩溃情况。