Fine-tuning do GPT-OSS-20B para raciocínio multilíngue no Amazon SageMaker HyperPod

O problema: um modelo que pensa só em inglês

O GPT-OSS-20B é um modelo de raciocínio de código aberto com um comportamento peculiar: independentemente do idioma em que você faz uma pergunta, ele raciocina em inglês. Peça para ele pensar em francês, alemão, espanhol ou italiano — ele reconhece a instrução, mas produz a cadeia de raciocínio em inglês mesmo assim.

Em um teste de avaliação com 50 exemplos multilíngues, o modelo base raciocinou no idioma-alvo correto em apenas 22% dos casos. Esse número, na prática, é quase inteiramente composto pelos exemplos em inglês: inglês atingiu 100%, enquanto francês e espanhol ficaram em 0%, alemão em 10% e italiano em 0%.

Para equipes que implantam modelos de raciocínio para usuários não anglófonos, essa lacuna é um problema real. Um usuário que pergunta em francês espera que o modelo raciocine e responda em francês — não que mude para o inglês no meio do pensamento. Fechar essa lacuna é um desafio de fine-tuning.

A questão que toda equipe enfrenta ao customizar modelos de peso aberto é: usar Ajuste Fino Supervisionado (SFT), Aprendizado por Reforço com Otimização de Política Relativa de Grupo (GRPO), ou os dois? A AWS publicou um artigo técnico detalhando exatamente essa decisão, com experimentos reais no Amazon SageMaker HyperPod.

Visão geral da solução

A tarefa e o conjunto de dados

O dataset utilizado foi o HuggingFaceH4/Multilingual-Thinking — um conjunto de dados de raciocínio com cadeia de pensamento traduzido para francês, espanhol, alemão e italiano, além do inglês, com aproximadamente 949 exemplos de treino e 50 de validação. O modelo recebe uma instrução de sistema indicando o idioma de raciocínio e deve tanto raciocinar quanto responder nesse idioma-alvo, sobrepondo-se ao forte viés pelo inglês.

A avaliação consiste em 50 casos de teste: 10 problemas de matemática e lógica em cada um dos 5 idiomas (EN/FR/DE/ES/IT). Duas métricas são acompanhadas:

Raciocínio no idioma correto — a cadeia de pensamento está no idioma solicitado?
Resposta no idioma correto — a resposta final está no idioma solicitado?

Como SFT e GRPO se complementam

A forma mais clara de entender as duas técnicas é esta:

SFT ensina um comportamento que o modelo ainda não tem.
GRPO otimiza um comportamento que o modelo já consegue produzir às vezes.

O GRPO funciona amostrando um grupo de respostas por prompt e empurrando a política em direção às respostas acima da média. Se o comportamento desejado tem probabilidade zero na distribuição base, há muito pouco para o GRPO reforçar. Essa ideia central explica praticamente todos os resultados do experimento.

A infraestrutura de treinamento

O treinamento foi executado no Amazon SageMaker HyperPod com orquestração pelo Serviço de Kubernetes Elástico da Amazon (Amazon EKS). Para os experimentos iniciais de SFT e GRPO no GPT-OSS-20B, foi utilizada a biblioteca PyTorch TRL. Para o experimento de controle com Qwen, a equipe migrou para o veRL como trainer, com Ray e KubeRay para orquestração distribuída e vLLM para geração rápida de rollouts.

Checkpoints e datasets ficaram armazenados no Amazon FSx for Lustre, e o fabric de comunicação entre nós utilizou o Adaptador de Malha Elástica (EFA). Todo o fine-tuning usou Adaptação de Baixo Posto (LoRA), que é eficiente em parâmetros e adequada tanto para SFT quanto para GRPO.

O diagrama abaixo ilustra o caminho completo dos dados de treinamento de ponta a ponta:

O fluxo começa com uma submissão via AWS CLI ou Amazon SageMaker Studio. O plano de controle do Amazon EKS no cluster HyperPod aloca pods em nós GPU e inicializa um Ray head com pods worker para o trainer veRL e os workers de rollout vLLM. O Amazon FSx for Lustre fornece armazenamento compartilhado de alta vazão para o modelo base, os dados de treinamento e os checkpoints. Ao final do treinamento, o adapter LoRA é gravado de volta no FSx para avaliação ou serving.

Configurando o ambiente de treinamento

Todos os comandos, manifestos e scripts estão no repositório GitHub awsome-distributed-ai, no diretório 3.test_cases/pytorch/trl/gpt-oss-lora-grpo/. Para começar:

git clone https://github.com/aws-samples/awsome-distributed-ai
cd awsome-distributed-ai/3.test_cases/pytorch/trl/gpt-oss-lora-grpo

Copie o template de variáveis de ambiente, edite o arquivo env_vars com seu registro ECR, nome da imagem, tag e token do HuggingFace, e construa a imagem de treinamento:

cp env_vars.example env_vars
source env_vars
cd artifacts && ./build_push.sh && cd ..

Em seguida, aponte o kubectl para o cluster e aplique o pod de gerenciamento de armazenamento FSx:

aws eks update-kubeconfig --name <YOUR_CLUSTER_NAME> --region us-east-2
envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/fsx-storage-manager.yaml | kubectl apply -f -

Fine-tuning com SFT e LoRA

A lacuna do modelo base

O GPT-OSS-20B base na avaliação de 50 casos marcou 22% de raciocínio no idioma correto (11/50). Por idioma: inglês 100%, francês 0%, alemão 10%, espanhol 0%, italiano 0%. O modelo ignora a instrução de idioma e raciocina em inglês de qualquer forma.

Configuração do SFT

O SFT utilizou LoRA r=8, α=16, treinado por 1.000 passos em 4× ml.g6e.12xlarge (16× L40S 48GB) com Paralelismo de Dados Totalmente Fragmentado (FSDP) full-shard, precisão bf16 e gradient checkpointing. O job roda como um HyperPodPyTorchJob elástico gerenciado pelo Kueue, iniciado com hyperpodrun.

envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < hyperpod-eks/train-lora-hyperpod-elastic-g6e.yaml | kubectl apply -f -

Uma execução saudável apresenta logs como este:

{'loss': 1.035, 'grad_norm': 0.163, 'learning_rate': 0.0002, 'mean_token_accuracy': 0.716, 'epoch': 3.25}
{'loss': 1.003, 'grad_norm': 0.171, 'lr': 0.0002, 'acc': 0.722, 'epoch': 3.32}
{'loss': 0.976, 'grad_norm': 0.159, 'lr': 0.0002, 'acc': 0.726, 'epoch': 3.57}

Resultados do SFT e a curva de aprendizado em estágios

Após o SFT, o modelo raciocina no idioma correto em 100% dos casos (50/50) nos cinco idiomas. O aprendizado ocorreu em dois estágios:

Por volta do passo 200, o modelo aprendeu a raciocinar no idioma-alvo (métrica de raciocínio atinge 100%).
Por volta do passo 1.000, o modelo aprendeu também a dar a resposta final no idioma-alvo.

Avaliação final do SFT (passo 1.000): raciocínio no idioma correto 100% (50/50); resposta no idioma correto 62% (31/50); ambos corretos 62% (31/50). Por idioma na resposta: FR 60%, DE 60%, ES 30%, IT 60%, EN 100%.

Conclusão: o SFT não é opcional aqui. O modelo base não consegue seguir a instrução de idioma de raciocínio, e o SFT eleva a métrica de 22% para 100%. A lacuna restante está na resposta final (62%) — a parte mais difícil, pois a resposta é curta e o viés pelo inglês é mais forte ali.

Refinando com GRPO

A acurácia da resposta no idioma correto era o alvo óbvio para o GRPO. O SFT já havia produzido raciocínio correto 100% das vezes, e o modelo às vezes produzia a resposta no idioma certo — havia sinal para reforçar.

Executando GRPO sobre o checkpoint SFT

O GRPO foi executado com o GRPOTrainer do TRL, partindo do checkpoint SFT. A configuração: carregar o LoRA do SFT, chamar merge_and_unload e adicionar um novo LoRA (r=8, α=16, cerca de 0,068% de parâmetros treináveis, 13,6M de 20B). O treinamento usou K=8 gerações por prompt, device_map=auto, taxa de aprendizado 1e-6 e 10 épocas.

A função de recompensa pontuou: idioma da resposta (±5,0), idioma do raciocínio (±1,5) e brevidade (+0,5 ou −1,0), com intervalo combinado de −7,5 a +7,0.

envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/train-grpo-singlenode.yaml | kubectl apply -f -
kubectl logs -f grpo-single

Resultados do GRPO

Partindo do checkpoint SFT: raciocínio no idioma correto 100% (mantido — o GRPO não quebrou o que o SFT ensinou); resposta no idioma correto: de 62% para 70% (+8 pontos, 35/50); ambos corretos: 70%. Por idioma na resposta: FR 70%, DE 80%, ES 60%, IT 40%, EN 100%.

O GRPO entregou um ganho real de +8 pontos na métrica de resposta no idioma correto, mantendo o raciocínio em 100%. O ganho veio onde havia margem para crescer.

Comparando base, SFT e GRPO na prática

Para visualizar a diferença qualitativa, é possível implantar três pods de inferência que compartilham o mesmo modelo base, mas carregam adapters diferentes:

envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/inference-g6e-base.yaml | kubectl apply -f -
envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/inference-g6e-trained.yaml | kubectl apply -f -
envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/inference-g6e-grpo.yaml | kubectl apply -f -

Rodando o mesmo prompt de raciocínio em francês em cada variante:

O pod base raciocina em inglês, ignorando a instrução.
O pod SFT raciocina em francês, mas frequentemente responde em inglês.
O pod GRPO raciocina e responde em francês.

Experimento de controle com Qwen2.5-7B

Para testar a tese “adequar a técnica à lacuna de capacidade” em um modelo diferente, a AWS rodou experimentos controlados no Qwen2.5-7B-Instruct em duas tarefas com lacunas distintas.

Raciocínio multilíngue (50 casos)

Resultados comparando base (A), SFT apenas (B) e SFT+GRPO (C):

Formato estruturado: A=0%, B=92%, C=98% (Δ GRPO: +6)
Idioma do raciocínio: A=0%, B=96%, C=98% (Δ GRPO: +2)
Idioma da resposta: A=0%, B=38%, C=44% (Δ GRPO: +6)
Ambos corretos: A=0%, B=38%, C=44% (Δ GRPO: +6)

GSM8K — uma tarefa sem margem de melhoria

No GSM8K, o Qwen2.5-7B base já marcava 90,5% (181/200). O GRPO no passo 30 também marcou 90,5%. Quando a maioria dos grupos amostrados está toda correta, a vantagem relativa do grupo colapsa para zero — não há sinal de aprendizado. Esse é o outro lado da tese: o GRPO precisa de margem (pass@k muito maior que pass@1) para ter algo a otimizar. Sem margem, sem ganho.

Boas práticas

Comece com SFT. O GRPO não consegue ensinar um comportamento com probabilidade zero na distribuição base. O SFT elevou o raciocínio no idioma correto de 22% para 100%, e cada ganho do GRPO neste trabalho foi construído sobre esse checkpoint SFT.
Use LoRA para SFT e GRPO. LoRA r=8/α=16 foi eficiente em parâmetros (0,068% treináveis) e permaneceu numericamente estável durante o GRPO no experimento com Qwen (grad norm 0,23, cerca de 11,7 GB por GPU).
Adeque a técnica à lacuna. Use SFT para comportamentos ausentes. Use GRPO para refinar comportamentos que o modelo já produz às vezes. O teto realista do GRPO aqui foi de +6 a +8 pontos.
Verifique se há margem antes de gastar horas de GPU. Se a maioria dos grupos está toda correta (como no GSM8K com 90,5% base), a vantagem relativa colapsa para zero e o GRPO não recebe gradiente.

Limpeza de recursos

Para evitar cobranças contínuas, exclua os recursos de workload e depois a infraestrutura subjacente:

kubectl delete hyperpodpytorchjob lora-hyperpod-elastic --ignore-not-found
kubectl delete pod grpo-single inference-base inference-trained inference-grpo \
  eval-grpo fsx-storage-manager --ignore-not-found

Em seguida, reduza ou exclua o cluster Amazon SageMaker HyperPod ou seus grupos de nós, exclua o sistema de arquivos Amazon FSx for Lustre após exportar os checkpoints desejados e remova quaisquer buckets ou objetos do Amazon Simple Storage Service (Amazon S3) usados para dados de treinamento ou arquivos de checkpoint.

Conclusão

O experimento da AWS demonstra como fazer fine-tuning do GPT-OSS-20B para raciocínio multilíngue no Amazon SageMaker HyperPod e como decidir entre SFT e GRPO. O SFT fez a maior parte do trabalho: o raciocínio no idioma correto foi de 22% para 100%. O GRPO adicionou um refinamento real e modesto (+8 pontos na resposta) onde havia margem.

O experimento de controle com Qwen2.5-7B confirmou o padrão: GRPO adicionou +6% no raciocínio multilíngue, onde havia espaço para melhorar, e não adicionou nada no GSM8K, onde o modelo já estava no pico de acurácia (90,5%).

Framework de decisão:

Comportamento que o modelo não consegue fazer: use SFT.
Comportamento que ele faz às vezes, com margem: use SFT e depois GRPO para mais alguns pontos.
Modelo já no pico de acurácia, ou alvo com probabilidade zero na base: GRPO não vai ajudar.

Para experimentar por conta própria, o repositório awsome-distributed-ai contém todos os casos de teste. Para mais informações sobre o Amazon SageMaker HyperPod, consulte a página de detalhes do serviço e a documentação oficial do Amazon SageMaker HyperPod.

Fonte

Fine-tuning GPT-OSS-20B for multilingual reasoning on Amazon SageMaker HyperPod (https://aws.amazon.com/blogs/machine-learning/fine-tuning-gpt-oss-20b-for-multilingual-reasoning-on-amazon-sagemaker-hyperpod/)