💬 NLP

Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow

Este artigo apresenta o Jet-RL, um framework de aprendizado por reforço unificado em FP8 que resolve a instabilidade de treinamento e o colapso de precisão causados por estratégias de precisão mista ao aplicar precisão FP8 consistente tanto para o rollout quanto para o treinamento, alcançando assim acelerações significativas enquanto mantém uma convergência estável.

Autores originais: Haocheng Xi, Charlie Ruan, Peiyuan Liao, Yujun Lin, Han Cai, Yilong Zhao, Shuo Yang, Kurt Keutzer, Song Han, Ligeng Zhu

Publicado 2026-01-27

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Haocheng Xi, Charlie Ruan, Peiyuan Liao, Yujun Lin, Han Cai, Yilong Zhao, Shuo Yang, Kurt Keutzer, Song Han, Ligeng Zhu

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

O Grande Problema: A "Caminhada Lenta" no Treinamento de IA

Imagine que você está treinando um robô muito inteligente (um Large Language Model) para resolver problemas matemáticos complexos. Para se tornar realmente bom nisso, o robô precisa praticar "pensar em voz alta" (gerar longas cadeias de raciocínio) antes de dar uma resposta.

No mundo do treinamento de IA, essa sessão de prática é chamada de Rollout.

O Gargalo: O artigo aponta que essa fase de "pensar em voz alta" é incrivelmente lenta. Ela ocupa 70% do tempo total gasto treinando o robô. É como um estudante que passa 7 horas estudando para uma prova, mas 5 dessas horas são apenas encarando a página em branco, tentando descobrir a primeira frase.
O Objetivo: Queremos acelerar essa fase de "pensamento" sem tornar o robô mais burro.

O Atalho que Falhou: "O Erro da Contabilidade Dupla"

Para acelerar o processo, engenheiros tentaram um truque simples: Quantização.

A Analogia: Imagine que o cérebro do robô geralmente trabalha com números de "Alta Definição" (HD) de 16 bits (BF16). Para ir mais rápido, eles tentaram mudar o modo de "pensamento" do robô para um modo de "Baixa Definição" (SD) de 8 bits (FP8) apenas para a sessão de prática, mantendo a sessão de "avaliação" em HD.
A Ideia: "Vamos usar o modo rápido e de baixa resolução para as partes longas de pensamento, mas manter o modo de alta resolução para as atualizações reais de aprendizado."
O Desastre: O artigo descobriu que essa abordagem é um desastre para tarefas longas ou difíceis.
- O Descompasso: É como pedir a um estudante para escrever uma redação de 10 páginas a lápis (baixa resolução), mas avaliá-la como se ele tivesse escrito à caneta (alta resolução). Os pequenos erros do esboço a lápis se acumulam. Quando a redação fica longa (16.000 palavras), a versão a "lápis" não se parece em nada com a versão à "caneta".
- O Resultado: O robô fica confuso. Ele começa a alucinar, o treinamento trava e o desempenho do robô entra em colapso. O artigo chama isso de "colapso catastrófico de precisão".

A Solução: Jet-RL (A Abordagem da "Linguagem Unificada")

Os autores propõem um novo sistema chamado Jet-RL. Em vez de trocar de idioma entre a prática e a avaliação, eles fazem o robô falar a mesma língua (FP8) para tudo.

A Analogia: Imagine que o robô agora escreve suas redações de prática a lápis, e o professor também avalia as redações usando uma rubrica baseada em lápis.
- Consistência: Como o "pensamento" e o "aprendizado" acontecem exatamente no mesmo formato de baixa resolução, não há confusão. O robô aprende exatamente o que praticou.
- Estabilidade: Mesmo para redações muito longas (rollouts longos) ou tópicos muito difíceis, o robô permanece estável porque os erros do "lápis" são consistentes durante todo o processo.

Como Funciona (A Magia Técnica)

Para fazer isso funcionar, a equipe teve que ser muito inteligente na forma como lidou com a matemática:

Fluxo Unificado: Eles garantiram que os dados que fluem através dos chips do computador sejam tratados como FP8 (baixa resolução) desde o primeiro passo do pensamento até o último passo da atualização do cérebro.
Agrupamento Inteligente: Eles não apenas encolheram tudo cegamente. Eles agruparam números juntos (como organizar livros em uma prateleira) para que a versão de "baixa resolução" permaneça precisa o suficiente para aprender.
Sem Calibração: Normalmente, quando você troca de formato, precisa gastar tempo "calibrando" (testando e ajustando) para garantir que funcione. O Jet-RL pula essa etapa inteira porque o sistema é projetado para ser consistente desde o início.

Os Resultados: Rápido e Preciso

O artigo testou isso em vários modelos e descobriu:

Velocidade: Tornou a fase de "pensamento" 33% mais rápida e todo o processo de treinamento 16% mais rápido.
Precisão: Ao contrário do atalho que falhou (que causou o fracasso do robô), o Jet-RL manteve a inteligência do robô quase exatamente a mesma da versão lenta de alta definição. A queda no desempenho foi insignificante (menos de 1%).
Confiabilidade: Funcionou mesmo quando o robô tinha que escrever respostas muito longas (16.000 palavras) ou resolver problemas matemáticos muito difíceis, onde o método antigo teria travado.

Resumo

Jet-RL é uma nova maneira de treinar IA que impede o robô de ficar confuso ao alternar entre o "modo rápido" e o "modo inteligente". Ao falar a mesma "linguagem rápida" tanto para o pensamento quanto para o aprendizado, ele torna o treinamento de IA muito mais rápido sem tornar a IA mais burra. Ele transforma um atalho instável e quebrado em uma rodovia confiável e de alta velocidade.

Resumo Técnico: Jet-RL

Declaração do Problema

O Aprendizado por Reforço (RL) tornou-se fundamental para aprimorar as capacidades de raciocínio complexo de Grandes Modelos de Linguagem (LLMs), particularmente através da geração de Cadeia de Pensamento (Chain-of-Thought - CoT). No entanto, os pipelines de treinamento de RL são computacionalmente ineficientes, com a fase de rollout (geração de respostas) consumindo mais de 70% do tempo total de treinamento, especialmente quando o comprimento da geração excede 8.000 tokens.

Para abordar esse gargalo, uma estratégia comum envolve o uso de quantização FP8 durante a fase de rollout, mantendo a precisão BF16 para a fase de treinamento (BF16-train + FP8-rollout). Embora essa abordagem seja amplamente adotada em frameworks como VeRL e OpenRLHF, este artigo identifica dois modos de falha críticos:

Instabilidade de Longo Rollout: À medida que o comprimento da geração aumenta (ex: >8k tokens), o descompasso numérico entre a distribuição de treinamento BF16 e a distribuição de rollout FP8 se acumula. Isso leva a uma instabilidade severa no treinamento e a um colapso catastrófico de precisão.
Falha em Tarefas Desafiadoras: A estratégia falha quando aplicada a tarefas de raciocínio difíceis ou modelos base mais fracos. Nesses cenários, a menor confiança do modelo amplifica os erros de quantização, fazendo com que a curva de treinamento divirja rapidamente em comparação com as linhas de base de precisão total.

A causa raiz é identificada como uma natureza off-policy: o ator aprende com dados gerados por uma política (rollout FP8) que difere numericamente da política sendo atualizada (treinamento BF16), violando a suposição on-policy essencial para um RL estável.

Metodologia: Jet-RL

Os autores propõem o Jet-RL, um framework projetado para permitir o treinamento de RL FP8 robusto e on-policy, ao impor um fluxo de precisão unificado entre o treinamento e o rollout.

Arquitetura Central

O Jet-RL elimina o descompasso de precisão garantindo que tanto o passo de propagação direta (forward pass) de treinamento quanto o rollout de inferência utilizem o mesmo fluxo de quantização FP8.

Grafo de Precisão Unificado: O framework modela o modelo como um grafo direcionado onde as arestas representam a propagação de tensores. O Jet-RL garante que o grafo de inferência ( $\mathcal{G}_{infer}$ ) seja um subgrafo do grafo de treinamento forward ( $\mathcal{G}^{fwd}_{train}$ ), compartilhando a mesma granularidade de quantização e precisão.
Pesos Mestres (Master Weights): Embora o forward pass e o rollout operem em FP8, os pesos mestres são mantidos em BF16 para estabilizar as atualizações de treinamento, de forma semelhante ao treinamento de precisão mista padrão.

Esquema de Quantização

Para manter a estabilidade enquanto acelera a computação, o Jet-RL emprega uma estratégia de quantização mista para Camadas Lineares (GEMMs):

Ativações: Quantizadas usando quantização 1×128 per-group.
Pesos: Quantizados usando quantização 128×128 per-block.
Passo de Retropropagação (Backward Pass):
- As ativações salvas para o backward pass são armazenadas em FP8 (fundidas com o operador anterior).
- Os gradientes transportados entre operadores são mantidos em BF16 para evitar o subfluxo de gradiente (underflow) e a degradação da convergência.
- As operações de backward GEMM (DGrad e WGrad) são realizadas em FP8 para aceleração, utilizando configurações de kernel específicas (ex: 1×128 e 128×1 para WGrad) para alinhar-se aos requisitos de hardware.

Implementação

O Jet-RL é implementado usando vLLM para o motor de inferência e VeRL para o framework de treinamento de RL. Ele aproveita kernels de GEMM FP8 de última geração do DeepGEMM e kernels Triton customizados para quantização, transposição e operações fundidas de ativação/RMSNorm.

Principais Contribuições

Identificação do Problema: O artigo demonstra que a estratégia padrão de BF16-train + FP8-rollout é fundamentalmente falha para longos rollouts e tarefas desafiadoras devido a descompassos numéricos off-policy, levando ao colapso do treinamento.
Framework Jet-RL: A proposta de um fluxo de precisão FP8 unificado que impõe consistência on-policy entre treinamento e rollout, eliminando a necessidade de calibração entre etapas e resolvendo o descompasso de política.
Desempenho e Estabilidade: A demonstração de que o Jet-RL alcança acelerações significativas mantendo a convergência e a precisão próximas das linhas de base BF16 em diversos modelos e tarefas.

Resultados Experimentais

Os autores avaliaram o Jet-RL em modelos incluindo Llama3.1-8B, Qwen2.5-7B e Qwen3-8B-Base em datasets como GSM8K, MATH e DeepMATH.

Precisão e Convergência

Estabilidade: Ao contrário da linha de base BF16-train + FP8-rollout, que falhou em convergir no Qwen2.5-7B e mostrou degradação significativa em rollouts de 16k, o Jet-RL convergiu com sucesso em todas as configurações.
Gap de Precisão: O Jet-RL reduziu o gap de desempenho em relação ao treinamento BF16 para aproximadamente 1% (ex: 1,1% de degradação no Qwen3-8B-Base vs. 2,9% para a linha de base). Em alguns casos (Llama3.1-8B), o Jet-RL até superou a linha de base BF16 em 2,0%.
Robustez: O Jet-RL manteve o desempenho em tarefas difíceis (DeepMATH) e contextos longos (16k) onde a estratégia de rollout FP8 ingênua colapsou.

Eficiência

Aceleração de Rollout: A quantização FP8 proporcionou um aumento de velocidade de 1,07× a 1,33× na fase de rollout (tokens/s), com modelos maiores (32B) apresentando os maiores ganhos.
Aceleração de Treinamento: O fluxo FP8 unificado acelerou a fase de atualização do ator em 1,54× e a inferência do modelo de referência em 1,80×, resultando em um aumento de velocidade de 1,41× para a fase de treinamento.
End-to-End: O framework alcançou um aumento de velocidade end-to-end de 1,16× para o treinamento de modelos 8B, com expectativas de ganhos maiores para modelos maiores.

Significância e Alegações

O artigo afirma que o Jet-RL fornece um caminho confiável e eficiente para aplicar a computação FP8 para acelerar o treinamento de RL em larga escala. Ao abordar a instabilidade fundamental causada por descompassos de precisão, o Jet-RL permite o uso de FP8 não apenas para inferência, mas para todo o loop de treinamento de RL. Isso permite uma aceleração significativa (até 33% de aceleração de rollout e 41% de aceleração de treinamento) sem sacrificar a estabilidade de convergência ou a precisão necessária para treinar modelos de raciocínio complexos. O trabalho estabelece que um fluxo de precisão unificado é essencial para escalar o treinamento de RL com hardware de baixa precisão.