NCCL 엔진

NCCL 엔진 (NCCL Engine)

NCCL 가중치 전송 엔진은 NCCL broadcast 연산으로 트레이너에서 추론 워커로 가중치를 전송합니다. 트레이너와 추론 엔진이 별도 GPU에서 실행되는 멀티 노드·멀티 GPU 구성을 지원합니다.

출처: 문서

본문

NCCL을 언제 사용하나 (When to Use NCCL)

  • 학습과 추론이 별도 GPU에서(가능하면 노드 간)
  • 갱신된 가중치가 모두 필요한 여러 워커가 있는 tensor-parallel 추론
  • NVLink나 InfiniBand 위에서 고대역폭·저지연 가중치 전송이 필요한 경우

동작 방식 (How It Works)

  1. 트레이너와 모든 추론 워커가 StatelessProcessGroup(vLLM의 torch.distributed 무관 그룹 추상화)을 사용해 공유 NCCL 프로세스 그룹에 참여합니다. 트레이너는 rank 0, 워커는 rank_offset(1)에서 시작합니다.
  2. 트레이너가 가중치를 모든 워커에 동시에 broadcast합니다. 각 워커는 가중치를 받아 로드합니다.
  3. 선택적으로, packed tensor broadcasting이 여러 작은 텐서를 더 큰 버퍼로 배치해 더블/트리플 버퍼링과 CUDA 스트림 오버랩으로 더 높은 처리량을 얻습니다. 이 구현은 NeMo-RL의 packed tensor에 기반합니다.

워커의 update_weights와 트레이너의 broadcast는 동시에 실행됩니다 — 양쪽이 같은 NCCL 호출 안에서 rendezvous합니다. 트레이너 엔진이 그 동시성을 내부적으로 소유합니다.

추론 측 (Inference Side)

추론 측은 일반 백엔드 선택자를 받습니다. rendezvous 파라미터와 packing wire 파라미터는 init handshake에서 트레이너가 보냅니다.

from vllm import LLM
from vllm.config import WeightTransferConfig

llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="nccl"))

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

그 외에는 필요하지 않습니다: init_weight_transfer_engine, start_weight_update, update_weights, finish_weight_update 모두 트레이너 엔진이 원격으로 구동합니다.

트레이너 측 (Trainer Side)

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

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"),  # or RayVLLMWeightSyncClient(llm)
    source=ModuleSource(model),
)

engine.send_weights()   # once per sync

trainer_init은 전체 handshake를 구동합니다: 사이드 스레드에서 추론 측의 init_weight_transfer_engine을(클라이언트를 통해) 시작하면서 트레이너 자신의 rank-0 엔드포인트를 여는데, 양쪽이 함께 rendezvous해야 하기 때문입니다. rank_offset=1과 같은 packed 파라미터로 워커의 init info를 자체 구성하므로, 양쪽이 불일치할 수 없습니다.

send_weights()는 그다음 start_weight_update, update_weights(broadcast와 동시에 실행됨), finish_weight_update를 구동하고, 모든 전송이 drain된 후에만 반환합니다.

NCCLTrainerInitInfo

필드 기본값 설명
master_address — Rendezvous 호스트
master_port — Rendezvous 포트
world_size — 전체 트레이너 + 워커 NCCL 그룹 크기
rank — 키워드 전용. 이 트레이너 프로세스의 랭크; 0이 송신자
packed True packed broadcasting 사용
packed_buffer_size_bytes 1 GiB Packed 버퍼 크기
packed_num_buffers 2 회전 버퍼 수(더블/트리플 버퍼링)

여기서 packed는 기본적으로 True입니다. 워커 측 기본은 False지만, 트레이너가 값을 보내지 않을 때만 적용되며 이 경로에서는 결코 그런 일이 없습니다.

Packed Tensor Broadcasting

packed=True이면 가중치 텐서를 broadcast 전에 큰 연속 버퍼로 포장합니다. 이로써 NCCL 연산 수를 줄이고, 전용 CUDA 스트림과 더블/트리플 버퍼링으로 packing·broadcasting·unpacking을 겹칩니다.

packed, packed_buffer_size_bytes, packed_num_buffers는 NCCLTrainerInitInfo에서 만 설정합니다. 트레이너가 trainer_init 안에서 워커에 전파하고, 워커는 handshake에서 기록하며, receive_weights는 트레이너가 인코딩한 값과 정확히 같은 값으로 디코딩합니다. 이들은 라운드별 update_weights 필드가 아닙니다.

메모리

회전 버퍼는 전체 전송 동안 살아 있습니다: 양쪽에 packed_buffer_size_bytes * packed_num_buffers(기본값에서 2 GiB). 여유 공간이 너무 크면 packed_buffer_size_bytes를 낮추세요.

두 WeightSource 채널은 일치해야 합니다

Dense NCCL은 두 WeightSource 채널을 모두 읽는 백엔드라, 두 채널의 불일치가 치명적인 곳입니다. 엔진은 metadata()에서 라운드별 update info를 만들어 바이트보다 앞서 보냅니다. 워커는 그 info로 수신 버퍼 크기를 정하고, packed 모드에서는 그 info로 청크 경계를 자릅니다. 바이트 자체는 소스를 반복하며 나옵니다.

반복이 metadata()가 선언한 것과 불일치하면 — 재정렬·생략·dtype 변경 — 양쪽이 같은 바이트 스트림을 다르게 나눕니다. 그러면 전송이 도착하지 않는 길이를 기다리며 NCCL에서 hang하거나, 모델에 가비지를 로드합니다.

따라서 송신자는 진행하면서 각 쌍을 선언된 metadata와 검사합니다(파라미터당 비교 1회), wire에 닿게 하는 대신 첫 발산 파라미터의 이름을 밝히며 예외를 던집니다. ModuleSource는 구성상 이를 충족합니다. 커스텀 소스를 쓴다면 — Megatron export, MoE re-fusing 패스 — 이것이 가장 먼저 테스트할 불변식입니다.

참고

IPC는 metadata()를 전혀 읽지 않습니다: 진행하면서 iteration에서 update info를 도출하므로 발산을 관찰할 수 없습니다.

Sparse NCCL

희소·flat-index 가중치 패치는 WeightTransferConfig(backend="sparse_nccl")를 사용합니다. 이름·전체 형태·flat 인덱스는 checkpoint/Hugging Face 좌표에서 해석됩니다. 모든 추론 랭크가 같은 checkpoint 전역 패치를 받은 뒤, 모델의 네이티브 load_weights()가 rank-local TP/EP 및 packed 런타임 파라미터로 매핑합니다.

Tensor parallelism은 지원되며 트레이너 레이아웃과 일치할 필요가 없습니다. Pipeline-parallel 동작은 현재 sparse NCCL GPU 테스트가 다루지 않습니다.

Sparse는 delta 백엔드입니다: 각 호출이 안정적인 모델 파라미터 스트림 대신 교체 패치를 공급합니다. 엔진은 WeightSource를 받지 않고, 패치가 send_weights(patches)로 바로 갑니다. 빈 패치 목록은 no-op입니다.

from vllm.distributed.weight_transfer import (
    RayVLLMWeightSyncClient,
    WeightTransferTrainerFactory,
)
from vllm.distributed.weight_transfer.sparse_nccl_engine import (
    SparseNCCLTrainerInitInfo,
    SparseWeightPatch,
)

client = RayVLLMWeightSyncClient(llm)
engine = WeightTransferTrainerFactory.trainer_init(
    init_info=SparseNCCLTrainerInitInfo(
        master_address=master_address,
        master_port=master_port,
        world_size=world_size,
        rank=0,
    ),
    client=client,
)

patches = [
    SparseWeightPatch(
        name="model.layers.0.mlp.down_proj.weight",
        indices=flat_indices,           # int32, 1-D
        values=new_values,              # same length as indices
        full_shape=tuple(param.shape),  # required when sending via the engine
    )
]
engine.send_weights(patches)

각 send_weights(patches) 호출은 완전한 one-shot start/update/finish 수명주기를 소유합니다.

제네릭 워커 수명주기를 소유하는 RL 인프라는 트레이너 측 세션 상태를 추가하지 않고 유계 청크에 걸쳐 논리적 업데이트 하나를 열어둘 수 있습니다:

client.start_weight_update()
for patches in patch_chunks:
    engine.send_weight_chunk(patches)
client.finish_weight_update()

Sparse NCCL은 O(nnz) 인덱스·값만 보냅니다. Checkpoint 적용은 네이티브 로더를 통한 O(N) 스테이징을 여전히 사용하고, 호출자가 export/diff 상태와 부분 실패 후 재시작·재시드를 소유합니다.

rlhf_sparse_nccl.py는 Qwen3 MoE TP2/EP2 엔진으로 per-expert checkpoint 패치를 보여줍니다.

예시 (Examples)

더 알아보기 (Learn more)