CS336 on Lorn

[CS336] Assignment 2：Systems

Fri, 09 Oct 2026 10:30:00 +0800

在上一个Assignment中，我们重头搭建了一个简单的大语言模型以及其训练到推理的全过程，在本节中，我们将进一步从计算机系统的角度，对其核心部件进行系统优化。

环境搭建

我们首先将仓库clone下来:

1

git clone https://github.com/stanford-cs336/assignment2-systems.git

然后其中有两个选择，一个是使用自己在Assignment1中实现的LM框架，另一个是使用课程提供的LM框架。这里建议使用自己的，更加熟悉一点。

若选用自己的LM框架，则需要在 pyproject.toml 中做一个简单的更改:

1

cs336-basics = { path = "./cs336-basics", editable = true } # Change this path to your assignment1-basics repo you want to use your own implementation!

将这一行中的地址定位到你实现的LM框架即可。

搭建完毕后，可以通过下面的指令进行测试:

1
2


uv run python
>>> import cs336_basics

若不报错，说明没有问题

Benchmarking

在这一小节中，我们将进行模型的性能分析与基准测试。

具体而言，是如下三种性能评估路径：

简单的End to End基准测试，即使用Python标准库对前向和反向传播进行计时
计算性能分析，我们将使用NVIDIA Nsight Systems工具分析计算过程，了解时间是如何分布在CPU和GPU的各个操作上的。
内存使用分析：对内存使用情况进行性能分析

此外，我们使用的模型规模如下：

词表大小为：10,000

batch_size:4

简单的端到端基准测试

现在我们实现一个简单的性能评估脚本。由于我们会测试模型的多种变体(例如替换精度，替换层结构等)，因此我们考虑使用命令行参数的方式来支持这些变体，以便后续运行更加方便。

该测试只对模型进行性能分析：即对前向/反向传播进行计时

由于我们只测试速度和内存，因此可以使用随机权重和随机数据。这里在测试时需要特别注意：

对于GPU代码的基准测试，一个重要的注意点是:CUDA调用是异步的

也就是说，当我们调用一个CUDA Kernel时，函数会立即返回控制权，而不会等待Kernel真正执行完成。这样CPU可以继续执行其他的代码，而GPU在后台执行Kernel。

这意味着，如果我们直接测量这个Kernel调用返回所花费的时间，我们就不能得到GPU实际执行该操作的时间。

在Pytorch中，我们可以通过

1

torch.cuda.synchronize()

该函数会等待所有 GPU kernel 执行完成，从而得到更准确的 CUDA kernel 运行时间。

具体的测试代码见代码仓库~

下面展示本人测试的结果。(机器为RTX3090，显存大小为24G)

warmup_steps = 10

Size	d_model	d_ff	num_layers	num_heads	Parameters	Forward (s)	Forward+Backward (s)	Peak Forward Mem (MiB)	Peak Train Mem (MiB)
small	768	3072	12	12	128.625M	0.023376	0.096487	1634.44	3170.83
medium	1024	4096	24	16	423.183M	0.070973	0.284677	4934.53	8745.55
large	1280	5120	36	20	969.412M	0.144661	0.603096	11503.8	18567.7
xl	1600	6400	48	25	1.998B	0.2913
7B	2560	10240	32	32	3.407B	0.450041

Nsight Systems Profile

上面的Benchmark测试不能反映在前向传播和反向传播过程中，模型的时间和显存究竟用在了哪些地方，因此我们也难以找到性能瓶颈从而找到具体的优化方法。

要了解程序在每个组件上花了多少时间，我们可以用Profiler。Profiler会在函数开始和结束时插入检测点，因此能够提供函数级别的详细执行统计信息，例如调用次数，平均耗时，在该函数上累计花费的时间等。

NVIDIA 提供了一个 profiler，我们可以通过命令行工具 nsys 来使用它。我们在本节将使用nsys 来分析 Transformer 模型的运行时间。

nsys的使用方法很直接：只要在上一小节的Python 脚本前加上nsys profile 即可。

例如，你可以对脚本 benchmark.py 进行分析，并将输出写入文件 result.nsys.rep：

1

~$ uv run nsys profile -o result python benchmark.py

随后，你可以在本地机器上使用 NVIDIA Nsight Systems 桌面应用查看这个 profile。

如果使用的是服务器环境，可以使用如下指令查看结果:

nsys stats --report nvtx_sum fwd_step_profile.nsys-rep

在 profile 的 CUDA API 行中选择某个特定的 CUDA API 调用（CPU 侧），会高亮显示 CUDA HW 行中所有对应的 kernel 执行（GPU 侧）。

此外，我们可以用nvtx range对代码进行标注，这些标注会以块的形式显示在 profile 的 NVTX 行中，并涵盖其中所有的 CUDA API 调用及其对应的 kernel 执行。

具体而言，假设我们想针对 forward阶段进行profiling，那么我们可以:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


from contextlib import contextmanager
import torch.cuda.nvtx as nvtx

@contextmanager
def nvtx_range(name):
 nvtx.range_push(name)
 try:
 yield
 finally:
 nvtx.range_pop()
def benchmark_forward(model, x, warmup_steps, steps):
 model.eval()
 for _ in range(warmup_steps):
 with torch.no_grad():
 _ = model(x)
 torch.cuda.synchronize()

 start_time = timeit.default_timer()
 for _ in range(steps):
 with torch.no_grad():
 with nvtx_range("FWD_STEP"):
 _ = model(x)
 torch.cuda.synchronize()
 return (timeit.default_timer() - start_time) / steps

然后使用指令

1
2
3
4


nsys profile -t cuda,nvtx,cublas,cudnn,osrt \
 -o fwd_step_profile --force-overwrite=true \
 uv run python Train/end2end_bench.py --config Train/config.yaml
nsys stats --report nvtx_sum fwd_step_profile.nsys-rep

然后我们可以得到如下结果（SMALL模型）：

Time (%)	Total Time (s)	Instances	Avg (s)	Med (s)	Min (s)	Max (s)	StdDev (s)	Style	Range
46.4	6.0856754600	100	0.0608567546	0.0549357350	0.0489394220	0.0925647510	0.0101878140	PushPop	`:BWD_STEP`
22.8	2.9956320640	100	0.0299563206	0.0314317430	0.0211455360	0.0349865510	0.0035540577	PushPop	`:FWD_STEP`
21.8	2.8684238990	100	0.0286842390	0.0270241125	0.0252365100	0.0404667660	0.0039458590	PushPop	`:OPTIM_STEP`
6.7	0.8811120340	48396	0.0000182063	0.0000168720	0.0000067570	0.0023475050	0.0000141058	PushPop	`cuBLAS:cublasLtSSSMatmul`
2.1	0.2773034430	48396	0.0000057299	0.0000042410	0.0000020940	0.0593009870	0.0002695509	PushPop	`cuBLAS:cublasLtSSSMatmulAlgoGetHeuristic`
0.2	0.0197736070	2	0.0098868035	0.0098868035	0.0003988660	0.0193747410	0.0134179699	PushPop	`cuBLAS:cublasCreate_v2`

通过这个指令我们就可以知道FWD，BWD，OPT各个阶段花费的时间，可以看到与用Python Cli测试出来的时间存在一些微小的差异。

如果我们要看各个阶段占用GPU时间最长的CUDA Kernel的信息，如调用次数，我们可以采用这个指令:

1

nsys stats --report nvtx_kern_sum fwd_step_profile.nsys-rep | grep ':FWD_STEP'

在我们的实验中，ampere_sgemm_128x64_tn这个Kernel占用时间最长Total Time = 1,282,045,469 ns（约 1.282 s），调用次数为8500，这是一个矩阵乘法Kernel(GEMM)

假设我们想知道scaled_dot_product_attention中内部Softmax操作和矩阵乘法操作的运行时间，它们之间的运行时间差异，与Flops的差异，我们同样可以使用nvtx，具体而言，我们可以实现一个注释版的scale_dot_product_attention，然后做一个替换就行:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def annotated_scaled_dot_product_attention(Q, K, V, mask=None):
 d_k = Q.shape[-1]
 with nvtx_range("ATTN_QK_MATMUL"):
 scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
 if mask is not None:
 scores = scores.masked_fill(~mask, float("-inf"))
 with nvtx_range("ATTN_SOFTMAX"):
 attn_weights = torch.softmax(scores, dim=-1)
 with nvtx_range("ATTN_AV_MATMUL"):
 return torch.matmul(attn_weights, V)

def install_annotated_attention():
 attention_module.run_scaled_dot_product_attention = annotated_scaled_dot_product_attention

然后像之前一样进行profiling就行，最后的结果为:

Range	Total Time (s)	Instances	Avg (s)	Med (s)	Min (s)	Max (s)	StdDev (s)	Time (%)
`ATTN_QK_MATMUL`	0.380551201	2664	0.0001428495	0.0001381315	0.000108449	0.000943515	0.0000377748	3.0
`ATTN_AV_MATMUL`	0.202245954	2664	0.0000759182	0.0000759685	0.000057080	0.000795635	0.0000197206	1.6
`ATTN_SOFTMAX`	0.083674056	2664	0.0000314092	0.0000260260	0.000019989	0.012475856	0.0002412466	0.7

混合精度

目前为止，我们一直都在使用FP32精度。然而现代的NVIDIA GPU 包含专门的 GPU 核心（Tensor Cores），用于在更低精度下加速矩阵乘法。例如，NVIDIA A100 的规格说明显示，它在 FP32 下的最大吞吐量是 19.5 TFLOP/s，而在 FP16（半精度浮点） 或 BF16（brain floating point） 下的最大吞吐量则高得多，可达 312 TFLOP/s。因此，使用更低精度的数据类型应当有助于加速训练和推理。

不过，如果只是简单地把模型直接转换成更低精度格式，可能会带来模型精度下降的问题。

例如，实际中的许多梯度值往往太小，无法用 FP16 表示，因此在直接用 FP16 训练时会变成 0。

为了解决这个问题，在使用 FP16 训练时，通常会采用 loss scaling（损失缩放）：即把 loss 乘上一个缩放因子，从而增大梯度幅值，避免它们下溢为 0。

此外，FP16 的动态范围比 FP32 更小，也可能导致数值溢出，表现为 loss 变成 NaN。

完整使用 BF16 训练通常会更稳定，因为 BF16 与 FP32 具有相同的动态范围；不过，与 FP32 相比，它仍然可能影响模型的最终性能。

为了同时利用低精度数据类型带来的速度提升，又尽量避免数值问题，通常会采用 混合精度训练（mixed-precision training）。在 PyTorch 中，这通过 torch.autocast 上下文管理器实现。

在这种模式下，某些操作（例如矩阵乘法）会使用较低精度执行，而另一些需要 FP32 完整动态范围的操作（例如累加和归约）则保持原样。

例如，下面的代码会在前向传播过程中自动识别哪些操作适合使用低精度，并将这些操作转换到指定的数据类型：

1
2
3
4
5
6


model: torch.nn.Module = ... # 例如你的 Transformer 模型
dtype: torch.dtype = ... # 例如 torch.float16
x: torch.Tensor = ... # 输入数据

with torch.autocast(device="cuda", dtype=dtype):
 y = model(x)

根据实验指导书，我们对一个ToyModel使用autocast来看看各个输出的类型，具体代码如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


class ToyModel(nn.Module):
 def __init__(self,in_features,out_features):
 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))
 print(f"fc1 output dtype: {x.dtype}")
 x = self.ln(x)
 print(f"LayerNorm output dtype: {x.dtype}")
 x = self.fc2(x)
 return x
def example2():
 model = ToyModel(20, 5).cuda()
 x = torch.randn(4, 20).cuda()
 y = torch.randint(0, 5, (4,)).cuda()
 with torch.autocast(device_type='cuda', dtype=torch.float16):
 logits = model(x)
 print(f"Final output dtype: {logits.dtype}")
 loss = nn.CrossEntropyLoss()(logits, y)
 print(f"Loss dtype: {loss.dtype}")
 loss.backward()
 print(f"Gradients dtype: {model.fc1.weight.grad.dtype}")

最后的输出为:

1
2
3
4
5


fc1 output dtype: torch.float16
LayerNorm output dtype: torch.float32
Final output dtype: torch.float16
Loss dtype: torch.float32
Gradients dtype: torch.float32

Memory Profiling

我们现在来看看内存，pytorch自带了一个功能强大的内存分析器，它可以持续追踪一段时间内的内存分配情况。

使用方法也很简单：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


# ... 你的 benchmarking script 中的 warm-up 阶段

# 开始记录内存历史
torch.cuda.memory._record_memory_history(max_entries=1000000)

# ... 你的 benchmarking script 中想要分析的部分

# 保存一个 pickle 文件，以便加载到 PyTorch 的在线工具中
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

# 停止记录历史
torch.cuda.memory._record_memory_history(enabled=None)

这会输出一个名为 memory_snapshot.pickle 的文件，你可以把它加载到下面这个在线工具中： https://pytorch.org/memory_viz

这个工具可以让你查看整体内存使用时间线，以及每一次单独的内存分配，包括它的大小和一条栈追踪（stack trace），从而定位这块内存分配源自哪段代码。要使用这个工具，你需要在浏览器中打开上面的链接，然后把你的 Pickle 文件拖放到页面中。

我们接下来看看在large模型中，在上下文为256时，forward pass，backward pass和optimizer step的内存使用情况。

从图中我们可以明显看到显存峰值，在FWD阶段大约是3.8GiB，而在TRAIN的话则会来到16GB

接下来我们看看在MIX_precision(BF16)的情况下的显存变化

可以发现一个奇怪的点在于开启了混合精度后，显存峰值不但没减小，反而更大了，这令人有点匪夷所思。

但是这其实是因为AMP 只会减少部分激活显存；参数和 AdamW 状态仍主要是 FP32，所以如果 batch 不大，节省不明显甚至被额外开销盖过。（在1.3中我们有输出过数据的type）

而autocast缓存会带来额外的开销，从而导致节省的内存不如额外的开销，从而带来混合精度占用更多内存的感觉。

Attention的优化– FlashAttention2

Attention Benchmark

在进入具体的优化之前我们做了一个基本的测试，对单头的 causal self-attention 做端到端基准测试，固定 batch size 为 8，然后遍历 d_model 和 seq_len 的所有组合。对每个配置，它会先做 warmup，再执行 100 次 forward 并统计平均耗时，然后执行 100 次 backward 并统计平均耗时，同时在 backward 开始前采样显存，并分别记录 forward 和 backward 阶段的峰值显存。如果某个配置触发 out-of-memory，会把该配置标记为 OOM 后继续跑后续配置，最后把所有配置的时间和显存结果统一汇总成一个 Markdown 表格输出。

d_model	seq_len	fwd_ms	bwd_ms	fwd_peak_mem_mib	bwd_mem_before_backward_mib	bwd_peak_mem_mib	status
16	256	0.98	6.71	15.17	21.28	29.23	ok
16	1024	1.52	9.04	117.92	85.42	212.92	ok
16	4096	12.55	37.81	1598.91	1072.91	3118.91	ok
16	8192	46.48	143.74	6317.57	4209.57	12397.6	ok
16	16384						oom
32	256	1.25	7.3	24.2	22.39	30.14	ok
32	1024	1.63	9.51	121.49	89.49	216.5	ok
32	4096	12.58	38.35	1613.17	1089.17	3133.18	ok
32	8192	47.4	146.02	6346.08	4242.08	12426.1	ok
32	16384						oom
64	256	1.01	5.89	26.03	24.46	31.98	ok
64	1024	1.62	9.51	128.66	97.66	223.68	ok
64	4096	12.98	39.55	1641.72	1121.72	3161.74	ok
64	8192	48.63	149.26	6403.12	4307.13	12483.1	ok
64	16384						oom
128	256	0.98	9.25	29.78	28.71	35.78	ok
128	1024	1.55	9.29	143.1	114.1	238.17	ok
128	4096	14.15	42.6	1698.91	1186.91	3218.97	ok
128	8192	53	159.16	6517.31	4437.31	12597.4	ok
128	16384						oom

除此以外，自 PyTorch 2.0 起，PyTorch 还内置了一个强大的 即时编译器（just-in-time compiler），它会自动尝试对 PyTorch 函数应用多种优化。

特别地，它会通过动态分析你的计算图，自动尝试生成融合后的 Triton kernel。使用 PyTorch 编译器的接口非常简单。例如，如果我们想把它应用到模型中的某一层，可以这样写：

1
2


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

现在，compiled_layer 在功能上与 layer 完全一致（例如，它的 forward 和 backward 行为相同）。

我们也可以用 torch.compile(model) 来编译整个 PyTorch 模型，甚至也可以编译一个调用了 PyTorch 操作的 Python 函数。

我们用上面介绍的即时编译器对我们的attention重新做了一次Benchmark，实验配置与之前一致，唯一不同在于使用的Attention是compile过后的。结果如下：

d_model	seq_len	fwd_ms	bwd_ms	fwd_peak_mem_mib	bwd_mem_before_backward_mib	bwd_peak_mem_mib	status
16	256	1.28	7.34	15.17	21.28	29.23	ok
16	1024	1.64	9.23	117.92	85.42	212.92	ok
16	4096	12.54	37.93	1598.91	1072.91	3118.91	ok
16	8192	46.49	143.7	6317.57	4209.57	12397.6	ok
16	16384						oom
32	256	1.26	8.7	24.2	22.39	30.14	ok
32	1024	1.62	8.86	121.49	89.49	216.5	ok
32	4096	12.7	38.46	1613.17	1089.17	3133.18	ok
32	8192	47.29	146.08	6346.08	4242.08	12426.1	ok
32	16384						oom
64	256	1.05	6.77	26.03	24.46	31.98	ok
64	1024	1.65	8.73	128.66	97.66	223.68	ok
64	4096	12.93	39.44	1641.72	1121.72	3161.74	ok
64	8192	48.62	149.38	6403.12	4307.13	12483.1	ok
64	16384						oom
128	256	1.01	8.77	29.78	28.71	35.78	ok
128	1024	1.54	8.77	143.1	114.1	238.17	ok
128	4096	14.09	42.53	1698.91	1186.91	3218.97	ok
128	8192	52.98	159.19	6517.31	4437.31	12597.4	ok
128	16384						oom

可以看到，虽然在较小的模型上，性能有所衰减，但是随着模型增大，带来的性能提升便逐渐开始显现。

我们在我们的End2End的Benchmark上进行一下测试，我们直接将整个Transformer模型进行编译，得到下表:

Forward (s)	Forward+Backward (s)
0.140853	0.578436

Forward (s)	Forward+Backward (s)
0.12706	0.578973

可以看到优势在Forward上还是很显著的。

Triton编程

这一小节实际上是Assignment 2的核心，也是是否能够理解FlashAttention的关键。因此我将这一小节的内容用markdown的形式进行了留档，更新至博客，见下面这篇：

Triton

FlashAttention V2

High level

在学习完Triton编程的基础，并解决了Triton Puzzles中的所有问题后，我们已经对Triton编程，特别是Block Pointer的用法有了一个清晰的认知，并对Online Softmax有了一个初步的了解，在此基础上，我们接下来将实现一个完整的FlashAttention的Forward和Backward流程，并将我们在Assignment1中的Attention用这个Triton算子进行替换。

我们先简单回顾一下Attention操作，并理解其低效之处。

Attention的前向传播过程可以写作:

$$ S = \frac{QK^\top}{\sqrt{d}}\\P_{ij} = \text{softmax}_j(S)_{ij}\\O=PV $$

其标准的反向传播过程可以写作:

$$ dV = P^\top dO\quad dP = dO V^\top\\dS_i = d\text{softmax}(dP_i)=(\text{diag}(P_i) - P_iP_i^\top)dP_i\\dQ = \frac{dS K}{\sqrt{d}} \quad dK = \frac{dS^\top Q}{\sqrt{d}} $$

正如我们可以看到的，反向传播依赖于前向传播中的一些非常大的激活值矩阵。例如上式中计算dV需要用到P,而P的形状为(batch_size,n_heads,seq_len,seq_len)的Attention Score。这个激活值矩阵的大小随序列长度呈现二次增长，这也解释了我们之前在对长序列进行Attention Benchmark时需要的OOM问题。

在普通 attention 的前向和反向传播中，我们都要付出显著的内存 I/O 成本，用于在片上 SRAM 和 GPU HBM 之间传输 P 以及其他大型激活值。标准实现中会进行多次这类传输。

那么我们实现的FlashAttention的主要目标就是避免将Attention矩阵写入和读出HBM，从而减少I/O和显存峰值开销。

我们将通过三种技术来实现这一点:

分块:为了避免将attention矩阵写入和读出HBM，我们需要在无法访问完整输入的情况下完成softmax归约。具体来说，我们会重构Attention的计算方式，把输入拆分为多个tile，并对这些块进行多次访问遍历，从而以增量的形式执行Softmax的规约
重计算：我们避免在HBM中存储(batch_size,n_heads,seq_len,seq_len)的大型中间注意力矩阵。取而代之的是，我们会在HBM中保存某些“激活检查点”，然后在反向传播时重新计算前向传播的一部分过程，以恢复计算梯度所需的其他激活值。FlashAttention-2还会存储注意力分数的logsumexp记作L,它将用于简化反向传播计算。L的表达式为:$L_i = \log(\sum_j \exp(S_{ij}))$.在最终的Kernel中，我们将以Online的方式计算它，但最终的结果应当保持一致。通过结合分块和重计算，我们的内存IO和峰值内存将不再依赖于Seq_len^2,因此可以支持更长的序列长度。
算子融合：最后我们通过在单个Kernel中完成操作，避免对注意力矩阵以及其他中间激活进行重复的内存IO。我们将编写一个单独的 Triton kernel 来执行前向传播，在注意力机制涉及的所有操作中，尽量减少 HBM 与 SRAM 之间的数据传输。算子融合部分得益于重计算，因为这样我们可以避免将每个中间激活都写入 HBM 所产生的常规内存 IO 开销。

下面简单介绍一下带重计算的BackWard Pass。

借助L，我们可以进行适当的重计算，并高效地完成反向传播。在开始反向传播之前，我们会先在全局内存中预计算数值

$$ D=\text{rowsum}(O\odot dO) $$

其中 $\odot$为逐元素乘法。由于: $\text{rowsum}(P\odot dP)=D$,这是因为

$$ PdP^\top = P(dOV^\top)^\top = (PV)dO^\top =OdO^\top $$

且，对于任意矩阵有 $\text{rowsum}(A\odot B) = \text{diag}(AB^\top)$

有了向量L和D后，反向传播的过程可以在不显式执行softmax的情况下完成。此时，完整的反向传播计算过程如下。

$$ S = \frac{QK^\top}{\sqrt{d}},P_{ij}=\exp(S_{ij}-L_i)\\dV = P^\top dO,dP=dOV^\top \\dS_{ij}=P_{ij}\odot (dP_{ij}-D_i)\\dQ=dSK/\sqrt{d},dK = dS^\top Q/\sqrt{d} $$

可以看到，这一系列操作并不要求我们在前向传播期间将注意力分数 P 存储在 HBM 中.

Forward Pass

现在我们已经对FlashAttention v2有了一个High Level的认知，接下来我们来实现具体的Kernel

为了避免将注意力矩阵在HBM中来回读写，我们希望采用tile，也就是让每个tile都能够独立于其他tile进行运算。这要求我们能够计算P的各个tile，并且最好能够在两个维度上都进行分块（query和key）。

然而，当我们对S应用softmax时，需要对S的整行进行规约，以计算Softmax的分母。这意味着我们不能直接按tile独立计算P。我们可以通过Online Softmax来解决这个问题。

在下面的描述中，我们用下标i来表示当前的query tile，用上标j来表示当前的key tile。沿query维度的tile大小为B_q,沿key维度的tile大小为B_k。我们不会沿隐藏维度d进行分块。

我们还会维护一些中间值: $m_i^{(j)}\in \mathbb{R}^{B_q},l_i^{(j)}\in \mathbb{R}^{B_q}$

其中前者表示运行中的最大值，我们跟踪它是为了能够以数值稳定的方式计算 softmax ；每当我们处理一个新的S的按行tile（即j增加时），我们都会更新它。

借助这个最大值，我们可以计算未归一化的softmax值（即分子）： $P_i^{(j)} = \exp(S_{ij}-m_i^{(j)})\cdot l_i^{(j)}$

而后者则是softmax分母的一个运行代理值，它会利用这些未归一化的softmax值进行更新。最终当我们写出输出结果时，还需要用它来进行归一化。

我们首先实现一版用torch来模拟的FlashAttention v2

代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


def forward(ctx,Q,K,V,is_causal=False):
 batch_size,seq_len,d_model = Q.shape
 device = Q.device
 dtype = Q.dtype
 # S = torch.zeros(batch_size,TILE_SIZE,TILE_SIZE)
 sqrt_d = math.sqrt(d_model)
 O = torch.zeros(batch_size,seq_len,d_model,dtype=dtype,device=device)
 L = torch.zeros(batch_size,seq_len,dtype=dtype,device=device)
 for i in range(0,seq_len // TILE_SIZE):
 m = torch.full((batch_size,TILE_SIZE),float("-inf"),device=device)
 l = torch.zeros(batch_size,TILE_SIZE,dtype=dtype,device=device)
 B_q = Q[:,i*TILE_SIZE:(i + 1) * TILE_SIZE,:]
 for j in range(0,seq_len // TILE_SIZE):
 B_k = K[:,j * TILE_SIZE:(j + 1) * TILE_SIZE,:]
 S = einsum(B_q,B_k,"... B_q d,... B_k d -> ... B_q B_k") / sqrt_d
 last_m,last_l = m.clone(),l.clone()
 m = torch.maximum(m,torch.amax(S,dim=2))
 l = last_l * torch.exp(last_m - m) + torch.sum(torch.exp(S - m[:,:,None]),dim=2) 
 S = torch.exp(S - m[:,:,None])
 O[:,i*TILE_SIZE:(i+1)*TILE_SIZE,:] = O[:,i*TILE_SIZE:(i+1)*TILE_SIZE,:] * torch.exp(last_m - m).unsqueeze(-1) + einsum(S,V[:,j*TILE_SIZE:(j+1)*TILE_SIZE,:],"... B_q B_k,... B_k d_model->... B_q d_model")
 O[:,i*TILE_SIZE:(i+1)*TILE_SIZE,:] = O[:,i*TILE_SIZE:(i+1)*TILE_SIZE,:] / l.unsqueeze(-1)
 L[:,i*TILE_SIZE:(i+1)*TILE_SIZE] = torch.log(l) + m
 ctx.save_for_backward(L,Q,K,V,O)
 return O

在实现的时候，可能会因为i和j晕头转向，因此推荐用纸和笔画一下Q，K，V，S，O之间的对应关系。

实现了Torch版本的forward pass之后实现Triton版本的就不是很困难了，我们只需要将在torch中对i维度想象成并行的就行。然后框架上，我们参考一下Weighted Sum的实现即可。

完整的Triton代码等到最后FWD+BWD一起实现了再进行展示。

Backward Pass

我们接下来实现一下Backward Pass。我们回顾一下之前在High Level中介绍的Backward的计算流程（带重计算）

$$ S = \frac{QK^\top}{\sqrt{d}},P_{ij}=\exp(S_{ij}-L_i)\\dV = P^\top dO,dP=dOV^\top \\dS_{ij}=P_{ij}\odot (dP_{ij}-D_i)\\dQ=dSK/\sqrt{d},dK = dS^\top Q/\sqrt{d} $$

这部分我们同样通过分块进行，从而避免直接计算出大小为Batch_size x N_q x Nk的注意力分数。

分块的大小和Forward Pass统一。为了方便理解，我画了一些图。（这也是我写Triton算子时的一个小习惯，可以让思路更加清晰）

相信你从上面这个图能对这个分块运算有了更加直观的认识，对于dP，dS，dQ，dK的计算和dV的分块逻辑类似，不再过多赘述，还是一样，建议在写代码之前在纸上画一下上面这样的块图。

代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


def backward(ctx,grad_output):
 L,Q,K,V,O = ctx.saved_tensors
 batch_size,N_QUERIES,d_model = Q.shape
 _,N_KEYS,_ = K.shape
 scale = 1.0 / (d_model ** 0.5)
 dQ = torch.zeros_like(Q)
 dK = torch.zeros_like(K)
 dV = torch.zeros_like(V)
 D = torch.sum(grad_output * O,dim=2)
 for i in range(0,N_QUERIES // TILE_SIZE):
 Q_b = Q[:,i*TILE_SIZE:(i+1)*TILE_SIZE,:]
 L_b =L[:,i*TILE_SIZE:(i+1)*TILE_SIZE]
 dO_b = grad_output[:,i*TILE_SIZE:(i+1)*TILE_SIZE,:]
 D_b = D[:,i*TILE_SIZE:(i+1)*TILE_SIZE]
 for j in range(0,N_KEYS // TILE_SIZE):
 K_b = K[:,j * TILE_SIZE:(j + 1) * TILE_SIZE,:]
 V_b = V[:,j*TILE_SIZE:(j+1)*TILE_SIZE,:]
 P_ij = torch.exp(torch.matmul(Q_b,K_b.transpose(-1,-2)) * scale - L_b[:,:,None])
 dV[:,j*TILE_SIZE:(j+1)*TILE_SIZE,:] += torch.matmul(P_ij.transpose(-1,-2),dO_b)
 dP_ij = torch.matmul(dO_b,V_b.transpose(-1,-2))
 dS_ij = P_ij * (dP_ij - D_b[:,:,None])
 dQ[:,i*TILE_SIZE:(i+1)*TILE_SIZE,:] += torch.matmul(dS_ij,K_b) * scale
 dK[:,j*TILE_SIZE:(j+1)*TILE_SIZE,:] += torch.matmul(dS_ij.transpose(-1,-2),Q_b) * scale
 return dQ,dK,dV,None

我们在用Triton实现的时候我们需要注意一个事情，就是dQ和dK，dV的计算要分为两个Phase，这是因为累计的方向不一致。

在计算dQ时，我们需要固定Q_block,遍历所有的K/V block，然后累计dQ

在计算dK，dV时，我们则需要固定K_Block,V_block,遍历所有的Q_Block,然后累计dK和dV

意思是我们需要重计算两遍P，从而提高并行度并减少不必要的通信和同步开销。

接下来给出完整的Triton实现的FlashAttention-2

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376


import triton
import triton.language as tl
import math

import torch
Q_TILE_SIZE = 16
K_TILE_SIZE = 16
@triton.jit
def flash_fwd_kernel(
 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,
 is_causal: tl.constexpr,
):
 query_tile_index = tl.program_id(0)
 batch_index = tl.program_id(1)
 Q_block_ptr = tl.make_block_ptr(
 base = 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),
 )
 K_block_ptr = tl.make_block_ptr(
 base = K_ptr + batch_index * stride_kb,
 shape = (N_KEYS,D),
 strides = (stride_kk,stride_kd),
 offsets = (0,0),
 block_shape = (K_TILE_SIZE,D),
 order = (1,0),
 )
 V_block_ptr = tl.make_block_ptr(
 base = V_ptr + batch_index * stride_vb,
 shape = (N_KEYS,D),
 strides = (stride_vk,stride_vd),
 offsets = (0,0),
 block_shape = (K_TILE_SIZE,D),
 order =(1,0),
 )
 L_block_ptr = tl.make_block_ptr(
 base = L_ptr + batch_index * stride_lb,
 shape = (N_QUERIES,),
 strides = (stride_lq,),
 offsets = (query_tile_index * Q_TILE_SIZE,),
 block_shape = (Q_TILE_SIZE,),
 order = (0,),
 )
 O_block_ptr = tl.make_block_ptr(
 base = O_ptr + batch_index * stride_ob,
 shape = (N_QUERIES,D),
 strides = (stride_oq,stride_od),
 offsets = (query_tile_index * Q_TILE_SIZE,0),
 block_shape = (Q_TILE_SIZE,D),
 order = (1,0),
 )
 
 m = tl.full((Q_TILE_SIZE,),float("-inf"),dtype=tl.float32)
 l = tl.zeros((Q_TILE_SIZE,),dtype = tl.float32)
 O = tl.zeros((Q_TILE_SIZE,D),dtype=tl.float32)
 B_q = tl.load(Q_block_ptr,boundary_check=(0,1),padding_option="zero")
 for j in range(tl.cdiv(N_KEYS,K_TILE_SIZE)):
 B_k = tl.load(K_block_ptr,boundary_check=(0,1),padding_option="zero")
 B_v = tl.load(V_block_ptr,boundary_check=(0,1),padding_option="zero")
 S = tl.dot(B_q,tl.trans(B_k)) * scale
 if is_causal:
 q_idx = query_tile_index * Q_TILE_SIZE + tl.arange(0,Q_TILE_SIZE)
 k_idx = j * K_TILE_SIZE + tl.arange(0,K_TILE_SIZE)
 causal_mask = q_idx[:,None] < k_idx[None,:]
 S = tl.where(causal_mask,float("-inf"),S)
 last_m,last_l = m,l
 m = tl.maximum(last_m,tl.max(S,axis=1))
 S_ = tl.exp(S - m[:,None])
 l = last_l * tl.exp(last_m - m) + tl.sum(S_,axis=1)
 O = tl.dot(S_.to(B_v.dtype),B_v,acc=O * tl.exp(last_m - m)[:,None])
 K_block_ptr = tl.advance(K_block_ptr,(K_TILE_SIZE,0))
 V_block_ptr = tl.advance(V_block_ptr,(K_TILE_SIZE,0))
 O = O / l[:,None]
 L = tl.log(l) + m
 tl.store(O_block_ptr, O.to(O_block_ptr.type.element_ty), boundary_check=(0, 1))
 tl.store(L_block_ptr,L,boundary_check=(0,))

@triton.jit
def flash_bwd_kernel_phase1(
 Q_ptr,K_ptr,V_ptr,L_ptr,O_ptr,
 dO_ptr,dQ_ptr,
 stride_qb,stride_qq,stride_qd,
 stride_kb,stride_kk,stride_kd,
 stride_vb,stride_vk,stride_vd,
 stride_lb,stride_lq,
 stride_ob,stride_oq,stride_od,
 stride_dOb,stride_dOq,stride_dOd,
 N_QUERIES,N_KEYS,
 scale,
 D:tl.constexpr,
 Q_TILE_SIZE: tl.constexpr,
 K_TILE_SIZE: tl.constexpr,
 is_causal: tl.constexpr,
):
 query_tile_index = tl.program_id(0)
 batch_index = tl.program_id(1)
 Q_block_ptr = tl.make_block_ptr(
 base = 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),
 )
 K_block_ptr = tl.make_block_ptr(
 base = K_ptr + batch_index * stride_kb,
 shape = (N_KEYS,D),
 strides = (stride_kk,stride_kd),
 offsets = (0,0),
 block_shape = (K_TILE_SIZE,D),
 order = (1,0),
 )
 V_block_ptr = tl.make_block_ptr(
 base = V_ptr + batch_index * stride_vb,
 shape = (N_KEYS,D),
 strides = (stride_vk,stride_vd),
 offsets = (0,0),
 block_shape = (K_TILE_SIZE,D),
 order =(1,0),
 )
 L_block_ptr = tl.make_block_ptr(
 base = L_ptr + batch_index * stride_lb,
 shape = (N_QUERIES,),
 strides = (stride_lq,),
 offsets = (query_tile_index * Q_TILE_SIZE,),
 block_shape = (Q_TILE_SIZE,),
 order = (0,),
 )
 O_block_ptr = tl.make_block_ptr(
 base = O_ptr + batch_index * stride_ob,
 shape = (N_QUERIES,D),
 strides = (stride_oq,stride_od),
 offsets = (query_tile_index * Q_TILE_SIZE,0),
 block_shape = (Q_TILE_SIZE,D),
 order = (1,0),
 )
 dO_block_ptr = tl.make_block_ptr(
 base = dO_ptr + batch_index * stride_dOb,
 shape = (N_QUERIES,D),
 strides = (stride_dOq,stride_dOd),
 offsets = (query_tile_index * Q_TILE_SIZE,0),
 block_shape = (Q_TILE_SIZE,D),
 order = (1,0),
 )
 dQ_block_ptr = tl.make_block_ptr(
 base = dQ_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)
 )
 B_dO = tl.load(dO_block_ptr,boundary_check=(0,1),padding_option="zero")
 O = tl.load(O_block_ptr,boundary_check=(0,1),padding_option="zero")
 D0 = tl.sum(B_dO * O,axis=1)
 B_q = tl.load(Q_block_ptr,boundary_check=(0,1),padding_option="zero")
 B_l = tl.load(L_block_ptr,boundary_check=(0,),padding_option="zero")
 dQ = tl.zeros((Q_TILE_SIZE,D),dtype=tl.float32)
 for j in range(tl.cdiv(N_KEYS,K_TILE_SIZE)):
 B_k = tl.load(K_block_ptr,boundary_check=(0,1),padding_option="zero")
 B_v = tl.load(V_block_ptr,boundary_check=(0,1),padding_option="zero")
 S_ij = tl.dot(B_q,B_k.T) * scale
 if is_causal:
 q_idx = query_tile_index * Q_TILE_SIZE + tl.arange(0,Q_TILE_SIZE)
 k_idx = j * K_TILE_SIZE + tl.arange(0,K_TILE_SIZE)
 causal_mask = q_idx[:,None] < k_idx[None,:]
 S_ij = tl.where(causal_mask,float("-inf"),S_ij)
 P_ij = tl.exp(S_ij - B_l[:,None])
 dP_ij = tl.dot(B_dO,B_v.T)
 dS_ij = P_ij * (dP_ij - D0[:,None])
 dQ += tl.dot(dS_ij,B_k) * scale
 K_block_ptr = tl.advance(K_block_ptr,(K_TILE_SIZE,0))
 V_block_ptr = tl.advance(V_block_ptr,(K_TILE_SIZE,0))
 tl.store(dQ_block_ptr,dQ.to(dQ_block_ptr.type.element_ty),boundary_check=(0,1))

@triton.jit
def flash_bwd_kernel_phase2(
 Q_ptr,K_ptr,V_ptr,L_ptr,O_ptr,
 dO_ptr,dK_ptr,dV_ptr,
 stride_qb,stride_qq,stride_qd,
 stride_kb,stride_kk,stride_kd,
 stride_vb,stride_vk,stride_vd,
 stride_lb,stride_lq,
 stride_ob,stride_oq,stride_od,
 stride_dOb,stride_dOq,stride_dOd,
 N_QUERIES,N_KEYS,
 scale,
 D:tl.constexpr,
 Q_TILE_SIZE:tl.constexpr,
 K_TILE_SIZE:tl.constexpr,
 is_causal:tl.constexpr,
):
 key_tile_index = tl.program_id(0)
 batch_index = tl.program_id(1)
 Q_block_ptr = tl.make_block_ptr(
 base = Q_ptr + batch_index * stride_qb,
 shape = (N_QUERIES,D),
 strides = (stride_qq,stride_qd),
 offsets = (0,0),
 block_shape = (Q_TILE_SIZE,D),
 order = (1,0),
 )
 K_block_ptr = tl.make_block_ptr(
 base = K_ptr + batch_index * stride_kb,
 shape = (N_KEYS,D),
 strides = (stride_kk,stride_kd),
 offsets = (key_tile_index * K_TILE_SIZE,0),
 block_shape = (K_TILE_SIZE,D),
 order = (1,0),
 )
 V_block_ptr = tl.make_block_ptr(
 base = V_ptr + batch_index * stride_vb,
 shape = (N_KEYS,D),
 strides = (stride_vk,stride_vd),
 offsets = (key_tile_index * K_TILE_SIZE,0),
 block_shape = (K_TILE_SIZE,D),
 order =(1,0),
 )
 L_block_ptr = tl.make_block_ptr(
 base = L_ptr + batch_index * stride_lb,
 shape = (N_QUERIES,),
 strides = (stride_lq,),
 offsets = (0,),
 block_shape = (Q_TILE_SIZE,),
 order = (0,),
 )
 O_block_ptr = tl.make_block_ptr(
 base = O_ptr + batch_index * stride_ob,
 shape = (N_QUERIES,D),
 strides = (stride_oq,stride_od),
 offsets = (0,0),
 block_shape = (Q_TILE_SIZE,D),
 order = (1,0),
 )
 dO_block_ptr = tl.make_block_ptr(
 base = dO_ptr + batch_index * stride_dOb,
 shape = (N_QUERIES,D),
 strides = (stride_dOq,stride_dOd),
 offsets = (0,0),
 block_shape = (Q_TILE_SIZE,D),
 order = (1,0),
 )
 dK_block_ptr = tl.make_block_ptr(
 base = dK_ptr + batch_index * stride_kb,
 shape = (N_KEYS,D),
 strides = (stride_kk,stride_kd),
 offsets = (key_tile_index * K_TILE_SIZE,0),
 block_shape = (K_TILE_SIZE,D),
 order = (1,0)
 )
 dV_block_ptr = tl.make_block_ptr(
 base = dV_ptr + batch_index * stride_vb,
 shape = (N_KEYS,D),
 strides = (stride_vk,stride_vd),
 offsets = (key_tile_index * K_TILE_SIZE,0),
 block_shape = (K_TILE_SIZE,D),
 order = (1,0)
 )
 B_k = tl.load(K_block_ptr,boundary_check=(0,1),padding_option="zero")
 B_v = tl.load(V_block_ptr,boundary_check=(0,1),padding_option="zero")
 dK = tl.zeros((K_TILE_SIZE,D),dtype=tl.float32)
 dV = tl.zeros((K_TILE_SIZE,D),dtype=tl.float32)
 for i in range(tl.cdiv(N_QUERIES,Q_TILE_SIZE)):
 B_dO = tl.load(dO_block_ptr,boundary_check=(0,1),padding_option="zero")
 B_O = tl.load(O_block_ptr,boundary_check=(0,1),padding_option="zero")
 B_q = tl.load(Q_block_ptr,boundary_check=(0,1),padding_option="zero")
 B_l = tl.load(L_block_ptr,boundary_check=(0,),padding_option="zero")
 D0 = tl.sum(B_dO * B_O,axis=1)
 S_ij = tl.dot(B_q,B_k.T) * scale
 if is_causal:
 q_idx = i * Q_TILE_SIZE + tl.arange(0,Q_TILE_SIZE)
 k_idx = key_tile_index * K_TILE_SIZE + tl.arange(0,K_TILE_SIZE)
 causal_mask = q_idx[:,None] < k_idx[None,:]
 S_ij = tl.where(causal_mask,float("-inf"),S_ij)
 P_ij = tl.exp(S_ij - B_l[:,None])
 dV += tl.dot(P_ij.T,B_dO)
 dP_ij = tl.dot(B_dO,B_v.T)
 dS_ij = P_ij * (dP_ij - D0[:,None])
 dK += tl.dot(dS_ij.T,B_q) * scale
 dO_block_ptr = tl.advance(dO_block_ptr,(Q_TILE_SIZE,0))
 O_block_ptr = tl.advance(O_block_ptr,(Q_TILE_SIZE,0))
 Q_block_ptr = tl.advance(Q_block_ptr,(Q_TILE_SIZE,0))
 L_block_ptr = tl.advance(L_block_ptr,(Q_TILE_SIZE,))
 tl.store(dK_block_ptr,dK.to(dK_block_ptr.type.element_ty),boundary_check = (0,1))
 tl.store(dV_block_ptr,dV.to(dV_block_ptr.type.element_ty),boundary_check = (0,1))

class FlashAttention(torch.autograd.Function):
 @staticmethod
 def forward(ctx,Q,K,V,is_causal=False):
 batch_size,N_QUERIES,D = Q.shape
 N_KEYS = K.shape[1]
 ctx.is_causal = is_causal
 ctx.Q_TILE_SIZE = Q_TILE_SIZE
 ctx.K_TILE_SIZE = K_TILE_SIZE
 ctx.D = D

 O = torch.zeros((batch_size,N_QUERIES,D),device=Q.device,dtype=Q.dtype)
 L = torch.zeros((batch_size,N_QUERIES),device=Q.device,dtype=Q.dtype)

 grid = ((N_QUERIES + Q_TILE_SIZE - 1) // Q_TILE_SIZE,batch_size)
 scale = 1.0 / math.sqrt(D)
 flash_fwd_kernel[grid](
 Q,K,V,
 O,L,
 Q.stride(0),Q.stride(1),Q.stride(2),
 K.stride(0),K.stride(1),K.stride(2),
 V.stride(0),V.stride(1),V.stride(2),
 O.stride(0),O.stride(1),O.stride(2),
 L.stride(0),L.stride(1),
 N_QUERIES,N_KEYS,
 scale,
 D,
 Q_TILE_SIZE,K_TILE_SIZE,
 is_causal,
 )
 ctx.save_for_backward(L,K,Q,V,O)
 return O
 def backward(ctx,grad_output):
 L,K,Q,V,O = ctx.saved_tensors
 is_causal = ctx.is_causal
 Q_TILE_SIZE = ctx.Q_TILE_SIZE
 K_TILE_SIZE = ctx.K_TILE_SIZE
 D = ctx.D
 batch_size,N_QUERIES,_ = Q.shape
 _,N_KEYS,_ = K.shape
 dQ = torch.zeros((batch_size,N_QUERIES,D),device=Q.device,dtype=Q.dtype)
 dK = torch.zeros((batch_size,N_KEYS,D),device=K.device,dtype=Q.dtype)
 dV = torch.zeros((batch_size,N_KEYS,D),device=V.device,dtype=V.dtype)
 grid = ((N_QUERIES + Q_TILE_SIZE - 1) // Q_TILE_SIZE,batch_size)
 scale = 1.0 / (D ** 0.5)
 flash_bwd_kernel_phase1[grid](
 Q,K,V,L,O,
 grad_output,dQ,
 Q.stride(0),Q.stride(1),Q.stride(2),
 K.stride(0),K.stride(1),K.stride(2),
 V.stride(0),V.stride(1),V.stride(2),
 L.stride(0),L.stride(1),
 O.stride(0),O.stride(1),O.stride(2),
 grad_output.stride(0),grad_output.stride(1),grad_output.stride(2),
 N_QUERIES,N_KEYS,
 scale,
 D,
 Q_TILE_SIZE,K_TILE_SIZE,
 is_causal
 ) 
 grid = ((N_KEYS + K_TILE_SIZE - 1) // K_TILE_SIZE,batch_size)
 flash_bwd_kernel_phase2[grid](
 Q,K,V,L,O,
 grad_output,dK,dV,
 Q.stride(0),Q.stride(1),Q.stride(2),
 K.stride(0),K.stride(1),K.stride(2),
 V.stride(0),V.stride(1),V.stride(2),
 L.stride(0),L.stride(1),
 O.stride(0),O.stride(1),O.stride(2),
 grad_output.stride(0),grad_output.stride(1),grad_output.stride(2),
 N_QUERIES,N_KEYS,
 scale,
 D,
 Q_TILE_SIZE,K_TILE_SIZE,
 is_causal
 ) 
 return dQ,dK,dV,None

对比测试&部分优化

这部分先暂时跳过了qwq

分布式数据并行训练

我们接下来探索的内容是如何在多个GPU上训练一个大语言模型。

Single-Node Distributed Communication in PyTorch

我们先看看一个简单的Pytorch分布式应用示例，其目标是生成四个随机整数张量，并计算他们的和。

在下面这个场景中，我们会启动四个工作进程，每个进程都会生成一个随机张量。为了在这些工作进程之间对这些张量进行求和，我们会调用all-reduce集体通信操作。这个操作会将每个进程上的原始数据张量替换为all-reduce后的结果，下面是示例代码:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20


import os
import torch
import torch.distrubuted 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)

运行该脚本后我们得到如下输出：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


$ uv run python assignment2-systems/distribute_example.py
[Gloo] Rank 1 is connected to 3 peer ranks. Expected number of connected peer ranks is : 3
[Gloo] Rank 3 is connected to 3 peer ranks. Expected number of connected peer ranks is : 3
[Gloo] Rank 0 is connected to 3 peer ranks. Expected number of connected peer ranks is : 3
[Gloo] Rank 2 is connected to 3 peer ranks. Expected number of connected peer ranks is : 3
rank 0 data (before all-reduce): tensor([4, 4, 2])
rank 1 data (before all-reduce): tensor([5, 3, 5])
rank 2 data (before all-reduce): tensor([2, 5, 3])
rank 3 data (before all-reduce): tensor([4, 9, 1])
rank 0 data (after all-reduce): tensor([15, 21, 11])
rank 1 data (after all-reduce): tensor([15, 21, 11])
rank 2 data (after all-reduce): tensor([15, 21, 11])
rank 3 data (after all-reduce): tensor([15, 21, 11])

正如我们所预期的一样，每个工作进程一开始都持有不同的数据张量。在执行完all-reduce后，这些张量会在所有工作进程之间进行求和，并且每个工作进程中的data都会被原地修改为all-reduce的结果。

我们现在可以再看看上面的脚本。命令mp.spawn会启动nprocs个进程,这些进程会使用提供的args参数去运行函数fn。此外函数fn会以fn(rank,*args)的形式被调用，其中rank为工作进程的索引。因此，我们的工作函数接受的第一个参数必须是这个rank。

这些工作进程都属于一个进程组(process group)，这个进程组可以通过dist.init_process_group进行初始化。进程组表示多个工作进程，它们会通过一个共享的 master 来进行协调和通信。master 由它的 IP 地址和端口定义，而 rank 为 0 的进程就是 master 所在的进程。

像 all-reduce 这样的集体通信操作，会作用于进程组中的每一个进程。

在上面的例子中我们使用的后端是gloo，实际上还有多种选择，特别是nccl，它会使用NVIDIA 的 NCCL 集体通信库，对于 CUDA 张量来说，通常会有更高的性能。不过，NCCL 只能在带有 GPU 的机器上使用，而 Gloo 可以运行在仅有 CPU 的机器上。一个很实用的经验法则是：分布式 GPU 训练使用 NCCL，分布式 CPU 训练和/或本地开发使用 Gloo。在这个示例中我们选择 Gloo，是因为它支持在仅有 CPU 的机器上进行本地运行和开发。

在运行多 GPU 任务时，要确保不同的 rank 使用不同的 GPU。一种实现方式是在 setup 函数中调用 torch.cuda.set_device(rank)，这样 tensor.to("cuda") 就会自动把张量移动到指定的设备上。另一种方式是显式地为每个 rank 创建一个设备字符串（例如 device = f"cuda:{rank}"），然后在进行任何数据移动时，把这个设备字符串作为目标设备来使用（例如 tensor.to(f"cuda:{rank}")）

在原有代码的基础上做了一定改进，作为一个benchmark代码在CPU+Gloo环境下进行测试，测试结果如下：

target	world_size	tensor_size	median_ms	mean_ms	algo_GiBps
gloo-cpu	2	1MB	8.476	8.517	0.115
gloo-cpu	2	10MB	50.797	51.094	0.192
gloo-cpu	2	100MB	523.528	526.105	0.187
gloo-cpu	2	1GB	5398.399	5398.969	0.185
gloo-cpu	4	1MB	18.153	18.488	0.081
gloo-cpu	4	10MB	109.498	110.754	0.134
gloo-cpu	4	100MB	901.176	911.375	0.163
gloo-cpu	4	1GB	7817.207	7812.762	0.192
gloo-cpu	6	1MB	19.964	20.654	0.082
gloo-cpu	6	10MB	97.604	97.988	0.167
gloo-cpu	6	100MB	879.720	877.797	0.185
gloo-cpu	6	1GB	20764.952	20206.942	0.080
gloo-cpu	8	1MB	19.493	23.385	0.088
gloo-cpu	8	10MB	189.436	190.334	0.090
gloo-cpu	8	100MB	1887.865	1912.338	0.091
gloo-cpu	8	1GB	23195.781	22370.482	0.075

由图可以看到

整体趋势是：tensor size 变大后，latency 基本随之上升；world size 从 2 增加到 8 后，整体延迟变大、扩展性变差。小消息（1MB）时固定开销和同步开销占主导，所以带宽利用率低；中等消息（10MB 到 100MB）时效率最好；超大消息（1GB）时，多进程配置尤其是 p=6/8 明显退化。

几个具体点：

p=2 最稳定，algo_GiBps 大致维持在 0.18~0.19 GiB/s，说明两进程下通信效率最好。
p=4 虽然延迟比 p=2 高，但在 1GB 时带宽还能到 0.192 GiB/s，表现还可以。
p=6 和 p=8 在 10MB/100MB 还能接受，但到 1GB 时延迟陡增到 20~23s，带宽掉到 0.08 左右，说明大规模下通信瓶颈非常明显。
图里 p=6 在 10MB、100MB 甚至比 p=4 略好，这更像是测试波动、拓扑/调度差异，而不是稳定规律；但 1GB 的恶化很明显，说明趋势仍然是进程数越多越难扩展。

简单解释就是：gloo-cpu 在当前环境下更适合较小规模或中等消息量，随着参与进程增多，通信同步、链路竞争和 CPU 端开销会迅速放大，导致大张量场景下吞吐下降明显。可以概括成一句：

gloo-cpu 在 2~4 个进程时扩展较平稳，但在 6~8 个进程、尤其 1GB 大张量下出现明显的通信退化，说明系统已进入带宽竞争和同步开销主导的区间。

后续在学校的集群上利用多卡进行了一个小测试：

target	world_size	tensor_size	median_ms	mean_ms	algo_GiBps
nccl-cuda	2	1MB	0.125	0.125	7.838
nccl-cuda	2	10MB	0.660	0.687	14.789
nccl-cuda	2	100MB	5.596	5.593	17.452
nccl-cuda	2	1GB	51.969	52.684	19.242
nccl-cuda	4	1MB	0.108	0.109	13.614
nccl-cuda	4	10MB	0.296	0.302	49.437
nccl-cuda	4	100MB	2.134	2.137	68.647
nccl-cuda	4	1GB	20.860	20.900	71.908

A Naïve Implementation of Distributed Data Parallel Training

我们现在已经了解了如何在Pytorch中编写分布式应用的基础，接下来我们来构建一个分布式数据并行训练（DDP）的最小实现。

数据并行会把一个Batch的数据切分到多个设备上，从而支持使用单个设备无法容纳的大batch进行训练。例如，我们现在有4个设备，每个设备最多只能处理的batch size为32，那么数据并行训练就能实现等效batch size 为32 x 4 = 128

下面是一种朴素方式实现分布式数据并行的步骤:起初，每个设备都会构建一个模型(随机初始化的)。我们使用broadcast集体通信操作，将rank 0 上的模型参数发送给其他所有rank。训练开始时，每个设备都持有完全相同的一份模型参数和优化器状态(比如Adam中累积的梯度等)

给定一个包含n个样本的batch，将该batch切分后，每个设备拟接收到n/d个互不重叠的样本（其中d是用于数据并行训练的设备数量）。n应当可以被d整除。
每个设备使用自己本地的一份模型参数，对其收到的n/d个样本进行前向反向传播，以计算梯度。需要注意的是此时每个设备只持有基于自己收到的这 n/d 个样本计算得到的梯度。
然后，我们使用all-reduce集体通信操作，对不同设备上的梯度求平均，这样每个设备都会持有基于全部n个样本平均后的梯度。
接下来每个设备执行一次优化器更新步骤，更新自己那份参数副本——从优化器的角度看，它只是在优化一个本地模型。由于所有设备都从相同的初始模型和优化器状态开始，并且在每次迭代中都使用相同的平均梯度，因此各设备上的参数和优化器状态会始终保持同步。至此，我们就完成了一次训练迭代，然后可以重复这一过程。

按照上面的逻辑，我编写了如下代码，并与baseline的模型进行了对比：

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108


import torch
import torch.nn as nn
from einops import rearrange

import torch.distributed as dist
import torch.multiprocessing as mp
from torch.optim import AdamW

import os
from copy import deepcopy

class ToyMLP(nn.Module):
 def __init__(self, d_in=16, d_hidden=32, d_out=4):
 super().__init__()
 self.fc1 = nn.Linear(d_in, d_hidden)
 self.act = nn.ReLU()
 self.fc2 = nn.Linear(d_hidden, d_out)

 def forward(self, x):
 x = self.fc2(self.act(self.fc1(x)))
 return x

def setup(rank, world_size, backend):
 os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29500"
 init_kwargs = {"backend": backend, "rank": rank, "world_size": world_size}
 if backend == "nccl":
 init_kwargs["device_id"] = torch.device(f"cuda:{rank}")
 dist.init_process_group(**init_kwargs)

def get_device(rank, backend):
 if backend == "nccl":
 torch.cuda.set_device(rank)
 return torch.device(f"cuda:{rank}")
 return torch.device("cpu")

def distributed_train(rank, world_size, backend, x_rand, y_rand):
 torch.manual_seed(rank)
 setup(rank, world_size, backend)
 device = get_device(rank, backend)
 model = ToyMLP().to(device)
 x_rand = x_rand.to(device)
 y_rand = y_rand.to(device)

 with torch.no_grad():
 for param in model.parameters():
 dist.broadcast(param, src=0)

 if rank == 0:
 model_baseline = deepcopy(model)
 baseline_opt = AdamW(
 model_baseline.parameters(),
 lr=1e-4,
 betas=(0.9, 0.999),
 eps=1e-8,
 weight_decay=1e-2,
 )
 loss_fn = nn.MSELoss(reduction="mean")
 baseline_opt.zero_grad()
 pred = model_baseline(x_rand)
 loss = loss_fn(pred, y_rand)
 loss.backward()
 baseline_opt.step()

 x_local = rearrange(x_rand, "(d b) f -> d b f", d=world_size)[rank]
 y_local = rearrange(y_rand, "(d b) f -> d b f", d=world_size)[rank]
 optimizer = AdamW(
 model.parameters(),
 lr=1e-4,
 betas=(0.9, 0.999),
 eps=1e-8,
 weight_decay=1e-2,
 )
 loss_fn = nn.MSELoss(reduction="mean")
 optimizer.zero_grad()
 pred = model(x_local)
 loss = loss_fn(pred, y_local)
 loss.backward()

 for param in model.parameters():
 if param.grad is None:
 continue
 dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
 param.grad /= world_size

 optimizer.step()

 dist.barrier()

 if rank == 0:
 for i, (p_base, p_ddp) in enumerate(zip(model_baseline.parameters(), model.parameters())):
 max_diff = (p_base - p_ddp).abs().max().item()
 print(f"param {i}: max_diff = {max_diff:.8e}")

 dist.destroy_process_group()

if __name__ == "__main__":
 torch.manual_seed(0)
 if torch.cuda.is_available():
 world_size = torch.cuda.device_count()
 backend = "nccl"
 else:
 world_size = 4
 backend = "gloo"
 batch_size = 128
 x_rand = torch.randn(batch_size, 16)
 y_rand = torch.randn(batch_size, 4)
 mp.spawn(fn=distributed_train, args=(world_size, backend, x_rand, y_rand), nprocs=world_size, join=True)

最终输出如下：

1
2
3
4


param 0: max_diff = 3.72529030e-09
param 1: max_diff = 5.82076609e-11
param 2: max_diff = 0.00000000e+00
param 3: max_diff = 0.00000000e+00

可以看到几乎是没有差距的，说明我们的DDP Train的实现是正确的。

Improving Upon the Minimal DDP Implementation

我们目前实现的Naive DDP的实现有几个关键的局限：

它会对每个张量分别执行一次all-reduce操作。每次通信调用都会产生开销，因此将多个通信调用进行批处理以减少这类开销，可能会更有利。
它会等到整个反向传播完成之后才开始通信梯度。但实际上，反向传播是逐步计算的。因此，当某个参数的梯度已经就绪时，就可以立刻对它进行通信，而不必等待其他参数的梯度也全部就绪。这使我们可以将梯度通信与反向传播计算重叠，从而减少分布式数据并行训练的开销。

Reducing the Number of Communication Calls

与其为每个参数都发起一次通信调用，不如看看能否通过批量执行all-reduce来提升性能。具体而言，我们会将需要进行all-reduce的梯度拼接成一个单独的张量，然后在所有的rank上对这个合并后的梯度张量执行一次all-reduce。此时两个比较重要的API是:torch._utils._flatten_dense_tensors和torch._utils._unflatten_dense_tensors

我们首先简单学习一下这两个API的作用：

他们的作用分别是：

torch._utils._flatten_dense_tensors(tensor_list)输入一个由多个dense tensor组成的list，输出一个由这个list中的dense tensor拼接成的连续大tensor，通常是一维的

torch._utils._unflatten_dense_tensors(tensor,tensor_list)输入就是刚才拼接成的大tensor，以及原始的tensor列表，用来提供每个张量的 shape / numel / dtype 信息。输出就是拆分后的 tensor 列表，形状和 tensor_list 一一对应

因此，我们可以在对梯度进行reduce前将他们先利用flatten拼接起来，然后在reduce后unflatten为原始的形状从而大大减少通信开销

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61


def reduce_less_distributed_train(rank, world_size, backend, x_rand, y_rand):
 torch.manual_seed(rank)
 setup(rank, world_size, backend)
 device = get_device(rank, backend)
 model = ToyMLP().to(device)
 x_rand = x_rand.to(device)
 y_rand = y_rand.to(device)

 params = [p.data for p in model.parameters()]
 flat_params = _flatten_dense_tensors(params)
 dist.broadcast(flat_params, src=0)
 params = _unflatten_dense_tensors(flat_params, params)
 with torch.no_grad():
 for p, p_new in zip(model.parameters(), params):
 p.data.copy_(p_new)
 if rank == 0:
 model_baseline = deepcopy(model)
 baseline_opt = AdamW(
 model_baseline.parameters(),
 lr=1e-4,
 betas=(0.9, 0.999),
 eps=1e-8,
 weight_decay=1e-2,
 )
 loss_fn = nn.MSELoss(reduction="mean")
 baseline_opt.zero_grad()
 pred = model_baseline(x_rand)
 loss = loss_fn(pred, y_rand)
 loss.backward()
 baseline_opt.step()
 
 x_local = rearrange(x_rand, "(d b) f -> d b f", d=world_size)[rank]
 y_local = rearrange(y_rand, "(d b) f -> d b f", d=world_size)[rank]
 optimizer = AdamW(
 model.parameters(),
 lr=1e-4,
 betas=(0.9, 0.999),
 eps=1e-8,
 weight_decay=1e-2,
 )
 loss_fn = nn.MSELoss(reduction="mean")
 optimizer.zero_grad()
 pred = model(x_local)
 loss = loss_fn(pred, y_local)
 loss.backward()
 params = [p for p in model.parameters() if p.grad is not None]
 grads = [p.grad for p in params]
 flat_grads = _flatten_dense_tensors(grads)
 dist.all_reduce(flat_grads, op=dist.ReduceOp.SUM)
 flat_grads /= world_size
 grads = _unflatten_dense_tensors(flat_grads, grads)
 with torch.no_grad():
 for p, g in zip(params, grads):
 p.grad.copy_(g)
 optimizer.step()
 dist.barrier()

 if rank == 0:
 for i, (p_base, p_ddp) in enumerate(zip(model_baseline.parameters(), model.parameters())):
 max_diff = (p_base - p_ddp).abs().max().item()
 print(f"param {i}: max_diff = {max_diff:.8e}")

后续我们直接在我们在assignment1中搭建的Transformer中进行了实验，实验代码见仓库Navie_DDP.py

结果如下，可以看到，每步确实更快了。

per-parameter all-reduce: 0.218 s/step

model config: d_model=1600, d_ff=6400, num_layers=48, num_heads=25

flattened all-reduce: 0.207 s/step

model config: d_model=1600, d_ff=6400, num_layers=48, num_heads=25

Overlapping Computation with Communication of Individual Parameter Gradients

虽然我们对通信调用进行了批处理，这或许有助于降低大量小型all-reduce操作，但所有通信时间仍然会直接构成总开销。

为了解决这个问题，我们可以利用这样一个事实：反向传播会逐层增量式地计算梯度（从损失开始，朝输入方向移动）——因此，我们可以在参数梯度一旦准备好之后就立刻对其执行 all-reduce，通过将反向传播计算与梯度通信重叠，来降低数据并行训练的开销。

我们将先实现并基准测试一个分布式数据并行包装器：当某个参数张量在反向传播中一旦准备好时，就异步地对该参数张量执行 all-reduce。下面这些提示可能会有帮助：

反向传播Hook

为了在某个参数的梯度在反向传播过程中被累积完成后，自动调用一个函数，我们可以使用：

register_post_accumulate_grad_hook函数：

其作用是在反向传播时，等某个参数的梯度都累加完，param.grad已经写好后，再执行我们注册的回调。

一个简单的使用示例如下:

1
2
3
4


p = torch.nn.Parameter(torch.randn(3))
def hook(param):
 param.data -= 0.01 * param.grad
h = p.register_post_accumulate_grad_hook(hook)

那么这段代码的语义就是，当p的梯度在反向传播中被计算出来后，hook就会被调用，此时我们可以直接读param.grad并更新param

异步通信

Pytorch中所有集合通信操作都支持同步执行和异步执行(通过参数async_op进行区分前者为false后者为true)

同步调用会阻塞，直到该集合通信操作被排入 GPU 队列中。这并不意味着 CUDA 操作本身已经完成，因为 CUDA 操作是异步的。尽管如此，后续依赖该输出的函数调用会按预期工作。
异步调用则会返回一个分布式请求句柄；因此，当函数返回时，该集合通信操作并不保证已经被排入 GPU，更不用说已经完成。若要等待该操作被排入 GPU（从而使得输出可以被后续操作安全使用），你可以对返回的通信句柄调用 handle.wait()。

我们可以通过下面的例子进行学习：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23


tensors = [torch.rand(5) for _ in range(10)]

# 同步方式：阻塞直到操作被排入 GPU。
for tensor in tensors:
 dist.all_reduce(tensor, async_op=False)

# 异步方式：每次调用后立即返回，
# 最后统一等待结果。
handles = []
for tensor in tensors:
 handle = dist.all_reduce(tensor, async_op=True)
 handles.append(handle)

# ...
# 此处可以执行其他不依赖 all_reduce 结果的操作
# ...

# 确保所有 all-reduce 调用都已经被排入队列，
# 从而后续依赖 all-reduce 输出的其他操作
# 也可以被排入队列。
for handle in handles:
 handle.wait()
handles.clear()

这段代码可以明显的看出同步方式和异步方式的差异，对于同步而言，只有上一个tensor的all_reduce完成了以后才会发送下一个all_reduce。

而异步的话就是先把很多个all_reduce发送出去，每次返回一个handle，我们可以在中间穿插一些不需要这个all_reduce结果的工作，等到最终需要用这些结果的时候，再统一wait。

在 DDP 中，反向传播时参数梯度是逐步变为 ready 的。某个参数（先不考虑 bucket）的梯度一旦 ready，DDP 通过预先注册的 autograd hook 发起对应的异步梯度同步；与此同时，backward 继续去计算后续尚未完成的梯度。等所有待同步的梯度都已经发起同步后，DDP 再在 backward 的同步边界上等待这些通信完成，随后各 rank 上对应参数的 grad 保持一致。

实现

那么我们整体的实现流程就很明确了，那么我们需要实现的分布式数据并行包装器的功能就是：

在训练开始前广播权重
发起用于梯度平均的通信调用（且给每个参数插入一个hook）

在整体流程上，我们希望的是:

1
2
3
4
5
6
7
8
9


ddp_model = MyDDP(model)

optimizer.zero_grad()
logits = ddp_model(tokens)
loss = ...
loss.backward()

ddp_model.finish_gradient_synchronization()
optimizer.step()

在loss.backward()中autograd会逐步计算出各个参数的梯度，一旦某个参数的p.grad计算完成，就会触发我们注册的register_post_accumulate_grad_hook，在hook中对发起all_reduce。

当backward全部结束之后，我们再统一finish_gradient_synchronization()

按照上述逻辑，以及指导手册中推荐的类函数，我们可以写作如下形式：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30


class DDP(nn.Module):
 def __init__(self, module):
 super().__init__()
 self.module = module
 self.world_size = dist.get_world_size()
 self.handles = []

 with torch.no_grad():
 for param in self.module.parameters():
 dist.broadcast(param, src=0)

 for param in self.module.parameters():
 if not param.requires_grad:
 continue
 def make_hook(p):
 def hook(_):
 if p.grad is None:
 return
 handle = dist.all_reduce(p.grad, op=dist.ReduceOp.SUM, async_op=True)
 self.handles.append((handle, p))
 return hook
 param.register_post_accumulate_grad_hook(make_hook(param))
 def forward(self, *inputs, **kwargs):
 return self.module(*inputs, **kwargs)
 
 def finish_gradient_synchronization(self):
 for handle, param in self.handles:
 handle.wait()
 param.grad /= self.world_size
 self.handles = []

如下是对比结果:

per-parameter all-reduce: 0.504 s/step

model config: d_model=1600, d_ff=6400, num_layers=48, num_heads=25

flattened all-reduce: 0.501 s/step

model config: d_model=1600, d_ff=6400, num_layers=48, num_heads=25

async per-parameter all-reduce: 0.497 s/step

model config: d_model=1600, d_ff=6400, num_layers=48, num_heads=25

Overlapping Computation with Communication of Bucketed Parameter Gradients

目前为止我们成功将反向传播计算与单个参数梯度的通信进行了重叠。然而，我们之前已经观察到，对通信调用进行批处理通常可以提升性能，尤其是在拥有大量参数张量的时候。

我们之前的处理方式就是一次性发送所有的梯度，但这要求等到整个反向传播结束后才能开始。

那么我们一个自然的想法就是，能否将两种方法的优点结合起来: 把参数组织到若干个bucket中，并在每个bucket中所有张量的梯度都准备好后，立即对该bucket执行all-reduce操作。

此外，最好的实现方式是按照parameters的逆序入桶，因为实际反向传播的过程中梯度的准备顺序也是这样的。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72


class Bucket_DDP(nn.Module):
 def __init__(self, module,bucket_size_mb=10):
 super().__init__()
 self.module = module
 self.world_size = dist.get_world_size()
 self.bucket_size_mb = bucket_size_mb

 with torch.no_grad():
 for param in self.module.parameters():
 dist.broadcast(param, src=0)
 params = [p for p in self.module.parameters() if p.requires_grad]
 params = list(reversed(params))
 
 self.buckets = []
 current_bucket = []
 current_size = 0

 for param in params:
 p_size = param.numel()
 if current_bucket and current_size + p_size > bucket_size_mb * 1024 * 1024 // 4:
 self.buckets.append(current_bucket)
 current_bucket = []
 current_size = 0
 current_bucket.append(param)
 current_size += p_size
 if current_bucket:
 self.buckets.append(current_bucket)
 self.param2bucket = {}
 for bucket_idx, bucket in enumerate(self.buckets):
 for p in bucket:
 self.param2bucket[p] = bucket_idx

 self.bucket_ready_count = [0 for _ in self.buckets]
 self.bucket_handles = [None for _ in self.buckets]
 self.bucket_flat_grads = [None for _ in self.buckets]

 for p in params:
 if not p.requires_grad:
 continue
 def make_hook(p):
 def hook(_):
 if p.grad is None:
 return
 bucket_idx = self.param2bucket[p]
 self.bucket_ready_count[bucket_idx] += 1
 if self.bucket_ready_count[bucket_idx] == len(self.buckets[bucket_idx]):
 grads = [param.grad for param in self.buckets[bucket_idx]]
 flat_grads = _flatten_dense_tensors(grads)
 handle = dist.all_reduce(flat_grads, op=dist.ReduceOp.SUM, async_op=True)
 self.bucket_flat_grads[bucket_idx] = flat_grads
 self.bucket_handles[bucket_idx] = handle
 return hook
 p.register_post_accumulate_grad_hook(make_hook(p))

 def forward(self, *inputs, **kwargs):
 return self.module(*inputs, **kwargs)
 
 def finish_gradient_synchronization(self):
 for bucket_idx, handle in enumerate(self.bucket_handles):
 handle = self.bucket_handles[bucket_idx]
 if handle is None:
 continue
 handle.wait()
 flat_grads = self.bucket_flat_grads[bucket_idx]
 flat_grads /= self.world_size
 grads = [p.grad for p in self.buckets[bucket_idx]]
 synced_grads = _unflatten_dense_tensors(flat_grads, grads)
 for param, grad in zip(self.buckets[bucket_idx], synced_grads):
 param.grad.copy_(grad)
 self.bucket_ready_count[bucket_idx] = 0
 self.bucket_handles[bucket_idx] = None
 self.bucket_flat_grads[bucket_idx] = None

4D Parallelism

事实上，在工业界，并行比我们了解的要更多。目前主流的做法是4D Parallelism:

数据并行（DP），也就是我们目前了解的，将同一批次的数据切分到多个设备上，每个设备对自己那一份batch计算梯度。随后必须以某种方式在设备之间对这些梯度求平均。
全分片数据并行（FSDP），优化器状态、梯度和权重都在设备之间切分。如果我们只使用 DP 和 FSDP，那么在执行前向传播或反向传播之前，每个设备都需要从其他所有设备收集权重分片。
张量并行（TP）——激活值沿一个新的维度被分片，每个设备只计算自己那一片对应的输出结果。在张量并行中，我们既可以沿被切分操作的输入维度进行分片，也可以沿输出维度进行分片。如果权重和激活值沿相对应的维度进行切分，那么张量并行可以和 FSDP 高效结合使用。
流水线并行（PP）——模型按层切分为多个阶段，每个阶段运行在不同的设备上。
专家并行（EP）——将专家（在混合专家模型，即 Mixture-of-Experts，MoE 中）分布到不同设备上，每个设备只计算自己所负责专家的输出结果。

通常，我们总是将 FSDP 和 TP 结合使用，因此可以把它们看作并行的同一个轴。这样就剩下 4 个并行轴：DP、FSDP/TP、PP 和 EP。我们还将重点关注稠密模型（而不是 MoE），因此不再进一步讨论 EP。

在分析分布式训练时，我们经常把集群描述为一个设备网格（mesh），其中网格的各个轴就对应我们定义并行方式的那些轴。比如说，如果我们有 16 张 GPU，而模型又远大于单个设备所能容纳的规模，那么我们可能会倾向于把设备网格组织成一个 4 × 4 的 GPU 网格，其中第一个维度表示 DP，第二个维度表示组合后的 FSDP 和 TP。

优化器状态共享

DDP在概念上十分简单，而且通常十分有效，但是我们能很明显地发现它存在的缺陷，它需要每个rank都持有一份独立的模型参数和优化器状态副本。这种冗余会带来显著的内存开销。比如，AdamW 优化器会为每个参数维护两个浮点数，这意味着它占用的内存是模型权重的两倍。

Rajbhandari 等人提出了几种方法，通过将以下内容在各个 rank 之间进行分片来减少数据并行训练中的这种冗余：（1）优化器状态，（2）梯度，以及（3）参数；并在需要时在各个 worker 之间进行通信。

接下来，我们将实现一个简单的优化器分片器来降低每个rank的内存损耗。与其为所有参数都保存优化器状态，不如让每个 rank 上的优化器实例只处理其中一部分参数（大约是 1 / world_size）。当每个 rank 的优化器执行一次 optimizer step 时，它只会更新自己分片中的那部分模型参数。然后，每个 rank 会将自己更新后的参数广播给其他 rank，以确保每次 optimizer step 之后，所有 rank 上的模型参数仍然保持同步。

根据实验手册，我们实现的分片器的实现如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69


from __future__ import annotations

from typing import Any, Type

import torch
import torch.distributed as dist
import torch.optim as optim

class OptimizerStateSharding(optim.Optimizer):
 def __init__(self, params, optimizer_cls: Type[optim.Optimizer], **kwargs: Any):
 if not dist.is_initialized():
 raise RuntimeError("OptimizerStateSharding requires torch.distributed to be initialized.")

 self.optimizer_cls = optimizer_cls
 self.optimizer_kwargs = dict(kwargs)
 self.world_size = dist.get_world_size()
 self.rank = dist.get_rank()
 self._owned_rank_by_param = {}
 self._local_param_groups = []
 self._next_param_index = 0
 self.local_optimizer = None
 self._initializing = True

 super().__init__(params, kwargs)

 self._initializing = False
 non_empty_local_groups = [group for group in self._local_param_groups if group["params"]]
 if non_empty_local_groups:
 self.local_optimizer = optimizer_cls(non_empty_local_groups, **kwargs)

 @torch.no_grad()
 def step(self, closure=None, **kwargs):
 loss = None
 if self.local_optimizer is not None:
 if closure is not None:
 with torch.enable_grad():
 loss = self.local_optimizer.step(closure=closure, **kwargs)
 else:
 loss = self.local_optimizer.step(**kwargs)

 for group in self.param_groups:
 for param in group["params"]:
 dist.broadcast(param.data, src=self._owned_rank_by_param[param])

 return loss

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

 added_group = self.param_groups[-1]
 local_params = []
 for param in added_group["params"]:
 owner_rank = self._next_param_index % self.world_size
 self._owned_rank_by_param[param] = owner_rank
 if owner_rank == self.rank:
 local_params.append(param)
 self._next_param_index += 1

 local_group = {key: value for key, value in added_group.items() if key != "params"}
 local_group["params"] = local_params
 self._local_param_groups.append(local_group)

 if self._initializing or not local_params:
 return

 if self.local_optimizer is None:
 self.local_optimizer = self.optimizer_cls([local_group], **self.optimizer_kwargs)
 else:
 self.local_optimizer.add_param_group(local_group)

[CS336] Assignment 1：Basics

Fri, 09 Oct 2026 10:00:00 +0800

环境配置

首先克隆课程仓库到本地的WSL中：

1

git clone https://github.com/stanford-cs336/assignment1-basics.git

接着按照仓库要求安装环境即可，分为环境安装和数据下载，环境的话:

1
2
3
4


#先在WSL上安装uv
curl -LsSf https://astral.sh/uv/install.sh | sh 
source $HOME/.local/bin/env
uv run pytest tests

数据的话就按照readme中的介绍正常下载就行。

下面是课程实验指导手册：

cs336_assignment1_basics.pdf

本人项目地址：https://github.com/July-h5kf3/CS336/tree/main

文本编码和Tokenizer

ASCII Unicode与UTF-8编码

ASCII编码在我们C++课程中就已经介绍了，这里不多赘述，见下表（0-31以及127是控制字符）

Unicode则是一个文本编码标准，不同于ASCII，Unicode会将诸如汉字，emoji等等的字符都有一个整形数字与之对应(如"牛"对应的数字为29275)，这个数字我们称为码点(code point).

在最新的Unicode 16.0中，已经包含了154,998个字符。假设从0开始编码，那么用二进制表示的话至少需要18位，假设每个字符定长，那么一个字符我们需要3B来表示。

实际上Unicode的码点到了0x10FFFF，也就是说如果采取定长的话就需要21bit。此外，为了兼容ASCII，每个B都需要牺牲一定的bit来做标识

因此最终如果采用定长的话，一个字符就需要4B来表示。这种编码方式我们称为UTF-32

这种方式显然会浪费大量的空间，一个改良的方式就是采用变长存储。比如0-255的用1B等，这种编码方式为UTF-8.

综上，Unicode标准定义了从字符到码点(整数)的映射，但是由于词表的规模过于庞大且稀疏，直接在Unicode码点上训练分词器是不现实的。因此通常的做法是将一个Unicode字符转换为一系列字节，也就是采用UTF-8编码。

我们接下来看看是如何进行这个过程的。上面我们在说为什么是UTF-8的时候说到之所以是4B而不是3B是因为需要一定的bit来进行标识，标识的作用就是用于判断该字符占几个字节。

其中前缀可以分为:

控制位：0,110,1110,11110就分别告知后面还有多少个字节(B)
延续位：10开头则表示这是一个”从属字节“，不是新字符的开头

我们以牛(unicode为29275，十六进制为0x725B)这个字符为例进行编码的模拟

将其转化为二进制后有:0111001001011011,且我们可以判断UTF-8需要3字节，因此使用3字节的UTF-8

带入3B的模板:

11100111 10001001 10011011

转化为十六进制即为:

牛 = E7 89 9B(231,137,155)

BPE Tokenizer的实现

BPE算法原理与训练实现

虽然字节级分词(如UTF-8)可以缓解词级分词器所面临的词表外(OOV)问题，但直接将文本分解为字节会导致冗长的输入序列。这会减慢模型训练速度，因为在词级语言模型中，一个包含 10 个词的句子可能只需要 10 个 token，而在字符级模型中（取决于词的长度），同样的句子可能需要 50 个甚至更多的 token。处理更长的序列会使模型在每一步都需要更多的计算量。

此外，在字节序列上进行语言建模是比较困难的，因为更长的输入序列会在数据中引入更强的长期依赖问题。

子词分词(subword tokenization)位于词级分词和字节级分词之间，是一种折中方案。需要注意的是，字节级分词器的词表大小固定为 256（字节值范围为 0 到 255）。子词分词器通过使用更大的词表，来换取对输入字节序列更好的压缩效果。例如，如果字节序列 b'the' 在原始训练数据中频繁出现，那么为它在词表中分配一个条目，就可以将原本由 3 个 token 组成的序列压缩为 1 个 token。

那么我们如何选择要加入词表的子词单元呢？

目前的主流方法是字节对编码即（BPE），这是一种压缩算法，它通过迭代的方式，将出现频率最高的一对字节替换（“合并”）为一个新的、尚未使用的索引。需要注意的是，该算法通过向词表中加入子词 token 来最大化输入序列的压缩率——如果某个词在输入文本中出现得足够频繁，它最终就会被表示为一个单独的子词单元。

使用 BPE 构建词表的子词分词器通常被称为 BPE 分词器。

这里举个例子来说明，假设一个字符的unicode转化为UTF-8后为[239,165,32],此时我们的词表中有:

{256:[239,165],257:[256,32]}

那么由于最终该字符的表示就是[257].这个字典就是我们得到的BPE

想要训练一个BPE分词器需要三个步骤。

词表初始化

分词器的词表是从字节串 token 到整数 ID 的一一映射。由于我们训练的是字节级 BPE 分词器，初始词表就是所有可能的字节集合。因为字节一共有 256 种可能取值，所以初始词表大小为 256。

预分词

在拥有词表之后，理论上我们可以直接统计语料中哪些字节经常相邻出现，并从出现频率最高的字节对开始进行合并。然而，这样做的计算代价非常高，因为每进行一次合并，都需要对整个语料做一次完整遍历。

此外，直接在整个语料上合并字节，可能会产生只在标点符号上有所不同的 token（例如 dog! 和 dog.）。尽管它们在语义上非常相近（只差一个标点），却会被分配完全不同的 token ID。

为了解决这些问题，我们会先对语料进行预分词。可以将其理解为一种粗粒度的分词方式，用于帮助我们更高效地统计字符对出现的频率。例如，单词 "text" 可能作为一个预分词单元出现了 10 次。那么在统计字符 't' 和 'e' 相邻出现的次数时，我们只需知道 "text" 中 't' 和 'e' 是相邻的，就可以一次性将它们的计数增加 10，而不需要逐字遍历整个语料。

由于我们训练的是字节级 BPE 模型，每一个预分词单元都会被表示为一串 UTF-8 字节序列。

在本项目中我们将使用一种基于正则表达式的预分词器，其定义如下:

1
2
3
4
5
6
7
8


PAT = r"""
'(?:[sdmt]|ll|ve|re) #英文缩写如 's,'d等
| ?\p{L}+ #单词
| ?\p{N}+ #阿拉伯数字
| ?[^\s\p{L}\p{N}]+ #标点符号
|\s+(?!\S) #行尾空白
|\s+ # 其他空白
"""

为了更好理解这个预分词器的行为，可以看下面的程序的运行结果：

1
2
3


# 需要安装 `regex` 包
import regex as re
re.findall(PAT, "some text that i'll pre-tokenize")

其输出为:

1

['some', ' text', ' that', ' i', "'ll", ' pre', '-', 'tokenize']

计算BPE合并

在将输入文本转化为预分词，并将每个预分词表示为UTF-8字节序后，我们就可以开始计算BPE合并操作（即训练BPE分词器）。

从整体上来看，BPE算法会反复统计所有字节对的出现频率，并找出出现次数最高的一对字节(“A”,“B”)。然后，将语料中所有该字节对(“A”,“B”)的出现位置合并，替换为一个新的token “AB”.这个新的合并token会被加入到词表中。

因此BPE训练完成后的最终词表大小，等于初始词表大小加上训练过程中BPE合并的次数。

为了提升训练效率，在BPE训练过程中，我们不考虑跨越预分词边界的字节对。当多个字节对具有相同的最高频率时，需要以确定性的方式打破平局，我们采用的方式是选择字典序更大的那一对。

特殊Token的处理

在实际应用中，某些字符串(例如 <|endoftext|>)常用于编码元数据（如文档之间的边界）。在对文本进行编码时，通常希望将这些字符串视为“特殊 token”，它们永远不应被拆分成多个 token，而是始终作为一个整体保留下来。

例如，序列结束标记 <|endoftext|> 应始终对应一个单独的 token（即一个整数 ID），这样语言模型才能明确知道何时停止生成文本。这些特殊 token 必须被显式加入词表，并分配固定的 token ID。

接下来我们进行具体实现：

首先进行词表的初始化，按照上面的内容，我们的词表初始时只有256个字节串token到整形的映射以及规定的特殊的token，具体实现代码如下:

我们找到test/adapters.py中的run_train_bpe函数，这是我们实现bpe分词器的主要部分。其中bytes是python内置的字节序列类型

1
2
3
4


vocab = {i: bytes([i]) for i in range(256)}
 for token in special_tokens:
 if token not in vocab.values():
 vocab[len(vocab)] = token.encode("utf-8")

接下来我们实现预分词器，项目中pretokenization_example.py中提供了参考，我们按照那个思路去做就行。

这里的预分词器采取了并行优化，也就是将文本拆分为了多个chunk同时进行预分词，example代码中提供了边界划分的代码。我们只需要设计每个chunk的预分词方法，并使用python中的 multiprocessing 库进行并行即可。

1
2
3
4
5
6
7
8
9


with open(input_path, "rb") as f:
 num_processes = os.cpu_count() or 4
 boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
 ranges = list(zip(boundaries[:-1], boundaries[1:]))
 with mp.Pool(processes=num_processes) as pool:
 results = pool.map(
 _pretokenize_range,
 [(input_path,s,e) for s,e in ranges],
 )

此外，在进行预分词前，我们还需要在预分词前删除特殊的token，在上面我们说到，我们是采用正则表达式进行的预分词，在此之前我们需要删除语料中的所有特殊token。

为了能够在后面的tokenizer中复用这里的预分词函数，我们选择采用如下方式进行处理。

1
2
3
4
5
6
7
8
9


def split_keep_special(text: str, special_tokens: list[str]) -> list[str]:
 #这里需要补充的一点在于，这里之所以对special token进行保留，是因为后续在tokenizer中我们可以直接复用这个函数
 if not special_tokens:
 return [text]
 special_tokens = sorted(special_tokens, key=lambda x: -len(x))
 #需要避免special token中有|等正则符号
 pattern = "(" + "|".join(re.escape(token) for token in special_tokens) + ")"
 parts = re.split(pattern, text)
 return [p for p in parts if p != ""]

这里需要注意，是我在后面实现tokenizer的时候发现的bug，就是由于会出现special token串联的情况，因此需要先将special token按照长度降序排列

比如:Special token = ["<|endoftext|>", “<|endoftext|><|endoftext|>”].我们需要先匹配后者，才能保证正确split

具体而言就是我们以special token为分割将一个chunk划分为若干个part，这样每个part要么是special token要么为单独的不含special token的文本。我们就可以对每个文本采用正则表达式的方式进行预分词了，预分词之后我们再统计每个token出现的次数即可。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


def _pretokenize_range(args: tuple[str | os.PathLike, int, int]) -> dict[bytes, int]:
 #首先将chunk按照special token进行split,避免出现跨doc的合并问题
 input_path, start, end = args
 token_counts = Counter()
 PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
 with open(input_path, "rb") as f:
 f.seek(start)
 to_read = end - start
 data = f.read(to_read)
 parts = split_keep_special(data.decode("utf-8", errors="ignore"), ["<|endoftext|>"])
 for part in parts:
 if part == "<|endoftext|>":
 #token_counts[part.encode("utf-8")] += 1
 continue
 for m in re.finditer(PAT, part):
 token_counts[m.group(0).encode("utf-8")] += 1
 return token_counts

我们得到了预分词的结果（即每个预分词token的频数）后，将预分词得到的每个token转化为UTF-8编码就能进行BPE的训练了。

BPE的训练简单来说就是对于每个预分词的token，我们会计算每个字节对出现的频率。注意是以token为单位进行字节对频率的统计的！(例如 hello world，我们统计的字节对就不会出现ow的统计)

然后将出现频率最高的字节对进行合并为一个并加入到词表中，并更新预分词的结果。不断迭代直到达到我们目标的词表大小（每次迭代词表大小增加1）

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85


def run_train_bpe(
 input_path: str | os.PathLike,
 vocab_size: int,
 special_tokens: list[str],
 **kwargs,
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
 """Given the path to an input corpus, run train a BPE tokenizer and
 output its vocabulary and merges.

 Args:
 input_path (str | os.PathLike): Path to BPE tokenizer training data.
 vocab_size (int): Total number of items in the tokenizer's vocabulary (including special tokens).
 special_tokens (list[str]): A list of string special tokens to be added to the tokenizer vocabulary.
 These strings will never be split into multiple tokens, and will always be
 kept as a single token. If these special tokens occur in the `input_path`,
 they are treated as any other string.

 Returns:
 tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
 vocab:
 The trained tokenizer vocabulary, a mapping from int (token ID in the vocabulary)
 to bytes (token bytes)
 merges:
 BPE merges. Each list item is a tuple of bytes (<token1>, <token2>),
 representing that <token1> was merged with <token2>.
 Merges are ordered by order of creation.
 """
 #词表的初始化，初始时词表应该只有从字节串 token 到整数 ID 的一一映射,以及规定的special tokens
 vocab = {i: bytes([i]) for i in range(256)}
 merges = []
 for token in special_tokens:
 b = token.encode("utf-8")
 if b not in vocab.values():
 vocab[len(vocab)] = b
 
 with open(input_path, "rb") as f:
 num_processes = os.cpu_count() or 4
 boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
 ranges = list(zip(boundaries[:-1], boundaries[1:]))
 total_token_counts = Counter()
 with mp.Pool(processes=num_processes) as pool:
 results = pool.map(
 _pretokenize_range,
 [(input_path,s,e) for s,e in ranges],
 )
 for token_counts in results:
 total_token_counts.update(token_counts)
 word_freqs = total_token_counts
 word_symbols = {
 token: tuple(vocab[b] for b in token)
 for token in word_freqs.keys()
 }
 while len(vocab) < vocab_size:
 pair_counts = Counter()
 for token,freq in word_freqs.items():
 symbols = word_symbols[token]
 if len(symbols) < 2:
 continue
 for a,b in zip(symbols, symbols[1:]):
 pair_counts[(a,b)] += freq
 if not pair_counts:
 break
 
 #合并出现频次最高的pair,如果有多个pair出现频次相同，则选择字典序最大的那个
 best_pair = max(pair_counts.items(), key=lambda kv: (kv[1], kv[0]))[0]
 #更新vocab和word_symbols
 merges.append(best_pair)
 merged_token = best_pair[0] + best_pair[1]
 vocab[len(vocab)] = merged_token

 for token in list(word_symbols.keys()):
 symbols = word_symbols[token]
 if len(symbols) < 2:
 continue
 new_symbols = []
 i = 0
 while i < len(symbols):
 if i < len(symbols) - 1 and (symbols[i], symbols[i+1]) == best_pair:
 new_symbols.append(merged_token)
 i += 2
 else:
 new_symbols.append(symbols[i])
 i += 1
 word_symbols[token] = tuple(new_symbols)
 return vocab, merges

至此，我们已经完成了BPE的训练，进行测试发现能通过所有测试点！

但是事实上是存在一定的优化空间的，因为每次合并发生后，我们会遍历所有token进行pair的统计，这个过程太浪费时间了，因此我们可以通过反向索引的方式进行进一步地优化。

具体而言，只有存在发生合并的pair的token才会出现统计值的变化，因此我们可以建立一个pair2token的索引，每次发生合并后只更新对应的token即可。

 1
 2
 3
 4
 5
 6
 7
 8
 9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102


def run_train_bpe(
 input_path: str | os.PathLike,
 vocab_size: int,
 special_tokens: list[str],
 **kwargs,
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
 """Given the path to an input corpus, run train a BPE tokenizer and
 output its vocabulary and merges.

 Args:
 input_path (str | os.PathLike): Path to BPE tokenizer training data.
 vocab_size (int): Total number of items in the tokenizer's vocabulary (including special tokens).
 special_tokens (list[str]): A list of string special tokens to be added to the tokenizer vocabulary.
 These strings will never be split into multiple tokens, and will always be
 kept as a single token. If these special tokens occur in the `input_path`,
 they are treated as any other string.

 Returns:
 tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
 vocab:
 The trained tokenizer vocabulary, a mapping from int (token ID in the vocabulary)
 to bytes (token bytes)
 merges:
 BPE merges. Each list item is a tuple of bytes (<token1>, <token2>),
 representing that <token1> was merged with <token2>.
 Merges are ordered by order of creation.
 """
 #词表的初始化，初始时词表应该只有从字节串 token 到整数 ID 的一一映射,以及规定的special tokens
 vocab = {i: bytes([i]) for i in range(256)}
 merges = []
 for token in special_tokens:
 b = token.encode("utf-8")
 if b not in vocab.values():
 vocab[len(vocab)] = b
 
 with open(input_path, "rb") as f:
 num_processes = os.cpu_count() or 4
 boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
 ranges = list(zip(boundaries[:-1], boundaries[1:]))
 total_token_counts = Counter()
 with mp.Pool(processes=num_processes) as pool:
 results = pool.map(
 _pretokenize_range,
 [(input_path,s,e) for s,e in ranges],
 )
 for token_counts in results:
 total_token_counts.update(token_counts)
 word_freqs = total_token_counts
 word_symbols = {
 token: tuple(vocab[b] for b in token)
 for token in word_freqs.keys()
 }
 pair2token = defaultdict(set)
 pair_counts = Counter()
 for token,freq in word_freqs.items():
 symbols = word_symbols[token]
 if len(symbols) < 2:
 continue
 for a,b in zip(symbols, symbols[1:]):
 pair_counts[(a,b)] += freq
 pair2token[(a,b)].add(token)
 while len(vocab) < vocab_size:
 if not pair_counts:
 break
 
 #合并出现频次最高的pair,如果有多个pair出现频次相同，则选择字典序最大的那个
 best_pair = max(pair_counts.items(), key=lambda kv: (kv[1], kv[0]))[0]
 #更新vocab和word_symbols
 merges.append(best_pair)
 merged_token = best_pair[0] + best_pair[1]
 vocab[len(vocab)] = merged_token

 tokens_to_update = pair2token[best_pair]

 for token in list(tokens_to_update):
 freq = word_freqs[token]
 symbols = word_symbols[token]
 if len(symbols) < 2:
 continue
 new_symbols = []
 #这里和原来不同了，我们需要减去旧的pair的计数
 i = 0
 while i < len(symbols) - 1:
 p = (symbols[i], symbols[i+1])
 pair_counts[p] -= freq
 pair2token[p].discard(token)
 i += 1
 #然后进行合并
 i = 0
 while i < len(symbols):
 if i < len(symbols) - 1 and (symbols[i], symbols[i+1]) == best_pair:
 new_symbols.append(merged_token)
 i += 2
 else:
 new_symbols.append(symbols[i])
 i += 1
 word_symbols[token] = tuple(new_symbols)
 for i in range(len(new_symbols) - 1):
 p = (new_symbols[i], new_symbols[i+1])
 pair_counts[p] += freq
 pair2token[p].add(token)
 return vocab, merges

至此我们完成了BPE分词的训练，下面我们在具体的数据集上进行训练，得到词表，并将词表存储到磁盘上。

这里选用的就是TinyStories数据集了，另外一个实在太大，懒得弄了

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47


import multiprocessing as mp
import regex as re
import os
import json
import pickle
from pathlib import Path
from typing import BinaryIO
from collections import Counter
from adapters import run_train_bpe
from common import gpt2_bytes_to_unicode

def main():
 filepath = "data/TinyStoriesV2-GPT4-train.txt"
 # filepath = "data/owt_train.txt"
 vocab_size = 10000
 special_tokens = ["<|endoftext|>"]
 vocab, merges = run_train_bpe(filepath, vocab_size, special_tokens)
 print(f"The longest tokens in the vocabulary are:{sorted(vocab.values(), key=len, reverse=True)[:10]}")
 output_dir = Path(__file__).resolve().parent / "outputs"
 output_dir.mkdir(parents=True, exist_ok=True)

 bytes_to_unicode = gpt2_bytes_to_unicode()

 def encode_token(token_bytes: bytes) -> str:
 return "".join(bytes_to_unicode[b] for b in token_bytes)

 vocab_items = sorted(vocab.items(), key=lambda kv: kv[0])
 vocab_json = {encode_token(token_bytes): token_id for token_id, token_bytes in vocab_items}

 vocab_path = output_dir / "trained_vocab.json"
 merges_path = output_dir / "trained_merges.txt"
 serialized_path = output_dir / "trained_vocab_merges.pkl"

 with open(vocab_path, "w", encoding="utf-8") as f:
 json.dump(vocab_json, f, ensure_ascii=False, indent=2)

 with open(merges_path, "w", encoding="utf-8") as f:
 for token_a, token_b in merges:
 f.write(encode_token(token_a))
 f.write(" ")
 f.write(encode_token(token_b))
 f.write("\n")

 with open(serialized_path, "wb") as f:
 pickle.dump({"vocab": vocab, "merges": merges}, f, protocol=pickle.HIGHEST_PROTOCOL)
if __name__ == "__main__":
 main()

BPE Tokenizer：Encoder & Decoder

接下来我们需要实现的就是完整的BPE Tokenizer了，这包含两个部分一个是Encoder，一个是Decoder。

其中Encoder的作用就是使用我们训练好的BPE进行编码的过程，这与训练BPE词表是相对应的，主要包括如下步骤

预分词

首先，我们需要对输入的序列进行预分词，并将每个预分词得到的Token表示为一个UTF-8字节序列。接下来我们会在每个Token内部，将这些字节合并成词表中的元素。

应用合并规则

然后我们按照BPE训练过程中生成合并规则的顺序，将这些词表元素的合并规则依次应用到预分词上。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40


def encode(self,text):
 from .adapters import split_keep_special
 #对单个文本进行BPE编码，返回token id列表
 #首先对文本进行预分词(此时假设text为"Hello <PAD> world!")
 parts = split_keep_special(text, self.special_tokens)
 #此时parts = [""Hello ","<PAD>"," world!"]
 token_ids = []
 for part in parts:
 if part in self.special_tokens:
 token_ids.append(self.vocab_inv[part.encode("utf-8")])
 continue
 for m in re.finditer(self.PAT, part):
 token = m.group(0).encode("utf-8")
 #此时Token为b"Hello"
 symbols = [bytes([b]) for b in token]
 #此时symbols为[b"H",b"e",b"l",b"l",b"o"]

 #接下来进行BPE合并
 while True:
 pairs = [(symbols[i], symbols[i+1]) for i in range(len(symbols)-1)]
 if not pairs:
 break
 candidate_pairs = [p for p in pairs if p in self.merges]
 if not candidate_pairs:
 break
 merged_token = min(candidate_pairs, key=lambda p: self.merges[p])
 new_symbols = []
 i = 0
 while i < len(symbols):
 if i < len(symbols) - 1 and (symbols[i], symbols[i+1]) == merged_token:
 new_symbols.append(merged_token[0] + merged_token[1])
 i += 2
 else:
 new_symbols.append(symbols[i])
 i += 1
 symbols = new_symbols
 #将合并后的symbols转换为token ids
 for sym in symbols:
 token_ids.append(self.vocab_inv[sym])
 return token_ids

此时有一个问题，事实上是要求encode的内存开销是小于1MB的，我们可以发现，主要的开销是在于对于每个Token的处理，我们都会新建pairs以及symbols。事实上我们的文本中有大量的重复token出现，因此可以考虑使用采用LRU的Cache机制。

由于我们通常需要encode的文本很长，我们做不到一次性将所有的文本加载到内存中，因此有时我们需要流式处理：

1
2
3
4
5
6


def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
 #对输入的流式文本的指定范围进行BPE编码，返回token id生成器
 for text in iterable:
 token_ids = self.encode(text)
 for tid in token_ids:
 yield tid

另外就是Decoder了，其作用就是将一串整数形式的token ID解码回原始文本，我们只需要查找每个ID在词表中对应的条目，将这些字节序列依次拼接起来，然后再将得到的字节序列解码为一个Unicode字符串即可。

另外，需要注意的是，输入的TokenID并不保证一定能够映射成合法的 Unicode 字符串。若输入的tokenID不能生成有效的Unicode字符串，那么我们还需要将格式错误的字节替换为官方的Unicode替换字符U+FFFD(按照指导手册的方法，我们使用 errors="replace" 即可)

1
2
3
4
5


def decode(self,token_ids):
 #将token id列表解码为文本,很简单，遍历一遍就行
 bytes_list = [self.vocab[tid] for tid in token_ids]
 text = b"".join(bytes_list).decode("utf-8",errors="replace")
 return text

Transformer模块的构建

接下来我们会具体构建一个Transformer语言模型。

语言模型以一批(batch)整数形式的token ID序列作为输入即形如(batch_size,sequence_length)的pytorch Tensor。其中，对于每一个输入的Token，模型都会预测其下一个词的概率分布。

在训练语言模型时，我们使用这些下一个词的预测结果，来计算真实下一个词与预测下一个词之间的交叉熵损失（cross-entropy loss）。

在推理阶段从语言模型生成文本时，我们取最后一个时间步（即序列中的最后一个位置）得到的下一个词概率分布，用它来生成序列中的下一个 token（例如，选择概率最大的 token、从分布中进行采样等），然后将生成的 token 加入到输入序列中，并重复这一过程。

模型架构介绍

下图为语言模型的架构图:

具体而言，给定一段token ID序列，Transformer语言模型首先使用输入嵌入(红色块，Input Embedding)将Token ID转化为稠密向量，然后将这些Embedding后的token依次送入 num_layers 个Transformer模块，最后通过一个可学习的线性投影层（称为"Output Embedding"或"LM Head"）来产生对下一个Token的预测Logits。

Token Embedding

在最开始的一步中，Transformer会批量地将token ID序列嵌入为一系列向量，这些向量包含了关于Token身份的信息。

更具体地说，给定一个 token ID 序列，Transformer 语言模型使用一个 token embedding 层 来生成一系列向量。该嵌入层接收一个形状为(batch_size, sequence_length) 的整数张量作为输入，并输出一个形状为 (batch_size, sequence_length, d_model) 的向量序列。

我们为什么要做这么一个Embedding？

经过 BPE Tokenizer 后，我们得到的是离散的 token ID。这些 ID 只是符号编号，本身不具备任何语义或数学结构，因此无法直接用于衡量 token 之间的相似性或进行连续建模。

为此，我们将离散的 token 映射到一个连续的高维向量空间（Embedding space），使模型可以通过向量运算来学习和表达语义关系。

在训练完成后，这样的向量空间通常会呈现出良好的语义结构。例如，在该空间中，“公主”和“女人”对应的向量在方向上更为接近，而与“小狗”的向量差异较大。

在实际应用中，我们常使用余弦相似度来衡量这种向量间的语义相似性，其定义为：
$$\cos(\theta) = \frac{A\cdot B}{||A||||B||}$$

Pre-Norm Transformer Block

为什么是Pre-Norm 而非原始论文的post-Norm？

这是一个比较经验主义的结论。大家发现使用Pre-Norm之后训练的梯度更加稳定。

一个比较合理的解释在于：

在Post-Norm结构中，梯度在反向传播时必须经过LayerNorm和子层变换，这会削弱残差连接为梯度提供的直通路径，从而在深层网络中引发梯度消失或梯度不稳定，训练更加困难

相比之下，Pre-Norm将LayerNorm放在子层之前，使残差连接成为一条更加接近恒等映射的路径。这样在反向传播时，梯度可以更直接地通过残差连接传递，从而显著提升训练稳定性，尤其是在深层 Transformer 中。

在完成嵌入之后，激活值会被送入若干个结构完全相同的神经网络层进行处理。一个标准的Decoder-Only Transformer LM由 num_layers 个相同的层组成（通常称为Transformer Block）

每一个Transformer Block都接收一个形状为(batch_size, sequence_length, d_model) 的输入，并输出一个同样形状的张量(batch_size, sequence_length, d_model)。

在每个模块中，模型一方面通过自注意力机制（self-attention）在整个序列范围内聚合信息，另一方面通过前馈网络（feed-forward layers）对这些信息进行非线性变换。

Output Normalization and Embedding

在经过了num_layers 个 Transformer 模块之后，我们将取最终的激活值，并将其转换为在整个词表上的概率分布。

在我们将要实现的Transformer Block中，我们要在最后一个block使用Layer Normalization，以确保其输出具有合适的尺度（Scale）

在完成归一化之后，我们将使用一个标准的可学习线性变换，把 Transformer 模块的输出转换为预测下一个 token 的 logits

编程优化小技巧

在整个Transformer的构建过程中，我们会对许多Batch-like的输入执行相同的操作。下面是一些例子:

batch elements:我们对每个batch元素都应用相同的Transformer前向计算
Sequence length:像RMSNorm和前馈网络这样的“按位置”(position-wise)操作，会对序列中的每一个位置执行完全相同的计算
Attention heads: 注意力操作会在多个注意力头之间以批次处理的方式进行即MHA(Multi-Head Attention)

为了充分利用GPU的并行能力，且让代码易读，我们需要一种高效的方式来执行这些操作。

许多pytorch操作都可以在张量前端接受额外的“类批次”维度，并在这些维度上高效地重复或广播计算。

例如，假设我们要执行一个按位置，按批次的操作。我们有一个形状为(batch_size, sequence_length, d_model) 的“数据张量” D，并希望将其与一个形状为(d_model, d_model) 的矩阵 A 进行批量向量-矩阵乘法。

在这种情况下，直接使用D @ A 就可以完成批量矩阵乘法。这是 PyTorch 中一个高效的基础操作，其中(batch_size, sequence_length) 这两个维度会被自动当作批处理维度。

正因如此，在编写函数的时候，假设输入可能包含额外的类批次维度，并将这些维度放在张量形状的最前面是很有帮助的。为了能够让张量能够以这种方式进行批处理，往往需要多次使用view、reshape 和 transpose 来调整形状。但这样做通常比较繁琐，而且代码会变得难以阅读，也不容易直观理解张量的形状变化。

一种更加符合人类直观理解的方式是选择使用torch.einsum 中的 einsum 记号，或者使用与框架无关的库，如 einops 或 einx。其中两个关键操作是：

einsum：用于在任意维度的输入张量之间执行张量收缩
rearrange：用于对张量维度进行重新排列，拼接或者拆分

下面我们通过一些具体的例子来进行学习。

Example 1

1
2
3
4
5
6
7
8


import torch
from einops import rearrange, einsum
Y = D @ A.T
#很难看出输入输出的张量形状以及具体的含义
#若我们采用einsum，就很直观了
Y = einsum(D,A,"batch sequence d_in,d_out d_in -> batch sequence d_out")
#我们还有一个更加简便的例子:
Y = einsum(D,A,"... d_in,d_out d_in -> ... d_out")

这里通过einsum我们清楚地表明了每个维度的语义，即说明了张量的结构，也说明了输出张量的结构。

Example 2

假设我们有一批图像，并且希望为每一张图像生成 10 个不同“变暗”程度的版本，这些变化由一个缩放因子控制：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


images = torch.randn(64,128,128,3)#(batch,height,width,channel)
dim_dy = torch.linspace(start = 0.0,end = 1.0,steps = 10)
## 通过reshape和逐元素相乘实现
dim_value = rearrange(dim_dy,"dim_value -> 1 dim_value 1 1")
image_rearr = rearrange(images,"b height width channel -> b 1 height width channel")
dimmed_images = images_rearr * dim_value
## 若我们通过enisum实现:
dimmed_images = enisum(
 images,dim_dy,
 "batch height width channel,dim_value -> batch dim_value height width channel"
)

Example3

假设我们有一批图像，其张量形状为 (batch, height, width, channel)。我们希望对图像中的所有像素进行一次线性变换，但这个变换在每个通道（channel）上是相互独立的。该线性变换由一个矩阵 B 表示，其形状为 (height × width, height × width)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


channels_last = torch.randn(64,32,32,3)#(batch,height,width,channel)
B = torch.randn(32*32,32*32)

#传统实现方法就是通过view + transpose
channels_last_flat = channels_last.view(
 -1,channels_last.size(1) * channels_last.size(2),channels_last.size(3)
)
channels_first_flat_transformed = channels_first_flat @ B.T
channels_last_flat_transformed = channels_first_flat_transformed.transpose(1, 2)
channels_last_transformed = channels_last_flat_transformed.view(*channels_last.shape)

#如果我们用enisum
height = width = 32
channels_last_transformed = einsum(
 channels_last,B,
 "batch h_in w_in channel,(h_out w_out)(h_in w_in) -> batch h_out w_out channel"
)

模型基本块的搭建：线性层和嵌入层

参数的初始化

要有效地训练神经网络，通常需要谨慎地进行模型参数初始化。

Pre-Norm Transformer对初始化异常地robust，但初始化方式仍然会对训练速度和收敛性产生显著影响。

在本任务中，我们采用如下初始化方式:

线性层权重

$N(\mu = 0,\sigma^2 = \frac{2}{d_{in} + d_{out}})$，并截断在区间 $[-3\sigma,3\sigma]$内

嵌入层权重

$N(\mu = 0,\sigma^2 = 1)$，并截断在区间$[-3,3]$内

RMSNorm

初始化为1

我们需要使用torch.nn.init.trunc_normal_ 来对截断正态分布权重进行初始化。

线性层模块

线性层是Transformer以及神经网络中最基本，最核心的构建模块之一。首先，我们需要实现一个自定义的Linear类，它继承自torch.nn.Module，并执行如下线性变换:

$$ y = Wx $$

需要注意的是，我们不包含bias，这与现代大多数大语言模型的设计是一致的，这是出自减少访存的考虑。

这里我们需要设计一个Linear类，其中不包含bias，且使用规定的初始化方法。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


import torch
import math
from einops import rearrange, einsum
class Linear(torch.nn.Module):
 def __init__(self,in_features,out_features,device = None,dtype = None):
 super().__init__()
 self.in_features = in_features
 self.out_features = out_features
 self.W = torch.nn.Parameter(torch.empty((out_features,in_features),device = device,dtype = dtype))
 sigma = math.sqrt(2 / (in_features + out_features))
 torch.nn.init.trunc_normal_(self.W,mean = 0.0,std = sigma,a = -3.0 * sigma,b = 3.0 * sigma)
 def forward(self,x):
 return einsum(x,self.W,"... d_in,d_out d_in->... d_out")

Embedding模块

如前所述，Transformer的第一层就是一个Embedding层，它将整数形式的Token ID映射到维度为 $d_{model}$的向量空间中。我们将实现一个自定义的Embedding类，该类继承自torch.nn.Module

forward 方法应当通过索引（indexing）操作，从一个形状为(vocab_size, d_model) 的嵌入矩阵中，为每一个 token ID 选取对应的嵌入向量。输入的 token ID 是一个torch.LongTensor，其形状为(batch_size, sequence_length)。

同样的，按照指导书的要求实现一个Embedding类即可，forward方式其实我们可以理解为查表，因此直接索引就行。同样需要注意初始化方法！

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


import torch
from einops import rearrange, einsum

class Embedding(torch.nn.Module):
 def __init__(self,num_embeddings,embedding_dim,device = None,dtype = None):
 """
 num_embeddings: int 表示词表大小
 embedding_dim: int 表示每个词向量的维度
 """
 super().__init__()
 self.num_embeddings = num_embeddings
 self.embedding_dim = embedding_dim
 self.weight = torch.nn.Parameter(torch.empty((num_embeddings,embedding_dim),device = device,dtype = dtype))
 torch.nn.init.trunc_normal_(self.weight,mean = 0.0,std = 1.0,a = -3.0,b = 3.0)
 def forward(self,token_ids):
 """
 根据给定的token_ids返回对应的Embedding向量
 """
 return self.weight[token_ids]

Pre-Norm Transformer Block

每个Transformer模块包括两个子层，MHA(多头注意力机制)以及按位置的前馈网络

在最初的Transformer论文中，模型在每一个子层外部都使用了残差连接，并在其后接层归一化（Layer Normalization）。这种结构通常被称为Post-norm Transformer.

然而，已有多项研究发现，将层归一化从每个子层的输出端移动到每个子层的输入端（并在最后一个block之后再额外加一层归一化），可以显著提升Transformer训练的稳定性。

如今，Pre-Norm Transformer已成为语言模型中的标准配置（例如 GPT-3、LLaMA、PaLM 等），因此我们也将实现这一变体。接下来，我们将依次介绍并实现预归一化 Transformer 模块中的各个组成部分。

均方根层归一化(RMSNorm)

最初的Transformer论文中采用Layer Normalization来对激活值进行归一化。在本项目中，我们采用RMSNorm的公式来进行归一化。

为什么使用RMSNorm？

我们先来看看原始的LayerNorm:
$$y = \frac{x - E[x]}{\sqrt{Var[x]+\epsilon}}*\gamma + \beta$$
其中， $\gamma$和 $\beta$ 为可训练的参数

一般而言，这出于两个考虑：

Fewer operation：RMSNorm无需计算均值和方差，减少了算术运算和内存访问

Fewer parameter：去掉了偏置参数，减少了参数量以及通讯开销

虽然在LLM或者说神经网络的训练中，矩阵乘法占据了大部分的计算开销，但是访存开销以及通讯开销同样是不能忽视的存在。

具体原理如下：

给定一个激活向量 $a \in \mathbb{R}^{d_{model}}$,RMSNorm会对每一个激活分量 $a_i$进行如下的缩放:

$$ RMSNorm(a_i) = \frac{a_i}{RMS(a)}g_i $$

其中:

$$ RMS(a) = \sqrt{\frac{1}{d_{model}}\sum_{i = 1}^{d_{model}}a_i^2 + \epsilon} $$

这里g是一个可学习的增益(gain)参数，而 $\epsilon$则是一个用于数值稳定的超参数，通常固定为1e-5。

在计算平方时，为了数值稳定，我们应该将输入的张量上转为float32.

同样实现一个类就行，注意初始化的gain全1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


import torch
from einops import rearrange, einsum
import math
class RMSNorm(torch.nn.Module):
 def __init__(self, d_model, eps = 1e-5,device = None,dtype = None):
 super().__init__()
 self.d_model = d_model
 self.eps = eps
 #依据实验指导书，初始化为全1
 self.weight = torch.nn.Parameter(torch.ones(d_model, device=device, dtype=dtype))
 def forward(self,x):
 #为了数值稳定，先将类型转化为float32
 in_dtype = x.dtype
 x = x.to(torch.float32)
 rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
 x_normed = x / rms
 x_normed = x_normed * self.weight
 return x_normed.to(in_dtype)

Position-Wise Feed-Forward Network

在最初的Transformer论文中，Transformer的前馈神经网络由两个线性变换组成，中间使用ReLU函数。其中前馈神经网络内部隐藏层的维度通常设为输入维度的4倍。

然而，现代语言模型相较于这一原始设计，通常会引入两项主要改动:即使用不同的激活函数以及引入门控机制。

具体而言，我们将实现在目前一些主流大模型中(LLaMA3，Qwen2.5等)的采用的SwiGLU激活函数。SwiGLU激活函数将SiLU激活函数与一种称为门控线性单元的机制结合在一起。

此外，我们将省略线性层中的偏置项。

SiLU/Swish激活函数

SiLU(也叫Swish)激活函数定义如下:

$$ SiLU(x) = x \cdot \sigma(x) = \frac{x}{1 + e^{-x}} $$

如下图所示，SiLU激活函数在形态上类似于ReLU，但在零点处是平滑的。

门控线性单元（GLU）

Gated Linear Units的定义为：

一个线性变换经过sigmoid函数后的结果，与另一个线性变换的结果进行逐元素相乘:

$$ GLU(x,W1,W2) = \sigma(W_1x)⊙ W_2x $$

GLU被认为可以通过为梯度提供一条线性传播路径，在保持非线性表达能力的同时，缓解深层网络中的梯度消失问题。

将二者结合，就得到了SwiGLU前馈网络：

$$ FFN(x) = SwiGLU(x,W_1,W_2,W_3) = W_2(SiLU(W_1x)⊙W_3x) $$

其中

$x \in \mathbb{R}^{d_{model}}$
$W_1,W_3 \in \mathbb{R}^{d_{ff}\times d_{model}}$
$W_2 \in \mathbb{R}^{d_{model}\times d_{ff}}$
$d_{ff} = \frac{8}{3}d_{model}$

我们按照要求直接实现就行，可以直接引入之前我们设计的Linear层，但是我这里还是手写了一遍。

需要注意d_ff肯定得是整数，并且实验手册中也强调了它得是64的整数倍以提升性能；然后我们的SiLU激活函数在实现的时候可以使用torch.sigmoid

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29


import torch
import torch.nn as nn
import math
from einops import rearrange, einsum

class SwiGLU(nn.Module):
 def __init__(self,d_model,device = None,dtype = None):
 super().__init__()
 self.d_model = d_model
 #根据指导手册要求，但是需要保证是64的整数倍
 d_ff = int(math.ceil((8/3) * d_model / 64) * 64)
 
 self.W1 = torch.nn.Parameter(torch.empty((d_ff,d_model),device=device,dtype=dtype))
 self.W2 = torch.nn.Parameter(torch.empty((d_model,d_ff),device=device,dtype=dtype))
 self.W3 = torch.nn.Parameter(torch.empty((d_ff,d_model),device=device,dtype=dtype))

 sigma1 = (2 / (d_model + d_ff)) ** 0.5
 sigma2 = (2 / (d_ff + d_model)) ** 0.5

 torch.nn.init.trunc_normal_(self.W1,mean = 0.0,std = sigma1,a = -3.0 * sigma1,b = 3.0 * sigma1)
 torch.nn.init.trunc_normal_(self.W2,mean = 0.0,std = sigma2,a = -3.0 * sigma2,b = 3.0 * sigma2)
 torch.nn.init.trunc_normal_(self.W3,mean = 0.0,std = sigma1,a = -3.0 * sigma1,b = 3.0 * sigma1)
 def forward(self,x):
 x_proj1 = einsum(x,self.W1,"... d_model,d_ff d_model->... d_ff")
 x_proj1 = x_proj1 * torch.sigmoid(x_proj1) #计算SiLU激活
 x_proj2 = einsum(x,self.W3,"... d_model,d_ff d_model->... d_ff")
 x_glu = x_proj1 * x_proj2
 out = einsum(x_glu,self.W2,"... d_ff,d_model d_ff->... d_model")
 return out

RoPE，旋转位置编码

为了向模型中注入位置信息，我们将实现旋转位置编码。

为什么需要位置编码？旋转位置编码有何优点?

这里的问题比较深入，我将会在自己的博客中从数学的角度进行学习介绍。

具体而言，给定位于位置i的给定Query token(这后面会介绍，我们目前专注于RoPE的实现):

$$ q(i) = W_qx(i)\in \mathbb{R}^{d} $$

我们将应用一个成对的旋转矩阵 $R_i$从而得到

$$ q'(i) = R_i q(i) = R_i W_q x(i) $$

这里， $R_i$会将embedding 向量中的元素成对地旋转（想想我们在二维坐标系中的旋转）：

我们将 $q(i)_{2k-1:2k}$视为二维向量，并按角度:$\theta_{i,k} = \frac{i}{\Theta^{\frac{2k-2}{d}}}$进行旋转，其中 $k\in \{1,\dots,d/2\},\Theta$为某个常数。

因此我们可以将 $R_i$看作一个大小为 $d \times d$的块对角矩阵，其中第k个块为 $R_i^k$,其中:

$$ R_i^k = \left[\begin{matrix}\cos(\theta_{i,k}) & -\sin(\theta_{i,k})\\ \sin(\theta_{i,k})& \cos(\theta_{i,k})\end{matrix}\right ] $$

于是完整的旋转矩阵为:

$$ R_i = \left[\begin{matrix}R_i^1 & 0 & 0 & \dots & 0\\0 & R_i^2 & 0 &\dots & 0\\0 & 0 & R_i^3 &\dots & 0\\\vdots & \vdots & \vdots & \ddots & \vdots\\0 & 0 & 0 &\dots & R_i^{d/2}\end{matrix} \right] $$

虽然我们可以显式构造完整的 d×d 矩阵，但一个更好的实现应当利用该矩阵的结构性质，以更高效的方式完成变换。由于我们仅仅关心同一序列内token的相对旋转关系，因此可以在不同层，不同batch之间复用已经计算好的 $\cos(\theta_{i,k}),\sin(\theta_{i,k})$值。

具体而言，我们可以实现一个被所有层共享的RoPE模块，并在函数初始化时通过self.register_buffer(persistent=False) 预先创建一个大小为 2d 的 sin 和 cos 值缓存，而不是使用 nn.Parameter（因为我们不希望学习这些固定的正弦和余弦值）

代码实现如下: 简单来说，对于每个输入x (…,seq_len,dim)，实际上每个位置的角度都是固定的，因此我们在初始化的时候就把每个位置的角度以及对应三角函数计算出来就行。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54


from einops import rearrange, einsum
import torch
import torch.nn as nn

class RotaryPositionalEmbedding(nn.Module):
 def __init__(self,theta,d_k,max_seq_len,device):
 """
 d_k: int, 维度大小，必须为偶数
 theta: float, RoPE中的\Theta值
 max_seq_len: int, 最大序列长度
 device: torch.device, 设备
 """
 super().__init__()
 assert d_k % 2 == 0, "d_k must be even"
 self.theta = theta
 self.d_k = d_k
 self.max_seq_len = max_seq_len
 self.device = device

 #一共有d / 2个频率
 half_dk = d_k // 2
 k = torch.arange(0,half_dk,device=device).float()
 inv_freq = 1.0 / (self.theta ** (2.0 * k / d_k))

 positions = torch.arange(0,max_seq_len,device = device).float()

 angles = einsum(positions,inv_freq,"max_seq_len,half_dk->max_seq_len half_dk")
 cos = torch.cos(angles)
 sin = torch.sin(angles)

 self.register_buffer("cos",cos,persistent = False)
 self.register_buffer("sin",sin,persistent = False)

 def forward(self,x,token_positions):
 """
 inputs:
 x: ...,seq_len,d_k
 token_positions:...,seq_len
 returns:
 x_rotated: ...,seq_len,d_k
 """
 cos = self.cos[token_positions] # ...,seq_len,half_dk
 sin = self.sin[token_positions] # ...,seq_len,half_dk

 x_even = x[...,0::2]
 x_odd = x[...,1::2]

 x_rot_even = x_even * cos - x_odd * sin
 x_rot_odd = x_even * sin + x_odd * cos

 out = torch.empty_like(x)
 out[...,0::2] = x_rot_even
 out[...,1::2] = x_rot_odd
 return out

Scaled dot-product Attention

我们接下来将实现缩放点积注意力，也就是Transformer原始论文中的Attention机制。

在此之前，我们需要实现Softmax，这是一种将未归一化的分数向量转化为归一化分布的操作。

$$ \text{Softmax}(v_i) = \frac{\exp (v_i)}{\sum_j \exp(v_j)} $$

虽然这个看似简单，但是我们需要特别注意数值稳定的问题，因为指数求和一般会很大很大。因此我们可以通过注意到softmax操作对所有输入上加上任意常数c是不变的来避免这个问题。

通常的做法是从向量 $o_i$的所有元素中减去其中最大的那个值，使其新的最大值为0.

具体实现没有什么额外需要说明的，直接应用就行:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18



def run_softmax(in_features: Float[Tensor, " ..."], dim: int) -> Float[Tensor, " ..."]:
 """
 Given a tensor of inputs, return the output of softmaxing the given `dim`
 of the input.

 Args:
 in_features (Float[Tensor, "..."]): Input features to softmax. Shape is arbitrary.
 dim (int): Dimension of the `in_features` to apply softmax to.

 Returns:
 Float[Tensor, "..."]: Tensor of with the same shape as `in_features` with the output of
 softmax normalizing the specified `dim`.
 """
 max_num = torch.max(in_features,dim=dim,keepdim=True).values
 exp_tensor = torch.exp(in_features - max_num)
 sum_exp = torch.sum(exp_tensor,dim=dim,keepdim=True)
 return exp_tensor / sum_exp

我们接下来可以进行Attention的实现，在数学上将Attention操作定义如下:

$$ \text{Attention}(Q,K,V) = \text{softmax}(\frac{Q^\top K}{\sqrt{d_k}})V $$

其中 $Q\in \mathbb{R}^{n\times d_k},K\in \mathbb{R}^{m\times d_k},V\in \mathbb{R}^{m\times d_v}$.这些都是该操作的输入。

有时我们需要对注意力操作的输出进行掩码。掩码应具有形状 $M\in \{\text{True},\text{False}\}^{n\times m}$,这是一个布尔矩阵，其中第i行表示第i个查询可以关注哪些键。

按照惯例，在位置(i,j)上取值为True表示查询i可以关注键j，而取值为False表示不能。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


def run_scaled_dot_product_attention(
 Q: Float[Tensor, " ... queries d_k"],
 K: Float[Tensor, " ... keys d_k"],
 V: Float[Tensor, " ... values d_v"],
 mask: Bool[Tensor, " ... queries keys"] | None = None,
) -> Float[Tensor, " ... queries d_v"]:
 """
 Given key (K), query (Q), and value (V) tensors, return
 the output of your scaled dot product attention implementation.

 Args:
 Q (Float[Tensor, " ... queries d_k"]): Query tensor
 K (Float[Tensor, " ... keys d_k"]): Key tensor
 V (Float[Tensor, " ... values d_v"]): Values tensor
 mask (Bool[Tensor, " ... queries keys"] | None): Mask tensor
 Returns:
 Float[Tensor, " ... queries d_v"]: Output of SDPA
 """
 d_k = Q.shape[-1]
 scores = einsum(Q,K,"... q d_k,... k d_k -> ... q k") / (d_k ** 0.5)
 if mask is not None:
 #对于mask的，我们直接将mask为False的位置设置为负无穷
 scores = scores.masked_fill(~mask,float("-inf"))
 attn_weights = run_softmax(scores,dim=-1)
 return einsum(attn_weights,V,"... q k,... k d_v -> ... q d_v")

Causal Multi-Head Self-Attention

接下来我们将按照Transformer原始论文中的描述来实现多头注意力机制。

具体而言：

$$ \text{MultiHead}(Q,K,V) = \text{Concat}(head_1,head_2,\dots,head_n) $$

其中:

$$ \text{head}_i = \text{attention}(Q_i,K_i,V_i) $$

基于此，我们可以得到多头注意力操作的形式:

$$ \text{MultiHeadSelfAttention(x)} = W_O \text{MultiHead}(W_Qx,W_Kx,W_Vx) $$

其中W均为可学习的参数。一般而言，这里得到Q，K，V需要3次矩阵乘法，但是我们可以尝试将key,query和value的投影合并到一个单一的权重矩阵中，从而只需要一次矩阵乘法。

此外，我们还需要实现因果掩码(Causal Masking).

其目的在于防止模型关注到序列中未来的Token。换言之，如果给定模型一个token序列 $t_1,t_2,\dots,t_n$而我们希望为前缀 $t_1,\dots,t_i$计算下一个词的预测，那么模型不应该访问的位置就是 $t_{i+1}\dots t_n$.

因为在推理阶段生成文本时，模型无法获取这些未来的Token，而这些Token会泄露关于真实下一个词的信息，从而使语言建模的预训练目标变得平凡。

事实上，我们可以通过对序列中每个不同的前缀分别运行一次多头注意力，从而防止访问未来token。但是这样效率太低，我们使用因果注意力掩码，它允许第i个token关注序列中所有满足 $j \leq i$的位置。

在实现上，我们可以通过torch.triu或基于广播的索引比较来构造这个掩码，并且在上面的Attention中我们已经支持了掩码。

此外，这里还需要应用我们先前实现的RoPE（针对Q，K）。此外，head维度应当被视为一个类batch维度进行处理，因为在MHA中，每个head的计算是相互独立的。

在这个的实现中，我们需要注意以下四个矩阵的维度。以及RoPE是只针对Q和K使用的就行。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52


class CausalMultiHeadSelfAttention(nn.Module):
 def __init__(self,d_model,num_heads,seq_len,device,use_rope=True):
 """
 d_model: int,模型维度
 num_heads: int,注意力头数
 device: torch.device,设备
 use_rope: bool,是否使用RoPE
 """
 super().__init__()
 assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
 self.d_model = d_model
 self.num_heads = num_heads
 self.d_k = d_model // num_heads
 self.device = device
 self.theta = 10000.0

 self.W_q = Linear(self.d_k,self.d_model,self.device)
 self.W_k = Linear(self.d_k,self.d_model,self.device)
 self.W_v = Linear(self.d_k,self.d_model,self.device)
 self.W_o = Linear(self.d_model,self.d_k,self.device)
 if use_rope:
 self.use_rope = True
 self.RoPE = RotaryPositionalEmbedding(theta=self.theta,d_k = self.d_k, max_seq_len = seq_len, device = device)
 else:
 self.use_rope = False

 def forward(self,x,token_positions):
 """
 inputs:
 x: Float[Tensor, "batch_size seq_len d_model"]
 token_positions: Long[Tensor, "batch_size seq_len"]
 returns:
 out: Float[Tensor, "batch_size seq_len d_model"]
 """
 batch_size,seq_len,_ = x.shape

 Q = self.W_q(x)
 K = self.W_k(x)
 V = self.W_v(x)
 
 Q = rearrange(Q,"b s (h d_k) -> b h s d_k",h=self.num_heads)
 K = rearrange(K,"b s (h d_k) -> b h s d_k",h=self.num_heads)
 V = rearrange(V,"b s (h d_k) -> b h s d_k",h=self.num_heads)
 if self.use_rope:
 Q = self.RoPE(Q,token_positions)
 K = self.RoPE(K,token_positions)
 
 mask = torch.tril(torch.ones((seq_len,seq_len),device=self.device)).bool()
 attn_output = run_scaled_dot_product_attention(Q,K,V,mask=mask)
 attn_output = rearrange(attn_output,"b h s d_k -> b s (h d_k)")
 out = self.W_o(attn_output)
 return out

Transformer Block

接下来组装Transformer Block，为了方便implement，我们看这张图:

如你所见，一个Transformer block包含了两个部分，一个用于多头注意力，另一个则用于前馈网络。

在每一个部分之前都会先执行RMSNorm，然后是主要运算，最后加上残差连接。

实现方面，我们之前已经把积木准备好了，只剩下积木的拼接啦！按照这个图进行拼接就可以咯！

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


import torch.nn as nn
from cs336_basics.CausalMultiHeadSelfAttention import CausalMultiHeadSelfAttention
from cs336_basics.Linear import Linear
from cs336_basics.RMSNorm import RMSNorm
from cs336_basics.SwiGLU import SwiGLU

class TransformerBlock(nn.Module):
 def __init__(self,d_model,num_heads,seq_len,ffn_hidden_dim,device,use_rope=True):
 super().__init__()
 self.attention = CausalMultiHeadSelfAttention(d_model,num_heads,seq_len,device,use_rope)
 self.ffn = SwiGLU(d_model,device)
 # pass device explicitly to avoid treating it as eps
 self.attn_norm = RMSNorm(d_model, device=device)
 self.ffn_norm = RMSNorm(d_model, device=device)
 def forward(self,x,token_positions):
 x_norm = self.attn_norm(x)
 attn_out = self.attention(x_norm,token_positions)
 x = x + attn_out
 x_norm = self.ffn_norm(x)
 ffn_out = self.ffn(x_norm)
 x = x + ffn_out
 return x

The Full Transformer LM！

最后我们可以把实现的Transformer Block进行组装了!

具体而言我们参考下图:

正常实现就行啦！

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32


from cs336_basics.Transformer_Block import TransformerBlock
from cs336_basics.Embedding import Embedding
from cs336_basics.RMSNorm import RMSNorm
from cs336_basics.Linear import Linear
import torch.nn as nn
import torch
from einops import rearrange, einsum

class Transformer:
 def __init__(self,vocab_size,context_length,num_layers,d_model,num_heads,device):
 self.TokenEmbedding = Embedding(vocab_size,d_model,device)
 self.TransformerBlocks = nn.ModuleList([
 TransformerBlock(d_model,num_heads,context_length,ffn_hidden_dim=None,device=device,use_rope=True)
 for _ in range(num_layers)
 ])
 self.FinalNorm = RMSNorm(d_model,device=device)
 self.OutputLayer = Linear(d_model,vocab_size,device)
 def forward(self,token_ids):
 """
 inputs:
 token_ids: Long[Tensor, "batch_size seq_len"]
 returns:
 logits: Float[Tensor, "batch_size seq_len vocab_size"]
 """
 batch_size,seq_len = token_ids.shape
 x = self.TokenEmbedding(token_ids) # [batch_size, seq_len, d_model]
 token_positions = torch.arange(seq_len,device=x.device).unsqueeze(0).expand(batch_size,-1) # [batch_size, seq_len]
 for block in self.TransformerBlocks:
 x = block(x,token_positions)
 x = self.FinalNorm(x)
 logits = self.OutputLayer(x)
 return logits

Transformer LM的训练

我们已经完成了对数据(Tokenizer)和模型(Transformer)进行预处理的步骤。剩下的工作就是编写所有支持训练的代码，主要包括以下几个部分:

Loss(损失函数):交叉熵
Optimizer(优化器):用于最小化该损失函数的优化器 AdamW
Training loop(训练循环):我们需要所有支撑训练的基础设施，包括数据的加载，保存checkpoint以及管理训练过程。

交叉熵损失(Cross-entropy loss)

在先前的Pipeline介绍中，我们知道LM会对每个长度为m+1的序列x，以及每一个 $i = 1,\dots,m$,定义分布:

$p_{\theta}(x_{i+1}|x_{1:i})$.

给定一个训练集D，其中包含长度为m的序列，我们定义标准的交叉熵损失函数:

$$ l(\theta;D) = \frac{1}{|D|m}\sum_{x\in D}\sum_{i=1}^m -\log p_{\theta}(x_{i+1}|x_{1:i}) $$

（需要注意的是，Transformer每一次前向就能同时得到所有 $i=1,\dots,m$的 $p_\theta(x_{i+1}|x_{1:i})$）

具体而言，Transformer会对每个位置i计算logits: $o_i\in \mathbb{R}^{vocab\_size}$

从而得到:

$$ p(x_{i+1}|x_{1:i}) = \text{softmax}(o_i)[x_{i+1}] = \frac{\exp(o[x_{i+1}])}{\sum_{a=1}^{vocab\_size}\exp(o_i[a])} $$

在交叉熵的实现中，与softmax一样也需要注意数值稳定的问题。

这里会出现两种数值稳定问题：

上溢，也就是之前softmax所需要解决的，通过减去max就行
下溢，如果logits很小，那么log操作后就会出现下溢，这里则需要通过log_sum_exp来解决

具体而言，在具体实现中，我们可以拆分为两部分来计算交叉熵:

第一部分是分子 $\log \exp(o[x_{i+1}]) = o[x_{i+1}]$

第二部分是分母 $\sum_{a=1}^{vocab\_size}\exp(o_i[a])$,在计算求和的时候我们需要类似于softmax一样处理，即减去最大logits:

$$ \log (\sum_{a=1}^{vocab\_size}\exp(o_i[a]-\max)) + \max $$

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


import torch
from einops import rearrange, einsum

def cross_entropy(logits, targets):
 """
 logits: Float[Tensor, "batch vocab_size"]
 targets: Int[Tensor, "batch"]
 """
 max_logits = logits.max(dim = -1,keepdim = True).values #shape: [batch_size,1]
 #如果直接对logits减去最大值，最后会因为log操作出现下溢的情况
 log_sum_exp = torch.log(torch.exp(logits - max_logits).sum(dim=-1,keepdim=True)) + max_logits # [batch,1]
 # [batch,1] -> [batch]
 log_sum_exp = rearrange(log_sum_exp,'b 1 -> b')
 loss = log_sum_exp - logits[torch.arange(logits.shape[0]),targets]
 return loss.mean()

优化器(SGD,AdamW)

我们已经设计好了损失函数，接下来就要实现优化器。最简单的基于梯度的优化器是随机梯度下降。我们从随机初始化的参数 $\theta_0$开始。随后，对于每一个步长 $t = 0,\dots,T-1$执行如下更新:

$$ \theta_{t+1} =\theta_t - \alpha_t \nabla L(\theta_t;B_t) $$

其中 $\alpha_t$为学习率， $B_t$是从数据集D中随机采样的批次数据。批次大小和学习率是超参数

在本项目中，我们不实现SGD，而是实现在现代LM中更加常用且更加复杂的优化器。

近期使用的大多数优化器都是Adam优化器的变体。我们将使用AdamW，在近期的工作中被广泛采用。AdamW 对 Adam 提出了一种改进，通过以一种与梯度更新解耦的方式添加权重衰减（在每次迭代中，将参数向 0 推拉）来增强正则化效果。

AdamW是有状态的：对于每个参数，它都会跟踪其一阶矩和二阶矩的运行估计。因此，AdamW使用额外的内存来换取更好的稳定性和收敛性。除了学习率外，AdamW 还有一对控制矩估计更新的超参数 $\beta_1,\beta_2$，以及一个权重衰减率 $\lambda$。典型的应用将 $\beta_1, \beta_2$ 设置为 (0.9, 0.999)，但像 LLaMA 和 GPT-3 这样的大语言模型通常使用 (0.9, 0.95) 进行训练。算法如下所示，其中 $\epsilon$ 是一个极小值（例如 $10^{-8}$），用于在 v 出现极小值时提高数值稳定性：

我们只需要按照上述算法流程，按照SGD章节提供的框架实现即可:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42


import torch
from einops import rearrange, einsum
import math

class Adamw(torch.optim.Optimizer):
 def __init__(self,params,lr=1e-3,betas=(0.9,0.999),eps=1e-8,weight_decay=0.01):
 defaults = dict(lr=lr,betas=betas,eps=eps,weight_decay=weight_decay)
 super().__init__(params,defaults)
 def step(self,closure=None):
 loss = None
 for group in self.param_groups:
 lr = group['lr']
 betas = group['betas']
 eps = group['eps']
 weight_decay = group['weight_decay']
 
 # step 放在 group 级别，所有参数共享
 step = group.get('step', 0) + 1
 group['step'] = step
 
 # 计算偏差修正后的学习率
 lr_t = lr * math.sqrt(1 - betas[1] ** step) / (1 - betas[0] ** step)
 
 for p in group['params']:
 if p.grad is None:
 continue
 state = self.state[p] # 只包含一阶矩、二阶矩
 m = state.get('m', torch.zeros_like(p))
 v = state.get('v', torch.zeros_like(p))
 
 # 更新一阶矩、二阶矩
 m = betas[0] * m + (1 - betas[0]) * p.grad
 v = betas[1] * v + (1 - betas[1]) * p.grad * p.grad

 # 更新参数
 p.data = p.data - lr_t * (m / (torch.sqrt(v) + eps))
 # 参数衰减
 p.data = p.data * (1 - weight_decay * lr)
 
 state['m'] = m
 state['v'] = v
 return loss

学习率调度(learning rate scheduling)

在训练过程中，能够导致损失函数下降最快的学习率通常是不断变化的。在训练Transformer模型时，通常会使用学习率调度策略：初期使用较大的学习率以实现快速更新，并随着模型的训练将其缓慢衰减至较小值

在本项目中，我们将实现用于训练LLaMA的余弦退火调度。

调度器本质上是一个函数，它接收当前步数和其他相关参数，并返回第t步执行梯度更新时应使用的学习率。最简单的调度策略是常数函数。

余弦退火调度接受以下参数:（i）当前迭代步数t,（ii）最大学习率 $\alpha_{\max}$ （iii）最小学习率 $\alpha_{\min}$(iv)预热迭代次数 $T_w$(v)余弦退火迭代次数 $T_c$

第t次迭代的学习率定义如下:

若 $t < T_w$,则: $\alpha_t = \frac{t}{T_w}\alpha_{\max}$
若 $T_w \leq t \leq T_c$,则： $\alpha_t = \alpha_{\min} + \frac{1}{2}(\alpha_{\max} - \alpha_{\min})(1 + \cos(\pi \frac{t - T_w}{T_c - T_w}))$
若 $t \geq T_c$,则 $\alpha_t = \alpha_{\min}$

按部就班实现就行，没有坑点

1
2
3
4
5
6
7
8
9


import math

def lr_cosine_schedule(t,a_max,a_min,t_w,t_c):
 if t < t_w:
 return t / t_w * a_max
 elif t <= t_c:
 return a_min + (a_max - a_min) * (1 + math.cos(math.pi * (t - t_w) / (t_c - t_w))) / 2
 else:
 return a_min

梯度裁剪(Gradient clipping)

在训练过程中，我们有时会遇到产生极大梯度的训练样本，这可能会导致训练过程变得不稳定。为了缓解这个问题，实践中通常采用一种技术是梯度裁剪。其核心思想在每次反向传播结束后，执行优化器步之前，对梯度的范数设定一个上限。

具体而言，给定所有参数的梯度g，我们计算其l2范数 $||g||_2$（所有参数）.若该范数小于最大值M，则保持g不变；否则我们将g按比例缩小，其中缩放因子为 $\frac{M}{||g||_2+\epsilon}$。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


import math
import torch

def gradient_clip(params, max_norm, epsilon=1e-6):
 total_norm = torch.sqrt(sum([torch.norm(p.grad.data, p=2) ** 2 for p in params if p.grad is not None]))
 scale = max_norm / (total_norm + epsilon)
 if scale < 1:
 for p in params:
 if p.grad is not None:
 p.grad.data *= scale

Training Loop(data_utils)

我们现在需要搭建整个模型的训练Pipeline，这需要把我们前面搭建的配件整合在一起

DataLoader

标记后的数据是一个单一的标记序列 $x=(x_1,x_2,\dots,x_n)$.尽管原始的数据可能由不同的文档组成，通用的做法是将它们全部连接成一个单一的标记序列，并在它们之间添加分割符。

DataLoader的作用是将此序列转化为批次流，其中每个批次包含B个长度为m的序列，并配对相应的长度为m的下一个标记作为目标。例如，当B = 1, m = 3 时， $([x_2, x_3, x_4], [x_3, x_4, x_5])$ 就是一个可能的批次

以这种方式加载数据可以简化训练，原因如下：首先，任何满足 $1 \le i < n - m$ 的 $i$ 都能产生一个有效的训练序列，因此采样过程变得非常简单。其次，由于所有训练序列长度相同，无需对输入序列进行填充（padding），这提高了硬件利用率。最后，我们不需要为了采样而将整个数据集完整加载到内存中，这使得处理无法放入内存的大规模数据集变得容易。

这里我们通过torch提供的两个api实现，一个是randint，它能让我们生成若干个随机数；另一个是stack，它能将多个形状相同的张量沿着新的维度堆叠起来，例如我们生成了batch_size个随机开头，并截取了batch_size个序列，那么将其用stack堆叠起来就得到了我们需要的形为[batch_size,context_len]的张量

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23


from torch.utils.data import Dataset
import numpy as np
import torch

def get_batch(x, batch_size, context_length, device):
 """
 input:
 x: Int[Tensor, "seq_len"] 或 numpy array
 batch_size: int
 context_length: int
 device: torch.device
 output:
 xb: Int[Tensor, "batch_size context_length"]
 yb: Int[Tensor, "batch_size context_length"]
 """
 if not isinstance(x, torch.Tensor):
 x = torch.tensor(x)
 
 idx = torch.randint(0, len(x) - context_length, size=(batch_size,))
 xb = torch.stack([x[i:i+context_length] for i in idx])
 yb = torch.stack([x[i+1:i+context_length+1] for i in idx])
 xb, yb = xb.to(device), yb.to(device)
 return xb, yb

Checkpoints

除了加载数据，我们还需要在训练过程中保存模型。在运行作业时，我们经常希望能够恢复由于某种原因中途停止的训练任务（例如，由于作业超时、机器故障等）。即使一切顺利，我们稍后也可能希望访问中间模型（例如，事后研究训练动态、从不同训练阶段的模型中提取样本等）。

一个检查点（Checkpoint）应该包含恢复训练所需的所有状态。我们至少需要能够恢复模型权重。如果使用有状态的优化器（如 AdamW），我们还需要保存优化器的状态（例如 AdamW 的矩估计值）。最后，为了恢复学习率调度，我们需要知道停止时的迭代次数。

PyTorch 使得保存这些内容变得非常简单：每个 nn.Module 都有一个 state_dict() 方法，返回一个包含所有可学习权重的字典；稍后我们可以通过其姊妹方法 load_state_dict() 来恢复这些权重。对于任何 nn.optim.Optimizer 也是如此。最后，torch.save(obj, dest) 可以将一个对象（例如，一个在某些值中包含张量的字典，也可以是像整数这样的普通 Python 对象）转储到文件（路径）或类文件对象中，随后可以通过 torch.load(src) 将其重新加载到内存中。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


def save_checkpoint(model,optimizer,epoch, out):
 torch.save({
 "epoch": epoch,
 "model_state_dict": model.state_dict(),
 "optimizer_state_dict": optimizer.state_dict(),
 }, out)

def load_checkpoint(src,model,optimizer):
 checkpoint = torch.load(src)
 model.load_state_dict(checkpoint["model_state_dict"])
 optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
 return checkpoint["epoch"]

Inference

现在我们来完成最终的推理，也就是生成文本。

LM接收一个长度为 sequence_length 的整数序列,并产生一个大小为 (sequence_length × vocab size) 的矩阵。在这个矩阵中，序列中的每个元素都是一个概率分布，用于预测该位置后的下一个词。现在，我们将编写一些函数，将其转化为新序列的采样方案。

按照标准惯例，语言模型的输出是最后一个线性层的输出（即 “logits”）。因此，我们必须通过 Softmax 操作将其转化为归一化概率，这在之前的公式中已经出现过。

为了从模型中生成文本，我们需要为模型提供一个前缀Token(即提示词Prompt)，并要求它生成一个词汇表上的概率分布，以此预测序列中的下一个词。然后，我们从这个分布中采样，以确定下一个输出的 Token。

具体而言，解码过程的一个步骤应该是输入一个序列 $x_{1\dots t}$,并通过如下方程返回一个 $Token\ x_{t+1}$：

$$ P(x_{t+1}=i|x_{1\dots t}) = \frac{\exp (v_i)}{\sum_{j}\exp (v_j)}\\v = \text{TransformerLM}(x_{1\dots t})_t \in \mathbb{R}^{\text{vocab\_size}} $$

其中TransformerLM是我们的模型，其输入为长度为 sequence_length 的整数序列，输出为大小为 (sequence_length × vocab size) 的矩阵。我们取该矩阵的最后一行元素，因为我们正在寻找的就是第t个位置后的下一个词预测。

通过反复从这些单步条件分布中采样(即将前一步生成的输出Token附加到下一步解码的输入中)，直到生成序列结束标记,我们就能得到一个基础的解码器

我们实验中使用的是小型模型，而小模型有时会生成质量非常低的文本。两个简单的解码技巧可以帮助解决这些问题：

温度缩放 (Temperature Scaling)：我们引入温度参数来修正 Softmax，新的 Softmax 公式为：

$\text{softmax}(v,\tau) = \frac{\exp (v_i / \tau)}{\sum_j \exp (v_j / \tau)}$

需要注意的是，当 $\tau \to 0$ 时，向量 v 中最大的元素将占据主导地位，Softmax 的输出将变成一个集中在该最大元素上的独热向量 (one-hot vector)。

Nucleus sampling 或 Top-p采样:另一种技巧则是通过截断低概率词来修正采样分布。假设q是经过温度放缩后的Softmax得到的大小为vocab_size的概率分布。具有超参数p的Nucleus sampling按照下式产生下一个Token:

$$ P(x_{t+1}=i|q) = \left\{\begin{aligned}\frac{q_i}{\sum_{j\in V(p)} q_j} &if \quad i\in V(p)\\0 \quad &else\end{aligned}\right . $$

其中V(p)是满足 $\sum_{j\in V(p)}q_j \geq p$的最小索引集合。可以通过先按大小对概率分布 q 进行排序，然后依次选择最大的词汇元素直到达到目标水平 p 来轻松计算此值。

这里就不贴具体的代码了，训练推理代码在仓库中都有~