KellerJordan/modded-nanogpt

解決する課題

言語モデルを可能な限り迅速にトレーニングするという課題に取り組みます。具体的には、8枚のNVIDIA H100 GPU上で、特定のパフォーマンス目標（FineWeb検証セットで3.28のクロスエントロピー損失）を達成できる最速のアルゴリズムを見つけることを目的としています。

仕組み

このプロジェクトは、トレーニングパイプラインを反復的に最適化する共同「スピードラン」です。NanoGPTとllm.cに基づくPyTorchトレーナーから始まり、その後、以下を含む数十のアーキテクチャおよびシステム最適化を統合してきました：

• Optimizers: MuonおよびNorMuonオプティマイザの実装。
• Architecture: Rotary embeddings、QK-Norm、ReLU²、および埋め込みから各ブロックへのスキップ接続の使用。
• Precision: ヘッドとMLPのフォワードパスにFP8を、クロスエントロピー計算にBF16を活用。
• Attention: Flash Attention 3、長短スライディングウィンドウパターン、およびYaRNの統合。
• Systems: 勾配のall-reduce/reduce-scatterの最適化、および計算と通信のオーバーラップ。

対象者

極限のトレーニング効率、LLMの最適化、およびH100 GPUにおけるハードウェア利用率の限界に興味のあるAI研究者、エンジニア、および愛好家向けです。

ハイライト

• 大幅な高速化: 同じ目標損失に対して、トレーニング時間を45分（ベースライン）から2.2分未満に短縮。
• データ効率: 必要トークン数を100億から4億未満に削減。
• 共同進化: トレーニング時間を数秒削るために使用された技術の進展を記録した、詳細な世界記録履歴。