NanoEuler: 使用 C 和 CUDA 從零開始構建的 GPT-2 規模 LLM

NanoEuler: 使用 C 和 CUDA 從零開始構建的 GPT-2 規模 LLM

NanoEuler 是一個研究與教育專案，它完全使用 C 和 CUDA 從零開始實現了一個 GPT-2 級別的語言模型。透過捨棄 PyTorch 或 autograd 等高階 ML 函式庫，該專案提供了一個透明且端到端的流程，涵蓋了從手寫的 byte-level BPE tokenizer 到監督式微調 (SFT)，並透過嚴格的梯度檢查進行驗證。

核心架構與設計

NanoEuler 實現了僅解碼器 (decoder-only) 的 transformer 架構，並利用了幾種現代 LLM 優化技術，以在消費級硬體上實現效率最大化。該模型在設計上不包含偏置 (biases)，以簡化架構並提高訓練穩定性。

技術規格

  • RMSNorm: 作為不含偏置的預歸一化 (pre-normalization)。
  • Rotary Position Embeddings (RoPE): 應用於 queries 和 keys 以處理位置資訊。
  • SwiGLU Feed-Forward: 實現了 down(silu(gate(x)) * up(x)) 激活模式。
  • Grouped-Query Attention (GQA): 透過在多個 query heads 之間共享較小的一組 key/value heads，來優化記憶體與計算。
  • Multi-Token Prediction (MTP): 採用 K 個輸出頭來預測下一個 K 個 token，這增強了學習到的表示，並支援投機解碼 (speculative decoding)（儘管生成主要使用 head 0）。
  • Byte-level BPE Tokenizer: 使用 GPT-2 風格的 pretokenization 的自定義實現，確保空格不會被浪費為獨立的 token。

模型配置

Version Dimension Q/KV Heads Layers Context Vocab Parameters
Small (CPU) 128 4 / 2 4 128 512 ~1.05M
GPU Pipeline 768 12 / 4 16 512 4096 ~116M

CUDA 引擎與性能

為了在單個消費級 GPU (例如 RTX 4070) 上進行訓練，NanoEuler 包含了完整的 CUDA 移植版本，負責處理前向傳播與反向傳播、訓練與推理。

Kernel 實現

GPU 引擎利用 cuBLAS 透過 TF32 tensor cores 進行矩陣乘法。一個關鍵的性能優化是手寫的 FlashAttention kernel，它使用 tiling 和 online softmax 來避免在記憶體中儲存 $T imes T$ 矩陣。此實現使訓練速度提升了約 3 倍。

驗證與穩定性

由於手寫反向傳播容易產生細微錯誤，NanoEuler 採用了嚴格的驗證流程。每個解析梯度 (analytic gradient) 都會與雙精度下的中心有限差分 (central finite difference) 進行比較，以防止浮點數抵消 (floating-point cancellation) 掩蓋錯誤。GPU kernel 亦會進一步與 CPU 參考值進行驗證，梯度匹配度在 $1e-6$ 以內。

訓練流程：從預訓練到 SFT

NanoEuler 展示了語言模型的完整生命週期，從原始數據到一個功能性（雖然規模較小）的聊天介面。

數據獲取

該專案避免了數據收集過程中的 Python 依賴，使用 shell scripts 和 DuckDB CLI 來收集：

  • Books: 來自 Project Gutenberg 的公有領域經典著作。
  • Web: 來自 FineWeb-Edu 數據集的優質教育切片。
  • Instructions: 用於監督式微調的 Alpaca 數據集。

兩階段訓練流程

  1. Pretraining: 約 116M 參數的基礎模型是在結合了書籍與網路語料庫上進行訓練的。此階段教導模型語法、語域 (register) 與基礎語言模式。
  2. Supervised Fine-Tuning (SFT): 預訓練的基礎模型使用 Alpaca 數據集進行微調。損失函數 (loss) 僅對回應 (response) tokens 進行遮罩 (masking)，這意味著模型被訓練去預測答案，同時忽略提示詞 (prompt) 與填充 (padding)。這教導了模型作為助手的「形式」(instruction $ ightarrow$ response)，而不需要深厚的知識量所需的龐大規模。

理論基礎：與「Euler」的聯繫

該專案以 Leonhard Euler 命名，是因為殘差網路 (residual networks) 與常微分方程 (ODE) 之間的數學關係。

在一個殘差塊中，操作 $x = x + f(x)$ 鏡像了數值積分中的 forward-Euler method，其中 $x(t+\Delta t) = x(t) + \Delta t \cdot f(x(t))$。當步長 $\Delta t = 1$ 時，一個深層殘差網路可以被視為一個離散化的 ODE，其中深度代表積分時間，而每一層都在將隱藏狀態向前積分一個 Euler 步長。

Sources