跳过正文
  1. Posts/

从参数量到 Peak Memory:模型显存到底应该怎么算?

目录

起因
#

上一篇讨论了模型调度中的 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 的粗估:

dtypebytes / param
FP324
FP16 / BF162
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 ComputeActive Experts
MoE Persistent MemoryResident 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 → C

B 可能不再完整 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:

$$E \times D = 50\text{M} \times 128 \times 2\ (\text{BF16}) = 12.8\text{GB}$$

仅一个中间 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 为例:

Componentbytes / 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 对比
#

ComponentInferenceTraining
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 场景特别值得注意的。

假设:

ModelPersistentRuntime Peak
A10GB6GB
B10GB8GB

如果简单把两者 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?。这两个问题共同构成模型调度最基础的两个资源约束。

回顾全文的认知链:

  1. Parameter memory 只是模型的 static footprint,并不等于模型运行时的 peak memory。
  2. GPU 是否 OOM,真正关心的是某一个时刻"同时还活着"的 tensor 有多少,而不是整个执行过程中一共产生过多少 tensor。
  3. 对于 scheduler,“模型占多少显存"不是一个固定常数,而是 Model + Workload + Implementation + Runtime 的函数。
  4. 多模型显存调度不仅取决于"每个模型多大”,还取决于它们的 runtime working set 是否在时间上重叠。

“这个模型到底占多少显存?“本身就是一个不够准确的问题。

更准确的问题应该是:在什么 workload、什么 implementation、什么 runtime 和什么 concurrency 条件下,它的 peak GPU footprint 是多少?

这才是 model scheduler 真正需要建模的东西。

参考资料
#

相关文章