大语言模型中的记忆瓶颈 | AI生成和翻译
问题： 为什么内存（以及统一内存）对LLM训练和推理如此关键？
回答：
根本原因在于计算与内存带宽之间的不匹配。现代GPU/加速器的计算能力远超内存带宽，因此LLM训练/推理期间的瓶颈几乎总是内存移动，而非数学运算。
核心问题
像RTX 4070或MI300X这样的GPU可执行约千万亿次浮点运算，但内存带宽有限：
| 设备 | 峰值计算能力 | 内存带宽 | 计算与带宽比 |
|---|---|---|---|
| RTX 4070 | ~29 TF32-TFLOPS | 432 GB/s | ~67 |
| MI300X | ~360 TFLOPS (FP8) | 5.3 TB/s | ~68 |
要让GPU保持忙碌状态，你需要足够快地将数据移入和移出内存，以喂饱计算单元。
如果每个FLOP对应67字节的数据，而每次计算耗时1纳秒，那么你需要67字节 × 1纳秒 = 67纳秒的内存延迟容忍度。GDDR6无法维持这一水平。HBM（高带宽内存）正是为解决此问题而生：MI300X上的HBM3可提供5.3 TB/s的速度，而GDDR6约为400 GB/s——带宽提升了10倍。
训练：激活内存暴增
在前向传播过程中，你必须将所有中间激活值保留在内存中，以供反向传播使用。对于Transformer：
激活内存 = 批次大小 × 序列长度 × 隐藏维度 × 层数 × 每参数字节数
示例（13B模型，bsz=1，seq=4096）：
= 1 × 4096 × 5120 × 40 × 2 字节（fp16）
≈ 每层仅1.6 GB激活值
× 40层 ≈ 64 GB
这与序列长度呈线性关系。上下文越长，所需内存越多。这就是为什么：
- Flash Attention通过融合内核将其最小化（无需中间存储）
- 梯度检查点用计算换内存（重新计算激活值，而非存储）
- 你只需HBM或大容量VRAM就能容纳标准训练
推理：KV缓存占主导
推理期间，你会缓存键/值向量，以避免重复计算注意力：
KV缓存大小 = 批次大小 × 序列长度 × 隐藏维度 × 2 × 层数 × 字节数
对于13B模型，bsz=32，seq=4096：
= 32 × 4096 × 5120 × 2 × 40 × 2 字节（fp16）
≈ 52 GB
这随序列长度和批次大小呈二次方增长。如果你要服务128个并发用户且上下文为8K，就需要大量内存。这就是为什么：
- 量化（int8/fp8）将KV缓存大小减少50-75%
- MQA（多查询注意力）按头部数量因子减少KV缓存
- 推理期间你受内存限制，而非计算限制
为什么统一内存很重要
NVIDIA的统一内存架构（UVA）允许CPU和GPU访问同一虚拟地址空间，并支持两者间的自动分页。这听起来很棒，但有一个关键缺陷：
分页速度很慢。 如果激活值溢出到系统RAM，你将通过PCIe（PCIe 4.0为16 GB/s）而非HBM（5.3 TB/s）移动数据。带宽损失300倍。
然而，统一内存在以下场景中很有用：
- 梯度累积工作流，你通过精细控制显式移动数据
- CPU-GPU异构计算（在LLM训练中很少见）
- 动态内存管理，当你事先不知道数据大小时
对于纯LLM训练，显式内存管理（在GPU上分配，并保持在那里）比依赖统一内存的自动分页更优。
针对你的实际情况
MI300X（192 GB HBM3）： 由于巨大的带宽和容量，你可以在全精度下训练更大模型并容纳更大批次。这就是它为何对760M GPT-2训练如此强大——你拥有充足的余量。
RTX 4070（12 GB VRAM）： 你受到限制。对于严肃工作，你需要：
- 混合精度（fp16/tf32）
- 梯度检查点
- 更小的批次大小或序列长度
- 量化（QLoRA用于微调）
优化优先级：
- 最小化数据移动（融合内核、Flash Attention）
- 使用最快可用内存（HBM > GDDR6 > 系统RAM）
- 选择合适的数据类型（推理用fp8/int8 KV缓存，训练用fp16）
- 批次大小调优（在不超出内存的情况下最大化GPU利用率）
统一内存的陷阱：现代推理引擎（vLLM、TensorRT）会显式管理内存，而非依赖自动分页——它们了解内存访问模式并会进行战略性预取。