Gemma 2 日本語特化モデルをGPUで動かす〜 速度比較編〜

はじめに

前回、Gemma 2 2B 日本語特化モデルをローカルで動かす〜キャッシュ活用編〜 でCPU版の動作確認をしましたが、今回はその続きとして、実際に組んでいる「ニュース編集AIエージェント」を GPU版として構築し直し、CPU版と速度を比較してみました。

このエージェントは、interviewフォルダに置いた記事をGemma 2 2B jpn-itで要約し、構造化JSONとしてnewspaperフォルダに蓄積保存する、というものです。前回までは動作確認レベルでしたが、今回は実運用を見据えた本体スクリプトをそのままGPU化しています。

今回もCPUのみ環境と同じく、あくまで備忘録的な立ち位置のメモです。

検証環境

項目	スペック
OS	Windows 11
GPU	NVIDIA GeForce RTX 4060 Ti(VRAM 8GB)
ドライババージョン	591.86
ドライバ対応CUDA	13.1(nvidia-smi 表示)
CUDA Toolkit(nvcc)	11.8
torch	2.5.1+cu121(torch.version.cuda は 12.1)
Python	venv による仮想環境(venv_gpu)
モデル	google/gemma-2-2b-jpn-it(ローカルキャッシュ済み)

nvidia-smi の「CUDA Version」はドライバが対応する最大バージョンであり、実際にインストールされているCUDA Toolkit(nvcc --version で11.8)やtorchがビルドされたCUDAバージョン(+cu121 = 12.1)とは別物、という点は混同しやすいので注意が必要でした。今回のように3つとも数字が違う、というのは普通に起こるようです。

torch側から見えているGPU情報は以下のスクリプトで確認しました。

import torch
print(torch.__version__)
print(torch.version.cuda)
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))
print(torch.cuda.get_device_properties(0).total_memory / 1024**3, "GB")

2.5.1+cu121
12.1
True
NVIDIA GeForce RTX 4060 Ti
7.99560546875 GB

GPU版torchのインストール

pip install torch --index-url https://download.pytorch.org/whl/cu121

CPU版のときは「CUDA関連の巨大ファイルを避けるためCPU版を明示的に指定」しましたが、GPU版では逆に、CUDA対応のビルドを明示的に指定してインストールする形になります。

つまずいたポイント: Anacondaとの衝突

環境構築時、AnacondaがインストールされているとPATHの優先順位の関係で venv 側のtorchではなくAnaconda側のPython/pipを掴んでしまうことがありました。CPU版のときは意識していなかったポイントですが、GPU版でCUDAが正しく認識されない原因になり得るため、今回は明示的に Anaconda環境とは別に venv_gpu という専用の仮想環境を切って作業しました。同じ問題に当たった場合は、where python でどちらの実行ファイルを掴んでいるか確認するのがおすすめです。

コードの変更点(CPU版 → GPU版)

前回までの動作確認スクリプトから、実運用スクリプトnews_editor_agent.pyにステップアップしたタイミングで、以下の変更を行いました。

変更点	CPU版	GPU版
デバイス指定	device_map="cpu"	device_map="cuda"
dtype	dtype=torch.bfloat16	dtype=torch.float16
推論の決定性・繰り返し抑制	なし	do_sample=False + repetition_penalty=1.05
KVキャッシュ	未指定	use_cache=True

float16 にしたのは、多くのNVIDIA GPUで float16 の演算スループットが bfloat16 より最適化されているケースが多いためです(CPU版で bfloat16 を使っていたのは、CPU側の対応状況・数値安定性を優先した結果でした)。

推論時間の計測コードを追加

GPU化の効果を数字で見たかったので、run_inference() に計測処理を追加しました。

# --- 計測開始 -----------------------------------------
# CUDA使用時はカーネル起動が非同期のため、synchronize()で
# 実際にGPU上の計算が完了するタイミングまで待ってから計測する。
if torch.cuda.is_available():
    torch.cuda.synchronize()
start_time = time.perf_counter()

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        repetition_penalty=1.05,
        use_cache=True,
        pad_token_id=tokenizer.eos_token_id,
    )

if torch.cuda.is_available():
    torch.cuda.synchronize()
elapsed_time = time.perf_counter() - start_time
# --- 計測終了 -----------------------------------------

generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
generated_token_count = generated_ids.shape[0]

print(
    f"推論時間: {elapsed_time:.2f}秒"
    f"（生成トークン数: {generated_token_count}、"
    f"{generated_token_count / elapsed_time:.2f} tok/s）"
)

ポイントは torch.cuda.synchronize() です。CUDAのカーネル実行は非同期なので、これを入れずに時間を測ると「Python側の処理は終わっているが、実際はまだGPU上で計算中」という状態のまま計測してしまい、実態より短い時間が出てしまいます。CPU実行時は torch.cuda.is_available() が False になるため、この処理は自動的にスキップされ、CPU版・GPU版で同じコードのまま公平に比較できます。

実行結果

同じinterview記事(リコーのGENIAC採択に関するプレスリリース記事、約830文字)を、GPU版・CPU版それぞれで1回ずつ実行しました。

環境	推論時間	生成トークン数	速度
GPU(RTX 4060 Ti)	4.81秒	139	28.90 tok/s
CPU	22.65秒	139	6.14 tok/s

約4.7倍の高速化という結果になりました。同じトークン数を生成しているので、秒数の比率とtok/sの比率がほぼ一致しており、素直な比較ができています。

RAM 8GBの低スペックPCでCPU実行に20秒以上かかっていたことを思うと、VRAM 8GBのミドルクラスGPUに載せ替えるだけでここまで差が出るのは体感としても分かりやすい結果でした。

メモ

• dtype=torch.float16 + device_map="cuda" の組み合わせで、VRAM使用量は問題なく8GB内に収まった(モデル自体が2Bと小さいため)
• CUDAの非同期実行を考慮せずに計測すると、GPU側の数値が不当に良く出てしまう。速度比較記事を書く際は要注意
• nvidia-smi のCUDA Version・nvcc --version・torch.version.cuda の3つはそれぞれ意味が違うため、記事や手順書に書く際はどれを指しているか明記したほうが親切

注意点メモ

事象	原因	対処
GPUが認識されず torch.cuda.is_available() が False になる	Anacondaとの競合でPATHがconda側を優先してしまう	Anacondaとは別に専用の venv(venv_gpu)を作成して作業する
GPU側の推論時間が実態より短く出る	CUDAカーネルの非同期実行を考慮せずに計測している	model.generate() の前後で torch.cuda.synchronize() を挟む
JSON解析失敗 がGPU/CPU両方で発生	プロンプトテンプレート側の出力フォーマットの問題(デバイス起因ではない)	生成テキストをそのまま出力して、JSON前後の余計な文字列や途中で途切れていないか確認する(次回TODO)

次回TODO

• 今回発生した JSON解析失敗 の原因調査(プロンプトテンプレートの見直し、max_new_tokens の調整など)
• 複数記事でのバッチ処理時のGPUメモリ推移の確認
• 将来的なマルチエージェント構成(オーケストレーターによる複数ローカルAIの統括)への組み込み検討