가중치 전송

가중치 전송 (Weight Transfer)

vLLM은 강화학습(RL) 워크플로우에서 트레이닝 프로세스의 모델 가중치를 추론 엔진으로 동기화하는 플러그형 가중치 전송 시스템을 제공합니다. RLHF, GRPO, 그리고 정책 모델이 학습 중 반복 갱신되고 그 갱신 가중치가 rollout 생성을 위해 추론 엔진에 반영되어야 하는 다른 온라인 RL 방법에 필수적입니다.

출처: 문서

본문

vLLM은 RL 워크플로우에서 트레이닝 프로세스의 모델 가중치를 추론 엔진으로 동기화하는 플러그형 가중치 전송 시스템을 제공합니다. 이는 정책 모델이 학습 중 반복적으로 갱신되고, 그 갱신 가중치가 rollout 생성을 위해 추론 엔진에 반영되어야 하는 RLHF, GRPO 및 기타 온라인 RL 방법에 필수적입니다.

아키텍처 (Architecture)

가중치 전송에는 프로세스당 하나씩 두 개의 엔진이 있으며 대칭적입니다:

트레이너 엔진은 상태를 가집니다(stateful): 자체 communicator와 wire 파라미터를 소유하고, WeightSource에서 가중치를 가져와 VLLMWeightSyncClient를 통해 추론 측을 구동합니다. 트레이너 코드는 전송 방식을 알 필요가 없고, 매 라운드 전송 상태를 다시 전달할 필요도 없습니다 — 동기화마다 send_weights() 호출 한 번이면 됩니다.

내부적으로 매 라운드는 같은 네 단계 프로토콜이며, 트레이너 엔진이 대신 구동합니다:

1. 초기화 (init_weight_transfer_engine): 트레이너와 추론 워커 사이 통신 채널을 설정. 학습 루프 전에 trainer_init에서 한 번 호출됨.
2. 시작 (start_weight_update): 가중치 업데이트를 위해 추론 엔진을 준비.
3. 가중치 업데이트 (update_weights): 갱신 가중치 전송. 한 번 또는 여러 번 호출 가능(예: 청크 전송).
4. 완료 (finish_weight_update): 업데이트 종료(예: checkpoint 형식 가중치의 후처리 실행). 모든 가중치 전송 후 한 번 호출됨.

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

빠른 시작 (Quickstart)

추론 측 (Inference Side)

추론 측은 백엔드 이름만 받습니다. 전송의 나머지는 트레이너가 결정하고 init handshake로 보냅니다.

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의 경우 트레이너 broadcast와 동시에 워커의 update_weights를 실행해야 하며, 둘 다 같은 NCCL 호출 안에서 rendezvous합니다).

트레이너 측에는 backend= 인자가 없습니다: 각 TrainerInitInfo 서브클래스가 자체 backend를 선언하고, 팩토리가 이를 디스패치합니다.

각 설정이 어디에 있는가 (Where Each Setting Lives)

트레이너는 양쪽이 동의해야 하는 모든 것에 대한 단일 진실 원천(source of truth)입니다. 이것이 기억해야 할 핵심 규칙입니다:

참고

Wire 파라미터는 의도적으로 WeightTransferConfig에도, 라운드별 update info에도 없습니다. 트레이너가 init 시에 보내고 워커가 들은 대로 읽으므로, 양쪽의 불일치한 packed 플래그는 단지 권장되지 않는 것이 아니라 표현 불가능합니다.

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

샤딩된 트레이너(FSDP, TP/PP/EP)의 경우 모든 트레이너 랭크가 엔진을 만들고 send_weights()를 호출합니다. 랭크 0이 송신자입니다: communicator를 보유하고, 클라이언트와 통신하며, 바이트를 wire에 올리는 것은 오직 랭크 0뿐입니다. 비송신자 랭크도 WeightSource를 반복합니다. 파라미터 물질화는 보통 자체가 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 엔드포인트

vLLM을 HTTP 서버로 실행할 때 가중치 전송용 다음 엔드포인트가 제공됩니다. HTTPVLLMWeightSyncClient는 처음 네 개를 대신 호출해 줍니다.

참고

HTTP 가중치 전송 엔드포인트는 VLLM_SERVER_DEV_MODE=1 설정이 필요합니다.

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

시스템 확장 (Extending the System)

시스템의 모든 조각은 교체 가능합니다: 보내는 가중치(WeightSource), 컨트롤 플레인 전송(VLLMWeightSyncClient), 그리고 전송 자체(두 엔진 ABC, 각각 자체 팩토리 레지스트리 보유). Base Classes를 참고하세요.

더 알아보기 (Learn more)

• Base Classes and Custom Engines — 네 가지 추상화와 커스텀 엔진
• 가중치 전송: NCCL — NCCL 엔진
• 가중치 전송: IPC — IPC 엔진
• 가중치 전송: sharded_rdt — sharded RDT 엔진
• 비동기 RL — pause/resume API