Entangle: Formal Verification for Distributed Deep Learning
It Takes Two to Entangle — ASPLOS 2026
* Equal contribution. New York University and ByteDance Seed.
Entangle is a static verification system for detecting correctness bugs in distributed deep-learning implementations. It checks whether a distributed implementation correctly refines its sequential reference model by inferring relationships between tensors in their computation graphs. Entangle uses iterative rewriting and equality saturation to reconstruct sequential outputs from distributed outputs, scales to large neural-network models, and provides actionable information for bug localization.
Published in the Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, pp. 1022–1039.
Paper abstract
Distributed machine learning training and inference is common today because today’s large models require more memory and compute than can be provided by a single GPU. Distributed models are generally produced by programmers who take a sequential model specification and apply several distribution strategies to distribute state and computation across GPUs. Unfortunately, bugs can be introduced in the process, and a distributed model implementation’s outputs might differ from the sequential model’s outputs. In this paper, we describe an approach to statically identify such bugs by checking model refinement, that is, can the sequential model’s outputs be reconstructed from the distributed model’s outputs? Our approach, implemented in Entangle, uses iterative rewriting to prove model refinement. Our approach can scale to today’s large models and deployments: we evaluate it using GPT and Llama-3. Further, it provides actionable outputs that aids in bug localization.
What is Entangle?
Entangle brings formal verification to distributed machine learning. Given sequential and distributed computation graphs and a user-provided relation between their input tensors, it checks model refinement: whether every sequential output can be reconstructed by rearranging or combining the distributed outputs. This establishes that the distributed implementation preserves the reference model’s semantics under the supplied input relation, even when tensors are partitioned across devices.
What problem does Entangle solve?
Distributing a neural network requires partitioning model state and computation and adding communication operations. Incorrect loss scaling, incompatible sharding configurations, or missing collective operations can change the model’s results. Some of these correctness bugs preserve tensor shapes, making them difficult to detect with shape checks alone.
For developers asking whether a distributed or tensor-parallel deep-learning implementation preserves a sequential reference, Entangle checks distributed model correctness before deployment. Its inferred tensor relationships help localize bugs to the part of the computation where refinement fails.
How does Entangle work?
Entangle traverses the sequential graph in topological order, processing one operator at a time. It uses already inferred tensor relations, rewrite lemmas, and the distributed graph to infer relations for the operator’s outputs. Equality saturation with e-graphs, implemented using egg, represents equivalent expressions during this iterative rewriting process.
The goal is a complete clean output relation: a mapping for every sequential output using tensor rearrangements and reductions that combine distributed tensors. This makes sequential versus distributed implementation equivalence precise through model refinement. Processing individual operators keeps the search manageable and produces intermediate relations useful for bug detection and localization.
The paper establishes soundness of the refinement check, but the method is not complete: a correct implementation can trigger a false alarm if the method’s assumptions are violated. Section 3.3 describes these assumptions, including matching optimizations and operator order between the two graphs.
Interactive relation inference
Step through Entangle’s relation inference, from the input relation Rᵢ to the final output relation Rₒ. The first tab follows the matrix multiplication and subtraction example in Figures 1–2 and the iterative exploration in Section 4.3.1. Each step advances an exploration iteration for the current operator in Gₛ, followed by the output or failure check.
In the correct example, C = A × B and F = C − E. The distributed graph splits the contraction dimension of A and B across two ranks, then uses reduce-scatter to sum the partial products and split the resulting rows.
The second tab illustrates Bug 4: incompatible configurations for model components from Appendix A.1. A ByteDance MoE model used sequence parallelism, but expert weights were mistakenly sharded instead of replicated. For the sequential computation X × A × B, the distributed ranks computed X₁ × A₁ × B₁ and X₂ × A₂ × B₂. Output shapes still matched, but required off-diagonal products were missing. Entangle detected the bug at the first matrix multiplication.
On narrow screens, scroll each diagram sideways to see both ranks.
Relations found
… Equivalent rewritings are omitted; the useful clean relations are shown.
Rewriting in this iteration
What kinds of systems does it apply to?
The paper evaluates Entangle on distributed LLM systems and other neural-network implementations: GPT-2 in Megatron-LM, Qwen2 in vLLM, Llama-3 in Transformers-Neuron, a Hugging Face regression model, and a proprietary ByteDance model. The evaluated strategies include tensor parallelism, sequence parallelism, expert parallelism, and gradient accumulation.
Entangle analyzes computation graphs captured from frameworks such as PyTorch using TorchDynamo, and the evaluation also includes HLO graphs from AWS NeuronX. The study checks forward passes for the open-source models and both forward and backward passes for the ByteDance model. Applying it to additional frameworks or operators may require graph converters and rewrite lemmas; see the code repository’s guide to verifying your own models.
Cite It Takes Two to Entangle
Zhanghan Wang, Ding Ding, Hang Zhu, Haibin Lin, and Aurojit Panda. 2026. It Takes Two to Entangle. ASPLOS 2026, Volume 2, 1022–1039. https://doi.org/10.1145/3779212.3790178
@inproceedings{wang2026entangle,
author = {Wang, Zhanghan and Ding, Ding and Zhu, Hang and Lin, Haibin and Panda, Aurojit},
title = {It Takes Two to Entangle},
year = {2026},
isbn = {9798400723599},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3779212.3790178},
doi = {10.1145/3779212.3790178},
abstract = {Distributed machine learning training and inference is common today because today's large models require more memory and compute than can be provided by a single GPU. Distributed models are generally produced by programmers who take a sequential model specification and apply several distribution strategies to distribute state and computation across GPUs. Unfortunately, bugs can be introduced in the process, and a distributed model implementation's outputs might differ from the sequential model's outputs. In this paper, we describe an approach to statically identify such bugs by checking model refinement, that is, can the sequential model's outputs be reconstructed from the distributed model's outputs? Our approach, implemented in Entangle, uses iterative rewriting to prove model refinement. Our approach can scale to today's large models and deployments: we evaluate it using GPT and Llama-3. Further, it provides actionable outputs that aids in bug localization.},
booktitle = {Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2},
pages = {1022–1039},
numpages = {18},
keywords = {formal verification, distributed deep learning, equality saturation, bug localization},
location = {Pittsburgh, PA, USA},
series = {ASPLOS '26}
}