Singularity: Elastic Scheduling for AI Workloads

Singularity is Microsoft’s globally distributed scheduling service for running deep learning training and inference workloads across a large fleet of AI accelerators. The system makes jobs preemptable, migratable, and elastically resizable by default, while preserving the user’s existing training code and framework choices.

The main challenge is that AI jobs expect stable ownership of accelerator state. Singularity hides preemption and elasticity behind a device proxy layer, allowing the scheduler to move or resize jobs while the application continues to see a normal accelerator interface.

My work

I worked on mechanisms that made transparent elasticity practical for distributed training jobs.

I built the buffer swapping mechanism used while time-slicing between ranks during elasticity. This used a GPU memcpy kernel and a dependency graph over copies, so independent transfers could run in parallel while copy cycles were detected and broken safely.

I also worked on 3D parallelism support for elasticity, covering data, tensor, and pipeline parallel jobs. I co-designed and built the mechanism that lets the device proxy identify which parallelism dimension a communicator is used for. That distinction matters because Singularity packs only data-parallel replicas onto a single GPU: replica splicing relies on data-parallel replicas having identical parameter and optimizer state buffers at the end of a minibatch. The same property does not hold for model-parallel shards or pipeline stages, so collectives outside the data-parallel ring remain inter-GPU communication and do not need to be time-sliced by the proxy.

I worked on placement strategies and topology handling for 3D parallel jobs, including cases where multiple ranks share a device during elastic execution. I implemented pending collectives across different allreduce reductions and allgather operations. The allgather path was especially subtle: gathered results can arrive shuffled in the final buffer, so the output has to be reordered in a topology-aware way before the job can continue correctly.

Time-slicing elastic jobs also required careful correctness work. The device proxy had to preserve CUDA stream semantics, including operation ordering across intercepted calls. I worked on deadlock avoidance during time-slices, and on overlapping execution during time-slices to reduce the overhead of elasticity.

I implemented an A/B testing sanity framework to verify that device proxy interception did not introduce correctness issues. Jobs first ran without interception from a controlled RNG state, and the final loss values after a fixed number of iterations were saved. The job was then transparently restarted with interception enabled and the same RNG state restored across Python, PyTorch, cuDNN, and cuBLAS. To make host randomness deterministic, the framework recorded and replayed reads from /dev/urandom.

Technical approach

Singularity relies on transparent preemption and migration. A running job can be paused, moved to a different node, cluster, data center, or region, and resumed from the point of interruption. Jobs can also be resized up or down across accelerators of the same type.

The device proxy layer is central to this design. It intercepts accelerator API calls and manages the mapping between a job’s logical device state and the physical accelerator resources assigned by the scheduler. This lets the scheduler reclaim, move, or share accelerators without requiring users to modify training code.

For distributed jobs, elasticity has to preserve communication semantics across data, tensor, and pipeline parallelism. Collective operations, rank placement, and device-local buffers all need to remain consistent even when rank-to-device mappings change.

Impact

The Singularity paper shows that transparent preemption and elastic resizing can improve accelerator utilization and workload reliability with negligible steady-state performance overhead. The system is designed to work across model architectures and parallelism strategies, including data, pipeline, and model parallel training.