Modal 发布 GPU 内存快照，冷启动最高提速 10 倍

在 Modal，我们对冷启动延迟极为关注。今年早些时候，我们推出了内存快照，将启动时间缩短了一半以上。今天，我们非常激动地宣布下一次进化：GPU 内存快照——将同样的检查点/恢复魔法带给 GPU 加速工作负载。

消除冷启动瓶颈

自成立以来，我们一直从三个角度攻克冷启动问题：

1. 为冷启动优化的自定义文件系统
2. CPU 内存快照
3. GPU 内存快照

我们的分布式文件系统使用一系列缓存，将 Modal 用户中最常用的文件直接存储在 worker 内存中。这非常棒，因为例如，如果某个程序导入了 torch，另一个程序也会受益，因为 torch 文件现在已在 worker 缓存中。这对性能有显著影响，通常比没有缓存时下载文件快 3-5 倍。

Modal Function 的生命周期涉及几个阶段：容器冷启动和运行输入。冷启动通常意味着两件事：下载程序文件和将程序读入内存。

将程序读入内存并启动 Function 需要时间——有时需要很多时间！如果我们能获取程序的内存表示并将其保存为镜像会怎样？这样就能跳过读取文件并在每次冷启动时在内存中重新创建程序，从而节省时间。

事实证明，从镜像重新创建程序确实更快，因此，我们在 2025 年 1 月推出了内存快照。我们在 Function 调用输入之前为其创建内存快照。然后 Function 被“冻结”，保存为优化格式，并缓存在我们的分布式文件系统中。每次程序冷启动时，都从这个冻结状态开始。

在我们的上一篇博客文章 Memory Snapshots: Checkpoint/Restore for Sub-second Startup 中阅读更多细节。

GPU 内存挑战

虽然内存快照显著改善了许多函数的冷启动时间，但它们对 GPU 工作负载有一个重要限制——GPU 状态无法包含在快照中。直到现在。

正如我们在上一篇文章中所解释的，NVIDIA GPU 状态必须在恢复后创建，要求你在恢复完成后填充 GPU 内存、文件描述符和 CUDA 会话。例如，这意味着你必须在恢复后将模型权重从 CPU 复制到 GPU：

这种两步方法有效——你可以将容器启动速度提高多达 3 倍——但并不理想。你需要采用多阶段方法：先将数据复制到 CPU，然后移到 GPU。否则，快照会失效。你还必须确保没有程序尝试创建 CUDA 会话或检查 GPU 可用性（例如，通过调用 torch.cuda.is_available()），因为这也会破坏快照。

这种方法对于预热 GPU 的程序尤其无效，例如使用 torch.compile 的程序。在这种情况下，你仍然需要在将模型加载到 GPU 后编译模型，因为优化后的代码依赖于硬件。由于 torch.compile 是一项重要的优化技术，我们需要更好的解决方案。

GPU 内存快照通过在操作执行后复制 GPU 内存来解决这些限制。使用这种方法，torch.compile 不需要再次运行，因为我们恢复的是已经编译好的模型。这同样适用于已加载的 CUDA 内核、捕获的 CUDA 图以及其他昂贵的冷启动操作。

进入 CUDA checkpoint API：一个用于管理 GPU 内存的接口

在 CUDA checkpoint/restore API 发布之后，该 API 现已在 570 和 575 分支的驱动上可用，我们现在能够为许多工作负载透明地对 GPU 内存进行 checkpoint 和 restore。该 API 允许我们对 CUDA 状态进行 checkpoint 和 restore，包括：

• 设备内存内容（GPU vRAM），例如模型权重
• CUDA 内核
• CUDA 对象，例如流和上下文
• 内存映射及其地址

与我们的 CPU 内存快照类似，GPU 内存快照会在容器即将接受请求之前保存其整个状态。但现在我们还会通过以下步骤捕获 GPU 状态：

• 锁定 CUDA 进程（cuCheckpointProcessLock()）：所有新的 CUDA 调用都会被锁定且永不返回，并会等待所有正在运行的调用（包括 CUDA Streams）完成
• Checkpoint（cuCheckpointProcessCheckpoint()）：我们将 GPU 内存和 CUDA 状态复制到主机内存，释放 GPU 资源，并终止 CUDA 会话

为了实现可靠的内存快照，我们必须首先枚举所有活动的 CUDA 会话及其关联的 PID，然后锁定每个会话以防止在 checkpoint 期间状态发生变化。我们通过 CUDA API 持续监控进程状态（状态定义见 CUDA API 文档），以检测错误、识别 CUDA API 死锁，并为失败的 checkpoint 尝试实现重试逻辑。只有在满足两个条件后，系统才会继续进行完整的程序内存快照：所有进程都已达到 CU_PROCESS_STATE_CHECKPOINTED 状态，且没有活动的 CUDA 会话残留，从而确保整个操作期间的内存一致性。

此时我们创建一份 CPU 内存的记忆——这一次同时包含 CPU 和 GPU 内存！在 restore 期间，我们使用 cuCheckpointProcessRestore() 和 cuCheckpointProcessUnlock() 以相反的顺序执行该过程。

此过程与我们现有的 gVisor checkpoint/restore 系统完全集成。与 CPU 快照一样，我们处理不同 worker 主机之间的兼容性问题，确保在一台机器上创建的快照可以安全地在另一台具有兼容 GPU 硬件的机器上 restore。

性能：冷启动速度提升 10 倍

我们在各种工作负载中测试了 GPU 内存快照，整体结果都非常出色。我们观察到 Functions 的启动速度比基线快最多 10 倍。这一点非常重要，因为你现在可以真正利用 serverless 并缩容到零，同时仍保持出色的用户体验。例如，假设你使用 NVIDIA NeMo 框架运行音频转录模型 Parakeet。一个 Function 冷启动大约需要 20 秒（P0）。使用 GPU 内存快照后，同一个 Function 现在最快只需 2 秒（P0）。

一个满载的 ViT 推理 function 之前仅使用 CPU 快照时需要 8.5 秒（P0），torch.compile 现在只需 2.25 秒（P0）。我们完全跳过了 torch.compile 操作，并使用已编译的产物。同样，运行 Qwen2.5-0.5B-Instruct 的 vLLM 之前启动需要 45 秒（P0），现在只需 5 秒（P0）。

使用 GPU 内存快照

将 GPU 内存快照添加到你的 Modal Apps 就像设置一个新标志一样简单。如果你已经在使用内存快照，只需添加 experimental_options={"enable_gpu_snapshot": True}。

最重大的 API 变化是，你不再需要单独的 snap=True 和 snap=False 生命周期方法。你的模型可以直接加载到 GPU，整个 GPU 状态都会被快照保存。

立即体验

GPU 内存快照功能现已在 Modal 上以 alpha 版本提供。我们仍在探索此功能的局限性；亲自试用一下，并告诉我们效果如何。

致谢

非常感谢 NVIDIA / CUDA 为 CUDA API 添加了检查点 / 恢复功能，也感谢 cuda-checkpoint 项目的维护者，这使得在不同工作负载中测试新 API 变得十分容易。同时，感谢 Google gVisor 团队出色的人们，他们持续开发运行时，使其具备可扩展性、高性能和安全性。