Elastic Training

"I went to sleep with eight peers and woke up with four. Nobody asked me. We counted heads, agreed on a number, and got back to work before the gradient cooled."

A Worker That Survived a Preemption

The previous sections of this chapter built the parts an elastic job stands on. Section 18.2 made checkpoints sharded and cheap to write, and Section 18.3 made a restart replay deterministically from one. Those give you a job that can recover from a failure by stopping and starting again at fixed size. Elastic training asks for something stronger: do not stop. When a node disappears, the remaining workers should notice, regroup, and continue within seconds, and when a node returns, they should absorb it without a human relaunching the job. The job's identity persists across membership changes; only its size fluctuates. That shift, from a fixed-size run that survives failure by restarting to a variable-size run that survives failure by reconfiguring, is what this section is about.

Two questions organize the material. First, how do the survivors agree on who is still present and resume coherently? That is a control-plane consensus problem, and it reuses machinery you met in Section 2.6. Second, what does changing the number of workers do to the math of the optimizer, given that the global batch is the sum of the per-worker batches? That ties directly to the large-batch scaling rules of Section 10.8. We take them in turn, then run a job through a real membership storm.

1. From Fixed Membership to a Variable World Size Beginner

In every data-parallel job so far, the number of workers, called the world size, was a constant fixed at launch. Each rank held an identical copy of the model and processed its own shard of every batch; the all-reduce after the backward pass averaged their gradients into one. That design has a brittle assumption baked in: the all-reduce expects exactly $K$ participants. If one of them dies mid-step, the collective hangs waiting for a tensor that will never arrive, and the standard remedy is a timeout that kills the entire job. At a thousand GPUs, where Section 18.1 showed that something is almost always broken, this turns routine hardware faults into full-run failures.

Elastic training removes the assumption. It treats the world size as a quantity that can change between steps, written $K_t$ for the number of live workers at step $t$. The model replica on each surviving worker is unchanged; what changes is how many replicas participate in the next all-reduce and how the data is divided among them. The job keeps a single logical state, the model weights and optimizer state, and that state is portable across any $K_t$ because it lives in the checkpoint, not in the membership. When the set of workers changes, the job does not lose its place; it reloads from the most recent checkpoint (the cheap sharded one from Section 18.2) and resumes with whatever workers are present.

2. Rendezvous: Agreeing on Who Is Here Intermediate

The hard part of "continue with whoever is present" is the word agree. The survivors must reach a single, consistent answer to three questions at the same instant: who is in the group, how big is the group, and what rank does each member hold. If two workers disagreed about the world size, their all-reduce would mismatch and corrupt the gradient; if two claimed rank $0$, the collective would be malformed. Reaching one consistent view of membership among machines that are themselves unreliable is exactly the consensus problem on the cluster control plane that Section 2.6 framed: a set of nodes electing a configuration through a coordination service that tolerates faults.

This agreement step is called rendezvous. Workers register their presence with a shared coordination service (an etcd or ZooKeeper-style store, or a built-in store keyed by a job id), and the rendezvous protocol periodically takes a census. When it detects that the membership has changed (a worker stopped heartbeating, or a new worker joined and is waiting), it closes the current round, counts the participants, assigns each a fresh contiguous rank from $0$ to $K_t - 1$, and signals everyone to re-form the process group. Each worker then rebuilds its communicator for the new world size, reloads the latest checkpoint so all replicas start from identical weights, and the training loop continues. Figure 18.4.1 traces this cycle for a worker that is preempted and a new worker that joins.

The rendezvous in Figure 18.4.1 is deliberately the same shape as a leader-election or configuration-change round on a control plane: a fault-tolerant store holds the authoritative membership, and a protocol drives all participants to one agreed view before any of them proceeds. This is why elastic training is a distributed-systems problem before it is a deep-learning one. The deep-learning part, reloading weights and resuming the optimizer, is straightforward once everyone agrees on $K_t$; getting that agreement reliably, without split-brain or duplicate ranks, is the engineering that the rendezvous backend exists to provide.

# A FIXED job pins the world size; losing a node aborts the run.
torchrun --nnodes=8 --nproc_per_node=8 train.py

# An ELASTIC job accepts a RANGE of nodes and a rendezvous backend.
# It runs with 4 to 8 nodes, shrinking on preemption and growing as capacity returns.
torchrun \
  --nnodes=4:8 \
  --nproc_per_node=8 \
  --max-restarts=100 \
  --rdzv-backend=c10d \
  --rdzv-endpoint=$HEAD_NODE:29400 \
  --rdzv-id=my-elastic-job \
  train.py

3. The Consequence: Changing World Size Moves the Global Batch Intermediate

Elasticity is not free even when the rendezvous works perfectly, because the world size is not just an operational number; it is an optimization hyperparameter in disguise. In synchronous data-parallel training, every worker processes a per-worker batch of size $b$, and the all-reduce averages their gradients, so the effective batch the optimizer sees, the global batch, is

When $K_t$ drops from $8$ to $4$ because four nodes were preempted, $B_t$ halves if you keep $b$ fixed. A smaller global batch means a noisier gradient estimate and, under the linear scaling rule from Section 10.8, a learning rate that is now too large for the batch it is paired with. Left uncorrected, an elastic shrink can destabilize a run that was tuned for the larger batch. So elasticity inherits the entire large-batch scaling discussion: every time membership changes, you have implicitly changed the batch, and you must decide what to hold constant.

There are two clean strategies, and real systems pick one deliberately. The first holds the per-worker batch $b$ fixed and lets $B_t = K_t b$ float; then you must rescale the learning rate with $B_t$, typically linearly, so $\eta_t = \eta_0 \cdot (B_t / B_0)$. The second holds the global batch $B$ fixed and rebalances the per-worker batch to $b_t = B / K_t$; then the optimizer sees a constant batch and the learning rate needs no adjustment, at the cost of each surviving worker doing proportionally more work per step. Table 18.4.1 contrasts the two.

The fix-the-batch strategy is usually the safer default for a run whose accuracy was tuned at a target batch, because it keeps the optimizer's view invariant under membership changes; the cost is that a shrunk job runs slower per step but mathematically identical. The float-the-batch strategy is simpler to implement and fine when the run is robust to batch changes, provided you actually rescale the learning rate. The demo below takes the fix-the-batch route so that the loss curve is comparable across world sizes.

4. A Job That Rides Out a Membership Storm Intermediate

The code below simulates both jobs on one synthetic regression problem so the difference is visible in a single run. A scripted timeline drives the available worker count through $8 \to 4 \to 2 \to 8$, modelling a preemption wave followed by a recovery. The elastic job re-forms at each change and holds the global batch fixed at $512$ by rebalancing the per-worker batch ($64$ at $K=8$, $256$ at $K=2$), so the optimizer sees a constant batch and the loss falls monotonically across every membership change. The rigid job pins the world size at $8$ and aborts the instant a rank goes missing. No deep-learning framework is used; the all-reduce is a plain sum of per-shard partial gradients, exactly as in Chapter 15.

import numpy as np

rng = np.random.default_rng(7)
N, d = 20_000, 16
X = rng.standard_normal((N, d))
w_true = rng.standard_normal(d)
y = X @ w_true + 0.05 * rng.standard_normal(N)

GLOBAL_BATCH = 512          # the batch size the run was tuned for
BASE_LR = 0.05              # learning rate tuned for GLOBAL_BATCH
STEPS = 12

# Scripted membership timeline: live world size per step (nodes fail, then recover).
world_timeline = [8, 8, 8, 4, 4, 4, 4, 2, 4, 8, 8, 8]   # preempt at step 3, regrow by step 9

def mse_loss(w):
    r = X @ w - y
    return float(r @ r / N)

def grad_on_indices(w, idx):
    Xs, ys = X[idx], y[idx]
    return (2.0 / len(idx)) * (Xs.T @ (Xs @ w - ys))

def run_elastic():
    """Re-form each step; hold GLOBAL_BATCH fixed by rebalancing the per-worker batch."""
    w = np.zeros(d)
    print("=== ELASTIC job: membership changes, run continues ===")
    print(f"{'step':>4}  {'world':>5}  {'per-worker':>10}  {'global':>6}     {'loss':>6}")
    for t in range(STEPS):
        world = world_timeline[t]                       # rendezvous agrees on K_t
        per_worker = GLOBAL_BATCH // world              # rebalance so K_t * b = GLOBAL_BATCH
        global_batch = per_worker * world
        idx_all = rng.integers(0, N, size=global_batch)
        shards = np.array_split(idx_all, world)         # one shard per surviving worker
        partials = [per_worker * grad_on_indices(w, s) for s in shards]   # unnormalized
        g = np.sum(partials, axis=0) / global_batch     # all-reduce SUM, then mean
        w = w - BASE_LR * g                             # constant batch => constant LR
        print(f"{t:>4}  {world:>5}  {per_worker:>10}  {global_batch:>6}  {mse_loss(w):>9.4f}")
    print(f"elastic final loss   : {mse_loss(w):.4f}  (survived 8 -> 4 -> 2 -> 8)")
    return mse_loss(w)

def run_rigid():
    """Fixed world size 8; if any rank is missing the process group cannot form."""
    w = np.zeros(d)
    FIXED, per_worker = 8, GLOBAL_BATCH // 8
    print("\n=== RIGID job: fixed world size of 8 ===")
    print(f"{'step':>4}  {'world':>5}               status")
    for t in range(STEPS):
        if world_timeline[t] < FIXED:                    # cannot form the group: abort
            print(f"{t:>4}  {world_timeline[t]:>5}  CRASH: need 8 ranks")
            print("rigid job aborted: process group could not be formed")
            return None
        idx_all = rng.integers(0, N, size=GLOBAL_BATCH)
        shards = np.array_split(idx_all, FIXED)
        partials = [per_worker * grad_on_indices(w, s) for s in shards]
        g = np.sum(partials, axis=0) / GLOBAL_BATCH
        w = w - BASE_LR * g
        print(f"{t:>4}  {world_timeline[t]:>5}      ok loss={mse_loss(w):.4f}")
    return mse_loss(w)

run_elastic()       # rides the 8 -> 4 -> 2 -> 8 storm to completion
run_rigid()         # aborts the moment the first node is lost

=== ELASTIC job: membership changes, run continues ===
step  world  per-worker  global       loss
   0      8          64     512    15.8416
   1      8          64     512    12.8740
   2      8          64     512    10.4142
   3      4         128     512     8.5130
   4      4         128     512     6.9097
   5      4         128     512     5.5668
   6      4         128     512     4.4785
   7      2         256     512     3.6760
   8      4         128     512     2.9728
   9      8          64     512     2.4372
  10      8          64     512     1.9335
  11      8          64     512     1.5450
elastic final loss   : 1.5450  (survived 8 -> 4 -> 2 -> 8)

=== RIGID job: fixed world size of 8 ===
step  world               status
   0      8      ok loss=15.4426
   1      8      ok loss=12.8529
   2      8      ok loss=10.3900
   3      4  CRASH: need 8 ranks
rigid job aborted: process group could not be formed

The contrast in Output 18.4.2 is the whole argument in miniature. The two jobs face the identical sequence of node losses. The rigid job treats the first loss as fatal and throws away three steps of progress plus the cost of a human relaunch; the elastic job treats each loss as a reconfiguration, holds its optimization invariant by rebalancing the batch, and reaches a final loss of $1.5450$ without interruption. Notice that the elastic job's rebalancing is what kept the learning rate valid: had it left the per-worker batch at $64$, the global batch would have collapsed to $128$ at $K=2$ and the fixed learning rate would have been four times too large for that batch, exactly the failure mode Section 10.8 warns about.

5. When Elasticity Earns Its Complexity Intermediate

Elastic training is not the right default for every job, because it adds real complexity: a rendezvous backend to operate, a checkpoint cadence frequent enough that a restart loses little, and the batch-and-learning-rate bookkeeping of Section 3. That overhead pays off precisely when the worker set is genuinely volatile or shared. The two clearest cases are spot and preemptible instances, where the provider reclaims capacity on its own schedule, and shared multi-tenant clusters, where a scheduler may preempt your low-priority job to make room for a higher-priority one. In both, the world size will change whether or not you planned for it, so a job that cannot tolerate change will spend its life restarting.

The alternative is a fixed-size, gang-scheduled run: the scheduler reserves all $K$ workers at once and runs the job to completion without preemption, so the world size never changes and elasticity buys nothing. Gang scheduling is the right choice on a dedicated reservation or a capacity block where you own the nodes for the duration, and it avoids both the rendezvous machinery and the batch bookkeeping. The decision is therefore about the resource, not the model: is your capacity stable and reserved, or volatile and shared? That scheduling question is the subject of Chapter 33, which develops gang scheduling and preemption policies in full; elastic training is the job-side counterpart to the cluster-side policies described there.

With elasticity in hand, the job survives the planned and unplanned arrival and departure of whole workers. But a worker need not leave to hurt the run: one that merely runs slow, a straggler, can stall a synchronous all-reduce just as effectively as a missing rank stalls a rigid one, without ever triggering a rendezvous. Detecting and mitigating those slow-but-present workers is the subject of the next section.