はじめに

私は機械学習用にRTX3090のデスクトップPCを1台持っており、定期的にメンテナンスを行なっています。
今回、RTX3090環境へUbuntu24.04をインストールしたので、その作業内容をメモとして公開します。

前回は基本的なRTX3090環境でのUbuntu24.04セットアップを行いました。
今回はNVIDIAの環境を整えます。

なお、関連する記事は下記の通りです。

RTX3090環境へのUbuntu24.04セットアップ
(本記事)RTX3090/Ubuntu24.04でNVIDIA環境をセットアップ
Githubの初期設定
画像認識の環境構築とVSCodeの設定
私のVSCodeのカスタマイズ

目次

ドライバのインストール
CUDA Toolkitのインストール
pytorchのインストール
最後に
参考文献

ドライバのインストール

nouveauを無効化

$ echo 'blacklist nouveau' | sudo tee -a /etc/modprobe.d/blacklist-nouveau.conf
$ echo 'options nouveau modeset=0' | sudo tee -a /etc/modprobe.d/blacklist-nouveau.conf
$ cat /etc/modprobe.d/blacklist-nouveau.conf
$ sudo update-initramfs -u

グラフィックボードが認識されているか確認

$ lspci | grep -i nvidia
01:00.0 VGA compatible controller: NVIDIA Corporation GA102 [GeForce RTX 3090] (rev a1)
01:00.1 Audio device: NVIDIA Corporation GA102 High Definition Audio Controller (rev a1)

nvidia ドライバのRecommendを確認

$ ubuntu-drivers devices
 	== /sys/devices/pci0000:00/0000:00:01.0/0000:01:00.0 ==
modalias : pci:v000010DEd00002204sv00001569sd0000F278bc03sc00i00
vendor   : NVIDIA Corporation
model    : GA102 [GeForce RTX 3090]
driver   : nvidia-driver-535-open - distro non-free
driver   : nvidia-driver-535-server-open - distro non-free
driver   : nvidia-driver-555-open - third-party non-free
driver   : nvidia-driver-470-server - distro non-free
driver   : nvidia-driver-535-server - distro non-free
driver   : nvidia-driver-470 - distro non-free
driver   : nvidia-driver-560-open - third-party non-free
driver   : nvidia-driver-550-open - third-party non-free
driver   : nvidia-driver-560 - third-party non-free recommended
driver   : nvidia-driver-535 - distro non-free
driver   : nvidia-driver-555 - third-party non-free
driver   : nvidia-driver-550 - third-party non-free
driver   : xserver-xorg-video-nouveau - distro free builtin

今入っているドライバのバージョンを確認

$ dpkg -l | grep nvidia

ここで推奨ではない古いドライバが入っていた場合
- 古いドライバの削除
```
$ sudo apt --purge remove nvidia-*
$ sudo apt autoremove
```
- 推奨されているドライバをインストール
```
$ sudo apt -y install nvidia-driver-560
```
- 再起動
```
$ reboot
```

GPUのサマリを表示

 $ nvidia-smi
 +-----------------------------------------------------------------------------------------+
 | NVIDIA-SMI 560.28.03              Driver Version: 560.28.03      CUDA Version: 12.6     |
 |-----------------------------------------+------------------------+----------------------+
 | 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  NVIDIA GeForce RTX 3090        Off |   00000000:01:00.0 Off |                  N/A |
 |  0%   38C    P8             28W /  370W |      15MiB /  24576MiB |      0%      Default |
 |                                         |                        |                  N/A |
 +-----------------------------------------+------------------------+----------------------+
                                                                                          
 +-----------------------------------------------------------------------------------------+
 | Processes:                                                                              |
 |  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
 |        ID   ID                                                               Usage      |
 |=========================================================================================|
 |    0   N/A  N/A      1842      G   /usr/lib/xorg/Xorg                              4MiB |
 +-----------------------------------------------------------------------------------------+

以上でドライバのインストールが完了しました。
次に、CUDA Toolkitをインストールします。

CUDA Toolkitのインストール

CUDA 12.6のTookitを確認する
- こちらで自分の環境を選択し、実行手順を確認
aptに追加するキーをダウンロード（Ubuntu24.04用のアドレスであることに注意）

$ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb

キーをaptに追加

$ sudo dpkg -i cuda-keyring_1.1-1_all.deb
$ sudo apt update

CUDA Toolkitのインストール

$ sudo apt -y install cuda-toolkit-12-6

CUDA Toolkitのパスを.bashrcに追加（これが通っていないと、この後のnvccコマンドが通りません）

$ echo "export PATH=/usr/local/cuda:/usr/local/cuda/bin:$PATH" >> .bashrc
$ echo "export LD_LIBRARY_PATH=/usr/local/lib:/usr/local/cuda/lib64:$LD_LIBRARY_PATH" >> .bashrc
$ source ~/.bashrc

パスが通ったか確認（バージョンが表示されればOK）

$ nvcc -V
vcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Fri_Jun_14_16:34:21_PDT_2024
Cuda compilation tools, release 12.6, V12.6.20
Build cuda_12.6.r12.6/compiler.34431801_0

pytorchのインストール

venvの環境を作る

今回は環境名をtestとします。

pythonのバージョンを確認し、venvをインストール

$ python3 -V
Python 3.12.3
$ sudo apt install python3.12-venv

workディレクトリを作成し、その中に環境名のディレクトリを作って移動

$ mkdir work
$ mkdir work/test
$ cd work

仮装環境を作成

$ python3 -m venv test

アクティベート

$ source test/bin/activate

仮装環境testにインストール

$ pip install --upgrade pip
$ pip install torch torchvision torchaudio

GPUが使える状態になっているか確認(Trueが帰って来ればOK)

$ python3
Python 3.12.3 (main, Jul 31 2024, 17:43:48) [GCC 13.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.cuda.is_available())
True

GPU確認に続けて、計算をしてみる(下記の通り'cuda:0'表示されればOK)

>>> array = torch.zeros(4) #長さ4の配列を定義（これはcpu配列）
>>> print(array)
tensor([0., 0., 0., 0.])
>>> array_gpu = array.cuda()
>>> print(array_gpu)
tensor([0., 0., 0., 0.], device='cuda:0')
>>> exit()

ディアクティベートして終了する

$ deactivate

これで一通りpytorchを使ったGPU計算環境の動作が確認できました。

最後に

RTX3090/Ubunt24.04の計算環境を構築しました。
確認はpytorchだけ行いましたが、基本的には動作できているので他のプラットフォームでも大丈夫だと思います。

参考文献