비동기 강화학습의 안정성을 좌우하는 커널 정렬

들어가며

비동기 강화학습은 rollout이 샘플을 생성하는 동안 learner가 모델을 업데이트할 수 있다는 점에서 매력적이다. 동기식 파이프라인처럼 한 작업이 끝날 때까지 다른 작업이 기다리지 않으므로, 추론과 학습 자원의 활용률을 높일 수 있다. 하지만 learner가 학습하는 샘플은 현재 정책이 아니라 이전 버전의 정책으로 생성된다.

vLLM Blog가 소개한 NVIDIA Research와 Inferact의 REINFORCE Zero는 이 문제를 단순한 정책 지연으로만 보지 않는다. 비동기 학습에는 서로 다른 성격을 가진 두 종류의 불일치가 함께 존재한다는 것이 핵심 주장이다.

정책 지연과 커널 불일치

샘플이 가중치 버전 v에서 생성되고 버전 k에서 학습된다면, learner가 계산한 토큰 확률을 rollout 시점에 기록된 확률로 나눠 중요도 샘플링을 적용할 수 있다. 이 비율은 현재 정책과 샘플 생성 정책 사이의 차이를 보정하기 위한 전형적인 오프폴리시 처리다.

그러나 양쪽이 같은 가중치를 사용하더라도 같은 logprob를 계산한다는 보장은 없다. 연산 커널, 수치 정밀도, 배치 구성, 실행 경로가 다르면 동일한 토큰에 대한 확률이 달라질 수 있다. 이를 커널 불일치라고 한다.

두 현상은 시간적 특성도 다르다. 정책 지연은 업데이트가 작을 때 정책 변화의 흐름을 따라 비교적 부드럽게 변한다. 반면 커널 불일치는 배치와 수치 실행에 따라 흔들리며, 중요도 가중치에 큰 이상치를 만들 수 있다. 소재에 인용된 실험에서 중요도 편차 지표의 최대값과 중앙값의 비율은 커널을 정렬했을 때 1.50이었지만, 표준 커널에서는 442.54까지 커졌다.

REINFORCE Zero의 구성

REINFORCE Zero는 다음 세 가지 원칙으로 정리된다.

• 같은 가중치와 토큰에 대해 rollout과 learner가 동일한 logprob를 계산하도록 커널을 정렬한다.
• 토큰 단위 중요도 샘플링을 적용하되, 가중치를 자르거나 시퀀스를 거부하는 게이트를 사용하지 않는다.
• 가치 모델 없이 그룹 평균 보상과 각 응답 보상의 차이를 어드밴티지로 사용해 일반적인 REINFORCE를 최적화한다.

중요도 비율을 클리핑하지 않는 이유는 정확한 비율을 사용해야 기대 정책 그래디언트의 보정이 유지되기 때문이다. 클리핑은 분산을 줄일 수 있지만 일반적으로 편향을 만든다. 다만 이 방식이 안정적으로 작동하려면 먼저 커널에서 발생하는 더 큰 변동을 제거해야 한다.

실험 결과와 의미

표준 커널을 사용한 비동기 실험에서는 rollout과 learner 사이의 KL이 계속 증가했고, AIME24 pass@1은 34.27%에서 13.96%로 하락했다. 최대 그래디언트 노름은 38.27에 이르렀다. 반대로 커널을 정렬한 REINFORCE Zero에서는 KL이 안정적으로 유지됐고 샘플 거부가 발생하지 않았다. 시퀀스 게이트를 추가하면 표준 커널도 안정화할 수 있었지만, 마지막 구간에는 시퀀스의 24.48%를 거부했다.

Qwen3-4B에 Python 도구와 16K 컨텍스트를 사용한 실험에서는 정책 지연이 64에 이르러도 5,000회 업데이트 동안 학습이 이어졌다. 평균 보상은 초기 100회 업데이트의 0.287에서 약 0.78까지 상승했다. 16개 GPU와 12,000회 업데이트를 비교한 결과, 정렬된 비동기 방식은 정렬된 동기 방식보다 학습 스텝 처리량이 2.40배 높았다. DeepSeek-V4.1-Flash, Kimi K3.1, AMD MI355X로의 적용도 제시됐다.

이 작업의 중요한 시사점은 중요도 샘플링만으로 모든 비동기 오차를 해결할 수는 없다는 것이다. 먼저 확률 계산을 일치시키고, 그 다음 정책 지연을 통계적으로 보정해야 한다. 도구 사용 에이전트나 긴 컨텍스트 학습에서는 유용한 시스템 설계 원칙이 될 수 있지만, 모델과 하드웨어별 커널 정렬 비용도 함께 고려해야 한다.

출처: vLLM Blog