Workaround PyTorch IPC handle issue (opendatahub-io#161)

prarit · Aug 28, 2024 · 056baed · 056baed
1 parent 07b6b14
commit 056baed
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py
@@ -264,11 +264,20 @@ def capture(self):
                 self.register_graph_buffers()
 
     def _get_ipc_meta(self, inp: torch.Tensor):
-        data = inp.untyped_storage()._share_cuda_()
-        shard_data = (
-            data[1],  # ipc handle to base ptr
-            data[3],  # offset of base ptr
-        )
+        if is_hip():
+            # _share_cuda_() doesn't accept meta buffer not allocated from
+            # PyTorch cache allocator, use direct HIP call to get IPC handle
+            handle = custom_ar.get_meta_buffer_ipc_handle(inp)
+            shard_data = (
+                bytes(handle),  # ipc handle to base ptr
+                0,  # offset of base ptr
+            )
+        else:
+            data = inp.untyped_storage()._share_cuda_()
+            shard_data = (
+                data[1],  # ipc handle to base ptr
+                data[3],  # offset of base ptr
+            )
         return self._gather_ipc_meta(shard_data)
 
     def _gather_ipc_meta(self, shard_data):