본문으로 건너뛰기

© 2026 Molayo

Zenn헤드라인2026. 05. 31. 18:31

RLVR 시대의 Inference Framework: Weight Syncing 편

요약

RLVR(Verifiable Reward 강화학습) 환경에서 추론 프레임워크의 Weight Syncing 기능의 중요성을 다룹니다. vLLM이 기존 학습 프레임워크의 복잡성을 줄이기 위해 네이티브 API로 가중치 동기화를 지원하게 된 배경과 이점을 설명합니다.

핵심 포인트

  • RL 학습 시 Trainer와 Generator 간의 정책 격차(Policy Lag) 최소화 필요
  • 기존 방식은 학습 프레임워크가 추론 엔진의 확장을 직접 관리하여 복잡성 증가
  • vLLM의 네이티브 Weight Syncing 지원으로 중복 구현 및 버전 종속성 문제 해결
  • 효율적인 RLVR 및 Reasoning LLM 개발을 위한 필수 인프라 기술

서론

본 기사는 vLLM이나 SGLang 등의 Inference Framework가 RL(Reinforcement Learning, 강화학습) 전성기인 현대(2026년)에 RL을 위해 어떠한 기능을 보유하고 있는지, 또한 어떠한 의도/배경으로 해당 기능을 확장해 왔는지에 대해 해설하는 시리즈 기사의 일환으로 집필되었습니다. 이번에는 Claude Code나 Codex 등의 Coding Agent를 구현하기 위해 필요한 Reasoning LLM 개발에서 매우 중요한 역할을 하는 RLVR(Reinforcement Learning with Verifiable Reward) 에서 핵심적인 Weight Syncing에 초점을 맞추어, Inference Framework 측이 어떤 배경 아래 어떤 기능을 지원하고 있는지 해설합니다.

Weights Syncing

본 기사에서는 주요 Inference Engine 중 하나인 vLLM에서의 Weight Syncing 기능 구현 배경과 vLLM이 기능으로서 이를 지원함으로써 어떤 이점이 있는지에 대해 해설합니다.

Weight Syncing Background

On-policy / synchronous RL 설정에서는 Trainer 측에서 업데이트되는 policy와 Generator 측에서 rollout을 생성하는 policy 사이의 괴리를 가능한 한 작게 유지해야 합니다. 이는 rollout generation과 training을 비동기적으로 진행하여 일정 수준의 policy lagstale rollouts를 허용하는 asynchronous RL 설정에서도 마찬가지입니다. Generator에서 이용되는 model weights를 최신 상태 또는 충분히 최근의 것으로 업데이트함으로써, 현재의 Trainer policy에 대해 극단적으로 오래된 rollouts가 섞이는 것을 방지하고, 학습 신호(training feedback)의 유용성을 유지하기 쉽게 만듭니다.

전제 설명이 길어졌습니다만, RL에서는 weight sync가 필요하다는 뜻입니다. 그리고 이 weight sync는 전통적으로 Inference Engine이 아니라, 각각의 RL Training Framework가 별도로 담당해 왔습니다. vLLM을 예로 들면, 많은 경우 vLLM workers가 weight를 전달받고 해당 weight를 load하는 방식으로 확장을 수행함으로써 실현되어 왔습니다. 물론 그러한 vLLM 확장은 기능하지만, 다음과 같은 문제도 존재합니다.

Training Framework 측의 복잡성 증가: Training Framework 개발자에게 vLLM/SGLang의 custom worker extension 구현/유지보수 비용을 전가하게 됩니다. 만약 Weight sync를 위한 기능을 Inference Engine 측에서 보유하고 있다면 이를 해결할 수 있습니다. -
Training Framework마다 중복되는 작업 발생: 많은 RL framework에서 packed tensor transfer나 RPC endpoints 등 유사한 구현이 채택되고 있습니다. NeMo-RL, slime, verl 등 유명한 RL Framework만 보더라도 비슷한 구현이 중복되는 것은 매우 낭비가 심하다고 할 수 있습니다. -
특정 버전에 종속(locking)됨: 많은 Framework에서 특정 vLLM/SGLang 버전에 대해 patch 처리로서 weights를 수신(receiving)한 후 pre/post-processing하는 처리를 ad-hoc하게 구현하고 있습니다. 그 때문에 Inference Engine 버전이 다르면 제대로 동작하지 않을 우려가 있어, 이용자와 개발자 양측 모두에게 비용이 발생하고 있습니다.

이러한 배경으로부터 vLLM은 weight syncing을 위한 API를 native support하게 되었습니다. 다음 섹션에서 어떤 API가 지원되었는지, 그리고 그 개요에 대해 해설하겠습니다.

Native weight transfer

New APIs in vLLM

vLLM이 지원하는 weight syncing을 위한 API는 주로 다음의 4가지로 분류됩니다.

initialization:init_weight_transfer_engine

Trainer workers(프로세스)와 Inference workers 사이에 Training Loop를 시작하기 전에 통신 채널 (communication channel)을 확립하는 역할을 담당합니다. -
Start weight update:start_weight_update

Training Loop 내에서 각 training step 이후, 또는 여러 step마다 호출되어 weight update 처리의 시작을 vLLM workers에 통지합니다. 이 API를 호출함으로써 vLLM workers는 새로운 weights를 받을 준비를 합니다. -
Update weights:update_weights

은(는) Trainer로부터 받은 weights로 Inference Engine 측의 weights를 업데이트하는 API입니다. 업데이트 대상은 모델 전체여도 되고, 일부 weights여도 상관없습니다. 또한, 큰 모델의 경우 한 번에 모든 weights를 전송하지 않고, chunked weights 단위로 여러 번 update_weights를 호출함으로써 단계적으로 weights를 로드할 수 있습니다. -
Finish weight update:finish_weight_update

weight update를 종료하기 위한 API입니다. Quantization (양자화) 등의 새로 받은 weight를 바탕으로 post-process를 수행해야 할 경우, 이를 수행하기 위한 기점이 됩니다.

참고로, 이 API들은 vLLM의 API server mode 및 engine mode 모두에서 구현되어 있습니다.

또한, weight syncing (weight transfer)을 위한 communication backend로서 현재는 다음 두 가지를 지원합니다.

NCCL: Trainer workers로부터의 weight transfer에 NCCL broadcast operations를 이용하는 backend입니다. Trainer workers와 Inference workers는 일반적으로 서로 다른 GPU 상에 존재하기 때문에, NCCL을 통한 통신으로 weight transfer를 실현합니다. -
IPC: 동일한 device 내에서의 weight transfer에는 shared memory handles를 이용한 CUDA IPC를 사용하여 weight transfer를 수행합니다.

이상과 같은 weight transfer를 실현하기 위한 core transport logic은 플러그인 가능한 추상화 구현인 WeightTransferEngine 형태로 구현되어 있습니다. (WeightTransferEngine 클래스를 상속하여 NCCLWeightTransferEngine이나 IPCWeightTransferEngine을 구현하는 방식입니다) 이와 같이 worker implementation (워커 구현)과 weight transport (중량 전송) 구현을 분리함으로써, vLLM을 이용하는 사용자가 독자적인 구현을 쉽게 통합할 수 있도록 하고 있습니다.

NCCL Weight Transfer Engine

IPC Weight Transfer Engine

Initialization (초기화) 및 Update Weights (중량 업데이트) 단계는 일반적으로 RL Framework 개발자에 의해 커스텀되며, transport logic 또한 마찬가지로 그 일환으로서 이용된다는 가정하에 설계되었습니다. 하지만 기존처럼 모든 것을 RL Framework 측에서 담당하는 것이 아니라, vLLM의 API를 이용하여 구현함으로써 RL Framework 간에 중복되었던 구현을 간소화할 수 있다(=통일적인 API 이용으로 대체할 수 있다)는 의미입니다.

다음 섹션에서는 NCCL을 이용한 weight transfer와 post processing에서 FP8 Quantization을 수행하는 경우에 대해 해설합니다.

Weight transfer Example

위 그림을 참고하여 vLLM의 API를 어떻게 호출하는지와 대조하며 해설을 진행하겠습니다.

  • weight transfer engine 설정

다음과 같이 통신에 사용할 backend로 이번에 사용할 NCCL을 지정하여, Weight Transfer 준비를 수행합니다.

from vllm import LLM
from vllm.config import WeightTransferConfig
lm = LLM(
...
  • communication state (통신 상태)의 초기화

위 그림에서는 /init_weight_transfer_engine이 이에 해당합니다. 통신에 사용할 MASTER_PORT, MASTER_ADDRESS, 그리고 NCCL process group (프로세스 그룹)에 참여하는 rank (랭크) 수를 나타내는 world_size를 설정합니다. 이 설정을 통해 Trainer worker와 Inference worker 간의 communication state (통신 상태)를 초기화합니다. 위 그림의 init NCCL Process Group은 바로 통신 Process Group의 초기화에 해당합니다.

from vllm.distributed.weight_transfer.base import WeightTransferInitRequest
# Initialization for inference
llm.init_weight_transfer_engine(
...
  • Trainer worker로부터 weight (가중치) 전송

다음과 같이 weight transfer (가중치 전송)를 시작합니다. trainer_send_weights 메서드를 사용하여 iterator=model.named_parameters()로 받은 parameters (파라미터)를 전송합니다. 또한 NCCLTrainerSendWeightsArgs는 효율화를 위해 packed tensor broadcasting (패킹된 텐서 브로드캐스팅)을 사용할지 여부 등의 옵션을 지정할 수 있습니다.

from vllm.distributed.weight_transfer.nccl_engine import (
NCCLTrainerSendWeightsArgs,
NCCLWeightTransferEngine,
...
  • Inference Engine에서 weight (가중치) 수신

다음은 update_weights가 호출된 시점부터 finish_weight_update가 호출될 때까지의 과정으로, weight를 수신하여 parameter (파라미터) 업데이트를 수행합니다. 또한 구현상에서는 생략하였으나, finish_weight_update를 호출하면 위 그림과 같이 Quantization (양자화)을 실시하고 있는 경우 FP8 Post-process (후처리) 등이 이어서 수행됩니다.

from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
# executed asynchronously while trainer sends weights
llm.start_weight_update()
...

Customizing Weight Transfer (가중치 전송 커스텀)

지금까지 Weight Transfer (가중치 전송)가 왜 필요한지, 그리고 vLLM이 이를 어떻게 지원하는지에 대해 설명했습니다. 마지막으로, custom weight transfer (커스텀 가중치 전송)를 수행할 때 어떻게 확장할 수 있는지에 대한 개요를 설명하겠습니다.

Custom WeightTransferEngine을 정의하려면 WeightTransferInitInfoWeightTransferUpdateInfo를 상속받은 독자적인 data class (데이터 클래스)를 생성해야 합니다. 참고로 WeightTransferInitInfoWeightTransferUpdateInfo 자체는 Abstract class (추상 클래스)이며 실체가 거의 없으므로, 구현 시에는 NCCLWeightTransferUpdateInfo, NCCLWeightTransferInitInfo 등을 참고하는 것을 권장합니다.

다음 예시에서는 MyInitInfoMyUpdateInfo를 생성하고, 이를 사용하여 MyWeightTransferEngine을 정의하고 있습니다. weight transfer 기능에 대한 더 자세한 내용은 이 README를 참조해 주세요.

from dataclasses import dataclass
from typing import Iterator, Callable, Any
from torch import Tensor
...

마치며

본 기사에서는 vLLM에서 지원되기 시작한 Weight Syncing (가중치 동기화) 기능에 대해 해설하였습니다. RLVR에서 왜 Weight Syncing이 필요한지에 대한 설명부터, Weight Syncing이 어떻게 구현되어 있는지까지 전체적으로 살펴보았습니다. 참고로, **Asynchronous RL (비동기 강화학습)**에서의 vLLM 가중치 업데이트(weight update) 기능에 대해서는 후속되는 별도의 기사를 통해 해설할 예정입니다.

저의 다른 기사에서는 LLM 개발을 진행하며 얻은 지식이나, 학습 프레임워크(learning framework)의 이면에서 어떤 일들이 일어나고 있는지에 대해 해설하고 있으니, 괜찮으시다면 그쪽도 확인해 주시기 바랍니다.

Discussion

AI 자동 생성 콘텐츠

본 콘텐츠는 Zenn ML의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.

원문 바로가기
0

댓글

0