device_map配置指定显卡

device_map配置指定显卡

在Transformers库中，device_map参数用于指定模型中的各个部分应该加载到哪个具体的计算设备上，这对于实现资源的有效分配和利用非常有用。以下是关于如何配置device_map以指定使用特定显卡的详细说明：

1. 配置device_map

• 使用device_map字典：

  device_map = {
      "transformer.h.0": "cuda:0", # 第一部分放在GPU 0
      "transformer.h.1": "cuda:1", # 第二部分放在GPU 1
      #... 根据模型结构继续分配
  }
  model = AutoModelForCausalLM.from_pretrained(model_dir, device_map=device_map)
  

• 使用os.environ['CUDA_VISIBLE_DEVICES']：

  os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'
  model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map="auto")
  

2. 不同device_map选项

• auto``和balanced``： 在所有的GPU上平衡切分模型，主要是有可能发现更高效的分配策略。

  • balanced参数的功能则保持稳定，但可能不是最佳选择。

• `balanced_low_0``： 在除了第一个GPU上的其它GPU上平衡划分模型，并且在第一个 GPU 上占据较少资源。这个选项符合需要在第一个 GPU 上进行额外操作的需求，例如需要在第一个 GPU 执行 generate 函数（迭代过程）。

• `sequential``： 按照GPU的顺序分配模型分片，从 GPU 0 开始，直到最后的 GPU。但是，GPU 0 可能会直接爆显存。

3. 注意事项

• 确保device_map在import torch之前设置： 如果device_map在import torch之后设置，它可能无法生效。这是因为import torch会优先执行，可能导致device_map设置被忽略。

• 确保显卡的可用性： 在配置device_map之前，确保指定的显卡是可用的。例如，在运行之前，请确保没有其他程序占用指定的显卡。

4. 示例代码

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

# 加载模型
model_path = "./model/chatglm2-6b"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map="auto")

# 输入文本
text = '什么是机器学习。'
inputs = tokenizer(text, return_tensors="pt")

# 输出结果
outputs = model.generate(**inputs, max_new_tokens=50