vLLM V1 마이그레이션: 강화 학습에서 백엔드 정확성 보장

ServiceNow AI는 목표 측면의 수정보다 백엔드 정확성을 우선시하여 RL 추론 엔진을 vLLM V0에서 V1으로 성공적으로 마이그레이션했습니다. 팀은 네 가지 주요 차이를 해결함으로써 vLLM 0.8.5 (V0)와 vLLM 0.18.1 (V1) 간의 동등성을 달성했습니다: 처리된 롤아웃 로그확률, V1 전용 런타임 기본값, 진행 중인 가중치 업데이트 경로, 그리고 lm_head 투영의 정밀도.

마이그레이션 목표

PipelineRL과 같은 시스템에서 롤아웃 생성을 위한 추론 엔진으로 vLLM을 사용할 때, 엔진은 토큰을 샘플링하고 트레이너가 정책 비율, KL 발산, 클립 비율, 엔트로피 등 필수 RL 지표를 계산하는 데 사용하는 로그확률을 반환합니다. 이러한 로그확률이 계산되는 방식에 차이가 있으면 학습-추론 불일치가 발생하여 학습 역학을 근본적으로 바꿀 수 있습니다.

vLLM V1으로 마이그레이션하기 위해 ServiceNow AI는 좁은 목표를 설정했습니다: V1이 트레이너가 기대하는 형태로 롤아웃 로그확률을 반환하는지 확인하고, V0 레퍼런스와 워크로드를 다시 실행하며, 백엔드 동등성이 복원된 후에만 목표 수준의 변화를 평가합니다. V1을 사용한 초기 시도에서는 보상, 엔트로피, 클립 비율이 V0 레퍼런스에 비해 크게 차이 나는 것이 나타났습니다.

백엔드 실패 모드

팀은 학습-추론 불일치의 잠재적 원인을 세 가지 계층으로 분류했습니다:

Semantic mismatch: 백엔드가 트레이너가 기대하는 것과 다른 의미의 로그확률을 반환합니다.
Inference-path mismatch: 캐싱, 스케줄링 또는 요청 처리에 대한 런타임 기본값 차이로 인해 프롬프트가 다른 실행 경로를 따르게 됩니다.
Objective mismatch: RL 목표는 남아 있는 오래된 상태나 백엔드 불일치를 보정해야 합니다.

첫 번째와 두 번째를 백엔드 동작 문제로 간주하고 먼저 배제함으로써, 팀은 백엔드 오류를 가리기 위해 목표 측면의 수정을 조기에 적용하는 실수를 피했습니다.

vLLM V1 백엔드 수정

로그확률 의미

기본적으로 vLLM V1은 로짓 후처리(예: 온도 스케일링 및 top-k/top-p 필터링) 이전의 원시 모델 출력에서 로그확률을 반환합니다. PipelineRL은 샘플러가 사용하는 처리된 분포의 로그확률을 필요로 합니다. 이를 해결하기 위해 logprobs-mode=processed_logprobs를 설정했으며, 이로써 롤아웃 로그확률의 평균 오프셋이 사라지고 평균 정책 비율이 1.0에 가깝게 중심을 잡았습니다.

런타임 기본값

동등성을 보장하기 위해 팀은 V0 레퍼런스 경로와 차이가 나는 V1 전용 기본값을 명시적으로 비활성화했습니다:

Prefix Caching: 비활성화 (enable-prefix-caching: false)하여 가중치 업데이트 이전에 계산된 상태가 재사용되는 것을 방지하고, 이는 캐시 수명에 V1 전용 차이를 도입할 수 있습니다.
Async Scheduling: 비활성화 (async-scheduling: false)하여 또 다른 V1 전용 자유도를 제거합니다.

진행 중인 가중치 업데이트

V0와 달리 명시적인 캐시 상태 무효화 없이 새로운 가중치를 로드하는 동작을 맞추기 위해 팀은 다음과 같은 V1 업데이트 경로를 구현했습니다:

await engine.pause_generation(mode="keep", clear_cache=False)
await engine_client.collective_rpc_async(
    "receive_weight_update",
    args=(request.model_dump_json(),),
)
await engine.resume_generation()

mode="keep"와 clear_cache=False를 사용함으로써 가중치 동기화 모델이 V0 레퍼런스와 일치하도록 보장했으며, 이는 학습 후반에 지속적인 지연을 감소시켰습니다.

수치 동등성: fp32 lm_head

백엔드 수정 후에도 최종 동등성을 위해 로짓을 계산하는 수치 경로를 일치시켜야 했습니다. 트레이너는 최종 투영에 fp32 lm_head를 사용했으며, 롤아웃 백엔드도 이 정밀도에 맞게 업데이트되어야 했습니다.

이 정밀도는 RL 업데이트가 토큰 로그확률을 직접 소비하기 때문에 중요합니다; 로짓의 작은 변화도 정책 비율과 KL 발산에 크게 영향을 미칠 수 있습니다. 이 발견은 MiniMax-M1 기술 보고서와 ScaleRL 논문과 일치하며, 두 문서 모두 대규모 RL에서 훈련/추론 토큰-확률 불일치를 방지하기 위해 fp32 헤드 계산이 필요한 설계 선택이라고 밝히고 있습니다.

왜 백엔드 정확성이 목표 수정보다 먼저인지

잘린 중요도 샘플링이나 중요도 비율 재가중치와 같은 도구는 오래되었거나 비동기적인 롤아웃을 보정할 수 있지만, 백엔드를 먼저 수정하지 않고 이를 적용하면 결과가 혼란스러워집니다. 목표 측면의 보정을 사용해 손상된 추론-백엔드 동작을 보완하면 학습 곡선을 해석하기 어려워집니다.

팀의 접근 방식은 남은 불일치가 추론 엔진의 버그가 아니라 RL 목표 설계(예: 비동기/오프-정책 문제)에서 비롯된 것임을 보장합니다. 향후 개선은 비동기/오프-정책 정리에 초점을 맞출 것이며, 예를 들어 롤아웃 시점의 명시적 행동-정책 로그확률을 유지하고 최적화 시점에 트레이너 측의 이전 정책 로그확률을 재계산하는 방식을 포함합니다.

SUMMARY: ServiceNow AI는 로그확률 의미, 런타임 기본값, 가중치 업데이트 경로를 수정하고 fp32 lm_head를 구현함으로써 RL 롤아웃 생성에서 vLLM V0와 V1 간의 동등성을 달성했습니다.

TITLE: vLLM V1 마이그레이션: 강화 학습에서 백엔드 정확성 보장