cs 336 第一章作业

原文与逐句翻译
CS336 Assignment 1 (basics): Building a Transformer LM Version 1.0.4
CS336 作业1（基础）：从零构建 Transformer 语言模型 版本1.0.4
CS336 Staff Spring 2025
CS336 教学团队，2025年春季
1 Assignment Overview
1 作业概述
In this assignment, you will build all the components needed to train a standard Transformer language model (LM) from scratch and train some models.
在本次作业中，你将从零开始构建训练一个标准 Transformer 语言模型所需的全部组件，并实际训练一些模型。
What you will implement
你将实现以下内容：
Byte-pair encoding (BPE) tokenizer (§2)
字节对编码（BPE）分词器（第2节）
Transformer language model (LM) (§3)
Transformer 语言模型（LM）（第3节）
The cross-entropy loss function and the AdamW optimizer (§4)
交叉熵损失函数与 AdamW 优化器（第4节）
The training loop, with support for serializing and loading model and optimizer state (§5)
训练循环，支持序列化与加载模型及优化器状态（第5节）
What you will run
你将运行以下任务：
Train a BPE tokenizer on the TinyStories dataset.
在 TinyStories 数据集上训练一个 BPE 分词器。
Run your trained tokenizer on the dataset to convert it into a sequence of integer IDs.
使用训练好的分词器处理数据集，将其转换为整数 ID 序列。
Train a Transformer LM on the TinyStories dataset.
在 TinyStories 数据集上训练一个 Transformer 语言模型。
Generate samples and evaluate perplexity using the trained Transformer LM.
使用训练好的 Transformer 语言模型生成样本，并计算困惑度。
Train models on OpenWebText and submit your attained perplexities to a leaderboard.
在 OpenWebText 上训练模型，并将你获得的困惑度提交到排行榜。
What you can use
你可以使用的工具：
We expect you to build these components from scratch.
我们要求你从零开始构建这些组件。
In particular, you may not use any definitions from torch.nn, torch.nn.functional, or torch.optim except for the following:
特别地，你不得使用 torch.nn、torch.nn.functional 或 torch.optim 中的任何定义，以下情况除外：
• torch.nn.Parameter
torch.nn.Parameter（参数类）
• Container classes in torch.nn (e.g., Module, ModuleList, Sequential, etc.)
torch.nn 中的容器类（例如 Module、ModuleList、Sequential 等）
• The torch.optim.Optimizer base class
torch.optim.Optimizer 基类
You may use any other PyTorch definitions.
你可以使用任何其他 PyTorch 定义。
If you would like to use a function or class and are not sure whether it is permitted, feel free to ask on Slack.
如果你想使用某个函数或类，但不确定是否允许，请随时在 Slack 上提问。
When in doubt, consider if using it compromises the “from-scratch” ethos of the assignment.
如有疑问，请考虑使用该函数是否违背本作业“从零开始”的宗旨。
Statement on AI tools
关于 AI 工具的声明：
Prompting LLMs such as ChatGPT is permitted for low-level programming questions or high-level conceptual questions about language models, but using it directly to solve the problem is prohibited.
允许使用 ChatGPT 等 LLM 回答低级编程问题或关于语言模型的高层次概念性问题，但禁止使用它们直接解决问题。
We strongly encourage you to disable AI autocomplete (e.g., Cursor Tab, GitHub CoPilot) in your IDE when completing assignments (though non-AI autocomplete, e.g., autocompleting function names is totally fine).
我们强烈建议你在完成作业时禁用 IDE 中的 AI 自动补全功能（例如 Cursor Tab、GitHub Copilot），但非 AI 的自动补全（如补全函数名）是完全允许的。
We have found that AI autocomplete makes it much harder to engage deeply with the content.
我们发现 AI 自动补全会显著降低你对内容的深入理解。
What the code looks like
代码结构如下：
All the assignment code as well as this writeup are available on GitHub at:
所有作业代码与本说明文档均可在 GitHub 获取：
github.com/stanford-cs336/assignment1-basics
github.com/stanford-cs336/assignment1-basics
Please git clone the repository.
请使用 git clone 克隆该仓库。
If there are any updates, we will notify you so you can git pull to get the latest.
如有更新，我们会通知你，届时请使用 git pull 获取最新内容。
cs336_basics/*: This is where you write your code.
cs336_basics/*：这是你编写代码的目录。
Note that there’s no code in here—you can do whatever you want from scratch!
注意：该目录初始为空，你可以完全从零开始实现。
adapters.py: There is a set of functionality that your code must have.
adapters.py：你的代码必须实现一组特定功能。
For each piece of functionality (e.g., scaled dot product attention), fill out its implementation (e.g., run_scaled_dot_product_attention) by simply invoking your code.
对于每个功能（例如缩放点积注意力），请通过调用你自己的代码来完成其实现（例如 run_scaled_dot_product_attention）。
Note: your changes to adapters.py should not contain any substantive logic; this is glue code.
注意：你对 adapters.py 的修改不应包含任何实质性逻辑，它只是胶水代码。
test_*.py: This contains all the tests that you must pass (e.g., test_scaled_dot_product_attention), which will invoke the hooks defined in adapters.py.
test_*.py：包含你必须通过的所有测试（例如 test_scaled_dot_product_attention），这些测试会调用 adapters.py 中定义的钩子函数。
Don’t edit the test files.
不要修改测试文件。
How to submit
如何提交：
You will submit the following files to Gradescope:
你需要向 Gradescope 提交以下文件：
• writeup.pdf: Answer all the written questions. Please typeset your responses.
writeup.pdf：回答所有书面问题，请使用排版工具（如 LaTeX）撰写。
• code.zip: Contains all the code you’ve written.
code.zip：包含你编写的所有代码。
To submit to the leaderboard, submit a PR to:
要将结果提交到排行榜，请向以下仓库提交 PR：
github.com/stanford-cs336/assignment1-basics-leaderboard
github.com/stanford-cs336/assignment1-basics-leaderboard
See the README.md in the leaderboard repository for detailed submission instructions.
请参考排行榜仓库中的 README.md 获取详细的提交说明。
Where to get datasets
如何获取数据集：
This assignment will use two pre-processed datasets: TinyStories [Eldan and Li, 2023] and OpenWebText [Gokaslan et al., 2019].
本作业将使用两个预处理过的数据集：TinyStories（Eldan 和 Li, 2023）与 OpenWebText（Gokaslan 等, 2019）。
Both datasets are single, large plaintext files.
这两个数据集均为单个大型纯文本文件。
If you are doing the assignment with the class, you can find these files at /data of any non-head node machine.
若你随课程一起完成作业，可在任意非头节点的 /data 目录下找到这些文件。
If you are following along at home, you can download these files with the commands inside the README.md.
若你在家中自行完成，可按照 README.md 中的命令下载这些文件。
Low-Resource/Downscaling Tip: Init
低资源/降规模提示：初始化
Throughout the course’s assignment handouts, we will give advice for working through parts of the assignment with fewer or no GPU resources.
在课程的所有作业说明中，我们将提供如何在没有或较少 GPU 资源下完成作业的建议。
For example, we will sometimes suggest downscaling your dataset or model size, or explain how to run training code on a MacOS integrated GPU or CPU.
例如，我们有时会建议缩小数据集或模型规模，或说明如何在 macOS 集成 GPU 或 CPU 上运行训练代码。
You’ll find these “low-resource tips” in a blue box (like this one).
你会在蓝色方框（如本框）中找到这些“低资源提示”。
Even if you are an enrolled Stanford student with access to the course machines, these tips may help you iterate faster and save time, so we recommend you to read them!
即使你是在读的斯坦福学生、可使用课程服务器，这些提示也可能帮助你更快迭代并节省时间，因此我们建议你阅读它们！
Low-Resource/Downscaling Tip: Assignment 1 on Apple Silicon or CPU
低资源/降规模提示：在 Apple Silicon 或 CPU 上完成作业1
With the staff solution code, we can train an LM to generate reasonably fluent text on an Apple M3 Max chip with 36 GB RAM, in under 5 minutes on Metal GPU (MPS) and about 30 minutes using the CPU.
使用助教提供的参考代码，我们可在配有 36GB RAM 的 Apple M3 Max 芯片上，于 Metal GPU（MPS）下 5 分钟内、于 CPU 下约 30 分钟内，训练出能生成较流畅文本的语言模型。
If these words don’t mean much to you, don’t worry!
如果这些术语对你而言不太熟悉，不必担心！
Just know that if you have a reasonably up-to-date laptop and your implementation is correct and efficient, you will be able to train a small LM that generates simple children’s stories with decent fluency.
只需知道：如果你有一台较新的笔记本电脑，并且你的实现正确且高效，你就能训练出一个小型语言模型，它能用不错的流畅度生成简单的儿童故事。
Later in the assignment, we will explain what changes to make if you are on CPU or MPS.
在作业后文，我们将说明如果你使用 CPU 或 MPS，需要做出哪些调整。

2 字节对编码（BPE）分词器
在这一部分中，我们将训练并实现一个基于字节的字节对编码（BPE）分词器（Sennrich 等人，2016；Wang 等人，2019）。特别地，我们将任意 Unicode 字符串表示为字节序列，并在此字节序列上训练 BPE 分词器。稍后，我们将使用此分词器将文本（字符串）编码为整数 ID 序列，以便进行语言建模。
2.1 Unicode 标准
Unicode 是一种文本编码标准，将字符映射为整数码位。截至 2024 年 9 月发布的 Unicode 16.0，标准共定义了 154 998 个字符，涵盖 168 种文字。例如，字符 “s” 的码位是 115（通常写作 U+0073，其中 U+ 是约定前缀，0073 是十六进制的 115），字符 “牛” 的码位是 29275。在 Python 中，可以使用 ord() 函数将单个 Unicode 字符转换为对应的整数表示，使用 chr() 函数将整数码位转换为对应字符。
复制
ord('牛')
29275
chr(29275)
'牛'
问题 (unicode1)：理解 Unicode（1 分）
(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")
回答：一句话即可。
2.2 Unicode 编码
尽管 Unicode 标准定义了字符到码位（整数）的映射，但直接在 Unicode 码位上训练分词器并不实际，因为词汇表会过于庞大（约 150 K 项）且稀疏（许多字符非常罕见）。相反，我们将使用 Unicode 编码，它将一个 Unicode 字符转换为一个字节序列。Unicode 标准本身定义了三种编码：UTF-8、UTF-16 和 UTF-32，其中 UTF-8 是互联网的主导编码（超过 98% 的网页使用）。
在 Python 中，我们可以使用 encode() 函数将 Unicode 字符串编码为 UTF-8；要访问 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'>
获取编码字符串的字节值（0 到 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]
一个字节并不一定对应一个 Unicode 字符！
复制
print(len(test_string))
13
print(len(utf8_encoded))
23
print(utf8_encoded.decode("utf-8"))
hello! こんにちは !
通过将 Unicode 码位转换为字节序列（例如通过 UTF-8 编码），我们实际上将一组码位（范围 0 到 154 997 的整数）转换为字节值序列（范围 0 到 255 的整数）。长度为 256 的字节词汇表更易于处理。使用字节级分词时，我们无需担心未登录词，因为任何输入文本都可以表示为 0 到 255 之间的整数序列。
问题 (unicode2)：Unicode 编码（3 分）
(a) 与 UTF-16 或 UTF-32 相比，基于 UTF-8 字节训练分词器有哪些优势？建议比较这些编码对各种输入字符串的输出。
回答：一到两句话即可。
(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'
回答：给出一个输入字节串示例，使得 decode_utf8_bytes_to_str_wrong 产生错误输出，并用一句话解释该函数为何错误。
(c) 给出一个两字节的序列，该序列无法解码为任何 Unicode 字符。
回答：给出一个示例，并用一句话解释。
2.3 子词分词
虽然字节级分词可以缓解词级分词器面临的未登录词问题，但将文本分词为字节会导致输入序列极长。这会拖慢模型训练，因为一个 10 词的句子在词级语言模型中可能只有 10 个 token，但在字符级模型中可能长达 50 个或更多 token（取决于词长）。处理这些更长的序列需要在模型的每一步进行更多计算。此外，由于输入序列更长，在字节序列上进行语言建模更加困难，因为数据中存在长期依赖关系。
子词分词是词级分词器与字节级分词器之间的折中。注意，字节级分词器的词汇表有 256 个条目（字节值为 0 到 255）。子词分词器以更大的词汇表为代价，换取对输入字节序列更好的压缩。例如，如果字节序列 b'the' 在原始文本训练数据中频繁出现，那么为其分配一个词汇表条目可以将这个 3-token 序列缩减为单个 token。
我们如何选择这些子词单元以添加到词汇表中？Sennrich 等人（2016）提出使用字节对编码（BPE；Gage，1994），这是一种压缩算法，迭代地用单个新的未使用索引替换（“合并”）最频繁的字节对。注意，该算法向词汇表中添加子词 token，以最大化输入序列的压缩——如果一个词在输入文本中出现得足够频繁，它将被表示为单个子词单元。
通过 BPE 构建词汇表的子词分词器通常称为 BPE 分词器。在本作业中，我们将实现一个基于字节的 BPE 分词器，其词汇表项为字节或合并后的字节序列，从而在解决未登录词问题和管理输入序列长度之间取得最佳平衡。构建 BPE 分词器词汇表的过程被称为“训练”BPE 分词器。
2.4 BPE 分词器训练
BPE 分词器训练过程包括三个主要步骤。
词汇表初始化 分词器词汇表是一个从字节串 token 到整数 ID 的一对一映射。由于我们训练的是基于字节的 BPE 分词器，初始词汇表就是所有字节的集合。由于有 256 个可能的字节值，初始词汇表大小为 256。
预分词 一旦有了词汇表，原则上可以统计字节在文本中相邻出现的频率，并从最频繁的字节对开始合并。然而，这在计算上非常昂贵，因为每次合并都需要对整个语料进行一次遍历。此外，直接跨语料合并字节可能导致 token 仅在标点符号上有所不同（例如 dog! 与 dog.）。这些 token 将获得完全不同的 token ID，尽管它们可能具有高度的语义相似性（因为它们仅在标点符号上不同）。
为了避免这种情况，我们对语料进行预分词。你可以将其视为对语料的粗粒度分词，帮助我们统计字符对出现的频率。例如，单词 “text” 可能是一个出现 10 次的预分词。在这种情况下，当我们统计字符 “t” 和 “e” 相邻出现的频率时，会看到单词 “text” 中 “t” 和 “e” 相邻，从而将其计数增加 10，而无需遍历整个语料。由于我们训练的是基于字节的 BPE 模型，每个预分词都表示为 UTF-8 字节序列。
Sennrich 等人（2016）的原始 BPE 实现通过简单地按空格拆分（即 s.split(" ")）来进行预分词。相比之下，我们将使用基于正则表达式的预分词器（GPT-2 使用；Radford 等人，2019），来自 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+"""
你可能需要与这个预分词器交互式地拆分一些文本，以更好地理解其行为：
复制
# 需要 regex 包
import regex as re
re.findall(PAT, "some text that i'll pre-tokenize")
['some', ' text', ' that', ' i', "'ll", ' pre', '-', 'tokenize']
然而，在代码中使用它时，应使用 re.finditer，以避免在构建预分词到其计数的映射时存储预分词后的单词。
计算 BPE 合并 现在，我们将输入文本转换为预分词，并将每个预分词表示为 UTF-8 字节序列，我们可以计算 BPE 合并（即训练 BPE 分词器）。从高层次来看，BPE 算法迭代地统计每个字节对的出现频率，并确定频率最高的字节对（“A”，“B”）。然后，每次出现这个最频繁的字节对（“A”，“B”）都会被合并，即被一个新的 token “AB” 替换。这个新的合并 token 被添加到我们的词汇表中；因此，BPE 训练后的最终词汇表大小为初始词汇表大小（我们这里是 256）加上训练期间执行的 BPE 合并操作的数量。
为了提高 BPE 训练的效率，我们不考虑跨越预分词边界的字节对。当合并时，如果字节对的频率相同，我们通过优先选择字典序更大的字节对来打破平局。例如，如果字节对（“A”，“B”）、（“A”，“C”）、（“B”，“ZZ”）和（“BA”，“A”）的频率最高，我们将合并（“BA”，“A”）：
复制
max([("A", "B"), ("A", "C"), ("B", "ZZ"), ("BA", "A")])
('BA', 'A')
特殊 token 通常，某些字符串（例如 <|endoftext|>）用于编码元数据（例如文档之间的边界）。在编码文本时，通常希望将某些字符串视为“特殊 token”，这些特殊 token 永远不应被拆分为多个 token（即始终作为单个 token 保留）。例如，序列结束字符串 <|endoftext|> 应始终作为单个 token（即单个整数 ID）保留，这样我们就知道何时停止从语言模型生成文本。这些特殊 token 必须添加到词汇表中，以便它们具有对应的固定 token ID。
Sennrich 等人（2016）的算法1 提供了一个低效的 BPE 分词器训练实现（基本上遵循我们上面概述的步骤）。作为第一个练习，实现并测试这个函数以验证你的理解可能会有所帮助。
示例 (bpe_example)：BPE 训练示例
这是 Sennrich 等人（2016）的一个简化示例。考虑一个由以下文本组成的语料库：
复制
low low low low low
lower lower widest widest widest
newest newest newest newest newest newest
词汇表中有一个特殊 token <|endoftext|>。
词汇表 我们用特殊 token <|endoftext|> 和 256 个字节值初始化词汇表。
预分词 为了简化并专注于合并过程，我们在此示例中假设预分词只是按空格拆分。当我们预分词并计数时，我们得到以下频率表：
{low: 5, lower: 2, widest: 3, newest: 6}
为了方便，我们可以将其表示为 dict[tuple[bytes], int]，例如 {(l,o,w): 5, ...}。注意，在 Python 中，即使是单个字节也是一个 bytes 对象。Python 中没有 byte 类型来表示单个字节，就像没有 char 类型来表示单个字符一样。
合并 我们首先查看所有连续的字节对，并统计它们出现的总频率：
{lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 9, st: 9, ne: 6, ew: 6}
字节对 (‘es’) 和 (‘st’) 的频率并列最高，我们选择字典序更大的 (‘st’)。然后，我们合并预分词，最终得到：
{(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 个字节字符..., st, est, ow, low, west, ne]
使用这个词汇表和合并集合，单词 newest 将被分词为 [ne, west]。
2.5 BPE 分词器训练实验
让我们在 TinyStories 数据集上训练一个基于字节的 BPE 分词器。数据集的获取/下载说明见第 1 节。开始之前，我们建议你查看 TinyStories 数据集，以了解数据内容。
并行化预分词 你会发现预分词步骤是主要瓶颈。你可以使用内置的 multiprocessing 库并行化你的代码以加速预分词。具体而言，我们建议在并行实现预分词时，将语料分块，同时确保分块边界出现在特殊 token 的开头。你可以直接使用以下链接中的示例代码来获取分块边界，然后将其用于在进程间分配工作：
https://github.com/stanford-cs336/assignment1-basics/blob/main/cs336_basics/pretokenization_example.py
这种分块始终是有效的，因为我们从不希望跨文档边界进行合并。就本作业而言，你总是可以这样拆分。不必担心收到一个非常大的语料且不包含 <|endoftext|> 的极端情况。
预分词前移除特殊 token 在使用正则表达式模式进行预分词（使用 re.finditer）之前，你应该从语料（或分块）中剥离所有特殊 token。确保你根据特殊 token 进行拆分，使得不会在它们所界定的文本之间发生合并。例如，如果你的语料（或分块）类似于 [Doc 1]<|endoftext|>[Doc 2]，你应该在特殊 token <|endoftext|> 处拆分，并分别预分词 [Doc 1] 和 [Doc 2]，使得不会在文档边界之间发生合并。这可以通过使用 re.split 并以 “|”.join(special_tokens) 为分隔符（需用 re.escape 转义，因为 | 可能出现在特殊 token 中）来完成。测试 test_train_bpe_special_tokens 将对此进行测试。
优化合并步骤 上述简化示例中的 BPE 训练朴素实现速度较慢，因为每次合并时，它都会遍历所有字节对以确定最频繁的字节对。然而，每次合并后，仅有与合并字节对重叠的字节对的计数会发生变化。因此，通过索引所有字节对的计数并增量更新这些计数，而不是显式遍历每对字节来统计频率，可以提高 BPE 训练速度。通过这种缓存机制，你可以获得显著的加速，尽管我们指出，BPE 训练的合并部分在 Python 中不可并行化。
低资源/降规模提示：性能分析
你应该使用 cProfile 或 scalene 等性能分析工具来识别实现中的瓶颈，并专注于优化这些瓶颈。
低资源/降规模提示：“降规模”
不要急于在整个 TinyStories 数据集上训练分词器，我们建议你先在一个小的数据子集（即“调试数据集”）上进行训练。例如，你可以在 TinyStories 验证集（22 K 个文档，而不是 2.12 M）上训练分词器。这说明了尽可能降规模的一般策略：例如，使用更小的数据集、更小的模型规模等。选择调试数据集或超参数配置的大小需要仔细考虑：你希望调试集足够大，以具有与完整配置相同的瓶颈（以便你进行的优化能够泛化），但又不能太大，以至于运行时间过长。
问题 (train_bpe)：BPE 分词器训练（15 分）
交付物：编写一个函数，给定一个输入文本文件路径，训练一个（基于字节的）BPE 分词器。你的 BPE 训练函数应至少处理以下输入参数：
input_path: str 指向用于 BPE 分词器训练的文本文件路径。
vocab_size: int 正整数，定义最大最终词汇表大小（包括初始字节词汇表、通过合并产生的词汇表项以及任何特殊 token）。
special_tokens: list[str] 要添加到词汇表中的字符串列表。这些特殊 token 在其他方面不影响 BPE 训练。
你的 BPE 训练函数应返回结果词汇表和合并：
vocab: dict[int, bytes] 分词器词汇表，从 int（词汇表中的 token ID）到 bytes（token 字节）的映射。
merges: list[tuple[bytes, bytes]] 训练产生的 BPE 合并列表。每个列表项是一个字节元组 (a, b)，表示 a 与 b 合并。合并应按创建顺序排序。
为了使用我们提供的测试测试你的 BPE 训练函数，你首先需要实现 test adapter 中的 [adapters.run_train_bpe]。然后运行 uv run pytest tests/test_train_bpe.py。你的实现应能通过所有测试。可选地（这可能需要大量时间），你可以使用系统语言实现训练方法的关键部分，例如 C++（考虑 cppyy）或 Rust（使用 PyO3）。如果你这样做，请注意哪些操作需要复制与直接从 Python 内存读取，并确保留下构建说明，或确保它仅使用 pyproject.toml 构建。另请注意，GPT-2 正则表达式在大多数正则引擎中支持不佳，且速度较慢。我们已验证 Oniguruma 速度尚可且支持负向前瞻，但 Python 的 regex 包速度更快。
问题 (train_bpe_tinystories)：在 TinyStories 上训练 BPE（2 分）
(a) 在 TinyStories 数据集上训练一个基于字节的 BPE 分词器，最大词汇表大小为 10 000。确保将 TinyStories 的特殊 token <|endoftext|> 添加到词汇表中。将生成的词汇表和合并序列化到磁盘以便进一步检查。训练耗时多少小时和内存？词汇表中最长的 token 是什么？是否合理？
资源需求：≤ 30 分钟（无 GPU），≤ 30 GB RAM
提示：通过以下两个事实，你应该能在预分词阶段使用多进程将 BPE 训练时间控制在 2 分钟以内：
(a) <|endoftext|> token 在数据文件中分隔文档。
(b) <|endoftext|> token 在应用 BPE 合并前作为特殊情况处理。
交付物：一到两句话的回答。
(b) 对你的代码进行性能分析。分词器训练过程中哪一步耗时最多？
交付物：一到两句话的回答。
接下来，我们将尝试在 OpenWebText 数据集上训练一个基于字节的 BPE 分词器。与之前一样，我们建议你先查看数据集以更好地了解其内容。
问题 (train_bpe_expts_owt)：在 OpenWebText 上训练 BPE（2 分）
(a) 在 OpenWebText 数据集上训练一个基于字节的 BPE 分词器，最大词汇表大小为 32 000。将生成的词汇表和合并序列化到磁盘以便进一步检查。词汇表中最长的 token 是什么？是否合理？
资源需求：≤ 12 小时（无 GPU），≤ 100 GB RAM
交付物：一到两句话的回答。
(b) 比较并对比你在 TinyStories 和 OpenWebText 上训练得到的分词器。
交付物：一到两句话的回答。
2.6 BPE 分词器：编码与解码
在上一部分中，我们实现了一个函数，用于在输入文本上训练 BPE 分词器以获得分词器词汇表和 BPE 合并列表。现在，我们将实现一个 BPE 分词器，它加载提供的词汇表和合并列表，并使用它们将文本编码为/从 token ID 解码。
2.6.1 文本编码
通过 BPE 编码文本的过程与训练 BPE 词汇表的过程类似。主要有以下几个步骤。
步骤 1：预分词。我们首先对序列进行预分词，并将每个预分词表示为 UTF-8 字节序列，就像我们在 BPE 训练中所做的那样。我们将在每个预分词内部将这些字节合并为词汇表元素，独立处理每个预分词（不跨预分词边界进行合并）。
步骤 2：应用合并。我们然后采用在 BPE 训练期间创建的词汇表元素合并序列，并按创建顺序将其应用于我们的预分词。
示例 (bpe_encoding)：BPE 编码示例
例如，假设我们的输入字符串是 'the cat ate'，词汇表为 {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'}，我们学到的合并为 [(b't', b'h'), (b' ', b'c'), (b' ', b'a'), (b'th', b'e'), (b' a', b't')]。首先，我们的预分词器会将该字符串拆分为 ['the', ' cat', ' ate']。
然后，我们将查看每个预分词并应用 BPE 合并。
第一个预分词 'the' 初始表示为 [b't', b'h', b'e']。查看我们的合并列表，我们确定第一个适用的合并是 (b't', b'h')，我们用它将预分词转换为 [b'th', b'e']。然后，我们回到合并列表，确定下一个适用的合并是 (b'th', b'e')，它将预分词转换为 [b'the']。最后，查看合并列表，我们发现没有更多适用的合并（因为整个预分词已合并为单个 token），因此我们完成了 BPE 合并的应用。对应的整数序列为 [9]。
对剩余的预分词重复此过程，我们看到预分词 ' cat' 在应用 BPE 合并后表示为 [b' c', b'a', b't']，变为整数序列 [7, 1, 5]。最后一个预分词 ' ate' 在应用 BPE 合并后为 [b' at', b'e']，变为整数序列 [10, 3]。因此，我们输入字符串的最终编码结果为 [9, 7, 1, 5, 10, 3]。
特殊 token。你的分词器应能够正确处理用户定义的特殊 token（在构建分词器时提供）。
内存考虑。假设我们想对一个无法装入内存的大型文本文件进行分词。为了高效地对这种大文件（或任何其他数据流）进行分词，我们需要将其分割为可管理的块并依次处理每个块，从而使内存复杂度保持恒定，而不是随文本大小线性增长。在此过程中，我们需要确保一个 token 不会跨越块边界，否则我们将得到与将整个序列一次性装入内存进行分词不同的结果。
2.6.2 文本解码
要将整数 token ID 序列解码回原始文本，我们可以简单地查找词汇表中每个 ID 对应的条目（一个字节序列），将它们连接起来，然后将字节解码为 Unicode 字符串。注意，输入的 ID 序列不一定映射到有效的 Unicode 字符串（因为用户可以输入任何整数 ID 序列）。如果输入的 token ID 不产生有效的 Unicode 字符串，你应使用官方 Unicode 替换字符 U+FFFD 替换格式错误的字节。bytes.decode 的 errors 参数控制如何处理 Unicode 解码错误，使用 errors='replace' 将自动用替换标记替换格式错误的数据。
问题 (tokenizer)：实现分词器（15 分）
交付物：实现一个 Tokenizer 类，给定一个词汇表和合并列表，将文本编码为整数 ID，并将整数 ID 解码为文本。你的分词器还应支持用户提供的特殊 token（如果它们尚不在词汇表中，则附加到词汇表中）。我们建议以下接口：
def __init__(self, vocab, merges, special_tokens=None)
从给定的词汇表、合并列表和（可选的）特殊 token 列表构建分词器。该函数应接受以下参数：
vocab: dict[int, bytes]
merges: list[tuple[bytes, bytes]]
special_tokens: list[str] | None = None
复制
@classmethod
def from_files(cls, vocab_filepath, merges_filepath, special_tokens=None)
类方法，从序列化的词汇表和合并列表（与你的 BPE 训练代码输出的格式相同）以及（可选的）特殊 token 列表构建并返回一个分词器。该方法应接受以下额外参数：
vocab_filepath: str
merges_filepath: str
special_tokens: list[str] | None = None
def encode(self, text: str) -> list[int]
将输入文本编码为 token ID 序列。
def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]
给定一个字符串可迭代对象（例如 Python 文件句柄），返回一个生成器，该生成器惰性地产出 token ID。这对于内存高效地对无法直接装入内存的大型文件进行分词是必需的。
def decode(self, ids: list[int]) -> str
将 token ID 序列解码为文本。
为了使用我们提供的测试测试你的分词器，你首先需要实现 test adapter 中的 [adapters.get_tokenizer]。然后运行 uv run pytest tests/test_tokenizer.py。你的实现应能通过所有测试。
2.7 实验
问题 (tokenizer_experiments)：分词器实验（4 分）
(a) 从 TinyStories 和 OpenWebText 中各采样 10 篇文档。使用你之前训练的 TinyStories 和 OpenWebText 分词器（词汇表大小分别为 10 K 和 32 K），将这些采样文档编码为整数 ID。每个分词器的压缩比（字节/token）是多少？
交付物：一到两句话的回答。
(b) 如果你用 TinyStories 分词器对 OpenWebText 样本进行分词，会发生什么？比较压缩比或定性描述结果。
交付物：一到两句话的回答。
(c) 估计你的分词器的吞吐量（例如，字节/秒）。对 Pile 数据集（825 GB 文本）进行分词需要多长时间？
交付物：一到两句话的回答。
(d) 使用你的 TinyStories 和 OpenWebText 分词器，将相应的训练和开发数据集编码为整数 token ID 序列。我们稍后将使用这些来训练语言模型。我们建议将 token ID 序列化为 uint16 数据类型的 NumPy 数组。为什么是 uint16 的合适选择？
交付物：一到两句话的回答。
3 Transformer 语言模型架构
语言模型将一批整数 token ID 序列（即形状为 (batch_size, sequence_length) 的 torch.Tensor）作为输入，并返回（批量的）词汇表上的归一化概率分布（即形状为 (batch_size, sequence_length, vocab_size) 的 PyTorch 张量），其中每个输入 token 的预测分布是对下一个词的预测。在训练语言模型时，我们使用这些下一个词预测来计算实际下一个词与预测下一个词之间的交叉熵损失。在推理期间从语言模型生成文本时，我们取最后一个时间步（即序列中的最后一项）的预测下一个词分布来生成序列中的下一个 token（例如，通过取概率最高的 token，从分布中采样等），将生成的 token 追加到输入序列中，并重复此过程。
在本作业的这一部分中，你将从零开始构建这个 Transformer 语言模型。我们将从模型的高级描述开始，逐步详细介绍各个组件。
3.1 Transformer LM
给定一个 token ID 序列，Transformer 语言模型使用输入嵌入将 token ID 转换为密集向量，将嵌入的 token 通过 num_layers 个 Transformer 块，然后应用学习的线性投影（“输出嵌入”或“LM 头”）以产生预测的下一个 token 的 logits。见图 1 的示意图。
3.1.1 Token 嵌入
在第一步中，Transformer 将（批量的）token ID 序列嵌入到包含 token 标识信息的向量序列中（图 1 中的红色块）。
更具体地说，给定一个 token ID 序列，Transformer 语言模型使用 token 嵌入层产生一个向量序列。每个嵌入层接收一个形状为 (batch_size, sequence_length) 的整数张量，并产生一个形状为 (batch_size, sequence_length, d_model) 的向量序列。
3.1.2 预归一化 Transformer 块
嵌入后，激活被多个结构相同的神经网络层处理。标准的仅解码器 Transformer 语言模型由 num_layers 个相同层（通常称为 Transformer “块”）组成。每个 Transformer 块接收一个形状为 (batch_size, sequence_length, d_model) 的输入，并返回一个形状为 (batch_size, sequence_length, d_model) 的输出。每个块通过自注意力聚合跨序列的信息，并通过前馈层非线性地转换它。
3.2 输出归一化与嵌入
经过 num_layers 个 Transformer 块后，我们将获取最终激活并将其转换为词汇表上的分布。
我们将实现“预归一化” Transformer 块（详见 §3.5），这还需要在最终 Transformer 块之后使用层归一化，以确保其输出被正确缩放。
在此归一化之后，我们将使用标准的学习线性变换将 Transformer 块的输出转换为预测的下一个 token 的 logits（例如，见 Radford 等人，2018，公式 2）。
3.3 批处理、爱因斯坦求和与高效计算
在整个 Transformer 中，我们将对许多批处理输入执行相同的计算。以下是一些示例：
批处理的元素：我们对每个批处理元素应用相同的 Transformer 前向操作。
序列长度：像 RMSNorm 和前馈这样的“逐位置”操作在序列的每个位置上相同地操作。
注意力头：注意力操作在“多头”注意力操作中跨注意力头进行批处理。
拥有一种符合人体工程学的方式来执行这样的操作，充分利用 GPU，并且易于阅读和理解，这是非常有用的。许多 PyTorch 操作可以在张量的开头接受额外的“批处理”维度，并高效地跨这些维度重复/广播操作。
例如，假设我们正在执行逐位置、批处理操作。我们有一个“数据张量” D，形状为 (batch_size, sequence_length, d_model)，我们希望对形状为 (d_model, d_model) 的矩阵 A 执行批处理的向量-矩阵乘法。在这种情况下，D @ A 将执行批处理的矩阵乘法，这是 PyTorch 中的一个高效原语，其中 (batch_size, sequence_length) 维度被批处理。
因此，假设你的函数可能被赋予额外的批处理维度，并将这些维度保留在 PyTorch 形状的开头，这是有帮助的。为了以这种方式批处理组织张量，它们可能需要使用 view、reshape 和 transpose 的多个步骤进行整形。这可能有点痛苦，并且通常很难阅读代码在做什么以及你的张量的形状是什么。
一个更符合人体工程学的选择是在 torch.einsum 中使用爱因斯坦求和符号，或者使用框架无关的库，如 einops 或 einx。这两个关键操作是 einsum，它可以对输入张量的任意维度进行张量收缩，以及 rearrange，它可以重新排序、连接和拆分任意维度。事实证明，机器学习中的几乎所有操作都是维度操作和张量收缩的某种组合，偶尔伴随着（通常是逐元素的）非线性函数。这意味着当你使用爱因斯坦求和符号时，你的很多代码可以更易于阅读和灵活。我们强烈建议为本课程学习和使用爱因斯坦求和符号。之前没有接触过爱因斯坦求和符号的学生应使用 einops（文档在此），已经熟悉 einops 的学生应学习更通用的 einx（在此）。这两个包都已在我们提供的环境中安装。
这里我们给出一些如何使用爱因斯坦求和符号的示例。这些是对 einops 文档的补充，你应该先阅读 einops 文档。
示例 (einstein_example1)：使用 einops.einsum 的批处理矩阵乘法
复制
import torch
from einops import rearrange, einsum

# 基本实现
Y = D @ A.T

# 很难判断输入和输出形状以及它们的含义
# 形状 D 和 A 可以是什么，是否有任何意外行为？

# Einsum 是自文档化且健壮的
# D A --> Y
Y = einsum(D, A, "batch sequence d_in, d_out d_in -> batch sequence d_out")

# 或者，一个批处理版本，其中 D 可以有任意前导维度，但 A 受约束
Y = einsum(D, A, "... d_in, d_out d_in -> ... d_out")
示例 (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)

# 重塑并相乘
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

# 或者一步到位：
dimmed_images = einsum(
    images, dim_by,
    "batch height width channel, dim_value -> batch dim_value height width channel"
)
示例 (einstein_example3)：使用 einops.rearrange 的像素混合
假设我们有一批图像，表示为形状为 (batch, height, width, 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)

# 重塑图像张量以混合所有像素
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 替换了笨拙的 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 一步到位
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
)
第一个实现可以通过在前后放置注释来改进，以指示……
爱因斯坦求和符号可以处理任意输入批处理维度，并且具有自文档化的关键优势。在使用爱因斯坦求和符号的代码中，输入和输出张量的相关形状更加清晰。对于剩余的张量，你可以考虑使用张量类型提示，例如使用 jaxtyping 库（不特定于 Jax）。
我们将在作业 2 中进一步讨论使用爱因斯坦求和符号的性能影响，但现在只需知道它们几乎总是比替代方案更好！
3.3.1 数学符号与内存顺序
许多机器学习论文在其符号中使用行向量，这使得表示与 NumPy 和 PyTorch 默认使用的行主序内存顺序很好地结合在一起。使用行向量时，线性变换看起来像这样：
y = x W^T          (1)
对于行主序的 W ∈ R^{d_out × d_in} 和行向量 x ∈ R^{1 × d_in}。
在线性代数中，通常更常见的是使用列向量，其中线性变换看起来像这样：
y = W x            (2)
对于行主序的 W ∈ R^{d_out × d_in} 和列向量 x ∈ R^{d_in}。我们将在此作业中使用列向量进行数学符号表示，因为通常更容易以这种方式遵循数学。你应该记住，如果你想使用纯矩阵乘法符号，你将不得不使用行向量约定来应用矩阵，因为 PyTorch 使用行主序内存顺序。如果你使用 einsum 进行矩阵操作，这应该不是问题。
3.4 基本构建块：线性与嵌入模块
3.4.1 参数初始化
有效地训练神经网络通常需要仔细初始化模型参数——糟糕的初始化可能会导致诸如梯度消失或爆炸等不良行为。预归一化 transformer 对初始化异常健壮，但它们仍可能对训练速度和收敛产生重大影响。由于本作业已经很长了，我们将把细节留到作业 3，而是给你一些近似初始化，这些初始化在大多数情况下应该效果很好。现在，使用：
嵌入：N(μ = 0, σ² = 1) 截断于 [-3, 3]
RMSNorm：1
你应该使用 torch.nn.init.trunc_normal_ 来初始化截断正态权重。
3.4.2 线性模块
线性层是 Transformer 和一般神经网络的基石之一。首先，你将实现自己的 Linear 类，该类继承自 torch.nn.Module，并执行线性变换：
y = W x          (3)
注意，我们不包括偏置项，遵循大多数现代 LLM。
问题 (linear)：实现线性模块（1 分）
交付物：实现一个继承自 torch.nn.Module 的 Linear 类，并执行线性变换。你的实现应遵循 PyTorch 内置 nn.Linear 模块的接口，除了没有偏置参数或权重。我们建议以下接口：
def __init__(self, in_features, out_features, device=None, dtype=None)
构建线性变换模块。该函数应接受以下参数：
in_features: int 输入的最终维度
out_features: int 输出的最终维度
device: torch.device | None = None 存储参数的设备
dtype: torch.dtype | None = None 参数的数据类型
def forward(self, x: torch.Tensor) -> torch.Tensor
对输入应用线性变换。
确保：
继承 nn.Module
调用超类构造函数
构造并将参数存储为 W（不是 W^T）以符合内存顺序，将其放入 nn.Parameter
不要使用 nn.Linear 或 nn.functional.linear
对于初始化，使用上面的设置以及 torch.nn.init.trunc_normal_ 来初始化权重。
为了测试你的 Linear 模块，在 test adapter 中实现 [adapters.run_linear]。adapter 应将给定权重加载到你的 Linear 模块中。你可以为此使用 Module.load_state_dict。然后运行 uv run pytest -k test_linear。
3.4.3 嵌入模块
如上所述，Transformer 的第一层是嵌入层，它将整数 token ID 映射到维度为 d_model 的向量空间。我们将实现一个自定义的 Embedding 类，该类继承自 torch.nn.Module（因此你不应使用 nn.Embedding）。forward 方法应通过索引到形状为 (vocab_size, d_model) 的嵌入矩阵中，使用形状为 (batch_size, sequence_length) 的 token ID torch.LongTensor 来选择每个 token ID 的嵌入向量。
问题 (embedding)：实现嵌入模块（1 分）
交付物：实现一个继承自 torch.nn.Module 的 Embedding 类，并执行嵌入查找。你的实现应遵循 PyTorch 内置 nn.Embedding 模块的接口。我们建议以下接口：
def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None)
构建嵌入模块。该函数应接受以下参数：
num_embeddings: int 词汇表大小
embedding_dim: int 嵌入向量的维度，即 d_model
device: torch.device | None = None 存储参数的设备
dtype: torch.dtype | None = None 参数的数据类型
def forward(self, token_ids: torch.Tensor) -> torch.Tensor
查找给定 token ID 的嵌入向量。
确保：
继承 nn.Module
调用超类构造函数
将嵌入矩阵初始化为 nn.Parameter
将嵌入矩阵存储为 d_model 为最终维度
不要使用 nn.Embedding 或 nn.functional.embedding
同样，使用上面的设置进行初始化，并使用 torch.nn.init.trunc_normal_ 来初始化权重。
为了测试你的实现，在 test adapter 中实现 [adapters.run_embedding]。然后运行 uv run pytest -k test_embedding。
3.5 预归一化 Transformer 块
每个 Transformer 块有两个子层：多头自注意力机制和逐位置前馈网络（Vaswani 等人，2017，第 3.1 节）。
在原始 Transformer 论文中，模型在每个两个子层周围使用残差连接，然后使用层归一化。这种架构通常被称为“后归一化” Transformer，因为层归一化应用于子层输出。然而，多项工作发现，将层归一化从每个子层的输出移到每个子层的输入（在最终 Transformer 块之后额外使用层归一化）可以提高 Transformer 训练的稳定性（Nguyen 和 Salazar，2019；Xiong 等人，2020）——见图 2 的“预归一化” Transformer 块的视觉表示。然后，每个 Transformer 块子层的输出通过残差连接加到子层输入上（Vaswani 等人，2017，第 5.4 节）。预归一化的直觉是，从 Transformer 的输入嵌入到最终输出之间有一条干净的“残差流”，没有任何归一化，这被认为可以改善梯度流。这种预归一化 Transformer 现在是当今语言模型的标准（例如 GPT-3、LLaMA、PaLM 等），因此我们将实现这种变体。我们将依次介绍预归一化 Transformer 块的组件，逐一实现它们。
3.5.1 均方根层归一化
原始 Transformer 实现（Vaswani 等人，2017）使用层归一化（Ba 等人，2016）来归一化激活。按照 Touvron 等人（2023）的做法，我们将使用均方根层归一化（RMSNorm；Zhang 和 Sennrich，2019，公式 4）进行层归一化。给定一个激活向量 a ∈ R^{d_model}，RMSNorm 将每个激活 a_i 重新缩放如下：
a_i = a_i / sqrt((1/d_model) * sum_{j=1}^{d_model} a_j^2 + eps) * g_i
其中 g ∈ R^{d_model} 是可学习的参数（与原始层归一化中的 2d_model 个参数相比，只有 d_model 个参数），eps 是一个通常固定为 1e-5 的超参数。
你应该将输入上转换为 torch.float32，以防止在平方输入时溢出。总体上，你的 forward 方法应如下所示：
复制
in_dtype = x.dtype
x = x.to(torch.float32)

# 在这里执行 RMSNorm
...
result = ...

# 以原始数据类型返回结果
return result.to(in_dtype)
问题 (rmsnorm)：实现均方根层归一化（1 分）
交付物：将 RMSNorm 实现为 torch.nn.Module。我们建议以下接口：
def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None)
构建 RMSNorm 模块。该函数应接受以下参数：
d_model: int 模型的隐藏维度
eps: float = 1e-5 用于数值稳定性的 epsilon 值
device: torch.device | None = None 存储参数的设备
dtype: torch.dtype | None = None 参数的数据类型
def forward(self, x: torch.Tensor) -> torch.Tensor
处理形状为 (batch_size, sequence_length, d_model) 的输入张量，并返回相同形状的张量。
注意：记住在执行归一化之前将输入上转换为 torch.float32（然后再下转换为原始数据类型），如上所述。
为了测试你的实现，在 test adapter 中实现 [adapters.run_rmsnorm]。然后运行 uv run pytest -k test_rmsnorm。
3.5.2 逐位置前馈网络
复制
SiLU:   f(x) = x * sigmoid(x)
Identity: f(x) = x
ReLU:   f(x) = max(0, x)
在原始 Transformer 论文（Vaswani 等人，2017，第 3.3 节）中，Transformer 前馈网络由两个线性变换组成，中间使用 ReLU 激活（ReLU(x) = max(0, x)）。内部前馈层的维度通常是输入维度的 4 倍。
然而，与原始设计相比，现代语言模型倾向于引入两个主要变化：它们使用另一个激活函数并采用门控机制。具体来说，我们将实现 LLM（如 Llama 3（Grattafiori 等人，2024）和 Qwen 2.5（Yang 等人，2024））采用的“SwiGLU”激活函数，该函数将 SiLU（通常称为 Swish）激活与门控线性单元（GLU）相结合。我们还将省略有时在线性层中使用的偏置项，遵循大多数现代 LLM（如 PaLM（Chowdhery 等人，2022）和 LLaMA（Touvron 等人，2023））的做法。
SiLU 或 Swish 激活函数（Hendrycks 和 Gimpel，2016；Elfwing 等人，2017）定义如下：
SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x))          (5)
如图 3 所示，SiLU 激活函数类似于 ReLU 激活函数，但在零处平滑。
门控线性单元（GLU）最初由 Dauphin 等人（2017）定义为通过 sigmoid 函数的线性变换与另一个线性变换的逐元素乘积：
GLU(x, W_1, W_2) = sigmoid(W_1 x) * W_2 x          (6)
其中 * 表示逐元素乘法。门控线性单元被认为“通过为梯度提供线性路径同时保留非线性能力，减少了深度架构中的梯度消失问题”。
将 SiLU/Swish 与 GLU 结合，我们得到 SwiGLU，我们将用于我们的前馈网络：
FFN(x) = SwiGLU(x, W_1, W_2, W_3) = W_2 (SiLU(W_1 x) * W_3 x)          (7)
其中 x ∈ R^{d_model}，W_1, W_3 ∈ R^{d_ff × d_model}，W_2 ∈ R^{d_model × d_ff}，通常 d_ff = (8/3) * d_model。
Shazeer（2020）首先提出了将 SiLU/Swish 激活与 GLU 结合，并进行实验表明 SwiGLU 在语言建模任务上优于 ReLU 和 SiLU（无门控）等基线。在作业后面，你将比较 SwiGLU 和 SiLU。虽然我们提到了这些组件的一些启发式论据（论文提供了更多支持证据），但保持实证视角是好的：Shazeer 论文中一句著名的话是：
我们不为这些架构为何有效提供解释；我们将它们的成功归因于神圣的仁慈，就像所有其他事情一样。
问题 (positionwise_feedforward)：实现逐位置前馈网络（2 分）
交付物：实现由 SiLU 激活函数和 GLU 组成的 SwiGLU 前馈网络。
注意：在这种特定情况下，你可以自由地在实现中使用 torch.sigmoid 以确保数值稳定性。
你应将 d_ff 设置为近似于 (8/3) * d_model，同时确保内部前馈层的维度是 64 的倍数，以充分利用你的硬件。为了使用我们提供的测试测试你的实现，你需要在 test adapter 中实现 [adapters.run_swiglu]。然后运行 uv run pytest -k test_swiglu 来测试你的实现。
3.5.3 相对位置嵌入
为了将位置信息注入到模型中，我们将实现旋转位置嵌入（Su 等人，2021），通常称为 RoPE。对于给定的查询 token q^{(i)} = W_q x^{(i)} ∈ R^{d} 在 token 位置 i，我们将应用成对旋转矩阵 R_i，得到 q'^{(i)} = R_i q^{(i)} = R_i W_q x^{(i)}。这里，R_i 将嵌入元素对 q_{2k-1:2k} 作为二维向量旋转角度 θ_{i,k} = i / (10000^{2k/d})，对于 k ∈ {1, ..., d/2} 和某个常数 10000。
因此，我们可以将 R_i 视为一个大小为 d × d 的块对角矩阵，对于 k ∈ {1, ..., d/2}，块为：
复制
R_i^{(k)} = [
    [cos(θ_{i,k}), -sin(θ_{i,k})],
    [sin(θ_{i,k}),  cos(θ_{i,k})]
]
因此，我们得到完整的旋转矩阵：
R_i = diag(R_i^{(1)}, R_i^{(2)}, ..., R_i^{(d/2)})
其中 0 表示 2 × 2 零矩阵。虽然可以构造完整的 d × d 矩阵，但一个好的解决方案应该利用该矩阵的属性来更高效地实现变换。由于我们只关心给定序列内 token 的相对旋转，我们可以在层之间以及不同批处理之间重用我们为 cos(θ_{i,k}) 和 sin(θ_{i,k}) 计算的值。如果你想优化它，你可以使用一个被所有层引用的单一 RoPE 模块，并且它可以有一个在 init 期间使用 self.register_buffer(persistent=False) 创建的二维预计算 sin 和 cos 值的缓冲区，而不是 nn.Parameter（因为我们不想学习这些固定的余弦和正弦值）。然后，我们对 k^{(j)} 执行与 q^{(i)} 完全相同的旋转过程，按相应的 R_j 旋转。注意，该层没有可学习的参数。
问题 (rope)：实现 RoPE（2 分）
交付物：实现一个 RotaryPositionalEmbedding 类，将 RoPE 应用于输入张量。
建议的接口如下：
def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None)
构建 RoPE 模块，并在需要时创建缓冲区。
theta: float RoPE 的 theta 值
d_k: int 查询和键向量的维度
max_seq_len: int 将输入的最大序列长度
device: torch.device | None = None 存储缓冲区的设备
def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor
处理形状为 (..., seq_len, d_k) 的输入张量，并返回相同形状的张量。
注意，你应容忍 x 具有任意数量的批处理维度。你应假设 token 位置是形状为 (..., seq_len) 的张量，指定 x 沿序列维度的 token 位置。
你应使用 token 位置沿序列维度切片你的（可能预计算的）cos 和 sin 张量。
为了测试你的实现，完成 [adapters.run_rope] 并确保它能通过 uv run pytest -k test_rope。
3.5.4 缩放点积注意力
我们现在将实现 Vaswani 等人（2017）（第 3.2.1 节）中描述的缩放点积注意力。作为初步步骤，注意力操作的定义将使用 softmax，这是一个将未归一化的分数向量转换为归一化分布的操作：
softmax(x)_i = exp(x_i) / sum_j exp(x_j)          (10)
注意，exp(v_i) 对于大值可能变为 inf（然后 inf/inf = NaN）。我们可以通过注意到 softmax 操作对于向所有输入添加任何常数 c 是不变的来避免这种情况。我们可以利用此属性以确保数值稳定性——通常，我们将从 o_i 的所有元素中减去 o_i 的最大条目，使新的最大条目为 0。你现在将实现 softmax，使用此技巧以确保数值稳定性。
问题 (softmax)：实现 softmax（1 分）
交付物：编写一个函数，对张量应用 softmax 操作。你的函数应接受两个参数：一个张量和一个维度 i，并将 softmax 应用于输入张量的第 i 个维度。输出张量应具有与输入张量相同的形状，但其第 i 个维度现在将具有归一化的概率分布。使用从第 i 个维度的所有元素中减去最大值以避免数值稳定性问题的技巧。
完成 [adapters.run_softmax]，然后运行 uv run pytest -k test_softmax_matches_pytorch 以测试你的实现。
我们现在可以如下数学地定义注意力操作：
Attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V          (11)
其中 Q ∈ R^{n × d_k}，K ∈ R^{m × d_k}，V ∈ R^{m × d_v}。这里，Q、K 和 V 都是此操作的输入——注意，这些不是可学习的参数。如果你想知道为什么不是 Q K^T，请参见 3.3.1。
掩码：有时方便屏蔽注意力操作的输出。掩码应具有形状 M ∈ {True, False}^{n × m}，并且该布尔矩阵的每一行 i 指示查询 i 应关注哪些键。规范地（且稍微令人困惑地），位置 (i, j) 处的 True 值指示查询 i 关注键 j，False 值指示查询 i 不关注键 j。换句话说，“信息”在值为 True 的 (i, j) 对处“流动”。例如，考虑一个 1 × 3 掩码矩阵，条目为 [[True, True, False]]。单个查询向量仅关注前两个键。
在计算上，使用掩码比计算子序列的注意力更有效，我们可以通过将预 softmax 值添加到掩码矩阵中值为 False 的任何条目中 -∞ 来实现这一点。
问题 (scaled_dot_product_attention)：实现缩放点积注意力（5 分）
交付物：实现缩放点积注意力函数。你的实现应处理形状为 (batch_size, ..., seq_len, d_k) 的查询和键，以及形状为 (batch_size, ..., seq_len, d_v) 的值，其中 ... 表示任意数量的其他批处理维度（如果提供）。实现应返回形状为 (batch_size, ..., d_v) 的输出。有关批处理维度的讨论，请参见第 3.3 节。
你的实现还应支持可选的、用户提供的形状为 (seq_len, seq_len) 的布尔掩码。掩码值为 True 的位置的注意力概率应共同和为 1，掩码值为 False 的位置的注意力概率应为零。
为了使用我们提供的测试测试你的实现，你需要在 test adapter 中实现 [adapters.run_scaled_dot_product_attention]。
uv run pytest -k test_scaled_dot_product_attention 在第三阶输入张量上测试你的实现，而 uv run pytest -k test_4d_scaled_dot_product_attention 在第四阶输入张量上测试你的实现。
3.5.5 因果多头自注意力
我们将实现 Vaswani 等人（2017）第 3.2.2 节中描述的多头自注意力。回想一下，在数学上，应用多头注意力的操作定义如下：
MultiHead(Q, K, V) = Concat(head_1, ..., head_h)          (12)
其中
head_i = Attention(Q_i, K_i, V_i)          (13)
其中 Q_i, K_i, V_i 是 Q、K 和 V 的嵌入维度的第 i ∈ {1, ..., h} 个切片，大小为 d_k 或 d_v。Attention 是第 3.5.4 节中定义的缩放点积注意力操作。由此我们可以形成多头自注意力操作：
MultiHeadSelfAttention(x) = W_O MultiHead(W_Q x, W_K x, W_V x)          (14)
这里，可学习的参数是 W_Q ∈ R^{h d_k × d_model}，W_K ∈ R^{h d_k × d_model}，W_V ∈ R^{h d_v × d_model}，和 W_O ∈ R^{d_model × h d_v}。由于 Q、K 和 V 在多头注意力操作中被切片，我们可以将 W_Q、W_K 和 W_V 视为沿输出维度为每个头分开的。当你实现这个功能时，你应该总共通过三个矩阵乘法来计算键、值和查询投影。
作为延伸目标，尝试将键、查询和值投影组合成单个权重矩阵，这样你只需要一个矩阵乘法。

Python