Add torch-latest and torch-nightly CI workflows

.github/workflows/nv-lightning-v100.yml

@@ -17,7 +17,7 @@ concurrency:
 
 jobs:
   unit-tests:
-    runs-on: [self-hosted, nvidia, torch18, v100]
+    runs-on: [self-hosted, nvidia, cu111, v100]
 
     steps:
       - uses: actions/checkout@v2
@@ -29,16 +29,20 @@ jobs:
           python --version
           which nvcc
           nvcc --version
+          pip install --upgrade pip
+          pip uninstall --yes torch torchvision
           pip install torch==1.8.2+cu111 torchvision==0.9.2+cu111 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
           python -c "import torch; print('torch:', torch.__version__, torch)"
           python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
       - name: Install deepspeed
         run: |
+          pip uninstall --yes deepspeed
           pip install .[dev,autotuning]
           ds_report
       - name: PyTorch Lightning Tests
         run: |
           if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
+          pip uninstall --yes pytorch-lightning
           pip install pytorch-lightning
           pip install "protobuf<4.21.0"
           cd tests

.github/workflows/nv-torch-latest-v100.yml

@@ -0,0 +1,60 @@
+name: nv-torch-latest-v100
+
+on:
+  push:
+    branches:
+      - 'master'
+      - 'staging**'
+    paths-ignore:
+      - 'docs/**'
+  pull_request:
+    paths-ignore:
+      - 'docs/**'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  unit-tests:
+    runs-on: [self-hosted, nvidia, cu113, v100]
+
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: environment
+        run: |
+          nvidia-smi
+          which python
+          python --version
+          which nvcc
+          nvcc --version
+          pip install --upgrade pip
+          pip uninstall --yes torch torchvision
+          pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113
+          python -c "import torch; print('torch:', torch.__version__, torch)"
+          python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
+
+      - name: Install transformers
+        run: |
+          git clone https://github.com/huggingface/transformers
+          cd transformers
+          # if needed switch to the last known good SHA until transformers@master is fixed
+          # git checkout 1cc453d33
+          git rev-parse --short HEAD
+          pip uninstall --yes transformers
+          pip install .
+
+      - name: Install deepspeed
+        run: |
+          pip uninstall --yes deepspeed
+          pip install .[dev,1bit,autotuning,sparse_attn]
+          ds_report
+
+      - name: Unit tests
+        run: |
+          unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
+          if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
+          cd tests
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 -m 'not sequential' unit/
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/

.github/workflows/nv-torch-nightly-v100.yml

@@ -0,0 +1,53 @@
+name: nv-torch-nightly-v100
+
+on:
+  schedule:
+    - cron: "0 0 * * *"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  unit-tests:
+    runs-on: [self-hosted, nvidia, cu113, v100]
+
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: environment
+        run: |
+          nvidia-smi
+          which python
+          python --version
+          which nvcc
+          nvcc --version
+          pip install --upgrade pip
+          pip uninstall --yes torch torchvision
+          pip install --pre torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cu113
+          python -c "import torch; print('torch:', torch.__version__, torch)"
+          python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
+
+      - name: Install transformers
+        run: |
+          git clone https://github.com/huggingface/transformers
+          cd transformers
+          # if needed switch to the last known good SHA until transformers@master is fixed
+          # git checkout 1cc453d33
+          git rev-parse --short HEAD
+          pip uninstall --yes transformers
+          pip install .
+
+      - name: Install deepspeed
+        run: |
+          pip uninstall --yes deepspeed
+          pip install .[dev,1bit,autotuning,sparse_attn]
+          ds_report
+
+      - name: Unit tests
+        run: |
+          unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
+          if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
+          cd tests
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 -m 'not sequential' unit/
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/

.github/workflows/nv-torch12-p40.yml

@@ -17,7 +17,7 @@ concurrency:
 
 jobs:
   unit-tests:
-    runs-on: [self-hosted, nvidia, torch12, p40]
+    runs-on: [self-hosted, nvidia, cu101, p40]
 
     steps:
       - uses: actions/checkout@v2
@@ -29,6 +29,9 @@ jobs:
           python --version
           which nvcc
           nvcc --version
+          pip install --upgrade pip
+          pip uninstall --yes torch torchvision
+          pip install torch==1.2.0 torchvision==0.4.0
           python -c "import torch; print('torch:', torch.__version__, torch)"
           python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
 
@@ -39,11 +42,13 @@ jobs:
           # if needed switch to the last known good SHA until transformers@master is fixed
           # git checkout 1cc453d33
           git rev-parse --short HEAD
+          pip uninstall --yes transformers
           pip install .
 
       - name: Install deepspeed
         run: |
-          pip install .[dev,autotuning]
+          pip uninstall --yes deepspeed
+          pip install .[dev,1bit,autotuning,sparse_attn]
           ds_report
 
       - name: Unit tests

.github/workflows/nv-torch18-v100.yml

@@ -17,7 +17,7 @@ concurrency:
 
 jobs:
   unit-tests:
-    runs-on: [self-hosted, nvidia, torch18, v100]
+    runs-on: [self-hosted, nvidia, cu111, v100]
 
     steps:
       - uses: actions/checkout@v2
@@ -29,6 +29,8 @@ jobs:
           python --version
           which nvcc
           nvcc --version
+          pip install --upgrade pip
+          pip uninstall --yes torch torchvision
           pip install torch==1.8.2+cu111 torchvision==0.9.2+cu111 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
           python -c "import torch; print('torch:', torch.__version__, torch)"
           python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
@@ -40,10 +42,12 @@ jobs:
           # if needed switch to the last known good SHA until transformers@master is fixed
           # git checkout 1cc453d33
           git rev-parse --short HEAD
+          pip uninstall --yes transformers
           pip install .
 
       - name: Install deepspeed
         run: |
+          pip uninstall --yes deepspeed
           pip install .[dev,1bit,autotuning,sparse_attn]
           ds_report
 

.github/workflows/nv-transformers-v100.yml

@@ -17,7 +17,7 @@ concurrency:
 
 jobs:
   unit-tests:
-    runs-on: [self-hosted, nvidia, torch18, v100]
+    runs-on: [self-hosted, nvidia, cu111, v100]
 
     steps:
       - uses: actions/checkout@v2
@@ -29,11 +29,14 @@ jobs:
           python --version
           which nvcc
           nvcc --version
+          pip install --upgrade pip
+          pip uninstall --yes torch torchvision
           pip install torch==1.8.2+cu111 torchvision==0.9.2+cu111 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
           python -c "import torch; print('torch:', torch.__version__, torch)"
           python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
       - name: Install deepspeed
         run: |
+          pip uninstall --yes deepspeed
           pip install .[dev,autotuning]
           ds_report
       - name: HF transformers tests

setup.py

@@ -221,17 +221,28 @@ def create_dir_symlink(src, dest):
     version_str += f'+{git_hash}'
 
 torch_version = ".".join([TORCH_MAJOR, TORCH_MINOR])
+bf16_support = False
 # Set cuda_version to 0.0 if cpu-only
 cuda_version = "0.0"
+nccl_version = "0.0"
 # Set hip_version to 0.0 if cpu-only
 hip_version = "0.0"
 if torch_available and torch.version.cuda is not None:
     cuda_version = ".".join(torch.version.cuda.split('.')[:2])
+    if isinstance(torch.cuda.nccl.version(), int):
+        # This will break if minor version > 9
+        nccl_version = ".".join(str(torch.cuda.nccl.version())[:2])
+    else:
+        nccl_version = ".".join(map(str, torch.cuda.nccl.version()[:2]))
+    if hasattr(torch.cuda, 'is_bf16_supported'):
+        bf16_support = torch.cuda.is_bf16_supported()
 if torch_available and hasattr(torch.version, 'hip') and torch.version.hip is not None:
     hip_version = ".".join(torch.version.hip.split('.')[:2])
 torch_info = {
     "version": torch_version,
+    "bf16_support": bf16_support,
     "cuda_version": cuda_version,
+    "nccl_version": nccl_version,
     "hip_version": hip_version
 }
 

tests/unit/util.py

@@ -13,20 +13,15 @@ def required_torch_version():
 
 
 def bf16_required_version_check():
-    TORCH_MAJOR = int(torch.__version__.split('.')[0])
-    TORCH_MINOR = int(torch.__version__.split('.')[1])
-
-    if type(torch.cuda.nccl.version()) != tuple:
-        return False
-    else:
-        NCCL_MAJOR = torch.cuda.nccl.version()[0]
-        NCCL_MINOR = torch.cuda.nccl.version()[1]
+    split_version = lambda x: map(int, x.split('.')[:2])
+    TORCH_MAJOR, TORCH_MINOR = split_version(torch_info['version'])
+    NCCL_MAJOR, NCCL_MINOR = split_version(torch_info['nccl_version'])
+    CUDA_MAJOR, CUDA_MINOR = split_version(torch_info['cuda_version'])
 
-    CUDA_MAJOR = int(torch_info['cuda_version'].split('.')[0])
     if (TORCH_MAJOR > 1 or
         (TORCH_MAJOR == 1 and TORCH_MINOR >= 10)) and (CUDA_MAJOR >= 11) and (
             NCCL_MAJOR > 2 or
-            (NCCL_MAJOR == 2 and NCCL_MINOR >= 10)) and torch.cuda.is_bf16_supported():
+            (NCCL_MAJOR == 2 and NCCL_MINOR >= 10)) and torch_info['bf16_support']:
         return True
     else:
         return False