Mixed-Precision Communication-Avoiding SGD for Generalized Linear Models on GPUs Aditya Devarakonda1 , Irene Simó Muñoz2 , and Giulia Guidi2
arXiv:2606.18463v1 [cs.DC] 16 Jun 2026
1
Department of Computer Science, Wake Forest University, Winston-Salem, NC, USA, [email protected] 2 Department of Computer Science, Cornell University, Ithaca, NY, USA, {is449, gguidi}@cornell.edu
Abstract Distributed stochastic gradient descent (SGD) is limited by communication rather than computation, since each iteration requires an AllReduce across processes. Communicationavoiding SGD (CA-SGD) amortizes communication over s iterations by replacing s consecutive AllReduces with a single AllReduce of an sb × sb Gram matrix, trading more computation and bandwidth for fewer synchronization points. Modern GPUs with matrix hardware and reducedprecision 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 finiteprecision 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× speedup over FP32 SGD on epsilon, SUSY, HIGGS, synth, and Poisson-synth. Our software is available at https://doi.org/10.5281/zenodo. 20448273
Keywords: Communication-avoiding algorithms, mixed precision, stochastic gradient descent, generalized linear models, GPU, NCCL, BF16
1
Introduction
Distributed mini-batch stochastic gradient descent (SGD) is one of the most widely used optimization methods for large-scale empirical risk minimization. Each iteration of SGD requires a synchronizing AllReduce across processes to aggregate partial gradients into a single solution update vector [2]. In the 1D block-column setting used here, the AllReduce synchronization costs 2α log P + β W · 2(P − 1)/P , where the latency term α log P dominates for the small message sizes typical of mini-batch SGD [11, 30]. Because local computation scales with P but communication latency does not, scalable mini-batch SGD requires algorithmic changes that reduce synchronization frequency.
1
Communication-avoiding SGD Communication-avoiding SGD (CA-SGD) addresses this gap by amortizing communication over s consecutive mini-batch iterations [14, 15]. Rather than performing s separate vector AllReduces of length b, CA-SGD samples a single outer block of sb rows (Y ∈ Rsb×nloc ), where nloc is the length of the locally stored features. CA-SGD then computes the local margin contribution r = Y x ∈ Rsb and the local Gram matrix G = Y Y ⊤ ∈ Rsb×sb , after which a single grouped AllReduce is performed to obtain the global margin and Gram matrix. CA-SGD then performs s inner residual-correction updates independently and redundantly on each process without further communication. The reformulation trades fewer synchronization points for a larger Gram GEMM and a larger grouped AllReduce. On CPU implementations, this trade-off limits the useful range of s because the Gram GEMM and grouped AllReduce payload grow as (sb)2 . Tensor-core GPUs change this balance by accelerating the Gram GEMM and reducing traffic with low-precision storage. The performance opportunity does not imply a uniform BF16 recipe, however, because the outer iteration combines operations with different rounding-error sensitivities. A single outer CA-SGD iteration involves nine distinct operations (Section 3.3) whose roundingerror tolerances vary by orders of magnitude: the Gram GEMM and margin GEMV are lengthnloc inner products whose accumulator precision sets the leading error, the AllReduced quantities carry a collective factor C(P ) set by the reduction topology and summation order, and the inner correction sum of length up to (s−1)b requires qu < 1 for the Higham factor and therefore FP32 storage. Uniform low precision provides insufficient precision for error-sensitive operations. Our work adapts this finite-precision program to stochastic s-step methods by bounding the local error of one CA-SGD outer iteration and lifting the resulting perturbation into a boundedregime convergence statement. We focus on generalized linear models, including linear, logistic, and Poisson regression in the empirical study. Because Poisson regression has an unbounded inverse-link derivative, the theorem-covered loss family is the Lipschitz-residual subset instantiated by logistic and linear regression. Contributions. This paper makes four contributions. (C1) It presents the first finite-precision analysis of CA-SGD (Theorem 5.7), where the local forward error of one outer iteration decomposes into nine additive terms. Each term depends on a single precision choice and can be computed from the algorithmic parameters and the bounds in Assumptions 5.1 to 5.5. The decomposition exposes a constraint: the inner s-step correction sum must be performed in FP32, which a uniform lowprecision choice would overlook. (C2) A concrete A100 mixed-precision recipe (Table 3) requires BF16 storage and BF16-input/FP32-accumulate tensor-core math for the Gram GEMM, margin, and outer-gradient GEMVs. It retains FP32 for the inner correction, residual, master weights, and Gram AllReduce, and uses BF16 for the margin AllReduce. The recipe is based directly on the nine per-choice coefficients α∗ rather than selected through experimentation. (C3) A perturbed-SGD convergence guarantee (Theorem 5.9) gives, in a bounded regime with local strong convexity and smoothness, a deterministic radius ρ around the exact SGD trajectory. (C4) In empirical validation on NERSC Perlmutter A100 GPUs with up to P = 256 GPUs, Recipe C matches the FP32 SGD baseline within 0.5% and reaches up to ∼16× speedup over FP32 SGD on synth at (b, s) = (1, 256) and 5.2–6.2× speedup over FP32 SGD on epsilon, SUSY, HIGGS, synth, and Poisson-synth at (b, s) = (8, 64). The Recipe D variant, which casts only the Gram AllReduce to BF16, reaches up to 6.8× on these datasets but trades a larger relative loss gap for the speedup and falls outside the deterministic budget that the theory verifies for Recipe C. The Recipe C Gram GEMM reaches 92% of the A100 BF16 tensor-core peak.
2
2
Background and Related Work
2.1
Generalized linear models and distributed mini-batch SGD
m×n and labels y minimizes the empirical A generalized linear Pmmodel⊤ (GLM) for features A ∈ R −1 risk F (x) = m i=1 L(ai x, yi ), with L depending on the weights only through the margin ri = a⊤i x [24]. We focus on three canonical losses: logistic (L(r, y) = log(1 + e−yr ), Lσ = 1/4), linear (L(r, y) = (r − y)2 /2, Lσ = 1), and Poisson (L(r, y) = er − yr, Lσ unbounded). The gradient is ∇F (x) = −m−1 A⊤ δ, where the residual is δi = −∂L/∂r (ri , yi ). Both logistic and linear models have Lipschitz residuals. The logistic residual is globally bounded, and the linear residual is bounded in the bounded-margin regimes used in Assumption 5.2. Distributed mini-batch SGD samples I ⊆ {1, . . . , m} of size b and computes gI (x) = −1 −b YI⊤ δ(YI x, yI ). Under the bounded-trajectory, local-curvature model of Assumptions 5.5 and 5.6 and for η ≤ 1/LF , gI is unbiased with variance σg2 /b, and the iterate satisfies the canonical contraction [1, 2]. In this work, we use 1D block-column partitioning (nloc = n/P columns per rank, Section 4.2), which requires a single vector AllReduce of length b for each mini-batch step. The 2D and hybrid variants of [15, 31] further reduce communication volume but still require one synchronization step per mini-batch and are complementary to the s-step amortization used here.
2.2
Communication-avoiding and s-step methods
The s-step approach [10, 12, 13, 22] replaces s communication-blocking linear algebra primitives with a single grouped collective on a larger composite object, followed by s local steps of the original recurrence. CA-SGD [14, 15, 31] applies this to mini-batch SGD: sample sb rows, compute the rank-(sb) Gram matrix G = Y Y ⊤ and the base margin r = Y x in parallel, fuse them into one grouped AllReduce, then run s local residual-correction updates z (j) = rj + (η/b) Gj,<j δ<j on the inner iterate before forming the outer gradient g = Y ⊤ δhist and updating the weights. The per-block cost (Table 2) replaces s vector AllReduces with one grouped AllReduce, at the cost of a Gram GEMM, an (sb)2 -word payload, and an O(b2 s2 ) inner correction sum. On CPUs, the γ(sb)2 nloc term breaks the trade-off at modest s. On tensor-core GPUs, this term is suppressed by up to 16× in compute peak and 2× in traffic, so the crossover shifts to smaller s (Table 6 and Figure 5). s-step error analysis The numerical stability of s-step Krylov methods has been studied extensively. The residual-replacement analysis of s-step CG and BiCG [4] identifies the basis condition number κ(Vs ) as the main amplifier of the gap between the recursively updated and true residual. This finite-precision analysis of classical Lanczos was extended to s-step Lanczos [5] and led to adaptive techniques [3] that stabilize the algorithm at runtime. Recent work by Carson, Gergelits, and Yamazaki [9] extends prior work to non-uniform floating-point storage choices within the same error analysis framework. Here, we extend this line of finite-precision analysis to s-step SGD by bounding the per-outer-iteration error and proving that the resulting perturbation remains within the bounded-regime SGD convergence neighborhood of Theorem 5.9.
2.3
Mixed-precision arithmetic and finite-precision analysis on modern GPUs
The NVIDIA A100 exposes FP32, TF32, FP16, and BF16 with unit roundoffs ≈ 6×10−8 , 5×10−4 , 5 × 10−4 , and 4 × 10−3 , respectively [26]. TF32 and BF16 share the FP32 exponent range but have truncated mantissas, while FP16 uses the narrower IEEE half-precision range. Tensor cores accept BF16 or FP16 inputs with FP32 accumulation at 312 TFLOP/s for both formats, and A100 also supports FP16-input/FP16-accumulate HMMA for Recipe F. The dominant approach 3
(u )
uses low-precision inputs with higher-precision accumulation, with Higham factor 2u16 + γn 32 (Equation (8)) in the BF16-input/FP32-accumulate rows of Theorem 5.7. (u) The deterministic Higham factor γq = qu/(1 − qu) used throughout Section 5 is standard [16, √ 18]. The Higham–Mary probabilistic refinement [19, 20] would substitute λ(ρ) u q, but this is left for future work. The closest stylistic precedents for our recipe table (Table 3) are the threeprecision iterative-refinement analyses by Carson and Higham [6–8]. We have nine slots instead of three because the CA-SGD outer iteration includes more kernel categories. A parallel study on GPU tensor-core mixed precision for dense linear system solvers [7, 17] shows that aggressive low-precision factorization preserves FP64 backward stability when the refinement loop is wellconditioned. Higham et al. [21] provide a software simulation methodology to inform effective roundoff configurations.
2.4
Mixed-precision machine-learning training
Mixed precision in machine learning has primarily been motivated by deep learning. Micikevicius et al. [25] introduced FP32 master weights with FP16 forward and backward computations and loss scaling. Kalamkar et al. [23] showed that BF16’s wider range can eliminate loss scaling in most training runs. Both approaches have been integrated into production software and frameworks such as PyTorch AMP, NVIDIA Apex, and cuBLAS Tensor Op. This work develops a finite-precision analysis for CA-SGD and uses it to select a mixed-precision recipe, with its accuracy and speedup measured in Section 7.
3
Notation and Finite-Precision Model
We fix the notation, finite-precision model, and storage-precision mapping that the algorithms of Section 4 and the error bounds of Section 5 both depend on.
3.1
Algorithmic quantities
Let A ∈ Rm×n be the row-standardized data matrix and y ∈ Rm the labels ({−1, +1}m for logistic, sb×n be the sb-row subsampled matrix drawn Rm for linear, Zm ≥0 for Poisson). Let Y = SA ∈ R during one outer iteration, where S ∈ {0, 1}sb×m is a sampling matrix with exactly one non-zero per row that selects a row of A. Rows of A may be selected multiple times among the sb outer samples (i.i.d. uniform sampling with replacement, see Assumption 5.6 and Algorithm 1), so a column of S may contain zero, one, or several nonzeros. Let yI = Sy ∈ Rsb be the corresponding label vector, and x ∈ Rn the weight iterate, with x(h) denoting its value after h outer iterations. The CA-SGD recurrence also maintains the Gram block G = Y Y ⊤ ∈ Rsb×sb , the base margin r = Y x ∈ Rsb , the per-sample residual δ(r, y) = −∂L/∂r ∈ Rsb (the negative loss gradient with respect to the margin), and an inner iterate that propagates through s residual corrections per outer iteration. The quantity g := Y ⊤ δ ∈ Rn assembled by the outer gradient GEMV is the (sample-summed) descent direction. The sample-mean true gradient estimator used in Assumption 5.6 and Theorem 5.9 is b ∇xL(x) = −b−1 YI⊤ δ = −b−1 gI , and the weight update x e ← x e + (η/b) g in Algorithm 1 is the b standard SGD step x ← x − η ∇xL(x) written in δ-form. We use the unqualified term gradient to refer to g when discussing the kernel and its precision slot ug (defined in Section 3.3), with the understanding that the sign and b−1 normalization are absorbed into the step size. We refer to δ as the residual, since δ is an intermediate per-sample quantity in Rsb whose storage precision is tracked separately (ur , uσ , and uc ). The residual is loss-specific with the following definitions: δi = yi (1 − σ(yi ri )) for logistic regression (Lσ = 1/4), δi = yi − ri for linear regression (Lσ = 1), 4
and δi = yi − exp(ri ) for Poisson regression. Poisson regression has an unbounded Lipschitz value (Lσ = ∞), which places this loss outside the analysis. We report its empirical behavior as boundarystress evidence. We use b for the mini-batch size, s for the CA inner-loop length, P for the number of ranks (e.g., GPUs or CPUs), and nloc = n/P for the number of local features stored per rank.
3.2
Finite-precision model
Here, we work under the standard IEEE-754 rounding model: every elementary floating-point operation op ∈ {+, −, ×, ÷} evaluated in a format with unit roundoff u satisfies flu (x op y) = (x op y)(1 + θ) with |θ| ≤ u [18]. For a stored vector or matrix z in unit-roundoff u∗ , we write ze ∈ flu∗ (R· ) for ze = z+e with |e| ≤ u∗ |z| componentwise, which yields ∥e∥ ≤ u∗ ∥z∥ for any vector pnorm or matrix Frobenius norm. Composite kernels (GEMV, GEMM, AllReduce, sigmoid) inherit Higham-style entrywise bounds that add to O(u∗ ), and second-order terms O(u2∗ ) are dropped. The analysis in Section 5 accumulates rounding error along inner-product reductions of various lengths q: the margin and Gram kernels with q = nloc , the outer gradient with q = sb, and the inner correction block at iteration j with q = jb. The standard deterministic Higham bound for the floating-point inner product of two length-q vectors x, y accumulated in precision u is flu (x⊤y) − x⊤y ≤ γq(u) |x|⊤ |y| ,
γq(u) :=
qu = q u + O((qu)2 ), 1−qu
(1)
(u)
valid whenever qu < 1 [18]. We use γq throughout Section 5 for all GEMV/GEMM reductions. The AllReduce constant C(P ) of Assumption 5.4 is an empirical hardware parameter depending on the reduction order and routing algorithm (e.g. tree vs. ring). It appears as C(P ) u∗ when bounding error associated with communicated quantities. For low-precision-input/FP32-accumulate kernels (u) we replace γq with the input-cast plus FP32-accumulation factor. The BF16 specialization is √ Equation (8). The Higham–Mary probabilistic-rounding model [19, 20] would substitute λ(ρ) u∗ q for the deterministic qu factor, which we leave for future work.
3.3
The nine unit roundoffs
A single outer iteration of CA-SGD executes nine distinct operations, each with its precision tracked separately. The first-order analysis of Theorem 5.7 separates the leading finite-precision error among these nine precision choices, allowing each precision to be independently selected based on the desired error budget. The nine symbols are uA for storage of A (and of the gathered subsampled matrix Y ), uG for the Gram GEMM G = Y Y ⊤ , ur for the base margin GEMV r = Y x, uc for the inner-loop correction sum (the CA-SGD correction recurrence), uσ for the elementwise nonlinearity kernel (e.g., sigmoid for logistic regression), ug for the outer gradient GEMV g = Y ⊤ δ, uAR,r for the margin AllReduce datatype, uAR,G for the Gram AllReduce datatype, and ux for weight storage and update. These nine subscripts appear in the algorithms of Section 4 and in every error term E∗ of Section 5.
3.4
Storage-precision mapping
Table 1 lists the quantities referenced by the algorithm floats of Section 4, along with their storageprecision variable. We use ze for the finite-precision quantity associated with an exact-arithmetic quantity z. For example, Ye ∈ fluA (Rsb×n ) is the finite-precision version (defined by uA ) of the sampled submatrix Y . Wherever an algorithm or theorem statement uses e·, the storage precision is the one listed in Table 1. 5
Table 1: Storage-precision mapping for the quantities in one outer CA-SGD iteration. The middle column gives the storage-precision variable. The right column gives the concrete A100 realization in our experiments. Labels y are stored independently of A at FP32 throughout. Symbol
Quantity
Storage precision
Typical format
e Ye A, e G re δe y, yI ge x e -
feature matrix, subsampled rows Gram block G = Y Y ⊤ base margin Y x residual labels, sampled labels outer gradient Y ⊤ δ weights margin / Gram AllReduce
uA uG ur uc , uσ FP32 ug ux uAR,r , uAR,G
BF16 BF16 in; FP32 accum/out BF16 in; FP32 accum; BF16 out FP32 FP32 BF16 in; FP32 accum/out FP32 BF16, FP32
sample sb rows, gather Ye ∈ fluA (Rsb×nloc ) e = flu (Ye Ye⊤ ) margin re = flur (Ye x e) ∥ Gram G G e at (uAR,r , uAR,G ) grouped AllReduce {e r, G} cast δe(0) ← re : uAR,r → uc (omitted if uAR,r = uc ) inner s-step loop at fluc , j = 0, . . . , s−1: e j,<j δe<j ; δej ← residual at uσ δej += (η/b) G cast δe → flug : uc → ug (omitted if uc = ug ) e x e at ux update x e += (η/b) flug (Ye⊤ δ);
Figure 1: Mixed-precision CA-SGD for one outer iteration on each rank. Each kernel and storage box is labeled by its precision slot from Table 1. Dashed boxes are precision casts before and after the inner s-step loop (omitted when input/output precisions agree). Recipe C (Table 3) instantiates uA = ur = uG = uAR,r = ug = uBF16 and uc = uσ = uAR,G = ux = uFP32 .
4
Algorithm Design
We present the mixed-precision CA-SGD outer iteration that the analysis of Section 5 tracks one precision slot at a time, then state its communication cost and place it on the roofline. The FP32 baseline is the special case of the algorithm in which every precision slot u∗ of Table 1 is set to uFP32 .
4.1
Mixed-precision algorithm
Algorithm 1 states the mixed-precision CA-SGD outer iteration. Every assignment that rounds in finite precision is wrapped in flu∗ (·), where the subscript u∗ names the slot that carries the rounding error. Each numbered line corresponds to one of the nine E∗ terms that appear in Theorem 5.7.
6
Algorithm 1: Mixed-precision CA-SGD outer iteration. ep ∈ flu (Rm×nloc ) (rank-p block-column feature slab, BF16); labels y ∈ Rm stored Input: A A in FP32; weights x e ∈ flux (Rnloc ); step size η, batch size b, CA parameter s. Output: Updated x e ∈ flux (Rnloc ). 1 Sample sb batch indices I = I0 ∪ · · · ∪ Is−1 , |I| = sb, drawn i.i.d. uniformly with
replacement from {0, . . . , m − 1} eI ∈ flu (Rsb×nloc ) 2 Gather sampled rows Y A e 3 r eloc ← flur YeI x e loc ← flu YeI Ye ⊤ 4 G G I e ← AllReduce reloc ; G e loc ; uAR,r , uAR,G 5 r e; G
// base margin GEMV; uA → ur // Gram GEMM; uA → uG // grouped NCCL, split datatype
ehist ← [ ] 6 δ 7 for j ← 0 to s − 1 do 8 9 10 11 12 13 14
if j = 0 then ze(0) ← re0 else e j,<j δehist ze(j) ← fluc rej + ηb G end δe(j) ← fluσ δ(e z (j) , yIj ) δehist ← δehist ; δe(j)
// corrected margin at uc // per-batch residual at uσ
15 end 16 g eloc ← flug YeI⊤ δehist
17 x e ← flux
4.2
x e + ηb geloc
// outer gradient GEMV // weight update
Communication cost and roofline
To model communication cost, we use the standard α-β-γ model for parallel runtime. Per rank, the runtime is T = γF + βW + αL, where F is the number of flops, W is the number of words, and L is the number of messages. In this model, mixed precision is used in two places. The word size is reduced from FP32 to BF16, resulting in βBF16 = βFP32 /2. The compute constant also decreases, with γBF16 up to 16× smaller than γFP32 on tensor cores, and we assume 1D blockcolumn partitioning across P ranks, with nloc = n/P columns per rank. The AllReduce cost under this partitioning is LAR (W ) = 2α log P + β W · 2(PP−1) . (2) Table 2 compares the cost of s classical mini-batch SGD iterations against the cost of one outer CASGD iteration. CA-SGD trades s small AllReduces for one grouped AllReduce of the pair (r, G). This latency improvement comes at the cost of (sb)2 nloc additional Gram flops and an (sb)2 -word increase in the AllReduce message size. The two costs have a crossover at α(s−1)·2 log P = γ 2(sb)2 nloc +b2 s(s−1) +β(sb)2 ·2(P −1)/P . Mixed precision shrinks the right-hand-side Gram terms by decreasing message size and by increasing computation throughput through the use of tensor cores. The inner-correction term γb2 s(s − 1) stays unchanged because we keep this kernel at FP32 throughout (Table 3) as specified by the analysis. Therefore, in the case where the Gram terms dominate the right-hand side, the crossover moves to smaller s under mixed precision.
7
Table 2: α-β-γ costs of s vanilla SGD iterations versus one outer CA-SGD iteration under 1D block-column partitioning with AllReduce cost (2). Cost component SGD (s iters) CA-SGD (one outer iter) Flops F Words W Messages L
5
4sb nloc sb s
2(sb)2 nloc + b2 s(s − 1) + 4sb nloc sb + (sb)2 1
Finite-Precision Analysis
This section presents the finite-precision error analysis of CA-SGD under the model described in Section 3. First, we bound the local forward error of a single outer CA-SGD iteration for the nine precision choices specified in Section 3.3. Uniform precision and Recipe C then appear as special cases of this bound. Finally, a perturbed-SGD argument raises the per-block error budget to a convergence guarantee—convergence to a neighborhood of the optimum, valid as long as the iterates remain within the bounded region XR . Our analysis covers the family of Lipschitz-residual GLMs, which includes logistic regression, linear regression, and any member whose inverse-link derivative is globally Lipschitz. Poisson regression lies outside this family because its inverse-link derivative is unbounded (Assumption 5.2). In addition, we include Poisson regression as empirical boundary-stress evidence beyond the loss family covered by the theorem.
5.1
Assumptions
Assumption 5.1√(Data normalization). Every row of A satisfies ∥ai ∥2 ≤ 1, which yields ∥Y ∥F ≤ √ sb and ∥Y ∥2 ≤ sb for every sampled block Y ∈ Rsb×n . Assumption 5.2 (Lipschitz and bounded residual). Let XR = {x : ∥x∥2 ≤ R} and MR = {a⊤ i x: i ≤ m, x ∈ XR }. The residual δ(r, y) = −∂L/∂r is separable componentwise and globally Lσ Lipschitz in r. For the bounded-margin regime used to instantiate the local error budget, the residual also satisfies |δ| ≤ Dδ and |flu (δ) − δ| ≤ Cδ uDδ on the exact and finite-precision margins encountered in the local comparison. Logistic and linear regression satisfy this with (Lσ , Dδ ) = (1/4, 1) and (1, max |yi | + R) respectively. Assumption 5.3 (Bounded Gram). ∥G∥F ≤ sb and ∥GL ∥F ≤ sb for every sampled G = Y Y ⊤ and its lower-triangular part GL (deterministic under Assumption 5.1). P Assumption 5.4 (AllReduce model). For any AllReduce summing rank-local vectors v = p vp P into collective unit roundoff u, the returned vb satisfies |b v − v| ≤ C(P ) u p |vp | + O(u2 ). The constant C(P ) depends on processor count, topology, and the sign structure of the reduction. C(P ) absorbs the in-network reduction rounding and any input-cast rounding that occurs when the rank-local accumulator is finer than the collective datatype. The bounds below hold for any finite C(P ). When evaluating constants numerically we use the worst-case trace bound C(P ) ≲ 2.4P . The direct NCCL measurements of Figure 2 reach this worst case (C(256) ≈ 6×102 ), while their empirical power-law fit 4.62P 0.56 grows more slowly (to ≈ 102 at P = 256), and both are datatypeindependent to within 4%. Assumption 5.5 (Bounded-trajectory regime). All exact and finite-precision iterates and exact mini-batch substeps considered in the local and global comparisons remain in XR .
8
C(P ) observed (max, L ≥ 131072)
FP32 BF16 FP16 0.56 4.62 P (empirical fit, signed) 3 + log2 (P/4) (prior assumption)
600
400
200
0 4
8
16
32
64
128
256
P (number of ranks)
Figure 2: AllReduce constant C(P ) measured directly on Perlmutter NCCL collectives versus rank count P , summing rank-local vectors of length L ≥ 131072 against an FP64-accumulated reference. The observed factor reaches the worst-case C(P ) ≲ 2.4P and is datatype-independent to within 4% across BF16, FP16, and FP32. The dotted curve is the power-law fit 4.62P 0.56 and the dashed curve is the prior 3 + log2 (P/4) assumption. P Assumption 5.6 (Optimization model). F (x) = m−1 i L(a⊤ i x, yi ) is µF -strongly convex and LF smooth on XR , with unique unconstrained minimizer x⋆ = arg minx∈Rn F (x) lying in the interior of XR . In particular ∇F (x⋆ ) = 0. Mini-batches are i.i.d. uniform. For x ∈ XR , the stochastic gradient gI (x) is unbiased with E ∥gI (x) − ∇F (x)∥22 ≤ σg2 /b. The exact CA-SGD block with parameter s is the composition of s mini-batch SGD steps with η ≤ 1/LF .
5.2
Local error bound
Source quantities and amplification factor. First, we fix the notation used in the local bound. Let α = η/b denote the effective step size, and let d = (d0 ; . . . ; ds−1 ) ∈ Rsb be the stacked exact residual history across the s inner steps. Let GL denote the strict block-lower triangular part of G = Y Y ⊤ . GL is nilpotent within the block (its s-th power vanishes), so the matrix I − Lσ α |GL | has a finite Neumann inverse given by a terminating power series. We define the amplification factor as: Ks := (I − Lσ α |GL |)−1 2 . (3) In addition, we need a non-negative block-lower matrix that replaces the Gram entries with their magnitudes to prevent sign cancellation. Thus, we define HL as the strict block-lower matrix with Hj,i = |Yj | |Yi |⊤ for i < j. Finally, we collect the four data-dependent source quantities that each precision slot multiplies: ⊤
Tg := |Y | |d|
2
, Tr := ∥|Y | |x|∥2 , TG := ∥HL |d|∥2 , Tc :=
s−1 X
jb ∥Gj,0:j−1 ∥F ∥d0:j−1 ∥2
2 1/2 . (4)
j=1
Theorem 5.7 (Mixed-precision local error). Let each operation in the outer iteration carry its own unit roundoff u∗ ∈ {uA , uG , ur , uc , uσ , ug , uAR,r , uAR,G , ux }, with raw or effective inner-product Higham factors γr , γG , γg for the margin, Gram, and gradient ker(u ) nels (single-format kernels use γq ∗ . Low-precision-input/FP32-accumulate kernels use the 9
analogous input-cast plus FP32-accumulation factor, with the BF16 case given in Equation (8)). Assume all deterministic Higham factors appearing below are in their validity domain. Under Assumptions 5.1 to 5.5 for the local one-step comparison in which the exact and finite-precision recurrences start from the same x ∈ XR and use the same sampled (uc ) rows and labels, with ω∗ = max{uA , γr , γG , γ(s−1)b , γg , C(P )uAR,r , C(P )uAR,G , uσ , ux } and τG,∗ = Ks Lσ α(2uA + γG + C(P )uAR,G ) ∥HL ∥2 both at most 1/2, the local forward error admits the first-order additive decomposition X (2) E∗ + Rmix , (5) ≤ x b(s) − x(s) 2
∗
with the nine leading terms EuA = α ∥Y ∥2 Ks Lσ uA Tr + 2α2 ∥Y ∥2 Ks Lσ uA TG + αuA Tg , Eur = α ∥Y ∥2 Ks Lσ γr Tr , EuAR,r = α ∥Y ∥2 Ks Lσ C(P )uAR,r Tr , EuG = α2 ∥Y ∥2 Ks Lσ γG TG , EuAR,G = α2 ∥Y ∥2 Ks Lσ C(P )uAR,G TG ,
(6)
2
2
Euc = α ∥Y ∥2 Ks Lσ uc Tr + 2α ∥Y ∥2 Ks Lσ uc TG + α ∥Y ∥2 Ks Lσ uc Tc , √ Euσ = α ∥Y ∥2 Ks Cδ uσ Dδ sb, Eug = αγg Tg , Eux = ux
x(s)
2
+ α Y ⊤d
2
,
P (2) and remainder Rmix ≤ Cloc Ω∗ ∗ E∗ with Ω∗ = max{ω∗ , τG,∗ }. Per Assumption 5.4, the constant C(P ) absorbs any input-cast rounding that occurs when a recipe casts its rank-local accumulator down to a coarser collective datatype. Corollary 5.8 (Uniform-precision local error). Set all nine slots in Theorem 5.7 to the same unit (u) (u) roundoff u, and take γr = γG = γnloc and γg = γsb . Then x b(s) − x(s)
(2)
2
≤ Estep + Er + EG + Einner + Eσ + Runif ,
(7)
where the five terms are obtained by grouping the nine rows of Equation (6) by operation family, (2) and Runif ≤ Cloc Ω(u)(Estep + Er + EG + Einner + Eσ ) with Ω(u) = max{ω(u), τG (u)}, ω(u) = (u) (u) (u) (u) max{γnloc , γsb , γ(s−1)b , C(P )u, u}, and τG (u) = Ks Lσ α(γnloc + (C(P ) + 2)u) ∥HL ∥2 . In Recipe C, the Gram GEMM and the margin and gradient GEMVs use BF16 inputs with an FP32 accumulator on tensor cores. This kernel has effective Higham factor (u
)
γ en r,G,g := 2u16 + γn(u32 ) ≈ 2u16 + nu32 , (u )
(8)
where 2u16 charges the two BF16 input conversions and γn 32 charges the FP32 accumulator. The experiments use reduction lengths nloc ≤ 104 and sb ≤ 16384 (Section 7), at which nu32 ≤ 10−3 . Therefore the 2u16 ≈ 8×10−3 term dominates, and the BF16 GEMM/GEMV rows reduce to first order in the input conversion. The correction row Euc is the only row that mixes three sources. The Tc term carries the jb factor from the correction GEMV. Its length-(s−1)b deterministic Higham tail is absorbed into 10
(u )
c R(2) , since γ(s−1)b ≤ Ω∗ . The Tr and TG summands enter from the AXPY and the scalar-multiply roundings inside the correction step. The hypothesis (s−1)b uc < 1 ensures that the deterministic (uc ) factor γ(s−1)b multiplying Tc is valid. Our experiments are in the regime where (s−1)b u16 > 1, so uc is set to FP32. The same FP32 choice keeps the Tr and TG summands at first order.
5.3
Global convergence
Theorem 5.9 (SGD convergence). Under Assumptions 5.1, 5.2, 5.5 and 5.6, let Φh be the exact CA-SGD block map for outer block h and let the finite-precision block satisfy x bh+1 = Φh (b xh )+eh with P (2) ∥eh ∥2 ≤ ρ deterministically, where ρ is any uniform upper bound on the local budget ∗ E∗ + Rmix from Theorem 5.7. For a run of length H = M s, set a = (1 − µF η)s and νs = ησg2 (1 − a)/(µF b). For any χ > 0 with q = (1 + χ)a < 1, E ∥b xM − x⋆ ∥22 ≤ q M ∥b x0 − x⋆ ∥22 +
1 − qM (1 + χ)νs + (1 + χ−1 )ρ2 , 1−q
(9)
and E[F (b xM )−F (x⋆ )] ≤ (LF /2)E ∥b xM − x⋆ ∥22 . Because Assumption 5.5 fixes XR a posteriori, from the smallest ball containing the observed iterates rather than through an enforced projection, Theorem 5.9 certifies the reported runs rather than predicting boundedness. Subject to that conditioning, a precision recipe that keeps the local block budget ρ small inherits the exact mini-batch SGD rate up to a variance/rounding neighborhood while the bounded-trajectory and local-curvature hypotheses hold. Recipe C neighborhood. We specialize Theorem 5.9 to Recipe C (Table 3). The radius becomes ρC = (1 + Cloc ΩC )SC , where SC is the sum of the nine leading terms of Equation (6) under the Recipe C precision assignment. On logistic regression with the parameters fixed to (b, s, P, nloc , η) = (32, 64, 16, 256, 0.5), we obtain ΩC ≤ 1.6×10−1 , with the bound ρdet C = (1 + Cloc ΩC )SC finite when ΩC ≤ 1/2. The dominant contribution to ΩC is the BF16 margin-AllReduce term C(P ) uAR,r , which grows with P , and this validation fixes P = 16, where C(16) u16 ≈ 1.5 × 10−1 . The measured growth C(P ) ≈ 4.62P 0.56 (Figure 2) keeps C(P ) u16 < 1/2 through P = 256, so the ΩC ≤ 1/2 hypothesis of Theorem 5.7 holds at every scale we run. The worst-case bound C(P ) ≲ 2.4P reaches the threshold near P ≈ 60 and the measured fit near P ≈ 330. Beyond that scale Recipe C must promote the margin AllReduce uAR,r to FP32 (the f column of Table 3) to keep the hypothesis. A direct validation of the local budget at P = 256, where the accuracy comparison of Table 5 is not run, is left for future work.
6
Implementation on NVIDIA A100
The kernel DAG of Figure 1 and the pseudocode of Algorithm 1 fix the algorithmic structure at the kernel level. This section describes the concrete CUDA [29], cuBLAS [28], and NCCL [27] choices that realize the mixed-precision implementation on NVIDIA A100 benchmarked in Section 7. We run every benchmark of Section 7 on the GPU partition of NERSC Perlmutter, an HPE Cray EX system. Each Perlmutter GPU node pairs a single AMD EPYC 7763 (Milan, 64 cores, 256 GB DDR4) host with four NVIDIA A100 SXM4-40GB devices [26]. Each A100 exposes 108 SMs, 40 GB of HBM2e at 2039 GB/s, and the FP32, TF32, and BF16 compute roofs of Section 4.2. Every device runs at the default 400 W power cap throughout. Within a node, the four GPUs form a fully-connected NVLink 3 mesh with twelve 25 GB/s links per device, giving 600 GB/s aggregate 11
Table 3: Precision symbols for the nine slots used in Section 7. For storage, collectives, and non-GEMM/GEMV rows, the symbols denote unit roundoffs. In the uG , ur , and ug rows, b denotes BF16 input with FP32 accumulation, h denotes FP16 input with FP32 accumulation, and ha denotes FP16 input/output with FP16 accumulation. These kernel rows are evaluated (u ) through the corresponding Higham factors in Theorem 5.7. The factor ha uses γq FP16 rather than (u ) 2uFP16 + γq FP32 . H and I are low-precision only baselines. Symbols: f = uFP32 ≈ 6 × 10−8 , t = uTF32 ≈ 5×10−4 , h = uFP16 ≈ 5×10−4 , b = uBF16 ≈ 4×10−3 . u∗
A
B
C
D
E
F
G
H
I
uA uG ur uc uσ ug uAR,r uAR,G ux
f f f f f f f f f
b f f f f f f f f
b b b f f b b f f
b b b f f b b b f
h h h f f h h f f
h ha ha f f ha h f f
f t t t f t f f f
h ha ha h h ha h h h
b b b b b b b b b
peer-to-peer bidirectional bandwidth. This mesh keeps the intra-node AllReduce of Section 4.2 on NVLink and off the PCIe path. Across nodes, communication runs over HPE Slingshot 11, with four 200 Gbps Cassini NICs per GPU node on a three-hop dragonfly fabric. The implementation applies no explicit projection or norm clip on x e. The local budgets are instantiated with the radius R of Assumption 5.5, taken a posteriori as the smallest ball containing the observed iterates of the reported run. Assumption 5.5 held on every run reported here, so the convergence theorem is a conditional guarantee that certifies these runs rather than predicting boundedness. Were an iterate to leave XR , the local-curvature constants µF , LF of Assumption 5.6 and hence the radius ρ would no longer be controlled, and a projection step would be required to restore the hypothesis. We fuse the outer gradient GEMV and the weight update of Algorithm 1 (lines 16–17) into a single BF16-input/FP32-accumulate GEMM with α = η/b and β = 1. This GEMM adds the gradient directly into the FP32 master weights. The fusion leaves the precision-slot accounting of Theorem 5.7 unchanged. Every BF16-input GEMM and GEMV uses FP32 accumulation. The Gram GEMM avoids cuBLAS’s TF32 “fast-FP32” fallback so that the s2 off-diagonal entries retain true FP32 precision. We issue a single grouped AllReduce that carries both r and G at their distinct datatypes (BF16 and FP32 in Recipe C). This collective costs one α · 2 log p latency per outer iteration. Compared to an unfused pair of AllReduces, the grouped form halves the latency cost (Section 4.2). Each optimizer records two CUDA-graph executables that bracket the collective. The precollective graph runs sampling, gather, and the local cuBLAS kernels. The post-collective graph runs the residual, correction, and weight update. The grouped AllReduce is issued between the two graph launches.
12
7
Numerical Experiments
7.1
Setup
We partition the feature matrix 1D column-block with nloc = n/P . Table 4 lists the benchmarks. We ℓ2 -normalize every row to satisfy Assumption 5.1. For each dataset we lock the step size η ⋆ from a FP32 SGD sweep at (b, s) = (32, 1), and we reuse this η ⋆ across every CA-SGD timing. Unless noted otherwise, every run uses three seeds (42, 43, 44). Table 4: Benchmark datasets. synth and Poisson-synth share the same generator and feature matrix.
7.2
Dataset
m
n
synth Poisson-synth epsilon SUSY HIGGS
6
9,984 9,984 2,000 18 28
10 106 4×105 5×106 1.1×107
Loss
Source
logistic Poisson logistic logistic logistic
synthetic synthetic LIBSVM LIBSVM LIBSVM
Accuracy gap
We compare Recipe C against Recipe A under a 0.5% empirical validation threshold on the final loss. The validation cells use (m, n) = (65,536, 4096), b = 32, s ∈ {16, 64}, H = 200 outer iterations, and P = 16. Table 5 reports the results. The Lipschitz-residual logistic and linear cells covered by the local finite-precision theorem pass the threshold. The largest listed gap is 1.1×10−4 . The Poisson rows also pass, but they sit outside Assumption 5.2 and are reported only as boundary-stress evidence. Table 5: Recipe C verification on relative final-loss difference |LC − LA |/|LA | at P = 16, (m, n) = (65,536, 4096), b = 32, H = 200, three seeds. † Poisson does not satisfy Assumption 5.2. Loss logistic linear Poisson†
s = 16
s = 64
0 7×10−5 1.3×10−4
0 1.1×10−4 3.3×10−4
Figure 3 extends the parity check to a long-horizon synthetic logistic regression stress run. Recipes A–G stay at the FP32 baseline to within 1.4×10−6 , which sits two orders of magnitude below the FP16 unit roundoff and over three orders below the BF16 unit roundoff. The end-to-end Recipe H lands at 8.3×10−6 and Recipe I at 1.19×10−3 , both above the mixed-precision recipes. We repeat the accuracy check on mini-batch SGD run end to end in BF16 and FP16, measuring the worst final-loss gap to FP32 mini-batch SGD over three seeds and s ∈ {16, 64} at the operating point of Table 5. The logistic gap is 7×10−6 in both BF16 and FP16, the linear gap is 1.4×10−4 (BF16) and 2.9×10−4 (FP16), and the Poisson gap, outside Assumption 5.2, is 8.8×10−4 (BF16) and 5.4×10−4 (FP16). Every cell stays below the 0.5% validation threshold.
13
relative loss gap to A
10−2
BF16 ε FP16 ε
10−4 10−6 FP32 ε −8
10 3.2 × 103
3.2 × 104
3.2 × 105
outer iteration H B (BF16-A storage)
C (recipe)
D (aggressive BF16)
E (FP16 storage)
F (FP16 + FP16-accum HMMA)
G (FP32 + TF32 TC)
H (end-to-end FP16)
I (end-to-end BF16)
Figure 3: Relative loss gap to Recipe A on a long-horizon synthetic logistic regression stress run with P = 4, b = 32, s = 16, H = 320,000, and η = 0.10. Values below 10−8 are plotted at 10−8 . Recipes A–G follow the precision choices in Table 3. Recipes H and I intentionally perform end-to-end FP16 and BF16 training, respectively.
7.3
Speedup vs (b, s) parameter grid
We perform a sweep on (b, s) to identify both the peak speedup and the working point we use for the real-data evaluation. Table 6 reports the BF16 Recipe C speedup over FP32 SGD on synth. The fastest point in the grid is (b, s) = (1, 256) at 16.57×. We use (b, s) = (8, 64) as the real-data operating point, which keeps b = 8 for the smaller feature dimensions in Table 4. The locating sweep of Table 6 (three seeds, H = 200) records 7.15× on synth here, while the five-seed median end-to-end speedup for Recipe C is 6.19× (Table 7). Table 6: CA-SGD BF16 (Recipe C) speedup over FP32 SGD (baseline: b = 32, s = 1) on synth (m = 106 , n = 9,984, P = 64), averaged over three seeds at H = 200 to locate the operating point. The five-seed median end-to-end speedups are in Table 7. Best cell highlighted per row. b = 128 rows omitted.
7.4
b
s=1
s=4
s = 16
s = 64
s = 256
1 8 32
0.92× 0.87× 0.67×
3.13× 2.32× 1.46×
9.20× 4.95× 2.34×
14.62× 7.15× 2.29×
16.57× 6.29× 1.28×
Real-data speedup
We evaluate Recipes A, G, C, D, and F against FP32 SGD on five datasets at (b, s) = (8, 64) from Section 7.3. We report the median speedup over five random number generator seeds. Table 7 reports the results. Across these rows, the variants shown reach 5.1–6.8× speedup over FP32 SGD. The theory-covered Recipe C reaches 5.2–6.2×. Recipe D gives the largest speedup on epsilon, synth, and Poisson-synth and occupies the upper end of this range on the large-n datasets, exceeding Recipe C by at most 12%. Recipe D differs from Recipe C only in the Gram AllReduce, which D casts to BF16 (uAR,G = b in Table 3). We recommend Recipe C as the theory-covered default and Recipe D as a faster large-n variant whose Gram-collective error needs a per-problem check. 14
Table 7: Speedup at (b, s) = (8, 64) over FP32 SGD (b = 32, s = 1), median runtime is reported over five seeds. Recipe C is the theory-covered default. † Poisson-synth does not satisfy Assumption 5.2.
7.5
Dataset
P
SGD (ms)
A
G
C
D
F
epsilon SUSY HIGGS synth Poisson-synth†
16 4 4 64 64
0.055 0.027 0.028 0.085 0.085
5.40× 5.31× 5.33× 5.85× 5.97×
5.24× 5.31× 5.33× 6.01× 6.01×
5.56× 5.21× 5.23× 6.19× 6.10×
5.86× 5.31× 5.23× 6.68× 6.84×
5.56× 5.11× 5.13× 6.24× 6.19×
Strong and weak scaling
2−1
2
Weak-scaling efficiency
Runtime (ms/equiv-iter)
Strong scaling fixes the total problem at (m, n) = (65,536, 524,288) and partitions it across P ∈ {4, 8, 16, 32, 64, 128, 256} GPUs, so nloc shrinks from 131,072 to 2,048. Weak scaling fixes the perrank work at (m, nloc ) = (65,536, 65,536) and grows the total n with P , from 262,144 to 1.7×107 . Both experiments use synthetic logistic regression data with ℓ2 -normalized rows (Assumption 5.1), b = 32, and H = 100 outer iterations per run. In Figure 4, CA-SGD with BF16 storage reaches 0.062 ms per equivalent SGD iteration at P = 128. At P = 256, CA-SGD with BF16 storage reaches 0.741 weak-scaling efficiency, compared with 0.668 for FP32 SGD. The BF16-storage-only mini-batch SGD variant tracks the FP32 mini-batch SGD baseline to within 0.9%, separating the scaling effect of the CA reformulation from the precision change alone.
−4
2−7 4
8
16
32
64
128
256
1 0.8 0.6 0.4 0.2 0
4
8
16
32
64
GPUs (P )
GPUs (P )
(a) Strong scaling.
(b) Weak scaling.
Ideal A (FP32) D (aggressive BF16)
SGD (FP32) C (recipe) F (FP16 + FP16-accum HMMA)
128
256
SGD (BF16) G (FP32 + TF32 TC)
Figure 4: Scaling of CA-SGD and SGD on NERSC Perlmutter A100 under various precision settings. Each panel includes an ideal-scaling reference (linear in (a), E = 1 in (b)).
7.6
Roofline
Figure 5 shows a per-kernel roofline analysis on a single A100 GPU on a synthetic dataset. We report analytical FLOP counts divided by measured kernel time for the Gram GEMM and margin GEMV kernels, which exercise different hardware instructions across the recipes in Table 3. We vary s to measure kernel performance as the s-step length increases. For Gram GEMM, FP32 SIMT delivers 17.6 TFLOP/s (90% of the FP32 roof), TF32 tensor cores deliver 78.4 TFLOP/s (50% of the TF32 roof), and BF16 with an FP32 accumulator delivers 223.3 TFLOP/s (72% of the 15
103 Kernel performance [TFLOP/s]
s = 128: 92% of BF16 peak 102
101 GEMVs (memory-bound) 100
10−1 10−1
100
101
102
103
104
Arithmetic intensity [FLOP/Byte] FP32 scalar roof
TF32 TC roof
BF16 TC roof
Gram GEMM A (FP32)
Gram GEMM G (TF32)
Gram GEMM C (BF16)
Gram GEMM F (FP16-acc)
Margin GEMV A/G
Margin GEMV C/F
Figure 5: Per-kernel measured roofline at P = 1, m = 8192, nloc = 16,384, b = 32, H = 20. Theoretical roofs (FP32 scalar 19.5, TF32 156, BF16 312 TFLOP/s) against 2039 GB/s HBM2e (knees at 9.56, 76.5, 153 FLOP/B). BF16 roof), rising to 285.9 TFLOP/s (92% of peak) at s = 128. The 4× jump from TF32 to BF16 at modest s is the empirical reason Recipe C uses BF16-input tensor-core GEMM. Recipe F uses FP16 inputs with an FP16 accumulator. At s = 16 it sits below Recipe C, and at s = 128 the two converge. The margin and gradient GEMVs sit at arithmetic intensities of 0.5–1.0 FLOP/B and below 1.3 TFLOP/s, bandwidth-bound at every precision, which matches the analytic placement of Section 4.2. Their only mixed-precision benefit is a 2× drop in the Y reads when storage moves from FP32 to BF16.
8
Discussion
Mixed-precision CA-SGD is effective only when the s-step reformulation and precision placement are designed together. Our nine-slot error analysis justifies Recipe C (Table 3) as the theory-covered default and gives a bounded-regime convergence neighborhood for s-step SGD, extending the finiteprecision program for s-step Krylov methods [3–5, 9]. On Perlmutter A100 GPUs, Recipe C matches FP32 SGD while reaching up to ∼16× on synth and 5.2–6.2× on the main datasets. Recipe D reaches 6.8× by casting the Gram AllReduce to BF16 but lies outside the verified budget.
Acknowledgments This work was supported by the U.S. Department of Energy, Office of Science, Advanced Scientific Computing Research (ASCR) under Award No. DE-SC0025394. This research used resources of the National Energy Research Scientific Computing Center, a DOE Office of Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DEAC02-05CH11231, using NERSC award ASCR-ERCAP0030076. ISM was supported by a Cornell fellowship. 16
Generative AI Statement. We used Claude (Anthropic) to produce initial drafts of proofs and to develop software and experimental tooling. The authors directed the work, accepted, rejected, or corrected all generated material, and independently verified all software, empirical results, and proofs.
References [1] Ahmad Ajalloeian and Sebastian U. Stich. On the convergence of SGD with biased gradients. arXiv preprint arXiv:2008.00051, 2020. [2] Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223–311, 2018. doi: 10.1137/16M1080173. [3] Erin Carson. The adaptive s-step conjugate gradient method. SIAM Journal on Matrix Analysis and Applications, 39(3):1318–1338, 2018. doi: 10.1137/16M1107942. [4] Erin Carson and James Demmel. A residual replacement strategy for improving the maximum attainable accuracy of s-step Krylov subspace methods. SIAM Journal on Matrix Analysis and Applications, 35(1):22–43, 2014. doi: 10.1137/120893057. [5] Erin Carson and James Demmel. Accuracy of the s-step Lanczos method for the symmetric eigenproblem in finite precision. SIAM Journal on Matrix Analysis and Applications, 36(2): 793–819, 2015. doi: 10.1137/140990735. [6] Erin Carson and Nicholas J. Higham. A new analysis of iterative refinement and its application to accurate solution of ill-conditioned sparse linear systems. SIAM Journal on Scientific Computing, 39(6):A2834–A2856, 2017. doi: 10.1137/17M1122918. [7] Erin Carson and Nicholas J. Higham. Accelerating the solution of linear systems by iterative refinement in three precisions. SIAM Journal on Scientific Computing, 40(2):A817–A847, 2018. doi: 10.1137/17M1140819. [8] Erin Carson, Nicholas J. Higham, and Srikara Pranesh. Three-precision GMRES-based iterative refinement for least squares problems. SIAM Journal on Scientific Computing, 42(6): A4063–A4083, 2020. doi: 10.1137/20M1316822. [9] Erin Carson, Tomáš Gergelits, and Ichitaro Yamazaki. Mixed precision s-step Lanczos and conjugate gradient algorithms. Numerical Linear Algebra with Applications, 29(3):e2425, 2022. doi: 10.1002/nla.2425. [10] Erin Claire Carson. Communication-Avoiding Krylov Subspace Methods in Theory and Practice. PhD thesis, University of California, Berkeley, Berkeley, CA, 2015. [11] Ernie Chan, Marcel Heimlich, Avi Purkayastha, and Robert van de Geijn. Collective communication: Theory, practice, and experience. Concurrency and Computation: Practice and Experience, 19(13):1749–1783, 2007. doi: 10.1002/cpe.1206. [12] James Demmel, Laura Grigori, Mark Hoemmen, and Julien Langou. Communication-optimal parallel and sequential QR and LU factorizations. SIAM Journal on Scientific Computing, 34 (1):A206–A239, 2012. doi: 10.1137/080731992.
17
[13] James W. Demmel, Michael T. Heath, and Henk A. van der Vorst. Parallel numerical linear algebra. In Acta Numerica, volume 2, pages 111–197, Cambridge, UK, 1993. Cambridge University Press. doi: 10.1017/S096249290000235X. [14] Aditya Devarakonda and James Demmel. Avoiding communication in logistic regression. In 2020 IEEE 27th International Conference on High Performance Computing, Data, and Analytics (HiPC), pages 91–100. IEEE, 2020. doi: 10.1109/HiPC50609.2020.00023. [15] Aditya Devarakonda and Ramakrishnan Kannan. Communication-efficient, 2d parallel stochastic gradient descent for distributed-memory optimization. arXiv preprint arXiv:2501.07526, 2025. [16] Gene H. Golub and Charles F. Van Loan. Matrix Computations. Johns Hopkins University Press, Baltimore, MD, 4 edition, 2013. [17] Azzam Haidar, Stanimire Tomov, Jack Dongarra, and Nicholas J. Higham. Harnessing GPU tensor cores for fast FP16 arithmetic to speed up mixed-precision iterative refinement solvers. In SC18: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 603–613. IEEE Press, 2018. doi: 10.1109/SC.2018.00050. [18] Nicholas J. Higham. Accuracy and Stability of Numerical Algorithms. SIAM, Philadelphia, PA, 2nd edition, 2002. doi: 10.1137/1.9780898718027. [19] Nicholas J. Higham and Theo Mary. A new approach to probabilistic rounding error analysis. SIAM Journal on Scientific Computing, 41(5):A2815–A2835, 2019. doi: 10.1137/18M1226312. [20] Nicholas J. Higham and Theo Mary. Sharper probabilistic backward error analysis for basic linear algebra kernels with random data. SIAM Journal on Scientific Computing, 42(5):A3427– A3446, 2020. doi: 10.1137/20M1314355. [21] Nicholas J. Higham and Srikara Pranesh. Simulating low precision floating-point arithmetic. SIAM Journal on Scientific Computing, 41(5):C585–C602, 2019. doi: 10.1137/19M1251308. [22] Mark Frederick Hoemmen. Communication-Avoiding Krylov Subspace Methods. PhD thesis, University of California, Berkeley, Berkeley, CA, 2010. [23] Dhiraj Kalamkar, Dheevatsa Mudigere, Naveen Mellempudi, Dipankar Das, Kunal Banerjee, Sasikanth Avancha, Dharma Teja Vooturi, Nataraj Jammalamadaka, Jianyu Huang, Hector Yuen, Jiyan Yang, Jongsoo Park, Alexander Heinecke, Evangelos Georganas, Sudarshan Srinivasan, Abhisek Kundu, Misha Smelyanskiy, Bharat Kaul, and Pradeep Dubey. A study of BFLOAT16 for deep learning training. arXiv:1905.12322, 2019. [24] Peter McCullagh and John A. Nelder. Generalized Linear Models. Chapman & Hall, London, UK, 2nd edition, 1989. doi: 10.1007/978-1-4899-3242-6. [25] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. In International Conference on Learning Representations (ICLR), 2018. URL https://openreview.net/forum?id=r1gs9JgRZ. [26] NVIDIA Corporation. NVIDIA A100 tensor core GPU: Data sheet. https://www.nvidia. com/content/dam/en-zz/Solutions/Data-Center/a100/pdf/nvidia-a100-datasheet. pdf, 2021. 18
[27] NVIDIA Corporation. NVIDIA Collective Communications Library (NCCL) documentation. https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html, 2024. [28] NVIDIA Corporation. cuBLAS documentation. https://docs.nvidia.com/cuda/cublas/, 2026. [29] NVIDIA Corporation. CUDA Programming Guide. cuda-programming-guide/, 2026.
https://docs.nvidia.com/cuda/
[30] Rajeev Thakur, Rolf Rabenseifner, and William Gropp. Optimization of collective communication operations in MPICH. International Journal of High Performance Computing Applications, 19(1):49–66, 2005. doi: 10.1177/1094342005051521. [31] Yixiao Wang, Zishan Shao, Ting Jiang, and Aditya Devarakonda. Enhanced cyclic coordinate descent methods for elastic net penalized linear models. In Advances in Neural Information Processing Systems (NeurIPS), 2025. URL https://openreview.net/forum?id=duunKHvWKz.
19