기반 클래스와 커스텀 엔진

기반 클래스와 커스텀 엔진 (Base Classes and Custom Engines)

가중치 전송 시스템은 네 개의 추상화로 구성되며 각각 독립적으로 교체 가능합니다. 어떤 형태의 트레이너 가중치든 WeightSource 어댑터로 추출하고, VLLMWeightSyncClient로 엔진에 도달하며, 트레이너·추론 양쪽 엔진이 각각 전송·수신을 담당합니다.

출처: 문서

본문

가중치 전송 시스템은 각각 독립적으로 교체 가능한 네 가지 추상화로 구성됩니다:

추상화	측	담당
WeightSource	트레이너	어떤 가중치를 보낼지 (What)
VLLMWeightSyncClient	트레이너	추론 엔진에 어떻게 도달할지 — RL 스택의 자체 vLLM 래퍼용 어댑터
TrainerWeightTransferEngine	트레이너	바이트를 어떻게 전송할지 (How to transmit)
WeightTransferEngine	추론	바이트를 어떻게 수신하고 로드할지 (How to receive)

두 엔진은 두 개의 별도 팩토리, WeightTransferTrainerFactory와 WeightTransferEngineFactory에 등록됩니다. 관례상 백엔드 이름을 공유하지만, 트레이너 프로세스는 워커 엔진을 인스턴스화하거나 그 반대를 하지 않으므로 레지스트리는 독립적으로 유지됩니다.

트레이너 측 (Trainer Side)

WeightSource

이것은 트레이너 가중치가 어떤 형태든 그에 맞춘 어댑터입니다. WeightSource는 특정 프레임워크에서 가중치를 추출하는 방법을 정의합니다.

엔진이 받는 것은 항상 동일합니다: HF 형식 파라미터 이름과 이미 전체(unsharded) 형태로 물질화된 텐서. 필요한 gathering·re-fusing·dequantizing·renaming은 모두 소스 안에 속합니다.

WeightSource는 **재반복 가능(re-iterable)**하며 두 필수 채널이 있습니다:

• metadata() -> list[ParamMeta] — 아무것도 전송하지 않고 모든 파라미터의 이름·wire dtype·전체 형태. 형태가 로컬로 알려져 있으면 저렴(FSDP DTensor는 전역 형태를 앎); 형태를 알기 위해 물질화해야 하는 프로듀서(Megatron-Bridge export)에서는 첫 호출이 비쌀 수 있으므로 캐시해야 함.
• 반복(iteration) — 완전히 물질화된 (name, tensor) 쌍을 한 번에 하나씩 산출.

@dataclass(frozen=True)
class ParamMeta:
    name: str
    dtype: torch.dtype
    shape: tuple[int, ...]

두 채널은 요소별로 일치해야 합니다

metadata()는 반복이 산출할 것을 정확히 선언해야 합니다: 같은 파라미터, 같은 순서, 같은 dtype·형태. 이것은 특정 백엔드가 아니라 ABC의 불변식입니다 — 두 채널 사이에서 파라미터를 재정렬·생략·각 dtype 변경하는 소스는, 테스트하는 백엔드가 심지어 알아차리지 못해도 손상된 것입니다. 백엔드는 두 채널을 모두 읽고 일치를 신뢰해도 됩니다. dense NCCL은 그렇게 하며 강제합니다.

물질화는 보통 collective이므로 모든 트레이너 랭크가 같은 소스를 같은 순서로, lockstep으로 반복해야 하며, 그렇지 않으면 랭크가 deadlock됩니다. metadata() 자체도 커스텀 프로듀서의 collective일 수 있으므로 모든 랭크에서 실행됩니다 — 결과를 보내는 것은 송신자뿐입니다.

iter(source)는 매 라운드 새로운 패스를 산출해야 합니다.

ModuleSource

ModuleSource(module)는 module.named_parameters() 위에서의 일반적인 경우입니다. plain·FSDP-sharded 모듈을 특별 케이스 없이 처리합니다: 반복은 full_tensor()로 각 DTensor를 all-gather하는 반면, metadata()는 전역 .shape/.dtype을 읽어 gather를 결코 트리거하지 않습니다.

from vllm.distributed.weight_transfer import ModuleSource

source = ModuleSource(model)

커스텀 소스 (Custom sources)

가중치가 HF 형식이 되기 위해 작업이 필요할 때(프레임워크 특정 export, re-fusing 단계, dtype 캐스트) WeightSource를 서브클래싱하세요.

from vllm.distributed.weight_transfer import ParamMeta, WeightSource

class MegatronBridgeSource(WeightSource):
    """Megatron model -> HF names, via a bridge that gathers TP/PP/EP internally
    and returns full tensors on every rank."""

    def __init__(self, bridge, module, dtype):
        self._bridge, self._module, self._dtype = bridge, module, dtype
        self._meta: list[ParamMeta] | None = None

    def _export(self):
        return self._bridge.export_hf_weights(self._module)

    def metadata(self) -> list[ParamMeta]:
        # Cache: for producers that must materialize to learn shapes, this is
        # the expensive channel. Runs on every rank (it may be a collective).
        if self._meta is None:
            self._meta = [
                ParamMeta(name, self._dtype, tuple(t.shape))
                for name, t in self._export()
            ]
        return self._meta

    def __iter__(self):
        # Must yield exactly what metadata() declared, in the same order.
        for name, tensor in self._export():
            yield name, tensor.to(self._dtype).detach().contiguous()

held_names(): 부분 소유권 (partial ownership)

기본적으로 모든 랭크가 모든 파라미터를 생산할 수 있다고 가정합니다 — 위 소스가 브리지를 통해 모든 병렬성을 gather한 뒤 산출하므로 그렇습니다. 단순하고 항상 옳지만, gather 비용을 모든 랭크가 전액 지불한다는 뜻입니다.

랭크가 분할되어 각자 모델의 일부만 소유할 때 held_names()를 선택적으로 재정의할 수 있습니다. 이 랭크가 소유한 파라미터 이름(또는 기본값인 None, 전부)을 반환합니다:

def held_names(self):
    # This pipeline stage's layers, and within them only this EP rank's experts.
    return self._my_stage_names - self._foreign_expert_names

이는 다양한 트레이너 레이아웃을 다룹니다 — 파이프라인 스테이지(랭크가 일부 레이어 소유), expert parallelism(랭크가 일부 전문가 소유), 둘 다, 또는 어느 쪽에도 맞지 않는 형태. 파라미터별 라우팅이 가능한 백엔드(sharded RDT 참고)가 실제로 소유한 랭크에서 각 이름을 풀링합니다.

재정의 시 세 가지 요구사항이 따릅니다:

• metadata()는 모든 랭크에서 여전히 전체 모델을 기술해야 합니다. 송신자의 metadata만 추론 측에 도달하므로, 자체 몫만 보고한 랭크는 나머지를 조용히 전송되지 않은 채로 남깁니다. Sharded RDT는 init 시 랭크 간에 이를 교차 검증합니다.
• 모든 이름은 최소 하나의 랭크가 소유해야 합니다. 그렇지 않으면 결코 서빙될 수 없습니다. 엔진은 첫 orphan의 이름을 밝히며 init에서 예외를 던집니다.
• 반복은 이 랭크가 소유하지 않은 이름에 대해 None을 산출합니다. 이름은 metadata 순서로 여전히 나타나 순서 검사가 랭크 간 정렬을 유지합니다 — 데이터만 없을 뿐입니다. 이름을 주장하고 None을 산출하는 것은 엔진이 이름을 밝히며 보고하는 오류입니다.

부분 소유권은 sharded RDT에서만 동작

파라미터별 라우팅을 하는 백엔드만 held_names()를 존중합니다. Broadcast 백엔드는 이를 무시하고 모든 랭크에서 모든 이름을 보내므로, 거기서 부분 소유권을 선언해도 아무 변화가 없습니다.

Gather 그룹 (Gather groups)

일부 백엔드는 모델 단위가 아니라 레이어 단위로 전송하므로 metadata()를 gather group으로 분할합니다. layerwise_groups는 각 이름을 그것이 담긴 가장 바깥쪽 인덱스 세그먼트로 키잉하므로, 그룹 하나가 decoder 레이어 한 개이며, 인덱스 없는 이름들의 연속(임베딩, 마지막 norm, lm_head)은 나타나는 곳마다 자체 그룹을 이룹니다:

group 0     model.embed_tokens.weight
group 1     model.layers.0.*          <- one decoder layer
group 2     model.layers.1.*
...
group N+1   model.norm.weight, lm_head.weight

리터럴 접두어 대신 인덱스로 키잉한다는 것은 아키텍처별 테이블이 없다는 뜻입니다: model.layers.0., model.language_model.layers.0.(최근 Qwen 텍스트 checkpoint), transformer.h.0.(GPT-2, Falcon), backbone.layers.0.(Mamba), 그리고 비전 타워의 visual.blocks.0. 모두 같은 방식으로 분할됩니다. 취해지는 인덱스는 가장 바깥쪽 것이라 MoE 레이어가 온전히 유지됩니다 — model.layers.3.mlp.experts.7.w1 같은 per-expert 이름은 전문가가 아니라 레이어로 키잉됩니다.

그룹 인덱스 g는 모든 랭크와 모든 소비자에서 같은 레이어를 의미합니다. 모든 측이 한 랭크의 metadata()에서 도출하기 때문입니다. 그 일치 덕분에 백엔드는 버퍼를 한 레이어로 한정하고 모두가 끝나면 레이어를 해제할 수 있습니다.

리프 모듈의 소스는 모두 한 그룹에 있어야 합니다

sharded-RDT 엔진은 마지막 청크가 도착하는 즉시 그룹을 해제하므로, 그룹 간에 분할된 모듈은 stall watchdog이 발동할 때까지 pull을 정체시킵니다. 기본 파티션은 이를 보장합니다. groups() 재정의는 이를 유지해야 합니다.

여기서 두 개의 훅이 나오며, 둘 다 동작하는 기본값이 있습니다:

• groups() — metadata 순서의 이 랭크 그룹. 기본은 layerwise_groups(metadata())를 소유한 이름이 하나 이상 있는 그룹으로 제한한 것. 여기서 소유한 게 없는 그룹은 완전히 건너뜁니다.
• iter_groups() — 같은 스트림을 한 번에 한 그룹씩 배치. 기본은 __iter__를 구동하고 그 출력을 배치하며 이름이 metadata 순서로 도착하는지 검사. 프레임워크가 전체 그룹을 한 단계로 생산할 수 있을 때 재정의하세요: 물질화는 보통 collective이고, 텐서 단위가 아니라 그룹 단위로 구동하면 per-expert MoE 모델에서 ~37k generator 재개가 ~95개로 줄어듭니다.

metadata() 순서가 파티션을 정의하므로, 레이어 인덱스를 공유하는 모든 이름은 그 안에서 연속이어야 합니다. 자연스러운 export 순서가 레이어를 인터리브하는 소스(예: 모든 MoE 전문가를 함께 버킷)는 반환 전에 재정렬해야 합니다.

VLLMWeightSyncClient

이것은 RL 스택이 vLLM에 도달하는 방식에 대한 어댑터입니다. 많은 RL 프레임워크가 추론 엔진을 자체 추상화로 감싸고, 각자 vLLM에 다른 방식으로 도달합니다. VLLMWeightSyncClient는 그 특수한 형태가 적응되는 단일 이음매(seam)이므로, 가중치 동기화 엔진은 컨트롤 플레인에 무관하게 유지됩니다.

계약은 이것뿐입니다: 래퍼가 어떤 형태든 반드시 같은 네 호출로 귀결되어야 합니다 — 설정 시 init_weight_transfer_engine 한 번, 그 다음 라운드마다 start_weight_update → 하나 이상의 update_weights → finish_weight_update. 트레이너 엔진이 추론 측에서 필요로 하는 모든 것이 이들을 통해 갑니다.

class VLLMWeightSyncClient(Protocol):
    def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: ...
    def start_weight_update(self) -> None: ...
    def update_weights(self, update_info: dict[str, Any]) -> None: ...
    def finish_weight_update(self, weight_version: str | None = None) -> None: ...

그것은 @runtime_checkable 구조적 Protocol(PEP 544)이라 적응이 저렴합니다: 네 메서드가 있는 어떤 객체든 이미 그것을 충족합니다. 프레임워크의 기존 래퍼는 보통 네 개의 전달(forwarding) 메서드를 추가하면 클라이언트가 됩니다.

vLLM과 함께 제공되는 두 구현:

클라이언트	수신 대상
RayVLLMWeightSyncClient(handle)	하나 이상의 AsyncLLM/LLM Ray 액터. 리스트를 받아 각 호출을 모든 핸들에 전파하고 전부에서 블록하므로, 멀티 액터(예: 멀티-DP) 배포를 한 단위로 구동
HTTPVLLMWeightSyncClient(base_url, timeout=300)	RLHF HTTP 라우트 위의 vLLM 서버

커스텀 가중치 동기화 클라이언트는 다음과 같이 구현할 수 있습니다:

class MyFrameworkWeightSyncClient:
    """Adapts one RL framework's rollout pool to the four weight-sync calls."""

    def __init__(self, rollout_pool):
        self.pool = rollout_pool          # whatever your stack already has

    def init_weight_transfer_engine(self, init_info):
        # Fan out to every replica and block: all of them receive weights.
        self.pool.broadcast_rpc("init_weight_transfer_engine", init_info=init_info)

    def start_weight_update(self):
        self.pool.broadcast_rpc("start_weight_update")

    def update_weights(self, update_info):
        self.pool.broadcast_rpc("update_weights", update_info=update_info)

    def finish_weight_update(self, weight_version=None):
        self.pool.broadcast_rpc("finish_weight_update")
        if weight_version is not None:
            self.pool.broadcast_rpc("update_weight_version", weight_version)

어떤 어댑터에서도 올바르게 해야 할 두 가지:

• 모든 복제본에 도달하고, 모두 끝날 때까지 블록하세요. 가중치 업데이트는 로드밸런싱된 요청이 아닙니다: 모델 복사본을 보유한 모든 워커가 받아야 합니다. 모두 끝나기 전에 반환하면 트레이너가 아직 로딩 중인 워커보다 앞서나갑니다. (내장 클라이언트 둘 다 이렇게 합니다 — Ray는 핸들에 전파하고, HTTP는 서버의 DP 클라이언트가 내부적으로 broadcast하기 때문.)
• 실패 시 예외를 던지세요. 트레이너 엔진은 추론 측 오류를 드러내기 위해 예외에 의존합니다. 예외를 삼키는 클라이언트는 실패한 동기화를 조용히 낡은 가중치로 만들거나, 전송이 워커와 rendezvous하는 백엔드에서 hang으로 만듭니다.

참고

HTTP는 원시 CUDA IPC 핸들을 운반할 수 없으므로, HTTPVLLMWeightSyncClient는 이를 피클링하고 base64 인코딩해 ipc_handles_pickled 필드에 넣습니다. 워커는 VLLM_ALLOW_INSECURE_SERIALIZATION=1일 때만 역직렬화합니다. 페이로드가 JSON 네이티브인 백엔드(NCCL)는 그대로 통과합니다.

TrainerWeightTransferEngine

트레이너 측 엔진: 전송 상태(NCCL communicator, IPC 디바이스 정보, 전송 계획)를 보유하고, WeightSource에서 가중치를 끌어오며, VLLMWeightSyncClient를 통해 추론 측을 구동합니다. init info 타입에 대해 제네릭이며, trainer_init 클래스메서드 팩토리로 만들어지고, send_weights()로 구동됩니다.

메서드	설명
trainer_init(init_info, *, client, source=None)	클래스메서드. 추론 측과 rendezvous하고 준비된 인스턴스 반환
send_weights()	가중치를 밀어 넣고 전체 업데이트 라운드 트립 구동
shutdown()	communicator/프로세스 그룹 정리. 기본 no-op

trainer_init과 send_weights 모두 모든 트레이너 랭크에서 호출됩니다. is_sender는 trainer_init에서 init_info.rank로 한 번 해석됩니다. 각 엔진은 모든 랭크에서 실제 클라이언트를 보유하지만 컨트롤 플레인 RPC와 전송을 self.is_sender로 가드합니다. 그래서 wire에 닿는 것은 송신자뿐입니다. 비송신자 랭크도 모든 collective를 실행해 그룹 정렬을 유지합니다.

트레이너 측은 WeightTransferConfig를 받지 않습니다. 백엔드는 init info의 backend ClassVar에서 오고, wire 파라미터도 init info에 실립니다.

TrainerInitInfo

위 trainer_init에 전달되는 init_info. 호출자가 전송을 구성하는 방법입니다: 백엔드를 선택하고, 이 프로세스가 어떤 랭크인지 말하며, wire 파라미터를 운반합니다. 각 백엔드가 이를 서브클래싱하며, 기본 클래스는 모든 백엔드가 필요로 하는 필드 하나를 보유합니다.

@dataclass
class TrainerInitInfo:
    backend: ClassVar[str]        # factory dispatch key
    rank: int = field(kw_only=True)

    @property
    def is_sender(self) -> bool:
        return self.rank == 0

• **rank**는 이 트레이너 프로세스의 랭크로 명시적으로 제공됩니다. 엔진은 전역 프로세스 그룹에서 읽지 않습니다. 여러 그룹(FSDP / TP / PP / EP)이 존재하면 모호하기 때문입니다. 랭크 0은 항상 송신자입니다 — trainer_init이 is_sender로 해석하는 것이 바로 이것입니다. 키워드 전용이라 백엔드 서브클래스가 위치 필드를 자유롭게 추가할 수 있습니다.
• **backend**는 __init__ 필드가 아니라 ClassVar입니다: 팩토리가 디스패치에 읽는 백엔드별 고정 상수라, 호출자가 backend= 인자를 절대 전달하지 않는 이유입니다. 모든 서브클래스는 이를 반드시 설정해야 하며 — __init_subclass__가 그렇지 않으면 예외를 던집니다.

서브클래스는 전송의 wire 파라미터(packed, 버퍼 크기)도 운반합니다. 송신자는 trainer_init 안에서 이를 워커에 전파하므로 양쪽이 불일치할 수 없습니다. 구체 필드는 NCCLTrainerInitInfo와 IPCTrainerInitInfo를 참고하세요.

Full-Resync vs. Delta 백엔드

source는 선택적이며, 백엔드를 두 형태로 나눕니다:

• Full resync (NCCL, IPC) — 안정적인 WeightSource가 trainer_init에서 고정되고 매 라운드 재반복됩니다. send_weights()는 인자를 받지 않습니다. 이 백엔드들은 스스로 source가 non-null임을 검증합니다.
• Delta (sparse NCCL) — 페이로드가 매 라운드 달라져 안정적인 소스가 없습니다. 엔진은 source를 받지 않고, 각 라운드 페이로드가 send_weights(patches)로 직접 전달됩니다.

커스텀 트레이너 엔진 구현

from dataclasses import dataclass
from typing import ClassVar

from typing_extensions import Self

from vllm.distributed.weight_transfer.base import (
    TrainerInitInfo,
    TrainerWeightTransferEngine,
    VLLMWeightSyncClient,
    WeightSource,
)

@dataclass
class MyTrainerInitInfo(TrainerInitInfo):
    backend: ClassVar[str] = "my_backend"

    endpoint: str
    chunk_size_bytes: int = 256 * 1024 * 1024   # a wire param: shipped to the worker

class MyTrainerWeightTransferEngine(TrainerWeightTransferEngine[MyTrainerInitInfo]):
    init_info_cls = MyTrainerInitInfo

    def __init__(self, *, client, source, is_sender=True, chunk_size_bytes=0):
        super().__init__(client=client, source=source, is_sender=is_sender)
        self.chunk_size_bytes = chunk_size_bytes

    @classmethod
    def trainer_init(
        cls,
        init_info: MyTrainerInitInfo,
        *,
        client: VLLMWeightSyncClient,
        source: WeightSource | None = None,
    ) -> Self:
        if source is None:
            raise ValueError("my_backend requires a WeightSource.")
        engine = cls(
            client=client,
            source=source,
            is_sender=init_info.is_sender,
            chunk_size_bytes=init_info.chunk_size_bytes,
        )
        if engine.is_sender:
            # Ship the must-agree wire params so the worker decodes exactly as
            # this trainer encodes, then open the trainer-side endpoint.
            engine.client.init_weight_transfer_engine(
                {"chunk_size_bytes": init_info.chunk_size_bytes}
            )
        return engine

    def send_weights(self) -> None:
        assert self.source is not None
        meta = self.source.metadata()      # every rank: may be a collective
        if not self.is_sender:
            for _ in self.source:          # stay in the trainer-side collective
                pass
            return

        self.client.start_weight_update()
        self.client.update_weights(
            {
                "names": [m.name for m in meta],
                "dtype_names": [str(m.dtype).split(".")[-1] for m in meta],
                "shapes": [list(m.shape) for m in meta],
            }
        )
        for name, tensor in self.source:
            ...                            # transmit
        self.client.finish_weight_update()

올바르게 해야 할 두 가지가 있으며 둘 다 내장 백엔드를 물었습니다:

• 반환 전에 drain하세요. send_weights는 전송이 여전히 진행 중인 채로 반환해서는 안 됩니다. 전송 버퍼를 살려두는 어떤 것도 프레임과 함께 죽고, 추론 측의 finish_weight_update 후처리가 그렇지 않으면 아직 도착하지 않은 가중치를 완료할 수 있습니다.
• 오류 경로에서 컨트롤 플레인 스레드를 join하지 마세요. NCCL처럼 전송과 동시에 사이드 스레드에서 update_weights를 실행하는데 전송이 예외를 던지면, 워커는 여전히 일치하는 collective에서 블록되어 결코 돌아오지 않습니다. 대기 없이 executor를 종료해 실제 예외가 hang 대신 드러나게 하세요.

WeightTransferTrainerFactory

from vllm.distributed.weight_transfer import WeightTransferTrainerFactory

# Lazy loading (recommended): the module is imported only when the backend is used
WeightTransferTrainerFactory.register_engine(
    "my_backend",
    "my_package.my_module",
    "MyTrainerWeightTransferEngine",
)

# Or register the class directly
WeightTransferTrainerFactory.register_engine("my_backend", MyTrainerWeightTransferEngine)

engine = WeightTransferTrainerFactory.trainer_init(
    init_info=MyTrainerInitInfo(rank=0, endpoint="..."),  # `backend` selects the engine
    client=client,
    source=source,
)

추론 측 (Inference Side)

WeightTransferEngine

두 dataclass 타입으로 파라미터화된 제네릭 추상 클래스:

• TInitInfo (WeightTransferInitInfo 확장): 백엔드별 초기화 파라미터.
• TUpdateInfo (WeightTransferUpdateInfo 확장): 백엔드별 가중치 업데이트 메타데이터.

서브클래스는 다섯 개 메서드를 구현해야 합니다:

메서드	설명
init_transfer_engine(init_info)	각 추론 워커의 통신 채널 초기화, 트레이너가 제공한 wire 파라미터 기록
start_weight_update()	업데이트 준비(예: layerwise reload 시작); in-place 엔진은 no-op
finish_weight_update()	업데이트 종료(예: layerwise reload 완료); in-place 엔진은 no-op
receive_weights(update_info)	가중치를 받아 self.model에 로드
shutdown()	리소스 정리

기본 클래스가 제공하는 것:

1. __init__, config(WeightTransferConfig), vllm_config(VllmConfig), device(torch.device), model(nn.Module)를 받음.
2. update_weights(update_info_dict), receive_weights의 얇은 래퍼: dict를 타이핑된 dataclass로 파싱하고 receive_weights를 호출하며 디바이스를 동기화 — 엔진이 아래 defers_processing을 설정하지 않는 한.
3. parse_init_info/parse_update_info, API 레벨 dict를 타이핑된 dataclass로 변환하고 잘못된 페이로드에 ValueError를 던짐.
4. set_weight_update_target/reset_weight_update_target, 업데이트를 speculative draft 모델로 재타깃하는 데 사용.

wire 파라미터는 페이로드가 아니라 handshake에서 읽으세요

양쪽이 동의해야 하는 것 — packed, 버퍼 지오메트리 — 은 init info에 도착하며 init_transfer_engine에서 self에 저장한 뒤 receive_weights에서 self에서 읽어야 합니다. 라운드별 update info는 라운드별 메타데이터만 운반합니다. 이것이 트레이너/워커 불일치를 표현 불가능하게 만드는 이유입니다.

defers_processing: 반환된 업데이트가 적용이 아닌 큐잉을 의미할 때

GPU 후처리를 백그라운드 스레드로 파이프라이닝하는 엔진은 update_weights가 디바이스를 동기화하게 놔둘 수 없습니다 — 그 스레드들에서 블록되어 파이프라인을 직렬화하기 때문입니다. 그런 엔진은 클래스 속성 defers_processing = True를 설정하고, 업데이트별 동기화를 생략하며, 대신 finish_weight_update에서 완료를 보장합니다.

finish_weight_update로 진행하는 호출자는 아무것도 할 필요가 없습니다. 엔진이 거기서 drain합니다. 꼬리를 직접 구동하는 호출자 — 자체 finalize_layerwise_reload를 실행하는 — 는 먼저 플래그를 확인하고 drain_pending()을 호출해야 합니다. 플래그가 설정되면 반환된 update_weights는 적용이 아니라 큐잉을 의미하기 때문입니다. drain_pending()은 멱등이며, 동기식으로 처리하는 엔진에서는 no-op이라 항상 호출해도 안전합니다.

Sharded RDT가 이를 설정하는 내장 엔진입니다: 자체 CUDA 스트림을 가진 백그라운드 스레드에서 분산·양자화하므로, drain_pending()이 두 큐를 모두 조인하고 finalize_layerwise_reload가 실행되기 전에 두 스트림을 동기화합니다.

요청 클래스 (Request Classes)

API 레벨 요청 클래스는 일반 딕셔너리를 사용한 백엔드 무관 직렬화를 제공합니다.

from vllm.distributed.weight_transfer.base import (
    WeightTransferInitRequest,
    WeightTransferUpdateRequest,
)

# Init request (dict is converted to backend-specific TInitInfo)
init_request = WeightTransferInitRequest(
    init_info={"master_address": "10.0.0.1", "master_port": 29500, ...}
)

# Update request (dict is converted to backend-specific TUpdateInfo)
update_request = WeightTransferUpdateRequest(
    update_info={"names": [...], "dtype_names": [...], "shapes": [...]}
)

내장 클라이언트를 사용하면 이것들을 직접 만들지 않습니다 — RayVLLMWeightSyncClient가 dict를 감싸고, HTTPVLLMWeightSyncClient가 JSON으로 게시합니다.

LLM/API 레이어에서 speculative draft 모델을 타깃하려면 start_weight_update() 대신 start_draft_weight_update()를 호출하세요. update_weights/finish_weight_update는 변하지 않습니다. 이를 지원할 수 없는 엔진은 supports_draft_weight_update = False를 설정합니다.

커스텀 엔진 구현

1. Info dataclass 정의

from dataclasses import dataclass
from vllm.distributed.weight_transfer.base import (
    WeightTransferEngine,
    WeightTransferInitInfo,
    WeightTransferUpdateInfo,
)

@dataclass
class MyInitInfo(WeightTransferInitInfo):
    endpoint: str
    chunk_size_bytes: int = 256 * 1024 * 1024   # must-agree wire param

@dataclass
class MyUpdateInfo(WeightTransferUpdateInfo):
    names: list[str]
    dtype_names: list[str]
    shapes: list[list[int]]
    # Per-round metadata only.

2. 엔진 구현

class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
    init_info_cls = MyInitInfo
    update_info_cls = MyUpdateInfo

    def init_transfer_engine(self, init_info: MyInitInfo) -> None:
        # Record the trainer's wire params, then set up the connection.
        self.chunk_size_bytes = init_info.chunk_size_bytes
        ...

    def start_weight_update(self) -> None:
        # Checkpoint-format engines: run initialize_layerwise_reload(self.model).
        # In-place engines: no-op
        ...

    def finish_weight_update(self) -> None:
        # Checkpoint-format engines: run finalize_layerwise_reload(...).
        # In-place engines: no-op
        ...

    def receive_weights(self, update_info: MyUpdateInfo) -> None:
        weights = []
        for name, dtype_name, shape in zip(
            update_info.names, update_info.dtype_names, update_info.shapes
        ):
            dtype = getattr(torch, dtype_name)
            weight = self._fetch_weight(name, shape, dtype)
            weights.append((name, weight))
        self.model.load_weights(weights)

    def shutdown(self) -> None:
        # Clean up resources
        ...

3. 팩토리에 등록

from vllm.distributed.weight_transfer import WeightTransferEngineFactory

# Option 1: Lazy loading (recommended for built-in engines)
WeightTransferEngineFactory.register_engine(
    "my_backend",
    "my_package.my_module",
    "MyWeightTransferEngine",
)

# Option 2: Direct class registration
WeightTransferEngineFactory.register_engine(
    "my_backend",
    MyWeightTransferEngine,
)

등록 후 사용자는 WeightTransferConfig(backend="my_backend")로 백엔드를 선택합니다.

WeightTransferEngineFactory

팩토리는 lazy loading이 있는 레지스트리 패턴을 사용합니다. 내장 엔진(nccl, ipc, sparse_nccl, sharded_rdt)은 import 시 등록되지만 해당 모듈은 백엔드가 실제로 요청될 때만 로드됩니다. 불필요할 때 NCCL communicator 같은 무거운 의존성을 import하지 않게 합니다.

from vllm.distributed.weight_transfer import WeightTransferEngineFactory

# Create an engine from config
engine = WeightTransferEngineFactory.create_engine(
    config=weight_transfer_config,
    vllm_config=vllm_config,
    device=device,
    model=model,
)

vLLM은 워커 시작 중에 이것을 대신 호출합니다. 자체 워커에 엔진을 임베드할 때만 직접 필요합니다.

더 알아보기 (Learn more)

• 가중치 전송 개요 — 네 단계 프로토콜과 백엔드
• 가중치 전송: NCCL — NCCL 엔진
• 가중치 전송: IPC — IPC 엔진
• 가중치 전송: sharded_rdt — sharded RDT 엔진