Refinamiento por refuerzo (RFT) con modelos de Amazon Nova

Descripción general

¿Qué es el RFT?

El refinamiento por refuerzo (RFT) mejora el rendimiento del modelo al entrenarlo con señales de comentarios (puntuaciones medibles o recompensas que indican el rendimiento del modelo) en lugar de con respuestas exactas correctas. A diferencia del refinamiento supervisado (SFT), que aprende a partir de pares de entrada y salida, RFT utiliza funciones de recompensa para evaluar las respuestas del modelo y optimiza el modelo de forma iterativa para maximizar esas recompensas. Este enfoque destaca al definir el resultado exacto correcto, pero le permite medir de forma fiable la calidad de la respuesta.

Cuándo se debe usar el RFT

Use RFT cuando pueda definir criterios de éxito claros y medibles, pero resulte difícil proporcionar resultados correctos exactos para el entrenamiento. RFT es ideal para lo siguiente:

Los mejores casos de uso

El RFT sobresale en dominios en los que la calidad de los resultados se puede medir objetivamente, pero las respuestas óptimas son difíciles de definir por adelantado:

Modelos compatibles

Nova Lite 2.0

Información general sobre el formato de datos

Los datos de entrenamiento RFT deben seguir el formato de refinamiento por refuerzo de OpenAI. Cada ejemplo de entrenamiento es un objeto JSON que contiene lo siguiente:

Limitaciones actuales

Ejemplos de formato de datos

Cada ejemplo debe estar en una sola línea del archivo JSONL, con un objeto JSON por línea.

Chemistry problem

{
  "id": "chem-01",
  "messages": [
    {
      "role": "system",
      "content": "You are a helpful chemistry assistant"
    },
    {
      "role": "user",
      "content": "Calculate the molecular weight of caffeine (C8H10N4O2)"
    }
  ],
  "reference_answer": {
    "molecular_weight": 194.19,
    "unit": "g/mol",
    "calculation": "8(12.01) + 10(1.008) + 4(14.01) + 2(16.00) = 194.19"
  }
}

Math problem

{
  "id": "sample-001",  // Optional
  "messages": [
    {
      "role": "system",
      "content": "You are a math tutor"
    },
    {
      "role": "user",
      "content": "Solve: 2x + 5 = 13"
    }
  ],
  "reference_answer": {
    "solution": "x = 4",
    "steps": ["2x = 13 - 5", "2x = 8", "x = 4"]
  }
}

Code problem

{
  "id": "code-002",
  "messages": [
    {
      "role": "system",
      "content": "You are a helpful programming assistant"
    },
    {
      "role": "user",
      "content": "Write a Python function that reverses a string without using built-in reverse methods"
    }
  ],
  "reference_answer": {
    "code": "def reverse_string(s):  \n    result = ''  \n    for i in range(len(s) - 1, -1, -1):  \n        result += s[i]  \n    return result",
    "test_cases": [
      {
        "input": "hello",
        "expected_output": "olleh"
      },
      {
        "input": "",
        "expected_output": ""
      },
      {
        "input": "a",
        "expected_output": "a"
      },
      {
        "input": "Python123",
        "expected_output": "321nohtyP"
      }
    ],
    "all_tests_pass": true
  }
}

El campo reference_answer contiene el resultado esperado o los requisitos de evaluación que la función de recompensa utiliza para puntuar la respuesta del modelo. No se limita a resultados estructurados, sino que puede contener cualquier formato que ayude a la función de recompensas a evaluar la calidad.

Recomendaciones de tamaño de conjunto de datos

Punto de partida

Enfoque de priorización de la evaluación

Antes de invertir en un entrenamiento RFT a gran escala, evalúe el rendimiento de línea de base del modelo:

Comenzar con un conjunto de datos pequeño le permite:

Una vez validado, puede ampliarlo a conjuntos de datos más grandes para mejorar aún más el rendimiento.

Características de los datos de entrenamiento eficaces

Claridad y coherencia

Los buenos ejemplos del RFT requieren datos de entrada claros e inequívocos que permitan un cálculo preciso de las recompensas en las diferentes salidas del modelo. Evite el ruido en los datos, lo que incluye lo siguiente:

Cualquier ambigüedad inducirá a error durante el proceso de entrenamiento y provocará que el modelo aprenda comportamientos no deseados.

Diversidad

El conjunto de datos debe capturar toda la diversidad de casos de uso de producción para garantizar un rendimiento real sólido. Incluya:

Consideraciones sobre la función de recompensas

Diseñe la función de recompensas para un entrenamiento eficiente:

Las funciones de recompensas rápidas y escalables permiten una iteración rápida y una experimentación rentable.

Propiedades adicionales

El formato de datos del RFT admite campos personalizados que van más allá de los requisitos básicos del esquema (messages y reference_answer). Esta flexibilidad le permite agregar cualquier dato adicional que su función de recompensas necesite para una evaluación adecuada.

Propiedades adicionales comunes

Campos de metadatos de ejemplo:

Ejemplo con propiedades adicionales

{
  "messages": [
    {
      "role": "system",
      "content": "You are a math tutor"
    },
    {
      "role": "user",
      "content": "Solve: 2x + 5 = 13"
    }
  ],
  "reference_answer": {
    "solution": "x = 4",
    "steps": ["2x = 13 - 5", "2x = 8", "x = 4"]
  },
  "task_id": "algebra_001",
  "difficulty_level": "easy",
  "domain": "algebra",
  "expected_reasoning_steps": 3
}

Estos campos adicionales se transfieren a la función de recompensas durante la evaluación, lo que permite aplicar una lógica de puntuación sofisticada adaptada a su caso de uso específico.

Configuración de entrenamiento

Fórmula de muestra

# Note:
# This recipe can run on p5.48xlarge, p5e.48xlarge, and p5en.48xlarge instance types.
run:
  name: "my-rft-run"                           # Unique run name (appears in logs and artifacts).
  model_type: amazon.nova-2-lite-v1:0:256k
  model_name_or_path: nova-lite-2/prod
  data_s3_path: s3://<bucket>/<data-file>      # Training dataset in JSONL format.
  replicas: 4                                   # Number of total training instances.
  generation_replicas: 2                        # Number of total instances dedicated to response generation.
  reward_lambda_arn: arn:aws:lambda:<region>:<account-id>:function:<function-name>

  ## MLFlow configs
  mlflow_tracking_uri: "" # Required for MLFlow
  mlflow_experiment_name: "my-rft-experiment" # Optional for MLFlow. Note: leave this field non-empty
  mlflow_run_name: "my-rft-run" # Optional for MLFlow. Note: leave this field non-empty

## SMTJ RFT training configs
training_config:
  max_length: 8192                              # Context window (tokens) for inputs and prompt.
  global_batch_size: 32                         # Total samples per optimizer step across all replicas (16/32/64/128/256).
  reasoning_effort: high                        # Reasoning mode: high, low, or null for non-reasoning.

  data:
    shuffle: true                               # Shuffle training data each epoch.

  rollout:                                      # Controls how responses are generated for advantage calculation.
    rollout_strategy:
      type: off_policy_async                    # Asynchronous rollout for higher throughput.
      age_tolerance: 2                          # Maximum policy age before regeneration.
    advantage_strategy:
      number_generation: 4                      # Samples per prompt to estimate advantages (higher = lower variance but higher cost).
    generator:
      max_new_tokens: 6000                      # Cap on tokens generated per sample.
      set_random_seed: true                     # Seed generation for reproducibility across runs.
      temperature: 1                            # Softmax temperature for sampling.
    rewards:
      preset_reward_function: null              # Preset reward functions: exact_match or null for custom.
      api_endpoint:
        lambda_arn: arn:aws:lambda:<region>:<account-id>:function:<function-name>
        lambda_concurrency_limit: 12             # Max concurrent Lambda invocations (throughput vs. throttling).
        lambda_batch_size: 128                  # Number of samples per Lambda invocation.

  trainer:
    max_steps: 2                                # Steps to train for. One step = global_batch_size samples.
    save_steps: 5                               # Save a checkpoint every N steps.
    test_steps: 1                               # Run validation every N reference model updates.
    refit_freq: 4                               # Frequency of reference model updates.
    clip_ratio_high: 0.2                        # PPO clip ratio for policy updates.
    loss_scale: 1.0                             # Scaling factor for the policy loss.

    # RL parameters
    ent_coeff: 0.0                              # Entropy bonus added to the policy loss (higher = more exploration).
    kl_loss_coef: 0.0                           # Weight on the KL penalty between the current and reference policy.

    optim_config:                               # Optimizer settings.
        lr: 1e-6                                # Learning rate.
        weight_decay: 0.0                       # L2 regularization strength (0.0 to 1.0).
        adam_beta1: 0.9
        adam_beta2: 0.95

    peft:                                       # Parameter-efficient fine-tuning (LoRA).
        peft_scheme: "lora"                     # Enable LoRA for PEFT.
        lora_tuning:
            alpha: 64                           # LoRA scaling factor.
            lora_plus_lr_ratio: 64.0            # LoRA+ learning rate scaling factor (0.0 to 100.0).

Entrenamiento RFT con LLM como juez

Descripción general

Los modelos de lenguaje de gran tamaño (LLM) se utilizan cada vez más como jueces en los flujos de trabajo de refinamiento por refuerzo (RFT), ya que proporcionan señales de recompensas automatizadas que guían la optimización de los modelos. Con este enfoque, el LLM evalúa las salidas del modelo en función de requisitos específicos (ya sea la corrección, la calidad, la adherencia al estilo o la equivalencia semántica) y asigna recompensas que impulsan el proceso de aprendizaje por refuerzo.

Esto es particularmente valioso para tareas en las que las funciones de recompensas tradicionales son difíciles de definir mediante programación, como determinar si diferentes representaciones (como “1/3”, “0,333” y “un tercio”) son semánticamente equivalentes, o evaluar cualidades matizadas como la coherencia y la relevancia. Al utilizar jueces basados en LLM como funciones de recompensas, puede escalar el RFT a dominios complejos sin necesidad de una amplia anotación humana, lo que permite una iteración rápida y una mejora continua de los modelos en varios casos de uso, más allá de los problemas de alineación tradicionales.

Selección del modo de razonamiento

Modos disponibles

Cuándo se debe usar cada modo

Utilice el razonamiento alto para:

Utilice ninguno (omita reasoning_effort) o un razonamiento bajo para:

Compensaciones de costo y rendimiento

Los modos de razonamiento alto aumentan lo siguiente:

Validación del juez de LLM

Antes de implementar un LLM como juez en producción, compruebe que las evaluaciones del modelo del juez se ajusten al juicio humano. Esto implica:

Este paso de validación ayuda a garantizar que el proceso de evaluación automatizada produzca modelos que cumplan con los requisitos de calidad de producción.

Configuración de Lambda para juez de LLM

El uso de un LLM como juez es una extensión del uso de las funciones de Lambda para el aprendizaje por refuerzo con recompensas verificables (RLVR). Dentro de la función de Lambda, puede hacer una llamada a uno de los modelos alojados en Amazon Bedrock.

Requisitos de configuración importantes:

Creación y ejecución de trabajos

Inicio de un trabajo de entrenamiento

Utilice la plantilla de cuaderno de trabajos de entrenamiento de SageMaker: https://docs.aws.amazon.com/sagemaker/latest/dg/nova-fine-tuning-training-job.html#nova-model-training-jobs-notebook.

Requisitos de instancias

El contenedor admite tanto el entrenamiento de rango completo como el de LoRA:

Supervisión del entrenamiento

Los registros de entrenamiento incluyen métricas completas en cada paso. Categorías de métricas clave:

Métricas de recompensas

Comportamiento del modelo

Estado del entrenamiento

Características de la respuesta

Desempeño

Uso de recursos

Uso de modelos de ajuste fino

Después de finalizar el entrenamiento, el punto de control final del modelo se guarda en la ubicación de salida especificada. La ruta del punto de control está disponible en:

Limitaciones y prácticas recomendadas

Limitaciones

Prácticas recomendadas

Empezar poco a poco

Evaluación del entrenamiento previo

Supervisión del entrenamiento

Optimización de las funciones de recompensas

Estrategia de iteración

Si las recompensas no mejoran:

Aprendizaje curricular adaptativo

El aprendizaje curricular adaptativo es una característica opcional que selecciona dinámicamente qué peticiones de entrenamiento presentar al modelo durante RFT. En lugar de entrenar de manera uniforme con todas las peticiones, el entrenador utiliza el propio modelo para predecir la dificultad de cada petición y selecciona peticiones dentro del rango de dificultad productivo, donde el modelo a veces tiene éxito y otras veces falla. Esto maximiza la variación de resultados dentro de cada grupo de despliegue de GRPO, lo que genera una señal de ventaja más sólida, una convergencia más rápida y una mayor estabilidad del entrenamiento de RL al reducir las actualizaciones ruidosas de gradientes provenientes de peticiones demasiado fáciles o demasiado difíciles.

Cómo funciona el aprendizaje curricular adaptativo

Cuando el aprendizaje curricular adaptativo está habilitado, el ciclo de entrenamiento agrega una fase de predicción y selección antes de cada paso de despliegue:

Cuándo usar el aprendizaje curricular adaptativo

El aprendizaje curricular adaptativo es más eficaz en los siguientes escenarios:

Configuración del aprendizaje curricular adaptativo

Agregue el bloque adaptive_curriculum en trainer dentro de la receta para habilitar el aprendizaje curricular adaptativo:

training_config:
  trainer:
    adaptive_curriculum:
      enable: true                               # Enable adaptive curriculum prompt selection.
      selection_pool_multiplier: 8               # Score 8 x global_batch_size candidates, keep best global_batch_size.
      prediction_mode: pass_rate                 # "pass_rate" for discrete rewards; "spread" for continuous rewards.
      exemplar_history_steps: 1                  # Previous training steps kept in the rolling exemplar history buffer.
      reinforce_coef: 0.01                       # Scale factor for the REINFORCE loss that trains the predictor (0 disables).
      predictor_prompt_column: predictor_prompt  # Dataset field with clean problem text used by the predictor.
      selection_lookahead_steps: 4               # Future training batches pre-approved per curriculum screening pass.

La siguiente tabla describe cada parámetro del aprendizaje curricular adaptativo:

Recomendaciones para conjuntos de datos de contexto largo

El aprendizaje curricular adaptativo funciona mediante la ejecución de un predictor ligero de tasa de aprobación sobre un conjunto de peticiones candidatas y la selección del lote más productivo para el despliegue. Cuando max_prompt_length es corto (unos pocos miles de tokens o menos), el predictor se ejecuta en unos segundos por ronda de evaluación y la sobrecarga del aprendizaje curricular es insignificante. A medida que aumenta la longitud de las peticiones, el tiempo de inferencia del predictor aumenta aproximadamente de forma cuadrática (la atención es O(n²) en función de la longitud de la secuencia), por lo que la evaluación puede llegar a dominar el tiempo por paso en conjuntos de datos donde las peticiones superan aproximadamente los 8000 tokens.

La siguiente configuración mantiene el aprendizaje curricular adaptativo práctico y rentable en conjuntos de datos de contexto largo. Aplíquela en conjunto; cada configuración aborda distintos componentes del costo de evaluación.

Ejemplo de configuración ajustada para un conjunto de datos de contexto largo (peticiones de aproximadamente 24 000 a 32 000 tokens):

training_config:
  max_length: 32768
  global_batch_size: 32

  trainer:
    adaptive_curriculum:
      enable: true
      selection_pool_multiplier: 4        # Smaller pool keeps predictor prefill bounded.
      selection_lookahead_steps: 1        # Predictor runs once per training step.
      prediction_mode: pass_rate
      exemplar_history_steps: 1
      reinforce_coef: 0.01
      predictor_prompt_column: predictor_prompt

Preparación de datos para el aprendizaje curricular adaptativo

Cuando utilice aprendizaje curricular adaptativo, los datos de entrenamiento deben incluir un campo predictor_prompt (o el nombre de campo especificado en predictor_prompt_column) que contenga una versión concisa del texto del problema. Este campo lo utiliza el predictor de tasa de aprobación para evaluar rápidamente la dificultad de una petición sin procesar el contexto completo de la conversación.

Ejemplo de entrada JSONL con una petición para el predictor:

{
  "messages": [
    {
      "role": "system",
      "content": "You are a math tutor. Show your work step by step."
    },
    {
      "role": "user",
      "content": "A train travels 120 miles in 2 hours. If it then increases speed by 50%, how far will it travel in the next 3 hours?"
    }
  ],
  "reference_answer": "270 miles",
  "predictor_prompt": "A train travels 120 miles in 2 hours. Speed increases 50%. Distance in next 3 hours?"
}

Si el campo predictor_prompt no está presente, el sistema vuelve a utilizar la petición completa del campo messages.

Ejemplo de receta RFT completa con aprendizaje curricular adaptativo

El siguiente ejemplo muestra una receta LoRA RFT completa con aprendizaje curricular adaptativo habilitado:

run:
  name: "my-rft-adaptive-curriculum"
  model_type: amazon.nova-2-lite-v1:0:256k
  model_name_or_path: nova-lite-2/prod
  data_s3_path: s3://<bucket>/<data-file>
  replicas: 4
  generation_replicas: 2
  reward_lambda_arn: arn:aws:lambda:<region>:<account-id>:function:<function-name>

training_config:
  max_length: 8192
  global_batch_size: 32
  reasoning_effort: null                        # Non-reasoning mode.

  data:
    shuffle: true

  rollout:
    rollout_strategy:
      type: off_policy_async
      age_tolerance: 2
    advantage_strategy:
      number_generation: 16                     # Higher n for better advantage estimates.
    generator:
      max_new_tokens: 6000
      temperature: 1.0
    rewards:
      preset_reward_function: exact_match       # Or null for custom Lambda reward.
      api_endpoint:
        lambda_arn: ${oc.select:run.reward_lambda_arn}   # Reuse the top-level run.reward_lambda_arn so the two stay in sync.
        lambda_concurrency_limit: 12
        lambda_batch_size: 128

  trainer:
    max_steps: 500
    save_steps: 50
    test_steps: 25
    refit_freq: 4
    clip_ratio_high: 0.2
    ent_coeff: 0.0
    kl_loss_coef: 0.0

    optim_config:
      lr: 1e-6
      weight_decay: 0.0

    peft:
      peft_scheme: "lora"
      lora_tuning:
        alpha: 64
        lora_plus_lr_ratio: 64.0

    adaptive_curriculum:
      enable: true
      selection_pool_multiplier: 8
      prediction_mode: pass_rate
      exemplar_history_steps: 1
      reinforce_coef: 0.01
      predictor_prompt_column: predictor_prompt

Supervisión del aprendizaje curricular adaptativo

Cuando el aprendizaje curricular adaptativo está habilitado, se registran métricas adicionales en cada paso de entrenamiento:

Capacidades avanzadas: Nova Forge

Para los usuarios que requieren capacidades avanzadas más allá de las limitaciones del RFT estándar, Nova Forge está disponible como un servicio de suscripción de pago que ofrece:

Nova Forge se ejecuta en SageMaker HyperPod y se ha diseñado para ayudar a los clientes empresariales a crear sus propios modelos de frontera.

Comandos y consejos útiles

Hay disponible una colección de scripts de observabilidad para ayudar a supervisar el estado y el progreso de los trabajos de entrenamiento.

Los scripts disponibles le permiten:

Instalación

pip install boto3
git clone https://github.com/aws-samples/amazon-nova-samples.git
cd amazon-nova-samples/customization/SageMakerUilts/SageMakerJobsMonitoring/

Uso básico

# Enabling email notifications for training job status updates
python enable_sagemaker_job_notifs.py --email test@amazon.com test2@gmail.com --region us-east-1 --platform SMTJ

Creating resources........
Please check your email for a subscription confirmation email, and click 'Confirm subscription' to start receiving job status email notifications!
You'll receive the confirmation email within a few minutes.

# Obtaining training time estimates based on job configurations
python get_training_time_estimate.py

# Obtaining approximations for how long training is expected to take for in-progress jobs
python get-training-job-progress.py --region us-east-1 --job-name my-training-job --num-dataset-samples 1000

Consulte en esta página los detalles y ejemplos adicionales.