From 9696cc4ef7b8859b830b9e4cb1be1cf881945bcd Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Tue, 5 Dec 2023 14:28:28 +0100 Subject: [PATCH] precompile deepspeed to avoid timeout during tests --- docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile index 36a2d8fafd8e..466af2fcbf88 100644 --- a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile +++ b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile @@ -21,10 +21,13 @@ RUN apt update && \ apt clean && \ rm -rf /var/lib/apt/lists/* -RUN python3 -m pip install --no-cache-dir --upgrade pip +RUN python3 -m pip install --no-cache-dir --upgrade pip ninja "pydantic<2" RUN python3 -m pip uninstall -y apex torch torchvision torchaudio RUN python3 -m pip install torch==$PYTORCH torchvision==$TORCH_VISION torchaudio==$TORCH_AUDIO --index-url https://download.pytorch.org/whl/rocm$ROCM --no-cache-dir +# Pre-build **latest** DeepSpeed, so it's be ready for testing (to avoid timeout) +RUN DS_BUILD_CPU_ADAM=1 DS_BUILD_FUSED_ADAM=1 python3 -m pip install deepspeed --global-option="build_ext" --global-option="-j8" --no-cache-dir -v --disable-pip-version-check 2>&1 + ARG REF=main WORKDIR / @@ -32,10 +35,10 @@ WORKDIR / ADD https://api.github.com/repos/huggingface/transformers/git/refs/heads/main version.json RUN git clone https://github.com/huggingface/transformers && cd transformers && git checkout $REF -RUN python3 -m pip install --no-cache-dir ./transformers[deepspeed-testing] +RUN python3 -m pip install --no-cache-dir ./transformers[accelerate,testing,sentencepiece] # When installing in editable mode, `transformers` is not recognized as a package. # this line must be added in order for python to be aware of transformers. RUN cd transformers && python3 setup.py develop -RUN python3 -c "from deepspeed.launcher.runner import main" +RUN python3 -c "from deepspeed.launcher.runner import main" \ No newline at end of file