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

인플라이트 (inflight) 중량 업데이트는 RL(강화학습) 훈련 동안 발생하는 현상입니다 (그림 1 참조). 이는 PipelineRL이 항상 높은 추론 성능을 달성하고, 롤아웃에 사용된 중량과 최신 업데이트 모델 중량 간의 지연을 최소화할 수 있게 합니다. 결과: 대규모 언어 모델의 빠르고 안정적인 RL 훈련.

이 블로그 포스트에서 우리는 다음 두 가지를 보여줍니다: 1) 인플라이트 중량 업데이트는 훈련 과정에 해를 끼치지 않으며, 2) PipelineRL은 더 간단한 RL 알고리즘을 사용하면서도 Open-Reasoner-Zero 와 경쟁력 있는 결과를 달성합니다. 또한 새로운 추론/트레이너 조합을 용이하게 하는 모듈러 PipelineRL 아키텍처도 제시합니다.

전통적인 RL 접근법 (그림 1a) 에서는 높은 성능의 추론과 온-폴리시 데이터 수집 사이에서 트레이드오프가 존재합니다. 이 트레이드오프를 설명하기 위해 먼저 전통적인 RL 알고리즘을 알고리즘적으로 정의해 보겠습니다:

current_policy = initial_policy
opt_state = init_optimizer(current_policy)
while True:
...

높은 성능의 추론을 달성하려면 인플레서버는 큰 배치 크기를 사용해야 하며, 따라서 여러 정책 최적화 단계에 대한 데이터를 생성해야 합니다. 그러나 각 최적화 단계는 현재 정책과 추론 정책으로 수집된 데이터 사이의 지연을 증가시켜, 점차적으로 수집된 데이터가 오프-폴리시 (off-policy) 가 되고 훈련에 더 효과적이지 않게 됩니다. 온-폴리시 학습은 단일 최적화 단계의 데이터를 필요로 합니다. 그러나 많은 GPU 를 사용하여 소량의 데이터를 생성하는 것은 비효율적이기 때문입니다. 이는 GPU 당 배치 크기가 작음을 의미하기 때문입니다. 또한 인플레서버가 짧은 시퀀스를 완료하고 몇 가지 가장 긴 시퀀스만 진행 중일 때, 배치 크기는 감소합니다.

PipelineRL (그림 1b) 은 인플라이트 중량 업데이트를 통해 이 트레이드오프를 해결합니다. 우리는 추론 서버에서 각 최적화 단계 후 중량을 업데이트하며, 추론을 결코 중단하지 않습니다. 우리는 모든 인플레서버에서 추론을 일시적으로 멈추는 데 필요한 시간만큼만 새로운 중량을 수신하기 위해 일시정지합니다. 인플라이트 중량 업데이트는 추론 서버가 항상 최적의 배치 크기를 유지할 수 있게 하며, 동시에 데이터가 온-폴리시 또는 근접 온-폴리시를 보장하여 더 나은 GPU 활용도와 더 효과적인 학습을 가능하게 합니다.

PipelineRL 의 효과성과 인플라이트 중량 업데이트의 이점을 입증하기 위해 우리는 Open-Reasoner-Zero 데이터셋에서 7B 모델과 32B 모델을 훈련시켰습니다. 학습 곡선을 보면, PipelineRL 은 인기 있는 추론 테스트 벤치마크인 AIME 2024 와 MATH 500 에서 Open-Reasoner 의 성능을 맞거나 초과함을 볼 수 있습니다 (그림 2 참조).

특히, 우리의 RL 구현은 Open-Reasoner-Zero 보다 훨씬 단순합니다. Open-Reasoner-Zero 는 가치 함수를 사용하지만, 우리의 구현은 GRPO 의 단순화된 버전입니다. 특히 우리는 안정적인 훈련을 위해 신뢰 영역 중요도 중량 클램핑이 필요하지 않다고 발견했습니다. DAPO 논문에서 제시된 과장 시퀀스 필터링이나 보상 형성도 필요하지 않았습니다. 손실을 정규화할 때 우리는 배치의 시퀀스 수를 분모로 사용하여 모든 토큰에 균등한 가중치를 부여합니다. 우리는 KL 페널티와 엔트로피 보상을 사용하지 않았으며 (물론 우리의 구현은 참조 모델 KL 을 지원합니다). 우리의 구현의 단순성에도 불구하고, 또는 아마도 그 덕택으로 훈련이 매우 안정적임을 이 wandb 보고서에서 볼 수 있습니다.

이전 모델 버전으로 계산된 KV 캐시 (sequence generation proceeds with stale keys and values in the KV cache) 를 사용하여 시퀀스 생성이 진행되므로, 비행 중 (inflight) 가중치 업데이트가 불안정한 훈련 프로세스를 초래할 것이라 예상할 수 있습니다. 그러나 우리의 실험은 이것이 안정성에 부정적인 영향을 미치지 않는다는 것을 보여줍니다.

PipelineRL 은 모듈러 구조로 설계되어 고도로 전문화된 추론 및 훈련 소프트웨어 (SGLang, vLLM, Nvidia Dynamo, DeepSpeed, FSDP, TorchTitan, FastLLM 등) 의 빠른 개선점을 활용할 수 있도록 구축되었습니다. 우리는 추론 및 훈련 컴포넌트 간의 명확한 계약을 제안하여 새로운 추론 및 훈련 솔루션이 사용할 때마다 쉽게 통합할 수 있습니다.

추론 소프트웨어는 PipelineRL[1] 에 다음 API 를 노출해야 합니다:

프로세스 그룹 초기화: 시작 시간 (start-up time) 에, Trainer 0 (지정된 코디네이터) 는 모든 추론 서버에 HTTP POST /init_process_group 요청을 보냅니다. 이 요청은 가중치 업데이트를 위한 프로세스 그룹을 초기화합니다.가중치 업데이트 트리거: 트레이너들이 학습 단계를 완료한 후 (옵티마이저 단계 및 가중치 수집), Trainer 0 는 추론 엔드포인트에 HTTP POST /request_weight_update 요청을 제출합니다. 요청에는 메인 트레이너 프로세스가 NCCL 을 통해 전송할 가중치의 순서 및 모양에 대한 세부 정보가 포함되어 있습니다. 추론 서버는 추론을 일시 중지하고 가중치 브로드캐스트를 받아야 합니다.채팅 완성: 액터 프로세스는 HTTP POST /v1/chat/completion 요청을 사용하여 액터 LLM 과 상호작용합니다.

만약 init_process_group 및 request_weight_update API 가 산업 표준이 된다면, PipelineRL 을 사용하여 다른 추론 구현을 플러그 앤 플레이로 시도할 수 있습니다.

PipelineRL 훈련 코드는 각 트레이너 워크러가 올바른数量的의 훈련 토큰이 누적될 때 바로 새로 생성된 훈련 데이터를 워크러에게 공급합니다. 이러한 Python API 를 노출하는 모든 훈련 소프트웨어를 PipelineRL 과 함께 사용할 수 있습니다:

워크러 초기화: 훈련 가중치 및 옵티마이저 상태를 로드하고 쉐어드 (shard) 합니다.프워드 패스 (Forward pass): 입력에 대한 토큰 로그 확률을 생성합니다.백워드 스텝 (Backward step): 선택된 RL 목표량을 나타내는 스칼라의 기울기를 계산하고 누적합니다.옵티마이저 단계 (Optimizer Step): 옵티마이저 단계를 실행합니다.가중치 수집 및 브로드캐스팅: 옵티마이저 단계 후, 트레이너 소프트웨어는 추론 서버에 브로드캐스트할 준비를 위해 업데이트된 모델 가중치를 레이어별로 수집해야 합니다.

PipelineRL 은 현재 HuggingFace accelerate 라이브러리를 사용하여 DeepSpeed 와 FSDP 사이에서 사용자의 선택을 제공합니다. 그러나 우리는 accelerate 계약이 너무 유연하고 혼란스러울 수 있다고 발견했습니다. 우리는 위의 엄격한 계약을 따르는 다른 트레이너를 사용하기 쉬워질 것입니다.

향상 기능 (Upcoming features). 우리의 구현은 여전히 실험적이며 중요한 기능을 누락하고 있습니다. 우리의 최우선 순위는 더 정확한 추론 배치 크기 제어, 멀티모달 지원 및 시퀀스 병렬 훈련을 위한 코루틴 사용입니다. 우리는 또한 더 많은 추론 서버 및 트레이너 통합에 대한 기여를 환영합니다. 그러나 우리는 pipeline-rl 저장소를 모든 가능한 알고리즘 및 보상 함수를 지원하는 프레임워크로 만들려고 시도하지 않습니다. 우리의 견해는 pipeline-rl 이 GRPO 의 해커블하고 빠른 참조 구현이어야 하며 쉽게 검증 가능한 보상을 가져야 한다는 것입니다. PipelineRL 을 사용하여 연구 프로젝트를 수행하고 싶다면, 저장소를 포크하여 코드를 해킹하는 것을 즐기세요!

곧 더 많은 연구가 공개됩니다. Inflight weight updates 가 훈련 동역학에 미치는 영향을 이해하고 PipelineRL 이 가져온 속도 향량을 정밀하게 측정하기 위해서는 더 많은 분석이 필요합니다. 또한, PipelineRL 과 비동기 강화학습 (Asynchronous Reinforcement Learning) 을 위한 LLM 에 대한 관련성 높은 기존 작업 간의 유사성에 대해 많은 것을 말할 수 있습니다. 모든 것들과 더 많은 내용을 위해 곧 발표될 연구 논문을 기다려주세요!

Alexandre Piché 는 TapeAgents 를 개발하는 동안 우리 RL 코드의 첫 번째 동기 버전 (synchronous version) 을 작성했습니다. Dzmitry Bahdanau 는 코드를 비동기 및 분산으로 재작성 (refactored) 하고 inflight weight updates 를 구현했습니다. Rafael Pardinas 는 sequence packing 을 구현했습니다. Ehsan Kamaloo 는 실험을 실행하는 데 도움을 주었고, Xiaoyin Chen 은 프레임워크 디버깅에 도움을 주었습니다.

우리는 TRL, OpenRLHF, veRL 과 같은 기존 LLM RL 구현에서 많은 트릭 (tricks) 을 빌려온 것을 인정합니다. Simple-RL, Deepscaler, DAPO, OpenReasoner 와 같은 다른 오픈소스 추론 프로젝트의 아티팩트 (artifacts) 는 PipelineRL 의 안정화에 필수적이었습니다. 우리는 thoughtful comments 를 주신 Christopher Manning 과 Michael Noukhovitch 를 인정하고 싶습니다. 마지막으로, ServiceNow Research 팀과 ServiceNow CoreLLM 팀에 훌륭한 동료로 감사드립니다.

[1] 코드 내의 현재 계약 (contract) 은 약간 다르지만, 위에서 설명한 대로 재작성 (refactoring) 하고 있습니다.

우리는 여기서 보고하는 7B 와 32B 실험 모두 동일한 하이퍼파라미터를 사용했습니다:

• 배치 크기 (batch size) 4096
• 학습률 (learning rate) 1e-6
• 생성된 토큰 최대 개수 (max number of generated tokens) 8192
• OpenReasoner 실험에서는 16K 토큰의 생성을 허용했습니다.

보고된 실험에 사용된 컴퓨팅 자원:

• 7B 모델: 2 노드에서 약 3.5 일
• 32B 모델: 4 노드에서 약 6 일