CS336 LecTure 2 笔记补充

核心思维：Mechanics vs. Mindset

• Mechanics（机械层面）：如何用 PyTorch 的 API 搭建张量、模块、优化器和训练循环。
• Mindset（思维层面）：每写一段代码，都要有显存与 FLOPs 的概念，知道这段代码“贵不贵”，瓶颈在哪里。

Memory Accounting（显存核算）

显存是大模型训练的“硬通货”。你需要清楚显存被谁吃了

显存四大消耗

1. Parameters（模型权重）：存储模型参数。
2. Gradients（梯度）：反向传播时存储，大小与参数相同。
3. Optimizer States（优化器状态）：AdamW 需要为每个参数存 2 个状态（Momentum 和 Variance）。
4. Activations（激活值）：前向传播的中间结果，用于反向传播。这玩意儿很费显存，随 Batch Size 和 Sequence Length 线性/平方增长。

精度与字节数

• FP32：4 Bytes。通常用于主权重和优化器状态，防止累积误差。（保精度）
• BF16：2 Bytes。用于前向/反向计算。动态范围大，不易溢出，是 H100/A100 的首选。（保速度）
• FP16：2 Bytes。容易 Underflow，需配合 Loss Scaling。
• FP8：1 Byte。前沿探索，需特殊硬件支持。

经验公式：AdamW 的显存开销

如果你使用标准的 AdamW 优化器，每个参数大约需要 16 Bytes：

• 权重 (FP32): 4B
• 梯度 (FP32): 4B
• Momentum (FP32): 4B
• Variance (FP32): 4B
• 总计: 16 Bytes/param

---

Compute Accounting（算力核算）

FLOPs vs. FLOP/s

• FLOPs (Floating Point Operations)：浮点运算总量，衡量任务有多重。
• FLOP/s (FLOPS)：每秒浮点运算次数，衡量硬件有多快。

黄金法则：6NP 公式

训练一个 Transformer 的总计算量估算：

\[\text{Total FLOPs} \approx 6 \cdot N \cdot D \]

• \(N\)：模型参数量
• \(D\)：Token 数量

系数 6 怎么来的？

1. 前向传播：矩阵乘法 \(Y=XW\)，每个元素 1次乘 + 1次加 = 2 FLOPs。
2. 反向传播：计算 \(\frac{\partial L}{\partial X}\) 和 \(\frac{\partial L}{\partial W}\)，计算量是前向的 2倍 = 4 FLOPs。
3. 总计：\(2 + 4 = 6\) FLOPs per parameter per token。

训练时长估算

\[\text{Days} = \frac{6 \cdot N \cdot D}{\text{GPU Count} \cdot \text{Peak FLOP/s} \cdot \text{MFU} \cdot 86400} \]

• MFU (Model FLOPs Utilization)：模型浮点利用率。理论峰值很难达到，MFU \(\ge\) 0.5 就算优秀。
• 案例：70B 模型，15T Tokens，1024张 H100 (MFU=0.5) \(\approx\) 144 天。

---

PyTorch 工程细节

Tensor 内存陷阱

• View vs. Reshape：
  • view()：要求内存连续，不拷贝数据，零开销。
  • reshape()：如果不连续，会隐式调用 contiguous() 触发内存拷贝！
• Slicing：x[:, 1:] 通常返回 View，不拷贝；但 x[:, 1:].contiguous() 会拷贝。
• 非连续张量：transpose 或 permute 后，张量变为非连续，后续操作可能变慢或报错。

Einops

使用 einops 库，少写这种 x.permute(0, 2, 1, 3).contiguous().view(...) ：

# 传统写法
y = x.permute(0, 2, 1, 3).contiguous().view(b, t, h * d)

# Einops 写法 (清晰安全的同时可以自动处理 contiguous)
from einops import rearrange
y = rearrange(x, 'b t h d -> b t (h d)')

混合精度训练 (AMP)

不要全程 FP32，也不要全程 BF16（保精度的同时保速度）：

• 前向/反向：用 BF16（快、省显存）。
• 优化器状态/主权重：用 FP32（稳）。
• Loss：用 FP32（防止梯度下溢）。
• PyTorch 的 torch.cuda.amp.autocast(dtype=torch.bfloat16) 会自动处理大部分转换。

---

训练循环与检查点

标准 Loop

for step, batch in enumerate(dataloader):
    # 1. 前向 (BF16)
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        loss = model(batch)
    
    # 2. 反向 (BF16 梯度，FP32 主权重更新)
    scaler.scale(loss).backward()
    
    # 3. 更新
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # 必须梯度裁剪！
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

Checkpointing

大模型训练崩，须定期保存

1. model.state_dict()
2. optimizer.state_dict()
3. scheduler.state_dict()
4. random_states \((Python, NumPy, Torch, CUDA)\)
5. step 和 epoch

---

总结

1. 显存是硬约束：算不清显存，代码跑不起来。
2. 算力是时间成本：6NP 公式是“预算表”。
3. 细节决定成败：一个 .contiguous() 可能导致 20% 的性能损失。