起因#
上一篇讨论了模型调度中的 Performance Constraint:一个模型放到 GPU 上以后,到底跑多快?
这一篇讨论另一个最基础的约束:Capacity Constraint。
也就是:一组模型到底能不能同时放进这张 GPU?
flowchart TD
S["Model Scheduling"]
S --> P["Performance Constraint
How fast?"]
S --> M["Capacity Constraint
Does it fit?"]
最直觉的 memory constraint 是:
$$\sum_i \text{Memory}_i \le \text{GPU Capacity}$$看起来和普通 Bin Packing 很像。但这里真正困难的是:\(\text{Memory}_i\) 到底是什么?
最开始,人很容易把它理解成:
$$\text{ParameterCount} \times \text{BytesPerParameter}$$比如:一个 BF16 模型有 1B 参数,大约占 2GB。
这个估算当然有价值。但是——参数显存只能回答"模型静态放在那里需要多少空间",不能回答"模型真正跑起来时峰值需要多少显存"。
这就是全文的起点。
Parameter Memory 只是 Static Footprint#
先从最简单的部分开始。
假设模型有 \(P\) 个参数,每个参数 \(s\) bytes,那么:
$$M_{\text{weights}} \approx P \times s$$常见 dtype 的粗估:
| dtype | bytes / param |
|---|---|
| FP32 | 4 |
| FP16 / BF16 | 2 |
| INT8 | ~1 |
| INT4 | ~0.5 |
实际量化模型还有 scale、zero point、alignment 等额外开销,所以这些只是 first-order estimation。
但到这里就可以得到第一个重要结论:
Parameter memory 只是模型的 static footprint,并不等于模型运行时的 peak memory。
Persistent Memory#
可以把 weights、persistent buffers、constant tensors 统称为 Persistent Memory。
它的特点是:模型一旦部署在 GPU 上,这部分显存通常长期存在。
对 scheduler 来说,这部分特别像传统 Bin Packing 中真正的 item size——多个模型 co-location 时,\(\sum_i \text{Persistent}_i\) 通常都必须同时存在。
MoE:Compute 和 Memory 是两个独立的资源维度#
沿用上一篇的 MoE 例子。假设 64 experts、top-2 routing。
上一篇讨论 FLOPs 时,一个 token 只执行 top-2 experts,所以 compute 主要取决于 active experts。
但显存不一样。如果所有 64 个 expert 都 resident 在 GPU 上,persistent memory 取决于 resident experts,而不是 top-k。
| 主要由什么决定 | |
|---|---|
| MoE Compute | Active Experts |
| MoE Persistent Memory | Resident Experts |
MoE 很好地说明了:执行多少计算和需要常驻多少显存,是两个独立的资源维度。
到这里,即使 persistent memory 完全已知,也仍然不能判断是否会 OOM。因为真正执行模型时,还会产生大量额外的 runtime 显存。
Runtime Working Set 与 Tensor Lifetime#
执行模型时还会产生:input、activations、intermediate tensors、temporary workspace。可以 high level 地把它们统称为 Runtime Working Set。
于是模型显存变成:
$$\text{Persistent Memory} + \text{Runtime Working Set}$$但这里不能简单把所有 runtime tensors 加起来。因为 tensor 有生命周期。
假设执行过程中依次产生 A、B、C 三个 tensor:
- A = 1GB
- B = 2GB
- C = 3GB
如果简单求和:\(1 + 2 + 3 = 6\text{GB}\)。
但实际执行时:
- 计算 B 时,需要 A + B = 3GB
- 到计算 C 时,如果 A 已经不会再使用,就可以释放
- 此时只需要 B + C = 5GB
所以:
$$\text{Peak Memory} = \max(1+2,\ 2+3) = 5\text{GB}$$而不是 6GB。
gantt
title Tensor Lifetime
dateFormat X
axisFormat %s
section Tensors
A_1GB :a, 0, 2
B_2GB :b, 1, 4
C_3GB :c, 3, 5
Peak memory 出现在 overlap 最大的位置。在这个例子里,是 B 和 C 同时存活的那段时间。
正式地说:
$$M_{\text{peak-live}} = \max_t \sum_{x \in \text{Live}(t)} \text{Size}(x)$$GPU 是否 OOM,真正关心的是某一个时刻"同时还活着"的 tensor 有多少,而不是整个执行过程中一共产生过多少 tensor。
那什么决定 tensor 的 lifetime?
Computation Graph 决定 Lifetime#
看一个最经典的例子:residual connection。
$$Y = F(X) + X$$
flowchart LR
X --> F
F --> Add
X --> Add
Add --> Y
这里 X 虽然已经被 F 使用过一次,但不能释放——因为后面的 Add 还要使用它。
X 的 lifetime 被 Add 节点拉长了。
Residual 只是最简单的例子。任何 branch、multi-output、skip connection 都可能产生类似的效果。核心是同一个:
Tensor lifetime 不仅由 tensor size 决定,也由 computation graph 的 dependency 决定。
Memory Reuse#
一旦一个 tensor 的 lifetime 结束,它占用的 memory 就有机会被后面的 tensor 复用。所以 \(\sum \text{TensorSize}\) 往往远大于实际需要的 peak memory。
这和 compiler 里的 liveness analysis 非常类似:编译器分析每个变量什么时候 “活着”,然后把 lifetime 不重叠的变量分配到同一个寄存器。GPU memory planning 做的事情本质上一样——让不同 tensor 的 storage 在不重叠的 lifetime 之间复用。
Operator Fusion:两个调度维度的交叉点#
上一篇讨论过 operator fusion 可以减少 memory traffic、改善 latency。这里补充另一面。
假设有一个执行链:
1A → op1 → B → op2 → C如果 B 必须完整 materialize 到 HBM,它会增加 runtime working set。
如果 op1 和 op2 fuse 成一个 kernel:
1A → fused_op → CB 可能不再完整 materialize。于是 peak memory 也下降了。
Operator fusion 有时同时改善两个调度维度:既减少数据搬运(latency),也减少中间 tensor 的 materialization(memory)。
但并不是所有显存都能从 computation graph 的 tensor shape 中直接看到。
Workspace 与 Implementation Matters#
某些 operator / library 在执行时还需要额外的 scratch buffer,通常叫 workspace。比如某些 GEMM algorithm、sort、reduce、sparse operation、MoE routing 的 dispatch buffer。
这些 workspace 不出现在模型的数学定义里,但确实占用显存。
同一个数学模型,换一个实现或 kernel,peak memory 也可能变化,因为它们使用的 temporary workspace 不一样。
GNN:workload 决定显存#
GNN 是解释这件事的最好例子。
GNN 的 parameters 可能并不大——几层 Linear 而已。但运行时需要保存 graph 的完整表示:node features、edge indices、edge features、node embeddings、temporary messages。
算一笔账。假设 \(N = 1\text{M}\) nodes,feature dimension \(D = 256\),BF16:
$$1\text{M} \times 256 \times 2 \approx 512\text{MB}$$再假设 \(E = 50\text{M}\) edges,COO 格式用两个 int64:
$$50\text{M} \times 2 \times 8 = 800\text{MB}$$仅 graph input 就已经超过 1GB。模型 parameter 反而可能只有几 MB。
对于 GNN,显存经常主要由 workload,而不是模型参数决定。
Message Materialization#
再看一个非常典型的 GNN aggregation 实现:
1msg = src_feature[edge_src] # gather
2dst.scatter_add_(edge_dst, msg) # scatter如果真的 materialize msg 这个中间 tensor:
仅一个中间 tensor 就可能爆显存。
而如果 implementation 可以 fuse gather + aggregation,就可能避免完整生成 \(E \times D\) 的 message tensor。
同一个数学模型,并不意味着同一个 Peak Memory;实现方式会改变中间 tensor 是否 materialize。
Inference Memory Mental Model#
到这里可以给出一个整体 conceptual model:
$$M_{\text{infer, peak}} \approx M_{\text{persistent}} + M_{\text{peak runtime working set}} + M_{\text{runtime overhead}}$$其中:
$$M_{\text{peak runtime working set}} = \max_t \left( \text{Input}_t + \text{Activation}_t + \text{Workspace}_t \right)$$核心 mental model 就是:
$$\boxed{\text{Peak Memory} = \text{Persistent Memory} + \text{Peak Live Working Set}}$$需要注意的是:这是 conceptual model,不是 allocator accounting identity。实际系统里 persistent 和 runtime 的边界不总是清晰的。
flowchart TD
PM["Parameter Memory"] --> SF["Static Footprint"]
SF --> PersM["Persistent Memory"]
PersM --> RWS["+ Runtime Working Set"]
RWS --> PLM["Peak Live Memory"]
PLM --> AR["Allocator / Runtime"]
AR --> PF["Physical GPU Footprint"]
但到这里还没结束。Logical memory demand 还不等于 physical footprint。
从 Logical Memory 到 Physical GPU Footprint#
到这里读者可能觉得:那把 peak live tensor memory 算出来就完事了。
但实际跑一个模型然后看 nvidia-smi,显示的数字经常和自己估算的 tensor memory 对不上。
因为中间还隔着 allocator 和 runtime。
PyTorch Allocated vs Reserved#
如果以 PyTorch 为例:
1torch.cuda.memory_allocated()表示当前 live tensors 真正在使用的 allocator memory。
1torch.cuda.memory_reserved()表示 PyTorch caching allocator 已经向 CUDA 获取、目前仍然管理的 memory。
概念上:
$$\text{allocated} \subseteq \text{reserved}$$不是 allocated + reserved。
flowchart TD
R["Reserved Memory
PyTorch 从 CUDA 获取的总量"]
R --> A["Allocated
当前 live tensors 使用"]
R --> C["Cached Free Blocks
已释放但未还给 CUDA"]
也就是说:tensor 已经释放,不代表 allocator 一定马上把对应 GPU memory 还给 driver。PyTorch 的 caching allocator 会把释放的 block 留在自己的 free list 里,方便下次分配时快速复用。
nvidia-smi 为什么又不一样#
除了 PyTorch allocator 管理的部分,GPU 上还有:
- CUDA context 本身的开销
- CUDA libraries(cuBLAS、cuDNN 等)的 internal allocation
- NCCL communication buffer
- CUDA Graph pool
- 其他 custom extension / runtime allocations
所以 nvidia-smi 显示的 footprint 往往比 memory_reserved() 还要大。
“模型逻辑上需要多少显存"和"这个进程实际上占住多少 GPU 显存"不是同一个问题。
由此可以定义两个 high-level 概念:
- Logical Memory Demand:用于理解模型和 workload 本身。对应前面讨论的 persistent memory + peak live working set。
- Physical GPU Footprint:用于真正做 deployment、co-location、scheduler capacity decision。对应
nvidia-smi看到的那个数字。
对于 scheduler,最终更重要的是 Physical GPU Footprint。
Headroom#
即使理论上 estimated footprint < GPU capacity,也不意味着一定可以把 GPU 塞到 100%。
真实系统还有 allocator fragmentation、runtime fluctuations、workspace variation、dynamic workload、library allocation 等不确定性。所以 scheduler 通常需要预留一定的 headroom。
但不要拍脑袋给一个"永远预留 10%“的 universal rule。
Headroom 应该来自真实 workload profiling,而不是拍脑袋给一个固定比例。
Training:为什么 Memory Model 明显不同#
以上讨论都以 inference 为主。Training 的 memory model 有明显不同,值得单独说。
Model State#
训练比推理多出来的 model state 主要是 gradients 和 optimizer states。以 mixed-precision Adam 为例:
| Component | bytes / param |
|---|---|
| BF16 weight | ~2 |
| Gradient | ~2 |
| Adam m(first moment) | 4 |
| Adam v(second moment) | 4 |
| FP32 master weight(如果有) | +4 |
所以经常看到 12~16 bytes/parameter 这样的经验估算。
但需要注意:12~16 bytes/parameter 只是某些 mixed-precision Adam 配置下的粗估,不是 universal law。 换 optimizer(比如 SGD、Adafactor)、换 precision scheme、加上 ZeRO/FSDP sharding,单卡 footprint 都会变。这里不展开。
Saved Activation 才是关键#
比 optimizer state 更重要的一件事:backward 改变了 activation 的 lifetime。
Inference 时,一个 activation 最后一次被使用后通常可以立即释放。
Training 时,backward 可能还需要 forward 中保存的中间结果来计算 gradient。
flowchart LR
F1["Forward L1"] --> F2["Forward L2"] --> F3["Forward L3"] --> F4["Forward L4"]
F4 --> B4["Backward L4"] --> B3["Backward L3"] --> B2["Backward L2"] --> B1["Backward L1"]
很多 forward activation 不是到下一层结束就死亡,而是要一直活到对应的 backward pass。L1 的 activation 可能要活到 B1 才能释放——也就是整个 forward + 大半个 backward 期间都得常驻。
Training 显存大,不只是因为多了 gradient 和 optimizer state,更重要的是 backward 改变了 activation 的 lifetime。
Activation Checkpointing:Compute 和 Memory 的 Trade-off#
这里自然引出 activation checkpointing(或 gradient checkpointing)。
正常训练:保存更多 activation,backward 直接使用。
Checkpointing:少保存一些 activation,backward 时重新做一次 forward 来重新计算。
$$\text{Memory} \downarrow \quad \Longleftrightarrow \quad \text{Compute} \uparrow$$这正好把上一篇讨论的 FLOPs / Latency 和这一篇讨论的 Memory 连接起来。
ML 系统优化经常不是单独优化某一种资源,而是在 Compute 和 Memory 之间做 trade-off。
Training vs Inference 对比#
| Component | Inference | Training |
|---|---|---|
| Weights | 有 | 有 |
| Runtime activations | 有 | 有 |
| Saved forward state | 少 | 多 |
| Gradients | 无 | 有 |
| Optimizer states | 无 | 有 |
| Workspace | 有 | 有 |
| Backward workspace | 无 | 有 |
回到 Scheduler:Does it fit?#
model.memory 是一个危险的抽象#
最开始做调度时可能会写:
1class ModelProfile:
2 memory_bytes: int但经过前面的讨论应该很清楚:这个字段本身语义不清。它可能表示:
- parameter memory
- persistent memory
- batch=1 时的 peak
- max allocated
- reserved
- nvidia-smi footprint
完全不是一回事。
在 scheduler 里,任何 model.memory 字段都应该明确它到底代表哪一种 memory metric。
显存需求不是常数#
对于 scheduler,“模型占多少显存"不是一个固定常数,而是 Model + Workload + Implementation + Runtime 的函数。
一个更合理的 mental model:
1class ModelMemoryProfile:
2 persistent_memory: int
3 runtime_memory: Callable[[Workload], int]不同类型的模型,runtime memory 依赖的 workload 参数不一样:
- MLP:\(\text{Memory} = f(\text{batch})\)
- MoE:\(\text{Memory} = f(\text{batch}, \text{routing})\)
- GNN:\(\text{Memory} = f(N, E, D)\)
模型显存需求本质上是 Model + Workload 的函数,而不是 Model 的一个固定属性。再加上 implementation 和 allocator/runtime 的影响,情况更复杂。
Concurrency:多模型 peak 不能机械相加#
最后一点,也是 scheduler 场景特别值得注意的。
假设:
| Model | Persistent | Runtime Peak |
|---|---|---|
| A | 10GB | 6GB |
| B | 10GB | 8GB |
如果简单把两者 peak memory 相加:\((10+6) + (10+8) = 34\text{GB}\)。
但如果两个模型永远不会同时 inference:
- persistent 必须同时存在:\(10 + 10 = 20\text{GB}\)
- runtime working set 可以错峰使用:\(\max(6, 8) = 8\text{GB}\)
整体可能接近 \(20 + 8 = 28\text{GB}\),而不是 34GB。
多模型显存调度不仅取决于"每个模型多大”,还取决于它们的 runtime working set 是否在时间上重叠。
最终 Scheduler Constraint#
flowchart TD
GPU["GPU Capacity"]
GPU --> Pers["Persistent Region"]
GPU --> RW["Runtime Working Set"]
GPU --> H["Headroom"]
Pers --> MA["Model A Persistent"]
Pers --> MB["Model B Persistent"]
RW --> W["Peak of Concurrent
Runtime Workloads"]
这就是全文最终的 scheduler conceptual constraint:
$$\boxed{\sum_i \text{Persistent}_i + \text{PeakRuntime}(\text{ConcurrentWorkloads}) + \text{Headroom} \le \text{GPU Capacity}}$$Persistent memory 更像长期占用的空间;runtime memory 更像执行过程中的 working set;是否并发决定 working set 如何叠加。
结论#
上一篇回答 How fast?,这一篇回答 Does it fit?。这两个问题共同构成模型调度最基础的两个资源约束。
回顾全文的认知链:
- Parameter memory 只是模型的 static footprint,并不等于模型运行时的 peak memory。
- GPU 是否 OOM,真正关心的是某一个时刻"同时还活着"的 tensor 有多少,而不是整个执行过程中一共产生过多少 tensor。
- 对于 scheduler,“模型占多少显存"不是一个固定常数,而是 Model + Workload + Implementation + Runtime 的函数。
- 多模型显存调度不仅取决于"每个模型多大”,还取决于它们的 runtime working set 是否在时间上重叠。
“这个模型到底占多少显存?“本身就是一个不够准确的问题。
更准确的问题应该是:在什么 workload、什么 implementation、什么 runtime 和什么 concurrency 条件下,它的 peak GPU footprint 是多少?
这才是 model scheduler 真正需要建模的东西。
参考资料#
- 从 FLOPs 到 Latency:GPU 推理性能到底由什么决定?(本文姊妹篇)
- CS336: Language Modeling from Scratch — Assignment 1 (Resource Accounting)
- NVIDIA. CUDA C++ Programming Guide — Memory Management.
- Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020.(关于 model state memory 的分析)
- Chen et al. Training Deep Nets with Sublinear Memory Cost. 2016.(activation checkpointing 的经典论文)