sys(1): 内存墙与LLM 计算

作者: 紫气东来
发布时间: 编辑于 2024-05-21 22:07・上海
原文链接: https://zhuanlan.zhihu.com/p/699027700

文章概要

本文深入分析了大语言模型（LLM）计算中的内存墙问题，从硬件发展历史、AI模型需求、到Transformer架构的具体分析，系统阐述了为什么内存带宽成为LLM性能的关键瓶颈，以及FlashAttention等优化技术的实际效果。

核心内容

一、内存墙问题的本质

1.1 什么是内存墙

定义：当处理器的计算速度超过数据传输速度时，处理器必须等待数据，导致性能受限的现象。

核心矛盾：计算能力增长远快于内存带宽增长

1.2 硬件发展的不平衡（过去20年）

结论：算力提升了100万倍，但DRAM带宽只提升了100倍，互联带宽仅提升30倍。这种严重的不平衡是内存墙问题的根源。

1.3 多级缓存架构

解决方案：通过多级缓存（L1/L2/L3 → DRAM）缓解内存墙

主流GPU内存规格对比：

关键观察：

• H100使用HBM3，带宽是A100的近2倍
• RTX 4090虽然算力高，但带宽相对较低
• L2缓存容量差异巨大（40MB vs 96MB）

1.4 Roofline模型

计算强度（Arithmetic Intensity）：

• 定义：每读取1字节数据能执行多少次浮点运算（FLOPs/Byte）
• 意义：衡量算法对内存带宽的需求

Roofline曲线特点：

• 左侧（低计算强度）：Memory-bound，受带宽限制
• 右侧（高计算强度）：Compute-bound，受算力限制
• 拐点：平衡点，带宽和算力都被充分利用

四款GPU的Roofline对比：

• H100：最高算力和带宽，拐点最靠右
• A100：平衡型，适合多种workload
• RTX 4090/A6000：算力高但带宽相对低，拐点靠左

二、AI模型发展对计算的挑战

2.1 模型规模的指数增长

训练算力需求：

• CV/NLP/语音：每2年×15
• Transformer模型：每2年×750（增长最快）

模型参数量：

• 大语言模型：每2年×410
• 推荐系统：已达O(10) TB级别

硬件内存容量：

• 仅以每2年×2的速度增长

矛盾：模型增长速度远超硬件内存增长速度

2.2 训练内存需求

内存占用构成：

1. 模型参数：基础占用
2. 梯度：与参数量相同
3. 优化器状态：Adam需要2倍参数量（momentum + variance）
4. 激活值：前向传播的中间结果

总内存需求：通常是模型参数量的3-4倍

示例（GPT-3，175B参数）：

• 参数：175B × 2 bytes (FP16) = 350GB
• 梯度：350GB
• 优化器：700GB
• 激活值：取决于batch size
• 总计：>1.4TB（单卡无法容纳）

2.3 为什么带宽成为瓶颈

单芯片内部：

• Tensor Core等专用单元可在几个周期内完成大量计算
• 需要快速输入数据才能让计算单元持续工作
• 芯片内存带宽成为瓶颈

多芯片分布式：

• 模型超过单卡内存时，必须使用分布式并行
• 片间带宽 < 片内带宽
• 面临更严重的内存墙问题

三、Transformer架构的内存分析

3.1 Encoder vs Decoder

Encoder（如BERT）：

• 可并发处理所有token
• 计算强度较高
• 主要是Compute-bound

Decoder（如GPT）：

• 每次迭代生成一个token（自回归）
• 计算强度低
• 主要是Memory-bound

3.2 Decoder的两个阶段

Prefill阶段：

• 任务：并行处理输入的所有token
• 操作：计算并存储KV Cache
• 特点：
  • 计算量大
  • 主要是Compute-bound
  • 可以充分利用GPU算力

Decode阶段：

• 任务：每步生成一个新token
• 操作：读取完整的KV Cache
• 特点：
  • 计算量小（只计算1个token）
  • 严重Memory-bound
  • 瓶颈是显存带宽

关键问题：Decode阶段的计算强度极低，GPU算力严重浪费

四、优化技术的实际效果

4.1 FlashAttention在Prefill阶段

测试结果（实际算力利用率）：

关键发现：

• H100/A100虽然算力高，但因为带宽相对不足，利用率反而更低
• A6000/4090的带宽/算力比更合理，利用率更高
• 所有GPU都处于IO-bound区间

趋势：

• 实际算力随序列长度增加快速上升
• 很快达到稳定值（受带宽限制）

4.2 FlashAttention在Decode阶段

带宽利用率测试：

结果：

• 带宽利用率随长度增加快速上升
• 最大利用率约80%
• 明显处于IO-bound区间

意义：

• Decode阶段几乎完全受带宽限制
• 优化重点应该是减少内存访问
• FlashAttention通过融合操作减少HBM访问

4.3 PagedAttention优化

核心思想：高效存储和访问KV Cache

FlashInfer vs vLLM PagedAttention：

性能对比：

• FlashInfer在更短的序列长度就能达到更高的带宽利用率
• 意味着对GPU的整体利用率更高
• 在长序列场景下优势更明显

优化原理：

• 分页管理KV Cache
• 减少内存碎片
• 提高缓存命中率

技术要点总结

1. 内存墙的三个层次

硬件层：

• 算力增长 >> 带宽增长
• 20年：算力×1M，带宽×100

架构层：

• 单卡：片内带宽瓶颈
• 多卡：片间带宽瓶颈

应用层：

• Prefill：Compute-bound（但仍受带宽影响）
• Decode：严重Memory-bound

2. 为什么高端GPU利用率反而低

H100/A100利用率~30%的原因：

1. 算力提升太快（Tensor Core）
2. 带宽提升相对较慢
3. 带宽/算力比失衡
4. 大部分时间在等待数据

A6000/4090利用率~60%的原因：

1. 算力相对较低
2. 带宽/算力比更合理
3. 更容易达到平衡点

3. 优化策略

减少内存访问：

• FlashAttention：融合操作，减少HBM读写
• PagedAttention：高效管理KV Cache

提高数据复用：

• 利用Shared Memory
• 提高L2缓存命中率

算法层面：

• 降低计算复杂度
• 减少中间结果存储

实践意义

1. 硬件选择

训练场景：

• 优先选择高带宽GPU（H100 > A100）
• 考虑HBM vs GDDR6

推理场景：

• Decode阶段完全受带宽限制
• 带宽比算力更重要
• H100的3000 GB/s带宽是关键优势

2. 模型设计

考虑内存约束：

• 模型大小受单卡内存限制
• 需要分布式训练时，通信成为瓶颈

优化方向：

• 减少激活值存储（Gradient Checkpointing）
• 降低KV Cache大小（MQA、GQA）

3. 系统优化

Prefill阶段：

• 使用FlashAttention减少内存访问
• 批处理提高吞吐量

Decode阶段：

• PagedAttention管理KV Cache
• Continuous Batching提高利用率

总结

这篇文章揭示了LLM计算中的核心矛盾：算力增长远超带宽增长，导致高端GPU的算力无法充分利用。文章的价值在于：

1. 量化分析：用具体数据说明内存墙的严重性（算力×1M vs 带宽×100）
2. 架构分析：深入Transformer的Prefill和Decode阶段，指出各自的瓶颈
3. 实测数据：FlashAttention等优化技术的实际效果，H100利用率仅30%
4. 指导意义：为硬件选择、模型设计、系统优化提供依据

核心启示：在LLM时代，带宽比算力更重要。优化的重点应该是减少内存访问，而不是提高计算速度。