如何将HuggingFace预训练Transformer模型直接加载至GPU？

当然可以直接在调用from_pretrained时把模型加载到GPU，以下是两种常用方案：

1. 使用device参数（适用于单GPU场景） #

调用from_pretrained时直接指定device='cuda'，模型会跳过CPU直接加载到GPU显存中：

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("bert-base-uncased", device='cuda')

2. 使用device_map参数（支持多GPU/自动分配） #

如果是多GPU环境，或者希望让库自动根据显存情况分配模型权重到可用设备，设置device_map='auto'即可：

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("bert-base-uncased", device_map='auto')

注意事项 #

• 以上方法要求你的环境已正确配置CUDA（GPU可用），若CUDA不可用，代码会自动回退到CPU加载。
• 针对较旧版本的transformers库，若device参数不可用，可结合torch_dtype和.to('cuda')优化加载流程，虽会加载少量元数据到CPU，但模型权重会直接写入GPU，能大幅节省CPU内存：

import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("bert-base-uncased", torch_dtype=torch.float32).to('cuda')

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