chore: Merge main into dev

.github/actions/action.yml

@@ -78,7 +78,7 @@ runs:
         export PYTHONPATH=$(pwd)
         export NEMORUN_HOME=$(pwd)
         pip install --no-cache-dir uv
-        uv sync --only-group test 
+        uv sync --only-group test
         uv run python tests/test_utils/python_scripts/launch_nemo_run_workload.py \
           --scope unit-tests \
           --model unit-tests \
@@ -90,7 +90,7 @@ runs:
 
         RUN_TEST_EOF
         )
-        echo "$cmd" | tee "job.sh"        
+        echo "$cmd" | tee "job.sh"
         echo "::endgroup::"
 
     - name: Get PR info
@@ -125,23 +125,34 @@ runs:
         #!/bin/bash
         set -euxo pipefail
 
+        if [ "${{ steps.has-run-tests-label.outputs.main }}" == "true" ]; then
+            ARGS=(
+              --scope mr-github
+              --enable-lightweight-mode
+            )
+          else
+            ARGS=(
+              --scope mr-slim
+              --enable-lightweight-mode
+            )
+          fi
+
         export PYTHONPATH=$(pwd)
         export NEMORUN_HOME=$(pwd)
         pip install --no-cache-dir uv
-        uv sync --only-group test 
+        uv sync --only-group test
         uv run python tests/test_utils/python_scripts/launch_nemo_run_workload.py \
-          --scope mr \
+          ${ARGS[@]} \
           --model ${{ inputs.model }} \
           --test-case ${{ inputs.test_case }} \
           --environment dev \
           --platform dgx_h100 \
           --container-image ${{ inputs.container-image }} \
           --data-dir /mnt/datadrive/TestData/megatron-lm/artifacts \
-          --enable-lightweight-mode
 
         RUN_TEST_EOF
         )
-        echo "$cmd" | tee "job.sh"        
+        echo "$cmd" | tee "job.sh"
         echo "::endgroup::"
 
     - name: Set timeout

.github/copy-pr-bot.yaml

@@ -1,3 +1,4 @@
 enabled: true
 auto_sync_draft: false
 auto_sync_ready: true
+# trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "gautham-kollu", "hxbai", "jaredcasper", "jkamalu", "jon-barker", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mkhona-nvidia", "pablo-garay", "pthombre", "rogerwaleffe", "sanandaraj5597", "santhnm2", "sbak5", "shanmugamr1992", "sidsingh-nvidia", "skyw", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"]

.github/workflows/_build_test_publish_wheel.yml

@@ -0,0 +1,156 @@
+on:
+  workflow_call:
+    secrets:
+      TWINE_USERNAME:
+        required: true
+      TWINE_PASSWORD:
+        required: true
+
+jobs:
+  build-and-test-wheels:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - PACKAGE: megatron-core
+            PLATFORM: arm64
+            IMAGE: quay.io/pypa/manylinux_2_28_aarch64
+          - PACKAGE: megatron-core
+            PLATFORM: amd64
+            IMAGE: quay.io/pypa/manylinux_2_28_x86_64
+          - PACKAGE: megatron-fsdp
+            IMAGE: quay.io/pypa/manylinux_2_28_x86_64
+            PLATFORM: amd64
+    runs-on: ${{ matrix.PLATFORM == 'amd64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
+    env:
+      PACKAGE: ${{ matrix.PACKAGE }}
+      IMAGE: ${{ matrix.IMAGE }}
+      PLATFORM: ${{ matrix.PLATFORM }}
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Build wheel
+        id: build-wheel
+        run: |
+          set -x
+
+          PUBLISH_DRYRUN=yes
+
+          if [ "$PACKAGE" = "megatron-core" ]; then
+            ROOTDIR="megatron/core"
+            BUILD_DIR="."
+          elif [ "$PACKAGE" = "megatron-fsdp" ]; then
+            ROOTDIR="megatron/core/distributed/fsdp/src/megatron_fsdp"
+            BUILD_DIR="megatron/core/distributed/fsdp/src"
+          else
+            echo Unknown package: $PACKAGE
+            exit 1
+          fi
+
+          if [ "$PUBLISH_DRYRUN" = "yes" ]; then
+            PRE_RELEASE=$(sed -n "s/.*PRE_RELEASE = '\(.*\)'/\1/p" $ROOTDIR/package_info.py)
+            sed -i "/^PRE_RELEASE/c\PRE_RELEASE = '${PRE_RELEASE}.dev$((RANDOM % 900000 + 100000))'" $ROOTDIR/package_info.py
+          fi
+
+          pushd $BUILD_DIR
+            rm LICENSE || true
+            docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE bash -c '\
+              for python_version in cp310 cp311 cp312 cp313; do \
+                /opt/python/${python_version}-${python_version}/bin/pip install --upgrade "setuptools>=80.0.0" build; \
+              done && \
+              for python_version in cp310 cp311 cp312 cp313; do \
+                /opt/python/${python_version}-${python_version}/bin/python -m build; \
+              done \
+            '
+
+            PLATFORM_WHEELS=$(find dist -name "*.whl" -not -name "*-none-any.whl")
+            if [ -n "$PLATFORM_WHEELS" ]; then
+                echo "Found platform wheels to repair: $PLATFORM_WHEELS"
+                docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE auditwheel repair $PLATFORM_WHEELS
+                docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE rm -rf dist/*.whl
+                docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE cp -a wheelhouse/* dist/
+            fi
+          popd
+
+          pushd $ROOTDIR
+            EXPECTED_RELEASE_NUMBER=$(python -c "import package_info; print(package_info.__version__)")
+          popd
+
+          echo "expected-release-number=$EXPECTED_RELEASE_NUMBER" | tee -a "${GITHUB_OUTPUT}"
+
+          if [ "$PACKAGE" = "megatron-fsdp" ]; then
+            mkdir -p dist/
+            cp -a megatron/core/distributed/fsdp/src/dist/* dist/
+          fi
+
+          ls -al dist/
+
+      - name: Test wheels
+        run: |
+          ls -al dist/
+
+          if [ "$PACKAGE" = "megatron-core" ]; then
+            ROOTPATH="megatron.core"
+            WHEEL_PREFIX="megatron_core"
+          elif [ "$PACKAGE" = "megatron-fsdp" ]; then
+            ROOTPATH="megatron_fsdp"
+            WHEEL_PREFIX="megatron_fsdp"
+          else
+            echo Unknown package: $PACKAGE
+            exit 1
+          fi
+
+          if [ "$PACKAGE" = "megatron-core" ]; then
+            if [[ "$PLATFORM" == "arm64" ]]; then
+              for file in dist/$WHEEL_PREFIX*cp310*aarch64.whl; do
+                pip install --no-cache-dir "$file"
+              done
+            else
+              for file in dist/$WHEEL_PREFIX*cp310*x86_64.whl; do
+                pip install --no-cache-dir "$file"
+              done
+            fi
+          else
+            pip install --no-cache-dir dist/$WHEEL_PREFIX*.whl
+          fi
+
+          sudo rm -rf megatron/
+
+          RELEASE_NUMBER=$(python -c "import $ROOTPATH; print($ROOTPATH.__version__)")
+          test "${{ steps.build-wheel.outputs.expected-release-number }}" == "$RELEASE_NUMBER"
+
+      - name: Upload wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheels-${{ matrix.PACKAGE }}-${{ matrix.PLATFORM }}
+          path: dist/
+
+  publish-wheels:
+    needs: [build-and-test-wheels]
+    runs-on: ubuntu-latest
+    environment: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/r')) && 'main' || 'public' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - PACKAGE: megatron_core
+          - PACKAGE: megatron_fsdp
+    env:
+      PACKAGE: ${{ matrix.PACKAGE }}
+    steps:
+      - name: Download wheels
+        uses: actions/download-artifact@v4
+        with:
+          path: dist/
+          merge-multiple: true
+
+      - name: Publish wheels
+        env:
+          TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }}
+          TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }}
+          TWINE_REPOSITORY: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/r')) && 'pypi' || 'testpypi' }}
+        run: |
+          ls -al dist/$PACKAGE*
+          pip install twine
+          twine upload -r $TWINE_REPOSITORY -u $TWINE_USERNAME -p $TWINE_PASSWORD dist/$PACKAGE*

.github/workflows/auto-update-copy-pr-bot.yml

@@ -0,0 +1,60 @@
+name: Auto Update Copy PR Bot
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "0 0 * * *"
+
+jobs:
+  auto-update-copy-pr-bot:
+    runs-on: ubuntu-latest
+    environment: nemo-ci
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Fetch list of members in mcore-reviewers team
+        shell: bash -euxo pipefail {0}
+        env:
+          GH_TOKEN: ${{ secrets.PAT }}
+        run: |
+          #!/bin/bash
+
+          get_members() {
+              local org=$1 team=$2 seen_file=$3    
+
+              gh api "/orgs/$org/teams/$team/members" --paginate --jq '.[].login' >> "$seen_file"
+
+              gh api "/orgs/$org/teams/$team/teams" --paginate --jq '.[].slug' | while read -r child; do
+                  get_members "$org" "$child" "$seen_file"
+              done
+
+              cat "$seen_file"
+          }
+
+          tmp=$(mktemp)
+          echo "" > final.txt
+          get_members "NVIDIA" "mcore-engineers" "$tmp" | sort -u >> final.txt && rm "$tmp"
+
+          tmp=$(mktemp)
+          get_members "NVIDIA" "mcore-reviewers" "$tmp" | sort -u >> final.txt && rm "$tmp"
+
+          cat final.txt | jq -sR 'split("\n") | map(select(. != "")) | flatten | unique'
+
+          export TRUSTEES=$(cat final.txt | jq -csR 'split("\n") | map(select(. != "")) | flatten | unique')
+          yq '.trustees_override = env(TRUSTEES)' .github/copy-pr-bot.yaml | yq -o yaml > .github/copy-pr-bot.yaml.new
+
+          mv .github/copy-pr-bot.yaml.new .github/copy-pr-bot.yaml
+
+      - name: Commit changes
+        run: |
+          git remote set-url origin https://x-access-token:${{ secrets.PAT }}@github.com/NVIDIA/Megatron-LM.git
+          git config --global user.name "GitHub Actions"
+          git config --global user.email "github-actions[bot]@users.noreply.github.com"
+          git add .github/copy-pr-bot.yaml
+          if git diff --cached --exit-code --quiet; then
+            echo "No changes to commit. Exiting gracefully."
+            exit 0
+          fi
+          git commit -m "Update copy-pr-bot.yaml [skip ci]"
+          git push

.github/workflows/build-test-publish-wheel.yml

@@ -37,32 +37,15 @@ jobs:
     uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/[email protected]
     if: github.repository == 'NVIDIA/Megatron-LM'
 
-  build-test-publish-wheel:
+  build-test-publish-wheels:
     needs: [pre-flight]
-    if: |
-      !(needs.pre-flight.outputs.docs_only == 'true'
-      || needs.pre-flight.outputs.is_merge_group == 'true'
-      || needs.pre-flight.outputs.is_deployment_workflow == 'true')
-      && github.repository == 'NVIDIA/Megatron-LM'
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/[email protected]
-    with:
-      dry-run: true
-      python-package: megatron.core
-      python-version: "3.10"
-      packaging: uv
-      no-publish: ${{ !(github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/r')) }}
-      custom-container: nvcr.io/nvidia/pytorch:25.05-py3
-      no-build-isolation: true
-      runner: linux-amd64-cpu16
+    uses: ./.github/workflows/_build_test_publish_wheel.yml
     secrets:
       TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }}
       TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }}
-      SLACK_WEBHOOK: ${{ secrets.SLACK_RELEASE_ENDPOINT }}
-      SLACK_WEBHOOK_ADMIN: ${{ secrets.SLACK_WEBHOOK_ADMIN }}
-      GH_TOKEN: ${{ secrets.PAT }}
 
   build-test-publish-wheel-summary:
-    needs: [pre-flight, build-test-publish-wheel]
+    needs: [pre-flight, build-test-publish-wheels]
     if: |
       (
         needs.pre-flight.outputs.docs_only == 'true'
@@ -74,7 +57,14 @@ jobs:
       && !cancelled()
     runs-on: ubuntu-latest
     steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
       - name: Result
+        env:
+          GH_TOKEN: ${{ github.token }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          SKIPPING_IS_ALLOWED: ${{ needs.pre-flight.outputs.docs_only == 'true' || needs.pre-flight.outputs.is_deployment_workflow == 'true' || needs.pre-flight.outputs.is_merge_group == 'true' || needs.pre-flight.outputs.is_ci_workload == 'true' }}
         run: |
           FAILED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion != "success")] | length') || echo 0