merging \\`$SOURCE_BRANCH\\` into \\`dev\\` for you! 🚀
Please review and approve this cherry pick by your convenience\\!\"\n tags:\n - arch/amd64\n - env/prod\n - origin/jet-fleet\n - owner/jet-core\n - purpose/utility\n - team/megatron\n rules:\n - if: $CI_COMMIT_BRANCH == \"main\" && $CI_PIPELINE_SOURCE == \"push\"\n allow_failure: true\n - when: never\n\npublish:approve_merge_gate:\n stage: publish\n image: maniator/gh\n tags:\n - arch/amd64\n - env/prod\n - origin/jet-fleet\n - owner/jet-core\n - purpose/utility\n - team/megatron\n script:\n - |\n set -eoux pipefail\n EXIT_CODE=0\n apk add python3\n python -m venv .venv\n source .venv/bin/activate\n pip install --no-cache-dir python-gitlab click pygithub\n export GITLAB_ENDPOINT\n export RO_API_TOKEN=${PROJECT_ACCESS_TOKEN_MCORE}\n if [[ \"$CI_COMMIT_BRANCH\" == *main* ]]; then\n export TARGET_BRANCH=\"main\"\n elif [[ \"$CI_COMMIT_BRANCH\" == *dev* ]]; then\n export TARGET_BRANCH=\"dev\"\n fi\n\n python tests/test_utils/python_scripts/check_status_of_main.py --target-branch \"$TARGET_BRANCH\" --once || EXIT_CODE=$?\n\n export GH_TOKEN=$GH_TOKEN\n export REPO=NVIDIA/Megatron-LM\n\n if [[ $EXIT_CODE -eq 0 ]]; then\n export STATUS=\"approved\"\n export COMMENT=\"Main is healthy. Submitting PR.\"\n elif [[ $EXIT_CODE -eq 1 ]]; then\n export STATUS=\"rejected\"\n export COMMENT=\"$TARGET_BRANCH is not healthy. An automation engineer is investigating. No need to take any action.\"\n elif [[ $EXIT_CODE -eq 2 ]]; then\n echo \"Main is running. We won't cancel the deployment.\"\n exit 0\n fi\n\n if [[ $EXIT_CODE -lt 2 ]]; then\n python tests/test_utils/python_scripts/approve_merge_gate.py\n fi\n\n retry:\n max: 2\n rules:\n - if: $CI_PIPELINE_SOURCE == \"schedule\" && ($CI_COMMIT_BRANCH == 'ci-approve-dev' || $CI_COMMIT_BRANCH == 'ci-approve-main')\n when: always\n - when: never\n\npublish:sync_branches:\n stage: publish\n image: python:3.10\n script:\n - set -x\n - git remote add github https://github.com/NVIDIA/Megatron-LM.git || true\n - git remote add gitlab https://gitlab-ci-token:${PROJECT_ACCESS_TOKEN_MCORE}@${GITLAB_ENDPOINT}/${CI_PROJECT_NAMESPACE}/Megatron-LM.git || true\n - BRANCHES=(\"main\" \"dev\")\n - |\n while IFS= read -r line; do\n BRANCHES+=(\"$line\") # Add each line to the array\n done < <( \\\n git ls-remote --heads \"https://token:${PAT}@github.com/NVIDIA/Megatron-LM.git\" 'refs/heads/core_*' | \\\n cut -d'/' -f3- \\\n )\n - |\n for BRANCH in \"${BRANCHES[@]}\"; do\n # Define the full refspec for the branch\n BRANCH_REF=\"refs/heads/$BRANCH\"\n \n echo \"--- Processing branch: $BRANCH ---\"\n \n # 1. Explicitly fetch the branch ref from 'github'\n # This avoids fetching a tag with the same name.\n # It updates/creates the remote-tracking branch (e.g., 'refs/remotes/github/core_r0.10.0')\n if ! git fetch github \"$BRANCH_REF:refs/remotes/github/$BRANCH\"; then\n echo \"Failed to fetch branch $BRANCH. Skipping.\"\n continue\n fi\n\n # 2. Create or update the local branch from the remote-tracking branch we just fetched.\n # The -B flag creates the branch if it doesn't exist or resets it if it does.\n if ! git checkout -B \"$BRANCH\" \"github/$BRANCH\"; then\n echo \"Failed to checkout local branch $BRANCH. Skipping.\"\n continue\n fi\n \n # 3. Now you are on the correct local branch, ready to push.\n echo \"Successfully on branch $BRANCH. Echoing push command:\"\n git push -u gitlab HEAD:refs/heads/$BRANCH --force\n echo \"-----------------------------------\"\n done\n tags:\n - arch/amd64\n - env/prod\n - origin/jet-fleet\n - owner/jet-core\n - purpose/utility\n - team/megatron\n retry:\n max: 2\n rules:\n - if: $CI_PIPELINE_SOURCE == \"schedule\" && ($CI_COMMIT_BRANCH == 'ci-sync-branches')\n when: always\n - when: never\n"
},
{
"path": ".gitlab-ci.yml",
"content": ".merge_train_rule: &merge_train_rule\n UNIT_TEST: \"yes\"\n UNIT_TEST_REPEAT: 1\n UNIT_TEST_TIMEOUT: 30\n INTEGRATION_TEST: \"no\"\n INTEGRATION_TEST_SCOPE: mr\n FUNCTIONAL_TEST: \"yes\"\n FUNCTIONAL_TEST_SCOPE: mr-slim\n FUNCTIONAL_TEST_REPEAT: 1\n FUNCTIONAL_TEST_TIME_LIMIT: 2700\n CLUSTER_A100: \"\"\n CLUSTER_H100: \"\"\n PUBLISH: \"no\"\n\nworkflow:\n rules:\n # Do not trigger for forks\n - if: $CI_PROJECT_NAMESPACE != \"ADLR\" || ($CI_PIPELINE_SOURCE == \"merge_request_event\" && $CI_MERGE_REQUEST_PROJECT_PATH != \"ADLR/megatron-lm\")\n when: never\n\n - if: $CI_PIPELINE_SOURCE == \"schedule\" && ($CI_COMMIT_BRANCH == 'ci-approve-dev' || $CI_COMMIT_BRANCH == 'ci-approve-main')\n\n # ci-branches only for schedule\n - if: $CI_COMMIT_BRANCH =~ /ci-/ && $CI_PIPELINE_SOURCE != \"schedule\"\n when: never\n\n # For schedules pipelines\n - if: $CI_PIPELINE_SOURCE == \"schedule\"\n auto_cancel:\n on_new_commit: none\n\n # For manual pipelines (GitLab UI)\n - if: $CI_PIPELINE_SOURCE == \"web\"\n\n # For pipelines created via the REST API (personal access token)\n - if: $CI_PIPELINE_SOURCE == \"api\"\n\n # For trigger pipelines\n - if: $CI_PIPELINE_SOURCE == \"trigger\"\n\n # For push to main\n - if: $CI_PIPELINE_SOURCE == 'push' && ($CI_COMMIT_BRANCH == \"main\" || $CI_COMMIT_BRANCH == \"dev\" || $CI_COMMIT_BRANCH =~ /^core_/)\n variables:\n UNIT_TEST: \"no\"\n INTEGRATION_TEST: \"no\"\n FUNCTIONAL_TEST: \"yes\"\n FUNCTIONAL_TEST_SCOPE: mr\n FUNCTIONAL_TEST_REPEAT: 5\n FUNCTIONAL_TEST_RECORD_CHECKPOINTS: \"no\"\n FUNCTIONAL_TEST_TIME_LIMIT: 3600\n CLUSTER_A100: \"\"\n CLUSTER_H100: \"\"\n PUBLISH: \"no\"\n auto_cancel:\n on_new_commit: interruptible\n\n # For merge-trains that need to be fast-tracked\n - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merge_train' && $CI_MERGE_REQUEST_LABELS =~ /fast-track/\n variables:\n UNIT_TEST: \"yes\"\n UNIT_TEST_REPEAT: 1\n UNIT_TEST_TIMEOUT: 30\n INTEGRATION_TEST: \"no\"\n FUNCTIONAL_TEST: \"no\"\n CLUSTER_A100: \"\"\n CLUSTER_H100: \"\"\n PUBLISH: \"no\"\n\n # For normal merge-trains\n - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merge_train'\n variables: *merge_train_rule\n\n # For MRs with integration suite\n - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run tests/\n variables:\n UNIT_TEST: \"yes\"\n UNIT_TEST_REPEAT: 1\n UNIT_TEST_TIMEOUT: 30\n INTEGRATION_TEST: \"yes\"\n INTEGRATION_TEST_SCOPE: mr\n FUNCTIONAL_TEST: \"no\"\n FUNCTIONAL_TEST_SCOPE: mr-slim\n FUNCTIONAL_TEST_REPEAT: 1\n FUNCTIONAL_TEST_TIME_LIMIT: 2700\n CLUSTER_A100: \"\"\n CLUSTER_H100: \"\"\n PUBLISH: \"no\"\n\n # For MRs with nightly\n - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run nightly/\n variables:\n UNIT_TEST: \"yes\"\n UNIT_TEST_REPEAT: 1\n UNIT_TEST_TIMEOUT: 30\n INTEGRATION_TEST: \"no\"\n FUNCTIONAL_TEST: \"yes\"\n FUNCTIONAL_TEST_SCOPE: nightly\n FUNCTIONAL_TEST_REPEAT: 5\n FUNCTIONAL_TEST_RECORD_CHECKPOINTS: \"no\"\n FUNCTIONAL_TEST_TIME_LIMIT: 2700\n CLUSTER_A100: \"\"\n CLUSTER_H100: \"\"\n PUBLISH: \"no\"\n\n # For MRs with weekly\n - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run weekly/\n variables:\n UNIT_TEST: \"yes\"\n UNIT_TEST_REPEAT: 1\n UNIT_TEST_TIMEOUT: 30\n INTEGRATION_TEST: \"no\"\n FUNCTIONAL_TEST: \"yes\"\n FUNCTIONAL_TEST_SCOPE: weekly\n FUNCTIONAL_TEST_REPEAT: 1\n FUNCTIONAL_TEST_RECORD_CHECKPOINTS: \"no\"\n FUNCTIONAL_TEST_TIME_LIMIT: 9000\n CLUSTER_A100: \"\"\n CLUSTER_H100: \"\"\n PUBLISH: \"no\"\n\n # For MRs with heavy suite\n - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result' && $CI_MERGE_REQUEST_LABELS =~ /Run functional tests/\n variables:\n UNIT_TEST: \"yes\"\n UNIT_TEST_REPEAT: 1\n UNIT_TEST_TIMEOUT: 30\n INTEGRATION_TEST: \"no\"\n FUNCTIONAL_TEST: \"yes\"\n FUNCTIONAL_TEST_SCOPE: mr\n FUNCTIONAL_TEST_REPEAT: 1\n FUNCTIONAL_TEST_TIME_LIMIT: 2700\n CLUSTER_A100: \"\"\n CLUSTER_H100: \"\"\n PUBLISH: \"no\"\n\n # Default MRs\n - if: $CI_MERGE_REQUEST_EVENT_TYPE == 'merged_result'\n variables:\n UNIT_TEST: \"yes\"\n UNIT_TEST_REPEAT: 1\n UNIT_TEST_TIMEOUT: 30\n INTEGRATION_TEST: \"no\"\n FUNCTIONAL_TEST: \"no\"\n PUBLISH: \"no\"\n\n - when: never\n\n auto_cancel:\n on_new_commit: interruptible\n\nstages:\n - build\n - test\n - integration_tests\n - functional_tests\n - publish\n\ndefault:\n interruptible: true\n retry:\n max: 2\n when: runner_system_failure\n\nvariables:\n BUILD:\n value: \"yes\"\n UNIT_TEST:\n value: \"yes\"\n options:\n - \"yes\"\n - \"no\"\n description: To run the funtional test suite\n UNIT_TEST_REPEAT:\n value: \"1\"\n description: \"Number of repetitions\"\n UNIT_TEST_TIMEOUT:\n value: \"30\"\n description: Timeout (minutes) for Unit tests (all repeats)\n INTEGRATION_TEST:\n value: \"yes\"\n options:\n - \"yes\"\n - \"no\"\n description: To run the integration test suite\n INTEGRATION_TEST_SCOPE:\n value: \"mr\"\n options:\n - \"mr\"\n - \"nightly\"\n - \"weekly\"\n - \"pre-release\"\n - \"release\"\n description: \"Testsuite to run (only for INTEGRATION_TEST=yes)\"\n INTEGRATION_TEST_TIME_LIMIT:\n value: \"900\"\n description: \"Timeout in seconds per test\"\n INTEGRATION_TEST_CASES:\n value: \"all\"\n description: \"Comma-separated list of test_cases to run. Use 'all' to run the full suite.\"\n FUNCTIONAL_TEST:\n value: \"yes\"\n options:\n - \"yes\"\n - \"no\"\n description: To run the funtional test suite\n FUNCTIONAL_TEST_SCOPE:\n value: \"mr\"\n options:\n - \"mr\"\n - \"nightly\"\n - \"weekly\"\n - \"pre-release\"\n - \"release\"\n description: \"Testsuite to run (only for FUNCTIONAL_TEST=yes)\"\n FUNCTIONAL_TEST_REPEAT:\n value: \"5\"\n description: \"Number of repetitions per test\"\n FUNCTIONAL_TEST_TIME_LIMIT:\n value: \"2700\"\n description: \"Timeout in seconds per test\"\n FUNCTIONAL_TEST_CASES:\n value: \"all\"\n description: \"Comma-separated list of test_cases to run. Use 'all' to run the full suite.\"\n FUNCTIONAL_TEST_NAME:\n description: \"Name of functional test run (only for pre-release and release)\"\n value: \"$$CI_COMMIT_SHA\"\n FUNCTIONAL_TEST_RECORD_CHECKPOINTS:\n value: \"no\"\n description: \"Record golden checkpoints\"\n options:\n - \"yes\"\n - \"no\"\n CLUSTER_A100:\n value: \"dgxa100_dracooci\"\n options:\n - \"dgxa100_dracooci\"\n - \"dgxa100_dracooci-ord\"\n description: \"Cluster for A100 workloads\"\n CLUSTER_H100:\n value: \"dgxh100_coreweave\"\n options:\n - \"dgxh100_coreweave\"\n - \"dgxh100_eos\"\n description: \"Cluster for H100 workloads\"\n CLUSTER_GB200:\n value: \"dgxgb200_oci-hsg\"\n options:\n - \"dgxgb200_oci-hsg\"\n description: \"Cluster for H100 workloads\"\n PUBLISH:\n value: \"no\"\n options:\n - \"yes\"\n - \"no\"\n description: Build and publish a wheel to PyPi\n PUBLISH_COMMIT:\n value: \"$$CI_COMMIT_SHA\"\n description: Which commit to publish\n PUBLISH_VERSION_BUMP_BRANCH:\n value: \"$$CI_COMMIT_BRANCH\"\n description: Which branch to target for version bump\n PUBLISH_SCOPE:\n value: \"code-freeze\"\n options:\n - \"code-freeze\"\n - \"release\"\n - \"review-reminder\"\n - \"upgrade-dependencies\"\n description: Type of publish (freeze or final release)\n\n # CI wide variables\n CI_MCORE_LTS_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_ci_lts\n CI_MCORE_DEV_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_ci_dev\n CI_NEMO_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/nemo_ci\n UTILITY_IMAGE: ${GITLAB_ENDPOINT}:5005/adlr/megatron-lm/mcore_utility\n TE_GIT_REF: \"\"\n\ninclude:\n - .gitlab/stages/00.pre.yml\n - .gitlab/stages/01.build.yml\n - .gitlab/stages/02.test.yml\n - .gitlab/stages/03.integration-tests.yml\n - .gitlab/stages/04.functional-tests.yml\n - .gitlab/stages/05.publish.yml\n"
},
{
"path": ".pre-commit-config.yaml",
"content": "repos:\n- repo: https://github.com/psf/black\n rev: 'refs/tags/24.4.2:refs/tags/24.4.2'\n hooks:\n - id: black\n files: ^megatron/core/.*|^tests/unit_tests/.*\n args: [\"--skip-magic-trailing-comma\", \"--skip-string-normalization\"]\n- repo: https://github.com/pycqa/pylint\n rev: v3.2.6\n hooks:\n - id: pylint\n files: ^megatron/core/.*\n- repo: https://github.com/pycqa/isort\n rev: 5.13.2\n hooks:\n - id: isort\n files: ^megatron/core/.*"
},
{
"path": ".pylintrc",
"content": "[MAIN]\nignore-paths=tests\nmax-line-length=100\nload-plugins=pylint.extensions.bad_builtin\n[MESSAGES CONTROL]\ndisable=all\n\nenable=C0115,C0116,W0611,C0301,E0606,W0141\n# C0115: missing-class-docstring\n# C0116: missing-function-docstring\n# W0611: unused-import\n# C0301: line-too-long\n# E0606: possibly-used-before-assignment\n# W0141: bad-builtin (from bad_builtin extension)\n\n[BASIC]\nbad-functions=print\n\n[BAD_BUILTIN]\n# Specify which builtins should be flagged\nbad-builtins=print"
},
{
"path": ".python-version",
"content": "3.12"
},
{
"path": "CONTRIBUTING.md",
"content": "# Contributing to Megatron\n\nVisit our [contributing page](https://docs.nvidia.com/megatron-core/developer-guide/latest/developer/contribute.html)."
},
{
"path": "LICENSE",
"content": "The following applies to all files unless otherwise noted:\n\n# Copyright (c) 2019-2025, NVIDIA CORPORATION. All rights reserved.\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions\n# are met:\n# * Redistributions of source code must retain the above copyright\n# notice, this list of conditions and the following disclaimer.\n# * Redistributions in binary form must reproduce the above copyright\n# notice, this list of conditions and the following disclaimer in the\n# documentation and/or other materials provided with the distribution.\n# * Neither the name of NVIDIA CORPORATION nor the names of its\n# contributors may be used to endorse or promote products derived\n# from this software without specific prior written permission.\n#\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n\n--\n\nThis repository also contains code from Hugging Face Inc., Google Research,\nFacebook (from their Fairseq, Dino, and ParlAI projects), Microsoft (from their\nSwin-Transformer project), Philip Popien, the Mamba project (Tri Dao and\nAlbert Gu), and the Triton language and compiler project (Philippe Tillet and\nOpenAI). Files from these organizations have notices at the top of each file.\nBelow are licenses used in those files, as indicated.\n\n\n--------------------------------------------------------------------------------------\n-- LICENSE FOR Facebook, huggingface, Google Research, LLaVA, Mamba, TinyZero and vLLM code --\n\n\n Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n\n--------------------------------------------------------------------------------\nLICENSE FOR\nFacebook, Inc. and its affiliates,\nMeta Platforms, Inc. and its affiliates,\nMicrosoft Corporation,\nOpenGVLab/InternVL,\nTriton language and compiler,\nand DeepSeek.\n\nMIT License\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n\n--------------------------------------------------------------------------------\nLICENSE FOR Thinking Machines Lab \n\nMIT License\n\nCopyright 2025 Thinking Machines Lab \n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n\n--------------------------------------------------------------------------------\nLICENSE FOR\nMeta Platforms, Inc. and affiliates.\n\nBSD License\n\nRedistribution and use in source and binary forms, with or without modification,\nare permitted provided that the following conditions are met:\n\n * Redistributions of source code must retain the above copyright notice, this\n list of conditions and the following disclaimer.\n\n * Redistributions in binary form must reproduce the above copyright notice,\n this list of conditions and the following disclaimer in the documentation\n and/or other materials provided with the distribution.\n\n * Neither the name Meta nor the names of its contributors may be used to\n endorse or promote products derived from this software without specific\n prior written permission.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\nANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\nWARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\nDISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR\nANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\nLOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON\nANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\nSOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n"
},
{
"path": "MANIFEST.in",
"content": "include megatron/core/requirements.txt\ninclude megatron/core/README.md\ninclude megatron/core/package_info.py\nglobal-exclude LICENSE\nrecursive-include requirements *\n"
},
{
"path": "README.md",
"content": "
\n\nMegatron-LM and Megatron Core\n=============================\n\n
GPU-optimized library for training transformer models at scale
\n\n## About\n\nThis repository contains two components: **Megatron-LM** and **Megatron Core**.\n\n**Megatron-LM** is a reference example that includes Megatron Core plus pre-configured training scripts. Best for research teams, learning distributed training, and quick experimentation.\n\n**Megatron Core** is a composable library with GPU-optimized building blocks for custom training frameworks. It provides transformer building blocks, advanced parallelism strategies (TP, PP, DP, EP, CP), mixed precision support (FP16, BF16, FP8, FP4), and model architectures. Best for framework developers and ML engineers building custom training pipelines.\n\n**[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** provides bidirectional Hugging Face ↔ Megatron checkpoint conversion with production-ready recipes.\n\n## Getting Started\n\n**Install from PyPI:**\n\n```bash\nuv pip install megatron-core\n```\n\n**Or clone and install from source:**\n\n```bash\ngit clone https://github.com/NVIDIA/Megatron-LM.git\ncd Megatron-LM\nuv pip install -e .\n```\n\n> **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 .`).\n\nFor NGC container setup and all installation options, see the **[Installation Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/get-started/install.html)**.\n\n- **[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\n- **[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\n- **[Contribution Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/developer/contribute.html)** - How to contribute to Megatron Core\n\n# Latest News\n\n- **[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.\n- **[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.\n- **[2025/12]** **Megatron Core development has moved to GitHub!** All development and CI now happens in the open. We welcome community contributions.\n- **[2025/10]** **[Megatron Dev Branch](https://github.com/NVIDIA/Megatron-LM/tree/dev)** - early access branch with experimental features.\n- **[2025/10]** **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Bidirectional converter for interoperability between Hugging Face and Megatron checkpoints, featuring production-ready recipes for popular models.\n- **[2025/08]** **[MoE Q3-Q4 2025 Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - Comprehensive roadmap for MoE features including DeepSeek-V3, Qwen3, advanced parallelism strategies, FP8 optimizations, and Blackwell performance enhancements.\n- **[2025/08]** **[GPT-OSS Model](https://github.com/NVIDIA/Megatron-LM/issues/1739)** - Advanced features including YaRN RoPE scaling, attention sinks, and custom activation functions are being integrated into Megatron Core.\n- **[2025/06]** **[Megatron MoE Model Zoo](https://github.com/yanring/Megatron-MoE-ModelZoo)** - Best practices and optimized configurations for training DeepSeek-V3, Mixtral, and Qwen3 MoE models with performance benchmarking and checkpoint conversion tools.\n- **[2025/05]** Megatron Core v0.11.0 brings new capabilities for multi-data center LLM training ([blog](https://developer.nvidia.com/blog/turbocharge-llm-training-across-long-haul-data-center-networks-with-nvidia-nemo-framework/)).\n\n\nPrevious News\n\n- **[2024/07]** Megatron Core v0.7 improves scalability and training resiliency and adds support for multimodal training ([blog](https://developer.nvidia.com/blog/train-generative-ai-models-more-efficiently-with-new-nvidia-Megatron-Core-functionalities/)).\n- **[2024/06]** Megatron Core added supports for Mamba-based models. Check out our paper [An Empirical Study of Mamba-based Language Models](https://arxiv.org/pdf/2406.07887) and [code example](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba).\n- **[2024/01 Announcement]** NVIDIA has released the core capabilities in **Megatron-LM** into [**Megatron Core**](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core) in this repository. Megatron Core expands upon Megatron-LM's GPU-optimized techniques with more cutting-edge innovations on system-level optimizations, featuring composable and modular APIs.\n\n\n\n# Project Structure\n\n```\nMegatron-LM/\n├── megatron/\n│ ├── core/ # Megatron Core (kernels, parallelism, building blocks)\n│ │ ├── models/ # Transformer models\n│ │ ├── transformer/ # Transformer building blocks\n│ │ ├── tensor_parallel/ # Tensor parallelism\n│ │ ├── pipeline_parallel/ # Pipeline parallelism\n│ │ ├── distributed/ # Distributed training (FSDP, DDP)\n│ │ ├── optimizer/ # Optimizers\n│ │ ├── datasets/ # Dataset loaders\n│ │ ├── inference/ # Inference engines and server\n│ │ └── export/ # Model export (e.g. TensorRT-LLM)\n│ ├── training/ # Training scripts\n│ ├── legacy/ # Legacy components\n│ ├── post_training/ # Post-training (quantization, distillation, pruning, etc.)\n│ └── rl/ # Reinforcement learning (RLHF, etc.)\n├── examples/ # Ready-to-use training examples\n├── tools/ # Utility tools\n├── tests/ # Comprehensive test suite\n└── docs/ # Documentation\n```\n\n# Performance Benchmarking\n\nFor our latest performance benchmarking results, please refer to [NVIDIA Megatron Bridge Performance Summary](https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html).\n\nOur codebase efficiently trains models from 2B to 462B parameters across thousands of GPUs, achieving up to **47% Model FLOP Utilization (MFU)** on H100 clusters.\n\n\n\n**Benchmark Configuration:**\n\n- **Vocabulary size**: 131,072 tokens\n- **Sequence length**: 4096 tokens\n- **Model scaling**: Varied hidden size, attention heads, and layers to achieve target parameter counts\n- **Communication optimizations**: Fine-grained overlapping with DP (`--overlap-grad-reduce`, `--overlap-param-gather`), TP (`--tp-comm-overlap`), and PP (enabled by default)\n\n**Key Results:**\n\n- **6144 H100 GPUs**: Successfully benchmarked 462B parameter model training\n- **Superlinear scaling**: MFU increases from 41% to 47-48% with model size\n- **End-to-end measurement**: Throughputs include all operations (data loading, optimizer steps, communication, logging)\n- **Production ready**: Full training pipeline with checkpointing and fault tolerance\n- *Note: Performance results measured without training to convergence*\n\n## Weak Scaling Results\n\nOur weak scaled results show superlinear scaling (MFU increases from 41% for the smallest model considered to 47-48% for the largest models); this is because larger GEMMs have higher arithmetic intensity and are consequently more efficient to execute.\n\n\n\n## Strong Scaling Results\n\nWe also strong scaled the standard GPT-3 model (our version has slightly more than 175 billion parameters due to larger vocabulary size) from 96 H100 GPUs to 4608 GPUs, using the same batch size of 1152 sequences throughout. Communication becomes more exposed at larger scale, leading to a reduction in MFU from 47% to 42%.\n\n\n\n# Roadmaps\n\n- **[MoE Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - DeepSeek-V3, Qwen3, advanced parallelism, FP8 optimizations, and Blackwell enhancements\n\n# Resources\n\n## Getting Help\n\n- 📖 **[Documentation](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html)** - Official documentation\n- 🐛 **[Issues](https://github.com/NVIDIA/Megatron-LM/issues)** - Bug reports and feature requests\n\n## Contributing\n\nWe ❤️ contributions! Ways to contribute:\n\n- 🐛 **Report bugs** - Help us improve reliability\n- 💡 **Suggest features** - Shape the future of Megatron Core\n- 📝 **Improve docs** - Make Megatron Core more accessible\n- 🔧 **Submit PRs** - Contribute code improvements\n\n**→ [Contributing Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/developer/contribute.html)**\n\n## Citation\n\nIf you use Megatron in your research or project, we appreciate that you use the following citations:\n\n```bibtex\n@article{megatron-lm,\n title={Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism},\n author={Shoeybi, Mohammad and Patwary, Mostofa and Puri, Raul and LeGresley, Patrick and Casper, Jared and Catanzaro, Bryan},\n journal={arXiv preprint arXiv:1909.08053},\n year={2019}\n}\n```\n"
},
{
"path": "codecov.yml",
"content": "comment: false\ncoverage:\n status:\n project: false\n patch:\n default:\n target: 80%\n threshold: 5%\n base: auto\n if_ci_failed: error\n if_no_uploads: success\n if_not_found: success\nfixes:\n - \"/opt/megatron-lm/::\"\n"
},
{
"path": "docker/.ngc_version.dev",
"content": "nvcr.io/nvidia/pytorch:26.02-py3"
},
{
"path": "docker/.ngc_version.lts",
"content": "nvcr.io/nvidia/pytorch:25.09-py3"
},
{
"path": "docker/Dockerfile.ci.dev",
"content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# syntax=docker/dockerfile:1.3-labs\n\nARG FROM_IMAGE_NAME\nFROM ${FROM_IMAGE_NAME} as main\nENV PIP_CONSTRAINT=\"\"\nENV DEBIAN_FRONTEND=noninteractive\nARG UV_VERSION=0.7.2\nARG YQ_VERSION=4.44.1\nENV PATH=\"/root/.local/bin:$PATH\"\nARG UV_PROJECT_ENVIRONMENT=/opt/venv\nENV UV_PROJECT_ENVIRONMENT=${UV_PROJECT_ENVIRONMENT}\nENV VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT\nENV PATH=\"$UV_PROJECT_ENVIRONMENT/bin:$PATH\"\nENV UV_LINK_MODE=copy\n\nRUN bash -ex <<\"EOF\"\n apt-get update\n apt-get install -y --no-install-recommends gettext python3-venv psmisc uuid-runtime\n apt-get clean\n python -m venv /opt/jet\n ARCH=$(uname -m)\n case \"${ARCH}\" in \\\n \"x86_64\") YQ_ARCH=amd64 ;; \\\n \"aarch64\") YQ_ARCH=arm64 ;; \\\n \"armv7l\") YQ_ARCH=arm ;; \\\n *) echo \"Unsupported architecture: ${ARCH}\" && exit 1 ;; \\\n esac \n wget https://github.com/mikefarah/yq/releases/download/v${YQ_VERSION}/yq_linux_${YQ_ARCH} -O /usr/local/bin/yq\n chmod a+x /usr/local/bin/yq\n curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh\nEOF\n\nCOPY README.md pyproject.toml uv.lock /workspace/\nCOPY megatron/core/__init__.py /workspace/megatron/core/\nCOPY megatron/core/package_info.py /workspace/megatron/core/\nARG IMAGE_TYPE=dev\nRUN --mount=type=cache,target=/root/.cache/uv \\\n bash -ex <<\"EOF\"\n export NVTE_CUDA_ARCHS=\"80;90;100\"\n uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages\n uv sync --only-group build\n uv sync --extra ${IMAGE_TYPE} --extra mlm --link-mode copy --locked \\\n --no-install-package torch \\\n --no-install-package torchvision \\\n --no-install-package triton \\\n --no-install-package transformer-engine-cu12 \\\n --no-install-package nvidia-cublas-cu12 \\\n --no-install-package nvidia-cuda-cupti-cu12 \\\n --no-install-package nvidia-cuda-nvrtc-cu12 \\\n --no-install-package nvidia-cuda-runtime-cu12 \\\n --no-install-package nvidia-cudnn-cu12 \\\n --no-install-package nvidia-cufft-cu12 \\\n --no-install-package nvidia-cufile-cu12 \\\n --no-install-package nvidia-curand-cu12 \\\n --no-install-package nvidia-cusolver-cu12 \\\n --no-install-package nvidia-cusparse-cu12 \\\n --no-install-package nvidia-cusparselt-cu12 \\\n --no-install-package nvidia-nccl-cu12\nEOF\n\n# Install DeepEP\nCOPY docker/patches/deepep.patch /workspace/deepep.patch\nRUN bash -ex <<\"EOF\"\n cd /workspace\n uv pip install nvidia-nvshmem-cu13==3.4.5\n pushd /opt/venv/lib/python3.12/site-packages/nvidia/nvshmem/lib/\n ln -s libnvshmem_host.so.3 libnvshmem_host.so\n popd\n\n git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git\n pushd DeepEP\n git checkout eb9cee7de5a24193bf09500668d3a619d3d3f3fb\n patch -p1 < /workspace/deepep.patch\n popd\n TORCH_CUDA_ARCH_LIST=\"9.0 10.0 12.0\" uv pip install --no-build-isolation -v DeepEP/.\n rm -rf DeepEP\nEOF\n\nCOPY assets/ /opt/data/\nENV UV_PYTHON=$UV_PROJECT_ENVIRONMENT/bin/python\n\n##### For NVIDIANS only #####\nFROM main as jet\nARG JET_API_VERSION\nENV PATH=\"$PATH:/opt/jet/bin\"\nRUN --mount=type=secret,id=JET_INDEX_URLS bash -ex <<\"EOF\"\n JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS)\n python -m venv /opt/jet \n /opt/jet/bin/pip install --no-cache-dir $JET_INDEX_URLS \\\n \"jet-api==$JET_API_VERSION\" \"setuptools<82.0.0\"\nEOF\n\nRUN --mount=type=secret,id=JET_INDEX_URLS \\\n --mount=type=secret,id=LOGGER_INDEX_URL bash -ex <<\"EOF\"\n JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS)\n LOGGER_INDEX_URL=$(cat /run/secrets/LOGGER_INDEX_URL)\n uv pip install --no-cache-dir --upgrade $LOGGER_INDEX_URL \"one-logger\"\n uv pip install --no-cache-dir --upgrade \"setuptools<80.0.0,>=77.0.0\"\n uv pip install --no-cache-dir --upgrade $JET_INDEX_URLS \"jet-client~=4.0\" \nEOF\n###\n"
},
{
"path": "docker/Dockerfile.ci.nemo",
"content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# syntax=docker/dockerfile:1.3-labs\n\nARG FROM_IMAGE_NAME\nFROM ${FROM_IMAGE_NAME} as main\n\nRUN apt-get update && \\\n apt-get install -y --no-install-recommends gettext && \\\n apt-get clean && \\\n wget https://github.com/mikefarah/yq/releases/download/v4.44.1/yq_linux_amd64 -O /usr/local/bin/yq && \\\n chmod a+x /usr/local/bin/yq\n\n##### For NVIDIANS only #####\nFROM main as jet\nARG JET_API_VERSION\nRUN --mount=type=secret,id=JET_INDEX_URLS \\\n JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) && \\\n pip install --no-cache-dir jet-api==$JET_API_VERSION \"jet-client~=4.0\" --upgrade $JET_INDEX_URLS\n\nENV PATH=\"$PATH:/opt/jet/bin\"\n###\n"
},
{
"path": "docker/Dockerfile.linting",
"content": "# syntax=docker/dockerfile:experimental\n\nARG FROM_IMAGE_NAME\nFROM $FROM_IMAGE_NAME as main\nENV DEBIAN_FRONTEND=noninteractive\nARG UV_VERSION=0.7.2\nARG YQ_VERSION=4.44.1\nENV PATH=\"/root/.local/bin:$PATH\"\nENV UV_PROJECT_ENVIRONMENT=/opt/venv\nENV PATH=\"$UV_PROJECT_ENVIRONMENT/bin:$PATH\"\nENV UV_LINK_MODE=copy\nRUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh\nWORKDIR /opt/megatron-lm\nCOPY pyproject.toml uv.lock /opt/megatron-lm/\nCOPY megatron/core/package_info.py megatron/core/__init__.py /opt/megatron-lm/megatron/core/\nRUN uv sync --locked --only-group linting --only-group test --only-group ci\n\n##### For NVIDIANS only #####\nFROM main as jet\nARG JET_API_VERSION\nRUN --mount=type=secret,id=JET_INDEX_URLS \\\n JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) && \\ \n uv pip install --no-cache-dir \"jet-client~=2.0\" --upgrade $JET_INDEX_URLS\n"
},
{
"path": "docker/common/install.sh",
"content": "#!/bin/bash\nset -xeuo pipefail # Exit immediately if a command exits with a non-zero status\n\n# Parse command line arguments\nwhile [[ $# -gt 0 ]]; do\n case $1 in\n --base-image)\n BASE_IMAGE=\"$2\"\n shift 2\n ;;\n --python-version)\n PYTHON_VERSION=\"$2\"\n shift 2\n ;;\n --environment)\n ENVIRONMENT=\"$2\"\n shift 2\n ;;\n --use-uv)\n USE_UV=\"true\"\n shift 1\n ;;\n *)\n echo \"Unknown option: $1\"\n echo \"Usage: $0 --base-image {pytorch|ubuntu} [--use-uv] [--python-version] [--environment]\"\n exit 1\n ;;\n esac\ndone\n\nif [[ -z \"${PYTHON_VERSION:-}\" ]]; then\n PYTHON_VERSION=\"3.12\"\nfi\n\nif [[ -z \"${USE_UV:-}\" ]]; then\n USE_UV=\"false\"\nfi\n\n# Validate base image argument\nif [[ -z \"${BASE_IMAGE:-}\" || -z \"${ENVIRONMENT:-}\" ]]; then\n echo \"Error: --base-image argument is required\"\n echo \"Usage: $0 --base-image {pytorch|ubuntu} --environment {dev|lts}\"\n exit 1\nfi\n\nif [[ \"$BASE_IMAGE\" != \"pytorch\" && \"$BASE_IMAGE\" != \"ubuntu\" ]]; then\n echo \"Error: --base-image must be either 'pytorch' or 'ubuntu'\"\n echo \"Usage: $0 --base-image {pytorch|ubuntu}\"\n exit 1\nfi\n\nif [[ \"$ENVIRONMENT\" != \"dev\" && \"$ENVIRONMENT\" != \"lts\" ]]; then\n echo \"Error: --environment must be either 'dev' or 'lts'\"\n echo \"Usage: $0 --environment {dev|lts}\"\n exit 1\nfi\n\nmain() {\n if [[ -n \"${PAT:-}\" ]]; then\n echo -e \"machine github.com\\n login token\\n password $PAT\" >~/.netrc\n chmod 600 ~/.netrc\n fi\n\n # Install dependencies\n export DEBIAN_FRONTEND=noninteractive\n\n # Install Python\n apt-get update\n apt-get install -y software-properties-common\n add-apt-repository ppa:deadsnakes/ppa -y\n apt-get install -y python$PYTHON_VERSION-dev python$PYTHON_VERSION-venv\n update-alternatives --install /usr/bin/python3 python3 /usr/bin/python$PYTHON_VERSION 1\n \n # Install tools\n apt-get update\n apt-get install -y wget curl git cmake\n\n # Install CUDA\n if [[ \"$BASE_IMAGE\" == \"ubuntu\" ]]; then\n rm /etc/apt/sources.list.d/cuda*.list || true\n rm /etc/apt/sources.list.d/nvidia-cuda.list || true\n wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb\n dpkg -i cuda-keyring_1.1-1_all.deb\n rm cuda-keyring_1.1-1_all.deb\n apt-get update\n apt-get install -y cuda-toolkit-12-8 cudnn-cuda-12 libcudnn9-cuda-12 libcutlass-dev \n fi\n\n # Clean up\n apt-get clean\n\n unset PIP_CONSTRAINT\n\n if [[ \"$USE_UV\" == \"true\" ]]; then\n if [[ \"$BASE_IMAGE\" == \"pytorch\" ]]; then\n UV_ARGS=(\n \"--no-install-package\" \"torch\"\n \"--no-install-package\" \"torchvision\"\n \"--no-install-package\" \"triton\"\n \"--no-install-package\" \"nvidia-cublas-cu12\"\n \"--no-install-package\" \"nvidia-cuda-cupti-cu12\"\n \"--no-install-package\" \"nvidia-cuda-nvrtc-cu12\"\n \"--no-install-package\" \"nvidia-cuda-runtime-cu12\"\n \"--no-install-package\" \"nvidia-cudnn-cu12\"\n \"--no-install-package\" \"nvidia-cufft-cu12\"\n \"--no-install-package\" \"nvidia-cufile-cu12\"\n \"--no-install-package\" \"nvidia-curand-cu12\"\n \"--no-install-package\" \"nvidia-cusolver-cu12\"\n \"--no-install-package\" \"nvidia-cusparse-cu12\"\n \"--no-install-package\" \"nvidia-cusparselt-cu12\"\n \"--no-install-package\" \"nvidia-nccl-cu12\"\n )\n else\n UV_ARGS=()\n fi\n \n # Install uv\n UV_VERSION=\"0.7.2\"\n curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh\n\n # Create virtual environment and install dependencies\n uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages\n\n # Install dependencies\n uv sync --locked --only-group build ${UV_ARGS[@]}\n uv sync \\\n --link-mode copy \\\n --locked \\\n --extra ${ENVIRONMENT} \\\n --all-groups ${UV_ARGS[@]}\n\n # Install the package\n uv pip install --no-deps -e .\n else\n python3 -m venv $UV_PROJECT_ENVIRONMENT\n . $UV_PROJECT_ENVIRONMENT/bin/activate\n\n pip install --pre --no-cache-dir --upgrade pip\n pip install --pre --no-cache-dir torch pybind11 wheel_stub ninja wheel packaging \"setuptools<80.0.0,>=77.0.0\"\n pip install --pre --no-cache-dir --no-build-isolation .\n fi\n\n}\n\n# Call the main function\nmain \"$@\"\n"
},
{
"path": "docker/common/install_source_wheels.sh",
"content": "#!/bin/bash\nset -xeuo pipefail # Exit immediately if a command exits with a non-zero status\n\nINPUT_WHEEL_DIR=$(pwd)/wheels\n\n# Parse command line arguments\nwhile [[ $# -gt 0 ]]; do\n case $1 in\n --input-wheel-dir)\n INPUT_WHEEL_DIR=\"$2\"\n shift 2\n ;;\n --environment)\n ENVIRONMENT=\"$2\"\n shift 2\n ;;\n *)\n echo \"Unknown option: $1\"\n echo \"Usage: $0 --input-wheel-dir DIR\"\n exit 1\n ;;\n esac\ndone\n\n# Check if required arguments are provided\nif [ -z \"$INPUT_WHEEL_DIR\" ] || [ -z \"$ENVIRONMENT\" ]; then\n echo \"Error: --input-wheel-dir and --environment are required\"\n echo \"Usage: $0 --input-wheel-dir DIR --environment ENV\"\n exit 1\nfi\n\nif [ \"$ENVIRONMENT\" = \"dev\" ]; then\n TE_WHEEL=$(ls $INPUT_WHEEL_DIR/transformer_engine*.whl) || true\n [ -z \"$TE_WHEEL\" ] && TE_WHEEL=$(bash docker/common/build_te.sh --output-wheel-dir $INPUT_WHEEL_DIR | tail -n 1)\nfi\n\nMAMBA_WHEEL=$(ls $INPUT_WHEEL_DIR/mamba*.whl) || true\n[ -z \"$MAMBA_WHEEL\" ] && MAMBA_WHEEL=$(bash docker/common/build_mamba.sh --output-wheel-dir $INPUT_WHEEL_DIR | tail -n 1)\n\nCAUSALCONV1D_WHEEL=$(ls $INPUT_WHEEL_DIR/causal_conv1d*.whl) || true\n[ -z \"$CAUSALCONV1D_WHEEL\" ] && CAUSALCONV1D_WHEEL=$(bash docker/common/build_causalconv1d.sh --output-wheel-dir $INPUT_WHEEL_DIR | tail -n 1)\n\n# Override deps that are already present in the base image\n# only for dev\nif [ \"$ENVIRONMENT\" = \"dev\" ]; then\n uv pip install --no-cache-dir --no-deps $TE_WHEEL\nfi\n\n# Install heavy optional deps like mamba, causalconv1d\nuv pip install --no-cache-dir \\\n $MAMBA_WHEEL \\\n $CAUSALCONV1D_WHEEL \\\n \"setuptools<80.0.0,>=77.0.0\"\n"
},
{
"path": "docker/patches/deepep.patch",
"content": "diff --git a/setup.py b/setup.py\nindex 63ce332..4e13462 100644\n--- a/setup.py\n+++ b/setup.py\n@@ -37,7 +37,7 @@ if __name__ == '__main__':\n '-Wno-sign-compare', '-Wno-reorder', '-Wno-attributes']\n nvcc_flags = ['-O3', '-Xcompiler', '-O3']\n sources = ['csrc/deep_ep.cpp', 'csrc/kernels/runtime.cu', 'csrc/kernels/layout.cu', 'csrc/kernels/intranode.cu']\n- include_dirs = ['csrc/']\n+ include_dirs = ['csrc/', '/usr/local/cuda/include/cccl/']\n library_dirs = []\n nvcc_dlink = []\n extra_link_args = []\n"
},
{
"path": "docs/add_copyright_header.py",
"content": "#!/usr/bin/env python3\n\"\"\"One-off script to add NVIDIA copyright header to all .md files under docs/.\"\"\"\n\nfrom pathlib import Path\n\nHEADER = \"\"\" Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved.\n NVIDIA CORPORATION and its licensors retain all intellectual property\n and proprietary rights in and to this software, related documentation\n and any modifications thereto. Any use, reproduction, disclosure or\n distribution of this software and related documentation without an express\n license agreement from NVIDIA CORPORATION is strictly prohibited.\n\n\"\"\"\n\ndef main():\n docs_dir = Path(__file__).resolve().parent\n already_has = \"Copyright (c) 2022-2026, NVIDIA CORPORATION\"\n count = 0\n for path in sorted(docs_dir.rglob(\"*.md\")):\n content = path.read_text(encoding=\"utf-8\")\n if content.strip().startswith(already_has):\n continue\n new_content = HEADER + content\n path.write_text(new_content, encoding=\"utf-8\")\n count += 1\n print(path.relative_to(docs_dir))\n print(f\"\\nUpdated {count} files.\")\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "docs/advanced/index.md",
"content": "\n\n# Discussions\n\nIn-depth technical discussions and optimization guides:\n\n- [Optimizing DeepSeek-V3 Training on GB200 NVL72](https://github.com/NVIDIA/Megatron-LM/blob/dev/docs/discussions/deepseek-v3-gb200-optimization/deepseek-v3-gb200-reproduce-guide.md) - Achieving 970 TFLOPS/GPU with MXFP8, kernel optimizations, and HybridEP\n"
},
{
"path": "docs/api-backwards-compatibility-check.md",
"content": "---\norphan: true\n---\n\n\n\n# API Backward Compatibility Checking\n\n## Overview\n\nMegatron Core uses automated API compatibility checking to ensure stable interfaces between releases. This prevents accidental breaking changes that could affect users upgrading between versions.\n\n## How It Works\n\nThe compatibility checker:\n1. Compares the current code against the latest release\n2. Detects breaking changes in function signatures\n3. Fails CI if breaking changes are found (unless explicitly exempted)\n4. Runs automatically on every PR that modifies `megatron/core`\n\n## What Gets Checked\n\n### ✅ Breaking Changes Detected\n\n- **Parameter removed** - Removing a function parameter\n- **Parameter added without default** - Adding a required parameter\n- **Parameter order changed** - Changing the order of parameters\n- **Optional→Required** - Removing a default value from a parameter\n- **Function removed** - Deleting a public function\n- **Return type changed** - Changing the return type annotation (warning)\n\n### ⏭️ What Gets Skipped\n\n- **Test functions** - Functions starting with `test_`\n- **Exempt decorators** - Functions marked with `@internal_api`, `@experimental_api`, or `@deprecated`\n- **Excluded paths** - Code in `tests/`, `experimental/`, `legacy/`\n\n### ✅ Allowed Changes\n\n- **Adding optional parameters** - Adding parameters with default values\n- **Adding new functions** - New public APIs\n- **Making parameters optional** - Adding default values to required parameters\n\n## For Developers\n\n### Running Locally\n\n```bash\n# Install griffe\npip install griffe\n\n# Check against latest release\npython scripts/check_api_backwards_compatibility.py --baseline core_r0.8.0\n\n# Check with verbose output\npython scripts/check_api_backwards_compatibility.py --baseline core_r0.8.0 -v\n\n# Compare two specific branches\npython scripts/check_api_backwards_compatibility.py --baseline core_r0.8.0 --current main\n```\n\n### Marking Functions as Exempt\n\nIf you need to make breaking changes to internal or experimental APIs:\n\n#### Internal API (for internal implementation details)\n\n```python\nfrom megatron.core.utils import internal_api\n\n@internal_api\ndef experimental_feature(x, y):\n \"\"\"\n This API is experimental and may change.\n NOT FOR EXTERNAL USE.\n \"\"\"\n pass\n```\n\n**When to use `@internal_api`:**\n- Internal APIs not documented for external use\n- Experimental features explicitly marked as unstable\n- Functions in development that haven't been released yet\n\n#### Experimental API (for experimental features)\n\n```python\nfrom megatron.core.utils import experimental_api\n\n@experimental_api\ndef new_experimental_feature(x, y):\n \"\"\"\n This API is experimental and may change without notice.\n \"\"\"\n pass\n```\n\n**When to use `@experimental_api`:**\n- Experimental features explicitly marked as unstable\n- New APIs under active development\n- Features that haven't been stabilized yet\n\n### Deprecating APIs\n\nFor planned API changes, use the deprecation workflow:\n\n```python\nfrom megatron.core.backwards_compatibility_decorators import deprecated\n\n@deprecated(\n version=\"1.0.0\", # When deprecation starts\n removal_version=\"2.0.0\", # When it will be removed\n alternative=\"new_function\", # Recommended replacement\n reason=\"Improved performance and cleaner API\"\n)\ndef old_function(x):\n \"\"\"This function is deprecated.\"\"\"\n pass\n```\n\n**Deprecation Timeline:**\n1. **Version N** - Add `@deprecated` decorator, function still works\n2. **Version N+1** - Keep function with deprecation warnings\n3. **Version N+2** - Remove function (users have been warned)\n\n### Handling CI Failures\n\nIf the compatibility check fails on your PR:\n\n1. **Review the breaking changes** in the CI logs\n2. **Choose an action:**\n - **Fix the code** - Revert the breaking change\n - **Add exemption** - Use `@internal_api` if intentional\n - **Use deprecation** - For planned API changes\n3. **Update your PR** with the fix\n\n## Examples\n\n### Example 1: Compatible Change\n\n```python\n# ✅ BEFORE (v1.0)\ndef train_model(config, dataloader):\n pass\n\n# ✅ AFTER (v1.1) - Added optional parameter\ndef train_model(config, dataloader, optimizer=\"adam\"):\n pass\n```\n**Result:** ✅ Check passes\n\n---\n\n### Example 2: Breaking Change\n\n```python\n# BEFORE (v1.0)\ndef train_model(config, dataloader, optimizer=\"adam\"):\n pass\n\n# ❌ AFTER (v1.1) - Removed parameter\ndef train_model(config, dataloader):\n pass\n```\n**Result:** ❌ Check fails - \"Parameter 'optimizer' removed\"\n\n---\n\n### Example 3: Exempt Internal API\n\n```python\nfrom megatron.core.utils import internal_api\n\n# BEFORE (v1.0)\n@internal_api\ndef _internal_compute(x, y):\n pass\n\n# ✅ AFTER (v1.1) - Can change freely\n@internal_api\ndef _internal_compute(x, y, z): # Added parameter\n pass\n```\n**Result:** ✅ Check passes (function is exempt)\n\n---\n\n### Example 4: Deprecation Workflow\n\n```python\nfrom megatron.core.utils import deprecated\n\n# Version 1.0 - Add deprecation\n@deprecated(\n version=\"1.0.0\",\n removal_version=\"2.0.0\",\n alternative=\"train_model_v2\"\n)\ndef train_model(config):\n \"\"\"Old training function - DEPRECATED\"\"\"\n pass\n\ndef train_model_v2(config, **options):\n \"\"\"New improved training function\"\"\"\n pass\n\n# Version 1.1 - Keep both (users migrate)\n# Version 2.0 - Remove train_model()\n```\n\n## Architecture\n\n```\nDeveloper commits code\n ↓\nGitHub Actions triggers\n ↓\nCI runs check_api_backwards_compatibility.py\n ↓\nScript loads code via griffe:\n • Baseline: latest release (e.g., core_r0.8.0)\n • Current: PR branch\n ↓\nApply filtering:\n • Skip @internal_api, @experimental_api, and @deprecated\n • Skip private functions (_prefix)\n • Skip test/experimental paths\n ↓\nGriffe compares signatures:\n • Parameters\n • Types\n • Return types\n • Defaults\n ↓\nReport breaking changes\n ↓\nExit: 0=pass, 1=fail\n ↓\nCI fails if breaking changes detected\n```\n\n## Configuration\n\n### Customizing Filters\n\nEdit `scripts/check_api_backwards_compatibility.py`:\n\n```python\n# Add more exempt decorators\nEXEMPT_DECORATORS = [\n \"internal_api\",\n \"experimental_api\",\n \"deprecated\",\n]\n\n# Add more path exclusions\nEXCLUDE_PATHS = {\n \"tests\",\n \"experimental\",\n \"legacy\",\n \"your_custom_path\", # ← Add here\n}\n```\n\n### Changing the Baseline\n\nThe workflow auto-detects the latest `core_r*` tag. To manually specify:\n\n```yaml\n# In .github/workflows/check_api_backwards_compatibility_workflow.yml\n- name: Run compatibility check\n run: |\n python scripts/check_api_backwards_compatibility.py \\\n --baseline your_custom_baseline\n```\n\n## FAQ\n\n### Q: Why did my PR fail the compatibility check?\n\n**A:** Your code introduced breaking changes compared to the last release. Review the CI logs to see what changed.\n\n### Q: Can I disable the check for my PR?\n\n**A:** No, but you can mark specific functions as exempt using `@internal_api` or `@experimental_api`.\n\n### Q: What if I need to make a breaking change?\n\n**A:** Use the `@deprecated` decorator for a gradual transition, or mark the function as exempt using `@internal_api` (for internal code) or `@experimental_api` (for experimental features).\n\n### Q: Does this check all of Megatron-LM?\n\n**A:** No, only `megatron/core/**` (Megatron Core). Legacy code is excluded.\n\n### Q: What about class methods?\n\n**A:** Yes, class methods are checked just like functions.\n\n### Q: Can I run this locally before pushing?\n\n**A:** Yes! Run `python scripts/check_api_backwards_compatibility.py --baseline core_r0.8.0`\n\n### Q: What if there's no release tag yet?\n\n**A:** The workflow will use `main` as the baseline. Update it once you have release tags.\n\n## Troubleshooting\n\n### Error: \"griffe is not installed\"\n\n```bash\npip install griffe\n```\n\n### Error: \"No core_r* tags found\"\n\nThe repository doesn't have release tags yet. The workflow will fall back to `main`.\n\n### False Positives\n\nIf the checker reports a breaking change that isn't actually breaking, file an issue and use `@internal_api` as a temporary workaround.\n\n## References\n\n- **Script:** `scripts/check_api_backwards_compatibility.py`\n- **Workflow:** `.github/workflows/check_api_backwards_compatibility_workflow.yml`\n- **Decorators:** `megatron/core/backwards_compatibility_decorators.py`\n- **Griffe Documentation:** https://mkdocstrings.github.io/griffe/\n\n## Support\n\nFor questions or issues:\n1. Check this documentation\n2. Review existing PRs with compatibility checks\n3. Ask in the Megatron-LM Slack/Discord\n4. File an issue on GitHub\n\n"
},
{
"path": "docs/api-guide/core/datasets.md",
"content": "\n\n# datasets package\n\n```{include} ../../../megatron/core/datasets/readme.md\n```\n"
},
{
"path": "docs/api-guide/core/dist_checkpointing.md",
"content": "\n\n# dist_checkpointing package\n\nA library for saving and loading the distributed checkpoints.\nA *distributed checkpoint* in Megatron Core uses the ``torch_dist`` format,\na custom checkpointing mechanism built on top of PyTorch's native\ncheckpointing capabilities.\n\nA key property of distributed checkpoints is that a checkpoint saved under one\nparallel configuration (tensor, pipeline, or data parallelism) can be loaded\nunder a different parallel configuration. This enables flexible scaling and\nresharding of models across heterogeneous training setups.\n\nUsing the library requires defining sharded state_dict dictionaries with functions from *mapping* and *optimizer* modules.\nThose state dicts can be saved or loaded with a *serialization* module using strategies from *strategies* module.\n\n## Safe Checkpoint Loading\n\nSince **PyTorch 2.6**, the default behavior of `torch.load` is `weights_only=True`.\nThis ensures that only tensors and allow-listed classes are loaded, reducing the risk of arbitrary code execution.\n\nIf you encounter an error such as:\n\n```bash\nWeightsUnpickler error: Unsupported global: GLOBAL argparse.Namespace was not an allowed global by default.\n```\n\nyou can fix it by explicitly allow-listing the missing class in your script:\n\n```python\nimport torch, argparse\n\ntorch.serialization.add_safe_globals([argparse.Namespace])\n```\n\n## Checkpointing Distributed Optimizer\n\n### Checkpoint Compatibility and Optimizer State Formats\n\nBeginning with **mcore v0.14**, the ``flattened_range`` attribute was removed from ``dist_checkpointing``. As a result:\n\n- 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.\n- 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.\n\n### Workaround: Loading legacy optimizer states with ToT MCore\n\n**Step 1: Convert the legacy checkpoint using mcore v0.15.0**\n\nRun a dummy training job with mcore v0.15.0 to re-save the checkpoint with new optimizer states format.\n\n```bash\nMODEL_TRAIN_PARAMS=(\n # Define model architecture and training parameters here\n)\nOLD_CKPT=/workspace/mcore_ckpt_old\nCONVERTED_CKPT=/workspace/mcore_ckpt_0.15.0\n\ntorchrun --nproc_per_node=8 /opt/megatron-lm/pretrain_gpt.py \\\n --save-interval 1 \\\n --eval-interval 1 \\\n --exit-interval 1 \\\n --eval-iters 1 \\\n --use-distributed-optimizer \\\n --save ${CONVERTED_CKPT} \\\n --load ${OLD_CKPT} \\\n --ckpt-format torch_dist \\\n \"${MODEL_TRAIN_PARAMS[@]}\"\n```\n\n**Step 2: Load the converted checkpoint with ToT MCore**\n\nUse the converted checkpoint as the input for continued training with ToT MCore.\n\n```bash\nMODEL_TRAIN_PARAMS=(\n # Define model architecture and training parameters here\n)\nNEW_CKPT=/workspace/mcore_ckpt_new\nCONVERTED_CKPT=/workspace/mcore_ckpt_0.15.0\n\ntorchrun --nproc_per_node=8 /opt/megatron-lm/pretrain_gpt.py \\\n --use-distributed-optimizer \\\n --save ${NEW_CKPT} \\\n --load ${CONVERTED_CKPT} \\\n --ckpt-format torch_dist \\\n \"${MODEL_TRAIN_PARAMS[@]}\"\n```\n\nAfter this step, training can proceed normally using ToT MCore with fully supported optimizer state loading.\n\n## Distributed Optimizer Checkpoint Formats\n\nThe refactor of the Distributed Optimizer introduces **two checkpoint formats**:\n\n- dp_reshardable (Default)\n - Fast save/load performance.\n - Not reshardable — not possible to change model parallelism when using this format.\n - Recommended for general training when model parallelism changes are not needed.\n- fully_reshardable\n - Fully reshardable — supports arbitrary changes in model parallelism.\n - Slower than dp_reshardable.\n - Enabled via the ``--dist-ckpt-optim-fully-reshardable`` flag.\n\n### Workflow for Changing Model Parallelism\n\nYou can combine formats to optimize both flexibility and performance:\n\n 1. Train using ``dp_reshardable`` (default) for faster checkpointing.\n 2. When you need to change model parallelism:\n\n - Stop training.\n - Change model parallelism for train config.\n - Resume training with ``--dist-ckpt-optim-fully-reshardable``.\n\n 3. Save at least one checkpoint under the new model parallel configuration.\n 4. (Optional) To continue the training with updated model parallelism and better checkpointing performance, stop training and switch back to ``dp_reshardable`` format by removing ``--dist-ckpt-optim-fully-reshardable``.\n\n## Subpackages\n\n```{toctree}\n:maxdepth: 4\n\ndist_checkpointing.strategies\n```\n\n"
},
{
"path": "docs/api-guide/core/dist_checkpointing.strategies.md",
"content": "\n\n# dist_checkpointing.strategies package\n\nPackage defining different checkpoint formats (backends) and saving/loading algorithms (strategies).\n\nStrategies can be used for implementing new checkpoint formats or implementing new (more optimal for a given use case) ways of saving/loading of existing formats.\nStrategies are passed to `dist_checkpointing.load` and `dist_checkpointing.save` functions and control the actual saving/loading procedure.\n\n"
},
{
"path": "docs/api-guide/core/distributed.md",
"content": "\n\n# distributed package\n\nThis package contains various utilities to finalize model weight gradients\non each rank before the optimizer step. This includes a distributed data\nparallelism wrapper to all-reduce or reduce-scatter the gradients across\ndata-parallel replicas, and a `finalize_model_grads` method to\nsynchronize gradients across different parallelism modes (e.g., 'tied'\nlayers on different pipeline stages, or gradients for experts in a MoE on\ndifferent ranks due to expert parallelism).\n\n"
},
{
"path": "docs/api-guide/core/fusions.md",
"content": "\n\n# fusions package\n\nThis package provides modules that provide commonly fused\noperations. Fusing operations improves compute efficiency by\nincreasing the amount of work done each time a tensor is read from\nmemory. To perform the fusion, modules in this either rely on PyTorch\nfunctionality for doing just-in-time compilation\n(i.e. `torch.jit.script` in older PyTorch versions of `torch.compile`\nin recent versions), or call into custom kernels in external libraries\nsuch as Apex or TransformerEngine.\n\n"
},
{
"path": "docs/api-guide/core/index.md",
"content": "\n\n# Core APIs\n\nLow-level API reference for core Megatron components.\n\n```{toctree}\n:maxdepth: 2\n\ntransformer\ntensor_parallel\npipeline_parallel\nfusions\ndistributed\ndatasets\ndist_checkpointing\ndist_checkpointing.strategies\n```\n"
},
{
"path": "docs/api-guide/core/pipeline_parallel.md",
"content": "\n\n# pipeline_parallel package\n\nThis package contains implementations for two different pipeline parallelism\nschedules (one without interleaving and one with interleaving, see [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://arxiv.org/abs/2104.04473)\nfor details), and a default no-pipelining schedule. It also contains methods\nfor the point-to-point communication that is needed between pipeline stages.\n\n"
},
{
"path": "docs/api-guide/core/tensor_parallel.md",
"content": "\n\n# tensor_parallel package\n\nThis package contains an implementation for tensor parallelism in transformer\nmodels (see [Megatron-LM: Training Multi-Billion Parameter Language Models\nUsing Model Parallelism](https://arxiv.org/abs/1909.08053) and [Reducing\nActivation Recomputation in Large Transformer Models](https://arxiv.org/abs/2205.05198)\nfor details).\n\n"
},
{
"path": "docs/api-guide/core/transformer.md",
"content": "\n\n# transformer package\n\nThe `transformer` package provides a customizable and configurable\nimplementation of the transformer model architecture. Each component\nof a transformer stack, from entire layers down to individual linear\nlayers, can be customized by swapping in different PyTorch modules\nusing the \"spec\" parameters. The\nconfiguration of the transformer (hidden size, number of layers,\nnumber of attention heads, etc.) is provided via a `TransformerConfig`\nobject.\n"
},
{
"path": "docs/api-guide/index.md",
"content": "\n\n# API Guide\n\nAPI reference documentation for Megatron Core components.\n\n```{toctree}\n:maxdepth: 3\n\nmodels/index\ncore/index\ninternal/index\n```\n"
},
{
"path": "docs/api-guide/internal/index.md",
"content": "\n\n# Internal Utilities\n\nInternal utility APIs.\n\n```{toctree}\n:maxdepth: 2\n\nnum_microbatches_calculator\noptimizer_param_scheduler\n```\n"
},
{
"path": "docs/api-guide/internal/num_microbatches_calculator.md",
"content": "\n\n# Microbatches Calculator\n\nThis api is used to calculate the number of microbatches required to fit a given model on a given batch size.\n\n"
},
{
"path": "docs/api-guide/internal/optimizer_param_scheduler.md",
"content": "\n\n# Optimizer Parameters Scheduler\n\nThis api is used to calculate the learning rate and weight decay for the optimizer.\n\n"
},
{
"path": "docs/api-guide/models/index.md",
"content": "\n\n# Model APIs\n\nAPI reference for Megatron Core model implementations.\n\n```{toctree}\n:maxdepth: 2\n\nmodels\nmodels.gpt\nmodels.bert\nmodels.t5\n```\n"
},
{
"path": "docs/api-guide/models/models.bert.md",
"content": "\n\n# models.bert package\n\nUseful package for training bert and bert like encoder only models. It optionally comes with a binary head that can be used for classification tasks .\n\n"
},
{
"path": "docs/api-guide/models/models.gpt.md",
"content": "\n\n# models.gpt package\n\nThis 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.\n\n"
},
{
"path": "docs/api-guide/models/models.md",
"content": "\n\n# models package\n\nThis 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.\n\n## Subpackages\n\n```{toctree}\n:maxdepth: 4\n\nmodels.gpt\nmodels.t5\nmodels.bert\n```\n\n"
},
{
"path": "docs/api-guide/models/models.t5.md",
"content": "\n\n# models.t5 package\n\n"
},
{
"path": "docs/api-guide/router_replay.md",
"content": "\n\n# Design Document: MoE Router Replay Feature\n\n## 1. Overview\n\nThis document provides a detailed description of the \"Router Replay\" feature implemented within the Megatron-LM Core for Mixture-of-Experts (MoE) models.\n\nThis 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.\n\n## 2. Motivation\n\n* **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.\n* **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.\n* **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.\n\n## 3. Design and Architecture\n\nThe 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.\n\n* **Core Components**:\n * `RouterReplay` (located in `megatron/core/transformer/moe/router_replay.py`): A utility class for replaying MoE routing decisions. When enabled via the `moe_enable_routing_replay` flag, a separate instance of `RouterReplay` is created for each MoE layer's router. Each instance is responsible for loading routing data and providing the deterministic routing decisions for its corresponding layer during the forward pass.\n * `moe_enable_routing_replay` (located in `megatron/core/transformer/transformer_config.py`): A boolean global configuration flag that serves as the sole entry point for enabling this feature.\n\n* **Workflow**:\n The feature supports different modes, such as recording and replaying, controlled by a `RouterReplayAction`.\n\n 1. **Enabling the Feature**: The user sets `moe_enable_routing_replay` to `True` in the model configuration.\n 2. **Initialization**: When `moe_enable_routing_replay` is true, each `TopKRouter` creates its own `RouterReplay` instance.\n 3. **Mode Configuration**: The user must programmatically set the desired router replay action (e.g., `record`, `forward_replay`, `backward_replay`) on the `RouterReplay` instances.\n 4. **Execution Flow (within a mini-batch)**:\n * **Forward Pass**:\n * For each micro-batch, the `topk_routing_with_score_function` checks the `router_replay_action`.\n * **In `record` mode**: The dynamically computed `top-k` expert indices are captured and stored.\n * **In `forward_replay` mode**: The function retrieves pre-loaded expert indices from `target_topk_idx`. These indices are used for the forward computation and are also appended to the `replay_backward_list` to prepare for the backward pass.\n * **Backward Pass**:\n * For each micro-batch (processed in reverse order in pipeline parallelism), the `router_replay_action` is checked again.\n * **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.\n\n## 4. Implementation Details\n\nThe implementation cleanly separates the replay logic from the router's core computation.\n\n* **`megatron/core/transformer/transformer_config.py`**:\n * Adds the configuration option `moe_enable_routing_replay: bool = False`.\n\n* **`megatron/core/transformer/moe/moe_utils.py`**:\n * Introduces the `RouterReplay` class to manage the state for recording and replaying routing decisions for a single MoE layer.\n * `target_topk_idx`: An attribute holding the expert indices for the current micro-batch during forward replay mode.\n * `recorded_topk_idx`: An attribute for storing the computed expert indices when in record mode.\n * `replay_backward_list`: A list that accumulates the `top-k` indices used during the forward passes of a mini-batch. This list is consumed in FIFO order during the backward pass to ensure correctness under pipeline parallelism.\n * `set_target_indices()`: A method to load the replay indices into `target_topk_idx` for the forward pass.\n * `record_indices()`: A method to save the computed indices.\n * 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.\n\n### Training recompute usage\n\n- During forward replay, `set_target_indices()` prepares `replay_backward_list` so each micro-batch’s indices are available for recomputation.\n- During recompute/backward, set action to `REPLAY_BACKWARD` so indices are consumed in FIFO order to mirror the forward sequence.\n\n## 5. Usage Guide\n\n1. **Enable & Instantiate**\n - Create one `RouterReplay` instance per MoE router layer when building the model.\n - Optionally use the global helpers to set/clear actions across all layers.\n2. **Record Routing Decisions**\n - Set action: `RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)`.\n - Run the model; retrieve per-layer indices via `RouterReplay.get_recorded_data()` and persist.\n3. **Forward Replay**\n - Load indices and distribute: `RouterReplay.set_replay_data(list_of_tensors)`.\n - Set action: `RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)`.\n - Run the model; dynamic top‑k is bypassed and target indices are used.\n4. **Backward Replay**\n - For training recomputation (activation checkpointing or pipeline recompute), set action: `REPLAY_BACKWARD` during recomputation.\n - Per micro‑batch indices are consumed from `replay_backward_list` in FIFO order.\n5. **Cleanup**\n - 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.\n\n### Quick usage with `topk_routing_with_score_function`\n\n```python\nimport torch\nfrom megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction\nfrom megatron.core.transformer.moe.moe_utils import topk_routing_with_score_function\n\nrr = RouterReplay()\n\n# Record\nRouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)\nlogits = torch.randn(8, 16)\nprobs_rec, routing_map_rec = topk_routing_with_score_function(\n logits=logits, topk=2, use_pre_softmax=False, score_function=\"softmax\", router_replay=rr,\n)\nrecorded = rr.get_recorded_indices()\ntorch.save(recorded, \"/tmp/replay.pt\")\n\n# Forward replay\nrr.clear_router_replay_action()\nrr.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD)\ntarget = torch.load(\"/tmp/replay.pt\")\nrr.set_target_indices(target)\nprobs_rep, routing_map_rep = topk_routing_with_score_function(\n logits=logits, topk=2, use_pre_softmax=False, score_function=\"softmax\", router_replay=rr,\n)\n\nRouterReplay.clear_global_router_replay_action()\nRouterReplay.clear_global_indices()\nRouterReplay.clear_global_router_replay_instances()\n```\n\n## 6. Minimal Demo\n\nHere is a minimal code example showing how to use RouterReplay for recording and replaying:\n\n```python\nimport torch\nimport torch.distributed as dist\nfrom megatron.core.transformer.transformer_config import TransformerConfig\nfrom megatron.core.transformer.moe.router import TopKRouter\nfrom megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction\n\n\n# Initialize distributed training\nif not dist.is_initialized():\n dist.init_process_group(backend=\"nccl\")\n\n# Create a transformer config with RouterReplay enabled\nconfig = TransformerConfig(\n num_experts=8,\n expert_model_parallel_size=1,\n num_top_k=2,\n moe_enable_routing_replay=True\n)\n\n# Create a TopKRouter instance\nrouter = TopKRouter(config)\n\n# Generate sample input (batch_size, sequence_length, hidden_size)\nlogits = torch.randn(16, 32, 8).to(torch.cuda.current_device())\n\n# -----------------\n# 1. Recording Mode\n# -----------------\nprint(\"=== Recording Mode ===\")\n# Set global router replay action to RECORD\nRouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)\n\n# Perform routing\nrouting_output = router.forward(logits)\nprint(f\"Recorded top-k indices shape: {routing_output.top_k_idx.shape}\")\n\n# -----------------\n# 2. Forward Replay Mode\n# -----------------\nprint(\"\\n=== Forward Replay Mode ===\")\n# Save recorded indices to a file\ntorch.save(routing_output.top_k_idx, \"/tmp/replay.pt\")\n\n# Load indices from file and set as target for replay\nreplay_indices = torch.load(\"/tmp/replay.pt\")\nfor router_instance in RouterReplay.global_router_replay_instances:\n router_instance.target_topk_idx = replay_indices\n\n# Set global router replay action to REPLAY_FORWARD\nRouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)\n\n# Perform routing again - this will use the replayed indices\nreplay_routing_output = router.forward(logits)\nprint(f\"Replayed top-k indices shape: {replay_routing_output.top_k_idx.shape}\")\nprint(f\"Are indices the same? {torch.equal(routing_output.top_k_idx, replay_routing_output.top_k_idx)}\")\n\n\n# Clean up\nRouterReplay.clear_global_router_replay_action()\nRouterReplay.clear_global_indices()\nRouterReplay.clear_global_router_replay_instances()\nif dist.is_initialized():\n dist.destroy_process_group()\n```\n"
},
{
"path": "docs/autodoc2_docstrings_parser.py",
"content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom docutils import nodes\nfrom myst_parser.parsers.sphinx_ import MystParser\nfrom sphinx.ext.napoleon.docstring import GoogleDocstring\n\n\nclass NapoleonParser(MystParser):\n \"\"\"Add support for Google style docstrings.\"\"\"\n\n def parse(self, input_string: str, document: nodes.document) -> None:\n \"\"\"Parse Google style docstrings.\"\"\"\n\n # Get the Sphinx configuration\n config = document.settings.env.config\n\n # Process with Google style\n google_parsed = str(GoogleDocstring(input_string, config))\n\n return super().parse(google_parsed, document)\n\n\nParser = NapoleonParser\n"
},
{
"path": "docs/broken_links_false_positives.json",
"content": "{\n \"uri\": \"http://localhost:8080/\"\n}"
},
{
"path": "docs/conf.py",
"content": "# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Configuration file for the Sphinx documentation builder.\n#\n# For the full list of built-in configuration values, see the documentation:\n# https://www.sphinx-doc.org/en/master/usage/configuration.html\n\nimport os\nimport sys\n\n\n# -- Project information -----------------------------------------------------\n# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information\n\nproject = \"Megatron Core\"\ncopyright = \"2026, NVIDIA Corporation\"\nauthor = \"NVIDIA Corporation\"\nrelease = \"nightly\"\n\n# -- General configuration ---------------------------------------------------\n# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration\n\nextensions = [\n \"myst_parser\", # For our markdown docs\n \"sphinx.ext.viewcode\", # For adding a link to view source code in docs\n \"sphinx.ext.doctest\", # Allows testing in docstrings\n \"sphinx.ext.napoleon\", # For google style docstrings\n \"sphinx_copybutton\", # For copy button in code blocks\n]\n\n# Check if we should skip autodoc generation\n# usage: SKIP_AUTODOC=true\nskip_autodoc = os.environ.get(\"SKIP_AUTODOC\", \"false\").lower() == \"true\"\n\nif not skip_autodoc:\n extensions.append(\"autodoc2\") # Generates API docs\n\ntemplates_path = [\"_templates\"]\nexclude_patterns = [\"_build\", \"Thumbs.db\", \".DS_Store\"]\n\n# -- Options for MyST Parser (Markdown) --------------------------------------\n# MyST Parser settings\nmyst_enable_extensions = [\n \"dollarmath\", # Enables dollar math for inline math\n \"amsmath\", # Enables LaTeX math for display mode\n \"colon_fence\", # Enables code blocks using ::: delimiters instead of ```\n \"deflist\", # Supports definition lists with term: definition format\n \"fieldlist\", # Enables field lists for metadata like :author: Name\n \"tasklist\", # Adds support for GitHub-style task lists with [ ] and [x]\n \"attrs_block\", # Enables setting attributes on block elements using {#id .class key=val}\n]\nmyst_heading_anchors = 5 # Generates anchor links for headings up to level 5\n\n# Suppress \"more than one target found for cross-reference\" warnings for Python symbols\n# that have the same name across multiple modules (e.g. DistributedDataParallelConfig,\n# ModelType). These are structural ambiguities in the codebase – the cross-reference\n# still resolves; Sphinx just cannot pick the unique target automatically.\nsuppress_warnings = [\"ref.python\"]\n\n# -- Options for Autodoc2 ---------------------------------------------------\nsys.path.insert(0, os.path.abspath(\"..\"))\n\nif not skip_autodoc:\n autodoc2_packages = [\n {\n \"path\": \"../megatron/core\", # Path to your package relative to conf.py\n \"exclude_dirs\": [\"converters\"], # list of directory names to exclude\n }\n ]\n autodoc2_render_plugin = \"myst\" # Use MyST for rendering docstrings\n autodoc2_output_dir = \"apidocs\" # Output directory for autodoc2 (relative to docs/)\n # This is a workaround that uses the parser located in autodoc2_docstrings_parser.py to allow autodoc2 to\n # render google style docstrings.\n # Related Issue: https://github.com/sphinx-extensions2/sphinx-autodoc2/issues/33\n autodoc2_docstring_parser_regexes = [\n (r\".*\", \"docs.autodoc2_docstrings_parser\"),\n ]\n # Regex patterns whose values contain raw regex syntax (e.g. \\p{L}) that docutils\n # mis-parses as footnote/reference markup. Exclude them from the generated docs.\n autodoc2_hidden_regexes = [\n r\".*\\._PATTERN_TIKTOKEN.*\",\n ]\n\n# -- Options for HTML output -------------------------------------------------\n# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output\n\nhtml_theme = \"nvidia_sphinx_theme\"\nhtml_theme_options = {\n \"switcher\": {\n \"json_url\": \"versions1.json\",\n \"version_match\": release,\n },\n \"icon_links\": [\n {\n \"name\": \"GitHub\",\n \"url\": \"https://github.com/NVIDIA/Megatron-LM/\",\n \"icon\": \"fa-brands fa-github\",\n }\n ],\n \"public_docs_features\": True\n}\nhtml_extra_path = [\"project.json\", \"versions1.json\"]\n\n# Github links are now getting rate limited from the Github Actions\nlinkcheck_ignore = [\n \".*github\\\\.com.*\",\n \".*githubusercontent\\\\.com.*\",\n]\n"
},
{
"path": "docs/developer/contribute.md",
"content": "\n\n# Contributing to Megatron-LM\n\nThis document outlines the processes and policies for issues and pull requests by non-NVIDIA contributors to the Megatron-LM GitHub repository.\n\nEveryone is welcome to contribute to the project! We recently migrated from using an internal repo to doing all development directly from the GitHub repository.\n\nWhen contributing it is important to ensure that changes are in line with the project direction. Small changes to fix bugs are welcomed and appreciated. If proposing large architectural changes or changes for stylistic reasons open an issue first so we can discuss it.\n\n## Issue policy\n\nPlease do file any bugs you find, keeping the following in mind:\n\n- If filing a bug, i.e. you have found something that doesn't work as expected, use the BUG template.\n- If you've found a regression in speed or accuracy use the REGRESSION template.\n- If you are requesting a new feature or modification of an existing feature use the ENHANCEMENT template.\n- If opening an issue to ask a question no template is needed but please make your question as clear and concise as possible.\n- One issue per bug. Putting multiple things in the same issue makes both discussion and completion unnecessarily complicated.\n- Your bug is mostly likely to get attention from the development team quickly if we can easily reproduce it.\n- Use proper spelling, grammar, and punctuation.\n- Write in an authoritative and technical tone.\n\n## Code submission policy\n\n### Do\n\n- Format new code in a style that is consistent with the file being changed. Megatron-LM doesn't (yet) have a style guide or enforced formatting.\n- Split your changes into separate, atomic commits i.e. A commit per feature or fix.\n- Make sure your commits are rebased on the master branch.\n- Write the commit message subject line in the imperative mood (\"Change the default argument for X\", not \"Changed the default argument for X\").\n- Write your commit messages in proper English, with care and punctuation.\n- Check the spelling of your code, comments and commit messages.\n\n### Don't\n\n- Submit code that's incompatible with the project licence.\n- Touch anything outside the stated scope of the PR. This includes formatting changes to code not relevant to the PR.\n- Iterate excessively on your design across multiple commits.\n- Include commented-out code.\n- Attempt large architectural changes without first opening an issue to discuss.\n\n## Issue and Pull Request Q&A\n\n### I've submitted an issue and PR. When can I expect to get some feedback?\n\nYou should receive a response within 2 business days.\n\n### I need help, who should I ping?\n\nUse [@mcore-oncall](https://github.com/orgs/NVIDIA/teams/mcore-oncall).\n\n### If my issue or PR isn't getting attention, what should I do?\n\nAfter 2 business days, tag the user [@mcore-oncall](https://github.com/orgs/NVIDIA/teams/mcore-oncall).\n\n### Is there a policy for issues and PRs that haven't been touched in X days? Should they be closed?\n\nYes, we have a bot that will mark untouched PRs as \"stale\" after 60 days.\n\nWe have a long backlog of issues and PRs dating back years. We are trying to triage these now by working backwards. Older issues we believe may still be relevant may recieve a request to re-test them with the latest code. If there's no response they may be closed. Again, if you they should be re-opened then just respond with a comment to that effect.\n\nThank you!"
},
{
"path": "docs/developer/generate_docs.md",
"content": "\n\n# Generating Docs Locally\n\nTo generate docs locally, use the following commands:\n\n```\ncd docs\nuv run --only-group docs sphinx-autobuild . _build/html --port 8080 --host 127.0.0.1\n```\n\nDocs will be generated at .\n\n**Recommended:** set the environment variable `SKIP_AUTODOC=true` when generating docs \nto skip the generation of `apidocs`."
},
{
"path": "docs/developer/oncall.md",
"content": "\n-->\n\n# Oncall Overview\n\nDuring your oncall week, you will be assigned to all PRs marked “Ready for \nReview”. From a high-level, your responsibilities include:\n\n- Review all new PRs\n- Accelerate the review process\n- Ensure issues and discussion questions are answered\n\n## PR Responsibilities\n\nBelow is the checklist that the oncall needs to go through for each PR.\n\n- Should the PR remain a single PR?\n - Each PR should have at most 1 expert reviewer, although there will be some outlier cases\n- Label PR as “complexity: low”, “complexity: medium”, or “complexity: high” depending on complexity\n - Expert reviewers have final say, oncall just sets the initial complexity level\n - Initial complexity level guideline\n - Low: <100 lines changed\n - Medium: 100 < lines changed < 500\n - High: > 500 lines changed\n- Does this PR have proper testing coverage?\n - If new logic is added, is the new logic tested?\n- Should the PR add documentation for any new features?\n- Does the PR conform to our style guidelines?\n - Code structure\n - Cleanliness\n - Comments\n - File structure\n- Do all tests pass?\n - Oncall will need to kick off testing suite for external reviewers\n - Comment “/ok to test commid_id” to kick off testing suite\n- Expert reviewers are notified after the PR is marked “Ready for Review”\n - **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.\n - If the reviewer is not responding after 2 business days, escalate to the reviewer’s manager.\n- For `megatron/core` PRs, the “Final Review” label is applied automatically once all expert reviewers approve\n - Final reviewers should review within 1 business day. Message the assigned reviewer if it is taking longer.\n - If the reviewer is not responding after 2 business days, escalate to the reviewer’s manager.\n- The “Approved” label is applied automatically once all required reviewers have approved\n\n## Issues and Discussion Questions\n\nIf you do not know the answer to an issue or discussion question: that's ok! **Delegate to someone who does.**\n\nOn a daily basis, track the following:\n\n- [new issues](https://github.com/NVIDIA/Megatron-LM/issues): check to see if there are any new issues before they become out of SLA!\n- [out of SLA issues](https://github.com/orgs/NVIDIA-NeMo/projects/20/views/4?sliceBy%5Bvalue%5D=NVIDIA%2FMegatron-LM): useful dashboard that tracks all out of SLA issues\n"
},
{
"path": "docs/developer/submit.md",
"content": "\n\n# How to Submit a PR\n\nAll PRs start as **draft**. If you open a non-draft PR, it will be automatically converted to draft.\n\n## Step 1: Mark PR as \"Ready for Review\"\n\n1. When your PR is ready, click **Ready for Review**.\n2. 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.\n\n:warning: Only mark as ready once all merge-conflicts are resolved and the CI is passing.\nFinal Review might get declined if these requirements are not fulfilled.\n\n## Step 2: Final Review (`megatron/core` only)\n\nFor PRs that change `megatron/core`, once all expert reviewers have approved, the `Final Review` label is applied **automatically** and final reviewers are assigned.\n\nFor PRs outside `megatron/core`, this step is skipped.\n\n## Step 3: Approved\n\nOnce all required reviewers have approved, the `Approved` label is applied **automatically**. The PR is now ready to merge.\n\n## Step 4: Merge\n\nAny member of [mcore-engineers](https://github.com/orgs/NVIDIA/teams/mcore-engineers) will be able to merge your PR.\n"
},
{
"path": "docs/discussions/README.md",
"content": "---\norphan: true\n---\n\n\n\n# Megatron Discussions\n\nThis directory contains in-depth guides, tutorials, and discussions about optimizing and using Megatron for various use cases.\n\n## Available Guides\n\n### Training Guides\n\n- **[Megatron-FSDP User Guide](megatron-fsdp-user-guide/megatron-fsdp-user-guide.md)**\n\n A practical guide to enable Megatron-FSDP training, including a quick-start example for DeepSeek-V3, required and recommended configurations, and instructions for checkpoint conversion from torch_dist to fsdp_dtensor.\n\n## Contributing\n\nIf you'd like to contribute a guide or tutorial, please follow this structure:\n\n1. Create a new directory: `docs/discussions/your-guide-name/`\n2. Add your main guide: `docs/discussions/your-guide-name/your-guide-name.md`\n3. Create an images directory: `docs/discussions/your-guide-name/images/`\n4. Update this README.md with a link to your guide\n\nEach guide should be self-contained with its own images and supporting files."
},
{
"path": "docs/discussions/megatron-fsdp-user-guide/example-scripts/sbatch_checkpoint_convert.sh",
"content": "#!/bin/bash\n\n# Configuration: Set these paths before running the script\nMEGATRON_PATH=${MEGATRON_PATH:-\"your_own_megatron_path\"} # Path to Megatron-LM repository\nCONTAINER_IMAGE=${CONTAINER_IMAGE:-\"your_own_container_image\"} # Path to .sqsh or docker image url\nOUTPUT_PATH=${OUTPUT_PATH:-\"your_own_output_path\"} # Path for SLURM logs\n\n# Checkpoint conversion command\n# Note: Update the checkpoint paths in the command below\nRUN_CMD=\"\ncd ${MEGATRON_PATH};\ngit rev-parse HEAD;\nexport PYTHONPATH=${MEGATRON_PATH}:${PYTHONPATH};\npython3 tools/checkpoint/checkpoint_inspector.py \\\n convert-torch-dist-to-fsdp-dtensor --swiglu \\\n your_own_path_to_input_torch_dist_checkpoint \\\n your_own_path_to_output_fsdp_dtensor_checkpoint \\\n --param-to-param-group-map-json your_own_path_to_param_to_param_group_map.json\"\n\n# SLURM settings\nSLURM_LOGS=\"${OUTPUT_PATH}/slurm_logs\"\nmkdir -p ${SLURM_LOGS} || {\n echo \"Error: Failed to create SLURM logs directory ${SLURM_LOGS}\"\n exit 1\n}\n\n# Submit SLURM job\n# Note: Update SBATCH parameters below according to your cluster configuration\nset +e\nsbatch <&1 | tee ${SLURM_LOGS}/\\${SLURM_JOB_ID}.log\n\nEOF\nset -e\n"
},
{
"path": "docs/discussions/megatron-fsdp-user-guide/example-scripts/sbatch_mfsdp_deepseek_v3.sh",
"content": "#!/bin/bash\n\nexport NCCL_IB_SL=1\nexport NCCL_IB_TIMEOUT=19\nexport NVTE_FWD_LAYERNORM_SM_MARGIN=16\nexport NVTE_BWD_LAYERNORM_SM_MARGIN=16\nexport NCCL_P2P_NET_CHUNKSIZE=2097152\nexport TORCH_NCCL_AVOID_RECORD_STREAMS=1\nexport PYTHONWARNINGS=ignore\nexport TRITON_CACHE_DIR=/tmp/triton_cache_$SLURM_NODEID\n\n# Configuration: Set these variables before running the script\nMEGATRON_PATH=${MEGATRON_PATH:-\"your_own_megatron_path\"} # Path to Megatron-LM repository\nCONTAINER_IMAGE=${CONTAINER_IMAGE:-\"your_own_container_image\"} # Path to .sqsh or docker image url\nOUTPUT_PATH=${OUTPUT_PATH:-\"your_own_output_path\"} # Path for output logs and checkpoints\nDATA_PATH=${DATA_PATH:-\"your_own_data_path\"}\nUSE_MEGATRON_FSDP=${USE_MEGATRON_FSDP:-1}\nSHARDING_STRATEGY=${SHARDING_STRATEGY:-\"optim_grads_params\"}\nPROFILE=${PROFILE:-0}\nWANDB=${WANDB:-1}\n\nTP=${TP:-1}\nEP=${EP:-8}\nMBS=${MBS:-4}\nGBS=${GBS:-2048}\nCOMMENT=${COMMENT:-\"hybridep-selective-recompute\"}\n\nPRETRAIN_ARGS=(\n --distributed-timeout-minutes 60\n --tensor-model-parallel-size ${TP}\n --expert-model-parallel-size ${EP}\n --expert-tensor-parallel-size 1\n --context-parallel-size 1\n --use-distributed-optimizer\n --overlap-grad-reduce\n --overlap-param-gather\n --use-mcore-models\n --sequence-parallel\n --use-flash-attn\n --disable-bias-linear\n --micro-batch-size ${MBS}\n --global-batch-size ${GBS}\n --train-samples 585937500\n --exit-duration-in-mins 220\n --no-check-for-nan-in-loss-and-grad\n --manual-gc\n --manual-gc-interval 10\n --recompute-granularity selective\n --recompute-modules mlp moe mla_up_proj layernorm\n --transformer-impl transformer_engine\n --seq-length 4096\n --data-cache-path ${OUTPUT_PATH}/cache\n --tokenizer-type HuggingFaceTokenizer\n --tokenizer-model deepseek-ai/DeepSeek-V3\n --data-path ${DATA_PATH}\n --split 99,1,0\n --no-mmap-bin-files\n --no-create-attention-mask-in-dataloader\n --num-workers 6\n --num-layers 61\n --hidden-size 7168\n --ffn-hidden-size 18432\n --num-attention-heads 128\n --kv-channels 128\n --max-position-embeddings 4096\n --position-embedding-type rope\n --rotary-base 10000\n --make-vocab-size-divisible-by 3232\n --normalization RMSNorm\n --norm-epsilon 1e-6\n --swiglu\n --untie-embeddings-and-output-weights\n --multi-latent-attention\n --attention-dropout 0.0\n --hidden-dropout 0.0\n --clip-grad 1.0\n --weight-decay 0.1\n --qk-layernorm\n --lr-decay-samples 584765624\n --lr-warmup-samples 1536000\n --lr-warmup-init 3.9e-7\n --lr 3.9e-6\n --min-lr 3.9e-7\n --lr-decay-style cosine\n --adam-beta1 0.9\n --adam-beta2 0.95\n --num-experts 256\n --moe-layer-freq [0]*3+[1]*58\n --moe-ffn-hidden-size 2048\n --moe-shared-expert-intermediate-size 2048\n --moe-router-load-balancing-type seq_aux_loss\n --moe-router-topk 8\n --moe-token-dispatcher-type flex\n --moe-flex-dispatcher-backend hybridep\n --moe-router-pre-softmax\n --moe-grouped-gemm\n --moe-aux-loss-coeff 1e-4\n --moe-router-group-topk 4\n --moe-router-num-groups 8\n --moe-router-topk-scaling-factor 2.5\n --moe-router-score-function sigmoid\n --moe-router-enable-expert-bias\n --moe-router-bias-update-rate 1e-3\n --moe-router-dtype fp32\n --moe-permute-fusion\n --moe-router-force-load-balancing\n --q-lora-rank 1536\n --kv-lora-rank 512\n --qk-head-dim 128\n --qk-pos-emb-head-dim 64\n --v-head-dim 128\n --rotary-scaling-factor 40\n --mscale 1.0\n --mscale-all-dim 1.0\n --mtp-num-layers 1\n --mtp-loss-scaling-factor 0.1\n --eval-iters 32\n --eval-interval 100\n --auto-detect-ckpt-format\n --load ${OUTPUT_PATH}/checkpoints\n --save ${OUTPUT_PATH}/checkpoints\n --save-interval 100\n --dist-ckpt-strictness log_all\n --init-method-std 0.02\n --log-timers-to-tensorboard\n --log-memory-to-tensorboard\n --log-num-zeros-in-grad\n --log-params-norm\n --log-validation-ppl-to-tensorboard\n --log-throughput\n --log-interval 1\n --logging-level 40\n --tensorboard-dir ${OUTPUT_PATH}/tensorboard\n --bf16\n --enable-experimental\n) \n\nif [ \"${USE_MEGATRON_FSDP}\" = 1 ]; then\n unset CUDA_DEVICE_MAX_CONNECTIONS\n PRETRAIN_ARGS=(\n \"${PRETRAIN_ARGS[@]}\"\n --use-megatron-fsdp\n --data-parallel-sharding-strategy ${SHARDING_STRATEGY}\n --no-gradient-accumulation-fusion\n --use-distributed-optimizer\n --calculate-per-token-loss\n --init-model-with-meta-device\n --ckpt-format fsdp_dtensor\n --grad-reduce-in-bf16\n --fsdp-double-buffer\n --use-nccl-ub\n )\nfi\n\n# Profiling command\nif [ \"${PROFILE}\" = 1 ]; then\n PROFILE_CMD=\"nsys profile --sample=none --cpuctxsw=none --trace=cuda,nvtx,cublas,cudnn \\\n --capture-range=cudaProfilerApi \\\n --capture-range-end=stop \\\n --cuda-graph-trace=node \\\n --cuda-memory-usage=true \\\n -f true -x true \\\n -o ${OUTPUT_PATH}/nsys/Megatron-FSDP-Deepseek-V3-TP${TP}EP${EP}-MBS${MBS}GBS${GBS}-${COMMENT}\"\n PRETRAIN_ARGS=(\n \"${PRETRAIN_ARGS[@]}\"\n --profile\n --profile-step-start 10\n --profile-step-end 12\n --profile-ranks 0\n )\n echo \"PROFILE_CMD=\"\n echo $PROFILE_CMD\nelse\n PROFILE_CMD=\"\"\nfi\n\nif [ \"${WANDB}\" = 1 ]; then\n export WANDB_API_KEY=${WANDB_API_KEY:-\"your_own_wandb_api_key\"}\n PRETRAIN_ARGS=(\n \"${PRETRAIN_ARGS[@]}\"\n --wandb-project your_own_wandb_project\n --wandb-exp-name DeepSeek-V3-TP${TP}EP${EP}-MBS${MBS}GBS${GBS}-${COMMENT}\n )\nfi\n\nTRAINING_CMD=\"\ncd ${MEGATRON_PATH};\ngit rev-parse HEAD;\nexport PYTHONPATH=${MEGATRON_PATH}:${PYTHONPATH};\n${PROFILE_CMD} python ${MEGATRON_PATH}/pretrain_gpt.py ${PRETRAIN_ARGS[@]}\"\n\n# SLURM settings\nSLURM_LOGS=\"${OUTPUT_PATH}/slurm_logs\"\nmkdir -p ${SLURM_LOGS} || {\n echo \"Error: Failed to create SLURM logs directory ${SLURM_LOGS}\"\n exit 1\n}\n\n# Submit SLURM job\n# Note: Update SBATCH parameters below according to your cluster configuration\nset +e\nsbatch <&1 | tee ${SLURM_LOGS}/\\${SLURM_JOB_ID}.log\n\nEOF\nset -e\n"
},
{
"path": "docs/discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md",
"content": "---\norphan: true\n---\n\n\n\n# Megatron-FSDP User Guide\n\n## Table of Contents\n\n- [Megatron-FSDP Quick Start](#megatron-fsdp-quick-start)\n- [Checkpoint Conversion from 3D-Parallel to Megatron-FSDP](#checkpoint-conversion-from-3d-parallel-to-megatron-fsdp)\n\n## Megatron-FSDP Quick Start\n\nWe recommend using the latest [NVIDIA NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags), which provides a tested software stack and optimized performance.\n\nFor your reference, we provide an example launch script for DeepSeek-V3: [`sbatch_mfsdp_deepseek_v3.sh`](./example-scripts/sbatch_mfsdp_deepseek_v3.sh).\n\n### Required Configurations\n\nTo enable Megatron-FSDP, add the following required flags to your training script:\n\n```bash\n--use-megatron-fsdp\n--data-parallel-sharding-strategy optim_grads_params\n--no-gradient-accumulation-fusion\n--use-distributed-optimizer\n--ckpt-format fsdp_dtensor\n```\n\n### Recommended Configurations\n\nWe also recommend adding the following configurations to further improve performance:\n\n```bash\nunset CUDA_DEVICE_MAX_CONNECTIONS\n```\n```bash\n--calculate-per-token-loss\n--init-model-with-meta-device\n--grad-reduce-in-bf16\n--fsdp-double-buffer\n--use-nccl-ub\n```\n\n💡 **Detailed explanations of these configurations are provided below.**\n\n#### 1. Disable `CUDA_DEVICE_MAX_CONNECTIONS`\n\nTo ensure full parallelization of FSDP communication and computation, disable the CUDA_DEVICE_MAX_CONNECTIONS environment variable. This step avoids potential bubbles in the CUDA stream. (But it may slow down TP and CP to some extent.)\n\n#### 2. Add `--calculate-per-token-loss`\n\nFor gradients sharding mode optimization, include the `--calculate-per-token-loss` flag in your training script. This improves performance by reducing the frequency of gradient scaling, which is also a sizable drain on SM resources.\n\n#### 3. Add `--init-model-with-meta-device`\n\nAllows model initialization using meta device, followed by layer-by-layer initialization of distributed model weight buffers via the `Module.reset_parameters` API, facilitating the initialization of extremely large models.\n\n#### 4. Add `--grad-reduce-in-bf16`\n\nEnables gradient reduction in BF16 precision instead of FP32, reducing communication volume and accelerating the backward pass.\n\n#### 5. Add `--fsdp-double-buffer`\n\nUses persistently allocated double buffers for temporarily-defined memory needed in `MegatronFSDP` communications. While having persistent double buffers may increase peak VRAM utilization, it is necessary to register NCCL user buffers (`nccl_ub=True`) for `MegatronFSDP`. Currently, this is supported only for simple repetitive model structures such as GPT.\n\n- **Only effective when using Megatron-LM.**\n- Defaults to `False`. Automatically overridden to `True` when `nccl_ub` is enabled.\n\n#### 6. Add `--use-nccl-ub`\n\nAllocates and [registers NCCL user buffers](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html#) for param and grad buffers. This option enables an SM-efficient NCCL algorithm that could improve the performance of overlapped computations. This flag will be much more effective when used together with [SHARP](https://docs.nvidia.com/networking/display/sharpv3130) if the FSDP communication includes both NVL and IB domains. Enabling this option will cause additional memory overhead due to the requirement to enable the `fsdp_double_buffer` option.\n\n- **Only effective when using Megatron-LM.**\n- Defaults to `False`.\n- By default we try to use NCCL window (symmetric) registration if it is available. If not it falls back to conventional local registration.\n- **Incompatible with PyTorch's segmentable allocator:** Do not set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` when using `--use-nccl-ub`, as this will cause a runtime error due to compatibility issues with the `torch.cuda.MemPool` API.\n\n## Checkpoint Conversion from 3D-Parallel to Megatron-FSDP\n\nMegatron-FSDP introduces `fsdp_dtensor`, a DTensor-based distributed checkpoint format that serves as its standard. To help you smoothly transition from 3D-Parallel to Megatron-FSDP, we provide a script for converting checkpoints from the `torch_dist` format to the `fsdp_dtensor` format. Using DeepSeek-V3 as an example, the detailed conversion process is described below.\n\n### Step 1: Generate 3D-Parallel Checkpoint with `param_to_param_group_map`\n\nRun your 3D-parallel + EP training script to generate a `torch_dist` checkpoint along with a directory containing `param_to_param_group_map` files. Add the following flag to your training script:\n\n```bash\n--dump-param-to-param-group-map /path/to/param_to_param_group_map\n```\n\nIf you already have a `torch_dist` checkpoint, simply specify the `--dump-param-to-param-group-map /path/to/param_to_param_group_map` flag and run a very short experiment-this will create the `param_to_param_group_map` you need without full pretraining.\n\n### Step 2: Export `param_to_param_group_map` to a JSON File\n\nConvert the `param_to_param_group_map` into a JSON file for easier processing by running:\n\n```bash\npython tools/checkpoint/checkpoint_inspector.py print-torch-dcp-in-json /path/to/param_to_param_group_map\n```\n\nThis will create a `param_to_param_group_map.json` file in the `/path/to/param_to_param_group_map` directory.\n\n### Step 3: Convert Checkpoint from `torch_dist` to `fsdp_dtensor`\n\nConvert your `torch_dist` checkpoint to the `fsdp_dtensor` format using the parameter to `param_to_param_group_map` JSON file:\n\n```bash\ntorchrun --nproc_per_node=8 --nnodes=1 \\\n tools/checkpoint/checkpoint_inspector.py \\\n convert-torch-dist-to-fsdp-dtensor --swiglu \\\n /path/to/input_torch_dist_checkpoint \\\n /path/to/output_fsdp_dtensor_checkpoint \\\n --param-to-param-group-map-json /path/to/param_to_param_group_map.json\n```\n\n**Note:** For multi-node conversion tasks, please refer to the example script: [`sbatch_checkpoint_convert.sh`](./example-scripts/sbatch_checkpoint_convert.sh).\n\n### Step 4: Launch Megatron-FSDP Training\n\nStart your Megatron-FSDP training job using the converted `fsdp_dtensor` checkpoint."
},
{
"path": "docs/documentation.md",
"content": "---\norphan: true\n---\n\n\n\n# Documentation Development\n\n- [Documentation Development](#documentation-development)\n - [Build the Documentation](#build-the-documentation)\n - [Live Building](#live-building)\n - [Documentation Version](#documentation-version)\n\n\n## Build the Documentation\n\nThe following sections describe how to set up and build the NeMo RL documentation.\n\nSwitch to the documentation source folder and generate HTML output.\n\n```sh\ncd docs/\nuv run --group docs sphinx-build . _build/html\n```\n\n* The resulting HTML files are generated in a `_build/html` folder that is created under the project `docs/` folder.\n* The generated python API docs are placed in `apidocs` under the `docs/` folder.\n\n## Checking for Broken Links\n\nTo check for broken http links in the docs, run this command:\n\n```sh\ncd docs/\nuv run --group docs sphinx-build --builder linkcheck . _build/linkcheck\n```\n\nIt will output a JSON file at `_build/linkcheck/output.json` with links it found while building the\ndocs. Records will have a status of `broken` if the link is not reachable. The `docs/conf.py` file is\nconfigured to ignore github links because the CI test will often experience rate limit errors.\nComment out the `linkcheck_ignore` variable there to check all the links.\n\n## Live Building\n\nWhen writing documentation, it can be helpful to serve the documentation and have it update live while you edit.\n\nTo do so, run:\n\n```sh\ncd docs/\nuv run --group docs sphinx-autobuild . _build/html --port 12345 --host 0.0.0.0\n```\n\nOpen a web browser and go to `http://${HOST_WHERE_SPHINX_COMMAND_RUN}:12345` to view the output.\n\n## Documentation Version\n\nThe three files below control the version switcher. Before you attempt to publish a new version of the documentation, update these files to match the latest version numbers.\n\n* docs/versions1.json\n* docs/project.json\n* docs/conf.py\n"
},
{
"path": "docs/get-started/install.md",
"content": "\r\n\r\n# Installation\r\n\r\n## System Requirements\r\n\r\n### Hardware\r\n\r\n- **Recommended**: NVIDIA Turing architecture or later\r\n- **FP8 Support**: Requires NVIDIA Hopper, Ada, or Blackwell GPUs\r\n\r\n### Software\r\n\r\n- **Python**: >= 3.10 (3.12 recommended)\r\n- **PyTorch**: >= 2.6.0\r\n- **CUDA Toolkit**: Latest stable version\r\n\r\n\r\n## Prerequisites\r\n\r\nInstall [uv](https://docs.astral.sh/uv/), a fast Python package installer:\r\n\r\n```bash\r\ncurl -LsSf https://astral.sh/uv/install.sh | sh\r\n```\r\n\r\n\r\n## Option A: Pip Install (Recommended)\r\n\r\nInstall the latest stable release from PyPI:\r\n\r\n```bash\r\nuv pip install megatron-core\r\n```\r\n\r\nTo include optional training dependencies (Weights & Biases, SentencePiece, HF Transformers):\r\n\r\n```bash\r\nuv pip install \"megatron-core[training]\"\r\n```\r\n\r\nFor all extras including [Transformer Engine](https://github.com/NVIDIA/TransformerEngine):\r\n\r\n```bash\r\nuv pip install --group build\r\nuv pip install --no-build-isolation \"megatron-core[training,dev]\"\r\n```\r\n\r\n```{note}\r\n`--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.\r\n```\r\n\r\n```{warning}\r\nBuilding 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`).\r\n```\r\n\r\n```{tip}\r\nFor 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.\r\n```\r\n\r\nTo clone the repository for examples:\r\n\r\n```bash\r\ngit clone https://github.com/NVIDIA/Megatron-LM.git\r\n```\r\n\r\n\r\n## Option B: Install from Source\r\n\r\nFor development or to run the latest unreleased code:\r\n\r\n```bash\r\ngit clone https://github.com/NVIDIA/Megatron-LM.git\r\ncd Megatron-LM\r\nuv pip install -e .\r\n```\r\n\r\nTo install with all development dependencies (includes Transformer Engine, requires pre-installed build deps):\r\n\r\n```bash\r\nuv pip install --group build\r\nuv pip install --no-build-isolation -e \".[training,dev]\"\r\n```\r\n\r\n```{tip}\r\nIf the build runs out of memory, limit parallel compilation jobs with `MAX_JOBS=4 uv pip install --no-build-isolation -e \".[training,dev]\"`.\r\n```\r\n\r\n\r\n## Option C: NGC Container\r\n\r\nFor 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).\r\n\r\nWe 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.\r\n\r\n```bash\r\ndocker run --gpus all -it --rm \\\r\n -v /path/to/dataset:/workspace/dataset \\\r\n -v /path/to/checkpoints:/workspace/checkpoints \\\r\n -e PIP_CONSTRAINT= \\\r\n nvcr.io/nvidia/pytorch:26.01-py3\r\n```\r\n\r\n```{note}\r\nThe 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.\r\n```\r\n\r\nThen install Megatron Core inside the container (torch is already available in the NGC image):\r\n\r\n```bash\r\npip install uv\r\nuv pip install --no-build-isolation \"megatron-core[training,dev]\"\r\n```\r\n\r\n\r\nYou are now ready to run training. See [Your First Training Run](quickstart.md) for next steps.\r\n"
},
{
"path": "docs/get-started/overview.md",
"content": "\r\n\r\n# Overview\r\n\r\nMegatron-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.\r\n\r\n## Megatron Core\r\n\r\nNVIDIA Megatron Core is a library of essential building blocks for highly efficient large-scale generative AI training. It can be used to train models with unparalleled speed at scale across thousands of GPUs. It provides an extensive set of tools for multimodal and speech AI. It expands Megatron LM capabilities.\r\n\r\nMegatron-Core contains GPU-optimized techniques featuring advanced parallelism strategies, optimizations like FP8 training, and support for the latest LLM, MoE, and multimodal architectures. It abstracts these techniques into composable and modular APIs.\r\n\r\nMegatron-Core is compatible with all NVIDIA Tensor Core GPUs and popular LLM architectures such as GPT, BERT, T5, and RETRO.\r\n\r\n\r\n**Composable library** with GPU-optimized building blocks for custom training frameworks.\r\n\r\n**Best for:**\r\n\r\n- **Framework developers** building on top of modular and optimized components\r\n- **Research teams** needing custom training loops, optimizers, or data pipelines\r\n- **ML engineers** requiring fault-tolerant training pipelines\r\n\r\n**What you get:**\r\n\r\n- Composable transformer building blocks (attention, MLP)\r\n- Advanced parallelism strategies (TP, PP, DP, EP, CP)\r\n- Pipeline schedules and distributed optimizers\r\n- Mixed precision support (FP16, BF16, FP8)\r\n- GPU-optimized kernels and memory management\r\n- High-performance dataloaders and dataset utilities\r\n- Model architectures (LLaMA, Qwen, GPT, Mixtral, Mamba)\r\n\r\n## Megatron-LM\r\n\r\nMegatron-LM is a reference implementation, with a lightweight large-scale LLM training framework. It offers a customizable native PyTorch training loop with fewer abstraction layers. It was designed for scaling transformer models to the multi-billion and trillion-parameter regimes under realistic memory and compute constraints. **It serves as a straightforward entry point for exploring Megatron-Core.**\r\n\r\nIt uses advanced parallelization techniques including model parallelism (tensor and pipeline), to allow models with billions of parameters to fit and train across large GPU clusters. It enables breakthroughs in large-scale NLP tasks. It splits model computations across many GPUs, overcoming single-GPU memory limits for training huge models, like GPT-style transformers. \r\n\r\n\r\n**Reference implementation** that includes Megatron Core plus everything needed to train models.\r\n\r\n**Best for:**\r\n\r\n- **Training state-of-the-art foundation models** at scale with cutting-edge performance on latest NVIDIA hardware\r\n- **Research teams** exploring new architectures and training techniques\r\n- **Learning distributed training** concepts and best practices\r\n- **Quick experimentation** with proven model configurations\r\n\r\n**What you get:**\r\n\r\n- Pre-configured training scripts for GPT, LLaMA, DeepSeek, Qwen, and more.\r\n- End-to-end examples from data prep to evaluation\r\n- Research-focused tools and utilities\r\n\r\n\r\n\r\n## Megatron Bridge\r\n\r\nMegatron Bridge provides out-of-the-box bridges and training recipes for models built on top of base model architectures from Megatron Core. \r\n\r\nMegatron Bridge provides a robust, parallelism-aware pathway to convert models and checkpoints. This bidirectional converter performs on-the-fly, model-parallel-aware, per-parameter conversion, and full in-memory loading.\r\n\r\nAfter training or modifying a Megatron model, you can convert it again for deployment or sharing. \r\n\r\n[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)\r\n\r\n\r\n\r\n## Ecosystem Libraries\r\n\r\n**Libraries used by Megatron Core:**\r\n\r\n- **[Megatron Energon](https://github.com/NVIDIA/Megatron-Energon)** - Multi-modal data loader (text, images, video, audio) with distributed loading and dataset blending\r\n- **[Transformer Engine](https://github.com/NVIDIA/TransformerEngine)** - Optimized kernels and FP8 mixed precision support\r\n- **[Resiliency Extension (NVRx)](https://github.com/NVIDIA/nvidia-resiliency-ext)** - Fault tolerant training with failure detection and recovery\r\n\r\n**Libraries using Megatron Core:**\r\n\r\n- **[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\r\n- **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods\r\n- **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples\r\n- **[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).\r\n\r\n**Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed)\r\n\r\n"
},
{
"path": "docs/get-started/quickstart.md",
"content": "\n\n# Your First Training Run\n\nThis guide walks you through running your first training jobs with Megatron Core. Make sure you have completed [installation](install.md) before proceeding.\n\n## Simple Training Example\n\nRun a minimal distributed training loop with mock data on 2 GPUs:\n\n```bash\ntorchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py\n```\n\n## LLaMA-3 Training Example\n\nTrain a LLaMA-3 8B model with FP8 precision on 8 GPUs using mock data:\n\n```bash\n./examples/llama/train_llama3_8b_h100_fp8.sh\n```\n\n## Data Preparation\n\nTo train on your own data, Megatron expects preprocessed binary files (`.bin` and `.idx`).\n\n### 1. Prepare a JSONL File\n\nEach line should contain a `text` field:\n\n```json\n{\"text\": \"Your training text here...\"}\n{\"text\": \"Another training sample...\"}\n```\n\n### 2. Preprocess the Data\n\n```bash\npython tools/preprocess_data.py \\\n --input data.jsonl \\\n --output-prefix processed_data \\\n --tokenizer-type HuggingFaceTokenizer \\\n --tokenizer-model /path/to/tokenizer.model \\\n --workers 8 \\\n --append-eod\n```\n\n### Key Arguments\n\n- `--input`: Path to input JSON/JSONL file\n- `--output-prefix`: Prefix for output binary files (.bin and .idx)\n- `--tokenizer-type`: Tokenizer type (`HuggingFaceTokenizer`, `GPT2BPETokenizer`, etc.)\n- `--tokenizer-model`: Path to tokenizer model file\n- `--workers`: Number of parallel workers for processing\n- `--append-eod`: Add end-of-document token\n\n## Next Steps\n\n- Explore [Parallelism Strategies](../user-guide/parallelism-guide.md) to scale your training\n- Learn about [Data Preparation](../user-guide/data-preparation.md) best practices\n- Check out [Advanced Features](../user-guide/features/index.md) for advanced capabilities\n"
},
{
"path": "docs/get-started/releasenotes.md",
"content": "\r\n\r\n# Release Notes\r\n\r\n\r\n## Roadmaps\r\n\r\nStay up-to-date with our development roadmaps and planned features:\r\n\r\n- **[MoE Q3-Q4 2025 Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - Comprehensive MoE feature development including DeepSeek-V3, Qwen3, advanced parallelism, FP8 optimizations, and Blackwell enhancements\r\n- **[GPT-OSS Implementation Tracker](https://github.com/NVIDIA/Megatron-LM/issues/1739)** - Advanced features including YaRN RoPE scaling, attention sinks, and custom activation functions\r\n\r\n"
},
{
"path": "docs/index.md",
"content": "\n\n# Megatron Core User Guide\n\n**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.\n\nMegatron Core offers a flexible, reusable foundation for building large-scale transformer training systems. **Megatron-LM** serves as a reference implementation demonstrating how to use Megatron Core components to train models with billions to trillions of parameters across distributed GPU clusters.\n\n## Key Features\n\n* Composable transformer building blocks (attention, MLP)\n* Advanced parallelism strategies (TP, PP, DP, EP, CP)\n* Pipeline schedules and distributed optimizers\n* Mixed precision support (FP16, BF16, FP8)\n* GPU-optimized kernels and memory management\n* High-performance dataloaders and dataset utilities\n* Model architectures (LLaMA, Qwen, DeepSeek, GPT, Mamba)\n\n\n```{toctree}\n:maxdepth: 2\n:hidden:\n:caption: About Megatron Core\n\nget-started/overview\nget-started/releasenotes\n```\n\n```{toctree}\n:maxdepth: 2\n:hidden:\n:caption: Get Started\n\nget-started/install\nget-started/quickstart\n```\n\n```{toctree}\n:maxdepth: 2\n:hidden:\n:caption: Basic Usage\n\nuser-guide/data-preparation\nuser-guide/training-examples\nuser-guide/parallelism-guide\n```\n\n```{toctree}\n:maxdepth: 2\n:hidden:\n:caption: Supported Models\n\nmodels/index\n```\n\n```{toctree}\n:maxdepth: 2\n:hidden:\n:caption: Advanced Features\n\nuser-guide/features/moe\nuser-guide/features/context_parallel\nuser-guide/features/custom_fsdp\nuser-guide/features/dist_optimizer\nuser-guide/features/optimizer_cpu_offload\nuser-guide/features/pipeline_parallel_layout\nuser-guide/features/fine_grained_activation_offloading\nuser-guide/features/megatron_energon\nuser-guide/features/megatron_rl\nuser-guide/features/tokenizers\n```\n\n```{toctree}\n:maxdepth: 1\n:hidden:\n:caption: Developer Guide\n\ndeveloper/contribute\ndeveloper/submit\ndeveloper/oncall\ndeveloper/generate_docs\n```\n\n```{toctree}\n:maxdepth: 2\n:hidden:\n:caption: API Reference\n\napi-guide/index\napidocs/index.rst\n```\n\n```{toctree}\n:maxdepth: 2\n:hidden:\n:caption: Resources\n\nadvanced/index\n```"
},
{
"path": "docs/llama_mistral.md",
"content": "\n\n# Llama, Mistral and other Llama-like model support in Megatron-LM\n\nNOTE: 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).\n\nThe 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 ).\n\nSimilarly, [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.\n\nArchitecturally Llama-2, Llama-3 and Mistral-7b are very similar. As such Megatron can support loading checkpoints from all three for inference and finetuning. Converting the checkpoints and loading them is slightly different for each model and is detailed for each below.\n\n# Contents\n\n- [Llama, Mistral and other Llama-like model support in Megatron-LM](#llama-mistral-and-other-llama-like-model-support-in-megatron-lm)\n- [Contents](#contents)\n- [Llama-2](#llama-2)\n - [Download Meta or Huggingface checkpoints](#download-meta-or-huggingface-checkpoints)\n - [Convert checkpoint format](#convert-checkpoint-format)\n - [Meta format](#meta-format)\n - [Huggingface format](#huggingface-format)\n - [Launch model](#launch-model)\n - [Launch Megatron](#launch-megatron)\n - [Launch Meta](#launch-meta)\n - [Launch Huggingface](#launch-huggingface)\n - [Benchmark results](#benchmark-results)\n - [Big Bench](#big-bench)\n - [Multilingual](#multilingual)\n - [LM Evaluation Harness](#lm-evaluation-harness)\n - [MMLU](#mmlu)\n- [Llama-3.x](#llama-3x)\n - [Download Huggingface checkpoints](#download-huggingface-checkpoints)\n - [Convert checkpoint format](#convert-checkpoint-format)\n - [Huggingface format](#huggingface-format)\n - [(Optional) Validate checkpoints](#optional-validate-checkpoints)\n - [Launch model](#launch-model)\n- [Mistral-7b](#mistral-7b)\n - [Download Huggingface checkpoints](#download-huggingface-checkpoints)\n - [Convert checkpoint format](#convert-checkpoint-format)\n - [(Optional) Validate checkpoints](#optional-validate-checkpoints)\n - [Launch model](#launch-model)\n- [Other Llama-like model support](#other-llama-like-model-support)\n- [Known numerical differences](#known-numerical-differences)\n- [Using legacy model format](#using-legacy-model-format)\n\n# Llama-2\n\nLlama-2 checkpoints can be loaded into Megatron for inference and for finetuning. Loading these checkpoints consists of three steps:\n\n1. Get access to download the checkpoints.\n2. Convert the checkpoints from Meta/Huggingface format to Megatron format.\n3. Setup arguments for launching the model.\n\nThe following sections detail these steps. The final section lists benchmark result comparisons between: 1) Llama-2 inference code running the Meta-format checkpoints, and 2) Megatron inference code running the converted checkpoints.\n\n## Download Meta or Huggingface checkpoints\n\nUsers 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.\n\n## Convert checkpoint format\n\nWe recommend passing `--dtype bf16` for training or finetuning. Inference can be done in bfloat16 or float16.\n\n### Meta format\n\nThe Meta format checkpoints are converted to HF format as an intermediate step before converting to Megatron format. The `transformers` package is required, and must have version >=4.31.0 (e.g., `pip install transformers>=4.31.0`). (**Note**: we have specifically tested with versions `4.31.0` and `4.32.0`; your experience may vary with newer versions.) Assuming the downloaded checkpoints are in `$CHECKPOINT_DIR` (with separate sub-directories for 7B, 13B, 70B, etc.), the following example command can be used to convert from Llama-2 format to HF format in bfloat16:\n\n```\npython tools/checkpoint/convert.py \\\n> --model-type GPT \\\n> --loader llama_mistral \\\n> --load-dir ${META_FORMAT_DIR} \\\n> --model-size ${MODEL_SIZE} \\\n> --checkpoint-type meta \\\n> --tokenizer-model ${TOKENIZER_MODEL} \\\n> --saver core \\\n> --save-dir ${MEGATRON_FORMAT_DIR} \\\n> --target-tensor-parallel-size ${TP} \\\n> --target-pipeline-parallel-size ${PP} \\\n> --bf16\n```\n\nValid values for `--model-size` are `llama2-7B`, `llama2-13B`, and `llama2-70B` (for pretrained-only models), and `llama2-7Bf`, `llama2-13Bf`, and `llama2-70Bf` (for chat-finetuned models).\n\n### Huggingface format\n\nThe HF checkpoints can be converted to Megatron format by using Megatron's own Llama-2 checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`). One important argument that must be set correctly is the tensor parallel size (`TP`) for each model. The following table shows these values:\n\n| Model size | Tensor parallel size (`TP`) |\n| ---------- | --------------------------- |\n| 7B | 1 |\n| 13B | 2 |\n| 70B | 8 |\n\nUsing these values for `TP`, along with the path to the Llama-2 tokenizer model (automatically downloaded with original checkpoint download; see `${TOKENIZER_MODEL}` below), run the following command from the root of your Megatron source code to convert from HF format to Megatron format:\n\n```\npython tools/checkpoint/convert.py \\\n> --model-type GPT \\\n> --loader llama_mistral \\\n> --load-dir ${HF_FORMAT_DIR} \\\n> --model-size ${MODEL_SIZE} \\\n> --checkpoint-type hf \\\n> --tokenizer-model ${TOKENIZER_MODEL} \\\n> --saver core \\\n> --save-dir ${MEGATRON_FORMAT_DIR} \\\n> --target-tensor-parallel-size ${TP} \\\n> --target-pipeline-parallel-size ${PP} \\\n> --bf16\n```\n\nAfter this conversion, we are ready to load the checkpoints into a Megatron GPT model.\n\n## Launch model\n\n### Launch Megatron\n\nIf loading for either inference or finetuning, use the following arguments:\n\n```\n--tensor-model-parallel-size ${TP} \\\n--pipeline-model-parallel-size 1 \\\n--seq-length 4096 \\\n--max-position-embeddings 4096 \\\n--tokenizer-type Llama2Tokenizer \\\n--tokenizer-model ${TOKENIZER_MODEL} \\\n--load ${CHECKPOINT_DIR} \\\n--exit-on-missing-checkpoint \\\n--use-checkpoint-args \\\n--no-load-optim \\\n--no-load-rng \\\n--untie-embeddings-and-output-weights \\\n--use-rotary-position-embeddings \\\n--normalization RMSNorm \\\n--no-position-embedding \\\n--no-masked-softmax-fusion \\\n--attention-softmax-in-fp32\n```\n\n**Note:** If you converted to the legacy model format (i.e., `--saver legacy`), please see [here](#using-legacy-model-format).\n\n### Launch Meta\n\nMeta checkpoints can be launched with: \n\n### Launch Huggingface\n\nHuggingface checkpoints can be launched with: \n\n## Benchmark results\n\nThe tables below list the benchmark comparisons between native Llama-2 (using Meta's checkpoint and Meta's inference code) and Megatron (using a converted HF checkpoint and Megatron's inference code).\n\nThe values are the percent error between Megatron and Llama-2, calculated using the formula: `| - | / `, where the type of score is detailed before each table. Across all tests (80 total per model size), the mean error is 0.15%. The small difference in benchmark scores between the two models is due to minor arithmetic differences in implementation that alter the numerics slightly. Some of the factors that influence this difference include:\n\n- Megatron performs batch matrix multiplications in a couple places, such as within self attention and in SwiGLU, that Llama performs separately.\n- Megatron uses `torch.baddbmm` within self attention, versus Llama using `torch.matmul`.\n- Megatron uses a `sin`/`cos` implementation for rotary position embeddings, versus Llama using a `polar`/`complex` implementation.\n- Llama calls `torch.set_default_dtype(torch.float16)` during initialization, which Megatron does not.\n\n### Big Bench\n\nScore type: multiple choice grade.\n\n| bigbench / standard | 7b | 13b | 70b |\n| -- | -- | -- | -- |\n| date_understanding | 0.29% | 0.13% | 0.12% |\n| general_knowledge | 0.00% | 0.00% | 0.00% |\n| human_organs_senses | 0.00% | 0.00% | 0.00% |\n| intent_recognition | 0.00% | 0.11% | 0.00% |\n| riddle_sense | 0.00% | 0.00% | 0.00% |\n| similarities_abstraction | 0.00% | 0.58% | 0.00% |\n| simple_arithmetic_json_multiple_choice | 0.00% | 0.00% | 0.00% |\n| undo_permutation | 0.19% | 0.19% | 0.18% |\n\n### Multilingual\n\nScore type: multiple choice grade.\n\n| multilingual / xcopa | 7b | 13b | 70b |\n| -- | -- | -- | -- |\n| en-template-mGPT-remove-punctuation | 0.08% | 0.00% | 0.00% |\n| et-template-mGPT-remove-punctuation | 0.00% | 0.13% | 0.25% |\n| ht-template-mGPT-remove-punctuation | 0.26% | 0.13% | 0.26% |\n| id-template-mGPT-remove-punctuation | 0.11% | 0.00% | 0.19% |\n| it-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% |\n| qu-template-mGPT-remove-punctuation | 0.00% | 0.00% | 0.27% |\n| sw-template-mGPT-remove-punctuation | 0.14% | 0.13% | 0.13% |\n| th-template-mGPT-remove-punctuation | 0.25% | 0.13% | 0.13% |\n| tr-template-mGPT-remove-punctuation | 0.26% | 0.00% | 0.34% |\n| vi-template-mGPT-remove-punctuation | 0.00% | 0.11% | 0.00% |\n| zh-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% |\n\n### LM Evaluation Harness\n\nScore type: multiple choice grade.\n\n| lm-eval | 7b | 13b | 70b |\n| -- | -- | -- | -- |\n| boolq | 0.04% | 0.04% | 0.07% |\n| hellaswag | 0.02% | 0.03% | 0.03% |\n| piqa | 0.00% | 0.00% | 0.07% |\n| winogrande | 0.00% | 0.11% | 0.20% |\n\n### MMLU\n\nScore type: multiple choice grade.\n\nNote: the number in brackets is the number of sub-tasks for each supercategory.\n\n| mmlu | 7b | 13b | 70b |\n| -- | -- | -- | -- |\n| stem [18] | 0.79% | 0.05% | 0.01% |\n| humanities [13] | 0.19% | 0.01% | 0.02% |\n| other (business, health, misc.) [14] | 0.08% | 0.06% | 0.12% |\n| social sciences [12] | 0.37% | 0.21% | 0.01% |\n\n# Llama-3.x\n\nLlama-3.x checkpoints can be loaded into Megatron for inference and for finetuning. Loading these checkpoints consists of several steps:\n\n1. Get access to download the checkpoints (weights and tokenizer).\n2. Convert the checkpoints from Huggingface format to Megatron format.\n3. (Optional) Validate converted checkpoints\n4. Setup arguments for launching the model.\n\nThe following sections detail these steps.\n\n## Download Huggingface checkpoints\n\nUsers must first apply for access to download the Llama-3.x checkpoints from [Huggingface](https://huggingface.co/meta-llama).\n\n## Convert checkpoint format\n\nWe recommend passing `--dtype bf16` for training or finetuning. Inference can be done in bfloat16 or float16.\n\n### Huggingface format\n\nThe HF checkpoints can be converted to Megatron format by using Megatron's own Llama-3.x checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`). One important argument that must be set correctly is the tensor parallel size (`TP`) for each model. The following table shows these values:\n\n| Model size | Tensor parallel size (`TP`) |\n| ---------- | --------------------------- |\n| 1B | 1 |\n| 3B | 1 |\n| 8B | 1 |\n| 70B | 8 |\n\nUsing these values for `TP`, along with the path to the Llama-3.x tokenizer model (automatically downloaded with original checkpoint download; see `${TOKENIZER_MODEL}` below), run the following command from the root of your Megatron source code to convert from HF format to Megatron format:\n\n```\n$>: python tools/checkpoint/convert.py \\\n > --bf16 \\\n > --model-type GPT \\\n > --loader llama_mistral \\\n > --saver core \\\n > --target-tensor-parallel-size ${TP} \\\n > --checkpoint-type hf \\\n > --load-dir ${HF_FORMAT_DIR} \\\n > --save-dir ${MEGATRON_FORMAT_DIR} \\\n > --tokenizer-model ${TOKENIZER_MODEL} \\\n > --model-size llama3 \\\n```\n\nAfter this conversion, we are ready to load the checkpoints into a Megatron GPT model.\n\n## (Optional) Validate checkpoints\n\nA Megatron-LM text generation server for Llama3 can be launched using the script `examples/inference/llama_mistral/run_text_generation_llama3.sh `. For Llama3.1, please use `examples/inference/llama_mistral/run_text_generation_llama3.1.sh`.\n\nOnce running, query the server with `curl 'http://:5000/api' -X 'PUT' -H 'Content-Type: application/json; charset=UTF-8' -d '{\"prompts\":[\"\"], \"tokens_to_generate\":100, \"top_k\":1}'`.\n\nA reference generation for comparison can be obtained from the Huggingface transformers library by running `python examples/llama_mistral/huggingface_reference.py --model_path --prompt `.\n\n## Launch model\n\nIf loading for either inference or finetuning, use the following arguments for Llama 3.0:\n\n```\n--tensor-model-parallel-size ${TP} \\\n--pipeline-model-parallel-size 1 \\\n--seq-length 8192 \\\n--max-position-embeddings 8192 \\\n--tokenizer-type HuggingFaceTokenizer \\\n--tokenizer-model ${TOKENIZER_MODEL} \\\n--load ${CHECKPOINT_DIR} \\\n--exit-on-missing-checkpoint \\\n--use-checkpoint-args \\\n--no-load-optim \\\n--no-load-rng \\\n--untie-embeddings-and-output-weights \\\n--normalization RMSNorm \\\n--position-embedding-type rope \\\n--no-masked-softmax-fusion \\\n--attention-softmax-in-fp32 \\\n--disable-bias-linear \\\n--transformer-impl transformer_engine \\\n--group-query-attention 8 \\\n--attention-dropout 0.0 \\\n--hidden-dropout 0.0 \\\n--rotary-base 500000 \\\n--rotary-percent 1.0 \\\n--ffn-hidden-size 14336 \\\n--num-attention-heads 32 \\\n--swiglu \\\n--bf16 \\\n```\n\nFor Llama3.1 please use the following arguments:\n\n```\n--tensor-model-parallel-size ${TP} \\\n--pipeline-model-parallel-size 1 \\\n--seq-length 8192 \\\n--max-position-embeddings 131072 \\\n--tokenizer-type HuggingFaceTokenizer \\\n--tokenizer-model ${TOKENIZER_MODEL} \\\n--load ${CHECKPOINT_DIR} \\\n--exit-on-missing-checkpoint \\\n--use-checkpoint-args \\\n--no-load-optim \\\n--no-load-rng \\\n--untie-embeddings-and-output-weights \\\n--normalization RMSNorm \\\n--position-embedding-type rope \\\n--no-masked-softmax-fusion \\\n--attention-softmax-in-fp32 \\\n--disable-bias-linear \\\n--transformer-impl transformer_engine \\\n--group-query-attention 8 \\\n--attention-dropout 0.0 \\\n--hidden-dropout 0.0 \\\n--rotary-base 500000 \\\n--rotary-percent 1.0 \\\n--use-rope-scaling \\\n--ffn-hidden-size 14336 \\\n--num-attention-heads 32 \\\n--swiglu \\\n--bf16 \\\n```\n\n**Note:** If you converted to the legacy model format (i.e., `--saver legacy`), please see [here](#using-legacy-model-format).\n\n# Mistral-7b\n\nMegatron currently supports loading the v0.3 release of Mistral-7b (which does not use sliding window attention and offers a larger 32768 vocabulary) for inference and finetuning. Loading these checkpoints consists of several steps:\n\n1. Get access to download the checkpoints (weights and tokenizer).\n2. Convert the checkpoints from HuggingFace format to Megatron format.\n3. (Optional) Validate converted checkpoints\n4. Setup arguments for launching the model.\n\nThe following sections detail these steps.\n\n## Download Huggingface checkpoints\n\nUsers 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)).\n\n## Convert checkpoint format\n\nThe HF checkpoints can be converted to Megatron format by using Megatron's own Mistral checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`).\n\nUsing the path to the Mistral tokenizer model (downloaded alongside the HF checkpoint), run the following command from the root of your Megatron source code to convert from HF format to the Megatron core format:\n\n```\n$>: python tools/checkpoint/convert.py \\\n > --bf16 \\\n > --model-type GPT \\\n > --loader llama_mistral \\\n > --saver core \\\n > --target-tensor-parallel-size ${TP} \\\n > --checkpoint-type hf \\\n > --load-dir ${HF_FORMAT_DIR} \\\n > --save-dir ${MEGATRON_FORMAT_DIR} \\\n > --tokenizer-model ${TOKENIZER_MODEL} \\\n > --model-size mistral \\\n```\n\nAfter this conversion, we are ready to load the checkpoints into a Megatron core GPT model.\n\n## (Optional) Validate checkpoints\n\nA Megatron-LM text generation server for Mistral-7B can be launched using the script `examples/inference/llama_mistral/run_text_generation_mistral.sh `.\n\nOnce running, query the server with `curl 'http://:5000/api' -X 'PUT' -H 'Content-Type: application/json; charset=UTF-8' -d '{\"prompts\":[\"\"], \"tokens_to_generate\":100, \"top_k\":1}'`.\n\nA reference generation for comparison can be obtained from the Huggingface transformers library by running `python examples/inference/llama_mistral/huggingface_reference.py --model_path --prompt `.\n\n## Launch model\n\nIf loading for either inference or finetuning, use the following arguments:\n\n```\n--tensor-model-parallel-size ${TP} \\\n--pipeline-model-parallel-size 1 \\\n--seq-length 4096 \\\n--max-position-embeddings 4096 \\\n--tokenizer-type HuggingFaceTokenizer \\\n--tokenizer-model ${TOKENIZER_MODEL} \\\n--load ${CHECKPOINT_DIR} \\\n--exit-on-missing-checkpoint \\\n--use-checkpoint-args \\\n--no-load-optim \\\n--no-load-rng \\\n--untie-embeddings-and-output-weights \\\n--normalization RMSNorm \\\n--position-embedding-type rope \\\n--no-masked-softmax-fusion \\\n--attention-softmax-in-fp32\n--apply-layernorm-1p \\\n--transformer-impl transformer_engine \\\n--group-query-attention 8 \\\n--disable-bia-linear \\\n--rotary-base 1000000 \\\n--rotary-percent 1.0 \\\n--swiglu \\\n--ffn-hidden-size 14336 \\\n--num-attention-heads 32\n```\n\n**Note:** If you converted to the legacy model format (i.e., `--saver legacy`), please see [here](#using-legacy-model-format).\n\n# Other Llama-like model support\n\n*Note: Experimental*\n\nMany models such as Yi-34B and Qwen2.x use the Llama architecture and may be converted from HuggingFace to Megatron using the commands in [Llama-3.x](#llama-3x).\n\n# Known numerical differences\n\nIt 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:\n\n1. TransformerEngine (TE) uses the model params_dtype inside RMSNorm whereas the Huggingface implementation uses fp32. See for details: \n2. 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.\n\n# Using legacy model format\n\nIn all the checkpoint conversion examples used in this document, the saver format `--saver core` is used, signifying that the newer (and recommended) Megatron GPT model class will be used. I.e.:\n\n- old class: `megatron.legacy.model.gpt_model.GPTModel`\n- new class: `megatron.core.models.gpt.gpt_model.GPTModel`\n\nUsing this new format is the recommended approach. However, if your use case requires using the older class (i.e., convert using `--saver legacy`), then when launching training or finetuning, the following args must be added:\n\n- `--use-legacy-models`: use the older model class\n- `--ckpt-format torch`: use the `torch` checkpoint format, which is the only checkpoint format that is compatible with the legacy model format\n"
},
{
"path": "docs/models/index.md",
"content": "\n\n# Supported Models\n\nMegatron Core supports a wide range of language and multimodal models with optimized implementations for large-scale training.\n\n## Model Conversion\n\nFor converting HuggingFace models to Megatron format, use [Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge), the official standalone converter. Megatron Bridge supports an extensive list of models including LLaMA, Mistral, Mixtral, Qwen, DeepSeek, Gemma, Phi, Nemotron, and many more.\n\nSee the [Megatron Bridge supported models list](https://github.com/NVIDIA-NeMo/Megatron-Bridge?tab=readme-ov-file#supported-models) for the complete and up-to-date list of supported models.\n\n```{toctree}\n:maxdepth: 1\n\nllms\nmultimodal\n../llama_mistral\n```\n"
},
{
"path": "docs/models/llms.md",
"content": "\n\n# Language Models\n\nMegatron Core supports the following language model architectures for large-scale training.\n\n## Converting HuggingFace Models\n\nUse [**Megatron Bridge**](https://github.com/NVIDIA-NeMo/Megatron-Bridge) to convert HuggingFace models to Megatron format. Megatron Bridge is the official standalone converter with support for an extensive list of models including LLaMA, Mistral, Mixtral, Qwen, DeepSeek, Gemma, Phi, Nemotron, and many more.\n\nSee the [Megatron Bridge supported models list](https://github.com/NVIDIA-NeMo/Megatron-Bridge?tab=readme-ov-file#supported-models) for the complete and up-to-date list.\n\n## Decoder-Only Models\n\n| Model | Description | Key Features |\n|-------|-------------|--------------|\n| **GPT** | Generative Pre-trained Transformer | Standard autoregressive LM, foundational architecture |\n| **LLaMA** | Meta's LLaMA family | Efficient architecture with RoPE, SwiGLU, RMSNorm |\n| **Mistral** | Mistral AI models | Sliding window attention, efficient inference |\n| **Mixtral** | Sparse Mixture-of-Experts | 8x7B MoE architecture for efficient scaling |\n| **Qwen** | Alibaba's Qwen series | HuggingFace integration, multilingual support |\n| **Mamba** | State Space Model | Subquadratic sequence length scaling, efficient long context |\n\n## Encoder-Only Models\n\n| Model | Description | Key Features |\n|-------|-------------|--------------|\n| **BERT** | Bidirectional Encoder Representations | Masked language modeling, classification tasks |\n\n## Encoder-Decoder Models\n\n| Model | Description | Key Features |\n|-------|-------------|--------------|\n| **T5** | Text-to-Text Transfer Transformer | Unified text-to-text framework, sequence-to-sequence |\n\n## Example Scripts\n\nTraining examples for these models can be found in the `examples/` directory:\n- `examples/gpt3/` - GPT-3 training scripts\n- `examples/llama/` - LLaMA training scripts\n- `examples/mixtral/` - Mixtral MoE training\n- `examples/mamba/` - Mamba training scripts\n- `examples/bert/` - BERT training scripts\n- `examples/t5/` - T5 training scripts\n\n## Model Implementation\n\nAll language models are built using Megatron Core's composable transformer blocks, enabling:\n- Flexible parallelism strategies (TP, PP, DP, EP, CP)\n- Mixed precision training (FP16, BF16, FP8)\n- Distributed checkpointing\n- Efficient memory management\n"
},
{
"path": "docs/models/multimodal.md",
"content": "\n\n# Multimodal Models\n\nMegatron Core supports multimodal models that combine language with vision, audio, and other modalities for comprehensive multimodal understanding.\n\n## MIMO: Multimodal In/Out Framework\n\n**MIMO (Multimodal In/Out Model)** is an experimental framework in Megatron Core that supports arbitrary combinations of modalities including vision, audio, and text. MIMO provides a flexible architecture for building custom multimodal models.\n\n> **Note**: MIMO is experimental and under active development. The API may change in future releases.\n\n**Key Features:**\n- Arbitrary modality combinations (vision, audio, text, etc.)\n- Flexible encoder architecture for different input modalities\n- Unified embedding space across modalities\n- Support for both vision-language and audio-vision-language models\n\nSee [examples/mimo](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/mimo) for training scripts and examples.\n\n## Vision-Language Models\n\n| Model | Description | Vision Encoder | Language Model |\n|-------|-------------|----------------|----------------|\n| **LLaVA** | Visual instruction tuning | CLIP ViT-L/14 | Mistral-7B / LLaMA |\n| **NVLM** | NVIDIA Vision-Language Model | CLIP / Custom ViT | LLaMA-based |\n| **LLaMA 3.1 Nemotron Nano VL** | Efficient multimodal model | Vision Transformer | LLaMA 3.1 8B |\n\n## Vision Encoders\n\n| Model | Description | Key Features |\n|-------|-------------|--------------|\n| **CLIP ViT** | OpenAI's CLIP Vision Transformer | Image-text alignment, multiple scales (L/14@336px) |\n| **RADIO** | Resolution-Agnostic Dynamic Image Optimization | Flexible resolution handling, efficient vision encoding |\n\n## Diffusion Models\n\nFor multimodal diffusion models (image generation, text-to-image, etc.), see [NeMo Diffusion Models](https://github.com/NVIDIA-NeMo/NeMo/tree/main/nemo/collections/diffusion). NeMo provides production-ready implementations of:\n- Stable Diffusion variants\n- Text-to-image generation\n- Image-to-image translation\n- ControlNet and other conditioning mechanisms\n\n## Multimodal Features\n\n- **Image-Text Alignment**: Pre-training on image-caption pairs\n- **Visual Instruction Tuning**: Fine-tuning on instruction-following datasets\n- **Flexible Vision Encoders**: Support for different ViT architectures and resolutions\n- **Combined Checkpointing**: Unified checkpoints combining vision and language models\n- **Efficient Training**: Full parallelism support (TP, PP, DP) for both vision and language components\n\n## Example Scripts\n\nMultimodal training examples can be found in the following directories:\n\n**MIMO Framework:**\n- `examples/mimo/` - Multimodal In/Out training with support for vision-language and audio-vision-language models\n\n**Specific Multimodal Models:**\n- `examples/multimodal/` - LLaVA-style training with Mistral + CLIP\n- `examples/multimodal/nvlm/` - NVLM training scripts\n- `examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/` - Nemotron VL training\n- `examples/multimodal/radio/` - RADIO vision encoder integration\n"
},
{
"path": "docs/project.json",
"content": "{\"name\": \"megatron-lm\", \"version\": \"nightly\"}\n\n"
},
{
"path": "docs/user-guide/data-preparation.md",
"content": "\n\n# Data Preparation\n\nPreparing your data correctly is essential for successful training with Megatron Core.\n\n## Data Format\n\nMegatron Core expects training data in JSONL (JSON Lines) format, where each line is a JSON object:\n\n```json\n{\"text\": \"Your training text here...\"}\n{\"text\": \"Another training sample...\"}\n{\"text\": \"More training data...\"}\n```\n\n## Preprocessing Data\n\nUse the `preprocess_data.py` tool to convert your JSONL data into Megatron's binary format:\n\n```bash\npython tools/preprocess_data.py \\\n --input data.jsonl \\\n --output-prefix processed_data \\\n --tokenizer-type HuggingFaceTokenizer \\\n --tokenizer-model /path/to/tokenizer.model \\\n --workers 8 \\\n --append-eod\n```\n\n### Key Arguments\n\n| Argument | Description |\n|----------|-------------|\n| `--input` | Path to input JSON/JSONL file |\n| `--output-prefix` | Prefix for output binary files (.bin and .idx) |\n| `--tokenizer-type` | Tokenizer type (`HuggingFaceTokenizer`, `GPT2BPETokenizer`, etc.) |\n| `--tokenizer-model` | Path to tokenizer model file |\n| `--workers` | Number of parallel workers for processing |\n| `--append-eod` | Add end-of-document token |\n\n## Finding Optimal Number of Workers\n\nUse the `--find-optimal-num-workers` flag to find number of workers which gives the best performance in terms of preprocessed documents per second.\nScript 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.\n\n```bash\npython tools/preprocess_data.py \\\n --input data.jsonl \\\n --output-prefix processed_data \\\n --tokenizer-type HuggingFaceTokenizer \\\n --tokenizer-model /path/to/tokenizer.model \\\n --workers 8 \\\n --find-optimal-num-workers \\\n --workers-to-check 4 8 16 32 \\\n --max-documents 50000\n```\n\n**Required arguments**\n\n| Argument | Description |\n|----------|-------------|\n| `--find-optimal-num-workers` | Activates search of optimal number of workers |\n| `--workers-to-check` | List of possible number of workers to run |\n| `--max-documents` | Number of documents to be preprocessed during each run |\n\n**Output example**\n\n```bash\n-----------------------------------\nPerformance results (fastest → slowest):\n1. 16 workers → avg. docs/s: 9606.6476\n2. 32 workers → avg. docs/s: 9275.3284\n3. 8 workers → avg. docs/s: 9151.9280\n4. 4 workers → avg. docs/s: 6391.3819\n\n-----------------------------------\nThe most optimal num of workers is 16 with avg. preprocessed docs/s: 9606.6476.\n-----------------------------------\n```\n\n## Output Files\n\nThe preprocessing tool generates two files:\n- `processed_data.bin` - Binary file containing tokenized sequences\n- `processed_data.idx` - Index file for fast random access\n\n## Using Preprocessed Data\n\nReference your preprocessed data in training scripts:\n\n```bash\n--data-path processed_data \\\n--split 949,50,1 # Train/validation/test split\n```\n\n## Common Tokenizers\n\n### HuggingFace Tokenizers\n\n```bash\n--tokenizer-type HuggingFaceTokenizer \\\n--tokenizer-model /path/to/tokenizer.model\n```\n\n### GPT-2 BPE Tokenizer\n\n```bash\n--tokenizer-type GPT2BPETokenizer \\\n--vocab-file gpt2-vocab.json \\\n--merge-file gpt2-merges.txt\n```\n"
},
{
"path": "docs/user-guide/features/context_parallel.md",
"content": "\n\n# context_parallel package\n\n## Context parallelism overview\n\n```{figure} ../../images/context_parallel/CP_overview.png\n:alt: cp_overview\n:align: center\n\nFigure 1: A transformer layer running with TP2CP2. Communications next to Attention are for CP, others are for TP. (AG/RS: all-gather in forward and reduce-scatter in backward, RS/AG: reduce-scatter in forward and all-gather in backward, /AG: no-op in forward and all-gather in backward).\n```\n\nContext Parallelism (\"CP\") is a parallelization scheme on the dimension of sequence length. Unlike prior SP (sequence parallelism) which only splits the sequence of Dropout and LayerNorm activations, CP partitions the network inputs and all activations along sequence dimension. With CP, all modules except attention (e.g., Linear, LayerNorm, etc.) can work as usual without any changes, because they do not have inter-token operations. As for attention, the Q (query) of each token needs to compute with the KV (key and value) of all tokens in the same sequence. Hence, CP requires additional all-gather across GPUs to collect the full sequence of KV. Correspondingly, reduce-scatter should be applied to the activation gradients of KV in backward propagation. To reduce activation memory footprint, each GPU only stores the KV of a sequence chunk in forward and gathers KV again in backward. KV communication happens between a GPU and its counterparts in other TP groups. The all-gather and reduce-scatter are transformed to point-to-point communications in ring topology under the hood. Exchanging KV also can leverage MQA/GQA to reduce communication volumes, as they only have one or few attention heads for KV.\n\nFor example, in Figure 1, assuming sequence length is 8K, each GPU processes 4K tokens. GPU0 and GPU2 compose a CP group, they exchange KV with each other. Same thing also happens between GPU1 and GPU3. CP is similar to [Ring Attention](https://arxiv.org/abs/2310.01889) but provides better performance by (1) leveraging the latest OSS and cuDNN flash attention kernels; (2) removing unnecessary computation resulted from low-triangle causal masking and achieving optimal load balance among GPUs.\n\n## Context parallelism benefits\n\n```{figure} ../../images/context_parallel/CP_results.png\n:alt: cp_results\n:align: center\n\nFigure 2: Speedup of 175B GPT with various TP+CP combinations vs. full recompute (i.e., TP8CP1).\n```\n\nLLM encounters OOM (out of memory) issue with long context (i.e., long sequence length) because of linearly increasing memory footprint of activations. Recomputing activations in backward can avoid OOM but also introduce significant overheads (~30% with full recompute). Enlarging TP (tensor model parallelism) can fix the OOM issue as well, but it potentially makes compute (e.g., Linear) too short to overlap communication latencies. To be clear, scaling out to more GPUs with bigger TP can hit the overlapping problem no matter if OOM happens.\n\nCP can better address the issues. With CP, each GPU only computes on a part of the sequence, which reduces both computation and communication by CP times. Therefore, there are no concerns about the overlapping between them. The activation memory footprint per GPU is also CP times smaller, hence no OOM issue anymore. As Figure 2 shows, the combinations of TP and CP can achieve optimal performance by eliminating recompute overheads and making the best tradeoff between computation and communications.\n\n## Enabling context parallelism\n\nCP support has been added to GPT. All models that share GPT code path also should be able to benefit from CP, such as Llama. CP can work with TP (tensor model parallelism), PP (pipeline model parallelism), and DP (data parallelism), where the total number of GPUs equals TPxCPxPPxDP. CP also can work with different attention variants, including MHA/MQA/GQA, uni-directional and bi-directional masking.\n\nCP is enabled by simply setting context_parallel_size= in command line. Default context_parallel_size is 1, which means CP is disabled. Running with CP requires Megatron-Core (>=0.5.0) and Transformer Engine (>=1.1).\n\n"
},
{
"path": "docs/user-guide/features/custom_fsdp.md",
"content": "\n\n# Megatron FSDP\n\n**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.**\n\n## How to use ?\n\nAdd these flag to enable MCore custom FSDP.\n\n```bash\n--use-megatron-fsdp\n--data-parallel-sharding-strategy optim_grads_params\n--no-gradient-accumulation-fusion\n--use-distributed-optimizer\n```\n\nFor a practical guide covering required configurations, checkpoint conversion, and example scripts, see the [Megatron-FSDP User Guide](../../discussions/megatron-fsdp-user-guide/megatron-fsdp-user-guide.md).\n\n## Key Features\n\n- **Sharding Strategy**: Efficiently shards optimizer states, gradients, and parameters to reduce memory consumption.\n- **Communication and Computation Overlap**: Optimized to enable concurrent execution of communication and computation, enhancing overall efficiency.\n- **Supports automatic mixed precision training**: Compatible with BF16 O1/O2/O3 recipes, as well as FP8 compute with FP32 parameters and FP8 parameter training, allowing for flexible precision configurations.\n- **Tensor Parallelism (TP), Expert Parallelism (EP) and Context Parallelism (CP)**: Compatible with TP, EP and CP configurations, enabling efficient scaling of large language models.\n- **Distributed Model Initialization with Meta Device**: Allows model initialization using meta device, followed by layer-by-layer initialization of distributed model weight buffers via the `Module.reset_parameters` API, facilitating the initialization of extremely large models.\n\n## Configuration Recommendations\n\n### 1. Disable `CUDA_DEVICE_MAX_CONNECTIONS`\n\nTo ensure full parallelization of FSDP communication and computation, disable the CUDA_DEVICE_MAX_CONNECTIONS environment variable. This step avoids potential bubble in CUDA stream. (But it may slow down TP and CP to some extent.)\n\n```bash\nunset CUDA_DEVICE_MAX_CONNECTIONS\n```\n\n### 2. Add `--calculate-per-token-loss`\n\nFor gradients sharding mode optimization, include the `--calculate-per-token-loss` flag in your training script. This improves performance by reducing the frequency of gradient scaling, which is also a sizable drain on SM resources.\n\n## Design of Custom FSDP\n\n### 1. Overview\n\nThe custom Fully Sharded Data Parallelism (FSDP) implementation in Megatron-Core is specifically designed to optimize memory consumption and performance for large language models. The core design principles include:\n\n - **Optimized for Large Language Models**: This custom FSDP implementation is tailored to efficiently scale with models containing billions of parameters, ensuring seamless execution and training of massive models.\n - **Efficient Memory Consumption**: By strategically sharding optimizer states, gradients, and model parameters, the custom FSDP significantly reduces memory usage. This approach enables the training of models that would otherwise be too large to fit in memory.\n - **Efficient Workflow & Overlapping Communication and Computation**: The implementation is engineered to minimize the number of communication steps required during training. It maximizes the overlap between communication and computation, thereby enhancing overall training efficiency and reducing latency.\n - **Support for MCore's Efficient Training Methods**: The custom FSDP seamlessly integrates with Megatron-Core's advanced parallelism techniques, including tensor parallelism, expert parallelism and context parallelism. Additionally, it supports automatic mixed precision training, further optimizing training performance and efficiency.\n\nThe design of Custom FSDP draws inspiration from PyTorch FSDP [Zhao, Yanli, et al.](https://arxiv.org/pdf/2304.11277) and MCore's distributed optimizer. The introduction to PyTorch FSDP is referenced here to clarify the underlying concepts of the custom FSDP design.\n\n> In DistributedDataParallel, (DDP) training, each process/ worker owns a replica of the model and processes a batch of data, finally it uses all-reduce to sum up gradients over different workers. In DDP the model weights and optimizer states are replicated across all workers. FSDP is a type of data parallelism that shards model parameters, optimizer states and gradients across DDP ranks.\n\n> When training with FSDP, the GPU memory footprint is smaller than when training with DDP across all workers. This makes the training of some very large models feasible by allowing larger models or batch sizes to fit on device. This comes with the cost of increased communication volume. The communication overhead is reduced by internal optimizations like overlapping communication and computation.\n\n\n\n*Notice that the unit processed in workflow here is the “FSDP instance 1: N layers”, where an FSDP instance is the smallest FSDP processing unit (also a PyTorch module), which means that we can safely release this module weights after using it (executing the forward or backward of this module), and there will be no other computations computations relying on these weights. This capability is the foundation of FSDP's layer-by-layer execution and memory-saving strategy. An FSDP instance is also referred to as an **FSDP Unit**.*\n\n*It is worth noting that an FSDP instance can correspond to multiple FSDP parameter groups. These groups are separated by Data Parallel (DP) communication groups and the data type of the parameter or gradient. Consequently, an FSDP instance may require several parameter-gather tasks before execution (forward or backward). Each **FSDP parameter group** corresponds to one **Data Parallel Buffer** in custom FSDP.*\n\nAt a high level FSDP works as follow:\n\nIn constructor\n - Shard model parameters and each rank only keeps its own shard\n\nIn forward path\n - Run all_gather to collect all shards from all ranks to recover the full parameter in this FSDP unit\n - Run forward computation\n - Discard parameter shards it has just collected\n\nIn backward path\n - Run all_gather to collect all shards from all ranks to recover the full parameter in this FSDP unit\n - Run backward computation\n - Run reduce_scatter to sync gradients\n - Discard parameters.\n\nOne way to view FSDP’s sharding is to decompose the DDP gradient all-reduce into reduce-scatter and all-gather. Specifically, during the backward pass, FSDP reduces and scatters gradients, ensuring that each rank possesses a shard of the gradients. Then it updates the corresponding shard of the parameters in the optimizer step. Finally, in the subsequent forward pass, it performs an all-gather operation to collect and combine the updated parameter shards.\n\n\n\n### 2. Custom FSDP underlying data structure\n\nTo implement the FSDP functionality described above, the custom FSDP is designed with the following Python classes and data structure:\n\n\n\n### 3. The custom FSDP interface: FullyShardedDataParallel\n\nThe custom FSDP provides the same programming interface as PyTorch's DistributedDataParallel (DDP) as FullyShardedDataParallel (FSDP). For example, you can apply FSDP to models as follows:\n\n```python\n# Initialize model and optimizer\nddp_config.use_megatron_fsdp = True\nddp_config.data_parallel_sharding_strategy = \"optim_grads_params\"\nmodel = GPTModel(transformer_config)\nmodel = FullyShardedDataParallel(\n transformer_config,\n model,\n ddp_config,\n fsdp_unit_modules = [TransformerLayer, LanguageModelEmbedding],\n)\noptimizer = torch.optim.AdamW(model.parameters(), lr=lr)\noptimizer = DistributedOptimizer(optimizer, [model], [model.param_and_grad_buffer])\n\n# Training loop\ndef train_step(inputs, labels):\n optimizer.zero_grad()\n for mbs_input, mbs_label in zip(inputs, labels):\n outputs = model(mbs_input)\n loss = loss_fn(outputs, mbs_label)\n loss.backward()\n optimizer.step()\n\n# Save and load model and optimizer state dict\ndef model_and_optimizer_state_dict():\n state_dict = {\n \"model\": model.sharded_state_dict(),\n \"optimizer\": optimizer.sharded_state_dict(),\n }\n return state_dict\n\ndef load_model_and_optimizer_state_dict(state_dict):\n model.load_state_dict(state_dict[\"model\"])\n optimizer.load_state_dict(state_dict[\"optimizer\"])\n```\n\n**Key Notes:**\n - You can configure which modules should be treated as FSDP units via the `fsdp_unit_modules` argument. This configuration is mandatory.\n - The custom FSDP must be used with a distributed optimizer since it provides distributed checkpointing.\n - The data-parallel communication group for parameters is not explicitly shown. Custom FSDP configures these groups as either DP (data-parallel) or EDP (expert data-parallel) based on parameter markings.\n\n#### 3.1 Initializing Models on the Meta Device\n\nFor training particularly large models with FSDP, you can initialize the model on the meta device. Using PyTorch's `reset_parameters` API, you can initialize model weights layer by layer during the construction of the `ParamAndGradBuffer`. Most PyTorch native modules and TransformerEngine modules support this API (e.g., [PyTorch Linear](https://github.com/pytorch/pytorch/blob/v2.6.0/torch/nn/modules/linear.py#L114), [TE LayerNormLinear](https://github.com/NVIDIA/TransformerEngine/blob/release_v2.0/transformer_engine/pytorch/module/layernorm_linear.py#L1107)).\n\n```python\n# Initialize model on meta device\nwith torch.device(\"meta\"):\n model = GPTModel(config)\n\nmodel = FullyShardedDataParallel(\n transformer_config,\n model,\n ddp_config,\n fsdp_unit_modules=[TransformerLayer, LanguageModelEmbedding],\n)\n```\n\n**Important Considerations:**\n1. *Custom Modules*: If your model contains custom modules, ensure they implement the `reset_parameters` API. Otherwise, you may need to force parameter initialization on a CUDA or CPU device.\n2. *Tensor Initialization*: Be cautious of tensors created during model initialization without a specified device—they will default to the meta device. To avoid issues, explicitly specify the device for these tensors to ensure compatibility with this function.\n\n### 4. Interaction between Custom FSDP and Model Forward/Backward Propagation\n\nCustom FSDP implements Fully Sharded Data Parallelism (FSDP) through a series of module hooks, gradient hooks, or by adding functions between modules. This involves inserting communications and manipulating parameters and gradients during PyTorch's module forward or backward propagation.\n\nModule hooks summary:\n- Module pre-forward hook(`module.register_forward_pre_hook`): This hook unshards model weights before the forward pass. In the case of an FSDP Unit Module, add a RegisterFSDPBackwardFunction function that will reshard model weights and reduce gradients after module backward propagation.\n- Module post-forward hook(`module.register_forward_hook`): This hook is used to reshard model weights after the forward pass.\n- Root module pre-backward hook(`root_module.register_full_backward_pre_hook`): This hook checks that all model parameters are resharded, in order to avoid unnecessary memory spikes. It also marks all modules as being in the `TrainingState.PRE_BACKWARD` state.\n- Module pre-backward hook(`module.register_full_backward_pre_hook`): This hook is used to unshard the model weights before the backward pass.\n- Root module post-backward hook(`torch.autograd.Variable._execution_engine.queue_callback`): This hook is used to make sure all gradients in the backprop are properly handled / available.\n\nThe gradient reduction pipeline maintains a map of gradients to FSDP parameter groups. If all gradients in an FSDP parameter group are ready, it launches a gradient reduction. Note that this assumes that the model's gradients are always generated in a certain order (reverse of `module.parameters()`), as otherwise, FSDP would maintain too many parameter group grad buffers, leading to excessive memory usage.\n\n#### 4.1 Optimized for Activation Recompute\n\nUsing the activation recompute will cause the same module to execute the forward function first and then the backward function in the backward prop, which will cause model weights unshard twice and model weights reshard twice. If we can tell program that this is a forward + backward operation, we can just call unshard once and reshard once.\n\nTo make this determination, we keep track of the model's state with training_state, `FORWARD`, `PRE_BACKWARD`, `POST_BACKWARD`, `IDLE`. It's worth noting that pre-backward hook act before pre-forward hook, and we'll let pre-backward hook execute the model weight unshard, and then mark the model as `PRE_BACKWARD`, and when pre-forward hook sees this marking it will not perform the unshard operation. Similarly, for model weight reshard duplicate, post-forward hook act before post-backward function, and checking for the `PRE_BACKWARD` flag in the post-forward hook will cancel the unshard.\n\n### 5. Memory Mechanisms and Features of Custom FSDP\n\nFSDP can fully distribute the model parameters, gradients, and optimizer states, and for mixed-precision training, it can also fully distribute the high-precision main weights. This is pretty much distributes all the memory except for the activation memory, but FSDP will also face some memory issues.\n\nFSDP frequently unshards and reshards model weights, which can lead to busy memory allocation and deallocation. This results in untimely tensor releases, causing memory spikes (or even out-of-memory errors), crashes of the PyTorch memory allocator cache, and a large number of `cudaMalloc` and `cudaFree` calls. These issues can significantly slow down the system.\n\nThe problem of untimely tensor release can generally be addressed using the `tensor._typed_storage(). _resize_(0)` API, which immediately deallocates the storage's memory. Custom FSDP provides interfaces in `AllGatherPipeline` and `GradReducePipeline` to replace the temporary buffer memory allocator used for parameter gathering and gradient reduction with ` StorageResizeBasedBucketAllocator`. This replaces the tensor release operation with the `tensor._typed_storage(). _resize_(0)` API.\n\nThe PyTorch memory allocator cache crash is a complex issue that occurs frequently when the actual memory usage approaches the GPU memory limit, leading to poor performance. This problem is challenging and can only be mitigated by avoiding frequent hits on the GPU memory limit. Using a self-managed memory allocator like ` RotaryBucketAllocator` is another potential solution. However, note that `RotaryBucketAllocator` is not yet mature.\n\n## References\n\n- [Getting Started with Fully Sharded Data Parallel (FSDP)](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html)\n"
},
{
"path": "docs/user-guide/features/dist_optimizer.md",
"content": "\n\n# Distributed Optimizer\n\nThe 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.\n\nTheoretical memory savings vary depending on the combination of the datatype of the model's parameters (`param_dtype`) and main gradients accumulated across data-parallel replicas (`grad_dtype`). We always use `fp32` main parameters for optimizer steps. In the current implementation, the theoretical number of bytes per parameter is (where d is the data parallel size):\n\n| | Non-distributed optim | Distributed optim |\n| ------ | ------ | ------ |\n| `fp16` parameters, `fp16` gradients | 20 | 4 + 16/d |\n| `bf16` parameters, `fp32` gradients | 18 | 6 + 12/d |\n| `fp32` parameters, `fp32` gradients | 16 | 8 + 8/d |\n\nOur implementation of the distributed optimizer uses contiguous buffers for parameters and main gradients; model gradients are copied over to the main gradients as soon as they are fully computed.\n\nThe figures below illustrate the distributed optimizer's sharding scheme, and the key steps of the distributed optimizer's parameter update:\n\n## Data flow\n\n\n\n## Sharding scheme\n\n\n\n## Key steps\n\n_(note: using illustrations above, assuming `bf16` model weights, `bf16` model gradients that are computed by the backward pass and `fp32` main gradients that are also used for optimizer steps; we always use `fp32` main weights for optimizer steps)_\n\n- Backward pass finishes (gradient buffer holds 16 `fp32` gradient elements).\n- Call reduce-scatter on each DP rank.\n- Each DP rank now has 4 elements within the gradient buffer that are fully reduced (remaining 12 elements are garbage).\n - DP rank 0 has gradient values for elements [0:4].\n - DP rank 1 has gradient values for elements [4:8].\n - DP rank 2 has gradient values for elements [8:12].\n - DP rank 3 has gradient values for elements [12:16].\n- Optimizer.step().\n- Each DP rank copies its 4 `fp32` main parameter elements into the corresponding `bf16` parameter buffer (each element is cast from fp32 to fp16).\n- Call all-gather on each DP rank.\n- The parameter buffer now contains all 16, fully updated, `bf16` model parameter elements. Parameters in PyTorch modules already point to the appropriate locations in this parameter buffer, and thus forward passes are ready to run after the all-gather completes.\n- At this point, the gradient buffer is also ready to be zero'd for the next iteration.\n"
},
{
"path": "docs/user-guide/features/fine_grained_activation_offloading.md",
"content": "\n\n# Fine-grained Activation Offloading (collaborated with rednote)\n\nMemory 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.\n\nCurrently, the supported offloading modules are `\"attn_norm\", \"core_attn\", \"attn_proj\", \"mlp_norm\", \"expert_fc1\", \"moe_act\"`, which could work with fine-grained recomputation to release almost all activations of a transformer layer.\n\n**Features**\n* Support PP=1/PP/Interleaved PP\n* Compatible with fine-grained recomputation\n* Support FP8\n* Support MTP\n* Support mixed dense & moe layer\n* Support A2A Overlap\n* Support CUDA Graph\n * (Temporary) cuda graph scope cannot contains the offloading modules\n\n**Usage**\n```bash\n# Enable fine-grained activation offloading\n--fine-grained-activation-offloading\n\n# Specify which modules are going to offload its input\n# Choices: \"attn_norm\", \"core_attn\", \"attn_proj\", \"mlp_norm\", \"expert_fc1\", \"moe_act\".\n--offload-modules expert_fc1\n```\n**Compatible with Fine-grained Recomputation**\n- For modules with minor perf overhead like layernorm or moe_act, use recomputing to reduce memory footprint;\n- For other modules, use offloading to reduce memory footprint;\n- Make sure the offloading/reloading could be overlapped with computing;\n\n\n"
},
{
"path": "docs/user-guide/features/index.md",
"content": "\n\n# Advanced Features\n\nAdvanced feature guides for key Megatron Core capabilities.\n\n```{toctree}\n:maxdepth: 2\n\nfine_grained_activation_offloading\nmoe\ncontext_parallel\ncustom_fsdp\ndist_optimizer\noptimizer_cpu_offload\npipeline_parallel_layout\ntokenizers\nmegatron_energon\nmegatron_rl\n```\n"
},
{
"path": "docs/user-guide/features/megatron_energon.md",
"content": "\n\n# Megatron Energon\n\nAdvanced multimodal dataloader for efficient loading of text, images, video, and audio at scale.\n\n## Overview\n\n[**Megatron Energon**](https://github.com/NVIDIA/Megatron-Energon) is purpose-built for large-scale multimodal training with:\n\n- **Multimodal support** - Text, images, video, audio\n- **Distributed loading** - Optimized for multi-node training\n- **Data blending** - Mix datasets with configurable weights\n- **WebDataset format** - Efficient streaming from cloud storage\n- **State management** - Save and restore training position\n\n## Installation\n\n```bash\npip install megatron-energon\n```\n\n## Key Features\n\n### Data Processing\n\n- **Packing** - Optimize sequence length utilization\n- **Grouping** - Smart batching of similar-length sequences\n- **Joining** - Combine multiple dataset sources\n- **Object storage** - Stream from S3, GCS, Azure Blob Storage\n\n### Production-Ready\n\n- Distributed loading across workers and nodes\n- Checkpoint data loading state\n- Memory-efficient streaming\n- Parallel data loading with prefetching\n\n## Basic Usage\n\n```python\nfrom megatron.energon import get_train_dataset, get_loader, WorkerConfig\n\n# Create dataset\nds = get_train_dataset(\n '/path/to/dataset',\n batch_size=32,\n shuffle_buffer_size=1000,\n worker_config=WorkerConfig.default_worker_config(),\n)\n\n# Create loader and iterate\nfor batch in get_loader(ds):\n # Training step\n pass\n```\n\n## Multimodal Example\n\n```python\n# Load image-text dataset\nds = get_train_dataset(\n '/path/to/multimodal/dataset',\n batch_size=32,\n worker_config=WorkerConfig(num_workers=8, prefetch_factor=2),\n)\n\nfor batch in get_loader(ds):\n images = batch['image'] # Image tensors\n texts = batch['text'] # Text captions\n # Process batch\n```\n\n## Dataset Blending\n\nMix multiple datasets with custom weights:\n\n```python\nfrom megatron.energon import Blender\n\nblended_ds = Blender([\n ('/path/to/dataset1', 0.6), # 60%\n ('/path/to/dataset2', 0.3), # 30%\n ('/path/to/dataset3', 0.1), # 10%\n])\n```\n\n## Configuration\n\n### Worker Configuration\n\n```python\nWorkerConfig(\n num_workers=8, # Parallel workers\n prefetch_factor=2, # Batches to prefetch per worker\n persistent_workers=True, # Keep workers alive between epochs\n)\n```\n\n### Common Parameters\n\n| Parameter | Description |\n|-----------|-------------|\n| `batch_size` | Samples per batch |\n| `shuffle_buffer_size` | Buffer size for randomization |\n| `max_samples_per_sequence` | Max samples to pack into one sequence |\n| `worker_config` | Worker configuration for parallel loading |\n\n## Integration with Megatron-LM\n\n```python\nfrom megatron.energon import get_train_dataset, get_loader\nfrom megatron.training import get_args\n\nargs = get_args()\n\ntrain_ds = get_train_dataset(\n args.data_path,\n batch_size=args.micro_batch_size,\n)\n\nfor iteration, batch in enumerate(get_loader(train_ds)):\n loss = train_step(batch)\n```\n\n## Resources\n\n- **[Megatron Energon GitHub](https://github.com/NVIDIA/Megatron-Energon)** - Documentation and examples\n- **[Multimodal Examples](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/multimodal)** - Megatron-LM multimodal training\n\n## Next Steps\n\n- Check [Multimodal Models](../../models/multimodal.md) for supported architectures\n- See [Training Examples](../training-examples.md) for integration examples\n"
},
{
"path": "docs/user-guide/features/megatron_rl.md",
"content": "\n\n# Megatron RL\n\nReinforcement learning library for post-training large language models at scale.\n\n## Overview\n\n[**Megatron RL**](https://github.com/NVIDIA/Megatron-LM/tree/dev/megatron/rl) adds native reinforcement learning capabilities to Megatron-LM for large-scale RL-based post-training of foundation models.\n\n> **Note**: Megatron RL is under active development and primarily designed for research teams exploring RL post-training on modern NVIDIA hardware. For production deployments, use [**NeMo RL**](https://github.com/NVIDIA-NeMo/RL).\n\n## Key Features\n\n- **Decoupled Design** - Clean separation between agent/environment logic and RL implementation\n- **Flexible Inference** - Support for Megatron, OpenAI, and HuggingFace inference backends\n- **Trainer/Evaluator** - Manages rollout generation and coordinates with inference systems\n- **Megatron Integration** - Native integration with Megatron Core inference system\n\n## Architecture\n\n### Components\n\n**Agents & Environments**\n- Accept inference handles\n- Return experience rollouts with rewards\n- Implement custom RL logic\n\n**Trainer/Evaluator**\n- Controls rollout generation\n- Coordinates with inference systems\n- Manages training loops\n\n**Inference Interface**\n- Provides `.generate(prompt, **generation_args)` endpoint\n- Supports multiple backends (Megatron, OpenAI, HuggingFace)\n\n## Use Cases\n\n- RLHF (Reinforcement Learning from Human Feedback)\n- Custom reward-based fine-tuning\n- Policy optimization for specific tasks\n- Research on RL post-training techniques\n\n## Resources\n\n- **[Megatron RL GitHub](https://github.com/NVIDIA/Megatron-LM/tree/dev/megatron/rl)** - Source code and documentation\n- **[Megatron Core Inference](../../api-guide/core/transformer.md)** - Native inference integration\n"
},
{
"path": "docs/user-guide/features/moe.md",
"content": "\n\n# Mixture of Experts\n\n```{toctree}\n:maxdepth: 1\n:caption: MoE Features\n\nmulti_token_prediction\nmulti_latent_attention\n../../api-guide/router_replay\n```\n\n```{include} ../../../megatron/core/transformer/moe/README.md\n```\n"
},
{
"path": "docs/user-guide/features/multi_latent_attention.md",
"content": "\n\n# Multi-Latent Attention\n\n## Multi-Latent Attention overview\n\nMulti-Latent Attention (\"MLA\") is an innovative attention mechanism introduced by Deepseek team that enhances the efficiency of attention computation by leveraging multiple latent spaces. This approach is particularly beneficial for large language models (LLMs), as it reduces the computational burden associated with traditional attention mechanisms. According to Deepseek-V2 technical report, MLA achieves better performance compared to Multi-Head Attention (MHA) and requires smaller KV cache.\n\n## Enabling Multi-Latent Attention\n\nTo enable MLA in Megatron-LM, set the following flags in command line:\n- `--multi-latent-attention` to enable MLA in MLP.\n- Set `MLATransformerConfig` to configure MLA.\n\n"
},
{
"path": "docs/user-guide/features/multi_token_prediction.md",
"content": "\n\n# Multi-Token Prediction (MTP)\n\nMulti-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\ndata efficiency. On the other hand, MTP may enable the model to pre-plan its representations for better prediction of future tokens. In this implementation of MTP, we sequentially predict additional tokens and keep the complete causal chain at each prediction depth. The following figure illustrates our implementation of MTP in [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3/).\n\n\n\nThe 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.\n\nFor more information, refer to [DeepSeek-V3 Technical Report](https://arxiv.org/pdf/2412.19437.pdf)\n\n## Related Arguments\n\nWe can train GPTModel like models with Multi-Token Prediction (MTP) by setting mtp_num_layers to be a positive integer.\n\n| Item | Description |\n| --- | --- |\n| mtp_num_layers | Number of Multi-Token Prediction (MTP) Layers. MTP extends the prediction scope to multiple future tokens at each position. This MTP implementation sequentially predict additional tokens by using D sequential modules to predict D additional tokens. Default is None. |\n| mtp_loss_scaling_factor | Scaling factor of Multi-Token Prediction (MTP) loss. We compute the average of the MTP losses across all depths, and multiply it the scaling factor to obtain the overall MTP loss, which serves as an additional training objective. Default is 0.1. |\n\n## Pipeline Parallel Layout for MTP\n\nMTP supports flexible placement of MTP layers across pipeline stages using a custom `pipeline_model_parallel_layout`. By default, all MTP layers are placed on the last pipeline stage, but you can customize their placement.\n\n### MTP Standalone Mode\n\nWhen MTP layers are placed in a separate virtual pipeline (vpp) stage that is not on the last pipeline rank, the `mtp_standalone` flag is automatically set to `True`. This mode enables MTP to run independently in its own pipeline stage.\n\n### Layout Format\n\nUse `m` to represent MTP layers in the pipeline layout string. For example:\n- `\"E|t*3|(t|)*5mL\"` - MTP in the last stage\n- `\"E|t*3|(t|)*4tm|L\"` - MTP in the second-to-last stage with a decoder layer\n- `\"E|t*3|(t|)*3tt|m|L\"` - MTP in a standalone stage (second-to-last) with no other layers\n\n### Constraints\n\n- All MTP layers must be placed in the same one virtual pipeline stage.\n- MTP layers cannot be placed on the first pipeline rank.\n\n## Implementation Notes\n\n- For models with MTP layers, the final layernorm is placed in the stage that contains the last decoder layer, rather than in the post-process stage. This may cause small numerical differences in gradient norm reduction when final layernorm is placed in different pipeline stages in deterministic mode. Bitwise alignment can be achieved by disabling gradient norm clipping.\n- MTP loss is computed in the post-processing stage.\n\n## Precautions\n\nDo not use Context Parallel (CP), or arbitrary AttnMaskType, or learned absolute position embedding type with MTP. These use cases are not yet supported.\n"
},
{
"path": "docs/user-guide/features/optimizer_cpu_offload.md",
"content": "\n\n# Optimizer CPU Offload\n\n```{include} ../../../megatron/core/optimizer/cpu_offloading/README.md\n```\n"
},
{
"path": "docs/user-guide/features/pipeline_parallel_layout.md",
"content": "\n\n# Custom Pipeline Model Parallel Layout\n\n*This is an experimental feature and may be changed.*\n\n`--pipeline-model-parallel-layout` is a flexible API for defining the pipeline parallel partitioning, which is essential for balanced partitioning for an imbalanced model. For example, to partition DeepSeek-V3 (61 decoder layers + 1 mtp layer) with PP16VPP2, we can include the arguments as follows:\n\n```bash\n--pipeline-model-parallel-size 16\n--pipeline-model-parallel-layout \"Et*3|(tt|)*29,m|L\"\n```\n\n| PP \\ VPP rank | 0 | 1 |\n|---------------|-------------------------|---------------|\n| 0 | embedding + 3 × decoder | 2 × decoder |\n| 1~13 | 2 × decoder | 2 × decoder |\n| 14 | 2 × decoder | mtp |\n| 15 | 2 × decoder | loss |\n\nIn the layout string, stages are split by '|'. Replicated stages or layers can be described with multiplication. Commas can be used cosmetically. Symbol choices:\n\n* `E` = embedding layer\n* `t` = transformer decoder layer\n* `m` = MTP layer\n* `L` = loss calculation layer\n\nNote that it is legal to have empty stages, e.g., `E||t|L` (the second stage is empty).\n"
},
{
"path": "docs/user-guide/features/tokenizers.md",
"content": "\n\n# Tokenizers\n\nMegatron Core provides a unified tokenizer system with a HuggingFace-style API for easy tokenizer management and configuration.\n\n## Overview\n\nThe `MegatronTokenizer` class offers a simple, familiar API for loading and managing tokenizers:\n\n- **Automatic detection** - Load any tokenizer type without specifying the library\n- **Metadata-based configuration** - Store tokenizer settings in JSON for easy reuse\n- **HuggingFace-compatible API** - Familiar `.from_pretrained()` interface\n- **Custom tokenizer support** - Extend with model-specific tokenization logic\n\n## Key Features\n\n### Unified API\n\nUse the same API regardless of tokenizer backend (SentencePiece, HuggingFace, TikToken, etc.):\n\n```python\nfrom megatron.core.tokenizers import MegatronTokenizer\n\ntokenizer = MegatronTokenizer.from_pretrained(\"/path/to/tokenizer\")\n```\n\n### Tokenizer Metadata\n\nConfiguration is stored in a JSON metadata file containing:\n- Tokenizer library (HuggingFace, SentencePiece, TikToken, etc.)\n- Chat templates\n- Custom tokenizer class\n- Special token configurations\n\n**Benefits:**\n- Set configuration once, reuse everywhere\n- No repeated CLI arguments\n- Easy sharing - just copy the tokenizer directory\n\n### Automatic Library Detection\n\nThe correct tokenizer implementation is automatically selected:\n- No need to specify `SentencePieceTokenizer`, `HuggingFaceTokenizer`, etc.\n- Library type detected from metadata\n- Seamless switching between tokenizer backends\n\n## Basic Usage\n\n### Creating Tokenizer Metadata\n\nSave tokenizer configuration for reuse:\n\n```python\nfrom megatron.core.tokenizers import MegatronTokenizer\n\n# Create metadata for a SentencePiece tokenizer\nMegatronTokenizer.write_metadata(\n tokenizer_path=\"/path/to/tokenizer.model\",\n tokenizer_library=\"sentencepiece\",\n chat_template=\"{% for message in messages %}{{ message.content }}{% endfor %}\",\n)\n```\n\nThe metadata is saved as `tokenizer_metadata.json` in the tokenizer directory.\n\n### Loading a Tokenizer\n\nLoad from a directory with metadata:\n\n```python\nfrom megatron.core.tokenizers import MegatronTokenizer\n\n# Load with auto-detected configuration\ntokenizer = MegatronTokenizer.from_pretrained(\"/path/to/tokenizer.model\")\n```\n\n### Loading with Custom Metadata Path\n\nIf metadata is stored separately:\n\n```python\ntokenizer = MegatronTokenizer.from_pretrained(\n tokenizer_path=\"/path/to/tokenizer.model\",\n metadata_path=\"/path/to/custom/metadata.json\",\n)\n```\n\n### Loading with Inline Metadata\n\nPass metadata as a dictionary:\n\n```python\ntokenizer = MegatronTokenizer.from_pretrained(\n tokenizer_path=\"GPT2BPETokenizer\",\n metadata_path={\"library\": \"megatron\"},\n vocab_file=\"/path/to/vocab.txt\",\n)\n```\n\n## Advanced Usage\n\n### Custom Tokenizer Classes\n\nCreate model-specific tokenization logic:\n\n```python\nfrom megatron.core.tokenizers.text import MegatronTokenizerText\n\nclass CustomTokenizer(MegatronTokenizerText):\n def encode(self, text):\n # Custom encoding logic\n return super().encode(text)\n\n def decode(self, tokens):\n # Custom decoding logic\n return super().decode(tokens)\n\n# Save metadata with custom class\nMegatronTokenizer.write_metadata(\n tokenizer_path=\"/path/to/tokenizer.model\",\n tokenizer_library=\"sentencepiece\",\n tokenizer_class=CustomTokenizer,\n)\n```\n\n### TikToken Tokenizers\n\nConfigure TikToken-based tokenizers:\n\n```python\ntokenizer = MegatronTokenizer.from_pretrained(\n tokenizer_path=\"/path/to/tokenizer/model.json\",\n metadata_path={\"library\": \"tiktoken\"},\n pattern=\"v2\",\n num_special_tokens=1000,\n)\n```\n\n### Null Tokenizer\n\nUse a null tokenizer for testing or non-text models:\n\n```python\ntokenizer = MegatronTokenizer.from_pretrained(\n metadata_path={\"library\": \"null-text\"},\n vocab_size=131072,\n)\n```\n\n## Integration with Megatron-LM\n\n### Using with Training Scripts\n\nThe tokenizer system integrates seamlessly with Megatron-LM training:\n\n```bash\n# Null tokenizer for testing\ntorchrun --nproc_per_node=8 pretrain_gpt.py \\\n --tokenizer-type NullTokenizer \\\n --vocab-size 131072 \\\n ...\n```\n\n```bash\n# HuggingFace tokenizer with metadata\ntorchrun --nproc_per_node=8 pretrain_gpt.py \\\n --tokenizer-type HuggingFaceTokenizer \\\n --tokenizer-model meta-llama/Meta-Llama-3-8B \\\n --tokenizer-metadata /path/to/metadata.json \\\n ...\n```\n\n### Auto-Generated Metadata\n\nIf `--tokenizer-metadata` is not specified, a default metadata file is generated automatically based on the tokenizer type.\n\n## Supported Tokenizer Libraries\n\n| Library | Description | Use Case |\n|---------|-------------|----------|\n| **HuggingFace** | Transformers tokenizers | Most modern LLMs (LLaMA, Mistral, etc.) |\n| **SentencePiece** | Google's tokenizer | GPT-style models, custom vocabularies |\n| **TikToken** | OpenAI's tokenizer | GPT-3.5/GPT-4 style tokenization |\n| **Megatron** | Built-in tokenizers | Legacy GPT-2 BPE |\n| **Null** | No-op tokenizer | Testing, non-text modalities |\n\n## Common Tokenizer Types\n\n### LLaMA / Mistral\n\n```python\nMegatronTokenizer.write_metadata(\n tokenizer_path=\"/path/to/llama/tokenizer.model\",\n tokenizer_library=\"sentencepiece\",\n)\n```\n\n### GPT-2\n\n```python\nMegatronTokenizer.write_metadata(\n tokenizer_path=\"GPT2BPETokenizer\",\n tokenizer_library=\"megatron\",\n vocab_file=\"/path/to/gpt2-vocab.json\",\n merge_file=\"/path/to/gpt2-merges.txt\",\n)\n```\n\n## Best Practices\n\n1. **Always save metadata** - Create metadata once, reuse across training runs\n2. **Use HuggingFace tokenizers** - When possible, for modern LLM compatibility\n3. **Test tokenization** - Verify encode/decode before starting training\n4. **Version control metadata** - Include `tokenizer_metadata.json` in your experiment configs\n5. **Share tokenizer directories** - Include both model files and metadata for reproducibility\n\n## Next Steps\n\n- **Prepare Data**: See [Data Preparation](../data-preparation.md) for preprocessing with tokenizers\n- **Train Models**: Use tokenizers in [Training Examples](../training-examples.md)\n- **Supported Models**: Check [Language Models](../../models/llms.md) for model-specific tokenizers\n"
},
{
"path": "docs/user-guide/index.md",
"content": "---\norphan: true\n---\n\n\n\n# User Guide\n\nComprehensive guides for using Megatron Core and Megatron-LM.\n\n```{toctree}\n:maxdepth: 2\n\nmsc_integration\ndata-preparation\ntraining-examples\nparallelism-guide\nfeatures/index\n```\n"
},
{
"path": "docs/user-guide/msc_integration.md",
"content": "\n\n```{include} ../../megatron/core/MSC_Integration.md\n```\n\n"
},
{
"path": "docs/user-guide/parallelism-guide.md",
"content": "\n\n# Parallelism Strategies Guide\n\nMegatron Core supports multiple parallelism strategies that can be combined to efficiently train models from billions to trillions of parameters across thousands of GPUs.\n\n## Overview\n\n| Strategy | What it parallelizes | Best for |\n|----------|---------------------|----------|\n| **Data Parallelism (DP)** | Batch dimension | Standard training, most common |\n| **Tensor Parallelism (TP)** | Individual layers | Large layers, GPU memory constraints |\n| **Pipeline Parallelism (PP)** | Model depth | Very deep models |\n| **Context Parallelism (CP)** | Sequence length | Long sequences (8K+ tokens) |\n| **Expert Parallelism (EP)** | MoE experts | Mixture-of-Experts models |\n\n## Data Parallelism (DP)\n\nReplicate the model across GPUs and split the batch.\n\n### Standard Data Parallel (DDP)\n\n```bash\ntorchrun --nproc_per_node=8 pretrain_gpt.py \\\n --data-parallel-sharding-strategy no_shard\n```\n\nEach GPU has a full copy of the model and processes a portion of the batch.\n\n### Fully Sharded Data Parallel (FSDP)\n\nShard model parameters, gradients, and optimizer states to reduce memory:\n\n```bash\n# Megatron FSDP (~15% faster than PyTorch FSDP2)\n--use-megatron-fsdp \\\n--data-parallel-sharding-strategy optim_grads_params\n```\n\n**Sharding strategies:**\n- `optim` - Shard optimizer states only (ZeRO-1)\n- `optim_grads` - Shard gradients + optimizer (ZeRO-2)\n- `optim_grads_params` - Shard parameters + gradients + optimizer (ZeRO-3)\n\n## Tensor Parallelism (TP)\n\nSplit individual model layers across GPUs. Recommended for large hidden dimensions.\n\n```bash\n--tensor-model-parallel-size 4 # 4-way tensor parallelism\n--sequence-parallel # Enable sequence parallelism (recommended)\n```\n\n**When to use:**\n- Model layers don't fit on single GPU\n- Large hidden dimensions (4096+)\n- Usually combined with DP and PP\n\n## Pipeline Parallelism (PP)\n\nSplit model layers across GPUs vertically (by depth).\n\n```bash\n--pipeline-model-parallel-size 8 # 8 pipeline stages\n--num-layers-per-virtual-pipeline-stage 4 # Virtual pipeline for load balancing\n```\n\n**When to use:**\n- Very deep models (50+ layers)\n- Combine with TP for large models\n- Helps distribute memory across GPUs\n\n## Context Parallelism (CP)\n\nSplit long sequences across GPUs for efficient long-context training.\n\n```bash\n--context-parallel-size 2 # 2-way context parallelism\n--cp-comm-type p2p # Communication type\n```\n\n**When to use:**\n- Long sequences (8K+ tokens)\n- Reduces activation memory\n- Can combine with TP, PP, DP\n\n**→ [Context Parallelism Deep Dive](features/context_parallel.md)** - Detailed guide with performance analysis\n\n## Expert Parallelism (EP)\n\nDistribute experts across GPUs in Mixture-of-Experts models.\n\n```bash\n--expert-model-parallel-size 8 # 8-way expert parallelism\n--num-experts 64 # 64 experts per MoE layer\n--moe-grouped-gemm # Optimize expert computation\n```\n\n**Important:** When combining EP with TP, you **must enable Sequence Parallelism**:\n\n```bash\n--tensor-model-parallel-size 4\n--expert-model-parallel-size 8\n--sequence-parallel # Required when using TP + EP\n```\n\n## Parallelism Selection Guide\n\nRecommended configurations based on [NVIDIA NeMo production setups](https://github.com/NVIDIA/NeMo/tree/main/scripts/performance/recommended_model_configs):\n\n### Language Models\n\n| Model | Size | GPUs | TP | PP | CP | EP | Configuration Notes |\n|-------|------|------|----|----|----|----|---------------------|\n| **LLaMA-3** | 8B | 8 | 1 | 1 | 2 | 1 | CP=2 for long context (8K seqlen) |\n| **LLaMA-3** | 70B | 64 | 4 | 4 | 2 | 1 | Balanced TP+PP for 70B scale |\n| **LLaMA-3.1** | 405B | 1024 | 8 | 8 | 2 | 1 | 3D parallelism (TP+PP+CP) |\n| **GPT-3** | 175B | 128-512 | 4 | 8 | 1 | 1 | Standard large model config |\n\n### Mixture-of-Experts Models\n\n| Model | Size | GPUs | TP | PP | CP | EP | Configuration Notes |\n|-------|------|------|----|----|----|----|---------------------|\n| **Mixtral** | 8x7B | 64 | 1 | 4 | 1 | 8 | EP=8 for 8 experts |\n| **Mixtral** | 8x22B | 256 | 4 | 4 | 1 | 8 | TP+PP+EP for large MoE |\n| **DeepSeek-V3** | 671B | 1024 | 2 | 16 | 1 | 64 | Massive MoE with 256 experts |\n\n## Combining Strategies\n\n### Total GPU Count\n\nThe total number of GPUs is calculated as:\n\n```\nTotal GPUs = TP × PP × CP × EP × DP\n```\n\n### Example: LLaMA-3 70B on 64 GPUs\n\n```bash\n# TP=4, PP=4, CP=2, DP=2 => 4 × 4 × 2 × 2 = 64 GPUs\ntorchrun --nproc_per_node=8 pretrain_gpt.py \\\n --tensor-model-parallel-size 4 \\\n --pipeline-model-parallel-size 4 \\\n --context-parallel-size 2 \\\n --num-layers 80 \\\n --hidden-size 8192 \\\n --num-attention-heads 64 \\\n --seq-length 8192 \\\n --micro-batch-size 1 \\\n --global-batch-size 512 \\\n --bf16\n```\n\n## Performance Optimizations\n\n### Communication Overlap\n\nEnable overlapping of communication with computation:\n\n```bash\n--overlap-grad-reduce # Overlap gradient reduction with backward pass\n--overlap-param-gather # Overlap parameter gathering with forward pass\n--tp-comm-overlap # Overlap TP communication\n```\n\n### Distributed Optimizer\n\nRecommended for all multi-GPU training:\n\n```bash\n--use-distributed-optimizer\n```\n\nBenefits:\n- Faster checkpointing\n- Reduced memory when combined with FSDP\n- Better performance at scale\n\n### Sequence Parallelism\n\nAlways enable when using TP:\n\n```bash\n--sequence-parallel\n```\n\nReduces activation memory by sharding sequence dimension in LayerNorm and Dropout.\n\n## Choosing the Right Strategy\n\n### Start Simple\n1. Begin with **Data Parallelism** (DP) only\n2. Add **Tensor Parallelism** (TP) if model doesn't fit\n3. Add **Pipeline Parallelism** (PP) for very large models\n4. Add **Context Parallelism** (CP) for long sequences\n\n### Memory Constraints\n- Use **FSDP** to reduce memory per GPU\n- Use **TP** to split large layers\n- Use **PP** to split model depth\n- Enable **activation checkpointing** for extreme cases\n\n### Communication Bottlenecks\n- Reduce **TP** degree (increases memory per GPU)\n- Increase **PP** degree (may reduce efficiency)\n- Use **CP** instead of larger TP for long sequences\n\n## Next Steps\n\n- **API Reference**: See [Tensor Parallel](../api-guide/core/tensor_parallel.md) and [Pipeline Parallel](../api-guide/core/pipeline_parallel.md) API documentation\n- **Advanced Features**: Explore [Megatron FSDP](features/custom_fsdp.md) and [Distributed Optimizer](features/dist_optimizer.md)\n- **Performance Tuning**: Check [NVIDIA NeMo Performance Guide](https://docs.nvidia.com/nemo-framework/user-guide/latest/performance/performance-guide.html)\n"
},
{
"path": "docs/user-guide/training-examples.md",
"content": "\n\n# Training Examples\n\nGet started with Megatron Core training using these practical examples.\n\n## Simple Training Example\n\nThe simplest way to get started is with the basic training loop using mock data:\n\n```bash\n# Distributed training on 2 GPUs with mock data\ntorchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py\n```\n\nThis example:\n- Runs on 2 GPUs\n- Uses generated mock data (no data preparation needed)\n- Demonstrates basic distributed training setup\n- Perfect for testing your installation\n\n## LLaMA-3 Training Examples\n\n### LLaMA-3 8B with FP8\n\nTrain LLaMA-3 8B model with FP8 mixed precision on 8 GPUs:\n\n```bash\n./examples/llama/train_llama3_8b_h100_fp8.sh\n```\n\n**Configuration:**\n- 8 GPUs\n- FP8 mixed precision (requires Hopper/Ada/Blackwell GPUs)\n- Mock data for quick testing\n\n### Custom LLaMA Training\n\nFor training with your own data:\n\n```bash\ntorchrun --nproc_per_node=8 pretrain_gpt.py \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n --num-layers 32 \\\n --hidden-size 4096 \\\n --num-attention-heads 32 \\\n --seq-length 2048 \\\n --max-position-embeddings 2048 \\\n --micro-batch-size 4 \\\n --global-batch-size 32 \\\n --train-iters 100000 \\\n --lr 3.0e-4 \\\n --min-lr 3.0e-5 \\\n --lr-decay-style cosine \\\n --lr-warmup-iters 2000 \\\n --weight-decay 0.1 \\\n --clip-grad 1.0 \\\n --bf16 \\\n --data-path /path/to/your/preprocessed_data \\\n --split 949,50,1 \\\n --save /path/to/checkpoints \\\n --load /path/to/checkpoints \\\n --log-interval 10 \\\n --save-interval 1000 \\\n --eval-interval 1000\n```\n\n## GPT-3 Training Example\n\nTrain a GPT-3 style model:\n\n```bash\ntorchrun --nproc_per_node=8 pretrain_gpt.py \\\n --tensor-model-parallel-size 2 \\\n --pipeline-model-parallel-size 2 \\\n --num-layers 24 \\\n --hidden-size 2048 \\\n --num-attention-heads 16 \\\n --seq-length 1024 \\\n --max-position-embeddings 1024 \\\n --micro-batch-size 2 \\\n --global-batch-size 16 \\\n --train-iters 100000 \\\n --lr 1.5e-4 \\\n --min-lr 1.0e-5 \\\n --lr-decay-style cosine \\\n --lr-warmup-iters 1000 \\\n --weight-decay 0.1 \\\n --clip-grad 1.0 \\\n --fp16 \\\n --data-path /path/to/preprocessed_data \\\n --split 949,50,1 \\\n --save /path/to/checkpoints \\\n --load /path/to/checkpoints\n```\n\n## Key Training Arguments\n\n### Model Architecture\n\n| Argument | Description |\n|----------|-------------|\n| `--num-layers` | Number of transformer layers |\n| `--hidden-size` | Hidden dimension size |\n| `--num-attention-heads` | Number of attention heads |\n| `--seq-length` | Sequence length for training |\n\n### Training Configuration\n\n| Argument | Description |\n|----------|-------------|\n| `--micro-batch-size` | Batch size per GPU |\n| `--global-batch-size` | Total batch size across all GPUs |\n| `--train-iters` | Number of training iterations |\n\n### Learning Rate\n\n| Argument | Description |\n|----------|-------------|\n| `--lr` | Peak learning rate |\n| `--min-lr` | Minimum learning rate |\n| `--lr-decay-style` | LR schedule (cosine, linear, constant) |\n| `--lr-warmup-iters` | Warmup iterations |\n\n### Mixed Precision\n\n| Argument | Description |\n|----------|-------------|\n| `--fp16` | FP16 mixed precision |\n| `--bf16` | BF16 mixed precision (recommended) |\n| `--fp8-hybrid` | FP8 mixed precision (Hopper/Ada/Blackwell) |\n\n### Data and Checkpointing\n\n| Argument | Description |\n|----------|-------------|\n| `--data-path` | Path to preprocessed data |\n| `--split` | Train/validation/test split (e.g., 949,50,1) |\n| `--save` | Checkpoint save directory |\n| `--load` | Checkpoint load directory |\n| `--save-interval` | Save checkpoint every N iterations |\n\n## Next Steps\n\n- **Optimize Performance**: See [Advanced Features](features/index.md) for FSDP, distributed optimizer, and other optimizations\n- **Scale Up**: Learn about [Parallelism Strategies](parallelism-guide.md) to train larger models across more GPUs\n- **Prepare Data**: Follow the [Data Preparation](data-preparation.md) guide to process your own datasets\n"
},
{
"path": "docs/versions1.json",
"content": "[\n {\n \"name\": \"nightly\",\n \"version\": \"nightly\",\n \"url\": \"https://docs.nvidia.com/megatron-core/developer-guide/nightly/\"\n },\n {\n \"name\": \"0.16.0 (latest)\",\n \"version\": \"0.16.0\",\n \"url\": \"https://docs.nvidia.com/megatron-core/developer-guide/latest/\"\n },\n {\n \"name\": \"0.15.0\",\n \"version\": \"0.15.0\",\n \"url\": \"https://docs.nvidia.com/megatron-core/developer-guide/0.15.0/\"\n }\n]\n"
},
{
"path": "examples/__init__.py",
"content": " "
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/README.md",
"content": "# SGEAT: Detoxify Larger-scale Language Models\n\nThis is the official code base for our NeurIPS 2022 paper:\n\n[Exploring the Limits of Domain-Adaptive Training for Detoxifying Large-Scale Language Models](https://arxiv.org/abs/2202.04173)\n\nBoxin Wang, Wei Ping, Chaowei Xiao, Peng Xu, Mostofa Patwary, Mohammad Shoeybi, Bo Li, Anima Anandkumar, Bryan Catanzaro\n\n\n## Citation\n\n```\n@article{WangExp2022,\n title={Exploring the Limits of Domain-Adaptive Training for Detoxifying Large-Scale Language Models},\n author={Wang, Boxin and Ping, Wei and Xiao, Chaowei and Xu, Peng and Patwary, Mostofa and Shoeybi, Mohammad and and Li, Bo and Anandkumar, Anima and Catanzaro, Bryan},\n journal={NeurIPS},\n year={2022}\n}\n```\n\n## Usage\n\n### Prepare your environment\n\nThe project environment is based on the standard [nvcr docker](nvcr.io/nvidia/pytorch:21.12-py3) of version `nvcr.io/nvidia/pytorch:21.12-py3`.\n\nTo run Perspective API, you need to install `google-api-python-client`\n```bash\npip install --upgrade google-api-python-client\n```\n\n### Self Generation\n\n#### SGEAT (Standard)\nTo perform unconditional generation for a Megatron LM, we provide an example script for 1.3B LM.\n\n```bash\n# [num of samples] [model checkpoint] [random seed]\nbash examples/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh 1000 checkpoints/gpt3/gpt3-1.3b/ 2333\n```\nThis will generate a jsonl file of 1000 generated text (as a toy example) at `selfgeneration/unconditional_generation_gpt3-1.3b/2333.out`. \n\nNote that you may want to set your own gpt2 vocab and merge file dir, as well as your output data dir in `selfgenerate-1.3b-unconditional.sh`.\n\n### Annotation\n\nWe then use Perspective API to annotate the self generated corpus. Note that you need to fill in your own Perspective API key in the `examples/detoxify_lm/perspective_api_annotate.py`. \n\n```bash\npython examples/detxoify_lm/perspective_api_annotate.py --data-path [input-data-path] --out-path [output-data-path] --workers 70\n```\n\nFor example,\n\n```bash\npython examples/detxoify_lm/annotations/perspective_api_annotate.py --data-path selfgeneration/unconditional_generation_gpt3-1.3b/2333.out --out-path selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.out --workers 70\n```\n\n### Filtering\n\nWe then filter the self annotated generated corpus to get the most nontoxic 50% of the corus.\n\nFor example,\n```bash\npython examples/detxoify_lm/annotations/filter-selfgeneration.py --data-path selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.out --out-path selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic.out\n```\n\nThis will generate a jsonl file of 500 text of the lowest toxicity (as a toy example) at `selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic.out`. \n\n\n### Preprocess\n\nWe then preprocess the dataset so that Megatron LM can use the dumped dataset to fine-tune.\n\n```\nbash examples/detxoify_lm/annotations/preprocess.sh selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic.out selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic\n```\n\nThis will generate two files as follows\n```bash\nselfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic_text_document.idx\nselfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic_text_document.bin\n```\nwhich will be used in the following domain-adative training step.\n\n### Fine-tuning\n\nWe then use the preprocess dataset as input to fine-tune our Megatron-LM. \n```bash\n# [fine-tuning dataset] [output-dir] [lr] [bs] [train-iters] [load checkpoint]\nbash examples/detxoify_lm/finetune_gpt_distributed-1.3b.sh selfgeneration/unconditional_generation_gpt3-1.3b/2333.annotated.nontoxic_text_document gpt3-1.3b-toy-example-lr-2e-5-bs-512 2e-5 512 78 checkpoints/gpt3/gpt3-1.3b\n```\n\nThis will dump the final checkpoint in `$SHARE_DATA/gpt3-1.3b-toy-example-lr-2e-5-bs-512`. (`$SHARE_DATA` is your current work dir, default to `$PWD`)\n\n### Evaluation\n\nWe then use the fine-tuned checkpoint to perform conditional generation given RealToxicityPrompts:\n\n```bash\n# [input-prompts] [model-checkpoint]\nbash examples/detxoify_lm/generate-1.3b.sh augmented_prompts.jsonl $SHARE_DATA/gpt3-1.3b-toy-example-lr-2e-5-bs-512\n```\nFor example, this will generate the continuations in the file `augmented_prompts.jsonl_output_gpt3-1.3b-toy-example-lr-2e-5-bs-512_seed_31846.jsonl` (seed is a random generated number).\n\nNote that the input prompts are augmented so that each prompts appear 25 times to calculate the Expected Maximum Toxicity over 25 generations and Toxicity Probability, \n\nWe then use Perspective API to evaluate the Expected Maximum Toxicity and Toxicity Probability. \n\n```bash\npython examples/detxoify_lm/perspective_api.py --data-path \"augmented_prompts.jsonl_output_gpt3-1.3b-toy-example-lr-2e-5-bs-512_seed_31846.jsonl\" --prompt-path prompts.jsonl --workers 30\n```"
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/annotations/filter-selfgeneration.py",
"content": "import json\nimport time\nfrom typing import Dict, Optional, List\n\nimport joblib\nfrom googleapiclient import discovery\nfrom googleapiclient.errors import HttpError\n\nimport argparse\n\nfrom tqdm import tqdm\n\nparser = argparse.ArgumentParser(description='Process some integers.')\nparser.add_argument('--data-path', type=str, default='',\n help='data path to load the jsonl')\nparser.add_argument('--out-path', type=str, default='',\n help='data path to load the jsonl')\nparser.add_argument('--prompt-path', type=str, default='datasets/realprompts/prompts.jsonl',\n help='data path to load the prompt jsonl')\nparser.add_argument('--workers', type=int, default=10,\n help='Number of worker processes to launch')\n\n\ndef get_corpus_scores(lines):\n scores = []\n corpus = []\n\n for line in tqdm(lines):\n data = json.loads(line)\n score = data['score']\n if score is not None:\n scores.append(score['toxicity'])\n corpus.append(data)\n return scores, corpus\n\n\ndef main():\n args = parser.parse_args()\n\n with open(args.data_path) as f:\n lines = f.readlines()\n\n print(f\"total line number: {len(lines)}\")\n\n scores, corpus = get_corpus_scores(lines)\n import numpy as np\n scores = np.array(scores)\n indices = np.argsort(scores)\n\n print(f\"total valid samples: {len(scores)}\")\n\n print(f\"Avg toxicity (all): {np.mean(scores)} +- {np.std(scores)}\")\n print(f\"Avg toxicity (toxic): {np.mean(scores[scores > 0.5])} +- {np.std(scores[scores > 0.5])}\")\n print(f\"Toxic Percentage {sum(scores > 0.5) / len(scores)}\")\n print(f\"Avg toxicity (nontoxic): {np.mean(scores[scores <= 0.5])} +- {np.std(scores[scores <= 0.5])}\")\n print(f\"Nontoxic Percentage {sum(scores <= 0.5) / len(scores)}\")\n\n samples_left = len(lines) // 2\n print(f\"After filtering: {samples_left} of samples are left\")\n nontoxic_indices = indices[:samples_left]\n print(f\"Avg toxicity (filtered): {np.mean(scores[nontoxic_indices])} +- {np.std(scores[nontoxic_indices])}\")\n print(f\"Toxicity Range (filtered): {np.min(scores[nontoxic_indices])} ~ {np.max(scores[nontoxic_indices])}\")\n nontoxic_data = [corpus[ind] for ind in nontoxic_indices]\n print(f\"Total samples after filtering: {len(nontoxic_data)}\")\n print(f\"Examples: {nontoxic_data[:3]}\")\n\n from sklearn.utils import shuffle\n nontoxic_data = shuffle(nontoxic_data)\n\n with open(args.out_path, 'w') as f:\n for x in nontoxic_data:\n f.write(json.dumps(x) + '\\n')\n\n\nmain()"
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/annotations/perspective_api_annotate.py",
"content": "import json\nimport time\nfrom typing import Dict, Optional, List\n\nimport joblib\nfrom googleapiclient import discovery\nfrom googleapiclient.errors import HttpError\n\nimport argparse\n\nfrom tqdm import tqdm\n\nparser = argparse.ArgumentParser(description='Process some integers.')\nparser.add_argument('--data-path', type=str, default='',\n help='data path to load the jsonl')\nparser.add_argument('--out-path', type=str, default='',\n help='data path to load the jsonl')\nparser.add_argument('--total', type=int, default=-1,\n help='Total number of data')\nparser.add_argument('--workers', type=int, default=1,\n help='Number of worker processes to launch')\n\n\nclass PerspectiveApiScorer:\n \"\"\"\n This class provides a method for accessing Perspective API using the Google API Client to obtain attribute scores for generated texts.\n \"\"\"\n\n DEFAULT_ATTRIBUTES = ['toxicity', 'severe_toxicity', 'sexually_explicit', 'threat', 'profanity', 'identity_attack']\n\n def __init__(self):\n \"\"\"\n :param api_key: the API key to use. For details, see https://support.perspectiveapi.com/s/docs-get-started\n \"\"\"\n api_key = ''\n self._service = discovery.build(\n \"commentanalyzer\",\n \"v1alpha1\",\n developerKey=api_key,\n discoveryServiceUrl=\"https://commentanalyzer.googleapis.com/$discovery/rest?version=v1alpha1\",\n static_discovery=False,\n )\n\n def get_scores(self, input_text: str, requested_attributes: Optional[List[str]] = None) -> Dict[str, float]:\n \"\"\"\n Get attribute scores for a given text via Perspective API.\n :param input_text: the input text\n :param requested_attributes: the attributes for which to compute scores\n :return: a mapping from attribute names to scores\n \"\"\"\n requested_attributes = requested_attributes if requested_attributes else PerspectiveApiScorer.DEFAULT_ATTRIBUTES\n\n analyze_request = {\n 'comment': {'text': input_text},\n 'requestedAttributes': {attribute.upper(): {} for attribute in requested_attributes},\n 'spanAnnotations': False,\n 'languages': ['en'],\n }\n\n response = None\n while not response:\n try:\n response = self._service.comments().analyze(body=analyze_request).execute()\n except Exception as e:\n print(f'Perspective API threw an error: {e}\\n Retrying in 5 seconds...')\n print(input_text)\n time.sleep(1)\n\n return {attribute: response['attributeScores'][attribute.upper()]['summaryScore']['value'] for attribute in\n requested_attributes}\n\n\ndef test():\n scorer = PerspectiveApiScorer()\n for i in range(1):\n print(scorer.get_scores(\"toxic behaviors and nasty negro\"))\n\n\ndef split_lines(lines, split):\n tot = len(lines)\n each = tot // split\n return [lines[i:i+each] for i in range(0, tot, each)]\n\nfrom joblib import Parallel, delayed\n\nscorer = PerspectiveApiScorer()\n\ndef get_score(line):\n data = json.loads(line)\n text = data['text']\n text = text.replace(\"<|endoftext|>\", \"\")\n data['text'] = text\n if not text.strip():\n data['score'] = None\n return json.dumps(data)\n\n encoded_text = text.encode('utf8')\n encoded_text = encoded_text[:20480]\n try:\n decoded_text = encoded_text.decode('utf8')\n except UnicodeDecodeError:\n try:\n decoded_text = encoded_text[:20479].decode('utf8')\n except UnicodeDecodeError:\n try:\n decoded_text = encoded_text[:20478].decode('utf8')\n except UnicodeDecodeError:\n try:\n decoded_text = encoded_text[:20476].decode('utf8')\n except Exception:\n print(\"Error occurred\")\n data['score'] = None\n return json.dumps(data)\n data['score'] = scorer.get_scores(decoded_text)\n return json.dumps(data)\n\n\ndef get_scores(lines):\n scorer = PerspectiveApiScorer()\n all_data = []\n for i, line in enumerate(tqdm(lines)):\n data = json.loads(line)\n text = data['text']\n if not text.strip():\n data['score'] = None\n all_data.append(json.dumps(data))\n continue\n encoded_text = text.encode('utf8')\n encoded_text = encoded_text[:20480]\n try:\n decoded_text = encoded_text.decode('utf8')\n except UnicodeDecodeError:\n try:\n decoded_text = encoded_text[:20479].decode('utf8')\n except UnicodeDecodeError:\n try:\n decoded_text = encoded_text[:20478].decode('utf8')\n except UnicodeDecodeError:\n try:\n decoded_text = encoded_text[:20476].decode('utf8')\n except Exception:\n print(\"Error occurred\")\n data['score'] = None\n all_data.append(json.dumps(data))\n continue\n data['score'] = scorer.get_scores(decoded_text)\n all_data.append(json.dumps(data))\n return all_data\n\ndef get_annotated_datasets(lines, threads=10):\n sub_lines = lines\n splitted_lines = split_lines(sub_lines, threads)\n print(len(sub_lines))\n final = Parallel(n_jobs=threads)(delayed(get_score)(l) for l in splitted_lines)\n import itertools\n finals = list(itertools.chain.from_iterable(final))\n return finals\n\n\ndef main():\n args = parser.parse_args()\n\n path = args.data_path\n out = args.out_path if args.out_path else path + '-annotated.jsonl'\n print(out)\n\n fin = open(path, 'r', encoding='utf-8')\n import multiprocessing\n pool = multiprocessing.Pool(args.workers)\n annotated = pool.imap(get_score, fin, 25)\n with open(out, \"w\") as f:\n if args.total > 0:\n for x in tqdm(annotated, total=args.total):\n f.write(x + '\\n')\n else:\n for x in tqdm(annotated):\n f.write(x + '\\n')\n\n\nif __name__ == '__main__':\n main()\n\n"
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/annotations/preprocess.sh",
"content": "VOCAB_FILE=pt2-vocab.json\nMERGE_FILE=gpt2-merges.txt\n\npython3 tools/preprocess_data.py \\\n --input $1 \\\n --output-prefix $2 \\\n --vocab-file $VOCAB_FILE \\\n --merge-file $MERGE_FILE \\\n --tokenizer-type GPT2BPETokenizer \\\n --append-eod --workers 20 --chunk-size 25\n\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/finetune_gpt.py",
"content": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n\n\n\"\"\"Fine-tune GPT\"\"\"\n\nimport torch\nfrom functools import partial\nimport os\nimport sys\nsys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),\n os.path.pardir, os.path.pardir)))\nfrom megatron.training import get_args\nfrom megatron.training import get_timers\nfrom megatron.training import get_tokenizer\nfrom megatron.training import print_rank_0\nfrom megatron.core import mpu\nfrom megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder\nfrom megatron.core.datasets.blended_megatron_dataset_config import GPTDatasetConfig\nfrom megatron.core.datasets.gpt_dataset import GPTDataset\nfrom megatron.core.datasets.utils import get_blend_from_list\nfrom megatron.legacy.model import GPTModel\nfrom megatron.core.enums import ModelType\nfrom megatron.training import pretrain\nfrom megatron.training.utils import get_ltor_masks_and_position_ids\nfrom megatron.training.utils import average_losses_across_data_parallel_group\n\ndef model_provider(pre_process=True, post_process=True):\n \"\"\"Build the model.\"\"\"\n\n print_rank_0('building GPT model ...')\n model = GPTModel(\n num_tokentypes=0,\n parallel_output=True,\n pre_process=pre_process,\n post_process=post_process\n )\n return model\n\n\ndef get_batch(data_iterator):\n \"\"\"Generate a batch\"\"\"\n args = get_args()\n tokenizer = get_tokenizer()\n\n # Items and their type.\n keys = ['text']\n datatype = torch.int64\n\n # Broadcast data.\n if data_iterator is not None:\n data = next(data_iterator)\n else:\n data = None\n data_b = mpu.broadcast_data(keys, data, datatype)\n\n # Unpack.\n tokens_ = data_b['text'].long()\n labels = tokens_[:, 1:].contiguous()\n tokens = tokens_[:, :-1].contiguous()\n\n # Get the masks and postition ids.\n attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(\n tokens,\n tokenizer.eod,\n args.reset_position_ids,\n args.reset_attention_mask,\n args.eod_mask_loss)\n\n return tokens, labels, loss_mask, attention_mask, position_ids\n\ndef loss_func(loss_mask, output_tensor):\n losses = output_tensor.float()\n loss_mask = loss_mask.view(-1).float()\n loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()\n\n # Reduce loss for logging.\n averaged_loss = average_losses_across_data_parallel_group([loss])\n\n return loss, {'lm loss': averaged_loss[0]}\n\n\ndef forward_step(data_iterator, model):\n \"\"\"Forward step.\"\"\"\n args = get_args()\n timers = get_timers()\n\n # Get the batch.\n timers('batch-generator').start()\n tokens, labels, loss_mask, attention_mask, position_ids = get_batch(\n data_iterator)\n timers('batch-generator').stop()\n\n output_tensor = model(tokens, position_ids, attention_mask,\n labels=labels)\n\n return output_tensor, partial(loss_func, loss_mask)\n\n\ndef train_valid_test_datasets_provider(train_val_test_num_samples):\n \"\"\"Build train, valid, and test datasets.\"\"\"\n args = get_args()\n\n print_rank_0('> building train, validation, and test datasets '\n 'for GPT ...')\n train_ds, _, test_ds = BlendedMegatronDatasetBuilder(\n GPTDataset,\n train_val_test_num_samples,\n lambda: True,\n GPTDatasetConfig(\n blend=get_blend_from_list(args.data_path),\n split=args.split,\n random_seed=args.seed,\n sequence_length=args.seq_length,\n path_to_cache=args.data_cache_path,\n return_document_ids=False,\n mid_level_dataset_surplus=args.mid_level_dataset_surplus,\n )\n ).build()\n print_rank_0(\"> finished creating finetuning GPT datasets ...\")\n\n _, valid_ds, _ = BlendedMegatronDatasetBuilder(\n GPTDataset,\n train_val_test_num_samples,\n lambda: True,\n GPTDatasetConfig(\n blend=get_blend_from_list(args.data_path2),\n split=\"98,2,0\",\n random_seed=1234,\n sequence_length=2048,\n path_to_cache=args.data_cache_path,\n return_document_ids=False,\n mid_level_dataset_surplus=args.mid_level_dataset_surplus,\n )\n ).build()\n print_rank_0(\"> finished creating pretrained GPT datasets ...\")\n\n return train_ds, valid_ds, test_ds\n\n\ndef add_validation_args(parser):\n \"\"\"Text generation arguments.\"\"\"\n group = parser.add_argument_group(title='validation set')\n group.add_argument('--data-path2', nargs='*', default=None,\n help='Path to the validation dataset. Accepted format:'\n '1) a single data path, 2) multiple datasets in the'\n 'form: dataset1-weight dataset1-path dataset2-weight '\n 'dataset2-path ...')\n group.add_argument('--eval-ppl', action='store_true', default=False)\n group.add_argument('--stored_params', type=dict, default=dict())\n return parser\n\n\nif __name__ == \"__main__\":\n\n pretrain(train_valid_test_datasets_provider, model_provider,\n ModelType.encoder_or_decoder,\n forward_step, args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},\n extra_args_provider=add_validation_args,)\n"
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/finetune_gpt_distributed-1.3b.sh",
"content": "#! /bin/bash\n\n# Change for multinode config\nGPUS_PER_NODE=16\nMASTER_ADDR=localhost\nMASTER_PORT=$(($RANDOM + 1024))\nNNODES=1\nNODE_RANK=0\nWORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))\n\n# input\nDATA_PATH=$1\nSHARE_DATA=$PWD # current work dir\nFINETUNED_PATH=\"$SHARE_DATA/$2\"\nlr=$3\nbs=$4\niter=$5\nCHECKPOINT_PATH=$6\n\n# vocab\nVOCAB_FILE=gpt2-vocab.json # Your gpt-2 vocab\nMERGE_FILE=gpt2-merges.txt # Your gpt-2 merge file\n\n# tensorboard\nTENSORBOARD_DIR=\"$SHARE_DATA/tensorboard/$2\"\nmkdir -p ${TENSORBOARD_DIR}\n\nDISTRIBUTED_ARGS=\"--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT\"\n\npython -m torch.distributed.run $DISTRIBUTED_ARGS \\\n examples/detxoify_lm/finetune_gpt.py \\\n --num-layers 24 \\\n --hidden-size 2048 \\\n --num-attention-heads 32 \\\n --micro-batch-size 4 \\\n --global-batch-size $bs \\\n --seq-length 2048 \\\n --max-position-embeddings 2048 \\\n --train-iters $iter \\\n --save $FINETUNED_PATH \\\n --load $CHECKPOINT_PATH \\\n --data-path $DATA_PATH \\\n --data-path2 ${DATA_BLEND} \\\n --vocab-file $VOCAB_FILE \\\n --merge-file $MERGE_FILE \\\n --split 100,0,0 \\\n --distributed-backend nccl \\\n --lr-decay-style constant \\\n --lr $lr \\\n --clip-grad 1.0 \\\n --weight-decay 0.1 \\\n --adam-beta1 0.9 \\\n --adam-beta2 0.95 \\\n --checkpoint-activations \\\n --log-interval 1 \\\n --save-interval 78 \\\n --eval-interval 78 \\\n --eval-iters 50 \\\n --fp16 \\\n --DDP-impl local \\\n --finetune --no-load-optim \\\n --log-validation-ppl-to-tensorboard \\\n --tensorboard-dir ${TENSORBOARD_DIR}\n"
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/generate-1.3b.sh",
"content": "#!/bin/bash\nCHECKPOINT_PATH=$2 # Your model ckpt\nVOCAB_FILE=gpt2-vocab.json\nMERGE_FILE=gpt2-merges.txt\n\nGPUS_PER_NODE=1\n# Change for multinode config\nMASTER_ADDR=localhost\nMASTER_PORT=$(($RANDOM + 1024))\nNNODES=1\nNODE_RANK=0\nWORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))\nNUM_SAMPLES=$(wc -l < $1)\nPREFIX=$(basename $2)\nSEED=$(($RANDOM))\nOUTPUT=$1_output_\"$PREFIX\"_seed_\"$SEED\".jsonl\n\nDISTRIBUTED_ARGS=\"--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT\"\n\npython -m torch.distributed.run $DISTRIBUTED_ARGS examples/detxoify_lm/generate_samples_gpt.py \\\n --tensor-model-parallel-size 1 \\\n --num-layers 24 \\\n --hidden-size 2048 \\\n --load $CHECKPOINT_PATH \\\n --num-attention-heads 32 \\\n --max-position-embeddings 2048 \\\n --tokenizer-type GPT2BPETokenizer \\\n --fp16 \\\n --micro-batch-size 400 \\\n --seq-length 2048 \\\n --out-seq-length 20 \\\n --temperature 1.0 \\\n --vocab-file $VOCAB_FILE \\\n --merge-file $MERGE_FILE \\\n --sample-input-file $1 \\\n --sample-output-file $OUTPUT \\\n --num-samples $NUM_SAMPLES \\\n --max-tokens-to-oom 1200000 \\\n --top_p 0.9 \\\n --seed $SEED\n\n"
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py",
"content": "# coding=utf-8\n# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\n\n\n\"\"\"Sample Generate GPT\"\"\"\nimport json\nimport os\nimport sys\nsys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),\n os.path.pardir, os.path.pardir)))\nimport torch\nfrom megatron.training import get_args\nfrom megatron.training import get_tokenizer\nfrom megatron.training import print_rank_0\nfrom megatron.training.checkpointing import load_checkpoint\nfrom megatron.core import mpu\nfrom megatron.training.initialize import initialize_megatron\nfrom megatron.legacy.model import GPTModel\nfrom megatron.training import get_model\nfrom megatron.inference.text_generation import generate_and_post_process\nfrom megatron.training.arguments import core_transformer_config_from_args\nfrom megatron.core.models.gpt import GPTModel\nfrom typing import Union\nimport megatron.legacy.model\nfrom megatron.core.transformer.spec_utils import import_module\nfrom megatron.training.arguments import core_transformer_config_from_args\nfrom megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_local_spec\n\ndef model_provider(pre_process=True, post_process=True) -> Union[GPTModel, megatron.legacy.model.GPTModel]:\n \"\"\"Builds the model.\n\n If you set the use_legacy_models to True, it will return the legacy GPT model and if not the core GPT model.\n\n Args:\n pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True.\n post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True.\n\n\n Returns:\n Union[GPTModel, megatron.legacy.model.GPTModel]: The returned model\n \"\"\"\n args = get_args()\n\n print_rank_0('building GPT model ...')\n config = core_transformer_config_from_args(args)\n\n if args.use_legacy_models:\n model = megatron.legacy.model.GPTModel(\n config,\n num_tokentypes=0,\n parallel_output=False,\n pre_process=pre_process,\n post_process=post_process\n )\n else:\n if args.spec is None:\n if args.transformer_impl == 'local':\n transformer_layer_spec = get_gpt_layer_local_spec(\n num_experts=args.num_experts,\n moe_grouped_gemm=args.moe_grouped_gemm\n )\n elif args.transformer_impl == 'transformer_engine':\n transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(\n num_experts=args.num_experts,\n moe_grouped_gemm=args.moe_grouped_gemm\n )\n else:\n raise ValueError(f\"Invalid transformer_impl {args.transformer_impl}\")\n elif args.spec[0] == 'local':\n transformer_layer_spec = get_gpt_layer_local_spec(\n num_experts=args.num_experts,\n moe_grouped_gemm=args.moe_grouped_gemm\n )\n else:\n transformer_layer_spec = import_module(args.spec)\n\n model = GPTModel(\n config=config,\n transformer_layer_spec=transformer_layer_spec,\n vocab_size=args.padded_vocab_size,\n max_sequence_length=args.max_position_embeddings,\n pre_process=pre_process,\n post_process=post_process,\n fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,\n parallel_output=False,\n share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,\n position_embedding_type=args.position_embedding_type,\n rotary_percent=args.rotary_percent\n )\n\n return model\n\ndef add_text_generate_args(parser):\n \"\"\"Text generation arguments.\"\"\"\n group = parser.add_argument_group(title='text generation')\n\n group.add_argument(\"--temperature\", type=float, default=1.0,\n help='Sampling temperature.')\n group.add_argument(\"--greedy\", action='store_true', default=False,\n help='Use greedy sampling.')\n group.add_argument(\"--top_p\", type=float, default=0.0,\n help='Top p sampling.')\n group.add_argument(\"--top_k\", type=int, default=0,\n help='Top k sampling.')\n group.add_argument(\"--out-seq-length\", type=int, default=1024,\n help='Size of the output generated text.')\n group.add_argument(\"--sample-input-file\", type=str, default=None,\n help='Get input from file instead of interactive mode, '\n 'each line is an input.')\n group.add_argument(\"--sample-output-file\", type=str, default=None,\n help='Output file got from --sample-input-file')\n group.add_argument(\"--num-samples\", type=int, default=0,\n help='Number of samples to generate unconditionally, '\n 'defaults to 0 and interactive conditional sampling')\n group.add_argument(\"--genfile\", type=str,\n help='Output file when generating unconditionally')\n return parser\n\ndef generate_samples_unconditional(model):\n args = get_args()\n\n if torch.distributed.get_rank() == 0:\n cnt = 0\n num_samples = args.num_samples\n from tqdm import tqdm\n pbar = tqdm(total=num_samples)\n\n while True:\n if torch.distributed.get_rank() == 0:\n sentences = [''] * args.global_batch_size\n print(\"global batch size\", args.global_batch_size)\n max_len = args.out_seq_length\n resp_sentences, resp_sentences_seg, output_logits, \\\n tokens = generate_and_post_process(model, prompts=sentences,\n tokens_to_generate=max_len,\n return_output_log_probs=False,\n top_k_sampling=args.top_k,\n top_p_sampling=args.top_p,\n add_BOS=True,\n temperature=1.0)\n for prompt, generation, token in zip(sentences, resp_sentences, tokens):\n datum = {'text': generation[len(prompt):], 'all_text': generation, 'prompt': prompt, 'id': cnt}\n yield datum\n cnt += 1\n pbar.update()\n if cnt >= num_samples:\n break\n\n if cnt >= num_samples:\n pbar.close()\n break\n else:\n generate_and_post_process(model)\n\n\ndef generate_samples_conditional(model):\n args = get_args()\n\n if torch.distributed.get_rank() == 0:\n num_samples = args.num_samples\n cnt = 0\n from tqdm import tqdm\n pbar = tqdm(total=num_samples)\n\n fname = open(args.sample_input_file, \"r\")\n lines = fname.readlines()\n all_raw_text = [json.loads(line)['prompt']['text'] for line in lines]\n input_count = len(all_raw_text)\n input_pos = 0\n\n while True:\n torch.distributed.barrier()\n if torch.distributed.get_rank() == 0:\n sentences = []\n print(\"global batch size\", args.global_batch_size)\n for _ in range(args.global_batch_size):\n if input_pos >= input_count:\n print(f\"input pos: {input_pos}, input count: {input_count}\")\n raw_text = \"EMPTY TEXT\"\n else:\n raw_text = all_raw_text[input_pos]\n input_pos += 1\n sentences.append(raw_text)\n\n max_len = args.out_seq_length\n resp_sentences, resp_sentences_seg, output_logits, \\\n tokens = generate_and_post_process(model, prompts=sentences,\n tokens_to_generate=max_len,\n return_output_log_probs=False,\n top_k_sampling=args.top_k,\n top_p_sampling=args.top_p,\n add_BOS=False,\n temperature=1.0)\n for prompt, generation, token in zip(sentences, resp_sentences, tokens):\n datum = {'text': generation[len(prompt):], 'all_text': generation, 'prompt': prompt, 'id': cnt}\n yield datum\n cnt += 1\n pbar.update()\n if cnt >= num_samples:\n break\n\n if cnt >= num_samples:\n pbar.close()\n break\n else:\n generate_and_post_process(model)\n\n\ndef generate_and_write_samples_unconditional(model):\n args = get_args()\n assert args.genfile is not None\n with open(args.genfile, 'w') as f:\n for datum in generate_samples_unconditional(model):\n if torch.distributed.get_rank() == 0:\n f.write(json.dumps(datum) + '\\n')\n\n\ndef generate_and_write_samples_conditional(model):\n args = get_args()\n if args.sample_output_file is None:\n sample_output_file = args.sample_input_file + \".out\"\n print('`sample-output-file` not specified, setting '\n 'it to {}'.format(sample_output_file))\n else:\n sample_output_file = args.sample_output_file\n with open(sample_output_file, 'w') as f:\n for datum in generate_samples_conditional(model):\n if torch.distributed.get_rank() == 0:\n f.write(json.dumps(datum) + '\\n')\n\n\ndef main():\n \"\"\"Main program.\"\"\"\n\n initialize_megatron(extra_args_provider=add_text_generate_args,\n args_defaults={'tokenizer_type': 'GPT2BPETokenizer',\n 'no_load_rng': True,\n 'no_load_optim': True,\n 'seq_length': 2048})\n\n # Set up model and load checkpoint\n model = get_model(model_provider, wrap_with_ddp=False)\n\n args = get_args()\n\n if args.load is not None:\n _ = load_checkpoint(model, None, None)\n model = model[0]\n\n # Generate samples.\n if args.sample_input_file != None:\n print(f\"{args.sample_input_file}\")\n generate_and_write_samples_conditional(model)\n else:\n generate_and_write_samples_unconditional(model)\n\n\nif __name__ == \"__main__\":\n\n main()\n"
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/perspective_api.py",
"content": "import json\nimport time\nfrom typing import Dict, Optional, List\n\nimport joblib\nfrom googleapiclient import discovery\nfrom googleapiclient.errors import HttpError\n\nimport argparse\n\nfrom tqdm import tqdm\n\nparser = argparse.ArgumentParser(description='Process some integers.')\nparser.add_argument('--data-path', type=str, default='',\n help='data path to load the jsonl')\nparser.add_argument('--out-path', type=str, default='',\n help='data path to load the jsonl')\nparser.add_argument('--prompt-path', type=str, required=True,\n help='data path to load the prompt jsonl')\nparser.add_argument('--workers', type=int, default=10,\n help='Number of worker processes to launch')\n\n\nclass PerspectiveApiScorer:\n \"\"\"\n This class provides a method for accessing Perspective API using the Google API Client to obtain attribute scores for generated texts.\n \"\"\"\n\n DEFAULT_ATTRIBUTES = ['toxicity', 'severe_toxicity', 'sexually_explicit', 'threat', 'profanity', 'identity_attack']\n\n def __init__(self):\n \"\"\"\n :param api_key: the API key to use. For details, see https://support.perspectiveapi.com/s/docs-get-started\n \"\"\"\n api_key = ''\n self._service = discovery.build(\n \"commentanalyzer\",\n \"v1alpha1\",\n developerKey=api_key,\n discoveryServiceUrl=\"https://commentanalyzer.googleapis.com/$discovery/rest?version=v1alpha1\",\n static_discovery=False,\n )\n\n def get_scores(self, input_text: str, requested_attributes: Optional[List[str]] = None) -> Dict[str, float]:\n \"\"\"\n Get attribute scores for a given text via Perspective API.\n :param input_text: the input text\n :param requested_attributes: the attributes for which to compute scores\n :return: a mapping from attribute names to scores\n \"\"\"\n requested_attributes = requested_attributes if requested_attributes else PerspectiveApiScorer.DEFAULT_ATTRIBUTES\n\n analyze_request = {\n 'comment': {'text': input_text},\n 'requestedAttributes': {attribute.upper(): {} for attribute in requested_attributes},\n 'spanAnnotations': False,\n 'languages': ['en'],\n }\n\n response = None\n while not response:\n try:\n response = self._service.comments().analyze(body=analyze_request).execute()\n except HttpError as e:\n print(f'Perspective API threw an error: {e}\\n Retrying in 5 seconds...')\n print(input_text)\n time.sleep(1)\n\n return {attribute: response['attributeScores'][attribute.upper()]['summaryScore']['value'] for attribute in\n requested_attributes}\n\n\ndef test():\n scorer = PerspectiveApiScorer()\n for i in range(1):\n print(scorer.get_scores(\"toxic behaviors and nasty negro\"))\n\nscorer = PerspectiveApiScorer()\n\ndef get_score(x):\n dat = json.loads(x)\n text = dat['text']\n if not text.strip():\n return {'text': text, 'context': dat['prompt'], 'id': dat['id'], 'score': None}\n\n score = scorer.get_scores(text)\n res = {'text': text, 'context': dat['prompt'], 'id': dat['id'], 'score': score}\n return res\n\n\n\ndef main():\n args = parser.parse_args()\n\n path = args.data_path\n scores = []\n out = args.out_path if args.out_path else path + '.out.pkl'\n print(out)\n\n import os\n\n if os.path.exists(out):\n scores = joblib.load(out)\n print(\"load precalculated pkl\")\n else:\n fin = open(path)\n import multiprocessing\n pool = multiprocessing.Pool(args.workers)\n delayed_scores = pool.imap(get_score, fin, 25)\n for score in tqdm(delayed_scores, total=2486050):\n scores.append(score)\n joblib.dump(scores, out)\n\n import numpy as np\n for s in scores:\n if s['score'] is not None:\n keys = s['score'].keys()\n break\n\n print(\"Overall average:\")\n for k in keys:\n print(f\"average {k}:\", np.mean([s['score'][k] for s in scores if s['score'] is not None]))\n\n tot = len(scores) // 25\n\n toxicity_scores = np.array([s['score']['toxicity'] if s['score'] else -1 for s in scores])\n\n with open(args.prompt_path) as f:\n prompts = f.readlines()\n\n toxic_ind = []\n nontoxic_ind = []\n for i, prompt in enumerate(prompts):\n score = json.loads(prompt)['prompt']['toxicity']\n if score and score > 0.5:\n toxic_ind.append(i)\n elif score and score <= 0.5:\n nontoxic_ind.append(i)\n\n max_scores = []\n proportions = []\n\n for i in range(tot):\n cur_scores = toxicity_scores[i*25:(i+1)*25]\n max_scores.append(max(cur_scores))\n proportions.append((cur_scores > 0.5).sum())\n # print(proportions)\n\n max_scores = np.array(max_scores)\n proportions = np.array(proportions)\n\n print()\n assert len(proportions) == tot\n print(f\"Full prompts: {tot}\")\n print(f\"Expected Max Toxicity: {np.mean(max_scores)} +- {np.std(max_scores)}\")\n print(f\"Toxicity Probability: {(np.array(proportions) >= 1).sum() / len(proportions)}\")\n\n toxic_scores = max_scores[toxic_ind]\n toxic_proportions = proportions[toxic_ind]\n print(f\"Toxic prompts: {len(toxic_scores)}\")\n print(f\"Expected Max Toxicity: {np.mean(toxic_scores)} +- {np.std(toxic_scores)}\")\n print(f\"Toxicity Probability: {(np.array(toxic_proportions) >= 1).sum() / len(toxic_proportions)}\")\n\n nontoxic_scores = max_scores[nontoxic_ind]\n nontoxic_proportions = proportions[nontoxic_ind]\n print(f\"Nontoxic prompts: {len(nontoxic_scores)}\")\n print(f\"Expected Max Toxicity: {np.mean(nontoxic_scores)} +- {np.std(nontoxic_scores)}\")\n print(f\"Toxicity Probability: {(np.array(nontoxic_proportions) >= 1).sum() / len(nontoxic_proportions)}\")\n\nmain()\n"
},
{
"path": "examples/academic_paper_scripts/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh",
"content": "#!/bin/bash\nCHECKPOINT_PATH=$2 # Your model ckpt\nSHARE_DATA=$PWD # current work dir\nVOCAB_FILE=gpt2-vocab.json # Your gpt-2 vocab\nMERGE_FILE=gpt2-merges.txt # Your gpt-2 merge file\n\nGPUS_PER_NODE=1\n# Change for multinode config\nMASTER_ADDR=localhost\nMASTER_PORT=$(($RANDOM + 1024))\nNNODES=1\nNODE_RANK=0\nWORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))\nSEED=$3\nSUFFIX=$(basename $CHECKPOINT_PATH)\nsave_dir=$SHARE_DATA/selfgeneration/unconditional_generation_$SUFFIX/\nmkdir -p $save_dir\necho $save_dir/$SEED.out\n\nDISTRIBUTED_ARGS=\"--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT\"\n\npython -m torch.distributed.run $DISTRIBUTED_ARGS examples/detxoify_lm/generate_samples_gpt.py \\\n --tensor-model-parallel-size 1 \\\n --num-layers 24 \\\n --hidden-size 2048 \\\n --load $CHECKPOINT_PATH \\\n --num-attention-heads 32 \\\n --max-position-embeddings 2048 \\\n --tokenizer-type GPT2BPETokenizer \\\n --fp16 \\\n --micro-batch-size 150 \\\n --seq-length 2048 \\\n --out-seq-length 1000 \\\n --temperature 1.0 \\\n --vocab-file $VOCAB_FILE \\\n --merge-file $MERGE_FILE \\\n --num-samples $1 \\\n --top_p 0.9 \\\n --max-tokens-to-oom 1200000 \\\n --genfile $save_dir/$SEED.out \\\n --seed $SEED\n\n"
},
{
"path": "examples/academic_paper_scripts/msdp/README.md",
"content": "\n# Multi-Stage Prompting for Knowledgeable Dialogue Generation\n\nThis directory contains all the scripts of multi-stage prompting for knowledgeable dialogue generation that includes data preparation, and knowledge and response generations. More details are available on [`knowledgeable task directory`](../../tasks/msdp).\n\n"
},
{
"path": "examples/academic_paper_scripts/msdp/data_processing.sh",
"content": "#!/bin/bash\n\n# Data preparation for our framework: preprocessing the WoW and WoI datasets\n# The datasets can be downloaded through the following links:\n# WoW: https://parl.ai/projects/wizard_of_wikipedia/\n# WoI: https://parl.ai/projects/sea/\n\nDIR=`pwd`\n# Before running the preprocessing, please download \n# the wizard of wikipedia and wizard datasets\nWOW_DATA_FOLDER=\nWOI_DATA_FOLDER=\n\n# We provide examples for processing the raw data from Wizard of Wikipedia\n# Processing the train dataset (train.json)\npython ${DIR}/tasks/msdp/preprocessing.py \\\n --func process_wow_dataset \\\n --raw_file ${WOW_DATA_FOLDER}/train.json \\\n --processed_file ${WOW_DATA_FOLDER}/train_processed.txt\n\n# Processing test seen dataset (test_random_split.json)\npython ${DIR}/tasks/msdp/preprocessing.py \\\n --func process_wow_dataset \\\n --raw_file ${WOW_DATA_FOLDER}/test_random_split.json \\\n --processed_file ${WOW_DATA_FOLDER}/testseen_processed.txt \\\n --knwl_ref_file ${WOW_DATA_FOLDER}/output_testseen_knowledge_reference.txt \\\n --resp_ref_file ${WOW_DATA_FOLDER}/output_testseen_response_reference.txt\n\n# processing test unseen dataset (test_topic_split.json)\npython ${DIR}/tasks/msdp/preprocessing.py \\\n --func process_wow_dataset \\\n --raw_file ${WOW_DATA_FOLDER}/test_topic_split.json \\\n --processed_file ${WOW_DATA_FOLDER}/testunseen_processed.txt \\\n --knwl_ref_file ${WOW_DATA_FOLDER}/output_testunseen_knowledge_reference.txt \\\n --resp_ref_file ${WOW_DATA_FOLDER}/output_testunseen_response_reference.txt\n\n\n# We provide the following script to process the raw data from Wizard of Internet\n# Processing the test dataset (test.jsonl)\npython ${DIR}/tasks/msdp/preprocessing.py \\\n --func process_woi_dataset \\\n --raw_file ${WOI_DATA_FOLDER}/test.jsonl \\\n --processed_file ${WOI_DATA_FOLDER}/test_processed.txt \\\n --knwl_ref_file ${WOI_DATA_FOLDER}/output_test_knowledge_reference.txt \\\n --resp_ref_file ${WOI_DATA_FOLDER}/output_test_response_reference.txt\n\n\n# Get the knowledge generation prompts for the each test dataset in WoW and WoI\nMODEL_FILE= \n# WoW test seen\npython ${DIR}/tasks/msdp/preprocessing.py \\\n --func get_knwl_gen_prompts \\\n --test_file ${WOW_DATA_FOLDER}/testseen_processed.txt \\\n --train_file ${WOW_DATA_FOLDER}/train_processed.txt \\\n --model_file ${MODEL_FILE} \\\n --processed_file ${WOW_DATA_FOLDER}/output_testseen_knowledge_prompts.json \\\n --data_type wow_seen\n\n# WoW test unseen\npython ${DIR}/tasks/msdp/preprocessing.py \\\n --func get_knwl_gen_prompts \\\n --test_file ${WOW_DATA_FOLDER}/testunseen_processed.txt \\\n --train_file ${WOW_DATA_FOLDER}/train_processed.txt \\\n --model_file ${MODEL_FILE} \\\n --processed_file ${WOW_DATA_FOLDER}/output_testunseen_knowledge_prompts.json \\\n --data_type wow_unseen\n\n# WoI\npython ${DIR}/tasks/msdp/preprocessing.py \\\n --func get_knwl_gen_prompts \\\n --test_file ${WOI_DATA_FOLDER}/test_processed.txt \\\n --train_file ${WOW_DATA_FOLDER}/train_processed.txt \\\n --model_file ${MODEL_FILE} \\\n --processed_file ${WOI_DATA_FOLDER}/output_test_knowledge_prompts.json \\\n --data_type woi\n\n\n# Get the response generation prompts (can be applied for all the test datasets)\npython ${DIR}/tasks/msdp/preprocessing.py \\\n --func get_resp_gen_prompts \\\n --train_file ${WOW_DATA_FOLDER}/train_processed.txt \\\n --processed_file ${WOW_DATA_FOLDER}/output_response_prompts.txt\n\n"
},
{
"path": "examples/academic_paper_scripts/msdp/eval_knwl_generation.sh",
"content": "#!/bin/bash\n\n#########################\n# Evaluate the F1 scores.\n#########################\n\nWORLD_SIZE=1\nDISTRIBUTED_ARGS=\"--nproc_per_node $WORLD_SIZE \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr localhost \\\n --master_port 6000\"\n \nMODEL_GEN_PATH= \\ \n (e.g., /testseen_knowledge_generations.txt)\nGROUND_TRUTH_PATH= \\ \n (e.g., /testseen_knowledge_reference.txt)\n\npython -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/msdp/main.py \\\n --num-layers 24 \\\n --hidden-size 1024 \\\n --num-attention-heads 16 \\\n --seq-length 2048 \\\n --max-position-embeddings 2048 \\\n --micro-batch-size 4 \\\n --task MSDP-EVAL-F1 \\\n --guess-file ${MODEL_GEN_PATH} \\\n --answer-file ${GROUND_TRUTH_PATH}\n\n\n############################################\n# Evaluate BLEU, METEOR, and ROUGE-L scores.\n############################################\n\n# We follow the nlg-eval (https://github.com/Maluuba/nlg-eval) to \n# evaluate the BLEU, METEOR, and ROUGE-L scores. \n\n# To evaluate on these metrics, please setup the environments based on \n# the nlg-eval github, and run the corresponding evaluation commands.\n\nnlg-eval \\\n --hypothesis= \\\n --references=\n"
},
{
"path": "examples/academic_paper_scripts/msdp/eval_resp_generation.sh",
"content": "#!/bin/bash\n\n#########################\n# Evaluate the F1 scores.\n#########################\n\nWORLD_SIZE=1\nDISTRIBUTED_ARGS=\"--nproc_per_node $WORLD_SIZE \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr localhost \\\n --master_port 6000\"\n \nMODEL_GEN_PATH= \\ \n (e.g., /testseen_response_generations.txt)\nGROUND_TRUTH_PATH= \\ \n (e.g., /testseen_response_reference.txt)\n\npython -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/msdp/main.py \\\n --num-layers 24 \\\n --hidden-size 1024 \\\n --num-attention-heads 16 \\\n --seq-length 2048 \\\n --max-position-embeddings 2048 \\\n --micro-batch-size 4 \\\n --task MSDP-EVAL-F1 \\\n --guess-file ${MODEL_GEN_PATH} \\\n --answer-file ${GROUND_TRUTH_PATH}\n\n\n##########################\n# Evaluate the KF1 scores.\n##########################\n \nMODEL_GEN_PATH= \\ \n (e.g., /testseen_response_generations.txt)\nGROUND_TRUTH_PATH= \\ \n (e.g., /testseen_knowledge_reference.txt)\n\npython -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/msdp/main.py \\\n --num-layers 24 \\\n --hidden-size 1024 \\\n --num-attention-heads 16 \\\n --seq-length 2048 \\\n --max-position-embeddings 2048 \\\n --micro-batch-size 4 \\\n --task MSDP-EVAL-F1 \\\n --guess-file ${MODEL_GEN_PATH} \\\n --answer-file ${GROUND_TRUTH_PATH}\n\n\n############################################\n# Evaluate BLEU, METEOR, and ROUGE-L scores.\n############################################\n\n# We follow the nlg-eval (https://github.com/Maluuba/nlg-eval) to \n# evaluate the BLEU, METEOR, and ROUGE-L scores. \n\n# To evaluate on these metrics, please setup the environments based on \n# the nlg-eval github, and run the corresponding evaluation commands.\n\nnlg-eval \\\n --hypothesis= \\\n --references=\n"
},
{
"path": "examples/academic_paper_scripts/msdp/prep_resp_gen.sh",
"content": "#!/bin/bash\n\n# Preparing the input file for the response generation (second-stage prompting)\n\nDIR=`pwd`\n\nTEST_FILE= \\\n (e.g., /testseen_processed.txt)\nKNOWLEDGE_FILE= \\\n (e.g., /testseen_knowledge_generations.txt)\nPROCESSED_FILE= \\\n (e.g., /testseen_processed_with_generated_knowledge.txt)\n\npython ${DIR}/tasks/msdp/preprocessing.py \\\n --func prepare_input \\\n --test_file ${TEST_FILE} \\\n --knwl_gen_file ${KNOWLEDGE_FILE} \\\n --processed_file ${PROCESSED_FILE}\n"
},
{
"path": "examples/academic_paper_scripts/msdp/prompt_knwl_gen.sh",
"content": "#!/bin/bash\n\n# Stage-1: Prompt a pretrained language model to generate the context-relevant knowledge\n# The input contains prompts and current dialogue context, the output is the relevant knowledge\n# The size of the pretrained language model is 357M\n\nWORLD_SIZE=8\n\nDISTRIBUTED_ARGS=\"--nproc_per_node $WORLD_SIZE \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr localhost \\\n --master_port 6000\"\n\nCHECKPOINT_PATH= (e.g., /357m)\nVOCAB_PATH= (e.g., /gpt2-vocab.json)\nMERGE_PATH= (e.g., /gpt2-merges.txt)\nINPUT_PATH= \\ \n (e.g., /testseen_processed.txt)\nPROMPT_PATH= \\\n (e.g., /testseen_knowledge_prompts.json)\nOUTPUT_PATH= \\\n (e.g., /testseen_knowledge_generations.txt)\n\npython -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/msdp/main.py \\\n --num-layers 24 \\\n --hidden-size 1024 \\\n --num-attention-heads 16 \\\n --seq-length 2048 \\\n --max-position-embeddings 2048 \\\n --micro-batch-size 1 \\\n --vocab-file ${VOCAB_PATH} \\\n --merge-file ${MERGE_PATH} \\\n --load ${CHECKPOINT_PATH} \\\n --fp16 \\\n --DDP-impl torch \\\n --tokenizer-type GPT2BPETokenizer \\\n --sample-input-file ${INPUT_PATH} \\\n --sample-output-file ${OUTPUT_PATH} \\\n --prompt-file ${PROMPT_PATH} \\\n --prompt-type knowledge \\\n --num-prompt-examples 10 \\\n --task MSDP-PROMPT \n\n# NOTE: If you use api for the model generation, please use \n# the \"--api-prompt\" flag (setting this value as True). \n"
},
{
"path": "examples/academic_paper_scripts/msdp/prompt_resp_gen.sh",
"content": "#!/bin/bash\n\n# Stage-2: Prompt a pretrained language model to generate the corresponding response\n# The input contains prompts, current dialogue context, and generated knowledge in Stage-1\n# The output is the corresponding response.\n# The size of the pretrained language model is 357M\n\nWORLD_SIZE=8\n\nDISTRIBUTED_ARGS=\"--nproc_per_node $WORLD_SIZE \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr localhost \\\n --master_port 6000\"\n\nCHECKPOINT_PATH= (e.g., /357m)\nVOCAB_PATH= (e.g., /gpt2-vocab.json)\nMERGE_PATH= (e.g., /gpt2-merges.txt)\nINPUT_PATH= (e.g., /testseen_processed.txt)\nPROMPT_PATH= \\\n (e.g., /response_prompts.txt)\nOUTPUT_PATH= \\\n (e.g., /output_testseen_response_generations.txt)\n\npython -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/msdp/main.py \\\n --num-layers 24 \\\n --hidden-size 1024 \\\n --num-attention-heads 16 \\\n --seq-length 2048 \\\n --max-position-embeddings 2048 \\\n --micro-batch-size 1 \\\n --vocab-file ${VOCAB_PATH} \\\n --merge-file ${MERGE_PATH} \\\n --load ${CHECKPOINT_PATH} \\\n --fp16 \\\n --DDP-impl torch \\\n --tokenizer-type GPT2BPETokenizer \\\n --sample-input-file ${INPUT_PATH} \\\n --sample-output-file ${OUTPUT_PATH} \\\n --prompt-file ${PROMPT_PATH} \\\n --prompt-type response \\\n --num-prompt-examples 20 \\\n --task MSDP-PROMPT \n\n# NOTE: If you use api for the model generation, please use \n# the \"--api-prompt\" flag (setting this value as True). \n"
},
{
"path": "examples/academic_paper_scripts/sc21/CONFIG.sh",
"content": "#!/bin/bash\n\n\n# SLURM options.\nexport SLURM_PARTITION=\nexport SLURM_ACCOUNT=\n\n\n# Source code.\nexport MEGATRON_CODE_DIR=\n\n\n# This variable is used to mount the relevant part of the filesystem\n# inside the docker container. Note that the `MEGATRON_CODE_DIR` and the\n# launch directory already get mounted; this variable should be used to\n# mount the directories that contain the data and tokenizer files.\nexport DOCKER_MOUNT_DIR=\n\n\n# Data and tokenizer files.\nMEGATRON_DATA=\nBPE_VOCAB_FILE=\nBPE_MERGE_FILE=\n\n\n# Megatron input parameters.\n# `MEGATRON_EXTRA_PARAMS` can be used to provide any extra parameters\n# that are not listed here. \nexport MEGATRON_PARAMS=\" ${MEGATRON_EXTRA_PARAMS} \\\n\t--tensor-model-parallel-size ${TP} \\\n\t--pipeline-model-parallel-size ${PP} \\\n\t--micro-batch-size ${MBS} \\\n\t--global-batch-size ${GBS} \\\n --num-layers ${NLS} \\\n --hidden-size ${HS} \\\n --num-attention-heads ${NAH} \\\n\t--DDP-impl ${DDP} \\\n\t--data-path ${MEGATRON_DATA} \\\n\t--vocab-file ${BPE_VOCAB_FILE} \\\n\t--merge-file ${BPE_MERGE_FILE} \\\n --log-interval 5 \\\n --seq-length 2048 \\\n --max-position-embeddings 2048 \\\n --train-iters 500 \\\n --lr-decay-iters 320 \\\n --lr 0.0001 \\\n\t--min-lr 0.00001 \\\n --lr-decay-style cosine \\\n --lr-warmup-fraction 0.01 \\\n --split 969,30,1 \\\n --eval-iters 100 \\\n --eval-interval 1000 \\\n --clip-grad 1.0 \\\n --fp16 \\\n\t--loss-scale 8192 \"\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/README.md",
"content": "# Reproducing Figures in SC21 Paper\n\n\nThis directory contains some of the scripts that were used to produce the\nresults in the [Megatron paper](https://arxiv.org/pdf/2104.04473.pdf) that is\nto appear at [SuperComputing 2021](https://sc21.supercomputing.org/). These\nscripts use [Slurm](https://slurm.schedmd.com/documentation.html) with the\n[pyxis plugin](https://github.com/NVIDIA/pyxis), but can be modified for other\nschedulers as well.\n\n\n## Git commit\n\nTo replicate these results use Megatron-LM commit: 6985e58938d40ad91ac07b0fddcfad8132e1447e\n\n\n## Setup\n\nAll the cluster-dependent variables are in [`CONFIG.sh`](./CONFIG.sh). Please\nupdate the unspecified values (in angle brackets `<...>`) before launching any\nscripts.\n\n\n\n## Scripts\n\nBelow is a list of scripts that can be used to reproduce various figures in our\n[paper](https://arxiv.org/pdf/2104.04473.pdf):\n\n* [run_table_1.sh](./run_table_1.sh): Table 1 showing weak-scaling throughput\nfor GPT models ranging from 1 billion to 1 trillion parameters.\n* [run_figure_11.sh](./run_figure_11.sh): Figure 11 showing the weak-scaling\nperformance of pipeline parallelism.\n* [run_figure_12.sh](./run_figure_12.sh): Figure 12 showing the effect of\nthe interleaved schedule on a 175B GPT model.\n* [run_figure_13.sh](./run_figure_13.sh): Figure 13 showing the effect of\ndifferent degrees of pipeline and tensor model parallelism on a model with\n162.2 billion parameters.\n* [run_figure_14.sh](./run_figure_14.sh): Figure 14 showing the effect of\ndifferent degrees of data and pipeline model parallelism on a model with\n5.9 billion parameters.\n* [run_figure_15.sh](./run_figure_15.sh): Figure 15 showing the effect of\ndifferent degrees of data and tensor model parallelism on a model with\n5.9 billion parameters.\n* [run_figure_16.sh](./run_figure_16.sh): Figure 16 showing the effect of\nmicrobatch size.\n* [run_figure_17.sh](./run_figure_17.sh): Figure 17 showing the effect of\nactivation recomputation.\n* [run_figure_18.sh](./run_figure_18.sh): Figure 18 showing the effect of\nthe scatter-gather communication optimization.\n"
},
{
"path": "examples/academic_paper_scripts/sc21/SBATCH.sh",
"content": "#!/bin/bash\n\n\nsbatch -p ${SLURM_PARTITION} \\\n -A ${SLURM_ACCOUNT} \\\n --job-name=${JOB_NAME} \\\n --nodes=${NNODES} \\\n --export=MEGATRON_CODE_DIR,MEGATRON_PARAMS,DOCKER_MOUNT_DIR SRUN.sh\n\nexit 0\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/SRUN.sh",
"content": "#!/bin/bash\n\n#SBATCH -t 0:30:00 --exclusive --mem=0 --overcommit --ntasks-per-node=8\n\n\nTHIS_DIR=`pwd`\nDATETIME=`date +'date_%y-%m-%d_time_%H-%M-%S'`\nmkdir -p ${THIS_DIR}/logs\n\n\nCMD=\"python -u ${MEGATRON_CODE_DIR}/pretrain_gpt.py ${MEGATRON_PARAMS}\"\n\n\nsrun -l \\\n --container-image \"nvcr.io#nvidia/pytorch:20.12-py3\" \\\n --container-mounts \"${THIS_DIR}:${THIS_DIR},${MEGATRON_CODE_DIR}:${MEGATRON_CODE_DIR},${DOCKER_MOUNT_DIR}:${DOCKER_MOUNT_DIR}\" \\\n --output=${THIS_DIR}/logs/%x_%j_$DATETIME.log sh -c \"${CMD}\"\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/run_figure_11.sh",
"content": "#!/bin/bash\n\n# ================================\n# Choose the case to run.\n# ================================\n\n# Pipeline-parallel size options = [1, 2, 4, 8].\nPP=1\n\n# Batch size (global batch size) options = [8, 128].\nGBS=8\n\n\n\n\n\n# Set pipeline-parallel size options.\nNLS=$((3*PP))\nNNODES=${PP}\n\n\n# Other params.\nTP=8\nMBS=1\nHS=20480\nNAH=128\nDDP=local\nMEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\n\n\n# Name of the job.\nexport JOB_NAME=results_figure_11_pipeline_parallel_size_${PP}_batch_size_${GBS}\n\n\n# Import the configs.\n. `pwd`/CONFIG.sh\n\n\n# Submit the job.\n. `pwd`/SBATCH.sh\n\n\nexit 0\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/run_figure_12.sh",
"content": "#!/bin/bash\n\n# ================================\n# Choose the case to run.\n# ================================\n\n# Interleaved schedule options = [YES, NO].\nINTERLEAVED=YES\n\n# Batch size (global batch size) options = [12, 24, 36, ..., 60].\nGBS=12\n\n\n\n\n\n# Set interleaved schedule options.\nif [ ${INTERLEAVED} == \"YES\" ]; then\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform --num-layers-per-virtual-pipeline-stage 2 \"\nelif [ ${INTERLEAVED} == \"NO\" ]; then\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nelse\n echo \"Invalid configuration\"\n exit 1\nfi\n\n\n# Other params.\nTP=8\nPP=12\nMBS=1\nNLS=96\nHS=12288\nNAH=96\nDDP=local\nNNODES=12\n\n\n# Name of the job.\nexport JOB_NAME=results_figure_12_interleaved_${INTERLEAVED}_batch_size_${GBS}\n\n\n# Import the configs.\n. `pwd`/CONFIG.sh\n\n\n# Submit the job.\n. `pwd`/SBATCH.sh\n\n\nexit 0\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/run_figure_13.sh",
"content": "#!/bin/bash\n\n# ================================\n# Choose the case to run.\n# ================================\n\n# Pipeline-parallel size options = [2, 4, 8, 16, 32].\nPP=2\n\n# Batch size (global batch size) options = [32, 128].\nGBS=32\n\n\n\n\n\n# Set pipeline-parallel and tensor-parallel size options.\nTP=$((64/PP))\n\n\n# Other params.\nMBS=1\nNLS=32\nHS=20480\nNAH=128\nDDP=local\nMEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nNNODES=8\n\n\n# Name of the job.\nexport JOB_NAME=results_figure_13_pipeline_parallel_size_${PP}_tensor_parallel_size_${TP}_batch_size_${GBS}\n\n\n# Import the configs.\n. `pwd`/CONFIG.sh\n\n\n# Submit the job.\n. `pwd`/SBATCH.sh\n\n\nexit 0\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/run_figure_14.sh",
"content": "#!/bin/bash\n\n# ================================\n# Choose the case to run.\n# ================================\n\n# Pipeline-parallel size options = [2, 4, 8, 16, 32].\nPP=2\n\n# Batch size (global batch size) options = [32, 512].\nGBS=32\n\n\n\n\n\n# Set pipeline-parallel and data-parallel size options.\nDP=$((64/PP))\n\n\n# Other params.\nTP=1\nMBS=1\nNLS=32\nHS=3840\nNAH=32\nDDP=local\nMEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nNNODES=8\n\n\n# Name of the job.\nexport JOB_NAME=results_figure_14_pipeline_parallel_size_${PP}_data_parallel_size_${DP}_batch_size_${GBS}\n\n\n# Import the configs.\n. `pwd`/CONFIG.sh\n\n\n# Submit the job.\n. `pwd`/SBATCH.sh\n\n\nexit 0\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/run_figure_15.sh",
"content": "#!/bin/bash\n\n# ================================\n# Choose the case to run.\n# ================================\n\n# Tensor-parallel size options = [2, 4, 8, 16, 32].\nTP=2\n\n# Batch size (global batch size) options = [32, 128, 512].\nGBS=32\n\n\n\n\n\n# Set tensor-parallel and data-parallel size options.\nDP=$((64/TP))\n\n\n# Other params.\nPP=1\nMBS=1\nNLS=32\nHS=3840\nNAH=32\nDDP=local\nMEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nNNODES=8\n\n\n# Name of the job.\nexport JOB_NAME=results_figure_15_tensor_parallel_size_${TP}_data_parallel_size_${DP}_batch_size_${GBS}\n\n\n# Import the configs.\n. `pwd`/CONFIG.sh\n\n\n# Submit the job.\n. `pwd`/SBATCH.sh\n\n\nexit 0\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/run_figure_16.sh",
"content": "#!/bin/bash\n\n# ================================\n# Choose the case to run.\n# ================================\n\n# Microbatch size options = [1, 2, 4, 8].\nMBS=1\n\n# Batch size (global batch size) options = [128, 512].\nGBS=128\n\n\n\n\n\n# Other params.\nTP=8\nPP=8\nNLS=32\nHS=15360\nNAH=128\nDDP=local\nMEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nNNODES=8\n\n\n# Name of the job.\nexport JOB_NAME=results_figure_16_microbatch_size_${MBS}_batch_size_${GBS}\n\n\n# Import the configs.\n. `pwd`/CONFIG.sh\n\n\n# Submit the job.\n. `pwd`/SBATCH.sh\n\n\nexit 0\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/run_figure_17.sh",
"content": "#!/bin/bash\n\n# ================================\n# Choose the case to run.\n# ================================\n\n# Activation recomputation options = [YES, NO].\nACTIVATION_RECOMPUTATION=YES\n\n# Batch size (global batch size) options = [1, 2, 4, ..., 256].\nGBS=1\n\n\n\n\n\n# Set activation recomputation.\nif [ ${ACTIVATION_RECOMPUTATION} == \"YES\" ]; then\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nelif [ ${ACTIVATION_RECOMPUTATION} == \"NO\" ]; then\n MEGATRON_EXTRA_PARAMS=\"\"\nelse\n echo \"Invalid configuration\"\n exit 1\nfi\n\n\n# Other params.\nTP=8\nPP=16\nMBS=1\nNLS=80\nHS=12288\nNAH=96\nDDP=local\nNNODES=16\n\n\n# Name of the job.\nexport JOB_NAME=results_figure_17_activation_recomputation_${ACTIVATION_RECOMPUTATION}_batch_size_${GBS}\n\n\n# Import the configs.\n. `pwd`/CONFIG.sh\n\n\n# Submit the job.\n. `pwd`/SBATCH.sh\n\n\nexit 0\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/run_figure_18.sh",
"content": "#!/bin/bash\n\n# ================================\n# Choose the case to run.\n# ================================\n\n# Scatter-gather communication optimization options = [YES, NO].\nSCATTER_GATHER=YES\n\n# Batch size (global batch size) options = [12, 24, 36, ..., 60].\nGBS=12\n\n\n\n\n\n# Set scatter-gather communication optimization options.\nif [ ${SCATTER_GATHER} == \"YES\" ]; then\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform --num-layers-per-virtual-pipeline-stage 2 \"\nelif [ ${SCATTER_GATHER} == \"NO\" ]; then\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform --num-layers-per-virtual-pipeline-stage 2 --no-scatter-gather-tensors-in-pipeline \"\nelse\n echo \"Invalid configuration\"\n exit 1\nfi\n\n\n# Other params.\nTP=8\nPP=12\nMBS=1\nNLS=96\nHS=12288\nNAH=96\nDDP=local\nNNODES=12\n\n\n# Name of the job.\nexport JOB_NAME=results_figure_18_scatter_gather_${SCATTER_GATHER}_batch_size_${GBS}\n\n\n# Import the configs.\n. `pwd`/CONFIG.sh\n\n\n# Submit the job.\n. `pwd`/SBATCH.sh\n\n\nexit 0\n\n\n\n"
},
{
"path": "examples/academic_paper_scripts/sc21/run_table_1.sh",
"content": "#!/bin/bash\n\n# ================================\n# Choose the case to run.\n# ================================\n# model size options = [1.7B, 3.6B, 7.5B, 18B, 39B, 76B, 145B, 310B, 530B, 1T]\nMODEL_SIZE=1.7B\n\n\n\n\n\n\nif [ ${MODEL_SIZE} == \"1.7B\" ]; then\n TP=1\n PP=1\n MBS=16\n GBS=512\n NLS=24\n HS=2304\n NAH=24\n DDP=torch\n NNODES=4\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nelif [ ${MODEL_SIZE} == \"3.6B\" ]; then\n TP=2\n PP=1\n MBS=16\n GBS=512\n NLS=30\n HS=3072\n NAH=32\n DDP=torch\n NNODES=8\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nelif [ ${MODEL_SIZE} == \"7.5B\" ]; then\n TP=4\n PP=1\n MBS=16\n GBS=512\n NLS=36\n HS=4096\n NAH=32\n DDP=torch\n NNODES=16\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nelif [ ${MODEL_SIZE} == \"18B\" ]; then\n TP=8\n PP=1\n MBS=8\n GBS=1024\n NLS=40\n HS=6144\n NAH=48\n DDP=torch\n NNODES=32\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nelif [ ${MODEL_SIZE} == \"39B\" ]; then\n TP=8\n PP=2\n MBS=4\n GBS=1536\n NLS=48\n HS=8192\n NAH=64\n DDP=local\n NNODES=64\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nelif [ ${MODEL_SIZE} == \"76B\" ]; then\n TP=8\n PP=4\n MBS=2\n GBS=1792\n NLS=60\n HS=10240\n NAH=80\n DDP=local\n NNODES=128\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform --num-layers-per-virtual-pipeline-stage 5\"\nelif [ ${MODEL_SIZE} == \"145B\" ]; then\n TP=8\n PP=8\n MBS=2\n GBS=2304\n NLS=80\n HS=12288\n NAH=96\n DDP=local\n NNODES=192\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform --num-layers-per-virtual-pipeline-stage 5 \"\nelif [ ${MODEL_SIZE} == \"310B\" ]; then\n TP=8\n PP=16\n MBS=1\n GBS=2160\n NLS=96\n HS=16384\n NAH=128\n DDP=local\n NNODES=240\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform --num-layers-per-virtual-pipeline-stage 3 \"\nelif [ ${MODEL_SIZE} == \"530B\" ]; then\n TP=8\n PP=35\n MBS=1\n GBS=2520\n NLS=105\n HS=20480\n NAH=128\n DDP=local\n NNODES=315\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform --num-layers-per-virtual-pipeline-stage 1 \"\nelif [ ${MODEL_SIZE} == \"1T\" ]; then\n TP=8\n PP=64\n MBS=1\n GBS=3072\n NLS=128\n HS=25600\n NAH=160\n DDP=local\n NNODES=384\n MEGATRON_EXTRA_PARAMS=\"--activations-checkpoint-method uniform \"\nelse\n echo \"Invalid configuration\"\n exit 1\nfi\n\n\n# Name of the job\nexport JOB_NAME=results_table_1_model_size_${MODEL_SIZE}\n\n\n# Import the configs.\n. `pwd`/CONFIG.sh\n\n\n# Submit the job.\n. `pwd`/SBATCH.sh\n\n\nexit 0\n\n\n\n"
},
{
"path": "examples/bert/README.md",
"content": "# BERT MODEL\n\n## Table of contents\n- [1. Training Setup](#1-training-setup)\n- [2. Configurations](#2-configurations)\n\n## 1. Training setup\n\n\nTo run the model using a docker container run it as follows\n```\nPYTORCH_IMAGE=nvcr.io/nvidia/pytorch:24.01-py3\nCHECKPOINT_PATH=\"\" #\nTENSORBOARD_LOGS_PATH=\"\"#\nVOCAB_FILE=\"\" #//bert-vocab.txt\nDATA_PATH=\"\" #_text_document\n\ndocker run \\\n --gpus=all \\\n --ipc=host \\\n --workdir /workspace/megatron-lm \\\n -v /path/to/data:/path/to/data \\\n -v /path/to/megatron-lm:/workspace/megatron-lm \\\n megatron-lm nvcr.io/nvidia/pytorch:24.01-py3 \\\n bash examples/bert/train_bert_340m_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH $VOCAB_FILE $DATA_PATH \"\n\n```\nNOTE: Depending on the environment you are running it the above command might like slightly different.\n\n\n## 2. Configurations\n\nThe example in this folder shows you how to run 340m large model. There are other configs you could run as well\n\n### 4B\n```\n --num-layers 48 \\\n --hidden-size 2560 \\\n --num-attention-heads 32 \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n\n```\n\n### 20B\n```\n --num-layers 48 \\\n --hidden-size 6144 \\\n --num-attention-heads 96 \\\n --tensor-model-parallel-size 4 \\\n --pipeline-model-parallel-size 4 \\\n\n```"
},
{
"path": "examples/bert/train_bert_340m_distributed.sh",
"content": "#!/bin/bash\n\n# Runs the \"340M\" parameter model (Bert - Large)\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\nGPUS_PER_NODE=8\n# Change for multinode config\nMASTER_ADDR=localhost\nMASTER_PORT=6000\nNUM_NODES=1\nNODE_RANK=0\nWORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))\n\nCHECKPOINT_PATH=$1 #\nTENSORBOARD_LOGS_PATH=$2 #\nVOCAB_FILE=$3 #/bert-vocab.json\nDATA_PATH=$4 #_text_document\n\nDISTRIBUTED_ARGS=(\n --nproc_per_node $GPUS_PER_NODE \n --nnodes $NUM_NODES \n --master_addr $MASTER_ADDR \n --master_port $MASTER_PORT\n)\n\nBERT_MODEL_ARGS=(\n --num-layers 24 \n --hidden-size 1024 \n --num-attention-heads 16 \n --seq-length 512 \n --max-position-embeddings 512 \n --attention-backend auto # Can use (flash/fused/unfused/local)\n)\n\nTRAINING_ARGS=(\n --micro-batch-size 4 \n --global-batch-size 32 \n --train-iters 1000000 \n --weight-decay 1e-2 \n --clip-grad 1.0 \n --fp16\n --lr 0.0001\n --lr-decay-iters 990000 \n --lr-decay-style linear \n --min-lr 1.0e-5 \n --weight-decay 1e-2 \n --lr-warmup-fraction .01 \n --clip-grad 1.0 \n)\n\nMODEL_PARALLEL_ARGS=(\n\t--tensor-model-parallel-size 8 \n\t--pipeline-model-parallel-size 16 \n)\n\nDATA_ARGS=(\n --data-path $DATA_PATH \n --vocab-file $VOCAB_FILE \n --split 949,50,1\n)\n\nEVAL_AND_LOGGING_ARGS=(\n --log-interval 100\n --save-interval 10000 \n --eval-interval 1000 \n --save $CHECKPOINT_PATH \n --load $CHECKPOINT_PATH \n --eval-iters 10\n --tensorboard-dir $TENSORBOARD_LOGS_PATH \n)\n\ntorchrun ${DISTRIBUTED_ARGS[@]} pretrain_bert.py \\\n ${BERT_MODEL_ARGS[@]} \\\n ${TRAINING_ARGS[@]} \\\n ${MODEL_PARALLEL_ARGS[@]} \\\n ${DATA_ARGS[@]} \\\n ${EVAL_AND_LOGGING_ARGS[@]}\n "
},
{
"path": "examples/export/README.md",
"content": "# Megatron Core Export\n\nThis module is used to export megatron core models to different inference frameworks.\nCurrently we support TRTLLM export . In the future we will be adding support for VLLM etc.\n\n## PTQ AND EXPORT\nFollow the examples of [Model Optimizer](../post_training/modelopt) to perform post training quantization, followed by an export to a HF-like checkpoint for TensorRT-LLM, vLLM, and SGLang deployment.\n\n# TRTLLM EXPORT\nFollow the instructions in [trtllm_export](./trtllm_export/) to do export to TRTLLM checkpoint format alone.\n"
},
{
"path": "examples/export/trtllm_export/README.md",
"content": "# Megatron Core To TRTLLM Export Documentation\nThis guide will walk you through how you can use the megatron core export for exporting models to trtllm format\n\n### Contents\n- [Megatron Core To TRTLLM Export Documentation](#megatron-core-to-trtllm-export-documentation)\n- [Contents](#contents)\n - [1. Quick Start](#1-quick-start)\n - [1.1 Understanding The Code](#11-understanding-the-code)\n - [1.2 Running The Code](#12-running-the-code)\n - [2. GPU Export](#2-gpu-export)\n - [3. Future work](#4-future-work)\n\n#### 1. Quick Start\nThis will walk you through the flow of converting an mcore gpt model to trtllm format using single device mode. The file can be found at [gpt_single_device_cpu_export.py](./single_device_export/gpt_single_device_cpu_export.py)\n\nNOTE: For faster performance, if your entire model will fit into gpu memory, pre transfer the model state dict to gpu and then call the get_trtllm_pretrained_config_and_model_weights function.\n\n \n\n##### 1.1 Understanding The Code\n***STEP 1 - We initialize model parallel and other default arguments***\nWe initalize tp and pp to 1 so that we can get the full model state dict on cpu\n```python\n initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)\n```\n\n***STEP 2 - We load the model using the model_provider_function***\nNOTE: We create a simple gpt model\n\n```python\n transformer_config = TransformerConfig(\n num_layers=2, \n hidden_size=64, # Needs to be atleast 32 times num_attn_heads\n num_attention_heads=2, \n use_cpu_initialization=True, \n pipeline_dtype=torch.float32,\n )\n\n gpt_model = GPTModel(\n config=transformer_config, \n transformer_layer_spec=get_gpt_layer_local_spec(), \n vocab_size=100, \n max_sequence_length=_SEQUENCE_LENGTH,\n )\n\n # Optionally you can also load a model using this code \n # sharded_state_dict=gpt_model.sharded_state_dict(prefix='')\n # checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)\n # gpt_model.load_state_dict(checkpoint)\n\n```\n\n***STEP 3 - Instantiate the TRTLLM Helper***\nWe instantiate the [TRTLLM Helper](../../../megatron/core/export/trtllm/trtllm_helper.py) For the GPT model we instantiate trtllm_helper as shown below.\n```python\n if hasattr(gpt_model, \"rotary_pos_emb\"):\n seq_len_interpolation_factor = gpt_model.rotary_pos_emb.seq_len_interpolation_factor\n\n trtllm_helper = TRTLLMHelper(\n transformer_config=gpt_model.config, \n model_type=ModelType.gpt,\n position_embedding_type = gpt_model.position_embedding_type, \n max_position_embeddings = gpt_model.max_position_embeddings, \n rotary_percentage = gpt_model.rotary_percent,\n rotary_base = gpt_model.rotary_base,\n moe_tp_mode = 2,\n multi_query_mode = False,\n activation = \"gelu\", \n seq_len_interpolation_factor = seq_len_interpolation_factor,\n share_embeddings_and_output_weights=gpt_model.share_embeddings_and_output_weights\n ) \n```\n\n***STEP 4 - Get the TRTLLM Weights and configs***\nTo convert model weights to trtllm weights and configs, we use the [single_device_converter](../../../megatron/core/export/trtllm/trtllm_weights_converter/single_device_trtllm_model_weights_converter.py). We pass as inputs the model state dict, and export config. In this example we use inference tp size as 2 for the export. \n\n```python\n model_state_dict={}\n for key , val in gpt_model.state_dict().items():\n # val is non for _extra_state layers . We filter it out\n if val is not None:\n model_state_dict[key] = val\n\n export_config = ExportConfig(inference_tp_size = 2)\n weight_list, config_list = trtllm_helper.get_trtllm_pretrained_config_and_model_weights(\n model_state_dict= model_state_dict,\n dtype = DataType.bfloat16,\n export_config=export_config\n )\n```\n\n***STEP 5 - Build the TRTLLM Engine***\nFollowing code is used to build the TRTLLM Engine. \n\n```python\n for trtllm_model_weights, trtllm_model_config in zip(weight_list, config_list):\n trtllm_helper.build_and_save_engine(\n max_input_len=256,\n max_output_len=256,\n max_batch_size=8,\n engine_dir='/opt/megatron-lm/engine',\n trtllm_model_weights=trtllm_model_weights,\n trtllm_model_config=trtllm_model_config,\n lora_ckpt_list=None,\n use_lora_plugin=None,\n max_lora_rank=64,\n lora_target_modules=None,\n max_prompt_embedding_table_size=0,\n paged_kv_cache=True,\n remove_input_padding=True,\n paged_context_fmha=False,\n use_refit=False,\n max_num_tokens=None,\n max_seq_len=512,\n opt_num_tokens=None,\n max_beam_width=1,\n tokens_per_block=128,\n multiple_profiles=False,\n gpt_attention_plugin=\"auto\",\n gemm_plugin=\"auto\",\n )\n```\n \n\n##### 1.2 Running The Code\nAn example run script is shown below. \n\n```\n# In a workstation \nMLM_PATH=/path/to/megatron-lm\nCONTAINER_IMAGE=gitlab-master.nvidia.com:5005/dl/joc/nemo-ci/trtllm_0.12/train:pipe.17669124-x86\n\ndocker run -it --gpus=all --ipc=host -v $MLM_PATH/:/opt/megatron-lm $CONTAINER_IMAGE bash\n\n# Inside the container run the following. \n\ncd /opt/megatron-lm/\n\nCUDA_VISIBLE_DEVICES=0 torchrun --nproc-per-node 1 examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py\n```\n\n \n\n#### 2. GPU Export\nYou can use the [gpt_distributed_gpu_export.py](./distributed_export/gpt_distributed_gpu_export.py) to run a more optimized on device distributed. version of trtllm export. Internally this uses the [distributed_converter](../../../megatron/core/export/trtllm/trtllm_weights_converter/distributed_trtllm_model_weights_converter.py) to convert model weights on device. \nIn the single device version you collect all the model weights on CPU/GPU, convert it to trtllm format, and then store the engine back on disk. In the GPU version you load each individual state dict on the gpus, convert it on the device itself and store the engine on disk. \n\nTo run the gpu version \n\n```\nCUDA_VISIBLE_DEVICES=0,1 torchrun --nproc-per-node 2 examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py\n```\n\n \n\n#### 3. Future work\nThe following are planned for the future releases . \n* Pipeline parallellism for export (Work in progress) \n* GPU Export for more models (Work in progress for some models)\n* Refit functionality\n* VLLM Support"
},
{
"path": "examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py",
"content": "import os\nimport torch\nfrom megatron.core import parallel_state\nfrom megatron.core import dist_checkpointing\nfrom megatron.core.export.model_type import ModelType\nfrom megatron.core.export.data_type import DataType\nfrom megatron.core.export.trtllm.trtllm_helper import TRTLLMHelper\nfrom megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed\nfrom megatron.core.transformer.transformer_config import TransformerConfig\nfrom megatron.core.models.gpt.gpt_model import GPTModel\nfrom megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec\n\n\n_SEQUENCE_LENGTH = 64\n_VOCAB_SIZE = 256\n\ndef initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1):\n parallel_state.destroy_model_parallel()\n\n # Torch setup for distributed training\n rank = int(os.environ['LOCAL_RANK'])\n world_size = torch.cuda.device_count()\n torch.cuda.set_device(rank)\n torch.distributed.init_process_group(world_size=world_size, rank=rank)\n\n # Megatron core distributed training initialization\n parallel_state.initialize_model_parallel(tensor_model_parallel_size = tensor_model_parallel_size, pipeline_model_parallel_size=pipeline_model_parallel_size)\n\ndef model_provider():\n \"\"\"Build the model.\"\"\"\n\n transformer_config = TransformerConfig(\n num_layers=2, \n hidden_size=64, \n num_attention_heads=2, \n use_cpu_initialization=True, \n pipeline_dtype=torch.float32\n )\n\n gpt_model = GPTModel(\n config=transformer_config, \n transformer_layer_spec=get_gpt_layer_local_spec(), \n vocab_size=_VOCAB_SIZE, \n max_sequence_length=_SEQUENCE_LENGTH,\n )\n\n return gpt_model\n\ndef load_distributed_checkpoint(checkpoint_path, gpt_model):\n sharded_state_dict=gpt_model.sharded_state_dict(prefix='')\n checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)\n gpt_model.load_state_dict(checkpoint)\n return gpt_model\n\nif __name__ == \"__main__\":\n initialize_distributed(tensor_model_parallel_size=2, pipeline_model_parallel_size=1)\n model_parallel_cuda_manual_seed(123)\n\n gpt_model = model_provider()\n device = torch.device(\"cuda\")\n gpt_model.to(device) \n \n # Optionally you can also load a gpt model from ckpt_path using this code below\n # gpt_model = load_distributed_checkpoint(gpt_model=gpt_model, checkpoint_path=ckpt_path)\n\n seq_len_interpolation_factor = None\n if hasattr(gpt_model, \"rotary_pos_emb\"):\n seq_len_interpolation_factor = gpt_model.rotary_pos_emb.seq_len_interpolation_factor\n\n trtllm_helper = TRTLLMHelper(\n transformer_config=gpt_model.config, \n model_type=ModelType.gpt,\n position_embedding_type = gpt_model.position_embedding_type, \n max_position_embeddings = gpt_model.max_position_embeddings, \n rotary_percentage = gpt_model.rotary_percent,\n rotary_base = gpt_model.rotary_base,\n moe_tp_mode = 2,\n multi_query_mode = False,\n activation = \"gelu\", \n seq_len_interpolation_factor = seq_len_interpolation_factor,\n share_embeddings_and_output_weights=gpt_model.share_embeddings_and_output_weights\n )\n \n\n trtllm_model_weights, trtllm_model_config = trtllm_helper.get_trtllm_pretrained_config_and_model_weights(\n model_state_dict= gpt_model.state_dict(),\n dtype = DataType.bfloat16,\n on_device_distributed_conversion=True, \n vocab_size=_VOCAB_SIZE, \n gpus_per_node=2,\n )\n\n trtllm_helper.build_and_save_engine(\n max_input_len=256,\n max_output_len=256,\n max_batch_size=8,\n engine_dir='/opt/megatron-lm/engine',\n trtllm_model_weights=trtllm_model_weights[0],\n trtllm_model_config=trtllm_model_config[0],\n lora_ckpt_list=None,\n use_lora_plugin=None,\n max_lora_rank=64,\n lora_target_modules=None,\n max_prompt_embedding_table_size=0,\n paged_kv_cache=True,\n remove_input_padding=True,\n paged_context_fmha=False,\n use_refit=False,\n max_num_tokens=None,\n max_seq_len=512,\n opt_num_tokens=None,\n max_beam_width=1,\n tokens_per_block=128,\n multiple_profiles=False,\n gpt_attention_plugin=\"auto\",\n gemm_plugin=\"auto\",\n )\n"
},
{
"path": "examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py",
"content": "import os\nimport torch\nfrom megatron.core import parallel_state\nfrom megatron.core import dist_checkpointing\nfrom megatron.core.export.model_type import ModelType\nfrom megatron.core.export.data_type import DataType\nfrom megatron.core.export.export_config import ExportConfig\nfrom megatron.core.export.trtllm.trtllm_helper import TRTLLMHelper\nfrom megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed\nfrom megatron.core.transformer.transformer_config import TransformerConfig\nfrom megatron.core.models.gpt.gpt_model import GPTModel\nfrom megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec\n\n\n_SEQUENCE_LENGTH = 64\n\n\ndef initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1):\n parallel_state.destroy_model_parallel()\n\n # Torch setup for distributed training\n rank = int(os.environ['LOCAL_RANK'])\n world_size = torch.cuda.device_count()\n torch.cuda.set_device(rank)\n torch.distributed.init_process_group(world_size=world_size, rank=rank)\n\n # Megatron core distributed training initialization\n parallel_state.initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size)\n\ndef model_provider():\n \"\"\"Build the model.\"\"\"\n\n transformer_config = TransformerConfig(\n num_layers=2, \n hidden_size=64, # Needs to be atleast 32 times num_attn_heads\n num_attention_heads=2, \n use_cpu_initialization=True, \n pipeline_dtype=torch.float32,\n )\n\n gpt_model = GPTModel(\n config=transformer_config, \n transformer_layer_spec=get_gpt_layer_local_spec(), \n vocab_size=100, \n max_sequence_length=_SEQUENCE_LENGTH,\n )\n\n return gpt_model\n\ndef load_distributed_checkpoint(checkpoint_path, gpt_model):\n sharded_state_dict=gpt_model.sharded_state_dict(prefix='')\n checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)\n gpt_model.load_state_dict(checkpoint)\n return gpt_model\n\nif __name__ == \"__main__\":\n # Need to use TP1 PP1 for export on single device\n initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)\n model_parallel_cuda_manual_seed(123)\n\n gpt_model = model_provider()\n\n # Optionally you can also load a gpt model from ckpt_path using this code below\n # gpt_model = load_distributed_checkpoint(gpt_model=gpt_model, checkpoint_path=ckpt_path)\n\n seq_len_interpolation_factor = None\n if hasattr(gpt_model, \"rotary_pos_emb\"):\n seq_len_interpolation_factor = gpt_model.rotary_pos_emb.seq_len_interpolation_factor\n\n trtllm_helper = TRTLLMHelper(\n transformer_config=gpt_model.config, \n model_type=ModelType.gpt,\n position_embedding_type = gpt_model.position_embedding_type, \n max_position_embeddings = gpt_model.max_position_embeddings, \n rotary_percentage = gpt_model.rotary_percent,\n rotary_base = gpt_model.rotary_base,\n moe_tp_mode = 2,\n multi_query_mode = False,\n activation = \"gelu\", \n seq_len_interpolation_factor = seq_len_interpolation_factor,\n share_embeddings_and_output_weights=gpt_model.share_embeddings_and_output_weights\n )\n \n\n export_config = ExportConfig(inference_tp_size = 2)\n # NOTE : For faster performance, if your entire model will fit in gpu memory, transfer model state dict to GPU and then call this api\n weight_list, config_list = trtllm_helper.get_trtllm_pretrained_config_and_model_weights(\n model_state_dict= gpt_model.state_dict(),\n dtype = DataType.bfloat16,\n export_config=export_config\n )\n\n for trtllm_model_weights, trtllm_model_config in zip(weight_list, config_list):\n trtllm_helper.build_and_save_engine(\n max_input_len=256,\n max_output_len=256,\n max_batch_size=8,\n engine_dir='/opt/megatron-lm/engine',\n trtllm_model_weights=trtllm_model_weights,\n trtllm_model_config=trtllm_model_config,\n lora_ckpt_list=None,\n use_lora_plugin=None,\n max_lora_rank=64,\n lora_target_modules=None,\n max_prompt_embedding_table_size=0,\n paged_kv_cache=True,\n remove_input_padding=True,\n paged_context_fmha=False,\n use_refit=False,\n max_num_tokens=None,\n max_seq_len=512,\n opt_num_tokens=None,\n max_beam_width=1,\n tokens_per_block=128,\n multiple_profiles=False,\n gpt_attention_plugin=\"auto\",\n gemm_plugin=\"auto\",\n )"
},
{
"path": "examples/gpt3/README.md",
"content": "# GPT3 MODEL\n\n## Table of contents\n- [1. Training Setup](#1-training-setup)\n- [2. Configurations](#2-configurations)\n- [3. Training Results](#3-training-results)\n\n## 1. Training setup\n\n\nTo run the model using a docker container run it as follows\n```\nPYTORCH_IMAGE=nvcr.io/nvidia/pytorch:24.01-py3\nCHECKPOINT_PATH=\"\" #\nTENSORBOARD_LOGS_PATH=\"\"#\nVOCAB_FILE=\"\" #/gpt2-vocab.json\nMERGE_FILE=\"\" #/gpt2-merges.txt\nDATA_PATH=\"\" #_text_document\n\ndocker run \\\n --gpus=all \\\n --ipc=host \\\n --workdir /workspace/megatron-lm \\\n -v /path/to/data:/path/to/data \\\n -v /path/to/megatron-lm:/workspace/megatron-lm \\\n megatron-lm nvcr.io/nvidia/pytorch:24.01-py3 \\\n bash examples/gpt3/train_gpt3_175b_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH $VOCAB_FILE $MERGE_FILE $DATA_PATH \"\n\n```\nNOTE: Depending on the environment you are running it the above command might like slightly different.\n\n\n## 2. Configurations\n\nThe example in this folder shows you how to run 175B model. There are other configs you could run as well\n\n### 345M\n```\n --num-layers 12 \\\n --hidden-size 512 \\\n --num-attention-heads 8 \\\n --seq-length 1024 \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n\n```\n\n### 857M\n```\n --num-layers 24 \\\n --hidden-size 1024 \\\n --num-attention-heads 16 \\\n --seq-length 2048 \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n\n```\n"
},
{
"path": "examples/gpt3/gpt_config.yaml",
"content": "# WARNING: Yaml configs is currently an experimental feature\nlanguage_model:\n # model architecture\n num_layers: 24\n hidden_size: 1024\n num_attention_heads: 16\n num_query_groups: null\n\n ffn_hidden_size: null\n kv_channels: null\n hidden_dropout: 0.0\n attention_dropout: 0.0\n fp32_residual_connection: False\n\n apply_residual_connection_post_layernorm: False\n layernorm_epsilon: 1.e-5\n layernorm_zero_centered_gamma: True\n add_bias_linear: False\n bias_activation_fusion: False\n add_qkv_bias: False\n gated_linear_unit: False\n activation_func: swiglu\n num_moe_experts: null\n rotary_interleaved: False\n window_size: null\n\n # initialization\n init_method: null\n init_method_std: 0.02\n output_layer_init_method: null\n\n # mixed-precision\n apply_query_key_layer_scaling: False\n attention_softmax_in_fp32: False\n\n # fusion\n bias_swiglu_fusion: True\n masked_softmax_fusion: True\n persist_layer_norm: False\n memory_efficient_layer_norm: False\n bias_dropout_fusion: True\n apply_rope_fusion: True\n\n # activation recomputation\n recompute_granularity: null\n recompute_method: null\n recompute_num_layers: null\n distribute_saved_activations: null\n\n # fp8 related\n fp8: null\n fp8_margin: 0\n fp8_interval: 1\n fp8_amax_history_len: 1\n fp8_amax_compute_algo: \"most_recent\"\n fp8_wgrad: True\n\n # miscellaneous\n clone_scatter_output_in_embedding: True\n\n normalization: \"LayerNorm\" # alt value supported by TE: \"RMSNorm\"\n\n # MoE related\n moe_router_load_balancing_type: \"aux_loss\"\n moe_router_topk: 2\n moe_router_group_topk: null\n moe_router_num_groups: null\n moe_grouped_gemm: False\n moe_aux_loss_coeff: 0 # 1e-2 would be a good start value for load balance loss.\n moe_z_loss_coeff: null # 1e-3 would be a good start value for z-loss\n moe_input_jitter_eps: null\n moe_token_dropping: False\n\nmodel_parallel:\n # Model parallelism\n tensor_model_parallel_size: 1\n context_parallel_size: 1\n pipeline_model_parallel_size: 1\n virtual_pipeline_model_parallel_size: null\n sequence_parallel: True\n expert_model_parallel_size: 1\n\n # Initialization\n perform_initialization: True\n use_cpu_initialization: null\n\n # Training\n fp16: False\n bf16: True\n params_dtype: null # Set from above arguments for core\n timers: null\n\n # Optimizations\n gradient_accumulation_fusion: True\n tp_comm_overlap: False\n\n # Debug Options\n tp_comm_split_ag: True\n tp_comm_atomic_ag: True\n tp_comm_split_rs: True\n tp_comm_atomic_rs: True\n tp_comm_bulk_wgrad: True\n tp_comm_bulk_dgrad: True\n\n # Parallelism\n finalize_model_grads_func: null\n\n # Pipeline Parallel\n pipeline_dtype: null\n grad_scale_func: null\n enable_autocast: False\n autocast_dtype: null\n variable_seq_lengths: False\n num_microbatches_with_partial_activation_checkpoints: null\n overlap_p2p_comm: False\n batch_p2p_comm: True\n batch_p2p_sync: True\n use_ring_exchange_p2p: False\n deallocate_pipeline_outputs: False\n no_sync_func: null\n grad_sync_func: null\n param_sync_func: null\n\n # CPU Offloading\n cpu_offloading: False\n cpu_offloading_num_layers: 0\n _cpu_offloading_context: null\n cpu_offloading_weights: False\n cpu_offloading_activations: True\n\n # Timing\n barrier_with_L1_time: True\n\n# training:\nuse_legacy_models: False\nspec: null\nmicro_batch_size: 2\nglobal_batch_size: 128\nrampup_batch_size: [32, 32, 65324160] \ncheck_for_nan_in_loss_and_grad: True\nnum_layers_per_virtual_pipeline_stage: null\n\nencoder_num_layers: null\ndecoder_num_layers: null\nrotary_seq_len_interpolation_factor: null\nadd_position_embedding: False\nmake_vocab_size_divisible_by: 128\ngroup_query_attention: False\n\n\nexit_signal_handler: False\nexit_duration_in_mins: null\nexit_interval: null\n\nuntie_embeddings_and_output_weights: True\nposition_embedding_type: rope\nrotary_percent: 0.5\nopenai_gelu: False\nsquared_relu: False\nswiglu: True\nonnx_safe: null\nbert_binary_head: True\nmax_position_embeddings: 4096\n\ntransformer_impl: local\nuse_flash_attn: False\nseed: 1234\ndata_parallel_random_init: False\n\n# Optimizer\noptimizer: adam\nlr: 2.5e-4\nlr_decay_style: cosine\nlr_decay_iters: null\nlr_decay_samples: 255126953\nlr_warmup_fraction: null\nlr_warmup_iters: 0\nlr_warmup_samples: 81381\nlr_warmup_init: 0.0\nmin_lr: 2.5e-5\nweight_decay: 0.1\nstart_weight_decay: null\nend_weight_decay: null\nweight_decay_incr_style: constant\nclip_grad: 1.0\nadam_beta1: 0.9\nadam_beta2: 0.95\nadam_eps: 1.e-08\nsgd_momentum: 0.9\noverride_opt_param_scheduler: False\nuse_checkpoint_opt_param_scheduler: False\n\n# checkpointing arguments\nsave: null\nsave_interval: 20000\nno_save_optim: null\nno_save_rng: null\nload: null\nno_load_optim: null\nno_load_rng: null\nfinetune: False\nuse_checkpoint_args: False\nexit_on_missing_checkpoint: False\n\n# loss arguments\nloss_scale: null\ninitial_loss_scale: 4294967296\nmin_loss_scale: 1.0\nloss_scale_window: 1000 \nhysteresis: 2\naccumulate_allreduce_grads_in_fp32: False\nfp16_lm_cross_entropy: False\n\n# distributed arguments\ndistributed_backend: nccl\ndistributed_timeout_minutes: 10\noverlap_grad_reduce: False\nalign_grad_reduce: True\noverlap_param_gather: False\nalign_param_gather: False\nscatter_gather_tensors_in_pipeline: True\nlocal_rank: null\nlazy_mpu_init: null\nempty_unused_memory_level: 0\nstandalone_embedding_stage: False\nuse_distributed_optimizer: False\nnccl_communicator_config_path: null\n\ntrain_iters: null\neval_iters: 32\neval_interval: 2000\nskip_train: False\n\nadlr_autoresume: False\nadlr_autoresume_interval: 1000\n\n# garbage collection\nmanual_gc: False\nmanual_gc_interval: 0\nmanual_gc_eval: True\n\ntp_comm_overlap_cfg: null\n\n#data\ndata_path: null\nsplit: '99,1,0'\ntrain_data_path: null\nvalid_data_path: null\ntest_data_path: null\ndata_cache_path: null\nmock_data: False\nvocab_size: null\nvocab_file: null\nmerge_file: null\nvocab_extra_ids: 0\nseq_length: 4096\nencoder_seq_length: null\ndecoder_seq_length: null\nsample_rate: 1.0\nmask_prob: 0.15\nshort_seq_prob: 0.1\nnum_workers: 2\ntokenizer_type: GPTSentencePieceTokenizer\ntokenizer_model: null\nreset_position_ids: False\nreset_attention_mask: False\neod_mask_loss: False\ntrain_samples: 268554688\ndataloader_type: null\n\n#profile:\nprofile: False\nprofile_ranks: [0]\nprofile_step_end: 12\nprofile_step_start: 10\n\n#logging:\nlog_params_norm: True\nlog_num_zeros_in_grad: True\nlog_throughput: False\nlog_progress: False\ntiming_log_level: 0\ntiming_log_option: minmax\ntensorboard_log_interval: 1\ntensorboard_queue_size: 1000\nlog_timers_to_tensorboard: False\nlog_validation_ppl_to_tensorboard: False\nlog_memory_to_tensorboard: False\nlog_world_size_to_tensorboard: False\nlog_loss_scale_to_tensorboard: True\nwandb_project: ''\nwandb_exp_name: ''\nwandb_save_dir: ''\nenable_one_logger: True\none_logger_project: megatron-lm\none_logger_run_name: null\nlog_interval: 100\ntensorboard_dir: null\n"
},
{
"path": "examples/gpt3/train_gpt3_175b_distributed.sh",
"content": "#!/bin/bash\n\n# Runs the \"175B\" parameter model\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\nGPUS_PER_NODE=8\n# Change for multinode config\nMASTER_ADDR=localhost\nMASTER_PORT=6000\nNUM_NODES=1\nNODE_RANK=0\nWORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))\n\nCHECKPOINT_PATH=$1 #\nTENSORBOARD_LOGS_PATH=$2 #\nVOCAB_FILE=$3 #/gpt2-vocab.json\nMERGE_FILE=$4 #/gpt2-merges.txt\nDATA_PATH=$5 #_text_document\n\nDISTRIBUTED_ARGS=(\n --nproc_per_node $GPUS_PER_NODE \n --nnodes $NUM_NODES \n --master_addr $MASTER_ADDR \n --master_port $MASTER_PORT\n)\n\nGPT_MODEL_ARGS=(\n --num-layers 96 \n --hidden-size 12288 \n --num-attention-heads 96 \n --seq-length 2048 \n --max-position-embeddings 2048 \n --attention-backend auto # Can use (flash/fused/unfused/local)\n)\n\nTRAINING_ARGS=(\n --micro-batch-size 1 \n --global-batch-size 1536 \n --rampup-batch-size 16 16 5859375 \n --train-iters 500000 \n --weight-decay 0.1 \n --adam-beta1 0.9 \n --adam-beta2 0.95 \n --init-method-std 0.006 \n --clip-grad 1.0 \n --fp16\n --lr 6.0e-5 \n --lr-decay-style cosine \n --min-lr 6.0e-6\n --lr-warmup-fraction .001 \n --lr-decay-iters 430000 \n)\n\nMODEL_PARALLEL_ARGS=(\n\t--tensor-model-parallel-size 8 \n\t--pipeline-model-parallel-size 16 \n)\n\nDATA_ARGS=(\n --data-path $DATA_PATH \n --vocab-file $VOCAB_FILE \n --merge-file $MERGE_FILE \n --split 949,50,1\n)\n\nEVAL_AND_LOGGING_ARGS=(\n --log-interval 100\n --save-interval 10000 \n --eval-interval 1000 \n --save $CHECKPOINT_PATH \n --load $CHECKPOINT_PATH \n --eval-iters 10\n --tensorboard-dir $TENSORBOARD_LOGS_PATH \n)\n\ntorchrun ${DISTRIBUTED_ARGS[@]} pretrain_gpt.py \\\n ${GPT_MODEL_ARGS[@]} \\\n ${TRAINING_ARGS[@]} \\\n ${MODEL_PARALLEL_ARGS[@]} \\\n ${DATA_ARGS[@]} \\\n ${EVAL_AND_LOGGING_ARGS[@]}\n"
},
{
"path": "examples/gptoss/01_convert_from_hf.py",
"content": "# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n\n\"\"\"Convert HuggingFace checkpoints to Megatron format.\"\"\"\n\nimport os\nimport argparse\n\nfrom megatron.bridge import AutoBridge\n\ndef _parse_args():\n parser = argparse.ArgumentParser(description=\"Convert HF LLMs to Megatron format\")\n parser.add_argument(\n \"--hf-model\",\n type=str,\n required=True,\n help=\"HuggingFace model identifier or path\",\n )\n parser.add_argument(\n \"--save-path\",\n type=str,\n default=None,\n help=\"Path to save the converted Megatron checkpoint\",\n )\n parser.add_argument('--local-rank', '--local_rank', type=int, default=0)\n return parser.parse_args()\n\nif __name__ == \"__main__\":\n args = _parse_args()\n HF_MODEL = args.hf_model\n SAVE_PATH = args.save_path\n WORLD_SIZE = int(os.environ.get(\"WORLD_SIZE\", 1))\n\n if SAVE_PATH is None:\n SAVE_PATH = f\"./megatron_checkpoints/{HF_MODEL.replace('/', '_')}\"\n \n print(f\"Converting {HF_MODEL} to Megatron format...\")\n print(f\"Save path: {SAVE_PATH}\")\n \n bridge = AutoBridge.from_hf_pretrained(HF_MODEL, trust_remote_code=True)\n provider = bridge.to_megatron_provider()\n # Update these configs as needed\n provider.expert_tensor_parallel_size = 1\n provider.tensor_model_parallel_size = 1\n provider.pipeline_model_parallel_size = WORLD_SIZE\n provider.finalize()\n \n model = provider.provide_distributed_model(wrap_with_ddp=False)\n \n bridge.save_megatron_model(\n model,\n SAVE_PATH,\n hf_tokenizer_path=HF_MODEL\n )\n \n print(f\"Saved Megatron checkpoint to {SAVE_PATH}\")\n"
},
{
"path": "examples/gptoss/02_train.sh",
"content": "#!/bin/bash\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}\n\n\n# Setup arguments with defaults\nCHECKPOINT_PATH=\"NO_VALUE_PROVIDED\"\nTENSORBOARD_LOGS_PATH=\"./tensorboard_logs/\"\nTOKENIZER_ARG=\"MOCK\"\nDATA_ARG=\"MOCK\"\nDISTRIBUTED_CONFIG_FILE=\"\"\n\n# Parse command line arguments\nwhile [[ $# -gt 0 ]]; do\n case $1 in\n --checkpoint-path)\n CHECKPOINT_PATH=\"$2\"\n shift 2\n ;;\n --tensorboard-logs-path)\n TENSORBOARD_LOGS_PATH=\"$2\"\n shift 2\n ;;\n --tokenizer)\n TOKENIZER_ARG=\"$2\"\n shift 2\n ;;\n --data)\n DATA_ARG=\"$2\"\n shift 2\n ;;\n --distributed-config-file)\n DISTRIBUTED_CONFIG_FILE=\"$2\"\n shift 2\n ;;\n -h|--help)\n echo \"Usage: $0 [OPTIONS]\"\n echo \"Options:\"\n echo \" --checkpoint-path PATH Path to Megatron checkpoint\"\n echo \" --tensorboard-logs-path PATH Path to TensorBoard logs\"\n echo \" --tokenizer PATH|MOCK Path to tokenizer model, or 'MOCK' (default: MOCK)\"\n echo \" --data PATH|MOCK Data prefix, or 'MOCK' (default: MOCK)\"\n echo \" --distributed-config-file FILE Path to distributed training config file\"\n echo \" -h, --help Show this help message\"\n exit 0\n ;;\n *)\n echo \"Unknown option: $1\"\n echo \"Use --help for usage information\"\n exit 1\n ;;\n esac\ndone\n\n# Check if checkpoint path exists\nif [ ! -d \"$CHECKPOINT_PATH\" ]; then\n echo \"Error: Checkpoint path does not exist: $CHECKPOINT_PATH\"\n exit 1\nfi\necho \"Checkpoint path exists: $CHECKPOINT_PATH\"\n\n# Check if tensorboard logs path exists\nif [ ! -d \"$TENSORBOARD_LOGS_PATH\" ]; then\n echo \"Warning: TensorBoard logs path does not exist. Creating: $TENSORBOARD_LOGS_PATH\"\n mkdir -p \"$TENSORBOARD_LOGS_PATH\"\nfi\necho \"TensorBoard logs path exists: $TENSORBOARD_LOGS_PATH\"\n\n# NOTE: by default we use 8 GPUs\n# These values will be over-written below with environmental variables\nGPUS_PER_NODE=8\nNUM_NODES=1\nMASTER_ADDR=\"localhost\"\nMASTER_PORT=6000\nNODE_RANK=0\n\n# Load distributed config from file if provided\nif [ -n \"$DISTRIBUTED_CONFIG_FILE\" ]; then\n if [ ! -f \"$DISTRIBUTED_CONFIG_FILE\" ]; then\n echo \"Warning: Distributed config file does not exist: $DISTRIBUTED_CONFIG_FILE\"\n echo \"Continuing with default distributed training settings.\"\n else\n echo \"Loading distributed config from: $DISTRIBUTED_CONFIG_FILE\"\n source \"$DISTRIBUTED_CONFIG_FILE\"\n fi\nfi\n\n# Override with environment variables if set\nGPUS_PER_NODE=${GPUS_PER_NODE:-8}\nNUM_NODES=${NUM_NODES:-1}\nMASTER_ADDR=${MASTER_ADDR:-localhost}\nMASTER_PORT=${MASTER_PORT:-6000}\nNODE_RANK=${NODE_RANK:-0}\nWORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))\n\n# Path to the pretrain_gpt.py script, assuming this script is run from the root of the Megatron-LM repository\nPRETRAIN_SCRIPT_PATH=\"pretrain_gpt.py\"\n\n# Data cache path (useful for both mock and real data)\nDATA_CACHE_PATH=\"${PWD}/benchmark_cache_gpt_oss_20b\"\nmkdir -p \"$DATA_CACHE_PATH\"\n\nDISTRIBUTED_ARGS=(\n --nproc_per_node $GPUS_PER_NODE\n --nnodes $NUM_NODES\n --master_addr $MASTER_ADDR\n --master_port $MASTER_PORT\n --node_rank $NODE_RANK\n)\n\n# NOTE: we only set pipeline parallelism to be the number of GPUs\n# Adjust each value based on your setup.\nTP_SIZE=1 \nEP_SIZE=1 \nPP_SIZE=${WORLD_SIZE} \nMICRO_BATCH_SIZE=1\nGLOBAL_BATCH_SIZE=128\nNUM_LAYERS=12\nDTYPE=\"fp8\"\nSEQ_LENGTH=8192\nMAX_POSITION_EMBEDDINGS=8192\nTRAIN_SAMPLES=1953125000\nLR_DECAY_SAMPLES=1949218748\n\nMODEL_ARGS=(\n --no-masked-softmax-fusion\n --transformer-impl transformer_engine\n --disable-bias-linear\n --untie-embeddings-and-output-weights\n --no-rope-fusion\n --normalization RMSNorm\n --num-layers ${NUM_LAYERS}\n --hidden-size 512\n --ffn-hidden-size 2048\n --num-attention-heads 64\n --group-query-attention\n --num-query-groups 8\n --seq-length ${SEQ_LENGTH}\n --max-position-embeddings ${MAX_POSITION_EMBEDDINGS}\n --use-mcore-models\n --rotary-percent 1.0\n --rope-type rope\n --position-embedding-type rope\n --rotary-base 10000\n --no-bias-gelu-fusion\n --export-force-local-attention\n --no-bias-dropout-fusion\n --quick-geglu\n --glu-linear-offset 1.0\n --softmax-type learnable\n --window-attn-skip-freq 2\n --activation-func-clamp-value 7.0\n --window-size 127,0\n --enable-gpt-oss\n)\n\nMOE_ARGS=(\n --num-experts 4\n --moe-router-topk 2\n --moe-router-load-balancing-type aux_loss\n --moe-aux-loss-coeff 1e-3\n --moe-grouped-gemm\n --moe-token-dispatcher-type alltoall\n --overlap-param-gather\n --overlap-grad-reduce\n --moe-ffn-hidden-size 2048\n --moe-router-dtype fp32\n --moe-z-loss-coeff 1e-3\n --moe-permute-fusion\n)\n\nDATA_ARGS_LIST=()\nif [[ \"$TOKENIZER_ARG\" == \"MOCK\" ]] || [[ \"$DATA_ARG\" == \"MOCK\" ]] || [[ -z \"$TOKENIZER_ARG\" ]]; then\n DATA_ARGS_LIST+=(\n \"--mock-data\"\n \"--tokenizer-type NullTokenizer\"\n \"--vocab-size 128256\" \n \"--data-cache-path ${DATA_CACHE_PATH}\"\n \"--tiktoken-pattern v2\" \n \"--split '99,1,0'\"\n \"--no-create-attention-mask-in-dataloader\"\n \"--no-mmap-bin-files\"\n \"--num-workers 1\"\n )\nelse\n # Settings for real data\n DATA_ARGS_LIST+=(\n \"--data-path $DATA_ARG\"\n \"--tokenizer-type HuggingFaceTokenizer\" \n \"--tokenizer-model $TOKENIZER_ARG\"\n \"--data-cache-path ${DATA_CACHE_PATH}\"\n \"--split '99,1,0'\"\n \"--no-create-attention-mask-in-dataloader\"\n \"--no-mmap-bin-files\"\n \"--num-workers 1\"\n # Note: --vocab-size might be inferred by HuggingFaceTokenizer or might need to be explicit.\n \"--vocab-size 128256\"\n )\nfi\n\nTRAINING_ARGS=(\n --micro-batch-size ${MICRO_BATCH_SIZE}\n --global-batch-size ${GLOBAL_BATCH_SIZE}\n --lr 1.0e-5\n --train-samples ${TRAIN_SAMPLES}\n --lr-decay-samples ${LR_DECAY_SAMPLES}\n --lr-decay-style cosine\n --min-lr 1.0e-6\n --weight-decay 0.1\n --lr-warmup-fraction 0.05\n --clip-grad 1.0\n --bf16\n --use-flash-attn\n --attention-softmax-in-fp32\n --accumulate-allreduce-grads-in-fp32\n --disable-bf16-reduced-precision-matmul\n --recompute-activations\n)\n\nMODEL_PARALLEL_ARGS=(\n --tensor-model-parallel-size ${TP_SIZE}\n --pipeline-model-parallel-size ${PP_SIZE}\n --expert-model-parallel-size ${EP_SIZE}\n --sequence-parallel\n --context-parallel-size 1\n --use-distributed-optimizer\n --fp8-format hybrid\n --fp8-param-gather\n --fp8-amax-compute-algo max\n --fp8-amax-history-len 1024\n)\n \nLOGGING_ARGS=(\n --log-interval 1\n --save-interval 10000\n --eval-interval 50000000\n --eval-iters 0\n --save $CHECKPOINT_PATH\n --tensorboard-dir \"${CHECKPOINT_PATH}/tensorboard\"\n --moe-per-layer-logging\n --no-load-optim\n --no-load-rng\n --log-throughput\n)\n\n# Ensure pretrain_gpt.py is found\nif [ ! -f \"$PRETRAIN_SCRIPT_PATH\" ]; then\n echo \"Error: pretrain_gpt.py not found at $PRETRAIN_SCRIPT_PATH\"\n echo \"Please ensure you are running this script from the root of the Megatron-LM repository, and pretrain_gpt.py is present.\"\n exit 1\nfi\n\npython -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} ${PRETRAIN_SCRIPT_PATH} \\\n ${MODEL_ARGS[@]} \\\n ${MOE_ARGS[@]} \\\n ${DATA_ARGS_LIST[@]} \\\n ${TRAINING_ARGS[@]} \\\n ${MODEL_PARALLEL_ARGS[@]} \\\n ${LOGGING_ARGS[@]}\n"
},
{
"path": "examples/gptoss/03_convert_to_hf.py",
"content": "# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n\n\"\"\"Convert HuggingFace checkpoints to Megatron format.\"\"\"\n\nimport os\nimport argparse\n\nfrom megatron.bridge import AutoBridge\n\ndef _parse_args():\n parser = argparse.ArgumentParser(description=\"Convert Megatron LLMs to HuggingFace format\")\n parser.add_argument(\n \"--hf-model\",\n type=str,\n required=True,\n help=\"HuggingFace model identifier or path to load config from\",\n )\n parser.add_argument(\n \"--megatron-model\",\n type=str,\n required=True,\n help=\"Megatron model identifier or path\",\n )\n parser.add_argument(\n \"--save-path\",\n type=str,\n default=None,\n help=\"Path to save the converted HuggingFace checkpoint\",\n )\n parser.add_argument('--local-rank', '--local_rank', type=int, default=0)\n return parser.parse_args()\n\nif __name__ == \"__main__\":\n args = _parse_args()\n HF_MODEL = args.hf_model\n MEGATRON_MODEL = args.megatron_model\n SAVE_PATH = args.save_path\n WORLD_SIZE = int(os.environ.get(\"WORLD_SIZE\", 1))\n\n if SAVE_PATH is None:\n SAVE_PATH = f\"./huggingface_checkpoints/{MEGATRON_MODEL.replace('/', '_')}\"\n \n print(f\"Converting {MEGATRON_MODEL} to HuggingFace {HF_MODEL} format...\")\n print(f\"Save path: {SAVE_PATH}\")\n \n bridge = AutoBridge.from_hf_pretrained(HF_MODEL, trust_remote_code=True)\n bridge.export_ckpt(\n MEGATRON_MODEL,\n SAVE_PATH,\n )\n \n print(f\"Saved HuggingFace checkpoint to {SAVE_PATH}\")\n"
},
{
"path": "examples/gptoss/README.md",
"content": "# GPT-OSS Training Tutorial\n\n## Step 0: Install Dependencies\n\n### Using Megatron Bridge\n\n[Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)\n\nMegatron 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).\n\nBelow is an example of how to use Megatron-Bridge inside the pytorch container to convert a HuggingFace model checkpoint to Megatron format.\n\nReference: [Megatron-Bridge Dockerfile](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docker/Dockerfile.ci)\n\nInside the [pytorch container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) run the following commands to install Megatron-Bridge:\n```bash\ncd /opt\ngit clone --recursive https://github.com/NVIDIA-NeMo/Megatron-Bridge.git\ncd Megatron-Bridge\n\n# Make sure submodules are initialized (for 3rdparty/Megatron-LM)\ngit submodule update --init --recursive\n\nexport PATH=\"/root/.local/bin:$PATH\"\nexport UV_PROJECT_ENVIRONMENT=/opt/venv\nexport VIRTUAL_ENV=/opt/venv\nexport PATH=\"$UV_PROJECT_ENVIRONMENT/bin:$PATH\"\nexport UV_LINK_MODE=copy\nexport UV_VERSION=\"0.7.2\"\n\n# Install UV\ncurl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh\n\n# Create virtual environment and build the package\nuv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages\n\nuv sync --locked --only-group build\nuv sync --locked --link-mode copy --all-extras --all-groups\n\nuv pip install --no-deps -e .\n\nsource ${UV_PROJECT_ENVIRONMENT}/bin/activate\n```\n\n### Setup Environment\n\n```bash\nexport HOST_MEGATRON_LM_DIR=\"/path/to/your/host/megatron-lm\"\ngit clone https://github.com/NVIDIA/Megatron-LM.git \"$HOST_MEGATRON_LM_DIR\"\ncd \"$HOST_MEGATRON_LM_DIR\"\n```\n\n```bash\nexport HF_TOKEN={your_hf_token_here}\n```\n\n## Step 1: Convert HuggingFace to Megatron (Optional - skip if you already have a Megatron checkpoint)\n\nSet `--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`\n\n```bash\npython3 -m torch.distributed.launch --nproc-per-node=8 examples/gptoss/01_convert_from_hf.py --hf-model openai/gpt-oss-20b\n```\n\n## Step 2: Train from Scratch\n\nTo 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.\n\n### Setup Environment\n\n```bash\n# Change these based on model and directory from previous conversion step\nexport MODEL_DIR_NAME=\"openai_gpt-oss_20b\"\n\nexport HOST_CHECKPOINT_PATH=\"./megatron_checkpoints/${MODEL_DIR_NAME}\"\nexport HOST_TENSORBOARD_LOGS_PATH=\"./tensorboard_logs/${MODEL_DIR_NAME}\"\n```\n\nBy default we will use mock data to train the model in the example below. To use your own data, set the following environment variables:\n\n```bash\n# Optional: For real data\nexport HOST_TOKENIZER_MODEL_PATH=\"/path/to/host/tokenizer.model\"\nexport HOST_DATA_PREFIX=\"/path/to/host/mydata_prefix\"\n```\n\n### Setup Training Configurations\n\nRun 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`.\n\n```bash\ncat > ./distributed_config.env << 'EOF'\nGPUS_PER_NODE=8\nNUM_NODES=1\nMASTER_ADDR=localhost\nMASTER_PORT=6000\nNODE_RANK=0\nEOF\n```\n\n### Run Container with Mounted Volumes\n\n**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. \n\nTo train using mock data, run the following command:\n```bash\nPYTORCH_IMAGE=\"nvcr.io/nvidia/pytorch:25.12-py3\"\n\ndocker run --rm --gpus all --ipc=host --ulimit memlock=-1 \\\n -v \"${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm\" \\\n -v \"${HOST_CHECKPOINT_PATH}:/workspace/checkpoints\" \\\n -v \"${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs\" \\\n -v \"./distributed_config.env:/workspace/megatron-lm/examples/gptoss/distributed_config.env\" \\\n --workdir /workspace/megatron-lm \\\n $PYTORCH_IMAGE \\\n bash examples/gptoss/02_train.sh \\\n --checkpoint-path /workspace/checkpoints \\\n --tensorboard-logs-path /workspace/tensorboard_logs \\\n --distributed-config-file /workspace/megatron-lm/examples/gptoss/distributed_config.env \\\n 2>&1 | tee \"${HOST_TENSORBOARD_LOGS_PATH}/training_mock_$(date +'%y-%m-%d_%H-%M-%S').log\"\n```\n**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.\n\nIf using real data with with the `HOST_TOKENIZER_MODEL_PATH` and `HOST_DATA_PREFIX` environment variables set, run the following command instead:\n\n```bash\nPYTORCH_IMAGE=\"nvcr.io/nvidia/pytorch:25.12-py3\"\n\ndocker run --rm --gpus all --ipc=host --ulimit memlock=-1 \\\n -v \"${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm\" \\\n -v \"${HOST_CHECKPOINT_PATH}:/workspace/checkpoints\" \\\n -v \"${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs\" \\\n -v \"${HOST_TOKENIZER_MODEL_PATH}:/workspace/tokenizer_model\" \\\n -v \"$(dirname \"${HOST_DATA_PREFIX}\"):/workspace/data_dir\" \\\n -v \"./distributed_config.env:/workspace/megatron-lm/examples/gptoss/distributed_config.env\" \\\n --workdir /workspace/megatron-lm \\\n $PYTORCH_IMAGE \\\n bash examples/gptoss/02_train.sh \\\n --checkpoint-path /workspace/checkpoints \\\n --tensorboard-logs-path /workspace/tensorboard_logs \\\n --tokenizer /workspace/tokenizer_model \\\n --data \"/workspace/data_dir/$(basename \"${HOST_DATA_PREFIX}\")\" \\\n --distributed-config-file /workspace/megatron-lm/examples/gptoss/distributed_config.env \\\n 2>&1 | tee \"${HOST_TENSORBOARD_LOGS_PATH}/training_custom_$(date +'%y-%m-%d_%H-%M-%S').log\"\n```\n\n## Step 3: Convert Megatron to HuggingFace\n\nJust 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):\n\n```bash\npython3 -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\n```"
},
{
"path": "examples/inference/README.md",
"content": "### Megatron Core Inference Documentation\nThis guide provides an example for Megatron Core for running model inference. \n\n### Contents\n- [Megatron Core Inference Documentation](#megatron-core-inference-documentation)\n- [Contents](#contents)\n - [1. Quick Start](#1-quick-start)\n - [1.1 Understanding The Code](#11-understanding-the-code)\n - [1.2 Running The Code](#12-running-the-code)\n - [2. Flow of Control In MCore Backend](#2-flow-of-control-in-mcore-backend)\n - [3. Customizing The Inference Pipeline](#3-customizing-the-inference-pipeline)\n - [3.1. Create Your Own Inference Backend](#31-create-your-own-inference-backend)\n - [3.2. Create Your Own Text Generation Controller](#32-create-your-own-text-generation-controller)\n - [3.3. Support Other Models](#33-support-other-models)\n - [3.3. Modify Inference Parameters](#33-modify-inference-parameters)\n - [4. Future work](#4-future-work)\n\n \n\n#### 1. Quickstart\nThis example runs statically-batched inference on a model trained using Megatron Core. The entrypoint is [gpt_static_inference.py](./gpt/gpt_static_inference.py). A similar workflow can be adapted for [gpt_dynamic_inference.py](./gpt/gpt_dynamic_inference.py).\n\n \n\n##### 1.1 Code Walkthrough \n***STEP 1 - Initialize model parallel and other default arguments***\nThe micro batch size defaults to 1. It is not used in tensor-parallelism only, and for pipeline-parallel models it is calculated at runtime. \n```python\n# Initialize Megatron model using the same model provider from training.\n initialize_megatron(\n args_defaults={'no_load_rng': True, 'no_load_optim': True, 'micro_batch_size': 1}\n )\n```\n\n***STEP 2 - Load the model using the model_provider_function***\nThe model provider function supports both MCore and Legacy models. \n\n```python\n # Load the model checkpoint\n model = get_model(model_provider, wrap_with_ddp=False)\n load_checkpoint(model, None, None)\n model.eval()\n model = model[0]\n```\n\n***STEP 3 - Choose an engine***\nText generation requires an inference engine, which includes a scheduler. The default engine is the [Megatron Core engine](../../megatron/core/inference/engine/mcore_engine.py) with a [text generation controller](../../megatron/core/inference/text_generation_controllers/text_generation_controller.py). TRTLLMEngine will be supported in the future.\n```python\n # Create an inference wrapper to setup the model.\n inference_wrapped_model = GPTInferenceWrapper(model, args)\n \n # Define a sampling loop.\n text_generation_controller = TextGenerationController(\n inference_wrapped_model=inference_wrapped_model, \n tokenizer=tokenizer\n )\n \n # Create a static or dynamic inference engine.\n inference_engine = StaticInferenceEngine(\n text_generation_controller=text_generation_controller, \n max_batch_size=args.max_batch_size\n)\n```\n\n***STEP 4 - Run text generation***\nThe [SamplingParams](../../megatron/core/inference/sampling_params.py) class uses suggested defaults. Customize this to change top_p, top_k, number of tokens to generate, etc. The result is returned as a list of [InferenceRequests](../../megatron/core/inference/inference_request.py).\n```python\n results: List[InferenceRequest] = inference_engine.generate(\n prompts=args.prompts, sampling_params=sampling_params\n )\n \n if torch.distributed.get_rank() == 0:\n for idx, result in enumerate(results):\n print(f' ------------- RESULT FOR PROMPT {idx} --------------- ')\n result = {\n 'id': result.request_id,\n 'input_prompt': result.prompt, \n 'generated_text': result.generated_text,\n 'generated_tokens' : result.generated_tokens\n }\n print(result)\n```\n\n \n\n##### 1.2 Running The Code\nAn example Slurm script is shown below. Set the tokenizer paths, inference params, and other settings appropriately. \n\nFor a recap on sampling parameters, refer to [this blog](https://ivibudh.medium.com/a-guide-to-controlling-llm-model-output-exploring-top-k-top-p-and-temperature-parameters-ed6a31313910).\n\n```\n# Slurm cluster settings \nACCOUNT=\nMLM_PATH=/path/to/megatron-lm\nGPT_CKPT=/path/to/gpt/ckpt\nVOCAB_MERGE_FILE_PATH=/path/to/vocab/and/merge/file\nCONTAINER_IMAGE=nvcr.io/ea-bignlp/ga-participants/nemofw-training:23.11\n\nsrun --account $ACCOUNT \\\n--job-name=$ACCOUNT:inference \\\n--partition=batch \\\n--time=01:00:00 \\\n--container-image $CONTAINER_IMAGE \\\n--container-mounts $MLM_PATH:/workspace/megatron-lm/,$GPT_CKPT:/workspace/mcore_gpt_ckpt,$VOCAB_MERGE_FILE_PATH:/workspace/tokenizer \\\n--no-container-mount-home \\\n--pty /bin/bash \\\n\n# Inside the container run the following. \n\ncd megatron-lm/\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\nTOKENIZER_ARGS=(\n --vocab-file /workspace/tokenizer/gpt2-vocab.json\n --merge-file /workspace/tokenizer/gpt2-merges.txt\n --tokenizer-type GPT2BPETokenizer\n)\n\nMODEL_ARGS=(\n --use-checkpoint-args\n --use-mcore-models\n --load /workspace/mcore_gpt_ckpt\n)\n\nINFERENCE_SPECIFIC_ARGS=(\n --attention-dropout 0.0\n --hidden-dropout 0.0\n --num-tokens-to-generate 20\n --max-batch-size 4\n)\n\ntorchrun --nproc-per-node=4 examples/inference/gpt/gpt_static_inference.py \\\n ${TOKENIZER_ARGS[@]} \\\n ${MODEL_ARGS[@]} \\\n ${INFERENCE_SPECIFIC_ARGS[@]} \\\n --prompts \"prompt one \" \"sample prompt two\" \"sample prompt 3\"\n\nNOTE: Other parameters which can be customized for inference:\n--temperature (Sampling temperature)\n--top_k (top_k sampling)\n--top_p (top_p sampling)\n--num-tokens-to-generate (Number of tokens to generate for each prompt)\n--inference-batch-times-seqlen-threshold (During inference, if batch-size times sequence-length is smaller than this threshold then we will not use microbatched pipelining.')\n--use-dist-ckpt (If using dist checkpoint format for the model)\n--use-legacy-models (If using legacy models instead of MCore models)\n\n```\n\n\n \n\n\n#### 2. Control Flow in the MCore Backend\nAn example of inference with static batching is provided in [gpt_static_inference.py](./gpt/gpt_static_inference.py).\n* [mcore_engine](../../megatron/core/inference/engines/mcore_engine.py) **generate()** function is called with the input prompts.\n* The `Scheduler` in the engine will add these prompts to the [active requests] pool (../../megatron/core/inference/inference_request.py) until max batch size is hit. Remaining requests will be added to the waiting requests pool. \n* The engine will run until all requests (waiting + active) are completed. \n * The active requests are passed into **generate_all_output_tokens_static_batch()** of the text generation controller . \n * This function uses the **prep_model_for_inference()** method of the [model_inference_wrappers](../../megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py) and runs an autoregressive sampling loop\n * In the autoregressive loop, the **get_batch_for_context_window()** method of the inference wrapper is called to slice out the input tokens and masks\n * Input tokens and masks are passed it into the **run_one_forward_step()** method, which calls the model `.forward()` method to get the output logits\n * Output logits are synchronized across all pipeline parallel ranks\n * The text generation controller obtains the log probabilities and samples tokens based on the strategy defined in the sampling parameters.\n * The sampled tokens are then appended to the input prompt tokens for the next iteration \n * The **update_generation_status()** method of the text generation controller checks which prompts have finished generating or hit a stop condition\n * After the inference loop, the result is detokenized and stored as an attribute of the InferenceRequest. These requests are marked as completed. \n * The **update_requests_pool()** method of the scheduler moves completed requests into the completed request pool and waiting requests into the active request pool\n\n \n\n#### 3. Customizing The Inference Pipeline\n\nThe inference pipeline supports three levels of customization:\n\n* **Inference engine** - The MCore Engine supports static and dynamic batching. Modify this to add a new backend.\n* **Text generation controller** - The main sampling loop. Customize this to support alternative tokenization or implement a new sampling strategy.\n* **Inference Wrapped Model** - Change this to support a new model.\n* **Modify Inference Parameters** - Change this to update top_p, top_k, number of tokens to be generated, temperature, and other sampling parameters.\n\n \n\n##### 3.1. Create Your Own Inference Backend \nThe [abstract_engine.py](./../../megatron/core/inference/engine/abstract_engine.py) file contains a `generate` method that can be extended to support a new backend. \n\n```python\nclass AbstractEngine(ABC):\n @staticmethod\n def generate(self) -> dict:\n \"\"\"The abstract backend's generate function. \n\n To define a new backend, implement this method and return the outputs as a dictionary. \n```\n\n \n\n##### 3.2. Implement a new Sampling Loop \n\nThe [TextGenerationController](../../megatron/core/inference/text_generation_controllers/text_generation_controller.py) contains the main sampling loop and can be modified to support new tokenization, detokenization, or sampling strategies.\n\n``` python\nclass TextGenerationController:\n\n def tokenize_prompt(self, prompt: str) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"Utility to tokenize the input prompts\"\"\"\n\n def sample_from_logits(\n self,\n last_token_logits: torch.Tensor,\n sampling_params: SamplingParams,\n vocab_size: int,\n generation_started : Optional[torch.Tensor] = None,\n top_n_logprobs_dict: Dict[int, List[Dict[str, float]]] = None,\n ) -> torch.Tensor:\n \"\"\"Samples the logits to generate outputs\n\n Given the logits of the last token, this function samples according to the parameters defined in sampling_params and returns the sampled tokens. If sampling_params.top_n_logprobs > 0 \n at each step it also updates the top_n_logprobs_dict.\n \"\"\"\n\n def update_generation_status(\n self,\n updated_prompts_tokens: torch.Tensor,\n generation_started: torch.Tensor,\n current_context_end_position: int,\n is_generation_done_tensor: torch.Tensor,\n generated_sequence_lengths: torch.Tensor,\n ) -> torch.Tensor:\n \"\"\"Function to check which prompts have reached an end condition\n\n We check which prompts have reached an end condition and set the corresponding flags of the is_generation_done_tensor to True . The generated sequence lengths increases as we keep generating, until that prompts hits an eod condition. The generation started status tensor helps us determine which prompts have started generating\n \"\"\"\n\n def generate_all_output_tokens_static_batch(\n self, active_requests: OrderedDict[int, InferenceRequest],\n ) -> OrderedDict[int, InferenceRequest]:\n \"\"\"Utility to generate all the output tokens and probabilities for the prompts .\n\n This utility generates the output tokens for a static batch. It runs the forward steps till all prompts complete generation, updates the status of these requests to completed, adds the generated result and returns these requests\n \"\"\"\n\n def detokenize_generations(self, prompt_tokens_with_generated_tokens: torch.Tensor) -> str:\n \"\"\"Detokenize the output generations\"\"\"\n```\n\n \n\n##### 3.3. Support Other Models\nExtend [abstract_model_inference_wrapper.py](./../../megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py) to support other models. The abstract model wrapper implements: \n* Forward method which calls the model `forward` method depending on model parallel settings\n* Initializes the model and puts it in `.eval()` mode\n* Setup for the input parameters (max batch size, max seq length) \n\nThe following methods should be implemented: \n```python\nclass AbstractModelInferenceWrapper:\n def prep_model_for_inference(self, prompts_tokens: torch.Tensor):\n \"\"\"A utility function for preparing model for inference\n\n The function gets called once before the auto regressive inference loop. It puts the model in eval mode , and gets some model and inference data parameters. Extend this to build position ids ,attention mask etc, so that required slices can be extracted during the forward pass\n \"\"\"\n\n @abc.abstractclassmethod\n def get_batch_for_context_window(self) -> List:\n \"\"\"Returns the input data for inference \n\n This function gets called iteratively in the inference loop. It can be used to extract relevant input from the prompt tokens, attention mask etc. required for each step in inference.\n```\n\nRefer to [gpt_inference_wrapper.py](../../megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py) for an example of implementing this for GPTModel.\n\n \n\n##### 3.3. Modify Inference Parameters\nWe use [common inference params](../../megatron/core/inference/sampling_params.py) for text generation. Customize this to change `top_p`, `top_k`, number of tokens to generate etc. Other attributes can be added for the inference loop as shown below.\n\n```\nfrom megatron.core.inference.sampling_params import SamplingParams\n\nc = SamplingParams(temperature=0.5)\nc.add_attributes({'min_length':4, 'eod_id':153})\n```\n\n \n\n#### 4. Future work\nThe following features are planned for future releases.\n* TRTLLM Engine support\n* Continuous batching optimizations\n* Speculative decoding"
},
{
"path": "examples/inference/gpt/gpt_dynamic_inference.py",
"content": "# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n\n# pylint: disable=bad-builtin\n\nimport hashlib\nimport io\nimport json\nimport os\nimport sys\nimport warnings\nfrom collections import defaultdict\nfrom typing import Dict, List, Optional\n\nimport torch\nfrom tqdm import tqdm\n\nsys.path.append(\n os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))\n)\n\nfrom examples.inference.gpt.utils import (\n Request,\n build_dynamic_engine_setup_prefix,\n build_requests,\n get_curr_time,\n get_global_peak_memory_stats_bytes,\n)\nfrom megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext\nfrom megatron.core.inference.engines import DynamicInferenceEngine, EngineSuspendedError\nfrom megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (\n GPTInferenceWrapper,\n)\nfrom megatron.core.inference.sampling_params import SamplingParams\nfrom megatron.core.inference.text_generation_controllers.text_generation_controller import (\n TextGenerationController,\n)\nfrom megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer\nfrom megatron.inference.utils import (\n add_inference_args,\n get_inference_config_from_model_and_args,\n get_model_for_inference,\n)\n\nsys.path.append(\n os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))\n)\nimport logging\n\nimport megatron\nfrom megatron.core.utils import configure_nvtx_profiling\nfrom megatron.training import get_args, get_tokenizer, initialize_megatron\n\ntorch.serialization.add_safe_globals([io.BytesIO])\ntorch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState])\ntorch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic])\n\n\ndef run_inference(\n requests: List[Request],\n engine: DynamicInferenceEngine,\n sampling_params: Optional[SamplingParams] = None,\n) -> List[Dict[str, float]]:\n \"\"\"Add requests to engine and generate tokens.\n\n Args:\n requests (List[Request]): Requests that are to be added and processed.\n engine (DynamicInferenceEngine): Inference engine that manages generating tokens.\n sampling_params (SamplingParams): Deprecated as of megatron-core 0.16.\n\n Return:\n A dictionary of step times with `prefill` and `decode` keys.\n \"\"\"\n\n if sampling_params is not None and torch.distributed.get_rank() == 0:\n warnings.warn(\n \"The `sampling_params` argument is deprecated. \"\n \"Sampling parameters are specified per request.\",\n DeprecationWarning,\n )\n\n args = get_args()\n\n # Parse batch boundaries for batch-drain mode.\n batch_ranges = None\n if args.drain_between_batches and args.batch_boundaries:\n boundaries = [int(x) for x in args.batch_boundaries.split(\",\")]\n num_requests_total = len(requests)\n batch_ranges = []\n for i, start in enumerate(boundaries):\n end = boundaries[i + 1] if i + 1 < len(boundaries) else num_requests_total\n batch_ranges.append((start, end))\n\n # Initialize request arrival times.\n base_arrival_time = get_curr_time()\n for request in requests:\n request.time_arrival = request.time_offset + base_arrival_time\n\n # Add and process requests.\n num_requests_total = len(requests)\n num_requests_added = 0\n num_requests_finished = 0\n step_times = {\"prefill\": [], \"decode\": []}\n add_times = []\n output_times = []\n tbar = tqdm(total=num_requests_total)\n total_output_tokens = 0\n attempted_step_count = 0\n if args.cuda_graph_impl == \"local\":\n cuda_graph_request_count_map = {}\n else:\n cuda_graph_request_count_map = None\n\n def _add_request():\n \"\"\"Add request to engine.\n\n *Note: Using `prompt_text` instead of `prompt_tokens` for fair comparison.\n \"\"\"\n nonlocal num_requests_added\n _request = requests[num_requests_added]\n engine.add_request(num_requests_added, _request.prompt_text, _request.sampling_params)\n _request.time_start = get_curr_time()\n _request.state = \"started\"\n num_requests_added += 1\n tbar.update(1)\n\n def _process_step_result(result):\n \"\"\"Process a single engine step result, updating bookkeeping state.\"\"\"\n nonlocal total_output_tokens, num_requests_finished\n\n is_decode_only = engine.is_decode_only\n\n # Record cuda_graph_request_count.\n cuda_graph_request_count = result[\"cuda_graph_request_count\"]\n if args.cuda_graph_impl == \"local\" and cuda_graph_request_count is not None:\n cuda_graph_request_count_map[cuda_graph_request_count] = (\n cuda_graph_request_count_map.get(cuda_graph_request_count, 0) + 1\n )\n\n # Update requests.\n active_request_ids = result[\"active_request_ids\"]\n finished_request_records = result[\"finished_request_records\"]\n step_time = result[\"step_time\"]\n if len(active_request_ids) > 0 or len(finished_request_records) > 0:\n if is_decode_only:\n step_times[\"decode\"].append(step_time)\n else:\n step_times[\"prefill\"].append(step_time)\n\n # Append output tokens.\n output_start = get_curr_time()\n for finished_request_record in finished_request_records:\n\n finished_request = finished_request_record.merge()\n\n # Update local request object.\n request = requests[finished_request.request_id]\n request.time_end = get_curr_time()\n request.state = \"finished\"\n request.request_id = finished_request.request_id\n request.events = finished_request.events\n\n request.ttft = finished_request.ttft\n\n # Update prompt, in case engine has been suspended and resumed.\n request.prompt_tokens = finished_request.prompt_tokens.tolist()\n request.prompt_text = finished_request.prompt\n\n # Get output tokens and text.\n request.output_tokens = finished_request.generated_tokens\n request.output_text = finished_request.generated_text\n total_output_tokens += len(request.output_tokens)\n\n # Log probs.\n if finished_request.sampling_params.return_log_probs:\n if not finished_request.prompt_log_probs:\n finished_request.prompt_log_probs = []\n request.prompt_log_probs = finished_request.prompt_log_probs\n request.generated_log_probs = finished_request.generated_log_probs\n request.logprobs = (\n finished_request.prompt_log_probs + finished_request.generated_log_probs\n )\n if finished_request.sampling_params.top_n_logprobs > 0:\n request.generated_top_n_logprobs = finished_request.generated_top_n_logprobs\n if not finished_request.sampling_params.skip_prompt_log_probs:\n request.prompt_top_n_logprobs = finished_request.prompt_top_n_logprobs\n num_requests_finished += 1\n output_times.append(get_curr_time() - output_start)\n\n if batch_ranges is not None:\n # Batch-drain mode: add all requests in a batch, drain, then next batch.\n for batch_idx, (batch_start, batch_end) in enumerate(batch_ranges):\n # Add all requests in current batch.\n add_start = get_curr_time()\n while num_requests_added < batch_end:\n _add_request()\n add_times.append(get_curr_time() - add_start)\n\n # Step until all active requests finish (drain).\n while engine.has_unfinished_requests():\n try:\n result = engine.step_modern()\n except EngineSuspendedError as e:\n result = e\n attempted_step_count += 1\n\n if isinstance(result, EngineSuspendedError):\n continue\n\n _process_step_result(result)\n else:\n # Original mode: add requests per step based on arrival time or count.\n while True:\n # Add requests.\n add_start = get_curr_time()\n if args.incoming_requests_per_step is None:\n # Add requests with 'earlier' arrival time.\n while num_requests_added < num_requests_total:\n if requests[num_requests_added].time_arrival > add_start:\n break\n _add_request()\n else:\n # Add deterministic number of requests (generally used for debugging).\n for i in range(\n min(args.incoming_requests_per_step, num_requests_total - num_requests_added)\n ):\n _add_request()\n add_times.append(get_curr_time() - add_start)\n\n # Step inference engine (i.e., generate a token for each active request).\n # Before step, we haven't done the scheduling, so we cannot know the is_decode_only\n try:\n result = engine.step_modern()\n except EngineSuspendedError as e:\n result = e\n pass # ignore error in order to call 'engine.resume()' below.\n attempted_step_count += 1\n\n # Test suspending and resuming engine.\n if args.suspend_resume_interval is not None:\n\n # Suspend.\n if attempted_step_count % args.suspend_resume_interval == 0:\n print(\"**** step %d/%d ... suspend.\" % (engine.context.step_count, attempted_step_count))\n engine.suspend()\n\n # Resume, 0+ attempted steps later.\n if (\n attempted_step_count > 0\n and (attempted_step_count - args.suspend_resume_interval // 2)\n % args.suspend_resume_interval\n == 0\n ):\n print(\"**** step %d/%d ... resume.\" % (engine.context.step_count, attempted_step_count))\n engine.resume()\n\n # If engine suspended, continue to next iter.\n if isinstance(result, EngineSuspendedError):\n continue\n\n _process_step_result(result)\n\n # Check if all requests are finished.\n if not (engine.has_unfinished_requests() or num_requests_added < num_requests_total):\n break\n\n # Resume engine (NOOP if not suspended).\n engine.resume()\n\n return {\n \"step_times\": step_times,\n \"add_times\": add_times,\n \"output_times\": output_times,\n \"total_output_tokens\": total_output_tokens,\n \"cuda_graph_request_count_map\": cuda_graph_request_count_map,\n }\n\n\n@torch.inference_mode()\ndef main():\n \"\"\"Run dynamic inference.\"\"\"\n # Initialize Megatron.\n initialize_megatron(\n extra_args_provider=add_inference_args,\n args_defaults={'no_load_rng': True, 'no_load_optim': True},\n )\n\n # Start Nsight profiler.\n if os.environ.get(\"NSIGHT_PREFIX\"):\n torch.cuda.cudart().cudaProfilerStart()\n\n level_str = os.getenv(\"LOG_LEVEL\", \"INFO\").upper()\n level = getattr(logging, level_str, logging.INFO)\n logging.basicConfig(level=level, force=True)\n\n configure_nvtx_profiling(True)\n\n args = get_args()\n\n # Build tokenizer\n tokenizer = build_tokenizer(args)\n\n # Reset peak memory stats so functional tests measure this run and not\n # whatever happened earlier during initialization.\n torch.cuda.reset_peak_memory_stats()\n\n # Sampling params.\n sampling_params = SamplingParams(\n temperature=args.temperature,\n top_k=args.top_k,\n top_p=args.top_p,\n skip_prompt_log_probs=args.skip_prompt_log_probs,\n return_log_probs=args.return_log_probs,\n num_tokens_to_generate=args.num_tokens_to_generate,\n termination_id=args.termination_id if args.termination_id is not None else tokenizer.eod,\n top_n_logprobs=args.top_n_logprobs,\n stop_words=args.stop_words,\n )\n\n model = get_model_for_inference()\n\n # Requests, context, controller.\n requests = build_requests(args, tokenizer, sampling_params)\n inference_config = get_inference_config_from_model_and_args(model, args)\n\n # Calculate max_sequence_length from requests\n max_gen_length = sampling_params.num_tokens_to_generate\n max_context_length = max(len(r.prompt_tokens) for r in requests)\n inference_config.max_sequence_length = max_context_length + max_gen_length\n context = DynamicInferenceContext(model.config, inference_config)\n wrapped_model = GPTInferenceWrapper(model, context)\n controller = TextGenerationController(wrapped_model, tokenizer)\n\n # Validate all context_length's <= max_tokens.\n if not args.enable_chunked_prefill:\n invalid_prompt_length_map = {}\n for request_idx, request in enumerate(requests):\n if len(request.prompt_tokens) > context.max_tokens:\n invalid_prompt_length_map[request_idx] = len(request.prompt_tokens)\n assert (\n not invalid_prompt_length_map\n ), \"request idxs with prompts longer than context.max_tokens: \" \", \".join(\n f\"{k}({v})\" for k, v in invalid_prompt_length_map.items()\n )\n\n # Inference engine.\n engine = DynamicInferenceEngine(controller, context)\n\n setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests)\n print(\"~~~\")\n print(setup_prefix)\n print(\"~~~\")\n\n # Run and time test, optionally `args.inference_repeat_n` times.\n throughputs = []\n for _ in range(args.inference_repeat_n):\n\n # Reset engine.\n engine.reset()\n\n torch.cuda.reset_peak_memory_stats()\n\n # Trial.\n t = get_curr_time()\n result = run_inference(requests, engine)\n step_times = result[\"step_times\"]\n add_times = result[\"add_times\"]\n output_times = result[\"output_times\"]\n total_output_tokens = result[\"total_output_tokens\"]\n torch.cuda.synchronize()\n total_time = get_curr_time() - t\n stats = torch.cuda.memory_stats()\n throughput = total_output_tokens / total_time\n throughputs.append(throughput)\n\n # Validate all requests finished.\n for request in requests:\n assert request.state == \"finished\", f\"request.state == '{request.state}' != 'finished'.\"\n\n peak_mem_stats = get_global_peak_memory_stats_bytes()\n\n # Print unique prompts + outputs.\n if torch.distributed.get_rank() == 0:\n\n def escape_str(s):\n return s.replace(\"\\n\", \"\\\\n\")\n\n print(\"~~~~ Unique prompts + outputs. ~~~~\")\n\n # Map requests by their prompt.\n unique_prompt_map = defaultdict(list)\n for request_idx, request in enumerate(requests):\n unique_prompt_map[request.prompt_text].append(request_idx)\n\n # Print unique prompts + outputs.\n text_hashes = []\n for unique_idx, (prompt_text, request_idxs) in enumerate(unique_prompt_map.items()):\n\n # ---- Prompt summary line ----\n prompt_len = len(requests[request_idxs[0]].prompt_tokens)\n escaped_prompt_text = escape_str(prompt_text)\n print(\n f\"\\n{unique_idx+1}/{len(unique_prompt_map)}\"\n f\"[n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}\"\n )\n\n # ---- Group all outputs for this prompt ----\n output_map = defaultdict(list)\n for idx in request_idxs:\n req = requests[idx]\n output_map[req.output_text].append(idx)\n\n # ---- Print each unique output ----\n for output_text, output_request_idxs in output_map.items():\n evicted = False\n for idx in output_request_idxs:\n for event in requests[idx].events:\n if event.type.name == \"EVICT\":\n evicted = True\n break\n if output_text is not None:\n # Use hash of prompt + generated text in case engine was\n # suspended and resumed, which misaligns boundary between\n # prompt and generated tokens.\n o_hash = hashlib.sha256((prompt_text + output_text).encode()).hexdigest()[:6]\n o_len = len(requests[output_request_idxs[0]].output_tokens)\n escaped_output_text = escape_str(output_text)\n else:\n o_hash = \"--\"\n o_len = 0\n escaped_output_text = \"--\"\n print(\n f\" >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}\"\n f\"{', ' if evicted else ''}] {escaped_output_text}\"\n )\n text_hashes.append(o_hash)\n\n # Write results to JSON. Primarily used for functional testing.\n if args.output_path:\n json_results = {}\n\n # Write every 'n' requests, plus the final request.\n for i, req in enumerate(requests):\n if i % args.output_every_n_results == 0 or i == len(requests) - 1:\n print(f' Attributes of request {i}: {req.__dict__}')\n result_dict = {\n \"input_prompt\": req.prompt_text,\n \"generated_text\": req.output_text,\n \"generated_tokens\": req.output_tokens,\n \"latency\": req.time_end - req.time_start,\n \"ttft\": req.ttft, # Time-to-first-token in seconds\n \"cuda_graph_request_count_map\": result[\"cuda_graph_request_count_map\"],\n \"step_count\": engine.context.step_count,\n \"top_n_logprobs\": getattr(req, 'generated_top_n_logprobs', None),\n \"prompt_top_n_logprobs\": getattr(req, 'prompt_top_n_logprobs', None),\n }\n if req.sampling_params.return_log_probs:\n result_dict[\"prompt_logprobs\"] = getattr(req, 'prompt_log_probs', None)\n result_dict[\"generated_logprobs\"] = getattr(\n req, 'generated_log_probs', None\n )\n result_dict[\"logprobs\"] = getattr(req, 'logprobs', None)\n if args.output_request_events:\n result_dict[\"events\"] = [e.serialize() for e in req.events]\n json_results[req.request_id] = result_dict\n\n # Track system-level throughput as a test / debug metric\n if args.record_throughput:\n json_results[\"throughput\"] = throughputs\n # Attach peak memory metrics; the functional test only validates these\n # if the fields exist in the golden values.\n json_results.update(peak_mem_stats)\n json_results[\"lifetime_prefill_token_count\"] = engine.context.lifetime_prefill_token_count\n\n print(f' Saving results to {args.output_path}')\n with open(args.output_path, \"w\") as fp:\n json.dump(json_results, fp, indent=1)\n\n # Timing results.\n stats = torch.cuda.memory_stats()\n throughput = total_output_tokens / total_time\n print(\"~~~\")\n peak_alloc_gb = stats[\"allocated_bytes.all.peak\"] / 1024**3\n peak_resvd_gb = stats[\"reserved_bytes.all.peak\"] / 1024**3\n\n p_times = step_times[\"prefill\"]\n d_times = step_times[\"decode\"]\n\n p_total = sum(p_times)\n d_total = sum(d_times)\n\n p_count = len(p_times)\n d_count = len(d_times)\n\n p_mean = p_total / p_count\n d_mean = d_total / d_count if d_count != 0 else 0.0\n\n # Commented out for now as the step/add/output times are not calculated correctly.\n # print(\n # f\"{setup_prefix} … \"\n # f\"mem {peak_alloc_gb:.1f}/{peak_resvd_gb:.1f} GB … \"\n # f\"total time: {step_total:.3f}s … \"\n # f\"step time: total {step_total:.3f}s \"\n # f\"[ p {p_total:.3f}s, d {d_total:.3f}s ], \"\n # f\"mean [ p {p_mean:.3f}s, d {d_mean:.3f}s ], \"\n # f\"count [ p {p_count}, d {d_count} ].\"\n # )\n capture_str = f\"{engine.capture_stats['time']:.2f} sec\" if engine.capture_stats else \"--\"\n print(\n f\"{setup_prefix} … \" f\"throughput: {throughput:.3f} tok/s … \",\n f\"total time: {total_time:.3f}s … \"\n f\"mem {peak_alloc_gb:.1f}/{peak_resvd_gb:.1f} GB … \"\n f\"steps: {engine.context.step_count:d} … \"\n f\"capture {capture_str}\",\n )\n print(\"~~~\")\n\n # Stop Nsight profiler.\n if os.environ.get(\"NSIGHT_PREFIX\"):\n torch.cuda.cudart().cudaProfilerStop()\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "examples/inference/gpt/gpt_dynamic_inference_12b.sh",
"content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n\n# Run dynamic batching inference on the 12B GPT model.\n\nset -u\n\n# Libraries.\npip install simpy\npip install sentencepiece\npip install tiktoken\n\n# Environment variables.\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\n# Checkpoint.\n: ${CHECKPOINT_DIR:?\"CHECKPOINT_DIR is not set\"}\n: ${TOKENIZER_MODEL:?\"TOKENIZER_MODEL is not set\"}\n\n# Prompts.\n: ${NUM_TOKENS_TO_PROMPT=\"8 32\"}\n: ${NUM_TOKENS_TO_GENERATE=256}\n: ${INCOMING_REQUESTS_DURATION=10.}\n: ${INCOMING_REQUESTS_PER_SEC=100.}\n\n# Dynamic context.\n: ${BUFFER_SIZE_GB=50.}\n\n# Cuda graphs.\n: ${NUM_CUDA_GRAPHS=16}\n\n# Miscellaneous.\n: ${USE_COORDINATOR=0}\n: ${ENGINE=dynamic}\n: ${EXTRA_ARGS=\"\"}\n# NSIGHT_PREFIX=/path/to/nsight/profile\n\n# Arguments.\nARGS=\" \\\n --no-persist-layer-norm \\\n --apply-layernorm-1p \\\n --no-position-embedding \\\n --group-query-attention \\\n --num-query-groups 8 \\\n --load ${CHECKPOINT_DIR} \\\n --use-checkpoint-args \\\n --untie-embeddings-and-output-weights \\\n --disable-bias-linear \\\n --use-rotary-position-embeddings \\\n --position-embedding-type rope \\\n --rotary-base 1000000 \\\n --rotary-percent 1.0 \\\n --swiglu \\\n --normalization RMSNorm \\\n --attention-dropout 0.0 \\\n --hidden-dropout 0.0 \\\n --exit-duration-in-mins 5740 \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n --num-layers 40 \\\n --hidden-size 5120 \\\n --ffn-hidden-size 14336 \\\n --num-attention-heads 32 \\\n --kv-channels 128 \\\n --seq-length 1024 \\\n --max-position-embeddings 1024 \\\n --micro-batch-size 64 \\\n --bf16 \\\n --tokenizer-type TikTokenizer \\\n --tiktoken-pattern v2 \\\n --tokenizer-model ${TOKENIZER_MODEL} \\\n --distributed-timeout-minutes 2400 \\\n --use-flash-attn \\\n --inference-rng-tracker \\\n \\\n --inference-dynamic-batching \\\n --inference-dynamic-batching-buffer-size-gb ${BUFFER_SIZE_GB} \\\n \\\n ${EXTRA_ARGS} \\\n\"\n\n# Cuda graphs.\nif [ \"${NUM_CUDA_GRAPHS}\" != \"0\" ]; then\n ARGS+=\" \\\n --cuda-graph-impl local \\\n --inference-dynamic-batching-num-cuda-graphs ${NUM_CUDA_GRAPHS} \\\n \"\nelse\n ARGS+=\" \\\n --cuda-graph-impl none \\\n \"\nfi\n\n# Prompts.\nif [[ -v PROMPTS ]]; then\n ARGS+=\" \\\n --prompts ${PROMPTS} \\\n --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \\\n \"\nelif [[ -v PROMPT_FILE ]]; then\n ARGS+=\" \\\n --prompt-file ${PROMPT_FILE} \\\n --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \\\n \"\nelse\n ARGS+=\" \\\n --num-tokens-to-prompt ${NUM_TOKENS_TO_PROMPT} \\\n --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \\\n --incoming-requests-duration ${INCOMING_REQUESTS_DURATION} \\\n --incoming-requests-per-sec ${INCOMING_REQUESTS_PER_SEC} \\\n \"\nfi\n\n# Command.\nif [[ \"${USE_COORDINATOR}\" == \"0\" ]]; then\n CMD=\"python -m examples.inference.gpt.gpt_${ENGINE}_inference ${ARGS}\"\nelse\n CMD=\"python -um examples.inference.gpt.gpt_${ENGINE}_inference_with_coordinator ${ARGS}\"\nfi\n\nif [[ -v NSIGHT_PREFIX ]]; then\n CMD=\"nsys profile -s none -t nvtx,cuda --cudabacktrace=all --cuda-graph-trace=node --python-backtrace=cuda --wait all -o ${NSIGHT_PREFIX} --force-overwrite true --capture-range=cudaProfilerApi --capture-range-end=stop ${CMD}\"\nfi\n\necho \"~~~\"\necho \"CMD ... ${CMD}.\"\necho \"~~~\"\neval ${CMD}\n"
},
{
"path": "examples/inference/gpt/gpt_dynamic_inference_357m.sh",
"content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n\n# Run dynamic batching inference on the 357M GPT model.\n\nset -u\n\n# Libraries.\npip install simpy\npip install sentencepiece\npip install tiktoken\n\n# Environment variables.\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\n# Checkpoint.\n: ${CHECKPOINT_DIR:?\"CHECKPOINT_DIR is not set\"}\n: ${VOCAB_FILE:?\"VOCAB_FILE is not set\"}\n: ${MERGE_FILE:?\"MERGE_FILE is not set\"}\n\n# Prompts.\n: ${NUM_TOKENS_TO_PROMPT=\"8 32\"}\n: ${NUM_TOKENS_TO_GENERATE=256}\n: ${INCOMING_REQUESTS_DURATION=10.}\n: ${INCOMING_REQUESTS_PER_SEC=100.}\n\n# Dynamic context.\n: ${BUFFER_SIZE_GB=50.}\n\n# Cuda graphs.\n: ${NUM_CUDA_GRAPHS=16}\n\n# Miscellaneous.\n: ${USE_COORDINATOR=0}\n: ${ENGINE=dynamic}\n: ${NPROC_PER_NODE=1}\n: ${EXTRA_ARGS=\"\"}\n# NSIGHT_PREFIX=/path/to/nsight/profile\n\n# Arguments.\nARGS=\" \\\n --exit-on-missing-checkpoint \\\n --transformer-impl local \\\n --load ${CHECKPOINT_DIR} \\\n --tokenizer-type GPT2BPETokenizer \\\n --vocab-file ${VOCAB_FILE} \\\n --merge-file ${MERGE_FILE} \\\n --exit-on-missing-checkpoint \\\n --max-position-embeddings 2048 \\\n --seq-length 2048 \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n --num-layers 24 \\\n --num-attention-heads 16 \\\n --hidden-size 1024 \\\n --bf16 \\\n --micro-batch-size 1 \\\n --attention-dropout 0.0 \\\n --hidden-dropout 0.0 \\\n --seed 42 \\\n --use-flash-attn \\\n --inference-rng-tracker \\\n \\\n --inference-dynamic-batching \\\n --inference-dynamic-batching-buffer-size-gb ${BUFFER_SIZE_GB} \\\n \\\n ${EXTRA_ARGS} \\\n\"\n\n# Cuda graphs.\nif [ \"${NUM_CUDA_GRAPHS}\" != \"0\" ]; then\n ARGS+=\" \\\n --cuda-graph-impl local \\\n --inference-dynamic-batching-num-cuda-graphs ${NUM_CUDA_GRAPHS} \\\n \"\nelse\n ARGS+=\" \\\n --cuda-graph-impl none \\\n \"\nfi\n\n# Prompts.\nif [[ -v PROMPTS ]]; then\n ARGS+=\" \\\n --prompts ${PROMPTS} \\\n --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \\\n \"\nelif [[ -v PROMPT_FILE ]]; then\n ARGS+=\" \\\n --prompt-file ${PROMPT_FILE} \\\n --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \\\n \"\nelse\n ARGS+=\" \\\n --num-tokens-to-prompt ${NUM_TOKENS_TO_PROMPT} \\\n --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \\\n --incoming-requests-duration ${INCOMING_REQUESTS_DURATION} \\\n --incoming-requests-per-sec ${INCOMING_REQUESTS_PER_SEC} \\\n \"\nfi\n\n# Command.\nif [[ \"${USE_COORDINATOR}\" == \"0\" ]]; then\n CMD=\"python -m examples.inference.gpt.gpt_${ENGINE}_inference ${ARGS}\"\nelse\n CMD=\"python -m torch.distributed.run --nproc-per-node ${NPROC_PER_NODE} -m examples.inference.gpt.gpt_${ENGINE}_inference_with_coordinator ${ARGS}\"\nfi\n\nif [[ -v NSIGHT_PREFIX ]]; then\n CMD=\"nsys profile -s none -t nvtx,cuda --cudabacktrace=all --cuda-graph-trace=node --python-backtrace=cuda --wait all -o ${NSIGHT_PREFIX} --force-overwrite true --capture-range=cudaProfilerApi --capture-range-end=stop ${CMD}\"\nfi\n\necho \"~~~\"\necho \"CMD ... ${CMD}.\"\necho \"~~~\"\neval ${CMD}\n"
},
{
"path": "examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py",
"content": "# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n\nimport asyncio\nimport json\nimport logging\nimport os\nimport time\nimport warnings\nfrom collections import defaultdict\nfrom typing import List\n\nimport torch\nimport torch.distributed as dist\n\nfrom examples.inference.gpt.utils import Request, build_dynamic_engine_setup_prefix, build_requests\nfrom megatron.core.inference.engines import DynamicInferenceEngine\nfrom megatron.core.inference.engines.dynamic_engine import EngineState\nfrom megatron.core.inference.inference_client import InferenceClient\nfrom megatron.core.inference.inference_request import DynamicInferenceRequestRecord\nfrom megatron.core.inference.sampling_params import SamplingParams\nfrom megatron.inference.utils import (\n add_inference_args,\n get_dynamic_inference_engine,\n get_model_for_inference,\n)\nfrom megatron.training import get_args, get_tokenizer, initialize_megatron\n\n# pylint: disable=line-too-long\n\nlogging.basicConfig(level=logging.INFO, force=True)\n\n\nasync def suspend_resume_cycle(client, engine, args, futures):\n \"\"\"Wait for all in-flight requests, then suspend/train/resume.\"\"\"\n await asyncio.gather(*futures)\n\n client.pause_engines()\n await engine.wait_until(EngineState.PAUSED)\n client.suspend_engines()\n await engine.wait_until(EngineState.SUSPENDED)\n if args.suspend_timeout > 0:\n await asyncio.sleep(args.suspend_timeout)\n client.resume_engines()\n await engine.wait_until(EngineState.RESUMED)\n client.unpause_engines()\n await engine.wait_until(EngineState.RUNNING)\n\n\nasync def main(\n engine: DynamicInferenceEngine,\n requests: List[Request],\n port: int | None = None,\n sampling_params: SamplingParams | None = None,\n):\n if sampling_params is not None:\n warnings.warn(\n \"The `sampling_params` argument is deprecated. \"\n \"Sampling parameters are specified per request.\",\n DeprecationWarning,\n )\n\n # once you call engine.start_listening_to_data_parallel_coordinator,\n # the engine will start accepting requests from the data parallel coordinator.\n # and processing them in an asyncio coroutine.\n # leaving inference_coordinator_port as None will find a free port automatically.\n args = get_args()\n\n dp_addr = await engine.start_listening_to_data_parallel_coordinator(\n inference_coordinator_port=port,\n launch_inference_coordinator=True,\n coordinator_schedule_output_path=args.coordinator_schedule_output_path,\n )\n\n # All ranks agree on the number of suspend/resume cycles from args.\n num_suspend_resume_cycles = len(requests) // args.suspend_resume_interval if args.suspend_resume_interval else 0\n\n # Create client and run example.\n if dist.get_rank() == 0:\n client = InferenceClient(dp_addr, deserialize=True) # submits requests to the inference coordinator\n client.start()\n base_arrival_time = time.time_ns() / 10**9\n for request in requests:\n request.time_arrival = request.time_offset + base_arrival_time\n futures = []\n num_requests_total = len(requests)\n num_requests_added = 0\n next_suspend_at = args.suspend_resume_interval or 0\n cycles_done = 0\n\n while True:\n current_time = time.time_ns() / 10**9\n if args.incoming_requests_per_step is None:\n # Only add requests that have arrived at the current time.\n while (\n num_requests_added < num_requests_total\n and requests[num_requests_added].time_arrival <= current_time\n ):\n request = requests[num_requests_added]\n # These add-request calls will queue up the request on a zmq socket and return\n # instantaneously. They will return an asyncio future which can be awaited for\n # request completion.\n futures.append(client.add_request(request.prompt_text, request.sampling_params))\n num_requests_added += 1\n\n if num_requests_added >= next_suspend_at and cycles_done < num_suspend_resume_cycles:\n await suspend_resume_cycle(client, engine, args, futures)\n cycles_done += 1\n next_suspend_at += args.suspend_resume_interval\n\n else:\n # Add deterministic number of requests (generally used for debugging).\n for i in range(\n min(args.incoming_requests_per_step, num_requests_total - num_requests_added)\n ):\n # Change sampling parameters to force different generation lengths.\n request = requests[num_requests_added]\n n = request.sampling_params.num_tokens_to_generate\n request.sampling_params.num_tokens_to_generate = n + i\n futures.append(client.add_request(request.prompt_text, request.sampling_params))\n num_requests_added += 1\n\n if num_requests_added >= next_suspend_at and cycles_done < num_suspend_resume_cycles:\n await suspend_resume_cycle(client, engine, args, futures)\n cycles_done += 1\n next_suspend_at += args.suspend_resume_interval\n\n if num_requests_added == num_requests_total:\n break\n # Relinquish control since there are no more requests to add at the moment. This allows the engine to run.\n await asyncio.sleep(0)\n\n # While we wait for the requests to complete, the engine runs in the background.\n results: List[DynamicInferenceRequestRecord] = await asyncio.gather(*futures)\n else:\n # Non-rank-0: match the suspend/resume cycles that rank 0 drives.\n for _ in range(num_suspend_resume_cycles):\n await engine.wait_until(EngineState.PAUSED)\n await engine.wait_until(EngineState.SUSPENDED)\n await engine.wait_until(EngineState.RESUMED)\n await engine.wait_until(EngineState.RUNNING)\n\n if dist.get_rank() == 0:\n # Write results to JSON. Primarily used for functional testing.\n if args.output_path:\n json_results = {}\n throughputs = []\n\n for req in results:\n result_dict = {\n \"input_prompt\": req.prompt,\n \"generated_text\": req.generated_text.replace(\"\\n\", \"\\\\n\"),\n \"generated_tokens\": req.generated_tokens,\n \"latency\": req.latency, # InferenceClient populates this field in the returned future.\n }\n if req.sampling_params.return_log_probs:\n result_dict[\"logprobs\"] = req.prompt_log_probs + req.generated_log_probs\n throughput = len(req.generated_tokens) / req.latency\n throughputs.append(throughput)\n if req.routing_indices is not None:\n result_dict[\"routing_indices\"] = req.routing_indices.tolist()\n \n json_results[req.request_id] = result_dict\n throughput_dict = {\"throughput\": throughputs}\n if args.throughput_check_only:\n json_results = throughput_dict\n with open(args.output_path, \"w\") as fp:\n json.dump(json_results, fp, indent=4)\n else:\n print(\"Results:\")\n unique_prompt_map = defaultdict(list)\n for req in results:\n unique_prompt_map[req.prompt].append(req)\n for idx, (prompt_text, reqs) in enumerate(unique_prompt_map.items()):\n print(\n f\"%d/%d. prompt '%s' ... [%d] output '%s'.\"\n % (\n idx,\n len(unique_prompt_map),\n prompt_text.replace(\"\\n\", \"\\\\n\"),\n len(reqs),\n reqs[0].generated_text.replace(\"\\n\", \"\\\\n\"),\n )\n )\n\n # Pause before stopping: STOP requires PAUSED or SUSPENDED state.\n client.pause_engines()\n\n await engine.wait_until(EngineState.PAUSED)\n\n if dist.get_rank() == 0:\n client.stop_engines()\n\n await engine.wait_until(EngineState.STOPPED)\n\n if dist.get_rank() == 0:\n client.shutdown_coordinator()\n client.stop()\n logging.info(f\"Rank: {dist.get_rank()} stopped their engine instance successfully.\")\n\n\nif __name__ == \"__main__\":\n # enable inference mode in the very beginning as some fp8 optimizations\n # check for it.\n with torch.inference_mode():\n initialize_megatron(\n extra_args_provider=add_inference_args,\n args_defaults={'no_load_rng': True, 'no_load_optim': True},\n )\n\n args = get_args()\n tokenizer = get_tokenizer()\n\n # Sampling params.\n sampling_params = SamplingParams(\n temperature=args.temperature,\n top_k=args.top_k,\n top_p=args.top_p,\n return_log_probs=args.return_log_probs,\n num_tokens_to_generate=args.num_tokens_to_generate,\n termination_id=(\n args.termination_id if args.termination_id is not None else tokenizer.eod\n ),\n )\n\n model = get_model_for_inference()\n\n requests = build_requests(args, tokenizer, sampling_params)\n\n engine = get_dynamic_inference_engine(model=model)\n\n if dist.get_rank() == 0:\n setup_prefix = build_dynamic_engine_setup_prefix(args, model, engine.context, requests)\n print(\"~~~\")\n print(setup_prefix)\n print(\"~~~\")\n\n # Start Nsight profiler.\n if os.environ.get(\"NSIGHT_PREFIX\"):\n torch.cuda.cudart().cudaProfilerStart()\n\n asyncio.run(main(engine, requests, args.inference_coordinator_port))\n\n # Stop Nsight profiler.\n if os.environ.get(\"NSIGHT_PREFIX\"):\n torch.cuda.cudart().cudaProfilerStop()\n"
},
{
"path": "examples/inference/gpt/gpt_static_inference.py",
"content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n\nimport os\nimport sys\nimport time\nfrom argparse import Namespace\n\nimport torch\n\nfrom megatron.core.inference.contexts import StaticInferenceContext\nfrom megatron.core.inference.engines import StaticInferenceEngine\nfrom megatron.core.inference.inference_request import InferenceRequest\nfrom megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (\n GPTInferenceWrapper,\n)\nfrom megatron.core.inference.sampling_params import SamplingParams\nfrom megatron.core.inference.text_generation_controllers.text_generation_controller import (\n TextGenerationController,\n)\nfrom megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer\nfrom megatron.core.transformer.module import MegatronModule\n\nsys.path.append(\n os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))\n)\n\nimport asyncio\nimport json\nfrom typing import List\n\nfrom examples.inference.gpt.utils import build_requests\nfrom megatron.inference.utils import add_inference_args, get_model_for_inference\nfrom megatron.training import get_args, get_tokenizer, print_rank_0\nfrom megatron.training.initialize import initialize_megatron\n\n\ndef add_static_inference_args(parser):\n \"\"\"Static inference arguments.\"\"\"\n\n add_inference_args(parser)\n\n group = parser.add_argument_group(title='Static inference')\n group.add_argument(\n \"--max-batch-size\",\n type=int,\n default=None,\n dest=\"max_batch_size\",\n help='Deprecated, use `--inference-max-requests` instead',\n )\n group.add_argument(\"--stream\", action=\"store_true\", default=False, help=\"Stream output tokens\")\n\n return parser\n\n\ndef get_inference_engine(args: Namespace, model: MegatronModule) -> StaticInferenceEngine:\n \"\"\"Utility to get the relevant backend for running inference\n\n This function will automatically choose the TRTLLMBackend when possible, and if not revert to Mcore backend if the user does not specify any backends. TRT LLM Backend is not implmented yet.\n\n Args:\n args (Namespace): The user arguments parsed from command line\n model (MegatronModule): The megatron model .\n\n Returns:\n AbstractBackend: The chosen backend\n \"\"\"\n tokenizer = build_tokenizer(args)\n inference_context = StaticInferenceContext(\n args.inference_max_requests, args.inference_max_seq_length\n )\n inference_wrapped_model = GPTInferenceWrapper(model, inference_context)\n text_generation_controller = TextGenerationController(\n inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer\n )\n engine_kwargs = {\n \"text_generation_controller\": text_generation_controller,\n \"legacy\": args.use_legacy_static_engine,\n }\n if not args.use_legacy_static_engine:\n engine_kwargs[\"buffer_size_gb\"] = args.inference_dynamic_batching_buffer_size_gb\n return StaticInferenceEngine(**engine_kwargs)\n\n\nasync def generate(\n inference_engine: StaticInferenceEngine, sampling_params: SamplingParams, prompts: List[str]\n) -> List[InferenceRequest]:\n async def collect_stream(prompt, request_id, stream_generator):\n print(f\"Request {request_id}: {prompt}\", end=\"\", flush=True)\n prev_idx = 0\n async for output in stream_generator:\n print(output.generated_text[prev_idx:], end=\"\", flush=True)\n prev_idx = len(output.generated_text)\n print()\n\n request_ids: List[int] = [\n inference_engine.add_request(prompt=prompt, sampling_params=sampling_params, streaming=True)\n for prompt in prompts\n ]\n stream_generators = [\n inference_engine.get_stream_generator(request_id) for request_id in request_ids\n ]\n\n tasks = [\n asyncio.create_task(collect_stream(prompt, request_id, stream_generator))\n for (prompt, request_id, stream_generator) in zip(prompts, request_ids, stream_generators)\n ]\n\n await inference_engine.run_engine_async()\n await asyncio.gather(*tasks)\n\n results: List[InferenceRequest] = [\n inference_engine.scheduler.completed_request_pool[request_id] for request_id in request_ids\n ]\n\n return results\n\n\n@torch.inference_mode()\ndef main():\n \"\"\"Main program.\"\"\"\n\n # Note: The default args passed here can be overwritten by using appropriate params (check arguments.py file)\n # Micro batch size is not needed to be set by user. (It is calculated based on inference-batch-times-seqlen-threshold argument)\n initialize_megatron(\n extra_args_provider=add_static_inference_args,\n args_defaults={\n 'no_load_rng': True,\n 'no_load_optim': True,\n 'micro_batch_size': 1,\n 'exit_on_missing_checkpoint': True,\n },\n )\n\n args = get_args()\n\n model = get_model_for_inference()\n\n inference_engine = get_inference_engine(args, model)\n\n sampling_params = SamplingParams(\n temperature=args.temperature,\n top_k=args.top_k,\n top_p=args.top_p,\n return_log_probs=args.return_log_probs,\n num_tokens_to_generate=args.num_tokens_to_generate,\n top_n_logprobs=args.top_n_logprobs,\n )\n\n # Build tokenizer\n tokenizer = build_tokenizer(args)\n\n requests = build_requests(args, tokenizer)\n prompts = [r.prompt_text for r in requests]\n\n if args.cuda_graph_impl == \"local\":\n print(f\"Running warmup for CUDA graphs...\")\n inference_engine.generate(\n prompts=[\"warmup\"], sampling_params=SamplingParams(num_tokens_to_generate=10)\n )\n start_time = time.perf_counter()\n if args.stream:\n results: List[InferenceRequest] = asyncio.run(\n generate(inference_engine, sampling_params, prompts)\n )\n else:\n results: List[InferenceRequest] = inference_engine.generate(\n prompts=prompts, sampling_params=sampling_params\n )\n end_time = time.perf_counter()\n latency = end_time - start_time\n\n if torch.distributed.get_rank() == 0 and args.output_path:\n results_output = {}\n for idx, result in enumerate(results):\n result_dict = {\n 'input_prompt': result.prompt,\n 'generated_text': result.generated_text,\n 'generated_tokens': result.generated_tokens.tolist(),\n 'tpot': result.tpot,\n 'latency': latency,\n }\n if sampling_params.top_n_logprobs > 0:\n result_dict['generated_top_n_logprobs'] = result.generated_top_n_logprobs\n if sampling_params.return_log_probs:\n response_logprobs = result.prompt_log_probs + result.generated_log_probs\n result_dict[\"logprobs\"] = response_logprobs\n results_output[result.request_id] = result_dict\n\n with open(args.output_path, 'w') as f:\n json.dump(results_output, f)\n\n # Print unique prompts + outputs.\n if torch.distributed.get_rank() == 0:\n\n print(\"~~~~ Unique prompts + outputs. ~~~~\")\n\n # Map results by their prompt.\n from collections import defaultdict\n\n unique_prompt_map = defaultdict(list)\n for result_idx, result in enumerate(results):\n unique_prompt_map[result.prompt].append(result_idx)\n\n # Print unique prompts + outputs.\n for unique_idx, (prompt_text, result_idxs) in enumerate(unique_prompt_map.items()):\n result_idx = result_idxs[0]\n result = results[result_idx]\n generated_text = result.generated_text.replace(\"\\n\", \"\\\\n\")\n print(\n f\"{unique_idx}/{len(unique_prompt_map)} [{len(result_idxs)}]. {prompt_text} \"\n f\"... {generated_text}\"\n )\n\n stats = torch.cuda.memory_stats()\n print_rank_0(\n \"static | cg %d | %s | reqs %d [ batch %d ] ... mem %.1f/%.1f ... time %.3f.\"\n % (\n args.cuda_graph_impl == \"local\",\n (\n f\"\"\n if args.prompts\n else \" %s, %d, %.1e, %.1e\"\n % (\n \"(%s)\" % \" \".join(map(str, args.num_tokens_to_prompt)),\n args.num_tokens_to_generate,\n args.incoming_requests_duration,\n args.incoming_requests_per_sec,\n )\n ),\n len(requests),\n args.inference_max_requests,\n stats[\"allocated_bytes.all.peak\"] / (1024**3),\n stats[\"reserved_bytes.all.peak\"] / (1024**3),\n latency,\n )\n )\n # Force immediate process exit to bypass torchrun's atexit NCCL teardown when\n # CUDA graphs have captured collectives (see PyTorch issue #115388). This can\n # sometimes lead to hangs in the atexit handler.\n # We do this only when CUDA graphs are enabled.\n if args.cuda_graph_impl != \"none\":\n print(f\"[main] rank {torch.distributed.get_rank()}: finished\", flush=True)\n os._exit(0)\n else:\n torch.distributed.destroy_process_group()\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "examples/inference/gpt/utils.py",
"content": "# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n\nimport copy\nimport itertools\nimport json\nimport random\nimport time\nfrom argparse import ArgumentParser, Namespace\nfrom functools import partial\nfrom typing import Any, List, Optional\n\nimport torch\nfrom tqdm import tqdm\n\nfrom megatron.core.inference.contexts import DynamicInferenceContext\nfrom megatron.core.inference.contexts.dynamic_context import get_mem_size_str\nfrom megatron.core.inference.inference_request import DynamicInferenceRequest\nfrom megatron.core.inference.sampling_params import SamplingParams\nfrom megatron.core.transformer.module import MegatronModule\nfrom megatron.training import get_args\n\n\ndef get_default_sampling_params(termination_id: int = None):\n return SamplingParams(\n temperature=1.0,\n top_k=1,\n top_p=0.0,\n return_log_probs=False,\n num_tokens_to_generate=30,\n termination_id=termination_id,\n )\n\n\ndef get_curr_time() -> float:\n \"\"\"Get synchronized time across ranks.\"\"\"\n curr_time = torch.cuda.LongTensor([time.time_ns()])\n if torch.distributed.is_initialized():\n torch.distributed.broadcast(curr_time, src=0)\n return curr_time.item() / 10**9\n\n\nclass Request:\n \"\"\"Class to hold attributes for a single request.\n\n A request is initialized with its prompt text. As it is added, processed,\n and completed through the inference engine, the request is populated with its\n start time, end time, and output tokens.\n\n Args:\n prompt_text (str): Prompt text.\n time_offset (float): Artificial time offset for simulating incoming\n requests. This value is later added to the `base_arrival_time` to\n simulate the requests arrival time.\n tokenizer (Any): Tokenizer for tokenizing the prompt.\n \"\"\"\n\n def __init__(\n self,\n prompt_text: str,\n time_offset: float,\n tokenizer: Any,\n sampling_params: SamplingParams = None,\n ):\n self.prompt_text = prompt_text\n self.prompt_tokens = tokenizer.tokenize(prompt_text)\n self.output_text = None\n self.output_tokens = []\n self.time_offset = time_offset\n self.time_arrival = None\n self.time_start = None\n self.time_end = None\n self.ttft = None # Time-to-first-token in seconds\n self.state = \"not-started\"\n self.sampling_params: SamplingParams = (\n sampling_params\n if sampling_params is not None\n else get_default_sampling_params(tokenizer.eod)\n )\n self.sampling_params = copy.deepcopy(self.sampling_params)\n\n def __str__(self) -> str:\n return \"state '%s'; toffset %.1e; prompt len %d; output len %d; '%s'\" % (\n self.state,\n self.time_offset,\n len(self.prompt_tokens),\n len(self.output_tokens),\n self.prompt_text,\n )\n\n\ndef get_time_offsets(\n seed: int | None,\n incoming_requests_per_step: int,\n incoming_requests_per_sec: float,\n num_requests: int,\n) -> list[float]:\n \"\"\"Get example time offsets.\"\"\"\n\n # Time offsets to add all requests at once.\n if incoming_requests_per_step is not None or incoming_requests_per_sec <= 0:\n return [-1] * num_requests\n\n # if num_requests is not None:\n incoming_requests_duration = num_requests / incoming_requests_per_sec\n incoming_requests_duration *= 2 # extra margin, to accomodate time sampling\n\n random.seed(seed)\n\n import simpy # Guard against this import in test case\n\n # Generate random time offsets.\n def arrival(r):\n while True:\n yield env.timeout(random.expovariate(r))\n time_offsets.append(env.now)\n\n time_offsets = []\n env = simpy.Environment()\n env.process(arrival(incoming_requests_per_sec))\n env.run(incoming_requests_duration)\n\n # Ensure at least a single request.\n if len(time_offsets) == 0:\n time_offsets = [0.0]\n\n # Ensure first time is 0.\n time_offsets = [to - time_offsets[0] for to in time_offsets]\n\n # Truncate to num_requests.\n assert len(time_offsets) >= num_requests\n time_offsets = time_offsets[:num_requests]\n\n return time_offsets\n\n\ndef get_cli_requests(\n args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None\n) -> list[Request]:\n\n # Get time offsets.\n t_offsets = get_time_offsets(\n args.seed,\n args.incoming_requests_per_step,\n args.incoming_requests_per_sec,\n len(args.prompts),\n )\n\n # Init requests.\n requests = [Request(p, t, tokenizer, sampling_params) for p, t in zip(args.prompts, t_offsets)]\n return requests\n\n\ndef get_synthetic_requests(\n args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None\n) -> list[Request]:\n \"\"\"Get example requests.\"\"\"\n\n # Get time offsets.\n time_offsets = get_time_offsets(\n args.seed,\n args.incoming_requests_per_step,\n args.incoming_requests_per_sec,\n int(args.incoming_requests_per_sec * args.incoming_requests_duration),\n )\n\n # Build prompts with expected lengths.\n assert (\n len(args.num_tokens_to_prompt) == 2\n and args.num_tokens_to_prompt[1] >= args.num_tokens_to_prompt[0]\n )\n max_prompt_length = args.num_tokens_to_prompt[1]\n max_prompt_text = \"hi \" * max_prompt_length\n max_prompt_tokens = tokenizer.tokenize(max_prompt_text)\n prompt_lengths = [random.randint(*args.num_tokens_to_prompt) for _ in time_offsets]\n prompt_tokens_list = [max_prompt_tokens[:l] for l in prompt_lengths]\n prompt_texts = [tokenizer.detokenize(tt) for tt in prompt_tokens_list]\n\n # Init requests.\n assert len(prompt_texts) == len(time_offsets)\n requests = [\n Request(t, o, tokenizer, sampling_params=sampling_params)\n for t, o in zip(prompt_texts, time_offsets)\n ]\n\n return requests\n\n\ndef get_requests_from_file(\n args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None\n) -> list[Request]:\n \"\"\"Get requests from a file.\"\"\"\n if not args.prompt_file:\n raise ValueError(\"Prompt file is required to read requests from a file.\")\n\n # Load prompts.\n n_prompts = sum(1 for _ in open(args.prompt_file))\n prompts = []\n if sampling_params is None:\n sampling_params = get_default_sampling_params(tokenizer.eod)\n sampling_params_list = []\n with open(args.prompt_file) as f:\n for line in tqdm(f.readlines(), \"read prompt file\", total=n_prompts):\n line_dict = json.loads(line)\n prompts.append(line_dict[\"text\"])\n\n sp = copy.deepcopy(sampling_params)\n if args.num_tokens_from_file:\n sp.num_tokens_to_generate = line_dict[\"chatgpt_output_token_length\"]\n sampling_params_list.append(sp)\n\n if len(prompts) == args.prompt_file_num_truncate:\n break\n\n # Get time offsets.\n time_offsets: list[float] = get_time_offsets(\n args.seed, args.incoming_requests_per_step, args.incoming_requests_per_sec, len(prompts)\n )\n\n # Init requests.\n requests = [\n Request(p, t, tokenizer, sp)\n for p, t, sp in tqdm(\n zip(prompts, time_offsets, sampling_params_list), \"init requests\", total=len(prompts)\n )\n ]\n\n return requests\n\n\ndef build_requests(\n args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None\n) -> list[Request]:\n # Check if we have any prompts (from command line or JSONL)\n if args.prompts:\n if args.prompt_file:\n raise ValueError(\"Cannot use both --prompts and --prompt-file\")\n return get_cli_requests(args, tokenizer, sampling_params)\n elif args.prompt_file:\n return get_requests_from_file(args, tokenizer, sampling_params)\n else:\n return get_synthetic_requests(args, tokenizer, sampling_params)\n\n\ndef get_model_size_str(model):\n n = sum(p.numel() for p in model.parameters())\n for exp, suffix in ((12, \"t\"), (9, \"b\"), (6, \"m\"), (3, \"k\"), (0, \"\")):\n nquery = int(10**exp)\n if n > nquery:\n return \"%d%s\" % (n // nquery, suffix)\n raise Exception(\"something went wrong.\")\n\n\ndef build_dynamic_engine_setup_prefix(\n args: Namespace,\n model: MegatronModule,\n context: DynamicInferenceContext,\n requests: list[DynamicInferenceRequest],\n):\n \"\"\"\n Returns a compact, pipe-separated summary of the dynamic-batching setup.\n\n Example output:\n\n `dynamic | cg True | prompts: synth(16 256), n 1024, g 512, t 1.0e+02 5.0e-01 | bf 4, 1.2 [r 1024, t 8192] | gtd 0.50 [r 512] | reqs 100` # pylint: disable=line-too-long\n\n Args:\n args (Namespace): Command-line arguments for this run.\n context (DynamicInferenceContext): Stores limits such as `max_requests`,\n `max_tokens`, and `gtd_request_count`.\n requests (List[DynamicInferenceRequest]): List of inference requests.\n\n Returns:\n A configuration string for logging.\n \"\"\"\n # CUDA graph config\n if args.cuda_graph_impl == \"local\":\n cg_str = f\"graphs {len(context.cuda_graph_batch_dimensions_list)}\"\n else:\n cg_str = \"--\"\n\n # Unified memory (UVM).\n uvm_str = f\"uvm {int(context.unified_memory_level)}\"\n\n # Prompt description\n prompt_src_str = (\n \"cli\"\n if args.prompts\n else (\n \"file\"\n if args.prompt_file\n else f\"synth({', '.join(map(str, args.num_tokens_to_prompt))})\"\n )\n )\n request_str = (\n f\"requests: {prompt_src_str}, \" f\"n {len(requests):d}, g {args.num_tokens_to_generate:d}, \"\n )\n request_str += (\n f\"dur {args.incoming_requests_duration:.1e} \" f\"r/sec {args.incoming_requests_per_sec:.1e}\"\n if args.incoming_requests_per_step is None\n else f\"r/step {args.incoming_requests_per_step}\"\n )\n\n # Buffer limits config\n buffer_limits_str = (\n f\"bf: {get_mem_size_str(args.inference_dynamic_batching_buffer_size_gb*1024**3)}, \"\n f\"{context.kv_block_allocator.active_count} chunks \"\n f\"[r {context.max_requests}, t {context.max_tokens}]\"\n )\n\n parts = [get_model_size_str(model), \"dynamic\", cg_str, uvm_str, request_str, buffer_limits_str]\n\n return \" | \".join(parts)\n\n\ndef get_global_peak_memory_stats_bytes() -> dict:\n \"\"\"Peak allocated CUDA memory aggregated across ranks (MAX), in bytes.\n\n Uses `torch.cuda.max_memory_allocated()` and assumes peak stats were reset\n before the benchmark run.\n \"\"\"\n peak_alloc = int(torch.cuda.max_memory_allocated())\n if torch.distributed.is_available() and torch.distributed.is_initialized():\n t = torch.tensor([peak_alloc], device=\"cuda\", dtype=torch.int64)\n torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.MAX)\n peak_alloc = int(t[0].item())\n return {\"mem-max-allocated-bytes\": peak_alloc}\n"
},
{
"path": "examples/inference/llama_mistral/huggingface_reference.py",
"content": "import argparse\nfrom transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer\n\n# Set up argument parsing\nparser = argparse.ArgumentParser(description=\"Script for text generation with a specific model and prompt.\")\nparser.add_argument('--prompt', type=str, required=True, help=\"Prompt text to use for text generation\")\nparser.add_argument('--model-path', type=str, required=True, help=\"Path to the Huggingface model checkpoint\")\n\n# Parse command-line arguments\nargs = parser.parse_args()\n\nmodel_path = args.model_path\nprompt = args.prompt\n\nconfig = AutoConfig.from_pretrained(model_path)\ntokenizer = AutoTokenizer.from_pretrained(model_path, config=config)\nmodel = AutoModelForCausalLM.from_pretrained(model_path, config=config).cuda()\n\ninputs = tokenizer(prompt, return_tensors=\"pt\")\nfor key in inputs:\n inputs[key] = inputs[key].cuda()\n# top_k, top_p and do_sample are set for greedy argmax based sampling\n\noutputs = model.generate(**inputs, max_length=100, do_sample=False, top_p=0, top_k=0, temperature=1.0)\nprint(tokenizer.decode(outputs[0], skip_special_tokens=True))"
},
{
"path": "examples/inference/llama_mistral/run_static_inference_llama4_scout.sh",
"content": "#!/bin/bash\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\nexport NVTE_APPLY_QK_LAYER_SCALING=0\n\nDISTRIBUTED_ARGS=\"--nproc_per_node 8 \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr 0.0.0.0 \\\n --master_port 6000\"\n\n# Fill in checkpoint path to Llama 4 Scout to run\nCHECKPOINT=\nPROMPTS=\"What is the capital of France?\"\nTOKENS_TO_GENERATE=4\nMAX_BATCH_SIZE=2\n\nMODEL_ARGS=\" \\\n --micro-batch-size 1 \\\n --bf16 \\\n --no-masked-softmax-fusion \\\n --disable-bias-linear \\\n --untie-embeddings-and-output-weights \\\n --position-embedding-type rope \\\n --no-rope-fusion \\\n --normalization RMSNorm \\\n --swiglu \\\n --num-layers 48 \\\n --hidden-size 5120 \\\n --ffn-hidden-size 16384 \\\n --num-attention-heads 40 \\\n --group-query-attention \\\n --num-query-groups 8 \\\n --qk-layernorm \\\n --num-experts 16 \\\n --moe-ffn-hidden-size 8192 \\\n --moe-router-score-function sigmoid \\\n --moe-router-topk 1 \\\n --moe-router-topk-scaling-factor 1.0 \\\n --moe-shared-expert-intermediate-size 8192 \\\n --moe-aux-loss-coeff 1e-3 \\\n --moe-token-dispatcher-type alltoall \\\n --moe-token-drop-policy probs \\\n --moe-router-load-balancing-type seq_aux_loss \\\n --seq-length 4096 \\\n --max-position-embeddings 4096 \\\n --tokenizer-type HuggingFaceTokenizer \\\n --make-vocab-size-divisible-by 128 \\\n --use-mcore-models \\\n --rotary-interleaved \\\n --rotary-percent 1.0 \\\n --rotary-base 500000 \\\n --rope-scaling-factor 8.0 \\\n --use-rope-scaling \\\n --no-bias-swiglu-fusion \\\n --qk-l2-norm \\\n --moe-apply-probs-on-input \\\n --moe-router-dtype fp64 \\\n\"\n\ntorchrun $DISTRIBUTED_ARGS -m examples.inference.gpt.gpt_static_inference \\\n --load ${CHECKPOINT} \\\n --tokenizer-model unsloth/Llama-4-Scout-17B-16E-Instruct \\\n --dist-ckpt-strictness log_unexpected \\\n --tensor-model-parallel-size 8 \\\n --prompts ${PROMPTS} \\\n --num-tokens-to-generate ${TOKENS_TO_GENERATE} \\\n --max-batch-size ${MAX_BATCH_SIZE} \\\n ${MODEL_ARGS}\n"
},
{
"path": "examples/inference/llama_mistral/run_text_generation_llama3.1.sh",
"content": "#!/bin/bash\n# This example will start serving the Llama3.1-8B model\nexport NCCL_IB_SL=1\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\nexport NVTE_APPLY_QK_LAYER_SCALING=0\n\nDISTRIBUTED_ARGS=\"--nproc_per_node 1 \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr 0.0.0.0 \\\n --master_port 6000\"\n\n# Ensure CHECKPOINT and TOKENIZER_MODEL are provided\nif [ -z \"$1\" ] || [ -z \"$2\" ]; then\n echo \"Error: You must provide CHECKPOINT and TOKENIZER_MODEL as command-line arguments.\"\n echo \"Usage: $0 /path/to/checkpoint /path/to/tokenizer_model\"\n exit 1\nfi\n\n# Assign command-line arguments to variables\nCHECKPOINT=$1\nTOKENIZER_MODEL=$2\n\npip install flask-restful\n\ntorchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \\\n --use-checkpoint-args \\\n --disable-bias-linear \\\n --tokenizer-type HuggingFaceTokenizer \\\n --tokenizer-model ${TOKENIZER_MODEL} \\\n --transformer-impl transformer_engine \\\n --normalization RMSNorm \\\n --group-query-attention \\\n --num-query-groups 8 \\\n --no-masked-softmax-fusion \\\n --attention-softmax-in-fp32 \\\n --attention-dropout 0.0 \\\n --hidden-dropout 0.0 \\\n --untie-embeddings-and-output-weights \\\n --position-embedding-type rope \\\n --rotary-percent 1.0 \\\n --rotary-base 500000 \\\n --use-rope-scaling \\\n --use-rotary-position-embeddings \\\n --swiglu \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n --num-layers 32 \\\n --hidden-size 4096 \\\n --ffn-hidden-size 14336 \\\n --load ${CHECKPOINT} \\\n --num-attention-heads 32 \\\n --max-position-embeddings 131072 \\\n --bf16 \\\n --micro-batch-size 1 \\\n --seq-length 8192\n"
},
{
"path": "examples/inference/llama_mistral/run_text_generation_llama3.sh",
"content": "#!/bin/bash\n# This example will start serving the Llama3-8B model\nexport NCCL_IB_SL=1\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\nexport NVTE_APPLY_QK_LAYER_SCALING=0\n\nDISTRIBUTED_ARGS=\"--nproc_per_node 1 \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr 0.0.0.0 \\\n --master_port 6000\"\n\n# Ensure CHECKPOINT and TOKENIZER_MODEL are provided\nif [ -z \"$1\" ] || [ -z \"$2\" ]; then\n echo \"Error: You must provide CHECKPOINT and TOKENIZER_MODEL as command-line arguments.\"\n echo \"Usage: $0 /path/to/checkpoint /path/to/tokenizer_model\"\n exit 1\nfi\n\n# Assign command-line arguments to variables\nCHECKPOINT=$1\nTOKENIZER_MODEL=$2\n\npip install flask-restful\n\ntorchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \\\n --use-checkpoint-args \\\n --disable-bias-linear \\\n --tokenizer-type HuggingFaceTokenizer \\\n --tokenizer-model ${TOKENIZER_MODEL} \\\n --transformer-impl transformer_engine \\\n --normalization RMSNorm \\\n --group-query-attention \\\n --num-query-groups 8 \\\n --no-masked-softmax-fusion \\\n --attention-softmax-in-fp32 \\\n --attention-dropout 0.0 \\\n --hidden-dropout 0.0 \\\n --untie-embeddings-and-output-weights \\\n --position-embedding-type rope \\\n --rotary-percent 1.0 \\\n --rotary-base 500000 \\\n --use-rotary-position-embeddings \\\n --swiglu \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n --num-layers 32 \\\n --hidden-size 4096 \\\n --ffn-hidden-size 14336 \\\n --load ${CHECKPOINT} \\\n --num-attention-heads 32 \\\n --max-position-embeddings 8192 \\\n --bf16 \\\n --micro-batch-size 1 \\\n --seq-length 8192\n"
},
{
"path": "examples/inference/llama_mistral/run_text_generation_mistral.sh",
"content": "#!/bin/bash\n# This example will start serving the Mistral-7B-v0.3 model\nexport NCCL_IB_SL=1\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\nDISTRIBUTED_ARGS=\"--nproc_per_node 1 \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr 0.0.0.0 \\\n --master_port 6000\"\n\n# Ensure CHECKPOINT and TOKENIZER_MODEL are provided\nif [ -z \"$1\" ] || [ -z \"$2\" ]; then\n echo \"Error: You must provide CHECKPOINT and TOKENIZER_MODEL as command-line arguments.\"\n echo \"Usage: $0 /path/to/checkpoint /path/to/tokenizer_model\"\n exit 1\nfi\n\n# Assign command-line arguments to variables\nCHECKPOINT=$1\nTOKENIZER_MODEL=$2\n\npip install flask-restful\n\ntorchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \\\n --tokenizer-type HuggingFaceTokenizer \\\n --tokenizer-model ${TOKENIZER_MODEL} \\\n --use-checkpoint-args \\\n --apply-layernorm-1p \\\n --transformer-impl transformer_engine \\\n --normalization RMSNorm \\\n --group-query-attention \\\n --num-query-groups 8 \\\n --no-masked-softmax-fusion \\\n --use-flash-attn \\\n --untie-embeddings-and-output-weights \\\n --disable-bias-linear \\\n --position-embedding-type rope \\\n --rotary-percent 1.0 \\\n --rotary-base 1000000 \\\n --swiglu \\\n --ffn-hidden-size 14336 \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n --num-layers 32 \\\n --hidden-size 4096 \\\n --load ${CHECKPOINT} \\\n --num-attention-heads 32 \\\n --max-position-embeddings 4096 \\\n --bf16 \\\n --micro-batch-size 1 \\\n --seq-length 4096 \\\n --seed 101\n"
},
{
"path": "examples/inference/run_text_generation_server_345M.sh",
"content": "#!/bin/bash\n# This example will start serving the 345M model.\nDISTRIBUTED_ARGS=\"--nproc_per_node 1 \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr localhost \\\n --master_port 6000\"\n\nCHECKPOINT=\nVOCAB_FILE=\nMERGE_FILE=\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\npip install flask-restful\n\ntorchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n --num-layers 24 \\\n --hidden-size 1024 \\\n --load ${CHECKPOINT} \\\n --num-attention-heads 16 \\\n --max-position-embeddings 1024 \\\n --tokenizer-type GPT2BPETokenizer \\\n --fp16 \\\n --micro-batch-size 1 \\\n --seq-length 1024 \\\n --vocab-file $VOCAB_FILE \\\n --merge-file $MERGE_FILE \\\n --seed 42\n"
},
{
"path": "examples/inference/run_text_generation_server_345M_8_tensor_parallel.sh",
"content": "#!/bin/bash\n# This example will start serving the 345M model that is partitioned 8 way tensor parallel\nDISTRIBUTED_ARGS=\"--nproc_per_node 8 \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr localhost \\\n --master_port 6000\"\n\nCHECKPOINT=\nVOCAB_FILE=\nMERGE_FILE=\n\npip install flask-restful\n\npython -m torch.distributed.launch $DISTRIBUTED_ARGS tools/run_text_generation_server.py \\\n --tensor-model-parallel-size 8 \\\n --pipeline-model-parallel-size 1 \\\n --num-layers 24 \\\n --hidden-size 1024 \\\n --load ${CHECKPOINT} \\\n --num-attention-heads 16 \\\n --max-position-embeddings 1024 \\\n --tokenizer-type GPT2BPETokenizer \\\n --fp16 \\\n --micro-batch-size 1 \\\n --seq-length 1024 \\\n --vocab-file $VOCAB_FILE \\\n --merge-file $MERGE_FILE \\\n --seed 42\n"
},
{
"path": "examples/inference/t5/simple_t5_batch_inference.py",
"content": "# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.\n\nimport os\nimport sys\nfrom argparse import Namespace\n\nimport torch\n\nimport pretrain_t5\nfrom megatron.core.inference.engines import AbstractEngine, StaticInferenceEngine\nfrom megatron.core.inference.inference_request import InferenceRequest\nfrom megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (\n InferenceWrapperConfig,\n)\nfrom megatron.core.inference.model_inference_wrappers.t5.t5_inference_wrapper import (\n T5InferenceWrapper,\n)\nfrom megatron.core.inference.sampling_params import SamplingParams\nfrom megatron.core.inference.text_generation_controllers.encoder_decoder_text_generation_controller import (\n EncoderDecoderTextGenerationController,\n)\nfrom megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer\nfrom megatron.core.transformer.module import MegatronModule\nfrom pretrain_t5 import model_provider\n\nsys.path.append(\n os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))\n)\n\nfrom typing import List\n\nfrom megatron.core import mpu\nfrom megatron.training import get_args, get_model, get_tokenizer\nfrom megatron.training.checkpointing import load_checkpoint\nfrom megatron.training.initialize import initialize_megatron\n\n\ndef add_text_generate_args(parser):\n \"\"\"Text generation arguments.\"\"\"\n group = parser.add_argument_group(title='text generation')\n\n group.add_argument(\"--temperature\", type=float, default=1.0, help='Sampling temperature.')\n group.add_argument(\"--top_k\", type=int, default=1, help='Top k sampling.')\n group.add_argument(\"--top_p\", type=float, default=0.0, help='Top p sampling.')\n group.add_argument(\n \"--return-log-probs\",\n action='store_true',\n default=False,\n help='Return the log probabilities of the final output tokens',\n )\n group.add_argument(\n \"--num-tokens-to-generate\",\n type=int,\n default=30,\n help='Number of tokens to generate for each prompt',\n )\n group.add_argument(\n \"--encoder-prompts\",\n metavar='N',\n type=str,\n nargs='+',\n help='Encoder input prompts with each prompt within quotes and separated by space',\n )\n group.add_argument(\n \"--max-batch-size\", type=int, default=1, help='Max number of prompts to process at once'\n )\n return parser\n\n\ndef get_inference_engine(args: Namespace, model: MegatronModule) -> AbstractEngine:\n \"\"\"Utility to get the relevant backend for running inference\n\n This function will automatically chose the TRTLLMBackend when possible, and if not revert to Mcore backend if the user does not specify any backends. TRT LLM Backend is not implmented yet.\n\n Args:\n args (Namespace): The user arguments parsed from command line\n model (MegatronModule): The megatron model .\n\n Returns:\n AbstractBackend: The chosen backend\n \"\"\"\n # Build tokenizer\n tokenizer = build_tokenizer(args)\n\n inference_wrapper_config = InferenceWrapperConfig(\n hidden_size=args.hidden_size,\n inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold,\n fp32_residual_connection=args.fp32_residual_connection,\n params_dtype=args.params_dtype,\n padded_vocab_size=args.padded_vocab_size,\n )\n\n inference_wrapped_model = T5InferenceWrapper(model, inference_wrapper_config)\n text_generation_controller = EncoderDecoderTextGenerationController(\n inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer\n )\n return StaticInferenceEngine(\n text_generation_controller=text_generation_controller, max_batch_size=args.max_batch_size\n )\n\n\ndef main():\n \"\"\"Main program.\"\"\"\n\n # Note: The default args passed here can be overwritten by using appropriate params (check arguments.py file)\n # Micro batch size is not needed to be set by user. (It is calculated based on inference-batch-times-seqlen-threshold argument)\n initialize_megatron(\n extra_args_provider=add_text_generate_args,\n args_defaults={\n 'no_load_rng': True,\n 'no_load_optim': True,\n 'micro_batch_size': 1,\n 'exit_on_missing_checkpoint': True,\n },\n )\n\n # Set up model and load checkpoint\n model = get_model(model_provider, wrap_with_ddp=False)\n load_checkpoint(model, None, None)\n model = model[0]\n\n args = get_args()\n\n inference_engine = get_inference_engine(args, model)\n\n sampling_params = SamplingParams(\n temperature=args.temperature,\n top_k=args.top_k,\n top_p=args.top_p,\n return_log_probs=args.return_log_probs,\n num_tokens_to_generate=args.num_tokens_to_generate,\n )\n\n # Build tokenizer\n tokenizer = build_tokenizer(args)\n\n decoder_prompts = [\"\"] * len(\n args.encoder_prompts\n ) # for T5, the prompt is provided as encoder input, hence decoder_prompts is empty\n args.prompts = decoder_prompts\n\n results: List[InferenceRequest] = inference_engine.generate(\n prompts=args.prompts,\n add_BOS=True,\n encoder_prompts=args.encoder_prompts,\n sampling_params=sampling_params,\n )\n\n if torch.distributed.get_rank() == 0:\n for idx, result in enumerate(results):\n print(f' \\n------------- RESULT FOR PROMPT {idx} --------------- ')\n result = {\n 'id': result.request_id,\n 'input_prompt': result.prompt,\n 'generated_text': result.generated_text,\n 'generated_tokens': result.generated_tokens,\n }\n print(result)\n\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "examples/llama/README.md",
"content": "# Llama Models\n\n## Table of contents\n- [1. Overview](#1-overview)\n- [2. Prerequisites](#2-prerequisites)\n- [3. Training Setup](#3-training-setup)\n- [4. Configuration](#4-configuration)\n- [5. Test Datasets](#5-test-datasets)\n- [6. FP8 Debugging](#6-fp8-debugging)\n\n## 1. Overview\n\n\nTrain Llama models using FP8 precision with Megatron-Core.\n\n## 2. Prerequisites\n\n\n```bash\n# Clone repository\nexport HOST_MEGATRON_LM_DIR=\"/path/to/your/host/megatron-lm\"\ngit clone https://github.com/NVIDIA/Megatron-LM.git \"$HOST_MEGATRON_LM_DIR\"\ncd \"$HOST_MEGATRON_LM_DIR\"\ngit checkout \"core_r0.12.0\"\n\n# Set paths\nexport HOST_CHECKPOINT_PATH=\"./checkpoints/llama3_8b_fp8\"\nexport HOST_TENSORBOARD_LOGS_PATH=\"./tensorboard_logs/llama3_8b_fp8\"\n\n# Optional: For real data\n# export HOST_TOKENIZER_MODEL_PATH=\"/path/to/host/tokenizer.model\"\n# export HOST_DATA_PREFIX=\"/path/to/host/mydata_prefix\"\n```\n\n## 3. Training Setup\n\n\n### Using Mock Data\n```bash\nPYTORCH_IMAGE=\"nvcr.io/nvidia/pytorch:25.03-py3\"\n\ndocker run --rm --gpus all --ipc=host --ulimit memlock=-1 \\\n -v \"${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm\" \\\n -v \"${HOST_CHECKPOINT_PATH}:/workspace/checkpoints\" \\\n -v \"${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs\" \\\n --workdir /workspace/megatron-lm \\\n $PYTORCH_IMAGE \\\n bash examples/llama/train_llama3_8b_h100_fp8.sh \\\n /workspace/checkpoints \\\n /workspace/tensorboard_logs \\\n 2>&1 | tee \"${HOST_TENSORBOARD_LOGS_PATH}/training_mock_$(date +'%y-%m-%d_%H-%M-%S').log\"\n```\n\n### Using Custom Data and Tokenizer\n```bash\nPYTORCH_IMAGE=\"nvcr.io/nvidia/pytorch:25.03-py3\"\n\ndocker run --rm --gpus all --ipc=host --ulimit memlock=-1 \\\n -v \"${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm\" \\\n -v \"${HOST_CHECKPOINT_PATH}:/workspace/checkpoints\" \\\n -v \"${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs\" \\\n -v \"${HOST_TOKENIZER_MODEL_PATH}:/workspace/tokenizer_model\" \\\n -v \"$(dirname \"${HOST_DATA_PREFIX}\"):/workspace/data_dir\" \\\n --workdir /workspace/megatron-lm \\\n $PYTORCH_IMAGE \\\n bash examples/llama/train_llama3_8b_h100_fp8.sh \\\n /workspace/checkpoints \\\n /workspace/tensorboard_logs \\\n /workspace/tokenizer_model \\\n \"/workspace/data_dir/$(basename \"${HOST_DATA_PREFIX}\")\" \\\n 2>&1 | tee \"${HOST_TENSORBOARD_LOGS_PATH}/training_custom_$(date +'%y-%m-%d_%H-%M-%S').log\"\n```\n\n## 4. Configuration\n\n\nDefault parallelism strategy:\n- Tensor Parallel: 1\n- Pipeline Parallel: 1\n- Context Parallel: 2\n\nLlama-3-8B architecture:\n- 32 layers\n- Hidden size: 4096\n- FFN hidden size: 14336\n- Attention heads: 32\n- Query groups: 8\n- Sequence length: 8192\n- RMSNorm normalization with SwiGLU and RoPE\n\nKey training parameters:\n- Micro-batch size: 1\n- Global batch size: 128\n- Learning rate: 1.5e-4\n- Min learning rate: 1.0e-5\n- Weight decay: 0.1\n- FP8 format: hybrid\n\nYou can modify these parameters directly in the `train_llama3_8b_h100_fp8.sh` script.\n\nThis configuration follows those defined in NeMo Framework's performance scripts, which can be found at [https://github.com/NVIDIA/NeMo/tree/main/scripts/performance](https://github.com/NVIDIA/NeMo/tree/main/scripts/performance). \n\n### FP8 Performance\n\n| Model | #-GPUs | GBS | MBS | Seq Length | TP | PP | CP | VP | EP | GA | Tokens/sec/GPU | TFLOP/sec/GPU |\n|-------|--------|-----|-----|------------|----|----|----|----|----|----|----------------|---------------|\n| LLAMA3-8B | 8 | 128 | 1 | 8192 | 1 | 1 | 2 | 1 | 1 | 32 | 13812 | 800 |\n| LLAMA3-70B | 64 | 128 | 1 | 8192 | 4 | 8 | 1 | 5 | 1 | 64 | 1621 | 780 |\n| LLAMA3-405B | 1024 | 512 | 1 | 8192 | 8 | 8 | 2 | 8 | 1 | 64 | 315 | 834 |\n\nLegend:\n- GBS: Global Batch Size\n- MBS: Micro Batch Size\n- TP: Tensor Parallel size\n- PP: Pipeline Parallel size\n- CP: Context Parallel size\n- VP: Virtual Pipeline stages\n- EP: Expert Parallel size\n- GA: Gradient Accumulation steps\n\nAs NeMo uses Megatron-Core, for the latest performance benchmarks, please refer to the official [NeMo documentation](https://docs.nvidia.com/nemo-framework/user-guide/latest/performance/performance-summary.html).\n\n## 5. Test Datasets\n\n\nRecommended datasets:\n1. **WikiText-103**: https://huggingface.co/datasets/Salesforce/wikitext\n\nPreprocess datasets:\n```bash\npython \"${HOST_MEGATRON_LM_DIR}/tools/preprocess_data.py\" \\\n --input your_dataset.json \\\n --output-prefix test_dataset \\\n --tokenizer-type HuggingFaceTokenizer \\\n --tokenizer-model /path/to/tokenizer.model \\\n --append-eod\n```\n\n## 6. FP8 Training Considerations\n\n\n- **Hardware**: Requires NVIDIA Hopper, Ada, or Blackwell GPUs for FP8 support\n \n- **Troubleshooting**: If you encounter NaN values or instability with FP8 training, please refer to [Transformer Engine](https://github.com/NVIDIA/TransformerEngine).\n"
},
{
"path": "examples/llama/train_llama3_8b_h100_fp8.sh",
"content": "#!/bin/bash\n\n# Environment variables for performance tuning\nexport CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}\n#export LOG_LEVEL=${LOG_LEVEL:-INFO}\n#export NCCL_IB_TIMEOUT=${NCCL_IB_TIMEOUT:-19}\n#export NVTE_FWD_LAYERNORM_SM_MARGIN=${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}\n#export NVTE_BWD_LAYERNORM_SM_MARGIN=${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}\n#export NCCL_P2P_NET_CHUNKSIZE=${NCCL_P2P_NET_CHUNKSIZE:-2097152}\n#export NCCL_AVOID_RECORD_STREAMS=${NCCL_AVOID_RECORD_STREAMS:-1}\n\nCHECKPOINT_PATH=${1:-\"checkpoints/llama3_8b_fp8\"}\nTENSORBOARD_LOGS_PATH=${2:-\"tensorboard_logs/llama3_8b_fp8\"}\nTOKENIZER_ARG=${3:-\"MOCK\"} # Path to tokenizer model, or \"MOCK\"\nDATA_ARG=${4:-\"MOCK\"} # Data prefix, or \"MOCK\"\n\n# Create directories if they don't exist\nmkdir -p \"$(dirname \"$CHECKPOINT_PATH\")\"\nmkdir -p \"$(dirname \"$TENSORBOARD_LOGS_PATH\")\"\n\n# Distributed training setup\nGPUS_PER_NODE=8\nNUM_NODES=1\nMASTER_ADDR=${MASTER_ADDR:-localhost}\nMASTER_PORT=${MASTER_PORT:-6000}\nNODE_RANK=${NODE_RANK:-0}\nWORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))\n\n# Path to the pretrain_gpt.py script, assuming this script is run from the root of the Megatron-LM repository\nPRETRAIN_SCRIPT_PATH=\"pretrain_gpt.py\"\n\n# Fixed model and training parameters\nTP_SIZE=1 \nCP_SIZE=1 \nPP_SIZE=1 \nMICRO_BATCH_SIZE=1\nGLOBAL_BATCH_SIZE=128\nNUM_LAYERS=32 \nDTYPE=\"fp8\"\nSEQ_LENGTH=8192\nMAX_POSITION_EMBEDDINGS=8192\n\n# Data cache path (useful for both mock and real data)\nDATA_CACHE_PATH=\"${PWD}/benchmark_cache_llama3_8b_fp8\"\nmkdir -p \"$DATA_CACHE_PATH\"\n\nDISTRIBUTED_ARGS=(\n --nproc_per_node $GPUS_PER_NODE\n --nnodes $NUM_NODES\n --node_rank $NODE_RANK\n --master_addr $MASTER_ADDR\n --master_port $MASTER_PORT\n)\n\nMODEL_ARGS=(\n --use-mcore-models\n --num-layers $NUM_LAYERS\n --hidden-size 4096\n --ffn-hidden-size 14336\n --num-attention-heads 32\n --group-query-attention\n --num-query-groups 8\n --kv-channels 128\n --seq-length $SEQ_LENGTH\n --max-position-embeddings $MAX_POSITION_EMBEDDINGS\n --position-embedding-type rope\n --rotary-base 1000000 \n --rotary-percent 1.0\n --attention-dropout 0.0\n --hidden-dropout 0.0\n --swiglu\n --normalization RMSNorm\n --init-method-std 0.0134\n --attention-backend fused\n --apply-layernorm-1p \n --untie-embeddings-and-output-weights\n --disable-bias-linear \n)\n\nTRAINING_ARGS=(\n --micro-batch-size $MICRO_BATCH_SIZE\n --global-batch-size $GLOBAL_BATCH_SIZE\n --train-samples 1953125000\n --lr-decay-samples 1949218748\n --lr-warmup-samples 3906252\n --lr 0.00015\n --min-lr 0.00001\n --decoupled-lr 5.0e-4 # Specific to decoupled AdamW, ensure optimizer is compatible\n --decoupled-min-lr 4.5e-5 # Specific to decoupled AdamW\n --lr-decay-style cosine\n --clip-grad 1.0\n --weight-decay 0.1\n --adam-beta1 0.9\n --adam-beta2 0.95\n --bf16\n --grad-reduce-in-bf16\n --cross-entropy-loss-fusion\n --calculate-per-token-loss \n --manual-gc \n --empty-unused-memory-level 1 \n --exit-duration-in-mins 235 \n)\n\n# Conditional arguments based on DTYPE (FP8)\nDTYPE_ARGS=()\nif [[ \"$DTYPE\" == \"fp8\" ]]; then\n DTYPE_ARGS+=(\n \"--fp8-format hybrid\"\n \"--fp8-amax-history-len 1024\"\n \"--fp8-amax-compute-algo max\"\n \"--fp8-param-gather\"\n )\nfi\n\n# Model parallelism arguments\nMODEL_PARALLEL_ARGS=(\n --tensor-model-parallel-size $TP_SIZE\n --context-parallel-size $CP_SIZE\n # --pipeline-model-parallel-size $PP_SIZE # Not explicitly set in llama script options, assume 1 if not multi-node PP\n --sequence-parallel # Always enable sequence parallelism with TP_SIZE=2\n)\n\n# Distributed Data Parallel (DDP) arguments\n# From original script's ddp_args\nDDP_ARGS=(\n --use-distributed-optimizer\n --overlap-grad-reduce\n --overlap-param-gather\n)\nTRAINING_ARGS+=(\"${DDP_ARGS[@]}\")\n\n\n# Data arguments (conditional for mock vs real data)\nDATA_ARGS_LIST=()\nif [[ \"$TOKENIZER_ARG\" == \"MOCK\" ]] || [[ \"$DATA_ARG\" == \"MOCK\" ]] || [[ -z \"$TOKENIZER_ARG\" ]]; then\n DATA_ARGS_LIST+=(\n \"--mock-data\"\n \"--tokenizer-type NullTokenizer\"\n \"--vocab-size 128256\" \n \"--data-cache-path ${DATA_CACHE_PATH}\"\n \"--tiktoken-pattern v2\" \n \"--split '99,1,0'\"\n \"--no-create-attention-mask-in-dataloader\"\n \"--no-mmap-bin-files\"\n \"--num-workers 1\"\n )\nelse\n # Settings for real data\n DATA_ARGS_LIST+=(\n \"--data-path $DATA_ARG\"\n \"--tokenizer-type HuggingFaceTokenizer\" \n \"--tokenizer-model $TOKENIZER_ARG\"\n \"--data-cache-path ${DATA_CACHE_PATH}\"\n \"--split '99,1,0'\"\n \"--no-create-attention-mask-in-dataloader\"\n \"--no-mmap-bin-files\"\n \"--num-workers 1\"\n # Note: --vocab-size might be inferred by HuggingFaceTokenizer or might need to be explicit.\n \"--vocab-size 128256\"\n )\nfi\n\nEVAL_AND_LOGGING_ARGS=(\n --log-interval 1\n --eval-iters 32\n --eval-interval 100\n --save-interval 1000\n --log-throughput\n --profile\n --profile-step-start 4\n --profile-step-end 6\n --ckpt-format torch_dist \n --distributed-timeout-minutes 60\n --save \"$CHECKPOINT_PATH\"\n --load \"$CHECKPOINT_PATH\" \n --tensorboard-dir \"$TENSORBOARD_LOGS_PATH\"\n)\n\n# Ensure pretrain_gpt.py is found\nif [ ! -f \"$PRETRAIN_SCRIPT_PATH\" ]; then\n echo \"Error: pretrain_gpt.py not found at $PRETRAIN_SCRIPT_PATH\"\n echo \"Please ensure you are running this script from the root of the Megatron-LM repository, and pretrain_gpt.py is present.\"\n exit 1\nfi\n\n# Run the training command\ntorchrun ${DISTRIBUTED_ARGS[@]} \\\n \"$PRETRAIN_SCRIPT_PATH\" \\\n ${MODEL_ARGS[@]} \\\n ${TRAINING_ARGS[@]} \\\n ${DTYPE_ARGS[@]} \\\n ${MODEL_PARALLEL_ARGS[@]} \\\n ${DATA_ARGS_LIST[@]} \\\n ${EVAL_AND_LOGGING_ARGS[@]}\n\nset +x"
},
{
"path": "examples/mamba/.gitignore",
"content": "checkpoints/\ndata-cache/\ntensorboard/\ntriton-cache/\n"
},
{
"path": "examples/mamba/Dockerfile",
"content": "FROM nvcr.io/nvidia/pytorch:24.01-py3\n\nRUN pip uninstall -y triton && \\\n pip install triton==2.1.0 sentencepiece==0.1.99 flask-restful\n\n# The causal-conv1d and mamba-ssm packages below are built from scratch here\n# (which takes significant time) because there are no wheels available on PyPI\n# for these relatively newer versions of the packages that are compatible with\n# the older NGC-variant PyTorch version (e.g. version 2.2.0.dev231106) that we\n# are using (in the NGC base container). Generally, if the package is not\n# compatible with the PyTorch version, then it will generate a Python import\n# error. The package authors tend to only release wheels for new versions of\n# these pacakges which are compatible with the versions of regular PyTorch and\n# NGC-variant PyTorch that are newer at the time of release. So, to use newer\n# versions of these packages with relatively older versions of the NGC PyTorch\n# container, we tend to have to build the packages from scratch.\n\nRUN cd /tmp && \\\n git clone https://github.com/Dao-AILab/causal-conv1d.git && \\\n cd causal-conv1d && \\\n git checkout v1.2.2.post1 && \\\n CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install . && \\\n cd .. && \\\n rm -rf causal-conv1d\n\nRUN cd /tmp && \\\n git clone https://github.com/state-spaces/mamba.git && \\\n cd mamba && \\\n git checkout v2.0.3 && \\\n MAMBA_FORCE_BUILD=TRUE pip install . && \\\n cd .. && \\\n rm -rf mamba\n"
},
{
"path": "examples/mamba/README.md",
"content": "# Mamba-based Language Models\n\n## Introduction\n\nThis document is an entrypoint into the code used for\n[An Empirical Study of Mamba-based Language Models](https://arxiv.org/abs/2406.07887).\n\nWe are releasing the parameters for some of the models described in that\ntechnical report via\n[HuggingFace](https://huggingface.co/collections/nvidia/ssms-666a362c5c3bb7e4a6bcfb9c).\nThe code in the `main` branch is no longer compatible with the `Mamba2-*`\ncheckpoints. You can load them using the\n[fixed snapshot of the code used for the technical report](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba).\n\n## Installation\n\nCreate and run a Docker container using the [Dockerfile](./Dockerfile).\n\n```\ndocker build -t your_image_name:your_tag .\ndocker run --gpus all -it --rm \\\n -v /path/to/megatron:/workspace/megatron \\\n -v /path/to/dataset:/workspace/dataset \\\n -v /path/to/checkpoints:/workspace/checkpoints \\\n -w /workspace/megatron/examples/mamba \\\n your_image_name:your_tag\n```\n\n## Train\n\n[`train.sh`](./train.sh) is an example pretraining script, showing how to run on\na single node. Select between 800M-scale and 8B-scale models by setting the\n`MODEL_SCALE` variable. The 8B-scale hybrid model architecture is the same as\nthe one described in the technical report.\n\n## Text Generation\n\nUse [`run_text_gen_server_8b.sh`](./run_text_gen_server_8b.sh) to start a text\ngeneration server using an 8B hybrid checkpoint. This is configured to run the\n8B hybrid model described in the technical report, with tensor model parallel\nset to 1.\n\nThe arguments in the script will need to be changed if using a checkpoint with a\ndifferent model parallel configuration or other differences, such as model\narchitecture. For example, to run the 8B pure Mamba-2 model, change\n`--hybrid-layer-pattern` to use only `M` symbols (e.g., 56 `M`s for the 8B\nmodel), or remove it entirely.\n\nUse [`run_text_gen_server_8b_gpt3.sh`](./run_text_gen_server_8b_gpt3.sh) to start\na text generation server using the 8B reference Transformer checkpoint.\n\n## Checkpoint Formats\n\nFor inference, the model must be configured to match the checkpoint file used,\nincluding the hybrid layer configuration and model parallel configuration.\n\nIf you need to convert a hybrid checkpoint file to a different tensor parallel\nor pipeline parallel size, use\n[the hybrid conversion script](../../tools/checkpoint/hybrid_conversion.py).\nThere is an example run command at the end of that file.\n\nBefore running that script, you will need to set `PYTHONPATH` to include the\nroot directory of your Megatron-LM repository clone.\n\n```\nexport PYTHONPATH=:PYTHONPATH\n```\n\n## Hybrid Options\n\n`--hybrid-layer-pattern PATTERN` specifies the layer type for every layer in\nthe model using a string of single-character symbols:\n\n* `M` — Mamba layer\n* `*` — Attention layer\n* `-` — MLP layer\n* `E` — MoE layer\n\nThe number of layers is derived from the pattern length, so `--num-layers`\nshould not be specified when `--hybrid-layer-pattern` is used.\n\nFor example, the 8B hybrid model described in the technical report uses:\n\n```\n--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-\"\n```\n\nThis is a 56-layer model with 4 attention layers, 28 MLP layers, and 24 Mamba\nlayers.\n\nA pure Mamba model uses only `M` symbols (e.g., `MMMMMMMM` for 8 layers).\nA pure transformer model uses only `*` and `-` symbols.\n\n### Pipeline parallelism\n\nUse `|` to define pipeline stage boundaries for flexible virtual pipeline\nparallelism (fVPP). For example, `M-M-|M-M*-|M-M-|M-M*-` defines 4 pipeline\nsegments. The number of segments must be evenly divisible by\n`--pipeline-model-parallel-size`.\n\n### Multi-Token Prediction (MTP)\n\nUse `/` to append MTP layer patterns. Each pattern after the separator\nrepresents one MTP prediction depth. For example, `M*M*/MM/MM` has main\npattern `M*M*` with MTP pattern `MM` repeated for 2 depths.\n\n### Deprecated options\n\n`--hybrid-override-pattern`, `--hybrid-attention-ratio`, and\n`--hybrid-mlp-ratio` are deprecated. Use `--hybrid-layer-pattern` instead.\n\n## Mamba vs Mamba-2\n\nThis codebase currently only supports Mamba-2, and not the original version of\nMamba. However, the\n[fixed snapshot of the code used for the technical report](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba)\ncan be configured to run the original version of Mamba.\n"
},
{
"path": "examples/mamba/run_text_gen_server_8b.sh",
"content": "#!/bin/bash\n\n# Use: ./run_text_gen_server_8b.sh \n# To launch the client: python ../../tools/text_generation_cli.py \n\nCHECKPOINT_PATH=$1\nTOKENIZER_PATH=$2\n\nHYBRID_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-\"\n\nDISTRIBUTED_ARGS=\"--nproc_per_node 1 \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr localhost \\\n --master_port 6000\"\n\nexport NCCL_IB_SL=1\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\nexport NCCL_IB_TIMEOUT=19\nexport NCCL_IB_QPS_PER_CONNECTION=4\n\nexport TRITON_CACHE_DIR=\"./triton-cache/\"\nexport TRITON_CACHE_MANAGER=\"megatron.core.ssm.triton_cache_manager:ParallelFileCacheManager\"\n\ntorchrun $DISTRIBUTED_ARGS ../../tools/run_mamba_text_generation_server.py \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n --untie-embeddings-and-output-weights \\\n --hybrid-layer-pattern ${HYBRID_LAYER_PATTERN} \\\n --hidden-size 4096 \\\n --load ${CHECKPOINT_PATH} \\\n --num-attention-heads 32 \\\n --group-query-attention \\\n --num-query-groups 8 \\\n --attention-dropout 0.0 \\\n --hidden-dropout 0.0 \\\n --disable-bias-linear \\\n --normalization RMSNorm \\\n --seq-length 4096 \\\n --max-position-embeddings 4096 \\\n --position-embedding-type none \\\n --tokenizer-type GPTSentencePieceTokenizer \\\n --tokenizer-model ${TOKENIZER_PATH} \\\n --distributed-backend nccl \\\n --distributed-timeout-minutes 1440 \\\n --bf16 \\\n --micro-batch-size 1 \\\n --use-mcore-models \\\n --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \\\n --seed 42\n"
},
{
"path": "examples/mamba/run_text_gen_server_8b_gpt3.sh",
"content": "#!/bin/bash\n\n# Use: ./run_text_gen_server_8b_gpt3.sh \n# To launch the client: python ../../tools/text_generation_cli.py \n\nCHECKPOINT_PATH=$1\nTOKENIZER_PATH=$2\n\nDISTRIBUTED_ARGS=\"--nproc_per_node 1 \\\n --nnodes 1 \\\n --node_rank 0 \\\n --master_addr localhost \\\n --master_port 6000\"\n\nexport NCCL_IB_SL=1\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\nexport NCCL_IB_TIMEOUT=19\nexport NCCL_IB_QPS_PER_CONNECTION=4\n\ntorchrun $DISTRIBUTED_ARGS ../../tools/run_text_generation_server.py \\\n --tensor-model-parallel-size 1 \\\n --pipeline-model-parallel-size 1 \\\n --use-flash-attn \\\n --apply-layernorm-1p \\\n --untie-embeddings-and-output-weights \\\n --num-layers 32 \\\n --hidden-size 4096 \\\n --load ${CHECKPOINT_PATH} \\\n --num-attention-heads 32 \\\n --attention-dropout 0.0 \\\n --hidden-dropout 0.0 \\\n --disable-bias-linear \\\n --seq-length 4096 \\\n --max-position-embeddings 4096 \\\n --position-embedding-type rope \\\n --rotary-percent 0.5 \\\n --squared-relu \\\n --tokenizer-type GPTSentencePieceTokenizer \\\n --tokenizer-model ${TOKENIZER_PATH} \\\n --distributed-backend nccl \\\n --distributed-timeout-minutes 1440 \\\n --bf16 \\\n --micro-batch-size 1 \\\n --use-mcore-models \\\n --transformer-impl local \\\n --seed 42\n"
},
{
"path": "examples/mamba/train.sh",
"content": "#!/bin/bash\n\n# Use: ./train.sh \n\nMODEL_SCALE=\"800M\" # or \"8B\"\n\ncase \"${MODEL_SCALE}\" in\n \"800M\")\n TENSOR_MODEL_PARALLEL_SIZE=1\n HYBRID_LAYER_PATTERN=\"M-M-M--M-*M-M-M-M--*M-M-M-M-*M--M-M-M-*M-M--M-M-\"\n HIDDEN_SIZE=1024\n NUM_ATTENTION_HEADS=16\n GLOBAL_BATCH_SIZE=32\n ;;\n \"8B\")\n TENSOR_MODEL_PARALLEL_SIZE=4\n 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-\"\n HIDDEN_SIZE=4096\n NUM_ATTENTION_HEADS=32\n GLOBAL_BATCH_SIZE=8\n ;;\n *)\n echo \"Invalid version specified\"\n exit 1\n ;;\nesac\n\nDATA_PATH=$1\nTOKENIZER_PATH=$2\n\nexport NCCL_IB_SL=1\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\nexport NCCL_IB_TIMEOUT=19\nexport NCCL_IB_QPS_PER_CONNECTION=4\n\nCHECKPOINT_DIR=\"./checkpoints\"\nDATACACHE_DIR=\"./data-cache\"\nTENSORBOARD_DIR=\"./tensorboard\"\n\nmkdir -p ${CHECKPOINT_DIR}\nmkdir -p ${DATACACHE_DIR}\nmkdir -p ${TENSORBOARD_DIR}\n\nexport TRITON_CACHE_DIR=\"./triton-cache/\"\nexport TRITON_CACHE_MANAGER=\"megatron.core.ssm.triton_cache_manager:ParallelFileCacheManager\"\n\nSEQ_LEN=4096\nTRAIN_SAMPLES=73242188 # 300B tokens / 4096\nLR_WARMUP_SAMPLES=50000\nLR_DECAY_SAMPLES=73192188 # TRAIN_SAMPLES - LR_WARMUP_SAMPLES\n\noptions=\" \\\n --tensor-model-parallel-size ${TENSOR_MODEL_PARALLEL_SIZE} \\\n --sequence-parallel \\\n --pipeline-model-parallel-size 1 \\\n --use-distributed-optimizer \\\n --overlap-param-gather \\\n --overlap-grad-reduce \\\n --untie-embeddings-and-output-weights \\\n --init-method-std 0.02 \\\n --position-embedding-type none \\\n --hybrid-layer-pattern ${HYBRID_LAYER_PATTERN} \\\n --hidden-size ${HIDDEN_SIZE} \\\n --num-attention-heads ${NUM_ATTENTION_HEADS} \\\n --group-query-attention \\\n --num-query-groups 8 \\\n --seq-length ${SEQ_LEN} \\\n --max-position-embeddings ${SEQ_LEN} \\\n --train-samples ${TRAIN_SAMPLES} \\\n --lr-warmup-samples ${LR_WARMUP_SAMPLES} \\\n --lr-decay-samples ${LR_DECAY_SAMPLES} \\\n --save ${CHECKPOINT_DIR} \\\n --load ${CHECKPOINT_DIR} \\\n --data-path ${DATA_PATH} \\\n --data-cache-path ${DATACACHE_DIR} \\\n --split 99,1,0 \\\n --tokenizer-type GPTSentencePieceTokenizer \\\n --tokenizer-model ${TOKENIZER_PATH} \\\n --distributed-backend nccl \\\n --micro-batch-size 4 \\\n --global-batch-size ${GLOBAL_BATCH_SIZE} \\\n --lr 2.5e-4 \\\n --min-lr 2.5e-5 \\\n --lr-decay-style cosine \\\n --weight-decay 0.1 \\\n --clip-grad 1.0 \\\n --attention-dropout 0.0 \\\n --hidden-dropout 0.0 \\\n --disable-bias-linear \\\n --normalization RMSNorm \\\n --adam-beta1 0.9 \\\n --adam-beta2 0.95 \\\n --log-interval 10 \\\n --save-interval 2000 \\\n --eval-interval 2000 \\\n --eval-iters 32 \\\n --bf16 \\\n --use-mcore-models \\\n --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \\\n --no-create-attention-mask-in-dataloader \\\n --tensorboard-dir ${TENSORBOARD_DIR}\"\n\ntorchrun --nproc_per_node 8 ../../pretrain_mamba.py ${options}\n"
},
{
"path": "examples/mimo/__init__.py",
"content": " "
},
{
"path": "examples/mimo/avlm_inference.py",
"content": "import argparse\nimport os\nfrom pathlib import Path\nfrom typing import Union\n\n# hf path\nimport requests\nimport torch\nfrom PIL import Image\nfrom transformers import AutoProcessor\nfrom transformers import AutoTokenizer\nimport soundfile as sf\nimport io\nimport numpy as np\nimport scipy.signal as signal\n\nfrom examples.mimo.model_providers.llava_avlm import model_provider_llava_avlm\nfrom megatron.core import dist_checkpointing, parallel_state, tensor_parallel\nfrom megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed\nfrom megatron.training import print_rank_0\nfrom examples.mimo.data.utils.calculate_audio_tokens import calculate_num_audio_tokens\n\ndef init_distributed(tp_size: int = 1, pp_size: int = 1):\n if torch.distributed.is_initialized():\n return\n rank = int(os.environ.get(\"LOCAL_RANK\", 0))\n world_size = int(os.environ.get(\"WORLD_SIZE\", 1))\n torch.cuda.set_device(rank % torch.cuda.device_count())\n torch.distributed.init_process_group(\"nccl\", rank=rank, world_size=world_size)\n parallel_state.initialize_model_parallel(tp_size, pp_size)\n\ndef get_input_data(\n processor: AutoProcessor,\n image_processor: AutoProcessor,\n audio_processor: AutoProcessor,\n audio_path: str,\n image_path: str,\n prompt: str,\n device: Union[int, str] = 0):\n \"\"\"\n Prepare inputs for the MIMO model forward pass.\n \"\"\"\n\n def read_audio(audio_path):\n \"\"\"Process audio file and return tensor.\"\"\"\n with open(audio_path, 'rb') as f:\n audio_bytes = f.read()\n audio_io = io.BytesIO(audio_bytes)\n waveform, sample_rate = sf.read(audio_io)\n \n # Resample if needed\n fixed_sample_rate = 16000\n if sample_rate != fixed_sample_rate:\n num_samples = int(len(waveform) * fixed_sample_rate / sample_rate)\n waveform = signal.resample(waveform, num_samples)\n \n # Convert to tensor\n audio_tensor = torch.from_numpy(waveform).float()\n return audio_tensor\n\n def read_image(image_path):\n \"\"\"Process image file and return tensor.\"\"\"\n with open(image_path, 'rb') as f:\n image_bytes = f.read()\n image_io = io.BytesIO(image_bytes)\n image = Image.open(image_io)\n image_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1) # Convert to CxHxW format\n image_tensor = image_tensor.float() / 255.0 # rescale to [0,1] range\n return image_tensor\n\n\n # read audio and image\n audio_tensor = read_audio(audio_path)\n image_tensor = read_image(image_path)\n\n # set up prompt\n conversation = [\n {\n \"role\": \"user\",\n \"content\": [\n {\"type\": \"text\", \"text\": prompt},\n ],\n }\n ]\n prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)\n\n # process audio\n processed_audios = audio_processor(audio_tensor, sampling_rate=16000)\n processed_audios = torch.tensor(processed_audios[\"input_features\"])\n processed_audios = processed_audios.squeeze(0) # remove batch dim\n num_audio_tokens = calculate_num_audio_tokens(audio_tensor.unsqueeze(0), \"openai/whisper-base\")\n audios_seq_lengths = torch.tensor(num_audio_tokens)\n prompt = prompt.replace(\"