Let Tensors Fly — Accelerating Large Model Weight Loading with R-Fork

发表时间: 2025-12 · Blog post by Ant Group DeepXPU Team, SGLang Team (lmsys.org)

原文: https://lmsys.org/blog/2025-12-10-rfork

作者/机构：Ant Group DeepXPU Team, SGLang Team
发布日期：2025年12月10日

速读

一句话结论
本文提出了一种名为 Tensor R-Fork 的大模型权重加载方法，通过节点间设备到设备的高效互联技术，实现了从运行中的 SGLang 实例向新实例零拷贝加载张量，将 DeepSeek-R1 模型的加载时间从数分钟缩短至数秒，并节省了约 600GB 的本地存储空间。

要解决什么问题
随着大语言模型服务规模和权重体积的不断膨胀，SGLang 实例的冷启动时间成为了生产效率的核心卡点，其中最耗时的环节就是权重加载。原有的权重加载做法受限于数据流转路径中的硬件带宽瓶颈：如果从远程存储中心加载，数据流需要经历“远程存储至远程网卡，经以太网至本地网卡，再到本地内存，最后至本地显存”的漫长路径，其卡点在于 NVMe 存储或以太网卡的带宽；如果从本地磁盘加载，数据流为“磁盘至内存再至显存”，卡点在于 NVMe 带宽；即使从本地内存加载，数据流为“内存至显存”，也受限于 CPU 与 GPU 之间的 PCIe 带宽。这种受限的带宽导致像 DeepSeek-R1 这样的大模型从本地磁盘加载需要数分钟，从远程存储加载甚至需要数十分钟。随着模型参数量呈指数级增长，这种基于传统数据流的初始化和数据传输耗时将进一步恶化，严重拖慢了集群的弹性扩缩容效率。

怎么做的
核心思路是利用 GPU-Direct RDMA 技术构建一个点对点（P2P）的权重存储架构，让所有运行中的 SGLang 实例都能作为新实例的权重数据源。这种设计直接利用了节点间设备到设备互联所提供的高达每秒数百 GB 的吞吐量，通过让本地 GPU 与配对的远程 GPU 直接传输数据，彻底绕开了传统路径中 CPU 和主机内存之间的 PCIe 带宽卡点。具体实现上，该框架包含两种后端方案。第一种是基于 NCCL 的后端：目标实例初始化时向源实例发送 HTTP 请求建立通信组，源实例和目标实例的张量并行工作进程（TPWorker）按序号一一配对（例如源端的 rank 0 对应目标端的 rank 0），随后源实例通过 NCCL 广播将显存中的权重直接发送到目标实例的显存中，全程无中间内存拷贝。第二种是为解决源端干扰问题而引入的 TransferEngine 后端：它是一个轻量级的传输运行时库，与源实例的每个 TPWorker 并行运行。源实例初始化时，TransferEngine 会将权重的显存地址注册到 RDMA 通道；目标实例启动时通过 HTTP 获取包含 RDMA 密钥的元数据，随后直接通过 RDMA 读取源实例的显存数据。TransferEngine 能够智能识别 RDMA、NVLink、GDS 或 TCP 等最优后端，并负责维护内存注册的生命周期以及协调并发读取的流量控制。为了在工业级集群中落地，作者还设计了 Tensor R-Fork Planner 集群调度器。它负责收集并追踪各个实例的运行模型、并行策略（如张量并行或流水线并行）以及服务健康状态。新实例启动时会先向调度器查询，匹配到模型和并行策略一致且健康的种子实例后，再执行点对点权重拉取，若无匹配则回退到默认加载方式。

效果如何
实验在配备 8 张 NVIDIA H20 GPU 的硬件环境下，针对 DeepSeek-R1 模型进行了新 SGLang 实例启动的性能评估。对比的基线方法为工业界常用的传统加载路线，即代表本地 NVMe 带宽瓶颈的“从本地磁盘读取”路线，以及代表网络与存储双重瓶颈的“从远程存储中心读取”路线。量化结果显示，使用 Tensor R-Fork 方法后，DeepSeek-R1 模型的权重加载时间从原先的数分钟（本地磁盘）或数十分钟（远程存储）大幅缩减至仅仅数秒钟。同时，由于实现了从运行实例直接拉取，该方法在本地磁盘或内存中减少了约 600GB 的冗余模型权重存储开销。在代价与局限性方面，两种后端各有取舍。使用 NCCL 后端无需额外依赖且建立通信组仅需数百毫秒，但源实例必须主动参与通信组创建，且 NCCL 广播会触发 CUDA 内核执行，这会与正常的推理生成任务竞争 GPU 资源，导致源实例出现延迟毛刺，破坏了推理服务的稳定性。相比之下，TransferEngine 后端实现了纯网卡驱动的传输，全程不触发源端 CUDA 内核，保证了对源实例推理服务的零干扰，但其代价是需要引入额外的 mooncake 依赖库，且受限于 RDMA 驱动，其内存区域注册过程较慢，可能耗时数秒（尽管作者指出该耗时可与其他初始化阶段重叠来掩盖）。

核心贡献

随着大语言模型（LLM）服务规模和模型权重尺寸的不断扩大，SGLang实例的冷启动时间已成为生产效率的关键瓶颈。在冷启动的各个阶段中，权重加载（weight loading）仍然是最耗时的任务。例如，从本地磁盘加载Deepseek-R1模型通常需要几分钟，而从远程存储系统加载可能长达数十分钟。

为了解决这一问题，本文提出了一种名为 Tensor R-Fork（代表 Tensor Remote Fork）的新型权重加载方法。该方法利用高效的节点间设备到设备（device-to-device）互连技术，以零拷贝（zero-copy）的方式将张量从正在运行的SGLang实例加载到新实例中。

该方法提供了三个核心优势：

  1. 显著加速了权重加载性能；
  2. 消除了本地磁盘和/或DRAM上冗余的模型权重存储；
  3. 确保对正在运行的推理服务无干扰（non-disturbing）。

在Deepseek-R1模型上的应用结果表明，该方法将加载时间从几分钟缩短至仅仅几秒钟，同时将本地磁盘和/或DRAM的存储使用量减少了约600GB，并且在模型传输期间维持了推理服务的质量。

背景知识与设计原则

权重加载的瓶颈分析
随着模型尺寸呈指数级增长，初始化和数据传输所需的时间将进一步恶化。优化权重加载性能最直接的方法是最大化权重数据流中的瓶颈带宽。业界常用的模型加载方法的数据流及其相关的瓶颈带宽如下：

  • 从远程存储中心加载：数据流为 远程存储 -> 远程以太网网卡 -> 以太网 -> 本地以太网网卡 -> 本地DRAM -> 本地GPU内存。其瓶颈在于 NVMe/以太网网卡。
  • 从本地磁盘加载：数据流为 磁盘 -> DRAM -> GPU内存。其瓶颈在于 NVMe。
  • 从本地DRAM加载：数据流为 DRAM -> GPU内存。其瓶颈在于 PCIe。

利用高带宽互连的观察与挑战
节点间设备到设备的互连能够提供每秒数百千兆字节（数百 GB/s）的吞吐量。因此，一个关键问题在于：如何充分利用这种互连带宽来实现SGLang中高效的权重加载？为了应对这一挑战，作者开发了 Tensor R-Fork 权重加载框架，该框架已达到生产就绪（production-ready）状态。

方法细节

Tensor R-Fork 架构设计

利用GPU-Direct RDMA构建P2P权重存储架构
Tensor R-Fork【0，Tensor R-Fork Documentation+2025+Github+https://github.com/sgl-project/sglang/blob/main/docs/advanced_features/rfork.md】的核心概念是利 用GPU-Direct RDMA构建点对点（P2P）权重存储架构。传统方法使用的数据传输性能较低，因为整个路径中始终存在瓶颈，其带宽远小于节点间设备到设备的互连带宽。通过数据流分析可以观察到，权重张量存储在每个GPU上，并且可以通过GPU-direct RDMA在节点之间直接传输。

单GPU对单GPU的数据传输策略
为了最大限度地利用支持RDMA的网卡带宽，设计了一种针对每个GPU对（per GPU-pair）的数据传输策略：本地GPU直接从/向其配对的远程GPU传输数据。

绕过PCIe瓶颈
这种设计有效地绕过了GPU和CPU之间的PCIe瓶颈，实现了高吞吐量通信，而不依赖于CPU或主机内存。从远程SGLang实例加载权重的数据流如图1所示。
从远程SGLang实例加载权重的数据流

双后端实现机制

为了使每个正在运行的实例都能作为需要相同模型的新实例的模型权重源，同时最小化（或理想情况下消除）对正在运行的实例推理服务的干扰，该框架实现了两个后端选项：NCCL 和 TransferEngine。假设有一个正在运行的实例A（称为源实例）和一个要启动的新实例B（目标实例），具体实现机制如下：

NCCL后端实现机制
当使用 NCCL 作为后端【1，Tensor R-Fork with NCCL backend+2025+Github PR+https://github.com/sgl-project/sglang/pull/8215】时，整个过程包括两个阶段 ：

  1. 在源实例和目标实例之间建立通信组。
  2. 通过这些组将权重从源实例传输到目标实例。

NCCL通信组建立过程
在目标实例初始化期间，它向指定的源实例发送HTTP请求以启动通信组的创建。目标实例的每个TPWorker（张量并行工作器）与其对应的源实例的TPWorker建立NCCL通信组（即源rank 0与目标rank 0配对，依此类推）。每个通信组恰好由两个成员组成：源TPWorker和目标TPWorker。

NCCL权重广播与接收
一旦建立通信组，每个源TPWorker就通过NCCL广播机制在组内广播其位于GPU内存上的权重张量。目标TPWorker直接将其接收到其GPU内存中，无需任何中间内存复制。

NCCL后端的局限性
虽然NCCL通过利用GPU-Direct RDMA作为Tensor R-Fork后端，但它确实有一个关键的限制：权重传输会干扰源实例的推理服务，主要由于两个关键因素：

  1. 通信组建立：源实例必须积极参与创建通信组。
  2. CUDA内核干扰：NCCL广播机制触发CUDA内核执行，这会争夺GPU资源并在生成任务期间引入延迟峰值。

TransferEngine后端引入
为了实现无干扰的权重传输，引入了另一种后端：TransferEngine【2，Tensor R-Fork with TransferEngine backend+2025+Github PR+https://github.com/sgl-project/sglang/pull/14997】，它利 用GPU-Direct RDMA进行高效的数据移动。TransferEngine (TE) 是一个轻量级的基于RDMA的传输运行时，它与源实例上的每个TPWorker一起运行，并将驻留在GPU的权重张量暴露给远程读取器，而无需在源上调用CUDA内核。

源实例初始化过程
在源SGLang实例初始化期间：

  1. 每个TPWorker生成一个TransferEngine实例。
  2. TransferEngine将它的权重的GPU内存地址注册到RDMA通道。

目标实例初始化过程
在初始化目标实例时：

  1. 它发送一个HTTP请求，以检索源实例的TransferEngine元数据，包括映射到相应GPU内存地址的RDMA密钥（RDMA keys）。
  2. 使用这些RDMA密钥，目标实例直接从源的GPU内存中加载权重，而不会中断源实例正在进行的推理服务。

NCCL与TransferEngine的对比

  • 部署复杂性：NCCL无额外依赖；TransferEngine需要额外的 mooncake 库。
  • 传输设置开销：NCCL建立通信组需要数百毫秒；TransferEngine将内存区域注册到RDMA通道可能需要几秒钟，但可以与其他初始化阶段重叠。
  • 对GPU工作负载的无干扰性：NCCL张量传输将启动CUDA内核，产生干扰；TransferEngine传输权重不启动CUDA内核，实现无干扰。

使用方法

详细用法请参考R-Fork官方文档。

使用NCCL作为后端
种子实例（seed instance）启动命令：

python -m sglang.launch_server [args]

客户端实例（client instance）启动命令：

python -m sglang.launch_server [args] \
  --load-format remote_instance	\
  --remote-instance-weight-loader-seed-instance-ip [seed_instance_ip] \
  --remote-instance-weight-loader-seed-instance-service-port [seed_instance_service_port] \
  --remote-instance-weight-loader-send-weights-group-ports [send_weights_nccl_group_ports_list]  \
  --remote-instance-weight-loader-backend nccl

使用TransferEngine作为后端
种子实例启动命令：

python -m sglang.launch_server [args] \
  --remote-instance-weight-loader-start-seed-via-transfer-engine

客户端实例启动命令：

python -m sglang.launch_server [args] \
  --load-format remote_instance	\
  --remote-instance-weight-loader-seed-instance-ip [seed_instance_ip] \
  --remote-instance-weight-loader-seed-instance-service-port [seed_instance_service_port] \
  --remote-instance-weight-loader-backend transfer_engine

工业实践：Tensor R-Fork Planner

手动配置的局限性
在前面的部分中，演示了如何在SGLang服务器参数中为Tensor R-Fork手动配置种子实例。然而，这种手动方法在现实世界的工业部署中是不切实际的，因为识别可用的种子实例需要大量的运维开销。

集群调度器设计
为了应对这一挑战，提出了 Tensor R-Fork Planner【4，Tensor R-Fork Planner SGLang RFC+2025+Github Issue+https://github.com/sgl-project/sglang/issues/12910】，这是一个旨在编排源实例元数据的集群调度器。Planner跟踪的关键信息包括 ：

  1. 模型兼容性：实例上当前正在运行哪个模型。
  2. 并行配置：采用的并行策略（例如，张量并行、流水线并行）。
  3. 服务健康状态：实例是否健康且适合作为种子实例。

工作流程
在完成初始化后，每个实例都会向Planner注册自己，提供其模型元数据和并行配置。当新实例启动时，它首先查询Planner以识别与它的模型和并行策略都匹配的合格种子实例。如果找到兼容的种子实例，新实例将直接从种子加载权重；否则，它将回退到默认的加载格式。

实验环境

  • 硬件配置：配备8张 NVIDIA H20 GPU 的节点。
  • 模型：DeepSeek-R1 模型。
  • 软件配置：SGLang 实例。

实验结果

不同数据源权重加载性能评估
实验评估了在启动配备8张 NVIDIA H20 GPU 的新SGLang实例时，从不同来源加载DeepSeek-R1模型的性能表现。
实验结果如图2所示，相较于传统的加载方式，使用 Tensor R-Fork 技术极大地缩短了模型权重的加载时间。对于 TransferEngine 后端，注册内存区域（Registering the memory region）的操作可以与其他初始化阶段重叠，从而能够进一步优化总启动时间。
加载DeepSeek-R1模型的性能对比图

结论与未来工作

Tensor R-Fork 成功解决了大规模 LLM 服务中 SGLang 实例冷启动时的权重加载瓶颈问题，显著缩短了加载时间并节省了存储开销。
R-Fork 的实践开启了更多富有想象力的可能性：R-Fork 的核心概念是使所有 SGLang 实例都能够充当其他实例的数据存储中心。未来，团队计划从权重张量开始，通过 Tensor R-Fork 机制管理额外的张量（additional tensors），使 GPU 集群不仅可以作为计算中心，还可以作为存储中心运行。

附录：TransferEngine 细节

TransferEngine【5，TransferEngine+2025+URL+https://kvcache-ai.github.io/Mooncake/design/transfer-engine.html】提供了以下关键优势 ：

  • 多后端支持：TE支持多个后端，包括RDMA（带有GPUDirect）、NVLink、GDS和TCP。它可以智能地识别每个请求的最佳后端，从而达到最高性能。
  • 直接RDMA读取：使用发布的地址和rkeys，目标端直接将RDMA操作（通常是RDMA READ）执行到其自己预先分配的GPU缓冲区中，利用GPU-Direct RDMA，因此不需要主机到设备（host-device）或设备到主机（device-host）的中间内存拷贝。
  • 无干扰（Non-disturbing）：TE执行纯网卡驱动（NIC-driven）的传输，避免在源GPU上启动CUDA内核。
  • 生命周期与内务管理（Lifecycle & housekeeping）：TE维持注册的生命周期，直到张量被驱逐或进程退出。
  • 并发与流量控制：TE协调并发读取（来自一个或多个目标），并可以应用节流（throttling）或速率限制，以避免使实例的网卡饱和或影响推理延迟。

当前 TransferEngine 实现的已知限制：

  • 内存注册（register_mr）速度慢：这是由于RDMA驱动程序造成的。开发团队呼吁社区提供见解或解决方案，并渴望共同探索创新方法。

补充细节：致谢

本研究的成功离不开以下团队的贡献：

  • 蚂蚁集团 DeepXPU 团队：Anqi Shen, Tianyu Zhou, Zehuan Li, Tiwei Bie, Mingliang Gong, Jianfeng Tan。
  • SGLang 团队：Chenyang Zhao, Liangsheng Yin, Lianmin Zheng。
  • TransferEngine 团队：Teng Ma, Feng Ren, Shangming Cai。