가중치 전송

가중치 전송 (Weight Transfer)

RLHF, GRPO 같은 온라인 RL 기법에서는 훈련 중 정책 모델이 반복적으로 갱신되고, 그 갱신된 가중치가 롤아웃 생성용 추론 엔진에 반영되어야 해요. 문제는 "훈련 프로세스의 가중치를 어떻게 추론 엔진까지 안전하게 전달할까"라는 거죠. 이 문서는 vLLM이 그 동기화를 처리하는 플러그형 가중치 전송(weight transfer) 시스템을 설명해요.

출처: vLLM 공식 문서 — weight_transfer

아키텍처 (Architecture)

가중치 전송은 프로세스당 두 개의 엔진을 가지며 서로 대칭이에요.

트레이너 엔진은 상태를 가져요. 자신의 communicator와 wire params를 소유하고, WeightSource에서 가중치를 끌어오며, VLLMWeightSyncClient를 통해 추론 쪽을 구동합니다. 트레이너 코드는 전송 방식을 알 필요가 없고, 매 라운드마다 전송 상태를 다시 넘길 필요도 없어요. 동기화당 send_weights() 호출 하나만으로 충분하죠.

내부적으로 모든 라운드는 같은 4단계 프로토콜이며, 트레이너 엔진이 대신 구동해줘요.

1. 초기화 (init_weight_transfer_engine): 트레이너와 추론 워커 사이의 통신 채널을 구축. 훈련 루프가 시작되기 전에 trainer_init에서 한 번 호출.
2. 시작 (start_weight_update): 가중치 갱신을 위해 추론 엔진을 준비.
3. 가중치 갱신 (update_weights): 갱신된 가중치를 전송. 한 번 이상 호출될 수 있어요(예: 청크 전송).
4. 종료 (finish_weight_update): 갱신을 마무리(예: 체크포인트 형식 가중치의 사후 처리). 모든 가중치가 전송된 뒤 한 번 호출.

사용 가능한 백엔드 (Available Backends)

퀵스타트 (Quickstart)

추론 쪽 (Inference Side)

추론 쪽은 백엔드 이름만 받아요. 전송에 관한 나머지는 트레이너가 정해서 초기화 핸드셰이크에서 보내줍니다.

from vllm import LLM
from vllm.config import WeightTransferConfig

llm = LLM(
    model="my-model",
    weight_transfer_config=WeightTransferConfig(backend="nccl"),  # or "ipc", "sparse_nccl", "sharded_rdt"
)

온라인 서빙이라면 다음과 같이요.

vllm serve my-model \
    --weight-transfer-config '{"backend": "nccl"}'

트레이너 쪽 (Trainer Side)

엔진을 한 번 만들고, 매 동기화마다 send_weights()를 호출하면 돼요.

from vllm.distributed.weight_transfer import (
    ModuleSource,
    HTTPVLLMWeightSyncClient,
    WeightTransferTrainerFactory,
)
from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo

# Once, before the training loop.
engine = WeightTransferTrainerFactory.trainer_init(
    init_info=NCCLTrainerInitInfo(
        master_address=master_address,
        master_port=master_port,
        world_size=world_size,   # trainer + all inference workers
        rank=0,                  # this trainer rank; rank 0 is the sender
        packed=True,
    ),
    client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
    source=ModuleSource(model),
)

# Once per weight sync.
for step in range(num_steps):
    train_one_step(model)
    engine.send_weights()

send_weights()는 추론 쪽의 start → update → finish 그리고 데이터 플레인 전송을 구동해요. 여기에는 백엔드가 필요로 하는 동시성도 포함되죠(NCCL은 예를 들어 워커의 update_weights와 트레이너의 broadcast를 동시에 실행해야 합니다. 둘 다 같은 NCCL 호출 안에서 rendezvous 하니까요).

트레이너 쪽에는 backend= 인자가 없어요. 각 TrainerInitInfo 서브클래스가 자신의 backend를 선언하고, 팩토리가 그걸로 디스패치합니다.

각 설정이 사는 곳 (Where Each Setting Lives)

트레이너는 양쪽이 동의해야 하는 모든 것의 단일 진실 소스(single source of truth)예요. 이것이 이 시스템에서 기억해야 할 핵심 규칙입니다.

참고: wire params는 의도적으로 WeightTransferConfig에도, 라운드별 update info에도 두지 않아요. 트레이너가 초기화 때 보내고 워커가 전달받은 값을 읽기 때문에, 양쪽의 packed 플래그 불일치는 단순히 권장되지 않는 게 아니라 표현 불가능합니다.

멀티 랭크 트레이너 (Multi-Rank Trainers)

샤딩된 트레이너(FSDP, TP/PP/EP)에서는 모든 트레이너 랭크가 엔진을 만들고 send_weights()를 호출해요. 랭크 0이 송신자(sender)이고, communicator를 보유하고 클라이언트와 통화하며 바이트를 와이어에 올리는 건 랭크 0뿐입니다. 송신자가 아닌 랭크도 여전히 WeightSource를 순회하는데, 파라미터를 실체화(materialize)하는 것이 보통 집단 연산(collective)이기 때문이에요(FSDP의 full_tensor() all-gather, Megatron export 등). 일부 랭크가 건너뛰면 교착(deadlock)이 일어날 수 있죠.

각 프로세스의 rank를 init info에 직접 넘기세요. 전역 프로세스 그룹에서 읽지 않는데, FSDP / TP / PP / EP처럼 여러 그룹이 존재하면 모호해지기 때문이에요.

engine = WeightTransferTrainerFactory.trainer_init(
    init_info=NCCLTrainerInitInfo(..., rank=torch.distributed.get_rank()),
    client=client,
    source=ModuleSource(model),
)
engine.send_weights()  # called on every rank

API 엔드포인트 (API Endpoints)

vLLM을 HTTP 서버로 실행할 때 가중치 전송을 위한 다음 엔드포인트를 사용할 수 있어요. HTTPVLLMWeightSyncClient가 처음 네 개를 대신 호출해줍니다.

참고: HTTP 가중치 전송 엔드포인트는 VLLM_SERVER_DEV_MODE=1이 설정되어야 해요.

Rust 프론트엔드의 선택적 gRPC Control 서비스는 신뢰할 수 있는 sidecar를 위해 동일한 pause, sleep, weight-transfer, weight-version 생애주기를 노출합니다. ServerInfo.rl_capabilities 응답은 weight transfer와 sleep mode가 구성되었는지 보고해요. 백엔드별 init_info와 update_info는 JSON 메타데이터로 남고, 모델 텐서는 계속 구성된 NCCL, IPC, sparse-NCCL 또는 sharded-RDT 전송을 통해 이동합니다.

시스템 확장 (Extending the System)

시스템의 모든 조각을 교체할 수 있어요. 보내는 가중치(WeightSource), 컨트롤 플레인 전송(VLLMWeightSyncClient), 그리고 전송 자체(각자 팩토리 레지스트리를 가진 두 엔진 ABC)까지요. 자세한 내용은 Base Classes 문서를 참고하세요.

더 알아보기 (Learn more)

• Base Classes와 커스텀 엔진 (Base Classes and Custom Engines)
• NCCL 엔진
• IPC 엔진
• Sharded RDT
• 비동기 강화학습 (Async RL)