Team 甲(きのえ) そして、やり直しへ

GW明け・・・

Phase1が開始して2週間という時間は地獄でしたが、今思うとまだまだ序の口でした。モデルがGPUに乗らない、学習を始めたらloss nanが起こる、そして、ノードの調子がおかしくて交換・・・。
様々なイベントが発生し「LLM開発は難しい」を経験したのですが、まだまだ魔界の王は解放してくれませんでした。

同時に起こる課題

GW明けの時には、以下のような状況でした。

• 2ノードで学習を進めている
• node-7の復帰待ち

そこに、新たな課題が持ち上がります。

なんか遅くないか・・・？

イテレーション速度が、当初想定していたよりも遅い・・・、このままで大丈夫か？という疑念がチーム内で持ち上がります。(以下は5/6の時点でのログの状態です。ログの見方は改めて別の記事で記載します。)

slurm0-a3-ghpc-6: [2024-05-05 21:22:01,117] [INFO] [logging.py:96:log_dist] [Rank 0] step=25280, skipped=0, lr=[0.0001275338142173688, 0.0001275338142173688], mom=[(0.9, 0.95), (0.9, 0.95)]
slurm0-a3-ghpc-6: [2024-05-05 21:22:01,997] [INFO] [timer.py:260:stop] epoch=0/micro_step=4880/global_step=4880, RunningAvgSamplesPerSec=442.13444082258474, CurrSamplesPerSec=433.88114192728943, MemAllocated=19.08GB, MaxMemAllocated=43.5GB
slurm0-a3-ghpc-8:  iteration    25280/  286102 | consumed samples:     25886720 | consumed tokens:  53016002560 | elapsed time per iteration (ms): 16585.5 | learning rate: 1.275E-04 | global batch size:  1024 | lm loss: 2.193328E+00 | loss scale: 1.0 | curriculum seqlen:  2048 | actual seqlen:  2048 | number of skipped iterations:   0 | number of nan iterations:   0 | samples per second: 61.741 | tokens per gpu per second (tgs): 7902.807 | TFLOPs: 391.94 |

3行目のiteration 25280/ 286102を見ると、6日間ぐらい実行してまだ25000回しか回っていない状態です。

3ノード実行へ変更するか？

2ノードで実行している学習を3ノードにしたら早くなるのだろうか。単体ノードのものを複数ノードにするのであれば通信によるオーバーヘッドが新たに発生しますが、2ノードで学習をしているものを3ノードにするのであれば、既に通信によるオーバーヘッドは内包されており、速度劣化は起こらないとは思います。
また、Phase2ではモデルサイズのターゲットが50Bと大きくなり、使えるノードが増えるので、どのようにノードを追加していくのかは課題になりそうですが、いずれにしても何かしら方法があるに違いありません。
とりあえず検討している間にもnode-7が復帰しましたので、3ノードでの学習にチャレンジします。しかし、チェックポイントから学習を再開できない事態となりました。

HuggingFace変換ができない

大体同じぐらいのタイミングで、HuggingFace形式への変換が難しいという課題もあがりました。どうも設定したパラメータでは、保存してあるモデルデータを分割するなり再結合するなりをして変換していかないとHuggingFace形式にならないらしく、相当難しいとのことでした。
また、その分割や再結合などは2ノードから3ノードへの変更にも絡んでいて、暗雲が立ち込めてきました。

DeepSpeedのバージョンが古い

次に問題になってきたのはno_pp=falseにして実行したいが、bf16が使えないというものでした。これは標準コードで指定しているDeepSpeedのバージョンは未対応でバージョンを上げれば対応が可能というものでした。ただし、DeepSpeedは根幹のライブラリであるため、上げた後のバージョンで動作が保証されるのか否かを賭けることになります。

最初から学習をやり直すか？

色々と課題が上がる中、現状のパラメータの設定のままでは、仮に学習を再開できても解決が難しい課題が残ります。それであれば学習を最初からやり直すことで、課題が解決しやすい形にするのが良いという意見も出てきました。
既に2週間という時間を使ってしまっていて、また学習をやり直すとなると相当勇気がいることです。保存されているチェックポイントを捨てるかどうか。胃が痛くなる選択です。

5月23日 12時過ぎ・・・

このタイミングで我々の事前学習は無事に処理を終了しました。
様々な課題が発生しましたが、対応を行った後は、些細な問題を除いて大きな問題は起こらず事前学習は最後まで走り切りました。
では、我々がどのような決断をしたのか、詳細は以下の通りです。

• DeepSpeedのバージョンアップ
  • バージョンを上げたあとはパラメータを調整しています
  • 例えばglobal batch sizeを1024 → 1536、fp16 → bf16、no_pp=true → false
• 最初から事前学習を3ノードでやり直し
  • やり直したとしても、見積もった学習速度から計算して、スケジュールには間に合うという想定でした
  • また、途中まで保存されたファイルの変換を頑張っても、変換自体がうまくいくのかどうかや変換後の検証まで実施する場合の工数が未知数でしたので、変換自体に積極的ではありませんでした

また、最初の 2 ノードでの学習時には多くの課題が重なりましたが、3 ノードでやり直した後の学習では幸いにも安定した動きとなりました。もしも、トラブルが再び発生するようなことがあれば、心が折れていたかもしれません。

そして、事後学習へ

事前学習が終わってからは事後学習へと推移していきます。
この辺の話は、また別の機会に記載します。

この成果は、NEDO（国立研究開発法人新エネルギー・産業技術総合開発機構）の助成事業「ポスト５Ｇ情報通信システム基盤強化研究開発事業」（JPNP20017）の結果得られたものです。