如何让EleutherAI/gpt-neo-1.3B模型运行在GPU上？

核心问题：未将模型与输入数据迁移至GPU #

你的代码默认将模型和张量加载到CPU，加上可能未确认PyTorch的CUDA支持，导致GPU未被利用。

1. 确认PyTorch的CUDA支持 #

先运行以下代码验证CUDA是否可用：

import torch
print(torch.cuda.is_available())

如果输出False，说明你安装的是CPU版PyTorch，需要卸载后重新安装适配CUDA 11.6的PyTorch版本。

2. 修改代码，强制使用GPU #

修改后的完整代码如下，核心是通过to(device)将模型和输入张量移至GPU：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 自动选择可用设备（优先GPU）
device = "cuda" if torch.cuda.is_available() else "cpu"

# 加载模型并迁移到指定设备
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B").to(device)
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")

prompt = "What is the capital of France?"

# 将输入张量也迁移到GPU
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
gen_tokens = model.generate(input_ids, do_sample=True, temperature=0.9, max_length=50)
gen_text = tokenizer.batch_decode(gen_tokens)[0]
print(gen_text)

3. 显存优化（可选） #

如果想进一步节省显存，可使用半精度加载模型，显存占用会减少近一半：

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neo-1.3B",
    torch_dtype=torch.float16
).to(device)

验证GPU使用情况 #

运行代码后，执行nvidia-smi命令，若GPU显存占用明显上升、GPU-Util数值增加，说明GPU已正常工作。

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