CS336 任务一 问答题

一、理解 Unicode

题目

(a) chr(0) 返回的是什么 Unicode 字符？
(b) 这个字符的字符串表示（__repr__()）与它的打印表示有什么区别？
(c) 当这个字符出现在文本中时会发生什么？可以在 Python 解释器中尝试以下代码，看看是否符合你的预期：

>>> chr(0)
>>> print(chr(0))
>>> "this is a test" + chr(0) + "string"
>>> print("this is a test" + chr(0) + "string")

解答

(a) chr(0) 返回的是 NUL 字符（U+0000），即 ASCII 码表中的第一个字符，也叫空字符（null character）。
(b) 区别：

• __repr__() 显示转义形式：'\x00'，让你能"看到"这个不可见字符。

• print() 输出时，NUL 字符不可见（不产生任何可视输出），终端上什么都不显示。

>>> chr(0)        # repr → '\x00'
>>> print(chr(0)) # 打印 → 什么都看不到（空白）

(c) NUL 字符嵌入文本中时，Python 字符串本身不会被截断（Python 的 str 是长度计数的，不像 C 语言用 \0 做终止符）。但在某些终端或程序中，NUL 会导致显示截断或异常行为：

>>> "this is a test" + chr(0) + "string"
# repr → 'this is a test\x00string'  （完整显示，中间有 \x00）

>>> print("this is a test" + chr(0) + "string")
# 取决于终端：有些会显示 "this is a teststring"（NUL 不可见）
# 有些终端可能在 NUL 处截断，只显示 "this is a test"

总结：NUL 字符在 Python 字符串内部是合法的，不会截断字符串，但它在显示层面是不可见的，某些外部工具（C 库、终端、文件系统）可能会把它当作字符串终止符。

二、Unicode 编码

问题

(a) 为什么我们更倾向于在 UTF-8 编码的字节上训练分词器，而不是 UTF-16 或 UTF-32？比较不同输入字符串在这些编码下的输出可能会有帮助。

(b) 考虑下面这个（错误的）函数，它的目的是将 UTF-8 字节串解码为 Unicode 字符串。为什么这个函数是错误的？请提供一个会产生错误结果的输入字节串示例。

def decode_utf8_bytes_to_str_wrong(bytestring: bytes):
    return "".join([bytes([b]).decode("utf-8") for b in bytestring])

>>> decode_utf8_bytes_to_str_wrong("hello".encode("utf-8"))
'hello'

(c) 给出一个不能解码为任何 Unicode 字符的两字节序列。

解答

(a) UTF-8 对 ASCII 字符只用 1 个字节（英文文本零膨胀），而 UTF-16 固定至少 2 字节、UTF-32 固定 4 字节，这意味着对于常见文本 UTF-8 序列更紧凑，分词器的词表能覆盖更多有意义的模式。此外，UTF-8 的字节词表大小固定为 256，而 UTF-16 基本单元是 65536 种可能值，词表基数过大不利于 BPE 合并效率。

(b) 反例："你".encode("utf-8") → b'\xe4\xbd\xa0'（3 个字节）

>>> decode_utf8_bytes_to_str_wrong("你".encode("utf-8"))
# 报错：UnicodeDecodeError

原因：该函数逐字节单独解码，但 UTF-8 中多字节字符（如中文占 3 字节）的每个后续字节（0x80–0xBF）单独来看不是合法的 UTF-8 序列。这个函数错在把多字节编码的字节拆开来单独解码，破坏了 UTF-8 的多字节结构。
(c) 示例：b'\xc0\x80'
解释：0xC0 0x80 虽然形式上符合两字节 UTF-8 模式（110xxxxx 10xxxxxx），但它是 U+0000 的"过长编码"（overlong encoding），UTF-8 标准明确禁止过长编码，因此这不是合法的 UTF 字符序列。
另一个更直接的例子：b'\xfe\xff' — 这两个字节都不是任何合法 UTF-8 序列的起始字节（0xFE 和 0xFF 在 UTF-8 中永远不会出现）。

三、分词器实验

题目

(a) 从 TinyStories 和 OpenWebText 中各采样 10 篇文档。使用你之前训练好的 TinyStories 和 OpenWebText 分词器（词表大小分别为 10K 和 32K），将这些采样文档编码为整数 ID。每个分词器的压缩比（bytes/token）是多少？

(b) 如果用 TinyStories 分词器去分词你的 OpenWebText 样本，会发生什么？比较压缩比并/或定性描述现象。

(c) 估算你的分词器的吞吐量（例如，bytes/秒）。对 Pile 数据集（825GB 文本）进行分词需要多长时间？

(d) 使用你的 TinyStories 和 OpenWebText 分词器，将各自的训练集和开发集编码为整数 token ID 序列。我们后续会用这些数据来训练语言模型。我们建议将 token ID 序列化为 uint16 数据类型的 NumPy 数组。为什么 uint16 是合适的选择？

解答

(a) 典型结果：

• TinyStories 分词器在 TinyStories 数据上：压缩比约 3.5–4.0 bytes/token

• OpenWebText 分词器在 OpenWebText 数据上：压缩比约 4.0–4.5 bytes/token

词表越大，每个 token 能表示越多字节，压缩比越高。

(b) 用 TinyStories 分词器去分词 OpenWebText 时，压缩比会显著下降（可能降到 2.0–3.0 bytes/token）。因为 TinyStories 词表是在简单儿童故事上训练的，缺少 OpenWebText 中常见的复杂词汇、专有名词和技术术语的合并规则，导致很多词只能被拆成细粒度的字节或短 token。

(c) 纯 Python 实现的 BPE 吞吐量大约在 0.5–2 MB/s 左右。对 825GB 的 Pile 数据集进行分词大约需要：

$\frac{825 \times 1024 \text{ MB}}{1 \text{ MB/s}} \approx 844,800 \text{ 秒} \approx 9.8 \text{ 天}$

如果用 regex 包 + 多进程并行，可以提速到几十 MB/s，将时间压缩到数小时。

(d) uint16 合适的原因：uint16 能表示 [0, 65535] 范围内的整数，而我们的词表大小（10K 或 32K）都远小于 65536，所以 uint16 足以容纳所有 token ID。相比 int32/int64，uint16 每个 token 只占 2 字节，存储空间减半或更多，对于数百 GB 的数据集这个节省非常可观。

四、Transformer 语言模型资源核算

问题

(a) 考虑 GPT-2 XL，其配置如下：

• vocab_size: 50,257

• context_length: 1,024

• num_layers: 48

• d_model: 1,600

• num_heads: 25

• d_ff: 6,400

假设我们用这个配置构建模型，模型有多少可训练参数？假设每个参数用单精度浮点数（float32）表示，仅加载这个模型需要多少内存？

(b) 找出完成 GPT-2 XL 形状模型一次前向传播所需的矩阵乘法。这些矩阵乘法总共需要多少 FLOPs？假设输入序列有 context_length 个 token。

(c) 基于上述分析，模型的哪些部分消耗最多 FLOPs？

(d) 对 GPT-2 small（12 层，768 d_model，12 头）、GPT-2 medium（24 层，1024 d_model，16 头）和 GPT-2 large（36 层，1280 d_model，20 头）重复你的分析。随着模型规模增大，Transformer LM 的哪些部分在总 FLOPs 中占比增加或减少？

(e) 将 GPT-2 XL 的上下文长度增加到 16,384。一次前向传播的总 FLOPs 如何变化？各模型组件的 FLOPs 相对占比如何变化？

解答

(a) 参数量计算

GPT-2 XL 的参数组成：

ps：上述都是针对gpt-2的，现代大模型有一些改动

• Q,K,V 以及 LayerNorm 都没有偏置项

• FFN 的参数如果使用swiglu也是有变化。
  每层参数：7,684,800 + 2,561,600 + 20,488,000 + 6,400 = 30,740,800

48 层总计：48 × 30,740,800 = 1,475,558,400

加上 Embedding + Position + 最终 LN：

• 80,411,200 + 1,638,400 + 3,200 = 82,052,800

总参数（不共享 LM Head）：1,475,558,400 + 82,052,800 + 80,411,200 ≈ 1,638,022,400 ≈ 1.56B

如果 LM Head 与 Embedding 共享权重：≈ 1.56B（GPT-2 实际共享权重，约 1.5B 参数）

内存：1.56B × 4 bytes ≈ 6.2 GB

(b) 矩阵乘法及 FLOPs

设 T = 1024（序列长度），d = 1600，d_{ff} = 6400，L = 48 层。

矩阵乘法的 FLOPs 公式：形状为$(m \times k) \times (k \times n)$的矩阵乘约 2mkn FLOPs。

每层的矩阵乘法：

每层 FLOPs：

• 线性投影部分：$6Td^2 + 2Td^2 + 2 \times 2Td \cdot d_{ff} = 8Td^2 + 4Td \cdot d_{ff}$

• 注意力得分部分：$4T^2d$

代入数值（每层）：

• 线性：$8 \times 1024 \times 1600^2 + 4 \times 1024 \times 1600 \times 6400 = 8 \times 1024 \times 2.56M + 4 \times 1024 \times 10.24M = 20,971,520,000 + 41,943,040,000 = 62,914,560,000 ≈ 62.9 GFLOPs$

• 注意力：$4 \times 1024^2 \times 1600 = 6,710,886,400 ≈ 6.7 GFLOPs$

每层总计 ≈ 69.6 GFLOPs

48 层总计 ≈ 48 × 69.6 = 3,341 GFLOPs ≈ 3.34 TFLOPs

加上 LM Head$（T \times d \times V）：2 \times 1024 \times 1600 \times 50257 ≈ 164.5 GFLOPs$

前向传播总 FLOPs ≈ 3.5 TFLOPs

(c) 消耗最多 FLOPs 的部分

FFN 层消耗最多 FLOPs。在每层中，FFN 贡献了约$4Td \cdot d_{ff}（因为 d_{ff} = 4d，即 16Td^2）$，占每层线性部分的 $16Td^2 / (8Td^2 + 16Td^2) = 2/3$。注意力的$QK^T 和 \text{scores} \times V$在序列较短时占比较小。

(d) 不同规模模型对比

随着模型增大（d 增大），注意力得分计算$（4T^2d）$相对于线性投影$（\propto d^2）$的占比下降，因为线性部分按 $d^2$ 增长而注意力得分只按 d 线性增长。FFN 和注意力投影的相对比例保持稳定（约 2:1）。

(e) 上下文长度增加到 16,384

将 T 从 1,024 增加到 16,384（16 倍）：

• 线性投影部分：$\propto T$，增加 16 倍

• 注意力得分部分：$\propto T^2$，增加 256 倍

原来注意力得分占比约$\frac{4T^2d}{8Td^2 + 4Td_{ff} + 4T^2d}$。当 T = 16384、d = 1600 时：

• 注意力得分每层：$4 \times 16384^2 \times 1600 ≈ 1.72 TFLOPs$

• 线性每层：$(8 \times 1600^2 + 4 \times 1600 \times 6400) \times 16384 ≈ 1.03 TFLOPs$

注意力得分现在超过了线性部分，从原来约 10% 的占比跃升到约 63%。总 FLOPs 大幅增加，注意力得分成为主要瓶颈——这也是为什么长上下文模型需要 Flash Attention 等优化的原因。

五、调整学习率

问题

我们将看到，学习率是对训练影响最大的超参数之一。让我们在玩具示例中实际验证这一点。用三个不同的学习率值：1e1、1e2 和 1e3，运行上面的 SGD 示例，仅训练 10 个迭代步。每个学习率下 loss 会怎样？它是更快衰减、更慢衰减，还是发散（即在训练过程中不断增大）？

解答

• lr = 1e1 (10)：loss 衰减较快，比默认学习率收敛更迅速，因为步长更大能更快接近最优解。

• lr = 1e2 (100)：loss 开始震荡，可能先下降后上升，呈现不稳定的锯齿状行为，处于收敛与发散的临界点。

• lr = 1e3 (1000)：loss 直接发散，从第一步开始就不断增大（可能变成 NaN 或 Inf），因为过大的步长使参数直接跳过最优解并越飞越远。

总结：学习率过小收敛慢，适中时收敛快，过大则导致训练发散。存在一个临界阈值，超过后优化器无法稳定在最优解附近。

六、AdamW 训练的资源核算

问题

让我们计算运行 AdamW 需要多少内存和计算量。假设每个张量都使用 float32。

(a) 运行 AdamW 需要多少峰值内存？根据参数、激活值、梯度和优化器状态的内存用量来分解你的答案。用 batch_size 和模型超参数（vocab_size、context_length、num_layers、d_model、num_heads）表示你的答案。假设 d_ff = 4 × d_model。

为简化起见，计算激活值内存时只考虑以下组件：

• Transformer block

  • RMSNorm(s)

  • 多头自注意力子层：QKV 投影、Q^TK 矩阵乘法、softmax、值的加权求和、输出投影

  • 逐位置前馈网络：W_1 矩阵乘法、SiLU、W_2 矩阵乘法

• 最终 RMSNorm

• 输出 embedding

• logits 上的交叉熵

交付物：分别给出参数、激活值、梯度和优化器状态的代数表达式，以及总计。

(b) 将你的答案代入 GPT-2 XL 的配置，得到一个仅依赖 batch_size 的表达式。在 80GB 内存下你能用的最大 batch size 是多少？

交付物：形如 $a \cdot \text{batchsize} + b$的表达式，以及最大 batch size 的数值。

(c) 运行 AdamW 一步需要多少 FLOPs？

交付物：一个代数表达式，附简要说明。

(d) 模型 FLOPs 利用率（MFU）定义为观测到的吞吐量（tokens/秒）与硬件理论峰值 FLOP 吞吐量的比值。NVIDIA A100 GPU 的 float32 理论峰值为 19.5 teraFLOP/s。假设你能达到 50% MFU，在单张 A100 上以 batch size 为 1024 训练 GPT-2 XL 400K 步需要多长时间？根据 Kaplan et al. [2020] 和 Hoffmann et al. [2022]，假设反向传播的 FLOPs 是前向传播的两倍。

交付物：训练所需天数，附简要说明。

解答

(a) 内存分解

设 B = batchsize，T = context_length，d = d_model，L = num_layers，V = vocab_size，$d{ff} = 4d$。

参数（Parameters）：

• Token Embedding：Vd

• Position Embedding：Td（如果有）

• 每层：$QKV (3d^2) + Output (d^2) + FFN (2 \times d \times 4d = 8d^2) + RMSNorm (2 \times d) = 12d^2 + 2d$

• 最终 LN：d

• LM Head（假设不共享）：Vd

总参数：$P = 2Vd + Td + L(12d^2 + 2d) + d \approx L \cdot 12d^2 + 2Vd$

参数内存：4P bytes

梯度（Gradients）：与参数大小相同 = 4P bytes

优化器状态（Optimizer State）：AdamW 存储一阶动量 m 和二阶动量 v，各与参数同大小 = $2 \times 4P = 8P bytes$

激活值（Activations）：每层需保存中间结果用于反向传播：

每层激活 $≈ BT(11d + 2HT)$（元素数）

总激活：$L \cdot BT(11d + 2HT) + BTV（最后输出层 logits）$

激活内存：$4 \times [L \cdot BT(11d + 2HT) + BTV] bytes$

总峰值内存：

$\text{Memory} = 16P + 4 \cdot [L \cdot BT(11d + 2HT) + BTV]$

其中 16P = 参数(4) + 梯度(4) + 优化器(8)。

(b) 代入 GPT-2 XL

配置：V = 50257，T = 1024，L = 48，d = 1600，H = 25。

参数量： $P \approx 48 \times 12 \times 1600^2 + 2 \times 50257 \times 1600 ≈ 1,474,560,000 + 160,821,600 ≈ 1.56B$

静态内存（参数 + 梯度 + 优化器）： $16 \times 1.56 \times 10^9 = 24.9 \text{ GB}$（约 25 GB，固定不随 batch size 变化）

激活内存（每个样本）：

• 每层：$1024 \times (11 \times 1600 + 2 \times 25 \times 1024) = 1024 \times (17600 + 51200) = 1024 \times 68800 = 70,451,200$

• 48 层：$48 \times 70,451,200 = 3,381,657,600$

• Logits：$1024 \times 50257 = 51,463,168$

• 总计每样本：3,433,120,768 元素 × 4 bytes ≈ 13.7 GB/sample

等等，这对单个样本来说太大了。让我重新检查——这是包含了所有层所有中间激活的全部保存。实际上很多框架会做 activation checkpointing。

简化表达：

$\text{Memory (bytes)} \approx 13.7B \times \text{batchsize} + 25B$

即形如$a \cdot \text{batchsize} + b$：

$\text{Memory} \approx 13.7 \times \text{batchsize} + 25 \text{ (GB)}$

最大 batch size（80 GB 内存）：

$\text{batchsize} = \lfloor(80 - 25) / 13.7\rfloor = \lfloor 55 / 13.7 \rfloor = \lfloor 4.01 \rfloor = 4$

最大 batch size ≈ 4（不使用梯度累积和 activation checkpointing 的情况下）。

(c) AdamW 一步的 FLOPs

一步 AdamW 包含：前向传播 + 反向传播 + 优化器更新。

• 前向传播：$C_{fwd} \approx 2 \times P \times T \times B$（近似公式：每个参数每个 token 约 2 FLOPs）

  • 更精确：$\approx B \times T \times (24Ld^2 + 4LT d + 2Vd)$

• 反向传播：$\approx 2 \times C_{fwd}$

• 优化器更新：$\approx$ 几P（相对于前向/反向可忽略）

总计：

$C_{step} \approx 6 \times B \times T \times P$

（经典近似：$6NTB$，其中 N=参数量）

代入 $GPT-2 XL：C_{step} \approx 6 \times B \times 1024 \times 1.56 \times 10^9$

(d) 训练时间计算

配置：$B = 1024，400K 步，单张 A100，50$

每步 FLOPs： $C_{step} = 6 \times 1024 \times 1024 \times 1.56 \times 10^9 \approx 9.83 \times 10^{15} \text{ FLOPs}$

（这里第一个 1024 是 batch_size，第二个 1024 是 context_length）

400K 步总 FLOPs： $C_{total} = 400000 \times 9.83 \times 10^{15} = 3.93 \times 10^{21} \text{ FLOPs}$

A100 有效吞吐（50% MFU）： $19.5 \times 10^{12} \times 0.5 = 9.75 \times 10^{12} \text{ FLOP/s}$

所需时间：$t = \frac{3.93 \times 10^{21}}{9.75 \times 10^{12}} = 4.03 \times 10^{8} \text{ 秒} \approx 4667 \text{ 天} \approx 12.8 \text{ 年}$

训练需要约 4,667 天（约 12.8 年），这就是为什么实际训练 GPT-2 XL 规模的模型需要数百张 GPU 并行。如果用 256 张 A100，则约 18 天完成。