A2 · Systems（让模型跑得又快又省）
年份 2026（Spring 2026）—— 往年（2025 / 2024）题面不同，见「作业说明」
分值 未记录 ✗（官方主页没有公布各项作业的比重 ✓ 不推断 ✗）
发布 / 截止 4 月 15 日发布 / 4 月 29 日截止（L10）（官方日程标注 tentative / 暂定 ✓）
算力 基准测试需 1×GPU（写 kernel 与 profiling 本身可在本机先调通 ✓）
配哪几讲 候选（机器按术语匹配，供参考）：第5讲 §第二部分：让 GPU 跑得快的六个技巧；第5讲 §第一部分：GPU 的硬件模型；第6讲 §五、Softmax：行内归约；第6讲 §六、行归约：当一行装不进一个块；第7讲 §四、PyTorch 动手：用代码实现集合通信；第7讲 §六、数据并行 (DDP)；第8讲 §五、模型并行；第8讲 §四、数据并行：从 DDP 到 ZeRO 三级
代码仓库 https://github.com/stanford-cs336/assignment2-systems
官方题面 PDF（本地原件在 _local/A2/）

官方原文：Profile and benchmark the model and layers from Assignment 1 using advanced tools, optimize Attention with your own Triton implementation of FlashAttention2. Build a memory-efficient, distributed version of the Assignment 1 model training code.

做什么

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

第 1 页

CS336 Assignment 2 (systems)：Systems and Parallelism 版本 26.1.3 CS336 Staff 2026 年春季

1 作业概览

在这个作业中，你将获得一些实践经验：提升单 GPU 训练速度，并将训练扩展到多块 GPU。

你将实现什么

  1. Benchmarking 与 profiling harness
  2. Activation checkpointing
  3. Flash Attention 2 Triton kernel
  4. Distributed data parallel 训练
  5. Optimizer state sharding
  6. Fully sharded data parallel 训练

代码是什么样子

作业代码和这份 write-up 可在 GitHub 上获取：

github.com/stanford-cs336/assignment2-systems

请使用 Git clone 该 repository。如果有任何更新，我们会通知你，你可以用 git pull 获取最新版本。

  1. cs336-basics/：在这个作业中，你将对我们在 assignment 1 中构建的一些组件进行 profiling。这个文件夹包含 assignment 1 的 staff solution 代码，所以你会在这里找到 cs336-basics/pyproject.toml 和 cs336-basics/cs336_basics package。如果你想使用自己的模型实现，可以修改 base directory 中的 pyproject.toml 文件，使其指向你自己的 package。
  2. /：cs336-systems 的 base directory。我们创建了一个名为 cs336_systems 的空 module。注意这里没有代码，所以你可以从零开始做任何你想做的事。
  3. tests/*.py：这个目录包含你必须通过的所有测试。这些测试会调用 tests/adapters.py 中定义的 hooks。你将实现 adapters，把你的代码连接到测试。编写更多测试和/或修改测试代码可能有助于调试你的代码，但你的实现应当通过最初提供的 test suite。
  4. README.md：这个文件包含关于期望的目录结构的更多细节，以及一些设置环境的基本说明。

如何提交

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

运行 test_and_make_submission.sh 脚本以创建 code.zip 文件。

第 2 页

2 Profiling 与 Benchmarking

在作业的第一部分，我们将探索如何优化 Transformer 模型的性能，以最高效地利用 GPU。我们将对模型进行 profiling，以了解它在 forward 和 backward pass 期间把时间和内存花在哪里，然后用自定义 GPU kernels 优化 self-attention 操作，使其比常规 PyTorch 所能达到的更快。在作业的后续部分，我们将利用多块 GPU，并理解如何跨 cluster 训练模型。

2.1 Profiling

在实现任何优化之前，先对程序进行 profiling 会很有帮助，以了解它把资源（例如时间和内存）花在哪里。否则，我们可能会优化模型中那些并不占据显著时间或内存的部分，从而看不到可测量的端到端改进。

我们将实现三条性能评估路径：

  1. 使用 Python 标准库进行简单的端到端 benchmarking，为我们的 forward 和 backward pass 计时
  2. 使用 NVIDIA Nsight Systems 工具进行 compute profiling，以了解该时间如何在 CPU 和 GPU 上的操作之间分配
  3. Memory profiling
2.1.1 Setup——导入你的 basics Transformer 模型

首先，让我们确保你能加载之前的 assignment 中的模型。在之前的 assignment 中，我们把模型放在一个 Python package 中，以便之后可以轻松导入。我们在 ./cs336-basics 文件夹中添加了模型的 reference implementation，并在 pyproject.toml 文件中指向它。像往常一样调用 uv run [command] 时，uv 会自动定位这个本地 cs336-basics package。如果你想使用自己的模型实现，可以修改 pyproject.toml 文件，使其指向你自己的 package。

你可以用以下方式测试能否导入你的模型：

$ uv run python
Using CPython 3.13.13
Creating virtual environment at: /path/to/uv/env/dir
      Built cs336-systems @ file:///path/to/systems/dir
      Built cs336-basics @ file:///path/to/basics/dir
Installed 78 packages in 168ms
Python 3.13.13 (main, Apr  7 2026, 20:49:46) [Clang 22.1.1 ] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import cs336_basics
...

assignment 1 中的相关 modules 现在应该可用了（例如，对于 model.py，你可以用 import cs336_basics.model 导入它）。

2.1.2 Model Sizing

在整个作业中，我们将对模型进行 benchmarking 和 profiling，以更好地理解它们的性能。为了感受规模变化带来的影响，我们将使用并引用以下模型配置。对于除 leaderboard 之外的所有模型，我们将使用 10,000 的 vocabulary size 和 4 的 batch size，并使用不同的 context lengths。这份作业（以及后续作业）将需要把大量结果以表格和图表呈现。我们强烈建议你用代码自动构建表格，供

第 3 页

你的 writeup 使用，因为在 LaTeX 或 Typst 中排版表格会非常繁琐。参见 pandas.DataFrame.to_latex() 和 pandas.DataFrame.to_typst()，或者编写你自己的函数，从你偏好的表格表示生成它们。

Size d_model d_ff num_layers num_heads
small 768 3072 12 12
medium 1024 4096 24 16
large 1280 5120 36 20
xl 2560 10240 32 32
10B 4608 12288 50 36

Table 1：不同 model size 的规格。这些主要基于 GPT-2 configs。

除非另有说明，否则使用 context length 512。

2.1.3 End-to-End Benchmarking

我们现在将实现一个简单的性能评估脚本。我们将测试模型的许多变体（改变 precision、替换 layers 等），因此让你的脚本通过 command-line arguments 支持这些变体是值得的，这样以后运行起来会很容易。

首先，让我们通过对 forward pass、backward pass 和 optimizer step 计时，对我们的模型做尽可能简单的 profiling。由于我们只测量速度和内存，使用随机 weights 和 data 是没问题的。

测量性能很微妙——一些常见的陷阱会导致我们没有测量到我们想要测量的东西。对于 GPU 代码的 benchmarking，一个注意事项是 CUDA 调用是异步的。当你调用一个 CUDA kernel（例如调用 torch.matmul）时，PyTorch 函数调用会把控制权返回给你的代码，而不等待矩阵乘法完成。这样，CPU 可以在 GPU 完成矩阵乘法的同时继续超前运行并调度新的操作，这是一个巨大的性能收益。另一方面，这意味着天真地测量 torch.matmul 调用返回需要多长时间，并不能告诉我们 GPU 实际运行矩阵乘法需要多长时间。在 PyTorch 中，我们可以调用 torch.cuda.synchronize() 来等待所有已调度的 GPU kernels 完成，从而让我们获得更准确的 CUDA kernel runtime 测量。这个操作中的同步指的是让 CPU runtime 与 GPU runtime 同步。考虑到这一点，让我们编写基本的 profiling 基础设施。

Problem (benchmarking_script)：Benchmarking Script（4 分）

(a) 编写一个脚本，对你的模型中的 forward pass、backward pass 和 optimizer step 执行基本的端到端 benchmarking。具体来说，你的脚本应支持以下内容：

第 4 页

交付物： 一个脚本，它用给定的 hyperparameters 初始化一个 basics Transformer 模型，创建一个随机数据 batch，并对仅 forward、forward 加 backward，以及包含 optimizer step 的完整 training step 计时。

(b) 对 Section 2.1.2 中描述的 model sizes，分别为 forward、backward 和 optimizer step 计时。使用 5 个 warmup steps，并在 10 个 measurement steps 上计算 timings 的平均值和标准差。一次 forward pass 需要多长时间？一次 backward pass 呢？你在各次 measurement 之间看到很大的波动，还是标准差很小？

交付物： 一段 1-2 句话的回答，附上你的 timings。

(c) benchmarking 的一个 caveat 是不执行 warm-up steps。在不使用 warm-up steps 的情况下重复你的分析。这如何影响你的结果？你认为为什么会发生这种情况？还可以尝试用 1 个或 2 个 warm-up steps 运行该脚本。为什么结果仍可能不同？

交付物： 一段 2-3 句话的回答。

2.1.4 Nsight Systems Profiler

End-to-end benchmarking 不会告诉我们模型在 forward 和 backward passes 期间把时间和内存花在哪里，因此也不会暴露具体的优化机会。要了解程序在每个 component（例如 function）上花了多少时间，我们可以使用 profiler。execution profiler 通过在函数开始和结束运行时插入 guard 来对代码进行 instrumentation，因此能给出 function 级别的详细执行统计（例如调用次数、平均耗时、在该 function 上累计花费的时间等）。

标准 Python profilers（例如 CProfile）无法对 CUDA kernels 做 profile，因为这些 kernels 是在 GPU 上异步执行的。幸运的是，NVIDIA 提供了一个我们可以通过 CLI nsys 使用的 profiler。我们建议你从包管理器获取最新版本，或使用其下载页面上的安装程序。在作业的这一部分，你将使用 nsys 分析你的 Transformer 模型的 runtime。

使用 nsys 很直接：在上一节的 Python 脚本前加上 nsys profile 来运行它。例如，你可以用以下命令为脚本 benchmark.py 运行一个基本 profile：

$ uv run nsys profile -- python benchmark.py

然后你可以在本地机器上使用 NVIDIA Nsight Systems desktop application 查看该 profile。在 profile 的 CUDA API 行中选择某个 CUDA API call（在 CPU 上），会高亮显示 CUDA HW 行中所有对应的 kernel executions（在 GPU 上）。

更完整的 profiling 运行可能如下所示：

$ uv run nsys profile --trace=cuda,cudnn,cublas,osrt,nvtx --pytorch=functions-trace,autograd-shapes-nvtx --cudabacktrace=all --python-backtrace=cuda --gpu-metrics-devices=0 -- python benchmark.py

第 5 页

图 1：一份详细的 Nsight Systems trace

在这个例子中，--trace 指定要记录哪些 APIs，--pytorch 在 module calls 和 autograd 期间插入 nvtx labels，--cudabacktrace 和 --python-backtrace 提供更好的 backtraces，以便理解某个 kernel 是从你代码中的哪里被调用的，而 --gpu-metrics-devices 指定要测量哪块 GPU 的 utilization。

给一次运行加上 profiling 不是没有代价的，并且总体上会拖慢你的运行。通常只启用你在某次运行中真正想看的 features 是值得的。具体来说，当不需要 tracebacks 时，你可能会想移除 --cudabacktrace=all 和 --python-backtrace=cuda，因为它们有格外大的 overhead。

我们鼓励你试验 nsys profile 的各种 command-line options，以了解它能做什么。你也可以用 NVTX ranges 标注你的代码，这些 ranges 会在 profile 的 NVTX 行中显示为 blocks，捕获所有 CUDA API calls 以及相关的 kernel executions。特别是，你应该使用 NVTX ranges 来忽略 benchmarking 脚本中的 warm-up steps（通过在 profile 中对 nvtx label 应用 --nvtx-capture filter）。你还可以隔离出哪些 kernels 负责你模型的 forward 和 backward passes，甚至可以通过如下标注你的实现，隔离出哪些 kernels 负责 self-attention layer 的不同部分：

...
import torch.cuda.nvtx as nvtx

@nvtx.range("scaled dot product attention")
def annotated_scaled_dot_product_attention(
    ... # Q, K, V, mask
):
    ...
    with nvtx.range("computing attention scores"):
        ... # compute attention scores between Q and K
    with nvtx.range("computing softmax"):
        ... # compute softmax of attention scores
    with nvtx.range("final matmul"):
        ... # compute output projection

第 6 页

return ...

你可以通过以下方式，在 benchmarking 脚本中把原来的实现替换成带标注的版本：

cs336_basics.model.scaled_dot_product_attention = annotated_scaled_dot_product_attention

最后值得注意，torch.compile 会让把时间和资源归因到代码的特定部分变得困难。你可能需要在 torch.compile 和 nvtx annotations 中包裹并剥离代码的各个部分，才能正确地把时间和资源使用归因到源码的各个部分。

Problem (nsys_profile): Nsight Systems Profiling（5 分）

使用 nsys 对你的 forward pass、backward pass 和 optimizer step 进行 profile：从 Table 1 中任选两种 model sizes，并选三个大于 128 的 2 的幂次 context lengths，其中最大的可用尺寸应是你能够在内存中容纳的最长 context length。挑选你认为最值得观察的组合。对每个 profile 回答以下问题：

(a) 你的 forward pass 总共花了多少时间？它是否与之前用 Python 标准库测得的数值一致？

交付物： 一段 1-2 句话的回答。

(b) 在 forward pass 期间，哪个 CUDA kernel 占用的累计 GPU 时间最多？在你的模型的一次 forward pass 中，这个 kernel 被调用了多少次？当你同时执行 forward 和 backward passes 时，耗时最多的还是同一个 kernel 吗？（提示：查看 “Stats System View” 下的 “CUDA GPU Kernel Summary”，并用 NVTX ranges 过滤，以分辨模型的哪些部分对应哪些 kernels。）

交付物： 一段 1-2 句话的回答。

(c) 尽管绝大多数 FLOPs 发生在 matrix multiplications 中，你会注意到仍有若干其他 kernels 在总 runtime 中占据不可忽略的比例。除了 matrix multiplication 之外，在 forward pass 中你还看到哪些 kernels 产生了不可忽略的 CUDA runtime？

交付物： 一段 1-2 句话的回答。

(d) 对你的 AdamW 实现完整运行一个 training step 进行 profile（即像训练时那样：forward pass、计算 loss 并运行 backward pass，最后执行 optimizer step）。与做 inference（仅 forward pass）相比，花在 matrix multiplication 上的时间占比如何变化？其他 kernels 呢？

交付物： 一段 1-2 句话的回答。

(e) 比较在一次 forward pass 中，你的模型 self-attention layer 内 softmax operation 与 matrix multiplication operations 的 runtime。runtime 上的差异与 FLOPs 上的差异相比如何？

交付物： 一段 1-2 句话的回答。

第 7 页

2.1.5 Mixed Precision（混合精度）

在本作业中，到目前为止我们一直使用 FP32 精度——所有模型参数和 activations 都是 torch.float32 数据类型。然而，现代 NVIDIA GPU 包含专门的 GPU 核心（Tensor Cores），用于在较低精度下加速矩阵乘法。例如，NVIDIA B200 的规格表称，其 FP32 最大吞吐量为 80 TFLOPS，而 FP16（半精度浮点数）或 BF16（bfloat16）的最大吞吐量则显著更高，达到惊人的 2500 TFLOPS。因此，使用较低精度的数据类型应有助于加速训练和推理。

然而，直接（naïvely）将模型转换为较低精度格式，可能带来模型精度下降。例如，实际中许多 gradient 值往往太小，无法在 FP16 中表示，因此在直接用 FP16 精度训练时会变成零。为应对这一点，使用 FP16 训练时常见做法是使用 loss scaling——只需将 loss 乘以一个缩放因子，增大 gradient 的数值幅度，使其不会下溢为零（flush to zero）。此外，FP16 的动态范围比 FP32 低，这可能导致溢出，并表现为 NaN loss。完全使用 bfloat16 训练通常更稳定（因为 BF16 与 FP32 具有相同的动态范围），但与 FP32 相比仍可能影响最终模型性能。

为了利用较低精度数据类型带来的加速，常见做法是使用混合精度训练。在 PyTorch 中，这通过 torch.autocast 上下文管理器实现。在这种情况下，某些操作（例如矩阵乘法）以较低精度数据类型执行，而其他需要 FP32 完整动态范围的操作（例如累加和归约）则保持原样。例如，以下代码会在前向传播中自动识别哪些操作应以较低精度执行，并将这些操作转换为指定的数据类型：

model : torch.nn.Module = ... # e.g. your Transformer model
dtype : torch.dtype = ... # e.g. torch.bfloat16
x : torch.Tensor = ... # input data
with torch.autocast(device_type="cuda", dtype=dtype):
    y = model(x)

正如上文所暗示的，即使被累加的 tensor 本身已经被降精度，通常也最好将累加保持在较高精度。下面的练习将帮助你建立直觉，理解为什么会这样。

Problem (mixed_precision_accumulation): 混合精度累加（1 分）

运行以下代码，并评论结果的准确性。

s = torch.tensor(0,dtype=torch.float32)
for i in range(1000):
    s += torch.tensor(0.01,dtype=torch.float32)
print(s)
s = torch.tensor(0,dtype=torch.float16)
for i in range(1000):
    s += torch.tensor(0.01,dtype=torch.float16)
print(s)
s = torch.tensor(0,dtype=torch.float32)
for i in range(1000):

第 8 页

    s += torch.tensor(0.01,dtype=torch.float16)
    print(s)
    s = torch.tensor(0,dtype=torch.float32)
    for i in range(1000):
        x = torch.tensor(0.01,dtype=torch.float16)
        s += x.type(torch.float32)
    print(s)

Deliverable: 2–3 句话的回答。

现在，我们会先将混合精度应用到一个 toy model 上以建立直觉，然后再应用到我们的 benchmarking 脚本中。

Problem (benchmarking_mixed_precision): 混合精度基准测试（2 分）

(a) 考虑以下模型：

class ToyModel(nn.Module):
    def __init__(self, in_features: int, out_features: int):
        super().__init__()
        self.fc1 = nn.Linear(in_features, 10, bias=False)
        self.ln = nn.LayerNorm(10)
        self.fc2 = nn.Linear(10, out_features, bias=False)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.ln(x)
        x = self.fc2(x)
        return x

假设我们在 GPU 上训练该模型，且模型参数最初为 FP32。我们希望使用 FP16 的 autocasting 混合精度。以下各项的数据类型是什么：

Deliverable: 上述每个组件的数据类型。

(b) 你应该已经看到，FP16 混合精度 autocasting 对 layer normalization 层的处理与对前馈层不同。layer normalization 的哪些部分对混合精度敏感？如果我们使用 BF16 而不是 FP16，是否仍然需要区别对待 layer normalization？为什么需要或为什么不需要？

Deliverable: 2–3 句话的回答。

第 9 页

(c) 修改你的 benchmarking 脚本，使其可以选择使用 BF16 的混合精度来运行模型。对 Section 2.1.2 中描述的每种 language model 大小，分别对使用和不使用混合精度的前向和后向传播计时。比较使用完整精度与混合精度的结果，并评论随着模型大小变化出现的任何趋势。你可能会发现 nullcontext 这个 no-op 上下文管理器很有用。

Deliverable: 2–3 句话的回答，附上你的计时结果和评论。

2.1.6 Profiling Memory（内存分析）

到目前为止，我们一直在关注计算性能。现在我们将把注意力转向内存——语言模型训练和推理中的另一项主要资源。PyTorch 还附带一个强大的 memory profiler，它可以跟踪随时间推移的内存分配。

要使用 memory profiler，你可以按如下方式修改 benchmarking 脚本：

... # warm-up phase in your benchmarking script
# Start recording memory history.
torch.cuda.memory._record_memory_history(max_entries=1000000)
... # what you want to profile in your benchmarking script
# Save a pickle file to be loaded by PyTorch's online tool.
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")
# Stop recording history.
torch.cuda.memory._record_memory_history(enabled=None)

这会输出一个 memory_snapshot.pickle 文件，你可以将其加载到以下在线工具中：pytorch.org/memory_viz。该工具可让你查看总体内存使用时间线，以及每一次单独的内存分配——包括其大小，以及指向其来源代码的 stack trace。要使用该工具，你应在网页浏览器中打开上面的链接，然后将你的 pickle 文件拖放到页面上。

你现在将使用 PyTorch profiler 来分析模型的内存使用情况。

Problem (memory_profiling): 内存分析（4 分）

对 Table 1 中的 xl 模型，在 context length 为 128 和 2048 时，对其完整训练步骤（前向传播、后向传播和 optimizer step）进行 profiling。

(a) 为你的 profiling 脚本添加一个选项，使模型可通过 memory profiler 运行。复用一些之前的 infrastructure（例如，用于启用混合精度、加载特定模型大小等）可能会有帮助。然后运行脚本，获取 xl 模型在仅进行推理（仅前向传播）或执行完整训练步骤时的 memory profile。你的内存时间线看起来是什么样的？你能根据看到的峰值判断当前正在运行哪个阶段吗？

Deliverable: 两张来自 memory_viz 工具的 xl 模型 “Active memory timeline” 图像：一张对应前向传播，一张对应运行完整训练步骤（前向和后向传播，然后是 optimizer step），以及 2–3 句话的回答。

第 10 页

(b) 做一次 forward pass 时，每个 context length 的峰值内存占用是多少？做一次完整训练 step 时呢？ Deliverable: 一张表，每个 context length 对应两个数字。

(c) 找出 xl 模型在使用 mixed-precision 时，forward pass 和完整训练 step 的峰值内存占用。mixed-precision 会显著影响内存占用吗？ Deliverable: 2–3 句话的回答。

(d) 考虑 xl 模型。在给定的参考超参数下，Transformer residual stream 中一个 activation tensor 在 single-precision 下的大小是多少？请用 MiB 表示该大小（即用字节数除以 $1024^2$）。 Deliverable: 1–2 句话的回答，并给出你的推导。

(e) 现在仔细查看 pytorch.org/memory_viz 中 xl 模型做 forward pass 的内存快照的 “Active Memory Timeline”。当你降低 “Detail” 级别时，该工具会隐藏小于对应级别的最小分配（例如，把 “Detail” 设为 10% 只显示最大的 10% allocations）。显示的 largest allocations 的大小是多少？查看 stack trace，你能看出这些 allocations 来自哪里吗？ Deliverable: 1–2 句话的回答。

(f) Nsight Systems 也有用于 memory profiling 的 flags。你可以把这些 flags 与之前的 Nsight flags 结合使用，以理解模型生命周期中不同步骤正在发生哪些 allocations。使用 PyTorch 提供的 NVTX labels，确定你模型中的单个 TransformerBlock 为 backward 保存了多少内存（这些 tensors 通常称为 residuals）。记下贡献最大的 5 个 operations，以及它们对总内存的贡献百分比。

在 backward pass 期间，所有这些 tensors 都会被释放，但同时会生成新的 gradient tensors。根据你的 profiles，显示 forward pass 期间分配了多少内存，以及 backward pass 中每个 TransformerBlock 的内存使用变化量，计算一个 TransformerBlock 产生的 gradient tensors 占用多少内存。结果符合你的预期吗？ Deliverable: 来自 Nsight Systems 的截图，以及 1–2 段回答。

3 Single-GPU Memory

本作业后面的部分会探讨把 tensors 分片（shard）到多个 GPUs 上的技巧，但也有一些技巧即使用于 single-GPU training 也能应用。其中最常见的是 gradient checkpointing（也称为 activation checkpointing）。

3.1 Autograd Residuals

回想一下，为了在你的模型中执行 backward pass，我们需要保存 forward pass 中产生的 activations。虽然对某些 operations 来说这显然是必要的，但默认情况下，发生这种情况的范围会比你预想的多得多。为 backward pass 保存的 tensors 被称为 “residuals”，或者简称为 “saved tensors”。

让我们先建立一些理解：网络中到底保存了什么。从我们不起眼的 RMSNorm 函数开始（为简单起见，使用纯 FP32），让我们添加一些 hooks，用于在 autograd 保存或取回 tensors 时触发。

第 11 页

import torch
from torch import nn
x = torch.randn((4, 512, 2560), requires_grad=True)
class RMSNorm(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        eps: float = 1e-5,
        device=None,
    ):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size, device=device))
        self.eps = eps
    def forward(self, x):
        rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        x = x * rms
        return self.weight * x
def pack_hook(t):
    shape, dtype, grad_fn = t.shape, t.dtype, t.grad_fn
    print(f"Saving residual: {shape=}, {dtype=}, {grad_fn=}")
    return t
def unpack_hook(t):
    shape, dtype, grad_fn = t.shape, t.dtype, t.grad_fn
    print(f"Loading residual: {shape=}, {dtype=}, {grad_fn=}")
    return t
ln = RMSNorm(x.shape[-1])
with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
    y = ln(x)
    y.sum().backward()

输出显示写出的 tensors 数量多得令人担忧，其中好几个都是完整 activation 大小！

$ uv run scripts/autograd_experiment.py
Saving residual: shape=torch.Size([4, 512, 2560]), dtype=torch.float32, grad_fn=None
Saving residual: shape=torch.Size([4, 512, 1]), dtype=torch.float32, grad_fn=<RsqrtBackward0 object at 0x7f7dd319b5e0>
Saving residual: shape=torch.Size([4, 512, 1]), dtype=torch.float32, grad_fn=<RsqrtBackward0 object at 0x7f7dd319b5e0>
Saving residual: shape=torch.Size([4, 512, 2560]), dtype=torch.float32, grad_fn=None
Saving residual: shape=torch.Size([4, 512, 2560]), dtype=torch.float32, grad_fn=<MulBackward0 object at 0x7f7dd319b5e0>
Saving residual: shape=torch.Size([2560]), dtype=torch.float32, grad_fn=None
Loading residual: shape=torch.Size([4, 512, 2560]), dtype=torch.float32, grad_fn=<MulBackward0 object at 0x7f7cf14e6740>

第 12 页

Loading residual: shape=torch.Size([2560]), dtype=torch.float32, grad_fn=None
Loading residual: shape=torch.Size([4, 512, 1]), dtype=torch.float32, grad_fn=<RsqrtBackward0 object at 0x7f7cf14e6740>
Loading residual: shape=torch.Size([4, 512, 2560]), dtype=torch.float32, grad_fn=None
Loading residual: shape=torch.Size([4, 512, 1]), dtype=torch.float32, grad_fn=<RsqrtBackward0 object at 0x7f7cf14e6740>
Loading residual: shape=torch.Size([4, 512, 2560]), dtype=torch.float32, grad_fn=None
3.1.1 Operator Fusion

在这个例子中，很明显所用 operations 的粒度太细了。我们希望有一个单独的 op，接收 RMSNorm 的 weights 和 activation，直接输出结果，并且该 operation 在 backward pass 中作为单一操作处理。这就是 kernel fusion 的一个动机。由于 RMSNorm 的行为相当规整，我们甚至可以用 torch.compile 自动融合它。

...
ln = torch.compile(RMSNorm(x.shape[-1]))
with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
    y = ln(x)
    y.sum().backward()

新的输出明显更好：

Saving residual: shape=torch.Size([4, 512, 2560]), dtype=torch.float32, grad_fn=None
Saving residual: shape=torch.Size([2560]), dtype=torch.float32, grad_fn=None
Saving residual: shape=torch.Size([4, 512, 1]), dtype=torch.float32, grad_fn=None
Loading residual: shape=torch.Size([4, 512, 2560]), dtype=torch.float32, grad_fn=None
Loading residual: shape=torch.Size([2560]), dtype=torch.float32, grad_fn=None
Loading residual: shape=torch.Size([4, 512, 1]), dtype=torch.float32, grad_fn=None

我们只需要为 backward pass 保存一个完整大小的 activation tensor——即 RMSNorm 函数的输入。还要注意，加载顺序不再是保存顺序的逆序，而且每个 residual 不再有 grad_fn 依赖——PyTorch 正把我们的整个 RMSNorm 当作单个函数来处理。

3.2 Activation Checkpointing

虽然 fusion 无疑很有用，但它在节省内存方面能做的有限。例如，让我们把 size xl 下的单个 TransformerBlock 融合起来。

import torch
from cs336_basics.model import RotaryEmbedding, TransformerBlock
# num_layers for this model is 32
d_model, d_ff, num_heads, context_length = 2560, 10240, 16, 2048
block = TransformerBlock(d_model=d_model, d_ff=d_ff, num_heads=num_heads,
                         positional_encoder=RotaryEmbedding(dim=d_model // num_heads, context_length=context_length))
# Fuse as much torch.compile will allow
block = torch.compile(block, fullgraph=True)
x = torch.randn((4, context_length, d_model), requires_grad=True)

第 13 页

...
# Now logs the number of bytes saved
total_size_bytes = 0
def pack_hook(t):
    if isinstance(t, torch.nn.Parameter):  # Skip logging parameters to avoid double counting
        return t
    global total_size_bytes
    shape, dtype, grad_fn = t.shape, t.dtype, t.grad_fn
    total_size_bytes += t.numel() * t.element_size()
    print(f"Saving residual: {shape=}, {dtype=}, {grad_fn=}")
    return t
...
# Run forward pass, saving for backward
with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
    y = block(x)
print(f"Total size of saved tensors in single TransformerBlock: {total_size_bytes / (1024**2):.2f} MiB")

这个脚本显示了我们在为 backward 保存多少内存：

Total size of saved tensors in single TransformerBlock: 3651.31 MiB

每一层就是 3.6 GiB。如果对所有层都这样做，我们仅为了 backward 保存的 activation 就会达到 114 GiB！attention 运算的 residuals 中存在相当可观的浪费，我们会在 Section 4 修复；但即使有了这个修复，内存使用仍会随 batch size、sequence length 和 embedding size 线性增长。

3.2.1 Recomputation

与其保留我们生成的每一个 tensor，不如只定期保存结果的 checkpoint，并在中间重计算这些值。PyTorch 有一个可以调用的接口，能以简单的方式处理这件事。torch.utils.checkpoint.checkpoint 接收一个函数，以及传给该函数的参数。然后它通过以下方式修改传入函数的行为：

  1. 在 forward pass 中：
  2. 保存传给该函数的输入值
  3. 抑制 forward pass 中 tensor 的保存
  4. 在 backward pass 中：
  5. 在前面加入一个重计算步骤：用之前保存的输入重新执行 forward pass，并为 backward 保存值
  6. 运行 backward pass，然后所有 tensor 都可以被释放

在依次运行 4 个 transformer block 的简单情形中，我们看到内存会如预期般累加。

...
def four_blocks(x):
    x = block(x)

第 14 页

    x = block(x)
    x = block(x)
    x = block(x)
    return x

with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
    y = four_blocks(x)

print(f"Total size of saved tensors in four TransformerBlocks: {total_size_bytes / (1024**2):.2f} MiB")

Total size of saved tensors in four TransformerBlocks: 14605.25 MiB

但我们可以像下面这样使用 gradient checkpointing：

from torch.utils.checkpoint import checkpoint

def two_blocks(x):
    x = block(x)
    x = block(x)
    return x

def four_blocks_checkpoint(x):
    # checkpoint throws out all the saved tensors until the backward pass
    # when getting to the checkpointed block in the backward pass,
    # it reruns a forward pass to produce the saved tensors,
    # then completes normal backward pass.
    x = checkpoint(two_blocks, x, use_reentrant=False)
    x = checkpoint(two_blocks, x, use_reentrant=False)
    return x

with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
    y = four_blocks_checkpoint(x)

print(f"Total size of saved tensors in four TransformerBlocks with checkpointing: {total_size_bytes / (1024**2):.2f} MiB")

Saving residual: shape=torch.Size([0]), dtype=torch.float32, grad_fn=None
Saving residual: shape=torch.Size([4, 2048, 2560]), dtype=torch.float32, grad_fn=None
Saving residual: shape=torch.Size([0]), dtype=torch.float32, grad_fn=None
Saving residual: shape=torch.Size([4, 2048, 2560]), dtype=torch.float32, grad_fn=<torch.autograd.function.CompiledFunctionBackward object at 0x7aa0657a19d0>
Total size of saved tensors in four TransformerBlocks with checkpointing: 160.00 MiB

请记住，这并没有消除内存使用。更准确地说，它把我们的内存使用分成了两类：长期存储，即我们为在每个 checkpoint 调用入口处准备重计算而保存的内容（checkpoint 本身）；以及短期内存，即在 checkpointed block 内执行重计算时产生的内存，用于支持穿过该 block 的 backward pass。由于我们的主要考虑是 peak memory usage，我们希望平衡已保存 checkpoint 的内存开销与物化（materializing）一整个 block 的 residuals 所需的内存开销。我们使用的 checkpoint 越多，单个 block 内物化的内存量就越小，但 checkpoint 本身需要的内存就越多。

第 15 页

在上面的例子中，很明显重计算 residuals 的内存占主导（部分原因是我们的 checkpoint 很小），因此把平衡更多地转向 checkpoint 内存会是有益的。于是我们会希望减小被 checkpoint 的 block 的范围。

被 checkpoint 的 block 更大或更小并不会影响重计算的计算开销，但我们可以通过 recursive checkpointing 继续以计算为代价降低所需内存；也就是说，把 checkpoint 调用嵌套在其他 checkpoint 调用内部。

Problem (gradient_checkpointing)：Memory-Optimal Gradient Checkpointing（内存最优梯度检查点，4 分）

考虑一个由 $N$ 个相同 block 顺序堆叠而成的 Transformer。在没有任何 checkpointing 的情况下，全部 $N$ 个 block 的 residuals 会同时保持存活，带来 $O(N)$ 的峰值激活内存（peak activation memory）。我们可以自由地把 forward pass 的任意子集包裹在 checkpoint 中，包括把 checkpoint 调用相互嵌套。

(a) 在忽略计算开销的情况下，哪种 checkpointing 策略能最小化峰值激活内存？描述你会如何安排 checkpoint 调用（给一个 code sketch 即可），并给出你的策略的渐近峰值激活内存和计算量（作为 $N$ 的函数）。假设单个 block 保存的 residuals 支配任何 per-checkpoint 的簿记开销。

Deliverable: 用 3-5 句话描述该策略及其渐近峰值内存，并附上一段简短的 code sketch。

(b) 考虑如上所述的 xl model config，batch size 为 4、sequence length 为 2048。如果你只有时间/计算预算运行一步重计算（意味着你不能嵌套 checkpoint 调用），那么降低峰值内存的最佳 checkpointing 策略是什么？对你的运行进行 profile，测出峰值内存以验证你的假设。比较相邻更小和更大的 checkpointing block size 的峰值内存，以确保判断正确。

Deliverable: 用 3-5 句话描述你的推理，并给出你的策略测得的峰值内存。

4 GPU Kernels

4.1 Optimizing Attention with FlashAttention-2

4.1.1 Benchmarking PyTorch Attention

你的 profiling 很可能表明，你的 attention 层在内存和计算两方面都有优化机会。从高层来看，attention 运算由一个矩阵乘法、随后 softmax、然后再一个矩阵乘法组成：

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

naïve attention 实现需要为每个 batch/head 元素保存形状为 $\text{seq\_len} \times \text{seq\_len}$ 的 attention score 矩阵；随着 sequence length 变长，这个矩阵会变得非常大，从而导致任何输入或输出较长的任务出现 out-of-memory 错误。我们将按照 FlashAttention-2 论文实现一个 attention kernel，它按 tile 计算 attention，并避免显式地物化（materialize）$\text{seq\_len} \times \text{seq\_len}$ 的 attention score 矩阵，从而能够扩展到长得多的 sequence length。

第 16 页

Problem（pytorch_attention）：PyTorch Attention Benchmarking（2 分）

(a) 在不同规模下对你的 attention 实现进行基准测试。编写一个脚本，该脚本将：

根据你的 GPU，其中一些配置预计会遇到 out-of-memory。报告你针对这些配置得到的计时（或 out-of-memory 错误）。在什么规模下你会遇到 out-of-memory 错误？对你发现的某个最小的、会耗尽内存的配置，计算 attention 的内存使用量（你可以使用 Assignment 1 中 Transformers 的内存使用公式）。为 backward 保存的内存如何随序列长度变化？你会怎样消除这项内存开销？

Deliverable: 一张包含你的计时、内存使用量计算，以及 1-2 段回答的表格。

4.2 Benchmarking JIT-Compiled Attention（JIT 编译 Attention 的基准测试）

自 2.0 版本起，PyTorch 还附带一个强大的 just-in-time compiler，它会自动尝试对 PyTorch 函数应用若干优化：入门介绍见 https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html。特别地，它会通过动态分析你的 computation graph，尝试自动生成 fused Triton kernels。使用 PyTorch compiler 的接口非常简单。例如，如果我们想把它应用到模型中的单层，可以这样用：

layer = SomePyTorchModule(...)
compiled_layer = torch.compile(layer)

现在，compiled_layer 在功能上就像 layer 一样（例如，其 forward 和 backward passes）。我们也可以使用 torch.compile(model) 编译整个 PyTorch 模型，甚至编译一个调用 PyTorch 操作的 Python 函数。

Problem（torch_compile）：Torch Compile（2 分）

(a) 扩展你的 attention benchmarking 脚本，加入你的 PyTorch attention 实现的编译版本，并在与上面的 pytorch_attention 问题相同的配置下，将其性能与未编译版本进行比较。

第 17 页

Deliverable: 一张表格，比较你的 compiled attention module 与上面 pytorch_attention 问题中未编译版本的前向和反向传播计时。

(b) 现在，在你的端到端 benchmarking 脚本中编译整个 Transformer 模型。前向传播的性能会如何变化？前向与反向传播以及 optimizer steps 合计起来又如何？

Deliverable: 一张比较 vanilla 与 compiled Transformer 模型的表格。

鉴于我们在序列长度方面看到的扩展行为，我们需要显著改进才能处理长序列。即使使用 torch.compile，当前实现在长序列长度下仍然受困于非常差的内存访问模式。为此，我们将编写 FlashAttention-2 的 Triton 实现，在其中我们将对如何访问内存以及何时计算什么拥有明显更多的控制。

4.2.1 Example - Weighted Sum（示例——加权求和）

为了介绍你需要了解的 Triton 知识以及它如何与 PyTorch 互操作，我们将讲解一个用于“weighted sum”（加权求和）操作的示例 kernel。要进一步上手 Triton，参见 Triton 的 tutorials。我们注意到这些 tutorials 没有使用新的、方便的 block pointer 抽象，而我们将在下面逐步讲解它。

给定输入矩阵 $X$，我们将其元素乘以按列组织的权重向量 $w$，并对每一行求和，从而得到 $X$ 与 $w$ 的矩阵-向量乘积。我们将先讲解该操作的前向传播，然后为反向传播编写 Triton kernel。

Forward pass（前向传播）

我们 kernel 的前向传播就是下面这个广播后的内积。

def weighted_sum(x, weight):
    # Here, assume that x has n-dim shape [..., D], and weight has 1D shape [D]
    return (weight * x).sum(axis=-1)

在编写 Triton kernel 时，我们会让每个 program instance（可能并行运行）计算 $x$ 的一个行 tile 的 weighted sum，并将对应的标量输出写入输出张量。在 Triton 中，一个 program instance 是一组运行同一程序的线程块，这些线程块可以在 GPU 上并行运行。我们不是把张量作为参数传入，而是传入指向其首元素的指针，以及每个张量的 strides，用来告诉我们应该如何沿各轴移动。

我们可以使用 strides 加载与当前运行 instance 中正在求和的 $x$ 的行 tile 相对应的张量，并使用 program ID 来划分工作（即 instance $i$ 将处理 $x$ 的第 $i$ 个行 tile）。在这个简单情形下，Triton 前向传播与 PyTorch 前向传播的主要区别在于需要做 pointer arithmetic 以及显式的 loads/stores。我们将使用 tl.make_block_ptr 的 block pointer 抽象来大大简化 pointer arithmetic，不过这意味着我们需要做一些设置来准备 block pointers。

第 18 页

图中标签：D、ROWS、ROWS_TILE、D_TILE、.advance((ROWS_TILE, 0))、.advance((0, D_TILE))。

Figure 2（图 2）： weighted sum kernel 示例（第 4.2.1 节）中的 tiling 与 block pointers 推进。

关于 tiling 以及 block pointers 如何推进的示意图，参见 Figure 2。上面的 weighted sum 函数看起来如下：

import triton
import triton.language as tl
@triton.jit
def weighted_sum_fwd(
    x_ptr, weight_ptr,  # Input pointers
    output_ptr,  # Output pointer
    x_stride_row, x_stride_dim,  # Strides tell us how to move one element in each axis of a tensor
    weight_stride_dim,  # Likely 1
    output_stride_row,  # Likely 1
    NUM_ROWS, D,
    ROWS_TILE_SIZE: tl.constexpr, D_TILE_SIZE: tl.constexpr,  # Tile shapes must be known at compile time
):
    # Each instance will compute the weighted sum of a tile of rows of x.
    # `tl.program_id` gives us a way to check which thread block we're running in
    row_tile_idx = tl.program_id(0)
    # Block pointers give us a way to select from an ND region of memory
    #  and move our selection around.
    # The block pointer must know:
    # - The pointer to the first element of the tensor
    # - The overall shape of the tensor to handle out-of-bounds access
    # - The strides of each dimension to use the memory layout properly
    # - The ND coordinates of the starting block, i.e., "offsets"
    # - The block shape to load/store at a time
    # - The order of the dimensions in memory from major to minor
    #    axes (= np.argsort(strides)) for optimizations, needed for
    #    TMA support on >=Hopper

第 19 页

    x_block_ptr = tl.make_block_ptr(
        x_ptr,
        shape=(NUM_ROWS, D,),
        strides=(x_stride_row, x_stride_dim),
        offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),
        block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
        order=(1, 0),
    )
    weight_block_ptr = tl.make_block_ptr(
        weight_ptr,
        shape=(D,),
        strides=(weight_stride_dim,),
        offsets=(0,),
        block_shape=(D_TILE_SIZE,),
        order=(0,),
    )
    output_block_ptr = tl.make_block_ptr(
        output_ptr,
        shape=(NUM_ROWS,),
        strides=(output_stride_row,),
        offsets=(row_tile_idx * ROWS_TILE_SIZE,),
        block_shape=(ROWS_TILE_SIZE,),
        order=(0,),
    )
    # Initialize a buffer to write to
    output = tl.zeros((ROWS_TILE_SIZE,), dtype=tl.float32)
    for i in range(tl.cdiv(D, D_TILE_SIZE)):
        # Load the current block pointer
        # Since ROWS_TILE_SIZE might not divide NUM_ROWS, and D_TILE_SIZE might not divide D,
        # we need boundary checks for both dimensions
        row = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero")  # (ROWS_TILE_SIZE, D_TILE_SIZE)
        weight = tl.load(weight_block_ptr, boundary_check=(0,), padding_option="zero")  # (D_TILE_SIZE,)
        # Compute the weighted sum of the row.
        output += tl.sum(row * weight[None, :], axis=1)
        # Move the pointers to the next tile.
        # These are (rows, columns) coordinate deltas
        x_block_ptr = x_block_ptr.advance((0, D_TILE_SIZE))  # Move by D_TILE_SIZE in the last dimension
        weight_block_ptr = weight_block_ptr.advance((D_TILE_SIZE,))  # Move by D_TILE_SIZE
    # Write output to the output block pointer (a single scalar per row).
    # Since ROWS_TILE_SIZE might not divide NUM_ROWS, we need boundary checks
    tl.store(output_block_ptr, output, boundary_check=(0,))

现在我们把这个 kernel 包装进一个 PyTorch Autograd function，使其能与 PyTorch 互操作（即接收 Tensors 作为输入、输出一个 Tensor，并且在 backward pass 中还能与 autograd engine 配合工作）：

class WeightedSumFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight):
        # Cache x and weight to be used in the backward pass, when we
        # only receive the gradient wrt. the output tensor, and
        # need to compute the gradients wrt. x and weight.
        D, output_dims = x.shape[-1], x.shape[:-1]
        # Reshape input tensor to 2D
        input_shape = x.shape
        x = rearrange(x, "... d -> (...) d")
        ctx.save_for_backward(x, weight)

第 20 页

        assert len(weight.shape) == 1 and weight.shape[0] == D, "Dimension mismatch"
        assert x.is_cuda and weight.is_cuda, "Expected CUDA tensors"
        assert x.is_contiguous(), "Our pointer arithmetic will assume contiguous x"
        ctx.D_TILE_SIZE = triton.next_power_of_2(D) // 16  # Roughly 16 loops through the embedding dimension
        ctx.ROWS_TILE_SIZE = 16  # Each thread processes 16 batch elements at a time
        ctx.input_shape = input_shape
        # Need to initialize empty result tensor. Note that these elements are not necessarily 0!
        y = torch.empty(output_dims, device=x.device)
        # Launch our kernel with n instances in our 1D grid.
        n_rows = y.numel()
        weighted_sum_fwd[(triton.cdiv(n_rows, ctx.ROWS_TILE_SIZE),)](
            x, weight,
            y,
            x.stride(0), x.stride(1),
            weight.stride(0),
            y.stride(0),
            NUM_ROWS=n_rows, D=D,
            ROWS_TILE_SIZE=ctx.ROWS_TILE_SIZE, D_TILE_SIZE=ctx.D_TILE_SIZE,
        )
        return y.view(input_shape[:-1])

注意，当我们用 weighted_sum_fwd[(triton.cdiv(n_rows, ctx.ROWS_TILE_SIZE),)] 调用 Triton kernel 时，我们通过传入元组 (triton.cdiv(n_rows, ctx.ROWS_TILE_SIZE),) 定义所谓的线程块 “launch grid”。然后，我们可以在 kernel 中用 tl.program_id(0) 访问线程块索引。

反向传播（Backward pass）

因为我们在定义自己的 kernel，所以也需要编写自己的 backward 函数。

在 forward pass 中，我们拿到层的输入，并需要计算它的输出。在 backward pass 中，回忆一下：我们会拿到目标函数关于我们输出的梯度，并需要计算关于每个输入的梯度。在我们的情形中，操作的输入是一个矩阵 $x: \mathbb{R}^{n \times h}$ 和一个权重向量 $w: \mathbb{R}^{h}$。为简洁起见，把我们的操作记作 $f(x,w)$，其值域为 $\mathbb{R}^{n}$。那么，假设给定 $\nabla_{f(x,w)}\mathcal{L}$，即损失 $\mathcal{L}$ 关于我们层输出的梯度，我们可以应用多元链式法则，得到关于 $x$ 和 $w$ 的梯度的如下表达式：

$$(\nabla_x \mathcal{L})_{ij} = \sum_{k=1}^{n} \frac{\partial f(x,w)_k}{\partial x_{ij}} (\nabla_{f(x,w)}\mathcal{L})_k = w_j \cdot (\nabla_{f(x,w)}\mathcal{L})_i \tag{2}$$

$$(\nabla_w \mathcal{L})_j = \sum_{i=1}^{n} \frac{\partial f(x,w)_i}{\partial w_j} (\nabla_{f(x,w)}\mathcal{L})_i = \sum_{i=1}^{n} x_{ij} \cdot (\nabla_{f(x,w)}\mathcal{L})_i \tag{3}$$

这给出了计算 backward pass 的简单公式。为了得到关于 $x$ 的 backward step，我们应用 Equation 2，并取 $w$ 与 $\nabla_{f(x,w)}$ 的外积。为了计算关于 $w$ 的 backward step（即 $(\nabla_w \mathcal{L})_j$），我们必须把输入梯度与对应的输出行相乘。

我们的 backward pass kernel 会先定义所有 block pointer，然后计算 $\nabla_x \mathcal{L}$：

@triton.jit
def weighted_sum_backward(
    x_ptr, weight_ptr,  # Input
    grad_output_ptr,  # Grad input
    grad_x_ptr, partial_grad_weight_ptr,  # Grad outputs
    stride_xr, stride_xd,

第 21 页

    stride_wd,
    stride_gr,
    stride_gxr, stride_gxd,
    stride_gwb, stride_gwd,
    NUM_ROWS, D,
    ROWS_TILE_SIZE: tl.constexpr, D_TILE_SIZE: tl.constexpr,
):
    row_tile_idx = tl.program_id(0)
    n_row_tiles = tl.num_programs(0)
    # Inputs
    grad_output_block_ptr = tl.make_block_ptr(
        grad_output_ptr,
        shape=(NUM_ROWS,), strides=(stride_gr,),
        offsets=(row_tile_idx * ROWS_TILE_SIZE,),
        block_shape=(ROWS_TILE_SIZE,),
        order=(0,),
    )
    x_block_ptr = tl.make_block_ptr(
        x_ptr,
        shape=(NUM_ROWS, D,), strides=(stride_xr, stride_xd),
        offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),
        block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
        order=(1, 0),
    )
    weight_block_ptr = tl.make_block_ptr(
        weight_ptr,
        shape=(D,), strides=(stride_wd,),
        offsets=(0,), block_shape=(D_TILE_SIZE,),
        order=(0,),
    )
    grad_x_block_ptr = tl.make_block_ptr(
        grad_x_ptr,
        shape=(NUM_ROWS, D,), strides=(stride_gxr, stride_gxd),
        offsets=(row_tile_idx * ROWS_TILE_SIZE, 0),
        block_shape=(ROWS_TILE_SIZE, D_TILE_SIZE),
        order=(1, 0),
    )
    partial_grad_weight_block_ptr = tl.make_block_ptr(
        partial_grad_weight_ptr,
        shape=(n_row_tiles, D,), strides=(stride_gwb, stride_gwd),
        offsets=(row_tile_idx, 0),
        block_shape=(1, D_TILE_SIZE),
        order=(1, 0),
    )
    for i in range(tl.cdiv(D, D_TILE_SIZE)):
        grad_output = tl.load(grad_output_block_ptr, boundary_check=(0,), padding_option="zero")  # (ROWS_TILE_SIZE,)
        # Outer product for grad_x
        weight = tl.load(weight_block_ptr, boundary_check=(0,), padding_option="zero")  # (D_TILE_SIZE,)
        grad_x_row = grad_output[:, None] * weight[None, :]
        tl.store(grad_x_block_ptr, grad_x_row, boundary_check=(0, 1))
        # Reduce as many rows as possible for the grad_weight result
        row = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option="zero")  # (ROWS_TILE_SIZE, D_TILE_SIZE)
        grad_weight_row = tl.sum(row * grad_output[:, None], axis=0, keep_dims=True)
        tl.store(partial_grad_weight_block_ptr, grad_weight_row, boundary_check=(1,))  # Never out of bounds for dim 0
        # Move the pointers to the next tile along D
        x_block_ptr = x_block_ptr.advance((0, D_TILE_SIZE))
        weight_block_ptr = weight_block_ptr.advance((D_TILE_SIZE,))
        partial_grad_weight_block_ptr = partial_grad_weight_block_ptr.advance((0, D_TILE_SIZE))
        grad_x_block_ptr = grad_x_block_ptr.advance((0, D_TILE_SIZE))

第 22 页

计算梯度 $\nabla x$ 很简单，我们把结果写到输出 tensor 的对应 tile 中。

但是，计算 $\nabla w$ 要稍微难一些。每个 kernel instance 负责 $x$ 的一个 row tile，但我们现在需要对 $x$ 的行求和。我们不在 backward pass 中直接做这个求和，而是假设 partial_grad_weight_ptr 包含一个 n_row_tiles × $H$ 的 matrix，其中第一维只在来自 $x$ 的一个 row tile 内归约。我们在写入这个 tensor 之前，先在当前 row tile 内归约。在 kernel 之外，我们用 torch.sum 归约 $\nabla w$，把每个 row tile 的结果加起来。[注 1]

autograd.Function 的最后一部分则相对简单：

class WeightedSumFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight):
        # ... (defined earlier)
    @staticmethod
    def backward(ctx, grad_out):
        x, weight = ctx.saved_tensors
        ROWS_TILE_SIZE, D_TILE_SIZE = ctx.ROWS_TILE_SIZE, ctx.D_TILE_SIZE  # These don't have to be the same
        n_rows, D = x.shape
        # Our strategy is for each thread block to first write to a partial buffer,
        # then we reduce over this buffer to get the final gradient.
        partial_grad_weight = torch.empty((triton.cdiv(n_rows, ROWS_TILE_SIZE), D), device=x.device, dtype=x.dtype)
        grad_x = torch.empty_like(x)
        weighted_sum_backward[(triton.cdiv(n_rows, ROWS_TILE_SIZE),)](
            x, weight,
            grad_out,
            grad_x, partial_grad_weight,
            x.stride(0), x.stride(1),
            weight.stride(0),
            grad_out.stride(0),
            grad_x.stride(0), grad_x.stride(1),
            partial_grad_weight.stride(0), partial_grad_weight.stride(1),
            NUM_ROWS=n_rows, D=D,
            ROWS_TILE_SIZE=ROWS_TILE_SIZE, D_TILE_SIZE=D_TILE_SIZE,
        )
        grad_weight = partial_grad_weight.sum(axis=0)
        return grad_x, grad_weight

最后，我们现在可以得到一个用法很像 torch.nn.functional 中那些函数的函数：

f_weightedsum = WeightedSumFunc.apply

现在，对两个 PyTorch tensor $x$ 和 $w$ 调用 f_weightedsum，会得到类似下面的 tensor：

tensor([ 90.8563, -93.6815, -80.8884,  ..., 103.4840, -21.4634, -24.0192],
       device='cuda:0', grad_fn=<WeightedSumFuncBackward>)

注意 tensor 上附带的 grad_fn —— 这说明当这个 tensor 出现在 computation graph 中时，PyTorch 知道在 backward pass 中该调用什么。至此，我们完成了 weighted sum 操作的 Triton 实现。

4.2.2 FlashAttention-2 前向传播

你将会用遵循 FlashAttention-2 [T. Dao, 2023] 的、显著改进的 Triton 实现，替换你的 PyTorch attention 实现。FlashAttention-2 使用一些技巧来按 tile 计算 forward pass，从而允许高效的内存访问模式，并避免在 global memory 中物化完整的 attention matrix。

[注 1] 当然，我们也可以为此写自己的 kernel。

第 23 页

在进入本节之前，我们强烈建议至少阅读原始的 FlashAttention 论文 [T. Dao et al., 2022]，它会让你对 FlashAttention 实现高效 attention 的核心技术产生直觉：跨 tile 以 online 方式计算 softmax（该技术由 [M. Milakov et al., 2018] 提出）。我们还建议看看 H. He [4]，以进一步理解 GPU 实际是如何执行 PyTorch 代码的。

理解 vanilla attention 中的低效之处

回忆一下，attention 的 forward pass（暂时忽略 masking）可以写成：

$$ \boldsymbol{S} = \frac{\boldsymbol{Q}\boldsymbol{K}^\top}{\sqrt{d}} \tag{4} $$

$$ P_{ij} = \operatorname{softmax}_j(\boldsymbol{S})_{ij} \tag{5} $$

$$ \boldsymbol{O} = \boldsymbol{P}\boldsymbol{V} \tag{6} $$

标准的 backward pass 是

$$ \boldsymbol{dV} = \boldsymbol{P}^\top \boldsymbol{dO} \tag{7} $$

$$ \boldsymbol{dP} = \boldsymbol{dO}\boldsymbol{V}^\top \tag{8} $$

$$ \boldsymbol{dS}_i = \operatorname{dsoftmax}(\boldsymbol{dP}_i) = (\operatorname{diag}(\boldsymbol{P}_i) - \boldsymbol{P}_i\boldsymbol{P}_i^\top)\boldsymbol{dP}_i \tag{9} $$

$$ \boldsymbol{dQ} = \frac{\boldsymbol{dS}\boldsymbol{K}}{\sqrt{d}} \tag{10} $$

$$ \boldsymbol{dK} = \frac{\boldsymbol{dS}^\top \boldsymbol{Q}}{\sqrt{d}} \tag{11} $$

正如我们所见，backward pass 依赖于 forward pass 中一些非常大的 activations。例如，在 Equation 7 中计算 $\boldsymbol{dV}$ 需要 $\boldsymbol{P}$，也就是 shape 为 (batch_size, n_heads, seq_len, seq_len) 的 attention scores —— 这个 activation matrix 的大小随 sequence length 二次方增长，这解释了我们在上面于大 sequence lengths 下 benchmark attention 时遇到的内存问题。在 vanilla attention 的 forward pass 和 backward pass 中，我们都要付出显著的内存 IO 开销，在 on-chip SRAM 和 GPU HBM 之间传输 $\boldsymbol{P}$ 和其他大型 activations。标准实现中会做若干次这样的传输：例如，标准的 backward pass 实现在 Equation 7 和 Equation 9 的计算中都会从 HBM 读取 $\boldsymbol{P}$。

FlashAttention 的主要目标是避免在 HBM 和 attention matrix 之间来回读写，以降低 IO 和峰值内存开销。我们通过三种技术实现这一点：tiling、recomputation 和 operator fusion。

Tiling

为了避免在 HBM 和 attention matrix 之间来回读写，我们在无法访问整个输入的情况下计算 softmax reduction。具体来说，我们重构 attention 计算，把输入切分成 tile，并对输入 tile 做多次遍历，从而增量式地执行 softmax reduction。

Recomputation

我们避免把 shape 为 (batch_size, n_heads, seq_len, seq_len) 的大型中间 attention matrices 存到 HBM 中。取而代之，我们会在 HBM 中保存某些 “activation checkpoints”，然后在 backward pass 中重新计算 forward pass 的一部分，以获得计算梯度所需的其他 activations。FlashAttention-2 还会存储 attention scores 的 logsumexp，即 $L$，它将用于简化 backward pass 的计算。$L$ 的表达式为：

$$ L_i = \log\left(\sum_j \exp(S_{ij})\right) \tag{12} $$

第 24 页

在我们最终的 kernel 中，我们会以 online 方式计算它，但最终结果应该相同。

把 tiling 和 recomputation 结合起来后，我们的内存 IO 和峰值用量不再依赖于 sequence_length 的平方，因此可以使用更大的 sequence lengths。

Operator fusion

最后，我们通过把所有操作放在单个 kernel 中执行，避免对 attention matrix 和其他中间 activations 的重复内存 IO —— 这被称为 operator fusion 或 kernel fusion。我们会为 forward pass 写一个单独的 Triton kernel，在 HBM 和 SRAM 之间数据转移有限的情况下，执行 attention 涉及的所有操作。Operator fusion 部分得益于 recomputation，因为我们可以避免通常为了把每个中间 activation 存到 HBM 而付出的内存 IO。

如果想对这些技术有更多直觉，可以看看 FlashAttention 论文 [T. Dao et al., 2022; T. Dao, 2023]。

使用 recomputation 的 backward pass

使用 $L$，我们可以做适当的 recomputation，并高效地计算 backward pass。在开始 backward pass 之前，我们在 global memory 中预计算 $D = \operatorname{rowsum}(\boldsymbol{O} \circ \boldsymbol{dO})$（其中 $\circ$ 是逐元素乘法），它等于 $\operatorname{rowsum}(\boldsymbol{P} \circ \boldsymbol{dP})$，因为 $\boldsymbol{P}\boldsymbol{dP}^\top = \boldsymbol{P}(\boldsymbol{dO}\boldsymbol{V}^\top)^\top = (\boldsymbol{P}\boldsymbol{V})\boldsymbol{dO}^\top = \boldsymbol{O}\boldsymbol{dO}^\top$（并且对任意矩阵 $\boldsymbol{A}$ 和 $\boldsymbol{B}$，有 $\operatorname{rowsum}(\boldsymbol{A} \circ \boldsymbol{B}) = \operatorname{diag}(\boldsymbol{A}\boldsymbol{B}^\top)$）。有了 $L$ 和 $D$ 向量，backward pass 就可以在不使用 softmax 的情况下计算。现在 backward pass 的完整计算如下：

$$ \boldsymbol{S} = \frac{\boldsymbol{Q}\boldsymbol{K}^\top}{\sqrt{d}} \tag{13} $$

$$ P_{ij} = \exp(S_{ij} - L_i) \tag{14} $$

$$ \boldsymbol{dV} = \boldsymbol{P}^\top \boldsymbol{dO} \tag{15} $$

$$ \boldsymbol{dP} = \boldsymbol{dO}\boldsymbol{V}^\top \tag{16} $$

$$ dS_{ij} = P_{ij}(dP_{ij} - D_i) \tag{17} $$

$$ \boldsymbol{dQ} = \frac{\boldsymbol{dS}\boldsymbol{K}}{\sqrt{d}} \tag{18} $$

$$ \boldsymbol{dK} = \frac{\boldsymbol{dS}^\top \boldsymbol{Q}}{\sqrt{d}} \tag{19} $$

可以看到，这一系列操作不要求我们在 forward pass 期间把 attention scores $\boldsymbol{P}$ 存到 HBM 中 —— 我们在 Equation 13 和 Equation 14 中从 activations $\boldsymbol{Q}$、$\boldsymbol{K}$ 和 $L$ 重新计算它们。

FlashAttention forward pass 的细节

现在我们已经对 FlashAttention-2 中使用的技术有了高层认识，接下来会深入你将实现的 FA2 forward pass kernel 的细节。为了避免在 HBM 和 attention matrix 之间来回读写，我们希望使用 tiling，即独立于其他 tile 计算输出的每个 tile。这要求我们能够计算 $\boldsymbol{P}$ 的 tile，理想情况下在 query 和 key 两个维度上都分 tile。

然而，当我们对 $\boldsymbol{S}$ 应用 softmax 时，需要归约 $\boldsymbol{S}$ 的整行来计算 softmax 分母，这意味着我们不能直接按 tile 计算 $\boldsymbol{P}$。FlashAttention-2 使用 online softmax 解决这个问题。在下面的文字中，我们用下标 $i$ 表示当前 query tile，用上标 $(j)$ 表示当前 key tile。沿 query 维度的 tile 大小为 $B_q$，沿 key 维度的 tile 大小为 $B_k$。我们不会沿 hidden dimension $d$ 分 tile。

我们还保留一些按行维护的 running values，$m_i^{(j)} \in \mathbb{R}^{B_q}$ 和 $l_i^{(j)} \in \mathbb{R}^{B_q}$。按行的 $m_i^{(j)}$ 值是一个 running maximum，跟踪它是为了能以数值稳定的方式计算 softmax（回忆一下 Assignment 1 中 softmax 实现里的这个技巧）。我们会用每个新的 row-

第 25 页

当 $j$ 增加时，$S$ 的逐行 tile。利用 running maximum，我们可以按如下方式计算未归一化的 softmax 值（分子）：$\tilde{P}_i^{(j)} = \exp(S_i^{(j)} - m_i^{(j)})$。$l_i^{(j)}$ 是 softmax 分母的 running proxy，并会用未归一化的 softmax 值更新为 $l_i^{(j)} = \exp(m_i^{(j-1)} - m_i^{(j)}) \cdot l_i^{(j-1)} + \operatorname{rowsum}(\tilde{P}_i^{(j)})$。当我们最终写出输出时，需要用 $l_i^{(T_k)}$ 完成归一化，它是在处理完所有 key tiles 之后 $l_i^{(j)}$ 的最终值。Algorithm 1 展示了 GPU 上的 forward pass。

Algorithm 1: FlashAttention-2 forward pass

  1. Require: $Q \in \mathbb{R}^{N_q \times d}$，$K, V \in \mathbb{R}^{N_k \times d}$，tile sizes $B_q, B_k$
  2. 将 $Q$ 切成 $T_q = \lceil \frac{N_q}{B_q} \rceil$ 个大小为 $B_q \times d$ 的 tiles $Q_1, \ldots, Q_{T_q}$
  3. 将 $K, V$ 切成 $T_k = \lceil \frac{N_k}{B_k} \rceil$ 个大小为 $B_k \times d$ 的 tiles $K^{(1)}, \ldots, K^{(T_k)}$ 和 $V^{(1)}, \ldots, V^{(T_k)}$
  4. for $i = 1, \ldots, T_q$ do
  5. 从 global memory 加载 $Q_i$
  6. 初始化 $O_i^{(0)} = \mathbf{0} \in \mathbb{R}^{B_q \times d}$，$l_i^{(0)} = 0 \in \mathbb{R}^{B_q}$，$m_i^{(0)} = -\infty \in \mathbb{R}^{B_q}$
  7. for $j = 1, \ldots, T_k$ do
  8. 从 global memory 加载 $K^{(j)}, V^{(j)}$
  9. 计算 pre-softmax attention scores 的 tile $S_i^{(j)} = \frac{Q_i (K^{(j)})^\top}{\sqrt{d}} \in \mathbb{R}^{B_q \times B_k}$
  10. 计算 $m_i^{(j)} = \max(m_i^{(j-1)}, \operatorname{rowmax}(S_i^{(j)})) \in \mathbb{R}^{B_q}$
  11. 计算 $\tilde{P}_i^{(j)} = \exp(S_i^{(j)} - m_i^{(j)}) \in \mathbb{R}^{B_q \times B_k}$
  12. 计算 $l_i^{(j)} = \exp(m_i^{(j-1)} - m_i^{(j)}) l_i^{(j-1)} + \operatorname{rowsum}(\tilde{P}_i^{(j)}) \in \mathbb{R}^{B_q}$
  13. 计算 $O_i^{(j)} = \operatorname{diag}(\exp(m_i^{(j-1)} - m_i^{(j)})) O_i^{(j-1)} + \tilde{P}_i^{(j)} V^{(j)}$
  14. end for
  15. 计算 $O_i = \operatorname{diag}(l_i^{(T_k)})^{-1} O_i^{(T_k)}$
  16. 计算 $L_i = m_i^{(T_k)} + \log(l_i^{(T_k)})$
  17. 将 $O_i$ 写入 global memory，作为 $O$ 的第 $i$ 个 tile。
  18. 将 $L_i$ 写入 global memory，作为 $L$ 的第 $i$ 个 tile。
  19. end for
  20. 返回输出 $O$ 和 logsumexp $L$。

Triton Tips and Tricks

第 26 页

kernel_fn[(launch_grid_d1, launch_grid_d2, ...)](...arguments...)

在 torch.autograd.Function 子类的方法中，如我们在 weighted sum 示例中所见。

Problem (flash_forward): FlashAttention-2 Forward Pass (15 points)

(a)

写一个纯 PyTorch（不用 Triton）的 autograd.Function，实现 FlashAttention-2 的 forward pass。这会比常规 PyTorch 实现慢很多，但会帮助你调试 Triton kernel。

你的实现应接收输入 $Q$、$K$、$V$ 以及标志 is_causal，并产生输出 $O$ 和 logsumexp 值 $L$。本任务中可以忽略 is_causal 标志。autograd.Function 的 forward 随后应为 backward pass 保存 $L, Q, K, V, O$，并返回 $O$。记住，autograd.Function 的 forward 方法总是把 context 作为第一个参数。任何 autograd.Function 类都需要实现 backward 方法，但现在可以只让它 raise NotImplementedError。如果你需要对照物，可以在 PyTorch 中实现 Equation 4 到 Equation 6 以及 Equation 12，并比较你的输出。

接口为 def forward(ctx, Q, K, V, is_causal=False)。自行确定 tile sizes，但确保它们至少为 $16 \times 16$。我们测试你的代码时使用的维度总是 2 的幂且至少为 16，因此你不需要担心越界访问。

Deliverable: 一个 torch.autograd.Function 子类，在 forward pass 中实现 FlashAttention-2。要测试你的代码，实现 adapters.get_flashattention_autograd_function_pytorch。然后运行测试 uv run pytest -k test_flash_forward_pass_pytorch，并确保你的实现通过。

(b)

写一个 Triton kernel，按照 Algorithm 1 实现 FlashAttention-2 的 forward pass。然后写另一个 torch.autograd.Function 的子类，在 forward pass 中调用这个（融合的）kernel，而不是在 PyTorch 中计算结果。几个针对本题的提示：

@triton.jit
def flash_fwd_kernel(

第 27 页

Q_ptr, K_ptr, V_ptr,
    O_ptr, L_ptr,
    stride_qb, stride_qq, stride_qd,
    stride_kb, stride_kk, stride_kd,
    stride_vb, stride_vk, stride_vd,
    stride_ob, stride_oq, stride_od,
    stride_lb, stride_lq,
    N_QUERIES, N_KEYS,
    scale,
    D: tl.constexpr,
    Q_TILE_SIZE: tl.constexpr,
    K_TILE_SIZE: tl.constexpr,
):
    # Program indices
    query_tile_index = tl.program_id(0)
    batch_index = tl.program_id(1)
    # Offset each pointer with the corresponding batch index
    # multiplied with the batch stride for each tensor
    Q_block_ptr = tl.make_block_ptr(
        Q_ptr + batch_index * stride_qb,
        shape=(N_QUERIES, D),
        strides=(stride_qq, stride_qd),
        offsets=(query_tile_index * Q_TILE_SIZE, 0),
        block_shape=(Q_TILE_SIZE, D),
        order=(1, 0),
    )
    ...

其中 scale 是 $1/\sqrt{d}$，Q_TILE_SIZE 和 K_TILE_SIZE 分别是 $B_q$ 和 $B_k$。你可以之后调优这些值。

这些额外准则可能帮助你避免精度问题：

Deliverable: 一个 torch.autograd.Function 子类，使用你的 Triton kernel 在 forward pass 中实现 FlashAttention-2。实现 adapters.get_flash_autograd_function_triton。然后运行测试 uv run pytest -k test_flash_forward_pass_triton，并确保你的实现通过。

(c)

在 autograd.Function 实现的最后加一个用于 causal masking 的 flag。这应是一个布尔 flag，当设为 True 时启用用于 causal masking 的索引比较。你的 Triton kernel 应有对应的额外参数 is_causal: tl.constexpr（这是必需的类型注解）。在 Triton 中，为 queries 和 keys 构造合适的 index vectors，并比较它们以形成大小为 $B_q \times B_k$ 的方形 mask。对于被 mask 掉的元素，在 attention score matrix $S_i^{(j)}$ 的对应元素上加上常量值 -1e6。确保用 ctx.is_causal = is_causal 为 backward 保存 mask flag。

Deliverable: 为你的 torch.autograd.Function 子类增加一个额外 flag，使用你的 Triton kernel 实现带 causal masking 的 FlashAttention-2 forward pass。确保该 flag 是可选的，并且默认值为 False，以便之前的测试仍然通过。

第 28 页

使用重计算实现反向传播

注意，与 Equation 7 到 Equation 11 中的标准反向传播不同，我们可以使用重计算来避免 Equation 13 到 Equation 19 所示反向传播中的 softmax 操作。这意味着我们可以用一个 trivial kernel 来计算反向传播，不需要任何 online 技巧。因此，在这一部分，你可以通过对一个普通 PyTorch 函数调用 torch.compile 来实现反向传播（不使用 Triton）。

Problem (flash_backward): FlashAttention-2 反向传播（5 points）

使用 PyTorch（不是 Triton）和 torch.compile 为你自己的 FlashAttention-2 autograd.Function 实现反向传播。你的实现应接收 $Q$、$K$、$V$、$O$、$dO$ 和 $L$ tensor 作为输入，并返回 $dQ$、$dK$ 和 $dV$。记得计算并使用 $D$ 向量。你可以参照 Equation 13 到 Equation 19 的计算过程。

Deliverable：要测试你的实现，运行：

uv run pytest -k test_flash_backward

现在，让我们比较你的 FlashAttention-2（部分）Triton 实现与你的普通 Attention 的 PyTorch 实现的性能。

Problem (flash_benchmarking): FlashAttention-2 Benchmarking（5 points）

(a) 使用 triton.testing.do_bench 编写一个 benchmarking 脚本，比较你的 FlashAttention-2 前向和反向传播的（部分）Triton 实现与普通 PyTorch 实现（即不使用 FlashAttention）的性能。

具体来说，你要报告一个表格，包含你的 Triton 实现和 PyTorch 实现的前向、反向以及端到端前向-反向传播的延迟。在开始 benchmarking 之前随机生成任何必要的输入，并在单张 B200 上运行 benchmark。始终使用 batch size 1 和 causal masking。遍历以下三者的笛卡尔积：从 128 到 65536 的各种 2 的幂的序列长度、从 16 到 128 的各种 2 的幂的 embedding 维度大小，以及 torch.bfloat16 和 torch.float32 两种精度。你可能需要根据输入大小调整 tile 大小。

Deliverable：一个结果表格，使用上述设置比较你的 FlashAttention-2 实现与 PyTorch 实现，并报告前向、反向和端到端延迟。

4.2.3 OPTIONAL: Triton 反向传播

如果你有兴趣用 Triton 做更多练习，和/或想提交一个快速上榜的结果，我们在下面提供了可以让你用 Triton 实现的 tiled FlashAttention-2 反向传播。Algorithm 2 展示了应如何在 Triton 中实现 FlashAttention-2 反向传播。这里的一个关键技巧是计算 $P$ 两次：一次用于 $dQ$，一次用于 $dK$ 和 $dV$。这让我们可以跳过 thread block 之间的同步，也就是说可以避免缓慢的 atomic 操作。

第 29 页

Algorithm 2：Tiled FlashAttention-2 反向传播

5 Distributed Data Parallel Training（分布式数据并行训练）

在作业的下一部分，我们将探索使用多块 GPU 训练语言模型的方法，重点是数据并行。我们会先介绍 PyTorch 中的分布式通信基础。然后，我们会研究一个分布式数据并行训练的朴素实现，接着实现并 benchmark 各种通信效率改进。

第 30 页

5.1 Single-Node Distributed Communication in PyTorch（PyTorch 中的单节点分布式通信）

我们先看一个 PyTorch 中的简单分布式应用，目标是生成四个随机整数 tensor 并计算它们的和。

在下面的分布式情形中，我们会启动四个 worker process，每个进程生成一个随机整数 tensor。为了跨 worker process 对这些 tensor 求和，我们会调用 all-reduce 集合通信操作；该操作会把每个进程上的原始 data tensor 替换为 all-reduce 后的结果（即求和结果）。

现在来看一些代码。

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
def distributed_demo(rank, world_size):
    setup(rank, world_size)
    data = torch.randint(0, 10, (3,))
    print(f"rank {rank} data (before all-reduce): {data}")
    dist.all_reduce(data, async_op=False)
    print(f"rank {rank} data (after all-reduce): {data}")
if __name__ == "__main__":
    world_size = 4
    mp.spawn(fn=distributed_demo, args=(world_size, ), nprocs=world_size, join=True)

运行上面的脚本后，我们会得到下面的输出。正如预期，每个 worker process 最初持有不同的 data tensor。经过 all-reduce 操作（它对所有 worker process 中的 tensor 求和）后，每个 worker process 上的 data 会被原地修改，以保存 all-reduce 后的结果。[注 2]

$ uv run python distributed_hello_world.py
rank 3 data (before all-reduce): tensor([3, 7, 8])
rank 0 data (before all-reduce): tensor([4, 4, 7])
rank 2 data (before all-reduce): tensor([6, 0, 7])
rank 1 data (before all-reduce): tensor([9, 5, 3])
rank 1 data (after all-reduce): tensor([22, 16, 25])
rank 0 data (after all-reduce): tensor([22, 16, 25])
rank 3 data (after all-reduce): tensor([22, 16, 25])
rank 2 data (after all-reduce): tensor([22, 16, 25])

现在让我们更仔细地回顾上面的脚本。命令 mp.spawn 会启动 nprocs 个进程，这些进程用提供的 args 运行 fn。此外，函数 fn 会以 fn(rank, *args) 的形式被调用，其中 rank 是 worker process 的索引（取值在 0 到 nprocs-1 之间）。因此，我们的 distributed_demo 函数

[注 2] 如果你多次运行这个脚本，你会注意到打印输出的顺序是不确定的。由于这个应用运行在分布式环境中，我们无法控制命令执行的确切顺序——我们唯一的保证是：all-reduce 操作完成后，各个进程将持有按位相同（bitwise identical）的结果 tensor。

第 31 页

必须接受这个整数 rank 作为它的第一个位置参数。此外，我们传入 world_size，它指的是 worker 进程的总数。

每个 worker 进程都属于一个 process group，该 process group 通过 dist.init_process_group 初始化。process group 表示多个 worker 进程，它们将通过一个共享的 master 进行协调和通信。master 由其 IP 地址和端口定义，并且 master 运行 rank 为 0 的进程。像 all-reduce 这样的 collective communication 操作会作用于 process group 中的每个进程。

在本例中，我们用 gloo backend 初始化了 process group，但也可以使用其他 backend。特别是，nccl backend 会使用 NVIDIA NCCL collective communications 库，对于 CUDA tensor 通常性能更好。然而，NCCL 只能在带有 GPU 的机器上使用，而 Gloo 可以在仅有 CPU 的机器上运行。在分布式 GPU 训练中应始终使用 NCCL，只有在没有可用 GPU 的本地开发中才使用 Gloo。本例中使用 Gloo 是因为它支持在仅有 CPU 的机器上进行本地执行和开发。

运行 multi-GPU 作业时，请确保不同 rank 使用不同 GPU。一种方法是在 setup 函数中调用 torch.cuda.set_device(rank)，这样 tensor.to("cuda") 会自动把它移动到指定的 device。或者，你可以显式创建一个 per-rank 的 device 字符串（例如 device = f"cuda:{rank}"），然后将这个 device 字符串用作任何数据移动的目标 device（例如 tensor.to(f"cuda:{rank}")）。

Terminology（术语）

在作业的其余部分（以及你可能在网上看到的各种其他资源）中，你可能会在 PyTorch 分布式通信的语境下遇到以下术语。虽然本作业将聚焦于单 node、多进程的分布式训练，但这些术语对于理解一般的分布式训练很有用。可视化表示见图 3。

第 32 页

Figure 3：一个运行在 2 个 node 上、world size 为 8 的分布式应用的示意图。每个 worker 进程由一个 global rank（从 0 到 7）和一个 local rank（从 0 到 3）标识。图片来源：https://lightning.ai/docs/fabric/stable/advanced/distributed_communication.html

5.1.1 分布式应用基准测试的最佳实践

在作业的这一部分中，你将通过对分布式应用进行基准测试，来更好地理解通信带来的开销。以下是一些最佳实践：

在 backend 之间切换应当只需修改你的 init_process_group 调用和 tensor 的 device cast 即可。

Problem (distributed_communication_single_node)：分布式通信（单节点）（5 分）

编写一个脚本，对单 node 多进程设置中 all-reduce 操作的运行时间进行基准测试。上面的示例代码可以作为一个合理的起点。尝试改变以下设置：

Resource requirements：最多 6 块 GPU。每次基准测试运行应少于 5 分钟。

Deliverable：比较各种设置的图（一张或多张）和/或表（一张或多张），并用 2–3 句话评论你的结果以及你对各因素如何相互作用的思考。

[注 3] 更多细节见 https://github.com/pytorch/pytorch/issues/68112#issuecomment-965932386。

第 33 页

5.2 分布式数据并行训练的一个朴素实现

既然我们已经了解了在 PyTorch 中编写分布式应用的基础知识，接下来就构建一个分布式数据并行（distributed data parallel, DDP）训练的最小实现。

数据并行会把 batch 拆分到多个设备（例如 GPU）上，从而能够训练那些大到无法放入单个设备的 batch size。例如，假设有四块设备，每块最多能处理 batch size 为 32，那么数据并行训练就能得到 128 的有效 batch size。

下面是朴素地进行分布式数据并行训练的步骤。最初，每个设备各自构造一个（随机初始化的）模型。我们使用 broadcast 集合通信操作，把模型参数从 rank 0 发送到所有其他 rank。训练开始时，每个设备都持有一份完全相同的模型参数和 optimizer states（例如 Adam 中累积的梯度统计量）。

  1. 给定一个包含 $n$ 个样本的 batch，这个 batch 会被分片，每个设备收到 $n/d$ 个互不重叠的样本（其中 $d$ 是用于数据并行训练的设备数）。$d$ 应当能整除 $n$，否则一些 rank 会比另一些 rank 做更多工作，这一步就会被最慢的设备拖成瓶颈。
  2. 每个设备使用本地的模型参数副本，对自己收到的 $n/d$ 个样本执行 forward pass，并执行 backward pass 计算梯度。注意，此时每个设备持有的是根据自己收到的 $n/d$ 个样本计算出的梯度。
  3. 然后我们使用 all-reduce 集合通信操作，在不同设备之间对梯度求平均，因此每个设备都持有在所有 $n$ 个样本上平均后的梯度。
  4. 接下来，每个设备执行一次 optimizer step，更新自己的参数副本——从 optimizer 的视角看，它只是在优化一个本地模型。由于所有设备都从相同的初始模型和 optimizer state 出发，并且在每次迭代中使用相同的平均梯度，所以参数和 optimizer states 会在所有不同设备上保持同步。至此，我们完成了一次训练迭代，可以重复该过程。

Problem (naive_ddp)：朴素 DDP（5 分）

Deliverable: 实现一种朴素形式的分布式数据并行训练：在 backward pass 之后，对各个参数的梯度分别执行 all-reduce。要测试你的实现，请实现 [adapters.get_ddp] 和（可选）[adapters.ddp_on_after_backward]，然后运行：

uv run pytest tests/test_ddp.py

Problem (naive_ddp_benchmarking)：朴素 DDP 基准测试（3 分）

在这个朴素 DDP 实现中，参数梯度会在每次 backward pass 之后，在各个 rank 之间分别执行 all-reduce。为了更好地理解数据并行训练的开销，请编写一个脚本，对之前实现的语言模型在使用这种朴素 DDP 实现训练时进行基准测试。测量每个训练 step 的总时间，以及花在梯度通信上的时间占比。在 Section 2.1.2 所述 xl 模型规模下，于单节点设置（1 node x 2 GPUs）中收集测量结果。

Deliverable: 描述你的基准测试设置，并给出每种设置下测得的每次训练迭代时间，以及花在梯度通信上的时间。

第 34 页

5.3 改进最小 DDP 实现

我们在 Section 5.2 中看到的最小 DDP 实现有几个关键限制：

  1. 它为每一个参数 tensor 单独执行一次 all-reduce 操作。每次通信调用都会带来开销，因此把通信调用批量化可能有助于尽量减小这一开销。
  2. 它会等待 backward pass 结束后再通信梯度。然而，backward pass 是增量计算的。因此，当某个参数的梯度准备好时，就可以立即通信，而不必等待其他参数的梯度。这样我们就能让梯度通信与 backward pass 的计算重叠，从而降低分布式数据并行训练的开销。

在作业的这一部分，我们会依次处理这些限制，并测量它们对训练速度的影响。

5.3.1 减少通信调用次数

与其为每个参数 tensor 都发起一次通信调用，不如看看能否通过批量化 all-reduce 来提升性能。具体来说，我们会把想要 all-reduce 的梯度拼接成一个 tensor，然后在所有 rank 之间对合并后的梯度执行 all-reduce。使用 torch._utils._flatten_dense_tensors 和 torch._utils._unflatten_dense_tensors 可能会有帮助。

Problem (minimal_ddp_flat_benchmarking)：使用扁平梯度的最小 DDP 基准测试（2 分）

修改你的最小 DDP 实现，使其通信一个包含所有参数扁平化梯度的 tensor。在先前使用的条件下（1 node x 2 GPUs，Section 2.1.2 所述 xl 模型规模），将它的性能与为每个参数 tensor 单独发起 all-reduce 的最小 DDP 实现进行比较。

Deliverable: 在使用单个批量化的 all-reduce 调用进行分布式数据并行训练时，测得的每次训练迭代时间以及花在梯度通信上的时间。用 1–2 句话比较批量化通信梯度与分别通信梯度的结果。

5.3.2 让计算与单个参数梯度的通信重叠

虽然批量化通信调用可能有助于降低发起大量小型 all-reduce 操作所带来的开销，但所有通信时间仍然会直接计入开销。为了解决这个问题，我们可以利用这样一个观察：backward pass 会逐层增量计算梯度（从 loss 开始，向输入方向推进）——因此，参数梯度一旦准备好，我们就可以立即对其执行 all-reduce，从而通过让 backward pass 的计算与梯度通信重叠来降低数据并行训练的开销。

我们先实现并基准测试一个 distributed data parallel wrapper：在 backward pass 期间，当各个参数 tensor 的梯度准备好时，异步地对它们执行 all-reduce。下面的提示可能会有用：

Backward hooks 要想在 backward pass 中某个参数的梯度累积完成后自动调用一个函数，可以使用 register_post_accumulate_grad_hook 函数。[注 4]

[注 4] 参见 https://pytorch.org/docs/stable/generated/torch.Tensor.register_post_accumulate_grad_hook.html 获取更多信息与用法示例。

第 35 页

异步通信 (Asynchronous communication)

所有 PyTorch collective communication 操作都支持同步（async_op=False）和异步执行（async_op=True）。同步调用会阻塞，直到该 collective operation 被排入 GPU 队列。这并不意味着 CUDA 操作已经完成，因为 CUDA 操作是异步的。话虽如此，后续使用该输出的函数调用仍会按预期运行。[注 5] 相比之下，异步调用会返回一个 distributed request handle——因此，当函数返回时，不能保证该 collective communication 操作已被排入 GPU 队列，更不用说已经完成。若要等待该操作被排入 GPU 队列（并因此使输出可用于后续操作），可以对返回的 communication handle 调用 handle.wait()。

例如，下面两个例子分别用同步或异步调用对 tensor 列表中的每个 tensor 做 all-reduce：

tensors = [torch.rand(5) for _ in range(10)]
# Synchronous, block until operation is queued on the GPU.
for tensor in tensors:
    dist.all_reduce(tensor, async_op=False)
# Asynchronous, return immediately after each call and
# wait on results at the end.
handles = []
for tensor in tensors:
    handle = dist.all_reduce(tensor, async_op=True)
    handles.append(handle)
# ...
# Possibly execute other commands that don't rely on the all_reduce results
# ...
# Ensure that all-reduce calls were queued and
# therefore other operations depending on the
# all-reduce output can be queued.
for handle in handles:
    handle.wait()
handles.clear()

Problem (ddp_overlap_individual_parameters): DDP with Overlapping Individual Parameters（重叠单个参数的 DDP，5 分）

实现一个 Python 类来处理 distributed data parallel training。该类应包装任意 PyTorch nn.Module，并负责在训练前广播权重（使所有 rank 具有相同的初始参数），以及发起用于梯度平均的通信调用。我们推荐以下 public interface：

def __init__(self, module: torch.nn.Module):

给定一个已实例化的、待并行的 PyTorch nn.Module，构造一个 DDP 容器，由它处理跨 rank 的梯度同步。

def forward(self, *inputs, **kwargs):

用提供的位置参数和关键字参数调用被包装模块的 forward() 方法。

[注 5] 在高级场景中，如果你使用多个 CUDA stream，可能需要显式地在 stream 之间同步，以确保输出已准备好供后续操作使用。参见 https://pytorch.org/docs/stable/notes/cuda.html#cuda-streams。

第 36 页

def finish_gradient_synchronization(self):

当被调用时，等待异步通信调用在 GPU 上完成。

为了使用这个类执行 distributed training，我们会传入一个要包装的 module，然后在运行 optimizer.step() 之前添加对 finish_gradient_synchronization() 的调用，以确保 optimizer step 这个依赖于梯度的操作可以被安全地排入队列：

model = ToyModel().to(device)
ddp_model = DDP(model)
for _ in range(train_steps):
    x, y = get_batch()
    logits = ddp_model(x)
    loss = loss_fn(logits, y)
    loss.backward()
    ddp_model.finish_gradient_synchronization()
    optimizer.step()

Deliverable: 实现一个容器类来处理 distributed data parallel training。该类应重叠梯度通信与 backward pass 的计算。为了测试你的 DDP 类，首先实现 adapters [adapters.get_ddp] 和 [adapters.ddp_on_after_backward]（后者是可选的，取决于你的实现，你可能不需要它）。

然后，要执行测试，运行：

uv run pytest tests/test_ddp.py

我们建议多次运行测试（例如 5 次），以确保它能稳定通过。

Problem (ddp_overlap_individual_parameters_benchmarking): DDP Overlapping Individual Parameters Benchmarking（重叠单个参数的 DDP 基准测试，1 分）

(a)

当将 backward pass 计算与单个 parameter 梯度的通信重叠时，对你的 DDP 实现进行性能基准测试。在相同设置下（1 个 node、2 个 GPU，以及 Section 2.1.2 中描述的 xl model size），将其性能与我们之前研究过的设置进行比较（最小 DDP 实现：要么对每个 parameter tensor 发起一次 all-reduce，要么对所有 parameter tensor 的拼接发起一次 all-reduce）。

Deliverable: 重叠 backward pass 与单个 parameter 梯度通信时测得的每次训练迭代时间，并用 1-2 句话比较结果。

(b)

使用 Nsight profiler 对你的基准测试代码进行插桩（使用 1 个 node、2 个 GPU、xl model size 的设置），比较初始 DDP 实现与该重叠实现。目视比较两条 trace，并提供 profiler 截图，证明一个实现将 compute 与 communication 重叠，而另一个没有。

Deliverable: 2 张截图（一张来自初始 DDP 实现，另一张来自该将 compute 与 communication 重叠的 DDP 实现），直观显示 communication 是否与 backward pass 重叠。

第 37 页

6 Optimizer State Sharding（优化器状态分片）

Distributed data parallel training 在概念上很简单，而且通常非常有效，但它要求每个 rank 持有一份不同的模型参数和 optimizer state 副本。这种冗余可能带来显著的内存开销。例如，AdamW optimizer 为每个 parameter 维护两个 float，这意味着它消耗的内存是模型权重的两倍。S. Rajbhandari et al. [5] 描述了几种通过将 (1) optimizer state、(2) gradients 和 (3) parameters 划分到不同 rank 上来减少 data-parallel training 中这种冗余的方法，并在 worker 之间按需通信。

在作业的这一部分，我们将通过实现一个简化版的 optimizer state sharding 来减少每个 rank 的内存消耗。每个 rank 的 optimizer 实例不会保留所有 parameter 的 optimizer state，而只处理一部分 parameter（约为 $1/\text{world\_size}$）。当每个 rank 的 optimizer 执行一次 optimizer step 时，它只会更新其 shard 中的那部分模型 parameter。然后，每个 rank 会将其更新后的 parameter 广播给其他 rank，以确保每次 optimizer step 后模型 parameter 保持同步。

Problem (optimizer_state_sharding): Optimizer State Sharding（优化器状态分片，15 分）

实现一个 Python 类来处理 optimizer state sharding。该类应包装任意输入的 PyTorch optim.Optimizer，并负责在每次 optimizer step 后同步更新后的 parameter。我们推荐以下 public interface：

def __init__(self, params, optimizer_cls: Type[Optimizer], **kwargs: Any):

初始化 sharded state optimizer。params 是要优化的 parameter 集合（或 parameter group，以防用户想对模型不同部分使用不同 hyperparameter，例如 learning rate）；这些 parameter 会被分片到所有 rank 上。optimizer_cls 参数指定要包装的 optimizer 类型（例如 optim.AdamW）。最后，任何剩余的关键字参数都会转发给 optimizer_cls 的构造函数。请确保在此方法中调用 torch.optim.Optimizer 超类构造函数。

def step(self, closure, **kwargs):

用提供的 closure 和关键字参数调用被包装 optimizer 的 step() 方法。在更新 parameter 后，与其他 rank 同步。

def add_param_group(self, param_group: dict[str, Any]):

此方法应向 sharded optimizer 添加一个 parameter group。它会在超类构造函数构造 sharded optimizer 期间被调用，也可能在训练期间被调用（例如，用于逐步 unfreeze 模型中的 layer）。因此，此方法应负责在 rank 之间分配模型 parameter。

Deliverable: 实现一个容器类来处理 optimizer state sharding。为了测试你的 sharded optimizer，首先实现 adapter [adapters.get_sharded_optimizer]。然后，要执行测试，运行：

uv run pytest tests/test_sharded_optimizer.py

我们建议多次运行测试（例如 5 次），以确保它们能稳定通过。

现在我们已经实现了 optimizer state sharding，让我们分析它对训练期间峰值内存使用量及其运行时开销的影响。

第 38 页

Problem (optimizer_state_sharding_accounting)：Optimizer State Sharding Accounting（优化器状态分片核算，5 分）

(a)

创建一个脚本，用于分析在有和没有 optimizer state sharding 的情况下训练语言模型时的峰值内存使用。使用标准配置（1 个节点、2 块 GPU、xl 模型规模），报告模型初始化后、optimizer step 之前紧接着的时刻，以及 optimizer step 之后紧接着的时刻的峰值内存使用。结果是否符合你的预期？分解每种设置下的内存使用（例如，参数占多少内存、优化器状态占多少内存等）。

Deliverable: 用 2–3 句话回复，给出峰值内存使用结果，并说明内存如何在不同的模型组件和优化器组件之间划分。

(b)

我们的 optimizer state sharding 实现如何影响训练速度？在标准配置（1 个节点、2 块 GPU、xl 模型规模）下，测量有和没有 optimizer state sharding 时每次迭代所花的时间。

Deliverable: 用 2–3 句话回复你的计时结果。

(c)

我们的 optimizer state sharding 方法与 ZeRO stage 1（在 S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He [5] 中描述为 ZeRO-DP $P_{os}$）有何不同？

Deliverable: 用 2–3 句话总结任何差异，尤其是与内存和通信量相关的差异。

7 Fully-Sharded Data Parallel（全分片数据并行）

借助 optimizer state sharding 和 data parallel，我们能够沿 data-parallel 轴拆分优化器状态和激活。然而，模型权重仍然重复——我们在每块 GPU 上都存储了它们的完整副本。

我们可以通过将 data parallel (DP) 轴变为 fully-sharded data parallel (FSDP) 轴来解决这个问题。使用 FSDP 时，每块 GPU 只存储每个权重张量中属于自己的切片，但必须通过 all-gather 从其他 GPU 拉取切片来组成完整权重张量，为 forward 或 backward pass 做准备。

为了避免 GPU 计算等待通信完成，大多数 FSDP 实现会提前调度该层的 all-gather，使相关权重在需要之前就绪，从而防止通信阻塞计算。这使权重分片通信不在关键路径上，也就是说，只要通信能跟上计算且调度得当，它就没有代价。

有些层在内存和计算上足够小，传输的延迟开销不值得。你应当将这些层标记为不被 FSDP 分片。在我们的架构中，这种情况主要出现在 norms 上。这样剩下的是 embedding 层和每个 linear 层。

虽然有必要以 FP32 存储 master weights（任何会被反复累加的值都对精度敏感），但权重不一定需要以 FP32 使用。在 mixed precision 中，我们总是在使用前转换为低精度计算数据类型，因此为了节省带宽，我们不妨在权重被通信之前就进行转换。

Problem (fsdp)：Fully-Sharded Data Parallel（全分片数据并行，15 分）

第 39 页

实现一个用于 fully-sharded data parallel training 的 Python 类。该类应包装任意 PyTorch nn.Module（你的完整模型），并 hook 或包装其中的任何 Linear 或 Embedding 层。我们推荐以下公共接口：

def __init__(self, module: torch.nn.Module, compute_dtype: torch.dtype | None = None):

给定一个已实例化的、待并行化的 PyTorch nn.Module，构造一个 FSDP 模块，负责 weight all-gather 和 gradient reduce-scatter。确保你的 hooks 或模块包装器能及时完成 weight all-gather，以便 forward pass 使用。为限制内存使用，只有当前层前面两层完成 forward pass 后才开始 gather。在 backward pass 中，你的 hooks 或模块包装器应通过 all-gather 让权重可用于计算。梯度可用时，应将其 reduce-scatter 到适当的 rank。确保用完后释放 gather 到的权重。当提供 compute_dtype 时，在通信或用于计算前将权重转换为该 dtype，同时保持 master weights 和优化器更新为 FP32。

def forward(self, *inputs, **kwargs):

调用被包装模块的 forward() 方法，并传入提供的位置参数和关键字参数。

def finish_gradient_synchronization(self):

调用时，等待 GPU 上的异步通信调用完成。

Deliverable: 实现一个容器类来处理 fully sharded data parallel training。该容器的每个 shard 都应兼容 assignment 1 中的标准 AdamW 实现。为测试你的 FSDP 实现，实现适配器 [adapters.get_fsdp]。运行测试：

uv run pytest tests/test_fsdp.py

我们建议多次运行测试（例如 5 次），以捕获任何竞态条件。

Problem (fsdp_accounting)：FSDP Accounting（FSDP 核算，5 分）

(a)

基于你在第 6 节中的分析，实施 FSDP 后你预计能从峰值内存中节省多少内存？计算时可以忽略将权重 all-gather 到每块 GPU 所需的预分配缓冲区大小。

Deliverable: 用 2–3 句话回复你的发现。

(b)

在 2 块 GPU 上对 xl 模型进行 profile，并关注权重的 all-gather。通信是否能在 forward pass 之前及时完成？

Deliverable: 用 2–3 句话回复你的计时结果。附上 Nsight 截图以支持你的结论。

8 Analyzing Parallelism Strategies（分析并行策略）

我们还可以沿更多轴并行化训练过程。一些常见策略包括：

第 40 页

• Tensor Parallelism (TP) —— 权重矩阵沿输入维或输出维进行分片。各设备计算其分片对应的 activation，然后跨设备对 activation 做 reduce 或 gather。

• Pipeline Parallelism (PP) —— 模型按层切分为多个 stage，每个 stage 在不同设备上运行。

• Expert Parallelism (EP) —— Mixture-of-Experts 模型中的 experts 被分到不同设备上，每个设备为自己的 expert 计算输出。

在本节中，我们会在一个简化设定下做一些基础数学，以便在并行策略之间做选择，并决定如何组合它们。首先，我们会聚焦于 DP、FSDP、TP 及其组合。我们的方法是计算每种策略的通信开销，并将其与计算开销比较；这样就能知道在通信开销成为瓶颈之前，我们最多能扩展到多少设备。

若想更详细地了解 TPU/GPU 拓扑与并行策略，TPU Scaling Book（J. Austin et al. [6]）是很好的资源。若想更详细地讨论 pipeline parallel，参见 The Ultra-Scale Playbook Appendix（H. Z. P. N. M. M. L. W. T. W. Nouamane Tazi Ferdinand Mom [7]）。这些书的其余部分也有很多可能对你有用的信息。

8.1 Communication Primitives

我们的第一步是理解通信原语。在这个简化设定中，假设有 $N$ 台设备，编号为 $0,\dots,N-1$，并且每一对设备之间都由一条链路连接。我们还假设每台设备有 $W$ 的出带宽（egress bandwidth，即发送带宽）；换句话说，每台设备能以每秒 $W$ 字节的速率向另一台设备发送数据。我们该如何实现 gather 和 reduce？

实现 all-gather 操作的一种常见方式是 ring all-gather。回顾一下，在 all-gather 中，每台设备 $i$ 一开始持有大小为 $S/N$ 的 chunk $x_i$，最终得到整个 $\mathbf{x}=[x_0,\dots,x_{N-1}]$，其大小为 $S$（字节）。在 ring all-gather 中，我们把设备排成一个环。在每一步中，每台设备把它当前的 chunk 发送给它右边的下一台设备，并保存它从左边设备收到的 chunk。这个过程重复进行：每台设备把它刚刚收到的 chunk 继续向右传，同时从左边接收新的 chunk。经过 $N-1$ 步后，每台设备都拥有整个 tensor。

在我们理想化的设定中，每台设备在每一步同时发送一个大小为 $S/N$ 的 chunk，发送带宽为 $W$，总共有 $N-1$ 步，因此 ring all-gather 需要 $\frac{N-1}{N}\frac{S}{W}$ 秒。

接下来，我们来分析 ring reduce-scatter。在 reduce-scatter 中，每台设备 $i$ 一开始持有一个大小为 $S$ 的完整 tensor $\mathbf{x}^{(i)}$。然后我们想计算归约 $\mathbf{y}=\sum_{i=0}^{N-1}\mathbf{x}^{(i)}$，但每台设备 $i$ 最终只得到其中一个大小为 $S/N$ 的 chunk $y_i$。与 ring all-gather 类似，我们先把设备排成一个环。每台设备首先把自己的 tensor $\mathbf{x}^{(i)}$ 分成 $N$ 个 chunk $[\mathbf{x}_0^{(i)},\dots,\mathbf{x}_{N-1}^{(i)}]$，每个大小为 $S/N$。然后我们像 ring all-gather 那样传递 chunk，区别在于：在把 chunk 继续传下去之前，每台设备会把自己的贡献加到该 chunk 上（该 chunk 保存的是部分和）。具体来说：

对于步骤 $t=1,\dots,N-1$，设备 $i$ 执行以下操作：

经过 $N-1$ 步后，设备 $i$ 就得到了 chunk $y_i$ 的完整和，因此 ring reduce-scatter 需要 $\frac{N-1}{N}\frac{S}{W}$ 秒，就像 ring all-gather 一样。

第 41 页

最后，我们来实现 ring all-reduce。在 all-reduce 中，每台设备 $i$ 一开始持有一个大小为 $S$ 的完整 tensor $\mathbf{x}^{(i)}$，最终得到归约结果 $\mathbf{y}=\sum_{i=0}^{N-1}\mathbf{x}^{(i)}$。我们会把 all-reduce 实现为一次 ring reduce-scatter 后接一次 ring all-gather，因此 ring all-reduce 需要 $2\frac{N-1}{N}\frac{S}{W}$ 秒。

问题（alternate_ring_all_reduce）：另一种 ring all-reduce（1 分）

不要用“ring reduce-scatter 后接 ring all-gather”的方式实现 all-reduce，而是使用下面的算法：

对于步骤 $t=1,\dots,N-1$，设备 $i$ 执行以下操作：

在与上面相同的设定下（每台设备的发送带宽为 $W$，每个 $\mathbf{x}^{(i)}$ 的大小为 $S$），这个算法需要多长时间？

Deliverable：用 $S$、$N$ 和 $W$ 表示答案，并给出一句话的理由。

8.2 Analyzing Data Parallel

有了这些通信原语，我们就可以分析并行策略了。我们将分析单个 FFN 层的并行化。回顾一下，给定输入 $\mathbf{x}$，前向传播如下：

$$\mathbf{x}_1 = \mathbf{x}\mathbf{W}_1 \tag{20}$$

$$\mathbf{x}_2 = \mathbf{x}\mathbf{W}_2 \tag{21}$$

$$\mathbf{z} = f(\mathbf{x}_1) * \mathbf{x}_2 \tag{22}$$

$$\mathbf{y} = \mathbf{z}\mathbf{W}_3, \tag{23}$$

其中 $\mathbf{x}$ 的形状为 $(B,D)$，$\mathbf{W}_1$ 和 $\mathbf{W}_2$ 的形状为 $(D,D_{\text{FF}})$，$\mathbf{W}_3$ 的形状为 $(D_{\text{FF}},D)$。$f$ 是逐元素激活函数（例如 SiLU），$*$ 表示逐元素乘法。

显式写出反向传播也会很有用。回顾一下，给定形状为 $(B,D)$ 的 $\mathbf{dy}$，反向传播如下：

$$\mathbf{dz} = \mathbf{dy}\mathbf{W}_3^\top \tag{24}$$

$$\mathbf{dx}_2 = \mathbf{dz} * f(\mathbf{x}_1) \tag{25}$$

$$\mathbf{dx}_1 = \mathbf{dz} * f'(\mathbf{x}_1) * \mathbf{x}_2 \tag{26}$$

$$\mathbf{dx} = \mathbf{dx}_1\mathbf{W}_1^\top + \mathbf{dx}_2\mathbf{W}_2^\top \tag{27}$$

$$\mathbf{dW}_3 = \mathbf{z}^\top\mathbf{dy} \tag{28}$$

$$\mathbf{dW}_2 = \mathbf{x}^\top\mathbf{dx}_2 \tag{29}$$

$$\mathbf{dW}_1 = \mathbf{x}^\top\mathbf{dx}_1, \tag{30}$$

其中 $*$ 表示逐元素乘法。

回顾一下，在使用 $N_{\text{DP}}$ 台设备的数据并行中，我们把输入 $\mathbf{x}$ 分片为 $\mathbf{x}^{(i)}$，其大小为 $(\frac{B}{N_{\text{DP}}}, D)$。前向传播照常进行，不需要任何 collective，产生大小为 $(\frac{B}{N_{\text{DP}}}, D)$ 的 activation $\mathbf{y}^{(i)}$。在

第 42 页

反向传播中，按照 batch 分片后的 activation 照常进行，设备 $i$ 最终得到梯度

$$\mathbf{dW}_3^{(i)} = \mathbf{z}^{(i)\top}\mathbf{dy}^{(i)} \tag{31}$$

$$\mathbf{dW}_2^{(i)} = \mathbf{x}^{(i)\top}\mathbf{dx}_2^{(i)} \tag{32}$$

$$\mathbf{dW}_1^{(i)} = \mathbf{x}^{(i)\top}\mathbf{dx}_1^{(i)}, \tag{33}$$

这里不是对所有 $B$ 个外积累加，而只是对我们输入分片对应的 $B/N_{\text{DP}}$ 个外积做部分和。然后，我们需要跨设备做一次 all-reduce，以得到完整梯度 $\mathbf{dW}_3$、$\mathbf{dW}_2$ 和 $\mathbf{dW}_1$。

问题（data_parallel_calcs）：数据并行计算（3 分）

现在我们已经具备计算“数据并行何时会成为通信瓶颈”所需的一切。令 $C$（单位 FLOP/s）表示设备加速器速度，令 $W$（单位 字节/秒）表示每台设备的发送带宽。然后我们可以计算计算时间和通信时间。由于计算和通信可以重叠，当通信时间大于计算时间时，我们就遇到了瓶颈。我们假设所有权重和 activation 都使用 FP16（即两个字节）。

(a) 在 $N_{\text{DP}}$ 数据并行下，计算反向传播需要多少 FLOPs？你可以忽略所有非 matmul 操作。回顾一下，一个 matmul $(A,B)(B,C)\to(A,C)$ 需要 $2ABC$ FLOPs。

Deliverable：用 $B$、$D$、$D_{\text{FF}}$ 和 $N_{\text{DP}}$ 表示答案，并给出一句话的理由。

(b) 在 $N_{\text{DP}}$ 数据并行下，反向传播需要多少通信时间？

Deliverable：用 $B$、$D$、$D_{\text{FF}}$、$N_{\text{DP}}$ 和 $W$ 的一个子集表示答案，并给出一句话的理由。

(c) 固定其他参数，在遇到通信瓶颈之前，$N_{\text{DP}}$ 最大可以是多少？

Deliverable：给出一个不等式，一边是 $N_{\text{DP}}$，另一边是用 $B$、$D$、$D_{\text{FF}}$、$C$ 和 $W$ 的一个子集表示的表达式，并给出一句话的理由。

8.3 Analyzing Fully Sharded Data Parallel

接下来，我们来分析 FSDP。回顾一下，和 DP 一样，FSDP 会对输入和 activation 的 batch 维进行分片。此外，为了节省内存，我们还会对 optimizer states、梯度以及权重进行分片。我们可以沿任一维度对权重进行分片，在设备 $i$ 上产生分片 $\mathbf{W}_1^{(i)}$、$\mathbf{W}_2^{(i)}$ 和 $\mathbf{W}_3^{(i)}$，每个大小为 $\frac{D D_{\text{FF}}}{N_{\text{FSDP}}}$。

在前向传播中，我们只需在按 batch 分片后的输入上做数据并行的前向传播。但要做到这一点，回顾一下，我们需要先跨设备 all-gather 权重分片：

$$\mathbf{W}_1 = \text{all-gather}\left(\{\mathbf{W}_1^{(i)}\}_{i=0}^{N_{\text{FSDP}}-1}\right) \tag{34}$$

$$\mathbf{W}_2 = \text{all-gather}\left(\{\mathbf{W}_2^{(i)}\}_{i=0}^{N_{\text{FSDP}}-1}\right) \tag{35}$$

第 43 页

$$\boldsymbol{W}_3 = \text{all-gather}\left(\{\boldsymbol{W}_3^{(i)}\}_{i=0}^{N_{\text{FSDP}}-1}\right) \tag{36}$$

（do batch-sharded forward pass）(37)

注意，我们在何时执行 all-gather 上有一定的自由度：我们只需要在某个权重被使用之前完成它的 all-gather，然后就应该把它丢弃，以保持较低的内存开销。为简单起见，本节中我们把三个 all-gather 放在一起列出。

与 forward 一样，在 backward pass 中我们需要先跨设备 all-gather 权重分片。然后就可以做 data parallel 的 backward pass，只是不再需要做 all-reduce，因为每个设备只需要自己那个分片的梯度。因此，我们改为做 reduce-scatter：

$$\boldsymbol{W}_1 = \text{all-gather}\left(\{\boldsymbol{W}_1^{(i)}\}_{i=0}^{N_{\text{FSDP}}-1}\right) \tag{38}$$

$$\boldsymbol{W}_2 = \text{all-gather}\left(\{\boldsymbol{W}_2^{(i)}\}_{i=0}^{N_{\text{FSDP}}-1}\right) \tag{39}$$

$$\boldsymbol{W}_3 = \text{all-gather}\left(\{\boldsymbol{W}_3^{(i)}\}_{i=0}^{N_{\text{FSDP}}-1}\right) \tag{40}$$

（do batch-sharded backward pass）(41)

$$d\boldsymbol{W}_1^{(i)} = \text{reduce-scatter}\left(\{d\boldsymbol{W}_1^{(i)}\}_{i=0}^{N_{\text{FSDP}}-1}\right) \tag{42}$$

$$d\boldsymbol{W}_2^{(i)} = \text{reduce-scatter}\left(\{d\boldsymbol{W}_2^{(i)}\}_{i=0}^{N_{\text{FSDP}}-1}\right) \tag{43}$$

$$d\boldsymbol{W}_3^{(i)} = \text{reduce-scatter}\left(\{d\boldsymbol{W}_3^{(i)}\}_{i=0}^{N_{\text{FSDP}}-1}\right) \tag{44}$$

注意，分片记号（上标 $(i)$）在这里表示的是重载含义：reduce-scatter 的输入是完整权重上的部分和，而输出是分片权重上的完整和。

Problem (fsdp_calcs): Fully sharded data parallel calculations（3 分）

在与 data parallel calculations 相同的设定下，我们来计算 FSDP 何时会成为通信瓶颈。

(a)

使用 $N_{\text{FSDP}}$ 路 FSDP 时，计算 backward pass 需要多少 FLOPs？forward pass 呢？

Deliverable：两个答案，用 $B$、$D$、$D_{\text{FF}}$ 和 $N_{\text{FSDP}}$ 表示，并附两个一句话的论证。

(b)

使用 $N_{\text{FSDP}}$ 路 FSDP 时，backward pass 需要多少通信时间？forward pass 呢？

Deliverable：两个答案，用 $B$、$D$、$D_{\text{FF}}$、$N_{\text{FSDP}}$、$W$ 的一个子集表示，并附两个一句话的论证。

(c)

固定其他参数，$N_{\text{FSDP}}$ 最大可以取到多少，backward pass 才会成为通信瓶颈？forward pass 呢？

第 44 页

Deliverable（接上页 (c)）：两个不等式，一边是 $N_{\text{FSDP}}$，另一边是仅涉及 $B$、$D$、$D_{\text{FF}}$、$C$、$W$ 中一部分的表达式，并附两个一句话的论证。

8.4 分析 Tensor Parallel

在实践中，FSDP 常与一种称为 tensor parallelism (TP) 的并行策略结合使用。在 TP 中，我们把每个权重矩阵的输入维度或输出维度分片到各设备上。按输入维度分片通常称为 “row parallel”，按输出维度分片通常称为 “column parallel”。

具体来说，假设我们要对矩阵乘法 $\boldsymbol{x}\boldsymbol{W}$ 做分片，其中 $\boldsymbol{x}$ 的形状为 $(B, D)$，$\boldsymbol{W}$ 的形状为 $(D, D_{\text{FF}})$。在 column parallel 中，我们得到形状为 $(D, \frac{D_{\text{FF}}}{N_{\text{TP}}})$ 的分片 $\boldsymbol{W}^{(i)}$，并且有

$$\boldsymbol{x}\boldsymbol{W} = \text{all-gather}\left(\{\boldsymbol{x}\boldsymbol{W}^{(i)}\}_{i=0}^{N_{\text{TP}}-1}\right). \tag{45}$$

另一方面，对于 row parallel，我们得到形状为 $(\frac{D}{N_{\text{TP}}}, D_{\text{FF}})$ 的分片 $\boldsymbol{W}^{(i)}$，并且每个设备在做矩阵乘法之前，还会把输入 $\boldsymbol{x}$ 切分成形状为 $(B, \frac{D}{N_{\text{TP}}})$ 的切片 $\boldsymbol{x}^{(i)}$。于是有

$$\boldsymbol{x}\boldsymbol{W} = \text{all-reduce}\left(\{\boldsymbol{x}^{(i)}\boldsymbol{W}^{(i)}\}_{i=0}^{N_{\text{TP}}-1}\right). \tag{46}$$

为了并行化我们的 FFN，我们会使用一种特定的 tensor parallel 配置：$\boldsymbol{W}_1$ 和 $\boldsymbol{W}_2$ 采用 column parallel（按输出维度分片），而 $\boldsymbol{W}_3$ 采用 row parallel（按输入维度分片）。由于 row parallel 的权重只需要输入的一个切片，这种配置让我们可以跳过 column parallel 权重之后的 all-gather。给定大小为 $(B, D)$ 的输入 $\boldsymbol{x}$，该策略给出如下 forward pass：

$$\boldsymbol{x}_1^{(i)} = \boldsymbol{x}\boldsymbol{W}_1^{(i)} \tag{47}$$

$$\boldsymbol{x}_2^{(i)} = \boldsymbol{x}\boldsymbol{W}_2^{(i)} \tag{48}$$

$$\boldsymbol{z}^{(i)} = f(\boldsymbol{x}_1^{(i)}) * \boldsymbol{x}_2^{(i)} \tag{49}$$

$$\boldsymbol{y}^{(i)} = \boldsymbol{z}^{(i)}\boldsymbol{W}_3^{(i)} \tag{50}$$

$$\boldsymbol{y} = \text{all-reduce}\left(\{\boldsymbol{y}^{(i)}\}_{i=0}^{N_{\text{TP}}-1}\right), \tag{51}$$

其中 $\boldsymbol{W}_1^{(i)}$ 和 $\boldsymbol{W}_2^{(i)}$ 的形状为 $(D, \frac{D_{\text{FF}}}{N_{\text{TP}}})$，$\boldsymbol{W}_3^{(i)}$ 的形状为 $(\frac{D_{\text{FF}}}{N_{\text{TP}}}, D)$。

Problem (tp_calcs): Tensor parallel calculations（4 分）

在与 DP 和 FSDP calculations 相同的设定下，我们来计算 TP 何时会成为通信瓶颈。

(a)

给定大小为 $(B, D)$ 的输入 $d\boldsymbol{y}$，写出上述 tensor parallel 策略的 backward pass（其中 $\boldsymbol{W}_1^{(i)}$ 和 $\boldsymbol{W}_2^{(i)}$ 的形状为 $(D, \frac{D_{\text{FF}}}{N_{\text{TP}}})$，$\boldsymbol{W}_3^{(i)}$ 的形状为 $(\frac{D_{\text{FF}}}{N_{\text{TP}}}, D)$）。

Deliverable：一系列描述 backward pass 的方程，用 $d\boldsymbol{y}$、分片权重（$\boldsymbol{W}_1^{(i)}$、$\boldsymbol{W}_2^{(i)}$、$\boldsymbol{W}_3^{(i)}$）、forward pass 中保存的 activations（$\boldsymbol{x}$、$\boldsymbol{x}_1^{(i)}$、$\boldsymbol{x}_2^{(i)}$、$\boldsymbol{z}^{(i)}$、$\boldsymbol{y}^{(i)}$）、通信原语，以及你想定义的任何中间变量表示。这些方程应给出每个设备的梯度 $d\boldsymbol{W}_1^{(i)}$、$d\boldsymbol{W}_2^{(i)}$、$d\boldsymbol{W}_3^{(i)}$，以及 backward pass 的输出 $d\boldsymbol{x}$。可以随意参考 8.2 节中未分片的 backward pass 并加以修改。

(b)

使用 $N_{\text{TP}}$ 路 TP 时，计算 forward pass 需要多少 FLOPs？backward pass 呢？

第 45 页

Deliverable（接上页 (b)）：两个答案，用 $B$、$D$、$D_{\text{FF}}$ 和 $N_{\text{TP}}$ 表示，并附两个一句话的论证。

(c)

使用 $N_{\text{TP}}$ 路 TP 时，forward pass 需要多少通信时间？backward pass 呢？

Deliverable：两个答案，用 $B$、$D$、$D_{\text{FF}}$、$N_{\text{TP}}$、$W$ 的一个子集表示，并附两个一句话的论证。

(d)

固定其他参数，$N_{\text{TP}}$ 最大可以取到多少，backward pass 才会成为通信瓶颈？forward pass 呢？

Deliverable：两个不等式，一边是 $N_{\text{TP}}$，另一边是仅涉及 $B$、$D$、$D_{\text{FF}}$、$C$、$W$ 中一部分的表达式，并附两个一句话的论证。

8.5 2D Parallelism (FSDP + TP)

我们终于可以组合各种并行策略了！本节将讨论如何最优地组合 FSDP 与 TP。提示：在前面的小节中，你应该已经发现，batch size 和模型规模参数限制了你最多能扩展到多少设备；把这两者都调大，就能继续扩展设备而不至于成为通信瓶颈。不幸的是，batch size 超过某个点后开始降低性能，因为梯度噪声显著变小，失去了 SGD 的隐式正则化性质；这个点通常称为 “critical batch size”。而 scaling laws 常常告诉我们模型需要多大（那将是你下次作业的任务！）。

本节考虑一个简化设定：有人把全部问题参数（batch size、模型规模、带宽、加速器速度）都交给你。你的任务是选择一种 FSDP 与 TP 的配置，使其能扩展到尽可能多的设备，同时保持 compute-bound 而不是 communication-bound。

我们先走一遍 FSDP 与 TP 结合的机制。每个设备会有一个 TP rank $i = 0, \dots, N_{\text{TP}}-1$ 和一个 FSDP rank $j = 0, \dots, N_{\text{FSDP}}-1$，总共形成 $N = N_{\text{TP}}N_{\text{FSDP}}$ 个设备的 2D 网格。按照 TP，我们先把 $\boldsymbol{W}_1$ 和 $\boldsymbol{W}_2$ 沿输出维度分片，把 $\boldsymbol{W}_3$ 沿输入维度分片。因此，我们必须在 activations 上插入一个 TP 风格的 all-reduce。接着应用 FSDP：我们会把输入的 batch 维度分片，并且把每个权重矩阵沿未被 TP 分片的那个维度进一步分片。然后，在执行 TP 风格的 forward/backward pass 之前，需要在权重上插入 FSDP 风格的 all-gather；在 TP 风格的 backward pass 之后，需要在权重梯度上做 reduce-scatter。

结果是，每个设备 $(i, j)$ 持有形状为 $(\frac{D}{N_{\text{FSDP}}}, \frac{D_{\text{FF}}}{N_{\text{TP}}})$ 的权重分片 $\boldsymbol{W}_1^{(i,j)}$ 和 $\boldsymbol{W}_2^{(i,j)}$，以及形状为 $(\frac{D_{\text{FF}}}{N_{\text{TP}}}, \frac{D}{N_{\text{FSDP}}})$ 的 $\boldsymbol{W}_3^{(i,j)}$。给定大小为 $(\frac{B}{N_{\text{FSDP}}}, D)$ 的 batch 分片输入 $\boldsymbol{x}^{(j)}$，我们可以把 forward pass 写成如下形式：

$$\boldsymbol{W}_1^{(i)} = \text{all-gather}\left(\{\boldsymbol{W}_1^{(i,j)}\}_{j=0}^{N_{\text{FSDP}}-1}\right) \tag{52}$$

$$\boldsymbol{W}_2^{(i)} = \text{all-gather}\left(\{\boldsymbol{W}_2^{(i,j)}\}_{j=0}^{N_{\text{FSDP}}-1}\right) \tag{53}$$

$$\boldsymbol{W}_3^{(i)} = \text{all-gather}\left(\{\boldsymbol{W}_3^{(i,j)}\}_{j=0}^{N_{\text{FSDP}}-1}\right) \tag{54}$$

$$\boldsymbol{x}_1^{(i,j)} = \boldsymbol{x}^{(j)}\boldsymbol{W}_1^{(i)} \tag{55}$$

$$\boldsymbol{x}_2^{(i,j)} = \boldsymbol{x}^{(j)}\boldsymbol{W}_2^{(i)} \tag{56}$$

第 46 页

$$\begin{aligned}

z^{(i,j)} &= f(x_1^{(i,j)}) \ast x_2^{(i,j)} \tag{57}\ y^{(i,j)} &= z^{(i,j)} W_3^{(i)} \tag{58}\ y^{(j)} &= \text{all-reduce}\left({y^{(i,j)}}{i=0}^{N{\text{TP}}-1}\right), \tag{59}

\end{aligned}$$

最终得到 batch-sharded（按 batch 分片）输出 $y^{(j)}$，大小为 $(B/N_{\text{FSDP}}, D)$。为简洁起见，本节省略反向传播，只关注前向传播。但此时，你应该已经掌握了自行写出它所需的全部信息。

Problem (fsdp_tp_calcs)：2D 并行计算（6 分）

在与目前这些计算相同的设置下，我们来计算 2D 并行何时会成为瓶颈。

(a) 在 $N_{\text{FSDP}}$ 路 FSDP + $N_{\text{TP}}$ 路 TP 下，计算前向传播需要多少 FLOPs？

Deliverable: 用 $B$、$D$、$D_{\text{FF}}$、$N_{\text{FSDP}}$、$N_{\text{TP}}$ 表示的答案，并附一句话说明理由。

(b) 在 $N_{\text{FSDP}}$ 路 FSDP + $N_{\text{TP}}$ 路 TP 下，前向传播需要多少通信时间？假设沿每个轴的通信可以重叠（换言之，FSDP 轴上的 collective 可以与 TP 轴上的 collective 重叠）。

Deliverable: 用 $B$、$D$、$D_{\text{FF}}$、$N_{\text{FSDP}}$、$N_{\text{TP}}$、$W$ 的一个子集表示的答案，并附一句话说明理由。提示：答案应写成两个量（FSDP 与 TP 的 collective 开销）之间的 max，因为二者可以重叠。

(c) 在 $N_{\text{TP}}$ 和 $N_{\text{FSDP}}$ 的最优设置下，$N = N_{\text{TP}} N_{\text{FSDP}}$ 最大能到多少，之后前向传播会受通信瓶颈限制？

Deliverable: 一个不等式，一边是 $N$，另一边是用 $B$、$D$、$D_{\text{FF}}$、$C$、$W$ 的一个子集表示的表达式，并附几句话和公式作为理由。

(d) 现在假设 FSDP 轴和 TP 轴的 collective 不能重叠，因为它们共享相同的网络资源。在 $N_{\text{TP}}$ 和 $N_{\text{FSDP}}$ 的最优设置下，$N = N_{\text{TP}} N_{\text{FSDP}}$ 最大能到多少，之后前向传播会受通信瓶颈限制？不用担心把 $N_{\text{TP}}$ 和 $N_{\text{FSDP}}$ 截断为整数。

Deliverable: 一个不等式，一边是 $N$，另一边是用 $B$、$D$、$D_{\text{FF}}$、$C$、$W$ 的一个子集表示的表达式，并附几句话和公式作为理由。

9 Leaderboard（排行榜）

Assignment 2 的 leaderboard 将测试一个 8B 模型完整训练步的速度。我们挑战你对自己的代码做 benchmark，并优化内存和运行时间，使用你能想到的任何技巧。关键限制是：你不能改变模型的输入/输出行为。你的实现将对照 cs336_basics 目录中的模型进行测试。你的输入将在 BF16 和 causal masking 下测试，并且必须通过与常规实现相同的测试。实现也必须是你自己的，你不能使用或复制已有实现。你的计时应在两张 B200 GPU 上使用 batch size 2、sequence length 32,768 的样本测量。故意让模型难以放入内存。完整配置见下方代码。我们

第 47 页

我们将验证前 5–10 份提交的正确性和性能。我们用来给你的实现计时的测试如下：

class Config:
    ctx_len = 32768
    vocab_size = 151936
    d_model = 4096
    d_ff = 11008
    num_layers = 34
    num_heads = 32
    torch_dtype = torch.bfloat16
    is_causal = True
    batch_size = 2
cfg = Config()
def test_timing_forward_backward():
    labels, targets = torch.randint(high=cfg.vocab_size, size=(2, cfg.batch_size, cfg.ctx_len))
    model = BasicsTransformerLM(Config())
    optimizer = AdamW(model.parameters())
    def train_step():
        optimizer.zero_grad(set_to_none=True)
        res = model(labels)
        loss = cross_entropy(res, targets).sum()
        loss.backward()
        optimizer.step()
    timing_results = triton.testing.do_bench(train_step, rep=30_000, warmup=10_000)
    print(timing_results)

为便于测试，你可以把重复次数和 warmup 时间（单位为 ms）缩短一些。

一些改进思路，以及确保模型能放得下的思路：

第 48 页

Problem (leaderboard)：Leaderboard：最快的训练步（10 分）

基准测试将在两张 B200 GPU 上以 batch size 2 运行。你的提交将按完成一个完整训练步的 wall-clock 时间评估：前向传播、loss、反向传播和 AdamW 更新。

从空的 PyTorch/Triton cache 开始，你的 benchmark 运行必须在 10 分钟内完成，因此要小心过于激进的 torch.compile 和 Triton autotuning。

Deliverable: 使用 AdamW 的完整前向+反向训练步的最佳 wall-clock 时间。

我们期望 leaderboard 提交能超过 10 秒的 naïve baseline。

将你的结果提交到 leaderboard：github.com/stanford-cs336/assignment2-systems-leaderboard

Bibliography（参考文献）

材料

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

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

自查