使用模型并行性在多个 GPU 上加载 HuggingFace 模型进行推理

Question

我可以使用 6 个 24GB GPU。当我尝试加载一些 HuggingFace 模型时，例如以下

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
tokenizer = AutoTokenizer.from_pretrained("google/ul2")
model = AutoModelForSeq2SeqLM.from_pretrained("google/ul2")

我收到内存不足错误，因为该模型似乎只能在单个 GPU 上加载。然而，虽然整个模型无法装入单个 24GB GPU 卡，但我有 6 个 GPU 卡，并且想知道是否有一种方法可以将模型加载分布到多个卡上，以执行推理。

HuggingFace 似乎有一个网页，他们解释了如何执行此操作，但截至目前还没有任何有用的内容。

Answer 1

当您使用 加载模型时from_pretrained()，您需要指定要将模型加载到哪个设备。因此，添加以下参数，变压器库将处理其余的事情：

model = AutoModelForSeq2SeqLM.from_pretrained("google/ul2", device_map = 'auto')

此处传递“auto”将按照以下优先级顺序自动在硬件上分割模型：GPU > CPU (RAM) > 磁盘。

当然，这个答案假设您已经安装了 cuda 并且您的环境可以看到可用的 GPU。nvidia-smi从命令行运行将确认这一点。如果您遇到其他问题，请报告。