1. 项目概述当张量计算遇上爱因斯坦求和在搞大规模机器学习或者高性能数值计算的朋友估计都绕不开一个头疼的问题怎么把一个大得吓人的张量运算高效地拆开扔到一堆CPU或者GPU上去跑这事儿说起来简单不就是“分而治之”嘛但真做起来里头全是坑。数据怎么切切完以后各个计算节点之间要传多少数据通信开销会不会把并行带来的收益全吃掉了更别提那些复杂的计算图比如大语言模型里一层套一层的矩阵乘法和注意力机制手动去设计一个全局最优的分解策略简直就像在迷宫里摸黑找路。传统的路子无非是数据并行把数据批次分到不同设备或者模型并行把模型参数分到不同设备。像Megatron-LM这类工作就是针对Transformer模型精心设计的手动模型并行方案效果很好但通用性差换个模型结构可能就得重来。我们需要一种更“聪明”的、能自动寻找最优分解方案的方法。这就引出了EinDecomp算法。它的核心洞察非常巧妙借用物理学里爱因斯坦求和约定Einstein summation notation这套声明式的语言来描述张量运算。比如一个矩阵乘法C[i,k] A[i,j] * B[j,k]用爱因斯坦求和来写就是ik ij, jk。这种写法剥离了具体的循环顺序和内存布局只关心张量维度之间的对应和缩并关系正好为我们抽象并行分解问题提供了完美的数学框架。EinDecomp算法正是基于此将整个计算过程表达为一个EinGraph爱因斯坦求和图其中的节点就是一个个用爱因斯坦求和描述的运算。然后它通过一套动态规划Dynamic Programming算法在这个图上自动搜索为每个运算节点分配合适的张量“分块”策略即TaskGraph最终目标是在给定的处理器数量下最小化节点间的数据搬运通信开销。实验表明这套自动化的方法在矩阵链计算、大模型训练和推理任务上不仅能达到、甚至常常能超越那些需要大量专家经验的手工定制并行方案如Megatron、ScaLAPACK的性能。简单来说EinDecomp想做的就是你只管用爱因斯坦求和声明你要算什么它来帮你搞定怎么分、怎么算最省时间。这对于需要频繁尝试不同模型架构或并行配置的研究者和工程师来说无疑是个强大的利器。2. 核心原理从EinGraph到TaskGraph的转化逻辑要理解EinDecomp得先搞清楚它面对的问题是什么以及它是如何将一个大问题拆解成可计算的子问题的。2.1 问题形式化边界向量与划分向量首先EinDecomp算法处理的输入是一个EinGraph。这是一个有向无环图DAG每个顶点代表一个张量运算用爱因斯坦求和表达式定义。每条边代表张量数据的流动。每个张量都有两个关键属性边界向量Bound Vector, b定义了张量每个维度的实际大小。例如一个形状为[8, 8]的矩阵其边界向量b [8, 8]。这是问题的固有约束不可改变。划分向量Partitioning Vector, d这是我们要求解的核心。它定义了为了并行化我们将每个维度切分成多少份。d的每个元素必须是2的幂次这是为了简化搜索空间后文会详述并且所有元素的乘积等于我们拥有的处理器或计算核心总数p。例如对于b [8, 8]的矩阵如果我们有8个处理器p8那么可能的划分向量d可以是[8,1]只切分第一维、[1,8]只切分第二维、[4,2]或[2,4]等。一个爱因斯坦求和运算如矩阵乘会涉及输入张量和输出张量。EinDecomp算法的任务就是为计算图中的每一个张量包括中间结果确定一个最优的划分向量d使得整个计算图的总通信成本最低。这个确定了所有划分向量的计算图就称为TaskGraph它可以直接映射到具体的并行执行任务。2.2 通信成本建模重划分开销是核心并行计算中时间主要花在计算和通信上。对于密集的张量运算计算时间相对固定由浮点运算量决定因此优化的重点就落在了通信上。在EinGraph中通信发生在当一个运算节点的输出张量需要作为下一个运算节点的输入时。如果这两个节点对同一个张量约定的划分向量d不同那么数据就必须在处理器之间重新分布这个过程称为重划分Repartitioning。重划分成本的计算是EinDecomp算法的基石。原文中的公式和例子图4及相关论述清晰地阐述了这一点。我们来拆解一下假设一个生产者Producer张量Z其边界向量为bZ [4, 2]划分向量为dZ [2, 1]。这意味着Z被垂直切成了2块因为dZ[0]2每块形状是[2, 2]bZ/dZ [4/2, 2/1] [2,2]。现在一个消费者Consumer运算需要Z作为输入但它期望的划分向量是dX [1, 2]即水平切成2块每块形状[4, 1]。这就产生了不匹配。生产者有2块数据消费者需要2块数据但每一块数据的形状和分布方式完全不同。为了满足消费者的需求生产者的每一块数据一个子张量可能都需要被拆分然后部分数据发送到多个消费者处理器上反之消费者的每一块数据也需要从多个生产者那里收集数据。原文推导出了一个计算这种重划分通信量以浮点数传输量为度量上界的通用公式。它考虑了n_p生产者每个子张量的大小∏ (bZ / dZ)。n_c消费者每个子张量的大小∏ (bX / dX)这里bX bZ因为是同一个张量。n_int一个生产者子张量对一个消费者子张量的贡献大小∏ min(bZ/dZ, bX/dX)。n张量的总大小∏ bZ。最终的成本函数cost_repart(dX, dZ, bZ)封装了这个计算过程。这个函数是可计算的、确定的。它为动态规划算法提供了比较不同划分方案优劣的“价格标签”。注意这里计算的是通信量的一个上界而非精确值。这是工程上的一个巧妙折衷。精确计算最优通信模式可能极其复杂依赖于具体的网络拓扑和调度。而上界计算相对简单、确定且为优化提供了一个可靠的目标——最小化这个上界通常就能在实际中获得很好的通信性能。2.3 动态规划求解在状态空间中寻找最优路径有了为单条边数据依赖定价的能力整个问题就变成了在一个巨大的状态空间中寻找最优路径的问题。每个顶点运算都有多个可能的输出划分向量dZ选择不同的选择会影响它作为后续运算输入时的重划分成本。EinDecomp采用动态规划来解决这个问题。它定义了一个查找表MM[v, d]表示计算从所有输入顶点到顶点v为止的子图并且要求v的输出张量采用划分向量d时所需的最小累计通信成本。动态规划的过程遵循计算图的拓扑顺序初始化对于没有输入的计算节点即原始输入数据M[v, d] 0。因为输入数据可以预先按任意我们希望的d划分好不产生计算图中的通信成本。递推对于一个运算顶点v假设是二元运算如矩阵乘它有左输入vX和右输入vY。为了计算M[v, dZ]我们需要枚举所有能产生输出划分dZ的、该运算内部可能的并行执行方案对应一个中间划分向量d。对于每一种(d)我们再枚举左输入所有可能的划分dX和右输入所有可能的划分dY。总成本是M[vX, dX] M[vY, dY] cost_repart(dX, d, b) cost_repart(dY, d, b) cost_join(...) cost_agg(...)其中cost_join和cost_agg是该运算内部如矩阵乘的乘加运算可能涉及的、跨处理器的计算成本通常远小于跨节点的重划分通信成本。算法选择使总成本最小的那个(d, dX, dY)组合并将最小成本存入M[v, dZ]。回溯当处理完所有顶点后查看最终输出顶点所有可能的d对应的M值最小的那个就是全局最优解的成本。然后从这个值开始反向回溯每一步做出的(d, dX, dY)选择就能为整个EinGraph中的每个张量确定最优的划分向量从而构造出TaskGraph。为什么可行关键在于一个爱因斯坦求和运算内部可能的并行划分方式d虽然多但并非无限。原文8.1节指出在限制d的每个元素为2的幂次、且总处理器数p也是2的幂次的前提下可能的划分数量是一个组合数C(ND-1, D-1)其中Nlog2(p)D是运算中涉及的不重复维度标签数。对于常见的运算和处理器规模如p1024,D6这个数量是完全可以暴力枚举的。这使得动态规划在每个节点上的状态转移是计算可行的。3. 算法实现细节与工程化挑战理解了原理我们来看看要把EinDecomp实现为一个可用的系统需要处理哪些棘手的工程问题以及原文中提到的解决方案。3.1 处理通用有向无环图线性化近似标准的动态规划算法有一个关键假设每个非输入顶点的输出最多只有一个消费者。这样子问题M[v, d]才具有最优子结构——到达v的最优方式完全由到达其前驱的最优方式决定。然而真实的计算图尤其是神经网络经常是多扇出的。一个层的输出可能同时作为后续多个层的输入。这就破坏了最优子结构。因为为v选择划分d时不仅要考虑它到下一个节点v1的成本还要考虑它到另一个消费者v2的成本。状态空间会爆炸式增长从M[v, d]变成M[v, d_for_v1, d_for_v2, ...]变得不可计算。EinDecomp采用了一种线性化Linearization的启发式方法来处理这种情况。其核心思想是忽略跨路径的相互影响每次只优化一条最长的路径。在EinGraph中找出从某个源点出发的最长路径按拓扑顺序。在这条路径上运行之前描述的动态规划算法忽略从这条路径分支出去或汇入这条路径的边所带来的重划分成本。也就是说在计算路径上节点v的成本时如果它的某个输入来自路径外部我们假设那个输入已经以“最方便v”的方式划分好了不计入成本。为这条路径上的所有节点确定划分方案。将这条路径上的节点和边从图中“标记”或“固定”然后在剩余的子图中重复步骤1-3寻找下一个最长路径进行优化。这个过程如图6所示。它显然不能保证找到全局最优解因为它人为地割裂了路径之间的耦合关系。但原文指出在像大语言模型这样的实际场景中这种近似方法带来的性能损失很小。这是因为计算图虽然复杂但主干路径如Transformer层的顺序执行的通信成本通常占主导地位优先优化主干路径能抓住主要矛盾。3.2 搜索空间剪枝与策略即使对单个运算可能的划分向量d的集合也可能很大。动态规划需要高效地枚举和评估这些可能性。幂次限制强制要求d的每个分量是2的幂次且总乘积为p2的幂次。这极大地减少了搜索空间使其从连续的组合优化变为离散的组合问题。从实践角度看计算机系统的并行度GPU SM数量、CPU核心数通常是2的幂次内存对齐也偏好2的幂次所以这个限制是合理且有益的。基于代价的剪枝在动态规划递推时对于顶点v的某个输出划分dZ我们可能需要枚举大量的(d, dX, dY)三元组。可以设置一些启发式规则进行早期剪枝。例如如果某个dX对应的M[vX, dX]成本已经非常高那么它参与构成的任何方案成本都不可能低可以提前跳过。记忆化与缓存cost_repart等函数会被反复计算。需要预先计算并缓存这些值避免重复计算带来的开销。3.3 与执行引擎的集成EinDecomp算法产生的是一个TaskGraph这是一个并行执行计划。要实际运行它需要一个强大的底层执行引擎。原文中的Einsummable系统就是这样一个载体。CPU集群将TaskGraph中的每个“任务”编译成调用高度优化数学库如Intel MKL的批处理矩阵乘的内核。同时需要处理张量的“打包”Packing和“解包”Unpacking以匹配批处理库期望的内存布局。通信层使用UCX这样的高性能网络库。GPU服务器将TaskGraph编译成在GPU上执行的核函数。这里可以利用NVIDIA的cuTENSOR库它专门为张量收缩运算提供了优化实现。多GPU间的通信则通过NVLink或PCIe总线进行。执行引擎如文中提到的Turnip还需要负责任务的调度、数据的放置、通信的发起与同步确保TaskGraph描述的计算和数据流得以正确、高效地执行。4. 实验评估与性能分析解读原文通过四组实验全面评估了EinDecomp算法的有效性。我们来深入解读一下这些实验背后的含义和给我们带来的启示。4.1 实验一矩阵链运算——与经典方案的较量实验内容计算(A × B) (C × (D × E))这样的矩阵链。对比了EinsummableEinDecomp与以下方案SQRT一种简单的启发式方法将n个处理器分配到一个矩阵上时在行和列方向各分配√n份。对于方阵这等价于经典的3D矩阵乘法算法。ScaLAPACK分布式内存线性代数的标杆库。DASK基于任务的Python并行计算库。结果分析图7图8对非方阵的适应性在矩阵尺寸不均匀的情况下EinDecomp显著优于SQRT约有2倍性能优势。这是因为SQRT的固定√n × √n划分对于瘦高或扁平的矩阵非常低效而EinDecomp能自动调整每个维度的切分比例找到通信量更小的划分。超越经典库EinDecomp在CPU和GPU上均大幅领先ScaLAPACK和DASK。这不仅仅是分解算法的胜利更是声明式编程自动化优化定制化内核整个技术栈的胜利。ScaLAPACK需要用户手动进行网格划分和数据分布DASK则存在Python层任务调度的开销。而EinDecomp生成的TaskGraph可以直接编译成高效的低级内核并精准控制通信。一个有趣的现象即使在方阵情况下CPU上的EinDecomp也比SQRT略好。作者推测这可能是因为EinDecomp能更好地利用数据的预放置Pre-placement。在动态规划中它可以为输入矩阵选择一种划分使得在计算开始前数据就已经分布在最“合适”的位置减少了计算过程中的数据迁移。4.2 实验二大规模前馈神经网络训练——对阵数据并行实验内容训练一个超大规模分类器输入特征约60万隐藏层8192神经元输出类别近1.5万。对比EinsummableEinDecomp与PyTorch的数据并行。结果分析图9数据并行的劣势PyTorch数据并行在这里表现极差甚至不如单GPU。原因在于这是一个典型的“大模型、小批量”场景。模型参数巨大在每一步训练中数据并行需要将整个模型参数广播到所有GPU通信开销巨大完全抵消了多GPU带来的计算收益。EinDecomp的优势EinDecomp自动选择了模型并行与数据并行混合的策略。它可能将庞大的权重矩阵按行或列切分到不同GPU上模型并行同时对于可以独立计算的维度如批量维度仍采用数据并行。这种混合策略是手动设计极其困难的但却是EinDecomp动态规划搜索的自然结果从而实现了显著的性能提升。4.3 实验三大语言模型推理——自动化与手工设计的对决实验内容在8个GPU上运行LLaMA 7B模型的“首令牌生成”Prefill阶段。对比EinDecomp与三种手工设计的并行方案Megatron经典的张量并行模型并行方案将权重矩阵按列切分。Sequence序列并行将输入序列切分到不同GPU。Attention将注意力头分组到不同GPU。结果分析图10EinDecomp的竞争力EinDecomp在大多数配置下都匹配或超越了手工方案。这证明了其自动化搜索的有效性。它没有预设任何并行模式而是根据具体的模型结构、张量尺寸和硬件数量动态地寻找最优组合。序列并行的意外之喜“Sequence”方案表现不俗甚至经常优于Megatron。这揭示了在长序列推理场景下沿序列维度切分是一种通信效率很高的策略。EinDecomp的搜索空间包含了这种可能性并能根据实际情况将其与其他并行方式结合。自动化搜索的价值这个实验最有力的结论是对于复杂的计算图不存在一个“放之四海而皆准”的最优并行策略。最优策略依赖于具体的模型规模、输入尺寸和硬件配置。手工设计如Megatron是静态的、经验性的而EinDecomp是动态的、自适应的。4.4 实验四内存受限推理——系统级优势实验内容在有限GPU内存下运行超大模型LLaMA 65B的推理。对比支持CPU Offload的Einsummable与同样支持此功能的PyTorch生态工具ZeRO-Inference和FlexGen。结果分析图11超越手工优化系统EinsummableEinDecomp大幅领先ZeRO和FlexGen。这不仅仅是分解算法的胜利更是Turnip执行引擎与EinDecomp协同设计的胜利。Turnip的“非确定性”调度和高效的内存分页机制与EinDecomp产生的精细并行计划相结合实现了极致的内存利用和计算通信重叠。启示一个高效的自动并行系统需要算法如EinDecomp和运行时如Turnip的紧密耦合。算法产生一个理论上通信高效的计划运行时负责以最小的开销将其执行出来并处理内存瓶颈等实际问题。5. 实操启示与未来展望基于对EinDecomp原理和实验的深入分析我们可以提炼出一些对实际开发和研究的启示。5.1 何时考虑采用类似EinDecomp的方案计算图复杂多变如果你的应用涉及复杂的、非线性的张量计算图且图结构或张量尺寸经常变化手动设计并行策略成本过高。追求极致性能在硬件规模较大如数十上百个GPU或网络带宽成为瓶颈时通信开销的细微差别会被放大自动优化工具的价值凸显。研发原型阶段在探索新模型架构时使用声明式接口如爱因斯坦求和描述计算并依赖自动化工具寻找并行策略可以极大加速迭代周期让研究者更专注于算法本身。5.2 潜在挑战与注意事项规划时间开销动态规划搜索虽然对单个运算可行但对于超大规模计算图规划阶段本身可能耗时。需要考虑缓存规划结果、增量更新或采用更快的启发式算法。对不规则运算的支持爱因斯坦求和主要针对规整的乘加类收缩运算。对于带有复杂逐元素操作、条件判断或动态控制流的计算表达和优化会变得困难。硬件特性建模当前的通信成本模型相对抽象。要获得更精确的优化需要纳入更详细的硬件模型如GPU间NVLink与PCIe的带宽差异、计算单元间的层次化通信成本等。与现有生态集成如何将EinDecomp这样的算法无缝集成到PyTorch、TensorFlow等主流框架中让用户以最小的迁移成本获益是一个重要的工程问题。5.3 未来发展方向分层优化结合算子内并行Intra-operator EinDecomp所擅长的与算子间并行Inter-operator 如流水线并行。Alpa等工作在这方面已有探索与EinDecomp的思想结合有望实现更全面的自动化。代价模型的持续学习可以利用实际运行的性能数据反过来训练或校准通信成本模型使其更贴近真实硬件行为形成闭环优化。扩展到更广泛的领域爱因斯坦求和的抽象能力很强其思想可以扩展到科学计算、数据库的张量查询优化等领域实现跨领域的自动并行化。EinDecomp算法向我们展示了一条通往“编译期自动并行化”的清晰路径。它通过一个优雅的声明式抽象爱因斯坦求和将一个复杂的系统优化问题转化为了一个可形式化、可计算的组合优化问题。虽然目前仍有其局限但其核心思想——用声明性语言描述计算意图让系统自动探索巨大的实现空间——无疑是未来大规模分布式机器学习系统发展的一个重要方向。对于开发者而言理解其原理不仅能帮助我们更好地使用这类工具更能启发我们在设计自己的系统时思考如何将复杂性封装在编译器或运行时之内为用户提供更简洁、更强大的抽象。