# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Sparse NCCL weight transfer engine. A standalone engine (not a subclass of `NCCLWeightTransferEngine`) for applying sparse, flat-index weight patches in place. It shares only NCCL process-group initialization with the dense engine (via `nccl_common`); the update path applies index/value patches directly to existing model parameters and never runs layerwise reload. MVP limitations: * TP=1 and PP=1 only * uses runtime/kernel-format parameter names * not composable with checkpoint-format or packed updates """ from collections.abc import Iterator from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( WeightTransferEngine, WeightTransferUpdateInfo, ) from vllm.distributed.weight_transfer.nccl_common import ( NCCLWeightTransferInitInfo, trainer_init, worker_init_process_group, ) from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerSendWeightsArgs __all__ = [ "SparseWeightPatch", "SparseNCCLWeightTransferUpdateInfo", "SparseNCCLWeightTransferEngine", ] @dataclass class SparseWeightPatch: """A sparse in-place patch for one existing parameter.""" name: str indices: torch.Tensor values: torch.Tensor @dataclass class SparseNCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo): """Update info for the sparse NCCL weight transfer backend.""" names: list[str] dtype_names: list[str] shapes: list[list[int]] num_updates_list: list[int] """Number of sparse entries to receive for each parameter in ``names``.""" def __post_init__(self) -> None: num_params = len(self.names) if len(self.dtype_names) != num_params: raise ValueError( f"`dtype_names` should be of the same size as `names`: " f"got {len(self.dtype_names)} and {len(self.names)}" ) if len(self.shapes) != num_params: raise ValueError( f"`shapes` should be of the same size as `names`: " f"got {len(self.shapes)} and {len(self.names)}" ) if len(self.num_updates_list) == 0: raise ValueError("`num_updates_list` cannot be empty for sparse updates") if len(self.num_updates_list) != num_params: raise ValueError( f"`num_updates_list` should be of the same size as `names`: " f"got {len(self.num_updates_list)} and {len(self.names)}" ) if any(num_updates < 0 for num_updates in self.num_updates_list): raise ValueError("Sparse `num_updates_list` entries must be non-negative") class SparseNCCLWeightTransferEngine( WeightTransferEngine[NCCLWeightTransferInitInfo, SparseNCCLWeightTransferUpdateInfo] ): """ Sparse weight transfer engine using NCCL. Receives flat-index (indices, values) patches broadcast from the trainer (rank 0) and applies them in place to existing model parameters. Weights are applied directly without layerwise reload, so `start_weight_update` and `finish_weight_update` are no-ops. """ init_info_cls = NCCLWeightTransferInitInfo update_info_cls = SparseNCCLWeightTransferUpdateInfo 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 def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None: """Initialize the NCCL process group with the trainer.""" self.model_update_group = worker_init_process_group( init_info, self.parallel_config ) def start_weight_update(self) -> None: """No-op: sparse patches are applied in place, no layerwise reload.""" if self.parallel_config.world_size != 1: raise NotImplementedError( "Sparse weight updates currently require TP=1 and PP=1" ) def finish_weight_update(self) -> None: """No-op: sparse patches are applied in place, no layerwise reload.""" pass def receive_weights(self, update_info: SparseNCCLWeightTransferUpdateInfo) -> None: """Receive sparse flat-index patches from the trainer and apply them.""" if self.model_update_group is None: raise RuntimeError( "NCCL weight transfer not initialized. " "Call init_transfer_engine() first." ) # Use the worker's assigned device rather than the ambient current # device: the receive path is no longer wrapped in # `with torch.device(self.device)` by the caller, so the current device # is not guaranteed to match self.device. The recv buffers must live on # the same device as the NCCL communicator (created on self.device). device = self.device for name, dtype_name, num_updates in zip( update_info.names, update_info.dtype_names, update_info.num_updates_list, ): dtype = getattr(torch, dtype_name) indices = torch.empty(num_updates, dtype=torch.int32, device=device) values = torch.empty(num_updates, dtype=dtype, device=device) self.model_update_group.broadcast( indices, src=0, stream=torch.cuda.current_stream() ) self.model_update_group.broadcast( values, src=0, stream=torch.cuda.current_stream() ) self._apply_patch( SparseWeightPatch(name=name, indices=indices, values=values) ) del indices del values def _apply_patch(self, patch: SparseWeightPatch) -> None: """Apply a single sparse flat-index patch to an existing model param.""" param = self.model.get_parameter(patch.name) if not param.data.is_contiguous(): raise NotImplementedError( "Sparse weight updates currently require contiguous params: " f"{patch.name}" ) if patch.indices.dtype != torch.int32: raise ValueError( f"Sparse weight updates currently require int32 indices: {patch.name}" ) if patch.indices.ndim != 1 or patch.values.ndim != 1: raise ValueError( f"Sparse weight patches must be 1D flattened updates: {patch.name}" ) if patch.indices.numel() != patch.values.numel(): raise ValueError( f"`indices` and `values` must have matching lengths for {patch.name}" ) if patch.values.dtype != param.dtype: raise ValueError( f"Sparse values dtype {patch.values.dtype} does not match " f"parameter dtype {param.dtype} for {patch.name}" ) flat_param = param.data.view(-1) flat_param.index_copy_( 0, patch.indices.to(device=flat_param.device, dtype=torch.long), patch.values.to(device=flat_param.device), ) def shutdown(self) -> None: if self.model_update_group is not None: self.model_update_group = None @staticmethod def trainer_send_weights( iterator: Iterator[SparseWeightPatch], trainer_args: dict[str, Any] | NCCLTrainerSendWeightsArgs, ) -> None: """Broadcast sparse flat-index patches from trainer to vLLM workers.""" if isinstance(trainer_args, dict): args = NCCLTrainerSendWeightsArgs(**trainer_args) else: args = trainer_args if args.packed: raise ValueError( "Sparse NCCL updates cannot be combined with `packed=True`" ) stream = args.stream or torch.cuda.current_stream() for patch in iterator: args.group.broadcast(patch.indices, src=args.src, stream=stream) args.group.broadcast(patch.values, src=args.src, stream=stream) # Trainer-side process-group setup (shared with the dense engine). trainer_init = staticmethod(trainer_init)