Mixed-Precision Communication-Avoiding SGD for Generalized Linear Models on GPUs

2026-06-16

AI Summary

This work addresses the communication bottleneck in distributed stochastic gradient descent (SGD) for training generalized linear models. The authors propose Communication-Avoiding SGD (CA-SGD), which reduces communication overhead by aggregating multiple AllReduce operations into a single Gram matrix exchange. They further introduce a mixed-precision strategy compatible across GPU generations, decomposing local rounding errors into nine distinct precision configurations to jointly optimize computation and communication. The method integrates BF16/FP32 mixed precision, GEMM acceleration, high-precision accumulation, and low-precision storage, fully leveraging NVIDIA’s tensor core hardware. Evaluated on an A100 GPU, CA-SGD achieves 5.1–6.8× speedup over FP32 SGD across multiple datasets while maintaining solution accuracy within 0.5% loss deviation.

Abstract

Distributed stochastic gradient descent (SGD) is limited by communication rather than computation, since each iteration requires an AllReduce across processes. Communication-avoiding SGD (CA-SGD) amortizes communication over $s$ iterations by replacing $s$ consecutive AllReduces with a single AllReduce of an $sb\times sb$ Gram matrix, trading more computation and bandwidth for fewer synchronization points. Modern GPUs with matrix hardware and reduced-precision formats offset this by accelerating the Gram GEMM and shrinking BF16 traffic. We study mixed-precision CA-SGD for generalized linear models on NVIDIA GPUs. Our finite-precision analysis decomposes the local rounding error of one CA-SGD outer iteration into nine independent precision choices, depending on the hardware only through its low-precision unit roundoffs, so the resulting recipes transfer in principle across GPU generations. The recipe stores the input matrix and margin vector in low precision, computes the Gram matrix from low-precision inputs with high-precision accumulation, communicates it in high precision, and performs the inner recurrence and weight updates in high precision. On NERSC Perlmutter A100 GPUs, mixed-precision CA-SGD matches FP32 SGD loss within $0.5\%$ on logistic, linear, and Poisson problems and reaches $5.1$--$6.8\times$ speedup over FP32 SGD on epsilon, SUSY, HIGGS, synth, and Poisson-synth. Our software is available at https://doi.org/10.5281/zenodo.20448273

Problem

• mixed-precision
• communication-avoiding SGD
• generalized linear models
• GPU acceleration
• distributed optimization

Innovation

• mixed-precision
• communication-avoiding SGD
• Gram matrix
• generalized linear models
• GPU acceleration