Distributed ML Training Platform

Scheduling and running jobs on thousands of GPUs

Design a shared platform that schedules training jobs onto a GPU cluster, splits large models across thousands of GPUs, and keeps week-long runs going through the hardware failures that are certain to happen.

Overview

A training platform turns a shared pool of GPU servers into a service: researchers submit a job (code, data, how many GPUs), and the platform finds the machines, starts the processes, watches them, saves progress, and restarts the job when something breaks.

The job is hard for three reasons. Size: a large model and its optimizer state do not fit on one GPU, so the model itself must be split across many. Lockstep: in synchronous training every GPU waits for every other GPU at each step, so the job runs at the speed of the slowest one and stops when any one fails. Failures are routine: Meta reported 419 unexpected interruptions in a 54-day window while training Llama 3 405B on up to 16,384 GPUs, about one every three hours, with about 78% traced to hardware.

Questions to ask first

How big is the cluster? — 4,096 GPUs: 512 servers of 8, NVLink inside a server, a 400 Gbps fabric between servers
How big are jobs? — Mostly 8 to 64 GPUs; a few flagship runs use the whole cluster for weeks
Largest model? — 70 billion parameters
Many teams? — Yes: quotas per team, priorities, preemption
What must survive a failure? — The job, automatically, losing at most a few minutes of work
Hardware mix? — One GPU type, to keep the answer focused

What you will be able to do

Explain gang scheduling and write an all-or-nothing GPU allocation that is safe under concurrency.
Work out a model's training memory and choose data, fully sharded, tensor or pipeline parallelism for it.
Explain ring all-reduce and why each GPU sends about twice the gradient size no matter how many GPUs there are.
Size checkpoint frequency from a failure rate, and explain what asynchronous checkpoints buy.
Say how the platform finds stragglers and hangs, and how preemption keeps a shared cluster fair.

Key concepts

Gang scheduling
Data parallelism
FSDP / ZeRO
Tensor parallelism
Pipeline parallelism
Ring all-reduce
Checkpoint and restart
Async checkpoints
Stragglers
Preemption

Functional requirements

Submit a job with a container image, command, dataset, GPU count and priority
Queue jobs per team quota; start a job only when all its GPUs can start together
Launch distributed workers with ranks, a rendezvous point and the right network setup
Save checkpoints and restart failed jobs from the latest one automatically
Preempt lower-priority jobs cleanly, by checkpointing first
Stream metrics (loss, throughput, GPU use) to a dashboard and keep them per run
Register a chosen checkpoint as a model version

Non-functional requirements

Effective training time above 90% for large jobs (a target we chose; Meta reported achieving this for Llama 3)
Detect a dead or hung worker within a few minutes and restart without a human
Lose at most one checkpoint interval of work per failure
High cluster use: idle GPUs are the most expensive waste in the system
Fair sharing: a team always gets its quota back within a bounded time

Interview questions

What is gang scheduling and why does distributed training need it?
A 70B model does not fit on one GPU. How do you train it?
Explain ring all-reduce. Why does it scale?
Your job has 4,096 GPUs. How often will it fail, and how often should it checkpoint?
How do you detect and handle a straggler or a hung collective?
How do you keep a shared cluster fair without leaving GPUs idle?
Why synchronous training rather than asynchronous updates?