Scalify：基于e-graph与符号推理的分布式机器学习静默错误检测工具

1. 项目概述与核心挑战在分布式机器学习的世界里我们常常需要将一个庞大的模型拆解分散到成百上千个计算设备GPU、TPU、Neuron Core上并行执行以应对模型参数量和数据量的爆炸式增长。这个过程我们称之为模型并行化。听起来很美好但实际操作起来就像指挥一个庞大的交响乐团每个乐手计算设备都必须严格遵循乐谱计算图的指示稍有差池整场演出就会走调。这里的“走调”在机器学习中就是静默错误——程序不崩溃、不报错但悄无声息地输出了错误的结果。想象一下你训练了一个月的大模型最后发现因为一个张量在某个设备上被错误地reshape或transpose了一下导致所有努力付诸东流。这种错误极难通过常规测试发现因为它们不引发异常只导致模型精度下降或行为异常。Scalify正是为了解决这个痛点而生的工具。它的核心使命是自动验证一个模型在单设备上运行的计算图我们称之为基线图与它在分布式环境下经过各种并行化优化如张量并行、专家并行后产生的计算图在数学语义上是否完全等价。简单来说Scalify要回答的问题是“我做的这些为了加速而进行的复杂变换有没有改变模型原本要计算的东西” 这个问题的答案对于确保百亿、千亿参数大模型训练和推理的可靠性至关重要。接下来我将深入拆解Scalify是如何运用e-graph等式图这一形式化工具结合巧妙的符号推理来高效、精准地完成这项看似不可能的任务的。2. 核心原理当e-graph遇见张量布局要理解Scalify首先得弄明白它手中的两把“利器”e-graph和张量布局关系。2.1 e-graph将“等价”形式化e-graph等式图是一种数据结构它不直接存储计算表达式而是存储表达式之间的等价关系。你可以把它想象成一个“等价类”的集合。在e-graph中一个计算图里的每个操作节点都可以有多个不同的、但语义等价的表达形式这些形式被归入同一个等价类e-class中。例如对于矩阵乘法(A B) C根据结合律它等价于A (B C)。在e-graph中这两个表达式会被归入同一个e-class。Scalify利用这一点将基线计算图和分布式计算图都“喂”给同一个e-graph。工具内部预置了一系列元规则这些规则描述了在分布式环境下哪些变换是保持语义不变的。例如“一个All-Reduce操作后接一个加法可以等价于先在本地做加法再进行All-Redduce”如果数据布局允许。Scalify的工作流程就是不断地应用这些元规则对e-graph进行重写和合并。如果最终基线图中代表模型输出的节点和分布式图中对应的输出节点被合并到了同一个e-class中那么我们就证明了这两个图在给定规则下是语义等价的。反之如果它们始终无法合并就说明存在不等价即引入了静默错误。注意e-graph的重写和合并是“饱和”过程它会探索所有已知规则下的等价形式。这比单纯地按路径比较两个图要强大得多因为它能发现通过不同变换序列达到相同结果的等价性。2.2 张量布局关系跨越设备边界的桥梁然而在分布式场景下事情没那么简单。一个张量可能被分片到多个设备上。例如一个形状为[4096, 4096]的权重矩阵在4路张量并行下可能被按列切分成4个[4096, 1024]的碎片分布到4个设备上。这时设备上的局部张量Sharded Tensor和全局的逻辑张量Global Tensor之间就存在一种布局映射关系。Scalify的核心创新之一就是显式地建模并跟踪这种布局关系。它为e-graph中的每个张量节点都附加了布局信息。这个信息不是一个具体的形状而是一个符号化的轴映射表达式。举个例子假设基线张量B的形状是(i, j, k)表示三个维度的规模。在分布式图中对应的张量D可能因为一个reshape操作变成了(⊗(i, j), k)这里⊗表示将前两个维度合并。那么B和D之间的布局关系S(B, D)就可以表述为B的(i, j, k)轴分别映射到D的(⊗(i, j), k)轴。Scalify在e-graph中传播等价关系时会严格检查布局关系是否兼容。只有当两个操作如加法的输入张量不仅计算语义等价而且它们的布局关系也一致时Scalify才会认为这两个操作的结果是等价的并合并它们的输出节点。这就像确保交响乐团中不仅小提琴组和中提琴组演奏的旋律要对他们的音高基准布局也必须一致否则合奏就是混乱的。3. 关键技术拆解双射推断与错误定位理解了基本原理我们来看Scalify解决的两个最棘手的问题如何判断复杂的布局变换序列是否等价以及一旦发现不等价如何精准定位bug。3.1 符号化双射推断破解布局变换的迷宫分布式优化常常引入一连串的reshape和transpose操作来调整数据布局以适应不同的并行策略或硬件特性。如图9所示基线路径和分布式路径可能各自经过不同的变换序列最终得到形状相同的张量。Scalify需要判断这两个序列在数学上是否描述的是同一个张量重排它的方法是符号化双射推断其算法对应原文Algorithm 2可以分解为四步生成符号表达式将每条路径上的布局操作序列如reshape,transpose转化为对张量轴的符号化操作。例如形状(4, 64, 4096)被符号化为(i, j, k)。一个reshape((4,64,4096), (256,4096))操作会被转化为将轴(i, j)合并为⊗(i, j)得到表达式(⊗(i, j), k)。秩归一化比较两条路径的最终符号表达式。如果它们的“秩”即维度数量不同Scalify会尝试通过引入虚拟的、规模为1的维度来进行归一化使两者具有相同的秩以便于比较。如果无法合理归一化则直接判定为不等价。寻找置换双射这是核心步骤。Scalify试图在归一化后的分布式表达式Êd和基线表达式Êb之间找到一个轴的置换关系。它逐个检查Êb中的每个符号轴如i在Êd中寻找结构上相等的对应轴考虑布局映射关系M。如果能找到唯一且一一对应的映射就形成了一个置换索引如(1, 0, 2)这对应一个transpose操作。构造操作序列根据找到的置换关系Scalify可以反向构造出一个操作序列。这个序列作用在分布式路径的终点张量上能将其变换到与基线路径终点张量完全一致的布局。例如生成的序列可能是[reshape((256,4096), (64,4,4096)), transpose(1,0,2), reshape((64,4,4096), (256,4096))]。如果这个构造的序列能使得两条路径等价则双射推断成功。这个过程本质上是在符号层面进行图同构匹配。它避免了直接对具体数值进行昂贵的符号执行而是利用张量操作的代数性质进行推理效率极高。实操心得双射推断的成功高度依赖于对reshape操作语义的精确建模——即它只能合并或拆分连续的维度。Scalify目前将范围限定于此这覆盖了生产框架中绝大多数情况如Megatron-LM、DeepSpeed中的张量分组在实用性和复杂性之间取得了良好平衡。如果你想将其用于更诡异的维度重排可能需要扩展这里的规则。3.2 基于e-graph差异的代码级错误定位验证失败输出“不相等”只是一个开始。对开发者来说更重要的是知道bug在哪里。Scalify的另一个强大之处在于它能进行精确的错误定位。它通过在编译过程中插桩将中间表示IR图中的每个节点都与源代码的抽象语法树AST节点关联起来记录下文件名、函数名、行号等元数据。当e-graph重写完成并检测到不等价时Scalify不会简单地列出所有“未验证”的节点——在复杂的非等价图中这样的节点可能成千上万毫无帮助。Scalify采用了一种更聪明的策略它遍历那些“未验证”的节点但只关注那些输入节点已经被验证为等价的未验证节点。为什么因为如果一个节点的所有输入都是等价的但这个节点本身的输出不等价那么问题很可能就出在这个节点所代表的操作上或者其直接的变换上。如图10所示的例子一个add操作未被验证。Scalify检查发现它的两个输入张量虽然各自都能通过某种双射与基线图的对应输入对齐但对齐所需的双射序列却不相同。这意味着在add操作执行之前两个输入张量的数据布局已经不一致了因此add无法进行有效的元素级加法。Scalify会报告这个add节点及其源代码位置并指出其输入已验证但自身失败从而将开发者的注意力直接引向导致布局分歧的根源——很可能是不正确的reshape或通信操作。这种定位方式极大地缩小了调试范围将海量的IR节点排查变成了对少数几个“输入已验证但自身未验证”的关键节点的审查。4. 实现与评估在真实模型上的实战表现理论再优美也需要实践检验。Scalify的实现大约有9000行Python代码其中约6500行用于手动编码25个针对不同并行模式张量、专家、序列并行等的元规则。它构建在PyTorch XLA之上直接处理ML模型的中间表示并与egglog集成作为其e-graph引擎。4.1 验证能力与效率评估显示Scalify能够处理Llama-3.18B到405B参数和Mixtral8x7B, 8x22B这样的真实世界大模型。所有验证均在几分钟内完成见表2运行在普通的6核CPU和16GB内存的机器上。这与之前一些基于SMT求解器的方法如TrainVerify形成了鲜明对比后者对于405B的模型可能需要数天时间。关键效率来源张量级抽象Scalify在张量层面而非元素层面进行推理将分片张量视为一个整体实体极大减少了推理状态。层记忆化与图分区对于Transformer这类具有重复层结构的模型Scalify采用了层记忆化技术。它验证完一个Transformer层后会缓存该层的等价性证明后续相同结构的层直接复用结果避免了重复计算。同时它将大计算图分区成子图分别处理降低了单次e-graph重写的复杂度。复杂度与模型规模解耦如图11所示Scalify的验证时间与输入张量的具体形状序列长度、批大小以及并行度TP degree无关仅与模型的层数呈线性关系图11c。这是因为其推理完全在计算图的结构层面进行不涉及具体数值。4.2 错误检测效果Scalify被设计用于检测五类典型的静默错误错误的分布式操作例如使用了不必要的all-reduce或者该用all-gather时用了reduce-scatter。错误的分布式配置例如归约操作只在部分设备子集上进行。不一致的张量精度单机流水线和分布式流水线使用了不同的数值精度如FP16 vs BF16。错误的轴切分reshape操作错误地分割了张量破坏了分片关系。错误的布局优化与基线相比使用了无效的布局变换序列。在复现的19个历史真实bug中Scalify成功检测出17个并在1分钟内完成。更重要的是它能将其中许多错误定位到具体的源代码行或可疑函数。此外在评估过程中Scalify还在AWS Neuron SDK中发现了5个此前未知的新bug这些bug都可能导致严重的正确性问题并已提交给开发者修复。5. 局限性与未来方向当然Scalify并非万能。首先它专注于计算图级别的验证。那些在图形编译阶段之后出现的错误例如分布式流水线中的数据竞争、运行时内存错误等超出了它的检测范围。其次Scalify是可靠但不完备的。这意味着所有被它验证通过的图我们都可以相信是正确的可靠性但有些正确的图可能因为其使用的变换超出了当前预定义元规则或双射推断的范围而无法被验证不完备性。目前Scalify对Tensor Parallelism, Flash Decoding, Expert Parallelism等主流并行模式支持良好但对于更复杂的流水线并行或涉及动态控制流的模式需要额外的工程努力来扩展规则集。最后虽然Scalify能精确定位到出现差异的代码行但根因分析有时仍需开发者手动进行。未来的一个有趣方向是结合大语言模型LLM利用其代码理解能力对Scalify定位出的可疑代码片段进行自动分析推测可能的错误原因从而进一步降低调试门槛。从我个人的实践经验来看像Scalify这样的形式化验证工具正在成为大规模机器学习系统开发中不可或缺的“安全带”。它不能替代全面的测试但能为那些最隐蔽、代价最高的静默错误提供一道强有力的防线。在动辄消耗数百万美元计算资源的千亿模型训练中提前几分钟发现一个布局错误其价值不言而喻。将验证左移从运行时测试前置到编译时证明是提升ML系统可靠性的必然趋势。