iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[RL] Support per-parameter source layouts for NCCL M2N (#60934)

Browse Source 
Signed-off-by: pkousha <pkousha@nvidia.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
This commit is contained in:
pkousha
2026-10-16 21:41:17 -04:00
committed by GitHub
co-authored by OpenAI Codex aoshen02
parent f8112704d5
commit 6b6a6c5c98
8 changed files with 2590 additions and 695 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/distributed/weight_transfer/m2n_common.py
+646 -120
View File
@@ -9,16 +9,19 @@ runtime import, layout descriptors, and the conversion to `nccl.m2n` types —
so the engine module stays about the transfer itself.
"""
from collections.abc import Sequence
import hashlib
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, cast
import torch
from vllm.distributed.weight_transfer.base import ParamMeta
from vllm.utils.nccl import unpinned_nccl_env
if TYPE_CHECKING:
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
from vllm.distributed.utils import StatelessProcessGroup
# Placement code for a mesh axis that replicates. Any other (non-negative) code
# is the tensor dim that axis shards. Ints rather than `nccl.m2n` objects so
@@ -35,15 +38,15 @@ MESH_NDIMS = 2
DESTINATION_REPLICA_AXIS = 0
DESTINATION_SHARD_AXIS = 1
# m2n bounds how many source shards may feed one destination shard, and how many
# destination shards one source shard may feed, with compile-time arrays in
# `reshard_limits.h`. The bindings do not expose them, so they are mirrored
# here; a build with larger arrays makes these conservative, never wrong.
# Conservative capacities accepted by the supported nccl_m2n algorithms. They
# are not exposed by its Python API yet; remove these mirrors once nccl_m2n
# provides communication-free plan validation.
MAX_SOURCE_SHARDS = 16
MAX_DEST_SHARDS = 64
# Wire dtypes `ncclReshard` accepts. Notably excludes fp4 and any packed
# sub-byte type, so quantized checkpoints are out of scope for now.
# Increment when the JSON init payload changes incompatibly.
M2N_WIRE_SCHEMA_VERSION = 2
SUPPORTED_DTYPES: frozenset[torch.dtype] = frozenset(
dtype
for dtype in (
@@ -70,12 +73,62 @@ _IMPORT_HINT = (
"The nccl_m2n weight transfer backend requires the `nccl-extensions` "
"package (and its `nccl4py` dependency), which is not installed by vLLM. "
"See https://github.com/NVIDIA/nccl-extensions for build "
"and install instructions. It needs NCCL 2.30.5 or newer, and "
"VLLM_NCCL_SO_PATH must point at the same libnccl.so that libnccl_m2n.so "
"was linked against."
"and install instructions. It needs NCCL 2.30.5 or newer."
)
def prepare_m2n_local_runtime(m2n: Any, max_cta: int | None) -> tuple[Any, int, Any]:
"""Prepare rank-local resources before any rank enters NCCL init."""
import nccl.core as nccl
device = torch.accelerator.current_device_index()
if not isinstance(device, int) or isinstance(device, bool) or device < 0:
raise RuntimeError(f"invalid current accelerator device index: {device!r}")
handle = m2n.Handle.create(m2n.Config(max_cta=max_cta))
return nccl, device, handle
def create_m2n_communicator(
nccl: Any,
group: "StatelessProcessGroup",
unique_id_bytes: bytes | None,
) -> Any:
"""Create the M2N data-plane communicator through NCCL4Py."""
if unique_id_bytes is None:
envelope: dict[str, bytes | str | None] | None = None
if group.rank == 0:
try:
envelope = {
"unique_id": bytes(nccl.get_unique_id().as_bytes),
"error": None,
}
except Exception as exc:
envelope = {
"unique_id": None,
"error": f"{type(exc).__name__}: {exc}",
}
envelope = group.broadcast_obj(envelope, src=0)
if not isinstance(envelope, Mapping) or set(envelope) != {"unique_id", "error"}:
raise RuntimeError("invalid NCCL4Py unique-id envelope")
error = envelope["error"]
received_unique_id = envelope["unique_id"]
if error is not None:
if not isinstance(error, str) or not error:
raise RuntimeError("invalid NCCL4Py unique-id error")
raise RuntimeError(f"NCCL4Py unique-id generation failed: {error}")
if not isinstance(received_unique_id, bytes):
raise RuntimeError("NCCL4Py unique-id broadcast did not return bytes")
unique_id_bytes = received_unique_id
unique_id = nccl.UniqueId.from_bytes(unique_id_bytes)
with unpinned_nccl_env():
return nccl.Communicator.init(
nranks=group.world_size,
rank=group.rank,
unique_id=unique_id,
)
def import_m2n() -> Any:
"""Import `nccl.m2n` lazily, with an actionable error when it is missing.
@@ -97,17 +150,25 @@ class M2NMesh:
`[start_rank, start_rank + dims[0] * dims[1])`. There is no 1-D mesh; a
single-axis topology is spelled with a second axis of size 1.
One mesh describes every tensor on its side, so it is exchanged once at the
init handshake rather than per parameter.
A model may use a different factorization for each tensor. For example,
dense weights can use a DP x TP mesh while expert weights use EDP x EP.
"""
dims: tuple[int, int]
start_rank: int
def __post_init__(self) -> None:
if len(self.dims) != MESH_NDIMS or any(d <= 0 for d in self.dims):
object.__setattr__(self, "dims", tuple(self.dims))
if len(self.dims) != MESH_NDIMS or any(
not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0
for dim in self.dims
):
raise ValueError(f"mesh dims must be {MESH_NDIMS} positive ints")
if self.start_rank < 0:
if (
not isinstance(self.start_rank, int)
or isinstance(self.start_rank, bool)
or self.start_rank < 0
):
raise ValueError(
f"mesh start_rank must be non-negative, got {self.start_rank}"
)
@@ -124,16 +185,26 @@ Placements = tuple[int, int]
REPLICATED: Placements | None = None
def check_placements(placements: Placements, context: str = "placements") -> None:
"""Reject placement pairs m2n cannot express.
@dataclass(frozen=True)
class M2NLayout:
"""A tensor's mesh and placement on one side of a transfer."""
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.
"""
mesh: M2NMesh
placements: Placements | None
def __post_init__(self) -> None:
if not isinstance(self.mesh, M2NMesh):
raise TypeError("M2N layout mesh must be an M2NMesh")
if self.placements is not None:
object.__setattr__(self, "placements", tuple(self.placements))
def check_placements(placements: Placements, context: str = "placements") -> None:
"""Reject placement pairs nccl_m2n cannot express."""
if len(placements) != MESH_NDIMS:
raise ValueError(f"{context} must have {MESH_NDIMS} entries")
if any(not isinstance(code, int) or isinstance(code, bool) for code in placements):
raise TypeError(f"{context} must contain integer placement codes")
invalid = [code for code in placements if code < REPLICATE]
if invalid:
raise ValueError(
@@ -144,8 +215,7 @@ def check_placements(placements: Placements, context: str = "placements") -> Non
if num_sharded == 0:
raise ValueError(
f"{context}: a fully replicated tensor is carried as REPLICATED, "
"not as two "
"REPLICATE axes"
"not as two REPLICATE axes"
)
if num_sharded == MESH_NDIMS:
raise ValueError(
@@ -175,87 +245,46 @@ def resolve_layout(
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],
context: str,
) -> tuple[int, ...]:
"""Validate one resolved layout and return its required local shape."""
local_shape = list(shape)
for axis, tensor_dim in enumerate(placements):
if tensor_dim == REPLICATE:
continue
if tensor_dim >= len(shape):
raise ValueError(
f"{context} shards tensor dim {tensor_dim}, but parameter has "
f"rank {len(shape)}"
)
factor = mesh.dims[axis]
if shape[tensor_dim] % factor:
raise ValueError(
f"{context} shards dim {tensor_dim} (size {shape[tensor_dim]}) "
f"over {factor} ranks, which does not divide evenly"
)
local_shape[tensor_dim] //= factor
return tuple(local_shape)
def expected_local_shape(meta: "M2NParamMeta") -> tuple[int, ...]:
"""Validate a parameter's source layout and return its local shape."""
mesh, placements = resolve_layout(
meta.source_layout.mesh,
meta.source_layout.placements,
f"parameter '{meta.name}' source placements",
)
return validate_layout(
mesh, placements, meta.shape, f"parameter '{meta.name}' source layout"
)
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 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"
)
@dataclass(frozen=True)
class M2NParamMeta(ParamMeta):
"""`ParamMeta` extended with how the trainer places this tensor.
The base class carries only name / dtype / full shape, which is not enough
to plan a reshard. `placements` is relative to its side's `M2NMesh`, or
`REPLICATED` when every rank holds the whole tensor.
"""
placements: Placements | None
def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> None:
"""Reject tensors m2n cannot move at all, with the parameter named."""
"""Reject parameters nccl_m2n cannot transfer."""
if dtype not in SUPPORTED_DTYPES:
raise ValueError(
f"parameter '{name}' has dtype {dtype}, which nccl_m2n does not "
@@ -268,6 +297,514 @@ def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> N
)
def _shard_spec(mesh: M2NMesh, placements: Placements) -> tuple[int, int]:
return next(
(code, mesh.dims[axis])
for axis, code in enumerate(placements)
if code != REPLICATE
)
def _max_partition_overlaps(
extent: int,
partitions: int,
other_partitions: int,
) -> int:
"""Return the maximum other partitions intersected by any one partition."""
if extent == 0:
return 0
width = extent // partitions
other_width = extent // other_partitions
max_overlaps = 0
for index in range(partitions):
start = index * width
end = start + width
first_other = start // other_width
last_other = (end - 1) // other_width
max_overlaps = max(max_overlaps, last_other - first_other + 1)
return max_overlaps
def check_plan_limits(
src: tuple[M2NMesh, Placements],
dst: tuple[M2NMesh, Placements],
name: str,
shape: Sequence[int],
) -> None:
"""Reject plans whose per-shard fan-in or fan-out exceeds M2N capacity."""
src_dim, src_shards = _shard_spec(*src)
dst_dim, dst_shards = _shard_spec(*dst)
if any(extent == 0 for extent in shape):
max_sources = max_targets = 0
elif src_dim != dst_dim:
max_sources = src_shards
max_targets = dst_shards
else:
extent = shape[src_dim]
max_sources = _max_partition_overlaps(extent, dst_shards, src_shards)
max_targets = _max_partition_overlaps(extent, src_shards, dst_shards)
if max_sources > MAX_SOURCE_SHARDS:
raise ValueError(
f"parameter '{name}' would feed up to {max_sources} source shards "
f"into one destination shard, over nccl_m2n MAX_SOURCES="
f"{MAX_SOURCE_SHARDS}"
)
if max_targets > MAX_DEST_SHARDS:
raise ValueError(
f"parameter '{name}' would fan one source shard out to up to "
f"{max_targets} destination shards, over nccl_m2n "
f"MAX_TARGETS={MAX_DEST_SHARDS}"
)
def validate_transfer_plan(
metas: Sequence["M2NParamMeta"],
dst_mesh: M2NMesh,
dst_placements: Sequence[Placements | None],
) -> None:
"""Validate the complete ordered source-to-destination transfer plan."""
if len(dst_placements) != len(metas):
raise ValueError(
f"destination plan has {len(dst_placements)} parameters, expected "
f"{len(metas)}"
)
for meta, placement in zip(metas, dst_placements):
check_transferable(meta.name, meta.dtype, meta.shape)
src = resolve_layout(
meta.source_layout.mesh,
meta.source_layout.placements,
f"parameter '{meta.name}' source placements",
)
validate_layout(*src, meta.shape, f"parameter '{meta.name}' source layout")
dst = resolve_layout(
dst_mesh,
placement,
f"parameter '{meta.name}' destination placements",
)
validate_layout(*dst, meta.shape, f"parameter '{meta.name}' destination layout")
check_plan_limits(src, dst, meta.name, meta.shape)
@dataclass(frozen=True)
class M2NParamMeta(ParamMeta):
"""`ParamMeta` extended with how the trainer places this tensor.
The base class carries only name / dtype / full shape, which is not enough
to plan a reshard. `source_layout` carries both the per-parameter mesh and
its placement on that mesh.
"""
source_layout: M2NLayout
def __post_init__(self) -> None:
if not isinstance(self.source_layout, M2NLayout):
raise TypeError("source_layout must be an M2NLayout")
@dataclass(frozen=True)
class M2NWireParam:
"""JSON-safe source metadata for one parameter in the init handshake."""
name: str
dtype_name: str
shape: tuple[int, ...]
src_mesh_dims: tuple[int, int]
src_placements: Placements | None
def __post_init__(self) -> None:
shape = tuple(self.shape)
dims = tuple(self.src_mesh_dims)
placements = None if self.src_placements is None else tuple(self.src_placements)
object.__setattr__(self, "shape", shape)
object.__setattr__(self, "src_mesh_dims", dims)
object.__setattr__(self, "src_placements", placements)
if not isinstance(self.name, str) or not self.name:
raise ValueError("M2N wire parameter name must not be empty")
if not isinstance(self.dtype_name, str) or not self.dtype_name:
raise ValueError(f"parameter '{self.name}' has an empty wire dtype name")
if not shape or any(
not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0
for dim in shape
):
raise ValueError(
f"parameter '{self.name}' shape must contain positive integers"
)
M2NMesh(cast(tuple[int, int], dims), 0)
if placements is not None:
check_placements(
cast(Placements, placements),
f"parameter '{self.name}' source placements",
)
@classmethod
def from_dict(cls, value: Mapping[str, Any]) -> "M2NWireParam":
"""Parse one strict wire record from its JSON-decoded mapping."""
fields = {
"name",
"dtype_name",
"shape",
"src_mesh_dims",
"src_placements",
}
missing = fields - value.keys()
extra = value.keys() - fields
if missing or extra:
details = []
if missing:
details.append(f"missing {sorted(missing)}")
if extra:
details.append(f"unexpected {sorted(extra)}")
raise ValueError(f"invalid M2N wire parameter: {', '.join(details)}")
return cls(**value)
def to_dict(self) -> dict[str, Any]:
"""Serialize this record using only JSON-compatible values."""
return {
"name": self.name,
"dtype_name": self.dtype_name,
"shape": list(self.shape),
"src_mesh_dims": list(self.src_mesh_dims),
"src_placements": (
None if self.src_placements is None else list(self.src_placements)
),
}
def source_plan_digest(params: Sequence[M2NWireParam]) -> str:
"""Return a rank-independent digest of the ordered source wire plan."""
payload = json.dumps(
{
"schema_version": M2N_WIRE_SCHEMA_VERSION,
"params": [param.to_dict() for param in params],
},
sort_keys=True,
separators=(",", ":"),
).encode()
return hashlib.sha256(payload).hexdigest()
def complete_plan_digest(
source_digest: str,
dst_mesh: M2NMesh,
dst_placements: Sequence[Placements | None],
) -> str:
"""Return a digest of the agreed source and destination plan."""
if not isinstance(source_digest, str) or not source_digest:
raise ValueError("source digest must be a non-empty string")
payload = json.dumps(
{
"schema_version": M2N_WIRE_SCHEMA_VERSION,
"source_digest": source_digest,
"destination_mesh": {
"dims": list(dst_mesh.dims),
"start_rank": dst_mesh.start_rank,
},
"destination_placements": [
None if placement is None else list(placement)
for placement in dst_placements
],
},
sort_keys=True,
separators=(",", ":"),
).encode()
return hashlib.sha256(payload).hexdigest()
def check_source_plan_agreement(
group: "StatelessProcessGroup",
params: Sequence[M2NWireParam],
expected_digest: str | None = None,
local_error: str | None = None,
) -> str:
"""Verify every participant entered with the same source plan."""
prepared_error = local_error
digest: str | None = None
declared: str | None = expected_digest
if prepared_error is None:
try:
digest = source_plan_digest(params)
declared = digest if expected_digest is None else expected_digest
if not isinstance(declared, str) or not declared:
raise ValueError("declared source digest must be a non-empty string")
except Exception as exc:
prepared_error = f"{type(exc).__name__}: {exc}"
elif not isinstance(prepared_error, str) or not prepared_error:
prepared_error = "TypeError: local source error must be a non-empty string"
gathered = group.all_gather_obj(
{
"phase": "source",
"digest": digest,
"declared_digest": declared,
"error": prepared_error,
}
)
errors: list[str] = []
expected_fields = {"phase", "digest", "declared_digest", "error"}
if len(gathered) != group.world_size:
errors.append(
"control-plane all-gather returned "
f"{len(gathered)} records for world size {group.world_size}"
)
agreed: list[tuple[int, str, str]] = []
for rank, item in enumerate(gathered):
if not isinstance(item, Mapping) or set(item) != expected_fields:
errors.append(f"rank {rank}: invalid source preflight envelope")
continue
if item["phase"] != "source":
errors.append(f"rank {rank}: invalid source preflight phase")
continue
error = item["error"]
if error is not None:
if not isinstance(error, str) or not error:
errors.append(f"rank {rank}: invalid source preflight error")
else:
errors.append(f"rank {rank}: {error}")
continue
computed = item["digest"]
claimed = item["declared_digest"]
if (
not isinstance(computed, str)
or not computed
or not isinstance(claimed, str)
or not claimed
):
errors.append(f"rank {rank}: invalid source digest")
continue
agreed.append((rank, computed, claimed))
if errors:
raise RuntimeError(
"nccl_m2n source preflight rejected before NCCL init: " + "; ".join(errors)
)
canonical = agreed[0][1]
for rank, computed, claimed in agreed:
if computed != claimed:
errors.append(
f"rank {rank}: computed digest {computed} does not match "
f"declared digest {claimed}"
)
if computed != canonical:
errors.append(f"rank {rank}: source digest disagrees with rank 0")
if errors:
raise RuntimeError(
"nccl_m2n source preflight rejected before NCCL init: " + "; ".join(errors)
)
return canonical
def check_complete_plan_agreement(
group: "StatelessProcessGroup",
source_digest: str,
dst_mesh: M2NMesh,
dst_placements: Sequence[Placements | None],
local_error: str | None = None,
) -> str:
"""Require every rank to accept the same complete transfer plan."""
prepared_error = local_error
digest: str | None = None
if prepared_error is None:
try:
digest = complete_plan_digest(source_digest, dst_mesh, dst_placements)
except Exception as exc:
prepared_error = f"{type(exc).__name__}: {exc}"
elif not isinstance(prepared_error, str) or not prepared_error:
prepared_error = "TypeError: local plan error must be a non-empty string"
gathered = group.all_gather_obj(
{"phase": "complete_plan", "digest": digest, "error": prepared_error}
)
errors: list[str] = []
valid: list[tuple[int, str]] = []
expected_fields = {"phase", "digest", "error"}
if len(gathered) != group.world_size:
errors.append(
"control-plane all-gather returned "
f"{len(gathered)} records for world size {group.world_size}"
)
for rank, item in enumerate(gathered):
if (
not isinstance(item, Mapping)
or set(item) != expected_fields
or item.get("phase") != "complete_plan"
):
errors.append(f"rank {rank}: invalid complete-plan envelope")
continue
error = item["error"]
if error is not None:
if not isinstance(error, str) or not error:
errors.append(f"rank {rank}: invalid complete-plan error")
else:
errors.append(f"rank {rank}: {error}")
continue
rank_digest = item["digest"]
if not isinstance(rank_digest, str) or not rank_digest:
errors.append(f"rank {rank}: invalid complete-plan digest")
continue
valid.append((rank, rank_digest))
if not errors:
canonical = valid[0][1]
for rank, rank_digest in valid[1:]:
if rank_digest != canonical:
errors.append(
f"rank {rank}: complete-plan digest disagrees with rank 0"
)
if errors:
raise RuntimeError(
"nccl_m2n complete-plan preflight rejected before first reshard: "
+ "; ".join(errors)
)
return valid[0][1]
def check_runtime_ready_agreement(
group: "StatelessProcessGroup", local_error: str | None
) -> None:
"""Require every rank to prepare local resources before NCCL init."""
prepared_error = local_error
if prepared_error is not None and (
not isinstance(prepared_error, str) or not prepared_error
):
prepared_error = "TypeError: local runtime error must be a non-empty string"
gathered = group.all_gather_obj({"phase": "runtime_ready", "error": prepared_error})
errors: list[str] = []
if len(gathered) != group.world_size:
errors.append(
"control-plane all-gather returned "
f"{len(gathered)} records for world size {group.world_size}"
)
for rank, item in enumerate(gathered):
if (
not isinstance(item, Mapping)
or set(item) != {"phase", "error"}
or item.get("phase") != "runtime_ready"
):
errors.append(f"rank {rank}: invalid runtime readiness envelope")
continue
error = item.get("error")
if error is not None:
if not isinstance(error, str) or not error:
errors.append(f"rank {rank}: invalid runtime readiness error")
else:
errors.append(f"rank {rank}: {error}")
if errors:
raise RuntimeError(
"nccl_m2n local runtime preflight failed before NCCL init: "
+ "; ".join(errors)
)
def check_data_plane_agreement(
group: "StatelessProcessGroup",
unique_id_bytes: bytes | None,
local_error: str | None = None,
) -> None:
"""Agree on the NCCL rendezvous identity before NCCL init."""
prepared_error = local_error
uid_digest: str | None = None
if prepared_error is None:
try:
if unique_id_bytes is not None:
if not isinstance(unique_id_bytes, bytes):
raise TypeError("NCCL unique id must decode to bytes")
uid_digest = hashlib.sha256(unique_id_bytes).hexdigest()
except Exception as exc:
prepared_error = f"{type(exc).__name__}: {exc}"
elif not isinstance(prepared_error, str) or not prepared_error:
prepared_error = "TypeError: local data-plane error must be non-empty"
gathered = group.all_gather_obj(
{
"phase": "data_plane",
"mode": "uid" if unique_id_bytes is not None else "tcp",
"uid_digest": uid_digest,
"error": prepared_error,
}
)
expected_fields = {
"phase",
"mode",
"uid_digest",
"error",
}
errors: list[str] = []
valid: list[tuple[int, str, str | None]] = []
if len(gathered) != group.world_size:
errors.append(
"control-plane all-gather returned "
f"{len(gathered)} records for world size {group.world_size}"
)
for rank, item in enumerate(gathered):
if (
not isinstance(item, Mapping)
or set(item) != expected_fields
or item.get("phase") != "data_plane"
):
errors.append(f"rank {rank}: invalid data-plane preflight envelope")
continue
error = item["error"]
if error is not None:
if not isinstance(error, str) or not error:
errors.append(f"rank {rank}: invalid data-plane preflight error")
else:
errors.append(f"rank {rank}: {error}")
continue
mode = item["mode"]
uid = item["uid_digest"]
if mode not in {"tcp", "uid"} or (mode == "uid") != isinstance(uid, str):
errors.append(f"rank {rank}: invalid data-plane identity")
continue
valid.append((rank, mode, uid))
if not errors:
_, mode, uid = valid[0]
for rank, other_mode, other_uid in valid[1:]:
if (other_mode, other_uid) != (mode, uid):
errors.append(f"rank {rank}: NCCL data-plane identity disagrees")
if errors:
raise RuntimeError(
"nccl_m2n data-plane preflight rejected before NCCL init: "
+ "; ".join(errors)
)
def validate_local_buffer(m2n: Any, meta: M2NParamMeta, buffer: Any) -> None:
"""Check source dtype/shape while delegating buffer validity to M2N."""
# A raw pointer carries no inspectable metadata. The source owns its
# correctness; _send supplies the declared dtype and local shape to M2N.
if isinstance(buffer, int):
return
mesh, placements = resolve_layout(
meta.source_layout.mesh, meta.source_layout.placements
)
actual = m2n.DistTensor(
buffer,
mesh=to_mesh(m2n, mesh),
placements=to_placements(m2n, placements),
)
declared = m2n.DistTensor(
None,
local_shape=actual.local_shape,
dtype=meta.dtype,
mesh=to_mesh(m2n, mesh),
placements=to_placements(m2n, placements),
)
if actual.dtype != declared.dtype:
raise ValueError(
f"parameter '{meta.name}' source returned dtype {actual.dtype}, "
f"but its metadata declared {declared.dtype}"
)
expected = expected_local_shape(meta)
if actual.local_shape != expected:
raise ValueError(
f"parameter '{meta.name}' source returned local shape "
f"{actual.local_shape}, but layout {meta.source_layout} implies "
f"{expected} from global shape {meta.shape}"
)
def to_mesh(m2n: Any, mesh: M2NMesh) -> Any:
return m2n.Mesh(mesh.dims, start_rank=mesh.start_rank)
@@ -279,7 +816,9 @@ def to_placements(m2n: Any, placements: Placements) -> list[Any]:
def publish_destination_placements(
comm: "PyNcclCommunicator",
comm: Any,
device: torch.device,
stream: torch.cuda.Stream,
first_worker_rank: int,
placements: "Sequence[Placements | None] | None",
num_parameters: int,
@@ -294,15 +833,15 @@ def publish_destination_placements(
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.
rank split, identically on both sides. The plan is published on the data
plane before any M2N reshard collective.
"""
replicated_sentinel = REPLICATE - 1
encoded = torch.full(
(num_parameters, MESH_NDIMS),
replicated_sentinel,
dtype=torch.int8,
device=comm.device,
device=device,
)
if comm.rank == first_worker_rank:
if placements is None or len(placements) != num_parameters:
@@ -314,25 +853,12 @@ def publish_destination_placements(
for placement in placements
]
if rows:
encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=comm.device))
encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=device))
comm.broadcast(encoded, src=first_worker_rank)
comm.broadcast(encoded, encoded, root=first_worker_rank, stream=stream)
return [
REPLICATED
if all(code == replicated_sentinel for code in row)
else cast(Placements, tuple(row))
for row in encoded.tolist()
]
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)
vllm/distributed/weight_transfer/m2n_engine.py
+288 -144
View File
@@ -20,7 +20,7 @@ inside `receive_weights` while the trainer is sending. Driving that from the
trainer is the trainer engine's job.
"""
from collections.abc import Iterator
from collections.abc import Iterator, Mapping
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast
@@ -34,20 +34,28 @@ from vllm.distributed.weight_transfer.base import (
)
from vllm.distributed.weight_transfer.m2n_common import (
DESTINATION_SHARD_AXIS,
M2N_WIRE_SCHEMA_VERSION,
MESH_NDIMS,
REPLICATED,
M2NLayout,
M2NMesh,
M2NParamMeta,
M2NWireParam,
Placements,
check_plan_limits,
check_complete_plan_agreement,
check_data_plane_agreement,
check_runtime_ready_agreement,
check_source_plan_agreement,
check_transferable,
comm_ptr,
create_m2n_communicator,
expected_local_shape,
import_m2n,
prepare_m2n_local_runtime,
publish_destination_placements,
resolve_layout,
source_plan_digest,
to_mesh,
to_placements,
validate_layout,
validate_transfer_plan,
)
from vllm.distributed.weight_transfer.m2n_layout import (
M2NDestination,
@@ -56,12 +64,14 @@ from vllm.distributed.weight_transfer.m2n_layout import (
from vllm.distributed.weight_transfer.nccl_common import (
NCCLWeightTransferInitInfo,
decode_nccl_unique_id,
worker_init_process_group,
worker_init_metadata_group,
)
from vllm.logger import init_logger
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
logger = init_logger(__name__)
__all__ = [
"M2NWeightTransferInitInfo",
@@ -84,59 +94,99 @@ class M2NWeightTransferInitInfo(WeightTransferInitInfo):
plain JSON so the HTTP control plane carries it unchanged.
"""
master_address: str | None = None
master_port: int | None = None
schema_version: int
master_address: str
master_port: int
rank_offset: int
"""First worker rank, i.e. the number of trainer ranks."""
world_size: int
"""Trainer ranks + all inference workers."""
src_mesh_dims: list[int]
"""The trainer's mesh, shared by every parameter (`start_rank` is 0: the
trainer occupies the front of the communicator)."""
dst_mesh_dims: list[int]
"""The inference mesh, starting at `rank_offset`. Declared rather than
derived so both sides describe the destination identically."""
names: list[str]
dtype_names: list[str]
shapes: list[list[int]]
src_placements: list[list[int] | None]
"""Per parameter, relative to `src_mesh_dims`; `None` means replicated."""
source_digest: str
"""SHA-256 of the ordered, versioned source parameter plan."""
params: list[dict[str, Any]]
"""Ordered per-parameter source metadata and layouts."""
nccl_unique_id_b64: str | None = field(default=None, repr=False)
max_cta: int | None = None
def __post_init__(self) -> None:
_ = self.nccl_unique_id_bytes
num_params = len(self.names)
for label, values in (
("dtype_names", self.dtype_names),
("shapes", self.shapes),
("src_placements", self.src_placements),
if not isinstance(self.master_address, str) or not self.master_address:
raise ValueError("`master_address` must name the metadata rendezvous")
if (
not isinstance(self.master_port, int)
or isinstance(self.master_port, bool)
or not 0 < self.master_port < 65536
):
raise ValueError("`master_port` must be a valid metadata port")
if (
not isinstance(self.rank_offset, int)
or isinstance(self.rank_offset, bool)
or not isinstance(self.world_size, int)
or isinstance(self.world_size, bool)
or self.rank_offset < 1
or self.rank_offset >= self.world_size
):
if len(values) != num_params:
raise ValueError(
f"`{label}` should be of the same size as `names`: "
f"got {len(values)} and {num_params}"
)
if self.rank_offset < 1 or self.rank_offset >= self.world_size:
raise ValueError(
f"`rank_offset` ({self.rank_offset}) must leave at least one "
f"trainer rank and one worker in world_size {self.world_size}"
)
num_workers = self.world_size - self.rank_offset
dst = tuple(self.dst_mesh_dims)
if len(dst) != MESH_NDIMS or dst[0] * dst[1] != num_workers:
try:
dst_mesh = M2NMesh(cast(tuple[int, int], dst), self.rank_offset)
except (TypeError, ValueError) as exc:
raise ValueError(f"invalid `dst_mesh_dims` {dst}: {exc}") from exc
if dst_mesh.size != num_workers:
raise ValueError(
f"`dst_mesh_dims` {dst} must be {MESH_NDIMS} dims covering the "
f"{num_workers} inference workers"
)
def parse_wire_params(self) -> tuple[M2NWireParam, ...]:
"""Parse and authenticate the ordered source plan."""
if self.schema_version != M2N_WIRE_SCHEMA_VERSION:
raise ValueError(
"unsupported nccl_m2n init schema version "
f"{self.schema_version}; expected {M2N_WIRE_SCHEMA_VERSION}"
)
if not isinstance(self.params, list):
raise TypeError("`params` must be a list of wire parameter mappings")
if not isinstance(self.source_digest, str) or not self.source_digest:
raise ValueError("`source_digest` must be a non-empty string")
wire_params: list[M2NWireParam] = []
names: set[str] = set()
for index, value in enumerate(self.params):
if not isinstance(value, Mapping):
raise TypeError(
f"`params[{index}]` must be a mapping, got {type(value).__name__}"
)
param = M2NWireParam.from_dict(value)
if param.name in names:
raise ValueError(
f"nccl_m2n init params contain duplicate parameter '{param.name}'"
)
names.add(param.name)
wire_params.append(param)
parsed = tuple(wire_params)
computed_digest = source_plan_digest(parsed)
if self.source_digest != computed_digest:
raise ValueError(
"nccl_m2n init source digest does not match params: "
f"computed {computed_digest}, declared {self.source_digest}"
)
return parsed
@property
def nccl_unique_id_bytes(self) -> bytes | None:
if self.nccl_unique_id_b64 is None:
return None
return decode_nccl_unique_id(
master_address=self.master_address,
master_port=self.master_port,
master_address=None,
master_port=None,
nccl_unique_id_b64=self.nccl_unique_id_b64,
ctx="M2NWeightTransferInitInfo",
ctx="M2NWeightTransferInitInfo data plane",
)
@@ -182,33 +232,172 @@ class M2NWeightTransferEngine(
model: torch.nn.Module,
) -> None:
super().__init__(config, vllm_config, device, model)
self.model_update_group: PyNcclCommunicator | None = None
self.model_update_group: Any = None
self._m2n: Any = None
self._handle: Any = None
self._metas: list[M2NParamMeta] = []
self._src_mesh: M2NMesh | None = None
self._dst_mesh: M2NMesh | None = None
self._parameter_destinations: list[M2NDestination] = []
self._index: dict[str, int] = {}
self._uses_load_weights = False
def init_transfer_engine(self, init_info: M2NWeightTransferInitInfo) -> None:
"""Join the trainer's communicator and rebuild the transfer plan.
Every precondition (dtype, tensor rank, divisibility) is checked here so
a bad plan fails the init RPC instead of hanging the first collective.
"""
self._m2n = import_m2n()
self._metas = []
self._src_mesh = M2NMesh(
cast(tuple[int, int], tuple(init_info.src_mesh_dims)), 0
)
if self._src_mesh.size != init_info.rank_offset:
raise ValueError(
f"source mesh covers {self._src_mesh.size} ranks, but there are "
f"{init_info.rank_offset} trainer ranks"
"""Agree on source and local readiness before constructing NCCL."""
if self._handle is not None or self.model_update_group is not None:
raise RuntimeError(
"nccl_m2n weight transfer is already initialized; call "
"shutdown_weight_transfer_engine() before initializing again"
)
rendezvous_info = NCCLWeightTransferInitInfo(
master_address=init_info.master_address,
master_port=init_info.master_port,
rank_offset=init_info.rank_offset,
world_size=init_info.world_size,
)
metadata_group = worker_init_metadata_group(
rendezvous_info, self.parallel_config
)
wire_params: tuple[M2NWireParam, ...] = ()
source_error: str | None = None
try:
self._m2n = import_m2n()
wire_params = init_info.parse_wire_params()
self._prepare_source_plan(wire_params, init_info.rank_offset)
except Exception as exc:
source_error = f"{type(exc).__name__}: {exc}"
self._metas = []
check_source_plan_agreement(
metadata_group,
wire_params,
expected_digest=init_info.source_digest,
local_error=source_error,
)
destination_error: str | None = None
try:
self._prepare_destination_plan(init_info)
except Exception as exc:
destination_error = f"{type(exc).__name__}: {exc}"
self._parameter_destinations = []
unique_id_bytes: bytes | None = None
data_plane_error: str | None = None
try:
unique_id_bytes = init_info.nccl_unique_id_bytes
except Exception as exc:
data_plane_error = f"{type(exc).__name__}: {exc}"
check_data_plane_agreement(
metadata_group,
unique_id_bytes,
data_plane_error,
)
handle = None
group = None
try:
runtime_error = destination_error
nccl = None
device = None
if runtime_error is None:
try:
nccl, device, handle = prepare_m2n_local_runtime(
self._m2n, init_info.max_cta
)
except Exception as exc:
runtime_error = f"{type(exc).__name__}: {exc}"
check_runtime_ready_agreement(metadata_group, runtime_error)
assert nccl is not None and device is not None and handle is not None
group = create_m2n_communicator(nccl, metadata_group, unique_id_bytes)
# Keep the merged destination-placement publication unchanged: source
# agreement is independent of worker-side destination selection.
mine = [
destination.placements for destination in self._parameter_destinations
]
agreed = publish_destination_placements(
group,
self.device,
torch.cuda.current_stream(),
init_info.rank_offset,
mine,
len(self._metas),
)
complete_plan_error: str | None = None
try:
if agreed != mine:
mismatched = next(
meta.name
for meta, ours, theirs in zip(self._metas, mine, agreed)
if ours != theirs
)
raise RuntimeError(
"inference workers resolved different nccl_m2n "
f"destinations (first mismatch: '{mismatched}'); all "
"workers must run the same model and parallel config"
)
assert self._dst_mesh is not None
validate_transfer_plan(self._metas, self._dst_mesh, agreed)
except Exception as exc:
complete_plan_error = f"{type(exc).__name__}: {exc}"
assert self._dst_mesh is not None
check_complete_plan_agreement(
metadata_group,
init_info.source_digest,
self._dst_mesh,
agreed,
complete_plan_error,
)
except BaseException:
if group is not None:
try:
# Transport cleanup only; reload validity belongs to the
# model-update lifecycle, not this backend.
group.abort()
except BaseException:
logger.exception(
"failed to abort partially initialized nccl_m2n communicator"
)
if handle is not None:
try:
handle.destroy()
except BaseException:
logger.exception(
"failed to destroy partially initialized nccl_m2n handle"
)
raise
self._handle = handle
self.model_update_group = group
def _prepare_source_plan(
self,
wire_params: tuple[M2NWireParam, ...],
num_trainer_ranks: int,
) -> None:
"""Validate worker-visible source metadata without constructing NCCL."""
self._metas = []
for param in wire_params:
dtype = getattr(torch, param.dtype_name, None)
if not isinstance(dtype, torch.dtype):
raise ValueError(
f"parameter '{param.name}' has invalid dtype name "
f"'{param.dtype_name}'; expected the name of a torch.dtype"
)
check_transferable(param.name, dtype, param.shape)
layout = M2NLayout(M2NMesh(param.src_mesh_dims, 0), param.src_placements)
if layout.mesh.size != num_trainer_ranks:
raise ValueError(
f"parameter '{param.name}' source mesh covers "
f"{layout.mesh.size} ranks, but there are "
f"{num_trainer_ranks} trainer ranks"
)
meta = M2NParamMeta(param.name, dtype, param.shape, layout)
expected_local_shape(meta)
self._metas.append(meta)
def _prepare_destination_plan(self, init_info: M2NWeightTransferInitInfo) -> None:
"""Resolve the merged-#51520 in-place/full-fallback destination plan."""
# The inference topology, as declared by the trainer. This is the mesh
# a sharded destination is placed over; a replicated one is described
# by the alternate descriptor `resolve_layout` derives from it, which
@@ -217,33 +406,6 @@ class M2NWeightTransferEngine(
cast(tuple[int, int], tuple(init_info.dst_mesh_dims)),
init_info.rank_offset,
)
for name, name_dtype, shape, placements in zip(
init_info.names,
init_info.dtype_names,
init_info.shapes,
init_info.src_placements,
):
dtype = getattr(torch, name_dtype, None)
if not isinstance(dtype, torch.dtype):
raise ValueError(
f"parameter '{name}' has invalid dtype name "
f"'{name_dtype}'; expected the name of a torch.dtype"
)
check_transferable(name, dtype, shape)
codes = (
REPLICATED
if placements is None
else cast(Placements, tuple(placements))
)
src_mesh, src_placements = resolve_layout(
self._src_mesh, codes, f"parameter '{name}' source placements"
)
validate_layout(src_mesh, src_placements, shape, "source")
# Destination checks are deferred until loader-aware resolution;
# assuming replication here can reject a valid sharded plan.
self._metas.append(M2NParamMeta(name, dtype, tuple(shape), codes))
parallel_config = self.parallel_config
shard_axis_size = self._dst_mesh.dims[DESTINATION_SHARD_AXIS]
quantization = getattr(self.model_config, "quantization", None)
@@ -266,27 +428,11 @@ class M2NWeightTransferEngine(
shard_axis_size=shard_axis_size,
allow_direct=allow_direct,
)
# Check the resolved layout rather than a provisional replicated one.
# For example, 32 source shards feeding 4 destination shards may need
# only 8 sources per destination, while a replicated plan needs all 32.
for meta, destination in zip(self._metas, self._parameter_destinations):
src_mesh, src_placements = resolve_layout(
self._src_mesh,
meta.placements,
f"parameter '{meta.name}' source placements",
)
dst_mesh, dst_placements = resolve_layout(
self._dst_mesh,
destination.placements,
f"parameter '{meta.name}' destination placements",
)
validate_layout(dst_mesh, dst_placements, meta.shape, "destination")
check_plan_limits(
(src_mesh, src_placements),
(dst_mesh, dst_placements),
meta.name,
)
validate_transfer_plan(
self._metas,
self._dst_mesh,
[destination.placements for destination in self._parameter_destinations],
)
# Update requests carry names only, so cache their plan indices. Track
# whether any fallback entry requires the `load_weights` lifecycle.
self._index = {meta.name: i for i, meta in enumerate(self._metas)}
@@ -294,38 +440,6 @@ class M2NWeightTransferEngine(
not destination.direct for destination in self._parameter_destinations
)
self.model_update_group = worker_init_process_group(
NCCLWeightTransferInitInfo(
master_address=init_info.master_address,
master_port=init_info.master_port,
nccl_unique_id_b64=init_info.nccl_unique_id_b64,
rank_offset=init_info.rank_offset,
world_size=init_info.world_size,
),
parallel_config,
)
# Destinations are per-parameter and depend on the inference model, so
# the trainer cannot derive them. The first worker publishes the plan
# and every other rank checks its own against it — a silent
# disagreement between workers would mean mismatched collectives.
mine = [destination.placements for destination in self._parameter_destinations]
agreed = publish_destination_placements(
self.model_update_group, init_info.rank_offset, mine, len(self._metas)
)
if agreed != mine:
mismatched = next(
meta.name for meta, a, b in zip(self._metas, mine, agreed) if a != b
)
raise RuntimeError(
"inference workers resolved different nccl_m2n destinations "
f"(first mismatch: '{mismatched}'); all workers must run the "
"same model and parallel config"
)
self._handle = self._m2n.Handle.create(
self._m2n.Config(max_cta=init_info.max_cta)
)
def start_weight_update(self) -> None:
"""Set up layerwise reloading, but only if some parameter needs it.
@@ -335,7 +449,9 @@ class M2NWeightTransferEngine(
"""
if not self._uses_load_weights:
return
from vllm.model_executor.model_loader.reload import initialize_layerwise_reload
from vllm.model_executor.model_loader.reload import (
initialize_layerwise_reload,
)
initialize_layerwise_reload(self.model)
@@ -354,6 +470,10 @@ class M2NWeightTransferEngine(
"nccl_m2n weight transfer not initialized. "
"Call init_transfer_engine() first."
)
if not isinstance(update_info.names, list) or any(
not isinstance(name, str) for name in update_info.names
):
raise TypeError("nccl_m2n update names must be a list of strings")
requested = []
for name in update_info.names:
@@ -369,7 +489,6 @@ class M2NWeightTransferEngine(
disable_mtp_completeness_check,
)
comm = comm_ptr(self.model_update_group)
stream = torch.cuda.current_stream()
# Reshard lazily so one `load_weights` invocation sees the complete
@@ -385,7 +504,13 @@ class M2NWeightTransferEngine(
meta.shape, dtype=meta.dtype, device=self.device
)
self._reshard(comm, stream, meta, destination.placements, buffer)
self._reshard(
self.model_update_group,
stream,
meta,
destination.placements,
buffer,
)
if destination.tensor is None:
# `load_weights` reads on the host stream, so the transfer
@@ -410,7 +535,7 @@ class M2NWeightTransferEngine(
def _reshard(
self,
comm: int,
comm: Any,
stream: torch.cuda.Stream,
src_meta: M2NParamMeta,
dst_placements: Placements | None,
@@ -421,11 +546,10 @@ class M2NWeightTransferEngine(
`dst_placements` describes the buffer's logical placement over the
worker mesh; `dst_buffer` is this rank's physical storage.
"""
assert self._src_mesh is not None
assert self._dst_mesh is not None
m2n = self._m2n
src_mesh, resolved_src_placements = resolve_layout(
self._src_mesh, src_meta.placements
src_meta.source_layout.mesh, src_meta.source_layout.placements
)
dst_mesh, resolved_dst_placements = resolve_layout(
self._dst_mesh, dst_placements
@@ -444,12 +568,32 @@ class M2NWeightTransferEngine(
def shutdown(self) -> None:
"""Finish pending GPU work and release M2N communication state."""
if self._handle is not None:
# M2N does not synchronize caller streams on finalize.
torch.accelerator.synchronize()
self._handle.destroy()
self._handle = None
if self.model_update_group is not None:
self.model_update_group.destroy()
self.model_update_group = None
handle, self._handle = self._handle, None
group, self.model_update_group = self.model_update_group, None
self._m2n = None
self._metas = []
self._dst_mesh = None
self._parameter_destinations = []
self._index = {}
self._uses_load_weights = False
first_error: BaseException | None = None
def attempt(action: Any, message: str) -> None:
nonlocal first_error
try:
action()
except BaseException as exc:
if first_error is None:
first_error = exc
logger.exception(message)
if handle is not None or group is not None:
# M2N does not synchronize caller streams on finalize.
attempt(torch.accelerator.synchronize, "failed to synchronize nccl_m2n")
if handle is not None:
attempt(handle.destroy, "failed to destroy nccl_m2n handle")
if group is not None:
attempt(group.abort, "failed to abort nccl_m2n communicator")
if first_error is not None:
raise first_error
vllm/distributed/weight_transfer/m2n_source.py
+88 -36
View File
@@ -5,21 +5,22 @@
m2n plans a transfer from both sides' layouts, so the trainer has to say how it
holds each parameter — which the base `WeightSource` / `ParamMeta` pair does not
express. This module supplies the m2n flavor of both, and the DTensor-backed
implementation that covers the common trainer.
A source declares its mesh once (`mesh()`) and a placement per parameter, since
one topology describes every tensor on the side.
implementation that covers the common trainer. Each parameter owns its source
mesh and placements because dense and expert weights can use different
factorizations over the same trainer ranks.
"""
from collections.abc import Iterator
from collections.abc import Callable, Iterator, Sequence
from dataclasses import dataclass
from typing import Any
import torch
from vllm.distributed.weight_transfer.base import ParamMeta, WeightSource
from vllm.distributed.weight_transfer.base import WeightSource
from vllm.distributed.weight_transfer.m2n_common import (
REPLICATE,
REPLICATED,
M2NLayout,
M2NMesh,
M2NParamMeta,
Placements,
@@ -27,6 +28,8 @@ from vllm.distributed.weight_transfer.m2n_common import (
__all__ = [
"M2NWeightSource",
"M2NManifestEntry",
"ManifestM2NWeightSource",
"DTensorModuleSource",
"mesh_from_tensor",
"placements_from_tensor",
@@ -40,17 +43,76 @@ class M2NWeightSource(WeightSource):
materialized full tensor: gathering is exactly the cost m2n removes.
"""
def mesh(self) -> M2NMesh:
"""The trainer's rank topology, shared by every parameter."""
def metadata(self) -> list[M2NParamMeta]: # type: ignore[override]
"""Name, dtype, full shape, and source layout for each parameter."""
raise NotImplementedError
def metadata(self) -> list[ParamMeta]:
"""Name, dtype, full shape, and trainer placement for each parameter."""
def __iter__(self) -> Iterator[tuple[str, Any]]: # type: ignore[override]
"""Yield `(name, local buffer)` pairs in metadata order.
Providers must enqueue materialization on the caller's current CUDA
stream, or return a buffer whose writes are already visible to it.
Buffers may be torch tensors, CUDA Array Interface objects, or raw
device pointers accepted by nccl4py. Providers own correct
communicator-device placement. Raw-pointer providers additionally own
allocation lifetime, dtype, and shape correctness.
"""
raise NotImplementedError
def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]:
"""Yield `(name, local shard)` pairs in the same order as `metadata()`."""
raise NotImplementedError
@dataclass(frozen=True)
class M2NManifestEntry:
"""One immutable source descriptor with a lazy rank-local tensor view."""
key: str
dtype: torch.dtype
global_shape: tuple[int, ...]
source_layout: M2NLayout
local_tensor: Callable[[], Any]
def __post_init__(self) -> None:
object.__setattr__(self, "global_shape", tuple(self.global_shape))
if not self.key:
raise ValueError("M2N manifest key must not be empty")
if not isinstance(self.dtype, torch.dtype):
raise TypeError(
f"manifest entry '{self.key}' has invalid dtype {self.dtype!r}"
)
if not isinstance(self.source_layout, M2NLayout):
raise TypeError(
f"manifest entry '{self.key}' source_layout must be an M2NLayout"
)
if not callable(self.local_tensor):
raise TypeError(
f"manifest entry '{self.key}' local_tensor must be callable"
)
@property
def metadata(self) -> M2NParamMeta:
"""Return the serializable portion consumed during initialization."""
return M2NParamMeta(
self.key,
self.dtype,
self.global_shape,
self.source_layout,
)
class ManifestM2NWeightSource(M2NWeightSource):
"""An ordered M2N source backed by lazy, rank-local manifest entries."""
def __init__(self, entries: Sequence[M2NManifestEntry]) -> None:
self._entries = tuple(entries)
self._metadata = tuple(entry.metadata for entry in self._entries)
def metadata(self) -> list[M2NParamMeta]: # type: ignore[override]
"""Return a fresh list while preserving the immutable manifest order."""
return list(self._metadata)
def __iter__(self) -> Iterator[tuple[str, Any]]: # type: ignore[override]
"""Materialize one rank-local buffer at a time."""
for entry in self._entries:
yield entry.key, entry.local_tensor()
def _placement_code(placement: Any) -> int:
@@ -69,7 +131,8 @@ def _placement_code(placement: Any) -> int:
def mesh_from_tensor(tensor: torch.Tensor, num_trainer_ranks: int) -> M2NMesh:
"""The mesh a parameter lives on, as an `M2NMesh`.
A plain tensor is identical on every trainer rank, so it spans all of them.
Tensors without DeviceMesh metadata are treated as replicated. Implicitly
sharded tensors require a custom source with an explicit layout.
"""
device_mesh = getattr(tensor, "device_mesh", None)
if device_mesh is None:
@@ -122,29 +185,18 @@ class DTensorModuleSource(M2NWeightSource):
self._module = module
self._num_trainer_ranks = num_trainer_ranks
def mesh(self) -> M2NMesh:
"""The mesh every sharded parameter agrees on.
Replicated parameters span the trainer without a device mesh of their
own, so they never decide this; a model whose *sharded* parameters
disagree cannot be described by one side mesh and is rejected.
"""
meshes = {
mesh_from_tensor(p, self._num_trainer_ranks)
for _, p in self._module.named_parameters()
if placements_from_tensor(p) is not REPLICATED
}
if len(meshes) > 1:
raise ValueError(
"nccl_m2n needs one mesh per side, but this module's sharded "
f"parameters span several: {sorted(m.dims for m in meshes)}"
)
return meshes.pop() if meshes else M2NMesh((self._num_trainer_ranks, 1), 0)
def metadata(self) -> list[ParamMeta]:
"""Read global shape/dtype and placements without gathering shards."""
def metadata(self) -> list[M2NParamMeta]: # type: ignore[override]
"""Read each global shape and source layout without gathering shards."""
return [
M2NParamMeta(name, p.dtype, tuple(p.shape), placements_from_tensor(p))
M2NParamMeta(
name,
p.dtype,
tuple(p.shape),
M2NLayout(
mesh_from_tensor(p, self._num_trainer_ranks),
placements_from_tensor(p),
),
)
for name, p in self._module.named_parameters()
]
vllm/distributed/weight_transfer/m2n_trainer.py
+416 -203
View File
@@ -13,6 +13,8 @@ engine -- where the `TrainerWeightTransferEngine` contract puts it -- rather tha
in every caller.
"""
import socket
from collections.abc import Mapping
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar, Self, cast
@@ -26,39 +28,43 @@ from vllm.distributed.weight_transfer.base import (
WeightSource,
)
from vllm.distributed.weight_transfer.m2n_common import (
M2N_WIRE_SCHEMA_VERSION,
M2NMesh,
M2NParamMeta,
M2NWireParam,
Placements,
check_complete_plan_agreement,
check_data_plane_agreement,
check_runtime_ready_agreement,
check_source_plan_agreement,
check_transferable,
comm_ptr,
create_m2n_communicator,
expected_local_shape,
import_m2n,
prepare_m2n_local_runtime,
publish_destination_placements,
resolve_layout,
source_plan_digest,
to_mesh,
to_placements,
validate_layout,
validate_local_buffer,
validate_transfer_plan,
)
from vllm.distributed.weight_transfer.m2n_source import M2NWeightSource
from vllm.distributed.weight_transfer.nccl_common import (
decode_nccl_unique_id,
trainer_init,
uid_init_process_group,
stateless_init_metadata_group,
)
from vllm.logger import init_logger
if TYPE_CHECKING:
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
from vllm.distributed.utils import StatelessProcessGroup
logger = init_logger(__name__)
__all__ = ["M2NTrainerInitInfo", "M2NTrainerWeightTransferEngine"]
def _dtype_name(dtype: torch.dtype) -> str:
"""Short name for a dtype, e.g. `bfloat16`, for the worker-side init dict."""
return str(dtype).split(".")[-1]
@dataclass(kw_only=True)
class M2NTrainerInitInfo(TrainerInitInfo):
"""Trainer-side init info for nccl_m2n.
@@ -70,8 +76,8 @@ class M2NTrainerInitInfo(TrainerInitInfo):
backend: ClassVar[str] = "nccl_m2n"
master_address: str | None = None
master_port: int | None = None
master_address: str
master_port: int
world_size: int
"""Trainer ranks + all inference workers."""
num_trainer_ranks: int = 1
@@ -79,20 +85,34 @@ class M2NTrainerInitInfo(TrainerInitInfo):
"""How the inference ranks are laid out: axis 0 replicates and axis 1
shards. The trainer declares it so both sides describe the destination
identically; defaults to a flat `(num_workers, 1)`, which is all a
replicated destination needs."""
max_cta: int | None = None
replicated destination needs.
Destination meshes remain common to every parameter; per-parameter
destination meshes are intentionally future work.
"""
nccl_unique_id_b64: str | None = field(default=None, repr=False)
"""Optional UID for the NCCL data plane; TCP remains the metadata plane."""
max_cta: int | None = None
listen_socket: socket.socket | None = field(default=None, repr=False, compare=False)
"""Optional rank-zero TCPStore reservation; never sent to workers."""
@property
def destination_mesh_dims(self) -> tuple[int, int]:
"""`dst_mesh_dims`, or a flat mesh over every inference worker."""
if self.dst_mesh_dims is not None:
return self.dst_mesh_dims
return cast(tuple[int, int], tuple(self.dst_mesh_dims))
return (self.world_size - self.num_trainer_ranks, 1)
def __post_init__(self) -> None:
"""Reject rank counts or destination meshes that cannot form a group."""
_ = self.nccl_unique_id_bytes
if not isinstance(self.master_address, str) or not self.master_address:
raise ValueError("`master_address` must name the metadata rendezvous")
if (
not isinstance(self.master_port, int)
or isinstance(self.master_port, bool)
or not 0 < self.master_port < 65536
):
raise ValueError("`master_port` must be a valid metadata port")
if not 0 < self.num_trainer_ranks < self.world_size:
raise ValueError(
f"`num_trainer_ranks` ({self.num_trainer_ranks}) must leave at "
@@ -101,7 +121,11 @@ class M2NTrainerInitInfo(TrainerInitInfo):
num_workers = self.world_size - self.num_trainer_ranks
if self.dst_mesh_dims is not None:
dims = tuple(self.dst_mesh_dims)
if len(dims) != 2 or dims[0] * dims[1] != num_workers:
try:
mesh = M2NMesh(cast(tuple[int, int], dims), self.num_trainer_ranks)
except (TypeError, ValueError) as exc:
raise ValueError(f"invalid `dst_mesh_dims` {dims}: {exc}") from exc
if mesh.size != num_workers:
raise ValueError(
f"`dst_mesh_dims` {dims} must be two dims covering the "
f"{num_workers} inference workers"
@@ -111,14 +135,23 @@ class M2NTrainerInitInfo(TrainerInitInfo):
f"trainer `rank` ({self.rank}) must be below "
f"`num_trainer_ranks` ({self.num_trainer_ranks})"
)
if self.listen_socket is not None:
if self.rank != 0:
raise ValueError("only trainer rank zero may reserve the TCPStore")
if not isinstance(self.listen_socket, socket.socket):
raise TypeError("listen_socket must be a socket.socket")
if self.listen_socket.fileno() < 0:
raise ValueError("listen_socket is already closed")
@property
def nccl_unique_id_bytes(self) -> bytes | None:
if self.nccl_unique_id_b64 is None:
return None
return decode_nccl_unique_id(
master_address=self.master_address,
master_port=self.master_port,
master_address=None,
master_port=None,
nccl_unique_id_b64=self.nccl_unique_id_b64,
ctx="M2NTrainerInitInfo",
ctx="M2NTrainerInitInfo data plane",
)
@@ -136,22 +169,20 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI
*,
client: VLLMWeightSyncClient,
source: M2NWeightSource,
is_controller: bool,
num_trainer_ranks: int,
is_sender: bool,
) -> None:
"""Hold the client and source; `trainer_init` fills in the group."""
super().__init__(client=client, source=source, is_sender=is_controller)
self.source: M2NWeightSource = source
self.is_controller = is_controller
self.num_trainer_ranks = num_trainer_ranks
self.group: PyNcclCommunicator | None = None
super().__init__(client=client, source=source, is_sender=is_sender)
self.group: Any = None
self._m2n: Any = None
self._handle: Any = None
self._metas: list[M2NParamMeta] = []
self._src_mesh: M2NMesh | None = None
self._wire_params: list[M2NWireParam] = []
self._source_digest = ""
self._dst_mesh: M2NMesh | None = None
self._dst_placements: list[Placements | None] = []
self._executor: ThreadPoolExecutor | None = None
self._metadata_group: StatelessProcessGroup | None = None
self._num_trainer_ranks = 0
@classmethod
def trainer_init(
@@ -168,204 +199,351 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI
build local state and join the communicator. Every rank participates
in every reshard and sends its local shard.
The ordering here is not arbitrary -- see the comments inline. In
short: validate before anything blocks; start the init RPC without
waiting (workers join the communicator inside that RPC, so it cannot
return until this side has joined too); then join the NCCL
communicator; then create the M2N handle; and only then wait for the RPC
to finish.
Build a local validation envelope, start the worker RPC without
waiting, join the metadata rendezvous, and require unanimous source,
data-plane, and runtime readiness before constructing NCCL.
"""
# The base WeightSource carries name/dtype/shape but not layout, which
# is not enough to plan a reshard. Fail here rather than at the first
# missing attribute deep inside the send loop.
if not isinstance(source, M2NWeightSource):
raise TypeError(
"nccl_m2n needs per-parameter layouts, so its source must be an "
f"M2NWeightSource (e.g. DTensorModuleSource); got "
f"{type(source).__name__}"
)
# Trainer-side weight transfer engine for this rank. `__init__` only
# holds the client and source; the communicator, meshes, and M2N handle
# are filled in below.
engine = cls(
client=client,
source=source,
is_controller=init_info.rank == 0,
num_trainer_ranks=init_info.num_trainer_ranks,
source=cast(M2NWeightSource, source),
is_sender=init_info.is_sender,
)
engine._m2n = import_m2n()
engine._src_mesh = source.mesh()
if engine._src_mesh.size != init_info.num_trainer_ranks:
raise ValueError(
f"source mesh covers {engine._src_mesh.size} ranks, but there "
f"are {init_info.num_trainer_ranks} trainer ranks"
)
# Validate the whole plan before any rendezvous. A layout m2n cannot
# express is a mismatched collective at send time, which hangs rather
# than raises -- so every precondition is checked while an exception can
# still propagate to the caller.
engine._metas = cast(list[M2NParamMeta], source.metadata())
for meta in engine._metas:
check_transferable(meta.name, meta.dtype, meta.shape)
mesh, placements = resolve_layout(engine._src_mesh, meta.placements)
validate_layout(mesh, placements, meta.shape, "source")
rendezvous: Future | None = None
if engine.is_controller:
engine._executor = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="m2n-weight-sync"
)
# Only trainer rank 0 starts the control-plane RPC. Workers join
# the communicator inside it, so submit it first, join below, and
# wait for it at the end.
rendezvous = engine._executor.submit(
client.init_weight_transfer_engine,
engine._worker_init_info(init_info),
)
# Every trainer rank joins -- reshard is collective over the whole
# communicator. Trainer ranks take
# [0, num_trainer_ranks) and the workers follow, which is the
# contiguous-interval layout m2n requires of both meshes.
unique_id_bytes = init_info.nccl_unique_id_bytes
if unique_id_bytes is not None:
engine.group = uid_init_process_group(
unique_id_bytes,
rank=init_info.rank,
world_size=init_info.world_size,
device=torch.accelerator.current_device_index(),
)
else:
assert init_info.master_address is not None
assert init_info.master_port is not None
engine.group = trainer_init(
{
"master_address": init_info.master_address,
"master_port": init_info.master_port,
"world_size": init_info.world_size,
},
rank=init_info.rank,
)
# The trainer declares the inference topology; the worker uses exactly
# this, so neither side has to infer the other's factorization.
engine._dst_mesh = M2NMesh(
init_info.destination_mesh_dims, init_info.num_trainer_ranks
)
engine._dst_placements = engine._receive_destination_plan(
init_info.num_trainer_ranks
)
logger.info(
"nccl_m2n trainer ready: %d parameters, trainer ranks [0, %d), "
"inference ranks [%d, %d)",
len(engine._metas),
init_info.num_trainer_ranks,
init_info.num_trainer_ranks,
init_info.world_size,
)
# Created after the group so a failed rendezvous does not leave a live
# handle behind; shutdown() releases it once caller streams are synced.
engine._handle = engine._m2n.Handle.create(
engine._m2n.Config(max_cta=init_info.max_cta)
)
# Now safe to collect: both sides have joined, so the RPC has returned
# or raised. Doing this last means a worker-side init failure surfaces
# here, as an exception from trainer_init, rather than at the first send.
if rendezvous is not None:
rendezvous.result()
engine._initialize(init_info, source)
return engine
def _receive_destination_plan(
self, first_worker_rank: int
) -> list[Placements | None]:
"""Receive and validate the inference-side per-parameter placements."""
if self.group is None or self._dst_mesh is None:
raise RuntimeError(
"destination plan requires the communicator and destination mesh"
def _initialize(
self,
init_info: M2NTrainerInitInfo,
source: WeightSource | None,
) -> None:
"""Run source validation, metadata agreement, and NCCL rendezvous."""
self._num_trainer_ranks = init_info.num_trainer_ranks
local_source_error: str | None = None
try:
if not isinstance(source, M2NWeightSource):
raise TypeError(
"nccl_m2n needs per-parameter layouts, so its source must "
f"be an M2NWeightSource; got {type(source).__name__}"
)
self._m2n = import_m2n()
self._prepare_source_plan(source, init_info.num_trainer_ranks)
except Exception as exc:
local_source_error = f"{type(exc).__name__}: {exc}"
self._metas = []
self._wire_params = []
self._source_digest = source_plan_digest([])
executor: ThreadPoolExecutor | None = None
init_future: Future[None] | None = None
if self.is_sender:
executor = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="m2n-weight-sync"
)
init_future = executor.submit(
self.client.init_weight_transfer_engine,
self._worker_init_info(init_info),
)
handle = None
group = None
try:
reservation = init_info.listen_socket
try:
metadata_group = stateless_init_metadata_group(
init_info.master_address,
init_info.master_port,
init_info.rank,
init_info.world_size,
listen_socket=reservation,
)
finally:
init_info.listen_socket = None
if reservation is not None:
reservation.close()
plan_digest = check_source_plan_agreement(
metadata_group,
self._wire_params,
expected_digest=self._source_digest,
local_error=local_source_error,
)
# Retain the common destination mesh and the merged NCCL publication
# of per-parameter placements.
self._dst_mesh = M2NMesh(
init_info.destination_mesh_dims, init_info.num_trainer_ranks
)
unique_id_bytes: bytes | None = None
data_plane_error: str | None = None
try:
unique_id_bytes = init_info.nccl_unique_id_bytes
except Exception as exc:
data_plane_error = f"{type(exc).__name__}: {exc}"
check_data_plane_agreement(
metadata_group,
unique_id_bytes,
data_plane_error,
)
runtime_error: str | None = None
nccl = None
device = None
try:
nccl, device, handle = prepare_m2n_local_runtime(
self._m2n, init_info.max_cta
)
except Exception as exc:
runtime_error = f"{type(exc).__name__}: {exc}"
check_runtime_ready_agreement(metadata_group, runtime_error)
assert nccl is not None and device is not None and handle is not None
group = create_m2n_communicator(nccl, metadata_group, unique_id_bytes)
destination_placements = self._receive_destination_plan(
group,
init_info.num_trainer_ranks,
torch.device("cuda", device),
)
complete_plan_error: str | None = None
try:
assert self._dst_mesh is not None
validate_transfer_plan(
self._metas, self._dst_mesh, destination_placements
)
except Exception as exc:
complete_plan_error = f"{type(exc).__name__}: {exc}"
assert self._dst_mesh is not None
check_complete_plan_agreement(
metadata_group,
plan_digest,
self._dst_mesh,
destination_placements,
complete_plan_error,
)
if init_future is not None:
init_future.result()
logger.info(
"nccl_m2n trainer ready: %d parameters, trainer ranks [0, %d), "
"inference ranks [%d, %d)",
len(self._metas),
init_info.num_trainer_ranks,
init_info.num_trainer_ranks,
init_info.world_size,
)
logger.info("nccl_m2n source plan digest: %s", plan_digest)
except BaseException:
if group is not None:
try:
# Transport cleanup only; reload validity belongs to the
# model-update lifecycle, not this backend.
group.abort()
except BaseException:
logger.exception(
"failed to abort partially initialized nccl_m2n communicator"
)
if handle is not None:
try:
handle.destroy()
except BaseException:
logger.exception(
"failed to destroy partially initialized nccl_m2n handle"
)
raise
finally:
if executor is not None:
executor.shutdown(wait=False, cancel_futures=True)
self._metadata_group = metadata_group
self._handle = handle
self.group = group
self._dst_placements = destination_placements
def _receive_destination_plan(
self, group: Any, first_worker_rank: int, device: torch.device
) -> list[Placements | None]:
"""Receive the inference-side per-parameter placements."""
# The first inference worker publishes one placement per parameter.
# Every trainer rank receives the complete plan in this single,
# initialization-time broadcast.
placements = publish_destination_placements(
self.group, first_worker_rank, None, len(self._metas)
group,
device,
torch.cuda.current_stream(),
first_worker_rank,
None,
len(self._metas),
)
# The send loop matches placements to metadata by position, so reject
# an incomplete or oversized plan before entering any M2N collective.
if len(placements) != len(self._metas):
raise RuntimeError(
f"inference side planned {len(placements)} parameters, but "
f"this trainer announced {len(self._metas)}"
)
# Check every worker-selected destination against the corresponding
# global parameter shape before the first transfer.
for meta, placement in zip(self._metas, placements):
mesh, resolved = resolve_layout(self._dst_mesh, placement)
validate_layout(mesh, resolved, meta.shape, "destination")
return placements
def _worker_init_info(self, init_info: M2NTrainerInitInfo) -> dict[str, Any]:
"""Handshake payload: rendezvous, both meshes, and the transfer plan."""
assert self._src_mesh is not None
payload: dict[str, Any] = {
"schema_version": M2N_WIRE_SCHEMA_VERSION,
"master_address": init_info.master_address,
"master_port": init_info.master_port,
"nccl_unique_id_b64": init_info.nccl_unique_id_b64,
"rank_offset": init_info.num_trainer_ranks,
"world_size": init_info.world_size,
"src_mesh_dims": list(self._src_mesh.dims),
"dst_mesh_dims": list(init_info.destination_mesh_dims),
"names": [m.name for m in self._metas],
"dtype_names": [_dtype_name(m.dtype) for m in self._metas],
"shapes": [list(m.shape) for m in self._metas],
"src_placements": [
None if m.placements is None else list(m.placements)
for m in self._metas
],
"source_digest": self._source_digest,
"params": [param.to_dict() for param in self._wire_params],
"max_cta": init_info.max_cta,
}
return {key: value for key, value in payload.items() if value is not None}
def _prepare_source_plan(
self, source: M2NWeightSource, num_trainer_ranks: int
) -> None:
"""Build and validate this rank's immutable source plan."""
self._metas = source.metadata()
names: set[str] = set()
for meta in self._metas:
if not isinstance(meta, M2NParamMeta):
raise TypeError(
"nccl_m2n source metadata must contain M2NParamMeta "
f"entries, got {type(meta).__name__}"
)
if meta.name in names:
raise ValueError(
"nccl_m2n source metadata contains duplicate parameter "
f"'{meta.name}'"
)
names.add(meta.name)
check_transferable(meta.name, meta.dtype, meta.shape)
layout = meta.source_layout
if layout.mesh.start_rank != 0:
raise ValueError(
f"parameter '{meta.name}' source mesh must start at rank 0, "
f"got {layout.mesh.start_rank}"
)
if layout.mesh.size != num_trainer_ranks:
raise ValueError(
f"parameter '{meta.name}' source mesh covers "
f"{layout.mesh.size} ranks, but there are "
f"{num_trainer_ranks} trainer ranks"
)
expected_local_shape(meta)
self._wire_params = [
M2NWireParam(
name=meta.name,
dtype_name=str(meta.dtype).split(".")[-1],
shape=meta.shape,
src_mesh_dims=meta.source_layout.mesh.dims,
src_placements=meta.source_layout.placements,
)
for meta in self._metas
]
self._source_digest = source_plan_digest(self._wire_params)
def send_weights(self) -> None:
"""Drive one update round: start, reshard concurrently, then finish."""
if self._handle is None or self.group is None:
raise RuntimeError("nccl_m2n trainer engine is not initialized")
update: Future | None = None
if self.is_controller:
self.client.start_weight_update()
# The workers reshard from inside `update_weights`, concurrently
# with the sends below.
assert self._executor is not None
update = self._executor.submit(
self.client.update_weights,
{"names": [m.name for m in self._metas]},
)
if not self.is_sender:
try:
self._send()
except BaseException:
try:
self._receive_round_outcome()
except BaseException:
logger.exception(
"failed to consume nccl_m2n round outcome after local failure"
)
raise
self._receive_round_outcome()
return
try:
self._send()
finally:
if update is not None:
self.client.start_weight_update()
executor = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="m2n-weight-sync"
)
try:
update = executor.submit(
self.client.update_weights,
{"names": [m.name for m in self._metas]},
)
# Surface an immediate worker-side rejection before trainers
# enter the collective sequence.
if update.done():
update.result()
self._send()
update.result()
if self.is_controller:
finally:
executor.shutdown(wait=False, cancel_futures=True)
self.client.finish_weight_update()
except BaseException as exc:
try:
self._publish_round_outcome(exc)
except BaseException:
logger.exception("failed to publish nccl_m2n round failure")
raise
self._publish_round_outcome(None)
def _publish_round_outcome(self, failure: BaseException | None) -> None:
"""Tell non-sender trainer ranks whether worker finalization passed."""
if self._num_trainer_ranks <= 1:
return
if self._metadata_group is None:
raise RuntimeError("nccl_m2n metadata group is unavailable")
envelope = {
"ok": failure is None,
"error": (
None if failure is None else f"{type(failure).__name__}: {failure}"
),
}
for rank in range(1, self._num_trainer_ranks):
self._metadata_group.send_obj(envelope, rank)
# Keep rank zero's TCPStore alive until every trainer consumed its
# outcome. This also keeps the per-rank message counters aligned.
expected_ack = {"phase": "round_outcome_consumed"}
for rank in range(1, self._num_trainer_ranks):
ack = self._metadata_group.recv_obj(rank)
if ack != expected_ack:
raise RuntimeError(
f"trainer rank {rank} sent an invalid round-outcome acknowledgement"
)
def _receive_round_outcome(self) -> None:
"""Wait until trainer rank zero observed every worker's finish RPC."""
if self._metadata_group is None:
raise RuntimeError("nccl_m2n metadata group is unavailable")
outcome = self._metadata_group.recv_obj(0)
self._metadata_group.send_obj({"phase": "round_outcome_consumed"}, 0)
if not isinstance(outcome, Mapping) or set(outcome) != {"ok", "error"}:
raise RuntimeError("trainer rank zero sent an invalid round outcome")
ok = outcome["ok"]
error = outcome["error"]
if not isinstance(ok, bool) or ok != (error is None):
raise RuntimeError("trainer rank zero sent an invalid round outcome")
if not ok:
if not isinstance(error, str) or not error:
raise RuntimeError("trainer rank zero sent an invalid round error")
raise RuntimeError(
"nccl_m2n worker update failed on trainer rank zero: " + error
)
def _send(self) -> None:
"""Reshard each local shard into its worker-planned destination."""
assert self.group is not None
assert self._src_mesh is not None
assert self._dst_mesh is not None
assert self.source is not None
m2n = self._m2n
comm = comm_ptr(self.group)
stream = torch.cuda.current_stream()
for index, (name, tensor) in enumerate(self.source):
device_index = torch.accelerator.current_device_index()
# Optimizer work may have used another CUDA stream. Make prior writes
# visible before a provider materializes current-stream local shards.
torch.accelerator.synchronize(device_index)
inflight_buffers: list[Any] = []
sent = 0
for index, (name, buffer) in enumerate(self.source):
if index >= len(self._metas):
raise RuntimeError(
"weight source yielded more parameters than metadata() "
f"declared (first extra parameter: '{name}')"
)
meta = self._metas[index]
if name != meta.name:
raise RuntimeError(
@@ -373,28 +551,47 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI
f"its metadata declared '{meta.name}'; the source must "
"iterate in a stable order"
)
shard = tensor.detach().contiguous()
src_mesh, src_placements = resolve_layout(self._src_mesh, meta.placements)
validate_local_buffer(m2n, meta, buffer)
src_mesh, src_placements = resolve_layout(
meta.source_layout.mesh, meta.source_layout.placements
)
dst_mesh, dst_placements = resolve_layout(
self._dst_mesh, self._dst_placements[index]
)
raw_pointer_args = {}
if isinstance(buffer, int):
raw_pointer_args = {
"src_local_shape": expected_local_shape(meta),
"src_dtype": meta.dtype,
}
m2n.reshard(
shard,
buffer,
None,
comm,
self.group,
stream,
src_mesh=to_mesh(m2n, src_mesh),
src_placements=to_placements(m2n, src_placements),
dst_mesh=to_mesh(m2n, dst_mesh),
dst_placements=to_placements(m2n, dst_placements),
handle=self._handle,
**raw_pointer_args,
)
# Let allocator-aware buffers recycle storage in stream order.
# Other object owners must stay alive until the final synchronize;
# raw-pointer lifetime remains the source's responsibility.
record_stream = getattr(buffer, "record_stream", None)
if callable(record_stream):
record_stream(stream)
elif not isinstance(buffer, int):
inflight_buffers.append(buffer)
sent += 1
if sent != len(self._metas):
missing = self._metas[sent].name
raise RuntimeError(
f"weight source yielded {sent} parameters, but metadata() "
f"declared {len(self._metas)} (first missing: '{missing}')"
)
# A custom source may yield temporary storage allocated on another
# stream. M2N retains only its device pointer for asynchronous work,
# so tell the allocator not to reuse that storage until this stream
# has finished with it. For storage allocated on this stream,
# record_stream is effectively a no-op.
shard.record_stream(stream)
# As in the other weight-transfer backends, wait before returning
# rather than relying on same-stream ordering. The caller may then
@@ -402,14 +599,30 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI
stream.synchronize()
def shutdown(self) -> None:
"""Destroy the m2n handle, join the helper thread, and drop the group."""
if self._handle is not None:
torch.accelerator.synchronize()
self._handle.destroy()
self._handle = None
if self._executor is not None:
self._executor.shutdown()
self._executor = None
if self.group is not None:
self.group.destroy()
self.group = None
"""Finish pending GPU work and release M2N communication state."""
handle, self._handle = self._handle, None
group, self.group = self.group, None
self._metadata_group = None
self._dst_placements = []
self._dst_mesh = None
self._m2n = None
first_error: BaseException | None = None
def attempt(action: Any, message: str) -> None:
nonlocal first_error
try:
action()
except BaseException as exc:
if first_error is None:
first_error = exc
logger.exception(message)
if handle is not None or group is not None:
attempt(torch.accelerator.synchronize, "failed to synchronize nccl_m2n")
if handle is not None:
attempt(handle.destroy, "failed to destroy nccl_m2n handle")
if group is not None:
attempt(group.abort, "failed to abort nccl_m2n communicator")
if first_error is not None:
raise first_error
vllm/distributed/weight_transfer/nccl_common.py
+46 -9
View File
@@ -8,6 +8,7 @@ The dense (`NCCLWeightTransferEngine`) and sparse
sparse engine does not have to subclass the dense one.
"""
import socket
from dataclasses import asdict, dataclass, field
from typing import TYPE_CHECKING, Protocol
@@ -17,6 +18,7 @@ import torch
if TYPE_CHECKING:
from vllm.config.parallel import ParallelConfig
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
from vllm.distributed.utils import StatelessProcessGroup
from vllm.distributed.device_communicators.pynccl_wrapper import (
NCCL_UNIQUE_ID_BYTES,
@@ -136,6 +138,26 @@ class NCCLRendezvous(Protocol):
world_size: int
def stateless_init_metadata_group(
master_address: str,
master_port: int,
rank: int,
world_size: int,
*,
listen_socket: socket.socket | None = None,
) -> "StatelessProcessGroup":
"""Create only the TCPStore group used for metadata agreement."""
from vllm.distributed.utils import StatelessProcessGroup
return StatelessProcessGroup.create(
host=master_address,
port=master_port,
rank=rank,
world_size=world_size,
listen_socket=listen_socket,
)
def stateless_init_process_group(
master_address: str,
master_port: int,
@@ -181,6 +203,29 @@ def uid_init_process_group(
)
def _worker_group_rank(rank_offset: int, parallel_config: "ParallelConfig") -> int:
"""Return this worker's rank in a trainer-to-worker process group."""
dp_rank = parallel_config.data_parallel_index
world_size_per_dp = parallel_config.world_size # TP * PP
rank_within_dp = parallel_config.rank
return rank_offset + dp_rank * world_size_per_dp + rank_within_dp
def worker_init_metadata_group(
init_info: NCCLWeightTransferInitInfo,
parallel_config: "ParallelConfig",
) -> "StatelessProcessGroup":
"""Join the trainer-to-worker TCP metadata group without constructing NCCL."""
assert init_info.master_address is not None
assert init_info.master_port is not None
return stateless_init_metadata_group(
init_info.master_address,
init_info.master_port,
_worker_group_rank(init_info.rank_offset, parallel_config),
init_info.world_size,
)
def worker_init_process_group(
init_info: NCCLWeightTransferInitInfo,
parallel_config: "ParallelConfig",
@@ -190,15 +235,7 @@ def worker_init_process_group(
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
rank = _worker_group_rank(init_info.rank_offset, parallel_config)
device = torch.accelerator.current_device_index()
unique_id_bytes = init_info.nccl_unique_id_bytes