vllm.distributed.weight_transfer.m2n_common ¶

def check_placements(placements: Placements, context: str = "placements") -> None:
    """Reject placement pairs m2n cannot express.

    The header requires exactly one SHARD axis and one REPLICATE axis.
    `{REPLICATE, REPLICATE}` hits a degenerate prepare branch, which is why
    full replication is carried as `REPLICATED` and resolved separately;
    sharding both axes is not expressible at all.
    """
    if len(placements) != MESH_NDIMS:
        raise ValueError(f"{context} must have {MESH_NDIMS} entries")
    invalid = [code for code in placements if code < REPLICATE]
    if invalid:
        raise ValueError(
            f"{context} contains invalid placement code {invalid[0]}; valid "
            f"codes are {REPLICATE} (Replicate) and non-negative tensor dimensions"
        )
    num_sharded = sum(code != REPLICATE for code in placements)
    if num_sharded == 0:
        raise ValueError(
            f"{context}: a fully replicated tensor is carried as REPLICATED, "
            "not as two "
            "REPLICATE axes"
        )
    if num_sharded == MESH_NDIMS:
        raise ValueError(
            f"{context}: nccl_m2n needs one REPLICATE mesh axis, but "
            f"{placements} shards both"
        )

def check_plan_limits(
    src: tuple[M2NMesh, Placements],
    dst: tuple[M2NMesh, Placements],
    name: str,
) -> None:
    """Reject plans that exceed m2n's static per-shard fan-in / fan-out arrays.

    Only the unambiguous cases are checked: when one side is a single shard it
    is fed by (or feeds) every shard on the other side, so the count is exactly
    the other side's shard count. In the general sharded-to-sharded case the
    overlap depends on the library's chunking, and reproducing that arithmetic
    here would duplicate internals that can change under us.

    m2n re-checks authoritatively and, because the plan is derived from the
    shared descriptors, fails identically on every rank -- so this is about
    reporting at init with a message that names the parameter, not about
    avoiding a hang.
    """
    src_shards = shard_count(*src)
    dst_shards = shard_count(*dst)
    if dst_shards == 1 and src_shards > MAX_SOURCE_SHARDS:
        raise ValueError(
            f"parameter '{name}' would feed {src_shards} source shards into one "
            f"replicated destination, over m2n's MAX_SOURCES={MAX_SOURCE_SHARDS}. "
            "Shard the destination, replicate the source, or rebuild m2n with "
            "larger arrays."
        )
    if src_shards == 1 and dst_shards > MAX_DEST_SHARDS:
        raise ValueError(
            f"parameter '{name}' would fan one source shard out to {dst_shards} "
            f"destination shards, over m2n's MAX_TARGETS={MAX_DEST_SHARDS}. "
            "Rebuild m2n with larger arrays to raise the bound."
        )

def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> None:
    """Reject tensors m2n cannot move at all, with the parameter named."""
    if dtype not in SUPPORTED_DTYPES:
        raise ValueError(
            f"parameter '{name}' has dtype {dtype}, which nccl_m2n does not "
            f"support. Supported: {sorted(str(d) for d in SUPPORTED_DTYPES)}"
        )
    if not 1 <= len(shape) <= MAX_TENSOR_DIMS:
        raise ValueError(
            f"parameter '{name}' has rank {len(shape)}; nccl_m2n supports "
            f"rank 1..{MAX_TENSOR_DIMS}"
        )

def comm_ptr(comm: "PyNcclCommunicator") -> int:
    """Raw `ncclComm_t` behind vLLM's `PyNcclCommunicator`.

    m2n links its own NCCL, so the handle only means anything if vLLM loaded
    the same `libnccl.so` — set `VLLM_NCCL_SO_PATH` accordingly.
    """
    handle = comm.comm
    ptr = getattr(handle, "value", handle)
    if not ptr:
        raise RuntimeError("PyNcclCommunicator has no live NCCL communicator")
    return int(ptr)

def import_m2n() -> Any:
    """Import `nccl.m2n` lazily, with an actionable error when it is missing.

    Deferred so that importing vLLM — or any other weight transfer backend —
    never requires the m2n runtime to be present.
    """
    try:
        import nccl.m2n as m2n
    except ImportError as e:
        raise ImportError(f"{_IMPORT_HINT} (import failed: {e})") from e
    return m2n

def publish_destination_placements(
    comm: "PyNcclCommunicator",
    first_worker_rank: int,
    placements: "Sequence[Placements | None] | None",
    num_parameters: int,
) -> list[Placements | None]:
    """Share the worker-side destination plan with every rank in the group.

    The trainer must issue each reshard with the same destination the workers
    use, but once destinations are per-parameter they depend on the inference
    model, which only the workers can see. The first worker publishes them here,
    over the shared NCCL communicator; trainer ranks pass `None` and receive
    them, and the other workers pass their own so a disagreement is caught
    rather than deadlocking later.

    Only placements travel: the destination *mesh* is still derived from the
    rank split, identically on both sides. Use the NCCL communicator directly
    because unique-id rendezvous deliberately has no bootstrap process group.
    """
    replicated_sentinel = REPLICATE - 1
    encoded = torch.full(
        (num_parameters, MESH_NDIMS),
        replicated_sentinel,
        dtype=torch.int8,
        device=comm.device,
    )
    if comm.rank == first_worker_rank:
        if placements is None or len(placements) != num_parameters:
            raise ValueError(
                f"publishing rank needs {num_parameters} destination placements"
            )
        rows = [
            [replicated_sentinel] * MESH_NDIMS if placement is REPLICATED else placement
            for placement in placements
        ]
        if rows:
            encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=comm.device))

    comm.broadcast(encoded, src=first_worker_rank)
    return [
        REPLICATED
        if all(code == replicated_sentinel for code in row)
        else cast(Placements, tuple(row))
        for row in encoded.tolist()
    ]

def resolve_layout(
    mesh: M2NMesh,
    placements: Placements | None,
    context: str = "placements",
) -> tuple[M2NMesh, Placements]:
    """Pair one tensor's placement with the mesh m2n should see for it.

    A replicated tensor needs a size-1 mesh axis to carry a no-op shard, since
    m2n has no `{REPLICATE, REPLICATE}`. It is therefore described over the
    *same rank interval* re-factored as `(size, 1)`. That re-factoring is sound
    precisely because replication is order-independent — every rank holds the
    whole tensor, so it does not matter that `(a, b)` and `(size, 1)` walk the
    interval in a different order. A sharded tensor keeps its side's own
    factorization, where rank order decides who owns which shard.
    """
    if placements is None:
        return M2NMesh((mesh.size, 1), mesh.start_rank), (REPLICATE, 0)
    check_placements(placements, context)
    return mesh, placements

def shard_count(mesh: M2NMesh, placements: Placements) -> int:
    """How many pieces this layout splits the tensor into."""
    return next(
        (mesh.dims[axis] for axis, code in enumerate(placements) if code != REPLICATE),
        1,
    )

def validate_layout(
    mesh: M2NMesh, placements: Placements, shape: Sequence[int], side: str
) -> None:
    """Check a resolved layout can describe `shape` before any collective runs.

    Called on both sides at init so a bad plan surfaces as an error from the
    init RPC rather than as a hang inside the first reshard.
    """
    for axis, code in enumerate(placements):
        if code == REPLICATE:
            continue
        if code >= len(shape):
            raise ValueError(
                f"{side} layout shards tensor dim {code}, but the tensor "
                f"has rank {len(shape)}"
            )
        factor = mesh.dims[axis]
        if shape[code] % factor:
            raise ValueError(
                f"{side} layout shards dim {code} (size {shape[code]}) over "
                f"{factor} ranks, which does not divide evenly"
            )