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 数据集。
两阶段训练过程
- Pretraining: ~116M 参数的基础模型在结合了书籍和网络语料库上进行训练。这一阶段让模型学习语法、语域 (register) 和基础语言模式。
- 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 步。