PipelineRL: 인플라이트 가중치 업데이트를 통한 LLM 강화 학습 최적화

PipelineRL은 대규모 LLM 훈련에서 높은 추론 처리량과 온-폴리시 데이터 수집 사이의 근본적인 트레이드오프를 해결하도록 설계된 실험적 강화 학습(RL) 구현입니다. 인플라이트 가중치 업데이트를 구현함으로써, PipelineRL은 추론 서버가 추론 과정을 중단하지 않고도 업데이트된 모델 가중치를 받을 수 있게 하여 높은 GPU 활용도를 보장하면서 훈련 데이터를 온-폴리시 근처에 유지합니다.

추론 처리량과 온-폴리시 트레이드오프 해결

전통적인 RL 워크플로에서는 효율성과 데이터 신선도 사이에 갈등이 존재합니다. 높은 처리량을 달성하기 위해 추론 서버는 일반적으로 큰 배치 크기를 사용하며, 이는 여러 정책 최적화 단계에 대한 데이터를 생성합니다. 그러나 각 후속 최적화 단계는 데이터를 수집하는 데 사용된 가중치와 현재 정책 가중치 사이의 "지연"을 증가시켜, 데이터가 점점 오프-폴리시가 되고 훈련에 덜 효과적이게 됩니다.

PipelineRL은 추론을 중단하지 않고 각 옵티마이저 단계 후에 추론 서버의 가중치를 업데이트함으로써 이를 해결합니다. 시스템은 새로운 가중치를 받는 데 필요한 짧은 시간 동안만 추론 서버를 일시 중지합니다. 이 접근 방식은 추론 서버가 최적의 배치 크기를 유지하면서 데이터가 온-폴리시 또는 온-폴리시에 가깝게 유지되도록 하여 보다 안정적이고 효과적인 학습을 가능하게 합니다.

성능 및 안정성 결과

Open-Reasoner-Zero 데이터셋으로 훈련된 7B 및 32B 모델을 사용한 실험 결과, PipelineRL이 AIME 2024 및 MATH 500 추론 벤치마크에서 Open-Reasoner의 성능과 동등하거나 능가함을 보여줍니다.

단순화된 RL 알고리즘

  • Simplified GRPO: 값 함수 없이 Group Relative Policy Optimization(GRPO)의 단순화된 버전을 사용합니다.
  • No Complex Filtering: 구현에서 신뢰 영역 중요도 가중치 클램핑, 과도한 시퀀스 필터링, 보상 형태 지정 등을 생략합니다.
  • Basic Loss Normalization: 손실은 배치 내 시퀀스 수를 분모로 사용하여 정규화하며, 모든 토큰에 동일한 가중치를 부여합니다.
  • No Penalties: 시스템은 KL 패널티나 엔트로피 보너스를 사용하지 않으며(참조 모델 KL은 지원됩니다).

KV 캐시 오래된 데이터의 영향

인플라이트 가중치 업데이트와 관련된 주요 우려는 KV 캐시의 키와 값이 이전 모델 버전으로 계산된 상태로 남아 있어, 시퀀스 생성이 오래된 키와 값으로 계속 진행된다는 점입니다. 그러나 실험 결과는 이것이 훈련 안정성에 부정적인 영향을 미치지 않음을 보여줍니다.

모듈형 아키텍처 및 기술 계약

PipelineRL은 SGLang, vLLM과 같은 특수 추론 소프트웨어 및 DeepSpeed, FSDP, TorchTitan과 같은 훈련 소프트웨어와의 통합을 가능하게 하기 위해 모듈형으로 설계되었습니다. 이는 두 가지 주요 계약을 통해 구현됩니다:

추론 계약

PipelineRL과 통합하려면 추론 소프트웨어가 세 가지 특정 API를 제공해야 합니다:

  1. Process group initialization: 가중치 업데이트를 위한 프로세스 그룹을 초기화하는 HTTP POST /init_process_group 요청.
  2. Weight Update Trigger: 추론 서버에 일시 중지를 알리고 NCCL을 통해 가중치 방송을 받도록 하는 HTTP POST /request_weight_update 요청.
  3. Chat completion: 액터와의 상호작용을 위한 표준 HTTP POST /v1/chat/completion 요청.

트레이너 계약

훈련 소프트웨어는 다음 작업을 위한 Python API를 제공해야 합니다:

  • Worker initialization: 훈련 가중치와 옵티마이저 상태를 로드하고 샤딩합니다.
  • Forward pass: 토큰 로그 가능도를 생성합니다.
  • Backward step: RL 목표에 대한 그래디언트를 계산하고 누적합니다.
  • Optimizer Step: 옵티마이저 단계를 실행합니다.
  • Weight gathering and broadcasting: 업데이트된 가중치를 레이어별로 수집하여 추론 서버에 방송합니다.

실험 구성

PipelineRL은 다음 하이퍼파라미터를 사용하여 7B 및 32B 모델에서 테스트되었습니다:

  • Batch size: 4096
  • Learning rate: 1e-6
  • Max generated tokens: 8192

훈련에 필요한 컴퓨팅은 7B 모델의 경우 2노드에서 약 3.5일, 32B 모델의 경우 4노드에서 6일 정도였습니다.

Sources