PASO: Step Parallel Stochastic Optimization

This paper approaches the fundamental challenge of accelerating the inherently autoregressive nature of gradient descent (GD) like SGD and Adam through a dynamic system perspective. Specifically, we introduce a unified framework that recasts the autoregressive GD process as solving a system of triangular nonlinear equations (TNEs), thereby enabling \textit{step-parallel} training, where gradients for different GD steps are computed concurrently without sequential dependencies. Within this generic framework, we establish that: (1) the TNE system admits a unique solution corresponding precisely to the autoregressive GD iterative trajectory; (2) solving the TNEs system guarantees convergence to the GD iterative trajectory in at most the equal iterations. Building on these insights, we present \textit{PASO}, the first step-parallel optimizer for accelerating a broad class of GD-based optimizers like SGD and Adam. Extensive experiments (\textit{e.g.}, Llama-3.2-1B and diffusion model) validate that PASO achieves up to \textbf{21}$\times$ reduction in GD steps and \textbf{4.5}$\times$ speedup in wall-clock time, with no model quality loss. Source code is available at: \url{https://github.com/Jianrong-Lu/PASO.git}.

The Core Problem: The Waiting Game in AI Training Training large AI models can take months. This is largely because the standard training process is strictly sequential. The system must finish calculating step one before it can begin step two, creating a massive bottleneck that slows down the entire operation. --- ### The Innovation: PASO PASO (Step Parallel Stochastic Optimization) introduces a new way to break this sequential chain. Instead of waiting for previous steps to finish, PASO translates the training process into a system of mathematical equations. By solving these equations simultaneously, multiple training steps can be calculated at the exact same time across different computer processors. --- ### Key Benefits * Massive Speedups: It reduces the required sequential training steps by up to 21x and speeds up actual real-world training time by up to 4.5x. * Comparable Quality: The final AI model matches one trained the slow, traditional way. * Hardware Efficient: It does not demand any additional memory per processor. * Highly Compatible: It can be seamlessly combined with other existing software systems designed to speed up AI training.