Optimisation des politiques proximales (PPO) - Amazon SageMaker AI

Optimisation des politiques proximales (PPO)

L’optimisation des politiques proximales (PPO) est le processus qui consiste à utiliser plusieurs modèles de machine learning pour entraîner et évaluer un modèle. Les modèles suivants font partie du processus PPO :

  • Modèle de politique ou d’entraînement acteur : un modèle de peaufinage supervisé (SFT) qui est optimisé et mis à jour à chaque époque. Les mises à jour sont effectuées en échantillonnant des invites, en générant les réponses du modèle (complétions) et en mettant à jour les poids à l’aide d’un objectif de substitution tronqué. Cela limite l’évolution de la rentabilité logarithmique par jeton, de sorte que chaque étape de la politique est proche de la précédente, préservant ainsi la stabilité de l’entraînement.

  • Modèle de génération acteur : modèle qui génère des complétions d’invites ou des réponses à évaluer selon le modèle de récompense et le modèle critique. Les poids de ce modèle sont mis à jour à partir d’un entraînement de modèle acteur ou d’un modèle de politique à chaque époque.

  • Modèle de récompense : modèle avec des poids figés, utilisé pour évaluer le modèle de génération acteur.

  • Modèle critique : modèle avec des poids figés, utilisé pour évaluer le modèle de génération acteur. Ce score est souvent considéré comme une estimation de la récompense totale que l’acteur reçoit lorsqu’il génère les jetons restants.

  • Modèle d’ancrage : modèle de SFT avec des poids figés, utilisé pour calculer la divergence KL entre le modèle d’entraînement acteur et le modèle de base. Le modèle d’ancrage garantit que les mises à jour du modèle acteur ne sont pas trop drastiques par rapport au modèle de base. Des changements radicaux peuvent entraîner une instabilité ou une dégradation des performances.

Les données d’entraînement doivent être au format JSONL, où chaque ligne contient un seul objet JSON représentant un exemple d’entraînement. Voici un exemple :

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

  • turns est un tableau de chaînes de conversation représentant la séquence du dialogue. Cette ligne contient les invites du système, les messages des utilisateurs et les réponses des bots. Les messages utilisateur se terminent généralement par « Bot: » pour indiquer où commence la sortie du modèle. Par exemple, [["System prompt"], ["User: Question Bot:"], ["Bot response"]].

  • turns_to_mask est un tableau d’indices basés sur 0 qui identifie les tours de dialogue qui ne doivent pas recevoir de mise à jour du gradient. Les tours masqués sont généralement les invites du système et les tours de l’utilisateur. Par exemple, [0, 1, 3] masque l’invite du système et les messages de l’utilisateur (les premier et troisième messages).

  • reward_category est une chaîne qui identifie les aspects des performances du modèle à évaluer. Elle est utilisée pour sélectionner la catégorie de modèle de récompense appropriée pendant l’entraînement. La catégorie de récompense est disponible pour les catégories de récompenses suivantes : default, math, coding, if, rag et rai.

  • meta_data est un objet facultatif qui contient des informations contextuelles ou factuelles supplémentaires. Cela peut inclure des identifiants, des informations sur la source ou le contexte de la conversation. La structure est flexible en fonction des besoins de votre jeu de données.

Voici un exemple d’enregistrement :

{
    "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..."
            }
        ]
    }
}

Le cadre de modélisation des récompenses met en œuvre une optimisation multidimensionnelle selon des objectifs catégoriels distincts afin de faciliter une convergence robuste des modèles. La catégorie de récompense doit être sélectionnée en fonction de la tâche pour laquelle le modèle doit être optimisé.

Nous vous recommandons de suivre les directives suivantes pour sélectionner le cadre adapté à vos tâches :

  • default : optimiseur à usage général pour les tâches conversationnelles standard et les interactions de base. Utilisé pour les conversations et discussions générales, les tâches d’écriture de base, les réponses simples à des questions et les requêtes de connaissances non spécialisées.

    Voici un exemple :

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

  • math : optimiseur spécialisé pour les calculs mathématiques et les tâches de raisonnement numérique. Utilisé pour la résolution de problèmes mathématiques, les calculs arithmétiques, les équations algébriques, les problèmes géométriques et l’analyse statistique.

    Voici un exemple :

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

  • coding : catégorie dédiée aux requêtes liées à la programmation et au développement de logiciels. Utilisée pour l’implémentation de code, l’assistance au débogage, la conception d’algorithmes, la documentation technique et les questions d’architecture système.

    Voici un exemple :

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

  • if : catégorie pour les tâches qui nécessitent une exécution procédurale précise et des instructions détaillées. Utilisée pour les procédures en plusieurs étapes, les instructions séquentielles, la décomposition de tâches complexes et la documentation des processus.

    Voici un exemple :

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

  • rag : catégorie de récompense pour les tâches qui nécessitent de répondre à des requêtes basées spécifiquement sur les informations contextuelles extraites. Utilisée lorsque les réponses doivent provenir directement des documents de référence fournis, en synthétisant le contenu factuel sans dépasser le cadre des informations extraites, en veillant à ce que les réponses soient fondées sur le contexte fourni plutôt que sur des connaissances générales.

    Voici un exemple :

    {
            "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 : catégorie de récompense pour les tâches qui nécessitent l’application des principes de l’IA responsable tels que l’équité, la transparence et l’éthique. Utilisée pour évaluer les biais potentiels dans les systèmes d’IA, garantir les considérations de confidentialité, résoudre les dilemmes éthiques et promouvoir les principes de conception inclusifs.

    Voici un exemple :

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

Dans les jeux de données d’entraînement, le paramètre turns_to_mask est crucial pour contrôler les tours de conversation qui reçoivent des mises à jour du gradient pendant l’entraînement. Ce tableau d’indices détermine les parties du dialogue que le modèle doit apprendre à générer par rapport aux parties qui doivent être traitées uniquement en tant que contexte. Un masquage approprié garantit que le modèle apprend les modèles de réponse appropriés tout en évitant de s’entraîner sur des invites système ou des entrées utilisateur susceptibles de dégrader les performances.

Nous vous recommandons la procédure de masquage suivante :

  • Indice de masque toujours à 0 : les invites système ne devraient jamais recevoir de mises à jour du gradient.

  • Toujours masquer les tours utilisateur : empêchez le modèle d’apprendre à générer des entrées utilisateur.

  • Cohérence des modèles : utilisez des modèles de masquage identiques pour des structures de conversation similaires, telles que (0, 1, 3, 5) pour les dialogues à plusieurs tours.

  • Entraînement sélectif : masquez les premières réponses des bots pour concentrer l’entraînement sur l’amélioration des réponses finales.

  • Préservation de la chaîne de pensée : ne masquez les tours système et utilisateur que lors d’un entraînement sur les séquences de raisonnement.

  • Filtrage de qualité : masquez les réponses de faible qualité des assistants pour éviter toute dégradation des performances.

  • Optimisation du contexte : assurez-vous que les tours masqués ne suppriment pas le contexte essentiel nécessaire aux réponses suivantes.

La clé d’un masquage efficace consiste à surveiller les métriques d’entraînement et les performances de validation afin de déterminer si votre stratégie de masquage préserve le contexte nécessaire tout en concentrant les mises à jour des gradients sur les résultats souhaités du modèle.

Activation de la perte de divergence KL

Pour activer la perte de divergence KL, le serveur d’ancrage doit être activé pour calculer la divergence entre la politique actuelle et la distribution d’origine. Le type de perte KL doit être spécifié et les coefficients doivent être d’une valeur autre que zéro. Des valeurs de coefficient plus élevées aident le modèle à ne pas s’écarter beaucoup de la politique d’origine, ce qui entraîne moins de moindres modifications des performances générales. Des valeurs de coefficient plus faibles permettent des écarts plus importants par rapport à la politique précédente, ce qui améliore les performances des indicateurs cibles, mais a un impact sur les performances générales.

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
Taux d’apprentissage

Le taux d’apprentissage des modèles critiques et de politiques peut être ajusté, 3e-6 étant le choix équilibré par défaut. Des taux d’apprentissage plus élevés entraînent généralement des instabilités d’entraînement, qui peuvent être identifiées par des pics de divergence KL et un comportement erratique des politiques. La baisse des taux d’apprentissage peut entraîner des problèmes de convergence et un ralentissement de l’apprentissage, comme en témoignent la stagnation des récompenses et des mises à jour minimes des politiques. Une surveillance régulière de la divergence KL, du score de récompense et de la perte de valeur permet de déterminer s’il convient d’ajuster le taux d’apprentissage pendant l’entraînement.

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

ppo_actor_train:
  model:
    optim:
      lr: 3e-06
Taille globale de lot

La taille globale de lot a un impact significatif sur les performances de la PPO dans Amazon Nova, les lots plus importants améliorant généralement la stabilité de l’entraînement et l’estimation du gradient tout en permettant un traitement parallèle plus efficace. Cependant, les lots de très grande taille peuvent entraîner des retours décroissants et être limités par la mémoire disponible, ce qui nécessite un équilibre prudent avec le taux d’apprentissage et d’autres hyperparamètres.

ppo_actor_train:
  model:
    global_batch_size: 160

Les paramètres Amazon Nova disponibles pour l’ajustement de la PPO sont les suivants :

  • Configuration d’exécution

    • actor_train_replicas : nombre d’instances de calcul à utiliser pour le modèle d’entraînement acteur. Les valeurs disponibles varient en fonction du modèle choisi. Amazon Nova Micro prend en charge 1 ou 2 réplicas. Amazon Nova Lite prend en charge 1, 2 ou 4 réplicas. Amazon Nova Pro prend en charge 3, 6 ou 12 réplicas.

    • rm_replicas : nombre d’instances de calcul à utiliser pour le modèle de récompense. Nous vous recommandons d’utiliser un réplica pour toutes les tailles de modèle.

    • cm_replicas : nombre d’instances de calcul à utiliser pour le modèle critique. Nous vous recommandons d’utiliser un réplica pour toutes les tailles de modèle.

    • actor_generation_replicas : nombre d’instances de calcul à utiliser pour la génération d’acteur. Les valeurs disponibles varient en fonction du modèle choisi. Amazon Nova Micro prend en charge 1 réplica. Amazon Nova Lite prend en charge 1 ou 2 réplicas. Amazon Nova Pro prend en charge 1 ou 2 réplicas.

    • am_replicas : nombre d’instances de calcul à utiliser pour le modèle d’ancrage. Nous vous recommandons d’utiliser un réplica pour toutes les tailles de modèle.

  • Configuration de l’entraînement acteur (configuration de politique)

    • max_steps : nombre maximal d’étapes pour optimiser ou entraîner le modèle d’entraînement acteur. Ici, une étape est définie comme déploiement, suivie de l’entraînement du modèle d’entraînement acteur avec le nombre d’exemples global_batch_size. Une époque est définie comme global_batch_size * trajectory_buffer_scale.

      La valeur choisie ici variera en fonction de votre cas d’utilisation et de la complexité du jeu de données. Nous vous recommandons de commencer par 65 époques ou 520 étapes, soit le nombre d’époques multiplié par la valeur de trajectory_buffer_scale. Cependant, certaines tâches nécessitent un temps d’entraînement de PPO plus long pour obtenir les mêmes performances.

      Pour une PPO, les métriques d’entraînement, telles que le score du modèle de récompense saturant et la durée moyenne des actions provenant de la console ml-flow, peuvent aider à identifier les points optimaux pour l’évaluation.

    • actor_model_max_length : longueur maximale des données d’entrée envoyées au composant de génération acteur pour générer des complétions.

    • reward_model_max_length : longueur maximale des données d’entrée envoyées au serveur de récompense pour évaluer les complétions.

    • trajectory_buffer_scale : cette zone tampon représente le nombre de déploiements générés à l’aide de l’ancien modèle d’entraînement acteur (politique) avant de mettre à jour les poids et de générer les nouveaux déploiements. Les valeurs prises en charge sont 1, 2, 4, 8 et 16.

      Si le paramètre trajectory_buffer_scale est égal à 1, l’entraînement est conforme à la politique. Cela signifie que les déploiements sont générés avec les poids de modèle les plus récents, mais le débit en pâtit. S’il est égal à 16, le modèle est légèrement hors politique, mais le débit est plus élevé. Nous recommandons de commencer par une valeur de 8 pour chaque modèle.

    • kl_reward_penalty_coeff : il s’agit la condition de divergence KL qui garantit que les mises à jour ne sont pas trop drastiques et que la politique n’est pas rédigée à partir du modèle de base ou SFT.

    • kl_loss_coeff : cette valeur contrôle dans quelle mesure la pénalité de divergence KL influence l’objectif d’entraînement global dans la PPO.

    • kl_loss_type : cette valeur indique comment calculer la divergence entre les distributions des politiques actuelles et de référence. Les kl_loss_types disponibles sont kl (divergence KL standard), mse (erreur quadratique moyenne), abs (différence absolue entre les probabilités logarithmiques) et low_var_kl (approximation KL à faible variance).

    • model.clip_ratio : le ratio de troncature de l’acteur (ε) dans une PPO est un hyperparamètre qui limite le degré de modification de la politique lors de chaque mise à jour.

    • model.optim.lr : taux d’apprentissage utilisé pour l’entraînement à la perte du modèle de substitution dans le modèle acteur.

    • model.lam : fait partie du processus d’estimation des avantages. Un λ plus élevé donne plus de poids aux récompenses à long terme, mais avec une variance plus élevée, tandis qu’un λ faible se concentre davantage sur les récompenses immédiates avec une variance plus faible, mais plus de biais.

    • model.ent_coeff : la perte d’entropie dans une PPO encourage l’exploration en pénalisant la politique lorsqu’elle devient trop déterministe (c’est-à-dire en choisissant toujours les mêmes actions avec une grande confiance).

  • Configuration du modèle de récompense

    • global_batch_size : taille du lot pour évaluer les complétions à l’aide du modèle de récompense. Si ppo_actor_train.model.global_batch_size est supérieur à ppo_reward.model.global_batch_size, elles sont traitées en plusieurs lots. Notez que ppo_actor_train.model.global_batch_size % ppo_reward.model.global_batch_size doit être égal à 0.

    • max_length : longueur maximale du contexte du modèle de récompense. Devrait être la même que ppo_actor_train.model.max_length.

  • Configuration du modèle critique

    • global_batch_size : taille de lot correspondant à la valeur du modèle critique. Le modèle critique fournit des estimations de valeur pour chaque jeton dans les réponses fournies par le modèle acteur. La taille de lot est utilisée à la fois pour l’inférence et pour l’entraînement.

      Notez que ppo_actor_train.model.global_batch_size % ppo_critic.model.global_batch_size doit être égale à 0 et ppo_actor_train.model.global_batch_size * ppo_actor_train.model.trajectory_buffer_size % ppo_critic.model.global_batch_size == 0.

    • max_length : longueur maximale du contexte du modèle critique. Devrait être la même que ppo_actor_train.model.max_length.

    • model.optim.lr : taux d’apprentissage utilisé pour l’entraînement à la perte du modèle de substitution dans le modèle acteur.

  • Configuration du modèle d’ancrage

    • global_batch_size : taille du lot pour générer le journal du modèle SFT ou d’ancrage figé. Notez que ppo_actor_train.model.global_batch_size % ppo_anchor.model.global_batch_size doit être égal à 0.

    • max_length : longueur maximale du contexte du modèle de récompense. Devrait être la même que ppo_actor_train.model.max_length.

  • Configuration du modèle de génération acteur

    • actor_model_max_length : longueur de contexte maximale du composant de génération du modèle acteur. Devrait être la même que ppo_actor_train.model.max_length.

Formule de PPO

Ce qui suit est une formule de 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
Limites

La PPO comporte les limitations suivantes :

  • Les points de contrôle intermédiaires ne sont pas enregistrés pour évaluation et vous ne pouvez pas les reprendre à partir d’un point de contrôle intermédiaire. Seul le dernier point de contrôle est enregistré.

  • Les jeux de données multimodaux ne sont pas pris en charge.

  • Les tâches d’entraînement ne sont pas automatiquement interrompues. Vous devez les arrêter à l’aide de la CLI HyperPod.

  • La journalisation MLflow n’est pas prise en charge.

  • Les métriques de l’entraînement critique ne sont pas prises en charge sur TensorBoard.

  • Pour ajuster les hyperparamètres, suivez les instructions dans Selecting hyperparameters.