Merge branch 'ko3n1g/chore/bump-pyt' into 'main'

chore: Bump Pytorch container

See merge request ADLR/megatron-lm!2017

Author: ko3n1g

.gitlab/stages/01.tests.yml

@@ -23,7 +23,7 @@ build_image:
         TAG: mcore-docker-node-large
       - IMAGE: CI_MCORE_DEV_IMAGE
         FILE: Dockerfile.ci.dev
-        BASE_IMAGE: nvcr.io/nvidia/pytorch:24.01-py3
+        BASE_IMAGE: nvcr.io/nvidia/pytorch:24.07-py3
         TAG: mcore-docker-node-large
       - IMAGE: CI_NEMO_IMAGE
         FILE: Dockerfile.ci
@@ -92,10 +92,12 @@ unit_tests:
     matrix:
       - TAG: latest
         IMAGE: ${CI_MCORE_IMAGE}
-      # - TAG: latest
-      #   IMAGE: ${CI_MCORE_DEV_IMAGE}
+      - TAG: latest
+        IMAGE: ${CI_MCORE_DEV_IMAGE}
       - TAG: core_r0.9.0
         IMAGE: ${CI_MCORE_IMAGE}
+      - TAG: core_r0.9.0
+        IMAGE: ${CI_MCORE_DEV_IMAGE}
   tags: [8xL40S]
   variables:
     GIT_STRATEGY: clone
@@ -109,6 +111,9 @@ unit_tests:
       fi
   script:
     - |
+      export NVTE_FLASH_ATTN=0
+      export NVTE_FUSED_ATTN=0
+      
       cd /opt/megatron-lm
       if [[ $UNIT_TEST_REPEAT -eq 0 ]]; then
         exit 0
@@ -118,12 +123,9 @@ unit_tests:
         SEED=$((RANDOM % 9000 + 1000));
         ARGS=()
         if [[ $TAG != latest ]]; then
-          ARGS+=(-m "not internal")
+          ARGS+=(-m "not internal and not flaky and not flaky_in_dev")
         else
-          ARGS+=(-m "not flaky")
-        fi
-        if [[ $IMAGE == ${CI_MCORE_DEV_IMAGE} ]]; then
-          ARGS+=(-m "experimental")
+          ARGS+=(-m "not flaky and not flaky_in_dev")
         fi
         timeout ${UNIT_TEST_TIMEOUT}m torchrun --nproc_per_node=8 -m pytest --random-order --random-order-seed ${SEED} -xvs --cov-report=term --cov-report=html --cov=megatron/core --no-cov-on-fail "${ARGS[@]}" tests/unit_tests
       done

.gitlab/stages/02.functional-tests.yml

@@ -51,25 +51,56 @@ jet-generate:
       export PYTHONPATH=$(pwd)
       python tests/functional_tests/python_test_utils/jet/generate_jet_trigger_job.py \
         --scope $FUNCTIONAL_TEST_SCOPE \
+        --environment dev \
         --a100-cluster $A100_CLUSTER \
         --h100-cluster $H100_CLUSTER \
+        --container-image ${CI_MCORE_IMAGE} \
         --container-tag ${CI_PIPELINE_ID} \
+        --output-path "jet-trigger-job-dev.yaml" \
+        ${RELEASE_ARGS[@]}
+    
+    - |
+      export PYTHONPATH=$(pwd)
+      python tests/functional_tests/python_test_utils/jet/generate_jet_trigger_job.py \
+        --scope $FUNCTIONAL_TEST_SCOPE \
+        --environment lts \
+        --a100-cluster $A100_CLUSTER \
+        --h100-cluster $H100_CLUSTER \
         --container-image ${CI_MCORE_IMAGE} \
-        --container-image-dev ${CI_MCORE_DEV_IMAGE} \
-        --output-path "jet-trigger-job.yaml" \
+        --container-tag ${CI_PIPELINE_ID} \
+        --output-path "jet-trigger-job-lts.yaml" \
         ${RELEASE_ARGS[@]}
   artifacts:
     paths:
-      - jet-trigger-job.yaml
+      - jet-trigger-job-lts.yaml
+      - jet-trigger-job-dev.yaml
       - tests/functional_tests/local_recipes
 
-jet-trigger:
+jet-trigger-lts:
+  stage: functional_tests
+  needs: [jet-generate]
+  extends: [.jet_common]
+  trigger:
+    include:
+      - artifact: jet-trigger-job-lts.yaml
+        job: jet-generate
+    strategy: depend
+  variables:
+    RO_API_TOKEN: $PAT
+    CONTAINER_TAG: $CI_PIPELINE_ID
+    CI_MCORE_IMAGE: $CI_MCORE_IMAGE
+    GITLAB_ENDPOINT: $GITLAB_ENDPOINT
+    PARENT_PIPELINE_ID: $CI_PIPELINE_ID
+  inherit:
+    variables: true
+
+jet-trigger-dev:
   stage: functional_tests
   needs: [jet-generate]
   extends: [.jet_common]
   trigger:
     include:
-      - artifact: jet-trigger-job.yaml
+      - artifact: jet-trigger-job-dev.yaml
         job: jet-generate
     strategy: depend
   variables:
@@ -81,10 +112,10 @@ jet-trigger:
   inherit:
     variables: true
 
-jet-results-notify:
+jet-results-notify-lts:
   extends: [.jet_common]
   image: ${GITLAB_ENDPOINT}:5005/dl/jet/api:latest
-  needs: [jet-trigger]
+  needs: [jet-trigger-lts, jet-trigger-dev]
   tags:
     - mcore-docker-node-small
   before_script:
@@ -96,7 +127,7 @@ jet-results-notify:
     - export GITLAB_ENDPOINT
     - export CONTEXT=$FUNCTIONAL_TEST_SCOPE
     - export DATE=$(date +"%Y-%m-%d")
-    - bash tests/functional_tests/shell_test_utils/notify.sh ${CI_PIPELINE_ID}
+    - bash tests/functional_tests/shell_test_utils/notify.sh ${CI_PIPELINE_ID} lts
   artifacts:
     when: always
     paths:
@@ -106,3 +137,27 @@ jet-results-notify:
       when: always
     - when: never
 
+jet-results-notify-dev:
+  extends: [.jet_common]
+  image: ${GITLAB_ENDPOINT}:5005/dl/jet/api:latest
+  needs: [jet-trigger-lts, jet-trigger-dev]
+  tags:
+    - mcore-docker-node-small
+  before_script:
+    - jet secrets jwt-login jwt/nvidia/gitlab-master adlr-megatron-lm-ci $VAULT_JWT_TOKEN
+  script:
+    - env
+    - export WEBHOOK_URL=${MCORE_NOTIFICATION_HOOK}
+    - export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}
+    - export GITLAB_ENDPOINT
+    - export CONTEXT=$FUNCTIONAL_TEST_SCOPE
+    - export DATE=$(date +"%Y-%m-%d")
+    - bash tests/functional_tests/shell_test_utils/notify.sh ${CI_PIPELINE_ID} dev
+  artifacts:
+    when: always
+    paths:
+      - scripts
+  rules:
+    - if: $CI_PIPELINE_SOURCE == "schedule" && $FUNCTIONAL_TEST == "yes"
+      when: always
+    - when: never

Dockerfile.ci.dev

@@ -11,7 +11,7 @@ RUN pip3 wheel -v git+https://github.com/fanshiqing/[email protected]
 
 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 MAMBA_FORCE_BUILD=TRUE pip3 wheel -v git+https://github.com/state-spaces/mamba.git@v2.2.0
 
 FROM $FROM_IMAGE_NAME as main
 ENV DEBIAN_FRONTEND=noninteractive
@@ -23,9 +23,9 @@ RUN apt-get update && \
     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 ./
+COPY --from=build_causal_conv1d /opt/causal_conv1d-*.whl ./
+COPY --from=build_grouped_gemm /opt/grouped_gemm-*.whl ./
+COPY --from=build_mamba_ssm /opt/mamba_ssm-*.whl ./
 
 RUN pip3 install --no-cache-dir --upgrade-strategy only-if-needed -v \
 einops \
@@ -40,10 +40,9 @@ tiktoken \
 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 \
+causal_conv1d-*.whl \
+mamba_ssm-*.whl \
+grouped_gemm-*.whl \
 tensorstore==0.1.45 && \
 rm *.whl