CrypTorch: PyTorch-based Auto-tuning Compiler for Machine Learning with Multi-party Computation

TL;DR

CrypTorch is a PyTorch extension that auto-tunes approximations for MPC-based machine learning, significantly improving performance while maintaining accuracy, thus enabling more efficient privacy-preserving ML computations.

Contribution

It introduces a compiler that disentangles and optimizes approximations in MPC ML, providing automatic selection for better performance and accuracy, integrated into PyTorch 2.

Findings

Auto-tuning yields 1.20--1.7× speedup without accuracy loss.

Allowing some accuracy degradation increases speedup to 1.31--1.8×.

Overall framework achieves 3.22--8.6× end-to-end speedup over CrypTen.

Abstract

Machine learning (ML) involves private data and proprietary model parameters. MPC-based ML allows multiple parties to collaboratively run an ML workload without sharing their private data or model parameters using multi-party computing (MPC). Because MPC cannot natively run ML operations such as Softmax or GELU, existing frameworks use different approximations. Our study shows that, on a well-optimized framework, these approximations often become the dominating bottleneck. Popular approximations are often insufficiently accurate or unnecessarily slow, and these issues are hard to identify and fix in existing frameworks. To tackle this issue, we propose a compiler for MPC-based ML, CrypTorch. CrypTorch disentangles these approximations with the rest of the MPC runtime, allows easily adding new approximations through its programming interface, and automatically selects approximations to maximize both performance and accuracy. Built as an extension to PyTorch 2's compiler, we show that CrypTorch's auto-tuning alone provides 1.20--1.7$\times$ immediate speedup without sacrificing accuracy, and 1.31--1.8$\times$ speedup when some accuracy degradation is allowed, compared to our well-optimized baseline. Combined with better engineering and adoption of state-of-the-art practices, the entire framework brings 3.22--8.6$\times$ end-to-end speedup compared to the popular framework, CrypTen.