在分布式深度学习系统中，张量分片策略的正确性直接决定了模型并行训练的效率与稳定性。当我们在多设备间拆分大型张量时，需要准确预判每个算子前向传播后的分片状态，以及反向传播中梯度的分布情况。传统方法依赖于绘制矩阵示意图并手工追踪切分边界，这种方式在复杂计算图中极易出错且效率低下。本文将介绍一种基于 Einstein 求和（einsum）的系统化分片计算方法，该方法能够像心算乘法表一样快速准确地推导出任意算子的分片结果，并已在 PyTorch DTensor 中得到工程验证。
Einsum 表示法与分片语义的统一
Einstein 求和是一种紧凑的多维线性代数表示方法，能够统一表达矩阵乘法、逐元素运算、点积等常见张量操作。其核心思想是通过索引字母标注张量维度，隐式定义求和规则，从而避免记忆多种算子语义的差异。例如，传统数学中的矩阵乘法可表示为 mm(x: f32[A, B], y: f32[B, C]) -> f32[A, C]，而在 einsum 中只需写作 torch.einsum("ij,jk->ik", x, y)，输入与输出形状一目了然。这种表示法之所以适合分片推理，根源在于其对维度语义的显式标注：每个字母对应一个具体的物理维度，索引的出现与消失直接映射到分片状态的变化。
在 einsum 术语中，出现于输入张量但消失于输出张量的索引称为「收缩维度」（contraction dimension），这类维度在计算过程中会发生跨设备的求和操作。保留在输出中的索引称为「自由维度」（free dimension），若自由维度同时存在于所有输入张量，则其为「批次维度」，若仅存在于部分输入张量，则其余张量需通过广播机制补全该维度。理解这对概念是掌握 einsum 分片规则的关键，它将张量运算的语义与分片行为建立了明确的对应关系。
Einsum 分片规则的系统化推导
在分布式张量框架中，分片规则回答的核心问题是：在何种条件下，多个分片张量可以通过本地计算直接得到正确输出的分片，而无需额外的跨设备通信。Einsum 分片规则可归纳为四条基本准则，适用于任意 einsum 公式在特定网格维度上的分片配置。以公式 "abi,aoi->abo" 为运行示例，可推导出以下分片放置规则。
第一条规则是复制张量的组合：当所有输入张量在目标维度上均为复制（Replicate）状态时，输出张量同样保持复制。这意味着在该维度上未进行任何切分，计算可在本地完成而无需同步。第二条规则针对批次维度的分片：当批次维度被分片（Shard）时，输出的批次维度也自动分片，这是数据并行的基础 —— 各设备处理不同的批次样本，本地计算后输出保持分片状态。
第三条规则涉及自由维度的分片：当自由维度被分片时，输出自由维度同样分片，但任何需要广播的输入张量在该维度上必须保持复制。例如，在公式 "sbh,h->sbh" 中，若序列维度 s 被分片而隐藏层权重 h 被复制，则输出序列维度保持分片，权重无需额外通信。第四条规则是收缩维度的分片处理：当收缩维度被分片时，输出将产生部分聚合（Partial）状态，意味着本地计算仅完成部分求和，需后续通过 all-reduce 操作完成完整聚合。这四条规则构成了任意 einsum 运算分片行为分析的基础框架。
张量并行中的 Einsum 分片实例分析
以 Megatron-LM 风格的张量并行（Tensor Parallelism, TP）为例，分析 ColumnParallelLinear 层的分片策略及其梯度流动。假设输入形状为 [sequence, batch, in_features]，权重形状为 [in_features, out_features]，输出形状为 [sequence, batch, out_features]。使用 einsum 表示为 "sbi,io->sbo"，其中 s 为序列维度，b 为批次维度，i 为输入特征，o 为输出特征。
在 TP 网格维度上，典型的分片配置为：输入张量保持复制，权重张量沿输出特征维度分片，输出张量同样沿输出特征维度分片。此时前向传播可本地完成：每个设备持有权重的一个分片，与完整的输入相乘后直接产出输出的对应分片，无需跨设备通信。
然而在反向传播中，梯度流动揭示了潜在的通信需求。已知 grad_output 沿输出特征维度分片，梯度计算公式为 grad_input = torch.einsum("sbo,io->sbi", grad_output, weight) 和 grad_weight = torch.einsum("sbi,sbo->io", input, grad_output)。应用分片规则分析可知：grad_input 在输出特征维度产生 Partial 状态，因为该维度是收缩维度且被分片；而 grad_weight 在输出特征维度保持 Shard 状态，因为该维度是自由维度。这意味着在后续计算中，若下游算子期望接收复制状态的 grad_input，框架必须触发 all-reduce 操作将部分聚合结果归约为完整梯度，这正是 Megatron-LM 中 _CopyToModelParallelRegion 手动触发 all-reduce 的数学本质。
序列并行与复制缩放因子的分片平衡
序列并行（Sequence Parallel, SP）是另一种重要的分片策略，它将序列维度切分以节省激活内存，但保持权重参数复制。考虑一个带有可学习缩放因子的序列并行层：输入形状为 [sequence, batch, hidden]，缩放因子形状为 [hidden]，输出形状相同。Einsum 表示为 "sbh,h->sbh"。
在 SP 网格维度上的分片配置为：输入沿序列维度分片，缩放因子保持复制，输出沿序列维度分片。前向传播同样可本地完成，各设备处理序列的不同片段。本地计算后输出保持分片，无需额外通信。
反向传播的梯度计算公式为 grad_input = torch.einsum("sbh,h->sbh", grad_output, weight) 和 grad_weight = torch.einsum("sbh,sbh->h", input, grad_output)。分析分片规则可知：grad_input 在序列维度保持 Shard 状态，因为该维度是自由维度；而 grad_weight 在序列维度产生 Partial 状态，因为该维度是收缩维度且被分片。因此，必须对 grad_weight 执行 all-reduce 操作，才能得到完整的复制权重梯度。这一实例与张量并行的 ColumnParallelLinear 形成对照：虽然数学结构相似，但分片维度的角色互换导致了不同的通信模式与负载分布。
工程实践中的负载均衡与参数配置
在工程实践中应用 einsum 分片策略时，需要关注若干关键参数与配置选项，以确保计算效率与内存使用的均衡。首要考量是分片维度的选择：当张量某一维度较大且计算可并行化时，沿该维度分片通常能获得更好的负载均衡效果。但需注意收缩维度的分片将引入额外的 all-reduce 通信开销，因此对于计算密集型算子（如矩阵乘法），将收缩维度保持复制而分片自由维度往往是更优选择。
PyTorch DTensor 提供了细粒度的分片控制 API。典型配置模式如下：对于嵌入层或输出层，可沿词汇维度或输出维度分片以支持大规模词表处理；对于 Transformer 的注意力头，可沿头维度分片以配合张量并行策略；对于激活值较大的中间层，可沿序列或批次维度分片以节省峰值内存。建议在分片配置前使用 einsum 公式预演梯度流动，识别潜在的 Partial 状态并规划相应的通信算子时序。
监控指标方面，应重点关注各设备的计算耗时与通信耗时比值。当该比值过低（如低于 3:1）时，可能存在分片不均衡或通信过于频繁的问题。此外，需监控 all-reduce 操作的数据量与频率，它们直接影响分布式训练的可扩展性。在极端情况下，可考虑引入梯度累积来降低通信频率，但需权衡内存节省与训练吞吐量的取舍。
总结与工程建议
Einsum 表示法为分布式张量分片提供了一种系统化、可编程的推理框架。通过将张量运算映射为索引字母的收缩与保留规则，工程师可以像心算乘除法一样快速推导出任意算子的分片状态与梯度分布。这种方法已在 PyTorch DTensor 和 JAX pjit 等分布式框架的内部实现中得到验证，其核心思想可概括为：首先使用 einsum 明确标注所有输入输出的维度语义，然后应用四条基本分片规则逐条推导，最后根据 Partial 状态规划必要的通信算子。
在实际工程中，建议将 einsum 分片分析纳入模型并行策略的设计阶段，而非事后调试。预先推导可以避免运行时出现意外通信导致的性能下降，也有助于在数据并行与模型并行之间做出更明智的决策。对于复杂的多维张量运算，尝试将其拆解为多个 einsum 子公式并分别分析，是管理推理复杂度的有效策略。
资料来源：本文核心内容参考自 Edward Yang 的技术博客《Computing sharding with einsum》（2026 年 1 月），该文系统阐述了 einsum 分片规则的数学推导与工程应用。