1. 项目概述在分布式机器学习的世界里我们常常面临一个看似简单实则棘手的问题我写的这个并行化代码真的和单机版本在数学上等价吗这个问题背后是无数个深夜调试的工程师是那些在数百个GPU上跑了一周才发现结果不对的昂贵实验更是那些难以察觉却足以毁掉整个模型训练的“静默错误”。这些错误不会导致程序崩溃它们悄无声息地改变着张量的布局、精度或通信逻辑最终让训练出的模型性能远低于预期而你却可能把原因归结为超参数没调好。Scalify正是为了解决这个痛点而生。它不是一个运行时监控工具而是一个编译时验证器。它的核心任务是在你的分布式计算图比如为32个GPU设计的张量并行计算流被真正执行之前就证明它在数学语义上与一个已知正确的“基线”计算图通常是单机版本完全等价。这听起来像是形式化验证的领域但Scalify的巧妙之处在于它没有选择传统的、计算量巨大的SMT求解器路径而是引入了来自编程语言优化领域的“e-graph”等式图技术结合了关系传播和符号推理将验证时间从“天”级别压缩到了“分钟”级别。想象一下你正在为一个拥有4050亿参数的Llama-3.1模型实现Flash Attention的分布式版本。你写了一系列复杂的reshape和transpose操作来调整张量布局以适应多卡通信同时加入了all-reduce来同步梯度。代码能跑通loss也在下降但你怎么能百分之百确定经过这一系列眼花缭乱的变换后最终每个设备上的计算结果和你在单张卡上顺序执行的结果在数学上是一模一样的Scalify就是你的“数学证明助手”。它不关心具体的数值而是关心计算的“形状”和“关系”。它通过分析计算图的结构、张量的维度映射以及操作符的语义来推理两个图是否表达了同一个数学函数。2. 核心原理为什么是e-graph与关系传播要理解Scalify首先要抛开“逐元素比对”的直觉。在分布式场景下一个大的张量被切分Shard到多个设备上每个设备只持有数据的一部分。直接比较两个图上对应节点的具体数值是行不通的因为数据分布已经不同。Scalify的智慧在于它比较的是“关系”而非“数值”。2.1 e-graph将等价性搜索转化为图重构问题e-graphEquality Graph是一种数据结构它能够高效地表示一个表达式集合并记录这些表达式之间的等价关系。在传统编译器优化中e-graph被用来做等式饱和Equality Saturation穷举一个表达式所有可能的等价形式从而找到最优的实现。Scalify创新性地将e-graph用于验证。它的工作流程可以概括为以下几步图转换将输入的基线计算图单机和分布式计算图分别转换为两个初始的e-graph。图中的每个节点代表一个操作如dot,add,reshape边代表数据流。规则应用预定义一组“重写规则”Rewrite Rules。这些规则描述了在保持语义等价的前提下一个计算子图可以被另一个子图替换。例如一个transpose后接reshape的操作序列可能等价于另一个不同的reshape后接transpose的序列。Scalify将这些规则反复应用到两个e-graph上。等价类合并当规则应用表明两个节点或子图等价时e-graph会将它们合并到同一个“等价类”中。这个过程是自反、对称和传递的。最终判定经过多轮规则应用和合并后Scalify检查两个e-graph的“根节点”即整个计算图的输出是否被合并到了同一个等价类中。如果是则证明两个计算图语义等价否则不等价。注意这里的“规则”不是任意的而是基于严格的数学定义和分布式计算语义如集合通信all-reduce的语义手工编码的。Scalify论文中提到了编码了约25条这样的元规则覆盖了张量并行、专家并行等多种并行模式。2.2 关系传播在分布式上下文中建立桥梁e-graph解决了“如何判断等价”的问题但分布式图与基线图之间还存在一个根本差异数据布局。这就是“关系传播”要解决的问题。Scalify在验证之初会为两个图中对应的输入张量例如同一个权重矩阵在单机版本和切分后的版本建立一种“布局关系”。这种关系描述了分布式张量中的每个逻辑元素对应到基线张量中的哪个位置。例如对于一个形状为(1024, 4096)的矩阵进行按列切分的张量并行TP4每个设备上的张量形状为(1024, 1024)。布局关系会记录设备0上的张量A_shard[0]对应基线张量A的列[0:1024]。验证的核心就是让这种布局关系随着计算图的执行而“传播”下去。Scalify会沿着数据流图分析每一个操作符Operator对输入输出张量布局的影响逐元素操作如add,mul,sin不改变布局。如果输入张量A和A‘有布局关系R那么经过add操作后输出张量B和B‘也自动具有相同的布局关系R。改变形状的操作如reshape,transpose会改变布局关系。Scalify需要计算新的映射。这是最复杂、最容易出错的部分也是Scalify算法双射推断的核心应用场景。通信操作如all-reduce,all-gather会同步或重组数据。Scalify需要根据通信操作的语义例如all-reduce是求和后广播来更新布局关系。例如一个按行切分的张量经过all-reduce后每个设备都拥有了完整的行其布局关系就变成了“复制”关系。如果关系能够从输入一直无矛盾地传播到最终输出并且最终输出的布局关系表明它们在逻辑上是同一个完整张量或等价的复制品那么验证就通过了。如果在某个节点输入的关系无法推导出一致的输出关系或者两个本应等价的分支出现了关系冲突Scalify就会标记此处为“未验证”并定位到可能的错误源头。3. 核心挑战与突破张量布局变换的等价性验证在分布式机器学习优化中大量的“静默错误”并非源于通信原语用错而是隐藏在那些为了性能而引入的、复杂的张量布局变换序列中。不同的编译器优化Pass、不同的内核实现可能会生成不同的reshape和transpose操作序列但它们的目标是实现同一个逻辑上的数据重排。手动验证这些序列的等价性极其容易出错而这正是Scalify的“符号双射推断”算法大放异彩的地方。3.1 问题定义从具体维度到符号表达式让我们通过一个论文中的简化例子来理解这个问题。假设基线图中有一个操作序列张量A (4, 64, 4096) - reshape - (256, 4096)我们称之为布局序列Sb。在分布式图中为了实现某种并行优化编译器可能生成了另一个序列张量A‘ (4, 64, 4096) - transpose - (64, 4, 4096) - reshape - (256, 4096)我们称之为布局序列Sd。肉眼观察(4,64,4096)经过transpose(1,0,2)变成(64,4,4096)再reshape成(256,4096)似乎和直接reshape成(256,4096)是等价的但我们需要严格的证明。Scalify的第一步是符号化。它将具体的维度数字4, 64, 4096用符号轴代替例如(i, j, k)。这样操作就变成了对符号表达式的变换Sb:(i, j, k) - reshape - (⊗(i, j), k)。这里⊗表示合并merge操作。Sd:(i, j, k) - transpose - (j, i, k) - reshape - (⊗(j, i), k)。现在问题转化为表达式(⊗(i, j), k)和(⊗(j, i), k)是否在逻辑上表示同一个多维数组的索引空间换句话说是否存在一个双射bijection函数能将Sd产生的数据布局一一映射到Sb产生的布局上3.2 双射推断算法四步走Scalify的算法2Algorithm 2清晰地描述了这一过程我们可以拆解为四个步骤步骤1生成符号表达式与轴映射如上所述为两个布局序列Sb和Sd生成符号表达式Eb和Ed。同时根据输入张量b和d的切分关系建立一个初始的轴映射M。例如如果d是b按第0维切分的一部分那么M会记录d的轴i‘对应b的轴i的一个子区间。步骤2秩归一化Eb和Ed的“秩”rank即维度数量可能因为reshape的合并/拆分而不同。例如Eb可能是(⊗(i, j), k)秩为2而Ed可能是(j, i, k)秩为3。Scalify通过引入虚拟的、大小为1的维度或将合并的维度临时拆开将两个表达式归一化到相同的秩得到Êb和Êd。如果无法归一化例如总元素数不同则直接返回“无等价关系”。步骤3寻找置换双射这是算法的核心。现在Êb和Êd具有相同的秩。Scalify需要找到一个置换permutationp使得将Êd的轴按照p重新排列后其结构与Êb在轴映射M的意义下“相等”。 对于Êb中的每一个符号轴如i在Êd中寻找一个符号轴如j使得在映射M下j和i代表的是基线张量中同一个逻辑轴或它的一个部分。如果对于Êb中的每个轴都能在Êd中找到唯一且不重复的对应轴那么就成功构造了一个置换p。 在上述例子中Êb (i, j, k)Êd (j, i, k)。显然i对应Êd中的jj对应Êd中的ik对应k。因此得到的置换p [1, 0, 2]这正好对应一个transpose(1, 0, 2)操作。步骤4构造操作序列最后Scalify将推断出的双射具体化为一个可执行的操作序列。这个序列的作用是当应用于分布式路径的末尾时能将其张量布局转换为与基线路径兼容的布局。如果Êd与原始的Ed形状不同先添加一个reshape操作将Ed的形状变为Êd的形状即归一化后的形状。添加transpose(p)操作应用步骤3找到的置换。如果Êd经过置换后的形状与Eb的最终形状不同再添加一个reshape操作变到最终目标形状。对于我们的例子算法会推断出双射操作序列为[reshape(64, 4, 4096), transpose(1, 0, 2), reshape(256, 4096)]。这个序列恰好是Sd的逆过程将它应用到分布式路径的末尾就能“抵消”掉分布式图中额外的transpose使其输出布局与基线路径一致。3.3 实操心得理解算法的局限性这个算法非常强大但它也有明确的适用范围理解这一点对正确使用Scalify至关重要。注意Scalify的双射推断算法主要针对维度的合并与拆分即reshape操作以及维度的重排即transpose操作。它假设布局变换主要由这两类操作构成这在生产级ML框架如PyTorch XLA、TensorFlow中是非常普遍的。然而它可能无法处理更广义的、非线性的索引变换例如gather、scatter、strided_slice等复杂操作。在实现自己的规则时需要确保操作语义在算法定义的范畴内。一个常见的陷阱算法依赖于准确的轴映射M。如果初始的切分关系定义错误例如误以为张量是按行切分但实际是按列切分那么整个双射推断就会建立在错误的基础上导致误判。因此在配置Scalify验证任务时明确定义每个输入张量的切分策略是第一步也是最重要的一步。4. 实战使用Scalify定位静默错误验证工具的价值不仅在于说“是”或“否”更在于当答案是“否”时它能告诉你“哪里出了问题”。Scalify的“后处理基于E-Graph差异定位代码错误”模块正是为此设计。4.1 错误定位机制Scalify在将IR计算图转换为e-graph时会通过编译器的日志API注入元数据将图中的每个节点与源代码中的具体位置文件、函数、行号关联起来。当验证失败时Scalify不是简单地列出所有“未验证”的节点——在不等价的图中这样的节点可能非常多。Scalify采用了一种更智能的溯源策略分类节点在重写过程中将所有节点分为“已验证”和“未验证”两类。寻找分歧点遍历所有“未验证”的节点检查它的所有输入节点是否都是“已验证”的。报告根源如果一个节点的所有输入都已被验证为等价但这个节点本身却无法被验证那么这个节点就很可能是错误的根源。因为错误一定发生在这个节点或其内部操作上而不是更早的祖先节点。4.2 案例分析错误的All-to-All布局变换让我们剖析论文中图10Figure 10的经典案例。这是一个在混合并行如同时使用张量并行和序列并行中容易出现的错误。基线图Oracle路径相对简单A和B进行矩阵乘法dot后得到C。分布式图Buggy为了适应并行对A‘和B‘做了额外的reshape和transpose操作然后进行dot得到C‘。之后需要对C‘进行一个all-reduce通信操作来聚合结果。错误出现在哪里Scalify发现A‘和B‘的布局变换序列与C‘和C之间所需的布局关系不匹配。具体来说为了使A‘与A对齐需要应用一个双射[transpose(2,0,1), reshape(...)]。为了使B‘与B对齐需要应用同一个双射。但是为了使all-reduce之后的C‘与C对齐需要应用一个不同的双射[transpose(2,1,0), ...]。这就产生了一个矛盾C‘是由A‘和B‘计算得来的如果A‘和B‘遵循第一种变换关系那么计算得到的C‘理应也通过第一种变换关系与C对齐。但实际代码中all-reduce操作前后隐含的布局假设却是第二种关系。这种不一致导致Scalify无法将布局关系通过add操作图中add节点传播下去。Scalify的输出它会报告这个add节点是未验证的并且其所有输入节点即all-reduce的结果和另一个张量都是已验证的。这直接将开发者的注意力引向了这个add操作或者更准确地说引向了产生add操作输入的上游代码——即那个有问题的all-reduce布局变换。错误信息会附带源代码行号例如hlo.py:214让开发者能迅速定位到问题代码。4.3 常见错误类型与Scalify的检测能力根据论文评估Scalify可以有效检测以下几类分布式机器学习中典型的“静默错误”不正确的分布式操作例如该用all-gather的时候用了all-reduce或者多了一个不必要的all-reduce。Scalify通过通信原语的语义规则检测。不正确的分布式配置例如设备分组replica groups设置错误导致只在部分设备上进行规约。Scalify通过分析操作符的设备属性关联来发现。不一致的张量精度单机图使用FP32分布式图某个环节错误地使用了FP16或BF16。Scalify可以检查操作符的dtype属性。不正确的轴拆分reshape操作错误地拆分或合并了维度破坏了张量切分关系。这是双射推断算法的主要检测目标。不正确的布局优化编译器或手动编写的布局变换序列存在错误与基线逻辑不等价。同样是双射推断算法的检测目标。实操心得并非所有错误都能在计算图层面捕获。Scalify明确指出了其局限性它专注于计算图IR级别的验证。因此诸如运行时调度错误如数据竞争、在图形编译阶段之前发生的错误、或者那些不影响计算图语义但影响数值稳定性的超细微差别Scalify可能无法发现。它是一款强大的“图形逻辑验证器”但不能替代全面的集成测试和数值精度测试。5. 性能与评估真的能用于大模型吗一个验证工具如果速度太慢就无法集成到开发流程中。Scalify论文中的评估数据有力地证明了其实用性。5.1 验证时间从数天到数分钟论文对比了之前的SOTA工具TrainVerify。TrainVerify使用SMT求解器进行逐元素的推理验证一个405B参数的Llama-3.1模型需要数天时间。而Scalify通过对张量进行整体关系推理将验证时间缩短到了2分37秒在6核AMD Ryzen 5 5600U CPU16GB RAM的消费级机器上。对于8B和70B的模型验证时间更是低于2分钟。这个时间开销对于在代码提交前或CI/CD流水线中运行检查是完全可接受的。5.2 可扩展性分析Scalify的性能表现呈现出几个关键特征这些特征源于其设计与张量形状无关如图11a, 11b, 11e所示改变序列长度seqlen、批大小batch size或注意力头数heads几乎不影响验证时间。因为Scalify在符号层面操作计算图的节点和边数量不随具体维度大小变化。与并行度弱相关如图11d所示增加张量并行度TP degree并不会显著增加验证时间。因为增加核心数主要是在计算图中添加更多的通信节点而图的整体拓扑结构复杂度增长有限。与模型层数线性相关如图11c所示验证时间随模型层数增加而线性增长。这是因为更多的层意味着更长的计算图需要处理更多的节点。Scalify采用了层记忆化技术来优化这一点它将模型按层分割成子图验证完一层后缓存该层的等价类信息在验证下一层时可以直接复用避免了重复计算大幅提升了效率见图12对比。5.3 真实漏洞检测效果论文在AWS Transformers NeuronX框架上复现了19个真实世界中的历史bug并利用Scalify进行检测。结果令人印象深刻检测率17/19的bug被成功检测出约89.5%。定位精度对于其中13个bugScalify能精确定位到有问题的具体指令➸对于另外4个bug能定位到有问题的函数或数据结构✻。只有2个bug未被检测出原因是它们发生在图编译阶段之外如运行时KV缓存切片错误。发现新bug在评估过程中Scalify甚至在Amazon的SDK中发现了5个此前未知的bug包括不正确的布局优化、错误的all-to-all变换、张量切分错误等这些都已提交给开发者修复。这些数据强有力地证明了Scalify不仅是一个学术原型更是一个能在工业级复杂框架和模型上发现真实问题的实用工具。6. 实现与应用展望6.1 实现要点Scalify的核心实现大约9000行Python代码其中约6500行用于手工编码那25条关键的元规则。这些规则定义了不同并行模式张量、专家、序列并行下各种操作符的语义和等价变换。作者指出一旦核心框架和基础规则集搭建完成为新的并行技术添加规则支持所需的工作量是可控的例如为序列并行添加2条规则仅需30行代码。它构建在PyTorch XLA之上直接操作ML模型的中间表示IR。其核心的布局等价推理引擎是egglog一个e-graph库。虽然原型针对AWS Neuron SDK但其算法是框架无关的可以移植到其他基于IR的系统如TensorFlow XLA或Megatron-LM。6.2 局限性范围限制专注于计算图验证。运行时错误、调度竞争、编译前错误无法捕获。可靠性而非完备性Scalify是“可靠的”sound即它验证通过的图一定是正确的。但它不是“完备的”complete意味着可能存在一些正确的图由于规则集未覆盖或算法限制当前版本的Scalify无法验证。这需要不断扩展规则库。模式支持目前对Tensor Parallelism, Flash Decoding, Expert Parallelism支持良好但对更复杂的流水线并行Pipeline Parallelism等涉及复杂跨设备通信和运行时语义的模式需要更多工作来扩展支持。根因分析Scalify能精确定位到出现不一致的代码行但有时无法直接揭示错误的根本原因例如是哪个工程师在什么背景下引入了这个错误逻辑这仍需开发者人工分析。6.3 未来与启示Scalify为分布式机器学习系统的可靠性工程树立了一个新的标杆。它表明通过巧妙的抽象e-graph、关系传播、符号推理可以对大规模、复杂的计算图进行高效的形式化验证。这对于未来越来越庞大、并行策略越来越复杂的模型开发至关重要。对于开发者和团队而言可以考虑将此类验证工具集成到CI/CD流程中作为对分布式优化代码变更的强制性检查项从而在代码合并前就拦截可能导致静默错误的修改。同时工具揭示的“无法验证”区域也可以指导我们编写更完备的测试用例。从更广阔的视角看Scalify的成功是编程语言、形式化方法与系统工程交叉融合的典范。它解决的不是一个纯理论问题而是扎根于工业实践、具有明确度量标准和显著效用的实际问题。随着大模型训练和推理日益成为基础设施这类确保底层计算正确性的工具其价值只会与日俱增。
