PytorchがCUDA12.4に対応してたから環境構築して試す

いつの間にかPyTorchがCUDA12.4に対応していた

この前まで12.1までだった気がします。

作成した環境

• WSL2
• Dockerイメージ: nvidia/cuda:12.4.0-base-ubuntu22.04

手順

• CUDA12.4.0が入っているコンテナを作成してコンテナに入る

docker run -itd --gpus all --name sandbox1240 nvidia/cuda:12.4.0-base-ubuntu22.04
docker exec -it sandbox1240 /bin/bash

• 諸々アップデートしてPythonをインストール

apt update
apt upgrade -y
apt install python3.11 python3.11-venv -y

• 仮想環境を作成してActivateし、pipをアップグレード

python3.11 -m venv .env
source .env/bin/activate
python -m pip install -U pip

• ここでPyTorchをインストール

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

• PyTorchでGPUを使えているか確認

import torch
torch.cuda.is_available()
# -> True
torch.Tensor([0, 1, 2]).to('cuda')
# -> tensor([0., 1., 2.], device='cuda:0')

ちゃんと動いたっぽいです。

ローカルLLM(Gemma-1.1-2b-it)の動作確認

pip install transformers accelerate

import os
import getpass
os.environ['HF_TOKEN'] = getpass.getpass('HF_TOKEN:')

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = 'google/gemma-1.1-2b-it'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map='cuda',
)
chat = [
    {'role': 'user', 'content': '徳川幕府の歴代将軍を表形式で教えてください。'},
]
input_ids = tokenizer.apply_chat_template(chat, tokenize=True, return_tensors='pt').to(model.device)
outputs = model.generate(input_ids, max_new_tokens=1000)
print(tokenizer.decode(outputs[0]))

出力

<bos><start_of_turn>user
徳川幕府の歴代将軍を表形式で教えてください。<end_of_turn>
* 1333年 - 1335年：足利義昭
* 1335年 - 1346年：足利義元
* 1346年 - 1351年：足利義昭
* 1351年 - 1368年：足利義元
* 1368年 - 1374年：足利義昭
* 1374年 - 1380年：足利義元
* 1380年 - 1382年：足利義昭
* 1382年 - 1385年：足利義元
* 1385年 - 1394年：足利義昭<eos>

出力内容はともかくちゃんと動きました。

CUDA12.4.1でも動作確認

コンテナイメージがCUDA12.4.0のほかにCUDA12.4.1がありましたので、こちらでも試してみました。
結果としては、CUDA12.4.1でもちゃんと動作しました。

環境構築のや実行の方法は12.4.0のときとほぼ同じで、コンテナ作成のnvidia/cuda:12.4.0-base-ubuntu22.04をnvidia/cuda:12.4.1-base-ubuntu22.04に変えるだけです。

まとめ

いまだにCUDA12.1でPyTorchを動かしている人は、頃合いを見計らってCUDA12.4に乗り換えましょう。