Gradient Balancing (GraB) is a recently proposed technique that finds provably better data permutations when training models with multiple epochs over a finite dataset. It converges at a faster rate than the widely adopted Random Reshuffling, by minimizing the discrepancy of the gradients on adjacently selected examples. However, GraB only operates under critical assumptions such as small batch sizes and centralized data, leaving open the question of how to order examples at large scale -- i.e. distributed learning with decentralized data. To alleviate the limitation, in this paper we propose D-GraB, an algorithm that orders the examples in a parallel setting with negligible overhead, which enjoys linear speed up at rate $\tilde{O}((mnT)^{-2/3})$ on smooth non-convex objectives and $\tilde{O}((mnT)^{-2})$ under PL condition, where $n$ denotes the number of parallel workers, $m$ denotes the number of examples per worker and $T$ denotes the number of epochs. D-GraB benefits from both data ordering and parallelism. Empirically, we show on various applications including GLUE, CIFAR10 and WikiText-2 that D-GraB outperforms naive parallel GraB and Distributed Random Reshuffling in terms of both training and validation performance.

翻译：梯度平衡（GraB）是近期提出的一种技术，能在有限数据集上多轮训练模型时找到可证明更优的数据排列。通过最小化相邻选取样本梯度的差异，该方法比广泛采用的随机重排策略收敛速度更快。然而，GraB只能在批尺寸较小且数据集中存储等严格假设条件下运行，未能解决大规模场景（即去中心化数据的分布式学习）中样本排序问题。为突破这一局限，本文提出D-GraB算法，能以极低开销在并行环境中实现样本排序，在光滑非凸目标上获得$\tilde{O}((mnT)^{-2/3})$的线性加速比，在PL条件下达到$\tilde{O}((mnT)^{-2})$的加速效果（其中$n$为并行工作节点数，$m$为每个工作节点的样本数，$T$为训练轮数）。D-GraB同时受益于数据排序与并行计算。在GLUE、CIFAR10和WikiText-2等多个任务上的实验表明，无论训练性能还是验证性能，D-GraB均显著优于朴素并行GraB及分布式随机重排方法。