chore(docker): Adds a Dockerfile with DeepSpeed and Flash-Attention.

docker/Dockerfile.flash

@@ -1,20 +1,35 @@
-# Defines the base image
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license
+
+# Root image to be based
+# Available images: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags
 FROM nvcr.io/nvidia/pytorch:23.02-py3
 
-# Installs desired packages for the workload
+# Labels for the docker
+LABEL description="NVIDIA Docker with Archai, DeepSpeed and Flash-Attention" \
+      repository="archai" \
+      tag="latest" \
+      creator="microsoft" \
+      tooltype="archai" \
+      createtime="03/20/2023"
+
+# Exports environment variables
+ENV PATH="/root/.local/bin:$PATH"
+
+# Installs basic utilities
 RUN apt-get update && \
     apt-get install --no-install-recommends --no-install-suggests -yq && \
     rm -rf /var/lib/apt/lists/* && \
     apt-get purge --auto-remove && \
     apt-get clean
 
-# DeepSpeed and OpenMPI
+# Installs DeepSpeed and OpenMPI
 RUN pip install --upgrade pip && \
     pip uninstall -y xgboost && \
     DS_BUILD_UTILS=1 DS_BUILD_FUSED_LAMB=1 pip install deepspeed==0.8.1 && \
     CC=mpicc MPICC=mpicc pip install mpi4py --no-binary mpi4py
 
-# Flash-Attention and CUDA extensions for cross-entropy, fused dense, layer norm
+# Installs Flash-Attention and CUDA extensions for cross-entropy, fused dense, layer norm
 RUN pip install flash-attn==0.2.8
 RUN git clone https://github.com/HazyResearch/flash-attention \
     && cd flash-attention && git checkout v0.2.8 \
@@ -26,5 +41,7 @@ RUN git clone https://github.com/HazyResearch/flash-attention \
     # && cd csrc/ft_attention && pip install . && cd ../../ \
     && cd .. && rm -rf flash-attention
 
-# Archai (development)
-RUN pip install --user git+https://github.com/microsoft/archai.git#egg=archai[dev]
+# Installs Archai
+RUN git clone -b main --single-branch https://github.com/microsoft/archai.git
+WORKDIR /workspace/archai
+RUN pip install --user --no-cache-dir .[dev]

scripts/trainers/README.md

@@ -20,7 +20,7 @@ Before you can start training models, you need to install Archai. To do so, you
 
 *Please note that DeepSpeed is not compatible with Windows.*
 
-Alternatively, you can use Docker to build a Docker image with Archai and all the necessary dependencies. Simply follow the instructions in the `Dockerfile` provided in this folder.
+Alternatively, you can use Docker to build a Docker image with Archai and all the necessary dependencies. Simply follow the instructions in the [Dockerfile](https://github.com/microsoft/archai/blob/main/docker/Dockerfile.flash).
 
 ## Data Preparation