AMD GPUでGoogleのGemma2を実行した

Yuichiro Minato

2024/08/12 04:29

AMD GPUでGoogleのGemma2を実行した体験記

今回は、AMDのRX7900XTXを使って、Googleの大規模言語モデル「Gemma2」を実行してみました。手順から結果までご紹介します。

環境構築と準備

まず、ROCm対応のPyTorchをインストールします。これは、AMDのGPUを活用するために必要なライブラリです。インストールが完了したら、次にHugging FaceのTransformersライブラリをインストールします。これにより、Gemma2モデルを簡単に利用できるようになります。

pip install -U transformers

インストールが完了したら、次にHugging Face経由でモデルを利用するための手続きを行います。以下のリンクから利用規約に同意し、アクセストークンを取得します。

次に、取得したアクセストークンを使ってJupyter Notebook上でログインします。

from huggingface_hub import login
login("あなたのトークン")

これで、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でも高性能なモデルを動かすことができる点が非常に有用でした。出力の質も満足できるものであり、今後もさらに応用が広がると感じました。