LLM制作爆死回顧録、自作LLMを作るときに使った手法全まとめとどのようにして何の成果も得られなかったのかについて

この記事は動画を流す箇所があるのでイヤホンをして見てください。

何があったのか

• 2023夏, 俺は大学3年の夏休みだったのでLLMを自作したくなった

どうなったのか？

なんの成果も得られなかった

何故やったのか?

自作LLMを作ってHuggingFaceとかに公開して英雄になりたかった。

目次

1. Attention作ってTransformerを組んだ
2. GPTを作れる気がした
3. ぶち当たった困難
   • データがデカすぎる, HuggingFace DatasetAPIとmemmapを使って対応へ
   • tokenizerの仕方がBPEとかいう初見の何かだった
   • Positional Encodingも違う
   • 学習がうまくいかない, PostLNからPreLNへ
   • 生成がうまくいかないのでTemperatureとTopKも生成手法として追加
   • GPUが足りない1, float16で一部の演算を計算することに
   • おい, 損失がnanになったんだが, float16からbfloat16へ
   • 計算効率を高めるためにcompileとかbackendsを試しまくった
   • GPUが足りないので勾配累積, メモリ解放を試す
   • 結果的にGPUが足りない2
   • Instruction Tuning
4. 余談
   • AC6
   • 研究について
   • LTやります

1. Attention作ってTransformerを組んだ

2023年夏, 俺はDeepLearningにお熱だったのでAttentionをfrom scratchで実装し、Transformerも追加で実装することで†超絶DeepLearning人間†になろうと思った。もちろんAttentionのfrom scratch実装に興味はあったが, 市販のAIの本は「Attentionっていうのがあるよ」止まりで、その実装まで踏み切った本はマジで無かった(今はマシかも)。流石に, 計算機分野で実装してないで本読んでよくわかったとは言えないだろという思想の元, 怒りと共に始めた部分はあった。

ところで、DNNの実装記事が少ないのは日本特有なんじゃ無いかと思う。みんなDNNの理論的な解説だけで満足してばかりなので, 実際に実装に興味がある人は「{Model Name} from scratch」という魔法のキーワードで検索すると良い。

Attentionの一つであるSelfAttentionは実際海外記事を参考にするとすぐに組めた, あれは1次元の系列データ$(BatchSize, SequenceLength, Features)$から特徴量を取り出し, $(BatchSize, SequenceLength, Features)$のサイズにしてまた戻すものである。
SelfAttentionが理解できればSource-TargetAttentionもAttentionのKeyとValue, Queryが何かに注意して組めば簡単に組むことができる。このあたりは他の記事の方が詳しいので是非見てきてほしい。
そしてもし実装に興味があれば, Paper with Codeというサイトを絶対に参照してほしい。

少し理論の話をすると, Attentionは以下のモジュールで構成される。実装自体は使用ライブラリのドキュメントをしっかり読んで理解して実装をしよう。注意点として, 画像ではMatmulとあるが実際はGPUやAttentionの利点をしっかり活かすためにはtorchなどだとeinsum関数, einopsライブラリなどを使って実装する必要がある。

そしてこのAttentionを並列で処理させたものがMultiHeadAttentionであり, こちらが現在のTransformerの基本構成要素となっている。こちらもPaper with Codeを参照すればしっかり実装することができる。

そしてここまでAttentionの実装について述べてきたが, もはや最近のDNNではAttentionは基本パーツにすぎないのでPytorchやTensorflowでは上のような基本的すぎるAttentionはライブラリにて実装されている。
ここが驚いた点であった。もちろん最近はFlash Attentionといった派生Attentionも作られているため, from scratchでAttentionを組んだ経験は新しいモデルを考案するような研究には生かされるのでは無いかと思う。

Attentionが作り終わった俺はTransformerの制作に取り掛かった。

ここからは一度組んだModuleは再利用するという精神の元, nn.MultiheadAttention, nn.LayerNormなどのModuleを使った。Positional EncodingだけはPytorchに実装がなかったので, 外国の記事とかを参考にして組んだ。ちなみに,Positionalエンコーディングは系列データt番目の2k, 2k+1番目の変数に以下の値を加えるものである。今回はLLM作成爆死記事なので具体的な理論は伏せるが, Positional Encodingとは、Transformerにおいて系列データの時系列方向の情報を含ませるためのものです。
$$
PE(t, 2k) = \sin{(\dfrac{t}{T^{\frac{2k}{D}}})}, PE(t, 2k+1) = \cos{(\dfrac{t}{T^{\frac{2k}{D}}})}
$$

これを使ってTransformerを組み, Masked Attentionなども理解し, 実装を行いました。(Causal Attentionと呼ばれることの方が最近は多いかな?)
この時cos schedulerやHuggingFaceからデータを落としてhttps://huggingface.co/datasets/snow_simplified_japanese_corpusを使って日英翻訳モデルを構築しましたが, まあ段々と説明がめんどくさくなったので興味のある方はしらべてみてください。
そしてTransformerでの日英翻訳モデルの学習を始めました。
ん?

CUDA OUT OF MEMORY

なんか見えましたね....でもおそらく気のせいです。この時はTransformerのモデルサイズを落として学習を続行しました。
今思えばここで自分のマシンパワーが足りていないのではないのかと疑うべきでした。
出来上がったモデルの検証データでの出力はこんな感じでした。

予測0 Put ' cold outside to <eos> on your coat .
正解0 It is cold outdoors . put on your coat .
予測1 You get to get in touch with your parents . once .
正解1 You had better get in touch with your parents at once .
予測2 He will pay for dollars . most .
正解2 He will pay 20 dollars at most .
予測3 I is half ten minutes .
正解3 It leaves every thirty minutes .
予測4 He is every statesman in every of .
正解4 He is a politician in all senses .
予測5 Who will assure ? success ?
正解5 Who can guarantee his success ?
予測6 Our airport broke down on our way to the airport .
正解6 The car broke down on the way to the airport .
予測7 The are wish to it is s rising to leave early hours .
正解7 We all know that it ' s better to keep early hours .
予測8 I got a pair of shoes .
正解8 I bought a pair of shoes .
予測9 I hair stood on end .
正解9 His hair stood on end .

う〜〜んまあ学習時間が足りなかったからしょうがないかな？と思いました。
とにかく, 夏休み中盤の僕は自分でTransformerを組んで学習した事実が嬉しくて舞い上がっていました。

2. GPTを作れる気がした

この後すぐ, じゃあGPTってどうなってるんだろうと思いました。

あれ, Decoderだけやん‼️なんか今の俺ならいけんじゃね？
というわけで作り始めました。

GPTも海外のリポジトリを漁っていたら実装例は簡単に見つかりました。特に、nanoGPTは参考にさせていただきました。

3. ぶち当たった困難

ここからが本題です。しっかり行き詰まった困難, 対策などを必要に応じて算数をしながら話をしていきます。

データがデカすぎる, HuggingFace DatasetAPIとmemmapを使って対応へ

GPT1を再現実装したnanoGPTではopenwebtextというデータセットを用いていたので、僕もそれを使うことにしました。さてこちらのデータセットですが、全部で13GBあります。こんなものをメモリに全て載せてしまっていては、一般の家庭にある16GBPCでは持ちません。そこで、どう処理するかですが、幸いなことにopenwebtextはHugging Faceで提供されているため、Hugging FaceのData Process APIで前処理を行うことができます。内部の詳しいことは知りませんが、これを使うとプロセスが落ちずに処理を実行することができます。こうして前処理をしたところ800万個の文章データができました。
前処理まではHugging Faceでできても, 学習時にはGPUに乗せる必要があります。そのために一度Numpyなどで配列を作る必要がありますが、このデータをメモリに載せて処理するのはやはり無謀です。そこで使うのがnumpy.memmapです。これはC言語のmmapと同じ機能を提供し、UNIXのOSが行なっているページングやスワッピングを自プロセスで行なうようなものです。とりあえずこれを使えばデータをメモリに乗せることができると思ってください。というわけでHugging Face APIとmemmapによって、データがメモリに載らない問題を解決することができました。まだここでは挫折しません。

追伸: 最近の東大電気電子情報系のOSの授業でもmemmapが出てきましたが、こうして授業で学んだmemmapが使われているところを見るとCSを学んで活用できる人材になると一番ベストですよね。

tokenizerの仕方がBPEとかいう初見の何かだった

この記事を読んでいる人間全員が自然言語処理の人間ではないでしょうから、簡単にtokenizeという処理について解説したいと思います。自然言語はご存知の通り数字ではありませんから, 自然言語がPCに扱えるようにするためには数値にエンコードする必要があります。つまり、
$$
\text{自然言語} \rightarrow \text{int64}
$$
の処理を何かしらの形で与えてあげなければなりません。例として,文の始まりの記号を<BOS>, 終わりの記号を<EOS>, ピリオドも単語として区切ることを考えると
$$
\text{<BOS> I am called the Beast . <EOS>} \rightarrow \text{[0, 7, 2, 11, 91, 14, 51, 4]}
$$
このように数値化されます。このことをtokenizeといいます。そして自然言語の文章生成タスクは分類問題としてモデル化されますから、tokenizeされる語彙数だけ分類クラスがあるモデルを作らなければありません。 ここでこのことに関するヤバさについて話していきたいと思います。皆さんは大学生なら英語の辞書を持っていると思いますが, アレに入っている単語数は10万です。この場合, モデルの出力は1行10万もあるデータになります。機械学習ではよくfloat32型が使われます。これは1つの数を表すのに4バイト使いますが1行10万もあるデータでは1行に40万バイト, つまり一行に大体0.4GBも使ってしまいます。
というわけで単語ごとに区切るようなエンコード方法では効率が悪いわけです。GPT3では英語以外にも日本語や中国語にも対応しているので, 新たに効率の良いエンコード方法を使わなければなりません。これがBPE(Byte Pair Encoding) というエンコード方法になります。詳しくはこの辺の記事, SentencePieceを参照してください。この技術により, 未知語対応もうまくいきます。
BPEはfrom scratchで実装しても良かったのですが、夏休みの僕はとにかくDNNのモデル開発の方がしたかったのでtiktokenライブラリのGPT2Tokenizerを使いました。まだここでは挫折してません。

Positional Encodingも違う

実は, Transformer原論文と違い, GPTでは位置エンコードも学習済みパラメータになっています。 まあこれは簡単です。Embeddingされたデータの集まりを$U_v \in M_{BatchSize \times SequenceLength \times EmbeddingSize}$としたとき, $U_p \in M_{BatchSize \times SequenceLength \times EmbeddingSize}$というTensorを足し算すればOKでした。まだここでは挫折してません。

学習開始

というわけで学習を開始しました。GPT1はパラメーター数が約1億と言われています。実際に自分が組んだ時も, vocab_size(語彙サイズ) = 50257, embedding_size(埋め込みサイズ) = 768, attention_heads(attentionのhead数) = 12, num_layers(DecoderのLayer数) = 12, max_sequence_length(最大系列長) = 10000で組んだところ, パラメーター数が163035648となりました。これが全てfloat32だとすると, パラメーターだけで約0.650GBとなります。ここで僕は思いました。8GBのGPUに乗るやん‼️マジでGPTが作れるんじゃね!?

CUDA OUT OF MEMORY

はい、ダメです。そんなのでうまくいけば日本は個人開発者がLLMを作りまくってAI大国です。最初の訓練が原因を解説します。

学習データの分のメモリ割り当てと計算グラフの構築

まず当然のことですが, 学習パラメーターだけがGPUに乗る変数というわけではありません。学習する時のデータもGPUに乗っています。さらに, モデルの学習中は誤差逆伝播アルゴリズムを使うため, 常に計算グラフが作られており, その分の数もGPUが消費されます。 つまり, 学習中は推論とは別に計算グラフ分でGPUが圧迫されています。
というわけでモデルのサイズを小さくせざるを得ませんでした。

attention_headsを6に、num_layers = 6にしました。
(^ω^;;)雲行きが怪しくなってきたな...
でもまだ挫折してません。俺がこの世に生まれてきたからです。

学習がうまくいかない, PostLNからPreLNへ

はい、うまくいきませんでした。Lossが一向に収束しませんでした。
これも明確な原因があります。自然言語処理は捧げられた心臓の上に今があります。

PostLNとPreLN

Layer Normalization層がMultiheadAttentionやFFN層の後に置かれるというPostLN構造では勾配が消失してしまうという欠点があります。
そのため、最近のGPTでは主にPreLNという, MultiheadAttentionやFFN層の前にLayer Normalization層を置く構造が流行っているそうです。

生成がうまくいかないのでTemperatureとTopKも生成手法として追加

じゃあPreLNにすれば万事解決かと思いますが、そうもいきません。今までのように予測確率が一番高い単語を出力したとしましょう。その場合こんな感じになります。

'This is,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,'

明らかにこれは出力が上手くいってないですよね。そしてみんながChatGPTを使う時を思い出して欲しいんですけど、毎回返答が違いませんか? 自然言語による応答システムは毎回出力が変わっても許容されます。むしろ回答に多様性は持たせるべきです。
これを実装するための手法がTemperatureとTopKです。

Temperature

Temperatureはモデルの出力に補正を加えるものです。Temperatureによるトークンの予測確率$p_{i}$は、
$$
p_{i} \to \dfrac{\exp(\dfrac{p_i}{T})}{\sum_{j}\exp(\dfrac{p_j}{T})}
$$
となります。この操作により, Temperature: Tが大きいほど, 元々の予測確率が低かったものの確率が高くなります。

TopK

Temperatureで元々の予測確率が小さくなったとしても, ランダムに単語を出していては変な出力になってしまいます。そこで, TopKでは予測確率の一番高いトークンを選ぶのではなく、予測確率の高いトークンのうち、上位K個を出力対象としてTemperatureで計算した確率に従ってランダムに出力を選んでいます。
こうすることにより, 多様性が担保され、いい感じの出力が得られるようになります。

GPUが足りない1, float16で一部の演算を計算することに

さて, TemperatureやTopKという手法を導入しましたが, それでもGPUのメモリが足りなければ学習すらできません。
そこで, 最近のDeepLearningの学習では、一部の精度を求められないような計算はfloat32ではなく, float16などで計算されます。 こうすると計算速度もメモリ使用量も落ちます。これはPytorchならば, torch.amp.autocast, torch.cuda.amp.GradScalerを用いて, 以下のように実装されます。
実際, この手法によりついにまともに学習が開始できました。

scaler = torch.cuda.amp.GradScaler(enabled=True) 
###
省略
###
for batch_iter in range(batch_iteration):
        optimizer.zero_grad()
        with torch.amp.autocast(device_type=device, dtype=torch.float16):
            x,y = get_batch("train",batch_size=batch_size,device=device)
            padding_mask, mask = gpt.create_mask(x, 0, device)
            loss, pred = gpt(x,y,padding_mask,mask)
        scaler.scale(loss).backward() 
        scaler.step(optimizer) 
        scaler.update()

というわけで1週間ほど計算を回していました。夏場だったので家は暑くなるし, エアコンも冷やさなきゃだし, 電気代を払ってくれた親には申し訳なく思います。

おい, 損失がnanになったんだが, float16からbfloat16へ

さて、1週間ほど計算を回したところで, 損失がどうやら減っていないようだったので, 一旦コードを停止しました。すると, 損失がnanになっていました。 原因は一部の演算をfloat16にしていたためです。
簡単な例示をいたします。float16では$10^{-8}$以下で0としてみなされます。その0となった数が, 割り算に使われたとすると.......?
nanになります。 というわけで, 学習の際には損失がnanになったら停止するコードを書いておきましょう。float16かbfloat16かは訓練をしながら見極める必要がありそうです。
代わりに, bfloat16を使います。bfloat16は指数サイズがfloat32と同じなため、小さい数にも対応することができます。

しかしこれでも定期的にLossがnanになります。
これはLLMの学習などではLoss Spikeという名前で知られている, 急激に損失が大きくなることによる現象によって引き起こされます。定期的に学習のCheckPointを取って対策しておきましょう。

計算効率を高めるためにcompileとかbackendsを試しまくった

ここまで色々な対策を述べてきたんですけど, 何回か学習をリトライしているとだんだん学習の遅さが気になってきます。
というわけでNVIDIAの資料とかを眺めていたらこんなものが見つかりました。これはNVIDIAのPytorchベストプラクティスに関する講演スライドです。
そこには学習の高速化に関するものが書かれていました。

学習の高速化

1. モデルをコンパイルする
   モデルをtorch.compileでコンパイルすることで計算の効率化、最適化を行ってくれます
2. torch.backendsをいじる

• torch.backends.cudnn.benchmark
  torch.backends.cudnn.benchmarkをTrueにすることでネットワークの構成に対して最適なアルゴリズムを見つけて計算を行うため、計算が早くなります。
  (しかし、再現性はなくなります。)

• torch.backends.cudnn.allow_tf32
  torch.backends.cudnn.allow_tf32をTrueにすることで対応しているGPUならばTensorFloat32コアを使用して計算を行います。

これらの機能を使うことで学習が爆速になりました。

GPUが足りないので勾配累積, メモリ解放を試す

先ほど, bfloat16を使ってメモリの節約を行いましたが、まだいけます。

メモリ解放

1 epoch終わるごとに使わなくなったメモリ領域をCPU, GPU共に解放することで, 多少無茶な計算をギリギリ通すことができます。コードはこんな感じです。

gc.collect() #メモリを解放
torch.cuda.empty_cache() #GPUのメモリを解放

勾配累積

クソデカモデルの学習ではモデルが大きすぎるため, あまり大きな訓練データサイズで訓練ができません。そこで考案されたのが勾配累積という方法になります。例えば, pytorchではパラメーターの更新の際, こんな感じのコードを書くはずです。

for batch_iter in range(batch_iteration):
        data, label = get_data()
        optimizer.zero_grad()
        loss = model(data, label)
        loss.backward()
        optimizer.step()

もしdataのバッチサイズが小さかったらどうでしょう、例えば6とかの場合, 6個の文章データの予測に関してパラメーターを更新しますこれは小さなデータセットならばすべてのデータをバッチサイズ6で更新すればいいですが, 大きなデータセットの場合大きな問題となります。
そのため, 勾配累積では勾配を貯めておいて最後にパラメーターを更新します。

all_loss = 0
for batch_iter in range(batch_iteration):
        data, label = get_data()
        optimizer.zero_grad()
        loss = model(data, label)
        all_loss += loss / batch_iteration
all_loss.backward()
optimizer.step()

こうすることで, $\text{BatchSize} \times \text{BatchIteration}$のデータを一度に学習させた場合と同じ勾配でパラメーターが更新されます。大きなデータで学習する時の基本的な手法のようです。

結果的にGPUが足りない2

はい, 以上全てを試しましたがどうしても損失は減り切らず, 結果的に上手くいくことはできませんでした。心が折れました。

↓クリックせよ

理由はいくつか考えられると思います。

1. GPU不足のため, 学習データサイズや系列長のサイズを十分に取ることができなかった
   上で述べたように学習データサイズが大きければ大きいほどやはり学習は効率的になりますが, 8GBのGPUではそれを行うことはできません。また, 系列長も十分に取れなかったのは一番痛いかなと思います。系列長がNの時, 現時点からNトークンだけ前の情報を元にモデルは単語を予測します。実際に組んでいる時も系列長が長い方が予測精度が良かったですし, 最近のGPTモデルはGPT1に比べて系列長が非常に長くなっています。系列長が長い方が文脈を広く捉えられていいモデルになるはずです。

2. まだ試せてない手法がある
   今回, ラベルスムージングという手法を試せていません。
   これは正解ラベルをone hot encodingして確率1にするのではなく, ある程度他の単語の確率を割り振ってあげる分, 正解ラベルの確率を低くするというものです。
   
   これをpytorchで実装すると, やはりCUDA OUT OF MEMORYがでてしまいます。
   余談ですが、今回参考書としてIT Text 自然言語処理の基礎を使わさせていただきました。この参考書はChatGPT以前の自然言語処理がまとまっておりますが、今読んでも役に立つ手法が多いため、自然言語に興味のある方はぜひご購入ください。

最後に, 出来上がったGPT1モデルの出力を見てみましょう。

入力: He has two sisters, and 
出力: He has two sisters, and ents, but not a great-uncle, 
The Boy in all that is the same type-snow, both, but a good name. 
I know there's two very bad and they are all but a good reason 
they've probably got good in comparison, in fact I know that 
they haven't the only the most of these guys in our division of the first
names.'The same old guys; I'm not the opposite that this is that, either (although they just got better guys are that, 
in that team. You can probably have more of their name: they probably just 
the same types as those other names and

入力: Japan is a country, which
出力: Japan is a country, which has the longest-populifying population of about 
10 people, but most recently. But this trend has seen a trend in the global 
poverty. A drop by population of 10 is emerging growth since 2010. 
What are people have seen some of the trend that is now is rising across 
all the global food, but more of which has seen some is experiencing growth 
trends in recent years, and growing. For most is a percentage since then 
growth since 2010. Now that growth since 2006 and growing by growth means 
increasing growth over the over half is growing more than this increase is 
just one, in part, with over 10, a trend in 2007 is

Instruction Tuning

さて、ここまでGPT1制作記を読んでいただきありがとうございます。ここまで読んで気づいた方もいるかもしれません。あれ、こいつFineTuningしてなくね？ と その通りです。 最近のLLMでは大規模な事前学習を行い、次単語予測モデルを組み立てた後は予測を正確にするためにFineTuningを行なっております。これはInstruction Tuningと呼ばれ, SFTやらRLHFやらの手法があります。LLaMAはRedPajamaデータ(1TB)で学習されたという噂があります。これの実装は俺のGPUが24GBとかになったらやろうと思います。 すでに海外ではnanoChatGPTというリポジトリがあり、GPUのある人は組むことができます。次回作にご期待ください。

4.余談

AC6

夏休みにGPTを組んだのですが, その時期はアーマードコア6が発売された時期だったのですごい遊びたかったのにGPT作ってたから遊べなかったストレスがすごかったです。この経験から、人類はGPU付きのパソコンを2つは持つべきだと言えるでしょう。

研究について

最近, 8GBでLLM組めなかったり, OSSLLMを動かせなかったりする経験があったので、大学ではモデル軽量化の研究をやります。多分

LTやります

この記事の内容を2024年1月にLTすることになりました。詳しくは追って連絡します。1/17かも

LTした。

サーバー用のGPUを譲り受けることになったので奨学金でパソコン買います