비동기 강화학습

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

강화학습(RL) 훈련 루프를 짜다 보면 곧바로 부딪히는 낭비가 있어요. 정책이 롤아웃(rollout)을 생성하는 동안 훈련 가속기는 놀고 있고, 훈련이 돌아가는 동안엔 생성이 멈춰 있죠. 이 문서는 그 병목을 어떻게 해소하는지, 그리고 생성과 훈련을 겹쳤을 때 생기는 까다로운 문제(추론 엔진이 돌아가는 중간에 가중치를 갱신해야 하는 문제)를 vLLM이 어떻게 안전하게 다루는지 설명해요.

출처: vLLM 공식 문서 — async_rl

개요 (Overview)

표준 RL 훈련 루프에서는 생성과 훈련이 순차적으로 진행돼요. 정책이 롤아웃을 생성하고, 그 롤아웃으로 훈련을 돌리고, 다시 반복하는 식이죠. 생성이 돌아가는 동안 훈련 가속기는 놀고, 그 반대 상황도 마찬가지예요.

one-off pipelining 접근 방식은 생성과 훈련 단계를 두 개의 병렬 코루틴으로 분리해요. 이전에 생성한 데이터로 훈련하면서 동시에 모델이 새 샘플을 생성하게 할 수 있죠. 이렇게 하면 GPU 활용률이 좋아지고 훈련 처리량도 높아질 수 있습니다.

다만 이 겹침은 한 가지 복잡함을 불러와요. 요청이 아직 진행 중일 수 있는데, 추론 엔진의 가중치를 도중에 갱신해야 한다는 점이죠.

Pause and Resume API

추론 엔진이 실행되는 동안 가중치를 안전하게 갱신하기 위해 vLLM은 pause_generation과 resume_generation 메서드를 제공해요. 이 메서드들은 트레이너가 진행 중인 작업을 잃지 않으면서 가중치 동기화를 위한 깨끗한 창을 조율할 수 있게 해줍니다.

pause_generation

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

mode 파라미터는 진행 중인 요청을 어떻게 처리할지 제어해요.

clear_cache 파라미터는 일시정지 후 KV 캐시와 프리픽스 캐시를 지울지 제어해요.

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 - 스케줄러를 일시정지하지 않고 진행 중인 요청을 중단 ({}를 보내면 모두 중단, {"request_ids": [...]}로 특정 요청만)
• GET /weight_info - 최신 커밋된 weight_version 반환

데이터 병렬 처리: vLLM의 내부 로드 밸런서(즉 data_parallel_backend="ray")로 데이터 병렬 처리를 쓸 때는 pause와 resume이 모든 DP rank에 걸쳐 자동으로 처리돼요. 한 번만 호출하면 충분하죠. 외부 로드 밸런서(즉 프록시 뒤에 여러 독립 vLLM 인스턴스)를 쓸 때는 가중치 갱신 전후로 모든 엔진 인스턴스에 각각 pause와 resume 요청을 보내야 합니다.

일반적인 Async RL 흐름

가중치 동기화가 있는 일반적인 async RL 루프는 이렇게 생겼어요.

1. 현재 정책으로 롤아웃 생성 시작
2. 트레이너가 갱신할 새 가중치를 갖게 되면 mode="keep"으로 생성 일시정지
3. 트레이너에서 추론 엔진으로 갱신된 가중치 동기화 (자세한 내용은 Weight Transfer 참고)
4. 생성 재개 — 진행 중이던 요청이 새 가중치로 계속됨
5. 반복

핵심 통찰은 mode="keep"으로 일시정지된 요청이 일시정지 전에는 이전 가중치로 토큰을, 재개 후에는 새 가중치로 토큰을 만든다는 점이에요. clear_cache 파라미터는 일시정지 중에 KV 캐시를 무효화할지 제어합니다. clear_cache=True면 이전에 캐시된 key-value 항목이 버려져서, 재개 후 생성되는 모든 토큰이 전적으로 새 가중치로 계산돼요. clear_cache=False면 기존 KV 캐시 항목이 유지되므로 컨텍스트의 일부 토큰이 여전히 이전 가중치를 반영할 수 있어요(오래된 KV 캐시).

예시 (Example)

async RLHF 예시는 vllm.AsyncLLMEngine, NCCL 가중치 전송, 그리고 검증과 함께 도중 pause/resume을 사용하는 이 패턴을 보여줘요.

더 알아보기 (Learn more)

• 가중치 전송 (Weight Transfer)
• RLHF
• 표본 마스크 (Sampling Mask)