🤖 machine learning

Accelerating RL Post-Training Rollouts via System-Integrated Speculative Decoding

Este artigo apresenta um framework de decodificação especulativa integrado ao sistema no NeMo-RL que atua como um primitivo de aceleração sem perdas para rollouts de pós-treinamento em RL, alcançando uma melhoria de 1,8x no throughput na escala de 8B e projetando até uma aceleração de 2,5x no treinamento de ponta a ponta na escala de 235B quando combinado com execução assíncrona.

Autores originais: Hayate Iso, Tiyasa Mitra, Sudipta Mondal, Rasoul Shafipour, Venmugil Elango, Terry Kong, Yuki Huang, Seonjin Na, Izzy Putterman, Benjamin Chislett, Maor Ashkenazi, Joseph Guman, Gerald Shen, Tugrul Ko

Publicado 2026-04-30

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Hayate Iso, Tiyasa Mitra, Sudipta Mondal, Rasoul Shafipour, Venmugil Elango, Terry Kong, Yuki Huang, Seonjin Na, Izzy Putterman, Benjamin Chislett, Maor Ashkenazi, Joseph Guman, Gerald Shen, Tugrul Konuk, Ashwath Aithal, Ritika Borkar, Ran Zilberstein, Bita Rouhani

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está treinando um aluno brilhante, mas de raciocínio lento (o modelo de IA) para resolver problemas matemáticos complexos. Para ensiná-lo, você precisa fazer uma pergunta, esperar que ele escreva todo o seu processo de pensamento passo a passo e, em seguida, verificar se ele acertou.

O problema é que a parte de "verificar" é rápida, mas a parte de "escrever" é incrivelmente lenta. O aluno escreve uma palavra, pausa para pensar, escreve a próxima palavra, pausa novamente, e assim por diante. Essa escrita lenta, palavra por palavra, é o maior gargalo no treinamento desses modelos de IA.

Este artigo apresenta um truque inteligente chamado Decodificação Especulativa para acelerar esse processo de escrita sem alterar como o aluno pensa ou o que ele aprende.

A Ideia Central: O "Assistente de Rascunho"

Pense no modelo de IA como um chef de cozinha mestre, muito preciso, mas lento. Para acelerar as coisas, você contrata um sous-chef rápido e energético (o "modelo de rascunho").

O Jeito Antigo (Autoregressivo): O chef mestre escreve uma palavra, para, pensa, escreve a próxima, para, pensa. Leva uma eternidade.
O Jeito Novo (Decodificação Especulativa): O sous-chef rapidamente adivinha as próximas 3 ou 4 palavras que o chef poderia escrever. O sous-chef as escreve rapidamente.
A Verificação: O chef mestre então lança um olhar rápido nas anotações do sous-chef.
- Se as anotações estiverem corretas, o chef diz: "Ótimo!" e aceita todas as 4 palavras de uma só vez.
- Se as anotações estiverem erradas, o chef as risca, escreve a palavra correta e recomeça.

A Magia: Como o sous-chef é rápido, o chef consegue aceitar várias palavras de uma só vez na maioria das vezes. O resultado final é exatamente o mesmo que se o chef tivesse escrito sozinho, mas acontece muito mais rápido.

O Que o Artigo Realmente Fez

Os pesquisadores integraram esse sistema em uma estrutura de treinamento real chamada NeMo-RL. Eles não o testaram apenas em tarefas simples; testaram em tarefas de "raciocínio" onde a IA precisa pensar bastante (como resolver problemas matemáticos).

Aqui estão suas principais descobertas, traduzidas para termos do cotidiano:

Funciona Sem Trapacear: Alguns métodos de aceleração tentam cortar cantos (como usar matemática de menor qualidade ou pular etapas), o que pode arruinar o aprendizado do aluno. Este método é "sem perdas". Ele garante que a IA aprenda exatamente da mesma forma que aprenderia sem o assistente, apenas mais rápido.
O Aumento de Velocidade:
- Em um modelo de tamanho médio (8 bilhões de parâmetros), eles viram o processo de escrita acelerar em 1,5 a 1,8 vezes.
- Como a escrita consome cerca de 70% do tempo total de treinamento, todo o processo de treinamento ficou cerca de 1,35 a 1,4 vezes mais rápido.
O "Rascunho" Importa: O sous-chef precisa ser um bom adivinhador.
- Se você treinar o sous-chef nos exatos mesmos tipos de problemas matemáticos que o aluno está aprendendo, o aumento de velocidade é enorme.
- Se você usar um sous-chef genérico que só conhece conversas gerais, o aumento de velocidade é menor.
- Não adivinhe muito à frente: Se o sous-chef tentar adivinhar 7 palavras de uma vez, ele comete muitos erros, e o chef mestre gasta muito tempo corrigindo-os. Adivinhar 3 palavras de cada vez foi o "ponto ideal".
O Futuro (Modelos Grandes): Eles usaram um simulador de computador superpreciso para prever o que acontece com modelos massivos (235 bilhões de parâmetros) e milhares de computadores trabalhando juntos.
- Eles preveem que, para esses modelos gigantes, essa técnica poderia tornar todo o processo de treinamento 2,5 vezes mais rápido.

Por Que Isso Importa

No mundo da IA, tempo é dinheiro. Se você pode treinar um modelo 2,5 vezes mais rápido, você pode obter um modelo mais inteligente no mesmo período de tempo, ou obter o mesmo modelo por uma fração do custo.

O artigo prova que você não precisa alterar o "cérebro" da IA ou as regras do jogo para obter essa velocidade. Você só precisa adicionar um assistente inteligente e rápido para ajudar a escrever os rascunhos e deixar o modelo principal fazer a verificação final. É uma atualização de sistema que faz todo o pipeline de treinamento funcionar de forma mais suave sem quebrar nada.

1. Declaração do Problema

O pós-treinamento de Aprendizado por Reforço (RL) para Modelos de Linguagem de Grande Escala (LLMs) de fronteira está cada vez mais limitado pela geração autoregressiva de rollouts. Em cargas de trabalho orientadas a raciocínio (por exemplo, matemática, código), o tempo gasto na geração de trajetórias (rollouts) domina o tempo total de treinamento, frequentemente excedendo o tempo gasto no cálculo de gradientes.

Métodos de eficiência existentes tentam resolver isso alterando a dinâmica de treinamento, como:

Execução assíncrona: Sobreposição da geração com o aprendizado, mas introduz atraso na política.
Replay off-policy: Reutilização de trajetórias desatualizadas, mas requer correções de amostragem por importância.
Rollouts de baixa precisão: Redução de computação, mas introduz incompatibilidade de distribuição.

Esses métodos trocam efetividade (a qualidade do sinal de aprendizado) por vazão. O artigo identifica uma necessidade crítica de um método que acelere a geração sem alterar a distribuição de amostragem da política-alvo, preservando assim a semântica exata de treinamento necessária para a convergência estável do RL.

2. Metodologia

Os autores propõem integrar a Decodificação Especulativa como um primitivo de aceleração sem perdas diretamente no framework de treinamento NeMo-RL (usando um backend vLLM).

Mecanismo Central

Rascunho (Drafting): Um modelo (ou cabeça) de "rascunho" propõe múltiplos tokens ( $k$ ) em paralelo.
Verificação: O modelo "verificador" alvo (a política de RL) verifica esses tokens em uma única passagem forward.
Amostragem por Rejeição: Os tokens são aceitos ou rejeitados com base na distribuição de probabilidade do verificador. Isso garante que a distribuição final do rollout permaneça idêntica à política-alvo, assegurando nenhuma mudança no objetivo de otimização do RL.

Integração ao Sistema

A integração vai além do serviço de inferência padrão ao abordar restrições específicas do RL:

Sincronização de Pesos: O modelo de rascunho deve permanecer alinhado com os pesos da política em rápida mudança durante o treinamento de RL.
Caminhos Duplos de Rascunho:
- Caminho Geral (EAGLE-3): Usa um modelo de rascunho externo treinado nas saídas da política. Funciona para qualquer modelo pré-treinado.
- Caminho Nativo: Usa cabeças de Previsão de Múltiplos Tokens (MTP) embutidas, se a arquitetura do modelo as suportar.
Desacoplamento de Gradientes: Para evitar que o treinamento do modelo de rascunho interfira no sinal de gradiente da política, estados ocultos e log-probabilidades são armazenados em cache durante a passagem forward da política e roteados através de uma fronteira detach() para a cabeça de rascunho.
Compatibilidade de Pipeline: O sistema suporta pipelines de RL síncronos (a geração aguarda o treinamento) e assíncronos (a geração se sobrepõe ao treinamento).

3. Principais Contribuições

Integração de Sistema: Primeira implementação de decodificação especulativa em uma pilha de treinamento de RL de nível de produção e open-source (NeMo-RL) que preserva semânticas de treinamento exatas do verificador.
Análise Operacional: Um estudo abrangente dos fatores que determinam o aumento de velocidade, incluindo inicialização do rascunho, comprimento do rascunho e adaptação online.
Projeções de Escalabilidade: Um simulador de desempenho de alta fidelidade usado para projetar ganhos em escalas de fronteira (até 235B de parâmetros e 2048 GPUs), indo além das limitações das configurações experimentais atuais de 8B.

4. Resultados Experimentais

Configuração Experimental

Modelos: Qwen3-8B (RL-Think) e Qwen3-8B-Base (RL-Zero).
Tarefas: Raciocínio matemático (treinamento DAPO-Math-17K, validação AIME-2024).
Hardware: 8 nós GB200 NVL72 (32 GPUs no total).
Linha de Base: Decodificação autoregressiva padrão.

Principais Descobertas

Ganhos de Vazão:
- Latência de Geração: A decodificação especulativa reduziu o tempo de geração em 1,54× (RL-Think) e 1,77× (RL-Zero).
- Tempo de Passo Fim-a-Fim: O aumento de velocidade geral do passo de treinamento atingiu 1,35× e 1,41×, respectivamente. Os ganhos foram limitados por estágios não relacionados à geração (cálculo de log-probabilidade e treinamento) que a decodificação especulativa não pode acelerar.
- Precisão: As curvas de precisão de validação foram indistinguíveis entre decodificação autoregressiva e especulativa, confirmando a natureza sem perdas da aceleração.
Insights Operacionais:
- Inicialização: Inicializar o modelo de rascunho em dados de domínio específico (DAPO-Math) superou significativamente a inicialização genérica de domínio de chat (UltraChat), resultando em maiores comprimentos de aceitação e aumentos de velocidade.
- Comprimento do Rascunho ( $k$ ): Contrariando a intuição, rascunhos mais curtos ( $k=3$ ) produziram o melhor aumento de velocidade fim-a-fim. Rascunhos mais longos ( $k=5, 7$ ) aumentaram a sobrecarga de verificação, às vezes tornando o processo mais lento que a decodificação autoregressiva.
- Adaptação Online: Atualizar o modelo de rascunho online forneceu ganhos marginais apenas quando o rascunho inicial estava mal alinhado. Para rascunhos bem inicializados, o treinamento offline foi suficiente.
- RL Assíncrono: Em modos assíncronos, a decodificação especulativa permanece complementar. Embora a sobreposição assíncrona oculte parte do custo de geração, a especulação reduz ainda mais o tempo de geração "exposto" no caminho crítico, resultando em um aumento de velocidade de 1,24× em configurações assíncronas específicas.

Projeções de Simulação (Escala de Implantação)

Usando um simulador para um modelo de 235B em 2048 GPUs:

Aumento de Velocidade de Rollout: Pode exceder 3,5× sob condições ótimas (comprimentos específicos de rascunho/aceitação e atraso moderado da política).
Aumento de Velocidade Fim-a-Fim: Projetado para atingir ~2,5× devido à alta proporção de tempo gasto na geração em modelos de escala de fronteira.
Sensibilidade à Escala: Modelos maiores beneficiam-se mais da especulação em implantações de grande escala porque sofrem mais com latência de cauda longa e subutilização, que a especulação mitiga.

5. Significado

Este artigo demonstra que a decodificação especulativa é uma estratégia viável e sem perdas para acelerar o pós-treinamento de RL. Seu significado principal reside em:

Preservação de Semântica: Ao contrário de outras técnicas de aceleração, não compromete o sinal de treinamento do RL nem requer correções off-policy.
Otimização em Nível de Sistema: Preenche a lacuna entre técnicas de aceleração de inferência e os requisitos complexos de loops de treinamento de RL (sincronização de pesos, manipulação de gradientes).
Escalabilidade: O aumento de velocidade fim-a-fim projetado de 2,5× em escalas de fronteira sugere que a decodificação especulativa pode se tornar um componente padrão na infraestrutura de RL futura, reduzindo significativamente o tempo e o custo necessários para treinar modelos capazes de raciocínio.

O trabalho estabelece um caminho claro de implantação: usar rascunhos inicializados em domínio específico com comprimentos de rascunho moderados (por exemplo, $k=3$ ) para maximizar a vazão sem sacrificar a estabilidade do treinamento.