Como o Ai2 resolveu a briga por GPUs com orçamentos em vez de fila de prioridade

O Allen Institute for AI (Ai2, instituto de pesquisa em inteligência artificial nos Estados Unidos) gerencia milhares de GPUs (chips especializados em processar tarefas de IA) dos modelos H100, B200 e B300 da NVIDIA, organizadas em clusters (grupos de computadores trabalhando juntos) com até 1024 GPUs. Esses recursos atendem cerca de 150 pesquisadores internos que trabalham em áreas que vão do treinamento de grandes modelos de linguagem (LLMs, os sistemas por trás de assistentes como o ChatGPT) até robótica e simulações. O problema: a demanda por tempo de GPU supera em duas a três vezes a capacidade disponível — ou seja, cada hora disponível tem até três projetos diferentes competindo por ela.

Segundo informou a equipe de infraestrutura do Ai2 em post no blog da Hugging Face (uma plataforma onde desenvolvedores compartilham e baixam modelos de IA já treinados), o instituto usava antes um sistema de prioridades onde cada equipe tinha um limite de GPUs protegidas de preempção (interrupção forçada para liberar recursos a outro projeto). Workloads preemptíveis (tarefas que podiam ser interrompidas) podiam usar GPUs ociosas acima desse limite. O resultado foi uma "tragédia dos comuns" clássica: pesquisadores começaram a estacionar workloads vazios só para garantir acesso rápido quando precisassem, e eventualmente 100% dos projetos passou a usar prioridade ALTA, tornando o sistema inútil. A equipe de plantão gastava a maior parte do tempo negociando o desligamento manual de tarefas que ocupavam máquinas com problemas de manutenção.

A solução foi trocar o controle de prioridades por um sistema de orçamentos de tempo de GPU. Em vez de alocar máquinas fixas a equipes — o que deixava GPUs paradas quando um time não tinha trabalho pronto —, o instituto passou a distribuir fatias percentuais do tempo total de processamento. Gerentes de programa decidem antecipadamente, como investidores, quanto do orçamento de GPU vai para cada projeto com base no impacto esperado, não na urgência de quem grita mais alto. Um projeto pode saber que tem direito garantido a 35% da capacidade total, independentemente de quantos outros estejam na fila.

Esse modelo hierárquico usa um algoritmo de fair-share (compartilhamento justo) que não é novo — variantes existem desde o Hadoop Fair Scheduler de 2009 e são usadas hoje em sistemas como SLURM e YARN. A novidade está nos inputs: a árvore de alocação espelha a estrutura de programas de pesquisa do instituto, e os pesos são orçamentos definidos por líderes com contexto sobre as trocas envolvidas. Qualquer pedido de GPU precisa estar financiado por um orçamento, ou fica sujeito a preempção. Truques para burlar o sistema agora custam do próprio orçamento da equipe: estacionar um workload vazio gasta recursos sem retorno. A estratégia, segundo o Ai2, é tornar manipular o sistema mais caro do que simplesmente negociar um orçamento maior nas revisões periódicas com a gerência.