【解决方案】使用transformer指定显卡后,模型依然加载到默认第1张显卡上
问题描述
下面是这个问题的一个简化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,AutoModelForCausalLMmodel_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
包。