Skip to main content
  1. Posts/

从线上模型切换问题理解 PyTorch CUDA Caching Allocator:为什么显存管理需要理解生命周期

Table of Contents
Note: This article is available in Chinese only. 本文暂无英文版本。 View original

起因
#

我们的线上推理系统中存在大量模型——MLP、GNN、MoE,以及它们的各种组合。这个 workload 有几个重要特点:

  • 单张 GPU 上可能同时存在多个模型;
  • 很多模型的 architecture 相同或者高度相似;
  • 模型会随着业务时间轴不断上线、下线和切换;
  • 模型切换过程对 latency 比较敏感。

随着模型数量增加以及动态上下线越来越频繁,我们逐渐观察到两个问题:GPU memory fragmentation,以及 model loading / switching latency。

尤其是模型加载过程中,大量 weight tensor 都需要建立对应的 CUDA storage。直觉上可能首先怀疑:是不是频繁 cudaMalloc / cudaFree 导致的?

但很快会发现——PyTorch 本身早就有 CUDA Caching Allocator。也就是说,tensor allocation 并不等于每次 cudaMalloc。PyTorch 已经在帮助我们缓存和复用 GPU memory。

于是出现本文真正的问题:

既然 PyTorch 已经使用 Caching Allocator,为什么在频繁模型切换场景下,我们仍然会遇到 fragmentation 和 allocation latency?

为什么需要 Caching Allocator
#

先建立一个基本事实。

cudaMalloccudaFree 不是轻量操作。根据 NVIDIA 官方文档:

Managing memory allocations using cudaMalloc and cudaFree causes the GPU to synchronize across all executing CUDA streams.

也就是说,cudaMalloc / cudaFree 涉及 CUDA Runtime / Driver 层面的交互,可能导致 device-wide synchronization,带来较高的 latency 和 latency jitter。

如果深度学习框架真的按照这样的方式运行:

1Tensor 创建 → cudaMalloc
2Tensor 销毁 → cudaFree

那么训练和推理中大量临时 tensor 会产生巨大的 allocator overhead。

因此 PyTorch 在 tensor 与 CUDA Runtime 之间增加了一层 CUDACachingAllocator。核心思想是:

从 CUDA Driver 获取较大块显存,然后由 PyTorch 自己管理、切分、缓存和复用。

用一个简单类比:PyTorch 从 CUDA Driver 批发显存,然后自己负责零售。

PyTorch CUDA Caching Allocator
#

以下重点讨论 allocator 长期稳定的设计思想,而不是依赖某一个版本的内部常数。

Cache 的基本思想
#

第一次 allocation:

1Tensor 申请 memory
23CUDACachingAllocator
45没有合适的 cached block
67向 CUDA 获取 segment

Tensor 生命周期结束后:

1Tensor 释放
23block 返回 allocator
45标记为 cached / free

下一次有新的 tensor 申请 memory:

1new Tensor 申请 memory
23在 cache 中找到合适的 free block
45直接复用

而不是重新执行 cudaMalloc

整个过程用流程图表示:

flowchart TD
    R["Tensor 申请 N bytes"] --> C{"allocator 中有
合适的 cached block?"} C -->|"有"| Reuse["复用 cached block"] C -->|"没有"| Alloc["向 CUDA 申请新 Segment"] Alloc --> Split["从 Segment 中切出 block"] Split --> Ret["返回 block 给 Tensor"] Reuse --> Ret Free["Tensor 生命周期结束"] --> Mark["Block 标记为 free
留在 allocator cache"] Mark --> TryMerge["尝试与相邻 free block 合并"]

Caching Allocator 的核心收益是:把很多 CUDA-level allocation 转化为 allocator 内部的 block bookkeeping 和 reuse。

Segment 与 Block
#

Allocator 的基本抽象可以分成两层。

Segment 是 allocator 从 CUDA 获取的一大块连续 memory。可以通过 cudaMalloc,也可以通过 CUDA virtual memory API(cuMemAddressReserve + cuMemCreate + cuMemMap,即 expandable segments 机制)获取。

Block 是 segment 内部的子区域,每个 block 对应一次 tensor allocation。

1Segment
2┌───────────────────────────────────────────┐
3│ Block A (allocated) │ Block B (allocated) │ Free Block │
4└───────────────────────────────────────────┘

当一个 tensor 申请 memory 时,allocator 在 free block 中找到合适的。如果 free block 比请求大,就 split——前半部分分配出去,后半部分变成新的 free block:

1申请 80MB,现有 200MB free block:
2┌───────────┬──────────────────┐
3│ 80 MB     │ 120 MB           │
4│ allocated │ cached / free    │
5└───────────┴──────────────────┘

当相邻的 block 都变成 free 时,allocator 会尝试 merge / coalesce——把它们合并回一个更大的 free block。

PyTorch 还按 size 把 block 分成两个 pool:small_blocks(<= 1 MiB)和 large_blocks(> 1 MiB)。两个 pool 完全独立,各自有各自的 segment 和 free list。

allocated、reserved 与 empty_cache
#

PyTorch 中两个很容易混淆的指标:

1torch.cuda.memory_allocated()   # 当前 live tensors 使用的 memory
2torch.cuda.memory_reserved()    # PyTorch 从 CUDA 获取的、目前仍管理的 memory

概念上 allocated 是 reserved 的子集:

1Reserved(PyTorch 从 CUDA 获取的总量)
2┌──────────────────────────────────┐
3│ Allocated(当前 live tensors)    │
4│                                  │
5│ Cached Free(已释放但未还给 CUDA)│
6└──────────────────────────────────┘

之前的文章讨论过这个区别。这里再强调一次:

Tensor 生命周期结束,不代表显存立即返回 CUDA Driver。

所以 del tensor 以后,nvidia-smi 显示的显存不一定下降。

torch.cuda.empty_cache() 释放的是 allocator 当前没有使用的 cached memory。它不会释放正在被 tensor 使用的 memory,也不应该在正常 inference / training loop 中频繁调用——那样反而会破坏 caching 带来的性能收益。

Caching 没有消灭 Fragmentation
#

CachingAllocator 解决了频繁 cudaMalloc / cudaFree 的问题。但它带来了另一个经典 allocator 问题:fragmentation。

这里需要区分两种 fragmentation。

Internal fragmentation:分配出去的 block 比实际请求大,多出的部分被浪费。PyTorch 通过 size rounding 和 split 来控制这个问题。

External fragmentation:free memory 总量足够,但被分散成很多不连续的小 block,无法满足一个较大的 contiguous allocation。

重点讨论 external fragmentation。

假设一个 segment 里有这样的 block 布局:

1┌────┬────┬────┬────┬────┬────┬────┐
2│ A  │ B  │ C  │ D  │ E  │ F  │ G  │
3└────┴────┴────┴────┴────┴────┴────┘

其中各个 block 的 lifetime 完全不同。一段时间后,B、D、F 释放了:

1┌────┬──────┬────┬──────┬────┬──────┬────┐
2│ A  │ FREE │ C  │ FREE │ E  │ FREE │ G  │
3└────┴──────┴────┴──────┴────┴──────┴────┘

虽然总 free memory 不少(free1 + free2 + free3),但如果需要一个更大的 contiguous block——大于任何一个单独的 free hole——仍然无法直接分配。

这就是 external fragmentation。

需要特别注意:fragmentation 不等于 memory leak。 Leaked memory 是应用逻辑上不再需要但没有被释放的 memory。Fragmentation 是 free memory 确实存在,但布局不够 contiguous。两者完全不同。

Allocation Lifetime:全文最关键的 Big Idea
#

到这里大多数讨论会集中在 block size 上——best fit、size rounding、split threshold。但我们的线上 workload 暴露出了一个更重要的问题:

不同生命周期的 allocation 被混合放进了同一个 memory arena。

在在线推理中,GPU memory 中的对象有非常不同的 lifetime。

Long-lived 对象:

1model weights
2persistent buffers

可能存活数分钟、数小时,甚至整个服务时段。

Short-lived 对象:

1activation
2temporary tensor
3operator workspace

可能只存活几十微秒、几毫秒、一次 inference 的时间。

如果把这些东西全部交给同一个 allocator,segment 内部的 block 分布可能变成:

1┌────────┬────────────┬────────┬──────┬────────┬───────────┬────────┐
2│ Weight │ Activation │ Weight │ Temp │ Weight │ Workspace │ Weight │
3└────────┴────────────┴────────┴──────┴────────┴───────────┴────────┘

随着 short-lived allocation 不断创建和销毁:

1┌────────┬────────────┬────────┬──────┬────────┬───────────┬────────┐
2│ Weight │    FREE    │ Weight │ FREE │ Weight │   FREE    │ Weight │
3└────────┴────────────┴────────┴──────┴────────┴───────────┴────────┘

Long-lived weight 就像一根根"柱子",把整个 address space 分割成很多 hole。

Free memory 看起来不少,但每个 hole 都不大。下一次需要分配一个较大的 contiguous block 时,这些 hole 用不上。

PyTorch 2026 年 6 月的 CUDACachingAllocator DevLog 中有一句话恰好描述了这个问题:

But if your use case has long-lived allocations interleaved with short-lived ones in the same pool, expandable segments won’t save you from fragmentation within the segment.

也就是说,即使使用了 expandable segments——一种通过 virtual memory mapping 让同一 segment 内的 block 总是可以 merge 的机制——如果 long-lived 和 short-lived allocation 交错存在,fragmentation 仍然无法避免。

因为 merge 的前提是相邻 block 都是 free。只要中间有一个 long-lived block 还活着,两边的 free block 就无法合并。

gantt
    title Allocation Lifetime 对比
    dateFormat X
    axisFormat %s

    section Long-lived
    Weight_A :wa, 0, 20
    Weight_B :wb, 0, 20
    Weight_C :wc, 0, 20

    section Short-lived
    Activation_1 :a1, 1, 3
    Temp_1 :t1, 2, 4
    Activation_2 :a2, 5, 7
    Workspace_1 :ws1, 6, 8
    Activation_3 :a3, 10, 12
    Temp_2 :t2, 11, 13

Weight 从头活到尾。Activation 和 temporary tensor 来了又走。如果它们混合分配在同一段 memory 里,short-lived 对象释放后留下的 hole 永远被 long-lived 对象隔开。

生命周期相近的对象应该放在一起
#

这是经典 allocator / arena design 中一个非常重要的思想:

Objects with similar lifetimes should be allocated together.

生命周期相近的对象应该尽量放在同一个 allocator domain / arena / pool 中。

目的不是单纯让 allocation 更快,而是:

1相近时间创建 + 相近时间释放
23memory region 更容易整体回收
45减少长期 fragmentation

当一次 inference 结束时,如果所有 short-lived 对象都在同一个 arena 里,大量 block 可以一起变成 free,甚至整个 arena 可以被整体回收。而 weight arena 的布局保持稳定,不会被 short-lived allocation 的创建和销毁搅乱。

flowchart LR
    subgraph WA["Weight Arena(long-lived)"]
        W1["W1"] --- W2["W2"] --- W3["W3"] --- W4["W4"]
    end
    subgraph RA["Runtime Arena(short-lived)"]
        A1["Act"] --- T1["Tmp"] --- A2["Act"] --- WS["Wsp"]
    end
    WA --> Stable["布局稳定"]
    RA --> Reclaim["inference 结束后
大量 block 一起释放"]

很多 allocator 优化会首先关注 size class、best fit、rounding、split、merge。这些都是空间维度上的优化。

但我们的线上 workload 暴露出了另一个更重要的维度:时间维度,也就是 allocation lifetime。

两个 tensor 即使 size 完全相同——都是 1MB——但如果一个的 lifetime 是 3 小时,另一个的 lifetime 是 200 微秒,从 allocator 角度,它们并不是同一种对象。

这是一个非常重要的 insight:

Memory allocation 不只有 size 这个维度,还有 lifetime。

甚至可以说:

Fragmentation 很大程度上来源于不同生命周期对象的混合,而不仅仅是 block size 不匹配。

Fragmentation 本质上不仅是空间布局问题,也是时间行为问题。一个 allocator 如果只按 size 做 best-fit,但完全不考虑 allocation 会存活多久,就没有利用到 workload 中非常强的时间结构信息。

回到线上模型:Lifetime Separation
#

把前面的原理映射回我们的 workload。

在线推理中,GPU memory 中的对象恰好具有非常清晰的 lifetime 分类:

1Model Weight = long-lived
2Activation / Workspace / Temporary = short-lived

所以第一步不是发明一个更复杂的 block search algorithm,而是:首先把完全不同生命周期的 allocation 分开。

 1┌───────────────────────────┐
 2│ Weight Memory Pool        │
 3│                           │
 4│ long-lived                │
 5│ stable layout             │
 6└───────────────────────────┘
 7
 8┌───────────────────────────┐
 9│ Runtime Memory Pool       │
10│                           │
11│ activation                │
12│ workspace                 │
13│ temporary tensor          │
14└───────────────────────────┘

这个 separation 本身就能减少 short-lived allocation 对 long-lived weight layout 的"污染"。

这是全文从"理解 PyTorch allocator"走向"设计 workload-specific memory management"的转折点。

Weight Pool 内利用 Size Repetition
#

完成 lifetime separation 后,可以进一步观察 weight workload 的特点。

模型 architecture 决定了参数的 shape。如果大量模型属于相同或者相似结构:

1Model A:  weight1 = [256, 128]   weight2 = [128, 64]   weight3 = [64, 32]
2Model B:  weight1 = [256, 128]   weight2 = [128, 64]   weight3 = [64, 32]
3Model C:  weight1 = [256, 128]   weight2 = [128, 64]   weight3 = [64, 32]

那么 tensor allocation 的 size 并不是随机的,而是高度重复的。

flowchart TD
    subgraph Pool["Weight Pool"]
        S1["256KB cached blocks"]
        S2["512KB cached blocks"]
        S3["1MB cached blocks"]
    end
    MA["Model A weight"] -->|"256KB"| S1
    MB["Model B weight"] -->|"256KB"| S1
    MC["Model C weight"] -->|"1MB"| S3
    MD["Model D weight"] -->|"1MB"| S3

在大量同构模型中,相同的 allocation size 不断重复出现。可以针对性地建立 size 到 reusable blocks 的映射,进一步提高复用率。

这里要强调:这是建立在 lifetime separation 之后的第二层优化。 文章逻辑是先按 lifetime 分开,然后在 Weight Pool 内利用 size repetition——而不是反过来。

不是重新发明 Caching
#

技术表述一定要严谨。

PyTorch CachingAllocator 本身当然会 reuse cached block,也会做 size rounding、best-fit matching、split、merge。所以我们的优化不是"PyTorch 不 cache,所以我们 cache",也不是"PyTorch 不会复用相同 size,所以我们会"。

真正的区别是:PyTorch 面向 generic tensor allocation。而我们知道:

  • 这个 allocation 是 weight
  • 它通常是 long-lived
  • 这个 size 会频繁再次出现
  • 它属于某一类 model architecture

我们引入的是 workload semantics,而不是重新发明 caching。

Batch Suballocation 与 Arena
#

对于大量 weight tensor,还有一种进一步的优化方式:不要让所有 tensor 都成为完全独立的底层 allocation。

假设一个模型有 w1、w2、w3 …… wn 个 weight tensor,可以先规划 total size,然后一次获得一个 large segment:

1┌──────────────────────────────────────┐
2│ w1 │ w2 │ w3 │ w4 │ ... │ wn        │
3└──────────────────────────────────────┘

每个 tensor 使用 base_ptr + offset 形成 suballocation。

这就是 arena allocator 的思想。当对象属于同一个模型,并且生命周期相近时,把它们放在同一个 arena 里,本身就比把它们当成完全无关的 allocation 更符合 workload 的实际行为。

这里再次联系到 lifetime locality:同一个模型的 weight tensor 几乎总是同时创建、同时存在、同时销毁。

从 Tensor Reuse 到 Model Layout Reuse
#

完成 Weight Pool 后,还有一个更强的业务先验。

很多时候模型切换是这样的:

1Model A 下线
2Model B 上线
3
4architecture(A) == architecture(B)

也就是说:tensor shapes 完全相同,tensor sizes 完全相同,memory layout 可以完全相同。

如果只站在 allocator 角度,这是一个"free blocks → 重新 allocate blocks"的过程。但如果站在 model runtime 角度:

为什么还需要重新做一次 allocation planning?

可以进一步做到:Model A 的 memory layout 直接保留,Model B 直接复用。

flowchart TD
    subgraph ModelA["Model A(online)"]
        SA0["segment 0"] --- SA1["segment 1"] --- SA2["segment 2"]
    end
    ModelA -->|"A offline"| Hold["保留 memory layout
不立即拆散 segments"] Hold -->|"B online
architecture compatible"| Reuse["Ownership Transfer"] Reuse --> ModelB["Model B 直接写入原有 region"]

这里的核心不是单纯的 block reuse,而是 layout reuse

Block caching 只知道"这块内存目前空闲";model-level reuse 还知道"这套内存布局恰好适合下一个模型"。

这是从 Memory Allocator 走向 Memory Planner 的重要一步。

Lazy Reclaim 与 TTL
#

模型 offline 后,也不一定马上把 Weight Pool 中的 segment 真正释放。可以这样做:

1model offline
23mark reusable
45设置 TTL

如果很快出现 architecture compatible 的新模型,直接 reuse。否则 TTL 到期后再 reclaim。

这利用的是 temporal locality:最近刚刚使用过的 model memory layout,在不久以后可能再次有价值。类似 connection pool、object cache、CPU cache 的思想。

三个层次
#

不要把整个优化理解成"我们写了自己的 allocator"。更准确地说,这是三个层级的递进。

Level 1:Generic Caching
#

PyTorch CachingAllocator:

1free block → cache → reuse

解决的问题:频繁 cudaMalloc / cudaFree

Level 2:Lifetime-aware Allocation
#

首先识别:Weight、Activation、Workspace、Temporary Tensor 具有完全不同的 lifetime。然后把 long-lived 和 short-lived 分开管理。

解决的问题:mixed-lifetime fragmentation。

按生命周期分配,比单纯按大小分配更能解释我们为什么有优化空间。

这是全文最关键的一层。

Level 3:Semantic Reuse
#

进入 Weight Pool 后进一步利用 same tensor size、same architecture、same model layout:

1Tensor reuse → Segment reuse → Whole-model layout reuse

这已经属于 workload-aware memory planning。

flowchart TD
    L0["cudaMalloc / cudaFree
每次都走 CUDA"] --> L1["Caching Allocator
block cache + reuse"] L1 --> L2["Lifetime-aware Pools
Weight 与 Runtime 分离"] L2 --> L3["Tensor-level Reuse
利用 size repetition"] L3 --> L4["Model-level Layout Reuse
ownership transfer"] L4 --> L5["Memory Planner
compile/runtime co-design"]

空间局部性与时间局部性
#

可以把整篇文章稍微上升一个抽象层次。

传统 allocator 很容易关注 size、address、free block——这是空间问题。

但 ML workload 存在非常强的时间结构:

  • Weight 长时间存在
  • Activation 短时间存在
  • Model 有明确的上线和下线生命周期

因此好的 memory management 应同时考虑 spatial locality 和 temporal locality。

相同生命周期对象放在一起——这是一种 temporal grouping。同一模型的 tensor 放在连续 arena——这又同时改善 spatial organization。

可以形成一个更一般的观察:

Allocator 看到的是 byte;Runtime 看到的是 object;ML Runtime 还能看到 Tensor、Graph 和 Model。越往上层,能够利用的语义越丰富。

通用 allocator 不可能利用"这个 allocation 是 model weight"这样的信息——因为它根本不知道。而在我们的场景中,runtime 知道每一个 allocation 的语义身份和预期 lifetime,这就是优化空间的来源。

PyTorch CachingAllocator 仍然是正确设计
#

文章到这里必须明确一件事:不是 PyTorch allocator 设计得不好。

PyTorch 面向的场景是:

1arbitrary model
2arbitrary tensor shape
3dynamic graph
4training + inference
5multi-stream
6CUDA Graph
7distributed workload
8...

它不能假设下一个 model 一定与当前 model 同构,也不能假设某一个 allocation 永远是 model weight。因此它必须是一个 general-purpose allocator。

而我们的线上 runtime 可以利用的信息包括:

1known model lifecycle
2known model class / architecture
3known weight size distribution
4known online / offline transition

所以:

通用 allocator 和 domain-specific memory planner 并不是竞争关系,而是不同 abstraction layer 的优化。

这是全文需要保持的技术立场。

PyTorch Allocator Tuning
#

PyTorch 也提供了一些 allocator tuning 机制。通过环境变量 PYTORCH_CUDA_ALLOC_CONF(或其统一名称 PYTORCH_ALLOC_CONF)可以配置的参数包括:

  • max_split_size_mb:阻止 allocator split 超过指定大小的 block,可以在某些 workload 下减少 fragmentation
  • expandable_segments:使用 CUDA virtual memory API 让 segment 可以动态增长,同一 segment 内的 block 总是可以 merge
  • roundup_power2_divisions:控制 size rounding 策略
  • garbage_collection_threshold:控制何时触发 GC 回收

这些机制主要是在 block、segment、size、reuse policy 层面改善通用 allocator 的行为。

而 weight lifetime、model identity、architecture equivalence、model replacement 属于更高层的 runtime semantics。

调 allocator 参数和做 workload-specific memory management 并不矛盾。前者优化通用 allocator 的行为,后者引入 workload 层面的语义信息。

联系 CUDA Graph 和 Compiler Memory Planning
#

现代 ML runtime 的很多优化其实都在走同一个方向:

如果 tensor shape 和 lifetime 可以预测,就尽量减少完全动态的 memory allocation。

CUDA Graph 要求执行过程中的 memory address 保持稳定。TensorRT 在 build 阶段会做 memory planning,预先规划每一层的 tensor 使用哪一段 memory。XLA 和 AOTInductor 也有类似的 static memory planning pass。

这些技术细节不在本文范围内,但它们说明了一件事:我们在线上 inference 中做的 lifetime-aware memory management 不是孤立的技巧。它属于更大的趋势:

1dynamic allocator
23lifetime-aware allocator
45memory planner
67compile / runtime co-design

越能够在 allocation 发生之前就预测 tensor 的 shape、lifetime 和使用模式,就越有机会用 planning 替代 dynamic allocation。

结论
#

回顾全文的认知链:

第一层。 PyTorch CachingAllocator 告诉我们:tensor free 不等于 cudaFree。通过 block cache,可以避免大量 expensive CUDA allocation。

第二层。 Fragmentation 告诉我们:memory allocation 不只是 size 问题,同时也是 lifetime 问题。

第三层。 我们的线上 workload 告诉我们:weight、activation、workspace、temporary tensor 具有完全不同的 lifetime。首先按生命周期隔离 allocation,比单纯调整 block size 更重要。

第四层。 一旦把 weight 单独管理,就可以继续利用 tensor size repetition、architecture repetition、model lifecycle,将 block reuse 升级成 tensor reuse、layout reuse、model reuse。

最终可以用这样一句话概括:

通用 Caching Allocator 回答的是"这块显存还能不能复用",而 workload-aware memory manager 还可以回答"这块显存应该和哪些对象一起管理,以及下一次最可能被谁复用"。

全文最希望读者记住的两个 insight:

Memory allocation 不只有 size 这个维度,还有 lifetime。

Fragmentation 很大程度上来源于不同生命周期对象的混合,而不仅仅是 block size 不匹配。

对于 ML workload,越能够理解 tensor 和 model 的生命周期,就越有机会从通用动态 allocation 走向可预测的 memory planning。

参考资料
#

Related