From 5c0697cc9777b2e32adf13554aa7e0090252dcd5 Mon Sep 17 00:00:00 2001 From: Oliver Koenig Date: Thu, 19 Sep 2024 19:33:57 -0700 Subject: [PATCH] ADLR/megatron-lm!2113 - build: Use multi-stage for parallel builds --- .gitlab/stages/01.tests.yml | 74 +++++++++++----------- Dockerfile.ci | 118 ++++++++++++++---------------------- 2 files changed, 80 insertions(+), 112 deletions(-) diff --git a/.gitlab/stages/01.tests.yml b/.gitlab/stages/01.tests.yml index b3cefc0fde..d087425af9 100644 --- a/.gitlab/stages/01.tests.yml +++ b/.gitlab/stages/01.tests.yml @@ -35,48 +35,44 @@ build_image: variables: STAGE: main script: + - apk add bash - | - set -x - env - eval "IMAGE=\$$IMAGE" - - docker system prune -a --filter "until=24h" -f || true - - if [[ "$CI_COMMIT_BRANCH" == "$CI_DEFAULT_BRANCH" ]]; then - ADDITIONAL_PARAMS="--pull" - fi - - docker pull ${IMAGE}:${CI_PIPELINE_ID} || true - docker pull ${IMAGE}:${CI_MERGE_REQUEST_IID:-noop} || true - docker pull ${IMAGE}:buildcache || true - - docker build \ - --secret id=JET_INDEX_URLS \ - --target $STAGE \ - -f $FILE \ - -t ${IMAGE}:${CI_PIPELINE_ID} \ - -t ${IMAGE}:${CI_MERGE_REQUEST_IID:-noop} \ - --build-arg CACHEBUST=$(cat /proc/sys/kernel/random/uuid) \ - --cache-to type=inline \ - --cache-from type=registry,ref=${IMAGE}:buildcache \ - --cache-from type=registry,ref=${IMAGE}:${CI_PIPELINE_ID} \ - --cache-from type=registry,ref=${IMAGE}:${CI_MERGE_REQUEST_IID:-noop} \ - --build-arg FROM_IMAGE_NAME=$BASE_IMAGE \ - ${ADDITIONAL_PARAMS} . - - docker push ${IMAGE}:${CI_PIPELINE_ID} - docker push ${IMAGE}:${CI_MERGE_REQUEST_IID:-noop} - - if [[ "$CI_COMMIT_BRANCH" == "ci-nightly-a100" ]]; then - docker tag ${IMAGE}:${CI_PIPELINE_ID} ${IMAGE}:nightly - docker push ${IMAGE}:nightly - fi + bash -c ' + set -x + env + eval "IMAGE=\$$IMAGE" + + docker system prune -a --filter "until=24h" -f || true + + docker buildx create --name container --driver=docker-container + + ADDITIONAL_PARAMS=() + + if [[ "$CI_COMMIT_BRANCH" == "$CI_DEFAULT_BRANCH" ]]; then + ADDITIONAL_PARAMS+=("--pull") + ADDITIONAL_PARAMS+=("--cache-to type=registry,ref=${IMAGE}-buildcache:main") + fi - if [[ "$CI_COMMIT_BRANCH" == "$CI_DEFAULT_BRANCH" ]]; then - docker tag ${IMAGE}:${CI_PIPELINE_ID} ${IMAGE}:buildcache - docker push ${IMAGE}:buildcache - fi + if [[ "$CI_COMMIT_BRANCH" == "ci-nightly-a100" ]]; then + ADDITIONAL_PARAMS+=("-t ${IMAGE}:nightly") + fi + DOCKER_BUILDKIT=1 docker build \ + --secret id=JET_INDEX_URLS \ + --target $STAGE \ + -f $FILE \ + -t ${IMAGE}:${CI_PIPELINE_ID} \ + --builder=container \ + --build-arg CACHEBUST=$(cat /proc/sys/kernel/random/uuid) \ + --cache-to type=registry,ref=${IMAGE}-buildcache:${CI_PIPELINE_ID} \ + --cache-to type=registry,ref=${IMAGE}-buildcache:${CI_MERGE_REQUEST_IID:-noop} \ + --cache-from type=registry,ref=${IMAGE}-buildcache:main \ + --cache-from type=registry,ref=${IMAGE}-buildcache:${CI_PIPELINE_ID} \ + --cache-from type=registry,ref=${IMAGE}-buildcache:${CI_MERGE_REQUEST_IID:-noop} \ + --build-arg FROM_IMAGE_NAME=$BASE_IMAGE \ + --push \ + ${ADDITIONAL_PARAMS[@]} . + ' retry: max: 2 diff --git a/Dockerfile.ci b/Dockerfile.ci index dfcc7381f7..40c1464154 100644 --- a/Dockerfile.ci +++ b/Dockerfile.ci @@ -1,82 +1,54 @@ -# syntax=docker/dockerfile:experimental +# syntax=docker/dockerfile:1.3-labs ARG FROM_IMAGE_NAME -FROM $FROM_IMAGE_NAME as main -ENV DEBIAN_FRONTEND=noninteractive - -RUN sed -i -e 's/^APT/# APT/' -e 's/^DPkg/# DPkg/' \ - /etc/apt/apt.conf.d/docker-clean - -RUN apt-get update && \ - apt-get install -y --no-install-recommends gettext && \ - apt-get clean - -RUN wget https://github.com/mikefarah/yq/releases/download/v4.44.1/yq_linux_amd64 -O /usr/local/bin/yq && \ -chmod a+x /usr/local/bin/yq - -##### For Mamba begin ##### -RUN pip uninstall -y triton && \ - pip install triton==2.1.0 +FROM $FROM_IMAGE_NAME as build_causal_conv1d +WORKDIR /opt +RUN CAUSAL_CONV1D_FORCE_BUILD=TRUE pip3 wheel -v git+https://github.com/Dao-AILab/causal-conv1d.git@v1.2.2.post1 -# The causal-conv1d and mamba-ssm packages below are built from scratch here -# (which takes significant time) because there are no wheels available on PyPI -# for these relatively newer versions of the packages that are compatible with -# the older NGC-variant PyTorch version (e.g. version 2.2.0.dev231106) that we -# are using (in the NGC base container). Generally, if the package is not -# compatible with the PyTorch version, then it will generate a Python import -# error. The package authors tend to only release wheels for new versions of -# these pacakges which are compatible with the versions of regular PyTorch and -# NGC-variant PyTorch that are newer at the time of release. So, to use newer -# versions of these packages with relatively older versions of the NGC PyTorch -# container, we tend to have to build the packages from scratch. +FROM $FROM_IMAGE_NAME as build_grouped_gemm +WORKDIR /opt +RUN pip3 wheel -v git+https://github.com/fanshiqing/grouped_gemm@v1.1.2 -RUN cd /tmp && \ - pip uninstall -y causal-conv1d && \ - git clone https://github.com/Dao-AILab/causal-conv1d.git && \ - cd causal-conv1d && \ - git checkout v1.2.2.post1 && \ - CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install . && \ - cd .. && \ - rm -rf causal-conv1d +FROM $FROM_IMAGE_NAME as build_mamba_ssm +WORKDIR /opt +RUN MAMBA_FORCE_BUILD=TRUE pip3 wheel -v git+https://github.com/state-spaces/mamba.git@v2.0.3 -RUN cd /tmp && \ - pip uninstall -y mamba-ssm && \ - git clone https://github.com/state-spaces/mamba.git && \ - cd mamba && \ - git checkout v2.0.3 && \ - MAMBA_FORCE_BUILD=TRUE pip install . && \ - cd .. && \ - rm -rf mamba -##### For Mamba end ##### - -##### For JET-API start ##### -RUN apt-get update && \ - apt-get install -y python3-venv && \ - apt-get clean -y && \ - python -m venv /opt/jet -##### For JET-API end ##### - -RUN pip3 install --no-cache-dir \ - einops \ - flask-restful \ - nltk \ - pytest \ - pytest-cov \ - pytest_mock \ - pytest-random-order \ - sentencepiece \ - wrapt \ - git+https://github.com/fanshiqing/grouped_gemm@v1.1.2 \ - zarr \ - tensorstore==0.1.45 \ - wandb - -COPY . /workspace/megatron-lm - -COPY . /workspace/megatron-lm -RUN cp -r /workspace/megatron-lm /opt && \ - pip install /opt/megatron-lm +FROM $FROM_IMAGE_NAME as main +ENV DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y --no-install-recommends gettext python3-venv && \ + apt-get clean && \ + python -m venv /opt/jet && \ + wget https://github.com/mikefarah/yq/releases/download/v4.44.1/yq_linux_amd64 -O /usr/local/bin/yq && \ + chmod a+x /usr/local/bin/yq + +COPY --from=build_causal_conv1d /opt/causal_conv1d-1.2.2.post1-cp310-cp310-linux_x86_64.whl ./ +COPY --from=build_grouped_gemm /opt/grouped_gemm-1.1.2-cp310-cp310-linux_x86_64.whl ./ +COPY --from=build_mamba_ssm /opt/mamba_ssm-2.0.3-cp310-cp310-linux_x86_64.whl ./ + +RUN pip3 install --no-cache-dir --upgrade-strategy only-if-needed -v \ +einops \ +flask-restful \ +nltk \ +pytest \ +pytest-cov \ +pytest_mock \ +pytest-random-order \ +sentencepiece \ +wrapt \ +zarr \ +wandb \ +triton==2.1.0 \ +causal_conv1d-1.2.2.post1-cp310-cp310-linux_x86_64.whl \ +mamba_ssm-2.0.3-cp310-cp310-linux_x86_64.whl \ +grouped_gemm-1.1.2-cp310-cp310-linux_x86_64.whl \ +tensorstore==0.1.45 && \ +rm *.whl + +# Since megatron does not have any dependencies (and isn't a dependency to any other package), we can install it separately to make everything a bit quicker +COPY . /opt/megatron-lm +RUN pip install /opt/megatron-lm ##### For NVIDIANS only ##### FROM main as jet