Scalify：基于等式饱和与关系推理的分布式ML计算图形式化验证

1. 项目概述：当分布式ML框架的优化“静默”出错

在构建和训练现代大规模机器学习模型时，我们依赖PyTorch、TensorFlow、JAX等框架及其分布式扩展（如DeepSpeed、FSDP）将复杂的模型代码转化为高效的计算图，并应用一系列激进的优化：算子融合以减少内核启动开销，张量并行以拆分巨大参数，流水线并行以重叠计算与通信。这些优化是支撑千亿参数模型训练的基石。然而，一个长期被低估的危机是“静默错误”（Silent Errors）——优化后的计算图与原始图的数学语义并不等价，但程序不会崩溃，只是训练出的模型性能严重退化，损失曲线诡异，收敛困难。这种错误极其隐蔽，可能源于框架编译器的一个错误reshape、一次不当的all-reduce，或者对某种新颖并行策略支持的不完善。

传统的调试手段是“数值比对”：在相同输入下，逐层比对原始单卡执行与分布式优化后执行的张量输出。这方法不仅耗时（需要实际运行模型），更因浮点数精度、非确定性计算顺序、不同硬件内核实现差异而变得不可靠。微小的数值差异是常态，但如何区分“可接受的浮点误差”与“致命的语义错误”？更根本的是，当比对失败时，如何从海量算子中定位到那个出错的代码行？这常常让开发者陷入数天甚至数周的痛苦排查。

Scalify的提出，正是为了系统性地解决这个问题。它不依赖数值近似，而是直接在计算图的语义层面进行形式化验证。其核心思想令人耳目一新：将待验证的原始图（Baseline）与优化/分布式图（Target）同时注册到一个称为“等式图”（E-Graph）的数据结构中，然后应用一系列已知的、正确的语义等价规则（如交换律、结合律、分布式all-reduce与求和等价）进行并行重写。如果两个图在数学上等价，那么经过充分重写后，它们的输出节点最终会“合并”到E-Graph中的同一个等价类里；如果不等价，则它们会停留在不同的等价类，并且E-Graph的结构能清晰地揭示出从哪个子图开始分道扬镳。

这听起来像是一个完美的理论方案，但将其应用于生产级、超大规模模型（如Llama 3.1 405B）时，会面临几个毁灭性的工程挑战：1)状态爆炸：直接对完整模型图进行等式饱和，其搜索空间会随节点数指数增长，内存与时间瞬间耗尽。2)规则泛化与效率的权衡：过于通用的规则会导致匹配开销巨大；过于具体的规则又无法覆盖多样的优化模式，且难以维护。3)布局异构性：分布式优化中充斥着reshape、transpose、shard操作，许多不同的操作序列能产生形状相同的张量，但数据排布（内存布局）完全不同，语义可能天差地别。4)从“不等价”到“可调试”：仅仅告诉开发者“这两个图不等价”毫无用处，必须能精确定位到框架源代码中引入错误的具体算子或变换。

Scalify正是为攻克这些挑战而生的生产级验证框架。它通过一套组合拳——基于拓扑的图分区与层记忆化、可复用的重写模板、增强的关系推理与符号双射推断，以及基于差异的缺陷定位——实现了对超大规模计算图的分钟级验证，并成功在Amazon的生产框架中发现了此前未知的缺陷。接下来，我将深入拆解它的设计思路、核心实现与实战技巧。

2. 核心原理：等式饱和与关系推理如何为计算图“验算”

要理解Scalify，必须先吃透其两大理论基础：等式饱和（Equality Saturation）与Datalog风格的关系推理（Relational Reasoning）。这不是简单的算法应用，而是针对计算图验证这一特定问题的深度改造。

2.1 等式饱和：将“找不同”游戏转化为“找共同归宿”

等式饱和的核心数据结构是E-Graph。你可以把它想象成一个不断扩张的“等价宇宙”，这个宇宙里容纳了一个程序所有可能的、语义等价的表达形式。E-Graph由两类基本元素构成：

• E-Node（等价节点）：代表一个具体的操作或值，例如一个加法节点add(x, y)，一个常量2，或者一个张量节点A。
• E-Class（等价类）：一个集合，里面装着所有被认为是语义等价的E-Node。一个E-Node可以属于且仅属于一个E-Class。

初始时，我们将待验证的两个计算图（原始图和目标图）的节点分别放入E-Graph。此时，一个图中的节点各自属于独立的E-Class。然后，我们定义一系列重写规则。每条规则都是一个“如果…那么…”的语句，描述了一种语义等价变换。例如：

• 交换律：add(x, y) => add(y, x)
• 结合律：add(add(x, y), z) => add(x, add(y, z))
• 分布式求和等价：all_reduce(sum, shard(x, dim=0)) => x（在所有设备上对按第0维分片的张量进行求和式全归约，等价于未分片的原始张量x）

系统的工作就是不断扫描E-Graph，寻找可以应用这些规则的模式。一旦找到，它并不直接“替换”节点，而是在目标位置创建新的E-Node，并将其合并到与原始模式等价的E-Class中。这个过程是“饱和”的：规则被反复、并行地应用，直到没有新的等价关系可以被发现。

Scalify的巧妙应用在于，它将两个计算图视为同一个“程序”的两种可能表达。验证过程就是看它们的输出节点，在经过所有可能的、正确的语义变换后，能否最终归属于同一个E-Class。如果能，则证明二者语义等价；如果不能，则找到了语义分歧点。这比基于SMT求解器的方法更“轻量”，因为它利用了计算图大量结构相似的特点，避免了昂贵的非线性实数算术推理。

实操心得：规则的设计是灵魂规则库的构建是Scalify有效性的关键。我们不仅需要基础的数学定律（交换、结合、分配），更需要大量针对ML和分布式计算的领域特定规则。例如：

• 通信原语规则：all_gather(shard(x, dim, n), dim, n) => x
• 算子融合规则：gelu(matmul(x, w)) => fused_gelu_matmul(x, w)（需确保融合后的数值精度在可接受范围内）
• 布局变换规则：reshape(transpose(x, (0,2,1)), (s1, s2)) => transpose(reshape(x, (s2, s1)), (1,0))（某些特定形状下成立） 规则的设计需要在“通用性”和“性能”间做权衡。过于通用（如“任何不改变元素顺序的reshape-transpose序列等价”）的规则会引发组合爆炸；过于具体又难以覆盖所有情况。Scalify的策略是提供一套核心通用规则，并允许框架开发者通过“模板”机制添加针对其特定优化pass的规则。

2.2 关系推理：为张量赋予“身份”与“关系”

仅靠等式饱和处理纯计算语义还不够。分布式计算的核心是数据变换：一个张量被分片（Shard）到多个设备，被转置（Transpose）改变内存布局，被重塑（Reshape）改变视图。我们必须能形式化地描述和推理这些变换之间的关系。

Scalify引入了一个基于Datalog的小型关系语言。Datalog是一种声明式逻辑编程语言，非常适合描述“关系”和推导新的事实。Scalify定义了以下几种核心关系：

1. 分片关系：sharded(t_base, t_dist, dim, n_devices)。表示分布式张量t_dist是基线张量t_base在第dim维上被分片到n_devices个设备的结果。
2. 复制关系：duplicate(t_base, t_dist, n_devices)。表示t_dist是t_base在n_devices个设备上的完整副本。
3. 布局关系：layout(t_base, t_dist, L)。表示张量t_dist可以通过一个双射布局变换L（一系列reshape和transpose的组合）转换到t_base。这是处理布局异构性的关键。
4. 部分结果关系：partial(t_base, t_dist, n_devices, op)。表示t_dist是一个需要在n_devices个设备上通过操作op（如sum、max）进行归约的中间部分结果。

这些关系在验证开始时，通过编译器插桩自动注入。例如，当框架生成一个分片算子时，它会同时输出sharded关系事实。对于某些无法自动推导的辅助张量（如设备ID张量），则需要手动注解。

推理过程是增量式的。系统从已知的输入关系（如输入张量是如何被分片的）开始，随着计算图的遍历，应用预定义的Datalog规则来传播这些关系。例如，有一条规则是：

sharded(x, x', dim, c), sharded(y, y', dim, c), z = add(x, y), z' = add'(x', y') => sharded(z, z', dim, c)

这条规则是说：如果x和y都在同一维度被分片到c个设备，那么它们的逐元素加法结果z和z'也保持同样的分片关系。

通过这种关系推理，Scalify不仅知道两个图中的某个加法节点在数学上等价，还知道它们处理的数据在物理分布和内存布局上的对应关系，这是实现精确验证的基石。

3. 工程实现：如何让理论驾驭千亿参数模型

将上述原理应用于Llama 3.1 405B这样的模型，最大的拦路虎是规模。一个完整的模型计算图可能有数百万个节点。直接构建全局E-Graph进行等式饱和，无异于自寻死路。Scalify的工程核心在于“分而治之”与“智能记忆”。

3.1 图分区与层记忆化：化整为零，避免爆炸

Scalify采用了一个两阶段的分区策略，其算法核心如论文中Algorithm 1所示。

第一阶段：按神经网络层分区这是最自然、最有效的切割点。现代Transformer模型结构高度规整，由重复的Attention层和FFN层堆叠而成。Scalify以层为边界，将整个计算图切割成一系列子图。这样做的好处是：

• 语义完整性：单层内的计算通常自包含，优化也大多发生在层内或相邻层之间。
• 对齐优化粒度：框架的许多优化（如算子融合）正是以层或更小的算子组为单位进行的。
• 实现简单：通过解析框架IR（如Torch FX的节点、JAX的jaxpr）可以较容易地识别层边界。

第二阶段：层内拓扑遍历与阶段划分即使单层，其计算图也可能很复杂。Scalify会对每一层的子图进行拓扑排序遍历，并进一步将其划分为多个“阶段”（Stage）。划分的原则是：一个阶段包含的所有节点，其输入都依赖于前序阶段已处理完成的节点。阶段之间通过“边界节点”（Boundary Nodes）连接。

如图5所示，每个阶段内部，通常包含可以并行处理的子图（例如，同一个矩阵乘法被分片到多个设备上独立计算的部分）。Scalify会为每个阶段启动多个并行线程，同时进行该阶段内的等式重写和关系推理。由于阶段间依赖清晰，这种并行是安全的。

层记忆化：避免重复劳动大模型的不同层往往结构相同或高度相似（例如，Transformer的所有中间层）。如果对每一层都独立进行完整的等式饱和分析，将是巨大的浪费。Scalify引入了指纹机制。它为每个分区后的子图（单设备版本和分布式版本）计算一个指纹，该指纹综合了图结构、算子类型、张量形状（符号化的）和初始关系等信息。

当处理到新的层时，Scalify先计算其指纹，并在内存中查找。如果找到匹配的指纹，它就直接复用之前已推导出的最终等价类和关系集合，跳过耗时的重写过程。这对于拥有数百个相同层的模型来说，带来了数量级的性能提升。

注意事项：分区与记忆化的代价分区和记忆化不是免费的午餐，它们引入了近似性。

1. 可能漏检跨层优化：有些优化可能涉及多个层（如跨层的激活重计算优化）。按层分区可能会切断这种跨层等价关系，导致本应等价的图无法被验证。在实践中，这类优化相对较少，且Scalify可以通过允许用户定义“跨层分区组”来缓解。
2. 指纹冲突风险：两个语义不同但结构巧合相同的子图可能产生相同指纹，导致错误地认为它们等价。Scalify的指纹设计需要足够精细，通常包含符号化的形状信息和初始关系，这大大降低了冲突概率，但理论上无法完全杜绝。这是一种用极低风险换取巨大性能收益的权衡。

3.2 可复用的重写模板与符号双射推断

重写模板为了平衡规则的通用性和性能，Scalify没有使用硬编码的、具体的重写规则，而是引入了参数化的重写模板。一个模板可以匹配一类操作模式。例如，一个处理“分片后归约”的模板可能如下：

template ShardedReduction(op, reduce_op, dim): pattern: reduce_op'( shard(x, dim, n), dim, n ) rewrite: reduce_op(x) condition: op is element-wise and commutative (e.g., add, max)

这个模板可以匹配任何在分片维度上进行归约的操作，只要该操作是满足交换律的。在匹配时，op、reduce_op、dim会被具体化。这极大地减少了需要维护的规则数量，同时保持了表达能力。

符号双射推断：破解布局迷宫布局异构性是验证中最棘手的问题之一。考虑论文中的经典Bug（图1）：为了将注意力输出转换为Batch-Sequence-Head格式，错误代码先reshape再（或需要）transpose，而正确代码需要不同的reshape/transpose序列。两个序列最终输出的张量形状完全相同，但元素排布不同，语义错误。

仅靠关系传播无法处理这种“殊途同归”但“归的不是同一个地方”的情况。Scalify的解决方案是符号双射推断。

1. 符号化张量：系统不为张量赋予具体数值，而是为每个张量的每个维度赋予一个符号轴（如i, j, k）。一个形状为(S*B, H)的张量，其元素可以表示为(s*b, h)，其中s和b是符号。
2. 追踪布局变换：每个reshape或transpose操作都被视为一个从输入符号轴到输出符号轴的映射函数。例如，reshape(x, (B, S, H))会将轴映射从(s*b, h)变为(b, s, h)（假设它能整除）。transpose(x, (1,0,2))会将轴(b, s, h)变为(s, b, h)。
3. 推断双射：当遇到一个复杂的、由多个reshape/transpose组成的序列时，Scalify会符号化地执行这个序列，推导出从原始张量每个元素到最终张量每个元素的精确映射关系。如果这个映射是一个双射（一一对应，且保持元素顺序的某种一致性），那么这两个布局变换序列在语义上是等价的。
4. 生成布局关系：一旦推断出双射，系统就生成一个layout(t1, t2, bijection_fn)关系事实，其中bijection_fn封装了这个符号映射。这个关系可以被后续的规则使用，来证明即使操作序列不同，两个子图在考虑布局变换后是等价的。

这个过程自动化地解决了图1中的BSH Bug。系统能推断出错误序列的映射与正确序列的映射不同，从而阻止两个输出节点合并到同一等价类，并精准定位到产生分歧的reshape操作。

4. 从验证到调试：如何将“不等价”转化为可行动的缺陷报告

验证出“不等价”只是第一步。对开发者来说，真正有价值的是：“bug在哪里？我该怎么修？” Scalify的缺陷定位机制是其从学术工具迈向工程实用的关键一步。

4.1 差异溯源与根因定位

当等式饱和过程停止，而两个目标输出节点仍未合并时，Scalify不会简单地报错。它会启动一个差异溯源过程：

1. 定位分歧点：从无法合并的输出节点开始，在E-Graph中反向遍历。系统会寻找这样一对节点：它们分别来自原始图和目标图，计算相似的功能，有相同（或应等价）的输入，但本身却无法被任何规则证明等价。这就是最初的“分歧点”。
2. 提取差异子图：以该分歧点为中心，向前（向输入方向）追溯若干步，提取出两个图中导致分歧的“差异子图”。这个子图通常很小，只包含几个到几十个节点。
3. 映射回源代码：这是最关键的一步。框架的IR（中间表示）通常保留了与源代码的映射信息（如Python字节码位置、操作符的创建栈）。Scalify利用这些调试信息，将差异子图中的每个问题节点，直接映射回生成该节点的框架源代码文件、函数、乃至行号。

最终，Scalify生成的不是一个抽象的“图不等价”报告，而是一个如论文图1所示的、具体的代码差分（Diff）提示：

--- src/attention.py (Buggy Version) if is_bsh: # (b * s, h) => (b, s, h) result = hlo.reshape(result, (n_seqs, n_active_tokens, hidden_size)) +++ src/attention.py (Corrected Version) if is_bsh: # (s * b, h) => (b, s, h) result = hlo.reshape(result, (n_active_tokens, n_seqs, hidden_size)) result = hlo.transpose(result, 0, 1)

它明确指出了在src/attention.py文件的某个函数中，一个reshape操作的参数顺序错了，并给出了正确的变换序列。开发者几乎可以“照抄”这个提示来修复Bug。

4.2 实战中的验证流程与集成

在实际生产环境中集成Scalify，通常遵循以下流程：

1. 测试用例生成：针对一个模型和一种特定的分布式配置（如TP=2， PP=4），编写一个简单的测试脚本。这个脚本会分别用“基线模式”（如单设备模拟）和“目标模式”（启用全部分布式优化）运行框架，并捕获其计算图IR。通常这需要框架提供图导出的接口。
2. 配置规则与注解：根据所使用的并行策略（数据并行、张量并行、流水线并行、序列并行等），加载或配置对应的重写规则模板和关系推理规则。对于框架自定义的、复杂的通信原语，可能需要添加手动注解来声明其语义关系。
3. 执行验证：将两套IR、规则和注解输入Scalify。验证过程在CPU上离线进行，无需GPU，也无需实际运行模型前向传播。
4. 分析报告：如果验证通过，则可以高度确信该优化配置对此模型是语义安全的。如果失败，则仔细审查Scalify输出的差异报告和代码定位信息。

实操心得：将验证嵌入CI/CD管道最有效的使用方式是将Scalify集成到框架的持续集成（CI）系统中。每当提交新的优化Pass或修改分布式策略时，CI管道可以针对一组代表性的模型（从几亿到几百亿参数）和配置运行Scalify验证。这能在代码合并前就拦截引入语义错误的变更，将静默错误扼杀在摇篮里。由于Scalify验证速度很快（分钟级），这种检查是可行的。我们在内部实践中，将其作为关键PR的必选检查项，显著提升了框架的可靠性。

5. 效果评估与局限性：它真的能解决所有问题吗？

根据论文评估，Scalify在单台普通服务器上，能在数分钟内完成对Llama-3.1-405B模型计算图的等价性验证。这证明了其卓越的可扩展性。更重要的是，它在Amazon的生产框架Transformers NeuronX和NeuronX Distributed中发现了5个此前未知的Bug，并复现了之前研究中17/19的已知Bug。这些Bug大多与张量布局变换、归约通信在特定条件下的错误有关，正是传统测试难以覆盖的角落。

然而，Scalify并非银弹，它有明确的适用范围和局限性：

1. 对数值不敏感：Scalify进行的是严格的符号等价验证。它无法处理因浮点数舍入顺序、非确定性算法（如某些dropout实现）或不同硬件内核实现带来的微小数值差异。如果两个计算图在数学上等价，但因浮点顺序导致最终结果有微小不同，Scalify会认为它们等价。这通常是可接受的，因为这种差异不被视为“错误”。
2. 依赖准确的规则和注解：验证的可靠性完全建立在重写规则和关系事实的正确性上。如果规则本身有误，或者编译器未能正确注入关系注解，验证结果可能出错。因此，维护一个正确、完备的规则库需要深厚的领域知识。
3. 无法验证算法本身的正确性：Scalify只验证“优化后的图是否与原始图语义等价”。如果原始模型代码本身就有Bug，或者原始的单设备执行逻辑就是错的，Scalify无法发现。它验证的是变换的正确性，而非模型的正确性。
4. 对极端动态形状的支持有限：虽然支持符号形状，但如果模型的控制流极度复杂，或张量形状在运行时高度动态且无法用符号表达，验证会变得困难甚至不可行。

尽管如此，Scalify代表了一种根本性的范式转变：从依赖脆弱数值比对的“黑盒测试”，转向基于形式化语义的“白盒验证”。它为大规模、高复杂性分布式机器学习框架的可靠性保障，提供了一套强大且实用的理论基础和工程实现。对于框架开发者而言，投入时间构建和维护这样一套验证基础设施，长远来看是降低调试成本、提升软件质量的关键投资。