winglian commited on
Commit
21b7439
1 Parent(s): 3f11b47

fix typo and add apex

Browse files
Files changed (1) hide show
  1. docker/Dockerfile-base +8 -1
docker/Dockerfile-base CHANGED
@@ -51,8 +51,15 @@ RUN git clone https://github.com/microsoft/DeepSpeed.git && \
51
 
52
  FROM base-builder
53
 
 
 
 
 
 
 
54
  RUN mkdir /workspace/wheels
55
  COPY --from=deepspeed-builder /workspace/DeepSpeed/dist/deepspeed-*.whl wheels
56
  COPY --from=flash-attn-builder /workspace/flash-attention/dist/flash_attn-*.whl wheels
57
 
58
- RUN pip3 install.sh wheels/deepspeed-*.whl wheels/flash_attn-*.whl
 
 
51
 
52
  FROM base-builder
53
 
54
+ # recompile apex
55
+ RUN python3 -m pip uninstall -y apex
56
+ RUN git clone https://github.com/NVIDIA/apex
57
+ # `MAX_JOBS=1` disables parallel building to avoid cpu memory OOM when building image on GitHub Action (standard) runners
58
+ RUN cd apex && MAX_JOBS=1 python3 -m pip install --global-option="--cpp_ext" --global-option="--cuda_ext" --no-cache -v --disable-pip-version-check .
59
+
60
  RUN mkdir /workspace/wheels
61
  COPY --from=deepspeed-builder /workspace/DeepSpeed/dist/deepspeed-*.whl wheels
62
  COPY --from=flash-attn-builder /workspace/flash-attention/dist/flash_attn-*.whl wheels
63
 
64
+ RUN pip3 install wheels/deepspeed-*.whl wheels/flash_attn-*.whl
65
+ RUN git lfs install --skip-repo