| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | |
| 5 | """Benchmark script for device communicators: |
| 6 | CustomAllreduce (oneshot, twoshot), PyNcclCommunicator, |
| 7 | and SymmMemCommunicator (multimem, two-shot). |
| 8 | |
| 9 | for NCCL symmetric memory you need to set the environment variables |
| 10 | NCCL_NVLS_ENABLE=1 NCCL_CUMEM_ENABLE=1 VLLM_USE_NCCL_SYMM_MEM=1, otherwise NCCL does |
| 11 | not use fast NVLS implementation for all reduce. |
| 12 | |
| 13 | Usage: |
| 14 | torchrun --nproc_per_node=<N> benchmark_device_communicators.py [options] |
| 15 | |
| 16 | Example: |
| 17 | torchrun --nproc_per_node=2 benchmark_device_communicators.py |
| 18 | --sequence-lengths 512 1024 2048 --num-warmup 10 --num-trials 100 |
| 19 | |
| 20 | """ |
| 21 | |
| 22 | import json |
| 23 | import os |
| 24 | import time |
| 25 | from collections.abc import Callable |
| 26 | from contextlib import nullcontext |
| 27 | |
| 28 | import torch |
| 29 | import torch.distributed as dist |
| 30 | from torch.distributed import ProcessGroup |
| 31 | |
| 32 | from vllm._aiter_ops import rocm_aiter_ops |
| 33 | from vllm.distributed.device_communicators.aiter_custom_all_reduce import ( |
| 34 | AiterCustomAllreduce, |
| 35 | ) |
| 36 | from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce |
| 37 | from vllm.distributed.device_communicators.flashinfer_all_reduce import ( |
| 38 | FlashInferAllReduce, |
| 39 | ) |
| 40 | from vllm.distributed.device_communicators.pynccl import ( |
| 41 | PyNcclCommunicator, |
| 42 | register_nccl_symmetric_ops, |
| 43 | ) |
| 44 | from vllm.distributed.device_communicators.pynccl_allocator import ( |
| 45 | set_graph_pool_id, |
| 46 | ) |
| 47 | from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator |
| 48 | from vllm.logger import init_logger |
| 49 | from vllm.platforms import current_platform |
| 50 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 51 | |
| 52 | logger = init_logger(__name__) |
| 53 | |
| 54 | # Default sequence lengths to benchmark |
| 55 | DEFAULT_SEQUENCE_LENGTHS = [16, 64, 128, 512, 1024, 2048, 4096, 8192] |
| 56 | |
| 57 | # Fixed hidden size and dtype for all benchmarks |
| 58 | HIDDEN_SIZE = 8192 |
| 59 | BENCHMARK_DTYPE = torch.bfloat16 |
| 60 | |
| 61 | # CUDA graph settings |
| 62 | CUDA_GRAPH_CAPTURE_CYCLES = 10 |
| 63 | |
| 64 | # All-gather and reduce-scatter are distinct collectives from all-reduce, so |
| 65 | # results are compared within a collective (each against its own NCCL/RCCL |
| 66 | # baseline) rather than across collectives. |
| 67 | COMM_COLLECTIVE = { |
| 68 | "ca_1stage": "all-reduce", |
| 69 | "ca_2stage": "all-reduce", |
| 70 | "pynccl": "all-reduce", |
| 71 | "pynccl-symm": "all-reduce", |
| 72 | "symm_mem_multimem": "all-reduce", |
| 73 | "symm_mem_two_shot": "all-reduce", |
| 74 | "flashinfer_trtllm": "all-reduce", |
| 75 | "flashinfer_mnnvl": "all-reduce", |
| 76 | "aiter_ag": "all-gather", |
| 77 | "pynccl_ag": "all-gather", |
| 78 | "aiter_rs": "reduce-scatter", |
| 79 | "pynccl_rs": "reduce-scatter", |
| 80 | } |
| 81 | COLLECTIVE_ORDER = ["all-reduce", "all-gather", "reduce-scatter"] |
| 82 | COLLECTIVE_BASELINE = { |
| 83 | "all-reduce": "pynccl", |
| 84 | "all-gather": "pynccl_ag", |
| 85 | "reduce-scatter": "pynccl_rs", |
| 86 | } |
| 87 | |
| 88 | |
| 89 | class CommunicatorBenchmark: |
| 90 | """Benchmark class for testing device communicators.""" |
| 91 | |
| 92 | def __init__( |
| 93 | self, |
| 94 | rank: int, |
| 95 | world_size: int, |
| 96 | device: torch.device, |
| 97 | cpu_group: ProcessGroup, |
| 98 | sequence_lengths: list[int], |
| 99 | ): |
| 100 | self.rank = rank |
| 101 | self.world_size = world_size |
| 102 | self.device = device |
| 103 | self.cpu_group = cpu_group |
| 104 | |
| 105 | # Calculate max_size_override based on largest sequence length |
| 106 | max_seq_len = max(sequence_lengths) |
| 107 | max_tensor_elements = max_seq_len * HIDDEN_SIZE |
| 108 | self.max_size_override = max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1 |
| 109 | # AITER has effective max size of half the max size, so we double it, and |
| 110 | # account for AG output size (which is world_size larger) |
| 111 | self.aiter_max_size = ( |
| 112 | 2 * world_size * max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1 |
| 113 | ) |
| 114 | |
| 115 | # Initialize communicators |
| 116 | self.custom_allreduce = None |
| 117 | self.pynccl_comm = None |
| 118 | self.symm_mem_comm = None |
| 119 | self.symm_mem_comm_multimem = None |
| 120 | self.symm_mem_comm_two_shot = None |
| 121 | self.fi_ar_comm = None |
| 122 | self.aiter_comm = None |
| 123 | |
| 124 | self._init_communicators() |
| 125 | |
| 126 | def _init_communicators(self): |
| 127 | """Initialize all available communicators.""" |
| 128 | try: |
| 129 | self.custom_allreduce = CustomAllreduce( |
| 130 | group=self.cpu_group, |
| 131 | device=self.device, |
| 132 | max_size=self.max_size_override, |
| 133 | ) |
| 134 | if not self.custom_allreduce.disabled: |
| 135 | logger.info("Rank %s: CustomAllreduce initialized", self.rank) |
| 136 | else: |
| 137 | logger.info("Rank %s: CustomAllreduce disabled", self.rank) |
| 138 | except Exception as e: |
| 139 | logger.warning( |
| 140 | "Rank %s: Failed to initialize CustomAllreduce: %s", self.rank, e |
| 141 | ) |
| 142 | self.custom_allreduce = None |
| 143 | |
| 144 | try: |
| 145 | self.pynccl_comm = PyNcclCommunicator( |
| 146 | group=self.cpu_group, device=self.device |
| 147 | ) |
| 148 | if not self.pynccl_comm.disabled: |
| 149 | logger.info("Rank %s: PyNcclCommunicator initialized", self.rank) |
| 150 | register_nccl_symmetric_ops(self.pynccl_comm) |
| 151 | else: |
| 152 | logger.info("Rank %s: PyNcclCommunicator disabled", self.rank) |
| 153 | self.pynccl_comm = None |
| 154 | except Exception as e: |
| 155 | logger.warning( |
| 156 | "Rank %s: Failed to initialize PyNcclCommunicator: %s", self.rank, e |
| 157 | ) |
| 158 | self.pynccl_comm = None |
| 159 | |
| 160 | # Initialize variants for SymmMemCommunicator |
| 161 | try: |
| 162 | self.symm_mem_comm_multimem = SymmMemCommunicator( |
| 163 | group=self.cpu_group, |
| 164 | device=self.device, |
| 165 | force_multimem=True, |
| 166 | max_size_override=self.max_size_override, |
| 167 | ) |
| 168 | if not self.symm_mem_comm_multimem.disabled: |
| 169 | logger.info( |
| 170 | "Rank %s: SymmMemCommunicator (multimem) initialized", self.rank |
| 171 | ) |
| 172 | else: |
| 173 | self.symm_mem_comm_multimem = None |
| 174 | except Exception as e: |
| 175 | logger.warning( |
| 176 | "Rank %s: Failed to initialize SymmMemCommunicator (multimem): %s", |
| 177 | self.rank, |
| 178 | e, |
| 179 | ) |
| 180 | self.symm_mem_comm_multimem = None |
| 181 | |
| 182 | try: |
| 183 | self.symm_mem_comm_two_shot = SymmMemCommunicator( |
| 184 | group=self.cpu_group, |
| 185 | device=self.device, |
| 186 | force_multimem=False, |
| 187 | max_size_override=self.max_size_override, |
| 188 | ) |
| 189 | if not self.symm_mem_comm_two_shot.disabled: |
| 190 | logger.info( |
| 191 | "Rank %s: SymmMemCommunicator (two_shot) initialized", self.rank |
| 192 | ) |
| 193 | else: |
| 194 | self.symm_mem_comm_two_shot = None |
| 195 | except Exception as e: |
| 196 | logger.warning( |
| 197 | "Rank %s: Failed to initialize SymmMemCommunicator (two_shot): %s", |
| 198 | self.rank, |
| 199 | e, |
| 200 | ) |
| 201 | self.symm_mem_comm_two_shot = None |
| 202 | |
| 203 | try: |
| 204 | self.fi_ar_comm = FlashInferAllReduce( |
| 205 | group=self.cpu_group, |
| 206 | device=self.device, |
| 207 | ) |
| 208 | if not self.fi_ar_comm.disabled: |
| 209 | logger.info("Rank %s: FlashInferAllReduce initialized", self.rank) |
| 210 | else: |
| 211 | logger.info("Rank %s: FlashInferAllReduce disabled", self.rank) |
| 212 | self.fi_ar_comm = None |
| 213 | except Exception as e: |
| 214 | logger.warning( |
| 215 | "Rank %s: Failed to initialize FlashInferAllReduce: %s", self.rank, e |
| 216 | ) |
| 217 | self.fi_ar_comm = None |
| 218 | |
| 219 | if current_platform.is_rocm() and rocm_aiter_ops.is_custom_all_reduce_enabled(): |
| 220 | try: |
| 221 | # Also used for AG/RS for DP |
| 222 | self.aiter_comm = AiterCustomAllreduce( |
| 223 | group=self.cpu_group, |
| 224 | device=self.device, |
| 225 | max_size=self.aiter_max_size, |
| 226 | ) |
| 227 | if not self.aiter_comm.disabled: |
| 228 | logger.info("Rank %s: AITER custom AG/RS initialized", self.rank) |
| 229 | else: |
| 230 | logger.info("Rank %s: AITER custom AG/RS disabled", self.rank) |
| 231 | self.aiter_comm = None |
| 232 | except Exception as e: |
| 233 | logger.warning( |
| 234 | "Rank %s: Failed to initialize AITER custom AG/RS: %s", |
| 235 | self.rank, |
| 236 | e, |
| 237 | ) |
| 238 | self.aiter_comm = None |
| 239 | |
| 240 | def benchmark_allreduce( |
| 241 | self, sequence_length: int, num_warmup: int, num_trials: int |
| 242 | ) -> dict[str, float]: |
| 243 | """Benchmark allreduce operations for all available communicators.""" |
| 244 | results = {} |
| 245 | |
| 246 | # Define communicators with their benchmark functions |
| 247 | communicators = [] |
| 248 | |
| 249 | if self.custom_allreduce is not None: |
| 250 | comm = self.custom_allreduce |
| 251 | # CustomAllreduce one-shot |
| 252 | communicators.append( |
| 253 | ( |
| 254 | "ca_1stage", |
| 255 | lambda t, c=comm: c.custom_all_reduce(t), |
| 256 | lambda t, c=comm: c.should_custom_ar(t), |
| 257 | comm.capture(), |
| 258 | {"VLLM_CUSTOM_ALLREDUCE_ALGO": "1stage"}, |
| 259 | None, # no destroy function |
| 260 | ) |
| 261 | ) |
| 262 | # CustomAllreduce two-shot |
| 263 | communicators.append( |
| 264 | ( |
| 265 | "ca_2stage", |
| 266 | lambda t, c=comm: c.custom_all_reduce(t), |
| 267 | lambda t, c=comm: c.should_custom_ar(t), |
| 268 | comm.capture(), |
| 269 | {"VLLM_CUSTOM_ALLREDUCE_ALGO": "2stage"}, |
| 270 | None, # no destroy function |
| 271 | ) |
| 272 | ) |
| 273 | |
| 274 | if self.pynccl_comm is not None: |
| 275 | comm = self.pynccl_comm |
| 276 | communicators.append( |
| 277 | ( |
| 278 | "pynccl", |
| 279 | lambda t, c=comm: c.all_reduce(t), |
| 280 | lambda t: True, # Always available if initialized |
| 281 | nullcontext(), |
| 282 | {}, # no env variable needed |
| 283 | None, # no destroy function |
| 284 | ) |
| 285 | ) |
| 286 | communicators.append( |
| 287 | ( |
| 288 | "pynccl-symm", |
| 289 | lambda t: torch.ops.vllm.all_reduce_symmetric_with_copy(t), |
| 290 | lambda t: True, # Always available if initialized |
| 291 | nullcontext(), |
| 292 | {}, # no env variable needed |
| 293 | None, # no destroy function |
| 294 | ) |
| 295 | ) |
| 296 | |
| 297 | if self.symm_mem_comm_multimem is not None: |
| 298 | comm = self.symm_mem_comm_multimem |
| 299 | communicators.append( |
| 300 | ( |
| 301 | "symm_mem_multimem", |
| 302 | lambda t, c=comm: c.all_reduce(t), |
| 303 | lambda t, c=comm: c.should_use_symm_mem(t), |
| 304 | nullcontext(), |
| 305 | {}, # no env variable needed |
| 306 | None, # no destroy function |
| 307 | ) |
| 308 | ) |
| 309 | |
| 310 | if self.symm_mem_comm_two_shot is not None: |
| 311 | comm = self.symm_mem_comm_two_shot |
| 312 | communicators.append( |
| 313 | ( |
| 314 | "symm_mem_two_shot", |
| 315 | lambda t, c=comm: c.all_reduce(t), |
| 316 | lambda t, c=comm: c.should_use_symm_mem(t), |
| 317 | nullcontext(), |
| 318 | {}, # no env variable needed |
| 319 | None, # no destroy function needed |
| 320 | ) |
| 321 | ) |
| 322 | |
| 323 | if self.fi_ar_comm is not None: |
| 324 | comm = self.fi_ar_comm |
| 325 | communicators.append( |
| 326 | ( |
| 327 | "flashinfer_trtllm", |
| 328 | lambda t, c=comm: c.all_reduce(t), |
| 329 | lambda t, c=comm: c.should_use_fi_ar(t), |
| 330 | nullcontext(), |
| 331 | {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "trtllm"}, |
| 332 | lambda c=comm: c.destroy(), |
| 333 | ) |
| 334 | ) |
| 335 | communicators.append( |
| 336 | ( |
| 337 | "flashinfer_mnnvl", |
| 338 | lambda t, c=comm: c.all_reduce(t), |
| 339 | lambda t, c=comm: c.should_use_fi_ar(t), |
| 340 | nullcontext(), |
| 341 | {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "mnnvl"}, |
| 342 | lambda c=comm: c.destroy(), |
| 343 | ) |
| 344 | ) |
| 345 | |
| 346 | # Benchmark each communicator |
| 347 | for ( |
| 348 | name, |
| 349 | allreduce_fn, |
| 350 | should_use_fn, |
| 351 | context, |
| 352 | env_dict, |
| 353 | destroy_fn, |
| 354 | ) in communicators: |
| 355 | # Save original values and apply new environment variables |
| 356 | saved_env = {key: os.environ.get(key) for key in env_dict} |
| 357 | for key, value in env_dict.items(): |
| 358 | os.environ[key] = value |
| 359 | try: |
| 360 | latency = self.benchmark_allreduce_single( |
| 361 | sequence_length, |
| 362 | allreduce_fn, |
| 363 | should_use_fn, |
| 364 | context, |
| 365 | num_warmup, |
| 366 | num_trials, |
| 367 | ) |
| 368 | if latency is not None: |
| 369 | results[name] = latency |
| 370 | finally: |
| 371 | if destroy_fn is not None: |
| 372 | destroy_fn() |
| 373 | # Restore environment variables to their original state |
| 374 | for key, original_value in saved_env.items(): |
| 375 | if original_value is None: |
| 376 | os.environ.pop(key, None) |
| 377 | else: |
| 378 | os.environ[key] = original_value |
| 379 | |
| 380 | return results |
| 381 | |
| 382 | def benchmark_allreduce_single( |
| 383 | self, |
| 384 | sequence_length: int, |
| 385 | allreduce_fn: Callable[[torch.Tensor], torch.Tensor | None], |
| 386 | should_use_fn: Callable[[torch.Tensor], bool], |
| 387 | context, |
| 388 | num_warmup: int, |
| 389 | num_trials: int, |
| 390 | ) -> float | None: |
| 391 | """Benchmark method with CUDA graph optimization.""" |
| 392 | try: |
| 393 | # Create test tensor (2D: sequence_length x hidden_size) |
| 394 | tensor = torch.randn( |
| 395 | sequence_length, HIDDEN_SIZE, dtype=BENCHMARK_DTYPE, device=self.device |
| 396 | ) |
| 397 | if not should_use_fn(tensor): |
| 398 | return None |
| 399 | |
| 400 | torch.accelerator.synchronize() |
| 401 | stream = torch.cuda.Stream() |
| 402 | with torch.cuda.stream(stream): |
| 403 | graph_input = tensor.clone() |
| 404 | |
| 405 | # Warmup before capture |
| 406 | for _ in range(3): |
| 407 | allreduce_fn(graph_input) |
| 408 | |
| 409 | # Capture the graph using context manager |
| 410 | with context: |
| 411 | graph = torch.cuda.CUDAGraph() |
| 412 | graph_pool = torch.cuda.graph_pool_handle() |
| 413 | set_graph_pool_id(graph_pool) |
| 414 | with torch.cuda.graph(graph, pool=graph_pool, stream=stream): |
| 415 | for _ in range(CUDA_GRAPH_CAPTURE_CYCLES): |
| 416 | allreduce_fn(graph_input) |
| 417 | |
| 418 | torch.accelerator.synchronize() |
| 419 | for _ in range(num_warmup): |
| 420 | graph.replay() |
| 421 | torch.accelerator.synchronize() |
| 422 | |
| 423 | torch.accelerator.synchronize() |
| 424 | start_time = time.perf_counter() |
| 425 | |
| 426 | for _ in range(num_trials): |
| 427 | graph.replay() |
| 428 | torch.accelerator.synchronize() |
| 429 | |
| 430 | end_time = time.perf_counter() |
| 431 | |
| 432 | # Convert to ms and divide by CUDA_GRAPH_CAPTURE_CYCLES |
| 433 | return ( |
| 434 | (end_time - start_time) / num_trials / CUDA_GRAPH_CAPTURE_CYCLES * 1000 |
| 435 | ) |
| 436 | |
| 437 | except Exception as e: |
| 438 | logger.error("CUDA graph benchmark failed: %s", e) |
| 439 | raise RuntimeError( |
| 440 | f"CUDA graph benchmark failed for communicator: {e}" |
| 441 | ) from e |
| 442 | |
| 443 | def benchmark_ag_rs( |
| 444 | self, sequence_length: int, num_warmup: int, num_trials: int |
| 445 | ) -> dict[str, float]: |
| 446 | """Benchmark all-gather and reduce-scatter: custom (AITER) vs NCCL/RCCL. |
| 447 | |
| 448 | Each collective is measured for both the AITER custom kernel and the |
| 449 | PyNccl baseline so they compare like-for-like (print_results groups |
| 450 | columns by collective). All-gather maps ``(seq, hidden)`` -> |
| 451 | ``(world_size * seq, hidden)``; reduce-scatter maps ``(seq, hidden)`` -> |
| 452 | ``(seq // world_size, hidden)`` and so only runs when ``seq`` divides |
| 453 | evenly across ranks. Both reuse ``benchmark_allreduce_single`` timed on a |
| 454 | ``(seq, hidden)`` input; baselines write into a pre-allocated buffer. |
| 455 | """ |
| 456 | results: dict[str, float] = {} |
| 457 | entries = [] |
| 458 | |
| 459 | # ---- all-gather: (seq, hidden) -> (world_size * seq, hidden) ---- |
| 460 | if self.aiter_comm is not None: |
| 461 | c = self.aiter_comm |
| 462 | entries.append( |
| 463 | ( |
| 464 | "aiter_ag", |
| 465 | lambda t, c=c: c.custom_all_gather(t, dim=0), |
| 466 | lambda t, c=c: c.should_custom_ag(t), |
| 467 | c.capture(), |
| 468 | ) |
| 469 | ) |
| 470 | if self.pynccl_comm is not None: |
| 471 | c = self.pynccl_comm |
| 472 | ag_out = torch.empty( |
| 473 | sequence_length * self.world_size, |
| 474 | HIDDEN_SIZE, |
| 475 | dtype=BENCHMARK_DTYPE, |
| 476 | device=self.device, |
| 477 | ) |
| 478 | entries.append( |
| 479 | ( |
| 480 | "pynccl_ag", |
| 481 | lambda t, c=c, o=ag_out: c.all_gather(o, t), |
| 482 | lambda t: True, |
| 483 | nullcontext(), |
| 484 | ) |
| 485 | ) |
| 486 | |
| 487 | # ---- reduce-scatter: (seq, hidden) -> (seq // world_size, hidden) ---- |
| 488 | if sequence_length % self.world_size == 0: |
| 489 | rs_shape = (sequence_length // self.world_size, HIDDEN_SIZE) |
| 490 | if self.aiter_comm is not None: |
| 491 | c = self.aiter_comm |
| 492 | rs_out = torch.empty( |
| 493 | rs_shape, dtype=BENCHMARK_DTYPE, device=self.device |
| 494 | ) |
| 495 | entries.append( |
| 496 | ( |
| 497 | "aiter_rs", |
| 498 | lambda t, c=c, o=rs_out: c.custom_reduce_scatter(t, o, dim=0), |
| 499 | lambda t, c=c: c.should_custom_rs(t, dim=0), |
| 500 | c.capture(), |
| 501 | ) |
| 502 | ) |
| 503 | if self.pynccl_comm is not None: |
| 504 | c = self.pynccl_comm |
| 505 | rs_out = torch.empty( |
| 506 | rs_shape, dtype=BENCHMARK_DTYPE, device=self.device |
| 507 | ) |
| 508 | entries.append( |
| 509 | ( |
| 510 | "pynccl_rs", |
| 511 | lambda t, c=c, o=rs_out: c.reduce_scatter(o, t), |
| 512 | lambda t: True, |
| 513 | nullcontext(), |
| 514 | ) |
| 515 | ) |
| 516 | |
| 517 | for name, fn, should_use_fn, context in entries: |
| 518 | latency = self.benchmark_allreduce_single( |
| 519 | sequence_length, fn, should_use_fn, context, num_warmup, num_trials |
| 520 | ) |
| 521 | if latency is not None: |
| 522 | results[name] = latency |
| 523 | |
| 524 | return results |
| 525 | |
| 526 | |
| 527 | def _calculate_speedup_info(comm_results: dict[str, float], baseline: str) -> str: |
| 528 | """Fastest comm in a collective group and its speedup vs ``baseline``.""" |
| 529 | if not comm_results: |
| 530 | return "N/A" |
| 531 | |
| 532 | fastest_comm = min(comm_results.keys(), key=lambda k: comm_results[k]) |
| 533 | fastest_time = comm_results[fastest_comm] |
| 534 | |
| 535 | if baseline in comm_results and fastest_time > 0: |
| 536 | speedup = comm_results[baseline] / fastest_time |
| 537 | return f"{fastest_comm} ({speedup:.2f}x)" |
| 538 | return f"{fastest_comm} (N/A)" |
| 539 | |
| 540 | |
| 541 | def print_results( |
| 542 | results: dict[int, dict[str, float]], sequence_lengths: list[int], world_size: int |
| 543 | ): |
| 544 | """Print benchmark results as one table per collective. |
| 545 | |
| 546 | All-reduce, all-gather and reduce-scatter are different collectives, so each |
| 547 | gets its own table and its speedup is computed against that collective's |
| 548 | NCCL/RCCL baseline rather than across collectives. |
| 549 | """ |
| 550 | present: set[str] = set() |
| 551 | for size_results in results.values(): |
| 552 | present.update(size_results.keys()) |
| 553 | |
| 554 | print(f"\n{'=' * 130}") |
| 555 | print("Device Communicator Benchmark Results") |
| 556 | print( |
| 557 | f"World Size: {world_size}, Data Type: {BENCHMARK_DTYPE}, " |
| 558 | f"Hidden Size: {HIDDEN_SIZE}" |
| 559 | ) |
| 560 | print("Tensor Shape is the per-rank input; times are ms per operation.") |
| 561 | print(f"{'=' * 130}") |
| 562 | |
| 563 | for collective in COLLECTIVE_ORDER: |
| 564 | comms = sorted(c for c in present if COMM_COLLECTIVE.get(c) == collective) |
| 565 | if not comms: |
| 566 | continue |
| 567 | baseline = COLLECTIVE_BASELINE[collective] |
| 568 | |
| 569 | print(f"\n{collective} (speedup vs {baseline})") |
| 570 | header = f"{'Tensor Shape':<20}{'Tensor Size':<15}" |
| 571 | for comm in comms: |
| 572 | header += f"{comm:<20}" |
| 573 | header += f"{'Best':<30}" |
| 574 | print(header) |
| 575 | print("-" * len(header)) |
| 576 | |
| 577 | for seq_len in sequence_lengths: |
| 578 | if seq_len not in results: |
| 579 | continue |
| 580 | row_results = { |
| 581 | c: results[seq_len][c] for c in comms if c in results[seq_len] |
| 582 | } |
| 583 | |
| 584 | tensor_bytes = seq_len * HIDDEN_SIZE * BENCHMARK_DTYPE.itemsize |
| 585 | tensor_size_str = f"{tensor_bytes / (1024 * 1024):.2f} MB" |
| 586 | tensor_shape = f"({seq_len}, {HIDDEN_SIZE})" |
| 587 | |
| 588 | row = f"{tensor_shape:<20}{tensor_size_str:<15}" |
| 589 | for comm in comms: |
| 590 | if comm in row_results: |
| 591 | row += f"{row_results[comm]:<20.3f}" |
| 592 | else: |
| 593 | row += f"{'N/A':<20}" |
| 594 | row += f"{_calculate_speedup_info(row_results, baseline):<30}" |
| 595 | print(row) |
| 596 | |
| 597 | print(f"\n{'=' * 130}") |
| 598 | print("All times are in milliseconds (ms) per operation") |
| 599 | |
| 600 | |
| 601 | def main(): |
| 602 | parser = FlexibleArgumentParser(description="Benchmark device communicators") |
| 603 | |
| 604 | parser.add_argument( |
| 605 | "--sequence-lengths", |
| 606 | type=int, |
| 607 | nargs="+", |
| 608 | default=DEFAULT_SEQUENCE_LENGTHS, |
| 609 | help="Sequence lengths to benchmark (tensor shape: seq_len x hidden_size)", |
| 610 | ) |
| 611 | |
| 612 | parser.add_argument( |
| 613 | "--num-warmup", type=int, default=5, help="Number of warmup iterations" |
| 614 | ) |
| 615 | |
| 616 | parser.add_argument( |
| 617 | "--num-trials", type=int, default=50, help="Number of benchmark trials" |
| 618 | ) |
| 619 | |
| 620 | parser.add_argument("--output-json", type=str, help="Output results to JSON file") |
| 621 | |
| 622 | args = parser.parse_args() |
| 623 | |
| 624 | # Initialize distributed |
| 625 | if not dist.is_initialized(): |
| 626 | dist.init_process_group(backend="gloo") |
| 627 | rank = dist.get_rank() |
| 628 | world_size = dist.get_world_size() |
| 629 | |
| 630 | # Set device |
| 631 | device = torch.device(f"cuda:{rank}") |
| 632 | torch.accelerator.set_device_index(device) |
| 633 | |
| 634 | # Get CPU process group |
| 635 | cpu_group = dist.new_group(backend="gloo") |
| 636 | |
| 637 | # Disable USE_SYMM_MEM to avoid affecting the max_sizes |
| 638 | # in symm_mem and custom_all_reduce for benchmark |
| 639 | os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" |
| 640 | |
| 641 | # Initialize benchmark |
| 642 | benchmark = CommunicatorBenchmark( |
| 643 | rank, world_size, device, cpu_group, args.sequence_lengths |
| 644 | ) |
| 645 | |
| 646 | # Run benchmarks |
| 647 | all_results = {} |
| 648 | |
| 649 | for seq_len in args.sequence_lengths: |
| 650 | if rank == 0: |
| 651 | logger.info( |
| 652 | "Benchmarking sequence length: %s (tensor shape: %s x %s)", |
| 653 | seq_len, |
| 654 | seq_len, |
| 655 | HIDDEN_SIZE, |
| 656 | ) |
| 657 | |
| 658 | results = benchmark.benchmark_allreduce( |
| 659 | sequence_length=seq_len, |
| 660 | num_warmup=args.num_warmup, |
| 661 | num_trials=args.num_trials, |
| 662 | ) |
| 663 | results.update( |
| 664 | benchmark.benchmark_ag_rs( |
| 665 | sequence_length=seq_len, |
| 666 | num_warmup=args.num_warmup, |
| 667 | num_trials=args.num_trials, |
| 668 | ) |
| 669 | ) |
| 670 | |
| 671 | all_results[seq_len] = results |
| 672 | |
| 673 | # Synchronize between ranks |
| 674 | dist.barrier() |
| 675 | |
| 676 | # Print results (only rank 0) |
| 677 | if rank == 0: |
| 678 | print_results(all_results, args.sequence_lengths, world_size) |
| 679 | |
| 680 | # Save to JSON if requested |
| 681 | if args.output_json: |
| 682 | # Add per-collective speedup information to results. |
| 683 | enhanced_results = {} |
| 684 | for seq_len, comm_results in all_results.items(): |
| 685 | speedup_info = {} |
| 686 | for collective in COLLECTIVE_ORDER: |
| 687 | group = { |
| 688 | k: v |
| 689 | for k, v in comm_results.items() |
| 690 | if COMM_COLLECTIVE.get(k) == collective |
| 691 | } |
| 692 | if group: |
| 693 | speedup_info[collective] = _calculate_speedup_info( |
| 694 | group, COLLECTIVE_BASELINE[collective] |
| 695 | ) |
| 696 | enhanced_results[seq_len] = { |
| 697 | "timings": comm_results, |
| 698 | "speedup_info": speedup_info, |
| 699 | } |
| 700 | |
| 701 | output_data = { |
| 702 | "world_size": world_size, |
| 703 | "dtype": str(BENCHMARK_DTYPE), |
| 704 | "hidden_size": HIDDEN_SIZE, |
| 705 | "sequence_lengths": args.sequence_lengths, |
| 706 | "num_warmup": args.num_warmup, |
| 707 | "num_trials": args.num_trials, |
| 708 | "cuda_graph_capture_cycles": CUDA_GRAPH_CAPTURE_CYCLES, |
| 709 | "results": enhanced_results, |
| 710 | } |
| 711 | |
| 712 | with open(args.output_json, "w") as f: |
| 713 | json.dump(output_data, f, indent=2) |
| 714 | |
| 715 | logger.info("Results saved to %s", args.output_json) |
| 716 | |
| 717 | # Cleanup |
| 718 | if cpu_group != dist.group.WORLD: |
| 719 | dist.destroy_process_group(cpu_group) |
| 720 | |
| 721 | |
| 722 | if __name__ == "__main__": |
| 723 | main() |