跳转至

DeepSpeed Communication Compression and Hiding

导言

通信优化只有两条基本路线：少传，或让传输不再暴露在关键路径。1-bit Adam/LAMB 与 0/1 Adam 属于前者，Domino 属于后者。前者改变数值算法，后者改变调度；两者的正确性风险和验收方式完全不同。

DeepSpeed 通信优化的两条路线

自绘示意图：1-bit/0-1 系列减少通信量，Domino 通过切片重排隐藏暴露时延。

基线：Adam 为什么难压缩

Adam 对梯度 \(g_t\) 维护一阶矩与二阶矩：

\[ m_t=\beta_1m_{t-1}+(1-\beta_1)g_t,\qquad v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2. \]

数据并行通常先聚合梯度，再更新这些状态。直接把 \(g_t\) 压到 1 bit 会把量化误差送入 \(v_t\)，使预条件器长期积累偏差。因此 1-bit 系列不是把 AllReduce 的 dtype 改掉这么简单。

1-bit Adam：预热、冻结方差、误差反馈

机制

1-bit Adam 先运行普通 Adam 预热，让 \(v_t\) 接近稳定；进入压缩阶段后冻结方差，通信一阶方向的 1-bit 量化值，并把本轮未表达的残差累积到下一轮。1

if step < freeze_step:
    g = allreduce_fp16(local_grad)
    m, v = adam_moments(g, m, v)
else:
    compensated = local_grad + error
    q, scale = one_bit_quantize(compensated)
    q_global = compressed_allreduce(q, scale)
    g_hat = dequantize(q_global)
    error = compensated - g_hat
    m = beta1 * m + (1 - beta1) * g_hat
    # v remains frozen in the compressed phase
param -= lr * m / (sqrt(v_frozen) + eps)

误差反馈保证被量化丢掉的信息不是永久消失；方差冻结避免 1-bit 噪声持续污染二阶矩。但这也意味着模型必须容忍优化器在 freeze_step 发生阶段切换。

官方教程标题引用最高 5× 通信下降和 3.4× 加速。它们来自特定模型、网络和 scale，不等于任意训练的默认收益。

1-bit LAMB：还要保护逐层 trust ratio

LAMB 在 Adam 方向外增加逐层缩放：

\[ r_l=\frac{\lVert w_l\rVert}{\lVert u_l\rVert},\qquad w_l\leftarrow w_l-\eta r_lu_l. \]

大 batch 训练依赖这个 trust ratio。若每个 Rank 用不一致的压缩方向计算 \(r_l\)，层级更新会偏离。因此 1-bit LAMB 除了预热、压缩和误差反馈，还要让压缩通信与逐层缩放契约一致。2

local grad + residual
  -> 1-bit compressed synchronization
  -> reconstructed global direction
  -> layer norm / update norm
  -> common trust ratio
  -> parameter update

它适合 LAMB 已经证明对目标大 batch 有效、且梯度通信确实暴露的训练。若原任务根本不需要 LAMB，换优化器只为使用 1-bit 通信会扩大收敛风险。

0/1 Adam：减少“长期冻结方差”的代价

0/1 Adam 的目标是让方差状态能够周期性更新，同时保留通信高效阶段；它还引入局部更新来减少同步频率。3

可以把流程理解为两种通信粒度交替：

for step in training:
    if variance_refresh_step(step):
        g = full_precision_sync(local_grad)
        m, v = refresh_adam_state(g, m, v)
    else:
        g_hat = one_bit_sync_with_error_feedback(local_grad)
        m = update_first_moment(g_hat, m)
        # local steps may proceed before the next global synchronization
    param = adam_update(param, m, v)

名称中的 0/1 不是“0 bit”。它强调在近零额外通信的局部更新阶段与 1-bit 同步阶段之间组织 Adam。相对 1-bit Adam，它降低方差长期冻结的限制，但多了刷新周期、局部步数和一致性调参。

Domino：不减少字节，切片后隐藏通信

Domino 面向另一个问题：collective 字节可能无法再减少，但通信在 layer 边界串行等待。它把 Tensor Parallel 和 Data Parallel 相关计算切成更细的 tile，用后续 tile 的 GEMM 覆盖前一 tile 的通信。4

时间线

baseline:
GEMM(full) ------> TP collective ------> next GEMM

Domino:
GEMM(tile 0) -> collective(tile 0) ----------->
               GEMM(tile 1) -> collective(tile 1) --->
                              GEMM(tile 2) -> ...

机制完整性取决于三点：

  • 切片后数学结果与未切分图一致；
  • collective 能异步推进，且没有过早同步；
  • tile 足够大，GEMM 时间能覆盖通信，同时又不产生过多 launch 和 bubble。

论文在 DGX-H100 等实验中报告最高约 1.3×。当 batch/sequence 太小、网络已很快、切片导致小 GEMM 效率下降，Domino 可能没有收益。

1-bit 优化器状态机与 Domino 重叠时间线

自绘示意图：压缩状态机和重叠时间线需要不同的正确性与性能验收。

验收指标不能混用

方法 首要指标 正确性指标 常见失败
1-bit Adam 通信字节、compressed phase step time loss/下游指标跨阶段连续性 freeze 太早、error buffer 不稳定
1-bit LAMB 大 batch 扩展效率 每层 trust ratio 与收敛 压缩扰动放大层级更新
0/1 Adam 同步频率、通信暴露 refresh/local-step 敏感性 局部漂移或刷新成本吞掉收益
Domino exposed collective time 参数/梯度与基线一致 小 GEMM、隐式同步、流水 bubble

结论

如果 profiler 显示瓶颈是通信字节，评估 0/1 系列；如果字节不可避免但 collective 暴露，评估 Domino。压缩方法必须做训练曲线与最终质量复验，隐藏方法必须做 trace 级关键路径复验。只看 NCCL 总时间或只看 headline speedup，都会选错路线。

  1. 1-bit Adam tutorial；教程文件最早提交于 2020-09-09。论文：arXiv:2102.02888。 ↩

  2. 1-bit LAMB tutorial；教程文件最早提交于 2021-04-20。论文：arXiv:2104.06069。 ↩

  3. 0/1 Adam tutorial；教程文件最早提交于 2022-03-10。论文：arXiv:2202.06009。 ↩

  4. DeepSpeed Domino tutorial；教程文件最早提交于 2024-12-13。论文：arXiv:2409.15241。 ↩

评论