| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the NCCL M2N weight transfer backend. |
| 4 | |
| 5 | These cover how a transfer is *described* — layout encoding and the wire-type |
| 6 | validation that keeps a bad plan from reaching a collective, since a mismatch |
| 7 | there is a hang rather than an error. The transfer itself needs the m2n runtime |
| 8 | and multiple GPUs, so it is exercised separately. |
| 9 | """ |
| 10 | |
| 11 | import logging |
| 12 | from functools import wraps |
| 13 | from types import MethodType |
| 14 | from unittest.mock import Mock |
| 15 | |
| 16 | import pybase64 as base64 |
| 17 | import pytest |
| 18 | import ray |
| 19 | import torch |
| 20 | |
| 21 | from vllm.distributed.weight_transfer import ( |
| 22 | WeightTransferEngineFactory, |
| 23 | WeightTransferTrainerFactory, |
| 24 | ) |
| 25 | from vllm.distributed.weight_transfer.m2n_common import ( |
| 26 | REPLICATE, |
| 27 | REPLICATED, |
| 28 | M2NMesh, |
| 29 | M2NParamMeta, |
| 30 | check_placements, |
| 31 | check_transferable, |
| 32 | publish_destination_placements, |
| 33 | resolve_layout, |
| 34 | validate_layout, |
| 35 | ) |
| 36 | from vllm.distributed.weight_transfer.m2n_engine import ( |
| 37 | M2NWeightTransferEngine, |
| 38 | M2NWeightTransferInitInfo, |
| 39 | M2NWeightTransferUpdateInfo, |
| 40 | ) |
| 41 | from vllm.distributed.weight_transfer.m2n_layout import ( |
| 42 | M2NDestination, |
| 43 | resolve_parameter_destinations, |
| 44 | ) |
| 45 | from vllm.distributed.weight_transfer.m2n_source import ( |
| 46 | mesh_from_tensor, |
| 47 | placements_from_tensor, |
| 48 | ) |
| 49 | from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo |
| 50 | from vllm.model_executor.layers.linear import ( |
| 51 | ColumnParallelLinear, |
| 52 | RowParallelLinear, |
| 53 | ) |
| 54 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 55 | from vllm.platforms import current_platform |
| 56 | |
| 57 | VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode() |
| 58 | |
| 59 | |
| 60 | class TestLayout: |
| 61 | def test_replicated_splits_nothing(self): |
| 62 | """A replicated tensor imposes no divisibility constraint. The shape |
| 63 | dims are deliberately coprime with the mesh size: a layout that actually |
| 64 | split the tensor would reject them.""" |
| 65 | mesh = M2NMesh((2, 2), start_rank=1) |
| 66 | indivisible_shape = (7, 13) # neither dim divisible by any mesh axis |
| 67 | |
| 68 | resolved_mesh, placements = resolve_layout(mesh, REPLICATED) |
| 69 | validate_layout(resolved_mesh, placements, indivisible_shape, "destination") |
| 70 | |
| 71 | def test_replicated_keeps_the_same_ranks(self): |
| 72 | """Replication re-factors the mesh to get a size-1 axis for its no-op |
| 73 | shard. That is only sound if it still covers exactly the same GPUs.""" |
| 74 | mesh = M2NMesh((2, 3), start_rank=4) |
| 75 | resolved_mesh, _ = resolve_layout(mesh, REPLICATED) |
| 76 | assert resolved_mesh.size == mesh.size |
| 77 | assert resolved_mesh.start_rank == mesh.start_rank |
| 78 | |
| 79 | def test_sharded_keeps_its_own_factorization(self): |
| 80 | """Rank order decides who owns which shard, so a sharded tensor must |
| 81 | not be re-factored the way a replicated one is.""" |
| 82 | mesh = M2NMesh((2, 3), start_rank=4) |
| 83 | resolved_mesh, placements = resolve_layout(mesh, (REPLICATE, 0)) |
| 84 | assert resolved_mesh == mesh |
| 85 | assert placements == (REPLICATE, 0) |
| 86 | |
| 87 | def test_two_shard_axes_rejected(self): |
| 88 | """One axis has to replicate; a 2-D mesh that shards both is not |
| 89 | something a single reshard can express.""" |
| 90 | with pytest.raises(ValueError, match="shards both"): |
| 91 | check_placements((0, 1)) |
| 92 | |
| 93 | def test_placement_below_replicate_rejected(self): |
| 94 | with pytest.raises(ValueError, match=r"parameter 'w'.*-2"): |
| 95 | check_placements((-2, REPLICATE), "parameter 'w' source placements") |
| 96 | |
| 97 | def test_shard_dim_must_exist(self): |
| 98 | with pytest.raises(ValueError, match="rank 2"): |
| 99 | validate_layout(M2NMesh((1, 2), 0), (REPLICATE, 2), (8, 16), "source") |
| 100 | |
| 101 | def test_shard_must_divide_evenly(self): |
| 102 | with pytest.raises(ValueError, match="does not divide evenly"): |
| 103 | validate_layout(M2NMesh((1, 3), 0), (REPLICATE, 0), (8, 16), "source") |
| 104 | |
| 105 | |
| 106 | class TestSourceLayout: |
| 107 | def test_plain_tensor_is_replicated_across_trainer_ranks(self): |
| 108 | """A tensor with no DTensor metadata is the same on every trainer rank, |
| 109 | so the source describes it as replicated over all of them.""" |
| 110 | assert placements_from_tensor(torch.zeros(4)) is REPLICATED |
| 111 | assert mesh_from_tensor(torch.zeros(4), 4) == M2NMesh((4, 1), 0) |
| 112 | |
| 113 | |
| 114 | class TestTransferable: |
| 115 | def test_unsupported_dtype_names_the_parameter(self): |
| 116 | with pytest.raises(ValueError, match="'w'"): |
| 117 | check_transferable("w", torch.complex64, (4,)) |
| 118 | |
| 119 | def test_rank_four_rejected(self): |
| 120 | with pytest.raises(ValueError, match="rank 4"): |
| 121 | check_transferable("w", torch.bfloat16, (2, 2, 2, 2)) |
| 122 | |
| 123 | |
| 124 | class TestWireTypes: |
| 125 | def _init_info(self, **overrides): |
| 126 | fields = dict( |
| 127 | master_address="127.0.0.1", |
| 128 | master_port=1234, |
| 129 | rank_offset=1, |
| 130 | world_size=3, |
| 131 | src_mesh_dims=[1, 1], |
| 132 | dst_mesh_dims=[2, 1], |
| 133 | names=["w"], |
| 134 | dtype_names=["bfloat16"], |
| 135 | shapes=[[16, 16]], |
| 136 | src_placements=[None], |
| 137 | ) |
| 138 | fields.update(overrides) |
| 139 | return M2NWeightTransferInitInfo(**fields) |
| 140 | |
| 141 | def test_accepts_a_consistent_plan(self): |
| 142 | assert self._init_info().names == ["w"] |
| 143 | |
| 144 | def test_accepts_pre_shared_nccl_unique_id(self): |
| 145 | info = self._init_info( |
| 146 | master_address=None, |
| 147 | master_port=None, |
| 148 | nccl_unique_id_b64=VALID_UID_B64, |
| 149 | ) |
| 150 | assert info.nccl_unique_id_bytes == b"\x00" * 128 |
| 151 | assert VALID_UID_B64 not in repr(info) |
| 152 | |
| 153 | def test_destination_plan_uses_uid_communicator_without_group(self): |
| 154 | comm = Mock(rank=0, device=torch.device("cpu"), group=None) |
| 155 | |
| 156 | def receive_plan(tensor, src): |
| 157 | assert src == 1 |
| 158 | tensor.copy_(torch.tensor([[-1, 0], [-2, -2]], dtype=torch.int8)) |
| 159 | |
| 160 | comm.broadcast.side_effect = receive_plan |
| 161 | |
| 162 | placements = publish_destination_placements(comm, 1, None, 2) |
| 163 | |
| 164 | assert placements == [(REPLICATE, 0), REPLICATED] |
| 165 | |
| 166 | def test_rejects_both_rendezvous_modes(self): |
| 167 | with pytest.raises(ValueError, match="not both"): |
| 168 | self._init_info(nccl_unique_id_b64=VALID_UID_B64) |
| 169 | |
| 170 | def test_trainer_accepts_pre_shared_nccl_unique_id(self): |
| 171 | info = M2NTrainerInitInfo( |
| 172 | nccl_unique_id_b64=VALID_UID_B64, |
| 173 | world_size=3, |
| 174 | num_trainer_ranks=1, |
| 175 | rank=0, |
| 176 | ) |
| 177 | assert info.nccl_unique_id_bytes == b"\x00" * 128 |
| 178 | assert VALID_UID_B64 not in repr(info) |
| 179 | |
| 180 | def test_ragged_plan_rejected(self): |
| 181 | with pytest.raises(ValueError, match="`shapes`"): |
| 182 | self._init_info(shapes=[]) |
| 183 | |
| 184 | def test_destination_mesh_must_cover_the_workers(self): |
| 185 | """The trainer declares the inference mesh, so one that does not cover |
| 186 | the workers is a config error — and it has to fail the init RPC, since |
| 187 | a mismatched mesh would otherwise surface as a hung collective.""" |
| 188 | with pytest.raises(ValueError, match="dst_mesh_dims"): |
| 189 | self._init_info(dst_mesh_dims=[3, 1]) # 3 != the 2 workers |
| 190 | |
| 191 | def test_world_must_hold_a_trainer_and_a_worker(self): |
| 192 | with pytest.raises(ValueError, match="rank_offset"): |
| 193 | self._init_info(rank_offset=3, world_size=3) |
| 194 | |
| 195 | @pytest.mark.parametrize("dtype_name", ["not_a_dtype", "Tensor"]) |
| 196 | def test_invalid_dtype_name_names_parameter(self, monkeypatch, dtype_name): |
| 197 | monkeypatch.setattr( |
| 198 | "vllm.distributed.weight_transfer.m2n_engine.import_m2n", |
| 199 | lambda: object(), |
| 200 | ) |
| 201 | engine = object.__new__(M2NWeightTransferEngine) |
| 202 | |
| 203 | with pytest.raises(ValueError, match=r"parameter 'w'.*dtype"): |
| 204 | engine.init_transfer_engine(self._init_info(dtype_names=[dtype_name])) |
| 205 | |
| 206 | def _init_32_to_4_plan(self, monkeypatch, destination): |
| 207 | m2n = Mock() |
| 208 | m2n.Handle.create.return_value = object() |
| 209 | monkeypatch.setattr( |
| 210 | "vllm.distributed.weight_transfer.m2n_engine.import_m2n", lambda: m2n |
| 211 | ) |
| 212 | monkeypatch.setattr( |
| 213 | "vllm.distributed.weight_transfer.m2n_engine." |
| 214 | "resolve_parameter_destinations", |
| 215 | lambda *args, **kwargs: [destination], |
| 216 | ) |
| 217 | monkeypatch.setattr( |
| 218 | "vllm.distributed.weight_transfer.m2n_engine.worker_init_process_group", |
| 219 | lambda *args, **kwargs: object(), |
| 220 | ) |
| 221 | monkeypatch.setattr( |
| 222 | "vllm.distributed.weight_transfer.m2n_engine." |
| 223 | "publish_destination_placements", |
| 224 | lambda *args: args[2], |
| 225 | ) |
| 226 | |
| 227 | engine = object.__new__(M2NWeightTransferEngine) |
| 228 | engine.parallel_config = Mock(pipeline_parallel_size=1) |
| 229 | engine.model_config = Mock(quantization=None) |
| 230 | engine.model = torch.nn.Module() |
| 231 | engine.init_transfer_engine( |
| 232 | self._init_info( |
| 233 | rank_offset=32, |
| 234 | world_size=36, |
| 235 | src_mesh_dims=[1, 32], |
| 236 | dst_mesh_dims=[1, 4], |
| 237 | shapes=[[32, 8]], |
| 238 | src_placements=[[REPLICATE, 0]], |
| 239 | ) |
| 240 | ) |
| 241 | |
| 242 | def test_limits_use_resolved_sharded_destination(self, monkeypatch): |
| 243 | destination = M2NDestination("w", (REPLICATE, 0), torch.empty(8, 8)) |
| 244 | |
| 245 | self._init_32_to_4_plan(monkeypatch, destination) |
| 246 | |
| 247 | def test_limits_still_reject_replicated_destination(self, monkeypatch): |
| 248 | destination = M2NDestination("w", REPLICATED, None) |
| 249 | |
| 250 | with pytest.raises(ValueError, match=r"32 source shards.*MAX_SOURCES=16"): |
| 251 | self._init_32_to_4_plan(monkeypatch, destination) |
| 252 | |
| 253 | def test_update_preflights_all_names_before_reshard(self): |
| 254 | engine = object.__new__(M2NWeightTransferEngine) |
| 255 | engine._handle = object() |
| 256 | engine.model_update_group = object() |
| 257 | engine._index = {"valid": 0} |
| 258 | engine._metas = [M2NParamMeta("valid", torch.float32, (4,), REPLICATED)] |
| 259 | engine._reshard = Mock() |
| 260 | |
| 261 | with pytest.raises(ValueError, match=r"parameter 'unknown'"): |
| 262 | engine.receive_weights( |
| 263 | M2NWeightTransferUpdateInfo(names=["valid", "unknown"]) |
| 264 | ) |
| 265 | |
| 266 | engine._reshard.assert_not_called() |
| 267 | |
| 268 | def test_fallback_tensors_share_one_load_weights_invocation(self, monkeypatch): |
| 269 | class CoupledModel: |
| 270 | def __init__(self): |
| 271 | self.calls = 0 |
| 272 | self.loaded = [] |
| 273 | |
| 274 | def load_weights(self, weights): |
| 275 | self.calls += 1 |
| 276 | received = list(weights) |
| 277 | if [name for name, _ in received] == ["pair.weight", "pair.scale"]: |
| 278 | self.loaded = [(name, tensor.clone()) for name, tensor in received] |
| 279 | |
| 280 | model = CoupledModel() |
| 281 | direct = torch.zeros(1) |
| 282 | engine = object.__new__(M2NWeightTransferEngine) |
| 283 | engine._handle = object() |
| 284 | engine.model_update_group = object() |
| 285 | engine.device = torch.device("cpu") |
| 286 | engine.model = model |
| 287 | engine._index = {"pair.weight": 0, "direct": 1, "pair.scale": 2} |
| 288 | engine._metas = [ |
| 289 | M2NParamMeta(name, torch.float32, (1,), REPLICATED) |
| 290 | for name in engine._index |
| 291 | ] |
| 292 | engine._parameter_destinations = [ |
| 293 | M2NDestination("pair.weight", REPLICATED, None), |
| 294 | M2NDestination("direct", REPLICATED, direct), |
| 295 | M2NDestination("pair.scale", REPLICATED, None), |
| 296 | ] |
| 297 | reshard_order = [] |
| 298 | |
| 299 | def reshard(comm, stream, meta, placements, buffer): |
| 300 | reshard_order.append(meta.name) |
| 301 | buffer.fill_(len(reshard_order)) |
| 302 | |
| 303 | engine._reshard = reshard |
| 304 | stream = Mock() |
| 305 | monkeypatch.setattr( |
| 306 | "vllm.distributed.weight_transfer.m2n_engine.comm_ptr", lambda group: 0 |
| 307 | ) |
| 308 | monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream) |
| 309 | |
| 310 | engine.receive_weights( |
| 311 | M2NWeightTransferUpdateInfo(names=["pair.weight", "direct", "pair.scale"]) |
| 312 | ) |
| 313 | |
| 314 | assert model.calls == 1 |
| 315 | assert [name for name, _ in model.loaded] == ["pair.weight", "pair.scale"] |
| 316 | assert [tensor.item() for _, tensor in model.loaded] == [1, 3] |
| 317 | assert direct.item() == 2 |
| 318 | assert reshard_order == ["pair.weight", "direct", "pair.scale"] |
| 319 | assert stream.synchronize.call_count == 2 |
| 320 | |
| 321 | def test_trainer_rank_must_be_a_trainer_rank(self): |
| 322 | """A trainer rank must fall within the trainer portion of the group.""" |
| 323 | with pytest.raises(ValueError, match="num_trainer_ranks"): |
| 324 | M2NTrainerInitInfo( |
| 325 | master_address="127.0.0.1", |
| 326 | master_port=1234, |
| 327 | world_size=4, |
| 328 | num_trainer_ranks=2, |
| 329 | rank=2, |
| 330 | ) |
| 331 | |
| 332 | def test_trainer_destination_mesh_must_cover_the_workers(self): |
| 333 | """The destination mesh must include every inference worker.""" |
| 334 | with pytest.raises(ValueError, match="dst_mesh_dims"): |
| 335 | M2NTrainerInitInfo( |
| 336 | master_address="127.0.0.1", |
| 337 | master_port=1234, |
| 338 | world_size=6, |
| 339 | num_trainer_ranks=2, |
| 340 | dst_mesh_dims=(3, 1), # 3 != the 4 inference workers |
| 341 | rank=0, |
| 342 | ) |
| 343 | |
| 344 | def test_destination_mesh_defaults_to_flat(self): |
| 345 | """A replicated destination does not care how the mesh is factored, so |
| 346 | callers that do not shard it need not supply one.""" |
| 347 | info = M2NTrainerInitInfo( |
| 348 | master_address="127.0.0.1", |
| 349 | master_port=1234, |
| 350 | world_size=6, |
| 351 | num_trainer_ranks=2, |
| 352 | rank=0, |
| 353 | ) |
| 354 | assert info.destination_mesh_dims == (4, 1) |
| 355 | |
| 356 | def test_sender_is_trainer_rank_zero(self): |
| 357 | """Trainer rank 0 drives the inference control plane.""" |
| 358 | info = M2NTrainerInitInfo( |
| 359 | master_address="127.0.0.1", master_port=1234, world_size=4, rank=0 |
| 360 | ) |
| 361 | assert info.is_sender |
| 362 | |
| 363 | |
| 364 | class TestRegistration: |
| 365 | def test_both_registries_expose_the_backend(self): |
| 366 | """Both worker and trainer factories register nccl_m2n.""" |
| 367 | assert "nccl_m2n" in WeightTransferEngineFactory._registry |
| 368 | assert "nccl_m2n" in WeightTransferTrainerFactory._registry |
| 369 | |
| 370 | def test_init_info_dispatches_to_the_backend(self): |
| 371 | """Trainer init info selects the nccl_m2n factory entry.""" |
| 372 | assert M2NTrainerInitInfo.backend == "nccl_m2n" |
| 373 | |
| 374 | |
| 375 | # --------------------------------------------------------------------------- |
| 376 | # End-to-end transfer |
| 377 | # |
| 378 | # Transport-only, in the style of the NCCL/sparse tests in |
| 379 | # test_weight_transfer.py: two Ray tasks with one GPU each drive the two engines |
| 380 | # directly, so no HTTP server or LLM instance is involved. The worker is not an |
| 381 | # RPC endpoint here, so the trainer engine gets a no-op control-plane client -- |
| 382 | # the NCCL rendezvous and the reshard itself are the real thing. |
| 383 | # --------------------------------------------------------------------------- |
| 384 | |
| 385 | SHAPE = [64, 32] |
| 386 | DTYPE = "float32" |
| 387 | |
| 388 | |
| 389 | def _init_ray() -> None: |
| 390 | if ray.is_initialized(): |
| 391 | return |
| 392 | ray.init( |
| 393 | ignore_reinit_error=True, |
| 394 | runtime_env={ |
| 395 | "env_vars": { |
| 396 | "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", |
| 397 | "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1", |
| 398 | "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1", |
| 399 | } |
| 400 | }, |
| 401 | ) |
| 402 | |
| 403 | |
| 404 | def _assigned_device() -> "torch.device": |
| 405 | gpu_ids = ray.get_gpu_ids() |
| 406 | device = torch.device(f"cuda:{int(gpu_ids[0])}" if gpu_ids else "cuda:0") |
| 407 | current_platform.set_device(device) |
| 408 | return device |
| 409 | |
| 410 | |
| 411 | @ray.remote(num_gpus=1) |
| 412 | def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool: |
| 413 | """Send one parameter through the real trainer engine.""" |
| 414 | device = _assigned_device() |
| 415 | |
| 416 | from vllm.distributed.weight_transfer import WeightTransferTrainerFactory |
| 417 | from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource |
| 418 | from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo |
| 419 | |
| 420 | class NoopClient: |
| 421 | def init_weight_transfer_engine(self, init_info): |
| 422 | pass |
| 423 | |
| 424 | def start_weight_update(self): |
| 425 | pass |
| 426 | |
| 427 | def update_weights(self, update_info): |
| 428 | pass |
| 429 | |
| 430 | def finish_weight_update(self, weight_version=None): |
| 431 | pass |
| 432 | |
| 433 | class Tiny(torch.nn.Module): |
| 434 | def __init__(self): |
| 435 | super().__init__() |
| 436 | self.weight = torch.nn.Parameter( |
| 437 | torch.arange( |
| 438 | SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device |
| 439 | ).reshape(SHAPE) |
| 440 | ) |
| 441 | |
| 442 | engine = WeightTransferTrainerFactory.trainer_init( |
| 443 | init_info=M2NTrainerInitInfo( |
| 444 | nccl_unique_id_b64=nccl_unique_id_b64, |
| 445 | world_size=world_size, |
| 446 | num_trainer_ranks=1, |
| 447 | dst_mesh_dims=(1, world_size - 1), |
| 448 | rank=0, |
| 449 | ), |
| 450 | client=NoopClient(), |
| 451 | source=DTensorModuleSource(Tiny(), num_trainer_ranks=1), |
| 452 | ) |
| 453 | engine.send_weights() |
| 454 | torch.accelerator.synchronize() |
| 455 | engine.shutdown() |
| 456 | return True |
| 457 | |
| 458 | |
| 459 | @ray.remote(num_gpus=1) |
| 460 | def _m2n_worker_receive( |
| 461 | nccl_unique_id_b64: str, |
| 462 | world_size: int, |
| 463 | worker_rank: int = 0, |
| 464 | ) -> dict: |
| 465 | """Receive that parameter through the real worker engine.""" |
| 466 | import contextlib |
| 467 | from unittest.mock import MagicMock |
| 468 | |
| 469 | device = _assigned_device() |
| 470 | |
| 471 | from vllm.config.parallel import ParallelConfig |
| 472 | from vllm.config.weight_transfer import WeightTransferConfig |
| 473 | from vllm.distributed.weight_transfer.m2n_engine import ( |
| 474 | M2NWeightTransferEngine, |
| 475 | M2NWeightTransferInitInfo, |
| 476 | M2NWeightTransferUpdateInfo, |
| 477 | ) |
| 478 | |
| 479 | class Recorder(torch.nn.Module): |
| 480 | def __init__(self): |
| 481 | super().__init__() |
| 482 | self.received = [] |
| 483 | if world_size > 2: |
| 484 | self.weight = torch.nn.Parameter( |
| 485 | torch.zeros( |
| 486 | SHAPE[0] // (world_size - 1), |
| 487 | SHAPE[1], |
| 488 | dtype=torch.float32, |
| 489 | device=device, |
| 490 | ) |
| 491 | ) |
| 492 | self.weight.output_dim = 0 |
| 493 | self.weight.weight_loader = MethodType( |
| 494 | ColumnParallelLinear.weight_loader, |
| 495 | Mock(tp_size=world_size - 1), |
| 496 | ) |
| 497 | |
| 498 | def load_weights(self, weights): |
| 499 | for name, tensor in weights: |
| 500 | self.received.append((name, tensor.clone())) |
| 501 | |
| 502 | parallel_config = MagicMock(spec=ParallelConfig) |
| 503 | parallel_config.rank = worker_rank |
| 504 | parallel_config.world_size = world_size - 1 |
| 505 | parallel_config.data_parallel_rank = 0 |
| 506 | parallel_config.data_parallel_index = 0 |
| 507 | parallel_config.tensor_parallel_size = world_size - 1 |
| 508 | parallel_config.pipeline_parallel_size = 1 |
| 509 | vllm_config = MagicMock() |
| 510 | vllm_config.parallel_config = parallel_config |
| 511 | vllm_config.model_config = MagicMock() |
| 512 | vllm_config.model_config.quantization = None |
| 513 | |
| 514 | recorder = Recorder() |
| 515 | engine = M2NWeightTransferEngine( |
| 516 | WeightTransferConfig(backend="nccl_m2n"), vllm_config, device, recorder |
| 517 | ) |
| 518 | # Transport-only: receive_weights enters set_current_vllm_config, and |
| 519 | # vllm_config here is a mock. |
| 520 | import vllm.config as _vllm_config_mod |
| 521 | |
| 522 | _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext() |
| 523 | |
| 524 | engine.init_transfer_engine( |
| 525 | M2NWeightTransferInitInfo( |
| 526 | nccl_unique_id_b64=nccl_unique_id_b64, |
| 527 | rank_offset=1, # trainer occupies rank 0 |
| 528 | world_size=world_size, |
| 529 | src_mesh_dims=[1, 1], |
| 530 | dst_mesh_dims=[1, world_size - 1], |
| 531 | names=["weight"], |
| 532 | dtype_names=[DTYPE], |
| 533 | shapes=[SHAPE], |
| 534 | src_placements=[None], # replicated on the single trainer rank |
| 535 | ) |
| 536 | ) |
| 537 | engine.receive_weights(M2NWeightTransferUpdateInfo(names=["weight"])) |
| 538 | torch.accelerator.synchronize() |
| 539 | |
| 540 | full = torch.arange( |
| 541 | SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device |
| 542 | ).reshape(SHAPE) |
| 543 | direct = world_size > 2 |
| 544 | name: str | None |
| 545 | got: torch.Tensor | None |
| 546 | if direct: |
| 547 | name, got = "weight", recorder.weight |
| 548 | expected = full.chunk(world_size - 1, dim=0)[worker_rank] |
| 549 | else: |
| 550 | name, got = recorder.received[0] if recorder.received else (None, None) |
| 551 | expected = full |
| 552 | result = { |
| 553 | "count": len(recorder.received), |
| 554 | "name": name, |
| 555 | "shape": list(got.shape) if got is not None else None, |
| 556 | "exact": bool(torch.equal(got, expected)) if got is not None else False, |
| 557 | "direct": direct, |
| 558 | } |
| 559 | engine.shutdown() |
| 560 | return result |
| 561 | |
| 562 | |
| 563 | @pytest.mark.skipif( |
| 564 | torch.accelerator.device_count() < 2, |
| 565 | reason="Need at least 2 GPUs: one trainer rank and one inference worker.", |
| 566 | ) |
| 567 | def test_m2n_weight_transfer_between_processes(): |
| 568 | """A parameter survives a real reshard from a trainer process to a worker. |
| 569 | |
| 570 | This is the only test here that moves bytes: it builds both engines, joins |
| 571 | one NCCL communicator across two processes, and reshards. Everything else |
| 572 | in this file checks how a transfer is *described*. |
| 573 | """ |
| 574 | pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") |
| 575 | _init_ray() |
| 576 | |
| 577 | from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary |
| 578 | |
| 579 | nccl = NCCLLibrary() |
| 580 | nccl_unique_id_b64 = base64.b64encode( |
| 581 | bytes(nccl.ncclGetUniqueId().internal) |
| 582 | ).decode() |
| 583 | world_size = 2 # 1 trainer + 1 inference worker |
| 584 | |
| 585 | worker = _m2n_worker_receive.remote(nccl_unique_id_b64, world_size) |
| 586 | trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size) |
| 587 | trainer_ok, result = ray.get([trainer, worker]) |
| 588 | |
| 589 | assert trainer_ok, "trainer engine did not complete" |
| 590 | assert result["count"] == 1, f"expected one parameter, got {result['count']}" |
| 591 | assert result["name"] == "weight" |
| 592 | assert result["shape"] == SHAPE |
| 593 | assert result["exact"], "resharded tensor does not match what the trainer sent" |
| 594 | |
| 595 | |
| 596 | @pytest.mark.skipif( |
| 597 | torch.accelerator.device_count() < 3, |
| 598 | reason="Need at least 3 GPUs: one trainer rank and two inference workers.", |
| 599 | ) |
| 600 | def test_m2n_weight_transfer_to_tp2_shards(): |
| 601 | """Two inference workers receive their own TP shard in live model storage.""" |
| 602 | pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") |
| 603 | _init_ray() |
| 604 | |
| 605 | from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary |
| 606 | |
| 607 | nccl = NCCLLibrary() |
| 608 | nccl_unique_id_b64 = base64.b64encode( |
| 609 | bytes(nccl.ncclGetUniqueId().internal) |
| 610 | ).decode() |
| 611 | world_size = 3 |
| 612 | |
| 613 | workers = [ |
| 614 | _m2n_worker_receive.remote(nccl_unique_id_b64, world_size, rank) |
| 615 | for rank in range(2) |
| 616 | ] |
| 617 | trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size) |
| 618 | trainer_ok, *results = ray.get([trainer, *workers]) |
| 619 | |
| 620 | assert trainer_ok, "trainer engine did not complete" |
| 621 | for rank, result in enumerate(results): |
| 622 | assert result["count"] == 0, f"worker {rank} unexpectedly called load_weights" |
| 623 | assert result["name"] == "weight" |
| 624 | assert result["shape"] == [SHAPE[0] // 2, SHAPE[1]] |
| 625 | assert result["direct"] |
| 626 | assert result["exact"], f"worker {rank} received the wrong shard" |
| 627 | |
| 628 | |
| 629 | class _Model(torch.nn.Module): |
| 630 | """Stand-in for a TP=2 shard of a tiny model. |
| 631 | |
| 632 | `column` is sharded on dim 0 and `row` on dim 1 (the two shapes a TP linear |
| 633 | layer produces), `norm` is replicated, and `fused` has no checkpoint-name |
| 634 | counterpart — the shape a fused parameter presents to the resolver. |
| 635 | """ |
| 636 | |
| 637 | def __init__(self) -> None: |
| 638 | super().__init__() |
| 639 | self.column = torch.nn.Parameter(torch.zeros(8, 16)) |
| 640 | self.column.output_dim = 0 |
| 641 | self.column.weight_loader = MethodType( |
| 642 | ColumnParallelLinear.weight_loader, Mock(tp_size=2) |
| 643 | ) |
| 644 | self.row = torch.nn.Parameter(torch.zeros(16, 8)) |
| 645 | self.row.input_dim = 1 |
| 646 | self.row.weight_loader = MethodType( |
| 647 | RowParallelLinear.weight_loader, Mock(tp_size=2) |
| 648 | ) |
| 649 | self.norm = torch.nn.Parameter(torch.zeros(16)) |
| 650 | self.norm.weight_loader = default_weight_loader |
| 651 | self.fused = torch.nn.Parameter(torch.zeros(24, 16)) |
| 652 | |
| 653 | |
| 654 | class _MixedModel(torch.nn.Module): |
| 655 | def __init__(self) -> None: |
| 656 | super().__init__() |
| 657 | self.proj = torch.nn.Module() |
| 658 | self.proj.weight = torch.nn.Parameter(torch.zeros(16, 8)) |
| 659 | self.proj.weight.input_dim = 1 |
| 660 | self.proj.weight.weight_loader = MethodType( |
| 661 | RowParallelLinear.weight_loader, Mock(tp_size=2) |
| 662 | ) |
| 663 | self.proj.bias = torch.nn.Parameter(torch.zeros(16)) |
| 664 | |
| 665 | |
| 666 | def _resolve(names, dtypes, shapes, **kwargs): |
| 667 | defaults = dict(num_workers=2, shard_axis_size=2, allow_direct=True) |
| 668 | defaults.update(kwargs) |
| 669 | return resolve_parameter_destinations(_Model(), names, dtypes, shapes, **defaults) |
| 670 | |
| 671 | |
| 672 | class TestDestinationResolution: |
| 673 | @pytest.mark.parametrize( |
| 674 | "name,expected_dim", |
| 675 | [("column", 0), ("row", 1)], |
| 676 | ) |
| 677 | def test_sharded_parameter_resolves_to_its_tp_dim(self, name, expected_dim): |
| 678 | [destination] = _resolve([name], [torch.float32], [(16, 16)]) |
| 679 | assert destination.direct |
| 680 | assert destination.placements == (REPLICATE, expected_dim) |
| 681 | |
| 682 | def test_replicated_parameter_needs_no_placement(self): |
| 683 | """A parameter every rank holds in full is REPLICATED, so it works |
| 684 | whatever way the inference mesh happens to be factored.""" |
| 685 | [destination] = _resolve(["norm"], [torch.float32], [(16,)]) |
| 686 | assert destination.direct |
| 687 | assert destination.placements is REPLICATED |
| 688 | |
| 689 | def test_direct_destination_is_the_live_parameter(self): |
| 690 | model = _Model() |
| 691 | [destination] = resolve_parameter_destinations( |
| 692 | model, |
| 693 | ["column"], |
| 694 | [torch.float32], |
| 695 | [(16, 16)], |
| 696 | num_workers=2, |
| 697 | shard_axis_size=2, |
| 698 | allow_direct=True, |
| 699 | ) |
| 700 | assert destination.tensor.data_ptr() == model.column.data_ptr() |
| 701 | |
| 702 | def test_known_transforming_model_submodule_falls_back(self): |
| 703 | class GPT2Model(_Model): |
| 704 | pass |
| 705 | |
| 706 | GPT2Model.__module__ = "vllm.model_executor.models.gpt2" |
| 707 | model = torch.nn.Module() |
| 708 | model.transforming = GPT2Model() |
| 709 | [destination] = resolve_parameter_destinations( |
| 710 | model, |
| 711 | ["transforming.row"], |
| 712 | [torch.float32], |
| 713 | [(16, 16)], |
| 714 | num_workers=2, |
| 715 | shard_axis_size=2, |
| 716 | allow_direct=True, |
| 717 | ) |
| 718 | |
| 719 | assert not destination.direct |
| 720 | assert destination.placements is REPLICATED |
| 721 | |
| 722 | def test_unknown_name_falls_back(self): |
| 723 | """Fused parameters reach the worker under checkpoint names that do not |
| 724 | exist in the model; those must take the full-tensor path.""" |
| 725 | [destination] = _resolve(["mlp.gate_proj.weight"], [torch.float32], [(16, 16)]) |
| 726 | assert not destination.direct |
| 727 | assert destination.placements is REPLICATED |
| 728 | |
| 729 | def test_fallback_parameter_demotes_direct_sibling_in_same_module(self): |
| 730 | """Layerwise finalization must not overwrite a directly loaded weight |
| 731 | when a sibling bias takes the fallback path.""" |
| 732 | destinations = resolve_parameter_destinations( |
| 733 | _MixedModel(), |
| 734 | ["proj.weight", "proj.bias"], |
| 735 | [torch.float32, torch.float32], |
| 736 | [(16, 16), (16,)], |
| 737 | num_workers=2, |
| 738 | shard_axis_size=2, |
| 739 | allow_direct=True, |
| 740 | ) |
| 741 | |
| 742 | assert not any(destination.direct for destination in destinations) |
| 743 | |
| 744 | def test_unknown_loader_falls_back_despite_matching_name_and_shape(self): |
| 745 | model = _Model() |
| 746 | model.norm.weight_loader = lambda param, weight: param.data.copy_(weight) |
| 747 | [destination] = resolve_parameter_destinations( |
| 748 | model, |
| 749 | ["norm"], |
| 750 | [torch.float32], |
| 751 | [(16,)], |
| 752 | num_workers=2, |
| 753 | shard_axis_size=2, |
| 754 | allow_direct=True, |
| 755 | ) |
| 756 | assert not destination.direct |
| 757 | |
| 758 | def test_missing_loader_is_not_assumed_to_be_a_plain_copy(self): |
| 759 | model = _Model() |
| 760 | del model.norm.weight_loader |
| 761 | [destination] = resolve_parameter_destinations( |
| 762 | model, |
| 763 | ["norm"], |
| 764 | [torch.float32], |
| 765 | [(16,)], |
| 766 | num_workers=2, |
| 767 | shard_axis_size=2, |
| 768 | allow_direct=True, |
| 769 | ) |
| 770 | assert not destination.direct |
| 771 | |
| 772 | def test_wrapped_known_loader_falls_back(self): |
| 773 | model = _Model() |
| 774 | |
| 775 | @wraps(default_weight_loader) |
| 776 | def wrapped_loader(param, weight): |
| 777 | default_weight_loader(param, weight) |
| 778 | |
| 779 | model.norm.weight_loader = wrapped_loader |
| 780 | [destination] = resolve_parameter_destinations( |
| 781 | model, |
| 782 | ["norm"], |
| 783 | [torch.float32], |
| 784 | [(16,)], |
| 785 | num_workers=2, |
| 786 | shard_axis_size=2, |
| 787 | allow_direct=True, |
| 788 | ) |
| 789 | assert not destination.direct |
| 790 | |
| 791 | def test_replicated_loader_with_transform_metadata_falls_back(self): |
| 792 | model = _Model() |
| 793 | model.norm.is_transposed = True |
| 794 | [destination] = resolve_parameter_destinations( |
| 795 | model, |
| 796 | ["norm"], |
| 797 | [torch.float32], |
| 798 | [(16,)], |
| 799 | num_workers=2, |
| 800 | shard_axis_size=2, |
| 801 | allow_direct=True, |
| 802 | ) |
| 803 | assert not destination.direct |
| 804 | |
| 805 | def test_packed_dim_zero_falls_back(self): |
| 806 | model = _Model() |
| 807 | model.column.packed_dim = 0 |
| 808 | [destination] = resolve_parameter_destinations( |
| 809 | model, |
| 810 | ["column"], |
| 811 | [torch.float32], |
| 812 | [(16, 16)], |
| 813 | num_workers=2, |
| 814 | shard_axis_size=2, |
| 815 | allow_direct=True, |
| 816 | ) |
| 817 | assert not destination.direct |
| 818 | |
| 819 | def test_loader_declared_dim_wins_over_shape_guessing(self): |
| 820 | model = _Model() |
| 821 | model.column.output_dim = 1 |
| 822 | [destination] = resolve_parameter_destinations( |
| 823 | model, |
| 824 | ["column"], |
| 825 | [torch.float32], |
| 826 | [(16, 16)], |
| 827 | num_workers=2, |
| 828 | shard_axis_size=2, |
| 829 | allow_direct=True, |
| 830 | ) |
| 831 | assert not destination.direct |
| 832 | |
| 833 | def test_shape_the_tp_factor_cannot_explain_falls_back(self): |
| 834 | [destination] = _resolve(["fused"], [torch.float32], [(16, 16)]) |
| 835 | assert not destination.direct |
| 836 | |
| 837 | def test_dtype_mismatch_falls_back(self): |
| 838 | """A parameter stored in a different dtype than the wire dtype means a |
| 839 | quantized or otherwise transformed layout — never write into it.""" |
| 840 | [destination] = _resolve(["column"], [torch.bfloat16], [(16, 16)]) |
| 841 | assert not destination.direct |
| 842 | |
| 843 | def test_allow_direct_false_forces_every_parameter_to_fall_back(self): |
| 844 | destinations = _resolve( |
| 845 | ["column", "norm"], |
| 846 | [torch.float32] * 2, |
| 847 | [(16, 16), (16,)], |
| 848 | allow_direct=False, |
| 849 | ) |
| 850 | assert not any(d.direct for d in destinations) |
| 851 | |
| 852 | def test_plan_reports_byte_coverage(self, caplog_vllm): |
| 853 | """A small direct parameter must not hide a large fallback tensor.""" |
| 854 | with caplog_vllm.at_level( |
| 855 | logging.INFO, |
| 856 | logger="vllm.distributed.weight_transfer.m2n_layout", |
| 857 | ): |
| 858 | _resolve( |
| 859 | ["norm", "missing_fused_weight"], |
| 860 | [torch.float32, torch.bfloat16], |
| 861 | [(16,), (1024, 1024)], |
| 862 | ) |
| 863 | |
| 864 | assert ( |
| 865 | "1/2 parameters resharded directly into the model, 1 via " |
| 866 | "full-tensor fallback; direct byte coverage: 64/2097216 bytes " |
| 867 | "(0.0%)" in caplog_vllm.text |
| 868 | ) |