Merge branch 'main' into dnarayanan/latent_moe

Author: deepakn94

.github/CODEOWNERS

@@ -33,7 +33,7 @@ megatron/core/parallel_state.py @NVIDIA/core-adlr @NVIDIA/core-nemo
 
 megatron/core/post_training/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/post-training
 
-megatron/post_training/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/post-training
+megatron/post_training/ @NVIDIA/post-training
 
 .gitlab/ @NVIDIA/ci
 .github/ @NVIDIA/ci

.github/actions/action.yml

@@ -11,28 +11,28 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-name: "Test Template"
-description: "Template for running NeMo tests in a containerized environment"
+name: 'Test Template'
+description: 'Template for running NeMo tests in a containerized environment'
 
 inputs:
   container-image:
-    description: "Container image to use for test"
+    description: 'Container image to use for test'
     required: true
   timeout:
-    description: "Max runtime of test in minutes"
+    description: 'Max runtime of test in minutes'
     required: false
-    default: "30"
+    default: '30'
   script:
-    description: "Test script to execute"
+    description: 'Test script to execute'
     required: true
   is-optional:
-    description: "Pass this job on failure."
+    description: 'Pass this job on failure.'
     required: false
-    default: "false"
+    default: 'false'
   is_unit_test:
-    description: "Upload coverage as unit test"
+    description: 'Upload coverage as unit test'
     required: false
-    default: "false"
+    default: 'false'
   tag:
     description: Latest or legacy test suite
     required: true
@@ -43,11 +43,11 @@ inputs:
     description: Model to launch
     required: false
   PAT:
-    description: "GitHub Personal Access Token"
+    description: 'GitHub Personal Access Token'
     required: true
 
 runs:
-  using: "composite"
+  using: 'composite'
   steps:
     - name: Checkout repository
       uses: actions/checkout@v2
@@ -114,6 +114,16 @@ runs:
         HAS_RUN_TESTS_LABEL=$(gh pr view $PR_NUMBER --json labels | jq '[.labels[].name] | any(. == "Run tests")') || echo "false"
         echo "main=$HAS_RUN_TESTS_LABEL" | tee -a $GITHUB_OUTPUT
 
+    - name: Has Run functional tests label
+      shell: bash -x -e -u -o pipefail {0}
+      id: has-run-functional-tests-label
+      env:
+        GH_TOKEN: ${{ github.token }}
+      run: |
+        PR_NUMBER=${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number }}
+        HAS_RUN_FUNCTIONAL_TESTS_LABEL=$(gh pr view $PR_NUMBER --json labels | jq '[.labels[].name] | any(. == "Run functional tests")') || echo "false"
+        echo "main=$HAS_RUN_FUNCTIONAL_TESTS_LABEL" | tee -a $GITHUB_OUTPUT
+
     - name: Create run-script (e2e test)
       shell: bash -x -e -u -o pipefail {0}
       if: inputs.is_unit_test == 'false'
@@ -126,16 +136,19 @@ runs:
         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
+          ARGS=(
+            --scope mr-github
+            --enable-lightweight-mode
+          )
+        elif [ "${{ steps.has-run-functional-tests-label.outputs.main }}" == "true" ]; then
+          ARGS=(
+            --scope mr-github
+          )
+        else
+          ARGS=(
+            --scope mr-github-slim
+          )
+        fi
 
         export PYTHONPATH=$(pwd)
         export NEMORUN_HOME=$(pwd)

.github/actions/check-nvidia-sso-membership/action.yml

@@ -0,0 +1,139 @@
+name: 'Check NVIDIA SSO Membership'
+description: 'Check if a GitHub username exists in the NVIDIA SSO users list from github-audits'
+author: 'NVIDIA'
+
+inputs:
+  username:
+    description: 'GitHub username to check'
+    required: true
+  github_audits_repo:
+    description: 'Repository containing SSO users file'
+    required: false
+    default: 'NVIDIA-GitHub-Management/github-audits'
+  github_audits_version:
+    description: 'Release version tag'
+    required: false
+    default: 'v0.1.0'
+  sso_users_filename:
+    description: 'Filename of SSO users JSON'
+    required: false
+    default: 'users_sso.json'
+  github_token:
+    description: 'GitHub token with access to github-audits repo'
+    required: true
+
+outputs:
+  is_member:
+    description: 'Boolean - true if user is in NVIDIA SSO list, false otherwise'
+    value: ${{ steps.check-membership.outputs.is_member }}
+  is_org_member:
+    description: 'Boolean - true if user has NVIDIA or NVIDIA-NeMo in org_roles'
+    value: ${{ steps.check-membership.outputs.is_org_member }}
+  user_orgs:
+    description: 'Comma-separated list of orgs user is member of'
+    value: ${{ steps.check-membership.outputs.user_orgs }}
+  sso_file_available:
+    description: 'Boolean - true if SSO file was successfully downloaded'
+    value: ${{ steps.download-sso.outputs.sso_file_available }}
+  user_count:
+    description: 'Number of users in the SSO file (0 if download failed)'
+    value: ${{ steps.download-sso.outputs.user_count }}
+
+runs:
+  using: 'composite'
+  steps:
+    - name: Download NVIDIA SSO users from github-audits
+      id: download-sso
+      shell: bash
+      env:
+        GH_TOKEN: ${{ inputs.github_token }}
+      run: |
+        echo "Downloading ${{ inputs.sso_users_filename }} from ${{ inputs.github_audits_repo }} ${{ inputs.github_audits_version }} release..."
+
+        # Download the release asset using gh CLI
+        gh release download ${{ inputs.github_audits_version }} \
+          --repo ${{ inputs.github_audits_repo }} \
+          --pattern ${{ inputs.sso_users_filename }} \
+          --clobber 2>&1 || {
+            echo "ERROR: Failed to download ${{ inputs.sso_users_filename }} from github-audits release"
+            echo "sso_file_available=false" >> $GITHUB_OUTPUT
+            echo "user_count=0" >> $GITHUB_OUTPUT
+            exit 0
+          }
+
+        # Verify file was downloaded and is valid JSON
+        if [ ! -f ${{ inputs.sso_users_filename }} ]; then
+          echo "ERROR: ${{ inputs.sso_users_filename }} file not found after download"
+          echo "sso_file_available=false" >> $GITHUB_OUTPUT
+          echo "user_count=0" >> $GITHUB_OUTPUT
+          exit 0
+        fi
+
+        # Validate JSON structure
+        if ! jq -e 'type == "object"' ${{ inputs.sso_users_filename }} > /dev/null 2>&1; then
+          echo "ERROR: ${{ inputs.sso_users_filename }} is not a valid JSON object"
+          echo "sso_file_available=false" >> $GITHUB_OUTPUT
+          echo "user_count=0" >> $GITHUB_OUTPUT
+          exit 0
+        fi
+
+        USER_COUNT=$(jq 'length' ${{ inputs.sso_users_filename }})
+        echo "Successfully downloaded ${{ inputs.sso_users_filename }} with $USER_COUNT NVIDIA SSO users"
+        echo "sso_file_available=true" >> $GITHUB_OUTPUT
+        echo "user_count=$USER_COUNT" >> $GITHUB_OUTPUT
+
+    - name: Check if user is in SSO list
+      id: check-membership
+      shell: bash
+      run: |
+        USERNAME="${{ inputs.username }}"
+        SSO_FILE="${{ inputs.sso_users_filename }}"
+
+        echo "Checking if $USERNAME is in NVIDIA SSO users list..."
+
+        # Check if SSO file is available
+        if [ "${{ steps.download-sso.outputs.sso_file_available }}" != "true" ] || [ ! -f "$SSO_FILE" ]; then
+          echo "ERROR: $SSO_FILE not available - cannot check membership"
+          echo "is_member=false" >> $GITHUB_OUTPUT
+          echo "is_org_member=false" >> $GITHUB_OUTPUT
+          echo "user_orgs=" >> $GITHUB_OUTPUT
+          exit 0
+        fi
+
+        # Check if username exists as a key in the JSON object
+        if jq -e --arg user "$USERNAME" 'has($user)' "$SSO_FILE" > /dev/null 2>&1; then
+          echo "$USERNAME found in NVIDIA SSO users"
+          echo "is_member=true" >> $GITHUB_OUTPUT
+
+          # Extract and check org membership
+          IS_ORG_MEMBER=$(jq -r --arg user "$USERNAME" '
+            .[$user].org_roles // [] |
+            map(select(test("^(NVIDIA|NVIDIA-NeMo):Member$"))) |
+            length > 0
+          ' "$SSO_FILE")
+
+          USER_ORGS=$(jq -r --arg user "$USERNAME" '
+            .[$user].org_roles // [] |
+            map(split(":")[0]) |
+            unique |
+            join(",")
+          ' "$SSO_FILE")
+
+          echo "is_org_member=$IS_ORG_MEMBER" >> $GITHUB_OUTPUT
+          echo "user_orgs=$USER_ORGS" >> $GITHUB_OUTPUT
+
+          if [ "$IS_ORG_MEMBER" == "true" ]; then
+            echo "$USERNAME is a member of NVIDIA or NVIDIA-NeMo org"
+          else
+            echo "$USERNAME has @nvidia.com email but is not in NVIDIA or NVIDIA-NeMo org (orgs: $USER_ORGS)"
+          fi
+        else
+          echo "$USERNAME NOT found in NVIDIA SSO users"
+          echo "is_member=false" >> $GITHUB_OUTPUT
+          echo "is_org_member=false" >> $GITHUB_OUTPUT
+          echo "user_orgs=" >> $GITHUB_OUTPUT
+        fi
+
+branding:
+  icon: 'shield'
+  color: 'green'

.github/workflows/_build_test_publish_wheel.yml

@@ -1,5 +1,21 @@
 on:
   workflow_call:
+    inputs:
+      ref:
+        required: false
+        description: Ref (SHA or branch) to release
+        type: string
+        default: ${{ github.sha }}
+      dry-run:
+        required: false
+        description: Upload to PyPy Test instance
+        type: boolean
+        default: true
+      no-publish:
+        required: false
+        description: Do not publish the wheel
+        type: boolean
+        default: true
     secrets:
       TWINE_USERNAME:
         required: true
@@ -26,17 +42,18 @@ jobs:
       PACKAGE: ${{ matrix.PACKAGE }}
       IMAGE: ${{ matrix.IMAGE }}
       PLATFORM: ${{ matrix.PLATFORM }}
+      PUBLISH_DRYRUN: ${{ inputs.dry-run }}
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.ref }}
 
       - name: Build wheel
         id: build-wheel
         run: |
           set -x
 
-          PUBLISH_DRYRUN=yes
-
           if [ "$PACKAGE" = "megatron-core" ]; then
             ROOTDIR="megatron/core"
             BUILD_DIR="."
@@ -48,7 +65,7 @@ jobs:
             exit 1
           fi
 
-          if [ "$PUBLISH_DRYRUN" = "yes" ]; then
+          if [ "$PUBLISH_DRYRUN" = "true" ]; 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
@@ -123,26 +140,31 @@ jobs:
       - name: Upload wheels
         uses: actions/upload-artifact@v4
         with:
-          name: wheels-${{ matrix.PACKAGE }}-${{ matrix.PLATFORM }}
+          name: wheels-${{ matrix.PACKAGE }}-${{ matrix.PLATFORM }}-${{ inputs.dry-run && 'dry-run' || 'release' }}
           path: dist/
 
   publish-wheels:
     needs: [build-and-test-wheels]
     runs-on: ubuntu-latest
-    if: github.ref == 'refs/heads/main'
+    if: inputs.no-publish == false
     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
+          - PACKAGE: megatron-core
+            PLATFORM: arm64
+          - PACKAGE: megatron-core
+            PLATFORM: amd64
+          - PACKAGE: megatron-fsdp
+            IMAGE: quay.io/pypa/manylinux_2_28_x86_64
     env:
       PACKAGE: ${{ matrix.PACKAGE }}
     steps:
       - name: Download wheels
         uses: actions/download-artifact@v4
         with:
+          name: wheels-${{ matrix.PACKAGE }}-${{ matrix.PLATFORM }}-${{ inputs.dry-run && 'dry-run' || 'release' }}
           path: dist/
           merge-multiple: true