如何在Transformers中加载预训练模型并指定特定多GPU？

方法1：手动构建device_map字典 #

transformers的device_map支持传入自定义字典，手动指定每个模型层的GPU分配目标，完全满足指定特定GPU的需求：

• 先将模型加载到CPU，获取所有层的名称
• 按需求把层分配到目标GPU（比如cuda:4、5、6）
• 用自定义字典作为device_map参数重新加载模型

示例代码：

import torch
from transformers import LlamaForCausalLM

model_dir = '/models/Llama-2-13b-chat-hf'

# 先加载模型到CPU，提取所有层名称
model = LlamaForCausalLM.from_pretrained(model_dir, torch_dtype=torch.float32, device_map='cpu')
layer_names = [name for name, _ in model.named_parameters()]

# 把层平均分配到cuda:4、5、6
target_gpus = ['cuda:4', 'cuda:5', 'cuda:6']
layer_per_gpu = len(layer_names) // len(target_gpus)
device_map = {}

for i, gpu in enumerate(target_gpus):
    start_idx = i * layer_per_gpu
    # 最后一个GPU承接剩余所有层
    end_idx = start_idx + layer_per_gpu if i != len(target_gpus)-1 else len(layer_names)
    for name in layer_names[start_idx:end_idx]:
        device_map[name] = gpu

# 单独指定嵌入层和输出层的GPU（可按需调整）
device_map['model.embed_tokens'] = target_gpus[0]
device_map['lm_head'] = target_gpus[-1]

# 用自定义device_map加载模型
model = LlamaForCausalLM.from_pretrained(
    model_dir,
    device_map=device_map,
    torch_dtype=torch.float32
)

方法2：用accelerate限制可见GPU #

通过accelerate库限制当前进程仅能访问指定GPU，再用device_map='auto'让模型自动分配，这样只会用到你指定的GPU，操作更简洁：

示例代码：

import torch
from accelerate import PartialState
from transformers import LlamaForCausalLM

model_dir = '/models/Llama-2-13b-chat-hf'

# 限制进程可见GPU为4、5、6
with PartialState(devices=[4,5,6]):
    model = LlamaForCausalLM.from_pretrained(
        model_dir,
        device_map='auto',
        torch_dtype=torch.float32
    )

使用前需安装accelerate库：pip install accelerate

注意事项 #

• 确保指定的GPU有足够显存，避免出现OOM（显存不足）错误
• 对于Llama这类层显存占用均匀的模型，平均分配即可；若模型层显存差异大，可手动调整层的分配策略

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