CS336 Notes

[toc]

Lecture 2 笔记：从张量到训练循环

1. 分词技术 (Tokenization)

1.1 分词的背景与意义

自然语言模型的输入是**整数序列 (tokens)**，因此必须先将原始文本转为 tokens。
Tokenizer 的职责是：

• **编码 (encode)**：string → tokens
• **解码 (decode)**：tokens → string

分词器的选择直接影响模型的效率、表达能力和泛化能力。
在课程中，tokenization 被称为一种 必要的“恶”：目前主流模型必须使用分词，但理想情况下我们希望模型能直接处理原始字节序列。

---

1.2 分词策略对比

策略	词汇量大小	压缩率	OOV 处理	优点	缺点
字符分词 (Character)	巨大 (~15 万 Unicode 字符)	低 (~1.5)	不存在 OOV	能表示任何文本，实现简单	词汇表过大，序列过长，效率低
字节分词 (Byte)	小且固定 (256)	极低 (=1)	不存在 OOV	词表小，覆盖面全	序列最长，注意力计算极不友好
词分词 (Word)	巨大且不固定	高	需 UNK token	序列短，符合语言直觉	词表爆炸，无法优雅处理新词
BPE (Byte Pair Encoding)	可控 (可自定义)	中等	不存在 OOV	在词汇量与序列长度间平衡	贪心算法，非全局最优，训练复杂

结论

• 字符/字节分词：覆盖面广，但序列太长。
• 词分词：序列短，但 OOV 问题严重。
• BPE：折中方案，成为主流（GPT-2/3/4 使用）。

---

1.3 分词策略详解

1.3.1 字符分词 (Character Tokenizer)

• 思想：每个 Unicode 字符 → code point (整数)。

• 示例：

  assert ord("a") == 97
  assert ord("🌍") == 127757
  assert chr(97) == "a"
  assert chr(127757) == "🌍"

• 实现：

  class CharacterTokenizer(Tokenizer):
      def encode(self, string: str) -> list[int]:
          return list(map(ord, string))
      def decode(self, indices: list[int]) -> str:
          return "".join(map(chr, indices))

• 问题：

  • vocab 过大 (~150K)
  • 序列过长，稀有字符浪费容量

---

1.3.2 字节分词 (Byte Tokenizer)

• 思想：UTF-8 编码 → 字节序列 (0-255)。

• 示例：

  bytes("a", "utf-8")  == b"a"
  bytes("🌍", "utf-8") == b"\xf0\x9f\x8c\x8d"

• 实现：

  class ByteTokenizer(Tokenizer):
      def encode(self, string: str) -> list[int]:
          return list(map(int, string.encode("utf-8")))
      def decode(self, indices: list[int]) -> str:
          return bytes(indices).decode("utf-8")

• 优点：

  • 词表固定 256
  • 无 OOV

• 缺点：

  • 压缩率始终为 1，序列过长
  • 对 Transformer 的注意力计算 (O(n²)) 极为低效

---

1.3.3 词分词 (Word Tokenizer)

• 思想：用正则表达式拆分单词。

• 示例：

  import regex
  string = "I'll say supercalifragilisticexpialidocious!"
  segments = regex.findall(r"\w+|.", string)

• 优点：

  • 序列短，语义自然

• 缺点：

  • vocab 爆炸，新词需 UNK
  • OOV 导致困惑度计算困难

---

1.3.4 BPE (Byte Pair Encoding)

• 历史：

  • Gage (1994)：用于数据压缩
  • Sennrich (2015)：引入机器翻译
  • Radford (2019)：GPT-2 采用 BPE

• 核心思想：

  • 从字节序列开始
  • 统计相邻 token 对频率
  • 将最高频对合并为新 token
  • 迭代，直到词表大小达到预设值

• 训练算法 (train_bpe)

  def train_bpe(string: str, num_merges: int) -> BPETokenizerParams:
      indices = list(map(int, string.encode("utf-8")))
      merges, vocab = {}, {x: bytes([x]) for x in range(256)}
      for i in range(num_merges):
          counts = defaultdict(int)
          for a, b in zip(indices, indices[1:]):
              counts[(a, b)] += 1
          pair = max(counts, key=counts.get)
          new_index = 256 + i
          merges[pair] = new_index
          vocab[new_index] = vocab[pair[0]] + vocab[pair[1]]
          indices = merge(indices, pair, new_index)
      return BPETokenizerParams(vocab, merges)

• 编码 (encode)

  • 将字符串转为字节
  • 按 merges 规则合并

• 解码 (decode)

  • 查 vocab → 拼接字节 → UTF-8 解码

• 优势：

  • 高频片段 → 单 token，提高压缩率
  • 低频片段 → 多 token，保证泛化
  • 无 OOV

---

1.4 压缩率 (Compression Ratio)

定义：

$\text{compression ratio} = \frac{\text{num_bytes}}{\text{num_tokens}}$

• 越高说明 token 越“信息密集”。
• 比较：
  • 字符分词：~1.5
  • 字节分词：=1
  • 词分词：高，但 OOV 问题严重
  • BPE：折中，通常 1.5–3 左右

---

1.5 总结

• Tokenizer = encode + decode
• 字符/字节分词：覆盖全面，但序列过长，效率低
• 词分词：序列短，但 OOV 问题严重
• BPE：折中方案，兼顾 vocab 控制与序列长度，成为主流
• 趋势：
  • 探索 tokenizer-free 模型（直接从字节输入），但目前未能在大规模上取代 BPE

Lecture 2 笔记：从张量到训练循环

1 引言：从资源到模型——构建大模型的基石

• 本讲的目标：解析训练一个现代深度学习模型所需的所有核心技术原语。
• 方法：自下而上 (bottom-up)，从 张量 (Tensors) → 模型 (Models) → 优化器 (Optimizers) → **训练循环 (Training Loop)**。
• 核心理念：效率 (Efficiency)
  • **内存 (Memory)**：以 GB 为单位，决定模型与 batch 的大小。
  • **计算 (Compute)**：以 FLOPs 为单位，决定训练时间与成本。
• 知识框架：
  • Mechanics：PyTorch 原语，直观实现。
  • Mindset：资源核算，效率优先。
  • Intuitions：理解大体趋势与代价，而非细节调优。

---

2 动机：餐巾纸数学 (Napkin Math)

通过粗略估算快速感受大模型训练的规模：

问题一：70B 参数模型，15T tokens，1024 × H100 需要多久？

• 总计算量：

  Total FLOPs=6×N×D\text{Total FLOPs} = 6 \times N \times D

  • $N = 70 \times 10^9$ (参数)
  • $D = 15 \times 10^{12}$ (tokens)
  • 系数 6 来源：前向传播 $2N$ + 反向传播 $4N$ = $6N$ FLOPs。
  • 结果：$6.3 \times 10^{24}$ FLOPs。

• 每日计算量：

  • 单卡 H100 理论峰值：$1979 \times 10^{12}$/s
  • 假设 MFU (利用率) = 0.5
  • 1024 卡每天：约 $4.35 \times 10^{22}$ FLOPs

• 所需天数：

  Days=6.3×10244.35×1022≈145\text{Days} = \frac{6.3 \times 10^{24}}{4.35 \times 10^{22}} \approx 145

结论：训练 GPT-3 规模模型即便在千卡 H100 上也需 ~5 个月。

---

问题二：8 × H100，AdamW，能训练的最大模型多大？

• 单卡显存：80 GB

• 朴素 float32 AdamW 内存开销：

  • 参数：4 B
  • 梯度：4 B
  • 优化器状态 (m, v)：8 B
  • 总：16 B/参数

• 8 卡总内存：$80 \times 10^9 \times 8 = 640$ GB

• 可训练参数量：

  Num Params=640×10916≈40×109\text{Num Params} = \frac{640 \times 10^9}{16} \approx 40 \times 10^9

结论：约 40B 参数模型（未计激活内存）。

---

3 张量

3.1 张量基础 (Tensors Basics)

张量是深度学习中存储一切的基本单元：参数、梯度、优化器状态、数据、激活。

• 创建张量：

  x = torch.tensor([[1., 2, 3], [4, 5, 6]])
  x = torch.zeros(4, 8)
  x = torch.ones(4, 8)
  x = torch.randn(4, 8)
  x = torch.empty(4, 8)  # 未初始化
  nn.init.trunc_normal_(x, mean=0, std=1, a=-2, b=2)

用途：先分配内存，后用自定义逻辑初始化。

---

3.2 内存核算 (Memory Accounting)

几乎所有内容都存为浮点数。关键：值的数量 × 每个值的字节数。

3.2.1 float32 (单精度)

• 结构：1 符号位 + 8 指数位 + 23 尾数位
• 大范围，高精度，默认类型
• 内存示例：
  • torch.zeros(4, 8) → $4 \times 8 \times 4$ B = 128 B
  • GPT-3 前馈层矩阵 $12288 \times 12288$ → 2.3 GB

3.2.2 float16 (半精度)

• 结构：1 符号位 + 5 指数位 + 10 尾数位
• 内存减半 (2B/值)
• 缺陷：动态范围窄，小数易下溢 (如 $1e-8 \to 0$)
• 风险：训练不稳定

3.2.3 bfloat16 (脑浮点)

• Google Brain 2018 提出
• 结构：1 符号位 + 8 指数位 + 7 尾数位
• 动态范围 = float32，避免下溢
• 精度低于 float16，但足够用于深度学习
• 常用于计算，而参数/优化器仍存 float32

3.2.4 FP8 (8位浮点)

• NVIDIA H100 支持 (2022)
• 两种：E4M3、E5M2
• 内存更小，但需配合 Transformer Engine 软件栈
• 依赖自动缩放/转换策略

---

3.3 浮点格式对比

特性	float32	float16	bfloat16	FP8
总位数	32	16	16	8
指数位	8	5	8	4/5
尾数位	23	10	7	2/3
动态范围	极大	有限	与 fp32 相同	较窄
精度	高	中	低	极低
内存/计算开销	大	中	中	极小
训练稳定性	稳定	容易下溢	稳定	需谨慎

要点：

• float32：稳定可靠，内存开销大
• float16：节省内存，但风险高
• bfloat16：折中，常用
• FP8：极致效率，需硬件/软件支持

---

3.4 现代解决方案：混合精度训练 (Mixed Precision Training)

• 思路：结合高精度与低精度优势
• 参数/优化器：float32
• 前向/部分计算：bfloat16 / fp16
• 部分关键模块（如注意力）：保留 float32
• 工具：NVIDIA Transformer Engine 提供自动缩放和类型转换

结论：混合精度是当前主流，兼顾速度、内存与稳定性。

---

4 张量存储与视图 (Tensor Storage & Views)

• PyTorch 张量本质：指针 + 元数据

  • Shape：形状
  • Stride：跨距 (定位逻辑索引到内存)
  • Storage Offset：起始偏移量

• 示例：4×4 矩阵

  x = torch.tensor([
      [0., 1, 2, 3],
      [4, 5, 6, 7],
      [8, 9, 10, 11],
      [12, 13, 14, 15],
  ])
  assert x.stride(0) == 4  # 行方向跨距
  assert x.stride(1) == 1  # 列方向跨距

• 定位元素公式：

结论：张量是存储的“视图”，而不是独立的数据块。

5 零成本操作：视图 vs 拷贝

• **视图 (View)**：只改变元数据，不复制底层存储
  • x[0, :]、x[:, 1]
  • .view()、.transpose()
  • 修改原张量会同步影响视图
• **拷贝 (Copy)**：创建新存储，占用额外内存
  • 非连续张量必须 .contiguous() 后才能 .view()

要点：性能优化中尽量使用视图，避免无意识拷贝。

6 张量操作分类

6.1 逐元素操作 (Element-wise)

• 对每个元素独立运算

• 示例：

  x = torch.tensor([1, 4, 9])
  x.pow(2)  # 平方
  x.sqrt()  # 开方
  x + x
  x * 2

• 复杂示例：因果注意力 mask

  x = torch.ones(3, 3).triu()

• 复杂度：$O(mn)$

6.2 矩阵乘法 (Matrix Multiplication)

• 核心运算：“深度学习的面包黄油”

• 示例：

  x = torch.ones(16, 32)
  w = torch.ones(32, 2)
  y = x @ w

• 批处理：输入 $(B, S, D_{in})$，权重 $(D_{in}, D_{out})$

  • 输出 $(B, S, D_{out})$

• 复杂度：$O(BDK)$，FLOPs = $2BDK$

6.3 高级张量操作：Einops 库

• 动机：PyTorch 默认索引可读性差 (-1, -2 容易出错)

• 核心函数：

  • einsum：广义矩阵乘法

    z = einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")

  • reduce：聚合降维

    y = reduce(x, "batch seq hidden -> batch seq", "mean")

  • rearrange：拆分与合并维度

    x = rearrange(x, "b s (h d) -> b s h d", h=num_heads)

结论：einops 提升可读性与健壮性，适合复杂模型（如多头注意力）。

7 计算核算 (Compute Accounting)

7.1 FLOPs 与 FLOP/s

• FLOPs：总运算量 (e.g. GPT-3 训练 $~3.1 \times 10^{23}$)
• FLOP/s：硬件性能 (e.g. H100 $~ 10^{15}$/s)
• 注意区分：小写 s (总量) vs 大写 S (速率)

7.2 矩阵乘法 FLOPs

• 公式：$2 \times m \times n \times p$
• 前向传播 FLOPs ≈ $2 \times$ (tokens × 参数)
• 反向传播 FLOPs ≈ $4 \times$ (tokens × 参数)
• 总 FLOPs = 前向 + 反向 = $6$ × (tokens × 参数)

7.3 硬件利用率 (MFU)

• 定义：
• 指标：>0.5 已属优秀
• 瓶颈可能来源：I/O、通信、内存访问、非优化代码
• 提升 MFU = 更低的成本 + 更快的训练

8 梯度与反向传播 (Gradients & Backprop)

8.1 自动微分 (Autograd)

• 步骤：
  1. requires_grad=True
  2. 前向传播 → loss
  3. 调用 .backward()
  4. 梯度存于 .grad

8.2 梯度 FLOPs

• 两层线性模型：
  • $h1 = x @ w1$，$h2 = h1 @ w2$
  • loss = $h2^2$.mean()
• 前向 FLOPs：$2BD(D+K)$
• 反向 FLOPs：$4BD(D+K)$
• 总 FLOPs = 前向 + 反向 = $6B \times$ 参数量
• 结论：反向传播的计算量 ≈ 前向的两倍

小结

• 张量是指针 + 元数据，视图操作零成本，拷贝需谨慎。
• 逐元素运算 $O(n)$，矩阵乘法 $O(n^3)$，后者主导计算成本。
• Einops 提供声明式接口，简化复杂张量操作。
• FLOPs 核算：前向 2N，反向 4N，总 6N。
• 硬件性能通过 MFU 衡量，>0.5 已属良好。
• 关键认知：效率优化的根本在于理解并量化内存与计算代价。

9 nn.Parameter 与初始化

• nn.Parameter：torch.Tensor 的子类，被赋值到 nn.Module 属性时会自动注册到模型参数列表。优化器会自动追踪这些参数。

• 初始化问题：若输入维度过大，随机初始化可能导致输出方差随输入维度增长，引发梯度爆炸。

• 解决方案：Xavier 初始化（Glorot 初始化）。通过按输入维度平方根缩放，保证输出方差稳定。

  self.weight = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))

• 改进：可进一步使用截断正态分布，避免极端值。

9.1 自定义模型：Linear 与 Cruncher

• Linear 层：简单的线性变换，继承自 nn.Module。

  class Linear(nn.Module):
      def __init__(self, input_dim, output_dim):
          super().__init__()
          self.weight = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))
      def forward(self, x):
          return x @ self.weight

• Cruncher 模型：由多层 Linear 堆叠构成的深度线性网络。

  class Cruncher(nn.Module):
      def __init__(self, dim, num_layers):
          super().__init__()
          self.layers = nn.ModuleList([Linear(dim, dim) for _ in range(num_layers)])
          self.final = Linear(dim, 1)
      def forward(self, x):
          for layer in self.layers:
              x = layer(x)
          return self.final(x).squeeze(-1)

• 要点：使用 nn.ModuleList 而不是 Python list，确保子模块参数被追踪到。

9.2 随机性与可复现性

• 随机源：参数初始化、数据顺序、Dropout 等。

• 最佳实践：固定所有随机种子，保证实验可复现性。

  import torch, numpy as np, random
  seed = 0
  torch.manual_seed(seed)
  np.random.seed(seed)
  random.seed(seed)

9.3 数据加载与管道

• 数据表示：语言模型的数据通常是 tokenizer 输出的整数序列，可序列化为 numpy 数组。
• 大规模数据集：TB 级语料无法一次性加载，使用 np.memmap 进行懒加载。
• GPU 数据传输优化：
  • 使用 pin_memory() 固定内存。
  • 在 .to(device, non_blocking=True) 时实现异步数据传输。

9.5 优化器原理与实现

• SGD：参数更新公式 p -= lr * grad。

• Momentum：在梯度更新中引入指数加权平均，抑制震荡。

• AdaGrad：为每个参数维护历史梯度平方和，自适应调整学习率。

  p -= lr * grad / torch.sqrt(g2 + 1e-5)

• RMSProp：对梯度平方使用指数加权平均。

• Adam：结合 Momentum 与 RMSProp，成为主流优化器。

• AdamW：解耦权重衰减，提升泛化性能。

9.6 内存核算（Parameters, Activations, Gradients, Optimizer States）

• 参数：$D^2 \times L + D$

• 激活值：$B \times D \times L$

• 梯度：与参数量相同。

• 优化器状态：以 AdaGrad 为例，需要额外存储平方梯度副本。

• 总内存：

  （假设 float32，每个值 4 字节）

9.7 训练循环与最佳实践

• 训练循环核心步骤：
  1. 获取 batch 数据
  2. 前向传播计算 loss
  3. 反向传播计算梯度
  4. 优化器更新参数
  5. 清空梯度，进入下一轮
• Checkpointing：定期保存模型与优化器状态，防止长时间训练崩溃后损失全部进度。
• **混合精度训练 (AMP)**：
  • 参数与优化器状态：float32
  • 前向传播与激活：bfloat16 / fp16
  • 损失缩放：解决梯度下溢问题
  • 工具：torch.cuda.amp.autocast + GradScaler

10 总结

• 模型是张量与运算的组合，nn.Module 提供了结构化封装。
• 优化器从最简单的 SGD 演化到 AdamW，反映了对学习率调控与正则化理解的深化。
• 内存与计算核算帮助我们量化训练需求。
• 最终通过 训练循环 将所有要素集成，构成现代深度学习系统的基本执行框架。