## 分布式深度学习中的分片计算的爱因斯坦求和 本文提倡使用爱因斯坦求和(einsum)来简化对分布式张量计算中分片策略的推理,尤其是在像DTensor这样的框架中。虽然传统上它被认为适用于基本算术,但理解分片对于高效的分布式矩阵运算至关重要。 Einsum提供了一种简洁的方式来表示线性代数,消除了不同矩阵乘法函数之间的混淆。重要的是,它也简化了通过反向模式自动微分计算梯度——只需交换输入/输出索引集即可。 核心思想扩展到分片:一组规则决定了分片在einsum操作期间的行为。这些规则涵盖了复制、分片和广播维度,从而产生复制的输出、分片的输出或需要全归约操作的“待处理归约”。 本文通过张量并行(Megatron)和序列并行示例说明了这些规则,展示了einsum如何理清梯度计算并识别正确的分布式训练所需的通信步骤(如全归约)。使用einsum简化了分片逻辑,使其更容易理解和实现分布式计算。