Alibaba Qwen apresenta o TRACE: Rollout em FP4 fica até 5,4 vezes mais rápido

A equipe Alibaba Qwen apresentou o TRACE, um método de treinamento de grandes modelos com aprendizado por reforço que move a fase de Rollout para FP4: tanto os pesos e as ativações (W/A) quanto o KV-cache. Segundo os autores, o Rollout fica até 5,4 vezes mais rápido que em BF16, e o treinamento continua estável. A estabilidade pesa mais que a velocidade aqui, porque muita gente sabe rodar a geração em 4 bits, mas nem todo mundo consegue impedir que o RL desande nesse cenário.

Por que o Rollout virou a parte mais cara do RL

Em RL para modelos de linguagem, o ciclo funciona assim: o modelo gera muitas respostas, em geral cadeias longas de raciocínio, elas são avaliadas, e os pesos são atualizados com base nessas trajetórias. A geração acontece token a token, esbarra na largura de banda da memória, e o KV-cache cresce junto com o tamanho da resposta. Quanto mais longo o reasoning e maior o número de amostras por prompt, mais tempo e memória de vídeo o Rollout consome, muito mais do que o passo de descida do gradiente.

Daí o desejo natural de comprimir essa etapa. Em teoria, migrar de BF16 para FP4 reduz tanto o volume de computação quanto o consumo de memória. Menos bytes por peso e por elemento do cache significam mais requisições paralelas por placa e menos tempo de espera por dados.

Onde o FP4 prejudica o treinamento

O preço dessa compressão é a grosseria da grade numérica. Em 4 bits cabem apenas 16 valores, então duas grandezas muito próximas facilmente ficam em lados opostos do limite de quantização e caem em pontos de código diferentes. Se o treinador e o motor de Rollout calculam com precisões distintas, entradas quase idênticas produzem resultados bem diferentes. Uma diferença minúscula no nível do número vira uma divergência entre o que o modelo «gerou» e o que ele «pensa» na hora de atualizar os pesos.

Isso é perigoso para o RL. O algoritmo pressupõe que as trajetórias foram geradas por uma política próxima daquela que está sendo treinada. Quando gerador e treinador divergem, o gradiente passa a se apoiar em dados que o modelo atual produziria de outra forma. Normalmente isso aparece como ruído, deriva e instabilidade no treinamento. Com FP4 o problema é mais agudo que com FP8, porque há mais efeitos de borda e menos margem de precisão. É exatamente esse problema que o TRACE tenta resolver.

O que se sabe sobre o TRACE

Parâmetros principais do trabalho:

• Desenvolvedor: equipe Qwen (Alibaba).
• Área de aplicação: etapa de Rollout no treinamento de grandes modelos com RL.
• Formato: FP4 para pesos e ativações (W/A), mais FP4 para o KV-cache.
• Linha de base: BF16.
• Efeito declarado: aceleração do Rollout de até 5,4 vezes.
• Objetivo: manter a estabilidade do treinamento com geração em baixa precisão.

A formulação «até 5,4 vezes» indica o máximo, e não o resultado médio. O ganho real depende do tamanho da geração, do tamanho do batch, da arquitetura do modelo e do hardware. Quanto maior o contexto e mais pesado o KV-cache, mais perceptível é o efeito da quantização do cache, já que ele costuma ser o gargalo em raciocínios longos.

O que as equipes que rodam RL nos próprios clusters devem verificar

Para quem treina modelos de reasoning por conta própria, a principal conclusão é que o Rollout merece ser tratado como um problema de otimização à parte, e não como um processo secundário. Se a geração ocupa a maior parte do tempo do passo, reduzir a precisão rende mais do que ajustar o otimizador.

Antes de migrar para FP4, é preciso medir a divergência entre as probabilidades dos tokens no treinador e no motor de Rollout. Se ela já é visível em BF16 ou FP8, vai aumentar em 4 bits. Vale também checar separadamente a qualidade em trajetórias longas: o erro do KV-cache se acumula conforme a resposta cresce.

Por fim, o ganho depende de suporte de hardware. O FP4 nativo existe em gerações novas de aceleradores, como a arquitetura Blackwell, e em placas mais antigas o caminho de 4 bits não entrega essa aceleração. Para laboratórios com um parque limitado de GPUs, trabalhos como o TRACE interessam porque permitem fazer mais iterações de RL com o mesmo hardware. A próxima pergunta é até que ponto o resultado se repete em modelos de outros tamanhos e em outras tarefas, da matemática ao código.