1. 项目概述当张量计算遇上动态规划在分布式机器学习和大规模科学计算的战场上我们每天都在和“大”作斗争大模型、大数据、大矩阵。一个直观的解法是把计算任务拆开扔到不同的处理器CPU核心或GPU上并行处理这就是张量并行化。但问题来了怎么拆拆成什么形状数据在处理器之间怎么流动一个糟糕的拆分方案可能让处理器们花费在“聊天”通信上的时间远超过真正“干活”计算的时间最终效率不升反降。过去这活儿得靠资深系统工程师手搓。他们像经验丰富的裁缝针对特定的计算图比如Transformer模型的前向传播精心设计一套拆分策略比如Megatron-LM的模型并行。但这不仅费时费力而且换一个模型或计算模式策略可能就失效了。我们迫切需要一种“自动裁缝”——给定任意由爱因斯坦求和EinSum表达式构成的计算图它能自动找出通信开销最小的并行分解方案。这就是EinDecomp算法要解决的核心问题。它不是一个具体的系统实现而是一套基于动态规划的自动并行化分解技术。其输入是一个用EinSum语言描述的计算图EinGraph输出则是一个详细规定了每个操作如何拆分、数据如何重分布的TaskGraph。最妙的是它把寻找最优拆分方案的过程形式化成了一个可以通过动态规划高效求解的优化问题。下面我们就来拆解这套“自动裁缝”的剪刀和尺子。2. 核心思路将并行化转化为图上的最优路径搜索理解EinDecomp首先要跳出“如何并行执行一个操作”的细节上升到计算图的全局视角。它的核心思路可以概括为将整个计算图的自动并行化转化为一个在状态空间中进行最优路径搜索的问题并用动态规划来高效求解。2.1 从EinGraph到TaskGraph问题的形式化想象你有一个计算过程比如C A B矩阵乘然后D C E。用EinSum表示就是ik, kj - ij和ij, ij - ij。多个这样的操作按照数据依赖关系连接起来就形成了一个有向无环图DAG这就是EinGraph。图中的每个顶点Vertex代表一个EinSum操作边代表张量数据的流动。我们的目标是生成一个TaskGraph。TaskGraph在EinGraph的基础上为每个顶点操作的输出张量增加了一个关键的属性划分向量Partitioning Vector记为d。这个向量定义了该张量在每个维度上被切分成了多少份。例如对于一个形状为[8, 8]的矩阵如果d [2, 4]就意味着它在行维度上被切成2份列维度上被切成4份总共分布在2 * 4 8个处理器上。为什么是划分向量这是将并行度具体化的关键。d的每个分量通常被约束为2的幂次与处理器数量p为2的幂次对齐所有分量的乘积就等于该操作所激发的并行任务数内核调用数。确定了每个顶点的输出张量的d整个计算图的并行执行方案就确定了。那么如何为每个顶点选择最优的d呢这取决于一个成本模型Cost Model其目标是最小化整个计算图执行过程中的总通信开销。通信主要发生在两种情况下数据重分布Repartitioning当一个操作的输出作为下一个操作的输入时如果两者的划分向量d不匹配就需要在处理器间重新组织数据。操作内部的通信Join/Aggregation Cost某些EinSum操作如矩阵乘涉及的内积本身在分布式执行时也可能需要跨处理器的数据交换。EinDecomp将寻找最优d赋值方案的问题建模为在EinGraph上寻找最小成本路径的动态规划问题。2.2 动态规划的状态与子问题动态规划的精髓在于定义状态和子问题的最优解。EinDecomp定义了一个关键的状态函数M[v, dZ]它表示对于计算图中从所有输入顶点到顶点v的这部分子图在强制要求顶点v的输出张量采用划分向量dZ的前提下所能达到的最小累计通信成本。这个定义是算法的基石。它意味着我们不是为整个图一次性找一个最优解而是先解决所有“前缀子图”在特定约束下的最优解。最终整个图的最优解就是M[输出顶点, 某个dZ]中的最小值。那么如何递推地计算M[v, dZ]呢考虑一个顶点v它代表一个二元EinSum操作如矩阵乘有两个输入顶点vX和vY。要计算M[v, dZ]我们需要考虑所有能产生输出划分dZ的可能操作执行方式并综合其输入子图的最优成本。递推公式的核心思想如下M[v, dZ] min over all possible d (执行v操作本身的划分方案) { min over all possible dX (vX的输出划分) { min over all possible dY (vY的输出划分) { M[vX, dX] // 计算左子图的最小成本 M[vY, dY] // 计算右子图的最小成本 cost_repart(dX - d, b) // 将左输入重分布到操作v所需格式的成本 cost_repart(dY - d, b) // 将右输入重分布到操作v所需格式的成本 cost_join(d, b) // 操作v内部执行所需的通信成本如归约 cost_agg(d, b) // 操作v可能产生的聚合成本 } } }这里d是顶点v这个操作内部的并行划分方案它决定了操作如何被拆分成并行任务。而dZ是d经过操作语义如矩阵乘消去中间维度后得到的输出张量的划分。b是张量的边界形状向量。这个递推关系清晰地展示了动态规划的过程要计算顶点v在某种输出划分下的最优成本需要枚举它所有可能的内部执行方式d并为每一种方式组合其所有输入顶点在各种可能划分下的最优成本再加上数据重分布和操作本身带来的通信成本。2.3 拓扑排序与回溯有了状态定义和递推公式算法就可以按照计算图的拓扑顺序依次计算每个顶点的M[v, *]。初始化对于没有输入即原始输入数据的顶点其M[v, d]通常设为0或一个很小的常数因为输入数据可以被预先放置好不产生计算成本。递推按照拓扑序对于每个顶点v枚举其所有可能的输出划分dZ利用上述公式计算M[v, dZ]。在计算过程中需要查询其前驱顶点vX,vY的所有M[vX, dX]和M[vY, dY]值这些值已经在之前的步骤中计算好了。回溯当计算到最终的输出顶点时选择使M[输出顶点, dZ]最小的那个dZ。然后从这个状态开始反向回溯根据递推计算时记录下的“选择”即每一步取最小值时对应的d,dX,dY就可以还原出为图中每个顶点选择的划分向量从而构造出完整的TaskGraph。这个过程就像是在一个由“顶点”和“划分方案”构成的多层状态空间中寻找一条从起点输入到终点输出的成本最低路径。动态规划确保了这条路径是全局最优的。3. 成本模型详解如何量化通信开销动态规划离不开一个准确评估每个选择代价的成本函数。EinDecomp的成本模型核心是量化**数据重分布Repartitioning**的成本。这是通信开销的大头理解它才能理解算法的优化目标。3.1 一个直观的例子数据切片与搬运假设我们有一个生产者Producer操作产生一个形状为[4, 4]的张量Z其划分向量为dZ [2, 2]。这意味着Z被均匀地切成了4块2行×2列分布在4个处理器上。每个子张量包含(4/2) * (4/2) 4个元素。现在有一个消费者Consumer操作需要Z作为输入但它期望的输入划分是dX [1, 4]即整张张量在行维度不切分在列维度切成4份。这就产生了不匹配需要重分布。重分布发生了什么生产者输出的每个子张量例如左上角的2x2块其数据需要被拆分并发送到多个消费者处理器。具体来说生产者的一个子张量其数据可能对应消费者多个子张量的一部分。为了在消费者端组装出一个完整的、符合dX划分的子张量可能需要从多个生产者子张量那里收集数据。论文中给出了一个更复杂的例子对应原文Figure 2, 4并推导了重分布成本的上界公式。其核心思想是计算在重分布过程中每个浮点数平均需要被传输多少次。这个次数乘以张量的总元素数就得到了总通信量的上界。3.2 重分布成本公式解析设生产者张量边界为bZ划分向量为dZ消费者张量边界为bX通常bX bZ期望划分向量为dX。定义几个关键量n_p每个生产者子张量的元素数 Π(bZ / dZ)。Π表示连乘。n_c每个消费者子张量的元素数 Π(bX / dX)。n_int一个生产者子张量对一个消费者子张量的贡献元素数 Π(min(bZ/dZ, bX/dX))。这里min是按元素取最小值。n张量的总元素数 Π(bZ)。重分布的成本上界cost_repart(dX, dZ, bZ)可以计算为(n_c / n_int - 1) * (n / n_c) * (n_c n_p)如果n_p ! n_int还需要加上一项n_p * (n / n_c)。这个公式怎么理解我们可以把它看作一个“组装”过程。消费者端的每个处理器最终要持有一个完整的、大小为n_c的子张量。为了组装它它需要从多个生产者处理器那里收集数据块。n_int可以看作每次从某个生产者那里“搬来”的数据块大小。(n_c / n_int - 1)大致描述了平均需要“搬”多少次减去第一次。(n / n_c)是消费者子张量的总数。(n_c n_p)项则涵盖了数据从生产者发出和消费者接收的双向通信开销的一个上界估计。额外的项处理了生产者子张量第一次被发送时的成本。实操心得成本模型的简化与实用性这个成本模型是一个上界而非精确值。它假设了最坏情况下的通信模式比如一个生产者子张量可能需要发送给所有消费者处理器。在实际系统中通信可以优化例如只发送需要的部分。但作为优化目标最小化这个上界是有效的因为它惩罚那些导致数据被过度拆分、需要大量广播的划分方案。在实际实现时可以根据硬件特性如网络拓扑、带宽、延迟对这个成本模型进行加权或调整。3.3 操作内部的通信成本除了重分布某些EinSum操作本身在并行执行时也有通信。例如在分布式矩阵乘法C[i,k] A[i,j] * B[j,k]中如果按i和k维度划分A和B那么计算C的某个块时需要收集所有j维度上的A和B的对应块并进行归约求和。这引入了cost_join。cost_join的估算与具体的并行算法有关。对于常见的归约类操作其成本通常与需要通信的数据量成正比而这个数据量又取决于划分向量d和归约的维度。EinDecomp算法将cost_join和可选的cost_agg聚合成本也纳入动态规划的总成本计算中使得优化目标更加全面。4. 动态规划算法的实现与优化理解了原理我们来看EinDecomp动态规划算法如何具体实现并处理一些复杂情况。4.1 算法核心流程算法输入一个EinGraph和目标并行度p处理器数量假设为2的幂次方。输出一个为每个顶点标注了划分向量d的TaskGraph。拓扑排序首先对EinGraph进行拓扑排序得到一个顶点处理顺序列表。状态表初始化创建一个字典或表格M用于存储M[v, dZ]。动态规划递推按照拓扑顺序遍历每个顶点v。对于顶点v枚举其所有可能的输出划分向量dZ。对于一个EinSum操作其可能的输出划分由其输入划分和操作语义决定。算法需要枚举所有能产生p个并行任务即内部划分d各维度乘积为p且输出划分恰好为dZ的方案。对于每一个(v, dZ)对根据递推公式计算其最小成本。这需要 a. 枚举所有能产生输出dZ的内部划分方案d。 b. 对于每个d枚举所有可能的左输入划分dX和右输入划分dY从状态表M中已计算的结果里获取。 c. 对每种组合(d, dX, dY)计算总成本M[vX, dX] M[vY, dY] cost_repart(dX-d) cost_repart(dY-d) cost_join cost_agg。 d. 取所有组合中的最小值存入M[v, dZ]并记录下取得最小值时对应的d,dX,dY用于后续回溯。回溯与标注处理完所有顶点后在输出顶点v_out的所有M[v_out, dZ]中找到成本最小的dZ*。从这个状态开始利用步骤3中记录的选择信息反向遍历图为每个顶点确定其最终采用的内部划分方案d和输出划分dZ从而完成TaskGraph的构建。4.2 处理复杂DAG图的线性化上述算法有一个关键假设计算图中每个非输入顶点的输出最多只有一个消费者。也就是说图基本上是一条链或树。这是因为M[v, dZ]只记录了子图在一种输出划分下的最优成本。如果一个顶点v的输出被两个后续顶点v1和v2同时消费那么为了计算v1和v2的成本我们需要知道v的输出划分。但v的最优划分可能对v1和v2来说不是同一个。这就需要在状态中同时记录v输出到v1和v2的划分组合导致状态空间爆炸。对于一般的DAG有向无环图EinDecomp采用了一种近似策略图的线性化Linearization。寻找关键路径在EinGraph中找到一条最长的计算路径例如从输入到输出的主链。路径优先优化在这条路径上运行完整的动态规划算法忽略从这条路径分支出去或汇入这条路径的边所带来的重分布成本。也就是说在计算路径上顶点的成本时如果某个输入来自路径外部就假设其成本为0且其划分可以任意选择以适应路径上的需求。迭代处理完成第一条路径的优化和标注后将其从图中“固定”下来即这些顶点的划分方案已确定。然后在剩余的图中寻找下一个最长的路径重复此过程但此时路径中顶点的输入如果来自已固定的顶点其划分方案是已知的需要作为约束条件。这种方法是一种贪心近似。它放弃了全局最优解但大大降低了算法复杂度。论文指出在像大语言模型这样的实际计算图中这种线性化方法产生的解质量仍然很高与全局最优解相差不大。4.3 状态空间的剪枝与优化即使对于线性链状态空间也可能很大。一个顶点可能的输出划分向量dZ的数量取决于张量的维度和并行度p。论文中给出了一个组合公式如果张量有D个唯一维度标签并行度p2^N那么可能的划分方式数量是C(ND-1, D-1)组合数。当N和D不大时例如N10, D6组合数为3003暴力枚举是可行的。但在实际实现中还可以进行剪枝对称性剪枝对于某些操作不同的划分向量在通信成本上可能是等价的。启发式剪枝优先考虑那些将计算负载均匀分配到处理器上的划分方案。基于代价的剪枝在枚举过程中如果发现某种部分组合的成本已经超过当前已知的最小成本上界可以提前终止对该分支的搜索。5. 实战评估EinDecomp表现如何理论再优美也需要实践检验。论文通过一系列实验将基于EinDecomp的Einsummable系统与多种现有方案进行了对比回答了核心问题自动并行化能打过手工优化吗5.1 实验设置概览实验在CPU集群和GPU服务器上进行。Einsummable系统作为EinDecomp的载体将EinSum计算编译成底层内核CPU上用MKL批处理矩阵乘GPU上用CuTensor并使用UCX进行通信。对比对象包括经典HPC库ScaLAPACKCPU分布式线性代数。通用并行框架DASKPython。主流深度学习框架PyTorch数据并行。专用并行策略针对矩阵乘的“SQRT”分解将矩阵按行列各切sqrt(p)份即经典的2.5D算法、针对LLM的Megatron模型并行、序列并行Sequence Parallelism等。5.2 关键实验结果与解读实验1矩阵链运算任务并行计算(A × B) (C × (D × E))。对比EinsummableEinDecomp vs. EinsummableSQRT vs. ScaLAPACK (CPU) / DASK (GPU)。结果CPU上EinDecomp显著优于ScaLAPACK甚至在某些均匀矩阵情况下也优于简单的SQRT分解。原因可能是EinDecomp能更好地与数据预放置等系统特性结合。GPU上对于均匀矩阵EinDecomp与SQRT即经典2.5D算法性能相当但对于非均匀矩阵EinDecomp有近2倍的优势。这说明自适应划分对于不规则计算至关重要而固定策略SQRT无法适应。实验2大规模前馈神经网络训练任务训练一个超大规模分类器输入特征约60万隐藏层8192神经元输出类别近1.5万。对比EinsummableEinDecomp vs. PyTorch数据并行。结果EinDecomp完胜。在批大小较小128时优势尤其巨大。这是因为数据并行需要将整个大模型广播到所有GPU通信开销极大。而EinDecomp自动选择了比单纯数据并行更优的混合并行策略可能是模型并行与数据并行的结合大幅降低了通信量。实验3大语言模型LLaMA推理任务LLaMA模型70亿参数的“首令牌生成”Prefill阶段这是计算和内存密集型阶段。对比EinDecomp vs. Megatron模型并行 vs. 序列并行 vs. 注意力头并行。结果EinDecomp在绝大多数配置下不同批大小、GPU数量都匹配或超越了所有手工设计的专用并行策略。一个有趣的发现是简单的“序列并行”按输入序列维度切分表现也相当不错但依然不敌EinDecomp。这证明了自动搜索策略的有效性它能找到人类专家可能忽略的、更优的混合并行方案。实验4内存受限下的LLM推理任务在有限GPU内存下运行超大模型650亿参数推理。对比EinsummableEinDecomp vs. ZeRO-Inference vs. FlexGen。结果EinsummableEinDecomp大幅领先。ZeRO和FlexGen是PyTorch生态中专门为解决大模型内存问题而手工设计的系统如分层卸载参数。虽然这个对比更多体现了Einsummable底层Turnip引擎的“非确定性”内存卸载能力但也说明了由EinDecomp生成的并行计划能与先进的内存管理机制协同工作实现整体高性能。5.3 实验结论与启示综合来看EinDecomp驱动的自动并行化在从矩阵计算到LLM的多种任务上都展现出了与手工优化方案竞争甚至超越的实力。这带来了几个重要启示通用性战胜专用性手工优化的策略如Megatron为特定模型Transformer设计而EinDecomp适用于任何EinSum计算图更具通用性。混合并行是未来最优的并行策略往往是数据并行、模型并行、序列并行等的复杂混合。自动算法比人类更擅长探索这个庞大的组合空间。系统协同设计自动并行算法需要与底层执行引擎如Turnip的内存管理、编译优化等紧密集成才能发挥最大效力。EinDecomp提供了一个清晰的接口TaskGraph来连接优化器与运行时。6. 深入思考优势、局限与未来方向EinDecomp提出了一种优雅且有效的自动并行化框架但它并非银弹。在实际应用和系统设计中我们需要辩证地看待其优势和面临的挑战。6.1 核心优势声明式抽象基于EinSum的声明式编程模型将“计算什么”与“如何并行计算”彻底分离。程序员只需描述计算逻辑系统自动寻找最优并行方案极大提升了开发效率。全局优化动态规划方法从计算图全局视角进行优化能够协调不同操作间的划分策略避免局部最优导致的整体通信瓶颈。这是它优于许多逐层或启发式策略的关键。成本模型驱动以通信开销上界作为优化目标直接瞄准了分布式计算的主要性能瓶颈。该模型虽然简化但抓住了主要矛盾在实践中被证明是有效的。实用性通过图的线性化等近似方法算法能够处理实际中常见的复杂DAG并在可接受的时间内给出高质量解。6.2 当前局限与挑战搜索空间与可扩展性动态规划的状态空间随图的深度和宽度、张量维度、并行度指数级增长。虽然通过限制划分向量为2的幂次和线性化进行了约束但对于极端大规模或维度非常高的计算图搜索可能仍会变得昂贵。需要更高效的剪枝策略或启发式搜索。成本模型的精确性当前的通信成本模型是一个上界未考虑网络拓扑如NVLink、InfiniBand、硬件层次结构NUMA、GPU间互联带宽差异等具体细节。一个更精细的、硬件感知的成本模型能带来更好的优化效果。计算与通信的重叠现代高性能计算极度依赖计算与通信的重叠Overlap以隐藏延迟。当前的成本模型主要评估通信量未对计算-通信重叠的潜力进行建模。最优的划分方案可能需要考虑这种重叠的可能性。动态形状与条件分支当前算法假设计算图和张量形状是静态已知的。对于动态形状或包含条件分支的计算图如某些动态神经网络如何应用该框架是一个开放问题。与底层编译器的协同EinDecomp产出TaskGraph但每个“任务”内核本身的性能也至关重要。如何将划分方案与底层内核的代码生成如TVM、MLIR相结合实现从高层并行策略到底层代码优化的全栈协同是工程上的巨大挑战。6.3 未来可能的演进方向结合上述局限我们可以看到几个有潜力的发展方向分层优化与学习型优化器分层可以先使用轻量级启发式或基于历史数据的预测快速缩小搜索空间再在子空间内进行精确的动态规划搜索。学习可以利用强化学习或图神经网络学习计算图特征与最优并行策略之间的映射关系用推理时的快速预测替代耗时的在线搜索。硬件感知的成本建模将网络带宽、延迟、GPU间P2P带宽、甚至缓存层次纳入成本模型。例如在同一台服务器内的GPU间通信成本远低于跨服务器通信优化算法应倾向于将通信密集的操作放在机箱内。扩展至更复杂的执行模型考虑流水线并行Pipeline Parallelism。这需要将执行时间线Timeline纳入优化目标而不仅仅是通信量可能涉及更复杂的动态规划或整数规划。考虑异构计算CPUGPU其他加速器为不同算子选择不同的执行设备并优化跨设备的数据移动。与生态的集成将EinDecomp的思想集成到主流框架如PyTorch、JAX的编译器后端如TorchDynamo/XLA、OpenXLA中。作为其自动并行化策略的一个可选或补充方案。7. 总结与个人体会回顾EinDecomp算法它的核心贡献在于提供了一种系统化的、基于严谨优化理论的自动张量并行化框架。它将一个看似需要“艺术”和“经验”的并行策略设计问题转化为了一个可计算、可优化的数学问题。动态规划的应用确保了在链式或树状计算图上的最优性而线性化近似则为处理更通用的DAG提供了实用且高效的解决方案。从我个人的系统开发经验来看这类自动优化技术的价值不仅在于其最终的性能数据更在于它降低了分布式机器学习系统的使用和开发门槛。对于算法研究员他们可以更专注于模型结构本身而无需深陷并行化的细节泥潭。对于系统开发者EinDecomp提供了一个清晰的优化模块接口可以将其作为更庞大编译优化管道中的一环。当然没有哪个自动优化器是万能的。EinDecomp的成功依赖于EinSum这一简洁而强大的抽象。如果计算无法用EinSum清晰表达或者包含了大量非规则、控制流密集的操作它的适用性就会受限。此时可能仍需结合算子库手工优化或采用其他并行范式。最后分享一个实操中的关键点如果你正在自己的系统中尝试实现或借鉴类似思想成本模型的校准至关重要。论文中的公式是一个很好的起点但你必须针对你的硬件平台特别是网络进行微调。一个简单有效的方法是设计一组微基准测试Micro-benchmark实际测量不同形状张量在不同划分下的通信时间然后用这些数据来拟合或调整成本模型中的参数。一个贴合实际硬件的成本模型是自动优化器产出有效结果的基石。盲目套用理论公式可能会让优化器在错误的方向上狂奔结果还不如一个简单的启发式规则。