ZeRO-3 每个微批量都会应用 Muon 的 Newton-Schulz 算法,因此梯度累积会改变优化器。

为什么

stage3.py:

    def _apply_distributed_muon_update(self, communication_data_type, buffer_to_reduce):
        if not self.use_muon:
            return
        ...

和它唯一的调用位置是在 IPG 存储减少路径中：

        dist.all_reduce(buffer_to_reduce, group=process_group)
        ...
        self._apply_distributed_muon_update(communication_data_type, buffer_to_reduce)
        for param in params_in_bucket:
            grad = param.grad

该路径在每个微批中都运行。没有累积边界保护，并且 use_muon 是唯一的条件。ZeRO-1/2 在 get_flat_partition 中执行相同的工作，ipg_epilogue 只在 if self.is_gradient_accumulation_boundary() 条件下调用它 — 这就是为什么这两个阶段的结果正确的原因。后果按其重要性的顺序排列：

1. 动量衰减错误。 momentum.lerp_(grad, 1 - beta) 在每个优化器步骤中运行 n 次，因此有效保留率为 beta**n，而不是 beta。在 beta=0.95 和 gas=4 时，它为 0.81，而在 gas=16 时为 0.44 — 用户配置的动量不是他们得到的动量，而且差异与相关的参数无关。
2. 正交化算法看到部分梯度。 Newton-Schulz 不是线性的，因此正交化的微批梯度之和不是它们之和的正交化。由于 Muon 的更新是尺度不变的，每个微批都贡献一个单位尺度的更新，无论它看到的样本数量如何，这使得噪声较大的微批与其余微批的权重相等。
3. **n 次 Newton-Schulz 运算…