336-hw2 - How to?

2.1.3 End-to-End Benchmarking

Q: 写 benchmark 脚本：按超参建 Transformer，生成随机 batch，支持 forward only、forward+backward、forward+backward+optimizer step。每步后 torch.cuda.synchronize()。

对表中模型大小，用 5 warmup steps，测 10 measurement steps，报告均值/标准差。再测无 warmup、1 warmup、2 warmup。

A: | small | 768 | 3072 | 12 | 12 |

→(mean, std) ↓(forward, for+back, for+back+opt)

1
(np.float64(0.0545499186962843), np.float64(9.301123529883161e-05))
2
(np.float64(0.230664774030447), np.float64(0.020122250377148584))
3
(np.float64(0.24185883365571498), np.float64(0.030474347009177154))

| medium | 1024 | 4096 | 24 | 16 |:

1
(np.float64(0.16028596945106982), np.float64(0.00011035529741210624))
2
(np.float64(0.6570748724043369), np.float64(0.03049904281007697))
3
(np.float64(0.6851631578058004), np.float64(0.03967747669197852))

| large | 1280 | 5120 | 36 | 20 |:

1
(np.float64(0.3304705709218979), np.float64(0.0023439186241350617))
2
(np.float64(1.3425680793821813), np.float64(0.04551350800571191))
3
(np.float64(1.3851026877760888), np.float64(0.036186751000143125))

| xl | 2560 | 10240 | 32 | 32 |: 换成了 H200. 初始化很慢，大概要 30s，print 了一下发现没有瓶颈就是慢.

1
(np.float64(0.37269446402788164), np.float64(0.020297942534701095))
2
(np.float64(1.3736595837865024), np.float64(0.020598247928436356))
3
(np.float64(1.4584950204007328), np.float64(0.043012806955922626))

| 10B | 4608 | 12288 | 50 | 36 |: 爆显存.

warmup: medium 模型 H200

1
--warmup 0 --outer 10 --inner 1
2
(np.float64(0.330410421686247), np.float64(0.08759358961566299))
3
# 即使销毁模型，只要进程没销毁，warmup 就持续有效.
4

5
--warmup 0 --outer 1 --inner 1
6
(np.float64(0.5848983488976955), np.float64(0.0))
7
(np.float64(0.7329192743636668), np.float64(0.0))
8

9
--warmup 1 --outer 1 --inner 1
10
(np.float64(0.298882813192904), np.float64(0.0))
11
(np.float64(0.2970326286740601), np.float64(0.0))
12

13
--warmup 2 --outer 1 --inner 1
14
(np.float64(0.2952738180756569), np.float64(0.0))
15
(np.float64(0.35315717151388526), np.float64(0.0))
4 collapsed lines
16

17
--warmup 5 --outer 1 --inner 1
18
(np.float64(0.296027516014874), np.float64(0.0))
19
(np.float64(0.2957768542692065), np.float64(0.0))

full command

1
uv run cs336_systems/profiles/lm.py --d_model 1024 --d_ff 4096 --num_layers 24 --num_heads 16 --mode 2 --warmup 5 --outer 1 --inner 1