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

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

トレーニングデータは JSONL 形式にする必要があります。各行には、トレーニング例を表す 1 つの JSON オブジェクトが含まれます。以下がその例です。

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

以下はレコードの例です。

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

報酬モデリングフレームワークは、堅牢なモデル収束を容易にするために、異なるカテゴリ別目標にわたって多次元最適化を実装します。報酬カテゴリは、モデルを最適化する必要があるタスクに基づいて選択する必要があります。

タスクに適したフレームワークを選択するには、次のガイドラインをお勧めします。

マスキングターン

トレーニングデータセットでは、トレーニング中に勾配更新を受信する会話ターンを制御するうえで turns_to_maskパラメータが重要となります。このインデックスの配列は、生成する対話のどの部分をモデルが学習すべきか、どの部分をコンテキストとしてのみ扱うべきかを決定します。適切なマスキングにより、モデルはパフォーマンスを低下させる可能性のあるシステムプロンプトやユーザー入力のトレーニングを回避しながら、適切なレスポンスパターンを学習できます。

マスキングには次のガイダンスをお勧めします。

効果的なマスキングの鍵は、トレーニングメトリクスと検証パフォーマンスをモニタリングして、目的のモデル出力に勾配の更新を集中させながら、マスキング戦略が必要なコンテキストを保持するかどうかを特定することです。

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

学習率

Critic モデルとポリシーモデルの学習率は調整でき、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 パラメータは次のとおりです。

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 には次の制約事項があります。