底层实现：训练框架
模型训练通常被压缩成几行代码：执行前向计算、得到损失、反向传播梯度，再由优化器更新参数。表面上，这只是对几个接口的依次调用；实际上，每一步都依赖前一步留下的数据和状态。把它们连接起来的是训练框架：它需要定义张量如何存储，记录算子之间的依赖，按照链式法则传递并累加梯度，还要组织模型参数，使同一套训练过程能够在不同模型和计算设备上稳定运行。
本书以 Zero 为基础，从一个可运行的训练框架出发，逐层拆解它的底层实现。Zero 保留了完成真实训练所需的主要结构，又将代码控制在能够直接阅读和追踪的范围内。重点不是复刻完整的 PyTorch，也不是罗列接口用法，而是沿着一次计算的实际路径，观察张量、自动微分、算子、模型和参数更新如何组成一个完整系统。在这个过程中，既要说明每个抽象解决什么问题，也要关注它们之间的数据所有权、形状约束和调用关系，让数学规则能够落实到具体代码。
全书首先介绍 Zero 的整体训练流程，随后进入自动微分、张量存储和算子实现。这三部分构成训练框架的计算基础，模型组织与参数更新在其上形成训练闭环，Python 前端则展示如何在保留 C++ 核心的同时提供更灵活的模型表达方式。最后两篇讨论 GPU 计算与 GPT 训练，进一步引入设备内存、CUDA Kernel、注意力和更复杂的训练任务，展示同一套框架如何从 CPU 上的简单模型扩展到 GPU 上的语言模型。
Zero 框架介绍
引言
从使用者的角度看，模型训练通常表现为一个不断重复的循环。数据和模型准备完成以后，每一轮只需要清理旧梯度、执行前向计算、计算损失、反向传播，再更新参数，核心代码通常只有几行：
import zero
optimizer.zero_grad()
predictions = model(images)
loss = zero.cross_entropy(predictions, labels)
loss.backward()
optimizer.step()
但这几行背后存在一条完整的数据链路：张量保存输入和参数，模型通过算子完成前向计算，算子在产生结果的同时记录依赖，损失再从计算图末端发起反向传播。优化器读取参数梯度并完成更新，下一轮计算随即从新的参数开始。
输入与参数
↓
模型调用算子
↓
预测结果 → 损失
↓ ↓
计算图 ← 记录依赖
↓
反向传播 → 参数梯度
↓
参数更新
Zero 将这条链路实现为一个小型训练框架。沿着一次训练迭代向下追踪，就能看到上层接口如何落到张量存储、算子执行和梯度计算。
实现
Zero 的计算核心使用 C++17 编写，同时支持 CPU 和 CUDA；轻量的 Python 前端则从中暴露完成基础训练所需的 CPU 功能。它参考 PyTorch 的动态图设计，但不追求覆盖完整的框架功能，也不以替代现有训练框架为目标。
它的定位是教学：用一套规模可控、能够实际运行的代码，把张量、自动微分、算子、模型、参数更新和 GPU Kernel 放在一起观察。各部分不仅具有独立接口，还通过明确的数据所有权、形状约束和执行顺序彼此连接。
Zero 可以在 CPU 上训练 MNIST，也可以使用同一套基础设施在 GPU 上训练 GPT。MNIST 用于验证最基本的训练闭环，GPT 则展示同一套框架扩展到注意力和更多参数后的形态。
张量存储
前向计算从 Tensor 开始。一个张量除了数据，还要保存形状、步长、数据类型和所在设备。View 与 Transpose 可以共享底层存储，却以不同的元数据解释它；CPU 和 CUDA 张量则由不同的内存上下文负责分配。
张量还承担自动微分中的身份。前向表达式产生的结果需要记住由哪个算子生成、依赖哪些输入。存储共享、视图关系与梯度关系必须分别处理，否则一次看似普通的浅拷贝或转置就可能破坏反向传播。
算子执行
每个可微算子都实现两件事：根据输入计算结果，以及根据上游梯度计算各输入的梯度。例如乘法 \(z=x\cdot y\) 的反向规则为：
$$ \frac{\partial L}{\partial x} =\frac{\partial L}{\partial z}y, \qquad \frac{\partial L}{\partial y} =\frac{\partial L}{\partial z}x $$
自动微分
Zero 在前向执行算子的同时动态建立计算图。调用 backward() 后，底层先得到拓扑顺序，再逆序执行每个算子的局部反向规则。一个张量经过多条路径影响损失时，来自各路径的梯度还要累加到一起。
同一个算子接口可以派发到 CPU 或 CUDA Kernel。设备差异被限制在数值实现和内存管理中，计算图仍然使用相同的节点和反向规则。
模型训练
Parameter 是需要优化的张量，Module 将参数和子模块组织成模型。Linear、MLP、Embedding、LayerNorm 和 Attention 都建立在张量算子之上，而不是绕开自动微分单独实现训练逻辑。
优化器持有参数集合。zero_grad() 清理上一轮梯度，step() 根据 SGD 或 AdamW 等规则修改参数。底层更新过程不会记录计算图，否则参数更新本身也会进入下一轮反向传播。
由此，一次训练迭代形成闭环：
Parameter
↓ 前向
Loss
↓ backward
Gradient
↓ optimizer.step
新的 Parameter
训练任务
MNIST 中的 MLP 已经会经过矩阵乘法、Bias 广播、激活函数、交叉熵、反向传播和参数更新，因此足以验证训练框架的基础结构是否正确。
GPT 没有改变这套基本过程，而是提高了对框架的要求：Embedding 引入索引操作，多头注意力引入批量矩阵乘法和因果遮罩，LayerNorm 与 FFN 产生更多中间张量，GPU 执行则需要考虑 Kernel 启动、内存池、归约和算子融合。
所以本书不是从一个简单框架跳到另一套 GPT 框架，而是沿着同一套训练框架逐步增加模型和设备复杂度。
Autograd 自动微分
引言
深度学习训练的目标，是找到一组使损失尽可能小的模型参数。梯度给出了损失相对于各个参数的变化率，优化器据此更新参数。为了得到这些梯度，训练框架必须解决两个问题：如何记录前向计算，以及如何沿着计算过程反向传播梯度。
Zero 使用反向模式自动微分。它在前向计算时动态构建计算图；调用 Backward() 后，从损失的梯度开始，按照与前向相反的顺序依次执行各个算子的反向规则。每个算子接收输出端传来的梯度，计算各个输入的梯度并继续向前传递；如果同一个张量参与了多条计算路径，来自不同路径的梯度还需要累加。这个过程一直进行到模型参数，得到参数更新所需的梯度。
链式法则
考虑复合函数：
$$ y=f(g(x)) $$
它对 \(x\) 的导数可以拆成两个局部导数的乘积：
$$ \begin{aligned} \frac{\partial y}{\partial x} &= \frac{\partial y}{\partial g} \frac{\partial g}{\partial x} \end{aligned} $$
训练引擎不需要一次求出整个模型的导函数。每个算子只需知道自己的局部导数，然后将上游梯度乘上局部导数即可。
例如 \(z=x\cdot y\) 的上游梯度为 \(g\)，乘法算子返回：
$$ \frac{\partial L}{\partial x}=g\cdot y, \qquad \frac{\partial L}{\partial y}=g\cdot x $$
如果一个张量被多个后续算子使用，那么它的总梯度是每条路径贡献之和：
$$ \begin{aligned} \frac{\partial L}{\partial x} &= \sum_i \frac{\partial L}{\partial y_i} \frac{\partial y_i}{\partial x} \end{aligned} $$
这也是反向传播中梯度必须累加，而不能直接覆盖的原因。
计算图
表达式 \(z=x\cdot y+x\) 可以写成一张有向无环图：
x ──┐
├── Mul ──┐
y ──┘ ├── Add ── z
x ────────────┘
图中的张量保存数值，算子保存局部的前向和反向规则。前向计算得到 z 时，还需要记住 Add 的输入；否则反向阶段无法继续找到 Mul、x 和 y。
设 \(x=3\)、\(y=2\)，前向计算得到 \(x\cdot y=6\)、\(z=9\)。反向计算从 \(\partial z/\partial z=1\) 开始：
∂z/∂z = 1
│
Add::Backward
┌─────┴─────┐
传给 Mul：1 传给 x：1
│
│
Mul::Backward
┌────┴────┐
传给 x：1 × y 传给 y：1 × x
│
└──→ x 的两条梯度相加
Add 将上游梯度分别传给两个输入，Mul 再根据局部导数计算 \(x\) 和 \(y\) 的梯度。由于 \(x\) 同时经过 Mul 和 Add 到达 \(z\)，它需要累加两条路径的贡献。
最终得到：
$$ \frac{\partial z}{\partial x}=y+1=3, \qquad \frac{\partial z}{\partial y}=x=3 $$
这个例子同时说明了两件事：算子只计算局部梯度，训练引擎负责按照计算图组合并累加它们。
Zero 在 Tensor 中使用 Back 保存这条连接：
struct Back {
std::unique_ptr<TensorOp> op;
std::vector<Tensor> inputs;
};
op 是产生当前结果的算子，inputs 是该算子的输入。Back 被当前逻辑张量的所有副本共享，算子的生命周期则由 unique_ptr 明确管理。
实现
动态建图
每次运算都会创建一个新的算子实例，再通过 TensorOp::Apply 执行：
Tensor Tensor::operator*(const Tensor& other) const {
return RunOp<OpMul>({*this, other});
}
Apply 先计算结果，然后检查输入是否需要梯度：
Tensor result = op->OpValue_(inputs);
bool need_grad = std::any_of(
inputs.begin(), inputs.end(),
[](const Tensor& t) { return t.RequiresGrad(); });
if (need_grad) {
result.RequiresGrad(true);
result.back_->op = std::move(op);
result.back_->inputs = std::move(inputs);
}
只有至少一个输入需要梯度时，结果才会记录计算图。推理计算和普通数据处理因此不会携带不必要的反向状态。优化器更新参数时使用 RunOpNoGrad，也是为了防止更新过程自身进入计算图。
反向传播
一个节点的梯度可能来自多条路径。在执行该节点的反向函数之前，必须先收集完所有下游贡献。Zero 首先从损失张量出发进行深度优先遍历，得到拓扑序，再逆序执行每个算子的 Backward。
对上面的例子，前向的依赖顺序是：
x, y → Mul → Add → z
反向传播则按照相反顺序处理：
z → Add → Mul → x, y
拓扑排序的意义不只是“倒着遍历”，而是保证处理某个节点时，它的所有下游梯度已经准备完成。
根节点的起始梯度为全 1 张量，对应：
$$ \frac{\partial L}{\partial L}=1 $$
反向遍历的核心过程可以概括为：
for (auto it = topo.rbegin(); it != topo.rend(); ++it) {
Tensor* current = *it;
if (!current->back_ || !current->back_->op) continue;
Tensor grad = RowMajorGradOf(*current);
auto input_grads = current->back_->op->Backward(
grad, current->back_->inputs);
for (size_t i = 0; i < current->back_->inputs.size(); ++i) {
auto& input = current->back_->inputs[i];
if (input.RequiresGrad()) {
AccumulateGrad(input, input_grads[i]);
}
}
}
第一次写入某个梯度槽时，AccumulateGrad 直接接管算子产生的存储；如果该梯度已经存在，则就地累加。这既满足多路径的链式法则，又避免了每次累加都分配新张量。
张量身份
Back::inputs 保存的是 Tensor 值副本。如果单纯使用对象地址识别计算图节点，同一个逻辑张量的不同副本就会被当成多个节点，导致上游梯度重复计算。
Zero 为每个逻辑张量设置一个共享的梯度槽：
struct GradHolder {
std::shared_ptr<TensorNode> grad_node;
};
普通 Tensor 副本共享 GradHolder，因此反向传播通过图中副本写入的梯度，用户持有的原张量也能立即看到。拓扑排序也以 GradHolder 的身份去重，而不是使用 Tensor* 或底层数据地址。
视图算子是一个特例。Transpose 和 View 可以共享数据存储，但形状和步长已经改变，因此必须使用独立的 GradHolder。张量的存储共享与梯度共享并不是同一件事，这一点会在下一篇详细讨论。
示例
有了动态计算图和反向传播，最小训练循环已经可以表达：
auto w = Parameter(Tensor::From<float>({0.1f}, {1}));
auto b = Parameter(Tensor::From<float>({0.0f}, {1}));
for (int step = 0; step < 200; ++step) {
w.ZeroGrad();
b.ZeroGrad();
auto pred = w * x + b;
auto diff = pred - y;
auto loss = diff * diff;
loss.Backward();
w.Update(w.Grad() * learning_rate);
b.Update(b.Grad() * learning_rate);
}
前向表达式在计算数值的同时建立计算图，Backward() 将损失梯度传回 w 和 b，最后再由优化器或参数更新函数修改数值。这条链路是训练引擎的核心，后续各篇将在它之上逐步扩展数据规模、模型结构和执行设备。
Tensor 张量存储
引言
前一篇使用张量保存计算结果，并将它们连接成计算图，但张量本身还没有展开。最简单的实现似乎只需要一个 std::vector<float>：申请一段连续内存，再把数值依次放进去。然而模型处理的不只是向量，还包括批量输入、权重矩阵和多维激活。同一组数值可能需要被变形或转置，也可能位于 CPU 或 GPU，并在反向传播时拥有对应的梯度。单独一段内存无法表达这些信息。
因此，张量不只是一组数值。它还需要记录形状、步长、数据类型、设备和自动微分状态，并明确这些信息在复制与变形时怎样共享。Zero 将“数据存储”与“数据解释”分开：底层内存只保存数值，元数据决定如何把它解释成多维数组。理解这个区别，是理解零拷贝视图和张量共享语义的起点。
多维数组
考虑一个两行三列的矩阵：
[[1, 2, 3],
[4, 5, 6]]
从数学上看，它有两个维度；从内存角度看，它仍然只是连续排列的六个数值。张量需要把这两种视角连接起来。Zero 使用统一的 Tensor 接口表示向量、矩阵和更高维数组：
Tensor a = Tensor::Zeros({2, 3});
Tensor b = Tensor::Ones({2, 3});
Tensor c = Tensor::Randn({2, 3});
Tensor flat = a.View({6});
Tensor transposed = a.T();
Tensor product = a.MatMul(b.T());
Tensor shifted = a + 5.0f;
{2, 3} 表示两行三列。a、b 和 c 具有相同形状，却可以采用不同的初始化方式；View 将六个元素重新解释成一维数组，T() 交换两个维度，MatMul 和加法则在这些结构之上执行计算。
这些接口背后需要处理不同的关系：View 和 Transpose 应该复用已有数据，计算结果需要拥有新的存储，普通张量副本又必须共享正确的计算图和梯度。如果这些边界没有定义清楚，减少一次数据拷贝就可能换来错误的布局或反向传播。张量设计的关键，正是确定哪些状态应该共享，哪些状态必须彼此独立。
实现
无论张量具有多少个维度，底层内存最终都是一段一维的连续地址。多维坐标不能直接用于访问这段内存，框架必须先把它转换成相对于起始位置的下标。以一个矩阵为例，沿列移动通常只需访问下一个元素，沿行移动则要跨过一整行；转置以后，两个方向的移动方式还会互换。形状描述每个维度包含多少个元素，步长描述沿各个维度移动一次需要跨过多少个内存位置，二者共同建立多维坐标与线性存储之间的对应关系。
形状与步长
TensorMeta 保存张量的元数据：
struct TensorMeta {
std::vector<int64_t> sizes;
std::vector<int64_t> strides;
Device device;
DataType datatype;
};
sizes 定义每个维度的长度，strides 定义对应坐标增加 \(1\) 时，内存下标需要跨过多少个元素。
对一个行优先存储的 \(2\times3\) 矩阵：
矩阵：[[1, 2, 3],
[4, 5, 6]]
内存：[1, 2, 3, 4, 5, 6]
sizes: [2, 3]
strides: [3, 1]
位置 \((i,j)\) 对应的线性下标为：
$$ \operatorname{offset}(i,j)=i\cdot3+j\cdot1 $$
例如 \((1,0)\) 对应下标 \(3\)，所以读到数值 \(4\)。
TensorMeta::InitStrides 从最后一维开始计算默认步长：
void InitStrides() {
strides.resize(sizes.size());
int64_t stride = 1;
for (int i = sizes.size() - 1; i >= 0; --i) {
strides[i] = stride;
stride *= sizes[i];
}
}
对形状 [2, 3, 4]，结果是 [12, 4, 1]。NumElements() 将所有维度相乘，NumBytes() 再乘以单个元素的字节数，得到需要分配的内存大小。
形状回答“每个维度有多长”，步长回答“沿某个维度移动时怎样访问内存”。二者共同决定张量的布局，也使同一块内存可以拥有不止一种解释方式。
零拷贝视图
转置一个 \(2\times3\) 矩阵时，不需要移动数据，只需交换形状和步长：
转置前：sizes=[2, 3], strides=[3, 1]
转置后：sizes=[3, 2], strides=[1, 3]
转置后的 \((0,1)\) 对应内存下标：
$$ 0\cdot1+1\cdot3=3 $$
因此读到原矩阵 \((1,0)\) 位置的数值 \(4\)。底层数据没有改变，改变的只是访问方式。
Zero 的 Transpose 通过 ViewWithMeta 创建这种视图：
Tensor Tensor::ViewWithMeta(TensorMeta new_meta) const {
Tensor view;
view.node_ = node_;
view.meta_ = std::move(new_meta);
view.grad_holder_ = std::make_shared<GradHolder>();
return view;
}
视图共享 node_，因此指向同一块数据；但它拥有自己的 meta_ 和 GradHolder。原张量和转置视图的梯度形状不同，如果共享梯度槽，反向传播就会以错误的步长解释梯度。
View 同样可以共享数据，但它有一个限制：当前布局必须能够用新形状合法表示。这样的设计避免了形状操作中不易察觉的数据拷贝。
数据存储
视图能够共享数据，是因为原始内存没有直接放在某个 Tensor 对象中，而是由独立的 TensorNode 管理。它只负责一块原始内存：
class TensorNode {
public:
TensorNode(TensorContext* ctx, size_t num_bytes)
: ctx_(ctx), data_ptr_(ctx_->NewMem(num_bytes)) {}
~TensorNode() {
if (data_ptr_) ctx_->DeleteMem(data_ptr_);
}
private:
TensorContext* ctx_;
std::byte* data_ptr_;
};
它在构造时申请内存，析构时归还内存，并禁止自身拷贝和移动，避免多个对象重复管理同一指针。Tensor 通过 shared_ptr<TensorNode> 共享它，最后一个引用消失后才真正归还内存。
反向传播还会在某个中间结果的数据不再被后续算子使用时，调用 ReleaseData() 提前归还存储。这能避免所有前向激活都一直保留到整张计算图被销毁。
内存分配器
一次前向和反向计算会产生大量中间张量。矩阵乘法、激活函数和梯度计算都需要临时结果，如果每次都向系统申请并立即释放内存，分配成本可能超过一些小算子本身的计算成本。Zero 通过 TensorContext 抽象屏蔽具体分配策略：
class TensorContext {
public:
virtual std::byte* NewMem(size_t size) = 0;
virtual void DeleteMem(std::byte* ptr) = 0;
};
CPU 分配器
CPUTensorContext 使用空闲链表保存已释放的内存块。新请求先按 \(16\) 字节对齐，再查找能容纳请求的最小空闲块。如果找到的块过大，则重新分配，避免一个小张量长期占用大块内存。空闲块数量过多时，多余内存会真正返回系统。
CUDA 分配器
cudaMalloc 不仅是分配操作，还会引入驱动层同步开销。CUDATensorContext 因此将请求向上对齐到 \(512\) 字节，并按大小放入不同的缓存桶：
申请内存 → 对齐到桶大小 → 命中空闲块？
├─ 是：直接复用
└─ 否：cudaMalloc
释放内存 → 放回对应的空闲桶
缓存的块默认不会立即返回 CUDA。当训练从大 batch 切换到单样本生成时，张量形状发生大幅改变，旧桶往往无法复用。此时可以调用 TrimFreePool() 将空闲块真正返回设备，为新阶段腾出显存。
张量组成
至此，数值由 TensorNode 保存，形状与步长由 TensorMeta 解释，计算图和梯度又有各自的生命周期。它们不能简单地绑在同一个对象中，否则一次浅拷贝或视图变换就会同时改变所有状态。
当前 Tensor 因此由四部分组成：
class Tensor {
std::shared_ptr<TensorNode> node_;
TensorMeta meta_;
std::shared_ptr<Back> back_;
std::shared_ptr<GradHolder> grad_holder_;
};
它们分别遵循不同的共享规则：
| 成员 | 作用 | 拷贝 Tensor 时 |
|---|---|---|
node_ | 底层数据 | 共享存储 |
meta_ | 形状、步长、设备和类型 | 按值拷贝 |
back_ | 产生当前结果的算子与输入 | 共享计算图记录 |
grad_holder_ | 梯度存储槽 | 共享梯度 |
因此 Tensor b = a 是浅拷贝：a 和 b 代表同一个逻辑张量，共享数据、计算图位置和梯度槽。这不仅避免了大块数据拷贝，也保证计算图中的副本写入梯度后，用户持有的张量可以读到相同结果。
视图则不同：它与原张量共享 node_，却有自己的 meta_、back_ 和 grad_holder_。回到开头的 \(2\times3\) 矩阵，转置视图和原张量读取的是同样六个数值，却拥有不同的形状、步长和梯度关系。这组共享语义决定了零拷贝视图和自动微分能否同时正确工作。
设备与类型
TensorOption 用于创建张量时指定设备、数据类型和是否需要梯度：
TensorOption option;
option.SetDevice(Device::CUDA())
.SetDataType(F32)
.RequiresGrad(true);
Tensor x({2, 3}, option);
Zero 目前支持 F32 和 I32，设备支持 CPU 和 CUDA。DefaultTensorContext::Get(device) 为两类设备分别返回默认分配器，Tensor::To(device) 则通过 DeviceTransfer 搬运数据及已有梯度。
这种设计使张量的上层接口不需要区分 CPU 和 CUDA 存储。无论数据位于哪种设备，张量都使用形状和步长解释内存，并通过相同的共享规则参与计算图。
张量至此解决了“数据是什么”以及“数据怎样存放”的问题。下一篇将继续处理“数据怎样计算”：算子如何读取不同布局的张量，并根据设备选择对应的 Kernel。
Operator 算子实现
引言
张量解决了数据如何表示和存储，但它本身并不知道两个张量能否相加、矩阵乘法会产生什么形状，也不知道计算结果应当怎样参与反向传播。这些规则由算子定义。一次看似简单的乘法，既要检查输入形状、分配输出并完成数值计算，也要记录输入之间的依赖，并在反向阶段根据上游梯度计算每个输入的梯度。
如果把这些工作全部放进 Tensor，数据结构会同时承担存储、计算和设备实现；如果每个算子分别处理 CPU 与 CUDA，形状规则和反向逻辑又容易重复。Zero 因此将一次算子调用分成三层：
Tensor 接口
↓
TensorOp：形状检查、建图、前向/反向规则
↓
Ops：按 Device 派发到 CPU 或 CUDA Kernel
Tensor 提供面向使用者的运算接口，TensorOp 定义算子的形状语义、前向与反向规则，并决定是否记录计算图；最底层的 Ops 只负责把数值计算派发到对应设备。三层之间各自处理一种变化：用户表达式保持不变，算子规则与设备实现可以独立扩展。
因此，同一段 Tensor 表达式既可以运行在 CPU 上，也可以迁移到 GPU，而自动微分仍然沿用相同的计算图和反向规则。下面先从连接这三层的 TensorOp 开始。
实现
算子抽象
TensorOp 只要求子类实现两个核心函数：
class TensorOp {
public:
static Tensor Apply(
std::unique_ptr<TensorOp> op,
std::vector<Tensor> inputs,
bool no_grad = false);
std::vector<Tensor> Backward(
const Tensor& grad_output,
const std::vector<Tensor>& inputs) const;
protected:
virtual Tensor OpValue_(
const std::vector<Tensor>& inputs) const = 0;
virtual std::vector<Tensor> Backward_(
const Tensor& grad_output,
const std::vector<Tensor>& inputs) const = 0;
};
OpValue_ 计算前向结果，Backward_ 接收结果的上游梯度，返回与输入一一对应的梯度张量。Apply 在前向完成后判断是否需要建图，而公开的 Backward 还会包装算子性能计时。
用户接口通过 RunOp 为每次调用创建独立算子：
template <typename Op, typename... Args>
Tensor RunOp(std::vector<Tensor> inputs, Args&&... args) {
return TensorOp::Apply(
std::make_unique<Op>(std::forward<Args>(args)...),
std::move(inputs));
}
Zero 不设置全局算子注册表，而是为每次前向调用创建独立的算子对象。算子只保存少量参数，构造成本相比张量计算可以忽略；算子的状态与生命周期则明确归属于它产生的计算图节点。
设备派发
TensorOp 不直接实现数值循环，而是调用 ops 命名空间中的统一入口。例如逐元素乘法：
inline void mul(const Tensor& a, const Tensor& b, Tensor& out) {
on_gpu(a) ? mul_cuda(a, b, out)
: mul_cpu(a, b, out);
}
inline void mul_backward(
const Tensor& grad, const Tensor& a, const Tensor& b,
Tensor& grad_a, Tensor& grad_b) {
on_gpu(grad)
? mul_backward_cuda(grad, a, b, grad_a, grad_b)
: mul_backward_cpu(grad, a, b, grad_a, grad_b);
}
算子层只处理形状、参数和自动微分，CPU 与 CUDA 的具体算法分别放在 tensor_ops_cpu.cc 和 tensor_ops_gpu.cu。这样新增算子时，边界非常明确：
- 在
ops.h和ops.cc定义前向、反向及形状语义。 - 在
tensor_ops_cpu.*实现 CPU Kernel。 - 在
tensor_ops_gpu.*实现 CUDA Kernel。 - 在
tensor_ops.h增加设备派发入口。
逐元素算子
加、减、乘、除的直接情况是两个输入形状相同。以乘法为例，前向计算为：
$$ z_i=x_i y_i $$
反向计算为：
$$ \begin{aligned} \frac{\partial L}{\partial x_i} &= \frac{\partial L}{\partial z_i}y_i, \qquad \frac{\partial L}{\partial y_i} &= \frac{\partial L}{\partial z_i}x_i \end{aligned} $$
相应的算子实现只需要分配结果与梯度，再调用设备无关的派发入口：
Tensor OpMul::OpValue_(const std::vector<Tensor>& inputs) const {
Tensor out(inputs[0].Sizes(), Tensor::LikeOption(inputs[0]));
ops::mul(inputs[0], inputs[1], out);
return out;
}
std::vector<Tensor> OpMul::Backward_(
const Tensor& grad, const std::vector<Tensor>& inputs) const {
Tensor grad_a(inputs[0].Sizes(), Tensor::LikeOption(grad));
Tensor grad_b(inputs[1].Sizes(), Tensor::LikeOption(grad));
ops::mul_backward(grad, inputs[0], inputs[1], grad_a, grad_b);
return {grad_a, grad_b};
}
实际代码还会检查输入数量和广播关系。这些检查属于算子语义，不应散落在 CPU 和 CUDA Kernel 中。
广播
当两个张量的形状不同时，逐元素算子可以按广播规则对齐它们。规则从最后一维向前比较，每一维需要满足以下条件之一：
- 两个维度相等。
- 其中一个维度为
1。 - 其中一个维度不存在，视为
1。
A: [2, 3, 4]
B: [ 3, 1]
→ [2, 3, 4]
A: [2, 3, 4]
B: [2, 4, 5]
→ 不兼容
Zero 当前会将较小的输入显式广播到目标形状，再执行形状相同的 Kernel。广播 Kernel 将输出线性下标还原为多维坐标，再映射回输入；输入中大小为 1 的维度，坐标始终为 0。
例如将 [10, 20, 30] 从 [3] 广播到 [2, 3]：
[[10, 20, 30],
[10, 20, 30]]
前向过程中的重复意味着反向过程中的求和。如果原始元素 \(x_j\) 被映射到多个输出位置，它的梯度为：
$$ \begin{aligned} \frac{\partial L}{\partial x_j} &= \sum_{i\mapsto j} \frac{\partial L}{\partial y_i} \end{aligned} $$
OpSumToSize 使用 broadcast_backward 将大形状的梯度收缩回原输入形状。因此所有二元算子都可以共用同一组流程：
前向：广播小张量 → 执行算子
反向：计算大形状梯度 → SumToSize 回原形状
线性层的 Bias 加法是一个高频特例。对于最后一维为 \(C\) 的输入和形状 [C] 的 Bias，OpAdd 会走专用路径，避免先生成一个完整的 Bias 广播中间张量。
矩阵乘法
对于矩阵 \(A\in\mathbb{R}^{M\times K}\) 和 \(B\in\mathbb{R}^{K\times N}\)，前向计算为：
$$ C_{mn}=\sum_{k=1}^{K}A_{mk}B_{kn} $$
其反向公式为：
$$ \begin{aligned} \frac{\partial L}{\partial A} &= \frac{\partial L}{\partial C}B^T, \qquad \frac{\partial L}{\partial B} &= A^T\frac{\partial L}{\partial C} \end{aligned} $$
Zero 同时支持普通矩阵乘法、批量矩阵乘法和 batch 维度广播：
[M, K] @ [K, N] → [M, N]
[B, M, K] @ [B, K, N] → [B, M, N]
[1, M, K] @ [B, K, N] → [B, M, N]
当 batch 维度需要广播时，前向先对齐两个输入；反向计算完成后，再将梯度收缩回各自的原始形状。
步长处理
线性层常用 x.MatMul(weight.T())。T() 是零拷贝视图，所以 weight.T() 不是默认行优先布局。如果 Kernel 用 m * K + k 这类连续下标读取它，就会把转置前的存储误当成转置后的矩阵。
正确的读取方式必须使用张量步长：
$$ \begin{aligned} \operatorname{offset}_A(m,k) &= m\cdot s_{A,-2}+k\cdot s_{A,-1} \end{aligned} $$
CPU 和 CUDA 的 MatMul Kernel 都必须遵守这条规则。这不是一个可选优化，而是计算转置视图时的正确性要求。
视图反向
View 和 Transpose 虽然只改变张量的形状或步长，不拷贝前向数据，但它们仍然是计算图中的算子。反向传播时，需要将视图形状下的梯度转换回输入的形状和布局。
Transpose 的反向过程需要将梯度的对应维度再交换一次：
auto grad_view = grad_output.ViewWithMeta(
SwapDims(grad_output.Meta(), dim0_, dim1_));
return {OpBroadcastTo(grad_view, inputs[0].Sizes())};
这里的 OpBroadcastTo 不是在逻辑上扩展形状，而是利用通用的 stride-aware 拷贝将视图梯度实体化为连续存储。
View 在输入连续时只重新解释形状；如果输入来自 Transpose 等非连续视图，Zero 会先生成行优先的连续副本，再改变形状。这样才能保证 Attention 中的 Transpose(...).View(...) 按正确顺序解释底层数据。
融合
基础算子让模型能够自由组合，但组合得越细，执行时需要保存的中间张量和计算图节点就越多。在 GPU 上，每个小算子通常还对应一次 Kernel 启动和一次显存读写；计算本身很少时，这些固定开销会变得格外明显。
减少调度
以线性变换为例，普通实现需要先完成矩阵乘法，再单独执行 Bias 加法：
x @ Wᵀ → 中间结果 → Bias Add → y
OpLinearBias 使用 cuBLASLt 的 Bias Epilogue，在矩阵乘法结束时直接完成加法：
x、W、b → LinearBias → y
这样可以减少一次 Kernel 启动，也不必把矩阵乘法结果完整写回显存后再读取。OpLayerNorm 采用相同思路，将均值、方差、归一化和仿射变换放进一个算子，避免由多个基础算子反复产生临时结果。
控制存储
融合还可以改变中间结果的保存方式。语言模型输出层会先计算形状为 [N, V] 的 Logits，再计算交叉熵；当词表较大时，这个张量会占用大量显存。OpLinearCrossEntropy 在前向阶段只返回每个 Token 的损失，不让 Logits 长期留在计算图中，反向传播时再根据输入和权重重新计算。
OpFFN 和 OpFlashAttention 也采用这种策略。它们在反向阶段重新计算部分激活，以减少需要跨越整个前向过程保存的中间张量。这是一种计算与显存之间的取舍：反向计算有所增加，但训练能够使用更大的 Batch 或序列长度。
融合边界
Zero 在保留基础算子的同时，只为高频且边界稳定的组合提供融合实现：
OpLayerNormOpLinearBiasOpFFNOpLinearCrossEntropyOpFlashAttention
融合算子仍然遵守 TensorOp 的统一接口：前向接收张量并返回结果，反向接收上游梯度并返回各输入的梯度。它没有改变模型的数学定义，只是把多个小算子的计算和存储安排合并到更粗的执行单元中。
如果所有表达式都改写为专用融合算子，框架会失去组合能力，代码也难以维护。因此，基础算子负责表达完整语义，融合算子只处理已经成为性能或显存瓶颈的关键路径。张量和算子由此构成计算基础，下一篇将在这套接口之上组织参数与模型层。
Module 模型组织
引言
自动微分能够沿计算图求出梯度，却不知道哪些张量代表训练数据，哪些张量需要在每轮训练后更新。对于只有权重和 Bias 的线性模型，可以手动保存两个张量；当模型包含多个线性层、归一化、注意力和重复堆叠的 Block 时，逐个管理参数很快就会变得困难。
模型还需要保留自身的层级关系。一个参数不仅是一块需要梯度的存储，还属于某个具体的层；优化器需要遍历所有参数，模型内部则需要稳定地持有参数和子模块。如果参数散落在普通 C++ 成员中，每增加一层，都要同步修改参数收集和生命周期管理代码。
Zero 用 Parameter 区分可训练张量，用 Module 注册参数和子模块。小模块可以继续组合成更大的模型，而 Module 能够递归展开整棵结构，为优化器提供统一的参数列表。模型的前向计算仍然由张量算子完成，自动微分也不需要理解模型层级。
实现
模型组织从单个可训练张量开始。首先需要让参数具有明确的梯度和更新语义，再由模块接管参数的生命周期，并递归管理更小的子模块。在此基础上，线性层等具体网络层只需定义自己的参数和前向过程；训练模式则负责控制 Dropout 这类在训练与评估阶段行为不同的模块。
Parameter
Parameter 是一种默认需要梯度的 Tensor：
class Parameter : public Tensor {
public:
Parameter(const Tensor& tensor) : Tensor(tensor) {
RequiresGrad(true);
}
void Update(const Tensor& delta) {
*this = RunOpNoGrad<OpSub>({*this, delta});
RequiresGrad(true);
}
};
训练数据和模型参数都用 Tensor 存储，但语义不同：数据通常不需要梯度，参数则需要在每轮反向传播后更新。Parameter 在构造时开启梯度跟踪，让这个区别成为类型自身的不变式。
Update 在 no_grad 模式下执行减法。如果参数更新也记录自动微分图，后一轮训练就会继续引用前一轮的图，导致计算图和内存不断增长。
Module
Module 同时保存直接参数和子模块：
class Module {
protected:
std::unordered_map<std::string,
std::unique_ptr<Parameter>> parameters_;
std::unordered_map<std::string,
std::shared_ptr<Module>> modules_;
};
参数由所属模块独占，因此使用 unique_ptr；子模块可以被组合和复用，使用 shared_ptr。RegisterParameter 和 RegisterModule 建立所有权关系，Parameters() 则递归遍历整棵模块树：
model
├── block0.attention.weight_q
├── block0.attention.bias_q
├── block0.ffn.fc1.weight
├── block0.ffn.fc1.bias
└── ...
返回值中同时包含层级化名称和 Parameter*。优化器不需要理解模型结构，只需接收这个扁平参数列表。
Linear
线性层实现仿射变换：
$$ y=xW^T+b $$
Zero 将权重存储为 [out_features, in_features]，使用时通过零拷贝转置视图参与矩阵乘法：
class Linear : public Module {
public:
Tensor OpValue(const Tensor& input) {
return input.MatMul(weight_->T()) + *bias_;
}
private:
Parameter* weight_;
Parameter* bias_;
};
weight_ 和 bias_ 是不拥有对象的原始指针，只用于快速访问；参数的生命周期由 Module::parameters_ 中的 unique_ptr 管理。
权重使用 He 初始化：
$$ W_{ij}\sim\mathcal{N}\left(0,\frac{2}{\operatorname{fan_in}}\right) $$
这种初始化适合后续使用 ReLU 或 GeLU 的网络，能减少层数增加时激活和梯度尺度的剧烈变化。
训练模式
部分模块在训练和推理时行为不同。例如 Dropout 在训练时随机丢弃激活，在评估时应该直接返回输入：
Tensor Dropout::OpValue(const Tensor& x) {
if (p_ == 0.0f || !IsTraining()) return x;
return RunOp<OpDropout>({x}, p_);
}
Zero 默认处于训练模式。验证或生成时可以在局部作用域中创建 EvalGuard：
{
EvalGuard guard;
auto output = model.OpValue(input);
} // 离开作用域后恢复原模式
RAII 保证即使中途返回或抛出异常，训练状态也能正确恢复。
示例
复杂模型由小模块组合而成。例如一个单隐藏层 MLP：
class SimpleMLP : public Module {
public:
SimpleMLP(int64_t input_size,
int64_t hidden_size,
int64_t output_size) {
fc1_ = std::make_shared<Linear>(input_size, hidden_size);
fc2_ = std::make_shared<Linear>(hidden_size, output_size);
RegisterModule("fc1", fc1_);
RegisterModule("fc2", fc2_);
}
Tensor OpValue(const Tensor& input) {
return fc2_->OpValue(ReLU(fc1_->OpValue(input)));
}
private:
std::shared_ptr<Linear> fc1_;
std::shared_ptr<Linear> fc2_;
};
同样的组合方式也用于 FFN、MultiHeadAttention 和 TransformerBlock。各层只提供自己的前向表达式，参数收集和自动微分由统一机制完成。
调用 model.Parameters() 会递归得到：
fc1.weight
fc1.bias
fc2.weight
fc2.bias
模型内部仍然保留由子模块构成的层级结构，对外则提供一个带名称的扁平参数列表。下一篇中的优化器只需要遍历这个列表，无须了解参数属于线性层、注意力还是其他模块。
模型完成前向计算并收集参数以后，还需要由优化器统一读取梯度并更新这些参数。
Optimizer 参数更新
引言
反向传播只负责计算参数梯度，并不会自动修改参数。训练框架还需要决定怎样使用梯度，以及每次更新采用多大的步长，这些工作由优化器完成。
优化器不需要理解模型的层级结构，只需遍历 Module::Parameters() 返回的扁平参数列表。除了执行具体的更新公式，它还要清理上一轮梯度，并为 AdamW 这类算法保存跨越多个训练步骤的状态。参数更新本身不能进入自动微分图，否则不同训练步骤的计算图会被连接起来。
完整训练流程可以概括为：
清空梯度 → 前向传播 → 计算损失 → 反向传播 → 更新参数
实现
参数更新并不只是把某个优化公式翻译成代码。优化器需要为参数保存跨越多个训练步骤的状态，训练循环还可能根据当前进度调整学习率。二者共同决定参数在这一轮发生多大变化。
Zero 用 Optimizer 统一管理参数和梯度清理，并使用 AdamW 完成参数更新。AdamW 为每个参数保存一阶、二阶动量，其计算由张量算子实现，因此同一份代码可以运行在 CPU 和 CUDA 上。SetLr 则让训练循环能够在每一步改变学习率。
下面先从统一接口开始，再说明 AdamW 的更新过程和学习率变化。
Optimizer
Optimizer 基类持有扁平参数列表，并提供统一的梯度清零接口：
class Optimizer {
public:
void InsertParameters(
const std::vector<Module::NameParameter>& parameters);
virtual void ZeroGrad();
virtual void Step() = 0;
};
ZeroGrad() 遍历参数并清空上一轮的梯度，Step() 由具体优化算法实现。因为反向传播会累加梯度，每轮训练开始前必须先调用 ZeroGrad()。
AdamW
AdamW 为每个参数维护一阶动量 m 和二阶动量 v：
$$ \begin{aligned} m_t &= \beta_1m_{t-1}+(1-\beta_1)g_t \\ v_t &= \beta_2v_{t-1}+(1-\beta_2)g_t^2 \end{aligned} $$
使用偏差修正后的更新量为：
$$ \begin{aligned} \alpha_t &= \eta\frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \\ \Delta\theta &= \alpha_t\frac{m_t}{\sqrt{v_t+\epsilon}} +\eta\lambda\theta \\ \theta &\leftarrow\theta-\Delta\theta \end{aligned} $$
权重衰减项 \eta\lambda\theta 与梯度动量分开，这是 AdamW 与将 L2 正则项直接加入梯度的 Adam 之间的关键区别。
AdamW 的更新完全由张量算子构成，所以同一份代码可以运行在 CPU 和 CUDA 上。所有更新算子都使用 RunOpNoGrad，不记录到模型的自动微分图中。
学习率调度
AdamW::SetLr 允许训练循环在每步更新学习率。GPT 训练使用线性 warmup 与余弦衰减：
$$ \eta_t= \begin{cases} \eta_{\max}\dfrac{t+1}{T_w}, & t<T_w \\ \eta_{\min}+\dfrac{1}{2}(\eta_{\max}-\eta_{\min}) \left(1+\cos(\pi r_t)\right), & t\ge T_w \end{cases} $$
其中 \(T_w\) 是 warmup 步数，\(r_t\) 是衰减阶段当前进度。warmup 避免训练初期在动量统计尚不稳定时使用过大步长，余弦衰减则让后期更新逐渐变小。
示例
MNIST 可以用一个很小的模型验证参数更新是否真正有效：如果前向计算、反向传播或优化器中的任意一环出错，损失就无法稳定下降。这个示例不依赖复杂网络，能够把注意力集中在一轮训练如何完成。
数据
MNIST 的每张图像包含 \(28\times28\) 个灰度像素。数据加载器读取 IDX 二进制文件，将其中的大端整数转换为主机字节序，再把像素从 [0, 255] 归一化到 [0, 1]。一个 Batch 最终形成形状为 [B, 784] 的 F32 张量，标签则使用形状为 [B] 的 I32 张量。
图像：[B, 28, 28] → 展平 → [B, 784]
标签：[B]
F32 图像参与矩阵计算，I32 标签用于在交叉熵中索引正确类别。两者分开表示，也说明张量的数据类型是算子语义的一部分。
模型
示例使用一个单隐藏层 MLP：
[B, 784] → Linear(784, 512) → ReLU
→ Linear(512, 10) → Logits [B, 10]
它直接复用上一篇的 Module 和 Linear：
class SimpleMLP : public Module {
public:
SimpleMLP(int64_t input_size,
int64_t hidden_size,
int64_t output_size) {
fc1_ = std::make_shared<Linear>(input_size, hidden_size);
fc2_ = std::make_shared<Linear>(hidden_size, output_size);
RegisterModule("fc1", fc1_);
RegisterModule("fc2", fc2_);
}
Tensor OpValue(const Tensor& input) {
return fc2_->OpValue(ReLU(fc1_->OpValue(input)));
}
private:
std::shared_ptr<Linear> fc1_;
std::shared_ptr<Linear> fc2_;
};
model->Parameters() 会递归收集两个线性层的权重和 Bias，并将同一份参数列表交给设备迁移与优化器：
auto all_params = model->Parameters();
训练循环
训练循环的核心代码只有几步：
optimizer->ZeroGrad();
auto logits = model->OpValue(images);
auto loss = CrossEntropyLoss(logits, targets);
loss.Backward();
optimizer->Step();
ZeroGrad() 清除上一个 Batch 的梯度，前向过程将 [B, 784] 的输入转换成 [B, 10] 的 Logits，CrossEntropyLoss 根据标签产生损失。Backward() 将梯度传回四个参数，Step() 再根据当前优化算法更新它们。下一个 Batch 会使用更新后的参数重新开始这条链路。
设备选择
CPU 和 GPU 使用相同的 AdamW。如果选择 CUDA，需要先将模型参数迁移到 GPU，再创建优化器状态；每个 Batch 的输入也要迁移到同一设备：
if (use_cuda) {
for (auto& [name, parameter] : all_params) {
parameter->Cuda();
}
}
auto optimizer = std::make_unique<AdamW>(0.001f);
optimizer->InsertParameters(all_params);
if (use_cuda) images.Cuda();
AdamW::InsertParameters() 会为每个参数创建一阶和二阶动量，因此应在参数迁移之后调用，保证优化器状态与参数位于同一设备。除此之外，模型结构、损失函数和训练循环不需要因设备不同而改变。
训练过程中可以记录每个 Batch 的损失和分类准确率。损失持续下降说明参数更新方向有效，准确率上升则说明这些更新确实改善了模型输出。这个示例的价值不在于模型规模，而在于它打通了数据、模型、损失、自动微分和优化器之间的完整链路。
这条训练链路不依赖具体设备，下一篇将说明相同算子如何派发到 CUDA，并在 GPU 上高效执行。
Python 前端
引言
前面的章节已经在 C++ 中实现了张量、自动微分、算子和参数更新。这些能力适合承担数值计算，却不一定适合直接描述模型和训练流程。C++ 需要显式处理类型、所有权和编译过程；Python 则更适合快速组合网络层、准备输入并控制训练循环。
训练代码经常需要反复调整模型结构、损失函数、学习率和数据处理方式。如果这些变化都写在 C++ 中，每次修改后都要重新编译，再运行程序观察结果。Python 可以直接执行新的模型与训练逻辑，也便于在运行过程中打印张量形状、损失和梯度，缩短发现问题、修改代码和重新验证之间的周期。只有底层算子或绑定发生变化时，才需要重新编译 C++ 扩展。
Python 前端并不是重新实现一套训练框架，而是在 C++ 核心之上提供另一种表达方式。两层之间的关系可以概括为：
Python：模型组合与训练流程
↓
PyBind11：类型与调用转换
↓
C++：Tensor、Autograd 与 Operator
Python 创建张量并调用算子，实际的数值计算和计算图仍由 C++ 完成。这样既保留底层实现的性能和统一语义，又让上层训练代码保持简洁。
实现
Python 前端需要解决两个问题：首先将 C++ 的张量和算子安全地暴露给 Python，然后用这些基础接口组织模型与参数更新。绑定层只负责跨越语言边界，Module、Linear 和 SGD 则使用普通 Python 代码完成组合。
PyBind11
Zero 使用 PyBind11 将 C++ 类型封装为 _zero 扩展模块。绑定层暴露 Tensor、Parameter 和常用算子，并负责在两种语言之间转换参数与返回值：
py::class_<Tensor>(m, "Tensor")
.def("backward", &Tensor::Backward)
.def_property_readonly("grad", [](const Tensor& tensor) {
return tensor.Grad();
})
.def("__matmul__", [](const Tensor& lhs, const Tensor& rhs) {
return lhs.MatMul(rhs);
});
当 Python 执行 loss.backward() 时，调用会直接进入 C++ 的 Tensor::Backward()。矩阵乘法、加减乘、转置、变形、ReLU 和交叉熵也沿用同一套 C++ 算子，因此绑定前后共享相同的前向与反向规则。
Python 列表创建张量时，绑定层先检查嵌套列表是否具有规则形状，再将数值展开为连续数组。张量返回 Python 时，则根据 shape 重新构造嵌套列表。这个边界只负责数据表示的转换，不参与数值计算。
Module
Python 层的 Module 保持得很小。__call__ 将调用转发给 forward，parameters() 则递归收集对象属性中的 Parameter 和子模块：
class Module:
def parameters(self):
result = []
for value in vars(self).values():
if isinstance(value, Parameter):
result.append(value)
elif isinstance(value, Module):
result.extend(value.parameters())
return result
def __call__(self, *args, **kwargs):
return self.forward(*args, **kwargs)
Linear
Linear 由参数和已有张量运算组成：
class Linear(Module):
def __init__(self, in_features, out_features):
scale = math.sqrt(2.0 / in_features)
self.weight = Parameter(
randn([out_features, in_features]) * scale
)
self.bias = Parameter(randn([out_features]) * 0.0)
def forward(self, x):
return x @ self.weight.T + self.bias
前向过程不需要特殊的模型执行器。每一次 Python 运算都会调用对应的 C++ 算子，并动态建立计算图。
SGD
Python 版 SGD 保存 model.parameters() 返回的参数。zero_grad() 清除上一轮梯度，step() 按学习率生成更新量：
def step(self):
for parameter in self.parameters:
if parameter.grad is not None:
parameter.update(parameter.grad * self.lr)
Parameter.update() 最终仍在 C++ 中完成无梯度的减法。这样，Python 可以控制训练循环，而参数存储、梯度和更新语义仍由底层统一管理。
示例
下面的线性回归同时经过 Python 模型、C++ 算子、自动微分和参数更新：
import zero
model = zero.nn.Linear(2, 1)
optimizer = zero.optim.SGD(model.parameters(), lr=0.05)
x = zero.tensor([[1.0, 2.0], [2.0, 1.0]])
target = zero.tensor([[5.0], [4.0]])
for _ in range(100):
prediction = model(x)
difference = prediction - target
loss = (difference * difference).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()
model(x) 从 Python 进入绑定层，并调用 C++ 的矩阵乘法、转置和加法；loss.backward() 沿 C++ 计算图生成参数梯度；optimizer.step() 再通过 Parameter.update() 修改底层参数。循环仍由 Python 控制，但每一步张量计算都复用了同一套 C++ 实现。
Python 前端由此展示了训练框架的一条边界：底层提供稳定的计算能力，上层提供灵活的表达方式。接口不需要覆盖完整的 PyTorch，只要足以把前面实现的组件组合成一次可运行的训练即可。
GPU 计算
引言
CPU 适合复杂控制流和低延迟任务，GPU 则通过大量计算单元同时处理数据。深度学习中的逐元素运算、归约和矩阵乘法往往需要对大量元素执行相同或相似的计算，因此能够利用 GPU 的并行能力。
将训练迁移到 GPU 并不只是把张量复制到显存。框架还要把计算划分给线程，管理设备内存，选择适合的矩阵乘法实现，并处理 Kernel 的异步执行。与此同时，GPU 版本必须遵守与 CPU 相同的形状、步长和反向传播规则，否则同一模型会在两种设备上产生不同结果。
Zero 不在模型层区分 CPU 和 GPU。张量记录自己所在的设备，算子根据设备选择对应实现；模型、计算图和参数更新仍然使用相同接口。设备差异被限制在内存管理、数值 Kernel 和派发过程之中。
实现
GPU 执行从 CUDA 的线程模型开始，再逐步进入设备抽象、显存管理和算子派发。逐元素计算可以直接映射到大量线程，归约与矩阵乘法则需要更专门的并行策略；异步执行最终把这些操作连接成一条连续的设备任务流。
CUDA
CUDA Kernel 是一个由 CPU 发起、在 GPU 上并行执行的函数：
__global__ void scale_kernel(float* data, float scale, size_t size) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) {
data[idx] *= scale;
}
}
线程按两层结构组织：
Grid
├── Block 0
│ ├── Thread 0
│ ├── Thread 1
│ └── ...
├── Block 1
└── ...
Zero 的一维逐元素 Kernel 通常使用 256 个线程组成一个 Block，Block 数向上取整：
const int block = 256;
const int grid = (size + block - 1) / block;
scale_kernel<<<grid, block>>>(data, scale, size);
最后一个 Block 可能只有部分线程对应有效元素，因此 Kernel 内部必须检查 idx < size。GPU 还会以 \(32\) 个线程为一个 Warp 调度指令，同一 Warp 中过多的条件分支会降低执行效率。
并行模式
逐元素运算
如果每个输出元素只依赖相同位置的输入，就可以让每个线程处理一个元素。加法、乘法、GeLU 和梯度缩放都属于这一类：
template<typename T>
__global__ void add_kernel(
const T* a, const T* b, T* out, size_t size) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) out[idx] = a[idx] + b[idx];
}
归约
均值、方差、Softmax 和梯度范数需要将多个输入合并为更少的输出。归约 Kernel 通常让每个线程先计算局部结果，再在 Block 内使用共享内存或 CUB 完成合并。
输入元素
↓ 每个线程处理一部分
线程局部结果
↓ BlockReduce
Block 结果
↓ 必要时再次合并
最终结果
广播反向也是归约。多个大张量位置可能映射到同一个小张量元素，Zero 使用原子加法将这些梯度累加到同一位置。
Device
Device 记录设备类型和设备编号：
class Device {
public:
enum class DeviceType { CPU, CUDA };
static Device CPU();
static Device CUDA(int index = 0);
bool IsCpu() const;
bool IsCuda() const;
};
创建张量时，TensorOption 将设备信息传给 TensorMeta，再由 DefaultTensorContext::Get(device) 选择对应的内存分配器：
auto option = TensorOption()
.SetDevice(Device::CUDA())
.SetDataType(F32);
Tensor x({1024, 1024}, option);
Tensor::Cuda() 和 Tensor::Cpu() 通过 Tensor::To 在设备之间搬运数据。如果梯度已经存在，梯度存储也会一起迁移：
auto x = Tensor::Randn({1024});
x.Cuda();
auto y = x * 2.0f;
y.Cpu();
CPU 到 CUDA 使用 cudaMemcpyHostToDevice，CUDA 到 CPU 使用 cudaMemcpyDeviceToHost。同设备之间的迁移由 Tensor::To 直接忽略，不进入 DeviceTransfer。
GPU 内存
GPU 上的频繁 cudaMalloc 会引入显著的驱动和同步开销。CUDATensorContext 因此将申请大小对齐到 \(512\) 字节，并按大小缓存已释放的显存块。
NewMem(size)
→ RoundUp(size, 512)
→ 尝试从同尺寸空闲桶取块
→ 未命中时才调用 cudaMalloc
DeleteMem(ptr)
→ 放回对应尺寸的空闲桶
分配器同时统计 live_bytes、peak_bytes、reserved_bytes 和缓存命中次数。当训练切换到不同形状的生成阶段时，TrimFreePool() 可以释放已无法复用的旧桶。
这部分与前文的张量存储属于同一套机制：TensorNode 只负责持有指针，TensorContext 决定指针如何分配和复用。
Kernel 派发
算子通过 tensor_ops.h 中的统一入口调用 Kernel：
inline void add(const Tensor& a, const Tensor& b, Tensor& out) {
on_gpu(a) ? add_cuda(a, b, out)
: add_cpu(a, b, out);
}
公开的张量接口不包含 CUDA 分支：
auto a = Tensor::Ones({1024});
auto b = Tensor::Ones({1024});
auto cpu_result = a + b;
a.Cuda();
b.Cuda();
auto gpu_result = a + b;
派发器以第一个输入所在设备为准，因此同一算子的所有输入必须位于同一设备。如果混用 CPU 和 CUDA 张量，算子可能把主机指针传给 CUDA Kernel，造成非法内存访问。进入计算前应先显式将输入迁移到相同设备。
广播等高维 Kernel 需要获取形状和步长。Zero 将最多八维的定长 DimsArg 直接作为 Kernel 参数传递，避免每次运算都为元数据额外执行 cudaMalloc 和 cudaMemcpy。
OpAdd
把两个同形状的 CUDA 张量相加，可以看到一次算子调用怎样穿过前面的各层：
a + b
↓ Tensor::operator+
RunOp<OpAdd>
↓ TensorOp::Apply
OpAdd::OpValue_
↓ ops::add
add_cuda
↓
add_kernel<<<Grid, Block>>>
入口仍然是普通的 C++ 运算符：
Tensor Tensor::operator+(const Tensor& other) const {
return RunOp<OpAdd>({*this, other});
}
RunOp 为这次调用创建独立的 OpAdd，TensorOp::Apply 再执行它的前向过程。对于形状相同的输入，OpAdd::OpValue_ 创建一个同形状输出，并继承输入的设备与数据类型：
Tensor out(a.Sizes(), Tensor::LikeOption(a));
ops::add(a, b, out);
由于 a 位于 CUDA，输出内存由 CUDATensorContext 分配，ops::add 也会选择 add_cuda。CUDA 入口根据张量的数据类型选择模板实例，最后启动逐元素 Kernel：
template<typename T>
__global__ void add_kernel(
const T* a, const T* b, T* out, size_t size) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) out[idx] = a[idx] + b[idx];
}
每个线程只负责一个输出位置。Kernel 启动以后，TensorOp::Apply 不必等待 GPU 完成全部计算；如果输入需要梯度，它会把 OpAdd 和两个输入保存到结果的计算图记录中。
调用 Backward() 后，这条路径会反向执行：
Tensor::Backward
↓ OpAdd::Backward_
ops::add_backward
↓ add_backward_cuda
add_backward_kernel
↓
累加到 a.grad 与 b.grad
加法对两个输入的局部导数都是 \(1\)，因此反向 Kernel 将上游梯度分别写入 grad_a 和 grad_b。这两个梯度张量同样分配在 CUDA 上，随后由自动微分系统累加到输入的梯度槽。至此，一次 a + b 已经经过输出分配、设备派发、前向计算、计算图记录和反向传播，而模型层始终不需要出现 CUDA 分支。
矩阵乘法
逐元素算子适合自行编写 Kernel，矩阵乘法则需要更复杂的分块、数据搬运和硬件调度。Zero 使用 cuBLAS 执行 F32 MatMul：
- 普通矩阵使用
cublasGemmEx。 - 批量矩阵使用
cublasGemmBatchedEx。 - 线性层使用 cuBLASLt 的 Bias Epilogue，将矩阵乘法和 Bias 加法融合。
cuBLAS 以列优先布局解释矩阵，Zero 的张量默认为行优先。实现通过交换乘数和转置标志，在不拷贝连续数据的前提下完成两种布局之间的对应。
对于内层维度不是单位步长的视图，cuBLAS 无法用一个 leading dimension 表示其布局。Zero 会先将这类输入实体化为连续张量，再交给 cuBLAS。这与前文“算子必须正确处理 stride”的要求一致。
默认情况下，输入、输出和内部计算都使用 FP32。设置 MX_FAST_MATMUL=1 后，cuBLAS 可在保持 FP32 输入输出的同时，使用 FP16 Tensor Core 执行内部乘加。这能提高吞吐量，但会引入额外数值误差，因此需要显式开启。
异步执行
CUDA Kernel 启动默认是异步的。CPU 将 Kernel 放入默认 Stream 后即可继续运行，同一 Stream 中的 Kernel 会按提交顺序执行。Zero 不在每次 Kernel 启动后调用 cudaDeviceSynchronize()，否则 CPU 会在每个小算子后等待 GPU，丧失异步执行的优势。
cudaMemcpy 在 CPU 和 GPU 之间搬运数据时会形成必要的同步点。因此训练循环应尽量让数据和中间结果停留在 GPU 上，只在需要记录损失或输出结果时复制少量数据回 CPU。
至此，训练框架已经具备从张量到 GPU 执行的完整链路。最后一篇将用 GPT 把这些组件组合成一个真实的语言模型训练程序。
GPT 训练
引言
GPT 是一种自回归语言模型。给定一段 Token 序列，模型根据已经出现的上下文预测下一个 Token。Token 可以表示字符、子词或其他离散单元；模型只接收整数 ID，并不依赖词表采用哪种构造方式。
这一篇将前面的张量、自动微分、模块、优化器和 GPU Kernel 连成一条完整训练链路。
自回归目标
对序列 \(x_1,x_2,\ldots,x_T\)，联合概率可以按条件概率展开：
$$ P(x_1,x_2,\ldots,x_T) = \prod_{t=1}^{T}P(x_t\mid x_1,\ldots,x_{t-1}) $$
训练数据通过错位一个 Token 构造输入和目标：
原序列： [x0, x1, x2, x3, x4]
输入 x： [x0, x1, x2, x3]
目标 y： [x1, x2, x3, x4]
模型在每个位置输出词表上的 Logits，再使用交叉熵最小化正确下一 Token 的负对数似然：
$$ \begin{aligned} \mathcal{L} &= -\frac{1}{BT} \sum_{b=1}^{B}\sum_{t=1}^{T} \log P(y_{b,t}\mid x_{b,1:t}) \end{aligned} $$
\(B\) 是 Batch 大小，\(T\) 是序列长度。一次前向计算可以同时训练所有位置，但必须防止某个位置看到未来的 Token。
模型
GPT 接收一批 Token ID，输出每个位置对下一 Token 的概率预测。要完成这个转换，模型首先需要把离散 ID 映射为连续向量，再让每个位置从已有上下文中提取信息，最后把隐藏表示投影回词表空间。三个阶段分别对应 Embedding、Transformer Block 和输出层。
上下文建模是其中的核心。单个 Token 的向量只包含自身和位置信息，经过多头注意力后才能聚合前文；FFN 再对每个位置的表示进行非线性变换。多个 Block 重复这一过程，使当前位置逐层形成对已有序列的表示。因果掩码贯穿所有注意力层，保证这种表示始终只依赖当前位置及其之前的 Token。
整体结构
Zero 使用 Transformer Decoder 结构，将这些阶段连接起来：
Token ID
→ Token Embedding + Position Embedding
→ Dropout
→ TransformerBlock × N
→ LayerNorm
→ Token Embedding 权重的转置
→ Logits
输入张量形状为 [B, T]，其中每个元素是一个 Token ID。经过嵌入后，张量变为 [B, T, C]，\(C\) 是隐藏维度；所有 Transformer Block 都保持这个形状，因此残差分支可以直接相加。最终输出层将隐藏维度 \(C\) 投影到词表维度 \(V\)，得到形状为 [B, T, V] 的 Logits。
Embedding
Token ID 是离散整数，不能直接参与神经网络运算。Embedding 使用索引从参数矩阵中选取对应行：
$$ H_0=E_{\mathrm{token}}(x)+E_{\mathrm{position}}(0,1,\ldots,T-1) $$
Token 嵌入矩阵形状为 [V, C]，位置嵌入矩阵形状为 [T_max, C]。自注意力本身不区分 Token 的先后顺序，因此必须额外加入位置信息。
auto token_embeddings = token_embedding_->OpValue(tokens);
auto pos_ids = GetPositionIds(batch_size, seq_len, option);
auto position_embeddings = position_embedding_->OpValue(pos_ids);
auto x = dropout_emb_.OpValue(
token_embeddings + position_embeddings);
嵌入权重初始化为标准差 0.02 的正态分布。较小的初始尺度能避免与输出层权重共享时 Logits 过大，从而使 Softmax 过早饱和。
Attention
对输入 \(X\in\mathbb{R}^{B\times T\times C}\)，首先做三次线性投影：
$$ Q=XW^Q+b^Q, \qquad K=XW^K+b^K, \qquad V=XW^V+b^V $$
为了让投影转化为一次大矩阵乘法，Zero 先将 [B, T, C] 展平为 [B*T, C]，再使用融合的 OpLinearBias。得到 Q、K、V 后，将它们重塑为多头布局：
[B*T, C]
→ View [B, T, H, D]
→ Transpose [B, H, T, D]
其中 \(C=H\times D\)，\(H\) 是注意力头数，\(D\) 是单头维度。每个头的缩放点积注意力为：
$$ \operatorname{Attention}(Q,K,V) = \operatorname{softmax}\left( \frac{QK^T}{\sqrt{D}}+M \right)V $$
\(M\) 是因果掩码。对于位置 \(i\) 和 \(j\)：
$$ M_{ij}= \begin{cases} 0, & j\le i \\ -10^4, & j>i \end{cases} $$
对角线上方的大负数经 Softmax 后接近零，使第 \(i\) 个位置只能使用自己和之前的上下文：
x0 x1 x2 x3
x0 ✓ × × ×
x1 ✓ ✓ × ×
x2 ✓ ✓ ✓ ×
x3 ✓ ✓ ✓ ✓
有掩码时，Zero 使用 OpFlashAttention 将 Q @ K^T、缩放掩码 Softmax 和 probs @ V 组织成一个自动微分算子。它不在计算图中长期保留 [B, H, T, T] 的 Scores 和 Probabilities，而在反向传播时重新计算，用计算换取显存。
Transformer Block
Zero 使用 Pre-LN 结构，两个子层都在输入端做 LayerNorm，再通过残差连接加回主干：
$$ \begin{aligned} H’ &= H+\operatorname{Dropout}( \operatorname{Attention}(\operatorname{LN}_1(H))) \\ H’’ &= H’+\operatorname{Dropout}( \operatorname{FFN}(\operatorname{LN}_2(H’))) \end{aligned} $$
残差连接为梯度提供直接路径。LayerNorm 放在子层之前，可以在较深的网络中改善训练稳定性。
FFN 对每个 Token 位置独立执行两次线性变换，中间维度是隐藏维度的四倍：
$$ \operatorname{FFN}(x) = \operatorname{GeLU}(xW_1+b_1)W_2+b_2 $$
OpFFN 将两个线性层、Bias 和 GeLU 放在同一个算子内，减少中间张量和自动微分节点。Dropout 分别用在嵌入、注意力残差分支和 FFN 残差分支上。
输出层
所有 Transformer Block 之后，先应用最终 LayerNorm，再将隐藏向量投影到词表：
$$ \operatorname{logits}=H_{\mathrm{final}}E_{\mathrm{token}}^T $$
Zero 直接复用 Token Embedding 的权重，不为输出层另外创建 [V, C] 参数。这种权重共享不仅减少参数量，也使输入和输出使用同一个 Token 表示空间。
普通 OpValue 返回 [B, T, V] Logits，适合推理。训练使用 OpValueLoss，将输出投影与交叉熵融合：
auto x_2d = x.View({B * T, C});
return RunOp<OpLinearCrossEntropy>(
{x_2d, *token_weight, target_flat});
词表较大时，[B*T, V] Logits 会占用大量显存。融合算子不让这个张量进入长期保留的自动微分图，可以显著降低训练峰值显存。
训练
训练阶段需要先将原始文本转换为整数序列，再确定模型规模和更新策略。数据批次经过模型得到每个位置的损失，反向传播后还要完成梯度裁剪、学习率调整和参数更新。
数据准备
原始文本首先由 Tokenizer 转换成整数 ID。训练程序不需要知道一个 ID 对应字符还是子词，只要求训练集与验证集使用同一份词表，并能把 ID 还原成 Token。
整数 ID 按 int32 二进制格式写入 train.bin 和 val.bin。训练时从整个 ID 序列中随机选择 \(B\) 个起点，每个起点连续取 \(T+1\) 个 ID，再错位生成输入和目标。随机起点让不同 Batch 覆盖语料中的不同位置，连续片段则保留了语言模型需要学习的上下文关系。
训练配置
训练程序使用一组规模较小的 GPT 配置：
| Hyperparameter | Value |
|---|---|
| Sequence Length | 256 |
| Batch Size | 64 |
| Hidden Size | 384 |
| Transformer Blocks | 6 |
| Attention Heads | 6 |
| Dropout | 0.2 |
| AdamW \(\beta_1,\beta_2\) | 0.9, 0.99 |
| Weight Decay | 0.1 |
| Gradient Norm Threshold | 1.0 |
| Max / Min Learning Rate | \(10^{-3},10^{-4}\) |
| Warmup Steps | 100 |
训练循环
每一步的核心训练过程为：
opt->ZeroGrad();
auto per_token_loss = model->OpValueLoss(x, mask, y);
auto loss = per_token_loss.Mean(0, false);
loss.Backward();
float grad_norm = ClipGradNorm(all_params, 1.0f);
opt->SetLr(lr_for_iter(iter));
opt->Step();
GPT 的计算图较深，个别 Batch 可能产生异常大的梯度。ClipGradNorm 先计算所有参数梯度共同组成的 L2 范数；当范数超过阈值 \(c\) 时，将所有梯度按同一比例缩小：
$$ g_p\leftarrow g_p\cdot\frac{c}{\lVert g\rVert_2+10^{-6}} $$
统一缩放不会改变整体梯度方向。这里将阈值设为 1.0，裁剪完成后再由 AdamW 读取梯度并更新参数。
学习率前 100 步线性增长，之后余弦衰减。每 500 步在验证集上估计损失，验证和采样都使用 EvalGuard 关闭 Dropout。
推理
生成过程
生成时取最近 \(T\) 个 Token 作为上下文，对最后一个位置的 Logits 应用温度缩放：
$$ p_i = \frac{\exp(z_i/\tau)}{\sum_j\exp(z_j/\tau)} $$
实现中会先减去最大 logit，避免指数溢出，再用 std::discrete_distribution 根据概率采样下一 Token。温度 \(\tau\) 越小，分布越尖锐；温度越大，生成结果越随机。
采样前使用 EvalGuard 关闭 Dropout，并暂时关闭所有参数的梯度跟踪，避免为每一个生成步骤构建计算图。训练 Batch 与单样本采样的张量形状差异较大，因此采样前后还会调用 TrimFreePool() 清理无法复用的 CUDA 缓存桶。
当前生成路径没有 KV Cache，每产生一个 Token 都会重新计算整个上下文窗口。这个实现足以验证模型的训练结果，但不适合高吞吐推理。
示例
训练 StoneGPT
仓库中的 prepare_stone.py 将 UTF-8 文本转换成字符级训练数据。把语料保存为根目录下的 FourBooks.txt，再执行：
python3 prepare_stone.py
脚本收集语料中出现的所有 Unicode 字符，排序后将每个字符映射为一个整数 ID。这种方式不需要额外的子词 Tokenizer，便于直接观察训练流程；代价是序列较长，也难以直接使用更大粒度的语义单元。
脚本生成 data/stone/vocab.txt，并按照九比一划分训练集和验证集：
data/stone/
├── vocab.txt
├── train.bin
└── val.bin
构建 Zero 后，可以先在 CPU 上运行少量步骤，确认数据读取、前向计算和反向传播能够连通：
./build/apps/stone_gpt
默认只训练 30 步。完整训练更适合使用 GPU，并通过命令行指定训练步数和数据目录：
./build/apps/stone_gpt cuda 5000 data/stone
程序每隔固定步数输出训练损失，并定期在验证集上估计损失和生成文本。训练损失反映模型对当前 Batch 的拟合程度，验证损失用于观察模型对未参与更新的数据是否同样改善，生成结果则提供更直观的检查。三者需要结合观察：训练损失下降并不必然意味着模型已经学会稳定生成，也可能只是对训练语料拟合得更好。
如果希望使用其他语料，只需生成相同格式的 vocab.txt、train.bin 和 val.bin，再将目录作为第三个参数传给 stone_gpt。模型结构和训练流程不需要改变。