SGLang 推出 Tensor R-Fork 权重加载方法，DeepSeek-R1 加载时间从数分钟缩短到数秒

简而言之

我们推出 Tensor R-Fork（即 Tensor Remote Fork），这是一种全新的权重加载方法，利用高效的节点间设备到设备互连，以零拷贝方式将张量从正在运行的 SGLang 实例加载到新实例。

我们的方法提供三大关键优势：

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

例如，应用于 Deepseek-R1 模型时，加载时间从数分钟缩短至仅数秒，同时本地磁盘和/或 DRAM 存储占用减少约 600GB，且模型传输期间推理服务质量保持稳定。

背景

随着 LLM 服务规模和模型权重体积持续扩大，SGLang 实例的冷启动时间已成为生产效率的关键瓶颈。在冷启动各阶段中，权重加载仍是最耗时的任务。

以 Deepseek-R1 为例，从本地磁盘加载权重通常需要数分钟，而从远程存储系统加载则可能长达数十分钟。随着模型规模持续指数级增长，初始化和数据传输所需时间很可能进一步恶化。

我们如何优化权重加载性能？最直接的方法是最大化权重数据流中的瓶颈带宽。业界常用模型加载方法的数据流及其相关瓶颈带宽如下：

我们能否利用更高带宽的数据流来传输张量？答案是肯定的——节点间设备到设备互连可提供每秒数百 GB 的吞吐量。然而，关键问题仍然存在：我们如何充分利用这种互连的带宽，在 SGLang 中实现高效的权重加载？

为应对这一挑战，我们开发了一种全新的权重加载框架，称为 Tensor R-Fork（即 Tensor Remote Fork），它将 Deepseek-R1 模型加载时间缩短至仅数秒，且已具备生产可用性。

设计

<a href=https://github.com/sgl-project/sglang/blob/main/docs/advanced_features/rfork.md>Tensor R-Fork[0] 的核心概念是利用 GPU-Direct RDMA 构建点对点（P2P）权重存储架构。

使用传统方法的数据传输性能较低，因为整个路径中始终存在瓶颈，其带宽远小于节点间设备到设备互连。 从数据流分析中，我们观察到权重张量存储在每个 GPU 上，可通过 GPU-direct RDMA 在节点之间直接传输。

为了最大化利用支持 RDMA 的网卡带宽，我们设计了一种按 GPU 对进行数据传输的策略：本地 GPU 直接与其配对的远程 GPU 进行数据收发。这一设计有效绕过了 GPU 与 CPU 之间的 PCIe 瓶颈，实现了不依赖 CPU 或主机内存的高吞吐通信。 从远程 SGLang 实例加载权重的数据流如下：

实现

为了让每个正在运行的实例都能作为任何需要相同模型的新实例的权重来源——同时尽量减少（或理想情况下消除）对正在运行实例推理服务的干扰——我们为该框架实现了两种后端选项：NCCL 和 TransferEngine。假设有一个正在运行的实例 A（称为源实例）和一个待启动的新实例 B（目标实例）。下面，我们将详细说明使用这两种后端进行权重传输机制的实现。

NCCL 后端

当使用 <a href=https://github.com/sgl-project/sglang/pull/8215>NCCL 作为后端[1]时，该过程包括两个阶段：

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

在目标实例初始化期间，它会向指定的源实例发送 HTTP 请求以发起通信组创建。目标实例的每个 TPWorker 都会与源实例中对应的 TPWorker 建立 NCCL 通信组（即源 rank 0 与目标 rank 0 配对，依此类推）。每个通信组恰好由两个成员组成：源 TPWorker 和目标 TPWorker。

通信组建立后，每个源 TPWorker 通过该组使用 NCCL broadcast 广播其位于 GPU 内存上的权重张量。目标 TPWorker 直接将权重接收到其 GPU 内存中，无需任何中间内存拷贝。

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

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

TransferEngine 后端

为了实现无干扰的权重传输，我们引入了另一种后端：<a href=https://github.com/sgl-project/sglang/pull/14997>TransferEngine，它利用 GPU-Direct RDMA 进行高效的数据移动[2]。TransferEngine（TE）是一个轻量级的基于 RDMA 的传输运行时，它与源实例上的每个 TPWorker 一起运行，并将驻留在 GPU 上的权重张量暴露给远程读取方，而无需在源端调用 CUDA 内核。

在源 SGLang 实例初始化期间：

1. 每个 TPWorker（张量并行 worker）都会启动一个 TransferEngine 实例。
2. TransferEngine 将其权重的 GPU 内存地址注册到 RDMA 通道。

在初始化目标实例时：

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

*想进一步了解 TransferEngine？非常欢迎查看附录部分中的 TransferEngine 🚀

NCCL 与 TransferEngine 对比

如何使用

详细用法请参考 <a href=https://github.com/sgl-project/sglang/blob/main/docs/advanced_features/rfork.md>R-Fork 文档

使用 NCCL 作为后端

种子实例：

python -m sglang.launch_server [args]

客户端实例：

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

性能

我们评估了启动一个配备八块 NVIDIA H20 GPU 的新 SGLang 实例的性能，同时从不同来源加载 DeepSeek-R1 模型。

注册内存区域可以与其他初始化阶段重叠进行，以进一步优化总启动时间。

工业实践

在前面的章节中，我们演示了如何在 SGLang 服务器参数中为 Tensor R-Fork 手动配置种子实例。然而，这种手动方式在实际工业部署中并不实用，因为识别可用的种子实例需要大量的运维开销。

为了解决这一挑战，我们提出了 <a href=https://github.com/sgl-project/sglang/issues/12910>Tensor R-Fork Planner[4]，一个用于编排源实例元数据的集群调度器。该 Planner 跟踪关键信息，包括：

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

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

未来工作

R-Fork 的实践开启了更多富有想象力的可能性：R-Fork 的核心概念是使所有 SGLang 实例都能充当其他实例的数据存储中心。从权重张量开始，未来我们将通过 Tensor R-Fork 机制管理更多张量，使 GPU 集群不仅能作为计算中心，还能作为存储中心。

致谢

蚂蚁集团 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

参考文献

[0] Tensor R-Fork 文档：<a href=https://github.com/sgl-project/sglang/blob/main/docs/advanced_features/rfork.md>Documentation
[1] 使用 NCCL 后端的 Tensor R-Fork：<a href=https://github.com/sgl-project/sglang/pull/8215>PR#8215
[2] 使用 TransferEngine 后端的 Tensor R-Fork：<a href=https://github.com/sgl-project/sglang/pull/14997>PR#14997
[3] 从磁盘并发加载权重：<a href=https://github.com/sgl-project/sglang/pull/7943>PR#7943
[4] Tensor R-Fork Planner SGLang RFC：<a href=https://github.com/sgl-project/sglang/issues/12910>Issue#12910
[5] TransferEngine：<a href=https://kvcache-ai.github.io/Mooncake/design/transfer-engine.html>https://kvcache-ai.github.io/Mooncake/design/transfer-engine.html，
TransferEngine API：<a href=https://kvcache-ai.github.io/Mooncake/python-api-reference/transfer-engine.html>https://kvcache-ai.github.io/Mooncake/python-api-reference/transfer-engine.html

附录

TransferEngine

TransferEngine[5] 提供的关键优势：

• 多后端支持：TE 支持多种后端，包括 RDMA（配合 GPUDirect）、NVLink、GDS 和 TCP。它能够智能地为每个请求识别最佳后端，从而可以达到最高性能。
• 直接 RDMA 读取：使用已发布的地址和 rkey，目标端直接对其预分配的 GPU 缓冲区执行 RDMA 操作（通常是 RDMA READ），借助 GPU-Direct RDMA，因此无需主机到设备或设备到主机的中间拷贝。
• 无干扰：TE 执行纯网卡驱动的传输，避免在源 GPU 上启动 CUDA 内核。
• 生命周期与内务管理：TE 维护注册的生命周期，直到张量被逐出或进程退出。
• 并发与流控：TE 协调并发读取（来自一个或多个目标端），并可施加节流或速率限制，以避免占满实例的网卡或影响推理延迟。

当前 TransferEngine 实现中的已知限制：

• 内存注册（register_mr）速度慢：这是由 RDMA 驱动导致的。如果您对此问题有任何见解或解决方案，我们将非常感激能听到您的意见。我们重视多元视角，并热切期待与您共同探索创新方法。