diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 1d6fca4885..0ed630daab 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -373,6 +373,7 @@ steps: - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput - pytest -v -s tests/v1/distributed/test_dbo.py - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer_m2n.py - pytest -v -s tests/distributed/test_packed_tensor.py mirror: amd: diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index a389dc7165..77d2728345 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -35,6 +35,7 @@ engine drives on your behalf: | [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU | | [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Checkpoint-coordinate sparse weight patches | | [sharded_rdt](sharded_rdt.md) | NIXL / Ray Direct Transport (pull-based) | Very large models where each worker needs only its own slice (MoE with expert parallelism) | +| [nccl_m2n](m2n.md) | NCCL M2N reshard | Trainer and inference use different sharding layouts | ## Quickstart @@ -49,7 +50,8 @@ from vllm.config import WeightTransferConfig llm = LLM( model="my-model", - weight_transfer_config=WeightTransferConfig(backend="nccl"), # or "ipc", "sparse_nccl", "sharded_rdt" + # Other backends: "ipc", "sparse_nccl", "sharded_rdt", "nccl_m2n". + weight_transfer_config=WeightTransferConfig(backend="nccl"), ) ``` diff --git a/docs/training/weight_transfer/m2n.md b/docs/training/weight_transfer/m2n.md new file mode 100644 index 0000000000..61510978c8 --- /dev/null +++ b/docs/training/weight_transfer/m2n.md @@ -0,0 +1,101 @@ +# NCCL M2N Engine + +The `nccl_m2n` weight transfer engine uses [NCCL M2N](https://github.com/NVIDIA/nccl-extensions) to **reshard** weights between two disjoint meshes of GPUs — the trainer and the inference workers — in a single collective per parameter. Unlike the broadcast-based [NCCL engine](nccl.md), the trainer sends its local shard as-is: no full tensor is ever materialized on the training side, whatever its parallelism layout. + +## When to Use NCCL M2N + +- The trainer and the inference engine use **different parallelism layouts** (e.g. trainer FSDP/EP, inference TP) and gathering full tensors on the trainer is expensive +- Training and inference are on **separate GPUs**, possibly across nodes +- You want weight sync off the critical path of an async RL loop + +## Requirements + +This backend depends on the `nccl-extensions` Python package (and its `nccl4py` dependency), which vLLM does not install. It needs NCCL 2.30.5 or newer, and `libnccl_m2n.so` is linked against a specific NCCL build, so vLLM must load the same library: + +```bash +export VLLM_NCCL_SO_PATH=/path/to/the/same/libnccl.so +``` + +The backend imports the runtime lazily, so vLLM is unaffected unless you select it. + +## How It Works + +1. The trainer ranks and all inference workers join **one** NCCL communicator via `StatelessProcessGroup`. Trainer ranks occupy `[0, T)` and workers `[T, T + N)` — m2n meshes are contiguous rank intervals, which is exactly what vLLM's existing `rank_offset` convention produces. +2. At the init handshake the trainer ships the full transfer plan: both meshes — one per side, shared by every parameter — then each parameter's name, dtype, shape and placement. The workers resolve the matching **destination placements** and publish them back over the shared group; the meshes themselves are declared, not negotiated. +3. Each round, both sides issue one `reshard` per parameter in the same order. Every rank in the communicator participates in every reshard, so — as with the NCCL engine — the worker must be inside `update_weights` while the trainer sends. The trainer engine handles that concurrency itself. + +### Destination resolution + +A parameter is resharded **directly into the live vLLM parameter** when its checkpoint name maps 1:1 onto a model parameter whose shape is either identical to the checkpoint shape (replicated) or differs on exactly one dim by the tensor-parallel factor. That rank then only ever receives its own shard. + +Anything else falls back to receiving the full tensor and letting `load_weights` do the sharding, exactly as the NCCL engine does for every parameter. The fallback covers: + +- fused parameters (`qkv_proj`, `gate_up_proj`, MoE `w13`/`w2`), whose checkpoint names do not name a model parameter +- quantized models, where the weight loader does more than a sharded copy +- pipeline-parallel deployments + +Correctness never depends on a parameter resolving. The engine logs how many parameters took each path at init. + +## Configuration + +```python +from vllm import LLM +from vllm.config import WeightTransferConfig + +llm = LLM( + model="my-model", + weight_transfer_config=WeightTransferConfig(backend="nccl_m2n"), +) +``` + +```bash +vllm serve my-model --weight-transfer-config '{"backend": "nccl_m2n"}' +``` + +## Trainer Side + +The trainer uses the stateful trainer engine. m2n needs each parameter's source layout, which the base `WeightSource` does not carry, so it takes an `M2NWeightSource`. `DTensorModuleSource` covers the common case: it reads each parameter's device mesh and placements and yields **local shards** — it never calls `full_tensor()`. + +```python +from vllm.distributed.weight_transfer import ( + RayVLLMWeightSyncClient, + WeightTransferTrainerFactory, +) +from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource +from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo + +# Called on every trainer rank; rank 0 is the sender. +engine = WeightTransferTrainerFactory.trainer_init( + init_info=M2NTrainerInitInfo( + master_address=master_address, + master_port=master_port, + world_size=num_trainer_ranks + num_inference_workers, + num_trainer_ranks=num_trainer_ranks, + rank=my_rank, + ), + client=RayVLLMWeightSyncClient(llm_handle), + source=DTensorModuleSource(model, num_trainer_ranks), +) + +# Each round, on every trainer rank: +engine.send_weights() +``` + +`send_weights()` drives the whole round — `start_weight_update`, `update_weights`, `finish_weight_update`. Every rank in the communicator must be inside the same reshard, so `update_weights` has to be in flight while the trainer sends; the engine runs it on a helper thread so the trainer loop does not have to open-code that. + +Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass `M2NWeightSource` and supply layouts explicitly. + +See [`examples/rl/rlhf_m2n.py`](../../../examples/rl/rlhf_m2n.py) for a runnable FSDP → TP example. + +## Limitations + +- Tensor rank 1..3; 4-D parameters are rejected. +- dtypes: int8/uint8, fp8 e4m3/e5m2, fp16, bf16, int32/uint32, fp32, int64/uint64, fp64. No fp4. +- The trainer's device mesh must cover a contiguous rank interval starting at 0, and be 1-D or 2-D. +- `Partial` placements are not supported. +- Models with known model-level checkpoint preprocessing use the full-tensor + fallback for the affected module. +- A reshard cannot be issued from inside a CUDA graph capture (weight updates run outside capture, so this only matters for custom callers). +- The reshard plan is capped per *shard*, not per mesh: at most 16 source shards may feed one destination shard, and one source shard may feed at most 64 destination shards (`MAX_SOURCES` / `MAX_TARGETS` in the m2n build; raising either needs a rebuild). +- A directly-resharded parameter spreads over `tp_size` destination shards, so each one is fed by fewer source shards. A parameter on the fallback path is replicated, i.e. a single destination shard fed by every source shard, so a trainer that **sharded** its parameters is still limited to 16 ranks whenever any parameter falls back — which fused parameters always do. A replicated trainer is unaffected. +- Each communicator holds a cached staging pool (~256 MiB by default: 4 channels x 64 MiB). diff --git a/examples/rl/rlhf_m2n.py b/examples/rl/rlhf_m2n.py new file mode 100644 index 0000000000..e02229ed09 --- /dev/null +++ b/examples/rl/rlhf_m2n.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RLHF with FSDP2 training and vLLM tensor-parallel inference using **NCCL M2N** +sharding-aware weight transfer. + +Layout (4 GPUs, no colocation): + * GPUs 0-1: two FSDP2 training workers, one per GPU. + * GPUs 2-3: one vLLM ``LLM`` actor with ``tensor_parallel_size=2``. + +The trainer and the inference workers share one NCCL communicator of 4 ranks: +trainer ranks ``[0, 2)`` and inference ranks ``[2, 4)``. Every parameter moves +with a single ``reshard`` that redistributes it from the FSDP layout to the +inference layout — the trainer sends its local shards and never all-gathers a +full tensor, which broadcast-based weight sync would force it to do. + +Every FSDP rank builds an ``M2NTrainerWeightTransferEngine`` and calls +``send_weights()``; all ranks run every reshard, and only rank 0 drives the +inference side through its ``RayVLLMWeightSyncClient``. + +Requires the ``nccl-extensions`` package (NCCL 2.30.5+) and a +``VLLM_NCCL_SO_PATH`` pointing at the same ``libnccl.so`` that +``libnccl_m2n.so`` was linked against. + +This example was written for 4xH100. +""" + +from __future__ import annotations + +import os + +import pybase64 as base64 +import ray +import torch +import torch.distributed as dist +from huggingface_hub import snapshot_download +from torch.distributed.fsdp import fully_shard +from transformers import AutoModelForCausalLM + +from vllm import LLM, SamplingParams +from vllm.config import WeightTransferConfig +from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary +from vllm.distributed.weight_transfer import ( + RayVLLMWeightSyncClient, + WeightTransferTrainerFactory, +) +from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource +from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo +from vllm.utils.network_utils import get_ip, get_open_port + +MODEL_NAME = "Qwen/Qwen3-0.6B" + +FSDP_WORLD_SIZE = 2 +INFERENCE_TP_SIZE = 2 + + +class MyLLM(LLM): + """LLM subclass that keeps Ray from pinning it to a single device.""" + + def __init__(self, *args, **kwargs): + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + super().__init__(*args, **kwargs) + + def ready(self): + return True + + +@ray.remote(num_cpus=1, num_gpus=1) +class FSDPTrainWorker: + """One FSDP2 worker per GPU. Rank 0 is the weight-transfer sender.""" + + def __init__( + self, + model_name: str, + rank: int, + fsdp_master_addr: str, + fsdp_master_port: int, + ): + self.rank = rank + + os.environ["MASTER_ADDR"] = fsdp_master_addr + os.environ["MASTER_PORT"] = str(fsdp_master_port) + + dist.init_process_group(backend="nccl", rank=rank, world_size=FSDP_WORLD_SIZE) + torch.accelerator.set_device_index(0) + + model = AutoModelForCausalLM.from_pretrained( + model_name, torch_dtype=torch.bfloat16 + ).cuda() + for layer in model.model.layers: + fully_shard(layer) + fully_shard(model) + self.model = model + + def ready(self): + return True + + def setup_engine(self, llm_handle, nccl_unique_id_b64, world_size, num_workers): + """Build the trainer engine on every FSDP rank. + + `DTensorModuleSource` reads each parameter's FSDP device mesh and + placements, so the engine knows the source layout without gathering + anything. Rank 0 additionally drives the inference-side handshake. + """ + self.engine = WeightTransferTrainerFactory.trainer_init( + init_info=M2NTrainerInitInfo( + nccl_unique_id_b64=nccl_unique_id_b64, + world_size=world_size, + num_trainer_ranks=FSDP_WORLD_SIZE, + # One DP group of TP=2 workers; declared so both sides + # describe the destination the same way. + dst_mesh_dims=(num_workers // INFERENCE_TP_SIZE, INFERENCE_TP_SIZE), + rank=self.rank, # FSDP rank; sender is 0 + ), + client=RayVLLMWeightSyncClient(llm_handle), + source=DTensorModuleSource(self.model, FSDP_WORLD_SIZE), + ) + + def send_weights(self): + """Called on all ranks concurrently; every rank runs every reshard.""" + self.engine.send_weights() + + +def main(): + ray.init() + + local_model_path = snapshot_download(MODEL_NAME) + print(f"[init] Model downloaded to {local_model_path}") + + fsdp_master_addr = get_ip() + fsdp_master_port = get_open_port() + + fsdp_workers = [ + FSDPTrainWorker.remote( + local_model_path, rank, fsdp_master_addr, fsdp_master_port + ) + for rank in range(FSDP_WORLD_SIZE) + ] + ray.get([w.ready.remote() for w in fsdp_workers]) + print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.") + + llm = ray.remote(num_cpus=0, num_gpus=0)(MyLLM).remote( + model=local_model_path, + enforce_eager=True, + tensor_parallel_size=INFERENCE_TP_SIZE, + distributed_executor_backend="ray", + weight_transfer_config=WeightTransferConfig(backend="nccl_m2n"), + load_format="dummy", + ) + ray.get(llm.ready.remote()) + num_workers = ray.get(llm.get_world_size.remote()) + print(f"[init] vLLM ready with {num_workers} inference workers.") + + prompts = [ + "Hello, my name is", + "The capital of France is", + ] + sampling_params = SamplingParams(temperature=0) + + outputs = ray.get(llm.generate.remote(prompts, sampling_params)) + print("-" * 60) + print("BEFORE weight sync (dummy weights):") + for output in outputs: + print(f"Prompt: {output.prompt!r}") + print(f"Generated: {output.outputs[0].text!r}") + print("-" * 60) + + # Trainer ranks [0, FSDP_WORLD_SIZE) and inference ranks after them share + # one communicator, so the two meshes are contiguous rank intervals. + nccl = NCCLLibrary() + nccl_unique_id_b64 = base64.b64encode( + bytes(nccl.ncclGetUniqueId().internal) + ).decode() + world_size = FSDP_WORLD_SIZE + num_workers + + print("[transfer] Initializing nccl_m2n weight transfer (all FSDP ranks)...") + ray.get( + [ + w.setup_engine.remote(llm, nccl_unique_id_b64, world_size, num_workers) + for w in fsdp_workers + ] + ) + + print("[sync] Resharding FSDP -> vLLM...") + ray.get([w.send_weights.remote() for w in fsdp_workers]) + print("[sync] Weight transfer complete.") + + outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params)) + print("-" * 60) + print("AFTER weight sync (real weights):") + for output in outputs_updated: + print(f"Prompt: {output.prompt!r}") + print(f"Generated: {output.outputs[0].text!r}") + print("-" * 60) + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_weight_transfer_m2n.py b/tests/distributed/test_weight_transfer_m2n.py new file mode 100644 index 0000000000..0591adff4d --- /dev/null +++ b/tests/distributed/test_weight_transfer_m2n.py @@ -0,0 +1,868 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the NCCL M2N weight transfer backend. + +These cover how a transfer is *described* — layout encoding and the wire-type +validation that keeps a bad plan from reaching a collective, since a mismatch +there is a hang rather than an error. The transfer itself needs the m2n runtime +and multiple GPUs, so it is exercised separately. +""" + +import logging +from functools import wraps +from types import MethodType +from unittest.mock import Mock + +import pybase64 as base64 +import pytest +import ray +import torch + +from vllm.distributed.weight_transfer import ( + WeightTransferEngineFactory, + WeightTransferTrainerFactory, +) +from vllm.distributed.weight_transfer.m2n_common import ( + REPLICATE, + REPLICATED, + M2NMesh, + M2NParamMeta, + check_placements, + check_transferable, + publish_destination_placements, + resolve_layout, + validate_layout, +) +from vllm.distributed.weight_transfer.m2n_engine import ( + M2NWeightTransferEngine, + M2NWeightTransferInitInfo, + M2NWeightTransferUpdateInfo, +) +from vllm.distributed.weight_transfer.m2n_layout import ( + M2NDestination, + resolve_parameter_destinations, +) +from vllm.distributed.weight_transfer.m2n_source import ( + mesh_from_tensor, + placements_from_tensor, +) +from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.platforms import current_platform + +VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode() + + +class TestLayout: + def test_replicated_splits_nothing(self): + """A replicated tensor imposes no divisibility constraint. The shape + dims are deliberately coprime with the mesh size: a layout that actually + split the tensor would reject them.""" + mesh = M2NMesh((2, 2), start_rank=1) + indivisible_shape = (7, 13) # neither dim divisible by any mesh axis + + resolved_mesh, placements = resolve_layout(mesh, REPLICATED) + validate_layout(resolved_mesh, placements, indivisible_shape, "destination") + + def test_replicated_keeps_the_same_ranks(self): + """Replication re-factors the mesh to get a size-1 axis for its no-op + shard. That is only sound if it still covers exactly the same GPUs.""" + mesh = M2NMesh((2, 3), start_rank=4) + resolved_mesh, _ = resolve_layout(mesh, REPLICATED) + assert resolved_mesh.size == mesh.size + assert resolved_mesh.start_rank == mesh.start_rank + + def test_sharded_keeps_its_own_factorization(self): + """Rank order decides who owns which shard, so a sharded tensor must + not be re-factored the way a replicated one is.""" + mesh = M2NMesh((2, 3), start_rank=4) + resolved_mesh, placements = resolve_layout(mesh, (REPLICATE, 0)) + assert resolved_mesh == mesh + assert placements == (REPLICATE, 0) + + def test_two_shard_axes_rejected(self): + """One axis has to replicate; a 2-D mesh that shards both is not + something a single reshard can express.""" + with pytest.raises(ValueError, match="shards both"): + check_placements((0, 1)) + + def test_placement_below_replicate_rejected(self): + with pytest.raises(ValueError, match=r"parameter 'w'.*-2"): + check_placements((-2, REPLICATE), "parameter 'w' source placements") + + def test_shard_dim_must_exist(self): + with pytest.raises(ValueError, match="rank 2"): + validate_layout(M2NMesh((1, 2), 0), (REPLICATE, 2), (8, 16), "source") + + def test_shard_must_divide_evenly(self): + with pytest.raises(ValueError, match="does not divide evenly"): + validate_layout(M2NMesh((1, 3), 0), (REPLICATE, 0), (8, 16), "source") + + +class TestSourceLayout: + def test_plain_tensor_is_replicated_across_trainer_ranks(self): + """A tensor with no DTensor metadata is the same on every trainer rank, + so the source describes it as replicated over all of them.""" + assert placements_from_tensor(torch.zeros(4)) is REPLICATED + assert mesh_from_tensor(torch.zeros(4), 4) == M2NMesh((4, 1), 0) + + +class TestTransferable: + def test_unsupported_dtype_names_the_parameter(self): + with pytest.raises(ValueError, match="'w'"): + check_transferable("w", torch.complex64, (4,)) + + def test_rank_four_rejected(self): + with pytest.raises(ValueError, match="rank 4"): + check_transferable("w", torch.bfloat16, (2, 2, 2, 2)) + + +class TestWireTypes: + def _init_info(self, **overrides): + fields = dict( + master_address="127.0.0.1", + master_port=1234, + rank_offset=1, + world_size=3, + src_mesh_dims=[1, 1], + dst_mesh_dims=[2, 1], + names=["w"], + dtype_names=["bfloat16"], + shapes=[[16, 16]], + src_placements=[None], + ) + fields.update(overrides) + return M2NWeightTransferInitInfo(**fields) + + def test_accepts_a_consistent_plan(self): + assert self._init_info().names == ["w"] + + def test_accepts_pre_shared_nccl_unique_id(self): + info = self._init_info( + master_address=None, + master_port=None, + nccl_unique_id_b64=VALID_UID_B64, + ) + assert info.nccl_unique_id_bytes == b"\x00" * 128 + assert VALID_UID_B64 not in repr(info) + + def test_destination_plan_uses_uid_communicator_without_group(self): + comm = Mock(rank=0, device=torch.device("cpu"), group=None) + + def receive_plan(tensor, src): + assert src == 1 + tensor.copy_(torch.tensor([[-1, 0], [-2, -2]], dtype=torch.int8)) + + comm.broadcast.side_effect = receive_plan + + placements = publish_destination_placements(comm, 1, None, 2) + + assert placements == [(REPLICATE, 0), REPLICATED] + + def test_rejects_both_rendezvous_modes(self): + with pytest.raises(ValueError, match="not both"): + self._init_info(nccl_unique_id_b64=VALID_UID_B64) + + def test_trainer_accepts_pre_shared_nccl_unique_id(self): + info = M2NTrainerInitInfo( + nccl_unique_id_b64=VALID_UID_B64, + world_size=3, + num_trainer_ranks=1, + rank=0, + ) + assert info.nccl_unique_id_bytes == b"\x00" * 128 + assert VALID_UID_B64 not in repr(info) + + def test_ragged_plan_rejected(self): + with pytest.raises(ValueError, match="`shapes`"): + self._init_info(shapes=[]) + + def test_destination_mesh_must_cover_the_workers(self): + """The trainer declares the inference mesh, so one that does not cover + the workers is a config error — and it has to fail the init RPC, since + a mismatched mesh would otherwise surface as a hung collective.""" + with pytest.raises(ValueError, match="dst_mesh_dims"): + self._init_info(dst_mesh_dims=[3, 1]) # 3 != the 2 workers + + def test_world_must_hold_a_trainer_and_a_worker(self): + with pytest.raises(ValueError, match="rank_offset"): + self._init_info(rank_offset=3, world_size=3) + + @pytest.mark.parametrize("dtype_name", ["not_a_dtype", "Tensor"]) + def test_invalid_dtype_name_names_parameter(self, monkeypatch, dtype_name): + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.import_m2n", + lambda: object(), + ) + engine = object.__new__(M2NWeightTransferEngine) + + with pytest.raises(ValueError, match=r"parameter 'w'.*dtype"): + engine.init_transfer_engine(self._init_info(dtype_names=[dtype_name])) + + def _init_32_to_4_plan(self, monkeypatch, destination): + m2n = Mock() + m2n.Handle.create.return_value = object() + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.import_m2n", lambda: m2n + ) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine." + "resolve_parameter_destinations", + lambda *args, **kwargs: [destination], + ) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.worker_init_process_group", + lambda *args, **kwargs: object(), + ) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine." + "publish_destination_placements", + lambda *args: args[2], + ) + + engine = object.__new__(M2NWeightTransferEngine) + engine.parallel_config = Mock(pipeline_parallel_size=1) + engine.model_config = Mock(quantization=None) + engine.model = torch.nn.Module() + engine.init_transfer_engine( + self._init_info( + rank_offset=32, + world_size=36, + src_mesh_dims=[1, 32], + dst_mesh_dims=[1, 4], + shapes=[[32, 8]], + src_placements=[[REPLICATE, 0]], + ) + ) + + def test_limits_use_resolved_sharded_destination(self, monkeypatch): + destination = M2NDestination("w", (REPLICATE, 0), torch.empty(8, 8)) + + self._init_32_to_4_plan(monkeypatch, destination) + + def test_limits_still_reject_replicated_destination(self, monkeypatch): + destination = M2NDestination("w", REPLICATED, None) + + with pytest.raises(ValueError, match=r"32 source shards.*MAX_SOURCES=16"): + self._init_32_to_4_plan(monkeypatch, destination) + + def test_update_preflights_all_names_before_reshard(self): + engine = object.__new__(M2NWeightTransferEngine) + engine._handle = object() + engine.model_update_group = object() + engine._index = {"valid": 0} + engine._metas = [M2NParamMeta("valid", torch.float32, (4,), REPLICATED)] + engine._reshard = Mock() + + with pytest.raises(ValueError, match=r"parameter 'unknown'"): + engine.receive_weights( + M2NWeightTransferUpdateInfo(names=["valid", "unknown"]) + ) + + engine._reshard.assert_not_called() + + def test_fallback_tensors_share_one_load_weights_invocation(self, monkeypatch): + class CoupledModel: + def __init__(self): + self.calls = 0 + self.loaded = [] + + def load_weights(self, weights): + self.calls += 1 + received = list(weights) + if [name for name, _ in received] == ["pair.weight", "pair.scale"]: + self.loaded = [(name, tensor.clone()) for name, tensor in received] + + model = CoupledModel() + direct = torch.zeros(1) + engine = object.__new__(M2NWeightTransferEngine) + engine._handle = object() + engine.model_update_group = object() + engine.device = torch.device("cpu") + engine.model = model + engine._index = {"pair.weight": 0, "direct": 1, "pair.scale": 2} + engine._metas = [ + M2NParamMeta(name, torch.float32, (1,), REPLICATED) + for name in engine._index + ] + engine._parameter_destinations = [ + M2NDestination("pair.weight", REPLICATED, None), + M2NDestination("direct", REPLICATED, direct), + M2NDestination("pair.scale", REPLICATED, None), + ] + reshard_order = [] + + def reshard(comm, stream, meta, placements, buffer): + reshard_order.append(meta.name) + buffer.fill_(len(reshard_order)) + + engine._reshard = reshard + stream = Mock() + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.comm_ptr", lambda group: 0 + ) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream) + + engine.receive_weights( + M2NWeightTransferUpdateInfo(names=["pair.weight", "direct", "pair.scale"]) + ) + + assert model.calls == 1 + assert [name for name, _ in model.loaded] == ["pair.weight", "pair.scale"] + assert [tensor.item() for _, tensor in model.loaded] == [1, 3] + assert direct.item() == 2 + assert reshard_order == ["pair.weight", "direct", "pair.scale"] + assert stream.synchronize.call_count == 2 + + def test_trainer_rank_must_be_a_trainer_rank(self): + """A trainer rank must fall within the trainer portion of the group.""" + with pytest.raises(ValueError, match="num_trainer_ranks"): + M2NTrainerInitInfo( + master_address="127.0.0.1", + master_port=1234, + world_size=4, + num_trainer_ranks=2, + rank=2, + ) + + def test_trainer_destination_mesh_must_cover_the_workers(self): + """The destination mesh must include every inference worker.""" + with pytest.raises(ValueError, match="dst_mesh_dims"): + M2NTrainerInitInfo( + master_address="127.0.0.1", + master_port=1234, + world_size=6, + num_trainer_ranks=2, + dst_mesh_dims=(3, 1), # 3 != the 4 inference workers + rank=0, + ) + + def test_destination_mesh_defaults_to_flat(self): + """A replicated destination does not care how the mesh is factored, so + callers that do not shard it need not supply one.""" + info = M2NTrainerInitInfo( + master_address="127.0.0.1", + master_port=1234, + world_size=6, + num_trainer_ranks=2, + rank=0, + ) + assert info.destination_mesh_dims == (4, 1) + + def test_sender_is_trainer_rank_zero(self): + """Trainer rank 0 drives the inference control plane.""" + info = M2NTrainerInitInfo( + master_address="127.0.0.1", master_port=1234, world_size=4, rank=0 + ) + assert info.is_sender + + +class TestRegistration: + def test_both_registries_expose_the_backend(self): + """Both worker and trainer factories register nccl_m2n.""" + assert "nccl_m2n" in WeightTransferEngineFactory._registry + assert "nccl_m2n" in WeightTransferTrainerFactory._registry + + def test_init_info_dispatches_to_the_backend(self): + """Trainer init info selects the nccl_m2n factory entry.""" + assert M2NTrainerInitInfo.backend == "nccl_m2n" + + +# --------------------------------------------------------------------------- +# End-to-end transfer +# +# Transport-only, in the style of the NCCL/sparse tests in +# test_weight_transfer.py: two Ray tasks with one GPU each drive the two engines +# directly, so no HTTP server or LLM instance is involved. The worker is not an +# RPC endpoint here, so the trainer engine gets a no-op control-plane client -- +# the NCCL rendezvous and the reshard itself are the real thing. +# --------------------------------------------------------------------------- + +SHAPE = [64, 32] +DTYPE = "float32" + + +def _init_ray() -> None: + if ray.is_initialized(): + return + ray.init( + ignore_reinit_error=True, + runtime_env={ + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1", + } + }, + ) + + +def _assigned_device() -> "torch.device": + gpu_ids = ray.get_gpu_ids() + device = torch.device(f"cuda:{int(gpu_ids[0])}" if gpu_ids else "cuda:0") + current_platform.set_device(device) + return device + + +@ray.remote(num_gpus=1) +def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool: + """Send one parameter through the real trainer engine.""" + device = _assigned_device() + + from vllm.distributed.weight_transfer import WeightTransferTrainerFactory + from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource + from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo + + class NoopClient: + def init_weight_transfer_engine(self, init_info): + pass + + def start_weight_update(self): + pass + + def update_weights(self, update_info): + pass + + def finish_weight_update(self, weight_version=None): + pass + + class Tiny(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.arange( + SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device + ).reshape(SHAPE) + ) + + engine = WeightTransferTrainerFactory.trainer_init( + init_info=M2NTrainerInitInfo( + nccl_unique_id_b64=nccl_unique_id_b64, + world_size=world_size, + num_trainer_ranks=1, + dst_mesh_dims=(1, world_size - 1), + rank=0, + ), + client=NoopClient(), + source=DTensorModuleSource(Tiny(), num_trainer_ranks=1), + ) + engine.send_weights() + torch.accelerator.synchronize() + engine.shutdown() + return True + + +@ray.remote(num_gpus=1) +def _m2n_worker_receive( + nccl_unique_id_b64: str, + world_size: int, + worker_rank: int = 0, +) -> dict: + """Receive that parameter through the real worker engine.""" + import contextlib + from unittest.mock import MagicMock + + device = _assigned_device() + + from vllm.config.parallel import ParallelConfig + from vllm.config.weight_transfer import WeightTransferConfig + from vllm.distributed.weight_transfer.m2n_engine import ( + M2NWeightTransferEngine, + M2NWeightTransferInitInfo, + M2NWeightTransferUpdateInfo, + ) + + class Recorder(torch.nn.Module): + def __init__(self): + super().__init__() + self.received = [] + if world_size > 2: + self.weight = torch.nn.Parameter( + torch.zeros( + SHAPE[0] // (world_size - 1), + SHAPE[1], + dtype=torch.float32, + device=device, + ) + ) + self.weight.output_dim = 0 + self.weight.weight_loader = MethodType( + ColumnParallelLinear.weight_loader, + Mock(tp_size=world_size - 1), + ) + + def load_weights(self, weights): + for name, tensor in weights: + self.received.append((name, tensor.clone())) + + parallel_config = MagicMock(spec=ParallelConfig) + parallel_config.rank = worker_rank + parallel_config.world_size = world_size - 1 + parallel_config.data_parallel_rank = 0 + parallel_config.data_parallel_index = 0 + parallel_config.tensor_parallel_size = world_size - 1 + parallel_config.pipeline_parallel_size = 1 + vllm_config = MagicMock() + vllm_config.parallel_config = parallel_config + vllm_config.model_config = MagicMock() + vllm_config.model_config.quantization = None + + recorder = Recorder() + engine = M2NWeightTransferEngine( + WeightTransferConfig(backend="nccl_m2n"), vllm_config, device, recorder + ) + # Transport-only: receive_weights enters set_current_vllm_config, and + # vllm_config here is a mock. + import vllm.config as _vllm_config_mod + + _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext() + + engine.init_transfer_engine( + M2NWeightTransferInitInfo( + nccl_unique_id_b64=nccl_unique_id_b64, + rank_offset=1, # trainer occupies rank 0 + world_size=world_size, + src_mesh_dims=[1, 1], + dst_mesh_dims=[1, world_size - 1], + names=["weight"], + dtype_names=[DTYPE], + shapes=[SHAPE], + src_placements=[None], # replicated on the single trainer rank + ) + ) + engine.receive_weights(M2NWeightTransferUpdateInfo(names=["weight"])) + torch.accelerator.synchronize() + + full = torch.arange( + SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device + ).reshape(SHAPE) + direct = world_size > 2 + name: str | None + got: torch.Tensor | None + if direct: + name, got = "weight", recorder.weight + expected = full.chunk(world_size - 1, dim=0)[worker_rank] + else: + name, got = recorder.received[0] if recorder.received else (None, None) + expected = full + result = { + "count": len(recorder.received), + "name": name, + "shape": list(got.shape) if got is not None else None, + "exact": bool(torch.equal(got, expected)) if got is not None else False, + "direct": direct, + } + engine.shutdown() + return result + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 2, + reason="Need at least 2 GPUs: one trainer rank and one inference worker.", +) +def test_m2n_weight_transfer_between_processes(): + """A parameter survives a real reshard from a trainer process to a worker. + + This is the only test here that moves bytes: it builds both engines, joins + one NCCL communicator across two processes, and reshards. Everything else + in this file checks how a transfer is *described*. + """ + pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") + _init_ray() + + from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary + + nccl = NCCLLibrary() + nccl_unique_id_b64 = base64.b64encode( + bytes(nccl.ncclGetUniqueId().internal) + ).decode() + world_size = 2 # 1 trainer + 1 inference worker + + worker = _m2n_worker_receive.remote(nccl_unique_id_b64, world_size) + trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size) + trainer_ok, result = ray.get([trainer, worker]) + + assert trainer_ok, "trainer engine did not complete" + assert result["count"] == 1, f"expected one parameter, got {result['count']}" + assert result["name"] == "weight" + assert result["shape"] == SHAPE + assert result["exact"], "resharded tensor does not match what the trainer sent" + + +@pytest.mark.skipif( + torch.accelerator.device_count() < 3, + reason="Need at least 3 GPUs: one trainer rank and two inference workers.", +) +def test_m2n_weight_transfer_to_tp2_shards(): + """Two inference workers receive their own TP shard in live model storage.""" + pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") + _init_ray() + + from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary + + nccl = NCCLLibrary() + nccl_unique_id_b64 = base64.b64encode( + bytes(nccl.ncclGetUniqueId().internal) + ).decode() + world_size = 3 + + workers = [ + _m2n_worker_receive.remote(nccl_unique_id_b64, world_size, rank) + for rank in range(2) + ] + trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size) + trainer_ok, *results = ray.get([trainer, *workers]) + + assert trainer_ok, "trainer engine did not complete" + for rank, result in enumerate(results): + assert result["count"] == 0, f"worker {rank} unexpectedly called load_weights" + assert result["name"] == "weight" + assert result["shape"] == [SHAPE[0] // 2, SHAPE[1]] + assert result["direct"] + assert result["exact"], f"worker {rank} received the wrong shard" + + +class _Model(torch.nn.Module): + """Stand-in for a TP=2 shard of a tiny model. + + `column` is sharded on dim 0 and `row` on dim 1 (the two shapes a TP linear + layer produces), `norm` is replicated, and `fused` has no checkpoint-name + counterpart — the shape a fused parameter presents to the resolver. + """ + + def __init__(self) -> None: + super().__init__() + self.column = torch.nn.Parameter(torch.zeros(8, 16)) + self.column.output_dim = 0 + self.column.weight_loader = MethodType( + ColumnParallelLinear.weight_loader, Mock(tp_size=2) + ) + self.row = torch.nn.Parameter(torch.zeros(16, 8)) + self.row.input_dim = 1 + self.row.weight_loader = MethodType( + RowParallelLinear.weight_loader, Mock(tp_size=2) + ) + self.norm = torch.nn.Parameter(torch.zeros(16)) + self.norm.weight_loader = default_weight_loader + self.fused = torch.nn.Parameter(torch.zeros(24, 16)) + + +class _MixedModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.proj = torch.nn.Module() + self.proj.weight = torch.nn.Parameter(torch.zeros(16, 8)) + self.proj.weight.input_dim = 1 + self.proj.weight.weight_loader = MethodType( + RowParallelLinear.weight_loader, Mock(tp_size=2) + ) + self.proj.bias = torch.nn.Parameter(torch.zeros(16)) + + +def _resolve(names, dtypes, shapes, **kwargs): + defaults = dict(num_workers=2, shard_axis_size=2, allow_direct=True) + defaults.update(kwargs) + return resolve_parameter_destinations(_Model(), names, dtypes, shapes, **defaults) + + +class TestDestinationResolution: + @pytest.mark.parametrize( + "name,expected_dim", + [("column", 0), ("row", 1)], + ) + def test_sharded_parameter_resolves_to_its_tp_dim(self, name, expected_dim): + [destination] = _resolve([name], [torch.float32], [(16, 16)]) + assert destination.direct + assert destination.placements == (REPLICATE, expected_dim) + + def test_replicated_parameter_needs_no_placement(self): + """A parameter every rank holds in full is REPLICATED, so it works + whatever way the inference mesh happens to be factored.""" + [destination] = _resolve(["norm"], [torch.float32], [(16,)]) + assert destination.direct + assert destination.placements is REPLICATED + + def test_direct_destination_is_the_live_parameter(self): + model = _Model() + [destination] = resolve_parameter_destinations( + model, + ["column"], + [torch.float32], + [(16, 16)], + num_workers=2, + shard_axis_size=2, + allow_direct=True, + ) + assert destination.tensor.data_ptr() == model.column.data_ptr() + + def test_known_transforming_model_submodule_falls_back(self): + class GPT2Model(_Model): + pass + + GPT2Model.__module__ = "vllm.model_executor.models.gpt2" + model = torch.nn.Module() + model.transforming = GPT2Model() + [destination] = resolve_parameter_destinations( + model, + ["transforming.row"], + [torch.float32], + [(16, 16)], + num_workers=2, + shard_axis_size=2, + allow_direct=True, + ) + + assert not destination.direct + assert destination.placements is REPLICATED + + def test_unknown_name_falls_back(self): + """Fused parameters reach the worker under checkpoint names that do not + exist in the model; those must take the full-tensor path.""" + [destination] = _resolve(["mlp.gate_proj.weight"], [torch.float32], [(16, 16)]) + assert not destination.direct + assert destination.placements is REPLICATED + + def test_fallback_parameter_demotes_direct_sibling_in_same_module(self): + """Layerwise finalization must not overwrite a directly loaded weight + when a sibling bias takes the fallback path.""" + destinations = resolve_parameter_destinations( + _MixedModel(), + ["proj.weight", "proj.bias"], + [torch.float32, torch.float32], + [(16, 16), (16,)], + num_workers=2, + shard_axis_size=2, + allow_direct=True, + ) + + assert not any(destination.direct for destination in destinations) + + def test_unknown_loader_falls_back_despite_matching_name_and_shape(self): + model = _Model() + model.norm.weight_loader = lambda param, weight: param.data.copy_(weight) + [destination] = resolve_parameter_destinations( + model, + ["norm"], + [torch.float32], + [(16,)], + num_workers=2, + shard_axis_size=2, + allow_direct=True, + ) + assert not destination.direct + + def test_missing_loader_is_not_assumed_to_be_a_plain_copy(self): + model = _Model() + del model.norm.weight_loader + [destination] = resolve_parameter_destinations( + model, + ["norm"], + [torch.float32], + [(16,)], + num_workers=2, + shard_axis_size=2, + allow_direct=True, + ) + assert not destination.direct + + def test_wrapped_known_loader_falls_back(self): + model = _Model() + + @wraps(default_weight_loader) + def wrapped_loader(param, weight): + default_weight_loader(param, weight) + + model.norm.weight_loader = wrapped_loader + [destination] = resolve_parameter_destinations( + model, + ["norm"], + [torch.float32], + [(16,)], + num_workers=2, + shard_axis_size=2, + allow_direct=True, + ) + assert not destination.direct + + def test_replicated_loader_with_transform_metadata_falls_back(self): + model = _Model() + model.norm.is_transposed = True + [destination] = resolve_parameter_destinations( + model, + ["norm"], + [torch.float32], + [(16,)], + num_workers=2, + shard_axis_size=2, + allow_direct=True, + ) + assert not destination.direct + + def test_packed_dim_zero_falls_back(self): + model = _Model() + model.column.packed_dim = 0 + [destination] = resolve_parameter_destinations( + model, + ["column"], + [torch.float32], + [(16, 16)], + num_workers=2, + shard_axis_size=2, + allow_direct=True, + ) + assert not destination.direct + + def test_loader_declared_dim_wins_over_shape_guessing(self): + model = _Model() + model.column.output_dim = 1 + [destination] = resolve_parameter_destinations( + model, + ["column"], + [torch.float32], + [(16, 16)], + num_workers=2, + shard_axis_size=2, + allow_direct=True, + ) + assert not destination.direct + + def test_shape_the_tp_factor_cannot_explain_falls_back(self): + [destination] = _resolve(["fused"], [torch.float32], [(16, 16)]) + assert not destination.direct + + def test_dtype_mismatch_falls_back(self): + """A parameter stored in a different dtype than the wire dtype means a + quantized or otherwise transformed layout — never write into it.""" + [destination] = _resolve(["column"], [torch.bfloat16], [(16, 16)]) + assert not destination.direct + + def test_allow_direct_false_forces_every_parameter_to_fall_back(self): + destinations = _resolve( + ["column", "norm"], + [torch.float32] * 2, + [(16, 16), (16,)], + allow_direct=False, + ) + assert not any(d.direct for d in destinations) + + def test_plan_reports_byte_coverage(self, caplog_vllm): + """A small direct parameter must not hide a large fallback tensor.""" + with caplog_vllm.at_level( + logging.INFO, + logger="vllm.distributed.weight_transfer.m2n_layout", + ): + _resolve( + ["norm", "missing_fused_weight"], + [torch.float32, torch.bfloat16], + [(16,), (1024, 1024)], + ) + + assert ( + "1/2 parameters resharded directly into the model, 1 via " + "full-tensor fallback; direct byte coverage: 64/2097216 bytes " + "(0.0%)" in caplog_vllm.text + ) diff --git a/vllm/config/weight_transfer.py b/vllm/config/weight_transfer.py index 3e849753c2..99676526af 100644 --- a/vllm/config/weight_transfer.py +++ b/vllm/config/weight_transfer.py @@ -9,7 +9,9 @@ from vllm.config.utils import config class WeightTransferConfig: """Configuration for weight transfer during RL training.""" - backend: Literal["nccl", "ipc", "sparse_nccl", "sharded_rdt"] | str = "nccl" + backend: Literal["nccl", "ipc", "sparse_nccl", "sharded_rdt", "nccl_m2n"] | str = ( + "nccl" + ) """The backend to use for weight transfer. Validated against the `WeightTransferEngineFactory` registry at engine creation time. """ diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py index 48f9e4f212..82228bef34 100644 --- a/vllm/distributed/weight_transfer/factory.py +++ b/vllm/distributed/weight_transfer/factory.py @@ -243,6 +243,12 @@ WeightTransferEngineFactory.register_engine( "ShardedRDTWeightTransferEngine", ) +WeightTransferEngineFactory.register_engine( + "nccl_m2n", + "vllm.distributed.weight_transfer.m2n_engine", + "M2NWeightTransferEngine", +) + # Trainer-side engines, parallel to the worker registry above. WeightTransferTrainerFactory.register_engine( @@ -268,3 +274,9 @@ WeightTransferTrainerFactory.register_engine( "vllm.distributed.weight_transfer.sharded_rdt_trainer", "ShardedRDTTrainerWeightTransferEngine", ) + +WeightTransferTrainerFactory.register_engine( + "nccl_m2n", + "vllm.distributed.weight_transfer.m2n_trainer", + "M2NTrainerWeightTransferEngine", +) diff --git a/vllm/distributed/weight_transfer/m2n_common.py b/vllm/distributed/weight_transfer/m2n_common.py new file mode 100644 index 0000000000..a2c6b73844 --- /dev/null +++ b/vllm/distributed/weight_transfer/m2n_common.py @@ -0,0 +1,338 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared helpers for the NCCL M2N (`nccl_m2n`) weight transfer backend. + +M2N reshards a tensor between two disjoint meshes of ranks that live in one +communicator: the trainer occupies ranks `[0, T)` and the inference workers +`[T, T + N)`. This module holds everything both sides need — the optional +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 +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, cast + +import torch + +from vllm.distributed.weight_transfer.base import ParamMeta + +if TYPE_CHECKING: + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + +# 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 +# layouts survive the JSON init handshake without a custom encoder. +REPLICATE = -1 + +# `NCCL_RESHARD_MAX_TENSOR_DIMS`; 4-D and higher are rejected by the library. +MAX_TENSOR_DIMS = 3 + +MESH_NDIMS = 2 + +# Destination mesh convention. Keeping these roles named avoids scattering +# positional axis assumptions through the worker-side planner. +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. +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. +SUPPORTED_DTYPES: frozenset[torch.dtype] = frozenset( + dtype + for dtype in ( + getattr(torch, name, None) + for name in ( + "int8", + "uint8", + "float8_e4m3fn", + "float8_e5m2", + "float16", + "bfloat16", + "int32", + "uint32", + "float32", + "int64", + "uint64", + "float64", + ) + ) + if dtype is not None +) + +_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." +) + + +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 + + +@dataclass(frozen=True) +class M2NMesh: + """One side's rank topology — pure topology, no tensor placement. + + Mirrors `ncclMesh_t`: a 2-axis mesh owning the contiguous rank interval + `[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. + """ + + 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): + raise ValueError(f"mesh dims must be {MESH_NDIMS} positive ints") + if self.start_rank < 0: + raise ValueError( + f"mesh start_rank must be non-negative, got {self.start_rank}" + ) + + @property + def size(self) -> int: + return self.dims[0] * self.dims[1] + + +# A tensor's placement over its side's mesh: one code per mesh axis +# (`REPLICATE`, or the tensor dim that axis shards), or `REPLICATED` for a +# tensor every rank holds in full. +Placements = tuple[int, int] +REPLICATED: Placements | None = None + + +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 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 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.""" + 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 to_mesh(m2n: Any, mesh: M2NMesh) -> Any: + return m2n.Mesh(mesh.dims, start_rank=mesh.start_rank) + + +def to_placements(m2n: Any, placements: Placements) -> list[Any]: + return [ + m2n.Replicate() if code == REPLICATE else m2n.Shard(code) for code in placements + ] + + +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 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) diff --git a/vllm/distributed/weight_transfer/m2n_engine.py b/vllm/distributed/weight_transfer/m2n_engine.py new file mode 100644 index 0000000000..c61bac4f2d --- /dev/null +++ b/vllm/distributed/weight_transfer/m2n_engine.py @@ -0,0 +1,453 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-side weight transfer engine built on NCCL M2N (`nccl_m2n`). + +The trainer and the inference workers share one NCCL communicator: trainer ranks +occupy `[0, T)`, workers `[T, T + N)`. Each parameter is moved with a single +`nccl.m2n.reshard`, which redistributes it from the trainer's layout (FSDP / EP +/ arbitrary DTensor sharding) to the inference layout — so the trainer sends its +local shards and never all-gathers a full tensor, which is what the broadcast +NCCL backend forces it to do. + +When an incoming checkpoint parameter maps directly to a live vLLM parameter, +M2N reshards into that worker's local model storage. Parameters whose names or +layouts cannot be resolved receive a full tensor and fall back to +`load_weights`, preserving the existing backend's loading behavior. + +Both meshes participate in every reshard, in the same order, so this backend has +the same concurrency shape as the broadcast NCCL backend: the worker must be +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 dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, cast + +import torch + +from vllm.config.weight_transfer import WeightTransferConfig +from vllm.distributed.weight_transfer.base import ( + WeightTransferEngine, + WeightTransferInitInfo, + WeightTransferUpdateInfo, +) +from vllm.distributed.weight_transfer.m2n_common import ( + DESTINATION_SHARD_AXIS, + MESH_NDIMS, + REPLICATED, + M2NMesh, + M2NParamMeta, + Placements, + check_plan_limits, + check_transferable, + comm_ptr, + import_m2n, + publish_destination_placements, + resolve_layout, + to_mesh, + to_placements, + validate_layout, +) +from vllm.distributed.weight_transfer.m2n_layout import ( + M2NDestination, + resolve_parameter_destinations, +) +from vllm.distributed.weight_transfer.nccl_common import ( + NCCLWeightTransferInitInfo, + decode_nccl_unique_id, + worker_init_process_group, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + +__all__ = [ + "M2NWeightTransferInitInfo", + "M2NWeightTransferUpdateInfo", + "M2NWeightTransferEngine", +] + + +# --------------------------------------------------------------------------- +# Wire types +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class M2NWeightTransferInitInfo(WeightTransferInitInfo): + """Worker-side init info: the rendezvous plus the full transfer plan. + + Layouts are static for the whole run, so they ride the one-time init + handshake and the per-round update info stays a list of names. Everything is + plain JSON so the HTTP control plane carries it unchanged. + """ + + master_address: str | None = None + master_port: int | None = None + 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.""" + 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 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: + raise ValueError( + f"`dst_mesh_dims` {dst} must be {MESH_NDIMS} dims covering the " + f"{num_workers} inference workers" + ) + + @property + def nccl_unique_id_bytes(self) -> bytes | None: + return decode_nccl_unique_id( + master_address=self.master_address, + master_port=self.master_port, + nccl_unique_id_b64=self.nccl_unique_id_b64, + ctx="M2NWeightTransferInitInfo", + ) + + +@dataclass +class M2NWeightTransferUpdateInfo(WeightTransferUpdateInfo): + """Per-round update info: which parameters this chunk carries, in order. + + Shapes, dtypes and layouts were fixed at init, so a round only needs to say + what is coming and in what order — both sides must issue their reshards in + exactly that order. + """ + + names: list[str] + + +# --------------------------------------------------------------------------- +# Worker engine +# --------------------------------------------------------------------------- + + +class M2NWeightTransferEngine( + WeightTransferEngine[M2NWeightTransferInitInfo, M2NWeightTransferUpdateInfo] +): + """Inference-side engine: receives each parameter with one reshard. + + Resolvable parameters are resharded from the trainer layout directly into + each worker's live local model storage. Unresolvable parameters receive a + full tensor and use `load_weights`. In both cases, the trainer sends its + local shards without materializing a full tensor. + """ + + init_info_cls = M2NWeightTransferInitInfo + update_info_cls = M2NWeightTransferUpdateInfo + # The destination plan is resolved against the model this engine was built + # for, so it cannot be retargeted at a draft model mid-flight. + supports_draft_weight_update = False + + def __init__( + self, + config: WeightTransferConfig, + vllm_config: "VllmConfig", + device: torch.device, + model: torch.nn.Module, + ) -> None: + super().__init__(config, vllm_config, device, model) + self.model_update_group: PyNcclCommunicator | None = 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" + ) + # 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 + # covers the same ranks but is not this mesh. + self._dst_mesh = M2NMesh( + 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) + # A quantized weight loader may dequantize / transpose / requantize, and + # under PP a rank's local shape is not the checkpoint shape split over + # the shard axis. Both break shape-derived resolution, so take the + # fallback. + allow_direct = ( + quantization is None and parallel_config.pipeline_parallel_size == 1 + ) + # Match every incoming checkpoint parameter to its rank-local target. + # Resolvable parameters can be resharded directly into live model + # storage; the rest receive a full tensor for `load_weights`. + self._parameter_destinations = resolve_parameter_destinations( + self.model, + [m.name for m in self._metas], + [m.dtype for m in self._metas], + [m.shape for m in self._metas], + num_workers=self._dst_mesh.size, + 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, + ) + + # 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)} + self._uses_load_weights = any( + 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. + + Directly-resharded parameters are written in place and never go through + `load_weights`, so a plan with no fallback entries has nothing to + reload. + """ + if not self._uses_load_weights: + return + from vllm.model_executor.model_loader.reload import initialize_layerwise_reload + + initialize_layerwise_reload(self.model) + + def finish_weight_update(self) -> None: + """Finalize layerwise reloading when the plan uses fallback entries.""" + if not self._uses_load_weights: + return + from vllm.model_executor.model_loader.reload import finalize_layerwise_reload + + finalize_layerwise_reload(self.model, self.model_config) + + def receive_weights(self, update_info: M2NWeightTransferUpdateInfo) -> None: + """Receive each requested parameter using its initialization-time plan.""" + if self._handle is None or self.model_update_group is None: + raise RuntimeError( + "nccl_m2n weight transfer not initialized. " + "Call init_transfer_engine() first." + ) + + requested = [] + for name in update_info.names: + index = self._index.get(name) + if index is None: + raise ValueError( + f"parameter '{name}' was not declared at init; the " + "trainer must send the same parameter set it announced" + ) + requested.append((name, index)) + + from vllm.model_executor.model_loader.mtp_validation import ( + 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 + # fallback sequence without staging every full tensor at once. Direct + # destinations still execute in request order as the iterator advances. + def reshard_requested_weights() -> Iterator[tuple[str, torch.Tensor]]: + for name, index in requested: + meta = self._metas[index] + destination = self._parameter_destinations[index] + buffer = destination.tensor + if buffer is None: + buffer = torch.empty( + meta.shape, dtype=meta.dtype, device=self.device + ) + + self._reshard(comm, stream, meta, destination.placements, buffer) + + if destination.tensor is None: + # `load_weights` reads on the host stream, so the transfer + # has to have landed before it runs. + stream.synchronize() + yield name, buffer + + has_fallback = any( + self._parameter_destinations[index].tensor is None for _, index in requested + ) + with disable_mtp_completeness_check(): + received_weights = reshard_requested_weights() + if has_fallback: + # Some model loaders keep invocation-local state to combine + # related checkpoint tensors, so preserve one iterable scope. + self.model.load_weights(received_weights) + else: + # The iterator performs the reshards. With nothing to yield to + # `load_weights`, exhaust it here to execute direct transfers. + for _ in received_weights: + pass + + def _reshard( + self, + comm: int, + stream: torch.cuda.Stream, + src_meta: M2NParamMeta, + dst_placements: Placements | None, + dst_buffer: torch.Tensor, + ) -> None: + """Reshard one parameter from its trainer layout into a worker buffer. + + `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 + ) + dst_mesh, resolved_dst_placements = resolve_layout( + self._dst_mesh, dst_placements + ) + m2n.reshard( + None, + dst_buffer, + comm, + stream, + src_mesh=to_mesh(m2n, src_mesh), + src_placements=to_placements(m2n, resolved_src_placements), + dst_mesh=to_mesh(m2n, dst_mesh), + dst_placements=to_placements(m2n, resolved_dst_placements), + handle=self._handle, + ) + + 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 + self.model_update_group = None + self._parameter_destinations = [] diff --git a/vllm/distributed/weight_transfer/m2n_layout.py b/vllm/distributed/weight_transfer/m2n_layout.py new file mode 100644 index 0000000000..3296ef63fb --- /dev/null +++ b/vllm/distributed/weight_transfer/m2n_layout.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Destination-layout resolution for the NCCL M2N weight transfer backend. + +For every incoming checkpoint parameter the worker needs a destination buffer +plus its placement over the inference mesh. Two outcomes are possible: + +* **direct** — when the model is not quantized, the parameter maps 1:1 onto a + live vLLM parameter and its loader is a small, known-safe copy or tensor- + parallel loader. The loader's declared input/output dimension determines the + placement; shapes are validation, not an inference mechanism. The reshard + writes straight into the live parameter, so each rank receives only its own + shard and nothing is copied afterwards. +* **fallback** — anything else. The reshard delivers the whole tensor to every + rank and `load_weights` does the sharding, exactly as the broadcast NCCL + backend does. Fused parameters (`qkv_proj`, `gate_up_proj`, MoE `w13`/`w2`) + take this path: the checkpoint name does not name a vLLM parameter, so there + is nothing to resolve against. + +Correctness never depends on a parameter resolving — the fallback is always +available and is the same path the existing backend uses. +""" + +import inspect +import math +from collections.abc import Sequence +from dataclasses import dataclass +from functools import partial +from typing import cast + +import torch + +from vllm.distributed.weight_transfer.m2n_common import ( + DESTINATION_REPLICA_AXIS, + DESTINATION_SHARD_AXIS, + MESH_NDIMS, + REPLICATE, + REPLICATED, + Placements, +) +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_DIRECT_WEIGHT_LOADING_BLOCKLIST = { + "vllm.model_executor.models.gpt2.GPT2Model", +} + + +def _module_type_name(module: torch.nn.Module) -> str: + cls = type(module) + return f"{cls.__module__}.{cls.__name__}" + + +def _destination_placements(shard_dim: int) -> Placements: + """Place one tensor shard on the configured destination mesh axis.""" + placements = [REPLICATE] * MESH_NDIMS + placements[DESTINATION_REPLICA_AXIS] = REPLICATE + placements[DESTINATION_SHARD_AXIS] = shard_dim + return cast(Placements, tuple(placements)) + + +@dataclass +class M2NDestination: + """Where one checkpoint parameter lands on this worker.""" + + name: str + placements: Placements | None + """Placement over the inference mesh, or `REPLICATED` for the fallback.""" + tensor: torch.Tensor | None + """The live parameter view to reshard into, or None for the fallback path + (the engine allocates a full replica per round and calls `load_weights`).""" + + @property + def direct(self) -> bool: + return self.tensor is not None + + +def _base_loader_shard_dim( + param: torch.nn.Parameter, + shard_axis_size: int, +) -> int | None: + """Return a known-safe loader's declared shard dim, or reject it. + + Loader identity is deliberately fail-closed. A missing loader is not + treated as ``default_weight_loader`` because a model-level ``load_weights`` + implementation may transform the tensor before it reaches that default. + Wrappers and partials are also rejected: recognizing their wrapped callable + would ignore behavior added by the wrapper. + """ + loader = getattr(param, "weight_loader", None) + if loader is None or isinstance(loader, partial): + return None + + owner = loader.__self__ if inspect.ismethod(loader) else None + loader_fn = loader.__func__ if inspect.ismethod(loader) else loader + + # Lazy imports avoid making the weight-transfer registry import all model + # executor layers during normal vLLM startup. + from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, + ) + from vllm.model_executor.model_loader.weight_utils import default_weight_loader + + replicated_loaders = { + default_weight_loader, + ReplicatedLinear.weight_loader, + } + column_loaders = { + ColumnParallelLinear.weight_loader, + ColumnParallelLinear.weight_loader_v2, + } + row_loaders = { + RowParallelLinear.weight_loader, + RowParallelLinear.weight_loader_v2, + } + + # These attributes signal packing, a transform, or pre-sharded checkpoint + # storage. Such semantics cannot be reproduced by a plain M2N placement. + # Check them before accepting replicated loaders too: model-level loading + # may transform a tensor before handing it to an otherwise plain copier. + metadata_attrs = ("packed_dim", "packed_factor", "marlin_tile_size") + flag_attrs = ("needs_scalar_to_array", "is_sharded_weight", "is_transposed") + if any(getattr(param, attr, None) is not None for attr in metadata_attrs): + return None + if any(bool(getattr(param, attr, False)) for attr in flag_attrs): + return None + + if loader_fn in replicated_loaders: + return REPLICATE + if loader_fn not in column_loaders | row_loaders: + return None + + if owner is None or getattr(owner, "tp_size", None) != shard_axis_size: + return None + + attr = "output_dim" if loader_fn in column_loaders else "input_dim" + dim = getattr(param, attr, None) + return dim if isinstance(dim, int) and dim >= 0 else None + + +def _validated_shard_dim( + global_shape: Sequence[int], + local_shape: Sequence[int], + declared_dim: int, + shard_axis_size: int, +) -> int | None: + """Validate a loader-declared shard dimension against both shapes.""" + if len(global_shape) != len(local_shape): + return None + + if declared_dim == REPLICATE: + return REPLICATE if tuple(global_shape) == tuple(local_shape) else None + if declared_dim >= len(global_shape): + return None + for dim, (whole, local) in enumerate(zip(global_shape, local_shape)): + expected = local * shard_axis_size if dim == declared_dim else local + if whole != expected: + return None + return declared_dim + + +def resolve_parameter_destinations( + model: torch.nn.Module, + names: Sequence[str], + dtypes: Sequence[torch.dtype], + shapes: Sequence[Sequence[int]], + *, + num_workers: int, + shard_axis_size: int, + allow_direct: bool, +) -> list[M2NDestination]: + """Build this worker's destination plan, one entry per checkpoint parameter. + + Placements are relative to the inference mesh: axis 0 replicates and axis 1 + shards. `allow_direct=False` forces every parameter onto the fallback path; + the engine sets it for pipeline-parallel or quantized deployments, where a + parameter's local shape is not simply the checkpoint shape split across the + shard axis. Parameters below a model type with known model-level checkpoint + preprocessing also fall back, since bypassing `load_weights` would skip the + transformation. + """ + params = dict(model.named_parameters()) if allow_direct else {} + blocked_param_ids = { + id(param) + for module in model.modules() + if _module_type_name(module) in _DIRECT_WEIGHT_LOADING_BLOCKLIST + for param in module.parameters() + } + + destinations: list[M2NDestination] = [] + for name, dtype, shape in zip(names, dtypes, shapes): + param = params.get(name) + dim = None + if ( + param is not None + and id(param) not in blocked_param_ids + and param.dtype == dtype + and param.data.is_contiguous() + and num_workers % shard_axis_size == 0 + ): + declared_dim = _base_loader_shard_dim(param, shard_axis_size) + if declared_dim is not None: + dim = _validated_shard_dim( + shape, param.shape, declared_dim, shard_axis_size + ) + + if dim is None: + # A replicated parameter is identical on every rank, so it needs no + # placement of its own — REPLICATED lets resolve_layout spread it + # over all the workers regardless of how the mesh is factored. + destinations.append(M2NDestination(name, REPLICATED, None)) + elif dim == REPLICATE: + assert param is not None + destinations.append(M2NDestination(name, REPLICATED, param.data)) + else: + assert param is not None + destinations.append( + M2NDestination(name, _destination_placements(dim), param.data) + ) + + # Layerwise reload finalizes a module as a unit. Mixing direct and fallback + # parameters within one module can copy stale storage over the direct load. + fallback_modules = { + name.rpartition(".")[0] + for name, destination in zip(names, destinations) + if not destination.direct and name in params + } + destinations = [ + M2NDestination(destination.name, REPLICATED, None) + if destination.direct + and destination.name.rpartition(".")[0] in fallback_modules + else destination + for destination in destinations + ] + + num_direct = sum(d.direct for d in destinations) + parameter_bytes = [ + math.prod(shape) * dtype.itemsize for dtype, shape in zip(dtypes, shapes) + ] + total_bytes = sum(parameter_bytes) + direct_bytes = sum( + nbytes + for destination, nbytes in zip(destinations, parameter_bytes) + if destination.direct + ) + direct_byte_percentage = 100 * direct_bytes / total_bytes if total_bytes else 0 + logger.info( + "nccl_m2n destination plan: %d/%d parameters resharded directly into " + "the model, %d via full-tensor fallback; direct byte coverage: " + "%d/%d bytes (%.1f%%)", + num_direct, + len(destinations), + len(destinations) - num_direct, + direct_bytes, + total_bytes, + direct_byte_percentage, + ) + return destinations diff --git a/vllm/distributed/weight_transfer/m2n_source.py b/vllm/distributed/weight_transfer/m2n_source.py new file mode 100644 index 0000000000..dd08d24b56 --- /dev/null +++ b/vllm/distributed/weight_transfer/m2n_source.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Trainer-side weight sources for the NCCL M2N backend. + +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. +""" + +from collections.abc import Iterator +from typing import Any + +import torch + +from vllm.distributed.weight_transfer.base import ParamMeta, WeightSource +from vllm.distributed.weight_transfer.m2n_common import ( + REPLICATE, + REPLICATED, + M2NMesh, + M2NParamMeta, + Placements, +) + +__all__ = [ + "M2NWeightSource", + "DTensorModuleSource", + "mesh_from_tensor", + "placements_from_tensor", +] + + +class M2NWeightSource(WeightSource): + """A `WeightSource` that also describes how the trainer holds its weights. + + Unlike `ModuleSource`, iteration yields each rank's **local shard**, not a + materialized full tensor: gathering is exactly the cost m2n removes. + """ + + def mesh(self) -> M2NMesh: + """The trainer's rank topology, shared by every parameter.""" + raise NotImplementedError + + def metadata(self) -> list[ParamMeta]: + """Name, dtype, full shape, and trainer placement for each parameter.""" + raise NotImplementedError + + def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: + """Yield `(name, local shard)` pairs in the same order as `metadata()`.""" + raise NotImplementedError + + +def _placement_code(placement: Any) -> int: + """Map a `torch.distributed` placement onto an m2n placement code.""" + name = type(placement).__name__ + if name == "Replicate": + return REPLICATE + if name == "Shard": + return int(placement.dim) + raise ValueError( + f"nccl_m2n cannot express the {name} placement; only Replicate and " + "Shard are supported" + ) + + +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. + """ + device_mesh = getattr(tensor, "device_mesh", None) + if device_mesh is None: + return M2NMesh((num_trainer_ranks, 1), 0) + + grid = device_mesh.mesh + ranks = grid.flatten().tolist() + if ranks != list(range(len(ranks))): + raise ValueError( + "nccl_m2n requires the trainer's device mesh to cover the " + f"contiguous rank interval [0, {len(ranks)}); got {ranks}" + ) + if grid.ndim > 2: + raise ValueError( + f"nccl_m2n supports 1-D and 2-D device meshes, got {grid.ndim}-D" + ) + dims = tuple(grid.shape) + return M2NMesh(dims if len(dims) == 2 else (dims[0], 1), 0) + + +def placements_from_tensor(tensor: torch.Tensor) -> Placements | None: + """How a parameter is placed over its mesh, or `REPLICATED`. + + `REPLICATED` is returned rather than a `(REPLICATE, REPLICATE)` pair: m2n + cannot express that, and the size-1-axis encoding it needs instead is + applied later by `resolve_layout`, which owns that workaround. + """ + placements = getattr(tensor, "placements", None) + if placements is None: + return REPLICATED + codes = [_placement_code(p) for p in placements] + if all(code == REPLICATE for code in codes): + return REPLICATED + if len(codes) == 1: + return (codes[0], REPLICATE) + return (codes[0], codes[1]) + + +class DTensorModuleSource(M2NWeightSource): + """`M2NWeightSource` over `module.named_parameters()`. + + Covers both the FSDP/DTensor trainer (placement read off each parameter, + local shard yielded via `to_local()`) and the plain replicated trainer, with + no special casing. Trainers with a custom producer (a Megatron export, MoE + re-fusing) subclass `M2NWeightSource` instead. + """ + + def __init__(self, module: torch.nn.Module, num_trainer_ranks: int) -> None: + """Wrap `module`; `num_trainer_ranks` sizes the mesh for plain tensors.""" + 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.""" + return [ + M2NParamMeta(name, p.dtype, tuple(p.shape), placements_from_tensor(p)) + for name, p in self._module.named_parameters() + ] + + def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: + """Yield each parameter's local shard (`to_local()`), or the tensor itself.""" + for name, param in self._module.named_parameters(): + to_local = getattr(param, "to_local", None) + yield name, (to_local() if callable(to_local) else param) diff --git a/vllm/distributed/weight_transfer/m2n_trainer.py b/vllm/distributed/weight_transfer/m2n_trainer.py new file mode 100644 index 0000000000..68c62f72db --- /dev/null +++ b/vllm/distributed/weight_transfer/m2n_trainer.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Trainer-side weight transfer engine for the NCCL M2N backend. + +Symmetric to `M2NWeightTransferEngine` but in the training process. Every +trainer rank joins the shared communicator and runs every reshard, sending its +own local shard; only rank 0 touches the inference control plane. + +Workers join the communicator during init and run reshard from inside +`update_weights`, so those two RPCs overlap with work on this side. Both run on +a helper thread and are joined afterwards, which keeps that overlap inside the +engine -- where the `TrainerWeightTransferEngine` contract puts it -- rather than +in every caller. +""" + +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, ClassVar, cast + +import torch +from typing_extensions import Self + +from vllm.distributed.weight_transfer.base import ( + TrainerInitInfo, + TrainerWeightTransferEngine, + VLLMWeightSyncClient, + WeightSource, +) +from vllm.distributed.weight_transfer.m2n_common import ( + M2NMesh, + M2NParamMeta, + Placements, + check_transferable, + comm_ptr, + import_m2n, + publish_destination_placements, + resolve_layout, + to_mesh, + to_placements, + validate_layout, +) +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, +) +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + +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. + + `rank` (from `TrainerInitInfo`) is this trainer process's rank; it is also + its rank in the shared communicator, since the trainer occupies + `[0, num_trainer_ranks)`. Rank 0 drives the control plane. + """ + + backend: ClassVar[str] = "nccl_m2n" + + master_address: str | None = None + master_port: int | None = None + world_size: int + """Trainer ranks + all inference workers.""" + num_trainer_ranks: int = 1 + dst_mesh_dims: tuple[int, int] | None = None + """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 + nccl_unique_id_b64: str | None = field(default=None, repr=False) + + @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 (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 0 < self.num_trainer_ranks < self.world_size: + raise ValueError( + f"`num_trainer_ranks` ({self.num_trainer_ranks}) must leave at " + f"least one worker in world_size {self.world_size}" + ) + 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: + raise ValueError( + f"`dst_mesh_dims` {dims} must be two dims covering the " + f"{num_workers} inference workers" + ) + if not 0 <= self.rank < self.num_trainer_ranks: + raise ValueError( + f"trainer `rank` ({self.rank}) must be below " + f"`num_trainer_ranks` ({self.num_trainer_ranks})" + ) + + @property + def nccl_unique_id_bytes(self) -> bytes | None: + return decode_nccl_unique_id( + master_address=self.master_address, + master_port=self.master_port, + nccl_unique_id_b64=self.nccl_unique_id_b64, + ctx="M2NTrainerInitInfo", + ) + + +class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitInfo]): + """Trainer-side engine: sends every rank's local shard, once per parameter. + + Called on every trainer rank. All ranks join the communicator and run every + reshard; only rank 0 touches the control plane. + """ + + init_info_cls = M2NTrainerInitInfo + + def __init__( + self, + *, + client: VLLMWeightSyncClient, + source: M2NWeightSource, + is_controller: bool, + num_trainer_ranks: int, + ) -> 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 + 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._dst_placements: list[Placements | None] = [] + self._executor: ThreadPoolExecutor | None = None + + @classmethod + def trainer_init( + cls, + init_info: M2NTrainerInitInfo, + *, + client: VLLMWeightSyncClient, + source: WeightSource | None = None, + ) -> Self: + """Build the engine and rendezvous with the inference side. + + Runs on *every* trainer rank. Rank 0 additionally ships the transfer + plan to the workers and drives the control plane; the other ranks only + 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. + """ + # 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, + ) + 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() + 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" + ) + + # 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) + ) + # 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] = { + "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 + ], + "max_cta": init_info.max_cta, + } + return {key: value for key, value in payload.items() if value is not None} + + 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]}, + ) + + try: + self._send() + finally: + if update is not None: + update.result() + + if self.is_controller: + self.client.finish_weight_update() + + 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 + m2n = self._m2n + comm = comm_ptr(self.group) + stream = torch.cuda.current_stream() + + for index, (name, tensor) in enumerate(self.source): + meta = self._metas[index] + if name != meta.name: + raise RuntimeError( + f"weight source yielded '{name}' at position {index}, but " + 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) + dst_mesh, dst_placements = resolve_layout( + self._dst_mesh, self._dst_placements[index] + ) + m2n.reshard( + shard, + None, + comm, + 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, + ) + # 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 + # safely mutate parameters or continue work on another stream. + 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 + self.group = None diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py index 57a17b9830..a329e456d5 100644 --- a/vllm/distributed/weight_transfer/nccl_common.py +++ b/vllm/distributed/weight_transfer/nccl_common.py @@ -220,11 +220,11 @@ def worker_init_process_group( def trainer_init( init_info: NCCLRendezvous | dict, + rank: int = 0, ) -> "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()). + Uses the current CUDA device (torch.accelerator.current_device_index()). Args: init_info: Any object carrying the `NCCLRendezvous` fields (a trainer or @@ -232,6 +232,9 @@ def trainer_init( - master_address: str - master_port: int - world_size: int + rank: This trainer process's rank in the group. The broadcast backends + have a single trainer at rank 0; multi-rank trainers (m2n) occupy + ranks `[0, num_trainer_ranks)` and the workers start after them. Returns: PyNcclCommunicator for weight transfer. @@ -246,12 +249,11 @@ def trainer_init( 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, + rank, world_size, device, )