データ、ハードウェア、モデルの選択

GPT をゼロから組み立てて学習する

このシリーズでは、データを集め、言語モデルを組み立て、系列の次の単語を予測するように学習させる。微調整のあとは、{in: システムプロンプト} + {in: ある単語で途切れた対話} + {out: 次の単語の予測} という形になる（これはあとで図にする）。

このシリーズは Sebastian Raschka の『Build a Large Language Model (From Scratch)』に沿って進める。

前提知識:

• Python
• ニューラルネットワーク

読了時間: 1 週間。

この回では、データセットをダウンロードし、ハードウェアに合わせたモデルの選択を根拠づけ、学習の予算を見積もる。

データのダウンロード

ダウンロードには 3 日ほどかかり、回線速度には依存しない。アーカイブにはディスク 47 GB が必要なので、早めに始めておくとよい。データセットには Project Gutenberg を使う。50,000 冊以上、$3 \cdot 10^{9}$ 語以上。

git clone https://github.com/pgcorpus/gutenberg.git
cd gutenberg
pip install -r requirements.txt
python get_data.py

システム要件

• RAM: 16 GB
• VRAM: 最低 5 GB（1024 トークンの系列 1 本のバッチ）、16 GB ならバッチ 4〜6、24 GB ならバッチ 8〜12。
• ディスク: 50 GB

少しの歴史

• 2017.06: Attention is All You Need
• 2018.06: GPT-1 (Improving Language Understanding by Generative Pre-Training) - 1 億 1700 万パラメータ、約 10 億トークン
• 2019.02: GPT-2 (Language Models are Unsupervised Multitask Learners) - 15 億パラメータ、100 億トークン
• 2020.05: GPT-3 (Language Models are Few-Shot Learners) - 1750 億パラメータ、3000 億トークン

モデルとデータの選択

GPT-3 ではなく GPT-2 small を選ぶ。両者の違いは学習データの量とモデルの大きさだけで、アーキテクチャは同じだ。2026 年なら、bf16、flash attention、torch.compile を有効にすれば、GPT-2 small は民生用 GPU 一枚で一日で学習できる。品質はもちろん劣るが、微調整のあとなら会話はできる。

モデルの大きさはパラメータ数で測り、データセットの大きさはトークン数で測る。

資源を節約するため、英語だけのモデルを学習する。他の言語は単語あたり約 2 倍のトークンを使う。GPT-2 は 1 億 2400 万パラメータ。データセット中の英語の本は 51,668 冊のうち 39,478 冊。

初期の GPT は、データよりモデルを大きくするほうが得だという考えで作られていた。その後 Hoffmann が、最適点はパラメータあたり約 20 トークンだと実験的に示した。実験に使ったモデルを Chinchilla と名づけたので、「20 の法則」と Chinchilla-optimal という言葉が生まれた。

GPT-2 に 20 の法則を当てると、データセットの下限は $124 \cdot 10^{6} \times 20 = 2.48 \cdot 10^{9}$ トークン。クリーニング後の英語 Gutenberg は、私の見積もりで約 42 億トークン。最適点の 1.7 倍なので、データは足りる。

Chinchilla が答えるのは「計算予算が決まっているとき、モデルとデータをどの大きさにするか」という問いだ。

• GPU の性能。自分の GPU の技術ホワイトペーパーを探す。見るべきは fp32 累算のテンソルコア bf16。PyTorch の autocast はまさにこの経路で動く。グラフィックカードは Nvidia 4090。Peak BF16 Tensor TFLOPS with FP32 Accumulate: 165.2 TFLOPS、つまり $165.2 \cdot 10^{12}$ FLOP/s。
• 時間。時間を見積もるには演算の数が要る。学習では、パラメータ 1 個とトークン 1 個あたり、forward pass で 2 演算（乗算と加算）、誤差逆伝播で 4 演算（活性に関する勾配と重みに関する勾配、それぞれ 2 演算）、合計 6 演算かかる。パラメータは $N = 124 \cdot 10^{6}$、トークンは $D = 4.2 \cdot 10^{9}$。

1 エポックの学習に必要な演算数:

$$C = 6ND = 6 \cdot 124 \cdot 10^{6} \cdot 4.2 \cdot 10^{9} \approx 3.12 \cdot 10^{18}\ \text{FLOP}$$

このカードでは:

$$t = \frac{C}{165.2 \cdot 10^{12}} = \frac{3.12 \cdot 10^{18}}{165.2 \cdot 10^{12}} \approx 18\,886\ \mathrm{s} \approx 5.2\ \mathrm{h}$$

これは利用率 100% の場合。attention、正規化、オプティマイザ、データ転送のオーバーヘッドがあるので、実際の利用率は 40% 程度。0.4 で割ると $18\,886 / 0.4 \approx 47\,215$ 秒、13.1 時間。2 エポックで 26.2 時間、ほぼ一日。

• 消費電力。仕様では GPU の電力上限は 450 W、実際の学習中に nvidia-smi が示すのは約 400 W。2 エポックで 26 時間。$0.4\ \mathrm{kW} \times 26\ \mathrm{h} = 10.4\ \mathrm{kWh}$。ここで利用率を二重に数えてはいけない。利用率のせいで 10 時間ではなく 26 時間かかっているからだ。CPU、メモリ、電源の変換損失にさらに 100〜150 W、つまり 3〜4 kWh。コンセントで合計約 14 kWh。

まとめ: $N = 124 \cdot 10^{6}$ パラメータのモデル（GPT-2 small）を、クリーニング後の Gutenberg $D = 4.2 \cdot 10^{9}$ トークンで学習するのに、約 14 kWh と 26 時間。bf16、flash attention、torch.compile を使う前提で。

同じ本での 2 エポック目は、新しいデータとほぼ同じだけ役に立つ（Muennighoff et al. 2023、4 回の繰り返しまで）。

トークン数は開始前に知っておく必要がある。学習率スケジュールはそれを基に組むからだ。だから次は、まずクリーニングして数え、それから学習する。