Merge branch 'main' into gptoss-120b-card

Author: cuichenx

.dev.commit

@@ -1 +1 @@
-cae55c1f0991abd3e391d0029b15bcb35723d3f6
+b3b1e2d793c596c2fa61076d001ab93f294e4e87

.github/CODEOWNERS

@@ -5,8 +5,11 @@ pyproject.toml @nvidia-nemo/automation
 ruff.toml @nvidia-nemo/automation
 .pre-commit-config.yaml
 
-scripts/performance/argument_parser.py @nvidia-nemo/automation @malay-nagda @erhoo82
-scripts/performance/setup_experiment.py @nvidia-nemo/automation @malay-nagda @erhoo82
-scripts/performance/bootstrap.py @nvidia-nemo/automation @malay-nagda @erhoo82
-scripts/performance/run_script.py @nvidia-nemo/automation @malay-nagda @erhoo82
-scripts/performance/run_recipe.py @nvidia-nemo/automation @malay-nagda @erhoo82
+scripts/performance @malay-nagda @dingqingy-nv
+src/megatron/bridge/perf_recipes @malay-nagda @dingqingy-nv
+
+scripts/performance/argument_parser.py @nvidia-nemo/automation
+scripts/performance/setup_experiment.py @nvidia-nemo/automation
+scripts/performance/bootstrap.py @nvidia-nemo/automation
+scripts/performance/run_script.py @nvidia-nemo/automation
+scripts/performance/run_recipe.py @nvidia-nemo/automation

.github/actions/test-template/action.yml

@@ -43,8 +43,8 @@ inputs:
     description: "Has Azure credentials"
     required: false
     default: "false"
-  PAT:
-    description: "GitHub Personal Access Token"
+  github-token:
+    description: "GitHub token"
     required: true
   container-image:
     description: "Container image to use for test"
@@ -59,6 +59,14 @@ inputs:
   runner:
     description: "Runner to use for test"
     required: true
+  shard-id:
+    description: "Zero-based pytest shard index"
+    required: false
+    default: ""
+  num-shards:
+    description: "Total number of pytest shards"
+    required: false
+    default: ""
 
 runs:
   using: "composite"
@@ -114,9 +122,14 @@ runs:
         echo "::group::Start test container"
         if [[ "$MOUNT_FS" == "true" ]]; then
           VOLUME_ARGS="--volume ${{ inputs.test-data-path }}:/home/TestData --env HF_HUB_OFFLINE=1"
+          HF_HOME=/home/TestData/HF_HOME
+          NEMO_HOME=/home/TestData/nemo_home
         else
           VOLUME_ARGS=""
+          HF_HOME=/home/ubuntu/.cache/huggingface
+          NEMO_HOME=/home/ubuntu/.cache/nemo
         fi
+        HF_MODULES_CACHE=/home/ubuntu/.cache/huggingface/modules
 
         cmd=$(cat <<RUN_TEST_EOF
         #!/bin/bash
@@ -131,10 +144,13 @@ runs:
           --cpus=40 \
           --env GHA_RUNNER=${{ inputs.runner }} \
           --env HYDRA_FULL_ERROR=1 \
-          --env HF_HOME=/home/TestData/HF_HOME \
-          --env NEMO_HOME=/home/TestData/nemo_home \
+          --env HF_HOME=$HF_HOME \
+          --env HF_MODULES_CACHE=$HF_MODULES_CACHE \
+          --env NEMO_HOME=$NEMO_HOME \
           --env RUN_ID=${{ github.run_id }} \
           --env UV_NO_SYNC=1 \
+          --env SHARD_ID=${{ inputs.shard-id }} \
+          --env NUM_SHARDS=${{ inputs.num-shards }} \
           --workdir /opt/Megatron-Bridge \
           $VOLUME_ARGS \
           ${{ inputs.container-image }} \
@@ -159,7 +175,7 @@ runs:
 
         docker exec -t nemo_container_${{ github.run_id }} bash -c '
           set -e
-          export GH_TOKEN=${{ inputs.PAT }}
+          export GH_TOKEN=${{ inputs.github-token }}
 
           timeout $(( ${{ inputs.timeout }} * 60 ))s bash tests/${{ inputs.is_unit_test == 'true' && 'unit_tests' || format('functional_tests/launch_scripts/{0}', inputs.script_dir) }}/${{ inputs.script }}.sh || EXIT_CODE=$?
 

.github/scripts/test_docker_dependency_layers.sh

@@ -3,6 +3,7 @@ set -euo pipefail
 
 dockerfile="${1:-docker/Dockerfile.ci}"
 workflow="${2:-.github/workflows/cicd-main.yml}"
+fw_final_dockerfile="${3:-docker/Dockerfile.fw_final}"
 
 baseline_arg_line=$(grep -n '^ARG BASELINE_MCORE_REF$' "$dockerfile" | cut -d: -f1)
 baseline_clone_line=$(grep -n 'git clone --filter=blob:none --no-checkout' "$dockerfile" | cut -d: -f1)
@@ -24,14 +25,23 @@ trap 'rm -rf "$temporary_dir"' EXIT
 ((baseline_line < dispatched_copy_line))
 ((dispatched_copy_line < delta_line))
 
+fw_final_from_line=$(grep -n '^FROM ${NEMO_FW_FINAL_BASE_IMAGE} AS nemo_fw_final$' "$fw_final_dockerfile" | cut -d: -f1)
+fw_final_root_line=$(grep -n '^USER root$' "$fw_final_dockerfile" | cut -d: -f1)
+fw_final_clone_line=$(grep -n '^RUN git clone --depth 1 https://github.com/NVIDIA-NeMo/NeMo.git' "$fw_final_dockerfile" | cut -d: -f1)
+[[ -n "$fw_final_from_line" ]]
+[[ -n "$fw_final_root_line" ]]
+[[ -n "$fw_final_clone_line" ]]
+((fw_final_from_line < fw_final_root_line))
+((fw_final_root_line < fw_final_clone_line))
+
 grep -q -- '--mount=type=cache,target=/root/.cache/uv' "$dockerfile"
 if grep -q -- '--mount=type=secret,id=GH_TOKEN' "$dockerfile"; then
   echo "Baseline MCore clone must not require a GitHub token" >&2
   exit 1
 fi
 mcore_reinstall_line=$(grep -n 'uv pip install --no-deps --reinstall -e 3rdparty/Megatron-LM' "$dockerfile" | cut -d: -f1)
 helper_assertion_line=$(grep -n "find 3rdparty/Megatron-LM/megatron/core/datasets -maxdepth 1 -name 'helpers_cpp\*\.so'" "$dockerfile" | cut -d: -f1)
-final_copy_line=$(grep -n '^COPY --chmod=644 \. /opt/Megatron-Bridge$' "$dockerfile" | cut -d: -f1)
+final_copy_line=$(grep -nE '^COPY (--chmod=644|--chown=1000:1000) \. /opt/Megatron-Bridge$' "$dockerfile" | cut -d: -f1)
 [[ -n "$mcore_reinstall_line" ]]
 [[ -n "$helper_assertion_line" ]]
 [[ -n "$final_copy_line" ]]
@@ -232,17 +242,59 @@ if "$temporary_dir/revision-validator" "$revision_repo" not-a-full-sha; then
   exit 1
 fi
 
+if ! grep -A1 '^te = \[$' pyproject.toml | grep -Fxq '    "megatron-core[te]",'; then
+  echo "Bridge's TE extra must enable the selected MCore ref's TE extra" >&2
+  exit 1
+fi
+if ! grep -Fq \
+  '{ name = "megatron-core", extras = ["te"], marker = "extra == '\''te'\''", editable = "3rdparty/Megatron-LM" }' \
+  uv.lock; then
+  echo "Bridge's lock metadata must preserve the MCore TE extra" >&2
+  exit 1
+fi
 if grep -q 'transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@' pyproject.toml || \
   grep -q '^name = "transformer-engine"$' pyproject.toml; then
   echo "Bridge must inherit the TransformerEngine source and metadata from the selected MCore ref" >&2
   exit 1
 fi
 
+lock_package_field() {
+  local lock_file="$1"
+  local package="$2"
+  local field="$3"
+
+  awk -v package="$package" -v field="$field" '
+    $0 == "[[package]]" { in_package = 1; package_matches = 0; next }
+    in_package && $0 ~ /^\[\[/ { in_package = 0; package_matches = 0 }
+    in_package && $0 == "name = \"" package "\"" { package_matches = 1; next }
+    package_matches && index($0, field " = ") == 1 { print; exit }
+  ' "$lock_file"
+}
+
+mlm_te_source=$(lock_package_field 3rdparty/Megatron-LM/uv.lock transformer-engine source)
+bridge_te_source=$(lock_package_field uv.lock transformer-engine source)
+if [[ ! "$mlm_te_source" =~ ^source\ =\ \{\ git\ =\ \"https://github.com/NVIDIA/TransformerEngine\.git\?rev=[0-9a-f]{40}#[0-9a-f]{40}\"\ \}$ ]]; then
+  echo "Selected MCore lock must pin TransformerEngine to a full source revision" >&2
+  exit 1
+fi
+if [[ "$bridge_te_source" != "$mlm_te_source" ]]; then
+  echo "Bridge must lock the TransformerEngine source selected by MCore" >&2
+  exit 1
+fi
+
 composite_action=".github/actions/test-template/action.yml"
 if grep -qE 'mcore_(commit|ref)|MCORE_COMMIT|uv sync --all-extras --all-groups' "$composite_action"; then
   echo "Test template must use the already-validated MCore in the built image" >&2
   exit 1
 fi
+grep -Fq 'VOLUME_ARGS="--volume ${{ inputs.test-data-path }}:/home/TestData --env HF_HUB_OFFLINE=1"' "$composite_action"
+grep -q 'MOUNT_FS: ${{ inputs.is_unit_test == '\''false'\'' }}' "$composite_action"
+grep -q 'HF_HOME=/home/TestData/HF_HOME' "$composite_action"
+grep -q 'NEMO_HOME=/home/TestData/nemo_home' "$composite_action"
+grep -q 'HF_HOME=/home/ubuntu/.cache/huggingface' "$composite_action"
+grep -q 'NEMO_HOME=/home/ubuntu/.cache/nemo' "$composite_action"
+grep -q 'HF_MODULES_CACHE=/home/ubuntu/.cache/huggingface/modules' "$composite_action"
+grep -q -- '--env HF_MODULES_CACHE=\$HF_MODULES_CACHE' "$composite_action"
 
 # The baseline dependency layer must be structurally independent of the mutable
 # dispatched checkout. CI validates the ordering statically so this regression

.github/scripts/test_queue_bypass.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+workflow="${1:-.github/workflows/cicd-main.yml}"
+
+awk '
+  function check_job() {
+    if (!in_queue_chain || !has_success) {
+      return
+    }
+
+    checked++
+    if (!allows_queue_bypass) {
+      print job " uses success() but does not allow the intentional queue skip" > "/dev/stderr"
+      failed = 1
+    }
+  }
+
+  /^  [[:alnum:]_-]+:$/ {
+    check_job()
+    job = $1
+    sub(/:$/, "", job)
+    in_queue_chain = (in_queue_chain || job == "cicd-compute-build-matrix")
+    if (job == "Nemo_CICD_Test") {
+      in_queue_chain = 0
+    }
+    has_success = 0
+    allows_queue_bypass = 0
+    next
+  }
+
+  in_queue_chain && /success\(\)/ {
+    has_success = 1
+  }
+
+  in_queue_chain && /needs\.configure\.outputs\.skip_queue == '\''true'\'' && !failure\(\)/ {
+    allows_queue_bypass = 1
+  }
+
+  END {
+    check_job()
+    if (checked == 0) {
+      print "No queue-dependent success() conditions found" > "/dev/stderr"
+      exit 1
+    }
+    exit failed
+  }
+' "$workflow"