本文主要介绍了 AllReduce 和 Bucketing 分别是什么，和他们之间的联系。

一、AllReduce 是什么？

AllReduce 是分布式训练中的一种集体通信操作，
用于在多个 GPU（worker）之间同步张量（通常是梯度）。

典型流程如下：

  1. 每个 GPU 独立计算自己的梯度张量（如 grad）。
  2. 所有 GPU 通过 AllReduce 操作，将各自的张量求和/平均，获得全局一致的梯度。
  3. 每个 GPU 使用这个同步后的梯度更新模型参数。

AllReduce 是数据并行训练中实现模型同步的关键机制。

二、为什么 AllReduce 会成为性能瓶颈？

  • 模型中参数众多，梯度张量数量也很多。
  • 每个张量如果单独 AllReduce，通信次数极多。
  • 小张量通信无法充分利用带宽，且频繁启动通信带来显著延迟（latency）。

三、Bucketing 是什么？

Bucketing 是一种优化 AllReduce 通信效率的策略，
将多个小张量合并成一个大 “bucket”，再一次性执行 AllReduce。

核心思想：Batch Small Reduces into One Large Reduce

举例：

  • 原始做法：
    • grad1 → AllReduce
    • grad2 → AllReduce
    • grad3 → AllReduce
  • Bucketing 后：
    • [grad1 + grad2 + grad3] → 合并成 bucket → 一次 AllReduce

四、Bucketing 的优势

优势 说明
降低通信启动次数 小张量太多会频繁触发 AllReduce，带来启动延迟
提高带宽利用率 大张量通信更接近理论带宽上限，传输更高效
便于通信计算重叠 bucket 可以在张量准备好后提前通信，提升整体吞吐

五、在 PyTorch 等系统中的实际应用

PyTorch DDP（DistributedDataParallel）:

  • 默认启用 bucketing 策略。
  • 使用参数 bucket_cap_mb 控制 bucket 的大小（如 25MB）。
  • 当 bucket 填满或梯度 ready，即可触发一次 AllReduce。

DeepSpeed, Megatron 等:

  • 通常会设计多级 bucket，例如按照张量类型（权重/偏置）或位置（layer-wise）划分。
  • 有时还结合流式调度（streaming）和异步通信。

六、Bucketing 与 MoE 的关系

  • 在 MoE 训练中，专家参数（expert weights）通常不需要 AllReduce（它们是局部的）。
  • 但非专家参数（如 Attention, LayerNorm 等）仍然共享，需要 AllReduce 同步。
  • Bucketing 同样可用于这些参数的通信优化。

七、小结

  • AllReduce 是实现数据并行训练中参数同步的核心通信操作。
  • Bucketing 是一种提升 AllReduce 效率的关键优化策略，特别适用于张量数量多、小张量多的场景。
  • 二者关系是协作：Bucketing 让 AllReduce 更高效，AllReduce 是 Bucketing 的目标操作。