Ai2 replaces priority-based GPU scheduler with budget-driven system to prioritize high-impact research

Ai2’s AI Infrastructure team, responsible for providing GPU compute capacity for large, distributed training workloads, introduced a pyramid of four metrics to guide scheduling: availability, occupancy, impact, and utilization. The team replaced a priority-based scheduler with a system that includes GPU time budgets, hierarchical fair-share allocation, and a time-slicing contract to improve the impact of scheduling decisions.

Ai2 manages thousands of NVIDIA H100, B200, and B300 GPUs arranged in clusters ranging from 88 to 1,024 GPUs, serving about 150 internal researchers across domains including LLM and VLM training, robotics reinforcement learning simulation, and post-training for scientific agentic use cases. Demand for GPU time far exceeds supply, with outstanding requests at any moment for 2–3x more GPUs than are available.

Historically, Ai2 used a priority-based scheduler that allowed workloads to opt out of preemptability, leading to predictable issues such as GPU 'squatting'—where users would reserve no-op workloads to connect to when needed—and priority inflation, where 100% of scheduled workloads eventually used HIGH priority, starving lower priority levels.

The team identified these issues as a 'tragedy of the commons,' where individuals competing over a scarce shared resource achieved a non-optimal global result. The solution involved shifting from assigning teams monopolies over GPUs to allocating a portion of GPU time, enabling leadership to fund research efforts based on strategic judgment of likely impact.

The new system features a hierarchical allocation model where managers proportionally allocate GPU time to projects and researchers, with every request for GPU time required to be funded by a budget or left unprotected from preemption. This design aims to make gaming the scheduler more costly than engaging in budget debates, while ensuring frequent opportunities for researchers to advocate for needed time and decisions made by managers with contextual tradeoff knowledge.