怎么读这一卷

主线位置：三条主线的总览。本页先讲每个对照怎么跑、两边的数字按什么标准算「对得上」，再把张量、训练、部署三条主线的每一步对到讲它的页。

版本：PyTorch 2.15.0（CPU）、burn 0.22.1（Flex 设备）

flowchart TB
  py["train_step.py"] --> pj["pt.json"]
  rs["train-step crate"] --> bj["burn.json"]
  pj --> cmp["compare.py"]
  bj --> cmp
  cmp -->|"差值不超过容差"| ok["OK 退出码 0"]
  cmp -->|超出| bad["MISMATCH 退出码 1"]

锁定版本

这一卷的 burn 实验一般开 std、flex、autodiff、optim 这四个 feature，个别章再加 ndarray、linalg、metal、store、dataset，每页写明。设备一般是 Device::flex()（纯 Rust 的 CPU 后端，见 Flex：CPU 后端），要求导时是 Device::flex().autodiff()。

default-features = false 有一个副作用要先说：burn 的默认 feature 会沿依赖链打开 Flex 的 simd 和 rayon，关掉之后 Flex 退回标量循环、单线程（关了默认 feature，GPU 变慢）。数值对照照样成立，但求和的结果会变，下面会实测。为什么要转比二进制大小和启动时间时，另外给出了带默认 feature 的构建。

burn 源码只认 tag v0.22.1，链接一律是 https://github.com/tracel-ai/burn/blob/v0.22.1/...。PyTorch 一侧只写在虚拟环境里实际跑出来的行为和它的 docstring，不链接 PyTorch 的源码。

平台：M2 Max、macOS 27、rustc 1.99.0。

每章的结构

每个内容页的顺序固定：

1. 「主线位置」：这一页是哪条主线的哪一步。
2. 版本行和一张图。
3. 正文按对照点分节。每一节先给 PyTorch 代码，再给 burn 代码，然后是两边的真实输出。先说两边概念上相同的部分，再说哪里不同、为什么不同，burn 内部的机制链到 burn 源码解析。
4. 说「对得上」的地方，给出最大绝对误差和容差，不是 0 时说明原因。
5. ## 小结。

一对对照程序

全卷的主例子是 A1：两层 MLP（4→8→3，ReLU），交叉熵，SGD 学习率 0.1，训练一步。权重不用随机数，按 PyTorch Linear.weight 的 [out, in] 下标 k 生成 0.5·sin(k)（第一层）和 0.5·cos(k)（第二层），所以两边拿到的是同一组数。PyTorch 一侧：

def w(rows, cols, f, scale):
  return torch.tensor([[scale * f(r * cols + c + 1) for c in range(cols)] for r in range(rows)])

with torch.no_grad():
  # PyTorch 的 Linear.weight 形状是 [out, in]
  model.fc1.weight.copy_(w(HID, IN, math.sin, 0.5))
  model.fc1.bias.copy_(torch.tensor([0.01 * o for o in range(HID)]))

burn 的 Linear 权重是 [in, out]，所以按同样的下标生成，再转置一次：

/// 按 PyTorch 的 [out, in] 下标生成数值，再转置成 burn 的 [in, out]
fn weight(out: usize, inp: usize, f: fn(f64) -> f64, device: &Device) -> Tensor<2> {
  let v: Vec<f32> =
    (0..out * inp).map(|k| (0.5 * f((k + 1) as f64)) as f32).collect();
  Tensor::<2>::from_data(TensorData::new(v, [out, inp]), device).transpose()
}

输入 x 和偏置也一样：两边都先在 f64 里算（0.1 * k - 0.5、0.01 * o），最后舍入成 f32 一次。在 Rust 里直接写 0.1 * k as f32 - 0.5 会在 f32 里算，16 个输入里有 7 个和 Python 差 1 到 3 个 ulp，比较的就不再只是计算误差了。

两个程序都把 logits、loss、四组梯度、更新后的参数和更新后的 loss 写成同样结构的 JSON（burn 侧输出时把权重转回 [out, in]）。完整代码和逐步讲解在训练一步。跑法：

python -I train_step.py > pt.json
cargo run --release --offline > burn.json
python -I compare.py pt.json burn.json

compare.py 递归走两份 JSON，对每个字段把所有数摊平，打印最大绝对误差，任何一个字段超出容差就退出码 1：

a, b = (json.load(open(p)) for p in sys.argv[1:3])
tol = float(sys.argv[3]) if len(sys.argv) > 3 else 1e-6
worst = 0.0
def walk(x, y, path):
  global worst
  if isinstance(x, dict):
    for k in x:
      walk(x[k], y[k], f"{path}.{k}" if path else k)
    return
  fx, fy = list(flat(x)), list(flat(y))
  assert len(fx) == len(fy), (path, len(fx), len(fy))
  d = max(abs(p - q) for p, q in zip(fx, fy))
  worst = max(worst, d)
  print(f"{path:24} n={len(fx):3}  max|diff|={d:.2e}")
walk(a, b, "")
print(f"worst={worst:.2e} tol={tol:.0e} {'OK' if worst <= tol else 'MISMATCH'}")
sys.exit(0 if worst <= tol else 1)

本页实测（burn 侧 release，Device::flex().autodiff()）：

logits                   n= 12  max|diff|=1.49e-08
loss                     n=  1  max|diff|=1.19e-07
grads.fc1.weight         n= 32  max|diff|=3.73e-09
grads.fc1.bias           n=  8  max|diff|=7.45e-09
grads.fc2.weight         n= 24  max|diff|=7.45e-09
grads.fc2.bias           n=  3  max|diff|=1.49e-08
params_after.fc1.weight  n= 32  max|diff|=0.00e+00
params_after.fc1.bias    n=  8  max|diff|=1.86e-09
params_after.fc2.weight  n= 24  max|diff|=1.49e-08
params_after.fc2.bias    n=  3  max|diff|=9.31e-10
loss_after               n=  1  max|diff|=1.19e-07
worst=1.19e-07 tol=1e-06 OK

退出码 0。一个从来不报错的检查等于没有检查，所以反过来验一次：把 PyTorch 输出里 params_after.fc2.bias 的第一个数加 0.001，再比：

params_after.fc2.bias    n=  3  max|diff|=1.00e-03
loss_after               n=  1  max|diff|=1.19e-07
worst=1.00e-03 tol=1e-06 MISMATCH

退出码 1。

「对得上」是什么意思

上面的误差不是 0，但都在 f32 舍入的量级上。拿 ulp（相邻两个 f32 的间距）来量：

按每个字段的最大 ulp 数逐个数过一遍（实验脚本把两边的数都转成 f32，比较位模式）：logits 最多 4 个 ulp，grads.fc2.weight 最多 48 个，其余字段 0 到 8 个。48 个 ulp 听起来多，绝对误差却只有 5.59e-9。这个梯度元素是 batch 里两项相加的结果：-0.0207 和 0.0217，两项几乎抵消，剩下 0.00105。舍入误差按 0.02 那个量级产生，结果却小了 20 倍，它的 ulp 只有 1.16e-10，同样的误差折成 ulp 就成了 48 个。所以这一卷用绝对误差判断，不用 ulp 数或相对误差：后两者会在接近 0 的元素上报出没有意义的大数。

两边做的是同样的数学运算，结果却不逐位相同，原因是浮点加法不满足结合律：同样一组数，换一个顺序相加，舍入就落在不同的地方。matmul 和求和的内部顺序由各自的实现决定：burn 的 Flex 把 matmul 交给 gemm crate（Flex：CPU 后端），PyTorch 用它自己的 CPU 内核。只要有一处求和顺序不同，最后一位就可能不同。

下面用一个只有求和顺序不同的实验来看这个效应。同样 100 万个 f32（sin(k) 先用 f64 算再转 f32），在几种方式下求和：

n = 1_000_000
v = torch.sin(torch.arange(n, dtype=torch.float64)).float()
print(f"torch sum  {v.sum().item():.9f}")
print(f"f64 exact  {v.double().sum().item():.9f}")

let n = 1_000_000;
let v: Vec<f32> = (0..n).map(|k| (k as f64).sin() as f32).collect();
let forward: f32 = v.iter().fold(0.0, |acc, x| acc + x);
let backward: f32 = v.iter().rev().fold(0.0, |acc, x| acc + x);
let exact: f64 = v.iter().map(|&x| x as f64).sum();
let t = Tensor::<1>::from_data(TensorData::new(v, [n]), &Device::flex());
let burn_sum: f32 = t.sum().into_scalar();

PyTorch：

torch sum  0.232940197
f64 exact  0.232891951

burn，这一卷的默认配置（default-features = false，Flex 没有 simd）：

burn sum   0.232917666
loop fwd   0.232917666
loop rev   0.232904017
f64 exact  0.232891951

同一段代码，依赖改成 burn = { version = "=0.22.1", features = ["flex"] }（默认 feature 打开，Flex 有 simd 和 rayon）：

burn sum   0.232849628
loop fwd   0.232917666
loop rev   0.232904017
f64 exact  0.232891951

同一组 100 万个数，四种求和方式得到了四个不同的 f32：

• 顺序累加和倒序累加差 1.4e-5，只是加的顺序反过来了。
• 关掉 simd 时，burn 的 sum 和顺序累加逐位相同：sum_f32_contiguous 在没有 simd feature 时走 data.iter().copied().sum()。
• 打开 simd 时走 sum_body：8 个 SIMD 累加器各加一部分，最后再合起来，顺序变了，结果也变了。同一个库、同一份数据，只因为 Cargo feature 不同，结果差了 6.8e-5。
• PyTorch 的结果和哪一种都不同，离 f64 参考值 4.8e-5。

这些结果没有哪个是错的，每一个都是「f32 求和」的合法结果。所以这一卷判断「对得上」，用的是容差，不是逐位相等：

• A1 这种小规模的计算，容差 1e-6。 A1 里的数都小于 2，每个求和不超过 8 项，实测最大误差 1.19e-7，是 1 附近的一个 ulp。1e-6 留了大约 8 倍的余量；如果哪里真的算错了（比如上面改了 0.001），远远超出。
• 规模大的计算，容差要放大，并写明原因。 100 万项的求和，两个正确的实现之间就能差 5e-5。哪一页用了别的容差，那一页会给出依据。
• 整数、形状、dtype 要求完全相同。 这些不涉及舍入，差了就是行为不同，比如整数除法的取整方向（整数运算）。

三条主线落在哪一页

张量主线：建一个张量 → 计算 → 把结果读回主机。

训练主线：forward → loss → backward → 优化器更新一步。

部署主线：保存 → 加载 → 不带自动求导地推理。

阅读路线

• 按顺序读：01 到 05 是一条完整的训练主线，读完能把一个 PyTorch 训练脚本改写成 burn。
• 只关心部署：为什么要转 → 不求导的区域 → 08 存取与迁移 → 09 部署。
• 已经在写 burn、被 PyTorch 习惯绊住：直接查 10 常见坑，每条都从 PyTorch 的写法出发，给出 burn 的报错原文。

小结

• 版本锁定 PyTorch 2.15.0（CPU）和 burn 0.22.1；burn 侧 default-features = false，Flex 因此没有 simd 和 rayon。
• 每个对照是一个 .py 加一个 cargo crate，输出同样结构的 JSON，compare.py 逐字段报最大绝对误差，超出容差退出码 1，并用一个改动 0.001 的输入验过它会报错。
• A1 训练一步两边最大绝对误差 1.19e-7（1 附近的一个 ulp）；不是 0，是因为浮点求和的顺序不同。接近 0 的元素折成 ulp 可以差到 48 个，所以判断用绝对误差。
• 同样 100 万个 f32，顺序累加、倒序累加、burn 有 simd、PyTorch 四种求和得到四个不同的结果，最大相差约 9e-5；burn 无 simd 时和顺序累加逐位相同。
• 容差：A1 这种规模用 1e-6；大规模计算另定并写明原因；整数、形状、dtype 要求完全相同。