Lorn

现代GPU编程指南（开篇）：Hopper 新特性总览

Thu, 27 Aug 2026 09:00:00 +0800

这是「现代GPU编程指南」系列的开篇。之所以想单独开一个系列，是因为 Hopper（H100 / H800，计算能力 sm_90）相比 Ampere 并不只是"更大更快"，而是引入了一批需要改变编程范式才能吃满的新硬件特性。本篇先做一个总览，后续每篇再展开一个专题，并尽量用 GEMM / FlashAttention 这类真实算子把它们串起来。

Hopper 带来了什么

1. Thread Block Cluster 与分布式共享内存（DSMEM）

Hopper 在 grid → block 之间新增了一层 Thread Block Cluster。同一个 cluster 内的若干 block 会被调度到同一个 GPC 内的相邻 SM 上，从而可以：

通过 分布式共享内存（Distributed Shared Memory） 直接读写同 cluster 内其它 block 的 shared memory；
使用 cluster 级别的同步原语（cluster barrier）。

这让"比一个 block 更大、又比走 global memory 更快"的协作成为可能，对 tiling 策略的设计影响很大。

2. TMA：Tensor Memory Accelerator

TMA 是一个专用的异步批量拷贝引擎，负责 global memory 与 shared memory 之间的多维 tile 搬运。它的价值在于：

用一个 copy descriptor 描述多维张量的搬运，一条指令搬一整块，不再需要大量线程手写地址计算；
天然异步，把访存和计算重叠起来，同时省下寄存器和线程去做真正的计算。

3. WGMMA：warpgroup 级异步 Tensor Core

Hopper 的 Tensor Core 通过 WGMMA（warpgroup MMA） 以 warpgroup（128 线程） 为粒度、以异步方式发射矩阵乘累加，并原生支持 FP8（e4m3 / e5m2）。相比 Ampere 的 mma，它的吞吐更高，也更依赖异步流水线来喂数据。

4. 异步事务屏障（mbarrier）与流水线

要把 TMA 的异步搬运和 WGMMA 的异步计算真正重叠起来，就需要 异步事务屏障（mbarrier / async transaction barrier） 来构建生产者–消费者流水线：一边 TMA 往 shared memory 里灌数据、累加 transaction count，另一边 WGMMA 等待到齐后开算。

5. 其它值得关注的点

FP8 与 Transformer Engine：配合 WGMMA 的 FP8，训练/推理里低精度 GEMM 成为一等公民；
setmaxnreg：warpgroup 之间可以动态重新分配寄存器，让"搬运 warpgroup"少占、“计算 warpgroup"多占；
更大的 shared memory（每 SM 可配置到约 228KB）与 DPX 指令（加速动态规划类算子）。

本系列规划

后续计划按下面的顺序展开（会随写随调）：

Thread Block Cluster 与分布式共享内存
TMA 异步数据搬运
WGMMA 与异步 Tensor Core
用 mbarrier 构建生产者–消费者流水线
FP8 与 Transformer Engine 实战

最后会用一个从 naive 到 Hopper-optimized 的 GEMM / FlashAttention 案例，把上面这些特性组合起来，看看它们各自贡献了多少性能。

本篇为系列总览，具体的代码与实验会放在后续各篇。

GPU 体系结构 —— H100 的硬件架构

Thu, 27 Aug 2026 10:00:00 +0800

💡 建议搭配阅读：H100 硬件架构可视化演示 —— 一个交互式网页，完整演示了 H100 的硬件架构。阅读本文前可以先花几分钟浏览，快速建立对 H100 硬件的整体认识。

GPU 的基本组成原理

一个 GPU 芯片可以被划分为以下几个主要组成部分：

NVIDIA GigaThread Engine（巨线程引擎）
8 个 GPC（Graphics Processing Cluster，图形处理集群）
HBM3 Stacks（HBM3 显存堆栈）
HBM3 Memory Controllers（HBM3 内存控制器）
PCI Express 5.0 Host Interface（PCIe 5.0 主机接口）
NVLink Switches、Ports、Hub（NVLink 交换机、端口和 Hub）
L2 Cache Slices（L2 缓存分片）

NVIDIA GigaThread Engine

GigaThread Engine 译作巨线程引擎，这是 GPU 上负责调度 Kernel 以及 Thread Block 的硬件。

当一个 Kernel 启动后，GigaThread Engine 会负责将 Kernel 中的 Thread Block（CTA）分配给各个 SM，并跟踪每个 CTA 当前的状态，如哪些 CTA 还没开始执行、哪些 CTA 正在执行、哪些 CTA 执行完毕。当某个 SM 有足够的资源、可以容纳另一个 CTA 时，GigaThread Engine 以及相关的前端调度逻辑就会把下一个 CTA 分给这个 SM。

同时它还会处理和强制执行一些资源约束，如 Occupancy 限制、寄存器资源控制等。

在 Hopper 架构的显卡上，GigaThread Engine 还能够理解和处理 Cluster 相关的资源调度。

HBM3

HBM3 是一种 off-chip 的设备内存，也就是我们常说的 Global Memory 的物理介质，在 H100 上其容量为 80GB，理论带宽为 3.35TB/s。

其组织结构大致如下：

HBM3 Stack：在 H100 上总共有 6 个 HBM3 Stack，但是出于良率考虑，只使用了其中的 5 个。
总线宽度共 5120 bit（每个 HBM3 Stack 的位宽为 1024 bit）。
通过 10 个独立的 512-bit memory controller 连接 HBM（每个 HBM3 Stack 两个）。

通常来说，我们一次从 HBM 中请求数据的流程如下：

首先从 SM 进入 L1D 缓存或合并器，若没有相关数据则进入 L2 缓存，若还没有则通过 Memory Partition / Crossbar 来确定这次请求应该由哪一个 memory controller 处理，最后由 Memory Controller 在 HBM3 Stack 中取出对应的 cache line 返回数据并写入 L1D Cache。

其中需要强调的是，Memory Controller 负责把 GPU 的内存访问请求转换成符合 HBM/DRAM 协议和时序的实际读写操作，并进行调度以尽可能提高内存带宽。

L2 Cache

在 H100 上有 50MB 的 L2 Cache，它们被划分为了两个 25MB 的分区。由于这种分区的特性，同一个 GPC 中的 SM 与其直接相连的 L2 Cache 分区之间会具有更近的访问路径。因此，当我们的内存访问命中这些相邻的 L2 Cache 分区中的 Cache Line 时，可以获得更低的有效访问延迟和更高的带宽。

L2 Cache 采用 128 字节的 Cache Line，以及 32 字节的 Sector。也就是说一次内存访问最多可以访问一个 128 字节 Cache Line 的 1-4 个 sector。如果内存访问的 Coalescing（合并）做得不好，就可能导致每个请求涉及的 sector 数量急剧增加，从而降低内存的访问效率。

GPC

GPC 是一组 SM（18 个）。每个 GPC 与物理上邻近的 L2 Cache 分区之间访问延迟更低（即前文提到的亲和性），但任何 SM 都可以访问全部的 L2 Cache。当需要把数据取到 shared memory 时，数据路径是 HBM → L2 → L1。GPC 还能让各 SM 之间共享 DSMEM（Distributed Shared Memory，详见文末的 Thread Block Cluster）。

GPC 中有 TPC，每个 TPC 包含两个 SM。TPC 主要是硬件物理实现层面的分组概念，对 CUDA 编程模型基本不可见。

SM

如下图所示是一个 SM，它是 GPU 的基本执行单元，负责执行 CUDA Kernel 的 thread block。其基本由 4 个相同的子部分组成，并且这四个子部分我们称作 SMSP 或者 quadrant。

整体上，每个 SM 由 4 个 SMSP 组成，每个 SM 都包含 L1 指令缓存以及大小为 256KB 的数据缓存/shared memory、TMA，与共享的 texture unit。

而每个子部分包含的内容如下：

FP32 CUDA Core

它们是通用图形计算或 compute shading 的主力单元。H100 的每个 SM 上有 128 个 FP32 CUDA Core。因此总共有 16896 个活跃的 CUDA Core（132 个 SM）。

SFU

SFU 通常负责复杂数学函数的计算，如正余弦、对数、指数、平方根、倒数这类。每个 SM 有 16 个 SFU，也就是每个 SMSP 有 4 个。每个 SFU 每 cycle 可以完成 1 个线程的运算，所以每个 SM 的 SFU 吞吐是 16 个线程/cycle——也就是说一条 warp 级（32 线程）的 SFU 指令需要 2 个 cycle 才能执行完。也因此，如果一个 warp 里所有 thread 都在调用这些复杂的数学函数，SFU 就会成为瓶颈。

INT unit

这是标准的整数 ALU，用来做内存寻址、循环控制以及一般的整数运算。每个 SM 有 64 个 int unit。这些 unit 与 CUDA Core 以及其他计算 core 是可以并行执行的，也就是说 GPU 可以一边计算内存地址（INT），一边处理数据（FP），互不 stall。

FP64 Unit

H100 有独立于 FP32 的专用 FP64 core。每个 SM 有 64 个 FP64 Core。

一般用于科学计算等场景。

Tensor Core

Tensor Core 是专门加速矩阵乘加运算（D = A × B + C）的单元，深度学习中的绝大多数计算都发生在它上面。H100 的每个 SM 配有 4 个 Tensor Core（每个 SMSP 1 个），全卡共 528 个。

H100 搭载的是第四代 Tensor Core：新增了对 FP8 精度的支持，并引入了 WGMMA 指令——矩阵乘可以异步执行，且操作数可以直接来自 Shared Memory。TMA 把数据异步搬进 Shared Memory，WGMMA 再异步地从 Shared Memory 取数计算，两者配合就构成了 Hopper 异步流水线的核心。（Tensor Core 的详细机制会在后续章节单独介绍。）

Load/Store Unit（LSU）

它负责执行每个线程的内存指令如 store、load、atomic 等。

每个 SMSP 有 8 个专用的 LSU，每个 SM 一共 32 个，都是直连 L1、L2 缓存的。

当 warp 执行 ld/st/atom 时，LSU 会把 32 个 thread 的地址做 coalescing，拼成 cache-line / sector 请求，再去查 L1 data cache。命中就很快回到 register；miss 则继续往 L2 / DRAM 走，回来时也可能填回 L1。

访问能 coalesced 时，LSU 会把这个 warp 的 32 个地址合并成尽量少的 cache-line，发给 L1/L2 的请求就更少，replay / stall 减少，有效带宽也更高。

这里需要注意的是 TMA 会直接绕过整个 L1 Cache，进入 L2 Cache 或 DRAM 读取数据并将其放入 Smem 中。

Unified Shared Memory + L1 Cache

每个 SM 共 256 KB，全卡合计带宽约 33 TB/s，分成 32 个 bank，每个 bank 宽 32 bit（4 byte）。其中 shared memory 大小是 228 KB。

无冲突时，load / store 各 128 B/cycle。TMA 异步拷贝能打到接近峰值带宽，并和计算重叠。Bank 按连续的 4-byte word 交错分布在相邻 bank 上。（对于 bank 的理解，可以想象成体育场的门）

L1 cache 相当于 coalescing buffer：把 warp 要的数据收齐，再高效送出去。

Cache line = 128 byte，Sector = 32 byte，cache 可以按 sector 逐块填充。

每个 SM 可配置的 shared memory 上限是 228 KB。

每个 thread block 上限是 227 KB，因为 CUDA 预留了 1 KB。

Registers

每个 thread 都有一套私有的片上 register。带宽最高、延迟最低，每个 SM 256 KB。每个 SM 每 cycle 最多 128 次读/写。

Register 经常是瓶颈。如果 kernel 每个 thread 用 R 个 register，那最多常驻 thread 数大约是 floor(65536 / R)（再按 block 粒度和其它限制取整/封顶）。到 128 register/thread 时，每个 block 最多 512 个 thread。

访问 register file 比访问 shared memory 还快 30–50 倍，比访问 HBM3 快 1000 倍以上。

编译器 register 不够用时，会把数据 spill 到 local memory，比 register 慢很多。CUDA 13.0 加了优先 spill 到 shared memory 的支持，只有 shared memory 也不够时才退回 local memory。

哪怕只多几个（比如 63 → 65 register/thread），常驻 warp 数也可能掉下去，因为分配会按内部粒度取整。活跃 warp 变少，性能就下来。

Warp

我们通常称 32 个线程组成的一组是一个 warp，它们从同一个 thread block 里被创建，每个 cycle 共享同一个 warp scheduler 调度，各自私有 register，但执行相同的指令流。

而 Warp scheduler 通常是负责指令发射的单元，每个 cycle 选出就绪的 warp 来发射指令，并处理 warp 级别的分支和 divergence。维护一个 scoreboard，记录哪些 warp 真正就绪。发射前确保源 register 已经从 register file 读出，或已从 pipeline 转发过来。

所谓 divergence（分支发散），是指同一个 warp 内的 32 个线程走进了不同的分支（例如 if/else 的两侧）。由于整个 warp 共享同一条指令流，硬件无法让两部分线程同时走两条路径，只能把各分支串行执行：先执行 if 分支，同时用 active mask 屏蔽走 else 的线程，然后再反过来。原本一条指令的工作变成了多条，warp 内的分支越碎，性能损失就越大——因此写 kernel 时应尽量避免 warp 内的分支发散。

为什么 GPU 要让尽可能多的 warp 常驻在 SM 上？ 答案是为了隐藏内存访问的延迟（Latency Hiding）。一次 HBM 访问需要数百个 cycle，如果 SM 上只有少量 warp，一旦它们全部陷入内存等待，计算单元就只能空转。而常驻 warp 足够多时，Warp Scheduler 每个 cycle 都能挑出一个已就绪的 warp 发射指令——某个 warp 发起内存请求后转入等待，它空出的执行槽立刻被其他就绪的 warp 顶上，内存延迟就这样被"藏"在了别人的计算之下。这也是前文反复强调的 Occupancy（由 register、shared memory 用量共同决定）对性能如此关键的根本原因。

动手实验

本章的四个硬件概念（Coalescing、Occupancy 与 Latency Hiding、Bank Conflict、SFU 吞吐）各配了一个可编译运行的小实验，放在 Assignment 01。读完本章建议亲手跑一遍——纸上的硬件数字变成实测数据，理解会深很多。

H100 的重大变革

H100 拥有 80GB 的 HBM3 内存，内存带宽达到 3.35 TB/s。

拥有 132 个 SM，每个 SM 拥有 4 个 Tensor Core，总共有 528 个 Tensor Core。

Hopper 架构最大的特性在于引入了异步特性：

TMA：每个 SM 有一个 TMA 单元，这是一个用于将 Tensor 数据拷贝操作从 SM 中卸载出去的硬件单元。在先前的架构中，我们通常需要让每个线程执行相应的计算、遍历数据区域、并发射多条指令，才能完成 Global Memory 到 Shared Memory 的数据传输。而有了 TMA 之后，我们可以通过 TMA Descriptor 来异步地执行这些操作。只需要一个线程发起整个数据搬运操作，剩下的数据移动工作由硬件在后台完成。
- 同步拷贝：如果一个拷贝是同步的，那么发起这个拷贝操作的线程在拷贝完成前就无法继续执行接下来的指令了，也就是说当控制流重新回到这个线程时，这个数据拷贝操作已经完成了。
- 异步拷贝：如果一个拷贝是异步的，那么发起拷贝操作仅仅意味着启动了数据传输，发起者（这个线程）可以在数据传输进行时继续执行其他工作，之后在真正需要使用这些数据之前，再去等待或检查拷贝是否完成。这种方式比同步拷贝更加高效，能够让我们有效地掩盖内存访问带来的延迟。（例如在处理第 i 块数据时，预取第 i+1 块数据）
第四代 Tensor Core：FP8 的支持（1979 TFLOPs），WGMMA
Thread Block Cluster：Hopper 在 Thread Block 和 Grid 之间引入的新层级。同一个 Cluster 内的多个 Thread Block 会被硬件保证共同调度到同一个 GPC 内的不同 SM 上执行，这些 Block 之间可以通过 Distributed Shared Memory（DSMEM）直接读写彼此的 Shared Memory，并在 Cluster 范围内做同步。这让跨 SM 的线程协作第一次有了硬件级的支持，而不必再绕道 Global Memory。（更详细的内容会在后续章节介绍。）

[CS336] Assignment 2：Systems

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

Assignment 1

在上一个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 来轻松计算此值。

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

[从零开始的编译原理][实验] Exp Chapter 2：从 Koopa IR 到 RISC-V

Tue, 06 Oct 2026 20:24:45 +0800

接下来，在本节中，我们将在上一节工作的基础上，让编译器进一步支持生成 RISC-V 汇编代码。

目前我们需要编译的 SysY 程序仍然非常简单：

1
2
3
4
5


int main()
{
 // 摊牌了，我是注释
 return 0;
}

上一节中，我们已经可以将它编译为如下的 Koopa IR：

1
2
3
4


fun @main(): i32 {
%entry:
 ret 0
}

而本节的目标，是继续完成从 Koopa IR 到 RISC-V 汇编的转换。例如生成：

1
2
3
4
5


 .text
 .globl main
main:
 li a0, 0
 ret

当然，根据具体的代码生成策略，也可能生成：

1
2
3
4
5
6


 .text
 .globl main
main:
 li t0, 0
 mv a0, t0
 ret

两种写法在这里的效果是一样的：最终都将整数 0 放入返回值寄存器 a0 中，然后从 main 函数返回。

具体生成哪一种形式，取决于我们的代码生成器如何处理中间结果和寄存器。对于当前这个非常简单的程序来说，显然没有必要额外使用临时寄存器，因此我们直接生成第一种形式即可。

目标代码生成

目前，我们的编译器已经完成了这样一条编译链：

1
2
3
4
5
6
7


SysY 源代码
 ↓
Lexer / Parser
 ↓
AST
 ↓
Koopa IR

接下来需要继续完成：

1
2
3
4
5


Koopa IR
 ↓
CodeGen
 ↓
RISC-V Assembly

上一章中，我们已经自己设计并实现了 Koopa IR 的数据结构。因此，这里的目标代码生成与之前实现 Printer 输出 Koopa IR 文本其实非常类似：

遍历 IR 数据结构，根据不同的 IR 节点输出对应的 RISC-V 指令。

例如：

1

ret 0

对应：

1
2


li a0, 0
ret

因此，从整体结构来看，CodeGen 同样可以采用 Visitor 风格的实现：依次访问 Program、Function、BasicBlock 和 Instruction，最终将每条 IR 指令翻译成对应的汇编指令。

不过，在开始实现之前，我们首先需要理解生成出来的这几行 RISC-V 汇编究竟在做什么。

对于：

1
2
3
4
5


 .text
 .globl main
main:
 li a0, 0
 ret

整体上主要完成了三件事情：

定义函数 main 的入口；
将函数返回值 0 放入规定的返回值寄存器；
执行函数返回。

因此，我们首先需要回答三个问题。

如何定义一个函数？

在高级语言中，我们习惯把函数理解成：

1
2
3
4


int main()
{
 return 0;
}

但从处理器的视角来看，函数本质上只是一段连续或逻辑关联的指令序列。

调用函数时，处理器跳转到这段指令序列的入口开始执行；函数执行结束时，再通过返回指令跳回调用者。

因此，在汇编层面，并不存在类似 C/C++ 中：

1

int main()

这样的“函数定义语法”。

我们真正需要做的，只是给这段指令的入口位置定义一个符号：

1

main:

这里的 main 是一个 label（标签）。汇编器会把它关联到当前位置对应的地址。之后，无论是链接器还是其他代码，都可以通过符号 main 找到这段代码的入口。

同时：

1

.globl main

用于声明 main 是一个全局符号，使这个符号可以被链接器看到。

而：

1

.text

则表示接下来的内容属于 .text 段，也就是程序的代码段。

因此：

1
2
3


 .text
 .globl main
main:

可以简单理解成：

接下来是一段代码，并且我们定义了一个可以被外部找到的函数入口 main。

至于函数如何返回，并不是函数标签本身负责的，而是函数内部生成的指令负责的。

如何在 RISC-V 中设置返回值？

函数调用不仅涉及“跳转到哪里”，还涉及调用者和被调用者之间如何传递参数、返回值，以及哪些寄存器需要保存等问题。

这些规则由 RISC-V Calling Convention（调用约定）规定。

对于整数返回值，RISC-V 使用：

作为返回值寄存器。

在 RV32 中，一个通用寄存器的宽度为 32 bit，因此 a0 和 a1 可以用于传递最多两个 XLEN 宽度的整数返回值。

不过对于我们现在的程序：

1
2
3
4


int main()
{
 return 0;
}

只有一个 int 类型返回值，因此只需要：

1

a0

也就是说：

1

return 0;

最终需要确保：

1

a0 = 0

随后再执行函数返回即可。

如何将整数加载到寄存器中？

接下来的问题就是：如何让：

1

a0 = 0

RISC-V 汇编器提供了一个非常方便的伪指令：

1

li rd, imm

其中：

rd 表示目标寄存器；
imm 表示立即数。

例如：

1

li a0, 0

表示：

将立即数 0 加载到寄存器 a0 中。

需要注意的是，li 是一个伪指令（pseudo-instruction），它并不一定直接对应某一条真实的 RISC-V 机器指令。

汇编器会根据立即数的大小，将它展开为一条或多条真正的 RISC-V 指令。例如对于较小的立即数：

1

li a0, 0

可以被展开成类似：

1

addi a0, zero, 0

至于具体如何展开，则交给汇编器处理即可。对于编译器前端和目前的简单代码生成器而言，我们直接生成 li 会方便很多。

函数最后的：

1

ret

同样也是一个伪指令。

它表示：

从当前函数返回到调用者。

其底层实际上可以展开为类似：

1

jalr zero, 0(ra)

其中 ra 保存了函数调用完成后应该返回的位置。

因此，我们生成：

1
2
3
4
5


 .text
 .globl main
main:
 li a0, 0
 ret

其含义可以完整地理解为：

1
2
3
4
5
6


 .text # 接下来的内容放入代码段
 .globl main # 将 main 声明为全局符号

main: # main 函数的入口
 li a0, 0 # 将返回值 0 放入 a0
 ret # 返回调用者

这就完成了：

1
2
3
4


int main()
{
 return 0;
}

最基本的目标代码生成。

接下来，我们就可以仿照上一节 Printer 的实现方式，设计一个 CodeGen 类：

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


class CodeGen
{
public:
 std::string Generate(const Program& program) const;

private:
 void Visit(const Function& func, std::ostream& out) const;
 void Visit(
 const BasicBlock& block,
 const std::string& func_name,
 std::ostream& out
 ) const;

 void Visit(const Instruction& inst, std::ostream& out) const;
 void Visit(const Return& ret, std::ostream& out) const;

 void LoadOperand(
 const IRValue& value,
 const char* reg,
 std::ostream& out
 ) const;
};

它的整体结构和 Printer 非常相似。

我们仍然按照：

1
2
3
4
5
6
7


Program
 ↓
Function
 ↓
BasicBlock
 ↓
Instruction

这样的层次遍历 IR。

区别在于，之前的 Printer 只是把 IR 数据结构重新打印为 Koopa IR 文本，例如：

1

ret 0

而现在的 CodeGen 需要真正考虑目标机器的语义。

例如面对：

1

ret 0

我们不能简单地输出：

1

ret 0

而需要知道 RISC-V 的调用约定规定返回值应当放在 a0 中，因此要生成：

1
2


li a0, 0
ret

也就是说：

Printer 主要关心 IR 的语法，而 CodeGen 开始需要关心目标架构的 ABI、寄存器和指令。

因此，我们额外定义：

1

LoadOperand(...)

负责将一个 IR 操作数加载到指定寄存器。

目前我们的 IR 非常简单，操作数只有整数立即数，因此 LoadOperand 暂时只需要处理：

1

IRInteger

未来随着 IR 中出现临时变量、二元表达式、局部变量等内容，这个函数才会逐渐涉及寄存器和栈上的数据。

具体实现如下：

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


std::string CodeGen::Generate(const Program& program) const
{
 std::ostringstream out;

 out << " .text\n";

 for (const auto& func : program.funcs)
 {
 Visit(*func, out);
 }

 return out.str();
}

Generate 是整个目标代码生成过程的入口。

首先：

1

out << " .text\n";

声明接下来生成的是代码段。

然后依次遍历程序中的所有函数：

1
2
3
4


for (const auto& func : program.funcs)
{
 Visit(*func, out);
}

虽然目前我们的程序中只有一个：

1

main

但是从数据结构和代码生成器的设计上，我们仍然按照“一个程序可以包含多个函数”的方式实现。

接下来处理函数：

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


void CodeGen::Visit(const Function& func, std::ostream& out) const
{
 out << " .globl " << func.name << "\n";
 out << func.name << ":\n";

 for (const auto& block : func.bbs)
 {
 Visit(*block, func.name, out);
 }

 out << "\n";
}

对于每个函数，首先生成：

1
2


.globl main
main:

分别对应：

1
2


out << " .globl " << func.name << "\n";
out << func.name << ":\n";

随后遍历函数中的所有基本块：

1
2
3
4


for (const auto& block : func.bbs)
{
 Visit(*block, func.name, out);
}

目前我们的 main 只有一个基本块：

1

%entry

并且不存在跳转，因此还不需要真的为 %entry 生成汇编标签。

后面当程序出现：

1
2


if
while

等控制流结构时，一个函数中会出现多个基本块。此时我们就需要为不同的基本块生成标签，例如：

1
2
3
4


main_entry:
main_then:
main_else:
main_end:

这里将 func.name 一并传给 Visit(BasicBlock)，也正是为了给后续生成基本块标签预留上下文。

目前基本块本身只需要继续遍历其中的指令：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


void CodeGen::Visit(
 const BasicBlock& bb,
 const std::string& func_name,
 std::ostream& out
) const
{
 for (const auto& inst : bb.insts)
 {
 Visit(*inst, out);
 }
}

接下来根据指令类型进行分派：

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


void CodeGen::Visit(const Instruction& inst, std::ostream& out) const
{
 switch (inst.kind)
 {
 case ValueKind::Return:
 {
 Visit(static_cast<const Return&>(inst), out);
 return;
 }

 default:
 {
 throw std::runtime_error(
 "unsupported instruction kind: " +
 std::to_string(static_cast<int>(inst.kind))
 );
 }
 }
}

目前我们的 IR 中唯一可能出现的指令就是：

1

Return

因此只需要处理：

1

ValueKind::Return

即可。

对于：

1

ret 0

对应的 IR 大致可以理解为：

1
2
3


Return
└── operand
 └── Integer(0)

因此在处理 Return 时：

1
2
3
4
5


void CodeGen::Visit(const Return& ret, std::ostream& out) const
{
 LoadOperand(*ret.operand, "a0", out);
 out << " ret\n";
}

我们首先调用：

1

LoadOperand(*ret.operand, "a0", out);

含义就是：

将 return 的操作数加载到返回值寄存器 a0。

随后：

1

out << " ret\n";

生成函数返回指令。

最后实现 LoadOperand：

 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


void CodeGen::LoadOperand(
 const IRValue& value,
 const char* reg,
 std::ostream& out
) const
{
 switch (value.kind)
 {
 case ValueKind::Integer:
 {
 const auto& integer =
 static_cast<const IRInteger&>(value);

 out << " li " << reg << ", "
 << integer.value << "\n";

 return;
 }

 default:
 {
 throw std::runtime_error(
 "unsupported value kind: " +
 std::to_string(static_cast<int>(value.kind))
 );
 }
 }
}

目前 LoadOperand 只需要处理整数立即数。

例如：

1

IRInteger(0)

并要求加载到：

1

a0

最终就会生成：

1

li a0, 0

于是整个调用过程就是：

1
2
3
4
5
6
7
8
9


Generate(Program)
 ↓
Visit(Function main)
 ↓
Visit(BasicBlock entry)
 ↓
Visit(Return)
 ↓
LoadOperand(Integer(0), a0)

最终输出：

1
2
3
4
5


 .text
 .globl main
main:
 li a0, 0
 ret

至此，我们的编译器已经第一次完成了一条完整的编译链：

1
2
3
4
5
6
7


SysY
 ↓
AST
 ↓
Koopa IR
 ↓
RISC-V Assembly

虽然目前支持的程序只有最简单的：

1
2
3
4


int main()
{
 return 0;
}

但这里建立起来的 CodeGen 框架之后基本可以继续沿用。

随着后面的 SysY 语法逐渐复杂，我们主要需要不断扩展：

1

Visit(...)

来支持新的 IR 指令，并逐步完善：

1

LoadOperand(...)

对临时变量、寄存器和栈上数据的处理。

[从零开始的编译原理][理论] Chapter 2：语法分析

Sat, 03 Oct 2026 15:32:40 +0800

这一节我们主要学习的是语法分析。

语法分析

在进行完词法分析后，我们得到了一组Token序列，而语法分析则用来分析这些token之间的结构关系。

例如我们有一个表达式:

1

a = b + 3 * 4;

经过词法分析后，我们可以得到:

1

ID(a) = ID(b) + NUM(3) * NUM(4);

那么它应该被解析为:

1

b + (3 * 4)

对应这个AST:

1
2
3
4
5


 +
 / \
 b *
 / \
 3 4

语法分析器同样有三种实现方法：

语法分析器生成器
手写语法分析器
自动化语法分析器

我们按照同样的顺序，介绍如何用语法分析器生成器来进行语法分析，以及自动化语法分析器背后的基本原理。

语法分析器生成器

在本节中，我们继续用antlr来设计一个类C 语言的grammer: Cymbol.g4( github实现参考，我们下面的可能有细微的不同),这一次，我们重点关心其中的语法部分。

在生成语法分析器的同时，我们还会利用这个分析器去抽取函数调用图.

Cymbol 的语法规则实现

接下来我们先来尝试用antlr来描述Cymbol.g4这个语言的语法结构，我们采用从上到下的描述顺序。

首先，整个程序，我们可以认为是有若干个变量声明和函数声明构成的，因此我们可以写出第一条语法规则:

1

prog : (varDecl | functionDecl)* EOF ;

那么接下来我们就需要去写变量声明和函数声明的语法:

变量声明写作:

1
2


varDecl : type ID ('=' expr)? ';' ;
type : 'int' | 'double' | 'void' ;

其中括号中的内容是变量初始值声明，这是一个可选的，因此我们括号后面用？描述

函数声明则写作:

1
2
3


functionDecl : type ID '(' formalParameters ? ')' block ;
formalParameters : formalParameter (',' formalParameter)* ;
formalParameter : type ID ;

其中，formalParameters还可以写作递归形式:

1
2
3


formalParameters : formalParameters ',' formalParameter
 | formalParameter
 ;

接下来我们看看函数体block的语法规则:

1
2
3
4
5
6
7
8


block : '{' stat* '}' ;
stat : block
 | varDecl
 | 'if' expr 'then' stat ('else' stat)?
 | 'return' expr? ';'
 | expr '=' expr ';'
 | expr ';'
 ;

其中 stat定义的第一行中调用了block，而block又调用了stat，这是一种互递归。

接下来我们重点看看一个expr也就是表达式的语法规则应该是怎样的:

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


expr : ID '(' exprList ? ')' # 表达函数调用
 | expr '[' expr ']' # 表达下标
 | '-' expr
 | '!' expr
 | expr '^' expr
 | expr ('*' | '/') expr
 | expr ('+' | '-') expr
 | expr ('==' | '!=') expr
 | '(' expr ')'
 | ID
 | INT;
exprList : expr (',' expr)* ;

一些问题

在上一小节中，我们初步完成了Cymbol语言的语法规则定义，但是上述语法规则会存在一定的问题:

1

if a then if b then c else d

我们现在用上面的语法规则去解析这个语句，发现有两种解析方法:

1
2


if a then [if b then c] else d
if a then [if b then c else d]

这就是语言的二义性，在设计语法分析器时，我们应该消除这种二义性。

这种二义性叫做Dangling Else二义性

一种消除的方式是改写上述匹配规则:

简单来说，我们规定else总是和最近的，尚未匹配else的if匹配。

为了体现这种规则，我们把语句分为两类: matched和open,那么此时文法可以写作:

1
2
3
4
5
6
7
8
9


stat : matched_stat | open_stat;

matched_stat : 'if' expr 'then' matched_stat 'else' matched_stat
 | expr
 ;
 
open_stat : 'if' expr 'then' stat
 | 'if' expr 'then' matched_stat 'else' open_stat
 ;

在这种文法下，我们上面的例子，就只会匹配第二种解析方法了。

而在antlr这样的语法分析器生成器中，则巧妙的用最前优先匹配原则解决了这个问题。

除了悬空的Else带来的二义性外，运算符的结合性也有可能带来二义性:

1
2
3
4


expr : expr '*' expr
 | expr '-' expr
 | DIGIT
 ;

例如1-2-3,此时可能是(1-2)-3也可以被识别为1-(2-3)这实际上是一个左结合，右结合的问题，由于大部分的运算符都是左递归的，因此像antlr这样的语法分析器生成器都是默认左结合的，那自然就会带来一个问题，那右递归的运算符怎么办？
一般来说，右递归的运算符大多都是前缀运算符或后缀运算符，这种运算符只有一种匹配规则，因此不太需要考虑二义性，而真正需要考虑二义性的就是像^这样的右递归运算符。

对于这一类需要右结合的运算符，在antlr中，我们可以手动指定结合的方式:

1
2
3
4


expr : '!' expr
 | <assoc = right> expr '^' expr
 | DIGIT
 ;

其他语法分析器生成器也存在类似的手动指定左/右结合的功能。

利用语法分析器得到函数调用图

通过语法分析器，我们最终可以得到一个AST，而antlr提供了一个接口:ParseTreeWalker,它可以以DFS的方式自动遍历整个AST，并且我们可以利用Listener来负责监听进入，退出节点的事件。

假设我们要得到函数调用图，那么一个必要的事情就是得到每个函数的名称。那利用Listener就很好实现这一点了，我们只需要每次监听到事件:functionDecl Enter就记录下节点的函数名称就能获取所有的函数名称了。

得到了所有函数的名称之后，我们还需要得到函数之间的调用关系，根据我们的语法规则，我们可以知道，函数调用语句会被语法规则:

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


expr : ID '(' exprList ? ')'
 | expr '[' expr ']'
 | '-' expr
 | '!' expr
 | expr '^' expr
 | expr ('*' | '/') expr
 | expr ('+' | '-') expr
 | expr ('==' | '!=') expr
 | '(' expr ')'
 | ID
 | INT;
exprList : expr (',' expr)* ;

识别，因此我们只需要识别到函数调用事件之后，由调用函数向被调用函数连一条边即可。但是比较麻烦的点在于，除了函数调用，expr还有多条解释规则，而Listener并不会对上述规则作区分，我们能看到的只有expr Enter以及expr Exit,要在此基础上识别出函数调用，我们还需要补充大量的if语句做判断，这显然不是我们希望的。

事实上，在antlr中，是支持对每条规则用#来加标签的:

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


expr : ID '(' exprList ? ')' # functionCall
 | expr '[' expr ']' 
 | '-' expr
 | '!' expr
 | expr '^' expr
 | expr ('*' | '/') expr
 | expr ('+' | '-') expr
 | expr ('==' | '!=') expr
 | '(' expr ')'
 | ID
 | INT;
exprList : expr (',' expr)* ;

这样，我们就可以通过Listener得到事件:functionCall Enter/Exit了。

语法分析的基本原理（1） CFG

我们先来看文法的组成:

1

functionDecl : type ID '(' formalParameters? ')' block;

每个文法规则都由:分隔为两部分，其中前面部分我们称为头部(Head),后者为规则体(Body)。整个规则我们称作产生式(Production).

每个产生式的头部都是非终结符。而所有没有出现在头部的符号，我们称作终结符，他们对应于我们词法分析器中产生的词法单元。

这种文法，我们称作上下文无关文法（Context-Free Grammar CFG）。

一个上下文无关文法通常写作:

$$ G = (V,T,P,S) $$

分别代表:

V：非终结符
T：终结符
P：产生式
S：开始符号

用数学符号表示就是:

$$ A\in N\to \alpha \in (T \cup N)^* $$

接下来我们介绍一下CFG的语义:

这里会涉及到几个基本概念:

推导：推导顾名思义就是用产生式对终结符进行替换，例如:

我们用规则:

$$ E \to E + E | E * E | (E) | -E | id $$

从E得到字符串：-(id + id)

那么推导的流程我们可以写作 :

$$ E \to -E \to -(E)\to-(E+E)\to -(id+E)\to-(id + id) $$

这其中还会涉及到一个最左推导和最右推导的概念，二者的区别在于，前者在推导时总是选择最左侧的非终结符进行推导，而后者则是选择最右侧的非终结符。

另外，我们会标记:

$E \Rightarrow -E $：经过一步推导得出
$E\xRightarrow{+}-(id+E)$：经过一步或多步推导得出
$E\xRightarrow{*}-(id+E)$：经过零步或多步推导得出

句型: 如果 $S \xRightarrow{*} \alpha,\alpha \in (T\cup N)^*$,则称 $\alpha$是文法G的一个句型。
句子：如果 $S \xRightarrow{*} \omega,\omega \in T*$,则称 $\omega$为文法G的一个句子
文法G的语言L(G)是它能推导出的所有句子构成的集合 : $L(G) = \{\omega | S\xRightarrow{*} \omega\}$

关于文法G，我们主要关心两个主要问题:

Membership 问题，即给定一个字符串，该字符串是否属于该文法产生的语言L(G)？
L(G)究竟是什么？

其中，第一个问题就是编译器语法分析器的任务，为输入的词法单元流寻找推导，构建语法分析树或者报错。

而第二个问题则是程序设计者需要考虑的问题。

这个问题通常体现在这样的题目上: 请给出文法，满足: $\{x \in \{a,b\}^*| x中a,b数目相同\}$

一个可能的文法是: $V\to VV|aVb|bVa|\epsilon$

下面简单证明一下为什么？

我们可以先证明文法生成的串一定满足a,b数目相同:

我们记: $\#_a(x)$表示串x中字符a的数目。

我们证明采用归纳法。

$V \Rightarrow \epsilon$,显然，此时满足 $\#_a(\epsilon) = \#_b(\epsilon)=0$
若 $V\xRightarrow{*} x$,且x中字符a和b的数目相同，那么有: $V\Rightarrow aVb\xRightarrow{*}axb$,同时增加一个a和一个b，所以仍然相同，同理bVa以及VV仍然成立。

所以我们有: $L(G)\subseteq L$

接下来证明所有a,b数目相同的字符串都能被这个文法生成

对串长 $|x|$做归纳。

若$|x|=0$，则 $x=\epsilon$，有 $V \Rightarrow \epsilon$

假设所有长度小于n，且有: $\#_a(x) = \#_b(x)$的串都能由V生成。

考虑长度为n的串$x = x_1x_2\dots x_n$,且有 $\#_a(x) = \#_b(x)$。

首先讨论首尾字符不同的串，比如 x=ayb,由于整个x中a,b数目相同，去掉一个a和一个b后，y中仍然满足:

$\#_a(y) = \#_b(y)$.由归纳假设 $V \xRightarrow{*} y$.

因此 $V \Rightarrow aVb \xRightarrow{*}ayb=x$,若x=bya，同理。

然后考虑首尾字符相同的串，我们假设首尾字符都是a，且定义前缀的差值:

$$ d(k) = \#_a(x_1\dots x_k) - \#_b(x_1\dots x_k) $$

因为第一个字符是a所以有: $d(1)=1$,而最后一个字符为a，因此 $d(n-1)=-1$,而 $d(k)$每读一个字符只会变化+1或者-1。因此从 $d(1)=1$到 $d(n-1)=-1$的过程中肯定存在某个k满足: $d(k)=0$.

于是我们可以把字符串x分为:$x=yz$,其中y，z都非空，且分别满足:

$$ \#_a(y) = \#_b(y),\#_a(z)=\#_b(z) $$

又因为: $|y|,|z| < |x|$,由归纳假设我们知道:

$$ V\xRightarrow{*} y,V\xRightarrow{*}z $$

因此使用 $V \Rightarrow VV$,有:

$$ V\Rightarrow VV\xRightarrow{*}yz=x $$

对于首尾为b的串同理。因此，我们证明了:

$$ L \subseteq L(G) $$

综上

$$ L(G) = \{x\in\{a,b\}^*|\#_a(x)=\#_b(x)\} $$

语法分析的基本原理（2） LL

接下来我们重点考虑一下和我们语法分析器背后原理相关的问题：
Membership 问题，即给定一个字符串，该字符串是否属于该文法产生的语言L(G)

我们语法分析要做的事情，无非就是根据词法单元流构建语法分析树。在构建时，我们有两种思路，自顶向下和自底向上。

我们先来介绍第一种: 自顶向下的，递归下降的，基于预测分析表的，适用于LL(1)文法的LL(1)语法分析器。

其中自顶向下就是说我们在构建语法分析树时，是从根节点（文法的起始符号）往叶节点（词法单元）构建的。而中间每个中间节点表示对某个非终结符应用某个产生式进行推导。

而递归下降就是说，我们会为每个非终结符写一个递归函数，内部按需调用其他非终结符对应的递归函数，下降一层。

我们以下面这个文法的匹配来展示一下递归下降的过程:

$S \to F,S\to (S+F),F\to a$,匹配的文本是:((a+a)+a)

我们从起始符号S开始，首先第一步显然是匹配文法: $S\to (S+F)$

接下来遍历目前的词法单元，对于终结符直接跳过，否则进行递归展开，比如我们遇到第一个S时，将其展开: $S\to (S+ F)$,而第一个F时， $F\to a$,从而得到串: $((S+F)+a)$

按照上一步的流程，不断展开非终结符，最终匹配((a+a)+a)

那么在语法分析的过程中，自然就会产生一些问题，应该选择哪个终结符进行展开？应该选择采用哪个产生式进行推导？

对于LL(1)语法分析器而言，我们在推导的每一步都是选择最左边的非终结符进行展开。

而对于第二个问题，我们在LL(1)语法分析器中，通常借助预测分析表确定:

在上面的例子中，我们会通过某种算法得到如下预测分析表:

	(	)	a	+	$
S	2		1
F			3

这张表指明了每个非终结符在面对不同的词法单元或文件结束符时，该选择哪个产生式或者报错。

有了这张表，匹配就很简单了，我们只需要当需要对非终结符进行推导时，看看此时匹配的串的非终结符，然后选择相应的规则进行推导即可。

而重点是，我们如果根据文法，构建出预测分析表。

我们定义: $\text{FIRST}(\alpha)$是可以从 $\alpha$推导得到的句型的首个终结符的集合。形式化定义为:

$$ \text{FIRST}(\alpha)=\{ t\in T \cup \{\epsilon\} | \alpha \xRightarrow * t\beta \lor \alpha \xRightarrow * \epsilon \} $$

我们定义: $\text{FOLLOW}(A)$是可能在某些句型中紧跟在A右侧的终结符集合。形式化定义为:

$$ \text{FOLLOW}(A) = \{t\in T \cup \{\$\} | \exist s. S\xRightarrow* s \triangleq \beta At\gamma \} $$

我们接下来看如何计算这两个集合。

我们这样计算每个符号X的 $\text{FIRST}$集合:

若X是终结符，那么 $\text{FIRST}(X) = X$;

若X是非终结符，那么 $\text{FIRST}(X) \leftarrow \text{FIRST}(X)\cup \{\text{FIRST}(Y_1) / \epsilon\}$,其中 $X \to Y_1Y_2\dots Y_k$

另外，对于 $Y_2-Y_k$,若 $\epsilon \in L(Y_1,\dots,Y_i)$,那么有: $\text{FIRST}(X) \leftarrow \text{FIRST}(X) \cup \{\text{FIRST}(Y_i) / \epsilon\}$

特别地，若: $\epsilon \in L(Y_1,\dots,Y_k)$,那么 $\text{FIRST}(X) \leftarrow \text{FIRST}(X) \cup \{ \epsilon\}$

我们以下面的文法进行一次 $\text{FIRST}(X)$计算的演示(为了方便，我们用F(x)代替):

$$ (1) X \to Y, (2)X \to a,(3)Y\to \epsilon,(4)Y\to c,(5)Z\to d,(6)Z\to XYZ $$

我们先算 F(x),我们根据文法 $X\to Y$知道，我们需要先算F(Y).

由文法 $Y\to \epsilon,Y\to c$我们可以知道: $F(Y) = \{c,\epsilon\}$

计算完F(Y)后，我们可以进一步得到 $F(X) = \{a,c,\epsilon\}$，其中F(X)包含 $\epsilon$则是根据最后一条规则

最后我们计算F(Z),首先由 $Z\to d$，我们有 $\{d\}\subseteq F(Z)$,再看: $Z\to XYZ$

根据计算方法，我们先看X，此时我们有:
$F(Z)\leftarrow F(Z) \cup \{F(X) / \epsilon\} = \{a,c,d\}$

而由于 $X \Rightarrow \epsilon$,因此我们还需要看Y，此时我们得到 $\{a,c,d\}\subseteq F(Z)$,而由于Y也可以推出 $\epsilon$

因此我们最后还需要看Z，就得到了: $F(Z) = \{a,c,d\}$

需要注意的是，Z推不出 $\epsilon$，因此F(Z)中并不会含有它。

我们这样计算每个符号X的 $\text{FOLLOW}$(X)

若X是开始符号，那么有: $\text{FOLLOW}(X)\leftarrow \text{FOLLOW}(X)\cup \{\$\}$

若X是某个产生式右部的最后一个符号： $A \to \alpha X$： $\text{FOLLOW}(X)\leftarrow \text{FOLLOW}(X) \cup \text{FOLLOW}(A)$

若X是某个产生式的右部的中间的一个符号: $A\to \alpha X \beta$: $\text{FOLLOW}(X)\leftarrow \text{FOLLOW}(X)\cup (\text{FIRST}(\beta) / \{\epsilon\})$

特别地，若此时 $\epsilon \in \text{FIRST}(\beta)$,那么 $\text{FOLLOW}(X)\leftarrow \text{FOLLOW}(X)\cup \text{FOLLOW}(A)$

还是一样，我们以上面的文法为例，展示 $\text{FOLLOW}$集合的计算方法。

我们首先处理开始符号X，因为X是开始符号，我们显然有：

$\$ \in \text{FOLLOW}(X)$,所以目前： $\text{FOLLOW}(X) = \{\$\}$

接下来根据生产式: $X\to Y$，以及规则2，我们有: $\text{FOLLOW}(X)\subseteq \text{FOLLOW}(Y)$

因此目前: $\text{FOLLOW}(Y) = \{\$\}$

接下来看 $Z\to XYZ$.

对X而言，后面跟着YZ，所以要把 $\text{FIRST}(YZ) -\{\epsilon\}$,加入到 $\text{FOLLOW}(X)$中，因此我们需要计算:

$$ \text{FIRST}(YZ) = \{ a,c,d\} $$

因此: $\text{FOLLOW}(X) = \{\$,a,c,d\}$

对于Y而言，后面跟着Z，所以有:

$\text{FIRST}(Z) -\epsilon \subseteq \text{FOLLOW}(Y)$,而: $\text{FIRST}(Z) = \{a,c,d\}$

因此: $\text{FOLLOW}(Y) = \{\$,a,c,d\}$

而Z根据规则1知道是一个空集。

得到了FIRST集和FOLLOW集后，我们可以开始构造给定文法的预测分析表了。

通常而言，预测分析表写作: M[A,a],其中行为非终结符A，列尾终结符a，单元格内则填写应该选择哪个产生式。

填表规则只有两条。

对于产生式 $A\to \beta$

若 $a\in \text{FIRST}(\beta)-\{\epsilon\}$那么有: $M[A,a] = A\to \beta$;

若 $\epsilon \in \text{FIRST}(\beta)$,那么对于所有的: $b\in \text{FOLLOW}(A)$都有: $M[A,b] = A\to \beta$.

在实际计算中我们通常会先计算产生式 $A\to\beta$的 $\text{SELECT}$集合

我们定义: $\text{SELECT}(A\to\beta)$

若： $\epsilon \notin \text{FIRST}(\beta)$,那么 $\text{SELECT}(A\to\beta) = \text{FIRST}(\beta)$

若： $\epsilon \in \text{FIRST}(\beta)$,那么 $\text{SELECT}(A\to \beta) = (\text{FIRST}(\beta)-\{\epsilon\})\cup \text{FOLLOW}(A)$

得到 $A\to\beta$的SELECT集后，对该集合中每一个终结符a，把生产式 $A\to\beta$填入M[A,a]

还是上面那个例子，我们来手算一遍:

我们知道:

$$ \begin{aligned} (1)\quad &X\to Y\\ (2)\quad &X\to a\\ (3)\quad &Y\to \epsilon\\ (4)\quad &Y\to c\\ (5)\quad &Z\to d\\ (6)\quad &Z\to XYZ \end{aligned} $$$$ \begin{aligned} FIRST(X)&=\{a,c,\epsilon\}\\ FIRST(Y)&=\{c,\epsilon\}\\ FIRST(Z)&=\{a,c,d\} \end{aligned} $$$$ \begin{aligned} FOLLOW(X)&=\{a,c,d,\$\}\\ FOLLOW(Y)&=\{a,c,d,\$\}\\ FOLLOW(Z)&=\varnothing \end{aligned} $$

那么对于产生式（1） $X\to Y$,因为: $\text{FIRST}(Y) = \{c,\epsilon\}$,其中包含了 $\epsilon$，因此:

$$ \text{SELECT}(X\to Y) = \{a,c,d,\$\} $$

对于产生式（2） $X \to a$,这里 $\text{FIRST}(a) = \{a\}$，不包含 $\epsilon$,因此: $\text{SELECT}(X\to a) =\{a\} $

按照同样的思路，我们可以求出每个产生式的 $\text{SELECT}$集，从而得到最终的预测分析表:

非终结符	a	c	d	$
X	(1), (2)	(1)	(1)	(1)
Y	(3)	(3), (4)	(3)	(3)
Z	(6)	(6)	(5), (6)	error

[从零开始的编译原理][实验] Exp Chapter 1：从 SysY 到 Koopa IR

Tue, 29 Sep 2026 20:23:03 +0800

接下来，我们正式开始实现自己的编译器。

本阶段的目标是先完成一个最小可用版本：让编译器能够处理一个简单的 main 函数，并生成对应的 Koopa IR。

一个完整编译器的典型流程大致如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


词法分析
-> 语法分析
-> AST
-> 语义分析
-> 带有类型 / 符号信息的 AST
-> IR 生成
-> 优化
-> 指令选择
-> 寄存器分配
-> 指令调度

其中，词法分析、语法分析和 AST 构建属于编译器前端；

IR 生成通常位于前端与中端的衔接位置，而后续的优化、指令选择、寄存器分配和指令调度则会逐步进入编译器中端和后端。

根据指导书的安排，本次实验暂时不会实现完整的编译器流程，而是先完成从源代码到 Koopa IR 的这一部分，也就是：

1
2
3
4
5


源代码
-> 词法分析
-> 语法分析
-> AST
-> IR 生成

通过这一阶段，我们可以先打通一个最基本的编译流程，再在后续实验中逐步加入更完整的语义分析、优化以及代码生成等功能。

词法分析 & 语法分析

首先来看编译器前端最基础的两个阶段：词法分析（Lexical Analysis）和语法分析（Syntax Analysis）。

在本实验中，我们并不打算从零手写词法分析器和语法分析器，而是分别使用 Flex 和 Bison 来自动生成它们。我们只需要根据 SysY 的词法和语法规范，描述各种 Token 的形式以及程序的文法结构，之后 Flex 和 Bison 就可以据此生成对应的分析器。

SysY 的部分词法 / 语法规范

在本章中，我们暂时只实现一个能够处理 main 函数和 return 语句的简单编译器。也就是说，目前我们的编译器只需要能够处理类似下面这样的 SysY 程序：

1
2
3
4


int main() {
 // 注释也应该被忽略
 return 0;
}

最终，我们希望将它编译成对应的 Koopa IR：

1
2
3
4


fun @main(): i32 {
%entry:
 ret 0
}

在真正开始编写 Flex 和 Bison 代码之前，我们先来看一下这一阶段需要处理的 SysY 词法和语法规范。

词法规范主要描述源代码中存在哪些 Token，以及每种 Token 长什么样。例如：

1
2
3


int main() {
 return 0;
}

经过词法分析以后，大致可以得到：

1

INT IDENT '(' ')' '{' RETURN INT_CONST ';' '}'

其中，main 和 0 除了 Token 类型之外，还需要携带它们自身的值。

SysY 中标识符 IDENT 的规范如下：

1
2
3


identifier ::= identifier-nondigit
 | identifier identifier-nondigit
 | identifier digit;

其中 identifier-nondigit 为下划线 _、小写英文字母或大写英文字母，digit 为数字 0 到 9。

换句话说，标识符的首字符只能是字母或下划线，后续字符还可以包含数字，因此可以使用正则表达式：

1

[a-zA-Z_][a-zA-Z0-9_]*

来描述。

SysY 中的整型常量记作 INT_CONST，可以采用十进制、八进制或十六进制表示：

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


integer-const ::= decimal-const
 | octal-const
 | hexadecimal-const;

decimal-const ::= nonzero-digit
 | decimal-const digit;

octal-const ::= "0"
 | octal-const octal-digit;

hexadecimal-const ::= hexadecimal-prefix hexadecimal-digit
 | hexadecimal-const hexadecimal-digit;

hexadecimal-prefix ::= "0x" | "0X";

其中，nonzero-digit 为数字 1 到 9，octal-digit 为数字 0 到 7，hexadecimal-digit 为数字 0 到 9 或大小写字母 a-f。

例如：

1
2
3


123 十进制
077 八进制
0xFF 十六进制

SysY 中的注释规则与 C 语言基本一致：

单行注释以 // 开始，一直到换行符结束；
多行注释以 /* 开始，直到第一次出现 */ 时结束。

注释在后续编译阶段中没有意义，因此词法分析器识别出注释之后，可以直接将它们丢弃。

目前我们只需要支持如下语法：

1
2
3
4
5
6
7
8


CompUnit ::= FuncDef;

FuncDef ::= FuncType IDENT "(" ")" Block;
FuncType ::= "int";

Block ::= "{" Stmt "}";
Stmt ::= "return" Number ";";
Number ::= INT_CONST;

其中，开始符号为 CompUnit。

EBNF 介绍

EBNF，即 Extended Backus–Naur Form（扩展巴科斯范式），是一种用于描述编程语言语法的形式化方法。

基于 SysY 的 EBNF，我们可以从指定的开始符号出发，通过不断应用产生式规则，推导出所有符合该语法的程序。

EBNF 由若干条类似下面的规则组成：

1

A ::= B;

它表示，当我们遇到 A 时，可以按照这条规则将 A 替换成 B，这个过程称为一次推导（Derivation）。

例如：

1

CompUnit ::= FuncDef;

表示：

1

CompUnit

可以被替换成：

1

FuncDef

像 CompUnit、FuncDef、FuncType、Block、Stmt、Number 这样还能继续按照某条产生式展开的符号，被称为非终结符（Non-terminal）。

我们从开始符号 CompUnit 出发：

1

CompUnit

根据：

1

CompUnit ::= FuncDef;

得到：

1

FuncDef

继续根据：

1

FuncDef ::= FuncType IDENT "(" ")" Block;

得到：

1

FuncType IDENT "(" ")" Block

再利用：

1

FuncType ::= "int";

可以得到：

1

"int" IDENT "(" ")" Block

继续展开 Block、Stmt 和 Number，最终得到：

1

"int" IDENT "(" ")" "{" "return" INT_CONST ";" "}"

此时已经没有符号可以继续按照产生式展开。像 "int"、IDENT、"("、INT_CONST 这样的符号，就称为终结符（Terminal）。

这些终结符基本上就对应词法分析器产生的 Token。

因此可以粗略地理解为：

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


源代码
 │
 ▼
词法分析
 │
 ▼
Token / 终结符
 │
 ▼
语法分析
 │
 ▼
按照文法规则组合成程序结构

Flex 教程

在 C++ 中，我们使用 Flex 生成词法分析器。

Flex 主要负责描述 EBNF 中的终结符，也就是描述每种 Token 的形式，并在识别出 Token 后返回对应的类型和语义值。Token 的形式通常可以使用正则表达式来描述。

首先在 src 目录下创建：

1

sysy.l

Flex 文件大致被两个 %% 分成三个区域：

1
2
3
4
5


定义区
%%
规则区
%%
用户代码区

定义区位于第一个 %% 之前，主要用于设置 Flex 选项、插入 C++ 代码，以及定义可以复用的正则表达式：

 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


%option noyywrap
%option nounput
%option noinput

%{

#include <cstdlib>
#include <string>

// Flex 需要使用 Bison 中定义的 Token 和 yylval
#include "sysy.tab.hpp"

using namespace std;

%}

/* 空白符和注释 */
WhiteSpace [ \t\n\r]*
LineComment "//".*
MultiLineComment "/*"([^*]|\*+[^*/])*\*+"/"

/* 标识符 */
Identifier [a-zA-Z_][a-zA-Z0-9_]*

/* 整数字面量 */
Decimal [1-9][0-9]*
Octal 0[0-7]*
Hexadecimal 0[xX][0-9a-fA-F]+

最前面的：

1
2
3


%option noyywrap
%option nounput
%option noinput

是 Flex 的一些配置选项。默认情况下，Flex 会生成一些我们当前用不到的接口，因此这里直接将它们关闭。

被：

1
2
3


%{
...
%}

包围的代码会被原样插入 Flex 生成的源文件。

这里需要包含：

1

#include "sysy.tab.hpp"

因为稍后的规则中需要使用：

1
2
3
4


INT
RETURN
IDENT
INT_CONST

等 Token，而这些 Token 的编号，以及 yylval 的类型，都是由 Bison 生成的头文件定义的。

后面的：

1

Identifier [a-zA-Z_][a-zA-Z0-9_]*

相当于给一段正则表达式起一个名字。之后在规则区中就可以直接写：

1

{Identifier}

而不需要反复写完整的正则表达式。

多行注释稍微特殊一些。在某些支持非贪婪匹配的正则表达式引擎中，可以写成类似：

1

/\*.*?\*/

但 Flex 并不直接支持这种写法，而且 . 默认不能匹配换行符，因此可以写成：

1

"/*"([^*]|\*+[^*/])*\*+"/"

其中：

1

[^*]

表示任意不是 * 的字符，因此也可以匹配换行符；而：

1

\*+[^*/]

则用于处理注释内部出现的一个或多个 *，同时避免过早将它们识别为结束标记 */。

规则区位于两个 %% 之间。每一条规则都由“模式 + 动作”组成，当 Flex 匹配到某个模式时，就执行后面的 C++ 动作：

 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


{WhiteSpace} { /* 忽略 */ }
{LineComment} { /* 忽略 */ }
{MultiLineComment} { /* 忽略 */ }

"int" { return INT; }
"return" { return RETURN; }

{Identifier} {
 yylval.str_val = new string(yytext);
 return IDENT;
}

{Decimal} {
 yylval.int_val = strtol(yytext, nullptr, 0);
 return INT_CONST;
}

{Octal} {
 yylval.int_val = strtol(yytext, nullptr, 0);
 return INT_CONST;
}

{Hexadecimal} {
 yylval.int_val = strtol(yytext, nullptr, 0);
 return INT_CONST;
}

. {
 return yytext[0];
}

Flex 在存在多条可以匹配的规则时，会遵循两个原则：

优先选择能够匹配最长字符串的规则；
如果匹配长度相同，则选择写在前面的规则。

例如 int 既可以匹配：

1

"int"

也可以匹配：

1

{Identifier}

两者匹配长度都是 3，因此 Flex 会选择写在前面的 "int" 规则。这也是为什么关键字通常要写在标识符规则之前。

在动作中，我们还会经常使用几个 Flex 提供的变量：

yytext：当前刚刚匹配到的源代码文本，是一个 C 风格字符串；
yyleng：当前 yytext 的长度；
yylval：传递给语法分析器的语义值。

例如输入：

1

main

被 {Identifier} 匹配后：

1

yytext

的内容就是 "main"。

而：

1

yylval.str_val = new string(yytext);

则将真正的标识符名称保存下来并传递给 Bison。

对于整数：

1

yylval.int_val = strtol(yytext, nullptr, 0);

则把文本形式的整数转换为真正的数值。

这里的：

1
2


str_val
int_val

会在稍后的 Bison %union 中定义。

最后：

1

. { return yytext[0]; }

表示，如果前面的所有规则都没有匹配，就直接把当前单个字符作为 Token 返回。

因此：

1
2
3
4
5


(
)
{
}
;

这些字符不需要分别定义 LPAREN、RPAREN 等 Token，而是可以直接返回字符本身。

最后一个 %% 之后是用户代码区。虽然目前我们暂时没有使用这一部分，但以后可以在这里定义一些普通的 C++ 辅助函数。

Bison 教程

完成词法分析器之后，接下来使用 Bison 生成语法分析器。

首先在 src 目录中创建：

1

sysy.y

与 Flex 类似，Bison 文件同样使用两个 %% 分成三个部分：

1
2
3
4
5


声明区
%%
文法区
%%
用户代码区

声明区主要负责插入 C++ 代码、声明 Token、定义语义值类型、定义非终结符的语义类型，以及给 parser 增加额外参数：

 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


%code requires {
 #include <memory>
 #include <string>
}

%{

#include <iostream>
#include <memory>
#include <string>

int yylex();
void yyerror(std::unique_ptr<std::string> &ast, const char *s);

using namespace std;

%}

%parse-param { std::unique_ptr<std::string> &ast }

%union {
 std::string *str_val;
 int int_val;
}

%token INT RETURN
%token <str_val> IDENT
%token <int_val> INT_CONST

%type <str_val> FuncDef FuncType Block Stmt Number

其中：

1
2
3


%code requires {
 ...
}

中的代码会被放入 Bison 生成的头文件 sysy.tab.hpp 中。

而：

1
2
3


%{
...
%}

中的内容主要会进入 Bison 生成的解析器源文件。

这里：

1

int yylex();

用于声明 Flex 生成的词法分析器函数。Bison 在需要读取下一个 Token 时，就会调用 yylex()。

而：

1

void yyerror(...);

用于声明语法分析发生错误时调用的错误处理函数。

1

%parse-param { std::unique_ptr<std::string> &ast }

表示给生成的 yyparse() 增加一个额外参数。于是我们就可以：

1
2


unique_ptr<string> ast;
yyparse(ast);

解析完成之后，再把最终结果通过 ast 交回调用方。

这里暂时使用一个字符串保存语法分析的结果，在下一节中再把它真正替换成 AST。

在语法分析过程中，不同 Token 和非终结符可能携带不同类型的值。例如：

1

IDENT

需要携带一个：

1

std::string *

而：

1

INT_CONST

需要携带一个：

1

int

因此 Bison 使用：

1
2
3
4


%union {
 std::string *str_val;
 int int_val;
}

定义语义值可能采用的类型。

于是 Flex 中的：

1

yylval.str_val

和：

1

yylval.int_val

就是这里定义的两个字段。

这里暂时使用 std::string *，而不是直接使用 std::string，是因为传统 C 风格 union 对带有非平凡构造函数和析构函数的 C++ 类型处理起来比较麻烦。

接下来：

1

%token INT RETURN

声明 INT 和 RETURN 两种 Token。

而：

1
2


%token <str_val> IDENT
%token <int_val> INT_CONST

则表示 IDENT 携带 str_val 类型的语义值，INT_CONST 携带 int_val 类型的语义值。

这些 Token 名字必须与 Flex 中：

1
2
3
4


return INT;
return RETURN;
return IDENT;
return INT_CONST;

使用的名字一致。

最后：

1

%type <str_val> FuncDef FuncType Block Stmt Number

表示这些非终结符的语义值类型都是：

1

std::string *

也就是说，当这些非终结符完成规约后，会产生一个字符串结果并继续向上传递。

Bison 的文法规则基本形式如下：

1
2
3
4
5


非终结符
 : 产生式 {
 动作
 }
 ;

如果一个非终结符存在多个产生式，则可以使用 | 分隔。

目前我们需要解析的程序只有：

1
2
3


int main() {
 return 0;
}

整个程序只包含一个函数定义，因此：

1

CompUnit ::= FuncDef;

在 Bison 中可以写成：

1
2
3
4
5


CompUnit
 : FuncDef {
 ast = unique_ptr<string>($1);
 }
 ;

这里：

1

$1

表示产生式右侧第一个符号的语义值。

因为右侧只有一个 FuncDef，所以 $1 就是 FuncDef 的返回值。

当整个输入最终被成功规约为开始符号 CompUnit 时，语法分析就完成了，因此我们把 $1 保存到 ast 中作为最终结果。

接下来是：

1

FuncDef ::= FuncType IDENT "(" ")" Block;

对应：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


FuncDef
 : FuncType IDENT '(' ')' Block {
 auto type = unique_ptr<string>($1);
 auto ident = unique_ptr<string>($2);
 auto block = unique_ptr<string>($5);

 $$ = new string(
 *type + " " + *ident + "()" + *block
 );
 }
 ;

这里需要认识 Bison 中两个非常重要的记号。

$1、$2 等表示产生式右侧各个符号的语义值，例如：

1
2


FuncType IDENT '(' ')' Block
 $1 $2 $3 $4 $5

而：

1

$$

则表示当前产生式左侧非终结符的语义值。

因此：

1

$$ = new string(...);

表示这几个符号被规约成 FuncDef 后，FuncDef 自身的语义值就是这个新生成的字符串。

这里：

1
2
3


auto type = unique_ptr<string>($1);
auto ident = unique_ptr<string>($2);
auto block = unique_ptr<string>($5);

则是使用 unique_ptr 接管之前通过 new string(...) 创建的字符串。

例如 $1 原本只是一个裸指针：

1

string *

将它交给：

1

unique_ptr<string>

后，这块内存就由智能指针负责管理。当当前语义动作执行结束、局部变量离开作用域时，对应的字符串也会被自动释放，因此我们不需要手动调用 delete。

剩下的语法可以按照同样的方法实现：

1
2
3
4


FuncType ::= "int";
Block ::= "{" Stmt "}";
Stmt ::= "return" Number ";";
Number ::= INT_CONST;

对应：

 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


FuncType
 : INT {
 $$ = new string("int");
 }
 ;

Block
 : '{' Stmt '}' {
 auto stmt = unique_ptr<string>($2);
 $$ = new string("{" + *stmt + "}");
 }
 ;

Stmt
 : RETURN Number ';' {
 auto number = unique_ptr<string>($2);
 $$ = new string("return " + *number + ";");
 }
 ;

Number
 : INT_CONST {
 $$ = new string(to_string($1));
 }
 ;

随着 Bison 不断进行规约，这些字符串会逐层向上传递：

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


INT_CONST
 │
 ▼
 Number
 │
 ▼
 Stmt
 │
 ▼
 Block
 │
 ▼
 FuncDef
 │
 ▼
CompUnit

最终可以得到类似：

1

int main(){return 0;}

这样的结果。

需要注意的是，这里还没有真正构造 AST。我们只是暂时使用字符串来模拟语法分析结果，从而熟悉 Bison 的语义动作以及 $1、$$ 等机制。真正的 AST 会在下一节实现。

第二个 %% 之后是用户代码区。目前我们只需要定义错误处理函数：

1
2
3


void yyerror(unique_ptr<string> &ast, const char *s) {
 cerr << "error: " << s << endl;
}

当 Bison 在语法分析过程中发现输入不符合文法时，就会调用 yyerror() 输出错误信息。

生成词法 / 语法分析器

完成 sysy.l 和 sysy.y 后，Flex 和 Bison 就可以根据我们的规则生成真正的词法分析器和语法分析器。

接下来修改：

1

src/main.cpp

内容如下：

 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


#include <cassert>
#include <cstdio>
#include <iostream>
#include <memory>
#include <string>

using namespace std;

extern FILE *yyin;
extern int yyparse(unique_ptr<string> &ast);

int main(int argc, const char *argv[]) {
 assert(argc == 5);

 auto mode = argv[1];
 auto input = argv[2];
 auto output = argv[4];

 yyin = fopen(input, "r");
 assert(yyin);

 unique_ptr<string> ast;

 auto ret = yyparse(ast);
 assert(!ret);

 cout << *ast << endl;

 return 0;
}

其中：

1

extern FILE *yyin;

表示 Flex 应该从哪个文件读取输入。

因此：

1

yyin = fopen(input, "r");

就是将需要编译的 SysY 源文件交给词法分析器。

而：

1

yyparse(ast);

会启动 Bison 生成的语法分析器。当语法分析器需要新的 Token 时，就会调用 Flex 生成的：

1

yylex()

因此整个过程可以表示为：

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


SysY 源文件
 │
 ▼
 Flex
 │
 │ yylex()
 ▼
 Token
 │
 ▼
 Bison
 │
 │ 规约
 ▼
解析结果

接下来进入开发环境：

1
2
3


docker run -it --rm \
 -v /path/to/compiler:/root/compiler \
 maxxing/compiler-dev bash

编译项目：

1

cmake --build build

然后运行：

1

./build/compiler --koopa debug/lab1.sy -o debug/lab1.koopa

此时可以看到输出：

1

int main(){return 0;}

这说明我们的词法分析器和语法分析器已经能够正确识别：

1
2
3


int main() {
 return 0;
}

到这里，我们已经完成了最基本的：

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


源代码
 │
 ▼
词法分析
 │
 ▼
Token Stream
 │
 ▼
语法分析
 │
 ▼
程序结构

不过，目前所谓的“程序结构”仍然只是通过字符串拼接得到的临时结果。

在下一节中，我们会正式定义 AST（Abstract Syntax Tree），并让 Bison 在语法分析过程中直接构造真正的 AST。

解析 main 函数：构建 AST

在第一节中，我们借助 Flex 和 Bison 实现了一个能够解析简单 main 函数的编译器前端。

不过，此时语法分析器只是将解析结果重新拼接成一个字符串。虽然这足以帮助我们理解 Bison 的基本工作方式，但字符串并不适合作为后续编译阶段处理程序的数据结构。

因此，在本节中，我们将正式设计 AST，并让 Bison 在语法分析过程中直接构造 AST。

设计 AST

AST，即 Abstract Syntax Tree（抽象语法树），用于保存源程序中对后续编译阶段真正有意义的结构。

后续的语义分析、中间代码生成以及优化，都需要在 AST 上进行，因此 AST 的设计主要需要考虑两点：

能够表达程序中必要的语法和语义结构；
尽可能方便后续编译阶段处理。

在当前实验中，我们处理的语法如下：

1
2
3
4
5
6
7
8


CompUnit ::= FuncDef;

FuncDef ::= FuncType IDENT "(" ")" Block;
FuncType ::= "int";

Block ::= "{" Stmt "}";
Stmt ::= "return" Number ";";
Number ::= INT_CONST;

从文法结构来看，可以得到大致如下的层次关系：

1
2
3
4
5
6
7


CompUnit
└── FuncDef
 ├── FuncType
 ├── IDENT
 └── Block
 └── Stmt
 └── Number

因此，可以先按照这个结构设计 AST。

首先创建：

1

src/ast.hpp

并定义所有 AST 节点的公共基类：

1
2
3
4
5
6
7
8
9


#pragma once

#include <memory>
#include <string>

class BaseAST {
public:
 virtual ~BaseAST() = default;
};

这里：

1

#pragma once

用于避免同一个头文件在一个编译单元中被重复包含，从而防止类或函数被重复定义。

而：

1

virtual ~BaseAST() = default;

则定义了一个虚析构函数。

之后，我们会经常使用：

1

std::unique_ptr<BaseAST>

来保存不同类型的 AST 节点。例如：

1
2


std::unique_ptr<BaseAST> ast =
 std::make_unique<FuncDefAST>();

虽然指针类型是 BaseAST，但实际对象可能是 FuncDefAST、BlockAST 等派生类。因此，基类的析构函数必须是虚函数，这样通过 BaseAST 指针销毁对象时，才能正确调用实际派生类的析构函数。

接下来可以根据文法定义具体的 AST 节点。

例如：

1

CompUnit ::= FuncDef;

说明一个 CompUnit 中包含一个 FuncDef，因此可以写成：

1
2
3
4


class CompUnitAST : public BaseAST {
public:
 std::unique_ptr<BaseAST> func_def;
};

类似地：

1

FuncDef ::= FuncType IDENT "(" ")" Block;

可以对应：

1
2
3
4
5
6


class FuncDefAST : public BaseAST {
public:
 std::unique_ptr<BaseAST> func_type;
 std::string ident;
 std::unique_ptr<BaseAST> block;
};

这里需要注意，文法中的每个符号并不一定都要在 AST 中单独建立节点。

例如：

1
2
3
4
5


(
)
{
}
;

这些符号只是用来约束源代码的语法形式，一旦语法分析完成，它们就没有必要继续保留。

类似地：

1
2


int
return

通常也不需要作为独立节点保存。因为一旦我们已经构造了 FuncTypeAST 或 ReturnStmtAST，节点类型本身就已经表达了对应的语义。

而：

1
2


IDENT
INT_CONST

则不同。

例如：

1
2
3


int main() {
 return 123;
}

其中：

1
2


main
123

都是程序本身携带的数据，因此必须在 AST 中保存，例如：

1
2


std::string ident;
int value;

所以

是否需要建立节点或者保存字段，取决于这部分信息对后续编译阶段是否仍然有意义。

这也是 AST 与普通语法树的一个重要区别。

如果严格按照 EBNF 的每一层都建立节点，那么得到的结构更接近 CST（Concrete Syntax Tree，具体语法树）。

而 AST 会主动省略很多只服务于语法分析的细节。

例如：

1

Stmt ::= "return" Number ";";

理论上完全可以直接设计成：

1
2
3
4


class ReturnStmtAST : public BaseAST {
public:
 int value;
};

而不一定非要保留：

1
2


StmtAST
└── NumberAST

不过在当前实验中，为了更直观地理解 Bison 如何逐层构造 AST，我们暂时让 AST 的结构和 EBNF 保持得比较接近。

生成 AST

设计好 AST 之后，就可以在 Bison 的语法动作中真正构造这些节点了。

在上一节中，我们对每个语法单元的处理方式都是：

1

std::string *

每完成一次规约，就把对应的内容重新拼成字符串。

现在我们要把它替换成：

1

BaseAST *

也就是说：

Bison 每完成一次规约，就构造一个对应的 AST 节点，并将这个节点继续向上传递。

首先修改 src/sysy.y。

我们需要让 Bison 生成的头文件能够看到 BaseAST 的定义，因此：

1
2
3
4
5


%code requires {
 #include <memory>
 #include <string>
 #include "ast.hpp"
}

这里使用的是：

1

%code requires

而不是单纯的：

1
2
3


%{
...
%}

原因是 Bison 不仅会生成：

1

sysy.tab.cpp

还会生成：

1

sysy.tab.hpp

而 %union 等声明可能需要出现在生成的头文件中。

因为这些声明中会使用：

1

BaseAST *

所以 ast.hpp 也必须在生成的头文件中可见。

接下来，将 parser 的额外参数从：

1

%parse-param { std::unique_ptr<std::string> &ast }

修改为：

1

%parse-param { std::unique_ptr<BaseAST> &ast }

这样，语法分析完成以后，得到的根节点就可以通过：

1

std::unique_ptr<BaseAST> ast;

传回 main 函数。

接下来修改 %union：

1
2
3
4
5


%union {
 std::string *str_val;
 int int_val;
 BaseAST *ast_val;
}

在上一节中，非终结符的语义值都是：

1

std::string *

现在则改成 AST 节点指针：

1

%type <ast_val> FuncDef FuncType Block Stmt Number

这样，像 FuncDef、Block 这样的非终结符在完成规约以后，返回的就是一个 AST 节点。

首先来看：

1

CompUnit ::= FuncDef;

对应的 Bison 代码可以写成：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


CompUnit
 : FuncDef {
 auto comp_unit = std::make_unique<CompUnitAST>();

 comp_unit->func_def =
 std::unique_ptr<BaseAST>($1);

 ast = std::move(comp_unit);
 }
 ;

这里：

1

$1

表示产生式右侧第一个符号，也就是 FuncDef 的语义值。

因为：

1

%type <ast_val> FuncDef

所以 $1 的实际类型就是：

1

BaseAST *

接下来：

1

std::unique_ptr<BaseAST>($1)

将这个裸指针交给 unique_ptr 管理。

最后：

1

ast = std::move(comp_unit);

把当前构造出的 CompUnitAST 交给 parser 的输出参数 ast。

这里涉及 unique_ptr 的所有权转移。

例如：

1

auto comp_unit = std::make_unique<CompUnitAST>();

此时 AST 对象由：

1

comp_unit

独占管理。

而 unique_ptr 不能被复制，因此不能直接：

1

ast = comp_unit;

而需要：

1

ast = std::move(comp_unit);

表示：

1
2
3
4
5


comp_unit
 │
 │ 所有权转移
 ▼
 ast

执行之后，AST 对象由 ast 管理，而原来的 comp_unit 会变成空指针。

因此，这里的 std::move 并不是把 AST 对象本身“搬走”，而是把：

对这个 AST 对象的所有权从一个 unique_ptr 转移给另一个 unique_ptr。

接下来处理：

1

FuncDef ::= FuncType IDENT "(" ")" Block;

可以写成：

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


FuncDef
 : FuncType IDENT '(' ')' Block {
 auto ast = new FuncDefAST();

 ast->func_type =
 std::unique_ptr<BaseAST>($1);

 ast->ident =
 *std::unique_ptr<std::string>($2);

 ast->block =
 std::unique_ptr<BaseAST>($5);

 $$ = ast;
 }
 ;

这里：

分别表示：

1
2


FuncType IDENT '(' ')' Block
 $1 $2 $3 $4 $5

而：

1

$$

表示产生式左侧 FuncDef 的语义值。

因此：

1

$$ = ast;

表示：

当前这些符号规约成 FuncDef 后，FuncDef 对应的 AST 节点就是刚刚构造出的 FuncDefAST。

整个过程可以理解成：

1
2
3
4
5


FuncType IDENT '(' ')' Block
 │
 │ 规约
 ▼
 FuncDefAST

然后这个节点会继续作为 $1、$2 等语义值被更上层的规则使用。

对于：

1

FuncType ::= "int";

可以写成：

1
2
3
4
5
6


FuncType
 : INT {
 auto ast = new FuncTypeAST();
 $$ = ast;
 }
 ;

因为当前只支持 int 一种函数返回类型，所以 FuncTypeAST 暂时不需要保存额外的数据。

以后如果支持：

1
2
3


int
void
float

等不同类型，再在 FuncTypeAST 中添加对应字段即可。

其他规则也可以按照完全相同的方法构造 AST。

整个过程本质上就是：

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


Token
 │
 ▼
Bison 规约
 │
 ├── NumberAST
 │
 ▼
StmtAST
 │
 ▼
BlockAST
 │
 ▼
FuncDefAST
 │
 ▼
CompUnitAST

随着语法分析不断进行，较小的 AST 节点被逐步组合成较大的 AST 节点，最终形成整棵树。

parser 参数改变以后，yyerror 也需要同步修改：

1
2
3
4
5


void yyerror(
 std::unique_ptr<BaseAST> &ast,
 const char *s) {
 ...
}

同时，main.cpp 中原来的：

1

extern int yyparse(unique_ptr<string> &ast);

也需要改成：

1

extern int yyparse(unique_ptr<BaseAST> &ast);

这样，语法分析器最终返回的就不再是一个字符串，而是一棵真正的 AST。

检查生成结果

现在 AST 已经可以正确构造出来了，但它暂时只保存在内存中。

为了检查生成结果是否符合我们的预期，我们希望能够把整棵 AST 打印到终端。

可以利用 C++ 的虚函数机制，在 BaseAST 中增加一个统一的：

1

Dump()

接口：

1
2
3
4
5
6


class BaseAST {
public:
 virtual ~BaseAST() = default;

 virtual void Dump() const = 0;
};

这里：

1

= 0

表示 Dump() 是一个纯虚函数。

因此，BaseAST 本身不能直接实例化，而所有具体的 AST 节点都需要实现自己的 Dump()。

例如：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


class CompUnitAST : public BaseAST {
public:
 std::unique_ptr<BaseAST> func_def;

 void Dump() const override {
 std::cout << "CompUnitAST { ";
 func_def->Dump();
 std::cout << " }";
 }
};

类似地：

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


class FuncDefAST : public BaseAST {
public:
 std::unique_ptr<BaseAST> func_type;
 std::string ident;
 std::unique_ptr<BaseAST> block;

 void Dump() const override {
 std::cout << "FuncDefAST { ";

 func_type->Dump();

 std::cout << ", "
 << ident
 << ", ";

 block->Dump();

 std::cout << " }";
 }
};

虽然：

1
2
3


func_def
func_type
block

这些字段的静态类型都是：

1

std::unique_ptr<BaseAST>

但它们实际指向的对象可能分别是：

1
2
3


FuncDefAST
FuncTypeAST
BlockAST

由于 Dump() 是虚函数，所以调用：

1

func_def->Dump();

时，C++ 会根据对象的实际类型自动调用正确的 Dump() 实现。

这就是 C++ 中的运行时多态。

最后，在 main.cpp 中：

1
2
3
4
5
6
7


std::unique_ptr<BaseAST> ast;

auto ret = yyparse(ast);
assert(!ret);

ast->Dump();
std::cout << std::endl;

对于输入：

1
2
3


int main() {
 return 0;
}

最终可以得到类似：

1
2
3
4
5
6
7
8
9


CompUnitAST {
 FuncDefAST {
 FuncTypeAST { int },
 main,
 BlockAST {
 StmtAST { 0 }
 }
 }
}

这样的输出。

整个编译前端到这里已经形成：

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


源代码
 │
 ▼
Flex
 │
 ▼
Token Stream
 │
 ▼
Bison
 │
 ▼
AST

相比上一节单纯把内容重新拼接成字符串，现在的语法分析器已经真正把源程序转换成了一种结构化表示。

换句话说，之前我们主要是在判断：

1

这个程序是否符合语法？

而现在，在确认程序符合语法的同时，我们还得到了：

1

这个程序的结构是什么？

这棵 AST 也将成为后续生成 Koopa IR 的基础。

IR生成

在上一节中，我们已经成功让编译器能够将简单的 SysY 程序解析为 AST。在此基础上，就可以继续进行后续的编译工作，例如生成中间表示（IR）。

在实际的编译器实现中，通常没有必要从头实现完整的后端。借助 LLVM IR 等成熟的编译基础设施，我们可以将源程序转换为统一的中间表示，再把后续的优化、指令选择、寄存器分配以及机器码生成等工作交给现有工具链完成。

因此，本小节将继续完善我们的编译器，实现从 AST 到 IR 的转换。

Koopa IR 基础

在 Koopa IR 中，最大的结构单位是 Program，它代表一个完整的 Koopa IR 程序。

一个 Program 由若干全局值（Value）和函数（Function）组成。其中，每个 Function 又由若干基本块（BasicBlock）构成，而基本块中则包含一系列指令。需要注意的是，在 Koopa IR 中，指令本身也是一种 Value。

因此，一个 Koopa IR 程序的基本结构如下：

1
2
3
4
5
6
7
8


Program
├── Global Values
│
└── Functions
 └── Function
 └── BasicBlocks
 └── BasicBlock
 └── Values / Instructions

接下来，我们重点介绍其中涉及到的几个基本概念。

首先是基本块（Basic Block）。基本块可以理解为一系列连续执行的指令，它具有以下特点：

只有一个入口点：如果其他基本块要将控制流转移到当前基本块，只能跳转到基本块的开头，而不能直接跳到基本块中间的某条指令。
控制流只能在末尾发生转移：基本块中只有最后一条指令可以改变控制流，例如跳转到其他基本块，或者从当前函数中返回（执行 return）。

基本块的存在能够简化编译过程中大量与控制流相关的分析，因此 Koopa IR 要求函数中的指令按照基本块进行组织。

同时，Koopa IR 约定函数的第一个基本块为函数的入口基本块。也就是说，当函数开始执行时，会首先从第一个基本块开始执行。

在本次实验中，我们暂时不需要考虑全局变量。同时，也可以暂时认为 Program 的函数列表中只有一个 Function，而这个 Function 中又只有一个 BasicBlock，也就是函数的入口基本块。

基本块中包含的指令同样属于 Value。Koopa IR 中主要包含以下几类 Value（更详细的介绍可以参考：文档）：

各类常量：整数常量（Integer）、零初始化器（ZeroInit）等。
参数引用：例如函数参数引用（FuncArgRef），用于表示传入函数的参数。
内存分配：全局内存分配（GlobalAlloc，全局变量通过它表示）和局部内存分配（Alloc）。
访存指令：加载（Load）和存储（Store）。
指针运算：GetPtr 和 GetElemPtr。
二元运算：Binary，例如加、减、乘、除、取模以及各种比较运算。
控制流转移：条件分支（Branch）和无条件跳转（Jump）。
函数相关操作：函数调用（Call）和函数返回（Return）。

在本节实验中，我们实际会用到的只有两种：整数常量 Integer 和函数返回 Return。

至此，我们需要生成的 Koopa IR 就已经非常明确了：

生成一个 Koopa IR Program。
Program 中包含一个名为 main 的函数。
函数中包含一个基本块。
基本块中包含一条返回指令。
返回指令的返回值，就是 SysY 中 return 语句后面的整数常量。

如果手动编写对应的 Koopa IR，它大致如下：

1
2
3
4


fun @main(): i32 {
%entry:
 ret 0
}

这里还有几个需要注意的地方：

符号名称的规范。可以看到，这里的函数名是 @main，而不是之前 SysY 程序中的 main。这是因为 Koopa IR 规定，Function、BasicBlock 和具名 Value 的名字必须以 @ 或 % 开头。@ 和 % 在语义上并没有本质区别，但在使用习惯上，我们通常使用 @ 表示源程序中本身就存在的符号，例如函数名和全局变量名；而使用 % 表示编译器在生成 IR 的过程中产生的局部符号或临时值。因此：@main对应的是 SysY 程序中的 main.而：%entry则是我们在生成 IR 时人为创建的基本块名称。
Koopa IR 是一种强类型 IR。也就是说，函数参数、函数返回值以及各种 Value 都具有明确的类型。例如：fun @main(): i32中的 i32 表示 main 函数的返回值类型为 32 位整数，对应 SysY 中的 int。不过，我们并不需要在文本形式的 Koopa IR 中为每一个值都显式写出类型。很多情况下，Koopa IR 的解析器可以根据上下文推导出对应的类型，因此可以省略一部分类型标注。例如：ret 0中并没有显式写出 0 的类型，但根据当前函数的返回类型以及 ret 指令的语义，可以确定这里返回的是一个 i32 整数。
基本块的名字可以自行指定。基本块叫什么并不会影响程序语义，例如 %entry、%start 等都可以使用。不过，给基本块起一个具有实际含义的名字，通常能够让生成的 IR 更容易阅读，也更方便我们在后续调试编译器时定位问题。

生成 Koopa IR（1）——Koopa IR 数据结构

生成 Koopa IR 最直接的方式，其实和我们之前在命令行中输出 AST 很类似：给不同的 AST 节点实现对应的输出逻辑，然后遍历 AST，直接输出 Koopa IR 文本即可。

不过在正式实现之前，我们先区分 Koopa IR 的两种表示形式：

文本形式：也就是字符串形式，主要方便人阅读、调试以及在不同工具之间传递。
内存形式：也就是使用程序中的数据结构来表示 IR，方便编译器对其进行遍历、分析和修改。

我们的编译器最终需要输出文本形式的 Koopa IR。之后，这些文本 IR 可以被 koopac 等 Koopa IR 工具读取并解析为内存形式，以进行进一步处理。Koopa IR 框架本身也提供了相关接口，用于处理不同形式的 IR。

因此，考虑到上述情况，我们主要有两种实现思路：

遍历 AST，直接输出文本形式的 IR。这种方式最简单，也适用于使用任意语言实现的编译器。
像定义 AST 一样，定义一套表示 Koopa IR 的数据结构，例如指令、基本块和函数等。首先遍历 AST 生成这些 IR 数据结构，然后再遍历 IR，将其输出为字符串。

虽然第一种方案最简单，但从工程结构上看，它实际上把两个不同的过程：

1
2


AST -> IR
IR -> Text

杂糅在了一起，不利于后续扩展和维护。

因此，我们选择第二种方案。此时整个流程变成：

1
2
3
4
5
6
7
8
9


AST
 ↓
IR Generation
 ↓
Koopa IR 内存结构
 ↓
IR Printer
 ↓
Koopa IR 文本

也就是说，AST 不负责输出 Koopa IR 字符串，IR 数据结构也不需要理解 AST，二者之间通过 IR Generation，也就是 AST 到 IR 的转换过程连接起来。

目前，我们并不需要一开始就实现完整的 Koopa IR 数据结构。根据本次实验的要求，我们只需要支持如下结构：

1
2
3
4
5
6
7


Program
└── Function
 └── BasicBlock
 └── Value
 ├── Integer
 └── Instruction
 └── Return

下面我们来具体设计 Koopa IR 的数据结构。

根据 Koopa IR 的设计，一个完整的程序由若干 Function 和全局 Value 组成，而每个 Function 又由若干 BasicBlock 组成，BasicBlock 中包含一系列指令。

与此同时，指令本身也属于 Value 的一种。因此，Value 可以看作 Koopa IR 中最基础的抽象之一。

我们首先使用枚举表示不同种类的 Value。同时，为了表示 IR 中的数据类型，再额外定义 IRType：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


enum class ValueKind
{
 Integer,
 Return
};

enum class IRType
{
 I32,
};

其中需要注意：

ValueKind 表示当前 IR 节点“是什么”，例如整数常量还是 Return 指令。
IRType 表示数据本身的类型，例如这里的 I32。

有了 ValueKind 后，我们就可以定义所有 IR Value 的基类 IRValue：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


class IRValue
{
public:
 const ValueKind kind;
 virtual ~IRValue() = default;

protected:
 explicit IRValue(ValueKind kind)
 : kind(kind) {}
};

kind 用来记录当前 IRValue 的具体种类，而构造函数被声明为 protected，因为我们并不会直接创建一个普通的 IRValue，而是通过其派生类创建具体的 IR 节点。

在此基础上，可以进一步派生出当前实验需要的两类 Value：

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


class IRInteger : public IRValue
{
public:
 int value;

 explicit IRInteger(int value)
 : IRValue(ValueKind::Integer),
 value(value) {}
};

class Instruction : public IRValue
{
protected:
 explicit Instruction(ValueKind kind)
 : IRValue(kind) {}
};

其中，IRInteger 表示整数常量，而 Instruction 表示所有指令的公共基类。

这里体现了一个很重要的关系：

1
2
3


IRValue
├── IRInteger
└── Instruction

也就是说，所有指令都是 Value，但并不是所有 Value 都是指令。例如整数常量 IRInteger 是一个 Value，但它并不是一条指令。

而 Instruction 又可以进一步派生出具体的指令类型。例如本次实验中需要使用的 Return：

1
2
3
4
5
6
7
8
9


class Return : public Instruction
{
public:
 std::unique_ptr<IRValue> operand;

 explicit Return(std::unique_ptr<IRValue> operand)
 : Instruction(ValueKind::Return),
 operand(std::move(operand)) {}
};

这里的 operand 表示 Return 指令的返回值。

例如：

1

return 42;

对应的 IR 数据结构可以理解为：

1
2


Return
└── IRInteger(42)

需要注意的是，我们并没有把 operand 定义在 Instruction 基类中，而是让具体的指令自行维护自己的操作数。

这是因为不同指令拥有的操作数数量和含义并不相同。例如 Return 只有一个返回值，而后续的二元运算指令通常会有两个操作数。因此，把操作数交给具体指令管理会更加合理。

有了这些最基础的 IR 节点后，我们就可以进一步向上组织它们。

首先，多条指令共同组成一个 BasicBlock：

1
2
3
4
5
6
7
8
9


class BasicBlock
{
public:
 std::string name;
 std::vector<std::unique_ptr<Instruction>> insts;

 explicit BasicBlock(std::string name)
 : name(std::move(name)) {}
};

这里需要注意，BasicBlock 中保存的是：

1

std::vector<std::unique_ptr<Instruction>>

而不是所有 IRValue。

这是因为基本块本质上包含的是一系列需要依次执行的指令。虽然 Instruction 同时也是一种 IRValue，但像 IRInteger 这样的常量并不会单独作为一条指令出现在基本块的指令列表中。

若干个 BasicBlock 又共同组成一个 Function：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


class Function
{
public:
 std::string name;
 IRType ret_type;
 std::vector<std::unique_ptr<BasicBlock>> bbs;

 Function(std::string name, IRType ret_type)
 : name(std::move(name)),
 ret_type(ret_type) {}
};

其中：

name 表示函数名称；
ret_type 表示函数返回值类型；
bbs 保存函数中的所有基本块。

最后，若干个 Function 组成整个 Program：

1
2
3
4
5


class Program
{
public:
 std::vector<std::unique_ptr<Function>> funcs;
};

完整的 Koopa IR Program 实际上还可以包含全局 Value，不过在当前 Lv1 实验中并不会使用到，因此这里暂时只保存函数。

至此，我们就完成了当前实验所需要的 Koopa IR 数据结构设计：

1
2
3
4
5
6


Program
└── Function
 └── BasicBlock
 └── Instruction
 └── Return
 └── IRInteger

有了这套内存中的 IR 表示后，接下来就可以完成两件事情：

1
2
3
4
5
6
7


AST
 ↓
生成 Koopa IR 数据结构
 ↓
IR Printer
 ↓
Koopa IR 文本

也就是说，下一步真正需要实现的，就是 AST 到 Koopa IR 的转换。

生成 Koopa IR（2）——Koopa IR 转换与打印

从 AST 到 Koopa IR 的转换并不复杂，在实现思路上，我们完全可以参考之前输出 AST 时使用的 Dump：从根节点开始，递归遍历整棵 AST，并根据不同类型的 AST 节点生成对应的 IR 节点。

不过，为了保持上一节中确定的设计，我们并不会直接在 AST 中实现 IR 生成逻辑，而是单独设计一个 IRGenerator，负责：

1

AST -> Koopa IR

同时，再使用 IRPrinter 负责：

1

Koopa IR -> Text

这样做虽然相比直接在 AST 中添加一个 DumpIR 方法多了一层结构，但也带来了一个明显的好处：

如果我们想了解 AST 是如何转换成 IR 的，直接查看 IRGenerator 即可；
如果我们想了解 IR 是如何被打印成 Koopa IR 文本的，直接查看 IRPrinter 即可。

这样就不需要在 AST 的各种节点中频繁切换上下文，也使不同模块之间的职责更加清晰。

整个过程可以表示为：

1
2
3
4
5
6
7
8
9


AST
 ↓
IRGenerator
 ↓
Koopa IR 内存结构
 ↓
IRPrinter
 ↓
Koopa IR 文本

我们先来看AST 到 Koopa IR 的转换。实现上，我们可以单独创建：

1
2


src/ir_generator.hpp
src/ir_generator.cpp

然后采用类似 Dump 的递归思路，根据不同的 AST 节点生成对应的 Koopa IR 数据结构。

首先，从最外层的 CompUnitAST 开始：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


Program IRGenerator::Generate(const CompUnitAST& ast) const
{
 Program program;

 const auto& func =
 static_cast<const FuncDefAST&>(*ast.func_def);

 program.funcs.push_back(GenerateFunction(func));

 return program;
}

当前实验中，一个 CompUnitAST 中只包含一个函数定义，因此我们取出其中的 FuncDefAST，生成对应的 Function，并加入 Program 中。

接下来处理函数定义：

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


std::unique_ptr<Function>
IRGenerator::GenerateFunction(const FuncDefAST& ast) const
{
 auto func =
 std::make_unique<Function>(ast.ident, IRType::I32);

 auto entry =
 std::make_unique<BasicBlock>("entry");

 const auto& block =
 static_cast<const BlockAST&>(*ast.block);

 EmitBlock(block, *entry);

 func->bbs.push_back(std::move(entry));

 return func;
}

在当前实验中，函数返回类型一定是 int，因此可以直接将其转换为 Koopa IR 中的 I32。

与此同时，我们为函数创建一个名为 entry 的入口基本块，然后将函数体中的语句转换成对应的指令，并插入这个基本块。

对于 BlockAST：

1
2
3
4
5
6
7
8
9


void IRGenerator::EmitBlock(
 const BlockAST& ast,
 BasicBlock& bb) const
{
 const auto& stmt =
 static_cast<const StmtAST&>(*ast.stmt);

 EmitStmt(stmt, bb);
}

当前实验中，一个 Block 中只有一条 Stmt，因此继续递归处理即可。

接下来处理 StmtAST：

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


void IRGenerator::EmitStmt(
 const StmtAST& ast,
 BasicBlock& bb) const
{
 const auto& number =
 static_cast<const NumberAST&>(*ast.number);

 bb.insts.push_back(
 std::make_unique<Return>(
 GenerateNumber(number)
 )
 );
}

当前语法中的 Stmt 只有：

1

return Number;

因此，我们首先将 NumberAST 转换成一个 IRInteger，然后将其作为操作数构造 Return 指令，最后把这条指令加入当前基本块。

最后是最底层的 NumberAST：

1
2
3
4
5


std::unique_ptr<IRValue>
IRGenerator::GenerateNumber(const NumberAST& ast) const
{
 return std::make_unique<IRInteger>(ast.value);
}

这样，一个类似：

1
2
3


int main() {
 return 0;
}

的 AST：

1
2
3
4
5
6


CompUnitAST
└── FuncDefAST
 ├── main
 └── BlockAST
 └── StmtAST
 └── NumberAST(0)

经过 IRGenerator 后，就会得到：

1
2
3
4
5


Program
└── Function("main")
 └── BasicBlock("entry")
 └── Return
 └── IRInteger(0)

这就完成了：

1

AST -> Koopa IR

的转换。

接着我们看Koopa IR 的打印。有了内存形式的 Koopa IR 后，接下来只需要按照 Koopa IR 的文本格式，将这些数据结构依次输出即可。

其过程同样是从 Program 开始递归向下遍历：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


std::string IRPrinter::Print(const Program& program) const
{
 std::string result;

 for (const auto& func : program.funcs)
 {
 result += PrintFunction(*func);
 }

 return result;
}

然后打印 Function：

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


std::string
IRPrinter::PrintFunction(const Function& func) const
{
 std::string result;

 result += "fun @" + func.name + "(): ";

 if (func.ret_type == IRType::I32)
 {
 result += "i32";
 }

 result += " {\n";

 for (const auto& bb : func.bbs)
 {
 result += PrintBasicBlock(*bb);
 }

 result += "}\n";

 return result;
}

再依次打印函数中的基本块：

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


std::string
IRPrinter::PrintBasicBlock(const BasicBlock& bb) const
{
 std::string result;

 result += "%" + bb.name + ":\n";

 for (const auto& inst : bb.insts)
 {
 result += " ";
 result += PrintInstruction(*inst);
 result += "\n";
 }

 return result;
}

对于具体指令，我们根据 ValueKind 判断指令类型：

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


std::string
IRPrinter::PrintInstruction(const Instruction& inst) const
{
 switch (inst.kind)
 {
 case ValueKind::Return:
 {
 const auto& ret =
 static_cast<const Return&>(inst);

 return "ret " + PrintOperand(*ret.operand);
 }

 default:
 throw std::logic_error("invalid instruction kind");
 }
}

这里虽然传入的是 Instruction&，但在判断出它的 kind 为 Return 后，我们就知道这个对象实际是一个 Return，因此可以通过 static_cast 将其转换为具体的 Return 类型，从而访问其中的 operand。

同样，对于操作数：

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


std::string
IRPrinter::PrintOperand(const IRValue& value) const
{
 switch (value.kind)
 {
 case ValueKind::Integer:
 {
 const auto& integer =
 static_cast<const IRInteger&>(value);

 return std::to_string(integer.value);
 }

 default:
 throw std::logic_error("invalid operand kind");
 }
}

如果这个 IRValue 的类型为 Integer，那么它实际对应的是 IRInteger，因此将其转换成 IRInteger 后即可取得其中保存的整数值。

至此，我们就完成了：

1
2
3
4
5


Koopa IR 内存结构
 ↓
IRPrinter
 ↓
Koopa IR 文本

这一过程。

最后我们来看如何在主程序中使用。实现完 IRGenerator 和 IRPrinter 后，最后只需要在 main 中将它们串联起来：

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


IRGenerator generator;
auto program =
 generator.Generate(
 static_cast<const CompUnitAST&>(*ast)
 );

IRPrinter printer;
auto koopa = printer.Print(program);

std::ofstream koopa_file(output);
assert(koopa_file);

koopa_file << koopa;

这样，我们完整的编译流程就变成了：

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


SysY 源代码
 ↓
词法分析 / 语法分析
 ↓
AST
 ↓
IRGenerator
 ↓
Koopa IR 内存结构
 ↓
IRPrinter
 ↓
Koopa IR 文本

对于：

1
2
3


int main() {
 return 0;
}

最终就可以生成：

1
2
3
4


fun @main(): i32 {
%entry:
 ret 0
}

至此，我们就完成了从 SysY AST 到 Koopa IR 的第一次完整转换。

至此，Exp Chapter 1的实验就全部完成了！我们成功的将一个简单的带有main函数的SysY代码编译为了Koopa IR程序！

[从零开始的编译原理][理论] Chapter 1：词法分析

Mon, 28 Sep 2026 20:15:14 +0800

这一章我们学习的是词法分析。

我们知道编译器通常由以下几个部分组成:

前端: 通过词法分析和语法分析，将源代码解析为抽象语法树（AST）。通过语义分析，扫描抽象语法树，检查其是否存在语义错误
中端：将AST转化为中间表示IR，并在此基础上完成一些机器无关优化
后端：将中间表示IR转换为目标平台的汇编代码，并在此基础上完成一些机器相关优化

词法分析

词法分析简单来说就是把字节流转换为单词流 (token stream). 词法分析器(lexer)会按照某种规则读文件，并将文件的内容拆分成一个个 token 作为输出, 传递给语法分析器 (parser). 同时, lexer 还会忽略文件里的一些无意义的内容, 比如空格, 换行符和注释.

Lexer 生成的 token 会包含一些信息, 用来让 parser 区分 token 的种类, 以及在必要时获取 token 的内容.

例如我们有这样一个程序:

1
2
3
4
5


int main() 
{
// 我是注释诶嘿嘿
return 0;
}

那么Lexer可能会将他转化为如下token 流:

1
2
3
4
5
6
7
8
9


种类: 关键字, 内容: int.
种类: 标识符, 内容: main.
种类: 其他字符, 内容: (.
种类: 其他字符, 内容: ).
种类: 其他字符, 内容: {.
种类: 关键字, 内容: return.
种类: 整数字面量, 内容: 0.
种类: 其他字符, 内容: ;.
种类: 其他字符, 内容: }.

目前为止，我们实现词法分析从易到难有三种方法：

词法分析器生成器
手写词法分析器
自动化词法分析器

接下来我们重点看看前面两项

词法分析器生成器

在本课程中，我们使用的词法分析生成器是antlr。

在使用antlr时，我们的输入是一个包含词法单元规约的g4文件，那么antlr会自动生成一个词法分析器

如果我们使用antlr这样的语法分析器生成器，我们主要做的事情就是就是写清楚语法规则。

我们接下来尝试用antlr实现一个类似C++的词法分析器:

首先我们知道一个C++程序肯定是由若干条Statement组成的，那么我们就有第一条语法规则：

1
2
3


grammer SimpleSysY; // 注意在g4中第一行需要定义grammer，并且保持和文件名一致

prog: stat* EOF; //这里加入EOF来说明结束条件

而statement都是由表达式expr组成的，为了简便，我们先初步规定expr只支持赋值语句和输出，那么我们可以有下一条规则:

1
2
3
4


stat: expr ';' // 需要注意''内的分号说明一条statement要以分号结尾
 | ID '=' expr ';'
 | 'print' expr ';'
 ;

接下来我们就可以去描述一个表达式的语法规则了。那么显然，这是一个递归的描述:

1
2
3
4


expr: expr ('+' | '-' | '*' | '/' ) expr
 | '(' expr ')'
 | ID
 ;

但是这样写会有点问题，这里会涉及到符号优先级的问题，例如乘法和除法的优先级肯定是高于加法和减法的，想要在antlr中表达优先级我们可以:

1
2
3
4
5


expr: expr ('*' | '/' ) expr
 | expr ('+' | '-' ) expr
 | '(' expr ')'
 | ID
 ;

当然除了优先级以外，我们还要考虑结合性的影响，在g4中，如果我们不显式指定是左结合还是右结合，就会默认是左结合，这和我们的计算是符合的，但是我们后续还会继续探讨这部分的内容。

整体上来看，语法结果就是这样，接下来我们看看词法结构:

我们知道在C++中一个标识符是以下划线或字母开头，以字母数字下划线结尾的字符串:

1

ID: ('_' | [a-zA-Z])('_' | [a-zA-Z0-9])*;

而如果此时我们用上面的写好的语法规则和词法规则去对一个语句去做测试: a = b + c;

会发现有部分的报错，这是因为空格没有被识别到，因此我们还需要加入空格的词法规则:

1

WS: [ \t\r\n]+ -> skip;

其中我们用+是因为至少要有一个才能被识别，而后面的skip则是识别到了就跳过，否则我们实际上grammar中没有相关的规则，会导致语法规则识别不出来。

这样，上面那个语句我们就可以画出语法分析树了:

当然在我们实际的实验中，这样简单的语法规则和词法规则肯定是不够的，具体的我们可以在Lab中实现，这里只做简单的介绍。

另外这里可以补充一点写语法规则和词法规则时一些需要注意的点:

特殊的规则尽量写在前面，比如docs comment和mul comment的匹配规则分别写作:

1
2


DOCS_COMMENT : '/**' .*? '*/';
MUL_COMMENT : '/*' .*? '*/'; // 其中 .*?中的？表示非贪婪匹配模式

那么此时明显可以看到DOCS_COMMENT也可以被MUL_COMMENT给匹配到的，因此我们在写词法规则的时候应该把DOCS_COMMENT的规则写在MUL_COMMENT前面，让前者优先匹配。

在antler这类词法分析器生成器中，提供了fragment的功能来进行助记:

1
2
3
4


fragment LETTER: [a-zA-Z]
fragment NUMBER: [0-9]
fragment WORD: '_' | LETTER | NUMBER
ID: ('_' | LETTER)(WORD)*;

antlr中有三个比较重要的优先匹配规则来解决冲突，分别是最前优先匹配，最长优先匹配(例如1.23会被匹配为float而不是INT和FLOAT，>=不会被识别为>和=)，非贪婪匹配

词法分析生成器的基本原理（1）

在Antlr这类词法分析器生成器中，我们只需要在g4文件中描述词法单元的正则表达式就能自动生成词法分析器，这背后的流程可以分为若干个步骤。

比较经典的流程是:

Regix->NFA->DFA->Transition Table->Machine Code

接下来我们会逐一介绍上面的各个部分

正则表达式

首先我们给出正则表达式的定义:

给定字母表$\Sigma$，$\Sigma$上的正则表达式有且仅有以下规则定义:

$\epsilon$是正则表达式
$\forall a \in \Sigma $,a是正则表达式
如果r是正则表达式，则(r)也是正则表达式
如果r与s都是正则表达式，则 r|s,rs,r*也是正则表达式

此外，我们规定正则表达式中的运算优先级如下:

$$ () > * > connect > | $$

对应的正则语言如下:

我们规定每个正则表达式r对应一个正则语言L(r)

$L(\epsilon) = \{\epsilon\}$
$L(a) = \{a\},\forall a \in \Sigma$
$L((r)) = L(r)$
$L(r|s)=L(r)\cup L(s)$ $L(rs)=L(r)L(s)$ $L(r*)=(L(r))*$

下面是常用到的正则表达式符号以及对应的含义:

符号	含义	例子	例子含义
`a`	匹配字符本身	`a`	匹配字符 a
`.`	匹配除换行符以外的任意单个字符	`a.b`	如 acb、a1b
`r\|s`	或 / 选择	`a\|b`	匹配 a 或 b
`r*`	前面的表达式重复 0 次或多次	`a*`	“"、a、aa、…
`r+`	前面的表达式重复 1 次或多次	`a+`	a、aa、…
`r?`	前面的表达式出现 0 次或 1 次	`a?`	"” 或 a
`(r)`	和r相同	`(ab)*`	“"、ab、abab、…
`[s]`	字符集合，匹配其中任意一个字符	`[abc]`	a、b 或 c
`[^s]`	字符集合取反	`[^abc]`	除 a、b、c 外的字符
`^`	通常表示字符串/行开头	`^abc`	以 abc 开头
`$`	通常表示字符串/行结尾	`abc$`	以 abc 结尾
`\`	转义特殊字符	`\*`	匹配字符 * 本身
`r{m}`	恰好重复 m 次	`a{3}`	aaa
`r{m,n}`	重复 m 到 n 次	`a{2,4}`	aa、aaa、aaaa
`"s"`	串s的字面值	`"*a"`	*a

NFA 非确定性有穷自动机

我们定义非确定性有穷自动机A是一个五元组 $A = (\Sigma,S,s_0,\delta,F)$:

字母表 $\Sigma\ \ (\epsilon \notin \Sigma)$
有穷的状态集合 S
唯一的初始状态 $s_0$
状态转移函数 $\delta$：$\delta: S\times(\Sigma \cup \{\epsilon\})\to 2^S$
接受状态集合 $F \subseteq S$

NFA的主要特点是一个输入可以有多个下一状态，且允许 $\epsilon$转移

自动机A定义了一种语言 L(A),它能接受的所有字符串构成的集合

DFA 确定性有穷自动机

我们定义确定性有穷自动机A是一个五元组 $A = (\Sigma,S,s_0,\delta,F)$:

字母表 $\Sigma\ \ (\epsilon \notin \Sigma)$
有穷的状态集合 S
唯一的初始状态 $s_0$
状态转移函数 $\delta$：$\delta: S\times \Sigma \to S$
接受状态集合 $F \subseteq S$

与NFA不同的点在于DFA一个输入只有一个状态，不允许空转移

NFA简洁易于理解，便于描述语言L(A)，而DFA容易判断 $x \in L(A)$,适合产生词法分析器

因此我们通常用NFA描述语言，DFA实现词法分析器

词法分析生成器的基本原理（2）

在上一节中我们提到了Regix以及NFA，DFA的基本含义，我们接下来介绍他们之间的相互转化。

从RE到NFA: Thompson 构造法

Thompson构造法的基本思想是按结构归纳，即把正则表达式递归地拆分成小的子表达式，每个子表达式先构造一个小NFA，再通过 $\epsilon$边把这些小的NFA连接起来。

依据正则表达式的定义，我们可以拆开看:

基本字符:

对于一个字符a，我们可以构造:

1
2


 a
(q0) ---> (q1)

其中q0是入口，q1是出口

连接:rs

假设我们已经构造好了正则表达式r和s的NFA，那么对于连接我们可以通过将r的出口和s的入口以一个 $\epsilon$边连接即可:

1
2


 r s
 --> [ NFA ] --ε--> [ NFA ] -->

选择 r|s

我们此时希望NFA可以选择走r的入口，也可以走s的入口，那么我们只需要新建一个入口和一个新的出口，并用 $\epsilon$边将他们分别与r和s的入口，出口连接即可

1
2
3
4
5


 ε --> [NFA(r)] --ε
 / \
start --- ---> end
 \ /
 ε --> [NFA(s)] --ε

Kleene 星号: r*

对于Kleene星号，我们需要具备两种能力，即一次都不执行和执行完以后重新执行:

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


 ε
 +-------------> end
 |
 start
 |
 | ε
 v
 start_r ---- r ----> end_r
 ^ |
 |_______ ε _________|
 |
 ε
 |
 v
 end

前者我们通过将新的入口和新的出口直接通过 $\epsilon$边连接完成，后者我们则在r的NFA上加入一条由出口指向入口的 $\epsilon $边。

我们可以通过下面这个例子来展示Thompson构造法转化RE到NFA的完整流程:

假设我们有正则表达式a(b|c)*

我们可以先按照语法树来决定构造顺序：

1
2
3
4
5
6
7
8
9


 连接
 / \
 a *
 |
 |
 |
 b|c
 / \
 b c

然后我们可以自底向上构造，首先构造一下a,b,c：

1
2
3
4
5


a → NFA(a)

b → NFA(b)

c → NFA(c)

接下来构造b|c:

1
2
3


 /-- b --\
-- ε - - ε -->
 \-- c --/

然后给b|c套上Kleene星号:

1
2
3
4
5
6
7


 ε----------------------+
 | |
 v |
start --> [ b|c ] --> merge ---+
 | |
 +----------ε----------------> end
 merge --ε-----> end

最后再把a的按照连接的方式连起来:

1

start --> [a] --> [(b|c)*] --> end

从NFA到DFA:子集构造法

子集构造法的核心思想在于DFA的一个状态，对应于NFA的一组状态。

由于DFA在描述状态转移时只允许目标状态有且仅有一个，而NFA可以有多个，因此我们干脆把NFA中所有可能的状态整体看成DFA的一个状态。

在描述具体方法前，我们通过上面的NFA介绍两个名词:

$\epsilon-\text{closure}$：表示在NFA中从状态集合S出发，只走任意条 $\epsilon$边，能够到达的状态，包括S自己。

例如下图中，从状态 $q_0$出发，只走任意条 $\epsilon$边，能够到达的状态为: $\{q_0,q_1,q_2,q_4,q_7\}$,因此我们有: $\epsilon-\text{closure}(\{q_0\})=\{q_0,q_1,q_2,q_4,q_7\}$

Move: 表示在NFA中从状态集合S出发，读取字符 a一步后能到达的状态集合

借用这两个名词，以及上面这个例子，我们可以很好的描述子集构造法的流程了。

首先，从初始状态出发，计算: $\epsilon-\text{closure}(q_0)$,计算出来后记作初始状态 $D_0 = \{q_0,q_1,q_2,q_4,q_7\}$

之后从 $D_0$读字符a，算 $\epsilon-\text{closure}(\text{move}(D_0,a)) = D_1 = \{q_1,q_2,q_3,q_4,q_6,q_7,q_8\}$

接着从 $D_0$读字符b，得到 $D_2 = \{q_1,q_2,q_4,q_5,q_6,q_7\}$

再从 $D_1$读字符a，算 $\epsilon-\text{closure}(\text{move}(D_1,a)) = \{q_1,q_2,q_3,q_4,q_6,q_7,q_8\} = D_1$

再从 $D_1$读取字符b，算 $\epsilon-\text{closure}(\text{move}(D_1,b))=D_3=\{q_1,q_2,q_3,q_4,q_5,q_6,q_7,q_9\}$

然后从 $D_2$读字符a，得到 $D_1$，从 $D_2$读字符b，得到 $\epsilon-\text{closure}(\text{move}(D_2,b)) = D_2 = \{q_1,q_2,q_4,q_5,q_6,q_7\}$

最后从 $D_3$读字符a，得到 $D_1$ ，从 $D_3$读字符b，得到 $\epsilon-\text{closure}(\text{move}(D_3,b)) = D_4^* = \{q_1,q_2,q_4,q_5,q_6,q_7,q_{10}\}$,因为它包含终态 $q_{10}$，所以我们会给它一个特殊的标记。

而从 $D_4$读字符a和字符b分别可以得到 $D_1,D_2$

总结下来，我们可以得到如下状态转移图:

DFA 状态	对应 NFA 状态集合	a	b
$D_0$	{0,1,2,4,7}	$D_1$	$D_2$
$D_1$	{1,2,3,4,6,7,8}	$D_1$	$D_3$
$D_2$	{1,2,4,5,6,7}	$D_1$	$D_2$
$D_3$	{1,2,4,5,6,7,9}	$D_1$	$D_4$
$D_4^*$	{1,2,4,5,6,7,10}	$D_1$	$D_2$

转化为DFA，可以画成:

最小化DFA

DFA最小化的目标是在保持语言完全不变的前提下，把行为等价的DFA状态合并，得到状态数最少的DFA。

首先，我们需要明确什么叫两个状态是等价的:

我们假设DFA中存在两个状态p,q。从他们出发，如果输入任意的字符串w，最终要么都接受，要么都拒绝，那么我们可以认为: $p\equiv q $。

知道了怎样的状态是等价，也就是可以合并的，我们可以开始考虑怎么最小化DFA了。

首先我们可以知道的是，接受状态和非接受状态是不可合并的，因此我们可以把DFA的状态集合分为两组:

$$ P_0 = \{F,Q-F\} $$

接下来则可以进一步不断细分，我们对于同一组中的两个状态p和q，观察他们读入每个字符后的去向。如果:

$$ \delta(p,a) =\epsilon-\text{closure}(\text{move}(p,a)) \neq \delta(q,a) $$

那么此时p和q不能继续待在一个组，于是把他们拆开。

不断重复上述流程，直至无法继续拆分为止。

最终每一个分组中状态都是等价状态，可以合并为一个DFA状态。

为了方便我们理解，我们利用我们上面根据子集构造法得到的DFA来做一遍上述流程。

状态	a	b	是否接受
$D_0$	$D_1$	$D_2$	否
$D_1$	$D_1$	$D_3$	否
$D_2$	$D_1$	$D_2$	否
$D_3$	$D_1$	$D_4$	否
$D_4$	$D_1$	$D_2$	是

其中唯一接受状态： $F = \{D_4 \}$

第一步，按接受/不接受进行第一步划分，得到:

$$ P_0 = \{ \{D_4\},\{D_0,D_1,D_2,D_3\}\} $$

为了方便，我们记作: $P_0 = \{A,B\},A=\{ D_4\}$

第二步，检查B是否可以进一步拆分。

$D_0$: 根据表格，我们发现无论输入a还是b，最后仍然落入B中，因此我们记作: $D_0:(B,B)$
$D_1$: 类似的，我们可以得到: $D_1:(B,B)$
$D_2$: $D_2:(B,B)$
$D_3$: $D_3(B,A)$

于是我们发现只有 $D_3$与B中的其他状态不同，因此我们可以进一步划分出:

$$ P_1 = \{\{D_4\},\{D_3\},\{D_0,D_1,D_2\}\} $$

接着，我们继续检查 $C=\{D_0,D_1,D_2\}$,按照上面的流程，我们可以得到:
$D_0:(C,C)$ $D_1:(C,B)$ $D_2:(C,C)$

发现 $D_1$，不同因此继续拆分: $P_2 = \{\{D_4\},\{D_3\},\{D_1\},\{D_0,D_2\}\}$

最后，我们可以发现 $D_0,D_2$显然是不可拆分的，他们是等价的，因此我们得到了最小的DFA状态:

$$ S_0 = \{D_0,D_2\},S_1=\{D_1\},S_2=\{D_3\},S_3=\{D_4\} $$

其中 $S_3$是终态，我们可以得到状态转移表格:

状态	a	b
$S_0$	$S_1$	$S_0$
$S_1$	$S_1$	$S_2$
$S_2$	$S_1$	$S_3$
$S_3^*$	$S_1$	$S_0$

画成DFA就是

[从零开始的编译原理][实验] Exp Chapter 0：实验环境搭建

Mon, 28 Sep 2026 16:36:25 +0800

Docker配置

根据课程指导书的要求，我们通过Docker拉取对应的镜像:

1

docker pull maxxing/compiler-dev

下面简单介绍一下Docker的使用方法。

首先我们需要简单区分一下镜像(image)和容器(container),前者可以看作是一个硬盘，里面装好了操作系统，但它是静态的，不能直接拿来运行。而后者则是一台计算机，里面装了硬盘，就能运行对应的操作系统。

在实际情况中，我们可以在容器中对文件系统进行修改，例如创建和删除文件，但是镜像不会受到影响。

通常而言，我们只希望在镜像的环境做一些一次性的工作，比如用里面的测试脚本来测试自己的编译器，然后查看测试结果。在此之后这个临时容器就没有任何作用了。

例如我们希望在上面image的环境中执行如下命令:ls -l /，那么我可以输入指令:

1

docker run maxxing/compiler-dev ls -l /

然后此时终端就会显示命令ls -l /的结果，此时会生成一个临时容器，我们可以用指令:docker ps -a查看:

1
2
3


user@macbook compiler % docker ps -a 
CONTAINER ID IMAGE COMMAND CREATED STATUS PORTS NAMES
441f432b9731 maxxing/compiler-dev "ls -l" 3 seconds ago Exited (0) 3 seconds ago hopeful_shannon

那么我们可以通过指令:docker rm 441f432b9731来删除这个容器。

当然我们还可以将上述执行命令-删除临时容器的操作合并为一个指令:

1

docker run --rm maxxing/compiler-dev ls -l /

这样创造临时容器，运行完指令后，就会直接把临时容器给删掉了。

有时候，可能我们期望在容器中执行的命令很复杂，不能一行输入完整，那么此时我们可以使用

1

docker run -it --rm maxxing/compiler-dev bash

这样我们就创建了临时容器后，执行了bash指令，就可以在容器内的shell中工作输入任何指令了。其中-it参数是为了开启容器的stdin。

如果需要退出，直接输入exit即可

在大部分情况下，我们还希望Docker容器能够访问宿主机中的文件，如我们实现的编译器位于宿主机的/path/to/compiler目录下，我们希望 Docker 容器也能访问到这个目录里的内容, 这样你就可以使用容器中的测试脚本测试你的编译器了. 可以执行:

1

docker run -it --rm -v /path/to/compiler:/root/compiler maxxing/compiler-dev bash

这条命令多了一个-v参数，它的作用是把宿主机的某个目录挂载到容器的某个目录（例如上面指令是/root/compiler）。这样，在进入容器之后，就可以通过访问挂载的目录访问宿主机的目录了。

Koopa IR介绍

Koopa IR是一种简化后的中间表示，在设计上类似LLVM IR。同时我们在后续lab中会用到对应的框架:

https://github.com/pku-minic/koopa

Koopa IR 是一种强类型的 IR, IR 中的所有值 (Value) 和函数 (Function) 都具备类型 (Type).

Koopa IR 中, 基本块 (basic block) 必须是显式定义的. 即, 在描述函数内的指令时, 你必须把指令按照基本块分组, 每个基本块结尾的指令只能是分支/跳转/函数返回指令之一. 在 IR 的数据结构表示上, 指令也会被按照基本块分类. 这很大程度上方便了 IR 的优化, 因为许多优化算法都是在基本块的基础上对程序进行分析/变换的.

假设我们有一个用Koopa IR编写的程序（见指导书中的"hello world!"），我们将其保存在了文件hello.koopa中，那么我们可以在实验环境中运行这个koopa IR程序:

1
2
3


koopac hello.koopa | llc --filetype=obj -o hello.o
clang hello.o -L$CDE_LIBRARY_PATH/native -lsysy -o hello
./hello

RISC-V介绍

我们的编译器最终会生成RISC-V汇编。

RISC-V的指令系统由

基础指令系统 (base instruction set) 和指令系统扩展 (extension) 构成. 每个 RISC-V 处理器必须实现基础指令系统, 同时可以支持若干扩展. 常用的基础指令系统有两种:

RV32I: 32 位整数指令系统.
RV64I: 64 位整数指令系统. 兼容 RV32I.

常用的标准指令系统扩展包括:

M 扩展: 包括乘法和除法相关的指令.
A 扩展: 包括原子内存操作相关的指令.
F 扩展: 包括单精度浮点操作相关的指令.
D 扩展: 包括双精度浮点操作相关的指令.
C 扩展: 包括常用指令的 16 位宽度的压缩版本.

我们通常使用 RV32/64I + 扩展名称的方式来描述某个处理器/平台支持的 RISC-V 指令系统类型, 例如 RV32IMA 代表这个处理器是一个 32 位的, 支持 M 和 A 扩展的 RISC-V 处理器.

在我们的Lab中，我们编译器将生成RV32IM范围内的RISC-V汇编。

假设我们有一段RISC-V的汇编程序:hello.S，我们可以在实验环境中将这个RISC-V程序汇编并链接为可执行文件并运行起来:

1
2
3


clang hello.S -c -o hello.o -target riscv32-unknown-linux-elf -march=rv32im -mabi=ilp32
ld.lld hello.o -L$CDE_LIBRARY_PATH/riscv32 -lsysy -o hello
qemu-riscv32-static hello

编程语言选择

为了方便起见，我最终还是选择使用C++进行开发。

我们使用CMake模版:https://github.com/pku-minic/sysy-cmake-template

配置好环境后，我们可以进入docker环境中，执行命令:

1
2


cmake -DCMAKE_BUILD_TYPE=Debug -B build
cmake --build build

构建编译器，后续如果我们更改了代码，则仅仅需要通过:

1

cmake --build build

即可。

SOAR: Scale Optimization for Accurate Reconstruction in NVFP4 Quantization

Wed, 26 Aug 2026 09:00:00 +0800

对于NVFP4量化，我们可以将其量化行为用下述公式描述:

对于高精度张量X，设$M_{\text{FP4}},M_{\text{FP8}}$分别表示FP4(E2M1)与FP8(E4M3)能够表示的最大值，分别为6和448，$Q_{\text{E4M3}}$表示FP8(E4M3)的量化函数，$\alpha,\Delta_i$分别表示全局scale以及块级scale:

$$ \alpha = \frac{\max (\abs{X})}{M_{\text{FP4}}\cdot M_{\text{FP8}}}\\ \Delta_i = Q_{\text{E4M3}}(\frac{\max (\abs{X_i})}{\alpha \cdot M_{\text{FP4}}})\\ \hat X = \begin{cases} \frac{1}{2}「\frac{2X}{\alpha \Delta}」\quad \abs{\frac{X}{\alpha \Delta}} < 2\\ 「\frac{X}{\alpha \Delta}」 \quad 2\leq\abs{\frac{X}{\alpha \Delta}}\leq 4\\ 2「\frac{X}{2\alpha \Delta}」\quad 4\leq\abs{\frac{X}{\alpha \Delta}}\leq 6 \end{cases} $$

这里之所以$\hat X$是分段函数，是因为在FP4量化中，在数值0-2的区间内，采样点的间隔是0.5（0，0.5，1，1.5），而在2-4内采样点间隔是1，而4-6内为2。

那么反量化函数可以写作:

$$ \hat X = \hat X \cdot(\alpha \Delta) $$

然而目前的NVFP4量化方法普通通过简单的启发式公式来确定全局scale 和块级scale，例如基于最大值的缩放方法，或仅在有限候选集上进行离散搜索。这写方法往往会导致次优结果。尤其是对于FP32全局scale 而言，如果只将其优化限制在固定规则或粗粒度的离散空间内，就很难准确刻画大语言模型复杂的权重分布。

位了解决这个问题，作者提出了闭式联合Scale优化方法，即CJSO。简单来说就是直接通过最小化原始权重W和反量化权重$\hat W$之间的重构误差来联合优化全局scale和块级scale:

$$ \alpha,\Delta = \arg \min_{\alpha,\Delta}\sum_i ||W_i - Q_i\cdot(\alpha \Delta_i)||_2^2 $$

其中:

$$ Q_i = Q_{\text{fp4}}(\frac{W_i}{\alpha \Delta_i}) $$

由于直接最小化这个重构误差是比较困难的，因为$Q_i$是离散的，并且同时依赖于$\alpha,\Delta_i$.不过，在固定$Q_i$的分配结果时，重构目标可以看作是关于缩放因子的二次函数，因此我们可以求解得到:

全局Scale优化，在给定块级scale$\Delta_i$的情况下，最优化张量级scale$\alpha^*$可以表示为:

$$ \alpha^* = \frac{\sum_{i=1}^N \sum_{j\in \text{block}_i}W_{ij}Q_{ij}\alpha}{\sum_{i=1}^N\sum_{j\in\text{block}_i}Q_{ij}^2 \alpha^2} $$

块级Scale优化，反过来，在$\alpha$固定的情况下，每个块级scale $\Delta_i^*$都可以独立优化，以拟合对应块内的局部分布:

$$ \Delta_i^* = \frac{\sum_{j\in \text{block}_i} W_{ij}Q_{ij}\alpha}{\sum_{j\in\text{block}_i}Q_{ij}^2 \alpha^2} $$

那么基于此，可以得到CJSO的量化方案:

首先使用NVFP4中标准的基于最大值的规则来初始化$\alpha,\Delta_i$。随后，用上述公式迭代更新全局scale $\alpha$以及块级scale$\Delta_i$,以及量化矩阵$Q_i$:每次更新后，都会在新的scale下重新执行FP4量化从而重新计算Q

除此之外，在NVFP4中，块级缩放因子还受限于FP8精度的表示限制，因此它本身也会引入量化误差。这种量化误差会传播到量化以及反量化两个阶段，从而直接影响重构精度。

因此论文考虑将量化与反量化使用的scale 解耦，具体而言，将原本的单一尺度块级缩放因子解耦为两个独立变量:

量化尺度$\Delta_i^q$（高精度）:用于确定FP4权重分配，不受硬件精度限制
反量化尺度$\Delta_i^d$（E4M3）:实际存储的硬件兼容尺度，用于推理时重建

那么，我们将NVFP4量化重新表述为:

$$ \min_{\Delta_i^q \in \mathbb{R},\Delta_i^d \in \text{FP8}} L = \sum_i||W_i - Q_{\text{FP4}}(\frac{W_i}{\alpha \Delta_i^q})\cdot (\alpha \Delta_i^d)||_2^2 $$

基于上述解耦，我们对两个块scale进行联合优化。

初始化：使用CJSO结果初始化$\Delta_i^d$,并令$\Delta_i^q = \Delta_i^d$
构造候选空间:
- 对于$\Delta_i^q$,在$[0.5,1.5]$范围内以步长0.01进行连续乘性扰动(高精度搜索)
- 对于$\Delta_i^d$则限制为当前值附近最近两个E4M3表示值
对每个候选$(\Delta_i^q,\Delta_i^d)$计算重建误差，选择最优组合

[2026 ACL] ARCQuant: Boosting NVFP4 Quantization with Augmented Residual Channels for LLMs

Wed, 26 Aug 2026 08:57:00 +0800

目前把使用NVFP4对大语言模型进行PTQ有一下难点:

NVFP4格式与现有的权重-激活值PTQ方法不兼容:
1. 旋转类方法，如QuaRot等会把异常值能量扩散到所有块，破坏NVFP4本身块隔离优势，反而放大局部动态范围
2. 平滑类方法，如Smoothquant等在低比特下几乎失效
混合精度方案硬件不兼容，Tensor Core要求数据格式统一，采用混合精度必然导致模型吞吐量下降

如图所示，虽然基于旋转矩阵的方法可以降低全局峰值，但是会显著增加原本低幅值 block 的局部动态范围。这个影响抵消了细粒度 scaling 原本带来的离群值隔离优势。导致在NVFP4上效果不佳。

为了解决上述问题，ARCQuant提出了一个“不改动数值格式，不混合精度，不旋转矩阵”的NVFP4量化框架。

具体实现如下:

离线选出"必须补偿"的异常通道:用校准集统计每层激活的通道最大值，按绝对值降序重排；以FP8（E5M2）动态范围作为参考，设阈值$\tau = 2^{-3}M$，只取超过该阈值的前S条通道作为"残差通道"。
对输入X先重拍然后进行量化得到$Q_x$,然后对选出来的残差通道计算残差:$R_o = X_o - s_{X_o}\cdot Q(X_o)$,再把$R_o$用同一NVFP4格式量化为$Q_{R_o}$.最后把$Q_x$和$Q_{R_o}$在通道维度拼接，得到增广激活矩阵:$Q_{Xaug}=[Q_x|Q_{R_o}]\in \mathbb{R}^{N\times (K+S)}$.之后权重侧离线做对称处理:把对应 S 条权重复制一份，拼接成 $Q_{Waug}=[Q_W∣Q_{Wo}]$。
最后矩阵乘法格式与NVFP4 Tensor Core所需格式完全相同:$Y = s_{Xaug} \cdot Q_{Xaug}\cdot(s_{Waug}\cdot Q_{Waug})^\top$。全程保持数据格式一致，可直接调用 CUTLASS/cuBLAS，无需改内核循环。

[2026 ICML] VEQ: Modality-Adaptive Quantization for MoE Vision-Language Models

Wed, 26 Aug 2026 08:54:00 +0800

本篇文章主要针对MoE架构的VLMs的PTQ。

目前主流的MoE VLMs PTQ方法往往忽略了两种异质性：

跨模态异质性:视觉token和语言token在统计分布，敏感度和对模型输出的影响上存在显著差异:

a. 敏感度差异，文本token的梯度范数平均是视觉token的22.4倍，表明文本token包含更密集的信息，对量化误差更加敏感。

b. 数量不平衡，视觉Token远多于文本Token，但文本token在推理中起主导作用
专家间异质性:MoE架构中不同专家的重要性呈现高度不均匀分布：

a. 激活稀疏性，少数热专家被频繁激活，而大部分专家很少被使用

b. 功能分化，部分专家专门处理视觉特征，部分专门处理文本语义，还有部分作为跨模态通用处理器。

c. 路由偏差，路由器对少数专家赋予极高置信度，这些专家对最终输出起决定性作用。

基于此论文提出了VEQ-ME,VEQ-MA.

首先介绍VEQ-ME，即模态-专家感知量化。其核心在于将专家重要性引入误差最小化目标中。

我们定义第i个专家的重要性分数为$S_i$,用于平衡衡量该专家在不同模态下的贡献。由于视觉token天然比文本token数量多，因此单纯用原始频次统计会使结果过度偏向视觉主导型专家，因此重要性分数采取加权求和的方式:

$$ S_i = \gamma N_i^{\text{text}} + \beta N_i^{\text{vis}} $$

其中$N_i^{\text{text}},N_i^{\text{vis}}$表示被路由到第i个专家的某种模态token数目。令$T_{text},T_{vis}$分别表示校准集中所有文本 token 和视觉 token 的总数。系数:

$$ \beta = \frac{T_{text}}{T_{vis}} $$

作为数量归一化因子，用于缩小高频视觉激活的影响，使其能够与文本 token 的计数处于可比较的尺度。系数

$$ \gamma = \frac{||\nabla_{text}||}{||\nabla_{vis}||} $$

作为质量敏感性因子，用于反映文本 token 具有更高梯度影响这一事实。

得到每个专家的重要性分数后，后续在计算量化误差时，采用加权形式:

$$ \mathcal{L}_{\text{weighted}} = \sum_{i=1}^M S_i \cdot ||W_iX_i-\hat W_i X_i||_2^2 $$

接下来介绍VEQ-MA，即模态-亲和性感知量化。其核心在于构建增强的Hessian矩阵$\hat H = XCX^\top$

目前具有代表性的PTQ框架通常使用二阶信息(Hessian矩阵)来确定最优量化参数，而出于计算效率的考虑，往往采用如下近似:

$$ H = 2X^\top X $$

这种形式隐含地假设所有输入 token 对重构误差的贡献是相同的，也就是说，它将序列维度上的优化地形视为均匀的。

然而在MoE-VLM中，存在如下问题:

路由多样性，不同token与特定专家之间具有不同程度的亲和性，这由路由器输出的概率决定；
模态敏感性,尽管文本 token 的数量更少，但相比空间冗余较高的视觉 token，它们通常具有更高的梯度密度和信息价值。

如果直接将统一的 Hessian 计算方式应用于 MoE 层，就无法刻画这些细粒度差异，进而可能导致关键语义信息被忽视。

因此VEQ-MA采用了增强的Hessian矩阵:

设$X \in \mathbb{R}^{d\times N}$表示被路由到当前专家的输入 token，其中 N 是这些 token 的数量。我们根据每个 token 的模态相关亲和性，对其贡献进行缩放，从而重构 Hessian 矩阵 $\hat H$：

$$ \hat H = (X\cdot \sqrt{C})(X\cdot \sqrt{C})^\top = XCX^\top $$

其中$C\in \mathbb{R}^{N\times N}$是一个对角矩阵，表示每个 token 的重要性权重。对于第 j 个 token $x_j$，其对应的对角元素 $c_j$ 定义为：

$$ c_j = p_j \cdot \alpha_j $$

其中:

$$ \alpha_j = \begin{cases} \gamma, & x_j \text{ 是文本 token}, \\ 1, & x_j \text{ 是视觉 token}. \end{cases} $$

$p_j$表示$x_j$与专家的亲和性。项$\gamma$表示梯度缩放因子。

整体架构图如下:

[2026 NeurIPS] Breaking Modality Heterogeneity in Low-Bit Quantization for Large Vision-Language Models

Wed, 26 Aug 2026 08:51:00 +0800

作者团队通过可视化文本和视觉Token在不同通道分布上的分布，观察到以下现象:

不同模态的激活分布存在本质差异:视觉激活通常呈现长尾分布，即只有少数通道包含幅值极大的激活，也就是需要谨慎处理的离群值。相比之下，文本激活中的离群值在通道之间分布得更加均匀。
对于两种模态而言，离群通道都只占全部通道的一小部分
更重要的是，文本模态和视觉模态的激活离群值位于不同的通道中。

由观察1，可以指导VLM量化的一个困难在于目前所有通道的激活都通过一个共享的变换矩阵P进行优化。然而，Modality-specific的离群通道会彼此显著干扰，导致该共享的变换矩阵不是最优的。

这里可以简单解释一下，目前基于变换矩阵的PTQ方法普遍通过最小化如下损失获取:

$$ P = \arg \min ||Q(XP)Q(P^{-1}W) - XW||_2^2 $$

而这种方式获取的P，往往会忽视不同模态Token之间的差异，以及不同通道之间的差异，因此论文提到这个共享的变换矩阵不是最优的。

而受第2，3点的启发，文章提出了一种Modality-specific 的离群通道解耦技术（MOCD）。该方法通过从所有通道分离出文本特定和视觉特定的离群通道，来解耦模态特定离群值之间的干扰。这样，通道被划分为了三组: 视觉特定通道，文本特定通道和模态兼容通道

形式化而言:

我们令:$C = \{0,1,\dots,D_{in}\}$,表示所有通道索引。MOCD将输入通道划分为三个互不相交的集合:

$$ C = C_m \cup C_t \cup C_v,\quad C_t\cap C_m = \emptyset,\quad C_v\cap C_m = \emptyset $$

其中，$C_m,C_t,C_v$分别表示模态兼容通道，文本通道和视觉通道。由于离群通道数目远远小于$D_{in}$,他们还满足如下条件:

$$ \abs{C_m}\gg \abs{C_v},\quad \abs{C_m}\gg \abs{C_t} $$

相应地，激活矩阵和权重矩阵被拆分为:

$$ X \to \{X_m,X_t,X_v\},\quad W\to \{W_m,W_t,W_v\} $$

这三组矩阵对:

$$ \{X_m,W_m\},\quad \{X_v,W_v\},\quad \{X_t,W_t\} $$

会通过学习不同的变换进行不同方式的处理。

模态特定通道的选取方式如下:

首先根据幅度值选择视觉特定的离群通道，然后通过一种基于一致性的校准代理指标来识别文本特定的离群通道。令$T_v$和$T_t$分别代表采样得到的视觉Token和文本Token，C代表完整的通道集合。

对于视觉token，我们使用每个通道上的最大绝对激活值作为该通道的分数:

$$ s_v(c) = \max_{i \in T_v}\abs{X_{i,c}} $$

视觉特定的离群通道集合通过选择分数最高的$K_v$个通道得到:

$$ C_v = \text{TopK}_{c\in C}(s_v(c),K_v) $$

记剩下的通道为:$C' = C - C_v$

对于文本Token而言，使用每个通道上的最大激活值作为通道分数并不合理，因为文本Token的离群通道更多与不同 token 之间不稳定的相对响应有关。因此，我们使用每个 token 内部的百分位排名作为一种对尺度不敏感的通道重要性度量：

$$ r_{i,c} = \frac{1}{\abs{C'}}\sum_{j\in C'}\mathbb{I}(\abs{X_{i,j}}\leq \abs{X_{i,c}}),\quad i\in T_t,c\in C' $$

这种基于排名的度量可以抑制 token 之间的尺度变化，使得不同文本 token 之间的通道响应具有可比性。

对于每个通道，我们将其在所有文本 token 上的排名序列聚类为 K 组，并使用簇内方差作为响应不稳定性的度量。令$z_{i,c}$表示$r_{i,c}$的聚类分配，$\mu_{c,z_{i,c}}$表示对应的聚类中心，则:

$$ s_t(c) = \frac{1}{\abs{T_t}}\sum_{i\in T_t}(r_{i,c} - \mu_{c,z_{i,c}})^2 $$

更大的$s_t(c)$表明在该通道上，不同文本 token 之间的相对响应更加不稳定。我们选择$s_t(c)$最高的$K_t$个通道作为文本特定通道:

$$ C_t = \text{TopK}_{c\in C'}(s_t(c),K_t) $$

剩余的通道构成模态兼容的通道:

$$ C_m = C - C_t - C_v $$

接下来视觉和文本特定通道对应的矩阵对$\{W_v,X_v\},\quad \{W_t,X_t\}$会走各自的独立量化路径，而主通道路径则会走自适应跨模态校准(ACC)

我们令$\Delta M= M - Q(M)$,其中M为权重矩阵或激活值矩阵。

对于权重侧的量化误差，采用CWS,通过低秩分解吸收由跨模态偏移触发的权重变化:

$$ P_m^{-1}W_m = \underbrace{(P_m^{-1}W_m)-U_sV_s}_{\text{平滑后主权重}}+\underbrace{U_sV_s}_{\text{低秩自适应分量}} $$

量化形式为:

$$ Y_m = \hat{X}_mQ(P_m^{-1}W_m-U_sV_s)+\hat{X}Q(U_s)Q(V_s) $$

其中$U_s\in\mathbb{R}^{D_m\times r},\quad V_s\in\mathbb{R}^{r\times D_{out}}$为可学习低秩矩阵，用于隔离敏感跨模态成分，使主权重分布更加平滑。

而对于激活侧的激活残差，与静态权重不同，激活是输入相关的，并且会随模态动态变化。因此，很难通过结构分解来重新表示其量化敏感分量。

因此我们引入一个直接补偿分支，用于恢复由激活量化带来的输出偏差:

$$ \Delta Y_{m}^{\text{act}} = \Delta(X_mP_m)P_{m}^{-1}W_m $$

而有研究表明:由于文本激活具有稠密语义特性，因此它们对量化噪声更加敏感；相比之下，视觉激活通常具有较强冗余性。因此，文本补偿通常能够捕获更关键的激活侧误差。

因此引入补偿分支恢复文本token的输出偏差,同时为了高效计算，我们用可学习的低秩矩阵$U_c,V_c$来近似权重映射。

$$ Y_{m}^{\text{text}}\leftarrow Y_{m}^{\text{text}}+Q(\Delta(X_mP_m)^{\text{text}})Q(U_c)Q(V_c) $$

整体的框架如图所示:

激活侧和权重侧的补偿都依赖于低秩分支。然而，像 LoRA 或 QLoRA这样自由学习的参数很容易在小规模校准集上过拟合；而固定的基于 SVD 的分量又缺乏足够的灵活性，难以吸收由模态异构性导致的量化误差。

为了在两种选择之间平衡，避免自由参数过拟合，同时保持对跨模态异质性的适应能力，低秩矩阵采用锚定SVD结构：

$$ W_m \approx U_r\Sigma_r V_r^\top\\ U^* = P_m^{-1}W_m,V^*=\Sigma_rG_*V_r^\top,*\in \{s,c\} $$

其中$G_* \in \mathbb{R}^{r\times r}$为可学习的对角门控矩阵，在保留SVD结构先验的同时实现奇异方向自适应重加权。

可以参考下图:

[2026 ICML] ReSpinQuant: Efficient Layer-Wise LLM Quantization via Subspace Residual Rotation Approximation

Wed, 26 Aug 2026 08:48:00 +0800

在LLM的权重-激活值量化中，目前的主流是基于旋转的方法，总体而言可以分为两类，一种是以SpinQuant，QuaRot为代表的全局旋转方法，另一种是以FlatQuant，OSTQuant为代表的layer-wise 变换方法。

二者的区别在于前者全局共享旋转矩阵，可以实现激活旋转与权重的离线融合，这种方式推理时无额外开销，效率高，但是表达能力有限；而Layer-Wise变换方法为每层分配独特的旋转矩阵，可以通过局部适应实现更优的异常值抑制，但是会产生额外的推理开销，因为这个旋转矩阵在激活侧无法融合进前一层的权重层。

ReSpinQuant克服了这一限制。实现了可融合的Layer-Wise Rotation base PTQ。

具体方法如下:

上图是respinquant应用于标准Transformer层时的完整架构。

我们设L表示总层数，对于第i层:

$R_1^i$:用于旋转MHSA模块的激活输入，以及FFN模块的输出激活（来自于下一层）
$R_2^i$:用于旋转FFN模块的输入，以及MHSA的输出。
$R_3^i$:作用于注意力机制的中间旋转，如Value projection
$R_4,R_5$:通过快速 Hadamard 变换实现的结构化旋转。与SpinQuant中保持一致。

上述旋转矩阵均可被MHSA，FFN内部的线性变换吸收，例如$W_v^i$会被融合为:$\hat W_v^i = {R_1^i}^\top W_v^i R_3^i$

上述吸收可以离线进行，因此基本上在保证计算不变性的同时，几乎没有带来额外的推理开销。

但是上述公式仅在Transfomer Block内部成立，当使用逐层旋转时，残差连接会带来挑战。

我们设$R_{in}$和$R_{out}$分别表示每个MHSA或FFN block的输入和输出所对应的最优逐层旋转矩阵。原始残差连接为:

$$ x_{out} = x_{in} + \text{Block}(x_{in}) $$

在旋转后，我们可以写作:

$$ \hat{x_{out}} = R_{out}R_{in}^{\top}\hat{x_{in}}+R_{out}\text{Block}(R_{in}^{\top}\hat{x_{in}}) $$

如果采用类似于SpinQuant的全局旋转策略，那么有$R_{out} = R_{in}$,此时:

$$ T = R_{out}R_{in}^\top = I $$

因此可以消除残差连接中的额外计算开销（也就是我们不需要显式计算）。然而，这会限制模型的表达能力，因为它强制所有层共享同一个旋转基。

为了解决这个问题，我们采用完整大小的逐层旋转矩阵，以最大化表达能力；同时，通过对子空间旋转近似来逼近残差旋转矩阵 T。

作者团队通过实验发现，用Hadamard矩阵初始化旋转矩阵，并通过Caley Optimizer对其进行优化，学习到的旋转矩阵$(R_1,R_2)$在收敛后并不会显著偏离初始的Hadamard结构。因此，残差旋转矩阵T表现出很强的对角占优特性:

$$ T = R_{out}R_{in}^\top\approx HH^\top = I $$

我们将其相对于单位矩阵I的偏差记为:

$$ \Delta T = T - I $$

随后，对偏差矩阵进行SVD分解,以识别基空间不匹配的主要方向（按照我们对SVD的理解，经过矩阵T的线性变换后，向量所在空间以Q为基底，也就是Q所在的线性空间）:

$$ Q,S,V^\top = \text{SVD}(T-I) $$

我们截断分解，只保留前r个奇异向量，从而构造投影矩阵（即我们认为在空间上只有少数方向上残差不匹配）:

$$ Q \in \mathbb{R}^{D\times r} $$

推理出这个子空间基后，我们便可以在子空间内推导最优旋转矩阵:

$$ \hat{R}_{sub}\in \mathbb{R}^{r\times r} $$

首先，将完整的变换矩阵投影到该子空间中:

$$ T_{\text{sub}} = Q^\top T Q \in \mathbb{R}^{r\times r} $$

由于投影操作不严格保持正交性，因此我们通过极分解提取最接近的正交矩阵。具体而言，我们对投影后的分量进行 SVD：

$$ U_{sub},\Sigma_{sub},V_{sub}^\top = \text{SVD}(T_{sub}) $$

由此得到正交化后的子空间旋转矩阵:

$$ \hat{R}_{\text{sub}} = U_{\text{sub}}V_{\text{sub}}^\top $$

我们通过仅在识别出的子空间内施加变换，同时保持其正交补空间不变，来近似完整旋转矩阵T。近似后的变换矩阵$\hat T$定义为:

$$ \hat T = \underbrace{I-QQ^\top}_{(D-r)维恒等变换}+\underbrace{Q\hat R_{\text{sub}}Q^\top}_{子空间旋转} $$

那么残差流的整体流程如下：

投影: $y = Q^\top \hat x_{in} \in \mathbb{R}^{r}$
子空间变换,在r维子空间内应用可学习的稠密变换，我们定义有效子空间矩阵：

$$ M = \hat{R_{\text{sub}}} - I_r $$

以合并加法操作，因此有：

$$ z = My \in \mathbb{R}^{r} $$

重投影与残差相加:

将原始结果投影回原始维度，并与输入相加:

$$ \hat{x_{\text{out}}} = \hat{x_{in}}+Qz $$

[2026 ICLR] SERQ: Saliency-Aware Low-Rank Error Reconstruction For LLM Quantization

Wed, 26 Aug 2026 08:45:00 +0800

这篇文章将LLM PTQ做到了W4A4，是通过低秩误差重建的方式做到的。

传统的W4A4量化方法一般是通过Rotation-Based的方式进行的，这种方式虽然有效，但是一方面可能不存在鲁棒性，另一方面校准/训练成本高。

相较于Rotation-Base的PTQ方法，基于低秩误差重建的方法有不用重训练太多，也不需要复杂在线层的优势，但仅在W4A8下表现优异，在W4A4下会呈现出明显的性能损失。更加关键的是，传统低秩补偿是两个矩阵$L_1,L_2$,推理时需要进行:

$$ X_qW_q + X_qL_1L_2 $$

这意味着第二项$X_qL_1L_2$我们除了要对$X_q$进行量化以外，我们还需要在线对$X_qL_1$的结果进行在线量化，这带来的额外开销对低精度的kernel而言并不友好。

可以见下图:

此外，传统的低秩误差重构方法通常对整个误差矩阵$E = W - Q(W)$做截断SVD，这带来的问题是固定 rank budget 会被分散到整个矩阵的所有行列上，而造成真正大影响的可能只是少数几个权重，这样会稀释低秩补偿能力。

基于上述问题，论文提出了SERQ方法，一种显著性感知的误差重构算法，它在单个低秩矩阵中同时考虑权重显著性和激活显著性。

具体而言该方法的步骤如下：

静态激活平滑

激活值量化通常因为异常值的存在而十分脆弱，通常的方法是通过旋转变换或者辅助层对异常值进行在线处理，这些方法虽然有效，但是会引入额外的推理时延，因此这里选用的是SmoothQuant的方式，即采用静态的逐通道缩放来平滑激活分布。

具体来说，激活会被一个缩放因子s缩放，同时对应的缩放因子会被折叠进权重中。因此，线性层中的操作可以表示为：

$$ Y = XW = (X\cdot \text{diag}(S^{-1}))(\text{diag}(S)\cdot W) = XW $$

这些缩放因子在校准阶段获得，并在线下合并到相邻层中，因此不会带来运行时开销。

显著性感知的误差重建

逐通道静态平滑过程会把激活离群值的尺度转移到对应的权重中。假设原始权重符合正态分布，那么在折叠后的权重中，显著行可以直接通过它们的尺度来识别。

这些显著行在反复与激活矩阵相乘时，会累积较大的量化误差。为了缓解这一问题，我们引入了一个低秩补偿矩阵:

$$ R \in \mathbb{R}^{r \times d} $$

它用于修正r个显著权重行中的量化误差，记这些显著权重行为$W_s$

考虑将权重行按照显著性降序排列，即通过置换矩阵P进行重排，则折叠后的矩阵W和显著性感知的低秩矩阵R可以定义为:

$$ W = P\cdot \text{diag}(S)\cdot W = P\cdot W = [W_s;W_r]\\ R = W_s - Q(W_s) $$

其中$W_r$表示剩余的非显著行

随后，整体线性操作可以描述为:

$$ Y = (X\cdot \text{diag}(s^{-1})\cdot P^{-1})(P\cdot \text{diag}(s)W) = XW\\ Q(X)\cdot Q(W)=Q([X_s;X_r])\cdot Q([W_s;W_r])+Q(X_s)\cdot R \approx X_q\cdot W_q + X_{s,q}\cdot Q(R) $$

需要注意的是我们也会对低秩矩阵R进行量化，从而保证整个推理流程都可以使用低精度算子。

通过提取出敏感行的方法，我们将原本的在线量化成本省去，并保留了低秩乘法的便捷！此时残差分支只需要执行一个计算量较低的低秩乘法:

$$ \mathbb{R}^{s\times r}\times \mathbb{R}^{r\times d} $$

离线置换权重

我们为了提取敏感行将激活值和权重通过置换矩阵P变换为了:

$$ X = [X_s;X_r]\quad W = [W_s;W_r] $$

为了不引入额外的在线计算开销，会做一个简单的融合。

一共有两部分，一个是权重上的，$P\cdot \text{diag}(s)\cdot W$,那这个很简单，我们一般的做法是会把scale离线融到W中，按照同样的思路我们也把P融合进去。另一部分则是激活值上的$X\cdot \text{diag}(s^{-1})\cdot P^{-1}$因为激活值X在推理时才有，我们无法离线融合，因此选用的方法是不在当前层进行融合，而是把这个重排继续传播到前一层中。因为我们知道当前层l的激活值X:

$$ X = X^{l-1}W^{l-1} $$

因此我们令:

$$ W^{l-1} = W^{l-1}\cdot \text{diag}(s^{-1})\cdot P^{-1} $$

即可，这个操作可以离线执行，因此不会带来额外的计算开销

至此，完整的流程可见下图:

[2026 ICML] OSAQ: Outlier Self-Absorption for Accurate Low-bit LLM Quantization

Wed, 26 Aug 2026 08:42:00 +0800

目前针对大语言模型中存在的系统性异常值问题，现有方法主要依赖层内乘法变换来抑制异常值，包括:

缩放:如AWQ，SmoothQuant等方法通过激活分布特征对权重进行缩放
旋转:如QuIP等方法通过正交矩阵旋转权重矩阵

然而这些方法在极低比特量化时，性能仍远未达到理想水平，表明单一乘法策略在根本上不足以充分处理异常值问题。

本方法基于如下发现:

即任务损失关于权重的Hessian矩阵具有低秩一致性，即

特征值沿特定方向趋于0零
对应的特征向量构成稳定的零空间
该零空间在不同输入样本间保持高度一致

而我们知道，通过泰勒展开，我们可以知道权重收到扰动时，任务损失L关于权重的二阶泰勒展开可以写作:

$$ \mathbb{E}[L(w+\Delta w)-L(w)]\approx\frac{1}{2}\Delta w^\top H_w \Delta w $$

而我们又发现$H_w$具有低秩一致性，因此根据零空间的定义，用$H_w$乘以零空间中的任意向量都会得到零。因此，通过对这些零空间向量进行加权组合，我们可以构造出$\Delta w$。这使得一种加性变换成为可能，并且保证损失保持不变:

$$ W' = W +\Delta W\quad s.t. \quad \Delta w^\top H_w \Delta w = 0 $$

基于这个发现，我们旨在构建一个由低秩结构引导的$\Delta w$，对权重执行加性变换，从而实现异常值的子吸收，同时保持模型的性能。

给定一个权重矩阵$W\in \mathbb{R}^{M\times N}$,其中M表示输出通道维度，N表示输入通道维度，$\Delta W$的构造过程如下所述:

零空间提取: 首先我们对Hessian矩阵$H_w$进行特征分解，并按照特征值幅度的非递减顺序进行排序，如下所示:

$$ H_w = V \text{diag}(\lambda_1,\dots,\lambda_N)V^\top,\quad 0\leq \abs{\lambda_1}\leq \abs{\lambda_2}\leq \dots \leq \abs{\lambda_N} $$

其中$V\in \mathbb{R}^{N\times N}$是特征矩阵,$\lambda_1 ,\dots,\lambda_N$是矩阵的特征值。我们采取尾部能量累积的策略，从最小的特征值开始累加，得到前缀能量，并将零空间维度确定为满足累积尾部能量达到预设阈值时的最小K:

$$ \mathcal{N}=V^\top_{[:,0:K-1]},\text{where}\quad K=\min_k \{\sum_{i=1}^k \abs{\lambda_i} \geq \gamma \sum_{i=1}^N\} $$

其中$\gamma \in (0,1)$是尾部能量阈值，$\mathcal{N}\in \mathbb{R}^{N\times K}$表示矩阵$H_w$的零空间，其中每一行对应一个特征方向，在该方向上$H_w$表现出近似消失的曲率。

$\text{softmax}-\infty$目标近似：在获取了Hessian矩阵的零空间后，我们引入了一个权重系数矩阵$\beta \in \mathbb{R}^{N\times K}$,用于为每个零空间中的每个向量分配权重，从而构造$\Delta w$：

$$ \Delta W = \beta \mathcal{N} $$

我们希望构造出来的$\Delta W$能够最小化施加加性扰动后权重的数值范围，我们可以通过最小化下式达到目标:

$$ \min_{\beta}||W+\Delta W||_{\infty} = \min_{\beta}||W+\beta \mathcal{N}||_{\infty} $$

其中$x=[x_1,\dots,x_n]^\top,\quad ||x||_{\infty} = \max_{1\leq i \leq n}\abs{x_i}$.显然无穷范数不可微，为了解决这个问题，我们采用$\text{softmax}-\infty$近似:

我们沿着输出通道维度应用softmax操作:

$$ s_{ij} = \frac{\exp{(\abs{W_{ij}}/\tau)}}{\sum_{t=1}^N \exp(\abs{W_{it}}/\tau)} $$

其中,$i=1,\dots,M,\quad \tau > 0$是温度系数。当它较大时，它能够捕捉所有分量的平均行为；而当$\tau \to 0^+$时，它会越来越强调极端峰值。

在这种情况下，对这些被“峰值强调”的参数施加$\mathcal{l}_2$范数，便可以作为$l_{\infty}$的一种近似，从而有效地识别并抑制异常值。

$\beta$的显式解: 接下来我们来显式解决上述优化问题。经过$\text{softmax}-\infty$近似后我们可以把优化目标写作如下形式，特别地，由于量化的scale和zero-point都是沿着输出通道维度计算的，因此我们给出每个输出通道对应的$\mathcal{l}_2$范数优化目标:

$$ \min_{b_i}\frac{1}{2}\sum_{j=1}^{N}s_{ij}(W_{ij}+b_i^\top \mathcal{n}_j)^2 + \frac{\mu_1}{2}||b_i||_2 + \frac{\mu_2}{2}(b_i^\top v)^2 $$

其中:

$$ b_i = \beta[i,:] \in \mathbb{R}^{K},\\ n_j = \mathcal{N}[:,j] \in \mathbb{R}^{K},\\ v = \mathcal{N}1_{N} \in \mathbb{R}^{K} $$

上式中第一项是主要的优化目标，作用是最小化施加加性扰动后权重的数值范围；第二项是关于$b_i$的正则化项，防止过大的修正；第三项施加了一个反平移约束，用于惩罚整个通道沿同一方向发生一致平移。

Remark：这个第三项约束是为了避免第 i 个输出通道的整行权重整体发生同方向平移。它希望$\sum_{j=1}^N \Delta w_{ij} \approx 0$

求解上述最优化方程（对$b_i$求导，并令一阶最优性条件为零），可以得到:

$$ A_ib_i = -\rho_i $$

其中

$$ A_i^* = \sum_{j=1}^N s_{ij}n_j n_j^\top + \mu_1 I_K +\mu_2 v v^\top,\quad \rho_i = \sum_{j=1}^N s_{ij}W_{ij}n_j $$

因此，我们可以得到最优的系数矩阵$\beta$：

$$ \beta^* = [b_1^*,\dots,b_M^*]^\top,\quad b_i = -A_i^{-1}\rho_i,i = 1,\dots,M $$

[2026 ICLR] SliderQuant: Accurate Post-Training Quantization for LLMs

Wed, 26 Aug 2026 08:39:00 +0800

现有的PTQ通常采用顺序量化框架，将预训练模型分割为相同大小的部分并依次量化，且对所有层一视同仁。在低比特情况下，这种平等处理方式存在如下缺陷：

层间敏感度差异被忽略:实证研究表明，浅层和深层通常比中间层对量化更敏感，且第一层和最后一层的量化误差显著大于其他层
量化误差跨层累积:随着逐层量化进行，误差逐渐放大，而现有方法缺乏有效的跨层协同机制来抑制此问题

具体而言，作者做了两个简单的实验

将某个模型的特定层量化，来比较量化不同层带来的影响，以及量化不同数量的层数，来比较量化误差的积累效应，结果如下：

可以看到无论模型大小，对量化最敏感的永远是第一层和最后一层，且随着量化层数的增加，量化误差的积累也越来越明显。

为了解决上述问题，论文提出了SliderQuant框架，其核心包括:

Inter-layer sliding quantization: 针对浅层，中间层和深层分别设计自适应滑动窗口，建立跨层的智能化接力机制
Intra-layer sliding quantization: 在每个量化窗口内采用增量式量化策略，实现局部到全局的参数协同。

Inter-layer sliding quantization

一般的PTQ量化普遍选择Layer-Wise的量化，因此误差会随着层数加深而逐渐累积。作者选择使用滑动窗口的方式进行量化，每次在整个窗口中的layer会被当作一个block进行量化。为了减少跨层量化误差，总会保证两个连续的窗口之间存在重叠即$s-i\geq 1$。（size，stride）

然而使用固定大小的滑动窗口时，预训练大语言模型的所有层都会以相同的窗口大小和每一步相同的移动间隔进行量化。也就是说，浅层、中间层和深层在很大程度上仍然被同等对待，这与我们期望的量化设计之间仍存在较大差距。

因此SliderQuant选用的方式为:

对于$L_s$个浅层，采用渐进扩展滑动窗口。具体而言，从仅量化第一层开始，窗口大小被设置为1；随后以第一层作为锚定层，每一步将窗口大小增加1，直到窗口覆盖所有浅层$L_s$

对于$L_D$个深层，采用渐进收缩滑动窗口。具体而言，从量化所有深层开始，然后每一步将窗口大小减少 1，直到窗口中只包含最后一层。最后一层始终作为锚定层，并参与每一个收缩后的滑动窗口的量化过程。

对于中间$L_i$个中间层，采用固定大小的滑动窗口，其中设置$\{s=2,i=1\}$,并保证各层具有均匀的优化频率。具体而言，在浅层和中间层之间设置一个重叠层，同时在中间层和深层之间也设置一个重叠层。

根据消融实验得到$L_s = L_D =4$可以在效率和精度之间达到平衡

可以参考下面的gif

Intra-Layer Sliding Quantization

为了进一步利用之前发现两个性质，论文提出了一个与Inter-layer sliding 互补的组件。

具体而言，它将逐步扩展的滑动设计进一步扩展到层间滑动量化的每一个窗口内部。在窗口内，s个层会沿着权重和激活维度，以比例$\gamma$并行地执行逐步扩展滑动。因此，所有 s个层的联合量化会在$N = \frac{1}{\gamma}$个滑动阶段完成。

这里从语言上描述比较抽象，可以参考下面的GIF，简单来说就是对于权重/激活值矩阵，沿着某个维度逐渐进行量化（而非一次性量化）。

层内滑动量化在层间滑动量化当前滑动窗口内部，建立了一种从局部到全局的跨层参数协同关系，从而降低量化误差。

在层内对激活值和权重进行量化时，采用如下策略，该策略参考了目前主流的通道缩放以及低秩近似：

设$W_i \in \mathbb{R}^{n\times m}$表示滑动窗口中第i层的权重矩阵，$X_i \in \mathbb{R}^{k\times n}$表示其对应于一小组校准样本的输入特征，其中校准样本数为c，默认设置c=128,那么，量化过程定义为:

$$ \hat{X_i} = X_i \oslash \alpha_i\\ \hat{W_i} = W_i \odot \alpha_i + A_iB_i \\ \hat{X_{i+1}} = \text{quantizer}(\hat X_i)\cdot \text{quantizer}(\hat{W_i}) $$

其中$\alpha_i$表示一个可学习的通道缩放参数，$A_i,B_i$为两个低秩矩阵。

[2026 ICLR] Task-related Token Compression in Multi-modal Large Language Models from an Explainability Perspective

Wed, 26 Aug 2026 08:36:00 +0800

个人评价:很有意思的一篇文章，出发点是发现了一种较好的可解释性的剪枝方法，但是剪枝决策需要在推理完成后得到，因此通过加入可学习模块的方式进行改良。

现有的MLLMs通常将视觉token和文本token一起输入到LLM中进行跨模态对齐和整合。然而，这种方法由于视觉token数量庞大（尤其是处理高分辨率图像或高帧率视频时），导致了巨大的内存和计算开销。因此，迫切需要有效的token压缩技术来提高模型的效率。

基于此，论文作者提出了一种可解释性Token剪枝方法:

我们假设MLLMs一共有L层，并将生成的文本token序列记为:

$$ Y=\{y_0,y_1,\dots,y_{T-1}\} $$

具体而言，我们从最终生成的token来回溯原始视觉输入的贡献。对于第t个生成步骤中的每个$y_t$，首先将相关性图$R_t$初始化为单位矩阵，然后在各层之间迭代更新。

记$A_t^l,\nabla A_t^l$分别为第l层中的Multi-Head Attention Map以及对应的梯度，它们分别在前向反向传播中获取。那么$R_t$的迭代方式如下:

$$ R_t = R_t + E_h((A_t^l\odot \nabla A_t^l)^+)\cdot R_t \tag{1} $$

其中，$\odot$ 表示Hadamard积，$E_h$表示沿注意力头维度取平均。该更新从第0层一直进行到最后一层。

最终，可以通过索引$R_t$最后一行中相应位置来提取$y_t$与视觉信号之间的相关性，即:

$$ R_t[-1,N_s:N_s+N_v] $$

最后，我们对所有时间步 t的视觉相关性取平均，从而得到相对于当前响应的整体视觉相关性分数：

$$ R_v\in \mathbb{R}^{1\times N_v} $$

接下来可以根据

下面对式(1)进行简单的理论解释:

该式来源于Generic Attention Explainability, GAE 框架。这是一个用于解释Transformer架构预测结果的强大方法。

GAE之所以选择将Attention Map和其梯度的Hadamard积是因为:

Attention Map可以反映每个token从其他token接受了多少注意力
梯度可以反映哪些 token 需要获得更多注意力，才能有效影响当前输出

从数学上我们有:

对某一层、第 h 个注意力头来说，注意力矩阵可以写成：

$$ A_h^l \in \mathbb{R}^{N\times N} $$

我们把输出分数$s_t$看作是注意力矩阵的函数:

$$ s_t = f(A_h^l) $$

假设我们对某个注意力权重$A_{h,ij}^l$做一个小扰动，根据一阶泰勒展开我们有:

$$ \Delta s_t \approx \frac{\partial s_t}{\partial A_{h,ij}^l}\Delta A_{h,ij}^l $$

而在剪枝场景中，扰动量就是其本身，那么我们就可以把贡献写作:

$$ A_{h,ij}^l\odot \nabla A_{h,ij}^l $$

这就是式子(1)中采用Hadamard积的原因。

实验表明，使用这种方法作为剪枝依据，可以在仅保留50%视觉Token的情况下，保留99%的性能

然而在实际应用中却存在一个局限，$R_v$是输出已经生成后得到的，这与我们进行剪枝的初衷相违背。为了解决这一限制，作者提出了一个独立于MLLM训练的单独模块来近似$R_v$。

整体而言，模型架构如下：

在式子(1)中，我们的相关性图是通过聚合Attention Map得到的，这表明从Attention Map到相关性图的映射是有前景的。

作者通过实验发现仅对第一层注意力应用一个简单的卷积网络就足够了。形式上，令$A^0$表示第一层Attention Map，因为我们是对视觉Token进行剪枝，我们更加关心文本token对于视觉Token的注意力，因此，我们取其子图:$A_{u\to v}^0\in \mathbb{R}^{N_{u}\times N_v}$

随后，我们对每个视觉token的$N_v$个分数取平均，得到一个紧凑表示$A_v^0 \in \mathbb{R}^{1\times N_v}$

该平均注意力向量$A_v^0$随后被输入到一个一维卷积模型$f_{\theta}$中，用于预测视觉相关性:

$$ \hat R_{v} = f_{\theta}(A_v^0) $$

训练时，我们将真实计算出来的$R_v$处理成为$R_v^*$来作为GT：首先按照上面介绍的方法，屏蔽掉最低的50%的数值，然后将剩余部分归一化为概率分布。作者为了避免原始分数接近，softmax 产生近似均匀的数值，而采用将每个分数除以总和来进行归一化。

最后给定$R_v^*,\hat R_v$，通过KL散度来计算Loss。

Towards Joint Quantization and Token Pruning of Vision-Language Models

Wed, 26 Aug 2026 08:33:00 +0800

个人评价:其实和QAPruner大同小异，就多了一个层预算分配的问题，以及打分的依据不同

现有的两阶段量化剪枝方法，如先剪枝后量化或先量化后剪枝往往会低比特校准流形与剪枝执行形式之间引入不匹配，具体到实验中就是这种方法的效果不佳。

一种可能的解释是量化的噪声会干扰Token重要性信号的估计，而剪枝则会改变激活的统计特性，可能使低比特算子依赖的校准假设失效。且目前的剪枝方法没有充分考虑量化对重要性信号评估的可靠性的影响

对应上述问题，论文提出了一个协作式量化与剪枝框架，通过QUOTA机制将低比特校准敏感度转换为层级的Token分配计划，并在统一的低比特推理流程（包括量化KV缓存）中执行确定性剪枝，从而确保剪枝决策与部署时的量化操作制度保持一致。

首先需要解决的是剪枝候选层$L_c$的确定问题。论文中选择在校准集上对逐层注意力集中度和视觉Token冗余度进行分析。其中逐层注意力集中度从多模态注意力图中模态间注意力块测得，具体使用从文本query到视觉token的top-10注意力分数的中位数，并结合样本间的四分位距进行衡量。结果表明前两层的集中度较低，随后急剧上升，这表明从该深度开始，基于注意力的重要性排序变得更加可靠。与此同时，基于视觉 token 两两余弦相似度中位数的冗余度代理指标在早期模块中仍然较低，并且显著低于最后几个模块，这说明视觉 token 在该阶段仍具有较高多样性。基于这些趋势，我们排除前两层，选择一个连续的早期层范围作为 LcLc，并避免在最后几个模块中进行剪枝，因为此时累积节省较小，且 token 移除更加脆弱。所得 $L_c$ 将用于后续的预算分配和 token 选择步骤。

给定$L_c$后，QUOTA在低比特校准过程中通过分析校准集上的量化敏感度来推导逐层token预算。对于每个$l\in L_c$，我们将敏感度定义为全精度激活$x_l$与在部署的低比特算子下计算得到的低比特激活值$x_l^q$之间的相对偏差

$$ S_l = \text{median}_{x\sim D_{cal}}\frac{||x_l^q-x_l||^2}{||x_l||^2+\epsilon} $$

较大的$S_l$说明在低比特下该层更加敏感，因此会分配更大的token预算。

原始的敏感度在不同层之间可能呈现出重尾分布，因此采用基于百分位数的裁剪和归一化。

$$ \hat{S_l} =\text{clip}(\frac{S_l-P_{10}}{P_{90}-P_{10}},0.1,0.9) $$

接下来我们将$\hat S_{l}$映射为逐层保留比例调度。通过一个温度控制的softmax实现:

$$ \pi_i = \frac{\exp(\frac{1-\hat S_{l_i}}{\tau})}{\sum_{j=1}^m \exp(\frac{1-\hat S_{l_j}}{\tau})} $$

为了确保鲁棒性和单调调度，我们设置保留比例下限 $p_{\min}$，并令总丢弃预算 $B = 1-p_{\min}$。我们分配 $d_i=Bπ_i$，并形成一个非递增的保留比例表：

$$ r_{l_i} = \max(p_{\min},1-\sum_{j=1}^{i} d_j) $$

接下来问题就转变为了给定逐层保留比例$\{r_l\}_{l\in L_c}$,我们在每个候选层执行带预算约束的token选择。我们遵循量化的一致性原则:

所有重要性信号都在实际部署的低比特算子下计算。随后，我们可以得到一个确定性的综合评分，并应用$Top-K_l$选择：

在候选层$l\in L_c$处，令:

$$ V_l = \{v_i^l\}_{i=1}^{N_l} $$

和$T_l$分别表示视觉Token表征和文本Token表征。令$V_0$表示经过projector后的参考视觉Token长度，该长度在校准阶段测量，并存储在剪枝策略中。逐层预算定义为:

$$ K_l = [r_lV_0] $$

我们保留按照下述评分排序后排名前 $K_l$的视觉 token。对于每个视觉 token $V_{l_i}$，我们从实际部署的量化前向传播中计算四种重要性指标。我们约定注意力权重 $A_{qk}$按查询 token q 和键 token k索引。具体而言，

$$ m_{i,l}^{\text{mag}}=||v_i^l||_2,\quad m_{i,l}^{\text{inter}}=\frac{1}{H}\sum_{h=1}^H\sum_{j\in T}A_{ji}^{\text{inter},h}(l),\\ m_{i,l}^{res}=||Q(v_i^l)-v_i^l||_2,\quad m_{i,l}^{\text{intra}}=\frac{1}{H}\sum_{h=1}^H\sum_{k\in V}A_{ki}^{\text{intra},h}(l) $$

由于这些指标具有不同尺度，并且可能受到离群值影响，我们采用基于百分位裁剪与重缩放的逐层鲁棒归一化算子 GG。对于每一种指标类型：

$$ m \in \{\text{mag,inter,res,intra}\}\\ \hat{m_{i,l}}^{(m)}=G(m_{i,l}^{(m)}) $$

其中，GG 在每一层 ℓℓ 上独立应用，将数值映射到可比较的范围，同时降低极端 token 的影响。随后，我们使用层间共享的加权和来构造综合重要性评分：

$$ \text{score}_{i,l} = \sum_{m\in \{\text{mag,inter,res,intra}\}}w_m \hat m_{i,l}^{(m)} $$

然后选出Top—K。

QAPruner: Quantization-Aware Vision Token Pruning for Multimodal Large Language Models

Wed, 26 Aug 2026 08:30:00 +0800

个人评价:方法特别简单，就是在剪枝的时候考虑量化的影响

当我们简单地将基于语义的token剪枝应用于经PTQ优化的模型时，会丢弃对数值稳定性至关重要的激活异常值，从而在低比特位制度（如W4A4）下显著加剧量化误差。

论文提出的解决思路是提出一种量化感知的视觉Token剪枝方法QAPruner，具体而言，方法如下：

我们考虑构建一个联合考虑语义相关性和量化鲁棒性的Token选择机制。

对于每个视觉token，$v_i \in \mathbb{R}^{D}$，我们通过融合两个互相正交的指标来计算其敏感度:分组量化模拟和全局异常值强度。

在目前主流的PTQ量化方法中，通常会将激活值重新分组为若干个更小的组，以便计算局部缩放因子，从而减轻通道异常值的影响。为了模拟这个过程，我们将token特征重排为$M = D / G$个组，其中第i个token的第m组记为$v_{i,m}\in\mathbb{R}^{G}$。假设采用INT4量化，则局部缩放因子$s_{i,m}$及其量化后的表示$\hat{v_{i,m}}$可以写作：

$$ s_{i,m} = \frac{\max(|v_{i,m}|)}{7}\\ \hat{v_{i,m}} = \text{Round}(\frac{v_{i,m}}{s_{i,m}+\epsilon})\cdot s_{i,m} $$

随后将所有的$v_{i,m}$进行拼接得到$\hat{v_i}\in \mathbb{R}^{D}$。那么第i个Token的分组量化误差记为:$E_i$:

$$ E_i = ||v_i - \hat{v_i}||_2^2 $$

具有较高$E_i $的 token 在局部层面上本质上更难量化，并会遭受显著的信息损失，因此是应当优先保留的关键候选。

尽管上述指标$E_i$可以捕获局部量化困难，但是它可能无法显式惩罚那些包含极端全局异常值的token被移除的情况。这类携带异常值的token决定了整个张量的最大激活范围，对于保持大语言模型的涌现特性至关重要。

为了显式保护这些结构性异常值，我们将第 i 个 token 的全局异常值强度 $R_i$定义为其在全部 D 个通道上的激活值跨度：

$$ R_i = \max_{j\in\{1,\dots,D\}}(v_{i,j})-\min_{j\in\{1,\dots,D\}}(v_{i,j}) $$

较大的 $R_i$ 表明该 token 中存在严重的激活异常值，因此一旦被丢弃，就会对量化后的数值分布造成更大的扰动。

为了构建一个能够兼顾局部细节保留和全局异常值保护的综合度量，我们首先在一个 batch 内，对 N个视觉 token 的这两个指标分别独立归一化到 [0,1] 区间。最终的量化敏感度分数 $S_i^Q$定义为两项归一化指标的等权和：

$$ S_i^Q = \frac{1}{2}\cdot \frac{E_i - \min(E)}{\max(E)-\min(E)} + \frac{1}{2}\cdot \frac{R_i -\min(R)}{\max(R)-\min(R)} $$

最后，我们将这一量化敏感度与传统的视觉token剪枝方法得到的分数$S_i^P$结合起来，共同指导token的选择过程。为此，我们引入了超参数$\alpha \in [0,1]$,用于控制语义对齐和数值稳定性之间的权衡：

$$ S_i^{Final} = \alpha S_i^P + (1-\alpha)S_i^Q $$

通过这种方法重新校准了 token 选择准则，使得剪枝后的视觉序列不仅在语义上对查询保持足够的信息性，同时也能更好地抵抗低比特 PTQ 所带来的性能退化。

VLMQ: Efficient Post-Training Quantization for Large Vision-Language Models via Hessian Augmentation

Wed, 26 Aug 2026 08:27:00 +0800

个人评价:文章引出给不同token分配重要性的方式特别好，值得学习！后面的方法就和MBQ，QIG大同小异了

文章提出在VLM中，视觉被过度表征是一个被广泛接受的观点，然而目前基于Hessian的PTQ方法没有利用这种冗余性，而我们知道在GPTQ中近似Hessian是通过激活值的内积近似得到的，这就会导致GPTQ等基于Hessian的量化方法直接应用于VLM上时，构建的Hessian矩阵大部分贡献来自于视觉冗余Token，从而表现不佳。因此，给各个Token分配不同的重要性是至关重要的。

那么视觉token是如何影响Hessian的估计和模型精度的呢？

论文作者发现：

在对VLM进行量化时，视觉Token的纳入是必要的
过量的视觉Token可能会导致量化性能下降，为冗余token赋予较低的重要性，可以缓解由此带来的性能衰退

作者对这两个发现给出了一定的解释：

性能的波动实际上可以归因于Hessian的特征偏移，如下图所示

可以看到，当仅有文本输入时，Hessian矩阵的主成分空间中的分布较为紧凑，而加入了视觉Token后，校准分布变得更加多样化，这有助于缓解量化与推理之间的差距。然而，由于视觉过度表征问题，过量引入视觉 token 会带来 Hessian 向冗余视觉特征偏置的风险（即周围稀疏的点)。

为了解决这个问题，作者提出了一种面向VLM的精确PTQ方法，该方法主要由两方面组成:

重要性感知的量化目标
建立了分块损失扰动与逐层输出误差之间的理论联系，从而能够仅通过一次分块反向传播，高效计算由梯度驱动的重要性分数。

重要性感知的量化目标实际上就是我们先前阅读过的MBQ，QIG等方法使用的量化目标，即Token-Wise的加权。VLMQ的Token重要性通过矩阵形式给出:

令$G\in \mathbb{R}^{N\times N}$,其中G为对角矩阵，第i个对角元素表示分配给输出token$Z_{:,i}$的重要性。

将其纳入目标函数后，我们得到了如下改进形式的目标函数：

$$ \arg \min_{\hat w} ||(\Delta w X-r)\sqrt{G}||_2^2 \quad s.t. \Delta w e_q^\top +w_q - \hat w_q = 0 $$

利用拉格朗日算子法，我们有：

$$ L = ||(\Delta w X -r)\sqrt{G}||_2^2 +\lambda(\Delta w e_q^\top + w_q - \hat w_q) $$

解得:

$$ \Delta w = \frac{\hat w_q - w_q}{\hat H_{qq}^{-1}}\cdot \hat H_{q,:}^{-1} + \hat r \hat X^\top \hat H_{-q,:}^{-1} $$

其中$\hat H = XGX^\top ,\quad \hat r = r\sqrt{G} ,\quad \hat X = X\sqrt{G}$

该形式与原始的GPTAQ方式对齐，因此可以复用其中的效率技巧如Cholesky等。

那么重要性是如何得到的呢？

参见MBQ中的推导，可以得知，最终使用的是:

$$ G = \text{diag}([\overline{|P|}_0,\overline{|P|}_1,\dots,\overline{|P|}_{N-1}]) $$

其中第 n个 token 的重要性定义为梯度中第 n 列的 ℓ1 范数

$$ |P|_n = \sum_{i=0}^{C_0 - 1}|P|_{i,n} $$

[2026 CVPR] Fine-Grained Post-Training Quantization for Large Vision Language Models with Quantization-Aware Integrated Gradients

Wed, 26 Aug 2026 08:24:00 +0800

个人评价:比较有意思的研究思路，从MBQ的Modality-Specific出发，通过实验发现，相较于Modality-Specific，更加细粒度的Token-Wise进行区分效果会更好。基于此研究了多种Token敏感度估计方法，最终采用基于公理化归因的积分梯度方法进行规约，取得不错的效果。但是问题在于文章对理论的分析严重不足！

假设你阅读过MBQ，它是按照模态去对优化目标进行加权的，那么我们可以仔细想想，不同模态最终都会以Token的形式输入模型，既然不同模态之间存在对量化噪声敏感性差异，那么我们其实可以说本质上是Token内部就存在差异，这个差异不仅仅存在于不同模态之间，还可能存在同一个模态之中。那么也就是说，我们完全可以仿照MBQ的思路去做更加细粒度的加权。

一般而言，衡量这种Token之间的差异，可以通过敏感性估计进行。作者在文章中尝试了三种敏感性估计方式:

基于梯度:和MBQ一致，依据Token关于量化损失的梯度
基于注意力:用Attention Score
基于扰动:人为扰动token，然后观察block输出变化有多大。

注:实验方法大概是像MBQ一样对不同的Token对量化损失进行加权，然后在VizWiz数据集上进行测试。

最后结果表明，Token-Level的扰动法的效果在不同敏感度估计下表现最优(0.36%的微弱优势)

基于上述分析，我们知道，按Token的细粒度量化方法可能会有更好的效果。因此作者声称基于公理化归因的启发。下面简单介绍一下公理化归因，这个方法来源于可解释AI。

我们从经典的积分梯度（IG）出发。IG用来衡量从参考输入x’到真实输入x的真实路径上，每个Token的累积贡献，其中$f(\cdot,\cdot)$表示该Block的输出：

$$ \text{IG}(x) = (x - x')\int_{0}^{1}\frac{\partial f(x^\alpha ,w)}{\partial x^{\alpha}}d\alpha \tag{QIG 1} $$

其中$x^\alpha = \alpha(x-x')$,而$f(\cdot,w)$表示全精度模型。

我简单介绍一下这个是怎么来的吧，本质上我们是想知道某个输入的Token发生变化后会对模型的输出产生怎样的变化。由导数的定义我们知道:

$$ f(x)-f(x') = \int_{x'}^x\frac{\partial f(t)}{\partial t}dt $$

在一维的情况下，因为只有一个变量，因此我们上式就是该Token的归约。当我们将输入扩展到多维，我们自然想知道每个维度对变化的贡献。

IG的做法是，我们从参考输入x‘出发，沿一条直线走到真实输入x，可以将这条路径写作:

$$ x^\alpha = x' + \alpha(x - x'),\quad \alpha \in [0,1] $$

那么此时，我们可以把函数写作按照路径变化的形式:

$$ F(\alpha) = f(x^\alpha) $$

这是一个一维函数，自变量只有$\alpha$，于是输入变化带来的变化可以写作:

$$ f(x)-f(x') = F(1)-F(0) $$

写作积分形式:

$$ F(1)-F(0) = \int_0^1 \frac{dF(\alpha)}{d\alpha}d\alpha = \int_{0}^1 \frac{\partial f(x^\alpha)}{\partial x^\alpha}\cdot \frac{\partial x^\alpha}{\partial \alpha} $$

而:

$$ x_i^\alpha = x_i' + \alpha(x_i-x_i') $$

故:

$$ \frac{\partial x_i^\alpha}{\partial \alpha} = x_i-x_i' $$

代入有：

$$ \frac{dF(\alpha)}{d\alpha} = \sum_{i=1}^n \frac{\partial f(x^\alpha)}{\partial x^\alpha}(x_i-x_i') $$

因此:

$$ f(x)-f(x') = \int_{0}^1 \sum_{i=1}^n\frac{\partial f(x^\alpha)}{\partial x^\alpha}(x_i-x_i') d\alpha = \sum_{i=1}^n (x_i-x_i')\int_{0}^1 \frac{\partial f(x^\alpha)}{\partial x^\alpha}d\alpha $$

那么IG就定义第i维的规约为:

$$ \text{IG}_i(x) = (x_i-x_i')\int_{0}^1 \frac{\partial f(x^\alpha)}{\partial x^\alpha}d\alpha $$

写作向量形式就是式子QIG(1)中的形式。

对应到量化场景，我们取原始输入$x'$为量化后的输入$x'=x^q$,那么可以写出量化感知积分梯度:

$$ \text{QIG} = (x-x_q)\int_{0}^1 \frac{\partial f(x^\alpha)}{\partial x^\alpha}d\alpha $$

不过我们不能直接将QIG作为优化的权，因为它呈现了重尾分布，这会导致极少部分token主导优化过程。为了抑制这种现象，作者选择按照四分位距(IQR)进行裁剪，从而得到裁剪后的分数:

$$ C(QIG_i) = clip(QIG_i,Q_1-1.5\cdot IQR,Q_3+1.5\cdot IQR) $$

其中$Q_1,Q_3$分别表示第一和第三四分位数，且$IQR=Q_3-Q_1$.随后对这些分数进行归一化，得到最终的token重要系数:

$$ \lambda_i = \frac{C(QIG_i)}{\sum C(QIG_i)} $$

之后按照类似于MBQ的思路，将这个加权融入量化优化的目标函数即可。

[2025 CVPR] MBQ: Modality-Balanced Quantization for Large Vision-Language Models

Wed, 26 Aug 2026 08:21:00 +0800

个人评价:文章的Insight很不错，不同模态Token的影响差异确实很显著，Method is Simple but benefit a lot

当我们直接使用大语言模型的量化的方法对多模态大模型进行量化时(如AWQ，GPTQ等)，往往会忽略不同模态激活值带来的差异，在论文中这种差异被描述为模态之间敏感度的差异。MBQ的思路是在量化过程中平衡这些差异，从而提高VLMs的准确性。

MBQ首先通过实验发现将\，作者认为这源于对不同模态一视同仁的处理方式。原因主要有两点:

从数据角度来看，视觉数据具有较高的冗余性，因此对小扰动具有强抗干扰性。
从模型角度来看，目前VLM生成的内容主要受预训练LLM的影响，而非输入的图像本身。

作者做了一个小实验来验证上述猜想，他们将图像-文本对作为VLM的输入，并计算并计算监督微调损失函数相对于语言token和视觉token的梯度。这些梯度反映了当对语言（文本）或视觉（图像）token特征施加微小扰动时，对输出语言token（caption）的影响。

如下图所示，可以发现,语言Token的平均绝对值比视觉的大了一个数量级。这也就意味着，在相同的扰动下，视觉token对SFT损失的影响仅为语言token的0.1倍，因此我们不能发把语言Token和视觉Token同等对待

为了展示在校准过程考虑模态差异的重要性，作者进行了一个简单的小实验:在CWE校准中，对视觉Token的重建损失施加一个0.1的模态平衡因子。此时优化目标可以写作:

$$ E^* = \arg \min_{E}[||Q(W\cdot E)(E^{-1}X_l)-W\cdot X_l||^2 + 0.1 *||Q(W\cdot E)Q(E^{-1}X_v)-W\cdot x_v||^2] $$

实验结果表明: 即使仅使用一个启发式选择的模态平衡因子，balanced CWE 也能够显著超过原始 CWE 的性能。

为了进一步探索这个最优的平滑因子，论文提出了MBQ方法。

具体而言，该方法通过最小化SFT损失函数的变化，为每一层分配最优的模态平衡因子。具体而言，我们用下式描述每个线性层的输出激活Y收到一个小扰动$\Delta $时，SFT损失L的变化:

$$ \mathcal{L}(Y+\Delta W) \simeq \mathcal{L} + g^\top \cdot \Delta $$

其中$g^\top$表示输出激活Y的梯度。那么由量化引起的SFT损失可以表示为：

$$ \begin{align} ||\mathcal{L}(\hat Y)-\mathcal{L}(Y)||\simeq ||g^\top \cdot \Delta||\\ =||g_v^\top \cdot \Delta_v + g_l^\top \cdot \Delta_l||\\ \leq ||g_v^\top \cdot \Delta_v|| + ||g_l^\top \cdot \Delta_l||\\ \leq |g_v^\top|\cdot |\Delta_v| + |g_l^\top|\cdot |\Delta_l|\\ =\overline{|g_v|}\cdot ||\hat Y_v - Y_v|| + \overline{|g_l|} \cdot ||\hat Y_l - Y_l|| \end{align} $$

在一般的大语言模型的量化中，通常会分为两个阶段采用不同粒度的量化:

Prefill阶段：这个阶段主要是将整个Prompt并行地计算每一层的hidden state（即计算KV，并得到KV-Cache），这个阶段整个Prompt一次性并行计算，矩阵乘法很大，算力利用率很高。这时候如果把权重和激活都量化，就能直接减少 GEMM 的计算/带宽开销，所以 W8A8 / FP8 W8A8 往往比较合适。
Decode阶段：decode阶段是逐token生成，其瓶颈不在于计算，而是权重的访存，因此通常只会对权重进行量化

在MBQ中遵循了同样的方式，在Prefill阶段进行权重-激活值的量化，在Decode阶段进行权重的量化。二者的优化目标均为:

$$ \min_E{\mathbb{E}}[\overline{|g_v|}\cdot ||WX_v-Q(W\cdot E)Q(E^{-1}\cdot X_v)|| + \overline{|g_l|}\cdot ||WX_l-Q(W\cdot E)Q(E^{-1}X_l)||] $$

需要注意的是，这里的重建损失函数是基于MAE而非MSE的。

[2026 CVPR] MASQuant: Modality-Aware Smoothing Quantization for Multimodal Large Language Models

Wed, 26 Aug 2026 08:18:00 +0800

个人评价:很有价值的工作，从High-Level的角度来看，不同模态激活值的分布不同带来的量化挑战可以看作LLM Quant中激活值Outlier带来的挑战，因此从这个角度出发可以很好的理解文章的出发点。

总结：本文旨在解决基于通道级平滑的PTQ方法应用于多模态大模型时面临的一个核心挑战：Smoothing Misalignment。论文通过MAS为每个模态确定一个平滑因子来解决这个问题，并通过CMC方法来解决与之伴随而来的Cross-Modal Computation Invariance问题。

当将基于通道级（Per-channel）平滑的PTQ应用于多模态大模型（MLLMs）时会面临两个核心挑战：

Smoothing Misalignment（平滑错位）：不同模态的激活值幅度存在数量级的差异，例如视觉Token的激活范围通常比文本和音频大10-100倍。传统的Per-channel量化为每个通道计算单一的缩放因子，导致主导模态的较大激活决定平滑因子，而使非主导模态的激活被过度平滑，信号被严重压制，最终导致量化后的模型性能不佳。
Cross-Modal Computational Invariance（跨模态计算不变性）：直接为不同模态计算独立的平滑因子会破坏计算不变性（坐标系不同）。若严格保持模态特定的平滑，推理时需要为不同模态存储不同的量化权重矩阵，这违背了量化技术通过单一低精度权重表示来减少内存占用的根本目标。

论文提出了MASQuant框架来解决上述两个问题，该框架包含两个核心组件：MAS（Modality-Aware Smoothing）以及CMC（Cross-Modal Compensation ）

对于Smoothing Misalignment的问题，其核心还是在于不同模态的激活值幅度存在数量级的差异，因此MASQuant通过为每种模态维护模态特定的平滑因子来解决这个问题，从而从根本上解决了某一模态的主导效应。

(n这里有一点与SmoothQuant不同，MAS的平滑因子是通过学习得到的：

首先获得模态感知的平滑因子初始值如下

$$ S_m= \text{diag}(s_m),\quad s_{m,i} = \frac{\max_t |x_{t,i}^m|}{\max_j |w_{j,i}|},\quad m\in M $$

随后，我们在模态特定的数据上最小化MAE损失来优化$S_m$。我们记$\{S_m\}_{m\in M}$为$\{S_m\}$,我们有:

$$ \{S_m^*\} = \arg \min_{\{S_m\}_{m\in M}}(\lambda_m \cdot \mathcal{L}_{\text{MAE}}(S_m,X_m,W)) $$

其中$\lambda_m$表示模态m的损失权重，对于模态m，量化重建的MAE损失为:

$$ \mathcal{L}_{\text{MAE}}=||Q(X_m S_m^{-1})Q(S_mW)-X_mW|| $$

这保证了$S_m^*$能捕获模态特定的统计特性，同时避免跨模态干扰。

此外，论文中还给出了通过信噪比量化的收益，可以证明相较于之前的统一平滑(Unified Smoothing)，MAS使用的最优平滑(Optimal Smoothing)二者差值：

$$ \Delta =10\log_{10}(\frac{\sum_{i=1}^d \frac{1}{\alpha_i^2}}{d\cdot (\max_i \frac{1}{\alpha})^2})\leq 0 $$

这说明，在任何情况下MAS的Optimal Smoothing都不会比Unified Smoothing更差（虽然这是显而易见的）

在MAS中，我们为每个模态都存储了一个$S_m$，那么这意味着我们模型中每一层的权重W，对于每一个模态都要维护一个量化矩阵$S_mW$，这显然是我们无法接受的。MAS为了保证在PTQ过程中所有的模态之间共享一个量化权重，采用了如下方法（CMC）：

首先，我们仅存储一个量化权重$Q(S_tW)$,以文本模态为参考，并通过lora矫正来补偿其他模块。以视觉输入为例：理想情况下，我们计算：

$$ X_vS_v^{-1}\cdot(S_vW) $$

但使用共享权重则会产生残差：

$$ \Delta Y = X_vS_v^{-1}\cdot(\underbrace{S_vW-Q(S_tW)}_{\Delta W}) $$

那么我们可以对于每个非文本模态，我们都去存储这个残差，然后为了避免大矩阵的存储开销，我们可以使用低秩近似。

然而，我们不能直接对$\Delta W$使用SVD进行近似，因为事实上，我们需要近似的是残差$\Delta Y$,而非$\Delta W$（可以理解为$\Delta Y$是带权重的$\Delta W$）,且$\Delta W$不一定具有低秩结构（即前若干个大的奇异值不能解释大部分的能量）

我们现在来看我们的优化目标：

$$ \arg \min_{L} ||X_vS_v^{-1}(L-\Delta W)||^2_F $$

为了符号简便起见，我们令$A = X_vS_v^{-1}$,那么我们可以把最小化目标拆开写作:

$$ \arg \min_L \text{tr}((\Delta W - L)^\top A^\top A (\Delta W - L)) $$

那么一个自然的想法就是考虑能否通过某种线性变换消去这个权重的影响。在线性代数中我们知道，可以通过对一个矩阵进行白化（可将数据的协方差变为单位矩阵I）来达成我们的目的。

我们通过如下方式计算白化变换：

$$ \text{SVD}(A^\top A) = P\Lambda P^\top,T = (P\Lambda^{\frac{1}{2}})^\top $$

那么此时$AT^{-1}$是正交的。我们近似的目标可以描述为:

$$ \arg \min_L ||AT^{-1}T(\Delta W -L)||_F^2 = \arg \min_L ||T(\Delta W-L)||_F^2 $$

所以我们现在相当于在用一个rank为r的矩阵$TL$对矩阵$T\Delta W$进行逼近（经过实验验证，它具有低秩结构），因此我们可以对$T\Delta W$进行SVD截断:

$$ SVD(T(\Delta W)) = U\Sigma V^\top \approx U_r \Sigma_r V^\top_r $$

那么我们对白化进行逆变换后就可以得到低秩修正项：

$$ \Delta W = L_1L_2 \quad L_1 = T^{-1}U_r \quad L_2 = \Sigma_r V_r^\top $$

可以证明上述近似在秩r补偿近似下最优：

假设秩为r的矩阵$L = L_1L_2$,其中$L_1,L_2$由上式中定义的秩 r 截断 SVD 给出，它能够最小化重构损失。用形式化语言描述：

$$ \mathcal{L}=\sum_v ||X_vS_v^{-1}(\Delta W-L)||_F^2 \quad L^* = \arg \min_{\text{rank}(L)\leq r}\sum_v ||X_vS_v^{-1}(\Delta W-L)||_F^2 $$

证明如下：

只考虑两个模态，并且仅对权重进行量化，那么根据定义我们有：

$$ L^*=T^{-1}(\text{Trunc}_r(T\Delta W)) $$

由

$$ (X_vS^{-1}_v)^\top (X_vS_v^{-1}) = P\Lambda P^{\top} $$

可以推出：

$$ X_vS_v^{-1}=U\Lambda^{\frac{1}{2}}P^{\top}=UT $$

那么有

$$ \begin{align} \mathcal{L}(L^*) &= \left\| X_v S_v^{-1}(\Delta W - L^*) \right\|_F^2 \\ &= \left\| U T (\Delta W - L^*) \right\|_F^2 \\ &= \left\| U T \left(\Delta W - T^{-1}\operatorname{Trunc}_r(T\Delta W)\right) \right\|_F^2 \\ &= \left\| T\Delta W - \operatorname{Trunc}_r(T\Delta W) \right\|_F^2 \\ &= \sum_{i>r} \sigma_i(T\Delta W)^2 \\ &= \sum_{i>r} \sigma_i\!\left(U^{-1}X_vS_v^{-1}\Delta W\right)^2 \\ &= \sum_{i>r} \sigma_i\!\left(X_vS_v^{-1}\Delta W\right)^2 \\ &= L_{\min}^2. \end{align} $$

那么至此我们可以写出最终推理阶段将基础量化输出与模态特定修正结合起来：

$$ Y =\left \{ \begin{aligned} Q(X_mS_m^{-1})Q(S_tW), \quad m=\text{text}\\ Q(x_mS_m^{-1})Q(S_tW) + X_mS_m^{-1}\cdot L_1^mL_2^m,\quad m\neq \text{text} \end{aligned} \right. $$

MAS完整的流程如下图

[2026 ICLR] TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate

Wed, 26 Aug 2026 08:15:00 +0800

个人评价:这篇文章主要是针对大模型的KV-Cache的压缩，虽然是同一作者不同方法的浓缩（PolarQuant+QJL），但是补充了在结合方法下的量化误差上下界。目前实验停留在纯语言模型阶段，也许可以拓广到多模态阶段。

总结：本文主要介绍了一种针对高维向量量化的创新方法，旨在通过大幅度压缩数据规模来优化AI模型推理，KV Cache管理以及向量数据库检索的效率。其核心在于结合了随机旋转技术和最优标量量化器，能在极低的比特位宽下实现接近理论极限的MSE。针对内积检索中的偏置问题，作者设计了一个两阶段架构，利用 1-比特 QJL 变换补偿余数，从而确保了内积估算的无偏性。实验数据表明，该算法在 Llama-3.1 等大语言模型的长文本测试中，仅需 2.5 至 3.5 比特即可保持与全精度近乎一致的性能。此外，相较于传统的乘积量化 (PQ) 技术，TurboQuant在保持高召回率的同时，将索引构建时间降低至接近于零，展现出卓越的加速器友好性。

简单来说，进行向量量化（VQ）的目的是最小化下列两个误差：

$$ D_{MSE} = \mathbb{E}_{Q}[||x-Q^{-1}(Q(x))||_2^2]\tag{1} $$$$ D_{prod} = \mathbb{E}_Q[|\langle y,x\rangle-\langle y,Q^{-1}(Q(x))\rangle|^2]\tag{2} $$

此外，对于内积量化，在大模型推理中我们更希望向量的内积是无偏的，即满足：

$$ \mathbb{E}_Q[\langle y,Q^{-1}(Q(x))\rangle]=\langle y,x\rangle\tag{3} $$

而这两个优化目标是难以兼顾的，因此在VQ中通常会设计两个Quantizer，分别是$Q_{MSE}$,$Q_{prod}$。对应到KV-Cache量化场景下就是对K向量用$Q_{prod}$，对V向量用$Q_{MSE}$。

对于$Q_{MSE}$，我们的目标就是最小化公式(1)。在此之前，我们有如下假设:

待量化向量满足$||x||^2 = 1$，即$x\in \mathbb{S}^{d-1}$,即分布在d维球面上。若不满足这个条件，在实际中可以通过存储L2范数进行Scale使向量满足条件。

我们有如下引理（Lemma 1）：

若$x\in \mathbb{S}^{d-1}$,是在单位超球面上均匀分布的随机变量，那么对任意$j\in [d]$,坐标$x_j$服从（缩放/平移后的）Beta型分布：

$$ x_j ～ f_X(x)=\frac{\Gamma(\frac{d}{2})}{\sqrt{\pi}\Gamma(\frac{d-1}{2})}(1-x^2)^{\frac{d-3}{2}},\quad x\in[-1,1] $$

在高维情况下，该分布收敛到正态分布：

$$ f_X(.)\to N(0,\frac{1}{d}) $$

证明略.

High-Level层面上可以理解为固定球上一点的一个坐标相当于用一个平面去截这个高维球面，那么此时其余坐标构成的截面是一个维度为d-2，半径为$\sqrt{1-x^2}$的球面。（可以想象一下三维球面被平面截得到圆），那么这个分布自然就是中间多（x=0），两边少（$x=\pm 1$）.

在这个引理的支持下，我们对原始的向量x乘以一个随机的旋转矩阵$\Pi $（这里相当于做了一个极坐标变换），使其成为在单位超球面上均匀分布的随机变量$z = \Pi x$。那么此时根据Lemma 1，z的每个坐标都可以认为符合上述Beta型分布，且在高维情况下收敛为正态分布。此外，在高维下，z不同坐标之间会变得近似独立，因此我们可以对每个坐标独立地应用最优标量量化器。于是我们的问题转变为：

为服从如下分布的随机变量设计一个标量量化器。

$$ x_j ～ f_X(x)=\frac{\Gamma(\frac{d}{2})}{\sqrt{\pi}\Gamma(\frac{d-1}{2})}(1-x^2)^{\frac{d-3}{2}},\quad x\in[-1,1] $$

在随机变量分布给定的情况下的最优标量量化问题可以表述为一个一维连续K-means问题。更具体而言，我们希望把区间[-1,1]划分为$2^b$个簇。最优解需要满足：当所有质心按照升序排序时，区间边界应当是相邻质心的中间。因此，若记这些升序排序的质心为$c_i$那么，该标量量化问题可以描述为如下k-means优化问题：

$$ C(f_x,b) = \min_{-1 \leq c_1\leq c_2\leq \dots \leq c_{2^b}\leq 1}\sum_{1}^{2^b}|x-c_i|^2f_x(x)dx\tag{4} $$

该问题可以通过迭代数值方法进行求解（Lloyd-Max量化器，本质和K-means相似，可以看作一维的K-means）。此外，我们只需要针对一组实际有效的bit-width b离线求解一次，然后把结果存储下来，供量化器之后重复使用。

至此，$Q_{MSE}$的做法很明确：先计算$z=\Pi x$,然后对z的每个坐标找到最近的质心，并存储该质心的索引。对应的反量化流程则通过读取这些索引对应的质心来重建旋转后的向量，再乘以$\Pi^\top$从而得到原始向量。

论文中还给出了该量化器损失的上界，理解起来并不困难，这里不做过多赘述。

之所以$Q_{MSE}$需要与$Q_{prod}$不能统一，是因为前者不满足内积无偏性的特性，即式(3)。（论文中给出了证明）

为了保证$Q_{prod}$的内积无偏性，作者提出了将$Q_{MSE}$与QJL相结合的方案。具体而言，设$Q_{MSE}$是对应于位宽 b−1 的$ Q_{mse} $的量化映射。对于任意$x\in \mathbb{S}^{d-1}$，我们定义残差向量：

$$ r := x - Q_{mse}^{-1}(Q_{mse}(x)) $$

其L2范数很小，即在期望意义下(见式（4）)

$$ \mathbb{E}[||r||]=\sqrt{C(f_X,b-1)} $$

随后，我们可将QJL 的量化映射 $Q_{QJL}$ 应用于该残差向量，从而使总体位宽达到 b，并得到如下无偏内积估计器：

$$ \langle y,Q^{-1}_{MSE}(Q_{MSE}(x))\rangle + ||r||^2\cdot\langle y,Q^{-1}_{qjl}(Q_{qjl}(r))\rangle $$

更形式化的来说，我们可以定义：

$$ Q_{prod}(x) = [Q_{MSE}(x),Q_{qjl}(x-Q^{-1}_{MSE}(Q_{MSE}(x))),||x-Q^{-1}_{MSE}(Q_{MSE}(x))||_2] $$

论文还对该方法的误差下界进行了估计，具体的参考原文。

[2023 ICLR] GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers

Wed, 26 Aug 2026 08:12:00 +0800

这个文章是OWQ的前身，借着对这篇文章的分析，我们梳理一下这一系列的文章的intuition。大概的分析顺序为:

$$ \text{OBD} \to \text{OBS} \to \text{OBC} \to \text{GPTQ} $$

首先是OBD，这是由Yann LeCun在1990年提出的神经网络剪枝算法。该算法基于二阶导数信息，旨在通过去除目标函数影响较小的参数来降低模型复杂度，提高泛化能力。

具体而言，就是希望去除目标函数对目标函数E(即Loss)影响小的参数，我们记去除了若干参数的模型的参数为$\hat W = W + \Delta w$,有:

$$ \Delta E = L(x,y,W) - L(x,y,W + \Delta w) = \sum_{i}g_i \Delta w_i + \frac{1}{2}\sum_{i}h_{i,i}\Delta w_{i}^2 + \frac{1}{2}\sum_{i\neq j}h_{i,j}\Delta w_i\Delta w_j + O(\Delta w^3) $$

其中$g_i = \nabla L$,$h_{i,j}$为Hessian矩阵$H_{L}$的一个元素

其中由于剪枝发生在对于已经训练好的神经网络，因此一阶导项可以忽略不计，而高阶项由于模型会进行归一化，因此$\Delta w$较小，可以忽略不计。

此外，OBD做了一个有争议的假设，即删除任意一个参数后，其他参数对目标函数的影响不变，也就是说每个参数对目标函数的影响是独立的，因此可以忽略交叉项:

那么我们可以得到简化后的公式:

$$ \Delta E = \frac{1}{2}\sum_{i}h_{i,i}\Delta w_i^2 $$

因此，对神经网络进行剪枝，删除参数时，参数对目标函数的影响可以通过海森矩阵的对角项进行衡量。我们只需要在剪枝时求出海森矩阵，按对角项从小到大排序，即可确定参数剪枝的次序。

可以注意到，OBD的这个认为参数对目标函数的影响是独立的假设是很强的。OBS认为参数之间的独立性不成立，如果考虑交叉项，可以写作矩阵形式

$$ \Delta E = \frac{1}{2}\Delta w^\top H\Delta w $$

OBS希望在W每次迭代找到一个位置q(即准备剪枝的位置，后续会将该位置的$w_q = 0$),以及在获得位置q的同时，计算处一个与之相关的$\Delta w$对w进行补偿，使得$L(w+\Delta w)-L(w)$尽量小。

那么这个流程可以描述为一个带约束的凸优化问题:

$$ \arg \min_q \frac{1}{2}\Delta w^\top H \Delta w\\ s.t.\ e_q^\top\Delta w + w_q = 0 $$

这里$e^\top_q$是第q个值为1的列向量。

采用Lagrange乘子法进行求解:

$$ \mathcal{L} = \frac{1}{2}\Delta w^\top H \Delta w + \lambda(e^\top_q \Delta w + w_q) $$

对$\lambda$并置为0得到:

$$ e_q^\top \Delta w + w_q = 0\\ \Delta w^\top e_q + w_q = 0 $$

对$\Delta w$求导并置为0有:

$$ \Delta w^\top H + \lambda e^\top_q = 0\\ \Delta w^\top H H^{-1} + \lambda e_q^\top H^{-1}=0\\ \Delta w^\top + \lambda e^\top_q H^{-1} = 0 $$

有

$$ w_q = \lambda e_q^\top H^{-1} e_q\\ \lambda = \frac{w_q}{[H^{-1}]_{qq}} $$

其中用到了等式$e_q^\top H^{-1}e_q = [H^{-1}]_{qq}$

将$\lambda = \frac{w_q}{[H^{-1}]_{qq}}$带入等式$\Delta w^\top H + \lambda e_q^\top = 0$,得到

$$ \Delta w^{\top} = -\frac{w_q}{[H^{-1}]_{qq}}e_{q}^\top H^{-1}\\ \Delta w = -\frac{w_q}{[H^{-1}]_{qq}}(H^{-1})^\top e_q \\ \Delta w = -\frac{w_q}{[H^{-1}]_{qq}}H_{:,q}^{-1} $$

其中$H_{:,q}^{-1}$表示$H^{-1}$的第q列,且Hessian矩阵是一个对称矩阵.

将$\Delta w$带入$\Delta \mathcal{L}$我们有:

$$ \Delta\mathcal{L} = \frac{1}{2}\Delta w^\top H \Delta w = \frac{1}{2}(-\frac{w_q}{[H^{-1}]_{qq}}e_q^\top H^{-1})H(-\frac{w_q}{[H^{-1}]_{qq}}H^{-1}e_q)\\ =\frac{1}{2}(\frac{w_q}{[H^{-1}]_{qq}})^2e_q^\top H^{-1}e_q\\ =\frac{1}{2}(\frac{w_q}{[H^{-1}]_{qq}})^2[H^{-1}]_{qq}\\ =\frac{1}{2}\frac{w_q^2}{[H^{-1}]_{qq}} $$

由此我们得到:

$$ q = \arg \min_q \frac{w_q^2}{[H^{-1}]_{qq}} $$

我们不难发现，要进行k次剪枝，每一次剪枝都要求一次Hessian矩阵的逆(时间复杂度为$O(d^3)$),这样的时间复杂度明显是不能实际应用的，因此OBC对其进行了进一步的优化。

OBC主要做了两点优化，一个是对原始问题进行了拆分，另一个是对Hessian矩阵的计算进行了简化。

首先是对原始问题的拆分，对于Layer-wise的量化/剪枝而言，通常将对整个网络进行的量化/剪枝拆分为每一层独立的子问题。在先前对AdaQuant的分析中我们有提到，这种Layer-wise的拆分是高度可并行的。

在Layer-wise的量化/剪枝下，参数变化带来的损失可以描述为以下形式:

$$ \Delta \mathcal{L} = ||W_lX_l - \hat W_l X_l||_2^2 $$

其中$\hat W$表示经过量化/剪枝后的参数。

OBC将这个损失函数进行了按行拆分，即认为删掉某个权重$w_{ij}$只影响该行的输出，行与行之间的Hessian矩阵元素是没有耦合的。(这两个都是对按行拆分合理性的解释，前者是直观解释，后者是数学解释)

对于第一点，我们知道，改变某个权重$w_{ij}$它只会对输出的某一行的结果产生影响，即$Y_{i,:} = W_{i,:}X$,那么既然只对某一行的输出产生影响,那么对于整体误差而言,也只对这一行的误差产生影响,而误差是可以按行拆分的:

$$ \Delta \mathcal{L} = ||WX - \hat W X||_2^2 = \sum_{i=1}^{d_{row}}||W_{i,:}X-\hat W_{i,:}X||_2^2 $$

因此剪枝/量化是可以按行拆分并行处理的。

由于是Layer-wise的量化/剪枝，我们在这个尺度下的进行单行损失函数(二阶范数)的Hessian矩阵的计算从而对第二点进行证明，我们有:

$$ H_{pq} = \frac{\partial^2\Delta \mathcal{L}_l}{\partial w_{lp}\partial{w_{lq}}} = \frac{\partial}{\partial w_{lp}}\sum_{k=1}^N 2(\sum_{j=1}^{d_{col}}(w_{lj}-\hat w_{lj})x_{jk})\frac{\partial}{\partial w_{lq}} \sum_{j = 1}^{d_{col}}(w_{lj}-\hat{w_{lj}})x_{jk}\\ =\frac{\partial}{\partial w_{lp}}\sum_{k=1}^N2(\sum_{j=1}^{d_{col}}(w_{lj}-\hat w_{lj})x_{jk})x_{qk}\\ =2\sum_{k=1}^N x_{pk}x_{qk} $$

写成矩阵的形式就是

$$ H = 2XX^\top $$

发现每一行的损失的Hessian矩阵只跟输入数据X有关且相等，而与模型权重无关，因此我们认为行与行之间的Hessian矩阵元素是没有耦合的。

而我们知道通过泰勒展开可以得到参数变化对损失函数的影响的近似表示$\Delta w^\top H \Delta w$,那么结合上式我们可以知道行与行之间的损失是相互独立的(对于行而言$\Delta w$,行与行之间互相独立,对于Hessian矩阵而言，行与行之间相等且互不影响),由此可以从数学上说明按行拆分进行单独处理的方式是合理的。

有了这个证明，我们可以对每行进行单独处理进行量化剪枝。这种方式为我们提供了一个更加简单的Hessian矩阵形式$2X^\top X$但是每次更新参数仍然需要对其求逆，因此OBC提供了一个高效的求逆方法:

给定一个可逆矩阵H以及其逆矩阵$H^{-1}$,我们希望高效地计算删除H第q行第q列(删除权重$w_q$)后的逆矩阵$H^{-1}_{-q}$:

$$ H_{-q}^{-1} = (H^{-1} - \frac{1}{[H^{-1}]_{qq}}H^{-1}_{:,q}H^{-1}_{q,:})_{-q} $$

这个定理证明较为复杂，将在博客上更新详细证明与intuition。

接下来我们来描述OBC剪枝的完整流程:

给定一个神经网络层的权重行向量$w \in \mathbb{R}^d$,以及其对应的Hessian矩阵的逆$H^{-1}\in \mathbb{R}^{d\times d}$,要求切除其中k个权重，同时最小化输出误差。

初始化 $M \leftarrow \{0,1,2,\dots,d-1\}$,为尚未被剪枝的权重索引集合。
重复执行k次:

首先选择当前最优的剪枝目标q:$q \leftarrow \arg\min_{q\in M}\frac{w_q^2}{[H^{-1}]_{qq}}$,接着弥补剪掉$w_p$带来的误差$\Delta w \leftarrow \Delta w - \frac{w_q}{[H^{-1}]_{qq}}(H^{-1}_{:,q})^\top$,然后更新$H^{-1}\leftarrow H^{-1}_{-q}$,从候选集中移除该索引$M\leftarrow M - \{q\}$

而对于量化而言(OBQ)，相对剪枝我们需要做一些调整。首先是之前的最优化问题的限制条件需要改为:

$$ \Delta w \cdot e_q + w_q - \text{quant}(w_q) = 0 $$

同样使用Lanrange乘子法进行求解:

$$ \mathcal{L} = \frac{1}{2}\Delta w^\top H \Delta w + \lambda(\Delta w \cdot e_q + w_q - \text{quant}(w_q)) $$

对$\Delta w$求导并置为0可以得到:

$$ \Delta w^\top H + \lambda e_q = 0\\ \Delta w^\top = -\lambda e_q H^{-1}\\ \Delta w = -\lambda H^{-1}e_{q}^\top $$

对$\lambda$求导并置为0可以得到:

$$ \Delta w\cdot e_q + w_q - \text{quant}(w_q) = 0\\ w_q - \text{quant}(w_q) = \lambda H^{-1}e_q^\top e_q\\ \lambda = \frac{w_q - \text{quant}(w_q)}{[H^{-1}]_{qq}} $$

得到:

$$ \Delta w = -\frac{w_q - \text{quant}(w_q)}{[H^{-1}]_{qq}}(H^{-1}_{:,q})^\top\\ q = \arg \min_{q} \frac{(w_q-\text{quant}(w_q))^2}{[H^{-1}]_{qq}} $$

用上式替换OBC剪枝的流程，便可以得到量化流程。

GPTQ则是对OBQ进行了改进，GPTQ发现，在对权重的每一行量化时，按照贪心策略选择量化的q和按照任意固定顺序来量化每一行的权重最终的误差是相差不大的，那么可以直接让所有行都按照列序(0 $\rightarrow$ col),这样可以提高计算效率与存储效率。

这样做的另一个好处在于：每行的顺序一样，那么每一行对应的Hessian矩阵都是相同的，每次Hessian矩阵的逆只需要计算一次！

在固定量化顺序的前提下，我们不再需要求解$q = \arg\min_q$只需要关心$\Delta w$,此时$\Delta w$的更新公式为:

$$ \Delta w = -\frac{w_q - \text{quant}(w_q)}{[H_{q:,q:}]^{-1}_{0,0}}([H_{q:,q:}]^{-1}_{:,0})^\top $$

这个式子我们对比在贪心策略下的式子便不难发现，这个式子就是贪心策略式子在每次删除当前候选集里的第一个q时的式子的等价形式，同样地我们也能给出Hessian矩阵的逆的更新形式:

$$ [H_{q:,q:}]^{-1} = ([H_{q-1:,q-1:}]^{-1} - \frac{1}{[H_{q-1:,q-1:}^{-1}]_{0,0}}[H^{-1}_{q-1:,q-1:}]_{:,0}[H_{q-1:,q-1:}^{-1}]_{0,:})_{1:,1:} $$

从矩阵的角度来看我们在做的是这样一个变换:

$$ (H^{-1})^{(k)} = \left[ \begin{matrix} I_{k-1} & 0 & 0^\top\\ 0 & a_{k,k} & b_{k}^\top\\ 0 & b_k & B'^{(k)} \end{matrix} \right]\to (H^{-1})^{(k+1)} = \left[ \begin{matrix} I_{k} & 0 & 0\\ 0 & a_{k+1,k+1} & b_{k+1}^\top\\ 0 & b_{k+1} & B''^{(k+1)} \end{matrix} \right] $$

若我们不断更新Hessian的逆总会产生非正定的Hessian逆矩阵,其原因可能是由于数值误差的累积。为了解决这个问题，作者注意到每次从$H^{(-1)}$中删除一行一列，本质上和对称正定矩阵的Cholesky分解的逐步过程类似，因此作者对初始的$H^{-1}$进行了Cholesky分解，得到了一个上三角矩阵$T$。

Cholesky分解:假设一个正定矩阵$A\in \mathbb{R}^{n\times n}$是正定对称矩阵，那么必然存在一个对角元素为正数的下三角矩阵$L\in \mathbb{R}^{n\times n}$满足$A = LL^\top$

我们尝试模拟一次这个分解过程:

$$ A = \left[ \begin{matrix} a_{11} & A_{21}^\top\\ A_{21} & A_{22} \end{matrix} \right],L = \left[ \begin{matrix} l_{11} & 0\\ L_{21} & L_{22} \end{matrix} \right],L^\top = \left[ \begin{matrix} l_{11} & L_{21}^\top\\ 0 & L_{22}^\top \end{matrix} \right] $$

由于$A = LL^\top$,我们有:

$$ l_{11} = \sqrt{a_{11}},L_{21} = \frac{1}{l_{11}}A_{21},L_{22}L_{22}^\top = A_{22} - L_{21}L_{21}^\top $$

于是我们可以惊奇地发现$L_{22}L_{22}^\top$就是我们想要的$H_{q:,q:}^{-1}$!因此我们可以认为删去$[H_{q:,q:}^{-1}]$的第一行和第一列的过程与对该矩阵进行一次Cholesky分解是等价的。因为我们进行Cholesky分解得到的$L_{22}$恰好是更新了之后的$H^{-1}$进行Cholesky分解得到的下三角矩阵。

进一步地，GPTQ对初始的Hessian矩阵的逆进行了Cholesky分解得到一个上三角矩阵$L^\top$,这个矩阵还有一个特点在于，它的每一行刚好就等于逆矩阵每次更新迭代后的第一行乘以一个常数:

$$ C_qL_{q,q:}^\top = [H_{q:,q:}]^{-1}_{0,:} $$

这个我们可以通过Cholesky分解的式子知道，因为分解得到的$L$是一个下三角矩阵，那么它的第一行就只有一个常数，而这个常数乘以$L^\top$便可以得到A的第一行。

而恰好我们发现，$\Delta w$的更新公式只需要用到当前Hessian矩阵的逆的第一行，那么我们有:

$$ \Delta w = -\frac{w_{:,q}- \text{quant}(w_{:,q})}{C_q T_{qq}}C_qT_{q,q:} $$

其中常数可以直接约掉:

$$ \Delta w = -\frac{w_{:,q} - \text{quant}(w_{:,q})}{T_{qq}}T_{q,q:} $$

因此我们在进行量化时不用每次都更新Hessian矩阵的逆，而是直接对$H^{-1}$进行Cholesky分解，得到它的每一行便可以进行参数的量化。

此外，如果每行的量化并行计算，那么每次更新都要读写一次参数矩阵。若参数矩阵的维度为$d_{row}\times d_{col}$，那么量化这个参数矩阵就要读写$d_{col}$次参数，总共的读写量高达$d_{row}\times d_{col}^2$.(因为我们量化第i列的时候，后面的列相应地也要补偿更新)

那这样大量的IO开销将会成为瓶颈，因此GPTQ采用了Lazy Batch-Update技术。我们注意到对于列i，最终的量化决策并不会受到尚未更新的列的影响。这使得我们可以将后续列的更新推迟到后续步骤中，从而减少不必要的内存操作。

具体步骤如下:

每次处理B列的一个小块，限制该列更新的补偿更新只影响块内的列
当前块的权重会在量化过程中被更新，而其影响暂时不传播到矩阵的其他部分
对当前块中的每一列进行量化，同时计算误差并更新当前块剩余的列
这些更新仅在当前块内进行，而不会影响整个矩阵
当一个块内的所有列完成量化后，将该块的更新结果批量应用到矩阵的剩余部分

这部分通过语言可能难以描述清楚，可以见下图:

[2024 AAAI] OWQ: Outlier-Aware Weight Quantization for Efficient Fine-Tuning and Inference of Large Language Models

Wed, 26 Aug 2026 08:09:00 +0800

总结：本文提出了一个异常感知的权重量化方法OWQ，利用LLMs中的异常激活值挑选出Weak Column，对其采用全精度的方式在牺牲很小的性能的情况下提升了巨大的精度。此外为进一步提升其性能做了一定的硬件适配并提出了一个基于OWQ的WTC方案，简单来说就是在OWQ量化模型上微调只更新Weak Column的参数。

评价: 这篇文章思路很新颖，从大模型中间激活值的异常出发，结合Hessian矩阵分析，提出了一个简单但高校的方法，似乎可以进一步提升？

本文基于这样一个发现，LLMs在中间激活中表现出一些异常值，其值显著大于其他值，并且这些异常值集中在特定的特征维度上。保留这些异常值的值已知对于在激活量化之后保持准确性至关重要。此外，作者团队还发现激活异常值仍然会影响权重量化的敏感性。基于此，作者提出了一种称为异常值感知权重量化的概念(OWQ).

我们知道逐层权重量化的过程，实际上进行如下优化过程:

给定输入特征$X\in R^{C_{i,n} \times N}$,其中$C_{i,n}$表示输入的通道数，N是输入的序列长度，用于$C_{out}$输出特征的完整精度权重矩阵$W \in R^{c_{out}\times C_{i,n}}$被映射到低精度。

$$ \arg \min_{W} E= \arg \min_{\hat W}||WX - \hat W\hat x||^2_2 $$

在量化时我们从输入到输出逐层量化。此外大模型的量化中，Embedding层以及LM Head的权重通常不被量化，因为前者的权重量化误差会随着网络的传播不断放大，且Token向量较为稀疏而LM Head层直接决定logits，而Top-k词之间的分数差往往较小，低比特改写会对排序以及argmax产生显著影响。

接下来阐释权重敏感性和激活异常值之间的关系。

（这里与原论文中的证明方式不一致，原论文只考虑了对量化误差

$$||WX - \hat W\hat X||$$

的论证，这是单层量化误差，但是忽略了误差随着网络的放大的影响）

在AdaRound文章中有提到，对于权重的量化产生的误差我们有如下基于泰勒展开的近似:

$$ \mathbb{E}[L(x,y,w+\Delta w) - L(x,y,w)] =\Delta L \approx \nabla L^{\top}\Delta w + \frac{1}{2}\Delta w^\top H\Delta w \approx \Delta W^\top H\Delta W $$

由此我们知道，输出误差可以直接与海森矩阵和权重扰动的幅度相关。

将Hessian矩阵写作Kronecker积的矩阵形式，我们得到

$$ H(w^{(l)}) = \mathbb{E}[x^{(l-1)} x^{(l-1)\top}⊗ \nabla^2_{z^{(l)}}L] $$

由此我们可以从全局的视角看到，异常激活值(激活值的激增)使Hessian矩阵H的某些元素具有异常大的值。Hessian矩阵的这种异常激增增加了相应权重通道对量化的敏感性。具体来说，即使在相同的权重扰动下，由于一些H的一些大元素，输出的变化也会相当大。我们可以将这些易受量化影响的权重称为Weak Column，特别是那些与特定输入通道中的激活值异常值相关联的权重。

OWQ为了解决这个问题，实现了如下技术：首先，识别Weak Column并将他们从量化中排除。随后使用精心调整的量化参数将剩余的权重量化为极低的bit。

对于Weak Column的检索，OWQ遵循如下方法:

我们定义j-th权重列的敏感性为:

$$ sensitivity_j = \lambda_j||\Delta W_{:,j}||_2^2 $$

其中$\lambda_j$是Hessian矩阵的第j个对角元素。

(若考虑的Hessian矩阵是层内重构误差的，那么这里$\lambda_j = (X^\top X)_{j,j} = 2\sum_{n}x_{j,n}^2$)

可以注意到在我们写的Hessian矩阵是针对全局损失而言的，那么这个场景下的$\lambda_j$就有所改变。虽然在这个场景下我们无法像论文里一样因为layer-wise量化误差输出通道之间没有Hessian交互，从而Hessian是对角矩阵。

但是将其作对角近似是合理的。因为在这样一个大的模型下，进行Hessian矩阵的精确计算是不可行的。

在此场景下我们有:

对于神经网络中的第j个神经元的输入$a_j$对应权重为$w_{ji}$:

$$ \frac{\partial^2 E_n}{\partial w_{ji}^2} = \frac{\partial^2 E_n}{\partial a_j^2}z_i^2 $$

其中$z_i$是上一层神经元的输出。

而$\frac{\partial^2 E_n}{\partial a_j^2}$可以通过链式法则递归计算(类似于反向传播):

$$ \frac{\partial^2 E_n}{\partial a_j^2} = \underbrace{\frac{\partial}{\partial a_j}[h'(a_j)\sum_{k}w_{kj}\frac{\partial E_n}{\partial a_k}]}_{链式法则}=h'(a_j)^2 \sum_{k,k'}w_{kj}w_{k'j}\frac{\partial^2 E_n}{\partial w_k\partial w_{k'}} + h''(a_j)\sum_k w_{kj}\frac{\partial E_n}{\partial a_n} $$

忽略二阶导中的非对角线项$k\neq k'$:

$$ \frac{\partial^2 E_n}{\partial a_j^2}\approx \underbrace{h'(a_j)^2\sum_k w_{kj}^2 \frac{\partial^2 E_n}{\partial a_k^2}}_{链式法则} + \underbrace{h''(a_j)\sum_{k}w_{kj}\frac{\partial E_n}{\partial a_k}}_{链式法则} $$

从而一次反向传播便可以计算出来，时间复杂度为$O(W)$

由此我们可以计算出在全局损失下的$\lambda_j$

在实际的计算中，同样也是只需要一个小批量的校验集，但是坏处是要多进行一次反向传播得到曲率。这样的Trade off得到的性能提升应该不小，因为它会真正识别出在全局下的 Real Weak Column

我们根据权重列的敏感性挑选出top-k个作为weak column。然后其余权重被量化为低精度。（这里可以采用任何的量化方法）论文中采用了OPTQ的方法。

作者团队对OPTQ进行了重要修改，使用二维网格搜索来搜索量化配置，包括步长以及零点。通过四舍五入到最接近的截断来搜索使量化前后差异最小的参数的最优值。

文章指出了一个利用Weak Column进一步减轻误差的方法:将高精度的Weak Column重新排列到权重的末尾，OPTQ过程中其他列的量化误差可以主要由Weak Column得到补偿

在此之后，我们将Weak Column存储为fp16，并为每一列使用一个额外的整数，该整数用于索引Weak Column。此外存储一个低精度矩阵，其中Weak Column的位置采用0填充。

此外，作者还对OWQ格式在真实GPU上提供了专门的加速以及WTC微调方案。

具体而言，这个微调方案会将OWQ的量化模型进行微调但只对Weak Column进行参数更新。因为weak column的数量很少，所以总体微调参数量很少，同时又因为weak column的权重使用fp16进行存储，因此微调空间较大，能够实现较好的微调效果。

[2022 NeurIPS] ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers

Wed, 26 Aug 2026 08:06:00 +0800

评价: 这篇文章比较Solid，考虑了硬件适配的问题，这是模型量化中一个老大难的问题尤其是混合精度。但是实验的模型都是参数规模较小的模型，在大模型上的效果有待考究。

总结：这篇文章指出，低比特量化在大型 Transformer 架构模型中精度受限的主要原因是激活值和权重矩阵的值分布方差较大。针对这一问题，提出了 ZeroQuant 方案。该方案主要包括：对权重采用 Group-wise 量化、对激活值采用 Token-wise 量化，这种方法既能适配硬件架构，又能保持较高的精度；同时，通过 Layer-wise 知识蒸馏方法来减少量化带来的精度损失。

文章指出，在大模型的量化中，采用PTQ会面临以下挑战。

激活分布动态性强

论文通过展示每一层的激活值在不同Token语义下的分布，发现了其范围随输入token的语义上下文变化极大的特点。这一特点使得难以对所有的token使用固定的量化范围。
权重矩阵范围差异大

通过同样的方式，展示了不同层行权重的范围。同样可以看到权重矩阵的神经元范围差异较大。

这两个挑战使得Per-Tensor粒度的量化很难在大模型的量化中使用，而采用Per-channel粒度的量化会面临大的计算存储开销，且会导致在硬件级别的矩阵乘法优化难以执行。基于此作者提出采用Group-wise量化旨在对精度和实用性做出权衡。

对于权重矩阵而言，Group-wise量化就是将$W\in R^{n\times m}$划分为g个组，每个组单独量化。但是在最先提出这个量化方法的Q-BERT中仅将其用于QAT且没有考虑硬件效率约束，以及系统后端支持。基于此作者团队考虑了GPU的架构（Ampere架构）的硬件约束，特别是将Group-size与Tensor Core中的计算单元对齐。

具体而言，Tensor Core允许16*16大小的矩阵块在一个warp中并行处理，从而加速矩阵乘法和其他张量操作。如果我们让Group-size为16或32这样与Tensor core中矩阵乘法相适配的大小，这样就能在降低延迟的同时保持模型精度。

注：这部分文章在附录D中有详细介绍，简单来说Group-wise的group size是通过CUTLASS库和Profiler工具，根据输入尺寸和硬件特性动态确定的，以优化Tensor Core的计算效率。

对于激活值而言，在挑战中我们已经阐明了在不同的Token上下文语义下，激活值的范围存在巨大方差，因此解决这个问题的一个自然而然的想法是采用Token-wise的量化策略。但是直接采用DL框架中的Token级量化会导致显著的量化和反量化成本，因为引入了额外的操作。基于此作者采用了算子融合的方法等一系列优化。

知识蒸馏是缓解模型压缩后精度下降的最强大的办法之一。因此，论文提出了一种逐层的知识蒸馏技术来避免低比特量化带来的精度损失。

具体而言，在传统的知识蒸馏中，教师模型和学生模型的输出通常是整个模型的输出，但在逐层知识蒸馏(LKD)中，蒸馏的学习目标是逐层的，即学生模型要学习教师模型每一层的中间激活值。

假设我们要量化的是$L_k$层，其量化版本为$\hat L_k$,然后我们使用$L_{k-1}$层的输出作为$L_k,\hat L_k$的输入，测量差异，并更新模型

$$ L_{LKD,k} = MSE(L_kL_{k-1}\dots L_1(X)-\hat L_k L_{k-1}\dots L_1(X)) $$

因为使用相同的前k-1层，所以无需单独保留一个单独的教师模型，因此额外的模型成本仅仅是$L_k$。而每次只对一层进行蒸馏，所以内存和计算开销非常小，并且无需原始训练数据。

在前面Token-wise的量化处我们提到，作者对Kernel做了对应的优化，下面我们详细展开。

首先是针对Token-wise 的激活值量化做了一系列的kernel融合。作者将激活值量化与其相关的逐元素和或基于reduction的操作（如bias，GELU,LayerNorm等）的kernel进行了融合。这样减少了数据转移的开销。而将反量化与矩阵乘法做了相应的融合。具体而言见下面的流程图

flowchart LR
x -->B((LN/GeLU))
B --> C(Quantize)
C --> D[GEMM]
D --> E(DeQuantize)

经优化后

flowchart LR
x --> A((LN/GeLU + Quantize))
A --> B[GeMM + DeQuantize]

[2020 ICML] AdaRound: Up or Down? Adaptive Rounding for Post-Training Quantization

Wed, 26 Aug 2026 08:03:00 +0800

总结:本篇文章作者首先从数学角度证明了在模型量化过程中，直接将浮点数进行四舍五入round到最近定点数的方法并不是精度最优的。并且通过了一个简单的实验验证了猜想，随后基于此作者进行一系列的数学推导和数学近似推导除了最终的优化目标:最小化由于量化在预激活值中引入的均方误差，从而提出了自适应的Round方法:AdaRound.这种方法在进行量化时，自适应地决定将浮点值转到最近右定点还是左定点值。AdaRound可以在不需要QAT or finetune的情况下仅使用少量无标签的校准数据在精度上达到SOTA，甚至4bit量化也可以保留较好的精度。

个人评价:这篇文章的行文流惯，公式推导顶级，从完全理论的方式推导出了大部分量化论文中量化目标函数。

首先作者将量化过程定义为了一个对预训练模型权重w的微小扰动$\Delta w$.我们的目标为最小化这个扰动对损失函数$L(w)$造成的影响，即最小化$E[L(w+\Delta w)-L(w)]$,为了近似这个损失，采用了二阶泰勒展开有:

$$ L(w + \Delta w)\approx L(w) + \nabla L(w)^\top \Delta w + \frac{1}{2}\Delta w^{\top} H(w)\Delta w\\ \Delta L \approx \nabla L^{\top}\Delta w + \frac{1}{2}\Delta w^\top H\Delta w $$

由于模型经过预训练损失函数的梯度很小，可以忽略，而高阶项只要扰动$\Delta w$不是特别大，二阶近似往往就是准确的。对于4-bit或更高精度而言这个是成立的。

因此我们可以认为影响模型精度的主要是$\Delta w$以及损失函数的曲率$H(w)$相关。

令$\Delta w^\top = [\Delta w_1,\Delta w_2]$,$H^{(w)} = \begin{bmatrix}1 & 0.5\\0.5 &1\end{bmatrix} $,那么由此我们可以计算出，量化导致的损失为:

$$ \Delta w^\top H^{(w)}\Delta w = \Delta w_1^2 + \Delta w_2^2 + \Delta w_1\Delta w_2 $$

对于对角线项$\Delta w_1^2,\Delta w_2^2$而言四舍五入是最优的，最小化了误差，但是对于非对角线项$\Delta w_1\Delta w_2$采用四舍五入就不一定最优了。例如若二者符号取反乘积为负就可以抵消一部分损失的增量。

因此从理论上分析出了四舍五入方法的局限性。后续也从实验上进行了论证，作者采用四舍五入，全部向上，全部向下，随机舍入进行比较，发现在随机舍入中存在比四舍五入高出10%的取舍法，说明在取舍办法中，存在更优的方法。

这个取舍办法的选取可以通过如下问题描述。

假设每层权重量化，量化后的权重为$\hat w_i^{(l)}$

$$ \hat w_i^{(l)}\in \{w_i^{(l),floor},w_i^{(l),ceil}\} $$

$\Delta w_i^{(l)} = w^{(l)} - \hat w_i^{(l)}$,由此，最优的舍入过程可以描述为以下二元优化问题:

$$ \arg \min_{\Delta w} \mathbb{E}[L(x,y,w+\Delta w) - L(x,y,w)] $$

直接对这个式子进行优化并不现实，因为，每次调整$\Delta w$都需要进行一次前向传播，计算成本太高，我们采用前面理论分析时的泰勒展开近似。此外，忽略属于不同层之间权重的交互。优化目标近似为：

$$ \arg \min_{\Delta w^{(l)}} \mathbb{E}[\Delta w^{(l)}H^{(w^{(l)})}\Delta w^{(l)}] $$

但是这个优化过程受限于Hessian矩阵的计算困难以及问题本身是一个NP-Hard问题。因此无法将这个作为最终的优化目标。

我们从Hessian矩阵计算的复杂性来分析

$$ \frac{\partial^2 L}{\partial W^{(l)}_{i,j}\partial W^{(l)}_{m,o}} = \frac{\partial}{\partial W_{m,o}^{(l)}}[\frac{\partial L}{\partial z_i^{(l)}}\cdot x_j^{(l-1)}] = \frac{\partial^2 L}{\partial z^{(l)}_i\partial z_m^{(l)}}\cdot x_j^{(l-1)}x_i^{(l-1)} $$

写作矩阵的形式

$$ H(w^{(l)}) = \mathbb{E}[x^{(l-1)} x^{(l-1)\top}⊗ \nabla^2_{z^{(l)}}L] $$

其中⊗为Kronecker积。由此看出Hessian矩阵的复杂性主要来于二阶导的求取，它需要通过网路的后续层反向传播二阶导数(见对角近似)。

为了解决这个问题，我们采用Hessian矩阵的对角近似，即将其近似为对角矩阵，记作$diag(\Delta^2_{z^{(l)}}L)$。

$$ H(w^{(l)}) = \mathbb{E}[x^{(l-1)} x^{(l-1)\top}⊗ diag(\nabla^2_{z^{(l)}}L)] $$

将这个近似带入优化方程中有：

$$ \arg \min_{\Delta W_{k,:}^{(l)}} \mathbb{E}[\nabla^2_{z^{(l)}}L_{k,k}\cdot \Delta W_{k,:}^{(l)}x^{(l-1)}x^{(l-1)\top}\Delta W_{k,:}^{(l)\top}]\\ =\arg\min_{\Delta W_{k,:}^{(l)}} \Delta W_{k,:}^{(l)}\mathbb{E}[x^{(l-1)}x^{(l-1)\top}]\Delta W_{k,:}^{(l)\top}\\ =\arg \min_{\Delta W_{k,:}^{(l)}} \mathbb{E}[(\Delta W_{k,:}^{(l)}x^{(l-1)})^2] $$

这里是认为$\nabla^2_{z^{(l)}}L_{i,i}$是一个与输入样本数据无关的常量结果。

由此我们推导出，我们只要最小化由于量化而在激活函数$z^{(l)}$中引入的均方误差。这与大部分量化的论文中的结论一致（如AdaQuant）

想要通过直接求解上面的优化方程仍然是一件困难的事情，因为它是NP-Hard的，因此作者将优化目标放宽为如下形式

$$ \arg \min_{V}||Wx-\hat Wx||^2_{F}+\lambda f_{reg}(V) $$

其中$||\cdot||^2_F$为F范数，$\hat W$为优化的软量化权重

$$ \hat W = s\cdot clip([\frac{W}{s}] + h(V),n,p) $$

$h(V_{i,j})$可以是任何在0和1之间取值的可微函数，$f_{reg}(V)$是一个可微正则项，用于鼓励$h(V_{i,j})$收敛到0或1.

但是这个方法存在一个缺陷，无法避免量化误差的不断积累且没有考虑到激活函数，所以做了进一步优化

$$ \arg \min_{V}||f_a(Wx)-f_a(\hat W\hat x)||_F^2 + \lambda f_{reg}(V) $$

其中$fa(\cdot)$为激活函数$\hat x$为当前层的反量化输入，x为当前层的浮点输入

[2021 ICML] AdaQuant: Accurate Post Training Quantization With Small Calibration Sets

Wed, 26 Aug 2026 08:00:00 +0800

总结：本篇文章的主要贡献在于提出了一个基于小数据集（校验集）的训练后量化方法AdaQuant，AdaQuant通过提出一个block/layer-wise的损失函数，通过在校验集上的训练学习量化参数(重点包括了一个最优的权重扰动，类似于AdaRound来避免四舍五入的不足),实现了减少量化的精度损失；提出了基于PI(整数规划)的bit精度分配方案，但是并没有解释精确损失的累加合理性；提出量化对BN融合造成的统计量偏移问题，并提出了PN(Para-Normalization)来解决这个问题。并在Bert-base网络上实现了不到1%的损失(4-8bit)

吐槽: 这篇文章作者(符号和表达)有点混乱，得多读几遍才能理解作者想表达什么.中间bit分配假设成立存疑

在一般的Post-training 量化中，我们的优化目标可以用下式表示:

$$ \hat\Delta = \arg \min_{\Delta}||X - Q_{\Delta}(X)||^2\\ Q_{\Delta}(x) = \Delta[\frac{X}{\Delta}] $$

其中，$Q(\cdot)$是量化方程。这种方法对所有的量化损失都是平等处罚的，但是事实上我们更应该对影响分类的量化损失进行更多的惩罚。量化感知训练可以缓解这个问题但是它存在计算开销大的问题。

基于此，作者提出了AdaQuant，其核心思想是采用一个block/layer-wise的优化误差函数:

$$ (\hat\Delta_{w}\hat\Delta_x,\hat V) = \arg\min_{\Delta_w,\Delta_x,V}||WX-Q_{\Delta_w}(W')Q_{\Delta_x}(X)||^2 $$

其中，$W' = W + V$,这里V是引入的一个连续的可学习的“补偿”张量。被量化的对象是进行补偿后的权重W’:$W_q = Q_{\Delta_w}(W') = Q_{\hat\Delta_w}(W+V)$.

这里其实有点类似于AdaRound的思想了，因为AdaRound主张的是在量化时直接采用四舍五入是一个不明智地选择，因此这里采用一个补偿张量V来做这个选择

由于量化参数的得到依赖的是上一层全精度的激活值输出作为输入，因此各个网络之间互不干扰，所以可以并行处理。

但是在实际的推理过程中，网络的输入是上一层量化后的激活值（见训练后静态量化），因此，作者提出了串行版本的AdaQuant，此时它的优化误差函数为

$$ (\hat\Delta_{w_l},\hat\Delta_{x_l},\hat V_l) = \arg\min_{\Delta_{w_l},\Delta_{x_l},V_l}||W_lX_l - Q_{\Delta_{w_l}}(W_l')\cdot Q_{\Delta_{x_l}}(X_l^q)||^2\\ X^q_l = \sigma(Q_{\Delta_{w_{l-1}}}(W'_{l-1})\cdot Q_{\Delta_{x_l}}(X_{l-1}^q) ) $$

其中，$\sigma(\cdot)$是激活函数

需要注意的是，串行版本的AdaQuant得在比特分配之后进行，这是因为它的优化依赖于上一层的输入。

为了在性能和精度之间做权衡，在量化时我们往往会给不同层的网络分配不同的bit精度。AdaQuant在此思想上，提出了采用整数规划(PI)的方式。

作者将网路的bit分配描述为这样一个问题：

给定L层的神经网络。对于每一层l，我们都有需要与前一层$X_{l-1}$的激活值相乘的权重$W_l$。令$W_l^k$和$X_{l-1}^n$表示$W_l,X_{l-1}$精度为k和n位的量化版本。对于每一层i，低位宽乘法$W_l^k X_{l-1}^k$会带来$\Delta L_l^{k,n}$的准确损失和$\Delta P_{l}^{k,n}$的性能提升。

作者假设了准确损失以及性能提升是满足可加性的。

这里存疑，因为l-1层的准确损失势必会影响l层的准确损失，不能简单做加法来描述整个网络的准确性退化

那么问题可以描述为，在不超过总网络退化$\Delta L$的前提下，最大化总性能提升。我们设示性函数$I_l^{k,n}$来表示第l层是否采用k，n位量化版本（1表示使用）。问题可以被符号表示为:

$$ \max \sum_{l = 0}^{L-1}\Delta P_l\\ Subject\ to\ \sum_{l}\Delta L_{l}\leq \Delta L\\ \forall l\in \{1,\dots,L\}:\Delta P_l = \sum_{k,n}I_l^{k,n}\cdot \Delta P_{l}^{k,n},\Delta L_l = \sum_{k,n}I_{l}^{k,n}\cdot L_{l}^{k,n}\\ \forall l\in \{1,\dots,L\}:\sum_{k,n}I_l^{k,n}=1,I_l^{k,n}\in\{0,1\} $$

BatchNormalization在部署时，通常会与前面的卷积/全连接层进行融合，这样可以减少推理时的计算量。(这是因为BN的线性变换乘$\gamma/\sqrt{\sigma^2+\epsilon}$，加$\beta-\gamma\mu/\sqrt{\sigma^2+\epsilon}$可以直接合并到权重和偏置中)

然而当网络量化后，会导致激活值的分布发生偏移，即统计量均值$\mu$和方差$\sigma^2$会偏离在全精度模型中应有的值。但是由于BN层已经被融合到了前面的层中，这个偏移无法被校准。基于此，作者采用了一个名为Para-Normalization(PN)的方法来更新BN的统计量，以补偿这种偏差。

具体而言，假设我们知道了原始的BN参数$\gamma_0,\beta_0$.然后我们初始化一个新的BN层，初始化$\mu,\sigma^2$以及BN参数$\gamma_r,\beta_r$以便重建BN，使其满足:

$$ BN_r(x) = \gamma_r \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta_r \approx x\\ \mu = \beta_r = \beta_0;\sigma^2 = \gamma_0^2;\gamma_r = \sqrt{\gamma_0^2+\epsilon} $$

然后在校准集上收集运行时均值和方差更新$\mu,\sigma^2$,需要注意的是$\beta_r,\gamma_r$是不变的，因为不进行反向传播。

收集到新的均值和方差后，我们重新将BN进行融合，收集到的统计数据可以按以下方式融合回当前的量化尺度：

$$ W_i' = W_i\frac{\gamma_r}{\sigma};b_i' = \frac{\gamma_r}{\sigma}(b_i - \mu) + \beta_r;\Delta_{w_i}' = \frac{\gamma_r}{\sigma}\Delta_{w_i} $$

浅谈低精度浮点数

Wed, 10 Jun 2026 20:25:01 +0800

Motivation

最近在搭建 VLM real quantization baseline。最初的实验环境选在 RTX PRO 6000 Blackwell 上，但在实际部署过程中发现，当前主流推理框架对 Blackwell 设备上传统整数低比特量化的支持还不够成熟，尤其是常见的 W4A8、W4A16 等 INT 量化配置，在 vLLM、SGLang、TensorRT-LLM 等推理栈中的可用性、kernel 覆盖和性能表现，都与 Hopper/Ada 等架构存在明显差异。

进一步调研后发现，Blackwell 架构的低精度推理路线正在发生变化：相比传统 INT4/INT8，硬件与软件生态开始更多地围绕低精度浮点格式展开，例如 MXFP4、NVFP4 等 micro-scaling FP4 格式。这类格式在保持 4-bit 存储与计算效率的同时，通过局部 scale、浮点指数以及更细粒度的动态范围适配，试图缓解 INT 量化中常见的离群值敏感和动态范围不足问题。

与此同时，近两年的低精度量化研究也逐渐从“如何把 INT4 做稳”转向“如何利用硬件原生支持的 FP4 格式”。例如，围绕 MXFP4/NVFP4 的推理量化、FP4 全量训练、FP4 误差分析以及格式感知的旋转/缩放方法，已经在 ICLR、ICML、NeurIPS 以及 arXiv 上陆续出现。因此，我准备写下这篇 blog，作为对低精度浮点量化格式的学习笔记，也作为后续在 Blackwell 设备上构建 VLM real quant baseline 的技术背景梳理。

Note

IEEE 754 标准浮点数

我们先来看看在内存中，浮点数是如何按照 IEEE 754 标准存储的。

本质上，浮点数可以看作是用科学计数法来表示一个数：

$$ V = (-1)^S \times M \times R^E $$

其中 S 表示符号位，M 表示尾数，R 表示基数，E 表示指数。对于计算机中的二进制浮点数而言，基数 R = 2。

以 IEEE 754 单精度浮点数，也就是 float32 为例，它一共占 32 bit，具体划分如下：

符号位 S：1 bit
指数字段 Exponent：8 bit
尾数字段 Fraction：23 bit

为了在有限的位宽内尽可能扩大表示范围并提高精度，IEEE 754 对指数和尾数做了特殊规定：

对于规格化浮点数，尾数的整数部分总是 1，因此这个最高位的 1 不需要显式存储，称为隐藏位。于是，尾数字段虽然只有 23 bit，但实际可以表示 24 bit 的有效数字。
指数字段采用偏置（bias）表示。对于 float32，指数字段占 8 bit，存储值范围为 0 ~ 255，偏置值为 127。对于规格化数，真实指数为：

$$ E = \text{Exponent} - 127 $$

需要注意的是，Exponent = 0 和 Exponent = 255 有特殊含义，分别用于表示 0、非规格化数、无穷大和 NaN。因此，规格化单精度浮点数的真实指数范围是 $-126 \sim 127$。

单精度浮点数在内存中的布局如下：

 bit index
 31 30 23 22 0
 │ │ │ │ │
 ▼ ▼ ▼ ▼ ▼
┌───────┬────────────────────────────┬──────────────────────────────────┐
│ S │ Exponent │ Fraction │
│ 1 bit │ 8 bits │ 23 bits │
└───────┴────────────────────────────┴──────────────────────────────────┘
 │ │ │
 │ │ └─ fraction bits of mantissa
 │ └─ biased exponent, bias = 127
 └─ sign bit: 0 = positive, 1 = negative

例如，我们想用单精度浮点数表示 25.125。

首先将它转换为二进制：

整数部分：25(D) = 11001(B)
小数部分：0.125(D) = 0.001(B)

因此：

25.125(D) = 11001.001(B)

将其写成二进制科学计数法：

11001.001(B) = 1.1001001(B) × 2^4

所以：

S = 0
M = 1.1001001
Fraction = 10010010000000000000000
Exponent = 4 + 127 = 131(D) = 10000011(B)

因此，25.125 的 IEEE 754 单精度表示为：

S | Exponent | Fraction
0 | 10000011 | 10010010000000000000000

完整的 32 bit 表示为：

01000001110010010000000000000000

对应的十六进制表示为：

0x41C90000

浮点数表示会带来一定的精度损失，主要原因是并非所有十进制小数都能被有限位宽的二进制小数精确表示。例如，0.2 转换为二进制后是一个无限循环小数：

0.2(D) = 0.0011001100110011...(B)

由于计算机的位宽有限，存储时只能保留有限位，并按照 IEEE 754 的规则进行舍入，因此会产生精度误差。

在 IEEE 754 单精度浮点数下，常用的有限数表示范围大约为：

-3.4 × 10^38 ~ 3.4 × 10^38

更准确地说，单精度浮点数能够表示的最大有限正数约为：

3.4028235 × 10^38

最小正规格化数约为：

1.17549435 × 10^-38

如果考虑非规格化数，还可以表示更接近 0 的数，但精度会进一步降低。

在如今的 LLM 训练推理过程中，更通用的浮点数格式是 FP16 与 BF16。它们与 FP32 的区别在于：

FP16 的指数位为 5，尾数位为 10
BF16 的指数位为 8，尾数位为 7

我们知道，在总位宽固定的情况下，尾数越多则指数越少，能够表示的范围越小，但相应精度越高。因此 FP16 比 BF16 精度更高，但表示范围更小。

正因如此，BF16 通常用在训练场景（更大的动态范围更利于反向传播中的梯度表示），而 FP16 通常用在推理场景（主要是因为目前很多主流推理框架对 FP16 的 kernel 优化做得更好）。

FP8

FP8 是 Hopper 架构设备开始支持的低精度浮点数，通常有两种表示形式：E4M3 和 E5M2（其中 E 表示指数位数，M 表示尾数位数）。

同样地，E4M3 精度更高、范围更小，而 E5M2 范围更大、精度更低。

在支持 FP8 的硬件（Hopper 架构显卡）上，已经支持 FP8 格式的矩阵乘法（E4M3×E4M3、E5M2×E5M2、E4M3×E5M2、E5M2×E4M3）。

我们接下来从精度和性能两个方面，对比分析 FP8 与 INT8 在量化上的表现。

首先分析二者在精度上的表现。

由浮点数的定义可知，浮点数可以看作是在相邻 2 的幂（$2^E$）之间均匀采样的 $2^M$ 个样本。例如在 E5M2 中，2（一次幂）和 4（二次幂）之间有 4 个样本，4 和 8（三次幂）之间也有 4 个样本；而在 E4M3 中，2 和 4 之间有 8 个样本。由此可知：

E5M2 的量化误差比 E4M3 更大；
FP8 表示的数值是非均匀的，浮点量化的误差会随着数值的增大而增大。

这正是 FP8 与 INT8 量化的最大差别：对于 FP8 而言，越靠近 0 分布越稠密，越远离 0 分布越稀疏。

接下来从数学角度分析 FP8 与 INT8 在量化上的表现。

对于 INT8，量化可以写作：

$$ \hat x = s \cdot \text{round}\left(\frac{x}{s}\right) $$

其中 $s$ 是 scale。如果使用对称 INT8，那么通常有：

$$ q \in [-127, 127] $$

假设一组待量化值的最大绝对值是 $\alpha$，那么 scale 大约是：

$$ s = \frac{\alpha}{127} $$

量化绝对误差满足：

$$ |x - \hat x| \leq \frac{s}{2} = \frac{\alpha}{254} $$

而相对误差为：

$$ \frac{|x - \hat x|}{|x|} \leq \frac{\alpha}{254\,|x|} $$

可以看到，无论从绝对误差还是相对误差来看，INT8 量化都更容易受到 outlier 的影响。

而对于 FP8，假设尾数有 $M$ bit，那么在区间 $[2^e, 2^{e+1})$ 内约有 $2^M$ 个采样点，间隔为：

$$ \Delta_e = \frac{2^e}{2^M} = 2^{e-M} $$

因此量化绝对误差大约满足：

$$ |x - \hat x| \leq \frac{1}{2}\,2^{e-M} $$

而相对误差近似满足：

$$ \frac{|x - \hat x|}{|x|} \leq 2^{-(M+1)} $$

也就是说，虽然 FP8 的绝对误差也会随着数值变大而变大，但相对误差大致稳定。

如果是 E4M3，相对误差不会超过 $6.25\%$；而 E5M2 则不会超过 $12.5\%$。

若令 INT8 的相对量化误差等于 E4M3：

$$ \frac{\alpha}{254\,|x|} = \frac{1}{16} $$

可以得到：

$$ |x| \approx 0.063\,\alpha $$

也就是说，当待量化值与最大值的差距不超过约 16 倍时，INT8 的量化效果往往更好；一旦超出这个范围（即数值相对 $\alpha$ 过小），FP8 则是更优的选择。

而我们知道，权重和激活值虽然都存在 outlier，但二者的数值范围并不相同：权重的 outlier 一般是普通 token 的数十倍，而激活值的 outlier 则不一定，甚至可以达到数百倍。

因此在实际部署时，一个理想的选择是权重用 INT 量化、激活值用 FP8 量化，也就是常说的 W4AFP8（在 TensorRT-LLM 中记作 W4A8）。目前的推理框架普遍推荐这种量化方式，例如 TensorRT-LLM 中：

QUANTIZATION METHODS	PERFORMANCE IMPROVEMENT (BATCH SIZE <= 4)	PERFORMANCE IMPROVEMENT (BATCH SIZE >= 16)	ACCURACY IMPACT	CALIBRATION TIME**
FP8 (W8A8)	Medium	Medium	Very Low	Minutes
Int8 SQ (W8A8)	Medium	Medium	Medium	Minutes
Int8 weight-only (W8A16)	Medium	Low	Low	Not Required
Int4 weight-only (W4A16)	High	Low	High	Not Required
Int4 AWQ (W4A16)	High	Low	Low	Tens of Minutes
Int4 GPTQ	High	Low	Low	Tens of Minutes
Int4-FP8 AWQ (W4A8)	High	Medium	Low	Tens of Minutes

可以看到，W4AFP8 在性能提升最大的同时，还能做到较低的精度影响。

接下来从性能角度比较 FP8 与 INT8。

在正式比较之前，需要先明确一点：低精度 Tensor Core 的核心，并不是把整个矩阵乘法都维持在 8 bit 精度下完成，而是采用 mixed-precision MMA 的方式。也就是说，输入矩阵以 FP8 或 INT8 这样的低精度格式参与乘法，但乘积通常会在更高精度的累加器中累加，最后再根据需要转换为目标输出精度。

这样设计的原因在于，GEMM 中每个输出元素本质上都是一个长度为 K 的点积，需要经历大量乘加操作。如果累加阶段仍然使用过低的精度，那么许多较小的增量会因为舍入而被直接吞掉，累加误差也会随着 K 的增大迅速放大。因此，低精度矩阵乘法的“低精度”主要体现在输入和乘法阶段，而不是简单地意味着整个计算链路都只有 8 bit。

从公开硬件规格来看，在支持 FP8 的 NVIDIA GPU 上，FP8 Tensor Core 和 INT8 Tensor Core 往往具有相同的标称峰值吞吐。这说明在产品级实现中，厂商可以通过专门的数据通路设计、流水线化、tile 级并行以及 mixed-precision 累加机制，使 FP8 和 INT8 在峰值计算能力上达到相同水平。

但这并不意味着 FP8 和 INT8 的底层硬件代价完全相同。INT8 的计算本质上是整数/定点乘加，数据格式规则，乘法和累加逻辑相对简单；而 FP8 虽然同样只有 8 bit，但它包含符号位、指数位和尾数位，计算时需要处理指数对齐、尾数运算、规格化、舍入以及高精度浮点累加等逻辑。因此，从单个 MAC 单元或数据通路的实现复杂度来看，FP8 通常比 INT8 更复杂。

所以更准确的理解是：理论上，FP8 运算单元的硬件实现成本高于 INT8；但在 Hopper、Blackwell 这类专门面向 AI 负载设计的 Tensor Core 中，NVIDIA 通过专门的硬件设计把 FP8 和 INT8 的峰值吞吐做到了相同水平。也就是说，FP8 的优势并不是“硬件更简单”，而是在相近吞吐下提供了更大的动态范围，从而更适合大模型中激活值波动大、outlier 明显的场景。（简单来说：放心大胆地用 FP8 吧，NVIDIA 帮你兜底了。）

FP4 / MXFP4 / NVFP4

在 4-bit 下，想要像 FP8 那样直接使用 FP4 是不现实的，因为 FP4 能表示的数值范围实在太小了。在 OCP 标准中，FP4 一般取 E2M1，它能表示的值大概是：

$$ 0,\ \pm 0.5,\ \pm 1,\ \pm 1.5,\ \pm 2,\ \pm 3,\ \pm 4,\ \pm 6 $$

也就是说，单单使用 FP4，能够表示的数值范围相当有限。因此实际应用时，通常不会只存储一个 FP4 值，还会额外存储一个 scale：

$$ \text{real num} = \text{Encoding of FP4} \times \text{scale} $$

scale 的不同对应着目前两种主流的 FP4 格式（Blackwell 确实支持 FP4 这种数值格式，但没有针对 scale 的专用硬件加速）。

我们先来看 MXFP4，它的 scale 格式是每 32 个值共享一个 E8M0 scale，也就是 2 的幂形式的 scale。

接下来按照与 FP8 相同的思路分析它的量化误差（取 block 内最大绝对值为 $A$）：

$$ \hat x_i = s_{mx} \cdot Q_{\text{E2M1}}\left(\frac{x_i}{s_{mx}}\right) $$

一般而言理想的 scale 是 $\frac{A}{6}$；然而在 MXFP4 中 scale 必须是 E8M0 格式，因此需要将其量化为最近的 2 的幂：

$$ s_{mx} = 2^{\text{round}\left(\log_2 \frac{A}{6}\right)} $$

在 TensorRT 中一般会向上取整。为了计算方便，我们定义：

$$ \rho_{mx} = \frac{s_{mx}}{A/6} $$

由于 E8M0 是 2 的幂，向上取整时有：

$$ 1 \leq \rho_{mx} < 2 $$

由于任意量化器的量化误差不会超过一个量化间隔，因此我们有：

$$ |e_i| \leq s_{mx} = \rho_{mx}\,\frac{A}{6} $$

而 $\rho_{mx}$ 与 block 内最大值和 2 的幂的对齐程度相关：当 $A$ 恰好为 2 的幂时 $\rho_{mx}$ 接近 1，误差较小；当 $A$ 略大于某个 2 的幂时 $\rho_{mx}$ 接近 2，误差较大。

例如 $A/6 = 1.99$ 与 $A/6 = 2.01$：前者会被向上取整为 2，此时 $\rho_{mx}$ 接近 1；后者会被向上取整为 4，此时 $\rho_{mx}$ 约为 2。

若对齐得好，则有：

$$ |e_i| \leq \frac{A}{6} $$

否则为：

$$ |e_i| < \frac{A}{3} $$

除了绝对误差外，我们再看看它的归零阈值（即低于该值会被量化为 0）：

$$ |x| < 0.25\,s_{mx} = \rho_{mx}\,\frac{A}{24} $$

于是有：

$$ \frac{A}{24} \leq \text{zero threshold} \leq \frac{A}{12} $$

大约是 $4.17\% A$ 到 $8.33\% A$。

接下来看看 NVFP4，它与 MXFP4 的区别在于：

block 更小：每 16 个元素共享一个 scale；
scale 更细：NVFP4 的 block scale 采用 FP8 E4M3，而不是 E8M0。

此外，NVFP4 还会额外使用一个全局的（per-tensor）FP32 scale 来避免 overflow。一般而言：

$$ s_{\text{global}} = \frac{\text{global}_{amax}}{6 \times 448}\\[4pt] s_{\text{block}} = \frac{\text{block}_{amax}/6}{s_{\text{global}}} $$

其中 6 是 E2M1 的最大可表示值，448 是 E4M3 的最大可表示值，二者相乘保证 block scale 落在 E4M3 的表示范围内。

同样地，我们可以把 NVFP4 写作：

$$ \hat x_i = s_{\text{global}} \cdot s_{\text{block}} \cdot Q_{\text{E2M1}}\left(\frac{x_i}{s_{\text{global}} \cdot s_{\text{block}}}\right) $$

令

$$ s_{nv} = s_{\text{global}} \cdot s_{\text{block}} $$

那么同样地有：

$$ \rho_{nv} = \frac{s_{nv}}{A/6} $$

又因为此时 scale 为 E4M3，其相邻正规格化数的比例最多约为：

$$ 1 + \frac{1}{8} = \frac{9}{8} $$

因此取向上取整时：

$$ 1 \leq \rho_{nv} \leq \frac{9}{8} $$

于是 NVFP4 的绝对误差上界约为：

$$ |e_i| \leq \rho_{nv}\,\frac{A}{6} \leq \frac{3}{16}A $$

约 $18.75\% A$，这明显好于 MXFP4 在未对齐时的最坏情况。

而 NVFP4 的归零阈值为：

$$ \text{zero threshold} \leq 4.69\% A $$

至于 INT4，可以通过与 INT8 类似的方式分析：对称量化下 $q \in [-7, 7]$，scale 为 $s = A/7$，因此绝对误差上界为 $|e_i| \leq s/2 = A/14$，归零阈值同样为 $A/14 \approx 7.14\% A$。

需要注意的是，INT4 的绝对误差上界（约 $7.14\% A$）看似比 NVFP4（约 $18.75\% A$）更小，但这是均匀分布在整个量程上的误差；而对于神经网络中占绝大多数的小数值，FP4 的误差与数值大小成比例，这些小值上的实际误差远小于 INT4 的均匀误差。再结合 FP4 对长尾异常值更友好的动态范围，可以得出结论：在 4-bit 量化中，面对神经网络这种“大量小值 + 长尾异常值”的分布，NVFP4 的综合表现最好。

而从性能上分析，虽然 MXFP4 由于采用 32 元素 block 和 E8M0 power-of-two scale，理论上具有更低的 scale 存储与缩放复杂度，但在真实的 LLM/VLM 推理中，性能并不只由单个 GEMM 的格式开销决定。NVFP4 通过 16 元素 micro-block 和 E4M3 scale 显著降低了量化误差，使更多层能够稳定进入 FP4 路径，减少混合精度 fallback 和额外补偿。因此，在当前 Blackwell + TensorRT-LLM 生态下，NVFP4 往往是更实际的高性能 FP4 选择；而 MXFP4 更适合在对精度不敏感、或有专门误差补偿方法的场景中追求更高的硬件效率。

此外，目前主流推理框架都对 NVFP4 做了专门优化，是性能—精度权衡下的最佳选择。

Reference

Triton

Tue, 14 Apr 2026 21:00:00 +0800

Motivation

CS336课程中在介绍实现FlashAttention2的时候介绍了Triton，我发现这是一个比CUDA更加方便的GPU编程语言。

另一方面许多大厂高性能计算岗位包括但不限于算子开发都逐渐趋向于使用Triton，因此目前看来学习Triton是一件从投资未来的角度来看性价比很高的事情

Introduction

Triton 是 OpenAI 开发的一种面向深度学习与高性能计算的编程语言及编译器，旨在简化 GPU 高性能算子的开发过程，并帮助开发者更容易地写出高效代码。

Triton的官方文档指出,其项目动机来自于基于分块算法的编程范式可以促进构建用于神经网络的高性能计算内核。与传统 CUDA 编程更强调线程级的执行方式不同，Triton 提供了更高层次的块级抽象，开发者可以直接围绕数据块来组织计算，而不必过多关注底层线程调度细节。

我们可以通过矩阵乘法的例子来更加直观的理解。

CUDA编程模型:

#pragma parallel
for(int m = 0; m < M; m++)
#pragma parallel
for(int n = 0; n < N; n++){
 float acc = 0;
 for(int k = 0; k < K; k++)
 acc += A[m, k] * B[k, n];

 C[m, n] = acc;
}

Triton编程模型:

#pragma parallel
for(int m = 0; m < M; m += MB)
#pragma parallel
for(int n = 0; n < N; n += NB){
 float acc[MB, NB] = 0;
 for(int k = 0; k < K; k += KB)
 acc += A[m:m+MB, k:k+KB]
 @ B[k:k+KB, n:n+NB];
 C[m:m+MB, n:n+NB] = acc;
}

GPU

在具体介绍Triton之前，我们先简单了解一下GPU相关的知识，以便减少后续专业术语带来的知识bias。

以下内容来自于本人于南开大学计算机体系结构期末突击的笔记。

相较于CPU而言，GPU的设计思路可以概括为三点:

简化流水线，增加核数：具体而言，对流水线进行瘦身，去掉乱序执行，分支预测等复杂逻辑，节省的空间用来增加大量的计算核心，从而能够同时处理大量数据
多个计算单元共用一条指令：GPU 擅长处理“很多数据做同一种计算”的场景，例如对向量中的每个元素做相同变换。为此，它通常让一组线程(这样的一组线程称为Warp)按照相同的程序逻辑并行执行，只是每个线程处理的数据不同。这种执行方式通常称为 SIMT。
驻留大量线程：GPU 会同时准备很多线程。当一部分线程因为访存而需要等待时，硬件可以立即切换去执行另一部分线程，从而减少计算单元空闲的时间。

我们接下来以数组加法来介绍GPU编程模型中的几个基本概念，这些概念有助于理解后续的并行执行，分块计算和访存优化。

假设我们要计算数组加法:

for i in range(1024):
	C[i] = A[i] + B[i]

那么一个GPU Kernel会以大量线程的形式启动。所有线程执行相同的Kernel代码，但每个线程会根据自己的编号处理不同位置的数据。在本例中，每个线程可以分别负责不同的i

Thread:即线程，这是程序员视角下最基本的并行单位。每个线程执行同一份kernel代码，但处理的数据不同
Block:多个Thread会组成一个Block。其意义在于协作，即同一个block内的线程可以通过共享内存交换数据，也可以使用同步原语协调执行不同 block 之间通常不能直接同步或共享局部数据。
Grid:一次Kernel启动产生的所有线程构成一个grid。grid只是一次kernel启动时的逻辑组织方式，本身并不对应某个具体的硬件结构
Warp:在硬件执行层面，线程通常不会以单个 thread 为单位独立调度，而是会进一步组成 warp。以 NVIDIA GPU 为例，一个 warp 通常包含 32 个线程。warp 内的线程会执行相同的指令流程，但处理各自的数据。这意味着，虽然我们在编程时使用的是 thread 和 block 这些抽象概念，但在硬件上，block 内的线程最终仍会被拆分为多个 warp 来执行。
SM:流处理器，GPU上的block会被分配到SM上执行。一个SM可以同时驻留多个warp，甚至是多个block。当某个warp因访存等高延迟操作暂时停顿时，硬件调度器可以快速切换到其他就绪的 warp 继续执行，从而隐藏延迟、提升整体吞吐。

Weighted Sum

我们接下来通过“加权求和”的例子来具体了解Triton的知识以及它是如何与Pytorch进行协同工作。

Forward Pass

“加权求和”的前向传播过程可以通过如下形式进行描述:

给定一个输入矩阵x,我们将它的每个元素乘以一个按列加权的向量w,然后对每一行求和，最终得到矩阵x与向量w的加权和。

用python代码描述则是:

def weighted_sum(x,weight):
	return (weight*x).sum(axis=-1)

而我们想把这个过程用Triton进行并行化，那么首先需要分析可并行性，我们可以发现矩阵与向量的乘积在最后求和得到结果之前各个元素的运算都是相互独立的，因此每个元素的计算都是可并行的。

因此在编写Triton Kernel时，我们会让每个program instance(Triton中一组执行同一段程序的线程块)负责x中某一块行的加权和，并把对应的标量输出写入输出张量。

例如，假设x的形状是[ROWS,D],weight的形状是[D]，我们会让每个program instance负责x中若干行的加权和的计算，得到每个块行的加权和结果，最后将结果均写入输出张量。

与直接把tensor作为参数不同，我们会在kernel中传入:

指向张量首元素的指针
每个tensor的stride，其含义是告诉我们如何沿着各个维度移动

我们可以利用这些 stride，结合 program ID，将工作分配给不同实例（例如，第 i 个实例处理 x 的第 i 个行块），从而加载当前实例所负责的 x 的一块数据。

在本例中，Triton的Forward Pass和PyTorch的主要区别在于:Triton需要显式地进行指针运算以及load/store操作。

我们来看具体的代码:

import triton
import triton.language as tl
@triton.jit 
def weighted_sum_fwd(
	x_ptr,weight_ptr,			#输入指针
	output_ptr,					#输出指针
	x_stride_row,x_stride_dim,	#stride告诉我们在tensor的每个轴上移动一个元素需要跳多远
	weight_stride_dim,			#通常是1
	output_stride_dim,			#通常是1
	ROWS,D,
	ROWS_TILE_SIZE,D_TILE_SIZE,	#tile形状必须在编译期已知
):
	#每个instance负责计算x中一个行块的加权和
	#`tl.program_id`用于查看当前运行的是哪个线程块
	row_tile_idx = tl.program_id(0)
	
	#block pointer 允许我们从一个N维内存区域中选取数据块
	#并在这个区域中移动所选的数据块
	#在使用block pointer时必须知道:
	#- tensor 第一个元素的指针
	#- tensor 的整体形状
	#- 每个维度的stride
	#- 起始block的N维坐标
	#- 每次load/store的block形状
	#- 内存中维度从主到次的顺序
	x_block_ptr = tl.make_block_ptr(
		x_ptr,
		shape=(ROWS,D),
		strides=(x_stride_row,x_stride_dim),
		offsets=(row_tile_idx * ROW_TILE_SIZE,0),
		block_shape=(ROW_TILE_SIZE,D_TILE_SIZE),
		order(1,0),
	)
	weight_block_ptr=tl.make_block_ptr(
		weight_ptr,
		shape=(D,),
		strides=(weight_stride_dim,),
		offsets=(0,),
		block_shape=(D_TILE_SIZE,),
		order(0,),
	)
	output_block_ptr=tl.make_block_ptr(
		output_ptr,
		shape=(ROWS,),
		strides=(output_stride_row,),
		offsets=(row_tile_idx*ROWS_TILE_SIZE,),
		block_shape=(ROWS_TILE_SIZE,),
		order(0,),
	)
	#初始化一个用于写入的buffer
	output = tl.zeros((ROWS_TILE_SIZE,),dtype=tl.float32)
	for i in range(tl.cdiv(D,D_TILE,SIZE)):
		#加载当前block pointer指向的数据块
		#因为ROW_TILE_SIZE可能不整除ROWS，
		#D_TILE_SIZE 也可能不能整除 D，
		#所以两个维度都需要进行边界检查
		row = tl.load(x_block_ptr,boundary_check=(0,1),padding_option="zero")
		weight = tl.load(weight_block_ptr,boundary_check=(0,),padding_option="zero")
		#计算当前行块的加权和
		output += tl.sum(row * weight[None,:],axis=1)
		#将指针移动到下一个tile,这里参数时(行偏移，列偏移)
		x_block_ptr = x_block_ptr.advance(0,D_TILE_SIZE)	#在最后一个维度上前进D_TILE_SIZE
		weight_block_ptr = weight_block_ptr.advance((D_TILE_SIZE,))
	# 将输出写入 output block pointer（每行一个标量）
 # 因为 ROWS_TILE_SIZE 可能不能整除 ROWS，所以需要边界检查
	tl.store(output_block_ptr,output,boundary_check(0,))

接下来我们将这个实现的Kernel封装到一个Pytorch AutoGrad函数中，使其能够与PyTorch生态协同工作(即接收 Tensor 输入、输出 Tensor，并在 backward pass 中和 autograd 引擎协作):

class WeightedSumFunc(torch.autograd.Function):
	@staticmethod
	def forward(ctx,x,weight):#这里ctx指上下文
		#将x和weight缓存起来，以便在backward中使用
		#backward时我们只会接收到输出张量的梯度
		# 需要利用这些缓存来计算 x 和 weight 的梯度
		D,output_dims = x.shape[-1],x.shape[:-1]
		input_shape = x.shape
		x = rearrange(x,"... d -> (...) d")
		ctx.save_for_backward(x,weight)
		
		assert len(weight.shape) == 1 and weight.shape[0] == D, "维度不匹配"
 assert x.is_cuda and weight.is_cuda, "期望输入是 CUDA tensor"
 assert x.is_contiguous(), "我们的指针运算默认 x 是连续的"
 
 ctx.D_TILE_SIZE = triton.next_power_of_2(D) # 16 #大致让dim维度循环16次
 ctx.ROWS_TILE_SIZE = 16						# 每个线程一次处理16个batch元素
 ctx.input_shape = input_shape
 # 需要初始化一个空的结果 tensor
 # 注意：这里的元素未必初始化为 0
 y = torch.empty(output_dims,device = x.device)
 
 # 以 1D launch grid 启动 kernel
 n_rows = y.numel()
 weighted_sum_fwd[(cdiv(n_rows,ctx.ROWS_TILE_SIZE),)](	#和CUDA的<<<>>>有点像，作用是定义grid大小
 	x,weight,
 	y,
 	x.stride(0),x.stride(1),
 	weight.stride(0),
 	y.stride(0),
 	ROWS = n_rows,D=D,
 	ROWS_TILE_SIZE=ctx.ROWS_TILE_SIZE, D_TILE_SIZE=ctx.D_TILE_SIZE,
 )
 return y.view(input_shape[:-1])

Backward Pass

由于我们定义了自己的kernel，因此也需要自己编写backward函数。

在forward pass中，我们拿到了层的输入，并计算其输出。而在backward pass中我们会拿到目标函数对该层输出的梯度，然后需要计算目标函数对每个输入的梯度。

在这个例子中，我们的操作有两个输入:

矩阵$x \in \mathbb{R}^{n\times h}$
权重向量$w \in \mathbb{R}^h$

记我们的操作为$f(x,w)$，其输出属于$\mathbb{R}^n$

若给定损失函数L对该层输出的梯度$\nabla_{f(x,w)}L$,则根据多元链式法则，可以得到关于x和w的梯度:

$$ (\nabla_xL)_{ij}=\sum_{k=1}^n\frac{\partial f(x,w)_k}{\partial x_{ij}}(\nabla_{f(x,w)}L)_k = w_j \cdot (\nabla_{f(x,w)}L)_i\\ (\nabla_w L)_{j} = \sum_{i=1}^n\frac{\partial f(x,w)_i}{\partial w_j}(\nabla_{f(x,w)}L)_i=\sum_{i=1}^nx_{ij}\cdot(\nabla_{f(x,w)}L)_i $$

这里给出一个很简单的backward计算公式。

为了得到关于x的backward结果，我们根据上式计算w和$\nabla f(x,w)$的外积。而为了得到关于w的backward结果(即$(\nabla_w L)_j$),我们必须把输入与对应输出的梯度逐行相乘并求和。

按照同样的思路，我们来实现这个kernel:

@triton.jit
def weighted_sum_backward(
	x_ptr,weight_ptr,
	grad_output_ptr,
	grad_x_ptr,partial_grad_weight_ptr,
	stride_xr,stride_xd,
	stride_wd,
	stride_gr,
	stride_gxr,stride_gxd,
	stride_gwb, stride_gwd,
 NUM_ROWS, D,
 ROWS_TILE_SIZE, D_TILE_SIZE,
):
	row_tile_idx = tl.program_id(0)
	n_row_tiles = tl.num_programs(0)
	#输入
	grad_output_block_ptr = tl.make_block_ptr(
		grad_output_ptr,
		shape=(NUM_ROWS,),strides=(stride_gr,),
		offsets=(row_tile_idx*ROW_TILE_SIZE,),
		block_shape=(ROWS_TILE_SIZE,),
		order=(0,),
	)
	x_block_ptr = tl.make_block_ptr(
		x_ptr,
		shape=(NUM_ROWS,D,),strides=(stride_xr,stride_xd),
		offsets=(row_tile_idx * ROW_TILE_SIZE,0),
		block_shape=(ROWS_TILE_SIZE,D_TILE_SIZE),
		order=(1,0),
	)
	weight_block_ptr = tl.make_block_ptr(
		weight_ptr,
		shape=(D,),strides=(stride_wd,),
		offsets=(0,),block_shape=(D_TILE_SIZE,),
		order=(0,),
	)
	grad_x_block_ptr = tl.make_block_ptr(
		grad_x_ptr,
		shape=(NUM_ROWS,D,),strides=(stride_gxr,stride_gxd),
		offsets=(row_tile_idx*ROW_TILE_SIZE,D_TILE_SIZE),
		block_shape=(ROW_TILE_SIZE,D_TILE_SIZE),
		order=(1,0),
	)
	partial_grad_weight_block_ptr = tl.make_block_ptr(
		partial_grad_weight_ptr,
		shape=(n_row_tiles,D,),strides=(stride_gwb,stride_gwd),
		offsets=(row_tile_idx,0),
		block_shape=(1,D_TILE_SIZE),
		order=(1,0),
	)
	for i in range(tl.cdiv(D,D_TILE_SIZE)):
		grad_output = tl.load(
			grad_output_block_ptr,
			boundary_check=(0,),
			padding_option="zero",
		)
		weight = tl.load(
			weight_block_ptr,
			boundary_check=(0,),
			padding_option="zero",
		)
		grad_x_row = grad_output[:,None] * weight[None,:]
		tl.store(grad_x_block_ptr,grad_x_row,boundary_chekc=(0,1))
		
		row = tl.load(
			x_block_ptr,
			boundary_check=(0,1),
			padding_option="zero",
		)
		grad_weight_row = tl.sum(row * grad_output[:,None],axis=0,keep_dims=True)
		tl.store(partial_grad_weight_block_ptr,grad_weight_row,boundary_check=(1,))
		
		x_block_ptr = x_block_ptr.advance((0,D_TILE_SIZE))
		weight_block_ptr = weight_block_ptr.advance((D_TILE_SIZE,))
		partial_grad_weight_block_ptr = parital_grad_weight_block_ptr.advance((0,D_TILE_SIZE))
		grad_x_block_ptr = grad_x_block_ptr.advance((0,D_TILE_SIZE))

每个 kernel instance 只负责处理 x 的一个行块，但我们现在需要沿着 x 的行方向求和。因此，在 backward kernel 内，我们不直接完成这个总和，而是假设 partial_grad_weight_ptr 是一个形状为 n_row_tiles × H 的矩阵，其中第一维表示来自每个行块的局部归约结果。

也就是说：

kernel 内部：只在当前行块内完成局部归约
kernel 外部：再用 torch.sum 把所有行块的结果加起来，得到最终的 $\nabla_w$

这样一来，autograd.Function 的最后一部分就相对简单了：

class WeightSumFunc(torch.autograd.Function):
	@staticmethod
	def forward(ctx,x,weight):
		#这里省略，前面定义过了
	@staticmethod
	def backward(ctx,grad_out):
		x,weight = ctx.saved_tensors
		ROWS_TILE_SIZE,D_TILE_SIZE = ctx.ROWS_TILE_SIZE,ctx.D_TILE_SIZE
		n_rows,D = x.shape()
		
		#我们的策略是：每个线程块先写入一个partial buffer，然后对这个buffer做归约得到最终梯度
		parital_grad_weight = torch.empty(
			(cdiv(n_rows,ROWS_TILE_SIZE),D),
			device=x.device,
			dtype=x.dtype
		)
		grad_x = torch.empty_like(x)
		weighted_sum_backward[(cdiv(n_rows,ROWS_TILE_SIZE),)](
			x,weight,
			grad_out,
			grad_x,partial_grad_weight,
			x.stride(0),x.stride(1),
			weight.stride(0),
			grad_out.stride(0),
			grad_x.stride(0), grad_x.stride(1),
 partial_grad_weight.stride(0), partial_grad_weight.stride(1),
 NUM_ROWS=n_rows, D=D,
 ROWS_TILE_SIZE=ROWS_TILE_SIZE, D_TILE_SIZE=D_TILE_SIZE,
		)
		grad_weight = partial_grad_weight.sum(axis=0)
		return grad_x,grad_weight

最后我们就可以得到一个函数，它的使用方法和torch.nn.functional里的函数很相似:

f_weightedsum = WeightedSumFunc.apply

我们可以对两个PyTorch tensor x和w调用这个函数，就会得到类似如下的结果:

tensor([ 90.8563, -93.6815, -80.8884, ..., 103.4840, -21.4634, -24.0192],
 device='cuda:0', grad_fn=<WeightedSumFuncBackward>)

请注意输出 tensor 上附带的 grad_fn —— 这表明 PyTorch 已经知道：当这个 tensor 出现在计算图中并需要做 backward pass 时，应该调用什么函数。

至此，我们就完成了这个 weighted sum 操作的 Triton 实现。

对于其它算子，我们可以按照同样的思路进行。实际上重点还是在于我们应该分析出哪些是可以并行的，program instance该如何划分。

简单来说，在设计Triton算子的时候我们需要考虑三个问题:

输出的自然分块是什么？例如在Weighted Sum中Y[0],Y[1]相互独立按行产生，那么Program Instance可以负责生成一个ROW_TILE_SIZE大小的行块的结果。
一个块的输出需要哪些输入？
这些输入能不能以连续，重复利用的方式被加载？

这三个问题虽然简单，但是在Triton编程中很重要，在实际问题中也很难立即给出标准答案，需要我们不断试错积累经验。

建议参考仓库(Assignment2-systems/Triton/WeightedSum.py)中的代码，不要直接复制粘贴，不然会出现一些奇奇怪怪的错误，这里的代码是参考了CS336 Assignment 2中的代码.

Triton Puzzles (Lite)

为了更深一步的学习Triton,我在网上查找到了一份通过做题来学习Triton的repo。

原始的repo是srush/Triton-Puzzles: Puzzles for learning Triton

但是有一个更好的实现版本是:SiriusNEO/Triton-Puzzles-Lite: Puzzles for learning Triton, play it with minimal environment configuration!

因此我们还是以这个仓库为题面来进行解答。

另外需要强调的一点是，正如Assignment2实验手册中提到的那样，目前大多教程没有使用更新的，更方便的block pointer抽象。因此在此基础上，在本blog中，我将以block pointer抽象来完成这些任务。

之所以说block pointer更加方便，是因为在Triton中，我们的Program Instance的操作对象是划分后的Tile，直观来看就是一个矩形区域。在使用block pointer之前，我们需要得到这个处理对象需要通过Program ID以及SIZE用公式去计算。例如我们想知道处理一个x[ROWS,D]的tile，我们需要写:

row_offsets = row_tile_idx * ROWS_TILE_SIZE + tl.arrange(0,ROWS_TILE_SIZE)
col_offsets = tl.arrange(0,D_TILE_SIZE)
#计算出当前块负责哪些行和哪些列
x_ptrs = x_ptr + row_offsets[:,None] * x_stride_row + col_offsets[None,:] * x_stride_dim
#然后把他们拼接为二维地址
mask = (row_offsets[:, None] < ROWS) & (col_offsets[None, :] < D)
#然后再load
x = tl.load(x_ptrs, mask=mask, other=0.0)

而使用block pointer的话，我们就可以写:

x_block_ptr = tl.make_block_ptr(...)
x = tl.load(x_block,boundary_check=(0,1),padding="zero")

这样做就省去了大量的地址计算，一方面更加便捷，另一方面代码的可读性也大大提升了。接下来我们看看具体的问题以及如何解决吧？

Puzzles 1/2 Constant Add

问题描述:

对向量中的每个元素都加上常数10.使用一个程序ID轴。

问题解答：

我们按照上面Weighted Sum中make_block_ptr的方式定义好x_block_ptr以及z_block_ptr，然后tl.load出x，进行运算得到z最后进行tl.store即可。

由于N0不一定能被B0整除，因此需要进行boundary check

@triton.jit
def add_kernel(x_ptr, z_ptr, N0, B0: tl.constexpr):
 row_tile_idx = tl.program_id(0)
 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(row_tile_idx * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(row_tile_idx * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 x = tl.load(x_block_ptr,boundary_check=(0,),padding_option='zero')
 z = x + 10.0
 tl.store(z_block_ptr, z, boundary_check=(0,))
 # Finish me!
 return

Puzzles 3/4 Outer Vector Add

问题描述:

把两个向量按广播方式相加，得到一个二维结果矩阵z。

$$ z_{j,i} = x_i+y_j \quad i=1,\dots,B_0,j=1,\dots,B_1 $$

问题解答:

这个可以利用Python中的broadcast机制:x[None,:]+y[:,None]。在计算时会先将x的形状变为(1,N0),y变为(N1,1)然后相加的时候触发广播机制，两者的 shape 均广播到 (N1, N0)。

在Triton中也有同样的机制，我们可以利用这个得到最终的结果。

考虑到N0,N1,B0,B1的差异，我们的grid相应地也要变成二维的，第一个维度是走x的tile，第二个维度走y的tile

@triton.jit
def add_vec_block_kernel(
 x_ptr, y_ptr, z_ptr, N0, N1, B0: tl.constexpr, B1: tl.constexpr
):
 x_tile_idx = tl.program_id(0)
 y_tile_idx = tl.program_id(1)

 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(x_tile_idx * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 y_block_ptr = tl.make_block_ptr(
 base=y_ptr,
 shape=(N1,),
 strides=(1,),
 offsets=(y_tile_idx * B1,),
 block_shape=(B1,),
 order=(0,)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N1, N0),
 strides=(N0, 1),
 offsets=(y_tile_idx * B1, x_tile_idx * B0),
 block_shape=(B1, B0),
 order=(0, 1)
 )
 x = tl.load(x_block_ptr,boundary_check=(0,),padding_option='zero')
 y = tl.load(y_block_ptr,boundary_check=(0,),padding_option='zero')
 z = x[None, :] + y[:, None]
 tl.store(z_block_ptr, z, boundary_check=(0,1))
 # Finish me!
 return

Puzzles 5 Fused Outer Multiplication

问题描述:

将行向量x与列向量y做外积，然后对结果矩阵的每个元素应用ReLU.

$$ z_{j,i} = \max(0,x_i\cdot y_j),\quad i = 1,\dots,N_0,j=1,\dots,N_1 $$

问题解答：

这一问实际上是与Puzzle 3/4类似的，只有符号从加法变成了乘法。对于ReLU操作直接将z与0取max就行:z = tl.maximum(0,z)

Puzzles 6 Fused Outer Multiplication - Backwards

问题描述：

对如下函数操作进行反向传播:矩阵x与向量y按行相乘，再经过ReLU。

$$ f(x,y) = \text{relu}(x_{j,i}\times y_j),\quad i = 1 \dots N_0,j = 1 \dots N_1\\ dx_{j,i} = f'_x(x,y)_{j,i}\times dz_{j,i} = 1(x_{j,i}\cdot y_j >0)\cdot y_j \cdot dz_{j,i} $$

其中$1(\cdots)$表示指示函数，括号内条件满足时为1，否则为0

问题解答：

按照公式实现即可，没有较为复杂的逻辑，因此不贴代码了。

Puzzles 7 Long Sum

问题描述:

给定一个二维张量x，对其每一行进行求和。

问题解答:

这是一个二维的张量，我们在分块计算的时候虽然直观的想法是每行进行并行计算，但是考虑到每行有T个元素，而我们设计的块的大小是$B_0\times B_1$，而B1远小于B0,因此我们需要沿着第二维分块累加，这也是为什么题目提醒我们需要使用for循环。

当我们使用block pointer时，我们可以直接用tl.advance(pointer,(stride_dim1,...))的方式移动指针。

@triton.jit
def sum_kernel(x_ptr, z_ptr, N0, N1, T, B0: tl.constexpr, B1: tl.constexpr):
 # Finish me!
 tile_id = tl.program_id(0)
 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0, T),
 strides=(T, 1),
 offsets=(tile_id * B0, 0),
 block_shape=(B0, B1),
 order=(1, 0)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(tile_id * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 z = tl.zeros((B0,), dtype=tl.float32)
 for i in range(0, T, B1):
 x = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option='zero')
 z += tl.sum(x,1)
 x_block_ptr = tl.advance(x_block_ptr,(0,B1))
 tl.store(z_block_ptr, z, boundary_check=(0,))
 return

Puzzles 8 Long Softmax

问题描述:

对一批logits做Softmax。需要保证数值稳定，即减去最大值后求Softmax

$$ z_{i,j} = \text{softmax}(x_{i,1},\dots,x_{i,T}) \quad i = 1,\dots,N_0 $$

另外需要注意在Triton中建议不要直接用exp而是用exp2.$\text{exp}(x) = 2^{\log_2(e)x}$

问题解答:

我们先来做没有优化的三个for的版本。实际上我们可以将任务做一个拆分。因为我们是分行块处理的，但是tile的大小[B0,B1]其中$B1 < T$，因此我们需要for循环来得到:

每行的max
每行的sum
每行的结果

因此需要3个for循环。因为后续我们将通过Online Softmax算法来实现更加高效实用的版本，因此这里就不贴出具体的代码了，思路和Puzzle 7类似。

Online Softmax的思路其实比较简单，就是第一个for和第二个for可以合并!

考虑两个tile，第一个tile的max我们记为$m_1$,那么此时局部求和的结果为:

$$ sum = \sum_{i = 0}^{B_0-1} \exp(x_i - m_1) $$

第二个tile的max记为$m_2(m_2 > m_1)$,我们记$a = m_1 - m_2$,那么此时我们可以更新局部求和结果为:

$$ sum' = sum \times \exp(a) + \sum_{B_0}^{2B_0-1}\exp(x_i - m_2) $$

由此我们可以通过一个for循环，便得到了求最终结果需要的max以及sum。

@triton.jit
def softmax_kernel(
 x_ptr, z_ptr, N0, N1, T, B0: tl.constexpr, B1: tl.constexpr
):
 """3 loops ver."""
 block_id_i = tl.program_id(0)
 log2_e = 1.44269504
 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0, T),
 strides=(T, 1),
 offsets=(block_id_i * B0, 0),
 block_shape=(B0, B1),
 order=(1, 0)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0, T),
 strides=(T, 1),
 offsets=(block_id_i * B0, 0),
 block_shape=(B0, B1),
 order=(1, 0)
 )
 z = tl.zeros((B0, B1), dtype=tl.float32)
 x_max = tl.full((B0,), float("-inf"), dtype=tl.float32)
 sum = tl.zeros((B0,), dtype=tl.float32)
 x_block_ptr1 = x_block_ptr
 for i in range(0,T,B1):
 x = tl.load(x_block_ptr1, boundary_check=(0, 1), padding_option='zero')
 new_x_max = tl.maximum(x_max, tl.max(x,1))
 if i == 0:
 x_max = new_x_max
 sum = tl.sum(tl.exp2(log2_e * (x - x_max[:, None])),1)
 else:
 scale = tl.exp2(log2_e * (x_max - new_x_max))
 sum = sum * scale + tl.sum(tl.exp2(log2_e * (x - new_x_max[:, None])),1)
 x_max = new_x_max
 x_block_ptr1 = tl.advance(x_block_ptr1,(0,B1))
 for i in range(0,T,B1):
 x = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option='zero')
 x = x - x_max[:, None]
 x_exp = tl.exp2(log2_e * x)
 z = x_exp / sum[:, None]
 tl.store(z_block_ptr, z, boundary_check=(0, 1))
 x_block_ptr = tl.advance(x_block_ptr,(0,B1))
 z_block_ptr = tl.advance(z_block_ptr,(0,B1))
 # Finish me!
 return

Puzzles 9 Simple FlashAttention

问题描述:

不使用 program。大小 B0 表示在总共 N0 个 q 中，本次要处理的一批。序列长度是 T。每次处理 B1 < T 个元素（k、v），其中 B1 是某个块大小。

$$ z_i = \sum_{j=1}^T \text{softmax}(q_ik_1,\dots,q_ik_T)_j v_j $$

这个问题可以通过类似于Online Softmax的思路来解决。

问题解答:(为便于理解，推荐读者用纸笔画一下)

考虑q的一个大小为[B0=2]的tile中元素为[q_0,q_1],k和v的一个大小为[B1=2]的tile中的元素为[k_1,k_2],[v_1,v_2]，那么对应的注意力矩阵元素为一个大小为[B0,B1]的二维矩阵，其中元素为[s_11,s_12,s_21,s_22]。

仅考虑结果中的第一个元素$z_1$,此时其累加和为:$z_1 = \text{softmax}(s_{11})v_1+\text{softmax}(s_{12})v_2 = \frac{\exp(s_{11}-m_1)}{l_1}v_1+\frac{\exp(s_12-m_1)}{l_1}v_2$,其中$m_1$为此时的局部最大值，$l_1$为此时的累计指数和.

计算完毕后，指针沿着axis=1的维度进行移动(步长为B1)，此时k和v的一个大小为[B1=2]的tile中的元素为[k_3,k_4],[v_3,v_4]，那么对应的注意力矩阵元素为一个大小为[B0,B1]的二维矩阵，其中元素为[s_13,s_14,s_23,s_24]。

同样考虑结果中的第一个元素$z_1$,此时最大值更新为$m_2$,指数累计和应更新为$l_2 = l_1\times(\text{scale}=\exp(m_1-m_2))+\sum_{j=3}^4 \exp(s_{1j}-m_2)$.

那么我们分两部分考虑此时的结果$z_1$,一部分是新的注意力权重$s_{13},s_{14}$与$v_3,v_4$的加权和，这部分直接和之前一样相加就行，而另一部分则为之前计算出来的贡献，这部分由于m和l的更新，也要对应更新。

更新可以从分子和分母考虑。分子上和Online Softmax类似，需要乘以$\exp(m_1-m_2)$,而分母则需要进行替换为$l_2$,因此最终我们可以更新局部加权和$z_1$为:

$$ z_1' = z_1\times scale\times \frac{l_1}{l_2} + \sum_{j=3}^4\frac{\exp(s_{1j}-m_2)}{l_2}v_j $$

由此不断迭代，便能得到最终结果。

@triton.jit
def flashatt_kernel(
 q_ptr, k_ptr, v_ptr, z_ptr, N0, T, B0: tl.constexpr, B1: tl.constexpr
):
 block_id_i = tl.program_id(0)
 log2_e = 1.44269504
 myexp = lambda x: tl.exp2(log2_e * x)
 q_block_ptr = tl.make_block_ptr(
 base=q_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(block_id_i * B0,),
 block_shape=(B0,),
 order=(0,)
 )
 k_block_ptr = tl.make_block_ptr(
 base=k_ptr,
 shape=(T,),
 strides=(1,),
 offsets=(0,),
 block_shape=(B1,),
 order=(0,)
 )
 v_block_ptr = tl.make_block_ptr(
 base=v_ptr,
 shape=(T,),
 strides=(1,),
 offsets=(0,),
 block_shape=(B1,),
 order=(0,)
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0,),
 strides=(1,),
 offsets=(block_id_i * B0,),
 block_shape=(B0,),
 order=(0,)
 )

 qk_max = tl.full((B0,), float("-inf"), dtype=tl.float32)
 qk_sum = tl.zeros((B0,), dtype=tl.float32)
 o = tl.zeros((B0,), dtype=tl.float32)
 for i in range(0, T, B1):
 q = tl.load(q_block_ptr, boundary_check=(0,), padding_option='zero')
 k = tl.load(k_block_ptr, boundary_check=(0,), padding_option='zero')
 v = tl.load(v_block_ptr, boundary_check=(0,), padding_option='zero')
 qk = q[:, None] * k[None, :]
 new_qk_max = tl.maximum(qk_max,tl.max(qk, 1))
 if i == 0:
 qk_max = new_qk_max
 qk_exp = myexp(qk - qk_max[:,None])
 qk_sum = tl.sum(qk_exp, 1)
 o = tl.sum(qk_exp * v[None, :] / qk_sum[:,None], 1)
 else:
 qk_exp = myexp(qk - new_qk_max[:,None])
 scale = myexp(qk_max - new_qk_max)
 scale1 = scale * qk_sum
 qk_sum = qk_sum * scale + tl.sum(qk_exp, 1)
 o = o * scale1 / qk_sum + tl.sum(qk_exp * v[None, :] / qk_sum[:,None], 1)
 qk_max = new_qk_max
 k_block_ptr = tl.advance(k_block_ptr,(B1,))
 v_block_ptr = tl.advance(v_block_ptr,(B1,))
 tl.store(z_block_ptr, o, boundary_check=(0,))
 # Finish me!
 return

Puzzles 10 Two Dimensional Convolution

问题描述:

实现一个带批处理的二维卷积。使用一个program id 轴。块大小为B0表示在N0中一次处理多少个batch。图像x的大小为H x W，并且只有一个通道；卷积核k的大小为kH x kW

$$ z_{i, j, l} = \sum_{oj, ol}^{j+oj\le H, l+ol\le W} k_{oj,ol} \times x_{i,j + oj, l + ol} \text{ for } i = 1\ldots N_0 \text{ for } j = 1\ldots H \text{ for } l = 1\ldots W $$

问题解决：(为了方便理解，最好对照题面的图来理解)

我们考虑如何分块，题目中要求的是按照批维度进行分块拆分并行。那么块的第一维就是B0,注意到卷积运算每次只会用到图像矩阵中KW*KH大小的子矩阵，因此块最合适的大小就是[B0,KH,KW]。这个块与卷积核进行卷积运算，得到一个标量，因此结果的块的大小就是[B0,1,1]，那么这样就能写出block pointer的定义。

接下来模拟二维卷积运算就行了，根据相关知识，我们知道卷积运算就是卷积核在图像矩阵上进行平移(这里默认步长为1).那么我们用两个for循环来模拟这个平移计算的流程就行，然后在内层for进行一次后沿着dim=1平移一格，在外层for进行一次后沿dim=0平移一格。

@triton.jit
def conv2d_kernel(
 x_ptr, k_ptr, z_ptr, N0, H, W, KH: tl.constexpr, KW: tl.constexpr, B0: tl.constexpr
):
 block_id_i = tl.program_id(0)
 k_block_ptr = tl.make_block_ptr(
 base=k_ptr,
 shape=(KH, KW),
 strides=(KW, 1),
 offsets=(0, 0),
 block_shape=(KH, KW),
 order=(1, 0),
 )
 k = tl.load(k_block_ptr, boundary_check=(0, 1), padding_option="zero")

 x_row_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N0, H, W),
 strides=(H * W, W, 1),
 offsets=(block_id_i * B0, 0, 0),
 block_shape=(B0, KH, KW),
 order=(2, 1, 0),
 )
 z_row_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N0, H, W),
 strides=(H * W, W, 1),
 offsets=(block_id_i * B0, 0, 0),
 block_shape=(B0, 1, 1),
 order=(2, 1, 0),
 )

 for _ in range(H):
 x_block_ptr = x_row_block_ptr
 z_block_ptr = z_row_block_ptr
 for _ in range(W):
 x = tl.load(x_block_ptr, boundary_check=(0, 1, 2), padding_option="zero")
 z = (x * k[None, :, :]).sum(1).sum(1)
 tl.store(z_block_ptr, z[:, None, None], boundary_check=(0, 1, 2))
 x_block_ptr = tl.advance(x_block_ptr, (0, 0, 1))
 z_block_ptr = tl.advance(z_block_ptr, (0, 0, 1))
 x_row_block_ptr = tl.advance(x_row_block_ptr, (0, 1, 0))
 z_row_block_ptr = tl.advance(z_row_block_ptr, (0, 1, 0))

 return

Puzzles 11 Matrix Multiplication

问题描述：

使用三条program id轴。块大小B2表示在N2中要处理的Batch数量。块大小B0表示在N0中要处理的行数，块大小B1表示在N1中要处理的列数。

中间维度的大小为MID

$$ z_{i, j, k} = \sum_{l} x_{i,j, l} \times y_{i, l, k}\quad \text{for } i = 1\ldots N_2,\ j = 1\ldots N_0,\ k = 1\ldots N_1 $$

Hint:可以使用tl.dot，它可以计算一个更小规模的矩阵乘法

问题解答:

我们知道两个矩阵的形状分别为shape(x) = [N2,N0,MID],shape(y) = [N2,MID,N1].我们分块的话，显然需要在MID维度上进行拆分，那么两个矩阵对应的块大小为:[B2,B0,B_MID],[B2,B_MID,B1],然后将这些小矩阵的乘积进行累加，便能得到最终的结果。

@triton.jit
def dot_kernel(
 x_ptr,
 y_ptr,
 z_ptr,
 N0,
 N1,
 N2,
 MID,
 B0: tl.constexpr,
 B1: tl.constexpr,
 B2: tl.constexpr,
 B_MID: tl.constexpr,
):
 block_id_j = tl.program_id(0)
 block_id_k = tl.program_id(1)
 block_id_i = tl.program_id(2)
 # Finish me!
 x_block_ptr = tl.make_block_ptr(
 base=x_ptr,
 shape=(N2,N0,MID),
 strides=(N0 * MID,MID,1),
 offsets=(block_id_i * B2, block_id_j * B0, 0),
 block_shape=(B2, B0, B_MID),
 order=(0, 1, 2),
 )
 y_block_ptr = tl.make_block_ptr(
 base=y_ptr,
 shape=(N2,MID,N1),
 strides=(MID * N1,N1,1),
 offsets=(block_id_i * B2, 0, block_id_k * B1),
 block_shape=(B2, B_MID, B1),
 order=(0, 1, 2),
 )
 z_block_ptr = tl.make_block_ptr(
 base=z_ptr,
 shape=(N2,N0,N1),
 strides=(N0 * N1,N1,1),
 offsets=(block_id_i * B2, block_id_j * B0, block_id_k * B1),
 block_shape=(B2, B0, B1),
 order=(0, 1, 2),
 )
 z = tl.zeros((B2, B0, B1), dtype=tl.float32)
 for i in range(0, MID, B_MID):
 x = tl.load(x_block_ptr, boundary_check=(0,1,2),padding_option='zero')
 y = tl.load(y_block_ptr, boundary_check=(0,1,2),padding_option='zero')
 z += tl.dot(x,y)
 x_block_ptr = tl.advance(x_block_ptr,(0,0,B_MID))
 y_block_ptr = tl.advance(y_block_ptr,(0,B_MID,0))
 tl.store(z_block_ptr, z, boundary_check=(0,1,2))
 return

Puzzles 12 Quantized Matrix Mult

问题描述:

我们将对一个矩阵乘法进行量化压缩，具体而言是将权重矩阵以更低精度存储，并额外配合一个偏移项和缩放项。

在本问题中，我们的weight将以4bit存储。一个32位整数中可以存放FPINT个这样的值。此外，对每个连续group个权重，我们还会额外存储:

1个scale浮点值
1个shift的4bit值

这些scale和shift是按weight的列来存储的。而activation则单独以普通浮点数格式存储。

$$ z_{j, k} = \sum_{l} sc_{j, \frac{l}{g}} (w_{j, l} - sh_{j, \frac{l}{g}}) \times y_{l, k} \quad \text{for } j = 1\ldots N_0,\ k = 1\ldots N_1 $$

其中g表示分组大小。

问题解答:

本题的主要难点，不在于计算，而是张量的形状处理。我们先来梳理以下完整的操作流程。

首先我们还是按照正常的矩阵乘法在MID维度上进行拆分，然后由于权重是以INT8格式存储的，我们还需要将其提取出来(利用位运算即可)。然后每个GROUP我们都会存储一个浮点型的scale和INT4类型的shift.

为了方便后续理解，我们看看各个分块的大小(仅看输入):

激活值,这个最好理解，它是以浮点型存储的，然后因为是在MID维度做了拆分，因此块大小为[B_MID,B_1],需要注意访问的顺序是列优先的，因此order = (0,1)
权重矩阵，这个需要考虑到INT32存储INT8带来的形状差异，我们记一个INT32可以存储FPINT个INT4,那么权重矩阵实际的大小为[N0,MID//FPINT],那么对应的块大小就是[B0,B_MID//FPINT]
scale,它是浮点型存储的，并且是一个GROUP(在MID维度分组)一个,那么其形状就是[N0,GROUP],对应块大小为[B0,GROUP]
offset,它与Scale类似，是一个GROUP一个，但是它是以INT4存储的，因此也需要考虑FPINT的事情，它的实际大小就是:[B0,GROUP//FPINT]

分完块之后就需要把用int4的块提取为int32格式的，方便后续计算。

然后需要考虑形状的问题了，我们提取出来的INT32格式的Weight的张量形状为[B0,B_MID // FPINT,FPINT],但是我们的scale的形状为[B0,GROUP]，我们无法将其直接与Weight进行运算得到dequant的权重，因此我们需要将其reshape为[B0,GROUP,B_MID//GROUP].offsets也同理，需要从[B0, GROUP // FPINT, FPINT]处理为[B0,GROUP]。之后借助广播进行dequant后，再进行矩阵乘法。

@triton.jit
def quant_dot_kernel(
 scale_ptr,
 offset_ptr,
 weight_ptr,
 activation_ptr,
 z_ptr,
 N0,
 N1,
 MID,
 B0: tl.constexpr,
 B1: tl.constexpr,
 B_MID: tl.constexpr,
):
 block_id_j = tl.program_id(0)
 block_id_k = tl.program_id(1)

 def extract(x):
 over = tl.arange(0,8) * 4
 mask = 2**4 - 1
 return (x[:,:, None] >> over) & mask

 activation_block_ptr = tl.make_block_ptr(
 base = activation_ptr,
 shape = (MID,N1),
 strides = (N1,1),
 offsets = (0, block_id_k * B1),
 block_shape = (B_MID, B1),
 order = (0,1)
 )
 weight_block_ptr = tl.make_block_ptr(
 base = weight_ptr,
 shape = (N0,MID // FPINT),
 strides = (MID // FPINT,1),
 offsets = (block_id_j * B0, 0),
 block_shape = (B0, B_MID // FPINT),
 order = (1,0)
 )
 scale_block_ptr = tl.make_block_ptr(
 base = scale_ptr,
 shape = (N0, GROUP),
 strides = (GROUP,1),
 offsets = (block_id_j * B0, 0),
 block_shape = (B0, GROUP),
 order = (1,0)
 )
 offset_block_ptr = tl.make_block_ptr(
 base = offset_ptr,
 shape = (N0,GROUP // FPINT),
 strides = (GROUP // FPINT,1),
 offsets = (block_id_j * B0, 0),
 block_shape = (B0, GROUP // FPINT),
 order = (1,0)
 )
 z_block_ptr = tl.make_block_ptr(
 base = z_ptr,
 shape = (N0,N1),
 strides = (N1,1),
 offsets = (block_id_j * B0, block_id_k * B1),
 block_shape = (B0, B1),
 order = (1,0)
 )
 z = tl.zeros((B0, B1), dtype=tl.float32)
 for i in range(0, MID, B_MID):
 scale_fp32 = tl.load(scale_block_ptr, boundary_check=(0,1), padding_option='zero')#[B0,GROUP]
 
 offset_int32 = extract(tl.load(offset_block_ptr, boundary_check=(0,1), padding_option='zero'))#[B0, GROUP // FPINT, FPINT]
 offset_int32 = offset_int32.reshape(B0, GROUP)#[B0, GROUP]
 
 weight_int32 = extract(tl.load(weight_block_ptr, boundary_check=(0,1), padding_option='zero'))#[B0, B_MID // FPINT, FPINT]
 weight_int32 = weight_int32.reshape(B0, B_MID)#[B0, B_MID]
 weight_int32 = weight_int32.reshape(B0,GROUP,B_MID // GROUP)#[B0, GROUP, B_MID // GROUP]

 activation_fp32 = tl.load(activation_block_ptr, boundary_check=(0,1), padding_option='zero')
 
 weight_fp32 = scale_fp32[:,:,None] * (weight_int32 - offset_int32[:,:,None])
 weight_fp32 = weight_fp32.reshape(B0,B_MID)

 z += tl.dot(weight_fp32, activation_fp32)
 scale_block_ptr = tl.advance(scale_block_ptr,(0, GROUP))
 offset_block_ptr = tl.advance(offset_block_ptr,(0, GROUP // FPINT))
 weight_block_ptr = tl.advance(weight_block_ptr,(0, B_MID // FPINT))
 activation_block_ptr = tl.advance(activation_block_ptr,(B_MID, 0))
 tl.store(z_block_ptr, z, boundary_check=(0,1))

 # Finish me!
 return

Reference

namoe的解答
Lite作者的解答
先进编译实验室的讲解
CS336 Assignment2-systems的实验指导手册
CSE559M Note

Transformer中的位置编码的探索与思考

Tue, 14 Apr 2026 20:48:13 +0800

Motivation

最近在完成CS336的Assignment1，从头搭建一个大模型，其中实现使用的是主流的旋转位置编码(RoPE)。虽然对照指导书完成了相关实验，但是对位置编码的选择以及原理存在一定的知识缺漏，因此写本博客以加强理解。

Note

我们在下面的数学推到中使用列向量推导

为什么Transformer需要位置编码?

设输入序列长度为L，模型的维度为d，那么我们输入的序列第i个token的表示为$x_i \in \mathbb{R}^{d}$,将整个序列拼接起来我们有:

$$ X = [x_1,x_2,\dots,x_L]\in \mathbb{R}^{d\times L} $$

那么Transformer Block中的MHA(由于最终结果是每一个Head通过相同的过程得到的不同结果进行concat，因此我们这里只考虑一个Head)在计算Attention Score时首先通过线性映射得到:

$$ Q = W_Q X \quad K = W_KX\quad V = W_V X $$

其中:

$$ W_Q,W_K,W_V \in \mathbb{R}^{d\times d} $$

因此:

$$ Q,K,V \in \mathbb{R}^{d\times L} $$

接下来会计算第i个query和第j个key的相似度:

$$ s_{ij} = \frac{q_i^\top k_j}{\sqrt{d}} $$

写成矩阵形式就是:

$$ S = \frac{Q^\top K}{\sqrt{d}}\in \mathbb{R}^{L\times L} $$

然后进行Softmax操作(沿key维度):

$$ A = \text{softmax}(S)\in \mathbb{R}^{L\times L} $$

最后进行加权求和:

$$ \text{Attention(Q,K,V)} = VA^\top \in \mathbb{R}^{d\times L} $$

令$P \in \mathbb{R}^{L\times L}$为置换矩阵，表示对序列顺序的重排。我们有:

$$ X' = XP $$

那么我们重新计算Attention:

$$ Q' = W_Q X' = W_QXP = QP \quad K' = W_KX'=KP \quad V' = VP $$$$ S' = \frac{Q'^\top K'}{\sqrt{d}} = \frac{P^\top Q^\top K P}{\sqrt{d}} $$$$ A' = \text{softmax}(S') = P^\top AP $$$$ \text{Attention}(Q',K',V') = VPA'^\top = VPP^\top A^\top P = VA^\top P = \text{Attention}(Q,K,V)P $$

上式表明，在不引入任何位置信息的情况下，Attention计算具有置换同变性。但是我们需要注意到，P是一个置换矩阵，它只是索引的重定向，我们计算出来的Attention的数值没有发生任何变化。如果我们把Attention输出结果看作带有语义信息的Embedding向量，那么它的语义将不会携带位置语义的信息。

举一个具体的例子:

我们输入的文本序列是["猫","吃","鱼"]，对应的输入序列为[x_1,x_2,x_3],而对应的Attention输出为[y1,y2,y3],而我们进行交换得到:["鱼","吃","猫"]，对应的输入序列为[x_3,x_2,x_1],根据我们上面的证明,Attention的输出将是[y3,y2,y1]，在数值上没有任何变换，只是简单的位置交换，也就是说Attention的输出没有提取出在语义结构(如主谓宾等)的信息。

位置编码为什么可以解决这个问题？

我们现在需要解决的问题是打破这种数值不变性，也就是交换了输入位置后，我们希望最终的输出为[y3',y2',y1'].我们将Attention视为一个输出为集合的函数f,那么上述不变性用数学语言表述就是:

$$ f(\dots,x_m,\dots,x_n,\dots) = f(\dots,x_n,\dots,x_m,\dots) $$

因此，我们要做的事情就是打破这种不变性，比如在每个位置都加上一个不同的编码向量:

$$ \hat f(\dots,x_m,\dots,x_n,\dots) = f(\dots,x_m+p_m,\dots,x_n+p_n,\dots) $$

一般来说，只要每个位置的编码向量不同，那么这种全对称性就被打破了，即可以用f~代替f来处理有序的输入。

我们写作矩阵的形式:

我们将位置编码矩阵定义为$E = [e_1,e_2,\dots,e_L]\in \mathbb{R}^{d\times L}$。每一个列向量$e_i$仅与位置索引i有关。

融入后的输入矩阵为:

$$ X_{pos} = X + E $$

此时我们重新推导Q,K,V的生成过程(以Q为例):

$$ Q = W_Q(X + E) = W_Q X + W_Q E = Q_X + Q_E $$

那么同理:

$$ K = K_X + K_E,V = V_X + V_E $$

我们可以将$\{K,Q,V\}_{E}$视为位置信息的“特征表达”

接下来我们计算$S = \frac{Q^\top K}{\sqrt{d}}$,我们具体来看其中一个点积项$s_{ij}$:

$$ s_{ij} = \frac{1}{\sqrt{d}}(q_{x,i}+q_{e,i})^\top(k_{x,j}+k_{e,j})\\ =\frac{1}{\sqrt{d}}(q^\top_{x,i}k_{x,j}+q^\top_{x,i}k_{e,j}+q^\top_{e,i}k_{x,j}+q^\top_{e,i}k_{e,j}) $$

那么此时就多出了位置-内容，位置-位置的信息。

我们再按照同样的推理,令$X' = XP$,那么此时:

$$ Q' = W_Q(X'+E) = W_Q(XP+E) = Q_XP+Q_E $$

同理:

$$ K' = K_XP+K_E \quad V' = V_XP+V_E $$

那么:

$$ S' = \frac{Q'^\top K'}{\sqrt{d}} = \frac{(Q_XP+Q_E)^\top(K_XP+K_E)}{\sqrt{d}}\\ =\frac{1}{\sqrt{d}}(P^\top Q_X^\top K_XP + P^\top Q_X^\top K_E + Q_E^\top K_XP+Q_E^\top K_E) $$

显而易见,$S' \neq P^\top SP$,且其中的$Q_E^\top K_E$项没有被P作用，也就意味着原本的置换对称性也就被打破了，因此我们可以用$\hat f $来代替f来处理有序的输入。

也就是说，位置编码的引入可以解决我们说的Attention置换数值不变性的问题。

怎样的位置编码是好的？

我们现在想要进一步分析位置编码的性质，从而设计更好的位置编码。我们将$\hat f$展开至二阶项(为了简化考虑，写作矩阵形式):

$$ \hat f(X) =f(X+E) \approx f(X) + \nabla f(X)\cdot E+\frac{1}{2}E^\top H_f(X)E $$

那么我们来看与位置编码有关的项:

一阶项$\nabla f(X)\cdot E = \sum_{i=1}^L e^\top_i \frac{\partial f}{\partial x_i}$:这项依赖于单一位置，所以是绝对位置信息
二阶项$\frac{1}{2}E^\top H_{f(X)}E = \sum_{i=1}^L\sum_{j=1}^L e_i^\top \frac{\partial^2 f}{\partial{x_i}\partial x_j}e_j$:这项与包含了任意两个位置的交互，所以是相对位置信息

对于一阶项，只要位置编码在每个位置都是独特的就能表述绝对位置的信息。而对于二阶信息，对于一个理想的位置编码而言，应该让这个二阶项满足某种平移不变性，即对于任意位移k，位置i与位置i+k的交互模式应该是稳定的。

我们先从最简单的情况入手，假设$\mathbf{H} = I$为单位矩阵，那么此时$E^\top E$是两个位置编码的内积，我们希望在这个简单的例子中该项表达的是相对位置信息，即存在某个函数g使得:

$$ \langle p_m,p_n\rangle = g(m-n) $$

这里的$p_m,p_n$为d维向量，这里我们从最简单的$d = 2$入手。我们称上式是一个位置编码为一个合理位置编码的条件式。

对于2维向量，我们借助复数来推导，视向量[x,y]为复数$x+yi$,那么我们有:

$$ \langle p_m,p_n\rangle=a_xb_x+a_yb_y = \text{Re}(p_m\hat p_n) $$

其中$\hat w$为$w$的共轭复数。

为了满足上式，我们可以假设存在复数$q_{m-n}$,使得:

$$ p_m\hat p_n = q_{m-n} $$

这样两边取实部就得到条件式。为了解这个方程，我们可以使用复数的指数形式，假设$p_m = r_me^{i\phi_m},\hat p_n = r_ne^{-i\phi_n},q_{m-n} = R_{m-n}e^{i\Phi_{m-n}}$,那么有:

$$ r_mr_ne^{i(\phi_m-\phi_n)}=R_{m-n}e^{i\Phi_{m-n}} $$

于是我们得到等式:

$$ \left\{ \begin{array}{l} r_m r_n = R_{m-n} \\ \phi_m - \phi_n = \Phi_{m-n} \end{array} \right. $$

对于第一个方程，带入m = n，可以得到$r_m^2 = R_0$,即$r_m$为一个常数，为了简单，我们令其为1；

对于第二个方程，显然等差数列满足上述性质(令m = 0有$\Phi_m = \phi_m$)，设公差为$\theta$，则通项为$\phi_m = \Phi_m = m\theta$,由此我们得到二维情况下的位置编码的解:

$$ p_m = e^{im\theta} \to p_m = \left(\begin{align}\cos m\theta\newline \sin m\theta\end{align}\right) $$

由于内积满足线性叠加性，我们可以由二维的情况直接扩展到更高偶数维的情况:

$$ p_m = \left(\begin{matrix}e^{im\theta_0}\\e^{im\theta_1}\\\vdots\\e^{im\theta_{d/2-1}}\end{matrix}\right) \to p_m = \left(\begin{matrix} \cos m\theta_0 \\ \sin m\theta_0 \\ \cos m\theta_1 \\ \sin m\theta_1 \\ \vdots\\ \cos m\theta_{d/2-1} \\ \sin m\theta_{d/2-1} \\ \end{matrix}\right) $$

这样我们就求出了满足条件式的一组解，显然解不唯一。

此外，一个好的位置编码应该满足远程衰减的性质，即随着$|m-n|$的增大,$\langle p_m,p_n\rangle$有趋于0的趋势。

那么有:

$$ \langle p_m,p_n\rangle = \text{Re}[e^{i(m-n)\theta_0} + e^{i(m-n)\theta_1}+\dots+e^{i(m-n)\theta_{d/2-1}}]\\ =\sum_{j=0}^{d/2-1}\cos (k\theta_j) \quad k = m-n $$

由于在LM中d通常为768，是一个较大值，因此我们可以将离散的索引j映射到连续变量$t\in[0,1]$上。设$\theta_j$是某个光滑单调函数$f(t)$生成的，即$\theta_j = f(2j/d)$。利用Euler-Maclaurin的一阶近似，我们可以将求和转化为积分:

$$ \langle p_m,p_n\rangle \approx \frac{d}{2}\int_0^1 \cos(k\cdot f(t))dt $$

那么现在的问题就转化为了，寻找一个函数$f(t)$,使得上述震荡积分在k很大时具有较好的衰减性质(足够快)。

根据黎曼-勒贝格引理，只要频率分布函数f(t)满足光滑且严格单调的条件，当$k\to \inf$时，被积函数的高频震荡将导致正负面积互相抵消，使得积分值必然趋近于0.

在Transformer2017的论文中的Sinusoidal位置编码选择的是$\theta_t = 10000^{-t}$.

由此，我们便推导出了Sinusoidal位置编码的形式:

$$ \left\{ \begin{array}{l} p_{k,2i} = \sin(k/10000^{2i/d})\\ p_{k,2i+1} = \cos(k / 10000^{2i/d}) \end{array} \right. $$

我们只能说明它的合理性，但是无法说明它的最优性，因为它并不一定最优[Lol]

事实上，一个可行的方案是将位置编码中的$\theta_i$设为可学习的参数，其初始值为$\theta_i = 10000^{-2i/d}$.

在上述推导中，都是基于H = I这个简单情况，对于一般的H，使用上述Sinusoidal位置编码，还能具备我们理想的性质吗？

事实上，有研究表明$^{[3]}$ 在网络规模足够大的情况下，其Hessian矩阵将呈现出块对角的形式，因此我们考虑H是一个对角阵的情况，此时:

$$ p_m^\top H p_n = \sum_{i=1}^{d/2}H_{2i,2i}\cos m\theta_i \cos n\theta_i+H_{2i+1,2i+1}\sin m\theta_i\sin n\theta_i $$

由和差化积有:

$$ \sum_{i=1}^{d/2}\frac{1}{2}(H_{2i,2i}+H_{2i+1,2i+1})\cos(m-n)\theta_i + \frac{1}{2}(H_{2i,2i}-H_{2i+1,2i+1})\cos(m+n)\theta_i $$

可以看到其中是包含了相对位置项(m-n)的，只是会出现m+n项。

因此我们可以认为Sinusoidal位置编码是一个有效的位置编码。

为什么RoPE$^{[4]}$比Sinusoidal位置编码更好？

在先前分析为什么需要位置编码的过程中，我们知道位置编码的作用是在计算Attention Score时，让QK点积中包含文本的结构语义信息，具体而言是下式中的后三项包含位置信息:

$$ S' = \frac{Q'^\top K'}{\sqrt{d}} = \frac{(Q_XP+Q_E)^\top(K_XP+K_E)}{\sqrt{d}}\\=\frac{1}{\sqrt{d}}(P^\top Q_X^\top K_XP + P^\top Q_X^\top K_E + Q_E^\top K_XP+Q_E^\top K_E) $$

我们可以发现，除了最后一项是相对位置信息外，另外两项则是绝对位置信息与文本语义信息的耦合，这实际上是一种冗余，因为绝对位置信息的作用相较于相对位置信息而言一方面具有误导性，例如在短文本中结尾出现在绝对位置10中，而10在长文本中可能仅仅是文本开头；另一方面模型还会额外去学习已经得到的相对位置信息造成资源的浪费.

为了解决这个问题，工业界进行了许多尝试$^{[5]}$,但是这些尝试普遍会带来大量的额外计算存储开销。由于绝对位置编码具有实现简单，计算速度快的特点，并且在Sinusoidal位置编码中我们也能看到通过绝对位置编码在一定程度上是可以得到相对位置信息的,如果可以通过绝对位置编码的方式实现相对位置编码，那么就是“集各家之所长”，“鱼和熊掌兼得”了。

为了实现这个目标，我们假设通过下述运算来给q,k添加绝对位置信息:

$$ \hat q_m = f(q,m)\quad \hat k_n = f(k,n) $$

而我们希望得到如下恒等关系:

$$ \langle f(q,m),f(k,n)\rangle = g(q,k,m-n) $$

为了求解方便，我们令:

$$ f(q,0) = q \quad f(k,0) = k $$

求解思路与之前推导Sinusoidal类似，我们先考虑二维的情况，然后借助复数来求解。

$$ \text{Re}[f(q,m)\hat f(k,n)] = g(q,k,m-n) $$

设:

$$ \begin{array}{rcl} f(q,m) &=& R_f(q,m)e^{i\theta_f(q,m)}\\ \hat f(k,n) &=& R_f(k,n)e^{-i\theta_f(k,n)}\\ g(q,k,m-n) &=& R_g(q,k,m-n)e^{i\theta_g(q,k,m-n)} \end{array} $$

那么带入方程求解得到:

$$ \left\{ \begin{array}{l} R_f(q,m)R_f(k,n) = R_g(q,k,m-n)\\ \theta_f(q,m) - \theta_f(k,n) = \theta_g(q,k,m-n) \end{array} \right . $$

对于第一个方程，令m = n有:

$$ R_f(q,m)R_f(k,n) = R_g(q,k,0) = R_f(q,0)R_f(k,0)=||q||||k|| $$

那么我们可以直接令$R_f(q,m) = ||q||$,即它不依赖于m。

对于第二个方程，同样地令m=n有:

$$ \theta_f(q,m)-\theta_f(k,n)=\theta_g(q,k,0)=\theta_f(q,0)-\theta_f(k,0) = \theta_q-\theta_k $$

这里的$\theta_q,\theta_k$是q,k本身的辐角。

由此我们得到:

$$ \theta_f(q,m) - \theta_q = \theta_f(k,n) - \theta_k $$

因此$\theta_f(q,m)-\theta_q$应该是一个只与m有关的而与q无关的函数，因为我们希望恒等式对任意的q,k都成立，所以右侧必须与q,k无关，因此记为常数，记为$\theta$，该式记为$\phi(m)$,即

$$ \theta_f(q,m) = \theta_q + \phi(m) $$

令n = m - 1 有:

$$ \phi(m) - \phi(m-1) = \theta_g(q,k,1)+\theta_k-\theta_q $$

故$\{\phi(m)\}$为等差数列，公差为$\theta$,得到$\phi(m) = m\theta$

由此，我们得到了二维情况下用复数表示的RoPE:

$$ f(q,m) = R_f(q,m)e^{i\theta_f(q,m)} = ||q||e^{i(\theta_q+m\theta)} = \vec{q}e^{im\theta} $$

根据复数乘法的集合意义，该变换实际上对应着向量的旋转，我们可以将其写为矩阵形式:

$$ f(q,m) = \left( \begin{matrix} \cos m\theta \quad -\sin m\theta\\ \sin m \theta \quad \cos m\theta \end{matrix} \right) \left( \begin{matrix} q_0\\ q_1 \end{matrix} \right) $$

由内积的线性叠加性，我们可以得到任意偶数维度的RoPE:

$$ \left( \begin{matrix} \cos m\theta_0 & -\sin m\theta_0 & 0 & 0 & \dots & 0 & 0 & \\ \sin m\theta_0 & \cos m\theta_0 & 0 & 0 & \dots & 0 & 0 & \\ 0 & 0 & \cos m\theta_1 & -\sin m\theta_1 & \dots & 0& 0\\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \dots &\cos m\theta_{d/2-1} & -\sin m \theta_{d/2-1}\\ 0 & 0 & 0 & 0 & \dots & \sin m\theta_{d/2-1} & \cos m \theta_{d/2-1} \end{matrix} \right) \left( \begin{matrix} q_0\\q_1\\q_2\\q_3\\ \vdots\\q_{d-2}\\q_{d-1} \end{matrix} \right) $$

我们称左侧的矩阵为旋转矩阵记为$R_m$.也就是说给位置为m的向量q乘上矩阵$R_m$,位置为n的向量k乘上矩阵$R_n$,用变换后的Q，K序列做Attention，那么Attention就自动包含相对位置信息了，因为恒等式成立:

$$ (R_mq)^\top(R_nk) = q^\top R_m^\top R_nk = q^\top R_{n-m}k $$

值得指出的是，Rm是一个正交矩阵，它不会改变向量的模长，因此通常来说它不会改变原模型的稳定性。

此外，在具体实现时，我们并不会拿这个大矩阵去乘以向量，而是采用逐位相乘的方式:

$$ \left( \begin{matrix} q_0\\q_1\\q_2\\q_3\\ \vdots\\q_{d-2}\\q_{d-1} \end{matrix} \right)\otimes \left( \begin{matrix} \cos m\theta_0\\\cos m\theta_0 \\\cos m\theta_1 \\ \cos m\theta_1\\ \vdots \\ \cos m\theta_{d/2-1} \\ \cos m\theta_{d/2-1} \end{matrix} \right)+ \left( \begin{matrix} -q_1\\q_0\\-q_3\\q_2\\ \vdots\\-q_{d-1}\\q_{d-2} \end{matrix} \right)\otimes \left( \begin{matrix} \sin m\theta_0\\\sin m\theta_0\\\sin m\theta_1\\\sin m\theta_1\\ \vdots\\\sin m\theta_{d/2-1}\\\sin m\theta_{d/2-1} \end{matrix} \right) $$

对于$\theta$的选择，作者选择了与Sinusoidal位置编码一样的$\theta_i = 10000^{-2i/d}$,从而带来远程衰减性。

至此，我们完成了对RoPE的推导，并通过推导成功说明了为何RoPE相较于Sinusoidal位置编码更好，因为它通过绝对位置的注入方式，实现了相对位置的注入，而不带来其它冗余。

最后，下面是我在CS336中实现的一个RoPE:

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

参考

Hessian矩阵系列串讲

Mon, 24 Nov 2025 13:28:52 +0800

定义

假设有一实值函数$f(x_1,x_2,\dots,x_n)$，若$f$的所有二阶偏导数都存在且在定义域里连续，那么我们定义函数$f$的Hessian矩阵为如下一个$n\times n$的方阵:

$$ \mathbf{H} = \left[ \begin{matrix} \frac{\partial^2 f}{\partial x_1^2} & \frac{\partial^2 f}{\partial x_1\partial x_2} & \dots&\frac{\partial^2 f}{\partial x_1 \partial x_n}\\ \frac{\partial^2 f}{\partial x_2 \partial x_1} & \frac{\partial^2 f}{\partial x_2^2} & \dots& \frac{\partial^2 f}{\partial x_2 \partial x_n}\\ \vdots & \vdots & \ddots & \vdots\\ \frac{\partial^2 f}{\partial x_n \partial x_1} & \frac{\partial^2 f}{\partial x_n \partial x_2} & \dots & \frac{\partial^2 f}{\partial x_n^2} \end{matrix} \right] $$

或使用下标标记表示为:

$$ \mathbf{H}_{ij} = \frac{\partial^2 f}{\partial x_i \partial x_j} $$

一般性质

对称性:

对于机器学习中遇到的大多数函数(特别是那些具有连续二阶偏导数的函数),混合偏导数的求导顺序无关紧要。这被称为克莱罗定理或施瓦茨定理,它说明了混合偏导数的相等性:

$$ \frac{\partial^2 f}{\partial x_i\partial x_j} = \frac{\partial^2 f}{\partial x_j \partial x_i} $$

这意味着Hessian矩阵是对称的，即$\mathbf{H} = \mathbf{H}^\top$

正定性与局部最优值的相关性

通过Hessian矩阵的正定性，我们可以判断该点是局部最小值，局部最大值还是鞍点

正定:若Hessian矩阵$H(x^*)$是正定的(意味着$\forall v \neq 0,v^\top\mathbf{H}v>0$)，则函数在$x^*$处有局部最小值。函数在该点周围向所有方向向上弯曲，类似于碗的底部
负定:若Hessian矩阵$H(x^{*})$是负定的(意味着$\forall v \neq 0,v^\top\mathbf{H}v<0$),则函数在$x^*$处有局部最大值。函数在该点周围向所有方向向下弯曲，类似于圆顶的顶部
不定:若Hessian矩阵$H(x^{*})$是不定的,则函数在$x^*$处有鞍点。函数在某些方向向上弯曲，在另一些方向向下弯曲，就像马鞍

Hessian在神经网络中的计算

Hessian矩阵在神经网络计算的许多方面有着重要作用，包括:

⼀些⽤来训练神经⽹络的⾮线性最优化算法是基于误差曲⾯的⼆阶性质的，这些性质

由Hessian矩阵控制(比如牛顿法和拟牛顿法)
对于训练数据的微⼩改变，Hessian矩阵构成了快速重新训练前馈⽹络的算法的基础
Hessian矩阵的逆矩阵⽤来鉴别神经⽹络中最不重要的权值，这是⽹络“剪枝”算法的⼀部分(LeCun的OBD)

对于Hessian矩阵的众多应用而言，一个重要的需要考虑的问题是计算效率。在神经网络中有$W$个参数(包括权值和偏置),那么Hessian矩阵的大小就是$W\times W$,那么计算Hessian矩阵的计算量为$O(W^2)$。这在具有大量参数的神经网络中是难以接受的，因此我们需要进行一些高效的近似。

对角近似

这个方法最早由Yan LeCun在OBD剪枝方法中提出

我们只保留Hessian矩阵的对角线元素$H_{i,i}$，把非对角线元素置为0.(这个方法在很大程度可以可以方便求逆)

对于神经网络中的第j个神经元的输入$a_j$对应权重为$w_{ji}$:

$$ \frac{\partial^2 E_n}{\partial w_{ji}^2} = \frac{\partial^2 E_n}{\partial a_j^2}z_i^2 $$

其中$z_i$是上一层神经元的输出。

而$\frac{\partial^2 E_n}{\partial a_j^2}$可以通过链式法则递归计算(类似于反向传播):

忽略二阶导中的非对角线项$k\neq k'$:

从而一次反向传播便可以计算出来，时间复杂度为$O(W)$

但是如你所见，在这个方法下，Hessian矩阵完全退化为了一个diag，这在神经网络中是不合理的，因为非线性层的引入会让Hessian矩阵的交叉项不为0，更重要的是这种误差会随着网络层数的堆砌不断放大。

因此针对这个问题在对角近似上近年来不断有相关文章发表，下面对我阅读过的一些进行简单介绍:

HesScale:在LeCun的方法的骨架上，针对常见网络的最后一层的softmax+CE结构可以通过一个简单的公式求出diag(H)的精确值:
$$ \frac{\partial^2 E_n}{\partial^2 a_j^2} = p - p\circ p $$
其中p为$\text{softmax}(a_j)$,$\circ$表示向量的逐元素乘积,为方便书写我们将$\{a_1,\dots,a_n\}$写作z

证明如下:

预测$\mathbf{p} = \text{softmax}(z)$,即$p_i = \frac{\exp(z_i)}{\sum_{j = 1}^k \exp(z_j)}$,设目标分布为$t\in \Delta^{K-1},\sum_{i}t_i = 1$，对于单样本而言，其损失可以写作:
$$ \mathcal{L}(z) = -\sum_{j=1}^Kt_i\log p_i = -t^\top z + \log \sum_{j=1}^K \exp(z_j) $$
求一阶导我们有:
$$ \nabla_z \mathcal{L} = -t + \nabla_z \log \sum_{j=1}^K \exp{z_j} = -t + \mathbf{p} $$
求二阶导相当于$\mathbf{p}$对z求导。我们知道$\text{softmax}$的Jocobian为
$$ \frac{\partial p_i}{\partial z_j} = p_i(\delta_{ij} - p_j),J_{\text{softmax}}(z) = \text{diag}(\mathbf{p}) - \mathbf{p}\mathbf{p}^\top $$
其中$\delta_{i,j}$为Kronecker函数，写成矩阵形式，则只有对角线元素为1，其余为0.

因此我们有:
$$ \nabla^2_z\mathcal{L} = \text{diag}(\mathbf{p}) - \mathbf{pp}^\top $$
于是我们可以得到:
$$ \frac{\partial^2 L}{\partial^2 a_j} = p_j(1-p_j) $$
AdaHessian:

具体而言，AdaHessian在进行Hessian矩阵的对角近似的时候利用了Hutchinson估计，它通常也被用来对矩阵的迹进行估计:

对任意矩阵$\mathbf{A}\in\mathbb{R}^{d\times d}$,若随机向量$z = (z_1,\dots,z_d)^\top$满足:
$$ \mathbb{E}[z_i] = 0\\ \mathbb{E}[z_iz_j] = \delta_{ij} $$
那么有:
$$ \mathbb{E}[z\odot (Az)] = \text{diag}(A) $$
下面进行证明:

记估计量$d = z\odot (Az)$的第i个分量:
$$ d_i = z_i(Az)_i = z_i\sum_{j=1}^d A_{ij}z_j = \sum_{j = 1}^d A_{ij}z_iz_j\\ \mathbb{E}[d_i] = \sum_{j=1}^dA_{ij}\mathbb{E}[z_iz_j] = \sum_{j=1}^d A_{ij}\delta_{ij} = A_{ii} $$
由此我们证明了$\mathbb{E}[z\odot(Az)] = \text{diag}(A)$

在神经网络中，相较于直接计算完整的Hessian矩阵以及其OBD方式的对角近似，计算其矩阵向量积(HVP,Hessian Vector product)是并不困难的，它只需要一次反向传播：
$$ Hz = \frac{\partial (g^\top z)}{\partial \theta} = \frac{\partial g^\top}{\partial \theta}z + g^\top \frac{\partial z}{\partial \theta} = \frac{\partial g^\top}{\partial \theta}z $$
由此，我们可以通过多次在满足Rademacher分布的向量取样计算$z\odot (Az)$的期望就能得到$\text{diag}(A)$的无偏估计，在实际的应用中，只进行一次取样就能得到较为不错的结果。

外积近似

在神经网络应用于回归问题时，通常采用下面形式的平方和误差

$$ E = \frac{1}{2}\sum_{n=1}^N(y_n-t_n)^2 $$

那么Hessian矩阵可以写成如下形式:

$$ H = \nabla\nabla E = \nabla(\sum_{n=1}^N (y_n-t_n)\nabla y_n) = \sum_{n = 1}^N \nabla y_n(\nabla y_n)^\top + \sum_{n=1}^N(y_n-t)\nabla\nabla y_n $$

在网络已经训练好的情况下，输出$y_n$与$t_n$接近，因此第二项可以忽略，由此我们得到了Hessian矩阵的外积近似

$$ H\approx \sum_{n=1}^N b_n b_n^\top $$

其中，$b_n = \nabla y_n = \nabla a_n$(输出单元的激活函数就是恒等函数)。这种方法中的Hessian矩阵可以跟随反向传播算法在$O(W)$个步骤内高效地求出误差函数地一阶导数。再通过简单地乘法就可以在$O(W^2)$步骤内求出矩阵元素。

Hessian矩阵逆的计算

使用外积近似，我们可以提出一个计算Hessian矩阵的逆的高效办法,首先我们有:

$$ \mathbf{H}_N = \sum_{i=1}^n b_n b_n^\top $$

其中，$b_n = \nabla w a_n$时数据点n产生的输出单元激活对梯度的贡献。我们现在推导一个建立Hessian矩阵的顺序步骤，每次处理一个数据点。假设我们已经使用了前L个数据点得到了Hessian矩阵的逆。通过将第L+1个数据点的贡献单独写出来，我们有:

$$ \mathbf{H}_{L+1} = \mathbf{H}_L + b_{L+1}b_{L+1}^\top $$

为了计算Hessian矩阵的逆，我们考虑下面这个矩阵的恒等式:

$$ (M + vv^\top)^{-1} = M^{-1} - \frac{(M^{-1}v)(v^\top M^{-1})}{1+v^\top M^{-1}v} $$

若我们令$H_L = M$,且$b_{L+1} = v$,我们有:

$$ H_{L+1}^{-1} = H_{L}^{-1} - \frac{H_L^{-1}b_{L+1}b_{L+1}^\top H_{L}^{-1}}{1 + b_{L+1}^\top H_L^{-1}b_{L+1}} $$

通过这种方式，数据点可以依次使用，直到L+1=N,整个数据集处理完毕。于是，这个结果表示一个计算Hessian矩阵的逆的算法。这个算法只需对数据集扫描一次。最开始的矩阵$\mathbf{H_0}$被选为$\alpha I$,其中$\alpha$是一个较小的量，从而算法实际找的是$\mathbf{H} + \alpha I$的逆。结果对于$\alpha$的精确值不敏感。

Hessian矩阵的逆计算的K-FAC分解方法：

在特定场景下，Hessian矩阵的逆有一个利用Kronecker分解的近似方法，它利用了在loss为log-like形式下$\mathbb{E}[H] = F$的特点，其中$F$为Fisher信息矩阵:

$$ F = \mathbb{E}[\nabla \log p(x)\nabla \log p(x)^\top] $$

其中上式中做了如下形式的简写:

$$ \nabla \log{p(x)} = \nabla_\theta \log{(p|\theta)}\\ \mathbb{E}[p(x)] = E_{x\sim p(x|\theta)} $$

首先我们有:

$$ \mathbb{E}[\nabla \log p(x)] = \int (\nabla \log p(x))p(x) dx\\ =\int \frac{\nabla \log p(x)}{p(x)}p(x) dx\\ =\int \nabla \log p(x)dx\\ =\nabla 1 = 0 $$

那么当loss为log-like形式下时:

$$ L = -\log p(x) $$

我们求它的Hessian:

$$ \nabla^2 L = -\nabla^2 \log p(x)\\ =-\nabla \frac{\nabla p(x)}{p(x)}\\ =\frac{\nabla p(x)}{p(x)}^\top \frac{\nabla P(x)}{p(x)} - \frac{P(x)^2}{p(x)}\\ =\nabla \log p(x)^\top \nabla \log p(x) - \frac{\nabla^2 P(x)}{p(x)} $$

对该式求期望，我们有:

$$ \mathbb{E}[H]=\mathbb{E}[\nabla \log p(x)^\top \nabla \log p(x)] - \mathbb{E}[\frac{\nabla^2 p(x)}{p(x)}] = F - \int \frac{\nabla^2 p(x)}{p(x)} p(x) dx\\ =F - \int \nabla^2 p(x)dx\\ =F - \nabla^2 \int p(x)\\ =F $$

由于$\theta$本质为所有层$\mathbf{W}$的拼接，我们有:

$$ d\theta = \nabla_{\theta}L(x)\\ F = \mathbb{E}[\nabla L(x)\nabla L(x)^\top] = \mathbb{E}[d\theta d\theta^\top]\\ \theta = [\text{vec}(W_0)^\top,\text{vec}(W_1)^\top,\dots,\text{vec}(W_n)^\top]^\top $$

带入展开得到:

$$ F_{ij} = \mathbb{E}[\text{vec}(dW_i)\text{vec}(dW_j)^\top] $$

令$a_i,g_i$分别为第i层的前向输入和反向传播梯度，由反向传播算法有:

$$ d W_i = g_ia_i^\top $$

带入有:

$$ \text{vec}(dW_i) = \text{vec}(g_ia_i^\top) = a_i\otimes g_i $$

这里可能不太直观，这是因为平常对kronecker积($\otimes$)的接触较少，具体而言，我们这样定义Kronecker积:

$$ A\otimes B = \left[ \begin{matrix} a_{11}B &\dots &a_{1n}B\\ \vdots & \ddots& \vdots\\ a_{m1}B & \dots & a_{mn}B \end{matrix} \right] $$

在这里，由于我们将$dW_i$即$\nabla W_i$第i层参数的梯度进行了向量化，原本的$\nabla W_i$ 是如下形式:

$$ \nabla W_i = g_i a_{i}^\top = \left[\begin{matrix}g_1a_1 & \dots &g_1a_n\\ \vdots & \ddots & \vdots\\ g_ma_1 &\dots&g_ma_n\end{matrix} \right] $$

按列堆叠(vec)就有:

$$ \text{vec}(dW_i) = \left[ \begin{matrix} g_1a_1\\ \vdots\\ g_1a_n\\ \vdots\\ g_ma_n \end{matrix} \right] = a_i\otimes g_i $$

那么我们可以将Fisher矩阵写作如下形式:

$$ F_{ij} = \mathbb{E}[\text{vec}(dW_i)\text{vec}(dW_j)]\\ =\mathbb{E}[(a_i\otimes g_i)(a_j\otimes g_j)^\top]\\ =\mathbb{E}[(a_ia_j^\top)\otimes (g_ig_j^\top)]\\ \approx \mathbb{E}[a_ia_j^\top]\otimes\mathbb{E}[g_ig_j^\top] $$

这个近似相当于我们忽略了$\text{Cov}_\otimes(a_ia_j^\top,g_ig_j^\top)$,这实际上是合理的，尤其是在网络较深的情况下。这里给出一个链接提供一个比较详细的说明。

虽然直接求$F_{ij}$的复杂度仍然不变，但是利用Kronecker积的性质，我们在求逆时可以得到较大的性能提升:

$$ (A\otimes B)^{-1} = A^{-1}\otimes G^{-1} $$

并且在实际运算中，并不是整个网络的Fisher进行计算，而是按层做块对角运算:

$$ \mathbf{F} = \text{blockdiag}(F_1,F_2,\dots),F_l\approx A_l\otimes G_l $$

这样每层求逆的时候只需要求两个小矩阵的逆。

一些统计量的计算：

除了直接对Hessian矩阵的近似计算，我们有时候仅仅需要对Hessian矩阵的统计量进行计算，例如矩阵的迹，最大特征值等。

在计算这些统计量时，我们需要一个重要的算子:Hv,即Hessian矩阵与任意向量的乘积，这个乘积我们在上面已经证明了通过一次简单的反向传播算法可以得到。(PyHessian是一个Python库，它实现了这个算子的高效计算)

顶部若干特征值

利用幂迭代法 + HVP可以高效地求解特征值。

简单来说幂迭代法的流程很简单:

给定矩阵A以及初始非零向量$x_0$,

迭代:
$$ y_{k+1} = Ax_k\\ x_{k+1} = \frac{y_{k+1}}{||y_{k+1}||}\\ \mu_{k+1} = \frac{x_{k+1}Ax_{k+1}^\top}{x_{k+1}x_{k+1}^\top} $$
具体的收敛性证明见教材
矩阵的迹

在上面近似计算中，我们曾用Hutchinson估计对Hessian矩阵进行了对角近似，事实上，我们还可以进行迹的估计。

简单来说，方法一样，从Rademacher分布中取一随机向量$v$,然后有恒等式:
$$ \text{Tr}(H) = \text{Tr}(HI) = \text{Tr}(H\mathbb{E}[vv^\top]) = \mathbb{E}[\text{Tr}(Hvv^\top)] = \mathbb{E}[v^\top H v] $$
Hessian矩阵在神经网络下的特殊结构
本小节内容主要基于Towards Quantifying the Hessian Structure of Neural Networks,B站上有作者的讲解视频[FAI] 港中深张雨舜 | 浅谈神经网络Hessian矩阵的特殊结构

这篇文章主要说明了在神经网络中，Hessian矩阵往往具有近块对角结构，这表明曲率信息(二阶信息)主要在层内耦合，层与层的二阶交互很弱。这间接的说明了Layer-wise的量化/剪枝的合理性。

此外，文章还进一步从理论和实验上阐释了造成这个现象的原因：

静态力量:即使在随机初始化阶段，即在训练开始之前，神经网络的Hessian矩阵就已经呈现出近块对角结构。这种结构的形成与网络的架构设计有关，因此被称为“静态力量”。具体来说，对于线性模型和单隐藏层网络，无论是使用均方误差（MSE）损失还是交叉熵（CE）损失，Hessian矩阵的对角块和非对角块在随机初始化时就已经表现出明显的差异。
动态力量:在训练过程中，Hessian矩阵的结构会进一步发生变化。特别是在使用CE损失时，训练过程会逐渐消除初始时存在于跨层Hessian分量（Hwv）中的“块循环”（block-circulant）模式，而对角块和非对角块的近块对角结构则保持稳定。这种由训练过程引起的结构变化被称为“动态力量”。

另外一点，文章证明了类别数C是影响Hessian矩阵结构的主要因素。(在实验中，隐藏层Hessian的非对角块与对角块的比值以$\frac{1}{\sqrt{C}}$的速度衰减，输出层Hessian的衰减速率为$\frac{1}{C}$).这个结果对于大模型而言是友好的，因为在神经网络中C的大小往往是$1e3\sim 1e4$级别的，这说明大模型的Hessian是具有强对角块的结构的！这一点实际上在传统的方法上人们意识or无意识的用到了(对角近似)，但是更丰富层面以及基于这个发现的在计算的可行性和性能的Trade-off做的工作是比较少的，在优化器那边做的比较多。

此外还有不少文献揭露了Hessian矩阵具有低秩特征谱，即只有少数的特征值显著大，其余大多接近0.