| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # can only run on machines with p2p access across GPUs |
| 5 | # can only run with torchrun: |
| 6 | # torchrun --nproc_per_node=2 tests/distributed/test_ca_buffer_sharing.py |
| 7 | |
| 8 | import ctypes |
| 9 | |
| 10 | import torch |
| 11 | import torch.distributed as dist |
| 12 | |
| 13 | from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary |
| 14 | from vllm.distributed.device_communicators.custom_all_reduce import ( # noqa |
| 15 | CustomAllreduce, |
| 16 | ) |
| 17 | |
| 18 | # create a cpu process group for communicating metadata (ipc handle) |
| 19 | dist.init_process_group(backend="gloo") |
| 20 | rank = local_rank = dist.get_rank() |
| 21 | world_size = dist.get_world_size() |
| 22 | |
| 23 | # every process sets its own device (differently) |
| 24 | lib = CudaRTLibrary() |
| 25 | lib.cudaSetDevice(rank) |
| 26 | |
| 27 | buffer_size_in_bytes = 1024 |
| 28 | byte_value = 2 # the value we write to the buffer for verification |
| 29 | |
| 30 | pointers = CustomAllreduce.create_shared_buffer(buffer_size_in_bytes) |
| 31 | |
| 32 | print(f"Rank {rank} has pointers {pointers}") |
| 33 | |
| 34 | dist.barrier() |
| 35 | torch.accelerator.synchronize() |
| 36 | |
| 37 | if rank == 0: |
| 38 | # the first rank tries to write to all buffers |
| 39 | for p in pointers: |
| 40 | pointer = ctypes.c_void_p(p) |
| 41 | lib.cudaMemset(pointer, byte_value, buffer_size_in_bytes) |
| 42 | |
| 43 | dist.barrier() |
| 44 | torch.accelerator.synchronize() |
| 45 | |
| 46 | host_data = (ctypes.c_char * buffer_size_in_bytes)() |
| 47 | |
| 48 | # all ranks read from all buffers, and check if the data is correct |
| 49 | for p in pointers: |
| 50 | pointer = ctypes.c_void_p(p) |
| 51 | lib.cudaMemcpy(host_data, pointer, buffer_size_in_bytes) |
| 52 | for i in range(buffer_size_in_bytes): |
| 53 | assert ord(host_data[i]) == byte_value, ( |
| 54 | f"Rank {rank} failed" |
| 55 | f" to verify buffer {p}. Expected {byte_value}, " |
| 56 | f"got {ord(host_data[i])}" |
| 57 | ) |
| 58 | |
| 59 | print(f"Rank {rank} verified all buffers") |
| 60 | |
| 61 | dist.barrier() |
| 62 | torch.accelerator.synchronize() |
| 63 | |
| 64 | CustomAllreduce.free_shared_buffer(pointers) |