近似ポリシー最適化 (PPO) - Amazon SageMaker AI

翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

近似ポリシー最適化 (PPO)

近似ポリシー最適化 (PPO) は、複数の機械学習モデルを使用してモデルをトレーニングおよびスコアリングするプロセスです。次のモデルは PPO プロセスの一部です。

  • アクタートレーニングまたはポリシーモデル: エポックごとに微調整および更新される教師ありファインチューニング (SFT) モデル。更新は、プロンプトをサンプリングし、補完を生成し、クリップされた代理目的を使用して重みを更新することによって行われます。これにより、トークンごとのログ収益性の変更が制限され、その結果、各ポリシーステップが前のステップに近似し、トレーニングの安定性が維持されます。

  • Actor 生成モデル: Reward モデルと Critic モデルによって判断されることになるプロンプトの補完またはレスポンスを生成するモデル。このモデルの重みは、各エポックごとに Actor トレーニングモデルまたはポリシーモデルから更新されます。

  • 報酬モデル: アクター生成モデルのスコア付けに使用されるフリーズウェイトのモデル。

  • 批判モデル: アクター生成モデルのスコア付けに使用される凍結されていない重みを持つモデル。このスコアは、多くの場合、残りのトークンの生成時にアクターが受け取る合計報酬の見積もりとして見なされます。

  • アンカーモデル: アクタートレーニングモデルとベースモデル間の KL 発散を計算するために使用されるフリーズウェイトの SFT モデル。Anchor モデルでは、Actor モデルの更新がベースモデルと比較してそれほど劇的にならないようになっています。大幅な変更は、不安定な状態やパフォーマンスの低下につながる可能性があります。

トレーニングデータは JSONL 形式である必要があります。各行には、トレーニング例を表す 1 つの 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 ベースのインデックスの配列です。マスクされたターンは通常、システムプロンプトとユーザーターンです。たとえば、 はシステムプロンプトとユーザーメッセージ (最初と 3 番目のメッセージ) を[0, 1, 3]マスクします。

  • reward_category は、評価するモデルパフォーマンスの側面を識別する文字列です。トレーニング中に適切な報酬モデルカテゴリを選択するために使用されます。報酬カテゴリは、、default、math、、codingifrag、および の報酬カテゴリで使用できますrai。

  • meta_data は、追加のコンテキスト情報またはグラウンドトゥルース情報を含むオプションのオブジェクトです。これには、識別子、ソース情報、会話コンテキストが含まれます。構造は、データセットのニーズに応じて柔軟です。

レコードの例を次に示します。

{
    "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 ダイバージェンス損失を有効にする

KL 分散損失を有効にするには、アンカーサーバーを有効にして、元のディストリビューションからの現在のポリシーの分散を計算する必要があります。KL 損失タイプを指定し、係数を 0 以外の値にする必要があります。係数値が大きいほど、モデルは元のポリシーから大きく逸脱せず、一般的なパフォーマンスへの変更が少なくなります。係数値が低いと、以前のポリシーからの偏差が大きくなり、ターゲットメトリクスのパフォーマンスは向上しますが、一般的なパフォーマンスには影響します。

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: 報酬モデルに使用されるコンピューティングインスタンスの数。任意のモデルサイズに 1 つのレプリカを使用することをお勧めします。

    • cm_replicas: 批評家モデルに使用されるコンピューティングインスタンスの数。任意のモデルサイズに 1 つのレプリカを使用することをお勧めします。

    • actor_generation_replicas: アクターの生成に使用されるコンピューティングインスタンスの数。使用可能な値は、選択したモデルによって異なります。Amazon Nova Micro は 1 つのレプリカをサポートしています。Amazon Nova Lite は 1 つまたは 2 つのレプリカをサポートしています。Amazon Nova Pro は 1 つまたは 2 つのレプリカをサポートしています。

    • am_replicas: アンカーモデルに使用されるコンピューティングインスタンスの数。任意のモデルサイズに 1 つのレプリカを使用することをお勧めします。

  • アクタートレーニング設定 (ポリシー設定）

    • max_steps: アクタートレーニングモデルを微調整またはトレーニングするステップの最大数。ここでは、1 つのステップをロールアウトとして定義し、次にサンプルglobal_batch_size数を使用してアクタートレーニングモデルをトレーニングします。1 つのエポックは として定義されますglobal_batch_size * trajectory_buffer_scale。

      ここで選択する値は、ユースケースとデータセットの複雑さによって異なります。65 エポックまたは 520 ステップから始めることをお勧めします。これは、エポックの数に の値を乗じたものです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: これは、更新がそれほど劇的ではなく、ポリシーがベースモデルまたは SFT モデルからドラフトされないようにする KL 分岐項です。

    • 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、複数のバッチで処理されます。は 0 に等しくppo_actor_train.model.global_batch_size % ppo_reward.model.global_batch_sizeなければなりません。

    • max_length: 報酬モデルの最大コンテキスト長。これは と同じである必要がありますppo_actor_train.model.max_length。

  • クリティカルモデル設定

    • global_batch_size: 批判モデル値のバッチサイズ。批評モデルは、アクターモデルによって提供されるレスポンス内の各トークンの値の見積もりを提供します。バッチサイズは、推論とトレーニングの両方に使用されます。

      は 0 と に等しくppo_actor_train.model.global_batch_size % ppo_critic.model.global_batch_sizeなければなりません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 またはアンカーモデルのログを生成するためのバッチサイズ。は 0 に等しくppo_actor_train.model.global_batch_size % ppo_anchor.model.global_batch_sizeなければなりません。

    • 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 ではサポートされていません。

  • ハイパーパラメータを調整するには、「ハイパーパラメータの選択」のガイダンスに従います。