336-hw2

See also README in hw2 repo for commands.

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))

一次 & 多次 warmup 似乎没区别.

2.1.4 nsys

(我只用了 medium, context = 512.)

(a) What is the total time spent on your forward pass? Does it match what we had measured before with the Python standard library?

1. forward: 约 250ms. 很 match (python 是 254ms).

(b) What CUDA kernel takes the most cumulative GPU time during the forward pass? How many times is this kernel invoked during a single forward pass of your model? Is it the same kernel that takes the most runtime when you do both forward and backward passes?

1. forward 中，最耗时 ampere_sgemm_128x128_tn (49ms). 143 times.
2. forward + backward 的话最耗时就是 backward 的 indexing_backward_kernel_stride_1 (149ms, 24 times)

(c) What other kernels besides matrix multiplies do you see accounting for non-trivial CUDA runtime in the forward pass?

1. 主要是 ampere_sgemm_128x32_tn (49ms), write_indices (10.4ms), vectorized_elementwise_kernel (7.2ms)

(d) Profile running one complete training step with your implementation of AdamW. How does the fraction of time spent on matrix multiplication change, compared to doing inference (forward pass only)? How about other kernels?

1. full step 相比 forward 矩乘占比明显降低. indexing_backward_kernel_stride_1 (梯度累加回原始张量）显著上升.

(e) Compare the runtime of the softmax operation versus the matrix multiplication operations within the self-attention layer of your model during a forward pass. How does the difference in runtimes compare to the difference in FLOPs?

1. attn softmax 耗时 9ms. attn matmul 耗时 10ms. FLOPS 则是 4.29e9 vs 8.39e7（手算）. 由于 softmax 不是矩阵乘法而是多次 log sum exp 运算，Arithmetic Intensity 很低.

2.1.5 mixed_precision_accumulation

Q: Run the following code and comment on the accuracy of the results.

1. float32 直接计算全部，是 (10.0001) 用 float16 存储总和的精度非常糟糕 (9.9531). 用 float16 存储中间变量而 float32 存储总和就还行. (10.0021).

2.1.5 benchmarking_mixed_precision

1
from torch import nn
2
class ToyModel(nn.Module):
3
    def __init__(self, in_features: int, out_features: int):
4
        super().__init__()
5
        self.fc1 = nn.Linear(in_features, 10, bias=False)
6
        self.ln = nn.LayerNorm(10)
7
        self.fc2 = nn.Linear(10, out_features, bias=False)
8
        self.relu = nn.ReLU()
9
    def forward(self, x):
10
        x = self.relu(self.fc1(x))
11
        x = self.ln(x)
12
        x = self.fc2(x)
13
        return x

(a) Using autocast, what are the data types of:

• the model parameters within the autocast context? fp32
• the output of the first feed-forward layer (ToyModel.fc1)? input=[torch.float32], output=[torch.bfloat16]
• the output of layer norm (ToyModel.ln)? output=[torch.float32] 因为涉及均值和方差计算
• the model’s predicted logits? bf16
• the loss? fp32
• the model’s gradients? fp32

(b) What parts of layer normalization are sensitive to mixed precision? 均值和方差.

If we use BF16 instead of FP16, do we still need to treat layer normalization differently? Why or why not? 还是 FP32，因为方差对尾数也敏感 (例如方差可能很小)

(c) … Compare the results of using full precision versus mixed precision, and comment on any trends as model size changes.

nullcontext: 0.877, 0.191

autocast bf16: 0.481, 0.023

2.1.6

(a) What do your memory timelines look like? Can you tell which stage is running based on the peaks you see?

Deliverable: Two images of the “Active memory timeline” of an xl model, from the memory_viz tool: one for the forward pass, and one for running a full training step.

full 上下文 128: 我用了 large | 1280 | 5120 | 36 | 20 | 模型，上下文 128 & 1024.

full 上下文 1024: 此图中尖端是 forward. 随着 backward 会逐渐释放.

forward(no grad) 上下文 1024: 有无数尖峰.

forward(no grad) 上下文 128:

(b) What is the peak memory usage of each context length when doing a forward pass? What about when doing a full training step?

full step:

• context 128: 18G
• context 1024: 60G

forward no grad:

• context 128: 3.7G
• context 1024: 5.6G

(c) Does mixed-precision significantly affect memory usage?

• full: mix 50G. 普通 60G.
• forward: mix 5.3G. 普通 5.6G. 确实降低了很多

(d) What is the size of a tensor of activations in the Transformer residual stream, in single-precision?

考虑 large 模型. context 1024. [b=4, l=1024, d=1280] * 4 / 1024 / 1024 = 20MB

(e) What is the size of the largest allocations shown? Looking through the stack trace, can you tell where those allocations come from?

最大分配 1.2G. 主要是 scaled_dot_product_attention 中创建 attn. （这里找到个 bug，我把 mask expand_as attn 浪费了很多显存，修复后峰值显存 61.8G -> 58.6G）

(f) How much memory was allocated during the forward pass, and how much memory usage changes for every TransformerBlock in the backward pass, calculate how much memory the produced gradient tensors for a TransformerBlock take. Does the result match what you expect?

1. 根据统计结果，with grad: 一个 transformer block forward: 20.83 - 19.54 = 1.29 GB
2. 和期望结果 1244MB 差不多