ATA: Adaptive Task Allocation for Efficient Resource Management in Distributed Machine Learning

TL;DR

The paper tackles efficient resource management for distributed machine learning under asynchronous execution by learning how to allocate a fixed task budget $B$ across $n$ heterogeneous workers with unknown compute times. It recasts the problem as a non-linear stochastic bandit with partial feedback and introduces Adaptive Task Allocation (ATA), a lower-confidence-bound strategy that minimizes a proxy loss $\ell(\bm{a},\bm{\mu}) = \max_i a_i \mu_i$ and yields a near-optimal total computation time. The authors prove that ATA achieves a total time within a factor $(1+4\eta\ln B)$ of the optimum with full knowledge of arm distributions, plus a logarithmic additive term in the horizon, and provide an empirical variant (ATA-Empirical) with data-dependent concentration for better practical performance. Through simulations and a CIFAR-100 CNN experiment, the method demonstrates reduced resource waste and competitive runtime compared with greedy and oracle-fixed allocations, highlighting its potential for cost-effective distributed learning. The work advances resource allocation for asynchronous ML by bridging online bandit methods with non-linear, combinatorial task assignment and providing rigorous performance guarantees.

Abstract

Asynchronous methods are fundamental for parallelizing computations in distributed machine learning. They aim to accelerate training by fully utilizing all available resources. However, their greedy approach can lead to inefficiencies using more computation than required, especially when computation times vary across devices. If the computation times were known in advance, training could be fast and resource-efficient by assigning more tasks to faster workers. The challenge lies in achieving this optimal allocation without prior knowledge of the computation time distributions. In this paper, we propose ATA (Adaptive Task Allocation), a method that adapts to heterogeneous and random distributions of worker computation times. Through rigorous theoretical analysis, we show that ATA identifies the optimal task allocation and performs comparably to methods with prior knowledge of computation times. Experimental results further demonstrate that ATA is resource-efficient, significantly reducing costs compared to the greedy approach, which can be arbitrarily expensive depending on the number of workers.

Figures (7)

• Figure 1: Each row increases the number of workers by a factor of 3, starting from $17$, that is, $n = 17, 51, 153, 459$ from top to bottom. The first column shows runtime vs. suboptimality. The second column also plots suboptimality, but against total worker time, i.e., $\sum_{i=1}^n T_{i,k}$ in \ref{['alg:ata']}. The third column presents the average iteration time, given by $C_k / k$ over all iterations $k$. The last column displays the averaged cumulative regret, as defined in \ref{['eq:proxy_loss']}.
• Figure 2: We use the same setup as in \ref{['fig:sqrt']}, with each row tripling the number of workers, starting from $n=17$.
• Figure 3: Each row corresponds to an increasing number of workers, with $n = 15, 45, 150$ from top to bottom. We consider five distributions—Exponential, Uniform, Half-Normal, Lognormal, and Gamma—grouping them to have the same mean and then varying the mean across different groups. The results demonstrate that the algorithms remain robust across different distributions. The columns represent the same as in \ref{['fig:sqrt']}.
• Figure 4: Regret growth over iterations.
• Figure 5: We use the CIFAR-100 dataset krizhevsky2009learning. The model is a CNN with three convolutional layers and two fully connected layers, totaling 160k parameters. The Adam optimizer kingma2014adam is used with a constant step size of $8 \cdot 10^{-5}$. The computation time of the workers follows the same setup as in \ref{['fig:linear']}, where the mean time increases linearly. The batch size remains the same at $B=23$. Each row corresponds to a different number of workers, with $n = 17, 51, 153$ from top to bottom.

Theorems & Definitions (26)

• Remark 4.1
• Theorem 4.2: Proof in \ref{['sec:proof_2']}
• Remark 4.3
• Theorem 6.1: Proof in \ref{['proof:thm:main']}
• Lemma 6.2
• Theorem 6.3: Proof in \ref{['proof:thm:main2']}
• Remark 3.1
• Lemma 3.2
• proof
• Lemma 3.3