근접 정책 최적화(PPO)

근접 정책 최적화(PPO)는 여러 기계 학습 모델을 활용해 하나의 모델을 학습하고 점수를 매기는 과정입니다. 다음 모델은 PPO 프로세스의 일부입니다.

훈련 데이터는 JSONL 형식이어야 하며, 각 줄에 훈련 예시를 나타내는 단일 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

학습률

비평가 및 정책 모델의 학습률을 조정할 수 있으며, 기본값은 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에는 다음과 같은 제한 사항이 있습니다.