2. 长序列建模的显存瓶颈：注意力机制的O(n²)复杂度，KV Cache的显存占用，激活值显存分析
好，咱们直接切入正题。做长序列建模，说白了就是在跟显存打仗。我这些年调过不少模型，从BERT到GPT，再到现在的LLaMA系列，每次序列长度一上去，显存就像开了闸的水一样哗哗往外流。今天咱们就把这三个“吃显存”的元凶揪出来，一个一个分析清楚。
2.1 注意力机制的O(n²)复杂度——显存爆炸的根源
先问个问题：为什么序列一长，显存就扛不住？
答案就在注意力机制的计算公式里。标准的Scaled Dot-Product Attention，它的计算复杂度是O(n²)，这里的n是序列长度。什么意思呢？就是序列长度翻倍，计算量和显存占用要翻四倍。
我举个例子你就明白了。假设序列长度是1024，注意力分数矩阵的大小就是1024×1024。如果序列长度变成4096，这个矩阵就变成了4096×4096——大小直接翻了16倍。你想想看，这谁顶得住？
核心问题：注意力分数矩阵需要显式存储，其大小与序列长度的平方成正比。对于长度为L的序列，单层单头的注意力分数矩阵占用内存为O(L² × 精度)。
我在项目中遇到过最夸张的一次，是把序列长度从2048提到8192，结果显存直接爆了。当时我还纳闷，明明模型参数量没变啊？后来一查，好家伙，注意力分数矩阵占了将近70%的显存。
2.2 KV Cache的显存占用——推理时的隐形杀手
推理阶段，还有个更隐蔽的显存杀手——KV Cache。
自回归生成的时候，每生成一个token，都需要重新计算注意力。但你不觉得每次都从头算一遍很浪费吗？其实前面已经算过的Key和Value是可以复用的。这就是KV Cache的由来——把之前所有token的Key和Value缓存起来，避免重复计算。
但问题来了：这个缓存有多大？
| 参数 | 计算公式 | 示例（L=4096, d=128, h=32, b=1） |
|---|---|---|
| 单层KV Cache | 2 × L × d × h × 精度 | 2 × 4096 × 128 × 32 × 2B = 64MB |
| 全部层KV Cache | 层数 × 单层大小 | 32层 × 64MB = 2GB |
| 批量推理 | batch_size × 全部层大小 | 8 × 2GB = 16GB |
看到没？一个32层的模型，序列长度4096，单batch推理就要吃掉2GB显存。如果batch_size是8，那就是16GB。我刚开始做推理优化的时候，就吃过这个亏——模型参数量才7B，结果KV Cache占的显存比模型本身还大。
我的经验：KV Cache的显存占用跟序列长度成正比，跟层数成正比，跟batch_size成正比。这三个维度任何一个放大，显存都会线性增长。所以做长序列推理时，KV Cache往往是第一个需要优化的目标。
2.3 激活值显存分析——训练时的“隐形账单”
训练阶段，还有个更让人头疼的东西——激活值（Activations）。
你可能觉得，模型参数才是显存大头。其实不是。训练的时候，反向传播需要用到前向传播的中间结果，这些中间结果就是激活值。它们占的显存，往往比模型参数还多。
咱们来算一笔账。一个标准的Transformer层，需要存储的激活值包括：
- 注意力分数矩阵：O(L² × h)
- 注意力输出：O(L × d × h)
- FFN中间激活：O(L × 4d × h)
- LayerNorm的均值和方差：O(L × d × h)
- Dropout mask：O(L × d × h)
把这些加起来，你会发现一个惊人的事实：激活值显存跟序列长度L的关系，主要是O(L²)主导的。也就是说，序列一长，激活值显存比模型参数显存大得多。
注意：很多人只关注模型参数量，忽略了激活值。我曾经在训练一个长序列模型时，把batch_size从1调到2，结果显存直接翻倍。后来才发现，激活值占了总显存的80%以上。所以，做长序列训练，激活值管理是第一优先级。
2.4 知识体系总览
为了让你更直观地理解这三个瓶颈之间的关系，我画了一张图：
这张图把三个瓶颈的关系理清楚了。你看，注意力O(n²)是根源，KV Cache是推理的痛点，激活值是训练的痛点。三者相互叠加，共同限制了你能处理的最大序列长度。
2.5 一个具体的显存估算例子
光说理论不够，咱们来算个实际的。假设你有一个7B参数的LLaMA模型，32层，隐藏维度4096，注意力头数32，每个头的维度128。
训练时，序列长度L=4096，batch_size=1：
- 模型参数：7B × 2B = 14GB（FP16）
- 优化器状态：14GB × 2 = 28GB（Adam，需要存储动量和方差）
- 激活值：约40-60GB（取决于具体实现）
- 总计：约82-102GB
你看，激活值占了将近一半。如果序列长度翻倍到8192，激活值会飙升到160-240GB，直接爆掉A100的80GB显存。
我的建议：在做长序列建模之前，先做个显存预算。把模型参数、优化器状态、激活值、KV Cache这几项都算清楚。我习惯用这个公式快速估算：总显存 ≈ 模型参数 × (2 + 优化器系数) + 激活值系数 × L²。这样心里就有底了。
嗯，这一章咱们把三个显存瓶颈都讲清楚了。下一章，我会介绍SwiGLU如何通过门控机制来缓解这些问题。不过在那之前，建议你先把自己模型的显存占用算一算，看看哪个瓶颈最严重。
公众号:蓝海资料掘金营，微信deep3321