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 启动、内存池、归约和算子融合。