新媒网 · 跨境数据社区

近年来，随着人工智能技术的飞速发展，特别是大型语言模型（LLM）的兴起，对分布式训练框架的性能、可扩展性和鲁棒性提出了前所未有的要求。在这一背景下，PyTorch分布式（PyTorch Distributed，简称PTD）生态系统迎来了一项重要更新：全新的实验性轻量级通信API torchcomms 正式推出。这一创新旨在为PyTorch用户提供更高效、更灵活的分布式通信解决方案，以应对超大规模模型训练的挑战。

torchcomms 的发布，不仅仅是一个新API的问世，它更代表着PyTorch在分布式训练领域迈向新一代架构的关键一步。其核心目标是解决现有PTD通信机制在处理超大规模并行计算时面临的瓶颈，并为未来的技术演进奠定基础。

事件概述：torchcomms 与 NCCLX 的发布

近日，PyTorch团队与Meta公司共同宣布，正式发布实验性的轻量级通信API torchcomms。伴随 torchcomms 的推出，一个名为 NCCLX 的新后端也同步开源。NCCLX 是Meta为支持超十万级别GPU集群训练所开发的通信后端，已在其内部用于Llama3和Llama4等大型语言模型的生产级训练与推理任务，其成熟度与生产实践经验为 torchcomms 提供了坚实支撑。

torchcomms 的发布，旨在提供一套全新的核心通信原语，以确保在当前2025年大规模分布式训练场景中的可靠性和高性能表现。这一初始版本聚焦于分布式训练所需的基础API和后端支持，为研究人员和开发者提供了进行超大规模模型训练所需的关键工具。

各方回应与事态进展

此次发布引发了PyTorch社区和广大AI开发者的广泛关注。作为PyTorch分布式通信体系的革新者，torchcomms 采取了完全开放的开发模式，旨在邀请社区参与反馈，共同推动API的成熟与演进。尽管初期阶段可能会存在一些破坏性变更，但这正是开放式创新过程的特点。

未来一年，torchcomms 将持续完善，计划引入更多新功能，例如更便捷的原型开发新集合通信操作、内置故障容忍能力的无缝扩展、以及设备中心化通信模式的优化。长远来看，torchcomms 有望成为下一代分布式技术验证的平台，并最终逐步取代现有的PyTorch Distributed功能，成为PyTorch中可扩展、容错且以设备为中心的分布式训练的核心支柱。

核心目标与创新方向

torchcomms 的推出并非偶然，它承载着解决当前分布式训练痛点、引领未来技术发展的多重战略目标。

1. 通信原语的快速原型开发能力：机器学习研究人员需要频繁试验新的通信范式。torchcomms 将通信层与PyTorch的核心数值计算原语解耦，使得通信层可以独立迭代，方便添加新的集合通信操作、API或后端，而不会影响现有功能。这种设计还支持树外（out-of-tree）后端，允许研究人员和硬件供应商轻松集成针对其设备和功能优化的专用通信栈。

2. 支持超十万级别GPU扩展：将现代训练工作负载扩展到数十万个GPU，需要重新思考通信资源的管理方式。当前的一些方法，如惰性初始化和点对点操作的有限并发语义，限制了NCCL等库的可扩展性。torchcomms 引入了即时初始化（eager initialization），用户可以显式管理后端资源，并通过模型特定的提示来优化通信器、NVLink缓冲区和RoCE资源的分配与共享，为真正的超大规模分布式作业铺平道路。

3. 异构硬件支持：现有集合通信后端通常针对单一供应商或硬件系列进行优化。torchcomms 在设计之初就考虑了异构系统，支持在单一训练作业中跨多个硬件代际和供应商进行混合部署。随着生态系统从同构GPU集群向多样化硬件发展，这种灵活性至关重要。

4. 大规模故障容忍能力：当前开源的PyTorch Distributed缺乏强大的容错进程组，这限制了如torchft等高级库的可靠性。torchcomms 旨在弥补这一差距，通过开源一个能够支持故障容忍HSDP和故障容忍Streaming DiLoCo等算法的后端，在不牺牲性能的前提下，实现大规模训练的韧性。

5. 单边通信能力：单边通信（例如RDMA风格语义）对于强化学习、检查点和大型语言模型中的异步工作流越来越重要。torchcomms 将提供对单边通信的一流支持，实现分布式进程之间高效、低开销的消息传递和数据交换。

6. 设备中心化集合通信：为了实现推理和训练的超低延迟，通信和计算必须紧密耦合。torchcomms 正在开发设备中心化集合通信API，这些API允许通信元数据和逻辑直接驻留在设备（如GPU）上。这包括来自GPU的直接RDMA操作（如IBGDA）和基于CPU代理的设计。这些能力使开发者能够无缝融合计算和通信操作，解锁新的性能水平。

为何需要全新的API？

PyTorch团队表示，开发 torchcomms 的原因在于追求一系列宏大目标，引入了现有通信库中尚不存在的功能。为了快速推进，团队需要在开放环境中自由迭代和演进设计，而不受现有接口的束缚。这意味着，在早期阶段，API可能会随着社区的实验和改进而发生破坏性变更。

现有的PyTorch Distributed c10d API存在显著的技术债务，使其难以扩展或现代化。随着 torchcomms API的稳定，团队计划逐步淘汰旧的 c10d::Backend 接口，并采用 torchcomms 作为PyTorch Distributed的底层实现。这一过渡将循序渐进，尽量减少对现有用户的干扰，大多数用户和模型将继续照常工作，同时自动受益于新后端带来的性能、可扩展性和灵活性。

核心数据与行业背景关联

torchcomms 和 NCCLX 的发布，直接响应了当前AI行业对超大规模计算能力和效率的迫切需求。新媒网跨境了解到，Meta公司在生产环境中，使用 NCCLX 成功支撑了如Llama3和Llama4这类万亿参数级别的大型语言模型的训练与推理。这并非单纯的技术演示，而是真实的生产实践数据，它有力地证明了 torchcomms 及其配套后端在扩展性（支持十万级别GPU）和性能上的卓越表现。

当前的AI发展趋势显示，模型规模持续扩大，参数量从数十亿到万亿级别不断攀升，使得单设备训练已不可行，分布式训练成为主流。torchcomms 正是在这样的行业背景下应运而生，它通过解决现有分布式通信框架的瓶颈，为AI模型突破算力限制、实现更大规模和更复杂功能的训练提供了基础设施支持。

基础用法示例

torchcomms 的设计理念是提供一个围绕底层后端和通信器的轻量级封装。其核心API直接映射到后端方法，并采用完全面向对象的设计。

以下代码片段展示了 torchcomms 的基本使用方式，例如如何初始化通信器、执行 all_reduce 操作以及分割通信组：

import torchcomms

# 使用torchrun提供的MASTER_PORT/MASTER_ADDR/RANK/WORLD_SIZE环境变量即时初始化一个通信器。
# 此通信器绑定到单个设备。
comm = torchcomms.new_comm("ncclx", torch.device("cuda"), name="my_comm")
print(f"I am rank {comm.get_rank()} of {comm.get_size()}!")

t = torch.full((10, 20), value=comm.rank, dtype=torch.float)

# 在当前流上运行all_reduce
comm.allreduce(t, torchcomms.ReduceOp.SUM, async_op=False)

# 在后台流上运行all_reduce
work = comm.allreduce(t, torchcomms.ReduceOp.SUM, async_op=True)
work.wait()

# 将通信器分成8个一组
split_groups = torch.arange(comm.get_size()).view(-1, 8).tolist()
tp_comm = comm.split(split_groups)

torchcomms 也支持与DeviceMesh兼容，以支持PyTorch的并行库，例如FSDP2。

import torchcomms
from torchcomms.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard

comm = torchcomms.new_comm("ncclx", torch.device("cuda:0"), name="global")

mesh = init_device_mesh(
    mesh_dim_comms=(comm,),
    mesh_dim_names=("global",),
)

fully_shard(model, device_mesh=mesh)

初始后端支持

伴随 torchcomms API的发布，同步推出了多个针对不同硬件平台的后端。

1. NCCLX：NCCLX 是Meta对流行的NCCL库的扩展版本，已在生产环境中经过验证，用于Llama3和Llama4等大型语言模型的大规模训练和推理。当前，Meta所有的生成式AI服务均由 NCCLX 提供支持。其关键特性包括可扩展初始化、零拷贝和SM-free通信、自定义集合通信算法、网络流量负载均衡、单边通信、GPU驻留和低延迟集合通信，以及故障分析和定位功能。为了承载Meta的内部优化和定制功能，团队还并行开发了一个独立的自定义传输（Custom Transport，简称CTran）栈。CTran包含NVLink、IB/RoCE和TCP传输，通过不同的硬件例程支持底层通信原语，并为各种通信语义（如集合通信、点对点、RMA）构建通信算法。NCCLX 和 CTran 都已随 torchcomms 一同开源。

2. NCCL和RCCL：除了 NCCLX，torchcomms 还支持上游的NCCL。当前PyTorch Distributed的NCCL用户可以轻松尝试 torchcomms，无需更改底层通信库设置。在当前的PyTorch Distributed中，对AMD RCCL的支持是通过NCCL进程组实现的。作为 torchcomms 发布的一部分，团队还包含了一个原生的RCCL后端。这使得 torchcomms 能够从发布之初就提供原生的多供应商GPU支持，允许不同库更独立地发展。

3. Gloo：Gloo通常用于机器间CPU元数据传输或测试场景。但它也新增了一些高级功能，如Infiniband和单边操作。近期还新增了“惰性初始化”模式，使得Gloo可以扩展到十万甚至更多的工作节点。

新的API特性

torchcomms 对现有集合通信语义进行了多项更改，旨在使其高级语义与底层设备语义更好地匹配，并提高灵活性。

1. 所有操作均通过面向对象API完成，而非使用全局的 dist.* API。每个 torchcomm.TorchComm 对象映射到单个设备和通信器。
2. 后端是即时初始化的，并在创建时需要传入一个设备。
3. 所有操作都使用通信器内的排名（communicator ranks），而非“全局”排名。
4. 所有操作按发出顺序执行，并发操作必须使用批处理（batch）API运行。发送/接收操作除非通过批处理API发出，否则不会并发执行。

此外，torchcomms 还引入了以下实验性API：

1. 窗口API（Window APIs）：支持动态的远程内存“put/get”操作。对于检查点、异步操作等特定用例，这可以显著提高性能且更容易表达，因为只需单方参与，不同于传统的集合通信。窗口API允许用户在不同排名（rank）之间创建内存缓冲区（无论是GPU还是CPU内存），并自动注册，通过提供的Put和Get API访问，利用底层的RDMA或NVL传输实现零拷贝、单边通信。此外，窗口API还提供原子信号机制，进一步增强异步通信能力。目前窗口API仍处于积极开发和实验阶段。

2. 传输API（Transport APIs）：允许直接使用底层传输进行点对点操作。这提供了与窗口API类似的接口，但不受集合通信库的约束。初始阶段，仅提供基于专用网络的RDMA传输支持，主要用于RPC类操作，由CTran中的IB后端提供内部支持。RdmaTransport 提供了一个写入API，允许用户直接写入远程内存。用户需要注册内存并在进程间交换其句柄以促进写入。这有效地实现了零拷贝数据传输，可以对CPU或GPU内存进行零拷贝操作。这些API仅是传输API，不执行任何计算（如reduce等）。目前传输API仍处于积极开发和实验阶段。

3. 故障容忍API（Fault Tolerance APIs）：团队正致力于创建新的故障容忍集合通信后端。该后端完全基于CTran传输构建，提供故障检测、超时、错误恢复和错误发生后的安全重新配置功能。

可扩展性

torchcomms 在设计上强调了高度的可扩展性，方便开发者进行二次开发和定制。

1. 使用新集合通信操作扩展后端：torchcomms 旨在支持对底层后端的直接访问。这使得在标准化并将其添加到共享后端接口之前，可以快速原型开发新的API。例如，可以轻松添加自定义操作，如 quantized_all_reduce。
2. 编写新的torchcomm后端：torchcomms 的一个关键特性是它极大地简化了第三方后端的编写。这些后端不再需要作为PyTorch的一部分进行构建，可以通过 pip 等工具像其他Python扩展一样简单安装。

下一步展望

新媒网跨境观察到，torchcomms 作为一个全新的API，目前仍处于积极开发阶段。PyTorch团队邀请社区积极参与，共同改进。团队正努力推进本篇文章中描述的各项功能，并期望在不久的将来使其稳定，同时改进对更多硬件设备的兼容性支持。

欲了解更多文档信息，请访问：

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/pytorch-gets-upgrade-for-100k-gpu-llm-train.html