NanoEuler: CとCUDAでゼロから構築されたGPT-2スケールのLLM
NanoEuler: CとCUDAでゼロから構築されたGPT-2スケールのLLM
NanoEulerは、CとCUDAのみを使用してGPT-2クラスの言語モデルを完全にゼロから実装する研究および教育プロジェクトです。PyTorchやautogradのような高レベルなMLライブラリを回避することで、このプロジェクトは、手書きのバイトレベルBPE tokenizerから、厳密な勾配チェックによって検証された教師あり微調整（SFT）に至るまで、すべてを網羅する透明性の高いエンドツーエンドのパイプラインを提供します。
コアアーキテクチャと設計
NanoEulerは、コンシューマー向けハードウェアでの効率を最大化するために、いくつかの現代的なLLMの最適化を利用したデコーダーのみのTransformerアーキテクチャを実装しています。アーキテクチャを簡素化し、トレーニングの安定性を向上させるために、バイアスなしで設計されています。
技術仕様
- RMSNorm: バイアスなしの事前正規化として使用。
- Rotary Position Embeddings (RoPE): 位置情報を扱うために、queryとkeyに適用。
- SwiGLU Feed-Forward:
down(silu(gate(x)) * up(x))の活性化パターンを実装。 - Grouped-Query Attention (GQA): 複数のquery headsを共有する、より小さなsetのkey/value headsを使用することで、メモリと計算を最適化。
- Multi-Token Prediction (MTP): 次の
K個のトークンを予測するためにK個の出力ヘッドを使用。これにより、学習された表現が強化され、speculative decoding（ただし、生成は主にhead 0を使用）をサポートします。 - Byte-level BPE Tokenizer: GPT-2スタイルのpretokenizationを使用するカスタム実装。スペースが独立したトークンとして無駄に消費されないことを保証します。
モデル構成
| バージョン | 次元 | Q/KV Heads | レイヤー | コンテキスト | Vocab | パラメータ数 |
|---|---|---|---|---|---|---|
| Small (CPU) | 128 | 4 / 2 | 4 | 128 | 512 | ~1.05M |
| GPU Pipeline | 768 | 12 / 4 | 16 | 512 | 4096 | ~116M |
CUDAエンジンとパフォーマンス
単一のコンシューマー向けGPU（RTX 4070など）でのトレーニングを可能にするため、NanoEulerには、forward passとbackward pass、トレーニング、および推論を処理する完全なCUDAポートが含まれています。
カーネル実装
GPUエンジンは、TF32 tensor coresを介した行列乗算のために cuBLAS を利用しています。重要なパフォーマンス最適化は、手書きの FlashAttention カーネルであり、これはtilingとonline softmaxを使用して、$T imes T$ の行列をメモリに保存することを避けます。この実装により、トレーニング速度が約3倍向上しました。
検証と安定性
手書きのバックプロパゲーションは微細なエラーが発生しやすいため、NanoEulerは厳格な検証プロセスを採用しています。浮動小数点数のキャンセルによるエラーの隠蔽を防ぐため、すべての解析的勾配は、double精度での中心差分法によって比較されます。GPUカーネルはさらにCPUリファレンスと比較して検証され、勾配は $1e-6$ 以内に一致します。
トレーニングパイプライン: プリトレーニングからSFTまで
NanoEulerは、生のデータから機能的な（小規模な）チャットインターフェースに至るまでの、言語モデルの完全なライフサイクルを示しています。
データ取得
このプロジェクトは、データ収集におけるPython依存関係を回避するため、シェルスクリプトとDuckDB CLIを使用して以下を収集します：
- Books: Project Gutenbergからのパブリックドメインの古典作品。
- Web: FineWeb-Eduデータセットの高品質な教育用スライス。
- Instructions: 教師あり微調整のためのAlpacaデータセット。
2段階のトレーニングプロセス
- Pretraining: 約116Mのパラメータを持つベースモデルは、書籍とWebコーパスを組み合わせてトレーニングされます。この段階では、モデルに文法、レジスター、および基本的な言語パターンを学習させます。
- Supervised Fine-Tuning (SFT): プリトレーニング済みのベースは、Alpacaデータセットを使用して微調整されます。損失（loss）はレスポンス・トークンのみにマスクされます。つまり、モデルはプロンプトとパディングを無視して、回答を予測するようにトレーニングされます。これにより、モデルは、膨大な知識を必要とする深い世界知識を必要とせずに、アシスタントの「形式」を学習します（instruction $ ightarrow$ response）。
理論的基礎: 「Euler」のつながり
残差ネットワークと常微分方程式（ODE）の数学的関係から、このプロジェクトはLeonhard Eulerにちなんで名付けられました。
残差ブロックにおいて、$x = x + f(x)$ という操作は、数値積分における forward-Euler method を反映しています。ここで、$x(t+\Delta t) = x(t) + \Delta t imes f(x(t))$ です。ステップサイズが $\Delta t = 1$ の場合、深い残差ネットワークは離散化されたODEと見なすことができ、深さは積分時間を表し、各レイヤーは隠れ状態を1つのEulerステップ分だけ前方に積分します。