feat: 完整中文翻译 maths-cs-ai-compendium（数学·计算机科学·AI 知识大全）
翻译自英文原版 maths-cs-ai-compendium，共 20 章全部完成。 第01章 向量 | 第02章 矩阵 | 第03章 微积分 第04章 统计学 | 第05章 概率论 | 第06章 机器学习 第07章 计算语言学 | 第08章 计算机视觉 | 第09章 音频与语音 第10章 多模态学习 | 第11章 自主系统 | 第12章 图神经网络 第13章 计算与操作系统 | 第14章 数据结构与算法 第15章 生产级软件工程 | 第16章 SIMD与GPU编程 第17章 AI推理 | 第18章 ML系统设计 第19章 应用人工智能 | 第20章 前沿人工智能 翻译说明： - 所有数学公式 $...$ / $$...$$、代码块、图片引用完整保留 - mkdocs.yml 配置中文导航 + language: zh - README.md 已翻译为中文（兼 docs/index.md） - docs/ 目录包含指向各章文件的 symlink - 约 29,000 行中文内容，排除 .cache/ 构建缓存
This commit is contained in:
@@ -0,0 +1,339 @@
|
||||
# 量化
|
||||
|
||||
*量化降低模型权重和激活值的精度，使模型更小、更快、运行成本更低。本文涵盖数字格式、训练后量化、量化感知训练、仅权重量化方法（GPTQ、AWQ）、激活值量化、混合精度和KV缓存量化*
|
||||
|
||||
- 一个70B参数的float16模型需要140 GB内存，超过任何单张GPU。量化为INT4后，它可以装入35 GB（一张A100）甚至20 GB（带卸载的消费级RTX 4090）。量化不是一种可有可无的优化；它是让大模型部署在经济上可行的关键。
|
||||
|
||||
- 基本权衡：低精度意味着更少内存、更高吞吐量和更低功耗，但会引入**量化误差**，可能降低模型质量。量化的艺术在于最小化这种降级。
|
||||
|
||||
## 为什么要量化
|
||||
|
||||
- **内存减少**：INT8比FP16小2倍，INT4小4倍。对于LLM，模型权重占主导内存。精度减半意味着内存需求减半。
|
||||
|
||||
- **吞吐量提升**：低精度意味着每秒更多操作。NVIDIA Tensor Core（第16章）在FP16 vs FP32上实现2倍吞吐量，INT8 vs FP16再实现2倍，INT4 vs INT8再实现2倍。H100在FP8下达到989 TFLOPS，而FP32下只有67 TFLOPS——相差15倍。
|
||||
|
||||
- **带宽节省**：LLM推理通常是**内存带宽受限**的（第16章，屋顶模型）。瓶颈是从GPU内存加载权重，而不是计算。更小的权重意味着更少的传输字节，直接提高每秒token数。这就是量化通常能为LLM推理带来近乎线性加速的原因。
|
||||
|
||||
- **节能**：低精度每次操作消耗更少能量。在数据中心规模（数千GPU）下，这转化为显著的电力成本降低。
|
||||
|
||||
## 数字格式
|
||||
|
||||
- 我们在第13章（计算机体系结构）中介绍了IEEE 754浮点数。以下是ML的完整精度全景：
|
||||
|
||||
|
||||
|
||||
| 格式 | 位数 | 指数 | 尾数 | 范围 | 用途 |
|
||||
|--------|------|----------|----------|-------|----------|
|
||||
| FP32 | 32 | 8 | 23 | ±3.4×10³⁸ | 训练（黄金标准） |
|
||||
| TF32 | 19 | 8 | 10 | ±3.4×10³⁸ | Tensor Core训练（A100+） |
|
||||
| FP16 | 16 | 5 | 10 | ±65504 | 混合精度训练 |
|
||||
| BF16 | 16 | 8 | 7 | ±3.4×10³⁸ | 训练（与FP32相同的范围） |
|
||||
| FP8 E4M3 | 8 | 4 | 3 | ±448 | 前向传播（Hopper+） |
|
||||
| FP8 E5M2 | 8 | 5 | 2 | ±57344 | 梯度（更宽范围） |
|
||||
| INT8 | 8 | — | — | -128 到 127 | PTQ推理 |
|
||||
| INT4 | 4 | — | — | -8 到 7 | 仅权重量化 |
|
||||
| INT2/三值 | 2 | — | — | {-1, 0, 1} | 极限压缩 |
|
||||
|
||||
- **FP8**有两种变体：**E4M3**（4位指数，3位尾数，范围较窄但精度更高）用于前向传播，**E5M2**（5位指数，2位尾数，范围更宽但精度较低）用于梯度。Transformer Engine（第16章）在每个张量之间自动切换。
|
||||
|
||||
- **BF16 vs FP16**：BF16具有与FP32相同的指数范围（无溢出风险），但尾数精度较低。FP16精度更高但范围较窄（最大65504），训练时需要损失缩放。对于推理，两者都表现良好；对于训练，BF16更安全。
|
||||
|
||||
- **整数格式**没有指数——它们表示定点值。要在浮点和整数之间转换，需要一个**缩放因子**和一个可选的**零点**：$x_{\text{float}} = \text{scale} \times (x_{\text{int}} - \text{zero\_point})$。
|
||||
|
||||
## 量化方程
|
||||
|
||||
- 所有量化方法都将浮点值映射到整数并返回：
|
||||
|
||||
$$x_q = \text{clamp}\left(\text{round}\left(\frac{x}{\text{scale}}\right) + \text{zero\_point}, \; q_{\min}, \; q_{\max}\right)$$
|
||||
|
||||
$$\hat{x} = \text{scale} \times (x_q - \text{zero\_point})$$
|
||||
|
||||
- **缩放因子**决定分辨率：$\text{scale} = \frac{x_{\max} - x_{\min}}{q_{\max} - q_{\min}}$。对于INT8：$q_{\min} = -128$，$q_{\max} = 127$。
|
||||
|
||||
- **对称量化**设置$\text{zero\_point} = 0$，因此$\text{scale} = \frac{\max(|x|)}{127}$。更简单、更快（推理时无需减去零点）。
|
||||
|
||||
- **非对称量化**使用非零$\text{zero\_point}$来处理非对称分布（例如，ReLU输出全为非负）。将$[x_{\min}, x_{\max}]$映射到无符号INT8的$[0, 255]$。
|
||||
|
||||
|
||||
|
||||
- **量化粒度**：多少个值共享同一个缩放因子：
|
||||
- **逐张量**：整个张量一个缩放因子。最简单但精度最低（一个异常值就会扭曲整个张量的缩放因子）。
|
||||
- **逐通道**：每个输出通道（卷积）或每行（线性层）一个缩放因子。精度好得多，开销最小。
|
||||
- **逐组**：每$g$个元素一组（例如$g = 128$）一个缩放因子。精度最佳，用于现代仅权重量化（GPTQ、AWQ）。
|
||||
- **逐token**：每个token一个缩放因子用于激活值。处理不同token激活值幅度差异很大的情况。
|
||||
|
||||
## 训练后量化（PTQ）
|
||||
|
||||
- **PTQ**量化预训练模型而不需要重新训练。通过**校准集**（一个小的代表性数据集，通常128-512个样本）输入模型收集激活值统计信息，然后计算最优缩放因子。
|
||||
|
||||
### 校准方法
|
||||
|
||||
- **最小-最大**：基于观察到的最小值和最大值设置缩放因子。简单但容易受异常值影响（一个极端值将大部分量化范围浪费在很少使用的值上）。
|
||||
|
||||
- **百分位数**：使用99.99百分位数而不是绝对最大值。裁剪极端异常值，为大多数值提供更好的分辨率。裁剪后的值饱和到$q_{\min}$或$q_{\max}$。
|
||||
|
||||
- **MSE最优**：找到最小化原始张量和量化张量之间均方误差的缩放因子。这是一个一维优化（搜索可能的裁剪值），通常给出最好的PTQ精度。
|
||||
|
||||
- **基于熵**（KL散度）：找到最小化原始和量化值分布之间KL散度的缩放因子。用于TensorRT的INT8校准。
|
||||
|
||||
### PTQ实践
|
||||
|
||||
```python
|
||||
# 使用PyTorch的简化PTQ（概念性）
|
||||
import torch
|
||||
|
||||
def quantise_tensor_symmetric(tensor, bits=8):
|
||||
qmax = 2 ** (bits - 1) - 1 # INT8的127
|
||||
scale = tensor.abs().max() / qmax
|
||||
quantised = torch.clamp(torch.round(tensor / scale), -qmax, qmax).to(torch.int8)
|
||||
return quantised, scale
|
||||
|
||||
def dequantise(quantised, scale):
|
||||
return quantised.float() * scale
|
||||
|
||||
# 量化一个权重矩阵
|
||||
weight = torch.randn(512, 512) # 预训练权重
|
||||
weight_q, scale = quantise_tensor_symmetric(weight, bits=8)
|
||||
weight_reconstructed = dequantise(weight_q, scale)
|
||||
|
||||
# 量化误差
|
||||
error = (weight - weight_reconstructed).abs().mean()
|
||||
print(f"平均绝对误差: {error:.6f}")
|
||||
print(f"压缩比: {weight.numel() * 4 / (weight_q.numel() * 1 + 4):.1f}x") # +4字节用于缩放因子
|
||||
```
|
||||
|
||||
- PTQ在INT8上对大多数模型效果良好，精度下降<1%。对于INT4，PTQ质量显著下降——仅权重量化方法（见下文）处理INT4要好得多。
|
||||
|
||||
## 量化感知训练（QAT）
|
||||
|
||||
- **QAT**在训练图中插入伪量化操作：在前向传播中，权重和激活值被量化和反量化，但梯度像没有量化一样流过（**直通估计器**）。
|
||||
|
||||
$$\text{前向: } \hat{W} = \text{反量化}(\text{量化}(W))$$
|
||||
$$\text{反向: } \frac{\partial L}{\partial W} \approx \frac{\partial L}{\partial \hat{W}}$$
|
||||
|
||||
- 模型在训练过程中学会了抵抗量化噪声。QAT通常能恢复PTQ损失的全部或大部分精度，特别是在低位宽（INT4、INT2）下。
|
||||
|
||||
- **成本**：QAT需要重新训练（或微调）模型，这对大模型来说成本高昂。对于一个70B参数模型，QAT可能需要$10,000-$100,000的计算成本。PTQ基本上零成本（只需校准）。
|
||||
|
||||
- **何时使用QAT**：当PTQ质量不可接受时（通常是INT4或更低），当部署到有严格延迟预算的边缘设备时，或者当模型将被量化数百万次时（一次性QAT成本被摊销）。
|
||||
|
||||
## 仅权重量化
|
||||
|
||||
- 对于LLM推理，瓶颈是从内存加载权重，而不是计算（内存带宽受限模式）。**仅权重量化**将权重量化为INT4或INT3，而保持激活值为FP16。计算在FP16中进行（在运行时反量化权重），但内存消耗和带宽减少了4-8倍。
|
||||
|
||||
### GPTQ
|
||||
|
||||
- **GPTQ**（Frantar等人，2022）一次量化一列权重，通过调整后续列来补偿每列的误差。它使用**Hessian矩阵**（来自校准集的二阶信息）来确定最优量化顺序和误差补偿：
|
||||
|
||||
$$\hat{W}_{:,j} = \text{quant}(W_{:,j}), \quad W_{:,j+1:} \mathrel{-}= \frac{(\hat{W}_{:,j} - W_{:,j}) \cdot H_{j,j+1:}}{H_{j,j}}$$
|
||||
|
||||
- 关键洞察：量化第$j$列会引入误差。GPTQ立即通过调整所有剩余列来补偿，使得该层的整体输出（$XW$）变化尽可能小。这是应用于Transformer的**最优脑量化**（OBQ）。
|
||||
|
||||
- 使用4位组量化（组大小128）的GPTQ在大多数LLM上达到<1%的困惑度降级。在单GPU上，70B模型的量化大约需要1小时。
|
||||
|
||||
### AWQ
|
||||
|
||||
- **AWQ**（激活感知权重量化，Lin等人，2023）观察到一小部分权重通道（1-3%）比其他通道重要得多——它们对应于具有大幅度的激活通道。保护这些显著通道可以大幅降低量化误差。
|
||||
|
||||
- AWQ在量化前将这些重要通道乘以一个因子$s$（使它们变大，因此受舍入影响更小），并将相应的激活值乘以$1/s$（以保持输出不变）。缩放因子$s$按组优化，以最小化整体量化误差。
|
||||
|
||||
- AWQ比GPTQ更简单（无需Hessian计算），运行更快，并达到可比较的质量。它已成为许多开源LLM量化流程的默认选择。
|
||||
|
||||
### GGUF / llama.cpp量化
|
||||
|
||||
- **GGUF**（GGML通用格式）是llama.cpp用于CPU推理的格式。它支持多种量化方案：
|
||||
- **Q4_0**：4位，32元素块，对称。
|
||||
- **Q4_K_M**：4位，带混合精度重要通道（k-quants）。
|
||||
- **Q5_K_M**：5位，带k-quants（更高质量）。
|
||||
- **Q8_0**：8位，简单快速。
|
||||
|
||||
- "K"变体（k-quants）为重要的权重块分配更多位，类似于AWQ的洞察但实现在格式层面。Q4_K_M是大多数模型的最佳选择：平均4位，质量损失最小。
|
||||
|
||||
### QuIP和QuIP#
|
||||
|
||||
- **QuIP**（Chee等人，2023）引入了**非相干处理**：在量化之前使用随机正交变换旋转权重矩阵。这会将信息分散到所有权重上，防止少数异常权重主导量化误差。
|
||||
|
||||
- 直觉：如果一个权重是100，其余的大约是1，用相同的缩放因子量化所有权重会浪费INT4的大部分范围在异常值上。经过正交旋转（保持矩阵的数学性质）后，所有权重具有相似幅度，均匀量化效果更好。
|
||||
|
||||
- **QuIP#** 通过**格点码本**扩展了这一思想：不是映射到均匀整数网格，而是映射到最优格点中的点（8D中的E8格点）。格点编码在相同位数内打包更多量化点，实现了比均匀量化更好的率失真性能。QuIP#在**2位**精度下达到了可用质量——典型INT4方法的一半位数。
|
||||
|
||||
### SpQR
|
||||
|
||||
- **SpQR**（Dettmers等人，2023）观察到极小一部分权重（0.1-1%）是**异常值**，对输出质量的贡献不成比例。SpQR不是将所有内容量化到相同精度，而是：
|
||||
|
||||
1. 使用敏感性分析（量化这个权重会改变层输出多少？）识别异常权重。
|
||||
2. 以**全精度**（FP16）的稀疏格式存储异常值。
|
||||
3. 将所有剩余权重量化为INT3或INT4。
|
||||
|
||||
- 结果：~99%的权重被积极量化（小），而关键的1%保持全精度（准确）。稀疏异常值存储增加的开销最小（占总大小的<5%）。
|
||||
|
||||
### HQQ
|
||||
|
||||
- **HQQ**（半二次量化，Badri & Shaji，2023）是一种**零样本**权重量化方法，完全不需要校准数据。它将量化表述为一个半二次优化问题，迭代求解最优量化权重和缩放因子。
|
||||
|
||||
- 优势：无需校准集意味着没有数据依赖，即时量化，也没有校准数据不匹配的风险。HQQ对于无法获得代表性校准数据或数据敏感型的模型特别有用。
|
||||
|
||||
### AQLM
|
||||
|
||||
- **AQLM**（Egiazarian等人，2024）将**加法量化**（多码本向量量化）应用于LLM。AQLM不是独立量化每个权重，而是将权重分组为向量，并将每个向量表示为来自多个学习到的码本的条目之和：
|
||||
|
||||
$$\mathbf{w} \approx \mathbf{c}_1^{(1)} + \mathbf{c}_2^{(2)} + \cdots + \mathbf{c}_M^{(M)}$$
|
||||
|
||||
- 其中$\mathbf{c}_i^{(m)}$是来自码本$m$的一个条目。有$M = 2$个码本，每个有256个条目，一个8元素向量被编码为两个8位索引 = 8个权重2字节 = 每个权重有效**2位**。AQLM在2位精度下达到了最先进的质量，在这个极限压缩水平上优于GPTQ和AWQ。
|
||||
|
||||
### BitNet和1位LLM
|
||||
|
||||
- **BitNet**（Wang等人，2023）将量化推向极致：权重是三值的（$\{-1, 0, +1\}$），每个权重仅需约1.58位。矩阵乘法变成**只有加法和减法**——不需要浮点乘法。
|
||||
|
||||
- **BitNet b1.58**（Ma等人，2024）将每个权重约束为$\{-1, 0, +1\}$。"1.58位"来自$\log_2(3) \approx 1.58$。在这个精度下，一个70B模型适合约15 GB，推理不需要乘法运算——只需加、减和符号翻转。
|
||||
|
||||
- 矩阵乘法变成：
|
||||
|
||||
$$y_j = \sum_i W_{ij} \cdot x_i = \sum_{i: W_{ij}=+1} x_i - \sum_{i: W_{ij}=-1} x_i$$
|
||||
|
||||
- 这比在任何硬件上的FP16矩阵乘法都要便宜得多，并且可以在没有浮点单元的设备上实现LLM推理。对于当前模型，质量权衡是显著的，但随着规模和训练时量化感知能力的提高而改善。
|
||||
|
||||
### 微缩放（MX）格式
|
||||
|
||||
- **微缩放**（MX）格式是一种新的行业标准（由AMD、Arm、Intel、Meta、Microsoft、NVIDIA、Qualcomm支持），使用**块浮点**：一组元素共享一个指数，每个元素有自己的尾数。
|
||||
|
||||
| 格式 | 共享指数 | 元素位数 | 总计（每元素） | 等价 |
|
||||
|--------|----------------|-------------|--------------------|----|
|
||||
| MXFP8 | 每块8位 | 8（E4M3/E5M2） | ~8 | 类似FP8，范围更好 |
|
||||
| MXFP6 | 每块8位 | 6 | ~6.5 | 介于FP8和INT4之间 |
|
||||
| MXFP4 | 每块8位 | 4 | ~4.5 | 类似INT4，但有浮点行为 |
|
||||
| MXINT8 | 每块8位 | 8（整数） | ~8.5 | INT8，带共享缩放 |
|
||||
|
||||
- 共享指数将指数成本分摊到一个块（通常16-32个元素）。每个元素比单独指数时保留更多尾数位，每位的精度更好。MX格式预计将在未来硬件中替代单独的FP8和INT8格式。
|
||||
|
||||
### FP8训练
|
||||
|
||||
- 在FP8中训练（不仅仅是推理）现在在NVIDIA Hopper和Blackwell GPU上可行。方案如下：
|
||||
|
||||
- **前向传播**：权重和激活值使用E4M3（更高精度，更窄范围）。Transformer Engine使用延迟缩放（跟踪上一次迭代的统计信息，应用于当前迭代）动态计算每张量缩放因子。
|
||||
|
||||
- **反向传播**：梯度使用E5M2（更宽范围，更低精度）。梯度的值范围比权重/激活值更广，因此额外的指数位防止溢出。
|
||||
|
||||
- **主权重**：以FP32维护，用于优化器状态（就像使用FP16的标准混合精度训练，第6章）。FP8计算仅用于矩阵乘法，不用于权重更新。
|
||||
|
||||
- **损失缩放**：FP8仍然需要，就像FP16一样。动态损失缩放器调整缩放因子，使梯度值保持在FP8的可表示范围内。
|
||||
|
||||
- FP8训练在大多数模型规模上达到与BF16训练相当的质量，吞吐量提高约2倍。它是在H100集群上进行新的大规模训练运行的默认选择。
|
||||
|
||||
## 激活值量化
|
||||
|
||||
- 激活值（层之间流动的中间张量）也可以量化，实现完全INT8计算（权重和激活值都是INT8，INT32累加）。
|
||||
|
||||
- **动态量化**：在运行时根据实际激活值计算缩放因子。更准确（适应每个输入），但增加开销（每层计算最小值/最大值或百分位数）。
|
||||
|
||||
- **静态量化**：在校准期间计算一次缩放因子并固定。推理时更快（无需运行时统计），但如果校准数据不具代表性则精度较低。
|
||||
|
||||
- **逐token量化**：为序列中的每个token计算单独的缩放因子。对LLM至关重要，因为不同token的激活值幅度可能差异很大（某些token的激活值比其他token大100倍）。
|
||||
|
||||
- 激活值量化比权重量化更难，因为激活值依赖于数据（它们随每个输入变化），而权重是固定的。"异常值"问题尤其严重：少数激活通道具有极值（平均值的100倍），用与正常通道相同的缩放因子量化它们会浪费精度。
|
||||
|
||||
- **SmoothQuant**（Xiao等人，2022）通过数学上将量化难度从激活值（由于异常值难以量化）迁移到权重（易于量化）来解决异常值问题：将激活值乘以$1/s$，权重乘以$s$，其中$s$平衡难度。输出$XW = (X \cdot \text{diag}(s^{-1})) \cdot (\text{diag}(s) \cdot W)$保持不变。
|
||||
|
||||
## 混合精度量化
|
||||
|
||||
- 并非所有层对量化的敏感度相同。注意力层通常可以容忍INT4，而嵌入层和最终分类器需要更高精度。
|
||||
|
||||
- **敏感性分析**：逐层量化并测量精度影响。敏感性高的层获得更多位；不敏感的层获得更少位。
|
||||
|
||||
- Transformer Engine（第16章，NVIDIA Hopper）在操作级别实现动态混合精度：每个矩阵乘法根据张量统计信息在FP8和FP16之间选择，最大化吞吐量同时保持质量。
|
||||
|
||||
## KV缓存量化
|
||||
|
||||
- 在LLM生成过程中，**KV缓存**存储所有先前token的键和值张量。对于长序列，这主导了内存：
|
||||
|
||||
$$\text{KV缓存大小} = 2 \times n_{\text{layers}} \times n_{\text{heads}} \times d_{\text{head}} \times \text{seq\_len} \times \text{bytes\_per\_element}$$
|
||||
|
||||
- 一个70B模型，80层，64头，128维头，序列长度128K，FP16：$2 \times 80 \times 64 \times 128 \times 131072 \times 2 = 330$ GB。这超过了GPU内存。
|
||||
|
||||
- **KV缓存量化**通过将缓存的键和值以INT8或INT4而不是FP16存储来减少内存。量化误差在序列中累积（每个新token关注所有缓存的K/V），但使用逐通道或逐头量化后，降级是可以接受的。
|
||||
|
||||
- **KV缓存量化具有乘法级收益**：它支持更长的序列（更多上下文）、更大的批次大小（更多并发用户）和更快的推理（加载缓存所需的内存带宽更少）。这是LLM服务中影响最大的优化之一。
|
||||
|
||||
## 编程任务（使用CoLab或notebook）
|
||||
|
||||
1. 从头实现对称INT8量化。量化一个权重矩阵，反量化它，并测量作为值分布函数的重建误差。
|
||||
```python
|
||||
import jax.numpy as jnp
|
||||
import jax
|
||||
|
||||
def quantise_int8(tensor):
|
||||
scale = jnp.max(jnp.abs(tensor)) / 127.0
|
||||
quantised = jnp.clip(jnp.round(tensor / scale), -127, 127).astype(jnp.int8)
|
||||
return quantised, scale
|
||||
|
||||
def dequantise(quantised, scale):
|
||||
return quantised.astype(jnp.float32) * scale
|
||||
|
||||
# 正常权重（典型训练模型）
|
||||
key = jax.random.PRNGKey(0)
|
||||
weights = jax.random.normal(key, (1024, 1024)) * 0.02
|
||||
|
||||
q, s = quantise_int8(weights)
|
||||
recon = dequantise(q, s)
|
||||
|
||||
print(f"原始: {weights.nbytes / 1024:.0f} KB")
|
||||
print(f"量化后: {q.nbytes / 1024:.0f} KB ({weights.nbytes / q.nbytes:.0f}x 更小)")
|
||||
print(f"平均绝对误差: {jnp.abs(weights - recon).mean():.6f}")
|
||||
print(f"最大绝对误差: {jnp.abs(weights - recon).max():.6f}")
|
||||
print(f"相对误差: {jnp.abs(weights - recon).mean() / jnp.abs(weights).mean():.4%}")
|
||||
```
|
||||
|
||||
2. 演示异常值问题。创建具有几个极端通道的激活值，展示逐张量量化失败而逐通道量化成功。
|
||||
```python
|
||||
import jax.numpy as jnp
|
||||
import jax
|
||||
|
||||
key = jax.random.PRNGKey(42)
|
||||
|
||||
# 激活值：大多数通道正常，2个通道有100x异常值
|
||||
activations = jax.random.normal(key, (32, 512)) * 0.1
|
||||
activations = activations.at[:, 0].set(activations[:, 0] * 100) # 异常通道
|
||||
activations = activations.at[:, 1].set(activations[:, 1] * 50) # 异常通道
|
||||
|
||||
# 逐张量量化（整个张量一个缩放因子）
|
||||
scale_tensor = jnp.max(jnp.abs(activations)) / 127.0
|
||||
q_tensor = jnp.clip(jnp.round(activations / scale_tensor), -127, 127)
|
||||
recon_tensor = q_tensor * scale_tensor
|
||||
|
||||
# 逐通道量化（每通道一个缩放因子）
|
||||
scales_channel = jnp.max(jnp.abs(activations), axis=0) / 127.0
|
||||
q_channel = jnp.clip(jnp.round(activations / scales_channel), -127, 127)
|
||||
recon_channel = q_channel * scales_channel
|
||||
|
||||
err_tensor = jnp.abs(activations - recon_tensor).mean()
|
||||
err_channel = jnp.abs(activations - recon_channel).mean()
|
||||
|
||||
print(f"逐张量误差: {err_tensor:.6f}")
|
||||
print(f"逐通道误差: {err_channel:.6f}")
|
||||
print(f"逐通道好 {err_tensor / err_channel:.1f}x")
|
||||
print(f"\n异常通道浪费了 {(activations.shape[1] - 2) / activations.shape[1]:.0%} "
|
||||
f"的量化范围给 {2 / activations.shape[1]:.1%} 的通道")
|
||||
```
|
||||
|
||||
3. 计算不同模型大小和序列长度的KV缓存内存。展示为什么KV缓存量化对长上下文模型至关重要。
|
||||
```python
|
||||
def kv_cache_gb(n_layers, n_heads, d_head, seq_len, bytes_per_elem):
|
||||
return 2 * n_layers * n_heads * d_head * seq_len * bytes_per_elem / 1e9
|
||||
|
||||
models = [
|
||||
("Llama-7B", 32, 32, 128),
|
||||
("Llama-70B", 80, 64, 128),
|
||||
("GPT-4 (估计)", 120, 96, 128),
|
||||
]
|
||||
|
||||
print(f"{'模型':<15} {'序列长度':>8} {'FP16 (GB)':>10} {'INT8 (GB)':>10} {'INT4 (GB)':>10}")
|
||||
print("-" * 60)
|
||||
|
||||
for name, layers, heads, d_head in models:
|
||||
for seq_len in [4096, 32768, 131072]:
|
||||
fp16 = kv_cache_gb(layers, heads, d_head, seq_len, 2)
|
||||
int8 = kv_cache_gb(layers, heads, d_head, seq_len, 1)
|
||||
int4 = kv_cache_gb(layers, heads, d_head, seq_len, 0.5)
|
||||
print(f"{name:<15} {seq_len:>8} {fp16:>9.1f} {int8:>9.1f} {int4:>9.1f}")
|
||||
print()
|
||||
```
|
||||
@@ -0,0 +1,238 @@
|
||||
# 高效架构
|
||||
|
||||
*让模型更快不仅仅是降低精度，还在于设计更智能的架构，使每个token的计算量更少。本文涵盖StreamingLLM、稀疏和线性注意力、多查询和分组查询注意力、推理时的混合专家、知识蒸馏、剪枝和神经架构搜索*
|
||||
|
||||
- 量化（文件1）使每个操作更廉价。本文从源头上减少操作数量。两者互补：一个架构高效且量化的模型可以比原始模型快10-100倍。
|
||||
|
||||
## StreamingLLM：无限长度生成
|
||||
|
||||
- 标准Transformer将所有先前的token存储在KV缓存中，KV缓存随序列长度线性增长。在某一点上，缓存超过GPU内存，生成失败。**StreamingLLM**（Xiao等人，2023）使用固定大小的**滚动KV缓存**解决了这个问题。
|
||||
|
||||
- 关键观察：序列中的前几个token，无论其内容如何，都获得不成比例的高注意力分数。这些被称为**注意力汇聚点**。如果将它们从缓存中逐出，注意力分布会崩溃，生成质量灾难性下降。
|
||||
|
||||
- StreamingLLM的解决方案：在缓存中永久保留少量**汇聚token**（前1-4个token），加上最近$w$个token的**滚动窗口**。总缓存大小为$\text{sink} + w$，无论生成了多少token都是固定的。
|
||||
|
||||
$$\text{缓存} = [\text{token}_0, \text{token}_1, \text{token}_{t-w+1}, \ldots, \text{token}_t]$$
|
||||
|
||||
- 注意力汇聚点锚定softmax分布，滚动窗口提供最近的上下文。这实现了**无限长度生成**，内存恒定，代价是失去了访问序列中间上下文的能力。
|
||||
|
||||
- StreamingLLM无需重新训练即可用于自然形成注意力汇聚点的模型（大多数预训练LLM都会）。对于不形成汇聚点的模型，在训练期间添加一个可学习的汇聚token即可解决。
|
||||
|
||||
## 稀疏注意力
|
||||
|
||||
- 全自注意力在序列长度$n$上是$O(n^2)$，因为每个token关注所有其他token。对于$n = 128K$，注意力矩阵有$128K^2 = 160$亿个条目。**稀疏注意力**模式通过限制哪些token关注哪些token来减少这个数量。
|
||||
|
||||
|
||||
|
||||
- **滑动窗口注意力**（Mistral、Gemma）：每个token只关注之前$w$个token（例如$w = 4096$）。注意力是$O(n \cdot w)$而不是$O(n^2)$。信息通过多层在窗口之外传播：经过$L$层后，有效上下文为$L \times w$。
|
||||
|
||||
- **局部+全局注意力**（Longformer、BigBird）：大多数token使用滑动窗口注意力（局部），但少数指定token（例如[CLS]，每512个token）关注所有token（全局）。这同时捕获了局部模式和长距离依赖。
|
||||
|
||||
- **膨胀注意力**：关注窗口内每第$k$个token，创建一个覆盖更大范围但注意力分数数量相同的稀疏模式。跨层增加膨胀度创建类似于膨胀卷积的层次结构（第8章）。
|
||||
|
||||
- 现代LLM的实际胜者是**滑动窗口+全注意力交错**：某些层使用滑动窗口（廉价，处理局部上下文），某些层使用全注意力（昂贵，捕获长距离）。Mistral/Mixtral使用这种模式。
|
||||
|
||||
## 线性注意力和状态空间模型
|
||||
|
||||
- 我们能完全替换$O(n^2)$的注意力吗？**线性注意力**和**状态空间模型（SSM）**通过避免显式注意力矩阵，以$O(n)$时间处理序列。
|
||||
|
||||
- **线性注意力**用核近似替换softmax注意力：
|
||||
|
||||
$$\text{标准: } O = \text{softmax}(QK^T / \sqrt{d}) V$$
|
||||
$$\text{线性: } O = \phi(Q) (\phi(K)^T V)$$
|
||||
|
||||
- 通过先关联$K^T V$乘积（这是$d \times d$，与序列长度无关），计算变成$O(n \cdot d^2)$而不是$O(n^2 \cdot d)$。对于$n \gg d$的长序列，这是巨大的节省。
|
||||
|
||||
- **RWKV**结合了RNN和Transformer的思想。它使用循环公式顺序处理token（像RNN），但可以在训练时并行化（像Transformer）。推理是每个token $O(1)$（常量内存，KV缓存不增长）。
|
||||
|
||||
- **Mamba**（Gu & Dao，2023）是一种选择性状态空间模型。它通过学习到的状态转换处理序列：
|
||||
|
||||
$$h_t = \bar{A} h_{t-1} + \bar{B} x_t, \quad y_t = C h_t$$
|
||||
|
||||
- 其中$\bar{A}$和$\bar{B}$是依赖于输入的（选择性），允许Mamba动态关注或忽略输入的部分。与固定SSM不同，选择性使Mamba在语言任务上与Transformer具有竞争力，同时保持$O(n)$的扩展性。
|
||||
|
||||
- **权衡**：线性注意力和SSM对长序列更快，但对于需要精确长距离检索的任务，通常不如全注意力。混合架构（一些Transformer层+一些Mamba层）通常能提供两全其美的效果。
|
||||
|
||||
## 多查询和分组查询注意力
|
||||
|
||||
- 标准多头注意力（MHA，第7章）为每个头使用独立的$K$、$V$投影。对于$h$个head，KV缓存中有$h$个独立的键和值张量。**多查询注意力（MQA）**和**分组查询注意力（GQA）**减少了这个数量。
|
||||
|
||||
- **MQA**（Shazeer，2019）：所有头共享单组$K, V$投影。每个头仍然有自己的$Q$投影。KV缓存缩小了$h$倍（例如，32个头则缩小32倍）。
|
||||
|
||||
- **GQA**（Ainslie等人，2023）：一个中间方案。头被分组，每组共享一组$K, V$投影。有$h = 32$个头和$g = 8$个组，每组4个头共享K/V。KV缓存缩小了$h/g = 4$倍。
|
||||
|
||||
$$\text{MHA: } h \text{ 个头, } h \text{ 个 K/V 集} \quad \to \quad \text{GQA: } h \text{ 个头, } g \text{ 个 K/V 集} \quad \to \quad \text{MQA: } h \text{ 个头, } 1 \text{ 个 K/V 集}$$
|
||||
|
||||
|
||||
|
||||
- 大多数现代LLM使用GQA（Llama 2/3、Gemma、Mistral）。它减少了KV缓存内存和推理延迟，与MHA相比质量损失可以忽略不计。
|
||||
|
||||
### 多头潜在注意力（MLA）
|
||||
|
||||
- **MLA**（DeepSeek-V2，2024）通过将KV缓存压缩为**低秩潜在空间**，比GQA更进一步。MLA不是缓存完整的键和值向量，而是每个token缓存一个压缩后的潜在向量$\mathbf{c}_t$，并在注意力期间动态重构K/V：
|
||||
|
||||
$$\mathbf{c}_t = W_{\text{compress}} \cdot [\mathbf{k}_t; \mathbf{v}_t], \quad \mathbf{k}_t = W_K^{\text{up}} \cdot \mathbf{c}_t, \quad \mathbf{v}_t = W_V^{\text{up}} \cdot \mathbf{c}_t$$
|
||||
|
||||
- 压缩向量$\mathbf{c}_t$比原始K和V的组合小得多。DeepSeek-V2实现了与MHA相比**93.3%的KV缓存大小减少**，甚至优于MQA，同时保持MHA级别的质量。
|
||||
|
||||
- 权衡：从潜在向量重构K/V在每个注意力操作中增加了少量计算成本。但由于LLM解码是内存带宽受限的（而非计算受限），这总体上是个净收益：更少的内存加载 > 每token稍多计算。
|
||||
|
||||
### Flash Attention
|
||||
|
||||
- **Flash Attention**（Dao等人，2022，第16章文件05有详细论述）不是架构变化，而是一种实现优化，在任何高效注意力讨论中都不可或缺。它计算精确的标准注意力，具有以下特点：
|
||||
|
||||
- **O(n)内存**而不是O(n²)（注意力矩阵从未在HBM中具体化）。
|
||||
- **比标准注意力快2-4倍**（通过分块和在线softmax将数据保留在SRAM中）。
|
||||
- **无质量损失**——输出在数学上与标准注意力完全相同。
|
||||
|
||||
- Flash Attention现在是PyTorch（`torch.nn.functional.scaled_dot_product_attention`）、JAX和所有主要推理框架中默认的注意力实现。如果你在2024+年运行注意力，你几乎肯定在使用Flash Attention。
|
||||
|
||||
### Ring Attention
|
||||
|
||||
- **Ring Attention**（Liu等人，2023）将注意力计算分布到多个设备上，用于即使使用Flash Attention也无法装入单GPU内存的长序列。
|
||||
|
||||
- 思路：将序列分区到$N$个设备上。每个设备持有$n/N$个token的Q、K、V。设备排列成环形。每一步：
|
||||
1. 每个设备计算局部注意力（其Q对其局部K/V）。
|
||||
2. 每个设备将其K/V块发送到环中的下一个设备。
|
||||
3. 每个设备从上一个设备接收K/V，并针对这些K/V计算注意力。
|
||||
4. 经过$N$步后，每个设备已经关注过每个K/V块。
|
||||
|
||||
- 通信与计算**重叠**：在当前K/V块上计算注意力的同时，下一个块正在传输中。这几乎完全隐藏了通信延迟。
|
||||
|
||||
- Ring Attention通过将KV缓存分布在一圈GPU上，实现了**百万token上下文窗口**。每台设备的内存为O(n/N)，使得任意长序列都可行（仅受设备数量限制）。
|
||||
|
||||
## 推理时的混合专家
|
||||
|
||||
- MoE模型（第7章）每个token只激活其参数的一小部分（通常8个专家中的2个）。在推理时，独特的挑战是**专家缓存**：所有专家都必须在内存中（因为任何token可能路由到任何专家），但每个token只有2个活跃。
|
||||
|
||||
- 对于Mixtral 8x7B模型：总参数 = 47B（8 × 7B专家，但有共享组件）。每个token的活跃参数 ≈ 13B（2个专家 + 共享层）。该模型具有LLM-70B级别的质量，但推理成本为LLM-13B级别，不过需要在内存中保留47B参数。
|
||||
|
||||
- **专家卸载**：对于GPU内存受限的部署，将非活跃专家保留在CPU或SSD上，按需加载。这之所以有效，是因为token路由足够可预测，可以预取可能的专家。
|
||||
|
||||
- **专家缓存**：在GPU内存中维护最近使用的专家的LRU缓存。如果相同的专家被重复激活（在领域内数据中常见），缓存命中率很高。
|
||||
|
||||
## 知识蒸馏
|
||||
|
||||
- **蒸馏**（第6章）训练一个小的"学生"模型来模仿一个大的"教师"。学生从教师的软预测（类上的概率分布）中学习，这比单独的硬标签包含更多信息。
|
||||
|
||||
$$\mathcal{L} = \alpha \cdot \text{KL}(p_{\text{teacher}}^{T} \| p_{\text{student}}^{T}) + (1 - \alpha) \cdot \mathcal{L}_{\text{CE}}(y, p_{\text{student}})$$
|
||||
|
||||
- 其中$T$是温度（更高的$T$使分布变软，揭示教师的不确定性），$\alpha$平衡蒸馏损失与标准交叉熵损失。
|
||||
|
||||
- **对于LLM**：蒸馏用于从大型、能力强的模型创建小型、快速的模型。GPT-4 → 一个7B学生模型，在特定任务上捕获GPT-4的大部分行为。学生模型的推理成本可以低10-100倍。
|
||||
|
||||
- **任务特定蒸馏**：仅在与部署任务相关的数据上进行蒸馏。从70B教师模型在医疗问答上蒸馏出的7B模型，在该特定任务上可以超越70B模型（因为学生有限的容量完全集中在目标领域上）。
|
||||
|
||||
## 剪枝
|
||||
|
||||
- **剪枝**移除不必要的权重（将其设为零），减少模型大小和计算量。
|
||||
|
||||
- **非结构化剪枝**（基于幅值）：移除绝对值最小的单个权重。这创建了稀疏权重矩阵。简单有效用于压缩，但当前硬件（GPU）除非稀疏性遵循特定模式，否则无法高效加速稀疏操作。
|
||||
|
||||
- **结构化剪枝**：移除整个单元——注意力头、MLP神经元或层。这产生一个更小的稠密模型，可以在标准硬件上直接加速。权衡是粒度更粗（移除一个完整的头可能同时移除了有用和无用的权重）。
|
||||
|
||||
- **2:4稀疏性**（NVIDIA Ampere+）：一种硬件支持的稀疏模式，每4个权重中有2个为零。GPU的稀疏Tensor Core跳过零乘法，实现约2倍加速。这是目前唯一具有实际硬件加速的稀疏模式。
|
||||
|
||||
- **彩票假说**（Frankle & Carlin，2019）：在随机初始化的网络中，存在一个子网络（"中奖彩票"），可以单独训练以匹配完整网络的性能。找到这些子网络（通过训练、剪枝和重置）成本高昂，但这个洞察激励了剪枝研究。
|
||||
|
||||
## 神经架构搜索（NAS）
|
||||
|
||||
- **NAS**通过搜索可能的架构空间来自动化架构设计，找到在硬件约束（延迟、内存、功耗）下最大化精度的架构。
|
||||
|
||||
- **EfficientNet**（第8章）就是通过NAS找到的：复合缩放规则（平衡深度、宽度、分辨率）是从搜索中涌现的，而非人类直觉。
|
||||
|
||||
- 对于推理效率，NAS可以找到针对特定硬件目标优化的架构："找到一个在iPhone神经引擎上延迟<5ms且在ImageNet上精度>80%的模型。"搜索空间包括层类型、宽度、激活函数和注意力模式。
|
||||
|
||||
- **一次性网络**训练一个单个过参数化网络，为不同的部署目标提取子网络。一次训练运行产生针对云GPU、移动GPU和CPU优化的模型，每个都针对其目标进行了优化。
|
||||
|
||||
## 编程任务（使用CoLab或notebook）
|
||||
|
||||
1. 实现滑动窗口注意力，并与全注意力比较内存使用。
|
||||
```python
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
|
||||
def full_attention(Q, K, V):
|
||||
"""标准O(n²)注意力。"""
|
||||
scores = Q @ K.T / jnp.sqrt(Q.shape[-1])
|
||||
weights = jax.nn.softmax(scores, axis=-1)
|
||||
return weights @ V
|
||||
|
||||
def sliding_window_attention(Q, K, V, window_size=128):
|
||||
"""滑动窗口注意力：每个token关注前window_size个token。"""
|
||||
n = Q.shape[0]
|
||||
d = Q.shape[-1]
|
||||
output = jnp.zeros_like(Q)
|
||||
|
||||
for i in range(n):
|
||||
start = max(0, i - window_size + 1)
|
||||
k_window = K[start:i+1]
|
||||
v_window = V[start:i+1]
|
||||
scores = Q[i] @ k_window.T / jnp.sqrt(d)
|
||||
weights = jax.nn.softmax(scores)
|
||||
output = output.at[i].set(weights @ v_window)
|
||||
|
||||
return output
|
||||
|
||||
n, d = 512, 64
|
||||
key = jax.random.PRNGKey(0)
|
||||
Q = jax.random.normal(key, (n, d))
|
||||
K = jax.random.normal(jax.random.PRNGKey(1), (n, d))
|
||||
V = jax.random.normal(jax.random.PRNGKey(2), (n, d))
|
||||
|
||||
print(f"全注意力内存: O(n²) = {n*n} 个条目")
|
||||
print(f"窗口 (w=128) 内存: O(n*w) = {n*128} 个条目")
|
||||
print(f"减少: {n*n / (n*128):.1f}x")
|
||||
```
|
||||
|
||||
2. 比较MHA、GQA和MQA的KV缓存大小。展示为什么GQA是实际的最佳选择。
|
||||
```python
|
||||
def kv_cache_size(n_heads, n_kv_heads, d_head, seq_len, bytes=2):
|
||||
"""KV缓存大小（MB）。"""
|
||||
return 2 * n_kv_heads * d_head * seq_len * bytes / 1e6
|
||||
|
||||
n_heads = 32
|
||||
d_head = 128
|
||||
seq_len = 32768
|
||||
|
||||
mha = kv_cache_size(n_heads, n_heads, d_head, seq_len) # 32个KV头
|
||||
gqa = kv_cache_size(n_heads, 8, d_head, seq_len) # 8个KV头
|
||||
mqa = kv_cache_size(n_heads, 1, d_head, seq_len) # 1个KV头
|
||||
|
||||
print(f"MHA (32个KV头): {mha:.0f} MB 每层")
|
||||
print(f"GQA (8个KV头): {gqa:.0f} MB 每层 ({mha/gqa:.0f}x 更小)")
|
||||
print(f"MQA (1个KV头): {mqa:.0f} MB 每层 ({mha/mqa:.0f}x 更小)")
|
||||
```
|
||||
|
||||
3. 通过从随机注意力层中移除最不重要的注意力头并测量输出变化来模拟结构化剪枝。
|
||||
```python
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
|
||||
key = jax.random.PRNGKey(0)
|
||||
n_heads, seq_len, d_head = 8, 64, 32
|
||||
|
||||
# 随机多头注意力输出（每个头一个）
|
||||
head_outputs = jax.random.normal(key, (n_heads, seq_len, d_head))
|
||||
|
||||
# 完整输出：连接所有头
|
||||
full_output = head_outputs.reshape(seq_len, n_heads * d_head)
|
||||
|
||||
# 重要性：通过范数度量每个头的贡献
|
||||
head_norms = jnp.linalg.norm(head_outputs, axis=(1, 2))
|
||||
print("头重要性（按范数）:", jnp.round(head_norms, 2))
|
||||
|
||||
# 剪枝最不重要的头
|
||||
for n_keep in [8, 6, 4, 2]:
|
||||
top_heads = jnp.argsort(head_norms)[-n_keep:]
|
||||
pruned = head_outputs[top_heads].reshape(seq_len, n_keep * d_head)
|
||||
|
||||
# 填充到原始大小用于比较（将剪掉的头设为零）
|
||||
full_pruned = jnp.zeros_like(head_outputs)
|
||||
full_pruned = full_pruned.at[top_heads].set(head_outputs[top_heads])
|
||||
full_pruned = full_pruned.reshape(seq_len, n_heads * d_head)
|
||||
|
||||
error = jnp.linalg.norm(full_output - full_pruned) / jnp.linalg.norm(full_output)
|
||||
print(f"保留 {n_keep}/{n_heads} 个头: 相对误差 = {error:.4f}, "
|
||||
f"内存 = {n_keep/n_heads:.0%}")
|
||||
```
|
||||
@@ -0,0 +1,236 @@
|
||||
# 服务与批处理
|
||||
|
||||
*向数千并发用户提供LLM服务需要的不只是加载模型和运行推理。本文涵盖预填充-解码分离、连续批处理、PagedAttention和vLLM、调度策略、分离式服务、多模型和LoRA服务，以及关键指标*
|
||||
|
||||
- 单个LLM推理请求很简单：输入token，生成输出token。但要向10,000个并发用户提供低延迟、高吞吐量的LLM服务，这是一个系统工程问题。朴素方法（一次处理一个请求）浪费了90%以上的GPU容量。智能批处理和调度可以在不增加硬件的情况下将吞吐量提高10-50倍。
|
||||
|
||||
## 预填充 vs 解码：两个截然不同的阶段
|
||||
|
||||
- LLM推理有两个不同的阶段，具有根本不同的计算特征：
|
||||
|
||||
- **预填充**（提示处理）：同时处理所有输入token。这是一个单次大规模矩阵乘法：$O(\text{prompt\_length} \times d_{\text{model}}^2)$。提示可以并行处理（所有token都已知）。预填充是**计算受限**的：GPU的ALU是瓶颈。
|
||||
|
||||
- **解码**（token生成）：自回归地一次生成一个token。每个新token需要通过KV缓存关注所有先前的token。解码是**内存带宽受限**的：GPU大部分时间花在从内存加载模型权重和KV缓存上，而不是计算。每个解码步骤只产生一个token，但必须加载整个模型（70B FP16模型约140 GB）。
|
||||
|
||||
- 含义：
|
||||
|
||||
| | 预填充 | 解码 |
|
||||
|--|---------|--------|
|
||||
| 处理的token | 一次性全部（并行） | 一次一个（顺序） |
|
||||
| 瓶颈 | 计算（FLOPS） | 内存带宽 |
|
||||
| 算术强度 | 高 | 非常低 |
|
||||
| GPU利用率 | 高（50-80%） | 低（1-10%），无批处理时 |
|
||||
| 延迟指标 | **首token时间（TTFT）** | **每输出token时间（TPOT）** |
|
||||
|
||||
- TTFT影响用户体验（多久直到响应开始流式传输）。TPOT决定感知的生成速度。用户可以容忍较高的TTFT（1-5秒），但期望快速的TPOT（对话应用每token 30-100毫秒）。
|
||||
|
||||
## 静态批处理（朴素方法）
|
||||
|
||||
- 最简单的批处理：收集$B$个请求，填充到相同长度，作为单个批次处理。
|
||||
|
||||
- **问题1**：请求有不同的提示长度，并生成不同数量的输出token。短请求提前完成，但必须等待批次中最长的请求完成后才能开始下一个批次。GPU在为剩余的一个长请求生成token时处于空闲状态。
|
||||
|
||||
- **问题2**：填充浪费计算。如果最长提示是2000个token，最短是50个，批次被填充到2000。GPU为短请求处理了1950个填充token——纯属浪费。
|
||||
|
||||
|
||||
|
||||
## 连续批处理
|
||||
|
||||
- **连续批处理**（也称为迭代级批处理）通过在单个解码步骤的粒度上操作来解决这两个问题，而不是整个请求。
|
||||
|
||||
- 在每个解码步骤：
|
||||
1. 所有进行中的请求并行生成一个token（作为一个批次）。
|
||||
2. 完成的请求（生成EOS token）立即从批次中**移除**。
|
||||
3. 队列中的新请求立即**插入**到释放的槽位中。
|
||||
|
||||
- 批次大小每步动态变化。GPU从不等候落后者，也没有浪费的填充（每个请求只使用它需要的槽位）。
|
||||
|
||||
- **影响**：连续批处理通常比静态批处理提高吞吐量2-10倍，模型质量不变且延迟无明显增加。
|
||||
|
||||
## PagedAttention和vLLM
|
||||
|
||||
- KV缓存造成了一个内存管理噩梦。每个请求都有一个随着每个生成的token而增长的KV缓存。不同请求处于不同阶段（不同缓存大小）。为每个请求分配连续内存浪费空间（必须为最大可能长度分配，即使请求只生成几个token）。
|
||||
|
||||
|
||||
|
||||
- **PagedAttention**（Kwon等人，2023）将操作系统虚拟内存的概念（第13章）应用于KV缓存。缓存被划分为固定大小的**页**（token位置的块）。页按需分配，在物理GPU内存中可以是非连续的。
|
||||
|
||||
- 优势：
|
||||
- **无碎片**：页大小统一，因此请求之间没有浪费内存的"空洞"。
|
||||
- **惰性分配**：仅在token实际生成时分配内存，而不是预分配最大长度。
|
||||
- **写时复制**：共享共同前缀（例如系统提示）的请求共享相同的KV缓存页。仅当请求分叉时才复制页。
|
||||
|
||||
- **vLLM**是基于PagedAttention构建的推理引擎。通过几乎消除KV缓存内存浪费，它实现了比静态分配服务（如没有分页注意力的HuggingFace text-generation-inference）高2-4倍的吞吐量。
|
||||
|
||||
## 调度策略
|
||||
|
||||
- 当多个请求在等待且GPU只能处理有限批次时，**调度**决定服务哪些请求：
|
||||
|
||||
- **先来先服务（FCFS）**：按到达顺序处理请求。简单但不公平：一个提交10K-token生成的用户会阻塞所有后面的用户。
|
||||
|
||||
- **最短作业优先（SJF）**：处理最先完成的请求。最小化平均延迟，但惩罚长时间运行的请求（它们可能被饿死）。在实践中，估计输出长度未知，因此SJF使用启发式方法（提示长度、用户历史）。
|
||||
|
||||
- **抢占**：如果高优先级请求到达，暂停低优先级的进行中请求（将其KV缓存交换到CPU内存或SSD），服务高优先级请求，然后恢复暂停的请求。vLLM支持此功能。
|
||||
|
||||
- **基于优先级**：为用户或请求类型分配优先级。实时交互查询比批处理作业获得更高优先级。结合抢占，这确保高优先级流量的延迟SLO。
|
||||
|
||||
- **Token预算**：限制活跃批次中的总token数。这防止少量长请求独占GPU内存并饿死新请求。
|
||||
|
||||
## 分离式服务
|
||||
|
||||
- 预填充和解码具有相反的计算特征。在同一GPU上运行两者意味着GPU在计算受限（预填充）和内存带宽受限（解码）之间交替，从未充分利用任一资源。
|
||||
|
||||
- **分离式服务**将它们分开：
|
||||
- **预填充节点**：为计算优化的GPU（高FLOPS，可能内存较少）。处理所有传入提示。
|
||||
- **解码节点**：为内存带宽优化的GPU（大KV缓存容量，高内存带宽）。处理所有token生成。
|
||||
|
||||
- 预填充节点计算初始KV缓存并通过NVLink或网络将其发送到解码节点。解码节点使用接收到的缓存生成token。
|
||||
|
||||
- 这是**Mooncake**（月之暗面）的架构，并正在被多个LLM服务团队探索。好处：每个GPU类型与其工作负载特征匹配，提高整体利用率。
|
||||
|
||||
## 多模型和LoRA服务
|
||||
|
||||
- 在生产中，你通常服务多个模型（不同层级的模型大小不同，不同任务的微调变体不同）。
|
||||
|
||||
- **模型复用**：在同一GPU上加载多个模型，将请求路由到相应模型。GPU内存共享：一个40 GB GPU可能同时持有一个13B模型（26 GB）和一个7B模型（14 GB）。
|
||||
|
||||
- **LoRA服务**：不是部署单独的微调模型，而是部署一个基础模型并带有多个**LoRA适配器**（第6章）。每个适配器增加<1%的参数。请求在推理时路由到相应的适配器。
|
||||
|
||||
- **S-LoRA**（Sheng等人，2023）：从一个基础模型服务数千个LoRA适配器。适配器存储在CPU上，按需分页到GPU内存。基础模型的KV缓存和权重被共享；只有小的LoRA矩阵因请求而异。
|
||||
|
||||
- **Punica**（Chen等人，2023）：通过使用自定义CUDA内核在同一批次中为不同请求应用不同的LoRA矩阵，跨不同LoRA适配器对请求进行批处理。这避免了每个请求切换适配器的开销。
|
||||
|
||||
## 受限和引导生成
|
||||
|
||||
- 许多应用需要LLM以特定格式产生输出：有效的JSON、SQL查询、特定语言的代码或遵循模式的响应。**受限生成**保证输出符合语法或模式。
|
||||
|
||||
- **语法受限解码**：在每个解码步骤，屏蔽会违反语法的token。如果到目前为止的输出是`{"name": "Alice", "age":`且语法要求接下来是整数，则屏蔽除数字外的所有token。LLM的概率分布在有效token上重新归一化。
|
||||
|
||||
- **Outlines**（Willard & Louf，2023）：将JSON模式或正则表达式编译成有限状态机（FSM）。在每个解码步骤，FSM确定哪些token是有效的后续。无效token获得概率0。这保证了100%的模式合规，零重试。
|
||||
|
||||
- **SGLang**原生集成受限生成：你用Python指定输出结构，引擎高效处理token掩码和缓存。这与RadixAttention（前缀缓存）结合，使得结构化输出重用缓存的公共前缀。
|
||||
|
||||
- **为什么重要**：没有受限生成，你自由生成然后解析输出，失败时重试。对于复杂JSON模式，重试率通常为10-30%，浪费计算。受限生成完全消除了重试。
|
||||
|
||||
## 请求路由
|
||||
|
||||
- 并非每个查询都需要最大的模型。**请求路由**根据估计的难度将查询定向到不同的模型：
|
||||
|
||||
- **级联**：先尝试小模型。如果小模型的置信度低于阈值（例如，top token的softmax概率<0.8），则升级到更大的模型。简单查询（80%+的流量）由小模型廉价服务；只有困难查询使用昂贵模型。
|
||||
|
||||
- **学习型路由**：训练一个轻量级分类器（或使用小模型的困惑度）来预测查询需要哪个模型层级。将"2+2等于多少？"路由到3B模型，将"解释量子纠缠的数学基础"路由到70B模型。
|
||||
|
||||
- **影响**：如果80%的查询可以由成本低10倍的模型处理，平均每查询成本下降约70%。这是多模型部署中影响最大的成本优化之一。
|
||||
|
||||
- **设备端+云混合路由**：**Cactus**（[github.com/cactus-compute/cactus](https://github.com/cactus-compute/cactus)）在设备级别实现请求路由。它通过自定义ARM SIMD内核在设备端（手机、笔记本电脑、可穿戴设备）运行小模型，并在本地模型置信度低或查询超出设备能力时自动路由到云端模型。应用为两条路径使用OpenAI兼容API——路由是透明的。这是在基础设施级别的级联：第一层是免费的（设备端），第二层花钱（云API）。对于大多数查询简单的应用（助手问答、自动补全、转录），设备端处理覆盖70-90%的流量，边际成本为零。
|
||||
|
||||
## 推理指标
|
||||
|
||||
- 正确的指标取决于用例：
|
||||
|
||||
| 指标 | 测量内容 | 目标（对话式） | 目标（批处理） |
|
||||
|--------|-----------------|------------------------|-----------------|
|
||||
| **TTFT** | 首token时间 | <1 s | 不太重要 |
|
||||
| **TPOT** | 每输出token时间 | <100 ms | 不太重要 |
|
||||
| **吞吐量** | token/秒（总计） | 不太重要 | 最大化 |
|
||||
| **p99延迟** | 最差的1%请求 | <5 s | <30 s |
|
||||
| **每token成本** | $/100万token | 最小化 | 最小化 |
|
||||
| **SLO合规率** | 满足延迟目标的请求百分比 | >99% | >95% |
|
||||
|
||||
- **TTFT vs TPOT权衡**：激进的批处理增加吞吐量（总token数/秒更多），但增加TPOT（每个token耗时更长，因为GPU处理更多请求）。调度策略必须平衡吞吐量（收入）与延迟（用户体验）。
|
||||
|
||||
- **每token成本**是生产的最终指标。它结合了硬件成本（GPU租金）、吞吐量（token/秒）和利用率。运行在50% GPU利用率的系统比100%利用率的系统每token成本高2倍。这就是批处理、调度和PagedAttention如此重要的原因——它们提高了利用率。
|
||||
|
||||
## 编程任务（使用CoLab或notebook）
|
||||
|
||||
1. 模拟连续vs静态批处理并测量吞吐量差异。
|
||||
```python
|
||||
import random
|
||||
import time
|
||||
|
||||
def simulate_static_batching(requests, batch_size=8):
|
||||
"""在固定批次中处理请求。等待所有完成。"""
|
||||
total_tokens = 0
|
||||
total_time = 0
|
||||
|
||||
for i in range(0, len(requests), batch_size):
|
||||
batch = requests[i:i + batch_size]
|
||||
max_len = max(r['output_len'] for r in batch)
|
||||
# 批次中所有请求耗时等于最长请求
|
||||
batch_time = max_len * 0.01 # 每token 10ms
|
||||
total_time += batch_time
|
||||
total_tokens += sum(r['output_len'] for r in batch)
|
||||
|
||||
return total_tokens / total_time # token/秒
|
||||
|
||||
def simulate_continuous_batching(requests, max_batch=8):
|
||||
"""使用连续批处理处理。移除完成请求，添加新请求。"""
|
||||
total_tokens = 0
|
||||
total_time = 0
|
||||
active = []
|
||||
queue = list(requests)
|
||||
|
||||
while active or queue:
|
||||
# 填充批次
|
||||
while len(active) < max_batch and queue:
|
||||
active.append({'remaining': queue.pop(0)['output_len']})
|
||||
|
||||
if not active:
|
||||
break
|
||||
|
||||
# 一个解码步骤：所有活跃请求生成1个token
|
||||
for req in active:
|
||||
req['remaining'] -= 1
|
||||
total_tokens += len(active)
|
||||
total_time += 0.01 # 每步10ms
|
||||
|
||||
# 移除完成的请求
|
||||
active = [r for r in active if r['remaining'] > 0]
|
||||
|
||||
return total_tokens / total_time
|
||||
|
||||
# 生成具有不同输出长度的请求
|
||||
random.seed(42)
|
||||
requests = [{'output_len': random.randint(10, 500)} for _ in range(100)]
|
||||
|
||||
static_tps = simulate_static_batching(requests)
|
||||
continuous_tps = simulate_continuous_batching(requests)
|
||||
|
||||
print(f"静态批处理: {static_tps:.0f} tokens/s")
|
||||
print(f"连续批处理: {continuous_tps:.0f} tokens/s")
|
||||
print(f"加速比: {continuous_tps / static_tps:.1f}x")
|
||||
```
|
||||
|
||||
2. 计算PagedAttention的KV缓存内存节省。比较预分配（最坏情况）vs分页（实际使用）。
|
||||
```python
|
||||
def paged_vs_preallocated(n_requests, max_seq_len, avg_seq_len, page_size, kv_per_token_bytes):
|
||||
"""比较内存使用：预分配vs分页KV缓存。"""
|
||||
# 预分配：每个请求获得max_seq_len个槽位
|
||||
preallocated_gb = n_requests * max_seq_len * kv_per_token_bytes / 1e9
|
||||
|
||||
# 分页：只分配使用的部分（按页粒度）
|
||||
import math
|
||||
avg_pages = math.ceil(avg_seq_len / page_size)
|
||||
paged_gb = n_requests * avg_pages * page_size * kv_per_token_bytes / 1e9
|
||||
|
||||
waste_preallocated = (max_seq_len - avg_seq_len) / max_seq_len
|
||||
waste_paged = (avg_pages * page_size - avg_seq_len) / (avg_pages * page_size)
|
||||
|
||||
print(f"请求数: {n_requests}, 最大序列: {max_seq_len}, 平均序列: {avg_seq_len}")
|
||||
print(f" 预分配: {preallocated_gb:.1f} GB (浪费: {waste_preallocated:.0%})")
|
||||
print(f" 分页: {paged_gb:.1f} GB (浪费: {waste_paged:.0%})")
|
||||
print(f" 节省: {preallocated_gb - paged_gb:.1f} GB ({preallocated_gb/paged_gb:.1f}x)")
|
||||
print()
|
||||
|
||||
# Llama-70B：每层每token约1.3 KB，80层 = 每token约100 KB总计
|
||||
kv_bytes = 100_000
|
||||
|
||||
# 场景1：短请求，大最大值
|
||||
paged_vs_preallocated(256, max_seq_len=4096, avg_seq_len=256, page_size=16, kv_per_token_bytes=kv_bytes)
|
||||
|
||||
# 场景2：不同长度
|
||||
paged_vs_preallocated(256, max_seq_len=8192, avg_seq_len=1024, page_size=16, kv_per_token_bytes=kv_bytes)
|
||||
|
||||
# 场景3：长上下文
|
||||
paged_vs_preallocated(64, max_seq_len=131072, avg_seq_len=16000, page_size=16, kv_per_token_bytes=kv_bytes)
|
||||
```
|
||||
@@ -0,0 +1,212 @@
|
||||
# 边缘推理
|
||||
|
||||
*边缘推理在用户设备（手机、笔记本电脑、物联网传感器）上运行模型，无需将数据发送到云端。本文涵盖边缘限制、模型压缩流水线、设备端运行时、编译器栈、硬件目标（NPU、神经引擎）、设备端LLM、联邦学习和延迟优化*
|
||||
|
||||
- 云端推理需要网络连接，增加延迟（50-200毫秒往返），每次请求花费金钱，并将用户数据发送到第三方服务器。**边缘推理**消除了所有四个问题：模型本地运行，即时响应，每次推理零成本，且数据保持私密。
|
||||
|
||||
- 权衡：边缘设备的计算和内存比数据中心GPU小100-1000倍。使模型在这些约束下运行需要在每个层面进行积极优化。
|
||||
|
||||
- **Cactus**（[github.com/cactus-compute/cactus](https://github.com/cactus-compute/cactus)) 是一个专为移动和可穿戴设备构建的低延迟AI引擎。它在生产中展示了本文涵盖的许多技术：自定义ARM SIMD内核用于注意力和矩阵运算（第16章）、KV缓存量化（第17章文件01）、分块预填充、Apple和Qualcomm芯片上的NPU加速推理、零拷贝内存映射实现10倍更低的RAM使用，以及在设备端计算不足时的自动云回退。Cactus支持跨iOS、Android、macOS和嵌入式Linux的多模态推理（LLM、视觉、语音），并提供Swift、Kotlin、Python、Flutter、React Native和Rust的SDK。其基准测试显示，在M4 Pro上1.2B INT4模型解码达到100 tokens/s，在iPhone 17 Pro上达到48 tokens/s——这是优化边缘推理的具体示例。
|
||||
|
||||
## 边缘约束
|
||||
|
||||
| 资源 | 云GPU（H100） | 笔记本电脑（M4） | 手机（Snapdragon 8 Gen 3） | IoT（ESP32） |
|
||||
|----------|-----------------|-------------|---------------------------|-------------|
|
||||
| 内存 | 80 GB HBM3 | 16-36 GB 统一内存 | 8-12 GB LPDDR5 | 520 KB |
|
||||
| 计算 | 989 TFLOPS（FP8） | 38 TOPS（神经引擎） | 45 TOPS（NPU） | 0.001 TOPS |
|
||||
| 功耗 | 700 W | 15-30 W | 5-10 W | 0.1 W |
|
||||
| 存储 | TB | 256 GB-2 TB | 128-512 GB | 4 MB |
|
||||
|
||||
- 云GPU和手机NPU之间的计算差距约为20倍。GPU和微控制器之间的差距约为1,000,000倍。不同设备需要不同程度的压缩和不同的模型架构。
|
||||
|
||||
## 模型压缩流水线
|
||||
|
||||
- 对于边缘部署，压缩不是单一技术——它是一个按顺序应用的互补技术**流水线**：
|
||||
|
||||
```
|
||||
完整模型（FP32，70B参数）
|
||||
↓ 知识蒸馏 → 更小模型（7B参数）
|
||||
↓ 结构化剪枝 → 移除冗余头/层（4B有效）
|
||||
↓ 量化（INT4） → 4倍更小（2 GB）
|
||||
↓ 编译器优化 → 融合内核，优化内存布局
|
||||
↓ 运行时 → 设备端执行
|
||||
```
|
||||
|
||||
- 每一步减少大小和延迟。顺序很重要：先蒸馏（减少架构），然后剪枝（移除结构），然后量化（降低精度），最后编译（为目标硬件优化）。在量化之后进行蒸馏会试图压缩已经损失质量的模型。
|
||||
|
||||
## 设备端运行时
|
||||
|
||||
- **运行时**加载模型、分配内存并在目标硬件上执行推理。每个平台有其偏好的运行时：
|
||||
|
||||
- **ONNX Runtime**：跨平台（Windows、Linux、macOS、iOS、Android）。支持CPU、GPU（CUDA、DirectML、CoreML、NNAPI）和许多加速器后端。最具可移植性的选项。模型从PyTorch/TensorFlow导出为ONNX格式。
|
||||
|
||||
- **TensorFlow Lite（TFLite）**：Google的边缘运行时。针对ARM CPU和Android NPU优化。二进制文件小巧（约1 MB）。支持INT8和float16。Android部署的标准。
|
||||
|
||||
- **Core ML**：Apple的iOS/macOS运行时。根据模型特征自动使用神经引擎、GPU或CPU。模型使用`coremltools`从PyTorch/TensorFlow转换。与Apple硬件紧密集成（统一内存、神经引擎）。
|
||||
|
||||
- **ExecuTorch**：Meta新推出的设备端PyTorch运行时。专为边缘部署设计，具有提前编译和操作级硬件加速器委派功能。PyTorch Mobile的继任者。
|
||||
|
||||
- **TensorRT**：NVIDIA的GPU推理优化运行时（第15章）。融合层、选择最优内核并自动量化。在NVIDIA GPU上比PyTorch eager模式快2-5倍。
|
||||
|
||||
- **llama.cpp**：用于LLM的单文件C++推理引擎。支持GGUF量化（Q4、Q5、Q8）、CPU（AVX/NEON）、Metal（Apple GPU）、CUDA和Vulkan。在消费级硬件上运行LLM的首选方案。
|
||||
|
||||
## 编译器栈
|
||||
|
||||
- 在高级模型（PyTorch图）和硬件（NPU指令）之间是**编译器栈**，它为特定目标优化模型：
|
||||
|
||||
```
|
||||
PyTorch模型
|
||||
↓ 导出（torch.export、ONNX、TorchScript）
|
||||
图IR（中间表示）
|
||||
↓ 图优化
|
||||
- 常量折叠（编译时计算常量表达式）
|
||||
- 死代码消除（移除未使用的操作）
|
||||
- 算子融合（conv + bn + relu → 单个融合操作）
|
||||
- 布局转换（NCHW → NHWC用于ARM，通道最后）
|
||||
↓ 降级
|
||||
硬件特定IR
|
||||
↓ 后端优化
|
||||
- 分块和循环排序（缓存友好的访问模式）
|
||||
- 向量化（SIMD，第16章）
|
||||
- 内存规划（重用缓冲区以最小化峰值内存）
|
||||
- 内核选择（为每个操作选择最佳实现）
|
||||
↓ 代码生成
|
||||
机器代码 / NPU指令
|
||||
```
|
||||
|
||||
- **算子融合**是影响最大的优化。一个Transformer块约有20个操作（矩阵乘、加法、层归一化、softmax等）。没有融合，每个操作将其输出写入内存，下一个操作再读回。有了融合，多个操作组合成一个内核，将数据保留在寄存器/缓存中。这可以使速度快2-5倍（第16章，屋顶模型）。
|
||||
|
||||
- **内存规划**：编译器分析模型图以确定哪些张量的生命周期重叠，可以共享相同的内存缓冲区。一个有100个中间张量的模型可能只需要10张量的内存，因为大多数张量在其他张量创建之前就被消耗和释放了。这在内存有限的设备上至关重要。
|
||||
|
||||
## 硬件目标
|
||||
|
||||
### 移动GPU
|
||||
|
||||
- **Qualcomm Adreno**（Android）：支持OpenCL、Vulkan计算（第16章）和Qualcomm专有的SNPE（Snapdragon神经处理引擎）。Adreno GPU具有256-1024个ALU，支持FP16和INT8。
|
||||
|
||||
- **ARM Mali**（Android）：支持OpenCL和Vulkan。Mali GPU使用基于图块的架构（与桌面GPU不同），这影响最优内存访问模式。
|
||||
|
||||
- **Apple GPU**（iOS/macOS）：通过Metal（Apple的GPU API）访问。统一内存架构意味着没有CPU↔GPU复制开销。Metal Performance Shaders（MPS）提供优化的ML原语。
|
||||
|
||||
### 神经处理单元（NPU）
|
||||
|
||||
- NPU是专门为ML推理设计的固定功能加速器。它们在标准ML操作（矩阵乘、卷积、激活）上比GPU节能得多。
|
||||
|
||||
- **Apple神经引擎**：16核，约38 TOPS（INT8）。通过Core ML访问。非常适合视觉模型和设备端扩散。不能运行任意代码——只支持Core ML支持的操作。
|
||||
|
||||
- **Qualcomm Hexagon NPU**：集成到Snapdragon SoC中。支持INT8和INT4推理。通过SNPE或ONNX Runtime（带QNN后端）访问。为设备端功能如背景虚化、语音识别和实时翻译提供支持。
|
||||
|
||||
- **Google Edge TPU**：云端TPU的小型低功耗版本。4 TOPS，2W。用于Coral设备进行设备端推理。仅支持INT8量化的TFLite模型。
|
||||
|
||||
- **委派模式**：运行时在NPU（用于支持的操作）和CPU（用于不支持的操作）之间拆分模型图。最大化在NPU上运行的部分是性能和能效的关键。
|
||||
|
||||
## 设备端LLM
|
||||
|
||||
- 在手机和笔记本电脑上运行LLM已变得可行，得益于小模型和积极的量化：
|
||||
|
||||
| 模型 | 参数 | 量化后大小 | 目标设备 | 性能 |
|
||||
|-------|--------|---------------|---------------|-------------|
|
||||
| Phi-3 Mini | 3.8B | ~2 GB（Q4） | 手机/笔记本 | iPhone 15上~15 tokens/s |
|
||||
| Gemma 2B | 2B | ~1.5 GB（Q4） | 手机 | Pixel 8上~20 tokens/s |
|
||||
| Llama 3.2 1B | 1B | ~700 MB（Q4） | 手机 | ~30 tokens/s |
|
||||
| Llama 3.2 3B | 3B | ~2 GB（Q4） | 手机/笔记本 | ~15 tokens/s |
|
||||
| Llama 3.1 8B | 8B | ~4.5 GB（Q4） | 笔记本 | M2上~20 tokens/s |
|
||||
|
||||
- **挑战**：
|
||||
- **内存**：3B Q4模型占2 GB，但长对话的KV缓存增加了显著额外内存。手机上的上下文长度通常限制在2-4K token。
|
||||
- **热节流**：持续推理使手机发热。连续生成30秒后，SoC会降低时钟速度以防止过热，性能下降30-50%。
|
||||
- **电池**：以15 tokens/s运行3B模型消耗约3-5W。30分钟的对话消耗典型手机电池约5%。偶尔使用可以接受，但始终在线应用存在问题。
|
||||
|
||||
- **llama.cpp**是设备端LLM的标准。它在CPU（AVX2、NEON、I8MM）、Apple GPU（Metal）、NVIDIA GPU（CUDA）、AMD GPU（ROCm/Vulkan）甚至手机上（通过Android上的Termux）运行。
|
||||
|
||||
## 联邦学习
|
||||
|
||||
- **联邦学习**在许多设备上训练模型，无需集中数据。每个设备在其本地数据上训练，计算梯度更新，并将只有更新（而非数据）发送到聚合更新的中央服务器。
|
||||
|
||||
- **算法**（FedAvg）：
|
||||
1. 服务器将当前模型发送给$K$个选定设备。
|
||||
2. 每个设备在其本地数据上微调模型几步。
|
||||
3. 每个设备将其更新后的模型（或差异）发送回服务器。
|
||||
4. 服务器平均更新：$W_{\text{new}} = \frac{1}{K} \sum_{k=1}^{K} W_k$。
|
||||
5. 重复。
|
||||
|
||||
- **隐私**：原始数据从不离开设备。服务器只看到聚合的模型更新。**差分隐私**向更新添加噪声，使得无法从梯度中逆向推断单个数据点。
|
||||
|
||||
- **通信效率**：模型更新很大（与模型相同大小）。压缩技术减少了这一点：**梯度量化**（发送INT8梯度而不是FP32）、**稀疏化**（只发送最大的梯度）和**梯度累积**（做更多本地步骤，发送更少频率）。
|
||||
|
||||
- **应用**：Google的键盘预测（Gboard）、Apple的语音识别、健康监测（在敏感健康数据上训练而不集中数据）。
|
||||
|
||||
## 延迟优化
|
||||
|
||||
- 除了压缩，还有几种技术减少端到端推理延迟：
|
||||
|
||||
- **提前退出**：在中间层添加分类头。如果模型在第6层（共24层）已经自信，则返回预测而不运行第7-24层。简单输入提前退出，困难输入使用完整模型。对于混合简单和困难输入的任务，平均延迟显著下降。
|
||||
|
||||
- **模型分区**：在NPU（对矩阵乘高效）、GPU（对不规则操作高效）和CPU（处理其他一切）之间拆分模型。编译器根据性能分析决定哪些操作去哪里。
|
||||
|
||||
- **缓存**：对于具有重复查询的应用（自动补全、代码补全），缓存最近的计算。如果用户输入"How do I"且模型最近生成了"How do I"的补全，可以重用缓存的KV缓存，完全跳过预填充阶段。
|
||||
|
||||
- **推测性预取**：预测用户下一步将做什么，在用户询问之前开始推理。聊天应用可能在用户阅读当前答案时开始生成可能后续问题的响应。
|
||||
|
||||
## 编程任务（使用CoLab或notebook）
|
||||
|
||||
1. 模拟模型压缩流水线。从float32模型开始，依次应用蒸馏（模拟）、剪枝和量化，并跟踪每一步的大小。
|
||||
```python
|
||||
def compression_pipeline(original_params_M, original_bits=32):
|
||||
size_mb = original_params_M * 1e6 * original_bits / 8 / 1e6
|
||||
|
||||
print(f"原始: {original_params_M}M 参数, {original_bits}-位 → {size_mb:.0f} MB")
|
||||
|
||||
# 步骤1：知识蒸馏（减少参数）
|
||||
distilled_params = original_params_M * 0.15 # 70B → ~10B 等价
|
||||
size_mb = distilled_params * 1e6 * original_bits / 8 / 1e6
|
||||
print(f"蒸馏后 ({distilled_params:.0f}M 参数): {size_mb:.0f} MB")
|
||||
|
||||
# 步骤2：结构化剪枝（移除剩余30%）
|
||||
pruned_params = distilled_params * 0.7
|
||||
size_mb = pruned_params * 1e6 * original_bits / 8 / 1e6
|
||||
print(f"剪枝后 ({pruned_params:.0f}M 参数): {size_mb:.0f} MB")
|
||||
|
||||
# 步骤3：INT4量化
|
||||
size_mb = pruned_params * 1e6 * 4 / 8 / 1e6
|
||||
print(f"INT4量化后: {size_mb:.0f} MB")
|
||||
|
||||
print(f"总压缩比: {original_params_M * 1e6 * original_bits / 8 / 1e6 / size_mb:.0f}x")
|
||||
|
||||
print("=== 从70B模型开始 ===")
|
||||
compression_pipeline(70000)
|
||||
|
||||
print("\n=== 从7B模型开始 ===")
|
||||
compression_pipeline(7000)
|
||||
```
|
||||
|
||||
2. 估计设备端推理延迟。给定模型的操作计数和硬件规格，计算是否满足延迟目标。
|
||||
```python
|
||||
def estimate_latency(model_name, params_M, bits, compute_tops, mem_bw_gbs, seq_len=256):
|
||||
"""估计内存带宽受限模型的token生成延迟。"""
|
||||
# 模型大小（字节）
|
||||
model_bytes = params_M * 1e6 * bits / 8
|
||||
|
||||
# 解码是内存受限的：每token必须加载整个模型
|
||||
time_per_token_ms = model_bytes / (mem_bw_gbs * 1e9) * 1000
|
||||
|
||||
# 每秒token数
|
||||
tokens_per_sec = 1000 / time_per_token_ms
|
||||
|
||||
print(f"{model_name}: {params_M/1000:.1f}B 参数 @ {bits}-位 = {model_bytes/1e9:.1f} GB")
|
||||
print(f" 内存带宽: {mem_bw_gbs} GB/s")
|
||||
print(f" 每token时间: {time_per_token_ms:.1f} ms")
|
||||
print(f" Tokens/秒: {tokens_per_sec:.0f}")
|
||||
print()
|
||||
|
||||
# Apple M2 Pro：200 GB/s 统一内存带宽
|
||||
print("=== Apple M2 Pro (200 GB/s) ===")
|
||||
estimate_latency("Llama-7B Q4", 7000, 4, 15.8, 200)
|
||||
estimate_latency("Llama-7B Q8", 7000, 8, 15.8, 200)
|
||||
estimate_latency("Llama-70B Q4", 70000, 4, 15.8, 200)
|
||||
|
||||
# 手机（Snapdragon 8 Gen 3）：~50 GB/s LPDDR5
|
||||
print("=== Snapdragon 8 Gen 3 (50 GB/s) ===")
|
||||
estimate_latency("Phi-3 Mini Q4", 3800, 4, 45, 50)
|
||||
estimate_latency("Llama-3B Q4", 3000, 4, 45, 50)
|
||||
```
|
||||
@@ -0,0 +1,251 @@
|
||||
# 扩缩与部署
|
||||
|
||||
*向数百万用户提供大模型服务需要跨多个GPU分布推理、在需要之前预测token、缓存共享上下文以及选择合适的框架。本文涵盖推理时的并行性、推测性解码、前缀缓存、推理框架、成本优化和监控*
|
||||
|
||||
- 单个H100 GPU服务一个70B模型可以处理约100个并发用户，交互延迟可接受。服务1000万用户需要100,000个GPU——云计算每年花费约30亿美元。每一个百分点的效率提升就能节省数千万美元。这就是推理优化不是学术问题的原因：它直接决定AI产品的经济性。
|
||||
|
||||
## 推理时的模型并行
|
||||
|
||||
- 当模型太大无法装入单张GPU时，必须跨多个GPU拆分。训练时的并行策略（第6章）在推理时适用，但权衡不同。
|
||||
|
||||
### 张量并行
|
||||
|
||||
- **张量并行**（Megatron风格，第6章）跨GPU拆分单个权重矩阵。对于线性层$Y = XW$，权重矩阵$W$跨$N$个GPU按列拆分。每个GPU计算部分结果，然后all-reduce聚合：
|
||||
|
||||
$$W = [W_1 | W_2 | \cdots | W_N], \quad Y_i = X W_i, \quad Y = \text{concat}(Y_1, \ldots, Y_N)$$
|
||||
|
||||
- 在推理时，张量并行是模型无法装入单张GPU时的默认选择。FP16的70B模型需要140 GB——跨2张80 GB GPU使用张量并行拆分。
|
||||
|
||||
- **延迟影响**：张量并行每层增加一个all-reduce通信步骤。在NVLink（900 GB/s）上，每层增加约0.1 ms。在PCIe（32 GB/s）上，每层增加约3 ms。对于80层的70B模型在2张GPU上：NVLink总增加约8 ms，PCIe总增加约240 ms。这就是NVLink对多GPU推理至关重要的原因。
|
||||
|
||||
### 流水线并行
|
||||
|
||||
- **流水线并行**将不同的层分配给不同的GPU。GPU 1处理第0-39层，GPU 2处理第40-79层。token顺序流过流水线。
|
||||
|
||||
- 在推理时，流水线并行的延迟高于张量并行（每个token必须遍历整个流水线），但通信开销更低（只有激活值在GPU之间传递，无需all-reduce）。当GPU通过慢速互连（不同节点，无NVLink）连接时，更倾向于使用流水线并行。
|
||||
|
||||
### 序列并行
|
||||
|
||||
- 对于非常长的序列，即使模型本身适合，KV缓存本身可能无法装入单张GPU。**序列并行**将KV缓存分片到多个GPU上：每个GPU存储序列缓存键和值的一部分。
|
||||
|
||||
- 在注意力期间，每个GPU在其缓存的段上计算部分注意力分数，然后通过规约合并结果。这用于长上下文推理（128K+ token），其中KV缓存超过单GPU内存。
|
||||
|
||||
## 推测性解码
|
||||
|
||||
- **推测性解码**是影响最大的LLM推理优化之一。其洞察：解码速度慢是因为一次只生成一个token，每个token需要大模型的完整前向传播。但小模型可以更快地生成候选token，而大模型可以**验证**多个候选token。
|
||||
|
||||
|
||||
|
||||
- **算法**：
|
||||
1. **草稿模型**（小型、快速——例如1B参数）自回归生成$k$个候选token。
|
||||
2. **目标模型**（大型、准确——例如70B）对整个草稿序列运行一次前向传播，计算每个候选token的概率。
|
||||
3. 如果目标模型同意（该token的概率足够高），每个候选被**接受**。被拒绝的候选从目标模型的分布中重新采样。
|
||||
4. 平均每个验证步骤接受多个token，加速比与接受率成正比。
|
||||
|
||||
$$\text{加速比} \approx \frac{k \times \text{acceptance\_rate}}{\text{cost\_ratio}} \approx 2\text{-}3\times$$
|
||||
|
||||
- **为什么无质量损失**：拒绝采样方案保证输出分布与目标模型完全匹配。推测性解码是无损的——输出在统计上与单独运行目标模型相同，只是更快。
|
||||
|
||||
- **变体**：
|
||||
- **Medusa**（Cai等人，2024）：不是独立的草稿模型，而是向目标模型添加多个轻量级"头"，同时预测多个未来token。无需独立模型。
|
||||
- **EAGLE**（Li等人，2024）：训练一个使用目标模型隐藏状态预测未来token的轻量级草稿头。接受率高于独立的草稿模型。
|
||||
- **自推测性解码**：目标模型本身使用提前退出生成草稿（仅运行前几层作为草稿，然后用完整模型验证）。
|
||||
- **并行解码**：并行生成多个延续（候选树）并一次性验证整棵树。吞吐量更高，但分支KV缓存使用更多内存。
|
||||
|
||||
## 前缀缓存
|
||||
|
||||
- 许多请求共享共同的前缀：系统提示、few-shot示例或常见查询模式。**前缀缓存**存储这些前缀的KV缓存并在请求之间重用。
|
||||
|
||||
- **系统提示缓存**：如果每个请求都以相同的2000-token系统提示开始，这2000个token的KV缓存被计算一次，并在所有请求之间共享。对于80层的70B模型，每次请求节省约200 MB。
|
||||
|
||||
- **基数树缓存**（SGLang）：将缓存的前缀组织在基数树（trie）中。当新请求到达时，找到最长的缓存前缀匹配，并从那里开始生成，跳过匹配前缀的计算。
|
||||
|
||||
- **影响**：对于具有长共享前缀的应用（带系统提示的聊天机器人、具有常见检索段落的RAG），前缀缓存将TTFT降低50-90%，并节省相应的GPU计算。
|
||||
|
||||
## KV缓存驱逐
|
||||
|
||||
- 除了量化KV缓存（文件01）和使用GQA/MLA减小其大小（文件02）之外，**KV缓存驱逐**策略选择性地移除不太可能在未来被关注的缓存token。
|
||||
|
||||
- **H2O**（重要token识别器，Zhang等人，2023）观察到注意力分数遵循幂律：一小部分token（"重要token"）获得大部分注意力，而大多数token获得的注意力可以忽略不计。H2O保留：
|
||||
|
||||
1. **最近token**（最后$w$个token的滑动窗口，类似StreamingLLM）。
|
||||
2. **重要token**（在所有过去解码步骤中累积注意力分数最高的前$k$个token）。
|
||||
|
||||
- 既不是最近也不是重要token的token被驱逐。这保持固定大小的KV缓存，同时保留实际影响生成的token。H2O仅使用20%的内存就实现了接近完整KV缓存的质量。
|
||||
|
||||
- **Scissorhands**（Liu等人，2023）采用类似方法，但使用更复杂的度量：在当前**步骤**中获得高注意力的token被保留，而已经$T$步没有被关注的token被驱逐。这适应了生成过程中注意力模式的变化。
|
||||
|
||||
- **动态驱逐+StreamingLLM**：结合注意力汇聚点（永久保留前几个token）和动态驱逐（保留最近+重要token）。这是最内存高效的方法，适用于非常长的生成，实现了无限长度生成，质量下降有限。
|
||||
|
||||
- 所有驱逐方法的核心洞察：LLM注意力在实践中是**稀疏的**——尽管架构会对所有缓存的token计算注意力，但实际注意力权重集中在小子集上。驱逐其余部分对输出质量影响极小。
|
||||
|
||||
## 推理框架
|
||||
|
||||
- LLM服务生态已收敛到几个主要框架：
|
||||
|
||||
| 框架 | 优势 | 最适合 |
|
||||
|-----------|-----------|----------|
|
||||
| **vLLM** | PagedAttention、连续批处理、高吞吐量 | 通用LLM服务，最高吞吐量 |
|
||||
| **TensorRT-LLM** | NVIDIA优化内核、FP8、飞行中批处理 | NVIDIA GPU上的最大性能 |
|
||||
| **SGLang** | 前缀缓存（RadixAttention）、快速结构化生成 | 具有共享前缀的应用，受限输出 |
|
||||
| **llama.cpp** | CPU/Metal/CUDA/Vulkan、GGUF量化、可移植 | 消费级硬件，设备端推理 |
|
||||
| **TGI**（HuggingFace） | 简单API，易于部署，模型中心集成 | 快速部署，HuggingFace生态 |
|
||||
| **Ollama** | 一键下载和提供服务 | 个人使用，本地开发 |
|
||||
| **ExLlamaV2** | 极致量化优化（EXL2格式） | 内存受限的GPU推理 |
|
||||
|
||||
- **vLLM**是生产级LLM服务的默认选择。它支持连续批处理、PagedAttention、张量并行、推测性解码、LoRA服务和大多数开源模型。
|
||||
|
||||
- **TensorRT-LLM**在NVIDIA硬件上实现最高的原始性能（在相同GPU上比vLLM快10-30%），但灵活性较低且更难以定制。
|
||||
|
||||
- **SGLang**在应用具有结构化输出（JSON、特定格式的代码）或共享前缀时表现出色，这得益于其基数注意力缓存和受限解码引擎。
|
||||
|
||||
## 成本优化
|
||||
|
||||
- 在规模上，推理成本主导ML预算。降低成本的策略：
|
||||
|
||||
- **合理选择GPU**：并非每个模型都需要H100。量化的7B模型在A10G（约$1/小时）上运行良好，而不是H100（约$8/小时）。匹配GPU到工作负载。
|
||||
|
||||
- **竞价实例**：云提供商提供未使用的GPU容量，折扣60-90%（AWS Spot、GCP Preemptible）。竞价实例可能被中断，因此适用于批处理推理而不是延迟关键型服务。结合抢占处理（保存状态，在新实例上恢复），竞价实例也可以服务交互式流量。
|
||||
|
||||
- **自动扩缩**：根据流量扩展GPU数量。高峰期扩展，夜间缩减。Kubernetes HPA（水平Pod自动扩缩器）或云原生自动扩缩（AWS SageMaker、GCP Vertex AI）处理此功能。
|
||||
|
||||
- **批处理+利用率**：30%和90% GPU利用率之间的差异是每token成本3倍。连续批处理、智能调度和PagedAttention都提高了利用率。
|
||||
|
||||
- **量化**：INT4 vs FP16是4倍更少内存 → 适合更小的GPU → 成本降低2-4倍。此外，更多请求适合同一批次 → 更高吞吐量 → 更低每token成本。
|
||||
|
||||
- **每token成本基准**（近似值，2026年）：
|
||||
|
||||
| 配置 | 每100万token成本 |
|
||||
|-------|-------------------|
|
||||
| GPT-4o API | $2.50 |
|
||||
| Claude 3.5 Sonnet API | $3.00 |
|
||||
| Llama-70B on H100（vLLM，FP16） | $0.50 |
|
||||
| Llama-70B on H100（TRT-LLM，INT8） | $0.25 |
|
||||
| Llama-8B on A10G（vLLM，INT4） | $0.05 |
|
||||
| Llama-3B 设备端（llama.cpp） | $0（硬件成本摊销） |
|
||||
|
||||
## 监控
|
||||
|
||||
- 生产推理需要持续监控，以便在用户受到影响之前发现降级：
|
||||
|
||||
- **延迟监控**：跟踪TTFT和TPOT的p50、p95和p99。设置告警，当p99超过SLO时触发。p99的尖峰通常指示：KV缓存内存压力（抖动）、长时间运行的请求垄断批次、或GPU热节流。
|
||||
|
||||
- **吞吐量监控**：跟踪每GPU每秒token数。下降指示：批次效率降低（许多短请求→批次利用率低）、序列长度增加（每个请求更多KV缓存内存）、或硬件问题（GPU处于ECC纠错模式，运行更慢）。
|
||||
|
||||
- **GPU利用率**：跟踪SM占用率、内存利用率和内存带宽。低SM占用率+高内存利用率=内存受限（需要更多带宽或量化）。高SM占用率+低内存利用率=计算受限（需要更多FLOPS或更小模型）。
|
||||
|
||||
- **模型质量监控**：跟踪每请求指标（响应长度、保留集上的困惑度、用户反馈信号）。模型质量可能因以下原因降级：数据漂移（传入请求的分布变化）、KV缓存量化误差在长对话中累积、或服务流水线中的错误。
|
||||
|
||||
- **成本监控**：跟踪每模型每GPU类型每token成本。如果成本增加而吞吐量没有增加，调查效率回归（新模型版本内存使用更高、批次配置次优、或GPU利用不足）。
|
||||
|
||||
- **工具**：Prometheus + Grafana（第15章）用于基础设施指标，vLLM/TRT-LLM的内置指标端点，以及用于模型级指标的自定义日志记录。
|
||||
|
||||
## 编程任务（使用CoLab或notebook）
|
||||
|
||||
1. 模拟推测性解码。使用快速的"草稿"函数和慢速的"目标"函数，测量一次生成和验证多个token的加速比。
|
||||
```python
|
||||
import random
|
||||
import time
|
||||
|
||||
def target_model(tokens):
|
||||
"""慢但准确的模型。返回每个候选token的概率。"""
|
||||
time.sleep(0.01) # 模拟每次前向传播10ms
|
||||
# 用于模拟：接受偶数token
|
||||
return [0.9 if t % 2 == 0 else 0.1 for t in tokens]
|
||||
|
||||
def draft_model():
|
||||
"""快但近似的模型。生成一个候选token。"""
|
||||
time.sleep(0.001) # 模拟每token 1ms
|
||||
return random.randint(0, 9)
|
||||
|
||||
def standard_decoding(n_tokens):
|
||||
"""一次生成一个token，使用目标模型。"""
|
||||
tokens = []
|
||||
for _ in range(n_tokens):
|
||||
time.sleep(0.01) # 目标模型生成1个token
|
||||
tokens.append(random.randint(0, 9))
|
||||
return tokens
|
||||
|
||||
def speculative_decoding(n_tokens, k=4):
|
||||
"""生成k个草稿token，用目标模型验证，接受/拒绝。"""
|
||||
tokens = []
|
||||
total_target_calls = 0
|
||||
|
||||
while len(tokens) < n_tokens:
|
||||
# 草稿：快速生成k个候选
|
||||
candidates = [draft_model() for _ in range(k)]
|
||||
|
||||
# 验证：一次目标模型调用验证所有k个候选
|
||||
probs = target_model(candidates)
|
||||
total_target_calls += 1
|
||||
|
||||
# 接受token，直到一个被拒绝
|
||||
for i, (tok, prob) in enumerate(zip(candidates, probs)):
|
||||
if random.random() < prob:
|
||||
tokens.append(tok)
|
||||
if len(tokens) >= n_tokens:
|
||||
break
|
||||
else:
|
||||
# 从目标分布重新采样
|
||||
tokens.append(tok + 1) # 简化重新采样
|
||||
break
|
||||
|
||||
return tokens, total_target_calls
|
||||
|
||||
n = 50
|
||||
|
||||
start = time.time()
|
||||
_ = standard_decoding(n)
|
||||
standard_time = time.time() - start
|
||||
|
||||
start = time.time()
|
||||
_, target_calls = speculative_decoding(n, k=5)
|
||||
spec_time = time.time() - start
|
||||
|
||||
print(f"标准: {standard_time:.2f}s ({n} 次目标调用)")
|
||||
print(f"推测性: {spec_time:.2f}s ({target_calls} 次目标调用)")
|
||||
print(f"加速比: {standard_time / spec_time:.1f}x")
|
||||
```
|
||||
|
||||
2. 估计应用于LLM服务部署的不同优化策略的成本节省。
|
||||
```python
|
||||
def serving_cost_analysis(
|
||||
model_name, params_B, precision_bits,
|
||||
gpu_name, gpu_mem_gb, gpu_cost_per_hr,
|
||||
target_throughput_tps,
|
||||
):
|
||||
"""估计LLM部署的服务成本。"""
|
||||
model_size_gb = params_B * 1e9 * precision_bits / 8 / 1e9
|
||||
gpus_for_model = max(1, int((model_size_gb * 1.2) / gpu_mem_gb + 0.99)) # 1.2x用于KV缓存
|
||||
|
||||
# 粗略吞吐量估计（内存带宽受限）
|
||||
tokens_per_gpu = 500 / (params_B * precision_bits / 16) # 归一化到7B FP16的500 tok/s
|
||||
total_throughput = tokens_per_gpu * gpus_for_model
|
||||
|
||||
replicas = max(1, int(target_throughput_tps / total_throughput + 0.99))
|
||||
total_gpus = gpus_for_model * replicas
|
||||
cost_per_hr = total_gpus * gpu_cost_per_hr
|
||||
cost_per_1M_tokens = cost_per_hr / (total_throughput * replicas * 3600 / 1e6)
|
||||
|
||||
print(f"{model_name} @ {precision_bits}-位 在 {gpu_name} 上:")
|
||||
print(f" 模型大小: {model_size_gb:.0f} GB → {gpus_for_model} GPU(s)/副本")
|
||||
print(f" 吞吐量: {total_throughput:.0f} tok/s/副本")
|
||||
print(f" 需达到{target_throughput_tps} tok/s的副本数: {replicas}")
|
||||
print(f" 总GPU数: {total_gpus}")
|
||||
print(f" 成本: ${cost_per_hr:.0f}/小时, ${cost_per_1M_tokens:.2f}/100万token")
|
||||
print()
|
||||
|
||||
print("=== 成本比较 ===\n")
|
||||
|
||||
# 基线：H100上的FP16
|
||||
serving_cost_analysis("Llama-70B", 70, 16, "H100", 80, 8.0, 1000)
|
||||
|
||||
# 量化后：H100上的INT8
|
||||
serving_cost_analysis("Llama-70B", 70, 8, "H100", 80, 8.0, 1000)
|
||||
|
||||
# 量化后：A100上的INT4
|
||||
serving_cost_analysis("Llama-70B", 70, 4, "A100", 80, 4.0, 1000)
|
||||
|
||||
# 小模型：A10G上的8B
|
||||
serving_cost_analysis("Llama-8B", 8, 4, "A10G", 24, 1.0, 1000)
|
||||
```
|
||||
Reference in New Issue
Block a user