如何配置device_map使模型在指定的两个GPU设备上运行？

如果只能使用服务器上的两个GPU，有两种常用配置方式：

方式一：环境变量限制可见GPU + 自动分配（推荐） #

先通过环境变量指定程序能访问的GPU，再让transformers自动把模型拆分到这两个GPU上，操作最简单且无需了解模型结构：

import os
# 替换成你有权限使用的两个GPU编号，比如要用到cuda:2和cuda:3就写"2,3"
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"

from transformers import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",  # 自动将模型层均衡分配到可见的两个GPU上
)

设置CUDA_VISIBLE_DEVICES后，程序会把指定的GPU识别为0和1号，device_map="auto"会自动拆分模型权重到这两个GPU上，充分利用显存资源。

方式二：手动指定层的GPU分配（精细控制） #

如果不需要修改环境变量，可以手动给模型的不同层分配目标GPU，适合需要精准控制层位置的场景：

from transformers import AutoTokenizer, AutoModelForCausalLM

# 根据模型实际层结构，指定各层对应的GPU（可通过print(model)查看层命名）
device_map = {
    "transformer.h.0": "cuda:2",
    "transformer.h.1": "cuda:2",
    "transformer.h.2": "cuda:3",
    "transformer.h.3": "cuda:3",
    "transformer.ln_f": "cuda:3",
    "lm_head": "cuda:3"
}

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map=device_map,
)

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