用 C++ 从零写一个能训练的 GPT：标量自动微分 + Arena 内存池

一、先说说这玩意儿是干嘛的

现在你想学 GPT 的原理，网上一搜全是 PyTorch 的教程，import 几个模块就完事了。但问题是，你虽然能跑通，却对"自动微分是怎么工作的"、"反向传播的梯度是怎么传回去的"一无所知------这些全被框架的层层封装盖住了。

这个项目走了一条完全不同的路：用纯 C++ 标准库，从零实现一个能训练的 GPT。不是调包，不是调用 PyTorch 的 C++ API，而是自己手写自动微分引擎、自己手写 Transformer 的每一层、自己手写 Adam 优化器。

最终结果是：一个 16 维嵌入、16 长度上下文、1 层 Transformer 的小网络，跑 10000 步训练只要 1.3 秒 。同样的逻辑用 Python CPython 跑要 20 多分钟，用 PyPy JIT 也要 3 分多钟------这个 C++ 版本比 PyPy 快了 150 倍。

更关键的是，整个项目只依赖 C++ 标准库，没有 PyTorch、没有 TensorFlow、没有 CUDA。所有东西------自动微分、矩阵乘法、注意力机制、优化器------全部自己手写，加起来不到 1000 行代码。

二、具体含义：它到底包含哪些东西？

这个项目把 GPT 的训练和推理链路拆成了几个自己动手做的核心模块：

模块	干啥的	为什么重要
Arena 内存池	用 SoA（Structure of Arrays）管理计算图的所有节点	内存连续、Cache 友好，是速度提升的关键
标量自动微分	每个标量操作都记录到计算图，反向传播时自动求梯度	不用手动写链式法则，代码量小且正确
Op 枚举优化	每个节点只存一个 `uint8_t` 的操作类型，反向时现场算梯度	省内存、省带宽，避免预存两个梯度指针
FMA 融合操作	把 `a*b+c` 三个节点合并成一个	减少计算图节点数，降低反向传播开销
单层 Transformer	多头因果注意力 + MLP + RMSNorm + 残差连接	GPT 的核心结构
KV Cache	推理时缓存历史 K/V，避免重复计算	放在栈上，零堆分配开销
Adam 优化器	带 bias correction 的自适应学习率	训练稳定收敛
Temperature 采样	控制生成随机性	让模型输出有变化，不死板

三、代码实现原理：拆开来看

1. 标量自动微分：每个数都是一个"节点"

传统的深度学习框架用张量作为基本单位，一个操作就是"两个矩阵相乘"。但这个项目走的是另一条路：标量自动微分 ------每个基本操作都是标量级别的，比如 a + b、a * b、exp(a)。

这样做的好处是代码极其直白 。你看 vadd(a, b) 的实现：

cpp 复制代码

inline int vadd(int a, int b) {
    return arena.push_binary_op(arena.data[a] + arena.data[b], a, b, OP_ADD);
}

它干了三件事：

计算 data[a] + data[b] 的值
在 Arena 里新建一个节点，记录这个值
标记这个节点的操作类型是 OP_ADD，子节点是 a 和 b

反向传播时，从 loss 节点倒着遍历，看到 OP_ADD 就知道：梯度要原封不动地分给两个子节点：

cpp 复制代码

case OP_ADD:
    p_grad[c0] += g;
    p_grad[p_c1[i]] += g;
    break;

看到 OP_MUL 就知道：梯度要乘上对方的值再传回去：

cpp 复制代码

case OP_MUL:
    p_grad[c0] += g * p_data[c1];
    p_grad[c1] += g * p_data[c0];
    break;

这就是自动微分的本质：前向时"记录操作"，反向时"根据操作类型现场算梯度"。

2. Arena 内存池：SoA 为什么比 AoS 快？

计算图里有成千上万个节点，每个节点需要存：

data：当前值
grad：梯度
i_child0, i_child1, i_child2：子节点索引
op：操作类型

传统的做法是 AoS（Array of Structs）：

cpp 复制代码

struct Node {
    float data, grad;
    int child0, child1, child2;
    Op op;
};
Node nodes[100000];

问题是：反向传播时，代码要遍历所有节点的 grad 字段。但 AoS 里 grad 是穿插在其他字段中间的，内存不连续，CPU Cache 命中率极低。

这个项目的做法是 SoA（Structure of Arrays）：

cpp 复制代码

float* data;      // data[0], data[1], data[2]... 连续
float* grad;      // grad[0], grad[1], grad[2]... 连续
int* i_child0;    // i_child0[0], i_child0[1]... 连续
int* i_child1;
int* i_child2;
Op* op;

反向传播时，用一个指针 p_grad 连续扫过 grad[] 数组，内存访问是线性的，Cache 命中率极高，编译器还能自动向量化。

3. Op 枚举：为什么只存一个 uint8_t？

很多自动微分框架会在每个节点里存两个梯度指针，指向子节点的梯度。但这样每个节点要多存 16 个字节（两个 float* 或 double*）。

这个项目的优化是：只存一个 uint8_t 的操作类型，反向传播时用 switch-case 现场算梯度。

cpp 复制代码

enum Op : uint8_t {
    OP_CONST, OP_ADD, OP_MUL, OP_DIV, OP_NEG,
    OP_LOG, OP_EXP, OP_RELU, OP_INV_SQRT,
    OP_SUB_CONST, OP_MUL_CONST, OP_DIV_CONST, OP_FMA
};

uint8_t 只占 1 个字节。反向传播时，根据 op 的值进不同的 case，现场计算局部梯度。这样省下来的内存带宽，直接转化成了速度。

4. FMA 融合操作：减少节点就是减少工作量

Transformer 里最常见的模式是 a * b + c------矩阵乘法的每个元素都是"两个数相乘，再加到累加器上"。如果拆成 vmul(a,b) 再 vadd(result, c)，计算图里会多出一个中间节点，反向传播时也要多走一趟。

这个项目定义了一个真正的 FMA（Fused Multiply-Add） 操作：

cpp 复制代码

inline int vmul_add(int a, int b, int c) {
    arena.ensure();
    int i = arena.size++;
    arena.data[i] = arena.data[a] * arena.data[b] + arena.data[c];
    arena.i_child0[i] = a;
    arena.i_child1[i] = b;
    arena.i_child2[i] = c;
    arena.op[i] = OP_FMA;
    return i;
}

一个节点干了三个节点的活。反向传播时：

cpp 复制代码

case OP_FMA:
    // result = a*b + c -> grad[a] += g*b, grad[b] += g*a, grad[c] += g
    p_grad[c0] += g * p_data[c1];
    p_grad[c1] += g * p_data[c0];
    p_grad[p_c2[i]] += g;
    break;

矩阵乘法里的内层循环全用这个 FMA，节点数直接砍半，反向传播的工作量也砍半。

5. Transformer 前向传播：单层也够用

虽然只实现了 1 层 Transformer，但结构是完整的：

复制代码

输入 token → Token嵌入 + 位置嵌入
    → RMSNorm
    → 多头因果自注意力（QKV投影 → 注意力计算 → 输出投影）
    → 残差相加
    → MLP（升维4x → ReLU → 降维）
    → 残差相加
    → LM Head 投影 → Softmax → 概率分布

训练时走交叉熵损失 -log(prob[target])，然后反向传播、Adam 更新权重。推理时加 Temperature 控制随机性，逐 token 采样生成。

If you need the complete source code, please add the WeChat number (c17865354792)

6. KV Cache 放在栈上

推理时的 KV Cache 是一个三维数组：

cpp 复制代码

using KVCache = int[N_LAYER][BLOCK_SIZE][N_EMBD];

注意它是固定大小的栈数组 ，不是 std::vector 动态分配的。栈内存的访问速度比堆内存快得多，而且不需要 malloc/free 的开销。

四、设计思路：为什么这样设计？

设计选择	为什么这么做	传统框架会怎么做
标量自动微分	代码最直白，每个操作一行就能看懂	张量级别的算子，背后几百行 C++
SoA Arena	Cache 友好，反向传播线性扫内存	AoS 或复杂的内存池，Cache Miss 多
Op 枚举 + 现场算梯度	1 字节 vs 16 字节，省内存带宽	每个节点存两个梯度指针
真正的 FMA	减少计算图节点数，反向工作量减半	拆成 mul + add 两个算子
f32 而不是 f64	速度更快，内存减半，精度够用	默认 f64（PyTorch）
栈上 KV Cache	零堆分配，访问速度最快	动态 vector，有扩容开销
`-Ofast -march=native`	让编译器大胆优化，生成 SIMD	保守的 `-O2`

这些优化的叠加效果非常惊人：从 Python 的 20 多分钟，压到了 1.3 秒。

五、性能长什么样？

在 Intel Core Ultra 7 165H 上跑 10000 步训练（16x16 网络）：

实现	时间	相对 PyPy
Python (CPython)	22 分 4 秒	~6.7x 慢
Python (PyPy JIT)	3 分 16 秒	1x（基准）
C++ 原始版	3.3 秒	~60x 快
C++ 增强版	2.2 秒	~88x 快
C++ 当前版	1.3 秒	~152x 快

如果把网络放大到 64x64（64 维嵌入、64 长度上下文），跑 1000 步：

实现	时间	相对 PyPy
Python (PyPy JIT)	6 分 47 秒	1x
C++ 当前版	1.6 秒	~249x 快

这说明优化的收益随网络变大而增加------内存带宽的优化在大网络上效果更明显。

六、代码实现用途

这个项目的用途很明确：

学习自动微分的底层原理：比 PyTorch 的 autograd 简单 100 倍，但核心思想完全一致
理解 Transformer 的每个操作：没有框架封装，每个矩阵乘法、每个 softmax 都是手写的
学习内存优化技巧：SoA、FMA、栈分配，这些技巧在生产代码里同样适用
作为更复杂项目的起点：在这个骨架上加多层、加 dropout、加更大的词表，就是一个完整的语言模型

七、怎么跑起来？一步步来

1. 运行

bash 复制代码

./microgpt

第一次运行会自动下载训练数据（一个英文名字数据集）：

复制代码

[LOG] Downloading input.txt ...
[LOG] We have 32033 names.
[LOG] Vocab size is: 27
[LOG] Number of params: 2432
[LOG] Step 1 / 1000 | loss 3.29581
[LOG] Step 2 / 1000 | loss 3.21874
...
[LOG] Step 1000 / 1000 | loss 0.82345

Time for inference---------------
[LOG] Sample: 0: emma
[LOG] Sample: 1: olivia
[LOG] Sample: 2: liam
...

2. 调整超参数

代码顶部有一堆 constexpr，你可以改这些来实验不同配置：

cpp 复制代码

constexpr int N_LAYER = 1;      // Transformer 层数
constexpr int N_EMBD = 16;      // 嵌入维度
constexpr int BLOCK_SIZE = 16;  // 上下文长度
constexpr int N_HEAD = 4;       // 注意力头数
constexpr int NUM_STEPS = 1000; // 训练步数

比如改成 64x64 网络：

cpp 复制代码

constexpr int N_EMBD = 64;
constexpr int BLOCK_SIZE = 64;
constexpr int NUM_STEPS = 1000;

3. 用自己的数据训练

把 input.txt 换成你自己的文本文件（每行一个样本），重新编译运行即可。注意：

字符集大小不能超过 MAX_VOCAB_SIZE（默认 27）
每行长度不要超过 BLOCK_SIZE

总结

这个项目最大的价值，是让你看到一个"能跑"的 GPT 可以有多小、多快、多直白。

它用不到 1000 行 C++ 代码，实现了：

完整的标量自动微分引擎
单层 Transformer（注意力 + MLP + RMSNorm）
Adam 优化器
KV Cache 推理
Temperature 采样生成

而且通过 SoA 内存布局、Op 枚举、FMA 融合等一系列优化，把训练速度推到了 Python 的 150 倍以上。

如果你想真正理解"深度学习框架的底层是怎么工作的"，而不是只会调包，这个项目就是最好的起点。把代码通读一遍，亲手改改嵌入维度、层数、学习率，你会对自动微分和 Transformer 有一个完全不一样的体感。

Welcome to follow WeChat official account【程序猿编码】