返回 文章 apply CMS 文章

NVIDIA ModelExpress:以光速分发模型权重,大幅缩短启动时间

ModelExpress 让模型权重在集群中像光一样快速传输,消除冷启动瓶颈。

模型分发NVIDIAP2P RDMA推理优化
成长分 / 100 79 综合收获、行动、留存与影响

NVIDIA ModelExpress:以光速分发模型权重,大幅缩短启动时间
为什么值得读了解如何将大模型启动时间从分钟级压缩到秒级,提升服务弹性。

掌握 P2P RDMA、ModelStreamer、GDS 等多种权重加载路径的适用场景与自动回退机制。

关键洞察
  1. ModelExpress 优先从已服务副本通过 P2P RDMA 直接传输权重,避免重复从存储加载。
  2. 通过池注册和 VMM 竞技场注册,NIXL 内存注册开销可降低 80% 到 99%。
  3. 内核缓存(如 Triton、DeepGEMM)可跨副本传输,进一步减少启动时间。
转成行动

深入阅读

正文与原文对照

原文保真覆盖:全文原文字符:16000

每移动一个字节都有成本。随着模型检查点(checkpoint)增长到数百 GB 乃至 1 TB,这笔成本会迅速累积。更糟的是,在集群中搬运这些模型权重是极其常见的操作:冷启动要把权重从远程存储拉进 GPU 显存;自动扩缩容和滚动更新要为每个新副本填充权重;RL 后训练则要不断把更新后的权重从训练器送到 rollout 工作节点。这些看似是不同的工作流,却都要反复缴纳同一笔“税”:在有效工作开始之前,先花时间搬运权重。

ModelExpress:加速模型权重的整个生命周期

NVIDIA ModelExpress(MX)围绕一个简单的想法构建:加载模型之前,先问一句——兼容的权重副本现在在哪里?MX 不再把每个副本都当作独立的冷启动,而是选出当前最快的来源和传输路径。

当某个在线服务的对等节点(peer)已在 GPU 中持有兼容的权重时,MX 会通过 NVIDIA Inference Xfer Library(NIXL)以 P2P RDMA 方式直接在 GPU 之间传输,绕开对对象存储、本地磁盘和主机内存的冗余访问。若没有可用的对等节点,MX 会从受支持的最快路径引导加载:要么从对象存储以流式方式读取而不落盘,要么把本地文件直接读入 GPU 显存。

MX 能在不到 10 秒内,把 DeepSeek-V4 Pro 的权重和 JIT Kernel 缓存产物从一个在线副本传输到全新副本,使总启动时间从 8 分钟缩短到 1 分 44 秒。本文余下部分将展示 MX 如何选出通往 GPU 显存的最快路径:优先使用来自在线副本的 P2P RDMA,并一路消除多余的下载和拷贝;随后再把同样的思路扩展到内核缓存复用和 RL 权重更新分发。

ModelExpress 概览。控制平面通过 Redis 或 Kubernetes 元数据发现兼容的来源。数据平面沿着探测得出的优先级链传输权重:通过 NIXL 以 GPUDirect RDMA 直接从在线对等节点搬运权重,通过 ModelStreamer 从对象存储以流式方式读取,或通过 GDS 从本地存储读取。
图 1. ModelExpress 概览

加速从远程存储到 GPU 显存的每一个环节

每个新的工作节点都必须从三个地方之一获取权重:远程存储(如 HF 或 S3)、本地存储,或已经在服务该模型的另一个工作节点。对第一个工作节点而言,还没有对等节点可用,只能从存储引导。MX 可以从对象存储以流式方式读取检查点,或从高速本地存储加载,并在这两条路径上都去掉可避免的拷贝。

一旦第一个工作节点开始提供服务,首选来源就变了。它的权重已经驻留在 GPU 显存中,完成了后处理并按推理所需排布好,因此之后的每个兼容工作节点都应通过 P2P RDMA 直接从它加载。MX 会自动完成这一转换:先从存储引导一次,然后在 GPU 之间横向扩展,只有在没有兼容对等节点时才回退到存储。

启动第一个工作节点:从存储引导

远程对象存储到 GPU:避开本地磁盘

当检查点存放在云存储桶中,而你又不想配置和管理磁盘缓存层时,MX 会使用 Model Streamer,经由可复用的 CPU 暂存缓冲区把 safetensors 拉入 GPU。检查点不会落到本地磁盘,从而省去了中间的下载、重新加载和存储卷。

Model Streamer 使用多线程张量读取器,在各个检查点分片之间并发获取张量区间。张量到达时,它把远程读取与 GPU 放置流水线化:已完成的张量交给推理引擎,同时后面的张量仍在获取中。这样存储、网络和 GPU 拷贝路径始终保持繁忙,同时只复用有限的主机内存。

在张量并行部署中,参与的各个 rank 会分担远程读取并共享结果(通常通过 NCCL),而不是让每个 rank 各自下载完整的检查点。MX 把这种分布式流直接接入推理引擎的权重加载器,让第一个工作节点做好准备,成为之后所有兼容副本的 P2P 来源。

集群入口:下载一次,而不是 N 次

当集群维护共享的磁盘缓存层(例如 K8s 中的持久卷)时,MX 确保整个集群只填充它一次。如果 10 个副本同时要获取 806 GiB 的 DeepSeek-V4 Pro 模型,它们将在争抢同一入口带宽的同时,通过网络拉取约 8 TiB 的相同数据。MX Model Cache Service 把这些请求合并成一次协调的下载:由 Metadata Store 中的原子认领选出下载者,其余副本则跟踪其进度并复用缓存的副本。集群只需支付一次外部下载的成本,之后每个副本都可以从同一份缓存的检查点开始。

本地存储到 GPU:绕过主机内存暂存

当系统支持 GPUDirect Storage(GDS)时,MX 通过 NIXL 的多线程 GDS 后端,把检查点文件直接从本地存储读入 GPU 显存。NIXL 并行执行批量张量读取,直接写入 GPU 显存,绕过主机内存以及传统加载器所需的暂存拷贝。用户无需显式启用 GDS:MX 会自动检测该能力,并在其不可用时回退到其他加载策略。

本地存储到 GPU:用 ModelStreamer 流水线化本地读取

MX 也可以通过 ModelStreamer 加载本地检查点。多个操作系统线程并发地把 safetensors 读入可配置的 CPU 缓冲区,同时已完成的张量被移入 GPU,后续读取继续并行进行。与 GDS 不同,这条路径仍需经过主机内存暂存,但它让磁盘 I/O 与 GPU 放置相互重叠,可受益于操作系统页缓存,并在无法直接从存储访问 GPU 时提供一条可移植的快速路径。

启动第一个之后的每个工作节点:从在线对等节点获取

这是 MX 的关键特性。一旦另一个副本已在服务同一个模型,权重就已经走完了大部分旅程:它们驻留在 GPU 显存中,完成了后处理,并按推理引擎的需要排布好。MX 把该副本视为一个实时的权重来源。确认兼容后,它将张量直接从源 GPU 传输到目标 GPU。新副本加载完权重后便加入来源池,为后续副本提供又一个可加载的对等节点。每成功完成一次传输,这个池子都会随部署规模一同增长,把扩容变成 GPU 到 GPU 的扇出,而不是反复的冷加载。

MX 控制平面负责发现兼容的对等节点、交换传输元数据并跟踪来源就绪状态,但从不经手权重字节本身。在数据平面上,MX 默认使用 NIXL 作为传输引擎,其可插拔的后端使其能在 Infiniband、RoCE、NVLink、EFA 等多种网络上发挥峰值性能。MX 提供一等公民级别的传输接口,允许 fabric-lib 和独立的 Mooncake 等库与 MX 集成。

示意图:新的引擎副本通过 NIXL 经 RDMA 直接从来源发现并获取模型权重,随后新副本加入来源池,从而实现扇出式扩容。
图 2. 通过 NIXL 实现的点对点 GPUDirect RDMA 权重传输

在任何传输开始之前,MX 会根据模型以及决定张量布局的运行时设置计算出一个 mx_source_id,然后只考虑 ID 匹配的对等节点。控制平面通过 Redis、Kubernetes CRD 或 k8s-service(无服务器)元数据后端来发现这些对等节点。

优化 NIXL 内存注册开销

NIXL 要对某个张量做 RDMA 之前,必须先注册其背后的 GPU 显存:这是一次 ibv_reg_mr 调用,返回用于远程访问的远程密钥(rkey)。一个大模型有数以万计的张量,逐个注册慢到足以在耗时预算中显现出来。默认情况下,MX 会逐个张量注册。有两种可选策略可降低这部分注册成本:

  • 池注册(Pool registration)对每个底层 cudaMalloc 分配只注册一次,而不是对每个张量注册,在典型模型上可将注册次数减少 80% 到 99%,且传输语义不变。
  • VMM arena 注册则更进一步。它安装一个 CUDAPluggableAllocator,把加载期间的所有分配都路由到单个 16 TiB 的虚拟地址 arena 中,然后在加载结束时把整个已使用范围作为一个基于 dmabuf 的内存区域一次性注册。注册次数从每个张量一次骤降为总共一次;每个张量描述符只需携带指向这一整块区域的偏移量。

我们在 vLLM 引擎上使用 DeepSeek-V4-Pro(TP=8),如下方图 3 所示,测量了每种方法的平均 NIXL 注册时间。

柱状图:比较 DeepSeek-V4-Pro(vLLM 上 TP=8)在三种策略下的 NIXL 内存注册时间,分别是逐张量注册(基线)、池注册和 VMM arena 注册(最快)。
图 3. NIXL 内存注册优化

运行时路径选择与安全回退

启动时,MX 会探测可用的能力,并自动跳过环境不支持的路径。第一个适用的策略按当前的优先级顺序运行:P2P RDMA -> ModelStreamer -> GDS -> default loader (host-staged POSIX I/O)。如果某条路径不可用,或在修改模型状态之前失败,MX 会自动向下一级回退。如果故障发生在权重已开始写入之后,它会先重新初始化模型再继续,因此不会提供只写了一部分的权重。

P2P 重试只会在传输开始前的元数据失败时切换对等节点,原生加载器始终是最后的兜底方案。这种由能力驱动的设计让 MX 核心与硬件和软件无关,平台特定的快速路径只在受支持的地方启用。

端到端结果

我们在配备 NVIDIA ConnectX-7 网卡的 8xB200 GPU 节点上运行 DeepSeek-V4-Pro,比较了不同冷启动场景下的总模型加载时间。每个副本使用 vLLM 0.23.0,TP=8,并启用 --enable-flashinfer-autotune。见下方图 4。

柱状图:比较 8xB200 节点上 DeepSeek-V4-Pro 在四条路径下的总模型加载时间,分别是 P2P RDMA、ModelStreamer、GDS 和默认的经主机暂存的 POSIX I/O。
图 4. 端到端冷启动模型加载时间对比:HF、ModelStreamer(S3)、磁盘与 P2P RDMA

不只是加载完,而是已预热:继承已编译的内核

把权重送进 GPU 显存至关重要,但加载完成的模型还不能立即对外服务。在最初几次前向传播中,引擎会针对确切的模型、dtype、量化方式和 GPU,对内核进行 JIT 编译和自动调优(例如 torch.compile、Triton、DeepGEMM、TileLang 等),并捕获 CUDA graph。对 DeepSeek-V4 Pro 这样的模型,这可能需要几分钟;一旦 MX 降低了权重加载延迟,它就可能成为启动成本的主要部分(见下方图 5)。

堆叠柱状图:在 ModelExpress 消除权重加载延迟之后,JIT 内核编译(torch.compile、Triton、DeepGEMM 等)成为 DeepSeek-V4-Pro 启动成本的主要部分。
图 5. DeepSeek-V4 Pro(TP=8 vLLM)的启动时间构成

这种重复的预热是可以避免的。当模型、软件栈和 GPU 架构都匹配时,只需一个副本承担编译成本,其余副本即可继承生成的缓存。

MX 的 Artifact Transfer API 会打包这些基于文件的产物,通过 NIXL 的 CPU 到 CPU RDMA 路径,在已注册的主机内存缓冲区之间直接传输,然后在目标引擎的缓存目录中进行校验和安装。这样在 Kubernetes 中就无需共享的 ReadWriteMany(RWX)卷,同时针对具体产物的 mx_source_id 可防止在不兼容的副本之间复用。配置了 Redis 或 Kubernetes 元数据后端时,MX 会自动检测标准的缓存位置。

我们使用相同的配置来测量内核产物传输能在多大程度上缩短启动时间。启用产物传输的运行传输了 Triton/DeepGEMM/TileLang/CuTe DSL/FlashInfer 缓存。图表比较了主要的启动阶段,以及从进程启动到 API 就绪的总挂钟时间。

分组柱状图:比较磁盘基线与 ModelExpress P2P RDMA 在启用和不启用产物传输机制时的总启动时间,显示继承 Triton、DeepGEMM、TileLang、CuTe DSL 和 FlashInfer 内核缓存能显著缩短到 API 就绪的挂钟时间。
图 6. ModelExpress 带来的总启动时间缩减

当权重每一步都在变:RL 后训练

前面所说的一切都假设模型权重一旦加载就固定不变。RL 后训练打破了这一假设。训练器每一步都会更新策略,而生成 rollout 的推理 actor 必须在下一轮生成之前拿到这些权重。与推理启动一样,权重搬运在 RL 中也处于关键路径上:rollout 工作节点要等待更新后的权重,从训练器的分布式布局(无论是 FSDP/DTensor 分片,还是 Megatron 的 TP、PP 和 EP 划分)搬运到推理引擎的布局中。

示意图:ModelExpress 的 RL refit 流程。训练器各 rank 向控制平面公布张量归属以供发现来源,rollout 工作节点通过 NIXL 直接从训练器各 rank 拉取更新后的权重字节。
图 7. ModelExpress 让 RL refit 由接收方驱动

MX 通过以下四个阶段驱动 refit:

  1. 发布:每个训练器 rank 向 MX 公布自己已持有的张量或分片,以及描述其形状、dtype、放置方式和参数映射的元数据。
  2. 发现:rollout 工作节点通过 MX 查找所请求的权重版本及其可用来源。
  3. 规划:接收方把已发布的归属信息映射到自己的目标布局上,并确定哪些来源包含所需的张量或区间。
  4. 拉取、转换与加载:接收方直接对这些来源发起单边读取。

MX 包含接收方驱动 refit 的核心构件,客户正在实际集成中对其进行评估。我们还在测试用于跨集群权重传输的增量权重差异(delta weight diff)refit,这是 Fireworks/Cursor、Cognition 等机构在近期 RL 训练中使用的技术。

为 Dynamo 做贡献与我们的路线图

MX 与 vLLM 和 SGLang 有原生集成,并支持包括 Dynamo 和 llm-d 在内的推理服务框架。

Dynamo 开源社区正积极推进更深入的 TensorRT-LLM 集成和更广泛的推理能力。欢迎查阅当前的 Dynamo 文档和路线图,在你自己的环境中试用现有工作流,并通过反馈帮助塑造项目的发展方向。

致谢
ModelExpress 是团队合作的成果。感谢 MX 团队的其他成员 Zhongdongming Dai 和 Tanushriya Singh 在项目上的核心工作。感谢 Itay Neeman、Anish Maddipoti、Istvan Haller 和 Omri Kahalon 对项目技术方向的指导,以及 Red Hat 的 Will Eaton 对 llm-d 集成的支持。