비동기 강화학습

비동기 강화학습 (Async Reinforcement Learning)

표준 RL 학습 루프에서는 생성과 학습이 순차적으로 일어나며, 생성 중에는 학습 가속기가, 학습 중에는 생성이 쉬게 됩니다. vLLM의 one-off pipelining 접근법은 이 두 단계를 두 개의 병렬 코루틴으로 분리해, 이전에 생성된 데이터로 학습하면서 동시에 새 샘플을 생성하게 합니다. 이는 더 나은 GPU 활용과 더 높은 학습 처리량을 낼 수 있습니다.

출처: 문서

본문

개요 (Overview)

표준 RL 학습 루프에서는 생성과 학습이 순차적으로 일어납니다: 정책이 rollout을 생성하고, 학습이 그 rollout에서 실행되며, 그 주기가 반복됩니다. 생성 중에는 학습 가속기가 유휴하고, 그 반대도 마찬가지입니다.

one-off pipelining 접근법은 생성과 학습 단계를 두 개의 병렬 코루틴으로 분리하여, 모델이 이전에 생성된 데이터로 학습하면서 동시에 새 샘플을 생성하게 합니다. 이는 더 나은 GPU 활용과 더 큰 학습 처리량으로 이어질 수 있습니다.

하지만 이 겹침은 복잡성을 도입합니다: 요청이 여전히 진행 중인 동안 가중치가 추론 엔진에서 mid-flight로 업데이트되어야 합니다.

일시정지/재개 API (The Pause and Resume API)

추론 엔진이 실행되는 동안 안전하게 가중치를 업데이트하기 위해 vLLM은 pause_generation과 resume_generation 메서드를 제공합니다. 이는 트레이너가 in-flight 작업을 잃지 않고 가중치 동기화를 위한 깨끗한 창을 조정하게 해 줍니다.

pause_generation

await engine.pause_generation(mode="keep", clear_cache=True)

mode 파라미터는 in-flight 요청을 처리하는 방식을 제어합니다:

모드 동작
"abort" 모든 in-flight 요청을 즉시 중단하고 부분 결과를 반환 (기본값)
"wait" 일시정지 전에 모든 in-flight 요청이 끝날 때까지 대기
"keep" 큐의 요청을 고정(freeze); resume_generation 호출 시 재개

clear_cache 파라미터는 일시정지 후 KV cache와 prefix cache를 비울지 제어합니다.

resume_generation

await engine.resume_generation()

일시정지 후 스케줄러를 재개합니다. mode="keep"로 고정된 요청은 계속 생성합니다.

HTTP 엔드포인트

VLLM_SERVER_DEV_MODE=1로 설정하면 vLLM HTTP 서버가 같은 기능을 노출합니다:

  • POST /pause?mode=keep — 생성 일시정지
  • POST /resume — 생성 재개
  • POST /abort_requests — 스케줄러를 일시정지하지 않고 in-flight 요청 중단 ({}를 보내면 모두 중단, 또는 {"request_ids": [...]})
  • GET /weight_info — 가장 최근에 커밋된 weight_version 반환

데이터 병렬

vLLM의 내부 로드밸런서(즉 data_parallel_backend="ray")로 데이터 병렬을 사용할 때는 pause/resume이 모든 DP 랭크에 걸쳐 자동으로 처리되므로, 단일 호출이면 충분합니다. 외부 로드밸런서(즉 프록시 뒤의 여러 독립 vLLM 인스턴스)를 사용할 때는 가중치 업데이트 전후에 모든 엔진 인스턴스에 각각 pause/resume 요청을 보내야 합니다.

일반적인 비동기 RL 흐름 (Typical Async RL Flow)

가중치 동기화를 포함한 일반적인 비동기 RL 루프는 다음과 같습니다:

  1. 현재 정책으로 rollout 생성 시작
  2. 트레이너가 업데이트할 새 가중치가 생기면 mode="keep"로 생성 일시정지
  3. 업데이트된 가중치를 트레이너에서 추론 엔진으로 동기화(Weight Transfer 참고)
  4. 생성 재개 — in-flight 요청은 새 가중치로 계속
  5. 반복

핵심 통찰은 mode="keep"로 일시정지된 요청이 일시정지 전에는 이전 가중치의 토큰을, 재개 후에는 새 가중치의 토큰을 만들어낸다는 것입니다. clear_cache 파라미터는 일시정지 중에 KV cache를 무효화할지 제어합니다. clear_cache=True면 이전에 캐시된 key-value 항목이 폐기되어, 재개 후 생성되는 모든 토큰이 완전히 새 가중치로 계산됩니다. clear_cache=False면 기존 KV cache 항목이 유지되어, 문맥의 일부 토큰이 여전히 이전 가중치(stale KV cache)를 반영할 수 있습니다.

예시 (Example)

async RLHF 예시는 vllm.AsyncLLMEngine, NCCL 가중치 전송, 검증을 포함한 mid-flight pause/resume으로 이 패턴을 보여줍니다.

더 알아보기 (Learn more)