Verifying Computational Graphs in Production-Grade Distributed Machine Learning Frameworks
This addresses the critical issue of debugging complex ML systems for practitioners, offering a scalable solution to prevent performance degradation in production environments.
The paper tackles the problem of silent errors in large-scale distributed ML frameworks by introducing Scalify, a lightweight verification framework that uses equality saturation and Datalog-style reasoning to check computational graphs, achieving verification of models like Llama-3.1-405B within minutes and exposing five unknown bugs in Amazon production frameworks.
Modern machine learning frameworks support very large models by incorporating parallelism and optimization techniques. Yet, these very techniques add new layers of complexity, introducing silent errors that severely degrade model performance. Existing solutions are either ad hoc or too costly for production. We present Scalify, a lightweight framework that exposes silent errors by verifying semantic equivalence of computational graphs using equality saturation and Datalog-style reasoning. To scale, Scalify partitions graphs with parallel rewriting and layer memoization, reuses rewrite templates, and augments equality saturation with relational reasoning and symbolic bijection inference. It further localizes discrepancies to precise code sites, turning verification results into actionable debugging guidance. Scalify verifies models as large as Llama-3.1-405B within minutes on a commodity machine and exposed five unknown bugs in Amazon production machine learning frameworks.