fix vllm>=0.14

Author: hjh0119

swift/pipelines/infer/rollout.py

@@ -143,7 +143,7 @@ def update_named_param(self, name: str, dtype: str, shape: Sequence[int]) -> Non
         weight = torch.empty(shape, dtype=dtype, device=self.communicator.device)
 
         # Use NCCL to broadcast the updated weights from the client (src) to all workers.
-        self.communicator.broadcast(weight, src=self.client_rank)
+        self.communicator.broadcast(weight, src=self.client_rank, stream=torch.cuda.current_stream())
         self.communicator.group.barrier()
 
         # Patch MoE weight_loader if needed
@@ -162,7 +162,7 @@ def update_adapter_flattened_param(self, lora_int_id: int, peft_config: Dict, me
         flatten_tensor_length = metadatas[-1].end_idx
         dtype = getattr(torch, metadatas[-1].dtype.split('.')[-1])
         flatten_tensor = torch.empty(flatten_tensor_length, dtype=dtype, device=self.communicator.device)
-        self.communicator.broadcast(flatten_tensor, src=self.client_rank)
+        self.communicator.broadcast(flatten_tensor, src=self.client_rank, stream=torch.cuda.current_stream())
         self.communicator.group.barrier()
         flattened_tensor_bucket = FlattenedTensorBucket(metadata=metadatas, flattened_tensor=flatten_tensor)
         named_params = flattened_tensor_bucket.reconstruct_tensors()
@@ -194,7 +194,7 @@ def update_adapter_param(self, lora_int_id: int, peft_config: Dict, lora_tensors
             dtype = getattr(torch, metadata['dtype'].split('.')[-1])
             shape = tuple(metadata['shape'])
             tensor = torch.empty(shape, dtype=dtype, device=self.communicator.device)
-            self.communicator.broadcast(tensor, src=self.client_rank)
+            self.communicator.broadcast(tensor, src=self.client_rank, stream=torch.cuda.current_stream())
             named_params[name] = tensor
 
         self.communicator.group.barrier()
@@ -222,7 +222,7 @@ def update_flattened_params(self, metadatas: list[Dict]) -> None:
         dtype = getattr(torch, metadatas[-1].dtype.split('.')[-1])
         flatten_tensor = torch.empty(flatten_tensor_length, dtype=dtype, device=self.communicator.device)
 
-        self.communicator.broadcast(flatten_tensor, src=self.client_rank)
+        self.communicator.broadcast(flatten_tensor, src=self.client_rank, stream=torch.cuda.current_stream())
         self.communicator.group.barrier()
 
         flattened_tensor_bucket = FlattenedTensorBucket(metadata=metadatas, flattened_tensor=flatten_tensor)

swift/rlhf_trainers/vllm_client.py

@@ -230,7 +230,10 @@ def _update_single_server(i):
                 if response.status_code != 200:
                     raise Exception(f'Server {i} update failed: {response.text}')
 
-                self.pynccl_comms[i].broadcast(weights, src=self.pynccl_comms[i].rank)
+                torch.cuda.synchronize()
+                self.pynccl_comms[i].broadcast(
+                    weights, src=self.pynccl_comms[i].rank, stream=torch.cuda.current_stream())
+                torch.cuda.synchronize()
                 self.pynccl_comms[i].group.barrier()
             except Exception as e:
                 errors[i] = e
@@ -275,7 +278,10 @@ def _update_single_server(i):
                 if response.status_code != 200:
                     raise Exception(f'Server {i} update adapter failed: {response.text}')
 
-                self.pynccl_comms[i].broadcast(flattened_tensor, src=self.pynccl_comms[i].rank)
+                torch.cuda.synchronize()
+                self.pynccl_comms[i].broadcast(
+                    flattened_tensor, src=self.pynccl_comms[i].rank, stream=torch.cuda.current_stream())
+                torch.cuda.synchronize()
                 self.pynccl_comms[i].group.barrier()
             except Exception as e:
                 errors[i] = e
@@ -333,8 +339,11 @@ def _update_single_server(i):
                     raise Exception(f'Server {i} update adapter failed: {response.text}')
 
                 # Broadcast each tensor individually
+                torch.cuda.synchronize()
                 for name, param in lora_params.items():
-                    self.pynccl_comms[i].broadcast(param, src=self.pynccl_comms[i].rank)
+                    self.pynccl_comms[i].broadcast(
+                        param, src=self.pynccl_comms[i].rank, stream=torch.cuda.current_stream())
+                torch.cuda.synchronize()
                 self.pynccl_comms[i].group.barrier()
             except Exception as e:
                 errors[i] = e
@@ -372,7 +381,10 @@ def _update_single_server(i):
                 if response.status_code != 200:
                     raise Exception(f'Server {i} update flattened params failed: {response.text}')
 
-                self.pynccl_comms[i].broadcast(flattened_tensor, src=self.pynccl_comms[i].rank)
+                torch.cuda.synchronize()
+                self.pynccl_comms[i].broadcast(
+                    flattened_tensor, src=self.pynccl_comms[i].rank, stream=torch.cuda.current_stream())
+                torch.cuda.synchronize()
                 self.pynccl_comms[i].group.barrier()
             except Exception as e:
                 errors[i] = e