update setup and workflow (#222)

.github/workflows/PR_CI.yml

@@ -43,7 +43,7 @@ jobs:
       - name: Install Colossal-AI      
         run: |
           pip install -r requirements/requirements.txt
-          pip install -v --no-cache-dir --global-option="--cuda_ext" .
+          pip install -v --no-cache-dir .
       - name: Unit Testing
         run: |
           pytest tests

README.md

@@ -17,9 +17,26 @@ An integrated large-scale model training system with efficient parallelization t
 
 ## Installation
 
-### Install From Source (Recommended)
+### PyPI
+
+```bash
+pip install colossalai
+```
+This command will install CUDA extension if your have installed CUDA, NVCC and torch. 
 
-> We **recommend** you to install from source as the Colossal-AI is updating frequently in the early versions. The documentation will be in line with the main branch of the repository. Feel free to raise an issue if you encounter any problem. :)
+If you don't want to install CUDA extension, you should add `--global-option="--no_cuda_ext"`, like:
+```bash
+pip install colossalai --global-option="--no_cuda_ext"
+```
+
+If you want to use `ZeRO`, you can run:
+```bash
+pip install colossalai[zero]
+```
+
+### Install From Source
+
+> The documentation will be in line with the main branch of the repository. Feel free to raise an issue if you encounter any problem. :)
 
 ```shell
 git clone https://github.com/hpcaitech/ColossalAI.git
@@ -31,19 +48,12 @@ pip install -r requirements/requirements.txt
 pip install .
 ```
 
-Install and enable CUDA kernel fusion (compulsory installation when using fused optimizer)
+If you don't want to install and enable CUDA kernel fusion (compulsory installation when using fused optimizer):
 
 ```shell
-pip install -v --no-cache-dir --global-option="--cuda_ext" .
+pip install --global-option="--no_cuda_ext" .
 ```
 
-### PyPI
-
-```bash
-pip install colossalai
-```
-
-
 ## Use Docker
 
 Run the following command to build a docker image from Dockerfile provided.

docker/Dockerfile

@@ -8,4 +8,4 @@ RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple \
 # install colossalai
 RUN git clone https://github.com/hpcaitech/ColossalAI.git \
  && cd ./ColossalAI \
- && pip install -v --no-cache-dir --global-option="--cuda_ext" .
+ && pip install -v --no-cache-dir .

requirements/requirements-zero.txt

+deepspeed
\ No newline at end of file

setup.py

@@ -2,12 +2,16 @@ import os
 import subprocess
 import sys
 
-import torch
-from setuptools import setup, find_packages
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
+from setuptools import find_packages, setup
 
 # ninja build does not work unless include_dirs are abs path
 this_dir = os.path.dirname(os.path.abspath(__file__))
+build_cuda_ext = True
+ext_modules = []
+
+if '--no_cuda_ext' in sys.argv:
+    sys.argv.remove('--no_cuda_ext')
+    build_cuda_ext = False
 
 
 def get_cuda_bare_metal_version(cuda_dir):
@@ -29,28 +33,22 @@ def check_cuda_torch_binary_vs_bare_metal(cuda_dir):
     print("\nCompiling cuda extensions with")
     print(raw_output + "from " + cuda_dir + "/bin\n")
 
-    if (bare_metal_major != torch_binary_major) or (bare_metal_minor != torch_binary_minor):
-        raise RuntimeError("Cuda extensions are being compiled with a version of Cuda that does " +
-                           "not match the version used to compile Pytorch binaries.  " +
-                           "Pytorch binaries were compiled with Cuda {}.\n".format(torch.version.cuda) +
-                           "In some cases, a minor-version mismatch will not cause later errors:  " +
-                           "https://github.com/NVIDIA/apex/pull/323#discussion_r287021798.  "
-                           "You can try commenting out this check (at your own risk).")
-
-
-def append_nvcc_threads(nvcc_extra_args):
-    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
-    if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
-        return nvcc_extra_args + ["--threads", "4"]
-    return nvcc_extra_args
-
+    if bare_metal_major != torch_binary_major:
+        print(
+            f'The detected CUDA version ({raw_output}) mismatches the version that was used to compile PyTorch ({torch.version.cuda}). CUDA extension will not be installed.')
+        return False
 
-def fetch_requirements(path):
-    with open(path, 'r') as fd:
-        return [r.strip() for r in fd.readlines()]
+    if bare_metal_minor != torch_binary_minor:
+        print("\nWarning: Cuda extensions are being compiled with a version of Cuda that does "
+              + "not match the version used to compile Pytorch binaries.  "
+              + "Pytorch binaries were compiled with Cuda {}.\n".format(torch.version.cuda)
+              + "In some cases, a minor-version mismatch will not cause later errors:  "
+              + "https://github.com/NVIDIA/apex/pull/323#discussion_r287021798. ")
+    return True
 
 
-if not torch.cuda.is_available():
+def check_cuda_availability(cuda_dir):
+    if not torch.cuda.is_available():
         # https://github.com/NVIDIA/apex/issues/486
         # Extension builds after https://github.com/pytorch/pytorch/pull/23408 attempt to query torch.cuda.get_device_capability(),
         # which will fail if you are compiling in an environment without visible GPUs (e.g. during an nvidia-docker build command).
@@ -62,38 +60,59 @@ if not torch.cuda.is_available():
               'If you wish to cross-compile for a single specific architecture,\n'
               'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n')
         if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None:
-        _, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME)
+            _, bare_metal_major, _ = get_cuda_bare_metal_version(cuda_dir)
             if int(bare_metal_major) == 11:
                 os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0"
             else:
                 os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5"
+        return False
 
-print("\n\ntorch.__version__  = {}\n\n".format(torch.__version__))
-TORCH_MAJOR = int(torch.__version__.split('.')[0])
-TORCH_MINOR = int(torch.__version__.split('.')[1])
+    if cuda_dir is None:
+        print(
+            "nvcc was not found. CUDA extension will not be installed. If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+        return False
+    return True
 
-if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 8):
-    raise RuntimeError("Colossal-AI requires Pytorch 1.8 or newer.\n" +
-                       "The latest stable release can be obtained from https://pytorch.org/")
 
-cmdclass = {}
-ext_modules = []
+def append_nvcc_threads(nvcc_extra_args):
+    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
+    if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
+        return nvcc_extra_args + ["--threads", "4"]
+    return nvcc_extra_args
 
-# Set up macros for forward/backward compatibility hack around
-# https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e
-# and
-# https://github.com/NVIDIA/apex/issues/456
-# https://github.com/pytorch/pytorch/commit/eb7b39e02f7d75c26d8a795ea8c7fd911334da7e#diff-4632522f237f1e4e728cb824300403ac
-version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']
 
-if "--cuda_ext" in sys.argv:
-    sys.argv.remove("--cuda_ext")
+def fetch_requirements(path):
+    with open(path, 'r') as fd:
+        return [r.strip() for r in fd.readlines()]
 
-    if CUDA_HOME is None:
-        raise RuntimeError(
-            "--cuda_ext was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
-    else:
-        check_cuda_torch_binary_vs_bare_metal(CUDA_HOME)
+
+if build_cuda_ext:
+    try:
+        import torch
+        from torch.utils.cpp_extension import (CUDA_HOME, BuildExtension,
+                                               CUDAExtension)
+        print("\n\ntorch.__version__  = {}\n\n".format(torch.__version__))
+        TORCH_MAJOR = int(torch.__version__.split('.')[0])
+        TORCH_MINOR = int(torch.__version__.split('.')[1])
+
+        if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 8):
+            raise RuntimeError("Colossal-AI requires Pytorch 1.8 or newer.\n"
+                               + "The latest stable release can be obtained from https://pytorch.org/")
+    except ImportError:
+        print('torch is not found. CUDA extension will not be installed')
+        build_cuda_ext = False
+
+if build_cuda_ext:
+    build_cuda_ext = check_cuda_availability(CUDA_HOME) and check_cuda_torch_binary_vs_bare_metal(CUDA_HOME)
+
+
+if build_cuda_ext:
+    # Set up macros for forward/backward compatibility hack around
+    # https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e
+    # and
+    # https://github.com/NVIDIA/apex/issues/456
+    # https://github.com/pytorch/pytorch/commit/eb7b39e02f7d75c26d8a795ea8c7fd911334da7e#diff-4632522f237f1e4e728cb824300403ac
+    version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']
 
     def cuda_ext_helper(name, sources, extra_cuda_flags):
         return CUDAExtension(name=name,
@@ -156,12 +175,9 @@ if "--cuda_ext" in sys.argv:
                                         'kernels/cuda_util.cu'],
                                        extra_cuda_flags + cc_flag))
 
-
-install_requires = fetch_requirements('requirements/requirements.txt')
-
 setup(
     name='colossalai',
-    version='0.0.1-beta',
+    version='0.0.2',
     packages=find_packages(exclude=('benchmark',
                                     'docker',
                                     'tests',
@@ -174,5 +190,8 @@ setup(
     description='An integrated large-scale model training system with efficient parallelization techniques',
     ext_modules=ext_modules,
     cmdclass={'build_ext': BuildExtension} if ext_modules else {},
-    install_requires=install_requires,
+    install_requires=fetch_requirements('requirements/requirements.txt'),
+    extras_require={
+        'zero': fetch_requirements('requirements/requirements-zero.txt'),
+    }
 )