diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 7496893749c..d9b619f9559 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 @@ -37,6 +39,9 @@ megatron/post_training/ @NVIDIA/post-training megatron/core/transformer/cuda_graphs.py @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/cuda-graphs +megatron/training/ @NVIDIA/training-adlr @NVIDIA/training-nemo +megatron/training/arguments.py + .gitlab/ @NVIDIA/ci .github/ @NVIDIA/ci .gitlab-ci.yml @NVIDIA/ci diff --git a/.github/actions/action.yml b/.github/actions/action.yml index f3e42e5843d..1236d694a99 100644 --- a/.github/actions/action.yml +++ b/.github/actions/action.yml @@ -45,10 +45,22 @@ inputs: PAT: description: "GitHub Personal Access Token" required: true - is_ci_workload: - description: "Is CI workload" - required: true - + scope: + description: "Test scope (e.g. mr-github, mr-github-slim)" + required: false + default: "mr-github-slim" + n_repeat: + description: "Number of test repetitions" + required: false + default: "5" + lightweight: + description: "Enable lightweight mode" + required: false + default: "false" + platform: + description: "Platform to run tests on (e.g. dgx_h100, dgx_gb200)" + required: false + default: "dgx_h100" runs: using: "composite" steps: @@ -56,40 +68,8 @@ runs: shell: bash -x -e -u -o pipefail {0} run: echo "node_name=$NODE_NAME" | tee -a "$GITHUB_OUTPUT" - - name: GPU Sanity Check - shell: bash -x -e -u -o pipefail {0} - run: | - echo "Starting GPU Sanity Check..." - - # 1. Check for active Compute Processes - # query-compute-apps returns a list of PIDs using the GPU. If empty, we are good. - OPEN_PROCESSES=$(docker run --rm --gpus all ubuntu nvidia-smi --query-compute-apps=pid,process_name --format=csv,noheader) - - if [ -n "$OPEN_PROCESSES" ]; then - echo "::error::❌ GPU is not clean! Found active processes:" - echo "$OPEN_PROCESSES" - else - echo "✅ No active compute processes found." - fi - - # 2. Check VRAM Usage (Optional but recommended) - # We allow a small buffer (e.g., < 300MiB) for driver overhead/Xorg, - # though on headless K8s nodes this should be very close to 0. - - MEMORY_USAGES=$(docker run --rm --gpus all ubuntu nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits) - - # Check each GPU visible to the container - for MEMORY in $MEMORY_USAGES; do - if [ "$MEMORY" -gt 300 ]; then - echo "::error::❌ GPU VRAM usage is suspiciously high: ${MEMORY} MiB" - fi - done - - echo "✅ GPU Memory is clean (all < 300 MiB)." - echo "Ready to start workflow." - - name: Checkout repository - uses: actions/checkout@v2 + uses: actions/checkout@v6 - name: Change ownership of /home/runner/ shell: bash @@ -117,55 +97,25 @@ 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 \ --test-case "${{ inputs.test_case }}" \ --environment dev \ - --platform dgx_h100 \ + --platform ${{ inputs.platform }} \ --tag ${{ inputs.tag }} \ - --container-image ${{ inputs.container-image }} + --container-image ${{ inputs.container-image }} \ + --hf-home /mnt/datadrive/TestData/nemo-fw/TestData/HF_HOME RUN_TEST_EOF ) echo "$cmd" | tee "job.sh" echo "::endgroup::" - - name: Get PR info - id: get-pr-info - if: startsWith(github.ref, 'refs/heads/pull-request/') - uses: nv-gha-runners/get-pr-info@main - - - name: Install GH CLI - shell: bash -x -e -u -o pipefail {0} - run: | - apt-get update - apt-get install -y gh - - - name: Has Run tests label - shell: bash -x -e -u -o pipefail {0} - id: has-run-tests-label - env: - GH_TOKEN: ${{ github.token }} - run: | - PR_NUMBER=${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number }} - 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 }} - IS_CI_WORKLOAD: ${{ inputs.is_ci_workload }} - 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 "$IS_CI_WORKLOAD" - HAS_RUN_FUNCTIONAL_TESTS_LABEL=${HAS_RUN_FUNCTIONAL_TESTS_LABEL:-$IS_CI_WORKLOAD} - 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' @@ -177,36 +127,29 @@ runs: #!/bin/bash set -euxo pipefail - if [ "${{ steps.has-run-tests-label.outputs.main }}" == "true" ]; then - ARGS=( - --scope mr-github - --enable-lightweight-mode - --n-repeat 1 - ) - elif [ "${{ steps.has-run-functional-tests-label.outputs.main }}" == "true" ]; then - ARGS=( - --scope mr-github - --n-repeat 5 - ) - else - ARGS=( - --scope mr-github-slim - --n-repeat 5 - ) + ARGS=( + --scope ${{ inputs.scope }} + --n-repeat ${{ inputs.n_repeat }} + ) + if [ "${{ inputs.lightweight }}" == "true" ]; then + ARGS+=(--enable-lightweight-mode) fi 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 }} \ --test-case ${{ inputs.test_case }} \ --environment dev \ - --platform dgx_h100 \ + --platform ${{ inputs.platform }} \ --container-image ${{ inputs.container-image }} \ --data-dir /mnt/datadrive/TestData/megatron-lm/artifacts \ + --hf-home /mnt/datadrive/TestData/nemo-fw/TestData/HF_HOME RUN_TEST_EOF ) diff --git a/.github/copy-pr-bot.yaml b/.github/copy-pr-bot.yaml index 72a5b915ecc..191a34902a3 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", "dingqingy-nv", "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", "CarlosGomes98", "ChenhanYu", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JRD971000", "Phlip79", "QiZhangNV", "RPrenger", "ShriyaRishab", "Victarry", "WanZzzzzz", "Wohox", "ZhiyuLi-Nvidia", "ahmadki", "aklife97", "ananthsub", "asolergi-nv", "buptzyb", "chtruong814", "cjld", "cspades", "cuichenx", "deepakn94", "dimapihtar", "dingqingy-nv", "duncanriach", "erhoo82", "ericharper", "fanshiqing", "faradawn", "frsun-nvda", "gautham-kollu", "gdengk", "guyueh1", "hexinw-nvidia", "huvunvidia", "hxbai", "ilml", "jalbericiola", "janEbert", "jaredcasper", "jenchen13", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kajalj22", "kanz-nv", "kevalmorabia97", "ko3n1g", "ksivaman", "kunlunl", "kvareddy", "kwyss-nvidia", "layalir", "lhb8125", "lmcafee-nvidia", "maanug-nv", "mathemakitten", "matthieule", "mchrzanowski", "mehraakash", "mkhona-nvidia", "nanz-nv", "parthmannan", "prajwal1210", "pthombre", "rhewett-nv", "rogerwaleffe", "sajadn", "sanandaraj5597", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "sharathts", "shengf-nv", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sudhakarsingh27", "tdene", "theothermike", "thomasdhc", "tomlifu", "trintamaki", "tylerpoon", "wdykas", "wplf", "xiaoyao0115", "xuwchen", "yanring", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yueshen2016", "yuzhongw-nvidia", "zhongbozhu"] diff --git a/.github/oncall_schedule.json b/.github/oncall_schedule.json index 5fa49e966bc..286fc3b1e13 100644 --- a/.github/oncall_schedule.json +++ b/.github/oncall_schedule.json @@ -1,50 +1,50 @@ [ { - "user": "dimapihtar", - "date": "2026-01-28" + "user": "janEbert", + "date": "2026-03-25" }, { "user": "gautham-kollu", - "date": "2026-02-04" + "date": "2026-04-01" }, { - "user": "janEbert", - "date": "2026-02-11" + "user": "ilml", + "date": "2026-04-08" }, { "user": "Phlip79", - "date": "2026-02-18" + "date": "2026-04-15" }, { "user": "asolergi-nv", - "date": "2026-02-25" + "date": "2026-04-22" }, { "user": "BoxiangW", - "date": "2026-03-04" + "date": "2026-04-29" }, { "user": "maanug-nv", - "date": "2026-03-11" + "date": "2026-05-06" }, { "user": "dimapihtar", - "date": "2026-03-18" + "date": "2026-05-13" }, { "user": "gautham-kollu", - "date": "2026-03-25" + "date": "2026-05-20" }, { - "user": "janEbert", - "date": "2026-04-01" + "user": "ilml", + "date": "2026-05-27" }, { - "user": "maanug-nv", - "date": "2026-04-08" + "user": "janEbert", + "date": "2026-06-03" }, { - "user": "BoxiangW", - "date": "2026-04-15" + "user": "maanug-nv", + "date": "2026-06-10" } ] diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 5cd5138eb69..d2825f9c34b 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -5,19 +5,8 @@ ## Contribution process -```mermaid -flowchart LR - A[Pre-checks] --> B[PR Tests] - subgraph Code Review/Approval - C1[Expert Review] --> C2[Final Review] - end - B --> C1 - C2 --> D[Merge] -``` - ### Pre-checks -- [ ] I want this PR in a versioned release and have added the appropriate Milestone (e.g., `Core 0.8`) - [ ] I have added relevant unit tests - [ ] I have added relevant functional tests - [ ] I have added proper typing to my code [Typing guidelines](https://docs.python.org/3/library/typing.html) @@ -26,33 +15,32 @@ flowchart LR ### Code review -The following process is enforced via the CODEOWNERS file for changes into `megatron/core`. For changes outside of `megatron/core`, it is up to the PR author whether or not to tag the Final Reviewer team. +Feel free to message or comment the [@mcore-oncall](https://github.com/orgs/NVIDIA/teams/mcore-oncall) to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged! -
-For MRs into `main` branch +All PRs start as **draft**. If you open a non-draft PR, it will be automatically converted to draft. -Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged! +#### Step 1: Mark PR as "Ready for Review" -#### (Step 1): Add PR label `Expert Review` +1. When your PR is ready, click **Ready for Review**. +2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes. + - Some PRs may jump straight to step 2. This is determined by `.github/CODEOWNERS`. -#### (Step 2): Collect the expert reviewers reviews +:warning: Only mark as ready once merge-conflicts are resolved and the CI is passing. +Final Review might get declined if these requirements are not fulfilled. -1. Attach the `Expert Review` label when your PR is ready for review. -2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon. +#### Step 2: Final Review -:warning: Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing. -Final Review might get declined if these requirements are not fulfilled. +For PRs that change `megatron/core`, once all expert reviewers have approved, the `Final Review` label is applied **automatically** and final reviewers are assigned. -#### (Step 3): Final Review +For PRs outside `megatron/core`, this step is skipped. -1. Add `Final Review` label -2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon. +#### Step 3: Approved -#### (Optional Step 4): Cherry-pick into release branch +Once all required reviewers have approved, the `Approved` label is applied **automatically**. -If this PR also needs to be merged into `core_r*` release branches, after this PR has been merged, select `Cherry-pick` to open a new PR into the release branch. +### Merge -
+Any member of [mcore-engineers](https://github.com/orgs/NVIDIA/teams/mcore-engineers) will be able to merge your PR.
For MRs into `dev` branch @@ -60,7 +48,3 @@ The proposed review process for `dev` branch is under active discussion. MRs are mergable after one approval by either `eharper@nvidia.com` or `zijiey@nvidia.com`.
- -### Merging your PR - -Any member of [core-adlr](https://github.com/orgs/teams/NVIDIA/core-adlr) and [`core-nemo`](https://github.com/orgs/teams/NVIDIA/core-nemo) will be able to merge your PR. diff --git a/.github/scripts/readme.sh b/.github/scripts/readme.sh new file mode 100644 index 00000000000..216d5224a28 --- /dev/null +++ b/.github/scripts/readme.sh @@ -0,0 +1,65 @@ +#!/bin/bash + +cat << 'EOF' +╔══════════════════════════════════════════════════════════════════════╗ +║ ║ +║ ███╗ ███╗██████╗ ██████╗ ██╗██████╗ ██████╗ ███████╗ ║ +║ ████╗ ████║██╔══██╗██╔══██╗██║██╔══██╗██╔════╝ ██╔════╝ ║ +║ ██╔████╔██║██████╔╝██████╔╝██║██║ ██║██║ ███╗█████╗ ║ +║ ██║╚██╔╝██║██╔══██╗██╔══██╗██║██║ ██║██║ ██║██╔══╝ ║ +║ ██║ ╚═╝ ██║██████╔╝██║ ██║██║██████╔╝╚██████╔╝███████╗ ║ +║ ╚═╝ ╚═╝╚═════╝ ╚═╝ ╚═╝╚═╝╚═════╝ ╚═════╝ ╚══════╝ ║ +║ ║ +║ H O W T O : M B R I D G E T E S T I N G ║ +╚══════════════════════════════════════════════════════════════════════╝ + + MBridge unit tests run automatically on every PR. To also trigger + functional tests, attach the label and re-run the workflow step. + + ┌─────────────────────────────────────────────────────────────────┐ + │ DEFAULT │ Unit tests run on every PR (no action needed) │ + ├─────────────────────────────────────────────────────────────────┤ + │ │ + │ Every PR ──► cicd-mbridge-testing ──► unit tests only │ + │ │ + └─────────────────────────────────────────────────────────────────┘ + + ┌─────────────────────────────────────────────────────────────────┐ + │ STEP 1 │ Attach the label to your PR (for functional tests) │ + ├─────────────────────────────────────────────────────────────────┤ + │ │ + │ PR Labels ──► [ + Add label ] ──► "Run MBridge tests" │ + │ │ + └─────────────────────────────────────────────────────────────────┘ + + ┌─────────────────────────────────────────────────────────────────┐ + │ STEP 2 │ Re-run this workflow step │ + ├─────────────────────────────────────────────────────────────────┤ + │ │ + │ Actions ──► [ Re-run jobs ] ──► Re-run failed jobs │ + │ │ + └─────────────────────────────────────────────────────────────────┘ + + ┌─────────────────────────────────────────────────────────────────┐ + │ RESULT │ Unit + functional tests run! │ + ├─────────────────────────────────────────────────────────────────┤ + │ │ + │ cicd-mbridge-testing ◄── unit + functional tests │ + │ │ + │ Tests run against MBridge using the merge commit │ + │ SHA of your pull request. │ + │ │ + └─────────────────────────────────────────────────────────────────┘ + + ┌────────────────────────────────────┐ + │ Label present? NO → unit │ + │ Label present? YES → unit + │ + │ functional│ + └────────────────────────────────────┘ + + NOTE: The label must be present BEFORE the re-run is triggered. + The CI checks for "Run MBridge tests" at runtime. + + NOTE: All MBridge test results are optional — failures do not + block merging your PR. +EOF diff --git a/.github/scripts/sync_team_usergroups.py b/.github/scripts/sync_team_usergroups.py index 429387fc6de..c3fa5d474ff 100644 --- a/.github/scripts/sync_team_usergroups.py +++ b/.github/scripts/sync_team_usergroups.py @@ -29,7 +29,12 @@ # Constants GITHUB_API_URL = "https://api.github.com" -PARENT_TEAM_SLUG = "mcore-reviewers" + +# Teams whose *children* are each synced to their own Slack usergroup +PARENT_TEAM_SLUGS = ["mcore-reviewers"] + +# Teams synced directly (the team itself, not its children) +DIRECT_TEAM_SLUGS = ["mcore-engineers"] # Caches for email and Slack lookups _email_cache = {} @@ -83,6 +88,8 @@ def github_team_to_slack_usergroup(team_slug): name = name[5:] # Remove "core-" elif name.startswith("megatron-"): name = name[9:] # Remove "megatron-" + elif name.startswith("mcore-"): + name = name[6:] # Remove "mcore-" # Remove "-and-" name = name.replace("-and-", "-") @@ -437,13 +444,13 @@ def sync_team_to_usergroup(team_slug, usergroup_handle, dry_run=False): return False -def get_team_to_usergroup_mapping(): - """Fetch child teams of mcore-reviewers and generate the mapping.""" +def get_team_to_usergroup_mapping(parent_team_slug): + """Fetch child teams of a parent team and generate the mapping.""" org = get_org() - child_teams = get_child_teams(org, PARENT_TEAM_SLUG) + child_teams = get_child_teams(org, parent_team_slug) if not child_teams: - print(f"Error: No child teams found under '{PARENT_TEAM_SLUG}'") + print(f"Error: No child teams found under '{parent_team_slug}'") return {} mapping = {} @@ -454,10 +461,30 @@ def get_team_to_usergroup_mapping(): return mapping -def sync_all_teams(dry_run=False): - """Sync all GitHub teams under mcore-reviewers to their Slack usergroups.""" - print(f"Fetching child teams of '{PARENT_TEAM_SLUG}'...") - team_to_usergroup = get_team_to_usergroup_mapping() +def sync_all_teams(dry_run=False, parent_teams=None, direct_teams=None): + """Sync GitHub teams to their Slack usergroups. + + Args: + parent_teams: List of team slugs whose *children* are each synced. + Defaults to PARENT_TEAM_SLUGS. + direct_teams: List of team slugs synced directly (not their children). + Defaults to DIRECT_TEAM_SLUGS. + """ + if parent_teams is None: + parent_teams = PARENT_TEAM_SLUGS + if direct_teams is None: + direct_teams = DIRECT_TEAM_SLUGS + + team_to_usergroup = {} + + for parent_slug in parent_teams: + print(f"Fetching child teams of '{parent_slug}'...") + mapping = get_team_to_usergroup_mapping(parent_slug) + team_to_usergroup.update(mapping) + + for team_slug in direct_teams: + usergroup_handle = github_team_to_slack_usergroup(team_slug) + team_to_usergroup[team_slug] = usergroup_handle if not team_to_usergroup: return False @@ -504,12 +531,40 @@ def main(): action="store_true", help="List all configured team-to-usergroup mappings", ) + parser.add_argument( + "--parent-team", + action="append", + dest="parent_teams", + metavar="SLUG", + help=( + "Sync all children of this GitHub team (can be repeated). " + f"Defaults to: {PARENT_TEAM_SLUGS}" + ), + ) + parser.add_argument( + "--team", + action="append", + dest="direct_teams", + metavar="SLUG", + help=( + "Sync this GitHub team directly (can be repeated). " + f"Defaults to: {DIRECT_TEAM_SLUGS}" + ), + ) args = parser.parse_args() + # Use CLI values when provided, otherwise fall back to module-level defaults + parent_teams = args.parent_teams if args.parent_teams is not None else PARENT_TEAM_SLUGS + direct_teams = args.direct_teams if args.direct_teams is not None else DIRECT_TEAM_SLUGS + if args.list: - print(f"Fetching child teams of '{PARENT_TEAM_SLUG}'...") - team_to_usergroup = get_team_to_usergroup_mapping() + team_to_usergroup = {} + for parent_slug in parent_teams: + print(f"Fetching child teams of '{parent_slug}'...") + team_to_usergroup.update(get_team_to_usergroup_mapping(parent_slug)) + for team_slug in direct_teams: + team_to_usergroup[team_slug] = github_team_to_slack_usergroup(team_slug) if not team_to_usergroup: sys.exit(1) print("\nTeam-to-usergroup mappings:") @@ -519,7 +574,9 @@ def main(): print(f"{team:<35} @{usergroup:<29}") return - success = sync_all_teams(dry_run=args.dry_run) + success = sync_all_teams( + dry_run=args.dry_run, parent_teams=parent_teams, direct_teams=direct_teams + ) sys.exit(0 if success else 1) diff --git a/.github/workflows/_build_test_publish_wheel.yml b/.github/workflows/_build_test_publish_wheel.yml index 9e9062827de..0c456d7164f 100644 --- a/.github/workflows/_build_test_publish_wheel.yml +++ b/.github/workflows/_build_test_publish_wheel.yml @@ -17,8 +17,6 @@ on: type: boolean default: true secrets: - TWINE_USERNAME: - required: true TWINE_PASSWORD: required: true @@ -45,7 +43,7 @@ jobs: PUBLISH_DRYRUN: ${{ inputs.dry-run }} steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: ref: ${{ inputs.ref }} @@ -73,10 +71,10 @@ jobs: pushd $BUILD_DIR rm LICENSE || true docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE bash -c '\ - for python_version in cp310 cp311 cp312 cp313; do \ + for python_version in cp311 cp312 cp313; do \ /opt/python/${python_version}-${python_version}/bin/pip install --upgrade "setuptools<80.0.0,>=77.0.0" build; \ done && \ - for python_version in cp310 cp311 cp312 cp313; do \ + for python_version in cp311 cp312 cp313; do \ /opt/python/${python_version}-${python_version}/bin/python -m build; \ done \ ' @@ -120,25 +118,23 @@ jobs: if [ "$PACKAGE" = "megatron-core" ]; then if [[ "$PLATFORM" == "arm64" ]]; then - for file in dist/$WHEEL_PREFIX*cp310*aarch64.whl; do - pip install --no-cache-dir "$file" - done + WHEEL_GLOB="dist/${WHEEL_PREFIX}*cp312*aarch64.whl" else - for file in dist/$WHEEL_PREFIX*cp310*x86_64.whl; do - pip install --no-cache-dir "$file" - done + WHEEL_GLOB="dist/${WHEEL_PREFIX}*cp312*x86_64.whl" fi else - pip install --no-cache-dir dist/$WHEEL_PREFIX*.whl + WHEEL_GLOB="dist/${WHEEL_PREFIX}*.whl" fi - sudo rm -rf megatron/ - - RELEASE_NUMBER=$(python -c "import $ROOTPATH; print($ROOTPATH.__version__)") - test "${{ steps.build-wheel.outputs.expected-release-number }}" == "$RELEASE_NUMBER" + docker run --rm -v $(pwd):/workspace -w /workspace $IMAGE bash -c "\ + /opt/python/cp312-cp312/bin/pip install --no-cache-dir $WHEEL_GLOB && \ + rm -rf megatron/ && \ + RELEASE_NUMBER=\$(/opt/python/cp312-cp312/bin/python -c 'import $ROOTPATH; print($ROOTPATH.__version__)') && \ + test '${{ steps.build-wheel.outputs.expected-release-number }}' == \"\$RELEASE_NUMBER\" \ + " - name: Upload wheels - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: wheels-${{ matrix.PACKAGE }}-${{ matrix.PLATFORM }}-${{ inputs.dry-run && 'dry-run' || 'release' }} path: dist/ @@ -147,7 +143,6 @@ jobs: needs: [build-and-test-wheels] runs-on: ubuntu-latest if: inputs.no-publish == false - environment: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/r')) && 'main' || 'public' }} strategy: fail-fast: false matrix: @@ -162,7 +157,7 @@ jobs: PACKAGE: ${{ matrix.PACKAGE }} steps: - name: Download wheels - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v7 with: name: wheels-${{ matrix.PACKAGE }}-${{ matrix.PLATFORM }}-${{ inputs.dry-run && 'dry-run' || 'release' }} path: dist/ @@ -170,7 +165,7 @@ jobs: - name: Publish wheels env: - TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }} + TWINE_USERNAME: __token__ TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }} TWINE_REPOSITORY: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/r')) && 'pypi' || 'testpypi' }} PLATFORM: ${{ matrix.PLATFORM }} diff --git a/.github/workflows/_release_library.yml b/.github/workflows/_release_library.yml index d39ee505c2a..954d81c8259 100644 --- a/.github/workflows/_release_library.yml +++ b/.github/workflows/_release_library.yml @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -name: 'Release' +name: "Release" defaults: run: @@ -38,17 +38,49 @@ on: description: Create a GitHub release type: boolean default: true + gh-release-use-changelog-builder: + required: false + description: Use release-changelog-builder-action to dynamically build changelog + type: boolean + default: true + gh-release-changelog-config: + required: false + description: Path to changelog builder configuration file + type: string + default: ".github/workflows/config/changelog-config.json" + gh-release-from-tag: + required: false + description: Starting tag for changelog builder (leave empty for auto-detect) + type: string + default: "" + publish-docs: + required: false + description: Publish documentation to S3 after release + type: boolean + default: true secrets: - TWINE_USERNAME: - required: true TWINE_PASSWORD: required: true - SLACK_WEBHOOK_ADMIN: - required: true SLACK_WEBHOOK: required: true PAT: required: true + AWS_ASSUME_ROLE_ARN: + required: true + AWS_ACCESS_KEY_ID: + required: true + AWS_SECRET_ACCESS_KEY: + required: true + AKAMAI_HOST: + required: true + AKAMAI_CLIENT_TOKEN: + required: true + AKAMAI_CLIENT_SECRET: + required: true + AKAMAI_ACCESS_TOKEN: + required: true + S3_BUCKET_NAME: + required: true permissions: contents: write # To read repository content @@ -62,12 +94,10 @@ jobs: ref: ${{ inputs.release-ref }} no-publish: true secrets: - TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }} TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }} bump-next-version: runs-on: ubuntu-latest - environment: main # ${{ inputs.dry-run == true && 'public' || 'main' }} needs: build-test-publish-wheels-dry-run if: | ( @@ -80,7 +110,7 @@ jobs: IS_DRY_RUN: ${{ inputs.dry-run }} steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: path: ${{ github.run_id }} token: ${{ secrets.PAT }} @@ -90,8 +120,8 @@ jobs: - name: Bump version MCore id: bump-version-mcore env: - SRC_DIR: '' - PYPROJECT_NAME: 'megatron.core' + SRC_DIR: "" + PYPROJECT_NAME: "megatron.core" run: | set +u cd ${{ github.run_id }} @@ -129,8 +159,8 @@ jobs: - name: Bump version MFSDP id: bump-version-mfsdp env: - SRC_DIR: 'megatron/core/distributed/fsdp/src/' - PYPROJECT_NAME: 'megatron_fsdp' + SRC_DIR: "megatron/core/distributed/fsdp/src/" + PYPROJECT_NAME: "megatron_fsdp" run: | set +u @@ -190,11 +220,8 @@ jobs: # Extract PR number from URL PR_NUMBER=$(echo $PR_URL | grep -o '[0-9]*$') - # Add comment to the newly created PR - echo gh pr comment $PR_NUMBER --body "/ok to test $(git rev-parse HEAD)" - - name: Wait for status checks on tmp branch - uses: actions/github-script@v7 + uses: actions/github-script@v8 id: wait-status with: github-token: ${{ secrets.PAT }} @@ -317,13 +344,11 @@ jobs: ref: ${{ inputs.release-ref }} no-publish: false secrets: - TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }} TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }} create-gh-release: needs: [build-test-publish-wheels, bump-next-version] runs-on: ubuntu-latest - environment: ${{ inputs.dry-run == true && 'public' || 'main' }} if: | ( success() || !failure() @@ -336,21 +361,60 @@ jobs: REPOSITORY: ${{ github.repository }} PROJECT_NAME: Megatron Core VERSION: ${{ needs.bump-next-version.outputs.release-version }} - TAG_PREFIX: ${{ inputs.gh-release-tag-prefix || '' }} + TAG_PREFIX: core_ steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: path: ${{ github.run_id }} ref: ${{ inputs.release-ref }} token: ${{ secrets.PAT || secrets.GITHUB_TOKEN }} + - name: Determine fromTag for changelog + id: determine-from-tag + if: inputs.gh-release-use-changelog-builder == true + run: | + cd ${{ github.run_id }} + + # If gh-release-from-tag is provided, use it + if [[ -n "${{ inputs.gh-release-from-tag }}" ]]; then + FROM_TAG="${{ inputs.gh-release-from-tag }}" + echo "Using provided fromTag: $FROM_TAG" + else + # Get the most recent tag + FROM_TAG=$(git describe --tags --abbrev=0 2>/dev/null || echo "") + if [[ -z "$FROM_TAG" ]]; then + echo "No previous tags found, leaving fromTag empty" + else + echo "Auto-detected most recent tag: $FROM_TAG" + fi + fi + + echo "from-tag=$FROM_TAG" >> $GITHUB_OUTPUT + + - name: Build Changelog + id: build-changelog + if: inputs.gh-release-use-changelog-builder == true + uses: mikepenz/release-changelog-builder-action@v6.1.0 + env: + GITHUB_TOKEN: ${{ secrets.PAT || secrets.GITHUB_TOKEN }} + with: + configuration: ${{ github.run_id }}/${{ inputs.gh-release-changelog-config }} + owner: ${{ github.repository_owner }} + repo: ${{ github.event.repository.name }} + ignorePreReleases: "false" + failOnError: "false" + fromTag: ${{ steps.determine-from-tag.outputs.from-tag }} + toTag: ${{ inputs.release-ref }} + mode: ${{ inputs.gh-release-changelog-mode }} + - name: Create release id: version-number env: SHA: ${{ inputs.release-ref }} GH_TOKEN: ${{ secrets.PAT }} IS_DRY_RUN: ${{ inputs.dry-run }} + BUILT_CHANGELOG: ${{ steps.build-changelog.outputs.changelog }} run: | cd ${{ github.run_id }} @@ -359,7 +423,10 @@ jobs: IS_PRERELEASE=$([[ "$IS_RELEASE_CANDIDATE" == "true" || "$IS_ALPHA" == "true" ]] && echo "true" || echo "false") NAME="NVIDIA $PROJECT_NAME ${VERSION}" - if [[ "$IS_RELEASE_CANDIDATE" == "true" ]]; then + # Use built changelog if available, otherwise fall back to CHANGELOG.md + if [[ -n "$BUILT_CHANGELOG" ]]; then + CHANGELOG="$BUILT_CHANGELOG" + elif [[ "$IS_RELEASE_CANDIDATE" == "true" ]]; then DATE=$(date +"%Y-%m-%d") CHANGELOG="Prerelease: $NAME ($DATE)" else @@ -402,18 +469,28 @@ jobs: eval "$CMD" fi + publish-docs: + needs: [bump-next-version, create-gh-release] + uses: ./.github/workflows/release-docs.yml + if: | + ( + success() || !failure() + ) + && inputs.publish-docs == true + && !cancelled() + with: + dry-run: ${{ inputs.dry-run }} + publish-as-latest: true + docs-version-override: ${{ needs.bump-next-version.outputs.release-version }} + build-docs-ref: ${{ inputs.release-ref }} + secrets: inherit + notify: - needs: [build-test-publish-wheels, create-gh-release] + needs: [build-test-publish-wheels, create-gh-release, bump-next-version] runs-on: ubuntu-latest - environment: ${{ inputs.dry-run == true && 'public' || 'main' }} - env: - GH_URL: https://github.com/${{ github.repository }}/releases/tag/v${{ needs.build-test-publish-wheels.outputs.version }} - PYPI_URL: https://${{ inputs.dry-run == true && 'test.' || '' }}pypi.org/project/${{ needs.build-test-publish-wheels.outputs.pypi-name }}/${{ needs.build-test-publish-wheels.outputs.version }}/ - PROJECT_NAME: Megatron Core - VERSION: ${{ needs.build-test-publish-wheels.outputs.version }} steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: repository: NVIDIA-NeMo/FW-CI-templates ref: v0.17.0 @@ -423,10 +500,10 @@ jobs: uses: ./send-slack-alert/.github/actions/send-slack-alert env: MESSAGE: | - ${{ inputs.dry-run == true && 'This is a dry-run, nothing actually happened: ' || '' }}We have released `${{ env.VERSION }}` of `NVIDIA ${{ env.PROJECT_NAME }}` 🚀✨🎉 + ${{ inputs.dry-run == true && 'This is a dry-run, nothing actually happened: ' || '' }}We have released `${{ needs.bump-next-version.outputs.release-version }}` of `NVIDIA Megatron Core` 🚀✨🎉 - • <${{ env.GH_URL }}|GitHub release> - • <${{ env.PYPI_URL }}|PyPi release> + • + • with: message: ${{ env.MESSAGE }} diff --git a/.github/workflows/_update_dependencies.yml b/.github/workflows/_update_dependencies.yml index 063b966b5de..903d773edbd 100644 --- a/.github/workflows/_update_dependencies.yml +++ b/.github/workflows/_update_dependencies.yml @@ -9,12 +9,6 @@ on: secrets: PAT: required: true - AZURE_CLIENT_ID: - required: true - AZURE_TENANT_ID: - required: true - AZURE_SUBSCRIPTION_ID: - required: true SSH_KEY: required: true SSH_PWD: @@ -32,36 +26,31 @@ jobs: run: echo "date=$(date +%F)" | tee -a "$GITHUB_OUTPUT" update-lockfile: - environment: nemo-ci runs-on: linux-amd64-cpu16 needs: [pre-flight] env: SOURCE_BRANCH: ${{ needs.pre-flight.outputs.bump-branch }} TARGET_BRANCH: ${{ inputs.target-branch }} steps: - - name: Install Azure CLI - run: curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash - - - name: Azure Login - uses: azure/login@v2 - with: - client-id: ${{ secrets.AZURE_CLIENT_ID }} - tenant-id: ${{ secrets.AZURE_TENANT_ID }} - subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }} - - - name: Azure ACR Login - run: az acr login --name nemoci - - name: Checkout repo - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: ref: ${{ env.TARGET_BRANCH }} + - name: Mock test data + run: mkdir -p assets/ + + - name: Fetch NGC Version + id: ngc-version + run: | + NGC_VERSION=$(cat docker/.ngc_version.dev) + echo "NGC_VERSION=${NGC_VERSION}" | tee -a "$GITHUB_OUTPUT" + - name: Build container env: GH_TOKEN: ${{ secrets.PAT }} run: | - docker build -f docker/Dockerfile.ci.dev --build-arg FROM_IMAGE_NAME="nvcr.io/nvidia/pytorch:25.06-py3" --target=main -t megatron-core . + docker build -f docker/Dockerfile.ci.dev --build-arg FROM_IMAGE_NAME="${{ steps.ngc-version.outputs.NGC_VERSION }}" --target=main -t megatron-core . - name: Create bump branch if not exists run: | @@ -71,7 +60,7 @@ jobs: fi - name: Checkout repo - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: ref: ${{ env.SOURCE_BRANCH }} @@ -88,7 +77,7 @@ jobs: bash -c 'uv lock --upgrade' - name: Upload lock file - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: lock-file-${{ env.SOURCE_BRANCH }} path: uv.lock @@ -96,29 +85,16 @@ jobs: create-pr: needs: [update-lockfile, pre-flight] runs-on: ubuntu-latest - environment: main env: SOURCE_BRANCH: ${{ needs.pre-flight.outputs.bump-branch }} TARGET_BRANCH: ${{ inputs.target-branch }} steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: token: ${{ secrets.PAT }} ref: ${{ env.TARGET_BRANCH }} - - name: Install GPG - run: sudo apt-get install -y gnupg2 - - - name: Import GPG key (for signing) - uses: crazy-max/ghaction-import-gpg@e89d40939c28e39f97cf32126055eeae86ba74ec - id: gpg-action - with: - gpg_private_key: ${{ secrets.SSH_KEY }} - passphrase: ${{ secrets.SSH_PWD }} - git_user_signingkey: true - git_commit_gpgsign: true - - name: Rebase against ${{ env.SOURCE_BRANCH }} run: | if git ls-remote --exit-code origin ${{ env.SOURCE_BRANCH }}; then @@ -127,7 +103,7 @@ jobs: fi - name: Download lock file - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v7 with: name: lock-file-${{ env.SOURCE_BRANCH }} @@ -150,4 +126,102 @@ jobs: 🙏 Please merge this PR only if the CI workflow completed successfully. commit-message: ${{ env.title }} signoff: true - committer: "${{ steps.gpg-action.outputs.name }} <${{ steps.gpg-action.outputs.email }}>" + committer: "github-actions[bot] " + + - name: Post /ok to test comment + env: + GH_TOKEN: ${{ secrets.PAT }} + run: | + PR_NUMBER="${{ steps.create-pull-request.outputs.pull-request-number }}" + if [ -z "$PR_NUMBER" ]; then + echo "No PR was created, skipping comment" + exit 0 + fi + SHA="${{ steps.create-pull-request.outputs.pull-request-head-sha }}" + gh pr comment "$PR_NUMBER" --body "/ok to test $SHA" + + - name: Wait for CI checks + env: + GH_TOKEN: ${{ secrets.PAT }} + run: | + PR_NUMBER="${{ steps.create-pull-request.outputs.pull-request-number }}" + if [ -z "$PR_NUMBER" ]; then + echo "No PR was created, skipping wait" + exit 0 + fi + + # Fetch required status checks from branch protection rules + REQUIRED_CHECKS=$(gh api \ + "repos/${{ github.repository }}/branches/${{ env.TARGET_BRANCH }}/protection/required_status_checks" \ + --jq '.checks[].context' 2>/dev/null \ + || gh api \ + "repos/${{ github.repository }}/branches/${{ env.TARGET_BRANCH }}/protection/required_status_checks" \ + --jq '.contexts[]' 2>/dev/null \ + || true) + + if [ -z "$REQUIRED_CHECKS" ]; then + echo "No branch protection rules found for ${{ env.TARGET_BRANCH }}, skipping wait" + exit 0 + fi + + echo "Required checks from branch protection:" + echo "$REQUIRED_CHECKS" + + echo "Waiting for required checks to complete on PR #$PR_NUMBER..." + i=0 + INITIALIZED=false + while true; do + i=$((i + 1)) + CHECKS_JSON=$(gh pr checks "$PR_NUMBER" --json name,state 2>/dev/null || echo "[]") + ALL_DONE=true + FAILED_CHECKS="" + while IFS= read -r check; do + CHECK_STATE=$(echo "$CHECKS_JSON" | jq -r --arg name "$check" '.[] | select(.name == $name) | .state // ""' | tr '[:upper:]' '[:lower:]') + case "$CHECK_STATE" in + *success*|*pass*|*skip*|*neutral*) ;; + *pending*|*queued*|*progress*|*waiting*|*request*|"") + ALL_DONE=false + INITIALIZED=true + break + ;; + *) + if [ "$INITIALIZED" = "true" ]; then + FAILED_CHECKS="${FAILED_CHECKS} - ${check} (${CHECK_STATE})"$'\n' + else + ALL_DONE=false + fi + ;; + esac + done <<< "$REQUIRED_CHECKS" + if [ "$ALL_DONE" = "true" ]; then + if [ -n "$FAILED_CHECKS" ]; then + echo "Required check(s) did not pass:" + echo "$FAILED_CHECKS" + exit 1 + fi + echo "All required checks passed!" + break + fi + echo "Checks not yet complete (attempt $i), retrying in 30s..." + sleep 30 + done + + - name: Merge PR + env: + title: "chore(beep boop 🤖): Bump `uv.lock` (${{ env.TARGET_BRANCH}}) (${{ needs.pre-flight.outputs.date }})" + run: | + PR_NUMBER="${{ steps.create-pull-request.outputs.pull-request-number }}" + if [ -z "$PR_NUMBER" ]; then + echo "No PR was created, skipping merge" + exit 0 + fi + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + git fetch origin ${{ env.SOURCE_BRANCH }} + git fetch origin ${{ env.TARGET_BRANCH }} + git checkout ${{ env.TARGET_BRANCH }} + git merge --squash origin/${{ env.SOURCE_BRANCH }} + git commit -m "${{ env.title }}" + git pull --rebase origin ${{ env.TARGET_BRANCH }} + git push origin ${{ env.TARGET_BRANCH }} + git push origin --delete ${{ env.SOURCE_BRANCH }} diff --git a/.github/workflows/auto-assign-milestone.yml b/.github/workflows/auto-assign-milestone.yml index 8153728f9fd..b972329bac1 100644 --- a/.github/workflows/auto-assign-milestone.yml +++ b/.github/workflows/auto-assign-milestone.yml @@ -13,7 +13,6 @@ permissions: jobs: assign-milestone: runs-on: ubuntu-latest - environment: nemo-ci if: github.repository == 'NVIDIA/Megatron-LM' steps: - name: Get PR info diff --git a/.github/workflows/auto-reminder-bot.yml b/.github/workflows/auto-reminder-bot.yml index c3aa8169b50..72a48e9539e 100644 --- a/.github/workflows/auto-reminder-bot.yml +++ b/.github/workflows/auto-reminder-bot.yml @@ -9,16 +9,15 @@ on: jobs: run-script: - environment: main name: Run Auto Reminder Bot runs-on: ubuntu-latest if: github.repository == 'NVIDIA/Megatron-LM' steps: - name: Check out repository code - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: "3.10" @@ -28,7 +27,7 @@ jobs: - name: Run Auto Reminder Bot run: | - export SLACK_TOKEN=${{ secrets.SLACK_TOKEN }} - export SLACK_WEBHOOK_URL=${{ secrets.SLACK_WEBHOOK_URL }} + export SLACK_TOKEN=${{ secrets.SLACK_BOT_TOKEN }} + export SLACK_WEBHOOK_URL=${{ secrets.SLACK_REVIEW_REMINDER_CHANNEL_WEBHOOK }} export GH_TOKEN=${{ secrets.PAT }} python tests/test_utils/python_scripts/auto_reminder_github.py diff --git a/.github/workflows/auto-swap-labels.yml b/.github/workflows/auto-swap-labels.yml index 5335026e2af..f1dd9757c8a 100644 --- a/.github/workflows/auto-swap-labels.yml +++ b/.github/workflows/auto-swap-labels.yml @@ -2,32 +2,74 @@ name: Auto Swap Labels on: - pull_request_review: - types: [submitted] + pull_request_target: + types: [ready_for_review, synchronize] + branches: + - main + workflow_run: + workflows: ["Review Trigger"] + types: [completed] permissions: pull-requests: write contents: read + actions: read jobs: check-approval: runs-on: ubuntu-latest - if: github.event.review.state == 'approved' && github.repository == 'NVIDIA/Megatron-LM' + if: >- + github.repository == 'NVIDIA/Megatron-LM' && ( + (github.event_name == 'pull_request_target' && + github.event.pull_request.base.ref == 'main' && + !github.event.pull_request.draft) || + (github.event_name == 'workflow_run' && + github.event.workflow_run.conclusion == 'success') + ) + steps: + - name: Get PR number from workflow_run + id: get-pr + if: github.event_name == 'workflow_run' + continue-on-error: true + uses: actions/download-artifact@v4 + with: + name: pr-number + path: pr-number + github-token: ${{ github.token }} + run-id: ${{ github.event.workflow_run.id }} + + - name: Set PR number + id: pr + run: | + if [ "${{ github.event_name }}" = "workflow_run" ]; then + if [ "${{ steps.get-pr.outcome }}" != "success" ]; then + echo "No approval artifact found — review was not an approval. Skipping." + exit 0 + fi + echo "number=$(cat pr-number/number)" >> $GITHUB_OUTPUT + else + echo "number=${{ github.event.pull_request.number }}" >> $GITHUB_OUTPUT + fi + - name: Check out repository code + if: steps.pr.outputs.number uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v5 + if: steps.pr.outputs.number + uses: actions/setup-python@v6 with: python-version: "3.10" - name: Install dependencies + if: steps.pr.outputs.number run: | pip install --no-cache-dir PyGithub slack-sdk - - name: Run Auto Reminder Bot + - name: Run Auto Swap Labels + if: steps.pr.outputs.number run: | - export GH_TOKEN=${{ github.token }} - export PR_NUMBER=${{ github.event.pull_request.number }} + export GH_TOKEN=${{ secrets.PAT }} + export PR_NUMBER=${{ steps.pr.outputs.number }} python tests/test_utils/python_scripts/swap_pr_labels.py diff --git a/.github/workflows/auto-update-copy-pr-bot.yml b/.github/workflows/auto-update-copy-pr-bot.yml index 5f6f1ade9e8..07fdcfbfbb8 100644 --- a/.github/workflows/auto-update-copy-pr-bot.yml +++ b/.github/workflows/auto-update-copy-pr-bot.yml @@ -3,16 +3,15 @@ name: Auto Update Copy PR Bot on: workflow_dispatch: schedule: - - cron: '0 0 * * *' + - cron: "0 0 * * *" jobs: auto-update-copy-pr-bot: runs-on: ubuntu-latest - environment: nemo-ci if: github.repository == 'NVIDIA/Megatron-LM' steps: - name: Checkout code - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: token: ${{ secrets.PAT }} ref: main diff --git a/.github/workflows/build-docs.yml b/.github/workflows/build-docs.yml new file mode 100644 index 00000000000..82aacad44d1 --- /dev/null +++ b/.github/workflows/build-docs.yml @@ -0,0 +1,65 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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: Build docs + +on: + push: + branches: + - main + - "pull-request/[0-9]+" + - "deploy-release/*" + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.event.label.name || 'main' }}-${{ github.event_name }} + cancel-in-progress: true + +jobs: + pre-flight: + uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2 + + build-docs: + needs: [pre-flight] + if: needs.pre-flight.outputs.is_deployment_workflow != 'true' + uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_build_docs.yml@v0.80.2 + + build-docs-summary: + needs: [pre-flight, build-docs] + if: | + ( + needs.pre-flight.outputs.is_deployment_workflow == 'true' + || always() + ) + && !cancelled() + runs-on: ubuntu-latest + steps: + - name: Get workflow result + id: result + shell: bash -x -e -u -o pipefail {0} + env: + GH_TOKEN: ${{ github.token }} + RUN_ID: ${{ github.run_id }} + SKIPPING_IS_ALLOWED: ${{ needs.pre-flight.outputs.docs_only == 'true' || needs.pre-flight.outputs.is_deployment_workflow == 'true' }} + run: | + FAILED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion != "success")] | length') || echo 0 + + if [ "${FAILED_JOBS:-0}" -eq 0 ] || [ "$SKIPPING_IS_ALLOWED" == "true" ]; then + echo "✅ All previous jobs completed successfully" + exit 0 + else + echo "❌ Found $FAILED_JOBS failed job(s)" + # Show which jobs failed + gh run view $GITHUB_RUN_ID --json jobs --jq '.jobs[] | select(.status == "completed" and .conclusion != "success") | .name' + exit 1 + fi diff --git a/.github/workflows/build-test-publish-wheel.yml b/.github/workflows/build-test-publish-wheel.yml index 2b2ea3dfc85..b26036ca939 100644 --- a/.github/workflows/build-test-publish-wheel.yml +++ b/.github/workflows/build-test-publish-wheel.yml @@ -18,8 +18,8 @@ on: push: branches: - main - - 'pull-request/[0-9]+' - - 'deploy-release/*' + - "pull-request/[0-9]+" + - "deploy-release/*" merge_group: types: [checks_requested] @@ -33,7 +33,7 @@ permissions: jobs: pre-flight: - uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.65.5 + uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2 if: github.repository == 'NVIDIA/Megatron-LM' build-test-publish-wheels: @@ -42,8 +42,7 @@ jobs: with: no-publish: true secrets: - TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }} - TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }} + TWINE_PASSWORD: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/r')) && secrets.SVC_PYPI_TOKEN || secrets.SVC_PYPI_TEST_TOKEN }} build-test-publish-wheel-summary: needs: [pre-flight, build-test-publish-wheels] @@ -59,22 +58,22 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Result env: GH_TOKEN: ${{ github.token }} GITHUB_RUN_ID: ${{ github.run_id }} - SKIPPING_IS_ALLOWED: ${{ needs.pre-flight.outputs.docs_only == 'true' || needs.pre-flight.outputs.is_deployment_workflow == 'true' || needs.pre-flight.outputs.is_merge_group == 'true' || needs.pre-flight.outputs.is_ci_workload == 'true' || github.ref != 'refs/heads/main' }} + SKIPPING_IS_ALLOWED: false run: | - FAILED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion != "success")] | length') || echo 0 + FAILED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion != "success" and (.name | test("build-and-test-wheels")))] | length') || echo 0 if [ "${FAILED_JOBS:-0}" -eq 0 ] || [ "$SKIPPING_IS_ALLOWED" == "true" ]; then - echo "✅ All previous jobs completed successfully" + echo "✅ All build-and-test-wheels jobs completed successfully" exit 0 else - echo "❌ Found $FAILED_JOBS failed job(s)" + echo "❌ Found $FAILED_JOBS failed build-and-test-wheels job(s)" # Show which jobs failed - gh run view $GITHUB_RUN_ID --json jobs --jq '.jobs[] | select(.status == "completed" and .conclusion != "success") | .name' + gh run view $GITHUB_RUN_ID --json jobs --jq '.jobs[] | select(.status == "completed" and .conclusion != "success" and (.name | test("build-and-test-wheels"))) | .name' exit 1 fi 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/cherry-pick-release-commit.yml b/.github/workflows/cherry-pick-release-commit.yml index 882b3f5b268..9da305f07e6 100644 --- a/.github/workflows/cherry-pick-release-commit.yml +++ b/.github/workflows/cherry-pick-release-commit.yml @@ -26,5 +26,5 @@ jobs: target-branches-pattern: 'core_(*dev_)?r[0-9]+\.[0-9]+\.[0-9]+' secrets: PAT: ${{ secrets.PAT }} - SLACK_WEBHOOK_ADMIN: ${{ secrets.SLACK_WEBHOOK_ADMIN }} - SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK }} + SLACK_WEBHOOK_ADMIN: ${{ secrets.SLACK_TEAM_GROUP_ID }} + SLACK_WEBHOOK: ${{ secrets.SLACK_CI_CHANNEL_WEBHOOK }} diff --git a/.github/workflows/cicd-approve-test-queue.yml b/.github/workflows/cicd-approve-test-queue.yml index f34657eb509..cfd94f02a7d 100644 --- a/.github/workflows/cicd-approve-test-queue.yml +++ b/.github/workflows/cicd-approve-test-queue.yml @@ -27,12 +27,13 @@ jobs: strategy: matrix: branch: [main, dev, others] + contributor_type: [internal, external] steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: "3.12" @@ -41,23 +42,45 @@ jobs: python -m pip install --upgrade pip pip install requests + - name: Download SSO users list + run: | + gh release download v0.1.0 \ + --repo NVIDIA-GitHub-Management/github-audits \ + --pattern users_sso.json \ + --output users_sso.json || echo '{}' > users_sso.json + env: + GH_TOKEN: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }} + - name: Approve waiting deployments env: GITHUB_TOKEN: ${{ secrets.PAT }} MAX_CONCURRENCY: ${{ vars.MAX_CONCURRENCY || 1 }} + MAX_CONCURRENCY_EXTERNAL: ${{ vars.MAX_CONCURRENCY_EXTERNAL || 1 }} + CONTRIBUTOR_TYPE: ${{ matrix.contributor_type }} + SSO_USERS_FILE: users_sso.json PYTHONUNBUFFERED: 1 shell: python run: | import os + import json import requests import re # GitHub API configuration GITHUB_TOKEN = os.environ["GITHUB_TOKEN"] REPO = os.environ["GITHUB_REPOSITORY"] - MAX_CONCURRENCY = int(os.environ["MAX_CONCURRENCY"]) // 2 + CONTRIBUTOR_TYPE = os.environ["CONTRIBUTOR_TYPE"] + if CONTRIBUTOR_TYPE == "external": + # Global limit across all branches — no division needed since we count globally. + MAX_CONCURRENCY = int(os.environ["MAX_CONCURRENCY_EXTERNAL"]) + else: + MAX_CONCURRENCY = int(os.environ["MAX_CONCURRENCY"]) // 2 API_BASE = f"https://api.github.com/repos/NVIDIA/Megatron-LM" + # Load SSO users for internal/external classification + with open(os.environ["SSO_USERS_FILE"]) as f: + sso_users = json.load(f) + # Headers for GitHub API headers = { "Authorization": f"token {GITHUB_TOKEN}", @@ -81,53 +104,92 @@ jobs: print(f"Response: {e.response.text}") return None - def is_pr_targeting_branch(workflow_run, target_branch): + def is_internal_contributor(pr_info): + """Return True if the PR author is a member of NVIDIA or NVIDIA-NeMo org (is_org_member).""" + login = pr_info.get("user", {}).get("login", "") + org_roles = sso_users.get(login, {}).get("org_roles", []) + return any(role in ("NVIDIA:Member", "NVIDIA-NeMo:Member") for role in org_roles) + + def get_pr_base_branch(workflow_run): """ - Check if a workflow run belongs to a PR targeting the given branch. - Extract PR number from head branch like 'pull-request/1913' and verify base branch. + Return the base branch of the PR associated with a workflow run, or None. + Extracts PR number from head branch like 'pull-request/1913' and fetches PR info. + Returns (base_branch, pr_info) tuple, or (None, None) if not a PR run. """ print(workflow_run.get("head_branch", "")) head_branch = workflow_run.get("head_branch", "") match = re.match(r"pull-request/(\d+)", head_branch) if not match: - return False # Not a PR branch pattern + return None, None # Not a PR branch pattern pr_number = int(match.group(1)) - + # Fetch PR info from GitHub API pr_info = make_request(f"pulls/{pr_number}") if not pr_info: print(f"Failed to fetch PR #{pr_number}") - return False + return None, None base_branch = pr_info.get("base", {}).get("ref") - if ( - (base_branch == target_branch) or - (base_branch != "main" and base_branch != "dev" and target_branch == "others") - ): - print(f"PR #{pr_number} targets {target_branch}") - return True + return base_branch, pr_info + + def matches_contributor(workflow_run, contributor_type): + """Return True if the workflow run matches the contributor type (ignores branch).""" + _, pr_info = get_pr_base_branch(workflow_run) + if pr_info is None: + return False + internal = is_internal_contributor(pr_info) + return (contributor_type == "internal") == internal + + def matches_queue(workflow_run, target_branch, contributor_type): + """ + Return True if the workflow run belongs to this queue cell: + matching target branch AND matching contributor type (internal/external). + """ + base_branch, pr_info = get_pr_base_branch(workflow_run) + if base_branch is None: + return False + + branch_match = ( + (base_branch == target_branch) or + (base_branch != "main" and base_branch != "dev" and target_branch == "others") + ) + if not branch_match: + return False - return False + pr_number = re.match(r"pull-request/(\d+)", workflow_run.get("head_branch", "")).group(1) + internal = is_internal_contributor(pr_info) + contributor_match = (contributor_type == "internal") == internal + if branch_match and contributor_match: + print(f"PR #{pr_number} targets {target_branch}, contributor_type={contributor_type} (internal={internal})") + return branch_match and contributor_match # Get current running and queued workflows print("Fetching workflow runs...") queued_workflow_runs = make_request("actions/runs?status=queued").get("workflow_runs", []) in_progress_workflow_runs = make_request("actions/runs?status=in_progress").get("workflow_runs", []) - # Filter for workflows belonging to PRs targeting ${{ matrix.branch }} - queued_workflow_runs = [run for run in queued_workflow_runs - if run["name"] == "CICD Megatron-LM" and is_pr_targeting_branch(run, "${{ matrix.branch }}")] - in_progress_workflow_runs = [run for run in in_progress_workflow_runs - if run["name"] == "CICD Megatron-LM" and is_pr_targeting_branch(run, "${{ matrix.branch }}")] + # For external contributors, enforce a single global concurrency limit across ALL branches. + # For internal contributors, enforce per-branch limits as before. + if CONTRIBUTOR_TYPE == "external": + queued_workflow_runs = [run for run in queued_workflow_runs + if run["name"] == "CICD Megatron-LM" and matches_contributor(run, CONTRIBUTOR_TYPE)] + in_progress_workflow_runs = [run for run in in_progress_workflow_runs + if run["name"] == "CICD Megatron-LM" and matches_contributor(run, CONTRIBUTOR_TYPE)] + else: + # Filter for workflows belonging to PRs targeting ${{ matrix.branch }} with matching contributor type + queued_workflow_runs = [run for run in queued_workflow_runs + if run["name"] == "CICD Megatron-LM" and matches_queue(run, "${{ matrix.branch }}", CONTRIBUTOR_TYPE)] + in_progress_workflow_runs = [run for run in in_progress_workflow_runs + if run["name"] == "CICD Megatron-LM" and matches_queue(run, "${{ matrix.branch }}", CONTRIBUTOR_TYPE)] # Count running and queued workflows queued_workflows = len(queued_workflow_runs) in_progress_workflows = len(in_progress_workflow_runs) total_workflows = queued_workflows + in_progress_workflows - print(f"Current queued workflows (PRs targeting ${{ matrix.branch }}): {queued_workflows}") - print(f"Current running workflows (PRs targeting ${{ matrix.branch }}): {in_progress_workflows}") + print(f"Current queued workflows (PRs targeting ${{ matrix.branch }}, {CONTRIBUTOR_TYPE}): {queued_workflows}") + print(f"Current running workflows (PRs targeting ${{ matrix.branch }}, {CONTRIBUTOR_TYPE}): {in_progress_workflows}") print(f"Total workflows: {total_workflows}") print(f"Max concurrency: {MAX_CONCURRENCY}") @@ -139,8 +201,8 @@ jobs: print("Fetching deployments...") pending_workflows = make_request("actions/runs?status=waiting").get("workflow_runs", []) print("Pending workflows:", len(pending_workflows)) - pending_workflows = [run for run in pending_workflows - if run["name"] == "CICD Megatron-LM" and is_pr_targeting_branch(run, "${{ matrix.branch }}")] + pending_workflows = [run for run in pending_workflows + if run["name"] == "CICD Megatron-LM" and matches_queue(run, "${{ matrix.branch }}", CONTRIBUTOR_TYPE)] # Sort deployments by creation date (oldest first) print("Sorting workflows...") @@ -181,8 +243,8 @@ jobs: steps: - name: Notify env: - SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK }} - SLACK_WEBHOOK_ADMIN: + SLACK_WEBHOOK: ${{ secrets.SLACK_CI_CHANNEL_WEBHOOK }} + SLACK_WEBHOOK_ADMIN: GITHUB_RUN_ID: ${{ github.run_id }} GITHUB_REPOSITORY: ${{ github.repository }} run: | diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 16e2051e4e2..0d2d5b9577e 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.event.merge_group.head_ref || github.ref }} cancel-in-progress: true permissions: @@ -36,6 +34,7 @@ permissions: env: container-registry: 766267172432.dkr.ecr.us-east-1.amazonaws.com + container-registry-gb200: us-east4-docker.pkg.dev/nv-projdgxchipp-20260113193621/megatron-lm jobs: is-not-external-contributor: @@ -45,6 +44,7 @@ jobs: is_external_contributor: ${{ github.event.pull_request.user.type == 'User' }} is_maintainer: ${{ steps.check-membership.outputs.is_maintainer }} selected_runner: ${{ steps.check-membership.outputs.is_maintainer == 'true' && 'nvidia-ci-aws-gpu-x8' || 'nvidia-ci-aws-gpu-x8-ephemeral' }} + selected_runner_gb200: ${{ steps.check-membership.outputs.is_maintainer == 'true' && 'nvidia-ci-gcp-gpu-x4' || 'ubuntu-latest' }} permissions: issues: write pull-requests: write @@ -54,13 +54,13 @@ jobs: DISABLE_EXTERNAL_CONTRIBUTOR: ${{ vars.DISABLE_EXTERNAL_CONTRIBUTOR }} steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: token: ${{ env.GITHUB_TOKEN }} - name: Get PR info id: get-pr-info - if: startsWith(github.ref, 'refs/heads/pull-request/') + if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push' uses: nv-gha-runners/get-pr-info@main - name: Check NVIDIA SSO membership @@ -69,6 +69,7 @@ jobs: with: username: ${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').user.login }} github_token: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }} + sso_users_filename: ${{ vars.SSO_USERS_FILENAME }} - name: Set maintainer status id: check-membership @@ -131,7 +132,114 @@ jobs: pre-flight: needs: [is-not-external-contributor] if: github.repository == 'NVIDIA/Megatron-LM' - uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.65.10 + uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2 + + configure: + runs-on: ubuntu-latest + needs: [pre-flight] + if: github.repository == 'NVIDIA/Megatron-LM' + outputs: + scope: ${{ steps.configure.outputs.scope }} + n_repeat: ${{ steps.configure.outputs.n_repeat }} + lightweight: ${{ steps.configure.outputs.lightweight }} + lts: ${{ steps.configure.outputs.lts }} + mbridge_suite: ${{ steps.configure.outputs.mbridge_suite }} + dev: ${{ steps.configure.outputs.dev }} + steps: + - name: Get PR info + id: get-pr-info + if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push' + uses: nv-gha-runners/get-pr-info@main + + - name: Configure + id: configure + shell: bash -x -e -u -o pipefail {0} + env: + GH_TOKEN: ${{ secrets.PAT }} + IS_CI_WORKLOAD: ${{ needs.pre-flight.outputs.is_ci_workload }} + IS_MERGE_GROUP: ${{ needs.pre-flight.outputs.is_merge_group }} + run: | + PR_NUMBER=${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number }} + + # Fetch all labels in a single API call; fall back to empty list if no PR + LABELS=$(gh pr view $PR_NUMBER --repo ${{ github.repository }} --json labels --jq '[.labels[].name]') || LABELS='[]' + + HAS_RUN_TESTS=$(echo "$LABELS" | jq 'any(. == "Run tests")') + HAS_RUN_FUNCTIONAL=$(echo "$LABELS" | jq 'any(. == "Run functional tests")') + HAS_LTS=$(echo "$LABELS" | jq 'any(. == "container::lts")') + HAS_MBRIDGE=$(echo "$LABELS" | jq 'any(. == "Run MBridge tests")') + + # Scheduled/CI workloads have no PR — treat as "Run functional tests" + [ "$IS_CI_WORKLOAD" == "true" ] && HAS_RUN_FUNCTIONAL=true + + if [ "$IS_MERGE_GROUP" == "true" ]; then + SCOPE=mr-github; N_REPEAT=1; LIGHTWEIGHT=false + elif [ "$HAS_RUN_TESTS" == "true" ]; then + SCOPE=mr-github; N_REPEAT=1; LIGHTWEIGHT=true + elif [ "$HAS_RUN_FUNCTIONAL" == "true" ]; then + SCOPE=mr-github; N_REPEAT=5; LIGHTWEIGHT=false + else + SCOPE=mr-github-slim; N_REPEAT=5; LIGHTWEIGHT=false + fi + + if [ "$HAS_MBRIDGE" == "true" || $IS_MERGE_GROUP == "true" ]; then + MBRIDGE_SUITE="L1" + else + MBRIDGE_SUITE="unit-only" + fi + + DEV=true + + echo "scope=$SCOPE" | tee -a $GITHUB_OUTPUT + echo "n_repeat=$N_REPEAT" | tee -a $GITHUB_OUTPUT + echo "lightweight=$LIGHTWEIGHT" | tee -a $GITHUB_OUTPUT + echo "lts=$HAS_LTS" | tee -a $GITHUB_OUTPUT + echo "mbridge_suite=$MBRIDGE_SUITE" | tee -a $GITHUB_OUTPUT + echo "dev=$DEV" | tee -a $GITHUB_OUTPUT + + # Pre-compute active row markers for the decision tree + _MG=$( [ "$IS_MERGE_GROUP" == "true" ] && echo "**→**" || echo "" ) + _RT=$( [ "$IS_MERGE_GROUP" != "true" ] && [ "$HAS_RUN_TESTS" == "true" ] && echo "**→**" || echo "" ) + _RF=$( [ "$IS_MERGE_GROUP" != "true" ] && [ "$HAS_RUN_TESTS" != "true" ] && [ "$HAS_RUN_FUNCTIONAL" == "true" ] && echo "**→**" || echo "" ) + _DF=$( [ "$SCOPE" == "mr-github-slim" ] && echo "**→**" || echo "" ) + _LTS=$( [ "$HAS_LTS" == "true" ] && echo "**→**" || echo "" ) + _DEV=$( [ "$HAS_LTS" != "true" ] && echo "**→**" || echo "" ) + + cat <> $GITHUB_STEP_SUMMARY + Beep boop 🤖 I have consulted the labels and decided to run **$SCOPE** $( [ "$LIGHTWEIGHT" == "true" ] && echo "in lightweight mode " || echo "" )against the **$( [ "$HAS_LTS" == "true" ] && echo "lts" || echo "dev" )** container with **$N_REPEAT** repetition(s). You are welcome. + + | Setting | Value | + |---|---| + | \`scope\` | \`$SCOPE\` | + | \`n_repeat\` | \`$N_REPEAT\` | + | \`lightweight\` | \`$LIGHTWEIGHT\` | + | \`lts\` | \`$HAS_LTS\` | + | \`dev\` | \`$DEV\` | + | \`mbridge_suite\` | \`$MBRIDGE_SUITE\` | + + ### Decision tree + + **Test scope** + + | | Trigger | \`scope\` | \`n_repeat\` | \`lightweight\` | + |---|---|---|---|---| + | $_MG | Merge group | \`mr-github\` | \`1\` | \`false\` | + | $_RT | Label: _Run tests_ | \`mr-github\` | \`1\` | \`true\` | + | $_RF | Label: _Run functional tests_ / CI workload | \`mr-github\` | \`5\` | \`false\` | + | $_DF | _(default)_ | \`mr-github-slim\` | \`5\` | \`false\` | + + **Container image** + + | | Trigger | \`image\` | + |---|---|---| + | $_LTS | Label: _container::lts_ | \`lts\` | + | $_DEV | _(default)_ | \`dev\` | + + ### Glossary + - **\`lightweight\`**: trains for 4 steps instead of 100 and skips comparison against golden values — faster feedback, no correctness guarantees + - **\`lts\`**: uses the Long Term Support container base image instead of the latest dev image + - **\`dev\`**: uses the latest development container base image (default) + SUMMARY linting: runs-on: ubuntu-latest @@ -147,7 +255,7 @@ jobs: ) steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 @@ -162,11 +270,11 @@ jobs: - name: Get PR info id: get-pr-info - if: startsWith(github.ref, 'refs/heads/pull-request/') + if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push' uses: nv-gha-runners/get-pr-info@main - name: Run linting - if: startsWith(github.ref, 'refs/heads/pull-request/') + if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push' run: | export PATH=".venv/bin:$PATH" export GITLAB_ENDPOINT=github.com @@ -179,10 +287,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 @@ -190,21 +299,164 @@ jobs: echo "Running CI tests" echo "is_merge_group: ${{ needs.pre-flight.outputs.is_merge_group }}" + cicd-parse-downstream-testing: + runs-on: ubuntu-latest + needs: + - pre-flight + - configure + - cicd-wait-in-queue + if: | + needs.pre-flight.result != 'cancelled' + && needs.configure.result != 'cancelled' + && needs.cicd-wait-in-queue.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' + ) + && !cancelled() + outputs: + mbridge-test-suite: ${{ needs.configure.outputs.mbridge_suite }} + steps: + - name: Checkout + uses: actions/checkout@v6 + + - name: How-To + run: bash .github/scripts/readme.sh + + cicd-mbridge-testing: + runs-on: ubuntu-latest + needs: + - pre-flight + - cicd-wait-in-queue + - cicd-parse-downstream-testing + if: | + needs.pre-flight.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && needs.cicd-parse-downstream-testing.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' + ) + && !cancelled() + steps: + - name: Get PR info + id: get-pr-info + if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push' + uses: nv-gha-runners/get-pr-info@main + + - name: Checkout MBridge and create testing branch + uses: actions/checkout@v6 + with: + ref: main + repository: NVIDIA-NeMo/Megatron-Bridge + path: megatron-bridge + token: ${{ secrets.PAT }} + + - name: Create testing branch + env: + MBRIDGE_BRANCH_NAME: mcore-testing-${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number || github.run_id }} + run: | + cd megatron-bridge + git fetch origin main + git checkout -b ${{ env.MBRIDGE_BRANCH_NAME }} origin/main + git push origin ${{ env.MBRIDGE_BRANCH_NAME }} --force + + - name: Get merge commit sha + shell: bash -x -e -u -o pipefail {0} + 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 + echo "main=${SHA}" | tee -a "$GITHUB_OUTPUT" + + - name: Trigger MBridge tests + uses: convictional/trigger-workflow-and-wait@v1.6.5 + env: + MBRIDGE_BRANCH_NAME: mcore-testing-${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number || github.run_id }} + with: + owner: NVIDIA-NeMo + repo: Megatron-Bridge + workflow_file_name: cicd-main.yml + github_token: ${{ secrets.PAT }} + ref: ${{ env.MBRIDGE_BRANCH_NAME }} + wait_interval: 60 + propagate_failure: true + client_payload: | + { + "mcore_ref": "${{ steps.sha.outputs.main }}", + "test_suite": "${{ needs.cicd-parse-downstream-testing.outputs.mbridge-test-suite }}", + "triggered_by": "https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}" + } + + - name: Delete testing branch + if: always() + env: + MBRIDGE_BRANCH_NAME: mcore-testing-${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number || github.run_id }} + run: | + cd megatron-bridge + git push origin --delete ${{ env.MBRIDGE_BRANCH_NAME }} + + cicd-compute-build-matrix: + runs-on: ubuntu-latest + needs: [is-not-external-contributor] + outputs: + matrix: ${{ steps.compute.outputs.matrix }} + steps: + - name: Compute build matrix + id: compute + env: + IS_MAINTAINER: ${{ needs.is-not-external-contributor.outputs.is_maintainer }} + SELECTED_RUNNER: ${{ needs.is-not-external-contributor.outputs.selected_runner }} + SELECTED_RUNNER_GB200: ${{ needs.is-not-external-contributor.outputs.selected_runner_gb200 }} + REGISTRY_AWS: ${{ env.container-registry }} + REGISTRY_GCP: ${{ env.container-registry-gb200 }} + run: | + AWS_ENTRY=$(jq -nc --arg registry "$REGISTRY_AWS" --arg runner "$SELECTED_RUNNER" \ + '{"cloud": "aws", "registry": $registry, "runner": $runner}') + if [ "$IS_MAINTAINER" == "true" ]; then + GCP_ENTRY=$(jq -nc --arg registry "$REGISTRY_GCP" --arg runner "$SELECTED_RUNNER_GB200" \ + '{"cloud": "gcp", "registry": $registry, "runner": $runner}') + MATRIX=$(jq -nc --argjson aws "$AWS_ENTRY" --argjson gcp "$GCP_ENTRY" \ + '{"include": [$aws, $gcp]}') + else + MATRIX=$(jq -nc --argjson aws "$AWS_ENTRY" '{"include": [$aws]}') + fi + echo "matrix=$MATRIX" | tee -a "$GITHUB_OUTPUT" + cicd-container-build: - needs: [is-not-external-contributor, pre-flight, cicd-wait-in-queue] - runs-on: ${{ needs.is-not-external-contributor.outputs.selected_runner }} + needs: [is-not-external-contributor, pre-flight, configure, cicd-wait-in-queue, cicd-compute-build-matrix] + strategy: + fail-fast: false + matrix: ${{ fromJson(needs.cicd-compute-build-matrix.outputs.matrix) }} + runs-on: ${{ matrix.runner }} if: | - ( + needs.is-not-external-contributor.result != 'cancelled' + && needs.pre-flight.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && needs.cicd-compute-build-matrix.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 id: get-pr-info - if: startsWith(github.ref, 'refs/heads/pull-request/') + if: startsWith(github.ref, 'refs/heads/pull-request/') && github.event_name == 'push' uses: nv-gha-runners/get-pr-info@main - name: Get merge commit sha @@ -212,21 +464,24 @@ 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 echo "main=${SHA}" | tee -a "$GITHUB_OUTPUT" - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: ref: ${{ steps.sha.outputs.main }} - name: Setup python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: 3.12 @@ -236,15 +491,6 @@ jobs: apt-get update apt-get install -y gh - - name: Has lts label - id: has-lts-label - env: - GH_TOKEN: ${{ secrets.PAT }} - run: | - PR_NUMBER=${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number }} - HAS_LTS_LABEL=$(gh pr view $PR_NUMBER --json labels | jq '[.labels[].name] | any(. == "container::lts")') || echo "false" - echo "main=$HAS_LTS_LABEL" | tee -a $GITHUB_OUTPUT - - name: Download test data shell: bash run: | @@ -274,7 +520,7 @@ jobs: } } }' | jq -r '.data.repository.pullRequests.nodes[].number' | while read -r number; do - echo "type=registry,ref=${{ env.container-registry }}/megatron-lm:$number-buildcache,mode=max" + echo "type=registry,ref=${{ matrix.registry }}/megatron-lm:$number-buildcache,mode=max" done) echo "LAST_PRS< integration-tests.json + cat integration-tests-h100.yaml | \ + yq -o json 'del(.default, .stages, .workflow) | to_entries | map({"model": .value.stage, "test_case": .key}) | sort_by(.model, .test_case)' | jq -c > integration-tests-h100.json - echo "integration-tests=$(cat integration-tests.json)" | tee -a "$GITHUB_OUTPUT" + echo "integration-tests-h100=$(cat integration-tests-h100.json)" | tee -a "$GITHUB_OUTPUT" - cicd-integration-tests-latest: + cicd-integration-tests-latest-h100: + timeout-minutes: 60 strategy: fail-fast: false matrix: - include: ${{ fromJson(needs.cicd-parse-integration-tests.outputs.integration-tests) }} + include: ${{ fromJson(needs.cicd-parse-integration-tests-h100.outputs.integration-tests-h100) }} needs: - is-not-external-contributor - pre-flight + - configure - cicd-wait-in-queue - - cicd-parse-integration-tests + - cicd-parse-integration-tests-h100 - cicd-unit-tests-latest runs-on: ${{ needs.is-not-external-contributor.outputs.selected_runner }} name: "${{ matrix.model }}/${{ matrix.test_case }} - latest" @@ -497,16 +720,22 @@ 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.configure.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && needs.cicd-parse-integration-tests-h100.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 - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: main uses: ./.github/actions with: @@ -517,13 +746,130 @@ jobs: is_unit_test: "false" PAT: ${{ secrets.PAT }} container-image: ${{ env.container-registry }}/megatron-lm:${{ github.sha }} - is_ci_workload: ${{ needs.pre-flight.outputs.is_ci_workload }} + scope: ${{ needs.configure.outputs.scope }} + n_repeat: ${{ needs.configure.outputs.n_repeat }} + lightweight: ${{ needs.configure.outputs.lightweight }} + + cicd-parse-integration-tests-gb200: + runs-on: ubuntu-latest + needs: + - is-not-external-contributor + - pre-flight + - configure + - cicd-wait-in-queue + - cicd-container-build + - cicd-unit-tests-latest + if: | + needs.is-not-external-contributor.outputs.is_maintainer == 'true' + && needs.pre-flight.result != 'cancelled' + && needs.configure.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' + || needs.pre-flight.outputs.is_merge_group == 'true' + ) + && !cancelled() + outputs: + integration-tests-gb200: ${{ steps.main.outputs.integration-tests-gb200 }} + steps: + - name: Checkout + uses: actions/checkout@v6 + + - name: Parse functional tests + id: main + env: + SCOPE: ${{ needs.configure.outputs.scope }} + LIGHTWEIGHT: ${{ needs.configure.outputs.lightweight }} + run: | + export PYTHONPATH=$(pwd) + + ARGS=(--scope $SCOPE) + [ "$LIGHTWEIGHT" == "true" ] && ARGS+=(--enable-lightweight-mode) + + python tests/test_utils/python_scripts/generate_jet_trigger_job.py \ + --n-repeat 5 \ + --time-limit 2700 \ + --test-cases all \ + --container-image mcore_ci_dev \ + --container-tag latest \ + --dependent-job functional:configure \ + --record-checkpoints false \ + --slurm-account gh \ + --no-enable-warmup \ + --environment dev \ + --platform dgx_gb200 \ + --cluster dgxgb200_oci-hsg \ + ${ARGS[@]} \ + --output-path integration-tests-gb200.yaml + + cat integration-tests-gb200.yaml | \ + yq -o json 'del(.default, .stages, .workflow) | to_entries | map({"model": .value.stage, "test_case": .key}) | sort_by(.model, .test_case)' | jq -c > integration-tests-gb200.json + + echo "integration-tests-gb200=$(cat integration-tests-gb200.json)" | tee -a "$GITHUB_OUTPUT" + + cicd-integration-tests-latest-gb200: + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + include: ${{ fromJson(needs.cicd-parse-integration-tests-gb200.outputs.integration-tests-gb200) }} + needs: + - is-not-external-contributor + - pre-flight + - configure + - cicd-wait-in-queue + - cicd-parse-integration-tests-gb200 + - cicd-unit-tests-latest + runs-on: ${{ needs.is-not-external-contributor.outputs.selected_runner_gb200 }} + name: "${{ matrix.model }}/${{ matrix.test_case }} - latest" + env: + PIP_DISABLE_PIP_VERSION_CHECK: 1 + PIP_NO_PYTHON_VERSION_WARNING: 1 + PIP_ROOT_USER_ACTION: ignore + if: | + needs.is-not-external-contributor.outputs.is_maintainer == 'true' + && needs.is-not-external-contributor.result != 'cancelled' + && needs.pre-flight.result != 'cancelled' + && needs.configure.result != 'cancelled' + && needs.cicd-wait-in-queue.result != 'cancelled' + && needs.cicd-parse-integration-tests-gb200.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' + ) + && !cancelled() + steps: + - name: Checkout + uses: actions/checkout@v6 + - name: main + uses: ./.github/actions + with: + test_case: ${{ matrix.test_case }} + model: ${{ matrix.model }} + tag: latest + timeout: ${{ matrix.timeout || 30 }} + is_unit_test: "false" + PAT: ${{ secrets.PAT }} + container-image: ${{ env.container-registry-gb200 }}/megatron-lm:${{ github.sha }} + scope: ${{ needs.configure.outputs.scope }} + n_repeat: ${{ needs.configure.outputs.n_repeat }} + lightweight: ${{ needs.configure.outputs.lightweight }} + platform: dgx_gb200 Nemo_CICD_Test: needs: - pre-flight + - is-not-external-contributor - cicd-unit-tests-latest - - cicd-integration-tests-latest + - cicd-integration-tests-latest-h100 + - cicd-integration-tests-latest-gb200 if: | ( needs.pre-flight.outputs.docs_only == 'true' @@ -538,7 +884,7 @@ jobs: permissions: write-all steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Get workflow result id: result @@ -546,19 +892,71 @@ jobs: env: GH_TOKEN: ${{ github.token }} GITHUB_RUN_ID: ${{ github.run_id }} - SKIPPING_IS_ALLOWED: ${{ needs.pre-flight.outputs.docs_only == 'true' || needs.pre-flight.outputs.is_deployment_workflow == 'true' || needs.pre-flight.outputs.is_merge_group == 'true' || needs.pre-flight.outputs.is_ci_workload == 'true' }} + DOCS_ONLY: ${{ needs.pre-flight.outputs.docs_only }} + IS_DEPLOYMENT: ${{ needs.pre-flight.outputs.is_deployment_workflow }} + IS_MAINTAINER: ${{ needs.is-not-external-contributor.outputs.is_maintainer }} + UNIT_RESULT: ${{ needs.cicd-unit-tests-latest.result }} + H100_RESULT: ${{ needs.cicd-integration-tests-latest-h100.result }} + GB200_RESULT: ${{ needs.cicd-integration-tests-latest-gb200.result }} 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 + # Docs-only and deployment workflows intentionally skip all tests + if [ "$DOCS_ONLY" == "true" ] || [ "$IS_DEPLOYMENT" == "true" ]; then + echo "✅ Docs-only or deployment workflow — test checks skipped" + exit 0 + fi + + FAILED=false + + # Unit tests must always succeed (never skipped or cancelled) + if [ "$UNIT_RESULT" != "success" ]; then + echo "❌ cicd-unit-tests-latest: $UNIT_RESULT" + FAILED=true + fi + + # H100 integration tests must always succeed + if [ "$H100_RESULT" != "success" ]; then + echo "❌ cicd-integration-tests-latest-h100: $H100_RESULT" + FAILED=true + fi + + # GB200 integration tests may be skipped only for non-maintainer PRs + # (no GB200 runners available); maintainer runs must always succeed + if [ "$GB200_RESULT" == "skipped" ] && [ "$IS_MAINTAINER" == "true" ]; then + echo "❌ cicd-integration-tests-latest-gb200: skipped unexpectedly for a maintainer run" + FAILED=true + elif [ "$GB200_RESULT" != "success" ] && [ "$GB200_RESULT" != "skipped" ]; then + echo "❌ cicd-integration-tests-latest-gb200: $GB200_RESULT" + FAILED=true + fi + + # Broad scan: catch any individual job failures or cancellations + # (e.g. a single matrix instance cancelled mid-run) + BAD_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq ' + [.jobs[] | select( + .status == "completed" + and (.conclusion == "failure" or .conclusion == "cancelled") + and .name != "merge-queue-notification" + and .name != "cicd-mbridge-testing" + )] | length + ') || BAD_JOBS=0 + + if [ "${BAD_JOBS:-0}" -gt 0 ]; then + echo "❌ Found ${BAD_JOBS} failed or cancelled job(s):" + gh run view $GITHUB_RUN_ID --json jobs --jq ' + .jobs[] | select( + .status == "completed" + and (.conclusion == "failure" or .conclusion == "cancelled") + and .name != "merge-queue-notification" + and .name != "cicd-mbridge-testing" + ) | .name + " → " + .conclusion + ' + FAILED=true + fi - if [ "${FAILED_JOBS:-0}" -eq 0 ] && ([ "${SKIPPED_JOBS:-0}" -eq 0 ] || [ "$SKIPPING_IS_ALLOWED" == "true" ]); then - echo "✅ All previous jobs completed successfully" - exit 0 + if [ "$FAILED" != "true" ]; then + echo "✅ All previous jobs completed successfully" else - echo "❌ Found $FAILED_JOBS failed job(s)" - # Show which jobs failed - gh run view $GITHUB_RUN_ID --json jobs --jq '.jobs[] | select(.status == "completed" and .conclusion == "failure") | .name' - exit 1 + exit 1 fi Coverage_Fake: @@ -575,7 +973,7 @@ jobs: && github.repository == 'NVIDIA/Megatron-LM' steps: - name: Generate fake coverage report - uses: actions/github-script@v6 + uses: actions/github-script@v8 with: github-token: ${{ secrets.PAT }} script: | @@ -603,10 +1001,10 @@ jobs: flag: [unit-test] steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Download coverage reports of current branch - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v7 with: pattern: coverage-${{ matrix.flag }}-* @@ -634,26 +1032,56 @@ jobs: flags: ${{ matrix.flag }} - name: Upload artifacts - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: coverage-${{ matrix.flag }}-aggregated path: | .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@v8 + 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: - is-not-external-contributor - cicd-container-build - cicd-unit-tests-latest - - cicd-integration-tests-latest + - cicd-integration-tests-latest-h100 + - cicd-integration-tests-latest-gb200 - Coverage - Coverage_Fake if: | always() && !cancelled() && contains(needs.is-not-external-contributor.outputs.selected_runner, 'ephemeral') + && !needs.pre-flight.outputs.is_deployment_workflow == 'true' steps: - name: Taint node for cleanup shell: bash diff --git a/.github/workflows/claude-complexity-label.yml b/.github/workflows/claude-complexity-label.yml new file mode 100644 index 00000000000..356eed2da29 --- /dev/null +++ b/.github/workflows/claude-complexity-label.yml @@ -0,0 +1,60 @@ +name: Claude Complexity Label + +on: + pull_request_target: + types: [ready_for_review] + +jobs: + label-complexity: + name: Label PR Complexity + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: write + issues: write + id-token: write + env: + GH_TOKEN: ${{ secrets.PAT }} + REPO: ${{ github.repository }} + PR_NUMBER: ${{ github.event.pull_request.number }} + steps: + - name: Checkout repository + uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: Run Claude Complexity Analysis + uses: anthropics/claude-code-action@v1 + with: + anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }} + github_token: ${{ secrets.PAT }} + prompt: | + REPO: ${{ env.REPO }} + PR NUMBER: ${{ env.PR_NUMBER }} + + You are a PR complexity analyzer. Your job is to analyze the diff of this PR and apply exactly one complexity label. + + STEPS: + 1. Get the PR diff by running: gh pr diff $PR_NUMBER --repo $REPO + 2. Analyze every changed line (added or removed) in the diff and classify each as one of: + - "docs-only": changes to docstrings, comments (lines starting with # or //), documentation files (.md, .rst, .txt), or similar non-functional text + - "test": changes in test files (files with "test" in the name/path, or inside a tests/ directory) + - "real code": all other changes (functional source code) + 3. Compute "real code line changes" using this formula: + real_code_line_changes = (number of real code lines changed) + (number of test lines changed / 10) + Count both added and removed lines. Do not count unchanged context lines. Do not count comments or docstrings. + 4. Remove any previously applied complexity or docs-only labels: + gh pr edit $PR_NUMBER --repo $REPO --remove-label "complexity: low,complexity: medium,complexity: high,docs-only" + 5. Apply exactly ONE label using the gh CLI: + - If there are ZERO real code lines and ZERO test lines (only docs-only changes), apply label "docs-only": + gh pr edit $PR_NUMBER --repo $REPO --add-label "docs-only" + - If real_code_line_changes < 100, apply label "complexity: low": + gh pr edit $PR_NUMBER --repo $REPO --add-label "complexity: low" + - If real_code_line_changes >= 100 and < 500, apply label "complexity: medium": + gh pr edit $PR_NUMBER --repo $REPO --add-label "complexity: medium" + - If real_code_line_changes >= 500, apply label "complexity: high": + gh pr edit $PR_NUMBER --repo $REPO --add-label "complexity: high" + + Do NOT post any comments on the PR. Only apply the label. + claude_args: | + --allowedTools "Bash(gh pr diff:*),Bash(gh pr edit:*),Bash(gh pr view:*)" diff --git a/.github/workflows/claude-copy-to-main.yml b/.github/workflows/claude-copy-to-main.yml new file mode 100644 index 00000000000..7bde3941bb8 --- /dev/null +++ b/.github/workflows/claude-copy-to-main.yml @@ -0,0 +1,122 @@ +name: Claude Copy PR to Main + +on: + issue_comment: + types: [created] + +jobs: + copy-to-main: + name: Copy PR to Main + if: | + github.event_name == 'issue_comment' && + github.event.issue.pull_request && + contains(github.event.comment.body, '/claude copy') + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + issues: write + id-token: write + env: + GH_TOKEN: ${{ secrets.PAT }} + REPO: ${{ github.repository }} + PR_NUMBER: ${{ github.event.issue.number }} + steps: + - name: Check commenter has write access + env: + COMMENTER: ${{ github.event.comment.user.login }} + run: | + PERMISSION=$(gh api repos/$REPO/collaborators/$COMMENTER/permission --jq .permission) + if [[ "$PERMISSION" != "admin" && "$PERMISSION" != "write" ]]; then + gh pr comment $PR_NUMBER --repo $REPO --body "❌ You do not have write access to use \`/claude copy\`." + exit 1 + fi + + - name: Check PR is merged and targets non-main + run: | + PR_JSON=$(gh pr view $PR_NUMBER --repo $REPO --json baseRefName,mergedAt) + PR_BASE=$(echo "$PR_JSON" | jq -r .baseRefName) + PR_MERGED=$(echo "$PR_JSON" | jq -r .mergedAt) + + if [ "$PR_BASE" = "main" ]; then + gh pr comment $PR_NUMBER --repo $REPO --body "❌ This PR already targets \`main\`. \`/claude copy\` only works on PRs targeting non-main branches." + exit 1 + fi + + if [ "$PR_MERGED" = "null" ] || [ -z "$PR_MERGED" ]; then + gh pr comment $PR_NUMBER --repo $REPO --body "❌ This PR has not been merged yet. \`/claude copy\` only works on merged PRs." + exit 1 + fi + + - name: Checkout repository + uses: actions/checkout@v6 + with: + fetch-depth: 0 + token: ${{ secrets.PAT }} + + - name: Fetch PR head ref from fork + run: | + git fetch origin pull/$PR_NUMBER/head:pr-$PR_NUMBER-head + + - name: Run Claude Copy to Main + uses: anthropics/claude-code-action@v1 + with: + anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }} + trigger_phrase: "/claude copy" + github_token: ${{ secrets.PAT }} + prompt: | + REPO: ${{ env.REPO }} + PR NUMBER: ${{ env.PR_NUMBER }} + + You are a PR copy assistant. Your job is to apply the final changes from a merged PR onto a new branch based on `main` and create a new PR targeting `main`. + + The PR's commits originated from a fork and have been fetched locally as the branch: pr-${PR_NUMBER}-head + + STEPS: + 1. Get the PR details (title, body, and base branch): + gh pr view $PR_NUMBER --repo $REPO --json title,body,baseRefName + + 2. Configure git for committing (use the svcnvidia-nemo-ci service account since secrets.PAT belongs to it): + git config user.name "svcnvidia-nemo-ci" + git config user.email "svcnvidia-nemo-ci@nvidia.com" + + 3. Create a new branch from `main`: + git checkout main + git pull origin main + git checkout -b copy-pr-${PR_NUMBER}-to-main + + 4. Generate a patch of the PR's final changes and apply it: + MERGE_BASE=$(git merge-base origin/ pr-${PR_NUMBER}-head) + git diff $MERGE_BASE pr-${PR_NUMBER}-head | git apply --3way + (Replace with the actual base branch name from step 1.) + + If the apply fails due to merge conflicts: + a. Identify conflicted files: git diff --name-only --diff-filter=U + b. For each conflicted file, read its contents to see the conflict markers + c. Resolve the conflicts by favoring the `main` branch side when there is a genuine + conflict between the two sides. The goal is to bring the PR's changes into main + without overriding what is already on main. + d. Stage the resolved files: git add + + 5. Commit the changes: + git add -A + git commit -m "Copy PR #${PR_NUMBER} to main" + + 6. Push the new branch: + git push origin copy-pr-${PR_NUMBER}-to-main + + 7. Create a new PR targeting `main`: + gh pr create --repo $REPO \ + --base main \ + --head copy-pr-${PR_NUMBER}-to-main \ + --title "[Copy to main] " \ + --body "🤖 **This PR was auto-generated by Claude** via the \`/claude copy\` command.\n\nCherry-picked from #${PR_NUMBER}.\n\n---\n\n" + + 8. Comment on the original PR with a link to the newly created PR. + + IMPORTANT: + - When resolving merge conflicts, favor `main` over the non-main branch. Do not override changes already on main. + - Do NOT force push. + claude_args: | + --allowedTools "Bash(git:*),Bash(gh:*),Read,Edit" + --model "claude-opus-4-6" diff --git a/.github/workflows/claude_review.yml b/.github/workflows/claude_review.yml new file mode 100644 index 00000000000..da182a9dc91 --- /dev/null +++ b/.github/workflows/claude_review.yml @@ -0,0 +1,71 @@ +name: Claude Code Review + +on: + issue_comment: + types: [created] + +jobs: + review-on-comment: + name: Claude Review (comment trigger) + if: | + github.event_name == 'issue_comment' && + github.event.issue.pull_request && + contains(github.event.comment.body, '/claude review') + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: write + issues: write + id-token: write + env: + GH_TOKEN: ${{ github.token }} + REPO: ${{ github.repository }} + PR_NUMBER: ${{ github.event.issue.number }} + steps: + - name: Get PR head commit + id: get-pr-head-commit + run: | + echo "sha=$(gh pr view $PR_NUMBER --repo $REPO --json headRefOid -q .headRefOid)" | tee -a $GITHUB_OUTPUT + + - name: Checkout repository + uses: actions/checkout@v6 + with: + fetch-depth: 1 + ref: ${{ steps.get-pr-head-commit.outputs.sha }} + + - name: Run Claude Code Review + uses: anthropics/claude-code-action@v1 + with: + anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }} + trigger_phrase: "/claude review" + show_full_output: true + claude_args: | + --allowedTools "mcp__github_inline_comment__create_inline_comment,Bash(gh pr comment:*),Bash(gh pr diff:*),Bash(gh pr view:*),Bash(gh pr review:*)" + --model "claude-opus-4-6" + prompt: | + REPO: ${{ env.REPO }} + PR NUMBER: ${{ env.PR_NUMBER }} + + You are doing a light code review. Keep it concise and actionable. + + Focus ONLY on: + - Critical bugs or logic errors + - Typos in code, comments, or strings + - Missing or insufficient test coverage for changed code + - Outdated or inaccurate documentation affected by the changes + + Do NOT comment on: + - Style preferences or formatting + - Minor naming suggestions + - Architectural opinions or refactoring ideas + - Performance unless there is a clear, measurable issue + + Only use inline ```suggestion blocks for simple, self-contained line replacements (typos, + renames, single-line fixes). For structural changes that add, remove, or reorganize blocks + of code (e.g. adding a new function, inserting a YAML step, reordering logic), use a + top-level PR comment with a code block showing the proposed change instead — inline + suggestions cannot express insertions or multi-block restructuring and will break the code + if applied. + + It's perfectly acceptable to not have anything to comment on. + If you do not have anything to comment on, approve the PR with: gh pr review $PR_NUMBER --repo $REPO --approve --body "LGTM" diff --git a/.github/workflows/config/changelog-config.json b/.github/workflows/config/changelog-config.json new file mode 100644 index 00000000000..19fb0e42364 --- /dev/null +++ b/.github/workflows/config/changelog-config.json @@ -0,0 +1,24 @@ +{ + "categories": [], + "ignore_labels": [ + "ignore" + ], + "sort": "ASC", + "template": "\n${{CHANGELOG}}\n\n
Changelog Details\n\n${{UNCATEGORIZED}}\n
\n", + "pr_template": "- ${{TITLE}} by @${{AUTHOR}} :: PR: #${{NUMBER}}", + "commit_template": "- ${{TITLE}} by @${{AUTHOR}}", + "empty_template": "${{OWNER}}\n${{REPO}}\n${{FROM_TAG}}\n${{TO_TAG}}", + "duplicate_filter": { + "pattern": ".+", + "on_property": "title", + "method": "match" + }, + "transformers": [], + "max_tags_to_fetch": 100, + "max_pull_requests": 1250, + "max_back_track_time_days": 365, + "exclude_merge_branches": [], + "tag_resolver": { + "method": "semver" + } +} diff --git a/.github/workflows/copyright-check.yml b/.github/workflows/copyright-check.yml index 9bbb7a1f201..33d30944f8d 100644 --- a/.github/workflows/copyright-check.yml +++ b/.github/workflows/copyright-check.yml @@ -17,14 +17,14 @@ name: Copyright check on: push: branches: - - 'pull-request/[0-9]+' - - 'deploy-release/*' + - "pull-request/[0-9]+" + - "deploy-release/*" merge_group: types: [checks_requested] jobs: pre-flight: - uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.65.10 + uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2 if: github.repository == 'NVIDIA/Megatron-LM' copyright-check: @@ -49,7 +49,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Result env: diff --git a/.github/workflows/dependabot.yml b/.github/workflows/dependabot.yml index 9dc1e6ac5a9..81a5cd57d12 100644 --- a/.github/workflows/dependabot.yml +++ b/.github/workflows/dependabot.yml @@ -11,7 +11,6 @@ permissions: jobs: get-release-branch-names: runs-on: ubuntu-latest - environment: nemo-ci outputs: mcore: ${{ steps.get-branch.outputs.mcore_release_branch }} if: github.repository == 'NVIDIA/Megatron-LM' @@ -21,11 +20,11 @@ jobs: env: PAT: ${{ secrets.PAT }} run: | - latest_branch=$(git ls-remote --heads https://token:${PAT}@github.com/NVIDIA-NeMo/Eval.git 'refs/heads/r*' | + latest_branch=$(git ls-remote --heads https://token:${PAT}@github.com/NVIDIA/Megatron-LM.git 'refs/heads/core_r*' | grep -o 'core_r[0-9]\+\.[0-9]\+\.[0-9]\+' | sort -V | tail -n1) - echo "mcore_release_branch=$latest_branch" >> $GITHUB_OUTPUT + echo "mcore_release_branch=$latest_branch" | tee -a $GITHUB_OUTPUT bump-tags: needs: [get-release-branch-names] @@ -41,9 +40,6 @@ jobs: target-branch: ${{ matrix.target-branch }} secrets: PAT: ${{ secrets.PAT }} - AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }} - AZURE_TENANT_ID: ${{ secrets.AZURE_TENANT_ID }} - AZURE_SUBSCRIPTION_ID: ${{ secrets.AZURE_SUBSCRIPTION_ID }} SSH_KEY: ${{ secrets.SSH_KEY }} SSH_PWD: ${{ secrets.SSH_PWD }} @@ -54,8 +50,8 @@ jobs: steps: - name: Notify env: - SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK }} - SLACK_WEBHOOK_ADMIN: + SLACK_WEBHOOK: ${{ secrets.SLACK_CI_CHANNEL_WEBHOOK }} + SLACK_WEBHOOK_ADMIN: GITHUB_RUN_ID: ${{ github.run_id }} GITHUB_REPOSITORY: ${{ github.repository }} run: | diff --git a/.github/workflows/force-draft-pr.yml b/.github/workflows/force-draft-pr.yml new file mode 100644 index 00000000000..d45dabf14b7 --- /dev/null +++ b/.github/workflows/force-draft-pr.yml @@ -0,0 +1,36 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +name: Force Draft PR + +on: + pull_request_target: + types: [opened] + branches: + - main + +permissions: + pull-requests: write + +jobs: + force-draft: + runs-on: ubuntu-latest + if: ${{ !github.event.pull_request.draft && github.repository == 'NVIDIA/Megatron-LM' }} + steps: + - name: Convert PR to draft + env: + GH_TOKEN: ${{ secrets.PAT }} + run: | + gh pr ready --undo ${{ github.event.pull_request.number }} --repo ${{ github.repository }} + + - name: Add comment explaining draft policy + env: + GH_TOKEN: ${{ github.token }} + run: | + gh pr comment ${{ github.event.pull_request.number }} --repo ${{ github.repository }} --body \ + "This PR has been automatically converted to **draft** because all PRs must start as drafts. + + When you are ready for review, click **Ready for Review** to begin the review process. This will: + 1. Add the oncall reviewer (optional reviewer) + 2. Add required review teams based on your changes + + See the [contribution guide](https://github.com/NVIDIA/Megatron-LM/blob/main/docs/developer/submit.md) for more details." diff --git a/.github/workflows/install-test.yml b/.github/workflows/install-test.yml index ece9184ee94..060e1c5ade0 100644 --- a/.github/workflows/install-test.yml +++ b/.github/workflows/install-test.yml @@ -22,14 +22,14 @@ on: branches: - dev - main - - 'pull-request/[0-9]+' - - 'deploy-release/*' + - "pull-request/[0-9]+" + - "deploy-release/*" merge_group: types: [checks_requested] jobs: pre-flight: - uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.65.5 + uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2 if: github.repository == 'NVIDIA/Megatron-LM' pip-test-pytorch: @@ -43,14 +43,13 @@ jobs: name: Pip - Python${{ matrix.python-version }} - AMD64/Linux - NGC PyTorch container: image: nvcr.io/nvidia/pytorch:25.05-py3 - environment: nemo-ci strategy: fail-fast: false matrix: - python-version: ['3.12'] + python-version: ["3.12"] steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set PATH run: | @@ -66,7 +65,7 @@ jobs: run: bash docker/common/install.sh --environment dev --base-image pytorch --python-version ${{ matrix.python-version }} - name: Checkout check-imports - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: repository: NVIDIA-NeMo/FW-CI-templates ref: v0.63.2 @@ -89,14 +88,13 @@ jobs: name: UV - Python${{ matrix.python-version }} - AMD64/Linux - NGC PyTorch container: image: nvcr.io/nvidia/pytorch:25.05-py3 - environment: nemo-ci strategy: fail-fast: false matrix: - python-version: ['3.12'] + python-version: ["3.12"] steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set PATH run: | @@ -115,7 +113,7 @@ jobs: # NGC PyTorch 25.05 has a version of triton that is broken on CPU only machines. # - name: Checkout check-imports - # uses: actions/checkout@v4 + # uses: actions/checkout@v6 # with: # repository: NVIDIA-NeMo/FW-CI-templates # ref: v0.63.2 @@ -141,7 +139,7 @@ jobs: && github.repository == 'NVIDIA/Megatron-LM' steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Get workflow result id: result diff --git a/.github/workflows/megatron-ci.yml b/.github/workflows/megatron-ci.yml index 7b1da5038e6..8e5c274c572 100644 --- a/.github/workflows/megatron-ci.yml +++ b/.github/workflows/megatron-ci.yml @@ -68,7 +68,7 @@ jobs: id: te-ref run: | set -e - BRANCH=release_v2.10_rocm + BRANCH=sudhu/v2.10_cherrypicks SHA=$(git ls-remote https://github.com/ROCm/TransformerEngine.git "refs/heads/$BRANCH" | cut -f1) if [ -z "$SHA" ]; then echo "::error::Failed to resolve TransformerEngine commit for branch $BRANCH" @@ -77,6 +77,17 @@ jobs: echo "Using TransformerEngine commit: $SHA" echo "sha=$SHA" >> "$GITHUB_OUTPUT" + - name: Determine whether to push image + id: push-image + run: | + # PR and workflow_dispatch jobs build locally only. Push (and cache export) + # runs on trusted branch pushes so registry tags stay branch-controlled. + if [ "${{ github.event_name }}" = "push" ] && [ "${{ github.ref }}" = "refs/heads/rocm_dev" ]; then + echo "should_push=true" >> "$GITHUB_OUTPUT" + else + echo "should_push=false" >> "$GITHUB_OUTPUT" + fi + - name: Login to Docker Hub uses: docker/login-action@v3 with: @@ -86,14 +97,13 @@ jobs: - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - - name: Build and push Docker image with registry cache + - name: Build Docker image with registry cache uses: docker/build-push-action@v5 with: context: . file: ./Dockerfile_rocm.ci - push: true - # Also write image to disk so `docker run` does not need `docker pull`. - # On some runners dockerd cannot reach auth.docker.io (timeout); BuildKit push may still work. + push: ${{ steps.push-image.outputs.should_push == 'true' }} + # Write image to disk so `docker run` does not need registry pull (works for PR builds). outputs: type=docker,dest=${{ runner.temp }}/megatron-ci-image.tar tags: | rocm/megatron-lm-private:${{ github.sha }} @@ -102,8 +112,7 @@ jobs: TE_COMMIT=${{ steps.te-ref.outputs.sha }} cache-from: | type=registry,ref=rocm/megatron-lm-private:cache - cache-to: | - type=registry,ref=rocm/megatron-lm-private:cache,mode=max + cache-to: ${{ steps.push-image.outputs.should_push == 'true' && 'type=registry,ref=rocm/megatron-lm-private:cache,mode=max' || '' }} - name: Load image for docker run run: docker load -i "${{ runner.temp }}/megatron-ci-image.tar" diff --git a/.github/workflows/multi-approval-bot.yml b/.github/workflows/multi-approval-bot.yml index e8507605aa7..c7477679201 100644 --- a/.github/workflows/multi-approval-bot.yml +++ b/.github/workflows/multi-approval-bot.yml @@ -9,13 +9,12 @@ on: jobs: pre-flight: - uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.65.5 + uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2 if: github.repository == 'NVIDIA/Megatron-LM' codeowners-approval: needs: [pre-flight] runs-on: ubuntu-latest - environment: nemo-ci if: | !(needs.pre-flight.outputs.docs_only == 'true' || needs.pre-flight.outputs.is_merge_group == 'true' @@ -27,7 +26,7 @@ jobs: uses: nv-gha-runners/get-pr-info@main - name: Checkout action - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: repository: noamelf/codeowner-multi-approval-action ref: v0.1 @@ -54,7 +53,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Result env: diff --git a/.github/workflows/oncall-assign.yml b/.github/workflows/oncall-assign.yml index d4cc47d5f9e..6da0776ffc2 100644 --- a/.github/workflows/oncall-assign.yml +++ b/.github/workflows/oncall-assign.yml @@ -16,7 +16,7 @@ name: Oncall Assign on: pull_request_target: - types: [opened, ready_for_review] + types: [ready_for_review] branches: - main @@ -30,10 +30,10 @@ jobs: if: ${{ !github.event.pull_request.draft }} steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: '3.10' diff --git a/.github/workflows/oncall-rotation.yml b/.github/workflows/oncall-rotation.yml index 46a45810ad1..0d5f774e441 100644 --- a/.github/workflows/oncall-rotation.yml +++ b/.github/workflows/oncall-rotation.yml @@ -17,7 +17,7 @@ name: Oncall Rotation on: schedule: # Runs at 09:00 UTC every Wednesday - - cron: '0 9 * * 3' + - cron: "0 9 * * 3" workflow_dispatch: permissions: @@ -25,18 +25,17 @@ permissions: jobs: rotate-schedule: - environment: main runs-on: ubuntu-latest steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: token: ${{ secrets.PAT }} - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: '3.10' + python-version: "3.10" - name: Rotate Schedule env: @@ -45,7 +44,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 @@ -56,4 +58,3 @@ jobs: git commit -m "chore: rotate oncall schedule" || echo "No changes to commit" git pull --rebase git push origin HEAD:main - diff --git a/.github/workflows/release-docs.yml b/.github/workflows/release-docs.yml index d15ea74f052..6d619a8a1bc 100644 --- a/.github/workflows/release-docs.yml +++ b/.github/workflows/release-docs.yml @@ -20,23 +20,62 @@ on: required: true type: boolean default: true - version-number: - description: Version number to release this as (use `latest` for main branch) - required: true + publish-as-latest: + description: Publish as Latest stable version. + required: false + type: boolean + default: true + docs-version-override: + description: Docs version if commit is not tagged + required: false type: string + default: "" + update-version-picker: + description: Update version picker. + required: false + type: boolean + default: true notify-emails: description: Email addresses to send the notification to. Format as "me@me.com,you@you.com". + required: false + type: string + workflow_call: + inputs: + dry-run: + description: Whether to run the workflow in dry-run mode required: true + type: boolean + default: true + publish-as-latest: + description: Publish as Latest stable version. + required: false + type: boolean + default: true + docs-version-override: + description: Docs version if commit is not tagged + required: false + type: string + default: "" + update-version-picker: + description: Update version picker. + required: false + type: boolean + default: true + notify-emails: + description: Email addresses to send the notification to. Format as "me@me.com,you@you.com". + required: false type: string - aws-region: - description: AWS region + build-docs-ref: + description: Reference to build the docs from required: false type: string - default: us-east-1 + default: ${{ github.sha }} jobs: build-docs: uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_build_docs.yml@v0.67.0 + with: + ref: ${{ inputs.build-docs-ref }} publish-docs: runs-on: ubuntu-latest @@ -45,7 +84,7 @@ jobs: - uses: actions/checkout@v6 with: repository: NVIDIA-NeMo/FW-CI-templates - ref: v0.67.2 + ref: v0.74.0 path: FW-CI-templates - uses: ./FW-CI-templates/.github/actions/publish-docs @@ -59,10 +98,12 @@ jobs: artifacts-name: docs-html artifacts-path: _build/html emails-csv: ${{ inputs.notify-emails && format('{0},{1}', vars.docs_release_emails, inputs.notify-emails) || vars.docs_release_emails }} - overwrite-latest-on-tag: false + overwrite-latest-on-tag: ${{ inputs.publish-as-latest }} + docs-version-override: ${{ inputs.docs-version-override }} + update-version-picker: ${{ inputs.update-version-picker }} run-on-version-tag-only: ${{ github.ref_name != 'main' }} request-name: megatron-core-publish-docs-${{ github.run_id }} - aws-region: ${{ inputs.aws-region }} + aws-region: ${{ vars.DOCS_AWS_REGION }} aws-role-to-assume: ${{ secrets.AWS_ASSUME_ROLE_ARN }} aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }} aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} diff --git a/.github/workflows/release-freeze.yml b/.github/workflows/release-freeze.yml index 82f26168bd6..dc4bad0a9a7 100644 --- a/.github/workflows/release-freeze.yml +++ b/.github/workflows/release-freeze.yml @@ -42,5 +42,5 @@ jobs: freeze-commit: ${{ inputs.freeze-commit }} dry-run: ${{ inputs.dry-run }} secrets: - SLACK_WEBHOOK: ${{ secrets.SLACK_RELEASE_ENDPOINT }} - SLACK_WEBHOOK_ADMIN: ${{ secrets.SLACK_WEBHOOK_ADMIN }} + SLACK_WEBHOOK: ${{ secrets.SLACK_MAIN_CHANNEL_WEBHOOK }} + SLACK_WEBHOOK_ADMIN: ${{ secrets.SLACK_TEAM_GROUP_ID }} diff --git a/megatron/training/tokenizer/__init__.py b/.github/workflows/release-nightly-docs.yml similarity index 61% rename from megatron/training/tokenizer/__init__.py rename to .github/workflows/release-nightly-docs.yml index 06edc3d6a5f..89ceb1fbcd8 100644 --- a/megatron/training/tokenizer/__init__.py +++ b/.github/workflows/release-nightly-docs.yml @@ -1,5 +1,4 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. -# Copyright (c) 2023 Alibaba PAI Team. +# Copyright (c) 2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,4 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .tokenizer import build_tokenizer +name: Release Nightly Docs + +on: + schedule: + - cron: "0 10 * * *" + +jobs: + call-release-docs: + uses: ./.github/workflows/release-docs.yml + with: + dry-run: false + publish-as-latest: false + docs-version-override: "nightly" + update-version-picker: false + secrets: inherit diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index aa04408689b..a756d49eb20 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -11,7 +11,7 @@ # 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: 'Release Megatron-Core' +name: "Release Megatron-Core" on: workflow_dispatch: @@ -30,10 +30,25 @@ on: required: true default: true type: boolean + generate-changelog: + description: Generate changelog + required: false + default: true + type: boolean + publish-docs: + description: Publish docs + required: false + default: true + type: boolean version-bump-branch: description: Branch for version bump required: true type: string + gh-release-from-tag: + description: Tag of previous release for changelog builder + required: false + type: string + default: "" permissions: contents: write # To read repository content @@ -47,9 +62,18 @@ jobs: dry-run: ${{ inputs.dry-run || false }} version-bump-branch: ${{ inputs.version-bump-branch || github.ref_name }} create-gh-release: ${{ inputs.create-gh-release || true }} + gh-release-use-changelog-builder: ${{ inputs.generate-changelog }} + publish-docs: ${{ inputs.publish-docs }} + gh-release-from-tag: ${{ inputs.gh-release-from-tag }} secrets: - TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }} - TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }} - SLACK_WEBHOOK_ADMIN: ${{ secrets.SLACK_WEBHOOK_ADMIN }} - SLACK_WEBHOOK: ${{ secrets.SLACK_RELEASE_ENDPOINT }} + TWINE_PASSWORD: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/r')) && secrets.SVC_PYPI_TOKEN || secrets.SVC_PYPI_TEST_TOKEN }} + SLACK_WEBHOOK: ${{ (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/heads/r')) && secrets.SLACK_MAIN_CHANNEL_WEBHOOK || secrets.SLACK_CI_CHANNEL_WEBHOOK }} PAT: ${{ secrets.PAT }} + AWS_ASSUME_ROLE_ARN: ${{ secrets.AWS_ASSUME_ROLE_ARN }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + AKAMAI_HOST: ${{ secrets.AKAMAI_HOST }} + AKAMAI_CLIENT_TOKEN: ${{ secrets.AKAMAI_CLIENT_TOKEN }} + AKAMAI_CLIENT_SECRET: ${{ secrets.AKAMAI_CLIENT_SECRET }} + AKAMAI_ACCESS_TOKEN: ${{ secrets.AKAMAI_ACCESS_TOKEN }} + S3_BUCKET_NAME: ${{ secrets.S3_BUCKET_NAME }} diff --git a/.github/workflows/review-trigger.yml b/.github/workflows/review-trigger.yml new file mode 100644 index 00000000000..28abf259882 --- /dev/null +++ b/.github/workflows/review-trigger.yml @@ -0,0 +1,28 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Lightweight workflow that triggers on review approval, otherwise there is no access to right secret. +# No secrets needed — just signals auto-swap-labels.yml via workflow_run. + +name: Review Trigger + +on: + pull_request_review: + types: [submitted] + +jobs: + signal: + runs-on: ubuntu-latest + if: >- + github.event.review.state == 'approved' && + github.event.pull_request.base.ref == 'main' && + github.repository == 'NVIDIA/Megatron-LM' + steps: + - name: Save PR number + run: | + mkdir -p pr + echo "${{ github.event.pull_request.number }}" > pr/number + - name: Upload PR number + uses: actions/upload-artifact@v4 + with: + name: pr-number + path: pr/ diff --git a/.github/workflows/sync-team-usergroups.yml b/.github/workflows/sync-team-usergroups.yml index 8b08182dceb..7f32ac55c57 100644 --- a/.github/workflows/sync-team-usergroups.yml +++ b/.github/workflows/sync-team-usergroups.yml @@ -16,24 +16,28 @@ name: Sync GitHub Teams to Slack User Groups on: workflow_dispatch: + schedule: + - cron: "0 0 * * *" jobs: sync-usergroups: - environment: main runs-on: ubuntu-latest steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: '3.10' + python-version: "3.10" - name: Sync Teams to User Groups env: 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/.github/workflows/trigger-mbridge-tests.yml b/.github/workflows/trigger-mbridge-tests.yml index b1a3aa0089d..023851e966a 100644 --- a/.github/workflows/trigger-mbridge-tests.yml +++ b/.github/workflows/trigger-mbridge-tests.yml @@ -2,182 +2,41 @@ # SPDX-License-Identifier: Apache-2.0 name: Trigger MBridge Tests -# Remote testing of MBridge from MCore -# Triggers MBridge CI tests with current MCore commit to verify backward compatibility - on: - # Manual trigger only workflow_dispatch: inputs: mbridge_ref: - description: 'MBridge branch/ref to trigger' + description: "MBridge branch/ref to trigger" required: false type: string - default: 'main' - run_cicd_main: - description: 'Run cicd-main.yml (full CI/CD)' - required: false - type: boolean - default: true - run_install_test: - description: 'Run install-test.yml (quick install check)' - required: false - type: boolean - default: true + default: "main" test_suite: - description: 'Test suite to run (for cicd-main)' + description: "Test suite to run" required: false type: choice options: - - 'all' - - 'unit-only' - - 'functional-only' - default: 'all' + - "all" + - "unit-only" + - "functional-only" + default: "all" jobs: - # First job: Get MCore commit info (shared by all matrix jobs) - get-mcore-info: + trigger-mbridge-tests: runs-on: ubuntu-latest - outputs: - sha: ${{ steps.mcore_info.outputs.sha }} - short_sha: ${{ steps.mcore_info.outputs.short_sha }} - branch: ${{ steps.mcore_info.outputs.branch }} - repo_url: ${{ steps.mcore_info.outputs.repo_url }} steps: - - name: Checkout MCore - uses: actions/checkout@v4 + - name: Trigger MBridge tests + uses: convictional/trigger-workflow-and-wait@v1.6.5 with: - fetch-depth: 0 - - - name: Get MCore commit info - id: mcore_info - run: | - echo "sha=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT - echo "short_sha=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT - echo "branch=${GITHUB_REF#refs/heads/}" >> $GITHUB_OUTPUT - - # Get repo URL from origin remote, fallback to constructing from github context - REPO_URL=$(git remote get-url origin 2>/dev/null || echo "${{ github.server_url }}/${{ github.repository }}.git") - echo "repo_url=${REPO_URL}" >> $GITHUB_OUTPUT - - echo "📦 MCore commit: $(git rev-parse --short HEAD)" - echo "🌿 Branch: ${GITHUB_REF#refs/heads/}" - echo "📍 Repo: ${REPO_URL}" - - # Matrix job: Trigger and monitor MBridge workflows in parallel - trigger-and-monitor: - needs: [get-mcore-info] - runs-on: ubuntu-latest - continue-on-error: true # Don't fail workflow if monitoring times out - strategy: - fail-fast: false # Continue other matrix jobs even if one fails - matrix: - include: - - workflow: install-test.yml - name: Install Test - - workflow: cicd-main.yml - name: CI/CD Main - - name: ${{ matrix.name }} - - steps: - - name: Check if workflow should run - id: should_run - run: | - if [[ "${{ matrix.workflow }}" == "install-test.yml" && "${{ inputs.run_install_test }}" == "true" ]]; then - echo "run=true" >> $GITHUB_OUTPUT - elif [[ "${{ matrix.workflow }}" == "cicd-main.yml" && "${{ inputs.run_cicd_main }}" == "true" ]]; then - echo "run=true" >> $GITHUB_OUTPUT - else - echo "run=false" >> $GITHUB_OUTPUT - echo "⏭️ Skipping ${{ matrix.workflow }} (not enabled)" - fi - - - name: Trigger ${{ matrix.workflow }} - if: steps.should_run.outputs.run == 'true' - id: trigger - env: - GH_TOKEN: ${{ secrets.PAT }} - run: | - echo "🚀 Triggering ${{ matrix.workflow }} | MCore: ${{ needs.get-mcore-info.outputs.short_sha }} | MBridge: ${{ inputs.mbridge_ref }}" - - gh workflow run ${{ matrix.workflow }} \ - --repo NVIDIA-NeMo/Megatron-Bridge --ref ${{ inputs.mbridge_ref }} \ - --field mcore_commit=${{ needs.get-mcore-info.outputs.sha }} \ - --field mcore_branch=${{ needs.get-mcore-info.outputs.branch }} \ - --field mcore_repo=${{ needs.get-mcore-info.outputs.repo_url }} \ - --field test_suite=${{ inputs.test_suite }} \ - --field triggered_by=mcore-ci - - - name: Get run ID - if: steps.should_run.outputs.run == 'true' - id: get_run_id - env: - GH_TOKEN: ${{ secrets.PAT }} - run: | - sleep 10 # Wait for run to appear - RUN_ID=$(gh run list \ - --repo NVIDIA-NeMo/Megatron-Bridge \ - --workflow=${{ matrix.workflow }} \ - --limit 5 \ - --json databaseId,createdAt \ - --jq "sort_by(.createdAt) | reverse | .[0] | .databaseId") - - echo "run_id=${RUN_ID}" >> $GITHUB_OUTPUT - echo "📋 Run ID: ${RUN_ID}" - - cat >> $GITHUB_STEP_SUMMARY << EOF - ## 🔄 ${{ matrix.name }} Triggered - - **MCore:** \`${{ needs.get-mcore-info.outputs.short_sha }}\` | **MBridge:** \`${{ inputs.mbridge_ref }}\` | **Suite:** \`${{ inputs.test_suite }}\` - - - 🔄 [${{ matrix.workflow }}](https://github.com/NVIDIA-NeMo/Megatron-Bridge/actions/runs/${RUN_ID}) - Running... - - ⏳ Monitoring every 5 minutes until completion - - > **Note:** Tests run without approval when triggered from MCore - EOF - - - name: Monitor workflow - if: steps.should_run.outputs.run == 'true' - id: monitor - continue-on-error: true - env: - GH_TOKEN: ${{ secrets.PAT }} - run: | - RUN_ID="${{ steps.get_run_id.outputs.run_id }}" - echo "📊 Monitoring ${{ matrix.workflow }} (Run ID: ${RUN_ID})" - - gh run watch ${RUN_ID} --repo NVIDIA-NeMo/Megatron-Bridge --exit-status - - CONCLUSION=$(gh run view ${RUN_ID} --repo NVIDIA-NeMo/Megatron-Bridge --json conclusion --jq -r .conclusion) - echo "workflow_status=${CONCLUSION}" >> $GITHUB_ENV - echo "✅ Completed: ${CONCLUSION}" - - - name: Report results - if: always() && steps.should_run.outputs.run == 'true' - run: | - CONCLUSION="${{ env.workflow_status || 'unknown' }}" - RUN_ID="${{ steps.get_run_id.outputs.run_id }}" - - case "$CONCLUSION" in - "success") ICON="✅"; MSG="passed" ;; - "failure") ICON="❌"; MSG="failed"; EXIT_CODE=1 ;; - "cancelled") ICON="🚫"; MSG="cancelled"; EXIT_CODE=0 ;; - *) ICON="⏳"; MSG="still running or timed out"; EXIT_CODE=0 ;; - esac - - cat >> $GITHUB_STEP_SUMMARY << EOF - ## 📊 ${{ matrix.name }} Results - - ### ${ICON} ${{ matrix.workflow }} - **Status:** \`${CONCLUSION}\` - - [View full results →](https://github.com/NVIDIA-NeMo/Megatron-Bridge/actions/runs/${RUN_ID}) - - --- - *Triggered from MCore \`${{ needs.get-mcore-info.outputs.short_sha }}\`* - EOF - - echo "${ICON} ${{ matrix.name }} ${MSG}" - exit ${EXIT_CODE:-0} - + owner: NVIDIA-NeMo + repo: Megatron-Bridge + workflow_file_name: cicd-main.yml + github_token: ${{ secrets.PAT }} + ref: ${{ inputs.mbridge_ref }} + wait_interval: 60 + propagate_failure: true + client_payload: | + { + "mcore_ref": "${{ github.sha }}", + "test_suite": "${{ inputs.test_suite }}", + "triggered_by": "https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}" + } diff --git a/.gitignore b/.gitignore index 13147d1cc88..d6c230d6bdc 100644 --- a/.gitignore +++ b/.gitignore @@ -21,4 +21,7 @@ runs/ # Sphinx documentation docs/_build -docs/apidocs \ No newline at end of file +docs/apidocs + +# Git worktrees +.worktrees/ \ No newline at end of file diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index a238f2c9999..2eb1b43be0c 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -1,16 +1,16 @@ .merge_train_rule: &merge_train_rule - UNIT_TEST: 'yes' + UNIT_TEST: "yes" UNIT_TEST_REPEAT: 1 UNIT_TEST_TIMEOUT: 30 - INTEGRATION_TEST: 'no' + INTEGRATION_TEST: "no" INTEGRATION_TEST_SCOPE: mr - FUNCTIONAL_TEST: 'yes' + FUNCTIONAL_TEST: "yes" FUNCTIONAL_TEST_SCOPE: mr-slim FUNCTIONAL_TEST_REPEAT: 1 FUNCTIONAL_TEST_TIME_LIMIT: 2700 - CLUSTER_A100: '' - CLUSTER_H100: '' - PUBLISH: 'no' + CLUSTER_A100: "" + CLUSTER_H100: "" + PUBLISH: "no" workflow: rules: @@ -29,36 +29,42 @@ workflow: auto_cancel: on_new_commit: none - # For manual pipelines + # For manual pipelines (GitLab UI) - if: $CI_PIPELINE_SOURCE == "web" + # For pipelines created via the REST API (personal access token) + - if: $CI_PIPELINE_SOURCE == "api" + + # For trigger pipelines + - if: $CI_PIPELINE_SOURCE == "trigger" + # For push to main - if: $CI_PIPELINE_SOURCE == 'push' && ($CI_COMMIT_BRANCH == "main" || $CI_COMMIT_BRANCH == "dev" || $CI_COMMIT_BRANCH =~ /^core_/) variables: - UNIT_TEST: 'no' - INTEGRATION_TEST: 'no' - FUNCTIONAL_TEST: 'yes' + UNIT_TEST: "no" + INTEGRATION_TEST: "no" + FUNCTIONAL_TEST: "yes" FUNCTIONAL_TEST_SCOPE: mr FUNCTIONAL_TEST_REPEAT: 5 - FUNCTIONAL_TEST_RECORD_CHECKPOINTS: 'no' + FUNCTIONAL_TEST_RECORD_CHECKPOINTS: "no" FUNCTIONAL_TEST_TIME_LIMIT: 3600 - CLUSTER_A100: '' - CLUSTER_H100: '' - PUBLISH: 'no' + CLUSTER_A100: "" + CLUSTER_H100: "" + PUBLISH: "no" auto_cancel: on_new_commit: interruptible # For merge-trains that need to be fast-tracked - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merge_train' && $CI_MERGE_REQUEST_LABELS =~ /fast-track/ variables: - UNIT_TEST: 'yes' + UNIT_TEST: "yes" UNIT_TEST_REPEAT: 1 UNIT_TEST_TIMEOUT: 30 - INTEGRATION_TEST: 'no' - FUNCTIONAL_TEST: 'no' - CLUSTER_A100: '' - CLUSTER_H100: '' - PUBLISH: 'no' + INTEGRATION_TEST: "no" + FUNCTIONAL_TEST: "no" + CLUSTER_A100: "" + CLUSTER_H100: "" + PUBLISH: "no" # For normal merge-trains - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merge_train' @@ -67,75 +73,75 @@ workflow: # For MRs with integration suite - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run tests/ variables: - UNIT_TEST: 'yes' + UNIT_TEST: "yes" UNIT_TEST_REPEAT: 1 UNIT_TEST_TIMEOUT: 30 - INTEGRATION_TEST: 'yes' + INTEGRATION_TEST: "yes" INTEGRATION_TEST_SCOPE: mr - FUNCTIONAL_TEST: 'no' + FUNCTIONAL_TEST: "no" FUNCTIONAL_TEST_SCOPE: mr-slim FUNCTIONAL_TEST_REPEAT: 1 FUNCTIONAL_TEST_TIME_LIMIT: 2700 - CLUSTER_A100: '' - CLUSTER_H100: '' - PUBLISH: 'no' + CLUSTER_A100: "" + CLUSTER_H100: "" + PUBLISH: "no" # For MRs with nightly - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run nightly/ variables: - UNIT_TEST: 'yes' + UNIT_TEST: "yes" UNIT_TEST_REPEAT: 1 UNIT_TEST_TIMEOUT: 30 - INTEGRATION_TEST: 'no' - FUNCTIONAL_TEST: 'yes' + INTEGRATION_TEST: "no" + FUNCTIONAL_TEST: "yes" FUNCTIONAL_TEST_SCOPE: nightly FUNCTIONAL_TEST_REPEAT: 5 - FUNCTIONAL_TEST_RECORD_CHECKPOINTS: 'no' + FUNCTIONAL_TEST_RECORD_CHECKPOINTS: "no" FUNCTIONAL_TEST_TIME_LIMIT: 2700 - CLUSTER_A100: '' - CLUSTER_H100: '' - PUBLISH: 'no' + CLUSTER_A100: "" + CLUSTER_H100: "" + PUBLISH: "no" # For MRs with weekly - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run weekly/ variables: - UNIT_TEST: 'yes' + UNIT_TEST: "yes" UNIT_TEST_REPEAT: 1 UNIT_TEST_TIMEOUT: 30 - INTEGRATION_TEST: 'no' - FUNCTIONAL_TEST: 'yes' + INTEGRATION_TEST: "no" + FUNCTIONAL_TEST: "yes" FUNCTIONAL_TEST_SCOPE: weekly FUNCTIONAL_TEST_REPEAT: 1 - FUNCTIONAL_TEST_RECORD_CHECKPOINTS: 'no' + FUNCTIONAL_TEST_RECORD_CHECKPOINTS: "no" FUNCTIONAL_TEST_TIME_LIMIT: 9000 - CLUSTER_A100: '' - CLUSTER_H100: '' - PUBLISH: 'no' + CLUSTER_A100: "" + CLUSTER_H100: "" + PUBLISH: "no" # For MRs with heavy suite - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run functional tests/ variables: - UNIT_TEST: 'yes' + UNIT_TEST: "yes" UNIT_TEST_REPEAT: 1 UNIT_TEST_TIMEOUT: 30 - INTEGRATION_TEST: 'no' - FUNCTIONAL_TEST: 'yes' + INTEGRATION_TEST: "no" + FUNCTIONAL_TEST: "yes" FUNCTIONAL_TEST_SCOPE: mr FUNCTIONAL_TEST_REPEAT: 1 FUNCTIONAL_TEST_TIME_LIMIT: 2700 - CLUSTER_A100: '' - CLUSTER_H100: '' - PUBLISH: 'no' + CLUSTER_A100: "" + CLUSTER_H100: "" + PUBLISH: "no" # Default MRs - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' variables: - UNIT_TEST: 'yes' + UNIT_TEST: "yes" UNIT_TEST_REPEAT: 1 UNIT_TEST_TIMEOUT: 30 - INTEGRATION_TEST: 'no' - FUNCTIONAL_TEST: 'no' - PUBLISH: 'no' + INTEGRATION_TEST: "no" + FUNCTIONAL_TEST: "no" + PUBLISH: "no" - when: never @@ -157,109 +163,109 @@ default: variables: BUILD: - value: 'yes' + value: "yes" UNIT_TEST: - value: 'yes' + value: "yes" options: - - 'yes' - - 'no' + - "yes" + - "no" description: To run the funtional test suite UNIT_TEST_REPEAT: - value: '1' - description: 'Number of repetitions' + value: "1" + description: "Number of repetitions" UNIT_TEST_TIMEOUT: - value: '30' + value: "30" description: Timeout (minutes) for Unit tests (all repeats) INTEGRATION_TEST: - value: 'yes' + value: "yes" options: - - 'yes' - - 'no' + - "yes" + - "no" description: To run the integration test suite INTEGRATION_TEST_SCOPE: - value: 'mr' + value: "mr" options: - - 'mr' - - 'nightly' - - 'weekly' - - 'pre-release' - - 'release' - description: 'Testsuite to run (only for INTEGRATION_TEST=yes)' + - "mr" + - "nightly" + - "weekly" + - "pre-release" + - "release" + description: "Testsuite to run (only for INTEGRATION_TEST=yes)" INTEGRATION_TEST_TIME_LIMIT: - value: '900' - description: 'Timeout in seconds per test' + value: "900" + description: "Timeout in seconds per test" INTEGRATION_TEST_CASES: - value: 'all' + value: "all" description: "Comma-separated list of test_cases to run. Use 'all' to run the full suite." FUNCTIONAL_TEST: - value: 'yes' + value: "yes" options: - - 'yes' - - 'no' + - "yes" + - "no" description: To run the funtional test suite FUNCTIONAL_TEST_SCOPE: - value: 'mr' + value: "mr" options: - - 'mr' - - 'nightly' - - 'weekly' - - 'pre-release' - - 'release' - description: 'Testsuite to run (only for FUNCTIONAL_TEST=yes)' + - "mr" + - "nightly" + - "weekly" + - "pre-release" + - "release" + description: "Testsuite to run (only for FUNCTIONAL_TEST=yes)" FUNCTIONAL_TEST_REPEAT: - value: '5' - description: 'Number of repetitions per test' + value: "5" + description: "Number of repetitions per test" FUNCTIONAL_TEST_TIME_LIMIT: - value: '2700' - description: 'Timeout in seconds per test' + value: "2700" + description: "Timeout in seconds per test" FUNCTIONAL_TEST_CASES: - value: 'all' + value: "all" description: "Comma-separated list of test_cases to run. Use 'all' to run the full suite." FUNCTIONAL_TEST_NAME: - description: 'Name of functional test run (only for pre-release and release)' - value: '$$CI_COMMIT_SHA' + description: "Name of functional test run (only for pre-release and release)" + value: "$$CI_COMMIT_SHA" FUNCTIONAL_TEST_RECORD_CHECKPOINTS: - value: 'no' - description: 'Record golden checkpoints' + value: "no" + description: "Record golden checkpoints" options: - - 'yes' - - 'no' + - "yes" + - "no" CLUSTER_A100: - value: 'dgxa100_dracooci' + value: "dgxa100_dracooci" options: - - 'dgxa100_dracooci' - - 'dgxa100_dracooci-ord' - description: 'Cluster for A100 workloads' + - "dgxa100_dracooci" + - "dgxa100_dracooci-ord" + description: "Cluster for A100 workloads" CLUSTER_H100: - value: 'dgxh100_coreweave' + value: "dgxh100_coreweave" options: - - 'dgxh100_coreweave' - - 'dgxh100_eos' - description: 'Cluster for H100 workloads' + - "dgxh100_coreweave" + - "dgxh100_eos" + description: "Cluster for H100 workloads" CLUSTER_GB200: - value: 'dgxgb200_oci-hsg' + value: "dgxgb200_oci-hsg" options: - - 'dgxgb200_oci-hsg' - description: 'Cluster for H100 workloads' + - "dgxgb200_oci-hsg" + description: "Cluster for H100 workloads" PUBLISH: - value: 'no' + value: "no" options: - - 'yes' - - 'no' + - "yes" + - "no" description: Build and publish a wheel to PyPi PUBLISH_COMMIT: - value: '$$CI_COMMIT_SHA' + value: "$$CI_COMMIT_SHA" description: Which commit to publish PUBLISH_VERSION_BUMP_BRANCH: - value: '$$CI_COMMIT_BRANCH' + value: "$$CI_COMMIT_BRANCH" description: Which branch to target for version bump PUBLISH_SCOPE: - value: 'code-freeze' + value: "code-freeze" options: - - 'code-freeze' - - 'release' - - 'review-reminder' - - 'upgrade-dependencies' + - "code-freeze" + - "release" + - "review-reminder" + - "upgrade-dependencies" description: Type of publish (freeze or final release) # CI wide variables @@ -267,7 +273,7 @@ variables: CI_MCORE_DEV_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_ci_dev CI_NEMO_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/nemo_ci UTILITY_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_utility - TE_GIT_REF: '' + TE_GIT_REF: "" include: - .gitlab/stages/00.pre.yml diff --git a/CHANGELOG.md b/CHANGELOG.md deleted file mode 100644 index babdc18b8a4..00000000000 --- a/CHANGELOG.md +++ /dev/null @@ -1,368 +0,0 @@ -# Changelog - -## NVIDIA Megatron Core 0.15.0 - -* Features - * Performance - * Fused QKV preprocessing with precomputed RoPE caches (3x preprocessing speedup, 10-14% E2E) ([MR \!3912](https://github.com/NVIDIA/Megatron-LM/commit/f0d9fa97fead9825ae3eada36ee2df568bfa415b)) - * Use new TE interface for user buffers ([MR \!3886](https://github.com/NVIDIA/Megatron-LM/commit/d47b83807142b6490c7a000e63d25a479b106fd9)) - * Add CPU activation offloading via TE ([MR \!4286](https://github.com/NVIDIA/Megatron-LM/commit/310671436c36e6bd198e92c4f30bc84469cc31d8)) - * Add configurable double buffering ([MR \!4026](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4026)) - * Add Muon optimizer and distributed optimizer support ([MR \!4106](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4106)) - * Add setting to support Adam or AdamW optimizer ([MR \!3866](https://github.com/NVIDIA/Megatron-LM/commit/03fd0b41b3840c6f19558161d98373a9242402e5)) - * MoE - * Add DTensor support for EP and DSv3 modules ([MR \!3955](https://github.com/NVIDIA/Megatron-LM/commit/268fda08592528b7bc1a21aadaed259980ca8efb)) - * Add HybridEP backend to Flex Dispatcher ([MR \!4237](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4237)) - * Support FP8 recomputation for MoE components ([MR \!4030](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4030)) - * Implement NVFP4 Zero Padding for MoE ([MR \!4225](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4225)) - * Compute shared experts before router ([MR \!4068](https://github.com/NVIDIA/Megatron-LM/commit/e8024d716f3036ebcef8c5254c7830ad09aaf41b)) - * Enable bias in expert MLP ([MR \!3858](https://github.com/NVIDIA/Megatron-LM/commit/a329dd6da586261a45a8f7d04c1e659ffedd80ae)) - * Model support - * Add YaRN support for GPT-OSS ([MR \!4044](https://github.com/NVIDIA/Megatron-LM/commit/2c1b77a9984bfa978e7cf1f58522e5f8e045d017)) - * Add support for Qwen3-Next arguments ([MR \!4070](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4070)) - * Add FP8 init for MTP ([MR \!3958](https://github.com/NVIDIA/Megatron-LM/commit/d6c6e54ec5eb43d4e196c7ae84e0e88f28613e6b)) - * Add fp8\_dpa option for FP8 scaling ([MR \!4053](https://github.com/NVIDIA/Megatron-LM/commit/61047e60e617e71ebe120ec293b62df6b0efc84f)) - * Add RADIO-g support to converter and tester ([MR \!4371](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4371)) - * Add audio semantic reasoning data for voice chat and speech instructions ([MR \!4397](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4397)) - * FSDP - * Enable joint training of parallel modules ([MR \!3850](https://github.com/NVIDIA/Megatron-LM/commit/53008b844f98886a2144c216ecd25952cb2dda58)) - * Add support for multimodule communication ([MR \!4235](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4235)) - * Inference - * Add CUDA Graph runner lookup table cache (up to 2x E2E speedup) ([MR \!4082](https://github.com/NVIDIA/Megatron-LM/commit/ab43252fdbedcc3662014ae0e110bd3278d844f4)) - * Add MoE dropping and padding router for CUDA Graph \+ decode ([MR \!3816](https://github.com/NVIDIA/Megatron-LM/commit/56818f9e5090ff9eb0f13f10bfe408aae4031c5c)) - * Dynamic audio shapes with variable sequence lengths (2.5x throughput improvement) ([MR \!4274](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4274)) - * Integrate unified memory for dynamic inference context ([MR \!3985](https://github.com/NVIDIA/Megatron-LM/commit/ef4ae4528a0924159069b9f3a2719616156bafa2)) - * Post-training - * Add GPT-OSS ModelOpt support with quantization, import/export ([MR \!4169](https://github.com/NVIDIA/Megatron-LM/commit/a2d8c806b35bc708b13e6c069e19e5dfb49b8481)) - * Enable KD support with hybrid training loop ([MR \!4021](https://github.com/NVIDIA/Megatron-LM/commit/48d7275062a8307f82bd0fa6c1504032c7f3af96)) - * Add ModelOpt pruning example ([MR \!4022](https://github.com/NVIDIA/Megatron-LM/commit/5a58976ebe007064c2ff5e76e815aa5fcf1a8787)) - * RL - * Add importance sampling and partial rollouts to Megatron RL ([MR \!4000](https://github.com/NVIDIA/Megatron-LM/commit/8399280ed3b72a183f44820896a67392c0a47e3e)) - * Add sequence packing for RL ([MR \!4191](https://github.com/NVIDIA/Megatron-LM/commit/ee8e9307f3ad655e6a46f98a483d8192995b02c2)) - * Ease of use - * Handle CUDA absence during import ([MR \!4120](https://github.com/NVIDIA/Megatron-LM/commit/ae44e49271dc45b51a7400ecf6debc598ba90b54)) - * Add granary dataloader functionality ([MR \!4291](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4291)) - * Enable SWA mixing with attention ([MR \!3855](https://github.com/NVIDIA/Megatron-LM/commit/e5bc9249d7ad34355f5db4c8ff7d7a9080f94dc2)) -* Bug fixes - * Fix convergence bug in MXFP8 parameter gradient buffer reuse ([MR \!3999](https://github.com/NVIDIA/Megatron-LM/commit/c2c36f77cf7a0476daee5bb2dec604c2764de320)) - * Fix loss mask cloning to prevent incorrect updates ([MR \!4164](https://github.com/NVIDIA/Megatron-LM/commit/c94d58f3260aa568588265e07b3c06bb58cbde41)) - * Fix metadata loss in checkpoints ([MR \!4182](https://github.com/NVIDIA/Megatron-LM/commit/d8c6aa4c0b5d4c15ec1196802bce292d4580ed4a)) - * Fix FSDP grad accum fusion support ([MR \!4018](https://github.com/NVIDIA/Megatron-LM/commit/9f72f4775509668173c75eaab5d58a49f4473748)) - * Fix non-TE optimizer checkpoint issue ([MR \!3931](https://github.com/NVIDIA/Megatron-LM/commit/2ebb6ee95af8b547e3c0ac394d494cb189b890bc)) - * Fix BERT virtual pipeline parallelism ([MR \!3993](https://github.com/NVIDIA/Megatron-LM/commit/18420b63408101fe5a49d125fb29625f1ad6ab26)) - * Fix gc.freeze() slowdown by adding gc.collect() on last layer ([MR \!4003](https://github.com/NVIDIA/Megatron-LM/commit/a3f9e566c9595753553a73d403b2a481ad283fc0)) - * Fix full iteration CUDA graph non-tensor handling ([MR \!4019](https://github.com/NVIDIA/Megatron-LM/commit/8479eb35fbca9631acb846c3ad5d868e02214227)) - * Fix model\_auto\_sync mis-set and add gradient assertion ([MR \!4062](https://github.com/NVIDIA/Megatron-LM/commit/03045f2d880813695f75707e3262a2bfb4206dfe)) - * Fix HF import dtype and checkpoint loading issues ([MR \!4095](https://github.com/NVIDIA/Megatron-LM/commit/435e7e0620ff870d99debd73b3c9113226622dde)) - * Fix missing initialization in ProcessGroupCollection ([MR \!4159](https://github.com/NVIDIA/Megatron-LM/commit/5f2becf232a85df8687dc539e604e00a6a875da1)) - * Fix sink attention TP ([MR \!4173](https://github.com/NVIDIA/Megatron-LM/commit/3b1b9b267193d72d4f8dc710561c2368de8c114c)) - * Fix num\_microbatches calculation ([MR \!4199](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4199)) - * Fix 1f1b overlap unit tests for MTP standalone ([MR \!4210](https://github.com/NVIDIA/Megatron-LM/commit/44bc753d69cf509c158bb261434498b141fe5130)) - * Fix stale state dict handling ([MR \!4226](https://github.com/NVIDIA/Megatron-LM/commit/0ba847081113a92ce01084f33cd4a0c1f31b327b)) - * Fix dataset divergence with tokenizer PAD handling ([MR \!4231](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4231)) - * Fix parameter initialization ([MR \!4296](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4296)) - * Ensure tensor-parallel attributes set regardless of initialization flag ([MR \!4312](https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/merge_requests/4312)) -* Known issues - -## NVIDIA Megatron Core 0.14.0 - -* Features - * Inference - * Add async support for DynamicInferenceEngine ([MR \!3187](https://github.com/NVIDIA/Megatron-LM/commit/05079d55a5bfcc7a43f4619e36a40a9e8db3f882)) - * Pad input tensors and enable FP8 weights for FP8 inference ([MR \!3341](https://github.com/NVIDIA/Megatron-LM/commit/6a6cd478839d90cf09a837adf8c79cbc844bc920)) - * Force inference to always gather logits with tensor parallelism ([MR \!3442](https://github.com/NVIDIA/Megatron-LM/commit/7c9cdcb794089968278c7272e0261a68edf5d369)) - * Multi batch size CUDA Graphs for Dynamic Inference ([MR \!3402](https://github.com/NVIDIA/Megatron-LM/commit/30aabe5e3133c6d70aa55aaabad4ea8cb04ce63c)) - * Post-training - * ModelOpt updates ([MR \!3268](https://github.com/NVIDIA/Megatron-LM/commit/550ed5243c3a18e39430c15e8918ee63e41d7eaf)) - * Add speculative decoding AR validation feature - * Add DeepSeek and Qwen model configs - * Performance - * ModelCommProcessGroup integration ([MR \!3391](https://github.com/NVIDIA/Megatron-LM/commit/26adc2dfde53fbc2b063e2fdd1d9ed26578811a6)) - * Add HyperCommGrid: N-Dimensional Communication Grid for Model Parallelism ([MR \!3398](https://github.com/NVIDIA/Megatron-LM/commit/45400df7da7fa23e3aff86804e5ac254d9a8d3c0)) - * Flexible creation and management of communication groups - * Add support for Spike No More embedding initializations and weight decay skipping ([MR \!3500](https://github.com/NVIDIA/Megatron-LM/commit/ee74aa66a06b24e511270f285db475941ef63bfd)) - * MoE - * We're actively optimizing large-scale fine-grained MoE performance on Blackwell Platform. - * Features: - * Support Expert Parallel A2A Overlapping ([MR \!3470](https://github.com/NVIDIA/Megatron-LM/commit/0c6c1176fb3e3e00534b3591f1ad023d4ecad6fb); [MR \!3074](https://github.com/NVIDIA/Megatron-LM/commit/4b30ec54aba97e16a083eca33d2df1dd48e1b48f)) - * Support CP and recompute for MTP ([MR \!3330](https://github.com/NVIDIA/Megatron-LM/commit/650ab87d04105869f197f2ddc441e3b18ca93724)) - * Add support for global aux loss ([MR \!3318](https://github.com/NVIDIA/Megatron-LM/commit/e58d9080ea212e005ccba0b6607bfcc86451285d)) - * Memory Optimization - * Support recomputation for FP8 layernorm/moe\_act/shared\_experts ([MR \!3465](https://github.com/NVIDIA/Megatron-LM/commit/6850cc6a739d168f8c84db6cdacf4fe2931c0c49)) - * Support optimizer offloading for DSV3 FP8 training ([MR \!3659](https://github.com/NVIDIA/Megatron-LM/commit/abbde02f54b62a5194ebe951218e98feceba6d42)) - * Performance Optimization - * Add MoE router fusion ([MR \!3809](https://github.com/NVIDIA/Megatron-LM/commit/d93743a9f11d5d17824b8b49868cc90f2904896f)) - * Updates for MoE cudagraph ([MR \!3631](https://github.com/NVIDIA/Megatron-LM/commit/95452706d7aa16dc174813e12639a8c8356fbe87)) - * Bug fixes: - * Fix router input jitter dtype ([MR \!3774](https://github.com/NVIDIA/Megatron-LM/commit/20b395424d2e2bbfaab57b2f954294eb57c90c82)) - * Model support - * Add MiMo video VLM train example ([MR \!3543](https://github.com/NVIDIA/Megatron-LM/commit/786f5629d3462aff2f8855f51db70e882c475116)) - * Add AVLM for MIMO ([MR \!3624](https://github.com/NVIDIA/Megatron-LM/commit/db41707430bff743f986b5779712c74242b99caa)) - * Ease of use - * Add uv support for source installs ([MR \!3615](https://github.com/NVIDIA/Megatron-LM/commit/164204cd7216e642bdef7299c569d95f02f9a79e)) - * Automated weekly prereleases ([MR \!3574](https://github.com/NVIDIA/Megatron-LM/commit/7e59266c70ef34a246438640af690b55c7ecac28)) -* Bug fixes - * Use mscale\_all\_dim for softmax\_factor ([MR \!2800](https://github.com/NVIDIA/Megatron-LM/commit/e96a358f60c82b8ac8d965d91c3cc4ad0230a4e0)) - * Fix FP8 param blockwise scaling unit test ([MR \!3480](https://github.com/NVIDIA/Megatron-LM/commit/57082f946a04c3390fcfc43634dc546ec3ded033)) - * Fix unit test blockwise scaling ([MR \!3491](https://github.com/NVIDIA/Megatron-LM/commit/6d95fe63658f967e56a3fda88a9c30a424fcb520)) - * Optimize prefill for token-less requests ([MR \!3499](https://github.com/NVIDIA/Megatron-LM/commit/daaa650a9ac4291d4027ca2fdeb4298ce024efd2)) - * Add default values for Fp8Padding and Fp8Unpadding ([MR \!3501](https://github.com/NVIDIA/Megatron-LM/commit/42b2b1d10a9cb699b7e5aa40f6bfba9c2a1348aa)) - * Fix CUDA graph logic for flexible pp layout ([MR \!3505](https://github.com/NVIDIA/Megatron-LM/commit/020d85e50ddf0f0282802002acb3662129a519c5)) - * Load FP8 models with strict=False ([MR \!3508](https://github.com/NVIDIA/Megatron-LM/commit/1ab876ddc4c1893c76f26d775226a8d1dcdfb3d2)) - * Skip rope check for torch \< 1.4.0 ([MR \!3528](https://github.com/NVIDIA/Megatron-LM/commit/d8180ef8ed0bb6f305dcdedf1b27d91304f361a3)) - * Disable Apex tests for stability ([MR \!3539](https://github.com/NVIDIA/Megatron-LM/commit/d1256277fe378add0a2cfd7251f5a350b6d126ec)) - * Fix typo in parallel\_state expert parallelism ([MR \!3548](https://github.com/NVIDIA/Megatron-LM/commit/5783ff32af759b8102cf0cb0bb82b30c48b9da26)) - * Guard modelopt on macOS ([MR \!3549](https://github.com/NVIDIA/Megatron-LM/commit/76144fe1106e4fb0e69aa75b7a6ab66e71e8f37f)) - * Retry on CUDA function failure ([MR \!3554](https://github.com/NVIDIA/Megatron-LM/commit/809aab68307a64c1386d68cc78ef70f8f4e12a80)) - * Fix NCCL mem pool creation error ([MR \!3557](https://github.com/NVIDIA/Megatron-LM/commit/b61e21153146a563309b5d44cb5d7f7425806072)) - * Fix get\_rotary\_seq\_len return type ([MR \!3559](https://github.com/NVIDIA/Megatron-LM/commit/1fa6bc83c7aeae95abc8e86ff0aac596985a01c3)) - * Retry on CUDA function failure ([MR \!3560](https://github.com/NVIDIA/Megatron-LM/commit/7da88d74865c3f1a59894173246f26e7b3bf91b9)) - * Fix NCCL allocator attribute error ([MR \!3565](https://github.com/NVIDIA/Megatron-LM/commit/6b656114795d74c3353cb007c59af49b1752f447)) - * Ensure multi-prompt inference works ([MR \!3568](https://github.com/NVIDIA/Megatron-LM/commit/0fae48931000c9c7af06f7dcf037b5b7d96e0cd6)) - * Fix MD5 on FIPS systems ([MR \!3577](https://github.com/NVIDIA/Megatron-LM/commit/83ee8c2848a3b1d42b40086a64da11e19f4b191f)) - * Fixes dynamic context and inference bugs ([MR \!3582](https://github.com/NVIDIA/Megatron-LM/commit/e9c1da60a1ccc85376666d58568ed1d3e5a4f9db)) - * Fix TE version for interleaved fused RoPE ([MR \!3586](https://github.com/NVIDIA/Megatron-LM/commit/b72b6cc161f5273b545bca09677382917cf20492)) - * Fix MTP with MoE and TP logging ([MR \!3594](https://github.com/NVIDIA/Megatron-LM/commit/9af96623b66693e058f6bfce8d0094dc976792d8)) - * Guard TE import fix ([MR \!3596](https://github.com/NVIDIA/Megatron-LM/commit/1bf946b1ec3f11e71459c7c0d06a97edbed96a1a)) - * Add assertion for NCCL UB case ([MR \!3599](https://github.com/NVIDIA/Megatron-LM/commit/e11d28592f19c122859be764b7afe7c208d9acc1)) - * Remove Encoder PP related Functions ([MR \!3604](https://github.com/NVIDIA/Megatron-LM/commit/9e49aa4446a58cc21c4dc0c5d0806551ad075ca7)) - * Fix segfaults in tests ([MR \!3605](https://github.com/NVIDIA/Megatron-LM/commit/f6492fe8164fd5b9ad55007d435ccfc66cb98cc7)) - * Fix TE error in distributed optimizer ([MR \!3625](https://github.com/NVIDIA/Megatron-LM/commit/e6c510ff3c1159f8955589b26f7c395bdf0607d9)) - * Remove redundant barrier in checkpoint flow ([MR \!3626](https://github.com/NVIDIA/Megatron-LM/commit/26869feb6a3ac7f5616cb7253c37a4244d107d70)) - * Support VPP MTP, fix logging ([MR \!3630](https://github.com/NVIDIA/Megatron-LM/commit/c351a473c7eedac2c43eab0815afb9759f4f8187)) - * Retry mechanism for free(): invalid pointer errors ([MR \!3632](https://github.com/NVIDIA/Megatron-LM/commit/ec35b41b2df145a7ccb84afc48d94e0786e094da)) - * Fix test\_replication.py issues ([MR \!3633](https://github.com/NVIDIA/Megatron-LM/commit/f7b50b271b2e0e396069e02551b21aa6fb374b43)) - * Fix typo in parallel\_state ([MR \!3634](https://github.com/NVIDIA/Megatron-LM/commit/3c79a2c330290df58804c33e28e7c197fcc1f0b9)) - * Fix CUDA graph logic determination ([MR \!3635](https://github.com/NVIDIA/Megatron-LM/commit/90efa3ef8a3c4f9e0f1db9f67ab9348bfa501387)) - * Fix TE installation error ([MR \!3636](https://github.com/NVIDIA/Megatron-LM/commit/7e7322c01c9cb8ec254ecd9042700b22b70fe5c8)) - * Ensure correct sharding type in local tests ([MR \!3643](https://github.com/NVIDIA/Megatron-LM/commit/946357f8dd7fdc12424b3a66bc999e6c0a02696c)) - * Fix cudagraphed backward buffer reuse for last layer ([MR \!3645](https://github.com/NVIDIA/Megatron-LM/commit/ee61cf450d24760952e8995aab045ab6d55b986e)) - * Set default for packed\_seq\_params in get\_rotary\_seq\_len ([MR \!3651](https://github.com/NVIDIA/Megatron-LM/commit/510d58c46664f44c556005ac928c5c531e12f761)) - * Fix dynamic example script errors ([MR \!3653](https://github.com/NVIDIA/Megatron-LM/commit/72e290bf1f4bbf0c8047bb10a51da6ea6372e163)) - * Guard TE import fix ([MR \!3666](https://github.com/NVIDIA/Megatron-LM/commit/ac198fc0d60a8c748597e01ca4c6887d3a7bcf3d)) -* Breaking changes: - * `megatron.core.distributed.custom_fsdp` refactored as breaking change to `megatron.core.distributed.fsdp.src.megatron_fsdp` -* Known issues - -## NVIDIA Megatron Core 0.13.0 - -* Support bf16 dtype for optimizer states to use precision-aware optimizer in TransformerEngine -* MoE - * Features: - * Flexible Asymmetric Virtual Pipeline Parallelism with Custom Pipeline Layout (--pipeline-model-parallel-layout) - * Add support to pass custom parallelism groups to MoE modules. - * Add Hybrid Shard Data-Parallel support for MoE models (--num-distributed-optimizer-instances) - * Support EP \+ custom FSDP training for DeepSeek-V3 - * FP8 support for Multi-Token-Prediction - * Memory Optimization - * Fine-grained recomputation to reduce activation memory. (--recompute-modules with \--recompute-granularity selective) - * Memory efficient token permutation by moving the probs multiplication from unpermutation to activation function of GroupedMLP. - * Performance Optimization - * MLA RoPE fusion kernel and YARN embedding cache. - * FP8 padding optimization of MoE models by padding the routing map. - * Bug fixes: - * Fix the aux loss calculation when expert\_bias or group limited routing is used. This leads to load\_balancing\_loss values change compared to the previous version. - * Fix packed sequence support for MLA - * Known Issues: - * MTP is not compatible with flexible pipeline layout, will be fixed at \!3594. - * MTP convergence issue with TP2, will be fixed at \!3594. - -## NVIDIA Megatron Core 0.12.0 - -* Add FP8 recipe selection to arguments (--fp8-recipe, --first-last-layers-bf16, --num-layers-at-start-in-bf16, --num-layers-at-end-in-bf16) -* Context parallel: fix loss scaling when calculate_per_token_loss=True -* Make the number of data parallel communication buckets configurable (--ddp-num-buckets, --ddp-pad-buckets-for-high-nccl-busbw) -* Inference - * Support in-flight batching and chunked KV cache - * Reduce memory usage, - * by not materializing full attention mask - * by only materializing logits for the last token during decode - * by removing an obsolete tensor reference -* Hybrid Model - * Inference - * Add CUDA graph support - * Change tools/run_mamba_text_generation_server.py to use megatron.core.inference - * Fix a shape issue when materializing logits for Mamba model - * Improve initialization of Mamba layers - * Add configuration switches (--mamba-state-dim, --mamba-head-dim, --mamba-num-groups, --is-hybrid-model) - * Make num_floating_point_operations work with hybrid model - * Make hybrid_conversion.py work with mixer that uses TE linear - * Add FP8 support - * Fix Mamba dt_bias tensor parallelism - * Support multimodal tokenizer - * Improve data parallelism scaling -* MoE - * Features: - * DeepEP support, compatible with all the parallelisms and token drop / dropless - * Important precision improvement: Enable FP32/FP64 routing and unpermutation using –moe-router-dtype. FP32 is recommended for all fine-grained MoE training - * CUDA Graph support for MoE - * Multi-Token Prediction (MTP) Support - * Fused indices_to_multihot kernel for DeepEP dispatcher - * Bug fixes: - * Fix Hang Issue with MoE+Dense Hybrid models - * Update theoretical memory and tflops estimation for MoE and MLA - * Fix MoE Aux loss scaling for per token loss - * Fixes for group limited routing and expert bias. We verified these fixes through dsv3 e2e verifications - * Known issues: - * The ckpt trained with Custom FSDP for MoE may not be compatible with 3D parallel training. - -## NVIDIA Megatron Core 0.11.0 - -* Add multi datacenter training support though N/S connection -* MoE - * Features - * Support DeepSeek-V3 fine-tuning - * Aux-loss-free load balancing strategy - * Node-limited routing and Device-limited routing support. - * Tensor Parallelism support for MLA and Sequence Auxiliary Loss - * MTP (with TP and PP support) is coming soon. - * Permutation / Unpermutation fusion kernel from TransformerEngine. - * Uneven virtual pipeline parallel split support in first and last PP stage. - * Bug fixes: - * Fix the grad scale when TP != expert-TP and average_in_collective is enabled in DDP. - * Fix TEGroupedMLP distckpt compatibility issue with FP8 padding/unpadding. - * Known Issues: - * When training the Dense+MoE hybrid model, the process will hang if any PP rank does not have expert params. -* Add MX-FP16 support for optimizer and master weights -* CUDA Graph memory optimizations -* Enable UCC backend for PP communication -* Optimizer CPU offload support for memory savings -* Models - * Initial RADIO/CRADIO implementation - * llama3.2 support -* Hybrid Model - * Support quantization via TensorRT Model Optimizer - -## NVIDIA Megatron Core 0.10.0 - -* Adding MLA to MCore -* Enable FP8 for GroupedMLP -* MoE Parallel Folding -* Enhance MoE Architecture: Support MoE Layer Frequency Patterns and Configurable MoE FFN Hidden Size -* Multimodal: NVLM training and evaluation support in MCore -* Mamba Hybrid - * Increase performance and reduce memory footprint of Triton language/compiler distributed caching - * Add more unit testing and fix bugs - -## NVIDIA Megatron Core 0.9.0 - -* Uneven pipeline parallelism - * Enable pipeline parallelism where first and last ranks have fewer transformer layers than the intermediate ranks -* Per layer CUDAGraph support for GPT training with Transformer Engine modules -* Enable different TP sizes for the vision encoder -* Enable pipeline parallelism for T5 & Llava models -* Support multi-tile multi-image input in Llava models -* MoE - * FP8 support - * Runtime upcycling support - * Dispatcher implementation optimizations - * Shared expert support with overlapping optimizations - * Qwen Model support -* Known Issues - * When using sequence parallel, during the transformer block forward pass, dropout is not using the appropriate rng context. -* NVRx / Fault tolerance - * fault and hang detection in addition to existing straggler detection - * graceful exit and auto restart - -## NVIDIA Megatron Core 0.8.0 - -* Multimodal - * Added initial support for training vision language models using the LLaVA architecture - * Added initial support for inference with multimodal inputs - * End-to-end multimodal example from data collection to training to evaluation is provided in examples/multimodal -* MoE - * Context Parallel support. - * Distributed checkpoint support for grouped GEMM. -* Mamba - -## NVIDIA Megatron Core 0.7.0 - -* MoE - * Token drop support - * Several efficiency optimizations - * Improved model parallelism - * Memory optimizations -* Distributed checkpointing - * Enabled for Retro - * Asynchronous checkpoint saving -* Several minor bug fixes, speed improvements, and memory optimizations - -## NVIDIA Megatron Core 0.6.0 - -* MoE (Mixture of Experts) - * Performance optimization - * Communication optimization for multi GPU and Single GPU - * 23% improvement (323 TFLOPS/GPU) over MCore 0.5.0 on Mixtral with Hopper BF16 - * GroupedMLP enhancement for Hopper - * DP Overlapping. Support overlapping computation with gradient reduction and parameter gathering. - * All-to-All based Token Dispatcher - * Layer-wise logging for load balancing loss. - * Improved expert parallel support including distributed optimizer. -* Distributed optimizer -* RETRO - * Data processing -* BERT - * Distributed checkpointing -* Dist checkpointing - * PyTorch native distributed backend - * Improved saving/loading speed -* TensorRT-LLM Export - * Integration with TensorRT Model Optimizer Post-training quantization (PTQ) - * Text generation driver to perform PTQ in Megatron-LM - * Llama2 and Nemotron3-8b examples to use TensorRT-LLM unified build API to build engine after training. -* Several minor enhancements, bug fixes, and documentation updates - -## NVIDIA Megatron Core 0.5.0 - -### Key Features and Enhancements - -Megatron core documentation is now [live!](https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/index.html#quick-start) - -### Model Features - -* MoE (Mixture of Experts) - * Support for Z-loss, Load balancing and Sinkhorn - * Layer and communications refactor - * Richer parallelism mappings and EP can be combined with other model parallel techniques for larger MoE variants, e.g. EP + TP + DP + SP + PP - * Token dropless architecture with Top-K routing - * Performance optimization with with GroupedGEMM when number of local experts is > 1 - * Distributed checkpointing -* Interleaved rotary embedding - -### Datasets - -* Masked WordPiece datasets for BERT and T5 -* Raw and mock datasets - -### Parallelism - -### Performance - -* Activation offloading to CPU -* Rope and Swiglu fusion -* Sliding window attention (via Transformer Engine) - -### General Improvements - -* Timers - -## NVIDIA Megatron Core 0.4.0 - -### Key Features and Enhancements - -#### Models - -* BERT -* RETRO -* T5 - -#### Parallelism - -* Mixture of Experts support for GPT -* Model parallel efficient Distributed Data Parallel (DDP) -* Context Parallel (2D Tensor Parallel) support - -#### Datasets - -* GPT Dataset -* Blended Dataset diff --git a/Dockerfile_rocm.ci b/Dockerfile_rocm.ci index 77525d43582..c5d596cabc1 100755 --- a/Dockerfile_rocm.ci +++ b/Dockerfile_rocm.ci @@ -42,14 +42,15 @@ wrapt \ zarr \ wandb \ tensorstore==0.1.45 \ -pytest_mock \ +mock \ pybind11 \ -setuptools==69.5.1 \ +'setuptools<82' \ datasets \ tiktoken \ pynvml \ fastapi \ -uvicorn +uvicorn \ +openai RUN pip install "nltk==3.8.1" @@ -65,7 +66,7 @@ RUN git clone https://github.com/Dao-AILab/causal-conv1d causal-conv1d &&\ # Install mamba WORKDIR ${STAGE_DIR} -RUN pip install mamba-ssm --no-build-isolation +RUN pip install "mamba-ssm==2.3.1" --no-build-isolation # Build flash-attention ARG BUILD_FA="1" @@ -107,6 +108,11 @@ RUN git clone https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git &&\ git checkout v0.1.0 &&\ pip install --no-build-isolation . +RUN git clone https://github.com/NVIDIA/nvidia-resiliency-ext && \ + cd nvidia-resiliency-ext && \ + git checkout v0.6.0 && \ + STRAGGLER_DET_SKIP_CUPTI_EXT_BUILD=1 pip install . + WORKDIR $WORKSPACE_DIR COPY . Megatron-LM WORKDIR $WORKSPACE_DIR/Megatron-LM diff --git a/README.md b/README.md index 3551e74762c..9a62f9bb750 100644 --- a/README.md +++ b/README.md @@ -5,8 +5,8 @@ Megatron-LM and Megatron Core

GPU-optimized library for training transformer models at scale

-[![Documentation](https://img.shields.io/badge/docs-latest-brightgreen.svg?style=flat)](https://docs.nvidia.com/Megatron-Core/developer-guide/latest/index.html) -[![version](https://img.shields.io/badge/release-0.12.0-green)](./CHANGELOG.md) +[![Documentation](https://img.shields.io/badge/docs-latest-brightgreen.svg?style=flat)](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html) +[![version](https://img.shields.io/badge/release-0.15.0-green)](./CHANGELOG.md) [![license](https://img.shields.io/badge/license-Apache-blue)](./LICENSE)
@@ -21,28 +21,33 @@ This repository contains two components: **Megatron-LM** and **Megatron Core**. **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** provides bidirectional Hugging Face ↔ Megatron checkpoint conversion with production-ready recipes. +## Getting Started -## Quick Start +**Install from PyPI:** -Install Megatron Core with pip: +```bash +uv pip install megatron-core +``` -1. Install Megatron Core with required dependencies: +**Or clone and install from source:** - ```bash - pip install --no-build-isolation megatron-core[mlm,dev] - ``` +```bash +git clone https://github.com/NVIDIA/Megatron-LM.git +cd Megatron-LM +uv pip install -e . +``` -2. Clone repository for examples: +> **Note:** Building from source can use a lot of memory. If the build runs out of memory, limit parallel compilation jobs by setting `MAX_JOBS` (e.g. `MAX_JOBS=4 uv pip install -e .`). - ```bash - git clone https://github.com/NVIDIA/Megatron-LM.git - cd Megatron-LM - pip install --no-build-isolation .[mlm,dev] - ``` +For NGC container setup and all installation options, see the **[Installation Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/get-started/install.html)**. +- **[Your First Training Run](https://docs.nvidia.com/megatron-core/developer-guide/latest/get-started/quickstart.html)** - End-to-end training examples with data preparation +- **[Parallelism Strategies](https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/parallelism-guide.html)** - Scale training across GPUs with TP, PP, DP, EP, and CP +- **[Contribution Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/developer/contribute.html)** - How to contribute to Megatron Core # Latest News +- **[2026/03]** **Deprecating Python 3.10 support:** We're officially dropping Python 3.10 support with the upcoming 0.17.0 release. Downstream applications must raise their lower boundary to 3.12 to stay compatible with MCore. - **[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. @@ -61,9 +66,6 @@ Install Megatron Core with pip: - - - # Project Structure ``` @@ -77,20 +79,18 @@ 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 └── docs/ # Documentation ``` - - # Performance Benchmarking 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). @@ -126,11 +126,6 @@ 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 @@ -139,7 +134,7 @@ We also strong scaled the standard GPT-3 model (our version has slightly more th ## Getting Help -- 📖 **[Documentation](https://docs.nvidia.com/Megatron-Core/)** - Official documentation +- 📖 **[Documentation](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html)** - Official documentation - 🐛 **[Issues](https://github.com/NVIDIA/Megatron-LM/issues)** - Bug reports and feature requests ## Contributing @@ -151,7 +146,7 @@ We ❤️ contributions! Ways to contribute: - 📝 **Improve docs** - Make Megatron Core more accessible - 🔧 **Submit PRs** - Contribute code improvements -**→ [Contributing Guide](./CONTRIBUTING.md)** +**→ [Contributing Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/developer/contribute.html)** ## Citation diff --git a/codecov.yml b/codecov.yml new file mode 100644 index 00000000000..aa37017f082 --- /dev/null +++ b/codecov.yml @@ -0,0 +1,14 @@ +comment: false +coverage: + status: + project: false + patch: + default: + target: 80% + threshold: 5% + base: auto + if_ci_failed: error + if_no_uploads: success + if_not_found: success +fixes: + - "/opt/megatron-lm/::" diff --git a/docker/.ngc_version.dev b/docker/.ngc_version.dev index 8e8108b9a9a..2c33440d4e2 100644 --- a/docker/.ngc_version.dev +++ b/docker/.ngc_version.dev @@ -1 +1 @@ -nvcr.io/nvidia/pytorch:25.11-py3 \ No newline at end of file +nvcr.io/nvidia/pytorch:26.02-py3 \ No newline at end of file diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev index bb9ca5fbe9a..7a8b69c1297 100644 --- a/docker/Dockerfile.ci.dev +++ b/docker/Dockerfile.ci.dev @@ -35,6 +35,7 @@ COPY README.md pyproject.toml uv.lock /workspace/ COPY megatron/core/__init__.py /workspace/megatron/core/ COPY megatron/core/package_info.py /workspace/megatron/core/ ARG IMAGE_TYPE=dev +ENV NVTE_BUILD_NUM_PHILOX_ROUNDS=3 RUN --mount=type=cache,target=/root/.cache/uv \ bash -ex <<"EOF" export NVTE_CUDA_ARCHS="80;90;100" @@ -88,7 +89,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/docker/common/install_source_wheels.sh b/docker/common/install_source_wheels.sh index 2f144a6ff0a..eaf601c6045 100644 --- a/docker/common/install_source_wheels.sh +++ b/docker/common/install_source_wheels.sh @@ -40,18 +40,14 @@ MAMBA_WHEEL=$(ls $INPUT_WHEEL_DIR/mamba*.whl) || true CAUSALCONV1D_WHEEL=$(ls $INPUT_WHEEL_DIR/causal_conv1d*.whl) || true [ -z "$CAUSALCONV1D_WHEEL" ] && CAUSALCONV1D_WHEEL=$(bash docker/common/build_causalconv1d.sh --output-wheel-dir $INPUT_WHEEL_DIR | tail -n 1) -GROUPEDGEMM_WHEEL=$(ls $INPUT_WHEEL_DIR/grouped_gemm*.whl) || true -[ -z "$GROUPEDGEMM_WHEEL" ] && GROUPEDGEMM_WHEEL=$(bash docker/common/build_groupedgemm.sh --output-wheel-dir $INPUT_WHEEL_DIR | tail -n 1) - # Override deps that are already present in the base image # only for dev if [ "$ENVIRONMENT" = "dev" ]; then uv pip install --no-cache-dir --no-deps $TE_WHEEL fi -# Install heavy optional deps like mamba, causalconv1d, groupedgemm +# Install heavy optional deps like mamba, causalconv1d uv pip install --no-cache-dir \ $MAMBA_WHEEL \ $CAUSALCONV1D_WHEEL \ - $GROUPEDGEMM_WHEEL \ "setuptools<80.0.0,>=77.0.0" diff --git a/docs/add_copyright_header.py b/docs/add_copyright_header.py new file mode 100644 index 00000000000..9694ef84819 --- /dev/null +++ b/docs/add_copyright_header.py @@ -0,0 +1,30 @@ +#!/usr/bin/env python3 +"""One-off script to add NVIDIA copyright header to all .md files under docs/.""" + +from pathlib import Path + +HEADER = """ Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. + NVIDIA CORPORATION and its licensors retain all intellectual property + and proprietary rights in and to this software, related documentation + and any modifications thereto. Any use, reproduction, disclosure or + distribution of this software and related documentation without an express + license agreement from NVIDIA CORPORATION is strictly prohibited. + +""" + +def main(): + docs_dir = Path(__file__).resolve().parent + already_has = "Copyright (c) 2022-2026, NVIDIA CORPORATION" + count = 0 + for path in sorted(docs_dir.rglob("*.md")): + content = path.read_text(encoding="utf-8") + if content.strip().startswith(already_has): + continue + new_content = HEADER + content + path.write_text(new_content, encoding="utf-8") + count += 1 + print(path.relative_to(docs_dir)) + print(f"\nUpdated {count} files.") + +if __name__ == "__main__": + main() diff --git a/docs/advanced/index.md b/docs/advanced/index.md index 573cb0ee81a..98ff0806cff 100644 --- a/docs/advanced/index.md +++ b/docs/advanced/index.md @@ -1,3 +1,12 @@ + + # Discussions In-depth technical discussions and optimization guides: diff --git a/docs/api-backwards-compatibility-check.md b/docs/api-backwards-compatibility-check.md index a417abfc2df..40f56ec0c00 100644 --- a/docs/api-backwards-compatibility-check.md +++ b/docs/api-backwards-compatibility-check.md @@ -1,3 +1,16 @@ +--- +orphan: true +--- + + + # API Backward Compatibility Checking ## Overview diff --git a/docs/api-guide/core/datasets.md b/docs/api-guide/core/datasets.md index e97e99ae1db..d80c0183375 100644 --- a/docs/api-guide/core/datasets.md +++ b/docs/api-guide/core/datasets.md @@ -1,3 +1,12 @@ + + # datasets package ```{include} ../../../megatron/core/datasets/readme.md diff --git a/docs/api-guide/core/dist_checkpointing.md b/docs/api-guide/core/dist_checkpointing.md index 959aa4b07e0..ee0e5562ef3 100644 --- a/docs/api-guide/core/dist_checkpointing.md +++ b/docs/api-guide/core/dist_checkpointing.md @@ -1,3 +1,12 @@ + + # dist_checkpointing package A library for saving and loading the distributed checkpoints. @@ -32,19 +41,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. + +### 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 -######################################## +## Distributed Optimizer Checkpoint Formats The refactor of the Distributed Optimizer introduces **two checkpoint formats**: @@ -57,8 +109,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/api-guide/core/dist_checkpointing.strategies.md b/docs/api-guide/core/dist_checkpointing.strategies.md index 7aab8609504..22fe3517a54 100644 --- a/docs/api-guide/core/dist_checkpointing.strategies.md +++ b/docs/api-guide/core/dist_checkpointing.strategies.md @@ -1,3 +1,12 @@ + + # dist_checkpointing.strategies package Package defining different checkpoint formats (backends) and saving/loading algorithms (strategies). diff --git a/docs/api-guide/core/distributed.md b/docs/api-guide/core/distributed.md index 1921c0bdd57..13da4285ec5 100644 --- a/docs/api-guide/core/distributed.md +++ b/docs/api-guide/core/distributed.md @@ -1,3 +1,12 @@ + + # distributed package This package contains various utilities to finalize model weight gradients diff --git a/docs/api-guide/core/fusions.md b/docs/api-guide/core/fusions.md index 396280ad7da..fdd358e813c 100644 --- a/docs/api-guide/core/fusions.md +++ b/docs/api-guide/core/fusions.md @@ -1,3 +1,12 @@ + + # fusions package This package provides modules that provide commonly fused diff --git a/docs/api-guide/core/index.md b/docs/api-guide/core/index.md index 150fd72cb1e..0d39e46e744 100644 --- a/docs/api-guide/core/index.md +++ b/docs/api-guide/core/index.md @@ -1,3 +1,12 @@ + + # Core APIs Low-level API reference for core Megatron components. diff --git a/docs/api-guide/core/pipeline_parallel.md b/docs/api-guide/core/pipeline_parallel.md index 42fac8cc449..35f3c5b5cc2 100644 --- a/docs/api-guide/core/pipeline_parallel.md +++ b/docs/api-guide/core/pipeline_parallel.md @@ -1,3 +1,12 @@ + + # pipeline_parallel package This package contains implementations for two different pipeline parallelism diff --git a/docs/api-guide/core/tensor_parallel.md b/docs/api-guide/core/tensor_parallel.md index 33a9160c82b..2d41c5f4467 100644 --- a/docs/api-guide/core/tensor_parallel.md +++ b/docs/api-guide/core/tensor_parallel.md @@ -1,3 +1,12 @@ + + # tensor_parallel package This package contains an implementation for tensor parallelism in transformer diff --git a/docs/api-guide/core/transformer.md b/docs/api-guide/core/transformer.md index d004381844b..d35144fda4f 100644 --- a/docs/api-guide/core/transformer.md +++ b/docs/api-guide/core/transformer.md @@ -1,11 +1,19 @@ + + # transformer package The `transformer` package provides a customizable and configurable implementation of the transformer model architecture. Each component of a transformer stack, from entire layers down to individual linear layers, can be customized by swapping in different PyTorch modules -using the "spec" parameters (see [here](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/nlp/nemo_megatron/mcore_customization.html)). The +using the "spec" parameters. The configuration of the transformer (hidden size, number of layers, number of attention heads, etc.) is provided via a `TransformerConfig` object. - diff --git a/docs/api-guide/index.md b/docs/api-guide/index.md index 851114d98e8..7afa2450dd0 100644 --- a/docs/api-guide/index.md +++ b/docs/api-guide/index.md @@ -1,3 +1,12 @@ + + # API Guide API reference documentation for Megatron Core components. diff --git a/docs/api-guide/internal/index.md b/docs/api-guide/internal/index.md index c216a976c77..312081ce70b 100644 --- a/docs/api-guide/internal/index.md +++ b/docs/api-guide/internal/index.md @@ -1,3 +1,12 @@ + + # Internal Utilities Internal utility APIs. diff --git a/docs/api-guide/internal/num_microbatches_calculator.md b/docs/api-guide/internal/num_microbatches_calculator.md index 470c9e49128..0c223588ce5 100644 --- a/docs/api-guide/internal/num_microbatches_calculator.md +++ b/docs/api-guide/internal/num_microbatches_calculator.md @@ -1,3 +1,12 @@ + + # Microbatches Calculator This api is used to calculate the number of microbatches required to fit a given model on a given batch size. diff --git a/docs/api-guide/internal/optimizer_param_scheduler.md b/docs/api-guide/internal/optimizer_param_scheduler.md index 13e1f77ccc0..45e5e4b7da1 100644 --- a/docs/api-guide/internal/optimizer_param_scheduler.md +++ b/docs/api-guide/internal/optimizer_param_scheduler.md @@ -1,3 +1,12 @@ + + # Optimizer Parameters Scheduler This api is used to calculate the learning rate and weight decay for the optimizer. diff --git a/docs/api-guide/models/index.md b/docs/api-guide/models/index.md index c6279d2409a..e5bb531454b 100644 --- a/docs/api-guide/models/index.md +++ b/docs/api-guide/models/index.md @@ -1,3 +1,12 @@ + + # Model APIs API reference for Megatron Core model implementations. diff --git a/docs/api-guide/models/models.bert.md b/docs/api-guide/models/models.bert.md index 3c53027c7c9..1543f4df865 100644 --- a/docs/api-guide/models/models.bert.md +++ b/docs/api-guide/models/models.bert.md @@ -1,3 +1,12 @@ + + # models.bert package Useful package for training bert and bert like encoder only models. It optionally comes with a binary head that can be used for classification tasks . diff --git a/docs/api-guide/models/models.gpt.md b/docs/api-guide/models/models.gpt.md index a7c254d348b..1c3cbb5484c 100644 --- a/docs/api-guide/models/models.gpt.md +++ b/docs/api-guide/models/models.gpt.md @@ -1,3 +1,12 @@ + + # models.gpt package This is the implementation of the popular GPT model. It supports several features like model parallelization (Tensor Parallel, Pipeline Parallel, Data Parallel) , mixture of experts, FP8 , Distributed optimizer etc. We are constantly adding new features. So be on the lookout or raise an issue if you want to have something added. diff --git a/docs/api-guide/models/models.md b/docs/api-guide/models/models.md index 69dfc80211d..a633546f0c9 100644 --- a/docs/api-guide/models/models.md +++ b/docs/api-guide/models/models.md @@ -1,3 +1,12 @@ + + # models package This package contains most of the popular LLMs . Currently we have support for GPT, Bert, and T5 . This is an ever growing list so keep an eye out. diff --git a/docs/api-guide/models/models.t5.md b/docs/api-guide/models/models.t5.md index 90952096b63..4694b80113a 100644 --- a/docs/api-guide/models/models.t5.md +++ b/docs/api-guide/models/models.t5.md @@ -1,2 +1,11 @@ + + # models.t5 package diff --git a/docs/api-guide/router_replay.md b/docs/api-guide/router_replay.md index 300a50db127..88476ea7537 100644 --- a/docs/api-guide/router_replay.md +++ b/docs/api-guide/router_replay.md @@ -1,18 +1,27 @@ + + # Design Document: MoE Router Replay Feature -### 1. Overview +## 1. Overview This document provides a detailed description of the "Router Replay" feature implemented within the Megatron-LM Core for Mixture-of-Experts (MoE) models. This feature is designed to enhance determinism and analyzability in MoE model training and inference. It enables the model to load routing decisions from a predefined file and enforce their use during the forward pass, thereby bypassing the real-time routing computation. -### 2. Motivation +## 2. Motivation * **Determinism & Reproducibility**: In distributed training, MoE routing decisions can exhibit minor variations due to factors like floating-point precision. By replaying a fixed routing table, the MoE computation path is guaranteed to be identical across runs, which facilitates debugging and reproducing experimental results. * **Performance Profiling**: The router's own computation (e.g., logits calculation, top-k selection) incurs overhead. In replay mode, this part of the computation can be completely skipped, allowing for more precise isolation and profiling of performance bottlenecks within the Expert Layers themselves. * **Debugging Aid**: When issues arise in the model, fixing the routing decisions helps to isolate variables, making it easier to determine whether the problem lies with the routing mechanism or the expert computations. -### 3. Design and Architecture +## 3. Design and Architecture The design follows the principles of being non-intrusive and on-demand, with the core idea of activating the replay logic only when explicitly requested by the user. @@ -35,7 +44,7 @@ The design follows the principles of being non-intrusive and on-demand, with the * For each micro-batch (processed in reverse order in pipeline parallelism), the `router_replay_action` is checked again. * **In `backward_replay` mode**: The function retrieves the expert indices for the corresponding micro-batch by popping them from the `replay_backward_list`. This mode is intended for training recomputation (e.g., activation checkpointing and pipeline recompute) so the same routing decisions are used during recompute/backward as in forward, ensuring determinism and correctness. -### 4. Implementation Details +## 4. Implementation Details The implementation cleanly separates the replay logic from the router's core computation. @@ -51,11 +60,12 @@ The implementation cleanly separates the replay logic from the router's core com * `record_indices()`: A method to save the computed indices. * The `topk_routing_with_score_function` is modified to contain the core logic. It checks the `router_replay_action` on the `router_replay` instance and accordingly performs one of the following actions: computes and records indices, replays indices from `target_topk_idx` (for forward), replays indices from `replay_backward_list` (for backward), or falls through to the default dynamic routing. -#### Training recompute usage +### Training recompute usage + - During forward replay, `set_target_indices()` prepares `replay_backward_list` so each micro-batch’s indices are available for recomputation. - During recompute/backward, set action to `REPLAY_BACKWARD` so indices are consumed in FIFO order to mirror the forward sequence. -### 5. Usage Guide +## 5. Usage Guide 1. **Enable & Instantiate** - Create one `RouterReplay` instance per MoE router layer when building the model. @@ -73,7 +83,7 @@ The implementation cleanly separates the replay logic from the router's core com 5. **Cleanup** - Use `RouterReplay.clear_global_indices()`, `RouterReplay.clear_global_router_replay_action()`, and `RouterReplay.clear_global_router_replay_instances()` to restore default behavior and prevent memory leaks. -#### Quick usage with `topk_routing_with_score_function` +### Quick usage with `topk_routing_with_score_function` ```python import torch @@ -105,7 +115,7 @@ RouterReplay.clear_global_indices() RouterReplay.clear_global_router_replay_instances() ``` -### 6. Minimal Demo +## 6. Minimal Demo Here is a minimal code example showing how to use RouterReplay for recording and replaying: diff --git a/docs/broken_links_false_positives.json b/docs/broken_links_false_positives.json new file mode 100644 index 00000000000..01377be5804 --- /dev/null +++ b/docs/broken_links_false_positives.json @@ -0,0 +1,3 @@ +{ + "uri": "http://localhost:8080/" +} \ No newline at end of file diff --git a/docs/conf.py b/docs/conf.py index a64da441084..9bf0b99c706 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -25,9 +25,9 @@ # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information project = "Megatron Core" -copyright = "2025, NVIDIA Corporation" +copyright = "2026, NVIDIA Corporation" author = "NVIDIA Corporation" -release = "latest" +release = "nightly" # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration @@ -63,6 +63,12 @@ ] myst_heading_anchors = 5 # Generates anchor links for headings up to level 5 +# Suppress "more than one target found for cross-reference" warnings for Python symbols +# that have the same name across multiple modules (e.g. DistributedDataParallelConfig, +# ModelType). These are structural ambiguities in the codebase – the cross-reference +# still resolves; Sphinx just cannot pick the unique target automatically. +suppress_warnings = ["ref.python"] + # -- Options for Autodoc2 --------------------------------------------------- sys.path.insert(0, os.path.abspath("..")) @@ -81,6 +87,11 @@ autodoc2_docstring_parser_regexes = [ (r".*", "docs.autodoc2_docstrings_parser"), ] + # Regex patterns whose values contain raw regex syntax (e.g. \p{L}) that docutils + # mis-parses as footnote/reference markup. Exclude them from the generated docs. + autodoc2_hidden_regexes = [ + r".*\._PATTERN_TIKTOKEN.*", + ] # -- Options for HTML output ------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output @@ -98,16 +109,7 @@ "icon": "fa-brands fa-github", } ], - "extra_head": { - """ - - """ - }, - "extra_footer": { - """ - - """ - }, + "public_docs_features": True } html_extra_path = ["project.json", "versions1.json"] diff --git a/docs/developer/contribute.md b/docs/developer/contribute.md index 859b5562f4b..aeb785f915d 100644 --- a/docs/developer/contribute.md +++ b/docs/developer/contribute.md @@ -1,3 +1,12 @@ + + # Contributing to Megatron-LM This document outlines the processes and policies for issues and pull requests by non-NVIDIA contributors to the Megatron-LM GitHub repository. diff --git a/docs/developer/generate_docs.md b/docs/developer/generate_docs.md index 52fa288122d..d985f542caa 100644 --- a/docs/developer/generate_docs.md +++ b/docs/developer/generate_docs.md @@ -1,3 +1,12 @@ + + # Generating Docs Locally To generate docs locally, use the following commands: diff --git a/docs/developer/oncall.md b/docs/developer/oncall.md index b88da7bb6df..0e5b38e2708 100644 --- a/docs/developer/oncall.md +++ b/docs/developer/oncall.md @@ -1,3 +1,13 @@ + +--> + # Oncall Overview During your oncall week, you will be assigned to all PRs marked “Ready for @@ -30,13 +40,13 @@ Below is the checklist that the oncall needs to go through for each PR. - Do all tests pass? - Oncall will need to kick off testing suite for external reviewers - Comment “/ok to test commid_id” to kick off testing suite -- Add the “Expert Review” label - - Select an expert reviewer from each expert group as a reviewer. If you’re unsure who to select, pick a “maintainer” or manager. +- Expert reviewers are notified after the PR is marked “Ready for Review” - **Expert reviewers should review within 1 business day.** Message the assigned reviewer if it is taking longer. The reviewer either needs to review the PR or suggest an alternate reviewer. - - If the reviewer is not responding after 2 business days, escalate to the reviewer's manager. -- Add the “Final Review” label after experts approve + - If the reviewer is not responding after 2 business days, escalate to the reviewer’s manager. +- For `megatron/core` PRs, the “Final Review” label is applied automatically once all expert reviewers approve - Final reviewers should review within 1 business day. Message the assigned reviewer if it is taking longer. - - If the reviewer is not responding after 2 business days, escalate to the reviewer's manager. + - If the reviewer is not responding after 2 business days, escalate to the reviewer’s manager. +- The “Approved” label is applied automatically once all required reviewers have approved ## Issues and Discussion Questions diff --git a/docs/developer/submit.md b/docs/developer/submit.md index a096312d21e..205e18cc52f 100644 --- a/docs/developer/submit.md +++ b/docs/developer/submit.md @@ -1,16 +1,34 @@ + + # How to Submit a PR -## Step 1: Add PR label `Expert Review` +All PRs start as **draft**. If you open a non-draft PR, it will be automatically converted to draft. -## Step 2: Collect the expert reviewers reviews +## Step 1: Mark PR as "Ready for Review" -1. Attach the `Expert Review` label when your PR is ready for review. -2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon. +1. When your PR is ready, click **Ready for Review**. +2. The oncall reviewer is auto-assigned and expert reviewers are notified based on your changes. They will get notified and pick up your PR soon. -:warning: Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing. +:warning: Only mark as ready once all merge-conflicts are resolved and the CI is passing. Final Review might get declined if these requirements are not fulfilled. -## Step 3: Final Review +## Step 2: Final Review (`megatron/core` only) + +For PRs that change `megatron/core`, once all expert reviewers have approved, the `Final Review` label is applied **automatically** and final reviewers are assigned. + +For PRs outside `megatron/core`, this step is skipped. + +## Step 3: Approved + +Once all required reviewers have approved, the `Approved` label is applied **automatically**. The PR is now ready to merge. + +## Step 4: Merge -1. Add `Final Review` label -2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon. +Any member of [mcore-engineers](https://github.com/orgs/NVIDIA/teams/mcore-engineers) will be able to merge your PR. diff --git a/docs/discussions/README.md b/docs/discussions/README.md index 4ac3c4e3254..e791ed57cd8 100644 --- a/docs/discussions/README.md +++ b/docs/discussions/README.md @@ -1,3 +1,16 @@ +--- +orphan: true +--- + + + # Megatron Discussions This directory contains in-depth guides, tutorials, and discussions about optimizing and using Megatron for various use cases. diff --git a/docs/discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md b/docs/discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md index c2354ad07f0..b5de090ab46 100644 --- a/docs/discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md +++ b/docs/discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md @@ -1,3 +1,16 @@ +--- +orphan: true +--- + + + # Megatron-FSDP User Guide ## Table of Contents diff --git a/docs/documentation.md b/docs/documentation.md index 16fbd7b9a7e..d2554157b45 100644 --- a/docs/documentation.md +++ b/docs/documentation.md @@ -2,6 +2,15 @@ orphan: true --- + + # Documentation Development - [Documentation Development](#documentation-development) diff --git a/docs/get-started/install.md b/docs/get-started/install.md index dd000500f58..5781d065fae 100644 --- a/docs/get-started/install.md +++ b/docs/get-started/install.md @@ -1,87 +1,123 @@ -# Megatron Core Installation + -Installation is supported using Docker and pip. +# Installation ## System Requirements -### Hardware Requirements +### Hardware -- **FP8 Support**: NVIDIA Hopper, Ada, Blackwell GPUs - **Recommended**: NVIDIA Turing architecture or later +- **FP8 Support**: Requires NVIDIA Hopper, Ada, or Blackwell GPUs -### Software Requirements +### Software -- **CUDA/cuDNN/NCCL**: Latest stable versions -- **PyTorch**: Latest stable version -- **Transformer Engine**: Latest stable version -- **Python**: 3.12 recommended +- **Python**: >= 3.10 (3.12 recommended) +- **PyTorch**: >= 2.6.0 +- **CUDA Toolkit**: Latest stable version -## Docker Installation (Recommended) +## Prerequisites -We strongly recommend using the previous releases of [PyTorch NGC Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) rather than the latest one for optimal compatibility with Megatron Core release and testing matrix. Our releases are always based on the previous month's NGC container, so this ensures compatibility and stability. +Install [uv](https://docs.astral.sh/uv/), a fast Python package installer: + +```bash +curl -LsSf https://astral.sh/uv/install.sh | sh +``` -**Note:** The NGC PyTorch container constraints the python environment globally via `PIP_CONSTRAINT`. In the following examples we will unset the variable. -This container comes with all dependencies pre-installed with compatible versions and optimized configurations for NVIDIA GPUs: +## Option A: Pip Install (Recommended) -- PyTorch (latest stable version) -- CUDA, cuDNN, NCCL (latest stable versions) -- Support for FP8 on NVIDIA Hopper, Ada, and Blackwell GPUs -- For best performance, use NVIDIA Turing GPU architecture generations and later +Install the latest stable release from PyPI: ```bash -# Run container with mounted directories -docker run --runtime --nvidia --gpus all -it --rm \ - -v /path/to/megatron:/workspace/megatron \ - -v /path/to/dataset:/workspace/dataset \ - -v /path/to/checkpoints:/workspace/checkpoints \ - -e PIP_CONSTRAINT= \ - nvcr.io/nvidia/pytorch:25.04-py3 +uv pip install megatron-core +``` + +To include optional training dependencies (Weights & Biases, SentencePiece, HF Transformers): + +```bash +uv pip install "megatron-core[training]" +``` + +For all extras including [Transformer Engine](https://github.com/NVIDIA/TransformerEngine): + +```bash +uv pip install --group build +uv pip install --no-build-isolation "megatron-core[training,dev]" +``` + +```{note} +`--no-build-isolation` requires build dependencies to be pre-installed in the environment. `torch` is needed because several `[dev]` packages (`mamba-ssm`, `nv-grouped-gemm`, `transformer-engine`) import it at build time to compile CUDA kernels. Expect this step to take **20+ minutes** depending on your hardware. If you prefer pre-built binaries, the [NGC Container](#option-c-ngc-container) ships with these pre-compiled. +``` + +```{warning} +Building from source can consume a large amount of memory. By default the build runs one compiler job per CPU core, which may cause out-of-memory failures on machines with many cores. To limit parallel compilation jobs, set the `MAX_JOBS` environment variable before installing (e.g. `MAX_JOBS=4`). +``` + +```{tip} +For a lighter set of development dependencies without Transformer Engine and ModelOpt, use `[lts]` instead of `[dev]`: `uv pip install --no-build-isolation "megatron-core[training,lts]"`. The `[lts]` and `[dev]` extras are mutually exclusive. ``` -## Pip Installation +To clone the repository for examples: -Megatron Core installation offers support for two NGC PyTorch containers: +```bash +git clone https://github.com/NVIDIA/Megatron-LM.git +``` -- `dev`: Moving head that supports the most recent upstream dependencies -- `lts`: Long-term support of NGC PyTorch 24.01 -Both containers can be combined with `mlm`, which adds package dependencies for Megatron-LM on top of Megatron Core. +## Option B: Install from Source +For development or to run the latest unreleased code: -1. Install the latest release dependencies +```bash +git clone https://github.com/NVIDIA/Megatron-LM.git +cd Megatron-LM +uv pip install -e . +``` - ```bash - pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2" - pip install --no-build-isolation megatron-core[dev] - ``` +To install with all development dependencies (includes Transformer Engine, requires pre-installed build deps): -2. Next choose one of the following options: +```bash +uv pip install --group build +uv pip install --no-build-isolation -e ".[training,dev]" +``` -* For running an Megatron LM application +```{tip} +If the build runs out of memory, limit parallel compilation jobs with `MAX_JOBS=4 uv pip install --no-build-isolation -e ".[training,dev]"`. +``` - ```bash - pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2" - pip install --no-build-isolation megatron-core[mlm,dev] - ``` -* Install packages for LTS support NGC PyTorch 24.01 - ```bash - pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2" - pip install --no-build-isolation megatron-core[lts] - ``` +## Option C: NGC Container -* For running an Megatron LM application +For a pre-configured environment with all dependencies pre-installed (PyTorch, CUDA, cuDNN, NCCL, Transformer Engine), use the [PyTorch NGC Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch). - ```bash - pip install "setuptools<80.0.0,>=77.0.0" "packaging>=24.2" - pip install --no-build-isolation megatron-core[mlm,lts] - ``` +We recommend using the **previous month's** NGC container rather than the latest one to ensure compatibility with the current Megatron Core release and testing matrix. -* For a version of Megatron Core with only Torch, run +```bash +docker run --gpus all -it --rm \ + -v /path/to/dataset:/workspace/dataset \ + -v /path/to/checkpoints:/workspace/checkpoints \ + -e PIP_CONSTRAINT= \ + nvcr.io/nvidia/pytorch:26.01-py3 +``` + +```{note} +The NGC PyTorch container constrains the Python environment globally via `PIP_CONSTRAINT`. The `-e PIP_CONSTRAINT=` flag above unsets this so that Megatron Core and its dependencies install correctly. +``` + +Then install Megatron Core inside the container (torch is already available in the NGC image): + +```bash +pip install uv +uv pip install --no-build-isolation "megatron-core[training,dev]" +``` - ```bash - pip install megatron-core - ``` +You are now ready to run training. See [Your First Training Run](quickstart.md) for next steps. diff --git a/docs/get-started/overview.md b/docs/get-started/overview.md index 883f40e0c61..b7f84ee22e5 100644 --- a/docs/get-started/overview.md +++ b/docs/get-started/overview.md @@ -1,3 +1,12 @@ + + # Overview Megatron-Core and Megatron-LM are open-source tools that are typically used together to train LLMs at scale across GPUs. Megatron-Core expands the capability of Megatron-LM. Megatron Bridge connects Megatron-Core and Megatron-LM to other popular training models, such as Hugging Face. @@ -78,7 +87,7 @@ After training or modifying a Megatron model, you can convert it again for deplo - **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Training library with bidirectional Hugging Face ↔ Megatron checkpoint conversion, flexible training loops, and production-ready recipes - **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods - **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples -- **[Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Checkout end-to-end examples in [examples/post_training/modelopt](./examples/post_training/modelopt/). +- **[Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Checkout end-to-end examples in [examples/post_training/modelopt](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/post_training/modelopt). **Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed) diff --git a/docs/get-started/quickstart.md b/docs/get-started/quickstart.md index 61868e7877c..c8797aeedd4 100644 --- a/docs/get-started/quickstart.md +++ b/docs/get-started/quickstart.md @@ -1,51 +1,46 @@ -# Quick Start + -## Quick Installation +# Your First Training Run -Install Megatron Core with pip: +This guide walks you through running your first training jobs with Megatron Core. Make sure you have completed [installation](install.md) before proceeding. -1. Install Megatron Core with required dependencies: +## Simple Training Example - ```bash - pip install --no-build-isolation megatron-core[mlm,dev] - ``` - -2. Clone repository for examples: - - ```bash - git clone https://github.com/NVIDIA/Megatron-LM.git - cd Megatron-LM - pip install --no-build-isolation .[mlm,dev] - ``` - -That's it! You're ready to start training. - -## Your First Training Run - -### Simple Training Example +Run a minimal distributed training loop with mock data on 2 GPUs: ```bash -# Distributed training example (2 GPUs, mock data) torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py ``` -### LLaMA-3 Training Example +## LLaMA-3 Training Example + +Train a LLaMA-3 8B model with FP8 precision on 8 GPUs using mock data: ```bash -# 8 GPUs, FP8 precision, mock data -./examples/llama/train_llama3_8b_fp8.sh +./examples/llama/train_llama3_8b_h100_fp8.sh ``` ## Data Preparation -### JSONL Data Format +To train on your own data, Megatron expects preprocessed binary files (`.bin` and `.idx`). + +### 1. Prepare a JSONL File + +Each line should contain a `text` field: ```json {"text": "Your training text here..."} {"text": "Another training sample..."} ``` -### Basic Preprocessing +### 2. Preprocess the Data ```bash python tools/preprocess_data.py \ diff --git a/docs/get-started/releasenotes.md b/docs/get-started/releasenotes.md index e2d77cf0070..e624de19f15 100644 --- a/docs/get-started/releasenotes.md +++ b/docs/get-started/releasenotes.md @@ -1,3 +1,12 @@ + + # Release Notes diff --git a/docs/index.md b/docs/index.md index 448a75e4c93..4b75ed2c0c8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,3 +1,12 @@ + + # Megatron Core User Guide **Megatron Core** is a GPU-optimized library for training large language models at scale. It provides modular, composable building blocks for creating custom training frameworks with state-of-the-art parallelism strategies and performance optimizations. @@ -6,13 +15,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} @@ -29,8 +38,8 @@ get-started/releasenotes :hidden: :caption: Get Started -get-started/quickstart get-started/install +get-started/quickstart ``` ```{toctree} @@ -62,6 +71,7 @@ user-guide/features/custom_fsdp user-guide/features/dist_optimizer user-guide/features/optimizer_cpu_offload user-guide/features/pipeline_parallel_layout +user-guide/features/fine_grained_activation_offloading user-guide/features/megatron_energon user-guide/features/megatron_rl user-guide/features/tokenizers @@ -81,16 +91,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/llama_mistral.md b/docs/llama_mistral.md index a79bb2c4bf9..95568adce78 100644 --- a/docs/llama_mistral.md +++ b/docs/llama_mistral.md @@ -1,8 +1,17 @@ + + # Llama, Mistral and other Llama-like model support in Megatron-LM NOTE: In order to simplify code we now only support converting llama-3.x and mistral checkpoints downloaded from Hugging Face. For converting other models, see [Megatron Bridge](models/index.md). -The [Llama-2](https://ai.meta.com/llama/) and [Llama-3.x](https://llama.meta.com/) family of models are an open-source set of pretrained & finetuned (for chat) models that have achieved strong results across a wide set of benchmarks. At their times of release, both Llama-2 and Llama-3 models achieved among the best results for open-source models, and were competitive with leading closed-source models (see https://arxiv.org/pdf/2307.09288.pdf and https://ai.meta.com/blog/meta-llama-3/). +The Llama-2 and Llama-3.x family of models are an open-source set of pretrained & finetuned (for chat) models that have achieved strong results across a wide set of benchmarks. At their times of release, both Llama-2 and Llama-3 models achieved among the best results for open-source models, and were competitive with leading closed-source models (see ). Similarly, [Mistral-7b](https://mistral.ai/news/announcing-mistral-7b/) is an open-source model with pretrained and finetuned (for chat) variants that achieve strong benchmark results. @@ -41,7 +50,6 @@ Architecturally Llama-2, Llama-3 and Mistral-7b are very similar. As such Megatr - [Known numerical differences](#known-numerical-differences) - [Using legacy model format](#using-legacy-model-format) - # Llama-2 Llama-2 checkpoints can be loaded into Megatron for inference and for finetuning. Loading these checkpoints consists of three steps: @@ -54,7 +62,7 @@ The following sections detail these steps. The final section lists benchmark res ## Download Meta or Huggingface checkpoints -Users must first apply for access to download the Llama-2 checkpoints either directly from [Meta](https://ai.meta.com/resources/models-and-libraries/llama-downloads/) or through [Huggingface](https://huggingface.co/docs/transformers/main/model_doc/llama2) (HF). The checkpoints are available in two formats, Meta's native format (available from both the Meta and HF links), and HF's format (available only from HF). Either format can be converted to Megatron, as detailed next. +Users must first apply for access to download the Llama-2 checkpoints either directly [Huggingface](https://huggingface.co/docs/transformers/main/model_doc/llama2) (HF). The checkpoints are available in two formats, Meta's native format (available from both the Meta and HF links), and HF's format (available only from HF). Either format can be converted to Megatron, as detailed next. ## Convert checkpoint format @@ -140,11 +148,11 @@ If loading for either inference or finetuning, use the following arguments: ### Launch Meta -Meta checkpoints can be launched with: https://github.com/facebookresearch/llama +Meta checkpoints can be launched with: ### Launch Huggingface -Huggingface checkpoints can be launched with: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py +Huggingface checkpoints can be launched with: ## Benchmark results @@ -352,7 +360,7 @@ The following sections detail these steps. ## Download Huggingface checkpoints -Users must first apply for access to download the Mistral-7b checkpoints through [Huggingface](https://huggingface.co/mistralai/Mistral-7B-v0.3) (HF). +Users must first apply for access to download the Mistral-7b checkpoints through Huggingface. Two variants are available: the base model ([Mistral-7B-v0.3](https://huggingface.co/mistralai/Mistral-7B-v0.3)) and the instruct model ([Mistral-7B-Instruct-v0.3](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3)). ## Convert checkpoint format @@ -428,7 +436,7 @@ Many models such as Yi-34B and Qwen2.x use the Llama architecture and may be con It is not expected that the megatron and Huggingface implementations of llama3.x and mistral models will produce numerically identical results. There are multiple points where small numerical differences are expected. This is a non-exhaustive list: -1. TransformerEngine (TE) uses the model params_dtype inside RMSNorm whereas the Huggingface implementation uses fp32. See for details: https://github.com/NVIDIA/TransformerEngine/issues/1132 +1. TransformerEngine (TE) uses the model params_dtype inside RMSNorm whereas the Huggingface implementation uses fp32. See for details: 2. Huggingface `transformers` implements the q, k and v projections in self-attention as separate GEMMs whereas Megatron core combines them into a single GEMM for efficiency. This leads to small numerical differences. # Using legacy model format diff --git a/docs/models/index.md b/docs/models/index.md index 6fabd1f582c..0ee379b01bd 100644 --- a/docs/models/index.md +++ b/docs/models/index.md @@ -1,3 +1,12 @@ + + # Supported Models Megatron Core supports a wide range of language and multimodal models with optimized implementations for large-scale training. diff --git a/docs/models/llms.md b/docs/models/llms.md index 6789a4c551c..f649673a2cc 100644 --- a/docs/models/llms.md +++ b/docs/models/llms.md @@ -1,3 +1,12 @@ + + # Language Models Megatron Core supports the following language model architectures for large-scale training. diff --git a/docs/models/multimodal.md b/docs/models/multimodal.md index 66ed8ccd9cb..dce977e261d 100644 --- a/docs/models/multimodal.md +++ b/docs/models/multimodal.md @@ -1,3 +1,12 @@ + + # Multimodal Models Megatron Core supports multimodal models that combine language with vision, audio, and other modalities for comprehensive multimodal understanding. diff --git a/docs/project.json b/docs/project.json index aa547ddd298..d5b9535338b 100644 --- a/docs/project.json +++ b/docs/project.json @@ -1,2 +1,2 @@ -{"name": "megatron-lm", "version": "latest"} +{"name": "megatron-lm", "version": "nightly"} diff --git a/docs/user-guide/data-preparation.md b/docs/user-guide/data-preparation.md index 3ff5eedba89..ea91bee4309 100644 --- a/docs/user-guide/data-preparation.md +++ b/docs/user-guide/data-preparation.md @@ -1,3 +1,12 @@ + + # Data Preparation Preparing your data correctly is essential for successful training with Megatron Core. @@ -37,6 +46,46 @@ python tools/preprocess_data.py \ | `--workers` | Number of parallel workers for processing | | `--append-eod` | Add end-of-document token | +## Finding Optimal Number of Workers + +Use the `--find-optimal-num-workers` flag to find number of workers which gives the best performance in terms of preprocessed documents per second. +Script will lauch a few short data preprocessing runs with a different number of workers to define the fastest run in respect to collected performance data. + +```bash +python tools/preprocess_data.py \ + --input data.jsonl \ + --output-prefix processed_data \ + --tokenizer-type HuggingFaceTokenizer \ + --tokenizer-model /path/to/tokenizer.model \ + --workers 8 \ + --find-optimal-num-workers \ + --workers-to-check 4 8 16 32 \ + --max-documents 50000 +``` + +**Required arguments** + +| Argument | Description | +|----------|-------------| +| `--find-optimal-num-workers` | Activates search of optimal number of workers | +| `--workers-to-check` | List of possible number of workers to run | +| `--max-documents` | Number of documents to be preprocessed during each run | + +**Output example** + +```bash +----------------------------------- +Performance results (fastest → slowest): +1. 16 workers → avg. docs/s: 9606.6476 +2. 32 workers → avg. docs/s: 9275.3284 +3. 8 workers → avg. docs/s: 9151.9280 +4. 4 workers → avg. docs/s: 6391.3819 + +----------------------------------- +The most optimal num of workers is 16 with avg. preprocessed docs/s: 9606.6476. +----------------------------------- +``` + ## Output Files The preprocessing tool generates two files: diff --git a/docs/user-guide/features/context_parallel.md b/docs/user-guide/features/context_parallel.md index 841c16326b3..c44366187be 100644 --- a/docs/user-guide/features/context_parallel.md +++ b/docs/user-guide/features/context_parallel.md @@ -1,3 +1,12 @@ + + # context_parallel package ## Context parallelism overview diff --git a/docs/user-guide/features/custom_fsdp.md b/docs/user-guide/features/custom_fsdp.md index 2f81eb0c5ef..ab1a1efc402 100644 --- a/docs/user-guide/features/custom_fsdp.md +++ b/docs/user-guide/features/custom_fsdp.md @@ -1,3 +1,12 @@ + + # Megatron FSDP **NOTE: In M-Core 0.14, the custom FSDP refactored its checkpoint implementation to use DTensor-based torch distributed checkpointing. The custom FSDP was also renamed Megatron FSDP. The relevant sections of this document are no longer applicable.** diff --git a/docs/user-guide/features/dist_optimizer.md b/docs/user-guide/features/dist_optimizer.md index ddb6079885c..4e47791c12f 100644 --- a/docs/user-guide/features/dist_optimizer.md +++ b/docs/user-guide/features/dist_optimizer.md @@ -1,3 +1,12 @@ + + # Distributed Optimizer The motivation for the distributed optimizer is to save memory by distributing the optimizer state evenly across data parallel ranks (https://arxiv.org/abs/1910.02054), versus the naive method of replicating the optimizer state across data parallel ranks. diff --git a/docs/user-guide/features/fine_grained_activation_offloading.md b/docs/user-guide/features/fine_grained_activation_offloading.md index 53211d1d06c..494674bd4f0 100644 --- a/docs/user-guide/features/fine_grained_activation_offloading.md +++ b/docs/user-guide/features/fine_grained_activation_offloading.md @@ -1,3 +1,12 @@ + + # Fine-grained Activation Offloading (collaborated with rednote) Memory capacity is more and more important with the rising of extreme sparse MoE models like DeepSeek-V3 and Qwen3-235B. Fine-grained recomputing reduces the memory footprint at the cost of extra recomputation, while offloading could utilize the host-device bandwidth to achieve nearly zero-overhead. Fine-grained Activation Offloading targets at offloading the activation at the granularity of specific modules, so that we can calibrate the amount of offloading activation to maximize the training throughput. diff --git a/docs/user-guide/features/index.md b/docs/user-guide/features/index.md index 7730443e91b..59cef95d574 100644 --- a/docs/user-guide/features/index.md +++ b/docs/user-guide/features/index.md @@ -1,3 +1,12 @@ + + # Advanced Features Advanced feature guides for key Megatron Core capabilities. @@ -5,6 +14,7 @@ Advanced feature guides for key Megatron Core capabilities. ```{toctree} :maxdepth: 2 +fine_grained_activation_offloading moe context_parallel custom_fsdp diff --git a/docs/user-guide/features/megatron_energon.md b/docs/user-guide/features/megatron_energon.md index d08bde21e38..9ebba72083a 100644 --- a/docs/user-guide/features/megatron_energon.md +++ b/docs/user-guide/features/megatron_energon.md @@ -1,3 +1,12 @@ + + # Megatron Energon Advanced multimodal dataloader for efficient loading of text, images, video, and audio at scale. diff --git a/docs/user-guide/features/megatron_rl.md b/docs/user-guide/features/megatron_rl.md index 128b41bdaf5..653ecb92459 100644 --- a/docs/user-guide/features/megatron_rl.md +++ b/docs/user-guide/features/megatron_rl.md @@ -1,3 +1,12 @@ + + # Megatron RL Reinforcement learning library for post-training large language models at scale. diff --git a/docs/user-guide/features/moe.md b/docs/user-guide/features/moe.md index 56aca8c6999..45efe291a21 100644 --- a/docs/user-guide/features/moe.md +++ b/docs/user-guide/features/moe.md @@ -1,3 +1,12 @@ + + # Mixture of Experts ```{toctree} @@ -6,6 +15,7 @@ multi_token_prediction multi_latent_attention +../../api-guide/router_replay ``` ```{include} ../../../megatron/core/transformer/moe/README.md diff --git a/docs/user-guide/features/multi_latent_attention.md b/docs/user-guide/features/multi_latent_attention.md index 5628b8cfee3..4310843557a 100644 --- a/docs/user-guide/features/multi_latent_attention.md +++ b/docs/user-guide/features/multi_latent_attention.md @@ -1,3 +1,12 @@ + + # Multi-Latent Attention ## Multi-Latent Attention overview diff --git a/docs/user-guide/features/multi_token_prediction.md b/docs/user-guide/features/multi_token_prediction.md index 891bf4c93c5..e16108bbcfa 100644 --- a/docs/user-guide/features/multi_token_prediction.md +++ b/docs/user-guide/features/multi_token_prediction.md @@ -1,3 +1,12 @@ + + # Multi-Token Prediction (MTP) Multi-Token Prediction (MTP) extends the prediction scope to multiple future tokens at each position. On the one hand, an MTP objective densifies the training signals and may improve @@ -7,7 +16,7 @@ data efficiency. On the other hand, MTP may enable the model to pre-plan its rep The k-th MTP module consists of a shared embedding layer, a projection matrix, a Transformer block, and a shared output head. For the i-th input token at the (k - 1)-th prediction depth, we first combine the representation of the i-th token and the embedding of the (i + K)-th token with the linear projection. The combined serves as the input of the Transformer block at the k-th depth to produce the output representation. -For more information, please refer to [DeepSeek-V3 Technical Report](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf) +For more information, refer to [DeepSeek-V3 Technical Report](https://arxiv.org/pdf/2412.19437.pdf) ## Related Arguments @@ -45,4 +54,4 @@ Use `m` to represent MTP layers in the pipeline layout string. For example: ## Precautions -Please do not use Context Parallel (CP), or arbitrary AttnMaskType, or learned absolute position embedding type with MTP. These use cases are not yet supported. +Do not use Context Parallel (CP), or arbitrary AttnMaskType, or learned absolute position embedding type with MTP. These use cases are not yet supported. diff --git a/docs/user-guide/features/optimizer_cpu_offload.md b/docs/user-guide/features/optimizer_cpu_offload.md index 408d7f6a788..1496bd0a91e 100644 --- a/docs/user-guide/features/optimizer_cpu_offload.md +++ b/docs/user-guide/features/optimizer_cpu_offload.md @@ -1,3 +1,12 @@ + + # Optimizer CPU Offload ```{include} ../../../megatron/core/optimizer/cpu_offloading/README.md diff --git a/docs/user-guide/features/pipeline_parallel_layout.md b/docs/user-guide/features/pipeline_parallel_layout.md index 30c8ce1a500..96b00eca004 100644 --- a/docs/user-guide/features/pipeline_parallel_layout.md +++ b/docs/user-guide/features/pipeline_parallel_layout.md @@ -1,3 +1,12 @@ + + # Custom Pipeline Model Parallel Layout *This is an experimental feature and may be changed.* diff --git a/docs/user-guide/features/tokenizers.md b/docs/user-guide/features/tokenizers.md index 0aecf8df8a7..bc1a47cec76 100644 --- a/docs/user-guide/features/tokenizers.md +++ b/docs/user-guide/features/tokenizers.md @@ -1,3 +1,12 @@ + + # Tokenizers Megatron Core provides a unified tokenizer system with a HuggingFace-style API for easy tokenizer management and configuration. @@ -141,7 +150,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 +182,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/docs/user-guide/index.md b/docs/user-guide/index.md index bbe85451582..d12f3e35af2 100644 --- a/docs/user-guide/index.md +++ b/docs/user-guide/index.md @@ -1,3 +1,16 @@ +--- +orphan: true +--- + + + # User Guide Comprehensive guides for using Megatron Core and Megatron-LM. @@ -5,7 +18,6 @@ Comprehensive guides for using Megatron Core and Megatron-LM. ```{toctree} :maxdepth: 2 -quickstart msc_integration data-preparation training-examples diff --git a/docs/user-guide/msc_integration.md b/docs/user-guide/msc_integration.md index fd73ac7e8f4..a197f25afc1 100644 --- a/docs/user-guide/msc_integration.md +++ b/docs/user-guide/msc_integration.md @@ -1,3 +1,12 @@ + + ```{include} ../../megatron/core/MSC_Integration.md ``` diff --git a/docs/user-guide/parallelism-guide.md b/docs/user-guide/parallelism-guide.md index 2baf518ae85..8d5cb8ff7c3 100644 --- a/docs/user-guide/parallelism-guide.md +++ b/docs/user-guide/parallelism-guide.md @@ -1,3 +1,12 @@ + + # Parallelism Strategies Guide Megatron Core supports multiple parallelism strategies that can be combined to efficiently train models from billions to trillions of parameters across thousands of GPUs. diff --git a/docs/user-guide/quickstart.md b/docs/user-guide/quickstart.md deleted file mode 100644 index 7baed06d6be..00000000000 --- a/docs/user-guide/quickstart.md +++ /dev/null @@ -1,3 +0,0 @@ -```{include} ../../megatron/core/QuickStart.md -``` - diff --git a/docs/user-guide/training-examples.md b/docs/user-guide/training-examples.md index 2824c608c36..5e7c0440073 100644 --- a/docs/user-guide/training-examples.md +++ b/docs/user-guide/training-examples.md @@ -1,3 +1,12 @@ + + # Training Examples Get started with Megatron Core training using these practical examples. @@ -24,7 +33,7 @@ This example: Train LLaMA-3 8B model with FP8 mixed precision on 8 GPUs: ```bash -./examples/llama/train_llama3_8b_fp8.sh +./examples/llama/train_llama3_8b_h100_fp8.sh ``` **Configuration:** diff --git a/docs/versions1.json b/docs/versions1.json index a524c5921a8..d55416e0a95 100644 --- a/docs/versions1.json +++ b/docs/versions1.json @@ -1,13 +1,17 @@ [ { - "name": "latest", - "version": "latest", - "url": "https://docs.nvidia.com/megatron-core/latest/" + "name": "nightly", + "version": "nightly", + "url": "https://docs.nvidia.com/megatron-core/developer-guide/nightly/" + }, + { + "name": "0.16.0 (latest)", + "version": "0.16.0", + "url": "https://docs.nvidia.com/megatron-core/developer-guide/latest/" }, { "name": "0.15.0", "version": "0.15.0", - "url": "https://docs.nvidia.com/megatron-core/0.15.0/" + "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.15.0/" } ] - diff --git a/examples/gpt3/gpt_config.yaml b/examples/gpt3/gpt_config.yaml index 18d305d9cb1..600f50221ce 100644 --- a/examples/gpt3/gpt_config.yaml +++ b/examples/gpt3/gpt_config.yaml @@ -92,7 +92,6 @@ model_parallel: # Optimizations gradient_accumulation_fusion: True - async_tensor_model_parallel_allreduce: True tp_comm_overlap: False # Debug Options diff --git a/examples/gptoss/01_convert_from_hf.py b/examples/gptoss/01_convert_from_hf.py new file mode 100644 index 00000000000..adee3358ec3 --- /dev/null +++ b/examples/gptoss/01_convert_from_hf.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Convert HuggingFace checkpoints to Megatron format.""" + +import os +import argparse + +from megatron.bridge import AutoBridge + +def _parse_args(): + parser = argparse.ArgumentParser(description="Convert HF LLMs to Megatron format") + parser.add_argument( + "--hf-model", + type=str, + required=True, + help="HuggingFace model identifier or path", + ) + parser.add_argument( + "--save-path", + type=str, + default=None, + help="Path to save the converted Megatron checkpoint", + ) + parser.add_argument('--local-rank', '--local_rank', type=int, default=0) + return parser.parse_args() + +if __name__ == "__main__": + args = _parse_args() + HF_MODEL = args.hf_model + SAVE_PATH = args.save_path + WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1)) + + if SAVE_PATH is None: + SAVE_PATH = f"./megatron_checkpoints/{HF_MODEL.replace('/', '_')}" + + print(f"Converting {HF_MODEL} to Megatron format...") + print(f"Save path: {SAVE_PATH}") + + bridge = AutoBridge.from_hf_pretrained(HF_MODEL, trust_remote_code=True) + provider = bridge.to_megatron_provider() + # Update these configs as needed + provider.expert_tensor_parallel_size = 1 + provider.tensor_model_parallel_size = 1 + provider.pipeline_model_parallel_size = WORLD_SIZE + provider.finalize() + + model = provider.provide_distributed_model(wrap_with_ddp=False) + + bridge.save_megatron_model( + model, + SAVE_PATH, + hf_tokenizer_path=HF_MODEL + ) + + print(f"Saved Megatron checkpoint to {SAVE_PATH}") diff --git a/examples/gptoss/02_train.sh b/examples/gptoss/02_train.sh new file mode 100755 index 00000000000..d129adc2b84 --- /dev/null +++ b/examples/gptoss/02_train.sh @@ -0,0 +1,259 @@ +#!/bin/bash + +export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} + + +# Setup arguments with defaults +CHECKPOINT_PATH="NO_VALUE_PROVIDED" +TENSORBOARD_LOGS_PATH="./tensorboard_logs/" +TOKENIZER_ARG="MOCK" +DATA_ARG="MOCK" +DISTRIBUTED_CONFIG_FILE="" + +# Parse command line arguments +while [[ $# -gt 0 ]]; do + case $1 in + --checkpoint-path) + CHECKPOINT_PATH="$2" + shift 2 + ;; + --tensorboard-logs-path) + TENSORBOARD_LOGS_PATH="$2" + shift 2 + ;; + --tokenizer) + TOKENIZER_ARG="$2" + shift 2 + ;; + --data) + DATA_ARG="$2" + shift 2 + ;; + --distributed-config-file) + DISTRIBUTED_CONFIG_FILE="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 [OPTIONS]" + echo "Options:" + echo " --checkpoint-path PATH Path to Megatron checkpoint" + echo " --tensorboard-logs-path PATH Path to TensorBoard logs" + echo " --tokenizer PATH|MOCK Path to tokenizer model, or 'MOCK' (default: MOCK)" + echo " --data PATH|MOCK Data prefix, or 'MOCK' (default: MOCK)" + echo " --distributed-config-file FILE Path to distributed training config file" + echo " -h, --help Show this help message" + exit 0 + ;; + *) + echo "Unknown option: $1" + echo "Use --help for usage information" + exit 1 + ;; + esac +done + +# Check if checkpoint path exists +if [ ! -d "$CHECKPOINT_PATH" ]; then + echo "Error: Checkpoint path does not exist: $CHECKPOINT_PATH" + exit 1 +fi +echo "Checkpoint path exists: $CHECKPOINT_PATH" + +# Check if tensorboard logs path exists +if [ ! -d "$TENSORBOARD_LOGS_PATH" ]; then + echo "Warning: TensorBoard logs path does not exist. Creating: $TENSORBOARD_LOGS_PATH" + mkdir -p "$TENSORBOARD_LOGS_PATH" +fi +echo "TensorBoard logs path exists: $TENSORBOARD_LOGS_PATH" + +# NOTE: by default we use 8 GPUs +# These values will be over-written below with environmental variables +GPUS_PER_NODE=8 +NUM_NODES=1 +MASTER_ADDR="localhost" +MASTER_PORT=6000 +NODE_RANK=0 + +# Load distributed config from file if provided +if [ -n "$DISTRIBUTED_CONFIG_FILE" ]; then + if [ ! -f "$DISTRIBUTED_CONFIG_FILE" ]; then + echo "Warning: Distributed config file does not exist: $DISTRIBUTED_CONFIG_FILE" + echo "Continuing with default distributed training settings." + else + echo "Loading distributed config from: $DISTRIBUTED_CONFIG_FILE" + source "$DISTRIBUTED_CONFIG_FILE" + fi +fi + +# Override with environment variables if set +GPUS_PER_NODE=${GPUS_PER_NODE:-8} +NUM_NODES=${NUM_NODES:-1} +MASTER_ADDR=${MASTER_ADDR:-localhost} +MASTER_PORT=${MASTER_PORT:-6000} +NODE_RANK=${NODE_RANK:-0} +WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES)) + +# Path to the pretrain_gpt.py script, assuming this script is run from the root of the Megatron-LM repository +PRETRAIN_SCRIPT_PATH="pretrain_gpt.py" + +# Data cache path (useful for both mock and real data) +DATA_CACHE_PATH="${PWD}/benchmark_cache_gpt_oss_20b" +mkdir -p "$DATA_CACHE_PATH" + +DISTRIBUTED_ARGS=( + --nproc_per_node $GPUS_PER_NODE + --nnodes $NUM_NODES + --master_addr $MASTER_ADDR + --master_port $MASTER_PORT + --node_rank $NODE_RANK +) + +# NOTE: we only set pipeline parallelism to be the number of GPUs +# Adjust each value based on your setup. +TP_SIZE=1 +EP_SIZE=1 +PP_SIZE=${WORLD_SIZE} +MICRO_BATCH_SIZE=1 +GLOBAL_BATCH_SIZE=128 +NUM_LAYERS=12 +DTYPE="fp8" +SEQ_LENGTH=8192 +MAX_POSITION_EMBEDDINGS=8192 +TRAIN_SAMPLES=1953125000 +LR_DECAY_SAMPLES=1949218748 + +MODEL_ARGS=( + --no-masked-softmax-fusion + --transformer-impl transformer_engine + --disable-bias-linear + --untie-embeddings-and-output-weights + --no-rope-fusion + --normalization RMSNorm + --num-layers ${NUM_LAYERS} + --hidden-size 512 + --ffn-hidden-size 2048 + --num-attention-heads 64 + --group-query-attention + --num-query-groups 8 + --seq-length ${SEQ_LENGTH} + --max-position-embeddings ${MAX_POSITION_EMBEDDINGS} + --use-mcore-models + --rotary-percent 1.0 + --rope-type rope + --position-embedding-type rope + --rotary-base 10000 + --no-bias-gelu-fusion + --export-force-local-attention + --no-bias-dropout-fusion + --quick-geglu + --glu-linear-offset 1.0 + --softmax-type learnable + --window-attn-skip-freq 2 + --activation-func-clamp-value 7.0 + --window-size 127,0 + --enable-gpt-oss +) + +MOE_ARGS=( + --num-experts 4 + --moe-router-topk 2 + --moe-router-load-balancing-type aux_loss + --moe-aux-loss-coeff 1e-3 + --moe-grouped-gemm + --moe-token-dispatcher-type alltoall + --overlap-param-gather + --overlap-grad-reduce + --moe-ffn-hidden-size 2048 + --moe-router-dtype fp32 + --moe-z-loss-coeff 1e-3 + --moe-permute-fusion +) + +DATA_ARGS_LIST=() +if [[ "$TOKENIZER_ARG" == "MOCK" ]] || [[ "$DATA_ARG" == "MOCK" ]] || [[ -z "$TOKENIZER_ARG" ]]; then + DATA_ARGS_LIST+=( + "--mock-data" + "--tokenizer-type NullTokenizer" + "--vocab-size 128256" + "--data-cache-path ${DATA_CACHE_PATH}" + "--tiktoken-pattern v2" + "--split '99,1,0'" + "--no-create-attention-mask-in-dataloader" + "--no-mmap-bin-files" + "--num-workers 1" + ) +else + # Settings for real data + DATA_ARGS_LIST+=( + "--data-path $DATA_ARG" + "--tokenizer-type HuggingFaceTokenizer" + "--tokenizer-model $TOKENIZER_ARG" + "--data-cache-path ${DATA_CACHE_PATH}" + "--split '99,1,0'" + "--no-create-attention-mask-in-dataloader" + "--no-mmap-bin-files" + "--num-workers 1" + # Note: --vocab-size might be inferred by HuggingFaceTokenizer or might need to be explicit. + "--vocab-size 128256" + ) +fi + +TRAINING_ARGS=( + --micro-batch-size ${MICRO_BATCH_SIZE} + --global-batch-size ${GLOBAL_BATCH_SIZE} + --lr 1.0e-5 + --train-samples ${TRAIN_SAMPLES} + --lr-decay-samples ${LR_DECAY_SAMPLES} + --lr-decay-style cosine + --min-lr 1.0e-6 + --weight-decay 0.1 + --lr-warmup-fraction 0.05 + --clip-grad 1.0 + --bf16 + --use-flash-attn + --attention-softmax-in-fp32 + --accumulate-allreduce-grads-in-fp32 + --disable-bf16-reduced-precision-matmul + --recompute-activations +) + +MODEL_PARALLEL_ARGS=( + --tensor-model-parallel-size ${TP_SIZE} + --pipeline-model-parallel-size ${PP_SIZE} + --expert-model-parallel-size ${EP_SIZE} + --sequence-parallel + --context-parallel-size 1 + --use-distributed-optimizer + --fp8-format hybrid + --fp8-param-gather + --fp8-amax-compute-algo max + --fp8-amax-history-len 1024 +) + +LOGGING_ARGS=( + --log-interval 1 + --save-interval 10000 + --eval-interval 50000000 + --eval-iters 0 + --save $CHECKPOINT_PATH + --tensorboard-dir "${CHECKPOINT_PATH}/tensorboard" + --moe-per-layer-logging + --no-load-optim + --no-load-rng + --log-throughput +) + +# Ensure pretrain_gpt.py is found +if [ ! -f "$PRETRAIN_SCRIPT_PATH" ]; then + echo "Error: pretrain_gpt.py not found at $PRETRAIN_SCRIPT_PATH" + echo "Please ensure you are running this script from the root of the Megatron-LM repository, and pretrain_gpt.py is present." + exit 1 +fi + +python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} ${PRETRAIN_SCRIPT_PATH} \ + ${MODEL_ARGS[@]} \ + ${MOE_ARGS[@]} \ + ${DATA_ARGS_LIST[@]} \ + ${TRAINING_ARGS[@]} \ + ${MODEL_PARALLEL_ARGS[@]} \ + ${LOGGING_ARGS[@]} diff --git a/examples/gptoss/03_convert_to_hf.py b/examples/gptoss/03_convert_to_hf.py new file mode 100644 index 00000000000..8089afec854 --- /dev/null +++ b/examples/gptoss/03_convert_to_hf.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Convert HuggingFace checkpoints to Megatron format.""" + +import os +import argparse + +from megatron.bridge import AutoBridge + +def _parse_args(): + parser = argparse.ArgumentParser(description="Convert Megatron LLMs to HuggingFace format") + parser.add_argument( + "--hf-model", + type=str, + required=True, + help="HuggingFace model identifier or path to load config from", + ) + parser.add_argument( + "--megatron-model", + type=str, + required=True, + help="Megatron model identifier or path", + ) + parser.add_argument( + "--save-path", + type=str, + default=None, + help="Path to save the converted HuggingFace checkpoint", + ) + parser.add_argument('--local-rank', '--local_rank', type=int, default=0) + return parser.parse_args() + +if __name__ == "__main__": + args = _parse_args() + HF_MODEL = args.hf_model + MEGATRON_MODEL = args.megatron_model + SAVE_PATH = args.save_path + WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1)) + + if SAVE_PATH is None: + SAVE_PATH = f"./huggingface_checkpoints/{MEGATRON_MODEL.replace('/', '_')}" + + print(f"Converting {MEGATRON_MODEL} to HuggingFace {HF_MODEL} format...") + print(f"Save path: {SAVE_PATH}") + + bridge = AutoBridge.from_hf_pretrained(HF_MODEL, trust_remote_code=True) + bridge.export_ckpt( + MEGATRON_MODEL, + SAVE_PATH, + ) + + print(f"Saved HuggingFace checkpoint to {SAVE_PATH}") diff --git a/examples/gptoss/README.md b/examples/gptoss/README.md new file mode 100644 index 00000000000..eeb92ad9953 --- /dev/null +++ b/examples/gptoss/README.md @@ -0,0 +1,153 @@ +# GPT-OSS Training Tutorial + +## Step 0: Install Dependencies + +### Using Megatron Bridge + +[Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) + +Megatron Bridge provides a quick and convenient way to convert HuggingFace checkpoints to the Megatron format used by Megatron-LM. Follow the instructions in the [Megatron-Bridge Installation](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/README.md#-installation) to run the nemo docker container and convert checkpoints (via mounted volumes - make sure that the huggingface cache location AND the megatron checkpoint locations are properly mounted, otherwise you may not be saving the converted model to disk correctly). + +Below is an example of how to use Megatron-Bridge inside the pytorch container to convert a HuggingFace model checkpoint to Megatron format. + +Reference: [Megatron-Bridge Dockerfile](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docker/Dockerfile.ci) + +Inside the [pytorch container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) run the following commands to install Megatron-Bridge: +```bash +cd /opt +git clone --recursive https://github.com/NVIDIA-NeMo/Megatron-Bridge.git +cd Megatron-Bridge + +# Make sure submodules are initialized (for 3rdparty/Megatron-LM) +git submodule update --init --recursive + +export PATH="/root/.local/bin:$PATH" +export UV_PROJECT_ENVIRONMENT=/opt/venv +export VIRTUAL_ENV=/opt/venv +export PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH" +export UV_LINK_MODE=copy +export UV_VERSION="0.7.2" + +# Install UV +curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh + +# Create virtual environment and build the package +uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages + +uv sync --locked --only-group build +uv sync --locked --link-mode copy --all-extras --all-groups + +uv pip install --no-deps -e . + +source ${UV_PROJECT_ENVIRONMENT}/bin/activate +``` + +### Setup Environment + +```bash +export HOST_MEGATRON_LM_DIR="/path/to/your/host/megatron-lm" +git clone https://github.com/NVIDIA/Megatron-LM.git "$HOST_MEGATRON_LM_DIR" +cd "$HOST_MEGATRON_LM_DIR" +``` + +```bash +export HF_TOKEN={your_hf_token_here} +``` + +## Step 1: Convert HuggingFace to Megatron (Optional - skip if you already have a Megatron checkpoint) + +Set `--nproc-per-node` to be the number of GPUs per node. Set `hf_model_name` to be the Huggingface model e.g. `openai/gpt-oss-20b` + +```bash +python3 -m torch.distributed.launch --nproc-per-node=8 examples/gptoss/01_convert_from_hf.py --hf-model openai/gpt-oss-20b +``` + +## Step 2: Train from Scratch + +To train from scratch first follow the steps below to setup the environment appropriately before running the training script in docker. Even though we are running the same container as before, it is better to restart the container to ensure a clean environment and that all environment and docker variables are set correctly. For the following example we used 8x GB300, but you should change the number of GPUs and nodes as needed. + +### Setup Environment + +```bash +# Change these based on model and directory from previous conversion step +export MODEL_DIR_NAME="openai_gpt-oss_20b" + +export HOST_CHECKPOINT_PATH="./megatron_checkpoints/${MODEL_DIR_NAME}" +export HOST_TENSORBOARD_LOGS_PATH="./tensorboard_logs/${MODEL_DIR_NAME}" +``` + +By default we will use mock data to train the model in the example below. To use your own data, set the following environment variables: + +```bash +# Optional: For real data +export HOST_TOKENIZER_MODEL_PATH="/path/to/host/tokenizer.model" +export HOST_DATA_PREFIX="/path/to/host/mydata_prefix" +``` + +### Setup Training Configurations + +Run the following to create a `distributed_config.env` file with the appropriate distributed training configurations. Change the values as needed for your setup. This file will override the default values in `02_train.sh`. + +```bash +cat > ./distributed_config.env << 'EOF' +GPUS_PER_NODE=8 +NUM_NODES=1 +MASTER_ADDR=localhost +MASTER_PORT=6000 +NODE_RANK=0 +EOF +``` + +### Run Container with Mounted Volumes + +**NOTE:** This container runs the example training script `02_train.sh` located in the `examples/gptoss` directory. By default, we have only set pipeline parallelism to be the number of GPUs. Adjust TP_SIZE, EP_SIZE, PP_SIZE, etc. in `02_train.sh`. You can also adjust modify `--hidden-size`, `--ffn-hidden-size`, `--num-attention-heads`, `NUM_LAYERS`, etc. + +To train using mock data, run the following command: +```bash +PYTORCH_IMAGE="nvcr.io/nvidia/pytorch:25.12-py3" + +docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \ + -v "${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm" \ + -v "${HOST_CHECKPOINT_PATH}:/workspace/checkpoints" \ + -v "${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs" \ + -v "./distributed_config.env:/workspace/megatron-lm/examples/gptoss/distributed_config.env" \ + --workdir /workspace/megatron-lm \ + $PYTORCH_IMAGE \ + bash examples/gptoss/02_train.sh \ + --checkpoint-path /workspace/checkpoints \ + --tensorboard-logs-path /workspace/tensorboard_logs \ + --distributed-config-file /workspace/megatron-lm/examples/gptoss/distributed_config.env \ + 2>&1 | tee "${HOST_TENSORBOARD_LOGS_PATH}/training_mock_$(date +'%y-%m-%d_%H-%M-%S').log" +``` +**Note:** If you run into issues generating mock data one solution might be to reduce the number of GPUs to 1 and try to generate the data again. + +If using real data with with the `HOST_TOKENIZER_MODEL_PATH` and `HOST_DATA_PREFIX` environment variables set, run the following command instead: + +```bash +PYTORCH_IMAGE="nvcr.io/nvidia/pytorch:25.12-py3" + +docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \ + -v "${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm" \ + -v "${HOST_CHECKPOINT_PATH}:/workspace/checkpoints" \ + -v "${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs" \ + -v "${HOST_TOKENIZER_MODEL_PATH}:/workspace/tokenizer_model" \ + -v "$(dirname "${HOST_DATA_PREFIX}"):/workspace/data_dir" \ + -v "./distributed_config.env:/workspace/megatron-lm/examples/gptoss/distributed_config.env" \ + --workdir /workspace/megatron-lm \ + $PYTORCH_IMAGE \ + bash examples/gptoss/02_train.sh \ + --checkpoint-path /workspace/checkpoints \ + --tensorboard-logs-path /workspace/tensorboard_logs \ + --tokenizer /workspace/tokenizer_model \ + --data "/workspace/data_dir/$(basename "${HOST_DATA_PREFIX}")" \ + --distributed-config-file /workspace/megatron-lm/examples/gptoss/distributed_config.env \ + 2>&1 | tee "${HOST_TENSORBOARD_LOGS_PATH}/training_custom_$(date +'%y-%m-%d_%H-%M-%S').log" +``` + +## Step 3: Convert Megatron to HuggingFace + +Just run the following command to change from the megatron checkpoint from training to the huggingface format to share with others (make sure you have the same virtual environment setup as in Step 0): + +```bash +python3 -m torch.distributed.launch --nproc-per-node=8 examples/gptoss/03_convert_to_hf.py --hf-model openai/gpt-oss-20b --megatron-model ./megatron_checkpoints/openai_gpt-oss_20b +``` \ No newline at end of file diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 88b744b3ac0..f02aae9c221 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, @@ -257,6 +80,16 @@ def run_inference( args = get_args() + # Parse batch boundaries for batch-drain mode. + batch_ranges = None + if args.drain_between_batches and args.batch_boundaries: + boundaries = [int(x) for x in args.batch_boundaries.split(",")] + num_requests_total = len(requests) + batch_ranges = [] + for i, start in enumerate(boundaries): + end = boundaries[i + 1] if i + 1 < len(boundaries) else num_requests_total + batch_ranges.append((start, end)) + # Initialize request arrival times. base_arrival_time = get_curr_time() for request in requests: @@ -284,72 +117,24 @@ 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 tbar.update(1) - while True: - # Add requests. - add_start = get_curr_time() - if args.incoming_requests_per_step is None: - # Add requests with 'earlier' arrival time. - while num_requests_added < num_requests_total: - if requests[num_requests_added].time_arrival > add_start: - break - _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, - )): - _add_request() - add_times.append(get_curr_time() - add_start) - - # Step inference engine (i.e., generate a token for each active request). - # Before step, we haven't done the scheduling, so we cannot know the is_decode_only - try: - result = engine.step_modern() - except EngineSuspendedError as e: - result = e - 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 - - # Test suspending and resuming engine. - if args.suspend_resume_interval is not None: - - # Suspend. - if attempted_step_count % args.suspend_resume_interval == 0: - print("**** step %d/%d ... suspend." % (engine.step_count, attempted_step_count)) - engine.suspend() - - # Resume, 0+ attempted steps later. - if ( - attempted_step_count > 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() - - # If engine suspended, continue to next iter. - if isinstance(result, EngineSuspendedError): - continue + def _process_step_result(result): + """Process a single engine step result, updating bookkeeping state.""" + nonlocal total_output_tokens, num_requests_finished + + is_decode_only = engine.is_decode_only # 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 +159,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 @@ -399,47 +186,118 @@ def _add_request(): num_requests_finished += 1 output_times.append(get_curr_time() - output_start) - # Check if all requests are finished. - if not (engine.has_unfinished_requests() or num_requests_added < num_requests_total): - break + if batch_ranges is not None: + # Batch-drain mode: add all requests in a batch, drain, then next batch. + for batch_idx, (batch_start, batch_end) in enumerate(batch_ranges): + # Add all requests in current batch. + add_start = get_curr_time() + while num_requests_added < batch_end: + _add_request() + add_times.append(get_curr_time() - add_start) + + # Step until all active requests finish (drain). + while engine.has_unfinished_requests(): + try: + result = engine.step_modern() + except EngineSuspendedError as e: + result = e + attempted_step_count += 1 + + if isinstance(result, EngineSuspendedError): + continue + + _process_step_result(result) + else: + # Original mode: add requests per step based on arrival time or count. + while True: + # Add requests. + add_start = get_curr_time() + if args.incoming_requests_per_step is None: + # Add requests with 'earlier' arrival time. + while num_requests_added < num_requests_total: + if requests[num_requests_added].time_arrival > add_start: + break + _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) + ): + _add_request() + add_times.append(get_curr_time() - add_start) + + # Step inference engine (i.e., generate a token for each active request). + # Before step, we haven't done the scheduling, so we cannot know the is_decode_only + try: + result = engine.step_modern() + except EngineSuspendedError as e: + result = e + pass # ignore error in order to call 'engine.resume()' below. + attempted_step_count += 1 + + # Test suspending and resuming engine. + if args.suspend_resume_interval is not None: + + # Suspend. + if attempted_step_count % args.suspend_resume_interval == 0: + print("**** step %d/%d ... suspend." % (engine.context.step_count, attempted_step_count)) + engine.suspend() + + # Resume, 0+ attempted steps later. + if ( + attempted_step_count > 0 + and (attempted_step_count - args.suspend_resume_interval // 2) + % args.suspend_resume_interval + == 0 + ): + print("**** step %d/%d ... resume." % (engine.context.step_count, attempted_step_count)) + engine.resume() + + # If engine suspended, continue to next iter. + if isinstance(result, EngineSuspendedError): + continue + + _process_step_result(result) + + # Check if all requests are finished. + if not (engine.has_unfinished_requests() or num_requests_added < num_requests_total): + break # Resume engine (NOOP if not suspended). - if engine.is_suspended: - engine.resume() + 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 +314,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 +374,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 +398,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 +421,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,15 +447,20 @@ 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.context.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) + if args.output_request_events: + result_dict["events"] = [e.serialize() for e in req.events] json_results[req.request_id] = result_dict # Track system-level throughput as a test / debug metric @@ -609,6 +469,7 @@ def escape_str(s): # Attach peak memory metrics; the functional test only validates these # if the fields exist in the golden values. json_results.update(peak_mem_stats) + json_results["lifetime_prefill_token_count"] = engine.context.lifetime_prefill_token_count print(f' Saving results to {args.output_path}') with open(args.output_path, "w") as fp: @@ -631,7 +492,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 +504,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"steps: {engine.context.step_count:d} … " + f"capture {capture_str}", ) print("~~~") diff --git a/examples/inference/gpt/gpt_dynamic_inference_12b.sh b/examples/inference/gpt/gpt_dynamic_inference_12b.sh index 4991d9d5177..ca21bb170a5 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_12b.sh +++ b/examples/inference/gpt/gpt_dynamic_inference_12b.sh @@ -69,7 +69,6 @@ ARGS=" \ --tiktoken-pattern v2 \ --tokenizer-model ${TOKENIZER_MODEL} \ --distributed-timeout-minutes 2400 \ - --transformer-impl local \ --use-flash-attn \ --inference-rng-tracker \ \ diff --git a/examples/inference/gpt/gpt_dynamic_inference_357m.sh b/examples/inference/gpt/gpt_dynamic_inference_357m.sh index 44abb575c63..cc99bdddec1 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_357m.sh +++ b/examples/inference/gpt/gpt_dynamic_inference_357m.sh @@ -32,6 +32,7 @@ export CUDA_DEVICE_MAX_CONNECTIONS=1 # Miscellaneous. : ${USE_COORDINATOR=0} : ${ENGINE=dynamic} +: ${NPROC_PER_NODE=1} : ${EXTRA_ARGS=""} # NSIGHT_PREFIX=/path/to/nsight/profile @@ -101,7 +102,7 @@ fi if [[ "${USE_COORDINATOR}" == "0" ]]; then CMD="python -m examples.inference.gpt.gpt_${ENGINE}_inference ${ARGS}" else - CMD="python -um examples.inference.gpt.gpt_${ENGINE}_inference_with_coordinator ${ARGS}" + CMD="python -m torch.distributed.run --nproc-per-node ${NPROC_PER_NODE} -m examples.inference.gpt.gpt_${ENGINE}_inference_with_coordinator ${ARGS}" fi if [[ -v NSIGHT_PREFIX ]]; then diff --git a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py index cbb7a1aa745..d34749e3a5a 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py +++ b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py @@ -2,43 +2,50 @@ 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.engines.dynamic_engine import EngineState 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 suspend_resume_cycle(client, engine, args, futures): + """Wait for all in-flight requests, then suspend/train/resume.""" + await asyncio.gather(*futures) + + client.pause_engines() + await engine.wait_until(EngineState.PAUSED) + client.suspend_engines() + await engine.wait_until(EngineState.SUSPENDED) + if args.suspend_timeout > 0: + await asyncio.sleep(args.suspend_timeout) + client.resume_engines() + await engine.wait_until(EngineState.RESUMED) + client.unpause_engines() + await engine.wait_until(EngineState.RUNNING) + + async def main( engine: DynamicInferenceEngine, requests: List[Request], @@ -51,54 +58,43 @@ 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. - + args = get_args() + dp_addr = await engine.start_listening_to_data_parallel_coordinator( inference_coordinator_port=port, launch_inference_coordinator=True, + coordinator_schedule_output_path=args.coordinator_schedule_output_path, ) - args = get_args() - - # Test suspend/resume intervals. - if dist.get_rank() == 0 and args.suspend_resume_interval is not None: - # 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, - )) - resume_idxs = set( - min(len(requests), i + suspend_resume_interval // 2) - for i in suspend_idxs - ) - else: - suspend_idxs = set() - resume_idxs = set() + # All ranks agree on the number of suspend/resume cycles from args. + num_suspend_resume_cycles = len(requests) // args.suspend_resume_interval if args.suspend_resume_interval else 0 # Create client and run example. if dist.get_rank() == 0: - client = InferenceClient(dp_addr) # submits requests to the inference coordinator - await client.start() + client = InferenceClient(dp_addr, deserialize=True) # submits requests to the inference coordinator + client.start() base_arrival_time = time.time_ns() / 10**9 for request in requests: request.time_arrival = request.time_offset + base_arrival_time futures = [] num_requests_total = len(requests) num_requests_added = 0 - # logging.info("Waiting for 20 seconds before starting to add requests. This is to mimic an RL style setup..") - # time.sleep(20) + next_suspend_at = args.suspend_resume_interval or 0 + cycles_done = 0 + while True: 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 @@ -106,18 +102,16 @@ async def main( futures.append(client.add_request(request.prompt_text, request.sampling_params)) num_requests_added += 1 - # Test suspend/resume. - if num_requests_added in suspend_idxs: - client.suspend_engines() - if num_requests_added in resume_idxs: - client.resume_engines() + if num_requests_added >= next_suspend_at and cycles_done < num_suspend_resume_cycles: + await suspend_resume_cycle(client, engine, args, futures) + cycles_done += 1 + next_suspend_at += args.suspend_resume_interval 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 @@ -125,19 +119,25 @@ async def main( futures.append(client.add_request(request.prompt_text, request.sampling_params)) num_requests_added += 1 - # Test suspend/resume. - if num_requests_added in suspend_idxs: - client.suspend_engines() - if num_requests_added in resume_idxs: - client.resume_engines() + if num_requests_added >= next_suspend_at and cycles_done < num_suspend_resume_cycles: + await suspend_resume_cycle(client, engine, args, futures) + cycles_done += 1 + next_suspend_at += args.suspend_resume_interval if num_requests_added == num_requests_total: 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) + else: + # Non-rank-0: match the suspend/resume cycles that rank 0 drives. + for _ in range(num_suspend_resume_cycles): + await engine.wait_until(EngineState.PAUSED) + await engine.wait_until(EngineState.SUSPENDED) + await engine.wait_until(EngineState.RESUMED) + await engine.wait_until(EngineState.RUNNING) if dist.get_rank() == 0: # Write results to JSON. Primarily used for functional testing. @@ -145,8 +145,7 @@ async def main( json_results = {} throughputs = [] - for record in results: - req = record.merge() + for req in results: result_dict = { "input_prompt": req.prompt, "generated_text": req.generated_text.replace("\n", "\\n"), @@ -157,6 +156,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: @@ -166,35 +168,42 @@ async def main( else: print("Results:") unique_prompt_map = defaultdict(list) - for record in results: - req = record.merge() + for req in results: 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"), - )) - - # kill the engines and suspend the client - # 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() + 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"), + ) + ) + + # Pause before stopping: STOP requires PAUSED or SUSPENDED state. + client.pause_engines() + + await engine.wait_until(EngineState.PAUSED) + + if dist.get_rank() == 0: + client.stop_engines() + + await engine.wait_until(EngineState.STOPPED) - # once the stop signal eventually makes its way to each GPU, the engines will stop. - await asyncio.gather(engine.engine_loop_task) + if dist.get_rank() == 0: + client.shutdown_coordinator() + client.stop() logging.info(f"Rank: {dist.get_rank()} stopped their engine instance successfully.") 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 +222,14 @@ async def main( ), ) - # Requests, context, conroller. - model = get_model() - mamba_inference_state_config = get_mamba_inference_state_config_from_model(model) - requests = ( - build_requests(args, tokenizer, sampling_params) if dist.get_rank() == 0 else None - ) + model = get_model_for_inference() - 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) + requests = build_requests(args, tokenizer, sampling_params) - # 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 +238,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..c9b1c05c544 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,36 +283,31 @@ 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 buffer_limits_str = ( f"bf: {get_mem_size_str(args.inference_dynamic_batching_buffer_size_gb*1024**3)}, " - f"{context.block_allocator.active_count} chunks " + f"{context.kv_block_allocator.active_count} chunks " 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/mamba/README.md b/examples/mamba/README.md index f8f6d796837..ce60f119ea5 100644 --- a/examples/mamba/README.md +++ b/examples/mamba/README.md @@ -43,7 +43,8 @@ set to 1. The arguments in the script will need to be changed if using a checkpoint with a different model parallel configuration or other differences, such as model architecture. For example, to run the 8B pure Mamba-2 model, change -`--hybrid-attention-ratio` and `--hybrid-mlp-ratio` to 0.0, or remove them. +`--hybrid-layer-pattern` to use only `M` symbols (e.g., 56 `M`s for the 8B +model), or remove it entirely. Use [`run_text_gen_server_8b_gpt3.sh`](./run_text_gen_server_8b_gpt3.sh) to start a text generation server using the 8B reference Transformer checkpoint. @@ -67,24 +68,46 @@ export PYTHONPATH=:PYTHONPATH ## Hybrid Options -`--hybrid-attention-ratio ATT` specifies a target ratio of attention layers -to total layers. For example, 4 attention layers out of 48 total layers is -specified by `--hybrid-attention-ratio 0.08`. +`--hybrid-layer-pattern PATTERN` specifies the layer type for every layer in +the model using a string of single-character symbols: -`--hybrid-mlp-ratio MLP` specifies a target ratio of MLP layers to total -layers. For example, 24 MLP layers out of 48 total layers is specified by -`--hybrid-mlp-ratio 0.5`. +* `M` — Mamba layer +* `*` — Attention layer +* `-` — MLP layer +* `E` — MoE layer -* (`ATT` + `MLP`) must be less than or equal to 1.0. -* (1.0 - `ATT` - `MLP`) is the hybrid mamba ratio, the ratio of mamba layers to -total layers. -* `ATT` = `MLP` = 0 is a pure Mamba model. -* `ATT` = `MLP` = 0.5 is a transfomer model. +The number of layers is derived from the pattern length, so `--num-layers` +should not be specified when `--hybrid-layer-pattern` is used. -If either `ATT` or `MLP` is greater than 0.0 or if `--hybrid-override-pattern` -is specified, the logfile will include information about the hybrid layer -pattern used. `--hybrid-override-pattern` can be used to specify a different -pattern than the default, algorithmically-generated one. +For example, the 8B hybrid model described in the technical report uses: + +``` +--hybrid-layer-pattern "M-M-M--M-M*-M-M-M-M--M*-M-M-M-M-M*--M-M-M-M-M*-M--M-M-M-" +``` + +This is a 56-layer model with 4 attention layers, 28 MLP layers, and 24 Mamba +layers. + +A pure Mamba model uses only `M` symbols (e.g., `MMMMMMMM` for 8 layers). +A pure transformer model uses only `*` and `-` symbols. + +### Pipeline parallelism + +Use `|` to define pipeline stage boundaries for flexible virtual pipeline +parallelism (fVPP). For example, `M-M-|M-M*-|M-M-|M-M*-` defines 4 pipeline +segments. The number of segments must be evenly divisible by +`--pipeline-model-parallel-size`. + +### Multi-Token Prediction (MTP) + +Use `/` to append MTP layer patterns. Each pattern after the separator +represents one MTP prediction depth. For example, `M*M*/MM/MM` has main +pattern `M*M*` with MTP pattern `MM` repeated for 2 depths. + +### Deprecated options + +`--hybrid-override-pattern`, `--hybrid-attention-ratio`, and +`--hybrid-mlp-ratio` are deprecated. Use `--hybrid-layer-pattern` instead. ## Mamba vs Mamba-2 diff --git a/examples/mamba/run_text_gen_server_8b.sh b/examples/mamba/run_text_gen_server_8b.sh index 8d3137f2442..d228e0c0edb 100755 --- a/examples/mamba/run_text_gen_server_8b.sh +++ b/examples/mamba/run_text_gen_server_8b.sh @@ -6,6 +6,8 @@ CHECKPOINT_PATH=$1 TOKENIZER_PATH=$2 +HYBRID_LAYER_PATTERN="M-M-M--M-M*-M-M-M-M--M*-M-M-M-M-M*--M-M-M-M-M*-M--M-M-M-" + DISTRIBUTED_ARGS="--nproc_per_node 1 \ --nnodes 1 \ --node_rank 0 \ @@ -24,14 +26,12 @@ torchrun $DISTRIBUTED_ARGS ../../tools/run_mamba_text_generation_server.py \ --tensor-model-parallel-size 1 \ --pipeline-model-parallel-size 1 \ --untie-embeddings-and-output-weights \ - --num-layers 56 \ + --hybrid-layer-pattern ${HYBRID_LAYER_PATTERN} \ --hidden-size 4096 \ --load ${CHECKPOINT_PATH} \ --num-attention-heads 32 \ --group-query-attention \ --num-query-groups 8 \ - --hybrid-attention-ratio 0.08 \ - --hybrid-mlp-ratio 0.5 \ --attention-dropout 0.0 \ --hidden-dropout 0.0 \ --disable-bias-linear \ diff --git a/examples/mamba/train.sh b/examples/mamba/train.sh index 3952a997d47..ba83f0d4e33 100755 --- a/examples/mamba/train.sh +++ b/examples/mamba/train.sh @@ -7,14 +7,14 @@ MODEL_SCALE="800M" # or "8B" case "${MODEL_SCALE}" in "800M") TENSOR_MODEL_PARALLEL_SIZE=1 - NUM_LAYERS=48 + HYBRID_LAYER_PATTERN="M-M-M--M-*M-M-M-M--*M-M-M-M-*M--M-M-M-*M-M--M-M-" HIDDEN_SIZE=1024 NUM_ATTENTION_HEADS=16 GLOBAL_BATCH_SIZE=32 ;; "8B") TENSOR_MODEL_PARALLEL_SIZE=4 - NUM_LAYERS=56 + HYBRID_LAYER_PATTERN="M-M-M--M-M*-M-M-M-M--M*-M-M-M-M-M*--M-M-M-M-M*-M--M-M-M-" HIDDEN_SIZE=4096 NUM_ATTENTION_HEADS=32 GLOBAL_BATCH_SIZE=8 @@ -59,13 +59,11 @@ options=" \ --untie-embeddings-and-output-weights \ --init-method-std 0.02 \ --position-embedding-type none \ - --num-layers ${NUM_LAYERS} \ + --hybrid-layer-pattern ${HYBRID_LAYER_PATTERN} \ --hidden-size ${HIDDEN_SIZE} \ --num-attention-heads ${NUM_ATTENTION_HEADS} \ --group-query-attention \ --num-query-groups 8 \ - --hybrid-attention-ratio 0.08 \ - --hybrid-mlp-ratio 0.5 \ --seq-length ${SEQ_LEN} \ --max-position-embeddings ${SEQ_LEN} \ --train-samples ${TRAIN_SAMPLES} \ 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 = "