-
-The models require vocabulary files to run. The BERT WordPiece vocab file can be extracted from Google's pretrained BERT models: [uncased](https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt), [cased](https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-vocab.txt). The GPT [vocab file](https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json) and [merge table](https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt) can be downloaded directly.
-
-# Usage
+
-After installation, there are several possible workflows. The most comprehensive is:
-1. Data preprocessing
-2. Pretraining
-3. Finetuning (Optional for zero-shot tasks)
-4. Downstream task evaluation or text generation
+Megatron-LM and Megatron Core
+=============================
-However, steps 1 and 2 can be replaced by using one of the pretrained models mentioned above.
+
GPU-optimized library for training transformer models at scale
-We've provided several scripts for pretraining both BERT and GPT in [`examples`](./examples) directory, as well as scripts for both zero-shot and fine-tuned downstream tasks including MNLI, RACE, WikiText103, and LAMBADA evaluation. There is also a script for GPT interactive text generation.
+[](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html)
+[](./CHANGELOG.md)
+[](./LICENSE)
-# Training
-## Data Preprocessing
-The training data requires preprocessing. First, place your training data in a loose json format, with one json containing a text sample per line. For example:
-
-{"src": "www.nvidia.com", "text": "The quick brown fox", "type": "Eng", "id": "0", "title": "First Part"}
-{"src": "The Internet", "text": "jumps over the lazy dog", "type": "Eng", "id": "42", "title": "Second Part"}
-
+
-The name of the `text` field of the json can be changed by using the `--json-key` flag in [`preprocess_data.py`](./tools/preprocess_data.py) The other metadata are optional and are not used in training.
+## About
-The loose json is then processed into a binary format for training. To convert the json into mmap, cached index file, or the lazy loader format use `preprocess_data.py`. Set the `--dataset-impl` flag to `mmap`, `cached`, or `lazy`, respectively (default is `mmap`). An example script to prepare data for BERT training is:
-
+This repository contains two components: **Megatron-LM** and **Megatron Core**.
-The output will be two files named, in this case, `my-bert_text_sentence.bin` and `my-bert_text_sentence.idx`. The `--data-path` specified in later BERT training is the full path and new filename, but without the file extension.
+**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.
-For T5 use the same preprocessing as BERT, perhaps renaming it to:
-
- --output-prefix my-t5 \
-
+**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.
-Some minor modifications are required for GPT data preprocessing, namely, the addition of a merge table, an end-of-document token, removal of sentence splitting, and a change to the tokenizer type:
-
+**[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** provides bidirectional Hugging Face ↔ Megatron checkpoint conversion with production-ready recipes.
-Here the output files are named `my-gpt2_text_document.bin` and `my-gpt2_text_document.idx`. As before, in GPT training, use the longer name without the extension as `--data-path`.
+## Getting Started
-Further command line arguments are described in the source file [`preprocess_data.py`](./tools/preprocess_data.py).
+**Install from PyPI:**
-## BERT Pretraining
-
-
-The [`examples/pretrain_bert.sh`](./examples/pretrain_bert.sh) script runs single GPU 345M parameter BERT pretraining. Debugging is the primary use for single GPU training, as the code base and command line arguments are optimized for highly distributed training. Most of the arguments are fairly self-explanatory. By default, the learning rate decays linearly over the training iterations starting at `--lr` to a minimum set by `--min-lr` over `--lr-decay-iters` iterations. The fraction of training iterations used for warmup is set by `--lr-warmup-fraction`. While this is single GPU training, the batch size specified by `--micro-batch-size` is a single forward-backward path batch-size and the code will perform gradient accumulation steps until it reaches `global-batch-size` which is the batch size per iteration. The data is partitioned into a 949:50:1 ratio for training/validation/test sets (default is 969:30:1). This partitioning happens on the fly, but is consistent across runs with the same random seed (1234 by default, or specified manually with `--seed`). We use `train-iters` as the training iterations requested. Alternatively, one can provide `--train-samples` which is total number of samples to train on. If this option is present, then instead of providing `--lr-decay-iters`, one will need to provide `--lr-decay-samples`.
+```bash
+uv pip install megatron-core
+```
-The logging, checkpoint-saving, and evaluation intervals are specified. Checkpointing the activations facilitates the training of larger models and/or batches. Note that the `--data-path` now includes the additional `_text_sentence` suffix added in preprocessing, but does not include the file extensions.
+**Or clone and install from source:**
-Further command line arguments are described in the source file [`arguments.py`](./megatron/arguments.py).
+```bash
+git clone https://github.com/NVIDIA/Megatron-LM.git
+cd Megatron-LM
+uv pip install -e .
+```
-To run `examples/pretrain_bert.sh`, make any desired modifications including setting the environment variables for `CHECKPOINT_PATH`, `VOCAB_FILE`, and `DATA_PATH`. Make sure to set these variables to their paths in the container. Then launch the container with Megatron and necessary paths mounted (as explained in [Setup](#setup)) and run the example script.
+> **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 .`).
-## GPT Pretraining
+For NGC container setup and all installation options, see the **[Installation Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/get-started/install.html)**.
-The `examples/pretrain_gpt.sh` script runs single GPU 345M parameter GPT pretraining. As mentioned above, single GPU training is primarily intended for debugging purposes, as the code is optimized for distributed training.
+- **[Your First Training Run](https://docs.nvidia.com/megatron-core/developer-guide/latest/get-started/quickstart.html)** - End-to-end training examples with data preparation
+- **[Parallelism Strategies](https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/parallelism-guide.html)** - Scale training across GPUs with TP, PP, DP, EP, and CP
+- **[Contribution Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/developer/contribute.html)** - How to contribute to Megatron Core
-It follows largely the same format as the previous BERT script with a few notable differences: the tokenization scheme used is BPE (which requires a merge table and a `json` vocabulary file) instead of WordPiece, the model architecture allows for longer sequences (note that the max position embedding must be greater than or equal to the maximum sequence length), and the `--lr-decay-style` has been set to cosine decay. Note that the `--data-path` now includes the additional `_text_document` suffix added in preprocessing, but does not include the file extensions.
+# Latest News
-Further command line arguments are described in the source file [`arguments.py`](./megatron/arguments.py).
+- **[2026/03]** **Deprecating Python 3.10 support:** We're officially dropping Python 3.10 support with the upcoming 0.17.0 release. Downstream applications must raise their lower boundary to 3.12 to stay compatible with MCore.
+- **[2026/01]** **[Dynamic Context Parallelism](https://developer.nvidia.com/blog/speeding-up-variable-length-training-with-dynamic-context-parallelism-and-nvidia-megatron-core/)** - Up to 1.48x speedup for variable-length sequence training with adaptive CP sizing.
+- **[2025/12]** **Megatron Core development has moved to GitHub!** All development and CI now happens in the open. We welcome community contributions.
+- **[2025/10]** **[Megatron Dev Branch](https://github.com/NVIDIA/Megatron-LM/tree/dev)** - early access branch with experimental features.
+- **[2025/10]** **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Bidirectional converter for interoperability between Hugging Face and Megatron checkpoints, featuring production-ready recipes for popular models.
+- **[2025/08]** **[MoE Q3-Q4 2025 Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - Comprehensive roadmap for MoE features including DeepSeek-V3, Qwen3, advanced parallelism strategies, FP8 optimizations, and Blackwell performance enhancements.
+- **[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.
+- **[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.
+- **[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/)).
-`examples/pretrain_gpt.sh` can be launched the same way as described for BERT. Set the env vars and make any other modifications, launch the container with appropriate mounts, and run the script.
+
+Previous News
-## T5 Pretraining
+- **[2024/07]** Megatron Core v0.7 improves scalability and training resiliency and adds support for multimodal training ([blog](https://developer.nvidia.com/blog/train-generative-ai-models-more-efficiently-with-new-nvidia-Megatron-Core-functionalities/)).
+- **[2024/06]** Megatron Core added supports for Mamba-based models. Check out our paper [An Empirical Study of Mamba-based Language Models](https://arxiv.org/pdf/2406.07887) and [code example](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba).
+- **[2024/01 Announcement]** NVIDIA has released the core capabilities in **Megatron-LM** into [**Megatron Core**](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core) in this repository. Megatron Core expands upon Megatron-LM's GPU-optimized techniques with more cutting-edge innovations on system-level optimizations, featuring composable and modular APIs.
-Very similar to BERT and GPT, the `examples/pretrain_t5.sh` script runs single GPU "base" (~220M parameter) T5 pretraining. The primary difference from BERT and GPT is the addition of the following arguments to accommodate the T5 architecture:
+
-* `--kv-channels` sets the inner dimension of the "key" and "value" matrices of all attention mechanisms in the model. For BERT and GPT this defaults to the hidden size divided by the number of attention heads, but can be configured for T5.
+# Project Structure
-* `--ffn-hidden-size` sets the hidden size in the feed-forward networks within a transformer layer. For BERT and GPT this defaults to 4 times the transformer hidden size, but can be configured for T5.
+```
+Megatron-LM/
+├── megatron/
+│ ├── core/ # Megatron Core (kernels, parallelism, building blocks)
+│ │ ├── models/ # Transformer models
+│ │ ├── transformer/ # Transformer building blocks
+│ │ ├── tensor_parallel/ # Tensor parallelism
+│ │ ├── pipeline_parallel/ # Pipeline parallelism
+│ │ ├── distributed/ # Distributed training (FSDP, DDP)
+│ │ ├── optimizer/ # Optimizers
+│ │ ├── datasets/ # Dataset loaders
+│ │ ├── inference/ # Inference engines and server
+│ │ └── export/ # Model export (e.g. TensorRT-LLM)
+│ ├── training/ # Training scripts
+│ ├── legacy/ # Legacy components
+│ ├── post_training/ # Post-training (quantization, distillation, pruning, etc.)
+│ └── rl/ # Reinforcement learning (RLHF, etc.)
+├── examples/ # Ready-to-use training examples
+├── tools/ # Utility tools
+├── tests/ # Comprehensive test suite
+└── docs/ # Documentation
+```
-* `--encoder-seq-length` and `--decoder-seq-length` set the sequence length for the encoder and decoder separately.
+# Performance Benchmarking
-All of the other arguments remain as they were for BERT and GPT pretraining. Run this example with the same steps described above for the other scripts.
+For our latest performance benchmarking results, please refer to [NVIDIA Megatron Bridge Performance Summary](https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html).
-## Distributed Pretraining
+Our codebase efficiently trains models from 2B to 462B parameters across thousands of GPUs, achieving up to **47% Model FLOP Utilization (MFU)** on H100 clusters.
-The `examples/pretrain_{bert,gpt,t5}_distributed.sh` scripts use the PyTorch distributed launcher for distributed training. As such, multi-node training can be achieved by properly setting environment variables. See the official PyTorch [documentation](https://pytorch.org/docs/stable/elastic/run.html#launcher-api) for further description of these [environment variables](https://pytorch.org/docs/stable/distributed.html#environment-variable-initialization). By default, multi-node training uses the [nccl](https://developer.nvidia.com/nccl) distributed backend. A simple set of additional arguments and the use of the PyTorch distributed module with the `torchrun` elastic launcher (equivalent to `python -m torch.distributed.run`) are the only additional requirements to adopt distributed training. See any of `examples/pretrain_{bert,gpt,t5}_distributed.sh` for more details.
+
-We use two types of parallelism: data and model parallelism. We facilitate two distributed data parallel implementations: a simple one of our own that performs gradient all-reduce at the end of back propagation step, and Torch's distributed data parallel wrapper that overlaps gradient reduction with back propagation computation. To switch between these two options use `--DDP-impl local` or `--DDP-impl torch`, respectively. As expected, Torch distributed data parallelism is more efficient at larger model sizes. For example, for the 8.3 billion parameters model running on 512 GPUs, the scaling increases from 60% to 76% when Torch's distributed data parallel is used. However, the overlapping method requires more memory and for some configurations (e.g., 2.5 billion parameters using 2-way model parallel and 1.2 billion parameters with no model parallel) can make the overall training slower as a result. We empirically found that using a smaller model in those cases improves the training time.
+**Benchmark Configuration:**
-Second, we developed a simple and efficient two-dimensional model-parallel approach. To use tensor model parallelism (splitting execution of a single transformer module over multiple GPUs, see Section 3 of [our paper](https://arxiv.org/pdf/1909.08053.pdf)), add the `--tensor-model-parallel-size` flag to specify the number of GPUs among which to split the model, along with the arguments passed to the distributed launcher as mentioned above. To use sequence parallelism specify `--sequence-parallel`, which requires tensor model parallel as it split among the same GPUs (more details in Section 4.2.2 of [our paper](https://arxiv.org/pdf/2205.05198.pdf)).
+- **Vocabulary size**: 131,072 tokens
+- **Sequence length**: 4096 tokens
+- **Model scaling**: Varied hidden size, attention heads, and layers to achieve target parameter counts
+- **Communication optimizations**: Fine-grained overlapping with DP (`--overlap-grad-reduce`, `--overlap-param-gather`), TP (`--tp-comm-overlap`), and PP (enabled by default)
-To use pipeline model parallelism (sharding the transformer modules into stages with an equal number of transformer modules on each stage, and then pipelining execution by breaking the batch into smaller microbatches, see Section 2.2 of [our paper](https://arxiv.org/pdf/2104.04473.pdf)), use the `--pipeline-model-parallel-size` flag to specify the number of stages to split the model into (e.g., splitting a model with 24 transformer layers across 4 stages would mean each stage gets 6 transformer layers each).
+**Key Results:**
-
+- **6144 H100 GPUs**: Successfully benchmarked 462B parameter model training
+- **Superlinear scaling**: MFU increases from 41% to 47-48% with model size
+- **End-to-end measurement**: Throughputs include all operations (data loading, optimizer steps, communication, logging)
+- **Production ready**: Full training pipeline with checkpointing and fault tolerance
+- *Note: Performance results measured without training to convergence*
-We have examples of how to use these two different forms of model parallelism the example scripts ending in `distributed_with_mp.sh`:
+## Weak Scaling Results
-Other than these minor changes, the distributed training is identical to the training on a single GPU.
+Our 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.
-The interleaved pipelining schedule (more details in Section 2.2.2 of [our paper](https://arxiv.org/pdf/2104.04473.pdf)) can be enabled using the `--num-layers-per-virtual-pipeline-stage` argument, which controls the number of transformer layers in a virtual stage (by default with the non-interleaved schedule, each GPU will execute a single virtual stage with `NUM_LAYERS / PIPELINE_MP_SIZE` transformer layers). The total number of layers in the transformer model should be divisible by this argument value. Additionally, the number of microbatches in the pipeline (computed as `GLOBAL_BATCH_SIZE / (DATA_PARALLEL_SIZE * MICRO_BATCH_SIZE)`) should be divisible by the `PIPELINE_MP_SIZE` when using this schedule (this condition is checked in an assertion in the code). The interleaved schedule is not supported for pipelines with 2 stages (`PIPELINE_MP_SIZE=2`).
+
-## Activation Checkpointing and Recomputation
+## Strong Scaling Results
-To reduce GPU memory usage so deploy a large model to a training system, we support activation checkpointing and recomputation. We support two levels of recompute granularity: `selective` and `full`. Selective recomputation is the default and recommended in almost all cases. It saves the activations that take less space and are expensive to recompute and recomputes activations that take a lot of space but are relatively cheap to recompute (see [our paper](https://arxiv.org/pdf/2205.05198) for details). To enable selective activation recompute simply use `--recompute-activations`.
+We 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%.
-For cases where memory is very tight, `full` checkpointing saves just the inputs to a transformer layer, or a block of transformer layers, and recomputes everything else. To turn on full activation recompute use `--recompute-granularity full`. When using full activation recomputation, there are two methods: `uniform` and `block`, chosen using the `--recompute-method` argument.
+
-* Uniform method uniformly divides the Transformer layers into groups of layers and stores the input activations of each group in the memory. The baseline group size is 1 and, in this case, the input activation of each Transformer layer is checkpointed. When the GPU memory is insufficient, increasing the number of layers per group reduces the memory usage thus enables running a bigger model. For example, when using the number of layers per group of 4, the input activation of each group of 4 Transformer layers is checkpointed.
+# Roadmaps
-* Block method checkpoints the input activations of a set number of individual Transformer layers per pipeline stage and do the rest of layers without any checkpointing. This method can be used to skip checkpointing some Transformer layers until the GPU memory is fully used, which is applicable only when there is unused GPU memory. Checkpointing fewer transformer layers avoids unnecessary activation recomputation in the backprop thus improves training performance. For example, when we specify 5 layers to checkpoint of 8 layers per pipeline stage, the input activations of only the first 5 Transformer layers are checkpointed and activation recomputation for the rest 3 layers is not needed in the backprop.
+- **[MoE Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/1729)** - DeepSeek-V3, Qwen3, advanced parallelism, FP8 optimizations, and Blackwell enhancements
+# Resources
-## Distributed Optimizer
+## Getting Help
-Usage: `--use-distributed-optimizer`. Compatible with all model and data types.
+- 📖 **[Documentation](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html)** - Official documentation
+- 🐛 **[Issues](https://github.com/NVIDIA/Megatron-LM/issues)** - Bug reports and feature requests
-The distributed optimizer is a memory savings technique, whereby the optimizer state is evenly distributed across data parallel ranks (versus the traditional method of replicating the optimizer state across data parallel ranks). As described in [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://arxiv.org/abs/1910.02054), our implementation distributes all optimizer state that does not overlap with the model state. For example, when using fp16 model params, the distributed optimizer maintains its own separate copy of fp32 main params & grads, which are distributed across DP ranks. When using bf16 model params, however, the distributed optimizer's fp32 main grads are the same as the model's fp32 grads, and so the grads in this case are not distributed (although the fp32 main params are still distributed, as they are separate from the bf16 model params).
+## Contributing
-Theoretical memory savings vary depending on the combination of the model's param dtype and grad dtype. In our implementation, the theoretical number of bytes per parameter is (where 'd' is the data parallel size):
+We ❤️ contributions! Ways to contribute:
-| | Non-distributed optim | Distributed optim |
-|-|-|-|
-| fp16 param, fp16 grads | 20 | 4 + 16/d |
-| bf16 param, fp32 grads | 18 | 6 + 12/d |
-| fp32 param, fp32 grads | 16 | 8 + 8/d |
+- 🐛 **Report bugs** - Help us improve reliability
+- 💡 **Suggest features** - Shape the future of Megatron Core
+- 📝 **Improve docs** - Make Megatron Core more accessible
+- 🔧 **Submit PRs** - Contribute code improvements
-## FlashAttention
+**→ [Contributing Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/developer/contribute.html)**
-Usage: `--use-flash-attn`. Support attention head dimensions at most 128.
+## Citation
-[FlashAttention](https://github.com/HazyResearch/flash-attention) is a fast and
-memory-efficient algorithm to compute exact attention. It speeds up model
-training and reduces memory requirement.
+If you use Megatron in your research or project, we appreciate that you use the following citations:
-To install FlashAttention:
-```sh
-pip install flash-attn
+```bibtex
+@article{megatron-lm,
+ title={Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism},
+ author={Shoeybi, Mohammad and Patwary, Mostofa and Puri, Raul and LeGresley, Patrick and Casper, Jared and Catanzaro, Bryan},
+ journal={arXiv preprint arXiv:1909.08053},
+ year={2019}
+}
```
-
-## GPT-3 Example
-
-In `examples/pretrain_gpt3_175B.sh` we have provided an example of how to configure Megatron to run [GPT-3](https://arxiv.org/abs/2005.14165) with 175 billion parameters on 1024 GPUs. The script is designed for [slurm](https://slurm.schedmd.com/documentation.html) with [pyxis](https://github.com/NVIDIA/pyxis) plugin but can be easily adopted to any other scheduler. It uses 8-way and 16-way tensor and pipeline parallelism, respectively. With options `global-batch-size 1536` and `rampup-batch-size 16 16 5859375`, the training will start with global batch size 16 and linearly increase the global batch size to 1536 over 5,859,375 samples with incrmeental steps 16. The training dataset can be either a single set or a multiple datasets combined with a set of weights.
-
-With full global batch size of 1536 on 1024 A100 GPUs, each iteration takes around 32 seconds resulting in 138 teraFLOPs per GPU which is 44% of the theoretical peak FLOPs.
-
-
-## Retro
-
-See:
-
-- `tools/retro/README.md` for an overview.
-- `tools/retro/examples/get_preprocess_cmd.sh` for an example of common preprocessing arguments.
-- `tools/retro/examples/preprocess_data.sh` for an example of how to preprocess data.
-- `tools/retro/examples/pretrain_model.sh` for an example of how to pretrain a model.
-
-Retro is a retrieval-enhanced model that is based on GPT. As described in [Improving language models by retrieving from trillions of tokens](https://arxiv.org/abs/2112.04426), Retro retrieves from a database of document chunks by performing locality search using a sample's tokens. The retrieval database can be large -- often billions or even trillions of tokens -- and provides a more efficient storage mechanism of factual knowledge, when compared to storing factual knowledge implicitly within the network's parameters.
-
-Using Retro requires two steps: 1) preprocessing the retrieval database and pretraining neighbors, and 2) pretraining a model using this data. Please see `tools/retro/README.md` for a detailed overview.
-
-
-
-# Evaluation and Tasks
-
-We provide several command line arguments, detailed in the scripts listed below, to handle various zero-shot and fine-tuned downstream tasks. However, you can also finetune your model from a pretrained checkpoint on other corpora as desired. To do so, simply add the `--finetune` flag and adjust the input files and training parameters within the original training script. The iteration count will be reset to zero, and the optimizer and internal state will be reinitialized. If the fine-tuning is interrupted for any reason, be sure to remove the `--finetune` flag before continuing, otherwise the training will start again from the beginning.
-
-Because evaluation requires substantially less memory than training, it may be advantageous to merge a model trained in parallel for use on fewer GPUs in downstream tasks. The following script accomplishes this. This example reads in a GPT model with 4-way tensor and 4-way pipeline model parallelism and writes out a model with 2-way tensor and 2-way pipeline model parallelism.
-
-
-
-Several downstream tasks are described for both GPT and BERT models below. They can be run in distributed and model parallel modes with the same changes used in the training scripts.
-
-## GPT Text Generation
-
-We have included a simple REST server to use for text generation in `tools/run_text_generation_server.py`. You run it much like you would start a pretraining job, specifying an appropriate pretrained checkpoint. There are also few optional parameters: `temperature`, `top-k`and `top-p`. See `--help` or the source file for more information. See [examples/run_text_generation_server_345M.sh](examples/run_text_generation_server_345M.sh) for an example of how to run the server.
-
-Once the server is running you can use `tools/text_generation_cli.py` to query it, it takes one argument which is the host the server is running on.
-
-
-tools/text_generation_cli.py localhost:5000
-
-
-You can also use CURL or any other tools to query the server directly:
-
-
-
-See [megatron/text_generation_server.py](megatron/text_generation_server.py) for more API options.
-
-### Detoxify GPT via Self-generation
-We include an example in `examples/detxoify_lm/` to detoxify language models by leveraging the generative power of language models.
-
-See [examples/detxoify_lm/README.md](examples/detxoify_lm/README.md) for step-by-step tutorials on how to perform domain-adaptive training and detoxify LM using self-generated corpus.
-
-
-## GPT Evaluation
-We include example scripts for GPT evaluation on WikiText perplexity evaluation and LAMBADA Cloze accuracy.
-
-### WikiText Perplexity Evaluation
-For even comparison with prior works, we evaluate perplexity on the word-level [WikiText-103 test dataset](https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-v1.zip), and appropriately compute perplexity given the change in tokens when using our subword tokenizer.
-
-We use the following command to run WikiText-103 evaluation on a 345M parameter model.
-
-
-
-### LAMBADA Cloze Accuracy
-To compute LAMBADA cloze accuracy (the accuracy of predicting the last token given the preceding tokens) we utilize a detokenized, processed version of the [LAMBADA dataset](https://github.com/cybertronai/bflm/blob/master/lambada_test.jsonl).
-
-We use the following command to run LAMBADA evaluation on a 345M parameter model. Note that the `--strict-lambada` flag should be used to require whole word matching. Make that `lambada` is part of the file path.
-
-
-
-Further command line arguments are described in the source file [`main.py`](./tasks/main.py)
-
-## BERT Task Evaluation
-### RACE Evaluation
-The following script finetunes the BERT model for evaluation on the [RACE dataset](http://www.cs.cmu.edu/~glai1/data/race/). The `TRAIN_DATA` and `VALID_DATA` directory contain the RACE dataset as separate `.txt` files. Note that for RACE, the batch size is the number of RACE query's to evaluate. Since each RACE query has four samples, the effective batch size passed through the model will be four times the batch size specified on the command line.
-
-
-
-### MNLI Evaluation
-The following script finetunes the BERT model for evaluation with the [MultiNLI sentence pair corpus](https://www.nyu.edu/projects/bowman/multinli/). Because the matching tasks are quite similar, the script can be quickly tweaked to work with the [Quora Question Pairs](https://www.kaggle.com/quora/question-pairs-dataset) (QQP) dataset as well.
-
-
-
-TRAIN_DATA="data/glue_data/MNLI/train.tsv"
-VALID_DATA="data/glue_data/MNLI/dev_matched.tsv \
- data/glue_data/MNLI/dev_mismatched.tsv"
-PRETRAINED_CHECKPOINT=checkpoints/bert_345m
-VOCAB_FILE=bert-vocab.txt
-CHECKPOINT_PATH=checkpoints/bert_345m_mnli
-COMMON_TASK_ARGS=<same as those in RACE Evaluation above>
-COMMON_TASK_ARGS_EXT=<same as those in RACE Evaluation above>
-
-python tasks/main.py \
- --task MNLI \
- $COMMON_TASK_ARGS \
- $COMMON_TASK_ARGS_EXT \
- --tokenizer-type BertWordPieceLowerCase \
- --epochs 5 \
- --micro-batch-size 8 \
- --lr 5.0e-5 \
- --lr-warmup-fraction 0.065
-
-
-# Datasets
-We do not host any datasets for GPT or BERT training, however, we detail their collection so that our results may be reproduced.
-
-## Collecting Wikipedia Training Data
-We recommend following the Wikipedia data extraction process specified by Google research: "the recommended pre-processing is to download [the latest dump](https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2), extract the text with [WikiExtractor.py](https://github.com/attardi/wikiextractor), and then apply any necessary cleanup to convert it into plain text."
-
-We recommend using the `--json` argument when using WikiExtractor, which will dump the Wikipedia data into loose json format (one json per line), making it more manageable on the file system and also readily consumable by our codebase. We recommend further preprocessing this json dataset by nltk punctuation standardization. For BERT training, use the `--split-sentences` flag to `preprocess_data.py` as described [above](#data-preprocessing) to include sentence breaks in the produced index. If you'd like to use Wikipedia data for GPT training you should still clean it with nltk/spacy/ftfy, but do not use the `--split-sentences` flag.
-
-## Collecting GPT Webtext Data
-We utilize the publicly available [OpenWebText](https://github.com/eukaryote31/openwebtext) library from [jcpeterson](https://github.com/jcpeterson/openwebtext) and [eukaryote31's](https://github.com/eukaryote31/openwebtext) work to download urls. We then filtered, cleaned, and deduplicated all downloaded content according to the procedure described in our [openwebtext](./tools/openwebtext) directory. For reddit URLs corresponding to content up to October 2018 we arrived at approximately 37GB of content.
diff --git a/SECURITY.md b/SECURITY.md
new file mode 100644
index 00000000000..728cdb4a1d2
--- /dev/null
+++ b/SECURITY.md
@@ -0,0 +1,25 @@
+## Security
+
+NVIDIA is dedicated to the security and trust of our software products and services, including all source code repositories managed through our organization.
+
+If you need to report a security issue, please use the appropriate contact points outlined below. **Please do not report security vulnerabilities through GitHub.** If a potential security issue is inadvertently reported via a public issue or pull request, NVIDIA maintainers may limit public discussion and redirect the reporter to the appropriate private disclosure channels.
+
+## Reporting Potential Security Vulnerability in an NVIDIA Product
+
+To report a potential security vulnerability in any NVIDIA product:
+
+- Web: [Security Vulnerability Submission Form](https://www.nvidia.com/object/submit-security-vulnerability.html)
+- E-Mail: psirt@nvidia.com
+ - We encourage you to use the following PGP key for secure email communication: [NVIDIA public PGP Key for communication](https://www.nvidia.com/en-us/security/pgp-key)
+ - Please include the following information:
+ - Product/Driver name and version/branch that contains the vulnerability
+ - Type of vulnerability (code execution, denial of service, buffer overflow, etc.)
+ - Instructions to reproduce the vulnerability
+ - Proof-of-concept or exploit code
+ - Potential impact of the vulnerability, including how an attacker could exploit the vulnerability
+
+While NVIDIA currently does not have a bug bounty program, we do offer acknowledgement when an externally reported security issue is addressed under our coordinated vulnerability disclosure policy. Please visit our [Product Security Incident Response Team (PSIRT)](https://www.nvidia.com/en-us/security/psirt-policies/) policies page for more information.
+
+## NVIDIA Product Security
+
+For all security-related concerns, please visit NVIDIA's Product Security portal at https://www.nvidia.com/en-us/security
diff --git a/codecov.yml b/codecov.yml
new file mode 100644
index 00000000000..aa37017f082
--- /dev/null
+++ b/codecov.yml
@@ -0,0 +1,14 @@
+comment: false
+coverage:
+ status:
+ project: false
+ patch:
+ default:
+ target: 80%
+ threshold: 5%
+ base: auto
+ if_ci_failed: error
+ if_no_uploads: success
+ if_not_found: success
+fixes:
+ - "/opt/megatron-lm/::"
diff --git a/docker/.ngc_version.dev b/docker/.ngc_version.dev
new file mode 100644
index 00000000000..3356f1f0bca
--- /dev/null
+++ b/docker/.ngc_version.dev
@@ -0,0 +1 @@
+nvcr.io/nvidia/pytorch:26.04-py3
diff --git a/docker/.ngc_version.lts b/docker/.ngc_version.lts
new file mode 100644
index 00000000000..6b72812b34f
--- /dev/null
+++ b/docker/.ngc_version.lts
@@ -0,0 +1 @@
+nvcr.io/nvidia/pytorch:25.09-py3
\ No newline at end of file
diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev
new file mode 100644
index 00000000000..377bd4a2005
--- /dev/null
+++ b/docker/Dockerfile.ci.dev
@@ -0,0 +1,103 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# syntax=docker/dockerfile:1.3-labs
+
+ARG FROM_IMAGE_NAME
+FROM ${FROM_IMAGE_NAME} as main
+ENV PIP_CONSTRAINT=""
+ENV DEBIAN_FRONTEND=noninteractive
+ARG UV_VERSION=0.7.2
+ARG YQ_VERSION=4.44.1
+ENV PATH="/root/.local/bin:$PATH"
+ARG UV_PROJECT_ENVIRONMENT=/opt/venv
+ENV UV_PROJECT_ENVIRONMENT=${UV_PROJECT_ENVIRONMENT}
+ENV VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT
+ENV PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH"
+ENV UV_LINK_MODE=copy
+
+RUN bash -ex <<"EOF"
+ apt-get update
+ apt-get install -y --no-install-recommends gettext python3-venv psmisc uuid-runtime
+ apt-get clean
+ python -m venv /opt/jet
+ ARCH=$(uname -m)
+ case "${ARCH}" in \
+ "x86_64") YQ_ARCH=amd64 ;; \
+ "aarch64") YQ_ARCH=arm64 ;; \
+ "armv7l") YQ_ARCH=arm ;; \
+ *) echo "Unsupported architecture: ${ARCH}" && exit 1 ;; \
+ esac
+ wget https://github.com/mikefarah/yq/releases/download/v${YQ_VERSION}/yq_linux_${YQ_ARCH} -O /usr/local/bin/yq
+ chmod a+x /usr/local/bin/yq
+ curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh
+EOF
+
+COPY README.md pyproject.toml uv.lock /workspace/
+COPY megatron/core/__init__.py /workspace/megatron/core/
+COPY megatron/core/package_info.py /workspace/megatron/core/
+ARG IMAGE_TYPE=dev
+ENV NVTE_BUILD_NUM_PHILOX_ROUNDS=3
+RUN --mount=type=cache,target=/root/.cache/uv \
+ bash -ex <<"EOF"
+ export NVTE_CUDA_ARCHS="80;90;100"
+ uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages
+ uv sync --only-group build
+ uv sync --extra ${IMAGE_TYPE} --extra mlm --extra ssm --extra te --link-mode copy --locked \
+ --no-install-package torch \
+ --no-install-package torchvision \
+ --no-install-package triton \
+ --no-install-package transformer-engine-cu12 \
+ --no-install-package nvidia-cublas-cu12 \
+ --no-install-package nvidia-cuda-cupti-cu12 \
+ --no-install-package nvidia-cuda-nvrtc-cu12 \
+ --no-install-package nvidia-cuda-runtime-cu12 \
+ --no-install-package nvidia-cudnn-cu12 \
+ --no-install-package nvidia-cufft-cu12 \
+ --no-install-package nvidia-cufile-cu12 \
+ --no-install-package nvidia-curand-cu12 \
+ --no-install-package nvidia-cusolver-cu12 \
+ --no-install-package nvidia-cusparse-cu12 \
+ --no-install-package nvidia-cusparselt-cu12 \
+ --no-install-package nvidia-nccl-cu12
+EOF
+
+# Install DeepEP
+COPY docker/patches/deepep.patch /workspace/deepep.patch
+RUN bash -ex <<"EOF"
+ cd /workspace
+ uv pip install nvidia-nvshmem-cu13==3.4.5
+ pushd /opt/venv/lib/python3.12/site-packages/nvidia/nvshmem/lib/
+ ln -s libnvshmem_host.so.3 libnvshmem_host.so
+ popd
+
+ git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git
+ pushd DeepEP
+ git checkout 34152ae28f80bcc3ee38d7a12cb2ad87cfd4ea72
+ patch -p1 < /workspace/deepep.patch
+ popd
+ TORCH_CUDA_ARCH_LIST="9.0 10.0 12.0" uv pip install --no-build-isolation -v DeepEP/.
+ rm -rf DeepEP
+EOF
+
+COPY assets/ /opt/data/
+ENV UV_PYTHON=$UV_PROJECT_ENVIRONMENT/bin/python
+
+##### For NVIDIANS only #####
+FROM main as jet
+ARG JET_API_VERSION
+ENV PATH="$PATH:/opt/jet/bin"
+RUN --mount=type=secret,id=JET_INDEX_URLS bash -ex <<"EOF"
+ JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS)
+ python -m venv /opt/jet
+ /opt/jet/bin/pip install --no-cache-dir $JET_INDEX_URLS \
+ "jet-api==$JET_API_VERSION" "setuptools<82.0.0"
+EOF
+
+RUN --mount=type=secret,id=JET_INDEX_URLS \
+ --mount=type=secret,id=LOGGER_INDEX_URL bash -ex <<"EOF"
+ JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS)
+ LOGGER_INDEX_URL=$(cat /run/secrets/LOGGER_INDEX_URL)
+ uv pip install --no-cache-dir --upgrade $LOGGER_INDEX_URL "one-logger"
+ uv pip install --no-cache-dir --upgrade "setuptools>=80"
+ uv pip install --no-cache-dir --upgrade $JET_INDEX_URLS "jet-client~=4.0"
+EOF
+###
diff --git a/docker/Dockerfile.ci.nemo b/docker/Dockerfile.ci.nemo
new file mode 100644
index 00000000000..b00349e101a
--- /dev/null
+++ b/docker/Dockerfile.ci.nemo
@@ -0,0 +1,21 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# syntax=docker/dockerfile:1.3-labs
+
+ARG FROM_IMAGE_NAME
+FROM ${FROM_IMAGE_NAME} as main
+
+RUN apt-get update && \
+ apt-get install -y --no-install-recommends gettext && \
+ apt-get clean && \
+ wget https://github.com/mikefarah/yq/releases/download/v4.44.1/yq_linux_amd64 -O /usr/local/bin/yq && \
+ chmod a+x /usr/local/bin/yq
+
+##### For NVIDIANS only #####
+FROM main as jet
+ARG JET_API_VERSION
+RUN --mount=type=secret,id=JET_INDEX_URLS \
+ JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) && \
+ pip install --no-cache-dir jet-api==$JET_API_VERSION "jet-client~=4.0" --upgrade $JET_INDEX_URLS
+
+ENV PATH="$PATH:/opt/jet/bin"
+###
diff --git a/docker/Dockerfile.linting b/docker/Dockerfile.linting
new file mode 100644
index 00000000000..259c0bbedcd
--- /dev/null
+++ b/docker/Dockerfile.linting
@@ -0,0 +1,23 @@
+# syntax=docker/dockerfile:experimental
+
+ARG FROM_IMAGE_NAME
+FROM $FROM_IMAGE_NAME as main
+ENV DEBIAN_FRONTEND=noninteractive
+ARG UV_VERSION=0.7.2
+ARG YQ_VERSION=4.44.1
+ENV PATH="/root/.local/bin:$PATH"
+ENV UV_PROJECT_ENVIRONMENT=/opt/venv
+ENV PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH"
+ENV UV_LINK_MODE=copy
+RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh
+WORKDIR /opt/megatron-lm
+COPY pyproject.toml uv.lock /opt/megatron-lm/
+COPY megatron/core/package_info.py megatron/core/__init__.py /opt/megatron-lm/megatron/core/
+RUN uv sync --locked --only-group linting --only-group test --only-group ci
+
+##### For NVIDIANS only #####
+FROM main as jet
+ARG JET_API_VERSION
+RUN --mount=type=secret,id=JET_INDEX_URLS \
+ JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) && \
+ uv pip install --no-cache-dir "jet-client~=2.0" --upgrade $JET_INDEX_URLS
diff --git a/docker/common/install.sh b/docker/common/install.sh
new file mode 100644
index 00000000000..3991fd41448
--- /dev/null
+++ b/docker/common/install.sh
@@ -0,0 +1,146 @@
+#!/bin/bash
+set -xeuo pipefail # Exit immediately if a command exits with a non-zero status
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --base-image)
+ BASE_IMAGE="$2"
+ shift 2
+ ;;
+ --python-version)
+ PYTHON_VERSION="$2"
+ shift 2
+ ;;
+ --environment)
+ ENVIRONMENT="$2"
+ shift 2
+ ;;
+ --use-uv)
+ USE_UV="true"
+ shift 1
+ ;;
+ *)
+ echo "Unknown option: $1"
+ echo "Usage: $0 --base-image {pytorch|ubuntu} [--use-uv] [--python-version] [--environment]"
+ exit 1
+ ;;
+ esac
+done
+
+if [[ -z "${PYTHON_VERSION:-}" ]]; then
+ PYTHON_VERSION="3.12"
+fi
+
+if [[ -z "${USE_UV:-}" ]]; then
+ USE_UV="false"
+fi
+
+# Validate base image argument
+if [[ -z "${BASE_IMAGE:-}" || -z "${ENVIRONMENT:-}" ]]; then
+ echo "Error: --base-image argument is required"
+ echo "Usage: $0 --base-image {pytorch|ubuntu} --environment {dev|lts}"
+ exit 1
+fi
+
+if [[ "$BASE_IMAGE" != "pytorch" && "$BASE_IMAGE" != "ubuntu" ]]; then
+ echo "Error: --base-image must be either 'pytorch' or 'ubuntu'"
+ echo "Usage: $0 --base-image {pytorch|ubuntu}"
+ exit 1
+fi
+
+if [[ "$ENVIRONMENT" != "dev" && "$ENVIRONMENT" != "lts" ]]; then
+ echo "Error: --environment must be either 'dev' or 'lts'"
+ echo "Usage: $0 --environment {dev|lts}"
+ exit 1
+fi
+
+main() {
+ if [[ -n "${PAT:-}" ]]; then
+ echo -e "machine github.com\n login token\n password $PAT" >~/.netrc
+ chmod 600 ~/.netrc
+ fi
+
+ # Install dependencies
+ export DEBIAN_FRONTEND=noninteractive
+
+ # Install Python
+ apt-get update
+ apt-get install -y software-properties-common
+ add-apt-repository ppa:deadsnakes/ppa -y
+ apt-get install -y python$PYTHON_VERSION-dev python$PYTHON_VERSION-venv
+ update-alternatives --install /usr/bin/python3 python3 /usr/bin/python$PYTHON_VERSION 1
+
+ # Install tools
+ apt-get update
+ apt-get install -y wget curl git cmake
+
+ # Install CUDA
+ if [[ "$BASE_IMAGE" == "ubuntu" ]]; then
+ rm /etc/apt/sources.list.d/cuda*.list || true
+ rm /etc/apt/sources.list.d/nvidia-cuda.list || true
+ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
+ dpkg -i cuda-keyring_1.1-1_all.deb
+ rm cuda-keyring_1.1-1_all.deb
+ apt-get update
+ apt-get install -y cuda-toolkit-12-8 cudnn-cuda-12 libcudnn9-cuda-12 libcutlass-dev
+ fi
+
+ # Clean up
+ apt-get clean
+
+ unset PIP_CONSTRAINT
+
+ if [[ "$USE_UV" == "true" ]]; then
+ if [[ "$BASE_IMAGE" == "pytorch" ]]; then
+ UV_ARGS=(
+ "--no-install-package" "torch"
+ "--no-install-package" "torchvision"
+ "--no-install-package" "triton"
+ "--no-install-package" "nvidia-cublas-cu12"
+ "--no-install-package" "nvidia-cuda-cupti-cu12"
+ "--no-install-package" "nvidia-cuda-nvrtc-cu12"
+ "--no-install-package" "nvidia-cuda-runtime-cu12"
+ "--no-install-package" "nvidia-cudnn-cu12"
+ "--no-install-package" "nvidia-cufft-cu12"
+ "--no-install-package" "nvidia-cufile-cu12"
+ "--no-install-package" "nvidia-curand-cu12"
+ "--no-install-package" "nvidia-cusolver-cu12"
+ "--no-install-package" "nvidia-cusparse-cu12"
+ "--no-install-package" "nvidia-cusparselt-cu12"
+ "--no-install-package" "nvidia-nccl-cu12"
+ )
+ else
+ UV_ARGS=()
+ fi
+
+ # Install uv
+ UV_VERSION="0.7.2"
+ curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh
+
+ # Create virtual environment and install dependencies
+ uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages
+
+ # Install dependencies
+ uv sync --locked --only-group build ${UV_ARGS[@]}
+ uv sync \
+ --link-mode copy \
+ --locked \
+ --extra ${ENVIRONMENT} \
+ --all-groups ${UV_ARGS[@]}
+
+ # Install the package
+ uv pip install --no-deps -e .
+ else
+ python3 -m venv $UV_PROJECT_ENVIRONMENT
+ . $UV_PROJECT_ENVIRONMENT/bin/activate
+
+ pip install --pre --no-cache-dir --upgrade pip
+ pip install --pre --no-cache-dir torch pybind11 wheel_stub ninja wheel packaging "setuptools>=80"
+ pip install --pre --no-cache-dir --no-build-isolation .
+ fi
+
+}
+
+# Call the main function
+main "$@"
diff --git a/docker/common/install_source_wheels.sh b/docker/common/install_source_wheels.sh
new file mode 100644
index 00000000000..7eaaef2e46f
--- /dev/null
+++ b/docker/common/install_source_wheels.sh
@@ -0,0 +1,53 @@
+#!/bin/bash
+set -xeuo pipefail # Exit immediately if a command exits with a non-zero status
+
+INPUT_WHEEL_DIR=$(pwd)/wheels
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --input-wheel-dir)
+ INPUT_WHEEL_DIR="$2"
+ shift 2
+ ;;
+ --environment)
+ ENVIRONMENT="$2"
+ shift 2
+ ;;
+ *)
+ echo "Unknown option: $1"
+ echo "Usage: $0 --input-wheel-dir DIR"
+ exit 1
+ ;;
+ esac
+done
+
+# Check if required arguments are provided
+if [ -z "$INPUT_WHEEL_DIR" ] || [ -z "$ENVIRONMENT" ]; then
+ echo "Error: --input-wheel-dir and --environment are required"
+ echo "Usage: $0 --input-wheel-dir DIR --environment ENV"
+ exit 1
+fi
+
+if [ "$ENVIRONMENT" = "dev" ]; then
+ TE_WHEEL=$(ls $INPUT_WHEEL_DIR/transformer_engine*.whl) || true
+ [ -z "$TE_WHEEL" ] && TE_WHEEL=$(bash docker/common/build_te.sh --output-wheel-dir $INPUT_WHEEL_DIR | tail -n 1)
+fi
+
+MAMBA_WHEEL=$(ls $INPUT_WHEEL_DIR/mamba*.whl) || true
+[ -z "$MAMBA_WHEEL" ] && MAMBA_WHEEL=$(bash docker/common/build_mamba.sh --output-wheel-dir $INPUT_WHEEL_DIR | tail -n 1)
+
+CAUSALCONV1D_WHEEL=$(ls $INPUT_WHEEL_DIR/causal_conv1d*.whl) || true
+[ -z "$CAUSALCONV1D_WHEEL" ] && CAUSALCONV1D_WHEEL=$(bash docker/common/build_causalconv1d.sh --output-wheel-dir $INPUT_WHEEL_DIR | tail -n 1)
+
+# Override deps that are already present in the base image
+# only for dev
+if [ "$ENVIRONMENT" = "dev" ]; then
+ uv pip install --no-cache-dir --no-deps $TE_WHEEL
+fi
+
+# Install heavy optional deps like mamba, causalconv1d
+uv pip install --no-cache-dir \
+ $MAMBA_WHEEL \
+ $CAUSALCONV1D_WHEEL \
+ "setuptools>=80"
diff --git a/docker/patches/deepep.patch b/docker/patches/deepep.patch
new file mode 100644
index 00000000000..070e4581481
--- /dev/null
+++ b/docker/patches/deepep.patch
@@ -0,0 +1,13 @@
+diff --git a/setup.py b/setup.py
+index 63ce332..4e13462 100644
+--- a/setup.py
++++ b/setup.py
+@@ -37,7 +37,7 @@ if __name__ == '__main__':
+ '-Wno-sign-compare', '-Wno-reorder', '-Wno-attributes']
+ nvcc_flags = ['-O3', '-Xcompiler', '-O3']
+ sources = ['csrc/deep_ep.cpp', 'csrc/kernels/runtime.cu', 'csrc/kernels/layout.cu', 'csrc/kernels/intranode.cu']
+- include_dirs = ['csrc/']
++ include_dirs = ['csrc/', '/usr/local/cuda/include/cccl/']
+ library_dirs = []
+ nvcc_dlink = []
+ extra_link_args = []
diff --git a/docs/add_copyright_header.py b/docs/add_copyright_header.py
new file mode 100644
index 00000000000..9694ef84819
--- /dev/null
+++ b/docs/add_copyright_header.py
@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+"""One-off script to add NVIDIA copyright header to all .md files under docs/."""
+
+from pathlib import Path
+
+HEADER = """ Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved.
+ NVIDIA CORPORATION and its licensors retain all intellectual property
+ and proprietary rights in and to this software, related documentation
+ and any modifications thereto. Any use, reproduction, disclosure or
+ distribution of this software and related documentation without an express
+ license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""
+
+def main():
+ docs_dir = Path(__file__).resolve().parent
+ already_has = "Copyright (c) 2022-2026, NVIDIA CORPORATION"
+ count = 0
+ for path in sorted(docs_dir.rglob("*.md")):
+ content = path.read_text(encoding="utf-8")
+ if content.strip().startswith(already_has):
+ continue
+ new_content = HEADER + content
+ path.write_text(new_content, encoding="utf-8")
+ count += 1
+ print(path.relative_to(docs_dir))
+ print(f"\nUpdated {count} files.")
+
+if __name__ == "__main__":
+ main()
diff --git a/docs/advanced/index.md b/docs/advanced/index.md
new file mode 100644
index 00000000000..98ff0806cff
--- /dev/null
+++ b/docs/advanced/index.md
@@ -0,0 +1,14 @@
+
+
+# Discussions
+
+In-depth technical discussions and optimization guides:
+
+- [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
diff --git a/docs/api-backwards-compatibility-check.md b/docs/api-backwards-compatibility-check.md
new file mode 100644
index 00000000000..40f56ec0c00
--- /dev/null
+++ b/docs/api-backwards-compatibility-check.md
@@ -0,0 +1,344 @@
+---
+orphan: true
+---
+
+
+
+# API Backward Compatibility Checking
+
+## Overview
+
+Megatron Core uses automated API compatibility checking to ensure stable interfaces between releases. This prevents accidental breaking changes that could affect users upgrading between versions.
+
+## How It Works
+
+The compatibility checker:
+1. Compares the current code against the latest release
+2. Detects breaking changes in function signatures
+3. Fails CI if breaking changes are found (unless explicitly exempted)
+4. Runs automatically on every PR that modifies `megatron/core`
+
+## What Gets Checked
+
+### ✅ Breaking Changes Detected
+
+- **Parameter removed** - Removing a function parameter
+- **Parameter added without default** - Adding a required parameter
+- **Parameter order changed** - Changing the order of parameters
+- **Optional→Required** - Removing a default value from a parameter
+- **Function removed** - Deleting a public function
+- **Return type changed** - Changing the return type annotation (warning)
+
+### ⏭️ What Gets Skipped
+
+- **Test functions** - Functions starting with `test_`
+- **Exempt decorators** - Functions marked with `@internal_api`, `@experimental_api`, or `@deprecated`
+- **Excluded paths** - Code in `tests/`, `experimental/`, `legacy/`
+
+### ✅ Allowed Changes
+
+- **Adding optional parameters** - Adding parameters with default values
+- **Adding new functions** - New public APIs
+- **Making parameters optional** - Adding default values to required parameters
+
+## For Developers
+
+### Running Locally
+
+```bash
+# Install griffe
+pip install griffe
+
+# Check against latest release
+python scripts/check_api_backwards_compatibility.py --baseline core_r0.8.0
+
+# Check with verbose output
+python scripts/check_api_backwards_compatibility.py --baseline core_r0.8.0 -v
+
+# Compare two specific branches
+python scripts/check_api_backwards_compatibility.py --baseline core_r0.8.0 --current main
+```
+
+### Marking Functions as Exempt
+
+If you need to make breaking changes to internal or experimental APIs:
+
+#### Internal API (for internal implementation details)
+
+```python
+from megatron.core.utils import internal_api
+
+@internal_api
+def experimental_feature(x, y):
+ """
+ This API is experimental and may change.
+ NOT FOR EXTERNAL USE.
+ """
+ pass
+```
+
+**When to use `@internal_api`:**
+- Internal APIs not documented for external use
+- Experimental features explicitly marked as unstable
+- Functions in development that haven't been released yet
+
+#### Experimental API (for experimental features)
+
+```python
+from megatron.core.utils import experimental_api
+
+@experimental_api
+def new_experimental_feature(x, y):
+ """
+ This API is experimental and may change without notice.
+ """
+ pass
+```
+
+**When to use `@experimental_api`:**
+- Experimental features explicitly marked as unstable
+- New APIs under active development
+- Features that haven't been stabilized yet
+
+### Deprecating APIs
+
+For planned API changes, use the deprecation workflow:
+
+```python
+from megatron.core.backwards_compatibility_decorators import deprecated
+
+@deprecated(
+ version="1.0.0", # When deprecation starts
+ removal_version="2.0.0", # When it will be removed
+ alternative="new_function", # Recommended replacement
+ reason="Improved performance and cleaner API"
+)
+def old_function(x):
+ """This function is deprecated."""
+ pass
+```
+
+**Deprecation Timeline:**
+1. **Version N** - Add `@deprecated` decorator, function still works
+2. **Version N+1** - Keep function with deprecation warnings
+3. **Version N+2** - Remove function (users have been warned)
+
+### Handling CI Failures
+
+If the compatibility check fails on your PR:
+
+1. **Review the breaking changes** in the CI logs
+2. **Choose an action:**
+ - **Fix the code** - Revert the breaking change
+ - **Add exemption** - Use `@internal_api` if intentional
+ - **Use deprecation** - For planned API changes
+3. **Update your PR** with the fix
+
+## Examples
+
+### Example 1: Compatible Change
+
+```python
+# ✅ BEFORE (v1.0)
+def train_model(config, dataloader):
+ pass
+
+# ✅ AFTER (v1.1) - Added optional parameter
+def train_model(config, dataloader, optimizer="adam"):
+ pass
+```
+**Result:** ✅ Check passes
+
+---
+
+### Example 2: Breaking Change
+
+```python
+# BEFORE (v1.0)
+def train_model(config, dataloader, optimizer="adam"):
+ pass
+
+# ❌ AFTER (v1.1) - Removed parameter
+def train_model(config, dataloader):
+ pass
+```
+**Result:** ❌ Check fails - "Parameter 'optimizer' removed"
+
+---
+
+### Example 3: Exempt Internal API
+
+```python
+from megatron.core.utils import internal_api
+
+# BEFORE (v1.0)
+@internal_api
+def _internal_compute(x, y):
+ pass
+
+# ✅ AFTER (v1.1) - Can change freely
+@internal_api
+def _internal_compute(x, y, z): # Added parameter
+ pass
+```
+**Result:** ✅ Check passes (function is exempt)
+
+---
+
+### Example 4: Deprecation Workflow
+
+```python
+from megatron.core.utils import deprecated
+
+# Version 1.0 - Add deprecation
+@deprecated(
+ version="1.0.0",
+ removal_version="2.0.0",
+ alternative="train_model_v2"
+)
+def train_model(config):
+ """Old training function - DEPRECATED"""
+ pass
+
+def train_model_v2(config, **options):
+ """New improved training function"""
+ pass
+
+# Version 1.1 - Keep both (users migrate)
+# Version 2.0 - Remove train_model()
+```
+
+## Architecture
+
+```
+Developer commits code
+ ↓
+GitHub Actions triggers
+ ↓
+CI runs check_api_backwards_compatibility.py
+ ↓
+Script loads code via griffe:
+ • Baseline: latest release (e.g., core_r0.8.0)
+ • Current: PR branch
+ ↓
+Apply filtering:
+ • Skip @internal_api, @experimental_api, and @deprecated
+ • Skip private functions (_prefix)
+ • Skip test/experimental paths
+ ↓
+Griffe compares signatures:
+ • Parameters
+ • Types
+ • Return types
+ • Defaults
+ ↓
+Report breaking changes
+ ↓
+Exit: 0=pass, 1=fail
+ ↓
+CI fails if breaking changes detected
+```
+
+## Configuration
+
+### Customizing Filters
+
+Edit `scripts/check_api_backwards_compatibility.py`:
+
+```python
+# Add more exempt decorators
+EXEMPT_DECORATORS = [
+ "internal_api",
+ "experimental_api",
+ "deprecated",
+]
+
+# Add more path exclusions
+EXCLUDE_PATHS = {
+ "tests",
+ "experimental",
+ "legacy",
+ "your_custom_path", # ← Add here
+}
+```
+
+### Changing the Baseline
+
+The workflow auto-detects the latest `core_r*` tag. To manually specify:
+
+```yaml
+# In .github/workflows/check_api_backwards_compatibility_workflow.yml
+- name: Run compatibility check
+ run: |
+ python scripts/check_api_backwards_compatibility.py \
+ --baseline your_custom_baseline
+```
+
+## FAQ
+
+### Q: Why did my PR fail the compatibility check?
+
+**A:** Your code introduced breaking changes compared to the last release. Review the CI logs to see what changed.
+
+### Q: Can I disable the check for my PR?
+
+**A:** No, but you can mark specific functions as exempt using `@internal_api` or `@experimental_api`.
+
+### Q: What if I need to make a breaking change?
+
+**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).
+
+### Q: Does this check all of Megatron-LM?
+
+**A:** No, only `megatron/core/**` (Megatron Core). Legacy code is excluded.
+
+### Q: What about class methods?
+
+**A:** Yes, class methods are checked just like functions.
+
+### Q: Can I run this locally before pushing?
+
+**A:** Yes! Run `python scripts/check_api_backwards_compatibility.py --baseline core_r0.8.0`
+
+### Q: What if there's no release tag yet?
+
+**A:** The workflow will use `main` as the baseline. Update it once you have release tags.
+
+## Troubleshooting
+
+### Error: "griffe is not installed"
+
+```bash
+pip install griffe
+```
+
+### Error: "No core_r* tags found"
+
+The repository doesn't have release tags yet. The workflow will fall back to `main`.
+
+### False Positives
+
+If the checker reports a breaking change that isn't actually breaking, file an issue and use `@internal_api` as a temporary workaround.
+
+## References
+
+- **Script:** `scripts/check_api_backwards_compatibility.py`
+- **Workflow:** `.github/workflows/check_api_backwards_compatibility_workflow.yml`
+- **Decorators:** `megatron/core/backwards_compatibility_decorators.py`
+- **Griffe Documentation:** https://mkdocstrings.github.io/griffe/
+
+## Support
+
+For questions or issues:
+1. Check this documentation
+2. Review existing PRs with compatibility checks
+3. Ask in the Megatron-LM Slack/Discord
+4. File an issue on GitHub
+
diff --git a/docs/api-guide/core/datasets.md b/docs/api-guide/core/datasets.md
new file mode 100644
index 00000000000..d80c0183375
--- /dev/null
+++ b/docs/api-guide/core/datasets.md
@@ -0,0 +1,13 @@
+
+
+# datasets package
+
+```{include} ../../../megatron/core/datasets/readme.md
+```
diff --git a/docs/api-guide/core/dist_checkpointing.md b/docs/api-guide/core/dist_checkpointing.md
new file mode 100644
index 00000000000..c3dfc7aa257
--- /dev/null
+++ b/docs/api-guide/core/dist_checkpointing.md
@@ -0,0 +1,175 @@
+
+
+# dist_checkpointing package
+
+A library for saving and loading the distributed checkpoints.
+A *distributed checkpoint* in Megatron Core uses the ``torch_dist`` format,
+a custom checkpointing mechanism built on top of PyTorch's native
+checkpointing capabilities.
+
+A key property of distributed checkpoints is that a checkpoint saved under one
+parallel configuration (tensor, pipeline, or data parallelism) can be loaded
+under a different parallel configuration. This enables flexible scaling and
+resharding of models across heterogeneous training setups.
+
+Using the library requires defining sharded state_dict dictionaries with functions from *mapping* and *optimizer* modules.
+Those state dicts can be saved or loaded with a *serialization* module using strategies from *strategies* module.
+
+## Safe Checkpoint Loading
+
+Since **PyTorch 2.6**, the default behavior of `torch.load` is `weights_only=True`.
+This ensures that only tensors and allow-listed classes are loaded, reducing the risk of arbitrary code execution.
+
+If you encounter an error such as:
+
+```bash
+WeightsUnpickler error: Unsupported global: GLOBAL argparse.Namespace was not an allowed global by default.
+```
+
+you can fix it by explicitly allow-listing the missing class in your script:
+
+```python
+import torch, argparse
+
+torch.serialization.add_safe_globals([argparse.Namespace])
+```
+
+## Checkpointing Distributed Optimizer
+
+### Checkpoint Compatibility and Optimizer State Formats
+
+Beginning with **mcore v0.14**, the ``flattened_range`` attribute was removed from ``dist_checkpointing``. As a result:
+
+- Optimizer states saved with mcore versions <= 0.14 can no longer be loaded directly. Loading these legacy optimizer states is not supported because the required sharded metadata is no longer available. If you need to continue training from older checkpoints, refer to the workaround described below.
+- Model weights from older checkpoints remain fully compatible. No extra steps are needed—model weights from checkpoints created by earlier versions load automatically; simply add the ``--no-load-optim`` flag.
+
+### Workaround: Loading legacy optimizer states with ToT MCore
+
+**Step 1: Convert the legacy checkpoint using mcore v0.15.0**
+
+Run a dummy training job with mcore v0.15.0 to re-save the checkpoint with new optimizer states format.
+
+```bash
+MODEL_TRAIN_PARAMS=(
+ # Define model architecture and training parameters here
+)
+OLD_CKPT=/workspace/mcore_ckpt_old
+CONVERTED_CKPT=/workspace/mcore_ckpt_0.15.0
+
+torchrun --nproc_per_node=8 /opt/megatron-lm/pretrain_gpt.py \
+ --save-interval 1 \
+ --eval-interval 1 \
+ --exit-interval 1 \
+ --eval-iters 1 \
+ --use-distributed-optimizer \
+ --save ${CONVERTED_CKPT} \
+ --load ${OLD_CKPT} \
+ --ckpt-format torch_dist \
+ "${MODEL_TRAIN_PARAMS[@]}"
+```
+
+**Step 2: Load the converted checkpoint with ToT MCore**
+
+Use the converted checkpoint as the input for continued training with ToT MCore.
+
+```bash
+MODEL_TRAIN_PARAMS=(
+ # Define model architecture and training parameters here
+)
+NEW_CKPT=/workspace/mcore_ckpt_new
+CONVERTED_CKPT=/workspace/mcore_ckpt_0.15.0
+
+torchrun --nproc_per_node=8 /opt/megatron-lm/pretrain_gpt.py \
+ --use-distributed-optimizer \
+ --save ${NEW_CKPT} \
+ --load ${CONVERTED_CKPT} \
+ --ckpt-format torch_dist \
+ "${MODEL_TRAIN_PARAMS[@]}"
+```
+
+After this step, training can proceed normally using ToT MCore with fully supported optimizer state loading.
+
+## Distributed Optimizer Checkpoint Formats
+
+The refactor of the Distributed Optimizer introduces **two checkpoint formats**:
+
+- dp_reshardable (Default)
+ - Fast save/load performance.
+ - Not reshardable — not possible to change model parallelism when using this format.
+ - Recommended for general training when model parallelism changes are not needed.
+- fully_reshardable
+ - Fully reshardable — supports arbitrary changes in model parallelism.
+ - Slower than dp_reshardable.
+ - Enabled via the ``--dist-ckpt-optim-fully-reshardable`` flag.
+
+### Workflow for Changing Model Parallelism
+
+You can combine formats to optimize both flexibility and performance:
+
+ 1. Train using ``dp_reshardable`` (default) for faster checkpointing.
+ 2. When you need to change model parallelism:
+
+ - Stop training.
+ - Change model parallelism for train config.
+ - Resume training with ``--dist-ckpt-optim-fully-reshardable``.
+
+ 3. Save at least one checkpoint under the new model parallel configuration.
+ 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``.
+
+## Async Checkpoint Saving Strategy
+
+The framework supports asynchronous checkpoint saving to improve training performance by offloading I/O operations.
+
+We are transitioning to a new async saving implementation based on the **NVRx (NVIDIA Resiliency Extension)** package. The legacy async strategy (referred to as **mcore**) is being deprecated.
+
+### Migration to NVRx
+
+- The **NVRx-based async strategy** will become the **default** in mcore v0.17.
+- The existing **mcore async strategy** is **deprecated** and will be removed in future versions.
+- A deprecation warning is emitted when using the mcore strategy.
+
+### Selecting Async Strategy
+
+`--async-strategy` flag is introduced to control the async strategy. To use legacy async strategy (**mcore**), set:
+
+```bash
+--async-strategy mcore
+```
+
+### NVRx Dependency
+
+To use the NVRx-based async strategy, you must install the `nvidia-resiliency-ext` package.
+
+```bash
+git clone https://github.com/NVIDIA/nvidia-resiliency-ext
+cd nvidia-resiliency-ext
+pip install .
+```
+
+> NOTE
+
+- If `nvidia-resiliency-ext` is not installed, the NVRx async strategy will not be available.
+- The `mcore` strategy will remain temporarily to ensure a smooth transition but will be removed in future releases.
+- It is strongly recommended to migrate to the NVRx strategy as soon as possible.
+
+### Async Saving for `fsdp_dtensor` and `torch_dcp` checkpoints
+
+Starting from mcore v0.17, asynchronous checkpoint saving is supported for `fsdp_dtensor` and `torch_dcp` formats.
+
+Note that async saving for these formats requires the `nvidia-resiliency-ext` package. As a result, the only supported `async_strategy` in this context is `nvrx`.
+
+## Subpackages
+
+```{toctree}
+:maxdepth: 4
+
+dist_checkpointing.strategies
+```
+
diff --git a/docs/api-guide/core/dist_checkpointing.strategies.md b/docs/api-guide/core/dist_checkpointing.strategies.md
new file mode 100644
index 00000000000..22fe3517a54
--- /dev/null
+++ b/docs/api-guide/core/dist_checkpointing.strategies.md
@@ -0,0 +1,16 @@
+
+
+# dist_checkpointing.strategies package
+
+Package defining different checkpoint formats (backends) and saving/loading algorithms (strategies).
+
+Strategies 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.
+Strategies are passed to `dist_checkpointing.load` and `dist_checkpointing.save` functions and control the actual saving/loading procedure.
+
diff --git a/docs/api-guide/core/distributed.md b/docs/api-guide/core/distributed.md
new file mode 100644
index 00000000000..13da4285ec5
--- /dev/null
+++ b/docs/api-guide/core/distributed.md
@@ -0,0 +1,19 @@
+
+
+# distributed package
+
+This package contains various utilities to finalize model weight gradients
+on each rank before the optimizer step. This includes a distributed data
+parallelism wrapper to all-reduce or reduce-scatter the gradients across
+data-parallel replicas, and a `finalize_model_grads` method to
+synchronize gradients across different parallelism modes (e.g., 'tied'
+layers on different pipeline stages, or gradients for experts in a MoE on
+different ranks due to expert parallelism).
+
diff --git a/docs/api-guide/core/fusions.md b/docs/api-guide/core/fusions.md
new file mode 100644
index 00000000000..fdd358e813c
--- /dev/null
+++ b/docs/api-guide/core/fusions.md
@@ -0,0 +1,20 @@
+
+
+# fusions package
+
+This package provides modules that provide commonly fused
+operations. Fusing operations improves compute efficiency by
+increasing the amount of work done each time a tensor is read from
+memory. To perform the fusion, modules in this either rely on PyTorch
+functionality for doing just-in-time compilation
+(i.e. `torch.jit.script` in older PyTorch versions of `torch.compile`
+in recent versions), or call into custom kernels in external libraries
+such as Apex or TransformerEngine.
+
diff --git a/docs/api-guide/core/index.md b/docs/api-guide/core/index.md
new file mode 100644
index 00000000000..0d39e46e744
--- /dev/null
+++ b/docs/api-guide/core/index.md
@@ -0,0 +1,25 @@
+
+
+# Core APIs
+
+Low-level API reference for core Megatron components.
+
+```{toctree}
+:maxdepth: 2
+
+transformer
+tensor_parallel
+pipeline_parallel
+fusions
+distributed
+datasets
+dist_checkpointing
+dist_checkpointing.strategies
+```
diff --git a/docs/api-guide/core/pipeline_parallel.md b/docs/api-guide/core/pipeline_parallel.md
new file mode 100644
index 00000000000..35f3c5b5cc2
--- /dev/null
+++ b/docs/api-guide/core/pipeline_parallel.md
@@ -0,0 +1,16 @@
+
+
+# pipeline_parallel package
+
+This package contains implementations for two different pipeline parallelism
+schedules (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)
+for details), and a default no-pipelining schedule. It also contains methods
+for the point-to-point communication that is needed between pipeline stages.
+
diff --git a/docs/api-guide/core/tensor_parallel.md b/docs/api-guide/core/tensor_parallel.md
new file mode 100644
index 00000000000..2d41c5f4467
--- /dev/null
+++ b/docs/api-guide/core/tensor_parallel.md
@@ -0,0 +1,17 @@
+
+
+# tensor_parallel package
+
+This package contains an implementation for tensor parallelism in transformer
+models (see [Megatron-LM: Training Multi-Billion Parameter Language Models
+Using Model Parallelism](https://arxiv.org/abs/1909.08053) and [Reducing
+Activation Recomputation in Large Transformer Models](https://arxiv.org/abs/2205.05198)
+for details).
+
diff --git a/docs/api-guide/core/transformer.md b/docs/api-guide/core/transformer.md
new file mode 100644
index 00000000000..03bc0f501f4
--- /dev/null
+++ b/docs/api-guide/core/transformer.md
@@ -0,0 +1,19 @@
+
+
+# transformer package
+
+The `transformer` package provides a customizable and configurable
+implementation of the transformer model architecture. Each component
+of a transformer stack, from entire layers down to individual linear
+layers, can be customized by swapping in different PyTorch modules
+using the "spec" parameters. The
+configuration of the transformer (hidden size, number of layers,
+number of attention heads) is provided using a `TransformerConfig`
+object.
diff --git a/docs/api-guide/index.md b/docs/api-guide/index.md
new file mode 100644
index 00000000000..7afa2450dd0
--- /dev/null
+++ b/docs/api-guide/index.md
@@ -0,0 +1,20 @@
+
+
+# API Guide
+
+API reference documentation for Megatron Core components.
+
+```{toctree}
+:maxdepth: 3
+
+models/index
+core/index
+internal/index
+```
diff --git a/docs/api-guide/internal/index.md b/docs/api-guide/internal/index.md
new file mode 100644
index 00000000000..312081ce70b
--- /dev/null
+++ b/docs/api-guide/internal/index.md
@@ -0,0 +1,19 @@
+
+
+# Internal Utilities
+
+Internal utility APIs.
+
+```{toctree}
+:maxdepth: 2
+
+num_microbatches_calculator
+optimizer_param_scheduler
+```
diff --git a/docs/api-guide/internal/num_microbatches_calculator.md b/docs/api-guide/internal/num_microbatches_calculator.md
new file mode 100644
index 00000000000..0c223588ce5
--- /dev/null
+++ b/docs/api-guide/internal/num_microbatches_calculator.md
@@ -0,0 +1,13 @@
+
+
+# Microbatches Calculator
+
+This api is used to calculate the number of microbatches required to fit a given model on a given batch size.
+
diff --git a/docs/api-guide/internal/optimizer_param_scheduler.md b/docs/api-guide/internal/optimizer_param_scheduler.md
new file mode 100644
index 00000000000..45e5e4b7da1
--- /dev/null
+++ b/docs/api-guide/internal/optimizer_param_scheduler.md
@@ -0,0 +1,13 @@
+
+
+# Optimizer Parameters Scheduler
+
+This api is used to calculate the learning rate and weight decay for the optimizer.
+
diff --git a/docs/api-guide/models/index.md b/docs/api-guide/models/index.md
new file mode 100644
index 00000000000..e5bb531454b
--- /dev/null
+++ b/docs/api-guide/models/index.md
@@ -0,0 +1,21 @@
+
+
+# Model APIs
+
+API reference for Megatron Core model implementations.
+
+```{toctree}
+:maxdepth: 2
+
+models
+models.gpt
+models.bert
+models.t5
+```
diff --git a/docs/api-guide/models/models.bert.md b/docs/api-guide/models/models.bert.md
new file mode 100644
index 00000000000..1543f4df865
--- /dev/null
+++ b/docs/api-guide/models/models.bert.md
@@ -0,0 +1,13 @@
+
+
+# models.bert package
+
+Useful package for training bert and bert like encoder only models. It optionally comes with a binary head that can be used for classification tasks .
+
diff --git a/docs/api-guide/models/models.gpt.md b/docs/api-guide/models/models.gpt.md
new file mode 100644
index 00000000000..1c3cbb5484c
--- /dev/null
+++ b/docs/api-guide/models/models.gpt.md
@@ -0,0 +1,13 @@
+
+
+# models.gpt package
+
+This is the implementation of the popular GPT model. It supports several features like model parallelization (Tensor Parallel, Pipeline Parallel, Data Parallel) , mixture of experts, FP8 , Distributed optimizer etc. We are constantly adding new features. So be on the lookout or raise an issue if you want to have something added.
+
diff --git a/docs/api-guide/models/models.md b/docs/api-guide/models/models.md
new file mode 100644
index 00000000000..a633546f0c9
--- /dev/null
+++ b/docs/api-guide/models/models.md
@@ -0,0 +1,23 @@
+
+
+# models package
+
+This package contains most of the popular LLMs . Currently we have support for GPT, Bert, and T5 . This is an ever growing list so keep an eye out.
+
+## Subpackages
+
+```{toctree}
+:maxdepth: 4
+
+models.gpt
+models.t5
+models.bert
+```
+
diff --git a/docs/api-guide/models/models.t5.md b/docs/api-guide/models/models.t5.md
new file mode 100644
index 00000000000..4694b80113a
--- /dev/null
+++ b/docs/api-guide/models/models.t5.md
@@ -0,0 +1,11 @@
+
+
+# models.t5 package
+
diff --git a/docs/api-guide/router_replay.md b/docs/api-guide/router_replay.md
new file mode 100644
index 00000000000..88476ea7537
--- /dev/null
+++ b/docs/api-guide/router_replay.md
@@ -0,0 +1,186 @@
+
+
+# Design Document: MoE Router Replay Feature
+
+## 1. Overview
+
+This document provides a detailed description of the "Router Replay" feature implemented within the Megatron-LM Core for Mixture-of-Experts (MoE) models.
+
+This feature is designed to enhance determinism and analyzability in MoE model training and inference. It enables the model to load routing decisions from a predefined file and enforce their use during the forward pass, thereby bypassing the real-time routing computation.
+
+## 2. Motivation
+
+* **Determinism & Reproducibility**: In distributed training, MoE routing decisions can exhibit minor variations due to factors like floating-point precision. By replaying a fixed routing table, the MoE computation path is guaranteed to be identical across runs, which facilitates debugging and reproducing experimental results.
+* **Performance Profiling**: The router's own computation (e.g., logits calculation, top-k selection) incurs overhead. In replay mode, this part of the computation can be completely skipped, allowing for more precise isolation and profiling of performance bottlenecks within the Expert Layers themselves.
+* **Debugging Aid**: When issues arise in the model, fixing the routing decisions helps to isolate variables, making it easier to determine whether the problem lies with the routing mechanism or the expert computations.
+
+## 3. Design and Architecture
+
+The design follows the principles of being non-intrusive and on-demand, with the core idea of activating the replay logic only when explicitly requested by the user.
+
+* **Core Components**:
+ * `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.
+ * `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.
+
+* **Workflow**:
+ The feature supports different modes, such as recording and replaying, controlled by a `RouterReplayAction`.
+
+ 1. **Enabling the Feature**: The user sets `moe_enable_routing_replay` to `True` in the model configuration.
+ 2. **Initialization**: When `moe_enable_routing_replay` is true, each `TopKRouter` creates its own `RouterReplay` instance.
+ 3. **Mode Configuration**: The user must programmatically set the desired router replay action (e.g., `record`, `forward_replay`, `backward_replay`) on the `RouterReplay` instances.
+ 4. **Execution Flow (within a mini-batch)**:
+ * **Forward Pass**:
+ * For each micro-batch, the `topk_routing_with_score_function` checks the `router_replay_action`.
+ * **In `record` mode**: The dynamically computed `top-k` expert indices are captured and stored.
+ * **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.
+ * **Backward Pass**:
+ * For each micro-batch (processed in reverse order in pipeline parallelism), the `router_replay_action` is checked again.
+ * **In `backward_replay` mode**: The function retrieves the expert indices for the corresponding micro-batch by popping them from the `replay_backward_list`. This mode is intended for training recomputation (e.g., activation checkpointing and pipeline recompute) so the same routing decisions are used during recompute/backward as in forward, ensuring determinism and correctness.
+
+## 4. Implementation Details
+
+The implementation cleanly separates the replay logic from the router's core computation.
+
+* **`megatron/core/transformer/transformer_config.py`**:
+ * Adds the configuration option `moe_enable_routing_replay: bool = False`.
+
+* **`megatron/core/transformer/moe/moe_utils.py`**:
+ * Introduces the `RouterReplay` class to manage the state for recording and replaying routing decisions for a single MoE layer.
+ * `target_topk_idx`: An attribute holding the expert indices for the current micro-batch during forward replay mode.
+ * `recorded_topk_idx`: An attribute for storing the computed expert indices when in record mode.
+ * `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.
+ * `set_target_indices()`: A method to load the replay indices into `target_topk_idx` for the forward pass.
+ * `record_indices()`: A method to save the computed indices.
+ * The `topk_routing_with_score_function` is modified to contain the core logic. It checks the `router_replay_action` on the `router_replay` instance and accordingly performs one of the following actions: computes and records indices, replays indices from `target_topk_idx` (for forward), replays indices from `replay_backward_list` (for backward), or falls through to the default dynamic routing.
+
+### Training recompute usage
+
+- During forward replay, `set_target_indices()` prepares `replay_backward_list` so each micro-batch’s indices are available for recomputation.
+- During recompute/backward, set action to `REPLAY_BACKWARD` so indices are consumed in FIFO order to mirror the forward sequence.
+
+## 5. Usage Guide
+
+1. **Enable & Instantiate**
+ - Create one `RouterReplay` instance per MoE router layer when building the model.
+ - Optionally use the global helpers to set/clear actions across all layers.
+2. **Record Routing Decisions**
+ - Set action: `RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)`.
+ - Run the model; retrieve per-layer indices via `RouterReplay.get_recorded_data()` and persist.
+3. **Forward Replay**
+ - Load indices and distribute: `RouterReplay.set_replay_data(list_of_tensors)`.
+ - Set action: `RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)`.
+ - Run the model; dynamic top‑k is bypassed and target indices are used.
+4. **Backward Replay**
+ - For training recomputation (activation checkpointing or pipeline recompute), set action: `REPLAY_BACKWARD` during recomputation.
+ - Per micro‑batch indices are consumed from `replay_backward_list` in FIFO order.
+5. **Cleanup**
+ - Use `RouterReplay.clear_global_indices()`, `RouterReplay.clear_global_router_replay_action()`, and `RouterReplay.clear_global_router_replay_instances()` to restore default behavior and prevent memory leaks.
+
+### Quick usage with `topk_routing_with_score_function`
+
+```python
+import torch
+from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+from megatron.core.transformer.moe.moe_utils import topk_routing_with_score_function
+
+rr = RouterReplay()
+
+# Record
+RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
+logits = torch.randn(8, 16)
+probs_rec, routing_map_rec = topk_routing_with_score_function(
+ logits=logits, topk=2, use_pre_softmax=False, score_function="softmax", router_replay=rr,
+)
+recorded = rr.get_recorded_indices()
+torch.save(recorded, "/tmp/replay.pt")
+
+# Forward replay
+rr.clear_router_replay_action()
+rr.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+target = torch.load("/tmp/replay.pt")
+rr.set_target_indices(target)
+probs_rep, routing_map_rep = topk_routing_with_score_function(
+ logits=logits, topk=2, use_pre_softmax=False, score_function="softmax", router_replay=rr,
+)
+
+RouterReplay.clear_global_router_replay_action()
+RouterReplay.clear_global_indices()
+RouterReplay.clear_global_router_replay_instances()
+```
+
+## 6. Minimal Demo
+
+Here is a minimal code example showing how to use RouterReplay for recording and replaying:
+
+```python
+import torch
+import torch.distributed as dist
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.moe.router import TopKRouter
+from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+
+
+# Initialize distributed training
+if not dist.is_initialized():
+ dist.init_process_group(backend="nccl")
+
+# Create a transformer config with RouterReplay enabled
+config = TransformerConfig(
+ num_experts=8,
+ expert_model_parallel_size=1,
+ num_top_k=2,
+ moe_enable_routing_replay=True
+)
+
+# Create a TopKRouter instance
+router = TopKRouter(config)
+
+# Generate sample input (batch_size, sequence_length, hidden_size)
+logits = torch.randn(16, 32, 8).to(torch.cuda.current_device())
+
+# -----------------
+# 1. Recording Mode
+# -----------------
+print("=== Recording Mode ===")
+# Set global router replay action to RECORD
+RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
+
+# Perform routing
+routing_output = router.forward(logits)
+print(f"Recorded top-k indices shape: {routing_output.top_k_idx.shape}")
+
+# -----------------
+# 2. Forward Replay Mode
+# -----------------
+print("\n=== Forward Replay Mode ===")
+# Save recorded indices to a file
+torch.save(routing_output.top_k_idx, "/tmp/replay.pt")
+
+# Load indices from file and set as target for replay
+replay_indices = torch.load("/tmp/replay.pt")
+for router_instance in RouterReplay.global_router_replay_instances:
+ router_instance.target_topk_idx = replay_indices
+
+# Set global router replay action to REPLAY_FORWARD
+RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+
+# Perform routing again - this will use the replayed indices
+replay_routing_output = router.forward(logits)
+print(f"Replayed top-k indices shape: {replay_routing_output.top_k_idx.shape}")
+print(f"Are indices the same? {torch.equal(routing_output.top_k_idx, replay_routing_output.top_k_idx)}")
+
+
+# Clean up
+RouterReplay.clear_global_router_replay_action()
+RouterReplay.clear_global_indices()
+RouterReplay.clear_global_router_replay_instances()
+if dist.is_initialized():
+ dist.destroy_process_group()
+```
diff --git a/docs/autodoc2_docstrings_parser.py b/docs/autodoc2_docstrings_parser.py
new file mode 100644
index 00000000000..14b722de65b
--- /dev/null
+++ b/docs/autodoc2_docstrings_parser.py
@@ -0,0 +1,35 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from docutils import nodes
+from myst_parser.parsers.sphinx_ import MystParser
+from sphinx.ext.napoleon.docstring import GoogleDocstring
+
+
+class NapoleonParser(MystParser):
+ """Add support for Google style docstrings."""
+
+ def parse(self, input_string: str, document: nodes.document) -> None:
+ """Parse Google style docstrings."""
+
+ # Get the Sphinx configuration
+ config = document.settings.env.config
+
+ # Process with Google style
+ google_parsed = str(GoogleDocstring(input_string, config))
+
+ return super().parse(google_parsed, document)
+
+
+Parser = NapoleonParser
diff --git a/docs/broken_links_false_positives.json b/docs/broken_links_false_positives.json
new file mode 100644
index 00000000000..01377be5804
--- /dev/null
+++ b/docs/broken_links_false_positives.json
@@ -0,0 +1,3 @@
+{
+ "uri": "http://localhost:8080/"
+}
\ No newline at end of file
diff --git a/docs/conf.py b/docs/conf.py
new file mode 100644
index 00000000000..c18f453490d
--- /dev/null
+++ b/docs/conf.py
@@ -0,0 +1,133 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Configuration file for the Sphinx documentation builder.
+#
+# For the full list of built-in configuration values, see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+import os
+import sys
+
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+project = "Megatron Core"
+copyright = "2026, NVIDIA Corporation"
+author = "NVIDIA Corporation"
+release = "nightly"
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+extensions = [
+ "myst_parser", # For our markdown docs
+ "sphinx.ext.viewcode", # For adding a link to view source code in docs
+ "sphinx.ext.doctest", # Allows testing in docstrings
+ "sphinx.ext.napoleon", # For google style docstrings
+ "sphinx_copybutton", # For copy button in code blocks
+]
+
+# Check if we should skip autodoc generation
+# usage: SKIP_AUTODOC=true
+skip_autodoc = os.environ.get("SKIP_AUTODOC", "false").lower() == "true"
+
+if not skip_autodoc:
+ extensions.append("autodoc2") # Generates API docs
+
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
+
+# -- Options for MyST Parser (Markdown) --------------------------------------
+# MyST Parser settings
+myst_enable_extensions = [
+ "dollarmath", # Enables dollar math for inline math
+ "amsmath", # Enables LaTeX math for display mode
+ "colon_fence", # Enables code blocks using ::: delimiters instead of ```
+ "deflist", # Supports definition lists with term: definition format
+ "fieldlist", # Enables field lists for metadata like :author: Name
+ "tasklist", # Adds support for GitHub-style task lists with [ ] and [x]
+ "attrs_block", # Enables setting attributes on block elements using {#id .class key=val}
+]
+myst_heading_anchors = 5 # Generates anchor links for headings up to level 5
+
+# Suppress "more than one target found for cross-reference" warnings for Python symbols
+# that have the same name across multiple modules (e.g. DistributedDataParallelConfig,
+# ModelType). These are structural ambiguities in the codebase – the cross-reference
+# still resolves; Sphinx just cannot pick the unique target automatically.
+suppress_warnings = ["ref.python"]
+
+# -- Options for Autodoc2 ---------------------------------------------------
+sys.path.insert(0, os.path.abspath(".."))
+
+if not skip_autodoc:
+ autodoc2_packages = [
+ {
+ "path": "../megatron/core", # Path to your package relative to conf.py
+ "exclude_dirs": ["converters"], # list of directory names to exclude
+ }
+ ]
+ autodoc2_render_plugin = "myst" # Use MyST for rendering docstrings
+ autodoc2_output_dir = "apidocs" # Output directory for autodoc2 (relative to docs/)
+ # This is a workaround that uses the parser located in autodoc2_docstrings_parser.py to allow autodoc2 to
+ # render google style docstrings.
+ # Related Issue: https://github.com/sphinx-extensions2/sphinx-autodoc2/issues/33
+ autodoc2_docstring_parser_regexes = [
+ (r".*", "docs.autodoc2_docstrings_parser"),
+ ]
+ # Regex patterns whose values contain raw regex syntax (e.g. \p{L}) that docutils
+ # mis-parses as footnote/reference markup. Exclude them from the generated docs.
+ autodoc2_hidden_regexes = [
+ r".*\._PATTERN_TIKTOKEN.*",
+ ]
+
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_theme = "nvidia_sphinx_theme"
+html_theme_options = {
+ "switcher": {
+ "json_url": "../versions1.json",
+ "version_match": release,
+ },
+ "icon_links": [
+ {
+ "name": "GitHub",
+ "url": "https://github.com/NVIDIA/Megatron-LM/",
+ "icon": "fa-brands fa-github",
+ }
+ ],
+ "public_docs_features": os.environ.get("SKIP_PUBLIC_DOCS_FEATURES", "false").lower() != "true",
+}
+html_extra_path = ["project.json", "versions1.json"]
+
+# Github links are now getting rate limited from the Github Actions
+linkcheck_ignore = [
+ ".*github\\.com.*",
+ ".*githubusercontent\\.com.*",
+ "http://localhost.*",
+]
+
+# PyTorch docs use a JS-rendered frontend; anchor IDs are injected at runtime
+# and are not present in the static HTML that linkcheck fetches.
+linkcheck_anchors_ignore_for_url = [
+ r"https://docs\.pytorch\.org/.*",
+]
+
+# PyTorch docs anchor IDs change between stable versions; verify the page
+# loads but skip anchor validation to avoid spurious failures on redirects.
+linkcheck_anchors_ignore_for_url = [
+ "https://docs.pytorch.org/.*",
+]
diff --git a/docs/developer/contribute.md b/docs/developer/contribute.md
new file mode 100644
index 00000000000..30a39e1cbc0
--- /dev/null
+++ b/docs/developer/contribute.md
@@ -0,0 +1,70 @@
+
+
+# Contributing to Megatron-LM
+
+This document outlines the processes and policies for issues and pull requests by non-NVIDIA contributors to the Megatron-LM GitHub repository.
+
+Everyone is welcome to contribute to the project! We recently migrated from using an internal repo to doing all development directly from the GitHub repository.
+
+When 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.**
+
+## Issue policy
+
+Please do file any bugs you find, keeping the following in mind:
+
+- If filing a bug, i.e. you have found something that doesn't work as expected, use the BUG template.
+- If you've found a regression in speed or accuracy use the REGRESSION template.
+- If you are requesting a new feature or modification of an existing feature use the ENHANCEMENT template.
+- If opening an issue to ask a question no template is needed but please make your question as clear and concise as possible.
+- One issue per bug. Putting multiple things in the same issue makes both discussion and completion unnecessarily complicated.
+- Your bug is mostly likely to get attention from the development team quickly if we can easily reproduce it.
+- Use proper spelling, grammar, and punctuation.
+- Write in an authoritative and technical tone.
+
+## Code submission policy
+
+### Do
+
+- 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.
+- Split your changes into separate, atomic commits i.e. A commit per feature or fix.
+- Make sure your commits are rebased on the master branch.
+- Write the commit message subject line in the imperative mood ("Change the default argument for X", not "Changed the default argument for X").
+- Write your commit messages in proper English, with care and punctuation.
+- Check the spelling of your code, comments and commit messages.
+
+### Don't
+
+- Submit code that's incompatible with the project licence.
+- Touch anything outside the stated scope of the PR. This includes formatting changes to code not relevant to the PR.
+- Iterate excessively on your design across multiple commits.
+- Include commented-out code.
+- Attempt large architectural changes without first opening an issue to discuss.
+
+## Issue and Pull Request Q&A
+
+### I've submitted an issue and PR. When can I expect to get some feedback?
+
+You should receive a response within 2 business days.
+
+### I need help, who should I ping?
+
+Use [@mcore-oncall](https://github.com/orgs/NVIDIA/teams/mcore-oncall).
+
+### If my issue or PR isn't getting attention, what should I do?
+
+After 2 business days, tag the user [@mcore-oncall](https://github.com/orgs/NVIDIA/teams/mcore-oncall).
+
+### Is there a policy for issues and PRs that haven't been touched in X days? Should they be closed?
+
+Yes, we have a bot that will mark untouched PRs as "stale" after 60 days.
+
+We 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.
+
+Thank you!
\ No newline at end of file
diff --git a/docs/developer/generate_docs.md b/docs/developer/generate_docs.md
new file mode 100644
index 00000000000..810c630681e
--- /dev/null
+++ b/docs/developer/generate_docs.md
@@ -0,0 +1,22 @@
+
+
+# Generating Docs Locally
+
+To generate docs locally, use the following commands:
+
+```
+cd docs
+SKIP_PUBLIC_DOCS_FEATURES=true uv run --only-group docs sphinx-autobuild . _build/html --port 8080 --host 127.0.0.1
+```
+
+Docs will be generated at .
+
+**Recommended:** set the environment variable `SKIP_AUTODOC=true` when generating docs
+to skip the generation of `apidocs`.
\ No newline at end of file
diff --git a/docs/developer/oncall.md b/docs/developer/oncall.md
new file mode 100644
index 00000000000..18d76f1436a
--- /dev/null
+++ b/docs/developer/oncall.md
@@ -0,0 +1,59 @@
+
+-->
+
+# Oncall Overview
+
+During your oncall week, you will be assigned to all PRs marked “Ready for
+Review”. From a high-level, your responsibilities include:
+
+- Review all new PRs
+- Accelerate the review process
+- Ensure issues and discussion questions are answered
+
+## PR Responsibilities
+
+Below is the checklist that the oncall needs to go through for each PR.
+
+- Should the PR remain a single PR?
+ - Each PR should have at most 1 expert reviewer, although there will be some outlier cases
+- Label PR as “complexity: low”, “complexity: medium”, or “complexity: high” depending on complexity
+ - Expert reviewers have final say, oncall just sets the initial complexity level
+ - Initial complexity level guideline
+ - Low: <100 lines changed
+ - Medium: 100 < lines changed < 500
+ - High: > 500 lines changed
+- Does this PR have proper testing coverage?
+ - If new logic is added, is the new logic tested?
+- Should the PR add documentation for any new features?
+- Does the PR conform to our style guidelines?
+ - Code structure
+ - Cleanliness
+ - Comments
+ - File structure
+- Do all tests pass?
+ - Oncall will need to kick off testing suite for external reviewers
+ - Comment “/ok to test commid_id” to kick off testing suite
+- Expert reviewers are notified after the PR is marked “Ready for Review”
+ - **Expert reviewers should review within 1 business day.** Message the assigned reviewer if it is taking longer. The reviewer either needs to review the PR or suggest an alternate reviewer.
+ - If the reviewer is not responding after 2 business days, escalate to the reviewer’s manager.
+- For `megatron/core` PRs, the “Final Review” label is applied automatically once all expert reviewers approve
+ - Final reviewers should review within 1 business day. Message the assigned reviewer if it is taking longer.
+ - If the reviewer is not responding after 2 business days, escalate to the reviewer’s manager.
+- The “Approved” label is applied automatically once all required reviewers have approved
+
+## Issues and Discussion Questions
+
+If you do not know the answer to an issue or discussion question, that's ok, **Delegate to someone who does.**
+
+On a daily basis, track the following:
+
+- [Dashboard for out of SLA issues](https://github.com/NVIDIA/Megatron-LM/issues?q=is%3Aissue%20state%3Aopen%20label%3Awaiting-on-maintainers).
+
+
diff --git a/docs/developer/submit.md b/docs/developer/submit.md
new file mode 100644
index 00000000000..205e18cc52f
--- /dev/null
+++ b/docs/developer/submit.md
@@ -0,0 +1,34 @@
+
+
+# How to Submit a PR
+
+All PRs start as **draft**. If you open a non-draft PR, it will be automatically converted to draft.
+
+## Step 1: Mark PR as "Ready for Review"
+
+1. When your PR is ready, click **Ready for Review**.
+2. The oncall reviewer is auto-assigned and expert reviewers are notified based on your changes. They will get notified and pick up your PR soon.
+
+:warning: Only mark as ready once all merge-conflicts are resolved and the CI is passing.
+Final Review might get declined if these requirements are not fulfilled.
+
+## Step 2: Final Review (`megatron/core` only)
+
+For PRs that change `megatron/core`, once all expert reviewers have approved, the `Final Review` label is applied **automatically** and final reviewers are assigned.
+
+For PRs outside `megatron/core`, this step is skipped.
+
+## Step 3: Approved
+
+Once all required reviewers have approved, the `Approved` label is applied **automatically**. The PR is now ready to merge.
+
+## Step 4: Merge
+
+Any member of [mcore-engineers](https://github.com/orgs/NVIDIA/teams/mcore-engineers) will be able to merge your PR.
diff --git a/docs/discussions/README.md b/docs/discussions/README.md
new file mode 100644
index 00000000000..aab65fc65ca
--- /dev/null
+++ b/docs/discussions/README.md
@@ -0,0 +1,31 @@
+---
+orphan: true
+---
+
+
+
+# Megatron Discussions
+
+This directory contains in-depth guides, tutorials, and discussions about optimizing and using Megatron for various use cases.
+
+## Available Guides
+
+### Training Guides
+
+## Contributing
+
+To contribute a guide or tutorial, follow this structure:
+
+1. Create a new directory: `docs/discussions/your-guide-name/`
+2. Add your main guide: `docs/discussions/your-guide-name/your-guide-name.md`
+3. Create an images directory: `docs/discussions/your-guide-name/images/`
+4. Update this README.md with a link to your guide
+
+Each guide should be self-contained with its own images and supporting files.
\ No newline at end of file
diff --git a/docs/distrib_optimizer.md b/docs/distrib_optimizer.md
deleted file mode 100644
index def23b20ebe..00000000000
--- a/docs/distrib_optimizer.md
+++ /dev/null
@@ -1,54 +0,0 @@
-# Distributed Optimizer
-
-The motivation for the distributed optimizer is to save memory by distributing the optimizer state evenly across data parallel ranks, versus the current method of replicating the optimizer state across data parallel ranks. As described in https://arxiv.org/abs/1910.02054, this branch specifically implements the following:
-
-- [yes] distribute all 'non-overlapping' optimizer state (i.e., model params already in fp32 are NOT distributed)
-- [no] distribute model gradients
-- [no] distribute model parameters
-
-Theoretical memory savings vary depending on the combination of the model's param dtype and grad dtype. In the current implementation, the theoretical number of bytes per parameter is (where 'd' is the data parallel size):
-
-| | Non-distributed optim | Distributed optim |
-| ------ | ------ | ------ |
-| float16 param, float16 grads | 20 | 4 + 16/d |
-| float16 param, fp32 grads | 18 | 6 + 12/d |
-| fp32 param, fp32 grads | 16 | 8 + 8/d |
-
-The implementation of the distributed optimizer is centered on using the contiguous grad buffer for communicating grads & params between the model state and the optimizer state. The grad buffer at any given moment either holds:
-
-1. all model grads
-2. a 1/d size _copy_ of the main grads (before copying to the optimizer state)
-3. a 1/d size _copy_ of the main params (after copying from the optimizer state)
-4. all model params
-5. zeros (or None), between iterations
-
-The grad buffer is used for performing reduce-scatter and all-gather operations, for passing grads & params between the model state and optimizer state. With this implementation, no dynamic buffers are allocated.
-
-The figures below illustrate the grad buffer's sharding scheme, and the key steps of the distributed optimizer's param update:
-
-## Data flow
-
-
-
-## Sharding scheme
-
-
-
-## Key steps
-
-_(note: using illustrations above, and assuming fp16 grads)_
-
-- Backward pass finishes (grad buffer holds 16 fp16 grad elements)
-- Call reduce-scatter on each DP rank
-- Each DP rank now has 4 elements within the grad buffer that are fully reduced (remaining 12 elements are garbage)
-- Each DP rank copies its relevant 4 fp16 grad elements from the grad buffer into 4 fp32 main grad elements (separate buffer, owned by the optimizer); i.e.
- - DP rank 0 copies elements [0:4]
- - DP rank 1 copies elements [4:8]
- - DP rank 2 copies elements [8:12]
- - DP rank 3 copies elements [12:16]
-- Optimizer.step()
-- Each DP rank copies its 4 fp32 main (/optimizer) param elements into the corresponding 4 fp16 elements in the grad buffer
-- Call all-gather on each DP rank
-- Grad buffer now contains all 16, fully updated, fp16 model param elements
-- Copy updated model params from grad buffer into their respective param tensors
-- (At this point, grad buffer is ready to be zero'd for the next iteration)
diff --git a/docs/documentation.md b/docs/documentation.md
new file mode 100644
index 00000000000..d2554157b45
--- /dev/null
+++ b/docs/documentation.md
@@ -0,0 +1,69 @@
+---
+orphan: true
+---
+
+
+
+# Documentation Development
+
+- [Documentation Development](#documentation-development)
+ - [Build the Documentation](#build-the-documentation)
+ - [Live Building](#live-building)
+ - [Documentation Version](#documentation-version)
+
+
+## Build the Documentation
+
+The following sections describe how to set up and build the NeMo RL documentation.
+
+Switch to the documentation source folder and generate HTML output.
+
+```sh
+cd docs/
+uv run --group docs sphinx-build . _build/html
+```
+
+* The resulting HTML files are generated in a `_build/html` folder that is created under the project `docs/` folder.
+* The generated python API docs are placed in `apidocs` under the `docs/` folder.
+
+## Checking for Broken Links
+
+To check for broken http links in the docs, run this command:
+
+```sh
+cd docs/
+uv run --group docs sphinx-build --builder linkcheck . _build/linkcheck
+```
+
+It will output a JSON file at `_build/linkcheck/output.json` with links it found while building the
+docs. Records will have a status of `broken` if the link is not reachable. The `docs/conf.py` file is
+configured to ignore github links because the CI test will often experience rate limit errors.
+Comment out the `linkcheck_ignore` variable there to check all the links.
+
+## Live Building
+
+When writing documentation, it can be helpful to serve the documentation and have it update live while you edit.
+
+To do so, run:
+
+```sh
+cd docs/
+uv run --group docs sphinx-autobuild . _build/html --port 12345 --host 0.0.0.0
+```
+
+Open a web browser and go to `http://${HOST_WHERE_SPHINX_COMMAND_RUN}:12345` to view the output.
+
+## Documentation Version
+
+The 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.
+
+* docs/versions1.json
+* docs/project.json
+* docs/conf.py
diff --git a/docs/get-started/install.md b/docs/get-started/install.md
new file mode 100644
index 00000000000..3e60a1fbb81
--- /dev/null
+++ b/docs/get-started/install.md
@@ -0,0 +1,123 @@
+
+
+# Installation
+
+## System Requirements
+
+### Hardware
+
+- **Recommended**: NVIDIA Turing architecture or later
+- **FP8 Support**: Requires NVIDIA Hopper, Ada, or Blackwell GPUs
+
+### Software
+
+- **Python**: >= 3.10 (3.12 recommended)
+- **PyTorch**: >= 2.6.0
+- **CUDA Toolkit**: Latest stable version
+
+
+## Prerequisites
+
+Install [uv](https://docs.astral.sh/uv/), a fast Python package installer:
+
+```bash
+curl -LsSf https://astral.sh/uv/install.sh | sh
+```
+
+
+## Option A: Pip Install (Recommended)
+
+Install the latest stable release from PyPI:
+
+```bash
+uv pip install megatron-core
+```
+
+To include optional training dependencies (Weights & Biases, SentencePiece, HF Transformers):
+
+```bash
+uv pip install "megatron-core[training]"
+```
+
+For all extras including [Transformer Engine](https://github.com/NVIDIA/TransformerEngine):
+
+```bash
+uv pip install --group build
+uv pip install --no-build-isolation "megatron-core[training,dev]"
+```
+
+```{note}
+`--no-build-isolation` requires build dependencies to be pre-installed in the environment. `torch` is needed because several `[dev]` packages (`mamba-ssm`, `nv-grouped-gemm`, `transformer-engine`) import it at build time to compile CUDA kernels. Expect this step to take **20+ minutes** depending on your hardware. If you prefer pre-built binaries, the [NGC Container](#option-c-ngc-container) ships with these pre-compiled.
+```
+
+```{warning}
+Building from source can consume a large amount of memory. By default the build runs one compiler job per CPU core, which can cause out-of-memory failures on machines with many cores. To limit parallel compilation jobs, set the `MAX_JOBS` environment variable before installing (for example, `MAX_JOBS=4`).
+```
+
+```{tip}
+For a lighter set of development dependencies without Transformer Engine and ModelOpt, use `[lts]` instead of `[dev]`: `uv pip install --no-build-isolation "megatron-core[training,lts]"`. The `[lts]` and `[dev]` extras are mutually exclusive.
+```
+
+To clone the repository for examples:
+
+```bash
+git clone https://github.com/NVIDIA/Megatron-LM.git
+```
+
+
+## Option B: Install from Source
+
+For development or to run the latest unreleased code:
+
+```bash
+git clone https://github.com/NVIDIA/Megatron-LM.git
+cd Megatron-LM
+uv pip install -e .
+```
+
+To install with all development dependencies (includes Transformer Engine, requires pre-installed build deps):
+
+```bash
+uv pip install --group build
+uv pip install --no-build-isolation -e ".[training,dev]"
+```
+
+```{tip}
+If the build runs out of memory, limit parallel compilation jobs with `MAX_JOBS=4 uv pip install --no-build-isolation -e ".[training,dev]"`.
+```
+
+
+## Option C: NGC Container
+
+For a pre-configured environment with all dependencies pre-installed (PyTorch, CUDA, cuDNN, NCCL, Transformer Engine), use the [PyTorch NGC Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch).
+
+We recommend using the **previous month's** NGC container rather than the latest one to ensure compatibility with the current Megatron Core release and testing matrix.
+
+```bash
+docker run --gpus all -it --rm \
+ -v /path/to/dataset:/workspace/dataset \
+ -v /path/to/checkpoints:/workspace/checkpoints \
+ -e PIP_CONSTRAINT= \
+ nvcr.io/nvidia/pytorch:26.01-py3
+```
+
+```{note}
+The NGC PyTorch container constrains the Python environment globally using `PIP_CONSTRAINT`. The `-e PIP_CONSTRAINT=` flag above unsets this so that Megatron Core and its dependencies install correctly.
+```
+
+Then install Megatron Core inside the container (torch is already available in the NGC image):
+
+```bash
+pip install uv
+uv pip install --no-build-isolation "megatron-core[training,dev]"
+```
+
+
+You are now ready to run training. Refer to [Your First Training Run](quickstart.md) for next steps.
diff --git a/docs/get-started/overview.md b/docs/get-started/overview.md
new file mode 100644
index 00000000000..5ceddcb1f41
--- /dev/null
+++ b/docs/get-started/overview.md
@@ -0,0 +1,88 @@
+
+
+# Overview
+
+Megatron-Core and Megatron-LM are open-source tools that are typically used together to train LLMs at scale across GPUs. Megatron-Core expands the capability of Megatron-LM. Megatron Bridge connects Megatron-Core and Megatron-LM to other popular training models, such as Hugging Face.
+
+## Megatron Core
+
+NVIDIA 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 high throughput at scale across thousands of GPUs. It provides an extensive set of tools for multimodal and speech AI. It expands Megatron-LM capabilities.
+
+Megatron-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.
+
+Megatron-Core is compatible with all NVIDIA Tensor Core GPUs and popular LLM architectures such as GPT, BERT, T5, and RETRO.
+
+
+**Composable library** with GPU-optimized building blocks for custom training frameworks.
+
+**Best for:**
+
+- **Framework developers** building on top of modular and optimized components
+- **Research teams** needing custom training loops, optimizers, or data pipelines
+- **ML engineers** requiring fault-tolerant training pipelines
+
+**What you get:**
+
+- Composable transformer building blocks (attention, MLP)
+- Advanced parallelism strategies (TP, PP, DP, EP, CP)
+- Pipeline schedules and distributed optimizers
+- Mixed precision support (FP16, BF16, FP8)
+- GPU-optimized kernels and memory management
+- High-performance dataloaders and dataset utilities
+- Model architectures (LLaMA, Qwen, GPT, Mixtral, Mamba)
+
+## Megatron-LM
+
+Megatron-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 direct entry point for exploring Megatron-Core.**
+
+It 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.
+
+**Reference implementation** that includes Megatron Core plus everything needed to train models.
+
+**Best for:**
+
+- **Training large foundation models** at scale with strong performance on the latest NVIDIA hardware
+- **Research teams** exploring new architectures and training techniques
+- **Learning distributed training** concepts and best practices
+- **Quick experimentation** with proven model configurations
+
+**What you get:**
+
+- Pre-configured training scripts for GPT, LLaMA, DeepSeek, Qwen, and more.
+- End-to-end examples from data prep to evaluation
+- Research-focused tools and utilities
+
+
+
+## Megatron Bridge
+
+Megatron Bridge provides out-of-the-box bridges and training recipes for models built on top of base model architectures from Megatron Core.
+
+Megatron Bridge provides a 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.
+
+After training or modifying a Megatron model, you can convert it again for deployment or sharing. Refer to the [Megatron Bridge repository](https://github.com/NVIDIA-NeMo/Megatron-Bridge) for the code and training recipes.
+
+## Ecosystem Libraries
+
+**Libraries used by Megatron Core:**
+
+- **[Megatron Energon](https://github.com/NVIDIA/Megatron-Energon)** - Multi-modal data loader (text, images, video, audio) with distributed loading and dataset blending
+- **[Transformer Engine](https://github.com/NVIDIA/TransformerEngine)** - Optimized kernels and FP8 mixed precision support
+- **[Resiliency Extension (NVRx)](https://github.com/NVIDIA/nvidia-resiliency-ext)** - Fault tolerant training with failure detection and recovery
+
+**Libraries using Megatron Core:**
+
+- **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Training library with bidirectional checkpoint conversion between Hugging Face and Megatron, customizable training loops, and production-ready recipes
+- **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods
+- **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples
+- **[Model Optimizer (ModelOpt)](https://github.com/NVIDIA/Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Check out end-to-end examples in [examples/post_training/modelopt](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/post_training/modelopt).
+
+**Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed)
+
diff --git a/docs/get-started/quickstart.md b/docs/get-started/quickstart.md
new file mode 100644
index 00000000000..9d68016ece7
--- /dev/null
+++ b/docs/get-started/quickstart.md
@@ -0,0 +1,68 @@
+
+
+# Your First Training Run
+
+This guide walks you through running your first training jobs with Megatron Core. Make sure you have completed [installation](install.md) before proceeding.
+
+## Minimal Training Example
+
+Run a minimal distributed training loop with mock data on 2 GPUs:
+
+```bash
+torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py
+```
+
+## LLaMA-3 Training Example
+
+Train an LLaMA-3 8B model with FP8 precision on 8 GPUs using mock data:
+
+```bash
+./examples/llama/train_llama3_8b_h100_fp8.sh
+```
+
+## Data Preparation
+
+To train on your own data, Megatron expects preprocessed binary files (`.bin` and `.idx`).
+
+### 1. Prepare a JSONL File
+
+Each line should contain a `text` field:
+
+```json
+{"text": "Your training text here..."}
+{"text": "Another training sample..."}
+```
+
+### 2. Preprocess the Data
+
+```bash
+python tools/preprocess_data.py \
+ --input data.jsonl \
+ --output-prefix processed_data \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model /path/to/tokenizer.model \
+ --workers 8 \
+ --append-eod
+```
+
+### Key Arguments
+
+- `--input`: Path to input JSON/JSONL file
+- `--output-prefix`: Prefix for output binary files (.bin and .idx)
+- `--tokenizer-type`: Tokenizer type (`HuggingFaceTokenizer`, `GPT2BPETokenizer`, and so on)
+- `--tokenizer-model`: Path to tokenizer model file
+- `--workers`: Number of parallel workers for processing
+- `--append-eod`: Add end-of-document token
+
+## Next Steps
+
+- Explore [Parallelism Strategies](../user-guide/parallelism-guide.md) to scale your training
+- Learn about [Data Preparation](../user-guide/data-preparation.md) best practices
+- Check out [Advanced Features](../user-guide/features/index.md)
diff --git a/docs/get-started/releasenotes.md b/docs/get-started/releasenotes.md
new file mode 100644
index 00000000000..e624de19f15
--- /dev/null
+++ b/docs/get-started/releasenotes.md
@@ -0,0 +1,19 @@
+
+
+# Release Notes
+
+
+## Roadmaps
+
+Stay up-to-date with our development roadmaps and planned features:
+
+- **[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
+- **[GPT-OSS Implementation Tracker](https://github.com/NVIDIA/Megatron-LM/issues/1739)** - Advanced features including YaRN RoPE scaling, attention sinks, and custom activation functions
+
diff --git a/docs/images/context_parallel/CP_overview.png b/docs/images/context_parallel/CP_overview.png
new file mode 100644
index 00000000000..38c55b371aa
Binary files /dev/null and b/docs/images/context_parallel/CP_overview.png differ
diff --git a/docs/images/context_parallel/CP_results.png b/docs/images/context_parallel/CP_results.png
new file mode 100644
index 00000000000..e0415ce86eb
Binary files /dev/null and b/docs/images/context_parallel/CP_results.png differ
diff --git a/docs/images/distrib_optimizer/data_flow.png b/docs/images/distrib_optimizer/data_flow.png
index d48fc134c40..01f5cfb2e7e 100644
Binary files a/docs/images/distrib_optimizer/data_flow.png and b/docs/images/distrib_optimizer/data_flow.png differ
diff --git a/docs/images/distrib_optimizer/sharding_scheme.png b/docs/images/distrib_optimizer/sharding_scheme.png
index b07c25b05f9..e48dd95024a 100644
Binary files a/docs/images/distrib_optimizer/sharding_scheme.png and b/docs/images/distrib_optimizer/sharding_scheme.png differ
diff --git a/docs/images/fine_grained_activation_offloading/offloading_and_recomputing.png b/docs/images/fine_grained_activation_offloading/offloading_and_recomputing.png
new file mode 100644
index 00000000000..6c8afa78bb1
Binary files /dev/null and b/docs/images/fine_grained_activation_offloading/offloading_and_recomputing.png differ
diff --git a/docs/images/megatron_fsdp/DDP_vs_FSDP.png b/docs/images/megatron_fsdp/DDP_vs_FSDP.png
new file mode 100644
index 00000000000..627821439e2
Binary files /dev/null and b/docs/images/megatron_fsdp/DDP_vs_FSDP.png differ
diff --git a/docs/images/megatron_fsdp/FSDP_Allreduce.png b/docs/images/megatron_fsdp/FSDP_Allreduce.png
new file mode 100644
index 00000000000..66e2391ed04
Binary files /dev/null and b/docs/images/megatron_fsdp/FSDP_Allreduce.png differ
diff --git a/docs/images/megatron_fsdp/fsdp_double_buffer.png b/docs/images/megatron_fsdp/fsdp_double_buffer.png
new file mode 100644
index 00000000000..fbfbcef9b28
Binary files /dev/null and b/docs/images/megatron_fsdp/fsdp_double_buffer.png differ
diff --git a/docs/images/megatron_fsdp/fsdp_streams.png b/docs/images/megatron_fsdp/fsdp_streams.png
new file mode 100644
index 00000000000..6b8840783c8
Binary files /dev/null and b/docs/images/megatron_fsdp/fsdp_streams.png differ
diff --git a/docs/images/megatron_fsdp/fsdp_v_hfsdp_streams.png b/docs/images/megatron_fsdp/fsdp_v_hfsdp_streams.png
new file mode 100644
index 00000000000..6f6e61dfb21
Binary files /dev/null and b/docs/images/megatron_fsdp/fsdp_v_hfsdp_streams.png differ
diff --git a/docs/images/megatron_fsdp/hfsdp.png b/docs/images/megatron_fsdp/hfsdp.png
new file mode 100644
index 00000000000..3c056d20689
Binary files /dev/null and b/docs/images/megatron_fsdp/hfsdp.png differ
diff --git a/docs/images/megatron_fsdp/lcm_dim0_shard.png b/docs/images/megatron_fsdp/lcm_dim0_shard.png
new file mode 100644
index 00000000000..910add676f1
Binary files /dev/null and b/docs/images/megatron_fsdp/lcm_dim0_shard.png differ
diff --git a/docs/images/megatron_fsdp/mixed_sharding.png b/docs/images/megatron_fsdp/mixed_sharding.png
new file mode 100644
index 00000000000..81cbc153f8a
Binary files /dev/null and b/docs/images/megatron_fsdp/mixed_sharding.png differ
diff --git a/docs/images/megatron_fsdp/quantized_param_gather.png b/docs/images/megatron_fsdp/quantized_param_gather.png
new file mode 100644
index 00000000000..e1908e66ad7
Binary files /dev/null and b/docs/images/megatron_fsdp/quantized_param_gather.png differ
diff --git a/docs/images/megatron_fsdp/sharded_quantization.png b/docs/images/megatron_fsdp/sharded_quantization.png
new file mode 100644
index 00000000000..c65bab5305a
Binary files /dev/null and b/docs/images/megatron_fsdp/sharded_quantization.png differ
diff --git a/docs/images/megatron_fsdp/uneven_sharding.png b/docs/images/megatron_fsdp/uneven_sharding.png
new file mode 100644
index 00000000000..0c34a51b026
Binary files /dev/null and b/docs/images/megatron_fsdp/uneven_sharding.png differ
diff --git a/docs/images/megatron_fsdp/zero3_model_state.png b/docs/images/megatron_fsdp/zero3_model_state.png
new file mode 100644
index 00000000000..84ad33ff779
Binary files /dev/null and b/docs/images/megatron_fsdp/zero3_model_state.png differ
diff --git a/docs/images/moe/token_drop.png b/docs/images/moe/token_drop.png
new file mode 100644
index 00000000000..1c335ee7aaf
Binary files /dev/null and b/docs/images/moe/token_drop.png differ
diff --git a/docs/images/multi_token_prediction/MTP_implementation.png b/docs/images/multi_token_prediction/MTP_implementation.png
new file mode 100644
index 00000000000..1f246c3e394
Binary files /dev/null and b/docs/images/multi_token_prediction/MTP_implementation.png differ
diff --git a/docs/index.md b/docs/index.md
new file mode 100644
index 00000000000..11337315588
--- /dev/null
+++ b/docs/index.md
@@ -0,0 +1,107 @@
+
+
+# Megatron Core User Guide
+
+**Megatron Core** is a GPU-optimized library for training large language models at scale. It provides modular, composable building blocks for creating custom training frameworks with state-of-the-art parallelism strategies and performance optimizations.
+
+Megatron 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.
+
+## Key Features
+
+* Composable transformer building blocks (attention, MLP)
+* Advanced parallelism strategies (TP, PP, DP, EP, CP)
+* Pipeline schedules and distributed optimizers
+* Mixed precision support (FP16, BF16, FP8)
+* GPU-optimized kernels and memory management
+* High-performance dataloaders and dataset utilities
+* Model architectures (LLaMA, Qwen, DeepSeek, GPT, Mamba)
+
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+:caption: About Megatron Core
+
+get-started/overview
+get-started/releasenotes
+```
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+:caption: Get Started
+
+get-started/install
+get-started/quickstart
+```
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+:caption: Basic Usage
+
+user-guide/data-preparation
+user-guide/training-examples
+user-guide/parallelism-guide
+```
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+:caption: Supported Models
+
+models/index
+```
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+:caption: Advanced Features
+
+user-guide/features/moe
+user-guide/features/context_parallel
+user-guide/features/megatron_fsdp
+user-guide/features/dist_optimizer
+user-guide/features/optimizer_cpu_offload
+user-guide/features/pipeline_parallel_layout
+user-guide/features/fine_grained_activation_offloading
+user-guide/data-loading
+user-guide/features/megatron_energon
+user-guide/features/megatron_rl
+user-guide/features/tokenizers
+```
+
+```{toctree}
+:maxdepth: 1
+:hidden:
+:caption: Developer Guide
+
+developer/contribute
+developer/submit
+developer/oncall
+developer/generate_docs
+```
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+:caption: API Reference
+
+api-guide/index
+apidocs/index.rst
+```
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+:caption: Resources
+
+advanced/index
+```
diff --git a/docs/llama_mistral.md b/docs/llama_mistral.md
new file mode 100644
index 00000000000..2754405610c
--- /dev/null
+++ b/docs/llama_mistral.md
@@ -0,0 +1,368 @@
+
+
+# Llama, Mistral and other Llama-like model support in Megatron-LM
+
+NOTE: In order to simplify code we now only support converting llama-3.x and mistral checkpoints downloaded from Hugging Face. For converting other models, see [Megatron Bridge](models/index.md).
+
+The Llama-2 and Llama-3.x family of models are an open-source set of pretrained & finetuned (for chat) models that have achieved strong results across a wide set of benchmarks. At their times of release, both Llama-2 and Llama-3 models achieved among the best results for open-source models, and were competitive with leading closed-source models (see ).
+
+Similarly, [Mistral-7b](https://mistral.ai/news/announcing-mistral-7b/) is an open-source model with pretrained and finetuned (for chat) variants that achieve strong benchmark results.
+
+Architecturally 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.
+
+# Contents
+
+- [Llama, Mistral and other Llama-like model support in Megatron-LM](#llama-mistral-and-other-llama-like-model-support-in-megatron-lm)
+- [Contents](#contents)
+- [Llama-2](#llama-2)
+ - [Download Huggingface checkpoints](#download-huggingface-checkpoints)
+ - [Convert checkpoint format](#convert-checkpoint-format)
+ - [Huggingface format](#huggingface-format)
+ - [Launch model](#launch-model)
+ - [Launch Megatron](#launch-megatron)
+ - [Launch Huggingface](#launch-huggingface)
+ - [Benchmark results](#benchmark-results)
+ - [Big Bench](#big-bench)
+ - [Multilingual](#multilingual)
+ - [LM Evaluation Harness](#lm-evaluation-harness)
+ - [MMLU](#mmlu)
+- [Llama-3.x](#llama-3x)
+ - [Download Huggingface checkpoints](#download-huggingface-checkpoints)
+ - [Convert checkpoint format](#convert-checkpoint-format)
+ - [Huggingface format](#huggingface-format)
+ - [(Optional) Validate checkpoints](#optional-validate-checkpoints)
+ - [Launch model](#launch-model)
+- [Mistral-7b](#mistral-7b)
+ - [Download Huggingface checkpoints](#download-huggingface-checkpoints)
+ - [Convert checkpoint format](#convert-checkpoint-format)
+ - [(Optional) Validate checkpoints](#optional-validate-checkpoints)
+ - [Launch model](#launch-model)
+- [Other Llama-like model support](#other-llama-like-model-support)
+- [Known numerical differences](#known-numerical-differences)
+
+# Llama-2
+
+Llama-2 checkpoints can be loaded into Megatron for inference and for finetuning. Loading these checkpoints consists of three steps:
+
+1. Get access to download the checkpoints.
+2. Convert the checkpoints from Huggingface format to Megatron format.
+3. Setup arguments for launching the model.
+
+The 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.
+
+## Download Huggingface checkpoints
+
+Users must first apply for access to download the Llama-2 checkpoints either directly [Huggingface](https://huggingface.co/docs/transformers/main/model_doc/llama2) (HF). The checkpoints are available in HF's format (available only from HF). HF format can be converted to Megatron, as detailed next.
+
+## Convert checkpoint format
+
+We recommend passing `--dtype bf16` for training or finetuning. Inference can be done in bfloat16 or float16.
+
+### Huggingface format
+
+The HF checkpoints can be converted to Megatron format by using Megatron-Bridge's checkpoint converter for HF format [see script](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py).
+
+```
+python Megatron-Bridge/examples/conversion/convert_checkpoints.py import \
+ --hf-model meta-llama/Llama-2-7B \
+ --megatron-path ./checkpoints/llama2_7b \
+ --torch-dtype bfloat16 \
+ --device-map auto
+```
+
+After this conversion, we are ready to load the checkpoints into a Megatron GPT model.
+
+## Launch model
+
+### Launch Megatron
+
+If loading for either inference or finetuning, use the following arguments:
+
+```
+--tensor-model-parallel-size ${TP} \
+--pipeline-model-parallel-size 1 \
+--seq-length 4096 \
+--max-position-embeddings 4096 \
+--tokenizer-type Llama2Tokenizer \
+--tokenizer-model ${TOKENIZER_MODEL} \
+--load ${CHECKPOINT_DIR} \
+--exit-on-missing-checkpoint \
+--use-checkpoint-args \
+--no-load-optim \
+--no-load-rng \
+--untie-embeddings-and-output-weights \
+--use-rotary-position-embeddings \
+--normalization RMSNorm \
+--no-position-embedding \
+--no-masked-softmax-fusion \
+--attention-softmax-in-fp32
+```
+
+### Launch Huggingface
+
+Huggingface checkpoints can be launched with:
+
+## Benchmark results
+
+The 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).
+
+The 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:
+
+- Megatron performs batch matrix multiplications in a couple places, such as within self attention and in SwiGLU, that Llama performs separately.
+- Megatron uses `torch.baddbmm` within self attention, versus Llama using `torch.matmul`.
+- Megatron uses a `sin`/`cos` implementation for rotary position embeddings, versus Llama using a `polar`/`complex` implementation.
+- Llama calls `torch.set_default_dtype(torch.float16)` during initialization, which Megatron does not.
+
+### Big Bench
+
+Score type: multiple choice grade.
+
+| bigbench / standard | 7b | 13b | 70b |
+| -- | -- | -- | -- |
+| date_understanding | 0.29% | 0.13% | 0.12% |
+| general_knowledge | 0.00% | 0.00% | 0.00% |
+| human_organs_senses | 0.00% | 0.00% | 0.00% |
+| intent_recognition | 0.00% | 0.11% | 0.00% |
+| riddle_sense | 0.00% | 0.00% | 0.00% |
+| similarities_abstraction | 0.00% | 0.58% | 0.00% |
+| simple_arithmetic_json_multiple_choice | 0.00% | 0.00% | 0.00% |
+| undo_permutation | 0.19% | 0.19% | 0.18% |
+
+### Multilingual
+
+Score type: multiple choice grade.
+
+| multilingual / xcopa | 7b | 13b | 70b |
+| -- | -- | -- | -- |
+| en-template-mGPT-remove-punctuation | 0.08% | 0.00% | 0.00% |
+| et-template-mGPT-remove-punctuation | 0.00% | 0.13% | 0.25% |
+| ht-template-mGPT-remove-punctuation | 0.26% | 0.13% | 0.26% |
+| id-template-mGPT-remove-punctuation | 0.11% | 0.00% | 0.19% |
+| it-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% |
+| qu-template-mGPT-remove-punctuation | 0.00% | 0.00% | 0.27% |
+| sw-template-mGPT-remove-punctuation | 0.14% | 0.13% | 0.13% |
+| th-template-mGPT-remove-punctuation | 0.25% | 0.13% | 0.13% |
+| tr-template-mGPT-remove-punctuation | 0.26% | 0.00% | 0.34% |
+| vi-template-mGPT-remove-punctuation | 0.00% | 0.11% | 0.00% |
+| zh-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% |
+
+### LM Evaluation Harness
+
+Score type: multiple choice grade.
+
+| lm-eval | 7b | 13b | 70b |
+| -- | -- | -- | -- |
+| boolq | 0.04% | 0.04% | 0.07% |
+| hellaswag | 0.02% | 0.03% | 0.03% |
+| piqa | 0.00% | 0.00% | 0.07% |
+| winogrande | 0.00% | 0.11% | 0.20% |
+
+### MMLU
+
+Score type: multiple choice grade.
+
+Note: the number in brackets is the number of sub-tasks for each supercategory.
+
+| mmlu | 7b | 13b | 70b |
+| -- | -- | -- | -- |
+| stem [18] | 0.79% | 0.05% | 0.01% |
+| humanities [13] | 0.19% | 0.01% | 0.02% |
+| other (business, health, misc.) [14] | 0.08% | 0.06% | 0.12% |
+| social sciences [12] | 0.37% | 0.21% | 0.01% |
+
+# Llama-3.x
+
+Llama-3.x checkpoints can be loaded into Megatron for inference and for finetuning. Loading these checkpoints consists of several steps:
+
+1. Get access to download the checkpoints (weights and tokenizer).
+2. Convert the checkpoints from Huggingface format to Megatron format.
+3. (Optional) Validate converted checkpoints
+4. Setup arguments for launching the model.
+
+The following sections detail these steps.
+
+## Download Huggingface checkpoints
+
+Users must first apply for access to download the Llama-3.x checkpoints from [Huggingface](https://huggingface.co/meta-llama).
+
+## Convert checkpoint format
+
+We recommend passing `--dtype bf16` for training or finetuning. Inference can be done in bfloat16 or float16.
+
+### Huggingface format
+
+The HF checkpoints can be converted to Megatron format by using Megatron-Bridge's checkpoint converter for HF format [see script](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py).
+
+```
+python Megatron-Bridge/examples/conversion/convert_checkpoints.py import \
+ --hf-model meta-llama/Llama-3.2-1B \
+ --megatron-path ./checkpoints/llama3_2_1b \
+ --torch-dtype bfloat16 \
+ --device-map auto
+```
+
+After this conversion, we are ready to load the checkpoints into a Megatron GPT model.
+
+## (Optional) Validate checkpoints
+
+A 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`.
+
+Once 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}'`.
+
+A reference generation for comparison can be obtained from the Huggingface transformers library by running `python examples/llama_mistral/huggingface_reference.py --model_path --prompt `.
+
+## Launch model
+
+If loading for either inference or finetuning, use the following arguments for Llama 3.0:
+
+```
+--tensor-model-parallel-size ${TP} \
+--pipeline-model-parallel-size 1 \
+--seq-length 8192 \
+--max-position-embeddings 8192 \
+--tokenizer-type HuggingFaceTokenizer \
+--tokenizer-model ${TOKENIZER_MODEL} \
+--load ${CHECKPOINT_DIR} \
+--exit-on-missing-checkpoint \
+--use-checkpoint-args \
+--no-load-optim \
+--no-load-rng \
+--untie-embeddings-and-output-weights \
+--normalization RMSNorm \
+--position-embedding-type rope \
+--no-masked-softmax-fusion \
+--attention-softmax-in-fp32 \
+--disable-bias-linear \
+--transformer-impl transformer_engine \
+--group-query-attention 8 \
+--attention-dropout 0.0 \
+--hidden-dropout 0.0 \
+--rotary-base 500000 \
+--rotary-percent 1.0 \
+--ffn-hidden-size 14336 \
+--num-attention-heads 32 \
+--swiglu \
+--bf16 \
+```
+
+For Llama3.1 please use the following arguments:
+
+```
+--tensor-model-parallel-size ${TP} \
+--pipeline-model-parallel-size 1 \
+--seq-length 8192 \
+--max-position-embeddings 131072 \
+--tokenizer-type HuggingFaceTokenizer \
+--tokenizer-model ${TOKENIZER_MODEL} \
+--load ${CHECKPOINT_DIR} \
+--exit-on-missing-checkpoint \
+--use-checkpoint-args \
+--no-load-optim \
+--no-load-rng \
+--untie-embeddings-and-output-weights \
+--normalization RMSNorm \
+--position-embedding-type rope \
+--no-masked-softmax-fusion \
+--attention-softmax-in-fp32 \
+--disable-bias-linear \
+--transformer-impl transformer_engine \
+--group-query-attention 8 \
+--attention-dropout 0.0 \
+--hidden-dropout 0.0 \
+--rotary-base 500000 \
+--rotary-percent 1.0 \
+--use-rope-scaling \
+--ffn-hidden-size 14336 \
+--num-attention-heads 32 \
+--swiglu \
+--bf16 \
+```
+
+# Mistral-7b
+
+Megatron 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:
+
+1. Get access to download the checkpoints (weights and tokenizer).
+2. Convert the checkpoints from HuggingFace format to Megatron format.
+3. (Optional) Validate converted checkpoints
+4. Setup arguments for launching the model.
+
+The following sections detail these steps.
+
+## Download Huggingface checkpoints
+
+Users must first apply for access to download the Mistral-7b checkpoints through Huggingface. Two variants are available: the base model ([Mistral-7B-v0.3](https://huggingface.co/mistralai/Mistral-7B-v0.3)) and the instruct model ([Mistral-7B-Instruct-v0.3](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3)).
+
+## Convert checkpoint format
+
+The HF checkpoints can be converted to Megatron format by using Megatron-Bridge's checkpoint converter for HF format [see script](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/conversion/convert_checkpoints.py).
+
+```
+python Megatron-Bridge/examples/conversion/convert_checkpoints.py import \
+ --hf-model mistralai/Mistral-7B-Instruct-v0.3 \
+ --megatron-path ./checkpoints/mistral_7b \
+ --torch-dtype bfloat16 \
+ --device-map auto
+```
+
+After this conversion, we are ready to load the checkpoints into a Megatron GPT model.
+
+## (Optional) Validate checkpoints
+
+A Megatron-LM text generation server for Mistral-7B can be launched using the script `examples/inference/llama_mistral/run_text_generation_mistral.sh `.
+
+Once 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}'`.
+
+A 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 `.
+
+## Launch model
+
+If loading for either inference or finetuning, use the following arguments:
+
+```
+--tensor-model-parallel-size ${TP} \
+--pipeline-model-parallel-size 1 \
+--seq-length 4096 \
+--max-position-embeddings 4096 \
+--tokenizer-type HuggingFaceTokenizer \
+--tokenizer-model ${TOKENIZER_MODEL} \
+--load ${CHECKPOINT_DIR} \
+--exit-on-missing-checkpoint \
+--use-checkpoint-args \
+--no-load-optim \
+--no-load-rng \
+--untie-embeddings-and-output-weights \
+--normalization RMSNorm \
+--position-embedding-type rope \
+--no-masked-softmax-fusion \
+--attention-softmax-in-fp32
+--apply-layernorm-1p \
+--transformer-impl transformer_engine \
+--group-query-attention 8 \
+--disable-bia-linear \
+--rotary-base 1000000 \
+--rotary-percent 1.0 \
+--swiglu \
+--ffn-hidden-size 14336 \
+--num-attention-heads 32
+```
+
+# Other Llama-like model support
+
+*Note: Experimental*
+
+Many 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).
+
+# Known numerical differences
+
+It is not expected that the megatron and Huggingface implementations of llama3.x and mistral models will produce numerically identical results. There are multiple points where small numerical differences are expected. This is a non-exhaustive list:
+
+1. TransformerEngine (TE) uses the model params_dtype inside RMSNorm whereas the Huggingface implementation uses fp32. See for details:
+2. Huggingface `transformers` implements the q, k and v projections in self-attention as separate GEMMs whereas Megatron core combines them into a single GEMM for efficiency. This leads to small numerical differences.
diff --git a/docs/models/index.md b/docs/models/index.md
new file mode 100644
index 00000000000..0ee379b01bd
--- /dev/null
+++ b/docs/models/index.md
@@ -0,0 +1,26 @@
+
+
+# Supported Models
+
+Megatron Core supports a wide range of language and multimodal models with optimized implementations for large-scale training.
+
+## Model Conversion
+
+For 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.
+
+See 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.
+
+```{toctree}
+:maxdepth: 1
+
+llms
+multimodal
+../llama_mistral
+```
diff --git a/docs/models/llms.md b/docs/models/llms.md
new file mode 100644
index 00000000000..f649673a2cc
--- /dev/null
+++ b/docs/models/llms.md
@@ -0,0 +1,59 @@
+
+
+# Language Models
+
+Megatron Core supports the following language model architectures for large-scale training.
+
+## Converting HuggingFace Models
+
+Use [**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.
+
+See 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.
+
+## Decoder-Only Models
+
+| Model | Description | Key Features |
+|-------|-------------|--------------|
+| **GPT** | Generative Pre-trained Transformer | Standard autoregressive LM, foundational architecture |
+| **LLaMA** | Meta's LLaMA family | Efficient architecture with RoPE, SwiGLU, RMSNorm |
+| **Mistral** | Mistral AI models | Sliding window attention, efficient inference |
+| **Mixtral** | Sparse Mixture-of-Experts | 8x7B MoE architecture for efficient scaling |
+| **Qwen** | Alibaba's Qwen series | HuggingFace integration, multilingual support |
+| **Mamba** | State Space Model | Subquadratic sequence length scaling, efficient long context |
+
+## Encoder-Only Models
+
+| Model | Description | Key Features |
+|-------|-------------|--------------|
+| **BERT** | Bidirectional Encoder Representations | Masked language modeling, classification tasks |
+
+## Encoder-Decoder Models
+
+| Model | Description | Key Features |
+|-------|-------------|--------------|
+| **T5** | Text-to-Text Transfer Transformer | Unified text-to-text framework, sequence-to-sequence |
+
+## Example Scripts
+
+Training examples for these models can be found in the `examples/` directory:
+- `examples/gpt3/` - GPT-3 training scripts
+- `examples/llama/` - LLaMA training scripts
+- `examples/mixtral/` - Mixtral MoE training
+- `examples/mamba/` - Mamba training scripts
+- `examples/bert/` - BERT training scripts
+- `examples/t5/` - T5 training scripts
+
+## Model Implementation
+
+All language models are built using Megatron Core's composable transformer blocks, enabling:
+- Flexible parallelism strategies (TP, PP, DP, EP, CP)
+- Mixed precision training (FP16, BF16, FP8)
+- Distributed checkpointing
+- Efficient memory management
diff --git a/docs/models/multimodal.md b/docs/models/multimodal.md
new file mode 100644
index 00000000000..07ff76d8d9a
--- /dev/null
+++ b/docs/models/multimodal.md
@@ -0,0 +1,71 @@
+
+
+# Multimodal Models
+
+Megatron Core supports multimodal models that combine language with vision, audio, and other modalities for comprehensive multimodal understanding.
+
+## MIMO: Multimodal In/Out Framework
+
+**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.
+
+> **Note**: MIMO is experimental and under active development. The API may change in future releases.
+
+**Key Features:**
+- Arbitrary modality combinations (vision, audio, text)
+- Flexible encoder architecture for different input modalities
+- Unified embedding space across modalities
+- Support for both vision-language and audio-vision-language models
+
+See [examples/mimo](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/mimo) for training scripts and examples.
+
+## Vision-Language Models
+
+| Model | Description | Vision Encoder | Language Model |
+|-------|-------------|----------------|----------------|
+| **LLaVA** | Visual instruction tuning | CLIP ViT-L/14 | Mistral-7B / LLaMA |
+| **NVLM** | NVIDIA Vision-Language Model | CLIP / Custom ViT | LLaMA-based |
+| **LLaMA 3.1 Nemotron Nano VL** | Efficient multimodal model | Vision Transformer | LLaMA 3.1 8B |
+
+## Vision Encoders
+
+| Model | Description | Key Features |
+|-------|-------------|--------------|
+| **CLIP ViT** | OpenAI's CLIP Vision Transformer | Image-text alignment, multiple scales (L/14@336px) |
+| **RADIO** | Resolution-Agnostic Dynamic Image Optimization | Flexible resolution handling, efficient vision encoding |
+
+## Diffusion Models
+
+For multimodal diffusion models (image generation, text-to-image). Refer to [Nvidia Diffusion Models](https://github.com/NVIDIA-NeMo/DFM/ ). The Developer Program, NIM, and NeMo can offer production-ready implementations of:
+
+- Stable Diffusion variants
+- Text-to-image generation
+- Image-to-image translation
+- ControlNet and other conditioning mechanisms
+
+## Multimodal Features
+
+- **Image-Text Alignment**: Pre-training on image-caption pairs
+- **Visual Instruction Tuning**: Fine-tuning on instruction-following datasets
+- **Flexible Vision Encoders**: Support for different ViT architectures and resolutions
+- **Combined Checkpointing**: Unified checkpoints combining vision and language models
+- **Efficient Training**: Full parallelism support (TP, PP, DP) for both vision and language components
+
+## Example Scripts
+
+Multimodal training examples can be found in the following directories:
+
+**MIMO Framework:**
+- `examples/mimo/` - Multimodal In/Out training with support for vision-language and audio-vision-language models
+
+**Specific Multimodal Models:**
+- `examples/multimodal/` - LLaVA-style training with Mistral + CLIP
+- `examples/multimodal/nvlm/` - NVLM training scripts
+- `examples/multimodal/llama_3p1_nemotron_nano_vl_8b_v1/` - Nemotron VL training
+- `examples/multimodal/radio/` - RADIO vision encoder integration
diff --git a/docs/project.json b/docs/project.json
new file mode 100644
index 00000000000..d5b9535338b
--- /dev/null
+++ b/docs/project.json
@@ -0,0 +1,2 @@
+{"name": "megatron-lm", "version": "nightly"}
+
diff --git a/docs/user-guide/data-loading.md b/docs/user-guide/data-loading.md
new file mode 100644
index 00000000000..b60cd685cf2
--- /dev/null
+++ b/docs/user-guide/data-loading.md
@@ -0,0 +1,152 @@
+
+
+# Data Loading at Scale
+
+This guide covers how Megatron's data pipeline works and how to configure it for efficient training at 256 nodes and beyond. At this scale, the primary bottlenecks are **index building** and **barrier synchronization** -- not raw data bandwidth.
+
+## How Data Loading Works
+
+Understanding the architecture helps explain why specific flags matter.
+
+Megatron builds three index arrays for each dataset: a **document index** (shuffled document order), a **sample index** (mapping samples to document offsets), and a **shuffle index** (final sample permutation). This happens once during initialization:
+
+1. **Rank 0** builds all three indices and writes them to a cache directory as `.npy` files.
+2. All ranks synchronize at a `torch.distributed.barrier()`.
+3. **All other ranks** load the cached indices via memory-mapped reads (`numpy.load(mmap_mode='r')`).
+
+After initialization, data access is **read-only and lock-free**. Each data-parallel rank consumes a disjoint subset of samples, and no cross-rank coordination is needed during training because all ranks derive the same deterministic permutation from a shared random seed.
+
+## The Problem at 256+ Nodes
+
+Three things break down at large node counts:
+
+1. **Barrier synchronization**: All ranks block while rank 0 builds indices. On a 512-node job, this means 4,095 GPUs sit idle.
+2. **Simultaneous memory-mapping**: All ranks `mmap` three large `.npy` files at once after the barrier, causing a burst of page faults and I/O.
+
+## Baseline: Establish Maximum Achievable Performance
+
+Before tuning data loading, establish a performance ceiling by running with `--mock-data`. This bypasses the data pipeline entirely and shows the maximum throughput your configuration can achieve without any dataloader overhead. The gap between `--mock-data` performance and real-data performance tells you exactly how much time the dataloader is costing you.
+
+## Recommended Configuration
+
+### Step 1: Consolidate dataset files
+
+A common issue at scale is having datasets split across many small file prefixes. Thousands of 100 MB files perform significantly worse than tens of 10 GB+ files, both for building dataset caches and for runtime file access.
+
+Use the merge tool to consolidate datasets stored as many small prefixes in one directory:
+
+```bash
+python tools/merge_datasets.py \
+ --input /path/to/input-directory \
+ --output-prefix /path/to/output/merged
+```
+
+**Target at least 10 GB per file.** This reduces the number of file descriptors, metadata lookups, and index-building work at initialization.
+
+### Step 2: Pre-build the dataset cache
+
+Build the GPT dataset cache as a separate step before training. This avoids the usual "rank 0 builds, everyone else waits" startup path and is the recommended workflow for large jobs:
+
+```bash
+python tools/prepare_cache.py \
+ --data-path \
+ --split 99,1,0 \
+ --data-cache-path /path/to/cache \
+ --global-batch-size \
+ --seq-length \
+ ...
+```
+
+If your later training job does not set `--global-batch-size`, or you are preparing the cache on a machine that does not match the future training topology, also pass:
+
+```bash
+--prepare-cache-world-size
+```
+
+This keeps the prepared cache aligned with the sample counts expected by training.
+
+> **Unsupported configurations:** `tools/prepare_cache.py` does not support `--mock-data`, `--sft`, `--fim-data`, or `--step-batch-size-schedule`. Using any of these will cause the script to exit with an error.
+
+### Step 3: Optionally pre-build per-dataset metadata
+
+When blending many datasets, generate the `--per-dataset-sequences-path` JSON ahead of time to avoid one metadata read per file prefix at startup:
+
+```bash
+python tools/build_sequences_per_dataset.py \
+ --data-path \
+ --per-dataset-sequences-path sequences.json
+```
+
+### Step 4: Launch training with optimized data loading
+
+Once the cache is ready, enable the fast-path flags:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=512 ... pretrain_gpt.py \
+ --dataloader-fast-cache-load \
+ --dataloader-defer-npy-index-mmap \
+ --per-dataset-sequences-path sequences.json \
+ --data-cache-path /path/to/cache \
+ --num-workers 2 \
+ ...
+```
+
+### Flag reference
+
+| Flag | Default | Recommendation | What it does |
+|------|---------|----------------|-------------|
+| `--dataloader-fast-cache-load` | off | **On** | Skips the rank-0 barrier by assuming the cache already exists. All ranks build their dataset views in parallel. This is the single biggest win at scale. |
+| `--dataloader-defer-npy-index-mmap` | off | **On** | Defers memory-mapping of `.npy` index files until first access. When combined with `--num-workers > 0`, index loading is overlapped with the training iteration rather than blocking startup. |
+| `--per-dataset-sequences-path` | None | **Set when blending many datasets** | Points to a JSON file mapping each dataset path to its `(sequence_count, document_count)`. Replaces per-file metadata reads with a single JSON lookup. Generate with `tools/build_sequences_per_dataset.py`. |
+| `--data-cache-path` | None | **Set** | Directory where index `.npy` files are cached. Must be on shared storage for multi-node jobs so all ranks can read it. |
+| `--num-workers` | 2 | **Keep as small as necessary** | Number of DataLoader worker processes. The goal is to satisfy: *time to process a batch > time to prepare a batch*. This hides dataloader work behind the training step. Increasing beyond what's needed wastes CPU and memory. |
+| `--no-mmap-bin-files` | mmap on | **Test both** | Memory-mapping `.bin` files leverages the OS page cache, but the optimal setting is filesystem-dependent. Some large-scale production configurations disable mmap. Test with and without to determine what works best for your storage. |
+
+### Object storage (S3 / Multi-Storage Client)
+
+When data lives on S3 or MSC rather than a POSIX filesystem:
+
+- **Index files** (`.idx`) are cached locally under `object_storage_cache_path`.
+- **Binary data files** (`.bin`) are streamed on-demand in 256 MB chunks, avoiding the need to download entire files.
+- Set `--no-mmap-bin-files` since memory-mapping doesn't apply to object storage.
+- Ensure the index-cache path is visible wherever the later dataset construction will run.
+
+## Scaling Characteristics
+
+| Aspect | Behavior | Why it works |
+|--------|----------|-------------|
+| **Cross-rank contention** | None after init | All index files are read-only; `numpy.memmap` uses OS page cache with no locking |
+| **Sampling determinism** | All ranks produce the same permutation | Shared `numpy.random.RandomState(seed)` with epoch-based seed variation |
+| **Data-parallel sharding** | Each DP rank gets a disjoint subset of samples | No overlap during training; assignment happens in the sampler rather than via extra dataset coordination |
+| **Index broadcast** | Via shared filesystem, not collectives | Rank 0 writes `.npy` files; other ranks read them. No explicit `torch.distributed.broadcast` |
+
+## Troubleshooting
+
+**Symptom: Training hangs at startup for minutes**
+- Likely cause: Rank 0 is building indices while all other ranks wait at the barrier.
+- Fix: Pre-build the cache with `tools/prepare_cache.py` and enable `--dataloader-fast-cache-load`.
+
+**Symptom: Spike in I/O at training start, then normal**
+- Likely cause: All ranks simultaneously memory-mapping index files after the barrier.
+- Fix: Enable `--dataloader-defer-npy-index-mmap` to overlap index loading with training.
+
+**Symptom: Slow data loading during training (not just startup)**
+- Run with `--mock-data` to confirm the dataloader is the bottleneck.
+- If startup, not steady-state throughput, is the main issue, try `--dataloader-defer-npy-index-mmap`.
+- If you are blending many dataset prefixes, try `--per-dataset-sequences-path`.
+- Test with `--no-mmap-bin-files` -- the optimal setting depends on your filesystem.
+
+## Related Resources
+
+- [PR #2445](https://github.com/NVIDIA/Megatron-LM/pull/2445): Original implementation of fast cache load, deferred mmap, and per-dataset sequences optimizations.
+- [PR #4080](https://github.com/NVIDIA/Megatron-LM/pull/4080): Adds `tools/prepare_cache.py` for offline GPT dataset cache preparation.
+- [`tools/prepare_cache.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/tools/prepare_cache.py): Pre-build GPT dataset caches ahead of training.
+- [`tools/merge_datasets.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/tools/merge_datasets.py): Merge multiple small dataset files into larger ones.
+- [`tools/build_sequences_per_dataset.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/tools/build_sequences_per_dataset.py): Generate the `--per-dataset-sequences-path` JSON file.
diff --git a/docs/user-guide/data-preparation.md b/docs/user-guide/data-preparation.md
new file mode 100644
index 00000000000..813f81501e7
--- /dev/null
+++ b/docs/user-guide/data-preparation.md
@@ -0,0 +1,126 @@
+
+
+# Data Preparation
+
+Preparing your data correctly is essential for successful training with Megatron Core.
+
+## Data Format
+
+Megatron Core expects training data in JSONL (JSON Lines) format, where each line is a JSON object:
+
+```json
+{"text": "Your training text here..."}
+{"text": "Another training sample..."}
+{"text": "More training data..."}
+```
+
+## Preprocessing Data
+
+Use the `preprocess_data.py` tool to convert your JSONL data into Megatron's binary format:
+
+```bash
+python tools/preprocess_data.py \
+ --input data.jsonl \
+ --output-prefix processed_data \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model /path/to/tokenizer.model \
+ --workers 8 \
+ --append-eod
+```
+
+### Key Arguments
+
+The following table summarizes the main preprocessor arguments:
+
+| Argument | Description |
+|----------|-------------|
+| `--input` | Path to input JSON/JSONL file |
+| `--output-prefix` | Prefix for output binary files (.bin and .idx) |
+| `--tokenizer-type` | Tokenizer type (`HuggingFaceTokenizer`, `GPT2BPETokenizer`, and so on) |
+| `--tokenizer-model` | Path to tokenizer model file |
+| `--workers` | Number of parallel workers for processing |
+| `--append-eod` | Add end-of-document token |
+
+## Finding Optimal Number of Workers
+
+Use the `--find-optimal-num-workers` flag to find the number of workers that gives the best performance in terms of preprocessed documents per second.
+The script launches a few short data preprocessing runs with different worker counts and identifies the fastest run using the collected performance data.
+
+```bash
+python tools/preprocess_data.py \
+ --input data.jsonl \
+ --output-prefix processed_data \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model /path/to/tokenizer.model \
+ --workers 8 \
+ --find-optimal-num-workers \
+ --workers-to-check 4 8 16 32 \
+ --max-documents 50000
+```
+
+**Required arguments**
+
+The following table lists the arguments required for worker optimization:
+
+| Argument | Description |
+|----------|-------------|
+| `--find-optimal-num-workers` | Activates search of optimal number of workers |
+| `--workers-to-check` | List of possible number of workers to run |
+| `--max-documents` | Number of documents to be preprocessed during each run |
+
+**Output example**
+
+The command prints performance results similar to the following:
+
+```bash
+-----------------------------------
+Performance results (fastest → slowest):
+1. 16 workers → avg. docs/s: 9606.6476
+2. 32 workers → avg. docs/s: 9275.3284
+3. 8 workers → avg. docs/s: 9151.9280
+4. 4 workers → avg. docs/s: 6391.3819
+
+-----------------------------------
+The most optimal num of workers is 16 with avg. preprocessed docs/s: 9606.6476.
+-----------------------------------
+```
+
+## Output Files
+
+The preprocessing tool generates two files:
+
+- `processed_data.bin` - Binary file containing tokenized sequences
+- `processed_data.idx` - Index file for fast random access
+
+## Using Preprocessed Data
+
+Reference your preprocessed data in training scripts:
+
+```bash
+--data-path processed_data \
+--split 949,50,1 # Train/validation/test split
+```
+
+## Common Tokenizers
+
+### HuggingFace Tokenizers
+
+```bash
+--tokenizer-type HuggingFaceTokenizer \
+--tokenizer-model /path/to/tokenizer.model
+```
+
+### GPT-2 BPE Tokenizer
+
+```bash
+--tokenizer-type GPT2BPETokenizer \
+--vocab-file gpt2-vocab.json \
+--merge-file gpt2-merges.txt
+```
diff --git a/docs/user-guide/features/context_parallel.md b/docs/user-guide/features/context_parallel.md
new file mode 100644
index 00000000000..890609ac7de
--- /dev/null
+++ b/docs/user-guide/features/context_parallel.md
@@ -0,0 +1,43 @@
+
+
+# Context Parallel Package
+
+## Context Parallelism Overview
+
+```{figure} ../../images/context_parallel/CP_overview.png
+:alt: Diagram of a transformer layer with tensor parallelism 2 and context parallelism 2, showing CP and TP communication patterns around attention and other blocks.
+:align: center
+
+Figure 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).
+```
+
+Context Parallelism (CP) is a parallelization scheme on the sequence-length dimension. Unlike prior SP (sequence parallelism), which only splits the sequence of Dropout and LayerNorm activations, CP partitions the network inputs and all activations along the sequence dimension. With CP, all modules except attention (for example, Linear and LayerNorm) can work as usual without any changes, because they do not have inter-token operations. For attention, the Q (query) of each token must combine with the KV (key and value) of all tokens in the same sequence. CP therefore requires an additional all-gather across GPUs to collect the full sequence of KV. Correspondingly, reduce-scatter is applied to the activation gradients of KV in backward propagation. To reduce activation memory footprint, each GPU stores only 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 implemented as point-to-point communications in a ring topology. Exchanging KV can also leverage MQA or GQA to reduce communication volume, because those variants use one or a few attention heads for KV.
+
+For example, in Figure 1, if the sequence length is 8K, each GPU processes 4K tokens. GPU0 and GPU2 form a CP group and exchange KV with each other; the same pattern applies between GPU1 and GPU3. CP is similar to [Ring Attention](https://arxiv.org/abs/2310.01889) but targets higher performance by (1) using current open-source and cuDNN flash attention kernels, and (2) avoiding extra work from lower-triangle causal masking while keeping load balanced across GPUs.
+
+## Context Parallelism Benefits
+
+```{figure} ../../images/context_parallel/CP_results.png
+:alt: Chart of speedup for 175B GPT with different tensor parallelism and context parallelism combinations compared with full activation recomputation.
+:align: center
+
+Figure 2: Speedup of 175B GPT with various TP+CP combinations compared to full recomputation (that is, TP8CP1).
+```
+
+An LLM can hit an out-of-memory (OOM) error on long contexts (long sequence lengths) because activation memory grows about linearly with sequence length. Recomputing activations in backward can avoid OOM but adds significant overhead (about 30 percent with full recomputation). Increasing TP (tensor model parallelism) can also fix OOM, but it can make compute in layers such as Linear too short to hide communication latency. Scaling to more GPUs with larger TP can hit that overlap limit even when OOM is not the driver.
+
+CP addresses these tradeoffs. With CP, each GPU computes on part of the sequence, which scales down both compute and communication by the CP degree. Overlap between them is less of a concern. The activation memory footprint per GPU is also smaller by the CP degree, which reduces OOM risk. As Figure 2 shows, TP and CP together can outperform full recomputation by removing most recompute overhead and balancing compute against communication.
+
+## Enabling Context Parallelism
+
+CP support is included on the GPT code path. Other models that share that path, such as LLaMA, can use CP as well. CP works with TP (tensor model parallelism), PP (pipeline model parallelism), and DP (data parallelism). The total GPU count is TP × CP × PP × DP. CP also works with different attention variants, including MHA, MQA, and GQA, with unidirectional or bidirectional masking.
+
+Enable CP by setting `context_parallel_size=` on the command line. The default `context_parallel_size` is 1, which disables CP. Running with CP requires Megatron Core (>=0.5.0) and Transformer Engine (>=1.1).
+
diff --git a/docs/user-guide/features/dist_optimizer.md b/docs/user-guide/features/dist_optimizer.md
new file mode 100644
index 00000000000..bfea3b63a66
--- /dev/null
+++ b/docs/user-guide/features/dist_optimizer.md
@@ -0,0 +1,49 @@
+
+
+# Distributed Optimizer
+
+The distributed optimizer saves memory by sharding optimizer state across data parallel ranks instead of replicating it on every rank, as described in the [ZeRO paper](https://arxiv.org/abs/1910.02054).
+
+Theoretical memory savings depend on the data types of the model parameters (`param_dtype`) and of the main gradients accumulated across data-parallel replicas (`grad_dtype`). Optimizer steps always use `fp32` main parameters. In the current implementation, the theoretical number of bytes per parameter is as follows (where *d* is the data parallel size):
+
+| | Non-distributed optim | Distributed optim |
+| ------ | ------ | ------ |
+| `fp16` parameters, `fp16` gradients | 20 | 4 + 16/d |
+| `bf16` parameters, `fp32` gradients | 18 | 6 + 12/d |
+| `fp32` parameters, `fp32` gradients | 16 | 8 + 8/d |
+
+This distributed optimizer uses contiguous buffers for parameters and main gradients. Model gradients copy into the main gradients as soon as they finish computing.
+
+The following figures show the sharding scheme and the main steps of the parameter update.
+
+## Data Flow
+
+
+
+## Sharding Scheme
+
+
+
+## Key Steps
+
+**Note:** The following steps match the illustrations above. They assume `bf16` model weights, `bf16` model gradients from the backward pass, and `fp32` main gradients for optimizer steps. Optimizer steps use `fp32` main weights.
+
+- Backward pass finishes (gradient buffer holds 16 `fp32` gradient elements).
+- Call reduce-scatter on each DP rank.
+- Each DP rank now has four elements within the gradient buffer that are fully reduced (remaining 12 elements are garbage).
+ - DP rank 0 has gradient values for elements [0:4].
+ - DP rank 1 has gradient values for elements [4:8].
+ - DP rank 2 has gradient values for elements [8:12].
+ - DP rank 3 has gradient values for elements [12:16].
+- Optimizer.step().
+- Each DP rank copies its four `fp32` main parameter elements into the corresponding `bf16` parameter buffer (each element is cast from `fp32` to `bf16`).
+- Call all-gather on each DP rank.
+- The parameter buffer now contains all 16 updated `bf16` model parameter elements. Parameters in PyTorch modules already point to the correct views in this buffer, so forward passes can start after the all-gather completes.
+- At this point, you can zero the gradient buffer for the next iteration.
diff --git a/docs/user-guide/features/fine_grained_activation_offloading.md b/docs/user-guide/features/fine_grained_activation_offloading.md
new file mode 100644
index 00000000000..7b97c17cef9
--- /dev/null
+++ b/docs/user-guide/features/fine_grained_activation_offloading.md
@@ -0,0 +1,46 @@
+
+
+# Fine-Grained Activation Offloading
+
+Contributed in collaboration with RedNote.
+
+Memory is often the limiting factor for very large sparse MoE models such as DeepSeek-V3 and Qwen3-235B. Fine-grained recomputation lowers activation memory at the cost of extra compute. Offloading can use host-device bandwidth so that reload overlaps compute and keeps overhead small in many setups. Fine-grained activation offloading moves activations at module granularity so you can tune how much activation memory leaves the device and adjust training throughput.
+
+Supported offloading modules are `"attn_norm"`, `"core_attn"`, `"attn_proj"`, `"mlp_norm"`, `"expert_fc1"`, and `"moe_act"`. They can be combined with fine-grained recomputation to free almost all activations for a transformer layer on the device.
+
+## Features
+
+- Pipeline parallelism: PP=1, PP, and interleaved PP
+- Compatible with fine-grained recomputation
+- FP8 training
+- MTP
+- Mixed dense and MoE layers
+- A2A overlap
+- CUDA graphs
+ - **Note:** A CUDA graph capture cannot include the offloading modules (temporary limitation).
+
+## Usage
+
+```bash
+# Enable fine-grained activation offloading
+--fine-grained-activation-offloading
+
+# Modules whose inputs are offloaded (refer to your training script for list or delimiter syntax).
+# Choices: "attn_norm", "core_attn", "attn_proj", "mlp_norm", "expert_fc1", "moe_act".
+--offload-modules expert_fc1
+```
+
+## Compatible With Fine-Grained Recomputation
+
+- For low-overhead modules such as LayerNorm or `moe_act`, use recomputation to save activation memory.
+- For other modules, use offloading to save activation memory.
+- Overlap offload and reload with compute when possible.
+
+
diff --git a/docs/user-guide/features/index.md b/docs/user-guide/features/index.md
new file mode 100644
index 00000000000..9dea6bd34a4
--- /dev/null
+++ b/docs/user-guide/features/index.md
@@ -0,0 +1,27 @@
+
+
+# Advanced Features
+
+Guides for Megatron Core training features.
+
+```{toctree}
+:maxdepth: 2
+
+fine_grained_activation_offloading
+moe
+context_parallel
+megatron_fsdp
+dist_optimizer
+optimizer_cpu_offload
+pipeline_parallel_layout
+tokenizers
+megatron_energon
+megatron_rl
+```
diff --git a/docs/user-guide/features/megatron_energon.md b/docs/user-guide/features/megatron_energon.md
new file mode 100644
index 00000000000..c32b2d8facd
--- /dev/null
+++ b/docs/user-guide/features/megatron_energon.md
@@ -0,0 +1,143 @@
+
+
+# Megatron Energon
+
+Multimodal dataloader for text, images, video, and audio at scale.
+
+## Overview
+
+[**Megatron Energon**](https://github.com/NVIDIA/Megatron-Energon) supports large-scale multimodal training with:
+
+- **Multimodal support** - Text, images, video, and audio
+- **Distributed loading** - Suited to multi-node training
+- **Data blending** - Mix datasets with configurable weights
+- **WebDataset format** - Streaming from cloud storage
+- **State management** - Save and restore training position
+
+## Installation
+
+```bash
+pip install megatron-energon
+```
+
+## Key Features
+
+### Data Processing
+
+- **Packing** - Packs samples to use sequence length capacity
+- **Grouping** - Batching of similar-length sequences
+- **Joining** - Combine multiple dataset sources
+- **Object storage** - Stream from S3, GCS, and Azure Blob Storage
+
+### Production Use
+
+- Distributed loading across workers and nodes
+- Checkpoint data loading state
+- Memory-efficient streaming
+- Parallel data loading with prefetching
+
+## Basic Usage
+
+```python
+from megatron.energon import get_train_dataset, get_loader, WorkerConfig
+
+# Create dataset
+ds = get_train_dataset(
+ '/path/to/dataset',
+ batch_size=32,
+ shuffle_buffer_size=1000,
+ worker_config=WorkerConfig.default_worker_config(),
+)
+
+# Create loader and iterate
+for batch in get_loader(ds):
+ # Training step
+ pass
+```
+
+## Multimodal Example
+
+```python
+# Load image-text dataset
+ds = get_train_dataset(
+ '/path/to/multimodal/dataset',
+ batch_size=32,
+ worker_config=WorkerConfig(num_workers=8, prefetch_factor=2),
+)
+
+for batch in get_loader(ds):
+ images = batch['image'] # Image tensors
+ texts = batch['text'] # Text captions
+ # Process batch
+```
+
+## Dataset Blending
+
+Mix multiple datasets with custom weights:
+
+```python
+from megatron.energon import Blender
+
+blended_ds = Blender([
+ ('/path/to/dataset1', 0.6), # 60%
+ ('/path/to/dataset2', 0.3), # 30%
+ ('/path/to/dataset3', 0.1), # 10%
+])
+```
+
+## Configuration
+
+### Worker Configuration
+
+```python
+WorkerConfig(
+ num_workers=8, # Parallel workers
+ prefetch_factor=2, # Batches to prefetch per worker
+ persistent_workers=True, # Keep workers alive between epochs
+)
+```
+
+### Common Parameters
+
+The following table summarizes frequently used dataset and loader parameters:
+
+| Parameter | Description |
+|-----------|-------------|
+| `batch_size` | Samples per batch |
+| `shuffle_buffer_size` | Buffer size for randomization |
+| `max_samples_per_sequence` | Max samples to pack into one sequence |
+| `worker_config` | Worker configuration for parallel loading |
+
+## Integration with Megatron-LM
+
+```python
+from megatron.energon import get_train_dataset, get_loader
+from megatron.training import get_args
+
+args = get_args()
+
+train_ds = get_train_dataset(
+ args.data_path,
+ batch_size=args.micro_batch_size,
+)
+
+for iteration, batch in enumerate(get_loader(train_ds)):
+ loss = train_step(batch)
+```
+
+## Resources
+
+- **[Megatron Energon GitHub](https://github.com/NVIDIA/Megatron-Energon)**: Documentation and examples
+- **[Multimodal Examples](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/multimodal)**: Megatron-LM multimodal training
+
+## Next Steps
+
+- Refer to [Multimodal Models](../../models/multimodal.md) for supported architectures
+- Refer to [Training Examples](../training-examples.md) for integration examples
diff --git a/docs/user-guide/features/megatron_fsdp.md b/docs/user-guide/features/megatron_fsdp.md
new file mode 100644
index 00000000000..36fcc68893c
--- /dev/null
+++ b/docs/user-guide/features/megatron_fsdp.md
@@ -0,0 +1,608 @@
+
+
+# Megatron-FSDP
+
+## ✨ Overview
+
+**Megatron-FSDP** is an NVIDIA-developed distributed parallelism library written in native PyTorch that provides a high-performance implementation of **Fully Sharded Data Parallelism (FSDP)**. It offers seamless cross-compatibility with various deep learning frameworks and parallelism libraries such as Megatron-Core, and is performance-optimized to support training and inference of extremely large PyTorch models at data-center scale on NVIDIA GPUs.
+
+- PyPI: https://pypi.org/project/megatron-fsdp/
+- Source Code: https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/distributed/fsdp/src
+
+### 🧩 Compatibility
+
+- PyTorch **[DeviceMesh](https://docs.pytorch.org/docs/2.11/distributed.html#torch.distributed.device_mesh.DeviceMesh)**, **[DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)**, and **[Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html)**
+- **[Megatron Core](https://github.com/NVIDIA/Megatron-LM)**
+- **[TransformerEngine](https://github.com/NVIDIA/TransformerEngine)**
+- **[NVIDIA NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo)**
+
+### 💡 Features
+
+- **Performant & Scalable**: Optimized for NVIDIA CUDA with efficient memory management and performance. Sports near-linear scaling up from single compute nodes to entire data-centers.
+- **Multiple Algorithms in One**: Supports sharding your choice of optimizer states, gradients, and model parameters (FSDP), including hierarchical data parallelism strategies such as **Hybrid-Sharded Data Parallelism (HSDP)** and **Hybrid-FSDP (HFSDP / Fully-Sharded Optimizer State)** for optimizing intra-node and inter-node memory, communication, and performance.
+- **"Bring Your Own Parallelism"**: Works seamlessly with PyTorch, Megatron-LM, Megatron-Bridge, and TransformerEngine, and can be plugged into other frameworks such as HuggingFace Transformers and TorchTitan.
+- **Simple & Powerful**: Similar to PyTorch FSDP, the `fully_shard` API doesn't depend on any complex training framework or distributed environment.
+
+### ⏱️ Optimizations
+
+- **[TransformerEngine](https://github.com/NVIDIA/TransformerEngine) Mixed-Precision & Fused Kernels**: Native performance- and memory-optimal _compatibility with MXFP8, NVFP4, and various other quantization recipes and fused kernels provided by TransformerEngine_.
+- **Advanced Bucketing**: `dtype`-customizable and precision-aware bucketing system to _tune the memory overhead, numerical accuracy, and latency of collectives_. Avoids redundant `COPY` operations before and after collectives, while remaining compatible with **[DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)** features such as **[Torch Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html)**.
+- **Buffer Management**: Efficient use of storage and [NCCL User Buffer Registration](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html#user-buffer-registration) enable _direct communication into NCCL-managed memory_, achieving true zero-`COPY` data movement. Introduced in NCCL `v2.27`, **NCCL Symmetric Memory** communications employ _symmetric kernels_ that drastically reduce SM utilization and include networking optimizations such as high-precision (`FP32`) reduction over-the-wire.
+- **Optimized Communication & SM Utilization via SHARP**: Leverages [**SHARP** (Scalable Hierarchical Aggregation and Reduction Protocol)](https://docs.nvidia.com/networking/display/sharpv3130) to _offload FSDP collectives to network switches (InfiniBand or NVLink-Switch)_ and significantly reduce utilization of GPU streaming multi-processors (SM) from 16-32 to 1-6 for **Multi-Node NVLink (MNNVL)** systems (Grace-Blackwell, Vera-Rubin, etc.), which lowers communication latency in large scaled-out workloads and frees up GPU-hosted processors for overlapped compute (GEMM) kernels. When FSDP sharding domains span both NVLink and InfiniBand, **hierarchical SHARP collectives** (NVL-SHARP and IB-SHARP) _optimize communication paths across the entire system topology_.
+- [**Hybrid-FSDP (HFSDP)**](#understanding-hybrid-fsdp-hfsdp), a variation of _Hybrid-Sharded Data Parallelism (HSDP)_ that further shards the optimizer state across intra- and inter-node data-parallel ranks, _bridges the memory-communication trade-off between HSDP and FSDP_, unlocking memory efficiency at minimal cost to performance.
+
+## 🚀 Quick Start
+
+### 📦 Installation
+
+#### NeMo Framework Container
+
+Megatron-FSDP is pre-installed with Megatron-Core in the [NVIDIA NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags).
+
+#### Megatron-Core
+
+Megatron-FSDP is bundled with Megatron-Core, which can be installed via `pip`:
+
+```
+# Install via PyPI
+pip install --no-build-isolation megatron-core[mlm,dev]
+
+# Install from Source
+git clone https://github.com/NVIDIA/Megatron-LM.git
+cd Megatron-LM
+pip install --no-build-isolation .[mlm,dev]
+```
+
+To import Megatron-FSDP in Python:
+```python
+import megatron.core.distributed.fsdp.src.megatron_fsdp
+```
+
+#### PyPI
+
+To install Megatron-FSDP as a standalone package to use the `fully_shard` API:
+
+```
+pip install megatron-fsdp
+```
+
+To import Megatron-FSDP in Python:
+
+```python
+import megatron_fsdp
+```
+
+### 🎛️ Megatron-FSDP `fully_shard`
+
+Megatron-FSDP supports a simple `fully_shard` API that seamlessly enables FSDP with very few lines of code.
+
+```python
+import torch
+from megatron_fsdp import (
+ fully_shard_model,
+ fully_shard_optimizer,
+)
+
+# Initialize Torch Distributed.
+torch.distributed.init_process_group()
+torch.cuda.set_device(torch.distributed.get_rank())
+
+# Fully-shard the model.
+model = torch.nn.Transformer()
+fsdp_model = fully_shard_model(
+ module=model,
+ fsdp_unit_modules=[
+ torch.nn.TransformerEncoder,
+ torch.nn.TransformerDecoder
+ ]
+)
+
+# Fully-shard the optimizer.
+toy_adam = torch.optim.AdamW(params=fsdp_model.parameters(), lr=0.01)
+optimizer = fully_shard_optimizer(optimizer=toy_adam)
+
+# Forward pass.
+inp = torch.randn(1, 512, 512).to("cuda")
+tgt = torch.randn(1, 512, 512).to("cuda")
+output = fsdp_model(inp, inp)
+
+# Backward pass.
+torch.nn.functional.mse_loss(output, tgt).backward()
+
+# Optimizer step.
+optimizer.step()
+optimizer.zero_grad()
+
+# Checkpoint the model and optimizer.
+torch.distributed.checkpoint.save({
+ "model": fsdp_model.state_dict(),
+ "optimizer": optimizer.state_dict(),
+}, checkpoint_id="ckpt/")
+
+# Load the saved checkpoint.
+ckpt = {
+ "model": fsdp_model.state_dict(),
+ "optimizer": optimizer.state_dict(),
+}
+torch.distributed.checkpoint.load(state_dict=ckpt, checkpoint_id="ckpt/")
+fsdp_model.load_state_dict(ckpt["model"], strict=False)
+optimizer.load_state_dict(ckpt["optimizer"])
+```
+
+> ℹ️ `fully_shard` is an _**experimental**_ API. Please check back for updates as we fine-tune our user experience! For more examples using `fully_shard` for Megatron-FSDP, refer to our suite of unit tests: [`tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_fully_shard.py`](../../../tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_fully_shard.py)
+
+### 🤖 Megatron-LM
+
+Megatron-FSDP is deeply integrated into Megatron-Core. To enable FSDP (where optimizer states, gradients, and compute parameters are sharded) in Megatron, use the following arguments:
+
+```
+# Train models in Megatron-LM using Megatron-FSDP.
+--use-megatron-fsdp
+--data-parallel-sharding-strategy {no_shard, optim, optim_grads, optim_grads_params}
+--ckpt-format fsdp_dtensor
+```
+
+Complete Llama-8B and DeepSeek-V3 training scripts using Megatron-FSDP with recommended settings can be found in [Megatron-LM/examples/megatron_fsdp](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/megatron_fsdp).
+
+#### Recommended Configuration for Megatron-LM
+
+Frequently-used options use with Megatron-FSDP include:
+
+```bash
+# Un-set CUDA_DEVICE_MAX_CONNECTIONS to ensure stream independence / full-parallelization of FSDP computation and communication. May slightly affect TP and CP performance though.
+unset CUDA_DEVICE_MAX_CONNECTIONS
+
+# Meta-Device Initialization - Load large model onto CUDA devices in shards to avoid OOM.
+--init-model-with-meta-device
+
+# Per-Token Loss / No Gradient Scaling - Deactivate DP scaling during gradient reduction, which can be a drain on SM resources.
+--calculate-per-token-loss
+
+# Decrease gradient reduction and accumulation precision to recommended data-types based on the precision of the model parameters, usually BF16. Reduces communication volume during the backwards pass. Can be further customized with `--megatron-fsdp-main-grads-dtype` and `--megatron-fsdp-grad-comm-dtype`, which are enabled by this argument.
+--grad-reduce-in-bf16
+
+# Register NCCL user buffers and Megatron-FSDP double buffers to enable zero-copy symmetric kernels and low-SM utilization via SHARP. Improves overall performance but increases memory overhead due to double-buffering and is NOT compatible with `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.
+--use-nccl-ub
+--fsdp-double-buffer
+--fsdp-manual-registration
+```
+
+### 🤖 Megatron-Core
+
+Megatron-FSDP has a lower-level `FullyShardedDataParallel` class API that can be used with a simplified version of Megatron-LM's training loop.
+
+```python
+# Initialize model and optimizer.
+ddp_config.use_megatron_fsdp = True
+# Megatron-FSDP Base Sharding Strategies:
+# no_shard, optim, optim_grads, optim_grads_params
+ddp_config.data_parallel_sharding_strategy = "optim_grads_params"
+model = GPTModel(transformer_config)
+model = FullyShardedDataParallel(
+ transformer_config,
+ model,
+ ddp_config,
+ fsdp_unit_modules = [TransformerLayer, LanguageModelEmbedding],
+)
+optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+optimizer = DistributedOptimizer(optimizer, [model], [model.param_and_grad_buffer])
+
+# Training loop
+def train_step(inputs, labels):
+ optimizer.zero_grad()
+ for mbs_input, mbs_label in zip(inputs, labels):
+ outputs = model(mbs_input)
+ loss = loss_fn(outputs, mbs_label)
+ loss.backward()
+ optimizer.step()
+
+# Save and load model and optimizer state dict
+def model_and_optimizer_state_dict():
+ state_dict = {
+ "model": model.sharded_state_dict(),
+ "optimizer": optimizer.sharded_state_dict(),
+ }
+ return state_dict
+
+def load_model_and_optimizer_state_dict(state_dict):
+ model.load_state_dict(state_dict["model"])
+ optimizer.load_state_dict(state_dict["optimizer"])
+```
+
+### 🔁 Checkpoint Conversion
+
+Megatron-FSDP checkpointing supports [PyTorch Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html). In Megatron-LM, this is the `--ckpt-format fsdp_dtensor` checkpointing format.
+
+#### Converting Torch DCP to Torch Save (Non-Distributed) Checkpoints
+
+PyTorch has utilities to convert Torch DCP checkpoints to and from regular Torch checkpoints:
+```shell
+python -m torch.distributed.checkpoint.format_utils --help
+usage: format_utils.py [-h] {torch_to_dcp,dcp_to_torch} src dst
+
+positional arguments:
+ {torch_to_dcp,dcp_to_torch}
+ Conversion mode
+ src Path to the source model
+ dst Path to the destination model
+
+options:
+ -h, --help show this help message and exit
+```
+For example:
+```shell
+python -m torch.distributed.checkpoint.format_utils dcp_to_torch dcp_ckpt/ torch_ckpt.pt
+```
+or:
+```python
+from torch.distributed.checkpoint.format_utils import (
+ dcp_to_torch_save,
+ torch_save_to_dcp,
+)
+
+# Convert DCP model checkpoint to torch.save format.
+dcp_to_torch_save(CHECKPOINT_DIR, TORCH_SAVE_CHECKPOINT_PATH)
+
+# Convert torch.save model checkpoint back to DCP format.
+torch_save_to_dcp(TORCH_SAVE_CHECKPOINT_PATH, f"{CHECKPOINT_DIR}_new")
+```
+Torch Save checkpoints can then be converted into HuggingFace SafeTensors or other checkpoint formats for distribution.
+
+> ℹ️ Megatron-FSDP checkpoints have a `module.` prefix pre-pended to all model parameter names in the state dictionary, and converting a Torch Save checkpoint to a Megatron-FSDP Torch DCP checkpoint requires testing. Work-in-progress!
+
+#### Converting N-D Parallel (`torch_dist`) to Megatron-FSDP (`fsdp_dtensor`) Checkpoints
+
+As a pre-requisite for checkpoint conversion, dump the parameter group mapping when training with 3D-parallel (DDP, TP, PP) and/or EP:
+
+```bash
+--dump-param-to-param-group-map /path/to/param_to_param_group_map
+```
+
+and convert the map to a `param_to_param_group_map.json` JSON file in the `/path/to/param_to_param_group_map` directory:
+
+```bash
+python tools/checkpoint/checkpoint_inspector.py print-torch-dcp-in-json /path/to/param_to_param_group_map
+```
+
+> ℹ️ If 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 trivial training or checkpointing experiment to create the `param_to_param_group_map` you need without full pretraining.
+
+Finally, convert your `torch_dist` checkpoint to the `fsdp_dtensor` format using the `param_to_param_group_map.json`:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 \
+ tools/checkpoint/checkpoint_inspector.py \
+ convert-torch-dist-to-fsdp-dtensor (--swiglu) \ # --swiglu for specific models.
+ /path/to/input_torch_dist_checkpoint/ \
+ /path/to/output_fsdp_dtensor_checkpoint/ \
+ --param-to-param-group-map-json /path/to/param_to_param_group_map.json
+```
+
+> ℹ️ For multi-node conversion tasks, please refer to the DeepSeek-V3 example script (`sbatch_checkpoint_convert.sh`) in [Megatron-LM/examples/megatron_fsdp](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/megatron_fsdp).
+
+## Megatron-FSDP Feature Guide & API
+
+| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config |
+|--------------|-------------|----------------------|----------------------|
+| **Megatron-FSDP** | Use Megatron-FSDP in Megatron-LM. | `--use-megatron-fsdp` | `fully_shard_model(module)` |
+| **Megatron-FSDP Checkpointing** | Save and load un-even DTensor checkpoints using [Torch Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html). | `--ckpt-format fsdp_dtensor` | `preproc_state_dict_for_dcp_ckpt=True` |
+| **Meta Device Initialization** | Megatron-FSDP initializes a meta-initialized model to the CUDA device in shards to avoid OOM on large models. Requires implementation of `Module.reset_parameters()` for per-Module sharded initialization. | `--init-model-with-meta-device` | `init_model_with_meta_device=True` |
+| **Distributed Optimizer** | Megatron-FSDP uses Megatron-Core's `DistributedOptimizer`. Automatically set when using Megatron-FSDP. | `--use-distributed-optimizer` | `fully_shard_optimizer(optimizer)` |
+
+### FSDP Fundamentals
+
+```{figure} ../../images/megatron_fsdp/DDP_vs_FSDP.png
+:alt: FSDP Pipeline
+:align: center
+
+Comparison between Distributed Data Parallelism (DDP) and Fully-Sharded Data Parallelism (FSDP). While gradients are all-reduced in DDP, they are sharded and reduce-scattered with FSDP.
+
+Source: Meta AI, Ott, Myle, et al. “Fully Sharded Data Parallel: Faster AI Training with Fewer GPUs.” _Facebook Engineering_, 15 July 2021, https://engineering.fb.com/2021/07/15/open-source/fsdp/.
+```
+
+**Fully Sharded Data Parallelism (FSDP)** is a type of distributed data parallelism (DDP) that shards optimizer state, weight gradients (`wgrad`), and model weights across devices that ingest data-parallel samples for data-parallel training or inference. Activations (`fprop`) and data gradients (`dgrad`) are not sharded or distributed, and are preserved for the backward pass, but can be recomputed during the backward pass, offloaded to CPU, or sharded / routed using other parallelisms such as tensor parallelism (TP), context parallelism (CP), or expert parallelism (EP).
+
+```{figure} ../../images/megatron_fsdp/zero3_model_state.png
+:alt: ZeRO-3 Model State
+:align: center
+
+Sharded memory profiles for ZeRO-1 (optimizer state), ZeRO-2 (optimizer state and gradients), and ZeRO-3 (optimizer state, gradients, and parameters).
+
+Source: Zero-Redundancy Optimizer Model State Partition Diagram. From _The Ultra-Scale Playbook: Training LLMs on GPU Clusters_ by Tazi, Nouamane, et al. HuggingFace, 2025, https://huggingface.co/spaces/nanotron/ultrascale-playbook.
+```
+
+The core principles of FSDP are:
+
+- Only a small depth-wise fraction of the model state can exist un-sharded at any point in time.
+- Communication should overlap computation.
+
+From these core principles, software requirements can be derived:
+
+0. Model states sharded by FSDP are directly initialized across devices in shards.
+1. Model parameters are all-gathered (AG) in pre-designated groups or modules pre-forward and pre-backward to un-shard a small fraction of the model state at any point in time during training or inference. After `fprop` and `dgrad` computation, the un-sharded weights are immediately de-allocated.
+2. `wgrad` are reduce-scattered (RS) and accumulated in pre-designated groups or modules immediately post-backward to limit the amount of un-sharded gradients at any point in time during training or inference.
+3. Distributed optimizers, optimizers that are initialized with respect to a sharded model state and support distributed mechanics, update the sharded model state using the reduced gradient shard to implement data parallelism (DP).
+4. Computation and communication are overlapped across multiple CUDA streams, expending multiple streaming multi-processors (SM). Weights from subsequent groups or modules are pre-fetched, which ideally hides the communication latency required for FSDP behind model computation kernels (GEMM).
+
+FSDP can also be visualized as a decomposition of the all-reduce collective used in DDP into a gradient reduce-scatter, distributed optimization step, and parameter all-gather.
+
+```{figure} ../../images/megatron_fsdp/FSDP_Allreduce.png
+:alt: FSDP RS & AG
+:align: center
+
+Source: Feng, Wei, Will Constable, and Yifan Mao. “Getting Started with Fully Sharded Data Parallel (FSDP2).” _PyTorch Tutorials_, 17 Mar. 2022, https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html.
+```
+
+### FSDP Unit Modules
+
+| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config |
+|--------------|-------------|----------------------|----------------------|
+| **FSDP Unit Modules** | A list of `str` or `class` import paths for `torch.nn.Module`(s) that are considered FSDP unit modules and sharded by Megatron-FSDP. Parameters and sub-modules that are not members of an FSDP unit are not sharded. | Defaults to supported Megatron-Core modules (`TransformerLayer`, etc.) in Megatron-LM. | `fsdp_unit_modules=[...]` |
+| **FSDP Double Buffer Allocator** | Megatron-FSDP uses the double-buffer allocator, which persistently allocates a buffer pair assigned to alternating FSDP units that temporarily stores parameters and gradients. Automatically used with NCCL user buffer registration. | `--fsdp-double-buffer` | `fsdp_double_buffer=True` |
+| **Param All-Gather Overlap** | Whether to overlap parameter all-gather with compute. Automatically activated for the ZeRO-3 sharding strategy. | `--overlap-param-gather` | `overlap_param_gather=True` |
+| **Gradient Reduce-Scatter Overlap** | Whether to overlap gradient reduce-scatter or all-reduce with compute. Automatically activated for ZeRO-2 and ZeRO-3 sharding strategies. | `--overlap-grad-reduce` | `overlap_grad_reduce=True` |
+| **FSDP Communication Size** | Customize the size (in `numel()` elements) of AG and RS communications in Megatron-FSDP, by limiting how many elements are concurrently pre-fetched or reduced for AG and RS. Effectively suggests how many FSDP units are processed concurrently, which may launch collectives earlier and improve performance. Optionally, tune this value depending on system memory and performance requirements. | `--suggested-communication-unit-size ` | N/A (Megatron-Core Only) |
+
+> Only a small depth-wise fraction of the model state can exist un-sharded at any point in time.
+
+**FSDP Unit Modules** represent fractions of the model state that are computed and communicated as a (coalesced) group, un-sharded when needed for computation, and re-sharded after computation to release memory for subsequent model states. Implicitly, an FSDP unit module is also a **_modeling contract_**, requiring that FSDP-managed unit module parameters are not accessed or modified beyond the scope of the forward pass, backward pass, or optimization step.
+
+Megatron-FSDP accepts a list of `str` or `class` paths representing FSDP unit modules via the `fsdp_unit_modules` argument, which is currently hard-coded to supported model classes (like `TransformerLayer`) in Megatron-Core. It performs a depth-first traversal of the model (via `torch.nn.Module.named_modules()`) and groups the parameters of each matching module for sharding and coalesced communication. Nested units are resolved by precedence: if a module matches an FSDP unit class but is already a sub-module of a previously registered FSDP unit, it is skipped, so the outermost (and necessarily largest) FSDP unit class in any module sub-tree becomes the effective FSDP unit module.
+
+> Communication should overlap computation.
+
+Once a model is partitioned into unit modules, computation is overlapped with communication based on the granularity of the FSDP unit module. Depending on the size of the compute and communication kernels, fine-tuning the unit module size and grouping configuration can impact performance and elicit trade-offs between overlap and memory when using FSDP.
+
+```{figure} ../../images/megatron_fsdp/fsdp_streams.png
+:alt: FSDP Streams
+:align: center
+
+Each color-coded block in the compute and communication streams, merged and categorized in the simplified (and worst-case) scenario where SM resources are under contention, correspond to a _single_ FSDP unit module.
+```
+
+Compute-communication overlaps are orchestrated using **CUDA streams** that capture and parallelize serial operations. All collectives associated with all combinations of `{DP-Inner, DP-Outer}` and `{AG, RS}` are scheduled and tracked with separate streams and communicators / `ProcessGroup`(s).
+
+- Parameters are un-sharded prior to `fprop` and `dgrad` computation. To overlap the pre-fetch all-gather with computation, at least two FSDP units worth of un-sharded weight memory is required at any point in time.
+- Gradients are reduced and sharded after `wgrad` computation. To overlap gradient reduce-scatter with `wgrad` computation, at least two FSDP units worth of un-sharded gradient memory is required at any point in time.
+
+#### FSDP Module Hooks
+
+To implement these "unit-periodic" mechanics, Megatron-FSDP uses `Module` hooks to install a variety of (pre- and post-) forward and backward operations:
+
+- **Pre-Forward**
+ - Un-shards the model parameters of the current and (via pre-fetching) forward-subsequent FSDP unit modules.
+ - When `MegatronFSDP.forward()` is invoked, Megatron-FSDP will swap all parameter references to point to the un-sharded `Tensor` compute weights for the forward and backward pass.
+- **Post-Forward**
+ - Re-shards model weights after the forward pass, if the module is an FSDP unit. Non-unit modules remain persistently un-sharded.
+ - When using activation recomputation during the backwards pass, computing both `fprop` and `dgrad` requires these parameters, so parameters are resharded during **Post-Backward**.
+ - Releases the transpose cache of quantized parameters (in FSDP / ZeRO-3) for specific quantization recipes in `TransformerEngine`.
+- **Pre-Backward**
+ - Un-shards the model parameters of the current and (via pre-fetching) backward-subsequent FSDP unit modules.
+ - Implemented as a `torch.autograd.graph.register_multi_grad_hook` triggered by the output `dgrad`, and installed via a `Module` _post-forward_ hook.
+- **Post-Backward**
+ - Re-shards model weights after the backward pass, if the module is an FSDP unit. Non-unit modules remain persistently un-sharded.
+ - Implemented by injecting an Autograd function (`RegisterFSDPBackwardFunction`) that is installed during a `Module` _pre-forward_ hook.
+ - Reduces gradients after the backward pass.
+ - Implemented using a `Tensor.register_post_accumulate_grad_hook` triggered by `param.grad`, as well as a root-level post-backward hook installed during **Pre-Backward** (`torch.autograd.Variable._execution_engine.queue_callback`).
+- **State Dictionary**
+ - When `module.state_dict()` (for any module managed by Megatron-FSDP) is invoked, Megatron-FSDP will swap all parameter references to point to sharded `DTensor` main weights for distributed optimization and checkpointing.
+ - When `MegatronFSDP.load_state_dict()` is invoked, both the main and compute weights are updated. When using quantized model compute, the main weights are quantized and sharded.
+
+#### Double Buffering
+
+Megatron-FSDP uses a `Tensor._typed_storage()._resize_(bytes)`-based allocator to instantly allocate and de-allocate memory without depending on the `CUDACachingAllocator` for un-sharded parameters and gradients by default. (Cache fragmentation and garbage collection can procrastinate large quantities of `cudaMalloc` and `cudaFree` operations that can block programs and spike memory, particularly when memory utilization is maxed out.) However, modifying the underlying storage of a buffer is not compatible with NCCL symmetric registration or CUDA graphability, which require a persistent state during runtime.
+
+To support these optimizations, Megatron-FSDP uses **double-buffering**, which assigns 2 persistently-allocated buffers to FSDP units in an alternating pattern, hard-limiting the memory overhead for parameter and gradient buffer allocation and ensuring that no more than 2 FSDP units are computed or communicated concurrently.
+
+```{figure} ../../images/megatron_fsdp/fsdp_double_buffer.png
+:alt: FSDP Double Buffering
+:align: center
+
+Visualization of double buffering in Megatron-FSDP. Even- and odd-indexed FSDP units share the same un-sharded parameter and gradient buffers, overwriting incumbent data as needed during runtime. Megatron-FSDP ensures that no more than two FSDP units are un-sharded at any point during runtime.
+```
+
+With double-buffering, Megatron-FSDP does not need to allocate memory after initialization, which can reduce memory fragmentation and improve performance. However, double-buffering requires _depth-wise model symmetry_, where even- and odd-indexed FSDP units have identical size during runtime. If double-buffering is utilized, Megatron-FSDP computes the **_mode_** of FSDP unit sizes as the symmetrical double-buffer size, and any FSDP units not symmetrical to the computed size will default to the `_resize_(bytes)`-based allocator (or persistently allocated for extremely large and asymmetrical layers that affect performance significantly like `torch.nn.Embedding` when the low-level argument `fsdp_db_use_persist_buf_on_alloc_fail` is set).
+
+### Data-Parallel Sharding Strategies
+
+| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config |
+|--------------|-------------|----------------------|----------------------|
+| **Data Parallel Sharding Strategy** | Primary data-parallel sharding strategy for FSDP, which supports DDP, ZeRO-1 (optimizer), ZeRO-2 (optimizer and gradients), and ZeRO-3 (optimizer, gradients, and parameters). Typically uses intra-node communications, i.e. "inner" or "intra" DP. | `--data-parallel-sharding-strategy {no_shard, optim, optim_grads, optim_grads_params}` | `zero_dp_strategy={no_shard, optim, optim_grads, optim_grads_params, 0, 1, 2, 3}` |
+| **DP-Outer Sharding Strategy** | Secondary data-parallel sharding strategy for HSDP, which supports Hybrid-Sharded Data Parallel (HSDP / `no_shard`) and Hybrid-FSDP (HFSDP / `optim`). Typically uses inter-node communications, i.e. "outer" or "inter" DP. | `--outer-dp-sharding-strategy {no_shard, optim}` | `outer_dp_sharding_strategy={no_shard, optim, 0, 1}` |
+| **Hybrid Data Parallelism Size** | Specify the DP-Outer / Inter-DP parallel size. DP-Inner / Intra-DP sizes will be deduced from the sizes of other parallelisms and `torch.distributed.get_world_size()`. | `--num-distributed-optimizer-instances ` | `dp_outer_dim=` (Cumulative DP groups `hybrid_fsdp_group` / `hybrid_fsdp_expt_group` are required for HFSDP.) |
+
+Megatron-FSDP supports a variety of sharding strategies over a variety of distributed topologies:
+
+- **Distributed Data Parallelism (DDP)**
+ - Model state is replicated across DP ranks.
+ - Gradient all-reduce is overlapped with backward compute and launched during the last backward pass before the optimization step.
+- **ZeRO-1**
+ - Optimizer state is sharded across DP ranks.
+ - Gradient reduce-scatter is overlapped with backward compute and launched during the last backward pass before the optimization step. (Reduce-scatter is used in lieu of all-reduce for performance, because only a shard of the gradient is needed for optimization.)
+- **ZeRO-2**
+ - Optimizer state and gradients are sharded across DP ranks.
+ - Gradient reduce-scatter is overlapped with backward compute and accumulated during every backward pass.
+- **Fully-Sharded Data Parallelism (FSDP / ZeRO-3)**
+ - Optimizer state, gradients, and parameters are sharded across DP ranks.
+ - Gradient reduce-scatter is overlapped with backward compute and accumulated during every backward pass.
+- **Hybrid-Sharded Data Parallelism (HSDP)**
+ - Optimizer state, gradients, and parameters are sharded across the "inner" or "intra" DP ranks.
+ - Model state is replicated across "outer" / "inter" DP ranks, and outer data-parallel gradients are all-reduced during the last backward pass before the optimization step.
+- **Hybrid-FSDP (HFSDP)**
+ - Optimizer state, gradients, and parameters are sharded across the "inner" or "intra" DP ranks.
+ - Optimizer state is _further_ sharded across "outer" / "inter" DP ranks.
+ - Outer data-parallel gradients are reduce-scattered after during the last backward pass before the optimization step.
+ - Outer data-parallel parameters are all-gathered during the first forward pass after the optimization step.
+ - FSDP primary sharding (`optim_grads_params`) is required for HFSDP secondary sharding (`optim`).
+ - Requires passing cumulative data-parallel groups (`hybrid_fsdp_group` / `hybrid_fsdp_expt_group`), which include ALL data-parallel ranks, to Megatron-FSDP.
+ - To create these using `DeviceMesh`, create a data-parallel `DeviceMesh` for the cumulative DP group and use `DeviceMesh._unflatten(dp_dim, mesh_sizes=(dp_outer_size, dp_inner_size), mesh_dim_names=("dp_outer_dim", "dp_shard_dim"))` to construct a `DeviceMesh` with DP-Inner and DP-Outer mesh dimensions for Hybrid-FSDP.
+
+#### Understanding Hybrid-FSDP (HFSDP)
+
+```{figure} ../../images/megatron_fsdp/hfsdp.png
+:alt: Hybrid-FSDP Topology
+:align: center
+
+Hybrid-FSDP (HFSDP) is a variation of HSDP where the optimizer state in particular is sharded across both DP-Inner and DP-Outer, i.e. all data-parallel ranks, which further reduces memory utilization. In other words, intra-node sharding and communication uses ZeRO-3, while inter-node sharding and communication uses ZeRO-1. Parameters and gradients are converted from and to the fully-sharded optimizer state during optimization steps only, reducing the frequency of inter-node communications.
+
+Inspired by the artistry in the DHEN (Zhang, Luo, Liu, Meta, et al., 2022) paper: https://arxiv.org/abs/2203.11014
+```
+
+**Hybrid-Fully Sharded Data Parallelism (HFSDP)** is a slight modification to HSDP that fully-shards the optimizer state across all data-parallel ranks and introduces outer-level all-gather and reduce-scatter collectives to map fully-sharded parameters and gradients into partially-sharded parameters and gradients.
+
+The memory profile of HFSDP is a "hybrid" of FSDP (optimizer state) and HSDP (gradients and model weights). Another elegant way to understand HFSDP functionality is ZeRO-1 composed with ZeRO-3.
+
+$$\text{Hybrid-FSDP Memory Profile} = \frac{\text{Optimizer State}}{\text{DP-Inner} \ \times \ \text{DP-Outer}} + \frac{\text{Gradient} + \text{Weight}}{\text{DP-Inner}}$$
+
+The modified algorithm has the following characteristics:
+
+- Megatron-FSDP maintains a view of the model parameters sharded across all data-parallel ranks.
+ - Distributed checkpoints save and load the fully-sharded model parameters.
+ - Distributed optimizer state is initialized on the fully-sharded model parameters.
+- During the first forward pass after checkpointing or optimization, fully-sharded model weights are all-gathered into partially-sharded model weights.
+- During the last backward pass before optimization, partially-sharded model gradients are reduce-scattered into fully-sharded model gradients.
+- Otherwise, FSDP is performed on the partially-sharded model weights and accumulated gradients. Because model weights and gradients are only updated and ingested once per optimization cycle, we can skip or postpone all expensive inter-node / DP-outer collectives until an optimization step.
+
+In addition to improved memory utilization, HFSDP communications are split in communication size (bytes communicated), communication topology (DP-Inner and DP-Outer groups), and communication domain (NVLink and InfiniBand) across two sharding stages.
+
+```{figure} ../../images/megatron_fsdp/fsdp_v_hfsdp_streams.png
+:alt: Hybrid-FSDP Streams
+:align: center
+
+Inter-node communications can also be parallelized with intra-node communications using separate CUDA streams.
+```
+
+#### Mixing FSDP & Model Parallelism
+
+Megatron-FSDP is also compatible with a variety of model parallelisms that shard the model state, such as **Tensor Parallelism (TP)** and **Expert Parallelism (EP)**. When sharding model states across multiple dimensions in the device topology, _**FSDP sharding is always performed last**_, because FSDP collectives un-shard and re-shard parameters and gradients immediately before and after computation. Thus, FSDP sharding mechanics are implemented over tensor and expert parallel (strided) shards.
+
+```{figure} ../../images/megatron_fsdp/mixed_sharding.png
+:alt: Mixed Model Parallelism
+:align: center
+
+Wheneveer FSDP is composed with other model parallelisms, FSDP sharding is always exercised last to seamlessly integrate with existing model shards.
+```
+
+Megatron-FSDP uses `torch.distributed.DeviceMesh` to describe and configure communications across devices in data-parallel group(s). Because heterogeneous models that have mixed layers, such as [Hybrid Mamba-Transformer](https://arxiv.org/abs/2504.03624) or [Mixture-of-Experts (MoE)](https://arxiv.org/abs/1701.06538) models, require different parallelism configurations, multiple `DeviceMesh`(s) may be required for specific layers that require distinct distributed topologies for optimal memory efficiency and performance.
+
+Currently, Megatron-FSDP supports two `DeviceMesh`(s), one for dense / non-expert `Module`(s) and another for Megatron-Core MoE sparse / expert `Module`(s). (Expert modules and parameters in Megatron-Core are automatically detected.)
+
+- Dense modules typically have a `DeviceMesh` with data parallel, tensor parallel, and context parallel dimensions, where the data parallel dimension is used for FSDP. Typically, both data-parallel and context-parallel ranks are used for sharding in FSDP.
+- Mixture-of-experts modules typically have a `DeviceMesh` with data parallel, tensor parallel, and expert parallel dimensions, where the data parallel dimension is used for FSDP.
+
+For more information about Mixture-of-Experts in Megatron-Core, refer to the [Megatron-Core User Guide - MoE](https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/features/moe.html).
+
+#### Non-Uniform / Un-Even Model Sharding
+
+While `torch.distributed.tensor.DTensor` defaults to per-parameter sharding, where Tensors are split evenly on `dim=0` across the data-parallel domain, Megatron-FSDP uses **non-uniform or un-even `DTensor` shards** of a (flattened) group of parameters associated with an FSDP unit.
+
+```{figure} ../../images/megatron_fsdp/uneven_sharding.png
+:alt: Non-Uniform Sharding
+:align: center
+
+Comparison of FSDP2 per-parameter sharding and Megatron-FSDP per-unit or per-module sharding. FSDP2 requires `COPY` operations to move parameters and gradients in and out of communication buffers to reduce the frequency of NCCL collective calls, while Megatron-FSDP assigns sliced views of contiguous communication buffers to parameters associated with an FSDP unit.
+```
+
+While complex and less user-intuitive, an un-evenly sharded data structure enables a few performance benefits without introducing expensive `COPY` operations to set up communication and computation buffers:
+
+- **Fewer NCCL calls**, reducing kernel launch and synchronization overhead. Only parameters in FSDP units that have different communication-related properties, such as their `dtype` or distributed topology, are coalesced into separate NCCL calls.
+- Flat communication and computation buffers are **contiguous-by-design**, supporting optimized CUDA kernels that require buffers backed by contiguous memory, such as grouped GEMMs used in MoE.
+
+Effectively, this implies that the same `DTensor`-sharded model parameters may have completely different shapes on different ranks, and if entire parameters are assigned to other ranks, the local `Tensor` will be empty.
+
+> ℹ️ Megatron-FSDP has a handy library ([`megatron_fsdp.uneven_dtensor`](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py)) for manipulating un-evenly sharded `DTensors`, focused on per-parameter operations like un-sharding or reducing parameters that have different shapes across ranks. While the parameter group is evenly-sharded for FSDP collectives, per-parameter collectives (that assume a symmetrical amount of bytes are communicated between devices) will hang waiting on bytes that will never arrive for un-evenly sharded `DTensors`.
+
+In particular, contiguous memory is only half the requirement for high-performance CUDA kernels. The other requirement is **locality**, which FSDP can violate, that introduces compatibility issues when combining FSDP with present and future optimizations. For example, block-wise quantization (scaling factor / `absmax` calculations for MXFP8, NVFP4, etc.) requires DP communication and custom max-reduce kernels if the block is sharded by FSDP.
+
+Megatron-FSDP supports `dim=0` sharding, which computes the _**least-common multiple (LCM) of `p.shape[1:]` for all parameters `p` in an FSDP unit**_ and _**pads the un-sharded buffer to the closest multiple of `DP x LCM(p.shape[1:])`**_, forming a "DP-LCM" partition with `LCM`-length parts to ensure that DP-sharding boundaries do not violate chunks of data for coordinates of `dim=0`.
+
+```{figure} ../../images/megatron_fsdp/lcm_dim0_shard.png
+:alt: Flat Buffer Sharding Algorithm
+:align: center
+
+Visualization of how parameters are assigned un-evenly to the flat per-unit buffer sharded across DP ranks. With the LCM algorithm, every slice of `dim=0` is never bisected by FSDP. Algorithms and compute kernels can leverage this locality and contiguity.
+```
+
+1. When a parameter is _divisble by the LCM_, it can be inserted at any index multiple of the LCM in the buffer that is free. `p[i]` chunks of this parameter by definition divide the LCM, and thus align with the DP-LCM sharding grid.
+2. When a parameter _is larger than but not divisible by the LCM_, the remainder `r` populates a fraction of another LCM part, so a "conjugate" parameter that also exceeds the LCM with a "conjugate" remainder `r'` that is less than or equal to `LCM - r` is installed to fill the remaining space and align with the DP-LCM sharding grid.
+3. When a parameter _is smaller than but not divisible by the LCM_, a post-assignment sweep on the leftover space in the flat buffer is run, and all gaps that are multiples of the LCM that are large enough to support the entire parameter are utilized. Once all gaps are filled, the final parameters are assigned to the tail of the buffer respecting the DP-LCM sharding grid.
+
+> ℹ️ Generalized support for contiguity and locality in Megatron-FSDP is a **_work-in-progress_** and will evolve with contribution from the OSS community and PyTorch. For more information about how kernel buffer requirements affect the design of FSDP data structures, refer to the [veScale: Consistent and Efficient Tensor Programming with Eager-Mode SPMD (Li, Youjie, ByteDance Seed, et al.)](https://arxiv.org/abs/2509.07003) paper that comprehensively analyzes these requirements.
+
+### Mixed-Precision & Quantization
+
+| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config |
+|--------------|-------------|----------------------|----------------------|
+| **Quantized Parameters** | Megatron-FSDP will shard and all-gather TransformerEngine-quantized parameters for computation. Quantized parameters are updated every optimization step, and both row-wise (FWD) and column-wise (BWD) data are managed for non-transposable 1-D quantization recipes like MXFP8. Otherwise, only activations are quantized. | `--fp8-param-gather` | TransformerEngine `quantized_model_init()` |
+| **Main Parameter (Optimization / Checkpoint) Data-Type** | Data-type for optimization and checkpointing parameters. If set to `auto`, model compute weights are utilized instead. Required for `--fp8-param-gather`. Defaults to FP32. | `--megatron-fsdp-main-params-dtype {fp32, bf16, fp16, auto}` | `MixedPrecisionPolicy(main_params_dtype=...)` |
+| **Main Gradient (Accumulation) Data-Type** | Data-type for gradient accumulation. If set to `auto`, main gradient precision will be derived from model parameter precision. Defaults to `auto`. | `--megatron-fsdp-main-grads-dtype {fp32, bf16, fp16, auto}` | `MixedPrecisionPolicy(main_grads_dtype=...)` |
+| **Gradient Communication (Reduction) Data-Type** | Data-type for gradient communication and reduction. If set to `auto`, the main gradient precision will be used for communication. (When using NCCL symmetric registration, low-precision gradients are reduced in FP32 over-the-wire.) Defaults to `auto`. | `--megatron-fsdp-grad-comm-dtype {fp32, bf16, fp16, auto}` | `MixedPrecisionPolicy(grad_comm_dtype=...)` |
+| **Weight Gradient Accumulation Fusion** | When using TransformerEngine modules, Megatron-FSDP implements `get_main_grad` to allocate un-sharded gradient buffers called by TransformerEngine, to avoid `COPY`-ing the gradient to Megatron-FSDP communication buffers. Used by default and can be deactivated with `--no-gradient-accumulation-fusion`. | `--no-gradient-accumulation-fusion` | N/A (Megatron-Core Only) |
+| **Precision-Aware Optimizer** | Use the TransformerEngine `FusedAdam` optimizer, and Megatron-FSDP will install the gradient in a temporary attribute `Parameter.decoupled_grad` which is consumed by `FusedAdam`. Megatron-FSDP manages the main parameters, but the optimizer state precision can be customized with `--exp-avg-dtype` and `--exp-avg-sq-dtype`, which both support `fp8` optimization state. | `--use-precision-aware-optimizer` | `use_decoupled_grad=True` |
+
+#### Quantization
+
+Quantization is an extremely important feature for Megatron-FSDP as it reduces memory utilization and communication size for both activations and parameters, which directly affects the viability and performance of FSDP.
+
+```{figure} ../../images/megatron_fsdp/quantized_param_gather.png
+:alt: Quantized Model Parameters & FSDP
+:align: center
+
+Visualization of Megatron-FSDP's training loop when using quantized weights from TransformerEngine. Every optimization step updates the quantized representation of sharded model weights, which have reduced communication size.
+```
+
+While TransformerEngine handles activation quantization, Megatron-FSDP shards quantized weights for AG.
+
+0. _**Quantized Model Initialization**_ - Model is initialized with quantized weights, e.g. MXFP8 or NVFP4. If using `meta` device initialization, Megatron-FSDP will call `reset_parameters()` to initialize quantized weights layer-by-layer. If row-wise and column-wise data are not transposable, Megatron-FSDP will shard and buffer both. Additionally, high-precision main weights are retrieved and sharded for distributed optimization, checkpointing, and quantization.
+0. _**Forward / Backward Pass**_ - Quantized weights are un-sharded for both the forward and backward pass. If row-wise and column-wise data aren't transposable, the row-wise weights are gathered for forward, and the column-wise weights are gathered for backward.
+0. _**Distributed Optimization Step**_ - Non-quantized accumulated gradient shards from quantized GEMMs are applied to high-precision main weight shards.
+0. _**Sharded Quantization**_ - Sharded main weights are quantized to update the quantized compute weights for subsequent training steps.
+
+```{figure} ../../images/megatron_fsdp/sharded_quantization.png
+:alt: Sharded Quantization
+:align: center
+
+Sharded quantization involves reducing maxima to compute a global set of scaling factors for local / sharded quantization.
+```
+
+In particular, _sharded quantization_ minimizes communication size and memory utilization by communicating scaling factors instead of main weights.
+
+1. _**Local Abs-Max**_ - For a group of parameters in an FSDP unit, compute local tensor-wise or block-wise maxima across the global un-sharded shape, with zero padding for non-local data.
+1. _**Global Abs-Max**_ - Globally all-reduce maxima and derive scaling factors from maxima.
+1. _**Local Quantization**_ - Locally quantize sharded main weights and install into compute weight buffers.
+
+#### Mixed-Precision
+
+Megatron-FSDP sharding and communication buffers support mixed-precision, such that users can customize the `dtype` used for main weights, gradient communication (reduction), and gradient accumulation in addition to the native or quantized `dtype` used for model computation. These options are wrapped in a `MixedPrecisionPolicy` dataclass.
+
+- _**Main Weight Precision**_ - Controls the data-type for parameters responsible for distributed optimization, distributed checkpointing, and quantization. If set to `auto` (`None`), the native model compute parameter data-type will be utilized. Required for parameter quantization with `--fp8-param-gather`. Defaults to `torch.float32`.
+- _**Main Gradient Precision**_ - Controls the data-type for `wgrad` accumulation and distributed optimization. Defaults to `auto` (`None`), the model native gradient data-type will be utilized. While `torch.float32` (or higher) is recommended for accuracy at scale, as `main_grads_dtype` controls the data-type for gradient accumulation, `auto` is more flexible and uses pre-determined parameter gradient logic in mixed-precision scenarios, such as `BF16` for `FP8`/`FP4` parameters quantized via TransformerEngine.
+- _**Gradient Communication Precision**_ - Controls the data-type for gradient communications when reducing gradients. Lower precision improves (communication) performance. Defaults to `auto` (`None`), in which the main gradient data-type will be utilized. If using `no_shard`, `optim`, HSDP, or HFSDP, allocating `dtype`-custom gradient communication buffers may increase per-unit memory overhead, so users should consider the performance-memory trade-off when using this feature.
+ - If using NCCL symmetric registration `v2.27+`, gradient reduction may be performed in high-precision depending on the network domain (NVLink or IB), and can enable mixed-precision communication and accumulation, e.g. setting grad_comm_dtype to `BF16` can support `FP32` reduction even though we have `BF16` input and output communication buffers. Otherwise, gradients will be reduced and accumulated in communication and accumulation precision as usual.
+
+### NCCL
+
+| Optimization | Description | `Megatron-Core` Config | `fully_shard` Config |
+|--------------|-------------|----------------------|----------------------|
+| **NCCL User Buffers** | Allocate and register Megatron-FSDP communication buffers with NCCL, which enables zero-`COPY`, high-precision reduction, copy-engine collectives, and symmetric kernels. Uses double buffering. | `--use-nccl-ub` | `nccl_ub=True` |
+| **NCCL Manual Registration** | Instead of registering NCCL user buffers on first allocation, batch registration of all communication buffers at the end of the initial training step. Reduces registration latency. | `--fsdp-manual-registration` | N/A (Megatron-Core Only) |
+| **Disable Symmetric Registration** | Disable symmetric registration with NCCL. Optional, as symmetric registration failure defaults to normal registration. | `--disable-symmetric-registration` | `disable_symmetric_registration=True` |
+
+[NVIDIA Collective Communications Library (NCCL)](https://developer.nvidia.com/nccl) implements multi-device and multi-node communication primitives optimized for CUDA devices and networking from NVIDIA. Megatron-FSDP communications are registered and deeply integrated with NCCL, which enables a variety of hardware-level networking optimizations such as copy-engine AG, high-precision RS, SHARP reduction offloading, and symmetric kernels.
+
+To leverage NCCL networking optimizations, **NCCL user buffer registration (UBR)** is required to inform NCCL of PyTorch Tensors ("user buffers") that act directly as the input and target of NCCL collectives for PyTorch `ProcessGroup`(s). Because registered communication buffers are known to NCCL, `COPY` operations that send collective inputs to NCCL buffers and collective outputs to PyTorch buffers are no longer required, which enables Megatron-FSDP to be zero-`COPY` end-to-end.
+
+NCCL (`v2.27+`) supports symmetric allocation or registration for communicators over the NVLink domain, which allow buffers that share identical virtual addresses across devices to benefit from optimized collectives:
+
+- **Symmetric Kernels** - On the NVLink domain, symmetric kernels operating on symmetric memory reduces the SM utilization for a single communication kernel to 1.
+- **NVSwitch SHARP Offloading** - To further minimize SM utilization for AG and RS collectives, NCCL SHARP offloads reduction and aggregation work to NVLink and IB Switch hardware that uses 1-6 SM depending on the domain: NVL, IB, or NVL + IB.
+- **Copy-Engine (CE) Collectives**: Instead of using SMs (or CTAs) for common non-computational collectives like AG in Megatron-FSDP, copy engines are instead used to perform all-gather collectives, dedicating SM resources to compute and reduction during FSDP. Requires NCCL `v2.28+`.
+- **High-Precision Reduction**: When training large models, high-precision gradient reduction and accumulation is desired for accuracy and convergence, but communicating FP32 gradients is expensive. With symmetric registration, FP32 accumulators enable gradients to be reduced in FP32 but communicated in BF16, which decreases gradient RS communication latency while maintaining high accuracy during training. Megatron-FSDP supports FP32 main gradient accumulation but BF16 gradient communication, customizable through `megatron_fsdp.MixedPrecisionPolicy`.
+
+These optimizations significantly reduce SM resource contention for overlapped compute and communication kernels in FSDP. Symmetric registration, allocation, and pooling is also supported in PyTorch: [`torch.distributed._symmetric_memory`](https://docs.pytorch.org/docs/stable/symmetric_memory.html).
diff --git a/docs/user-guide/features/megatron_rl.md b/docs/user-guide/features/megatron_rl.md
new file mode 100644
index 00000000000..9cd46d79ae2
--- /dev/null
+++ b/docs/user-guide/features/megatron_rl.md
@@ -0,0 +1,55 @@
+
+
+# Megatron RL
+
+Reinforcement learning library for post-training large language models at scale.
+
+## Overview
+
+[**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.
+
+> **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).
+
+## Key Features
+
+- **Decoupled Design** - Separates agent and environment logic from the core RL implementation
+- **Inference Backends** - Megatron, OpenAI, and Hugging Face inference stacks
+- **Trainer or Evaluator** - Manages rollout generation and coordinates with inference systems
+- **Megatron Integration** - Native integration with Megatron Core inference system
+
+## Architecture
+
+### Components
+
+**Agents and Environments**
+- Accept inference handles
+- Return experience rollouts with rewards
+- Implement custom RL logic
+
+**Trainer or Evaluator**
+- Controls rollout generation
+- Coordinates with inference systems
+- Manages training loops
+
+**Inference Interface**
+- Exposes a `.generate(prompt, **generation_args)` endpoint
+- Supports multiple backends (Megatron, OpenAI, Hugging Face)
+
+## Use Cases
+
+- RLHF (Reinforcement Learning from Human Feedback)
+- Custom reward-based fine-tuning
+- Policy optimization for specific tasks
+- Research on RL post-training techniques
+
+## Resources
+
+- **[Megatron RL GitHub](https://github.com/NVIDIA/Megatron-LM/tree/dev/megatron/rl)**: Source code and documentation
+- **[Megatron Core Inference](../../api-guide/core/transformer.md)**: Native inference integration
diff --git a/docs/user-guide/features/moe.md b/docs/user-guide/features/moe.md
new file mode 100644
index 00000000000..45efe291a21
--- /dev/null
+++ b/docs/user-guide/features/moe.md
@@ -0,0 +1,22 @@
+
+
+# Mixture of Experts
+
+```{toctree}
+:maxdepth: 1
+:caption: MoE Features
+
+multi_token_prediction
+multi_latent_attention
+../../api-guide/router_replay
+```
+
+```{include} ../../../megatron/core/transformer/moe/README.md
+```
diff --git a/docs/user-guide/features/multi_latent_attention.md b/docs/user-guide/features/multi_latent_attention.md
new file mode 100644
index 00000000000..65a1e573c6c
--- /dev/null
+++ b/docs/user-guide/features/multi_latent_attention.md
@@ -0,0 +1,21 @@
+
+
+# Multi-Latent Attention
+
+## Multi-Latent Attention Overview
+
+Multi-Latent Attention (MLA) is an attention variant from the DeepSeek team. It uses multiple latent spaces to change how attention is computed. That design often lowers cost for large language models (LLMs) compared with standard attention and can shrink the KV cache. The DeepSeek-V2 technical report compares MLA to Multi-Head Attention (MHA) on quality and cache size.
+
+## Enabling Multi-Latent Attention
+
+To enable MLA in Megatron-LM, set the following on the command line:
+
+- `--multi-latent-attention` to turn on MLA.
+- Use `MLATransformerConfig` for MLA-specific model settings when you build the training configuration.
diff --git a/docs/user-guide/features/multi_token_prediction.md b/docs/user-guide/features/multi_token_prediction.md
new file mode 100644
index 00000000000..e2d51c1b705
--- /dev/null
+++ b/docs/user-guide/features/multi_token_prediction.md
@@ -0,0 +1,58 @@
+
+
+# Multi-Token Prediction (MTP)
+
+Multi-Token Prediction (MTP) extends the prediction scope to several future tokens at each position. An MTP objective adds extra prediction targets, which can improve data efficiency. It may also encourage representations that anticipate later tokens. This implementation predicts additional tokens in sequence and preserves the causal dependency chain at each depth. The following figure illustrates MTP as used in [DeepSeek-V3](https://github.com/deepseek-ai/DeepSeek-V3/).
+
+
+
+The *k*-th MTP module includes a shared embedding layer, a projection matrix, a Transformer block, and a shared output head. For the *i*-th input token at depth *k - 1*, the implementation combines the representation of the *i*-th token and the embedding of the *(i + K)*-th token with a linear projection. That combined representation is the input to the Transformer block at depth *k*, which produces the output representation.
+
+For more detail, refer to the [DeepSeek-V3 technical report](https://arxiv.org/pdf/2412.19437.pdf).
+
+## Related Arguments
+
+Train `GPTModel`-style models with MTP by setting `mtp_num_layers` to a positive integer.
+
+The following table summarizes MTP configuration fields:
+
+| Item | Description |
+| --- | --- |
+| `mtp_num_layers` | Number of MTP layers. MTP extends prediction to multiple future tokens at each position. This stack uses `mtp_num_layers` sequential modules to predict that many additional tokens per position. Default: `None`. |
+| `mtp_loss_scaling_factor` | Weight for the MTP loss term. The implementation averages MTP losses across depths, multiplies by this factor, and adds the result to the training objective. Default: `0.1`. |
+
+## Pipeline Parallel Layout for MTP
+
+MTP supports user-defined placement of MTP layers across pipeline stages through `pipeline_model_parallel_layout`. By default, all MTP layers sit on the last pipeline stage; you can override placement in the layout string.
+
+### MTP Standalone Mode
+
+When 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`. MTP then runs in its own pipeline stage.
+
+### Layout Format
+
+Use `m` for MTP layers in the pipeline layout string. For example:
+- `"E|t*3|(t|)*5mL"` - MTP in the last stage
+- `"E|t*3|(t|)*4tm|L"` - MTP in the second-to-last stage with a decoder layer
+- `"E|t*3|(t|)*3tt|m|L"` - MTP in a standalone stage (second-to-last) with no other layers
+
+### Constraints
+
+- Place all MTP layers in the same virtual pipeline stage.
+- Do not place MTP layers on the first pipeline rank.
+
+## Implementation Notes
+
+- For models with MTP layers, the final LayerNorm sits in the stage that contains the last decoder layer, not in the post-process stage. That can change gradient norm reduction slightly in deterministic mode when LayerNorm would otherwise live in another stage. For bitwise alignment, disable gradient norm clipping.
+- MTP loss is computed in the post-processing stage.
+
+## Unsupported Combinations
+
+Context Parallel (CP), arbitrary `AttnMaskType`, and learned absolute position embeddings are not supported with MTP.
diff --git a/docs/user-guide/features/optimizer_cpu_offload.md b/docs/user-guide/features/optimizer_cpu_offload.md
new file mode 100644
index 00000000000..1496bd0a91e
--- /dev/null
+++ b/docs/user-guide/features/optimizer_cpu_offload.md
@@ -0,0 +1,13 @@
+
+
+# Optimizer CPU Offload
+
+```{include} ../../../megatron/core/optimizer/cpu_offloading/README.md
+```
diff --git a/docs/user-guide/features/pipeline_parallel_layout.md b/docs/user-guide/features/pipeline_parallel_layout.md
new file mode 100644
index 00000000000..69ffb63da5a
--- /dev/null
+++ b/docs/user-guide/features/pipeline_parallel_layout.md
@@ -0,0 +1,37 @@
+
+
+# Custom Pipeline Model Parallel Layout
+
+*This is an experimental feature and may be changed.*
+
+`--pipeline-model-parallel-layout` takes a string that defines pipeline parallel partitioning. Use it to balance partitioning for an imbalanced model. For example, to partition a DeepSeek-V3-style stack (61 decoder layers and one MTP layer) with PP16 and VPP2, pass arguments similar to the following:
+
+```bash
+--pipeline-model-parallel-size 16
+--pipeline-model-parallel-layout "Et*3|(tt|)*29,m|L"
+```
+
+The table below shows one possible rank map for that layout:
+
+| PP \ VPP rank | 0 | 1 |
+|---------------|-------------------------|---------------|
+| 0 | embedding + 3 × decoder | 2 × decoder |
+| 1~13 | 2 × decoder | 2 × decoder |
+| 14 | 2 × decoder | mtp |
+| 15 | 2 × decoder | loss |
+
+In the layout string, stages are split by `|`. Replicated stages or layers use multiplication (for example, `t*3`). Commas are optional for readability. Symbols:
+
+* `E`: embedding layer
+* `t`: transformer decoder layer
+* `m`: MTP layer
+* `L`: loss calculation layer
+
+**Note:** Empty stages are allowed, for example `E||t|L` (the second stage is empty).
diff --git a/docs/user-guide/features/tokenizers.md b/docs/user-guide/features/tokenizers.md
new file mode 100644
index 00000000000..1455d6e617e
--- /dev/null
+++ b/docs/user-guide/features/tokenizers.md
@@ -0,0 +1,261 @@
+
+
+# Tokenizers
+
+Megatron Core provides a unified tokenizer system with a Hugging Face-style API for configuration and loading.
+
+## Overview
+
+The `MegatronTokenizer` class uses the same entry points as many Hugging Face workflows for loading and managing tokenizers:
+
+- **Automatic detection** - Load tokenizer types without naming the backing library in code
+- **Metadata-based configuration** - Store tokenizer settings in JSON for reuse across runs
+- **Hugging Face-compatible API** - `.from_pretrained()`-style loading
+- **Custom tokenizer support** - Extend with model-specific tokenization logic
+
+## Key Features
+
+### Unified API
+
+Use the same API regardless of tokenizer backend (SentencePiece, Hugging Face, TikToken, and so on):
+
+```python
+from megatron.core.tokenizers import MegatronTokenizer
+
+tokenizer = MegatronTokenizer.from_pretrained("/path/to/tokenizer")
+```
+
+### Tokenizer Metadata
+
+Configuration is stored in a JSON metadata file containing:
+
+- Tokenizer library (Hugging Face, SentencePiece, TikToken, and so on)
+- Chat templates
+- Custom tokenizer class
+- Special token configurations
+
+**Benefits**
+
+- Set configuration once, reuse everywhere
+- No repeated CLI arguments
+- Share setups by copying the tokenizer directory
+
+### Automatic Library Detection
+
+The correct tokenizer implementation is selected automatically:
+
+- Avoids hard-coding `SentencePieceTokenizer`, `HuggingFaceTokenizer`, and related class names in user code
+- Library type is read from metadata
+- Change tokenizer backends by updating metadata and paths
+
+## Basic Usage
+
+### Creating Tokenizer Metadata
+
+Save tokenizer configuration for reuse:
+
+```python
+from megatron.core.tokenizers import MegatronTokenizer
+
+# Create metadata for a SentencePiece tokenizer
+MegatronTokenizer.write_metadata(
+ tokenizer_path="/path/to/tokenizer.model",
+ tokenizer_library="sentencepiece",
+ chat_template="{% for message in messages %}{{ message.content }}{% endfor %}",
+)
+```
+
+The metadata is saved as `tokenizer_metadata.json` in the tokenizer directory.
+
+### Loading a Tokenizer
+
+Load from a directory with metadata:
+
+```python
+from megatron.core.tokenizers import MegatronTokenizer
+
+# Load with auto-detected configuration
+tokenizer = MegatronTokenizer.from_pretrained("/path/to/tokenizer.model")
+```
+
+### Loading with Custom Metadata Path
+
+If metadata is stored separately:
+
+```python
+tokenizer = MegatronTokenizer.from_pretrained(
+ tokenizer_path="/path/to/tokenizer.model",
+ metadata_path="/path/to/custom/metadata.json",
+)
+```
+
+### Loading with Inline Metadata
+
+Pass metadata as a dictionary:
+
+```python
+tokenizer = MegatronTokenizer.from_pretrained(
+ tokenizer_path="GPT2BPETokenizer",
+ metadata_path={"library": "megatron"},
+ vocab_file="/path/to/vocab.txt",
+)
+```
+
+## Advanced Usage
+
+### Custom Tokenizer Classes
+
+Create model-specific tokenization logic:
+
+```python
+from megatron.core.tokenizers.text import MegatronTokenizerText
+
+class CustomTokenizer(MegatronTokenizerText):
+ def encode(self, text):
+ # Custom encoding logic
+ return super().encode(text)
+
+ def decode(self, tokens):
+ # Custom decoding logic
+ return super().decode(tokens)
+
+# Save metadata with custom class
+MegatronTokenizer.write_metadata(
+ tokenizer_path="/path/to/tokenizer.model",
+ tokenizer_library="sentencepiece",
+ tokenizer_class=CustomTokenizer,
+)
+```
+
+### TikToken Tokenizers
+
+Configure TikToken-based tokenizers:
+
+```python
+tokenizer = MegatronTokenizer.from_pretrained(
+ tokenizer_path="/path/to/tokenizer/model.json",
+ metadata_path={"library": "tiktoken"},
+ pattern="v2",
+ num_special_tokens=1000,
+)
+```
+
+### Null Tokenizer
+
+The Null tokenizer is a lightweight, zero-I/O tokenizer that requires no model files.
+It is useful in three scenarios:
+
+1. **Performance benchmarking** with `--mock-data` where real tokenization is unnecessary.
+2. **Testing** in functional tests and CI pipelines where tokenizer model files may not
+ be available. The Null tokenizer removes the dependency on external files, making
+ tests self-contained and portable.
+3. **Pretraining with pretokenized data** where all data is already tokenized into
+ `.bin`/`.idx` files. In this case the tokenizer is only needed for metadata
+ (`vocab_size`, `eod`, `pad`) — not for actual tokenization. Using the Null tokenizer
+ avoids redundant filesystem access at scale, which is particularly beneficial on
+ shared filesystems like Lustre where thousands of ranks would otherwise all load the
+ same tokenizer files.
+
+Properties derived from `--vocab-size N`:
+- `vocab_size` = `N` (the exact value passed)
+- `eod` = `N - 1` (last token in the vocabulary)
+- `pad` = `0`
+
+```python
+tokenizer = MegatronTokenizer.from_pretrained(
+ metadata_path={"library": "null-text"},
+ vocab_size=131072,
+)
+```
+
+## Integration with Megatron-LM
+
+### Using with Training Scripts
+
+The tokenizer system works with Megatron-LM training scripts:
+
+```bash
+# Null tokenizer for benchmarking with mock data
+torchrun --nproc_per_node=8 pretrain_gpt.py \
+ --tokenizer-type NullTokenizer \
+ --vocab-size 131072 \
+ --mock-data \
+ ...
+```
+
+```bash
+# Null tokenizer for pretraining with pretokenized data (no tokenizer files needed)
+torchrun --nproc_per_node=8 pretrain_gpt.py \
+ --tokenizer-type NullTokenizer \
+ --vocab-size 128256 \
+ --data-path /path/to/pretokenized_data \
+ ...
+```
+
+```bash
+# Hugging Face tokenizer with metadata
+torchrun --nproc_per_node=8 pretrain_gpt.py \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model meta-llama/Meta-Llama-3-8B \
+ --tokenizer-metadata /path/to/metadata.json \
+ ...
+```
+
+### Auto-Generated Metadata
+
+If `--tokenizer-metadata` is not specified, a default metadata file is generated automatically based on the tokenizer type.
+
+## Supported Tokenizer Libraries
+
+The following table lists supported tokenizer backends:
+
+| Library | Description | Use Case |
+|---------|-------------|----------|
+| **Hugging Face** | Transformers tokenizers | Most modern LLMs, such as LLaMA and Mistral |
+| **SentencePiece** | Google's tokenizer | GPT-style models, custom vocabularies |
+| **TikToken** | OpenAI's tokenizer | GPT-3.5/GPT-4 style tokenization |
+| **Megatron** | Built-in tokenizers | Legacy GPT-2 BPE |
+| **Null** | Zero-I/O tokenizer | Benchmarking, pretokenized data |
+
+## Common Tokenizer Types
+
+### LLaMA / Mistral
+
+```python
+MegatronTokenizer.write_metadata(
+ tokenizer_path="/path/to/llama/tokenizer.model",
+ tokenizer_library="sentencepiece",
+)
+```
+
+### GPT-2
+
+```python
+MegatronTokenizer.write_metadata(
+ tokenizer_path="GPT2BPETokenizer",
+ tokenizer_library="megatron",
+ vocab_file="/path/to/gpt2-vocab.json",
+ merge_file="/path/to/gpt2-merges.txt",
+)
+```
+
+## Recommendations
+
+1. **Save metadata** - Create metadata once, then reuse across training runs
+2. **Prefer Hugging Face tokenizers** - When the model ships one, it reduces integration work
+3. **Test tokenization** - Verify encode and decode before long training jobs
+4. **Version control metadata** - Track `tokenizer_metadata.json` with experiment configs
+5. **Share tokenizer directories** - Ship model files and metadata together for reproducibility
+
+## Next Steps
+
+- **Prepare data**: Refer to [Data Preparation](../data-preparation.md) for preprocessing with tokenizers
+- **Train models**: Refer to [Training Examples](../training-examples.md)
+- **Supported models**: Refer to [Language Models](../../models/llms.md) for model-specific tokenizers
diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md
new file mode 100644
index 00000000000..2a7ee2eeab9
--- /dev/null
+++ b/docs/user-guide/index.md
@@ -0,0 +1,26 @@
+---
+orphan: true
+---
+
+
+
+# User Guide
+
+Guides for using Megatron Core and Megatron-LM.
+
+```{toctree}
+:maxdepth: 2
+
+msc_integration
+data-preparation
+training-examples
+parallelism-guide
+features/index
+```
diff --git a/docs/user-guide/msc_integration.md b/docs/user-guide/msc_integration.md
new file mode 100644
index 00000000000..a197f25afc1
--- /dev/null
+++ b/docs/user-guide/msc_integration.md
@@ -0,0 +1,12 @@
+
+
+```{include} ../../megatron/core/MSC_Integration.md
+```
+
diff --git a/docs/user-guide/parallelism-guide.md b/docs/user-guide/parallelism-guide.md
new file mode 100644
index 00000000000..2540ca0a827
--- /dev/null
+++ b/docs/user-guide/parallelism-guide.md
@@ -0,0 +1,250 @@
+
+
+# Parallelism Strategies Guide
+
+Megatron Core supports multiple parallelism strategies that can be combined to efficiently train models from billions to trillions of parameters across thousands of GPUs.
+
+## Overview
+
+The following table summarizes supported parallelism strategies.
+
+| Strategy | Parallelism Objective | Best For |
+|----------|---------------------|----------|
+| **Data Parallelism (DP)** | Batch Dimension | Data Scalability, Standard Training |
+| **Tensor Parallelism (TP)** | Individual Layers | Large Layers & Activation, GPU Memory Constraints |
+| **Pipeline Parallelism (PP)** | Model Depth | Very Deep Models |
+| **Context Parallelism (CP)** | Sequence Length | Long Sequences (8K+ Tokens) |
+| **Expert Parallelism (EP)** | MoE Experts | Mixture-of-Experts Models |
+| **Fully-Sharded Data Parallelism (Megatron-FSDP)** | Model State | Extremely Large Models & DP Interchangeability |
+
+## Data Parallelism (DP)
+
+### Standard Distributed Data Parallel (DDP)
+
+Replicate the model across GPUs and split the batch.
+
+```bash
+torchrun --nproc_per_node=8 pretrain_gpt.py \
+ --data-parallel-sharding-strategy no_shard
+```
+
+Each GPU has a full copy of the model and processes a portion of the batch.
+
+### Megatron Fully-Sharded Data Parallel (Megatron-FSDP)
+
+Shard model parameters, gradients, and optimizer states across GPUs to reduce memory utilization.
+
+```
+--use-megatron-fsdp
+--data-parallel-sharding-strategy optim_grads_params
+--ckpt-format fsdp_dtensor
+--init-model-with-meta-device
+```
+
+**Sharding Strategies**
+
+`--data-parallel-sharding-strategy` supports the following options:
+
+- `optim` - Shard optimizer states only (ZeRO-1)
+- `optim_grads` - Shard gradients + optimizer (ZeRO-2)
+- `optim_grads_params` - Shard parameters + gradients + optimizer (ZeRO-3)
+
+If `--num-distributed-optimizer-instances` is > 1, then hierarchical data parallelism is enabled.
+
+`--outer-dp-sharding-strategy` supports the following options:
+
+- `no_shard` (**Hybrid-Sharded Data Parallelism**) - Replicate the model state across outer data parallel ranks.
+- `optim` (**Hybrid-FSDP**) - Shard the optimizer state across the outer data parallel ranks.
+ - Requires `--data-parallel-sharding-strategy optim_grads_params`.
+
+**When to Use**
+
+- Large models with large or fused compute kernels to hide communications under.
+- Integrated with TP, CP, EP, and easily composable with heterogeneous parallelisms.
+- With SM-reducing optimizations from NCCL and activation offloading from TransformerEngine.
+- Using `fully_shard` without depending on Megatron-LM.
+
+## Tensor Parallelism (TP)
+
+Split individual model layers across GPUs. Recommended for large hidden dimensions.
+
+```bash
+--tensor-model-parallel-size 4 # 4-way tensor parallelism
+--sequence-parallel # Enable sequence parallelism (recommended)
+```
+
+**When to Use**
+
+- Model layers do not fit on a single GPU
+- Large hidden dimensions (4096+)
+- Usually combined with DP and PP
+
+## Pipeline Parallelism (PP)
+
+Split model layers across GPUs vertically (by depth).
+
+```bash
+--pipeline-model-parallel-size 8 # 8 pipeline stages
+--num-layers-per-virtual-pipeline-stage 4 # Virtual pipeline for load balancing
+```
+
+**When to Use**
+
+- Very deep models (50+ layers)
+- Combine with TP for large models
+- Helps distribute memory across GPUs
+
+## Context Parallelism (CP)
+
+Split long sequences across GPUs for efficient long-context training.
+
+```bash
+--context-parallel-size 2 # 2-way context parallelism
+--cp-comm-type p2p # Communication type
+```
+
+**When to Use**
+
+- Long sequences (8K+ tokens)
+- Reduces activation memory
+- Can combine with TP, PP, DP
+
+Refer to [Context Parallelism Deep Dive](features/context_parallel.md) for a detailed guide with performance analysis.
+
+## Expert Parallelism (EP)
+
+Distribute experts across GPUs in Mixture-of-Experts models.
+
+```bash
+--expert-model-parallel-size 8 # 8-way expert parallelism
+--num-experts 64 # 64 experts per MoE layer
+--moe-grouped-gemm # Optimize expert computation
+```
+
+**Important:** When combining EP with TP, you **must enable Sequence Parallelism**:
+
+```bash
+--tensor-model-parallel-size 4
+--expert-model-parallel-size 8
+--sequence-parallel # Required when using TP + EP
+```
+
+## Parallelism Selection Guide
+
+For a list of supported configurations, refer to [Megatron Bridge Supported Models](https://github.com/NVIDIA-NeMo/Megatron-Bridge#supported-models).
+
+### Language Models
+
+Recommended language model configurations:
+
+| Model | Size | GPUs | TP | PP | CP | EP | Configuration Notes |
+|-------|------|------|----|----|----|----|---------------------|
+| **LLaMA-3** | 8B | 8 | 1 | 1 | 2 | 1 | CP=2 for long context (8K seqlen) |
+| **LLaMA-3** | 70B | 64 | 4 | 4 | 2 | 1 | Balanced TP+PP for 70B scale |
+| **LLaMA-3.1** | 405B | 1024 | 8 | 8 | 2 | 1 | 3D parallelism (TP+PP+CP) |
+| **GPT-3** | 175B | 128-512 | 4 | 8 | 1 | 1 | Standard large model config |
+
+### Mixture-of-Experts Models
+
+Recommended mixture-of-experts configurations:
+
+| Model | Size | GPUs | TP | PP | CP | EP | Configuration Notes |
+|-------|------|------|----|----|----|----|---------------------|
+| **Mixtral** | 8x7B | 64 | 1 | 4 | 1 | 8 | EP=8 for 8 experts |
+| **Mixtral** | 8x22B | 256 | 4 | 4 | 1 | 8 | TP+PP+EP for large MoE |
+| **DeepSeek-V3** | 671B | 1024 | 2 | 16 | 1 | 64 | Massive MoE with 256 experts |
+
+## Combining Strategies
+
+### Total GPU Count
+
+The total number of GPUs is calculated as:
+
+```
+Total GPUs = TP × PP × CP × EP × DP
+```
+
+### Example: LLaMA-3 70B on 64 GPUs
+
+```bash
+# TP=4, PP=4, CP=2, DP=2 => 4 × 4 × 2 × 2 = 64 GPUs
+torchrun --nproc_per_node=8 pretrain_gpt.py \
+ --tensor-model-parallel-size 4 \
+ --pipeline-model-parallel-size 4 \
+ --context-parallel-size 2 \
+ --num-layers 80 \
+ --hidden-size 8192 \
+ --num-attention-heads 64 \
+ --seq-length 8192 \
+ --micro-batch-size 1 \
+ --global-batch-size 512 \
+ --bf16
+```
+
+## Performance Optimizations
+
+### Communication Overlap
+
+Enable overlapping of communication with computation:
+
+```bash
+--overlap-grad-reduce # Overlap gradient reduction with backward pass
+--overlap-param-gather # Overlap parameter gathering with forward pass
+--tp-comm-overlap # Overlap TP communication
+```
+
+### Distributed Optimizer
+
+Recommended for all multi-GPU training:
+
+```bash
+--use-distributed-optimizer
+```
+
+**Benefits**
+
+- Faster checkpointing
+- Reduced memory when combined with FSDP
+- Better performance at scale
+
+### Sequence Parallelism
+
+Always enable when using TP:
+
+```bash
+--sequence-parallel
+```
+
+Reduces activation memory by sharding sequence dimension in LayerNorm and Dropout.
+
+## Choosing the Right Strategy
+
+### Start Simple
+1. Begin with **Data Parallelism** (DP) only.
+2. Add **Tensor Parallelism** (TP) if the model does not fit.
+3. Add **Pipeline Parallelism** (PP) for very large models.
+4. Add **Context Parallelism** (CP) for long sequences.
+
+### Memory Constraints
+- Use **FSDP** to split model state per GPU.
+- Use **TP** to split large layers.
+- Use **PP** to split model depth.
+- Enable **activation checkpointing or offloading** for extreme cases.
+
+### Communication Bottlenecks
+- Reduce **TP** degree (increases memory per GPU).
+- Increase **PP** degree (may reduce efficiency).
+- Use **CP** instead of larger TP for long sequences.
+
+## Next Steps
+
+- **API Reference**: Refer to [Tensor Parallel](../api-guide/core/tensor_parallel.md) and [Pipeline Parallel](../api-guide/core/pipeline_parallel.md) in the API documentation
+- **Advanced Features**: Refer to [Megatron-FSDP](features/megatron_fsdp.md), [MoE](features/moe.md), and [Distributed Optimizer](features/dist_optimizer.md)
+- **Performance Tuning**: Refer to the [NVIDIA NeMo Performance Guide](https://docs.nvidia.com/nemo-framework/user-guide/latest/performance/performance-guide.html)
diff --git a/docs/user-guide/training-examples.md b/docs/user-guide/training-examples.md
new file mode 100644
index 00000000000..ca8b182adc7
--- /dev/null
+++ b/docs/user-guide/training-examples.md
@@ -0,0 +1,159 @@
+
+
+# Training Examples
+
+Get started with Megatron Core training using these practical examples.
+
+## Basic Training Example
+
+Use the basic training loop with mock data to get started:
+
+```bash
+# Distributed training on 2 GPUs with mock data
+torchrun --nproc_per_node=2 examples/run_simple_mcore_train_loop.py
+```
+
+This example:
+
+- Runs on two GPUs
+- Uses generated mock data (no data preparation needed)
+- Demonstrates basic distributed training setup
+- Provides a quick way to verify your installation
+
+## LLaMA-3 Training Examples
+
+### LLaMA-3 8B with FP8
+
+Train the LLaMA-3 8B model with FP8 mixed precision on eight GPUs:
+
+```bash
+./examples/llama/train_llama3_8b_h100_fp8.sh
+```
+
+**Configuration**
+
+- Eight GPUs
+- FP8 mixed precision (requires Hopper/Ada/Blackwell GPUs)
+- Mock data for quick testing
+
+### Custom LLaMA Training
+
+For training with your own data:
+
+```bash
+torchrun --nproc_per_node=8 pretrain_gpt.py \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --num-attention-heads 32 \
+ --seq-length 2048 \
+ --max-position-embeddings 2048 \
+ --micro-batch-size 4 \
+ --global-batch-size 32 \
+ --train-iters 100000 \
+ --lr 3.0e-4 \
+ --min-lr 3.0e-5 \
+ --lr-decay-style cosine \
+ --lr-warmup-iters 2000 \
+ --weight-decay 0.1 \
+ --clip-grad 1.0 \
+ --bf16 \
+ --data-path /path/to/your/preprocessed_data \
+ --split 949,50,1 \
+ --save /path/to/checkpoints \
+ --load /path/to/checkpoints \
+ --log-interval 10 \
+ --save-interval 1000 \
+ --eval-interval 1000
+```
+
+## GPT-3 Training Example
+
+Train a GPT-3 style model:
+
+```bash
+torchrun --nproc_per_node=8 pretrain_gpt.py \
+ --tensor-model-parallel-size 2 \
+ --pipeline-model-parallel-size 2 \
+ --num-layers 24 \
+ --hidden-size 2048 \
+ --num-attention-heads 16 \
+ --seq-length 1024 \
+ --max-position-embeddings 1024 \
+ --micro-batch-size 2 \
+ --global-batch-size 16 \
+ --train-iters 100000 \
+ --lr 1.5e-4 \
+ --min-lr 1.0e-5 \
+ --lr-decay-style cosine \
+ --lr-warmup-iters 1000 \
+ --weight-decay 0.1 \
+ --clip-grad 1.0 \
+ --fp16 \
+ --data-path /path/to/preprocessed_data \
+ --split 949,50,1 \
+ --save /path/to/checkpoints \
+ --load /path/to/checkpoints
+```
+
+## Key Training Arguments
+
+The following tables group common training arguments by category.
+
+### Model Architecture
+
+| Argument | Description |
+|----------|-------------|
+| `--num-layers` | Number of transformer layers |
+| `--hidden-size` | Hidden dimension size |
+| `--num-attention-heads` | Number of attention heads |
+| `--seq-length` | Sequence length for training |
+
+### Training Configuration
+
+| Argument | Description |
+|----------|-------------|
+| `--micro-batch-size` | Batch size per GPU |
+| `--global-batch-size` | Total batch size across all GPUs |
+| `--train-iters` | Number of training iterations |
+
+### Learning Rate
+
+| Argument | Description |
+|----------|-------------|
+| `--lr` | Peak learning rate |
+| `--min-lr` | Minimum learning rate |
+| `--lr-decay-style` | LR schedule (cosine, linear, constant) |
+| `--lr-warmup-iters` | Warmup iterations |
+
+### Mixed Precision
+
+| Argument | Description |
+|----------|-------------|
+| `--fp16` | FP16 mixed precision |
+| `--bf16` | BF16 mixed precision (recommended) |
+| `--fp8-hybrid` | FP8 mixed precision (Hopper/Ada/Blackwell) |
+
+### Data and Checkpointing
+
+| Argument | Description |
+|----------|-------------|
+| `--data-path` | Path to preprocessed data |
+| `--split` | Train/validation/test split (for example, 949,50,1) |
+| `--save` | Checkpoint save directory |
+| `--load` | Checkpoint load directory |
+| `--save-interval` | Save checkpoint every N iterations |
+
+## Next Steps
+
+- **Optimize Performance**: Refer to [Advanced Features](features/index.md) for FSDP, the distributed optimizer, and other optimizations
+- **Scale Up**: Refer to [Parallelism Strategies](parallelism-guide.md) to train larger models across more GPUs
+- **Prepare Data**: Follow the [Data Preparation](data-preparation.md) guide to process your own datasets
diff --git a/docs/versions1.json b/docs/versions1.json
new file mode 100644
index 00000000000..b0bc489fa71
--- /dev/null
+++ b/docs/versions1.json
@@ -0,0 +1,22 @@
+[
+ {
+ "name": "nightly",
+ "version": "nightly",
+ "url": "https://docs.nvidia.com/megatron-core/developer-guide/nightly/"
+ },
+ {
+ "name": "0.17.0 (latest)",
+ "version": "0.17.0",
+ "url": "https://docs.nvidia.com/megatron-core/developer-guide/latest/"
+ },
+ {
+ "name": "0.16.0",
+ "version": "0.16.0",
+ "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.16.0/"
+ },
+ {
+ "name": "0.15.0",
+ "version": "0.15.0",
+ "url": "https://docs.nvidia.com/megatron-core/developer-guide/0.15.0/"
+ }
+]
diff --git a/examples/__init__.py b/examples/__init__.py
new file mode 100644
index 00000000000..0519ecba6ea
--- /dev/null
+++ b/examples/__init__.py
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/examples/detxoify_lm/README.md b/examples/academic_paper_scripts/detxoify_lm/README.md
similarity index 100%
rename from examples/detxoify_lm/README.md
rename to examples/academic_paper_scripts/detxoify_lm/README.md
diff --git a/examples/detxoify_lm/annotations/filter-selfgeneration.py b/examples/academic_paper_scripts/detxoify_lm/annotations/filter-selfgeneration.py
similarity index 100%
rename from examples/detxoify_lm/annotations/filter-selfgeneration.py
rename to examples/academic_paper_scripts/detxoify_lm/annotations/filter-selfgeneration.py
diff --git a/examples/detxoify_lm/annotations/perspective_api_annotate.py b/examples/academic_paper_scripts/detxoify_lm/annotations/perspective_api_annotate.py
similarity index 98%
rename from examples/detxoify_lm/annotations/perspective_api_annotate.py
rename to examples/academic_paper_scripts/detxoify_lm/annotations/perspective_api_annotate.py
index fd82c2a2ae9..9736db099a4 100644
--- a/examples/detxoify_lm/annotations/perspective_api_annotate.py
+++ b/examples/academic_paper_scripts/detxoify_lm/annotations/perspective_api_annotate.py
@@ -107,7 +107,7 @@ def get_score(line):
except UnicodeDecodeError:
try:
decoded_text = encoded_text[:20476].decode('utf8')
- except:
+ except Exception:
print("Error occurred")
data['score'] = None
return json.dumps(data)
@@ -138,7 +138,7 @@ def get_scores(lines):
except UnicodeDecodeError:
try:
decoded_text = encoded_text[:20476].decode('utf8')
- except:
+ except Exception:
print("Error occurred")
data['score'] = None
all_data.append(json.dumps(data))
diff --git a/examples/detxoify_lm/annotations/preprocess.sh b/examples/academic_paper_scripts/detxoify_lm/annotations/preprocess.sh
similarity index 100%
rename from examples/detxoify_lm/annotations/preprocess.sh
rename to examples/academic_paper_scripts/detxoify_lm/annotations/preprocess.sh
diff --git a/examples/detxoify_lm/finetune_gpt_distributed-1.3b.sh b/examples/academic_paper_scripts/detxoify_lm/finetune_gpt_distributed-1.3b.sh
similarity index 98%
rename from examples/detxoify_lm/finetune_gpt_distributed-1.3b.sh
rename to examples/academic_paper_scripts/detxoify_lm/finetune_gpt_distributed-1.3b.sh
index 62a36c0b79e..a212fbdf3f6 100755
--- a/examples/detxoify_lm/finetune_gpt_distributed-1.3b.sh
+++ b/examples/academic_paper_scripts/detxoify_lm/finetune_gpt_distributed-1.3b.sh
@@ -43,7 +43,6 @@ python -m torch.distributed.run $DISTRIBUTED_ARGS \
--data-path2 ${DATA_BLEND} \
--vocab-file $VOCAB_FILE \
--merge-file $MERGE_FILE \
- --data-impl mmap \
--split 100,0,0 \
--distributed-backend nccl \
--lr-decay-style constant \
diff --git a/examples/detxoify_lm/generate-1.3b.sh b/examples/academic_paper_scripts/detxoify_lm/generate-1.3b.sh
similarity index 100%
rename from examples/detxoify_lm/generate-1.3b.sh
rename to examples/academic_paper_scripts/detxoify_lm/generate-1.3b.sh
diff --git a/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py
new file mode 100644
index 00000000000..2a2b1d63a21
--- /dev/null
+++ b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py
@@ -0,0 +1,249 @@
+# coding=utf-8
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+
+
+"""Sample Generate GPT"""
+import json
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
+ os.path.pardir, os.path.pardir)))
+import torch
+from megatron.training import get_args
+from megatron.training import get_tokenizer
+from megatron.training import print_rank_0
+from megatron.training.checkpointing import load_checkpoint
+from megatron.core import mpu
+from megatron.training.arguments import parse_and_validate_args
+from megatron.training.initialize import initialize_megatron
+from megatron.training import get_model
+from megatron.inference.text_generation import generate_and_post_process
+from megatron.training.arguments import core_transformer_config_from_args
+from megatron.core.models.gpt import GPTModel
+from typing import Union
+from megatron.core.transformer.spec_utils import import_module
+from megatron.training.arguments import core_transformer_config_from_args
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_local_spec
+
+def model_provider(pre_process=True, post_process=True) -> GPTModel:
+ """Builds the model.
+
+ Args:
+ pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True.
+ post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True.
+
+
+ Returns:
+ GPTModel: The returned model
+ """
+ args = get_args()
+
+ print_rank_0('building GPT model ...')
+ config = core_transformer_config_from_args(args)
+
+ if args.spec is None:
+ if args.transformer_impl == 'local':
+ transformer_layer_spec = get_gpt_layer_local_spec(
+ num_experts=args.num_experts,
+ moe_grouped_gemm=args.moe_grouped_gemm
+ )
+ elif args.transformer_impl == 'transformer_engine':
+ transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+ num_experts=args.num_experts,
+ moe_grouped_gemm=args.moe_grouped_gemm
+ )
+ else:
+ raise ValueError(f"Invalid transformer_impl {args.transformer_impl}")
+ elif args.spec[0] == 'local':
+ transformer_layer_spec = get_gpt_layer_local_spec(
+ num_experts=args.num_experts,
+ moe_grouped_gemm=args.moe_grouped_gemm
+ )
+ else:
+ transformer_layer_spec = import_module(args.spec)
+
+ model = GPTModel(
+ config=config,
+ transformer_layer_spec=transformer_layer_spec,
+ vocab_size=args.padded_vocab_size,
+ max_sequence_length=args.max_position_embeddings,
+ pre_process=pre_process,
+ post_process=post_process,
+ fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+ parallel_output=False,
+ share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+ position_embedding_type=args.position_embedding_type,
+ rotary_percent=args.rotary_percent
+ )
+
+ return model
+
+def add_text_generate_args(parser):
+ """Text generation arguments."""
+ group = parser.add_argument_group(title='text generation')
+
+ group.add_argument("--temperature", type=float, default=1.0,
+ help='Sampling temperature.')
+ group.add_argument("--greedy", action='store_true', default=False,
+ help='Use greedy sampling.')
+ group.add_argument("--top_p", type=float, default=0.0,
+ help='Top p sampling.')
+ group.add_argument("--top_k", type=int, default=0,
+ help='Top k sampling.')
+ group.add_argument("--out-seq-length", type=int, default=1024,
+ help='Size of the output generated text.')
+ group.add_argument("--sample-input-file", type=str, default=None,
+ help='Get input from file instead of interactive mode, '
+ 'each line is an input.')
+ group.add_argument("--sample-output-file", type=str, default=None,
+ help='Output file got from --sample-input-file')
+ group.add_argument("--num-samples", type=int, default=0,
+ help='Number of samples to generate unconditionally, '
+ 'defaults to 0 and interactive conditional sampling')
+ group.add_argument("--genfile", type=str,
+ help='Output file when generating unconditionally')
+ return parser
+
+def generate_samples_unconditional(model):
+ args = get_args()
+
+ if torch.distributed.get_rank() == 0:
+ cnt = 0
+ num_samples = args.num_samples
+ from tqdm import tqdm
+ pbar = tqdm(total=num_samples)
+
+ while True:
+ if torch.distributed.get_rank() == 0:
+ sentences = [''] * args.global_batch_size
+ print("global batch size", args.global_batch_size)
+ max_len = args.out_seq_length
+ resp_sentences, resp_sentences_seg, output_logits, \
+ tokens = generate_and_post_process(model, prompts=sentences,
+ tokens_to_generate=max_len,
+ return_output_log_probs=False,
+ top_k_sampling=args.top_k,
+ top_p_sampling=args.top_p,
+ add_BOS=True,
+ temperature=1.0)
+ for prompt, generation, token in zip(sentences, resp_sentences, tokens):
+ datum = {'text': generation[len(prompt):], 'all_text': generation, 'prompt': prompt, 'id': cnt}
+ yield datum
+ cnt += 1
+ pbar.update()
+ if cnt >= num_samples:
+ break
+
+ if cnt >= num_samples:
+ pbar.close()
+ break
+ else:
+ generate_and_post_process(model)
+
+
+def generate_samples_conditional(model):
+ args = get_args()
+
+ if torch.distributed.get_rank() == 0:
+ num_samples = args.num_samples
+ cnt = 0
+ from tqdm import tqdm
+ pbar = tqdm(total=num_samples)
+
+ fname = open(args.sample_input_file, "r")
+ lines = fname.readlines()
+ all_raw_text = [json.loads(line)['prompt']['text'] for line in lines]
+ input_count = len(all_raw_text)
+ input_pos = 0
+
+ while True:
+ torch.distributed.barrier()
+ if torch.distributed.get_rank() == 0:
+ sentences = []
+ print("global batch size", args.global_batch_size)
+ for _ in range(args.global_batch_size):
+ if input_pos >= input_count:
+ print(f"input pos: {input_pos}, input count: {input_count}")
+ raw_text = "EMPTY TEXT"
+ else:
+ raw_text = all_raw_text[input_pos]
+ input_pos += 1
+ sentences.append(raw_text)
+
+ max_len = args.out_seq_length
+ resp_sentences, resp_sentences_seg, output_logits, \
+ tokens = generate_and_post_process(model, prompts=sentences,
+ tokens_to_generate=max_len,
+ return_output_log_probs=False,
+ top_k_sampling=args.top_k,
+ top_p_sampling=args.top_p,
+ add_BOS=False,
+ temperature=1.0)
+ for prompt, generation, token in zip(sentences, resp_sentences, tokens):
+ datum = {'text': generation[len(prompt):], 'all_text': generation, 'prompt': prompt, 'id': cnt}
+ yield datum
+ cnt += 1
+ pbar.update()
+ if cnt >= num_samples:
+ break
+
+ if cnt >= num_samples:
+ pbar.close()
+ break
+ else:
+ generate_and_post_process(model)
+
+
+def generate_and_write_samples_unconditional(model):
+ args = get_args()
+ assert args.genfile is not None
+ with open(args.genfile, 'w') as f:
+ for datum in generate_samples_unconditional(model):
+ if torch.distributed.get_rank() == 0:
+ f.write(json.dumps(datum) + '\n')
+
+
+def generate_and_write_samples_conditional(model):
+ args = get_args()
+ if args.sample_output_file is None:
+ sample_output_file = args.sample_input_file + ".out"
+ print('`sample-output-file` not specified, setting '
+ 'it to {}'.format(sample_output_file))
+ else:
+ sample_output_file = args.sample_output_file
+ with open(sample_output_file, 'w') as f:
+ for datum in generate_samples_conditional(model):
+ if torch.distributed.get_rank() == 0:
+ f.write(json.dumps(datum) + '\n')
+
+
+def main():
+ """Main program."""
+
+ parse_and_validate_args(extra_args_provider=add_text_generate_args,
+ args_defaults={'tokenizer_type': 'GPT2BPETokenizer',
+ 'no_load_rng': True,
+ 'no_load_optim': True,
+ 'seq_length': 2048})
+ initialize_megatron()
+
+ # Set up model and load checkpoint
+ model = get_model(model_provider, wrap_with_ddp=False)
+
+ args = get_args()
+
+ if args.load is not None:
+ _ = load_checkpoint(model, None, None)
+ model = model[0]
+
+ # Generate samples.
+ if args.sample_input_file != None:
+ print(f"{args.sample_input_file}")
+ generate_and_write_samples_conditional(model)
+ else:
+ generate_and_write_samples_unconditional(model)
+
+
+if __name__ == "__main__":
+
+ main()
diff --git a/examples/detxoify_lm/perspective_api.py b/examples/academic_paper_scripts/detxoify_lm/perspective_api.py
similarity index 100%
rename from examples/detxoify_lm/perspective_api.py
rename to examples/academic_paper_scripts/detxoify_lm/perspective_api.py
diff --git a/examples/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh b/examples/academic_paper_scripts/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh
similarity index 100%
rename from examples/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh
rename to examples/academic_paper_scripts/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh
diff --git a/examples/msdp/README.md b/examples/academic_paper_scripts/msdp/README.md
similarity index 100%
rename from examples/msdp/README.md
rename to examples/academic_paper_scripts/msdp/README.md
diff --git a/examples/msdp/data_processing.sh b/examples/academic_paper_scripts/msdp/data_processing.sh
similarity index 100%
rename from examples/msdp/data_processing.sh
rename to examples/academic_paper_scripts/msdp/data_processing.sh
diff --git a/examples/msdp/eval_knwl_generation.sh b/examples/academic_paper_scripts/msdp/eval_knwl_generation.sh
similarity index 100%
rename from examples/msdp/eval_knwl_generation.sh
rename to examples/academic_paper_scripts/msdp/eval_knwl_generation.sh
diff --git a/examples/msdp/eval_resp_generation.sh b/examples/academic_paper_scripts/msdp/eval_resp_generation.sh
similarity index 100%
rename from examples/msdp/eval_resp_generation.sh
rename to examples/academic_paper_scripts/msdp/eval_resp_generation.sh
diff --git a/examples/msdp/prep_resp_gen.sh b/examples/academic_paper_scripts/msdp/prep_resp_gen.sh
similarity index 100%
rename from examples/msdp/prep_resp_gen.sh
rename to examples/academic_paper_scripts/msdp/prep_resp_gen.sh
diff --git a/examples/msdp/prompt_knwl_gen.sh b/examples/academic_paper_scripts/msdp/prompt_knwl_gen.sh
similarity index 100%
rename from examples/msdp/prompt_knwl_gen.sh
rename to examples/academic_paper_scripts/msdp/prompt_knwl_gen.sh
diff --git a/examples/msdp/prompt_resp_gen.sh b/examples/academic_paper_scripts/msdp/prompt_resp_gen.sh
similarity index 100%
rename from examples/msdp/prompt_resp_gen.sh
rename to examples/academic_paper_scripts/msdp/prompt_resp_gen.sh
diff --git a/examples/sc21/CONFIG.sh b/examples/academic_paper_scripts/sc21/CONFIG.sh
similarity index 100%
rename from examples/sc21/CONFIG.sh
rename to examples/academic_paper_scripts/sc21/CONFIG.sh
diff --git a/examples/academic_paper_scripts/sc21/README.md b/examples/academic_paper_scripts/sc21/README.md
new file mode 100644
index 00000000000..ec922d153d6
--- /dev/null
+++ b/examples/academic_paper_scripts/sc21/README.md
@@ -0,0 +1,50 @@
+# Reproducing Figures in SC21 Paper
+
+
+This directory contains some of the scripts that were used to produce the
+results in the [Megatron paper](https://arxiv.org/pdf/2104.04473.pdf) that is
+to appear at [SuperComputing 2021](https://sc21.supercomputing.org/). These
+scripts use [Slurm](https://slurm.schedmd.com/documentation.html) with the
+[pyxis plugin](https://github.com/NVIDIA/pyxis), but can be modified for other
+schedulers as well.
+
+
+## Git commit
+
+To replicate these results use Megatron-LM commit: 6985e58938d40ad91ac07b0fddcfad8132e1447e
+
+
+## Setup
+
+All the cluster-dependent variables are in [`CONFIG.sh`](./CONFIG.sh). Please
+update the unspecified values (in angle brackets `<...>`) before launching any
+scripts.
+
+
+
+## Scripts
+
+Below is a list of scripts that can be used to reproduce various figures in our
+[paper](https://arxiv.org/pdf/2104.04473.pdf):
+
+* [run_table_1.sh](./run_table_1.sh): Table 1 showing weak-scaling throughput
+for GPT models ranging from 1 billion to 1 trillion parameters.
+* [run_figure_11.sh](./run_figure_11.sh): Figure 11 showing the weak-scaling
+performance of pipeline parallelism.
+* [run_figure_12.sh](./run_figure_12.sh): Figure 12 showing the effect of
+the interleaved schedule on a 175B GPT model.
+* [run_figure_13.sh](./run_figure_13.sh): Figure 13 showing the effect of
+different degrees of pipeline and tensor model parallelism on a model with
+162.2 billion parameters.
+* [run_figure_14.sh](./run_figure_14.sh): Figure 14 showing the effect of
+different degrees of data and pipeline model parallelism on a model with
+5.9 billion parameters.
+* [run_figure_15.sh](./run_figure_15.sh): Figure 15 showing the effect of
+different degrees of data and tensor model parallelism on a model with
+5.9 billion parameters.
+* [run_figure_16.sh](./run_figure_16.sh): Figure 16 showing the effect of
+microbatch size.
+* [run_figure_17.sh](./run_figure_17.sh): Figure 17 showing the effect of
+activation recomputation.
+* [run_figure_18.sh](./run_figure_18.sh): Figure 18 showing the effect of
+the scatter-gather communication optimization.
diff --git a/examples/sc21/SBATCH.sh b/examples/academic_paper_scripts/sc21/SBATCH.sh
similarity index 100%
rename from examples/sc21/SBATCH.sh
rename to examples/academic_paper_scripts/sc21/SBATCH.sh
diff --git a/examples/sc21/SRUN.sh b/examples/academic_paper_scripts/sc21/SRUN.sh
similarity index 100%
rename from examples/sc21/SRUN.sh
rename to examples/academic_paper_scripts/sc21/SRUN.sh
diff --git a/examples/sc21/run_figure_11.sh b/examples/academic_paper_scripts/sc21/run_figure_11.sh
similarity index 100%
rename from examples/sc21/run_figure_11.sh
rename to examples/academic_paper_scripts/sc21/run_figure_11.sh
diff --git a/examples/sc21/run_figure_12.sh b/examples/academic_paper_scripts/sc21/run_figure_12.sh
similarity index 100%
rename from examples/sc21/run_figure_12.sh
rename to examples/academic_paper_scripts/sc21/run_figure_12.sh
diff --git a/examples/sc21/run_figure_13.sh b/examples/academic_paper_scripts/sc21/run_figure_13.sh
similarity index 100%
rename from examples/sc21/run_figure_13.sh
rename to examples/academic_paper_scripts/sc21/run_figure_13.sh
diff --git a/examples/sc21/run_figure_14.sh b/examples/academic_paper_scripts/sc21/run_figure_14.sh
similarity index 100%
rename from examples/sc21/run_figure_14.sh
rename to examples/academic_paper_scripts/sc21/run_figure_14.sh
diff --git a/examples/sc21/run_figure_15.sh b/examples/academic_paper_scripts/sc21/run_figure_15.sh
similarity index 100%
rename from examples/sc21/run_figure_15.sh
rename to examples/academic_paper_scripts/sc21/run_figure_15.sh
diff --git a/examples/sc21/run_figure_16.sh b/examples/academic_paper_scripts/sc21/run_figure_16.sh
similarity index 100%
rename from examples/sc21/run_figure_16.sh
rename to examples/academic_paper_scripts/sc21/run_figure_16.sh
diff --git a/examples/sc21/run_figure_17.sh b/examples/academic_paper_scripts/sc21/run_figure_17.sh
similarity index 100%
rename from examples/sc21/run_figure_17.sh
rename to examples/academic_paper_scripts/sc21/run_figure_17.sh
diff --git a/examples/academic_paper_scripts/sc21/run_figure_18.sh b/examples/academic_paper_scripts/sc21/run_figure_18.sh
new file mode 100755
index 00000000000..c9e254200a1
--- /dev/null
+++ b/examples/academic_paper_scripts/sc21/run_figure_18.sh
@@ -0,0 +1,41 @@
+#!/bin/bash
+
+# ================================
+# Choose the case to run.
+# ================================
+
+# Batch size (global batch size) options = [12, 24, 36, ..., 60].
+GBS=12
+
+
+
+MEGATRON_EXTRA_PARAMS="--activations-checkpoint-method uniform --num-layers-per-virtual-pipeline-stage 2 "
+
+
+# Other params.
+TP=8
+PP=12
+MBS=1
+NLS=96
+HS=12288
+NAH=96
+DDP=local
+NNODES=12
+
+
+# Name of the job.
+export JOB_NAME=results_figure_18_batch_size_${GBS}
+
+
+# Import the configs.
+. `pwd`/CONFIG.sh
+
+
+# Submit the job.
+. `pwd`/SBATCH.sh
+
+
+exit 0
+
+
+
diff --git a/examples/sc21/run_table_1.sh b/examples/academic_paper_scripts/sc21/run_table_1.sh
similarity index 100%
rename from examples/sc21/run_table_1.sh
rename to examples/academic_paper_scripts/sc21/run_table_1.sh
diff --git a/examples/bert/README.md b/examples/bert/README.md
new file mode 100644
index 00000000000..6c1fe95bf06
--- /dev/null
+++ b/examples/bert/README.md
@@ -0,0 +1,53 @@
+# BERT MODEL
+
+## Table of contents
+- [1. Training Setup](#1-training-setup)
+- [2. Configurations](#2-configurations)
+
+## 1. Training setup
+
+
+To run the model using a docker container run it as follows
+```
+PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:24.01-py3
+CHECKPOINT_PATH="" #
+TENSORBOARD_LOGS_PATH=""#
+VOCAB_FILE="" #//bert-vocab.txt
+DATA_PATH="" #_text_document
+
+docker run \
+ --gpus=all \
+ --ipc=host \
+ --workdir /workspace/megatron-lm \
+ -v /path/to/data:/path/to/data \
+ -v /path/to/megatron-lm:/workspace/megatron-lm \
+ megatron-lm nvcr.io/nvidia/pytorch:24.01-py3 \
+ bash examples/bert/train_bert_340m_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH $VOCAB_FILE $DATA_PATH "
+
+```
+NOTE: Depending on the environment you are running it the above command might like slightly different.
+
+
+## 2. Configurations
+
+The example in this folder shows you how to run 340m large model. There are other configs you could run as well
+
+### 4B
+```
+ --num-layers 48 \
+ --hidden-size 2560 \
+ --num-attention-heads 32 \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+
+```
+
+### 20B
+```
+ --num-layers 48 \
+ --hidden-size 6144 \
+ --num-attention-heads 96 \
+ --tensor-model-parallel-size 4 \
+ --pipeline-model-parallel-size 4 \
+
+```
\ No newline at end of file
diff --git a/examples/bert/train_bert_340m_distributed.sh b/examples/bert/train_bert_340m_distributed.sh
new file mode 100644
index 00000000000..f0d9c87c8bf
--- /dev/null
+++ b/examples/bert/train_bert_340m_distributed.sh
@@ -0,0 +1,79 @@
+#!/bin/bash
+
+# Runs the "340M" parameter model (Bert - Large)
+
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+GPUS_PER_NODE=8
+# Change for multinode config
+MASTER_ADDR=localhost
+MASTER_PORT=6000
+NUM_NODES=1
+NODE_RANK=0
+WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
+
+CHECKPOINT_PATH=$1 #
+TENSORBOARD_LOGS_PATH=$2 #
+VOCAB_FILE=$3 #/bert-vocab.json
+DATA_PATH=$4 #_text_document
+
+DISTRIBUTED_ARGS=(
+ --nproc_per_node $GPUS_PER_NODE
+ --nnodes $NUM_NODES
+ --master_addr $MASTER_ADDR
+ --master_port $MASTER_PORT
+)
+
+BERT_MODEL_ARGS=(
+ --num-layers 24
+ --hidden-size 1024
+ --num-attention-heads 16
+ --seq-length 512
+ --max-position-embeddings 512
+ --attention-backend auto # Can use (flash/fused/unfused/local)
+)
+
+TRAINING_ARGS=(
+ --micro-batch-size 4
+ --global-batch-size 32
+ --train-iters 1000000
+ --weight-decay 1e-2
+ --clip-grad 1.0
+ --fp16
+ --lr 0.0001
+ --lr-decay-iters 990000
+ --lr-decay-style linear
+ --min-lr 1.0e-5
+ --weight-decay 1e-2
+ --lr-warmup-fraction .01
+ --clip-grad 1.0
+)
+
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size 8
+ --pipeline-model-parallel-size 16
+)
+
+DATA_ARGS=(
+ --data-path $DATA_PATH
+ --vocab-file $VOCAB_FILE
+ --split 949,50,1
+)
+
+EVAL_AND_LOGGING_ARGS=(
+ --log-interval 100
+ --save-interval 10000
+ --eval-interval 1000
+ --save $CHECKPOINT_PATH
+ --load $CHECKPOINT_PATH
+ --eval-iters 10
+ --tensorboard-dir $TENSORBOARD_LOGS_PATH
+)
+
+torchrun ${DISTRIBUTED_ARGS[@]} pretrain_bert.py \
+ ${BERT_MODEL_ARGS[@]} \
+ ${TRAINING_ARGS[@]} \
+ ${MODEL_PARALLEL_ARGS[@]} \
+ ${DATA_ARGS[@]} \
+ ${EVAL_AND_LOGGING_ARGS[@]}
+
\ No newline at end of file
diff --git a/examples/detxoify_lm/finetune_gpt.py b/examples/detxoify_lm/finetune_gpt.py
deleted file mode 100644
index 70b781e0eee..00000000000
--- a/examples/detxoify_lm/finetune_gpt.py
+++ /dev/null
@@ -1,145 +0,0 @@
-# coding=utf-8
-# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
-
-
-"""Fine-tune GPT"""
-
-import torch
-from functools import partial
-import os
-import sys
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
- os.path.pardir, os.path.pardir)))
-from megatron import get_args
-from megatron import get_timers
-from megatron import get_tokenizer
-from megatron import print_rank_0
-from megatron.core import mpu
-from megatron.data.blendable_dataset import BlendableDataset
-from megatron.data.gpt_dataset import build_train_valid_test_datasets
-from megatron.model import GPTModel
-from megatron.core.enums import ModelType
-from megatron.training import pretrain
-from megatron.utils import get_ltor_masks_and_position_ids
-from megatron.utils import average_losses_across_data_parallel_group
-
-def model_provider(pre_process=True, post_process=True):
- """Build the model."""
-
- print_rank_0('building GPT model ...')
- model = GPTModel(
- num_tokentypes=0,
- parallel_output=True,
- pre_process=pre_process,
- post_process=post_process
- )
- return model
-
-
-def get_batch(data_iterator):
- """Generate a batch"""
- args = get_args()
- tokenizer = get_tokenizer()
-
- # Items and their type.
- keys = ['text']
- datatype = torch.int64
-
- # Broadcast data.
- if data_iterator is not None:
- data = next(data_iterator)
- else:
- data = None
- data_b = mpu.broadcast_data(keys, data, datatype)
-
- # Unpack.
- tokens_ = data_b['text'].long()
- labels = tokens_[:, 1:].contiguous()
- tokens = tokens_[:, :-1].contiguous()
-
- # Get the masks and postition ids.
- attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
- tokens,
- tokenizer.eod,
- args.reset_position_ids,
- args.reset_attention_mask,
- args.eod_mask_loss)
-
- return tokens, labels, loss_mask, attention_mask, position_ids
-
-def loss_func(loss_mask, output_tensor):
- losses = output_tensor.float()
- loss_mask = loss_mask.view(-1).float()
- loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
-
- # Reduce loss for logging.
- averaged_loss = average_losses_across_data_parallel_group([loss])
-
- return loss, {'lm loss': averaged_loss[0]}
-
-
-def forward_step(data_iterator, model):
- """Forward step."""
- args = get_args()
- timers = get_timers()
-
- # Get the batch.
- timers('batch-generator').start()
- tokens, labels, loss_mask, attention_mask, position_ids = get_batch(
- data_iterator)
- timers('batch-generator').stop()
-
- output_tensor = model(tokens, position_ids, attention_mask,
- labels=labels)
-
- return output_tensor, partial(loss_func, loss_mask)
-
-
-def train_valid_test_datasets_provider(train_val_test_num_samples):
- """Build train, valid, and test datasets."""
- args = get_args()
-
- print_rank_0('> building train, validation, and test datasets '
- 'for GPT ...')
- train_ds, valid_ds1, test_ds = build_train_valid_test_datasets(
- data_prefix=args.data_path,
- data_impl=args.data_impl,
- splits_string=args.split,
- train_valid_test_num_samples=train_val_test_num_samples,
- seq_length=args.seq_length,
- seed=args.seed,
- skip_warmup=(not args.mmap_warmup))
- print_rank_0("> finished creating finetuning GPT datasets ...")
-
- _, valid_ds, _ = build_train_valid_test_datasets(
- data_prefix=args.data_path2,
- data_impl="mmap",
- splits_string="98,2,0",
- train_valid_test_num_samples=train_val_test_num_samples,
- seq_length=2048,
- seed=1234,
- skip_warmup=(not args.mmap_warmup))
- print_rank_0("> finished creating pretrained GPT datasets ...")
-
- return train_ds, valid_ds, test_ds
-
-
-def add_validation_args(parser):
- """Text generation arguments."""
- group = parser.add_argument_group(title='validation set')
- group.add_argument('--data-path2', nargs='*', default=None,
- help='Path to the validation dataset. Accepted format:'
- '1) a single data path, 2) multiple datasets in the'
- 'form: dataset1-weight dataset1-path dataset2-weight '
- 'dataset2-path ...')
- group.add_argument('--eval-ppl', action='store_true', default=False)
- group.add_argument('--stored_params', type=dict, default=dict())
- return parser
-
-
-if __name__ == "__main__":
-
- pretrain(train_valid_test_datasets_provider, model_provider,
- ModelType.encoder_or_decoder,
- forward_step, args_defaults={'tokenizer_type': 'GPT2BPETokenizer'},
- extra_args_provider=add_validation_args,)
diff --git a/examples/detxoify_lm/generate_samples_gpt.py b/examples/detxoify_lm/generate_samples_gpt.py
deleted file mode 100644
index 47e1590ea56..00000000000
--- a/examples/detxoify_lm/generate_samples_gpt.py
+++ /dev/null
@@ -1,199 +0,0 @@
-# coding=utf-8
-# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
-
-
-"""Sample Generate GPT"""
-import json
-import os
-import sys
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
- os.path.pardir, os.path.pardir)))
-import torch
-from megatron import get_args
-from megatron import get_tokenizer
-from megatron import print_rank_0
-from megatron.checkpointing import load_checkpoint
-from megatron.core import mpu
-from megatron.initialize import initialize_megatron
-from megatron.model import GPTModel
-from megatron.training import get_model
-from megatron.text_generation import generate_and_post_process
-
-
-def model_provider(pre_process=True, post_process=True):
- """Build the model."""
-
- print_rank_0('building GPT model ...')
- model = GPTModel(num_tokentypes=0, parallel_output=False,
- pre_process=pre_process, post_process=post_process)
-
- return model
-
-def add_text_generate_args(parser):
- """Text generation arguments."""
- group = parser.add_argument_group(title='text generation')
-
- group.add_argument("--temperature", type=float, default=1.0,
- help='Sampling temperature.')
- group.add_argument("--greedy", action='store_true', default=False,
- help='Use greedy sampling.')
- group.add_argument("--top_p", type=float, default=0.0,
- help='Top p sampling.')
- group.add_argument("--top_k", type=int, default=0,
- help='Top k sampling.')
- group.add_argument("--out-seq-length", type=int, default=1024,
- help='Size of the output generated text.')
- group.add_argument("--sample-input-file", type=str, default=None,
- help='Get input from file instead of interactive mode, '
- 'each line is an input.')
- group.add_argument("--sample-output-file", type=str, default=None,
- help='Output file got from --sample-input-file')
- group.add_argument("--num-samples", type=int, default=0,
- help='Number of samples to generate unconditionally, '
- 'defaults to 0 and interactive conditional sampling')
- group.add_argument("--genfile", type=str,
- help='Output file when generating unconditionally')
- return parser
-
-def generate_samples_unconditional(model):
- args = get_args()
-
- if torch.distributed.get_rank() == 0:
- cnt = 0
- num_samples = args.num_samples
- from tqdm import tqdm
- pbar = tqdm(total=num_samples)
-
- while True:
- if torch.distributed.get_rank() == 0:
- sentences = [''] * args.global_batch_size
- print("global batch size", args.global_batch_size)
- max_len = args.out_seq_length
- resp_sentences, resp_sentences_seg, output_logits, \
- tokens = generate_and_post_process(model, prompts=sentences,
- tokens_to_generate=max_len,
- return_output_log_probs=False,
- top_k_sampling=args.top_k,
- top_p_sampling=args.top_p,
- add_BOS=True,
- temperature=1.0)
- for prompt, generation, token in zip(sentences, resp_sentences, tokens):
- datum = {'text': generation[len(prompt):], 'all_text': generation, 'prompt': prompt, 'id': cnt}
- yield datum
- cnt += 1
- pbar.update()
- if cnt >= num_samples:
- break
-
- if cnt >= num_samples:
- pbar.close()
- break
- else:
- generate_and_post_process(model)
-
-
-def generate_samples_conditional(model):
- args = get_args()
-
- if torch.distributed.get_rank() == 0:
- num_samples = args.num_samples
- cnt = 0
- from tqdm import tqdm
- pbar = tqdm(total=num_samples)
-
- fname = open(args.sample_input_file, "r")
- lines = fname.readlines()
- all_raw_text = [json.loads(line)['prompt']['text'] for line in lines]
- input_count = len(all_raw_text)
- input_pos = 0
-
- while True:
- torch.distributed.barrier()
- if torch.distributed.get_rank() == 0:
- sentences = []
- print("global batch size", args.global_batch_size)
- for _ in range(args.global_batch_size):
- if input_pos >= input_count:
- print(f"input pos: {input_pos}, input count: {input_count}")
- raw_text = "EMPTY TEXT"
- else:
- raw_text = all_raw_text[input_pos]
- input_pos += 1
- sentences.append(raw_text)
-
- max_len = args.out_seq_length
- resp_sentences, resp_sentences_seg, output_logits, \
- tokens = generate_and_post_process(model, prompts=sentences,
- tokens_to_generate=max_len,
- return_output_log_probs=False,
- top_k_sampling=args.top_k,
- top_p_sampling=args.top_p,
- add_BOS=False,
- temperature=1.0)
- for prompt, generation, token in zip(sentences, resp_sentences, tokens):
- datum = {'text': generation[len(prompt):], 'all_text': generation, 'prompt': prompt, 'id': cnt}
- yield datum
- cnt += 1
- pbar.update()
- if cnt >= num_samples:
- break
-
- if cnt >= num_samples:
- pbar.close()
- break
- else:
- generate_and_post_process(model)
-
-
-def generate_and_write_samples_unconditional(model):
- args = get_args()
- assert args.genfile is not None
- with open(args.genfile, 'w') as f:
- for datum in generate_samples_unconditional(model):
- if torch.distributed.get_rank() == 0:
- f.write(json.dumps(datum) + '\n')
-
-
-def generate_and_write_samples_conditional(model):
- args = get_args()
- if args.sample_output_file is None:
- sample_output_file = args.sample_input_file + ".out"
- print('`sample-output-file` not specified, setting '
- 'it to {}'.format(sample_output_file))
- else:
- sample_output_file = args.sample_output_file
- with open(sample_output_file, 'w') as f:
- for datum in generate_samples_conditional(model):
- if torch.distributed.get_rank() == 0:
- f.write(json.dumps(datum) + '\n')
-
-
-def main():
- """Main program."""
-
- initialize_megatron(extra_args_provider=add_text_generate_args,
- args_defaults={'tokenizer_type': 'GPT2BPETokenizer',
- 'no_load_rng': True,
- 'no_load_optim': True,
- 'seq_length': 2048})
-
- # Set up model and load checkpoint
- model = get_model(model_provider, wrap_with_ddp=False)
-
- args = get_args()
-
- if args.load is not None:
- _ = load_checkpoint(model, None, None)
- model = model[0]
-
- # Generate samples.
- if args.sample_input_file != None:
- print(f"{args.sample_input_file}")
- generate_and_write_samples_conditional(model)
- else:
- generate_and_write_samples_unconditional(model)
-
-
-if __name__ == "__main__":
-
- main()
diff --git a/examples/evaluate_retriever_nq.sh b/examples/evaluate_retriever_nq.sh
deleted file mode 100644
index 16e937f4fd0..00000000000
--- a/examples/evaluate_retriever_nq.sh
+++ /dev/null
@@ -1,38 +0,0 @@
-#!/bin/bash
-
-# Evaluate natural question test data given Wikipedia embeddings and pretrained
-# ICT model or a finetuned model for Natural Question task
-
-# Datasets can be downloaded from the following link:
-# https://github.com/facebookresearch/DPR/blob/master/data/download_data.py
-
-EVIDENCE_DATA_DIR=
-EMBEDDING_PATH=
-CHECKPOINT_PATH=
-
-QA_FILE=
-
-python tasks/main.py \
- --task RETRIEVER-EVAL \
- --tokenizer-type BertWordPieceLowerCase \
- --num-layers 12 \
- --hidden-size 768 \
- --num-attention-heads 12 \
- --tensor-model-parallel-size 1 \
- --micro-batch-size 128 \
- --activations-checkpoint-method uniform \
- --seq-length 512 \
- --max-position-embeddings 512 \
- --load ${CHECKPOINT_PATH} \
- --evidence-data-path ${EVIDENCE_DATA_DIR} \
- --embedding-path ${EMBEDDING_PATH} \
- --retriever-seq-length 256 \
- --vocab-file bert-vocab.txt\
- --qa-data-test ${QA_FILE} \
- --faiss-use-gpu \
- --retriever-report-topk-accuracies 1 5 20 100 \
- --fp16 \
- --indexer-log-interval 1000 \
- --indexer-batch-size 128
-
-
diff --git a/examples/evaluate_zeroshot_gpt.sh b/examples/evaluate_zeroshot_gpt.sh
deleted file mode 100755
index f8c38dc01d4..00000000000
--- a/examples/evaluate_zeroshot_gpt.sh
+++ /dev/null
@@ -1,38 +0,0 @@
-#!/bin/bash
-
-WORLD_SIZE=8
-
-DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
- --nnodes 1 \
- --node_rank 0 \
- --master_addr localhost \
- --master_port 6000"
-
-TASK="LAMBADA"
-
-VALID_DATA=
-VOCAB_FILE=gpt2-vocab.json
-MERGE_FILE=gpt2-merges.txt
-CHECKPOINT=checkpoints/gpt2_345m
-
-
-python -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/main.py \
- --task $TASK \
- --valid-data $VALID_DATA \
- --tokenizer-type GPT2BPETokenizer \
- --strict-lambada \
- --vocab-file $VOCAB_FILE \
- --merge-file $MERGE_FILE \
- --load $CHECKPOINT \
- --tensor-model-parallel-size 1 \
- --num-layers 24 \
- --hidden-size 1024 \
- --num-attention-heads 16 \
- --batch-size 8 \
- --activations-checkpoint-method uniform \
- --seq-length 1024 \
- --max-position-embeddings 1024 \
- --log-interval 10 \
- --fp16 \
- --no-load-optim \
- --no-load-rng
diff --git a/examples/export/README.md b/examples/export/README.md
new file mode 100644
index 00000000000..c9539e8ab21
--- /dev/null
+++ b/examples/export/README.md
@@ -0,0 +1,10 @@
+# Megatron Core Export
+
+This module is used to export megatron core models to different inference frameworks.
+Currently we support TRTLLM export . In the future we will be adding support for VLLM etc.
+
+## PTQ AND EXPORT
+Follow 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.
+
+# TRTLLM EXPORT
+Follow the instructions in [trtllm_export](./trtllm_export/) to do export to TRTLLM checkpoint format alone.
diff --git a/examples/export/trtllm_export/README.md b/examples/export/trtllm_export/README.md
new file mode 100644
index 00000000000..52cad785838
--- /dev/null
+++ b/examples/export/trtllm_export/README.md
@@ -0,0 +1,161 @@
+# Megatron Core To TRTLLM Export Documentation
+This guide will walk you through how you can use the megatron core export for exporting models to trtllm format
+
+### Contents
+- [Megatron Core To TRTLLM Export Documentation](#megatron-core-to-trtllm-export-documentation)
+- [Contents](#contents)
+ - [1. Quick Start](#1-quick-start)
+ - [1.1 Understanding The Code](#11-understanding-the-code)
+ - [1.2 Running The Code](#12-running-the-code)
+ - [2. GPU Export](#2-gpu-export)
+ - [3. Future work](#4-future-work)
+
+#### 1. Quick Start
+This 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)
+
+NOTE: 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.
+
+
+
+##### 1.1 Understanding The Code
+***STEP 1 - We initialize model parallel and other default arguments***
+We initalize tp and pp to 1 so that we can get the full model state dict on cpu
+```python
+ initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
+```
+
+***STEP 2 - We load the model using the model_provider_function***
+NOTE: We create a simple gpt model
+
+```python
+ transformer_config = TransformerConfig(
+ num_layers=2,
+ hidden_size=64, # Needs to be atleast 32 times num_attn_heads
+ num_attention_heads=2,
+ use_cpu_initialization=True,
+ pipeline_dtype=torch.float32,
+ )
+
+ gpt_model = GPTModel(
+ config=transformer_config,
+ transformer_layer_spec=get_gpt_layer_local_spec(),
+ vocab_size=100,
+ max_sequence_length=_SEQUENCE_LENGTH,
+ )
+
+ # Optionally you can also load a model using this code
+ # sharded_state_dict=gpt_model.sharded_state_dict(prefix='')
+ # checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)
+ # gpt_model.load_state_dict(checkpoint)
+
+```
+
+***STEP 3 - Instantiate the TRTLLM Helper***
+We instantiate the [TRTLLM Helper](../../../megatron/core/export/trtllm/trtllm_helper.py) For the GPT model we instantiate trtllm_helper as shown below.
+```python
+ if hasattr(gpt_model, "rotary_pos_emb"):
+ seq_len_interpolation_factor = gpt_model.rotary_pos_emb.seq_len_interpolation_factor
+
+ trtllm_helper = TRTLLMHelper(
+ transformer_config=gpt_model.config,
+ model_type=ModelType.gpt,
+ position_embedding_type = gpt_model.position_embedding_type,
+ max_position_embeddings = gpt_model.max_position_embeddings,
+ rotary_percentage = gpt_model.rotary_percent,
+ rotary_base = gpt_model.rotary_base,
+ moe_tp_mode = 2,
+ multi_query_mode = False,
+ activation = "gelu",
+ seq_len_interpolation_factor = seq_len_interpolation_factor,
+ share_embeddings_and_output_weights=gpt_model.share_embeddings_and_output_weights
+ )
+```
+
+***STEP 4 - Get the TRTLLM Weights and configs***
+To 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.
+
+```python
+ model_state_dict={}
+ for key , val in gpt_model.state_dict().items():
+ # val is non for _extra_state layers . We filter it out
+ if val is not None:
+ model_state_dict[key] = val
+
+ export_config = ExportConfig(inference_tp_size = 2)
+ weight_list, config_list = trtllm_helper.get_trtllm_pretrained_config_and_model_weights(
+ model_state_dict= model_state_dict,
+ dtype = DataType.bfloat16,
+ export_config=export_config
+ )
+```
+
+***STEP 5 - Build the TRTLLM Engine***
+Following code is used to build the TRTLLM Engine.
+
+```python
+ for trtllm_model_weights, trtllm_model_config in zip(weight_list, config_list):
+ trtllm_helper.build_and_save_engine(
+ max_input_len=256,
+ max_output_len=256,
+ max_batch_size=8,
+ engine_dir='/opt/megatron-lm/engine',
+ trtllm_model_weights=trtllm_model_weights,
+ trtllm_model_config=trtllm_model_config,
+ lora_ckpt_list=None,
+ use_lora_plugin=None,
+ max_lora_rank=64,
+ lora_target_modules=None,
+ max_prompt_embedding_table_size=0,
+ paged_kv_cache=True,
+ remove_input_padding=True,
+ paged_context_fmha=False,
+ use_refit=False,
+ max_num_tokens=None,
+ max_seq_len=512,
+ opt_num_tokens=None,
+ max_beam_width=1,
+ tokens_per_block=128,
+ multiple_profiles=False,
+ gpt_attention_plugin="auto",
+ gemm_plugin="auto",
+ )
+```
+
+
+##### 1.2 Running The Code
+An example run script is shown below.
+
+```
+# In a workstation
+MLM_PATH=/path/to/megatron-lm
+CONTAINER_IMAGE=gitlab-master.nvidia.com:5005/dl/joc/nemo-ci/trtllm_0.12/train:pipe.17669124-x86
+
+docker run -it --gpus=all --ipc=host -v $MLM_PATH/:/opt/megatron-lm $CONTAINER_IMAGE bash
+
+# Inside the container run the following.
+
+cd /opt/megatron-lm/
+
+CUDA_VISIBLE_DEVICES=0 torchrun --nproc-per-node 1 examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py
+```
+
+
+
+#### 2. GPU Export
+You 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.
+In 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.
+
+To run the gpu version
+
+```
+CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc-per-node 2 examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py
+```
+
+
+
+#### 3. Future work
+The following are planned for the future releases .
+* Pipeline parallellism for export (Work in progress)
+* GPU Export for more models (Work in progress for some models)
+* Refit functionality
+* VLLM Support
\ No newline at end of file
diff --git a/examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py b/examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py
new file mode 100644
index 00000000000..57d44f9f628
--- /dev/null
+++ b/examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py
@@ -0,0 +1,117 @@
+import os
+import torch
+from megatron.core import parallel_state
+from megatron.core import dist_checkpointing
+from megatron.core.export.model_type import ModelType
+from megatron.core.export.data_type import DataType
+from megatron.core.export.trtllm.trtllm_helper import TRTLLMHelper
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
+
+
+_SEQUENCE_LENGTH = 64
+_VOCAB_SIZE = 256
+
+def initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1):
+ parallel_state.destroy_model_parallel()
+
+ # Torch setup for distributed training
+ rank = int(os.environ['LOCAL_RANK'])
+ world_size = torch.cuda.device_count()
+ torch.cuda.set_device(rank)
+ torch.distributed.init_process_group(world_size=world_size, rank=rank)
+
+ # Megatron core distributed training initialization
+ parallel_state.initialize_model_parallel(tensor_model_parallel_size = tensor_model_parallel_size, pipeline_model_parallel_size=pipeline_model_parallel_size)
+
+def model_provider():
+ """Build the model."""
+
+ transformer_config = TransformerConfig(
+ num_layers=2,
+ hidden_size=64,
+ num_attention_heads=2,
+ use_cpu_initialization=True,
+ pipeline_dtype=torch.float32
+ )
+
+ gpt_model = GPTModel(
+ config=transformer_config,
+ transformer_layer_spec=get_gpt_layer_local_spec(),
+ vocab_size=_VOCAB_SIZE,
+ max_sequence_length=_SEQUENCE_LENGTH,
+ )
+
+ return gpt_model
+
+def load_distributed_checkpoint(checkpoint_path, gpt_model):
+ sharded_state_dict=gpt_model.sharded_state_dict(prefix='')
+ checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)
+ gpt_model.load_state_dict(checkpoint)
+ return gpt_model
+
+if __name__ == "__main__":
+ initialize_distributed(tensor_model_parallel_size=2, pipeline_model_parallel_size=1)
+ model_parallel_cuda_manual_seed(123)
+
+ gpt_model = model_provider()
+ device = torch.device("cuda")
+ gpt_model.to(device)
+
+ # Optionally you can also load a gpt model from ckpt_path using this code below
+ # gpt_model = load_distributed_checkpoint(gpt_model=gpt_model, checkpoint_path=ckpt_path)
+
+ seq_len_interpolation_factor = None
+ if hasattr(gpt_model, "rotary_pos_emb"):
+ seq_len_interpolation_factor = gpt_model.rotary_pos_emb.seq_len_interpolation_factor
+
+ trtllm_helper = TRTLLMHelper(
+ transformer_config=gpt_model.config,
+ model_type=ModelType.gpt,
+ position_embedding_type = gpt_model.position_embedding_type,
+ max_position_embeddings = gpt_model.max_position_embeddings,
+ rotary_percentage = gpt_model.rotary_percent,
+ rotary_base = gpt_model.rotary_base,
+ moe_tp_mode = 2,
+ multi_query_mode = False,
+ activation = "gelu",
+ seq_len_interpolation_factor = seq_len_interpolation_factor,
+ share_embeddings_and_output_weights=gpt_model.share_embeddings_and_output_weights
+ )
+
+
+ trtllm_model_weights, trtllm_model_config = trtllm_helper.get_trtllm_pretrained_config_and_model_weights(
+ model_state_dict= gpt_model.state_dict(),
+ dtype = DataType.bfloat16,
+ on_device_distributed_conversion=True,
+ vocab_size=_VOCAB_SIZE,
+ gpus_per_node=2,
+ )
+
+ trtllm_helper.build_and_save_engine(
+ max_input_len=256,
+ max_output_len=256,
+ max_batch_size=8,
+ engine_dir='/opt/megatron-lm/engine',
+ trtllm_model_weights=trtllm_model_weights[0],
+ trtllm_model_config=trtllm_model_config[0],
+ lora_ckpt_list=None,
+ use_lora_plugin=None,
+ max_lora_rank=64,
+ lora_target_modules=None,
+ max_prompt_embedding_table_size=0,
+ paged_kv_cache=True,
+ remove_input_padding=True,
+ paged_context_fmha=False,
+ use_refit=False,
+ max_num_tokens=None,
+ max_seq_len=512,
+ opt_num_tokens=None,
+ max_beam_width=1,
+ tokens_per_block=128,
+ multiple_profiles=False,
+ gpt_attention_plugin="auto",
+ gemm_plugin="auto",
+ )
diff --git a/examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py b/examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py
new file mode 100644
index 00000000000..587e7cfdd32
--- /dev/null
+++ b/examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py
@@ -0,0 +1,118 @@
+import os
+import torch
+from megatron.core import parallel_state
+from megatron.core import dist_checkpointing
+from megatron.core.export.model_type import ModelType
+from megatron.core.export.data_type import DataType
+from megatron.core.export.export_config import ExportConfig
+from megatron.core.export.trtllm.trtllm_helper import TRTLLMHelper
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
+
+
+_SEQUENCE_LENGTH = 64
+
+
+def initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1):
+ parallel_state.destroy_model_parallel()
+
+ # Torch setup for distributed training
+ rank = int(os.environ['LOCAL_RANK'])
+ world_size = torch.cuda.device_count()
+ torch.cuda.set_device(rank)
+ torch.distributed.init_process_group(world_size=world_size, rank=rank)
+
+ # Megatron core distributed training initialization
+ parallel_state.initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size)
+
+def model_provider():
+ """Build the model."""
+
+ transformer_config = TransformerConfig(
+ num_layers=2,
+ hidden_size=64, # Needs to be atleast 32 times num_attn_heads
+ num_attention_heads=2,
+ use_cpu_initialization=True,
+ pipeline_dtype=torch.float32,
+ )
+
+ gpt_model = GPTModel(
+ config=transformer_config,
+ transformer_layer_spec=get_gpt_layer_local_spec(),
+ vocab_size=100,
+ max_sequence_length=_SEQUENCE_LENGTH,
+ )
+
+ return gpt_model
+
+def load_distributed_checkpoint(checkpoint_path, gpt_model):
+ sharded_state_dict=gpt_model.sharded_state_dict(prefix='')
+ checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)
+ gpt_model.load_state_dict(checkpoint)
+ return gpt_model
+
+if __name__ == "__main__":
+ # Need to use TP1 PP1 for export on single device
+ initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
+ model_parallel_cuda_manual_seed(123)
+
+ gpt_model = model_provider()
+
+ # Optionally you can also load a gpt model from ckpt_path using this code below
+ # gpt_model = load_distributed_checkpoint(gpt_model=gpt_model, checkpoint_path=ckpt_path)
+
+ seq_len_interpolation_factor = None
+ if hasattr(gpt_model, "rotary_pos_emb"):
+ seq_len_interpolation_factor = gpt_model.rotary_pos_emb.seq_len_interpolation_factor
+
+ trtllm_helper = TRTLLMHelper(
+ transformer_config=gpt_model.config,
+ model_type=ModelType.gpt,
+ position_embedding_type = gpt_model.position_embedding_type,
+ max_position_embeddings = gpt_model.max_position_embeddings,
+ rotary_percentage = gpt_model.rotary_percent,
+ rotary_base = gpt_model.rotary_base,
+ moe_tp_mode = 2,
+ multi_query_mode = False,
+ activation = "gelu",
+ seq_len_interpolation_factor = seq_len_interpolation_factor,
+ share_embeddings_and_output_weights=gpt_model.share_embeddings_and_output_weights
+ )
+
+
+ export_config = ExportConfig(inference_tp_size = 2)
+ # NOTE : For faster performance, if your entire model will fit in gpu memory, transfer model state dict to GPU and then call this api
+ weight_list, config_list = trtllm_helper.get_trtllm_pretrained_config_and_model_weights(
+ model_state_dict= gpt_model.state_dict(),
+ dtype = DataType.bfloat16,
+ export_config=export_config
+ )
+
+ for trtllm_model_weights, trtllm_model_config in zip(weight_list, config_list):
+ trtllm_helper.build_and_save_engine(
+ max_input_len=256,
+ max_output_len=256,
+ max_batch_size=8,
+ engine_dir='/opt/megatron-lm/engine',
+ trtllm_model_weights=trtllm_model_weights,
+ trtllm_model_config=trtllm_model_config,
+ lora_ckpt_list=None,
+ use_lora_plugin=None,
+ max_lora_rank=64,
+ lora_target_modules=None,
+ max_prompt_embedding_table_size=0,
+ paged_kv_cache=True,
+ remove_input_padding=True,
+ paged_context_fmha=False,
+ use_refit=False,
+ max_num_tokens=None,
+ max_seq_len=512,
+ opt_num_tokens=None,
+ max_beam_width=1,
+ tokens_per_block=128,
+ multiple_profiles=False,
+ gpt_attention_plugin="auto",
+ gemm_plugin="auto",
+ )
\ No newline at end of file
diff --git a/examples/finetune_mnli_distributed.sh b/examples/finetune_mnli_distributed.sh
deleted file mode 100755
index 9219e595dd2..00000000000
--- a/examples/finetune_mnli_distributed.sh
+++ /dev/null
@@ -1,44 +0,0 @@
-#!/bin/bash
-
-WORLD_SIZE=8
-
-DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
- --nnodes 1 \
- --node_rank 0 \
- --master_addr localhost \
- --master_port 6000"
-
-TRAIN_DATA="data/glue_data/MNLI/train.tsv"
-VALID_DATA="data/glue_data/MNLI/dev_matched.tsv \
- data/glue_data/MNLI/dev_mismatched.tsv"
-PRETRAINED_CHECKPOINT=checkpoints/bert_345m
-VOCAB_FILE=bert-vocab.txt
-CHECKPOINT_PATH=checkpoints/bert_345m_mnli
-
-python -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/main.py \
- --task MNLI \
- --seed 1234 \
- --train-data $TRAIN_DATA \
- --valid-data $VALID_DATA \
- --tokenizer-type BertWordPieceLowerCase \
- --vocab-file $VOCAB_FILE \
- --epochs 5 \
- --pretrained-checkpoint $PRETRAINED_CHECKPOINT \
- --tensor-model-parallel-size 1 \
- --num-layers 24 \
- --hidden-size 1024 \
- --num-attention-heads 16 \
- --micro-batch-size 8 \
- --activations-checkpoint-method uniform \
- --lr 5.0e-5 \
- --lr-decay-style linear \
- --lr-warmup-fraction 0.065 \
- --seq-length 512 \
- --max-position-embeddings 512 \
- --save-interval 500000 \
- --save $CHECKPOINT_PATH \
- --log-interval 10 \
- --eval-interval 100 \
- --eval-iters 50 \
- --weight-decay 1.0e-1 \
- --fp16
diff --git a/examples/finetune_race_distributed.sh b/examples/finetune_race_distributed.sh
deleted file mode 100755
index e7f70a70abe..00000000000
--- a/examples/finetune_race_distributed.sh
+++ /dev/null
@@ -1,47 +0,0 @@
-#!/bin/bash
-
-WORLD_SIZE=8
-
-DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
- --nnodes 1 \
- --node_rank 0 \
- --master_addr localhost \
- --master_port 6000"
-
-TRAIN_DATA="data/RACE/train/middle"
-VALID_DATA="data/RACE/dev/middle \
- data/RACE/dev/high"
-VOCAB_FILE=bert-vocab.txt
-PRETRAINED_CHECKPOINT=checkpoints/bert_345m
-CHECKPOINT_PATH=checkpoints/bert_345m_race
-
-python -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/main.py \
- --task RACE \
- --seed 1234 \
- --train-data $TRAIN_DATA \
- --valid-data $VALID_DATA \
- --tokenizer-type BertWordPieceLowerCase \
- --vocab-file $VOCAB_FILE \
- --epochs 3 \
- --pretrained-checkpoint $PRETRAINED_CHECKPOINT \
- --tensor-model-parallel-size 1 \
- --num-layers 24 \
- --hidden-size 1024 \
- --num-attention-heads 16 \
- --micro-batch-size 4 \
- --activations-checkpoint-method uniform \
- --lr 1.0e-5 \
- --lr-decay-style linear \
- --lr-warmup-fraction 0.06 \
- --seq-length 512 \
- --max-position-embeddings 512 \
- --save-interval 100000 \
- --save $CHECKPOINT_PATH \
- --log-interval 10 \
- --eval-interval 100 \
- --eval-iters 50 \
- --weight-decay 1.0e-1 \
- --clip-grad 1.0 \
- --hidden-dropout 0.1 \
- --attention-dropout 0.1 \
- --fp16
diff --git a/examples/finetune_retriever_distributed.sh b/examples/finetune_retriever_distributed.sh
deleted file mode 100755
index 535a2e053d4..00000000000
--- a/examples/finetune_retriever_distributed.sh
+++ /dev/null
@@ -1,56 +0,0 @@
-#!/bin/bash
-
-# Finetune a BERT or pretrained ICT model using Google natural question data
-# Datasets can be downloaded from the following link:
-# https://github.com/facebookresearch/DPR/blob/master/data/download_data.py
-
-WORLD_SIZE=8
-
-DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
- --nnodes 1 \
- --node_rank 0 \
- --master_addr localhost \
- --master_port 6000"
-
-CHECKPOINT_PATH=
-
-# Load either of the below
-BERT_LOAD_PATH=
-PRETRAINED_CHECKPOINT=
-
-python -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/main.py \
- --task RET-FINETUNE-NQ \
- --train-with-neg \
- --train-hard-neg 1 \
- --pretrained-checkpoint ${PRETRAINED_CHECKPOINT} \
- --num-layers 12 \
- --hidden-size 768 \
- --num-attention-heads 12 \
- --tensor-model-parallel-size 1 \
- --tokenizer-type BertWordPieceLowerCase \
- --train-data nq-train.json \
- --valid-data nq-dev.json \
- --save ${CHECKPOINT_PATH} \
- --load ${CHECKPOINT_PATH} \
- --vocab-file bert-vocab.txt \
- --bert-load ${BERT_LOAD_PATH} \
- --save-interval 5000 \
- --log-interval 10 \
- --eval-interval 20000 \
- --eval-iters 100 \
- --indexer-log-interval 1000 \
- --faiss-use-gpu \
- --DDP-impl torch \
- --fp16 \
- --retriever-report-topk-accuracies 1 5 10 20 100 \
- --seq-length 512 \
- --retriever-seq-length 256 \
- --max-position-embeddings 512 \
- --retriever-score-scaling \
- --epochs 80 \
- --micro-batch-size 8 \
- --eval-micro-batch-size 16 \
- --indexer-batch-size 128 \
- --lr 2e-5 \
- --lr-warmup-fraction 0.01 \
- --weight-decay 1e-1
diff --git a/examples/gpt3/README.md b/examples/gpt3/README.md
new file mode 100644
index 00000000000..8d6f2674163
--- /dev/null
+++ b/examples/gpt3/README.md
@@ -0,0 +1,57 @@
+# GPT3 MODEL
+
+## Table of contents
+- [1. Training Setup](#1-training-setup)
+- [2. Configurations](#2-configurations)
+- [3. Training Results](#3-training-results)
+
+## 1. Training setup
+
+
+To run the model using a docker container run it as follows
+```
+PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:24.01-py3
+CHECKPOINT_PATH="" #
+TENSORBOARD_LOGS_PATH=""#
+VOCAB_FILE="" #/gpt2-vocab.json
+MERGE_FILE="" #/gpt2-merges.txt
+DATA_PATH="" #_text_document
+
+docker run \
+ --gpus=all \
+ --ipc=host \
+ --workdir /workspace/megatron-lm \
+ -v /path/to/data:/path/to/data \
+ -v /path/to/megatron-lm:/workspace/megatron-lm \
+ megatron-lm nvcr.io/nvidia/pytorch:24.01-py3 \
+ bash examples/gpt3/train_gpt3_175b_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH $VOCAB_FILE $MERGE_FILE $DATA_PATH "
+
+```
+NOTE: Depending on the environment you are running it the above command might like slightly different.
+
+
+## 2. Configurations
+
+The example in this folder shows you how to run 175B model. There are other configs you could run as well
+
+### 345M
+```
+ --num-layers 12 \
+ --hidden-size 512 \
+ --num-attention-heads 8 \
+ --seq-length 1024 \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+
+```
+
+### 857M
+```
+ --num-layers 24 \
+ --hidden-size 1024 \
+ --num-attention-heads 16 \
+ --seq-length 2048 \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+
+```
diff --git a/examples/gpt3/train_gpt3_175b_distributed.sh b/examples/gpt3/train_gpt3_175b_distributed.sh
new file mode 100755
index 00000000000..be00d76120d
--- /dev/null
+++ b/examples/gpt3/train_gpt3_175b_distributed.sh
@@ -0,0 +1,82 @@
+#!/bin/bash
+
+# Runs the "175B" parameter model
+
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+GPUS_PER_NODE=8
+# Change for multinode config
+MASTER_ADDR=localhost
+MASTER_PORT=6000
+NUM_NODES=1
+NODE_RANK=0
+WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
+
+CHECKPOINT_PATH=$1 #
+TENSORBOARD_LOGS_PATH=$2 #
+VOCAB_FILE=$3 #/gpt2-vocab.json
+MERGE_FILE=$4 #/gpt2-merges.txt
+DATA_PATH=$5 #_text_document
+
+DISTRIBUTED_ARGS=(
+ --nproc_per_node $GPUS_PER_NODE
+ --nnodes $NUM_NODES
+ --master_addr $MASTER_ADDR
+ --master_port $MASTER_PORT
+)
+
+GPT_MODEL_ARGS=(
+ --num-layers 96
+ --hidden-size 12288
+ --num-attention-heads 96
+ --seq-length 2048
+ --max-position-embeddings 2048
+ --attention-backend auto # Can use (flash/fused/unfused/local)
+)
+
+TRAINING_ARGS=(
+ --micro-batch-size 1
+ --global-batch-size 1536
+ --step-batch-size-schedule "0:16 2.4B:320 4.8B:624 7.2B:928 9.6B:1232 12B:1536"
+ --train-iters 500000
+ --weight-decay 0.1
+ --adam-beta1 0.9
+ --adam-beta2 0.95
+ --init-method-std 0.006
+ --clip-grad 1.0
+ --fp16
+ --lr 6.0e-5
+ --lr-decay-style cosine
+ --min-lr 6.0e-6
+ --lr-warmup-fraction .001
+ --lr-decay-iters 430000
+)
+
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size 8
+ --pipeline-model-parallel-size 16
+)
+
+DATA_ARGS=(
+ --data-path $DATA_PATH
+ --vocab-file $VOCAB_FILE
+ --merge-file $MERGE_FILE
+ --split 949,50,1
+)
+
+EVAL_AND_LOGGING_ARGS=(
+ --log-interval 100
+ --save-interval 10000
+ --eval-interval 1000
+ --save $CHECKPOINT_PATH
+ --load $CHECKPOINT_PATH
+ --eval-iters 10
+ --tensorboard-dir $TENSORBOARD_LOGS_PATH
+)
+
+torchrun ${DISTRIBUTED_ARGS[@]} pretrain_gpt.py \
+ ${GPT_MODEL_ARGS[@]} \
+ ${TRAINING_ARGS[@]} \
+ ${MODEL_PARALLEL_ARGS[@]} \
+ ${DATA_ARGS[@]} \
+ ${EVAL_AND_LOGGING_ARGS[@]}
diff --git a/examples/gptoss/01_convert_from_hf.py b/examples/gptoss/01_convert_from_hf.py
new file mode 100644
index 00000000000..adee3358ec3
--- /dev/null
+++ b/examples/gptoss/01_convert_from_hf.py
@@ -0,0 +1,55 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Convert HuggingFace checkpoints to Megatron format."""
+
+import os
+import argparse
+
+from megatron.bridge import AutoBridge
+
+def _parse_args():
+ parser = argparse.ArgumentParser(description="Convert HF LLMs to Megatron format")
+ parser.add_argument(
+ "--hf-model",
+ type=str,
+ required=True,
+ help="HuggingFace model identifier or path",
+ )
+ parser.add_argument(
+ "--save-path",
+ type=str,
+ default=None,
+ help="Path to save the converted Megatron checkpoint",
+ )
+ parser.add_argument('--local-rank', '--local_rank', type=int, default=0)
+ return parser.parse_args()
+
+if __name__ == "__main__":
+ args = _parse_args()
+ HF_MODEL = args.hf_model
+ SAVE_PATH = args.save_path
+ WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1))
+
+ if SAVE_PATH is None:
+ SAVE_PATH = f"./megatron_checkpoints/{HF_MODEL.replace('/', '_')}"
+
+ print(f"Converting {HF_MODEL} to Megatron format...")
+ print(f"Save path: {SAVE_PATH}")
+
+ bridge = AutoBridge.from_hf_pretrained(HF_MODEL, trust_remote_code=True)
+ provider = bridge.to_megatron_provider()
+ # Update these configs as needed
+ provider.expert_tensor_parallel_size = 1
+ provider.tensor_model_parallel_size = 1
+ provider.pipeline_model_parallel_size = WORLD_SIZE
+ provider.finalize()
+
+ model = provider.provide_distributed_model(wrap_with_ddp=False)
+
+ bridge.save_megatron_model(
+ model,
+ SAVE_PATH,
+ hf_tokenizer_path=HF_MODEL
+ )
+
+ print(f"Saved Megatron checkpoint to {SAVE_PATH}")
diff --git a/examples/gptoss/02_train.sh b/examples/gptoss/02_train.sh
new file mode 100755
index 00000000000..d129adc2b84
--- /dev/null
+++ b/examples/gptoss/02_train.sh
@@ -0,0 +1,259 @@
+#!/bin/bash
+
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+
+
+# Setup arguments with defaults
+CHECKPOINT_PATH="NO_VALUE_PROVIDED"
+TENSORBOARD_LOGS_PATH="./tensorboard_logs/"
+TOKENIZER_ARG="MOCK"
+DATA_ARG="MOCK"
+DISTRIBUTED_CONFIG_FILE=""
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+ case $1 in
+ --checkpoint-path)
+ CHECKPOINT_PATH="$2"
+ shift 2
+ ;;
+ --tensorboard-logs-path)
+ TENSORBOARD_LOGS_PATH="$2"
+ shift 2
+ ;;
+ --tokenizer)
+ TOKENIZER_ARG="$2"
+ shift 2
+ ;;
+ --data)
+ DATA_ARG="$2"
+ shift 2
+ ;;
+ --distributed-config-file)
+ DISTRIBUTED_CONFIG_FILE="$2"
+ shift 2
+ ;;
+ -h|--help)
+ echo "Usage: $0 [OPTIONS]"
+ echo "Options:"
+ echo " --checkpoint-path PATH Path to Megatron checkpoint"
+ echo " --tensorboard-logs-path PATH Path to TensorBoard logs"
+ echo " --tokenizer PATH|MOCK Path to tokenizer model, or 'MOCK' (default: MOCK)"
+ echo " --data PATH|MOCK Data prefix, or 'MOCK' (default: MOCK)"
+ echo " --distributed-config-file FILE Path to distributed training config file"
+ echo " -h, --help Show this help message"
+ exit 0
+ ;;
+ *)
+ echo "Unknown option: $1"
+ echo "Use --help for usage information"
+ exit 1
+ ;;
+ esac
+done
+
+# Check if checkpoint path exists
+if [ ! -d "$CHECKPOINT_PATH" ]; then
+ echo "Error: Checkpoint path does not exist: $CHECKPOINT_PATH"
+ exit 1
+fi
+echo "Checkpoint path exists: $CHECKPOINT_PATH"
+
+# Check if tensorboard logs path exists
+if [ ! -d "$TENSORBOARD_LOGS_PATH" ]; then
+ echo "Warning: TensorBoard logs path does not exist. Creating: $TENSORBOARD_LOGS_PATH"
+ mkdir -p "$TENSORBOARD_LOGS_PATH"
+fi
+echo "TensorBoard logs path exists: $TENSORBOARD_LOGS_PATH"
+
+# NOTE: by default we use 8 GPUs
+# These values will be over-written below with environmental variables
+GPUS_PER_NODE=8
+NUM_NODES=1
+MASTER_ADDR="localhost"
+MASTER_PORT=6000
+NODE_RANK=0
+
+# Load distributed config from file if provided
+if [ -n "$DISTRIBUTED_CONFIG_FILE" ]; then
+ if [ ! -f "$DISTRIBUTED_CONFIG_FILE" ]; then
+ echo "Warning: Distributed config file does not exist: $DISTRIBUTED_CONFIG_FILE"
+ echo "Continuing with default distributed training settings."
+ else
+ echo "Loading distributed config from: $DISTRIBUTED_CONFIG_FILE"
+ source "$DISTRIBUTED_CONFIG_FILE"
+ fi
+fi
+
+# Override with environment variables if set
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+NUM_NODES=${NUM_NODES:-1}
+MASTER_ADDR=${MASTER_ADDR:-localhost}
+MASTER_PORT=${MASTER_PORT:-6000}
+NODE_RANK=${NODE_RANK:-0}
+WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
+
+# Path to the pretrain_gpt.py script, assuming this script is run from the root of the Megatron-LM repository
+PRETRAIN_SCRIPT_PATH="pretrain_gpt.py"
+
+# Data cache path (useful for both mock and real data)
+DATA_CACHE_PATH="${PWD}/benchmark_cache_gpt_oss_20b"
+mkdir -p "$DATA_CACHE_PATH"
+
+DISTRIBUTED_ARGS=(
+ --nproc_per_node $GPUS_PER_NODE
+ --nnodes $NUM_NODES
+ --master_addr $MASTER_ADDR
+ --master_port $MASTER_PORT
+ --node_rank $NODE_RANK
+)
+
+# NOTE: we only set pipeline parallelism to be the number of GPUs
+# Adjust each value based on your setup.
+TP_SIZE=1
+EP_SIZE=1
+PP_SIZE=${WORLD_SIZE}
+MICRO_BATCH_SIZE=1
+GLOBAL_BATCH_SIZE=128
+NUM_LAYERS=12
+DTYPE="fp8"
+SEQ_LENGTH=8192
+MAX_POSITION_EMBEDDINGS=8192
+TRAIN_SAMPLES=1953125000
+LR_DECAY_SAMPLES=1949218748
+
+MODEL_ARGS=(
+ --no-masked-softmax-fusion
+ --transformer-impl transformer_engine
+ --disable-bias-linear
+ --untie-embeddings-and-output-weights
+ --no-rope-fusion
+ --normalization RMSNorm
+ --num-layers ${NUM_LAYERS}
+ --hidden-size 512
+ --ffn-hidden-size 2048
+ --num-attention-heads 64
+ --group-query-attention
+ --num-query-groups 8
+ --seq-length ${SEQ_LENGTH}
+ --max-position-embeddings ${MAX_POSITION_EMBEDDINGS}
+ --use-mcore-models
+ --rotary-percent 1.0
+ --rope-type rope
+ --position-embedding-type rope
+ --rotary-base 10000
+ --no-bias-gelu-fusion
+ --export-force-local-attention
+ --no-bias-dropout-fusion
+ --quick-geglu
+ --glu-linear-offset 1.0
+ --softmax-type learnable
+ --window-attn-skip-freq 2
+ --activation-func-clamp-value 7.0
+ --window-size 127,0
+ --enable-gpt-oss
+)
+
+MOE_ARGS=(
+ --num-experts 4
+ --moe-router-topk 2
+ --moe-router-load-balancing-type aux_loss
+ --moe-aux-loss-coeff 1e-3
+ --moe-grouped-gemm
+ --moe-token-dispatcher-type alltoall
+ --overlap-param-gather
+ --overlap-grad-reduce
+ --moe-ffn-hidden-size 2048
+ --moe-router-dtype fp32
+ --moe-z-loss-coeff 1e-3
+ --moe-permute-fusion
+)
+
+DATA_ARGS_LIST=()
+if [[ "$TOKENIZER_ARG" == "MOCK" ]] || [[ "$DATA_ARG" == "MOCK" ]] || [[ -z "$TOKENIZER_ARG" ]]; then
+ DATA_ARGS_LIST+=(
+ "--mock-data"
+ "--tokenizer-type NullTokenizer"
+ "--vocab-size 128256"
+ "--data-cache-path ${DATA_CACHE_PATH}"
+ "--tiktoken-pattern v2"
+ "--split '99,1,0'"
+ "--no-create-attention-mask-in-dataloader"
+ "--no-mmap-bin-files"
+ "--num-workers 1"
+ )
+else
+ # Settings for real data
+ DATA_ARGS_LIST+=(
+ "--data-path $DATA_ARG"
+ "--tokenizer-type HuggingFaceTokenizer"
+ "--tokenizer-model $TOKENIZER_ARG"
+ "--data-cache-path ${DATA_CACHE_PATH}"
+ "--split '99,1,0'"
+ "--no-create-attention-mask-in-dataloader"
+ "--no-mmap-bin-files"
+ "--num-workers 1"
+ # Note: --vocab-size might be inferred by HuggingFaceTokenizer or might need to be explicit.
+ "--vocab-size 128256"
+ )
+fi
+
+TRAINING_ARGS=(
+ --micro-batch-size ${MICRO_BATCH_SIZE}
+ --global-batch-size ${GLOBAL_BATCH_SIZE}
+ --lr 1.0e-5
+ --train-samples ${TRAIN_SAMPLES}
+ --lr-decay-samples ${LR_DECAY_SAMPLES}
+ --lr-decay-style cosine
+ --min-lr 1.0e-6
+ --weight-decay 0.1
+ --lr-warmup-fraction 0.05
+ --clip-grad 1.0
+ --bf16
+ --use-flash-attn
+ --attention-softmax-in-fp32
+ --accumulate-allreduce-grads-in-fp32
+ --disable-bf16-reduced-precision-matmul
+ --recompute-activations
+)
+
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size ${TP_SIZE}
+ --pipeline-model-parallel-size ${PP_SIZE}
+ --expert-model-parallel-size ${EP_SIZE}
+ --sequence-parallel
+ --context-parallel-size 1
+ --use-distributed-optimizer
+ --fp8-format hybrid
+ --fp8-param-gather
+ --fp8-amax-compute-algo max
+ --fp8-amax-history-len 1024
+)
+
+LOGGING_ARGS=(
+ --log-interval 1
+ --save-interval 10000
+ --eval-interval 50000000
+ --eval-iters 0
+ --save $CHECKPOINT_PATH
+ --tensorboard-dir "${CHECKPOINT_PATH}/tensorboard"
+ --moe-per-layer-logging
+ --no-load-optim
+ --no-load-rng
+ --log-throughput
+)
+
+# Ensure pretrain_gpt.py is found
+if [ ! -f "$PRETRAIN_SCRIPT_PATH" ]; then
+ echo "Error: pretrain_gpt.py not found at $PRETRAIN_SCRIPT_PATH"
+ echo "Please ensure you are running this script from the root of the Megatron-LM repository, and pretrain_gpt.py is present."
+ exit 1
+fi
+
+python -m torch.distributed.run ${DISTRIBUTED_ARGS[@]} ${PRETRAIN_SCRIPT_PATH} \
+ ${MODEL_ARGS[@]} \
+ ${MOE_ARGS[@]} \
+ ${DATA_ARGS_LIST[@]} \
+ ${TRAINING_ARGS[@]} \
+ ${MODEL_PARALLEL_ARGS[@]} \
+ ${LOGGING_ARGS[@]}
diff --git a/examples/gptoss/03_convert_to_hf.py b/examples/gptoss/03_convert_to_hf.py
new file mode 100644
index 00000000000..8089afec854
--- /dev/null
+++ b/examples/gptoss/03_convert_to_hf.py
@@ -0,0 +1,52 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Convert HuggingFace checkpoints to Megatron format."""
+
+import os
+import argparse
+
+from megatron.bridge import AutoBridge
+
+def _parse_args():
+ parser = argparse.ArgumentParser(description="Convert Megatron LLMs to HuggingFace format")
+ parser.add_argument(
+ "--hf-model",
+ type=str,
+ required=True,
+ help="HuggingFace model identifier or path to load config from",
+ )
+ parser.add_argument(
+ "--megatron-model",
+ type=str,
+ required=True,
+ help="Megatron model identifier or path",
+ )
+ parser.add_argument(
+ "--save-path",
+ type=str,
+ default=None,
+ help="Path to save the converted HuggingFace checkpoint",
+ )
+ parser.add_argument('--local-rank', '--local_rank', type=int, default=0)
+ return parser.parse_args()
+
+if __name__ == "__main__":
+ args = _parse_args()
+ HF_MODEL = args.hf_model
+ MEGATRON_MODEL = args.megatron_model
+ SAVE_PATH = args.save_path
+ WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1))
+
+ if SAVE_PATH is None:
+ SAVE_PATH = f"./huggingface_checkpoints/{MEGATRON_MODEL.replace('/', '_')}"
+
+ print(f"Converting {MEGATRON_MODEL} to HuggingFace {HF_MODEL} format...")
+ print(f"Save path: {SAVE_PATH}")
+
+ bridge = AutoBridge.from_hf_pretrained(HF_MODEL, trust_remote_code=True)
+ bridge.export_ckpt(
+ MEGATRON_MODEL,
+ SAVE_PATH,
+ )
+
+ print(f"Saved HuggingFace checkpoint to {SAVE_PATH}")
diff --git a/examples/gptoss/README.md b/examples/gptoss/README.md
new file mode 100644
index 00000000000..eeb92ad9953
--- /dev/null
+++ b/examples/gptoss/README.md
@@ -0,0 +1,153 @@
+# GPT-OSS Training Tutorial
+
+## Step 0: Install Dependencies
+
+### Using Megatron Bridge
+
+[Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)
+
+Megatron Bridge provides a quick and convenient way to convert HuggingFace checkpoints to the Megatron format used by Megatron-LM. Follow the instructions in the [Megatron-Bridge Installation](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/README.md#-installation) to run the nemo docker container and convert checkpoints (via mounted volumes - make sure that the huggingface cache location AND the megatron checkpoint locations are properly mounted, otherwise you may not be saving the converted model to disk correctly).
+
+Below is an example of how to use Megatron-Bridge inside the pytorch container to convert a HuggingFace model checkpoint to Megatron format.
+
+Reference: [Megatron-Bridge Dockerfile](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/docker/Dockerfile.ci)
+
+Inside the [pytorch container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) run the following commands to install Megatron-Bridge:
+```bash
+cd /opt
+git clone --recursive https://github.com/NVIDIA-NeMo/Megatron-Bridge.git
+cd Megatron-Bridge
+
+# Make sure submodules are initialized (for 3rdparty/Megatron-LM)
+git submodule update --init --recursive
+
+export PATH="/root/.local/bin:$PATH"
+export UV_PROJECT_ENVIRONMENT=/opt/venv
+export VIRTUAL_ENV=/opt/venv
+export PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH"
+export UV_LINK_MODE=copy
+export UV_VERSION="0.7.2"
+
+# Install UV
+curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh
+
+# Create virtual environment and build the package
+uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages
+
+uv sync --locked --only-group build
+uv sync --locked --link-mode copy --all-extras --all-groups
+
+uv pip install --no-deps -e .
+
+source ${UV_PROJECT_ENVIRONMENT}/bin/activate
+```
+
+### Setup Environment
+
+```bash
+export HOST_MEGATRON_LM_DIR="/path/to/your/host/megatron-lm"
+git clone https://github.com/NVIDIA/Megatron-LM.git "$HOST_MEGATRON_LM_DIR"
+cd "$HOST_MEGATRON_LM_DIR"
+```
+
+```bash
+export HF_TOKEN={your_hf_token_here}
+```
+
+## Step 1: Convert HuggingFace to Megatron (Optional - skip if you already have a Megatron checkpoint)
+
+Set `--nproc-per-node` to be the number of GPUs per node. Set `hf_model_name` to be the Huggingface model e.g. `openai/gpt-oss-20b`
+
+```bash
+python3 -m torch.distributed.launch --nproc-per-node=8 examples/gptoss/01_convert_from_hf.py --hf-model openai/gpt-oss-20b
+```
+
+## Step 2: Train from Scratch
+
+To train from scratch first follow the steps below to setup the environment appropriately before running the training script in docker. Even though we are running the same container as before, it is better to restart the container to ensure a clean environment and that all environment and docker variables are set correctly. For the following example we used 8x GB300, but you should change the number of GPUs and nodes as needed.
+
+### Setup Environment
+
+```bash
+# Change these based on model and directory from previous conversion step
+export MODEL_DIR_NAME="openai_gpt-oss_20b"
+
+export HOST_CHECKPOINT_PATH="./megatron_checkpoints/${MODEL_DIR_NAME}"
+export HOST_TENSORBOARD_LOGS_PATH="./tensorboard_logs/${MODEL_DIR_NAME}"
+```
+
+By default we will use mock data to train the model in the example below. To use your own data, set the following environment variables:
+
+```bash
+# Optional: For real data
+export HOST_TOKENIZER_MODEL_PATH="/path/to/host/tokenizer.model"
+export HOST_DATA_PREFIX="/path/to/host/mydata_prefix"
+```
+
+### Setup Training Configurations
+
+Run the following to create a `distributed_config.env` file with the appropriate distributed training configurations. Change the values as needed for your setup. This file will override the default values in `02_train.sh`.
+
+```bash
+cat > ./distributed_config.env << 'EOF'
+GPUS_PER_NODE=8
+NUM_NODES=1
+MASTER_ADDR=localhost
+MASTER_PORT=6000
+NODE_RANK=0
+EOF
+```
+
+### Run Container with Mounted Volumes
+
+**NOTE:** This container runs the example training script `02_train.sh` located in the `examples/gptoss` directory. By default, we have only set pipeline parallelism to be the number of GPUs. Adjust TP_SIZE, EP_SIZE, PP_SIZE, etc. in `02_train.sh`. You can also adjust modify `--hidden-size`, `--ffn-hidden-size`, `--num-attention-heads`, `NUM_LAYERS`, etc.
+
+To train using mock data, run the following command:
+```bash
+PYTORCH_IMAGE="nvcr.io/nvidia/pytorch:25.12-py3"
+
+docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \
+ -v "${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm" \
+ -v "${HOST_CHECKPOINT_PATH}:/workspace/checkpoints" \
+ -v "${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs" \
+ -v "./distributed_config.env:/workspace/megatron-lm/examples/gptoss/distributed_config.env" \
+ --workdir /workspace/megatron-lm \
+ $PYTORCH_IMAGE \
+ bash examples/gptoss/02_train.sh \
+ --checkpoint-path /workspace/checkpoints \
+ --tensorboard-logs-path /workspace/tensorboard_logs \
+ --distributed-config-file /workspace/megatron-lm/examples/gptoss/distributed_config.env \
+ 2>&1 | tee "${HOST_TENSORBOARD_LOGS_PATH}/training_mock_$(date +'%y-%m-%d_%H-%M-%S').log"
+```
+**Note:** If you run into issues generating mock data one solution might be to reduce the number of GPUs to 1 and try to generate the data again.
+
+If using real data with with the `HOST_TOKENIZER_MODEL_PATH` and `HOST_DATA_PREFIX` environment variables set, run the following command instead:
+
+```bash
+PYTORCH_IMAGE="nvcr.io/nvidia/pytorch:25.12-py3"
+
+docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \
+ -v "${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm" \
+ -v "${HOST_CHECKPOINT_PATH}:/workspace/checkpoints" \
+ -v "${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs" \
+ -v "${HOST_TOKENIZER_MODEL_PATH}:/workspace/tokenizer_model" \
+ -v "$(dirname "${HOST_DATA_PREFIX}"):/workspace/data_dir" \
+ -v "./distributed_config.env:/workspace/megatron-lm/examples/gptoss/distributed_config.env" \
+ --workdir /workspace/megatron-lm \
+ $PYTORCH_IMAGE \
+ bash examples/gptoss/02_train.sh \
+ --checkpoint-path /workspace/checkpoints \
+ --tensorboard-logs-path /workspace/tensorboard_logs \
+ --tokenizer /workspace/tokenizer_model \
+ --data "/workspace/data_dir/$(basename "${HOST_DATA_PREFIX}")" \
+ --distributed-config-file /workspace/megatron-lm/examples/gptoss/distributed_config.env \
+ 2>&1 | tee "${HOST_TENSORBOARD_LOGS_PATH}/training_custom_$(date +'%y-%m-%d_%H-%M-%S').log"
+```
+
+## Step 3: Convert Megatron to HuggingFace
+
+Just run the following command to change from the megatron checkpoint from training to the huggingface format to share with others (make sure you have the same virtual environment setup as in Step 0):
+
+```bash
+python3 -m torch.distributed.launch --nproc-per-node=8 examples/gptoss/03_convert_to_hf.py --hf-model openai/gpt-oss-20b --megatron-model ./megatron_checkpoints/openai_gpt-oss_20b
+```
\ No newline at end of file
diff --git a/examples/inference/README.md b/examples/inference/README.md
new file mode 100644
index 00000000000..3259bf7f943
--- /dev/null
+++ b/examples/inference/README.md
@@ -0,0 +1,288 @@
+### Megatron Core Inference Documentation
+This guide provides an example for Megatron Core for running model inference.
+
+### Contents
+- [Megatron Core Inference Documentation](#megatron-core-inference-documentation)
+- [Contents](#contents)
+ - [1. Quick Start](#1-quick-start)
+ - [1.1 Understanding The Code](#11-understanding-the-code)
+ - [1.2 Running The Code](#12-running-the-code)
+ - [2. Flow of Control In MCore Backend](#2-flow-of-control-in-mcore-backend)
+ - [3. Customizing The Inference Pipeline](#3-customizing-the-inference-pipeline)
+ - [3.1. Create Your Own Inference Backend](#31-create-your-own-inference-backend)
+ - [3.2. Create Your Own Text Generation Controller](#32-create-your-own-text-generation-controller)
+ - [3.3. Support Other Models](#33-support-other-models)
+ - [3.3. Modify Inference Parameters](#33-modify-inference-parameters)
+ - [4. Future work](#4-future-work)
+
+
+
+#### 1. Quickstart
+This 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).
+
+
+
+##### 1.1 Code Walkthrough
+***STEP 1 - Initialize model parallel and other default arguments***
+The micro batch size defaults to 1. It is not used in tensor-parallelism only, and for pipeline-parallel models it is calculated at runtime.
+```python
+# Initialize Megatron model using the same model provider from training.
+ initialize_megatron(
+ args_defaults={'no_load_rng': True, 'no_load_optim': True, 'micro_batch_size': 1}
+ )
+```
+
+***STEP 2 - Load the model using the model_provider_function***
+The model provider function supports both MCore and Legacy models.
+
+```python
+ # Load the model checkpoint
+ model = get_model(model_provider, wrap_with_ddp=False)
+ load_checkpoint(model, None, None)
+ model.eval()
+ model = model[0]
+```
+
+***STEP 3 - Choose an engine***
+Text 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.
+```python
+ # Create an inference wrapper to setup the model.
+ inference_wrapped_model = GPTInferenceWrapper(model, args)
+
+ # Define a sampling loop.
+ text_generation_controller = TextGenerationController(
+ inference_wrapped_model=inference_wrapped_model,
+ tokenizer=tokenizer
+ )
+
+ # Create a static or dynamic inference engine.
+ inference_engine = StaticInferenceEngine(
+ text_generation_controller=text_generation_controller,
+ max_batch_size=args.max_batch_size
+)
+```
+
+***STEP 4 - Run text generation***
+The [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).
+```python
+ results: List[InferenceRequest] = inference_engine.generate(
+ prompts=args.prompts, sampling_params=sampling_params
+ )
+
+ if torch.distributed.get_rank() == 0:
+ for idx, result in enumerate(results):
+ print(f' ------------- RESULT FOR PROMPT {idx} --------------- ')
+ result = {
+ 'id': result.request_id,
+ 'input_prompt': result.prompt,
+ 'generated_text': result.generated_text,
+ 'generated_tokens' : result.generated_tokens
+ }
+ print(result)
+```
+
+
+
+##### 1.2 Running The Code
+An example Slurm script is shown below. Set the tokenizer paths, inference params, and other settings appropriately.
+
+For 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).
+
+```
+# Slurm cluster settings
+ACCOUNT=
+MLM_PATH=/path/to/megatron-lm
+GPT_CKPT=/path/to/gpt/ckpt
+VOCAB_MERGE_FILE_PATH=/path/to/vocab/and/merge/file
+CONTAINER_IMAGE=nvcr.io/ea-bignlp/ga-participants/nemofw-training:23.11
+
+srun --account $ACCOUNT \
+--job-name=$ACCOUNT:inference \
+--partition=batch \
+--time=01:00:00 \
+--container-image $CONTAINER_IMAGE \
+--container-mounts $MLM_PATH:/workspace/megatron-lm/,$GPT_CKPT:/workspace/mcore_gpt_ckpt,$VOCAB_MERGE_FILE_PATH:/workspace/tokenizer \
+--no-container-mount-home \
+--pty /bin/bash \
+
+# Inside the container run the following.
+
+cd megatron-lm/
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+TOKENIZER_ARGS=(
+ --vocab-file /workspace/tokenizer/gpt2-vocab.json
+ --merge-file /workspace/tokenizer/gpt2-merges.txt
+ --tokenizer-type GPT2BPETokenizer
+)
+
+MODEL_ARGS=(
+ --use-checkpoint-args
+ --use-mcore-models
+ --load /workspace/mcore_gpt_ckpt
+)
+
+INFERENCE_SPECIFIC_ARGS=(
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --num-tokens-to-generate 20
+ --max-batch-size 4
+)
+
+torchrun --nproc-per-node=4 examples/inference/gpt/gpt_static_inference.py \
+ ${TOKENIZER_ARGS[@]} \
+ ${MODEL_ARGS[@]} \
+ ${INFERENCE_SPECIFIC_ARGS[@]} \
+ --prompts "prompt one " "sample prompt two" "sample prompt 3"
+
+NOTE: Other parameters which can be customized for inference:
+--temperature (Sampling temperature)
+--top_k (top_k sampling)
+--top_p (top_p sampling)
+--num-tokens-to-generate (Number of tokens to generate for each prompt)
+--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.')
+--use-dist-ckpt (If using dist checkpoint format for the model)
+
+```
+
+
+
+
+
+#### 2. Control Flow in the MCore Backend
+An example of inference with static batching is provided in [gpt_static_inference.py](./gpt/gpt_static_inference.py).
+* [mcore_engine](../../megatron/core/inference/engines/mcore_engine.py) **generate()** function is called with the input prompts.
+* 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.
+* The engine will run until all requests (waiting + active) are completed.
+ * The active requests are passed into **generate_all_output_tokens_static_batch()** of the text generation controller .
+ * 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
+ * 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
+ * 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
+ * Output logits are synchronized across all pipeline parallel ranks
+ * The text generation controller obtains the log probabilities and samples tokens based on the strategy defined in the sampling parameters.
+ * The sampled tokens are then appended to the input prompt tokens for the next iteration
+ * The **update_generation_status()** method of the text generation controller checks which prompts have finished generating or hit a stop condition
+ * After the inference loop, the result is detokenized and stored as an attribute of the InferenceRequest. These requests are marked as completed.
+ * The **update_requests_pool()** method of the scheduler moves completed requests into the completed request pool and waiting requests into the active request pool
+
+
+
+#### 3. Customizing The Inference Pipeline
+
+The inference pipeline supports three levels of customization:
+
+* **Inference engine** - The MCore Engine supports static and dynamic batching. Modify this to add a new backend.
+* **Text generation controller** - The main sampling loop. Customize this to support alternative tokenization or implement a new sampling strategy.
+* **Inference Wrapped Model** - Change this to support a new model.
+* **Modify Inference Parameters** - Change this to update top_p, top_k, number of tokens to be generated, temperature, and other sampling parameters.
+
+
+
+##### 3.1. Create Your Own Inference Backend
+The [abstract_engine.py](./../../megatron/core/inference/engine/abstract_engine.py) file contains a `generate` method that can be extended to support a new backend.
+
+```python
+class AbstractEngine(ABC):
+ @staticmethod
+ def generate(self) -> dict:
+ """The abstract backend's generate function.
+
+ To define a new backend, implement this method and return the outputs as a dictionary.
+```
+
+
+
+##### 3.2. Implement a new Sampling Loop
+
+The [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.
+
+``` python
+class TextGenerationController:
+
+ def tokenize_prompt(self, prompt: str) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Utility to tokenize the input prompts"""
+
+ def sample_from_logits(
+ self,
+ last_token_logits: torch.Tensor,
+ sampling_params: SamplingParams,
+ vocab_size: int,
+ generation_started : Optional[torch.Tensor] = None,
+ top_n_logprobs_dict: Dict[int, List[Dict[str, float]]] = None,
+ ) -> torch.Tensor:
+ """Samples the logits to generate outputs
+
+ 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
+ at each step it also updates the top_n_logprobs_dict.
+ """
+
+ def update_generation_status(
+ self,
+ updated_prompts_tokens: torch.Tensor,
+ generation_started: torch.Tensor,
+ current_context_end_position: int,
+ is_generation_done_tensor: torch.Tensor,
+ generated_sequence_lengths: torch.Tensor,
+ ) -> torch.Tensor:
+ """Function to check which prompts have reached an end condition
+
+ 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
+ """
+
+ def generate_all_output_tokens_static_batch(
+ self, active_requests: OrderedDict[int, InferenceRequest],
+ ) -> OrderedDict[int, InferenceRequest]:
+ """Utility to generate all the output tokens and probabilities for the prompts .
+
+ 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
+ """
+
+ def detokenize_generations(self, prompt_tokens_with_generated_tokens: torch.Tensor) -> str:
+ """Detokenize the output generations"""
+```
+
+
+
+##### 3.3. Support Other Models
+Extend [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:
+* Forward method which calls the model `forward` method depending on model parallel settings
+* Initializes the model and puts it in `.eval()` mode
+* Setup for the input parameters (max batch size, max seq length)
+
+The following methods should be implemented:
+```python
+class AbstractModelInferenceWrapper:
+ def prep_model_for_inference(self, prompts_tokens: torch.Tensor):
+ """A utility function for preparing model for inference
+
+ 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
+ """
+
+ @abc.abstractclassmethod
+ def get_batch_for_context_window(self) -> List:
+ """Returns the input data for inference
+
+ 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.
+```
+
+Refer to [gpt_inference_wrapper.py](../../megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py) for an example of implementing this for GPTModel.
+
+
+
+##### 3.3. Modify Inference Parameters
+We 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.
+
+```
+from megatron.core.inference.sampling_params import SamplingParams
+
+c = SamplingParams(temperature=0.5)
+c.add_attributes({'min_length':4, 'eod_id':153})
+```
+
+
+
+#### 4. Future work
+The following features are planned for future releases.
+* TRTLLM Engine support
+* Continuous batching optimizations
+* Speculative decoding
\ No newline at end of file
diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
new file mode 100644
index 00000000000..02a257c1b46
--- /dev/null
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -0,0 +1,523 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+# pylint: disable=bad-builtin
+
+import hashlib
+import io
+import json
+import os
+import sys
+import warnings
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+from megatron.training.arguments import parse_and_validate_args
+import torch
+from tqdm import tqdm
+
+sys.path.append(
+ os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
+)
+
+from examples.inference.gpt.utils import (
+ Request,
+ build_dynamic_engine_setup_prefix,
+ build_requests,
+ get_curr_time,
+ get_global_peak_memory_stats_bytes,
+)
+from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
+from megatron.core.inference.engines import DynamicInferenceEngine, EngineSuspendedError
+from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
+ GPTInferenceWrapper,
+)
+from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.inference.text_generation_controllers.text_generation_controller import (
+ TextGenerationController,
+)
+from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
+from megatron.inference.utils import (
+ add_inference_args,
+ get_inference_config_from_model_and_args,
+ get_model_for_inference,
+)
+
+sys.path.append(
+ os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
+)
+import logging
+
+import megatron
+from megatron.core.utils import configure_nvtx_profiling
+from megatron.training import get_args, get_tokenizer, initialize_megatron
+
+torch.serialization.add_safe_globals([io.BytesIO])
+torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState])
+torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic])
+
+
+def run_inference(
+ requests: List[Request],
+ engine: DynamicInferenceEngine,
+ sampling_params: Optional[SamplingParams] = None,
+) -> List[Dict[str, float]]:
+ """Add requests to engine and generate tokens.
+
+ Args:
+ requests (List[Request]): Requests that are to be added and processed.
+ engine (DynamicInferenceEngine): Inference engine that manages generating tokens.
+ sampling_params (SamplingParams): Deprecated as of megatron-core 0.16.
+
+ Return:
+ A dictionary of step times with `prefill` and `decode` keys.
+ """
+
+ if sampling_params is not None and torch.distributed.get_rank() == 0:
+ warnings.warn(
+ "The `sampling_params` argument is deprecated. "
+ "Sampling parameters are specified per request.",
+ DeprecationWarning,
+ )
+
+ args = get_args()
+
+ # Parse batch boundaries for batch-drain mode.
+ batch_ranges = None
+ if args.drain_between_batches and args.batch_boundaries:
+ boundaries = [int(x) for x in args.batch_boundaries.split(",")]
+ num_requests_total = len(requests)
+ batch_ranges = []
+ for i, start in enumerate(boundaries):
+ end = boundaries[i + 1] if i + 1 < len(boundaries) else num_requests_total
+ batch_ranges.append((start, end))
+
+ # Initialize request arrival times.
+ base_arrival_time = get_curr_time()
+ for request in requests:
+ request.time_arrival = request.time_offset + base_arrival_time
+
+ # Add and process requests.
+ num_requests_total = len(requests)
+ num_requests_added = 0
+ num_requests_finished = 0
+ step_times = {"prefill": [], "decode": []}
+ add_times = []
+ output_times = []
+ tbar = tqdm(total=num_requests_total)
+ total_output_tokens = 0
+ attempted_step_count = 0
+ if args.cuda_graph_impl == "local":
+ cuda_graph_request_count_map = {}
+ else:
+ cuda_graph_request_count_map = None
+
+ def _add_request():
+ """Add request to engine.
+
+ *Note: Using `prompt_text` instead of `prompt_tokens` for fair comparison.
+ """
+ nonlocal num_requests_added
+ _request = requests[num_requests_added]
+ engine.add_request(num_requests_added, _request.prompt_text, _request.sampling_params)
+ _request.time_start = get_curr_time()
+ _request.state = "started"
+ num_requests_added += 1
+ tbar.update(1)
+
+ def _process_step_result(result):
+ """Process a single engine step result, updating bookkeeping state."""
+ nonlocal total_output_tokens, num_requests_finished
+
+ is_decode_only = engine.is_decode_only
+
+ # Record cuda_graph_request_count.
+ cuda_graph_request_count = result["cuda_graph_request_count"]
+ if args.cuda_graph_impl == "local" and cuda_graph_request_count is not None:
+ cuda_graph_request_count_map[cuda_graph_request_count] = (
+ cuda_graph_request_count_map.get(cuda_graph_request_count, 0) + 1
+ )
+
+ # Update requests.
+ active_request_ids = result["active_request_ids"]
+ finished_request_records = result["finished_request_records"]
+ step_time = result["step_time"]
+ if len(active_request_ids) > 0 or len(finished_request_records) > 0:
+ if is_decode_only:
+ step_times["decode"].append(step_time)
+ else:
+ step_times["prefill"].append(step_time)
+
+ # Append output tokens.
+ output_start = get_curr_time()
+ for finished_request_record in finished_request_records:
+
+ finished_request = finished_request_record.merge()
+
+ # Update local request object.
+ request = requests[finished_request.request_id]
+ request.time_end = get_curr_time()
+ request.state = "finished"
+ request.request_id = finished_request.request_id
+ request.events = finished_request.events
+
+ request.ttft = finished_request.ttft
+
+ # Update prompt, in case engine has been suspended and resumed.
+ request.prompt_tokens = finished_request.prompt_tokens.tolist()
+ request.prompt_text = finished_request.prompt
+
+ # Get output tokens and text.
+ request.output_tokens = finished_request.generated_tokens
+ request.output_text = finished_request.generated_text
+ total_output_tokens += len(request.output_tokens)
+
+ # Log probs.
+ if finished_request.sampling_params.return_log_probs:
+ if not finished_request.prompt_log_probs:
+ finished_request.prompt_log_probs = []
+ request.prompt_log_probs = finished_request.prompt_log_probs
+ request.generated_log_probs = finished_request.generated_log_probs
+ request.logprobs = (
+ finished_request.prompt_log_probs + finished_request.generated_log_probs
+ )
+ if finished_request.sampling_params.top_n_logprobs > 0:
+ request.generated_top_n_logprobs = finished_request.generated_top_n_logprobs
+ if not finished_request.sampling_params.skip_prompt_log_probs:
+ request.prompt_top_n_logprobs = finished_request.prompt_top_n_logprobs
+ num_requests_finished += 1
+ output_times.append(get_curr_time() - output_start)
+
+ if batch_ranges is not None:
+ # Batch-drain mode: add all requests in a batch, drain, then next batch.
+ for batch_idx, (batch_start, batch_end) in enumerate(batch_ranges):
+ # Add all requests in current batch.
+ add_start = get_curr_time()
+ while num_requests_added < batch_end:
+ _add_request()
+ add_times.append(get_curr_time() - add_start)
+
+ # Step until all active requests finish (drain).
+ while engine.has_unfinished_requests():
+ try:
+ result = engine.step_modern()
+ except EngineSuspendedError as e:
+ result = e
+ attempted_step_count += 1
+
+ if isinstance(result, EngineSuspendedError):
+ continue
+
+ _process_step_result(result)
+ else:
+ # Original mode: add requests per step based on arrival time or count.
+ while True:
+ # Add requests.
+ add_start = get_curr_time()
+ if args.incoming_requests_per_step is None:
+ # Add requests with 'earlier' arrival time.
+ while num_requests_added < num_requests_total:
+ if requests[num_requests_added].time_arrival > add_start:
+ break
+ _add_request()
+ else:
+ # Add deterministic number of requests (generally used for debugging).
+ for i in range(
+ min(args.incoming_requests_per_step, num_requests_total - num_requests_added)
+ ):
+ _add_request()
+ add_times.append(get_curr_time() - add_start)
+
+ # Step inference engine (i.e., generate a token for each active request).
+ # Before step, we haven't done the scheduling, so we cannot know the is_decode_only
+ try:
+ result = engine.step_modern()
+ except EngineSuspendedError as e:
+ result = e
+ pass # ignore error in order to call 'engine.resume()' below.
+ attempted_step_count += 1
+
+ # Test suspending and resuming engine.
+ if args.suspend_resume_interval is not None:
+
+ # Suspend.
+ if attempted_step_count % args.suspend_resume_interval == 0:
+ print("**** step %d/%d ... suspend." % (engine.context.step_count, attempted_step_count))
+ engine.suspend()
+
+ # Resume, 0+ attempted steps later.
+ if (
+ attempted_step_count > 0
+ and (attempted_step_count - args.suspend_resume_interval // 2)
+ % args.suspend_resume_interval
+ == 0
+ ):
+ print("**** step %d/%d ... resume." % (engine.context.step_count, attempted_step_count))
+ engine.resume()
+
+ # If engine suspended, continue to next iter.
+ if isinstance(result, EngineSuspendedError):
+ continue
+
+ _process_step_result(result)
+
+ # Check if all requests are finished.
+ if not (engine.has_unfinished_requests() or num_requests_added < num_requests_total):
+ break
+
+ # Resume engine (NOOP if not suspended).
+ engine.resume()
+
+ return {
+ "step_times": step_times,
+ "add_times": add_times,
+ "output_times": output_times,
+ "total_output_tokens": total_output_tokens,
+ "cuda_graph_request_count_map": cuda_graph_request_count_map,
+ }
+
+
+@torch.inference_mode()
+def main():
+ """Run dynamic inference."""
+ # Initialize Megatron.
+ args = parse_and_validate_args(
+ extra_args_provider=add_inference_args,
+ args_defaults={'no_load_rng': True, 'no_load_optim': True},
+ )
+ initialize_megatron()
+
+ # Start Nsight profiler.
+ if os.environ.get("NSIGHT_PREFIX"):
+ torch.cuda.cudart().cudaProfilerStart()
+
+ level_str = os.getenv("LOG_LEVEL", "INFO").upper()
+ level = getattr(logging, level_str, logging.INFO)
+ logging.basicConfig(level=level, force=True)
+
+ configure_nvtx_profiling(True)
+
+ # Build tokenizer
+ tokenizer = build_tokenizer(args)
+
+ # Reset peak memory stats so functional tests measure this run and not
+ # whatever happened earlier during initialization.
+ torch.cuda.reset_peak_memory_stats()
+
+ # Sampling params.
+ sampling_params = SamplingParams(
+ temperature=args.temperature,
+ top_k=args.top_k,
+ top_p=args.top_p,
+ skip_prompt_log_probs=args.skip_prompt_log_probs,
+ return_log_probs=args.return_log_probs,
+ num_tokens_to_generate=args.num_tokens_to_generate,
+ termination_id=args.termination_id if args.termination_id is not None else tokenizer.eod,
+ top_n_logprobs=args.top_n_logprobs,
+ stop_words=args.stop_words,
+ )
+
+ model = get_model_for_inference()
+
+ # Requests, context, controller.
+ requests = build_requests(args, tokenizer, sampling_params)
+ inference_config = get_inference_config_from_model_and_args(model, args)
+
+ # Calculate max_sequence_length from requests
+ max_gen_length = sampling_params.num_tokens_to_generate
+ max_context_length = max(len(r.prompt_tokens) for r in requests)
+ inference_config.max_sequence_length = max_context_length + max_gen_length
+ context = DynamicInferenceContext(model.config, inference_config)
+ wrapped_model = GPTInferenceWrapper(model, context)
+ controller = TextGenerationController(wrapped_model, tokenizer)
+
+ # Validate all context_length's <= max_tokens.
+ if not args.enable_chunked_prefill:
+ invalid_prompt_length_map = {}
+ for request_idx, request in enumerate(requests):
+ if len(request.prompt_tokens) > context.max_tokens:
+ invalid_prompt_length_map[request_idx] = len(request.prompt_tokens)
+ assert (
+ not invalid_prompt_length_map
+ ), "request idxs with prompts longer than context.max_tokens: " ", ".join(
+ f"{k}({v})" for k, v in invalid_prompt_length_map.items()
+ )
+
+ # Inference engine.
+ engine = DynamicInferenceEngine(controller, context)
+
+ setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests)
+ print("~~~")
+ print(setup_prefix)
+ print("~~~")
+
+ # Run and time test, optionally `args.inference_repeat_n` times.
+ throughputs = []
+ for _ in range(args.inference_repeat_n):
+
+ # Reset engine.
+ engine.reset()
+
+ torch.cuda.reset_peak_memory_stats()
+
+ # Trial.
+ t = get_curr_time()
+ result = run_inference(requests, engine)
+ step_times = result["step_times"]
+ add_times = result["add_times"]
+ output_times = result["output_times"]
+ total_output_tokens = result["total_output_tokens"]
+ torch.cuda.synchronize()
+ total_time = get_curr_time() - t
+ stats = torch.cuda.memory_stats()
+ throughput = total_output_tokens / total_time
+ throughputs.append(throughput)
+
+ # Validate all requests finished.
+ for request in requests:
+ assert request.state == "finished", f"request.state == '{request.state}' != 'finished'."
+
+ peak_mem_stats = get_global_peak_memory_stats_bytes()
+
+ # Print unique prompts + outputs.
+ if torch.distributed.get_rank() == 0:
+
+ def escape_str(s):
+ return s.replace("\n", "\\n")
+
+ print("~~~~ Unique prompts + outputs. ~~~~")
+
+ # Map requests by their prompt.
+ unique_prompt_map = defaultdict(list)
+ for request_idx, request in enumerate(requests):
+ unique_prompt_map[request.prompt_text].append(request_idx)
+
+ # Print unique prompts + outputs.
+ text_hashes = []
+ for unique_idx, (prompt_text, request_idxs) in enumerate(unique_prompt_map.items()):
+
+ # ---- Prompt summary line ----
+ prompt_len = len(requests[request_idxs[0]].prompt_tokens)
+ escaped_prompt_text = escape_str(prompt_text)
+ print(
+ f"\n{unique_idx+1}/{len(unique_prompt_map)}"
+ f"[n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}"
+ )
+
+ # ---- Group all outputs for this prompt ----
+ output_map = defaultdict(list)
+ for idx in request_idxs:
+ req = requests[idx]
+ output_map[req.output_text].append(idx)
+
+ # ---- Print each unique output ----
+ for output_text, output_request_idxs in output_map.items():
+ evicted = False
+ for idx in output_request_idxs:
+ for event in requests[idx].events:
+ if event.type.name == "EVICT":
+ evicted = True
+ break
+ if output_text is not None:
+ # Use hash of prompt + generated text in case engine was
+ # suspended and resumed, which misaligns boundary between
+ # prompt and generated tokens.
+ o_hash = hashlib.sha256((prompt_text + output_text).encode()).hexdigest()[:6]
+ o_len = len(requests[output_request_idxs[0]].output_tokens)
+ escaped_output_text = escape_str(output_text)
+ else:
+ o_hash = "--"
+ o_len = 0
+ escaped_output_text = "--"
+ print(
+ f" >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}"
+ f"{', ' if evicted else ''}] {escaped_output_text}"
+ )
+ text_hashes.append(o_hash)
+
+ # Write results to JSON. Primarily used for functional testing.
+ if args.output_path:
+ json_results = {}
+
+ # Write every 'n' requests, plus the final request.
+ for i, req in enumerate(requests):
+ if i % args.output_every_n_results == 0 or i == len(requests) - 1:
+ print(f' Attributes of request {i}: {req.__dict__}')
+ result_dict = {
+ "input_prompt": req.prompt_text,
+ "generated_text": req.output_text,
+ "generated_tokens": req.output_tokens,
+ "latency": req.time_end - req.time_start,
+ "ttft": req.ttft, # Time-to-first-token in seconds
+ "cuda_graph_request_count_map": result["cuda_graph_request_count_map"],
+ "step_count": engine.context.step_count,
+ "top_n_logprobs": getattr(req, 'generated_top_n_logprobs', None),
+ "prompt_top_n_logprobs": getattr(req, 'prompt_top_n_logprobs', None),
+ }
+ if req.sampling_params.return_log_probs:
+ result_dict["prompt_logprobs"] = getattr(req, 'prompt_log_probs', None)
+ result_dict["generated_logprobs"] = getattr(
+ req, 'generated_log_probs', None
+ )
+ result_dict["logprobs"] = getattr(req, 'logprobs', None)
+ if args.output_request_events:
+ result_dict["events"] = [e.serialize() for e in req.events]
+ json_results[req.request_id] = result_dict
+
+ # Track system-level throughput as a test / debug metric
+ if args.record_throughput:
+ json_results["throughput"] = throughputs
+ # Attach peak memory metrics; the functional test only validates these
+ # if the fields exist in the golden values.
+ json_results.update(peak_mem_stats)
+ json_results["lifetime_prefill_token_count"] = engine.context.lifetime_prefill_token_count
+
+ print(f' Saving results to {args.output_path}')
+ with open(args.output_path, "w") as fp:
+ json.dump(json_results, fp, indent=1)
+
+ # Timing results.
+ stats = torch.cuda.memory_stats()
+ throughput = total_output_tokens / total_time
+ print("~~~")
+ peak_alloc_gb = stats["allocated_bytes.all.peak"] / 1024**3
+ peak_resvd_gb = stats["reserved_bytes.all.peak"] / 1024**3
+
+ p_times = step_times["prefill"]
+ d_times = step_times["decode"]
+
+ p_total = sum(p_times)
+ d_total = sum(d_times)
+
+ p_count = len(p_times)
+ d_count = len(d_times)
+
+ p_mean = p_total / p_count
+ d_mean = d_total / d_count if d_count != 0 else 0.0
+
+ # Commented out for now as the step/add/output times are not calculated correctly.
+ # print(
+ # f"{setup_prefix} … "
+ # f"mem {peak_alloc_gb:.1f}/{peak_resvd_gb:.1f} GB … "
+ # f"total time: {step_total:.3f}s … "
+ # f"step time: total {step_total:.3f}s "
+ # f"[ p {p_total:.3f}s, d {d_total:.3f}s ], "
+ # f"mean [ p {p_mean:.3f}s, d {d_mean:.3f}s ], "
+ # f"count [ p {p_count}, d {d_count} ]."
+ # )
+ capture_str = f"{engine.capture_stats['time']:.2f} sec" if engine.capture_stats else "--"
+ print(
+ f"{setup_prefix} … " f"throughput: {throughput:.3f} tok/s … ",
+ f"total time: {total_time:.3f}s … "
+ f"mem {peak_alloc_gb:.1f}/{peak_resvd_gb:.1f} GB … "
+ f"steps: {engine.context.step_count:d} … "
+ f"capture {capture_str}",
+ )
+ print("~~~")
+
+ # Stop Nsight profiler.
+ if os.environ.get("NSIGHT_PREFIX"):
+ torch.cuda.cudart().cudaProfilerStop()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/inference/gpt/gpt_dynamic_inference_12b.sh b/examples/inference/gpt/gpt_dynamic_inference_12b.sh
new file mode 100644
index 00000000000..ca21bb170a5
--- /dev/null
+++ b/examples/inference/gpt/gpt_dynamic_inference_12b.sh
@@ -0,0 +1,127 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+# Run dynamic batching inference on the 12B GPT model.
+
+set -u
+
+# Libraries.
+pip install simpy
+pip install sentencepiece
+pip install tiktoken
+
+# Environment variables.
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+# Checkpoint.
+: ${CHECKPOINT_DIR:?"CHECKPOINT_DIR is not set"}
+: ${TOKENIZER_MODEL:?"TOKENIZER_MODEL is not set"}
+
+# Prompts.
+: ${NUM_TOKENS_TO_PROMPT="8 32"}
+: ${NUM_TOKENS_TO_GENERATE=256}
+: ${INCOMING_REQUESTS_DURATION=10.}
+: ${INCOMING_REQUESTS_PER_SEC=100.}
+
+# Dynamic context.
+: ${BUFFER_SIZE_GB=50.}
+
+# Cuda graphs.
+: ${NUM_CUDA_GRAPHS=16}
+
+# Miscellaneous.
+: ${USE_COORDINATOR=0}
+: ${ENGINE=dynamic}
+: ${EXTRA_ARGS=""}
+# NSIGHT_PREFIX=/path/to/nsight/profile
+
+# Arguments.
+ARGS=" \
+ --no-persist-layer-norm \
+ --apply-layernorm-1p \
+ --no-position-embedding \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --load ${CHECKPOINT_DIR} \
+ --use-checkpoint-args \
+ --untie-embeddings-and-output-weights \
+ --disable-bias-linear \
+ --use-rotary-position-embeddings \
+ --position-embedding-type rope \
+ --rotary-base 1000000 \
+ --rotary-percent 1.0 \
+ --swiglu \
+ --normalization RMSNorm \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --exit-duration-in-mins 5740 \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 40 \
+ --hidden-size 5120 \
+ --ffn-hidden-size 14336 \
+ --num-attention-heads 32 \
+ --kv-channels 128 \
+ --seq-length 1024 \
+ --max-position-embeddings 1024 \
+ --micro-batch-size 64 \
+ --bf16 \
+ --tokenizer-type TikTokenizer \
+ --tiktoken-pattern v2 \
+ --tokenizer-model ${TOKENIZER_MODEL} \
+ --distributed-timeout-minutes 2400 \
+ --use-flash-attn \
+ --inference-rng-tracker \
+ \
+ --inference-dynamic-batching \
+ --inference-dynamic-batching-buffer-size-gb ${BUFFER_SIZE_GB} \
+ \
+ ${EXTRA_ARGS} \
+"
+
+# Cuda graphs.
+if [ "${NUM_CUDA_GRAPHS}" != "0" ]; then
+ ARGS+=" \
+ --cuda-graph-impl local \
+ --inference-dynamic-batching-num-cuda-graphs ${NUM_CUDA_GRAPHS} \
+ "
+else
+ ARGS+=" \
+ --cuda-graph-impl none \
+ "
+fi
+
+# Prompts.
+if [[ -v PROMPTS ]]; then
+ ARGS+=" \
+ --prompts ${PROMPTS} \
+ --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \
+ "
+elif [[ -v PROMPT_FILE ]]; then
+ ARGS+=" \
+ --prompt-file ${PROMPT_FILE} \
+ --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \
+ "
+else
+ ARGS+=" \
+ --num-tokens-to-prompt ${NUM_TOKENS_TO_PROMPT} \
+ --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \
+ --incoming-requests-duration ${INCOMING_REQUESTS_DURATION} \
+ --incoming-requests-per-sec ${INCOMING_REQUESTS_PER_SEC} \
+ "
+fi
+
+# Command.
+if [[ "${USE_COORDINATOR}" == "0" ]]; then
+ CMD="python -m examples.inference.gpt.gpt_${ENGINE}_inference ${ARGS}"
+else
+ CMD="python -um examples.inference.gpt.gpt_${ENGINE}_inference_with_coordinator ${ARGS}"
+fi
+
+if [[ -v NSIGHT_PREFIX ]]; then
+ 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}"
+fi
+
+echo "~~~"
+echo "CMD ... ${CMD}."
+echo "~~~"
+eval ${CMD}
diff --git a/examples/inference/gpt/gpt_dynamic_inference_357m.sh b/examples/inference/gpt/gpt_dynamic_inference_357m.sh
new file mode 100644
index 00000000000..cc99bdddec1
--- /dev/null
+++ b/examples/inference/gpt/gpt_dynamic_inference_357m.sh
@@ -0,0 +1,115 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+# Run dynamic batching inference on the 357M GPT model.
+
+set -u
+
+# Libraries.
+pip install simpy
+pip install sentencepiece
+pip install tiktoken
+
+# Environment variables.
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+# Checkpoint.
+: ${CHECKPOINT_DIR:?"CHECKPOINT_DIR is not set"}
+: ${VOCAB_FILE:?"VOCAB_FILE is not set"}
+: ${MERGE_FILE:?"MERGE_FILE is not set"}
+
+# Prompts.
+: ${NUM_TOKENS_TO_PROMPT="8 32"}
+: ${NUM_TOKENS_TO_GENERATE=256}
+: ${INCOMING_REQUESTS_DURATION=10.}
+: ${INCOMING_REQUESTS_PER_SEC=100.}
+
+# Dynamic context.
+: ${BUFFER_SIZE_GB=50.}
+
+# Cuda graphs.
+: ${NUM_CUDA_GRAPHS=16}
+
+# Miscellaneous.
+: ${USE_COORDINATOR=0}
+: ${ENGINE=dynamic}
+: ${NPROC_PER_NODE=1}
+: ${EXTRA_ARGS=""}
+# NSIGHT_PREFIX=/path/to/nsight/profile
+
+# Arguments.
+ARGS=" \
+ --exit-on-missing-checkpoint \
+ --transformer-impl local \
+ --load ${CHECKPOINT_DIR} \
+ --tokenizer-type GPT2BPETokenizer \
+ --vocab-file ${VOCAB_FILE} \
+ --merge-file ${MERGE_FILE} \
+ --exit-on-missing-checkpoint \
+ --max-position-embeddings 2048 \
+ --seq-length 2048 \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 24 \
+ --num-attention-heads 16 \
+ --hidden-size 1024 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --seed 42 \
+ --use-flash-attn \
+ --inference-rng-tracker \
+ \
+ --inference-dynamic-batching \
+ --inference-dynamic-batching-buffer-size-gb ${BUFFER_SIZE_GB} \
+ \
+ ${EXTRA_ARGS} \
+"
+
+# Cuda graphs.
+if [ "${NUM_CUDA_GRAPHS}" != "0" ]; then
+ ARGS+=" \
+ --cuda-graph-impl local \
+ --inference-dynamic-batching-num-cuda-graphs ${NUM_CUDA_GRAPHS} \
+ "
+else
+ ARGS+=" \
+ --cuda-graph-impl none \
+ "
+fi
+
+# Prompts.
+if [[ -v PROMPTS ]]; then
+ ARGS+=" \
+ --prompts ${PROMPTS} \
+ --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \
+ "
+elif [[ -v PROMPT_FILE ]]; then
+ ARGS+=" \
+ --prompt-file ${PROMPT_FILE} \
+ --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \
+ "
+else
+ ARGS+=" \
+ --num-tokens-to-prompt ${NUM_TOKENS_TO_PROMPT} \
+ --num-tokens-to-generate ${NUM_TOKENS_TO_GENERATE} \
+ --incoming-requests-duration ${INCOMING_REQUESTS_DURATION} \
+ --incoming-requests-per-sec ${INCOMING_REQUESTS_PER_SEC} \
+ "
+fi
+
+# Command.
+if [[ "${USE_COORDINATOR}" == "0" ]]; then
+ CMD="python -m examples.inference.gpt.gpt_${ENGINE}_inference ${ARGS}"
+else
+ CMD="python -m torch.distributed.run --nproc-per-node ${NPROC_PER_NODE} -m examples.inference.gpt.gpt_${ENGINE}_inference_with_coordinator ${ARGS}"
+fi
+
+if [[ -v NSIGHT_PREFIX ]]; then
+ 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}"
+fi
+
+echo "~~~"
+echo "CMD ... ${CMD}."
+echo "~~~"
+eval ${CMD}
diff --git a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
new file mode 100644
index 00000000000..aa42f492ca4
--- /dev/null
+++ b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
@@ -0,0 +1,248 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import asyncio
+import json
+import logging
+import os
+import time
+import warnings
+from collections import defaultdict
+from typing import List
+
+from megatron.training.arguments import parse_and_validate_args
+import torch
+import torch.distributed as dist
+
+from examples.inference.gpt.utils import Request, build_dynamic_engine_setup_prefix, build_requests
+from megatron.core.inference.engines import DynamicInferenceEngine
+from megatron.core.inference.engines.dynamic_engine import EngineState
+from megatron.core.inference.inference_client import InferenceClient
+from megatron.core.inference.inference_request import DynamicInferenceRequestRecord
+from megatron.core.inference.sampling_params import SamplingParams
+from megatron.inference.utils import (
+ add_inference_args,
+ get_dynamic_inference_engine,
+ get_model_for_inference,
+)
+from megatron.training import get_args, get_tokenizer, initialize_megatron
+from megatron.core.utils import configure_nvtx_profiling
+
+# pylint: disable=line-too-long
+
+logging.basicConfig(level=logging.INFO, force=True)
+
+
+async def suspend_resume_cycle(client, engine, args, futures):
+ """Wait for all in-flight requests, then suspend/train/resume."""
+ await asyncio.gather(*futures)
+
+ client.pause_engines()
+ await engine.wait_until(EngineState.PAUSED)
+ client.suspend_engines()
+ await engine.wait_until(EngineState.SUSPENDED)
+ if args.suspend_timeout > 0:
+ await asyncio.sleep(args.suspend_timeout)
+ client.resume_engines()
+ await engine.wait_until(EngineState.RESUMED)
+ client.unpause_engines()
+ await engine.wait_until(EngineState.RUNNING)
+
+
+async def main(
+ engine: DynamicInferenceEngine,
+ requests: List[Request],
+ port: int | None = None,
+ sampling_params: SamplingParams | None = None,
+):
+ if sampling_params is not None:
+ warnings.warn(
+ "The `sampling_params` argument is deprecated. "
+ "Sampling parameters are specified per request.",
+ DeprecationWarning,
+ )
+
+ # once you call engine.start_listening_to_data_parallel_coordinator,
+ # the engine will start accepting requests from the data parallel coordinator.
+ # and processing them in an asyncio coroutine.
+ # leaving inference_coordinator_port as None will find a free port automatically.
+ args = get_args()
+
+ dp_addr = await engine.start_listening_to_data_parallel_coordinator(
+ inference_coordinator_port=port,
+ launch_inference_coordinator=True,
+ coordinator_schedule_output_path=args.coordinator_schedule_output_path,
+ )
+
+ # All ranks agree on the number of suspend/resume cycles from args.
+ num_suspend_resume_cycles = len(requests) // args.suspend_resume_interval if args.suspend_resume_interval else 0
+
+ # Create client and run example.
+ if dist.get_rank() == 0:
+ client = InferenceClient(dp_addr, deserialize=True) # submits requests to the inference coordinator
+ client.start()
+ base_arrival_time = time.time_ns() / 10**9
+ for request in requests:
+ request.time_arrival = request.time_offset + base_arrival_time
+ futures = []
+ num_requests_total = len(requests)
+ num_requests_added = 0
+ next_suspend_at = args.suspend_resume_interval or 0
+ cycles_done = 0
+
+ while True:
+ current_time = time.time_ns() / 10**9
+ if args.incoming_requests_per_step is None:
+ # Only add requests that have arrived at the current time.
+ while (
+ num_requests_added < num_requests_total
+ and requests[num_requests_added].time_arrival <= current_time
+ ):
+ request = requests[num_requests_added]
+ # These add-request calls will queue up the request on a zmq socket and return
+ # instantaneously. They will return an asyncio future which can be awaited for
+ # request completion.
+ futures.append(client.add_request(request.prompt_text, request.sampling_params))
+ num_requests_added += 1
+
+ if num_requests_added >= next_suspend_at and cycles_done < num_suspend_resume_cycles:
+ await suspend_resume_cycle(client, engine, args, futures)
+ cycles_done += 1
+ next_suspend_at += args.suspend_resume_interval
+
+ else:
+ # Add deterministic number of requests (generally used for debugging).
+ for i in range(
+ min(args.incoming_requests_per_step, num_requests_total - num_requests_added)
+ ):
+ # Change sampling parameters to force different generation lengths.
+ request = requests[num_requests_added]
+ n = request.sampling_params.num_tokens_to_generate
+ request.sampling_params.num_tokens_to_generate = n + i
+ futures.append(client.add_request(request.prompt_text, request.sampling_params))
+ num_requests_added += 1
+
+ if num_requests_added >= next_suspend_at and cycles_done < num_suspend_resume_cycles:
+ await suspend_resume_cycle(client, engine, args, futures)
+ cycles_done += 1
+ next_suspend_at += args.suspend_resume_interval
+
+ if num_requests_added == num_requests_total:
+ break
+ # Relinquish control since there are no more requests to add at the moment. This allows the engine to run.
+ await asyncio.sleep(0)
+
+ # While we wait for the requests to complete, the engine runs in the background.
+ results: List[DynamicInferenceRequestRecord] = await asyncio.gather(*futures)
+ else:
+ # Non-rank-0: match the suspend/resume cycles that rank 0 drives.
+ for _ in range(num_suspend_resume_cycles):
+ await engine.wait_until(EngineState.PAUSED)
+ await engine.wait_until(EngineState.SUSPENDED)
+ await engine.wait_until(EngineState.RESUMED)
+ await engine.wait_until(EngineState.RUNNING)
+
+ if dist.get_rank() == 0:
+ # Write results to JSON. Primarily used for functional testing.
+ if args.output_path:
+ json_results = {}
+ throughputs = []
+
+ for req in results:
+ result_dict = {
+ "input_prompt": req.prompt,
+ "generated_text": req.generated_text.replace("\n", "\\n"),
+ "generated_tokens": req.generated_tokens,
+ "latency": req.latency, # InferenceClient populates this field in the returned future.
+ }
+ if req.sampling_params.return_log_probs:
+ result_dict["logprobs"] = req.prompt_log_probs + req.generated_log_probs
+ throughput = len(req.generated_tokens) / req.latency
+ throughputs.append(throughput)
+ if req.routing_indices is not None:
+ result_dict["routing_indices"] = req.routing_indices.tolist()
+
+ json_results[req.request_id] = result_dict
+ throughput_dict = {"throughput": throughputs}
+ if args.throughput_check_only:
+ json_results = throughput_dict
+ with open(args.output_path, "w") as fp:
+ json.dump(json_results, fp, indent=4)
+ else:
+ print("Results:")
+ unique_prompt_map = defaultdict(list)
+ for req in results:
+ unique_prompt_map[req.prompt].append(req)
+ for idx, (prompt_text, reqs) in enumerate(unique_prompt_map.items()):
+ print(
+ f"%d/%d. prompt '%s' ... [%d] output '%s'."
+ % (
+ idx,
+ len(unique_prompt_map),
+ prompt_text.replace("\n", "\\n"),
+ len(reqs),
+ reqs[0].generated_text.replace("\n", "\\n"),
+ )
+ )
+
+ # Pause before stopping: STOP requires PAUSED or SUSPENDED state.
+ client.pause_engines()
+
+ await engine.wait_until(EngineState.PAUSED)
+
+ if dist.get_rank() == 0:
+ client.stop_engines()
+
+ await engine.wait_until(EngineState.STOPPED)
+
+ if dist.get_rank() == 0:
+ client.shutdown_coordinator()
+ client.stop()
+ logging.info(f"Rank: {dist.get_rank()} stopped their engine instance successfully.")
+
+
+if __name__ == "__main__":
+ # enable inference mode in the very beginning as some fp8 optimizations
+ # check for it.
+ with torch.inference_mode():
+ args = parse_and_validate_args(
+ extra_args_provider=add_inference_args,
+ args_defaults={'no_load_rng': True, 'no_load_optim': True},
+ )
+ initialize_megatron()
+ configure_nvtx_profiling(True)
+
+ tokenizer = get_tokenizer()
+
+ # Sampling params.
+ sampling_params = SamplingParams(
+ temperature=args.temperature,
+ top_k=args.top_k,
+ top_p=args.top_p,
+ return_log_probs=args.return_log_probs,
+ num_tokens_to_generate=args.num_tokens_to_generate,
+ termination_id=(
+ args.termination_id if args.termination_id is not None else tokenizer.eod
+ ),
+ )
+
+ model = get_model_for_inference()
+
+ requests = build_requests(args, tokenizer, sampling_params)
+
+ engine = get_dynamic_inference_engine(model=model)
+
+ if dist.get_rank() == 0:
+ setup_prefix = build_dynamic_engine_setup_prefix(args, model, engine.context, requests)
+ print("~~~")
+ print(setup_prefix)
+ print("~~~")
+
+ # Start Nsight profiler.
+ if os.environ.get("NSIGHT_PREFIX"):
+ torch.cuda.cudart().cudaProfilerStart()
+
+ asyncio.run(main(engine, requests, args.inference_coordinator_port))
+
+ # Stop Nsight profiler.
+ if os.environ.get("NSIGHT_PREFIX"):
+ torch.cuda.cudart().cudaProfilerStop()
diff --git a/examples/inference/gpt/gpt_static_inference.py b/examples/inference/gpt/gpt_static_inference.py
new file mode 100644
index 00000000000..d3dd619eaa1
--- /dev/null
+++ b/examples/inference/gpt/gpt_static_inference.py
@@ -0,0 +1,249 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import os
+import sys
+import time
+from argparse import Namespace
+
+from megatron.training.arguments import parse_and_validate_args
+import torch
+
+from megatron.core.inference.contexts import StaticInferenceContext
+from megatron.core.inference.engines import StaticInferenceEngine
+from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
+ GPTInferenceWrapper,
+)
+from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.inference.text_generation_controllers.text_generation_controller import (
+ TextGenerationController,
+)
+from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
+from megatron.core.transformer.module import MegatronModule
+
+sys.path.append(
+ os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
+)
+
+import asyncio
+import json
+from typing import List
+
+from examples.inference.gpt.utils import build_requests
+from megatron.inference.utils import add_inference_args, get_model_for_inference
+from megatron.training import get_args, get_tokenizer, print_rank_0
+from megatron.training.initialize import initialize_megatron
+
+
+def add_static_inference_args(parser):
+ """Static inference arguments."""
+
+ add_inference_args(parser)
+
+ group = parser.add_argument_group(title='Static inference')
+ group.add_argument(
+ "--max-batch-size",
+ type=int,
+ default=None,
+ dest="max_batch_size",
+ help='Deprecated, use `--inference-max-requests` instead',
+ )
+ group.add_argument("--stream", action="store_true", default=False, help="Stream output tokens")
+
+ return parser
+
+
+def get_inference_engine(args: Namespace, model: MegatronModule) -> StaticInferenceEngine:
+ """Utility to get the relevant backend for running inference
+
+ 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.
+
+ Args:
+ args (Namespace): The user arguments parsed from command line
+ model (MegatronModule): The megatron model .
+
+ Returns:
+ AbstractBackend: The chosen backend
+ """
+ tokenizer = build_tokenizer(args)
+ inference_context = StaticInferenceContext(
+ args.inference_max_requests, args.inference_max_seq_length
+ )
+ inference_wrapped_model = GPTInferenceWrapper(model, inference_context)
+ text_generation_controller = TextGenerationController(
+ inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer
+ )
+ engine_kwargs = {
+ "text_generation_controller": text_generation_controller,
+ "legacy": args.use_legacy_static_engine,
+ }
+ if not args.use_legacy_static_engine:
+ engine_kwargs["buffer_size_gb"] = args.inference_dynamic_batching_buffer_size_gb
+ return StaticInferenceEngine(**engine_kwargs)
+
+
+async def generate(
+ inference_engine: StaticInferenceEngine, sampling_params: SamplingParams, prompts: List[str]
+) -> List[InferenceRequest]:
+ async def collect_stream(prompt, request_id, stream_generator):
+ print(f"Request {request_id}: {prompt}", end="", flush=True)
+ prev_idx = 0
+ async for output in stream_generator:
+ print(output.generated_text[prev_idx:], end="", flush=True)
+ prev_idx = len(output.generated_text)
+ print()
+
+ request_ids: List[int] = [
+ inference_engine.add_request(prompt=prompt, sampling_params=sampling_params, streaming=True)
+ for prompt in prompts
+ ]
+ stream_generators = [
+ inference_engine.get_stream_generator(request_id) for request_id in request_ids
+ ]
+
+ tasks = [
+ asyncio.create_task(collect_stream(prompt, request_id, stream_generator))
+ for (prompt, request_id, stream_generator) in zip(prompts, request_ids, stream_generators)
+ ]
+
+ await inference_engine.run_engine_async()
+ await asyncio.gather(*tasks)
+
+ results: List[InferenceRequest] = [
+ inference_engine.scheduler.completed_request_pool[request_id] for request_id in request_ids
+ ]
+
+ return results
+
+
+@torch.inference_mode()
+def main():
+ """Main program."""
+
+ # Note: The default args passed here can be overwritten by using appropriate params (check arguments.py file)
+ # Micro batch size is not needed to be set by user. (It is calculated based on inference-batch-times-seqlen-threshold argument)
+ args = parse_and_validate_args(
+ extra_args_provider=add_static_inference_args,
+ args_defaults={
+ 'no_load_rng': True,
+ 'no_load_optim': True,
+ 'micro_batch_size': 1,
+ 'exit_on_missing_checkpoint': True,
+ },
+ )
+ initialize_megatron()
+
+ model = get_model_for_inference()
+
+ inference_engine = get_inference_engine(args, model)
+
+ sampling_params = SamplingParams(
+ temperature=args.temperature,
+ top_k=args.top_k,
+ top_p=args.top_p,
+ return_log_probs=args.return_log_probs,
+ num_tokens_to_generate=args.num_tokens_to_generate,
+ top_n_logprobs=args.top_n_logprobs,
+ )
+
+ # Build tokenizer
+ tokenizer = build_tokenizer(args)
+
+ requests = build_requests(args, tokenizer)
+ prompts = [r.prompt_text for r in requests]
+
+ if args.cuda_graph_impl == "local":
+ print(f"Running warmup for CUDA graphs...")
+ inference_engine.generate(
+ prompts=["warmup"], sampling_params=SamplingParams(num_tokens_to_generate=10)
+ )
+ start_time = time.perf_counter()
+ if args.stream:
+ results: List[InferenceRequest] = asyncio.run(
+ generate(inference_engine, sampling_params, prompts)
+ )
+ else:
+ results: List[InferenceRequest] = inference_engine.generate(
+ prompts=prompts, sampling_params=sampling_params
+ )
+ end_time = time.perf_counter()
+ latency = end_time - start_time
+
+ if torch.distributed.get_rank() == 0 and args.output_path:
+ results_output = {}
+ for idx, result in enumerate(results):
+ result_dict = {
+ 'input_prompt': result.prompt,
+ 'generated_text': result.generated_text,
+ 'generated_tokens': result.generated_tokens.tolist(),
+ 'tpot': result.tpot,
+ 'latency': latency,
+ }
+ if sampling_params.top_n_logprobs > 0:
+ result_dict['generated_top_n_logprobs'] = result.generated_top_n_logprobs
+ if sampling_params.return_log_probs:
+ response_logprobs = result.prompt_log_probs + result.generated_log_probs
+ result_dict["logprobs"] = response_logprobs
+ results_output[result.request_id] = result_dict
+
+ with open(args.output_path, 'w') as f:
+ json.dump(results_output, f)
+
+ # Print unique prompts + outputs.
+ if torch.distributed.get_rank() == 0:
+
+ print("~~~~ Unique prompts + outputs. ~~~~")
+
+ # Map results by their prompt.
+ from collections import defaultdict
+
+ unique_prompt_map = defaultdict(list)
+ for result_idx, result in enumerate(results):
+ unique_prompt_map[result.prompt].append(result_idx)
+
+ # Print unique prompts + outputs.
+ for unique_idx, (prompt_text, result_idxs) in enumerate(unique_prompt_map.items()):
+ result_idx = result_idxs[0]
+ result = results[result_idx]
+ generated_text = result.generated_text.replace("\n", "\\n")
+ print(
+ f"{unique_idx}/{len(unique_prompt_map)} [{len(result_idxs)}]. {prompt_text} "
+ f"... {generated_text}"
+ )
+
+ stats = torch.cuda.memory_stats()
+ print_rank_0(
+ "static | cg %d | %s | reqs %d [ batch %d ] ... mem %.1f/%.1f ... time %.3f."
+ % (
+ args.cuda_graph_impl == "local",
+ (
+ f""
+ if args.prompts
+ else " %s, %d, %.1e, %.1e"
+ % (
+ "(%s)" % " ".join(map(str, args.num_tokens_to_prompt)),
+ args.num_tokens_to_generate,
+ args.incoming_requests_duration,
+ args.incoming_requests_per_sec,
+ )
+ ),
+ len(requests),
+ args.inference_max_requests,
+ stats["allocated_bytes.all.peak"] / (1024**3),
+ stats["reserved_bytes.all.peak"] / (1024**3),
+ latency,
+ )
+ )
+ # Force immediate process exit to bypass torchrun's atexit NCCL teardown when
+ # CUDA graphs have captured collectives (see PyTorch issue #115388). This can
+ # sometimes lead to hangs in the atexit handler.
+ # We do this only when CUDA graphs are enabled.
+ if args.cuda_graph_impl != "none":
+ print(f"[main] rank {torch.distributed.get_rank()}: finished", flush=True)
+ os._exit(0)
+ else:
+ torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py
new file mode 100644
index 00000000000..c9b1c05c544
--- /dev/null
+++ b/examples/inference/gpt/utils.py
@@ -0,0 +1,326 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import copy
+import itertools
+import json
+import random
+import time
+from argparse import ArgumentParser, Namespace
+from functools import partial
+from typing import Any, List, Optional
+
+import torch
+from tqdm import tqdm
+
+from megatron.core.inference.contexts import DynamicInferenceContext
+from megatron.core.inference.contexts.dynamic_context import get_mem_size_str
+from megatron.core.inference.inference_request import DynamicInferenceRequest
+from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.transformer.module import MegatronModule
+from megatron.training import get_args
+
+
+def get_default_sampling_params(termination_id: int = None):
+ return SamplingParams(
+ temperature=1.0,
+ top_k=1,
+ top_p=0.0,
+ return_log_probs=False,
+ num_tokens_to_generate=30,
+ termination_id=termination_id,
+ )
+
+
+def get_curr_time() -> float:
+ """Get synchronized time across ranks."""
+ curr_time = torch.cuda.LongTensor([time.time_ns()])
+ if torch.distributed.is_initialized():
+ torch.distributed.broadcast(curr_time, src=0)
+ return curr_time.item() / 10**9
+
+
+class Request:
+ """Class to hold attributes for a single request.
+
+ A request is initialized with its prompt text. As it is added, processed,
+ and completed through the inference engine, the request is populated with its
+ start time, end time, and output tokens.
+
+ Args:
+ prompt_text (str): Prompt text.
+ time_offset (float): Artificial time offset for simulating incoming
+ requests. This value is later added to the `base_arrival_time` to
+ simulate the requests arrival time.
+ tokenizer (Any): Tokenizer for tokenizing the prompt.
+ """
+
+ def __init__(
+ self,
+ prompt_text: str,
+ time_offset: float,
+ tokenizer: Any,
+ sampling_params: SamplingParams = None,
+ ):
+ self.prompt_text = prompt_text
+ self.prompt_tokens = tokenizer.tokenize(prompt_text)
+ self.output_text = None
+ self.output_tokens = []
+ self.time_offset = time_offset
+ self.time_arrival = None
+ self.time_start = None
+ self.time_end = None
+ self.ttft = None # Time-to-first-token in seconds
+ self.state = "not-started"
+ self.sampling_params: SamplingParams = (
+ sampling_params
+ if sampling_params is not None
+ else get_default_sampling_params(tokenizer.eod)
+ )
+ self.sampling_params = copy.deepcopy(self.sampling_params)
+
+ def __str__(self) -> str:
+ return "state '%s'; toffset %.1e; prompt len %d; output len %d; '%s'" % (
+ self.state,
+ self.time_offset,
+ len(self.prompt_tokens),
+ len(self.output_tokens),
+ self.prompt_text,
+ )
+
+
+def get_time_offsets(
+ seed: int | None,
+ incoming_requests_per_step: int,
+ incoming_requests_per_sec: float,
+ num_requests: int,
+) -> list[float]:
+ """Get example time offsets."""
+
+ # Time offsets to add all requests at once.
+ if incoming_requests_per_step is not None or incoming_requests_per_sec <= 0:
+ return [-1] * num_requests
+
+ # if num_requests is not None:
+ incoming_requests_duration = num_requests / incoming_requests_per_sec
+ incoming_requests_duration *= 2 # extra margin, to accomodate time sampling
+
+ random.seed(seed)
+
+ import simpy # Guard against this import in test case
+
+ # Generate random time offsets.
+ def arrival(r):
+ while True:
+ yield env.timeout(random.expovariate(r))
+ time_offsets.append(env.now)
+
+ time_offsets = []
+ env = simpy.Environment()
+ env.process(arrival(incoming_requests_per_sec))
+ env.run(incoming_requests_duration)
+
+ # Ensure at least a single request.
+ if len(time_offsets) == 0:
+ time_offsets = [0.0]
+
+ # Ensure first time is 0.
+ time_offsets = [to - time_offsets[0] for to in time_offsets]
+
+ # Truncate to num_requests.
+ assert len(time_offsets) >= num_requests
+ time_offsets = time_offsets[:num_requests]
+
+ return time_offsets
+
+
+def get_cli_requests(
+ args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
+
+ # Get time offsets.
+ t_offsets = get_time_offsets(
+ args.seed,
+ args.incoming_requests_per_step,
+ args.incoming_requests_per_sec,
+ len(args.prompts),
+ )
+
+ # Init requests.
+ requests = [Request(p, t, tokenizer, sampling_params) for p, t in zip(args.prompts, t_offsets)]
+ return requests
+
+
+def get_synthetic_requests(
+ args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
+ """Get example requests."""
+
+ # Get time offsets.
+ time_offsets = get_time_offsets(
+ args.seed,
+ args.incoming_requests_per_step,
+ args.incoming_requests_per_sec,
+ int(args.incoming_requests_per_sec * args.incoming_requests_duration),
+ )
+
+ # Build prompts with expected lengths.
+ assert (
+ len(args.num_tokens_to_prompt) == 2
+ and args.num_tokens_to_prompt[1] >= args.num_tokens_to_prompt[0]
+ )
+ max_prompt_length = args.num_tokens_to_prompt[1]
+ max_prompt_text = "hi " * max_prompt_length
+ max_prompt_tokens = tokenizer.tokenize(max_prompt_text)
+ prompt_lengths = [random.randint(*args.num_tokens_to_prompt) for _ in time_offsets]
+ prompt_tokens_list = [max_prompt_tokens[:l] for l in prompt_lengths]
+ prompt_texts = [tokenizer.detokenize(tt) for tt in prompt_tokens_list]
+
+ # Init requests.
+ assert len(prompt_texts) == len(time_offsets)
+ requests = [
+ Request(t, o, tokenizer, sampling_params=sampling_params)
+ for t, o in zip(prompt_texts, time_offsets)
+ ]
+
+ return requests
+
+
+def get_requests_from_file(
+ args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
+ """Get requests from a file."""
+ if not args.prompt_file:
+ raise ValueError("Prompt file is required to read requests from a file.")
+
+ # Load prompts.
+ n_prompts = sum(1 for _ in open(args.prompt_file))
+ prompts = []
+ if sampling_params is None:
+ sampling_params = get_default_sampling_params(tokenizer.eod)
+ sampling_params_list = []
+ with open(args.prompt_file) as f:
+ for line in tqdm(f.readlines(), "read prompt file", total=n_prompts):
+ line_dict = json.loads(line)
+ prompts.append(line_dict["text"])
+
+ sp = copy.deepcopy(sampling_params)
+ if args.num_tokens_from_file:
+ sp.num_tokens_to_generate = line_dict["chatgpt_output_token_length"]
+ sampling_params_list.append(sp)
+
+ if len(prompts) == args.prompt_file_num_truncate:
+ break
+
+ # Get time offsets.
+ time_offsets: list[float] = get_time_offsets(
+ args.seed, args.incoming_requests_per_step, args.incoming_requests_per_sec, len(prompts)
+ )
+
+ # Init requests.
+ requests = [
+ Request(p, t, tokenizer, sp)
+ for p, t, sp in tqdm(
+ zip(prompts, time_offsets, sampling_params_list), "init requests", total=len(prompts)
+ )
+ ]
+
+ return requests
+
+
+def build_requests(
+ args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
+ # Check if we have any prompts (from command line or JSONL)
+ if args.prompts:
+ if args.prompt_file:
+ raise ValueError("Cannot use both --prompts and --prompt-file")
+ return get_cli_requests(args, tokenizer, sampling_params)
+ elif args.prompt_file:
+ return get_requests_from_file(args, tokenizer, sampling_params)
+ else:
+ return get_synthetic_requests(args, tokenizer, sampling_params)
+
+
+def get_model_size_str(model):
+ n = sum(p.numel() for p in model.parameters())
+ for exp, suffix in ((12, "t"), (9, "b"), (6, "m"), (3, "k"), (0, "")):
+ nquery = int(10**exp)
+ if n > nquery:
+ return "%d%s" % (n // nquery, suffix)
+ raise Exception("something went wrong.")
+
+
+def build_dynamic_engine_setup_prefix(
+ args: Namespace,
+ model: MegatronModule,
+ context: DynamicInferenceContext,
+ requests: list[DynamicInferenceRequest],
+):
+ """
+ Returns a compact, pipe-separated summary of the dynamic-batching setup.
+
+ Example output:
+
+ `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
+
+ Args:
+ args (Namespace): Command-line arguments for this run.
+ context (DynamicInferenceContext): Stores limits such as `max_requests`,
+ `max_tokens`, and `gtd_request_count`.
+ requests (List[DynamicInferenceRequest]): List of inference requests.
+
+ Returns:
+ A configuration string for logging.
+ """
+ # CUDA graph config
+ if args.cuda_graph_impl == "local":
+ cg_str = f"graphs {len(context.cuda_graph_batch_dimensions_list)}"
+ else:
+ cg_str = "--"
+
+ # Unified memory (UVM).
+ uvm_str = f"uvm {int(context.unified_memory_level)}"
+
+ # Prompt description
+ prompt_src_str = (
+ "cli"
+ if args.prompts
+ else (
+ "file"
+ if args.prompt_file
+ else f"synth({', '.join(map(str, args.num_tokens_to_prompt))})"
+ )
+ )
+ request_str = (
+ f"requests: {prompt_src_str}, " f"n {len(requests):d}, g {args.num_tokens_to_generate:d}, "
+ )
+ request_str += (
+ f"dur {args.incoming_requests_duration:.1e} " f"r/sec {args.incoming_requests_per_sec:.1e}"
+ if args.incoming_requests_per_step is None
+ else f"r/step {args.incoming_requests_per_step}"
+ )
+
+ # Buffer limits config
+ buffer_limits_str = (
+ f"bf: {get_mem_size_str(args.inference_dynamic_batching_buffer_size_gb*1024**3)}, "
+ f"{context.kv_block_allocator.active_count} chunks "
+ f"[r {context.max_requests}, t {context.max_tokens}]"
+ )
+
+ parts = [get_model_size_str(model), "dynamic", cg_str, uvm_str, request_str, buffer_limits_str]
+
+ return " | ".join(parts)
+
+
+def get_global_peak_memory_stats_bytes() -> dict:
+ """Peak allocated CUDA memory aggregated across ranks (MAX), in bytes.
+
+ Uses `torch.cuda.max_memory_allocated()` and assumes peak stats were reset
+ before the benchmark run.
+ """
+ peak_alloc = int(torch.cuda.max_memory_allocated())
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ t = torch.tensor([peak_alloc], device="cuda", dtype=torch.int64)
+ torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.MAX)
+ peak_alloc = int(t[0].item())
+ return {"mem-max-allocated-bytes": peak_alloc}
diff --git a/examples/inference/llama_mistral/huggingface_reference.py b/examples/inference/llama_mistral/huggingface_reference.py
new file mode 100644
index 00000000000..9d8f4465f65
--- /dev/null
+++ b/examples/inference/llama_mistral/huggingface_reference.py
@@ -0,0 +1,25 @@
+import argparse
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+# Set up argument parsing
+parser = argparse.ArgumentParser(description="Script for text generation with a specific model and prompt.")
+parser.add_argument('--prompt', type=str, required=True, help="Prompt text to use for text generation")
+parser.add_argument('--model-path', type=str, required=True, help="Path to the Huggingface model checkpoint")
+
+# Parse command-line arguments
+args = parser.parse_args()
+
+model_path = args.model_path
+prompt = args.prompt
+
+config = AutoConfig.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path, config=config)
+model = AutoModelForCausalLM.from_pretrained(model_path, config=config).cuda()
+
+inputs = tokenizer(prompt, return_tensors="pt")
+for key in inputs:
+ inputs[key] = inputs[key].cuda()
+# top_k, top_p and do_sample are set for greedy argmax based sampling
+
+outputs = model.generate(**inputs, max_length=100, do_sample=False, top_p=0, top_k=0, temperature=1.0)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
\ No newline at end of file
diff --git a/examples/inference/llama_mistral/run_static_inference_llama4_scout.sh b/examples/inference/llama_mistral/run_static_inference_llama4_scout.sh
new file mode 100755
index 00000000000..cc8cfac5e69
--- /dev/null
+++ b/examples/inference/llama_mistral/run_static_inference_llama4_scout.sh
@@ -0,0 +1,68 @@
+#!/bin/bash
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NVTE_APPLY_QK_LAYER_SCALING=0
+
+DISTRIBUTED_ARGS="--nproc_per_node 8 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr 0.0.0.0 \
+ --master_port 6000"
+
+# Fill in checkpoint path to Llama 4 Scout to run
+CHECKPOINT=
+PROMPTS="What is the capital of France?"
+TOKENS_TO_GENERATE=4
+MAX_BATCH_SIZE=2
+
+MODEL_ARGS=" \
+ --micro-batch-size 1 \
+ --bf16 \
+ --no-masked-softmax-fusion \
+ --disable-bias-linear \
+ --untie-embeddings-and-output-weights \
+ --position-embedding-type rope \
+ --no-rope-fusion \
+ --normalization RMSNorm \
+ --swiglu \
+ --num-layers 48 \
+ --hidden-size 5120 \
+ --ffn-hidden-size 16384 \
+ --num-attention-heads 40 \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --qk-layernorm \
+ --num-experts 16 \
+ --moe-ffn-hidden-size 8192 \
+ --moe-router-score-function sigmoid \
+ --moe-router-topk 1 \
+ --moe-router-topk-scaling-factor 1.0 \
+ --moe-shared-expert-intermediate-size 8192 \
+ --moe-aux-loss-coeff 1e-3 \
+ --moe-token-dispatcher-type alltoall \
+ --moe-token-drop-policy probs \
+ --moe-router-load-balancing-type seq_aux_loss \
+ --seq-length 4096 \
+ --max-position-embeddings 4096 \
+ --tokenizer-type HuggingFaceTokenizer \
+ --make-vocab-size-divisible-by 128 \
+ --use-mcore-models \
+ --rotary-interleaved \
+ --rotary-percent 1.0 \
+ --rotary-base 500000 \
+ --rope-scaling-factor 8.0 \
+ --use-rope-scaling \
+ --no-bias-swiglu-fusion \
+ --qk-l2-norm \
+ --moe-apply-probs-on-input \
+ --moe-router-dtype fp64 \
+"
+
+torchrun $DISTRIBUTED_ARGS -m examples.inference.gpt.gpt_static_inference \
+ --load ${CHECKPOINT} \
+ --tokenizer-model unsloth/Llama-4-Scout-17B-16E-Instruct \
+ --dist-ckpt-strictness log_unexpected \
+ --tensor-model-parallel-size 8 \
+ --prompts ${PROMPTS} \
+ --num-tokens-to-generate ${TOKENS_TO_GENERATE} \
+ --max-batch-size ${MAX_BATCH_SIZE} \
+ ${MODEL_ARGS}
diff --git a/examples/inference/llama_mistral/run_text_generation_llama3.1.sh b/examples/inference/llama_mistral/run_text_generation_llama3.1.sh
new file mode 100755
index 00000000000..06584f0917d
--- /dev/null
+++ b/examples/inference/llama_mistral/run_text_generation_llama3.1.sh
@@ -0,0 +1,56 @@
+#!/bin/bash
+# This example will start serving the Llama3.1-8B model
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NVTE_APPLY_QK_LAYER_SCALING=0
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr 0.0.0.0 \
+ --master_port 6000"
+
+# Ensure CHECKPOINT and TOKENIZER_MODEL are provided
+if [ -z "$1" ] || [ -z "$2" ]; then
+ echo "Error: You must provide CHECKPOINT and TOKENIZER_MODEL as command-line arguments."
+ echo "Usage: $0 /path/to/checkpoint /path/to/tokenizer_model"
+ exit 1
+fi
+
+# Assign command-line arguments to variables
+CHECKPOINT=$1
+TOKENIZER_MODEL=$2
+
+pip install flask-restful
+
+torchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \
+ --use-checkpoint-args \
+ --disable-bias-linear \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model ${TOKENIZER_MODEL} \
+ --transformer-impl transformer_engine \
+ --normalization RMSNorm \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --no-masked-softmax-fusion \
+ --attention-softmax-in-fp32 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --untie-embeddings-and-output-weights \
+ --position-embedding-type rope \
+ --rotary-percent 1.0 \
+ --rotary-base 500000 \
+ --use-rope-scaling \
+ --use-rotary-position-embeddings \
+ --swiglu \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --ffn-hidden-size 14336 \
+ --load ${CHECKPOINT} \
+ --num-attention-heads 32 \
+ --max-position-embeddings 131072 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --seq-length 8192
diff --git a/examples/inference/llama_mistral/run_text_generation_llama3.sh b/examples/inference/llama_mistral/run_text_generation_llama3.sh
new file mode 100755
index 00000000000..c5fc4103ab5
--- /dev/null
+++ b/examples/inference/llama_mistral/run_text_generation_llama3.sh
@@ -0,0 +1,55 @@
+#!/bin/bash
+# This example will start serving the Llama3-8B model
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NVTE_APPLY_QK_LAYER_SCALING=0
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr 0.0.0.0 \
+ --master_port 6000"
+
+# Ensure CHECKPOINT and TOKENIZER_MODEL are provided
+if [ -z "$1" ] || [ -z "$2" ]; then
+ echo "Error: You must provide CHECKPOINT and TOKENIZER_MODEL as command-line arguments."
+ echo "Usage: $0 /path/to/checkpoint /path/to/tokenizer_model"
+ exit 1
+fi
+
+# Assign command-line arguments to variables
+CHECKPOINT=$1
+TOKENIZER_MODEL=$2
+
+pip install flask-restful
+
+torchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \
+ --use-checkpoint-args \
+ --disable-bias-linear \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model ${TOKENIZER_MODEL} \
+ --transformer-impl transformer_engine \
+ --normalization RMSNorm \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --no-masked-softmax-fusion \
+ --attention-softmax-in-fp32 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --untie-embeddings-and-output-weights \
+ --position-embedding-type rope \
+ --rotary-percent 1.0 \
+ --rotary-base 500000 \
+ --use-rotary-position-embeddings \
+ --swiglu \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --ffn-hidden-size 14336 \
+ --load ${CHECKPOINT} \
+ --num-attention-heads 32 \
+ --max-position-embeddings 8192 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --seq-length 8192
diff --git a/examples/inference/llama_mistral/run_text_generation_mistral.sh b/examples/inference/llama_mistral/run_text_generation_mistral.sh
new file mode 100755
index 00000000000..4358fd494c7
--- /dev/null
+++ b/examples/inference/llama_mistral/run_text_generation_mistral.sh
@@ -0,0 +1,53 @@
+#!/bin/bash
+# This example will start serving the Mistral-7B-v0.3 model
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr 0.0.0.0 \
+ --master_port 6000"
+
+# Ensure CHECKPOINT and TOKENIZER_MODEL are provided
+if [ -z "$1" ] || [ -z "$2" ]; then
+ echo "Error: You must provide CHECKPOINT and TOKENIZER_MODEL as command-line arguments."
+ echo "Usage: $0 /path/to/checkpoint /path/to/tokenizer_model"
+ exit 1
+fi
+
+# Assign command-line arguments to variables
+CHECKPOINT=$1
+TOKENIZER_MODEL=$2
+
+pip install flask-restful
+
+torchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model ${TOKENIZER_MODEL} \
+ --use-checkpoint-args \
+ --apply-layernorm-1p \
+ --transformer-impl transformer_engine \
+ --normalization RMSNorm \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --no-masked-softmax-fusion \
+ --use-flash-attn \
+ --untie-embeddings-and-output-weights \
+ --disable-bias-linear \
+ --position-embedding-type rope \
+ --rotary-percent 1.0 \
+ --rotary-base 1000000 \
+ --swiglu \
+ --ffn-hidden-size 14336 \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --load ${CHECKPOINT} \
+ --num-attention-heads 32 \
+ --max-position-embeddings 4096 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --seq-length 4096 \
+ --seed 101
diff --git a/examples/run_text_generation_server_345M.sh b/examples/inference/run_text_generation_server_345M.sh
similarity index 92%
rename from examples/run_text_generation_server_345M.sh
rename to examples/inference/run_text_generation_server_345M.sh
index a151b984676..e8e61adb163 100755
--- a/examples/run_text_generation_server_345M.sh
+++ b/examples/inference/run_text_generation_server_345M.sh
@@ -26,9 +26,6 @@ torchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \
--fp16 \
--micro-batch-size 1 \
--seq-length 1024 \
- --out-seq-length 1024 \
- --temperature 1.0 \
--vocab-file $VOCAB_FILE \
--merge-file $MERGE_FILE \
- --top_p 0.9 \
--seed 42
diff --git a/examples/run_text_generation_server_345M_8_tensor_parallel.sh b/examples/inference/run_text_generation_server_345M_8_tensor_parallel.sh
similarity index 92%
rename from examples/run_text_generation_server_345M_8_tensor_parallel.sh
rename to examples/inference/run_text_generation_server_345M_8_tensor_parallel.sh
index 027ab421727..368cec3b312 100755
--- a/examples/run_text_generation_server_345M_8_tensor_parallel.sh
+++ b/examples/inference/run_text_generation_server_345M_8_tensor_parallel.sh
@@ -24,9 +24,6 @@ python -m torch.distributed.launch $DISTRIBUTED_ARGS tools/run_text_generation_s
--fp16 \
--micro-batch-size 1 \
--seq-length 1024 \
- --out-seq-length 1024 \
- --temperature 1.0 \
--vocab-file $VOCAB_FILE \
--merge-file $MERGE_FILE \
- --top_p 0.9 \
--seed 42
diff --git a/examples/inference/t5/simple_t5_batch_inference.py b/examples/inference/t5/simple_t5_batch_inference.py
new file mode 100644
index 00000000000..1aca74b3176
--- /dev/null
+++ b/examples/inference/t5/simple_t5_batch_inference.py
@@ -0,0 +1,162 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+import os
+import sys
+from argparse import Namespace
+
+import torch
+
+import pretrain_t5
+from megatron.core.inference.engines import AbstractEngine, StaticInferenceEngine
+from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
+ InferenceWrapperConfig,
+)
+from megatron.core.inference.model_inference_wrappers.t5.t5_inference_wrapper import (
+ T5InferenceWrapper,
+)
+from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.inference.text_generation_controllers.encoder_decoder_text_generation_controller import (
+ EncoderDecoderTextGenerationController,
+)
+from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
+from megatron.core.transformer.module import MegatronModule
+from pretrain_t5 import model_provider
+
+sys.path.append(
+ os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
+)
+
+from typing import List
+
+from megatron.core import mpu
+from megatron.training import get_args, get_model, get_tokenizer
+from megatron.training.checkpointing import load_checkpoint
+from megatron.training.initialize import initialize_megatron
+
+
+def add_text_generate_args(parser):
+ """Text generation arguments."""
+ group = parser.add_argument_group(title='text generation')
+
+ group.add_argument("--temperature", type=float, default=1.0, help='Sampling temperature.')
+ group.add_argument("--top_k", type=int, default=1, help='Top k sampling.')
+ group.add_argument("--top_p", type=float, default=0.0, help='Top p sampling.')
+ group.add_argument(
+ "--return-log-probs",
+ action='store_true',
+ default=False,
+ help='Return the log probabilities of the final output tokens',
+ )
+ group.add_argument(
+ "--num-tokens-to-generate",
+ type=int,
+ default=30,
+ help='Number of tokens to generate for each prompt',
+ )
+ group.add_argument(
+ "--encoder-prompts",
+ metavar='N',
+ type=str,
+ nargs='+',
+ help='Encoder input prompts with each prompt within quotes and separated by space',
+ )
+ group.add_argument(
+ "--max-batch-size", type=int, default=1, help='Max number of prompts to process at once'
+ )
+ return parser
+
+
+def get_inference_engine(args: Namespace, model: MegatronModule) -> AbstractEngine:
+ """Utility to get the relevant backend for running inference
+
+ 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.
+
+ Args:
+ args (Namespace): The user arguments parsed from command line
+ model (MegatronModule): The megatron model .
+
+ Returns:
+ AbstractBackend: The chosen backend
+ """
+ # Build tokenizer
+ tokenizer = build_tokenizer(args)
+
+ inference_wrapper_config = InferenceWrapperConfig(
+ hidden_size=args.hidden_size,
+ inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold,
+ fp32_residual_connection=args.fp32_residual_connection,
+ params_dtype=args.params_dtype,
+ padded_vocab_size=args.padded_vocab_size,
+ )
+
+ inference_wrapped_model = T5InferenceWrapper(model, inference_wrapper_config)
+ text_generation_controller = EncoderDecoderTextGenerationController(
+ inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer
+ )
+ return StaticInferenceEngine(
+ text_generation_controller=text_generation_controller, max_batch_size=args.max_batch_size
+ )
+
+
+def main():
+ """Main program."""
+
+ # Note: The default args passed here can be overwritten by using appropriate params (check arguments.py file)
+ # Micro batch size is not needed to be set by user. (It is calculated based on inference-batch-times-seqlen-threshold argument)
+ initialize_megatron(
+ extra_args_provider=add_text_generate_args,
+ args_defaults={
+ 'no_load_rng': True,
+ 'no_load_optim': True,
+ 'micro_batch_size': 1,
+ 'exit_on_missing_checkpoint': True,
+ },
+ )
+
+ # Set up model and load checkpoint
+ model = get_model(model_provider, wrap_with_ddp=False)
+ load_checkpoint(model, None, None)
+ model = model[0]
+
+ args = get_args()
+
+ inference_engine = get_inference_engine(args, model)
+
+ sampling_params = SamplingParams(
+ temperature=args.temperature,
+ top_k=args.top_k,
+ top_p=args.top_p,
+ return_log_probs=args.return_log_probs,
+ num_tokens_to_generate=args.num_tokens_to_generate,
+ )
+
+ # Build tokenizer
+ tokenizer = build_tokenizer(args)
+
+ decoder_prompts = [""] * len(
+ args.encoder_prompts
+ ) # for T5, the prompt is provided as encoder input, hence decoder_prompts is empty
+ args.prompts = decoder_prompts
+
+ results: List[InferenceRequest] = inference_engine.generate(
+ prompts=args.prompts,
+ add_BOS=True,
+ encoder_prompts=args.encoder_prompts,
+ sampling_params=sampling_params,
+ )
+
+ if torch.distributed.get_rank() == 0:
+ for idx, result in enumerate(results):
+ print(f' \n------------- RESULT FOR PROMPT {idx} --------------- ')
+ result = {
+ 'id': result.request_id,
+ 'input_prompt': result.prompt,
+ 'generated_text': result.generated_text,
+ 'generated_tokens': result.generated_tokens,
+ }
+ print(result)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/llama/README.md b/examples/llama/README.md
new file mode 100644
index 00000000000..9872185ab2f
--- /dev/null
+++ b/examples/llama/README.md
@@ -0,0 +1,144 @@
+# Llama Models
+
+## Table of contents
+- [1. Overview](#1-overview)
+- [2. Prerequisites](#2-prerequisites)
+- [3. Training Setup](#3-training-setup)
+- [4. Configuration](#4-configuration)
+- [5. Test Datasets](#5-test-datasets)
+- [6. FP8 Debugging](#6-fp8-debugging)
+
+## 1. Overview
+
+
+Train Llama models using FP8 precision with Megatron-Core.
+
+## 2. Prerequisites
+
+
+```bash
+# Clone repository
+export HOST_MEGATRON_LM_DIR="/path/to/your/host/megatron-lm"
+git clone https://github.com/NVIDIA/Megatron-LM.git "$HOST_MEGATRON_LM_DIR"
+cd "$HOST_MEGATRON_LM_DIR"
+git checkout "core_r0.12.0"
+
+# Set paths
+export HOST_CHECKPOINT_PATH="./checkpoints/llama3_8b_fp8"
+export HOST_TENSORBOARD_LOGS_PATH="./tensorboard_logs/llama3_8b_fp8"
+
+# Optional: For real data
+# export HOST_TOKENIZER_MODEL_PATH="/path/to/host/tokenizer.model"
+# export HOST_DATA_PREFIX="/path/to/host/mydata_prefix"
+```
+
+## 3. Training Setup
+
+
+### Using Mock Data
+```bash
+PYTORCH_IMAGE="nvcr.io/nvidia/pytorch:25.03-py3"
+
+docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \
+ -v "${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm" \
+ -v "${HOST_CHECKPOINT_PATH}:/workspace/checkpoints" \
+ -v "${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs" \
+ --workdir /workspace/megatron-lm \
+ $PYTORCH_IMAGE \
+ bash examples/llama/train_llama3_8b_h100_fp8.sh \
+ /workspace/checkpoints \
+ /workspace/tensorboard_logs \
+ 2>&1 | tee "${HOST_TENSORBOARD_LOGS_PATH}/training_mock_$(date +'%y-%m-%d_%H-%M-%S').log"
+```
+
+### Using Custom Data and Tokenizer
+```bash
+PYTORCH_IMAGE="nvcr.io/nvidia/pytorch:25.03-py3"
+
+docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \
+ -v "${HOST_MEGATRON_LM_DIR}:/workspace/megatron-lm" \
+ -v "${HOST_CHECKPOINT_PATH}:/workspace/checkpoints" \
+ -v "${HOST_TENSORBOARD_LOGS_PATH}:/workspace/tensorboard_logs" \
+ -v "${HOST_TOKENIZER_MODEL_PATH}:/workspace/tokenizer_model" \
+ -v "$(dirname "${HOST_DATA_PREFIX}"):/workspace/data_dir" \
+ --workdir /workspace/megatron-lm \
+ $PYTORCH_IMAGE \
+ bash examples/llama/train_llama3_8b_h100_fp8.sh \
+ /workspace/checkpoints \
+ /workspace/tensorboard_logs \
+ /workspace/tokenizer_model \
+ "/workspace/data_dir/$(basename "${HOST_DATA_PREFIX}")" \
+ 2>&1 | tee "${HOST_TENSORBOARD_LOGS_PATH}/training_custom_$(date +'%y-%m-%d_%H-%M-%S').log"
+```
+
+## 4. Configuration
+
+
+Default parallelism strategy:
+- Tensor Parallel: 1
+- Pipeline Parallel: 1
+- Context Parallel: 2
+
+Llama-3-8B architecture:
+- 32 layers
+- Hidden size: 4096
+- FFN hidden size: 14336
+- Attention heads: 32
+- Query groups: 8
+- Sequence length: 8192
+- RMSNorm normalization with SwiGLU and RoPE
+
+Key training parameters:
+- Micro-batch size: 1
+- Global batch size: 128
+- Learning rate: 1.5e-4
+- Min learning rate: 1.0e-5
+- Weight decay: 0.1
+- FP8 format: hybrid
+
+You can modify these parameters directly in the `train_llama3_8b_h100_fp8.sh` script.
+
+This 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).
+
+### FP8 Performance
+
+| Model | #-GPUs | GBS | MBS | Seq Length | TP | PP | CP | VP | EP | GA | Tokens/sec/GPU | TFLOP/sec/GPU |
+|-------|--------|-----|-----|------------|----|----|----|----|----|----|----------------|---------------|
+| LLAMA3-8B | 8 | 128 | 1 | 8192 | 1 | 1 | 2 | 1 | 1 | 32 | 13812 | 800 |
+| LLAMA3-70B | 64 | 128 | 1 | 8192 | 4 | 8 | 1 | 5 | 1 | 64 | 1621 | 780 |
+| LLAMA3-405B | 1024 | 512 | 1 | 8192 | 8 | 8 | 2 | 8 | 1 | 64 | 315 | 834 |
+
+Legend:
+- GBS: Global Batch Size
+- MBS: Micro Batch Size
+- TP: Tensor Parallel size
+- PP: Pipeline Parallel size
+- CP: Context Parallel size
+- VP: Virtual Pipeline stages
+- EP: Expert Parallel size
+- GA: Gradient Accumulation steps
+
+As 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).
+
+## 5. Test Datasets
+
+
+Recommended datasets:
+1. **WikiText-103**: https://huggingface.co/datasets/Salesforce/wikitext
+
+Preprocess datasets:
+```bash
+python "${HOST_MEGATRON_LM_DIR}/tools/preprocess_data.py" \
+ --input your_dataset.json \
+ --output-prefix test_dataset \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model /path/to/tokenizer.model \
+ --append-eod
+```
+
+## 6. FP8 Training Considerations
+
+
+- **Hardware**: Requires NVIDIA Hopper, Ada, or Blackwell GPUs for FP8 support
+
+- **Troubleshooting**: If you encounter NaN values or instability with FP8 training, please refer to [Transformer Engine](https://github.com/NVIDIA/TransformerEngine).
diff --git a/examples/llama/train_llama3_8b_h100_fp8.sh b/examples/llama/train_llama3_8b_h100_fp8.sh
new file mode 100644
index 00000000000..28227546bc7
--- /dev/null
+++ b/examples/llama/train_llama3_8b_h100_fp8.sh
@@ -0,0 +1,196 @@
+#!/bin/bash
+
+# Environment variables for performance tuning
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+#export LOG_LEVEL=${LOG_LEVEL:-INFO}
+#export NCCL_IB_TIMEOUT=${NCCL_IB_TIMEOUT:-19}
+#export NVTE_FWD_LAYERNORM_SM_MARGIN=${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}
+#export NVTE_BWD_LAYERNORM_SM_MARGIN=${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}
+#export NCCL_P2P_NET_CHUNKSIZE=${NCCL_P2P_NET_CHUNKSIZE:-2097152}
+#export NCCL_AVOID_RECORD_STREAMS=${NCCL_AVOID_RECORD_STREAMS:-1}
+
+CHECKPOINT_PATH=${1:-"checkpoints/llama3_8b_fp8"}
+TENSORBOARD_LOGS_PATH=${2:-"tensorboard_logs/llama3_8b_fp8"}
+TOKENIZER_ARG=${3:-"MOCK"} # Path to tokenizer model, or "MOCK"
+DATA_ARG=${4:-"MOCK"} # Data prefix, or "MOCK"
+
+# Create directories if they don't exist
+mkdir -p "$(dirname "$CHECKPOINT_PATH")"
+mkdir -p "$(dirname "$TENSORBOARD_LOGS_PATH")"
+
+# Distributed training setup
+GPUS_PER_NODE=8
+NUM_NODES=1
+MASTER_ADDR=${MASTER_ADDR:-localhost}
+MASTER_PORT=${MASTER_PORT:-6000}
+NODE_RANK=${NODE_RANK:-0}
+WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
+
+# Path to the pretrain_gpt.py script, assuming this script is run from the root of the Megatron-LM repository
+PRETRAIN_SCRIPT_PATH="pretrain_gpt.py"
+
+# Fixed model and training parameters
+TP_SIZE=1
+CP_SIZE=1
+PP_SIZE=1
+MICRO_BATCH_SIZE=1
+GLOBAL_BATCH_SIZE=128
+NUM_LAYERS=32
+DTYPE="fp8"
+SEQ_LENGTH=8192
+MAX_POSITION_EMBEDDINGS=8192
+
+# Data cache path (useful for both mock and real data)
+DATA_CACHE_PATH="${PWD}/benchmark_cache_llama3_8b_fp8"
+mkdir -p "$DATA_CACHE_PATH"
+
+DISTRIBUTED_ARGS=(
+ --nproc_per_node $GPUS_PER_NODE
+ --nnodes $NUM_NODES
+ --node_rank $NODE_RANK
+ --master_addr $MASTER_ADDR
+ --master_port $MASTER_PORT
+)
+
+MODEL_ARGS=(
+ --use-mcore-models
+ --num-layers $NUM_LAYERS
+ --hidden-size 4096
+ --ffn-hidden-size 14336
+ --num-attention-heads 32
+ --group-query-attention
+ --num-query-groups 8
+ --kv-channels 128
+ --seq-length $SEQ_LENGTH
+ --max-position-embeddings $MAX_POSITION_EMBEDDINGS
+ --position-embedding-type rope
+ --rotary-base 1000000
+ --rotary-percent 1.0
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --swiglu
+ --normalization RMSNorm
+ --init-method-std 0.0134
+ --attention-backend fused
+ --apply-layernorm-1p
+ --untie-embeddings-and-output-weights
+ --disable-bias-linear
+)
+
+TRAINING_ARGS=(
+ --micro-batch-size $MICRO_BATCH_SIZE
+ --global-batch-size $GLOBAL_BATCH_SIZE
+ --train-samples 1953125000
+ --lr-decay-samples 1949218748
+ --lr-warmup-samples 3906252
+ --lr 0.00015
+ --min-lr 0.00001
+ --decoupled-lr 5.0e-4 # Specific to decoupled AdamW, ensure optimizer is compatible
+ --decoupled-min-lr 4.5e-5 # Specific to decoupled AdamW
+ --lr-decay-style cosine
+ --clip-grad 1.0
+ --weight-decay 0.1
+ --adam-beta1 0.9
+ --adam-beta2 0.95
+ --bf16
+ --grad-reduce-in-bf16
+ --cross-entropy-loss-fusion
+ --calculate-per-token-loss
+ --manual-gc
+ --empty-unused-memory-level 1
+ --exit-duration-in-mins 235
+)
+
+# Conditional arguments based on DTYPE (FP8)
+DTYPE_ARGS=()
+if [[ "$DTYPE" == "fp8" ]]; then
+ DTYPE_ARGS+=(
+ "--fp8-format hybrid"
+ "--fp8-amax-history-len 1024"
+ "--fp8-amax-compute-algo max"
+ "--fp8-param-gather"
+ )
+fi
+
+# Model parallelism arguments
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size $TP_SIZE
+ --context-parallel-size $CP_SIZE
+ # --pipeline-model-parallel-size $PP_SIZE # Not explicitly set in llama script options, assume 1 if not multi-node PP
+ --sequence-parallel # Always enable sequence parallelism with TP_SIZE=2
+)
+
+# Distributed Data Parallel (DDP) arguments
+# From original script's ddp_args
+DDP_ARGS=(
+ --use-distributed-optimizer
+ --overlap-grad-reduce
+ --overlap-param-gather
+)
+TRAINING_ARGS+=("${DDP_ARGS[@]}")
+
+
+# Data arguments (conditional for mock vs real data)
+DATA_ARGS_LIST=()
+if [[ "$TOKENIZER_ARG" == "MOCK" ]] || [[ "$DATA_ARG" == "MOCK" ]] || [[ -z "$TOKENIZER_ARG" ]]; then
+ DATA_ARGS_LIST+=(
+ "--mock-data"
+ "--tokenizer-type NullTokenizer"
+ "--vocab-size 128256"
+ "--data-cache-path ${DATA_CACHE_PATH}"
+ "--tiktoken-pattern v2"
+ "--split '99,1,0'"
+ "--no-create-attention-mask-in-dataloader"
+ "--no-mmap-bin-files"
+ "--num-workers 1"
+ )
+else
+ # Settings for real data
+ DATA_ARGS_LIST+=(
+ "--data-path $DATA_ARG"
+ "--tokenizer-type HuggingFaceTokenizer"
+ "--tokenizer-model $TOKENIZER_ARG"
+ "--data-cache-path ${DATA_CACHE_PATH}"
+ "--split '99,1,0'"
+ "--no-create-attention-mask-in-dataloader"
+ "--no-mmap-bin-files"
+ "--num-workers 1"
+ # Note: --vocab-size might be inferred by HuggingFaceTokenizer or might need to be explicit.
+ "--vocab-size 128256"
+ )
+fi
+
+EVAL_AND_LOGGING_ARGS=(
+ --log-interval 1
+ --eval-iters 32
+ --eval-interval 100
+ --save-interval 1000
+ --log-throughput
+ --profile
+ --profile-step-start 4
+ --profile-step-end 6
+ --ckpt-format torch_dist
+ --distributed-timeout-minutes 60
+ --save "$CHECKPOINT_PATH"
+ --load "$CHECKPOINT_PATH"
+ --tensorboard-dir "$TENSORBOARD_LOGS_PATH"
+)
+
+# Ensure pretrain_gpt.py is found
+if [ ! -f "$PRETRAIN_SCRIPT_PATH" ]; then
+ echo "Error: pretrain_gpt.py not found at $PRETRAIN_SCRIPT_PATH"
+ echo "Please ensure you are running this script from the root of the Megatron-LM repository, and pretrain_gpt.py is present."
+ exit 1
+fi
+
+# Run the training command
+torchrun ${DISTRIBUTED_ARGS[@]} \
+ "$PRETRAIN_SCRIPT_PATH" \
+ ${MODEL_ARGS[@]} \
+ ${TRAINING_ARGS[@]} \
+ ${DTYPE_ARGS[@]} \
+ ${MODEL_PARALLEL_ARGS[@]} \
+ ${DATA_ARGS_LIST[@]} \
+ ${EVAL_AND_LOGGING_ARGS[@]}
+
+set +x
\ No newline at end of file
diff --git a/examples/mamba/.gitignore b/examples/mamba/.gitignore
new file mode 100644
index 00000000000..940f4797e4b
--- /dev/null
+++ b/examples/mamba/.gitignore
@@ -0,0 +1,4 @@
+checkpoints/
+data-cache/
+tensorboard/
+triton-cache/
diff --git a/examples/mamba/Dockerfile b/examples/mamba/Dockerfile
new file mode 100644
index 00000000000..2e194095b75
--- /dev/null
+++ b/examples/mamba/Dockerfile
@@ -0,0 +1,32 @@
+FROM nvcr.io/nvidia/pytorch:24.01-py3
+
+RUN pip uninstall -y triton && \
+ pip install triton==2.1.0 sentencepiece==0.1.99 flask-restful
+
+# The causal-conv1d and mamba-ssm packages below are built from scratch here
+# (which takes significant time) because there are no wheels available on PyPI
+# for these relatively newer versions of the packages that are compatible with
+# the older NGC-variant PyTorch version (e.g. version 2.2.0.dev231106) that we
+# are using (in the NGC base container). Generally, if the package is not
+# compatible with the PyTorch version, then it will generate a Python import
+# error. The package authors tend to only release wheels for new versions of
+# these pacakges which are compatible with the versions of regular PyTorch and
+# NGC-variant PyTorch that are newer at the time of release. So, to use newer
+# versions of these packages with relatively older versions of the NGC PyTorch
+# container, we tend to have to build the packages from scratch.
+
+RUN cd /tmp && \
+ git clone https://github.com/Dao-AILab/causal-conv1d.git && \
+ cd causal-conv1d && \
+ git checkout v1.2.2.post1 && \
+ CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install . && \
+ cd .. && \
+ rm -rf causal-conv1d
+
+RUN cd /tmp && \
+ git clone https://github.com/state-spaces/mamba.git && \
+ cd mamba && \
+ git checkout v2.0.3 && \
+ MAMBA_FORCE_BUILD=TRUE pip install . && \
+ cd .. && \
+ rm -rf mamba
diff --git a/examples/mamba/README.md b/examples/mamba/README.md
new file mode 100644
index 00000000000..ce60f119ea5
--- /dev/null
+++ b/examples/mamba/README.md
@@ -0,0 +1,117 @@
+# Mamba-based Language Models
+
+## Introduction
+
+This document is an entrypoint into the code used for
+[An Empirical Study of Mamba-based Language Models](https://arxiv.org/abs/2406.07887).
+
+We are releasing the parameters for some of the models described in that
+technical report via
+[HuggingFace](https://huggingface.co/collections/nvidia/ssms-666a362c5c3bb7e4a6bcfb9c).
+The code in the `main` branch is no longer compatible with the `Mamba2-*`
+checkpoints. You can load them using the
+[fixed snapshot of the code used for the technical report](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba).
+
+## Installation
+
+Create and run a Docker container using the [Dockerfile](./Dockerfile).
+
+```
+docker build -t your_image_name:your_tag .
+docker run --gpus all -it --rm \
+ -v /path/to/megatron:/workspace/megatron \
+ -v /path/to/dataset:/workspace/dataset \
+ -v /path/to/checkpoints:/workspace/checkpoints \
+ -w /workspace/megatron/examples/mamba \
+ your_image_name:your_tag
+```
+
+## Train
+
+[`train.sh`](./train.sh) is an example pretraining script, showing how to run on
+a single node. Select between 800M-scale and 8B-scale models by setting the
+`MODEL_SCALE` variable. The 8B-scale hybrid model architecture is the same as
+the one described in the technical report.
+
+## Text Generation
+
+Use [`run_text_gen_server_8b.sh`](./run_text_gen_server_8b.sh) to start a text
+generation server using an 8B hybrid checkpoint. This is configured to run the
+8B hybrid model described in the technical report, with tensor model parallel
+set to 1.
+
+The arguments in the script will need to be changed if using a checkpoint with a
+different model parallel configuration or other differences, such as model
+architecture. For example, to run the 8B pure Mamba-2 model, change
+`--hybrid-layer-pattern` to use only `M` symbols (e.g., 56 `M`s for the 8B
+model), or remove it entirely.
+
+Use [`run_text_gen_server_8b_gpt3.sh`](./run_text_gen_server_8b_gpt3.sh) to start
+a text generation server using the 8B reference Transformer checkpoint.
+
+## Checkpoint Formats
+
+For inference, the model must be configured to match the checkpoint file used,
+including the hybrid layer configuration and model parallel configuration.
+
+If you need to convert a hybrid checkpoint file to a different tensor parallel
+or pipeline parallel size, use
+[the hybrid conversion script](../../tools/checkpoint/hybrid_conversion.py).
+There is an example run command at the end of that file.
+
+Before running that script, you will need to set `PYTHONPATH` to include the
+root directory of your Megatron-LM repository clone.
+
+```
+export PYTHONPATH=:PYTHONPATH
+```
+
+## Hybrid Options
+
+`--hybrid-layer-pattern PATTERN` specifies the layer type for every layer in
+the model using a string of single-character symbols:
+
+* `M` — Mamba layer
+* `*` — Attention layer
+* `-` — MLP layer
+* `E` — MoE layer
+
+The number of layers is derived from the pattern length, so `--num-layers`
+should not be specified when `--hybrid-layer-pattern` is used.
+
+For example, the 8B hybrid model described in the technical report uses:
+
+```
+--hybrid-layer-pattern "M-M-M--M-M*-M-M-M-M--M*-M-M-M-M-M*--M-M-M-M-M*-M--M-M-M-"
+```
+
+This is a 56-layer model with 4 attention layers, 28 MLP layers, and 24 Mamba
+layers.
+
+A pure Mamba model uses only `M` symbols (e.g., `MMMMMMMM` for 8 layers).
+A pure transformer model uses only `*` and `-` symbols.
+
+### Pipeline parallelism
+
+Use `|` to define pipeline stage boundaries for flexible virtual pipeline
+parallelism (fVPP). For example, `M-M-|M-M*-|M-M-|M-M*-` defines 4 pipeline
+segments. The number of segments must be evenly divisible by
+`--pipeline-model-parallel-size`.
+
+### Multi-Token Prediction (MTP)
+
+Use `/` to append MTP layer patterns. Each pattern after the separator
+represents one MTP prediction depth. For example, `M*M*/MM/MM` has main
+pattern `M*M*` with MTP pattern `MM` repeated for 2 depths.
+
+### Deprecated options
+
+`--hybrid-override-pattern`, `--hybrid-attention-ratio`, and
+`--hybrid-mlp-ratio` are deprecated. Use `--hybrid-layer-pattern` instead.
+
+## Mamba vs Mamba-2
+
+This codebase currently only supports Mamba-2, and not the original version of
+Mamba. However, the
+[fixed snapshot of the code used for the technical report](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba)
+can be configured to run the original version of Mamba.
diff --git a/examples/mamba/run_text_gen_server_8b.sh b/examples/mamba/run_text_gen_server_8b.sh
new file mode 100755
index 00000000000..f183dea4ad1
--- /dev/null
+++ b/examples/mamba/run_text_gen_server_8b.sh
@@ -0,0 +1,50 @@
+#!/bin/bash
+
+# Use: ./run_text_gen_server_8b.sh
+# To launch the client: python ../../tools/text_generation_cli.py
+
+CHECKPOINT_PATH=$1
+TOKENIZER_PATH=$2
+
+HYBRID_LAYER_PATTERN="M-M-M--M-M*-M-M-M-M--M*-M-M-M-M-M*--M-M-M-M-M*-M--M-M-M-"
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr localhost \
+ --master_port 6000"
+
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NCCL_IB_TIMEOUT=19
+export NCCL_IB_QPS_PER_CONNECTION=4
+
+export TRITON_CACHE_DIR="./triton-cache/"
+export TRITON_CACHE_MANAGER="megatron.core.ssm.triton_cache_manager:ParallelFileCacheManager"
+
+torchrun $DISTRIBUTED_ARGS ../../tools/run_hybrid_text_generation_server.py \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --untie-embeddings-and-output-weights \
+ --hybrid-layer-pattern ${HYBRID_LAYER_PATTERN} \
+ --hidden-size 4096 \
+ --load ${CHECKPOINT_PATH} \
+ --num-attention-heads 32 \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --disable-bias-linear \
+ --normalization RMSNorm \
+ --seq-length 4096 \
+ --max-position-embeddings 4096 \
+ --position-embedding-type none \
+ --tokenizer-type GPTSentencePieceTokenizer \
+ --tokenizer-model ${TOKENIZER_PATH} \
+ --distributed-backend nccl \
+ --distributed-timeout-minutes 1440 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --use-mcore-models \
+ --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \
+ --seed 42
diff --git a/examples/mamba/run_text_gen_server_8b_gpt3.sh b/examples/mamba/run_text_gen_server_8b_gpt3.sh
new file mode 100644
index 00000000000..5413b245ed3
--- /dev/null
+++ b/examples/mamba/run_text_gen_server_8b_gpt3.sh
@@ -0,0 +1,46 @@
+#!/bin/bash
+
+# Use: ./run_text_gen_server_8b_gpt3.sh
+# To launch the client: python ../../tools/text_generation_cli.py
+
+CHECKPOINT_PATH=$1
+TOKENIZER_PATH=$2
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr localhost \
+ --master_port 6000"
+
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NCCL_IB_TIMEOUT=19
+export NCCL_IB_QPS_PER_CONNECTION=4
+
+torchrun $DISTRIBUTED_ARGS ../../tools/run_text_generation_server.py \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --use-flash-attn \
+ --apply-layernorm-1p \
+ --untie-embeddings-and-output-weights \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --load ${CHECKPOINT_PATH} \
+ --num-attention-heads 32 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --disable-bias-linear \
+ --seq-length 4096 \
+ --max-position-embeddings 4096 \
+ --position-embedding-type rope \
+ --rotary-percent 0.5 \
+ --squared-relu \
+ --tokenizer-type GPTSentencePieceTokenizer \
+ --tokenizer-model ${TOKENIZER_PATH} \
+ --distributed-backend nccl \
+ --distributed-timeout-minutes 1440 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --use-mcore-models \
+ --transformer-impl local \
+ --seed 42
diff --git a/examples/mamba/train.sh b/examples/mamba/train.sh
new file mode 100755
index 00000000000..f971242ff0b
--- /dev/null
+++ b/examples/mamba/train.sh
@@ -0,0 +1,103 @@
+#!/bin/bash
+
+# Use: ./train.sh
+
+MODEL_SCALE="800M" # or "8B"
+
+case "${MODEL_SCALE}" in
+ "800M")
+ TENSOR_MODEL_PARALLEL_SIZE=1
+ HYBRID_LAYER_PATTERN="M-M-M--M-*M-M-M-M--*M-M-M-M-*M--M-M-M-*M-M--M-M-"
+ HIDDEN_SIZE=1024
+ NUM_ATTENTION_HEADS=16
+ GLOBAL_BATCH_SIZE=32
+ ;;
+ "8B")
+ TENSOR_MODEL_PARALLEL_SIZE=4
+ HYBRID_LAYER_PATTERN="M-M-M--M-M*-M-M-M-M--M*-M-M-M-M-M*--M-M-M-M-M*-M--M-M-M-"
+ HIDDEN_SIZE=4096
+ NUM_ATTENTION_HEADS=32
+ GLOBAL_BATCH_SIZE=8
+ ;;
+ *)
+ echo "Invalid version specified"
+ exit 1
+ ;;
+esac
+
+DATA_PATH=$1
+TOKENIZER_PATH=$2
+
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NCCL_IB_TIMEOUT=19
+export NCCL_IB_QPS_PER_CONNECTION=4
+
+CHECKPOINT_DIR="./checkpoints"
+DATACACHE_DIR="./data-cache"
+TENSORBOARD_DIR="./tensorboard"
+
+mkdir -p ${CHECKPOINT_DIR}
+mkdir -p ${DATACACHE_DIR}
+mkdir -p ${TENSORBOARD_DIR}
+
+export TRITON_CACHE_DIR="./triton-cache/"
+export TRITON_CACHE_MANAGER="megatron.core.ssm.triton_cache_manager:ParallelFileCacheManager"
+
+SEQ_LEN=4096
+TRAIN_SAMPLES=73242188 # 300B tokens / 4096
+LR_WARMUP_SAMPLES=50000
+LR_DECAY_SAMPLES=73192188 # TRAIN_SAMPLES - LR_WARMUP_SAMPLES
+
+options=" \
+ --tensor-model-parallel-size ${TENSOR_MODEL_PARALLEL_SIZE} \
+ --sequence-parallel \
+ --pipeline-model-parallel-size 1 \
+ --use-distributed-optimizer \
+ --overlap-param-gather \
+ --overlap-grad-reduce \
+ --untie-embeddings-and-output-weights \
+ --init-method-std 0.02 \
+ --position-embedding-type none \
+ --hybrid-layer-pattern ${HYBRID_LAYER_PATTERN} \
+ --hidden-size ${HIDDEN_SIZE} \
+ --num-attention-heads ${NUM_ATTENTION_HEADS} \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --seq-length ${SEQ_LEN} \
+ --max-position-embeddings ${SEQ_LEN} \
+ --train-samples ${TRAIN_SAMPLES} \
+ --lr-warmup-samples ${LR_WARMUP_SAMPLES} \
+ --lr-decay-samples ${LR_DECAY_SAMPLES} \
+ --save ${CHECKPOINT_DIR} \
+ --load ${CHECKPOINT_DIR} \
+ --data-path ${DATA_PATH} \
+ --data-cache-path ${DATACACHE_DIR} \
+ --split 99,1,0 \
+ --tokenizer-type GPTSentencePieceTokenizer \
+ --tokenizer-model ${TOKENIZER_PATH} \
+ --distributed-backend nccl \
+ --micro-batch-size 4 \
+ --global-batch-size ${GLOBAL_BATCH_SIZE} \
+ --lr 2.5e-4 \
+ --min-lr 2.5e-5 \
+ --lr-decay-style cosine \
+ --weight-decay 0.1 \
+ --clip-grad 1.0 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --disable-bias-linear \
+ --normalization RMSNorm \
+ --adam-beta1 0.9 \
+ --adam-beta2 0.95 \
+ --log-interval 10 \
+ --save-interval 2000 \
+ --eval-interval 2000 \
+ --eval-iters 32 \
+ --bf16 \
+ --use-mcore-models \
+ --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \
+ --no-create-attention-mask-in-dataloader \
+ --tensorboard-dir ${TENSORBOARD_DIR}"
+
+torchrun --nproc_per_node 8 ../../pretrain_hybrid.py ${options}
diff --git a/examples/megatron_fsdp/README.md b/examples/megatron_fsdp/README.md
new file mode 100644
index 00000000000..eaf5eca1364
--- /dev/null
+++ b/examples/megatron_fsdp/README.md
@@ -0,0 +1,157 @@
+# Megatron-FSDP Examples
+
+Example scripts for training and checkpoint conversion using [Megatron-FSDP](../../docs/user-guide/features/megatron_fsdp.md). These demonstrate recommended configurations for Llama 3 8B and DeepSeek-V3 671B models, as well as checkpoint format conversion between `torch_dist` (N-D parallel) and `fsdp_dtensor` formats.
+
+## Scripts
+
+### `train_llama3_8b_fsdp_h100_fp8.sh`
+
+Single-node training script for **Llama 3 8B** using Megatron-FSDP with FP8 precision on H100 GPUs. Uses `torchrun` for local distributed training and supports both mock data (for benchmarking) and real datasets.
+
+#### Usage
+
+Run from the root of the Megatron-LM repository:
+
+```bash
+# With mock data (default, for benchmarking)
+bash examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh
+
+# With real data
+bash examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh \
+ checkpoints/llama3_8b_fsdp_fp8 \
+ tensorboard_logs/llama3_8b_fsdp_fp8 \
+ /path/to/tokenizer \
+ /path/to/data_prefix
+```
+
+| Positional Argument | Default | Description |
+|---------------------|---------|-------------|
+| `$1` — Checkpoint path | `checkpoints/llama3_8b_fsdp_fp8` | Directory for saving and loading checkpoints. |
+| `$2` — TensorBoard path | `tensorboard_logs/llama3_8b_fsdp_fp8` | Directory for TensorBoard logs. |
+| `$3` — Tokenizer | `MOCK` | Path to a tokenizer model, or `MOCK` for `NullTokenizer`. |
+| `$4` — Data path | `MOCK` | Data prefix for training data, or `MOCK` for mock data. |
+
+#### Environment Variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `USE_MEGATRON_FSDP` | `1` | Set to `1` to enable Megatron-FSDP. Set to `0` to train with standard DDP. |
+| `SHARDING_STRATEGY` | `optim_grads_params` | FSDP sharding strategy (ZeRO-3). Options: `no_shard`, `optim`, `optim_grads`, `optim_grads_params`. |
+| `OUTER_SHARDING_STRATEGY` | `no_shard` | DP-Outer sharding strategy for HSDP/HFSDP. Options: `no_shard`, `optim`. |
+| `MASTER_ADDR` | `localhost` | Master node address for distributed training. |
+| `MASTER_PORT` | `6000` | Master node port. |
+| `NODE_RANK` | `0` | Rank of the current node. |
+
+#### Configuration Summary
+
+- **Model**: Llama 3 8B (GQA with 32 heads / 8 KV groups, RoPE, SwiGLU, RMSNorm)
+- **Parallelism**: TP=1, CP=1, PP=1, 8 GPUs per node, FSDP ZeRO-3
+- **Precision**: FP8 (hybrid format) with BF16 training and BF16 gradient reduction
+- **Batch size**: micro-batch=1, global-batch=128, sequence length=8192
+- **Optimizations**: NCCL user buffers, FSDP double buffering, manual registration, meta-device initialization, per-token loss, overlapped grad-reduce and param-gather
+
+---
+
+### `sbatch_mfsdp_deepseek_v3.sh`
+
+Multi-node SLURM training script for **DeepSeek-V3** (671B MoE) using Megatron-FSDP. Submits an `sbatch` job with containerized execution via `srun`.
+
+#### Usage
+
+Set the required configuration variables and submit:
+
+```bash
+export MEGATRON_PATH=/path/to/Megatron-LM
+export CONTAINER_IMAGE=/path/to/container.sqsh # or docker image URL
+export OUTPUT_PATH=/path/to/output
+export DATA_PATH=/path/to/training/data
+
+bash examples/megatron_fsdp/sbatch_mfsdp_deepseek_v3.sh
+```
+
+Before running, update the `#SBATCH` directives and `--container-mounts` in the script to match your cluster configuration.
+
+#### Environment Variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MEGATRON_PATH` | *(required)* | Path to the Megatron-LM repository. |
+| `CONTAINER_IMAGE` | *(required)* | Container image (`.sqsh` file or Docker URL). |
+| `OUTPUT_PATH` | *(required)* | Base directory for checkpoints, TensorBoard logs, SLURM logs, and Nsight profiles. |
+| `DATA_PATH` | *(required)* | Training data prefix path. |
+| `USE_MEGATRON_FSDP` | `1` | Enable Megatron-FSDP. Set to `0` for standard DDP. |
+| `SHARDING_STRATEGY` | `optim_grads_params` | FSDP sharding strategy (ZeRO-3). |
+| `TP` | `1` | Tensor parallel size. |
+| `EP` | `8` | Expert parallel size. |
+| `MBS` | `4` | Micro-batch size. |
+| `GBS` | `2048` | Global batch size. |
+| `PROFILE` | `0` | Set to `1` to enable Nsight Systems profiling (steps 10–12). |
+| `WANDB` | `1` | Set to `1` to enable Weights & Biases logging. Requires `WANDB_API_KEY`. |
+| `COMMENT` | N/A | Tag appended to W&B experiment names and Nsight profile filenames. |
+
+#### Configuration Summary
+
+- **Model**: DeepSeek-V3 (61 layers, 256 routed experts, top-8 routing, Multi-Latent Attention, MTP)
+- **Parallelism**: TP=1, EP=8, CP=1, FSDP ZeRO-3
+- **Precision**: BF16
+- **MoE**: Flex dispatcher with HybridEP backend, grouped GEMM, sigmoid routing with expert bias, auxiliary sequence loss
+- **Recomputation**: Selective recomputation of `mlp`, `moe`, `mla_up_proj`, and `layernorm` modules
+- **Optimizations**: NCCL user buffers, FSDP double buffering, meta-device initialization, per-token loss, overlapped grad-reduce and param-gather
+- **Tokenizer**: `deepseek-ai/DeepSeek-V3` via HuggingFace
+
+---
+
+### `sbatch_checkpoint_convert.sh`
+
+SLURM batch script for converting checkpoints from **`torch_dist`** (N-D parallel) format to **`fsdp_dtensor`** (Megatron-FSDP) format. This enables resuming training under Megatron-FSDP from checkpoints originally saved with tensor/pipeline/expert parallelism.
+
+#### Prerequisites
+
+Before converting, you need a `param_to_param_group_map.json` file. Generate it by running a `torch_dist` training job with the `--dump-param-to-param-group-map` flag, then converting the output:
+
+```bash
+# 1. Run a training job (or trivial experiment) with the dump flag
+--dump-param-to-param-group-map /path/to/param_to_param_group_map
+
+# 2. Convert the dumped map to JSON
+python tools/checkpoint/checkpoint_inspector.py \
+ print-torch-dcp-in-json /path/to/param_to_param_group_map
+```
+
+See the [Checkpoint Conversion](../../docs/user-guide/features/megatron_fsdp.md#checkpoint-conversion) section in the Megatron-FSDP docs for details.
+
+#### Usage
+
+Set the required configuration variables, update the checkpoint paths in `RUN_CMD`, and submit:
+
+```bash
+export MEGATRON_PATH=/path/to/Megatron-LM
+export CONTAINER_IMAGE=/path/to/container.sqsh
+export OUTPUT_PATH=/path/to/output
+
+bash examples/megatron_fsdp/sbatch_checkpoint_convert.sh
+```
+
+Before running, you must edit the script to fill in:
+- The input `torch_dist` checkpoint path
+- The output `fsdp_dtensor` checkpoint path
+- The path to `param_to_param_group_map.json`
+- The `#SBATCH` directives and `--container-mounts` for your cluster
+
+#### Environment Variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `MEGATRON_PATH` | *(required)* | Path to the Megatron-LM repository. |
+| `CONTAINER_IMAGE` | *(required)* | Container image (`.sqsh` file or Docker URL). |
+| `OUTPUT_PATH` | *(required)* | Base directory for SLURM logs. |
+
+#### Conversion Command
+
+The script runs `checkpoint_inspector.py convert-torch-dist-to-fsdp-dtensor` with the `--swiglu` flag (for models using SwiGLU activations). Remove `--swiglu` if converting a non-SwiGLU model.
+
+## Further Reading
+
+- [Megatron-FSDP User Guide](../../docs/user-guide/features/megatron_fsdp.md) — full feature guide, API reference, and sharding strategy documentation.
+- [Megatron-FSDP on PyPI](https://pypi.org/project/megatron-fsdp/) — standalone `fully_shard` API.
+- [Megatron-FSDP Source](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/distributed/fsdp/src) — implementation source code.
diff --git a/examples/megatron_fsdp/sbatch_checkpoint_convert.sh b/examples/megatron_fsdp/sbatch_checkpoint_convert.sh
new file mode 100644
index 00000000000..9f302c93f8f
--- /dev/null
+++ b/examples/megatron_fsdp/sbatch_checkpoint_convert.sh
@@ -0,0 +1,50 @@
+#!/bin/bash
+
+# Configuration: Set these paths before running the script
+MEGATRON_PATH=${MEGATRON_PATH:-"your_own_megatron_path"} # Path to Megatron-LM repository
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-"your_own_container_image"} # Path to .sqsh or docker image url
+OUTPUT_PATH=${OUTPUT_PATH:-"your_own_output_path"} # Path for SLURM logs
+
+# Checkpoint conversion command
+# Note: Update the checkpoint paths in the command below
+RUN_CMD="
+cd ${MEGATRON_PATH};
+git rev-parse HEAD;
+export PYTHONPATH=${MEGATRON_PATH}:${PYTHONPATH};
+python3 tools/checkpoint/checkpoint_inspector.py \
+ convert-torch-dist-to-fsdp-dtensor --swiglu \
+ your_own_path_to_input_torch_dist_checkpoint \
+ your_own_path_to_output_fsdp_dtensor_checkpoint \
+ --param-to-param-group-map-json your_own_path_to_param_to_param_group_map.json"
+
+# SLURM settings
+SLURM_LOGS="${OUTPUT_PATH}/slurm_logs"
+mkdir -p ${SLURM_LOGS} || {
+ echo "Error: Failed to create SLURM logs directory ${SLURM_LOGS}"
+ exit 1
+}
+
+# Submit SLURM job
+# Note: Update SBATCH parameters below according to your cluster configuration
+set +e
+sbatch <&1 | tee ${SLURM_LOGS}/\${SLURM_JOB_ID}.log
+
+EOF
+set -e
diff --git a/examples/megatron_fsdp/sbatch_mfsdp_deepseek_v3.sh b/examples/megatron_fsdp/sbatch_mfsdp_deepseek_v3.sh
new file mode 100644
index 00000000000..22a8f22f68c
--- /dev/null
+++ b/examples/megatron_fsdp/sbatch_mfsdp_deepseek_v3.sh
@@ -0,0 +1,223 @@
+#!/bin/bash
+
+export NCCL_IB_SL=1
+export NCCL_IB_TIMEOUT=19
+export NVTE_FWD_LAYERNORM_SM_MARGIN=16
+export NVTE_BWD_LAYERNORM_SM_MARGIN=16
+export NCCL_P2P_NET_CHUNKSIZE=2097152
+export TORCH_NCCL_AVOID_RECORD_STREAMS=1
+export PYTHONWARNINGS=ignore
+export TRITON_CACHE_DIR=/tmp/triton_cache_$SLURM_NODEID
+
+# Configuration: Set these variables before running the script
+MEGATRON_PATH=${MEGATRON_PATH:-"your_own_megatron_path"} # Path to Megatron-LM repository
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-"your_own_container_image"} # Path to .sqsh or docker image url
+OUTPUT_PATH=${OUTPUT_PATH:-"your_own_output_path"} # Path for output logs and checkpoints
+DATA_PATH=${DATA_PATH:-"your_own_data_path"}
+USE_MEGATRON_FSDP=${USE_MEGATRON_FSDP:-1}
+SHARDING_STRATEGY=${SHARDING_STRATEGY:-"optim_grads_params"}
+PROFILE=${PROFILE:-0}
+WANDB=${WANDB:-1}
+
+TP=${TP:-1}
+EP=${EP:-8}
+MBS=${MBS:-4}
+GBS=${GBS:-2048}
+COMMENT=${COMMENT:-""}
+
+PRETRAIN_ARGS=(
+ --distributed-timeout-minutes 60
+ --tensor-model-parallel-size ${TP}
+ --expert-model-parallel-size ${EP}
+ --expert-tensor-parallel-size 1
+ --context-parallel-size 1
+ --use-distributed-optimizer
+ --overlap-grad-reduce
+ --overlap-param-gather
+ --use-mcore-models
+ --sequence-parallel
+ --use-flash-attn
+ --disable-bias-linear
+ --micro-batch-size ${MBS}
+ --global-batch-size ${GBS}
+ --train-samples 585937500
+ --exit-duration-in-mins 220
+ --no-check-for-nan-in-loss-and-grad
+ --manual-gc
+ --manual-gc-interval 10
+ --recompute-granularity selective
+ --recompute-modules mlp moe mla_up_proj layernorm
+ --transformer-impl transformer_engine
+ --seq-length 4096
+ --data-cache-path ${OUTPUT_PATH}/cache
+ --tokenizer-type HuggingFaceTokenizer
+ --tokenizer-model deepseek-ai/DeepSeek-V3
+ --data-path ${DATA_PATH}
+ --split 99,1,0
+ --no-mmap-bin-files
+ --no-create-attention-mask-in-dataloader
+ --num-workers 6
+ --num-layers 61
+ --hidden-size 7168
+ --ffn-hidden-size 18432
+ --num-attention-heads 128
+ --kv-channels 128
+ --max-position-embeddings 4096
+ --position-embedding-type rope
+ --rotary-base 10000
+ --make-vocab-size-divisible-by 3232
+ --normalization RMSNorm
+ --norm-epsilon 1e-6
+ --swiglu
+ --untie-embeddings-and-output-weights
+ --multi-latent-attention
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --clip-grad 1.0
+ --weight-decay 0.1
+ --qk-layernorm
+ --lr-decay-samples 584765624
+ --lr-warmup-samples 1536000
+ --lr-warmup-init 3.9e-7
+ --lr 3.9e-6
+ --min-lr 3.9e-7
+ --lr-decay-style cosine
+ --adam-beta1 0.9
+ --adam-beta2 0.95
+ --num-experts 256
+ --moe-layer-freq [0]*3+[1]*58
+ --moe-ffn-hidden-size 2048
+ --moe-shared-expert-intermediate-size 2048
+ --moe-router-load-balancing-type seq_aux_loss
+ --moe-router-topk 8
+ --moe-token-dispatcher-type flex
+ --moe-flex-dispatcher-backend hybridep
+ --moe-router-pre-softmax
+ --moe-grouped-gemm
+ --moe-aux-loss-coeff 1e-4
+ --moe-router-group-topk 4
+ --moe-router-num-groups 8
+ --moe-router-topk-scaling-factor 2.5
+ --moe-router-score-function sigmoid
+ --moe-router-enable-expert-bias
+ --moe-router-bias-update-rate 1e-3
+ --moe-router-dtype fp32
+ --moe-permute-fusion
+ --moe-router-force-load-balancing
+ --q-lora-rank 1536
+ --kv-lora-rank 512
+ --qk-head-dim 128
+ --qk-pos-emb-head-dim 64
+ --v-head-dim 128
+ --rotary-scaling-factor 40
+ --mscale 1.0
+ --mscale-all-dim 1.0
+ --mtp-num-layers 1
+ --mtp-loss-scaling-factor 0.1
+ --eval-iters 32
+ --eval-interval 100
+ --auto-detect-ckpt-format
+ --load ${OUTPUT_PATH}/checkpoints
+ --save ${OUTPUT_PATH}/checkpoints
+ --save-interval 100
+ --dist-ckpt-strictness log_all
+ --init-method-std 0.02
+ --log-timers-to-tensorboard
+ --log-memory-to-tensorboard
+ --log-num-zeros-in-grad
+ --log-params-norm
+ --log-validation-ppl-to-tensorboard
+ --log-throughput
+ --log-interval 1
+ --logging-level 40
+ --tensorboard-dir ${OUTPUT_PATH}/tensorboard
+ --bf16
+ --enable-experimental
+)
+
+if [ "${USE_MEGATRON_FSDP}" = 1 ]; then
+ unset CUDA_DEVICE_MAX_CONNECTIONS
+ PRETRAIN_ARGS=(
+ "${PRETRAIN_ARGS[@]}"
+ --use-megatron-fsdp
+ --data-parallel-sharding-strategy ${SHARDING_STRATEGY}
+ --no-gradient-accumulation-fusion
+ --use-distributed-optimizer
+ --calculate-per-token-loss
+ --init-model-with-meta-device
+ --ckpt-format fsdp_dtensor
+ --grad-reduce-in-bf16
+ --fsdp-double-buffer
+ --use-nccl-ub
+ )
+fi
+
+# Profiling command
+if [ "${PROFILE}" = 1 ]; then
+ PROFILE_CMD="nsys profile --sample=none --cpuctxsw=none --trace=cuda,nvtx,cublas,cudnn \
+ --capture-range=cudaProfilerApi \
+ --capture-range-end=stop \
+ --cuda-graph-trace=node \
+ --cuda-memory-usage=true \
+ -f true -x true \
+ -o ${OUTPUT_PATH}/nsys/Megatron-FSDP-Deepseek-V3-TP${TP}EP${EP}-MBS${MBS}GBS${GBS}-${COMMENT}"
+ PRETRAIN_ARGS=(
+ "${PRETRAIN_ARGS[@]}"
+ --profile
+ --profile-step-start 10
+ --profile-step-end 12
+ --profile-ranks 0
+ )
+ echo "PROFILE_CMD="
+ echo $PROFILE_CMD
+else
+ PROFILE_CMD=""
+fi
+
+if [ "${WANDB}" = 1 ]; then
+ export WANDB_API_KEY=${WANDB_API_KEY:-"your_own_wandb_api_key"}
+ PRETRAIN_ARGS=(
+ "${PRETRAIN_ARGS[@]}"
+ --wandb-project your_own_wandb_project
+ --wandb-exp-name DeepSeek-V3-TP${TP}EP${EP}-MBS${MBS}GBS${GBS}-${COMMENT}
+ )
+fi
+
+TRAINING_CMD="
+cd ${MEGATRON_PATH};
+git rev-parse HEAD;
+export PYTHONPATH=${MEGATRON_PATH}:${PYTHONPATH};
+${PROFILE_CMD} python ${MEGATRON_PATH}/pretrain_gpt.py ${PRETRAIN_ARGS[@]}"
+
+# SLURM settings
+SLURM_LOGS="${OUTPUT_PATH}/slurm_logs"
+mkdir -p ${SLURM_LOGS} || {
+ echo "Error: Failed to create SLURM logs directory ${SLURM_LOGS}"
+ exit 1
+}
+
+# Submit SLURM job
+# Note: Update SBATCH parameters below according to your cluster configuration
+set +e
+sbatch <&1 | tee ${SLURM_LOGS}/\${SLURM_JOB_ID}.log
+
+EOF
+set -e
diff --git a/examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh b/examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh
new file mode 100644
index 00000000000..ddd3f160fa7
--- /dev/null
+++ b/examples/megatron_fsdp/train_llama3_8b_fsdp_h100_fp8.sh
@@ -0,0 +1,212 @@
+#!/bin/bash
+
+CHECKPOINT_PATH=${1:-"checkpoints/llama3_8b_fsdp_fp8"}
+TENSORBOARD_LOGS_PATH=${2:-"tensorboard_logs/llama3_8b_fsdp_fp8"}
+TOKENIZER_ARG=${3:-"MOCK"} # Path to tokenizer model, or "MOCK"
+DATA_ARG=${4:-"MOCK"} # Data prefix, or "MOCK"
+
+# Create directories if they don't exist
+mkdir -p "$(dirname "$CHECKPOINT_PATH")"
+mkdir -p "$(dirname "$TENSORBOARD_LOGS_PATH")"
+
+# Distributed training setup
+GPUS_PER_NODE=8
+NUM_NODES=1
+MASTER_ADDR=${MASTER_ADDR:-localhost}
+MASTER_PORT=${MASTER_PORT:-6000}
+NODE_RANK=${NODE_RANK:-0}
+WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
+
+# Path to the pretrain_gpt.py script, assuming this script
+# is run from the root of the Megatron-LM repository.
+PRETRAIN_SCRIPT_PATH="pretrain_gpt.py"
+
+# Model & Training Parameters
+USE_MEGATRON_FSDP=${USE_MEGATRON_FSDP:-1}
+SHARDING_STRATEGY=${SHARDING_STRATEGY:-"optim_grads_params"}
+OUTER_SHARDING_STRATEGY=${OUTER_SHARDING_STRATEGY:-"no_shard"}
+TP_SIZE=1
+CP_SIZE=1
+PP_SIZE=1
+MICRO_BATCH_SIZE=1
+GLOBAL_BATCH_SIZE=128
+NUM_LAYERS=32
+DTYPE="fp8"
+SEQ_LENGTH=8192
+MAX_POSITION_EMBEDDINGS=8192
+
+# Data cache path (useful for both mock and real data)
+DATA_CACHE_PATH="${PWD}/benchmark_cache_llama3_8b_fsdp_fp8"
+mkdir -p "$DATA_CACHE_PATH"
+
+DISTRIBUTED_ARGS=(
+ --nproc_per_node $GPUS_PER_NODE
+ --nnodes $NUM_NODES
+ --node_rank $NODE_RANK
+ --master_addr $MASTER_ADDR
+ --master_port $MASTER_PORT
+)
+
+MODEL_ARGS=(
+ --use-mcore-models
+ --num-layers $NUM_LAYERS
+ --hidden-size 4096
+ --ffn-hidden-size 14336
+ --num-attention-heads 32
+ --group-query-attention
+ --num-query-groups 8
+ --kv-channels 128
+ --seq-length $SEQ_LENGTH
+ --max-position-embeddings $MAX_POSITION_EMBEDDINGS
+ --position-embedding-type rope
+ --rotary-base 1000000
+ --rotary-percent 1.0
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --swiglu
+ --normalization RMSNorm
+ --init-method-std 0.0134
+ --attention-backend fused
+ --apply-layernorm-1p
+ --untie-embeddings-and-output-weights
+ --disable-bias-linear
+)
+
+TRAINING_ARGS=(
+ --micro-batch-size $MICRO_BATCH_SIZE
+ --global-batch-size $GLOBAL_BATCH_SIZE
+ --train-samples 1953125000
+ --lr-decay-samples 1949218748
+ --lr-warmup-samples 3906252
+ --lr 0.00015
+ --min-lr 0.00001
+ --decoupled-lr 5.0e-4
+ --decoupled-min-lr 4.5e-5
+ --lr-decay-style cosine
+ --clip-grad 1.0
+ --weight-decay 0.1
+ --adam-beta1 0.9
+ --adam-beta2 0.95
+ --bf16
+ --cross-entropy-loss-fusion
+ --manual-gc
+ --empty-unused-memory-level 1
+ --exit-duration-in-mins 235
+)
+
+if [ "${USE_MEGATRON_FSDP}" = 1 ]; then
+ unset CUDA_DEVICE_MAX_CONNECTIONS
+ TRAINING_ARGS=(
+ "${TRAINING_ARGS[@]}"
+ --use-megatron-fsdp
+ --data-parallel-sharding-strategy ${SHARDING_STRATEGY}
+ --no-gradient-accumulation-fusion
+ --calculate-per-token-loss
+ --init-model-with-meta-device
+ --ckpt-format fsdp_dtensor
+ --grad-reduce-in-bf16
+ --use-nccl-ub
+ --fsdp-double-buffer
+ --fsdp-manual-registration
+ # To enable HFSDP, DP full-sharding of the optimizer state with
+ # hierarchical data parallelism (DP-Outer=2, DP-Inner=DP//2)...
+ # --num-distributed-optimizer-instances 2
+ # --outer-dp-sharding-strategy ${OUTER_SHARDING_STRATEGY}
+ # To further customize Megatron-FSDP data precision...
+ # --megatron-fsdp-main-params-dtype fp32
+ # --megatron-fsdp-main-grads-dtype auto
+ # --megatron-fsdp-grad-comm-dtype auto
+ )
+fi
+
+# Conditional arguments based on DTYPE (FP8)
+DTYPE_ARGS=()
+if [[ "$DTYPE" == "fp8" ]]; then
+ DTYPE_ARGS+=(
+ "--fp8-format hybrid"
+ "--fp8-amax-history-len 1024"
+ "--fp8-amax-compute-algo max"
+ "--fp8-param-gather"
+ )
+fi
+
+# Model parallelism arguments
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size $TP_SIZE
+ --context-parallel-size $CP_SIZE
+ --sequence-parallel
+)
+
+# Distributed Data Parallel (DDP) arguments
+# From original script's ddp_args
+DDP_ARGS=(
+ --use-distributed-optimizer
+ --overlap-grad-reduce
+ --overlap-param-gather
+)
+TRAINING_ARGS+=("${DDP_ARGS[@]}")
+
+
+# Data arguments (conditional for mock vs real data)
+DATA_ARGS_LIST=()
+if [[ "$TOKENIZER_ARG" == "MOCK" ]] || [[ "$DATA_ARG" == "MOCK" ]] || [[ -z "$TOKENIZER_ARG" ]]; then
+ DATA_ARGS_LIST+=(
+ "--mock-data"
+ "--tokenizer-type NullTokenizer"
+ "--vocab-size 128256"
+ "--data-cache-path ${DATA_CACHE_PATH}"
+ "--tiktoken-pattern v2"
+ "--split '99,1,0'"
+ "--no-create-attention-mask-in-dataloader"
+ "--no-mmap-bin-files"
+ "--num-workers 1"
+ )
+else
+ # Settings for real data
+ DATA_ARGS_LIST+=(
+ "--data-path $DATA_ARG"
+ "--tokenizer-type HuggingFaceTokenizer"
+ "--tokenizer-model $TOKENIZER_ARG"
+ "--data-cache-path ${DATA_CACHE_PATH}"
+ "--split '99,1,0'"
+ "--no-create-attention-mask-in-dataloader"
+ "--no-mmap-bin-files"
+ "--num-workers 1"
+ # Note: --vocab-size might be inferred by HuggingFaceTokenizer or might need to be explicit.
+ "--vocab-size 128256"
+ )
+fi
+
+EVAL_AND_LOGGING_ARGS=(
+ --log-interval 1
+ --eval-iters 32
+ --eval-interval 100
+ --save-interval 1000
+ --log-throughput
+ --profile
+ --profile-step-start 4
+ --profile-step-end 6
+ --distributed-timeout-minutes 60
+ --save "$CHECKPOINT_PATH"
+ --load "$CHECKPOINT_PATH"
+ --tensorboard-dir "$TENSORBOARD_LOGS_PATH"
+)
+
+# Ensure pretrain_gpt.py is found
+if [ ! -f "$PRETRAIN_SCRIPT_PATH" ]; then
+ echo "Error: pretrain_gpt.py not found at $PRETRAIN_SCRIPT_PATH"
+ echo "Please ensure you are running this script from the root of the Megatron-LM repository, and pretrain_gpt.py is present."
+ exit 1
+fi
+
+# Run the training command
+torchrun ${DISTRIBUTED_ARGS[@]} \
+ "$PRETRAIN_SCRIPT_PATH" \
+ ${MODEL_ARGS[@]} \
+ ${TRAINING_ARGS[@]} \
+ ${DTYPE_ARGS[@]} \
+ ${MODEL_PARALLEL_ARGS[@]} \
+ ${DATA_ARGS_LIST[@]} \
+ ${EVAL_AND_LOGGING_ARGS[@]}
+
+set +x
\ No newline at end of file
diff --git a/examples/merge_mp_bert.sh b/examples/merge_mp_bert.sh
deleted file mode 100755
index 1383433284b..00000000000
--- a/examples/merge_mp_bert.sh
+++ /dev/null
@@ -1,18 +0,0 @@
-#!/bin/bash
-
-TENSOR_MODEL_PARALLEL_SIZE=2
-
-VOCAB_FILE=bert-vocab.txt
-CHECKPOINT_PATH=checkpoints/bert_345m
-
-WORLD_SIZE=$TENSOR_MODEL_PARALLEL_SIZE python tools/merge_mp_partitions.py \
- --model-type BERT \
- --tensor-model-parallel-size $TENSOR_MODEL_PARALLEL_SIZE \
- --tokenizer-type BertWordPieceLowerCase \
- --vocab-file $VOCAB_FILE \
- --num-layers 24 \
- --hidden-size 1024 \
- --num-attention-heads 16 \
- --seq-length 512 \
- --max-position-embeddings 512 \
- --load $CHECKPOINT_PATH
diff --git a/examples/mimo/__init__.py b/examples/mimo/__init__.py
new file mode 100644
index 00000000000..0519ecba6ea
--- /dev/null
+++ b/examples/mimo/__init__.py
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/examples/mimo/avlm_inference.py b/examples/mimo/avlm_inference.py
new file mode 100644
index 00000000000..14ca2068cac
--- /dev/null
+++ b/examples/mimo/avlm_inference.py
@@ -0,0 +1,244 @@
+import argparse
+import os
+from pathlib import Path
+from typing import Union
+
+# hf path
+import requests
+import torch
+from PIL import Image
+from transformers import AutoProcessor
+from transformers import AutoTokenizer
+import soundfile as sf
+import io
+import numpy as np
+import scipy.signal as signal
+
+from examples.mimo.model_providers.llava_avlm import model_provider_llava_avlm
+from megatron.core import dist_checkpointing, parallel_state, tensor_parallel
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.training import print_rank_0
+from examples.mimo.data.utils.calculate_audio_tokens import calculate_num_audio_tokens
+
+def init_distributed(tp_size: int = 1, pp_size: int = 1):
+ if torch.distributed.is_initialized():
+ return
+ rank = int(os.environ.get("LOCAL_RANK", 0))
+ world_size = int(os.environ.get("WORLD_SIZE", 1))
+ torch.cuda.set_device(rank % torch.cuda.device_count())
+ torch.distributed.init_process_group("nccl", rank=rank, world_size=world_size)
+ parallel_state.initialize_model_parallel(tp_size, pp_size)
+
+def get_input_data(
+ processor: AutoProcessor,
+ image_processor: AutoProcessor,
+ audio_processor: AutoProcessor,
+ audio_path: str,
+ image_path: str,
+ prompt: str,
+ device: Union[int, str] = 0):
+ """
+ Prepare inputs for the MIMO model forward pass.
+ """
+
+ def read_audio(audio_path):
+ """Process audio file and return tensor."""
+ with open(audio_path, 'rb') as f:
+ audio_bytes = f.read()
+ audio_io = io.BytesIO(audio_bytes)
+ waveform, sample_rate = sf.read(audio_io)
+
+ # Resample if needed
+ fixed_sample_rate = 16000
+ if sample_rate != fixed_sample_rate:
+ num_samples = int(len(waveform) * fixed_sample_rate / sample_rate)
+ waveform = signal.resample(waveform, num_samples)
+
+ # Convert to tensor
+ audio_tensor = torch.from_numpy(waveform).float()
+ return audio_tensor
+
+ def read_image(image_path):
+ """Process image file and return tensor."""
+ with open(image_path, 'rb') as f:
+ image_bytes = f.read()
+ image_io = io.BytesIO(image_bytes)
+ image = Image.open(image_io)
+ image_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1) # Convert to CxHxW format
+ image_tensor = image_tensor.float() / 255.0 # rescale to [0,1] range
+ return image_tensor
+
+
+ # read audio and image
+ audio_tensor = read_audio(audio_path)
+ image_tensor = read_image(image_path)
+
+ # set up prompt
+ conversation = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": prompt},
+ ],
+ }
+ ]
+ prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+
+ # process audio
+ processed_audios = audio_processor(audio_tensor, sampling_rate=16000)
+ processed_audios = torch.tensor(processed_audios["input_features"])
+ processed_audios = processed_audios.squeeze(0) # remove batch dim
+ num_audio_tokens = calculate_num_audio_tokens(audio_tensor.unsqueeze(0), "openai/whisper-base")
+ audios_seq_lengths = torch.tensor(num_audio_tokens)
+ prompt = prompt.replace("