Shockwave: Fair and Efficient Cluster Scheduling for Dynamic Adaptation in Machine Learning

Pengfei Zheng, Rui Pan, Tarannum Khan, Shivaram Venkataraman, Aditya Akella | University of Wisconsin-Madison / UT Austin | arXiv:2210.00093v1 [cs.DC], 30 Sep 2022 | https://github.com/uw-mad-dash/shockwave

Problem

Shared GPU clusters must fairly arbitrate among competing DNN training jobs while packing them efficiently, but existing schedulers do neither well once jobs change their own configuration mid-run. Elastic ML jobs now scale batch size during training — OpenAI scaled GPT-3's batch size from 32 to 32M for a 500× speedup, and BERT-Large uses 256 to 4096 for a 2.5× speedup — which changes a job's throughput and remaining run time unpredictably. Prior market-theoretic schedulers (DRF, Themis, REF) all assume static, known resource requests, and filter-based schedulers (AlloX, Themis) need a hand-tuned fixed filter that is sub-optimal on both fairness and efficiency. Schedulers that are agnostic or reactive to adaptation overestimate remaining run time, mistakenly extend the finish-time-fairness deadline, and cannot compensate for early under-prioritization. Pollux does handle dynamism, but by auto-scaling batch size on behalf of jobs — which the authors show costs accuracy.

Core Insight

Extending the classic static Fisher market to a discrete-time Volatile Fisher Market (buyers with time-variant utilities over volatile goods), and proving that maximizing Nash Social Welfare Over Time solves its equilibrium, lets a scheduler co-optimize long-term fairness and efficiency under dynamic adaptation — provided the future adaptation trajectory is forecast rather than reacted to, which is feasible because real batch-size scaling policies (Accordion, GNS) have deterministic configuration transitions so only regime duration is random.

Method

Three components, composed into one round-based scheduler:

Component 1: Volatile Fisher Market (VFM)   [theory]
  - discrete rounds t = 1..T; all resources volatile (no carry-over)
  - job i's utility is a sequence of time-variant functions u_it
  - equilibrium = (a) Optimal Spending  (b) Work-conserving
  - maximizing NSW_OT  =>  PO_OT, Proportionality Over Time, and
    Corollary 4.0.1: minimizes PROD_i rho_i; equal budgets => rho_i <= 1
  - under uncertainty, extends to MNSW_OTE (max NSW over time in expectation)

Component 2: Bayesian adaptation predictor   [forecast]
  - regime R = (c, f): config c (batch size), duration f (fraction of epochs)
  - trajectory = sequence of regimes; only duration is random
  - Dirichlet prior Dir(N/K, ..., N/K)
  - RESTATEMENT RULE posterior:
      Dir(m_1, ..., m_k, S_k, ..., S_k),  S_k = (N - SUM m_k) / (K - k)
    (updates only completed epochs; assumes remaining regimes split evenly)
  - keeps ONE trajectory per job = posterior mean, for tractability

Component 3: Schedule solver + estimators   [systems]
  - long-term FAIRNESS estimator -> FTF rho-hat per job  (Eq. 9)
  - long-term EFFICIENCY estimator -> makespan lower bound H  (Eq. 10)
  - generalized Nash social welfare, solved by Gurobi

Solver objective (Equation 2):

Maximize over X: ( Σ_{j=1..N} ρ̂(j)^k · log Σ_t UTIL_j(X[j,t]) ) / (N·M) − (λ · H(X)) / Z₀

• Output is an N × T binary matrix X (X[j,t] = 1 schedules job J_j in round t). Planning window default is 20 two-minute rounds in §6.1 and 30 in Appendix G.
• FTF weights ρ̂(j)^k (default k = 5) act as the market budgets: a job predicted to be treated unfairly long-term gets a higher budget and is proactively prioritized.
• λ (default 1e−3) scales the makespan regularizer; Z₀ normalizes it and M is the cluster GPU count. Performance is consistent for k ∈ [1, 10] and λ ∈ [1e−4, 1e−2].
• If adaptation is predicted inside the window, a job's schedule is decomposed into regimes, each a micro-job with static throughput, and NSW is applied at regime level.
• Implemented in Python atop Gavel; round default 2 minutes; each worker binds one GPU; reactive re-solve mode by default; checkpoints on Linux NFS with < 3% checkpointing overhead.

Fairness metric used throughout — Finish Time Fairness:

ρ(G) = t_schedule / t_egalitarian, where t_egalitarian = t_exclusive · N

ρ > 1 means the job was scheduled unfairly.

Experimental Setup

Simulator fidelity vs. the physical cluster: makespan 4.97%, average JCT 4.62%, unfair fraction 3.83% — roughly 5% divergence, using the same scheduling code base and solver engine.

Headline Quantitative Results

Motivating measurements (why reactive scheduling fails):

• A job that doubles batch size three times (32 → 256) speeds up by up to 1.7×; reactive Themis misses its fairness deadline by 2.07×, and fully agnostic scheduling reaches FTF ρ = 3.07.
• For MILP makespan minimization, reactive scheduling gives 22.3% worse makespan and 28% worse cluster utilization than proactive; fully agnostic gives 30% worse makespan.
• Themis filter toy example: fixed f = 2/3 and f = 1 both give worst FTF-ρ = 1.1 (violating sharing incentive); fixed f = 1/3 gives avg. JCT 5.7 vs. 5.0 for an adaptive filter.

Predictor accuracy: regime-duration modeling error averages 6%, yielding 84% accuracy in run time prediction, with no prior training and only epoch-progress observation (200 Accordion/GNS jobs from the Gavel trace).

Physical cluster — 32 GPUs, 120 jobs (relative values, Shockwave = 1.0):

• Makespan 1.3× less than Themis, Gavel, AlloX on average; 37% better than MST; similar to OSSP. Cluster utilization 28% better on average.
• Shockwave's worst-case FTF ρ = 1.82 (2× better than Themis/Gavel/AlloX); unfair-job fraction 2.7× better. OSSP and MST reach worst-case ρ of 5.79 and 5.2, unfairly scheduling 70.8% and 25% of jobs.

Simulation at scale (64–256 GPUs, 220–900 jobs):

• Worst-case FTF ρ averages 1.32 — 2.5×, 2.4×, 3.1×, 3.9× better than Themis, Gavel, AlloX, Gandiva-Fair. Unfair-job fraction averages 4%, 6× better than the other fair schedulers.
• Gandiva-Fair is the only baseline whose average JCT is materially worse (16–22% longer).

Wins scale with dynamism. With all jobs static, Shockwave holds the unfair fraction < 5% (baselines < 18%) and improves makespan 18% with no JCT loss. As the dynamic-job fraction grows from 0.4 to 1.0, the makespan speedup over Gavel/Themis/AlloX rises to 1.3×; with all jobs dynamic, Themis schedules 28% of jobs unfairly and AlloX 22%, versus Shockwave's 9%.

Versus Pollux (same trace, same batch size schedule). Pollux wins average JCT by 3× because it scales worker counts and thereby reduces requested GPU hours per job by 2.4×; Shockwave holds worker counts fixed and so inherits 2.4× more contention. Makespan is similar, FTF is substantially better for Shockwave, and Pollux's auto-tuning costs 2% accuracy on ResNet18 and up to 4% on DeepSpeech.

Robustness and overhead. With ±100% noise injected into predicted run times, worst-case FTF is only 1.51× the oracle and average JCT 1.06×, though makespan degrades 1.36× (efficiency drops >30%, still on par with baselines). Solver bound gap at a 15 s timeout is 0.03% / 0.11% / 0.44% for 500 / 1000 / 2000 active jobs, and overall solver overhead stays < 12.5% of a two-minute round.

Limitations

• Worker count is held fixed in the evaluated prototype — only time-sharing is implemented. The market formulation is said to be re-parameterizable for worker scaling, but this is neither implemented nor measured; the cost is a 3× worse average JCT than Pollux.
• Only Accordion and GNS are modeled by the predictor; other batch-sizing policies are unsupported.
• A single trajectory per job (the posterior mean) is kept for tractability; alternative trajectories are discarded.
• Batch-size dynamism only — other adaptation forms such as model-structure changes (the lottery ticket hypothesis named in the abstract) are not modeled.
• Efficiency is sensitive to prediction error: 100% injected noise lowers scheduling efficiency by over 30%.
• The efficiency win is contention- and trace-dependent: the makespan gain falls from 35% at contention factor 3 to 8% at 1.5, and from 30–35% to 20% on the less duration-diverse Pollux trace.
• Solver quality degrades at very large job counts — a 0.44% bound gap at 2000 jobs exceeds Gurobi's recommended 0.1%.
• Makespan is 5–9% worse than OSSP, the pure-efficiency baseline, at scale.
• Internal inconsistency in reported defaults: the planning window is 20 rounds in §6.1 and 30 in Appendix G; the good λ range is [1e−4, 1e−2] in §6.1 and [1e−1, 1e1] in Appendix G.2.

Open Problems Called Out

1. Extending the predictor beyond Accordion and GNS to other gradient-state-driven batch sizing policies.
2. Re-parameterizing the Volatile Fisher Market so it allocates worker counts, not just time slices.
3. Supporting fairness metrics other than FTF — the market can encode priorities as budgets (weighted proportional fairness), but only FTF is evaluated.
4. Multi-resource allocation: VFM formally supports GPUs and CPUs, yet the evaluation covers only GPU allocation.
5. Finding a trustworthy adaptation signal — the authors report Pollux's statistical efficiency metric can be incorrect for Neural-MF models.