SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Given the statistical characteristics of the training setting, it is possible to estimate analytically the expected value of the effective speedup  $\mathbb{E}[S_{\mathrm{eff}}(\tau)]$  by using Equations 5 and 4. Moreover, when plugging in the asymptotic form of  $\mathbb{E}[T]$ , we find the expected speedup increases to infinity with  $N$
$$
\mathbb {E} [ T ] = \Theta (\sqrt {\log N}) \Rightarrow \mathbb {E} [ S _ {\mathrm {e f f}} (\tau) ] (N) \underset {N \to \infty} {\longrightarrow} \infty
$$
As shown in figure 3b, Equation 4 is less accurate when samples deviate from a normal distribution. To find the optimal compute threshold, we synchronize the empirical distribution of micro-batch
(c)
Figure 3: Statistical characteristics of the micro-batch computation latency  $t_{n}^{(m)}$  can provide a reliable estimate of the effective speedup  $S_{\mathrm{eff}}$ . The graphs depict  $S_{\mathrm{eff}}(\tau)$  based on Equation 6. Calculations rely on samples of  $t_{n}^{(m)}$  to calculate  $\tilde{M}$  and  $T$ , before plugging them into Equation 6. In the 'simulation' curves, we directly use the samples to calculate  $\tilde{M}$  and  $T$ . For the 'analytical' curves, only the mean and variance of  $t_{n}^{(m)}$  are used to approximate  $\tilde{M}$  and  $T$  using Equations 5 and 4, respectively. For 'analytical given  $\mathbb{E}[T]$ ',  $\tilde{M}$  is approximated using Equations 5, but  $\mathbb{E}[T]$  is calculated directly from the samples. More details in appendix C.2. Panels: (a) The  $t_{n}^{(m)}$  samples follow a normal distribution. (b) samples are taken from BERT1.5B pretraining with simulated delay as described in Section 5.2. (c) The optimal compute threshold can be found automatically. The effective speedup, micro-batch completion rate, and step speedup as a function of  $\tau$  in a simulated delay environment.
compute latency between all workers after a few iterations. Given this distribution, we find  $T$  and  $\tilde{M}$ , and search in a decentralized way for the threshold  $\tau^{*}$  that maximizes the effective speedup  $S_{\mathrm{eff}}(\tau)$
defined in Equation 6. Overall, the cost of synchronizing the empirical distribution and finding  $\tau^{*}$  is negligible, compared to a full training session, because it happens only once in a training session. Lowering the threshold leads to reduced compute time but higher compute drop rates. Figure 3c highlights this trade-off and the optimal  $\tau^{*}$  is marked.
4.5 Compensating for dropped samples
The effective speedup metric  $S_{\mathrm{eff}}$ , accounts for dropped samples by treating them as a source of slowdown in direct proportion to the drop rate. This consideration enables us to execute additional computations to achieve the theoretical speedup. The extent of extra time spent on redundant calculations can be as much as  $R = (M / \tilde{M} - 1)$  times the computational effort required when not applying DropCompute.
For instance, when  $10\%$  of the samples are dropped, we can expect to perform approximately  $11\%$  more calculations. Achieving this can be approached in several ways. One straightforward compensation method for LLM training involves adding an extra  $R \cdot I_{\mathrm{base}}$  steps to the training process, where  $I_{\mathrm{base}}$  represents the number of training steps conducted without using DropCompute. In practice, achieving the original accuracy often requires even fewer additional steps, as illustrated in Figure 5, resulting in an even higher effective speedup.
Another method of compensating for the dropped samples is to increase the maximal batch size. When increasing the batch by  $R$  and dropping in average  $1 - \tilde{M} / M$ , we keep the average batch size the same as without using DropCompute, hence compensating for the lost samples. A third method, orthogonal to the first two, is resampling dropped samples before starting a new epoch to diversify the overall samples seen by the model. These approaches are rigorously tested and compared in Table 1b
To be useful, DropCompute must possess two properties. First, it should not compromise the accuracy of the trained model. This property is put to test in section 5.1 where we fully train BERT-Large and ResNet-50 (Devlin et al., 2018; He et al., 2015), each on a different task, with different drop rates to compare accuracy. Second, DropCompute should maintain a high level of runtime performance, especially when compute variance or straggling workers exist and vanilla synchronous training time deteriorates. Section 5.2 tests runtime performance of DropCompute by training a 1.5 billion parameter language model, BERT1.5B (Devlin et al., 2018) with additive noise to the compute time of each worker.
Experimental setup. The analysis of all BERT models is performed on the same dataset as Devlin et al. (2018), which is a concatenation of Wikipedia and BooksCorpus with 2.5B and 800M words respectively. The finetuning of the pretrained models is performed on SQuAD-v1 (Rajpurkar et al., 2016). We verify the generality of DropCompute by additional evaluation of a ResNet-50 model for image classification on ImageNet (Deng et al., 2009). The experiments depicted in section 5.2 and section 5.1 are executed on Habana Gaudi-1 and Gaudi-2 accelerators, respectively, with high performance network (Habana, 2020).
5.1 Generalization performance
The sole difference in the optimization when DropCompute is applied is that the batch size is not deterministic, but stochastic, as explained in section 3.2. To complement theorem 4.1, we examine the generalization performance achieved with a stochastic batch size on two popular tasks.
Image classification. To evaluate the generality of stochastic batch size and DropCompute in particular, we evaluate the Top-1 accuracy of a ResNet-50 model on theImagenet dataset using our method. Since it is not common to use gradient accumulation in large scale training of this task, we simulate the drops such that each worker randomly drops its local batch, so the total batch size is stochastic. This simulated environment enables us to examine the extent of drop rate we can use without compromising accuracy. Figure 10 in appendix B.2.2 shows that up to  $10\%$  drop rate, which is more than what DropCompute operates on, there is a negligible deterioration in accuracy.
Large language model. Training LLMs is resource intensive, typically using large batch sizes, which makes DropCompute appealing. We evaluate DropCompute method on this task by fully pretraining
(a) Varying drop rate, no compensation
% Drop rate	F1 score on dev set
0%	91.32 ± 0.15
2.5-3%	91.34 ± 0.04
5.5-6%	91.44 ± 0.02
10-11%	91.19 ± 0.02
(b)  $10\%$  drop rate, with compensation
Compensation method	F1 score on dev set
None	91.19 ± 0.02
11% extra steps	91.40 ± 0.08
11% increased batch size	91.38 ± 0.08
re-computation	91.19 ± 0.11
Table 1: Maintaining the accuracy of BERT-Large pretraining. Fine-tuning results on SqUAD v1.1, where the F1 score is obtained by the pretrained model. (a) The effect of different drop rates during pretraining on the final accuracy, without compensating for the dropped samples. (b) When  $10\%$  drop rate is used during pretraining, with different methods of compensating for the dropped samples.
BERT-Large model several times, each with a different drop rate. We follow the optimization regime described in You et al. (2019) with a batch size of 64K for phase-1 and 32K for phase-2 (more details are provided in appendix B.2). Each of the pretrained models is fine-tuned on the SQuAD task 3 times with different initializations. Fine-tuning is performed without drops, as it is not a large scale resource consuming task. Table 1a shows the average accuracy ( $\pm$  standard deviation) obtained for each drop rate. As shown, DropCompute at drop rates of up to  $10\%$  have negligible accuracy difference. Higher values measured up to  $20\%$  of dropped gradients provide acceleration with a small yet discernible hit on accuracy. We note that these results are for a fixed budget of steps. In the presence of compute variance, the effective speedup indicates that additional steps can be executed while still maintaining competitive runtime performance. This notion is demonstrated in section 5.2.
5.2 Runtime performance
The main purpose of our proposed method is to maintain runtime performance when compute variance is present. We examine this by measuring the speedup of DropCompute over standard synchronous training in several settings. First, we measure the potential speedup for different drop rates and training settings by post analysis of synchronous training without drops. In addition, we introduce compute variance by training with additive noise, and measure actual speedups using DropCompute. The experiments in this section are performed on BERT1.5B. Details are provided in appendix B.1.
Training with different number of workers and micro-batches. We evaluate the potential speedup of DropCompute on several training settings with natural heterogeneity and no drops. For each setting, we post analyze what would have been the speedup for different drop rates. As can be seen in Figure 4, DropCompute exhibits increasing benefits with a growing number of workers and compute requirements. However, there are diminishing returns in terms of speedup with more accumulations. This could possibly be explained by the amortization time of a large number of micro-batches.
Figure 4: DropCompute exhibits increasing benefit on a large scale. Effective speedup versus drop rate with (left) 32 accumulations and varying workers, and (right) 112 workers and varying number of accumulations.
Simulated delay environment. Although DropCompute may have value when the workers' compute latency variance is low, its significance becomes crucial when the workers' compute latency exhibits high variability. To evaluate our method, we introduce a delay environment where random latency is added to each micro-batch computation. This additive noise follows a bounded log-normal distribution. Detailed information and motivation regarding the additive noise are in appendix B.1. The experiments are executed with 12 gradient accumulations and a local batch size of 192. In Figure 1, the negative impact of compute variance on scalability is demonstrated and mitigated
using DropCompute. The results in Figure 1 also correspond to section 4 and Equation 11, where a theoretical extrapolation follows the same trend line. When utilizing DropCompute in this setup, achieving the same training loss as the baseline might require additional training steps, however, it leads to a notable reduction in overall training time. Figure 5 demonstrates it in a training session with 64 workers, where approximately  $3\%$  more steps is needed to reach the same loss, in  $13\%$  less time.
Figure 5: DropCompute improves training time for workers with compute variance. Train loss curve of BERT1.5B pretraining, in a simulated delay environment. (left) Horizontal axis in training steps, and (right) horizontal axis in training time.
Summary. Efficient scalable systems are a key component to enable the continued development of deep learning models. To this day, state-of-the-art models rely on synchronous distributed optimization. The challenge to maintain synchronous training as an efficient solution grows larger with the quickly growing model sizes and data. Therefore, improving the robustness and scalability of distributed synchronous training is an important endeavor. This paper tackles the challenge of maintaining synchronous training scalable in the face of compute variance. We propose DropCompute to improve the robustness of synchronous training. Workers drop their remaining compute when they reach a compute threshold, determined by exchanging and analyzing the compute latency distribution. We find that for a small percentage of dropped data, a much larger percentage of time can be saved, depending on the compute latency distribution of the workers. In addition, we provide theoretical convergence guarantees and runtime predictions. We further discuss the motivation behind DropCompute and how it effectively solves the problem in appendix A.
Limitations. While DropCompute is simple and straightforward, it deals with system efficiency, and as such, the user-level implementation provided is not optimal. Mainly, the provided implementation is limited by using many gradient accumulations and integrating compute timeout in between them. However, we believe that this is not a major concern since having multiple gradient accumulations is a common practice in training LLM on a large scale and is used in state-of-the-art training configurations (Smith et al., 2022; Nvidia, 2023). In addition, DropCompute addresses variance that originates from the compute stage of the training iteration and does not solve the potential issue of network variance during the all-reduce stage.
Future directions. DropCompute is described and analyzed in this paper as a method built on top of synchronous training. However, this method can be integrated with other possibly asynchronous methods such as periodic synchronization. In appendix B.3, we implement DropCompute on top of Local-SGD (Lin et al., 2020) and show that DropCompute can also improve the robustness of Local-SGD to stragglers. A different extension for DropCompute is to apply it during the model backward calculation and save the partial gradients that were already calculated. This would generalize DropCompute for workloads that do not utilize gradient accumulations. However, it will require further study as it differs from the stochastic batch-size setting where the entire data sample is either saved or dropped.
We thank Itay Hubara for technical advising and valuable comments on the manuscript. The research of DS was Funded by the European Union (ERC, A-B-C-Deep, 101039436). Views and opinions expressed are however those of the author only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency (ERCEA). Neither the European Union nor the granting authority can be held responsible for them. DS also acknowledges the support of Schmidt Career Advancement Chair in AI.