[chore] add circleci testing of torch==1.5.1 (#172)

.circleci/config.yml

@@ -38,12 +38,12 @@ setup_venv: &setup_venv
         which pip
         pip install --upgrade pip
 
-install_dep_17: &install_dep_17
+install_dep_15: &install_dep_15
   - run:
       name: Install Dependencies
       command: |
         sudo apt-get install -y libopenmpi-dev
-        pip install --progress-bar off torch==1.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html
+        pip install --progress-bar off torch==1.5.1+cu101 -f https://download.pytorch.org/whl/torch_stable.html
         pip install --progress-bar off -r requirements-test.txt
         python -c 'import torch; print("Torch version:", torch.__version__)'
         python -m torch.utils.collect_env
@@ -58,6 +58,16 @@ install_dep_16: &install_dep_16
         python -c 'import torch; print("Torch version:", torch.__version__)'
         python -m torch.utils.collect_env
 
+install_dep_17: &install_dep_17
+  - run:
+      name: Install Dependencies
+      command: |
+        sudo apt-get install -y libopenmpi-dev
+        pip install --progress-bar off torch==1.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html
+        pip install --progress-bar off -r requirements-test.txt
+        python -c 'import torch; print("Torch version:", torch.__version__)'
+        python -m torch.utils.collect_env
+
 install_repo_cpu: &install_repo_cpu
   - run:
       name: Install Repository
@@ -136,7 +146,7 @@ jobs:
           keys:
             - cache-key-cpu-{{ checksum "setup.py"}}-{{ checksum "requirements-test.txt"}}
 
-      - <<: *install_dep_16
+      - <<: *install_dep_17
 
       - save_cache:
           paths:
@@ -167,7 +177,7 @@ jobs:
       - store_test_results:
           path: test-results
 
-  gpu_tests_17:
+  gpu_tests_15:
     <<: *gpu
 
     working_directory: ~/fairscale
@@ -184,14 +194,14 @@ jobs:
       # Cache the venv directory that contains dependencies
       - restore_cache:
           keys:
-            - cache-key-gpu17-{{ checksum "setup.py"}}-{{ checksum "requirements-test.txt"}}
+            - cache-key-gpu15-{{ checksum "setup.py"}}-{{ checksum "requirements-test.txt"}}
 
-      - <<: *install_dep_17
+      - <<: *install_dep_15
 
       - save_cache:
           paths:
             - ~/venv
-          key: cache-key-gpu17-{{ checksum "setup.py"}}-{{ checksum "requirements-test.txt"}}
+          key: cache-key-gpu15-{{ checksum "setup.py"}}-{{ checksum "requirements-test.txt"}}
 
       - <<: *install_repo_gpu
 
@@ -235,6 +245,39 @@ jobs:
       - store_test_results:
           path: test-results
 
+  gpu_tests_17:
+    <<: *gpu
+
+    working_directory: ~/fairscale
+
+    steps:
+      - checkout
+
+      - run: nvidia-smi
+
+      - run: pyenv global 3.7.0
+
+      - <<: *setup_venv
+
+      # Cache the venv directory that contains dependencies
+      - restore_cache:
+          keys:
+            - cache-key-gpu17-{{ checksum "setup.py"}}-{{ checksum "requirements-test.txt"}}
+
+      - <<: *install_dep_17
+
+      - save_cache:
+          paths:
+            - ~/venv
+          key: cache-key-gpu17-{{ checksum "setup.py"}}-{{ checksum "requirements-test.txt"}}
+
+      - <<: *install_repo_gpu
+
+      - <<: *run_unittests
+
+      - store_test_results:
+          path: test-results
+
   benchmarks:
     <<: *gpu
 
@@ -281,6 +324,7 @@ workflows:
   build:
     jobs:
       - cpu_tests
-      - gpu_tests_17
+      - gpu_tests_15
       - gpu_tests_16
+      - gpu_tests_17
       - benchmarks

README.md

@@ -106,7 +106,7 @@ if __name__ == "__main__":
 
 # Testing
 
-We use circleci to test on PyTorch versions 1.6.0 and 1.7.0 and CUDA version 10.1. Please create an [issue](https://github.com/facebookresearch/fairscale/issues) if you are having trouble with installation.
+We use circleci to test on PyTorch versions 1.5.1, 1.6.0 and 1.7.0 and CUDA version 10.1. Please create an [issue](https://github.com/facebookresearch/fairscale/issues) if you are having trouble with installation.
 
 ## Contributors
 

requirements-test.txt

@@ -7,7 +7,7 @@ pytest == 5.4.1
 pytest-cov == 2.10.0
 pytest-mpi == 0.4
 torchtext == 0.6.0
-torch >= 1.6.0
+torch >= 1.5.1
 torchvision >= 0.6.0
 # NOTE(msb) not a dependency but needed by torch
 numpy == 1.17.4

tests/nn/model_parallel/commons.py

@@ -64,7 +64,8 @@ def dist_init(rank, world_size, hostname=None):
 
     if version.parse(torch.__version__).release >= (1, 6, 0):
         init_method = f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}"
-        torch.distributed.init_process_group(backend="nccl", rank=rank, world_size=world_size, init_method=init_method)
+        backend = "nccl" if torch.cuda.is_available() else "gloo"
+        torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size, init_method=init_method)
         os.environ["MASTER_ADDR"] = hostname
         os.environ["MASTER_PORT"] = "10639"
         init_method = f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}"