Apex

https://github.com/NVIDIA/apex

WORKDIR /apex-master
RUN export TORCH_CUDA_ARCH_LIST="6.0;7.0"
RUN pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

deepspeed

https://github.com/microsoft/DeepSpeed