diff --git a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py index 722c0d89fa..c168970147 100644 --- a/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py +++ b/tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py @@ -13,6 +13,8 @@ from unittest.mock import MagicMock, patch import pytest import torch +from vllm.config import ParallelConfig +from vllm.distributed.ec_transfer.ec_connector.cpu import common from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import ( ECSharedRegion, _wait_for_file_size, @@ -95,6 +97,32 @@ def test_second_instance_opens_existing_file_and_shares_memory(): r1.cleanup() +def test_dense_dp_engines_get_distinct_regions(monkeypatch): + """Dense DP engines reset data_parallel_rank to 0 but share instance_id, + so the region must be keyed by data_parallel_index.""" + monkeypatch.setattr(common, "_get_encoder_cache_hidden_dim", lambda cfg: 8) + instance_id = str(uuid.uuid4()) + regions = [] + try: + for dp_index in range(2): + parallel_config = ParallelConfig( + data_parallel_size=2, data_parallel_rank=dp_index + ) + parallel_config.reconfigure_for_independent_dp_rank() + cfg = MagicMock() + cfg.instance_id = instance_id + cfg.parallel_config = parallel_config + cfg.model_config.dtype = torch.float16 + cfg.ec_transfer_config.ec_connector_extra_config = {"ec_cpu_bytes": 64} + regions.append(common.create_ec_shared_region(cfg)) + + assert all(r._is_creator for r in regions) + assert regions[0]._mmap_path != regions[1]._mmap_path + finally: + for r in regions: + r.cleanup() + + def test_only_creator_unlinks_file_on_cleanup(): """Critical contract: if the non-creator unlinks, the creator's mmap path becomes a dangling backing file and a third opener would create a new one diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index e576ebd031..537653bda3 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -11,7 +11,7 @@ import pytest import torch import torch.nn as nn -from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config import ParallelConfig, VllmConfig, get_current_vllm_config from vllm.lora.layers import BaseLayerWithLoRA from vllm.v1.worker.gpu_model_runner import _get_parameter_for_reload from vllm.v1.worker.gpu_worker import Worker @@ -136,12 +136,22 @@ def test_rank_local_update_selects_worker_payload(rank, expected): assert worker._weight_update_active is True -def test_rank_local_update_includes_data_parallel_rank(): +def test_rank_local_update_uses_data_parallel_index_after_reconfigure(): engine = _RecordingEngine() worker = _make_worker(engine) worker.rank = 0 - worker.vllm_config.parallel_config.data_parallel_size = 4 - worker.vllm_config.parallel_config.data_parallel_rank = 2 + parallel_config = ParallelConfig( + data_parallel_size=4, + data_parallel_rank=2, + ) + assert parallel_config.data_parallel_rank == 2 + assert parallel_config.data_parallel_index == 2 + + parallel_config.reconfigure_for_independent_dp_rank() + assert parallel_config.data_parallel_rank == 0 + assert parallel_config.data_parallel_index == 2 + + worker.vllm_config.parallel_config = parallel_config Worker.start_weight_update(worker) Worker.update_weights( diff --git a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py index e2c37b72a6..eff16e08b0 100644 --- a/vllm/distributed/ec_transfer/ec_connector/cpu/common.py +++ b/vllm/distributed/ec_transfer/ec_connector/cpu/common.py @@ -99,8 +99,8 @@ def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion: ec_config = vllm_config.ec_transfer_config assert ec_config is not None, "ec_transfer_config required to build region" - dp_rank = vllm_config.parallel_config.data_parallel_rank - engine_id = f"{vllm_config.instance_id}_dp{dp_rank}" + dp_index = vllm_config.parallel_config.data_parallel_index + engine_id = f"{vllm_config.instance_id}_dp{dp_index}" dtype = vllm_config.model_config.dtype hidden_dim = _get_encoder_cache_hidden_dim(vllm_config) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index a288a97cc1..02f599e57e 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1449,7 +1449,7 @@ class Worker(WorkerBase): if isinstance(update_info, list): parallel_config = self.vllm_config.parallel_config local_update_info = update_info[ - parallel_config.data_parallel_rank * parallel_config.world_size + parallel_config.data_parallel_index * parallel_config.world_size + self.rank ] else: