Method compute_capability_args

op_builder/builder.py:635–753

Returns nvcc compute capability compile flags. 1. Under ``jit_mode``, the precedence is: a. preserved ``TORCH_CUDA_ARCH_LIST`` captured by ``jit_load()`` b. live ``TORCH_CUDA_ARCH_LIST`` from the environment c. runtime device probing when the proces

(self, cross_compile_archs=None)

        return torch.cuda.get_device_properties(0).major

    def compute_capability_args(self, cross_compile_archs=None):
        """
        Returns nvcc compute capability compile flags.

        1. Under ``jit_mode``, the precedence is:
           a. preserved ``TORCH_CUDA_ARCH_LIST`` captured by ``jit_load()``
           b. live ``TORCH_CUDA_ARCH_LIST`` from the environment
           c. runtime device probing when the process is not in a bad-fork context
           d. an error when no explicit arch list exists in a bad-fork context

           JIT mode auto-adds ``+PTX`` to the highest compute capability when
           no entry already carries it, then sets ``TORCH_CUDA_ARCH_LIST`` so
           PyTorch can generate the ``-gencode`` flags itself.
        2. ``TORCH_CUDA_ARCH_LIST`` takes priority over ``cross_compile_archs``.
        3. If neither is set default compute capabilities will be used.

        Format:

        - ``TORCH_CUDA_ARCH_LIST`` may use ; or whitespace separators. Examples:

        TORCH_CUDA_ARCH_LIST="6.1;7.5;8.6;9.0;10.0" pip install ...
        TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 9.0 10.0+PTX" pip install ...

        - ``cross_compile_archs`` uses ; separator.

        """
        ccs = []
        if self.jit_mode:
            arch_string = getattr(self, '_jit_arch_list', None)
            if arch_string:
                arch_string = arch_string.replace(' ', ';')
                ccs = [cc.strip() for cc in arch_string.split(';') if cc.strip()]
            else:
                arch_string = os.environ.get('TORCH_CUDA_ARCH_LIST', '').strip()
                if arch_string:
                    arch_string = arch_string.replace(' ', ';')
                    ccs = [cc.strip() for cc in arch_string.split(';') if cc.strip()]
                else:
                    if hasattr(torch.cuda, '_is_in_bad_fork') and torch.cuda._is_in_bad_fork():
                        raise RuntimeError(
                            f"DeepSpeed JIT builder for '{self.name}' cannot probe CUDA device capabilities "
                            "in a forked subprocess where CUDA has already been initialized. Set "
                            "TORCH_CUDA_ARCH_LIST to specify target architectures explicitly.")
                    for i in range(torch.cuda.device_count()):
                        CC_MAJOR, CC_MINOR = torch.cuda.get_device_capability(i)
                        cc = f"{CC_MAJOR}.{CC_MINOR}"
                        if cc not in ccs:
                            ccs.append(cc)
                    if len(ccs) == 0:
                        raise RuntimeError(f"DeepSpeed JIT builder for '{self.name}' found no CUDA devices. Set "
                                           "TORCH_CUDA_ARCH_LIST or make GPUs visible.")

            ccs = sorted(ccs, key=lambda cc: tuple(int(part.split('+')[0]) for part in cc.split('.')))
            if not any('+PTX' in cc for cc in ccs):
                ccs[-1] += '+PTX'
        else:
            # Cross-compile mode, compile for various architectures
            # env override takes priority