Accelerating Secure Machine Learning Training on GPUs With Pipeline Parallelism

The proliferation of data-driven machine learning (ML) applications makes privacy and data security risks increasingly prominent. Secure multi-party computation (MPC) offers a privacy-preserving ML method by enabling joint model training without data disclosure, but its reliance on complex cryptography adds significant delay and resource demands, challenging its practicality in large-scale applications. To address the low GPU utilization within the MPC framework, we propose an innovative privacy-preserving ML training optimization framework that exploits pipeline parallelism. Through a detailed analysis of the underlying principles of MPC-based training, we identify computation and communication as the primary bottlenecks for linear and non-linear computations, respectively. Drawing from traditional ML optimization strategies, we design a sub-network partitioning and pipeline parallelism method, specifically tailored for MPC training. This method not only allows for simultaneous training computations across different network layers, but also strategically overlaps computation with communication to enhance GPU utilization and reduce latency. Additionally, we develop a distributed communication mechanism to further improve communication efficiency. We integrate our framework into two distinct, state-of-the-art secure training frameworks: CryptGPU and Piranha. Compared to their original versions, our enhancements boost training speeds by up to 51%, significantly increasing GPU utilization, with negligible impact on model convergence speed and accuracy.