KellerJordan/modded-nanogpt

何を解決するか

8台のNVIDIA H100 GPUを使用して、特定のパフォーマンス目標（FineWeb検証セットでの3.28の交差エントロピー損失）に達するための、可能な限り高速なアルゴリズムとシステム構成を見つけることを目指しています。

仕組み

このプロジェクトは「スピードラン」として動作し、アーキテクチャ、最適化、システムレベルの技術の膨大な組み合わせを反復的に適用・テストして、学習時間を短縮します。アンドレイ・カーパティの llm.c と NanoGPT から発展し、45分のベースラインから3分未満まで短縮しました。主な手法は以下の通りです：

• アーキテクチャの変更：Rotary埋め込み、QK-Norm、ReLU²、および埋め込みから各ブロックへのスキップ接続の実装。
• 最適化の革新：MuonおよびNorMuon最適化手法、およびツインレールモーメンタムを備えたANVIL最適化手法の利用。
• 精度とハードウェア：MLPの前方/後方伝播、QKV投影、およびヘッドでFP8を使用。CUDAグラフキャプチャとFlash Attention 3を併用。
• データと学習ロジック：バッチの開始をEnd-of-Sentence（EoS）トークンと同期、マルチトークン予測の実装、バッチサイズスケジュールの使用。

対象読者

極限的なLLM学習効率、ハードウェア加速、および最適化・アーキテクチャ設計の最先端に興味を持つ研究者やエンジニア。

主な成果

• 8台のH100 GPU上で学習時間を45分から3分未満に短縮。
• 目標損失に到達するためのトークン数を100億から3億3000万未満に削減。
• Polar Express、スパースアテンションゲート、ハッシュ化されたビグラム/トリグラム埋め込みテーブルなど、多数の最先端技術を実装。