Part II - Applications and libraries

Ray Train: Distributed Training

Lesson 09 built a streaming data pipeline whose whole job was to keep a GPU fed — and it ended on an open question: once the pipeline is fast enough, what consumes the batches? This lesson answers it. Ray Train takes the one distributed-training shape that recurs across PyTorch, the orientation table, and lesson 05's tree-reduce — N processes, each holding a full model replica, each chewing a different data shard, all agreeing on one gradient every step — and packages it as a job harness. The work here is to understand that shape precisely enough to predict when adding a GPU stops helping, and to see exactly where Ray Data hands its shards across to Train.

1 · The shape: one replica per GPU, one shard per replica, one gradient per step

Distributed training is not "run the script on more GPUs." It is a specific, tightly synchronized dance, and the most common form of it — the one Ray Train defaults to — is data parallelism. Hold this picture: you have N worker processes, each pinned to its own GPU. Every worker holds a complete copy of the model — same architecture, same weights, identical at the start of every step. We call each such process a rank (rank 0, rank 1, … rank N−1 — the integer identity PyTorch's distributed layer assigns each worker). PyTorch's DistributedDataParallel (DDP) is the machinery that keeps those replicas identical; a Ray Train worker is, mechanically, one DDP rank running inside a Ray actor.

Within one training step, each rank does the following, in lockstep with all the others:

Allreduce is the load-bearing word, so define it exactly. An allreduce is a collective operation in which every participant contributes a vector (here, its local gradient) and every participant ends up with the element-wise reduction (here, the sum, which is then divided by N to average) of all contributions. The naive way — send everything to rank 0, sum, broadcast back — makes rank 0 a bottleneck that moves N×model_bytes. The standard implementation is ring-allreduce: arrange the N ranks in a ring, split each gradient into N chunks, and pass chunks around the ring in two passes (a reduce-scatter then an all-gather). Its bandwidth cost per rank is

comm_bytes ≈ 2 × (N − 1) / N × model_bytes

which approaches 2 × model_bytes as N grows — crucially, independent of N for large N. That bandwidth-optimality is exactly why ring-allreduce, not the naive star reduce of lesson 05, is what production training uses. The deep mechanics of collectives and the wires they run over are the subject of the system_ml track — see collectives and interconnect, and DDP for the full data-parallel treatment. Ray Train is the orchestration layer that launches those ranks and hands them their data; it does not reinvent the collective.

2 · How Ray Train wraps the shape

Ray Train's contract is small on purpose. You write a function that trains on one worker — exactly the loop you already know — and Ray Train runs N copies of it as actors, sets up the DDP process group, and feeds each its shard. You hand it a TorchTrainer with a ScalingConfig(num_workers=N, use_gpu=True) describing the physical shape, then call .fit().

import ray.train.torch
from ray.train import ScalingConfig, Checkpoint, RunConfig
from ray.train.torch import TorchTrainer
import tempfile, os, torch

def train_loop_per_worker(config):
    # Ray Train has already set up the DDP process group and assigned this rank a GPU.
    model = build_model(config)
    model = ray.train.torch.prepare_model(model)   # wraps in DDP + moves to this rank's GPU
    opt = torch.optim.Adam(model.parameters(), lr=config["lr"])

    # Ray Data → Train handoff: this rank gets ITS shard, no one else's (see lesson 09).
    shard = ray.train.get_dataset_shard("train")

    for epoch in range(config["epochs"]):
        for batch in shard.iter_torch_batches(batch_size=config["bs"]):
            loss = step(model, opt, batch)         # forward/backward; DDP allreduces the grad here
        # report metrics AND a checkpoint at the epoch barrier (all ranks must call this)
        with tempfile.TemporaryDirectory() as d:
            if ray.train.get_context().get_world_rank() == 0:   # rank 0 writes weights once
                torch.save(model.module.state_dict(), os.path.join(d, "model.pt"))
            ray.train.report({"loss": loss.item(), "epoch": epoch},
                             checkpoint=Checkpoint.from_directory(d))

trainer = TorchTrainer(
    train_loop_per_worker,
    train_loop_config={"lr": 1e-3, "epochs": 10, "bs": 256},
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),   # 4 GPUs, 4 DDP ranks
    datasets={"train": ray_dataset},                             # from lesson 09
    run_config=RunConfig(storage_path="s3://my-bucket/runs"),    # durable, shared checkpoint store
)
result = trainer.fit()
best = result.checkpoint                                          # handoff to Tune/Serve

Three pieces do all the work, and each maps onto something from earlier lessons:

Fault tolerance falls out of the checkpoint. A checkpoint is a serialized snapshot of everything needed to resume — model weights, optimizer state, scheduler state, epoch counter — written to durable shared storage (S3, NFS), never worker-local scratch (a dead worker takes its scratch with it). If a worker actor dies mid-run, Ray Train tears down the process group, requests a fresh worker from the cluster, and restarts every rank from the last reported checkpoint. The blast radius of a single GPU failure is therefore "redo the work since the last checkpoint," not "lose the whole run" — which is precisely why checkpoint frequency becomes a tunable cost in §5.

3 · The central trade-off: scaling efficiency and the interconnect knee

The promise of data parallelism is linear speedup: 4 GPUs train 4× faster. It is never quite true, because step 3 — the allreduce — is pure overhead that grows with the cluster. The right mental model is that each step costs compute + comm, where compute is the forward/backward on one shard and comm is the allreduce. Define scaling efficiency as the fraction of wall-clock spent doing useful compute:

efficiency = compute / (compute + comm)

Work it with numbers. Take a model with 1 GB of gradients (a ~250M-parameter model in fp32, or a ~500M-param model in fp16). Per-rank allreduce volume is ≈ 2 × model_bytes = 2 GB moved across the interconnect each step. Suppose the per-GPU forward+backward takes compute = 200 ms.

Note that ring-allreduce's comm volume is roughly constant in N (it stays ≈ 2× model_bytes per rank), so on a fast fabric efficiency degrades gently. The real damage comes from crossing a slow boundary: 8 GPUs inside one NVLink-connected node stay at 97%, but spreading those same 8 ranks across two nodes joined by 10 GbE forces the slow link into the ring and collapses efficiency. This is why placement matters and why Ray Train's ScalingConfig should keep a job's workers on well-connected nodes — a lesson that connects directly to placement groups (lesson 04) and to cluster topology (lesson 12). The fix when the wire is the wall is either a faster fabric, or larger per-GPU batches so compute grows relative to the fixed comm (raise compute to 800 ms on 10 GbE and efficiency climbs from 11% to 33% — still poor, but it shows the lever).

4 · The memory wall: where data parallelism simply cannot help

Everything above assumes one thing: the model fits on one GPU. Data parallelism replicates the entire model — weights, gradients, and optimizer state — on every rank. Adding workers adds throughput, never capacity. So when the model itself is too big for a single accelerator, data parallelism does nothing at all; you cannot even start.

This is the boundary where you must switch strategies. Tensor parallelism splits each layer's matrices across GPUs so one layer lives on several devices; pipeline parallelism splits the model by depth so different GPUs hold different layers and micro-batches flow through like an assembly line. These are not Ray Train concepts — they are model-partitioning strategies you bring (via DeepSpeed, FSDP, Megatron), and Ray Train then orchestrates the resulting workers. The full treatment is in system_ml: tensor parallelism and pipeline parallelism. The practical decision rule: data parallel until the model stops fitting, then partition the model, and usually combine both (data-parallel groups of model-parallel shards). Knowing which wall you are on — the throughput wall (use more data-parallel ranks) or the memory wall (partition the model) — is the whole judgment of this lesson.

5 · Checkpoint frequency: a cost balance, not a habit

Checkpointing trades a known, recurring cost (the time to write a snapshot, during which all ranks barrier) against an unknown, occasional cost (the work you redo when a worker dies and you restart from the last checkpoint). Checkpoint too rarely and a failure late in an epoch throws away hours; too often and you spend the run writing to S3 instead of training.

Two refinements. Cheaper checkpoints (async writes, sharded/incremental snapshots) lower C and so let you checkpoint more often for the same overhead. And MTBF gets worse as you add workers — more GPUs means more things to fail — so a 512-GPU run on commodity hardware genuinely needs frequent checkpoints, while an 8-GPU run on reliable hardware barely needs them. The deeper failure-tolerance machinery (lineage, reconstruction, idempotency) is lesson 14 and system_ml's checkpointing & fault tolerance.

6 · Failure modes & checklist

Checkpoint exercise

Where this points next

Training produces a checkpoint; the checkpoint's whole reason to exist is to be served. Lesson 11 picks up exactly there: the trained weights now have to answer live requests under a latency SLO, which is a completely different shape from the synchronized lockstep of training. Where Train fans out replicas to agree on one gradient, Serve fans out replicas to handle independent requests, and the central knob flips from "allreduce overhead" to "how long do I wait to batch requests before the GPU is full." The handoff object is the same Checkpoint you just produced.

Interview prompts

• Walk through one data-parallel training step and define rank, replica, and allreduce. (§1 — each rank holds a full replica, processes its own shard, computes a local gradient; allreduce makes every rank end with the same summed/averaged gradient so replicas apply an identical update and never drift.)
• Why ring-allreduce instead of sending everything to rank 0? (§1 — the star reduce moves N×model_bytes through one bottleneck; ring-allreduce costs ≈ 2×(N−1)/N × model_bytes per rank, bandwidth-optimal and near-constant in N.)
• Your 4-GPU job is only 1.6× faster than 1 GPU. How do you diagnose it? (§3 — compute scaling efficiency = compute/(compute+comm); a slow interconnect makes allreduce dominate. Check data-loader throughput (§6, lesson 09) and the fabric before the code.)
• Same model, NVLink vs 10 GbE — what happens to efficiency and why? (§3 — comm = 2×model_bytes ÷ bandwidth; a 1 GB-gradient model goes from ~7 ms (97%) on NVLink to ~1600 ms (11%) on 10 GbE. The interconnect sets the knee.)
• You have a 70B model and 100 GPUs. Does data parallelism solve it? (§4 — no; data parallelism replicates the whole model per GPU (~1.1 TB with Adam state), so it scales throughput not capacity. You hit the memory wall and must use tensor/pipeline parallelism.)
• How does Ray Data hand data to Ray Train, and where does GPU starvation bite? (§2 — get_dataset_shard() gives each rank a disjoint shard; if the pipeline can't fill a rank, that rank's GPU idles and the step barrier makes every rank wait for it.)
• How often should you checkpoint? (§5 — minimize C/T + T/(2·MTBF) → optimal interval ≈ √(2 × checkpoint_cost × MTBF); checkpoint more often as MTBF drops with more workers or cheaper snapshots, not "every epoch" by reflex.)