Fix sender device alignment in vLLM weight synchronization

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_kunlun_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_kunlun_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)
 
@@ -41,6 +41,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.
@@ -250,10 +278,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
 
@@ -295,11 +322,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
 
@@ -353,12 +378,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
 
@@ -395,11 +418,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
 

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')