GRPO assíncrono com LoRA: como a Hugging Face treina LLMs distribuídos sem NCCL

A Hugging Face publicou um guia técnico detalhado mostrando como treinar um modelo de linguagem com GRPO assíncrono — técnica de aprendizado por reforço usada em modelos de raciocínio como o DeepSeek-R1 — distribuído entre máquinas separadas na nuvem, sem depender de NCCL, a biblioteca de comunicação que normalmente conecta as GPUs de um cluster. O segredo é simples e poderoso: sincronizar apenas o adaptador LoRA, que pesa alguns megabytes, em vez de mover o modelo inteiro, que pesa gigabytes.

Por que isso importa agora

Treinar modelos com aprendizado por reforço (RL) exige orquestrar duas etapas que competem entre si: a geração de respostas (os chamados rollouts) e o treinamento do modelo a partir dessas respostas. Até recentemente, isso exigia que as GPUs de geração e de treino compartilhassem uma mesma máquina ou um cluster com NCCL — infraestrutura cara e complexa de operar.

Com o suporte a LoRA que chegou ao AsyncGRPOTrainer da TRL (versão 1.14, via pull request #7017), esse cenário muda. Um adaptador LoRA de rank 1 para um modelo de 1,5 bilhão de parâmetros ocupa poucos megabytes, contra cerca de 3 GB do modelo completo. Em vez de enviar a política completa aos servidores de inferência a cada atualização, basta enviar o adaptador — e isso abre caminho para rodar treino e geração em máquinas totalmente separadas.

Há uma razão matemática por trás disso. Como o time aponta, a função de vantagem no RL fornece apenas ~1 bit de informação por episódio, então não há muito o que aprender a cada passo — um adaptador de rank 1 tem capacidade suficiente para absorver isso. É a mesma lógica defendida pelo Thinking Machines Lab no artigo LoRA Without Regret.

A arquitetura: três Jobs, um bucket e um proxy

O desafio era este: um Hugging Face Job é um contêiner rodando em uma única máquina virtual, limitado a no máximo 8 GPUs H200 por nó. Sem rede compartilhada e sem disco local compartilhado, como conectar um treinador a uma frota de servidores de inferência?

A resposta do time foi montar um setup surpreendentemente enxuto:

• Um Job treinador rodando o AsyncGRPOTrainer com LoRA;
• Dois Jobs vLLM, cada um servindo o modelo base mais o adaptador mais recente publicado pelo treinador;
• Um Storage Bucket montado nos três Jobs no mesmo caminho — é por onde o adaptador viaja do treinador até os servidores;
• Um proxy que roteia os rollouts e distribui as atualizações de adaptador.

A cada poucos passos de otimização, o treinador salva o adaptador em <output_dir>/.vllm_lora/trl-policy-v{N}, publica o diretório com uma renomeação atômica e envia o caminho ao endpoint /v1/load_lora_adapter do vLLM. Como o bucket é montado como um sistema de arquivos POSIX (via hf-mount) em todos os Jobs, o caminho é válido dentro de qualquer contêiner. Nenhuma mudança foi necessária na TRL ou no vLLM para isso funcionar.

O proxy: roteamento por prefixo de KV cache

O proxy resolve dois problemas. Primeiro, ele adiciona o cabeçalho de autenticação exigido pelas portas expostas dos Jobs, de modo que a TRL não precise saber disso. Segundo — e aqui está a parte engenhosa — ele distribui cada requisição para a réplica que já guarda o prefixo daquele prompt no KV cache.

Gerar uma resposta tem duas fases com perfis muito diferentes: o prefill, que processa o prompt inteiro de uma vez e computa as chaves e valores de atenção de cada token; e o decode, que produz um token por vez. Essas chaves e valores formam o KV cache. Como a atenção é causal, duas requisições que compartilham um prefixo também compartilham o KV daquele prefixo — e uma réplica que já o tem em cache pode pular essa parte do prefill.

No GRPO, o mesmo prompt é enviado G vezes (no experimento, G=8). Se todas as oito requisições caírem na mesma réplica, a primeira computa o prefill e as outras sete o reutilizam. Com roteamento round-robin, metade cairia em uma réplica sem o prefixo em cache e refaria o trabalho. O roteador do proxy rastreia quais blocos de 16 tokens cada réplica já serviu, usando hashes encadeados (o hash do bloco 3 representa os blocos 1, 2 e 3, espelhando a causalidade da atenção). No fim da execução, os contadores mostraram 84,5% de acertos de afinidade contra um mínimo teórico de 87,5% — muito perto do ideal.

Do gargalo ao resultado: 3,9× mais rápido

O ponto alto do artigo é a caça ao gargalo. Com as métricas do AsyncGRPOTrainer, o time identificou, em cinco execuções sucessivas, onde estava o problema a cada vez e corrigiu:

O resultado final: a mesma receita, os mesmos 500 passos e a mesma curva de recompensa, mas em 53 minutos em vez de 3h27 — uma aceleração de 3,9×, treinando 31% mais amostras. A lição mais importante é metodológica: não otimize estágios do pipeline isoladamente. Um estágio lento esconde o desempenho real de tudo o que vem antes dele. Só olhando para o sistema inteiro o time descobriu, por exemplo, que o terceiro servidor não adiantava nada porque o limite de requisições era do cliente, não do vLLM.

O que isso significa para quem trabalha com IA no Brasil

Para equipes e pesquisadores brasileiros que treinam ou ajustam modelos com recursos limitados, o artigo é um roteiro prático de como fazer RL em escala sem um cluster dedicado com NCCL. O código está disponível no repositório hfjobs-lora-buckets, com um script run_all.sh que reproduz o experimento completo — incluindo um teste de fumaça de 15 minutos com apenas 20 passos. O custo das três máquinas H200 ficou em torno de US$ 20 por hora.

É também uma demonstração de que LoRA não serve apenas para economizar memória em ajuste fino supervisionado: em aprendizado por reforço, ele muda a própria arquitetura de sistemas, permitindo separar fisicamente quem gera e quem aprende — e sincronizar os dois com arquivos de poucos megabytes.