| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import json |
| 6 | import os |
| 7 | import statistics |
| 8 | from collections.abc import Callable |
| 9 | |
| 10 | import torch |
| 11 | import torch.distributed as dist |
| 12 | |
| 13 | import vllm._custom_ops as ops |
| 14 | from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce |
| 15 | |
| 16 | |
| 17 | def parse_args() -> argparse.Namespace: |
| 18 | parser = argparse.ArgumentParser() |
| 19 | parser.add_argument("--tokens", type=int, nargs="+", default=[8, 32, 128, 1024]) |
| 20 | parser.add_argument("--hidden-size", type=int, default=7168) |
| 21 | parser.add_argument("--graph-repeats", type=int, default=20) |
| 22 | parser.add_argument("--warmup-replays", type=int, default=5) |
| 23 | parser.add_argument("--samples", type=int, default=15) |
| 24 | return parser.parse_args() |
| 25 | |
| 26 | |
| 27 | def capture_graph(op: Callable[[], None], repeats: int) -> torch.cuda.CUDAGraph: |
| 28 | stream = torch.cuda.Stream() |
| 29 | stream.wait_stream(torch.cuda.current_stream()) |
| 30 | with torch.cuda.stream(stream): |
| 31 | for _ in range(3): |
| 32 | op() |
| 33 | stream.synchronize() |
| 34 | |
| 35 | graph = torch.cuda.CUDAGraph() |
| 36 | with torch.cuda.graph(graph, stream=stream): |
| 37 | for _ in range(repeats): |
| 38 | op() |
| 39 | torch.cuda.current_stream().wait_stream(stream) |
| 40 | return graph |
| 41 | |
| 42 | |
| 43 | def max_rank_graph_time( |
| 44 | graph: torch.cuda.CUDAGraph, |
| 45 | repeats: int, |
| 46 | warmup_replays: int, |
| 47 | samples: int, |
| 48 | device_group: dist.ProcessGroup, |
| 49 | cpu_group: dist.ProcessGroup, |
| 50 | ) -> float: |
| 51 | for _ in range(warmup_replays): |
| 52 | graph.replay() |
| 53 | torch.accelerator.synchronize() |
| 54 | |
| 55 | timings = [] |
| 56 | start = torch.cuda.Event(enable_timing=True) |
| 57 | end = torch.cuda.Event(enable_timing=True) |
| 58 | for _ in range(samples): |
| 59 | dist.barrier(group=cpu_group) |
| 60 | start.record() |
| 61 | graph.replay() |
| 62 | end.record() |
| 63 | end.synchronize() |
| 64 | elapsed = torch.tensor( |
| 65 | start.elapsed_time(end) / repeats, |
| 66 | dtype=torch.float64, |
| 67 | device=torch.accelerator.current_device_index(), |
| 68 | ) |
| 69 | dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group) |
| 70 | timings.append(elapsed.item()) |
| 71 | return statistics.median(timings) |
| 72 | |
| 73 | |
| 74 | def check_outputs( |
| 75 | comm: CustomAllreduce, |
| 76 | local: torch.Tensor, |
| 77 | reduce_input: torch.Tensor, |
| 78 | device_group: dist.ProcessGroup, |
| 79 | ) -> None: |
| 80 | expected_gather = torch.empty( |
| 81 | (local.shape[0] * dist.get_world_size(), local.shape[1]), |
| 82 | dtype=local.dtype, |
| 83 | device=local.device, |
| 84 | ) |
| 85 | dist.all_gather_into_tensor(expected_gather, local, group=device_group) |
| 86 | gathered = comm.custom_all_gather(local) |
| 87 | assert gathered is not None |
| 88 | torch.testing.assert_close(gathered, expected_gather) |
| 89 | |
| 90 | expected_scatter = torch.empty_like(local) |
| 91 | dist.reduce_scatter_tensor( |
| 92 | expected_scatter, |
| 93 | reduce_input.clone(), |
| 94 | group=device_group, |
| 95 | ) |
| 96 | scattered = comm.custom_reduce_scatter(reduce_input) |
| 97 | assert scattered is not None |
| 98 | torch.testing.assert_close(scattered, expected_scatter) |
| 99 | |
| 100 | |
| 101 | def benchmark_shape( |
| 102 | comm: CustomAllreduce, |
| 103 | global_tokens: int, |
| 104 | hidden_size: int, |
| 105 | graph_repeats: int, |
| 106 | warmup_replays: int, |
| 107 | samples: int, |
| 108 | device_group: dist.ProcessGroup, |
| 109 | cpu_group: dist.ProcessGroup, |
| 110 | ) -> dict[str, float | int]: |
| 111 | world_size = dist.get_world_size() |
| 112 | rank = dist.get_rank() |
| 113 | padded_tokens = (global_tokens + world_size - 1) // world_size * world_size |
| 114 | local_tokens = padded_tokens // world_size |
| 115 | local = torch.full( |
| 116 | (local_tokens, hidden_size), |
| 117 | rank + 1, |
| 118 | dtype=torch.bfloat16, |
| 119 | device=torch.accelerator.current_device_index(), |
| 120 | ) |
| 121 | reduce_input = torch.full( |
| 122 | (padded_tokens, hidden_size), |
| 123 | rank + 1, |
| 124 | dtype=torch.bfloat16, |
| 125 | device=local.device, |
| 126 | ) |
| 127 | check_outputs(comm, local, reduce_input, device_group) |
| 128 | |
| 129 | custom_gather_out = torch.empty( |
| 130 | (padded_tokens, hidden_size), |
| 131 | dtype=local.dtype, |
| 132 | device=local.device, |
| 133 | ) |
| 134 | custom_scatter_out = torch.empty_like(local) |
| 135 | nccl_gather_out = torch.empty_like(custom_gather_out) |
| 136 | nccl_scatter_out = torch.empty_like(local) |
| 137 | |
| 138 | def custom_ag() -> None: |
| 139 | ops.mnnvl_lamport_all_gather( |
| 140 | comm._ptr, |
| 141 | local, |
| 142 | custom_gather_out, |
| 143 | comm.mnnvl_lamport_ag_local_ptr, |
| 144 | comm.mnnvl_lamport_ag_multicast_ptr, |
| 145 | comm.mnnvl_lamport_ag_epoch_ptr, |
| 146 | comm.mnnvl_buffer_size, |
| 147 | ) |
| 148 | |
| 149 | def custom_rs() -> None: |
| 150 | ops.mnnvl_lamport_reduce_scatter( |
| 151 | comm._ptr, |
| 152 | reduce_input, |
| 153 | custom_scatter_out, |
| 154 | comm.mnnvl_lamport_rs_local_ptr, |
| 155 | comm.mnnvl_lamport_rs_epoch_ptr, |
| 156 | comm.mnnvl_buffer_size, |
| 157 | ) |
| 158 | |
| 159 | def nccl_ag() -> None: |
| 160 | dist.all_gather_into_tensor(nccl_gather_out, local, group=device_group) |
| 161 | |
| 162 | def nccl_rs() -> None: |
| 163 | dist.reduce_scatter_tensor( |
| 164 | nccl_scatter_out, |
| 165 | reduce_input, |
| 166 | group=device_group, |
| 167 | ) |
| 168 | |
| 169 | graphs = { |
| 170 | "custom_ag_us": capture_graph(custom_ag, graph_repeats), |
| 171 | "nccl_ag_us": capture_graph(nccl_ag, graph_repeats), |
| 172 | "custom_rs_us": capture_graph(custom_rs, graph_repeats), |
| 173 | "nccl_rs_us": capture_graph(nccl_rs, graph_repeats), |
| 174 | } |
| 175 | times = { |
| 176 | name: max_rank_graph_time( |
| 177 | graph, |
| 178 | graph_repeats, |
| 179 | warmup_replays, |
| 180 | samples, |
| 181 | device_group, |
| 182 | cpu_group, |
| 183 | ) |
| 184 | * 1000 |
| 185 | for name, graph in graphs.items() |
| 186 | } |
| 187 | torch.testing.assert_close(custom_gather_out, nccl_gather_out) |
| 188 | torch.testing.assert_close(custom_scatter_out, nccl_scatter_out) |
| 189 | return { |
| 190 | "global_tokens": global_tokens, |
| 191 | "padded_tokens": padded_tokens, |
| 192 | "local_bytes": local.nbytes, |
| 193 | "full_bytes": reduce_input.nbytes, |
| 194 | **times, |
| 195 | "ag_speedup": times["nccl_ag_us"] / times["custom_ag_us"], |
| 196 | "rs_speedup": times["nccl_rs_us"] / times["custom_rs_us"], |
| 197 | } |
| 198 | |
| 199 | |
| 200 | def main() -> None: |
| 201 | args = parse_args() |
| 202 | local_rank = int(os.environ["LOCAL_RANK"]) |
| 203 | torch.accelerator.set_device_index(local_rank) |
| 204 | dist.init_process_group("nccl") |
| 205 | device_group = dist.group.WORLD |
| 206 | cpu_group = dist.new_group(backend="gloo") |
| 207 | |
| 208 | comm = CustomAllreduce( |
| 209 | group=cpu_group, |
| 210 | device=torch.device("cuda", local_rank), |
| 211 | ) |
| 212 | assert not comm.disabled |
| 213 | assert comm.world_size == 16 |
| 214 | assert comm.mnnvl_only |
| 215 | assert comm.mnnvl_multicast_ptr |
| 216 | |
| 217 | results = [ |
| 218 | benchmark_shape( |
| 219 | comm, |
| 220 | tokens, |
| 221 | args.hidden_size, |
| 222 | args.graph_repeats, |
| 223 | args.warmup_replays, |
| 224 | args.samples, |
| 225 | device_group, |
| 226 | cpu_group, |
| 227 | ) |
| 228 | for tokens in args.tokens |
| 229 | ] |
| 230 | if dist.get_rank() == 0: |
| 231 | print(json.dumps(results, indent=2), flush=True) |
| 232 | |
| 233 | comm.close() |
| 234 | dist.destroy_process_group(cpu_group) |
| 235 | dist.destroy_process_group() |
| 236 | |
| 237 | |
| 238 | if __name__ == "__main__": |
| 239 | main() |