| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from enum import Enum |
| 5 | |
| 6 | import torch |
| 7 | import torch.distributed as dist |
| 8 | from torch.distributed import ProcessGroup |
| 9 | |
| 10 | import vllm.envs as envs |
| 11 | from vllm import _custom_ops as ops |
| 12 | from vllm.config import get_current_vllm_config_or_none |
| 13 | from vllm.distributed.parallel_state import in_the_same_node_as |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.platforms import current_platform |
| 16 | |
| 17 | logger = init_logger(__name__) |
| 18 | |
| 19 | try: |
| 20 | ops.qr_max_size() |
| 21 | quick_ar = True |
| 22 | except Exception: |
| 23 | # For CPUs and CUDA |
| 24 | quick_ar = False |
| 25 | |
| 26 | |
| 27 | from vllm.distributed.utils import is_weak_contiguous # noqa: E402, F401 |
| 28 | |
| 29 | |
| 30 | class QuickReduceRegime(Enum): |
| 31 | # Keep integer ids aligned with csrc/quickreduce/quick_reduce.h |
| 32 | FP = 0 |
| 33 | INT8 = 1 |
| 34 | INT6 = 2 |
| 35 | INT4 = 3 |
| 36 | INT3 = 4 |
| 37 | NONE = 5 |
| 38 | |
| 39 | |
| 40 | KB = 1024 |
| 41 | MB = 1024 * KB |
| 42 | |
| 43 | |
| 44 | class QuickAllReduce: |
| 45 | _SUPPORTED_WORLD_SIZES = [2, 4, 8] |
| 46 | _SUPPORTED_DTYPES = [torch.float16, torch.bfloat16] |
| 47 | # The following data is based on kernel tests. |
| 48 | # In this order [FP, INT8, INT6, INT4, INT3]. |
| 49 | _QR_MIN_SIZE = { |
| 50 | (torch.float16, 2): [1 * MB, 2 * MB, 2 * MB, 1 * MB, 1 * MB], |
| 51 | (torch.float16, 4): [1 * MB, 16 * MB, 4 * MB, 2 * MB, 2 * MB], |
| 52 | (torch.float16, 8): [16 * MB, 4 * MB, 4 * MB, 2 * MB, 2 * MB], |
| 53 | (torch.bfloat16, 2): [2 * MB, 8 * MB, 8 * MB, 8 * MB, 8 * MB], |
| 54 | (torch.bfloat16, 4): [8 * MB, 64 * MB, 64 * MB, 16 * MB, 16 * MB], |
| 55 | (torch.bfloat16, 8): [ |
| 56 | 16 * MB, |
| 57 | 2048 * MB, |
| 58 | 2048 * MB, |
| 59 | 2048 * MB, |
| 60 | 2048 * MB, |
| 61 | ], |
| 62 | } |
| 63 | |
| 64 | def __init__(self, group: ProcessGroup, device: int | str | torch.device) -> None: |
| 65 | """Custom allreduce provides non-destructive acceleration and is |
| 66 | available for CUDA and ROCm MI300 series. |
| 67 | |
| 68 | Custom quick allreduce leverages quantization for further |
| 69 | acceleration on ROCm. It currently supports Q8, Q6, Q4, and Q3 |
| 70 | quantization formats and FP(float16, bfloat16). Q3 (INT3) is |
| 71 | restricted to TP2 (world_size == 2) due to poor performance on |
| 72 | larger world sizes. |
| 73 | |
| 74 | Quick allreduce is designed as a complement to custom allreduce. |
| 75 | Its initialization requires even stricter conditions. |
| 76 | |
| 77 | Only the ROCm MI300 series is supported for quick allreduce at |
| 78 | this time. |
| 79 | |
| 80 | Args: |
| 81 | group: the process group to work on. If None, it will use the |
| 82 | default process group. |
| 83 | device: the device to bind the CustomAllreduce to. If None, |
| 84 | it will be bound to f"cuda:{local_rank}". |
| 85 | It is the caller's responsibility to make sure each communicator |
| 86 | is bind to a unique device, and all communicators in this group |
| 87 | are in the same node. |
| 88 | |
| 89 | """ |
| 90 | self.disabled = True |
| 91 | if not self._rocm_arch_available(): |
| 92 | logger.debug( |
| 93 | "Custom quick allreduce is only supported on ROCm MI300 series." |
| 94 | ) |
| 95 | return |
| 96 | |
| 97 | if not quick_ar: |
| 98 | # disable because of missing quick reduce library |
| 99 | # e.g. in a cuda environment |
| 100 | logger.info( |
| 101 | "Custom quick allreduce is disabled because " |
| 102 | "of missing custom quick allreduce library" |
| 103 | ) |
| 104 | return |
| 105 | |
| 106 | self.group = group |
| 107 | assert dist.get_backend(group) != dist.Backend.NCCL, ( |
| 108 | "Custom quick allreduce should be attached to a non-NCCL group." |
| 109 | ) |
| 110 | if not all(in_the_same_node_as(group, source_rank=0)): |
| 111 | # No need to initialize custom quick allreduce for |
| 112 | # multi-node case. |
| 113 | logger.warning( |
| 114 | "Custom quick allreduce is disabled because this " |
| 115 | "process group spans across nodes." |
| 116 | ) |
| 117 | return |
| 118 | rank = dist.get_rank(group=self.group) |
| 119 | world_size = dist.get_world_size(group=self.group) |
| 120 | self.rank = rank |
| 121 | self.world_size = world_size |
| 122 | if world_size == 1: |
| 123 | # No need to initialize QuickReduce for single GPU case. |
| 124 | return |
| 125 | |
| 126 | if world_size not in QuickAllReduce._SUPPORTED_WORLD_SIZES: |
| 127 | logger.warning( |
| 128 | "Custom quick allreduce is disabled due to an " |
| 129 | "unsupported world size: %d. Supported world sizes: %s.", |
| 130 | world_size, |
| 131 | str(QuickAllReduce._SUPPORTED_WORLD_SIZES), |
| 132 | ) |
| 133 | return |
| 134 | |
| 135 | if isinstance(device, int): |
| 136 | device = torch.device(f"cuda:{device}") |
| 137 | elif isinstance(device, str): |
| 138 | device = torch.device(device) |
| 139 | assert isinstance(device, torch.device) |
| 140 | self.device = device |
| 141 | |
| 142 | # device.index is a visible ordinal, not a logical local ID. |
| 143 | physical_device_id = current_platform.visible_device_id_to_physical_device_id( |
| 144 | device.index |
| 145 | ) |
| 146 | tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu") |
| 147 | gather_list = [ |
| 148 | torch.tensor([0], dtype=torch.int, device="cpu") |
| 149 | for _ in range(self.world_size) |
| 150 | ] |
| 151 | dist.all_gather(gather_list, tensor, group=self.group) |
| 152 | physical_device_ids = [t.item() for t in gather_list] |
| 153 | |
| 154 | # test nvlink first, this will filter out most of the cases |
| 155 | # where custom quick allreduce is not supported |
| 156 | # this checks hardware and driver support for NVLink |
| 157 | assert current_platform.is_cuda_alike() |
| 158 | self.fully_connected = current_platform.is_fully_connected(physical_device_ids) |
| 159 | if self.world_size > 2 and not self.fully_connected: |
| 160 | logger.debug( |
| 161 | "Custom quick allreduce is disabled because it's not supported " |
| 162 | "on more than two PCIe-only GPUs. " |
| 163 | ) |
| 164 | return |
| 165 | |
| 166 | self.init_quick_all_reduce() |
| 167 | |
| 168 | def init_quick_all_reduce(self): |
| 169 | # On RocM, bfloat16 kernels are slower than fp16 |
| 170 | # due to slower match operations |
| 171 | # If environment variable is set to 1, we convert input to fp16 |
| 172 | self.use_fp16_kernels = envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 |
| 173 | regime_str = envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION |
| 174 | if regime_str not in QuickReduceRegime.__members__: |
| 175 | logger.warning( |
| 176 | "Custom quick allreduce:", |
| 177 | f"Invalid quantization level: {regime_str}. " |
| 178 | "Supported levels: " |
| 179 | f"{list(QuickReduceRegime.__members__.keys())}", |
| 180 | ) |
| 181 | return |
| 182 | |
| 183 | if regime_str == "NONE": |
| 184 | logger.debug( |
| 185 | "Custom quick allreduce is disabled based " |
| 186 | "on env variable " |
| 187 | "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION='NONE'" |
| 188 | ) |
| 189 | return |
| 190 | self.qr_quant_level = QuickReduceRegime[regime_str] |
| 191 | |
| 192 | # INT3 is only enabled for TP2 (world_size == 2). |
| 193 | # Kernel benchmarks show INT3 all-reduce on TP4/TP8 has poor |
| 194 | # performance (the extra ranks make the 3-bit codec's pack/unpack |
| 195 | # overhead outweigh the reduced communication volume), so INT3 is |
| 196 | # restricted to 2-GPU tensor parallelism. For TP4/TP8 use a wider |
| 197 | # codec (e.g. INT4) or NONE instead. |
| 198 | if self.qr_quant_level == QuickReduceRegime.INT3 and self.world_size != 2: |
| 199 | logger.warning( |
| 200 | "Custom quick allreduce is disabled: INT3 quantization is " |
| 201 | "only supported for TP2 (world_size == 2), but world_size " |
| 202 | "is %d. INT3 on TP4/TP8 is disabled due to poor kernel " |
| 203 | "performance. Use INT4/NONE for this world size.", |
| 204 | self.world_size, |
| 205 | ) |
| 206 | return |
| 207 | |
| 208 | self.qr_quantization_min_size = self._get_qr_quantization_min_size() |
| 209 | vllm_config = get_current_vllm_config_or_none() |
| 210 | if ( |
| 211 | vllm_config is not None |
| 212 | and hasattr(vllm_config, "model_config") |
| 213 | and hasattr(vllm_config.model_config, "dtype") |
| 214 | ): |
| 215 | dtype = vllm_config.model_config.dtype |
| 216 | if dtype not in [torch.float16, torch.bfloat16]: |
| 217 | logger.debug( |
| 218 | "Custom quick allreduce disabled: only supports " |
| 219 | "float16 and float16, but get %s.", |
| 220 | dtype, |
| 221 | ) |
| 222 | return |
| 223 | |
| 224 | if dtype == torch.bfloat16 and self.use_fp16_kernels: |
| 225 | logger.info( |
| 226 | "Custom quick allreduce: BF16 inputs will be converted " |
| 227 | "to FP16 to improve performance. set " |
| 228 | "envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16=0 " |
| 229 | "to turn off." |
| 230 | ) |
| 231 | |
| 232 | # VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is specified in MB |
| 233 | qr_max_size = envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB |
| 234 | if qr_max_size is not None: |
| 235 | if qr_max_size < 1: |
| 236 | logger.info( |
| 237 | "You should not set a max_size smaller than 1MB, which can " |
| 238 | "lead to error or degradation to custom allreduce or rccl." |
| 239 | ) |
| 240 | qr_max_size = qr_max_size * MB |
| 241 | effective_qr_max_size = ( |
| 242 | qr_max_size if qr_max_size is not None else ops.qr_max_size() |
| 243 | ) |
| 244 | qr_min_size = self._get_qr_min_size(effective_qr_max_size) |
| 245 | self._ptr = ops.init_custom_qr(self.rank, self.world_size, qr_max_size) |
| 246 | self.qr_max_size = effective_qr_max_size |
| 247 | self.qr_min_size = qr_min_size |
| 248 | if qr_min_size is not None: |
| 249 | logger.info( |
| 250 | "Custom quick allreduce: min size override = %d MB", |
| 251 | qr_min_size // MB, |
| 252 | ) |
| 253 | if self.qr_quantization_min_size is not None: |
| 254 | logger.info( |
| 255 | "Custom quick allreduce: quantization codec threshold = %d KB", |
| 256 | self.qr_quantization_min_size // KB, |
| 257 | ) |
| 258 | self.create_shared_buffer() |
| 259 | self.disabled = False |
| 260 | |
| 261 | @staticmethod |
| 262 | def _get_qr_min_size(qr_max_size: int | None) -> int | None: |
| 263 | qr_min_size = envs.VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB |
| 264 | if qr_min_size is None: |
| 265 | return None |
| 266 | if qr_min_size < 0: |
| 267 | raise ValueError( |
| 268 | "VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB must be non-negative, " |
| 269 | f"got {qr_min_size}" |
| 270 | ) |
| 271 | qr_min_size *= MB |
| 272 | if qr_max_size is not None and qr_min_size > qr_max_size: |
| 273 | raise ValueError( |
| 274 | "VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB must be less than or " |
| 275 | "equal to the effective QuickReduce max size" |
| 276 | ) |
| 277 | return qr_min_size |
| 278 | |
| 279 | @staticmethod |
| 280 | def _get_qr_quantization_min_size() -> int | None: |
| 281 | quantization_min_size = envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB |
| 282 | if quantization_min_size is None: |
| 283 | return None |
| 284 | if quantization_min_size < 0: |
| 285 | raise ValueError( |
| 286 | "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB must be " |
| 287 | f"non-negative, got {quantization_min_size}" |
| 288 | ) |
| 289 | return quantization_min_size * KB |
| 290 | |
| 291 | def _rocm_arch_available(self): |
| 292 | if not current_platform.is_rocm(): |
| 293 | return False |
| 294 | try: |
| 295 | props = torch.cuda.get_device_properties(0) |
| 296 | gcn_arch = getattr(props, "gcnArchName", "") |
| 297 | supported_archs = ["gfx94", "gfx95"] |
| 298 | return any(gfx in gcn_arch for gfx in supported_archs) |
| 299 | except Exception as e: |
| 300 | logger.warning("Failed to determine ROCm for quick allreduce: %s", e) |
| 301 | return False |
| 302 | |
| 303 | def create_shared_buffer(self): |
| 304 | """Creates a shared buffer for quickreduce. |
| 305 | Has to be called after init_custom_qr |
| 306 | """ |
| 307 | handle = ops.qr_get_handle(self._ptr) |
| 308 | world_size = dist.get_world_size(group=self.group) |
| 309 | handles = [None] * world_size |
| 310 | dist.all_gather_object(handles, handle, group=self.group) |
| 311 | ops.qr_open_handles(self._ptr, handles) |
| 312 | |
| 313 | def should_quick_allreduce(self, inp: torch.Tensor): |
| 314 | """Check if quickreduce is available.""" |
| 315 | if self.disabled: |
| 316 | return False |
| 317 | if inp.dtype not in self._SUPPORTED_DTYPES: |
| 318 | return False |
| 319 | inp_size = inp.numel() * inp.element_size() |
| 320 | # custom quick allreduce requires input byte size to be |
| 321 | # multiples of 16 |
| 322 | if inp_size % 16 != 0: |
| 323 | return False |
| 324 | if not is_weak_contiguous(inp): |
| 325 | return False |
| 326 | dtype = inp.dtype |
| 327 | if self.use_fp16_kernels: |
| 328 | dtype = torch.float16 |
| 329 | min_size = self.qr_min_size |
| 330 | if min_size is None: |
| 331 | min_size = self._QR_MIN_SIZE[(dtype, self.world_size)][ |
| 332 | self.qr_quant_level.value |
| 333 | ] |
| 334 | return inp_size <= self.qr_max_size and inp_size >= min_size |
| 335 | |
| 336 | def quick_all_reduce(self, inp: torch.Tensor, *, out: torch.Tensor = None): |
| 337 | """Performs an out-of-place custom quick all reduce.""" |
| 338 | # quick allreduce doesn't require a separate graph mode, |
| 339 | # as QR uses static IPC buffer. |
| 340 | if out is None: |
| 341 | out = torch.empty_like(inp) |
| 342 | ops.qr_all_reduce( |
| 343 | self._ptr, inp, out, self._get_qr_quant_level(inp), self.use_fp16_kernels |
| 344 | ) |
| 345 | return out |
| 346 | |
| 347 | def _get_qr_quant_level(self, inp: torch.Tensor) -> int: |
| 348 | quantization_min_size = self.qr_quantization_min_size |
| 349 | if ( |
| 350 | quantization_min_size is not None |
| 351 | and inp.numel() * inp.element_size() < quantization_min_size |
| 352 | ): |
| 353 | return QuickReduceRegime.FP.value |
| 354 | return self.qr_quant_level.value |
| 355 | |
| 356 | def close(self): |
| 357 | if not self.disabled and getattr(self, "_ptr", None): |
| 358 | if ops is not None: |
| 359 | ops.qr_destroy(self._ptr) |
| 360 | self._ptr = 0 |
| 361 | self.disabled = True |
| 362 | |
| 363 | def __del__(self): |
| 364 | self.close() |