Aprimorando gpt-oss para Precisão e Desempenho com Treinamento Consciente à Quantização

TL;DR

• A NVIDIA apresenta um fluxo prático de ajuste fino para o gpt-oss que usa fine-tuning supervisionado de alta precisão (SFT) seguido de treinamento consciente à quantização (QAT) para recuperar precisão em FP4, mantendo a eficiência de implantação.
• O fluxo troca para BF16 para o SFT e, em seguida, aplica QAT para retornar à precisão MXFP4, permitindo alinhamento e benefícios de implantação em baixa precisão.
• Na avaliação, a receita aumentou duas tarefas de 16% e 30% de base para 98% de taxa de aprovação após SFT + QAT.
• O NVFP4, um formato FP4 mais novo para treinamento e inferência no NVIDIA Blackwell, mostra perda de validação 2–3% melhor em média e promete convergência mais acentuada para tarefas com raciocínio mais profundo.
• A receita MXFP4 pode ser adaptada para NVFP4 com uma única linha de código, e o suporte NVFP4 em TensorRT-LLM deve ampliar a adoção em mais frameworks.
• O fluxo completo está implementado no repositório NVIDIA Model Optimizer, com um script de conveniência para exportar para checkpoints PyTorch padrão e caminhos de implantação via TensorRT-LLM.

Contexto e antecedentes

Lançamentos de modelos de base de código aberto têm impulsionado a comunidade de IA com inovações arquitetônicas e novas capacidades. A família gpt-oss representa a primeira suíte de modelos de código aberto lançada pela OpenAI desde o GPT-2, oferecendo um modelo avançado com arquitetura MoE, janela de contexto de 128K e habilidades de raciocínio profundo ajustáveis. O maior variante, gpt-oss-120B, alcança desempenho em benchmarks abertos comparable aos modelos o3 e o4 proprietários da OpenAI. Apesar do desempenho sólido, muitos modelos de base exigem técnicas de pós-treinamento para implantação em produção, especialmente em setores de baixa tolerância a falhas. O FP4 nativo do gpt-oss apresentou desafios únicos para ajuste fino. A NVIDIA observa que a precisão estável para o ajuste fino em FP4 ainda não está plenamente estabelecida, o que motiva a abordagem de duas etapas: upcast para maior precisão para estabilizar gradientes, seguido de SFT de alta precisão e uma aplicação subsequente de QAT para retornar ao FP4 mantendo o desempenho específico da tarefa. O objetivo é oferecer alinhamento e eficiência de implantação em baixa precisão. As ferramentas-chave incluem o repositório Model Optimizer da NVIDIA, com suporte de Hugging Face gpt-oss-recipes, datasets OpenAI Cookbook e as ferramentas de Transformer Engine. O objetivo é recuperar precisão em FP4 ao mesmo tempo em que se mantém a eficiência de inferência de baixa precisão para aplicações em produção.

O que há de novo

• A recomendação central é realizar o fine-tuning em alta precisão (BF16) para estabilizar gradientes, depois aplicar QAT para trazer o modelo de volta à precisão MXFP4 para implantação. Pular a etapa de alta precisão tende a reduzir a precisão final.
• Em duas tarefas de avaliação, foram observados ganhos drásticos: raciocínio em non-English com um dataset multilíngue do OpenAI Cookbook e redução de recusas desnecessárias de prompts seguros com o FalseReject da Amazon. As pontuações base iniciais eram 16% e 30%; após o fluxo, atingiram 98% de aprovação.
• A NVIDIA compara MXFP4 com NVFP4, indicando que NVFP4 tende a convergir com mais confiabilidade e oferece perda de validação 2–3% menor em várias tarefas. NVFP4 é projetado para formatos FP4 em treinamento e inferência, utilizando o segundo gera Transformer Engine para maior desempenho.
• A migração de MXFP4 para NVFP4 pode ser feita com uma única linha de código, demonstrando um caminho simples de migração conforme o suporte NVFP4 se expande no TensorRT-LLM e em outros frameworks.
• Com o NVIDIA Blackwell, NVFP4 permite até 15 PFLOPs de computação FP4 para Ultra compute, proporcionando convergência mais firme e margens maiores para requisitos mais rigorosos e raciocínio mais profundo. A escala FP8 E4M3 ajuda a reduzir erros de quantização durante o forward, auxiliando a adaptação dos pesos para a precisão alvo.
• Após completar a receita, um script de conveniência no repositório Model Optimizer exporta o checkpoint BF16 para MXFP4, e os checkpoints MXFP4 resultantes foram testados com SGLang, TensorRT-LLM e vLLM. A implantação pode ser realizada com TensorRT-LLM 1.1.0rc1.
• O desafio central continua: recuperar a precisão em FP4 mantendo as vantagens de eficiência da baixa precisão. O caminho proposto—upcast para BF16 para SFT, seguido de QAT—resolve essa lacuna ao adaptar os pesos para a precisão baixa, reforçando o comportamento específico da tarefa.
• Olhando para o futuro, o suporte NVFP4 no TensorRT-LLM e em outros frameworks de inferência de código aberto ampliará a adoção, permitindo NVFP4 com a mesma estratégia SFT + QAT para maior precisão em implantações gpt-oss.

Por que isso importa (impacto para desenvolvedores/empresas)

Para desenvolvedores e empresas, a capacidade de implantar modelos abertos poderosos como o gpt-oss em FP4, mantendo ou aumentando a precisão, apresenta um ROI atraente. A combinação de SFT e QAT ajuda a recuperar o desempenho específico da tarefa sem sacrificar os ganhos de eficiência da inferência de baixa precisão. Em domínios sensíveis à segurança, maior alinhamento e menos recusas resultam em IA mais utilizável e confiável. À medida que o hardware avança, a introdução do NVFP4 pode trazer ganhos adicionais de precisão quando emparelhado com QAT. A arquitetura Blackwell da NVIDIA e ferramentas associadas, como o Transformer Engine de segunda geração e o TensorRT-LLM, estão posicionados para entregar convergência mais firme e margens maiores para requisitos mais rigorosos e raciocínio mais profundo em implantações de produção. A capacidade de adaptar checkpoints MXFP4 para NVFP4 com alterações mínimas de código diminui a barreira à adoção e acelera prazos de implantação.

Detalhes técnicos ou Implementação

• Fluxo central: upcast para BF16 para SFT, seguido de QAT para MXFP4 para implantação. Essa sequência estabiliza gradientes em alta precisão e, em seguida, permite adaptar pesos para a precisão de baixa precisão alvo.
• Os hiperparâmetros e a duração do treinamento para QAT são otimizáveis; pular a etapa de alta precisão tende a reduzir a precisão final, portanto recomenda-se um fine-tuning em alta precisão antes do QAT.
• As duas tarefas de avaliação mostram o impacto prático do fluxo: | Tarefa | Base | Taxa de aprovação após fluxo |---|---|---| | Raciocínio não em inglês (multilíngue OpenAI Cookbook) | 16% | 98% |Recusa de prompts seguros (Amazon FalseReject) | 30% | 98% |
• Para migrar de MXFP4 para NVFP4, uma única linha de código é suficiente para adaptar a rota, após o que a validação mostra melhoria de 2–3% na perda de validação entre tarefas.
• NVFP4 introduz um formato FP4 pensado para treinamento e inferência, permitindo até 15 PFLOPs de compute FP4 com o NVIDIA Blackwell Ultra para maior eficiência e precisão. A escala FP8 E4M3 reduz erros de quantização durante o forward, ajudando a adaptar pesos para a precisão alvo.
• O fluxo padrão inclui exportar o checkpoint BF16 treinado para MXFP4 via o script de conveniência do Model Optimizer, seguido pela implantação com stacks validados como SGLang, TensorRT-LLM e vLLM.
• O fluxo descrito está alinhado com esforços contínuos para integrar o suporte NVFP4 ao gpt-oss em TensorRT-LLM e em outros frameworks de inferência de código aberto, sinalizando maior acessibilidade quando o NVFP4 estiver disponível.

Principais aprendizados

• Um caminho de ajuste fino em duas etapas (SFT em alta precisão seguido de QAT para FP4) recupera a precisão do gpt-oss em implantação.
• Os resultados em tarefas alvo são expressivos: 98% de aprovação após o recipe para as duas tarefas avaliadas.
• NVFP4 oferece potenciais ganhos de precisão e de convergência sobre MXFP4, com melhor perda de validação em diversas tarefas.
• A migração MXFP4 -> NVFP4 é simplificada, exigindo apenas uma linha de código.
• O repositório NVIDIA Model Optimizer fornece ferramentas completas para exportar, validar e implantar checkpoints em ambientes de produção.

References

• NVIDIA Dev Blog: Fine-Tuning gpt-oss for Accuracy and Performance with Quantization Aware Training