@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 202 lines 7.65 KB
Raw
1 #include "torch_utils.h"
2
3 #include <torch/csrc/stable/macros.h>
4 #include <torch/csrc/stable/accelerator.h>
5 #include <torch/csrc/stable/tensor.h>
6 #include <torch/csrc/stable/ops.h>
7 #include <torch/headeronly/core/ScalarType.h>
8 #include <torch/csrc/stable/device.h>
9
10 #include "custom_all_reduce.cuh"
11
12 // Fake pointer type, must match fptr_t type in ops.h.
13 // We use this type alias to indicate when pointers are passed in as int64_t.
14 using fptr_t = int64_t;
15 static_assert(sizeof(void*) == sizeof(fptr_t));
16
17 fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
18 torch::stable::Tensor& rank_data, int64_t rank,
19 bool fully_connected) {
20 int world_size = fake_ipc_ptrs.size();
21 if (world_size > vllm::kMaxCustomCollectiveRanks)
22 throw std::invalid_argument("world size > 16 is not supported");
23 if (world_size % 2 != 0)
24 throw std::invalid_argument("Odd num gpus is not supported for now");
25 if (rank < 0 || rank >= world_size)
26 throw std::invalid_argument("invalid rank passed in");
27
28 vllm::Signal* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
29 for (int i = 0; i < world_size; i++) {
30 ipc_ptrs[i] = reinterpret_cast<vllm::Signal*>(fake_ipc_ptrs[i]);
31 }
32 return (fptr_t) new vllm::CustomAllreduce(
33 ipc_ptrs, rank_data.mutable_data_ptr(), rank_data.numel(), rank,
34 world_size, fully_connected);
35 }
36
37 /**
38 * Make sure tensor t's data lies completely within ((char)t.data_ptr()) +
39 * t.numel() * t.element_size(). This is slightly weaker than t.is_contiguous()
40 * because it allows transpose of contiguous slice (i.e. slicing the first
41 * dimension). Currently, we require this because stride information is not
42 * passed into the kernels and we treat input tensors as flat.
43 *
44 * Examples
45 * A = torch.zeros(3, 3, 3)
46 * 1. A: OK
47 * 2. A[1:]: OK
48 * 3. A.permute(2, 0, 1): OK
49 * 4. A[1:].permute(2, 0, 1): OK
50 * 5. A[None].expand(2, -1, -1, -1): Not OK
51 * 6. A[:, 1:, 1:]: Not OK
52 */
53 bool _is_weak_contiguous(torch::stable::Tensor& t) {
54 if (t.is_contiguous()) {
55 return true;
56 }
57 int64_t storage_nbytes = 0;
58 TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_size(t.get(), &storage_nbytes));
59 return storage_nbytes - t.storage_offset() * t.element_size() ==
60 static_cast<int64_t>(t.numel() * t.element_size());
61 }
62
63 /**
64 * Performs an out-of-place allreduce and stores result in out.
65 *
66 * If _reg_buffer is null, assumes inp.data_ptr() is already IPC-registered.
67 * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first
68 * copied into _reg_buffer.
69 */
70 void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
71 torch::stable::Tensor& out, fptr_t _reg_buffer,
72 int64_t reg_buffer_sz_bytes) {
73 auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
74 const torch::stable::accelerator::DeviceGuard device_guard(
75 inp.get_device_index());
76 const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
77
78 STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
79 STD_TORCH_CHECK((inp.numel()) == (out.numel()));
80 STD_TORCH_CHECK(_is_weak_contiguous(out));
81 STD_TORCH_CHECK(_is_weak_contiguous(inp));
82 auto input_size = inp.numel() * inp.element_size();
83 auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
84 if (reg_buffer) {
85 STD_TORCH_CHECK((input_size) <= (reg_buffer_sz_bytes));
86 STD_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.const_data_ptr(), input_size,
87 cudaMemcpyDeviceToDevice, stream));
88 } else {
89 reg_buffer = inp.mutable_data_ptr();
90 }
91 switch (out.scalar_type()) {
92 case torch::headeronly::ScalarType::Float: {
93 fa->allreduce<float>(stream, reinterpret_cast<float*>(reg_buffer),
94 reinterpret_cast<float*>(out.mutable_data_ptr()),
95 out.numel());
96 break;
97 }
98 case torch::headeronly::ScalarType::Half: {
99 fa->allreduce<half>(stream, reinterpret_cast<half*>(reg_buffer),
100 reinterpret_cast<half*>(out.mutable_data_ptr()),
101 out.numel());
102 break;
103 }
104 #if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
105 case torch::headeronly::ScalarType::BFloat16: {
106 fa->allreduce<nv_bfloat16>(
107 stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
108 reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), out.numel());
109 break;
110 }
111 #endif
112 default:
113 throw std::runtime_error(
114 "custom allreduce only supports float32, float16 and bfloat16");
115 }
116 }
117
118 void dispose(fptr_t _fa) {
119 delete reinterpret_cast<vllm::CustomAllreduce*>(_fa);
120 }
121
122 int64_t meta_size() { return sizeof(vllm::Signal); }
123
124 void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
125 auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
126 STD_TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
127 void* ipc_ptrs[vllm::kMaxCustomCollectiveRanks];
128 for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
129 ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
130 }
131 fa->register_buffer(ipc_ptrs);
132 }
133
134 // Use vector<int64_t> to represent byte data for python binding compatibility.
135 std::tuple<std::vector<int64_t>, std::vector<int64_t>>
136 get_graph_buffer_ipc_meta(fptr_t _fa) {
137 auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
138 auto [handle, offsets] = fa->get_graph_buffer_ipc_meta();
139 std::vector<int64_t> bytes(handle.begin(), handle.end());
140 return std::make_tuple(bytes, offsets);
141 }
142
143 // Use vector<int64_t> to represent byte data for python binding compatibility.
144 void register_graph_buffers(fptr_t _fa,
145 const std::vector<std::vector<int64_t>>& handles,
146 const std::vector<std::vector<int64_t>>& offsets) {
147 auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
148 std::vector<std::string> bytes;
149 bytes.reserve(handles.size());
150 for (int i = 0; i < handles.size(); i++) {
151 bytes.emplace_back(handles[i].begin(), handles[i].end());
152 }
153 bytes.reserve(handles.size());
154 fa->register_graph_buffers(bytes, offsets);
155 }
156
157 std::tuple<fptr_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
158 int64_t size) {
159 int device_index;
160 STD_CUDA_CHECK(cudaGetDevice(&device_index));
161 const torch::stable::accelerator::DeviceGuard device_guard(device_index);
162 void* buffer;
163 cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed;
164 const cudaStream_t stream = get_current_cuda_stream(device_index);
165 STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
166
167 // Allocate buffer
168 #if defined(USE_ROCM)
169 // data buffers need to be "uncached" for signal on MI200
170 STD_CUDA_CHECK(
171 hipExtMallocWithFlags((void**)&buffer, size, hipDeviceMallocUncached));
172 #else
173 STD_CUDA_CHECK(cudaMalloc((void**)&buffer, size));
174 #endif
175 STD_CUDA_CHECK(cudaMemsetAsync(buffer, 0, size, stream));
176 STD_CUDA_CHECK(cudaStreamSynchronize(stream));
177 STD_CUDA_CHECK(cudaThreadExchangeStreamCaptureMode(&mode));
178
179 // Create IPC memhandle for the allocated buffer.
180 // Will use it in open_mem_handle.
181 auto handle = torch::stable::empty(
182 {static_cast<int64_t>(sizeof(cudaIpcMemHandle_t))},
183 torch::headeronly::ScalarType::Byte, std::nullopt,
184 torch::stable::Device(torch::stable::DeviceType::CPU));
185 STD_CUDA_CHECK(cudaIpcGetMemHandle(
186 (cudaIpcMemHandle_t*)handle.mutable_data_ptr(), buffer));
187
188 return std::make_tuple(reinterpret_cast<fptr_t>(buffer), handle);
189 }
190
191 fptr_t open_mem_handle(torch::stable::Tensor& mem_handle) {
192 void* ipc_ptr;
193 STD_CUDA_CHECK(cudaIpcOpenMemHandle(
194 (void**)&ipc_ptr,
195 *((const cudaIpcMemHandle_t*)mem_handle.const_data_ptr()),
196 cudaIpcMemLazyEnablePeerAccess));
197 return reinterpret_cast<fptr_t>(ipc_ptr);
198 }
199
200 void free_shared_buffer(fptr_t buffer) {
201 STD_CUDA_CHECK(cudaFree(reinterpret_cast<void*>(buffer)));
202 }