Pollux: Co-adaptive Cluster Scheduling for Goodput-Optimized Deep Learning

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, July 14–16, 2021

Problem

DL cluster schedulers make users specify how many GPUs each job gets. Allocating too many causes long queuing times and inefficient resource use; allocating too few causes long runtimes and unused resources. The decision is especially hard in a shared cluster because the optimal choice is dynamic and depends on cluster load while the job is running. Recent elastic schedulers can pick the resource amount for the user, 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 to train the model, their optimal values vary by task and architecture, and they depend strongly on the allocation. Because resource count, batch size, and learning rate are inter-dependent they must be configured jointly, and because clusters are dynamic their optimal values change over time. Existing schedulers are agnostic to the statistical efficiency of DL training and to this inter-dependence entirely.

Core Insight

A properly-configured DL job balances system throughput (examples processed per wall-clock second) against statistical efficiency (training progress per example processed) — raising the batch size increases the first and decreases the second, and how fast the second falls depends on training progress, since a job late in training can tolerate 10× or larger batch sizes without degrading statistical efficiency. Defining goodput as the product of the two makes both axes optimizable by a single scheduler, and both factors can be learned online from a running job and then used predictively for allocations and batch sizes never actually tried.

Method

Goodput (Definition 3.1) is the product of throughput and statistical efficiency at iteration t:

GOODPUT_t(⋆) = THROUGHPUT(⋆) × EFFICIENCY_t(M(⋆)) ... (Eqn. 4)

Pollux controls ⋆ = (a, m, s): the allocation vector a (GPUs per node), the per-GPU batch size m, and the number of gradient accumulation steps s, giving total batch size M(a,m,s) = SUM(a) × m × (s+1).

Statistical efficiency is modeled from the pre-conditioned gradient noise scale (PGNS) ϕ_t = tr(PΣPᵀ)/|Pg|² (Eqn. 5), a generalization of the gradient noise scale derived from pre-conditioned SGD so that Adam/AdaGrad are supported. Since it takes 1 + ϕ_t/M iterations to make equal progress across batch sizes:

EFFICIENCY_t(M) = (ϕ_t + M_0) / (ϕ_t + M) ... (Eqn. 6)

System throughput is THROUGHPUT = M / T_iter, with T_grad linear in per-GPU batch size (Eqn. 8) and T_sync linear in GPU count using separate parameter pairs for intra-node vs. inter-node placement (Eqn. 9). Compute/communication overlap is captured by a learnable exponent γ ≥ 1, and gradient accumulation adds s extra local passes:

T_iter(a,m,s) = s × T_grad(a,m) + ( T_grad(a,m)^γ + T_sync(a)^γ )^(1/γ) ... (Eqn. 11)

γ = 1 means no overlap (T_grad + T_sync); γ → ∞ means perfect overlap (max(T_grad, T_sync)).

Two-component architecture:

                    +-------------------------------------+
                    |            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 (per job) fits the 7-tuple θ_sys = (α_grad, β_grad, α_sync^local, β_sync^local, α_sync^node, β_sync^node, γ) by minimizing RMSLE against Eqn. 11 with L-BFGS-B (α, β non-negative; γ ∈ [1,10]), estimates ϕ_t, then solves (m*, s*) = argmax GOODPUT(a, m, s) and rescales the learning rate through a plug-in SCALE_LR(M_0, M) → λ interface supporting AdaScale, square-root, linear, and LEGW rules.
• Prior-driven exploration: each job starts on one GPU with priors that zero out the sync terms until larger configurations are tried, so it is initially assumed to scale perfectly; scale-out is capped at twice the maximum GPU count the job has ever held.
• PolluxSched (cluster-wide) maximizes a power-mean fitness FITNESS_p(A) = ((1/J) Σ_j SPEEDUP_j(A_j)^p)^(1/p) (Eqn. 14), where SPEEDUP_j is goodput under A_j over goodput under a fair 1/J share (Eqn. 15). p is a fairness knob — p = 1 is the arithmetic mean of speedups; p → −∞ approaches the minimum speedup. It uses a population-based search over allocation matrices, applies a re-allocation penalty REALLOC_FACTOR_j(δ) = (T_j − R_j·δ)/(T_j + δ), and enforces interference avoidance by allowing at most one distributed job per node.
• Non-adaptive jobs (fixed M = M_0) are supported by fixing EFFICIENCY_t = 1.

Experimental Setup

Baseline job configurations were manually tuned: a GPU count was deemed valid only if the optimal batch size for that count achieves 50%–80% of ideal linear scalability. The authors state this assumption of uniformly sophisticated users is unrealistically biased in favor of the baselines.

Headline Quantitative Results

Table 2 — testbed macrobenchmark:

• vs. well-tuned baselines: 50% and 37% shorter average JCT, 27% and 27% shorter tail JCT, 20% and 33% shorter makespan versus Optimus+Oracle+TunedJobs and Tiresias+TunedJobs respectively.
• vs. realistically-configured baselines (GPU counts exactly as in the Microsoft trace, batch size M_0 × GPUs): 72% and 73% shorter average JCT, 50% and 56% shorter tail JCT, 43% and 48% shorter makespan. Optimus+Oracle only slightly outperforms Tiresias here because it raises the GPU count without also raising the batch size.
• Simulator agreement: 48% and 32% average JCT reduction over Optimus+Oracle+TunedJobs and Tiresias+TunedJobs.
• Fairness (finish-time fairness ρ): Pollux at p = −1 gives 99% of jobs ρ < 2; at p = 1 fairness is poor with a long tail of ρ > 4. Max-ρ improvement over Tiresias and Optimus is 1.5× and 5.4×.
• Model accuracy: the fitted THROUGHPUT model had at most 10% average error across all tasks on a 64-GPU cluster. Validation metrics across batch sizes matched within ±1% for all tasks except DeepSpeech2 at ±4%.
• Cloud auto-scaling: with U = 2/3 and L = 1/2, Pollux trains ImageNet at 25% cheaper cost than Or et al.'s throughput-based auto-scaler, at only 6% longer completion time.
• HPO (Table 3): Pollux 95.4 ± 0.2 accuracy, 25 min avg JCT, 10 h makespan vs. baseline 95.5 ± 0.3, 34 min, 14 h — 30% faster HPO at equivalent accuracy.
• Overheads: 1 s on 1 vCPU per 60 s scheduling interval; re-allocation once every 7 minutes per job on average; 8% average run-time overhead from checkpoint-restarts; 0.2 s per throughput-model fit; 0.4 ms to find the optimal per-GPU batch size and accumulation steps. Checkpoint-restart delay itself is 15–120 s depending on model size.
• Interference avoidance: average JCT is unaffected by even severe injected slowdowns when enabled; 1.4× longer without it at 50% interference slowdown.
• Prior-driven exploration costs only 2–5% JCT for short jobs like CIFAR10 versus offline-fitted models, and nothing measurable for long jobs.
• Scheduling interval can go up to 2 minutes before performance degrades; queuing accounts for about half of the degradation beyond that.
• Batch sizes up to 32× larger than M_0 work well in most cases.

Limitations

• No accelerator heterogeneity in the current throughput model; the authors suggest extending with Gavel's cross-accelerator metric.
• Throughput model is deliberately narrow — only GPU count/co-locality, batch size, and accumulation steps. The linear assumptions of Eqn. 11 may diverge for specialized hardware, sophisticated synchronization algorithms, different parallelization strategies, larger scales, or hidden contention unrelated to the gradient-synchronization network. GOODPUT_t is modular so THROUGHPUT can be swapped out.
• Data parallelism only; other parallelization strategies are out of scope.
• Locality is two-level. Eqn. 9 distinguishes only intra-node vs. inter-node; rack-level locality is noted as an extension, not implemented.
• Batch-size ceiling depends on the LR scaling rule, so an application-defined maximum batch size limit is required when the rule breaks down before statistical efficiency degrades.
• Re-allocation is not free — 15–120 s checkpoint-restart, 8% average run-time overhead in the testbed.
• Exploration is heuristic (prior bias, not Bayesian optimization).
• Interference is a hard constraint, not a model — at most one distributed job per node, rather than modeling the contention.
• The Optimus baseline is modified: its parameter-server throughput model was replaced with Pollux's, and it was given oracle knowledge of exact iterations-to-completion.
• Workload is synthetic — the Microsoft trace gives only submission time, GPU count, and duration, so models and datasets were substituted.
• HPO and cloud auto-scaling evaluations are preliminary (one HPO algorithm on one task; a hand-set auto-scaling policy evaluated in simulation).
• The paper notes the reasons large batch sizes may degrade validation performance "are not completely understood at the time of this paper."

Open Problems Called Out

1. Heterogeneous-accelerator goodput — extend the throughput model with a Gavel-style cross-accelerator metric.
2. Alternative THROUGHPUT models for specialized hardware, sophisticated synchronization algorithms, alternative parallelization strategies, and larger scales, exploiting GOODPUT_t's modularity.
3. A full goodput-driven cloud auto-scaling system — the paper's auto-scaler is preliminary evidence only.
4. Pollux × HPO — a full evaluation across HPO algorithm types (Bayesian-optimization vs. bandit-based trial schedulers).
5. Better LR scaling rules incorporated via the SCALE_LR plug-in interface, potentially raising the usable batch-size ceiling.
6. Richer adaptation tooling — KungFu's mechanisms are suggested for implementing PolluxAgent's adaptive policies.

Note on NCCL Tuning

Pollux 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 abstracts the whole collective into two fitted parameter pairs, (α_sync^local, β_sync^local) for intra-node and (α_sync^node, β_sync^node) for inter-node placements (Eqn. 9), plus the overlap exponent γ. The cost of that abstraction shows up in the sharp per-iteration-time jump beyond 4 GPUs when inter-node synchronization becomes required, and in the finding that every measured model except ImageNet is highly sensitive to inter-node synchronization. The authors flag the limit themselves, noting the linear assumptions "may diverge from reality for ... sophisticated synchronization algorithms" — which places NCCL-level algorithm selection squarely in the gap Pollux's fitted T_sync leaves open.