We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[FEATURE]: Make CUDA extensions default to optional
### Describe the feature

When installing `Colossal-AI` with `pip install .` from source or `pip install colossalai` from pypi, we always build CUDA extensions by default. However, this brings many problems for users who do not have a well-built environment with CUDA and torch. As a result, they failed at the first step, which is installation. I would suggest not to build extensions by default. The users can build extensions in two ways:
1. export `CUDA_EXT=1` when running `pip install`
2. build CUDA kernels only when the kernel is used

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `setup.py`
Content:
```
1 import os
2 import re
3 
4 from setuptools import find_packages, setup
5 
6 from op_builder.utils import get_cuda_bare_metal_version
7 
8 try:
9     import torch
10     from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
11     print("\n\ntorch.__version__  = {}\n\n".format(torch.__version__))
12     TORCH_MAJOR = int(torch.__version__.split('.')[0])
13     TORCH_MINOR = int(torch.__version__.split('.')[1])
14 
15     if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 10):
16         raise RuntimeError("Colossal-AI requires Pytorch 1.10 or newer.\n"
17                            "The latest stable release can be obtained from https://pytorch.org/")
18     TORCH_AVAILABLE = True
19 except ImportError:
20     TORCH_AVAILABLE = False
21 
22 
23 # ninja build does not work unless include_dirs are abs path
24 this_dir = os.path.dirname(os.path.abspath(__file__))
25 build_cuda_ext = True
26 ext_modules = []
27 
28 if int(os.environ.get('NO_CUDA_EXT', '0')) == 1 or not TORCH_AVAILABLE:
29     build_cuda_ext = False
30 
31 
32 def check_cuda_torch_binary_vs_bare_metal(cuda_dir):
33     raw_output, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(cuda_dir)
34     torch_binary_major = torch.version.cuda.split(".")[0]
35     torch_binary_minor = torch.version.cuda.split(".")[1]
36 
37     print("\nCompiling cuda extensions with")
38     print(raw_output + "from " + cuda_dir + "/bin\n")
39 
40     if bare_metal_major != torch_binary_major:
41         print(f'The detected CUDA version ({raw_output}) mismatches the version that was used to compile PyTorch '
42               f'({torch.version.cuda}). CUDA extension will not be installed.')
43         return False
44 
45     if bare_metal_minor != torch_binary_minor:
46         print("\nWarning: Cuda extensions are being compiled with a version of Cuda that does "
47               "not match the version used to compile Pytorch binaries.  "
48               f"Pytorch binaries were compiled with Cuda {torch.version.cuda}.\n"
49               "In some cases, a minor-version mismatch will not cause later errors:  "
50               "https://github.com/NVIDIA/apex/pull/323#discussion_r287021798. ")
51     return True
52 
53 
54 def check_cuda_availability(cuda_dir):
55     if not torch.cuda.is_available():
56         # https://github.com/NVIDIA/apex/issues/486
57         # Extension builds after https://github.com/pytorch/pytorch/pull/23408 attempt to query
58         # torch.cuda.get_device_capability(), which will fail if you are compiling in an environment
59         # without visible GPUs (e.g. during an nvidia-docker build command).
60         print(
61             '\nWarning: Torch did not find available GPUs on this system.\n',
62             'If your intention is to cross-compile, this is not an error.\n'
63             'By default, Colossal-AI will cross-compile for Pascal (compute capabilities 6.0, 6.1, 6.2),\n'
64             'Volta (compute capability 7.0), Turing (compute capability 7.5),\n'
65             'and, if the CUDA version is >= 11.0, Ampere (compute capability 8.0).\n'
66             'If you wish to cross-compile for a single specific architecture,\n'
67             'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n')
68         if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None:
69             _, bare_metal_major, _ = get_cuda_bare_metal_version(cuda_dir)
70             if int(bare_metal_major) == 11:
71                 os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0"
72             else:
73                 os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5"
74         return False
75 
76     if cuda_dir is None:
77         print("nvcc was not found. CUDA extension will not be installed. If you're installing within a container from "
78               "https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
79         return False
80     return True
81 
82 
83 def append_nvcc_threads(nvcc_extra_args):
84     _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
85     if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
86         return nvcc_extra_args + ["--threads", "4"]
87     return nvcc_extra_args
88 
89 
90 def fetch_requirements(path):
91     with open(path, 'r') as fd:
92         return [r.strip() for r in fd.readlines()]
93 
94 
95 def fetch_readme():
96     with open('README.md', encoding='utf-8') as f:
97         return f.read()
98 
99 
100 def get_version():
101     setup_file_path = os.path.abspath(__file__)
102     project_path = os.path.dirname(setup_file_path)
103     version_txt_path = os.path.join(project_path, 'version.txt')
104     version_py_path = os.path.join(project_path, 'colossalai/version.py')
105 
106     with open(version_txt_path) as f:
107         version = f.read().strip()
108         if build_cuda_ext:
109             torch_version = '.'.join(torch.__version__.split('.')[:2])
110             cuda_version = '.'.join(get_cuda_bare_metal_version(CUDA_HOME)[1:])
111             version += f'+torch{torch_version}cu{cuda_version}'
112 
113     # write version into version.py
114     with open(version_py_path, 'w') as f:
115         f.write(f"__version__ = '{version}'\n")
116 
117     return version
118 
119 
120 if build_cuda_ext:
121     build_cuda_ext = check_cuda_availability(CUDA_HOME) and check_cuda_torch_binary_vs_bare_metal(CUDA_HOME)
122 
123 if build_cuda_ext:
124     # Set up macros for forward/backward compatibility hack around
125     # https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e
126     # and
127     # https://github.com/NVIDIA/apex/issues/456
128     # https://github.com/pytorch/pytorch/commit/eb7b39e02f7d75c26d8a795ea8c7fd911334da7e#diff-4632522f237f1e4e728cb824300403ac
129     version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']
130 
131     def cuda_ext_helper(name, sources, extra_cuda_flags, extra_cxx_flags=[]):
132         return CUDAExtension(
133             name=name,
134             sources=[os.path.join('colossalai/kernel/cuda_native/csrc', path) for path in sources],
135             include_dirs=[os.path.join(this_dir, 'colossalai/kernel/cuda_native/csrc/kernels/include')],
136             extra_compile_args={
137                 'cxx': ['-O3'] + version_dependent_macros + extra_cxx_flags,
138                 'nvcc': append_nvcc_threads(['-O3', '--use_fast_math'] + version_dependent_macros + extra_cuda_flags)
139             })
140 
141     #### fused optim kernels ###
142     from op_builder import FusedOptimBuilder
143     ext_modules.append(FusedOptimBuilder().builder('colossalai._C.fused_optim'))
144 
145     #### N-D parallel kernels ###
146     cc_flag = []
147     for arch in torch.cuda.get_arch_list():
148         res = re.search(r'sm_(\d+)', arch)
149         if res:
150             arch_cap = res[1]
151             if int(arch_cap) >= 60:
152                 cc_flag.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}'])
153 
154     extra_cuda_flags = [
155         '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '--expt-relaxed-constexpr',
156         '--expt-extended-lambda'
157     ]
158 
159     from op_builder import ScaledSoftmaxBuilder
160     ext_modules.append(ScaledSoftmaxBuilder().builder('colossalai._C.scaled_upper_triang_masked_softmax'))
161 
162     ext_modules.append(
163         cuda_ext_helper('colossalai._C.scaled_masked_softmax',
164                         ['scaled_masked_softmax.cpp', 'scaled_masked_softmax_cuda.cu'], extra_cuda_flags + cc_flag))
165 
166     from op_builder import MOEBuilder
167     ext_modules.append(MOEBuilder().builder('colossalai._C.moe'))
168 
169     extra_cuda_flags = ['-maxrregcount=50']
170 
171     ext_modules.append(
172         cuda_ext_helper('colossalai._C.layer_norm', ['layer_norm_cuda.cpp', 'layer_norm_cuda_kernel.cu'],
173                         extra_cuda_flags + cc_flag))
174 
175     ### MultiHeadAttn Kernel ####
176     from op_builder import MultiHeadAttnBuilder
177     ext_modules.append(MultiHeadAttnBuilder().builder('colossalai._C.multihead_attention'))
178 
179     ### Gemini Adam kernel ####
180     from op_builder import CPUAdamBuilder
181     ext_modules.append(CPUAdamBuilder().builder('colossalai._C.cpu_optim'))
182 
183 setup(name='colossalai',
184       version=get_version(),
185       packages=find_packages(exclude=(
186           'benchmark',
187           'docker',
188           'tests',
189           'docs',
190           'examples',
191           'tests',
192           'scripts',
193           'requirements',
194           '*.egg-info',
195       )),
196       description='An integrated large-scale model training system with efficient parallelization techniques',
197       long_description=fetch_readme(),
198       long_description_content_type='text/markdown',
199       license='Apache Software License 2.0',
200       url='https://www.colossalai.org',
201       project_urls={
202           'Forum': 'https://github.com/hpcaitech/ColossalAI/discussions',
203           'Bug Tracker': 'https://github.com/hpcaitech/ColossalAI/issues',
204           'Examples': 'https://github.com/hpcaitech/ColossalAI-Examples',
205           'Documentation': 'http://colossalai.readthedocs.io',
206           'Github': 'https://github.com/hpcaitech/ColossalAI',
207       },
208       ext_modules=ext_modules,
209       cmdclass={'build_ext': BuildExtension} if ext_modules else {},
210       install_requires=fetch_requirements('requirements/requirements.txt'),
211       entry_points='''
212         [console_scripts]
213         colossalai=colossalai.cli:cli
214     ''',
215       python_requires='>=3.6',
216       classifiers=[
217           'Programming Language :: Python :: 3',
218           'License :: OSI Approved :: Apache Software License',
219           'Environment :: GPU :: NVIDIA CUDA',
220           'Topic :: Scientific/Engineering :: Artificial Intelligence',
221           'Topic :: System :: Distributed Computing',
222       ],
223       package_data={'colossalai': ['_C/*.pyi']})
224
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -18,15 +18,22 @@
     TORCH_AVAILABLE = True
 except ImportError:
     TORCH_AVAILABLE = False
+    CUDA_HOME = None
 
 
 # ninja build does not work unless include_dirs are abs path
 this_dir = os.path.dirname(os.path.abspath(__file__))
-build_cuda_ext = True
+build_cuda_ext = False
 ext_modules = []
 
-if int(os.environ.get('NO_CUDA_EXT', '0')) == 1 or not TORCH_AVAILABLE:
-    build_cuda_ext = False
+if int(os.environ.get('CUDA_EXT', '0')) == 1:
+    if not TORCH_AVAILABLE:
+        raise ModuleNotFoundError("PyTorch is not found while CUDA_EXT=1. You need to install PyTorch first in order to build CUDA extensions")
+
+    if not CUDA_HOME:
+        raise RuntimeError("CUDA_HOME is not found while CUDA_EXT=1. You need to export CUDA_HOME environment vairable or install CUDA Toolkit first in order to build CUDA extensions")
+
+    build_cuda_ext = True
 
 
 def check_cuda_torch_binary_vs_bare_metal(cuda_dir):

{"golden_diff": "diff --git a/setup.py b/setup.py\n--- a/setup.py\n+++ b/setup.py\n@@ -18,15 +18,22 @@\n     TORCH_AVAILABLE = True\n except ImportError:\n     TORCH_AVAILABLE = False\n+    CUDA_HOME = None\n \n \n # ninja build does not work unless include_dirs are abs path\n this_dir = os.path.dirname(os.path.abspath(__file__))\n-build_cuda_ext = True\n+build_cuda_ext = False\n ext_modules = []\n \n-if int(os.environ.get('NO_CUDA_EXT', '0')) == 1 or not TORCH_AVAILABLE:\n-    build_cuda_ext = False\n+if int(os.environ.get('CUDA_EXT', '0')) == 1:\n+    if not TORCH_AVAILABLE:\n+        raise ModuleNotFoundError(\"PyTorch is not found while CUDA_EXT=1. You need to install PyTorch first in order to build CUDA extensions\")\n+\n+    if not CUDA_HOME:\n+        raise RuntimeError(\"CUDA_HOME is not found while CUDA_EXT=1. You need to export CUDA_HOME environment vairable or install CUDA Toolkit first in order to build CUDA extensions\")\n+\n+    build_cuda_ext = True\n \n \n def check_cuda_torch_binary_vs_bare_metal(cuda_dir):\n", "issue": "[FEATURE]: Make CUDA extensions default to optional\n### Describe the feature\n\nWhen installing `Colossal-AI` with `pip install .` from source or `pip install colossalai` from pypi, we always build CUDA extensions by default. However, this brings many problems for users who do not have a well-built environment with CUDA and torch. As a result, they failed at the first step, which is installation. I would suggest not to build extensions by default. The users can build extensions in two ways:\r\n1. export `CUDA_EXT=1` when running `pip install`\r\n2. build CUDA kernels only when the kernel is used\n", "before_files": [{"content": "import os\nimport re\n\nfrom setuptools import find_packages, setup\n\nfrom op_builder.utils import get_cuda_bare_metal_version\n\ntry:\n    import torch\n    from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension\n    print(\"\\n\\ntorch.__version__  = {}\\n\\n\".format(torch.__version__))\n    TORCH_MAJOR = int(torch.__version__.split('.')[0])\n    TORCH_MINOR = int(torch.__version__.split('.')[1])\n\n    if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 10):\n        raise RuntimeError(\"Colossal-AI requires Pytorch 1.10 or newer.\\n\"\n                           \"The latest stable release can be obtained from https://pytorch.org/\")\n    TORCH_AVAILABLE = True\nexcept ImportError:\n    TORCH_AVAILABLE = False\n\n\n# ninja build does not work unless include_dirs are abs path\nthis_dir = os.path.dirname(os.path.abspath(__file__))\nbuild_cuda_ext = True\next_modules = []\n\nif int(os.environ.get('NO_CUDA_EXT', '0')) == 1 or not TORCH_AVAILABLE:\n    build_cuda_ext = False\n\n\ndef check_cuda_torch_binary_vs_bare_metal(cuda_dir):\n    raw_output, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(cuda_dir)\n    torch_binary_major = torch.version.cuda.split(\".\")[0]\n    torch_binary_minor = torch.version.cuda.split(\".\")[1]\n\n    print(\"\\nCompiling cuda extensions with\")\n    print(raw_output + \"from \" + cuda_dir + \"/bin\\n\")\n\n    if bare_metal_major != torch_binary_major:\n        print(f'The detected CUDA version ({raw_output}) mismatches the version that was used to compile PyTorch '\n              f'({torch.version.cuda}). CUDA extension will not be installed.')\n        return False\n\n    if bare_metal_minor != torch_binary_minor:\n        print(\"\\nWarning: Cuda extensions are being compiled with a version of Cuda that does \"\n              \"not match the version used to compile Pytorch binaries.  \"\n              f\"Pytorch binaries were compiled with Cuda {torch.version.cuda}.\\n\"\n              \"In some cases, a minor-version mismatch will not cause later errors:  \"\n              \"https://github.com/NVIDIA/apex/pull/323#discussion_r287021798. \")\n    return True\n\n\ndef check_cuda_availability(cuda_dir):\n    if not torch.cuda.is_available():\n        # https://github.com/NVIDIA/apex/issues/486\n        # Extension builds after https://github.com/pytorch/pytorch/pull/23408 attempt to query\n        # torch.cuda.get_device_capability(), which will fail if you are compiling in an environment\n        # without visible GPUs (e.g. during an nvidia-docker build command).\n        print(\n            '\\nWarning: Torch did not find available GPUs on this system.\\n',\n            'If your intention is to cross-compile, this is not an error.\\n'\n            'By default, Colossal-AI will cross-compile for Pascal (compute capabilities 6.0, 6.1, 6.2),\\n'\n            'Volta (compute capability 7.0), Turing (compute capability 7.5),\\n'\n            'and, if the CUDA version is >= 11.0, Ampere (compute capability 8.0).\\n'\n            'If you wish to cross-compile for a single specific architecture,\\n'\n            'export TORCH_CUDA_ARCH_LIST=\"compute capability\" before running setup.py.\\n')\n        if os.environ.get(\"TORCH_CUDA_ARCH_LIST\", None) is None:\n            _, bare_metal_major, _ = get_cuda_bare_metal_version(cuda_dir)\n            if int(bare_metal_major) == 11:\n                os.environ[\"TORCH_CUDA_ARCH_LIST\"] = \"6.0;6.1;6.2;7.0;7.5;8.0\"\n            else:\n                os.environ[\"TORCH_CUDA_ARCH_LIST\"] = \"6.0;6.1;6.2;7.0;7.5\"\n        return False\n\n    if cuda_dir is None:\n        print(\"nvcc was not found. CUDA extension will not be installed. If you're installing within a container from \"\n              \"https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.\")\n        return False\n    return True\n\n\ndef append_nvcc_threads(nvcc_extra_args):\n    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)\n    if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:\n        return nvcc_extra_args + [\"--threads\", \"4\"]\n    return nvcc_extra_args\n\n\ndef fetch_requirements(path):\n    with open(path, 'r') as fd:\n        return [r.strip() for r in fd.readlines()]\n\n\ndef fetch_readme():\n    with open('README.md', encoding='utf-8') as f:\n        return f.read()\n\n\ndef get_version():\n    setup_file_path = os.path.abspath(__file__)\n    project_path = os.path.dirname(setup_file_path)\n    version_txt_path = os.path.join(project_path, 'version.txt')\n    version_py_path = os.path.join(project_path, 'colossalai/version.py')\n\n    with open(version_txt_path) as f:\n        version = f.read().strip()\n        if build_cuda_ext:\n            torch_version = '.'.join(torch.__version__.split('.')[:2])\n            cuda_version = '.'.join(get_cuda_bare_metal_version(CUDA_HOME)[1:])\n            version += f'+torch{torch_version}cu{cuda_version}'\n\n    # write version into version.py\n    with open(version_py_path, 'w') as f:\n        f.write(f\"__version__ = '{version}'\\n\")\n\n    return version\n\n\nif build_cuda_ext:\n    build_cuda_ext = check_cuda_availability(CUDA_HOME) and check_cuda_torch_binary_vs_bare_metal(CUDA_HOME)\n\nif build_cuda_ext:\n    # Set up macros for forward/backward compatibility hack around\n    # https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e\n    # and\n    # https://github.com/NVIDIA/apex/issues/456\n    # https://github.com/pytorch/pytorch/commit/eb7b39e02f7d75c26d8a795ea8c7fd911334da7e#diff-4632522f237f1e4e728cb824300403ac\n    version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']\n\n    def cuda_ext_helper(name, sources, extra_cuda_flags, extra_cxx_flags=[]):\n        return CUDAExtension(\n            name=name,\n            sources=[os.path.join('colossalai/kernel/cuda_native/csrc', path) for path in sources],\n            include_dirs=[os.path.join(this_dir, 'colossalai/kernel/cuda_native/csrc/kernels/include')],\n            extra_compile_args={\n                'cxx': ['-O3'] + version_dependent_macros + extra_cxx_flags,\n                'nvcc': append_nvcc_threads(['-O3', '--use_fast_math'] + version_dependent_macros + extra_cuda_flags)\n            })\n\n    #### fused optim kernels ###\n    from op_builder import FusedOptimBuilder\n    ext_modules.append(FusedOptimBuilder().builder('colossalai._C.fused_optim'))\n\n    #### N-D parallel kernels ###\n    cc_flag = []\n    for arch in torch.cuda.get_arch_list():\n        res = re.search(r'sm_(\\d+)', arch)\n        if res:\n            arch_cap = res[1]\n            if int(arch_cap) >= 60:\n                cc_flag.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}'])\n\n    extra_cuda_flags = [\n        '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '--expt-relaxed-constexpr',\n        '--expt-extended-lambda'\n    ]\n\n    from op_builder import ScaledSoftmaxBuilder\n    ext_modules.append(ScaledSoftmaxBuilder().builder('colossalai._C.scaled_upper_triang_masked_softmax'))\n\n    ext_modules.append(\n        cuda_ext_helper('colossalai._C.scaled_masked_softmax',\n                        ['scaled_masked_softmax.cpp', 'scaled_masked_softmax_cuda.cu'], extra_cuda_flags + cc_flag))\n\n    from op_builder import MOEBuilder\n    ext_modules.append(MOEBuilder().builder('colossalai._C.moe'))\n\n    extra_cuda_flags = ['-maxrregcount=50']\n\n    ext_modules.append(\n        cuda_ext_helper('colossalai._C.layer_norm', ['layer_norm_cuda.cpp', 'layer_norm_cuda_kernel.cu'],\n                        extra_cuda_flags + cc_flag))\n\n    ### MultiHeadAttn Kernel ####\n    from op_builder import MultiHeadAttnBuilder\n    ext_modules.append(MultiHeadAttnBuilder().builder('colossalai._C.multihead_attention'))\n\n    ### Gemini Adam kernel ####\n    from op_builder import CPUAdamBuilder\n    ext_modules.append(CPUAdamBuilder().builder('colossalai._C.cpu_optim'))\n\nsetup(name='colossalai',\n      version=get_version(),\n      packages=find_packages(exclude=(\n          'benchmark',\n          'docker',\n          'tests',\n          'docs',\n          'examples',\n          'tests',\n          'scripts',\n          'requirements',\n          '*.egg-info',\n      )),\n      description='An integrated large-scale model training system with efficient parallelization techniques',\n      long_description=fetch_readme(),\n      long_description_content_type='text/markdown',\n      license='Apache Software License 2.0',\n      url='https://www.colossalai.org',\n      project_urls={\n          'Forum': 'https://github.com/hpcaitech/ColossalAI/discussions',\n          'Bug Tracker': 'https://github.com/hpcaitech/ColossalAI/issues',\n          'Examples': 'https://github.com/hpcaitech/ColossalAI-Examples',\n          'Documentation': 'http://colossalai.readthedocs.io',\n          'Github': 'https://github.com/hpcaitech/ColossalAI',\n      },\n      ext_modules=ext_modules,\n      cmdclass={'build_ext': BuildExtension} if ext_modules else {},\n      install_requires=fetch_requirements('requirements/requirements.txt'),\n      entry_points='''\n        [console_scripts]\n        colossalai=colossalai.cli:cli\n    ''',\n      python_requires='>=3.6',\n      classifiers=[\n          'Programming Language :: Python :: 3',\n          'License :: OSI Approved :: Apache Software License',\n          'Environment :: GPU :: NVIDIA CUDA',\n          'Topic :: Scientific/Engineering :: Artificial Intelligence',\n          'Topic :: System :: Distributed Computing',\n      ],\n      package_data={'colossalai': ['_C/*.pyi']})\n", "path": "setup.py"}], "after_files": [{"content": "import os\nimport re\n\nfrom setuptools import find_packages, setup\n\nfrom op_builder.utils import get_cuda_bare_metal_version\n\ntry:\n    import torch\n    from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension\n    print(\"\\n\\ntorch.__version__  = {}\\n\\n\".format(torch.__version__))\n    TORCH_MAJOR = int(torch.__version__.split('.')[0])\n    TORCH_MINOR = int(torch.__version__.split('.')[1])\n\n    if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 10):\n        raise RuntimeError(\"Colossal-AI requires Pytorch 1.10 or newer.\\n\"\n                           \"The latest stable release can be obtained from https://pytorch.org/\")\n    TORCH_AVAILABLE = True\nexcept ImportError:\n    TORCH_AVAILABLE = False\n    CUDA_HOME = None\n\n\n# ninja build does not work unless include_dirs are abs path\nthis_dir = os.path.dirname(os.path.abspath(__file__))\nbuild_cuda_ext = False\next_modules = []\n\nif int(os.environ.get('CUDA_EXT', '0')) == 1:\n    if not TORCH_AVAILABLE:\n        raise ModuleNotFoundError(\"PyTorch is not found while CUDA_EXT=1. You need to install PyTorch first in order to build CUDA extensions\")\n\n    if not CUDA_HOME:\n        raise RuntimeError(\"CUDA_HOME is not found while CUDA_EXT=1. You need to export CUDA_HOME environment vairable or install CUDA Toolkit first in order to build CUDA extensions\")\n\n    build_cuda_ext = True\n\n\ndef check_cuda_torch_binary_vs_bare_metal(cuda_dir):\n    raw_output, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(cuda_dir)\n    torch_binary_major = torch.version.cuda.split(\".\")[0]\n    torch_binary_minor = torch.version.cuda.split(\".\")[1]\n\n    print(\"\\nCompiling cuda extensions with\")\n    print(raw_output + \"from \" + cuda_dir + \"/bin\\n\")\n\n    if bare_metal_major != torch_binary_major:\n        print(f'The detected CUDA version ({raw_output}) mismatches the version that was used to compile PyTorch '\n              f'({torch.version.cuda}). CUDA extension will not be installed.')\n        return False\n\n    if bare_metal_minor != torch_binary_minor:\n        print(\"\\nWarning: Cuda extensions are being compiled with a version of Cuda that does \"\n              \"not match the version used to compile Pytorch binaries.  \"\n              f\"Pytorch binaries were compiled with Cuda {torch.version.cuda}.\\n\"\n              \"In some cases, a minor-version mismatch will not cause later errors:  \"\n              \"https://github.com/NVIDIA/apex/pull/323#discussion_r287021798. \")\n    return True\n\n\ndef check_cuda_availability(cuda_dir):\n    if not torch.cuda.is_available():\n        # https://github.com/NVIDIA/apex/issues/486\n        # Extension builds after https://github.com/pytorch/pytorch/pull/23408 attempt to query\n        # torch.cuda.get_device_capability(), which will fail if you are compiling in an environment\n        # without visible GPUs (e.g. during an nvidia-docker build command).\n        print(\n            '\\nWarning: Torch did not find available GPUs on this system.\\n',\n            'If your intention is to cross-compile, this is not an error.\\n'\n            'By default, Colossal-AI will cross-compile for Pascal (compute capabilities 6.0, 6.1, 6.2),\\n'\n            'Volta (compute capability 7.0), Turing (compute capability 7.5),\\n'\n            'and, if the CUDA version is >= 11.0, Ampere (compute capability 8.0).\\n'\n            'If you wish to cross-compile for a single specific architecture,\\n'\n            'export TORCH_CUDA_ARCH_LIST=\"compute capability\" before running setup.py.\\n')\n        if os.environ.get(\"TORCH_CUDA_ARCH_LIST\", None) is None:\n            _, bare_metal_major, _ = get_cuda_bare_metal_version(cuda_dir)\n            if int(bare_metal_major) == 11:\n                os.environ[\"TORCH_CUDA_ARCH_LIST\"] = \"6.0;6.1;6.2;7.0;7.5;8.0\"\n            else:\n                os.environ[\"TORCH_CUDA_ARCH_LIST\"] = \"6.0;6.1;6.2;7.0;7.5\"\n        return False\n\n    if cuda_dir is None:\n        print(\"nvcc was not found. CUDA extension will not be installed. If you're installing within a container from \"\n              \"https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.\")\n        return False\n    return True\n\n\ndef append_nvcc_threads(nvcc_extra_args):\n    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)\n    if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:\n        return nvcc_extra_args + [\"--threads\", \"4\"]\n    return nvcc_extra_args\n\n\ndef fetch_requirements(path):\n    with open(path, 'r') as fd:\n        return [r.strip() for r in fd.readlines()]\n\n\ndef fetch_readme():\n    with open('README.md', encoding='utf-8') as f:\n        return f.read()\n\n\ndef get_version():\n    setup_file_path = os.path.abspath(__file__)\n    project_path = os.path.dirname(setup_file_path)\n    version_txt_path = os.path.join(project_path, 'version.txt')\n    version_py_path = os.path.join(project_path, 'colossalai/version.py')\n\n    with open(version_txt_path) as f:\n        version = f.read().strip()\n        if build_cuda_ext:\n            torch_version = '.'.join(torch.__version__.split('.')[:2])\n            cuda_version = '.'.join(get_cuda_bare_metal_version(CUDA_HOME)[1:])\n            version += f'+torch{torch_version}cu{cuda_version}'\n\n    # write version into version.py\n    with open(version_py_path, 'w') as f:\n        f.write(f\"__version__ = '{version}'\\n\")\n\n    return version\n\n\nif build_cuda_ext:\n    build_cuda_ext = check_cuda_availability(CUDA_HOME) and check_cuda_torch_binary_vs_bare_metal(CUDA_HOME)\n\nif build_cuda_ext:\n    # Set up macros for forward/backward compatibility hack around\n    # https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e\n    # and\n    # https://github.com/NVIDIA/apex/issues/456\n    # https://github.com/pytorch/pytorch/commit/eb7b39e02f7d75c26d8a795ea8c7fd911334da7e#diff-4632522f237f1e4e728cb824300403ac\n    version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']\n\n    def cuda_ext_helper(name, sources, extra_cuda_flags, extra_cxx_flags=[]):\n        return CUDAExtension(\n            name=name,\n            sources=[os.path.join('colossalai/kernel/cuda_native/csrc', path) for path in sources],\n            include_dirs=[os.path.join(this_dir, 'colossalai/kernel/cuda_native/csrc/kernels/include')],\n            extra_compile_args={\n                'cxx': ['-O3'] + version_dependent_macros + extra_cxx_flags,\n                'nvcc': append_nvcc_threads(['-O3', '--use_fast_math'] + version_dependent_macros + extra_cuda_flags)\n            })\n\n    #### fused optim kernels ###\n    from op_builder import FusedOptimBuilder\n    ext_modules.append(FusedOptimBuilder().builder('colossalai._C.fused_optim'))\n\n    #### N-D parallel kernels ###\n    cc_flag = []\n    for arch in torch.cuda.get_arch_list():\n        res = re.search(r'sm_(\\d+)', arch)\n        if res:\n            arch_cap = res[1]\n            if int(arch_cap) >= 60:\n                cc_flag.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}'])\n\n    extra_cuda_flags = [\n        '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '--expt-relaxed-constexpr',\n        '--expt-extended-lambda'\n    ]\n\n    from op_builder import ScaledSoftmaxBuilder\n    ext_modules.append(ScaledSoftmaxBuilder().builder('colossalai._C.scaled_upper_triang_masked_softmax'))\n\n    ext_modules.append(\n        cuda_ext_helper('colossalai._C.scaled_masked_softmax',\n                        ['scaled_masked_softmax.cpp', 'scaled_masked_softmax_cuda.cu'], extra_cuda_flags + cc_flag))\n\n    from op_builder import MOEBuilder\n    ext_modules.append(MOEBuilder().builder('colossalai._C.moe'))\n\n    extra_cuda_flags = ['-maxrregcount=50']\n\n    ext_modules.append(\n        cuda_ext_helper('colossalai._C.layer_norm', ['layer_norm_cuda.cpp', 'layer_norm_cuda_kernel.cu'],\n                        extra_cuda_flags + cc_flag))\n\n    ### MultiHeadAttn Kernel ####\n    from op_builder import MultiHeadAttnBuilder\n    ext_modules.append(MultiHeadAttnBuilder().builder('colossalai._C.multihead_attention'))\n\n    ### Gemini Adam kernel ####\n    from op_builder import CPUAdamBuilder\n    ext_modules.append(CPUAdamBuilder().builder('colossalai._C.cpu_optim'))\n\nsetup(name='colossalai',\n      version=get_version(),\n      packages=find_packages(exclude=(\n          'benchmark',\n          'docker',\n          'tests',\n          'docs',\n          'examples',\n          'tests',\n          'scripts',\n          'requirements',\n          '*.egg-info',\n      )),\n      description='An integrated large-scale model training system with efficient parallelization techniques',\n      long_description=fetch_readme(),\n      long_description_content_type='text/markdown',\n      license='Apache Software License 2.0',\n      url='https://www.colossalai.org',\n      project_urls={\n          'Forum': 'https://github.com/hpcaitech/ColossalAI/discussions',\n          'Bug Tracker': 'https://github.com/hpcaitech/ColossalAI/issues',\n          'Examples': 'https://github.com/hpcaitech/ColossalAI-Examples',\n          'Documentation': 'http://colossalai.readthedocs.io',\n          'Github': 'https://github.com/hpcaitech/ColossalAI',\n      },\n      ext_modules=ext_modules,\n      cmdclass={'build_ext': BuildExtension} if ext_modules else {},\n      install_requires=fetch_requirements('requirements/requirements.txt'),\n      entry_points='''\n        [console_scripts]\n        colossalai=colossalai.cli:cli\n    ''',\n      python_requires='>=3.6',\n      classifiers=[\n          'Programming Language :: Python :: 3',\n          'License :: OSI Approved :: Apache Software License',\n          'Environment :: GPU :: NVIDIA CUDA',\n          'Topic :: Scientific/Engineering :: Artificial Intelligence',\n          'Topic :: System :: Distributed Computing',\n      ],\n      package_data={'colossalai': ['_C/*.pyi']})\n", "path": "setup.py"}]}