Huggingface加载模型后，如何实现类似device_map="auto"的多GPU分配

问题背景 #

我通过以下代码从Hugging Face加载大语言模型：

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", trust_remote_code=True)

加载完成后，我对模型内部层做了修改，还新增了一些自定义层，但启动微调训练时发现，自定义层并没有像原模型层那样自动分布到多GPU上，导致部分层设备不一致。尝试直接设置model.device_map = "auto"没有效果，需要找到加载模型后重新实现device_map="auto"自动分配的方法。

解决方案 #

方法1：使用model.to(device_map="auto")（推荐，transformers≥4.28.0） #

修改完自定义层后，直接调用to方法并指定device_map="auto"，框架会重新计算各层的设备分配并迁移：

# 加载模型并完成自定义层修改后执行
model = model.to(device_map="auto")

方法2：使用accelerate库的dispatch_model #

如果你的transformers版本较低，可以借助accelerate库的工具重新分配设备：

from accelerate import dispatch_model
from accelerate.utils import get_balanced_memory

# 加载模型并完成自定义层修改后执行
# 计算各层内存占用，生成平衡的设备映射
max_memory = get_balanced_memory(model)
model = dispatch_model(model, device_map="auto", max_memory=max_memory)

注意事项 #

• 修改自定义层时，确保新层参数默认处于可训练状态（若手动设置requires_grad=False需调整）
• 重新分配设备后，可通过打印model.device_map查看各层设备分配情况，确认自定义层已被正确分配

内容的提问来源于stack exchange，提问作者Omar