From 4b8fb9903097bd67521831c76b49de89a9af4497 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E6=B5=B7?= <101489263+0KEAHA@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:33:04 +0800 Subject: [PATCH] Fix sender device alignment in vLLM weight synchronization --- swift/rlhf_trainers/vllm_client.py | 67 +++++++++----- tests/train/test_vllm_weight_sync.py | 127 +++++++++++++++++++++++++++ 2 files changed, 171 insertions(+), 23 deletions(-) create mode 100644 tests/train/test_vllm_weight_sync.py diff --git a/swift/rlhf_trainers/vllm_client.py b/swift/rlhf_trainers/vllm_client.py index ef585c676e..67b1ec8fe5 100644 --- a/swift/rlhf_trainers/vllm_client.py +++ b/swift/rlhf_trainers/vllm_client.py @@ -10,14 +10,14 @@ from pydantic import ValidationError from requests import ConnectionError from torch import nn -from typing import List, Optional, Union +from typing import Iterable, List, Optional, Union from urllib.parse import urlparse from swift.infer_engine import AdapterRequest, RequestConfig from swift.infer_engine.protocol import ChatCompletionResponse, RolloutInferRequest, RolloutOutput from swift.metrics import Metric -from swift.utils import (is_trl_available, is_vllm_ascend_available, is_vllm_available, is_vllm_metax_available, - synchronize) +from swift.utils import (get_torch_device, is_trl_available, is_vllm_ascend_available, is_vllm_available, + is_vllm_metax_available, synchronize) from .utils import (broadcast_tensor_for_vllm_weight_sync, format_host_for_url, is_valid_ipv6_address, peft_config_to_dict, resolve_hostname) @@ -38,6 +38,34 @@ logger = logging.getLogger(__name__) +def _broadcast_tensors_for_vllm_weight_sync(communicator, tensors: Iterable[torch.Tensor]) -> None: + """Broadcast outgoing tensors on the device owned by ``communicator``.""" + tensors = list(tensors) + if not tensors: + return + + # Exported tensors may originate from a different model-parallel device. + # Wait for each source device before starting a blocking cross-device copy. + source_devices = dict.fromkeys(tensor.device for tensor in tensors) + for source_device in source_devices: + if source_device.type != 'cpu': + synchronize(source_device) + + prepared_tensors = [ + tensor if tensor.device == communicator.device else tensor.to(device=communicator.device, non_blocking=False) + for tensor in tensors + ] + + # VLLM's communicator and its stream must use the same device. This is + # especially important in ThreadPoolExecutor workers, whose current device + # is not inherited from the caller thread. + device_module = get_torch_device() + with device_module.device(communicator.device): + for tensor in prepared_tensors: + broadcast_tensor_for_vllm_weight_sync(communicator, tensor, src=communicator.rank) + synchronize(communicator.device) + + class VLLMInferClient: """Inference-only vLLM client. Posts to /infer/ endpoint. No weight synchronization. Used for GKD teacher server etc. @@ -247,10 +275,9 @@ def _update_single_server(i): if response.status_code != 200: raise Exception(f'Server {i} update failed: {response.text}') - synchronize() - broadcast_tensor_for_vllm_weight_sync(self.pynccl_comms[i], weights, src=self.pynccl_comms[i].rank) - synchronize() - self.pynccl_comms[i].group.barrier() + comm = self.pynccl_comms[i] + _broadcast_tensors_for_vllm_weight_sync(comm, [weights]) + comm.group.barrier() except Exception as e: errors[i] = e @@ -292,11 +319,9 @@ def _update_single_server(i): if response.status_code != 200: raise Exception(f'Server {i} update adapter failed: {response.text}') - synchronize() - broadcast_tensor_for_vllm_weight_sync( - self.pynccl_comms[i], flattened_tensor, src=self.pynccl_comms[i].rank) - synchronize() - self.pynccl_comms[i].group.barrier() + comm = self.pynccl_comms[i] + _broadcast_tensors_for_vllm_weight_sync(comm, [flattened_tensor]) + comm.group.barrier() except Exception as e: errors[i] = e @@ -350,12 +375,10 @@ def _update_single_server(i): if response.status_code != 200: raise Exception(f'Server {i} update adapter failed: {response.text}') - # Broadcast each tensor individually - synchronize() - for name, param in lora_params.items(): - broadcast_tensor_for_vllm_weight_sync(self.pynccl_comms[i], param, src=self.pynccl_comms[i].rank) - synchronize() - self.pynccl_comms[i].group.barrier() + # Broadcast each tensor individually. + comm = self.pynccl_comms[i] + _broadcast_tensors_for_vllm_weight_sync(comm, lora_params.values()) + comm.group.barrier() except Exception as e: errors[i] = e @@ -392,11 +415,9 @@ def _update_single_server(i): if response.status_code != 200: raise Exception(f'Server {i} update flattened params failed: {response.text}') - synchronize() - broadcast_tensor_for_vllm_weight_sync( - self.pynccl_comms[i], flattened_tensor, src=self.pynccl_comms[i].rank) - synchronize() - self.pynccl_comms[i].group.barrier() + comm = self.pynccl_comms[i] + _broadcast_tensors_for_vllm_weight_sync(comm, [flattened_tensor]) + comm.group.barrier() except Exception as e: errors[i] = e diff --git a/tests/train/test_vllm_weight_sync.py b/tests/train/test_vllm_weight_sync.py new file mode 100644 index 0000000000..9502b2548b --- /dev/null +++ b/tests/train/test_vllm_weight_sync.py @@ -0,0 +1,127 @@ +import torch +from contextlib import contextmanager +from unittest.mock import patch + +from swift.rlhf_trainers.vllm_client import _broadcast_tensors_for_vllm_weight_sync + + +class _FakeTensor: + + def __init__(self, device): + self.device = device + self.to_calls = [] + + def to(self, *, device, non_blocking): + self.to_calls.append((device, non_blocking)) + return _FakeTensor(device) + + +class _FakeCommunicator: + + def __init__(self, device=torch.device('cuda:1')): + self.device = device + self.rank = 2 + self.events = None + + def broadcast(self, tensor, src, stream): + self.events.append(('broadcast', tensor.device, src, stream.device)) + + +class _FakeStream: + + def __init__(self, device): + self.device = device + + +class _FakeDeviceModule: + + def __init__(self, events): + self.events = events + # ThreadPoolExecutor workers can start with a different current device. + self.current_device = torch.device('cuda:0') + + @contextmanager + def device(self, device): + previous_device = self.current_device + self.current_device = device + self.events.append(('enter_device', device)) + try: + yield + finally: + self.events.append(('exit_device', device)) + self.current_device = previous_device + + def current_stream(self): + self.events.append(('current_stream', self.current_device)) + return _FakeStream(self.current_device) + + +def _run_broadcast(tensors, communicator): + events = [] + device_module = _FakeDeviceModule(events) + communicator.events = events + + def fake_synchronize(device): + events.append(('synchronize', device)) + + with patch( + 'swift.rlhf_trainers.vllm_client.get_torch_device', return_value=device_module), patch( + 'swift.rlhf_trainers.vllm_client.synchronize', side_effect=fake_synchronize), patch( + 'swift.rlhf_trainers.utils.get_torch_device', return_value=device_module), patch( + 'swift.rlhf_trainers.utils.is_torch_npu_available', return_value=False): + _broadcast_tensors_for_vllm_weight_sync(communicator, tensors) + + return events + + +def test_weight_sync_same_device_has_no_copy(): + communicator = _FakeCommunicator() + tensor = _FakeTensor(torch.device('cuda:1')) + + events = _run_broadcast([tensor], communicator) + + assert tensor.to_calls == [] + assert ('current_stream', torch.device('cuda:1')) in events + assert ('broadcast', torch.device('cuda:1'), communicator.rank, torch.device('cuda:1')) in events + assert events.count(('synchronize', torch.device('cuda:1'))) == 2 + + +def test_weight_sync_aligns_mismatched_tensor_to_communicator_device(): + communicator = _FakeCommunicator() + tensor = _FakeTensor(torch.device('cuda:0')) + + events = _run_broadcast([tensor], communicator) + + assert tensor.to_calls == [(torch.device('cuda:1'), False)] + assert ('synchronize', torch.device('cuda:0')) in events + assert ('current_stream', torch.device('cuda:1')) in events + assert ('broadcast', torch.device('cuda:1'), communicator.rank, torch.device('cuda:1')) in events + assert events[-2:] == [('synchronize', torch.device('cuda:1')), ('exit_device', torch.device('cuda:1'))] + + +def test_weight_sync_copies_cpu_tensor_without_accelerator_synchronize_on_cpu(): + communicator = _FakeCommunicator() + tensor = _FakeTensor(torch.device('cpu')) + + events = _run_broadcast([tensor], communicator) + + assert tensor.to_calls == [(torch.device('cuda:1'), False)] + assert ('synchronize', torch.device('cpu')) not in events + assert ('current_stream', torch.device('cuda:1')) in events + assert ('broadcast', torch.device('cuda:1'), communicator.rank, torch.device('cuda:1')) in events + + +def test_weight_sync_synchronizes_each_source_device_once(): + communicator = _FakeCommunicator() + tensors = [ + _FakeTensor(torch.device('cuda:0')), + _FakeTensor(torch.device('cuda:0')), + _FakeTensor(torch.device('cuda:1')), + ] + + events = _run_broadcast(tensors, communicator) + + first_broadcast = next(i for i, event in enumerate(events) if event[0] == 'broadcast') + assert events[:2] == [('synchronize', torch.device('cuda:0')), ('synchronize', torch.device('cuda:1'))] + assert all(event[1] == torch.device('cuda:1') and event[3] == torch.device('cuda:1') + for event in events[first_broadcast:] if event[0] == 'broadcast')