NCCL源码解读3.1：double binary tree双二叉树构建算法，相比ring环算法的优势

目录

一、双二叉树出现的原因

二、双二叉树介绍

三、双二叉树大规模性能

四、双二叉树源码解读

双二叉树注意事项

核心逻辑

源码速递

---

视频分享在这：

NCCL集合通信源码解读、案例、任务调度、拓扑_哔哩哔哩_bilibili

一、双二叉树出现的原因

ring环算法，延迟随 GPU 数量线性扩展，从而无法扩展到数百个 GPU 以上。

想象一下，使用数以万计的 GPU 来训练您的神经网络。使用多个 GPU 来训练神经网络在所有深度学习框架中都已变得非常普遍，提供优化的多 GPU 和多机器训练。用于对多个 GPU 上的梯度求和的 Allreduce 操作通常使用环 [1] [2] 来实现，以实现全带宽。环的缺点是延迟随 GPU 数量线性扩展，从而无法扩展到数百个 GPU 以上。进入 NCCL 2.4。

许多大规模实验已经用分层的 2D 环算法 [3] [4] [5] 取代了平面环，以获得相当好的带宽，同时降低延迟。

NCCL 2.4 现在增加了双二叉树，它提供完整的带宽和甚至低于 2D 环形延迟的对数延迟。

二、双二叉树介绍

双二叉树于 2009 年在 MPI 中引入 [6]，其优势是将广播和 reduce 操作的全带宽（可以组合成一个执行 reduce，然后执行广播的 allreduce）和对数延迟相结合，从而在中小型操作中实现良好的性能。

在 NCCL 中，我们使用一种易于实现的模式构建二叉树，该模式使局部性最大化，如图 1 所示。

图 1.二叉树

双二叉树依赖于这样一个事实，即二叉树中一半或更少的秩是节点，而一半（或更多）秩是叶。因此，我们可以用叶子作为节点来构建第二棵树，反之亦然。可能有一个等级在两棵树上都是叶子，但没有一个等级是两棵树上的节点。

图 2 显示了如何使用上面的模式通过翻转树来反转节点和叶子来构建双二叉树。

图 2.两个互补的二叉树，其中每个秩最多是一个树中的一个节点和另一个树中的一个叶子。

如果将这两个树叠加，则所有等级都同时具有两个父级和两个子级，但根级除外，根级只有一个父级和一个子级。如果我们用两棵树中的每一棵来处理一半的数据，那么每个 rank 最多会收到两次一半的数据，两次发送一半的数据，这在发送/接收的数据方面就像 rings 一样最优。

三、双二叉树大规模性能

我们在各种大型机器上测试了 NCCL 2.4，包括 Summit [7] 超级计算机，最高可达 24,576 个 GPU。如图 3 所示，使用树时延迟显著改善。与ring的差异随着规模的增加而增加，在 24k GPU 上提升高达 180 倍。

图 3.NCCL 延迟高达 24,576 个 GPU

我们确认该系统使用双二叉树保持全带宽。在规模上，当我们在 InfiniBand 结构中跨越 L3 交换机时，带宽会略有下降，我们认为这是由于 NCCL 通信模式和 InfiniBand 路由算法之间的效率低下造成的。

虽然并不完美，但将来可能会改进。即便如此，由于初始延迟较小，即使带宽受限，树仍然显示出明显的优势。但是，当该模式导致更大的带宽时，NCCL 会自动切换回环。

图 4.多达 24,576 个 GPU 上的 NCCL 总线带宽

四、双二叉树源码解读

4.1 双二叉树注意事项

在正式开始源码解读之前，有几点是需要大家注意的！

待补充。。。。

 4.2 connectTrees()

核心逻辑

1、connectTrees()中调用ncclGetDtree()构建双二叉树

2、ncclGetDtree()调用ncclGetBtree()构建树

这里重点想和大家强调ncclGetDtree()的输入，注意看是node不是rank。

源码速递

源码位置：nccl-2.19\src\graph\connect.cc

static ncclResult_t connectTrees(struct ncclComm* comm, int* treeToParent, int* treeToChild0, int* treeToChild1, int* treePatterns) {
  const int nChannels = comm->nChannels, nNodes = comm->nNodes, node = comm->node;

  // Compute tree depth. Not an exact value but a good approximation in most
  // cases
  int depth = comm->nRanks/nNodes - 1 + log2i(nNodes);

  int t0u, t0d0, t0d1, t0ChildType, t1u, t1d0, t1d1, t1ChildType;
  int* ttp, *ttc0, *ttc1;
  NCCLCHECK(ncclGetDtree(nNodes, node, &t0u, &t0d0, &t0d1, &t0ChildType, &t1u, &t1d0, &t1d1, &t1ChildType));

  //其它代码...................

}

 4.3 ncclGetDtree() 

核心逻辑

1、调用ncclGetBtree来生成第一棵树

2、nranks是奇数时，通过“移位”操作生成第二颗树。即，每个rank的排名都会向前移动一个位置（通过取模操作处理边界情况），再调用ncclGetBtree来生成第二棵树。

3、nranks是偶数时，通过“镜像”操作生成第二颗树。即，每个rank的排名都会取镜像值（nranks-1-rank），再调用ncclGetBtree来生成第二棵树。

源码速递

源码位置：nccl-2.19\src\graph\trees.cc

/* Build a double binary tree. Take the previous tree for the first tree.
 * For the second tree, we use a mirror tree (if nranks is even)
 *
 * 0---------------8                   3----------------11
 *          ______/ \                 / \______
 *         4         \               /         7
 *       /   \        \             /        /   \
 *     2       6       10         1        5      9
 *    / \     / \     /  \       / \      / \    / \
 *   1   3   5   7   9   11     0   2    4   6  8   10
 *
 * or shift it by one rank (if nranks is odd).
 *
 * 0---------------8            1---------------9
 *          ______/ \______              ______/ \______
 *         4               12           5                0
 *       /   \            /           /   \            /
 *     2       6       10           3       7       11
 *    / \     / \     /  \         / \     / \     /  \
 *   1   3   5   7   9   11       2   4   6   8  10   12
 */
// 定义一个函数，用于获取双通信树。输入参数包括总节点数、当前节点标识符，以及指向存储树信息的指针。

ncclResult_t ncclGetDtree(int nranks, int rank, int* s0, int* d0_0, int* d0_1, int* parentChildType0, int* s1, int* d1_0, int* d1_1, int* parentChildType1) {


  // 首先构建第一个通信树（基础树），使用B树（Binary Tree）的构建方法。
  // 该步骤填充了s0, d0_0, d0_1, parentChildType0指针所指向的数组。
  ncclGetBtree(nranks, rank, s0, d0_0, d0_1, parentChildType0);


  // 接下来构建第二个通信树。根据总节点数nranks的奇偶性，决定是构建镜像树还是位移树。
  // 如果nranks是奇数，则构建位移树。
  if (nranks % 2 == 1) {
    // 计算位移后的节点标识符。这里通过(rank-1+nranks) % nranks来确保结果是非负的。
    int shiftrank = (rank-1+nranks) % nranks;

    // 临时变量，用于存储B树构建过程中得到的父节点、两个子节点的标识符。
    int u, d0, d1;

    
    // 使用B树的构建方法，根据位移后的节点标识符来获取位移树的父节点和子节点。
    ncclGetBtree(nranks, shiftrank, &u, &d0, &d1, parentChildType1);

    // 由于是位移树，所以需要将得到的节点标识符进行模nranks运算，以确保它们在有效范围内。
    // 注意这里的+1操作，是因为位移树的构建是基于位移后的节点标识符，而我们需要的是原始标识符空间中的对应值。
    *s1 = u == -1 ? -1 : (u+1) % nranks;
    *d1_0 = d0 == -1 ? -1 : (d0+1) % nranks;
    *d1_1 = d1 == -1 ? -1 : (d1+1) % nranks;

  } else {
    // 如果nranks是偶数，则构建镜像树。
    // 临时变量，用途同上。
    int u, d0, d1;
  

    // 使用B树的构建方法，但这次是根据镜像后的节点标识符（即nranks-1-rank）来获取镜像树的父节点和子节点。
    ncclGetBtree(nranks, nranks-1-rank, &u, &d0, &d1, parentChildType1);

    
    // 由于是镜像树，所以需要将得到的节点标识符进行镜像处理，即使用nranks-1减去对应的标识符。
    *s1 = u == -1 ? -1 : nranks-1-u;
    *d1_0 = d0 == -1 ? -1 : nranks-1-d0;
    *d1_1 = d1 == -1 ? -1 : nranks-1-d1;
  }
 
  // 函数执行成功，返回ncclSuccess。
  return ncclSuccess;

}

 4.4 ncclGetBtree

核心逻辑

1、通过位操作找到rank中的第一个非零位（bit）。

2、rank0是树的根节点，父节点为-1（无父节点），并计算第一个子节点（如果存在）

3、对于非根节点，计算父节点的排名。

4、对于非根节点，计算两个子节点的排名。

源码速递

源码位置：

nccl-2.19\src\graph\trees.cc

/* Btree which alternates leaves and nodes.
 * Assumes root is 0, which conveniently builds a tree on powers of two,
 * (because we have pow2-1 ranks) which lets us manipulate bits.
 * Find first non-zero bit, then :
 * Find the parent :
 *   xx01[0] -> xx10[0] (1,5,9 below) or xx00[0] if xx10[0] is out of bounds (13 below)
 *   xx11[0] -> xx10[0] (3,7,11 below)
 * Find the children :
 *   xx10[0] -> xx01[0] (2,4,6,8,10,12) or -1 (1,3,5,7,9,11,13)
 *   xx10[0] -> xx11[0] (2,4,6,8,10) or xx101[0] (12) or xx1001[0] ... or -1 (1,3,5,7,9,11,13)
 *
 * Illustration :
 * 0---------------8
 *          ______/ \______
 *         4               12
 *       /   \            /  \
 *     2       6       10     \
 *    / \     / \     /  \     \
 *   1   3   5   7   9   11    13
 */

// 定义一个函数，用于获取B树的结构信息。输入参数包括总节点数、当前节点标识符，以及指向存储父节点和两个子节点标识符的指针。
ncclResult_t ncclGetBtree(int nranks, int rank, int* u, int* d0, int* d1, int* parentChildType) {

  int up, down0, down1; // 声明用于存储父节点和两个子节点标识符的变量
  int bit; // 声明用于位操作的变量

  // 通过位操作找到rank中第一个非零位的位置
  for (bit=1; bit<nranks; bit<<=1) {
    if (bit & rank) break; // 若bit位与rank的对应位均为1，则跳出循环
  }

  // 处理根节点（rank为0）的特殊情况
  if (rank == 0) {
    *u = -1; // 根节点没有父节点
    *d0 = -1; // 根节点可能没有第一个子节点（取决于nranks）

    // 根节点的第二个子节点（如果存在）是bit右移一位得到的节点
    *d1 = nranks > 1 ? bit >> 1 : -1;
    return ncclSuccess; // 函数执行成功
  }

 
  // 根据位操作计算父节点的标识符
  up = (rank ^ bit) | (bit << 1); // 若rank的bit位为1，则将其翻转并设置更高一位为1，得到父节点的候选值

  // 若候选父节点标识符超出范围，则使用另一个候选值（仅翻转bit位）
  if (up >= nranks) up = (rank ^ bit);
  *parentChildType = (rank < up) ? 0 : 1; // 根据当前节点与父节点的关系设置父子关系类型（0为第一个子节点，1为第二个子节点）

  *u = up; // 存储父节点标识符

 
  int lowbit = bit >> 1; // 计算低位的值（用于确定子节点的位置）
  // 第一个子节点总是在当前节点的低位侧，但需注意边界情况
  down0 = lowbit == 0 ? -1 : rank-lowbit;

 
  // 第二个子节点的计算需要考虑边界情况，确保其标识符在有效范围内内
  down1 = lowbit == 0 ? -1 : rank+lowbit;
  while (down1 >= nranks) { // 若第二个子节点标识符超出范围，则通过不断右移lowbit来尝试修正
    down1 = lowbit == 0 ? -1 : rank+lowbit;
    lowbit >>= 1;
  }

  // 存储两个子节点的标识符
  *d0 = down0;
  *d1 = down1; 

  return ncclSuccess; // 函数执行成功

}

参考资料

Massively Scale Your Deep Learning Training with NCCL 2.4 | NVIDIA Technical Blog