CS336 LLM from Scratch Lab2 writeup
Last updated on February 28, 2026 1:26 AM
概述
本系列为斯坦福 Stanford CS336 | Language Modeling from Scratch 课程的作业笔记。
该作业相比第一次作业而言更为繁琐，需要跑大量的 benchmark，对算力也有相当大的需求。在此衷心感谢北京大学 Linux 俱乐部提供的算力资源和组织讨论平台。
本 lab 的相关代码仓库在此处，将其 clone 到本地即可。该实验使用 uv 管理环境，代码需要在 cs336_systems/ 中完成，cs336-basics 中有官方实现的 lab1 的简单模型，test/ 中为测试点，在完成各个功能的时候需要同步实现 adapters.py 里面的接口。
这个 lab 需要实现/做的事：
- Benchmarking and profiling:
- End-to-end benchmarking
- Nsight systems profiler
- BF16 mixed precision benchmarking
- Memory profiling
- Flash Attention 2
- Benchmarking PyTorch attention
- JIT-compiled attention benchmarking
- 用 PyTorch 实现 FA2 的 forward 和 backward
- 用 Triton 实现 FA2 的 forward
- 各种 attention 的 benchmarking
- DDP 训练
- 试验单节点上的通信性能（6 GPUs）
- Naive DDP 实现（把 batch 分发）和 benchmarking
- 在 naive 的基础上做优化（flatten、异步通信、bucket）
- 4D 并行的计算分析
- 优化器状态分片的实现以及 resource accounting
整体来说任务量丝毫不亚于 lab1，不过这几种并行以及 FA2 的实现对初探 GPU 编程的人（比如我）来说还是比较有价值的。
Profiling and Benchmarking
End-to-End Benchmarking
这里需要写一个脚本，关于 model size、context length、warm-up 次数做消融实验，统计前向和后向传播的耗时。
需要注意的点基本在作业文档里面说清楚了，比如说在正式开始跑之前需要做几次 warm-up 以免一些初始化干扰计时结果，以及由于 GPU 是异步执行，所以在每一个 step 之后都需要 torch.cuda.synchronize() 以确保计时结果精准。
这种略为重复劳动的工作直接用 codex 来写就好了，codex 会把命令行参数什么的都写好，然后用一个 bash 脚本或者 sbatch 脚本（如果有 slurm 集群）来 sweep 就好了。
benchmarking_script
(a) 脚本实现
我实现了一个可参数化的 benchmark 脚本，支持按给定超参数初始化模型、生成随机输入 batch、执行 warm-up 步骤后再进行正式计时，并在每步结束后调用 torch.cuda.synchronize() 以避免 CUDA 异步导致的计时偏差。脚本支持 forward、forward-backward 和 train-step 三种模式，其中本题使用 forward-backward 来分别统计前向与反向耗时，并将结果输出为 JSON 便于后续汇总。
(b) benchmarking
实验环境：
NVIDIA H800 PCIe，batch_size=4，warm-up=5，measure=10单位：ms（均值 ± 标准差）
| Model Size | Context | Forward | Backward |
|---|---|---|---|
| small | 128 | 21.57 ± 0.13 | 23.02 ± 1.06 |
| medium | 128 | 45.03 ± 1.10 | 46.51 ± 0.47 |
| large | 128 | 64.98 ± 0.69 | 93.93 ± 0.05 |
| xl | 128 | 87.18 ± 1.23 | 175.88 ± 0.18 |
| 2.7b | 128 | 119.78 ± 0.71 | 264.10 ± 0.75 |
对于 context 变化的补充（同样 warm-up=5, measure=10）：
| Model Size | Context=256 (F/B) | Context=512 (F/B) |
|---|---|---|
| small | 21.86 / 26.00 | 27.24 / 54.23 |
| medium | 44.37 / 81.12 | 84.28 / 164.97 |
| large | 86.51 / 170.95 | 161.45 / 342.85 |
| xl | 152.54 / 321.49 | 326.48 / 665.57 |
| 2.7b | 232.57 / 487.16 | 490.07 / 1003.19 |
结论：随着模型规模和 context length 增大，forward/backward 耗时都显著上升；在小模型上 backward 仅略高于 forward，但在大模型（尤其 xl、2.7b）上 backward 明显更重。整体波动较小，大多数配置的标准差维持在低毫秒量级。
© warm-up
单位：ms（均值 ± 标准差）
实验环境：5090
| Model Size | Warm-up | Forward | Backward |
|---|---|---|---|
| small | 0 | 71.46 ± 138.07 | 30.60 ± 37.73 |
| small | 1 | 33.54 ± 43.82 | 22.82 ± 3.18 |
| small | 2 | 32.45 ± 39.19 | 25.24 ± 6.83 |
| medium | 0 | 104.75 ± 206.12 | 57.04 ± 48.73 |
| medium | 1 | 47.29 ± 9.32 | 47.87 ± 6.53 |
| medium | 2 | 32.89 ± 1.52 | 34.60 ± 1.13 |
| large | 0 | 113.16 ± 189.20 | 77.75 ± 48.34 |
| large | 1 | 52.67 ± 5.29 | 64.52 ± 3.04 |
| large | 2 | 53.73 ± 1.56 | 65.99 ± 0.82 |
| xl | 0 | 122.53 ± 182.62 | 128.62 ± 35.63 |
| xl | 1 | 63.73 ± 8.23 | 118.83 ± 0.23 |
| xl | 2 | 65.76 ± 4.05 | 121.82 ± 0.08 |
| 2.7b | 0 | 137.96 ± 181.38 | 179.41 ± 49.02 |
| 2.7b | 1 | 81.74 ± 0.73 | 165.86 ± 0.90 |
| 2.7b | 2 | 81.58 ± 0.09 | 165.02 ± 0.69 |
结论：不做 warm-up（warm-up=0）时，forward/backward 的平均耗时明显偏高，而且标准差显著变大，结果不稳定。主要原因是初始迭代会包含额外开销（如 CUDA 上下文初始化、kernel 选择与缓存、内存分配器预热等），导致首批 step 不能代表稳态性能。即使做 1-2 步 warm-up，结果仍可能不同，因为有些配置需要更多步才能完全进入稳态，且系统噪声也会带来残余波动。
Nsight Systems Profiler
为了更深入探究更具体的各个函数/内核的运行时间情况，需要使用更细致的 profiling 方法，且由于 GPU 是异步执行的所以需要用 CUDA 自己的工具来进行分析。这需要运行脚本的机子上有安装 nsys，注意有可能实际已经安装但没有在 PATH 里，具体可以让 AI 设法帮忙解决。总之，运行如下指令：
1 | |
就可以跑出 result.nsys-rep 这样的 profiling 文件。之后就需要在本地打开 Nsight Systems 的桌面客户端来进行分析。
同时，还可以用 nvtx 标记来单独区分代码中特定部分（比如说，标记出计算 attention score 那一部分的）。
跑出来后，可以直观看到如一次 forward 的耗时，然后将下方视图切换为 “Stats System View” 并选择 “CUDA GPU Kernel Summary” 便可以查看各个 CUDA Kernel 被调用的次数以及运行耗时占比。在上方的 timeline 中选择对应事件右击并 “Apply Filter” 后可以查看对应部分的 kernel summary，满足题目中分别分析 forward 和 backward 被调用次数最多的内核之类问题的需求。
nsys_profile
(a) forward pass
| Model Size | Context | Forward (Python) | Forward (nsight) |
|---|---|---|---|
| small | 128 | 15.57 ± 0.41 | 18.382 |
| medium | 128 | 29.02 ± 0.09 | 38.213 |
| large | 128 | 43.23 ± 0.21 | 54.855 |
| xl | 128 | 59.66 ± 0.03 | 74.066 |
| 2.7b | 128 | 81.04 ± 0.07 | 53.136 |
还是比较不一样的，推测是因为 kernel 的问题。
(b) kernel
| Model Size | Context | Kernel | Time | Invoked |
|---|---|---|---|---|
| small | 128 | void cutlass::Kernel2<cutlass_80_simt_sgemm_128x128_8x4_tn_align1>(T1::Params) |
55.0% | 60 |
| medium | 128 | void cutlass::Kernel2<cutlass_80_simt_sgemm_128x128_8x4_tn_align1>(T1::Params) |
45.2% | 120 |
| large | 128 | void cutlass::Kernel2<cutlass_80_simt_sgemm_128x256_8x4_tn_align1>(T1::Params) |
52.6% | 107 |
| xl | 128 | void cutlass::Kernel2<cutlass_80_simt_sgemm_128x256_8x4_tn_align1>(T1::Params) |
56.5% | 143 |
| 2.7b | 128 | void cutlass::Kernel2<cutlass_80_simt_sgemm_128x256_8x4_tn_align1>(T1::Params) |
90.2% | 148 |
确实还是同一个 kernel（各个配置下）在 backward 中耗最长时间，但是比例明显降低了。
© other kernel
比如 element-wise 乘法的 kernel。
(d) training
矩阵乘法的比例明显下降了，从之前的将近一半下降到只有 10-20%。处理 element-wise 的乘法的 kernel 比例上升。
(e) softmax vs. mm
可以发现 softmax 还是需要相当时间的，略低于计算 attention score 的时间，略低于 final matmul 的 二倍。