OmniLearn: A Framework for Distributed Deep Learning over Heterogeneous Clusters

TL;DR

OmniLearn tackles the challenge of training deep networks on heterogeneous distributed clusters by introducing throughput-aware dynamic batching that equalizes compute times across workers, addressing BSP stragglers and ASP staleness. The framework combines static resource-based batching, proportional-control based dynamic batching, and memory-aware safeguards, along with enhancements like weighted gradient aggregation and learning-rate scaling. Empirical results show substantial training-time reductions (14–85%) and accuracy improvements up to 6.9% in ASP, with an average speedup around 26% across models and datasets. This work provides a practical, zero-configuration approach that integrates with PyTorch and TensorFlow, enabling efficient large-scale training across edge, cloud, and HPC environments.

Abstract

Deep learning systems are optimized for clusters with homogeneous resources. However, heterogeneity is prevalent in computing infrastructure across edge, cloud and HPC. When training neural networks using stochastic gradient descent techniques on heterogeneous resources, performance degrades due to stragglers and stale updates. In this work, we develop an adaptive batch-scaling framework called OmniLearn to mitigate the effects of heterogeneity in distributed training. Our approach is inspired by proportional controllers to balance computation across heterogeneous servers, and works under varying resource availability. By dynamically adjusting worker mini-batches at runtime, OmniLearn reduces training time by 14-85%. We also investigate asynchronous training, where our techniques improve accuracy by up to 6.9%.

Figures (13)

• Figure 1: BSP worker compute-time distribution across HLs and rise in time-to-accuracy (TTA) relative to HL1 (in hours).
• Figure 2: ASP's iteration distribution of model updates from different workers across HLs. The accuracy degradation at high HLs is attributed to staleness from skewed distribution of training iterations.
• Figure 3: OmniLearn overview: Initially, fast and slow workers ($p$, $q$) train on the same mini-batch $b_{p} = b_{q}$ (breadth of 'Compute' block) that leads to stragglers (BSP) or staleness (ASP) as compute times $t_{p} < t_{q}$ (length of 'Compute'). Controller adjusts mini-batches to equalize compute times, i.e., $t^{'}_{p} \sim t^{'}_{q} \; : \; b^{'}_{p} > b^{'}_{q}$ since $p > q$ from computational standpoint.
• Figure 4: The median compute time in BSP is lower with variable-batching, but the min-max is more spread apart at high HLs as allocating batches based on core-count is not highly accurate for throughput estimation.
• Figure 5: ASP worker density of model updates across HLs. The accuracy degradation at high HLs is attributed to staleness from skewed distribution of training iterations.