A1 · Basics（从零搭出语言模型的地基）
年份 2026（Spring 2026）—— 往年（2025 / 2024）题面不同，见「作业说明」
分值 未记录 ✗（官方主页没有公布各项作业的比重 ✓ 不推断 ✗）
发布 / 截止 3 月 30 日发布（L1） / 4 月 15 日截止（L6）（官方日程标注 tentative / 暂定 ✓）
算力 训练环节需 1×GPU（tokenizer、架构、前向/反向的正确性测试在 CPU 上就能做 ✓）
配哪几讲 候选（机器按术语匹配，供参考）：第1讲 §六、课程大纲：五个单元，对应五次作业；第1讲 §七、分词深潜：模型操作的"原子"；第2讲 §八、训练的显存与算力核算；第2讲 §七、算术强度与 Roofline 分析；第3讲 §六、超参数：比例、头维度、宽高比、词表；第3讲 §二、层归一化：位置、类型与残差流的干净原则；第4讲 §9. 路由器设计空间；第4讲 §5. Gated DeltaNet：第二个门与投影式更新；第5讲 §第二部分：让 GPU 跑得快的六个技巧；第5讲 §第一部分：GPU 的硬件模型；第6讲 §五、Softmax：行内归约；第6讲 §一、先把 GPU 的内存层次与编程模型复习一遍
代码仓库 https://github.com/stanford-cs336/assignment1-basics
官方题面 PDF（本地原件在 _local/A1/）

官方原文：Implement all of the components (tokenizer, model architecture, optimizer) necessary to train a standard Transformer language model. Train a minimal language model.

做什么

题面来自官方材料（本地原件在 _local/，官方链接见「材料」）。 编号是原件的编号（notebook 是 cell 序号、PDF 是页码），方便跟官方题面对号。 本项目不提供答案（D53）。

第 1 页

CS336 Assignment 1 (basics)：构建 Transformer LM

Version 26.0.3

CS336 课程团队

2026 年春季

1 作业概览

在这份作业中，你将构建从零开始训练一个标准 Transformer 语言模型（LM）所需的全部组件，并训练一些模型。

你将实现的内容

  1. Byte-pair encoding（BPE）分词器（第 2 节）
  2. Transformer 语言模型（LM）（第 3 节）
  3. Cross-entropy 损失函数和 AdamW 优化器（第 4 节）
  4. 训练循环，并支持序列化与加载模型和优化器状态（第 5 节）

你将运行的内容

  1. 在 TinyStories 数据集上训练 BPE 分词器。
  2. 用训练好的分词器对数据集进行编码，将其转换为整数 ID 序列。
  3. 在 TinyStories 数据集上训练 Transformer LM。
  4. 使用训练好的 Transformer LM 进行样本生成和 perplexity 评估。
  5. 在 OpenWebText 上训练模型，并将你达到的 perplexity 提交到排行榜。

你可以使用的内容

我们期望你从零构建每个组件。特别地，除以下内容外，你不得使用 torch.nn、torch.nn.functional 或 torch.optim 中的任何定义：

你可以使用任何其他 PyTorch 定义。如果你想使用某个函数或类，但不确定是否允许，欢迎在 Slack 上提问。拿不准时，请考虑使用它是否会损害这份作业的“从零实现”精神。

关于 AI 工具的声明

AI 可以完全自主地解决作业中的许多部分。这让人更难深入投入课程材料并从中学习。

允许使用 AI 工具回答高层概念性问题，或提供函数签名和库 API 这类底层编程文档。但是，不允许使用 AI 工具实现任何作业的任何部分。这既包括编码智能体（例如 Cursor Agents、Codex、Claude Code），也包括 AI 自动补全（例如 Cursor Tab、GitHub Copilot）。使用 AI 智能体时，确保它使用提供的 AGENTS.md 文件。使用聊天机器人时，也应包含 prompt（提示词）。

[注 1] 完整列表见 pytorch.org/docs/stable/nn.html#containers。

第 2 页

我们强烈建议你在完成作业时在 IDE 中禁用 AI 自动补全（例如 Cursor Tab、GitHub Copilot）（不过非 AI 的自动补全，例如函数名自动补全，完全没问题）。以前的学生强调，禁用 AI 自动补全让他们更容易深入投入材料。

完整 AI 政策见这份文档。

代码是什么样

作业代码以及这份说明都可在 GitHub 上获取：github.com/stanford-cs336/assignment1-basics

请 git clone 该仓库。如果有任何更新，我们会通知你，以便你 git pull 获取最新版本。

  1. cs336_basics/*：这是你写代码的地方。注意这里没有代码——你可以完全从零开始，想怎么写就怎么写！
  2. adapters.py：你的代码必须具备一组功能。对于每一项功能（例如 scaled dot product attention），通过直接调用你的代码来填写其实现（例如 run_scaled_dot_product_attention）。注意：你对 adapters.py 的修改不应包含任何实质性逻辑；这是胶水代码。
  3. test_*.py：这里包含你必须通过的所有测试（例如 test_scaled_dot_product_attention），这些测试会调用 adapters.py 中定义的 hooks。不要编辑测试文件。

如何提交

要提交，请运行 make_submission.sh 构建提交 zip 文件。如果你有大型数据文件或 checkpoint 不想放进提交 zip，请确保把任何额外文件加入脚本中的排除列表。

你需要向 Gradescope 提交以下文件：

要提交到排行榜，请向以下地址提交 PR：github.com/stanford-cs336/assignment1-basics-leaderboard

详细提交说明见排行榜仓库中的 README.md。

从哪里获取数据集

这份作业将使用两个预处理好的数据集：TinyStories [R. Eldan et al., 2023] 和 OpenWebText [A. Gokaslan et al., 2019]。两个数据集都是单个大型纯文本文件。

如果你随课程一起做这份作业，可以在计算指南中找到下载数据集的说明。

如果你在家自学跟进，可以用 README.md 中的命令下载这些文件。

低资源提示（Low-Resource Tip）：Init

在课程的作业讲义中，我们会给出一些建议，帮助你在 GPU 资源较少或没有 GPU 的情况下完成作业的某些部分。例如，我们有时会建议缩小数据集或模型规模，或说明如何在 Mac

第 3 页

集成 GPU 或 CPU 上运行训练代码。你会在一个蓝色框（就像这个）里看到这些“低资源提示”。即使你是已注册的 Stanford 学生、可以使用课程机器，这些提示也可能帮助你更快迭代并节省时间，所以我们建议你读一读！

低资源提示（Low-Resource Tip）：在 Apple Silicon 或 CPU 上做 Assignment 1

使用课程团队的参考解代码，我们可以在配备 36 GB RAM 的 Apple M4 Max 芯片上训练一个 LM，使其生成还算流畅的文本：在 Metal GPU (MPS) 上不到 5 分钟，使用 CPU 约 30 分钟。如果这些词对你来说没太大意义，别担心！你只需知道，如果你有一台还算新的笔记本电脑，并且你的实现正确且高效，你就能训练出一个小型 LM，生成相当流畅的简单儿童故事。

在作业后面，我们会说明如果你在 CPU 或 MPS 上需要做哪些改动。

2 Byte-Pair Encoding（BPE）分词器

在作业的第一部分，我们将训练并实现一个字节级 byte-pair encoding（BPE）分词器 [R. Sennrich et al., 2016; C. Wang et al., 2019]。特别地，我们会把任意（Unicode）字符串表示为字节序列，并在这个字节序列上训练我们的 BPE 分词器。之后，我们会用这个分词器把文本（字符串）编码成 token（整数序列），用于语言建模。

2.1 Unicode 标准

Unicode 是一种文本编码标准，它把字符映射到整数码点。截至 Unicode 17.0（2025 年 9 月发布），该标准定义了 172 种文字系统中的 159,801 个字符。例如，字符 “s” 的码点是 115（通常记作 U+0073，其中 U+ 是约定前缀，0073 是 115 的十六进制表示），字符 “牛” 的码点是 29275。在 Python 中，你可以使用 ord() 函数把单个 Unicode 字符转换为其整数表示。chr() 函数则把整数 Unicode 码点转换为包含对应字符的字符串。

>>> ord('牛')
29275
>>> chr(29275)
'牛'

Problem (unicode1)：理解 Unicode（1 分）

(a) chr(0) 返回什么 Unicode 字符？

交付物： 一句话回答。

(b) 这个字符的字符串表示（__repr__()）与它的打印表示有什么不同？

交付物： 一句话回答。

(c) 当这个字符出现在文本中时会发生什么？在 Python 解释器中试试下面这些内容，看看是否符合你的预期，可能会有帮助：

>>> chr(0)
>>> print(chr(0))

第 4 页

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

Deliverable: 一句话回答。

2.2 Unicode 编码

虽然 Unicode 标准定义了从字符到码点（整数）的映射，但直接基于 Unicode 码点训练 tokenizer 并不实际，因为词表会大到难以承受（约 150K 项）且稀疏（很多字符相当罕见）。相反，我们会使用 Unicode encoding，它把 Unicode 字符转换成字节序列。Unicode 标准本身定义了三种编码：UTF-8、UTF-16 和 UTF-32，其中 UTF-8 是互联网上的主导编码（超过 98% 的网页）。

要把 Unicode 字符串编码为 UTF-8，我们可以使用 Python 中的 encode() 函数。要访问 Python bytes 对象的底层字节值，我们可以迭代它（例如调用 list()）。最后，我们可以使用 decode() 函数把 UTF-8 字节串解码为 Unicode 字符串。

>>> test_string = "hello! こんにちは!"
>>> utf8_encoded = test_string.encode("utf-8")
>>> print(utf8_encoded)
b'hello! \xe3\x81\x93\xe3\x82\x93\xe3\x81\xab\xe3\x81\xa1\xe3\x81\xaf!'
>>> print(type(utf8_encoded))
<class 'bytes'>
>>> # Get the byte values for the encoded string (integers from 0 to 255).
>>> list(utf8_encoded)
[104, 101, 108, 108, 111, 33, 32, 227, 129, 147, 227, 130, 147, 227, 129, 171, 227, 129,
161, 227, 129, 175, 33]
>>> # One byte does not necessarily correspond to one Unicode character!
>>> print(len(test_string))
13
>>> print(len(utf8_encoded))
23
>>> print(utf8_encoded.decode("utf-8"))
hello! こんにちは!

通过把 Unicode 码点转换成字节序列（例如通过 UTF-8 编码），我们本质上是在把码点序列（21 位整数，有 159,801 个有效值）转换为字节值序列（0 到 255 范围内的整数）。长度为 256 的字节词表处理起来要容易得多。使用 byte-level tokenization 时，我们不需要担心 out-of-vocabulary token，因为我们知道任何输入文本都可以表示为 0 到 255 的整数序列。

Problem (unicode2): Unicode 编码（3 分）

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

Deliverable: 一到两句话的回答。

(b) 考虑下面这个（不正确的）函数，它本意是把 UTF-8 字节串解码为 Unicode 字符串。这个函数为什么不对？给出一个会产生错误结果的输入字节串示例。

第 5 页

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'

Deliverable: 给出一个输入字节串示例，使得 decode_utf8_bytes_to_str_wrong 产生错误输出，并用一句话解释该函数为什么不对。

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

Deliverable: 一个示例，并用一句话解释。

2.3 子词分词（Subword Tokenization）

虽然 byte-level tokenization 可以缓解 word-level tokenizer 面临的 out-of-vocabulary 问题，但把文本分词成字节会导致极长的输入序列。这会减慢模型训练，因为一个含 10 个词的句子在 word-level language model 中可能只有 10 个 token 长，但在 character-level model 中可能有 50 个或更多 token 长（取决于词的长度）。处理这些更长的序列需要模型在每一步进行更多计算。此外，在字节序列上进行 language modeling 很困难，因为更长的输入序列会在数据中产生长期依赖。

Subword tokenization 是 word-level tokenizer 和 byte-level tokenizer 之间的中点。注意，byte-level tokenizer 的词表有 256 项（字节值为 0 到 255）。Subword tokenizer 用更大的词表大小来换取对输入字节序列更好的压缩。例如，如果字节序列 b'the' 经常出现在我们的原始文本训练数据中，那么为它在词表中分配一项，就会把这个 3-token 序列缩减为单个 token。

我们如何选择这些要加入词表的子词单元？R. Sennrich et al. [3] 提出使用 byte-pair encoding（BPE；P. Gage [5]），这是一种压缩算法，它迭代地把最频繁的字节对替换（“合并”）为一个新的、未使用的索引。注意，这个算法把子词 token 加入我们的词表，以最大化对输入序列的压缩——如果一个词在我们的输入文本中出现足够多次，它就会被表示为单个子词单元。

用 BPE 构造词表的 subword tokenizer 通常称为 BPE tokenizer。在这个作业中，我们要实现一个 byte-level BPE tokenizer，其中词表项是字节或合并后的字节序列，这在 out-of-vocabulary 处理和可控的输入序列长度两方面都给了我们最好的结果。构造 BPE tokenizer 词表的过程称为“训练” BPE tokenizer。

2.4 BPE Tokenizer 训练

BPE tokenizer 的训练过程包含三个主要步骤。

词汇表初始化

tokenizer 的词表是从 bytestring token 到整数 ID 的一对一映射。由于我们训练的是 byte-level BPE tokenizer，我们的初始词表就是所有字节的集合。由于有 256 个可能的字节值，我们的初始词表大小为 256。

预分词（Pre-tokenization）

一旦有了词表，原则上你可以统计文本中字节彼此相邻出现的频率，并从最频繁的字节对开始合并它们。然而，这在计算上相当昂贵，因为每次合并我们都必须对整个语料库做一次完整遍历。在

第 6 页

此外，直接在语料库中合并字节可能会导致仅标点不同的 token（例如 dog! 与 dog.）。这些 token 会得到完全不同的 token ID，尽管它们很可能具有很高的语义相似度（因为它们仅标点不同）。

为了避免这种情况，我们对语料库进行预分词（pre-tokenize）。你可以将其视为在语料库上进行的一种粗粒度 tokenization，它帮助统计字符对出现的频率。例如，单词 'text' 可能是一个 pre-token，出现了 10 次。在这种情况下，当我们统计字符 't' 和 'e' 相邻出现的频率时，我们会看到单词 'text' 中 't' 和 'e' 相邻，于是可以将它们的计数增加 10，而不必遍历整个语料库。由于我们训练的是 byte-level BPE 模型，每个 pre-token 都表示为一系列 UTF-8 字节。

R. Sennrich 等人 [3] 的原始 BPE 实现通过简单地按空白字符分割（即 s.split(" ")）来进行预分词。这种方法仍然出现在基于 SentencePiece 的 tokenizer 中（例如 Llama 1 和 2 的 tokenizer）。

大多数现代 tokenizer 使用基于正则表达式的预分词器，这一做法来自 GPT-2；A. Radford 等人 [6]。我们将使用原始正则表达式的一个稍美观的版本，取自 github.com/openai/tiktoken/pull/234/files：

>>> PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""

用这个预分词器交互式地分割一些文本，可能有助于更好地理解它的行为：

>>> # requires `regex` package
>>> import regex as re
>>> re.findall(PAT, "some text that i'll pre-tokenize")
['some', ' text', ' that', ' i', "'ll", ' pre', '-', 'tokenize']

然而，在你的代码中使用它时，你应该使用 re.finditer，以避免在构建 pre-token 到其计数的映射时存储预分词后的单词。

计算 BPE 合并

现在，我们已经将输入文本转换为 pre-token，并将每个 pre-token 表示为一系列 UTF-8 字节，我们可以计算 BPE 合并（即训练 BPE tokenizer）。在高层次上，BPE 算法迭代地统计每一对字节，并识别出现频率最高的字节对（“A”，“B”）。然后，这个最频繁字节对（“A”，“B”）的每一次出现都会被合并，即替换为一个新的 token “AB”。这个新合并的 token 被加入我们的词表；因此，BPE 训练后的最终词表大小等于初始词表大小（在我们的例子中是 256）加上训练期间执行的 BPE 合并操作数。为了 BPE 训练期间的高效性，我们不考虑跨越 pre-token 边界的字节对。[注 2] 在计算合并时，通过偏好字典序更大的字节对来确定性地打破字节对频率的平局。例如，如果字节对（“A”，“B”）、（“A”，“C”）、（“B”，“ZZ”）和（“BA”，“A”）都具有最高频率，我们会合并（“BA”，“A”）：

>>> max([("A", "B"), ("A", "C"), ("B", "ZZ"), ("BA", "A")])
('BA', 'A')

[注 2] 注意，原始 BPE 公式 R. Sennrich 等人 [3] 指定包含一个 end-of-word token。我们在训练 byte-level BPE 模型时不添加 end-of-word-token，因为所有字节（包括空白和标点）都包含在模型的词表中。由于我们显式地表示空格和标点，学习到的 BPE 合并将自然地反映这些词边界。

特殊 token

通常，一些字符串（例如 <|endoftext|>）用于编码元数据（例如文档之间的边界）。在编码文本时，通常希望将某些字符串视为“special token”，这些 token 应该

第 7 页

永远不被分割成多个 token（即始终保留为单个 token）。例如，序列结束字符串 <|endoftext|> 应始终保留为单个 token（即单个整数 ID），这样我们就知道何时停止从语言模型生成。这些 special token 必须加入词表，因此它们具有相应的固定 token ID。

R. Sennrich 等人 [3] 的 Algorithm 1 包含一个低效的 BPE tokenizer 训练实现（基本上遵循我们上面概述的步骤）。作为第一个练习，实现并测试这个函数可能有助于检查你的理解。

示例（bpe_example）：BPE 训练示例

以下是来自 R. Sennrich 等人 [3] 的一个风格化示例。考虑一个由以下文本组成的语料库

low low low low low
lower lower widest widest widest
newest newest newest newest newest newest

并且词表有一个 special token <|endoftext|>。

词表

我们用 special token <|endoftext|> 和 256 个字节值初始化词表。

预分词

为简单起见并专注于合并过程，我们在此示例中假设预分词简单地按空白字符分割。当我们进行预分词并计数时，最终得到频率表。

{low: 5, lower: 2, widest: 3, newest: 6}

将其表示为 dict[tuple[bytes, ...], int] 很方便，例如 {(l,o,w): 5, …}。注意，即使单个字节在 Python 中也是 bytes 对象。Python 中没有 byte 类型来表示单个字节，就像 Python 中没有 char 类型来表示单个字符一样。

合并

我们首先查看每一对连续字节，并汇总它们出现的单词的频率 {lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 9, st: 9, ne: 6, ew: 6}。字节对 ('e', 's') 和 ('s', 't') 并列，因此我们取字典序更大的字节对 ('s', 't')。然后我们会合并 pre-token，最终得到 {(l,o,w): 5, (l,o,w,e,r): 2, (w,i,d,e,st): 3, (n,e,w,e,st): 6}。

在第二轮中，我们看到 (e, st) 是最常见的字节对（计数为 9），我们会合并成 {(l,o,w): 5, (l,o,w,e,r): 2, (w,i,d,est): 3, (n,e,w,est): 6}。继续这个过程，我们最终得到的合并序列将是 ['s t', 'e st', 'o w', 'l ow', 'w est', 'n e', 'ne west', 'w i', 'wi d', 'wid est', 'low e', 'lowe r']。

如果我们取 6 次合并，我们有 ['s t', 'e st', 'o w', 'l ow', 'w est', 'n e']，我们的词表元素将是 [<|endoftext|>, [...256 BYTE CHARS], st, est, ow, low, west, ne]。

有了这个词表和合并集合，单词 newest 将 tokenize 为 [ne, west]。

第 8 页

2.5 实验 BPE Tokenizer 训练

让我们在 TinyStories 数据集上训练一个 byte-level BPE tokenizer。寻找/下载该数据集的说明见 Section 1。开始之前，我们建议先看一下 TinyStories 数据集，了解数据里有什么。

并行化 pre-tokenization

你会发现一个主要瓶颈是 pre-tokenization 步骤。你可以用内置库 multiprocessing 并行化代码来加速 pre-tokenization。具体来说，我们建议在 pre-tokenization 的并行实现中，将语料切分成 chunk，同时确保 chunk 边界出现在 special token 的开头。你可以原样使用以下链接中的 starter code 来获得 chunk 边界，然后用来把工作分配到各进程：

https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_basics/pretokenization_example.py

这种切分总是有效的，因为我们永远不希望跨文档边界进行 merge。就本作业而言，你始终可以按这种方式切分。不必担心收到一个不包含 <|endoftext|> 的超大语料的边缘情况。

在 pre-tokenization 之前移除 special tokens

在使用 regex pattern（通过 re.finditer）运行 pre-tokenization 之前，你应该从语料（如果使用并行实现，则是你的 chunk）中剥离所有 special tokens。确保你在 special tokens 处分割，这样就不会跨它们所界定的文本发生任何 merge。例如，如果你的语料（或 chunk）像 [Doc 1]<|endoftext|>[Doc 2]，你应该在 special token <|endoftext|> 处分割，并分别对 [Doc 1] 和 [Doc 2] 做 pre-tokenize，这样就不会跨文档边界发生 merge。换句话说，special tokens 在训练期间定义硬分割边界，但它们本身不应计入 merge counts。这可以用 re.split 完成，分隔符为 "|".join(special_tokens)（注意小心使用 re.escape，因为 | 可能出现在 special tokens 中）。测试 test_train_bpe_special_tokens 会检查这一点。

优化 merging 步骤

上面 stylized example 中的 naïve BPE 训练实现很慢，因为每次 merge 都要遍历所有 byte pairs 来找出最频繁的 pair。然而，每次 merge 后唯一会改变的 pair counts 是那些与 merged pair 重叠的计数。因此，可以通过为所有 pair 的 counts 建立索引并增量更新这些 counts 来提高 BPE 训练速度，而不是显式遍历每一对 byte 来计算 pair frequencies。使用这种 caching 过程可以获得显著加速，不过我们注意到 BPE 训练中的 merging 部分在 Python 中不可并行化。

低资源提示：Profiling

你应该使用 cProfile 或 py-spy 等 profiling 工具来找出实现中的瓶颈，并集中优化这些瓶颈。

低资源提示：“缩小规模”（Downscaling）

我们建议不要直接跳到完整 TinyStories 数据集上训练 tokenizer，而是先在一小部分数据上训练：一个“debug dataset”。例如，你可以改为在 TinyStories validation set 上训练 tokenizer，它是 22K 篇文档，而不是 2.12M 篇。这体现了一个通用策略：尽可能缩小规模以加速开发，例如

第 9 页

使用更小的数据集、更小的模型规模等。选择 debug dataset 或 hyperparameter config 的大小需要仔细考虑：你希望 debug set 足够大，以便具有与完整配置相同的瓶颈（这样你做的优化才能泛化），但也不能大到运行起来要花很久。

Problem (train_bpe)：BPE Tokenizer 训练（15 分）

Deliverable： 写一个函数：给定输入文本文件的路径，训练一个（byte-level）BPE tokenizer。你的 BPE 训练函数应处理（至少）以下输入参数：

Input

你的 BPE 训练函数应返回得到的 vocabulary 和 merges：

Output

要用我们提供的测试来测试你的 BPE 训练函数，你首先需要在 [adapters.run_train_bpe] 实现 test adapter。然后运行：

uv run pytest tests/test_train_bpe.py

你的实现应能通过所有测试。可选地（这可能是很大的时间投入），你可以用某种 systems language 实现训练方法的关键部分，例如 C++（考虑 cppyy 或 nanobind）或 Rust（使用 PyO3）。如果你这样做，要注意哪些操作需要复制、哪些需要直接从 Python 内存读取，并确保留下构建说明，或确保仅使用 pyproject.toml 就能构建。还要注意，GPT-2 regex 在大多数 regex engine 中支持得不好，并且在大多数支持它的 engine 中会太慢。我们已确认 Oniguruma 相当快并支持 negative lookahead，而 Python 中的 regex 包如果有区别的话，甚至更快。

Problem (train_bpe_tinystories)：在 TinyStories 上进行 BPE 训练（2 分）

(a) 在 TinyStories 数据集上训练一个 byte-level BPE tokenizer，最大词表大小为 10,000。确保将 TinyStories 的 <|endoftext|> special token 加入词表。将得到的 vocabulary 和 merges 序列化到磁盘，以便进一步检查。训练花了多少时间和内存？词表中最长的 token 是什么？它合理吗？

Resource requirements： ≤30 分钟（无 GPU），≤30 GB RAM

第 10 页

Hint 如果使用 multiprocessing 进行 pre-tokenization，并利用以下两个事实，你应该能让 BPE 训练耗时低于 2 分钟：

(a) <|endoftext|> token 在数据文件中分隔文档。 (b) <|endoftext|> token 在应用 BPE merges 之前作为特殊情况处理。

Deliverable： 一到两句话的回答。

(b) 对你的代码做 profile。tokenizer 训练过程中哪一部分耗时最多？

Deliverable： 一到两句话的回答。

接下来，我们将尝试在 OpenWebText 数据集上训练一个 byte-level BPE tokenizer。和之前一样，我们建议先看一下数据集，以更好地理解其内容。

Problem (train_bpe_expts_owt)：在 OpenWebText 上进行 BPE 训练（2 分）

(a) 在 OpenWebText 数据集上训练一个 byte-level BPE tokenizer，最大词表大小为 32,000。将得到的 vocabulary 和 merges 序列化到磁盘，以便进一步检查。词表中最长的 token 是什么？它合理吗？

Resource requirements： ≤12 小时（无 GPU），≤100 GB RAM

Deliverable： 一到两句话的回答。

(b) 比较并对照你在 TinyStories 上训练得到的 tokenizer 与在 OpenWebText 上训练得到的 tokenizer。

Deliverable： 一到两句话的回答。

2.6 BPE Tokenizer：编码与解码

在作业的前一部分，我们实现了一个函数，用于在输入文本上训练 BPE tokenizer，从而得到 tokenizer vocabulary 和 BPE merges 列表。现在，我们将实现一个 BPE tokenizer，它加载提供的 vocabulary 和 merges 列表，并使用它们将文本编码为 token IDs / 从 token IDs 解码文本。

2.6.1 编码文本

通过 BPE 编码文本的过程与训练 BPE vocabulary 的过程类似。主要有几个步骤。

Step 1: Pre-tokenize. 我们首先对序列做 pre-tokenize，并将每个 pre-token 表示为 UTF-8 bytes 序列，就像在 BPE 训练中那样。我们会在每个 pre-token 内将这些 bytes merge 成 vocabulary elements，独立处理每个 pre-token（不跨 pre-token 边界做 merges）。

Step 2: Apply the merges. 然后，我们取 BPE 训练期间创建的 vocabulary element merges 序列，并按相同的创建顺序将其应用到我们的 pre-tokens 上。

Example (bpe_encoding)：BPE 编码示例

例如，假设我们的输入字符串是 'the cat ate'，我们的 vocabulary 是：

{0: b' ', 1: b'a', 2: b'c', 3: b'e', 4: b'h', 5: b't', 6: b'th', 7: b' c', 8: b' a', 9: b'the', 10: b' at'}

而我们学到的 merges 是：

[(b't', b'h'), (b' ', b'c'), (b' ', b'a'), (b'th', b'e'), (b' a',

第 11 页

b't')]。首先，我们的 pre-tokenizer 会把该字符串切分为 ['the', ' cat', ' ate']。然后，我们会查看每个 pre-token 并应用 BPE merges。

第一个 pre-token 'the' 最初表示为 [b't', b'h', b'e']。查看我们的 merges 列表，我们确定第一个适用的 merge 是 (b't', b'h')，并用它把该 pre-token 转换为 [b'th', b'e']。然后，我们回到 merges 列表，确定下一个适用的 merge 是 (b'th', b'e')，它把该 pre-token 转换为 [b'the']。最后，再看 merges 列表，我们看到没有更多 merge 适用于该字符串（因为整个 pre-token 已经合并成单个 token），因此 BPE merges 的应用就完成了。对应的 integer 序列是 [9]。

对其余 pre-tokens 重复这一过程，我们看到 pre-token ' cat' 在应用 BPE merges 后表示为 [b' c', b'a', b't']，它得到 integer 序列 [7, 1, 5]。最后一个 pre-token ' ate' 在应用 BPE merges 后是 [b' at', b'e']，它得到 integer 序列 [10, 3]。因此，对输入字符串进行 encoding 的最终结果是 [9, 7, 1, 5, 10, 3]。

Special tokens（特殊 tokens）

你的 tokenizer 在 encoding 文本时应能正确处理用户定义的 special tokens（在构造 tokenizer 时提供）。

Memory considerations（内存考虑）

假设我们想 tokenize 一个无法放进内存的大文本文件。为了高效地 tokenize 这个大文件（或任何其他数据流），我们需要把它切成可管理的 chunks，并依次处理每个 chunk，使内存复杂度是常数，而不是与文本大小成线性关系。这样做时，我们需要确保一个 token 不会跨越 chunk 边界，否则得到的 tokenization 会与把整个序列在内存中 tokenize 的 naïve 方法不同。

2.6.2 Decoding text（解码文本）

要把一串 integer token IDs 解码回原始文本，我们可以简单地查找每个 ID 在 vocabulary 中对应的条目（一个 byte sequence），把它们拼接在一起，然后把 bytes 解码成 Unicode 字符串。注意，输入 IDs 不保证会映射到合法的 Unicode 字符串（因为用户可以输入任意 integer IDs 序列）。如果输入 token IDs 不能产生合法的 Unicode 字符串，你应该把格式错误的 bytes 替换为官方 Unicode replacement character U+FFFD。[注 3] bytes.decode 的 errors 参数控制 Unicode decoding 错误如何处理，使用 errors='replace' 会自动把格式错误的数据替换为 replacement marker。

Problem (tokenizer): Implementing the tokenizer (15 points)（问题（tokenizer）：实现 tokenizer（15 分））

Deliverable（交付物）： 实现一个 Tokenizer 类，给定一个 vocabulary 和一个 merges 列表，它能把文本编码为 integer IDs，并把 integer IDs 解码为文本。你的 tokenizer 还应支持用户提供的 special tokens（如果它们还不在 vocabulary 中，就把它们追加到 vocabulary）。我们推荐以下 interface：

[注 3] 关于 Unicode replacement character 的更多信息，见 en.wikipedia.org/wiki/Specials_(Unicode_block)#Replacement_character。

第 12 页

def __init__(self, vocab, merges, special_tokens=None)

从给定的 vocabulary、merges 列表以及（可选的）special tokens 列表构造一个 tokenizer。该函数应接受以下参数：

def from_files(cls, vocab_filepath, merges_filepath, special_tokens=None)

类方法：从序列化的 vocabulary 和 merges 列表（格式与你 BPE 训练代码输出的相同）以及（可选的）special tokens 列表构造并返回一个 Tokenizer。该方法应接受以下附加参数：

def encode(self, text: str) -> list[int]

把输入文本编码为 token IDs 序列。

def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]

给定一个字符串的 iterable（例如一个 Python file handle），返回一个惰性 yield token IDs 的 generator。这对于无法直接加载进内存的大文件的内存高效 tokenization 是必需的。

def decode(self, ids: list[int]) -> str

把 token IDs 序列解码为文本。

要用我们提供的测试来测试你的 Tokenizer，你首先需要在 [adapters.get_tokenizer] 实现 test adapter。然后运行：

$ uv run pytest tests/test_tokenizer.py

你的实现应能通过所有测试。

2.7 Experiments（实验）

Problem (tokenizer_experiments): Experiments with tokenizers (4 points)（问题（tokenizer_experiments）：tokenizer 实验（4 分））

(a) 从 TinyStories 和 OpenWebText 中采样 10 篇文档。使用你之前训练好的 TinyStories 和 OpenWebText tokenizers（vocabulary size 分别为 10K 和 32K），把这些采样文档编码为 integer IDs。每个 tokenizer 的 compression ratio（bytes/token）是多少？

Deliverable（交付物）： 一到两句话的回答。

(b) 如果你用 TinyStories tokenizer 来 tokenize 你的 OpenWebText 样本，会发生什么？比较 compression ratio 和/或用定性方式描述会发生什么。

Deliverable（交付物）： 一到两句话的回答。

(c) 估计你的 tokenizer 的 throughput（例如，以 bytes/second 计）。tokenize Pile 数据集（825GB 文本）需要多长时间？

Deliverable（交付物）： 一到两句话的回答。

(d) 使用你的 TinyStories 和 OpenWebText tokenizers，把各自的 training 和 development 数据集编码为 integer token IDs 序列。我们之后会用它来训练我们的

第 13 页

language model。我们建议把 token IDs 序列化为 dtype 为 uint16 的 NumPy array。为什么 uint16 是合适的选择？

Deliverable（交付物）： 一到两句话的回答。

3 Transformer Language Model Architecture（Transformer 语言模型架构）

language model 以一批 integer token IDs 序列作为输入（即形状为 $(\text{batch\_size}, \text{sequence\_length})$ 的 torch.Tensor），并返回 vocabulary 上的（批量的）归一化概率分布（即形状为 $(\text{batch\_size}, \text{sequence\_length}, \text{vocab\_size})$ 的 PyTorch Tensor），其中预测分布是针对每个输入 token 的下一个词。训练 language model 时，我们使用这些 next-word 预测来计算实际的下一个词与预测的下一个词之间的 cross-entropy loss。在 inference 期间从 language model 生成文本时，我们取最后一个时间步（即序列中的最后一项）的预测 next-word 分布来生成序列中的下一个 token（例如，取概率最高的 token、从分布中采样等），把生成的 token 加到输入序列中，然后重复。

在作业的这一部分，你将从头构建这个 Transformer 语言模型。我们会先对模型做高层描述，然后逐步详述各个组件。

3.1 Transformer LM（Transformer 语言模型）

给定一个 token IDs 序列，Transformer 语言模型使用一个 input embedding 把 token IDs 转换为稠密向量，把 embedding 后的 tokens 通过 num_layers 个 Transformer blocks，然后应用一个学习到的 linear projection（“output embedding”或“LM head”）来产生预测的 next-token logits。示意图见图 1。

图 1 中的标签（按原图顺序）：

图 1： 我们的 Transformer 语言模型概览。

图 2 中的标签（按原图顺序）：

图 2： 一个 pre-norm Transformer block。

第 14 页

Token Embeddings

在第一步中，Transformer 将（已分批的）token ID 序列 embed 成一组向量，这些向量包含 token 身份信息（图 1 中的红色块）。

更具体地说，给定一个 token ID 序列，Transformer language model 使用 token embedding layer 产生一组向量序列。每个 embedding layer 接收形状为 (batch_size, sequence_length) 的整数 tensor，并产生形状为 (batch_size, sequence_length, d_model) 的向量序列。

Pre-norm Transformer Block

embedding 之后，activations 会经过若干个结构相同的神经网络层处理。一个标准的 decoder-only Transformer language model 由 num_layers 个相同的层（通常称为 Transformer “blocks”）组成。每个 Transformer block 接收形状为 (batch_size, sequence_length, d_model) 的输入，并返回形状为 (batch_size, sequence_length, d_model) 的输出。每个 block 聚合整个序列上的信息（通过 self-attention），并对其进行非线性变换（通过 feed-forward layers）。

在 num_layers 个 Transformer blocks 之后，我们会取最终的 activations，并把它们变成 vocabulary 上的一个分布。

我们将实现 “pre-norm” Transformer block（详见 Section 3.4），它额外要求在最后一个 Transformer block 之后使用 layer normalization（详见下文），以确保其输出被正确地缩放。

在这次 normalization 之后，我们会使用一个标准的 learned linear transformation，把 Transformer blocks 的输出转换成预测的 next-token logits（例如见 A. Radford et al. [7] equation 2）。

3.2 Remark：Batching、Einsum 与高效计算

在整个 Transformer 中，我们会执行相同的计算，并将其应用到许多 batch-like 输入上。这里有几个例子：

有一种顺手的方式来执行这类操作是很有用的：既要充分利用 GPU，又要易于阅读和理解。许多 PyTorch 操作可以在 tensor 开头接收额外的 “batch-like” 维度，并在这些维度上高效地 repeat/broadcast 该操作。

例如，假设我们在做一个 position-wise、batched 的操作。我们有一个形状为 (batch_size, sequence_length, d_model) 的 “data tensor” D，并且想对形状为 (d_model, d_model) 的矩阵 A 做 batched vector-matrix multiply。在这种情况下，D @ A 会做 batched matrix multiply，这是 PyTorch 中的一个高效 primitive，其中 (batch_size, sequence_length) 维度会被 batch 起来。

因此，有帮助的做法是假设你的函数可能会被传入额外的 batch-like 维度，并把这些维度保留在 PyTorch shape 的开头。为了组织 tensor，使它们能以这种方式 batch，它们可能需要经过许多步 view、reshape 和 transpose 来塑形。这可能会有点麻烦，而且常常会变得很难看出代码在做什么、你的 tensor 形状是什么。

第 15 页

更顺手的选择是使用 torch.einsum 中的 einsum notation，或者使用像 einops 或 einx 这类与框架无关的库。两个关键操作是 einsum 和 rearrange：einsum 可以对输入 tensor 的任意维度做 tensor contractions，rearrange 可以重排、拼接和拆分任意维度。事实证明，机器学习中几乎所有的操作都是某种“维度腾挪”加 tensor contraction 的组合，偶尔还会加上一个（通常是 pointwise 的）非线性函数。这意味着，使用 einsum notation 时，你的很多代码可以更具可读性、更灵活。

我们强烈建议在本课程中学习并使用 einsum notation。之前没有接触过 einsum notation 的同学应该使用 einops（文档见此处），已经熟悉 einops 的同学应该学习更通用的 einx（此处）[注 4]。这两个包都已经安装在提供给我们的环境中。

下面给出一些如何使用 einsum notation 的例子。这些例子是 einops 文档的补充，你应该先阅读该文档。

Example (einstein_example1)：使用 einops.einsum 的批处理矩阵乘法

import torch
from einops import rearrange, einsum
## Basic implementation
Y = D @ A.T
# Hard to tell the input and output shapes and what they mean.
# What shapes can D and A have, and do any of these have unexpected behavior?
## Einsum is self-documenting and robust
#                          D                A     ->          Y
Y = einsum(D, A, "batch sequence d_in, d_out d_in -> batch sequence d_out")
## Or, a batched version where D can have any leading dimensions but A is constrained.
Y = einsum(D, A, "... d_in, d_out d_in -> ... d_out")

Example (einstein_example2)：使用 einops.rearrange 的广播操作

我们有一批图像，并且对于每张图像，我们想基于某个缩放因子生成 10 个变暗的版本：

images = torch.randn(64, 128, 128, 3)  # (batch, height, width, channel)
dim_by = torch.linspace(start=0.0, end=1.0, steps=10)
## Reshape and multiply
dim_value = rearrange(dim_by,    "dim_value              -> 1 dim_value 1 1 1")
images_rearr = rearrange(images, "b height width channel -> b 1 height width channel")
dimmed_images = images_rearr * dim_value
## Or in one go:
dimmed_images = einsum(
    images, dim_by,
    "batch height width channel, dim_value -> batch dim_value height width channel"
)

[注 4] 值得注意：尽管 einops 有大量支持，einx 还没有经过同等程度的实战检验。如果你发现 einx 有任何限制或 bug，可以放心回退到使用 einops 加一些更普通的 PyTorch。

第 16 页

Example (einstein_example3)：使用 einops.rearrange 进行像素混合

假设我们有一批图像，表示为一个形状为 (batch, height, width, channel) 的 tensor，并且我们想对图像的所有像素执行一个线性变换，但这个变换应该对每个 channel 独立进行。我们的线性变换表示为形状为 (height * width, height * width) 的矩阵 B。

channels_last = torch.randn(64, 32, 32, 3)  # (batch, height, width, channel)
B = torch.randn(32*32, 32*32)
## Rearrange an image tensor for mixing across all pixels
channels_last_flat = channels_last.view(
    -1, channels_last.size(1) * channels_last.size(2), channels_last.size(3)
)
channels_first_flat = channels_last_flat.transpose(1, 2)
channels_first_flat_transformed = channels_first_flat @ B.T
channels_last_flat_transformed = channels_first_flat_transformed.transpose(1, 2)
channels_last_transformed = channels_last_flat_transformed.view(*channels_last.shape)

改为使用 einops：

height = width = 32
## Rearrange replaces clunky torch view + transpose
channels_first = rearrange(
    channels_last,
    "batch height width channel -> batch channel (height width)"
)
channels_first_transformed = einsum(
    channels_first, B,
    "batch channel pixel_in, pixel_out pixel_in -> batch channel pixel_out"
)
channels_last_transformed = rearrange(
    channels_first_transformed,
    "batch channel (height width) -> batch height width channel",
    height=height, width=width
)

或者，如果你想玩点野的：用 einx.dot 一次性完成（einx 中等价于 einops.einsum 的写法）

height = width = 32
channels_last_transformed = einx.dot(
    "batch row_in col_in channel, (row_out col_out) (row_in col_in)"
    "-> batch row_out col_out channel",
    channels_last, B,
    col_in=width, col_out=width
)

这里的第一个实现可以通过在前后放置注释来说明输入和输出形状来改进，但这样很笨重，而且容易出 bug。使用 einsum notation 时，文档就是实现！

einsum notation 可以处理任意输入 batching 维度，但还有一个关键好处：它是自文档化的。在代码中，输入和输出 tensor 的相关形状会清楚得多。

第 17 页

使用 einsum 记号的张量。对于其余张量，可以考虑使用 Tensor 类型提示，例如使用 jaxtyping 库（并非 JAX 专用）。

我们会在 assignment 2 中进一步讨论使用 einsum 记号的性能影响，但现在你只需知道，它们几乎总是比替代方案更好！

3.2.1 Mathematical Notation and Memory Ordering

许多机器学习论文在记号中使用行向量（row vectors），由此得到的表示与 NumPy 和 PyTorch 默认使用的 row-major 内存顺序很好地契合。使用行向量时，线性变换的形式为

$$y = xW^\top \tag{1}$$

其中 $W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$ 为 row-major，$x \in \mathbb{R}^{1 \times d_{\text{in}}}$ 为行向量。注意，这让我们可以通过增大 $x$ 的最外层维度来对输入做批处理，也就是说，我们可以把向量输入 $x$ 替换为矩阵输入 $X \in \mathbb{R}^{\text{batch} \times d_{\text{in}}}$。

在线性代数中，更常见的做法是使用列向量（column vectors），此时线性变换的形式为

$$y = Wx \tag{2}$$

给定 row-major 的 $W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$ 和列向量 $x \in \mathbb{R}^{d_{\text{in}}}$。在这种设置下要对输入做批处理，$x$ 的 batch 维度必须放在最后，因此 $x$ 需要替换为矩阵 $\tilde{X} \in \mathbb{R}^{d_{\text{in}} \times \text{batch}}$。

在本 assignment 中，数学记号主要会使用列向量，因为数学通常遵循这种记号。你需要记住：如果你想使用普通的矩阵乘法记号，就必须对矩阵应用转置，如公式 1 中的行向量约定所示，因为 PyTorch 使用 row-major 内存顺序。如果你用 einsum 来做线性代数运算，那么只要正确标注各轴，这应该就不成问题。顺便一提，其他语言/线性代数包（如 Matlab、Julia 和 Fortran）都使用 column-major 内存顺序，这意味着 batch 维度放在最后，但 Python 及相关包采用了 C 标准的 row-major 顺序。

3.3 Basic Building Blocks: Linear and Embedding Modules

3.3.1 Parameter Initialization

有效地训练神经网络通常需要谨慎地初始化模型参数——糟糕的初始化可能导致不良行为，例如梯度消失或梯度爆炸。Pre-norm transformer 对初始化异常鲁棒，但初始化仍然会对训练速度和收敛产生显著影响。由于本 assignment 已经很长，我们会把细节留到 assignment 3，转而给你一些在大多数情况下都应该效果不错的近似初始化。目前请使用：

你应该使用 torch.nn.init.trunc_normal_ 来初始化截断正态（truncated normal）权重。

3.3.2 Linear Module

Linear 层是 Transformer 以及一般神经网络的基本构建块。首先，你要实现自己的 Linear 类，它继承自 torch.nn.Module，并执行线性变换：

第 18 页

$$y = Wx \tag{3}$$

注意，按照大多数现代 LLM 的做法，我们不包含 bias 项。

Problem (linear): Implementing the linear module (1 point)

Deliverable: 实现一个继承自 torch.nn.Module 并执行线性变换的 Linear 类。你的实现应遵循 PyTorch 内置 nn.Linear 模块的接口，只是不要有 bias 参数或 bias 项。我们推荐以下接口：

def __init__(self, in_features, out_features, device=None, dtype=None)

构造一个线性变换模块。该函数应接受以下参数：

def forward(self, x: torch.Tensor) -> torch.Tensor

将线性变换应用于输入。

请确保：

初始化时，使用上面的设置，并用 torch.nn.init.trunc_normal_ 初始化权重。

要测试你的 Linear 模块，请在 [adapters.run_linear] 处实现测试适配器。该适配器应把给定的权重加载到你的 Linear 模块中。为此你可以使用 Module.load_state_dict。然后运行：

uv run pytest -k test_linear

3.3.3 Embedding Module

如上所述，Transformer 的第一层是一个 embedding 层，它把整数 token ID 映射到维度为 $d_{\text{model}}$ 的向量空间。我们将实现一个自定义的 Embedding 类，它继承自 torch.nn.Module（因此你不应使用 nn.Embedding）。forward 方法应通过索引一个形状为 (vocab_size, d_model) 的 embedding 矩阵，并使用形状为 (batch_size, sequence_length) 的 token ID 的 torch.LongTensor，来为每个 token ID 选择对应的 embedding 向量。

Problem (embedding): Implement the embedding module (1 point)

Deliverable: 实现继承自 torch.nn.Module 并执行 embedding lookup 的 Embedding 类。你的实现应遵循 PyTorch 内置 nn.Embedding 模块的接口。我们推荐以下接口：

第 19 页

def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None)

构造一个 embedding 模块。该函数应接受以下参数：

def forward(self, token_ids: torch.Tensor) -> torch.Tensor

查找给定 token ID 对应的 embedding 向量。

请确保：

同样，使用上面的设置进行初始化，并使用 torch.nn.init.trunc_normal_ 初始化权重。

要测试你的实现，请在 [adapters.run_embedding] 处实现测试适配器。然后运行：

uv run pytest -k test_embedding

3.4 Pre-Norm Transformer Block

每个 Transformer block 有两个子层：multi-head self-attention 机制和一个 position-wise feed-forward network（[A. Vaswani et al., 2017]，第 3.1 节）。

在最初的 Transformer 论文中，模型在两个子层中的每一个周围都使用 residual connection，随后进行 layer normalization。这种架构通常被称为 “post-norm” Transformer，因为 layer normalization 被应用于子层输出。然而，多项工作发现，把 layer normalization 从每个子层的输出移到每个子层的输入（并在最后一个 Transformer block 之后额外加一个 layer normalization）可以提高 Transformer 的训练稳定性 [T. Q. Nguyen et al., 2019; R. Xiong et al., 2020]——这种 “pre-norm” Transformer block 的可视化表示见图 2。每个 Transformer block 子层的输出随后通过 residual connection 加到该子层的输入上（A. Vaswani et al. [8]，第 5.4 节）。对 pre-norm 的一种直觉是：从输入 embeddings 到 Transformer 的最终输出，存在一条干净的、没有任何 normalization 的 “residual stream”，据称这能改善梯度流动。这种 pre-norm Transformer 现在是当今语言模型（例如 GPT-3、LLaMA、PaLM 等）中使用的标准，因此我们将实现这一变体。我们会依次走完 pre-norm Transformer block 的每个组件，并按顺序实现它们。

3.4.1 Root Mean Square Layer Normalization

A. Vaswani et al. [8] 最初的 Transformer 实现使用 layer normalization [J. L. Ba et al., 2016] 来归一化激活值。遵循 H. Touvron et al. [12]，我们将使用 root mean square layer normalization（RMSNorm；B. Zhang et al. [13]，公式 4）来进行 layer normalization。给定一个激活值向量 $a \in \mathbb{R}^{d_{\text{model}}}$，RMSNorm 会按如下方式重新缩放每个激活值 $a_i$：

$$\text{RMSNorm}(a_i) = \frac{a_i}{\text{RMS}(a)} g_i \tag{4}$$

第 20 页

其中 $\mathrm{RMS}(a) = \sqrt{\frac{1}{d_{\text{model}}} \sum_{i=1}^{d_{\text{model}}} a_i^2 + \varepsilon}$。这里，$g_i$ 是一个可学习的“gain”参数（总共有 $d_{\text{model}}$ 个这样的参数），而 $\varepsilon$ 是一个超参数，通常固定为 1e-5。

在平方输入之前，你应将输入 upcast 到 torch.float32 以防止溢出。总体而言，你的 forward 方法应如下所示：

  in_dtype = x.dtype
  x = x.to(torch.float32)
  # Your code here performing RMSNorm
  ...
  result = ...
  # Return the result in the original dtype
  return result.to(in_dtype)

Problem (rmsnorm)：均方根层归一化（Root Mean Square Layer Normalization，1 分）

Deliverable: 将 RMSNorm 实现为一个 torch.nn.Module。我们推荐以下接口：

def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None)

构造 RMSNorm 模块。该函数应接受以下参数：

def forward(self, x: torch.Tensor) -> torch.Tensor

处理形状为 (batch_size, sequence_length, d_model) 的输入张量，并返回相同形状的张量。

Note: 记得在执行 normalization 之前将输入 upcast 到 torch.float32（之后再 downcast 回原始 dtype），如上所述。

为了测试你的实现，请在 [adapters.run_rmsnorm] 处实现 test adapter。然后运行 uv run pytest -k test_rmsnorm。

第 21 页

3.4.2 逐位置前馈网络（Position-Wise Feed-Forward Network）

图 3： 比较 SiLU（又称 Swish）与 ReLU 激活函数。

在原始 Transformer 论文（A. Vaswani 等人 [8] 的第 3.3 节）中，Transformer 前馈网络由两个线性变换组成，二者之间有一个 ReLU 激活（$\mathrm{ReLU}(x) = \max(0, x)$）。在那个原始架构中，内层前馈层的维度通常是输入维度的 4 倍。

然而，现代语言模型相较于这一原始设计通常会引入两个主要变化：它们使用另一种激活函数，并采用门控机制。具体来说，我们将实现被 Llama 3 [A. Grattafiori et al., 2024] 和 Qwen 2.5 [A. Yang et al., 2024] 等 LLM 采用的“SwiGLU”激活函数，它把 SiLU（常称为 Swish）激活与一种称为 Gated Linear Unit（GLU）的门控机制结合起来。我们还将省略线性层中有时使用的 bias 项，遵循自 PaLM [A. Chowdhery et al., 2022] 和 LLaMA [H. Touvron et al., 2023] 以来大多数现代 LLM 的做法。

SiLU 或 Swish 激活函数 [D. Hendrycks et al., 2016; S. Elfwing et al., 2017] 定义如下：

$$\operatorname{SiLU}(x) = x \cdot \sigma(x) = \frac{x}{1 + e^{-x}} \tag{5}$$

从图 3 可以看出，SiLU 激活函数与 ReLU 激活函数相似，但在零点处是平滑的。

门控线性单元（Gated Linear Units，GLU）最初由 Y. N. Dauphin 等人 [19] 定义为：一个经过 sigmoid 函数的线性变换与另一个线性变换的逐元素乘积：

$$\operatorname{GLU}(x, W_1, W_2) = \sigma(W_1 x) \odot W_2 x, \tag{6}$$

其中 $\odot$ 表示逐元素乘法。Gated Linear Units 被建议用来“通过为梯度提供线性路径，同时保留非线性能力，从而减少深度架构中的梯度消失问题”。

第 22 页

将 SiLU/Swish 与 GLU 结合起来，就得到 SwiGLU，我们将把它用于前馈网络：

$$\operatorname{FFN}(x) = \operatorname{SwiGLU}(x, W_1, W_2, W_3) = W_2(\operatorname{SiLU}(W_1 x) \odot W_3 x), \tag{7}$$

其中 $x \in \mathbb{R}^{d_{\text{model}}}$，$W_1, W_3 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}$，$W_2 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}$，并且按惯例 $d_{\text{ff}} = \frac{8}{3} d_{\text{model}}$。在具体实现中，为了硬件效率，将其舍入到附近的 64 的倍数是可以的。

N. Shazeer [20] 首次提出将 SiLU/Swish 激活与 GLU 结合，并通过实验表明 SwiGLU 在语言建模任务上优于 ReLU 和 SiLU（不带门控）等基线。在作业后续部分，你将比较 SwiGLU 和 SiLU。尽管我们已提到这些组件的一些启发式论据（论文也提供了更多支持证据），但保持经验视角是好的：Shazeer 论文中如今著名的一句话是

“对于这些架构为何似乎有效，我们不作解释；我们把它们的成功，如同其他一切，归因于神圣的恩典。”

Problem (positionwise_feedforward)：实现逐位置前馈网络（2 分）

Deliverable: 实现 SwiGLU 前馈网络，它由 SiLU 激活函数和一个 GLU 组成。

Note: 在这种特定情况下，为了数值稳定性，你在实现中可以放心使用 torch.sigmoid。

在你的实现中，你应将 $d_{\text{ff}}$ 设置为约 $\frac{8}{3} \times d_{\text{model}}$，同时确保内层前馈层的维度是 64 的倍数，以充分利用你的硬件。为了用我们提供的测试来测试你的实现，你需要在 [adapters.run_swiglu] 处实现 test adapter。然后运行 uv run pytest -k test_swiglu 来测试你的实现。

3.4.3 相对位置嵌入（Relative Positional Embeddings）

为了将位置信息注入模型，我们将实现 Rotary Position Embeddings [J. Su et al., 2021]，通常称为 RoPE。对于 token 位置 $i$ 处的给定 query token $q^{(i)} = W_q x^{(i)} \in \mathbb{R}^d$，我们将应用成对旋转矩阵 $R_i$，得到 $q'^{(i)} = R_i q^{(i)} = R_i W_q x^{(i)}$。这里，$R_i$ 将嵌入元素对 $q^{(i)}_{2k-1:2k}$ 作为 2d 向量旋转角度 $\theta_{i,k} = \frac{i}{\Theta^{(2k-2)/d}}$，其中 $k \in \{1, \ldots, d/2\}$，$\Theta$ 为某个常数。因此，我们可以把 $R_i$ 视为大小为 $d \times d$ 的块对角矩阵，其块为 $R_i^k$，其中 $k \in \{1, \ldots, d/2\}$，且有

$$R_i^k = \begin{pmatrix} \cos(\theta_{i,k}) & \sin(\theta_{i,k}) \\ -\sin(\theta_{i,k}) & \cos(\theta_{i,k}) \end{pmatrix} \tag{8}$$

于是我们得到完整的旋转矩阵

第 23 页

$$R_i = \begin{pmatrix} R_i^1 & 0 & 0 & \cdots & 0 \\ 0 & R_i^2 & 0 & \cdots & 0 \\ 0 & 0 & R_i^3 & \cdots & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & \cdots & R_i^{d/2} \end{pmatrix}, \tag{9}$$

其中 0 表示 $2 \times 2$ 零矩阵。虽然可以构造完整的 $d \times d$ 矩阵，但一个好的解决方案应利用该矩阵的性质来更高效地实现变换。由于我们只关心给定序列内 token 的相对旋转，我们可以跨层以及不同 batch 复用为 $\cos(\theta_{i,k})$ 和 $\sin(\theta_{i,k})$ 计算出的值。如果你想优化它，可以使用一个被所有层引用的单个 RoPE 模块，并且它可以有一个在初始化时用 self.register_buffer(persistent=False) 创建的、预先计算好的 sin 和 cos 值组成的二维 buffer，而不是 nn.Parameter（因为我们不想学习这些固定的余弦和正弦值）。然后，对我们 $q^{(i)}$ 所做的完全相同的旋转过程也会对 $k^{(j)}$ 执行，按相应的 $R_j$ 旋转。注意，这一层没有可学习参数。

Problem (rope)：实现 RoPE（2 分）

Deliverable: 实现一个类 RotaryPositionalEmbedding，对输入张量应用 RoPE。

推荐以下接口：

def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None)

构造 RoPE 模块，并在需要时创建 buffers。

def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor

处理形状为 (..., seq_len, d_k) 的输入张量，并返回相同形状的张量。注意，你应能容忍 $x$ 具有任意数量的 batch 维度。你应假设 token positions 是一个形状为 (..., seq_len) 的张量，指定 $x$ 沿序列维度的 token 位置。

你应使用 token positions 来沿序列维度对你的（可能是预先计算好的）cos 和 sin 张量进行切片。

为了测试你的实现，请完成 [adapters.run_rope]，并确保它通过 uv run pytest -k test_rope。

3.4.4 缩放点积注意力（Scaled Dot-Product Attention）

我们现在将实现 A. Vaswani 等人 [8]（第 3.2.1 节）中描述的缩放点积注意力。作为预备步骤，Attention 操作的定义会用到 softmax，该操作接收一个未归一化的分数向量，并将其转换为归一化分布：

$$\operatorname{softmax}(v)_i = \frac{\exp(v_i)}{\sum_{j=1}^{n} \exp(v_j)}. \tag{10}$$

第 24 页

注意，对于较大的值，$\exp(v_i)$ 可能变成 inf（此时，$\frac{\inf}{\inf} = \mathrm{NaN}$）。我们可以通过注意到 softmax 操作对给所有输入加上任意常数 $c$ 具有不变性来避免这一点。我们可以利用这一性质来提高数值稳定性——通常，我们会从 $v$ 的所有元素中减去 $v$ 中的最大项，使新的最大项为 0。现在你将实现 softmax，并利用这个技巧来保证数值稳定性。

Problem (softmax)：实现 softmax（1 分）

Deliverable: 编写一个对 tensor 应用 softmax 操作的函数。你的函数应接收两个参数：一个 tensor 和一个维度 $i$，并对输入 tensor 的第 $i$ 维应用 softmax。输出 tensor 的形状应与输入 tensor 相同，但其第 $i$ 维现在应具有归一化的概率分布。使用从第 $i$ 维的所有元素中减去该维最大值的技巧，以避免数值稳定性问题。

为了测试你的实现，完成 [adapters.run_softmax]，并确保它通过 uv run pytest -k test_softmax_matches_pytorch。

我们现在可以用数学方式定义 Attention 操作如下：

$$\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V \tag{11}$$

其中 $Q \in \mathbb{R}^{n \times d_k}$，$K \in \mathbb{R}^{m \times d_k}$，$V \in \mathbb{R}^{m \times d_v}$。这里，$Q$、$K$ 和 $V$ 都是该操作的输入——注意它们不是可学习参数。

Masking: 有时对 attention 操作的输出进行 masking 会很方便。一个 mask 的形状应为 $M \in \{\text{True}, \text{False}\}^{n \times m}$，这个布尔矩阵的每一行 $i$ 表示 query $i$ 应该关注哪些 key。按惯例（并且稍微有点容易混淆），位置 $(i,j)$ 处的值为 True 表示 query $i$ 确实关注 key $j$，值为 False 表示该 query 不关注该 key。换句话说，在值为 True 的 $(i,j)$ 对上，“信息流动”。例如，考虑一个 $1 \times 3$ 的 mask 矩阵，其元素为 $[[\text{True}, \text{True}, \text{False}]]$。这个唯一的 query 向量只关注前两个 key。

在计算上，使用 masking 比在子序列上计算 attention 要高效得多；我们可以通过取 softmax 前的值 $(QK^\top/\sqrt{d_k})$，并对 mask 矩阵中任何为 False 的元素加上 $-\infty$ 来实现。

Problem (scaled_dot_product_attention)：实现缩放点积 attention（5 分）

Deliverable: 实现缩放点积 attention 函数。你的实现应处理形状为 (batch_size, ..., seq_len, d_k) 的 key 和 query，以及形状为 (batch_size, ..., seq_len, d_v) 的 value，其中 ... 表示任意数量的其他 batch-like 维度（如果提供的话）。实现应返回形状为 (batch_size, ..., seq_len, d_v) 的输出。关于 batch-like 维度的讨论见 Section 3.2。

你的实现还应支持可选的、由用户提供的形状为 (seq_len, seq_len) 的布尔 mask。mask 值为 True 的位置的 attention 概率应合计为 1，而 mask 值为 False 的位置的 attention 概率应为 0。

第 25 页

为了用我们提供的测试来测试你的实现，你需要实现位于 [adapters.run_scaled_dot_product_attention] 的测试 adapter。uv run pytest -k test_scaled_dot_product_attention 会在三阶输入 tensor 上测试你的实现，而 uv run pytest -k test_4d_scaled_dot_product_attention 会在四阶输入 tensor 上测试你的实现。

3.4.5 因果多头自 attention（Causal Multi-Head Self-Attention）

我们将按照 A. Vaswani et al. [8] 的 section 3.2.2 所述实现 multi-head self-attention。回顾一下，从数学上讲，应用 multi-head attention 的操作定义如下：

$$\operatorname{MultiHead}(Q,K,V)=\operatorname{Concat}(\operatorname{head}_1,\ldots,\operatorname{head}_h) \tag{12}$$

其中

$$\operatorname{head}_i=\operatorname{Attention}(Q_i,K_i,V_i) \tag{13}$$

并且 $Q_i,K_i,V_i$ 分别是 $Q,K,V$ 的 embedding 维度中大小为 $d_k$ 或 $d_v$ 的第 $i\in\{1,\ldots,h\}$ 个切片。这里的 Attention 是 Section 3.4.4 中定义的缩放点积 attention 操作。由此我们可以构造 multi-head self-attention 操作：

$$\operatorname{MultiHeadSelfAttention}(x)=W^O\operatorname{MultiHead}(W^Q x, W^K x, W^V x) \tag{14}$$

这里，可学习参数为 $W^Q\in\mathbb{R}^{h d_k\times d_{\text{model}}}$，$W^K\in\mathbb{R}^{h d_k\times d_{\text{model}}}$，$W^V\in\mathbb{R}^{h d_v\times d_{\text{model}}}$，以及 $W^O\in\mathbb{R}^{d_{\text{model}}\times h d_v}$。由于 $Q$、$K$ 和 $V$ 在 multi-head attention 操作中会被切片，我们可以认为 $W^Q$、$W^K$ 和 $W^V$ 是沿着输出维度为每个 head 分开的。当这部分能正常工作时，你应该总共用三次矩阵乘法来计算 key、value 和 query 的投影。[注 5]

因果 masking（Causal masking）

你的实现应防止模型关注序列中的未来 token。换句话说，如果模型给定一个 token 序列 $t_1,\ldots,t_n$，并且我们想计算前缀 $t_1,\ldots,t_i$（其中 $itorch.triu 或广播的索引比较来构造这个 mask，并且你应该利用 Section 3.4.4 中的缩放点积 attention 实现已经支持 attention masking 这一事实。

应用 RoPE

RoPE 应应用于 query 和 key 向量，而不是 value 向量。此外，head 维度应作为 batch 维度处理，因为在 multi-head attention 中，attention 是独立地对每个 head 应用的。这意味着对于每个 head，应将完全相同的 RoPE 旋转应用于 query 和 key 向量。

Problem (multihead_self_attention)：实现因果多头自 attention（5 分）

[注 5] 作为扩展目标，尝试将 key、query 和 value 投影合并为一个单一的权重矩阵，这样你只需要一次矩阵乘法。

第 26 页

Deliverable: 将因果多头自 attention 实现为 torch.nn.Module。你的实现应接受（至少）以下参数：

按照 A. Vaswani et al. [8]，设置 $d_k = d_v = \frac{d_{\text{model}}}{h}$。为了用我们提供的测试来测试你的实现，实现位于 [adapters.run_multihead_self_attention] 的测试 adapter。然后运行 uv run pytest -k test_multihead_self_attention 来测试你的实现。

3.5 完整的 Transformer LM

让我们从组装 Transformer block 开始（回头参考 Figure 2 会很有帮助）。一个 Transformer block 包含两个“sub-layer”，一个用于 multihead self attention，另一个用于 SwiGLU feed-forward network。在每个 sub-layer 中，我们首先执行 RMSNorm，然后执行主要操作（MHA/FF），最后加入残差连接。

具体来说，Transformer block 的前半部分（第一个“sub-layer”）应实现以下一组更新，以从输入 $x$ 产生输出 $y$：

$$y=x+\operatorname{MultiHeadSelfAttention}(\operatorname{RMSNorm}(x)) \tag{15}$$

Problem (transformer_block)：实现 Transformer block（3 分）

实现 Section 3.4 中描述、Figure 2 中展示的 pre-norm Transformer block。你的 Transformer block 应接受（至少）以下参数：

为了测试你的实现，实现 adapter [adapters.run_transformer_block]。然后运行 uv run pytest -k test_transformer_block 来测试你的实现。

Deliverable: 能通过所提供测试的 Transformer block 代码。

现在我们把各个 block 组装起来，遵循 Figure 1 中的高层示意图。按照我们在 Section 3.1.0.1 中对 embedding 的描述，将其输入到 num_layers 个 Transformer block 中，然后将其传入最终的 layer norm 和 LM head，以获得词表上的未归一化分布（即 logits）。

Problem (transformer_lm)：实现 Transformer LM（3 分）

是时候把所有东西组合起来了！实现 Section 3.1 中描述、Figure 1 中展示的 Transformer 语言模型。你的实现至少应接受前面提到的 Transformer block 的所有构造参数，以及以下额外参数：

第 27 页

为了用我们提供的测试检验你的实现，你需要先在 [adapters.run_transformer_lm] 处实现 test adapter。然后运行 uv run pytest -k test_transformer_lm 来测试你的实现。

Deliverable: 一个能通过上述测试的 Transformer LM 模块。

Resource accounting

理解 Transformer 各个部分如何消耗算力和内存是很有用的。我们将一步步做一些基本的“FLOPs accounting”。Transformer 中绝大多数 FLOPs 来自矩阵乘法，所以我们的核心思路很简单：

  1. 写下 Transformer 前向传播中的所有矩阵乘法。
  2. 把每个矩阵乘法换算成所需的 FLOPs。

对于第二步，下面这个事实很有用：

Rule: 给定 $A\in\mathbb{R}^{m\times n}$ 和 $B\in\mathbb{R}^{n\times p}$，矩阵-矩阵乘积 $AB$ 需要 $2mnp$ 个 FLOPs。

要看出这一点，注意 $(AB)[i,j] = A[i,:]\cdot B[:,j]$，而这个点积需要 $n$ 次加法和 $n$ 次乘法（$2n$ FLOPs）。于是，由于矩阵-矩阵乘积 $AB$ 有 $m\times p$ 个元素，FLOPs 的总数为 $(2n)(mp) = 2mnp$。

现在，在做下一题之前，先把你的 Transformer block 和 Transformer LM 的每个组件过一遍，列出所有矩阵乘法及其相应的 FLOPs 开销，这会很有帮助。

Problem (transformer_accounting)：Transformer LM 资源核算（5 分）

(a)

考虑一个使用我们作业架构的 GPT-2 XL 规模模型，其配置如下：

vocab_size: 50,257
context_length: 1,024
num_layers: 48
d_model: 1,600
num_heads: 25
d_ff: 4,288    # 距 (8/3) × 1,600 最近的 64 的倍数

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

Deliverable: 一到两句话的回答。

(b)

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

第 28 页

Deliverable: 一份矩阵乘法清单（带说明），以及所需的 FLOPs 总数。

(c)

基于你上面的分析，模型的哪些部分需要最多的 FLOPs？

Deliverable: 一到两句话的回答。

(d)

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

Deliverable: 对每个模型，给出模型各组件的 breakdown 及其对应的 FLOPs（占一次前向传播所需总 FLOPs 的比例）。另外，用一到两句话描述改变模型规模如何影响各组件的比例 FLOPs。

(e)

取 GPT-2 XL，把 context length 增大到 16,384。一次前向传播的总 FLOPs 如何变化？模型各组件 FLOPs 的相对贡献如何变化？

Deliverable: 一到两句话的回答。

4 训练 Transformer LM

我们现在已经有了预处理数据（通过 tokenizer）和模型（Transformer）的步骤。剩下的就是搭建所有支持训练的代码。这包括以下几部分：

4.1 Cross-entropy loss

回忆一下，Transformer 语言模型对每条长度为 $m+1$ 的序列 $x$ 以及 $i = 1, \dots, m$ 定义了分布 $p_\theta(x_{i+1}\mid x_{1:i})$。给定由长度为 $m+1$ 的序列组成的训练集 $D$，我们定义标准的 cross-entropy（负对数似然）loss 函数：

$$\ell(\theta; D) = \frac{1}{|D|}\sum_{x\in D}\sum_{i=1}^{m} -\log p_\theta(x_{i+1}\mid x_{1:i}). \tag{16}$$

（注意，Transformer 中的一次前向传播会给出所有 $i = 1, \dots, m$ 的 $p_\theta(x_{i+1}\mid x_{1:i})$。）

具体来说，Transformer 为每个位置 $i$ 计算 logits $o_i\in\mathbb{R}^{\text{vocab\_size}}$，从而有：[注 6]

$$p(x_{i+1}\mid x_{1:i}) = \mathrm{softmax}(o_i)[x_{i+1}] = \frac{\exp(o_i[x_{i+1}])}{\sum_{a=1}^{\text{vocab\_size}}\exp(o_i[a])}. \tag{17}$$

cross-entropy loss 通常是针对 logits 向量 $o_i\in\mathbb{R}^{\text{vocab\_size}}$ 和目标 $x_{i+1}$ 定义的。[注 7]

[注 6] 注意 $o_i[k]$ 指的是向量 $o_i$ 中下标为 $k$ 的值。

[注 7] 这对应于 $x_{i+1}$ 上的 Dirac delta 分布与预测分布 $\mathrm{softmax}(o_i)$ 之间的 cross-entropy。

第 29 页

实现 cross-entropy loss 时，和在 softmax 的情形下一样，需要对数值问题多加小心。

Problem (cross_entropy)：实现 cross-entropy（1 分）

Deliverable: 写一个计算 cross-entropy loss 的函数，它接收预测的 logits（$o_i$）和目标（$x_{i+1}$），计算 cross-entropy $\ell_i = -\log \mathrm{softmax}(o_i)[x_{i+1}]$。你的函数应当处理以下几点：

实现 [adapters.run_cross_entropy]，然后运行 uv run pytest -k test_cross_entropy 来测试你的实现。

Perplexity

cross-entropy 足以用于训练，但在评估模型时，我们还想报告 perplexity。对于长度为 $m$、cross-entropy loss 分别为 $\ell_1, \dots, \ell_m$ 的序列：

$$\text{perplexity} = \exp\left(\frac{1}{m}\sum_{i=1}^{m}\ell_i\right). \tag{18}$$

4.2 The SGD Optimizer

现在我们有了 loss 函数，开始探索优化器。最简单的基于梯度的优化器是随机梯度下降（SGD）。我们从随机初始化的参数 $\theta_0$ 开始。然后对每一步 $t = 0, \dots, T-1$，执行如下更新：

$$\theta_{t+1} \leftarrow \theta_t - \alpha_t \nabla L(\theta_t; B_t), \tag{19}$$

其中 $B_t$ 是从数据集 $D$ 中采样得到的一个随机 batch，学习率 $\alpha_t$ 和 batch 大小 $|B_t|$ 是超参数。

4.2.1 Implementing SGD in PyTorch

要实现我们的优化器，我们将继承 PyTorch 的 torch.optim.Optimizer 类。Optimizer 子类必须实现两个方法：

def __init__(self, params, ...) 应当初始化你的优化器。这里 params 是将要被优化的参数集合（或参数组，以防用户想对模型的不同部分使用不同的超参数，比如学习率）。务必把 params 传给基类的 __init__ 方法，基类会保存这些参数供 step 使用。你可以根据优化器的不同接受额外的参数（例如学习率就是常见的一个），并把它们以字典的形式传给基类构造函数，字典的键是你为这些参数选择的名字（字符串）。

def step(self) 应当对参数做一次更新。在训练循环中，它会在 backward 之后被调用，所以你可以拿到最后一个 batch 上的梯度。这个方法应当遍历每个参数张量 p 并原地修改它们，即设置 p.data，它保存着

第 30 页

与该参数关联的 tensor（依据梯度 p.grad，如果存在的话），即表示损失对该参数的梯度的 tensor。

PyTorch optimizer API 有一些细节，所以用一个例子来解释会更容易。为了让例子更丰富，我们实现一个 SGD 的小变体：学习率在训练过程中衰减，从初始学习率 $\alpha$ 开始，并随时间依次取更小的步长：

$$\theta_{t+1} = \theta_t - \frac{\alpha}{\sqrt{t+1}} \nabla L(\theta_t; B_t) \tag{20}$$

下面看看这个版本的 SGD 如何实现为 PyTorch Optimizer：

from collections.abc import Callable, Iterable
from typing import Optional
import torch
import math
class SGD(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3):
        if lr < 0:
            raise ValueError(f"Invalid learning rate: {lr}")
        defaults = {"lr": lr}
        super().__init__(params, defaults)
    def step(self, closure: Optional[Callable] = None):
        loss = None if closure is None else closure()
        for group in self.param_groups:
            lr = group["lr"]  # Get the learning rate.
            for p in group["params"]:
                if p.grad is None:
                    continue
                state = self.state[p]  # Get state associated with p.
                t = state.get("t", 0)  # Get iteration number from the state, or 0.
                grad = p.grad.data  # Get the gradient of loss with respect to p.
                p.data -= lr / math.sqrt(t + 1) * grad  # Update weight tensor in-place.
                state["t"] = t + 1  # Increment iteration number.
        return loss

在 __init__ 中，我们把参数以及默认超参数传给基类构造函数（参数可能按组传入，每组有不同的超参数）。如果参数只是单个 torch.nn.Parameter 对象的集合，基类构造函数会创建一个组，并为其分配默认超参数。然后，在 step 中，我们遍历每个参数组，再遍历该组中的每个参数，并应用 Equation 20。这里，我们把迭代编号保存为与每个参数关联的 state：先读取这个值，在梯度更新中使用它，然后再更新它。API 规定用户可以传入一个可调用的 closure，在 optimizer step 之前重新计算 loss。我们要用的 optimizer 不需要这个，但为了符合 API，我们加上了它。

为了看看它如何工作，我们可以使用下面这个最小训练循环示例：

weights = torch.nn.Parameter(5 * torch.randn((10, 10)))
opt = SGD([weights], lr=1)
for t in range(100):
    opt.zero_grad()  # Reset the gradients for all learnable parameters.
    loss = (weights**2).mean() # Compute a scalar loss value.
    print(loss.cpu().item())

第 31 页

    loss.backward()  # Run backward pass, which computes gradients.
    opt.step()  # Run optimizer step.

这是训练循环的典型结构：每次迭代中，我们会计算 loss，并运行一步 optimizer。训练语言模型时，可学习参数来自模型（在 PyTorch 中，m.parameters() 给出这个集合）。loss 会在一批采样数据上计算，但训练循环的基本结构相同。

Problem (learning_rate_tuning): 调整学习率 (1 point)

正如我们将看到的，对训练影响最大的超参数之一就是学习率。让我们在玩具示例中实际看看。用另外三个学习率值运行上面的 SGD 示例：1e1、1e2 和 1e3，只训练 10 次迭代。每个学习率下 loss 会发生什么？它是衰减得更快、更慢，还是发散（即在训练过程中增大）？

Deliverable: 用一到两句话回应你观察到的行为。

4.3 AdamW

现代语言模型通常使用更复杂的 optimizer 训练，而不是 SGD。最近使用的大多数 optimizer 都是 Adam optimizer 的衍生版本 [D. P. Kingma et al., 2015]。我们将使用 AdamW [I. Loshchilov et al., 2019]，它在最近的工作中被广泛使用。AdamW 对 Adam 提出了一种修改，通过加入 weight decay（每次迭代中把参数拉向 0）来改进正则化，并且这种方式与梯度更新解耦。我们将按照 I. Loshchilov 等人 [23] 的算法 2 来实现 AdamW。

AdamW 是有状态的：对于每个参数，它维护该参数第一矩和第二矩的运行估计。因此，AdamW 使用额外内存来换取更好的稳定性和收敛性。除了学习率 $\alpha$，AdamW 还有一对超参数 $(\beta_1, \beta_2)$ 控制矩估计的更新，以及一个 weight decay 率 $\lambda$。典型应用把 $(\beta_1, \beta_2)$ 设为 $(0.9, 0.999)$，但像 LLaMA [H. Touvron et al., 2023] 和 GPT-3 [T. B. Brown et al., 2020] 这样的大语言模型常用 $(0.9, 0.95)$ 训练。该算法可以写成如下形式，其中 $\varepsilon$ 是一个小值（例如 $10^{-8}$），用于在我们得到极小的 $v$ 值时提高数值稳定性：

第 32 页

Algorithm 1: AdamW Optimizer

行 步骤 说明
1 $\text{init}(\theta)$ ▷ 初始化可学习参数
2 $m \leftarrow 0$ ▷ 第一矩向量的初始值；形状与 $\theta$ 相同
3 $v \leftarrow 0$ ▷ 第二矩向量的初始值；形状与 $\theta$ 相同
4 for $t = 1, \ldots, T$ do
5 采样一批数据 $B_t$
6 $g \leftarrow \nabla_\theta \ell(\theta; B_t)$ ▷ 计算 loss 的梯度
7 $\alpha_t \leftarrow \alpha \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t}$ ▷ 计算第 $t$ 次迭代调整后的 $\alpha$
8 $\theta \leftarrow \theta - \alpha \lambda \theta$ ▷ 应用 weight decay
9 $m \leftarrow \beta_1 m + (1-\beta_1) g$ ▷ 更新第一矩估计
10 $v \leftarrow \beta_2 v + (1-\beta_2) g^2$ ▷ 更新第二矩估计
11 $\theta \leftarrow \theta - \alpha_t \frac{m}{\sqrt{v}+\varepsilon}$ ▷ 应用经矩调整的权重更新
12 end for

Algorithm 1：AdamW optimizer 伪代码。

注意 $t$ 从 1 开始。现在你要实现这个 optimizer。

Problem (adamw): 实现 AdamW (2 points)

Deliverable: 实现 AdamW optimizer，作为 torch.optim.Optimizer 的子类。你的类应该在 __init__ 中接收学习率 $\alpha$，以及 $\beta$、$\varepsilon$ 和 $\lambda$ 超参数。为了帮你保存 state，基类 Optimizer 提供了一个字典 self.state，它把 nn.Parameter 对象映射到一个字典，该字典存储你为该参数需要的任何信息（对于 AdamW，就是矩估计）。实现 [adapters.get_adamw_cls]，并确保它通过 uv run pytest -k test_adamw。

Problem (adamw_accounting): 使用 AdamW 训练的资源核算 (2 points)

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

(a)

运行 AdamW 需要多少峰值内存？请根据参数、activations、梯度和 optimizer state 的内存使用来分解你的答案。用 batch_size 和模型超参数（vocab_size、context_length、num_layers、d_model、num_heads）表示你的答案。假设 $d_{\text{ff}} = \frac{8}{3} \times d_{\text{model}}$。

为简化起见，在计算 activations 的内存使用时，只考虑以下组件：

第 33 页

Deliverable: 分别给出 parameters、activations、gradients 和 optimizer state 的代数表达式，以及总和的代数表达式。

(b) 将你的答案代入一个 GPT-2 XL 形状的模型，得到一个只依赖于 batch_size 的表达式。在仍然能放入 80GB 内存的前提下，你能使用的最大 batch size 是多少？

Deliverable: 一个形如 $a \cdot \text{batch\_size} + b$ 的表达式，其中 $a$、$b$ 为数值，并给出一个表示最大 batch size 的数。

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

Deliverable: 一个代数表达式，并附简短理由。

(d) Model FLOPs utilization (MFU) 定义为观察到的吞吐量（tokens/second）相对于硬件理论峰值 FLOP 吞吐量的比值 [A. Chowdhery et al., 2022]。一块 NVIDIA H100 GPU 在“float32”（实际上是 TensorFloat-32，而现实中是“bfloat19”）运算上的理论峰值为 495 teraFLOP/s。假设你能达到 50% MFU，在单块 H100 上训练 GPT-2 XL 400K 步、batch size 为 1024，需要多长时间？按照 J. Kaplan et al. [25] 和 J. Hoffmann et al. [26]，假设 backward pass 的 FLOPs 是 forward pass 的两倍。

Deliverable: 训练所需的小时数，并附简短理由。

4.4 Learning rate scheduling

导致 loss 下降最快的学习率取值在训练过程中经常会变化。在训练 Transformers 时，通常会使用 learning rate schedule：开始时使用较大的 learning rate，在初期进行更快的更新，并随着模型训练慢慢将其衰减到较小的值。[注 8] 在本作业中，我们将实现用于训练 LLaMA [H. Touvron et al., 2023] 的 cosine annealing schedule。

[注 8] 有时也会使用让学习率回升（restarts）的 schedule，以帮助越过局部最小值。

scheduler 就是一个函数，它接收当前 step $t$ 和其他相关参数（例如初始和最终学习率），并返回在第 $t$ 步进行 gradient update 时要使用的学习率。最简单的 schedule 是常数函数，对任意 $t$ 都返回相同的学习率。

cosine annealing learning rate schedule 接收：(i) 当前 iteration $t$，(ii) 最大学习率 $\alpha_{\max}$，(iii) 最小（最终）学习率 $\alpha_{\min}$，(iv) warm-up iteration 数 $T_w$，以及 (v) cosine annealing 的最终 iteration $T_c$。第 $t$ 个 iteration 的学习率定义如下：

$$\alpha_t = \alpha_{\min} + \frac{1}{2}\left(1 + \cos\left(\frac{t-T_w}{T_c-T_w}\pi\right)\right)(\alpha_{\max} - \alpha_{\min})$$

Problem (learning_rate_schedule)：实现带 warmup 的 cosine learning rate schedule（1 分）

第 34 页

编写一个函数，接收 $t$、$\alpha_{\max}$、$\alpha_{\min}$、$T_w$ 和 $T_c$，并按照上面定义的 scheduler 返回学习率 $\alpha_t$。然后实现 [adapters.get_lr_cosine_schedule]，并确保它通过以下测试：

uv run pytest -k test_get_lr_cosine_schedule

4.5 Gradient clipping

训练过程中，有时会遇到产生大梯度的训练样本，这会使训练不稳定。为了缓解这个问题，实践中常用的一种技术是 gradient clipping。其思想是在每次 backward pass 之后、执行 optimizer step 之前，对梯度的 norm 施加限制。

给定梯度（针对所有 parameters）$g$，我们计算它的 $\ell_2$-norm $\lVert g \rVert_2$。如果该 norm 小于最大值 $M$，则保持 $g$ 不变；否则，将 $g$ 按因子 $\frac{M}{\lVert g \rVert_2 + \varepsilon}$ 缩小（其中加入一个很小的 $\varepsilon$，例如 $10^{-6}$，用于数值稳定性）。注意，得到的 norm 将略小于 $M$。

Problem (gradient_clipping)：实现 gradient clipping（1 分）

编写一个实现 gradient clipping 的函数。你的函数应接收一个 parameter 列表和一个最大 $\ell_2$-norm。它应就地修改每个 parameter 的梯度。使用 $\varepsilon = 10^{-6}$（PyTorch 默认值）。然后实现 adapter [adapters.run_gradient_clipping]，并确保它通过以下测试：

uv run pytest -k test_gradient_clipping

5 Training loop

现在我们将最终把目前为止构建的主要组件组合起来：tokenized data、model 和 optimizer。

5.1 Data Loader

tokenized data（例如你在 tokenizer_experiments 中准备好的数据）是一个单独的 token 序列 $x = (x_1, \dots, x_n)$。尽管源数据可能由多个独立文档组成（例如不同的网页或源代码文件），一种常见做法是将它们全部拼接成一个 token 序列，并在它们之间加入分隔符（例如 <|endoftext|> token）。

data loader 会将其转换为 batch 流，其中每个 batch 由 $B$ 个长度为 $m$ 的序列组成，并配上对应的 next tokens（长度也为 $m$）。例如，当 $B=1, m=3$ 时，$([x_2, x_3, x_4], [x_3, x_4, x_5])$ 可以是一个可能的 batch。

以这种方式加载数据会因多种原因简化训练。首先，任何 $1 \le i \le n-m$ 都给出一个有效的训练序列，因此采样训练序列很简单。由于所有训练序列长度相同，不需要对输入序列做 padding，这提高了硬件利用率（也通过增大 batch size $B$ 来提高）。最后，我们也不需要加载完整数据集来采样训练数据，这使得处理那些原本可能无法放入内存的大型数据集变得容易。

Problem (data_loading)：实现 data loading（2 分）

Deliverable: 编写一个函数，接收一个 numpy array $x$（包含 token IDs 的整数数组）、一个 batch_size、一个 context_length 和一个 PyTorch device string（例如 'cpu' 或 'cuda:0'），并返回一对 tensors：采样得到的输入序列以及对应的 next-token targets。两者（Both）

第 35 页

tensors 的形状都应为 (batch_size, context_length)，包含 token IDs，并且两者都应放在请求的 device 上。为了用我们提供的测试来测试你的实现，你首先需要实现测试 adapter [adapters.run_get_batch]。然后运行以下命令来测试你的实现：

uv run pytest -k test_get_batch

Low-Resource Tip：在 CPU 或 Apple Silicon 上加载数据

如果你计划在 CPU 或 Apple Silicon 上训练你的 LM，你需要将数据移动到正确的 device（同样，之后你也应该为模型使用相同的 device）。

如果你在 CPU 上，可以使用 'cpu' device string；在 Apple Silicon（M* 芯片）上，可以使用 'mps' device string。

有关 MPS 的更多信息，请查看以下资源：

如果数据集太大，无法加载到内存中怎么办？我们可以使用一个名为 mmap 的 Unix system call，它将磁盘上的文件映射到虚拟内存，并在访问该内存位置时惰性加载文件内容。因此，你可以“假装”整个数据集都在内存中。Numpy 通过 np.memmap（或者，如果你最初用 np.save 保存数组，则使用 np.load 的 mmap_mode='r' flag）实现这一点，它会返回一个类似 numpy array 的对象，在你访问条目时按需加载它们。在训练期间从数据集（即 numpy array）中采样时，请确保以 memory-mapped 模式加载数据集（通过 np.memmap，或根据你保存数组的方式使用 np.load 的 mmap_mode='r' flag）。还要确保指定一个与你正在加载的数组匹配的 dtype。显式验证 memory-mapped 数据看起来正确（例如，不包含超出预期 vocabulary size 的值）可能会有所帮助。

5.2 Checkpointing

除了加载数据之外，我们还需要在训练过程中保存模型。在运行 job 时，我们经常希望能够恢复中途停止的训练运行（例如，由于 job 超时、机器故障等）。即使一切顺利，我们之后也可能希望访问 intermediate models（例如，事后研究 training dynamics、从训练不同阶段的模型取样本等）。

一个 checkpoint 应包含恢复训练所需的所有 state。我们当然至少希望能够恢复 model weights。如果使用 stateful optimizer（例如 AdamW），我们还需要保存 optimizer 的 state（例如，对于 AdamW，是 moment estimates）。最后，为了恢复 learning rate schedule，我们需要知道停止时的 iteration number。PyTorch 让保存所有这些变得很容易：每个 nn.Module 都有一个 state_dict() 方法，返回一个包含所有可学习 weights 的 dictionary；我们之后可以用配套方法 load_state_dict() 恢复这些 weights。任何 torch.optim.Optimizer 也是如此。最后，torch.save(obj, dest) 可以将一个对象（例如，一个 dictionary，其中一些值是 tensors，但也有整数等普通 Python 对象）转储到文件（path）或 file-like object，之后可以用 torch.load(src) 将其加载回内存。

第 36 页

Problem (checkpointing): Implement model checkpointing（1 分）

实现下面两个函数，用来加载和保存 checkpoint：

def save_checkpoint(model, optimizer, iteration, out)

它应当把 model、optimizer 和 iteration 的全部状态 dump 到 file-like object out 里。你可以用 model 和 optimizer 各自的 state_dict 方法取出它们相关的状态，并用 torch.save(obj, out) 把 obj dump 到 out（PyTorch 在这里既支持路径，也支持 file-like object）。一个典型的选择是让 obj 是一个 dictionary，但你可以用任意格式，只要之后能把你这个 checkpoint 加载回来即可。

该函数期望以下参数：

def load_checkpoint(src, model, optimizer)

它应当从 src（路径或 file-like object）加载 checkpoint，然后从该 checkpoint 中恢复 model 和 optimizer 的状态。你的函数应当返回保存到该 checkpoint 里的 iteration 编号。你可以用 torch.load(src) 取回你在 save_checkpoint 实现里保存的内容，并用 model 和 optimizer 各自的 load_state_dict 方法把它们恢复到之前的状态。

该函数期望以下参数：

实现 [adapters.run_save_checkpoint] 和 [adapters.run_load_checkpoint] 两个 adapter，并确保它们能通过

uv run pytest -k test_checkpointing

5.3 Training loop

现在，终于到了把你实现的所有组件拼进主训练脚本的时候了。让“用不同的超参数启动训练 run”变得容易（例如把它们作为 command-line 参数传入）是会回本的，因为之后你会反复做很多次，来研究不同选择如何影响训练。

Problem (training_together): Put it together（4 分）

Deliverable: 写一个脚本，运行 training loop，在用户提供的输入上训练你的 model。具体来说，我们建议你的训练脚本（至少）支持以下各项：

第 37 页

6 Generating text

既然我们已经能训练模型了，最后还需要的一块能力是从模型中生成文本。回忆一下，一个 language model 接收一个长度为 sequence_length 的整数序列（可能带 batch），并产生一个大小为 (sequence_length, vocab_size) 的矩阵，其中序列的每个元素都是一个概率分布，预测该位置之后的下一个 token。现在我们要写几个函数，把它变成对新序列的 sampling 方案。

Softmax

按标准惯例，language model 的输出是最后一层 linear layer 的输出（即 “logits”），所以我们必须通过 softmax 操作把它变成归一化的概率，也就是我们之前在 Equation 10 中看到的那个。

Decoding

为了从我们的模型生成文本（decode），我们会给模型一个前缀 token 序列（即 “prompt”），并要求它产生一个在整个 vocabulary 上的概率分布，用来预测序列中的下一个 token。然后，我们从这个 vocabulary 各项上的分布中采样，来确定下一个输出的 token。

具体来说，decoding 过程的一步应当接收一个序列 $x_{1\ldots t}$，并通过下面的式子返回一个 token $x_{t+1}$：

$$P(x_{t+1}=i \mid x_{1\ldots t}) = \frac{\exp(v_i)}{\sum_j \exp(v_j)} \quad (21)$$

$$v = \text{TransformerLM}(x_{1\ldots t})_t \in \mathbb{R}^{\text{vocab\_size}} \quad (22)$$

其中 TransformerLM 是我们的模型，它接收长度为 sequence_length 的序列作为输入，产生大小为 (sequence_length, vocab_size) 的矩阵；我们取这个矩阵的最后一个元素，因为我们要找的是第 $t$ 个位置上对下一个 token 的预测。

这样我们就得到一个基础的 decoder：反复从这些 one-step conditional 中采样（把上一次生成的输出 token 追加到下一个 decoding timestep 的输入里），直到生成 end-of-sequence token <|endoftext|>（或者达到用户指定的最大生成 token 数）。

Decoder tricks

我们将要实验的是小模型，小模型有时会生成质量很低的文本。两个简单的 decoder trick 可以帮助解决这些问题。第一个是 temperature scaling：我们用一个温度参数 $\tau$ 修改 softmax，新的 softmax 为

$$\text{softmax}(v,\tau)_i = \frac{\exp(v_i/\tau)}{\sum_{j=1}^{\text{vocab\_size}} \exp(v_j/\tau)} \quad (23)$$

注意，令 $\tau \to 0$ 会使得 $v$ 中最大的元素占主导，softmax 的输出变成一个集中在该最大元素上的 one-hot 向量。

[注 9] wandb.ai

第 38 页

第二个 trick 是 nucleus sampling，或称 top-$p$ sampling：我们通过截断低概率 token 来修改采样分布。设 $q$ 是我们从（经温度缩放的）softmax 得到的、大小为 vocab_size 的概率分布。超参数为 $p$ 的 nucleus sampling 按下面的式子产生下一个 token：

$$P(x_{t+1}=i \mid q) = \begin{cases} \frac{q_i}{\sum_{j \in V(p)} q_j} & \text{if } i \in V(p) \\ 0 & \text{otherwise} \end{cases} \quad (24)$$

其中 $V(p)$ 是满足 $\sum_{j \in V(p)} q_j \ge p$ 的最小下标集合。你可以先按大小对概率分布 $q$ 排序，然后选取最大的那些 vocabulary 元素直到达到目标水平 $p$，从而很容易地算出这个量。

Problem (decoding): Decoding（3 分）

Deliverable: 实现一个从你的 language model 做 decode 的函数。我们建议你支持以下特性：

7 Experiments

现在是把所有东西拼起来、在预训练数据集上训练（小）language model 的时候了。

7.1 How to Run Experiments and Deliverables

理解 Transformer 各架构组件背后道理的最佳方式，是亲手修改它并运行它。动手经验是无可替代的。

为此，能够快速、一致地做实验，并记录下你做过什么，就很重要。为了快速实验，我们会在一个小规模模型（约 17M 总参数）和简单数据集（TinyStories）上跑很多实验。为了一致地做事，你要系统地 ablate 组件、改变超参数；为了保留记录，我们会要求你提交一份实验日志，以及每个实验对应的 learning curve。

为了能提交 loss curve，务必定期评估 validation loss，并同时记录步数和 wall-clock time。你可能会觉得 Weights and Biases 之类的 logging 基础设施很有用。

Problem (experiment_log): Experiment logging（3 分）

针对你的训练和评估代码，搭建 experiment tracking 基础设施，使你能追踪实验，以及 loss curve 随 gradient step 和 wall-clock time 的变化。

第 39 页

Deliverable（交付物）： 用于你实验的 logging infrastructure 代码，以及一份实验日志（记录你尝试过的所有内容的文档），对应本节下面的作业问题。

7.2 TinyStories

我们将从一个非常简单的数据集（TinyStories；R. Eldan et al. [1]）开始，模型会在其上训练得很快，我们也能看到一些有趣的行为。获取该数据集的说明见 Section 1。该数据集的样子示例如下。

Example (tinystories_example)： TinyStories 中的一个例子

从前有一个小男孩叫 Ben。Ben 喜欢探索周围的世界。他见过许多奇妙的东西，比如商店里陈列的漂亮花瓶。一天，Ben 走过商店时，遇到一个非常特别的花瓶。Ben 看到它时惊叹不已！他说：“哇，那真是一个了不起的花瓶！我可以买下它吗？”店主微笑着说：“当然可以。你可以把它带回家，向所有朋友展示它有多奇妙！”于是 Ben 把花瓶带回家，他非常自豪！他把朋友们叫来，给他们看这个奇妙的花瓶。他的朋友们都觉得花瓶很漂亮，简直不敢相信 Ben 有多幸运。就这样，Ben 在商店里发现了一个奇妙的花瓶！

7.2.1 超参数调优

我们会先告诉你一些非常基础的超参数作为起点，并要求你为其他超参数找到效果良好的设置。

你应该做一些试错，为以下其他超参数找到良好的默认值：learning rate、learning rate warmup、其他 AdamW 超参数（$\beta_1$、$\beta_2$、$\varepsilon$）以及 weight decay。你可以在 D. P. Kingma et al. [22] 中找到这类超参数的一些典型选择。

7.2.2 整合起来

现在你可以把一切整合起来：得到一个训练好的 BPE tokenizer，对训练数据集进行 tokenize，并在你写的 training loop 中运行。重要提示： 如果你的实现正确且高效，上述超参数应能在 1 块 B200

第 40 页

GPU 上带来大约 20–30 分钟的运行时间。如果你的运行时间远长于此，请检查并确保你的 dataloading、checkpointing 或 validation loss 代码没有成为运行时间的瓶颈，并且你的实现被正确地 batch 化。

7.2.3 调试模型架构的技巧与窍门

我们强烈建议你熟悉 IDE 内置的 debugger（例如 VSCode/Zed），与用 print 语句调试相比，这会节省你的时间。如果你使用文本编辑器，可以使用类似 ipdb 的工具。调试模型架构时，另外一些良好实践是：

Problem (learning_rate)：调优 learning rate（2 B200 hrs；3 points）

learning rate 是最需要调优的超参数之一。以你训练好的 base model 为基础，回答以下问题：

(a) 对 learning rate 进行超参数 sweep，并报告最终 loss（如果 optimizer 发散，则注明 divergence）。

Deliverable（交付物）： 与多个 learning rate 对应的 learning curves。解释你的超参数搜索策略。

Deliverable（交付物）： 一个在 TinyStories 上 validation loss（per-token）至多为 1.45 的模型。

Low-Resource Tip（低资源提示）：Train for a few steps on CPU or Apple Silicon

如果你在 cpu 或 mps 上运行，则应将处理的 token 总数减少到 40,000,000，这足以产生相当流畅的文本。你也可以把目标 validation loss 从 1.45 提高到 2.00。

我们在 M4 Max 芯片和 36 GB RAM 上使用调优后的 learning rate 运行我们的 solution code，使用 batch size × total step count × context length = $32 \times 5000 \times 256 = 40{,}960{,}000$ tokens，在 cpu 上耗时 1 小时 22 分钟，在 mps 上耗时 36 分钟。在第 5000 步时，我们达到 1.80 的 validation loss。

Some additional tips（一些额外提示）：

torch.set_float32_matmul_precision('high')

就像你可能在 cuda 设备上做的那样。我们在 mps 上尝试启用 TF32 kernels（torch version 2.9.0），发现后端有时会使用静默损坏的 kernels，导致训练不稳定。

第 41 页

截至 torch version 2.9.0，mps 上不支持使用 Inductor 进行编译。

(b) 民间经验法则是，最佳 learning rate 位于“稳定性的边缘”。研究 learning rate 开始发散的临界点与你的最佳 learning rate 之间有何关系。

Deliverable（交付物）： 随着 learning rate 增大而得到的 learning curves，其中至少包含一次发散运行，并分析这与收敛速度有何关系。

现在让我们改变 batch size，看看训练会发生什么。Batch size 很重要——它们让我们通过做更大的 matrix multiply 来提高 GPU 效率，但我们是否总是希望 batch size 很大？让我们跑一些实验来找出答案。

Problem (batch_size_experiment)：Batch size 变化（1 B200 hr；1 point）

把你的 batch size 从 1 一直变化到 GPU 内存上限。中间至少尝试几个 batch size，包括 64 和 128 这样的典型大小。

Deliverable（交付物）： 不同 batch size 运行的 learning curves。如有必要，应重新优化 learning rate。

Deliverable（交付物）： 用几句话讨论你对 batch size 及其对训练影响的发现。

有了你的 decoder，我们现在可以生成文本了！我们将从模型生成，并看看效果有多好。作为参考，你得到的输出至少应看起来和下面例子一样好。

Example (ts_generate_example)： 来自 TinyStories 语言模型的样本输出

从前有一个漂亮的女孩叫 Lily。她喜欢吃口香糖，尤其是那种又大又黑的。一天，Lily 的妈妈让她帮忙做晚饭。Lily 非常兴奋！她喜欢帮妈妈。Lily 的妈妈做了一大锅汤当晚饭。Lily 非常开心，说：“谢谢你，妈妈！我爱你。”她帮妈妈把汤倒进一个大碗里。晚饭后，Lily 的妈妈又做了一些美味的汤。Lily 很喜欢！她说：“谢谢你，妈妈！这汤太好吃了！”她妈妈微笑着说：“我很高兴你喜欢，Lily。”她们做完饭，又继续一起做饭。故事结束。

Low-Resource Tip（低资源提示）：Generate text on CPU or Apple Silicon

第 42 页

如果你改用处理了 40M tokens 的低资源（low-resource）配置，你应该会看到生成结果仍然像英语，但不如上面那样流畅。例如，下面是我们用 TinyStories 语言模型（只在这 40M tokens 上训练）得到的样例输出：

Once upon a time, there was a little girl named Sue. Sue had a tooth that she loved very much. It
was his best head. One day, Sue went for a walk and met a ladybug! They became good friends
and played on the path together.
“Hey, Polly! Let’s go out!” said Tim. Sue looked at the sky and saw that it was difficult to find a
way to dance shining. She smiled and agreed to help the talking!“
As Sue watched the sky moved, what it was. She

下面是确切的题目陈述，以及我们的要求：

Problem (generate)：生成文本（1 point）

使用你的 decoder 和训练好的 checkpoint，报告你的模型生成的文本。你可能需要调整 decoder 的参数（temperature、top-p 等）才能得到流畅的输出。

Deliverable: 至少 256 tokens 的文本 dump（或者到第一个 <|endoftext|> token 为止），以及一段关于这段输出流畅度的简短评论，并至少给出两个会影响这段输出好或坏的因素。

7.3 Ablations and architecture modification

理解 Transformer 最好的方式，就是真的去改动它，然后看它的行为如何变化。下面我们做几个简单的消融（ablation）和修改。

Ablation 1: layer normalization

人们常说 layer normalization 对 Transformer 训练的稳定性很重要。但也许我们就想冒点险。让我们把每个 Transformer block 里的 RMSNorm 都去掉，看看会发生什么。

Problem (layer_norm_ablation)：移除 RMSNorm 并训练（0.5 B200 hrs）（1 point）

把 Transformer 里所有的 RMSNorm 都移除并训练。在之前的最优学习率下会发生什么？你能通过使用更低的学习率来获得稳定性吗？

Deliverable: 移除 RMSNorm 并训练时的学习曲线，以及在最佳学习率下的学习曲线。

Deliverable: 几句关于 RMSNorm 影响的评论。

下面我们考察另一种乍看之下很随意的 layer normalization 选择。pre-norm 的 Transformer block 定义为

$$z = x + \mathrm{MultiHeadSelfAttention}(\mathrm{RMSNorm}(x)) \tag{25}$$

$$y = z + \mathrm{FFN}(\mathrm{RMSNorm}(z)). \tag{26}$$

第 43 页

这是对原始 Transformer 架构少数几个「共识」修改之一——原始架构用的是 post-norm 方式，即

$$z = \mathrm{RMSNorm}(x + \mathrm{MultiHeadSelfAttention}(x)) \tag{27}$$

$$y = \mathrm{RMSNorm}(z + \mathrm{FFN}(z)). \tag{28}$$

让我们改回 post-norm 方式，看看会发生什么。

Problem (pre_norm_ablation)：实现 post-norm 并训练（0.5 B200 hrs）（1 point）

把你的 pre-norm Transformer 实现改成 post-norm 的。用 post-norm 模型训练，看看会发生什么。

Deliverable: post-norm Transformer 的学习曲线，并与 pre-norm 的作对比。

我们看到 layer normalization 对 Transformer 的行为有重大影响，甚至 layer normalization 所处的位置也很重要。

Ablation 2: position embeddings

接下来我们考察 position embeddings 对模型性能的影响。具体来说，我们会把基础模型（用 RoPE）与完全不含 position embeddings 的模型（NoPE）作对比。结果表明，decoder-only transformer，也就是我们实现的这种带 causal mask 的模型，在理论上可以在不显式给定 position embeddings 的情况下推断出相对或绝对位置信息 [Y.-H. H. Tsai et al., 2019; A. Kazemnejad et al., 2023]。下面我们实证检验 NoPE 相比 RoPE 表现如何。

Problem (no_pos_emb)：实现 NoPE（0.5 B200 hrs）（1 point）

把你带 RoPE 的 Transformer 实现改成完全移除 position embedding 信息，看看会发生什么。

Deliverable: 比较 RoPE 与 NoPE 性能的学习曲线。

Ablation 3: SwiGLU vs. SiLU

接下来，我们沿用 N. Shazeer [20] 的做法，通过比较 SwiGLU feed-forward 网络与使用 SiLU 激活但不含 gated linear unit（GLU）的 feed-forward 网络，来检验 feed-forward 网络中 gating 的重要性：

$$\mathrm{FFN}_{\mathrm{SiLU}}(x) = W_2\,\mathrm{SiLU}(W_1 x). \tag{29}$$

回想一下，在我们的 SwiGLU 实现中，我们把内层 feed-forward 层的维度设成大约 $d_{\mathrm{ff}} = \frac{8}{3} d_{\mathrm{model}}$（同时保证 $d_{\mathrm{ff}} \bmod 64 = 0$，以便用上 GPU tensor core）。在这个消融基线里，你的 $\mathrm{FFN}_{\mathrm{SiLU}}$ 实现应当改为设 $d_{\mathrm{ff}} = 4 \times d_{\mathrm{model}}$，以大致匹配默认 SwiGLU feed-forward 网络的参数量（后者有三个权重矩阵，而不是两个）。

Problem (swiglu_ablation)：SwiGLU vs. SiLU（0.5 B200 hrs）（1 point）

第 44 页

Deliverable: SwiGLU 与 SiLU feed-forward 网络的性能对比学习曲线，两者参数量大致匹配。

Deliverable: 几句讨论你的发现的文字。

Low-Resource Tip: Online students with limited GPU resources should test modifications on TinyStories

在作业的剩余部分，我们将转向一个更大规模、噪声更多的网络数据集（OpenWebText），尝试各种架构修改，并（可选地）向课程 leaderboard 提交结果。

在 OpenWebText 上把 LM 训练到流畅需要很长时间，所以我们建议 GPU 资源有限的线上学生继续在 TinyStories 上测试修改（用验证 loss 作为评估性能的指标）。

7.4 Running on OpenWebText

现在我们要转向一个更标准的预训练数据集，它由网络爬取构建。OpenWebText [A. Gokaslan et al., 2019] 的一个小样本也以单个文本文件的形式提供：如何获取该文件见 Section 1。

下面是 OpenWebText 中的一个例子。注意这段文本要真实、复杂、多样得多。你也许想浏览一下训练数据集，感受一下网络爬取语料库的训练数据是什么样子。

Example (owt_example)：OWT 中的一个例子

Baseball Prospectus director of technology Harry Pavlidis took a risk when he hired Jonathan Judge.

Pavlidis knew that, as Alan Schwarz wrote in The Numbers Game, “no corner of American culture is more precisely counted, more passionately quantified, than performances of baseball players.” With a few clicks here and there, you can find out that Noah Syndergaard’s fastball revolves more than 2,100 times per minute on its way to the plate, that Nelson Cruz had the game’s highest average exit velocity among qualified hitters in 2016 and myriad other tidbits that seem ripped from a video game or science fiction novel. The rising ocean of data has empowered an increasingly important actor in baseball’s culture: the analytical hobbyist.

That empowerment comes with added scrutiny – on the measurements, but also on the people and publications behind them. With Baseball Prospectus, Pavlidis knew all about the backlash that accompanies quantitative imperfection. He also knew the site’s catching metrics needed to be reworked, and that it would take a learned mind – someone who could tackle complex statistical modeling problems – to complete the job.

“He freaks us out.” Harry Pavlidis

Pavlidis had a hunch that Judge “got it” based on the latter’s writing and their interaction at a site-sponsored ballpark event. […]

第 45 页

Note：你可能需要为这个实验重新调一下超参数，比如 learning rate 或 batch size。

Problem (main_experiment)：在 OWT 上做实验（2 B200 小时）（2 分）

在与 TinyStories 相同的模型架构和总训练迭代数下，在 OpenWebText 上训练你的 language model。这个模型表现如何？

7.5 你自己的修改 + leaderboard

恭喜你走到这一步，你快要完成了！现在你要尝试改进 Transformer 架构，并看看你的超参数和架构在全班同学中能排到什么位置。

Leaderboard 规则

除了下面几条，没有其他限制：

除此之外，你想做什么都可以。

如果你想找点实现思路，可以看看这些资源：

在跑完整的 45 分钟之前，你大概会想先在 OpenWebText 的一个小子集或 TinyStories 上测试这些修改。

需要提醒的是，你在这个 leaderboard 上发现效果不错的一些修改，未必能推广到更大规模的预训练。我们会在课程的 scaling laws 单元进一步探讨这个想法。

Problem (leaderboard)：Leaderboard（10 B200 小时）（6 分）

第 46 页

你将在上述 leaderboard 规则下训练一个模型，目标是在 0.75 B200 小时内最小化你的 language model 的 validation loss。

Deliverable：最终记录的 validation loss、一条能清楚看到横轴 wall-clock time 小于 45 分钟的 learning curve，以及对你做了什么的描述。我们希望 leaderboard 提交至少能超过 5.0 loss 这个朴素基线。在这里提交到 leaderboard：github.com/stanford-cs336/assignment1-basics-leaderboard。

Bibliography（参考文献）

第 47 页

材料

数据 / 仓库（只给链接，不下载 ✗）

按 D50：数据集与 starter code 只给链接 ✓（starter code 用 git clone 拿 ✓，数据集太大 ✗ 而且年年变 ✓）。这些链接故意不放在「材料」小节 ✓ —— 放了就会被 --fetch 下载 ✗。

自查