显存优化

训练大模型的显存困境 训练一个70B参数的模型，仅模型参数（FP16）就需要140GB显存。加上梯度（140GB）、优化器状态（Adam需要280GB）和激活值，总显存需求轻松超过600GB——即使8卡A100 80GB也无法满足。 在这些显存消耗中，激活值是最容易被忽视的部分。在标准反向传播中，前向传播的所有中间激活值都需要保存下来供反向传播计算梯度使用。对于70B模型、batch_size=4、seq_len=4096，激活值占用可达100GB以上。 梯度检查点（Gradient Checkpointing）通过"以时间换空间"的策略，大幅减少激活值的显存占用。 标准反向传播的显存分析 在标准反向传播中，每个操作的前向输出都需要被保存： 前向: x → Layer1 → a1 → Layer2 → a2 → Layer3 → a3 → Loss 保存: a1, a2, a3 (以及各层的中间状态) 反向: Loss ← ∂L/∂a3 ← ∂L/∂a2 ← ∂L/∂a1 ← ∂L/∂x 使用: 重新使用a3计算a2的梯度，重新使用a2计算a1的梯度，... 对于N层Transformer，每层的激活值大小为 O(batch × seq_len × hidden_dim)。所有层的激活值总和为 O(N × batch × seq_len × hidden_dim)，与层数线性增长。 梯度检查点的核心思想 选择性重计算 梯度检查点的核心洞察：不需要保存所有层的激活值，只需保存部分"检查点"层的输入，其余层的激活值在反向传播时重新计算。 检查点设置（每2层一个检查点）: 保存: x, a2, a4, a6, ... 前向: x → L1 → a1 → L2 → a2* → L3 → a3 → L4 → a4* → ... (* = 保存的检查点) 反向: 从a6*出发，重新前向计算a5, a4→计算梯度 从a4*出发，重新前向计算a3, a2→计算梯度 ... 显存-计算权衡 假设N层Transformer，每隔k层设置一个检查点： ...

引言 KV Cache（键值缓存）是大模型推理中显存占用的最大头。一个512K上下文的推理，KV Cache可能占用数十GB显存。2026年，KV Cache优化技术已经非常成熟，从vLLM的PagedAttention到DeepSeek的MLA，各种创新层出不穷。本文将全面解析KV Cache的优化技术，帮助开发者深入理解并合理选择。 KV Cache基础 为什么需要KV Cache？ Transformer的自注意力机制： Attention(Q, K, V) = softmax(QK^T / √d_k) × V 问题：生成第t个token时，需要K_1, K_2, …, K_{t-1}和V_1, …, V_{t-1}。如果每次都重新计算，复杂度O(n²)。 解决：缓存历史K和V → KV Cache KV Cache的显存占用 对于Llama 4 70B（32层，hidden_size=8192）： 上下文长度 FP16 KV Cache INT8 KV Cache INT4 KV Cache 2K 512MB 256MB 128MB 8K 2GB 1GB 512MB 32K 8GB 4GB 2GB 128K 32GB 16GB 8GB 512K 128GB 64GB 32GB 关键观察：上下文长度每增加4倍，KV Cache增大4倍。 优化技术全景 技术分类 技术路线 核心思路 压缩比 质量损失 注意力优化 减少KV Cache占用 - - ├─ Multi-Query Attention (MQA) 多查询共享KV 1/n_heads 0% ├─ Grouped-Query Attention (GQA) 分组共享KV 1/group_size <0.5% └─ Multi-Head Latent Attention (MLA) 低维潜在表示 8-10× <1% 分页管理 动态内存管理 - - ├─ PagedAttention 分页管理KV Cache 减少碎片 0% └─ RadixAttention 前缀树缓存 复用公共前缀 0% 量化压缩 降低精度 - - ├─ INT8量化 8-bit量化 2× <1% ├─ INT4量化 4-bit量化 4× 2-3% └─ FP8量化 8-bit浮点 2× <0.5% 稀疏化 只保留重要token - - ├─ StreamingLLM 保留初始+近期token 可变 1-3% ├─ H2O 动态淘汰低注意力token 2-4× 2-5% └─ ScissorHands 基于重要性的稀疏化 2-3× 2-4% 注意力优化 1. Multi-Query Attention (MQA) 原理：所有查询头共享同一组KV。 ...

QLoRA解决了什么问题？ LoRA已经大幅降低了微调成本，但对于70B模型，仍然需要2×A100(80GB)。QLoRA（Quantized LoRA）通过4-bit量化基座模型，将显存需求再降一个数量级： 模型 全量微调 LoRA QLoRA 7B 80GB 16GB 6GB 14B 160GB 32GB 12GB 70B 800GB 160GB 24GB 这意味着：用一张RTX 4090就能微调70B模型。 QLoRA核心原理 三大创新 1. NF4（NormalFloat 4-bit）量化 传统4-bit量化（INT4）假设数据均匀分布，但神经网络权重通常服从正态分布。NF4专门为正态分布设计量化区间： # NF4量化原理 import numpy as np # INT4: 均匀量化 int4_levels = np.linspace(-1, 1, 16) # 16个均匀区间 # NF4: 正态分布感知量化 from scipy.stats import norm # 16个分位点对应正态分布的分位 nf4_levels = norm.ppf(np.linspace(0, 1, 17)[1:-1]) nf4_levels = nf4_levels / np.max(np.abs(nf4_levels)) # 归一化到[-1, 1] print(f"INT4 levels: {int4_levels}") print(f"NF4 levels: {nf4_levels}") # NF4在0附近更密集，在两端更稀疏，匹配权重分布 2. Double Quantization 对量化后的缩放因子本身再量化，进一步节省显存： 原始权重 W → 4-bit量化 → 量化值 + 缩放因子(32bit) ↓ 缩放因子再量化为8-bit ↓ 额外节省约0.4bit/参数 3. Paged Optimizer 使用NVIDIA的统一内存，当显存不够时自动将优化器状态转移到CPU内存： ...