TL;DR 本讲将分布式通信概念落地为可执行代码，通过 PyTorch 和 NCCL 展示了集合通信的实际实现，并手写了三种并行策略的简化版本。从 All-Reduce 的基准测试到 MLP 的切分实现，揭示了通信开销与计算模式的核心逻辑。

通信原语：从接口到底层

集合通信在 PyTorch 中的接口简洁直白。初始化进程组后，每个进程执行相同的通信操作，NCCL 自动处理底层数据流动。

1
2
3
4
5
6
7
8
9

# 初始化进程组（每个进程都要执行）
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "15623"
dist.init_process_group("nccl", rank=rank, world_size=world_size)

# All-Reduce 示例
tensor = torch.tensor([0., 1, 2, 3], device=device) + rank
dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)  # 原地修改
print(f"Rank {rank}: {tensor}")  # 所有进程输出相同的求和结果

基准测试揭示了硬件带宽的实际限制。测量 All-Reduce 的有效带宽时，需要计算传输数据总量和总耗时：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

# 基准测试核心逻辑
start_time = time.time()
dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM)
torch.cuda.synchronize()
end_time = time.time()

duration = end_time - start_time
size_bytes = tensor.element_size() * tensor.numel()
sent_bytes = size_bytes * 2 * (world_size - 1)  # 发送和接收
bandwidth = sent_bytes / (world_size * duration)
print(f"带宽：{round(bandwidth / 1024**3)} GB/s")

数据并行：梯度同步的集体操作

数据并行的核心是在反向传播后同步梯度。每个进程处理部分数据，计算局部梯度，然后通过 All-Reduce 得到全局平均梯度。

 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

def data_parallelism_main(rank, world_size, data, num_layers):
    setup(rank, world_size)
    
    # 数据切片：每个进程获取部分批次
    batch_size = data.size(0)
    local_batch_size = batch_size // world_size
    local_data = data[rank*local_batch_size:(rank+1)*local_batch_size].to(device)
    
    # 完整模型副本
    params = [get_init_params(dim, dim, rank) for _ in range(num_layers)]
    optimizer = torch.optim.AdamW(params, lr=1e-3)
    
    # 前向传播（使用本地数据）
    x = local_data
    for param in params:
        x = x @ param
        x = F.gelu(x)
    loss = x.square().mean()
    
    # 反向传播
    loss.backward()
    
    # 关键：梯度同步
    for param in params:
        dist.all_reduce(tensor=param.grad, op=dist.ReduceOp.AVG, async_op=False)
    
    # 参数更新（各进程同步更新相同参数）
    optimizer.step()

张量并行：层内矩阵的切分与聚合

张量并行将权重矩阵按列切分，每个进程计算部分输出，然后通过 All-Gather 聚合完整结果。

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

def tensor_parallelism_main(rank, world_size, data, num_layers):
    setup(rank, world_size)
    
    data = data.to(device)
    num_dim = data.size(1)
    local_num_dim = num_dim // world_size  # 特征维度切分
    
    # 每个进程只持有部分参数
    params = [get_init_params(num_dim, local_num_dim, rank) for _ in range(num_layers)]
    
    # 前向传播
    x = data
    for i in range(num_layers):
        # 局部矩阵乘法
        x = x @ params[i]  # 输出形状：[batch_size, local_num_dim]
        x = F.gelu(x)
        
        # 聚合所有进程的部分结果
        activations = [torch.empty_like(x) for _ in range(world_size)]
        dist.all_gather(tensor_list=activations, tensor=x, async_op=False)
        
        # 拼接得到完整特征
        x = torch.cat(activations, dim=1)  # 形状：[batch_size, num_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

def pipeline_parallelism_main(rank, world_size, data, num_layers, num_micro_batches):
    setup(rank, world_size)
    
    # 层分配
    local_num_layers = num_layers // world_size
    local_params = [get_init_params(dim, dim, rank) for _ in range(local_num_layers)]
    
    # 微批次划分
    micro_batch_size = data.size(0) // num_micro_batches
    if rank == 0:
        micro_batches = data.chunk(num_micro_batches, dim=0)
    else:
        micro_batches = [torch.empty(micro_batch_size, dim, device=device) 
                        for _ in range(num_micro_batches)]
    
    # 前向流水线
    for micro_batch in micro_batches:
        # 从上一进程接收（如果不是第一个进程）
        if rank > 0:
            dist.recv(tensor=micro_batch, src=rank-1)
        
        # 本地计算
        for param in local_params:
            micro_batch = micro_batch @ param
            micro_batch = F.gelu(micro_batch)
        
        # 发送到下一进程（如果不是最后一个进程）
        if rank < world_size - 1:
            dist.send(tensor=micro_batch, dst=rank+1)

总结：代码揭示的模式

混合并行策略的选择在代码层面体现为通信原语的组合。数据并行的 All-Reduce、张量并行的 All-Gather、流水线并行的 Send/Recv，这些操作共同构成了分布式训练的通信骨架。

实际系统会在这些基础模式上添加优化：通信与计算的重叠、梯度检查点、异步执行等。但理解这些基础实现是分析和优化分布式训练性能的起点。

注：以上代码为教学用简化实现，省略了错误处理、设备管理、性能优化等工程细节。实际项目应使用 PyTorch 的 DistributedDataParallel、FullyShardedDataParallel 或第三方框架如 DeepSpeed。

• All reduce：根进程 $P_r$ 拥有一个大小为 $K$ 的消息缓冲区 $B$。操作结束后，通信组内的所有进程 $P_i$ 都拥有了 $B$ 的完全一致的副本。
• Scatter：根进程 $P_r$ 拥有一个大向量 $V$，将其均匀切分为 $N$ 个块，记为 $v_0, v_1, \dots, v_{N-1}$。操作结束后，$P_i$ 仅获得对应的分块 $v_i$。
• Gather：Scatter 的逆操作。每个进程 $P_i$ 拥有一个数据块 $v_i$。操作结束后，根进程 $P_r$ 将这些块按 Rank 顺序拼接，在其内存中形成完整向量 $V = [v_0, v_1, \dots, v_{N-1}]$。
• All-Gather：每个进程 $P_i$ 拥有数据块 $v_i$。操作相当于先执行 Gather，再执行 Broadcast。操作结束后，每一个进程 $P_i$ 都拥有完整的向量 $V = [v_0, v_1, \dots, v_{N-1}]$。
• Reduce：每个进程 $P_i$ 拥有一个向量 $X_i$。操作结束后，根进程 $P_r$ 获得结果 $R = X_0 \oplus X_1 \oplus \dots \oplus X_{N-1}$。
• All-Reduce：逻辑上等价于 Reduce + Broadcast。操作结束后，所有进程 $P_i$ 都获得了完全相同的规约结果 $R = \sum_{j=0}^{N-1} X_j$。
• Reduce-Scatter：逻辑上等价于先对全局数据做 Reduce 得到结果 $R$，然后将 $R$ Scatter 给各个进程。操作结束后，$P_i$ 获得结果向量的第 $i$ 个分块。
  

LLM 训练中不同的并行策略

朴素数据并行 DP

数据并行中每个 GPU 都有一个完整的模型副本，输入 Batch 会被切分为多个 Mini Batch 并喂给不同的 GPU，每个 GPU 正常计算出梯度后，通过 All-Reduce 计算出全局梯度，再对参数进行更新。

• 计算扩展性：每个 GPU 计算 B/M 个样本。
• 通信开销： 2x#params，来自 All-Reduce 的过程。
• 内存拓展性：无。每个 GPU 必须完整存储模型。

ZeRO

如下图所示，一个参数量为 $\Phi$ 的模型，进行 AMP 训练时，显存占用大致为：

• 参数占 2$\Phi$ （2 指的是字节，BF16 为两字节）
• 梯度占 2$\Phi$
• 优化器状态 12$\Phi$ （FP32 参数、FP32 动量、FP32 方差）
  

在朴素 DP 中， 每个 GPU 都要完全存储上面三个状态。微软 DeepSpeed 团队提出的 ZeRO 技术，逐步消除了上述三个部分的数据冗余。

• Stage 1：切分优化器状态
  在 Stage 1 中，优化器状态被切分到不同的 GPU 上，每个 GPU 负责对一部分参数进行更新，然后通过 All Gather 让每个 GPU 都有一完整的更新后的参数。

具体来说，其算法步骤为：

1. 每个 GPU 计算 Mini Batch 上的梯度；

2. Reduce Scatter 梯度，这里产生了 #params 的通信开销；
   

3. 每个 GPU 根据它们自己维护的优化器状态对一部分参数做更新；

4. All Gather 参数，这里产生了 #params 的通信开销。
   

对比朴素 DP，Stage 1 就是将一个 All Reduce 操作替换为了 Reduce Scatter + All Gather，并将参数更新的步骤插入到二者之间。之前提过，All Reduce 与 Reduce Scatter + All Gather 是等价的，也就是说，我们在没有新增通信开销的情况下拿到了显著的内存收益。

• Stage 2：切分梯度和优化器状态
  在 Stage 2 中，我们尝试在 Stage 1 的来切分梯度。切分梯度，也就是每块 GPU 保留一小片空间用来存放自己负责的那个切片的梯度。在 Stage 1 中，我们需要先计算完所有的梯度后再对梯度做 Reduce Scatter。为了减少梯度的占用空间，必须得当梯度算完之后立刻做 Reduce Scatter，将其发送到存储这片梯度的 GPU 上，然后当计算图用不到时扔掉不属于自己的梯度。

具体来说，其算法步骤为：

1. 计算一层梯度后，立刻做 Reduce Scatter 将其规约到正确的 GPU 上；
2. 一旦梯度在计算图中再也用不到了，立刻释放对应的空间；
3. 每个 GPU 更新参数；
4. 对参数做 All Reduce。

相比 Stage 1，其通信开销保持不变。在 Stage 1 中对梯度做的一次性 Reduce Scatter 在 Stage 2 中被拆分为多次，但是彼此没有重合，所以整体通信开销不变。

• Stage 3 (FSDP)：切分梯度 + 优化器状态 + 参数
  在 Stage 3 中，参数也被切分，因此进行前向计算时，需要通过通信获取完整参数。如下图所示，在进行前向和反向计算钱，使用 All Gather 获取完整的权重再进行下一步的计算。
  

每次计算前都要做一次 All Gather，这里面的通信开销非常大，但是也并非慢到无法接受。这得益于 Stage 3 将通信和计算折叠：如下图所示，除了第一个前向计算 FWD0 之外，其他的前向计算需要的通信都可以与前一此计算重叠，如果前一层的计算时间大于拉取下一层的时间，通信延迟还能够被完全隐藏。

Stage 3 的通信量是 3 x #params，包括前向和反向各一次的 All Gather 和计算梯度时的 Reduce Scatter，是 Stage 2 的 1.5 倍。

DP 的限制

Data Parallelism 不是万能，其有如下限制：

1. DP 并行数必须小于 Batch Size。Batch Size 不是越大越好，当 Batch Size 过大时，其所需要的训练步数并不能同比下降，也就是模型并不能学的更快。而 DP 的并行数必须得小于 Batch Size，这就约束了 DP 并行的上限。
   
2. ZeRO 各有优劣。Stage 1 和 2 几乎没有额外开销，但是没有切分模型；Stage 3 有额外的开销并且不能减少激活值的内存。

模型并行 Data Parallelism

在 ZeRO Stage 3 中我们将对激活值做切分，从而解决数据并行中难以扩展内存的痛点。在模型并行中，也会将参数切分到不同的卡上，但是卡间交换的是激活值而非参数。模型并行有两类实现，包括流水线并行和张量并行。

流水线并行 Pipeline Parallelism

流水线并行指的是将不同的层切分到不同的卡上，数据像流水线一样在 GPU 之间流动。这是一个很符合直觉但是实现比较困难的并行策略。如下图所示，如果再等待其它卡的过程中什么也不干，每张卡只有 1/n 的时间处于激活状态，效率极其低下。

对此的解决策略是将一个 Batch 分割多个小 Batch，并依次送入流水线中。如下图所示，将数据分割为四条送入流水线，仍会有一些空泡，但是整体效率相比上面的朴素实现得到了提升。整体的空泡率为：$(n_\text{stages}-1)/n_\text{micro}$，Micro Batch 越多，空泡率越少。但是，与模型并行一样，Micro Batch 同样受限于 Batch Size。

既然 PP 也没多好，为啥使用它？

1. 相比 DP，PP 能够节省显存，它不仅对参数进行了切分，还对激活值进行了切分。
2. PP 只需要进行点对点的通信，很适合部署在通信链路效率不高的机器之间。

最后还介绍了 Zero Bubble Pipeline。具体来说，反向计算可以分成两步，第一步是计算对输入的微分 B，第二步是计算对权重的微分 W。其中，对输入的微分是前一层依赖的 OutGrad，必须尽可能快地计算；后者是不被其它卡依赖的，只要在梯度更新前算完即可。Zero Bubble Pipeline 的核心思想就是优先完成 B，当流水线空闲时再完成 W。当然，这只是一个 High level 的介绍，老师一直在强调在工程实现中不管是这个 Zero Bubble Pipeline 还是单纯 PP，其实现起来都相当困难。

张量并行 Tensor Parallelism

流水线并行可以认为是沿着纵轴对模型切分，张量并行则是对横轴切分，即将一个大的矩阵运算切分到两张卡上，如下图所示，X@A = X1@A1 + X2@A2 = Y。

下图展示了一个 MLP 网络中的标准切分方式：第一个矩乘的权重按列切分，第二个按行切分。

前向的流程是：输入 X 通过恒等映射 f 复制到所有卡上，GPU 1 计算出 XA1，GPU 2 计算出 XA2，由于激活函数是逐元素独立计算的，因此无需与其它 GPU 通信。此时 GPU 1 正好持有按按列切分的 Y1，GPU 2 持有 Y2，这与行切分的 B 在维度上完美对应，不需要进行 All Gather 操作即可继续流转到第二个矩乘上。然后再经过 All Reduce 操作 g 对输出进行聚合，就可以得到完整的 Z。

反向的流程与前向类似，只是操作 g 变成了恒等映射，操作 f 变成了 All Reduce 操作。

TP 在每一层中都需要进行大量的通信，因此其对通信要求很高。如下图所示，当 TP 并行度大于 8 时，单卡吞吐量显著下降。同一个节点内的卡通信速度很快，通常 TP 在一个节点内的 8 张卡上部署。

序列并行 Sequence Parallelism

目前介绍的 TP 和 PP 都是对模型参数进行切分的，但是没有切分激活值。在训练过程中，模型的激活值可以通过以下公式来估算：

其中变量含义如下：

• $s$ (sequence length): 序列长度
• $b$ (microbatch size): 微批次大小
• $h$ (hidden dimension size): 隐藏层维度
• $a$ (number of attention heads): 注意力头数

激活值由两部构成，第一部分是线性项，来自 MLP 、Attention 中的线性投影和 LayerNorm；第二部分是二次项 $5bas^2$，这部分来自 Attention 机制，包括注意力分数、Softmax 结果等，这部分可以通过 Flash Attention 的重计算消除。

应用 TP 后，显存占比中线性部分中的系数 24 可以减少为 $1/t$，其中 t 是 TP 并行度。，没有变化的系数 10 来自非 Matmul 的部分，包括 LayerNorm（4）、Dropout（2）、Attention 和 MLP 的输入（4）。

以 LayerNorm 为例，其对每一条序列独立地做 Norm，因此可以通过在序列维度上进行切分，从而降低显存占用。如下图所示，在 LayerNorm 和 Dropout 上做了序列并行，在前向中，LayerNorm 的输出需要经过一次 $g$ All Gather 以将完整 Tensor 送入 Attention，Dropout 的输入需要经过一次 $\bar{g}$ Reduce Scatter 以将完整 Tensor 按序列切分。在反向中，$g$ 和 $\bar{g}$ 操作交换。

应用 SP + FA 后，模型整体的激活值可以降低为 $sbh \left(\frac{34}{t}\right)$。

Ring Attention 和 专家并行

最后简单介绍了下 Ring Attention 和专家并行。

Ring Attention 解决的痛点是超长序列中单卡装不下 KV Cache，其核心思想是将 KV 进行切分，GPUs 之间以环的形式流转 $KV_i$，并计算出相应的结果。

专家并行则是将专家散布到不同的 GPU 上，会面临来自负载均衡和计算通信 Overlap 的挑战。

使用并行策略扩大 LM 的规模并训练

3D 并行

3D 并行指的是混合使用 DP、PP 和 TP 策略，一个简单的原则是：

• 先将模型切分到单卡放得下
  • 使用 TP 策略，TP 并行度不大于每个节点内 GPU 的数量
  • 在节点之间使用 PP 或者 ZeRO Stage 3，取决于带宽大小
• 扩大 DP 规模直至达到 GPU 上限

TL;DR 本讲是 CS336 系列笔记的第六讲。本节从算子融合的必要性出发，横向对比了手写 CUDA、使用 Triton 以及 PyTorch 2.0 编译技术三种实现方式。重点解析了 Triton 如何通过“块级（Block-level）”抽象简化显存管理，并以 GeLU、Softmax 和 Matmul 为例，展示了利用共享内存（SRAM）和分块（Tiling）技术打破访存墙的关键技巧。

算子融合 (Kernel Fusion)

动机：Warehouse vs Factory

老师有一个非常经典的类比：

• DRAM (显存) 就像是一个巨大的仓库（Warehouse），容量大但存取慢。
• SRAM (共享内存/寄存器) 就像是工厂（Factory）的车间，容量小但处理快。

在执行深度学习模型时，如果不做算子融合，数据需要在仓库和工厂之间反复搬运：

1. 从 DRAM 读取数据 $x$ -> 计算乘法 -> 写回 DRAM
2. 从 DRAM 读取数据 -> 计算加法 -> 写回 DRAM
3. …

这种模式下，内存带宽（Memory Bandwidth） 成为了绝对瓶颈。算子融合的核心思想就是：将数据一次性搬运到工厂，完成一系列计算（乘、加、Tanh 等）后，再统一写回仓库。

实例：GeLU

以 GeLU 激活函数为例，其数学表达式包含乘法、加法、Tanh 等多个操作。

• Manual 实现：使用 PyTorch 基础算子拼凑（x * 0.5 * ...），每次操作都会触发一次 Kernel Launch 和全局显存读写。
• Fused 实现：PyTorch 官方的 F.gelu 是经过融合的，只触发一次 Kernel Launch。

性能测试显示，Fused 版本比 Manual 版本快得多（源码中约为 7-8 倍差距），且 Profiling 结果显示 Manual 版本充斥着大量琐碎的 Kernel 调用。

CUDA Kernels：打开黑盒

为了追求极致性能，我们可以直接编写 CUDA C++ 代码。

执行模型

CUDA 的执行层级映射了硬件结构：

• Grid：对应整个计算任务，由多个 Thread Block 组成。
• Thread Block：对应一个 SM，块内线程可以共享 Shared Memory 并同步。
• Thread：最小执行单元，处理单个数据点。

代码与限制

通过 torch.utils.cpp_extension.load_inline 可以方便地在 Python 中内联 CUDA 代码。虽然手动管理线程索引（blockIdx, threadIdx）和内存能够带来性能收益，但其开发门槛极高：

• 必须手动处理内存合并访问（Coalescing）。
• 必须手动管理 Shared Memory 的数据搬运。
• 代码冗长且容易出错（Off-by-one error）。

Triton：Python 时代的 GPU 编程

OpenAI 于 2021 年推出的 Triton 旨在降低 GPU 编程门槛。它引入了 Block-level 的编程抽象，让开发者关注“数据块”而非“单个线程”。

Triton vs CUDA

Triton 编译器自动处理了许多 CUDA 中需要手动优化的痛点：

实现 GeLU

在 Triton 中，我们通过 tl.program_id 获取 Block ID，并计算出该 Block 需要处理的数据指针偏移量。计算过程完全向量化，代码非常接近 Python 原生写法。

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

@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    # 1. 计算当前 Block 处理的数据范围
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    # 2. 加载数据 (Triton 自动处理合并访问)
    x = tl.load(x_ptr + offsets, mask=mask)
    
    # 3. 计算 (会被编译为高效的 PTX 指令)
    y = 0.5 * x * (1 + ... ) 
    
    # 4. 写回
    tl.store(y_ptr + offsets, y, mask=mask)

查看生成的 PTX 汇编代码可以看到，Triton 编译器自动进行了 Thread Coarsening（线程粗化），即一个线程可能处理多个元素以提高指令级并行度。

PyTorch Compilation

这是 PyTorch 2.0 的杀手级特性。

• 原理：通过 torch.compile(model)，PyTorch 会捕获计算图，分析算子间的依赖关系，并自动调用 Triton 后端生成融合后的 Kernel。
• 效果：在 GeLU 的例子中，torch.compile 生成的 Kernel 性能与手写 Triton 几乎一致，远超 Manual 实现。
• Profiling：在 Profiler 中可以看到类似 triton_poi_fused_add_mul_tanh... 的名称，这标志着自动融合生效了。

进阶计算：Softmax 与 Matmul

Triton Softmax：Reduce 操作

Softmax 是一个典型的 Row-wise 操作：$y_i = \frac{e^{x_i}}{\sum e^{x_j}}$。

• 朴素实现：需要多次遍历显存（求 Max -> 减 Max -> 求 Exp -> 求 Sum -> 除法）。对于 $M \times N$ 的矩阵，读写次数高达 $5MN + M$。
• Triton 实现：
  1. 每个 Block 处理矩阵的一行（Row）。
  2. 将整行数据加载到 SRAM。
  3. 在 SRAM 中完成 Max、Exp、Sum 的计算（利用 tl.max, tl.sum）。
  4. 写回结果。
  • 收益：显存读写降低至 $MN$ 次，实现数倍加速。

Triton Matmul：分块

矩阵乘法（$C = A \times B$）是计算密集型任务，但对于大矩阵，显存带宽依然可能成为瓶颈。

• 朴素做法：计算 $C$ 的每个元素都需要读取 $A$ 的一行和 $B$ 的一列，导致大量重复读取。
• Tiling 策略：
  1. 将 $A$ 和 $B$ 分割成小块（Tiles）。
  2. 将 Tiles 加载到 Shared Memory 中。
  3. 复用 Shared Memory 中的数据计算出 $C$ 的一部分局部和。

通过 Tiling，我们将对慢速 DRAM 的访问转换为了对快速 Shared Memory 的访问。此外，Triton 还支持 L2 Cache 优化（Grouped Ordering），通过调整 Block 的执行顺序，尽可能提高 L2 缓存的命中率。

TL;DR 本讲是 CS336 系列笔记的第五讲。本讲从 GPU 的设计理念出发，梳理了 SM/SP 计算架构与内存层级体系，并结合 Roofline 模型，重点解析了访存合并、算子融合、重计算及分块（Tiling）等核心优化策略。此外，还应用上述策略，粗略推导了 FlashAttention 如何利用 Online Softmax 技术打破显存带宽瓶颈，实现 IO 感知的极致加速。

GPU 架构

CPU 与 GPU 设计理念

从架构上看，CPU 有一个很大的控制单元，而 GPU 中计算单元占主导。这揭示了二者在设计理念上的区别：CPU 致力于降低执行延迟，GPU 致力于提高计算吞吐量。

GPU 计算单元架构

如下图所示，GPU 有很多流式多处理器 SM 单元， SM 单元是执行编程模型中的“Block”的执行载体，其具备控制单元，根据资源允许，一个 SM 会并发驻留多个 Block 块；单个 SM 由很多流式处理器 SP 构成，每个 SP 代表编程模型中的一个“thread”，SPs 在不同的数据上执行相同的指令。

GPU 内存架构

一言以蔽之，内存离 SM 越近，它的速度就越快，延迟就越小。L1 和共享内存在 SM 内部，其访存延迟最小；L2 缓存在片上，其延迟是共享内存的 10x；全局显存是 GPU 核心邻近的存储芯片，其延迟最大。

GPU 的执行模型

GPU 的执行模型分为 Block、Thread、Warp 三个级别：

• Block 由多个线程束 Warp 组成，具有独立的显存
• Warp 由 32 个连续的线程组成，线程以线程束为单位进行调度
• Thread 实际的执行单元，线程对不同的数据执行相同的指令，即所谓的 SIMT 模型
  

GPU 编程模型的优势

• 很容易扩大规模——堆叠 SMs 即可
• 得益于 SIMT 模型，容易（？）编程
• 线程是轻量化的，切换线程的开销很小

一些趋势

• GPU 算力在矩乘上做了高度特化
  从下图可以看出，NVIDIA 的 GPU 在发展过程中矩乘和非矩乘的算力提升幅度差异巨大。这意味着我们在构造神经网络时要尽可能使用基于矩乘的算子才能获得最大的硬件收益。
  
• 算力的发展速度快于通信和存储
  

让机器学习任务在 GPU 上跑的更快

Roofline 模型

以计算强度为横轴，吞吐量为纵轴，可以绘制出形如屋顶的曲线，这条曲线可以显示不同计算密度下性能的瓶颈。在上升阶段，瓶颈在于访存；在水平线上，瓶颈在于计算。

技巧零：控制流分歧

GPU 上同一个线程束内执行相同的指令，因此在控制语句中如果同一个线程束中的条件值不一样，那么这些线程就会依次进入所有必要的分支，这相当影响这个线程束的执行性能。这点在之前学习 PMPP 时有提到，参阅：Programming Massively Parallel Processors A Hands-on Approach 4th Edition 学习笔记 Part 1 | 周鑫的个人博客

这是本节中唯一一个与内存无关的技巧。

技巧一：低精度计算

低精度的数据类型占据的空间更小，随之而来的就是更高的计算密度。如下所示，低于 ReLU 算子，将数据类型从 FP32 切换为 FP16，能够将计算密度提升一倍。

低精度存在舍入导致的精度问题，因此在 Tensor Core 中矩乘会使用全精度的累加器来计算，从而同时获得低精度低内存带宽和高精度计算精度的收益。

技巧二：算子融合

如下图左，当有很多算子的时候，如果不做算子融合，就需要反复对同一块数据（作为输入输出）进行搬运。此时，可以将多个小算子融合成一个大算子，一次搬运，多次计算，只在必要时才对数据进行搬运。

技巧三：Recomputation

在依次前向计算过程中，我们需要保存每一层的输出以便在反向中计算梯度。使用重计算技术我们可以在求反向过程中重新计算一遍前向的激活值，从而减少数据搬运。如下所示，一个三层的 Softmax 的前向 + 反向需要 8 次访存，应用重计算后只需要 5 次，访存需求减少 38%。这个技术与模型训练过程中减少显存占用的重计算是同一个技术，只是二者的出发点不同，一个为了加速，一个为了节省显存。

技巧四：内存合并访问

这也是一个在 PMPP 中重点介绍的技术：Programming Massively Parallel Processors A Hands-on Approach 4th Edition 学习笔记 Part 1 | 周鑫的个人博客。
DRAM 的物理结构决定了其支持突发访存，即当访问某个位置的元素时，其周围连续的元素也会被一起读取。

为了充分利用突发访存的特性，CUDA 会自动将线程束中的多个线程连续的访存指令转换为突发访存指令，即如果线程束中 0-31 号线程的同一个访存指令访问的目标是全局内存中连续的 32 个位置，则该访存指令将通过突发访存来实现。具体请参考之前的学习笔记。

技巧五：分块

分块的动机是减少对全局内存的重复访存，将这些重复访存合并对片上内存的重复访问。

如下图所示，在朴素的矩阵实现中，每 P(0,0) 这个元素会被多个线程反复访问，并且这些线程访问的模型也不是合并访问的模式，这会大大降低执行效率。

将输入矩阵按照 2x2 分块以后，就可以将计算设计的小块使用合并访问的访存模式提前加载到共享内存中，并重复利用

如果我们按照步长 T 对 NxN 的矩阵分块，在非分块的情况在每个元素要被加载 N 次，而在分块的情况下每个元素只要被加载 N/T 次，这是 T 倍的加速比，相当可观。

在实际应用中，分块大小可能无法被原张量形状整除，这就会导致资源利用低效。分块大小受到很多因素的约束：

• 内存合并访问
• 共享内存大小
• 原始张量的形状
  

分片的另一大挑战是内存对齐，如果输入张量的形状不是很完美，分片后每个 Block 需要访问的数据横跨两个不同的突发访存块中，那么访存次数相比对齐的情况就会加倍。这类问题的解决方案是通过 padding 手段将其补全为一个内存对齐的布局。

矩乘性能图

下图展示了矩阵乘法计算吞吐量随矩阵形状的变化情况，本节将应用前面所讲的理论对此进行解释。

• 在 size 增加的前期，性能瓶颈在于访存，此阶段吞吐量随计算密度的提升而提升；在后期呈现出波浪线的阶段，性能瓶颈在于计算，即 roof line 的水平阶段。
• 曲线之间的性能差异来自内存对齐和内存合并访问
• 在波浪线阶段，以 K=2 为例，从 1792 到 1793 出现了陡峭的性能下降。这是因为 GPU 在计算矩乘是一般以 256x128 进行分块，1792 需要 98 个 SM，而 1793 需要 120 个 SM，A100 一共只有 108 个 SM，无法一轮全部执行所有的 120 个 Block 所以性能陡峭下降。

FlashAttention

FA 通过使用分块和重计算技术，显著提升了 Attention 的计算速度并降低了对 HBM 的访存需求：

分块一：对 KQV 矩乘分块

第一个分块操作是对 KQV 之间的矩阵乘法进行分块计算，这是之前讲过的常规操作：

分块二：Softmax

标准的 Softmax 需要至少对数据进行三次遍历：

• 遍历所有数据，找到最大值
• 再次遍历数据，计算分母
• 第三次遍历数据，计算结果

其痛点在于这个过程无法分块，需要先统计出全局最值之后才能计算最后的结果。

而 Online Softmax 则引入了一个技巧，在不知道全局最值的情况在边扫描数据，边维护 Softmax 统计量。具体来说，设我们处理到第 $j-1$ 个元素，当前的局部最大值是 $m_{j-1}$，当前的局部和是 $d_{j-1} = \sum_{k=1}^{j-1} e^{x_k - m_{j-1}}$。现在进来一个新的元素 $x_j$，需要进行两个操作：

• 更新最大值：新的最大值 $m_j = \max(m_{j-1}, x_j)$。
• 更新分母和：我们需要把旧的 $d_{j-1}$ 修正到基于新最大值 $m_j$ 的尺度上。
  利用指数运算性质：$e^{x - m_{new}} = e^{x - m_{old} + m_{old} - m_{new}} = e^{x - m_{old}} \times e^{m_{old} - m_{new}}$，所以新的和 $d_j$ 等于：

基于上述数学变换，我们可以在每个 tile 内计算局部最大值和求和值，最后再讲结果合并以修正误差。

FA

下图展示了 FA 的计算过程，首先对 QK 进行分块矩阵，算完立刻做 Online Softmax，并得到一组输出，最后再做一次修正得到最终的结果。这里只是理论介绍，忽略了很多细节，后面 Lab 2 中应该是需要手搓一个 FA 的。

TL;DR 本讲 CS336 系列笔记的第四讲。本讲梳理了 MoE 架构利用稀疏激活实现“高效扩参”的核心机制，并结合 DeepSeek 系列模型的演进路线，重点解析了细粒度专家、共享专家及无辅助损失负载均衡等策略，如何解决大规模训练中的路由坍塌与稳定性难题。

引入

什么是 MoE

MoE 模型指的是将 Transformer 架构中的 FFN（MLP）模块替换为多个稀疏的 FFN 模块（专家），前向计算时每次只稀疏激活一部分专家，由此实现在不增加计算量的情况加增加模型参数。

由于 MoE 架构同计算量在模型参数更多，最终学习到的模型的性能也更好，因此 2025 以来主流模型均转向了 MoE 架构。

当然，天下没有免费的午餐。MoE 的代价是整个模型尺寸变得更大，需要更多的存储空间，以及在并行化方面带来了更大的系统复杂度。

为什么 MoE 变成主流

• 相同的计算量下，参数越多，模型性能越好。
  
• MoE 训练更快。相比 Dense 架构，MoE 架构达到相同的 loss 只需要 1/7 的时间。
  
• 天然适合并行化。将 MoE 的不同专家部署到不同的设备上是一个很自然的想法。

MoE 的局限

• 依赖于基础设施建设
  MoE 模型尽管在计算方面有优势，但是模型本身很大，这就要求训练团队需要有一套强大的基建能够支撑起 MoE 模型在训练过程中对于多硬件、多数据、高度并行化的需求，这对于中小团队来说是很困难的。

• 训练过程依赖经验且不稳定
  MoE 的路由过程是不可微分的，并且路由策略的选择对训练结果影响巨大。因此相比稠密模型其训练过程更容易出现不稳定性，这对开发者的经验和技巧提出了更高的要求。

架构

通常 MoE 指的是下图左边的使用 MoE 替换 Transformer 中的 MLP 层，也有工作尝试使用相同思路将 MHA 替换为 MoE 的版本，但是这并不主流，并且此类模型的训练难度更大。

MoE 的变种

MoE 的各类变种可以归结为三类：路由算法、专家数量和损失函数。

路由算法

总览

路由算法指的是决定每个 token 去往哪个专家的路由机制。可以归结为三类：

• token 选择专家：以 token 为主体，每个 token 独立挑选最匹配的专家
• 专家选择 token：以专家为主体，每个专家从所有输入中挑选最匹配的 token
• 通过全局优化路由：全局角度最优地分配 Token 到专家的任务，通常涉及复杂的数学优化问题
  

常见算法

常见的路由算法是 Top-k 和哈希算法，后者使用一个哈希函数将不同的 token 映射到专家，常常作为 baseline 来评估其它路由算法。

Top-K 算法

如下图所示（从下往上），在 Top-K 路由算法中，有一个可学习的参数 e，首先将输入 u 与 e 做内积，得到二者相似度，对相似度做 Softmax 得到每个专家的分数，然后过滤出前 k 个专家作为门控，再根据每个专家的分数对每个（前 k 个）专家的输出做加权求和。

Fine-grained Expert Segmentation 、Shared Expert Isolation

国产大模型在专家路由架构方面有两个技术趋势。下图 (a) 是传统的 Tok-2 路由算法；(b) 是细粒度专家分割，即将单个专家规模缩小一半，专家总数增加一倍，激活专家数增加一倍，总计算量没有增加，但是专家数的增加能够带来更好的性能；(c) 是共享专家隔离，引入一个始终被激活的专家，从可解释性的角度来说，这个专家可以解决知识冗余的问题，即这个始终被激活的专家负责记忆那些“大家都得会”的知识。

从 DeepSeek 技术报告中的消融实验可以看出，上述架构的演进确实能够提升模型的性能。

训练 MoE

梯度实际上是可以顺着被选到的专家流动下来的。这里所谓的不可微分实际指的是只有被选到的专家才有梯度，但是选不到的专家很容易陷入饿死的状态，即强者越强，弱者越弱，永远得不到更新。

专家路由是一个离散不可微分的过程，对此有三种解决方案：

• 强化学习
• 随机扰动
• 启发式负载均衡损失

强化学习

强化学习可以对离散变量进行建模，理论上其完美适配这个问题，但是实验数据表明这个方法相比其它方案、甚至相比基线模型其性能优势并不显著，但是引入强化学习的同时还带了梯度的不确定性和整体复杂度的显著提升，因此目前没有工作在大规模训练上使用强化学习解决 MoE 的训练的问题。

随机扰动

随机扰动这个方案通过在路由的分数中加入正态分布的噪声和一个可学习的线性层控制整体噪声的幅度。随机性因素的加入使得模型有着更好的鲁棒性，但是此类方案在 LM 后期也由于其潜在的不稳定性被抛弃，业界投向效果更好的启发式方法。

启发式负载均衡损失

这是在 Switch Transformer 中提出的方案，引入一个辅助的损失函数，这个函数是对每个专家的 $f_i$ 和 $P_i$ 的积求均值得到，其中 $f_i$ 表示专家 $i$ 在一个 batch 中实际被选中的概率，$P_i$ 表示专家 $i$ 在路由算法中被分配到的概率总和的平均值。

引用自 Gemini 的答疑解惑：

为什么不直接优化 $f_i$ 使之趋向平均分布？
因为 $f_i$ 是由离散路由算法计算出来的，是一个不可微的值。

为什么不直接优化 $P_i$ 使之趋向平均分布？
因为 Top-K 算法只对数值的绝对排序敏感，只优化 $P_i$ 仍无法杜绝马太效应——只要模型让最强大的专家的分数略高于平均分布即可。

总结，这个 loss 设计得非常巧妙，它实际上是在说：如果一个专家实际上很忙（$f_i$ 大），那我们就惩罚它的预测概率（$P_i$），以此来减少它未来被选中的机会。

在 DeepSeek v1-2 中，还有一个结构类似的辅助损失函数，其作用是在设备之间实现负载均衡，将频率统计从按照专家统计改为按照设备统计即可：

如果没有负载均衡机制，可以看到模型的性能不如上了负载均衡的情况。此外，如果咩有负载均衡机制，只有两个专家被路由，其他专家专家都饿死了。

无需辅助损失的启发式负载均衡

在 DeepSeek V3 中提出了一种不需要引入辅助损失的负载均衡手段。其在 Softmax 计算得到的注意力分数中额外增加了一个偏置项，这个偏置项与这个专家的实际负载负相关，从而实现

MoE 的问题

随机性

当模型所在的卡无法容纳更多的 token 时，多出来的 token 会被丢弃，这在训练和推理阶段引入了随机性，使得模型的输出可能会在不同的 batch 之间表现不同。

稳定性

路由算法中的 Softmax 计算存在稳定性问题，尤其是在 BF16 下，一点数值扰动或者舍入可能对最终的输出有很大的影响。为此，需要在路由选择计算中采用 FP32 并在必要时引入上一讲中介绍的 z-loss 优化手段。

可以看到，z-loss 的引入可以有效抑制训练过程中损失函数的尖峰。

微调

MoE 模型在微调过程中很容易过拟合。一种解决思路是调整模型架构，在模型架构中交替使用 MoE 和 Dense 架构，在微调时只对 Dense 部分微调。另一种思路是大力出奇迹，增加数据量。

Upcycling

一种低成本训练 MoE 模型的方式是先训练一个 Dense 模型，然后将 Dense MLP 拷贝 n 份来构造出一个 MoE 架构，并以此为起点开始训练。通过这一方式可以高效地训练出一个 MoE 模型。

DeepSeek 演进路线

DeepSeek MoE V1

DeekSeek V1 是一个 16B 激活 2.8B 的模型，在架构方面选择了 2 共享专家 +64 细粒度专家，路由算法为 Top-6。使用了标准的辅助负载均衡损失函数。

DeepSeek V2

V2 是一个 236B 激活 21B 的模型，采用 2 共享专家，160 细粒度专家，路由算法为 Top-6 。专家切分的越细，激活的专家就越多，因此引发的通信成本就越大。为此，在路由算法方面他们采取了两步走的策略，先选取 Tok-M 个设备，然后在这些设备上再进行路由，从而控制整体的通信成本。

DeepSeek V3

V3 是一个 631 B 激活 37B 的模型，1 共享专家，258 细粒度专家，路由算法是 Top-8。在路由选择中，计算分数的函数从 Softmax 替换为 Sigmoid。在损失函数方面，其采用了之前提到的“无需辅助损失的启发式负载均衡”，同时采用了一个 sequence 级别的负载均衡损失函数，以确保 token 在推理阶段能够均匀派发给不同的设备。

MLA & MTP

这部分老师讲的比较仓促，直接参考知乎上大佬的解读文章：
deepseek技术解读(1)-彻底理解MLA（Multi-Head Latent Attention）
deepseek技术解读(2)-MTP（Multi-Token Prediction）的前世今生

TL;DR 本讲梳理了现代 LLM 架构设计的“事实标准”（Pre-Norm + RMSNorm + RoPE），并从系统视角解析了 GQA 与滑动窗口机制如何通过优化 KV Cache 访存，解决推理阶段的算术强度劣化问题。

如下图所示，LLM Dense 模型的架构演进始于 2017 年。经历早期的百家争鸣与多方探索后，目前已呈现出明显的收敛态势。本讲将总结当前大模型在架构设计与超参数选择上的经验与技巧，旨在帮助我们立足前人成果，少走弯路。

具体来说，本讲包括：

• 模型架构的变种
  • 激活层和全连接层的选择
  • Attention 的变种
  • 位置编码
• 超参数的选择
  • ff_dim
  • 多头的 head_dim 之和是否一定要与模型的特征维度相等
  • 词汇表大小
• 提升训练稳定性的技巧
• Attention 架构
  

模型架构

Pre-vs-Post Norm

在标准 Transformer 架构中（下图左），Norm 块被安排在每个残差加法结束后进行，即 $x_{t+1}=\text{LayerNorm}(x_t+F(x_t))$，这是 Post-Norm。而在 Pre-Norm 中（下图右），Norm 被放在残差内的子层之前，即 $x_{t+1}=x_t+F(\text{LayerNorm}(x_t))$。

从 2024 年起，使用 Pre Norm 替代 Post Norm 成为业界广泛共识， Pre Norm 在训练稳定性上优于 Post Norm。 在 Post-Norm 中，需要引入一些提升稳定性的 tricks ，例如 Warm Up ，来确保训练稳定进行；而在 Pre Norm 对此则没有那么敏感，并且最终的模型性能与 Post-Norm 不相上下。

为啥 Pre Norm 更好？主流的解释是在 Pre Norm 的残差路径上不存在 Norm 模块，梯度可以通过恒等映射直接从顶层传播到底层，从而避免了梯度消失或者爆炸问题。当下的模型越来越深，训练的稳定性就是刀乐，业界普遍转向 Pre Norm 也是个很符合直觉的现象。

既然 Pre Norm 是在残差的子层计算前进行 Norm，是不是有工作尝试在残差子层计算后进行 Norm 呢？有的，包有的。老师将其成为 Double Norm，即在残差子层计算前后都进行一次 Norm，近期的 Gemma 2 就采用这种架构。

LayerNorm vs RMSNorm

LayerNorm 与 RMSNorm 公式如下，前者在特征维度进行标准化后按照两个可学习的向量 $\gamma$ 和 $\beta$ 缩放和偏移，后者省去了对均值偏移的操作，

近年来，业界普遍从 LayerNorm 转向了 RMSNorm。 这一转变可以归结为：RMSNorm 计算量和参数量都更少。如下图所示，就理论计算量而言，Norm 模块占比只有 0.17% 并不高，但是由于需要搬运显存，导致了在实际运行中占比高达 25.5%，省略对均值的偏移可以降低对显存读写的需求，从而提升实际运行速度。与此同时，模型的精度并没有因为这一转换而降低。

在 Norm 模块，另一潮流是现代模型倾向于扔掉偏置项。 不只是 RMSNorm 的偏置项，还包括所有线性层中的偏置项。这一转变的底层动机与 RMSNorm 一致：减少内存开销。此外，尽管还说不清楚，但是确实有现象表明丢弃偏置项可以提升学习过程的稳定性。

激活函数

激活函数有很多很多，采用不同的激活函数，可以衍生出各种 MLP 的变种。这里直接用课程里的 PPT 来介绍：

ReLU 是最简单的激活函数，但是在负值区域没梯度，可能导致神经元饿死。GeLU 在负值区有微小梯度仍可以继续训练。

*GLU 系列的 MLP 层修改了第一个全连接层，将其额外逐元素乘上一个 $xV$，实现所谓“门控机制”。

根据激活函数的不同，产生了 ReGLU、GeGLU 和 SwiGLU 的变种，其中门控单元中特征维度长度为模型特征维度的 2/3 以确保整体参数量和计算量不变。

实验表明 GLU 系列确实具有更好的表现，但是也不是必须的，例如 GPT-3 就没有使用这一激活函数。

并行与串行架构

传统的 Transformer 架构属于串行架构，输入依次经过 LayerNorm、 Attention 、残差、MLP 和 MLP 模块；与之相对的变种是并行架构，输入同时经过 Attention 和 LayerNorm 模块，然后将两个模块的输出相加并送入 MLP。这一设计使得模型在大规模训练时能够提升一定的训练速度。

位置编码

在 LM 早期，有各种位置编码函数，包括 Sine、绝对位置编码、相对位置编码：

有了以上这些函数，为啥还需要 RoPE 呢？一个完美的相对位置编码应该满足如下的数学定义：对于一个位置嵌入函数 $f(x, i)$，我们希望两个被位置嵌入后的向量的点积（也就是注意力分数）只和这两个向量的相对位置有关，也就是：

上述性质也可以被描述为平移不变性。

为了实现平移不变性，RoPE 的思路是将一个每个 token 都根据其位置旋转对应角度，这样在计算二者的夹角时就仅与二者相对位置相关。每个 token 在特征维度构成的空间是高维的，而角度是一个在二维平面上的概念，为此 RoPE 的解决方案是将特征维度两两分组，每组视为一个二维向量进行旋转操作，并且不同组的旋转角度不同，从而同时捕获高频和低频位置信息。

小节

• Pre-vs-Post Norm
  • 选 PreNorm，更稳定
• LayerNorm vs RMSNorm
  • 果断选 RMSNorm，更快，模型性能有时更好
• 激活函数
  • GLU 系列似乎更好，但是差距并不显著
• 并行与串行架构
  • 没有严格的消融实验证明并行更好
  • 并行在计算速度上有一定的优势
• 位置编码
  • 果断选 RoPE，其具备平移不变性

超参数

MLP 特征数

在 MLP 模块中，会将输入的特征的维度从模型维度 $d_{\text{model}}$ 在内部投射到 MLP 特征数 $d_{\text{ff}}$。对于非 GLU 系列的 MLP 模块，一般将其放大 4 倍；对于 GLU 系列的模块，一般将其放大 8/3 倍。这是目前的共识。

有工作研究过上采样倍数与 loss 的关系，结果显示倍数在个位数这个数量级上效果最好。曾经有工作（T5）选择了 64 的倍数，但是那也是昙花一现，在其后续的 T5 v1.1 中也选择了常见的 2.5 倍作为上采样系数。

注意力特征数

本节讲的是 num_heads x head_dim 这一组超参数的选择，我将其描述为注意力超参数。

目前业界的共识是注意力特征数等于或者略多于模型特征数。

对于这个参数的选择，有工作认为当固定总特征数后，如果想要提升头数，就不得不降低每个头的特征数，这会导致每个头的表征能力下降，进而降低模型的整体性能。但是在实践中这一效应并不显著，业界还是倾向于使用模型特征数作为注意力的特征数。

横纵比

横纵比指的是模型的特征数与模型的深度 num_layers 之间的比值。主流模型选择 1xx 这个区间，这似乎也是一种共识。

相关工作也证实了 1xx 这个区间是比较好的一个甜点区，与模型参数量相关性并不显著。

横纵比影响着并行策略的选取，比较深的模型适合流水线并行，比较胖的模型适合张量并行。

词汇表数量

词表的数量取决于支持的语言数量和具体生产用途，单语言模型在 30-50k 这个量级，多语言模型在 100-250k 这个量级：

正则化

理论来说，在预训练阶段不需要正则化手段（Dropout、Weight Decay），因为预训练阶段在训练数据上跑一轮，很难产生过拟合。但事实上，大家倾向于使用一些正则化技术：

对于这一反直觉的现象，一篇工作对此的解释是：Weight Decay 不是为了防止过拟合，它是一种与学习率相互交互的技巧，将权重变小在某种程度上等价于将学习率变大，这可以让模型在训练后期学习率较小的时候获取一个更小的损失值。

小结

• MLP
  • GLU 系列选择 8/3 作为上采样系数
  • 非 GLU 系列选择 4 作为上采样系数
• 注意力特征数
  • 注意力特征数等于模型特征数
• 纵横比
  • 选择 1xx 作为纵横比
• 正则化
  • 使用正则化手段，但是是为了获得更小的损失

技巧

本节将介绍一些训练过程中提升稳定性的技巧。

Softmax

Softmax 计算中由于存在 exp 和除法操作，因此其对数值很敏感，例如除 0 等问题，很容易在训练过程中爆炸。在 Transformer 中有两处 Softmax，一处是模型最后对 logits 进行 Softmax 操作，另一处是 Attention 块中计算注意力得分时的 Softmax 操作。下面针对这两处分别介绍提升稳定性的技巧。

• Output Softmax
  记 Softmax 中的分母项为 Z，为了防止 logits 中 Z 过大，引入了一种叫做 Z-loss 的优化手段，即将 log(Z) 作为正则项加入到损失函数中，迫使 Z 优化到 1 附近，从而避免在大规模计算中的指数运算引发浮点数溢出和梯度异常。

• Attention Score
  在 MHA 计算中，一些视觉和多模态模型在 QK 计算注意力分数前分别对 QK 计算一次 LayerNorm，这一技术被称为 QK Norm，能够提升训练过程中的稳定性。
  

Attention 机制

GQA / MQA

如下图所示，在 LM 的训练阶段，Attention 模型具备比较高的算术强度 Arithmetic Intensity（下图中 k 应该是 d），即在训练阶段是计算受限的，这能够充分发挥 GPU 的计算能力。

但是在推理阶段，我们无法并行地去预测每一个 token，只能逐 token 生成。如下图所示，当生成 t+1 个 token 时，前 t 个 KV 值可以被 Cache，在 qkv 的 projection 阶段，只需要计算 input 中最后一个 token 的对应的 qkv。

带有 Cache 的 MHA 中，总的计算量并没有变（训练一次性计算出所有的 QKV，推理分 n 步计算出所有的 QKV），但是总的访存量显著增加，因为需要反复加载 Cache，最终导致算术强度在推理阶段劣化。

导致算术强度劣化的原因是 KV Cache 的反复加载，解决问题的思路是减少 KV Cache 的大小。KV Cache 来自 KV Projection，因此一个很自然的想法是减少 KV 的注意力头数。如果将头数减少到 1，就是 MQA（Multi-Query Attention），此时的计算强度随着注意力头数的增加而增加。如果将 KV 的头数降低为 Q 的 1/n，即每一个 KV 对应 n 个 Q，这就是 GQA（Group Query Attention），同样可以提升算术强度。

Sparse/Sliding Window Attention

在传统 Attention 中，每个 token 会与其之前所有的 token 交互生成注意力分数，这使得 Attention 计算量和 KV Cache 的长度随 N 增长而平方和线性增长。

如下图所示，在滑动窗口注意力中，每个 Token 只与其附近固定数量个 Token 交互，从而使得总计算量随 N 线性增长。

对于稀疏注意力，老师没有太多的介绍，Gemini 对此的概括是：为了解决 Sliding Window “目光短浅”的问题，Sparse Attention 引入了更复杂的连接模式，包括：滑动窗口，保留 Sliding Window，用于捕捉局部语法关系；全局节点，定几个特殊的 Token（比如 [CLS] 或者特定的 prompt token），让它们可以看到所有人，所有人也可以看到它们；空洞，像空洞卷积一样，每隔 $k$ 个词看一眼。增加了感受野，但不需要全看。

Trick

为了同时获取稠密 Attention 全局感受野和稀疏 Attention 减少计算量的收益，现在有一种架构同时使用上述两种 Attention。以 4 个 block 为一个周期，前三个使用稀疏注意力，第四个使用标准注意力。在某些实现中，仅在滑动窗口注意力中使用位置嵌入 RoPE，在全量注意力中不做位置嵌入

我的实现在 GitHub - LittleHeroZZZX/cs366 at assignment1 ，欢迎交流与指正。

Tips

提高开发效率以及避免踩坑的一些建议：

• 关闭 AI 补全（显著提高降低效率，但是能够了解到更多的细节处理）
• 使用 Python 类型标注系统，并将 Pylance 类型检查设置为标准，这样能在静态检查出绝大多数类型、参数不匹配的问题
• 实现 BPE 前首先厘清其中各个数据结构和流程的概念，例如语料库、pre-token、token、预分词，然后再动手
• 使用 logger 和 tqdm 随时随地打印进度，以便对各个组件耗时和瓶颈组件有个直观的认识

BPE 分词器

分词器的整个构建流程包括：

• 词表初始化：构造初始词表，包括 256 个 ASCII 字符和 special tokens
• pre-tokenize： 给定语料库，将语料库按照给定正则表达式划分为 pre-token，并统计词频
• 计算 BPE merges：给定 pre-token，合并出现频次最高相邻 pre-token 对作为一个 token，并不断重复这个过程直至词表达到目标

对于分词器，我将其组织为两个类：

• Pre-Tokenizer：负责将给定语料划分为 pre-token，并统计各个 pre-token 的词频
• TokenizerTrainer：基于给定词频结果运行 BPE 算法，并产出 merge pair 和 词表

Pre-Tokenizer

Pre-Tokenizer 类的接口如下所示，其对外提供一个 pre_tokenize 方法为给定字节流生成 pre-token 迭代器，一个 __call__ 方法将给定语料库文件转换为 pre-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

class PreTokenizer(ABC):
    @staticmethod
    def _merge_pre_token_counts(*pre_token_counts: PreTokenCount) -> PreTokenCount:
        """Merge multiple PreTokenCount dictionaries into one.

        Returns:
            PreTokenCount: The merged PreTokenCount.
        """

    def _process_chunk(self, chunk: Chunk, special_tokens: list[Token]) -> PreTokenCount:
        """
        Process a single chunk of text and return the pre-token counts.

        Args:
            chunk (Chunk): The chunk of text to process.

        Returns:
            PreTokenCount: A dictionary-like object mapping pre-tokens to their counts.
        """

    def pre_tokenize(self, str_bytes: bytes, special_token_list: list[Token]) -> Iterator[Token]:
        """
        Pre-tokenize the given bytes string.

        Args:
            str_bytes (bytes): The input bytes string to pre-tokenize.
            special_token_list (list[Token]): The list of special tokens.
        Returns:
            Iterator[Token]: An iterator over the pre-tokens.
        """

    @abstractmethod
    def __call__(self, corpos_path: str, split_special_token: Token, special_tokens: list[Token]) -> PreTokenCount:
        """
        Pre-tokenize the given corpus.

        Args:
            corpos_path (str): Path to the corpus file.
            split_special_token (token): The special token used to split the corpus.
            special_tokens (list[Token]): List of special tokens.

        Returns:
            PreTokenCount: A dictionary-like object mapping pre-tokens to their counts.
        """

在具体实现时，为了采用多进程以提高处理效率， _process_chunk 负责统计一个 chunk 中 pre-token 出现的频次。如下所示，一个语料库会被以 special token 为边界划分为多个 chunk，每个 chunk 由多篇被 special token 分隔的文档组成。

在多进程的实现中，每个进程每次处理一个 chunk。统计每个 chunk 中各个 pre-token 的出现次数，_process_chunk 先按照 special token 将 chunk 划分为 document，然后使用 Counter 模块和给定的 pre-token 的正则表达式统计单篇 document 中每个 pre-token 的出现次数。

在 _process_chunk 的帮助下，多进程的实现就变得很简单，只需将给定的语料库划分为指定数量的 chunk，然后使用进程池派发这些任务，最后将每个 chunk 的词频归约在一起即可。

参考实现：

 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 MultiProcessPreTokenizer(PreTokenizer):
    def _process_chunk_with_boundry(
        self, corpos_path: str, start: int, end: int, special_tokens: list[Token]
    ) -> PreTokenCount:
        with open(corpos_path, mode="br") as f:
            f.seek(start)
            chunk = f.read(end - start)
            pre_token_count = self._process_chunk(chunk, special_tokens)
        return pre_token_count

    def __call__(self, corpos_path: str, split_special_token: Token, special_tokens: list[Token]) -> PreTokenCount:
        final_pre_token_count: PreTokenCount = defaultdict(int)

        start_time = time.time()
        file_size = os.path.getsize(corpos_path)
        num_cpus = cpu_count()

        desired_chunks = num_cpus * 100

        chunk_boundaries = find_chunk_boundaries(
            file_path=corpos_path,
            desired_num_chunks=desired_chunks,
            split_special_token=split_special_token,
        )

        chunks_args = []
        for i in range(len(chunk_boundaries) - 1):
            start = chunk_boundaries[i]
            end = chunk_boundaries[i + 1]
            chunks_args.append((corpos_path, start, end, special_tokens))

        logger.info(f"Splitting task into {len(chunks_args)} chunks.")

        with Pool(processes=num_cpus * 2) as pool:
            chunk_iter = pool.imap_unordered(self._worker_wrapper, chunks_args)

            for chunk_result in tqdm.tqdm(chunk_iter, total=len(chunks_args), desc="Pre-tokenizing"):
                for token, count in chunk_result.items():
                    final_pre_token_count[token] += count

        end_time = time.time()
        logger.info(
            "Takes {:.2f} seconds to pre-tokenize, speed: {:.2f} bytes/second",
            end_time - start_time,
            file_size / (end_time - start_time),
        )
        return final_pre_token_count

    @staticmethod
    def _worker_wrapper(args):
        tokenizer_instance = MultiProcessPreTokenizer()
        return tokenizer_instance._process_chunk_with_boundry(*args)

TokenizerTrainer

TokenizerTrainer 负责在给定词频统计结果上运行 BPE 算法。BPE 算法简而言之，就是找出出现频次最高的组合并合并，为了高效实现 BEP 算法，首先要回答以下几个问题：

1. 怎么初始化各个组合（Pair）的出现频次？
   Pre-Tokenizer 提供的是每个 pre-token 的频次，可以通过遍历每个 pre-token 中可能出现的组合并累加频次即可。例如，对于 pre-token Hello ，可以贡献 (H, e), (e, l), (l, l), (l, o), (o, 空格) 这六组频次。

2. 合并后怎么维护各个组合的频次？
   每当一个组合被合并后，各个组合的频次需要更新，以 Hello 中 el 合并为例，需要更新的频次包括：

• 新产生组合： el 作为一个独立的 token，与其邻接的两个 token 将产生新的组合 (H, el) 和 (el, l)，这两个组合需要新增。
• 频次减少的组合：el 合并之后，这个组合不复存在，其词频需要设置为 0。除此之外，与被合并的 token 邻接的其它 token 组成的组合频次也许相应减少，即 (H, e) 和 (l, l) 的频次需要减少 Hello 这个单词的次数。

3. 合并后如何快速定位到受影响的 pre-token？
   在上一问中，我们解决了在 pre-token 已知的情况下词频的维护逻辑。但是如何快速找到受影响的 pre-token？朴素的方法是直接遍历整个 pre-token，显然每次都要遍历的方案完全不可接受。可行的方案是我们维护一张 pair 到 pre-token 的映射表，表示含有这个 pair 的 pre-token 列表。这张表在组合频次初始化时也一起被初始化，在组合被合并时也一起被更新，从而在 merge 过程中找到受影响的 pre-token 列表。

4. 如何记录当前 pre-token 的状态？
   pre-teken 的状态指的是当前的 pre-token 由哪些 token 组成。举个栗子，BPE 算法刚开始时，Hello 这个 pre-token 的是由 (H, e, l, l, o, 空格) 这六个 token 组成的，在算法后期，其可能是由 (Hel, l, o空格) 这三个 token 组成。此时如果需要合并 (l, o空格)，在更新组合的频次时就要知道 l 的前一个 token 是啥，而非简单查询 l 前一个字符是什么，因此我们还需要一个字典来维护每个 pre-token 当前的状态。

5. 如何获取频率最大的组合？
   朴素的方案是每次都遍历整个词频表，时间复杂度是 O(n)。我们可以使用最大堆来优化这一过程，从而可以将单次获取并维护最大值的时间复杂度降为 O(log n)。

回答上述问题后，TokenizerTrainer 在训练过程中需要维护的数据结构就呼之欲出了，包括：

• pair_counts：各个组合的出现频次（组合是一个两个 token 组成的元组 tuple(token, token)）
• pre_token_states：每个 pre-token 当前的状态（组成），即当前 pre-token 如何使用 token 来表示
• pair_to_pretokens：一个字典，表示含有指定组合的 pre-token 的 token 列表
• pair_heap ：一个列表，用于在 Python 中实现最大堆

train

在 train 方法中，首先调用 init 方法对 pre_token_states、pair_counts、pair_to_pretokens、pair_heap 这四个数据结构进行初始化。

在主循环中获取出现次数最多的组合作为新的 token，并合并和维护上述数据结构，直至词汇表达到目标值或者没有可合并的组合。

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

def train(self) -> tuple[Vocab, list[tuple[bytes, bytes]]]:
	self._init()

	num_merges_needed = self.target_vocab_size - self.vocab_size
	if num_merges_needed <= 0:
		return {}, self.merges

	while self.vocab_size < self.target_vocab_size:
		merge_pair = self._determine_merge_pair()
		if merge_pair is None:
			break

		self.merges.append(merge_pair)
		self._merge_pair(merge_pair)

	return self.vocab, self.merges

_merge_pair

TokenizerTrainer 的核心是合并逻辑。合并算法的流程图如下所示，首先根据 Pair 获取包含合并 Pair 的所有 pre-token 列表，然后对于其中的每个 pre-token，遍历其状态（token 组成），如果是 Pair，则将 (..., A, B, ...) 替换为 (..., AB, ...)，然后减少 (PrevA, A) 和 (B, NextB) 的频次，并增加 (PrevA, AB) 和 (AB, NextB) 的频次，最后设置 (A, B) 的出现频次为 0。

参考源码：cs366/assignment1-basics/src/tokenization/tokenizer_trainer.py at 8afb4e7e1973a3da6f0cede46d8d35a0959f2982 · LittleHeroZZZX/cs366 · GitHub

TokenizerTrainerC

上述 Python 版本的存在一定的性能瓶颈，所以使用 C++ 重构了 TokenizerTrainerC，核心逻辑与 Python 一致，提速十几倍，可以直接查阅代码：

• C++ 实现： cs366/assignment1-basics/csrc/tokenizer_trainer.cc at 2b4fb5a84a382aeafbd16775fcc23cba021f9831 · LittleHeroZZZX/cs366 · GitHub
• Python 接口： cs366/assignment1-basics/src/tokenization/tokenizer_trainer.py at 8afb4e7e1973a3da6f0cede46d8d35a0959f2982 · LittleHeroZZZX/cs366 · GitHub

构建 Transformer 语言模型组件

在本部分，我们需要使用基本的 Tensor 操作来构建 Transformer 各个组件模块。绝大多数组件难度不大，根据定义照抄即可。各组件实现思路：

• Linear：直接矩乘，注意权重需要转置
• Embedding：熟练运用高级索引
• RMSNorm：抄公式；为了数值稳定性，先转 float 再转回去
• SwiGLU：抄公式
• RoPE：这个难度较大，建议让 Gemini 辅助理解，强烈安利我的这份对话辅导 ‎Gemini - 直接体验 Google AI 黑科技
• SDPA：笔者最近在 Paddle 上改造这个 API，比较熟悉，刚上手可能需要多理解一下其中的 shape 变换
• MHA：引入了 mask，同样需要熟悉这里面的 shape 变换
• Transformer LM：搭积木，注意 LM Head 后不需要 softmax

参考实现在：

• 基础组件： cs366/assignment1-basics/src/nn/basic.py at 8d800aeb4942e710ac835b1be6f89aecc0bae483 · LittleHeroZZZX/cs366 · GitHub
• 前向函数： cs366/assignment1-basics/src/nn/functional.py at 8d800aeb4942e710ac835b1be6f89aecc0bae483 · LittleHeroZZZX/cs366 · GitHub
• 网络结构： cs366/assignment1-basics/src/nn/networks.py at d0d8413fd2e8e048450d893bc648baf70cbd4258 · LittleHeroZZZX/cs366 · GitHub

训练组件

本部分需要实现损失函数、AdamW 优化器、lr schedule 和梯度裁剪。同样，需要注意的点：

• 交叉熵损失
  • logits 转 float 以避免精度和溢出问题
  • target（label）转 int64，使用 gather 提高效率
• SGD：抄公式
• AdamW：抄公式
• lr schedule：抄公式
• 梯度裁剪：公式中的梯度的二范数指的是所有梯度拼在一起的二范数，即如果要裁剪，则裁剪所有参数的梯度，而非分参数判断这个参数的梯度是否要裁剪

参考实现：

• 交叉熵损失：cs366/assignment1-basics/src/nn/functional.py at 8afb4e7e1973a3da6f0cede46d8d35a0959f2982 · LittleHeroZZZX/cs366 · GitHub
• 优化器： cs366/assignment1-basics/src/nn/optimizer.py at 66c38c225aa1424feb0fc95bb30153391e5ae638 · LittleHeroZZZX/cs366 · GitHub
• lr schedule 和 梯度裁剪： cs366/assignment1-basics/src/nn/utils.py at 8afb4e7e1973a3da6f0cede46d8d35a0959f2982 · LittleHeroZZZX/cs366 · GitHub

训练循环

本部分我们需要实现 Data Loader、Checkpoint 和训练循环。

Data Loader

load batch 函数就是随机对给定 dataset 按照长度 context_size 采样 batch_size 次：

1
2
3
4
5
6
7
8

def load_batch(dataset: npt.NDArray, batch_size: int, context_size: int, device: torch.device):
    max_start_index = dataset.shape[0] - context_size + 1
    start_indices = np.random.randint(0, max_start_index - 1, size=batch_size)
    x_batch = np.stack([dataset[i : i + context_size] for i in start_indices])
    y_batch = np.stack([dataset[i + 1 : i + context_size + 1] for i in start_indices])
    x_tensor = torch.tensor(x_batch, device=device, dtype=torch.int)
    y_tensor = torch.tensor(y_batch, device=device, dtype=torch.int)
    return x_tensor, y_tensor

Checkpoint

保存点机制可以使用 torch.save 机制来实现，需要保存的状态包括：模型状态、优化器状态和迭代步数，将他们组织成一个字段让 torch 来保存即可：

 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 save_checkpoint(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    iteration: int,
    out: str | os.PathLike | typing.BinaryIO | typing.IO[bytes],
):
    states = {
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "iteration": iteration,
    }
    torch.save(states, out)


def load_checkpoint(
    checkpoint: str | os.PathLike | typing.BinaryIO | typing.IO[bytes],
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
) -> int:
    states = torch.load(checkpoint)
    model.load_state_dict(states["model_state_dict"])
    if optimizer is not None:
        optimizer.load_state_dict(states["optimizer_state_dict"])
    iteration = states["iteration"]
    return iteration

训练循环

整个训练过程流程图如下所示，根据配置文件解析训练配置，初始化模型和优化器后扫描输出目录，如果存在检查点，则从断点恢复模型和优化器状态，然后使用 mmap 加载数据集。每个训练 Step 依次计算学习率、加载数据集、计算 logits 和损失、反向传播、梯度裁剪，最后优化器步进。其中，每隔指定 steps 数量，都会保存训练状态到磁盘，并使用 wandb 记录训练指标。

• 配置文件
  通过使用 yaml 配置文件，可以高效地记录和管理模型的训练配置，以便后续进行大量的对比和消融实验，参考配置：

 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

model:
  vocab_size: 10000
  # 上下文长度/序列最大长度 (Context length)
  context_length: 256
  # Transformer 块层数 (Number of Transformer layers)
  num_layers: 4
  # 隐藏层维度/模型大小 (Hidden dimension)
  hidden_dim: 512
  # Transformer 内部前馈网络的维度 (Inner dimension of FFN, typically 4*hidden_dim)
  inner_dim: 1344
  # 多头注意力机制的头数 (Number of attention heads)
  num_heads: 16
  # RoPE (Rotary Position Embedding) 的旋转基数（通常为 10000）
  theta: 10000.0
  # 训练设备 (Device for training: 'cuda', 'mps', or 'cpu')
  device: "cuda"
  # 模型参数的数据类型 (Data type: 'float32' or 'bfloat16')
  dtype: "bfloat16"

optimizer:
  # 学习率 (Learning rate)
  learning_rate: 3.0e-4
  # 权重衰减 (Weight decay for AdamW)
  weight_decay: 1e-2
  # AdamW 优化器的一阶矩估计衰减率 (Beta1)
  beta1: 0.9
  # AdamW 优化器的二阶矩估计衰减率 (Beta2)
  beta2: 0.999
  # AdamW 优化器的数值稳定性参数 (Epsilon)
  eps: 1.0e-8

training:
  # 批次大小 (Batch size)
  batch_size: 96
  # 总训练迭代次数 (Total training iterations)
  total_iterations: 20000
  # 学习率预热的迭代次数 (Warmup iterations)
  warmup_iterations: 1000
  # 学习率余弦退火周期迭代次数 (Cosine cycle iterations)
  cosine_cycle_iterations: 20000
  # 梯度裁剪的最大 L2 范数 (Max L2 norm for gradient clipping)
  max_l2_norm: 1.0
  # 模型保存的间隔（迭代次数） (Checkpoint saving interval)
  checkpoint_interval: 1000
  # 模型和配置的输出目录 (Output directory)
  output_dir: "./output_full/"
  # 训练数据文件路径 (Path to training data .npy file)
  train_data: "save/data/TS-train.bin"
  # 验证数据文件路径 (Path to validation data .npy file)
  val_data: "save/data/owt_train.bin"
  # 保存检查点的步长（在 `train` 函数中用于判断是否保存检查点）
  save_step: 500
  log_step: 50

• WandB
  推荐使用 WandB 自动记录实验配置和性能指标，避免手动整理记录实验数据，提升参数搜索和消融实验的效率。

train 函数参考实现在： cs366/assignment1-basics/src/nn/train.py at d0d8413fd2e8e048450d893bc648baf70cbd4258 · LittleHeroZZZX/cs366 · GitHub

生成文本

文本生成的流程图如下所示，从指定目录加载模型权重和 Tokenizer，转换为 token id 后喂给模型，获取最后一个 token 的预测结果，应用 temperature 和 softmax 后按照概率从大到小排序，保留累计概率不小于 p 的候选词，并按比例从中随机选择一个作为预测结果。预测下一轮时，将上一轮的 token 加入输入序列，重复上述过程，直至输出 endoftext 或者达到用户指定的 token 上限。

decode 参考实现在 cs366/assignment1-basics/src/nn/decode.py at d78dc879e554f47094824deadcdc2801551d0158 · LittleHeroZZZX/cs366 · GitHub

本文是 CS336 系列笔记的第二讲，PyTorch 与资源计算。本讲比较简单，从两个问题出发，讨论如何计算计算量，并顺带讲一些 PyTorch 的入门概念。本讲的重点不在于 PyTorch 基础，而在于培养“资源计算”的思维模式。

本节课的目标就是回答两个问题：

• 使用 1024 张 H100 在 15 T token 的数据集上训练 70B 的模型需要多久？
• 在 8 张 H100 上使用 AdamW 优化器，最大可以训练多大的模型？

内存计算

Tensor 入门

深度学习的一切都是使用 Tensor 进行存储的，包括参数、梯度、优化器状态、数据、激活层。

入门结束。（逃）

Tensor 内存

几乎所有 Tensor 都以浮点数的形式存储，通常默认的数据格式为 FP32。

Tensor 占据的内存取决于数据类型和 Tensor 内元素的个数。

为了容纳更多的参数，业界还引入了 FP16 、 BF16 和 FP8 数据格式。

总而言之：

• FP32 对于训练已经够够的了，但是需要大量内存
• FP8 / FP16 / BF16 训练存在风险，可能会导致不稳定
• 解决方案：使用混合精度训练

Tensor 操作

课程介绍了一些 Tensor 的基本操作，此处不赘述。值得一提的是 stride 机制，若不熟悉可以查阅一下。

Tensor 操作 FLOPs 计算

一次浮点运算 Float-point Operation FLOP 指的是一次浮点数加法或者浮点数乘法。需要区分的两个概念：

• FLOPs：float-point operations，浮点计算量，用来表达需要进行多少次浮点计算
• FLOP/s：FLOP per second，每秒浮点计算量，用来衡量硬件的计算能力

一些直观感受：

• 训练 GPT-3（2020）使用了 3.14e23 FLOPs
• 训练 GPT-3（2023）被推测使用了 2e25 FLOPs
• A100 峰值算力 3.12e14 FLOP/s
• H100 系数

假定输入数据为 [B, D]，线性层将 D 映射到 K 即权重矩阵形状为 [D, K]。对于输出张量的每个元素，都需要通过将两个长度为 D 的向量逐元素相乘后相加，即 2D 个 FLOPs，一共有 B x K 个输出元素，所以一次矩乘的 FLOPs 为 2 x B x D x K。

mfu = 模型实际每秒FLOP / 硬件理论最大每秒FLOP，一般来说 MFU 大于 0.5 已经相当不错了。

模型

参数初始化

简单来说，参数初始化也是一门学问，否则会导致梯度爆炸或者消失，这门课对此没有详细介绍。CMU 10414 对此有从数学上的推导，见我当时的学习笔记：《CMU 10-414 deep learning system》学习笔记 | 周鑫的个人博客。

其它内容

这里还介绍了如何训练一个模型、优化器、检查点等内容，过于基础，此处不赘述。

混合精度训练

低精度可以减少显存占用并加快计算速度，同时避免低精度带来的不稳定，我们的策略是：默认使用 FP32，并尽可能使用 BF16、FP8，具体来说，在前向时使用低精度，在反向时使用高精度。

小结

本讲作为 PyTorch 的入门课，介绍的都是一些很基本的概念。涉及的 PyTorch 的用法和训练过程都是一个引子，重点还是与培养资源计算的思考模式，毕竟大模型时代，效率就是金钱。

本文是 CS336 第一讲的学习笔记，主要介绍开设这门课程的背景和动机，并对课程主要内容做了概览。还介绍了不同的分词器基本原理及其优缺点。

引入

为什么要学习这门课程

现象：研究人员与底层技术越来越远。八年前，他们需要自己实现和训练模型；六年前，他们需要下载一个模型，并进行微调；现如今，他们仅仅修改模型的提示词。

当然，上述现象并不是坏事，通过抽象程度的不断提高，研究人员才得进行更高效地工作。但是，基于 抽象漏洞定律，抽象不可能完美隐藏底层细节。并且，仍有许多基础性的工作亟待开发，这些工作需要对底层有着深入的洞察。

这门课程将通过从头建立语言模型的方式来让大家彻底理解语言模型各个层面的技术。

制约因素

学习现代大语言，有两大制约因素：

1. 在大模型工业化时代下，训练一个大模型的成本极高，需要大量的硬件、时间和金钱成本。
2. 最前沿模型（GPT 系列）几乎不公开详细的技术报告，缺乏学习材料。

小模型不具备代表性

我们可以通过训练小模型（参数量小于 1B）来弥补上述制约因素，但是小模型的经验并不能完全适用于大模型。

• 案例一：随着规模的扩大，注意力层和全连接层的计算量的比例将改变

  • 小规模下 FFN 不显著
  • 大规模下 FFN 主导
    

• 案例二：随着参数量的增加，模型能力显著提高
  

从这门课能学到什么？

有三方面的知识：

• 运行原理：大模型的各个部件的工作原理（例如什么是 Transformer、如何构建 GPU 上并行的模型）
• 思维模式（Mindset）：旨在将硬件性能发挥到极致，并认真对待规模化问题的理念（例如，扩展定律/scaling laws）
• 实践直觉（Intuitions）：凭经验判断哪些数据和建模决策能够带来良好的准确率。

惨痛的教训

有一种错误认知是：算法不重要，重要的是规模。
正确的认知应该是：可以被规模化扩大的算法最重要。

正确率=效率 x 资源。在大规模下，算法效率格外作用，因为承担不起试错成本。[2005.04305] Measuring the Algorithmic Efficiency of Neural Networks 这篇文章揭示了，在 ImageNet 上算法的不同可以带来 44x 的效能提升，远超摩尔定律带来的 11x 的硬件性能提升。

研究现状

• 基础组件
  • Sequence to Sequence 建模（2014）
  • Adam 优化器（2014）
  • 注意力机制（2014）
  • Transformer 架构（2017）
  • MoE（2017）
  • 模型并行技术（2018-2019）
• 早期基础模型
  • ElMo：基于 LSTMs 做预训练，在下游任务上微调（2018）
  • BERT：基于 Transformer 做预训练，在下游任务上微调（2018）
  • Google’s T5：将所有任务都映射成 text to text 任务（2019）
• 模型规模化探索（闭源模型）
  • OpenAI’s GPT-2 (1.5B)：生成流畅文本，展现了初步的零样本（zero-shot）学习能力，并采取分阶段发布策略（2019）
  • OpenAI’s GPT-3 (175B)：展示出强大的上下文学习（in-context learning）能力，但模型闭源（2020）
  • Google’s PaLM (540B)：进行了更大规模的训练，但后来被认为是欠训练的（undertrained）（2022）
  • DeepMind’s Chinchilla (70B)：提出了计算最优的缩放定律（compute-optimal scaling laws），认为在同等算力下，应该用更多数据训练更小的模型（2022）
• 开源模型发展
  • EleutherAI (The Pile & GPT-J)：发布了大规模开放数据集 The Pile 和开源模型 GPT-J，推动了开源生态的发展 (2020-2021)
  • Meta’s OPT (175B)：尝试复刻 GPT-3，但遇到了大量的硬件挑战 (2022)
  • Hugging Face / BigScience’s BLOOM：一个大型多语言开源模型，项目重点关注数据的来源和治理 (2022)
  • Meta’s Llama 系列：发布了多个版本的 Llama 模型，在开源社区产生了巨大影响，成为许多后续模型的基础 (2023-2024)
  • Alibaba’s Qwen (通义千问) 系列：阿里巴巴推出的一系列功能强大的开源模型 (2024)
  • DeepSeek’s models (深度求索)：由深度求索公司发布的一系列高性能开源模型 (2024)
  • AI2’s OLMo 2：由艾伦人工智能研究所（AI2）推出的完全开放模型，包括训练数据和代码 (2024)
• 模型的开源程度
  • 只开放 API，不开源模型和权重
  • 开源模型、权重和技术报告，但是不开源数据集
  • 开放模型、权重和数据集
• 当下最前沿的模型（2025）
  • OpenAI o3
  • Anthropic Claude Sonnet 3.7
  • xAI Grok 3
  • Google Gemini 2.5
  • Meta Llama 3.3
  • DeepSeek r1
  • Alibaba Qwen 2.5
  • Tencent Hunyuan-T1

课程总览

基础

目标：对模型全流程有一个大致了解。

在此阶段将学习分词器、模型架构和训练。与之匹配的任务一中，我们要实现 BPE 分词器、Transformer、交叉熵、AdamW 优化器和训练循环，并在特定数据集上做训练。

系统

目标：榨干硬件性能

在此阶段将学习：核函数、并行化和推理技术。与之匹配的任务二中，我们要实现融合 RMSNorm、分布式并行训练、优化器状态切分，并对实现进行性能测试和分析。

缩放定律

目标：在小规模上做实验，并据此预测大规模下的超参数和损失

问题：在给定计算量 FLOPs 的预算下，是应该使用更大的模型还是在更多的 token 上训练？

学习计算最优化缩放定律。与之匹配的任务三中，我们将定义一个 API，其能够基于之前的运行结果预测制定超参下模型的损失值。以及绘制一个曲线，拟合不同算力预算下的损失值。并给出指定预算下让损失最小的超参配置。

数据

问题：我们希望模型具备什么能力？多模态、代码、数学？这决定了我们需要什么样的数据。

与之匹配的任务四中，我们将把爬取的 html 文件转换为文本，训练分类器对文本的质量和是否有害进行分类，使用 MinHash 识别重复数据，并在置顶预算下最小化困惑度。

对齐

到此为止，我们得到的是一个基础模型，其擅长预测下一个词。但是仍需要通过对齐使得其在特定任务上表现优秀。

与之匹配的任务五中，我们将实现 SFT、DPO、GRPO。

分词器

引入

原始文本是一串使用 Unicode 编码的字符。而语言模型的输出是在一系列 token 表上的概率分布。

因此我们需要一个分词器将字符串编码为 token，以及将 token 解码为字符串。词汇表大小 vocabulary size 指的就是所有可能 token 的个数。

字符分词器 Character Tokenizer

最简单也是最符合直觉的分词器是字符分词器。每个字符都对应 Unicode 编码中的一个编号，通过查表可以将将字符串编码为 tokens。

这个方案的问题在于：

• 字符表会很大
• 很多词汇几乎不被使用，使得字符表的使用效率很低

字节分词器 Byte Tokenizer

Unicode 编码的文本可以表示为一系列字节流。在 UTF-8 编码中，每个字符都被编码为 1~4 个字节长度。使用这个方案，所有的字符串都被编码为最大 255 token 序列。

这个方案的问题在于：

• token 序列太长
• 每个 token 只能表示一个 1 字节

按词分类器 Word Tokenizer

这个分类器的思想是将字符串按照每个词进行分割，例如 "Hello world" 分为 ["Hello], " ", "world"]，然后这个词集合映射到整数序列上。

这个方案的问题在于：

• 字符表是相当有限的，无法编码碰到没见过的词
• 字符表同样可能很大

字节对编码分类器 Byte Pair Encoding BPE

BPE 在 1994 被提出用来进行数据压缩，在 2015 在 NLP 中被应用到机器翻译，并在 2019 年被 GPT-2 使用。

其基本动机是使用单个 token 来表示常见的序列，使用多个 token 来表示罕见的序列。在 GPT-2 中使用词分类器将原始文本分解为一个个片段，然后在每个片段上运行 BPE 算法。

BPE 算法：将字节流视作原始 token，然后将最常出现的相邻 token 对合并。

例如：

• 初始字节流 [22, 134, 245, 22, 134, 22, 245]
• 统计其中出现次数最多的相邻 token 对 (22, 134) 出现两次
• 将出现最多的 token 对 (22, 134) 合并并产生新 token 256 替换，得到 [256, 245, 256, 22, 245]
• 重复上述过程，只至达到置顶字符表大小或者没有重复的相邻字节对

欢迎来到我的个人技术博客！很高兴能在这里与你分享我的学习、思考和实践经验。

🧑‍💻 我是谁？

我是一名 INTJer，热衷于探索前沿技术，并渴望通过代码创造实际价值。

目前就读于中国科学技术大学软件学院研究生一年级，同时也通过 第八期飞桨黑客松活动 在百度飞桨进行远程实习。

🚀 这个博客是关于什么的？

在这个博客里，你将会看到关于以下主题的内容：

• [AI 框架特性解析与实践]
  • 百度飞桨「启航计划」小结——CINN后端Pass改造 | 周鑫的个人博客
  • Programming Massively Parallel Processors A Hands-on Approach 4th Edition 学习笔记 Part 1 | 周鑫的个人博客
  • Programming Massively Parallel Processors A Hands-on Approach 4th Edition 学习笔记 Part 2 | 周鑫的个人博客
• [C++ 学习笔记]:
  • Effective Cpp 第三版学习笔记 | 周鑫的个人博客
  • 如何在VSCode中“优雅”地配置CMake —— 以PaddlePaddle为例 | 周鑫的个人博客
  • 在VSCode中对CUDA和Python代码进行联合调试 | 周鑫的个人博客
• [国外公开课学习记录]
  • 《CMU 10-414 deep learning system》学习笔记 | 周鑫的个人博客
  • CS144 Lab 实验笔记 | 周鑫的个人博客
  • EfficientML | 周鑫的个人博客
• [实际问题的解决过程]
  • 搭建ZeroTier MOON服务器 | 周鑫的个人博客
  • 安装并切换指定gcc或者g++版本 | 周鑫的个人博客
  • 博客搭建日志 | 周鑫的个人博客
  • 在Hugo中使用KATEX渲染数学公式 | 周鑫的个人博客
  • 在PaperMod中引入侧边目录和阅读进度显示 | 周鑫的个人博客
• 以及任何我觉得有趣或值得分享的技术点滴。

我创建这个博客的初衷是记录自己的学习过程 同时与更多技术同好交流 。我希望通过文字，能够帮助到正在学习或工作中遇到相似问题的朋友们，同时也期待与大家共同进步。

感谢你的阅读！希望你在这个博客能有所收获。

本文为 EfficientML Lab 2 实验笔记，包含 K-Means 量化、K-Means QAT、线性量化等内容，难度不大，内容丰富。

Part 1: K-Means Quantization

Qustion 1

第一个问题是实现 K-means 量化的核心算法，其中 K-means 本身是调库实现的。

第一小问求簇数，n bit 可以表示 2^n 个簇。第二小问是根据已有的 Codebook 表示量化后的张量，使用 Tensor 的索引表示即可。需要注意的是，codebook 可能是调用者传入的而非一定是由我们计算得到的，因此在表示的时候要使用 codebook 的成员来获取 centroids 和 labels，否则在后面会报错。

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

    if codebook is None:
        ############### YOUR CODE STARTS HERE ###############
        # get number of clusters based on the quantization precision
        # hint: one line of code
        n_clusters = 2 ** bitwidth
        ############### YOUR CODE ENDS HERE #################
        # use k-means to get the quantization centroids
        kmeans = KMeans(n_clusters=n_clusters, mode='euclidean', verbose=0)
        labels = kmeans.fit_predict(fp32_tensor.view(-1, 1)).to(torch.long)
        centroids = kmeans.centroids.to(torch.float).view(-1)
        codebook = Codebook(centroids, labels)
    ############### YOUR CODE STARTS HERE ###############
    # decode the codebook into k-means quantized tensor for inference
    # hint: one line of code
    quantized_tensor = codebook.centroids[codebook.labels].view(fp32_tensor.shape)
    ############### YOUR CODE ENDS HERE #################

Question 2

略

Part 2: Trained K-Means Quantization

Question 3

在低比特量化后模型掉点很厉害，因此要进行 QAT。量化后的权重的梯度推导为：

但在本实验中，简单起见，我们使用相同簇的原始权重的均值作为量化后的该簇更新后的值。代码实现就一行：

1

codebook.centroids[k] = fp32_tensor[codebook.labels == k].mean()

最终不同量化位数得到的性能指标如下表所示，2 bits 掉点很夸张，QAT 后也难以恢复到原始性能。

Part 3: Linear Quantization

Question 4

本问实现的是线性量化的核心函数，即给定张量、量化位宽、缩放系数、零点，计算量化后的张量。根据线性量化公式：

计算即可：

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

############### YOUR CODE STARTS HERE ###############
    # Step 1: scale the fp_tensor
    scaled_tensor = fp_tensor / scale
    # Step 2: round the floating value to integer value
    rounded_tensor = scaled_tensor.round()
    ############### YOUR CODE ENDS HERE #################

    rounded_tensor = rounded_tensor.to(dtype)

    ############### YOUR CODE STARTS HERE ###############
    # Step 3: shift the rounded_tensor to make zero_point 0
    shifted_tensor = rounded_tensor + zero_point
    ############### YOUR CODE ENDS HERE #################

值得一提的是，在 step 4 中，执行了一步将溢出的结果压缩到范围之内的操作。

Question 5

计算缩放系数和零点的公式分别为：

代码照抄即可。

权重有一个特殊的性质：分布通常都是关于 0 点对称的，因此权重量化的零点可以直接设置为 0。

此外，经验表明，对卷积核进行量化时，按照输出通道逐通道量化能够取得更好的表现

Question 6-8

在此之前，实验文档首先推导了考虑线性量化的全连接层和卷积层的表达式，推导过程进行了一系列代入、假设和化简，主要包括：

最终得到的结论为：

Q7 和 Q8 的代码实现相同，需要注意的是在对 output 进行缩放时，由于是逐通道量化的，因此权重的缩放系数是个张量，需要处理好形状以便进行广播。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

    ############### YOUR CODE STARTS HERE ###############
    # Step 2: scale the output
    #         hint: 1. scales are floating numbers, we need to convert output to float as well
    #               2. the shape of weight scale is [oc, 1, 1, 1] while the shape of output is [batch_size, oc]
    output = output * (input_scale * weight_scale / output_scale).swapaxes(0, 1)

    # Step 3: shift output by output_zero_point
    #         hint: one line of code
    output = output + output_zero_point
    ############### YOUR CODE ENDS HERE #################

Question 9

要干的活文档基本都干好了，只剩下一个对输入进行量化的活需要我们完成。照猫画虎，使用 get_quantization_scale_and_zero_point 计算缩放系数和零点，使用 linear_quantize 进行线性量化即可。

1
2
3
4

############### YOUR CODE STARTS HERE ###############
x_scale, x_zero_point = get_quantization_scale_and_zero_point(x, 8)
return linear_quantize(x, 8, x_scale, x_zero_point)
############### YOUR CODE ENDS HERE #################

量化后的模型精度为 92.21% 几乎没掉点。

为什么在线性量化中没有 ReLU 层：ReLU 层被融合到前一层中网络中，可以减少数据的搬运次数。

Question 10

回答来自 deepseek：

小结

实验文档本身体量很大，知识点也很丰富，但是大多数代码都已经给出了，每个回答只需要写一行或者两行代码，并且周围也给出了充足的提示。这使得实验本身缺乏挑战性，有点鸡肋。

本文介绍了训练后量化（PTQ）和量化感知训练（QAT）技术，PTQ 通过 Per-Tensor/Channel/Vector 等不同粒度划分量化参数，结合动态范围裁剪（校准集统计或 KL 散度优化）和 AdaRound 学习式舍入来平衡精度与效率；QAT 则在前向传播中模拟量化并利用直通估计器（STE）绕过梯度断层，而二元/三元量化通过引入可学习缩放因子减少极低比特（1-2bit）下的精度损失，在压缩模型的同时实现硬件加速与内存优化。

训练后量化 Post-training Quantization (PTQ)

量化粒度

• Per-Tensor
  整个张量共享一个缩放因子。对于一些小模型，精度损失得比较厉害。绘制出一个 Tensor 不同通道的箱式图，可以发现有的通道之间数据范围差异很大，因此不适合在 Tensor 尺度上使用统一的量化参数。

• Per-Channel
  Per-Channel 为每个通道都计算了独立的放缩因子，该方案在反量化后能够得到更多样的权重表示，量化误差也相对应地更低。
  

• Group wise

Per-Vector在传统 Per-Tensor 量化的基础上，为每个子向量增加一个额外的整型缩放因子，即反量化公式变为：

整个 Tensor 共享同一个高精度的缩放因子，达成了精度和硬件效率之间的折中。

以 4-bit quantization with 4-bit per-vector scale for every 16 elements 为例，实际的量化位宽为 4+4/16 = 4.25 bits。

Per-vector 本质上是一种多级的缩放方案，通过使用不同的数据表示类型来表示缩放因子，能够实现 MX4/6/9 等缩放方法，从而达成不同程度的量化位宽。

Dynamic Range Clipping

权重的数据在模型训练结束后就已经固定，因此其范围也是固定的。而激活层随输入数据的变化而变化，数据范围会很大，确定其数据范围的技术需要专门拿出来讨论。

首先需要收集激活层的统计信息，分为在训练过程中和训练后收集。

• 训练时
  使用指数移动平均来统计训练过程中激活层的最大和最小值，其公式为：

• 训练后
  如果无法获取训练时的数据集，可以在模型中运行一些校准集。激活层的数据分布可能如下图所示，两侧的长尾是激活层中的极值，这些极值的存在会降低量化范围的表示能力（通过压缩量化表示范围，可以在有限的数据位宽中提高量化的表示精度）。
  
  优化目标可以设置为最小化量化的均方误差，即：

对于高斯分布或者拉普拉斯等已知分布，可以求出上述优化目标的数值解。但对于一些不常见的分布，优化目标就是 KL 散度来衡量用量化后的分布来近似激活层的分布的信息损失，即：

舍入

• AdaRound
  常见的舍入策略是舍入到最近的数，例如四舍五入。但是很多时候这并不是一个最佳策略。这里介绍一种基于学习的舍入方法，在舍入前给张量加上一个可学习的参数再进行四舍五入，即：

那么确定该参数的优化目标为：

其中，$V$ 为待学习参数，$h()$ 是一个将输入映射到 $[0, 1]$ 的函数，$f_{reg}$ 是一个正则项，使得 $V$ 接近于 0/1。

量化感知训练 Quantization-Aware Train (QAT)

在 K-means 量化中，将每个簇的梯度相加后作为该簇的梯度从而更新权重，线性量化的微调过程则要复杂得多。本节名字为 Quantization-Aware Train，即在训练过程中考虑量化。

伪量化 Fake Quantization

伪量化示意图如下所示，在前向传播过程中对权重和激活层进行模拟量化（先量化到低精度，再反量化到高精度），再反向传播中使用高精度更新权重。

直通估计器 Straight-Through Estimator, STE

随之而来会引入一个新的问题，量化是一种舍入，这就导致输出层会变成分段的台阶函数，这个函数的梯度为 0。这里使用 STE 技术，绕过不可导的计算，即认为量化操作对梯度没有影响：

加上 STE 后，整个 STE 的示意图如下所示，直接使用 Loss 对量化后的 W 和 Y 的梯度来估计对量化前的 W 和 Y 的梯度。

二元/三元量化 Binary/Tenary Quantization

所谓二元/三元就是将权重量化到 0 和±1，从而大大减少内存占用和加速计算。

二元化

• 确定性二元化
  确定性二元化根据一个阈值（常见为 0），小于该阈值的权重量化为 -1，大于该阈值的权重量化为 +1。可以预见的是，这个方案掉点很厉害。
• 随机二元化
  一种随机策略为 Binary Connect，权重 r 以概率 $\sigma(\min (\max ((r+1)/2, 0), 1))$ 被量化到 1，否则量化到 -1，该函数的图像如下所示。简单来说，若权重大于 1 或者小于 -1，则被量化到 1 或者 -1，否则按照概率线性增加的形式量化到 1。
  
  这个方案的缺点是硬件不友好，其要求硬件在量化过程中生成随机数。

可以遇见的，采用这个方案量化后模型掉点很厉害，下图达到了 21.2%。但是，如果给量化后的权重增加一个 32 位的 scale factor，确保量化前后权重的均值一致，尽管量化误差仍旧很大，但是模型精度掉点仅仅只有 0.2。

同时二元量化权重和激活层

过于震惊，后面用到在学。

三元量化

• Ternary Weight Networks, TWN
  三元量化将阈值之内的权重量化为 0，之外的量化为±1。阈值的选取一般为权重绝对值的均值乘上 0.7。同样这里需要一个 scale factor。

符合预期得，三元量化精度损失小于二元量化。

• Trained Ternary Quantization, TTQ
  为了进一步减少量化损失，TTQ 将量化后 scale factor 修改为两个可学习的参数，分别用来表示 +1 和 -1 对应的缩放系数，然后寻找这两个参数的最优值。

Why 启航？

为什么选择了启航计划？在回答这个问题之前，先介绍一下背景：当时学习了 CMU 10414 DLSys 课程，准备学习 TVM 或者 MLIR，但相关基础欠缺，一直苦于找不到切入口。在互联网上🏄‍♀️的时候无意中发现了启航计划，了解到其对新手相当友好：没有面试筛选、任务比较简单、有专门答疑研发老师，当时第三期正在进行，遂订阅了第三期的 ISSUE，蹲第四期的活动。

启航计划安排

刚开始有三个打卡任务，分别是编译 Paddle、跑通 Paddle Mix 和 为 Paddle 添加文档。第一个任务用来熟悉本地编译 Paddle 和单测，第三个任务用来熟悉 GitHub 工作流程。

理论上，完成这三个任务就能够达到最低结营条件，但我们参加这个活动肯定不是为了这张结营证书，而是想要提升自己的。这三个任务对于提升自己的作用聊胜于无。下一步，就可以选择几个的专项团，尝试一些低星任务。

由于启航计划面向新手，任务比较简单。低星任务基本是照葫芦画瓢，即照着样例基本就能完成，通过低星任务可以理解这个专项团的总体目标。高星任务则是一些推广，或者逻辑比较复杂，但也基本不涉及从 0 到 1 的创作，本质上还是模仿。

当然，任务简单并不意味着可以很轻松地完成。对于我们这种零经验的开发者来说，极大概率需要花上几天时间才能理解“1+1=2”，后期还会发现理解是不完备的或者根本就是错的😭。在完成的过程中，可以反复阅读任务文档和观看任务讲解视频，多与导师沟通，很多时候他们都能一语点醒梦中人。特别感谢 Hongqing-work 老师，CINN Pass 改造基本都是在向她请教，老师周末和晚上都能不厌其烦地答疑解惑，太感动了😭。

在训练营中，每两周都需要提交周报。这既是一个让我们回顾过去两周产出、规划未来的好机会，也能够了解其他同学的进度，保证自己不掉队。按照我的经验，1-2 周用于完成打卡任务，开始尝试低星任务；3-4 周继续完成某个专项团的任务，此时已经可以冲击一些高星任务了；5-8 周，渐臻佳境，对于某个专项团的任务已经能够做到游刃有余，并且尝试其它专项团任务。

CINN 后端 Pass 改造

在本次启航计划中，我一共完成 7 个 CINN 后端 Pass 改造任务。这里介绍一下这个专项团的收获。

背景

本次任务的背景是 CINN 升级了后端 IR 表示，将原来 Func-Expr 层级结构中的 Expr 进行了细化，重新划分为 Func-Block-Stmt-Expr，重新划分后的 IR 层次更加清晰。

与之对应地，后端 Pass 也被划分为 FuncPass、BlockPass、StmtPass、ExprPass 四个级别，使用配套的 PassManager 应用 Pass。其层次结构为：

此外，还提供了 IR 访问方法：

1. 类型不敏感的 Stmt 和 Block 级别的访问/修改方法，在遍历 Stmt 前后将会调用用户传入的回调方法：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

// Visitors
void Visit(const BlockRef &block,
           const std::function<void(const StmtRef &)> &pre_callback,
           const std::function<void(const StmtRef &)> &post_callback);

void Visit(const StmtRef &stmt,
           const std::function<void(const StmtRef &)> &pre_callback,
           const std::function<void(const StmtRef &)> &post_callback);
// Mutators
// ...

2. 类型敏感的 Stmt 和 Block 定制化访问模板类，用户可以通过重写 virtual StmtRetTy VisitStmt(const StmtRef &stmt, Args... args) 定制化访问不同的 Stmt：

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

template <typename StmtRetTy = void,
          typename BlockRetTy = void,
          typename... Args>
class StmtVisitor {
 public:
  virtual StmtRetTy VisitStmt(const StmtRef &stmt, Args... args) {
    CINN_CHECK_STMT_DEFINED(stmt)
    switch (stmt->stmt_type()) {
#define __(stmt__)                                \
  case ir::StmtNodeTy::stmt__:                    \
    return VisitStmt(stmt.as<stmt__>(), args...); \
    break;

      NODETY_FORALL_STMT(__)

      default:
        PADDLE_THROW(::common::errors::InvalidArgument(
            "Deadcode, not supported StmtNodeTy"));
#undef __
    }
  }
...

为什么要升级 IR？

从后端 Pass 的角度来看，IR 升级主要有两个好处：1. Pass 编写更加清晰和规范；2. Pass 便于管理。

旧 IR 下的的 Pass 大都通过继承 IRMutator/Visitor 在遍历整个 IR 的过程中修改来实现 Pass 的功能，但实际上其只需要针对某个特定类型的 Stmt/Block 处理即可。旧 IR 下的 IRMutator 为了便于开发者使用，提供了对各种类型的 Expr/Stmt/Block 默认遍历，例如对于 IfThenElse 默认实现版本会遍历条件和两个分支：

1
2
3
4
5
6
7
8
9

template <typename T>
void IRMutator<T>::Visit(const IfThenElse *expr, T op) {
  auto *node = op->template As<IfThenElse>();
  IRVisitorRequireReImpl<void, T>::Visit(&node->condition, &node->condition);
  IRVisitorRequireReImpl<void, T>::Visit(&node->true_case, &node->true_case);
  if (node->false_case.defined())
    IRVisitorRequireReImpl<void, T>::Visit(&node->false_case,
                                           &node->false_case);
}

这种默认实现在很多情况下是不必要的，比如在合并两个相同的 If 中，显然不需要对条件应用此 Pass，也不需要对 Expr 级别的表达式进行访问。

理论上说，开发者可以通过重写对应的 Visit 方法来及时进行截断，但一方面这样会使得 Pass 的代码比较臃肿，另一方面 Pass 在开发时并没有此规范，已经成为遗留问题。

在此次 IR 和 Pass 改造后，原有的 IRMutator 将只保留对于 Expr 级别的访问逻辑，对于 Stmt 和 Block 级别的遍历由 PassManager 完成。例如，StmtPassManager 将会遍历这个函数，并为每一条 Stmt 调用一次其管理的 StmtPass，而在 StmtPass 内部，其只需要处理符合其目标的逻辑。

此外，新版的 StmtVisitor 没有提供 VisitStmt 默认实现，这可以强迫开发者自定义遍历逻辑，并及时截断不需要的遍历。

Pass 编写范式

升级后的 IR 的编写范式一般为：1. 继承对应级别的 Pass 基类；2. 使用一个内部类对 Func/Block/Stmt 进行遍历实现核心功能，这个类可以继承 StmtMutator/IRMutator 或者调用 Visit/Mutate 方法来实现遍历；3. 返回 Success。

1. 继承对应级别的 Pass 基类
   第一步就是分析原 Pass 是什么级别 Pass，核心要义是抓住原 Pass 需要什么级别的信息以及是什么级别的修改。例如：

• IfFusionPass 是合并两个多个条件相同的 If，其要识别和删除多个 If，只有拿到这个 If 所在的 Block 能够实现多个语句的识别和单个语句的删除，这是一个 Block 级别的 Pass；
• RearrangeLoadInstruction 是将一个函数内的 Load 放到最前面执行，以提高指令的并行并行程度，为了确保 Load 是一个函数内最先执行的语句，以及对函数内所有相同的 Load 替换为本地变量，这是一个 Func 级别的 Pass；
• EliminateCommonFactorOfLocalIndex 需要获取当前 For 的嵌套信息，那必须由本 Pass 负责对 IR 的遍历，否则无法获取当前 For 的嵌套级别信息，因此这是一个 Func 级别的 Pass。

我的经验是：如果一个 Pass 仅仅需要当前 Stmt 的内部信息、不需要删除或者替换当前 Stmt、并且对于当前 Stmt 的嵌套级别没有要求（例如不要求当前的 For 是最外层/最内层的 For），那么其是一个 Stmt 级别的 Pass；如果一个 Pass 需要跨语句的信息，或者需要删除/替换/添加一条 Stmt，那么其是一个 BlockPass；如果一个 Pass 需要自己控制对 IR 的遍历过程，或者需要当前的嵌套上下文，那么这是一个 Func 级别的 Pass。

2. 编写实现类
   一些比较简单的 Pass 就是一个继承了 IRMutator 的实现类，此类 Pass 一般只需要额外继承 StmtMutator，如果不涉及 Expr 层面，则去掉对于 IRMutator，然后将原有逻辑迁移到新 IR 下即可。可参考 RemoveScheduleBlock。新 IR 下，很多变量都被设置为私有变量，必须通过 getter 和 setter 进行读写。

一些比较复杂的 Pass 可能有多个 Mutator 对 IR 进行多次访问，一般第一次是收集全局信息，之后再进行修改。读懂源码后再照葫芦画瓢修改即可。

更复杂的是调用了一些旧 IR 的方法，例如 ir::ir_utils::CollectIRNodesWithoutTensor，这种情况下可以判断一下传入的参数是否是 Expr，如果是 Expr 则还可以调用该方法（因为对 Expr 是封闭的，Expr 中不会有 Stmt 或者 Block），否则要根据这些方法的逻辑在新 IR 下进行实现。

3. 返回 Success
   这个没啥好说，返回 LogicalResult::success() 即可。

Tips

1. Pass 应该实现为无状态的
   无状态指的是 Pass 不应该依赖之前的信息，或者记录一些持久信息。例如，一个对于 For 进行处理的 Pass，其内部不应该记录当前 For 的名字以防止重复。如果想要避免重复访问，可以将其实现为 FuncPass 手动处理遍历逻辑。

2. PassManager 是按照 DFS 后序遍历的
   这一遍历顺序可以保证最内部的语句被最先访问。Pass 改造过程中是可以依赖这一行为的。

3. Pass 之间缺乏通信机制
   Pass 之间是缺乏通信机制的，一些 Pass 在应用是前是需要检查能否进行变换的，这些检查 Pass 可以作为变换 Pass 的内部的一部分，在变换 Pass 实例化一个 PassManager 应用检查的 Pass。

4. 可参考 Halide文档
   CINN 在很多设计上参考了 Halide 和 TVM，在如果碰到一些例如不知道 Stmt 的作用的疑问，可以参考这两个这两个文档更加丰富的社区，往往会有惊喜收获。

后记

作为第一次开源活动经历，我个人觉得还是收获颇丰的。纸上得来终觉浅，绝知此事要躬行，很多之前没有实操的技术都在这次活动中得到了锻炼，例如 Git 和 GitHub 的工作流、VSCode 和 CMake 的配套、GLOG 的使用等等，以及对于 CINN 中 Pass 改造的经验，更是很好的学习 AI Sys 的切入口。

鼓励没有尝试过的同学多多参加这类活动，一定能不虚此行！

Why 启航？

为什么选择了启航计划？在回答这个问题之前，先介绍一下背景：当时学习了 CMU 10414 DLSys 课程，准备学习 TVM 或者 MLIR，但相关基础欠缺，一直苦于找不到切入口。在互联网上🏄‍♀️的时候无意中发现了启航计划，了解到其对新手相当友好：没有面试筛选、任务比较简单、有专门答疑研发老师，当时第三期正在进行，遂订阅了第三期的 ISSUE，蹲第四期的活动。

启航计划安排

刚开始有三个打卡任务，分别是编译 Paddle、跑通 Paddle Mix 和 为 Paddle 添加文档。第一个任务用来熟悉本地编译 Paddle 和单测，第三个任务用来熟悉 GitHub 工作流程。

理论上，完成这三个任务就能够达到最低结营条件，但我们参加这个活动肯定不是为了这张结营证书，而是想要提升自己的。这三个任务对于提升自己的作用聊胜于无。下一步，就可以选择几个的专项团，尝试一些低星任务。

由于启航计划面向新手，任务比较简单。低星任务基本是照葫芦画瓢，即照着样例基本就能完成，通过低星任务可以理解这个专项团的总体目标。高星任务则是一些推广，或者逻辑比较复杂，但也基本不涉及从 0 到 1 的创作，本质上还是模仿。

当然，任务简单并不意味着可以很轻松地完成。对于我们这种零经验的开发者来说，极大概率需要花上几天时间才能理解“1+1=2”，后期还会发现理解是不完备的或者根本就是错的😭。在完成的过程中，可以反复阅读任务文档和观看任务讲解视频，多与导师沟通，很多时候他们都能一语点醒梦中人。特别感谢 Hongqing-work 老师，CINN Pass 改造基本都是在向她请教，老师周末和晚上都能不厌其烦地答疑解惑，太感动了😭。

在训练营中，每两周都需要提交周报。这既是一个让我们回顾过去两周产出、规划未来的好机会，也能够了解其他同学的进度，保证自己不掉队。按照我的经验，1-2 周用于完成打卡任务，开始尝试低星任务；3-4 周继续完成某个专项团的任务，此时已经可以冲击一些高星任务了；5-8 周，渐臻佳境，对于某个专项团的任务已经能够做到游刃有余，并且尝试其它专项团任务。

CINN 后端 Pass 改造

在本次启航计划中，我一共完成 7 个 CINN 后端 Pass 改造任务。这里介绍一下这个专项团的收获。

背景

本次任务的背景是 CINN 升级了后端 IR 表示，将原来 Func-Expr 层级结构中的 Expr 进行了细化，重新划分为 Func-Block-Stmt-Expr，重新划分后的 IR 层次更加清晰。

与之对应地，后端 Pass 也被划分为 FuncPass、BlockPass、StmtPass、ExprPass 四个级别，使用配套的 PassManager 应用 Pass。其层次结构为：

此外，还提供了 IR 访问方法：

1. 类型不敏感的 Stmt 和 Block 级别的访问/修改方法，在遍历 Stmt 前后将会调用用户传入的回调方法：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

// Visitors
void Visit(const BlockRef &block,
           const std::function<void(const StmtRef &)> &pre_callback,
           const std::function<void(const StmtRef &)> &post_callback);

void Visit(const StmtRef &stmt,
           const std::function<void(const StmtRef &)> &pre_callback,
           const std::function<void(const StmtRef &)> &post_callback);
// Mutators
// ...

2. 类型敏感的 Stmt 和 Block 定制化访问模板类，用户可以通过重写 virtual StmtRetTy VisitStmt(const StmtRef &stmt, Args... args) 定制化访问不同的 Stmt：

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

template <typename StmtRetTy = void,
          typename BlockRetTy = void,
          typename... Args>
class StmtVisitor {
 public:
  virtual StmtRetTy VisitStmt(const StmtRef &stmt, Args... args) {
    CINN_CHECK_STMT_DEFINED(stmt)
    switch (stmt->stmt_type()) {
#define __(stmt__)                                \
  case ir::StmtNodeTy::stmt__:                    \
    return VisitStmt(stmt.as<stmt__>(), args...); \
    break;

      NODETY_FORALL_STMT(__)

      default:
        PADDLE_THROW(::common::errors::InvalidArgument(
            "Deadcode, not supported StmtNodeTy"));
#undef __
    }
  }
...

为什么要升级 IR？

从后端 Pass 的角度来看，IR 升级主要有两个好处：1. Pass 编写更加清晰和规范；2. Pass 便于管理。

旧 IR 下的的 Pass 大都通过继承 IRMutator/Visitor 在遍历整个 IR 的过程中修改来实现 Pass 的功能，但实际上其只需要针对某个特定类型的 Stmt/Block 处理即可。旧 IR 下的 IRMutator 为了便于开发者使用，提供了对各种类型的 Expr/Stmt/Block 默认遍历，例如对于 IfThenElse 默认实现版本会遍历条件和两个分支：

1
2
3
4
5
6
7
8
9

template <typename T>
void IRMutator<T>::Visit(const IfThenElse *expr, T op) {
  auto *node = op->template As<IfThenElse>();
  IRVisitorRequireReImpl<void, T>::Visit(&node->condition, &node->condition);
  IRVisitorRequireReImpl<void, T>::Visit(&node->true_case, &node->true_case);
  if (node->false_case.defined())
    IRVisitorRequireReImpl<void, T>::Visit(&node->false_case,
                                           &node->false_case);
}

这种默认实现在很多情况下是不必要的，比如在合并两个相同的 If 中，显然不需要对条件应用此 Pass，也不需要对 Expr 级别的表达式进行访问。

理论上说，开发者可以通过重写对应的 Visit 方法来及时进行截断，但一方面这样会使得 Pass 的代码比较臃肿，另一方面 Pass 在开发时并没有此规范，已经成为遗留问题。

在此次 IR 和 Pass 改造后，原有的 IRMutator 将只保留对于 Expr 级别的访问逻辑，对于 Stmt 和 Block 级别的遍历由 PassManager 完成。例如，StmtPassManager 将会遍历这个函数，并为每一条 Stmt 调用一次其管理的 StmtPass，而在 StmtPass 内部，其只需要处理符合其目标的逻辑。

此外，新版的 StmtVisitor 没有提供 VisitStmt 默认实现，这可以强迫开发者自定义遍历逻辑，并及时截断不需要的遍历。

Pass 编写范式

升级后的 IR 的编写范式一般为：1. 继承对应级别的 Pass 基类；2. 使用一个内部类对 Func/Block/Stmt 进行遍历实现核心功能，这个类可以继承 StmtMutator/IRMutator 或者调用 Visit/Mutate 方法来实现遍历；3. 返回 Success。

1. 继承对应级别的 Pass 基类
   第一步就是分析原 Pass 是什么级别 Pass，核心要义是抓住原 Pass 需要什么级别的信息以及是什么级别的修改。例如：

• IfFusionPass 是合并两个多个条件相同的 If，其要识别和删除多个 If，只有拿到这个 If 所在的 Block 能够实现多个语句的识别和单个语句的删除，这是一个 Block 级别的 Pass；
• RearrangeLoadInstruction 是将一个函数内的 Load 放到最前面执行，以提高指令的并行并行程度，为了确保 Load 是一个函数内最先执行的语句，以及对函数内所有相同的 Load 替换为本地变量，这是一个 Func 级别的 Pass；
• EliminateCommonFactorOfLocalIndex 需要获取当前 For 的嵌套信息，那必须由本 Pass 负责对 IR 的遍历，否则无法获取当前 For 的嵌套级别信息，因此这是一个 Func 级别的 Pass。

我的经验是：如果一个 Pass 仅仅需要当前 Stmt 的内部信息、不需要删除或者替换当前 Stmt、并且对于当前 Stmt 的嵌套级别没有要求（例如不要求当前的 For 是最外层/最内层的 For），那么其是一个 Stmt 级别的 Pass；如果一个 Pass 需要跨语句的信息，或者需要删除/替换/添加一条 Stmt，那么其是一个 BlockPass；如果一个 Pass 需要自己控制对 IR 的遍历过程，或者需要当前的嵌套上下文，那么这是一个 Func 级别的 Pass。

2. 编写实现类
   一些比较简单的 Pass 就是一个继承了 IRMutator 的实现类，此类 Pass 一般只需要额外继承 StmtMutator，如果不涉及 Expr 层面，则去掉对于 IRMutator，然后将原有逻辑迁移到新 IR 下即可。可参考 RemoveScheduleBlock。新 IR 下，很多变量都被设置为私有变量，必须通过 getter 和 setter 进行读写。

一些比较复杂的 Pass 可能有多个 Mutator 对 IR 进行多次访问，一般第一次是收集全局信息，之后再进行修改。读懂源码后再照葫芦画瓢修改即可。

更复杂的是调用了一些旧 IR 的方法，例如 ir::ir_utils::CollectIRNodesWithoutTensor，这种情况下可以判断一下传入的参数是否是 Expr，如果是 Expr 则还可以调用该方法（因为对 Expr 是封闭的，Expr 中不会有 Stmt 或者 Block），否则要根据这些方法的逻辑在新 IR 下进行实现。

3. 返回 Success
   这个没啥好说，返回 LogicalResult::success() 即可。

Tips

1. Pass 应该实现为无状态的
   无状态指的是 Pass 不应该依赖之前的信息，或者记录一些持久信息。例如，一个对于 For 进行处理的 Pass，其内部不应该记录当前 For 的名字以防止重复。如果想要避免重复访问，可以将其实现为 FuncPass 手动处理遍历逻辑。

2. PassManager 是按照 DFS 后序遍历的
   这一遍历顺序可以保证最内部的语句被最先访问。Pass 改造过程中是可以依赖这一行为的。

3. Pass 之间缺乏通信机制
   Pass 之间是缺乏通信机制的，一些 Pass 在应用是前是需要检查能否进行变换的，这些检查 Pass 可以作为变换 Pass 的内部的一部分，在变换 Pass 实例化一个 PassManager 应用检查的 Pass。

4. 可参考 Halide文档
   CINN 在很多设计上参考了 Halide 和 TVM，在如果碰到一些例如不知道 Stmt 的作用的疑问，可以参考这两个这两个文档更加丰富的社区，往往会有惊喜收获。

后记

作为第一次开源活动经历，我个人觉得还是收获颇丰的。纸上得来终觉浅，绝知此事要躬行，很多之前没有实操的技术都在这次活动中得到了锻炼，例如 Git 和 GitHub 的工作流、VSCode 和 CMake 的配套、GLOG 的使用等等，以及对于 CINN 中 Pass 改造的经验，更是很好的学习 AI Sys 的切入口。

墙裂安利没有尝试过的同学多多参加这类活动，一定能不虚此行！

Python 环境

需要用到如下 Python 环境：

• PyTorch GPU 版本
• jupyter notebook
• tqdm
• matplotlib
• torchprofile

数据集准备

Lab 1 中用到了 CIFAR-10 数据集，可以使用 https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz 直接下载，并将整个 cifar-10-batched-py 文件夹解压到 data/cifar10 文件夹内。

Part 1: Fine-grained Pruning

Question 1

除最后一层分类头外，其它层均服从均值为 0 的无偏正态分布，这意味着占很大比例的参数是可以被移除的，这为模型压缩留下了很大的空间。

Question 2

第二个问题要求实现细粒度剪枝，即可以对权重矩阵中的单个元素进行剪枝，关于不同颗粒度的剪枝介绍，见 课程第三讲笔记。

这里使用每个参数的绝对值来表示其重要性，剪掉不重要的参数，保留重要的参数。

本问比较简单，根据稀疏度计算出需要剪去的参数总量，然后使用找到阈值并根据阈值得到 mask 矩阵。唯一的一个注意点是计算 mask 矩阵是使用大于而不是大于等于，这是由于计算得到的阈值也需要被剪掉。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

##################### YOUR CODE STARTS HERE #####################
# Step 1: calculate the #zeros (please use round())
num_zeros = round(sparsity * num_elements)
# Step 2: calculate the importance of weight
importance = tensor.abs()
# Step 3: calculate the pruning threshold
threshold = torch.kthvalue(importance.view(-1), num_zeros).values
# Step 4: get binary mask (1 for nonzeros, 0 for zeros)
mask = importance > threshold
##################### YOUR CODE ENDS HERE #######################

Question 3

问题三要求我们在一个 5 x 5 的矩阵中保留 10 个元素，相应的稀疏度为 $1-\frac{10}{25}$，此问就算结束了。

Question 4

此问对 VGG 网络每一层进行了灵敏度分析，建议将步长修改为 0.2 或者 0.1，以获得更加平滑的灵敏度曲线。

从图中可以看到大部分层中，随着稀疏度的增加，模型精度相应变低，不同层的敏感程度不同，第 0 个卷积层对稀疏度最敏感。

Question 5

第 5 问中，要求根据前面灵敏度分析结果和模型参数计算量，设置每一层剪枝时的稀疏度。❗️注意，最终整个模型的稀疏度很大程度上取决于参数量比较大的层的稀疏度，对于参数量比较大的层，可以考虑设置比较高的稀疏度。

我选择的稀疏度参数为：

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

sparsity_dict = {
##################### YOUR CODE STARTS HERE #####################
    # please modify the sparsity value of each layer
    # please DO NOT modify the key of sparsity_dict
    'backbone.conv0.weight': 0,
    'backbone.conv1.weight': 0.6,
    'backbone.conv2.weight': 0.5,
    'backbone.conv3.weight': 0.5,
    'backbone.conv4.weight': 0.5,
    'backbone.conv5.weight': 0.6,
    'backbone.conv6.weight': 0.6,
    'backbone.conv7.weight': 0.75,
    'classifier.weight': 0
##################### YOUR CODE ENDS HERE #######################
}

经过剪枝后，大小约为原始稠密模型的 38.48%，精度从 92.9% 降低到了 91.50%，在 5 轮的微调后，模型精度恢复为 92.95%。

Part 2: Channel Pruning

Question 6

第 6 问需要实现 Channel Pruning，剪枝标准是只保留前 k 个通道。问题本身时简单的，用好 Python 中的切片即可：

 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

def get_num_channels_to_keep(channels: int, prune_ratio: float) -> int:
    """A function to calculate the number of layers to PRESERVE after pruning
    Note that preserve_rate = 1. - prune_ratio
    """
    ##################### YOUR CODE STARTS HERE #####################
    return int(round(channels * (1 - prune_ratio)))
    ##################### YOUR CODE ENDS HERE #####################

@torch.no_grad()
def channel_prune(model: nn.Module,
                  prune_ratio: Union[List, float]) -> nn.Module:
    """Apply channel pruning to each of the conv layer in the backbone
    Note that for prune_ratio, we can either provide a floating-point number,
    indicating that we use a uniform pruning rate for all layers, or a list of
    numbers to indicate per-layer pruning rate.
    """
    # sanity check of provided prune_ratio
    assert isinstance(prune_ratio, (float, list))
    
    n_conv = len([m for m in model.backbone if isinstance(m, nn.Conv2d)])
    # note that for the ratios, it affects the previous conv output and next
    # conv input, i.e., conv0 - ratio0 - conv1 - ratio1-...
    if isinstance(prune_ratio, list):
        assert len(prune_ratio) == n_conv - 1
    else:  # convert float to list
        prune_ratio = [prune_ratio] * (n_conv - 1)
    # we prune the convs in the backbone with a uniform ratio
    model = copy.deepcopy(model)  # prevent overwrite
    # we only apply pruning to the backbone features
    all_convs = [m for m in model.backbone if isinstance(m, nn.Conv2d)]
    all_bns = [m for m in model.backbone if isinstance(m, nn.BatchNorm2d)]
    # apply pruning. we naively keep the first k channels
    assert len(all_convs) == len(all_bns)
    for i_ratio, p_ratio in enumerate(prune_ratio):
        prev_conv = all_convs[i_ratio]
        prev_bn = all_bns[i_ratio]
        next_conv = all_convs[i_ratio + 1]
        original_channels = prev_conv.out_channels  # same as next_conv.in_channels
        n_keep = get_num_channels_to_keep(original_channels, p_ratio)

        # prune the output of the previous conv and bn
        prev_conv.weight.set_(prev_conv.weight.detach()[:n_keep])
        prev_bn.weight.set_(prev_bn.weight.detach()[:n_keep])
        prev_bn.bias.set_(prev_bn.bias.detach()[:n_keep])
        prev_bn.running_mean.set_(prev_bn.running_mean.detach()[:n_keep])
        prev_bn.running_var.set_(prev_bn.running_var.detach()[:n_keep])

        # prune the input of the next conv (hint: just one line of code)
        ##################### YOUR CODE STARTS HERE #####################
        next_conv.weight.set_(next_conv.weight.detach()[:, :n_keep])
        ##################### YOUR CODE ENDS HERE #####################

    return model

记得一提的是框架已经给出的代码，所谓 Channel 是在卷积中才会出现的，剪枝也是对输出通道进行剪枝。例如，当前卷积核中本来有 k 个通道输出，剪枝后变成 l 个输出通道，那么下一层的卷积核的输入通道也要相对应地从 k 变成 l。此外一般 Conv 后都会有一个 Batch Norm，应该这个 Conv 的 weight、bias、running_mean 和 running_var 也要一起进行剪枝。

Question 7

改进 Channel Pruning，使用 Frobenius 范数来评估每一个通道的重要程度。本问的核心就是 Frobenius 范数的计算，说明中推荐使用 torch.norm 进行实现，但是官网文档中提到这个 API 已经被弃用，这里改用 torch.linalg.vector_norm()。根据文档，dim 指定为需要展开为向量的维度，即第 [0, 2, 3]。

 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

# function to sort the channels from important to non-important
def get_input_channel_importance(weight):
    in_channels = weight.shape[1]
    # importances = []
    # # compute the importance for each input channel
    # for i_c in range(weight.shape[1]):
    #     channel_weight = weight.detach()[:, i_c]
    #     ##################### YOUR CODE STARTS HERE #####################
    #     importance = torch.linalg.norm(channel_weight, ord="fro", dim
    #     ##################### YOUR CODE ENDS HERE #####################
    #     importances.append(importance.view(1))
    # return torch.cat(importances)
    return torch.linalg.vector_norm(weight, ord=2, dim=(0, 2, 3))

@torch.no_grad()
def apply_channel_sorting(model):
    model = copy.deepcopy(model)  # do not modify the original model
    # fetch all the conv and bn layers from the backbone
    all_convs = [m for m in model.backbone if isinstance(m, nn.Conv2d)]
    all_bns = [m for m in model.backbone if isinstance(m, nn.BatchNorm2d)]
    # iterate through conv layers
    for i_conv in range(len(all_convs) - 1):
        # each channel sorting index, we need to apply it to:
        # - the output dimension of the previous conv
        # - the previous BN layer
        # - the input dimension of the next conv (we compute importance here)
        prev_conv = all_convs[i_conv]
        prev_bn = all_bns[i_conv]
        next_conv = all_convs[i_conv + 1]
        # note that we always compute the importance according to input channels
        importance = get_input_channel_importance(next_conv.weight)
        # sorting from large to small
        sort_idx = torch.argsort(importance, descending=True)

        # apply to previous conv and its following bn
        prev_conv.weight.copy_(torch.index_select(
            prev_conv.weight.detach(), 0, sort_idx))
        for tensor_name in ['weight', 'bias', 'running_mean', 'running_var']:
            tensor_to_apply = getattr(prev_bn, tensor_name)
            tensor_to_apply.copy_(
                torch.index_select(tensor_to_apply.detach(), 0, sort_idx)
            )

        # apply to the next conv input (hint: one line of code)
        ##################### YOUR CODE STARTS HERE #####################
        next_conv.weight.copy_(torch.index_select(next_conv.weight.detach(), 1, sort_idx))
        ##################### YOUR CODE ENDS HERE #####################

    return model

相比没有计算重要性的通道剪枝，改进剪枝后的模型的准确率从 28.15% 提升到 36.81。经过微调后恢复为 92.41%。

Question 8

1. 为什么剪枝 30% 但是计算量减少了大约 50%。
   VGG 模型主要由卷积层构成，卷积层的计算量 FLOPs 为：

其中输入和输出通道都变为原来的 70%，因而总计算量变为原来的 49%。

2. 解释一下为什么延迟（latency）的减少比例略小于计算量的减少比例。
   延迟不仅仅来源于计算，还来自于数据的搬运，这部分时间在没做算子融合的情况下减少并不显著。

Question 9

1. 讨论一下 fine-grained pruning 和 channel pruning 的优缺点。
   细粒度剪枝：压缩率更高、对硬件不友好、延迟高；
   通道剪枝：压缩率低、硬件友好、延迟低、易于微调。

2. 如果想在智能手机上加速模型，使用哪种方案更合适。
   通道剪枝。智能手机上一般缺乏对于稀疏矩阵的支持，选取对硬件更友好的方案。

小结

第一个 Lab 本身比较简单，做完能够建立起对于剪枝的初步认识，希望后面的实验能够上点强度，代码量也太少了😂。

本讲开始介绍量化技术，首先介绍各种数据表示格式，然后介绍了两种量化技术：K-means和线性量化，最后提到了模型压缩的流水线。

数值数据类型

课程第一部分介绍了整型、定点小数、浮点数的数据表示格式，属于计算机组成原理的基本知识，此处不再赘述。

值得一提的是，在浮点数表示法中，阶码长度意味着该表示法能够表示数值的范围，尾数长度则决定该表示法能够表示的数据精度。如下图所示，Google 提出了一种 BF16 表示格式，其总共占用 2 字节，但是与 IEEE 754 格式中的 4 字节的单精度浮点数具有相同的阶码长度，从而对齐二者的表示范围。

NVIDIA 提出了两种不同的 8 位浮点数的表示格式，它们具有不同的阶数与尾数的长度：

而如果仅仅使用 4 比特来表示整型或者浮点型，那能够表示出的数可以简单在数轴上点出：

量化

什么是量化

量化指的是将连续值或者大量的可能的离散取值近似为少量取值的过程。例如下图展示了对连续信号和对高清图片的量化效果。

量化前后的差别被定义为量化误差，量化过程需要最小化这一误差。

K-means 量化

如下图所示，基于 K-means 的量化算法通过 K-means 算法对所有权重进行聚类，每个权重使用其聚类后所在簇的簇中心值作为量化后的值，并在权重中使用簇号来指代。

假设被量化为 $2^n$ 个簇，即每个权重需要使用 $n$ 比特整型表示其簇号，原始权重共有 $m$ 个参数，且 $m»2^n$。在量化前权重矩阵内存占用为 $32m$ 比特，量化后为 $nm+32\times 2^n$ 比特。考虑到 $m»2^n$，量化后近似为 $nm$ 比特，量化后内存相当于量化前的 $n/32$。

量化后的模型通过微调可以取的更好的效果，K-means 量化的微调算法将同一簇的梯度相加，作为这个簇共同的梯度，并进行权重更新。

通过结合剪枝和量化两个策略，可以实现 20x 的模型压缩。

一般来说，先做剪枝，再进行量化。通过剪枝可以只保留有用的参数，减少量化的参数量。

量化位宽一般取决于不同的模型，如下图所示，一般卷积层需要 4 比特，全连接层仅需要 2 比特。

此外，还可以使用哈夫曼编码来进一步压缩模型的大小，为频数高的参数使用较短的编码。

线性量化

线性量化利用了一种将低精度整型通过仿射变换转换为浮点数的技术，如下图通过仿射变换可以将量化后的 2 比特权重张量近似还原为原始权重张量。

上述变换中有两个参数需要确定，一个是整型参数零点 $Z$ 和一个浮点型参数放缩因数 $S$，放射变换过程可以使用公式表示为：

其中，q 为量化后权重矩阵的值，r 为经过还原后权重矩阵的值。其中，量化矩阵中恰好为 Z 的元素将被还原为 0，而 S 决定了还原后的范围。

记 $r^\prime$ 表示权重矩阵的真实值，那真实参数表示的范围为 $r_{max}^\prime-r_{min}^\prime$，而经过还原的参数表示范围为 $(q_{max}-q_{max})\times S$，让二者相等可以得到放缩因数的计算公式：

在课程中，计算 $Z$ 的方法是让原始矩阵和还原后的权重矩阵的最小值对齐，即 $r_{min}^\prime = (q_{min} - Z)\times S$，从而推导出零点的计算公式：

这里有点奇怪，为什么不是让整个权重矩阵的总误差最小。

使用线性量化计算矩乘

在矩乘计算 $\mathbf{Y} = \mathbf{WX}$ 中，三个矩阵均使用线性量化进行表示，通过恒等变换可以得到：

其中，与 $W$ 相关的变量都是常量，此外 $Z_W$ 也是常量（why???)，因此 $Z_X \mathbf{q_W} + Z_W Z_X$ 可以在量化时预先计算得到，在运行时没有计算开销。

对于放缩因子 $\frac{S_W S_X}{S_Y}$，经验上可以确定其范围在 $(0,1)$ 之间，因此，其可以表示为一个定点小数通过右移运算得到，如下图所示。因此只需要存储一个定点小数和右移次数，而不需要存储一个高精度浮点数。

考虑到权重一般是关于 0 对称，因此可以有理由假定 $Z_W=0$， 从而消去 $Z_W\mathbf{q_X}$ 和 $Z_W Z_X$ 这两项，从而计算公式变形为：

主要计算开销在于 $\mathbf{q_w q_x}$ 这一低精度整数矩乘计算。

模型压缩的流水线

模型压缩全流程如下所示，首先通过剪枝和微调来减少无效参数，然后通过量化来减少模型参数表示，最后使用哈夫曼编码进一步压缩模型大小。

Again，在有着极高压缩率的同时，仍旧能够保持精度不变。

环境说明

本文使用 WSL Ubuntu 22.04 作为演示环境，VSCode 版本为 1.95.2，使用项目为 PaddlePaddle。

VSCode 中需要安装如下插件：

• clangd：配合 clangd server 实现对 C/C++ 的代码高亮、补全、跳转、重构等
• CMake Tools：为 CMake 项目的提供支持
• Python：为项目提供 Pythjon 支持，包括高亮、跳转、调试等

还需要安装如下工具：

• clangd：clangd sever，为 C/C++ 的代码解析提供支持

CMake、编译器、调试器等工具默认可用。

配置 CMake 项目

在 VSCode 中安装 CMake Tools 插件后第一次打开 CMake 项目，VSCode 默认会自动进行配置，即默认执行 CMake: Configure 命令。如果检测到多个编译器，会提示用户选择一个。此时 CMake 插件还没有做任何配置，这时候进行 Configure 大概率是不符合用户预期的，我们可以使用 ESC 退出 Configure 过程。

我们首先对 VSCode 插件进行配置。在 VSCode 中打开 settings-Workspace，在 workspace 修改的设置内容将以 .vscode/settings.json 文件的形式保存在项目根文件夹中。通过为每个项目保存不同的配置，可以方便地且“优雅”地在不同项目之间切换。依据下图，在设置中找到 CMake Tools 插件的设置。

其中有几项值得关注，可以根据自己需要进行修改：

• Build Directory：指定 CMake 构建目录路径
• Build Environment & Configure Environment：指定配置和构建阶段环境变量
• Build Args & Configure Args：指定配置和构建阶段额外命令行参数
• Cmake Path：指定 Cmake 可执行文件路径
• Generator：指定生成器，例如 Ninja

插件配置完成后，在 .vscode/settings.json 文件中就可以看到对应的修改：

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

{
    "cmake.configureArgs": [
        "-DPY_VERSION=3.12",
        "-DWITH_GPU=OFF",
        "-DWITH_TESETING=ON",
        "-DPYTHON_EXECUTABLE=/home/zhouxin/miniconda3/envs/paddle-dev/bin/python"
    ],
    "cmake.configureSettings": {
        "CMAKE_EXPORT_COMPILE_COMMANDS": true
    },
    "cmake.buildDirectory": "${workspaceFolder}/build_mask",
    "cmake.automaticReconfigure": false,
    "cmake.configureOnOpen": false,
    "cmake.configureOnEdit": false,
    "cmake.generator": "Ninja",
}

接下来可以为整个项目指定编译工具链，使用快捷键 Ctirl+Shift+P 唤起 VSCode 命令面板，并搜索 cmake: kits 可以找到 CMake: Edit User-Local CMake Kits，在该文件内，可以配置多个不同的编译工具链，相关说明见：vscode-cmake-tools/docs/kits.md at main · microsoft/vscode-cmake-tools · GitHub。然后使用命令 CMake: Select a Kit 选择为本项目选择一套合适的工具链。

一切就绪，接下来就可以对 CMake 项目进行 Configure 操作。在命令面板找到 CMake Configure，执行之。在 VSCode Output 面板中切换到 CMake，可以看到输出的日志，如果有错误，可以根据错误 Debug。

一般来说，在排除 CMakeList 文件本身出错和环境没准备妥当之后，大概率是某些环境变量的问题。可以在 .vscode/settings.json 设置文件中修改某些环境变量值，或者传入某些参数以指定某些工具的路径就可以解决。

Code Intelligence

Code Intelligence 指的是一系列语法高亮、代码跳转、自动补全、错误检测等等功能的集合，一言以蔽之，就是让 IDE 理解你的代码。对于没有配置好 Code Intelligence 的项目，随意打开一个文件，可能存在大量头文件找不到的报错，函数调用之间跳转基本也都是失败的，IDE¹ 无法定位到函数源码的位置。

而构建工具，例如 CMake，完全掌握着文件之间的依赖关系 ²。在 CMake 配置过程中，可以使用参数 -DCMAKE_EXPORT_COMPILE_COMMANDS=1 或者在插件设置中添加如下内容以要求 CMake 在配置过程在构建目录生成包含文件依赖信息的文件 compile_commands.json。

1
2
3

"cmake.configureSettings": {
	"CMAKE_EXPORT_COMPILE_COMMANDS": true
},

构建工具给出信息之后，还得告诉 clangd 这些“信息”的具体位置。在 .vscode/settings.json 文件中添加如下内容：

1
2
3
4
5
6
7
8

"clangd.arguments": [
    "--compile-commands-dir=${workspaceFolder}/build", // 指定编译信息所在目录
    "-j=20",                                        // 设置并行任务数为20
    "--background-index",                           // 启用后台索引
    "--pch-storage=memory",                         // 将预编译头存储在内存中
    "--limit-results=500",                          // 限制结果数量为500
    "--log=info"                                    // 设置日志级别为info
],

其中第一行指定 compile_commands.json 文件所在目录，默认为 CMake 的构建目录。其余为推荐使用的一些其它配置，加速大型项目的解析。

接下来重新执行 CMake: Configure，并在配置完成后重启 clangd 服务器 clang: restart language server。经过此番折腾，IDE 已经能够对大部分头文件进行解析，并正确实现文件中的跳转。

然而，在 Paddle 的项目中，仍有许多函数无法被正确解析。这是由于在 Paddle/third_party 文件中有众多第三方工具的源码。这些工具将在 CMake build 过程中被安装到构建目录之内，这些目前仍无法解析的函数依赖这些第三方工具。使用命令 CMake: Build 构建整个项目，完成之后所有函数就能够被正常解析了！

Debug

Debug 取决于具体的项目，在 VSCode 使用 .vscode/launch.json 对 Debug 进行配置。关于对 Python 和 CUDA/C/C++ 代码进行联合调试的内容，可以查看另一篇文章：在VSCode中对CUDA和Python代码进行联合调试 | 周鑫的个人博客。

测试

在 VSCode 中，使用 Test Explorer 对测试进行管理和配置 ³，VSCode 本身不提供特定语言的测试配置，而是以插件的形式扩展特定语言的测试支持。在插件中搜索 @category:"testing" 可以查看所有测试插件。Python 和 CMake Tools 插件似乎自带对 Python Test 和 CTest 的支持，不需要额外安装测试插件。

在 VSCode 的测试面板，可以看到所有测试项目。下图展示了来自 CMake 的 CTest 和来自 Python 的 Python Tests 项目，Paddle 的那个测试还没确定怎么来的。Anyway，能跑就行😎。

参考

MIT 6.824 Distributed Systems 第三讲学习笔记，简单介绍了存储系统和一致性，主要介绍了 GFS 中的文件读写流程。

存储系统概述

存储系统在分布式系统中相当重要：如果能够建立一个可靠的存储系统，可以讲其它应用构建为无状态的，而在存储系统中持久存储状态，这能够大量简化应用设计。这种情况下，应用即使崩溃也可以迅速重启，并从存储系统中读取状态进行恢复。

这也意味着，存储系统本身必须具备良好的容错能力，存储系统自身设计相当困难：

• 高性能：需要跨服务器共享数据
• 多服务器：这意味着某个服务器挂掉是常态
• 容错设计：必须有 replication 冗余备份
• 冗余备份：可能导致潜在的不一致性
• 强一致性：需要一致性协议和网络通信，进而降低性能

Ideal consistensy

理想的一致性系统应该表现得像一个单机系统。

在一致性系统中，涉及到并发问题。在分布式系统会出现一些额外的并发问题。例如，当 A 和 B 依次分别发送像 X 写入 1 和 2 的请求，C 和 D 再依次发送读取 X 的请求。CD 读到的数据组合可能为 (1, 1)、(1, 2)、(2, 2)，而不可能出现 (2, 1)。

而在分布式系统中，同样是上述读写次序，A 的写入请求可能只背服务器 S1 接收到、B 的写入请求只被 S2 接受接受到，而 C 的读取请求被 S2 处理，D 的读取请求被 S1 处理，那么 CD 读到的数据组合为 (2, 1)。这里我们需要一种同步协议来协调读者和写者。本课程的后半段将花费大量篇幅介绍不同的同步协议，他们基本都是在容错和一致性之间的折中。

GFS: Google File System

GFS 是本课程第一个案例分析，其是一个以高性能作为设计目标的分布式文件系统，其具备副本、容错、一致性。

特点

这篇论文发表于 21 世纪初，彼时学界对于分布式系统已经有较为成熟的理论体系，但缺少实际可用的工业界产品。GFS 作为一个成功的分布式文件系统，其实际上并非是一个标准的学术界研究的分布式系统：

• 单一 master
• 可能存在不一致性

GFS 的特点：

• Big：具有大的数据集
• Fast：自动将文件分片到多个服务器上
• global：所有应用都能看到相同的文件
• fault tolerance：自动容错和恢复机制

整体设计

GFS 架构图如下所示：

每个文件被分为多个 chunk，每个 chunk 不是很小，约为 64MB。应用告知 master 需要访问的文件名和 chunk 好，master 返回对应的句柄和存储位置，应用再根据存储位置找到对应的 chunk server，发送访存请求。

Master

Master 维护一组从文件名到 chunk 句柄数组的映射表。对于每一个 chunk 句柄，Master 维护其版本号、持有该句柄的服务器列表、这些服务器的主从次序、服务器的租约长度（lease time）信息。此外，Master 还要负责日志和检查点保存。Master 在响应请求之前，首先将其写入日志，这意味着即便 Master 挂了，也可以从日志中恢复重建，并响应请求。

文件名到 chunk 句柄数组的映射表应当作为持久状态被 Master 定期保存到磁盘上，而与 chunk 相关的服务器信息则不必要，该信息可以在 Master 重启时主动要求其它服务器汇报其持有的 chunk 信息从而重建。chunk 版本号也需要作为持久状态，因为 Master 需要明确了解整个系统每个句柄的最新版本号，而不是由其它服务器汇报，以应对那些真正持有最新版本号的服务器也许也一起挂掉的情况。

读写流程

• 读文件流程

1. 客户端将文件名和偏移量发送给 Master。
2. Master 告知客户端句柄、持有该句柄的服务器列表、版本号。
3. 客户端缓存上述信息。这一步可以减少 Master 的压力、网络流量和客户端自身访问文件的延迟。
4. 客户端按照由近及远的顺序依次尝试从服务器列表中获取文件。
5. chunk 服务器检查版本号，通过则将文件数据发送给客户端。

• 写入流程之 append 操作

1. 客户端告知 Master 文件名。
2. Master 根据文件名找到对应 chunk 句柄。Master 根据 chunk 句柄找到持有该句柄的 chunk 服务器列表。 如果没有主 chunk 服务器，Master 选择一个作为主 chunk 服务器，该主 chunk 服务器被授权一段时间的 lease 即租约，时间内其可以对该句柄进行修改。同时 Master 增加版本号，该版本号信息将分发给所有持有该句柄的服务器，句柄服务器应当将版本号持久保存在磁盘上。
3. Master 告知客户端文件主从 chunk 服务器列表和版本号。
4. 客户端将文件由近及远尝试发送给 chunk 服务器。
5. 如果是从 chunk 服务器接收到了写入内容，其再转发给主 chunk 服务器，主服务器转发给从服务器。如果一切正常，写入完毕后主服务器告知客户端成功写入。如果错误，则告知客户端错误原因。
6. 如果客户端接受到错误信息。其会自动重试，直至成功。以确保 at-least-once RPC 语义。
7. 在第二次重试中，主 chunk 服务器将不会在原有磁盘 offset 上写入，该位置已经存放了第一次写入的内容，其会向后寻找新位置。因此，磁盘上可能存有重复的记录。

本讲延续上一讲继续介绍了两种确定剪枝比例的算法：灵敏度分析和强化学习。此外还介绍了为稀疏网络提供支持的硬件加速器，包括 EIE、NVIDIA Tensor Core、TorchSparse 等。

如无另外说明，图片均引用自 EfficientML 课程幻灯片。

Lecture 4: Pruning and sparsity 剪枝和稀疏性

剪枝率

如下图所示，研究指出对每一层采取不均匀的剪枝比例的效果显著优于均匀剪枝，问题在于如何确定每一层的剪枝比例。

灵敏度分析

通过对每一层灵敏度进行分析，即对每一层按照不同剪枝率进行剪枝，观察其对最后精度的镜像程度来判断精度对于每一层的敏感程度。如下图所示，可以发现精度对于 L1 层最不敏感，对于 L0 层最敏感。

此外，当我们实现剪枝算法时，灵敏度曲线也可以用于检查算法是否实现有误：前期几乎不掉点，后期显著掉点。

进行灵敏度分析之后，可以确定一个能够接受的掉点阈值，根据此阈值确定每一层的剪枝比例。

这里隐含了一个假设：这些层之间是相互独立的，即我们没有考虑层一层之间的交互作用。

自动剪枝

之前聊过的剪枝方案都是由人工来确定剪枝策略的，但是这种方案不够优雅，并且不具备可扩展性 (scalability)。

这里介绍一种基于强化学习的剪枝比例确定方案，其动机是我们训练一个模型，输入为每一层的信息，输入为相对应的剪枝比例。

笔者缺乏强化学习领域相关知识，以下翻译可能不太恰当。

我们的模型 setup 为：

• 状态（输入）
  • 描述每一层的特征，包括层序号、通道数、卷积核大小、FLOPs…
• Action （应该是模型输出？）
  • 0 到 1 之间的数，表示剪枝率
• 智能体
  • DDPG agent（不懂）
• 奖励 （目标函数？）
  • -Error 如果满足约束
  • -inf 如果不满足

如下图所示，相比人类耗时地手工调优，基于强化学习的方法耗时更短并且性能更优。

微调剪枝后的神经网络

前文提到，剪枝完成后需要对网络进行微调，以恢复网络性能或者取得更好的效果。

• 学习率：一般设置为原先学习率的十分之一到百分之一
• 采用迭代剪枝：第二讲提到对迭代进行剪枝 - 微调的效果显著优于一次性的剪枝 - 微调操作
• 正则化：在微调过程中需要使用 L1 或者 L2 正则化

为稀疏性提供支持的系统

EIE: Efficient Inference Engine

EIE 使用了三种优化方式：

• 权重稀疏化：能够节省 90% 的计算资源和 80% 的内存资源，内存节省略少是由于保存稀疏信息带来的额外开销
• 激活层稀疏化：能够节省 66% 的计算资源
• 量化

对于如上的稀疏矩阵 - 向量乘法 (Sparse Matrix Vector Multiplication, SpMV)，上半图表示 SpMV 的逻辑形式，权重矩阵被染为四种颜色，每个颜色代表一个处理单元 PE 即 processing element，这些 PEs 将并行执行。下半图表示第一个 PE 中保存的权重信息，PE0 中的稀疏矩阵采取的是 CSC 格式 ¹，即在内存中 Virtual Weight 只保存 5 个非零元素，以及这五个非零元素与前一个非零元素之间的距离 Relative Index【这里是列优先的距离】，和指向列的行指针 Column Pointer。Pointer[i] 表示从 Virtual Weight 中第 Pointer[i] 开始是原稀疏矩阵中第 i 行的元素。

注意，上图并非标准 CSC 格式，在标准格式中，采用的 Row Index 来直接记录行号，而非使用 relative Index 记录相邻元素之间的距离。

在计算过程中，每个 PE 对输入向量 $\vec{a}$ 逐元素遍历，直接跳过零元；对于非零元素，将其广播到所有非零元 PE 中并计算想用结果。

整个 PE 单元的微架构图如下所示：

首先是 Activation Queue 用于存储所有非零激活层元素，然后根据非零元的 Activation Index 来获取 Weight Column Pointer，确定与激活层元素相对应的权重元素的起始和结束索引，接着获取权重元素，并对其进行解码（解码相关将在后文说明），最后将结果暂存累加，并经过 ReLU 层后输出。

得益于激活层和权重矩阵的稀疏性，可以将其放在 SRAM 中加速存取。

NVIDIA Tensor Core

第三讲提到的 M:N 的稀疏策略，即连续 N 个元素中必定有 M 个零元素。例如，2:4 的稀疏矩阵可以表示为：

对于 2:4 的且形状为 RxC 的稀疏矩阵，在内存中可以只保存所有非零元，即 R x C/2 的矩阵。还需要保存额外的元数据来确定非零元在原始矩阵中的位置，每个非零元可能来自其原来小组的 0~3 中的某个位置，因此每个非零元需要 4 bit 信息来记录其在小组中的原始位置。

上图右侧展示了 Tensor core 中稀疏矩乘的实现，以 A 的第一行为例，A 的第一行需要与 B 的第一列做内积，使用 A 在稀疏化过程中保存的 indices 信息，就可以对 B 的第一列进行筛选，只索引出 A 中非零元对应位置的元素，并进行乘法和累加操作。

TorchSparse

这里介绍一种稀疏卷积：如下图左所示，当对稀疏矩阵进行传统卷积操作时，经过卷积后激活层的稀疏度会下降，非零元素会逐渐扩散到周围的零元中；而在右侧的稀疏卷积中，要求保持输入中的稀疏性，即卷积后的零元仍旧是零元。

在稀疏卷积的实现中，对于单次卷积，只有非零元参加与卷积核的矩乘计算。这里介绍的实现方式为：建立输入元素 - 输出元素 - 卷积核权重的三元组，按照权重对三元组进行排序（实质上找到所有需要与该权重做乘法的元素）。

分组结束后，就可以采用一种自适应分组算法将上述标量乘法 - 加法转换为矩阵乘法（这一转换过程我直觉上认为类似于 img2col 技术将卷积转换为矩乘）。当然，这一过程存在大量开销。

• 计算规律性 (computation regularity) 与计算开销之间的折中
  为了将多个向量乘法聚合为一个矩阵乘法，我们不得不对某些较短的向量进行补零操作，这使得计算过程呈现出更好的规律性，但是带来了额外的计算开销。如下图所示，这里采取的折中手段是对每个计算 batch 进行动态分组
  

TorchSparse++

如下图左所示，邻接的行在一组中进行计算，以第一行为例，为了计算出 B1，在计算 B0 中必须将 $W_{0,-1}$ 和 $W_{1,0}$ 也参与计算，这就是前文提到的额外计算开销。在 TorchSparse++ 中，提出了一种行重排算法，经过重新分组后可以减少此类冗余计算。此外，还可以对权重进行分割，以实现更加细粒度的冗余优化。

PointAcc: 稀疏卷积的硬件加速器

这里介绍了前面一直在用的“输入元素 - 输出元素 - 卷积核权重”三元组的构建算法，该算法可以使用硬件加速，但没太理解算法原理，似乎也不是重点，略。

参考

MIT 6.824 Distributed Systems 第二讲学习笔记，包括 Go 语言和并发编程的简单介绍，以及对 RPC 和故障情况下的 RPC 语义的说明。

Lecture 2: RPC and Threads

为什么选择 Go 语言

• 对于线程和 RPC 的良好支持
• 拥有垃圾回收机制
• 类型安全
• 简单
• 可以编译

并行编程入门

在 Go 中，线程被称为 go routine，线程具有独立的 PC、栈和寄存器。线程之间共享内存地址，处于同一片内存空间内。

线程还可以被视为由 Go 运行时提供 start、exit、stop、resume 操作的一组结构。

为什么要使用线程

在本课程中，使用线程是基于并发的考虑：

• IO 并发：当一个线程发起例如网络 IO 请求时，在等待回复时可以将该线程阻塞，并调度其它线程；
• 多核并行：在多核处理器上，可以在不同的核心上运行不同的线程，显著提升吞吐量；
• 方便：在实验中可能需要周期性地执行某一任务，go routine 可以轻松地实现这一功能。

线程中的挑战

• 条件竞争：程序执行结果取决于线程之间具体的执行顺序
  • 规避方案一：线程之间不要共享变量
  • 规避方案二：使用锁
  • Go 中具有一个竞争检测器，可以检测出许多潜在的竞争行为
• 协调（似乎是指同步）
  • 通道
  • 条件变量
• 死锁

Go 的解决方案

针对上述挑战，Go 中提供了对应的解决方案：

• 通道（不共享内存）
• 锁和条件变量

RPC 远程过程调用

• 目标
  • 将远程调用透明化，使之表现得像在进行本地调用。所谓过程调用就是指调用可执行代码，例如函数、方法、子程序等等。
• 实现
  • 在本地将有一个与远程调用对应的函数桩 stub，stub 记录了远程调用的必要信息，例如函数签名，其通过网络将调用信息发送给服务器。
  • 在服务器也有一个类似的函数桩，负责对通过网络接收到的调用信息进行解码，并调用服务器端函数实现，返回计算结果。
  • 本地 stub 接收到服务器返回的结果后，进行解码并返回给本地调用者。
  • 上述过程全部由编译器和运行时负责实现，对于程序员而言是透明的。

故障情况下的 RPC 语义/行为

故障情况下 RPC 可以有多种不多的行为：

• at least once
  • 如果服务器挂了，客户端将不断试直至至少执行一次
• at most once
  • 最多执行一次
  • 服务器需要确保不对重复的请求多次执行
  • Go 采用的模式
  • 使用 TCP 协议来实现最多发送一次 RPC
• exact once
  • 恰好一次
  • 这比较难以实现

MIT 6.5940 EfficientML 第三讲学习笔记，主要介绍剪枝的定义、效果和粗细程度，并详细介绍了多种剪枝标准。

Lecture 3: Pruning and sparsity 剪枝和稀疏性

剪枝的动机

在上一讲提到，内存操作的代价相当昂贵，因此为了加速模型的运行，一个思路就是减少模型中一切内存的占用，包括减小模型大小、减小激活层大小和数量。

剪枝定义

剪枝的数学定义如下所示：

具体来说，剪枝指的是移除神经网络中的某一些参数，使得神经网络成为一个相对稀疏网络。

剪枝效果

如下图所示，经过剪枝后的模型，其性能损失并不显著；相反，如果在剪枝后对模型进行微调，甚至迭代剪枝和微调的步骤，可以实现仅使用几分之一的参数量达到相同或者更优的准确率。

这一节的效果多少有点震撼到我：经过剪枝后，模型的参数量可以减少 90%，并且不掉点，甚至性能更优！

剪枝细粒程度

根据剪枝实施的细粒程度，可以对其进行分类。

• 细粒度剪枝
  细粒度剪枝允许对任意单个权重进行剪枝，其优点是剪枝程度更高，确定是参数张量作为稀疏矩阵，在硬件上的加速更加难以实现。

• 粗粒度剪枝
  粗粒度剪枝只能都对参数矩阵中的某一行进行全部剪枝，缺点是不够灵活，优点是经过剪枝后的参数仍旧是一个非稀疏矩阵（剪枝后的矩阵变得更小），易于加速。
  

前面举的是全连接层的例子，接下来讨论卷积操作的剪枝。卷积层的参数张量形状为 $[c_o, c_i, k_h, k_w]$。由于其具有四个维度，因此其剪枝的细粒程度种类多得多：

第一种方式是剪枝最细的，往往可以实现数量级的剪枝，然而其一般需要在特定硬件上才能取得比较好的推理速度，在 GPU 上设计算法比较困难。

第二种方式是基于模式的，采用固定的模式对卷积核进行剪枝。一种典型的模式是 N:M，即在连续 M 个参数中选择 N 的进行剪枝。

最后一种方式是对整个通道进行剪枝，其优点在于经过剪枝后仍旧是通用矩乘，缺点是剪枝率很比较低。

剪枝标准

• 基于大小的剪枝
  一种符合直觉、并且效果很好的剪枝方式是根据参数的绝对值来决定其重要性，剪掉那些绝对值很小的参数。如果要进行按行剪枝，则每一行的权重对于这一行向量的 L1 Norm，即这一行向量的绝对值的和。

• 基于缩放参数的剪枝
  例如在卷积的每一层中，有多个通道，可以给每个通道配置一个可学习的参数，作为该通道的放缩参数。通过训练来确定每个通道放缩参数，即该通道的重要性，并进行剪枝。

• 基于二阶的剪枝
  我们可以使用泰勒展开来表示经过剪枝后模型的误差：

其中，$g_i = \frac{\partial L}{\partial w_i}$，$h_{i,j} = \frac{\partial ^2 L}{\partial w_i \partial w_j}$。

基于二阶的剪枝方法假设：
损失函数是近似二次的，因此高阶误差项可以忽视；
神经网络在训练过程中已经收敛，因此 L 对 w 的一阶导为 0；
对不同参数的剪枝操作引发的误差是彼此独立的，因此 $\frac{1}{2} \sum_{i\neq j} h_{ij} \delta w_i \delta w_j$ 为 0。

因此，剪枝后模型的误差为：

为了最小化剪枝后的误差，应当保留权重更大的参数，因此重要性表示为：

其中，$h_{ii}$ 是 Hessian 矩阵。

• 对激活层进行剪枝
  对激活层进行剪枝其本质上就是粗粒度的权重剪枝。如下图所示，如果我们需要移除激活层的某个节点，其进行的操作就是在 FC 网络中移除权重举证的某一行，或者在卷积中移除某个通道。
  

• 基于 0 概率的剪枝
  ReLU 层的输出的激活层有概率为零，通过统计 batch 中每个位置的输出为 0 的频率，可以对 0 概率高的位置进行剪枝。

注意，这里是对激活层进行剪枝，而非之前提到的直接对参数进行剪枝。我的理解是，对激活层进行剪枝相当于直接对这激活层对应的两层网络的参数进行剪枝。

• 基于回归的剪枝
  如果直接评估整个模型剪枝前后的误差，这一代价可能很高。基于回归的剪枝对网络逐层评估和剪枝。

下图展示了对全连接层输入通道的剪枝算法示意图。首先将矩乘结果视为多个通道外积结果的和，为每个通道设置一个缩放系数 $\beta_{c}$，缩放系数越接近 0 说明该通道越不重要。

优化算法为：首先固定权重，对 $\beta$ 进行优化，对系数小的参数进行剪枝；然后固定 $\beta$，对 $W$ 进行优化。还可以对上述过程进行重复和迭代。

Lecture 2: Basics of neural networks 神经网络基础

神经网络

• 基本术语
  如下图所示，我们使用术语 Synapses（突触？）、权重、参数来指代网络中的参数，使用术语神经元、特征、激活层来指代网络中每一层的计算结果。
  

模型的宽度指的是隐藏层的维度，对于相同的参数量，宽而浅的模型相比窄而深的模型计算效率更高，因为其核函数调用次数更少，并且能够充分进行并行计算。然而后者在准确率上往往表现得更好，这需要进行折中。

• 全连接层
  全连接层是对输入进行加权求和并加上偏执项，如下所示：
  

• 2D 卷积
  

基本术语这一节略，大多是 DL 的入门知识。

神经网络效率的评价指标

• latency 延迟
  神经网络完成某一特定任务需要的时间。延迟主要取决于计算时间和内存时间，即：

内存操作消耗的资源比计算多的多的多，如下图所示 32 位的访存操作消耗的能量是常见其它计算操作的数百倍。

• throuthput 吞吐量
  单位时间内神经网络能够处理的数据量。

由于 batch 的存在，延迟很高的模型可能能够同时处理多个 batch，因此高延迟≠低吞吐量。在移动设备上更关心延迟，而在高性能设备上更关心吞吐量。

• number of parameters 参数量
  不同层的参数量计算公式如下图所示：
  
• model size 模型大小
  模型大小取决于参数量与参数字长，即：

• total / peak number of activations 激活层总量和峰值大小
  在模型推理中，通常瓶颈在于激活层大小而非参数量，如下图所示，相同性能的不同模型，模型参数可以进行大幅度优化，但参数量几乎差不多。而激活层的峰值大小则决定了这个模型在推理过程中消耗的最大内存数。
  
  在模型训练中，激活层更是瓶颈，其大小是参数量的数倍。
  
  在 CNN 的训练中，由于前期图像分辨率高，瓶颈在于激活层；后期特征通道数较高，瓶颈在于权重。
  

• MAC
  MAC 操作指的是 multiply-accumulate，即一次 $a\leftarrow a+b\cdot c$ 操作，在 GPU 上其通常可以被翻译为一条指令。

在如下所示的举证向量乘法即 MV 中，MAC 数为 $m\cdot n$，即每个结果都需要 $n$ 次 MAC 操作。 而在矩乘即 GEMM 中，MAC 数为 $m\cdot n\cdot k$，即每个结果需要 $k$ 次 MAC 操作。

不同层的 MAC 数如下所示：

• FLOP 浮点操作数
  一个 MAC 相当于两个 FLOP

• FLOPs 每秒浮点计算次数

• OP 操作数
  模型并不一定总是用浮点数进行表示和计算，对于非浮点计算，我们称之为 OP。

• OPs 每秒操作计算次数

算法原理

约定和前置知识

在优先级爬升法中，中缀表达式被分解为主表达式（primary expression）和运算符（operator），例如在表达式 a+b*c-d 中，主表达式为 ['a', 'b', 'c', 'd']，运算符包括 ['+', '*', '-']，每个运算符都有与之对应的优先级和结合性，优先级使用正整数表示相对大小，四则运算中乘除优先级高于加减，均为左结合。在本例中，约定加减的优先级为 10，乘除的为 20。

中缀表达式可以被解析为 表达式树，表达式树能够反映出计算的优先级。所谓运算符优先级解析本质上就是要解析出正确的表达式树。

本文仅限于对双目运算符进行讨论，即每个运算符接收两个输入参数，输出一个计算结果。

原理

以表达式 a+b*c-d 为例，当我们自左向右扫描到第一个运算符 + 时，由于 b*c 的优先级更高，因此不能直接在表达式树上构造出 a+b，而应该优先构造出 b*c。那优先级高的是不是一定要被优先计算呢？非也，例如表达式 a+b+c*d，即便在计算 c*d 之前优先计算了 a+b，也并不会妨碍我们构造出正确的表达式树。

那当我们扫描到一个运算符时，什么情况可以在表达式树上构造对应节点，什么时候需要先计算优先级更高的节点呢？我们知道，在中缀表达式中，每个主元周围最多有有两个运算符，主元需要与优先级更高的那个运算符进行结合，因此，当我们扫描的运算符时，我们可以先解析出下一个主元，以及下一个运算符，如果当前运算符的优先级高于下一个运算符（或者优先级一致，但是当前运算符是左结合的），那么说明下一个主元是当前运算符的第二个输入参数 rhs，否则说明下一个主元是下一个运算符的第一个输入参数 lhs，需要先将下一个运算符解析完毕，解析结果才是当前运算符的 rhs。

例如，在表达式 a+b*c-d 中，首先扫描到第一个主元 a，将其记录为 rhs；扫描到第一个运算符 +，继续向后扫描到一个主元 b，以及与其邻接的运算符 *，* 的优先级高于 +，因此开辟一个新的函数栈，lhs = b，对 * 进行解析；向后扫描到一个主元 c，以及与其邻接的运算符 -，* 的优先级高于 -，因此 c 是 * 的 rhs，构造出 b*c 并返回该结果；回到解析第一个 + 的函数栈中，其接收到 b*c 作为新的 rhs，并继续扫描下一个运算符 -，+ 的优先级与 - 一致，且 + 是左结合的，因此构造出 a+b*c，并作为新的 lhs；继续扫描到下一个主元 d，并且不存在下一个运算符，则继续构造出 a+b*c-d，并结束解析。

算法伪代码如下所示 ¹：

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

parse_expression()
    return parse_expression_1(parse_primary(), 0)

parse_expression_1(lhs, min_precedence)
    _lookahead_ := peek next token
    while _lookahead_ is a binary operator whose precedence is >= _min_precedence_
        _op_ := _lookahead_
        advance to next token
        _rhs_ := _parse_primary_ ()
        _lookahead_ := peek next token
        while _lookahead_ is a binary operator whose precedence is greater
                 than _op_'s, or a right-associative operator
                 whose precedence is equal to _op'_s
            _rhs_ := _parse_expression_1_ (_rhs_, precedence of _op_ + (1 if _lookahead_ precedence is greater, else 0))
            _lookahead_ := peek next token
        _lhs_ := the result of applying _op_ with operands _lhs_ and _rhs_
    return _lhs_

其中 parse_expression_1 是一个递归调用的函数，其接收两个参数 lhs 和 min_precedence，功能是对中缀表达式进行顺序解析，直至碰到优先级低于 min_precedence 返回解析结果。

而在 parse_expression_1 内部，其如果碰到了可以直接计算的情形，即当前运算符的优先级更高，或者相等且为左结合，则将 lhs __0p__ rhs 的结果作为新的 lhs，并进行下一趟外部循环；如果需要先计算下一个右结合，则递归调用自身，并将 min_precedence 参数设置为当前运算符的优先级 +0/1。

+0 是为了正确处理右结合的情况，当下一个运算符与当前运算符优先级相当且右结合时，min_precedence 需要设置为与当前运算符优先级相等的值，以确保递归调用时碰到 1^2^3（右结合运算）时能够继续递归解析，而非像左结合一样直接返回；+1 是为了正确处理左结合的情况，以确保当碰到与当前优先级相同的运算符时其能够及时返回，而不是继续向后解析。

括号处理

上述算法碰到括号就扑街了：一方面，括号不是双目运算符，无法在小修的情况在融入我们的算法；另一方面，括号拥有着最高的优先级，意味着我们需要对其进行特殊处理。好在 LLVM 的 tutorial 中提供了另一种解决思路：括号内的内容本身就是一个表达式，调用 parse_expression() 函数对括号内的内容进行解析即可，并将解析结果当作我们算法中的一个主元即可。

具体实现见下一章。

实现

本章主要讲解 LLVM tutorial 对优先级爬升法的实现，在教程中实现了对 Kaleidoscope 语言的词法和语法分析，该语言支持基本四则运算。在第二章中需要代码解析为抽象语法树 AST，首先介绍几个脚手架，用于进行词法分析。

ParsePrimary 用于解析前文中提到的主表达式，主表达式可能是标识符、数字、括号包围的表达式：

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

std::unique_ptr<ExprAST> ParsePrimary(){
    switch (CurTok) {
    default:
        return LogError("Unknown token when expecting an expression");
    case tok_identifier:
        return ParseIdentifierExpr();
    case tok_number:
        return ParseNumberExpr();
    case '(':
        return ParseParenExpr();
    }
}

运算符的优先级使用 map 来记录：

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

std::map<char, int> BinopPredence;

int GetTokPrecidence() {
    if(BinopPredence.empty())
        InstallBinopPredence();
    if (!isascii(CurTok)){
        return -1;
    }

    int TokPrec = BinopPredence[CurTok];
    if (TokPrec <= 0)   return -1;
    return TokPrec;
}

void InstallBinopPredence(){
    // 1 is the lowest predence
    BinopPredence['<'] = 10;
    BinopPredence['>'] = 10;
    BinopPredence['+'] = 20;
    BinopPredence['-'] = 20;
    BinopPredence['*'] = 30;
    BinopPredence['/'] = 30;
}