问题描述

下面是这个问题的一个简化demo：

import torch
from transformers import AutoTokenizer,AutoModelForCausalLM
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '2,3' 

model_path = "..."

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

tokenizer =  AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

按理说指定了环境变量为后面几张显卡，模型会自动加载到我们指定的显卡上。可是查看gpu使用情况，发现还是加载在了第一张显卡上，这是为啥呢？

原因及解决方案

参考这篇博客：【踩坑】PyTorch中指定GPU不生效和GPU编号不一致问题 ，发现如果在代码里指定CUDA_VISIBLE_DEVICES的话，指定的位置一定要在所有引用torch的代码之前。上面是单个文件的例子，如果是多个文件互相引用，也是一样的道理。所以推荐在脚本入口的一开始就设置该环境变量：

# 在入口文件的一开始设置该环境变量。原则是它的设置一定要早于所有模块对torch的引用
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '2,3' 
import torch
from transformers import AutoTokenizer,AutoModelForCausalLM


model_path = "..."

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

tokenizer =  AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

再查看显卡，模型会被自动加载到编号为2,3的两张显卡上。值得注意的是，若指定device_map = "auto"，需要在环境中安装accelerate包。