diff --git a/training/Dockerfile b/training/Dockerfile index 098132c..c56480f 100644 --- a/training/Dockerfile +++ b/training/Dockerfile @@ -2,7 +2,7 @@ # ARG COMPAT=0 ARG PERSONAL=0 # FROM nvidia/cuda:11.3.1-devel-ubuntu20.04 as base-0 -FROM nvcr.io/nvidia/pytorch:22.11-py3 as base +FROM nvcr.io/nvidia/pytorch:22.12-py3 as base ENV HOST docker ENV LANG=C.UTF-8 LC_ALL=C.UTF-8 @@ -67,30 +67,29 @@ ENV PIP_NO_CACHE_DIR=1 # RUN pip install --upgrade --force-reinstall --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_multihead_attn" --global-option="--fmha" --global-option="--fast_layer_norm" --global-option="--xentropy" git+https://github.com/NVIDIA/apex.git#egg=apex # xgboost conflicts with deepspeed -RUN pip uninstall -y xgboost && DS_BUILD_UTILS=1 DS_BUILD_FUSED_LAMB=1 pip install deepspeed==0.7.5 +RUN pip uninstall -y xgboost && DS_BUILD_UTILS=1 DS_BUILD_FUSED_LAMB=1 pip install deepspeed==0.7.7 # General packages that we don't care about the version # zstandard to extract the_pile dataset # psutil to get the number of cpu physical cores # twine to upload package to PyPI -# ninja is broken for some reason, it returns error code 245 -RUN pip uninstall -y ninja && pip install ninja -RUN pip install pytest matplotlib jupyter ipython ipdb gpustat scikit-learn spacy munch einops opt_einsum fvcore gsutil cmake pykeops zstandard psutil h5py twine \ +RUN pip install pytest matplotlib jupyter ipython ipdb gpustat scikit-learn spacy munch einops opt_einsum fvcore gsutil cmake pykeops zstandard psutil h5py twine gdown \ && python -m spacy download en_core_web_sm # hydra -RUN pip install hydra-core==1.2.0 hydra-colorlog==1.2.0 hydra-optuna-sweeper==1.2.0 pyrootutils rich +RUN pip install hydra-core==1.3.1 hydra-colorlog==1.2.0 hydra-optuna-sweeper==1.2.0 pyrootutils rich # Core packages -RUN pip install transformers==4.24.0 datasets==2.7.1 pytorch-lightning==1.7.7 triton==2.0.0.dev20221120 wandb==0.13.5 timm==0.6.12 torchmetrics==0.10.3 +RUN pip install transformers==4.25.1 datasets==2.8.0 pytorch-lightning==1.8.6 triton==2.0.0.dev20221202 wandb==0.13.7 timm==0.6.12 torchmetrics==0.11.0 # For MLPerf RUN pip install git+https://github.com/mlcommons/logging.git@2.1.0 # Install FlashAttention -RUN pip install flash-attn==0.2.2 +RUN pip install flash-attn==0.2.6.post # Install CUDA extensions for cross-entropy, fused dense, layer norm -RUN git clone https://github.com/HazyResearch/flash-attention \ - && cd flash-attention && git checkout v0.2.2 \ +# Install CUDA extensions for cross-entropy, fused dense, layer norm + RUN git clone https://github.com/HazyResearch/flash-attention \ + && cd flash-attention && git checkout v0.2.6 \ && cd csrc/fused_softmax && pip install . && cd ../../ \ && cd csrc/rotary && pip install . && cd ../../ \ && cd csrc/xentropy && pip install . && cd ../../ \