Spotlight: Synergizing Seed Exploration and Spot GPUs for DiT RL Post-Training

Resumo Técnico: SPOTLIGHT

Declaração do Problema

O pós-treinamento de Aprendizado por Reforço (RL) para Diffusion Transformers (DiTs) é proibitivamente caro, exigindo frequentemente milhares de GPUs de alto desempenho. O fluxo de trabalho padrão do estilo GRPO envolve duas fases: uma fase de rollout (geração de amostras com sementes aleatórias) e uma fase de treinamento (cálculo de gradientes e atualização de pesos).

As estratégias de redução de custo existentes enfrentam dois gargalos primários quando aplicadas a DiTs:

1. Sobrecarga de Exploração de Sementes: Para acelerar a convergência, métodos como DanceGRPO e TreeGRPO exploram múltiplas sementes aleatórias para encontrar amostras de alto contraste (recompensas altas e baixas). No entanto, esta fase de exploração reside no caminho crítico, adicionando tempo de computação a cada iteração e anulando potenciais economias de custo.
2. Ociosidade e Fragmentação de GPUs Spot: Embora as GPUs spot ofereçam custos 69–77% menores, elas são preemptíveis e possuem largura de banda inter-nós limitada.
   • Ociosidade: Diferente dos Large Language Models (LLMs), onde as latências de rollout variam (permitindo o pipelining com o treinamento), os rollouts de DiT completam-se em um tempo quase uniforme. Isso cria uma barreira de sincronização onde as GPUs spot ficam ociosas durante todo o processo de treinamento.
   • Fragmentação: Preempções frequentes de instâncias spot quebram grupos de Paralelismo de Sequência (SP), que são essenciais para executar grandes modelos DiT. Reconfigurar esses grupos normalmente requer minutos para recarregar modelos e reinicializar motores, desperdiçando recursos significativos.

Metodologia: SPOTLIGHT

Os autores apresentam o SPOTLIGHT, o primeiro sistema projetado para colher GPUs spot para o pós-treinamento de RL de DiT ao sinergizar a exploração de sementes com recursos spot. O sistema é construído sobre dois insights fundamentais e três técnicas principais.

Insights Principais

1. A Exploração Tolera Pesos Obsoletos: Os autores demonstram que o ranking relativo das sementes aleatórias (quais sementes produzem recompensas altas vs. baixas) permanece altamente consistente mesmo utilizando os pesos do modelo da iteração anterior em vez dos pesos atuais atualizados. Isso permite que a exploração de sementes seja executada em GPUs spot usando pesos obsoletos enquanto o cluster reservado realiza a fase de treinamento, efetivamente removendo a exploração do caminho crítico.
2. A Reconfiguração de SP Pode Reutilizar o Estado: O alto custo de reconfiguração de grupos SP decorre da inicialização do escalonador de CPU e do carregamento remoto do modelo. Ao desacoplar o escalonador e utilizar a cópia de pesos intra-nó (via NVLink) de pares co-localizados, o tempo de reconfiguração de SP pode ser reduzido de minutos para sub-segundos.

Arquitetura do Sistema e Técnicas

O SPOTLIGHT é implementado na plataforma de RL de código aberto ROLL e emprega três técnicas principais:

1. Planejador de Exploração Baseado em Bandidos:

   • O sistema formula o planejamento de exploração como um problema de multi-armed bandit.
   • Ele seleciona dinamicamente o número de sequências e os passos de denoising efetivos (usando estratégias de cache como Teache) para maximizar a variância da recompensa dentro do orçamento de tempo fornecido pela fase de treinamento.
   • Utiliza uma estratégia de Upper Confidence Bound (UCB) para equilibrar a exploração de novas configurações e a explotação de configurações de alta variância conhecidas.

2. Paralelismo de Sequência Elástico (SP):

   • Escalonador Persistente: O escalonador de CPU é desacoplado dos trabalhadores de GPU e mantido residente no nó através das mudanças de SP, eliminando custos repetidos de inicialização.
   • Carregamento de Pesos Intra-Nó: Quando um novo trabalhador de GPU entra em um grupo SP no mesmo nó, ele copia os pesos de um par co-localizado via NVLink em vez de puxá-los de um nó remoto, reduzindo drasticamente a latência de reconfiguração.

3. Escalonador de Requisições Baseado em Pull e Consciente de Preempção:

   • Trabalhadores de rollout puxam requisições de uma fila centralizada em vez de receberem tarefas estáticas, adaptando-se a velocidades heterogêneas.
   • Comprometimento de Estado: Estados intermediários de rollout (latentes e metadados) são comprometidos em um Tensor Store persistente em nós reservados. Se uma GPU spot for preemptada, o sistema pode recuperar o estado em curso e retomar a requisição em outro trabalhador, minimizando a perda de progresso.

Resultados

O sistema foi avaliado no pós-treinamento do Qwen-Image (20B) utilizando os datasets DeepSeek-OCR e Geneval em resoluções de 512×512 e 1280×1280.

• Redução de Custo: O SPOTLIGHT reduz o custo total de treinamento em 1.4× a 6.4× em comparação com baselines de última geração (incluindo sistemas que utilizam GPUs spot sem estas otimizações e sistemas apenas com recursos reservados).
• Velocidade de Convergência: Atinge o mesmo score de validação alvo 4× mais rápido que os baselines. Isso é alcançado reduzindo o número de iterações necessárias (devido à exploração de sementes eficaz) e sobrepondo a exploração com o treinamento.
• Qualidade: O SPOTLIGHT atinge scores de validação superiores em todos os datasets comparado aos baselines, atribuído à seleção dinâmica de sementes de alto contraste.
• Eficiência:
  • O overhead de exploração é mantido baixo (2–3% do tempo de iteração) ao encaixar a exploração dentro da janela de treinamento.
  • O SP elástico permite que o sistema acompanhe mudanças de capacidade spot em tempo quase real, enquanto abordagens ingênuas sofrem atrasos de vários minutos durante a reconfiguração.
  • O sistema demonstra escalabilidade quase linear no throughput conforme o número de GPUs spot aumenta de 8 para 64.

Significância e Alegações

O artigo afirma que o SPOTLIGHT é o primeiro sistema a colher com sucesso GPUs spot para o pós-treinamento de RL de DiT. Sua significância reside em:

1. Quebrar a Dependência do Caminho Crítico: Ao provar que a exploração de sementes pode tolerar pesos obsoletos, o sistema elimina o trade-off entre convergência mais rápida (via exploração) e latência de iteração.
2. Resolver a Ociosidade de Spot: Resolve o problema específico de ociosidade de rollouts de DiT ao utilizar a capacidade ociosa de spot da fase de treinamento para exploração, uma sinergia negligenciada por sistemas focados em LLM anteriores.
3. Robustez à Preempção: Através de SP elástico e persistência de estado, torna o treinamento de DiT resiliente às frequentes preempções inerentes aos mercados spot, transformando um pool de recursos fragmentado em uma infraestrutura de treinamento viável e de baixo custo.

Os autores concluem que estas contribuições permitem um sistema de pós-treinamento de RL de DiT eficiente em termos de custo que não sacrifica a qualidade da imagem, tornando o treinamento de modelos de difusão de alta qualidade mais acessível.