Repository: facebookresearch/fairscale

Issue:
Fail on install - CUDA 11.1.0
## 🐛 Bug

Hi,
pip install on below environment throws an error. I'm happy to provide more info if it would be useful. Thanks!

fatal error: multi_tensor_apply.cuh: No such file or directory

## Environment (NVIDIA-Python Docker: 20.10)
Python: 3.6
PyTorch: 1.7.0
CUDA: 11.1.0
cuDNN: 8.0.4

## Error
```
Error:
b'  ERROR: Command errored out with exit status 1:
   command: /usr/bin/python3 /usr/local/lib/python3.6/dist-packages/pip/_vendor/pep517/_in_process.py build_wheel /tmp/tmpl6xin7ge
       cwd: /tmp/pip-install-0_xiegmi/fairscale
  Complete output (158 lines):
  running bdist_wheel
  running build
  running build_py
  creating build
  creating build/lib.linux-x86_64-3.6
  creating build/lib.linux-x86_64-3.6/fairscale
  copying fairscale/__init__.py -> build/lib.linux-x86_64-3.6/fairscale
  creating build/lib.linux-x86_64-3.6/fairscale/optim
  copying fairscale/optim/grad_scaler.py -> build/lib.linux-x86_64-3.6/fairscale/optim
  copying fairscale/optim/oss.py -> build/lib.linux-x86_64-3.6/fairscale/optim
  copying fairscale/optim/__init__.py -> build/lib.linux-x86_64-3.6/fairscale/optim
  copying fairscale/optim/utils.py -> build/lib.linux-x86_64-3.6/fairscale/optim
  copying fairscale/optim/adam.py -> build/lib.linux-x86_64-3.6/fairscale/optim
  creating build/lib.linux-x86_64-3.6/fairscale/nn
  copying fairscale/nn/__init__.py -> build/lib.linux-x86_64-3.6/fairscale/nn
  creating build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/dependency.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/pipe.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/stream.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/batchnorm.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/__init__.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/pipeline.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/checkpoint.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/copy.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/phony.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/worker.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  copying fairscale/nn/pipe/microbatch.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe
  creating build/lib.linux-x86_64-3.6/fairscale/nn/model_parallel
  copying fairscale/nn/model_parallel/random.py -> build/lib.linux-x86_64-3.6/fairscale/nn/model_parallel
  copying fairscale/nn/model_parallel/initialize.py -> build/lib.linux-x86_64-3.6/fairscale/nn/model_parallel
  copying fairscale/nn/model_parallel/layers.py -> build/lib.linux-x86_64-3.6/fairscale/nn/model_parallel
  copying fairscale/nn/model_parallel/cross_entropy.py -> build/lib.linux-x86_64-3.6/fairscale/nn/model_parallel
  copying fairscale/nn/model_parallel/mappings.py -> build/lib.linux-x86_64-3.6/fairscale/nn/model_parallel
  copying fairscale/nn/model_parallel/__init__.py -> build/lib.linux-x86_64-3.6/fairscale/nn/model_parallel
  copying fairscale/nn/model_parallel/utils.py -> build/lib.linux-x86_64-3.6/fairscale/nn/model_parallel
  creating build/lib.linux-x86_64-3.6/fairscale/nn/data_parallel
  copying fairscale/nn/data_parallel/sharded_ddp.py -> build/lib.linux-x86_64-3.6/fairscale/nn/data_parallel
  copying fairscale/nn/data_parallel/__init__.py -> build/lib.linux-x86_64-3.6/fairscale/nn/data_parallel
  creating build/lib.linux-x86_64-3.6/fairscale/nn/moe
  copying fairscale/nn/moe/top2gate.py -> build/lib.linux-x86_64-3.6/fairscale/nn/moe
  copying fairscale/nn/moe/moelayer.py -> build/lib.linux-x86_64-3.6/fairscale/nn/moe
  copying fairscale/nn/moe/__init__.py -> build/lib.linux-x86_64-3.6/fairscale/nn/moe
  creating build/lib.linux-x86_64-3.6/fairscale/nn/pipe/balance
  copying fairscale/nn/pipe/balance/profile.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe/balance
  copying fairscale/nn/pipe/balance/__init__.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe/balance
  copying fairscale/nn/pipe/balance/blockpartition.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe/balance
  creating build/lib.linux-x86_64-3.6/fairscale/nn/pipe/skip
  copying fairscale/nn/pipe/skip/skippable.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe/skip
  copying fairscale/nn/pipe/skip/__init__.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe/skip
  copying fairscale/nn/pipe/skip/tracker.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe/skip
  copying fairscale/nn/pipe/skip/layout.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe/skip
  copying fairscale/nn/pipe/skip/portal.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe/skip
  copying fairscale/nn/pipe/skip/namespace.py -> build/lib.linux-x86_64-3.6/fairscale/nn/pipe/skip
  running egg_info
  writing fairscale.egg-info/PKG-INFO
  writing dependency_links to fairscale.egg-info/dependency_links.txt
  writing requirements to fairscale.egg-info/requires.txt
  writing top-level names to fairscale.egg-info/top_level.txt
  reading manifest file \'fairscale.egg-info/SOURCES.txt\'
  reading manifest template \'MANIFEST.in\'
  writing manifest file \'fairscale.egg-info/SOURCES.txt\'
  creating build/lib.linux-x86_64-3.6/fairscale/clib
  creating build/lib.linux-x86_64-3.6/fairscale/clib/fused_adam_cuda
  copying fairscale/clib/fused_adam_cuda/fused_adam_cuda.cpp -> build/lib.linux-x86_64-3.6/fairscale/clib/fused_adam_cuda
  copying fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.cu -> build/lib.linux-x86_64-3.6/fairscale/clib/fused_adam_cuda
  running build_ext
  building \'fairscale.fused_adam_cuda\' extension
  creating /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6
  creating /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6/fairscale
  creating /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6/fairscale/clib
  creating /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6/fairscale/clib/fused_adam_cuda
  Emitting ninja build file /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6/build.ninja...
  Compiling objects...
  Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
  [1/2] /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c -c /tmp/pip-install-0_xiegmi/fairscale/fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.cu -o /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6/fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options \'\'"\'"\'-fPIC\'"\'"\'\' -O3 --use_fast_math -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=fused_adam_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_61,code=sm_61 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_52,code=sm_52 -gencode=arch=compute_86,code=sm_86 -std=c++14
  FAILED: /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6/fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.o
  /usr/local/cuda/bin/nvcc -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c -c /tmp/pip-install-0_xiegmi/fairscale/fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.cu -o /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6/fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options \'\'"\'"\'-fPIC\'"\'"\'\' -O3 --use_fast_math -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=fused_adam_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_61,code=sm_61 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_52,code=sm_52 -gencode=arch=compute_86,code=sm_86 -std=c++14
  /tmp/pip-install-0_xiegmi/fairscale/fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.cu:12:10: fatal error: multi_tensor_apply.cuh: No such file or directory
   #include "multi_tensor_apply.cuh"
            ^~~~~~~~~~~~~~~~~~~~~~~~
  compilation terminated.
  [2/2] c++ -MMD -MF /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6/fairscale/clib/fused_adam_cuda/fused_adam_cuda.o.d -pthread -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.6/dist-packages/torch/include -I/usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include -I/usr/local/lib/python3.6/dist-packages/torch/include/TH -I/usr/local/lib/python3.6/dist-packages/torch/include/THC -I/usr/local/cuda/include -I/usr/include/python3.6m -c -c /tmp/pip-install-0_xiegmi/fairscale/fairscale/clib/fused_adam_cuda/fused_adam_cuda.cpp -o /tmp/pip-install-0_xiegmi/fairscale/build/temp.linux-x86_64-3.6/fairscale/clib/fused_adam_cuda/fused_adam_cuda.o -O3 -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=fused_adam_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14
  In file included from /usr/local/lib/python3.6/dist-packages/torch/include/ATen/Parallel.h:149:0,
                   from /usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include/torch/utils.h:3,
                   from /usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include/torch/nn/cloneable.h:5,
                   from /usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include/torch/nn.h:3,
                   from /usr/local/lib/python3.6/dist-packages/torch/include/torch/csrc/api/include/torch/all.h:12,
                   from /usr/local/lib/python3.6/dist-packages/torch/include/torch/extension.h:4,
                   from /tmp/pip-install-0_xiegmi/fairscale/fairscale/clib/fused_adam_cuda/fused_adam_cuda.cpp:1:
  /usr/local/lib/python3.6/dist-packages/torch/include/ATen/ParallelOpenMP.h:84:0: warning: ignoring #pragma omp parallel [-Wunknown-pragmas]
   #pragma omp parallel for if ((end - begin) >= grain_size)
  
  ninja: build stopped: subcommand failed.
  Traceback (most recent call last):
    File "/usr/local/lib/python3.6/dist-packages/torch/utils/cpp_extension.py", line 1522, in _run_ninja_build
      env=env)
    File "/usr/lib/python3.6/subprocess.py", line 438, in run
      output=stdout, stderr=stderr)
  subprocess.CalledProcessError: Command \'[\'ninja\', \'-v\']\' returned non-zero exit status 1.
```


Discussion:
What is the command-line you are using to install?

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -10,6 +10,8 @@
 import torch
 from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
 
+this_dir = os.path.dirname(os.path.abspath(__file__))
+
 
 def fetch_requirements():
     with open("requirements.txt") as f:
@@ -35,6 +37,7 @@ def find_version(version_file_path):
         [
             CUDAExtension(
                 name="fairscale.fused_adam_cuda",
+                include_dirs=[os.path.join(this_dir, "fairscale/clib/fused_adam_cuda")],
                 sources=[
                     "fairscale/clib/fused_adam_cuda/fused_adam_cuda.cpp",
                     "fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.cu",


Test patch: