SkyRL e vLLM refinam o Mistral Large 4 com RL: AIME sobe de 64% para 75%

A equipe do projeto vLLM e o laboratório NovaSky AI informaram que o framework SkyRL já suporta, desde o primeiro dia, o aprendizado por reforço (RL) para o Mistral Large 4. Um cluster de 8 nós Nvidia GB200 elevou o resultado do modelo no AIME 2024 de cerca de 64% para 75% em 75 passos de treinamento. O AIME 2024 reúne 30 problemas da olimpíada americana de matemática, em que a resposta é sempre um número inteiro. A métrica avg@8 significa que o modelo faz 8 tentativas para cada problema, e a taxa final de acertos é a média dessas tentativas. Assim, o resultado depende menos da sorte em uma única tentativa.

O que significa Day-0 e para que serve

Day-0 quer dizer que o suporte está pronto no dia do lançamento do modelo. Normalmente, uma nova modelo grande só pode ser executada de início, e o refinamento fica para depois, quando a comunidade escreve o código para a arquitetura dela. Aqui, a afirmação é que a combinação de SkyRL e vLLM funcionou com o Mistral Large 4 imediatamente.

O aprendizado por reforço funciona assim. O modelo recebe um problema e escreve várias versões de solução. Essas versões são chamadas de rollouts ("execuções"). Em seguida, cada rollout é verificado: a resposta bateu com a de referência? Por uma resposta correta, o modelo recebe uma recompensa, e os pesos são ajustados para tornar mais prováveis as cadeias de raciocínio bem-sucedidas. Em matemática, essa verificação é barata e inequívoca, por isso o AIME e conjuntos de problemas semelhantes costumam ser usados em RL.

Quem faz o quê na combinação

A principal ideia de engenharia é a divisão de trabalho. As duas etapas do ciclo exigem coisas diferentes do hardware, e cada uma fica a cargo de uma ferramenta própria:

• vLLM gera os rollouts. É um motor de inferência rápida, ou seja, de execução do modelo para produzir respostas. A geração ocupa a maior parte do tempo do ciclo de RL, e por isso é preciso um servidor veloz.
• FSDP2 ou Megatron cuidam do treinamento. O FSDP2 é a segunda versão da abordagem do PyTorch em que parâmetros e gradientes são fragmentados entre as GPUs. O Megatron é o framework da Nvidia para treinar modelos muito grandes. O SkyRL permite escolher qualquer um dos dois.
• SkyRL, da NovaSky AI, une tudo em um único fluxo: recebe os rollouts do vLLM, calcula as recompensas e envia os dados para o treinamento.

Na prática, o engenheiro não precisa escrever a camada de integração entre o motor de geração e o código de treinamento. Ele pega a solução pronta e encaixa o modelo.

Por que o treinamento usou

LoRA sobre um checkpoint de 4 bits

A publicação cita duas técnicas que barateiam o processo.

A primeira é o LoRA. Em vez de atualizar todos os pesos do modelo gigante, são treinadas pequenas matrizes adicionais, conectadas por cima dele. Elas ocupam pouca memória em comparação com o modelo em si, e o próprio Mistral Large 4 permanece congelado.

A segunda é o checkpoint no formato NVFP4. É uma representação numérica de 4 bits que a Nvidia suporta na arquitetura Blackwell, à qual pertencem os GB200. Os pesos nesse formato ocupam várias vezes menos espaço do que em 16 bits. Um modelo grande cabe em menos nós, e a geração fica mais rápida.

No total, resulta uma configuração modesta para os padrões de modelos grandes: 8 nós, 75 passos e um ganho de cerca de 11 pontos percentuais em um único benchmark.

Como interpretar esse resultado

Os números vêm da própria publicação do projeto vLLM, e ainda não há verificação independente. O ganho no AIME 2024 mostra que o ciclo de RL funciona com esse modelo e esse hardware. Ele não prova que o modelo ficou melhor de modo geral. Em problemas de matemática com verificação inequívoca, o RL costuma ter efeito notável, mas em tarefas abertas (escrita, cenários de agentes, código em um repositório grande) o resultado precisa ser medido separadamente.

Para os engenheiros, a conclusão é esta. Refinar um grande modelo aberto com aprendizado por reforço ficou mais simples: uma combinação pronta de componentes abertos funciona desde o primeiro dia de lançamento, e LoRA e NVFP4 reduzem as exigências de hardware. Para empresas que têm tarefas próprias com verificação automática da resposta, é um caminho para adaptar o modelo às suas necessidades sem retreiná-lo por completo. Resta testar com os próprios dados quantos desses 11 pontos dá para reproduzir.