Sharded RDT 엔진

Sharded RDT 엔진 (Sharded RDT Engine)

sharded RDT 가중치 전송 엔진은 Ray Direct Transport(RDT, Ray의 제로 카피 텐서 전송)를 사용해 NIXL 위에서 가중치를 point-to-point로 이동합니다. pull 기반입니다: 추론 워커가 모든 전송을 시작하고, 각자 tensor·expert parallelism 하에서 자신이 소비하는 슬라이스만 요청합니다. 따라서 대형 MoE 모델은 broadcast가 소비하는 total_bytes 대신 워커당 대략 total_bytes / num_workers를 이동합니다.

출처: 문서

본문

Sharded RDT를 언제 사용하나 (When to Use Sharded RDT)

전체 파라미터 broadcast가 병목인 매우 큰 모델 — 보통 expert parallelism으로 서빙되는 MoE, 각 워커가 전문가의 작은 일부를 소유
트레이너와 추론이 별도 GPU에서, NIXL이 지원하는 패브릭(InfiniBand, RoCE, EFA) 위에서
파이프라인 병렬 트레이너처럼 랭크가 모델의 일부만 보유하는 샤딩된 트레이너

요구사항:

distributed_executor_backend="ray" — 워커가 Ray 액터여야 함
트레이너와 워커 모두 Ray >= 2.56.0
트레이너·워커가 공유하는 환경에 nixl 설치
지원 연산 집합(아래) 안에 머무는 가중치 로더
EPLB(enable_eplb=true)는 거부됨: 런타임에 전문가를 재배열해 기록된 계획을 무효화

동작 방식 (How It Works)

슬라이스는 vLLM 자체 가중치 로더를 통해 추적됩니다

가중치 로더는 보통 전체 HF-format 텐서를 받아 이 워커가 필요로 하는 부분을 잘라냅니다. 엔진은 그 대신 FakeRDTTensor를 건네줍니다: .shape/.dtype/.size()는 답하지만 데이터는 없는 제로 스토리지 텐서입니다. 로더가 호출하는 모든 view 또는 slice 연산은 새 fake를 반환하고 그 연산을 기록된 체인에 추가하며, copy_가 그 체인을 끝내는 싱크입니다.

그 체인이 wire 형식입니다. ("model.layers.0.w", (("narrow", (0, 512, 512), ()), ("t", (), ())))는 트레이너에게 "이 텐서를 가져서 narrow하고, transpose하고, 결과를 보내라"고 말합니다. 트레이너는 getattr(tensor, op)(*args, **kwargs)로 재생합니다.

발견(discovery)은 비싸므로 한 번만, init_transfer_engine에서, 모든 파라미터를 meta에 둔 채 model.load_weights에 대한 dry run으로 발생합니다. 아무것도 전송되지 않고, 엔진은 리프 모듈별로 어떤 슬라이스가 어떤 목적지 영역을 공급하는지 기록만 합니다. 이후 모든 동기화는 순수 재생입니다.

로더가 실제 데이터를 필요로 하는 것(산술, .to(), .float(), .item(), .data, bool-mask 인덱싱)은 모두 allowlist 밖이라 init에서 예외를 던집니다. 이것은 의도적입니다: 설정 중 크게 실패하는 것이 잘못된 바이트를 조용히 전송하는 것보다 낫습니다. sharded_rdt_common.py의 SUPPORTED_OPS는 양쪽이 도출하는 단일 테이블이므로, 기록기와 재생기가 어긋날 수 없습니다.

수신 슬라이스는 레이어별 재로드 버퍼에 직접 들어갑니다

엔진은 start_weight_update/finish_weight_update에서 layerwise reload를 자체 구동합니다. dry run이 각 목적지를 파라미터의 as_strided 영역으로 이미 기록했으므로, 도착하는 슬라이스는 재로드 중인 레이어에 바로 복사됩니다 — 워커에서 전체 HF 텐서가 결코 물질화되지 않고, load_weights의 두 번째 패스도 실행되지 않습니다. 각 레이어는 마지막 슬라이스가 도착하는 즉시 양자화되어 지속 커널 저장소에 복사됩니다.

Gather와 pull은 파이프라이닝되며, 이것이 gather_lookahead가 한정하는 것입니다

트레이너는 보통 파라미터를 그 자리에서 서빙할 수 없습니다: FSDP가 샤딩하고, EP로 분할된 트레이너조차 전체 전문가를 조립해야 합니다. 따라서 각 동기화는 여전히 gather collective를 실행하지만 — 모델 단위가 아니라 레이어 단위로 합니다.

gather 그룹은 decoder 레이어 하나입니다. 파라미터 목록은 각 이름의 가장 바깥쪽 인덱스 세그먼트로 키잉됩니다. 인덱스 없는 이름들의 연속 — 첫 레이어 앞의 임베딩, 마지막 뒤의 final 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

레이어는 뒤따르는 모든 것의 단위입니다: 트레이너가 레이어를 gather하고, 게시하며(즉시 pull 가능), 소비자가 방금 게시한 것을 pull하는 동안 다음 레이어로 진행합니다. 모든 소비자가 레이어에 대한 처리를 끝냈다고 신호를 보내면 트레이너는 그것을 버리고 다른 레이어를 gather할 크레딧을 얻습니다. gather_lookahead는 소비자보다 그 루프가 얼마나 앞서 실행될 수 있는지로, 한 번에 최대 gather_lookahead + 1개 레이어가 트레이너에 상주합니다. 기본값 1은 현재 레이어가 pull되는 동안 다음 레이어를 gather된 채 pull 가능하게 유지합니다 — 트레이너 메모리를 두 배로 하지 않고 핸드오프를 숨기기에 충분합니다. 한 레이어의 gather가 pull보다 느릴 때만 높이세요.

레이어는 소비자가 해제하는 단위이자 수신 버퍼의 크기가 맞춰지는 단위이기도 해서, 양쪽에서 메모리를 한정합니다: 이것이 없으면 전체 모델이 한 번의 전송이 되고 양쪽이 그 전체 몫을 동시에 보유해야 합니다. 고정 model.layers. 접두어가 아니라 인덱스로 키잉하는 것이 이 성립을 이름 규약 전반에 걸쳐 만드는 이유입니다 — VLM의 model.language_model.layers., GPT-2의 transformer.h., 비전 타워의 visual.blocks.. 소스가 파티션을 제어할 수 있습니다. gather groups 참고.

소유권 (Ownership)

트레이너 랭크는 전체 모델을 보유할 필요가 없습니다. 각 랭크는 WeightSource.held_names()로 자기가 보유한 것을 선언하고, fleet이 trainer_init에서 그 선언들을 all-gather하며, 소비자는 각 pull을 실제로 그 이름을 보유한 랭크로 라우팅합니다. 파이프라인 스테이지, expert parallelism, 둘의 조합 모두 같은 선언입니다. 소비자는 이름을 보유한 랭크에 걸쳐 pull을 분산하므로 단일 트레이너 NIC가 병목이 되지 않습니다.

추론 측 (Inference Side)

from vllm import LLM
from vllm.config import WeightTransferConfig

llm = LLM(
    model="my-model",
    weight_transfer_config=WeightTransferConfig(backend="sharded_rdt"),
    distributed_executor_backend="ray",
)

vllm serve my-model \
  --distributed-executor-backend ray \
  --weight-transfer-config '{"backend": "sharded_rdt"}'

그 외의 모든 것 — 어떤 프로듀서가 존재하는지, 모델이 레이어 그룹으로 어떻게 나뉘는지, 소유권 테이블 — 은 init handshake에서 트레이너가 보냅니다.

gpu_memory_utilization을 정하기 전에 수신 버퍼를 측정하세요

각 워커는 num_rdt_buffers 수신 버퍼를 보유하며, 각각이 pull하는 가장 큰 단일 슬라이스 배치를 담을 크기입니다. NCCL·NIXL 내부처럼 이들은 gpu_memory_utilization에 계산되지 않으므로, 여유 공간이 없는 비율은 엔진이 정상 기동하더라도 첫 동기화에서 OOM이 납니다. 버퍼 크기는 가장 큰 원자 슬라이스에 의해 결정됩니다 — untied vocab 행렬을 슬라이스하지 않고 보유하는 워커라면 그게 전체 임베딩입니다.

트레이너 측 (Trainer Side)

from vllm.distributed.weight_transfer import (
    ModuleSource,
    HTTPVLLMWeightSyncClient,
    WeightTransferTrainerFactory,
)
from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
    ShardedRDTTrainerInitInfo,
)

engine = WeightTransferTrainerFactory.trainer_init(
    init_info=ShardedRDTTrainerInitInfo(
        rank=rank,                                # rank 0 is the sender
        num_consumers=8,                          # inference workers, fleet-wide
        trainer_actor_namespace="my_namespace",   # must be visible to the workers
    ),
    client=HTTPVLLMWeightSyncClient("http://localhost:8000"),
    source=ModuleSource(model),
)

engine.send_weights()   # once per sync, on every trainer rank

trainer_init과 send_weights는 모든 트레이너 랭크에서 실행됩니다: 각자가 serve 액터 하나를 소유하고 gather에 참여하는 반면, rank 0만 추론 측 handshake를 구동합니다. 어떤 VLLMWeightSyncClient든 동작합니다.

plain nn.Module이 아닌 트레이너(Megatron export, 원시 샤딩 checkpoint)를 적응시키려면 WeightSource를 서브클래싱하세요.

ShardedRDTTrainerInitInfo

필드	기본값	설명
`rank`	—	키워드 전용. 이 트레이너 랭크; 0이 송신자
`num_consumers`	—	전체 fleet의 추론 워커 수(TP × DP)
`trainer_actor_namespace`	`None`	serve 액터용 Ray 네임스페이스; 워커가 여기서 이름으로 해석
`num_rdt_buffers`	2	양쪽 링 깊이
`buffer_presize_gb`	0.0	각 버퍼 슬롯 사전 크기(GiB). 가장 큰 원자 슬라이스를 덮도록 설정
`gather_lookahead`	1	gather 루프가 앞서 실행하는 gather-되었지만 해제 안 된 레이어 수
`stall_timeout_s`	300.0	진행 없이 이 시간이 지나면 동기화 실패. mid-sync에 죽는 소비자용 liveness backstop이지 지연 목표는 아님

예시 (Examples)

Small MoE on 4 GPUs — 2 FSDP2 트레이너 랭크 → expert parallelism을 가진 2 vLLM DP 랭크, 한 노드. 트레이너 fleet을 별도 추론 fleet과 짝지으며, 이 백엔드가 지원하는 유일한 배열입니다. 그리고 동기화가 가중치를 옮겼고 두 번째 동기화가 생성을 바꾸지 않음을 단언하므로 CI에서 무인 실행됩니다

가중치를 실제로 만들기에 충분한 최소 FSDP2만으로 트레이너를 의도적으로 작게 유지해, 파일이 트레이너가 아니라 가중치 동기화에 집중하게 합니다. 전체 RL 트레이너의 경우 SkyRL이 FSDP와 함께 Megatron(PP-local gather, MoE용 expert-stack fusion)으로 이 백엔드를 통합합니다: NovaSky-AI/SkyRL#1753.

더 알아보기 (Learn more)

가중치 전송 개요 — 플러그형 백엔드 시스템
가중치 전송: NCCL — NCCL 엔진
가중치 전송: IPC — IPC 엔진
Base Classes — 커스텀 소스와 gather groups