rocm.py 20.8 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import os
5
from datetime import timedelta
6
from functools import cache, lru_cache, wraps
7
from typing import TYPE_CHECKING, Optional
8
9

import torch
10
11
from torch.distributed import PrefixStore, ProcessGroup
from torch.distributed.distributed_c10d import is_nccl_available
12

13
import vllm.envs as envs
14
from vllm.logger import init_logger
15
from vllm.utils import cuda_device_count_stateless
16

17
from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
18

zhuwenwen's avatar
zhuwenwen committed
19
from vllm.utils import is_kme, SUPPORT_TC
zhuwenwen's avatar
zhuwenwen committed
20
21
22
23
24
25

if not SUPPORT_TC:
    os.environ['VLLM_USE_V1'] = '0'
    os.environ['VLLM_USE_FLASH_ATTN_PA'] = '0'
    os.environ['VLLM_USE_FLASH_MLA'] = '0'
    
zhuwenwen's avatar
zhuwenwen committed
26
27
28
29
if is_kme:
    os.environ['VLLM_USE_FLASH_ATTN_PA'] = '0'
    
    
30
if TYPE_CHECKING:
31
    from vllm.config import ModelConfig, VllmConfig
32

33
34
logger = init_logger(__name__)

35
try:
36
37
38
    from amdsmi import (AmdSmiException, amdsmi_get_gpu_asic_info,
                        amdsmi_get_processor_handles, amdsmi_init,
                        amdsmi_shut_down, amdsmi_topo_get_link_type)
39
40
41
except ImportError as e:
    logger.warning("Failed to import from amdsmi with %r", e)

42
43
44
45
46
47
try:
    import vllm._C  # noqa: F401
except ImportError as e:
    logger.warning("Failed to import from vllm._C with %r", e)

# import custom ops, trigger op registration
zhuwenwen's avatar
zhuwenwen committed
48
49
50
51
# try:
#     import vllm._rocm_C  # noqa: F401
# except ImportError as e:
#     logger.warning("Failed to import from vllm._rocm_C with %r", e)
52

53

54
# Models not supported by ROCm.
55
_ROCM_UNSUPPORTED_MODELS: list[str] = []
56
57
58

# Models partially supported by ROCm.
# Architecture -> Reason.
zhuwenwen's avatar
zhuwenwen committed
59
60
61
62
# _ROCM_SWA_REASON = ("Sliding window attention (SWA) is not yet supported in "
#                     "Triton flash attention. For half-precision SWA support, "
#                     "please use CK flash attention by setting "
#                     "`VLLM_USE_TRITON_FLASH_ATTN=0`")
63
_ROCM_PARTIALLY_SUPPORTED_MODELS: dict[str, str] = {
zhuwenwen's avatar
zhuwenwen committed
64
65
66
67
68
69
    # "Qwen2ForCausalLM":
    # _ROCM_SWA_REASON,
    # "MistralForCausalLM":
    # _ROCM_SWA_REASON,
    # "MixtralForCausalLM":
    # _ROCM_SWA_REASON,
zhuwenwen's avatar
zhuwenwen committed
70
71
72
    # "PaliGemmaForConditionalGeneration":
    # ("ROCm flash attention does not yet "
    #  "fully support 32-bit precision on PaliGemma"),
zhuwenwen's avatar
zhuwenwen committed
73
74
75
76
    # "Phi3VForCausalLM":
    # ("ROCm Triton flash attention may run into compilation errors due to "
    #  "excessive use of shared memory. If this happens, disable Triton FA "
    #  "by setting `VLLM_USE_TRITON_FLASH_ATTN=0`")
77
}
78
_ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = {
79
80
81
82
83
84
85
86
    "0x74a0": "AMD_Instinct_MI300A",
    "0x74a1": "AMD_Instinct_MI300X",
    "0x74b5": "AMD_Instinct_MI300X",  # MI300X VF
    "0x74a5": "AMD_Instinct_MI325X",
    "0x74b9": "AMD_Instinct_MI325X",  # MI325X VF
    "0x74a9": "AMD_Instinct_MI300X_HF",
    "0x74bd": "AMD_Instinct_MI300X_HF",
}
87

88
# Prevent use of clashing `{CUDA/HIP}_VISIBLE_DEVICES``
zhuwenwen's avatar
zhuwenwen committed
89
90
91
92
93
94
# if "HIP_VISIBLE_DEVICES" in os.environ:
#     val = os.environ["HIP_VISIBLE_DEVICES"]
#     if cuda_val := os.environ.get("CUDA_VISIBLE_DEVICES", None):
#         assert val == cuda_val
#     else:
#         os.environ["CUDA_VISIBLE_DEVICES"] = val
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

# AMDSMI utils
# Note that NVML is not affected by `{CUDA/HIP}_VISIBLE_DEVICES`,
# all the related functions work on real physical device ids.
# the major benefit of using AMDSMI is that it will not initialize CUDA


def with_amdsmi_context(fn):

    @wraps(fn)
    def wrapper(*args, **kwargs):
        amdsmi_init()
        try:
            return fn(*args, **kwargs)
        finally:
            amdsmi_shut_down()

    return wrapper


115
116
117
118
119
120
121
122
123
def device_id_to_physical_device_id(device_id: int) -> int:
    if "CUDA_VISIBLE_DEVICES" in os.environ:
        device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
        physical_device_id = device_ids[device_id]
        return int(physical_device_id)
    else:
        return device_id


124
125
126
127
@cache
def on_gfx1x() -> bool:
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
    return any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
128

129

130
@cache
131
def on_mi3xx() -> bool:
132
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
133
    return any(arch in GPU_ARCH for arch in ["gfx942", "gfx950"])
134
135


136
@cache
137
138
139
def on_gfx9() -> bool:
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
    return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
140
141


142
@cache
143
144
145
146
147
148
149
150
151
def use_rocm_custom_paged_attention(
        qtype: torch.dtype,
        head_size: int,
        block_size: int,
        gqa_ratio: int,
        max_seq_len: int,
        sliding_window: int,
        kv_cache_dtype: str,
        alibi_slopes: Optional[torch.Tensor] = None) -> bool:
152

153
154
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
    ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
155
    ON_GFX11_GFX12 = any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
156

157
158
    # custom paged attn always supported on V0. On V1, requires sliding window
    # disabled due to observed numerical discrepancy.
zhuwenwen's avatar
zhuwenwen committed
159
160
161
162
163
164
165
    # if ON_GFX9:
    #     return ((not envs.VLLM_USE_V1 or sliding_window == 0
    #              or sliding_window == (-1, -1))
    #             and (qtype == torch.half or qtype == torch.bfloat16)
    #             and (head_size == 64 or head_size == 128)
    #             and (block_size == 16 or block_size == 32)
    #             and (gqa_ratio >= 1 and gqa_ratio <= 16)
zhuwenwen's avatar
zhuwenwen committed
166
167
    #             and max_seq_len <= 128 * 1024
    #             and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
zhuwenwen's avatar
zhuwenwen committed
168
169
170
171
172
173
174
175
176
    #             and not (envs.VLLM_ROCM_USE_AITER_PAGED_ATTN
    #                      and envs.VLLM_ROCM_USE_AITER))

    # else:
    #     return (ON_GFX11_GFX12 and (not envs.VLLM_USE_V1 or sliding_window == 0
    #                                 or sliding_window == (-1, -1))
    #             and (qtype == torch.half or qtype == torch.bfloat16)
    #             and head_size == 128 and block_size == 16
    #             and (gqa_ratio >= 3 and gqa_ratio <= 16)
zhuwenwen's avatar
zhuwenwen committed
177
    #             and max_seq_len <= 128 * 1024 and alibi_slopes is None
zhuwenwen's avatar
zhuwenwen committed
178
179
    #             and kv_cache_dtype == "auto"
    #             and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
180
    return False
181
182


183
184
class RocmPlatform(Platform):
    _enum = PlatformEnum.ROCM
185
    device_name: str = "rocm"
186
    device_type: str = "cuda"
187
    dispatch_key: str = "CUDA"
188
    ray_device_key: str = "GPU"
189
190
    # rocm shares the same device control env var as CUDA
    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
191

192
    supported_quantization: list[str] = [
193
        "awq", "gptq", "fp8", "compressed-tensors", "fbgemm_fp8", "gguf",
194
195
        "quark", "ptpc_fp8", "moe_wna16", "blockwise_int8","slimquant_w4a8","awq_marlin",
        "slimquant_w4a8_marlin","slimquant_compressed_tensors_marlin"
196
    ]
197

198
    @classmethod
199
    def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
200
201
                             kv_cache_dtype, block_size, use_v1,
                             use_mla) -> str:
202
        if use_mla:
zhuwenwen's avatar
zhuwenwen committed
203
204
205
206
207
            if selected_backend == _Backend.TRITON_MLA or block_size != 64:
                if use_v1:
                    logger.info_once("Using Triton MLA backend on V1 engine.")
                    return ("vllm.v1.attention.backends.mla."
                            "triton_mla.TritonMLABackend")
208

zhuwenwen's avatar
zhuwenwen committed
209
210
                else:
                    logger.info("Using Triton MLA backend.")
zhuwenwen's avatar
zhuwenwen committed
211
212
                    return "vllm.attention.backends.triton_mla.TritonMLABackend"   
            else:         
213
214
215
216
217
218
219
220
221
222
223
224
                if envs.VLLM_USE_FLASH_MLA:
                    from vllm.attention.backends.flashmla import (
                        is_flashmla_supported)
                    if not is_flashmla_supported()[0]:
                        logger.warning(
                            "FlashMLA backend is not supported due to %s",
                            is_flashmla_supported()[1])
                    elif block_size != 64:
                        logger.warning(
                            "FlashMLA backend is not supported for block size %d"
                            " (currently only supports block size 64).",
                            block_size)
zhuwenwen's avatar
zhuwenwen committed
225
                    else:
226
227
228
229
230
231
232
233
234
235
236
                        if use_v1:
                            logger.info_once(
                                "Using FlashMLA backend on V1 engine.")
                            return ("vllm.v1.attention.backends.mla."
                                    "flashmla.FlashMLABackend")
                        else:
                            logger.info("Using FlashMLA backend.")
                            return ("vllm.attention.backends."
                                    "flashmla.FlashMLABackend")
                else:
                    logger.info("Using Triton MLA backend (block size 64).")
zhuwenwen's avatar
zhuwenwen committed
237
                    return "vllm.attention.backends.triton_mla.TritonMLABackend"           
zhuwenwen's avatar
zhuwenwen committed
238
239
240
241
242
243
244
245
246
247
248

            # from vllm.attention.backends.rocm_aiter_mla import (
            #     is_aiter_mla_enabled)

            # if selected_backend is None:
            #     selected_backend = (_Backend.ROCM_AITER_MLA if
            #                         is_aiter_mla_enabled() or block_size == 1
            #                         else _Backend.TRITON_MLA)

            # if selected_backend == _Backend.TRITON_MLA:
            #     if block_size != 1:
zhuwenwen's avatar
zhuwenwen committed
249
250
251
252
253
254
255
256
            #         if use_v1:
            #             logger.info_once(
            #                 "Using Triton MLA backend on V1 engine.")
            #             return ("vllm.v1.attention.backends.mla."
            #                     "triton_mla.TritonMLABackend")
            #         else:
            #             logger.info("Using Triton MLA backend.")
            #             return "vllm.attention.backends.triton_mla.TritonMLABackend"  # noqa: E501
zhuwenwen's avatar
zhuwenwen committed
257
258
259
260
            #     else:
            #         raise ValueError(
            #             f" The selected backend, {selected_backend.name},"
            #             f"does not support block size {block_size}.")
zhuwenwen's avatar
zhuwenwen committed
261
262
            # elif selected_backend == _Backend.ROCM_AITER_MLA \
            #     or selected_backend == _Backend.ROCM_AITER_MLA_VLLM_V1:
zhuwenwen's avatar
zhuwenwen committed
263
            #     if block_size == 1:
zhuwenwen's avatar
zhuwenwen committed
264
265
266
267
268
269
            #         if use_v1:
            #             logger.info("Using AITER MLA backend on V1 engine.")
            #             return "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"  # noqa: E501
            #         else:
            #             logger.info("Using AITER MLA backend")
            #             return "vllm.attention.backends.rocm_aiter_mla.AiterMLABackend"  # noqa: E501
zhuwenwen's avatar
zhuwenwen committed
270
271
272
273
274
275
276
277
278
            #     else:
            #         raise ValueError(
            #             f" The selected backend, {selected_backend.name},"
            #             f"does not support block size {block_size}."
            #             "(currently only supports block size 1)")
            # else:
            #     raise ValueError(
            #         f" The selected backend, {selected_backend.name},"
            #         f"is not MLA type while requested for MLA backend.")
279
280
281
      
        if selected_backend is None or selected_backend == _Backend.FLASH_ATTN:
            selected_backend = _Backend.ROCM_FLASH
282

283
        if envs.VLLM_USE_V1:
zhuwenwen's avatar
zhuwenwen committed
284
285
286
287
288
289
            TRITON_ATTN_VLLM_V1 = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"  # noqa: E501
            FLASH_ATTN_V1 = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"  # noqa: E501
            # if selected_backend == _Backend.TRITON_ATTN_VLLM_V1:
            #     logger.info_once("Using Triton backend on V1 engine.")
            #     return TRITON_ATTN_VLLM_V1
            
zhuwenwen's avatar
zhuwenwen committed
290
            if envs.VLLM_USE_FLASH_ATTN_PA and block_size == 64:
zhuwenwen's avatar
zhuwenwen committed
291
292
                logger.info_once("Using Flash Attention backend on V1 engine. (only supports block size 64)")
                return FLASH_ATTN_V1
293
            else:
294
                os.environ['VLLM_USE_FLASH_ATTN_PA'] = '0'
zhuwenwen's avatar
zhuwenwen committed
295
296
297
                logger.info_once("Using Triton backend on V1 engine.")
                return TRITON_ATTN_VLLM_V1
            
298
299
300
301
        if selected_backend == _Backend.DUAL_CHUNK_FLASH_ATTN:
            logger.info("Using DualChunkFlashAttention backend.")
            return ("vllm.attention.backends.dual_chunk_flash_attn."
                    "DualChunkFlashAttentionBackend")
302
303
304
305
306
307
        if selected_backend == _Backend.ROCM_FLASH:
            if not cls.has_device_capability(90):
                # not Instinct series GPUs.
                logger.info("flash_attn is not supported on NAVI GPUs.")
        else:
            logger.info("%s is not supported in AMD GPUs.", selected_backend)
zhuwenwen's avatar
zhuwenwen committed
308
309
        if is_kme:
            os.environ['VLLM_USE_TRITON_FLASH_ATTN'] = '1'
310
311
        logger.info("Using ROCmFlashAttention backend.")
        return "vllm.attention.backends.rocm_flash_attn.ROCmFlashAttentionBackend"  # noqa: E501
312

313

314
    @classmethod
315
    @lru_cache(maxsize=8)
316
317
318
    def get_device_capability(cls,
                              device_id: int = 0
                              ) -> Optional[DeviceCapability]:
319
320
        major, minor = torch.cuda.get_device_capability(device_id)
        return DeviceCapability(major=major, minor=minor)
321

322
    @classmethod
323
    @with_amdsmi_context
324
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
zhuwenwen's avatar
zhuwenwen committed
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
        """
        Query if the set of gpus are fully connected by xgmi (1 hop)
        """
        handles = [
            amdsmi_get_processor_handles()[i] for i in physical_device_ids
        ]
        for i, handle in enumerate(handles):
            for j, peer_handle in enumerate(handles):
                if i < j:
                    try:
                        link_type = amdsmi_topo_get_link_type(
                            handle, peer_handle)
                        # type is 2 for XGMI
                        if link_type["hops"] != 1 or link_type["type"] != 2:
                            return False
                    except AmdSmiException as error:
                        logger.error("AMD 1 hop XGMI detection failed.",
                                     exc_info=error)
                        return False
        return True
345

346
    @classmethod
347
    @with_amdsmi_context
348
    @lru_cache(maxsize=8)
349
    def get_device_name(cls, device_id: int = 0) -> str:
350
        physical_device_id = device_id_to_physical_device_id(device_id)
351
        handle = amdsmi_get_processor_handles()[physical_device_id]
352
353
        # return amdsmi_get_gpu_asic_info(handle)["market_name"]
        return torch.cuda.get_device_name(device_id)
354

zhuwenwen's avatar
zhuwenwen committed
355
356
357
358
    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        device_props = torch.cuda.get_device_properties(device_id)
        return device_props.total_memory
359

360
361
362
363
364
365
366
367
368
369
    @classmethod
    def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
        if enforce_eager:
            logger.warning(
                "To see benefits of async output processing, enable CUDA "
                "graph. Since, enforce-eager is enabled, async output "
                "processor cannot be used")
            return False
        return True

370
    @classmethod
371
    def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
372
373
374
375
        cache_config = vllm_config.cache_config
        if cache_config and cache_config.block_size is None:
            cache_config.block_size = 16

376
377
378
379
        parallel_config = vllm_config.parallel_config
        scheduler_config = vllm_config.scheduler_config
        if parallel_config.worker_cls == "auto":
            if scheduler_config.is_multi_step:
380
381
382
                if envs.VLLM_USE_V1:
                    raise NotImplementedError(
                        "Multi-step scheduling is not supported (and not "
383
                        "needed) on vLLM V1. Please launch without "
384
385
386
387
                        "--num-scheduler-steps.")
                else:
                    parallel_config.worker_cls = \
                        "vllm.worker.multi_step_worker.MultiStepWorker"
388
            elif vllm_config.speculative_config:
389
                if envs.VLLM_USE_V1:
zhuwenwen's avatar
zhuwenwen committed
390
391
392
393
394
                    # raise NotImplementedError(
                    #     "Speculative decoding is not yet supported on vLLM V1."
                    # )
                    parallel_config.worker_cls = \
                            "vllm.v1.worker.gpu_worker.Worker"
395
396
397
398
399
                else:
                    parallel_config.worker_cls = \
                        "vllm.spec_decode.spec_decode_worker.create_spec_worker"
                    parallel_config.sd_worker_cls = \
                        "vllm.worker.worker.Worker"
400
            else:
401
402
403
404
405
                if envs.VLLM_USE_V1:
                    parallel_config.worker_cls = \
                            "vllm.v1.worker.gpu_worker.Worker"
                else:
                    parallel_config.worker_cls = "vllm.worker.worker.Worker"
406

407
408
409
410
411
412
413
414
415
416
417
418
    @classmethod
    def verify_model_arch(cls, model_arch: str) -> None:
        if model_arch in _ROCM_UNSUPPORTED_MODELS:
            raise ValueError(f"Model architecture '{model_arch}' is not "
                             "supported by ROCm for now.")

        if model_arch in _ROCM_PARTIALLY_SUPPORTED_MODELS:
            msg = _ROCM_PARTIALLY_SUPPORTED_MODELS[model_arch]
            logger.warning(
                "Model architecture '%s' is partially "
                "supported by ROCm: %s", model_arch, msg)

419
420
421
422
423
424
    @classmethod
    def verify_quantization(cls, quant: str) -> None:
        super().verify_quantization(quant)
        if quant == "awq" and not envs.VLLM_USE_TRITON_AWQ:
            logger.warning(
                "Using AWQ quantization with ROCm, but VLLM_USE_TRITON_AWQ"
425
426
                " is not set, disabling VLLM_USE_TRITON_AWQ.")
            envs.VLLM_USE_TRITON_AWQ = False
427
428
429
430

    @classmethod
    def get_punica_wrapper(cls) -> str:
        return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"
431
432
433
434
435
436

    @classmethod
    def get_current_memory_usage(cls,
                                 device: Optional[torch.types.Device] = None
                                 ) -> float:
        torch.cuda.reset_peak_memory_stats(device)
zhuwenwen's avatar
zhuwenwen committed
437
438
439
        # return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(
        #     device)[0]
        return torch.cuda.max_memory_allocated(device)
440
441
442
443

    @classmethod
    def get_device_communicator_cls(cls) -> str:
        return "vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
444

445
446
447
448
449
    @classmethod
    def supports_mx(cls) -> bool:
        gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
        return any(gfx in gcn_arch for gfx in ["gfx95"])

450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
    @classmethod
    def supports_fp8(cls) -> bool:
        gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
        return any(gfx in gcn_arch for gfx in ['gfx94', 'gfx95', 'gfx12'])

    @classmethod
    def is_fp8_fnuz(cls) -> bool:
        # only device 0 is checked, this assumes MI300 platforms are homogeneous
        return 'gfx94' in torch.cuda.get_device_properties(0).gcnArchName

    @classmethod
    def fp8_dtype(cls) -> torch.dtype:
        if cls.is_fp8_fnuz():
            return torch.float8_e4m3fnuz
        else:
            return torch.float8_e4m3fn
466
467

    @classmethod
468
    def supports_v1(cls, model_config: "ModelConfig") -> bool:
469
470
        # V1 support on AMD gpus is experimental
        return True
471
472
473
474
475

    @classmethod
    def use_custom_allreduce(cls) -> bool:
        # We only enable custom allreduce for MI300 series
        gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
476
        supported_archs = ['gfx94', 'gfx95']
477
        return any(gfx in gcn_arch for gfx in supported_archs)
478
479
480
481

    @classmethod
    def get_cu_count(cls, device_id: int = 0) -> int:
        return torch.cuda.get_device_properties(
482
            device_id).multi_processor_count
483
484
485
486

    @classmethod
    def is_navi(cls) -> bool:
        return 'gfx1' in torch.cuda.get_device_properties(0).gcnArchName
487
488
489
490

    @classmethod
    def get_piecewise_backend_cls(cls) -> str:
        return "vllm.compilation.cuda_piecewise_backend.CUDAPiecewiseBackend"  # noqa
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520

    @classmethod
    def stateless_init_device_torch_dist_pg(
        cls,
        backend: str,
        prefix_store: PrefixStore,
        group_rank: int,
        group_size: int,
        timeout: timedelta,
    ) -> ProcessGroup:
        assert is_nccl_available()
        pg: ProcessGroup = ProcessGroup(
            prefix_store,
            group_rank,
            group_size,
        )
        from torch.distributed.distributed_c10d import ProcessGroupNCCL

        backend_options = ProcessGroupNCCL.Options()
        backend_options._timeout = timeout

        backend_class = ProcessGroupNCCL(prefix_store, group_rank, group_size,
                                         backend_options)
        backend_type = ProcessGroup.BackendType.NCCL
        device = torch.device("cuda")
        pg._set_default_backend(backend_type)
        backend_class._set_sequence_number_for_group()

        pg._register_backend(device, backend_type, backend_class)
        return pg
521
522
523
524

    @classmethod
    def device_count(cls) -> int:
        return cuda_device_count_stateless()