sparse_attn.py 3 KB
Newer Older
Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
1
2
3
"""
Copyright 2020 The Microsoft DeepSpeed Team
"""
4
5
from .builder import OpBuilder

aiss's avatar
aiss committed
6
7
8
9
10
try:
    from packaging import version as pkg_version
except ImportError:
    pkg_version = None

11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

class SparseAttnBuilder(OpBuilder):
    BUILD_VAR = "DS_BUILD_SPARSE_ATTN"
    NAME = "sparse_attn"

    def __init__(self):
        super().__init__(name=self.NAME)

    def absolute_name(self):
        return f'deepspeed.ops.sparse_attention.{self.NAME}_op'

    def sources(self):
        return ['csrc/sparse_attention/utils.cpp']

    def cxx_args(self):
        return ['-O2', '-fopenmp']

aiss's avatar
aiss committed
28
    def is_compatible(self, verbose=True):
29
        # Check to see if llvm and cmake are installed since they are dependencies
aiss's avatar
aiss committed
30
31
32
33
34
        #required_commands = ['llvm-config|llvm-config-9', 'cmake']
        #command_status = list(map(self.command_exists, required_commands))
        #deps_compatible = all(command_status)

        if self.is_rocm_pytorch():
aiss's avatar
aiss committed
35
36
37
        #aiss debug
            #self.warning(f'{self.NAME} is not compatible with ROCM')
            #return False
aiss's avatar
aiss committed
38
            return True
aiss's avatar
aiss committed
39

aiss's avatar
aiss committed
40
41
42
43
44
        try:
            import torch
        except ImportError:
            self.warning(f"unable to import torch, please install it first")
            return False
45

46
47
48
49
50
51
        # torch-cpu will not have a cuda version
        if torch.version.cuda is None:
            cuda_compatible = False
            self.warning(f"{self.NAME} cuda is not available from torch")
        else:
            major, minor = torch.version.cuda.split('.')[:2]
aiss's avatar
aiss committed
52
53
            cuda_compatible = (int(major) == 10
                               and int(minor) >= 1) or (int(major) >= 11)
54
            if not cuda_compatible:
aiss's avatar
aiss committed
55
                self.warning(f"{self.NAME} requires CUDA version 10.1+")
56

57
58
59
60
61
62
63
64
        TORCH_MAJOR = int(torch.__version__.split('.')[0])
        TORCH_MINOR = int(torch.__version__.split('.')[1])
        torch_compatible = TORCH_MAJOR == 1 and TORCH_MINOR >= 5
        if not torch_compatible:
            self.warning(
                f'{self.NAME} requires a torch version >= 1.5 but detected {TORCH_MAJOR}.{TORCH_MINOR}'
            )

aiss's avatar
aiss committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
        try:
            import triton
        except ImportError:
            # auto-install of triton is broken on some systems, reverting to manual install for now
            # see this issue: https://github.com/microsoft/DeepSpeed/issues/1710
            self.warning(
                f"please install triton==1.0.0 if you want to use sparse attention")
            return False

        if pkg_version:
            installed_triton = pkg_version.parse(triton.__version__)
            triton_mismatch = installed_triton != pkg_version.parse("1.0.0")
        else:
            installed_triton = triton.__version__
            triton_mismatch = installed_triton != "1.0.0"

        if triton_mismatch:
            self.warning(
                f"using untested triton version ({installed_triton}), only 1.0.0 is known to be compatible"
            )
            return False

        return super().is_compatible(verbose) and torch_compatible and cuda_compatible