Sharded RDT 엔진

Sharded RDT 엔진 (Sharded RDT Engine)

매우 큰 모델, 특히 전문가 병렬(expert parallelism)로 서빙되는 MoE 모델에서는 전체 파라미터를 broadcast하는 게 병목이 돼요. 모든 워커가 전체 모델이 아니라 자기 슬라이스만 필요하다면, point-to-point로 필요한 것만 받는 게 훨씬 낫겠죠. Sharded RDT 엔진이 바로 그 방식이에요. NIXL 위에서 Ray Direct Transport(RDT)를 통해 가중치를 point-to-point로 이동합니다.

출처: vLLM 공식 문서 — weight_transfer/sharded_rdt

개요

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

Sharded RDT를 언제 쓸까 (When to Use Sharded RDT)

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

요구 사항:

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

어떻게 동작하는가 (How It Works)

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

가중치 로더는 보통 전체 HF 형식 텐서를 받아 이 워커가 필요한 부분을 잘라내요. 엔진은 그걸 하나 보내는 대신 FakeRDTTensor를 건네줍니다. .shape / .dtype / .size()에 답할 수 있지만 데이터는 없는 zero-storage 텐서죠. 로더가 호출하는 모든 view나 slice 연산은 그 연산이 기록된 체인에 추가된 새 fake를 반환하고, copy_가 그것을 끝내는 sink입니다.

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

발견(discovery)은 비싸서, init_transfer_engine에서 한 번, 모든 파라미터를 메타 디바이스에 올린 model.load_weights로 dry run을 수행해 일어나요. 아무것도 전송되지 않고, 엔진은 리프 모듈별로 어떤 슬라이스가 어떤 대상 영역을 먹이는지 기록만 합니다. 이후의 모든 동기화는 순수 재생(pure replay)이에요.

로더가 실제 데이터가 필요한 연산 — 산술, .to(), .float(), .item(), .data, bool-mask 인덱싱 — 을 하면 허용 목록(allowlist) 밖이라 init에서 예외가 발생해요. 이것은 의도적입니다. 설정 중에 크게 실패하는 것이 잘못된 바이트를 조용히 옮기는 것보다 낫죠. sharded_rdt_common.py의 SUPPORTED_OPS는 양쪽이 파생하는 단일 테이블이라서, 기록기(recorder)와 재생기(replayer)가 어긋날 수 없어요.

받은 슬라이스는 레이어별 재로드 버퍼에 직접 들어간다

엔진은 start_weight_update / finish_weight_update에서 레이어별 재로드(layerwise reload)를 스스로 구동해요. dry run이 각 대상을 그 파라미터의 as_strided 영역으로 이미 기록했으므로, 도착하는 슬라이스는 재로드 중인 레이어에 바로 복사됩니다. 워커에서 전체 HF 텐서가 실체화되는 일은 없고, load_weights를 두 번째로 도는 일도 없어요. 각 레이어는 마지막 슬라이스가 도착하자마자 양자화되어 영구 커널 저장소에 복사됩니다.

gather와 pull은 파이프라인된다 — gather_lookahead가 그것을 제한

트레이너는 보통 파라미터를 있는 그대로 서빙할 수 없어요. FSDP가 샤딩하고, EP로 쪼갠 트레이너도 전체 전문가를 조립해야 하죠. 그래서 각 동기화는 여전히 gather 집단 연산을 돌리지만, 모델 단위가 아니라 레이어 단위로 돌립니다.

gather group은 하나의 디코더 레이어예요. 파라미터 리스트는 각 이름의 가장 바깥 인덱스 세그먼트로 키잉되고, 인덱스 없는 이름의 연속(첫 레이어 앞의 embedding, 마지막 뒤의 최종 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을 실제로 이름을 가진 랭크로 라우팅합니다. 파이프라인 스테이지, 전문가 병렬, 둘의 조합 모두 같은 선언이에요. 소비자는 이름을 가진 랭크들에 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 핸드셰이크에서 트레이너로부터 옵니다.

gpu_memory_utilization을 고르기 전에 수신 버퍼 크기를 정하세요. 각 워커는 num_rdt_buffers개의 수신 버퍼를 가지며, 각각 pull하는 가장 큰 단일 슬라이스 배치만큼 큽니다. NCCL과 NIXL 내부처럼 gpu_memory_utilization에 계산되지 않으므로, 여유(headroom)를 남기지 않는 비율은 엔진이 정상적으로 떴더라도 첫 동기화에서 OOM이 나요. 버퍼 크기는 가장 큰 원자 슬라이스에 의해 결정됩니다 — 슬라이스되지 않은 채로 가진 워커의 untied vocab 행렬이라면 그게 전체 embedding이에요.

트레이너 쪽 (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 actor를 소유하고 gather에 참여하며, 추론 쪽 핸드셰이크를 구동하는 것은 랭크 0뿐입니다. 어떤 VLLMWeightSyncClient든 동작해요.

일반 nn.Module이 아닌 트레이너(Megatron export, 원시 sharded 체크포인트)를 적응시키려면 WeightSource를 서브클래싱하세요.

ShardedRDTTrainerInitInfo

예시 (Examples)

• 4개 GPU의 작은 MoE — 2 FSDP2 트레이너 랭크 → 전문가 병렬을 가진 2 vLLM DP 랭크, 단일 노드. 트레이너 fleet과 별도 추론 fleet을 짝짓는데, 이 백엔드가 지원하는 유일한 배치이며, 동기화가 가중치를 옮겼다는 것과 두 번째 동기화가 생성물을 바꾸지 않는다는 것을 검증해서 CI에서 무인 실행됩니다. 트레이너를 의도적으로 작게 유지해 — 가중치를 실재로 만드는 데 충분한 FSDP2 — 파일이 트레이너보다는 가중치 동기화에 집중하게 합니다.

전체 RL 트레이너를 위해서는 SkyRL이 FSDP와 함께 Megatron(PP-local gathering과 MoE 전문가 스택 융합)으로 이 백엔드를 통합해요: NovaSky-AI/SkyRL#1753.

더 알아보기 (Learn more)

• 가중치 전송 (Weight Transfer)
• Base Classes와 커스텀 엔진
• NCCL 엔진