What makes transformers heavy on computation and memory, and how can we address this
1️⃣ 考察意图
面试官想看你是否真正理解Transformer的“重”在哪里，而不是背复杂度公式。核心考察三点：一是对自注意力O(n²d)计算和O(n²)显存瓶颈的直觉（为什么n是杀手）；二是对优化技术的工程取舍（比如FlashAttention用IO感知换显存，但牺牲灵活性）；三是能否区分“训练重”和“推理重”的不同解法。刁钻点在于：候选人常只提稀疏注意力，却忽略FFN层的参数占比和KV Cache的推理瓶颈。答好了能展示系统级优化思维和落地经验。
2️⃣ 标准答
Transformer的计算和内存瓶颈分两大块：自注意力机制和前馈网络（FFN），且训练与推理的问题不同。
1. 计算瓶颈
- 自注意力：标准实现复杂度O(n²d)，n是序列长度，d是头维度。每个token需与所有token计算点积，n=4096时约1600万次操作，n=8192时翻4倍。这是“二次方墙”。
- FFN层：参数占比最大（约2/3）。例如GPT-3 175B中，FFN的权重矩阵尺寸为12288x49152，单次前向传播需约6亿次乘加。但FFN计算是O(nd²)，对短序列更重，长序列时注意力成瓶颈。
- 工程取舍：注意力是“内存带宽瓶颈”（需频繁读写K,V矩阵），FFN是“计算瓶颈”（矩阵乘法密集）。优化需对症下药。
2. 内存瓶颈
- 训练：需存储所有层的激活值（attention scores、K,V缓存、FFN中间结果）。n=4096、d=4096时，单层注意力得分矩阵占约64MB（FP16），32层模型仅此一项超2GB。加上梯度、优化器状态，显存易爆。
- 推理：核心是KV Cache。自回归生成时，需缓存所有历史token的K,V矩阵。n=2048、batch=1时约16MB，但batch=64时超1GB，且随序列长度线性增长。长上下文（128K tokens）下KV Cache可达数十GB。
- 实际坑：很多人忽略“中间激活”比模型权重更吃显存。例如训练LLaMA-65B时，激活显存是权重的3-5倍。
3. 优化方法（按场景选）
- 注意力优化：
- FlashAttention：IO感知算法，将Q,K,V分块加载到SRAM，避免全局显存读写。实测训练速度提升2-4倍，显存降低5-10倍（n=4096时）。但牺牲了掩码灵活性（如因果掩码需特殊处理）。
- 稀疏注意力：如Longformer的滑动窗口+全局token，复杂度降为O(nw)，w为窗口大小。适合文档级任务，但长距离依赖可能丢失。
- 线性注意力：如Performer用核方法近似softmax，复杂度O(nd²)。但精度损失明显（约0.5-1%），且对长序列收益递减。
- FFN优化：
- MoE（混合专家）：如Mixtral 8x7B，每个token只激活2个专家，计算量降为1/4。但需额外通信开销，且负载均衡难调。
- 量化：INT8/FP8量化FFN权重，显存减半，计算加速1.5-2倍。但需校准集避免精度崩坏。
- 训练技巧：
- 梯度检查点：不存中间激活，反向传播时重算。显存降30-50%，但训练时间增20-30%。适合显存受限场景。
- 序列并行：将长序列切分到多GPU，每个GPU只算子序列注意力。需跨GPU通信，但n=128K时显存可降8倍。
- 推理优化：
- KV Cache量化：INT8量化KV Cache，显存减半，精度损失<0.1%。需注意量化误差随序列长度累积。
- PagedAttention：如vLLM，将KV Cache分页管理，避免碎片化。吞吐量提升2-4倍。
4. 架构级改进
- 线性注意力变体：如Mamba（状态空间模型），复杂度O(n)，但牺牲了Transformer的并行训练优势。
- 混合架构：如RetNet，结合注意力与循环，训练用注意力（并行），推理用循环（O(1)显存）。但实现复杂，生态不成熟。
总结：没有银弹。长序列场景优先FlashAttention+梯度检查点；推理优先KV Cache量化+PagedAttention；超大规模模型考虑MoE。关键是根据n和d的比值选方案。
3️⃣ 答题模板（30 秒电梯版）
“这个问题我从计算、内存、优化三个层面回答。计算上，自注意力O(n²d)和FFN的O(nd²)是两大瓶颈，n是杀手。内存上，训练时中间激活比权重更吃显存，推理时KV Cache随序列线性增长。优化需分场景：长序列用FlashAttention降显存，推理用KV Cache量化+PagedAttention，超大规模用MoE。总结一句：没有通用解法，需根据序列长度和部署场景做工程取舍。”
4️⃣ 高频追问 & 应对
追问 1：FlashAttention为什么能省显存？它和稀疏注意力比有什么trade-off？
FlashAttention通过分块计算+重算避免存储完整注意力矩阵，核心是IO感知：将Q,K,V分块加载到SRAM，计算局部注意力后写回HBM。显存从O(n²)降到O(n)，但计算量略增（重算开销）。与稀疏注意力比：FlashAttention保持全注意力精度，适合需要全局依赖的任务（如代码生成）；稀疏注意力显存更低（O(nw)），但长距离依赖可能丢失，适合文档分类等局部任务。工程上，FlashAttention对GPU架构敏感（需A100+），稀疏注意力更通用。
追问 2：训练时梯度检查点怎么用？显存省多少，代价是什么？
梯度检查点只存储部分层激活值（如每4层存一个），反向传播时从最近检查点重算。显存省30-50%（取决于检查点间隔），但训练时间增20-30%（重算开销）。实际坑：检查点间隔太密（如每层存）显存省得少；太疏（如每8层）重算开销大。推荐每4-6层存一个，且对FFN层（计算密集）重算更划算。注意：梯度检查点对batch size敏感，batch=1时收益最大，batch=64时收益递减。
追问 3：推理时KV Cache为什么是瓶颈？怎么优化？
KV Cache存储所有历史token的K,V矩阵，随序列长度线性增长。n=128K、d=4096、batch=1时，单层KV Cache约1GB（FP16），32层模型超32GB。优化方法：1）KV Cache量化（INT8），显存减半，精度损失<0.1%，但需注意长序列下误差累积；2）PagedAttention，分页管理避免碎片化，吞吐量提升2-4倍；3）窗口KV Cache，只存最近N个token（如4K），适合流式应用，但长距离依赖丢失。工程取舍：量化+分页组合效果最好，但需额外算子支持。
5️⃣ 避坑 · 常见错误答法
- ❌ 只提“自注意力复杂度O(n²)”，不提FFN和KV Cache → ✅ 需区分训练和推理瓶颈：训练时FFN参数占比大，推理时KV Cache是显存杀手。
- ❌ 说“FlashAttention是银弹，所有场景都用它” → ✅ 需说明trade-off：FlashAttention牺牲灵活性（如不支持自定义掩码），且对GPU架构有要求（A100+），稀疏注意力在长文档分类中更高效。
- ❌ 把“梯度检查点”和“混合精度训练”混为一谈 → ✅ 梯度检查点省显存但增时间，混合精度训练省显存且加速，两者可组合但需调参。
6️⃣ 简历呼应
- 如果你有LLM训练项目：从“训练LLaMA-13B时显存爆了”切入，对比FlashAttention和梯度检查点的实际收益（如显存从80GB降到40GB，训练时间增15%），展示工程调优能力。
- 如果你只做过传统NLP：用“BERT推理时序列长度从512扩到2048，显存从4GB涨到16GB”类比，引出KV Cache优化，强调从短序列到长序列的迁移思路。
- 如果你是校招无项目：聚焦FlashAttention论文复现，用PyTorch实现分块注意力，对比标准实现的速度和显存，展示对IO感知的理解。
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022)
- Longformer: The Long-Document Transformer (Beltagy et al., 2020)
- PagedAttention: Efficient Memory Management for Large Language Model Serving (Kwon et al., 2023)
- Mixture of Experts: Outrageously Large Neural Networks (Shazeer et al., 2017)
- Gradient Checkpointing: Training Deep Nets with Sublinear Memory Cost (Chen et al., 2016)