Ajuste fino de GPT-OSS-20B para razonamiento multilingüe en Amazon SageMaker HyperPod

GPT-OSS-20B es un modelo de razonamiento de pesos abiertos que piensa en inglés. Si le pide razonar en francés, alemán, español o italiano, reconoce la instrucción, pero luego produce su cadena de pensamiento en inglés de todos modos. En una prueba de razonamiento multilingüe de 50 ejemplos, el modelo base razonó en el idioma objetivo correcto solo el 22 por ciento de las veces. Ese 22 por ciento correspondía casi por completo a los ejemplos en inglés (inglés 100 por ciento, francés 0 por ciento, alemán 10 por ciento, español 0 por ciento, italiano 0 por ciento).

Para los equipos que despliegan modelos de razonamiento para usuarios que no hablan inglés, esa brecha importa. Un usuario que hace una pregunta en francés espera que el modelo razone y responda en francés, no que cambie al inglés a mitad del pensamiento. Cerrar esa brecha es un problema de ajuste fino. Esto plantea una pregunta que todo equipo enfrenta cuando comienza a personalizar modelos de pesos abiertos: ¿ajuste fino supervisado (SFT), aprendizaje por refuerzo con Group Relative Policy Optimization (GRPO), o ambos?

En esta publicación, explicamos cómo ajustamos GPT-OSS-20B para razonamiento multilingüe en Amazon SageMaker HyperPod. Mostramos cuándo el SFT solo fue suficiente y cuándo GRPO añadió un valor medible por encima. Para entender de manera más amplia cuándo ayuda GRPO, también ejecutamos un experimento controlado con Qwen2.5-7B en dos tareas con diferentes brechas de capacidad. La respuesta corta: el SFT hace la mayor parte del trabajo, GRPO refina cuando hay margen para refinar, y emparejar la técnica con la brecha de capacidad es lo que decide el resultado.

Descripción general de la solución

Esta sección cubre la tarea y el conjunto de datos de razonamiento multilingüe, cómo se complementan el SFT y GRPO, y la pila de entrenamiento que ejecutamos en Amazon SageMaker HyperPod.

La tarea y el conjunto de datos

El conjunto de datos es HuggingFaceH4/Multilingual-Thinking —un conjunto de datos de razonamiento de cadena de pensamiento traducido al francés, español, alemán e italiano junto con el inglés, con aproximadamente 949 ejemplos de entrenamiento y 50 de validación. El modelo recibe una instrucción de reasoning language: system prompt y debe tanto razonar como responder en ese idioma objetivo, anulando su fuerte prior hacia el inglés. La evaluación consta de 50 casos de prueba: 10 problemas de matemáticas/lógica en 5 idiomas (EN/FR/DE/ES/IT).

Dos métricas importan:

• Correctitud del idioma de razonamiento – ¿está la cadena de pensamiento en el idioma solicitado?
• Correctitud del idioma de la respuesta – ¿está la respuesta final en el idioma solicitado?

Cómo se complementan el SFT y GRPO

La forma más clara de pensar en las dos técnicas: el SFT enseña un comportamiento que el modelo no tiene. GRPO optimiza un comportamiento que el modelo ya puede producir a veces. GRPO funciona muestreando un grupo de respuestas por prompt y empujando la política hacia las que están por encima del promedio. Si el comportamiento deseado tiene probabilidad cero en la distribución base, hay muy poco para que GRPO refuerce. Esa única idea explica casi todos los resultados de esta publicación.

La pila de entrenamiento

Ejecutamos en Amazon SageMaker HyperPod con orquestación de Amazon Elastic Kubernetes Service (Amazon EKS). Para las ejecuciones iniciales de SFT y GRPO en GPT-OSS-20B, usamos la biblioteca PyTorch TRL. Para el experimento de control con Qwen, cambiamos a veRL como entrenador, con Ray y KubeRay para la orquestación distribuida y vLLM para una generación rápida de rollouts. Los checkpoints y los conjuntos de datos se almacenaban en Amazon FSx for Lustre, y la estructura entre nodos usaba Elastic Fabric Adapter (EFA). Todo el ajuste fino usó Low-Rank Adaptation (LoRA), que es eficiente en parámetros y adecuado tanto para SFT como para GRPO.

El siguiente diagrama muestra la ruta de datos de entrenamiento de extremo a extremo. Un envío del usuario llega al plano de control de Amazon EKS del clúster de Amazon SageMaker HyperPod. Amazon EKS programa un pod de Ray head más pods de trabajadores que ejecutan el entrenador veRL y los trabajadores de rollout de vLLM. Ambos conjuntos de pods montan Amazon FSx for Lustre para el modelo base, los datos de entrenamiento y los checkpoints. Los colectivos entre nodos usan EFA. Los rollouts fluyen desde vLLM de vuelta al entrenador veRL, y el entrenador escribe los checkpoints de nuevo en Amazon FSx for Lustre.

Figura 1: Ruta de datos de entrenamiento de extremo a extremo para el ajuste fino con SFT y GRPO en Amazon SageMaker HyperPod con veRL, Ray, KubeRay, vLLM y Amazon FSx for Lustre para los experimentos de control con Qwen

Lo siguiente describe cada salto con más detalle. El remitente inicia la ejecución desde Amazon SageMaker Studio o la AWS Command Line Interface (AWS CLI). El plano de control de Amazon EKS en Amazon SageMaker HyperPod coloca los pods en nodos de GPU e inicia un Ray head con pods de trabajadores para el entrenador veRL y los trabajadores de rollout de vLLM. Amazon FSx for Lustre proporciona almacenamiento compartido de alto rendimiento para el modelo base, los datos de entrenamiento y los checkpoints. Cuando el entrenamiento se completa, el entrenador escribe el adaptador LoRA de nuevo en Amazon FSx for Lustre para su evaluación o servicio posterior.

Requisitos previos

Para seguir adelante, necesita:

• Un clúster de Amazon SageMaker HyperPod con orquestación de Amazon EKS. El SFT de GPT-OSS usó 4× ml.g6e.12xlarge (16× L40S 48GB). El experimento de control con Qwen usó 2× g6e.48xlarge (16× L40S).
• Un sistema de archivos de Amazon FSx for Lustre montado en el clúster con un PersistentVolumeClaim denominado fsx-claim en /fsx.
• Instancias habilitadas para EFA y una NVIDIA Collective Communications Library (NCCL) funcional y aws-ofi-nccl stack para SFT multinodo.
• Un registro de contenedores como Amazon Elastic Container Registry (Amazon ECR) para alojar la imagen de entrenamiento construida a partir del repositorio de referencia.
• Un token de acceso de HuggingFace para openai/gpt-oss-20b y el HuggingFaceH4/Multilingual-Thinking conjunto de datos.
• La AWS Command Line Interface (AWS CLI), kubectl, docker, y envsubst disponibles en su estación de trabajo.

Configure el entorno de entrenamiento

Todos los comandos, manifiestos y scripts mencionados a continuación se encuentran en el repositorio de GitHub awsome-distributed-ai bajo 3.test_cases/pytorch/trl/gpt-oss-lora-grpo/. Clone el repositorio y cámbiese al directorio del caso de prueba:

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

Copie la plantilla de entorno y luego edite env_vars para configurar su registro de ECR, el nombre de la imagen, la etiqueta y el token de HuggingFace:

cp env_vars.example env_vars
source env_vars

Construya la imagen de entrenamiento y envíela a su registro. La imagen incluye PyTorch, TRL, HuggingFace Transformers, PEFT y los scripts de entrenamiento e inferencia:

cd artifacts && ./build_push.sh && cd ..

Apunte kubectl a su clúster y aplique el pod del administrador de almacenamiento FSx. El administrador de almacenamiento le da un shell al sistema de archivos compartido para la conversión de checkpoints y la inspección de registros de evaluación:

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 -

Realice el fine-tuning de GPT-OSS-20B con SFT y LoRA

Comenzamos midiendo el modelo base, luego configuramos la ejecución de SFT y examinamos los resultados.

La brecha del modelo base

El GPT-OSS-20B base en la evaluación de 50 casos obtuvo un 22 por ciento de razonamiento-idioma correcto (11/50). Por idioma: inglés 100 por ciento, francés 0 por ciento, alemán 10 por ciento, español 0 por ciento, italiano 0 por ciento. El modelo recurre al razonamiento en inglés independientemente de la instrucción.

Configuración de SFT

La configuración de SFT usó LoRA r=8, α=16, entrenada durante 1,000 pasos en 4× ml.g6e.12xlarge (16× L40S 48GB) con Fully Sharded Data Parallel (FSDP) full-shard, precisión bf16 y gradient checkpointing. El trabajo se ejecuta como un HyperPodPyTorchJob elástico administrado por Kueue, lanzado con hyperpodrun.

Aplique el manifiesto de SFT:

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

Observe el progreso del entrenamiento:

kubectl get pods -l app=gpt-oss-lora-elastic
kubectl logs -f -l app=gpt-oss-lora-elastic

La pérdida y la norma del gradiente se imprimen en cada paso. Una ejecución saludable se ve así:

{'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}

Cuando el entrenamiento se complete, convierta el checkpoint de FSDP al formato Parameter-Efficient Fine-Tuning (PEFT) que el entrenador GRPO de TRL y los scripts de inferencia esperan:

kubectl exec -it fsx-storage-manager -- python /app/convert_fsdp_checkpoint.py \
  --checkpoint /fsx/checkpoints/checkpoint-1000 \
  --output    /fsx/checkpoints/converted-peft/lora-checkpoint-1000-peft

Resultados de SFT y la curva de aprendizaje por etapas

Después de SFT, el modelo razona en el idioma correcto el 100 por ciento del tiempo (50/50) en los cinco idiomas. SFT aprendió esto en dos etapas:

• Alrededor del paso 200, el modelo aprendió a razonar en el idioma objetivo (la métrica de idioma-de-razonamiento alcanza el 100 por ciento).
• Alrededor del paso 1000, el modelo aprendió además a dar la respuesta en el idioma objetivo.

Evaluación final de SFT (paso 1000):

• Correcto en idioma de razonamiento: 100 por ciento (50/50).
• Correcto en idioma de respuesta: 62 por ciento (31/50).
• Ambos correctos: 62 por ciento (31/50).
• Precisión de respuestas por idioma: FR 60 por ciento, DE 60 por ciento, ES 30 por ciento, IT 60 por ciento, EN 100 por ciento.

Conclusión: SFT no es opcional aquí. El modelo base no puede seguir en absoluto la instrucción de idioma de razonamiento, y SFT lleva el idioma de razonamiento del 22 por ciento al 100 por ciento. La brecha restante es el idioma de respuesta, en 62 por ciento. Esa es la mitad más difícil, porque la respuesta final es corta y el sesgo hacia el inglés es más fuerte ahí.

Refinar el idioma de respuesta con GRPO

El idioma de respuesta era el objetivo obvio para GRPO. SFT ya producía razonamiento correcto el 100 por ciento de las veces, y el modelo a veces producía la respuesta en el idioma correcto, así que había una señal que reforzar.

Ejecutar GRPO sobre el checkpoint de SFT

Ejecutamos GRPO con GRPOTrainerde TRL, partiendo del checkpoint de SFT. La configuración: cargar el LoRA de SFT, llamar a merge_and_unload, y añadir un nuevo LoRA (r=8, α=16, aproximadamente 0.068 por ciento de parámetros entrenables, 13.6M de 20B). El entrenamiento usó K=8 generaciones por prompt, device_map=auto, tasa de aprendizaje 1e-6, y 10 épocas. La función de recompensa puntuó el idioma de respuesta (±5.0), el idioma de razonamiento (±1.5) y la brevedad (+0.5 o −1.0), para un rango combinado de −7.5 a +7.0.

Lanzar GRPO desde el checkpoint de SFT convertido:

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

GRPO escribe checkpoints en /fsx/checkpoints/grpo-singlenode/. Elige el mejor checkpoint contra tu eval de reserva (la métrica de idioma de respuesta alcanzó su pico al principio de esta ejecución) y conviértelo al formato PEFT para inferencia y evaluación.

Resultado (desde el checkpoint de SFT):

• Correcto en idioma de razonamiento: 100 por ciento (reservado: GRPO no rompió lo que SFT enseñó).
• Correcto en idioma de respuesta: 62 a 70 por ciento (+8 puntos, 35/50).
• Ambos correctos: 70 por ciento.
• Respuesta por idioma: FR 70 por ciento, DE 80 por ciento, ES 60 por ciento, IT 40 por ciento, EN 100 por ciento.

GRPO logró una ganancia real de +8 puntos en la métrica de idioma de respuesta mientras mantenía el razonamiento en 100 por ciento. La ganancia vino de donde había margen para mejorar. El idioma de respuesta era el comportamiento a veces correcto en el que SFT no había convergido completamente, y GRPO lo reforzó.

Compara la inferencia de base, SFT y GRPO lado a lado

Para ver la diferencia cualitativa, despliega tres pods de inferencia que comparten el mismo modelo base pero cargan adaptadores 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 -

Ejecuta el mismo prompt de razonamiento en francés en cada variante:

kubectl exec -it inference-base    -- python /app/inference_g6e.py \
  --reasoning_language French

kubectl exec -it inference-trained -- python /app/inference_g6e.py \
  --checkpoint /fsx/checkpoints/converted-peft/lora-checkpoint-1000-peft \
  --reasoning_language French

kubectl exec -it inference-grpo    -- python /app/inference_grpo_new.py \
  --use_grpo \
  --grpo_checkpoint /fsx/checkpoints/converted-peft/checkpoint-1470-peft \
  --reasoning_language French

El pod base razona en inglés a pesar de la instrucción. El pod SFT razona en francés pero a menudo responde en inglés. El pod GRPO razona y responde en francés. Para reproducir la eval completa de 50 casos que generó los números anteriores, ejecuta el manifiesto de evaluación y lee el archivo de resultados desde el gestor de almacenamiento FSx:

envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/eval-grpo.yaml | kubectl apply -f -
kubectl logs -f eval-grpo
kubectl exec fsx-storage-manager -- \
  cat /fsx/experiments/grpo_eval_checkpoint1470_*.txt

Para entender de manera más amplia cuándo GRPO añade valor, realizamos un experimento controlado en Qwen2.5-7B en dos tareas con diferentes brechas de capacidad.

Prueba el patrón con un control de Qwen2.5-7B

Para probar la tesis de «emparejar la técnica con la brecha de capacidad» en un modelo diferente y tareas distintas, realizamos experimentos controlados en Qwen2.5-7B-Instruct. La primera tarea es la misma tarea de razonamiento multilingüe (donde existe margen). La segunda tarea es GSM8K, donde el modelo ya está en su precisión máxima. Este es un experimento interno y el repositorio público aún no ha sido publicado.

Infraestructura: 2× g6e.48xlarge (16× L40S), EFA 100Gbps, FSx Lustre, Ray 2.44 + KubeRay 1.4.2, veRL v0.6.1, vLLM 0.8.4.

Brazos: A = base, B = solo SFT, C = SFT+GRPO.

Razonamiento multilingüe (50 casos)

GSM8K, una tarea sin margen

En GSM8K, el Qwen2.5-7B base ya obtiene 90.5 por ciento (181/200). GRPO en el paso 30 también obtiene 90.5 por ciento. Cuando la mayoría de los grupos muestreados son todos correctos, la ventaja relativa al grupo colapsa a cero, por lo que no hay señal de aprendizaje. Esta es la otra mitad de la tesis: GRPO necesita margen (pass@k >> pass@1) para tener algo que optimizar. Sin margen no hay ganancia.

Buenas prácticas

1. Comenzar con SFT. GRPO no puede enseñar un comportamiento con probabilidad base cero. Solo puede reforzar respuestas que el modelo ya produce a veces. SFT llevó el idioma de razonamiento del 22 por ciento al 100 por ciento, y cada victoria de GRPO en este trabajo se construyó sobre ese checkpoint de SFT.
2. Usar LoRA tanto para SFT como para GRPO. LoRA r=8/α=16 fue eficiente en parámetros (0.068 por ciento entrenable), y en el control de Qwen se mantuvo numéricamente estable durante GRPO (norma de gradiente 0.23, unos 11.7 GB por GPU). Es una opción por defecto sólida cuando quieres que SFT y GRPO compartan la misma forma de adaptador.
3. Ajustar la técnica a la brecha. Usar SFT para comportamientos ausentes. Usar GRPO para refinar un comportamiento existente que a veces es correcto. El techo realista de GRPO aquí fue de +6 a +8 puntos.
4. Compruebe el margen de mejora (reward headroom) antes de gastar horas de GPU. Si la mayoría de los grupos son todos correctos (GSM8K 90.5 percent base), la ventaja relativa del grupo colapsa a cero y GRPO no obtiene gradiente. Verifique que exista margen (pass@k >> pass@1) antes de iniciar una ejecución de GRPO.

Limpieza de recursos

Para evitar cargos continuos, elimine los recursos del workload y luego la infraestructura subyacente:

1. Elimine los pods y jobs de entrenamiento e inferencia:

   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

2. Reduzca la escala o elimine el clúster de Amazon SageMaker HyperPod o sus grupos de nodos.
3. Elimine el sistema de archivos de Amazon FSx for Lustre una vez que haya exportado los checkpoints que desee conservar.
4. Elimine los buckets u objetos de Amazon Simple Storage Service (Amazon S3) utilizados para datos de entrenamiento o archivos de checkpoints.

Conclusión

Mostramos cómo ajustar (fine-tune) GPT-OSS-20B para razonamiento multilingüe en Amazon SageMaker HyperPod, y cómo decidir entre SFT y GRPO. SFT hizo la mayor parte del trabajo: el idioma de razonamiento pasó de 22 percent a 100 percent. GRPO añadió un refinamiento real y modesto por encima (answer 62 percent a 70 percent) donde existía margen de mejora. Un control paralelo con Qwen2.5-7B confirmó el patrón. GRPO añadió +6 percent en razonamiento multilingüe, donde había espacio para mejorar. No aportó ninguna ganancia en GSM8K, donde el modelo ya estaba en la precisión máxima (90.5 percent).

Marco de decisión:

• Comportamiento que el modelo no puede hacer en absoluto: use SFT.
• Comportamiento que realiza a veces, con margen de mejora: use SFT then GRPO para unos puntos más.
• Modelo ya en la precisión máxima, o el objetivo tiene probabilidad base cero: GRPO won’t help.

Para probar esto usted mismo, recorra los casos de prueba en el repositorio awsome-distributed-ai . Para más información sobre Amazon SageMaker HyperPod, consulte la service detail page y la Amazon SageMaker HyperPod documentation. Para un recorrido de Ray en Amazon SageMaker HyperPod, vea la publicación relacionada Ejecución de entrenamiento distribuido en Amazon SageMaker HyperPod con Ray y Amazon SageMaker Studio. El enlace se añadirá después de que se publique esa entrada. Comparta sus resultados en los comentarios.

Sobre los autores