Ottimizzazione prossimale delle policy (PPO) - Amazon SageMaker AI

Ottimizzazione prossimale delle policy (PPO)

L’ottimizzazione prossimale delle policy (PPO, Proximal Policy Optimization) è il processo di utilizzo di diversi modelli di machine learning per addestrare e assegnare un punteggio a un modello. I seguenti modelli fanno parte del processo PPO:

  • Modello di policy o addestramento attori: un modello di fine-tuning supervisionato (SFT) che viene sottoposto a fine-tuning e aggiornato a ogni epoch. Gli aggiornamenti vengono effettuati campionando prompt, generando completamenti e aggiornando i pesi utilizzando un obiettivo sostitutivo limitato. Ciò limita la modifica della redditività dei log per token in modo che ogni fase della policy sia prossimale a quella precedente, preservando la stabilità dell’addestramento.

  • Modello di generazione attori: un modello che genera completamenti o risposte a prompt da giudicare in base al modello di ricompensa e al modello di critica. I pesi di questo modello vengono aggiornati in base al modello di addestramento attori o al modello di policy di ogni epoch.

  • Modello di ricompensa: un modello con pesi ibernati che viene utilizzato per calcolare il punteggio modello di generazione attori.

  • Modello di critica: un modello con pesi non ibernati che viene utilizzato per calcolare il punteggio modello di generazione attori. Questo punteggio viene spesso visto come una stima della ricompensa totale che l’attore riceve quando genera i token rimanenti.

  • Modello di ancoraggio: un modello SFT con pesi ibernati utilizzato per calcolare la divergenza KL tra il modello di addestramento attori e il modello di base. Il modello di ancoraggio garantisce che gli aggiornamenti al modello di attore non siano troppo drastici rispetto al modello di base. Cambiamenti drastici possono portare a instabilità o a una riduzione delle prestazioni.

I dati di addestramento devono essere in formato JSONL, in cui ogni riga contiene un singolo oggetto JSON che rappresenta un esempio di addestramento. Ecco un esempio:

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

  • turns è un array di stringhe di conversazione che rappresentano la sequenza di dialogo. Questa riga contiene prompt di sistema, messaggi utente e risposte dei bot. I messaggi utente in genere terminano con “Bot:”, a indicare dove inizia l’output del modello. Ad esempio, [["System prompt"], ["User: Question Bot:"], ["Bot response"]].

  • turns_to_mask è un array di indici basati su 0 che identificano quali interazioni non devono ricevere aggiornamenti del gradiente. Le interazioni mascherate sono in genere prompt di sistema e interazioni utente. Ad esempio, [0, 1, 3] maschera il prompt di sistema e i messaggi utente (il primo e il terzo messaggio).

  • reward_category è una stringa che identifica gli aspetti delle prestazioni del modello da valutare. Viene utilizzato per selezionare la categoria del modello di ricompensa appropriata durante l’addestramento. Sono disponibili le seguenti categorie di ricompense: default, math, coding, if, rag e rai.

  • meta_data è un oggetto opzionale che contiene informazioni contestuali o dati reali di riferimento. Ciò può includere identificatori, informazioni sulla fonte o contesto della conversazione. La struttura è flessibile in base alle esigenze del set di dati.

Di seguito è riportato un record di esempio:

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

Il framework di modellazione delle ricompense implementa l’ottimizzazione multidimensionale su obiettivi di categorie distinti per favorire una convergenza efficace dei modelli. La categoria di ricompensa deve essere selezionata in base all’attività per cui il modello deve essere ottimizzato.

È consigliabile seguire queste linee guida per scegliere il framework più adatto alle attività:

  • default: un ottimizzatore generico per attività di conversazione standard e interazioni di base. Utilizzato per conversazioni e discussioni generiche, attività di scrittura di base, risposte a domande semplici e query non specializzate.

    Ecco un esempio:

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

  • math: un ottimizzatore specializzato per calcoli matematici e attività di ragionamento numerico. Utilizzato per la risoluzione di problemi matematici, calcoli aritmetici, equazioni algebriche, problemi geometrici e analisi statistiche.

    Ecco un esempio:

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

  • coding: una categoria dedicata per le query relative alla programmazione e allo sviluppo di software. Utilizzato per attività di implementazione di codice, assistenza al debug, progettazione di algoritmi, documentazione tecnica e domande sull’architettura di sistema.

    Ecco un esempio:

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

  • if: una categoria per attività che richiedono un’esecuzione procedurale precisa e una guida dettagliata. Utilizzata per procedure in più fasi, istruzioni sequenziali, scomposizione di attività complesse e documentazione dei processi.

    Ecco un esempio:

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

  • rag: una categoria di ricompensa per le attività che richiedono la risposta a query basate specificamente sulle informazioni contestuali recuperate. Utilizzata quando le risposte devono essere derivate direttamente dai materiali di riferimento forniti, sintetizzando contenuti fattuali senza andare oltre l’ambito delle informazioni recuperate, garantendo che le risposte siano fondate sul contesto fornito anziché su conoscenze generali.

    Ecco un esempio:

    {
            "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: una categoria di ricompensa per le attività che richiedono l’applicazione di principi di IA responsabile come equità, trasparenza ed etica. Utilizzata per valutare potenziali bias nei sistemi di intelligenza artificiale, garantire considerazioni sulla privacy, affrontare dilemmi etici e promuovere principi di progettazione inclusivi.

    Ecco un esempio:

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

Nei set di dati di addestramento, il parametro turns_to_mask è fondamentale per controllare quali interazioni di conversazione ricevono aggiornamenti in base al gradiente durante l’addestramento. Questo array di indici determina quali parti del dialogo il modello dovrebbe imparare a generare e quali parti devono essere trattate solo come contesto. Un mascheramento adeguato assicura che il modello apprenda i modelli di risposta appropriati, evitando al contempo l’addestramento su prompt di sistema o input dell’utente che potrebbero compromettere le prestazioni.

Per il mascheramento è consigliabile seguire queste indicazioni:

  • Mascherare sempre l’indice 0 - I prompt di sistema non dovrebbero mai ricevere aggiornamenti del gradiente.

  • Mascherare sempre le interazioni degli utenti - Impedire al modello di imparare a generare input da parte dell’utente.

  • Coerenza dei modelli - Utilizzare schemi di mascheramento identici per strutture di conversazione simili, ad esempio (0, 1, 3, 5) per dialoghi con più interazioni.

  • Addestramento selettivo - Mascherare le risposte iniziali dei bot per concentrare l’addestramento sul miglioramento delle risposte finali.

  • Preservazione della catena di ragionamento - Mascherare solo il sistema e le interazioni degli utenti durante l’addestramento sulle sequenze di ragionamento.

  • Filtro di qualità - Mascherare le risposte degli assistenti di bassa qualità per evitare una riduzione delle prestazioni.

  • Ottimizzazione del contesto - Verificare che le interazioni mascherate non rimuovano il contesto essenziale necessario per le risposte successive.

La chiave per un mascheramento efficace è il monitoraggio delle metriche di addestramento e delle prestazioni di convalida per determinare se la strategia di mascheramento preserva il contesto necessario, concentrando al contempo gli aggiornamenti del gradiente sugli output del modello desiderati.

Abilitare la perdita di divergenza KL

Per abilitare la perdita di divergenza KL, il server di ancoraggio deve essere abilitato a calcolare la divergenza della policy corrente rispetto alla distribuzione originale. È necessario specificare il tipo di perdita KL e i coefficienti devono essere di un valore diverso da zero. Valori di coefficienti più elevati evitano che il modello si discosti troppo dalla policy originale, con minori modifiche alle prestazioni generali. Valori di coefficienti più ridotti consentono deviazioni maggiori rispetto alla policy precedente, con conseguente miglioramento delle prestazioni delle metriche di riferimento, ma con un impatto sulle prestazioni generali.

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
Velocità di apprendimento

Il tasso di apprendimento per i modelli di critica e di policy può essere modificato, con 3e-6 come scelta bilanciata predefinita. Tassi di apprendimento più elevati determinano in genere instabilità nell’addestramento, che possono essere individuate tramite picchi nella divergenza KL e comportamenti irregolari delle policy. Tassi di apprendimento ridotti possono causare problemi di convergenza e rallentamento dell’apprendimento, evidenziati da ricompense stagnanti e aggiornamenti minimi delle policy. Il monitoraggio regolare della divergenza KL, del punteggio di ricompensa e della perdita di valore aiuta a capire se modificare il tasso di apprendimento durante l’addestramento.

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

ppo_actor_train:
  model:
    optim:
      lr: 3e-06
Dimensioni batch globali

Le dimensioni batch globali influiscono in modo significativo sulle prestazioni del processo PPO in Amazon Nova, con batch di dimensioni maggiori che generalmente migliorano la stabilità dell’addestramento e la stima del gradiente, consentendo al contempo un’elaborazione parallela più efficiente. Tuttavia, batch di dimensioni eccessive possono determinare rendimenti decrescenti ed essere limitati dalla memoria disponibile, il che richiede un attento equilibrio tra il tasso di apprendimento e altri iperparametri.

ppo_actor_train:
  model:
    global_batch_size: 160

I parametri di Amazon Nova disponibili per l’ottimizzazione con PPO includono:

  • Configurazione run

    • actor_train_replicas: il numero di istanze di calcolo utilizzate per il modello di addestramento attori. I valori disponibili variano in base al modello scelto. Amazon Nova Micro supporta 1 oppure 2 repliche. Amazon Nova Lite supporta 1, 2 oppure 4 repliche. Amazon Nova Pro supporta 3, 6 oppure 12 repliche.

    • rm_replicas: il numero di istanze di calcolo utilizzate per il modello di ricompensa. È consigliabile utilizzare una sola replica per qualsiasi dimensione del modello.

    • cm_replicas: il numero di istanze di calcolo utilizzate per il modello di critica. È consigliabile utilizzare una sola replica per qualsiasi dimensione del modello.

    • actor_generation_replicas: il numero di istanze di calcolo utilizzate per la generazione attori. I valori disponibili variano in base al modello scelto. Amazon Nova Micro supporta 1 replica. Amazon Nova Lite supporta 1 oppure 2 repliche. Amazon Nova Pro supporta 1 oppure 2 repliche.

    • am_replicas: il numero di istanze di calcolo utilizzate per il modello di ancoraggio. È consigliabile utilizzare una sola replica per qualsiasi dimensione del modello.

  • Configurazione di addestramento attori (configurazione di policy)

    • max_steps: il numero massimo di fasi per eseguire il fine-tuning o addestrare il modello di addestramento attori. Qui, una fase è definita come implementazione, seguita dall’addestramento del modello di addestramento attori con un numero global_batch_size di campioni. Un’epoch è definita come global_batch_size * trajectory_buffer_scale.

      Il valore scelto qui varia in base al caso d’uso e alla complessità del set di dati. È consigliabile iniziare con 65 epoch o 520 passaggi, ovvero il numero di epoch moltiplicato per il valore di trajectory_buffer_scale. Tuttavia, alcune attività richiedono un periodo di addestramento di PPO più esteso per ottenere le stesse prestazioni.

      Per PPO, le metriche di addestramento, come la saturazione del punteggio del modello di ricompensa e la durata media delle azioni rilevate dalla console ml-flow, possono aiutare a identificare i punti ottimali per la valutazione.

    • actor_model_max_length: la lunghezza massima dei dati di input che vengono inviati al componente di generazione attori per generare completamenti.

    • reward_model_max_length: la lunghezza massima dei dati di input inviati al server di ricompensa per ottenere punteggi di completamento.

    • trajectory_buffer_scale: questo buffer rappresenta il numero di implementazioni generate utilizzando il vecchio modello di addestramento attori (policy) prima di aggiornare i pesi e generare le nuove implementazioni. I valori supportati sono 1, 2, 4, 8 e 16.

      Se trajectory_buffer_scale è 1, l’addestramento segue la policy. Significa che le implementazioni vengono generate con i pesi dei modelli più aggiornati, ma il throughput ne risente. Se è 16, il modello si discosta leggermente dalla policy, ma il throughput è maggiore. È consigliabile iniziare con 8 per ogni modello.

    • kl_reward_penalty_coeff: questo è il termine di divergenza KL che garantisce che gli aggiornamenti non siano troppo drastici e che la policy non si discosti troppo dalla base o dal modello SFT.

    • kl_loss_coeff: questo valore controlla in che misura la penalità di divergenza KL influisce sull’obiettivo generale di addestramento in PPO.

    • kl_loss_type: questo valore specifica come calcolare la divergenza tra le distribuzioni di policy attuali e quelle di riferimento. I valori di kl_loss_types disponibili sono kl (divergenza KL standard), mse (errore quadratico medio), abs (differenza assoluta tra probabilità logaritmiche) e low_var_kl (approssimazione KL a bassa varianza).

    • model.clip_ratio: il coefficiente di limitazione dell’attore (ε) in PPO è un iperparametro che limita le modifiche che la policy può subire durante ogni aggiornamento.

    • model.optim.lr: il tasso di apprendimento utilizzato per l’addestramento basato sulla perdita del modello surrogato nel modello di attore.

    • model.lam: parte del processo di stima dei vantaggi. Un valore λ più elevato conferisce più peso ai premi a lungo termine ma con varianza maggiore, mentre un valore λ ridotto si concentra maggiormente sui premi immediati con varianza inferiore, ma maggiori bias.

    • model.ent_coeff: la perdita di entropia in PPO favorisce l’esplorazione penalizzando la policy quando diventa troppo deterministica (ovvero, quando sceglie sempre le stesse azioni con elevata fiducia).

  • Configurazione del modello di ricompensa

    • global_batch_size: la dimensione del batch per calcolare il punteggio dei completamenti utilizzando il modello di ricompensa. Se ppo_actor_train.model.global_batch_size è maggiore di ppo_reward.model.global_batch_size, vengono elaborati in più batch. ppo_actor_train.model.global_batch_size % ppo_reward.model.global_batch_size deve essere uguale a 0.

    • max_length: la lunghezza massima del contesto del modello di ricompensa. Deve essere uguale a ppo_actor_train.model.max_length.

  • Configurazione del modello di critica

    • global_batch_size: la dimensione di batch del valore del modello di critica. Il modello di critica fornisce stime di valore per ogni token nelle risposte fornite dal modello di attore. La dimensione di batch viene utilizzata sia per l’inferenza che per l’addestramento.

      ppo_actor_train.model.global_batch_size % ppo_critic.model.global_batch_size deve essere uguale a 0 e ppo_actor_train.model.global_batch_size * ppo_actor_train.model.trajectory_buffer_size % ppo_critic.model.global_batch_size == 0.

    • max_length: la lunghezza massima del contesto del modello di critica. Deve essere uguale a ppo_actor_train.model.max_length.

    • model.optim.lr: il tasso di apprendimento utilizzato per l’addestramento basato sulla perdita del modello surrogato nel modello di attore.

  • Configurazione del modello di ancoraggio

    • global_batch_size: la dimensione di batch per la generazione del logp del modello SFT o di ancoraggio ibernato. ppo_actor_train.model.global_batch_size % ppo_anchor.model.global_batch_size deve essere uguale a 0.

    • max_length: la lunghezza massima del contesto del modello di ricompensa. Deve essere uguale a ppo_actor_train.model.max_length.

  • Configurazione del modello di generazione attori

    • actor_model_max_length: la lunghezza massima del contesto del componente di generazione del modello di attore. Deve essere uguale a ppo_actor_train.model.max_length.

Ricetta per PPO

Di seguito è riportata una ricetta per 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
Limitazioni

PPO presenta le seguenti limitazioni:

  • I checkpoint intermedi non vengono salvati per la valutazione e non è possibile riprendere da un checkpoint intermedio. Viene salvato solo l’ultimo checkpoint.

  • I set di dati multimodali non sono supportati.

  • I job di addestramento non vengono arrestati automaticamente. È necessario arrestare il job mediante HyperPod CLI.

  • La registrazione di log MLflow non è supportata.

  • Le metriche di addestramento del modello di critica non sono supportate su TensorBoard.

  • Per regolare gli iperparametri, segui le istruzioni riportate in Selecting hyperparameters.