前回、NVIDIAのGPUドライバーのインストール

のところまで書いたので続きでドライバーインストールする済のUbuntu24.04内でGPUを使ったPYTHON環境の構築方法を記載していきます。ここからはWSLもローカル、クラウド型の純正UBUNTUも同じ手順です。
まずはOSを最新にしておきます

# apt update
# apt upgrade
# reboot

CUDAのインストール（公式サイトよりご自身の環境に適合するインストール方法を特定します）

https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=24.04&target_type=deb_local

混乱を避けるため、どこかに作業フォルダーを作り作業します。
そのあと、以下の手順で実行していきます。

# wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-ubuntu2404.pin
# mv cuda-ubuntu2404.pin /etc/apt/preferences.d/cuda-repository-pin-600
# wget https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda-repo-ubuntu2404-12-6-local_12.6.3-560.35.05-1_amd64.deb
# dpkg -i cuda-repo-ubuntu2404-12-6-local_12.6.3-560.35.05-1_amd64.deb
# cp /var/cuda-repo-ubuntu2404-12-6-local/cuda-*-keyring.gpg /usr/share/keyrings/
# apt update
# apt -y install cuda-toolkit-12-6

導入後PATHを通すためそれぞれのログインユーザー　.bashrc　などにPATHを追加することでそのユーザーがCUDAにアクセスできるようになりますが、マシン全体に適用したかったので/etc/environment や /etc/profileに追加していきます。

/etc/environment
PATHの最後に:/usr/local/cuda/binを追加
#LD_LIBRARY_PATH にも追加

PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/usr/local/cuda/bin"
LD_LIBRARY_PATH="/usr/local/cuda/lib64:$LD_LIBRARY_PATH"

/etc/profile
最後に以下追加

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

再起動して、GPUを利用して計算する実行ユーザー環境でnvccコマンドで確認します
以下のような出力が出れば問題ありません。

# nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Tue_Oct_29_23:50:19_PDT_2024
Cuda compilation tools, release 12.6, V12.6.85
Build cuda_12.6.r12.6/compiler.35059454_0

CUDNNのインストール（公式サイトよりご自身の環境に適合するインストール方法を特定します）

https://developer.nvidia.com/cudnn-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=24.04&target_type=deb_local

同様に推奨された方法でインストールを実施します

# wget https://developer.download.nvidia.com/compute/cudnn/9.6.0/local_installers/cudnn-local-repo-ubuntu2404-9.6.0_1.0-1_amd64.deb
# dpkg -i cudnn-local-repo-ubuntu2404-9.6.0_1.0-1_amd64.deb
# cp /var/cudnn-local-repo-ubuntu2404-9.6.0/cudnn-*-keyring.gpg /usr/share/keyrings/
# apt update
# apt -y install cudnn-cuda-12

インストール後一応インストールが行われたか確認しておきましょう。

# cat /usr/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
#define CUDNN_MAJOR 9
#define CUDNN_MINOR 6
#define CUDNN_PATCHLEVEL 0
--
#define CUDNN_VERSION (CUDNN_MAJOR * 10000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL)

/* cannot use constexpr here since this is a C-only file */

PYTHON, PYTORCHのインストールとCUDAの認識テスト
　仮想環境をインストールする場所や実行ファイルの場所などは整理しておくといいかもしれません。

# tree -L 2
.
├── cutest
├── llm_test
│   ├── benchmark.py
│   ├── elyza-tasks-100.jsonl
│   ├── llm-jp-3-13b-finetune-22_lora_output.jsonl
│   ├── models
│   └── suiron.py
└── pyenv
    ├── cutest
    └── llm_test

早速PYTHONをインストールをして仮想環境を作成します。（ANACONDAではなく今回はUbuntuの標準レポジトリトリでPYTHON３を導入しました）

# apt install python3-pip python3-venv
# cd pyenv
# python3 -m venv cutest
# source cutest/bin/activate

早速PYTOCHをいれようとしましたが、CUDA 12.6の記載はPYTHONもCONDAもなかったです。なので一番近い12.4を選択してインストールしました。
https://pytorch.org/get-started/locally/

あとは指示どうりインストールして、チェックをします。

(cutest)# pip install torch torchvision torchaudio
(cutest)# python -c "import torch; print(torch.cuda.is_available())"
True

TrueとなればGPUは認識されていることになります。
さらに、少し長い計算コードを実行してGPUが使われていることを確認します。

cutest.py

import torch
import torch.nn as nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 1)

    def forward(self, x):
        return self.linear(x)

model = SimpleModel().to(device)
model.eval()

input_tensor = torch.randn(1, 10).to(device)

with torch.no_grad():
    for _ in range(100000):  # 推論回数を多くして観察できるようにする
        output = model(input_tensor)

実行

(cutest)# python cutest.py

別のTERMINALでnvidia-smiで使われている様子を確認(PID 2333で使用されていることが確認できます)

# nvidia-smi
Tue Dec 17 18:27:45 2024
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.183.01             Driver Version: 535.183.01   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                       Off | 00000000:00:1E.0 Off |                    0 |
| N/A   20C    P0              27W /  70W |    139MiB / 15360MiB |     32%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      2333      C   python                                      134MiB |
+---------------------------------------------------------------------------------------+