AllReduce 的环与双二叉树
用不同通信拓扑组织梯度归约,在带宽利用与启动延迟之间取舍。
相关概念
工程问题#
多 GPU 训练必须交换并聚合梯度;设备数增加后,通信阶段可能拖慢整步训练。
实际实现#
NCCL 2.4 的双二叉树将数据分成两部分,交换两棵树的内部节点与叶节点角色;每棵树执行归约再广播。环则沿相邻设备流水传输。
处理流程#
- 依据设备连接组织通信拓扑
- 将张量分块分配给通信路径
- 聚合来自各 rank 的局部结果
- 把归约结果传回全部 rank
源码片段#
固定版本 · 原文摘录src/collectives/device/all_reduce.hv2.4.8-1 · NCCL 2.4 系列 · L102–L117
102 do {103 // Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)104 ncclPrimitives<UNROLL, 1, 1, T, NCCL_MAX_TREE_ARITY, 1, FUNC> prims(tid, nthreads, tree->down, &tree->up, NULL, stepSize, channel, comm, args->opCount);105 for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {106 // Up107 ssize_t offset = gridOffset + bid*chunkSize;108 int nelem = min(chunkSize, size-offset);109 if (tree->up == -1) {110 prims.recvReduceCopy(thisInput+offset, thisOutput+offset, nelem);111 } else if (tree->down[0] == -1) {112 prims.send(thisInput+offset, nelem);113 } else {114 prims.recvReduceSend(thisInput+offset, nelem);115 }116 }117 } while(0);ncclAllReduceTreeKernel · Reduce stage
树根接收并归约后保存结果,叶节点只上发本地输入,内部节点接收、归约并继续上发。后续独立的 Broadcast 阶段把结果下发;该片段只展示树形归约内核,不包含双树拓扑构建或算法选择。环内核在同一文件另行实现。
CPPBSD-3-Clause
设计解释引用 NCCL 2.4 发布文章;代码使用同一 2.4 系列的 v2.4.8-1,不宣称与发布文章当天的二进制逐行相同。
收益与取舍#
工程分析带来的收益
- 环可实现高带宽利用
- 树的对数层数降低大规模通信延迟
付出的代价
- 环的通信轮次随 rank 数增长
- 树的实际带宽仍受物理拓扑与网络路由影响
可比较的选择
分层归约可先在节点内聚合再跨节点通信;效果取决于链路差异与消息大小。
版本与适用范围#
解释 NCCL 2.4 已公开的环与树取舍;后续 NCCL 还有其他算法,不据此宣称当前默认算法或固定切换阈值。
参考来源#
源码、官方文档与原始论文01
NCCL 2.4 双二叉树设计
官方文档NCCL 2.4 · 2019-02-04Double binary trees / Performance at scale · AllReduce / Double binary trees