SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Niv Giladi $^{1,2*}$  Shahar Gottlieb $^{1,2*}$  Moran Shkolnik $^{1,2}$  Asaf Karnieli $^{2}$
Ron Banner² Elad Hoffer² Kfir Yehuda Levy¹ Daniel Soudry¹
$^{1}$ Technion - Israel Institute of Technology  $^{2}$ Habana-Labs
{giladiniv, moranshkolnik, elad.hoffer, kfiryehud, daniel.soudry}@gmail.com {sgottlieb, akarnieli, rbanner}@habana.ai
Background. Distributed training is essential for large scale training of deep neural networks (DNNs). The dominant methods for large scale DNN training are synchronous (e.g. All-Reduce), but these require waiting for all workers in each step. Thus, these methods are limited by the delays caused by straggling workers.
Results. We study a typical scenario in which workers are straggling due to variability in compute time. We find an analytical relation between compute time properties and scalability limitations, caused by such straggling workers. With these findings, we propose a simple yet effective decentralized method to reduce the variation among workers and thus improve the robustness of synchronous training. This method can be integrated with the widely used All-Reduce. Our findings are validated on large-scale training tasks using 200 Gaudi Accelerators. A reference implementation<sup>2</sup> is provided.
Deep Neural Networks (DNNs) training continues to scale over size and computational footprint, as a result of a higher number of trainable parameters, wider and deeper models, and growing amounts of training data. As improvements in model quality (measured by test loss, for example) (Kaplan et al., 2020) lead over hardware capabilities (Hooker, 2020), this scale-up translates into a need for a growing number of training devices working in tandem (Chowdhery et al., 2022), turning distributed training to the standard approach for training DNNs on a large scale.
Distributed training typically refers to three parallelism paradigms — data parallel, model parallel and layer pipelining (Ben-Nun & Hoefler, 2019). Several variants and hybrid solutions exist in modern implementations such as tensor parallel (Narayanan et al., 2021) and parameter sharding (Rajbhandari et al., 2020; Rasley et al., 2020). These can be used separately or combined as they are orthogonal to each other. Mainly, data parallelism is straightforward, where the data is sharded among workers, and all workers share the same global model state. At each step, workers compute gradients locally and then aggregate them before taking an optimization step. When training synchronously, workers update their parameters in lockstep. This ensures that all workers hold a consensus on the same model and that gradients are averaged over all workers before being applied to the model. This approach is easy to implement and allows for good convergence properties, and correspondingly is the prevalent optimization method.
$^{2}$ https://github.com/paper-submissions/dropcompute
Figure 1: DropCompute improves robustness and scalability of synchronous training. A scale graph, showcasing the proposed method runtime performance of synchronous training of a 1.5 billion parameter model with additive noise, simulating compute variance. The baseline represents existing synchronous training and the dashed black line is linear scaling. (left) Real measurements of up to 200 workers. (right) An extrapolation to 2048 workers using a theoretical estimation, also proposed in the paper. More details are provided in section 5.2.
Although state-of-the-art models use synchronous optimization for training, synchronous methods scale poorly, as stragglers and communication overhead might severely deteriorate system utilization. Many resources are invested in alleviating these issues in large-scale training systems since even minor improvements can be worth hundreds of thousands of dollars. These issues are exacerbated as the required scale grows, even in homogeneous high-performance computing clusters (Petrini et al., 2003; Hoefler et al., 2010). In this paper, we are interested in cases where significant computing variance between the workers exists. This includes (but is not limited to) straggling workers.
For instance, certain learning tasks entail heterogeneity in the required computation of data, such as varying sentence lengths in language processing, or different image sizes and frame numbers in computer vision. In addition, recent state-of-the-art models use all three parallelism paradigms (data (DP), tensor (TP), and pipeline (PP) parallelism), thus each data parallel node is a set of processing units (accelerators) communicating between them (via TP and PP) to calculate the model gradients collectively. This could potentially intensify compute variance between data parallel workers. Moreover, sub-optimal hardware systems can also lead to straggling workers. Although some sources of heterogeneity can be mitigated by non-trivial engineering work, they still persist as challenges that can introduce compute variance and substantial performance drops. This is further discussed in appendix A.
In this paper, we suggest a simple, yet effective method called DropCompute to improve the robustness and scalability of synchronous optimization in the face of compute variance. We model the compute time as a random variable and show that under reasonable assumptions, a tail of straggling workers slows down the system at a rate that is not proportional to the contributed compute by these straggling workers. We harness the gradient accumulation method widely used in Large Language Models (LLMs) (Ott et al., 2018; Liu et al., 2019) to implement the method in a few lines of code on a relevant large-scale learning task.
The contributions of our work include:
- DropCompute: a novel, decentralized method to better handle heterogeneity or stragglers without additional hyper-parameters. DropCompute is hardware and framework agnostic, runs on top of existing optimizers, and can also be combined with other methods that improve other aspects of robustness such as communication overhead.
- A theoretical convergence proof of SGD with stochastic batch size, as in DropCompute.
- A theoretical runtime analysis on standard synchronous training and the proposed method. We find an approximation of the expected speedup using DropCompute, and show this speedup goes to infinity as the number of workers grows.
- An empirical evaluation of the proposed method on a relevant large scale task, using up to 200 accelerators connected with high bandwidth communication. For example, when using DropCompute before and after system optimizations (for both software and hardware) we show accelerations of  $18\%$  and  $5\%$ , respectively.
The challenge of training deep neural networks on a large scale has been extensively explored. With rapidly growing models and data sizes, numerous works tackled the weaknesses in synchronous DNN training on a large scale and suggested methods to alleviate these weaknesses.
Redundancy methods. This line of work addresses the straggling worker problem using a redundancy mechanism. Redundant workers or redundant data are used such that straggling workers will not slow down the entire system (Chen et al., 2016; Bitar et al., 2020). These methods provide better robustness to synchronous training, even in the event of a complete failure of a subset of the workers or considerable communication slowdown. However, the robustness is limited by the redundancy factor, and more generally, more compute resources are required and full utilization cannot be achieved. In addition, some coordination method is required to keep the training synchronous, i.e., keeping consensus between the model replicas of the workers. In particular, Chen et al. (2016); Bitar et al. (2020) use a centralized approach of a parameter server to determine which workers are left out at each iteration. Modern large-scale systems use decentralized variants of All-Reduce (von Luxburg et al.; Patarasuk & Yuan, 2009), so it is not trivial to determine which workers should be considered at each step, given that each worker can see a different subset of straggling workers. Moreover, combining redundancy with communication primitive collectives (e.g., All-Reduce) requires adaptation to existing underlying frameworks (Sanders et al., 2019).
Asynchronous optimization. Another approach is introducing asynchrony to the optimization. Asynchronous training is inherently more scalable than synchronous training by being robust to all kinds of workers and communication faults. This includes periodic synchronization by exchanging parameters every  $\tau$  optimization steps (Stich, 2019; Lin et al., 2020; Wang & Joshi, 2021; Zhang et al., 2015; Wang et al., 2020; Li et al., 2020b,a), approximate distributed averaging where each worker communicates with a subset of workers each step (Jiang et al., 2017; Lian et al., 2017; Assran et al., 2019; Yang et al., 2020), and many more. These works provide better scale-up properties and improve time performance. The main drawback is asynchronous optimization itself. In practice, the convergence is less stable, and more optimization steps are needed to generalize as well as synchronous training. In addition, hyperparameters should be chosen more precisely to guarantee convergence and generalization properties (Giladi et al., 2019; Mitliagkas et al., 2016). Due to these issues, asynchronous methods are less commonly used on a large scale.
Sparse and compressed communication. Alternatively, several works addressed only the communication overhead. A common technique is to reduce the amount of data exchanged by the workers at each step. This can be done by gradient pruning (Xu et al., 2021), gradient compression (Seide et al., 2014; Chen et al., 2020; Tang et al., 2021) or low-rank approximation (Vogels et al., 2019). These works reduce the communication overhead in a deterministic form while ignoring any compute variance across processing workers. This makes these works orthogonal to ours and potentially can be combined with our method.
3 Reducing Compute Variance
This paper proposes a method called DropCompute that improves the robustness of synchronous training by reducing compute variance. First, we describe the vanilla synchronous training framework. Then, we introduce the proposed approach.
We start with a model formulation for data-parallel synchronous SGD training with  $N$  workers, where each worker holds a replica of the model parameters  $\theta$ . In parallelism paradigms that combine TP or PP with DP,  $N$  represents the count of data-parallel workers, not the total number of workers involved in training. Given a dataset  $\mathcal{D}$ , we are interested in minimizing the empirical loss
$$
\mathcal {L} (\mathcal {D}, \theta) = \frac {1}{| \mathcal {D} |} \sum_ {z \in \mathcal {D}} \ell (z, \theta),
$$
where  $\ell(z, \theta)$  is the loss with respect to data-point  $z$  and the model parameters  $\theta$ . At each step, the workers calculate gradients based on a local batch and then aggregate the gradients before taking an
optimization step. An equivalent strategy would be that the workers aggregate the parameters after taking an optimization step based on the local gradients. The aggregation is done in a decentralized fashion, such as AllReduce.
We also consider the use of gradient accumulation to increase the size of the local batch. Gradient accumulation breaks down each worker's local batch into  $M$  micro-batches for computation, which enables reaching a large global batch size beyond hardware capacity. This is common practice in training of LLM on large scale where a substantial number of accumulations are being utilized (Smith et al., 2022; Nvidia, 2023). In each iteration  $i$  and accumulation  $m$ , the gradients of the loss function with respect to the model parameters are computed, denoted by
$$
g _ {n} ^ {(m)} (\theta_ {i}) = \nabla \mathcal {L} (\mathcal {D} _ {i, n} ^ {(m)}, \theta_ {i}),
$$
where  $\mathcal{D}_{i,n}^{(m)}\subseteq \mathcal{D}$  is the micro-batch  $m$  of worker  $n$ , sampled without replacement from  $\mathcal{D}$ , and we assume a constant micro-batch size  $|\mathcal{D}_{i,n}^m|$ . The gradients accumulated by worker  $n$  at step  $i$  are
$$
g _ {n} (\theta_ {i}) = \frac {1}{M} \sum_ {m = 1} ^ {M} g _ {n} ^ {(m)} (\theta_ {i}).
$$
Finally, the workers aggregate and average the computed gradients to update the model parameters
$$
\theta_ {i + 1} = \theta_ {i} - \eta g \left(\theta_ {i}\right); \quad g \left(\theta_ {i}\right) = \frac {1}{N} \sum_ {n = 1} ^ {N} g _ {n} \left(\theta_ {i}\right), \tag {1}
$$
where  $\eta$  is the learning rate. Equation 1 requires all workers to receive all gradients before updating the parameters. This communication restriction is what makes the training synchronous and ensures the model parameters are the same for all workers. However, due to this restriction, the slowest worker at each step dictates the iteration time. More generally, any variation in computation time among workers will lead to workers with idle time. Therefore, to improve the efficiency and robustness of the synchronous training process, we need to reduce the computation time variance among workers (namely, the compute variance).
3.2 Our method: DropCompute
To mitigate the effect of compute variance and to improve the robustness of synchronous training, we propose a simple yet effective method called DropCompute. DropCompute reduces the compute variance by introducing a compute threshold after which workers have to stop calculating local gradients and move directly to the communication phase, i.e., Equation 1.
In each iteration  $i$ , each worker  $n$  measures the time while calculating its local batch, and compares it against a given threshold  $\tau$ , which is set as described in section 4.4. If that time exceeds the compute threshold, the worker stops and sends the gradients it has calculated so far. The rest of the training remains unchanged, thus synchronous training is intact. The pseudo-code of DropCompute is shown in Algorithm 1. This method maintains a decentralized approach while improving its robustness to compute variance and stragglers in particular. Since our method drops samples, the batch size is no longer deterministic. Specifically, the total batch size (the total number of samples used in one iteration by all workers) can be smaller than the maximal total batch size  $b_{\mathrm{max}} = NM|\mathcal{D}_{i,n}^{(m)}|$ . Nevertheless, we show both theoretically and empirically that this makes no difference to convergence or generalization when the number of computed samples remains the same. Next, we analyze this method in section 4 and evaluate it in section 5.
4 Method Analysis
After establishing the notion of compute variance and the formulation of the proposed method, we analyze synchronous training and the potential value of the proposed method on the training time performance. To do so, we start with convergence guarantees when using DropCompute. Then, we theoretically analyze the computation time with synchronous training and when using DropCompute. Through this analysis, we estimate the potential speedup over vanilla synchronous training.
Algorithm 1 DropCompute on worker  $n$  at iteration  $i$
1: Input: model parameters  $\theta_{i}$ ; total number of micro-batches  $M$ ; compute threshold  $\tau$   
2: local batch data  $\mathcal{D}_{i,n} = \{\mathcal{D}_{i,n}^{(1)}, \mathcal{D}_{i,n}^{(2)}, \dots, \mathcal{D}_{i,n}^{(M)}\}$ ; step time  $T_{n}$  to compute local batch  $\mathcal{D}_{i,n}$   
3: Initialize step time  $T_{n} = 0$  and accumulated gradients  $g_{n}(\theta_{i}) = 0$   
4: do (1) and (2) in parallel:  
5: (1) for  $m = 1, \dots, M$  do  
6:  $g_{n}^{(m)}(\theta_{i}) = \nabla \mathcal{L}(\mathcal{D}_{i,n}^{(m)}, \theta_{i})$   
7:  $g_{n}(\theta_{i}) \gets g_{n}(\theta_{i}) + g_{n}^{(m)}(\theta_{i}) / M$   
8: (2) wait for  $T_{n} > \tau$  and break for loop (1)  
9: Output:  $g_{n}(\theta_{i})$ $\triangleright$  AllReduce
4.1 Convergence analysis of DropCompute
Using DropCompute, the batch size is no longer fixed, but a random variable, which can also potentially depends on the data samples. To the best of our knowledge, this setting is somewhat different from existing convergence guarantees. Therefore, in this section, we provide convergence guarantees for DropCompute by analyzing the convergence of SGD with stochastic batch size.
Assumption 4.1. Following the notations in section 3.1, consider a possibly non-convex smooth loss function  $\mathcal{L}(\mathcal{D},\theta)$ , with a global minimum  $\theta^{*}\in \mathbb{R}^{d}$ , and the following (commonly used) assumptions
1.  $L$ -smooth: All functions  $\mathcal{L}(\cdot, \theta)$  are with  $L$ -Lipschitzian gradients.
2. Unbiased estimation: Each  $\nabla \ell (z,\theta_i),z\in \mathcal{D}$  is an unbiased estimate of  $\nabla \mathcal{L}(\mathcal{D},\theta_i)$  which is the true (full batch) gradient at  $\theta_{i}^{2}$ . Namely,  $\forall i:\mathbb{E}[\nabla \ell (z,\theta_i)|\theta_i] = \nabla \mathcal{L}(\mathcal{D},\theta_i)$ .
3. Bounded variance: The variance of the stochastic gradient is bounded by a constant  $\sigma$ ,
$$
\forall i: \mathbb {E} [ \| \nabla \ell (z, \theta_ {i}) - \nabla \mathcal {L} (\mathcal {D}, \theta_ {i}) \| ^ {2} | \theta_ {i} ] \leq \sigma^ {2}.
$$
Theorem 4.1. Under the above assumption, applying SGD with DropCompute (Algorithm 1), ensures
$$
\mathbb {E} \| \nabla \mathcal {L} (\mathcal {D}, \bar {\theta}) \| ^ {2} \leq \frac {2 L b _ {\max } \left(\mathcal {L} \left(\mathcal {D} , \theta_ {1}\right) - \mathcal {L} \left(\mathcal {D} , \theta^ {*}\right)\right)}{K} + \frac {2 \sigma \sqrt {L \left(\mathcal {L} \left(\mathcal {D} , \theta_ {1}\right) - \mathcal {L} \left(\mathcal {D} , \theta^ {*}\right)\right)}}{\sqrt {K}}, \tag {2}
$$
where  $b_{\mathrm{max}}$  is the maximal total batch size,  $K$  is the total number of samples that are used throughout the training,  $\theta_1$  is the initial model, and  $\bar{\theta}$  is a random sample of  $\theta_i$  from the trajectory obtained by Algorithm 1, where  $i$  is selected with probability proportional to the total batch size at iteration  $i$ . The expectation is with respect to the randomization introduced due to sampling from  $\mathcal{D}$  throughout the optimization process and with respect to choosing  $\bar{\theta}$ .
The bound in Theorem 4.1 is similar to existing fixed batch-size guarantees (Dekel et al., 2012). The second term in the bound does not degrade with the batch sizes, and behaves like  $O(1 / \sqrt{K})$ , while the first term behaves like  $O(b_{\max} / K)$ . This implies that as long as  $b_{\max} \leq O(\sqrt{K})$ , the second term is dominant and we are in the regime of linear speedup. This shows we can attain a linear speedup despite using changing batch sizes, as long as the maximal batch size is bounded.
Similarly, in Theorem D.1 in appendix D.1 we show that the loss itself converges, in the convex case along with a proof. The proof of Theorem 4.1 is provided in appendix D.2. Lastly, we discuss the impact of the stochastic batch size on generalization in appendix D.4.
4.2 Iteration time in standard synchronous distributed training
We start with finding a closed-form expression for the cumulative distribution function (CDF) of the iteration time, denoted as  $T$ , defined as
$$
T = \max (T _ {1}, T _ {2},..., T _ {N}),
$$
2This is easily satisfied when all workers can access all data.
where  $T_{n}$  represents the time taken by worker  $n$  to compute its local batch, which follows some cumulative probability function  $F_{T_n}(x) = \mathbb{P}(T_n < x)$ , which we assume is independent for each worker. Let  $F_{T}(x)$  represent the cumulative distribution function and  $f_{T}(x)$  represent the probability density function of the maximum iteration time  $T$ . The relation between  $F_{T}(x)$  and  $F_{T_{n}}(x)$  is
$$
F _ {T} (x) = \mathbb {P} \left(\max  \left(T _ {1}, \dots , T _ {N}\right) \leq x\right) = \prod_ {n = 1} ^ {N} F _ {T _ {n}} (x).
$$
Differentiating with respect to  $x$  and applying the chain rule gives:
$$
f _ {T} (x) = \frac {d F _ {T}}{d x} (x) = \sum_ {n = 1} ^ {N} f _ {T _ {n}} (x) \prod_ {n ^ {\prime} \neq n} ^ {N} F _ {T _ {n ^ {\prime}}} (x).
$$
In the special case where all of the workers' iteration time distributions are identically and independently distributed (i.i.d.), this reduces to the well-known formula:
$$
f _ {T} (x) = N \cdot f _ {T _ {n}} (x) \cdot F _ {T _ {n}} (x) ^ {N - 1} \tag {3}
$$
If the iteration time of each worker is distributed normally  $(\sim \mathcal{N}(\mu, \sigma^2))$ , the expected value of  $T$  can be approximated as shown by Bailey et al. (2014):
$$
\mathbb {E} (T) \approx \sigma \cdot \left((1 - \gamma) \cdot \Phi^ {- 1} \left(1 - \frac {1}{N}\right) + \gamma \cdot \Phi^ {- 1} \left(1 - \frac {1}{e \cdot N}\right)\right) + \mu \tag {4}
$$
where  $\Phi$  is the CDF of the standard normal distribution, and  $\gamma$  is the euler-mascheroni constant. Asymptotically the total iteration time is  $\mathbb{E}[T] = \Theta (\sqrt{\log N})$ . When the number of micro-batches  $M\gg 1$ , we can make a similar approximation under Central Limit Theorem (CLT) conditions. More details are in appendix C.2.
4.3 Iteration time and number of micro-batches with DropCompute
When using DropCompute with a constant threshold  $\tau$ , each worker is preempted at  $\tilde{T}_n = \min \{\tau, T_n\}$  and joins the AllReduce. Therefore, the total iteration time with DropCompute is
$$
\tilde {T} + T ^ {c} = \min  \left\{\tau , T \right\} + T ^ {c},
$$
where  $T^c$  is a serial latency present in each iteration, which includes the AllReduce step. This upper limit serves to clip extreme values of  $T_n$ , effectively constraining the range of potential outcomes for  $\tilde{T}$ . As a result, the compute time variability decreases, leading to a narrower distribution and enhanced compute efficiency. These effects are illustrated in Figure 2.
As a consequence of preempting each worker at  $\tilde{T}_n$ , the number of micro-batches computed in each step varies. Denote as  $t_n^{(m)}$  the compute latency of a single micro-batch  $m$  for worker  $n$ , and  $T_{n}^{(m)} = \sum_{j = 1}^{m}t_{n}^{(j)}$ . We can define the average number of micro-batches computed by each worker before reaching threshold  $\tau$  as
$$
\tilde {M} = \frac {1}{N} \sum_ {n = 1} ^ {N} \sum_ {m = 1} ^ {M} \left\{ \begin{array}{c c} 1, & \text {i f} T _ {n} ^ {(m)} <   \tau \\ 0, & \text {o t h e r w i s e} \end{array} \right\}.
$$
Under CLT conditions, the expected value for  $\tilde{M}$  can be approximated in a closed form:
$$
\mathbb {E} [ \tilde {M} ] \approx \sum_ {m = 1} ^ {M} \Phi \left(\frac {\tau - m \cdot \mu}{\sqrt {m \cdot \sigma^ {2}}}\right) \tag {5}
$$
where  $\mu, \sigma^2$  are the mean and variance for a single micro-batch  $t_n^{(m)}$  compute latency, and  $\Phi$  is the CDF of the standard normal distribution. This approximation closely fits the real value of  $\tilde{M}$  and can be used to analyze the expected gain from DropCompute. More details in appendix C.2.
Figure 2: Reduction in variance and mean iteration time using DropCompute. Iteration time distribution of 200 workers using DropCompute on BERT 1.5B. (left) Step time  $T_{n}$  distribution of all workers, without DropCompute. (right) Maximum iteration time  $T$ , across all workers, using DropCompute, with different drop rates. The dashed 'simulation' distribution is generated by drawing  $T_{n}$  randomly from an independent normal distribution separately for each worker, using the empiric mean and variance of that worker.
4.4 Choosing the threshold
The throughput of the system can be seen as the number of micro-batches computed per second. For  $N$  workers, this can be written as  $NM / (T + T^c)$ . To evaluate the effectiveness of DropCompute, we consider the difference in throughput between the baseline and when using DropCompute. Doing so, we can define the effective speedup for  $\tau$  as:
$$
S _ {\text {e f f}} (\tau) = \frac {\text {D r o p C o m p u t T h r o u g h p u t}}{\text {B a s e l i n e T h r o u g h p u t}} = \frac {N \tilde {M} / (\min  \left\{\tau , T \right\} + T ^ {c})}{N M / (T + T ^ {c})} = \frac {\tilde {M} (T + T ^ {c})}{M (\min  \left\{\tau , T \right\} + T ^ {c})} \tag {6}
$$