ServiceNow/PipelineRL
A scalable asynchronous reinforcement learning implementation with in-flight weight updates.
해결하는 문제
PipelineRL은 대규모 언어 모델(LLM)을 위한 강화학습(RL)에서 효율성의 트레이드오프를 해결합니다. 구체적으로, 높은 추론 처리량(여러 GPU에 걸쳐 대규모 배치가 필요)을 달성하는 것과 "온폴리시(On-policy)" 데이터의 최신성(모델이 최신 버전으로 생성한 데이터를 기반으로 학습되도록 보장)을 유지하는 것 사이의 갈등을 해결합니다.
작동 방식
PipelineRL은 "비행 중 가중치 업데이트"를 사용하는 확장 가능한 비동기 아키텍처를 사용합니다. 모델을 업데이트하기 위해 전체 샘플링 프로세스를 중단하지 않고, 최적화 단계 이후 바로 NCCL을 통해 업데이트된 가중치를 추론 서버에 브로드캐스트합니다. 이로 인해 정책 업데이트 중에도 로울아웃을 계속 생성할 수 있어 GPU 활용도를 희생하지 않으면서도 데이터를 온폴리시에 가깝게 유지할 수 있습니다.
시스템은 여섯 가지 모듈식 구성 요소로 구성됩니다:
- 오케스트레이터: GPU 배치를 관리하고 하위 프로세스를 시작합니다.
- 추론 서버: vLLM 기반 서버로 샘플링을 처리하고 가중치 업데이트를 수신합니다.
- 액터 프로세스: LLM에서 샘플링하여 원시 로울아웃을 생성하고 보상을 수집합니다.
- 프리프로세서: 시퀀스를 토큰화하고 이득을 계산합니다.
- 트레이너: RL 단계(간소화된 GRPO 알고리즘 사용)를 수행하고 가중치 업데이트를 트리거합니다.
- 검증자: 선택적 서버로, 모델 출력의 정확성을 검증합니다(예: 수학 문제).
대상 사용자
특히 수학이나 코딩과 같은 추론 작업에 초점을 맞추고 보상이 검증 가능한 연구자 및 개발자에게 적합합니다. 여러 GPU에 걸쳐 훈련을 확장해야 하는 사용자에게 특별히 설계되었습니다.
주요 특징
- 비행 중 가중치 업데이트: 샘플링 파이프라인을 중단하지 않고 모델 파라미터를 업데이트합니다.
- 에이전트 독립성:
load_problems및generate_rollout함수를 구현함으로써 어떤 에이전트 작업에도 적응 가능합니다. - 고성능: AIME-2024 및 MATH-500 벤치마크에서 Open-Reasoner-Zero와 동등하거나 이를 초과합니다.
- 유연한 백엔드: 프로세스 간 통신에 파일 시스템 또는 Redis 기반 스트리밍을 모두 지원합니다.
- 사전 통합: SandboxFusion을 지원하여 원격 사전에서 코드를 실행하고 검증할 수 있습니다.
관련
- Dispatch
- 프로젝트
- 프로젝트
- Dispatch
- 프로젝트