GPT-OSS Agentic RL 훈련: 실용적인 회고

TL;DR

Hugging Face와 LinkedIn 연구원들은 PPO 온-정책 무결성, attention-sink 백워드 패스, 그리고 MoE 메모리 물질화에서의 치명적인 불안정성을 해결함으로써 GPT-OSS 모델에 대한 에이전트 강화 학습(RL)을 성공적으로 잠금 해제했습니다. 이러한 수정 사항은 GPT-OSS가 환경과 도구와 상호작용하는 다단계 의사결정 에이전트의 안정적인 백본으로 작동하도록 합니다.

Agentic RL과 GPT-OSS

에이전트 RL은 전통적인 단일 턴 RL과 달리 전체 의사결정 과정을 최적화합니다. 정적 응답을 생성하는 대신, 모델은 다단계 궤적에서 행동을 계획하고, 도구를 호출하며, 행동을 적응하도록 학습합니다. 이는 에이전트가 롤아웃 궤적을 수집하고, 보상을 계산하며, PPO 또는 GRPO와 같은 알고리즘을 사용하여 정책을 반복적으로 업데이트하는 폐쇄 루프 시스템이 필요함을 의미합니다.

GPT-OSS는 OpenAI o3-mini와 o4-mini와 성능이 비슷함을 보였지만, 에이전트 RL에 대한 적합성은 이전에 검증되지 않았습니다. 연구원들은 verl 훈련 프레임워크를 사용하고 GSM8K, Retool(에이전트 코딩 작업), 그리고 검증 가능한 명령 따르기 작업에서 GPT-OSS-20B 모델을 테스트했습니다.

PPO 온-정책 불안정성 해결

초기 훈련 실행에서는 보상이 증가하지 않는 가운데 KL 발산과 엔트로피가 폭발하는 현상이 나타났습니다. 팀은 Mixture of Experts(MoE) 아키텍처로 인한 Proximal Policy Optimization(PPO) 온-정책 무결성의 실패를 식별했습니다.

MoE 로그-확률 불일치

순수 온-정책 PPO에서는 중요도 샘플링 비율이 정확히 1이어야 합니다. 그러나 GPT-OSS와 같은 MoE 아키텍처에서는 게이팅 네트워크가 부동소수점 차이 또는 확률성으로 인해 롤아웃 생성을 위한 순방향 패스와 훈련을 위한 순방향 패스 사이에서 입력을 서로 다른 전문가로 라우팅할 수 있습니다. 이로 인해 현재 로그-확률과 이전 로그-확률 사이에 불일치가 발생하여 PPO 클립을 잘못 트리거하고 온-정책 가정을 위반합니다.

수정: 팀은 환경이 온-정책임을 알고 있을 때(미니배치 크기가 글로벌 배치 크기와 같을 때) 계산을 재정의하는 로그-확률 대체법을 구현하여 old_log_prob = log_prob.detach()를 설정함으로써 중요도 비율을 1로 강제로 되돌렸습니다.

Attention Sinks를 통한 훈련-추론 불일치 수정

PPO 무결성을 수정한 후에도 그래디언트 노름이 계속 폭발했습니다. 연구원들은 추론 엔진(SGLang)과 훈련 스택(FSDP와 FlashAttention-v2)이 토큰 수준 확률에서 다른 결과를 생성하는 근본적인 훈련-추론 불일치를 식별했습니다.

Attention Sinks의 역할

GPT-OSS는 소프트맥스 계산에서 "가상 토큰" 역할을 하는 학습 가능한 스칼라 파라미터인 attention sinks를 사용합니다. 이러한 sinks는 모델이 콘텐츠 토큰에 강제로 집중하는 대신 학습된 파라미터에 주의 질량을 할당하도록 허용하여 스트리밍 추론의 안정성을 향상시킵니다.

FlashAttention v3에서의 구현

팀은 verl이 attention sinks를 지원하지 않는 FlashAttention v2를 하드코딩했으며, v2와 v3 모두 싱크 그래디언트에 필요한 백워드 패스를 지원하지 않는다는 사실을 발견했습니다. 이를 해결하기 위해 그들은 다음과 같이 했습니다:

  1. vLLM FlashAttention 포크의 순방향 패스를 활용했습니다.
  2. 싱크 그래디언트 $ rac{\partial L}{\partial S_{h}}$를 계산하는 백워드 패스를 구현했습니다.

결과: 이 수정으로 GSM8K, VerifyIf, 그리고 Retool 작업에서 수렴 속도가 크게 빨라지고 보상 개선이 안정적으로 이루어졌습니다.

장기 컨텍스트를 위한 메모리 효율성 확장

에이전트 RL은 환경 피드백이 궤적에 추가됨에 따라 컨텍스트 윈도우를 확장해야 합니다. 팀은 Out-of-Memory(OOM) 실패를 방지하기 위해 두 가지 주요 메모리 최적화를 구현했습니다.

MoE 전문가 물질화 완화

연구원들은 Hugging Face Transformers 추론 순방향 경로가 모든 전문가에 대해 숨은 상태를 복제하여 GPU 메모리에 extrêmement 큰 텐서를 물질화하는 것을 발견했습니다(예: 20B 모델에 대해 180 GiB 할당 시도). 그들은 전문가를 순차적으로 처리하는 더 메모리 효율적인 실행 경로를 사용하도록 구현을 패치했습니다.

FlashAttention v3를 사용한 시퀀스 병렬성

GPU당 활성화 메모리를 더욱 줄이기 위해 팀은 시퀀스 병렬성(컨텍스트 병렬성)을 구현했습니다. 이는 입력 시퀀스를 장치 간에 분할하여 피크 활성화 발자국을 줄입니다.

주의 레이어는 시퀀스의 모든 토큰이 하나의 GPU에 존재해야 하므로, 팀은 모든-모든 통신 전략을 구현했습니다:

  • Pre-attention: 시퀀스 요소를 모으고 attention-head 수준에서 분할합니다.
  • Post-attention: 출력을 원래의 시퀀스-병렬 레이아웃으로 재분배합니다.

이 설계는 attention sink를 인식하고 FlashAttention v3와 호환되어 다단계 에이전트에 필요한 장기 컨텍스트 윈도우를 처리할 수 있도록 모델을 가능하게 합니다.

Sources