[斯坦福CS336]作业二：系统与并行计算

1 作业概述

本次作业中，你将亲自动手实践提升单 GPU 训练速度和将训练扩展到多 GPU 的方法。

需实现的内容

1. 基准测试与性能分析工具
2. FlashAttention-2 的 Triton 内核
3. 分布式数据并行训练
4. 优化器状态分片

作业地址： Assignment2-systems GitHub仓库

接下来我将分享完成作业的一部分细节和心得。

2 性能分析与基准测试

2.1 动机

在进行任何优化之前，先对程序进行性能分析是很有必要的，这能帮我们明确资源（如时间和内存）的消耗重点。否则，我们可能会在对整体性能影响不大的模块上浪费精力，无法带来显著的端到端性能提升。

我们将实现三种性能评估方案：

(a) 基于 Python 标准库的简单端到端基准测试，用于统计前向和反向传播耗时；

(b) 使用 NVIDIA Nsight Systems 工具分析计算过程，明确 CPU 和 GPU 上各操作的耗时分布；

(c) 内存使用情况分析。

2.2 模型规格

在本次作业中，我们将对不同规格的模型进行基准测试和性能分析，以观察模型规模对性能的影响。所有模型的词汇量均设为 10000，批次大小设为 4，仅调整上下文长度。

不同模型规格参数如下：

2.3 端到端基准测试

首先，我们对模型进行最基础的性能分析 —— 统计前向和反向传播的耗时。由于我们只关注速度和内存，实验中可以使用随机初始化的权重和数据。

对于 GPU 代码进行基准测试时，需注意 CUDA 调用是异步的：调用内核后 CPU 会立即继续执行，因此直接测量函数返回时间无法反映 GPU 实际计算耗时。为了准确测量内核运行时间，应在测试前后调用 torch.cuda.synchronize() 以确保 GPU 操作完成。这是构建可靠性能分析工具的基础。

作业题（benchmarking_script）：4 分

(a) 编写一个脚本，对模型的前向和反向传播进行基础的端到端基准测试。具体来说，你的脚本需要支持以下功能：

• 根据给定的超参数（如层数）初始化模型；
• 生成随机批次的测试数据；
• 先执行 w 次预热步骤（不计入正式计时），然后统计 n 次迭代的耗时（可通过参数控制仅测试前向传播，或同时测试前向和反向传播）；
• 计时工具可使用 Python 的 timeit 模块（例如 timeit 函数，或 timeit.default_timer()——该函数调用系统最高精度的时钟，比 time.time() 更适合基准测试）；
• 每次迭代后调用 torch.cuda.synchronize()。

提交要求：一个脚本，能够根据给定超参数初始化基础 Transformer 模型，生成随机批次数据，并统计前向和反向传播的耗时。

答案：benchmarking_script.py

(b) 针对 1.1.2 节中描述的所有模型规格，统计其前向和反向传播的耗时。预热步骤设为 5 次，正式测试迭代 10 次，计算耗时的平均值和标准差。前向传播耗时多少？反向传播呢？测试结果的波动性大吗？标准差是否较小？

提交要求：1-2 句话，简要说明你的测试结果。

答案：前向与反向传播耗时随模型规格增大而显著增加（反向耗时约为前向的 1.8 倍）。测试结果中，仅 large 模型在前向传播阶段表现出异常高的波动（标准差约 15.6ms），其余所有模型标准差均低于 1ms，整体稳定性极高。

(c) 基准测试的一个常见误区是不执行预热步骤。请重复上述实验，但不进行预热，观察这会对结果产生什么影响？你认为原因是什么？另外，尝试将预热步骤设为 1 或 2 次，结果为何仍然会有差异？

提交要求：2-3 句话，回答上述问题。

答案：不执行预热会导致平均耗时显著增加且方差极大，这是因为测试数据包含了 CUDA 上下文初始化、显存分配及内核编译 等高昂的一次性开销。仅进行 1-2 次预热结果仍有差异，是因为 GPU 往往需要更多迭代才能 将时钟频率提升至高性能状态（Clock Boost） 并使硬件调度达到稳态，短期预热不足以完全消除这种硬件延迟。

未执行预热结果：

预热 2 次后结果：

2.4 Nsight Systems 性能分析工具