混合精度通信优化算法让分布式训练快6倍，精度几乎不损失。
Mixed-Precision Communication-Avoiding SGD for Generalized Linear Models on GPUs
- 用低精度存数据、高精度算矩阵，减少通信次数
- 在多个数据集上比单精度快5.1到6.8倍，误差小于0.5%
- 适配不同GPU，代码开源可直接用
分布式随机梯度下降（SGD）受限于通信而非计算，因每次迭代需跨进程执行AllReduce。通信避免型SGD（CA-SGD）通过将s次连续的AllReduce替换为一次sb×sb Gram矩阵的AllReduce，将通信分摊到s个迭代中，以增加计算和带宽开销换取更少的同步点。现代GPU具备矩阵硬件和低精度格式，可加速Gram GEMM并减小BF16通信量。本文研究了在NVIDIA GPU上针对广义线性模型的混合精度CA-SGD。有限精度分析将一次外层迭代的局部舍入误差分解为九种独立的精度选择，仅依赖硬件的低精度单位舍入误差，因此策略可在不同代GPU间迁移。具体方案为：输入矩阵与边距向量采用低精度存储，从低精度输入计算Gram矩阵时使用高精度累加，通信时以高精度传输，内层递推与权重更新均用高精度完成。在NERSC Perlmutter A100 GPU上，混合精度CA-SGD在逻辑回归、线性与泊松问题上损失值与FP32 SGD相差小于0.5%，在epsilon、SUSY、HIGGS、synth和Poisson-synth数据集上速度提升5.1至6.8倍。相关软件已开源，地址为https://doi.org/10.5281/zenodo.20448273。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。