vllm.distributed.weight_transfer.nccl_common ¶

def decode_nccl_unique_id(
    *,
    master_address: str | None,
    master_port: int | None,
    nccl_unique_id_b64: str | None,
    ctx: str,
) -> bytes | None:
    """Validate the rendezvous mode and decode a pre-shared unique id."""
    has_uid = bool(nccl_unique_id_b64)
    has_addr = master_address is not None
    has_port = master_port is not None

    if has_uid and (has_addr or has_port):
        raise ValueError(
            f"{ctx}: pass nccl_unique_id_b64 OR master_address/master_port, not both"
        )
    if not has_uid:
        if not (has_addr and has_port):
            raise ValueError(
                f"{ctx}: need nccl_unique_id_b64, or both master_address and "
                f"master_port (got master_address={master_address!r}, "
                f"master_port={master_port!r})"
            )
        return None
    try:
        decoded = base64.b64decode(nccl_unique_id_b64, validate=True)
    except (ValueError, TypeError) as e:
        raise ValueError(
            f"{ctx}: nccl_unique_id_b64 is not valid standard base64 "
            f"(URL-safe base64 is not accepted)"
        ) from e
    if len(decoded) != NCCL_UNIQUE_ID_BYTES:
        raise ValueError(
            f"{ctx}: nccl_unique_id_b64 must decode to {NCCL_UNIQUE_ID_BYTES} "
            f"bytes, got {len(decoded)}"
        )
    return decoded

def stateless_init_process_group(
    master_address: str,
    master_port: int,
    rank: int,
    world_size: int,
    device,
) -> "PyNcclCommunicator":
    """
    vLLM provides `StatelessProcessGroup` to create a process group
    without considering the global process group in torch.distributed.
    It is recommended to create `StatelessProcessGroup`, and then initialize
    the data-plane communication (NCCL) between external (train processes)
    and vLLM workers.
    """
    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
    from vllm.distributed.utils import StatelessProcessGroup

    pg = StatelessProcessGroup.create(
        host=master_address, port=master_port, rank=rank, world_size=world_size
    )
    return PyNcclCommunicator(pg, device=device)

def trainer_init(
    init_info: NCCLRendezvous | dict,
) -> "PyNcclCommunicator":
    """
    Initialize NCCL process group for trainer-side weight transfer.

    The trainer is always rank 0 in the process group. Uses the current
    CUDA device (torch.accelerator.current_device_index()).

    Args:
        init_info: Any object carrying the `NCCLRendezvous` fields (a trainer or
            worker NCCL init info), or a dict with keys:
            - master_address: str
            - master_port: int
            - world_size: int

    Returns:
        PyNcclCommunicator for weight transfer.
    """
    if isinstance(init_info, dict):
        master_address = init_info["master_address"]
        master_port = init_info["master_port"]
        world_size = init_info["world_size"]
    else:
        master_address = init_info.master_address
        master_port = init_info.master_port
        world_size = init_info.world_size

    # Trainer is always rank 0
    device = torch.accelerator.current_device_index()
    return stateless_init_process_group(
        master_address,
        master_port,
        0,
        world_size,
        device,
    )

def uid_init_process_group(
    nccl_unique_id_bytes: bytes,
    rank: int,
    world_size: int,
    device,
) -> "PyNcclCommunicator":
    """Join the NCCL group from pre-shared ``ncclUniqueId`` bytes.

    The torch-free rendezvous alternative to `stateless_init_process_group`: no
    TCPStore, and therefore no barrier -- every rank must enter concurrently.
    """
    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator

    return PyNcclCommunicator.from_unique_id_bytes(
        nccl_unique_id_bytes,
        rank=rank,
        world_size=world_size,
        device=device,
    )

def worker_init_payload(init_info: NCCLWeightTransferInitInfo) -> dict:
    """Serialize a worker init info for `init_weight_transfer_engine`, dropping
    the unset rendezvous field (the UID in TCP mode) so the wire payload carries
    only the mode actually in use. Shared by the dense and sparse trainer
    engines so the two cannot drift."""
    return {key: value for key, value in asdict(init_info).items() if value is not None}

def worker_init_process_group(
    init_info: NCCLWeightTransferInitInfo,
    parallel_config: "ParallelConfig",
) -> "PyNcclCommunicator":
    """Create the trainer<->worker NCCL group on an inference worker.

    Computes a unique rank for this worker across all data-parallel groups and
    joins the trainer via whichever rendezvous mode `init_info` carries.
    """
    # Calculate the global rank in the trainer-worker process group.
    # Must account for data parallel to get unique ranks across all workers.
    dp_rank = parallel_config.data_parallel_index
    world_size_per_dp = parallel_config.world_size  # TP * PP
    rank_within_dp = parallel_config.rank

    # Unique rank across all DP groups
    worker_rank = dp_rank * world_size_per_dp + rank_within_dp
    rank = worker_rank + init_info.rank_offset

    device = torch.accelerator.current_device_index()
    unique_id_bytes = init_info.nccl_unique_id_bytes
    if unique_id_bytes is not None:
        return uid_init_process_group(
            unique_id_bytes,
            rank,
            init_info.world_size,
            device,
        )
    # __post_init__ guarantees the TCP pair when there is no unique id.
    assert init_info.master_address is not None
    assert init_info.master_port is not None
    return stateless_init_process_group(
        init_info.master_address,
        init_info.master_port,
        rank,
        init_info.world_size,
        device=device,
    )