GPT-2(124M)をllm.cでGoogle Colab環境で再現チャレンジしてみる

今回は、下記の「GPT2(124M)をllm.cで90分間20ドルで再現してみましょう」という記事を参考に、Google Colabでも再現にチャレンジしました。

上記記事は、12層、12ヘッド、768次元の124Mトランスフォーマーモデルを、FineWebの100億トークンで訓練しています。

今回は、Google Colabで、上記コードを修正してみましたので紹介していきます。

今回のGoogle Colab環境は、ランタイムはA100 GPUを使用しています。

以下、Google Colab環境で実行してください。

最初に、Minicondaをインストールします。

# 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

次に、必要なライブラリをインストールします。

!pip install tqdm tiktoken requests datasets

次に、cudnn-frontendをインストールします。

!git clone https://github.com/NVIDIA/cudnn-frontend.git

次に、openmpi-bin, openmpi-doc, libopenmpi-devをインストールします。マルチGPU用だからたぶん不要ですが、おまじないとして。

!apt-get install openmpi-bin openmpi-doc libopenmpi-dev

次に、100億のFinewebのデータセットをダウンロードして、トークン化します。Shard 103までありますので時間がかかります。

!git clone https://github.com/karpathy/llm.c.git
!cd llm.c && python dev/data/fineweb.py --version 10B

次に、llm.cをコンパイルします。#USE_CUDNN=1をコメントアウトしていますが、備忘として残してます。

%cd /content
# Compile llm.c (mixed precision, with cuDNN flash-attention)
# This step might take a while
!cd llm.c && make train_gpt2cu
#USE_CUDNN=1 

次に、Single GPUで訓練します。train_gpu2cuのオプションで、-bのバッチのサイズを64から32に変更しています。64だとOut of Memoryで動かなかったからです。

%cd llm.c
!./train_gpt2cu \
  -i "dev/data/fineweb10B/fineweb_train_*.bin" \
  -j "dev/data/fineweb10B/fineweb_val_*.bin" \
  -o log124M \
  -e "d12" \
  -b 32 -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

すると、実行結果として次のように表示されます。

+-----------------------+----------------------------------------------------+
| Parameter             | Value                                              |
+-----------------------+----------------------------------------------------+
| train data pattern    | dev/data/fineweb10B/fineweb_train_*.bin            |
| val data pattern      | dev/data/fineweb10B/fineweb_val_*.bin              |
| output log dir        | log124M                                            |
| checkpoint_every      | 5000                                               |
| resume                | 0                                                  |
| micro batch size B    | 32                                                 |
| sequence length T     | 1024                                               |
| total batch size      | 524288                                             |
| learning rate (LR)    | 6.000000e-04                                       |
| warmup iterations     | 700                                                |
| final LR fraction     | 0.000000e+00                                       |
| weight decay          | 1.000000e-01                                       |
| max_steps             | -1                                                 |
| val_loss_every        | 250                                                |
| val_max_steps         | 20                                                 |
| sample_every          | 20000                                              |
| genT                  | 64                                                 |
| overfit_single_batch  | 0                                                  |
| use_master_weights    | enabled                                            |
| recompute             | 1                                                  |
+-----------------------+----------------------------------------------------+
| device                | NVIDIA A100-SXM4-40GB                              |
| precision             | BF16                                               |
+-----------------------+----------------------------------------------------+
| load_filename         | d12                                                |
| max_sequence_length T | 1024                                               |
| vocab_size V          | 50257                                              |
| padded_vocab_size Vp  | 50304                                              |
| num_layers L          | 12                                                 |
| num_heads NH          | 12                                                 |
| channels C            | 768                                                |
| num_parameters        | 124475904                                          |
+-----------------------+----------------------------------------------------+
| train_num_batches     | 19560                                              |
| val_num_batches       | 20                                                 |
+-----------------------+----------------------------------------------------+
| run hellaswag         | no                                                 |
+-----------------------+----------------------------------------------------+
| Zero Stage1 is enabled                                                     |
| num_processes         | 1                                                  |
| zero_stage            | 1                                                  |
+-----------------------+----------------------------------------------------+
HellaSwag eval not found at dev/data/hellaswag/hellaswag_val.bin, skipping its evaluation
You can run `python dev/data/hellaswag.py` to export and use it with `-h 1`.
num_parameters: 124475904 => bytes: 248951808
allocated 237 MiB for model parameters
batch_size B=32 * seq_len T=1024 * num_processes=1 and total_batch_size=524288
=> setting grad_accum_steps=16
---
WARNING: Failed to open the tokenizer file gpt2_tokenizer.bin
The Tokenizer is a new feature added April 14 2024.
Re-run `python train_gpt2.py` to write it
---
allocating 20715 MiB for activations
val loss 11.008764
allocating 237 MiB for parameter gradients
allocating 1008 MiB for activation gradients
allocating 474 MiB for AdamW optimizer state m
allocating 474 MiB for AdamW optimizer state v
allocating 474 MiB for master copy of params
step    1/19560 | train loss 11.009281 | norm 15.7032 | lr 8.57e-07 | 3462.81 ms | 41.7% A100 fp16 MFU | 151405 tok/s
step    2/19560 | train loss 10.959524 | norm 15.3318 | lr 1.71e-06 | 3458.26 ms | 41.8% A100 fp16 MFU | 151605 tok/s
step    3/19560 | train loss 10.857655 | norm 14.6833 | lr 2.57e-06 | 3466.00 ms | 41.7% A100 fp16 MFU | 151431 tok/s
step    4/19560 | train loss 10.715219 | norm 13.0278 | lr 3.43e-06 | 3464.53 ms | 41.7% A100 fp16 MFU | 151396 tok/s

今回は、Step数が19560あり、時間がかかるので、途中で止めました。時間とGoogle Colabのクレジットに余裕のある人は最後まで実行してみると良いかもしれません。