近端政策最佳化 (PPO) - Amazon SageMaker AI

本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。

近端政策最佳化 (PPO)

近端政策最佳化 (PPO) 是使用數個機器學習模型來訓練和評分模型的程序。下列模型是 PPO 程序的一部分:

  • 演員訓練或政策模型:監督式微調 (SFT) 模型,可微調和更新每個 epoch。更新是透過取樣提示、產生完成,以及使用剪輯代理目標更新權重來進行。這會限制每個金鑰的日誌獲利性變更,讓每個政策步驟都接近前一個政策步驟,以保持訓練穩定性。

  • 演員產生模型:產生提示完成或回應的模型,由獎勵模型和評論模型判斷。此模型的權重會從每個 epoch 的演員訓練或政策模型進行更新。

  • 獎勵模型:具有凍結權重的模型,用於對演員產生模型進行評分。

  • 關鍵模型:具有未凍結權重的模型,用於對演員產生模型進行評分。此分數通常被視為演員產生剩餘權杖時所收到總獎勵的預估值。

  • 錨點模型:具有凍結權重的 SFT 模型,用於計算演員訓練模型和基本模型之間的 KL 差異。錨點模型可確保演員模型的更新不會比基本模型過於劇烈。劇烈變更可能會導致不穩定或效能降低。

訓練資料必須是 JSONL 格式,其中每一行都包含代表訓練範例的單一 JSON 物件。請見此處範例:

{ "turns": ["string", "string", ...], // Required "turns_to_mask": [integer, integer, ...], // Required "reward_category": "string", // Required "meta_data": {} // Optional }
  • turns 是代表對話方塊序列的對話字串陣列陣列。此行包含系統提示、使用者訊息和機器人回應。使用者訊息通常會以「Bot:」結尾,以指出模型輸出的開始位置。例如 [["System prompt"], ["User: Question Bot:"], ["Bot response"]]

  • turns_to_mask 是以 0 為基礎的索引陣列,可識別哪些轉彎不應接收漸層更新。遮罩轉彎通常是系統提示和使用者轉彎。例如, 會[0, 1, 3]遮罩系統提示和使用者訊息 (第一個和第三個訊息)。

  • reward_category 是一個字串,可識別要評估的模型效能面向。用於在訓練期間選取適當的獎勵模型類別。獎勵類別適用於下列獎勵類別:defaultmathcodingifrag、 和 rai

  • meta_data 是選用的物件,其中包含其他內容或 Ground-truth 資訊。這可能包括識別符、來源資訊或對話內容。此結構可根據您的資料集需求進行彈性調整。

以下是範例記錄:

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

獎勵建模架構會跨不同的分類目標實作多維度最佳化,以促進強大的模型收斂。獎勵類別應根據模型必須最佳化的任務進行選取。

針對為您的任務選取正確的架構,我們建議您遵循下列準則:

  • default:標準對話任務和基本互動的一般用途最佳化工具。用於一般對話和討論、基本撰寫任務、簡單問題回答和非專業知識查詢。

    請見此處範例:

    { "turns": ["Write a summary of climate change"], "turns_to_mask": [0], "reward_category": "default" }
  • math:專用於數學運算和數值推理任務的最佳化工具。用於數學問題解決、算術計算、代數方程式、幾何問題和統計分析。

    請見此處範例:

    { "turns": ["Calculate the derivative of x²"], "turns_to_mask": [0], "reward_category": "math" }
  • coding:程式設計和軟體開發相關查詢的專用類別。用於程式碼實作、除錯協助、演算法設計、技術文件和系統架構問題。

    請見此處範例:

    { "turns": ["Write a function to check if a string is palindrome"], "turns_to_mask": [0], "reward_category": "coding" }
  • if:需要精確程序執行和step-by-step指導的任務類別。用於多步驟程序、循序指示、複雜任務分解和程序文件。

    請見此處範例:

    { "turns": ["Provide steps to deploy a web application"], "turns_to_mask": [0], "reward_category": "if" }
  • rag:特別需要根據擷取的內容資訊回應查詢之任務的獎勵類別。當回應應直接從提供的參考資料衍生時使用,合成事實內容而不超出擷取資訊的範圍,確保答案以提供的內容為基礎,而不是一般知識。

    請見此處範例:

    { "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:需要套用負責任 AI 原則的任務的獎勵類別,例如公平性、透明度和道德。用於評估 AI 系統中的潛在偏差、確保隱私權考量、解決道德困境,以及提升包容性設計原則。

    請見此處範例:

    { "turns": ["Identify potential bias concerns when developing a loan approval algorithm and suggest mitigation strategies"], "turns_to_mask": [0], "reward_category": "rai" }
遮罩轉彎

在訓練資料集中, turns_to_mask 參數對於控制哪些對話輪流在訓練期間接收漸層更新至關重要。此索引陣列會決定模型應學習產生對話的哪些部分,以及應僅將哪些部分視為內容。適當的遮罩可確保模型學習適當的回應模式,同時避免對可能降低效能的系統提示或使用者輸入進行訓練。

我們建議以下遮罩指引:

  • 一律遮罩索引 0 - 系統提示不應接收漸層更新。

  • 一律遮罩使用者轉彎 - 防止模型學習以產生使用者輸入。

  • 模式一致性 - 針對類似的對話結構使用相同的遮罩模式,例如 (0、1、3、5) 用於多迴轉對話。

  • 選擇性訓練 - 遮罩早期機器人回應,以專注於改善最終回應的訓練。

  • Chain-of-thought保留 - 只有在訓練推理序列時,才遮罩系統和使用者。

  • 品質篩選 - 遮罩低品質助理回應,以防止效能降低。

  • 內容最佳化 - 確保遮罩轉彎不會移除後續回應所需的基本內容。

有效遮罩的關鍵是監控訓練指標和驗證效能,以識別您的遮罩策略是否保留必要的內容,同時將梯度更新集中在所需的模型輸出上。

啟用 KL-divergence 遺失

若要啟用 KL-divergence 遺失,需要啟用錨點伺服器,才能計算目前政策與原始分佈的差異。必須指定 KL 損失類型,且係數必須是零以外的值。較高的係數值有助於模型與原始政策的差異不大,這會導致對一般效能的較少變更。較低的係數值允許與先前政策的較大偏差,進而提高目標指標的效能,但會影響一般效能。

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
學習速率

評論和政策模型的學習率可以調整,3e-6 是預設的平衡選擇。較高的學習率通常會導致訓練不穩定性,這可以透過 KL 差異尖峰和不規律的政策行為來識別。較低的學習率可能會導致收斂問題和緩慢學習,以停滯的獎勵和最少的政策更新表示。定期監控 KL 分歧、獎勵分數和價值損失,有助於判斷是否在訓練期間調整學習率。

ppo_critic: model: optim: lr: 3e-6 ppo_actor_train: model: optim: lr: 3e-06
全域批次大小

全域批次大小會大幅影響 Amazon Nova 中的 PPO 效能,較大的批次通常會改善訓練穩定性和梯度估算,同時實現更有效率的平行處理。不過,批次大小非常大可能會導致傳回減少,並可能受到可用記憶體的限制,因此需要謹慎平衡學習率和其他超參數。

ppo_actor_train: model: global_batch_size: 160

可使用 PPO 進行調校的 Amazon Nova 參數包括:

  • 執行組態

    • actor_train_replicas:要用於演員訓練模型的運算執行個體數目。可用的值會根據選擇的模型而有所不同。Amazon Nova Micro 支援 1 或 2 個複本。Amazon Nova Lite 支援 1、2 或 4 個複本。Amazon Nova Pro 支援 3、6 或 12 個複本。

    • rm_replicas:用於獎勵模型的運算執行個體數目。建議您針對任何模型大小使用一個複本。

    • cm_replicas:用於評論模型的運算執行個體數目。建議您針對任何模型大小使用一個複本。

    • actor_generation_replicas:用於產生演員的運算執行個體數量。可用的值會根據選擇的模型而有所不同。Amazon Nova Micro 支援 1 個複本。Amazon Nova Lite 支援 1 或 2 個複本。Amazon Nova Pro 支援 1 或 2 個複本。

    • am_replicas:用於錨點模型的運算執行個體數目。建議您針對任何模型大小使用一個複本。

  • 演員訓練組態 (政策組態)

    • max_steps:微調或訓練演員訓練模型的步驟數目上限。在這裡,一個步驟定義為推展,接著使用範例global_batch_size數目訓練演員訓練模型。一個 epoch 定義為 global_batch_size * trajectory_buffer_scale

      此處選擇的值會根據您的使用案例和資料集複雜性而有所不同。我們建議從 65 個 epoch 或 520 個步驟開始,這是 epoch 的數量乘以 的值trajectory_buffer_scale。不過,某些任務需要較長的 PPO 訓練時間才能達到相同的效能。

      對於 PPO,訓練指標,例如飽和獎勵模型分數和 ml-flow 主控台的平均動作長度,有助於識別評估的最佳點。

    • actor_model_max_length:傳送至演員產生元件以產生完成的輸入資料長度上限。

    • reward_model_max_length:傳送至獎勵伺服器以評分完成的輸入資料長度上限。

    • trajectory_buffer_scale:此緩衝區代表在更新權重和產生新的推展之前,使用舊演員訓練 (政策) 模型產生的推展次數。支援的值為 1、2、4、8 和 16。

      如果 trajectory_buffer_scale是 1,則表示訓練為政策。這表示使用最新的模型權重產生推展,但輸送量會受到影響。如果為 16,則模型稍微偏離政策,但輸送量較高。我們建議每個模型從 8 開始。

    • kl_reward_penalty_coeff:這是 KL 分歧術語,可確保更新不會過於劇烈,並且政策不會從基礎或 SFT 模型起草。

    • kl_loss_coeff:此值控制 KL 分歧懲罰對 PPO 中整體訓練目標的影響程度。

    • kl_loss_type:此值指定如何計算目前和參考政策分佈之間的差異。kl_loss_types 可用的是 kl(標準 KL 差異)、 mse (均方誤差)、 abs(日誌機率之間的絕對差異) 和 low_var_kl(低變異 KL 近似值)。

    • model.clip_ratio:PPO 中的演員剪輯比例 (ε) 是一種超參數,會限制政策在每次更新期間可以變更的程度。

    • model.optim.lr:用於演員模型中替代模型損失訓練的學習率。

    • model.lam:優勢估計程序的一部分。較高的 λ 為長期獎勵提供更多權重,但變異數較高,而較低的 λ 更專注於變異數較低但偏差較多的立即獎勵。

    • model.ent_coeff:在 PPO 中隱含損失,透過在政策過於確定時懲罰政策來鼓勵探勘 (也就是說,始終選擇具有高可信度的相同動作)。

  • 獎勵模型組態

    • global_batch_size:使用獎勵模型評分完成的批次大小。如果 ppo_actor_train.model.global_batch_size 大於 ppo_reward.model.global_batch_size,則會分多個批次處理它們。請注意, ppo_actor_train.model.global_batch_size % ppo_reward.model.global_batch_size 必須等於 0。

    • max_length:獎勵模型的最大內容長度。這應該與 相同ppo_actor_train.model.max_length

  • 關鍵模型組態

    • global_batch_size:批評模型值的批次大小。批評模型會在演員模型提供的回應中提供每個權杖的值估計。批次大小用於推論和訓練。

      請注意, ppo_actor_train.model.global_batch_size % ppo_critic.model.global_batch_size 必須等於 0 和 ppo_actor_train.model.global_batch_size * ppo_actor_train.model.trajectory_buffer_size % ppo_critic.model.global_batch_size == 0

    • max_length:評論模型的最大內容長度。這應該與 相同ppo_actor_train.model.max_length

    • model.optim.lr:用於演員模型中替代模型損失訓練的學習率。

  • 錨點模型組態

    • global_batch_size:產生凍結 SFT 或錨點模型之 logp 的批次大小。請注意, ppo_actor_train.model.global_batch_size % ppo_anchor.model.global_batch_size 必須等於 0。

    • max_length:獎勵模型的最大內容長度。這應該與 相同ppo_actor_train.model.max_length

  • 演員產生模型組態

    • actor_model_max_length:演員模型產生元件的最大內容長度。這應該與 相同ppo_actor_train.model.max_length

PPO 配方

以下是 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
限制

PPO 有下列限制:

  • 中繼檢查點不會儲存以供評估,而且您無法從中繼檢查點繼續。只會儲存最後一個檢查點。

  • 不支援多模態資料集。

  • 訓練任務不會自動停止。您必須使用 HyperPod CLI 停止任務。

  • 不支援 MLflow 記錄。

  • TensorBoard 不支援關鍵訓練指標。

  • 若要調整超參數,請遵循選取超參數中的指引。