Does deepspeed support aarch64?

[khayamgondal]

I am seeing the following error when trying to run it on aarch64 machine with H100.

Linux r8-u37 6.5.0-1019-nvidia-64k #19-Ubuntu SMP PREEMPT_DYNAMIC Tue May 7 12:54:40 UTC 2024 aarch64 aarch64 aarch64 GNU/Linux

anaconda3/envs/deepspeed/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp:10:10: fatal error: immintrin.h: No such file or directory
   10 | #include <immintrin.h>
      |          ^~~~~~~~~~~~~
compilation terminated.
[2/3] c++ -MMD -MF shm_interface.o.d -DTORCH_EXTENSION_NAME=deepspeed_shm_comm -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -I/home/khayam/anaconda3/envs/deepspeed/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/includes -isystem /home/khayam/anaconda3/envs/deepspeed/lib/python3.10/site-packages/torch/include -isystem /home/khayam/anaconda3/envs/deepspeed/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /home/khayam/anaconda3/envs/deepspeed/lib/python3.10/site-packages/torch/include/TH -isystem /home/khayam/anaconda3/envs/deepspeed/lib/python3.10/site-packages/torch/include/THC -isystem /home/khayam/anaconda3/envs/deepspeed/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++17 -O2 -fopenmp -c /home/khayam/anaconda3/envs/deepspeed/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp -o shm_interface.o 
/home/khayam/anaconda3/envs/deepspeed/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp: In function ‘void initialize(int, int)’:
/home/khayam/anaconda3/envs/deepspeed/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp:42:46: warning: ISO C++ forbids converting a string constant to ‘char*’ [-Wwrite-strings]
   42 |     if (addr_string == NULL) { addr_string = ""; }
      |                                              ^~
/home/khayam/anaconda3/envs/deepspeed/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm_interface.cpp:44:46: warning: ISO C++ forbids converting a string constant to ‘char*’ [-Wwrite-strings]
   44 |     if (port_string == NULL) { port_string = ""; }
      |                                              ^~
ninja: build stopped: subcommand failed.

[Sar-thak-3]

Hi @khayamgondal I am also trying to install Deepspeed for jetson nanoarchitecture. but I am encountering some different issues, can you let me know if you encountered the same error while installation?

error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [17 lines of output]
      fatal: not a git repository (or any of the parent directories): .git
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "/tmp/pip-install-2ye8_7hh/deepspeed_e72ce29d13824bf99346091714d95736/setup.py", line 247, in <module>
          if isinstance(torch.cuda.nccl.version(), int):
        File "/home/ies/miniconda3/envs/gdrnpp/lib/python3.10/site-packages/torch/cuda/nccl.py", line 35, in version
          ver = torch._C._nccl_version()
      AttributeError: module 'torch._C' has no attribute '_nccl_version'
      DS_BUILD_OPS=0
       [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
       [WARNING]  async_io: please install the libaio-dev package with apt
       [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
       [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
       [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.3
       [WARNING]  please install triton==1.0.0 if you want to use sparse attention
      Install Ops={'async_io': False, 'fused_adam': False, 'cpu_adam': False, 'cpu_adagrad': False, 'cpu_lion': False, 'evoformer_attn': False, 'fp_quantizer': False, 'fused_lamb': False, 'fused_lion': False, 'inference_core_ops': False, 'cutlass_ops': False, 'transformer_inference': False, 'quantizer': False, 'ragged_device_ops': False, 'ragged_ops': False, 'random_ltd': False, 'sparse_attn': False, 'spatial_inference': False, 'transformer': False, 'stochastic_transformer': False}
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

I think this might be the issue of nccl installation, but on the documentation page no nccl version satisfies aarch64 requirements

[tjruwase]

@khayamgondal and @Sar-thak-3, unfortunately we have not tested DeepSpeed in those environments due to lack of access.

[tjruwase]

anaconda3/envs/deepspeed/lib/python3.10/site-packages/deepspeed/ops/csrc/cpu/comm/shm.cpp:10:10: fatal error: immintrin.h: No such file or directory
10 | #include <immintrin.h>

@khayamgondal, immintrin.h is not available on arm64. However, I think you should not be compiling this extension if you are using CUDA. Can you share your build/installation command?

[khayamgondal]

Thanks for getting back - I installed deepspeed using pip install deepspeed. For cuda I am using NVIDIA NGC Pytorch container.

[tjruwase]

@khayamgondal, it is worthwhile building the source from the github to make it easier to see and fix the issues. Like I said earlier, we have not previously tested on arm.

[costin-eseanu]

@Sar-thak-3, it looks like the error you're seeing is unrelated to the original issue. We have not tested DeepSpeed on ARM, but the error you are seeing seems related to not having NCCL installed. Please install NCCL and then try installing DeepSpeed again. Alternately, you can also try building the source from github.

[loadams]

Closing this issue as stale and seems to be unrelated to aarch64.