在AMD GPU上运行Google的Gemma2：详细指南

在本文中，我将分享使用AMD RX7900XTX GPU运行Google的大型语言模型Gemma2的经验。从环境设置到实际执行，以下是完整的操作步骤。

环境设置与准备

首先，我安装了ROCm兼容的PyTorch，这是充分利用AMD GPU所需的关键组件。安装完成后，我接着安装了Hugging Face的Transformers库，这使得我们能够轻松访问Gemma2模型。

pip install -U transformers

安装完成后，我继续从Hugging Face获取访问令牌。为此，我访问了以下链接，阅读并同意使用条款后，获取了我的令牌：

https://huggingface.co/google/gemma-2-2b

接下来，通过Jupyter Notebook使用令牌登录：

from huggingface_hub import login
login("your-token-here")

这样，我就可以通过Hugging Face访问模型了。

模型运行

一切准备就绪后，我运行了以下代码来加载模型并生成文本：

import torch
from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="google/gemma-2-9b",
    device="cuda",  # 使用ROCm的AMD GPU时仍然使用"cuda"
    torch_dtype=torch.float16,
)

text = "你知道蓝猫公司吗？"
outputs = pipe(text, max_new_tokens=256)
response = outputs[0]["generated_text"]
print(response)

执行结果

运行代码后，我得到了以下输出：

输入文本：

你知道蓝猫公司吗？

生成的输出：

蓝猫公司成立于2015年10月，以“让与猫的生活更加有趣”为理念，开发和销售支持猫咪生活的产品。

生成输出的时间约为12秒。感谢AMD RX7900XTX的高性能处理能力，能够在相对较短的时间内得到结果。

结论

这次实验展示了使用AMD GPU可以顺利运行Google的Gemma2模型。通过安装ROCm兼容的PyTorch，即使在没有CUDA支持的情况下，AMD GPU也能发挥出色的性能。生成的输出质量也令人满意，显示出进一步应用的巨大潜力。

在AMD GPU上运行Google的Gemma2：详细指南

Yuichiro Minato