Pollux: Co-adaptive Cluster Scheduling for Goodput-Optimized Deep Learning — Detailed Summary
Aurick Qiao, Sang Keun Choe, Suhas Jayaram Subramanya, Willie Neiswanger, Qirong Ho, Hao Zhang, Gregory R. Ganger, Eric P. Xing | Petuum, Inc. / Carnegie Mellon University / UC Berkeley / MBZUAI | OSDI '21 (15th USENIX Symposium on Operating Systems Design and Implementation), July 14–16, 2021
Per-section summary organized by the paper's headings, with
paragraph-level bullets and exact quantitative results where the paper
provides them. Artifact:
https://github.com/petuum/adaptdl
Abstract
- Pollux improves DL cluster scheduling by adaptively co-optimizing inter-dependent factors at both the per-job level and the cluster-wide level.
- Most existing schedulers make users specify resources per job, often leading to inefficient use. Some recent schedulers choose resources for the user, but do so without awareness of how DL training can itself be re-optimized to better use those resources.
- Pollux considers both simultaneously: by monitoring each job during training it models how the job's goodput — a metric the paper introduces combining system throughput with statistical efficiency — would change if resources were added or removed.
- Pollux dynamically (re-)assigns resources to improve cluster-wide goodput while respecting fairness, and continually optimizes each job to better use those resources.
- In real DL-job experiments and trace-driven simulation, Pollux reduces average job completion time (JCT) by 37–50% relative to state-of-the-art DL schedulers, even when those baselines are given ideal resource and training configurations for every job.
- Pollux promotes fairness based on a more meaningful measure of useful job progress, and reveals a new opportunity for reducing DL cost in cloud environments.
1. Introduction
- DL training is now a dominant workload in shared resource environments (datacenters, cloud). Jobs are resource-intensive and long-running, often needing distributed execution on expensive accelerators; dedicated DL clusters are commonly provisioned with a scheduler mediating sharing between competing jobs.
- Existing schedulers require users to manually configure jobs; done improperly this greatly degrades training performance and resource efficiency. Too many GPUs → long queuing and inefficient use; too few → long runtimes and unused resources. These decisions are especially hard in a shared cluster because optimal choices are dynamic and depend on cluster load while the job runs.
- Recent elastic schedulers can pick resource amounts but do so blindly with respect to training-related configurations that matter just as much. Batch size and learning rate influence how much computation is needed; their optimal choices vary by task and architecture and depend strongly on the allocation.
- Because resource count, batch size, and learning rate are inter-dependent they should be configured jointly; because clusters are dynamic their optimal values change over time. This creates a complex web of considerations users must navigate, requiring expert knowledge of both cluster hardware performance and DL model architecture.
The fundamental trade-off. A properly-configured DL job balances two opposing desires:
- System throughput — training examples processed per unit wall-clock time.
- Statistical efficiency — training progress made per training example processed.
- Throughput can be raised by increasing the batch size (Fig. 1a) since a larger batch enables higher utilization of more compute. But even with an optimally re-tuned learning rate, a larger batch size often decreases statistical efficiency.
- For every distinct GPU allocation there is potentially a different batch size that best balances rising throughput against falling statistical efficiency (Fig. 1b).
- How fast statistical efficiency falls with batch size depends on training progress: a job late in training can potentially tolerate 10× or larger batch sizes without degrading statistical efficiency versus early training.
Contributions as stated. Pollux jointly manages number of GPUs, co-location of workers, per-GPU batch size, gradient accumulation, and learning rate scaling. Specifically:
- A formulation of goodput — a holistic training-performance measure accounting for both throughput and statistical efficiency.
- A demonstration that a job's goodput model can be learned by observing throughput and statistical behavior during training, then used to predict performance under different allocations and batch sizes.
- A scheduling architecture using such models to configure the right combination of allocation and training parameters for every pending and running job — locally tuning per job, globally optimizing cluster-wide, with local and global components actively cooperating toward goodput maximization.
- Evaluation on a cluster testbed with a Microsoft-trace-derived workload: versus Tiresias and Optimus, Pollux cuts average JCT by up to 73%; even when all jobs are manually tuned beforehand it cuts average JCT by 37%–50% while improving finish-time fairness by 1.5×–5.4×.
- In the cloud, goodput-driven auto-scaling based on Pollux can potentially reduce the cost of training large models by 25%.
2. Background: Distributed DL Training
- DL training minimizes a loss over the dataset, where w ∈ R^d are model parameters, X the dataset, x_i a sample, ℓ the per-sample loss:
L(w) = (1/|X|) Σ_{x_i ∈ X} ℓ(w, x_i) ... (Eqn. 1)
- The loss is minimized with SGD or variants (AdaGrad, Adam). SGD repeatedly applies w^(t+1) = w^(t) − η·ĝ^(t) until the loss converges; η is the learning rate and ĝ^(t) is a stochastic gradient estimate on a random mini-batch M^(t) ⊂ X:
ĝ^(t) = (1/M) Σ_{x_i ∈ M^(t)} ∇ℓ(w^(t), x_i) ... (Eqn. 2)
- η and the batch size M = |M^(t)| are training parameters typically chosen by the user.
2.1 System Throughput
- System throughput = training samples processed per unit wall-clock time; for a distributed job it is set by (1) allocation and placement of resources, (2) the method of distributed execution and synchronization, and (3) the batch size.
- Data-parallel execution. Synchronous data parallelism replicates w^(t) across GPUs 1..K and splits each mini-batch into equal partitions M_1^(t)...M_K^(t). Each GPU k computes a local gradient estimate, where m = |M_k^(t)| is the per-GPU batch size:
ĝ_k^(t) = (1/m) Σ_{x_i ∈ M_k^(t)} ∇ℓ(w^(t), x_i) ... (Eqn. 3)
Local estimates are averaged across all GPUs to obtain ĝ^(t); each node then applies the same update.
- Iteration time decomposition. Run-time has two main components: T_grad, the time computing each ĝ_k^(t); and T_sync, the time averaging ĝ_k^(t) (e.g. via collective all-reduce) and/or synchronizing w^(t) (e.g. via parameter servers) across all GPUs. T_sync is influenced by gradient size and network performance, and is typically shorter when GPUs are co-located within the same physical node or rack.
- Limitations due to batch size. More GPUs decreases T_grad (smaller per-GPU batch), but T_sync is typically independent of batch size and unchanged. By Amdahl's Law, no matter how many GPUs are used, iteration run-time is lower-bounded by T_sync. The common remedy is to increase the batch size, so local estimates are computed over more examples, raising the T_grad : T_sync ratio and enabling higher throughput when scaling out.
2.2 Statistical Efficiency
- Statistical efficiency = training progress per unit of training data processed, influenced by batch size and learning rate; a larger batch size normally decreases it. Predicting it is key to improving it, since predictions allow adapting batch sizes and learning rates.
- Gradient noise scale (GNS). Prior work relates statistical efficiency to the GNS, the noise-to-signal ratio of the stochastic gradient. A larger GNS means batch size and learning rate can be raised further with relatively less loss of statistical efficiency. GNS varies greatly between models, is non-constant, and tends to increase during training by up to 10× or more — so significantly better statistical efficiency at large batch sizes is attainable later in training.
- Intuition. With a low-noise stochastic gradient, adding examples to the mini-batch does not significantly improve the estimate, lowering statistical efficiency. With a high-noise gradient, adding examples reduces estimate noise, maintaining efficiency. Near convergence gradients have relatively lower signal than noise, so larger batches are more useful later.
- Learning rate scaling. When total batch size M grows, η must grow too or final model quality is significantly worse. How to scale varies by model and algorithm (SGD, Adam, AdamW). Established rules: linear scaling (η ∝ M), square-root scaling (η ∝ √M, commonly used with Adam), and adaptive rules such as AdaScale.
- Large batch sizes may also degrade final validation performance for reasons "not completely understood at the time of this paper." However, for each scaling rule there is usually a problem-dependent range of batch sizes achieving similar validation performance; within it the batch size can be chosen freely without significantly degrading model quality.
2.3 Existing DL Schedulers
Non-scale-adaptive (agnostic to a job's scalability w.r.t. allocated resources):
| System | Behavior |
|---|---|
| Tiresias | User specifies GPU count at submission; fixed for the job's lifetime |
| Gandiva | User specifies GPU count; improves utilization via fine-grained time sharing and job packing. May change GPU count dynamically, but opportunistically and without knowledge of job scalability |
Scale-adaptive (automatically decide resource amounts based on how well they speed up the job):
| System | Behavior |
|---|---|
| Optimus | Learns a predictive model of each job's system throughput given various resource amounts; optimizes cluster-wide allocations to minimize average JCT |
| SLAQ | Similar technique to minimize average loss for general ML models (not evaluated on DL) |
| Gavel | Schedules using a throughput metric comparable across accelerator types |
| AntMan | Dynamic scaling + fine-grained GPU sharing for utilization, fairness, and JCT |
| Themis | Introduces finish-time fairness; two-level scheduling architecture |
- Gap: existing schedulers are agnostic to the statistical efficiency of DL training and to the inter-dependence of resource decisions and training parameters. Pollux explicitly co-adapts these values to improve goodput.
- Footnote: Pollux's current throughput model does not consider accelerator heterogeneity; the authors believe extending with Gavel's metric would let Pollux co-adapt goodput in heterogeneous clusters.
3. The Goodput of DL Training and Pollux
Definition 3.1 (Goodput). The goodput of a DL training job at iteration t is the product of its system throughput and its statistical efficiency at iteration t:
GOODPUT_t(⋆) = THROUGHPUT(⋆) × EFFICIENCY_t(M(⋆)) ... (Eqn. 4)
where ⋆ represents any configuration parameters jointly influencing throughput and batch size, and M is total batch size summed across all allocated GPUs. Footnote: the notion is analogous to goodput in computer networks — the useful portion of throughput, benchmarked by training progress per unit wall-clock time.
The three parameters Pollux controls, ⋆ = (a, m, s):
| Symbol | Meaning |
|---|---|
| a ∈ Z^N | allocation vector; a_n = number of GPUs allocated from node n |
| m ∈ Z | per-GPU batch size |
| s ∈ Z | number of gradient accumulation steps |
Total batch size: M(a, m, s) = SUM(a) × m × (s+1).
- Pollux's approach. The user supplies initial batch size M_0 and learning rate η_0 at submission. Pollux starts each job on a single GPU with m = M = M_0, s = 0, η = η_0. As the job runs, Pollux profiles execution to learn and refine predictive models for THROUGHPUT (§3.2) and EFFICIENCY (§3.1), then periodically re-tunes (a, m, s) against cluster-wide resource availability and performance.
- EFFICIENCY_t is measured relative to M_0 and η_0, and Pollux only considers batch sizes M ≥ M_0, so EFFICIENCY_t(M) is a fraction between 0 and 1 relative to EFFICIENCY_t(M_0). Goodput is therefore the portion of throughput useful for training progress — equal to throughput if and only if perfect statistical efficiency is achieved.
- Plug-in learning rate scaling. Different jobs need
different LR scaling rules, so Pollux exposes a plug-in interface
SCALE_LR(M_0, M) → λ. SCALE_LR is called before every model update step and λ is used to scale the learning rate; implementations may use metrics collected during training such as the gradient noise scale. The interface supports AdaScale, square-root scaling, linear scaling, and LEGW.
3.1 Modeling Statistical Efficiency
- EFFICIENCY_t(M) is modeled as progress made per training example using M relative to M_0. For SGD-based training this is expressible via the gradient noise scale.
- To support adaptive SGD variants (Adam, AdaGrad) the authors use the pre-conditioned gradient noise scale (PGNS), derived by closely following the original GNS ("simple" noise scale) derivation but starting from pre-conditioned SGD rather than vanilla SGD:
ϕ_t = tr(P Σ P^T) / |P g|² ... (Eqn. 5)
where g is the true gradient, P is the pre-conditioning matrix of the adaptive SGD algorithm, and Σ is the covariance matrix of per-example stochastic gradients. Footnote: pre-conditioned SGD optimizes L(Pw) instead of L(w); Adam and AdaGrad may be viewed as vanilla SGD (with momentum) plus a particular P.
- The PGNS generalizes the GNS and is mathematically equivalent to it in the vanilla-SGD special case.
- As with the GNS, it takes 1 + ϕ_t/M training iterations to make a similar amount of training progress across different batch sizes M. Therefore:
EFFICIENCY_t(M) = (ϕ_t + M_0) / (ϕ_t + M) ... (Eqn. 6)
- Interpretation: Eqn. 6 measures each example's contribution to overall progress. If EFFICIENCY_t(M) = E then (1) 0 < E ≤ 1, and (2) training at batch size M must process 1/E times as many examples to make the same progress as at M_0.
- Pollux estimates ϕ_t during training and uses Eqn. 6 to predict EFFICIENCY_t at different batch sizes. Because measured ϕ_t varies with training progress, EFFICIENCY_t(M) reflects the lifetime-dependent trends of true statistical efficiency.
Empirical validation (Fig. 2, across the six tasks of Table 1):
- TOP row: validation metric vs. training progress for three batch sizes (M_0, an intermediate size, and the per-task maximum batch size limit). Progress is measured in "statistical epochs", defined as (M/|X|) Σ_t EFFICIENCY_t(M) with |X| the dataset size — iterations normalized by EFFICIENCY_t so each statistical epoch makes, as projected by the model, the same progress across batch sizes. (Similar to "scale-invariant iterations" in AdaScale.) The similarity of validation curves across batch sizes therefore indicates how accurate EFFICIENCY_t is as a predictor of actual progress.
- Although validation curves differ for several tasks (especially in earlier epochs), they achieve similar best values across batch sizes: ±1% relative difference for all tasks except DeepSpeech2 at ±4% — margins within the plateau of high-quality models expected from large-batch training.
- MIDDLE and BOTTOM rows: measured and predicted EFFICIENCY_t during training and across batch sizes. Generally, larger batch sizes have lower EFFICIENCY_t early in training but close the gap later. The exceptions are BERT, a fine-tuning task starting from an already pre-trained model, and recommendation, which uses a much smaller and shallower architecture.
- How EFFICIENCY_t changes over training varies by task and depends on properties like the LR schedule — e.g. ImageNet, which uses step-based LR annealing, shows sharp EFFICIENCY_t increases whenever the learning rate is annealed.
- The EFFICIENCY_t function accurately models observed values across batch sizes, meaning ϕ_t measured at batch size M can predict EFFICIENCY_t at a different batch size M′ without training at M′ ahead of time.
- Upper batch size limit. In some cases the chosen LR scaling rule breaks down as batch size grows before statistical efficiency decreases, degrading final model quality; applications may therefore define a maximum batch size limit that Pollux respects. Nevertheless the authors find a batch size up to 32× larger works well in most cases, and limits for common models are well-studied for popular LR scaling rules. Better rules can be added via the plug-in interface.
- Estimating ϕ_t. Estimated like the GNS but using the pre-conditioned gradient Pg instead of g — efficient with multiple data-parallel processes because the differing ĝ_k^(t) are already available on each GPU k. That method fails with a single GPU and gradient accumulation off (s = 0); in that case Pollux switches to a differenced variance estimator using consecutive gradient estimates ĝ^(t−1) and ĝ^(t).
3.2 Modeling System Throughput
- Goal: predict per-iteration time T_iter, then compute:
THROUGHPUT(a, m, s) = M(a, m, s) / T_iter(a, m, s) ... (Eqn. 7)
- Modeling T_grad. Back-propagation run-time scales linearly with per-GPU batch size m, with α_grad and β_grad fittable:
T_grad(m) = α_grad + β_grad · m ... (Eqn. 8)
- Modeling T_sync. With a single GPU no synchronization is needed (T_sync = 0). Otherwise T_sync is linear in the number of GPUs, because in data parallelism the data sent/received per replica typically depends only on gradient/parameter size; a linear factor accounts for performance retrogressions with three or more GPUs (increasing likelihood of stragglers or network delays). Co-location on the same node reduces network communication and can improve T_sync, so different parameters are used per placement. With K = SUM(a) and N the number of physical nodes occupied by at least one replica:
T_sync(a, m) = 0 if K = 1 α_sync^local + β_sync^local · (K − 2) if N = 1, K ≥ 2 α_sync^node + β_sync^node · (K − 2) otherwise ... (Eqn. 9)
α_sync^local / β_sync^local are the constant and retrogression parameters when all processes are co-located on one node; α_sync^node / β_sync^node are the analogues when at least two processes are on different nodes. The authors note the model can be extended to rack-level locality by adding a third parameter pair.
- Combining T_grad and T_sync. Modern DL frameworks can partially overlap the two by overlapping gradient computation with network communication; the degree of overlap depends on model structure such as the ordering and sizes of its layers. With no overlap T_iter = T_grad + T_sync; with perfect overlap T_iter = max(T_grad, T_sync). A realistic value lies in between:
T_iter(a, m, 0) = ( T_grad(a,m)^γ + T_sync(a)^γ )^(1/γ) ... (Eqn. 10)
with γ ≥ 1 a learnable parameter: T_iter = T_grad + T_sync at γ = 1, smoothly transitioning toward max(T_grad, T_sync) as γ → ∞.
- Gradient accumulation. GPU memory limits per-GPU batch size, and many models hit that limit before the batch size is large enough for T_grad to overcome T_sync (or before statistical efficiency starts diminishing), causing suboptimal scalability. Of the several techniques for overcoming the memory limit the authors focus on gradient accumulation, easily implemented in popular frameworks: per-GPU gradients are aggregated locally over s forward-backward passes before being synchronized during the (s+1)th pass. One SGD iteration thus spans s accumulation steps plus one synchronization step:
T_iter(a, m, s) = s × T_grad(a,m) + ( T_grad(a,m)^γ + T_sync(a)^γ )^(1/γ) ... (Eqn. 11)
Throughput model validation (Fig. 3):
- Each task was implemented in PyTorch, which overlaps backward-pass computation with communication. Gradients are synchronized with NCCL 2.7.8, which uses either ring all-reduce or tree all-reduce depending on the detected GPUs and their placements and its own internal performance estimates.
- Measurements were taken on AWS
g4dn.12xlargeinstances (4 NVIDIA T4 GPUs each) created within the same placement group; Eqn. 11 was fitted to the data appearing in each plot. - TOP row: time per training iteration vs. number of allocated GPUs (log-scaled) with per-GPU batch size held constant. GPUs are packed into as few 4-GPU nodes as possible, producing a sharp increase beyond 4 GPUs, where inter-node network synchronization becomes required.
- BOTTOM row: system throughput (examples/sec) vs. total batch size (log-scaled) with GPU count held constant. Left of the vertical dashed line the entire mini-batch fits in GPU memory; to the right the total batch size is achieved via gradient accumulation.
- The fitted model represents observed data closely while varying both resources and batch size. All measured models except ImageNet exhibited high sensitivity to inter-node synchronization, indicating they benefit from GPU co-location. YOLOv3 and BERT benefit from gradient accumulation to increase total batch size. These characteristics are well-represented by THROUGHPUT and can be optimized for by Pollux.
- Beyond Fig. 3's configurations, THROUGHPUT was fitted on a diverse set of GPU placements and batch sizes in a 64-GPU cluster; across all DL tasks the average error of the fitted model was at most 10%.
- Limits of the throughput model. Pollux models data-parallel throughput only in the dimensions it cares about: number and co-locality of GPUs, batch size, and gradient accumulation steps. The simple linear assumptions of Eqn. 11, though sufficiently accurate for the settings tested, may diverge from reality for specialized hardware, sophisticated synchronization algorithms, different parallelization strategies, larger scales, or hidden resource contention unrelated to the gradient-synchronization network. Rather than covering all scenarios with one model, GOODPUT_t (Eqn. 4) was designed to be modular so different THROUGHPUT equations can be plugged in without interfering with Pollux's core functionality.
4. Pollux Design and Architecture
Pollux adapts DL job execution at two granularities: job-level, dynamically tuning batch size and learning rate for best utilization of allocated resources; and cluster-wide, dynamically (re-)allocating resources driven by the goodput of all jobs combined with cluster-level goals including fairness and JCT.
+-------------------------------------+
| PolluxSched |
| (Kubernetes service, cluster-wide) |
| maximize FITNESS_p(A) over J jobs |
| + re-allocation penalty |
| + interference avoidance |
+-------------------------------------+
^ (theta_sys, phi_t) | allocation
| reported every 30s | matrix A
| v (every 60s)
+-----------------+-----------+ +-------+-------------------+
| PolluxAgent (job 1) | | PolluxAgent (job J) |
| fit EFFICIENCY_t, THROUGHPUT| | ... one per job ... |
| tune (m*, s*) + LR scaling | | |
+-----------------------------+ +---------------------------+
| |
v v
PyTorch training workers (all-reduce via NCCL 2.7.8)
- PolluxAgent runs with each job: fits that job's EFFICIENCY_t and THROUGHPUT functions, tunes its batch size and learning rate for efficient utilization of current resources, and periodically reports its goodput function to PolluxSched.
- PolluxSched periodically optimizes allocations for all jobs given each job's current goodput function and cluster-wide contention, also accounting for re-allocation overhead, slowdowns from network interference between jobs, and resource fairness.
- The two co-adapt: PolluxAgent adapts each job to use its allocation efficiently while PolluxSched re-allocates resources taking into account PolluxAgent's ability to tune the job.
4.1 PolluxAgent: Job-level Optimization
- One instance starts with each job. During training it continually measures gradient noise scale and system throughput, reports them to PolluxSched at a fixed interval, uses this information to determine the most efficient batch size for the current allocation, and adapts the learning rate via the appropriate plug-in rule (e.g. AdaScale for SGD, square-root scaling for Adam).
- Online model fitting. The system throughput parameters form a 7-tuple:
θ_sys = ( α_grad, β_grad, α_sync^local, β_sync^local, α_sync^node, β_sync^node, γ ) ... (Eqn. 12)
Together with the PGNS ϕ_t and initial batch size M_0, the triple (θ_sys, ϕ_t, M_0) fully specifies the GOODPUT function. M_0 is a user-provided constant, ϕ_t is computed per §3.1, and θ_sys is estimated by fitting THROUGHPUT to observed throughput values collected during training.
- PolluxAgent measures T_iter and records the tuple (a, m, s, T_iter) for all combinations of allocation, per-GPU batch size, and accumulation steps encountered in its lifetime. Periodically it fits θ_sys to all collected data by minimizing the root mean squared logarithmic error (RMSLE) between Eqn. 11 and the collected triples, using L-BFGS-B, with every α and β constrained non-negative and γ ∈ [1, 10]. It then reports updated θ_sys and ϕ_t to PolluxSched.
- Prior-driven exploration. At job start no throughput values exist, so priors bias θ_sys toward the belief that throughput scales perfectly with more resources until such configurations are explored. Specifically: α_sync^local = 0 while the job has not used more than one GPU; α_sync^local = β_sync^local = 0 while it has not used more than one node; β_sync^local = β_sync^node = 0 while it has not used more than two GPUs.
- Resulting behavior: each job starts on a single GPU and is initially assumed to scale perfectly, so PolluxSched is encouraged to allocate more GPUs and/or nodes naturally as part of its resource optimization until PolluxAgent can estimate θ_sys more accurately. To prevent immediate scale-out to arbitrarily many GPUs, the maximum allocatable GPU count is restricted to at most twice the maximum the job has ever been allocated.
- Other principled exploration approaches (e.g. Bayesian optimization) could apply, but the authors find this simple prior-driven strategy sufficient: §5.3.2 shows it performs within 2–5% of an idealized scenario where the model is fitted offline before submission.
- Training job tuning. With θ_sys, ϕ_t, and M_0 fully specifying GOODPUT at current progress, PolluxAgent determines, with a the current allocation:
(m*, s*) = argmax_{m,s} GOODPUT(a, m, s) ... (Eqn. 13)
The job then uses that configuration for subsequent iterations, adapting its learning rate via the plug-in rule. Because EFFICIENCY_t changes over time, PolluxAgent periodically re-evaluates the most efficient configuration.
4.2 PolluxSched: Cluster-wide Optimization
- PolluxSched periodically (re-)allocates resources for every job by maximizing a fitness function defined as a generalized (power) mean across per-job speedups:
FITNESS_p(A) = ( (1/J) Σ_{j=1}^{J} SPEEDUP_j(A_j)^p )^(1/p) ... (Eqn. 14)
A is an allocation matrix whose row A_j is job j's allocation vector (A_jn = GPUs on node n allocated to job j); J is the total number of running and pending jobs sharing the cluster.
- Speedup is the factor of goodput improvement from a given allocation over a fair-resource allocation a_f, defined as an exclusive 1/J share of the cluster, with GOODPUT_j evaluated at job j's current training iteration:
SPEEDUP_j(A_j) = max_{m,s} GOODPUT_j(A_j, m, s) / max_{m,s} GOODPUT_j(a_f, m, s) ... (Eqn. 15)
Footnote: SPEEDUP has similarities with finish-time fairness, but SPEEDUP concerns training performance at a moment in time whereas finish-time fairness concerns end-to-end job completion time.
- PolluxSched leverages the fitted, predictive GOODPUT function to maximize FITNESS via a search procedure, then applies the resulting allocations to the cluster.
- Fairness and the effect of p. At p = 1, FITNESS_p is the arithmetic mean of SPEEDUP values, causing PolluxSched to allocate more GPUs to jobs achieving high SPEEDUP with many GPUs (i.e. jobs that scale well). As p → −∞, FITNESS_p smoothly approaches the minimum SPEEDUP, so maximizing it promotes equal SPEEDUP across jobs but ignores overall cluster goodput and resource efficiency. p is therefore a "fairness knob" with larger negative values more fair; the operator selects a value by organizational priorities. The authors find p = −1 achieves most goodput improvements with reasonable fairness.
- Re-allocation penalty. Each re-allocation delays the job while the training process reconfigures. Using the popular checkpoint-restart method the authors measured between 15 and 120 seconds of delay depending on model size and other initialization tasks. To prevent excessive re-allocations, fitness evaluation applies a penalty per job needing re-allocation:
SPEEDUP_j(A_j) ← SPEEDUP_j(A_j) × REALLOC_FACTOR_j(δ)
REALLOC_FACTOR_j(δ) = (T_j − R_j·δ) / (T_j + δ)
where T_j is the job's age, R_j the number of re-allocations incurred so far, and δ an estimate of the re-allocation delay. Intuitively it scales SPEEDUP under the assumption that the job's historical average re-allocation rate will continue indefinitely, so a job with a historically higher rate is penalized more for future re-allocations.
- Interference avoidance. When multiple distributed DL jobs share a node, their network usage during gradient/parameter synchronization may interfere and slow both jobs; Xiao et al. report up to 50% slowdown for DL jobs competing for network resources. PolluxSched mitigates this by disallowing different distributed jobs (each using GPUs across multiple nodes) from sharing the same node, implemented as a constraint in the search algorithm ensuring at most one distributed job per node.
- Supporting non-adaptive jobs. A user may want a fixed batch size M = M_0. PolluxSched supports these by fixing EFFICIENCY_t = 1 for that job, and can still adapt its resource allocations based solely on system throughput.
4.3 Implementation
- PolluxAgent is a Python library imported into DL training code, integrated with PyTorch, which uses all-reduce for gradient synchronization. It inserts performance-profiling code measuring per-iteration time and calculating the gradient noise scale.
- At a fixed interval PolluxAgent fits the system throughput model (Eqn. 10) to metrics collected so far and reports the fitted parameters plus latest gradient statistics to PolluxSched, then updates the job's per-GPU batch size and gradient accumulation steps by optimizing the now-up-to-date goodput function (Eqn. 4) under the current allocation.
- PolluxSched is implemented as a service in Kubernetes. At a fixed interval it runs its search algorithm and applies the resulting allocation matrix by creating and terminating Kubernetes Pods that run the job workers.
- To find a good allocation matrix, PolluxSched uses a population-based search algorithm that perturbs and combines candidate allocation matrices to produce higher-value ones, then modifies them to satisfy node resource constraints and interference avoidance; the highest-fitness matrix is applied.
- Both components need a sub-procedure optimizing GOODPUT_t(a, m, s) at fixed a (Eqn. 13). It samples a range of candidate total batch sizes M, finds the smallest s such that m = ⌈M/s⌉ fits in GPU memory per a user-defined upper bound, and takes the configuration with the highest GOODPUT.
5. Evaluation
Scope: comparison against two state-of-the-art DL schedulers on a 64-GPU testbed, where even with well-tuned baseline job configurations Pollux cuts average JCT by 37–50%; a cluster simulator for workload intensity, prior-driven exploration, scheduling interval, and interference avoidance; finish-time fairness improvements of 1.5–5.4×; and a Pollux-based auto-scaler that can potentially cut cloud cost of training large models (e.g. ImageNet) by 25%. Pollux's gains come from dynamically trading off high-throughput/low-efficiency against low-throughput/high-efficiency training modes depending on cluster state and training progress.
5.1 Experimental Setup
Testbed:
| Component | Value |
|---|---|
| Nodes | 16 |
| GPUs per node | 4 × NVIDIA T4 |
| Total GPUs | 64 |
| Instance type | AWS EC2 g4dn.12xlarge |
| vCPUs / memory per node | 48 vCPUs, 192 GB |
| Local storage | 900 GB SSD |
| Placement | all instances launched in the same placement group |
| Orchestration | Kubernetes 1.18.2 |
| Shared storage | CephFS 14.2.8 (checkpoints for checkpoint-restart elasticity) |
Synthetic workload construction:
- 160 jobs were randomly sampled from the busiest 8-hour range (hours 3–10) of the published Microsoft DL cluster traces.
- The original trace records submission time, GPU count, and duration, but nothing about model architectures or dataset characteristics — so the synthetic workload substitutes the models/datasets of Table 1.
- Jobs in both the trace and Table 1 were categorized by total GPU-time: Small (0–1 GPU-hours), Medium (1–10), Large (10–100), XLarge (100–1000). For each trace job, a Table 1 job in the same category was chosen.
Table 1 — Models and datasets used in the evaluation workload:
| Task | Dataset | Model | Optimizer | LR Scaler | M_0 | Validation target | Size | Frac. Jobs |
|---|---|---|---|---|---|---|---|---|
| Image Classification | ImageNet | ResNet-50 | SGD | AdaScale | 200 imgs | 75% top-1 acc. | XL | 2% |
| Object Detection | PASCAL-VOC | YOLOv3 | SGD | AdaScale | 8 imgs | 84% mAP | L | 6% |
| Speech Recognition | CMU-ARCTIC | DeepSpeech2 | SGD | AdaScale | 20 seqs | 25% word err. | M | 10% |
| Question Answering | SQuAD | BERT (finetune) | AdamW | Square-Root | 12 seqs | 88% F1 score | M | 10% |
| Image Classification | Cifar10 | ResNet18 | SGD | AdaScale | 128 imgs | 94% top-1 acc. | S | 36% |
| Recommendation | MovieLens | NeuMF | Adam | Square-Root | 256 pairs | 69% hit rate | S | 36% |
Each training task achieves the provided validation metrics; the fraction of jobs per category follows the public Microsoft cluster traces.
Manually-tuned jobs for baseline schedulers:
- GPU count and batch size were manually tuned per job: per-iteration time was measured across a range of allocations and batch sizes, and each model was fully trained at a range of batch sizes.
- A GPU count was deemed valid if the optimal batch size for that count achieves 50%–80% of ideal (perfectly linear) scalability versus the optimal batch size on a single GPU. Each submitted job draws its GPU count and batch size randomly from its valid set. Rationale: below 50% of ideal the job under-utilizes resources; above 80% it could still efficiently use more GPUs.
- The authors emphasize this assumption of uniformly sophisticated, highly rational users is unrealistically biased in favor of the baseline schedulers, serving only to compare Pollux against the baselines' ideal performance.
Comparison of DL schedulers:
- Baselines are Tiresias and Optimus. Pollux co-adapts GPU count and batch size; Optimus adapts only GPU count; Tiresias adapts neither. For fairness all three scale the learning rate with AdaScale for SGD and the square-root rule for Adam and AdamW.
- Pollux configuration: 60 s scheduling interval; REALLOC_FACTOR computed with δ = 30 s; PolluxAgent reports throughput parameters and gradient statistics every 30 s; default fairness knob p = −1.
- Tiresias configuration: as in Gu et al.'s testbed experiments — two priority queues, PromoteKnob disabled, queue threshold manually tuned for this workload. Jobs were placed onto as few different nodes as possible to promote worker locality.
- Optimus+Oracle: Optimus's throughput prediction model is specific to parameter-server jobs, so the authors substituted their own throughput model (§3.2). Optimus also predicts iterations-to-convergence by fitting a function to the convergence curve; since this does not work consistently for all workload models, each job was run ahead of time and Optimus was given the exact number of iterations to completion — hence "Optimus+Oracle."
- Tiresias uses the GPU count and batch size from the synthetic workload; Optimus+Oracle uses the specified batch size but determines GPU count dynamically. Each job uses gradient accumulation if allocated too few GPUs to support the specified batch size.
5.2 Testbed Macrobenchmark Experiments
Table 2 — Summary of testbed experiments:
| Policy | Avg JCT | 99%tile JCT | Makespan |
|---|---|---|---|
| Pollux (p = −1) | 0.76 h | 11 h | 16 h |
| Optimus+Oracle+TunedJobs | 1.5 h | 15 h | 20 h |
| Tiresias+TunedJobs | 1.2 h | 15 h | 24 h |
| Optimus+Oracle | 2.7 h | 22 h | 28 h |
| Tiresias | 2.8 h | 25 h | 31 h |
| Pollux (p = +1) | 0.83 h | 10 h | 16 h |
| Pollux (p = −10) | 0.84 h | 12 h | 18 h |
- Well-tuned baselines. Against Optimus+Oracle+TunedJobs and Tiresias+TunedJobs respectively, Pollux (p = −1) achieved 50% and 37% shorter average JCT, 27% and 27% shorter tail (99th percentile) JCT, and 20% and 33% shorter makespan. This setting strongly favors the baselines, essentially mimicking users with expert knowledge of system throughput, statistical efficiency, and how they change with allocations and batch sizes.
- Source of improvement (Fig. 5). During low cluster contention Pollux allocates more GPUs and uses larger batch sizes to boost throughput even at the cost of lower statistical efficiency, because overall goodput is higher. During high cluster contention it instead uses smaller batch sizes to increase statistical efficiency. Fig. 5 also notes Tiresias+TunedJobs dips between hours 16 and 20 because a 24-GPU job blocks a 48-GPU job from running.
- Realistic baselines. Without a system like Pollux, users are likely to try various GPU counts and batch sizes before finding an efficient configuration; others may not invest the time at all. A more realistic baseline ran Optimus+Oracle and Tiresias with GPU counts exactly as specified in the Microsoft cluster trace and batch size set to M_0 × (number of GPUs) — how the authors expect most users to initially configure distributed jobs. These jobs typically use fewer GPUs and smaller batch sizes than their well-configured counterparts.
- Under this workload Pollux has 72% and 73% shorter average JCT, 50% and 56% shorter tail JCT, and 43% and 48% shorter makespan versus Optimus+Oracle and Tiresias respectively. Even though Optimus+Oracle can dynamically increase a job's GPU allocation, it only slightly outperforms Tiresias because it does not also increase the batch size to better utilize the additional GPUs.
- Co-adapted configurations over time (Fig. 6). For one ImageNet job: (A) during initial low contention more GPUs are allocated, causing a larger batch size and lower statistical efficiency; (B) during subsequent high contention fewer GPUs are allocated, causing a smaller batch size and raising statistical efficiency; (C) when contention falls again more GPUs are allocated and a larger batch size used — but the per-GPU batch size is much higher than in the first low-contention period, because the job is now in its final, high-statistical-efficiency phase. Similar trade-offs over time are observed for two YOLOv3 jobs.
- Effect of the fairness knob. Pollux was run at p = 1, −1, −10. Versus no fairness (p = 1), moderate fairness (p = −1) improved average JCT but degraded tail JCT — because in this workload the tail JCT comprises long but scalable jobs (e.g. ImageNet) which take a large number of GPUs away from other jobs in the absence of fairness. Further increasing fairness (p = −10) degraded average JCT, tail JCT, and makespan. Footnote: p = −1 (harmonic mean over speedups) may be more suitable than p = 1 (arithmetic mean) when optimizing for average JCT.
System overheads:
| Overhead | Measured value |
|---|---|
| PolluxSched FITNESS_p optimization per 60 s interval | 1 second on 1 vCPU |
| Average re-allocation frequency per job | once every 7 minutes |
| Average run-time overhead due to checkpoint-restarts | 8% |
| PolluxAgent throughput-model fit (every 30 s) | 0.2 seconds |
| Finding optimal per-GPU batch size + accumulation steps | 0.4 milliseconds |
5.3 Simulator Experiments
- A discrete-time cluster simulator was built to evaluate a broader set of workloads and settings, constructed by measuring performance and gradient statistics of each Table 1 model under many resource and batch size configurations and replaying them per simulated job — thus simulating both system throughput and statistical efficiency.
- Unless stated otherwise each experiment is repeated on 8 different workload traces generated with the same duration, job count, and job size distributions as §5.2; average results across all 8 are reported.
- Simulator construction. For each Table 1 job, per-iteration time was measured for 146 different GPU allocations + placements in the 16-node / 64-GPU testbed, and for each allocation a range of batch sizes up to the GPU memory limit. Throughput is simulated by querying a multi-dimensional linear interpolation over the measured configurations. The (pre-conditioned) gradient noise scale was also measured across a range of batch sizes and across every epoch; statistical efficiency at a given batch size is simulated by linearly interpolating the PGNS between the two nearest measured batch sizes.
- Simulator fidelity. The collected data lets the simulator reproduce several system effects including the performance impact of different GPU placements. Checkpoint-restart overhead is simulated by injecting a 30-second delay for each re-allocated job. Unless stated otherwise no network interference between jobs is simulated.
- The simulator obtains similar improvement factors to the testbed: Pollux reduces average JCT by 48% and 32% over Optimus+Oracle+TunedJobs and Tiresias+TunedJobs respectively.
5.3.1 Scheduling Fairness
- Fairness is evaluated using finish-time fairness (ρ): the ratio of a job's JCT on shared resources to its JCT running in an isolated, equally-partitioned cluster. ρ < 1 means treated better-than-fair; ρ > 1 means worse-than-fair.
- Pollux with p = 1 gives poor fairness, similar to Tiresias+TunedJobs — visible as a long tail of jobs with ρ > 4. Optimus+Oracle+TunedJobs obtains better fairness thanks to an allocation algorithm attempting to equalize per-job JCT improvement.
- Pollux with p = −1 provides the best fairness, with 99% of jobs achieving ρ < 2, while still delivering the significant performance gains of Table 2. p = −10 yields slightly worse fairness overall, because PolluxSched incurs a larger number of re-allocations by ignoring their cost in favor of equalizing speedups at all times.
- For context, the Tiresias and Optimus curves are consistent with those reported (for different workloads) by Mahajan et al. Themis was not available for direct comparison, but Pollux's ρ range at p = −1 is similar to the range reported for Themis, and the max-ρ improvements (1.5× and 5.4×) over Tiresias and Optimus are also similar.
5.3.2 Other Effects on Scheduling
- Sensitivity to job load (Fig. 8a). Under increasing workload intensity (rate of job submissions), all three policies suffer longer average JCT and makespan as expected. Across all job loads, Pollux maintains similar relative improvements over the baselines.
- Impact of prior-driven exploration. Pollux explores GPU allocations from scratch during training; seeding each job's throughput model with offline historical data yielded only a 2–5% JCT reduction for short jobs like CIFAR10 and no significant change for longer-running jobs, indicating low overhead from prior-driven exploration.
- Impact of scheduling interval (Fig. 8b). Pollux performs similarly well in average JCT for intervals up to 2 minutes; longer intervals degrade performance. Since newly-submitted jobs can only start at the next interval, longer intervals should raise average queuing time — but queuing contributed only about half of the observed degradation, indicating Pollux still benefits from relatively frequent adjustment of resource allocations.
- Impact of interference avoidance (Fig. 8c). Various degrees of slowdown were artificially injected for distributed jobs sharing a node. With interference avoidance enabled, average JCT is unaffected even by severe slowdowns, because network contention is completely mitigated. Without it, average JCT is 1.4× longer when the interference slowdown is 50%. In the ideal zero-slowdown scenario PolluxSched performs similarly with or without the constraint — indicating it can still find efficient allocations while obeying it.
5.4 More Applications of Pollux
5.4.1 Cloud Auto-scaling
- In the cloud, resources can be acquired and released on demand and users pay for holding time. Goodput-driven scheduling presents a unique opportunity: when a model's statistical efficiency increases during training, it may be more cost-effective to provision more resources and use larger batch sizes during later epochs rather than earlier. The authors present preliminary simulator evidence and note a full goodput-based auto-scaling design may be future work.
- Auto-scaling policy. Scale up the number of nodes whenever goodput exceeds a fraction U of the predicted ideal goodput assuming perfect scalability:
max_{m,s} GOODPUT_t(a,m,s) / SUM(a) > U · max_{m,s} GOODPUT_t(1,m,s)
The authors set U = 2/3 and increase to a node count such that predicted goodput is approximately L = 1/2 of predicted ideal goodput.
- Baseline: the auto-scaler of Or et al., which allows the batch size to increase during training but models job performance using system throughput rather than goodput. Because system throughput does not change with training progress, throughput-based auto-scaling quickly scales out to more nodes and a larger batch size, which then remains constant (Fig. 9a). Pollux instead starts with a small number of nodes and gradually increases as the effectiveness of larger batch sizes improves over time; Fig. 9b shows Pollux maintains high statistical efficiency throughout training.
- Result: Pollux trains ImageNet at 25% cheaper cost than Or et al.'s throughput-based auto-scaling, with only a 6% longer completion time.
5.4.2 Hyper-parameter Optimization (HPO)
- In HPO a user defines a search space over model hyper-parameters and an HPO algorithm (trial scheduler) submits many training jobs (trials) to evaluate particular hyper-parameters against objectives such as model accuracy or energy efficiency.
- HPO types manage trials differently: Bayesian optimization algorithms may submit a few jobs at a time and base future trials on fully-trained results of previous ones; bandit-based algorithms may launch many trials at once and early-stop unpromising ones. A full evaluation of Pollux's effect across HPO algorithm types is future work.
- Experiment: tuning a ResNet18 on CIFAR10 with Tree-structured Parzen Estimator (TPE), a popular Bayesian-optimization HPO algorithm; search space covers learning rate and annealing, momentum, weight decay, and network width. TPE was configured so 4 trials run concurrently and 100 trials run in total. Testbed: two NVIDIA DGX A100 nodes, each with 8 A100 GPUs. The baseline scheduler assigns a static allocation of 4 GPUs (all on the same node) per trial with a fixed per-GPU batch size for every trial.
Table 3 — Summary of HPO experiments:
| Policy | Accuracy (Top 5 trials) | Avg JCT | Makespan |
|---|---|---|---|
| Pollux | 95.4 ± 0.2 | 25 min | 10 h |
| Baseline | 95.5 ± 0.3 | 34 min | 14 h |
- As expected, similar accuracy values are achieved, but Pollux completes HPO 30% faster due to adaptive (re-)allocation of resources as trials progress and adaptive batch sizes.
5.5 Artifact
- A full artifact provides the Pollux implementation, benchmark model
implementations (Table 1), testbed experiment scripts (§5.2), and
cluster simulator implementation and results (§5.3), at
https://github.com/petuum/adaptdl/tree/osdi21-artifact. Raw testbed logs and analysis scripts are athttps://github.com/petuum/pollux-results.
6. Additional Related Work
Prior DL schedulers are covered in §2.3.
Adaptive batch size training:
- AdaBatch increases the batch size at pre-determined iterations during training while linearly scaling the learning rate.
- Smith et al. suggest that instead of decaying the learning rate during training, the batch size should be increased instead.
- CABS adaptively tunes batch size and learning rate during training using gradient statistics similar to Pollux's.
- These works share a common assumption that extra computing resources are available to parallelize larger batch sizes whenever desired — rarely true in shared-resource environments. Pollux complements existing adaptive batch size strategies by adapting batch size and learning rate in conjunction with the amount of resources currently available.
- Anytime minibatch instead adapts the batch size to mitigate stragglers in distributed training.
- KungFu supports adaptive training algorithms, including adaptive batch sizes, by letting applications define custom adaptation policies and enabling efficient adaptation and monitoring during training. KungFu targets single-job training and Pollux targets cluster scheduling; the authors believe KungFu offers useful tools for implementing PolluxAgent's adaptive policies.
Hyper-parameter tuning:
- A large body of work tunes hyper-parameters for ML/DL models, typically involving many training jobs. Although batch size and learning rate lie within the space those systems optimize, Pollux's goal is fundamentally different: HPO algorithms search for the highest model quality, whereas Pollux adapts batch size and learning rate for the most efficient execution of each job, while not degrading model quality.
7. Conclusion
- Pollux is a DL cluster scheduler that co-adaptively allocates resources while simultaneously tuning each training job to best utilize those resources.
- The paper presents a formulation of goodput combining system throughput and statistical efficiency for distributed DL training.
- Based on goodput maximization, Pollux automatically and jointly tunes resource allocations, batch sizes, and learning rates — configurations particularly difficult for users to set manually.
- Pollux outperforms and is more fair than recent DL schedulers even when users can configure their jobs well, and provides even bigger benefits under more realistic user knowledge.
Limitations (as stated or acknowledged by the paper)
- No accelerator heterogeneity. The current throughput model does not consider heterogeneous accelerator types; the authors suggest extending with Gavel's metric.
- Throughput model is deliberately narrow. It covers only GPU count and co-locality, batch size, and gradient accumulation steps. The linear assumptions in Eqn. 11 may diverge from reality for specialized hardware, sophisticated synchronization algorithms, different parallelization strategies, larger scales, or hidden resource contention unrelated to the gradient-synchronization network. GOODPUT_t is modular so an alternative THROUGHPUT model can be substituted.
- Data parallelism only. T_sync is modeled for synchronous data-parallel execution; other parallelization strategies are explicitly out of scope.
- Locality granularity is two-level. Eqn. 9 distinguishes only intra-node vs. inter-node placement; rack-level locality would need a third parameter pair (noted as an extension, not implemented).
- Batch-size ceiling depends on the LR scaling rule. Some rules break down as batch size grows before statistical efficiency degrades, so an application-defined maximum batch size limit is required.
- Re-allocation is not free. Checkpoint-restart costs 15–120 s depending on model size; in the testbed this produced an average 8% run-time overhead.
- Exploration is heuristic. Prior-driven exploration is a simple bias rather than a principled method such as Bayesian optimization; it costs 2–5% JCT for short jobs versus offline-fitted models.
- Interference handling is a hard constraint, not a model. PolluxSched simply disallows two distributed jobs from sharing a node rather than modeling the interference.
- Optimus baseline is modified. Its parameter-server throughput model was replaced with Pollux's, and it was given oracle knowledge of the exact iterations to completion, because its convergence-curve fitting did not work consistently for all workload models.
- Workload is synthetic. The Microsoft trace supplies only submission time, GPU count, and duration — model architectures and datasets were substituted from Table 1.
- HPO evaluation is preliminary. Only one HPO algorithm (TPE) on one task (ResNet18/CIFAR10) was tested; a full evaluation across HPO algorithm types is future work.
- Cloud auto-scaling is preliminary. The result is simulator-based evidence with a hand-set policy (U = 2/3, L = 1/2); a full auto-scaling system design is named as future work.
- Large-batch quality effects are not fully understood. The paper states the reasons large batch sizes may degrade validation performance "are not completely understood at the time of this paper."
Open Problems and Future Work Named by the Paper
- Heterogeneous-accelerator goodput. Extending the throughput model with a Gavel-style cross-accelerator metric to co-adapt goodput in heterogeneous DL clusters.
- Alternative THROUGHPUT models. Plugging in models for specialized hardware, sophisticated synchronization algorithms, alternative parallelization strategies, and larger scales, exploiting GOODPUT_t's modularity.
- A full goodput-driven cloud auto-scaling system. The paper's auto-scaler is preliminary evidence only.
- Pollux × HPO. A full evaluation of how Pollux affects different HPO algorithm types (Bayesian-optimization vs. bandit-based trial schedulers).
- Better LR scaling rules. As new rules are developed they can be incorporated via the SCALE_LR plug-in interface, potentially raising the usable batch-size ceiling.
- Richer adaptation tooling. KungFu's mechanisms are suggested as useful for implementing PolluxAgent's adaptive policies.
Note on NCCL Tuning
Pollux is one of the few cluster-scheduling papers that names the collective library explicitly: gradients are synchronized with NCCL 2.7.8, which per the paper "uses either ring all-reduce or tree all-reduce depending on the detected GPUs and their placements and its own internal performance estimates." Pollux does not touch that decision — it instead abstracts the entire collective into two fitted parameter pairs, (α_sync^local, β_sync^local) for intra-node placements and (α_sync^node, β_sync^node) for inter-node placements (Eqn. 9), plus the overlap exponent γ of Eqn. 10/11. The cost of that abstraction is visible in Fig. 3, where per-iteration time jumps sharply beyond 4 GPUs as inter-node synchronization becomes required, and in the finding that all measured models except ImageNet are highly sensitive to inter-node synchronization. The authors themselves flag the limit, stating the linear assumptions "may diverge from reality for ... sophisticated synchronization algorithms" — placing NCCL-level algorithm selection squarely in the gap that Pollux's fitted T_sync leaves open.