@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 124 lines 5.38 KB
Raw
1 #include <ATen/cuda/Exceptions.h>
2 #include <c10/cuda/CUDAGuard.h>
3 #include <c10/cuda/CUDAStream.h>
4 #include <torch/all.h>
5
6 #ifdef USE_ROCM
7
8 #include "quickreduce/quick_reduce.h"
9
10 quickreduce::fptr_t init_custom_qr(int64_t rank, int64_t world_size,
11 std::optional<int64_t> qr_max_size) {
12 if (world_size > 8)
13 throw std::invalid_argument("world size > 8 is not supported");
14 if (world_size == 6)
15 throw std::invalid_argument("world size == 6 is not supported");
16 if (world_size % 2 != 0)
17 throw std::invalid_argument("Odd num gpus is not supported for now");
18 if (rank < 0 || rank >= world_size)
19 throw std::invalid_argument("invalid rank passed in");
20 quickreduce::DeviceComms* fptr = new quickreduce::DeviceComms();
21 fptr->init(world_size, rank, qr_max_size);
22 return (quickreduce::fptr_t)fptr;
23 }
24
25 void qr_destroy(quickreduce::fptr_t _fa) {
26 if (_fa) {
27 auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
28 fa->destroy();
29 delete fa;
30 }
31 }
32
33 torch::Tensor qr_get_handle(quickreduce::fptr_t _fa) {
34 auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
35 hipIpcMemHandle_t handle = fa->get_handle();
36 auto options =
37 torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
38 auto data_handle =
39 torch::empty({static_cast<int64_t>(sizeof(hipIpcMemHandle_t))}, options);
40 std::memcpy(data_handle.data_ptr(), &handle, sizeof(hipIpcMemHandle_t));
41 return data_handle;
42 }
43
44 void qr_open_handles(quickreduce::fptr_t _fa,
45 const std::vector<torch::Tensor>& handles) {
46 auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
47 std::vector<hipIpcMemHandle_t> ipc_handles;
48 ipc_handles.reserve(handles.size());
49 for (auto& handle : handles) {
50 // Ensure the tensor is on the same device as the current device.
51 hipIpcMemHandle_t ipc_handle;
52 std::memcpy(&ipc_handle, handle.data_ptr(), sizeof(hipIpcMemHandle_t));
53 ipc_handles.push_back(ipc_handle);
54 }
55 fa->open_ipc_handles(ipc_handles);
56 }
57
58 void qr_all_reduce(quickreduce::fptr_t _fa, torch::Tensor& inp,
59 torch::Tensor& out, int64_t quant_level, bool cast_bf2half) {
60 auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
61 const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
62 auto stream = at::cuda::getCurrentHIPStreamMasqueradingAsCUDA();
63
64 TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
65 TORCH_CHECK_EQ(inp.numel(), out.numel());
66 TORCH_CHECK_LE(out.numel(), fa->kMaxProblemSize);
67 if (out.scalar_type() == at::ScalarType::Half) {
68 fa->allreduce<half, false>(reinterpret_cast<half*>(inp.data_ptr()),
69 reinterpret_cast<half*>(out.data_ptr()),
70 out.numel(), quant_level, stream);
71 } else if (out.scalar_type() == at::ScalarType::BFloat16) {
72 if (cast_bf2half) {
73 fa->allreduce<half, true>(reinterpret_cast<half*>(inp.data_ptr()),
74 reinterpret_cast<half*>(out.data_ptr()),
75 out.numel(), quant_level, stream);
76 } else {
77 fa->allreduce<quickreduce::nv_bfloat16, false>(
78 reinterpret_cast<quickreduce::nv_bfloat16*>(inp.data_ptr()),
79 reinterpret_cast<quickreduce::nv_bfloat16*>(out.data_ptr()),
80 out.numel(), quant_level, stream);
81 }
82 } else {
83 throw std::runtime_error(
84 "quick allreduce only supports float16 and bfloat16");
85 }
86 }
87
88 int64_t qr_max_size() {
89 // The default is 2GB (2,147,483,648 bytes)
90 return static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
91 }
92
93 #define INSTANTIATE_FOR_WORLDSIZE(T, Codec, cast_bf2half) \
94 template struct quickreduce::AllReduceTwoshot<T, Codec<T, 2>, \
95 cast_bf2half>; \
96 template struct quickreduce::AllReduceTwoshot<T, Codec<T, 4>, \
97 cast_bf2half>; \
98 template struct quickreduce::AllReduceTwoshot<T, Codec<T, 8>, cast_bf2half>;
99
100 // INT3 (CodecQ3) is restricted to TP2 only, so we only instantiate the
101 // world_size == 2 kernel for it.
102 #define INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(T, Codec, cast_bf2half) \
103 template struct quickreduce::AllReduceTwoshot<T, Codec<T, 2>, cast_bf2half>;
104
105 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, false)
106 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, false)
107 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, false)
108 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, false)
109 INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
110 quickreduce::CodecQ3, false)
111 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, true)
112 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, true)
113 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, true)
114 INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, true)
115 INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16,
116 quickreduce::CodecQ3, true)
117
118 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecFP, false)
119 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ4, false)
120 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ6, false)
121 INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ8, false)
122 INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(half, quickreduce::CodecQ3, false)
123
124 #endif // USE_ROCM