借助NVIDIA Snapshot,几秒钟内启动并恢复GPU工作负载

NVIDIA Snapshot是NVIDIA DSX AI Factory平台的新开源组件,可使已初始化的GPU工作负载在不同推理副本之间复用,从而加快启动与恢复速度。

当推理需求激增或某个副本发生故障时,可用的GPU并不会立即转化为服务能力。初始化一个大语言模型工作负载——这一过程被称为冷启动——可能需要数分钟,从而延迟扩容与恢复进程。

Snapshot是一套面向Kubernetes上GPU工作负载的检查点与恢复系统。它能够捕获一个已完全初始化的工作负载(包括CPU和GPU内存),并在兼容的基础设施上恢复该状态,从而将推理能力的上线时间从数分钟缩短至数秒。图1展示了启动流程。Snapshot跳过了重复的工作负载初始化过程;基础设施的配置和容器镜像的加载仍不在此次恢复路径的测量范围内。

本文将介绍Snapshot如何对GPU工作负载进行检查点与恢复操作、平台团队如何进行集成,以及在三种推理引擎中恢复时间与冷启动时间的对比情况。

面向AI工厂运营的Snapshot

Snapshot最初是为NVIDIA Dynamo推理工作节点开发的,它提供的检查点与恢复能力,也可供其他推理引擎、调度器、控制器及平台进行集成。

Snapshot为GPU Pod提供两种操作:

检查点:当工作负载完全初始化后,将其CPU和GPU内存状态捕获为一个可移植的工件。

恢复:在具有相同GPU型号和驱动的任意节点上,根据该工件启动一个新的Pod。该流程将从检查点所记录的状态继续执行,因此权重加载、编译及图捕获等步骤均可跳过。

这两项操作均使用Kubernetes自定义资源来实现。若要对一个已在运行的Pod进行检查点操作,可创建一个命名该Pod的PodSnapshot资源;若要从模板启动一个Pod并在其准备就绪后立即执行检查点操作,则可创建一个SnapshotJob资源。若要进行恢复,需在新Pod中添加注解nvidia.com/restore-from,并指定要恢复的PodSnapshot名称。工作负载只需进行简单的握手操作,即可在检查点前暂停,并在恢复后继续运行。该项目仓库中提供了针对vLLM、SGLang和NVIDIA TensorRT LLM的现成示例。

Snapshot能够将单个节点上跨多个GPU的兼容工作负载,作为跨计算单元(rank)的单一检查点工件进行捕获和恢复。这为采用张量并行(TP)、数据与专家并行(DEP)、解码上下文并行(DCP)等并行策略的大型混合专家(MoE)模型实现快速启动提供了支持。需要注意的是,工作负载必须满足Snapshot在检查点/恢复及运行时方面的兼容性要求,具体可参见相关示例配置。

基准测试

多家设计合作伙伴也已在自身的工作负载中对Snapshot进行了评估:

Crusoe作为Snapshot的上游贡献者,在NVIDIA HGX B300上将GLM-5.3 Pod的就绪时间从约8分钟缩短至31秒。

Saturn Cloud作为一家Token工厂平台,针对从零扩容的推理场景对NVIDIA Snapshot进行了基准测试,在NVIDIA HGX B300上将vLLM的启动时间从2至8.5分钟缩短至6至34秒。

以下基准测试展示了Snapshot在vLLM、SGLang和NVIDIA TensorRT LLM(TRT-LLM)这三种引擎上对推理启动的加速效果——测试模型涵盖从6亿参数模型到在8块GPU上运行、检查点大小达1 TB的混合专家(MoE)模型。在所有测试配置中,Snapshot均能在45秒内恢复一个副本,速度比冷启动快8至36倍。

测试方法

每种配置均在相同的硬件与软件环境下,采用两种策略进行测试:

冷启动:工作负载从零开始启动。计时从进程启动开始,到引擎初始化完成(即推理开始前)结束。该测量过程被划分为四个阶段(见表2):启动阶段,即启动引擎与工作进程、导入Python库并设置GPU间通信,此阶段也会进行部分即时(JIT)内核编译;模型加载阶段,即读取模型权重并构建模型张量,包括加载过程中进行的特定引擎转换或重打包操作;初始化阶段,即为模型提供服务做准备,包括内存分析、分配键值(KV)缓存与运行时缓冲区,以及其他内核和运行时准备工作;编译、自动调优与图捕获阶段,即编译内核(例如使用torch.compile)、针对GPU进行自动调优,并捕获CUDA图。

Snapshot恢复:工作负载从检查点恢复。计时从发起Snapshot恢复请求开始,到工作负载重新开始提供服务为止。该测量过程被划分为四个阶段(见表3):代理设置阶段,即节点上的Snapshot代理定位检查点工件,并为目标Pod和GPU做好准备;用户空间检查点/恢复(CRIU)恢复阶段,即CRIU(Linux标准进程检查点工具)将进程镜像从存储中读取至主机内存,此阶段速度主要取决于存储后端性能;CUDA恢复阶段,即将GPU状态重新复制回GPU内存;唤醒/重映射阶段,即使进程适配新节点并唤醒引擎。

两种测量方式均不包含容器镜像拉取及容器运行时设置的时间。KV缓存释放操作在每次检查点之前均会被触发,因此不计入工件内容。

表1还报告了检查点操作的耗时与大小:即捕获一个完全初始化的副本所需的时间,以及存储该工件所需的存储空间。这是针对每个模型与配置的一次性成本。

测试环境配置

GPU:NVIDIA B200

NVIDIA驱动版本:615.71.09

存储后端:VAST(NFS)

推理引擎:vLLM 0.31.0、SGLang 0.5.21、TensorRT LLM 1.3.0

测试结果

基准测试结果显示出四个主要发现。

恢复时间均保持在43秒以内:在所有测试配置中,Snapshot的副本恢复耗时为3.6至42.7秒,比冷启动快8至36倍。即便是检查点大小约为600 GB至1 TB、分布在八块GPU上的最大规模MoE模型,其恢复时间也低于43秒,而对应的冷启动时间则为6至19分钟。

恢复过程跳过的不仅是权重加载:冷启动时间包括权重加载、内核编译与自动调优,以及CUDA图捕获等环节。而Snapshot恢复的是已初始化的工作负载,可跳过这些重复步骤。

检查点操作在首次恢复副本时即可实现收益:在所有测试配置中,检查点耗时与恢复耗时之和均短于单次冷启动所需的时间。检查点创建耗时根据模型与配置的不同,为15秒至5.5分钟。

大规模恢复过程的耗时主要来自GPU状态迁移:在八GPU的MoE配置中,恢复GPU状态耗时18至29秒,而使用CRIU恢复CPU进程状态仅耗时4.2至10.8秒。

图2至图4对比了vLLM、SGLang和TensorRT LLM在所测试模型及并行配置下的冷启动时间与Snapshot恢复时间。表1提供了所有配置的详细数据,包括检查点耗时与大小,所有时间单位均为秒。在图表与数据表中,TEP8表示跨八块GPU进行的张量与专家并行。

表2对冷启动时间进行了细分,展示了恢复操作可以跳过的具体初始化阶段。

表3对恢复时间进行了细分,划分为代理设置、CRIU恢复、CUDA恢复以及唤醒/重映射四个阶段。

图5展示了vLLM场景下恢复时间的详细构成。

相关基准测试工具及复现说明均可在项目仓库中获取。

快速上手

若要在Kubernetes上针对GPU工作负载试用Snapshot,可参阅NVIDIA Snapshot项目仓库中的安装说明与使用指南。

可参照用户指南,结合NVIDIA Dynamo试用Snapshot,体验工作节点的恢复功能。

如需探讨集成方案或提供反馈意见,可在项目仓库中提交issue,或加入NVIDIA Dynamo Slack工作空间,访问#snapshot-dev频道。

欢迎前往KubeCon大会上的NVIDIA展位(170号展位),参与Kubernetes技术深度解析及现场演示活动。

NVIDIA