RLHF와 PPO의 N 구현 세부 사항 – Hugging Face 블로그 요약

Hugging Face 블로그 글은 OpenAI의 2019 RLHF 코드베이스를 재현하고, 스타일 작업에 대한 학습 곡선을 일치시키며, 재현 가능성에 영향을 미치는 구현 세부 사항을 나열합니다.

Matching Learning Curves

재현은 감정과 서술성 작업에 대해 OpenAI의 원래 코드베이스와 거의 동일한 학습 곡선을 생성합니다. 저자들은 기준 지표를 얻기 위해 AWS p3dn.24xlarge 인스턴스에서 원래의 TensorFlow 1.x 저장소를 실행했으며,その後 those metrics를 사용하여 PyTorch 기반 재현을 검증했습니다.

General Implementation Details

보상 모델과 정책 가치 헤드는 연결된 쿼리와 응답을 입력으로 받습니다. 시퀀스는 특수 패딩 토큰을 사용하여 고정 길이에 패딩되고, 너무 길면 잘립니다. 또한 로짓 계산 중 패딩 토큰에 대해 위치 인덱스가 조정됩니다. 응답 생성은 EOS 토큰에서 멈추지 않고 고정 길이 출력을 샘플링하며, 보상 모델과 정책 훈련 모두에 대해 학습률이 0으로 어닐링됩니다. 탐색을 장려하기 위해 GPU 프로세스마다 다른 랜덤 시드가 사용됩니다.

Reward Model Implementation Details

보상 모델은 연결된 쿼리‑응답 쌍의 마지막 토큰 값만 출력합니다. 보상 헤드 가중치는 variance가 1/(sqrt(d_model)+1)인 정규 분포에서 초기화되고 편향은 0으로 설정됩니다. 보상 정규화는 고정된 참조 정책에서의 보상의 경험적 평균과 표준편차를 계산하여 단위 분산과 0 평균으로 스케일링함으로써 훈련 전후에 수행됩니다.

Policy Training Implementation Details

로짓은 로그 확률을 계산하기 전에 샘플링 온도로 나뉩니다. 가치 헤드 가중치는 평균 0과 분산 0으로 초기화됩니다. 정책 훈련 중 드롭아웃은 비활성화됩니다. 거부 샘플링은 토큰 16과 24 사이에 기간을 강제하고, 그렇지 않은 경우 고정된 낮은 보상을 할당합니다. 할인 인자 γ는 1로 설정됩니다. 훈련은 그래디언트 누적을 위해 배치, 미니배치, 마이크로배치를 사용합니다. 토큰당 KL 페널티가 보상에 추가되고, 미니배치당 보상과 어드밴티지 화이트닝이 적용됩니다(어드밴티지에만 평균 이동). 가치 함수는 표준 PPO와 같이 클리핑되며, KL 페널티 계수 β는 관측된 KL 발산에 따라 적응됩니다.

PyTorch Adam Optimizer Numerical Issues w.r.t RLHF

PyTorch의 Adam 구현은 분모 내 epsilon의 위치에서 TensorFlow의 구현과 다릅니다. 이로 인해 학습 초기에 효과적인 정규화 항이 작아져 더 공격적인 그래디언트 업데이트가 발생합니다. 실험 결과, PyTorch Adam은 TensorFlow‑style Adam에 비해 로그‑확률 분산이 더 높고, 비율 극단이 더 크며, 근사 KL과 클립 분수가 증가합니다. 특히 gpt2‑xl과 같은 더 큰 모델에서 그러한 현상이 두드러집니다.

Limitations

이 작업은 높은 계산 비용과 훈련 불안정성으로 인해 CNN/DM 또는 TL;DR의 요약 작업을 재현하려고 하지 않습니다. 원래 코드는 V100 32GB GPU 8개가 필요하며, 재현은 GPU 활용률이 낮아(~30%) 실행이 비용이 많이 들고 취약해집니다.

Conclusion

OpenAI의 RLHF 코드베이스를 재현하고 그 구현 세부 사항을 문서화함으로써 저자들은 RLHF 엔지니어링 관행을 이해하기 위한 참조를 제공하고, 훈련 역학에 중대한 영향을 미칠 수 있는 미세한 옵티마이저 차이를 강조합니다.