Optimización de políticas proximales (PPO) - Amazon SageMaker AI

Optimización de políticas proximales (PPO)

La optimización de políticas proximales (PPO) es el proceso de utilizar varios modelos de machine learning para entrenar y otorgar una puntuación a un modelo. Los siguientes modelos forman parte del proceso de PPO:

  • Modelo de entrenamiento de actores o de políticas: modelo de refinamiento supervisado (SFT) que se refina y actualiza en cada época. Las actualizaciones se realizan mediante el muestreo de peticiones, la generación de respuestas completas y la actualización de las ponderaciones mediante un objetivo sustituto recortado. Esto limita el cambio de rentabilidad logarítmica por token, de modo que cada paso de la política es próximo al anterior, y se preserva la estabilidad del entrenamiento.

  • Modelo de generación de actores: un modelo que genera respuestas o finalizaciones de peticiones para ser juzgado por el modelo de recompensa y el modelo crítico. Las ponderaciones de este modelo se actualizan a partir del modelo de entrenamiento de actores o modelo de políticas en cada época.

  • Modelo de recompensa: modelo con ponderaciones congeladas que se utiliza para otorgar una puntuación al modelo de generación de actores.

  • Modelo crítico: modelo con ponderaciones descongeladas que se utiliza para otorgar una puntuación al modelo de generación de actores. Esta puntuación suele considerarse una estimación de la recompensa total que recibe el actor al generar los tokens restantes.

  • Modelo de anclaje: modelo SFT con ponderaciones congeladas que se utiliza para calcular la divergencia de KL entre el modelo de entrenamiento de actores y el modelo base. El modelo de anclaje garantiza que las actualizaciones del modelo de actor no sean demasiado drásticas en comparación con el modelo base. Los cambios drásticos pueden provocar inestabilidad o una degradación del rendimiento.

Los datos de entrenamiento deben estar en formato JSONL, donde cada línea contenga un único objeto JSON que represente un ejemplo de entrenamiento. A continuación se muestra un ejemplo:

{
    "turns": ["string", "string", ...], // Required
    "turns_to_mask": [integer, integer, ...], // Required
    "reward_category": "string", // Required
    "meta_data": {} // Optional
}

  • turns es una matriz de matrices de cadenas de conversación que representan la secuencia de diálogo. Esta línea contiene peticiones del sistema, mensajes del usuario y respuestas del bots. Los mensajes de usuario suelen terminar con “Bot:” para indicar dónde comienza la salida del modelo. Por ejemplo, [["System prompt"], ["User: Question Bot:"], ["Bot response"]].

  • turns_to_mask es una matriz de índices basados en 0 que identifica qué giros no deberían recibir actualizaciones de gradiente. Los giros enmascarados suelen ser peticiones del sistema y turnos del usuario. Por ejemplo, [0, 1, 3] enmascara la petición del sistema y los mensajes del usuario (el primer y el tercer mensaje).

  • reward_category es una cadena que identifica qué aspectos del rendimiento del modelo se evalúan. Se utiliza para seleccionar la categoría de modelo de recompensa adecuada durante el entrenamiento. La categoría de recompensa está disponible para las siguientes categorías de recompensa: default, math, coding, if, rag y rai.

  • meta_data es un objeto opcional que contiene información adicional de carácter contextual o de verdad fundamental. Puede incluir identificadores, información sobre fuentes o contexto de conversación. La estructura es flexible en función de las necesidades del conjunto de datos.

A continuación se ofrece un registro de ejemplo:

{
    "turns": ["You are a helpful AI assistant.",
        "User: What is ML? Bot:",
        "Machine learning is...", "User: Examples? Bot:",
        "Email spam filtering is..."
    ],
    "turns_to_mask": [0, 1, 3],
    "reward_category": "default",
    "meta_data": {
        "messages": [{
                "role": "system",
                "content": "You are a helpful AI assistant."
            },
            {
                "role": "user",
                "content": "What is ML?"
            },
            {
                "role": "assistant",
                "content": "Machine learning is..."
            },
            {
                "role": "user",
                "content": "Examples?"
            },
            {
                "role": "assistant",
                "content": "Email spam filtering is..."
            }
        ]
    }
}

El marco de modelado de recompensa implementa optimización multidimensional en distintos objetivos categóricos para facilitar una sólida convergencia de modelos. La categoría de recompensa debe seleccionarse en función de la tarea para la que se debe optimizar el modelo.

Se recomienda seguir estas directrices para seleccionar el marco adecuado para las tareas:

  • default: optimizador de uso general para tareas de conversación estándar e interacciones básicas. Se utiliza para conversaciones y debates generales, tareas básicas de escritura, respuesta sencilla a preguntas y consultas de conocimientos no especializados.

    A continuación se muestra un ejemplo:

    {
    "turns": ["Write a summary of climate change"],
    "turns_to_mask": [0],
    "reward_category": "default"
}

  • math: optimizador especializado para cálculos matemáticos y tareas de razonamiento numérico. Se utiliza para resolución de problemas matemáticos, cálculos aritméticos, ecuaciones algebraicas, problemas geométricos y análisis estadístico.

    A continuación se muestra un ejemplo:

    {
    "turns": ["Calculate the derivative of x²"],
    "turns_to_mask": [0],
    "reward_category": "math"
}

  • coding: categoría dedicada para consultas relacionadas con programación y desarrollo de software. Se utiliza para implementación de código, asistencia en depuración, diseño de algoritmos, documentación técnica y preguntas sobre arquitectura del sistema.

    A continuación se muestra un ejemplo:

    {
    "turns": ["Write a function to check if a string is palindrome"],
    "turns_to_mask": [0],
    "reward_category": "coding"
}

  • if: categoría para tareas que requieren ejecución procedimental precisa y orientación paso a paso. Se utiliza para procedimientos de varios pasos, instrucciones secuenciales, descomposición de tareas complejas y documentación de procesos.

    A continuación se muestra un ejemplo:

    {
    "turns": ["Provide steps to deploy a web application"],
    "turns_to_mask": [0],
    "reward_category": "if"
}

  • rag: categoría de recompensa para tareas que requieren respuesta a consultas basadas específicamente en información contextual recuperada. Se usa cuando las respuestas deben obtenerse directamente de los materiales de referencia proporcionados, sintetizando contenido fáctico sin ir más allá del alcance de la información recuperada, lo que garantiza que las respuestas se basen en el contexto proporcionado más que en el conocimiento general.

    A continuación se muestra un ejemplo:

    {
            "turns": ["The Synthesis Report integrates findings from all six IPCC assessment cycles, revealing that global surface temperature has increased 1.1°C from 1850-1900 to 2011-2020, with human activities unequivocally identified as the cause of this warming. Alarmingly, current policies put the world on track for 3.2°C warming by 2100. The document identifies 5 key climate system "tipping points" approaching and emphasizes that greenhouse gas emissions must decline 43% by 2030 (compared to 2019 levels) to limit warming to 1.5°C. Climate-related risks will escalate with every increment of warming, with loss and damage disproportionately affecting vulnerable populations. Despite some progress, climate adaptation remains uneven with significant gaps, and financial flows continue to fall below levels needed for mitigation goals.",
            "What were the key findings of the latest IPCC climate report?"],
            "turns_to_mask": [0, 0],
            "reward_category": "rag"
            }

  • rai: categoría de recompensa para tareas que requieren la aplicación de principios de IA responsable, como equidad, transparencia y ética. Se utiliza para evaluar posibles sesgos en los sistemas de IA, garantizar aspectos de privacidad, abordar dilemas éticos y promover principios de diseño inclusivos.

    A continuación se muestra un ejemplo:

    {
            "turns": ["Identify potential bias concerns when developing a loan approval algorithm and suggest mitigation strategies"],
            "turns_to_mask": [0],
            "reward_category": "rai"
            }
Enmascaramiento de turnos

En el entrenamiento de conjuntos de datos, el parámetro turns_to_mask es crucial para controlar qué turnos de conversación reciben actualizaciones graduales durante el entrenamiento. Esta matriz de índices determina qué partes del diálogo debe aprender a generar el modelo y qué partes deben tratarse únicamente como contexto. El enmascaramiento adecuado garantiza que el modelo aprenda patrones de respuesta adecuados, al tiempo que evita entrenamiento en peticiones del sistema o entradas del usuario que podrían degradar el rendimiento.

Se recomienda seguir estas instrucciones de enmascaramiento:

  • Enmascare siempre el índice 0: las peticiones del sistema nunca deberían recibir actualizaciones de gradiente.

  • Enmascare siempre los turnos del usuario: evite que el modelo aprenda a generar entradas de usuario.

  • Coherencia de patrones: utilice patrones de enmascaramiento idénticos para estructuras de conversación similares, como (0, 1, 3, 5) para diálogos de varios turnos.

  • Entrenamiento selectivo: enmascare las respuestas iniciales del bot para centrar el entrenamiento en mejorar las respuestas finales.

  • Preservación de la cadena de pensamiento: enmascare los giros del sistema y del usuario únicamente cuando entrene secuencias de razonamiento.

  • Filtrado de calidad: enmascare las respuestas de asistentes de baja calidad para evitar la degradación del rendimiento.

  • Optimización del contexto: asegúrese de que los turnos enmascarados no eliminen contexto esencial necesario para respuestas posteriores.

La clave de un enmascaramiento efectivo es supervisar métricas de entrenamiento y rendimiento de validación para identificar si la estrategia de enmascaramiento preserva el contexto necesario y, al mismo tiempo, centra las actualizaciones de gradiente en los resultados deseados del modelo.

Habilitación de pérdida de divergencia de KL

Para habilitar pérdida de divergencia de KL, el servidor de anclaje debe estar habilitado para calcular la divergencia entre la política actual y la distribución original. Es necesario especificar el tipo de pérdida de KL y los coeficientes deben tener un valor distinto de cero. Los valores de coeficientes más altos ayudan al modelo a no desviarse mucho de la política original, lo que se traduce en menos cambios en el rendimiento general. Los valores de coeficientes más bajos permiten desviaciones mayores respecto a la política anterior, lo que se traduce en un mejor rendimiento de las métricas objetivo, pero repercute en el rendimiento general.

ppo_anchor:
  max_length: 8192
  trainer:
    num_nodes: ${recipes.run.cm_replicas}
  model:
    global_batch_size: 32
    
ppo_actor_train:
  model:
    ######## Use KL in actor loss ########
    kl_loss_type: low_var_kl 
    kl_loss_coeff: 0.1 

    ######## Use KL in reward model ######
    kl_reward_penalty_coeff: 0.1
Tasa de aprendizaje

Se puede ajustar la tasa de aprendizaje de los modelos crítico y de políticas, y la opción equilibrada predeterminada es 3e-6. Las tasas de aprendizaje más altas suelen provocar inestabilidad en el entrenamiento, algo que puede identificarse a través de picos de divergencia de KL y del comportamiento errático de las políticas. Las tasas de aprendizaje más bajas pueden producir problemas de convergencia y aprendizaje lento, que se aprecian en el estancamiento de recompensas y en la reducción de las actualizaciones de las políticas. La supervisión regular de divergencia de KL, puntuación de recompensa y pérdida de valor ayuda a determinar si se debe ajustar la tasa de aprendizaje durante el entrenamiento.

ppo_critic:
  model:
    optim:
      lr: 3e-6

ppo_actor_train:
  model:
    optim:
      lr: 3e-06
Tamaño de lote global

El tamaño de lote global afecta significativamente al rendimiento de PPO en Amazon Nova, ya que los lotes más grandes suelen mejorar la estabilidad del entrenamiento y la estimación del gradiente, al tiempo que permiten un procesamiento en paralelo más eficiente. Sin embargo, los lotes muy grandes pueden provocar una disminución de los retornos y pueden verse limitados por la memoria disponible, lo que requiere un equilibrio cuidadoso con la tasa de aprendizaje y otros hiperparámetros.

ppo_actor_train:
  model:
    global_batch_size: 160

Parámetros de Amazon Nova disponibles para el ajuste con PPO:

  • Ejecutar configuración

    • actor_train_replicas: número de instancias de cómputo que se utilizarán para el modelo de entrenamiento de actores. Los valores disponibles varían en función del modelo elegido. Amazon Nova Micro admite 1 o 2 réplicas. Amazon Nova Lite admite 1, 2 o 4 réplicas. Amazon Nova Pro admite 3, 6 o 12 réplicas.

    • rm_replicas: número de instancias de cómputo utilizadas para el modelo de recompensa. Recomendamos utilizar una réplica para cualquier tamaño de modelo.

    • cm_replicas: número de instancias de cómputo utilizadas para el modelo crítico. Recomendamos utilizar una réplica para cualquier tamaño de modelo.

    • actor_generation_replicas: número de instancias de cómputo utilizadas para la generación de actores. Los valores disponibles varían en función del modelo elegido. Amazon Nova Micro admite 1 réplica. Amazon Nova Lite admite 1 o 2 réplicas. Amazon Nova Pro admite 1 o 2 réplicas.

    • am_replicas: número de instancias de cómputo utilizadas para el modelo de anclaje. Recomendamos utilizar una réplica para cualquier tamaño de modelo.

  • Configuración del entrenamiento de actores (configuración de políticas)

    • max_steps: número máximo de pasos para refinar o entrenar el modelo de entrenamiento de actores. En este caso, un paso se define como despliegue, seguido del entrenamiento del modelo de entrenamiento de actores con el número global_batch_size de muestras. Una época se define como global_batch_size * trajectory_buffer_scale.

      El valor elegido aquí variará en función del caso de uso y de la complejidad del conjunto de datos. Recomendamos empezar con 65 épocas o 520 pasos, que es el número de épocas multiplicado por el valor de trajectory_buffer_scale. Sin embargo, algunas tareas requieren un tiempo de entrenamiento de PPO más prolongado para lograr el mismo rendimiento.

      En el caso de PPO, las métricas de entrenamiento, como la saturación de la puntuación del modelo de recompensa y la duración media de las acciones de la consola ml-flow, pueden ayudar a identificar los puntos óptimos para la evaluación.

    • actor_model_max_length: longitud máxima de los datos de entrada que se envían al componente de generación de actores para generar finalizaciones.

    • reward_model_max_length: longitud máxima de los datos de entrada que se envían al servidor de recompensas para puntuar finalizaciones.

    • trajectory_buffer_scale: este búfer representa el número de despliegues generados con el antiguo modelo de entrenamiento de actores (políticas) antes de actualizar las ponderaciones y generar los nuevos despliegues. Los valores admitidos son 1, 2, 4, 8 y 16.

      Si trajectory_buffer_scale es 1, el entrenamiento se basa en política. Esto significa que los despliegues se generan con las ponderaciones de modelo más actualizadas, pero el rendimiento se ve afectado. Si es 16, significa que el modelo se aparta ligeramente de la política, pero el rendimiento es mayor. Le recomendamos empezar con 8 para cada modelo.

    • kl_reward_penalty_coeff: es el término de divergencia de KL que garantiza que las actualizaciones no sean demasiado drásticas y que la política no se base en el modelo base o SFT.

    • kl_loss_coeff: este valor controla en qué medida la penalización por divergencia de KL influye en el objetivo de entrenamiento general en PPO.

    • kl_loss_type: este valor especifica cómo calcular la divergencia entre distribuciones políticas actuales y de referencia. Los kl_loss_types disponibles son kl (divergencia de KL estándar), mse (error cuadrático medio), abs (diferencia absoluta entre probabilidades logarítmicas) y low_var_kl (aproximación de KL de baja varianza).

    • model.clip_ratio: proporción de clips de actores (ε) en PPO es un hiperparámetro que limita los cambios que puede experimentar la política en cada actualización.

    • model.optim.lr: tasa de aprendizaje utilizada para entrenamiento de pérdida del modelo sustituto en el modelo de actores.

    • model.lam: forma parte del proceso de estimación de ventajas. Una λ más alta otorga más ponderación a las recompensas a largo plazo, pero con mayor varianza, mientras que una λ más baja se centra más en las recompensas inmediatas con menor varianza pero con más sesgo.

    • model.ent_coeff: la pérdida de entropía en PPO fomenta la exploración, ya que penaliza la política cuando se vuelve demasiado determinista (es decir, siempre opta por las mismas acciones con mucha confianza).

  • Configuración del modelo de recompensa

    • global_batch_size: tamaño de lote para puntuar las finalizaciones utilizando el modelo de recompensa. Si ppo_actor_train.model.global_batch_size es mayor que ppo_reward.model.global_batch_size, se procesan en varios lotes. Tenga en cuenta que ppo_actor_train.model.global_batch_size % ppo_reward.model.global_batch_size debe ser igual a 0.

    • max_length: longitud máxima del contexto del modelo de recompensa. Debería ser igual que ppo_actor_train.model.max_length.

  • Configuración del modelo crítico

    • global_batch_size: tamaño de lote del valor del modelo crítico. El modelo crítico proporcionará estimaciones del valor de cada token en las respuestas proporcionadas por el modelo de actores. El tamaño de lote se utiliza tanto para inferencia como para entrenamiento.

      Tenga en cuenta que ppo_actor_train.model.global_batch_size % ppo_critic.model.global_batch_size debe ser igual a 0 y ppo_actor_train.model.global_batch_size * ppo_actor_train.model.trajectory_buffer_size % ppo_critic.model.global_batch_size == 0.

    • max_length: longitud máxima del contexto del modelo crítico. Debería ser igual que ppo_actor_train.model.max_length.

    • model.optim.lr: tasa de aprendizaje utilizada para entrenamiento de pérdida del modelo sustituto en el modelo de actores.

  • Configuración del modelo de anclaje

    • global_batch_size: tamaño de lote para generar el logp del modelo SFT o de anclaje congelado. Tenga en cuenta que ppo_actor_train.model.global_batch_size % ppo_anchor.model.global_batch_size debe ser igual a 0.

    • max_length: longitud máxima del contexto del modelo de recompensa. Debería ser igual que ppo_actor_train.model.max_length.

  • Configuración del modelo de generación de actores

    • actor_model_max_length: longitud máxima del contexto del componente de generación del modelo de actores. Debería ser igual que ppo_actor_train.model.max_length.

Fórmula de PPO

A continuación se ofrece una fórmula para PPO.

## Run config
run:
  name: ndry-ppo-pro
  model_type: amazon.nova-pro-v1:0:300k
  model_name_or_path: nova-pro/prod
  data_s3_path: s3://testing/train.jsonl # Your training data S3 path

  actor_train_replicas: 6 # Actor train model replicas
  rm_replicas: 1 # Reward model replicas
  cm_replicas: 1 # Critic model replicas
  actor_generation_replicas: 2 # Actor generation model replicas
  am_replicas: 1 # Anchor model replicas

## Training config for each PPO component
ppo_reward:
  max_length: 8192 # model architecture max length
  trainer:
    num_nodes: ${recipes.run.rm_replicas}
  model:
    global_batch_size: 16

ppo_critic:
  max_length: 8192
  trainer:
    num_nodes: ${recipes.run.cm_replicas}
  model:
    global_batch_size: 16
    optim:
      lr: 3e-6
      name: distributed_fused_adam
      adam_w_mode: true
      eps: 1e-06
      weight_decay: 0.0
      betas:
        - 0.9
        - 0.999

ppo_anchor:
  max_length: 8192
  trainer:
    num_nodes: ${recipes.run.am_replicas}
  model:
    global_batch_size: 16

ppo_actor_generation:
  actor_model_max_length: 8192
  trainer:
    num_nodes: ${recipes.run.actor_generation_replicas}

ppo_actor_train:
  max_length: 8192
  max_steps: 520 # Stopping criteria Desired epoch num * trajectory_buffer_scale
  actor_model_max_length: 8192 # truncate input data to max length
  reward_model_max_length: 8192 # truncate input data to max length
  trajectory_buffer_scale: 8
  trainer:
    num_nodes: ${recipes.run.actor_train_replicas}
  model:
    global_batch_size: 160
    ent_coeff: 0
    clip_ratio: 0.2
    lam: 1
    kl_loss_coeff: 0.0
    kl_loss_type: low_var_kl
    kl_reward_penalty_coeff: 0.0
    hidden_dropout: 0.0 # Dropout probability for hidden state transformer.
    attention_dropout: 0.0 # Dropout probability in the attention layer.
    ffn_dropout: 0.0 # Dropout probability in the feed-forward layer.
    optim:
      lr: 3e-06
      name: distributed_fused_adam # only this one is available for p0.
      adam_w_mode: true
      eps: 1e-08
      weight_decay: 0.0
      betas:
        - 0.9
        - 0.999
Limitaciones

PPO presenta las siguientes limitaciones:

  • Los puntos de comprobación intermedios no se guardan para su evaluación y no se pueden reanudar desde un punto de comprobación intermedio. Solo se guarda el último punto de comprobación.

  • No admite conjuntos de datos multimodales.

  • Los trabajos de entrenamiento no se detienen automáticamente. Debe detener el trabajo utilizando la CLI de HyperPod.

  • No se admite el registro de MLflow.

  • TensorBoard no admite métricas de entrenamiento crítico.

  • Para ajustar los hiperparámetros, siga las instrucciones de Selección de hiperparámetros.