NCCL 엔진

NCCL 엔진 (NCCL Engine)

훈련과 추론이 별도 GPU(때로는 별도 노드)에 있을 때, 가중치를 어떻게 전달할까요? 가장 일반적인 방법이 NCCL broadcast를 이용한 전송이에요. NCCL 엔진은 NCCL broadcast 연산으로 트레이너에서 추론 워커로 가중치를 전송합니다.

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

NCCL을 언제 쓸까 (When to Use NCCL)

• 훈련과 추론이 별도 GPU(노드를 넘어서도)에 있을 때
• 갱신된 가중치가 모두 필요한 여러 워커가 있는 텐서 병렬 추론
• NVLink 또는 InfiniBand를 통한 고대역폭·저지연 가중치 전송이 필요할 때

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

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

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

추론 쪽 (Inference Side)

추론 쪽은 단순한 백엔드 선택자만 받아요. rendezvous 파라미터와 패킹 wire params는 초기화 핸드셰이크에서 트레이너로부터 옵니다.

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이 전체 핸드셰이크를 구동해요. 추론 쪽의 init_weight_transfer_engine을(클라이언트를 통해) 사이드 스레드에서 시작하면서 트레이너 자신의 rank-0 엔드포인트를 여는 방식이죠. 양쪽 끝이 함께 rendezvous 해야 하기 때문입니다. 워커의 init info도 rank_offset=1과 같은 packed params로 직접 만들기 때문에, 두 쪽이 어긋날 수 없어요.

send_weights()는 start_weight_update, broadcast와 동시에 실행되는 update_weights, 그리고 finish_weight_update를 구동하고, 모든 전송이 소진된 뒤에만 반환합니다.

NCCLTrainerInitInfo

packed는 여기서 True가 기본이에요. 워커 쪽 기본값은 False지만, 트레이너가 값을 보내지 않을 때만 적용되고 이 경로에서는 그런 일이 발생하지 않습니다.

Packed Tensor Broadcasting

packed=True면 가중치 텐서가 broadcast 전에 큰 연속 버퍼로 포장돼요. 이렇게 하면 NCCL 연산 수가 줄고, 전용 CUDA 스트림으로 이중/삼중 버퍼링을 사용해 패킹·broadcast·언패킹을 겹칩니다.

packed, packed_buffer_size_bytes, packed_num_buffers는 NCCLTrainerInitInfo에서만 설정해요. 트레이너가 trainer_init 안에서 워커로 전파하고, 워커는 핸드셰이크에서 기록하며, 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 모드에서는 청크 경계를 거기서 자릅니다. 바이트 자체는 소스를 순회해서 나옵니다.

만약 순회가 metadata()가 선언한 것과 다르다면(재정렬, 생략, 재타입 파라미터), 양쪽이 같은 바이트 스트림을 다르게 나눠요. 그러면 전송이 도착하지 않는 길이를 NCCL에서 기다리며 멈추거나, 모델에 가비지를 로드하게 됩니다.

그래서 송신자는 진행하면서 각 쌍을 선언된 메타데이터와 비교해요. 파라미터당 한 번의 비교로, 첫 번째 다른 파라미터를 와이어에 올리기 전에 이름을 명시해 예외를 던집니다. ModuleSource는 구조적으로 이를 만족해요. 커스텀 소스를 쓴다면(Megatron export, MoE 재융합 패스) 가장 먼저 테스트할 불변식이 이거예요.

참고: IPC는 metadata()를 전혀 읽지 않아요. 순회에서 update info를 바로 파생하기 때문에 불일치를 관찰할 수 없습니다.

Sparse NCCL

희소, 플랫 인덱스 가중치 패치는 WeightTransferConfig(backend="sparse_nccl")를 사용해요. 이름, 전체 모양, 플랫 인덱스는 체크포인트/Hugging Face 좌표로 해석됩니다. 모든 추론 랭크가 같은 체크포인트 전역 패치를 받고, 모델의 네이티브 load_weights()가 이를 랭크 로컬 TP/EP 및 packed 런타임 파라미터로 매핑해요.

텐서 병렬은 지원되고 트레이너의 레이아웃과 일치하지 않아도 돼요. 파이프라인 병렬 동작은 현재 sparse NCCL GPU 테스트에서 다루지 않아요(확인 필요).

Sparse는 델타 백엔드예요. 각 호출이 모델 파라미터의 안정적인 스트림이 아니라 교체 패치를 공급하죠. 그래서 엔진은 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) 호출은 완전한 일회성 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) 인덱스와 값만 보내요. 체크포인트 적용은 여전히 네이티브 로더를 통한 O(N) 스테이징을 사용하고, 호출자가 export/diff 상태를 소유하며 부분 실패 후 재시작이나 리시드를 책임집니다.

rlhf_sparse_nccl.py는 Qwen3 MoE TP2/EP2 엔진으로 전문가별 체크포인트 패치를 보여줘요.

더 알아보기 (Learn more)

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