NanoEuler: 使用 C 和 CUDA 从零开始构建的 GPT-2 规模 LLM

NanoEuler: 使用 C 和 CUDA 从零开始构建的 GPT-2 规模 LLM

NanoEuler 是一个研究和教育项目,它完全使用 C 和 CUDA 从零开始实现了一个 GPT-2 级别的语言模型。通过避开 PyTorch 或 autograd 等高级 ML 库,该项目提供了一个透明的端到端流水线,涵盖了从手写的字节级 BPE 分词器到监督微调 (SFT) 的所有内容,并经过了严格的梯度检查验证。

核心架构与设计

NanoEuler 实现了一种仅解码器 (decoder-only) 的 transformer 架构,并利用了几种现代 LLM 优化技术,以在消费级硬件上实现效率最大化。该模型在设计时未包含偏置 (biases),以简化架构并提高训练稳定性。

技术规格

  • RMSNorm: 作为无偏置的预归一化使用。
  • Rotary Position Embeddings (RoPE): 应用于查询 (queries) 和键 (keys) 以处理位置信息。
  • SwiGLU Feed-Forward: 实现 down(silu(gate(x)) * up(x)) 激活模式。
  • Grouped-Query Attention (GQA): 通过在多个查询头之间共享较小的一组键/值头来优化内存和计算。
  • Multi-Token Prediction (MTP): 采用 K 个输出头来预测接下来的 K 个 token,这增强了学习到的表示并支持投机采样 (speculative decoding)(尽管生成过程主要使用 head 0)。
  • Byte-level BPE Tokenizer: 使用 GPT-2 风格的预分词实现的自定义实现,确保空格不会作为独立的 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 移植版本,负责处理前向和后向传播、训练和推理。

内核实现

GPU 引擎利用 cuBLAS 通过 TF32 tensor cores 进行矩阵乘法。一个关键的性能优化是手写的 FlashAttention 内核,它使用分块 (tiling) 和在线 softmax 来避免在内存中存储 $T imes T$ 矩阵。这种实现使训练速度提高了约 3 倍。

验证与稳定性

由于手写的反向传播容易出现细微错误,NanoEuler 采用了严格的验证过程。每个解析梯度都会与双精度下的中心有限差分进行比较,以防止浮点数抵消掩盖错误。GPU 内核还通过与 CPU 参考实现进行进一步验证,梯度匹配度在 $1e-6$ 以内。

训练流水线:从预训练到 SFT

NanoEuler 展示了一个语言模型的完整生命周期,从原始数据到功能性(尽管规模较小)的聊天界面。

数据获取

该项目避免了数据收集过程中的 Python 依赖,使用 shell 脚本和 DuckDB CLI 来收集:

  • Books: 来自 Project Gutenberg 的公有领域经典著作。
  • Web: 来自 FineWeb-Edu 数据集的高质量教育切片。
  • Instructions: 用于监督微调的 Alpaca 数据集。

两阶段训练过程

  1. Pretraining: ~116M 参数的基础模型在结合了书籍和网络语料库上进行训练。这一阶段让模型学习语法、语域 (register) 和基础语言模式。
  2. Supervised Fine-Tuning (SFT): 使用 Alpaca 数据集对预训练的基础模型进行微调。损失函数仅对响应 token 进行掩码 (masked),这意味着模型被训练为预测答案,同时忽略提示词 (prompt) 和填充 (padding)。这让模型学习到助手 (instruction $\rightarrow$ response) 的 形式,而不需要深厚的世界知识所需的庞大规模。

理论基础:“Euler” 的联系

该项目以 Leonhard Euler 命名,是因为残差网络与常微分方程 (ODEs) 之间的数学关系。

在残差块中,操作 $x = x + f(x)$ 镜像了数值积分中的 forward-Euler method,其中 $x(t+\Delta t) = x(t) + \Delta t imes f(x(t))$。当步长 $\Delta t = 1$ 时,一个深层残差网络可以被视为一个离散化的 ODE,其中深度代表积分时间,而每一层都在将隐藏状态向前推进一个 Euler 步。

Sources