Appearance
Q129 · 显存瓶颈：单卡无法容纳模型
假设团队要训练一个 100B 参数的语言模型，手里有一张标称 80 GB 显存的 GPU。有人说：“用 BF16，每个参数只占 2 字节，应该能训练吧？”先做一道乘法：100B 是 1000 亿个参数；1000 亿 × 2 字节 = 2000 亿字节，约 200 GB。仅把权重以 BF16 放入显存，就已经超过 80 GB；训练时还要放别的东西。这就是“单卡无法容纳模型”最直接的含义。
这里的 GB 用十进制计算：1 GB = 10 亿字节。某些工具按 GiB 显示，1 GiB = 2³⁰ 字节，因此同一份 200 GB 数据大约是 186 GiB。单位不同不能解释掉从 80 GB 到 200 GB 的差距。
术语：显存里究竟装着什么
| 名词 | 从零开始理解 | 在本文例子里对应什么 |
|---|---|---|
| 参数 / 权重 | 模型训练出来的数值。推理时也需要读取它们，才能计算下一步输出。 | 共 100B 个。 |
| BF16 | 用 16 位（2 字节）表示一个数的常见计算格式。它影响每个数占用的空间，并不改变参数的个数。 | 权重若全用 BF16 存放，约 200 GB。 |
| GPU 显存 | GPU 上供计算快速读取的内存，容量有限。 | 假设单卡标称 80 GB。 |
| 前向传播 | 输入经过模型，逐层算出预测结果。 | 训练时留下部分中间结果，以便之后算梯度。 |
| 激活值 | 前向传播的中间计算结果。反向传播通常要用到它们。 | 占用随批量、序列长度、层结构而变，不能只靠参数量算出。 |
| 反向传播 / 梯度 | 训练根据误差计算“每个参数该往哪个方向改”；梯度是这个方向及大小的信息。 | 通常需要额外存储，具体字节数取决于精度与实现。 |
| 优化器状态 | 优化器更新参数时保存的历史信息。Adam 常保存一阶、二阶矩，可以理解为两本关于过往梯度的账。 | 若这两份状态都为 FP32，合计每参数 8 字节。 |
| 主权重副本 | 一些混合精度训练方案中，用 FP32 保存的高精度参数版本，供稳定更新。 | 是否存在、存在哪里由训练栈决定，不能一概而论。 |
| 临时缓冲 | 通信、矩阵计算、框架分配器等在某一步暂时需要的空间。 | 峰值可能高于平时占用。 |
| OOM | Out of Memory，程序要分配显存却已没有足够空间。 | 常在加载、第一次前向、反向或优化器更新时出现。 |
图把每类东西画成书，是为了强调权重只是第一摞。实际计算时不一定所有数据永久同时驻留：实现会释放、重算、移动或分片。做容量估算时必须说明所采用的训练方案，不能把教学示意当成每个框架固定的内存账单。
为什么训练比“读取模型”占得多
先区分推理和完整参数训练。推理是拿已有权重回答问题，通常不用保存每个参数的梯度，也不用保存 Adam 状态；训练则要根据答案误差反向传播并更新权重。因此“一个模型能加载做推理”并不意味着“这张卡能训练它”。反过来，如果连 200 GB 的 BF16 权重都无法放进 80 GB 单卡，常规的单卡直接加载推理同样不成立，除非再用量化、卸载等额外方法。
给训练写一份教学账单。假设某套混合精度 Adam 方案对每个参数同时持有：
- BF16 计算权重：2 字节。
- FP32 主权重：4 字节。
- BF16 梯度：2 字节。
- Adam 的两份 FP32 状态：4 + 4 = 8 字节。
合计 16 字节/参数。100B × 16 字节约为 1.6 TB 的模型状态。再加激活、临时缓冲和框架开销，峰值更高。这里的“16 字节”只是这套明确假设的算术结果：有的实现将梯度存为 FP32，有的没有独立的 FP32 主权重，有的优化器或分片策略又不同。不能把它说成所有 BF16 训练的通用常数。DeepSpeed 的官方内存估算文档也把参数、梯度、优化器状态分开，并明确说明其 Adam 状态假设。DeepSpeed Memory Requirements
**激活值不能漏掉。**例如一次输入有很多 token，模型逐层处理时会生成中间结果；为了在反向传播中求梯度，部分结果要保留。批量大小变大、序列变长，激活开销往往也变大。它并非简单的“100B × 某固定字节数”，还和网络结构、注意力实现、是否重算等有关。PyTorch 官方对激活重算的解释正是：少存部分前向中间张量，反向时重新计算，以计算量换显存。PyTorch：Activation Checkpointing
多加几张卡，为什么仍可能 OOM
“8 张 80 GB 卡，合计 640 GB，权重 200 GB，够了吧？”关键要问：一份模型如何分布在八张卡上？
最容易想到的数据并行，是每张卡各放一整份相同模型，给不同卡分不同训练样本，再汇总梯度。这有利于处理更多数据，但每张卡仍要装完整权重；100B BF16 权重本身就约 200 GB，依旧塞不进单张 80 GB 卡。把八张卡的容量直接相加，只在训练状态确实可以跨卡分布时才有意义。
要解决这个问题，可以把模型状态分片：不是每张卡保存全部，而是按规则把参数、梯度和优化器状态分别放到多张卡上。ZeRO 的各阶段分片范围不同；PyTorch FSDP 的 FULL_SHARD 会分片参数、梯度与优化器状态，在需要计算某一部分时临时聚合参数，算完再释放或重新分片。DeepSpeed ZeRO 官方教程 · PyTorch FSDP 文档
按前述1.6 TB 教学账单做极其粗略的平均：若 8 卡均匀分，单卡仅模型状态也约 200 GB，仍超过 80 GB；若 32 卡均匀分，平均约 50 GB/卡，看似有空间，但激活、临时聚合的一层权重、通信缓冲、框架占用仍未计入，不能因此宣布一定能跑。真实分片并非“总字节数除以卡数”就结束，峰值和最大一层的大小也重要。DeepSpeed 官方的估算器特意要求最大的层参数量，正是因为计算中可能需要临时聚合。DeepSpeed Memory Requirements
这还带来新的代价：分片卡之间要交换参数和梯度。显存压力下降，不代表训练一定更快；网络带宽不足或频繁通信可能拖慢计算。需要把“每卡装得下”与“整体训练吞吐量可接受”分开验收。
从一次失败定位真正的瓶颈
设团队在 32 张卡上启动训练，程序初始化成功，跑到第一轮反向传播才 OOM。只看“平均每卡模型状态 50 GB”会误以为程序有故障，实际上可能是：前向留下的激活尚未释放，反向又需要临时聚合参数、产生梯度和通信缓冲，峰值越过 80 GB。
排查顺序可以这样走：
- **写清配置。**记录模型参数量、数值精度、优化器、卡数、每卡批量、输入序列长度、是否分片、是否重算激活。少一个条件，就难以解释两次运行的差异。
- **找出 OOM 阶段。**加载时失败，先怀疑权重/初始化方式；前向失败，注意激活与输入尺寸；反向失败，注意梯度、激活峰值及参数临时聚合；更新时失败，注意优化器状态首次创建。这里只是排查线索，不能凭阶段直接下结论。
- **记录峰值和分配器状态。**PyTorch 的
torch.cuda.max_memory_allocated()可查看经 PyTorch 分配的张量峰值；memory_reserved()反映缓存分配器管理的空间。快照工具能定位分配栈，但某些直接经 CUDA 或第三方库申请的显存不在 PyTorch 快照视野内，要结合设备监控。PyTorch CUDA 内存说明 · PyTorch 内存快照文档 - **按原因调整。**激活占主要峰值时，可缩小每卡批量或序列长度、使用激活重算；模型状态占主导时，考虑分片、卸载或适配更大显存；通信开销成为瓶颈时，则要看网络与分片粒度。调整一项后复测峰值和训练速度，确认收益与代价。
有时有人一看到 OOM 就调用 empty_cache()。它能释放分配器里未被张量使用的缓存，不能让仍在使用的权重、梯度或激活凭空消失。若实际活跃数据已经超过容量，清缓存解决不了根因。PyTorch CUDA 内存管理说明
什么情况下单卡仍能工作
“单卡放不下”指当前模型规模、精度、训练方式及显存容量组合下无法直接执行。换一种任务可能改变结论。例如只做推理且对权重量化，每参数占用可能降低；使用 CPU/NVMe 卸载，会把部分状态移出 GPU，代价是数据搬运和速度；只训练少量适配参数与完整参数训练也不是同一份账单。这些手段各有质量、延迟、吞吐与实现约束，不能只看容量数字。
因此面试中最好先问清“是推理、全量训练，还是参数高效微调？”以及“单卡显存多少、精度是什么、序列多长？”再给估算。若题目没有条件，给一个清楚标注假设的算例，比报出一个看似精确的固定数字更可靠。
面试时怎么回答
单卡显存瓶颈首先要算权重下限。例如 100B 参数按 BF16 每参数 2 字节，权重约 200 GB，单张 80 GB 卡连权重都装不下。训练还需要梯度、Adam 等优化器状态、前向激活与临时缓冲；某种带 BF16 权重、FP32 主权重、BF16 梯度和两份 FP32 Adam 状态的配置，模型状态约 16 字节/参数，但这个数字依赖实现。多张卡直接做普通数据并行仍是每卡一份完整模型，所以不能简单把显存相加。要让大模型训练起来，通常按瓶颈采用参数/梯度/优化器分片、激活重算或卸载，并检查临时聚合与反向传播的峰值显存、通信和吞吐量。我会先明确训练方式和精度，再估算静态状态，最后用内存峰值工具验证，而不是只凭参数量猜测。
如果被追问“32 张 80 GB 卡为什么还会 OOM”，说出上面的 1.6 TB ÷ 32 = 50 GB 只是模型状态平均值：激活、单层临时聚合、通信缓冲和不均匀分布可把峰值推过 80 GB。若被追问“数据并行为什么不够”，回答“它复制模型，增加处理样本的能力，但不自动拆开单卡必须装下的模型状态”。
资料依据
- DeepSpeed：Memory Requirements：参数、梯度、Adam 状态的估算与最大层影响。
- DeepSpeed：ZeRO 教程：分片训练状态的原理。
- PyTorch：FSDP 分片策略：
FULL_SHARD对参数、梯度和优化器状态的处理。 - PyTorch：激活重算技术：激活与计算量的权衡。
- PyTorch：CUDA 内存语义：allocated、reserved 与峰值监测。