From jameshan.net

Designing a Multi-Site GPU Training Platform: Gang Scheduling, Preemption, and Data-Path Throughput

A research team training models on multichannel neural recordings had about 50 A100 and H100 GPUs across three sites, partitioned by team and fed by loaders that preprocessed raw recordings during training. GPUs were therefore idle for two independent reasons: no job held them, or the job holding them was waiting on data. This note describes the platform built to address both, a submission interface, queue and scheduler on top of Ray, and a data path restructured around offline sharding and prefetch.

  • Allocation. A shared queue and placement scheduler cut mean queue wait from about 10 hours to 3 and raised fleet allocation from 68% to 85%.
  • Utilization. Moving preprocessing offline and prefetching from local NVMe roughly doubled utilization during training, from about 34%.
  • The two are distinct metrics: allocation measures whether a GPU is assigned to a job, and utilization whether an assigned GPU is computing.

Platform

Researchers submit jobs through a web interface or the same API from a command line. Jobs enter a queue stored in Postgres with a priority and a resource request. A scheduler polls the queue and the fleet’s free capacity every few seconds, decides placement, and hands placed jobs to Ray, which launches the processes and reports their state. Ray supplies mechanism, including gang reservation, so that a multi-GPU job acquires all of its devices at once or none of them; the scheduler supplies policy. Jobs checkpoint to an object store, which is what makes preemption safe.

Architecture: web UI and CLI submit to a Postgres job queue; the scheduler polls it, launches jobs through Ray onto GPU hosts at three sites, receives free-GPU and job state back, and writes checkpoints to an object store
Platform architecture. The scheduler is the only component that decides placement; Ray executes it.

Scheduling Policy

  • Single-host placement. Data-parallel training synchronizes gradients every step, over NVLink within a host and over the network between hosts, so a job is confined to one machine.
  • Packing. Small jobs are placed on the most occupied host that fits them, keeping whole hosts free for large jobs. Spreading them instead leaves free GPUs on every host and none with eight.
  • Preemption. A high-priority job that cannot fit displaces the lowest-priority job whose GPUs suffice; the displaced job checkpoints, requeues, and resumes from its checkpoint.⁠Timeline of two GPU groups: a low-priority job is preempted and checkpointed when a high-priority job arrives, waits in the queue, and resumes from its checkpoint Preemption with checkpoint and requeue. Illustrative.
  • Starvation control. Packing and priority together can starve a large low-priority job. A host is reserved and drained for it, and priority ages with waiting time.
Six hosts of eight GPUs with the same twelve jobs placed two ways: spread, leaving 24 GPUs free but no free host; and packed, leaving two whole hosts free
The same twelve jobs placed by spreading and by packing. Free capacity is identical; only packing leaves room for an eight-GPU job.
Queue-wait distributions before and after the shared scheduler, and weekly fleet allocation rising from about 68 to about 85 percent
(a) Time to first GPU. (b) Weekly fleet allocation. Illustrative: synthetic data shaped to the reported means of 10 h to 3 h and 68% to 85%.

Data Path

nvidia-smi reports utilization as the fraction of time any kernel is resident, so a low value reliably indicates idleness while a high value does not imply efficiency. Training on H100s showed utilization falling to zero between steps. Timing each step in two parts, batch retrieval and compute, with torch.cuda.synchronize() before the second timestamp so that asynchronous kernel launches are not mistaken for completed work, showed that retrieval exceeded compute.⁠Seconds per training step split into compute and waiting on the next batch, for online loading, offline shards with prefetch, and BF16 Per-step time split into compute and batch wait. BF16 shortens compute, so the step is faster even as the share of time spent waiting rises. Illustrative, synthetic.

GPU utilization over two minutes of training, alternating between full and zero with online loading, and mostly full with offline shards and prefetch

The loader queried a database and cleaned and windowed raw recordings on the CPU at every step, repeating identical work each epoch. Additional loader workers raise throughput only until the database saturates.⁠Samples per second per GPU against loader workers: online loading plateaus well below the GPU bound as the database saturates; offline shards reach the bound with five workers Throughput against loader workers. Online loading is bounded by the database, not by worker count. Illustrative, synthetic. The data path was instead restructured: cleaning runs once as an offline job that writes training-ready shards, shards are cached on local NVMe on each GPU host, and loader workers prefetch subsequent batches into pinned memory during compute. Because prefetching overlaps loading with computation, once retrieval is faster than compute it disappears from the step time, which is why a modest reduction in retrieval time produced a large gain in utilization. Faster accelerators make the problem worse: the same jobs showed higher utilization on older GPUs, which request batches less often.

Once the GPU is fed, the next lever is precision. BF16 halves the bytes moved per value and runs on tensor cores, and its exponent range matches FP32, so training remains stable without loss scaling.