我们如何通过 GPU 间广播权重而非从存储下载,将模型冷启动时间从分钟级缩短至秒级,以及在此过程中发现的棘手问题。冷启动代价
当 GPU 推理工作节点启动时,它做的第一件事(在能处理任何请求之前)就是加载模型权重。对于 Runway 运行的这类模型而言,这意味着从云存储拉取数 TB 的参数,将其从主机内存复制到 GPU 内存,并调用 load_state_dict。对于单个工作节点,这需要几分钟时间。对于同时推出新模型版本的一整批工作节点而言,这就像一场惊群效应:数十台机器独立地从同一个存储后端下载相同的字节,占满带宽,并让所有人的冷启动时间变得更长。我们每天部署数十次,而且通常不涉及模型权重变更。冷启动直接制约了自动扩缩容、部署速度和面向用户的延迟,并拉长了研究与工程之间的反馈循环。
由于模型权重变更相对少见,如此频繁地部署却伴随着漫长的冷启动,这给 GPU 推理工作节点的可用性带来了不必要的代价,并影响了我们的用户体验。这些权重其实已经存在于集群中某处的 GPU 上。五分钟前完成启动的工作节点正待在那里,权重已完全加载,等待任务。新的工作节点可以直接从对等节点接收这些权重,而无需从存储重新下载。这保证了生成任务的持续进行,并让我们的研究和工程团队能够不断交付新功能。该系统让我们的集群变得灵活:我们可以快速重新平衡工作负载、回滚有问题的代码,并最大化 GPU 用于执行推理的时间。各类方案概览
冷启动优化有几种方法,我们尝试了其中大部分。
更快的存储(NVMe 缓存、并行分块下载)能减少下载时间,但并未改变问题的本质:N 个工作节点仍然要进行 N 次独立的下载。
共享内存预加载(/dev/shm、CUDA IPC)消除了单节点上的冗余下载,但本质上局限于单机。
共享卷(NFS、EFS、Ceph)可以在多台机器间对下载内容进行去重,但会增加运维开销,而且通常仍远慢于 GPU 之间的直接传输。
分布式缓存(Memorystore、自定义权重服务器)会引入另一个需要运维、容量规划和监控的系统。
点对点传输则不同。一台工作节点正常下载并加载权重,后续每台工作节点直接从已加载权重的对等节点,通过 GPU 互联接收权重。下载只发生一次,广播可扩展到整个集群。我们围绕这一思路构建了一个名为 NCCLBack 的系统(其命名源于所依赖的 NCCL 集合通信,以及它“支撑”权重加载路径的功能,顺便也与某支加拿大摇滚乐队沾点边)。
点对点传输之所以能取得压倒性优势,原因在于物理规律。GCS 下载带宽即使经过优化,每台工作节点也大约只有 2–10 Gbps。而 GPU 节点间的 InfiniBand 或 RoCE 链路可提供 200–400 Gbps。节点内的 NVLink 速度更快,在 H100 SXM 系统上可达 900 GB/s。当你能用 GPU 互联替代存储 API 时,就相当于从“分钟级”世界进入了“秒级”世界。
在传统模式下,每台工作节点独立从存储下载,冷启动时间随集群规模线性增长。而使用 NCCLBack,一台工作节点下载,其余节点在数秒内通过 GPU 互联接收权重:
图 1. 冷启动时间 vs. 集群规模:传统存储下载的耗时随工作节点数量线性增长,而 NCCLBack 通过 GPU 互联广播权重,基本保持平稳。NCCLBack 设计:分层构建
NCCLBack 并非单一的巧妙技巧,而是一套子系统堆栈(发现、协调、传输与验证),每个子系统解决一个独立问题。以下是完整协议概览:
图 2. NCCLBack 协议端到端流程,展示接收方如何发现对等节点、完成存活握手、通过 NCCL 广播传输权重,并在提供服务前验证完整性。第一层:基本思路
NCCLBack 的核心是利用 NCCL 的广播原语,在两个 GPU 工作节点之间进行一对一的权重传输。
发送方是一个已在 GPU 上加载了权重的 worker。接收方是一个新启动的 worker,它已经构建好了模型结构(所有正确的张量形状、数据类型和设备),但尚未填充具体数值。我们将这种状态称为"骨架模型"。
传输过程本身很直接。双方按排序后的键顺序遍历 `model.state_dict()`,并对每个张量调用 `nccl.broadcast()`:
对键进行排序的重要性远超表面所见。`nccl.broadcast` 是一个集合操作:双方必须以完全相同的顺序、针对完全相同的张量大小和数据类型发起调用。如果发送方与接收方的遍历顺序不同,广播操作会静默地将错误数据写入错误的缓冲区,最终得到一个输出垃圾结果且不报任何错误的模型。
这就是完整的数据平面。NCCLBack 中的所有其他功能都是为了确保这个循环正确、安全且不卡死地运行。第二层:发现
在进行广播之前,接收方需要找到一个发送方。这是一个协调问题。
我们通过一个基于 Redis 的队列来解决这个问题。当一个 worker 完成设置并加载好权重后,它会向 Redis 发布两样东西:
- 模型可用性:一个表示"我已加载模型版本 X"的键。
- 权重哈希:逐层的完整性哈希(更多细节见第五层)。
当一个新的 worker 启动并希望接收权重时,它首先检查 Redis:是否存在一个拥有我所需模型版本的 peer?这个检查在分布式 worker 组内的所有 rank 之间必须达成一致。如果 rank 0 认为存在某个 peer,但 rank 3 认为不存在,那么当部分 rank 尝试 P2P 传输而其他 rank 尝试下载时,整个组就会死锁。
我们通过一个在 NCCLBack 中反复出现的模式来强制执行共识:
all_reduce(MIN) 技巧:每个 rank 进行投票,只有当最小投票数为 1 时我们才继续。如果任何一个 rank 在 Redis 中找不到该模型,所有 rank 都回退到下载。这是一个刻意保守的选择;我们宁愿冗余下载,也不愿冒险进行部分 P2P 传输,导致某些 rank 有权重而其他 rank 没有。
模型身份在此至关重要。“模型版本”是模型配置和检查点元数据的哈希值。实际操作中,每个检查点在部署时都会转换为 safetensors 格式,我们的云存储会为生成的文件提供 CRC32C 校验和。然后,模型配置(包含 safetensors URI 及其 CRC32C)会通过 MD5 进行哈希运算,生成版本标识符。这意味着,只有当两个工作节点对完全相同的检查点字节和模型配置达成一致时,它们才认为彼此兼容。如果这个哈希值出错,那么两个拥有细微差异模型状态的工作节点会尝试互相传输权重,这要么会直接失败(形状不匹配),要么会静默失败(加载了错误的权重)。第三层:握手
在 Redis 中找到一个对等节点是不够的。Redis 只能告诉你某个时间点存在一个对等节点;它无法告诉你该节点是否仍然存活并准备好立即发送数据。如果我们乐观地启动 NCCL 广播,而发送方已经崩溃或正忙,那么 `ncclCommInitRank` 将会无限期挂起。NCCL 的集合通信模型在通信器创建时没有内置超时机制。
因此,在接触 NCCL 之前,我们通过 Redis 运行一个双向的存活握手:
- 接收方在一个已知的 Redis 键 `receiver_ready:{peer_id}` 处存储一个随机整数。
- 发送方轮询该键。当它出现时,发送方将该值加 1 并存储回去。
- 同时,发送方已将自己的随机整数存储在 `sender_ready:{peer_id}` 处。
- 接收方轮询发送方的键,找到它,将其加 1,并存储回去。
双方进行验证:“我存储了 X,对方将其改为了 X+1。”这证明两个对等节点都存活、有响应,并且正在讨论同一个传输任务。整个握手过程使用较短的超时时间:发送方为 1 秒(它不应为不稳定的接收方等待太久),接收方为 10 秒(它可以等待,因为对等节点可能只是正忙)。
如果握手超时,接收方会干净地回退到从存储下载。不会创建任何 NCCL 通信器,不会占用任何 GPU 资源,也无需进行清理。第四层:带防护栏的传输
握手成功后,我们会创建一个 2 秩 NCCL 通信器(发送方 = 秩 0,接收方 = 秩 1)并运行广播循环。但 NCCL 没有内置超时机制——如果对端在传输过程中消失,调用将永久阻塞。我们将通信器创建和广播循环都包裹在守护线程超时中:如果操作在几秒内未完成,就中止通信器并抛出 TimeoutError。这确保了挂起的传输不会阻塞工作节点回退到下载方式。第 5 层:完整性校验
传输完成后,我们如何知道接收方得到了正确的权重?NCCL 在应用层不提供校验和或完整性保证。一次静默数据损坏(GPU 间传输中的比特翻转)会产生一个能运行但输出轻微错误结果的模型。
我们通过采样哈希解决这个问题。当第一个工作节点从存储下载并加载权重("基准真值"路径)时,它会为模型的每一层计算一个哈希值,并将这些哈希值存储在 Redis 中。对千兆字节级别的模型逐字节哈希速度很慢,因此我们改用固定种子对每个参数采样 100 个随机元素(这样发送方和接收方采样相同的索引),仅对这些样本进行哈希。哈希前我们将其归一化为 float16 格式,这样代码路径之间的 dtype 差异不会导致误报不匹配。这使原本需要数分钟的校验缩短为几秒钟,同时仍能捕获任何影响超过极小比例权重的损坏。
接收方通过 P2P 接收权重后,计算相同的哈希值并与 Redis 中存储的"真值"进行比较。如果任何层的哈希值不匹配,我们就抛出 WeightVerificationError,工作节点重新启动。骨架模型问题
NCCL 的广播不分配内存。它直接写入现有缓冲区。这意味着接收方必须在广播开始前预分配形状、dtype 和设备完全正确的张量。我们将这些预分配的张量称为"容器",将其所在的模型结构称为"骨架模型"。
构建骨架模型意味着运行模型的构造函数和初始化代码,同时跳过实际权重加载(传入 `skip_checkpoint=True`)。对于简单模型来说,这很直接。但对于实际生产环境中的模型,则要复杂得多。一个视频生成管线并非单一的 `nn.Module`,而是一个子模块图(包括扩散 Transformer、VAE 解码器、文本编码器、图像嵌入器等),每个子模块都有自己的初始化逻辑,其中一些会在构造过程中急切地加载权重。
在流水线架构中,不同的 GPU 秩(rank)托管着完全不同的子模块。一个秩可能运行 VAE 解码器,而其余秩运行扩散 Transformer。NCCLBack 必须为每个秩返回正确的子模块,因为每个秩都参与其自身独立的点对点(P2P)传输。接收方不仅需要一个骨架,它需要的是适合其秩的正确骨架,且容器必须与发送方将要广播的内容完全匹配。任何不匹配(例如不同的数据类型)……