結論から言うとハードウェアが対応しておらず動きませんでした。しかし環境構築の手順としては間違っていないと思うので書き残しておきます。

FP8で学習したい

FP8は、8bitで浮動小数点数を表現するデータフォーマットです。深層学習をFP8で行うことのメリットとして、以下が挙げられます。

学習速度の向上
メモリ使用量の削減
推論と同一の精度を用いることによる、量子化に伴う誤差の低減

FP8を用いてTransformerモデルの学習を効率的に行うためのNVIDIA謹製オープンソースライブラリがTransformer Engineです。NVIDIAの技術ブログにて詳細な解説がなされています。

導入

環境

Ubuntu 22.04 LTS amd64
GeForce RTX3090
CUDA 12.3

cuDNNのインストール

pipでいきなり入れようとしたらcuDNNが無いと怒られたので導入します。
cuDNN Archiveからダウンロードします(要NVIDIAアカウント)。OSおよびCUDAの環境に合わせてダウンロードしてください。以下はdeb系の場合のコマンドです。

sudo dpkg -i cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb 
sudo cp /var/cudnn-local-repo-*/cudnn-local-*-keyring.gpg /usr/share/keyrings/
sudo apt update
sudo apt install libcudnn8 libcudnn8-dev libcudnn8-samples

Transformer Engineのインストール

pipで導入する場合です。

pip3 install git+https://github.com/NVIDIA/TransformerEngine.git@stable

動かない

NVIDIAの技術ブログにあるサンプルコードを動かしてみると、

File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/fp8.py", line 539, in fp8_autocast
FP8GlobalStateManager.fp8_autocast_enter(
File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/fp8.py", line 394, in fp8_autocast_enter
assert fp8_available, reason_for_no_fp8
AssertionError: Device compute capability 8.9 or higher required for FP8 execution.

と表示されました。私の環境では動かないようです。私の調査によるとFP8がコンシューマ向けGPUに搭載されたのは4000番台かららしいので、3090は非対応ということになります。

~~いかがでしたか？~~ 研究室の先生におねだりしてFP8が使えるGPUを買ってもらおうと思います。