如何让Llama 2在NVIDIA A6000 GPU上运行而非CPU？

你当前代码未明确指定模型运行设备，导致模型默认加载到CPU，因此GPU占用始终为0%。以下是具体修改方案及完整代码：

核心修改点 #

• 加载模型时指定半精度数据类型+自动设备映射，让模型自动分配到GPU
• 确保文本生成pipeline绑定GPU设备

修改后的完整代码 #

import torch
import transformers
from transformers import LlamaForCausalLM, LlamaTokenizer
import setGPU

model_dir = "llama/llama-2-7b-chat-hf"

# 加载模型时指定半精度+自动设备映射，自动将模型分配到GPU
model = LlamaForCausalLM.from_pretrained(
    model_dir,
    torch_dtype=torch.float16,
    device_map="auto"
)
tokenizer = LlamaTokenizer.from_pretrained(model_dir)

# 初始化pipeline时指定设备为cuda（单GPU场景下也可用设备索引0）
pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    device=0
)

sequences = pipeline(
    'I wanna hear some news. What is up today',
    do_sample=True,
    top_k=10,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    max_length=400,
)

for seq in sequences:
    print(f"{seq['generated_text']}")

可选替代方案（手动指定设备） #

如果device_map="auto"不生效，可手动将模型移至GPU：

# 加载模型后手动迁移到cuda
model = LlamaForCausalLM.from_pretrained(model_dir, torch_dtype=torch.float16)
model = model.to("cuda")

验证方法 #

• 运行代码时，在终端执行nvidia-smi，可看到GPU显存占用和使用率上升
• 在代码中添加print(model.device)，输出应为cuda:0，证明模型已加载到GPU

内容的提问来源于stack exchange，提问作者Nasi Goreng Ayam