第 4 课：模块化重构 —— 项目结构分层与 Rc<RefCell> 架构
第 4 课：模块化重构 —— 项目结构分层与 Rc 架构
代码位置：src/tensor.rs（核心：Tensor 结构体与自动微分）、src/module.rs（Module trait）、src/main.rs（模块声明与演示编排） 参考实现：src/layers.rs（Module 的第一个实现者）
1. 本课要搞懂的问题
- 一个从零写的项目，文件该怎么切？每个文件各管什么？
Tensor里的Rc<RefCell<Vec<f32>>>是什么？为什么张量需要"共享可变数据"？- 为什么
parents要用Rc<Vec<Tensor>>而不是直接Vec<Tensor>？直接存会有什么灾难？ backward()里的 DFS 拓扑排序是怎么保证"梯度从 loss 流回叶子"的？Moduletrait 是干什么的？它和优化器、训练循环怎么配合？
2. 为什么要"重构"
经过前三课，tensor.rs 已经长成了一个功能完整的张量库：构造、形状、逐元素运算、广播、归约、softmax、matmul、自动微分……如果所有代码继续堆在一个文件里，很快就会失控：
- 职责不分：张量运算、自动微分、模型层、训练逻辑混在一起，改一处可能碰坏另一处
- 无法复用：后续每节课都要加新东西（层、损失、优化器、GPT 模型），没有清晰的模块边界就没法增量开发
- 依赖混乱：谁依赖谁说不清，改个接口要满项目搜索
所以这一课做两件事：
- 横向：把项目按职责拆成多个
.rs文件，形成清晰的依赖金字塔 - 纵向：深入理解
Tensor的Rc<RefCell>架构——它是一切上层建筑的地基，前三课里所有算子都建立在这套设计之上
3. src/ 各文件职责分工
main.rs 顶部用 mod 声明了全部模块：
// src/main.rs
mod attention;
mod autograd;
mod checkpoint;
mod cli;
mod config;
mod data;
mod gpu;
mod layers;
mod loss;
mod model;
mod module;
mod optim;
mod rng;
mod rope;
mod sample;
mod tensor;
mod tokenizer;
mod train;
每个文件的职责（与文件头注释一致）：
| 文件 | 职责 | 对应课程 |
|---|---|---|
tensor.rs |
张量数据结构 + 运算（前向 + 反向闭包） | 第 1-4 课 |
autograd.rs |
自动微分：backward 迭代式拓扑排序（从 tensor.rs 拆出） |
第 2 课 |
rope.rs |
RoPE 旋转位置编码（从 tensor.rs 拆出） | 第 19 课 |
attention.rs |
多头自注意力 + KV Cache（从 model.rs 拆出） | 第 9-10、18 课 |
gpu.rs |
GPU 计算后端（--features gpu）：WGSL 计算着色器 |
第 21 课 |
module.rs |
Module trait（参数收集的统一接口） |
第 5 课 |
layers.rs |
神经网络层：Linear / LayerNorm / Embedding + 激活函数 |
第 5、11、12 课 |
loss.rs |
损失函数：mse_loss / cross_entropy_loss |
第 6 课 |
optim.rs |
优化器：SGD / AdamW |
第 6、17 课 |
rng.rs |
xorshift64 伪随机数生成器（参数初始化、采样） | 第 5 课起 |
data.rs |
数据加载 DataLoader 与内置语料 CORPUS |
第 14 课 |
tokenizer.rs |
分词器：CharTokenizer / BPETokenizer |
第 8 课 |
model.rs |
GPT 模型（Transformer Block 堆叠） |
第 11-12 课 |
train.rs |
训练循环 + 学习率调度 LRScheduler |
第 13、20 课 |
sample.rs |
推理采样：temperature / top-k / top-p | 第 15 课 |
config.rs |
配置加载（serde 序列化） | 配套 |
checkpoint.rs |
checkpoint 保存/恢复 | 配套 |
cli.rs |
命令行定义（clap） | 配套 |
main.rs |
程序入口，依次演示各课成果 | 每课 |
依赖关系呈金字塔形——越往下越基础，越往上越具体：
| 层次 | 模块 | 说明 |
|---|---|---|
| 第 4 层（入口） | main.rs |
声明全部模块，编排演示流程 |
| 第 3 层（流程） | train.rs model.rs sample.rs checkpoint.rs cli.rs |
训练循环、GPT 组装、推理采样、checkpoint、命令行 |
| 第 2 层（组件） | layers.rs loss.rs optim.rs data.rs tokenizer.rs rng.rs config.rs |
层、损失、优化器、数据、分词、随机数、配置 |
| 第 1 层（协议） | module.rs |
只定义一个 trait，不依赖上层任何模块 |
| 第 0.5 层（拆分） | autograd.rs rope.rs attention.rs |
从 tensor.rs / model.rs 拆出的独立模块 |
| 第 0 层（地基） | tensor.rs gpu.rs |
张量与自动微分、GPU 计算后端，被所有模块依赖 |
要点：
tensor.rs在最底层：所有模块都use crate::tensor::Tensor，但它不依赖任何其他模块。module.rs只依赖tensor.rs：它只定义一个 trait，是个"协议层"，供layers.rs/model.rs实现。- 上层模块互相独立：
loss.rs不知道optim.rs的存在，改动一个不影响另一个；data.rs、tokenizer.rs、rng.rs与神经网络计算完全解耦。
4. 核心架构：Rc<RefCell<Vec>> 共享可变数据
先看 Tensor 的结构体（src/tensor.rs）：
#[derive(Clone)]
pub struct Tensor {
data: Rc<RefCell<Vec<f32>>>, // 数值（前向结果）
shape: Vec<usize>, // 形状
grad: Rc<RefCell<Vec<f32>>>, // 梯度（反向累加）
requires_grad: bool, // 是否需要梯度（参数=true，数据=false）
parents: Rc<Vec<Tensor>>, // 父节点列表（我是由谁算出来的）
backward: Option<BackwardFn>, // 反向函数：把梯度传给父节点
}
两个问题：为什么数值和梯度要共享？共享之后怎么才能"可变"？
4.1 Rc：多个句柄指向同一份数据
计算图里一个节点经常被多个运算同时引用。比如 z = x·y + w 里 y 只被用一次，但看下面这种图：
x ──┐
mul ──► t ──┐
x ──┘ add ──► z （x 被用了两次！）
w ───────────────┘
x 同时是两次运算的输入。如果 Tensor 直接 pub data: Vec<f32>，克隆一个 Tensor 就要复制整份数据；而 Rc<T>（引用计数）允许无数个 Tensor 句柄共享同一个底层数据，克隆只把引用计数 +1，是 O(1) 的操作。这也是 #[derive(Clone)] 能放心用的前提。
4.2 RefCell：共享的同时允许修改
Rc 只解决"共享"，不解决"可变"——Rc<T> 的内容默认不可变。但我们的张量既被多处引用、又需要更新（前向写 data、反向写 grad）。Rust 的常规借用规则（一个值同时只能有一个可变借用）会拒绝这种写法。
RefCell<T>（内部可变性）把"借用检查"从编译期推迟到运行期：borrow() 只读借用、borrow_mut() 可变借用，同一时刻只能有一个 borrow_mut，违反就 panic。这就是"共享可变数据"的标准解法：
| 组合 | 解决的问题 | 代价 |
|---|---|---|
Rc<T> |
多个句柄共享一份数据 | 只读共享 |
RefCell<T> |
共享的数据还能改 | 借用检查推迟到运行时 |
Rc<RefCell<T>> |
共享 + 可变 | 运行时 panic 风险、无线程安全 |
对应到代码里的经典模式：
let sd = self.data.borrow(); // 读：前向计算
// …… 用完立刻释放，避免借用重叠
drop(sd);
let mut g = self.grad.borrow_mut(); // 写：反向累加梯度
注意代码里几乎每个算子的前向都有 drop(sd); drop(od);——因为之后构造结果节点要 borrow_mut 自己的东西，如果 data 的只读借用还挂着，同一 RefCell 上"只读借用未释放就再取可变借用"会在运行时 panic。这是 RefCell 使用者的纪律。
4.3 为什么 data 和 grad 分开两个 Rc
细看：data 和 grad 是两个独立的 Rc<RefCell<Vec<f32>>>。除了"语义上一个是数值、一个是梯度"之外，还有两个实际好处：
- 可以同时读 data、写 grad：反向闭包里常见
let sd_b = sd.borrow(); let mut sgm = sg.borrow_mut();——两者不是同一个RefCell，互不冲突。 - data 的指针充当计算图节点 ID：
backward里用Rc::as_ptr(&t.data) as usize唯一标识一个节点（见第 6 节），而grad的指针用来判断"两个父节点是不是同一个张量"（Rc::ptr_eq(&sg, &og)，处理x*x自乘的梯度合并）。
4.4 为什么不是 Arc
Rc 是单线程引用计数（快，无锁），Arc<Mutex> 是多线程方案。我们的训练是单线程的，用 Rc + RefCell 更轻、更符合"从零实现"的教学目标。等将来要并行训练，再换成 Arc<Mutex> 或 Arc<RwLock> 不迟。
5. parents: Rc<Vec> —— 为什么不能直接用 Vec
这是全项目最重要的一个设计决策。看结构体注释（源码原文）：
/// 父节点列表。
/// 注意：用 `Rc<Vec<_>>` 而不是 `Vec<Tensor>`——
/// 若直接存 Vec，`derive(Clone)` 会递归深拷贝整棵祖先计算图，
/// 深层图上每次建节点都是 O(图深) 的灾难。用 Rc 共享后克隆是 O(1)。
parents: Rc<Vec<Tensor>>,
5.1 深拷贝灾难是怎么发生的
每个运算在构造结果时都要记录父节点：
result.parents = Rc::new(vec![self.clone(), other.clone()]); // 每个算子都这么干
问题出在 self.clone()。如果 parents 的类型是 Vec<Tensor>，那么 clone 一个 Tensor 时：
data：Rc共享 → O(1)shape：Vec<usize>→ O(维度数)grad：Rc共享 → O(1)parents：Vec<Tensor>→ 递归克隆所有父节点！
而父节点自己又有 parents……于是 clone 会递归复制整棵祖先子树。设计算图深度为 L：
| 方案 | 建一个新节点的开销 | 一次前向（L 个节点） | N 步训练 |
|---|---|---|---|
parents: Vec<Tensor> |
O(L)（递归深拷贝） | O(L²) | O(N·L²) |
parents: Rc<Vec<Tensor>> |
O(1)（Rc 引用计数 +1） | O(L) | O(N·L) |
线性回归的计算图只有几层，L 很小，第 2 课那样写还能跑；但到了 GPT（几十层 Transformer），每步前向要建几百上千个节点，O(L²) 的复制会直接拖垮训练——这就是"性能灾难"。
5.2 第 2 课的旧代码 vs 现在的写法
第 2 课文档里的版本（能跑，但有问题）：
// ❌ 旧版：parents 直接存 Vec<Tensor>
pub struct Tensor {
data: Rc<RefCell<Vec<f32>>>,
shape: Vec<usize>,
grad: Rc<RefCell<Vec<f32>>>,
requires_grad: bool,
parents: Vec<Tensor>, // clone 时递归复制祖先图！
backward: Option<BackwardFn>,
}
重构后的版本：
// ✅ 现在的版本：parents 用 Rc 共享
pub struct Tensor {
// ……
parents: Rc<Vec<Tensor>>, // clone 时只复制一个 Rc 句柄
backward: Option<BackwardFn>,
}
对比要点：
- 旧版：
parents: Vec<Tensor>里存的是"真实拥有的子图副本"，克隆 = 深拷贝，图越深越慢。 - 新版：
parents: Rc<Vec<Tensor>>里存的是"共享的子图句柄"，克隆 = 引用计数 +1，O(1)。所有节点的parents共享同一棵祖先图，图在内存里只有一份。
这也解释了为什么 backward 里能放心地 for p in t.parents.iter() 递归——遍历的是共享图，不会触发任何复制。
6. backward 的 DFS 拓扑排序
6.1 为什么需要拓扑排序
反向传播要求梯度从输出端向输入端逐层传递：z = (x·y) + w 里，必须先执行 z 的 backward（把梯度分给 t、w），t 的梯度就绪后才能执行 t 的 backward（再分给 x、y）。执行顺序由计算图的依赖关系决定，而 DFS 后序 + 逆序执行恰好给出这个顺序。
6.2 代码：迭代式 DFS + HashSet 去重
注：早期版本使用递归 DFS，后改为迭代式栈 DFS（三色标记法）， 因为深层计算图（4 层 Transformer + batch=16）的递归深度会超过 Windows 默认 1MB 栈。 见
src/autograd.rs。
pub fn backward(&self) {
assert_eq!(self.rank(), 0, "backward() 只支持标量（0 维）输出，当前形状 {:?}", self.shape);
{
let mut g = self.grad.borrow_mut();
g[0] = 1.0; // ① loss 的梯度 = 1（∂loss/∂loss = 1）
}
// 迭代式 DFS 拓扑排序（避免递归栈溢出）
// 用显式栈模拟递归：每个元素 (node, child_index) 表示"该节点的第 child_index 个子节点待访问"
let mut order: Vec<Tensor> = Vec::new();
let mut visited: HashSet<usize> = HashSet::new();
let mut stack: Vec<(Tensor, usize)> = Vec::new();
let key = Rc::as_ptr(&self.grad) as usize; // 用 grad 的指针作节点唯一 ID
if visited.insert(key) {
stack.push((self.clone(), 0));
}
// ② 迭代式后序 DFS：先展开所有子节点，再入队
while let Some((node, idx)) = stack.last_mut() {
if *idx < node.parents.len() {
let child = node.parents[*idx].clone();
*idx += 1;
let child_key = Rc::as_ptr(&child.grad) as usize;
if visited.insert(child_key) {
stack.push((child, 0)); // ③ 先递归父节点
}
} else {
let (node, _) = stack.pop().unwrap();
order.push(node); // ④ 后序入队：子在前、父在后
}
}
for t in order.iter().rev() { // ⑤ 逆序遍历：从根走向叶
if let Some(b) = &t.backward {
b(); // 执行各节点的反向闭包
}
}
}
三个关键设计：
| 设计 | 原因 |
|---|---|
Rc::as_ptr(&t.grad) 做节点 ID |
Tensor 没有自增 id；用梯度 grad 的堆地址当身份证最省事，同一张量的所有句柄地址相同 |
HashSet 去重 |
一个节点被多个运算共享（如 x 被用两次），DFS 可能从两条路径走到它。梯度只需也只需执行一次它的 backward，否则会重复累加 |
| 后序 DFS + 逆序执行 | 后序 DFS 得到"子在前、父在后"（叶 → 根）；逆序执行变成"根 → 叶"，正是梯度从输出流向输入的次序 |
6.3 走一遍例子：z = (x·y) + w
x ──┐
mul ──► t ──┐
y ──┘ add ──► z
w ───────────────┘
DFS 后序结果（子在前、父在后）：order = [x, y, t, w, z]，逆序执行（z → w → t → y → x）：
| 顺序 | 节点 | 执行的 backward | 效果 |
|---|---|---|---|
| 1 | z | add 的闭包（收到的 g = 1） | t += 1，w += 1 |
| 2 | w | （叶子，无 backward） | — |
| 3 | t = x·y | mul 的闭包（收到的 g = 1） | x += 1·y = 3，y += 1·x = 2 |
| 4 | y | （叶子，无 backward） | — |
| 5 | x | （叶子，无 backward） | — |
和 test_chain_rule 验证的结果一致：∂z/∂x = y = 3、∂z/∂y = x = 2、∂z/∂w = 1 ✓
7. module.rs：Module trait 与 zero_grad_all
7.1 trait 定义
深度学习里一切"可训练的结构"（Linear、LayerNorm、Transformer Block、整个 GPT……）都需要回答一个问题：你有多少参数？ module.rs 就把这个回答统一成接口：
// src/module.rs（全文）
use crate::tensor::Tensor;
/// 模块接口：任何可训练结构都实现它
pub trait Module {
/// 返回模块的所有参数（含嵌套子模块）
fn parameters(&self) -> Vec<Tensor>;
}
/// 便捷方法：清零所有参数的梯度
pub fn zero_grad_all(module: &dyn Module) {
for p in module.parameters() {
p.zero_grad();
}
}
Module 目前只要求实现一个方法 parameters()。别小看这一个方法——它是优化器、训练循环、梯度清零与模型结构之间的唯一桥梁。
7.2 三个实现者：叶子模块
layers.rs 里三个最简单的实现：
// Linear：线性层 y = xW + b
impl Module for Linear {
fn parameters(&self) -> Vec<Tensor> {
vec![self.weight.clone(), self.bias.clone()]
}
}
// LayerNorm：层归一化
impl Module for LayerNorm {
fn parameters(&self) -> Vec<Tensor> {
vec![self.gamma.clone(), self.beta.clone()]
}
}
// Embedding：词嵌入查表
impl Module for Embedding {
fn parameters(&self) -> Vec<Tensor> {
vec![self.table.clone()]
}
}
注意 clone() 在这里是安全的：因为 Tensor 内部全是 Rc（第 4、5 节），克隆只复制句柄，不复制数据——返回 Vec<Tensor> 给外部持有也不会有性能问题。
7.3 组合模式：大模块递归收集小模块的参数
model.rs 里的 GPT 由若干 TransformerBlock 组成，TransformerBlock 又由 LayerNorm、MultiHeadAttention、Linear 组成。每个组合模块实现 Module 时，把子模块的参数拼接起来：
// model.rs：GPT
impl Module for GPT {
fn parameters(&self) -> Vec<Tensor> {
let mut ps = self.tok_emb.parameters(); // 词嵌入
for block in &self.blocks {
ps.extend(block.parameters()); // 每层 Transformer Block
}
ps.extend(self.ln_f.parameters()); // 最终 LayerNorm
ps.extend(self.lm_head.parameters()); // 输出头
ps
}
}
因为 parameters() 是 trait 方法，block.parameters() 会自动递归调用 TransformerBlock 的实现，后者又调用 MultiHeadAttention 的实现……最终一棵参数树被拍平成一张参数表。
GPT.parameters()
├── tok_emb.parameters() → [table]
├── block[0].parameters() → ln1[γ,β] + attn[c_q,c_k,c_v,c_proj] + ln2[γ,β] + mlp[W1,b1,W2,b2]
├── block[1].parameters() → …
└── ln_f.parameters() → [γ, β]
（lm_head 通过权重绑定复用 tok_emb.table.T，不单独存储）
7.4 与优化器的配合（main.rs 的 demo_xor）
main.rs 里第 7 课的 XOR 演示展示了完整闭环：
// ① 收集所有参数，一次性交给优化器
let params: Vec<Tensor> = {
let mut ps = fc1.parameters();
ps.extend(fc2.parameters());
ps
};
let opt = SGD::new(0.5, params);
// ② 训练循环：前向 → 反向 → 更新 → 清梯度
for step in 0..1000 {
let h = relu(&fc1.forward(&x_data));
let logits = fc2.forward(&h);
let loss = cross_entropy_loss(&logits, &y_targets);
loss.backward(); // 自动算所有参数的梯度
opt.step(); // SGD 遍历 params 更新：θ -= lr·g
opt.zero_grad(); // 清零，防累加
}
分工很清晰：
| 角色 | 谁 | 干什么 |
|---|---|---|
| 参数清单 | Module::parameters() |
告诉外界"我的参数在哪" |
| 梯度计算 | Tensor::backward() |
自动微分填满每个参数的 grad |
| 参数更新 | optim.rs 的 SGD / AdamW |
按梯度更新 data |
| 梯度清零 | Tensor::zero_grad() / zero_grad_all() |
重置 grad 为 0 |
zero_grad_all 就是"把清梯度也统一成接口"：传入任意 &dyn Module，它调用 parameters() 拿到所有参数逐个清零。实际训练中优化器（如上面 SGD::zero_grad）自带同样的清零逻辑，而 zero_grad_all 的价值在于：不经过优化器也能直接清任意模块的梯度（比如想"只求梯度、手动更新"时）。任何新模块只要实现 Module，就能复用这套参数管理设施。
8. 运行与测试
cargo test # 全部测试通过
cargo run # 依次演示：XOR MLP → BPE 分词 → 小 GPT 训练与生成
至此，项目的"地基"（阶段一）全部完成。从下一课开始，所有新代码（层、损失、优化器、Transformer）都将站在 tensor.rs + module.rs 这两块基石之上。
9. 动手练习
- 给
Linear之外再写一个你自己的模块（比如BiasOnly { bias: Tensor }），实现Moduletrait，并用zero_grad_all清一次梯度，确认能编译通过。 - 思考：
parameters()返回Vec<Tensor>时我们敢clone，靠的是什么？如果data不是Rc而是裸Vec<f32>，clone一次要花多少钱？ - 画一画
z = (x·y) + w的计算图，标出 DFS 后序的order数组内容，再写出逆序执行的顺序（参考答案在第 6.3 节）。 - 挑战题：把
parents改成Vec<Tensor>，跑cargo test，确认功能仍然正确；再构造一个 100 层深的计算图，对比两种写法每次前向/反向的耗时，体会"O(1) vs O(L)"的差别。 - 阅读
model.rs里TransformerBlock的impl Module，数一数一个 Block 最终贡献多少个参数张量。
10. 本课总结
- 项目按职责拆成 12 个模块，依赖呈金字塔形，
tensor.rs是最底层地基 Rc<RefCell<Vec<f32>>>= 共享（Rc）+ 可变（RefCell），是单线程下"共享可变数据"的标准方案parents: Rc<Vec<Tensor>>让克隆从 O(图深) 降到 O(1)，避免深拷贝灾难backward用Rc::as_ptr做节点 ID、DFS 后序 + 逆序执行，保证梯度按拓扑序从 loss 流回叶子Moduletrait 统一"收集参数"，配合zero_grad_all、优化器、训练循环形成完整闭环- 下一课：在 Module 之上实现第一个神经网络层 Linear 与激活函数，搭建第一个真正的网络