diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 7496893749c..f02d8e1f5b5 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -9,6 +9,8 @@ megatron/core/ssm/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/hybrid-mamba megatron/core/datasets/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/datasets +megatron/core/tokenizers/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/tokenizers + megatron/core/distributed/fsdp/ @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/megatron-fsdp megatron/core/transformer/fsdp_dtensor_checkpoint.py @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/megatron-fsdp diff --git a/.github/actions/action.yml b/.github/actions/action.yml index f3e42e5843d..4a838b24d95 100644 --- a/.github/actions/action.yml +++ b/.github/actions/action.yml @@ -48,7 +48,9 @@ inputs: is_ci_workload: description: "Is CI workload" required: true - + is_merge_group: + description: "Is merge group" + required: true runs: using: "composite" steps: @@ -117,8 +119,10 @@ runs: export PYTHONPATH=$(pwd) export NEMORUN_HOME=$(pwd) export NCCL_DEBUG=INFO - pip install --no-cache-dir uv - uv sync --only-group test + pip install --no-cache-dir "uv<0.9.29" + uv venv .venv + uv cache clean + uv sync --no-cache --only-group test uv run python tests/test_utils/python_scripts/launch_nemo_run_workload.py \ --scope unit-tests \ --model unit-tests \ @@ -177,7 +181,12 @@ runs: #!/bin/bash set -euxo pipefail - if [ "${{ steps.has-run-tests-label.outputs.main }}" == "true" ]; then + if [ "${{ inputs.is_merge_group }}" == "true" ]; then + ARGS=( + --scope mr-github + --n-repeat 1 + ) + elif [ "${{ steps.has-run-tests-label.outputs.main }}" == "true" ]; then ARGS=( --scope mr-github --enable-lightweight-mode @@ -197,8 +206,10 @@ runs: export PYTHONPATH=$(pwd) export NEMORUN_HOME=$(pwd) - pip install --no-cache-dir uv - uv sync --only-group test + pip install --no-cache-dir "uv<0.9.29" + uv venv .venv + uv cache clean + uv sync --no-cache --only-group test uv run python tests/test_utils/python_scripts/launch_nemo_run_workload.py \ ${ARGS[@]} \ --model ${{ inputs.model }} \ diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml index e305bb25ce4..d198668660f 100644 --- a/.github/copy-pr-bot.yaml +++ b/.github/copy-pr-bot.yaml @@ -1,4 +1,4 @@ enabled: true auto_sync_draft: false auto_sync_ready: true -trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "jalbericiola", "janEbert", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "kwyss-nvidia", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "parthmannan", "prajwal1210", "pthombre", "rogerwaleffe", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "shengf-nv", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"] +trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "RPrenger", "ShriyaRishab", "Victarry", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cspades", "cuichenx", "deepakn94", "dimapihtar", "dingqingy-nv", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hxbai", "ilml", "jalbericiola", "janEbert", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kanz-nv", "kevalmorabia97", "ko3n1g", "kunlunl", "kvareddy", "kwyss-nvidia", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mehraakash", "mkhona-nvidia", "parthmannan", "prajwal1210", "pthombre", "rogerwaleffe", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "sharathts", "shengf-nv", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "trintamaki", "tylerpoon", "wdykas", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yuzhongw-nvidia", "zhongbozhu"] diff --git a/.github/oncall_schedule.json b/.github/oncall_schedule.json index 5fa49e966bc..6a47241a575 100644 --- a/.github/oncall_schedule.json +++ b/.github/oncall_schedule.json @@ -1,18 +1,10 @@ [ { - "user": "dimapihtar", - "date": "2026-01-28" - }, - { - "user": "gautham-kollu", - "date": "2026-02-04" - }, - { - "user": "janEbert", + "user": "Phlip79", "date": "2026-02-11" }, { - "user": "Phlip79", + "user": "janEbert", "date": "2026-02-18" }, { @@ -46,5 +38,13 @@ { "user": "BoxiangW", "date": "2026-04-15" + }, + { + "user": "Phlip79", + "date": "2026-04-22" + }, + { + "user": "asolergi-nv", + "date": "2026-04-29" } ] diff --git a/.github/workflows/check_api_backwards_compatibility_workflow.yml b/.github/workflows/check_api_backwards_compatibility_workflow.yml deleted file mode 100644 index 44340bdedc5..00000000000 --- a/.github/workflows/check_api_backwards_compatibility_workflow.yml +++ /dev/null @@ -1,276 +0,0 @@ -# Temporarily disable this check until we can enforce it on PRs -# -# name: API Compatibility Check - -# on: -# push: -# branches: -# - dev -# - main -# - 'pull-request/[0-9]+' -# - 'deploy-release/*' -# merge_group: -# types: [checks_requested] - -# # Allow manual trigger -# workflow_dispatch: -# inputs: -# baseline: -# description: 'Baseline git reference (tag/branch/commit)' -# required: true - -# jobs: -# pre-flight: -# name: Pre-flight check -# runs-on: ubuntu-latest -# outputs: -# should_skip: ${{ steps.check_files.outputs.should_skip }} -# steps: -# - name: Checkout code -# uses: actions/checkout@v4 -# with: -# fetch-depth: 0 - -# - name: Check if relevant files changed -# id: check_files -# run: | -# # For manual triggers, never skip -# if [ "${{ github.event_name }}" == "workflow_dispatch" ]; then -# echo "should_skip=false" >> $GITHUB_OUTPUT -# echo "Manual trigger - will run compatibility check" -# exit 0 -# fi - -# # Determine base SHA based on event type -# if [ "${{ github.event_name }}" == "merge_group" ]; then -# BASE_SHA="${{ github.event.merge_group.base_sha }}" -# echo "Merge group event - comparing against base: $BASE_SHA" -# else -# # For push events, use merge-base to find common ancestor -# # This ensures we only detect changes actually made in this PR branch, -# # not changes that happened in main after the branch was created -# BASE_SHA=$(git merge-base origin/main HEAD 2>/dev/null || echo "") -# if [ -z "$BASE_SHA" ]; then -# # Fallback for pull-request/* branches targeting dev -# BASE_SHA=$(git merge-base origin/dev HEAD 2>/dev/null || echo "") -# fi -# echo "Push event - comparing against merge-base: $BASE_SHA" -# fi - -# if [ -z "$BASE_SHA" ]; then -# echo "Could not determine base SHA - will run compatibility check" -# echo "should_skip=false" >> $GITHUB_OUTPUT -# exit 0 -# fi - -# # Check for changes in megatron/core Python files (excluding tests and legacy) -# # Note: Using both *.py and **/*.py to match files at root and in subdirectories -# CHANGED_FILES=$(git diff --name-only "$BASE_SHA" HEAD -- \ -# 'megatron/core/*.py' \ -# 'megatron/core/**/*.py' \ -# ':!megatron/core/tests/**' \ -# ':!megatron/legacy/**' 2>/dev/null || echo "") - -# if [ -z "$CHANGED_FILES" ]; then -# echo "should_skip=true" >> $GITHUB_OUTPUT -# echo "No relevant megatron/core files changed - will skip compatibility check" -# else -# echo "should_skip=false" >> $GITHUB_OUTPUT -# echo "Relevant files changed:" -# echo "$CHANGED_FILES" -# fi - -# check-compatibility: -# needs: [pre-flight] -# if: needs.pre-flight.outputs.should_skip != 'true' -# name: "OPTIONAL: Check API Backward Compatibility" -# runs-on: ubuntu-latest - -# # ============================================================================ -# # Configuration Parameters (modify here) -# # ============================================================================ -# env: -# # Default baseline for automatic PR checks -# # Can be: branch name (e.g., 'main'), commit hash, or tag -# # Will be resolved to commit hash during execution -# DEFAULT_BASELINE: '5ab481cb45efc72add12f8ba0378e849b3d2bc50' -# # Tag pattern for auto-detection (e.g., 'core_r*', 'core_v*') -# TAG_PATTERN: 'core_v*' -# # Tag regex filter (e.g., '^core_v[0-9]+\.[0-9]+\.[0-9]+$' for stable versions only) -# TAG_REGEX_FILTER: '^core_v[0-9]+\.[0-9]+\.[0-9]+$' -# # ============================================================================ - -# steps: -# - name: Checkout code -# uses: actions/checkout@v4 -# with: -# fetch-depth: 0 # Need full history to access baseline ref - -# - name: Set up Python -# uses: actions/setup-python@v5 -# with: -# python-version: '3.12' - -# - name: Install griffe -# run: | -# python -m pip install --upgrade pip -# python -m pip install griffe -# python -c "import griffe; print('Griffe installed successfully')" -# python -c "from griffe import Object; print('Object import successful')" || echo "Object import from griffe failed" -# python -c "from griffe.dataclasses import Object; print('Object import from dataclasses successful')" || echo "Object import from dataclasses failed" - -# - name: Determine baseline reference -# id: baseline -# run: | -# if [ "${{ github.event_name }}" == "workflow_dispatch" ]; then -# # Use manually specified baseline (branch, tag, or commit hash) -# BASELINE_REF="${{ github.event.inputs.baseline }}" -# else -# # Use the configured default baseline -# BASELINE_REF="${{ env.DEFAULT_BASELINE }}" - -# # Uncomment below to auto-detect from tags instead: -# # BASELINE_REF=$(git tag -l '${{ env.TAG_PATTERN }}' | grep -E '${{ env.TAG_REGEX_FILTER }}' | sort -V | tail -1) -# # if [ -z "$BASELINE_REF" ]; then -# # echo "Warning: No tags matching pattern found. Using default: ${{ env.DEFAULT_BASELINE }}" >&2 -# # BASELINE_REF="${{ env.DEFAULT_BASELINE }}" -# # fi -# fi - -# # Resolve baseline to commit hash (works for branches, tags, or commit hashes) -# BASELINE_HASH=$(git rev-parse "$BASELINE_REF") - -# echo "baseline=$BASELINE_HASH" >> $GITHUB_OUTPUT -# echo "Using baseline: $BASELINE_REF (resolved to commit: $BASELINE_HASH)" - -# - name: Run compatibility check -# id: compat_check -# run: | -# # Save output to file for later display -# python scripts/check_api_backwards_compatibility.py \ -# --baseline ${{ steps.baseline.outputs.baseline }} \ -# --verbose 2>&1 | tee compat_check_output.txt - -# # Capture exit code -# EXIT_CODE=${PIPESTATUS[0]} -# echo "exit_code=$EXIT_CODE" >> $GITHUB_OUTPUT -# exit $EXIT_CODE -# continue-on-error: true - -# - name: Fail job if breaking changes detected -# if: steps.compat_check.outcome == 'failure' -# run: | -# echo "" -# echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -# echo "🔍 WHAT IS THIS CHECK?" -# echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -# echo "" -# echo "This check ensures that changes to Megatron Core's public API do not" -# echo "break backward compatibility for users. It compares your PR against" -# echo "the latest stable release to detect breaking changes in:" -# echo "" -# echo " • Function signatures (parameters, order, types)" -# echo " • Class structures and methods" -# echo " • Return types and public interfaces" -# echo "" -# echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -# echo "🛠️ HOW TO FIX THIS" -# echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -# echo "" -# echo "Choose ONE of these resolution strategies:" -# echo "" -# echo "1️⃣ REVERT THE BREAKING CHANGE (Recommended)" -# echo " → Modify your code to preserve backward compatibility" -# echo " → Add new parameters as optional (with defaults)" -# echo " → Keep existing parameters in the same order" -# echo "" -# echo "2️⃣ MARK AS INTERNAL API (If this is internal code)" -# echo " → Add @internal_api decorator from megatron.core.utils" -# echo "" -# echo " Example (for classes):" -# echo " from megatron.core.utils import internal_api" -# echo "" -# echo " @internal_api" -# echo " class ExperimentalFeature:" -# echo " pass" -# echo "" -# echo " Example (for functions):" -# echo " from megatron.core.utils import internal_api" -# echo "" -# echo " @internal_api" -# echo " def internal_helper_function():" -# echo " pass" -# echo "" -# echo "3️⃣ MARK AS EXPERIMENTAL API (If this is experimental code)" -# echo " → Add @experimental_api decorator from megatron.core.utils" -# echo "" -# echo " Example:" -# echo " from megatron.core.utils import experimental_api" -# echo "" -# echo " @experimental_api" -# echo " class ExperimentalFeature:" -# echo " pass" -# echo "" -# echo "4️⃣ USE DEPRECATION (For gradual API changes)" -# echo " → Add @deprecated decorator for transition period" -# echo " → Example:" -# echo " from megatron.core.utils import deprecated" -# echo "" -# echo " @deprecated(version='1.0', removal_version='2.0'," -# echo " alternative='new_function')" -# echo " def old_function():" -# echo " pass" -# echo "" -# echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -# echo "📋 BREAKING CHANGES DETECTED" -# echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -# echo "" -# cat compat_check_output.txt -# echo "" -# echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -# echo "📚 MORE INFORMATION" -# echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -# echo "" -# echo "📖 Full documentation: docs/api-backwards-compatibility-check.md" -# echo "🔧 Checker script: scripts/check_api_backwards_compatibility.py" -# echo "❓ Questions? Check the docs or ask in #megatron-core" -# echo "" - -# echo "::error::Breaking API changes detected. Please review the output above and choose a resolution strategy." -# exit 1 - -# - name: Success message -# if: steps.compat_check.outcome == 'success' -# run: | -# echo "::notice::✅ No breaking API changes detected!" - -# api-backward-compatibility-summary: -# needs: [pre-flight, check-compatibility] -# runs-on: ubuntu-latest -# name: "OPTIONAL: API Backward Compatibility Check Summary" -# if: always() && !cancelled() -# steps: -# - name: Checkout -# uses: actions/checkout@v4 - -# - name: Validate workflow result -# shell: bash -x -e -u -o pipefail {0} -# env: -# GH_TOKEN: ${{ github.token }} -# SKIPPING_IS_ALLOWED: ${{ needs.pre-flight.outputs.should_skip == 'true' }} -# run: | -# FAILED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion != "success" and .name != "OPTIONAL: API Backward Compatibility Check Summary")] | length') || echo 0 - -# if [ "${FAILED_JOBS:-0}" -eq 0 ] || [ "$SKIPPING_IS_ALLOWED" == "true" ]; then -# if [ "$SKIPPING_IS_ALLOWED" == "true" ]; then -# echo "✅ Compatibility check was skipped (no relevant files changed)" -# else -# echo "✅ All checks passed successfully" -# fi -# exit 0 -# else -# echo "❌ Found $FAILED_JOBS failed job(s)" -# gh run view $GITHUB_RUN_ID --json jobs --jq '.jobs[] | select(.status == "completed" and .conclusion != "success" and .name != "OPTIONAL: API Backward Compatibility Check Summary") | .name' -# exit 1 -# fi - diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 16e2051e4e2..d6519cdb5ef 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -18,8 +18,6 @@ on: - cron: 0 0 * * * push: branches: - - dev - - main - "pull-request/[0-9]+" - "deploy-release/*" merge_group: @@ -27,7 +25,7 @@ on: workflow_dispatch: concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.event.label.name || 'main' }}-${{ github.event_name }} + group: ${{ github.workflow }}-${{ github.head_ref || github.ref || github.event.pull_request.number }} cancel-in-progress: true permissions: @@ -179,10 +177,11 @@ jobs: cicd-wait-in-queue: runs-on: ubuntu-latest needs: [pre-flight, linting] - environment: ${{ needs.pre-flight.outputs.is_merge_group == 'true' && 'merge-gate' || 'test' }} + environment: "test" if: | !(needs.pre-flight.outputs.is_ci_workload == 'true' || needs.pre-flight.outputs.is_deployment_workflow == 'true' + || needs.pre-flight.outputs.is_merge_group == 'true' || needs.pre-flight.outputs.docs_only == 'true') steps: - name: Running CI tests @@ -194,12 +193,15 @@ jobs: needs: [is-not-external-contributor, pre-flight, cicd-wait-in-queue] runs-on: ${{ needs.is-not-external-contributor.outputs.selected_runner }} if: | - ( + needs.is-not-external-contributor.result != 'cancelled' + && needs.pre-flight.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && ( success() || needs.pre-flight.outputs.is_ci_workload == 'true' + || needs.pre-flight.outputs.is_merge_group == 'true' || needs.pre-flight.outputs.force_run_all == 'true' ) - && needs.pre-flight.outputs.is_merge_group == 'false' && !cancelled() steps: - name: Get PR info @@ -212,9 +214,12 @@ jobs: id: sha env: IS_PR: ${{ startsWith(github.ref, 'refs/heads/pull-request/') }} + IS_MERGE_GROUP: ${{ github.event_name == 'merge_group' }} run: | if [[ "$IS_PR" == "true" ]]; then SHA=${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').merge_commit_sha }} + elif [[ "$IS_MERGE_GROUP" == "true" ]]; then + SHA=${{ github.event.merge_group.head_sha }} else SHA=${GITHUB_SHA} fi @@ -332,12 +337,15 @@ jobs: - cicd-wait-in-queue - cicd-container-build if: | - ( + needs.pre-flight.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && needs.cicd-container-build.result != 'cancelled' + && ( success() || needs.pre-flight.outputs.is_ci_workload == 'true' || needs.pre-flight.outputs.force_run_all == 'true' + || needs.pre-flight.outputs.is_merge_group == 'true' ) - && needs.pre-flight.outputs.is_merge_group == 'false' && !cancelled() steps: - name: Checkout @@ -363,12 +371,17 @@ jobs: timeout-minutes: 60 name: "${{ matrix.bucket }} - latest" if: | - ( + needs.is-not-external-contributor.result != 'cancelled' + && needs.pre-flight.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && needs.cicd-container-build.result != 'cancelled' + && needs.cicd-parse-unit-tests.result != 'cancelled' + && ( success() || needs.pre-flight.outputs.is_ci_workload == 'true' || needs.pre-flight.outputs.force_run_all == 'true' + || needs.pre-flight.outputs.is_merge_group == 'true' ) - && needs.pre-flight.outputs.is_merge_group == 'false' && !cancelled() env: PIP_DISABLE_PIP_VERSION_CHECK: 1 @@ -390,14 +403,17 @@ jobs: cicd-parse-integration-tests: runs-on: ubuntu-latest - timeout-minutes: 60 needs: - pre-flight - cicd-wait-in-queue - cicd-container-build - cicd-unit-tests-latest if: | - ( + needs.pre-flight.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && needs.cicd-container-build.result != 'cancelled' + && needs.cicd-unit-tests-latest.result != 'cancelled' + && ( success() || needs.pre-flight.outputs.is_ci_workload == 'true' || needs.pre-flight.outputs.force_run_all == 'true' @@ -439,7 +455,7 @@ jobs: id: main env: HAS_RUN_TESTS_LABEL: ${{ steps.has-run-tests-label.outputs.main }} - HAS_RUN_FUNCTIONAL_TESTS_LABEL: ${{ steps.has-run-functional-tests-label.outputs.main }} + HAS_RUN_FUNCTIONAL_TESTS_LABEL: ${{ steps.has-run-functional-tests-label.outputs.main == 'true' || needs.pre-flight.outputs.is_merge_group == 'true' }} run: | export PYTHONPATH=$(pwd) @@ -480,6 +496,7 @@ jobs: echo "integration-tests=$(cat integration-tests.json)" | tee -a "$GITHUB_OUTPUT" cicd-integration-tests-latest: + timeout-minutes: 60 strategy: fail-fast: false matrix: @@ -497,12 +514,17 @@ jobs: PIP_NO_PYTHON_VERSION_WARNING: 1 PIP_ROOT_USER_ACTION: ignore if: | - ( + needs.is-not-external-contributor.result != 'cancelled' + && needs.pre-flight.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && needs.cicd-parse-integration-tests.result != 'cancelled' + && needs.cicd-unit-tests-latest.result != 'cancelled' + && ( success() || needs.pre-flight.outputs.is_ci_workload == 'true' || needs.pre-flight.outputs.force_run_all == 'true' + || needs.pre-flight.outputs.is_merge_group == 'true' ) - && needs.pre-flight.outputs.is_merge_group == 'false' && !cancelled() steps: - name: Checkout @@ -518,6 +540,7 @@ jobs: PAT: ${{ secrets.PAT }} container-image: ${{ env.container-registry }}/megatron-lm:${{ github.sha }} is_ci_workload: ${{ needs.pre-flight.outputs.is_ci_workload }} + is_merge_group: ${{ needs.pre-flight.outputs.is_merge_group }} Nemo_CICD_Test: needs: @@ -548,8 +571,8 @@ jobs: 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 == "failure")] | length') || echo 0 - SKIPPED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion == "skipped")] | length') || echo 0 + FAILED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion == "failure" and .name != "merge-queue-notification")] | length') || echo 0 + SKIPPED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion == "skipped" and .name != "merge-queue-notification")] | length') || echo 0 if [ "${FAILED_JOBS:-0}" -eq 0 ] && ([ "${SKIPPED_JOBS:-0}" -eq 0 ] || [ "$SKIPPING_IS_ALLOWED" == "true" ]); then echo "✅ All previous jobs completed successfully" @@ -641,6 +664,34 @@ jobs: .coverage include-hidden-files: true + merge-queue-notification: + runs-on: ubuntu-latest + if: github.event_name == 'merge_group' + permissions: + pull-requests: write + steps: + - name: Extract PR number from merge group + id: get-pr-number + run: | + # Extract PR number from merge group head_ref (format: refs/heads/gh-readonly-queue/main/pr--) + PR_NUMBER=$(echo "${{ github.event.merge_group.head_ref }}" | sed -n 's/.*\/pr-\([0-9]*\)-.*/\1/p') + echo "pr_number=$PR_NUMBER" >> $GITHUB_OUTPUT + + - name: Comment on PR with action run URL + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.PAT }} + script: | + const prNumber = ${{ steps.get-pr-number.outputs.pr_number }}; + const runUrl = `https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}`; + + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: prNumber, + body: `🔄 Merge queue validation started!\n\nYou can track the progress here: ${runUrl}` + }); + cleanup-taint-node: runs-on: ${{ needs.is-not-external-contributor.outputs.selected_runner }} needs: diff --git a/.github/workflows/oncall-rotation.yml b/.github/workflows/oncall-rotation.yml index 46a45810ad1..71ae094e6c8 100644 --- a/.github/workflows/oncall-rotation.yml +++ b/.github/workflows/oncall-rotation.yml @@ -45,7 +45,10 @@ jobs: # Slack token for updating the Slack usergroup SLACK_TOKEN: ${{ secrets.ONCALL_SLACK_TOKEN }} run: | - pip install --no-cache-dir uv + pip install --no-cache-dir "uv<0.9.29" + uv venv .venv + uv cache clean + uv sync --no-cache uv run --with slack-sdk python .github/scripts/oncall_manager.py rotate - name: Commit and Push changes diff --git a/.github/workflows/sync-team-usergroups.yml b/.github/workflows/sync-team-usergroups.yml index 8b08182dceb..1c6cecaeb7a 100644 --- a/.github/workflows/sync-team-usergroups.yml +++ b/.github/workflows/sync-team-usergroups.yml @@ -35,5 +35,8 @@ jobs: GH_TOKEN: ${{ secrets.NVIDIA_MCORE_ONCALL_TOKEN || secrets.PAT || secrets.GITHUB_TOKEN }} SLACK_TOKEN: ${{ secrets.ONCALL_SLACK_TOKEN }} run: | - pip install --no-cache-dir uv + pip install --no-cache-dir "uv<0.9.29" + uv venv .venv + uv cache clean + uv sync --no-cache uv run --with slack-sdk python .github/scripts/sync_team_usergroups.py diff --git a/.gitlab/stages/01.build.yml b/.gitlab/stages/01.build.yml index 20252e7d045..61521295a93 100644 --- a/.gitlab/stages/01.build.yml +++ b/.gitlab/stages/01.build.yml @@ -121,6 +121,7 @@ test:build_image: KUBERNETES_SERVICE_MEMORY_LIMIT: 90Gi SHARED_PATH: /builds/$CI_PROJECT_PATH/shared script: + - apk add skopeo - | set -x @@ -132,6 +133,11 @@ test:build_image: ${IMAGE}:${CI_PIPELINE_ID}-arm64 docker manifest push ${IMAGE}:${CI_PIPELINE_ID} + + if [[ "$CI_COMMIT_BRANCH" == "ci-rebuild-mcore-nemo-image" || "$CI_COMMIT_BRANCH" == "main" || "$CI_COMMIT_BRANCH" == "dev" ]]; then + skopeo copy --all docker://${IMAGE}:${CI_PIPELINE_ID} docker://${IMAGE}:${CI_COMMIT_BRANCH} + fi + - echo "MCORE_MR_COMMIT=$CI_COMMIT_SHA" | tee -a build.env - echo "MCORE_BACKWARDS_COMMIT=$MCORE_BACKWARDS_COMMIT" | tee -a build.env - cat build.env diff --git a/CHANGELOG.md b/CHANGELOG.md index babdc18b8a4..478ee56ada0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,20 @@ # Changelog +## NVIDIA Megatron Core 0.15.3 + +This release addresses known security issues. For the latest NVIDIA Vulnerability Disclosure Information visit , for acknowledgement please reach out to the NVIDIA PSIRT team at + +## NVIDIA Megatron Core 0.15.2 + +* Bug fixes + * Various small fixes for Megatron-FSDP. [#2346](https://github.com/NVIDIA/Megatron-LM/pull/2346) + * [Megatron-FSDP] Support both old and new DeviceMesh APIs. [#2575](https://github.com/NVIDIA/Megatron-LM/pull/2575) + * [Megatron-FSDP] Build default FSDP DeviceMesh, and remove model arg from `fully_shard_optimizer()`. [#2471](https://github.com/NVIDIA/Megatron-LM/pull/2471) + +## NVIDIA Megatron Core 0.15.1 + +Yanked release. + ## NVIDIA Megatron Core 0.15.0 * Features diff --git a/README.md b/README.md index 6fa300a6d4d..a5cdf4ad438 100644 --- a/README.md +++ b/README.md @@ -13,14 +13,13 @@ Megatron-LM and Megatron Core ## About -**Megatron-Core (MCore)**: Composable library with GPU-optimized building blocks for custom training frameworks. -You can install this library using pip or use it within the Megatron-LM GitHub repository. +This repository contains two components: **Megatron-LM** and **Megatron Core**. -**Megatron-LM**: Reference implementation that includes end-to-end examples utilizing Megatron Core. +**Megatron-LM** is a reference example that includes Megatron Core plus pre-configured training scripts. Best for research teams, learning distributed training, and quick experimentation. -**Megatron-Bridge**: Training library with bidirectional Hugging Face ↔ Megatron checkpoint conversion, flexible training loops, and example model training recipes. +**Megatron Core** is a composable library with GPU-optimized building blocks for custom training frameworks. It provides transformer building blocks, advanced parallelism strategies (TP, PP, DP, EP, CP), mixed precision support (FP16, BF16, FP8, FP4), and model architectures. Best for framework developers and ML engineers building custom training pipelines. -For more information, refer to [Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge). +**[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** provides bidirectional Hugging Face ↔ Megatron checkpoint conversion with production-ready recipes. ## Quick Start @@ -44,7 +43,8 @@ Install Megatron Core with pip: # Latest News -- **[2025/12]** 🎉 **Megatron Core development has moved to GitHub!** All development and CI now happens in the open. We welcome community contributions. +- **[2026/01]** **[Dynamic Context Parallelism](https://developer.nvidia.com/blog/speeding-up-variable-length-training-with-dynamic-context-parallelism-and-nvidia-megatron-core/)** - Up to 1.48x speedup for variable-length sequence training with adaptive CP sizing. +- **[2025/12]** **Megatron Core development has moved to GitHub!** All development and CI now happens in the open. We welcome community contributions. - **[2025/10]** **[Megatron Dev Branch](https://github.com/NVIDIA/Megatron-LM/tree/dev)** - early access branch with experimental features. - **[2025/10]** **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Bidirectional converter for interoperability between Hugging Face and Megatron checkpoints, featuring production-ready recipes for popular models. - **[2025/08]** **[MoE Q3-Q4 2025 Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - Comprehensive roadmap for MoE features including DeepSeek-V3, Qwen3, advanced parallelism strategies, FP8 optimizations, and Blackwell performance enhancements. @@ -57,13 +57,11 @@ Install Megatron Core with pip: - **[2024/07]** Megatron Core v0.7 improves scalability and training resiliency and adds support for multimodal training ([blog](https://developer.nvidia.com/blog/train-generative-ai-models-more-efficiently-with-new-nvidia-Megatron-Core-functionalities/)). - **[2024/06]** Megatron Core added supports for Mamba-based models. Check out our paper [An Empirical Study of Mamba-based Language Models](https://arxiv.org/pdf/2406.07887) and [code example](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba). -- **[2024/01 Announcement]** NVIDIA has released the core capabilities in **Megatron-LM** into [**Megatron Core**](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core) in this repository. Megatron Core expands upon Megatron-LM's GPU-optimized techniques with more cutting-edge innovations on system-level optimizations, featuring composable and modular APIs. Explore the [Megatron Core intro](#Megatron Core) for more details. +- **[2024/01 Announcement]** NVIDIA has released the core capabilities in **Megatron-LM** into [**Megatron Core**](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core) in this repository. Megatron Core expands upon Megatron-LM's GPU-optimized techniques with more cutting-edge innovations on system-level optimizations, featuring composable and modular APIs. - - # Project Structure ``` @@ -77,12 +75,12 @@ Megatron-LM/ │ │ ├── distributed/ # Distributed training (FSDP, DDP) │ │ ├── optimizer/ # Optimizers │ │ ├── datasets/ # Dataset loaders -│ │ ├── inference/ # Inference engines +│ │ ├── inference/ # Inference engines and server │ │ └── export/ # Model export (e.g. TensorRT-LLM) │ ├── training/ # Training scripts -│ ├── inference/ # Inference server │ ├── legacy/ # Legacy components -│ └── post_training/ # Post-training (RLHF, etc.) +│ ├── post_training/ # Post-training (quantization, distillation, pruning, etc.) +│ └── rl/ # Reinforcement learning (RLHF, etc.) ├── examples/ # Ready-to-use training examples ├── tools/ # Utility tools ├── tests/ # Comprehensive test suite @@ -90,10 +88,9 @@ Megatron-LM/ ``` - # Performance Benchmarking -For our latest performance benchmarking results, please refer to [NVIDIA NeMo Framework Performance Summary](https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html). +For our latest performance benchmarking results, please refer to [NVIDIA Megatron Bridge Performance Summary](https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html). Our codebase efficiently trains models from 2B to 462B parameters across thousands of GPUs, achieving up to **47% Model FLOP Utilization (MFU)** on H100 clusters. @@ -127,8 +124,9 @@ We also strong scaled the standard GPT-3 model (our version has slightly more th ![Strong scaling](images/strong_scaling.png) +# Roadmaps - +- **[MoE Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - DeepSeek-V3, Qwen3, advanced parallelism, FP8 optimizations, and Blackwell enhancements # Resources diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev index bb9ca5fbe9a..fa214deeea5 100644 --- a/docker/Dockerfile.ci.dev +++ b/docker/Dockerfile.ci.dev @@ -88,7 +88,7 @@ RUN --mount=type=secret,id=JET_INDEX_URLS bash -ex <<"EOF" JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) python -m venv /opt/jet /opt/jet/bin/pip install --no-cache-dir $JET_INDEX_URLS \ - jet-api==$JET_API_VERSION + "jet-api==$JET_API_VERSION" "setuptools<82.0.0" EOF RUN --mount=type=secret,id=JET_INDEX_URLS \ diff --git a/docs/api-guide/core/dist_checkpointing.md b/docs/api-guide/core/dist_checkpointing.md index 959aa4b07e0..38aca198b1f 100644 --- a/docs/api-guide/core/dist_checkpointing.md +++ b/docs/api-guide/core/dist_checkpointing.md @@ -32,19 +32,62 @@ import torch, argparse torch.serialization.add_safe_globals([argparse.Namespace]) ``` -Checkpointing Distributed Optimizer ------------------------------------ +## Checkpointing Distributed Optimizer -Checkpoint Compatibility and Optimizer State Formats -#################################################### +### Checkpoint Compatibility and Optimizer State Formats Beginning with **mcore v0.14**, the ``flattened_range`` attribute was removed from ``dist_checkpointing``. As a result: -- Optimizer states saved with mcore versions < 0.14 are no longer loadable. Loading these legacy optimizer states is not supported because the required sharded metadata is no longer available. -- Model weights from older checkpoints remain fully compatible. No additional work is required—model weights from checkpoints produced by earlier versions are loaded automatically. +- Optimizer states saved with mcore versions <= 0.14 can no longer be loaded directly. Loading these legacy optimizer states is not supported because the required sharded metadata is no longer available. If you need to continue training from older checkpoints, refer to the workaround described below. +- Model weights from older checkpoints remain fully compatible. No extra steps are needed—model weights from checkpoints created by earlier versions load automatically; simply add the ``--no-load-optim`` flag. -Distributed Optimizer Checkpoint Formats -######################################## +### Workaround: Loading legacy optimizer states with ToT MCore + +**Step 1: Convert the legacy checkpoint using mcore v0.15.0** + +Run a dummy training job with mcore v0.15.0 to re-save the checkpoint with new optimizer states format. + +```bash +MODEL_TRAIN_PARAMS=( + # Define model architecture and training parameters here +) +OLD_CKPT=/workspace/mcore_ckpt_old +CONVERTED_CKPT=/workspace/mcore_ckpt_0.15.0 + +torchrun --nproc_per_node=8 /opt/megatron-lm/pretrain_gpt.py \ + --save-interval 1 \ + --eval-interval 1 \ + --exit-interval 1 \ + --eval-iters 1 \ + --use-distributed-optimizer \ + --save ${CONVERTED_CKPT} \ + --load ${OLD_CKPT} \ + --ckpt-format torch_dist \ + "${MODEL_TRAIN_PARAMS[@]}" +``` + +**Step 2: Load the converted checkpoint with ToT MCore** + +Use the converted checkpoint as the input for continued training with ToT MCore. + +```bash +MODEL_TRAIN_PARAMS=( + # Define model architecture and training parameters here +) +NEW_CKPT=/workspace/mcore_ckpt_new +CONVERTED_CKPT=/workspace/mcore_ckpt_0.15.0 + +torchrun --nproc_per_node=8 /opt/megatron-lm/pretrain_gpt.py \ + --use-distributed-optimizer \ + --save ${NEW_CKPT} \ + --load ${CONVERTED_CKPT} \ + --ckpt-format torch_dist \ + "${MODEL_TRAIN_PARAMS[@]}" +``` + +After this step, training can proceed normally using ToT MCore with fully supported optimizer state loading. + +## Distributed Optimizer Checkpoint Formats The refactor of the Distributed Optimizer introduces **two checkpoint formats**: @@ -57,8 +100,7 @@ The refactor of the Distributed Optimizer introduces **two checkpoint formats**: - Slower than dp_reshardable. - Enabled via the ``--dist-ckpt-optim-fully-reshardable`` flag. -Workflow for Changing Model Parallelism -####################################### +### Workflow for Changing Model Parallelism You can combine formats to optimize both flexibility and performance: diff --git a/docs/index.md b/docs/index.md index 448a75e4c93..c9df5264b17 100644 --- a/docs/index.md +++ b/docs/index.md @@ -6,13 +6,13 @@ Megatron Core offers a flexible, reusable foundation for building large-scale tr ## Key Features -* Composable transformer building blocks (attention, MLP, etc.) +* Composable transformer building blocks (attention, MLP) * Advanced parallelism strategies (TP, PP, DP, EP, CP) * Pipeline schedules and distributed optimizers * Mixed precision support (FP16, BF16, FP8) * GPU-optimized kernels and memory management * High-performance dataloaders and dataset utilities -* Model architectures (LLaMA, Qwen, DeepSeek, GPT, Mamba, etc.) +* Model architectures (LLaMA, Qwen, DeepSeek, GPT, Mamba) ```{toctree} @@ -81,16 +81,16 @@ developer/generate_docs ```{toctree} :maxdepth: 2 :hidden: -:caption: Discussions +:caption: API Reference -advanced/index +api-guide/index +apidocs/index.rst ``` ```{toctree} :maxdepth: 2 :hidden: -:caption: API Reference +:caption: Resources -api-guide/index -apidocs/index.rst +advanced/index ``` \ No newline at end of file diff --git a/docs/user-guide/features/tokenizers.md b/docs/user-guide/features/tokenizers.md index 0aecf8df8a7..77b165f5269 100644 --- a/docs/user-guide/features/tokenizers.md +++ b/docs/user-guide/features/tokenizers.md @@ -141,7 +141,7 @@ Use a null tokenizer for testing or non-text models: ```python tokenizer = MegatronTokenizer.from_pretrained( - metadata_path={"library": "null"}, + metadata_path={"library": "null-text"}, vocab_size=131072, ) ``` @@ -173,16 +173,6 @@ torchrun --nproc_per_node=8 pretrain_gpt.py \ If `--tokenizer-metadata` is not specified, a default metadata file is generated automatically based on the tokenizer type. -### Legacy Tokenizer Support - -The old tokenizer system is still supported for backward compatibility: - -```bash -torchrun --nproc_per_node=8 pretrain_gpt.py \ - --legacy-tokenizer \ - ... -``` - ## Supported Tokenizer Libraries | Library | Description | Use Case | diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 88b744b3ac0..2582f57e13e 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -1,40 +1,31 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# pylint: disable=bad-builtin + import hashlib import io import json -import math import os -import pickle import sys import warnings -import torch -from argparse import ArgumentParser from collections import defaultdict -from functools import partial +from typing import Dict, List, Optional + +import torch from tqdm import tqdm -from typing import Dict, List, Tuple, Optional sys.path.append( os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)) ) -import megatron from examples.inference.gpt.utils import ( Request, - add_common_inference_args, build_dynamic_engine_setup_prefix, build_requests, get_curr_time, get_global_peak_memory_stats_bytes, ) -from megatron.core.inference.contexts.dynamic_context import ( - ContextOverflowError, - DynamicInferenceContext, -) -from megatron.core.inference.contexts.attention_context.mamba_metadata import ( - MambaInferenceStateConfig, -) +from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext from megatron.core.inference.engines import DynamicInferenceEngine, EngineSuspendedError from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import ( GPTInferenceWrapper, @@ -43,195 +34,27 @@ from megatron.core.inference.text_generation_controllers.text_generation_controller import ( TextGenerationController, ) -from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer -from megatron.core.transformer.module import MegatronModule -from megatron.core.utils import get_mamba_inference_state_config_from_model +from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer +from megatron.inference.utils import ( + add_inference_args, + get_inference_config_from_model_and_args, + get_model_for_inference, +) sys.path.append( os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)) ) -from megatron.training import get_args, get_model as _get_model, get_tokenizer, initialize_megatron -from megatron.training.checkpointing import load_checkpoint -from model_provider import model_provider -from gpt_builders import gpt_builder -from mamba_builders import mamba_builder +import logging +import megatron from megatron.core.utils import configure_nvtx_profiling -import logging +from megatron.training import get_args, get_tokenizer, initialize_megatron torch.serialization.add_safe_globals([io.BytesIO]) torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState]) torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic]) -def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser: - """Dynamic inference arguments.""" - - add_common_inference_args(parser) - - group = parser.add_argument_group(title='Dynamic inference') - group.add_argument( - "--inference-ckpt-non-strict", - action="store_true", - help="Load checkpoint with `strict=False`.", - ) - group.add_argument( - "--termination-id", type=int, default=None, - help="Termination ID that overrides `tokenizer.eod`.", - ) - group.add_argument( - "--suspend-resume-interval", type=int, default=None, - help="Suspend and resume the dynamic engine every " - "`suspend_resume_interval` steps. This is used to tet the suspend/resume " - "system.", - ) - group.add_argument( - "--inference-repeat-n", type=int, default=1, - help="Repeat inference iterations N times for benchmarking." - ) - group.add_argument( - "--throughput-check-only", - action='store_true', - default=False, - help="If true, only run throughput check without verifying outputs." - ) - - return parser - - -def get_model() -> MegatronModule: - """Initialize model and load checkpoint.""" - - args = get_args() - - if args.model_provider == "gpt": - model_builder = gpt_builder - elif args.model_provider == "mamba": - model_builder = mamba_builder - else: - raise ValueError(f"Invalid model provider {args.model_provider}") - - # Build model. - model = _get_model( - partial(model_provider, model_builder), - wrap_with_ddp=False - ) - - # Load checkpoint. - assert args.load is not None - args.exit_on_missing_checkpoint = True - load_checkpoint( - ddp_model=model, - optimizer=None, - opt_param_scheduler=None, - strict=not args.inference_ckpt_non_strict, - ) - - # No virtual PP. - assert len(model) == 1, "Above condition should have caught this" - model = model[0] - - # Eval mode. - model.eval() - - return model - - -def get_inference_context( - requests: List[Request], - sampling_params: Optional[SamplingParams] = None, - calculate_max_sequence_length_from_requests: bool = True, - mamba_inference_state_config: Optional[MambaInferenceStateConfig] = None, -): - """The inference context manages the KV cache and other inference state.""" - - args = get_args() - - # Max sequence length. - if calculate_max_sequence_length_from_requests: - max_gen_length = sampling_params.num_tokens_to_generate - max_context_length = max(len(r.prompt_tokens) for r in requests) - max_sequence_length = max_context_length + max_gen_length - else: - max_sequence_length = args.inference_max_seq_length - - metrics_writer = None - if args.inference_logging_step_interval > 0 and args.inference_wandb_logging: - metrics_writer = get_wandb_writer() - - # Inference context. - context = DynamicInferenceContext( - params_dtype=args.params_dtype, - num_layers=args.num_layers // args.pipeline_model_parallel_size, - kv_channels=args.kv_channels, - num_attention_heads=( - args.num_query_groups if args.group_query_attention else args.num_attention_heads - ), - max_sequence_length=max_sequence_length, - num_cuda_graphs=( - args.inference_dynamic_batching_num_cuda_graphs - if args.cuda_graph_impl == "local" - else None - ), - block_size_tokens=args.inference_dynamic_batching_block_size, - buffer_size_gb=args.inference_dynamic_batching_buffer_size_gb, - paused_buffer_size_gb=args.inference_dynamic_batching_paused_buffer_size_gb, - max_requests=args.inference_dynamic_batching_max_requests, - max_tokens=args.inference_dynamic_batching_max_tokens, - tensor_model_parallel_size=args.tensor_model_parallel_size, - pipeline_model_parallel_size=args.pipeline_model_parallel_size, - materialize_only_last_token_logits=not args.return_log_probs, - mamba_inference_state_config=mamba_inference_state_config, - cache_mla_latent=args.multi_latent_attention and args.cache_mla_latents, - kv_lora_rank=args.kv_lora_rank if args.multi_latent_attention else None, - qk_pos_emb_head_dim=args.qk_pos_emb_head_dim, - use_cuda_graphs_for_non_decode_steps=not args.decode_only_cuda_graphs, - use_flashinfer_fused_rope=args.use_flashinfer_fused_rope, - unified_memory_level=args.inference_dynamic_batching_unified_memory_level, - cuda_graph_max_tokens=args.inference_dynamic_batching_cuda_graph_max_tokens, - cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count, - metrics_writer=metrics_writer, - offload_kv_cache=args.rl_offload_kv_cache_during_training - ) - - return context - - -def get_inference_controller( - model: MegatronModule, context: DynamicInferenceContext -) -> TextGenerationController: - """Buid text generation controller, which manages the model inference context. - - Args: - model (MegatronModule): Megatron GPT model. - context (DynamicInferenceContext): Context for managing KV cache blocks. - - Return: - (TextGenerationController) Inference text generation controller. - """ - - args = get_args() - if args.legacy_tokenizer: - tokenizer = get_tokenizer() - else: - tokenizer = build_tokenizer(args) - - # Wrap model in inference wrapper. - model = GPTInferenceWrapper(model, args, context) - - # Note: the following is taken from AbstractModelInferenceWrapper.prep_model_for_inference(). - from megatron.core import parallel_state - - model.model_is_pipeline_parallel = not ( - parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage() - ) - - # Text generation controller. - controller = TextGenerationController(model, tokenizer) - - return controller - - def run_inference( requests: List[Request], engine: DynamicInferenceEngine, @@ -284,11 +107,7 @@ def _add_request(): """ nonlocal num_requests_added _request = requests[num_requests_added] - engine.add_request( - num_requests_added, - _request.prompt_text, - _request.sampling_params, - ) + engine.add_request(num_requests_added, _request.prompt_text, _request.sampling_params) _request.time_start = get_curr_time() _request.state = "started" num_requests_added += 1 @@ -305,10 +124,9 @@ def _add_request(): _add_request() else: # Add deterministic number of requests (generally used for debugging). - for i in range(min( - args.incoming_requests_per_step, - num_requests_total - num_requests_added, - )): + for i in range( + min(args.incoming_requests_per_step, num_requests_total - num_requests_added) + ): _add_request() add_times.append(get_curr_time() - add_start) @@ -318,11 +136,12 @@ def _add_request(): result = engine.step_modern() except EngineSuspendedError as e: result = e - pass # ignore error in order to call 'engine.resume()' below. + pass # ignore error in order to call 'engine.resume()' below. attempted_step_count += 1 - # After step, we lost track of last iteration's is_decode_only, so we need to get it from the engine - is_decode_only = engine.is_decode_only + # After step, we lost track of last iteration's is_decode_only, + # so we need to get it from the engine + is_decode_only = engine.is_decode_only # Test suspending and resuming engine. if args.suspend_resume_interval is not None: @@ -335,9 +154,9 @@ def _add_request(): # Resume, 0+ attempted steps later. if ( attempted_step_count > 0 - and - (attempted_step_count - args.suspend_resume_interval // 2) - % args.suspend_resume_interval == 0 + and (attempted_step_count - args.suspend_resume_interval // 2) + % args.suspend_resume_interval + == 0 ): print("**** step %d/%d ... resume." % (engine.step_count, attempted_step_count)) engine.resume() @@ -349,7 +168,9 @@ def _add_request(): # Record cuda_graph_request_count. cuda_graph_request_count = result["cuda_graph_request_count"] if args.cuda_graph_impl == "local" and cuda_graph_request_count is not None: - cuda_graph_request_count_map[cuda_graph_request_count] = cuda_graph_request_count_map.get(cuda_graph_request_count, 0) + 1 + cuda_graph_request_count_map[cuda_graph_request_count] = ( + cuda_graph_request_count_map.get(cuda_graph_request_count, 0) + 1 + ) # Update requests. active_request_ids = result["active_request_ids"] @@ -374,6 +195,8 @@ def _add_request(): request.request_id = finished_request.request_id request.events = finished_request.events + request.ttft = finished_request.ttft + # Update prompt, in case engine has been suspended and resumed. request.prompt_tokens = finished_request.prompt_tokens.tolist() request.prompt_text = finished_request.prompt @@ -408,38 +231,37 @@ def _add_request(): engine.resume() return { - "step_times" : step_times, - "add_times" : add_times, - "output_times" : output_times, - "total_output_tokens" : total_output_tokens, - "cuda_graph_request_count_map" : cuda_graph_request_count_map, + "step_times": step_times, + "add_times": add_times, + "output_times": output_times, + "total_output_tokens": total_output_tokens, + "cuda_graph_request_count_map": cuda_graph_request_count_map, } @torch.inference_mode() def main(): - + """Run dynamic inference.""" # Initialize Megatron. initialize_megatron( - extra_args_provider=add_dynamic_inference_args, + extra_args_provider=add_inference_args, args_defaults={'no_load_rng': True, 'no_load_optim': True}, ) # Start Nsight profiler. if os.environ.get("NSIGHT_PREFIX"): torch.cuda.cudart().cudaProfilerStart() - - level_str = os.getenv("LOG_LEVEL", "INFO").upper() - level = getattr(logging, level_str, logging.INFO) + + level_str = os.getenv("LOG_LEVEL", "INFO").upper() + level = getattr(logging, level_str, logging.INFO) logging.basicConfig(level=level, force=True) configure_nvtx_profiling(True) args = get_args() - if args.legacy_tokenizer: - tokenizer = get_tokenizer() - else: - tokenizer = build_tokenizer(args) + + # Build tokenizer + tokenizer = build_tokenizer(args) # Reset peak memory stats so functional tests measure this run and not # whatever happened earlier during initialization. @@ -456,42 +278,36 @@ def main(): termination_id=args.termination_id if args.termination_id is not None else tokenizer.eod, top_n_logprobs=args.top_n_logprobs, stop_words=args.stop_words, - ) - - model = get_model() + ) - mamba_inference_state_config = get_mamba_inference_state_config_from_model(model) + model = get_model_for_inference() # Requests, context, controller. requests = build_requests(args, tokenizer, sampling_params) - context = get_inference_context( - requests, - sampling_params, - mamba_inference_state_config=mamba_inference_state_config, - ) - controller = get_inference_controller(model, context) + inference_config = get_inference_config_from_model_and_args(model, args) + + # Calculate max_sequence_length from requests + max_gen_length = sampling_params.num_tokens_to_generate + max_context_length = max(len(r.prompt_tokens) for r in requests) + inference_config.max_sequence_length = max_context_length + max_gen_length + context = DynamicInferenceContext(model.config, inference_config) + wrapped_model = GPTInferenceWrapper(model, context) + controller = TextGenerationController(wrapped_model, tokenizer) # Validate all context_length's <= max_tokens. - if args.disable_chunked_prefill: + if not args.enable_chunked_prefill: invalid_prompt_length_map = {} for request_idx, request in enumerate(requests): if len(request.prompt_tokens) > context.max_tokens: invalid_prompt_length_map[request_idx] = len(request.prompt_tokens) - assert not invalid_prompt_length_map, ( - "request idxs with prompts longer than context.max_tokens: " - ", ".join(f"{k}({v})" for k, v in invalid_prompt_length_map.items()) + assert ( + not invalid_prompt_length_map + ), "request idxs with prompts longer than context.max_tokens: " ", ".join( + f"{k}({v})" for k, v in invalid_prompt_length_map.items() ) # Inference engine. - engine = DynamicInferenceEngine( - controller, - context, - enable_cuda_graph=args.cuda_graph_impl == "local", - random_seed=args.seed, - track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events, - enable_chunked_prefill=not args.disable_chunked_prefill, - inference_logging_step_interval=args.inference_logging_step_interval, - ) + engine = DynamicInferenceEngine(controller, context) setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests) print("~~~") @@ -522,14 +338,13 @@ def main(): # Validate all requests finished. for request in requests: - assert request.state == "finished", ( - f"request.state == '{request.state}' != 'finished'." - ) + assert request.state == "finished", f"request.state == '{request.state}' != 'finished'." peak_mem_stats = get_global_peak_memory_stats_bytes() # Print unique prompts + outputs. if torch.distributed.get_rank() == 0: + def escape_str(s): return s.replace("\n", "\\n") @@ -547,7 +362,10 @@ def escape_str(s): # ---- Prompt summary line ---- prompt_len = len(requests[request_idxs[0]].prompt_tokens) escaped_prompt_text = escape_str(prompt_text) - print(f"\n{unique_idx+1}/{len(unique_prompt_map)} [n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}") + print( + f"\n{unique_idx+1}/{len(unique_prompt_map)}" + f"[n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}" + ) # ---- Group all outputs for this prompt ---- output_map = defaultdict(list) @@ -567,16 +385,17 @@ def escape_str(s): # Use hash of prompt + generated text in case engine was # suspended and resumed, which misaligns boundary between # prompt and generated tokens. - o_hash = hashlib.sha256( - (prompt_text + output_text).encode() - ).hexdigest()[:6] + o_hash = hashlib.sha256((prompt_text + output_text).encode()).hexdigest()[:6] o_len = len(requests[output_request_idxs[0]].output_tokens) escaped_output_text = escape_str(output_text) else: o_hash = "--" o_len = 0 escaped_output_text = "--" - print(f" >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}{', ' if evicted else ''}] {escaped_output_text}") + print( + f" >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}" + f"{', ' if evicted else ''}] {escaped_output_text}" + ) text_hashes.append(o_hash) # Write results to JSON. Primarily used for functional testing. @@ -592,14 +411,17 @@ def escape_str(s): "generated_text": req.output_text, "generated_tokens": req.output_tokens, "latency": req.time_end - req.time_start, - "cuda_graph_request_count_map" : result["cuda_graph_request_count_map"], - "step_count" : engine.step_count, - "top_n_logprobs" : getattr(req, 'generated_top_n_logprobs', None), - "prompt_top_n_logprobs" : getattr(req, 'prompt_top_n_logprobs', None), + "ttft": req.ttft, # Time-to-first-token in seconds + "cuda_graph_request_count_map": result["cuda_graph_request_count_map"], + "step_count": engine.step_count, + "top_n_logprobs": getattr(req, 'generated_top_n_logprobs', None), + "prompt_top_n_logprobs": getattr(req, 'prompt_top_n_logprobs', None), } if req.sampling_params.return_log_probs: result_dict["prompt_logprobs"] = getattr(req, 'prompt_log_probs', None) - result_dict["generated_logprobs"] = getattr(req, 'generated_log_probs', None) + result_dict["generated_logprobs"] = getattr( + req, 'generated_log_probs', None + ) result_dict["logprobs"] = getattr(req, 'logprobs', None) json_results[req.request_id] = result_dict @@ -631,7 +453,7 @@ def escape_str(s): d_count = len(d_times) p_mean = p_total / p_count - d_mean = d_total / d_count if d_count != 0 else 0. + d_mean = d_total / d_count if d_count != 0 else 0.0 # Commented out for now as the step/add/output times are not calculated correctly. # print( @@ -643,18 +465,13 @@ def escape_str(s): # f"mean [ p {p_mean:.3f}s, d {d_mean:.3f}s ], " # f"count [ p {p_count}, d {d_count} ]." # ) - capture_str = ( - f"{engine.capture_stats['time']:.2f} sec" - if engine.capture_stats else - "--" - ) + capture_str = f"{engine.capture_stats['time']:.2f} sec" if engine.capture_stats else "--" print( - f"{setup_prefix} … " - f"throughput: {throughput:.3f} tok/s … ", + f"{setup_prefix} … " f"throughput: {throughput:.3f} tok/s … ", f"total time: {total_time:.3f}s … " f"mem {peak_alloc_gb:.1f}/{peak_resvd_gb:.1f} GB … " f"steps: {engine.step_count:d} … " - f"capture {capture_str}" + f"capture {capture_str}", ) print("~~~") diff --git a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py index cbb7a1aa745..536f533eccd 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py +++ b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py @@ -2,43 +2,33 @@ import asyncio import json +import logging import os import time -import torch -import torch.distributed as dist +import warnings from collections import defaultdict -from tqdm import tqdm from typing import List -import warnings -import logging -from examples.inference.gpt.gpt_dynamic_inference import ( - add_dynamic_inference_args, - get_inference_context, - get_inference_controller, - get_model, -) -from examples.inference.gpt.utils import ( - Request, - build_dynamic_engine_setup_prefix, - build_requests, - add_common_inference_args -) +import torch +import torch.distributed as dist -from megatron.core import parallel_state +from examples.inference.gpt.utils import Request, build_dynamic_engine_setup_prefix, build_requests from megatron.core.inference.engines import DynamicInferenceEngine from megatron.core.inference.inference_client import InferenceClient from megatron.core.inference.inference_request import DynamicInferenceRequestRecord from megatron.core.inference.sampling_params import SamplingParams -from megatron.core.utils import get_mamba_inference_state_config_from_model - +from megatron.inference.utils import ( + add_inference_args, + get_dynamic_inference_engine, + get_model_for_inference, +) from megatron.training import get_args, get_tokenizer, initialize_megatron -from megatron.training.arguments import parse_args # pylint: disable=line-too-long logging.basicConfig(level=logging.INFO, force=True) + async def main( engine: DynamicInferenceEngine, requests: List[Request], @@ -51,12 +41,11 @@ async def main( "Sampling parameters are specified per request.", DeprecationWarning, ) - + # once you call engine.start_listening_to_data_parallel_coordinator, # the engine will start accepting requests from the data parallel coordinator. # and processing them in an asyncio coroutine. # leaving inference_coordinator_port as None will find a free port automatically. - dp_addr = await engine.start_listening_to_data_parallel_coordinator( inference_coordinator_port=port, launch_inference_coordinator=True, @@ -69,14 +58,11 @@ async def main( # Since the client doesn't directly call engine.async_step here, we test # the suspend-resume system ~4 times. suspend_resume_interval = max(1, len(requests) // 4) - suspend_idxs = set(range( - suspend_resume_interval, - len(requests) + 1, - suspend_resume_interval, - )) + suspend_idxs = set( + range(suspend_resume_interval, len(requests) + 1, suspend_resume_interval) + ) resume_idxs = set( - min(len(requests), i + suspend_resume_interval // 2) - for i in suspend_idxs + min(len(requests), i + suspend_resume_interval // 2) for i in suspend_idxs ) else: suspend_idxs = set() @@ -98,7 +84,10 @@ async def main( current_time = time.time_ns() / 10**9 if args.incoming_requests_per_step is None: # Only add requests that have arrived at the current time. - while num_requests_added < num_requests_total and requests[num_requests_added].time_arrival <= current_time: + while ( + num_requests_added < num_requests_total + and requests[num_requests_added].time_arrival <= current_time + ): request = requests[num_requests_added] # These add-request calls will queue up the request on a zmq socket and return # instantaneously. They will return an asyncio future which can be awaited for @@ -114,10 +103,9 @@ async def main( else: # Add deterministic number of requests (generally used for debugging). - for i in range(min( - args.incoming_requests_per_step, - num_requests_total - num_requests_added - )): + for i in range( + min(args.incoming_requests_per_step, num_requests_total - num_requests_added) + ): # Change sampling parameters to force different generation lengths. request = requests[num_requests_added] n = request.sampling_params.num_tokens_to_generate @@ -135,7 +123,7 @@ async def main( break # Relinquish control since there are no more requests to add at the moment. This allows the engine to run. await asyncio.sleep(0) - + # While we wait for the requests to complete, the engine runs in the background. results: List[DynamicInferenceRequestRecord] = await asyncio.gather(*futures) @@ -157,6 +145,9 @@ async def main( result_dict["logprobs"] = req.prompt_log_probs + req.generated_log_probs throughput = len(req.generated_tokens) / req.latency throughputs.append(throughput) + if req.routing_indices is not None: + result_dict["routing_indices"] = req.routing_indices.tolist() + json_results[req.request_id] = result_dict throughput_dict = {"throughput": throughputs} if args.throughput_check_only: @@ -170,16 +161,19 @@ async def main( req = record.merge() unique_prompt_map[req.prompt].append(req) for idx, (prompt_text, reqs) in enumerate(unique_prompt_map.items()): - print(f"%d/%d. prompt '%s' ... [%d] output '%s'." % ( - idx, - len(unique_prompt_map), - prompt_text.replace("\n", "\\n"), - len(reqs), - reqs[0].generated_text.replace("\n", "\\n"), - )) + print( + f"%d/%d. prompt '%s' ... [%d] output '%s'." + % ( + idx, + len(unique_prompt_map), + prompt_text.replace("\n", "\\n"), + len(reqs), + reqs[0].generated_text.replace("\n", "\\n"), + ) + ) # kill the engines and suspend the client - # Right now, we can only call stop when all requests are done. + # Right now, we can only call stop when all requests are done. # Todo: Make this explicit in the Client class.... await client.stop_engines() client.stop() @@ -190,11 +184,11 @@ async def main( if __name__ == "__main__": - # enable inference mode in the very beginning as some fp-8 optimizations + # enable inference mode in the very beginning as some fp8 optimizations # check for it. with torch.inference_mode(): initialize_megatron( - extra_args_provider=add_dynamic_inference_args, + extra_args_provider=add_inference_args, args_defaults={'no_load_rng': True, 'no_load_optim': True}, ) @@ -213,34 +207,16 @@ async def main( ), ) - # Requests, context, conroller. - model = get_model() - mamba_inference_state_config = get_mamba_inference_state_config_from_model(model) + model = get_model_for_inference() + requests = ( build_requests(args, tokenizer, sampling_params) if dist.get_rank() == 0 else None ) - context = get_inference_context( - None, - None, - calculate_max_sequence_length_from_requests=False, - mamba_inference_state_config=mamba_inference_state_config, - ) - - controller = get_inference_controller(model, context) - - # Inference engine. - engine = DynamicInferenceEngine( - controller, - context, - enable_cuda_graph=args.cuda_graph_impl == "local", - random_seed=args.seed, - enable_chunked_prefill=not args.disable_chunked_prefill, - inference_logging_step_interval=args.inference_logging_step_interval, - ) + engine = get_dynamic_inference_engine(model=model) if dist.get_rank() == 0: - setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests) + setup_prefix = build_dynamic_engine_setup_prefix(args, model, engine.context, requests) print("~~~") print(setup_prefix) print("~~~") @@ -249,13 +225,7 @@ async def main( if os.environ.get("NSIGHT_PREFIX"): torch.cuda.cudart().cudaProfilerStart() - asyncio.run( - main( - engine, - requests, - args.inference_coordinator_port, - ) - ) + asyncio.run(main(engine, requests, args.inference_coordinator_port)) # Stop Nsight profiler. if os.environ.get("NSIGHT_PREFIX"): diff --git a/examples/inference/gpt/gpt_static_inference.py b/examples/inference/gpt/gpt_static_inference.py index 03a60927ab2..17cf7c53b05 100644 --- a/examples/inference/gpt/gpt_static_inference.py +++ b/examples/inference/gpt/gpt_static_inference.py @@ -1,21 +1,11 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import os -from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import ( - InferenceWrapperConfig, -) -from model_provider import model_provider -from gpt_builders import gpt_builder -from mamba_builders import mamba_builder -import torch import sys import time -import warnings -from functools import partial from argparse import Namespace import torch -import tqdm from megatron.core.inference.contexts import StaticInferenceContext from megatron.core.inference.engines import StaticInferenceEngine @@ -23,17 +13,12 @@ from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import ( GPTInferenceWrapper, ) -from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import ( - InferenceWrapperConfig, -) from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.text_generation_controllers.text_generation_controller import ( TextGenerationController, ) -from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer +from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer from megatron.core.transformer.module import MegatronModule -from pretrain_gpt import model_provider as gpt_model_provider -from pretrain_mamba import model_provider as mamba_model_provider sys.path.append( os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)) @@ -41,18 +26,18 @@ import asyncio import json -from typing import Any, AsyncIterator, List +from typing import List -from examples.inference.gpt.utils import add_common_inference_args, build_requests -from megatron.core import mpu -from megatron.training import get_args, get_model, get_tokenizer, print_rank_0 -from megatron.training.checkpointing import load_checkpoint +from examples.inference.gpt.utils import build_requests +from megatron.inference.utils import add_inference_args, get_model_for_inference +from megatron.training import get_args, get_tokenizer, print_rank_0 from megatron.training.initialize import initialize_megatron + def add_static_inference_args(parser): """Static inference arguments.""" - add_common_inference_args(parser) + add_inference_args(parser) group = parser.add_argument_group(title='Static inference') group.add_argument( @@ -79,34 +64,17 @@ def get_inference_engine(args: Namespace, model: MegatronModule) -> StaticInfere Returns: AbstractBackend: The chosen backend """ - if args.legacy_tokenizer: - tokenizer = get_tokenizer() - else: - tokenizer = build_tokenizer(args) - inference_wrapper_config = InferenceWrapperConfig( - hidden_size=args.hidden_size, - inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold, - fp32_residual_connection=args.fp32_residual_connection, - params_dtype=args.params_dtype, - padded_vocab_size=args.padded_vocab_size, - inference_max_requests=args.inference_max_batch_size, - inference_max_seq_length=args.inference_max_seq_length, - nccl_all_reduce_for_prefill=args.nccl_all_reduce_for_prefill, - fp8=args.fp8, - moe_pad_experts_for_cuda_graph_inference = args.moe_pad_experts_for_cuda_graph_inference - ) - - inference_context = StaticInferenceContext.from_config(inference_wrapper_config) - - inference_wrapped_model = GPTInferenceWrapper( - model, inference_wrapper_config, inference_context + tokenizer = build_tokenizer(args) + inference_context = StaticInferenceContext( + args.inference_max_requests, args.inference_max_seq_length ) + inference_wrapped_model = GPTInferenceWrapper(model, inference_context) text_generation_controller = TextGenerationController( inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer ) engine_kwargs = { - "text_generation_controller" : text_generation_controller, - "legacy" : args.use_legacy_static_engine, + "text_generation_controller": text_generation_controller, + "legacy": args.use_legacy_static_engine, } if not args.use_legacy_static_engine: engine_kwargs["buffer_size_gb"] = args.inference_dynamic_batching_buffer_size_gb @@ -165,22 +133,7 @@ def main(): args = get_args() - if args.max_batch_size is not None: - warnings.warn( - f"`--max-batch-size` has been deprecated in favor of `--inference-max-requests`." - ) - args.inference_max_batch_size = max(args.max_batch_size, args.inference_max_batch_size) - - # Set up model and load checkpoint - if args.model_provider == "gpt": - model_builder = gpt_builder - elif args.model_provider == "mamba": - model_builder = mamba_builder - else: - raise ValueError(f"Invalid model provider {args.model_provider}") - model = get_model(partial(model_provider, model_builder), wrap_with_ddp=False) - load_checkpoint(model, None, None, strict=False) - model = model[0] + model = get_model_for_inference() inference_engine = get_inference_engine(args, model) @@ -193,10 +146,9 @@ def main(): top_n_logprobs=args.top_n_logprobs, ) - if args.legacy_tokenizer: - tokenizer = get_tokenizer() - else: - tokenizer = build_tokenizer(args) + # Build tokenizer + tokenizer = build_tokenizer(args) + requests = build_requests(args, tokenizer) prompts = [r.prompt_text for r in requests] @@ -276,7 +228,7 @@ def main(): ) ), len(requests), - args.inference_max_batch_size, + args.inference_max_requests, stats["allocated_bytes.all.peak"] / (1024**3), stats["reserved_bytes.all.peak"] / (1024**3), latency, @@ -293,6 +245,5 @@ def main(): torch.distributed.destroy_process_group() - if __name__ == "__main__": main() diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index a04b856c0a6..eceb5fa16fd 100644 --- a/examples/inference/gpt/utils.py +++ b/examples/inference/gpt/utils.py @@ -1,158 +1,23 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import copy -import json import itertools +import json import random import time -import torch from argparse import ArgumentParser, Namespace -from tqdm import tqdm +from functools import partial from typing import Any, List, Optional -from megatron.core.inference.inference_request import DynamicInferenceRequest +import torch +from tqdm import tqdm + from megatron.core.inference.contexts import DynamicInferenceContext from megatron.core.inference.contexts.dynamic_context import get_mem_size_str -from megatron.core.transformer.module import MegatronModule - +from megatron.core.inference.inference_request import DynamicInferenceRequest from megatron.core.inference.sampling_params import SamplingParams - - -def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser: - """Common inference arguments.""" - - group = parser.add_argument_group(title='Common inference') - - group.add_argument("--temperature", type=float, default=1.0, help='Sampling temperature.') - group.add_argument("--top_k", type=int, default=1, help='Top k sampling.') - group.add_argument("--top_p", type=float, default=0.0, help='Top p sampling.') - group.add_argument( - "--return-log-probs", - action='store_true', - default=False, - help='Return the log probabilities of the final output tokens', - ) - group.add_argument( - "--prompts", - metavar='N', - type=str, - nargs='+', - help='Input prompts with each prompt within quotes and seperated by space', - ) - group.add_argument( - "--num-tokens-to-prompt", - type=int, - nargs="+", - default=[64, 1024], - help='Number of tokens to use for simulated prompts. This should be a ' - 'space-separated pair of integers, and the generated prompt lengths will ' - 'be uniformly sampled within this range.', - ) - group.add_argument( - "--num-tokens-to-generate", - type=int, - default=30, - help='Number of tokens to generate for each prompt', - ) - group.add_argument( - "--num-tokens-from-file", - action='store_true', - default=False, - help='Use per-prompt num_tokens_to_generate from prompt file', - ) - group.add_argument( - "--top-n-logprobs", - type=int, - default=0, - help='Return the top n logprobs for the generated tokens and their corresponding token as a dictionary', - ) - group.add_argument( - "--incoming-requests-per-step", - type=int, default=None, - help="Add a deterministic number of requests per step. This arg is " - "prioritized over `--incoming-requests-per-sec` below (which is non-" - "deterministic). Note that the number of requests added per step is " - "additionally limited by the inference context's `max_requests`, " - "`max_tokens`, and KV buffer size.", - ) - group.add_argument( - "--incoming-requests-per-sec", - type=float, - default=100.0, - help="Simulated number of requests per second. Set to -1 to add all requests together.", - ) - group.add_argument( - "--incoming-requests-duration", - type=float, - default=10.0, - help="Total amount of time to simulate that requests are " - "arriving. Multiply this value with " - "`--incoming-requests-per-sec` to get the approximate " - "total number of requests. Set to -1 to add all requests together.", - ) - group.add_argument( - "--model-provider", - choices=["mamba", "gpt"], - default="gpt", - help="Model provider", - ) - group.add_argument( - "--skip-prompt-log-probs", - action='store_true', - default=False, - help='Skip prompt log probs.', - ) - group.add_argument( - "--stop-words", - metavar='WORD', - type=str, - nargs='+', - default=None, - help='Stop words to terminate generation. Each word should be quoted and ' - 'separated by space. Example: --stop-words "\\n\\n" "END" "###"', - ) - group.add_argument( - "--output-path", - type=str, - default=None, - help="Path to save generations as JSON", - ) - group.add_argument( - "--output-every-n-results", - type=int, - default=1, - help="To minimize the output file size of larger runs, only write the " - "results of every `n` requests.", - ) - group.add_argument( - "--prompt-file", - help='Jsonl file containing input prompts, where each item (i.e., line) ' - 'contains the field \'text\' where the value is the prompt. All other ' - 'fields within each item are ignored, and may be customized for each ' - 'application.', - ) - group.add_argument( - "--prompt-file-num-truncate", - type=int, - help='Number of samples to use from the loaded prompt file (see ' - '`--prompt-file` above). The first `--prompt-file-num-truncate` samples ' - 'will be used, in order.', - ) - group.add_argument( - "--use-flashinfer-fused-rope", - action='store_true', - default=False, - help='Use flashinfer fused rope implementation.', - ) - group.add_argument( - "--no-record-throughput", - action='store_false', - dest="record_throughput", - help="Disable throughput recording in --output-file" - - ) - - return parser +from megatron.core.transformer.module import MegatronModule +from megatron.training import get_args def get_default_sampling_params(termination_id: int = None): @@ -162,9 +27,10 @@ def get_default_sampling_params(termination_id: int = None): top_p=0.0, return_log_probs=False, num_tokens_to_generate=30, - termination_id = termination_id, + termination_id=termination_id, ) + def get_curr_time() -> float: """Get synchronized time across ranks.""" curr_time = torch.cuda.LongTensor([time.time_ns()]) @@ -188,7 +54,13 @@ class Request: tokenizer (Any): Tokenizer for tokenizing the prompt. """ - def __init__(self, prompt_text: str, time_offset: float, tokenizer: Any, sampling_params: SamplingParams = None): + def __init__( + self, + prompt_text: str, + time_offset: float, + tokenizer: Any, + sampling_params: SamplingParams = None, + ): self.prompt_text = prompt_text self.prompt_tokens = tokenizer.tokenize(prompt_text) self.output_text = None @@ -197,8 +69,13 @@ def __init__(self, prompt_text: str, time_offset: float, tokenizer: Any, samplin self.time_arrival = None self.time_start = None self.time_end = None + self.ttft = None # Time-to-first-token in seconds self.state = "not-started" - self.sampling_params: SamplingParams = sampling_params if sampling_params is not None else get_default_sampling_params(tokenizer.eod) + self.sampling_params: SamplingParams = ( + sampling_params + if sampling_params is not None + else get_default_sampling_params(tokenizer.eod) + ) self.sampling_params = copy.deepcopy(self.sampling_params) def __str__(self) -> str: @@ -225,10 +102,10 @@ def get_time_offsets( # if num_requests is not None: incoming_requests_duration = num_requests / incoming_requests_per_sec - incoming_requests_duration *= 2 # extra margin, to accomodate time sampling + incoming_requests_duration *= 2 # extra margin, to accomodate time sampling random.seed(seed) - + import simpy # Guard against this import in test case # Generate random time offsets. @@ -241,14 +118,14 @@ def arrival(r): env = simpy.Environment() env.process(arrival(incoming_requests_per_sec)) env.run(incoming_requests_duration) - + # Ensure at least a single request. if len(time_offsets) == 0: time_offsets = [0.0] # Ensure first time is 0. time_offsets = [to - time_offsets[0] for to in time_offsets] - + # Truncate to num_requests. assert len(time_offsets) >= num_requests time_offsets = time_offsets[:num_requests] @@ -257,7 +134,7 @@ def arrival(r): def get_cli_requests( - args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None + args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None ) -> list[Request]: # Get time offsets. @@ -269,7 +146,7 @@ def get_cli_requests( ) # Init requests. - requests = [Request(p, t, tokenizer, sampling_params) for p,t in zip(args.prompts, t_offsets)] + requests = [Request(p, t, tokenizer, sampling_params) for p, t in zip(args.prompts, t_offsets)] return requests @@ -289,18 +166,14 @@ def get_synthetic_requests( # Build prompts with expected lengths. assert ( len(args.num_tokens_to_prompt) == 2 - and - args.num_tokens_to_prompt[1] >= args.num_tokens_to_prompt[0] + and args.num_tokens_to_prompt[1] >= args.num_tokens_to_prompt[0] ) max_prompt_length = args.num_tokens_to_prompt[1] max_prompt_text = "hi " * max_prompt_length max_prompt_tokens = tokenizer.tokenize(max_prompt_text) - prompt_lengths = [ - random.randint(*args.num_tokens_to_prompt) - for _ in time_offsets - ] - prompt_tokens_list = [ max_prompt_tokens[:l] for l in prompt_lengths ] - prompt_texts = [ tokenizer.detokenize(tt) for tt in prompt_tokens_list ] + prompt_lengths = [random.randint(*args.num_tokens_to_prompt) for _ in time_offsets] + prompt_tokens_list = [max_prompt_tokens[:l] for l in prompt_lengths] + prompt_texts = [tokenizer.detokenize(tt) for tt in prompt_tokens_list] # Init requests. assert len(prompt_texts) == len(time_offsets) @@ -340,16 +213,15 @@ def get_requests_from_file( # Get time offsets. time_offsets: list[float] = get_time_offsets( - args.seed, - args.incoming_requests_per_step, - args.incoming_requests_per_sec, - len(prompts), + args.seed, args.incoming_requests_per_step, args.incoming_requests_per_sec, len(prompts) ) # Init requests. requests = [ Request(p, t, tokenizer, sp) - for p, t, sp in tqdm(zip(prompts, time_offsets, sampling_params_list), "init requests", total=len(prompts)) + for p, t, sp in tqdm( + zip(prompts, time_offsets, sampling_params_list), "init requests", total=len(prompts) + ) ] return requests @@ -411,19 +283,21 @@ def build_dynamic_engine_setup_prefix( # Prompt description prompt_src_str = ( - "cli" if args.prompts else - "file" if args.prompt_file else - f"synth({', '.join(map(str, args.num_tokens_to_prompt))})" + "cli" + if args.prompts + else ( + "file" + if args.prompt_file + else f"synth({', '.join(map(str, args.num_tokens_to_prompt))})" + ) ) request_str = ( - f"requests: {prompt_src_str}, " - f"n {len(requests):d}, g {args.num_tokens_to_generate:d}, " + f"requests: {prompt_src_str}, " f"n {len(requests):d}, g {args.num_tokens_to_generate:d}, " ) request_str += ( - f"dur {args.incoming_requests_duration:.1e} " - f"r/sec {args.incoming_requests_per_sec:.1e}" - if args.incoming_requests_per_step is None else - f"r/step {args.incoming_requests_per_step}" + f"dur {args.incoming_requests_duration:.1e} " f"r/sec {args.incoming_requests_per_sec:.1e}" + if args.incoming_requests_per_step is None + else f"r/step {args.incoming_requests_per_step}" ) # Buffer limits config @@ -433,14 +307,7 @@ def build_dynamic_engine_setup_prefix( f"[r {context.max_requests}, t {context.max_tokens}]" ) - parts = [ - get_model_size_str(model), - "dynamic", - cg_str, - uvm_str, - request_str, - buffer_limits_str, - ] + parts = [get_model_size_str(model), "dynamic", cg_str, uvm_str, request_str, buffer_limits_str] return " | ".join(parts) @@ -456,4 +323,4 @@ def get_global_peak_memory_stats_bytes() -> dict: t = torch.tensor([peak_alloc], device="cuda", dtype=torch.int64) torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.MAX) peak_alloc = int(t[0].item()) - return {"mem-max-allocated-bytes": peak_alloc} \ No newline at end of file + return {"mem-max-allocated-bytes": peak_alloc} diff --git a/examples/inference/t5/simple_t5_batch_inference.py b/examples/inference/t5/simple_t5_batch_inference.py index 4b15952e07f..1aca74b3176 100644 --- a/examples/inference/t5/simple_t5_batch_inference.py +++ b/examples/inference/t5/simple_t5_batch_inference.py @@ -1,3 +1,5 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + import os import sys from argparse import Namespace @@ -17,7 +19,7 @@ from megatron.core.inference.text_generation_controllers.encoder_decoder_text_generation_controller import ( EncoderDecoderTextGenerationController, ) -from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer +from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer from megatron.core.transformer.module import MegatronModule from pretrain_t5 import model_provider @@ -57,7 +59,7 @@ def add_text_generate_args(parser): metavar='N', type=str, nargs='+', - help='Encoder input prompts with each prompt within quotes and seperated by space', + help='Encoder input prompts with each prompt within quotes and separated by space', ) group.add_argument( "--max-batch-size", type=int, default=1, help='Max number of prompts to process at once' @@ -77,10 +79,8 @@ def get_inference_engine(args: Namespace, model: MegatronModule) -> AbstractEngi Returns: AbstractBackend: The chosen backend """ - if args.legacy_tokenizer: - tokenizer = get_tokenizer() - else: - tokenizer = build_tokenizer(args) + # Build tokenizer + tokenizer = build_tokenizer(args) inference_wrapper_config = InferenceWrapperConfig( hidden_size=args.hidden_size, @@ -131,10 +131,9 @@ def main(): num_tokens_to_generate=args.num_tokens_to_generate, ) - if args.legacy_tokenizer: - tokenizer = get_tokenizer() - else: - tokenizer = build_tokenizer(args) + # Build tokenizer + tokenizer = build_tokenizer(args) + decoder_prompts = [""] * len( args.encoder_prompts ) # for T5, the prompt is provided as encoder input, hence decoder_prompts is empty diff --git a/examples/llama/train_llama3_8b_h100_fp8.sh b/examples/llama/train_llama3_8b_h100_fp8.sh index f791996308e..28227546bc7 100644 --- a/examples/llama/train_llama3_8b_h100_fp8.sh +++ b/examples/llama/train_llama3_8b_h100_fp8.sh @@ -69,6 +69,7 @@ MODEL_ARGS=( --attention-dropout 0.0 --hidden-dropout 0.0 --swiglu + --normalization RMSNorm --init-method-std 0.0134 --attention-backend fused --apply-layernorm-1p diff --git a/examples/mimo/data/energon_avlm_task_encoder.py b/examples/mimo/data/energon_avlm_task_encoder.py index 32afb1b2cfb..a6a86761720 100644 --- a/examples/mimo/data/energon_avlm_task_encoder.py +++ b/examples/mimo/data/energon_avlm_task_encoder.py @@ -1,3 +1,5 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + import argparse import logging import os @@ -39,7 +41,6 @@ ) from megatron.energon.task_encoder.base import stateless from megatron.training import get_args -from megatron.training.tokenizer.multimodal_tokenizer import mistral_custom_template IMAGE_TOKEN = "" AUDIO_TOKEN = "