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


1. Introduction

The fundamental trade-off. A properly-configured DL job balances two opposing desires:

  1. System throughput — training examples processed per unit wall-clock time.
  2. Statistical efficiency — training progress made per training example processed.

Contributions as stated. Pollux jointly manages number of GPUs, co-location of workers, per-GPU batch size, gradient accumulation, and learning rate scaling. Specifically:


2. Background: Distributed DL Training

L(w) = (1/|X|) Σ_{x_i ∈ X} ℓ(w, x_i) ... (Eqn. 1)

ĝ^(t) = (1/M) Σ_{x_i ∈ M^(t)} ∇ℓ(w^(t), x_i) ... (Eqn. 2)

2.1 System Throughput

ĝ_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.

2.2 Statistical Efficiency

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

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).

3.1 Modeling Statistical Efficiency

ϕ_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.

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

Empirical validation (Fig. 2, across the six tasks of Table 1):

3.2 Modeling System Throughput

THROUGHPUT(a, m, s) = M(a, m, s) / T_iter(a, m, s) ... (Eqn. 7)

T_grad(m) = α_grad + β_grad · m ... (Eqn. 8)

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.

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 γ → ∞.

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):


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)

4.1 PolluxAgent: Job-level Optimization

θ_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.

(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

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_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.

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.

4.3 Implementation


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:

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:

Comparison of DL schedulers:

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

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

5.3.1 Scheduling Fairness

5.3.2 Other Effects on Scheduling

5.4 More Applications of Pollux

5.4.1 Cloud Auto-scaling

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.

5.4.2 Hyper-parameter Optimization (HPO)

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

5.5 Artifact


Prior DL schedulers are covered in §2.3.

Adaptive batch size training:

Hyper-parameter tuning:


7. Conclusion


Limitations (as stated or acknowledged by the paper)


Open Problems and Future Work Named by the Paper

  1. Heterogeneous-accelerator goodput. Extending the throughput model with a Gavel-style cross-accelerator metric to co-adapt goodput in heterogeneous DL clusters.
  2. Alternative THROUGHPUT models. Plugging in 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 of how Pollux affects different HPO algorithm types (Bayesian-optimization vs. bandit-based trial schedulers).
  5. 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.
  6. 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.