OpenAIの共同設立者が「GPT-2」をわずか90分の時間と3100円の費用で再現するレポートを公開

2019年に公開されるやいなや「危険すぎる」と話題になったOpenAIの「GPT-2」を、ゼロから作成したレポートが公開されました。作成にかかった時間は、1時間14ドル(約2200円)のクラウドコンピューティングサービスで1時間半だったため、総コストは約20ドル(約3100円)に収まっています。

Reproducing GPT-2 (124M) in llm.c in 90 minutes for $20 · karpathy/llm.c · Discussion #481 · GitHub
https://github.com/karpathy/llm.c/discussions/481

今回、GPT-2の再現を行ったのは、OpenAIの共同設立者のAndrej Karpathy氏です。あえてGPT-2を選んだのは、GPT-2が大規模言語モデル(LLM)の祖先であり、また2019年のモデルにもかかわらず、GPT-3やその他のモデルと基本的な仕組み自体はほとんど同じだからだというのが理由だとKarpathy氏は説明しています。

AIの作成にあたって最初に必要になるのはGPUですが、Karpathy氏はAI向けGPUをクラウドで提供しているサービスのLambdaを選択しました。

Karpathy氏によると、こうしたケースでよく使われる構成はLinux x86 64bit Ubuntu 22.04にライブラリとしてCUDA 12を導入したものとのこと。この環境で、Karpathy氏は以下の起動コマンドを使用しました。長いですが、要約すると「12層のGPT-2(124M)をゼロから、10億トークンのFineWebデータセットを使って、最大シーケンス長1024トークンでトレーニングする」という意味になります。

# install miniconda
mkdir -p ~/miniconda3
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
rm -rf ~/miniconda3/miniconda.sh
~/miniconda3/bin/conda init bash
source ~/.bashrc

# pytorch nightly (optional) https://pytorch.org/get-started/locally/
# conda install --yes pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch-nightly -c nvidia

# pip installs so we can tokenize the FineWeb dataset
yes | pip install tqdm tiktoken requests datasets

# install cudnn so we can use FlashAttention and run fast (optional)
# https://developer.nvidia.com/cudnn-downloads
# for me, CUDA 12 (run `nvcc --version`) running on Linux x86_64 Ubuntu 22.04
wget https://developer.download.nvidia.com/compute/cudnn/9.1.1/local_installers/cudnn-local-repo-ubuntu2204-9.1.1_1.0-1_amd64.deb
sudo dpkg -i cudnn-local-repo-ubuntu2204-9.1.1_1.0-1_amd64.deb
sudo cp /var/cudnn-local-repo-ubuntu2204-9.1.1/cudnn-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cudnn-cuda-12

# "install" cudnn-frontend to ~/
git clone https://github.com/NVIDIA/cudnn-frontend.git

# install MPI (optional, if you intend to use multiple GPUs)
sudo apt install openmpi-bin openmpi-doc libopenmpi-dev

# tokenize the FineWeb dataset 10B tokens sample (takes ~1 hour, get lunch?)
# writes ~19GB of raw GPT-2 tokens to dev/data/fineweb10B
# and ~46GB in ~/.cache/huggingface/datasets/HuggingFaceFW___fineweb
git clone https://github.com/karpathy/llm.c.git
cd llm.c
python dev/data/fineweb.py --version 10B

# compile llm.c (mixed precision, with cuDNN flash-attention)
# first compilation is ~1 minute, mostly due to cuDNN
make train_gpt2cu USE_CUDNN=1

# train on a single GPU
./train_gpt2cu \
-i "dev/data/fineweb10B/fineweb_train_*.bin" \
-j "dev/data/fineweb10B/fineweb_val_*.bin" \
-o log124M \
-e "d12" \
-b 64 -t 1024 \
-d 524288 \
-r 1 \
-z 1 \
-c 0.1 \
-l 0.0006 \
-q 0.0 \
-u 700 \
-n 5000 \
-v 250 -s 20000 \
-h 1

# if you have multiple GPUs (e.g. 8), simply prepend the mpi command, e.g.:
# mpirun -np 8 ./train_gpt2cu \ ... (the rest of the args are same)

その結果、GPT-2の最小モデルであるGPT-2 124Mを、8X A100 80GB SXMノードで約90分で作成できました。Lambdaでは、このノードは1時間当たり最大14ドルで提供されているので、総コストは約20ドルになる計算です。なお、GPUを1枚に減らしてトレーニングすることも可能ですが、それに比例して必要な時間が増えるため、GPU次第ですが4～24時間かかるようになるとのこと。

FineWebの検証用データセットでの評価(左)と、ベンチマークのHellaSwagの精度(右)は以下の通り。

さらに、Karpathy氏はGPT-2 350Mも作成してみました。その結果、かかったトレーニング時間は14時間で、費用はGPT-2 124Mの10倍の200ドル(約3万1000円)でした。

今後、Karpathy氏はGPT-2のサイズを740M、1558Mと増やしていくつもりですが、1558Mモデルを前述のノードで作るとなると1週間と2500ドル(約39万円)の費用がかかる見込みです。そのため、Karpathy氏は「できないこともありませんが、まずはコードを洗練させたり、テストを改善したり、マルチノードでトレーニングできるようにしたりするのに時間を割いてみようと思います」と述べました。