amt/src/extras/install_deepspeed.md

"""

not required on pytorch 2.0:latest container

pip install cupy-cuda11x -f https://pip.cupy.dev/aarch64

apt-get update apt-get install git apt-get install libaio-dev

DS_BUILD_OPS=1 pip install deepspeed ds_report

pip install deepspeed==0.7.7

git clone https://github.com/NVIDIA/apex cd apex pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

In case you have trouble building apex from source we recommend using the NGC containers from here which come with a pre-built PyTorch and apex release.

nvcr.io/nvidia/pytorch:23.01-py3

pip install deepspeed, pip install transformers[deepspeed] https://www.deepspeed.ai/docs/config-json/#autotuning

"""