-Megatron-LM and Megatron Core
-=============================
+Megatron-LM & Megatron Core
+===========================
-[](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html)
-[](./CHANGELOG.md)
+[](https://docs.nvidia.com/Megatron-Core/developer-guide/latest/index.html)
+[](./CHANGELOG.md)
[](./LICENSE)
-## About
+> ## 🚨 **DEVELOPMENT BRANCH**
+> ⚠️ **EXPERIMENTAL FEATURES** - This is the **dev branch** with experimental features.
+>
+> **→ For releases and comprehensive documentation, visit the [main branch](https://github.com/NVIDIA/Megatron-LM)**
-This repository contains two components: **Megatron-LM** and **Megatron Core**.
+## ⚡ Quickstart
-**Megatron-LM** is a reference example that includes Megatron Core plus pre-configured training scripts, ideal for research teams, learning distributed training, and quick experimentation.
+```bash
+# Clone the dev branch
+git clone -b dev https://github.com/NVIDIA/Megatron-LM.git
+cd Megatron-LM
+
+# Install from source with dev dependencies (includes transformer_engine)
+pip install -e .[mlm,dev]
+```
-**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, and CP), mixed precision support (FP16, BF16, FP8, and FP4), and model architectures, ideal for framework developers and ML engineers building custom training pipelines.
+**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.
**[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** provides bidirectional Hugging Face ↔ Megatron checkpoint conversion with production-ready recipes.
@@ -37,9 +47,9 @@ cd Megatron-LM
uv pip install -e .
```
-> **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` (for example, `MAX_JOBS=4 uv pip install -e .`).
+> **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 .`).
-For NVIDIA GPU Cloud (NGC) container setup and all installation options, review the **[Installation Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/get-started/install.html)**.
+For NGC container setup and all installation options, see the **[Installation Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/get-started/install.html)**.
- **[Your First Training Run](https://docs.nvidia.com/megatron-core/developer-guide/latest/get-started/quickstart.html)** - End-to-end training examples with data preparation
- **[Parallelism Strategies](https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/parallelism-guide.html)** - Scale training across GPUs with TP, PP, DP, EP, and CP
@@ -47,21 +57,31 @@ For NVIDIA GPU Cloud (NGC) container setup and all installation options, review
# Latest News
-- **[2026/05]** **[DeepSeek-V4 initial support](https://github.com/NVIDIA/Megatron-LM/issues/4468)** - Megatron Core's `dev` branch includes the initial DeepSeek-V4 implementation; Megatron Bridge provides [conversion, inference, and pretraining recipes](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/deepseek_v4).
-- **[2026/04]** **[Advancing Emerging Optimizers for Accelerated LLM Training with NVIDIA Megatron](https://developer.nvidia.com/blog/advancing-emerging-optimizers-for-accelerated-llm-training-with-nvidia-megatron/)** - Muon and other emerging optimizers are now supported in Megatron Core via the new **[Emerging-Optimizers](https://github.com/NVIDIA-NeMo/Emerging-Optimizers)** library.
-- **[2026/03]** **[Scalable Training of Mixture-of-Experts Models with Megatron Core](https://arxiv.org/abs/2603.07685)** - Technical report on scaling MoE training with integrated optimizations for memory, communication, and computation.
-- **[2026/03]** **[Implementing Falcon-H1 Hybrid Architecture in Megatron Core](https://developer.nvidia.com/blog/implementing-falcon-h1-hybrid-architecture-in-nvidia-megatron-core/)** - Technology Innovation Institute (TII) contributes Falcon-H1 hybrid transformer-Mamba architecture and BitNet ternary quantization support to Megatron Core.
-- **[2026/03]** **[Megatron Core Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/4003)** - Roadmap for upcoming Megatron Core features and improvements.
-- **[2026/03]** **Deprecating Python 3.10 support:** The upcoming 0.17.0 release drops Python 3.10 support. Downstream applications must raise their lower boundary to 3.12 to stay compatible with Megatron Core.
+- **[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 happen in the open, and community contributions are welcome.
-- **[2025/10]** **[Megatron Dev Branch](https://github.com/NVIDIA/Megatron-LM/tree/dev)** - Early access branch with experimental features.
+- **[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]** **[Mixture of Experts (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)** - Megatron Core integrates advanced features including YaRN RoPE scaling, attention sinks, and custom activation functions.
+- **[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/)).
+
+
+Table of Contents
+
+**Getting Started**
+- [⚡ Quick Start](#-quick-start)
+- [🧠 Dev Branch Philosophy](#-dev-branch-philosophy)
+- [📊 Performance & Benchmarking](#-performance--benchmarking)
+- [👥 Community & Support](#-community--support)
+
+**For Complete Documentation** → [Main Branch](https://github.com/NVIDIA/Megatron-LM) | [Official Docs](https://docs.nvidia.com/Megatron-Core/)
+
+
-[Previous News](docs/discussions/README.md#previous-news)
+
+## Dev Branch Philosophy
# Project Structure
@@ -77,11 +97,11 @@ Megatron-LM/
│ │ ├── optimizer/ # Optimizers
│ │ ├── datasets/ # Dataset loaders
│ │ ├── inference/ # Inference engines and server
-│ │ └── export/ # Model export (example: TensorRT-LLM)
+│ │ └── export/ # Model export (e.g. TensorRT-LLM)
│ ├── training/ # Training scripts
│ ├── legacy/ # Legacy components
│ ├── post_training/ # Post-training (quantization, distillation, pruning, etc.)
-│ └── rl/ # Reinforcement learning (including RLHF)
+│ └── rl/ # Reinforcement learning (RLHF, etc.)
├── examples/ # Ready-to-use training examples
├── tools/ # Utility tools
├── tests/ # Comprehensive test suite
@@ -90,66 +110,77 @@ Megatron-LM/
# Performance Benchmarking
-For the latest performance benchmarking results, refer to [NVIDIA Megatron Bridge Performance Summary](https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html).
+For our latest performance benchmarking results, please refer to [NVIDIA Megatron Bridge Performance Summary](https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html).
-The codebase efficiently trains models from 2B to 462B parameters across thousands of GPUs, achieving up to **47% Model FLOP Utilization (MFU)** on H100 clusters.
+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.
**Benchmark Configuration:**
- **Vocabulary size**: 131,072 tokens
-- **Sequence length**: 4,096 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)
**Key Results:**
-- **6,144 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, and logging).
-- **Production ready**: Full training pipeline with checkpointing and fault tolerance.
+- **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*
## Weak Scaling Results
-The 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.
+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.
## Strong Scaling Results
-This test strong scales the standard GPT-3 model (slightly more than 175 billion parameters due to larger vocabulary size) from 96 H100 GPUs to 4,608 GPUs, using the same batch size of 1,152 sequences throughout. Communication becomes more exposed at larger scale, leading to a reduction in MFU from 47% to 42%.
+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%.
# Roadmaps
-- **[2026 Q2 Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/4997)**
-- **[2026 Q2 MoE-Specific Roadmap](https://github.com/NVIDIA/Megatron-LM/issues/4815)** [`dev` branch first developments]
+### Fast Iteration
+- **Streamlined Review**: 1 code owner + 1 dev approver (can delegate review) + CI/CD
+
+### Feature Lifecycle (Coming Soon)
+- **6-Month Timeline**: Experimental features must graduate to stable or be deprecated
+- **Migration Support**: Assistance provided for feature transitions
+
+### Stability Expectations
+- **Experimental Nature**: Features may change or be removed as development progresses
+- **Testing**: All features will pass convergence and performance validation before inclusion
+- **Support**: Dev branch issues should include `[DEV]` prefix
# Resources
-## Getting Help
+## Performance & Benchmarking
-- 📖 **[Documentation](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html)** - Official guides and API reference
-- 🐛 **[Issues](https://github.com/NVIDIA/Megatron-LM/issues)** - Bug reports and feature requests
+- 🚀 [2025/11] [Optimizing DeepSeek-V3 Training Performance on NVIDIA GB200 NVL72](docs/discussions/deepseek-v3-gb200-optimization/deepseek-v3-gb200-optimization.md).
+- ⚡ [2025/11] [A Guide to Reproduce DeepSeek-V3 Pre-training Performance on GB200](docs/discussions/deepseek-v3-gb200-optimization/deepseek-v3-gb200-reproduce-guide.md).
-## Contributing
+## Community & Support
-Contributions are welcome. Ways to contribute:
+### Getting Help
+- 📖 **[Documentation](https://docs.nvidia.com/Megatron-Core/)** - Official documentation
+- 🐛 **[Issues](https://github.com/NVIDIA/Megatron-LM/issues)** - Bug reports and feature requests
+
+### Contributing
+We ❤️ contributions! Ways to contribute:
-- 🐛 **Report bugs** - Help improve reliability
+- 🐛 **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
-**→ [Contributing Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/developer/contribute.html)**
-
-## Citation
-
-If you use Megatron in your research or project, use the following citation:
+**→ [Contributing Guide](./CONTRIBUTING.md)**
+### Citation
```bibtex
@article{megatron-lm,
title={Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism},
diff --git a/codecov.yml b/codecov.yml
index aa37017f082..fb1e12b547c 100644
--- a/codecov.yml
+++ b/codecov.yml
@@ -12,3 +12,5 @@ coverage:
if_not_found: success
fixes:
- "/opt/megatron-lm/::"
+ignore:
+ - "megatron/core/fusions/*"
diff --git a/docker/.ngc_version.dev b/docker/.ngc_version.dev
index 3356f1f0bca..59a585ec262 100644
--- a/docker/.ngc_version.dev
+++ b/docker/.ngc_version.dev
@@ -1 +1 @@
-nvcr.io/nvidia/pytorch:26.04-py3
+nvcr.io/nvidia/pytorch:26.04-py3
\ No newline at end of file
diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev
index 5405de51142..7f3a5c0552a 100644
--- a/docker/Dockerfile.ci.dev
+++ b/docker/Dockerfile.ci.dev
@@ -1,8 +1,8 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
# syntax=docker/dockerfile:1.3-labs
-ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:26.04-py3
-FROM ${FROM_IMAGE_NAME} AS main
+ARG FROM_IMAGE_NAME
+FROM ${FROM_IMAGE_NAME} as main
ENV PIP_CONSTRAINT=""
ENV DEBIAN_FRONTEND=noninteractive
ARG UV_VERSION=0.7.2
@@ -16,7 +16,7 @@ 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 install -y --no-install-recommends gettext python3-venv psmisc uuid-runtime
apt-get clean
python -m venv /opt/jet
ARCH=$(uname -m)
@@ -31,18 +31,21 @@ RUN bash -ex <<"EOF"
curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh
EOF
+RUN ln -sf /usr/local/cuda/targets/x86_64-linux/include/cuda \
+ /usr/local/include/cuda
+RUN find /usr/local/cuda -name "utility" 2>/dev/null | head -5 && \
+ ls /usr/local/cuda/targets/x86_64-linux/include/ | head -20
+
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 IMAGE_TYPE=${IMAGE_TYPE}
-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 \
+ uv sync --extra ${IMAGE_TYPE} --extra mlm --group no_pypi_wheels --link-mode copy --locked \
--no-install-package torch \
--no-install-package torchvision \
--no-install-package triton \
@@ -62,44 +65,20 @@ RUN --mount=type=cache,target=/root/.cache/uv \
EOF
# Install DeepEP
-ARG DEEPEP_COMMIT=17cfb817bccec3a9c247013360cc550c2bac441e
-ENV DEEPEP_COMMIT=$DEEPEP_COMMIT
-ENV HYBRID_EP_MULTINODE=1
-ENV RDMA_CORE_HOME=/opt/rdma-core/build
-ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64/:$LD_LIBRARY_PATH
COPY docker/patches/deepep.patch /workspace/deepep.patch
RUN bash -ex <<"EOF"
- if [ "$IMAGE_TYPE" = "lts" ]; then
- echo "[DeepEP] skipping install for IMAGE_TYPE=lts"
- exit 0
- fi
- apt-get update
- apt-get install -y --allow-change-held-packages rdma-core libibverbs-dev
- apt-get clean
- ARCH_LIB=$(dpkg-architecture -qDEB_HOST_MULTIARCH)
- test -f /usr/lib/${ARCH_LIB}/libmlx5.so || ln -sf /usr/lib/${ARCH_LIB}/libmlx5.so.1 /usr/lib/${ARCH_LIB}/libmlx5.so
- mkdir -p ${RDMA_CORE_HOME}
- ln -sfn /usr/include ${RDMA_CORE_HOME}/include
- ln -sfn /usr/lib/${ARCH_LIB} ${RDMA_CORE_HOME}/lib
-
cd /workspace
uv pip install nvidia-nvshmem-cu13==3.4.5
pushd /opt/venv/lib/python3.12/site-packages/nvidia/nvshmem/lib/
- ln -sf libnvshmem_host.so.3 libnvshmem_host.so
+ ln -s libnvshmem_host.so.3 libnvshmem_host.so
popd
- git clone https://github.com/deepseek-ai/DeepEP.git
+ git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git
pushd DeepEP
- git fetch origin $DEEPEP_COMMIT
- git checkout FETCH_HEAD
+ git checkout eb9cee7de5a24193bf09500668d3a619d3d3f3fb
patch -p1 < /workspace/deepep.patch
- apt-get update
- apt-get install -y --no-install-recommends libnvidia-ml-dev
- TORCH_CUDA_ARCH_LIST="9.0 10.0 12.0" uv pip install --no-build-isolation -v .
- apt-get purge -y libnvidia-ml-dev
- apt-get autoremove -y
- rm -rf /var/lib/apt/lists/*
popd
+ TORCH_CUDA_ARCH_LIST="9.0 10.0 12.0" uv pip install --no-build-isolation -v DeepEP/.
rm -rf DeepEP
EOF
@@ -107,7 +86,7 @@ COPY assets/ /opt/data/
ENV UV_PYTHON=$UV_PROJECT_ENVIRONMENT/bin/python
##### For NVIDIANS only #####
-FROM main AS jet
+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"
@@ -122,7 +101,7 @@ RUN --mount=type=secret,id=JET_INDEX_URLS \
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 "setuptools<80.0.0,>=77.0.0"
uv pip install --no-cache-dir --upgrade $JET_INDEX_URLS "jet-client~=4.0"
EOF
###
diff --git a/docs/add_copyright_header.py b/docs/add_copyright_header.py
index 9694ef84819..9bc4481c506 100644
--- a/docs/add_copyright_header.py
+++ b/docs/add_copyright_header.py
@@ -1,3 +1,5 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
#!/usr/bin/env python3
"""One-off script to add NVIDIA copyright header to all .md files under docs/."""
diff --git a/docs/conf.py b/docs/conf.py
index c3055bec4ab..f354c1b8f69 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -20,7 +20,6 @@
import os
import sys
-
# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
@@ -84,24 +83,17 @@
# 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"),
- ]
+ 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.*",
- ]
+ 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,
- },
+ "switcher": {"json_url": "../versions1.json", "version_match": release},
"icon_links": [
{
"name": "GitHub",
@@ -114,23 +106,15 @@
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.*",
-]
+linkcheck_ignore = [".*github\\.com.*", ".*githubusercontent\\.com.*", "http://localhost.*"]
linkcheck_retries = 10
linkcheck_rate_limit_timeout = 600
linkcheck_workers = 1
# 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/.*",
-]
+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/.*",
-]
+linkcheck_anchors_ignore_for_url = ["https://docs.pytorch.org/.*"]
diff --git a/docs/discussions/deepseek-v3-gb200-optimization/deepseek-v3-gb200-reproduce-guide.md b/docs/discussions/deepseek-v3-gb200-optimization/deepseek-v3-gb200-reproduce-guide.md
new file mode 100644
index 00000000000..740980b9165
--- /dev/null
+++ b/docs/discussions/deepseek-v3-gb200-optimization/deepseek-v3-gb200-reproduce-guide.md
@@ -0,0 +1,358 @@
+---
+orphan: true
+---
+
+# A Guide to Reproduce DeepSeek-V3 Pre-training Performance on GB200
+
+## 1. Dockerfile
+
+Requirements:
+- Transformer Engine: We recommend using commit [d2945c6](https://github.com/NVIDIA/TransformerEngine/commit/d2945c6a571e3978677614d1fe08779966a5a4ef) with PR [2146](https://github.com/NVIDIA/TransformerEngine/pull/2146) and [2150](https://github.com/NVIDIA/TransformerEngine/pull/2150). You could prepare the branch by yourself, or use this [branch](https://github.com/hxbai/TransformerEngine/commits/dev_20251024/) based on TE v2.9 plus the above three commits/PRs.
+- cuDNN: v9.14 is required.
+- HybridEP: Install it from [here](https://github.com/deepseek-ai/DeepEP/commits/3f601f7ac1c062c46502646ff04c535013bfca00).
+
+Dockerfile for reference.
+
+```dockerfile
+FROM nvcr.io/nvidia/pytorch:25.09-py3 AS base
+
+ENV SHELL=/bin/bash
+
+# =========================
+# Install system packages
+# =========================
+RUN rm -rf /opt/megatron-lm && \
+ apt-get update && \
+ apt-get install -y sudo gdb bash-builtins git zsh autojump tmux curl gettext libfabric-dev && \
+ wget https://github.com/mikefarah/yq/releases/download/v4.27.5/yq_linux_arm64 -O /usr/bin/yq && \
+ chmod +x /usr/bin/yq
+
+# =========================
+# Install Python packages
+# =========================
+# NOTE: `unset PIP_CONSTRAINT` to install packages that do not meet the default constraint in the base image.
+# Some package requirements and related versions are from
+# https://github.com/NVIDIA/Megatron-LM/blob/core_v0.12.0/Dockerfile.linting.
+# https://github.com/NVIDIA/Megatron-LM/blob/core_v0.12.0/requirements_mlm.txt.
+# https://github.com/NVIDIA/Megatron-LM/blob/core_v0.12.0/requirements_ci.txt.
+RUN unset PIP_CONSTRAINT && pip install --no-cache-dir debugpy dm-tree torch_tb_profiler einops wandb \
+ sentencepiece tokenizers transformers torchvision ftfy modelcards datasets tqdm pydantic \
+ nvidia-pytriton py-spy yapf darker \
+ tiktoken flask-restful \
+ nltk wrapt pytest pytest_asyncio pytest-cov pytest_mock pytest-random-order \
+ black==24.4.2 isort==5.13.2 flake8==7.1.0 pylint==3.2.6 coverage mypy \
+ setuptools==69.5.1
+
+# =========================
+# Install cudnn 9.14.0.64 for correct mxfp8 quantization and layernorm fusion
+# =========================
+RUN apt-get update && \
+ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/sbsa/cuda-keyring_1.1-1_all.deb && \
+ dpkg -i cuda-keyring_1.1-1_all.deb && \
+ apt-get update && \
+ apt-get -y install libcudnn9-cuda-13
+
+# =========================
+# Install latest TE
+# Use a specific commit instead of main to make it more stable.
+# This is based on release_v2.9 branch and contains some CPU and quantization optimizations.
+# =========================
+ARG COMMIT="7dd3914726abb79bc99ff5a5db1449458ed64151"
+ARG TE="git+https://github.com/hxbai/TransformerEngine.git@${COMMIT}"
+RUN pip install nvidia-mathdx==25.1.1 && \
+ unset PIP_CONSTRAINT && \
+ NVTE_CUDA_ARCHS="100" NVTE_BUILD_THREADS_PER_JOB=8 NVTE_FRAMEWORK=pytorch pip install --no-build-isolation --no-cache-dir $TE
+
+# =========================
+# Install HybridEP
+# =========================
+WORKDIR /home/
+RUN git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git && \
+ cd DeepEP && git checkout 3f601f7ac1c062c46502646ff04c535013bfca00 && \
+ TORCH_CUDA_ARCH_LIST="10.0" pip install --no-build-isolation .
+
+# =========================
+# Clean cache
+# =========================
+RUN rm -rf /root/.cache /tmp/*
+```
+
+> [!Tip]
+>
+> If you prefer to use CUDA 12.9, please change the base container to `nvcr.io/nvidia/pytorch:25.06-py3` and the cuDNN to be installed to `libcudnn9-cuda-12`.
+
+## 2. Megatron-Core
+
+We recommend using the [dev branch](https://github.com/NVIDIA/Megatron-LM/tree/dev) after PR [1917](https://github.com/NVIDIA/Megatron-LM/pull/1917).
+
+```bash
+git clone https://github.com/NVIDIA/Megatron-LM.git && \
+cd Megatron-LM &&
+git checkout effebd81f410bc6566fffee6c320b6f8f762e06d
+```
+
+## 3. Cluster Configuration
+
+Since we're using EP 32 on NVL72, it's important to make sure
+
+> [!Important]
+> **Every 32 GB200 GPUs (8 nodes) are in the same NVL domain (or rack)**.
+
+Usually you can make it via your cluster workload manager. Taking Slurm as an example, you could pass `--segment 8` to the sbatch command to ensure that every segment of 8 nodes will be scheduled to a rack.
+
+## 4. Training scripts
+
+### Environment variables
+
+```bash
+CUDA_DEVICE_MAX_CONNECTIONS=1
+NVTE_FWD_LAYERNORM_SM_MARGIN=0
+NVTE_BWD_LAYERNORM_SM_MARGIN=0
+NVLINK_DOMAIN_SIZE=72
+NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+NCCL_NVLS_ENABLE=0
+NVTE_FUSED_ATTN=1
+NVTE_NORM_FWD_USE_CUDNN=1
+NVTE_NORM_BWD_USE_CUDNN=1
+PYTHONWARNINGS=ignore
+NCCL_DEBUG=VERSION
+NCCL_GRAPH_REGISTER=0
+```
+
+### bindpcie
+
+Download [bindpcie](https://github.com/NVIDIA/mlperf-common/blob/main/client/bindpcie) to your workdir, make it executable,
+
+```bash
+wget https://raw.githubusercontent.com/NVIDIA/mlperf-common/refs/heads/main/client/bindpcie &&
+chmod 755 bindpcie
+```
+
+and then
+
+> [!Important]
+> **Place it at the beginning of your launch command in every process.**
+
+Taking Slurm as an example, your script should look like
+
+```bash
+#!/bin/bash
+
+#SBATCH [... sbatch args]
+
+srun [... srun args] /path/to/bindpcie /path/to/pretrain_gpt.py [... mcore arguments]
+```
+
+This is a very important step on GB200.
+
+### Launch script
+
+```bash
+/path/to/bindpcie \
+/path/to/megatron-lm/pretrain_gpt.py \
+--distributed-timeout-minutes 60 \
+--tensor-model-parallel-size 1 \
+--pipeline-model-parallel-size 8 \
+--expert-model-parallel-size 32 \
+--context-parallel-size 1 \
+--expert-tensor-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 1 \
+--global-batch-size 2048 \
+--train-samples 585937500 \
+--exit-duration-in-mins 220 \
+--no-save-optim \
+--no-check-for-nan-in-loss-and-grad \
+--cross-entropy-loss-fusion \
+--cross-entropy-fusion-impl te \
+--manual-gc \
+--manual-gc-interval 10 \
+--enable-experimental \
+--transformer-impl transformer_engine \
+--seq-length 4096 \
+--data-cache-path /path/to/data_cache \
+--tokenizer-type HuggingFaceTokenizer \
+--tokenizer-model unsloth/DeepSeek-V3 \
+--data-path /path/to/data \
+--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-grouped-gemm \
+--moe-aux-loss-coeff 1e-4 \
+--moe-router-group-topk 4 \
+--moe-router-num-groups 8 \
+--moe-router-pre-softmax \
+--moe-router-padding-for-quantization \
+--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-fusion \
+--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 \
+--eval-iters 32 \
+--eval-interval 200 \
+--no-load-optim \
+--no-load-rng \
+--auto-detect-ckpt-format \
+--load None \
+--save /path/to/checkpoints \
+--save-interval 500 \
+--dist-ckpt-strictness log_all \
+--init-method-std 0.02 \
+--log-timers-to-tensorboard \
+--log-memory-to-tensorboard \
+--log-validation-ppl-to-tensorboard \
+--log-throughput \
+--log-interval 1 \
+--logging-level 40 \
+--tensorboard-dir /path/to/tensorboard \
+--wandb-project deepseek-v3-benchmarking-v0.15 \
+--wandb-exp-name DeepSeek-V3-TP1PP8EP32CP1VPP4-MBS1GBS2048-v0.15 \
+--bf16 \
+--enable-experimental \
+--recompute-granularity selective \
+--recompute-modules moe_act mlp \
+--cuda-graph-impl transformer_engine \
+--cuda-graph-modules attn moe_router moe_preprocess \
+--te-rng-tracker \
+--pipeline-model-parallel-layout "Et|(tt|)*30L" \
+--moe-router-force-load-balancing \
+--moe-token-dispatcher-type flex \
+--moe-flex-dispatcher-backend hybridep \
+--moe-hybridep-num-sms 32 \
+--fp8-recipe mxfp8 \
+--fp8-format e4m3 \
+--fp8-param-gather \
+--reuse-grad-buf-for-mxfp8-param-ag \
+--use-precision-aware-optimizer \
+--main-grads-dtype fp32 \
+--main-params-dtype fp32 \
+--exp-avg-dtype bf16 \
+--exp-avg-sq-dtype bf16 \
+```
+
+### Explanation of arguments
+
+The following arguments indicate key optimizations.
+
+- Pipeline parallel layout
+
+```bash
+--pipeline-model-parallel-layout "Et|(tt|)*30L"
+```
+
+`E` stands for embedding, `t` for transformer layer, `L` for Loss. So it's interpreted as a total of 32 stages, where the first stage is Embedding + 1 transformer layer, the last stage is Loss, and the middle 30 stages are 2 transformer layers.
+
+- Fine-grained recompute
+
+```bash
+--recompute-granularity selective \
+--recompute-modules moe_act mlp \
+```
+
+- Partial CUDA Graphs
+
+```bash
+--cuda-graph-impl transformer_engine \
+--cuda-graph-modules attn moe_router moe_preprocess \
+--te-rng-tracker \
+```
+
+- Force load balancing for performance benchmark
+
+```bash
+--moe-router-force-load-balancing \
+```
+
+- HybridEP
+
+```bash
+--moe-token-dispatcher-type flex \
+--moe-flex-dispatcher-backend hybridep \
+--moe-hybridep-num-sms 32 \
+```
+
+- MXFP8 recipe
+
+```bash
+--fp8-recipe mxfp8 \
+--fp8-format e4m3 \
+--fp8-param-gather \
+--reuse-grad-buf-for-mxfp8-param-ag \
+```
+
+- BF16 optimizer states
+
+```bash
+--use-precision-aware-optimizer \
+--main-grads-dtype fp32 \
+--main-params-dtype fp32 \
+--exp-avg-dtype bf16 \
+--exp-avg-sq-dtype bf16 \
+```
+
+- Kernel fusions
+
+```bash
+--cross-entropy-loss-fusion \
+--cross-entropy-fusion-impl te \
+--moe-permute-fusion \
+--moe-router-fusion \
+```
+
+- Manual GC to make ranks better synchronized
+
+```bash
+--manual-gc \
+--manual-gc-interval 10 \
+```
diff --git a/docs/discussions/deepseek-v3-gb200-optimization/images/image1.png b/docs/discussions/deepseek-v3-gb200-optimization/images/image1.png
new file mode 100644
index 00000000000..6e4dad685c4
Binary files /dev/null and b/docs/discussions/deepseek-v3-gb200-optimization/images/image1.png differ
diff --git a/docs/discussions/deepseek-v3-gb200-optimization/images/image2.png b/docs/discussions/deepseek-v3-gb200-optimization/images/image2.png
new file mode 100644
index 00000000000..920e3c57f94
Binary files /dev/null and b/docs/discussions/deepseek-v3-gb200-optimization/images/image2.png differ
diff --git a/docs/discussions/deepseek-v3-gb200-optimization/images/image3.png b/docs/discussions/deepseek-v3-gb200-optimization/images/image3.png
new file mode 100644
index 00000000000..f606dbfb744
Binary files /dev/null and b/docs/discussions/deepseek-v3-gb200-optimization/images/image3.png differ
diff --git a/docs/discussions/deepseek-v3-gb200-optimization/images/image4.png b/docs/discussions/deepseek-v3-gb200-optimization/images/image4.png
new file mode 100644
index 00000000000..04239401edd
Binary files /dev/null and b/docs/discussions/deepseek-v3-gb200-optimization/images/image4.png differ
diff --git a/docs/discussions/deepseek-v3-gb200-optimization/images/image5.png b/docs/discussions/deepseek-v3-gb200-optimization/images/image5.png
new file mode 100644
index 00000000000..0128fc7ae45
Binary files /dev/null and b/docs/discussions/deepseek-v3-gb200-optimization/images/image5.png differ
diff --git a/docs/discussions/deepseek-v3-gb200-optimization/images/image6.png b/docs/discussions/deepseek-v3-gb200-optimization/images/image6.png
new file mode 100644
index 00000000000..cb2ed2eb9ad
Binary files /dev/null and b/docs/discussions/deepseek-v3-gb200-optimization/images/image6.png differ
diff --git a/docs/discussions/deepseek-v3-gb200-optimization/images/image7.png b/docs/discussions/deepseek-v3-gb200-optimization/images/image7.png
new file mode 100644
index 00000000000..325d0fd4f52
Binary files /dev/null and b/docs/discussions/deepseek-v3-gb200-optimization/images/image7.png differ
diff --git a/docs/source/api-guide/router_replay.md b/docs/source/api-guide/router_replay.md
new file mode 100644
index 00000000000..b2e043b3065
--- /dev/null
+++ b/docs/source/api-guide/router_replay.md
@@ -0,0 +1,180 @@
+---
+orphan: true
+---
+
+# 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 `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.
+ * `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 `enable_routing_replay` to `True` in the model configuration.
+ 2. **Initialization**: When `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 `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,
+ 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/user-guide/features/cuda_graph.md b/docs/user-guide/features/cuda_graph.md
index 28a1a5575dc..3d5077e92ce 100644
--- a/docs/user-guide/features/cuda_graph.md
+++ b/docs/user-guide/features/cuda_graph.md
@@ -167,10 +167,9 @@ eager mode. Example using `transformer_engine` (`local` works the same way):
--cuda-graph-modules attn moe_router moe_preprocess
```
-With paged stash (currently available only on `dev`; see
-`docs/user-guide/features/paged_stash.md` on the `dev` branch), expert dispatch shapes become
-static (pre-sized via `--moe-expert-rank-capacity-factor`), which allows full-iteration CUDA
-graphs to be used on MoE models as well:
+With [paged stash](paged_stash.md), expert dispatch shapes become static (pre-sized via
+`--moe-expert-rank-capacity-factor`), which allows full-iteration CUDA graphs to be used on MoE
+models as well:
```bash
--cuda-graph-impl full_iteration \
diff --git a/docs/user-guide/features/fine_grained_activation_offloading.md b/docs/user-guide/features/fine_grained_activation_offloading.md
index 915926a6b9b..f83645d7ec4 100644
--- a/docs/user-guide/features/fine_grained_activation_offloading.md
+++ b/docs/user-guide/features/fine_grained_activation_offloading.md
@@ -42,13 +42,13 @@ Supported offloading modules are `"attn_norm"`, `"qkv_linear"`, `"core_attn"`, `
```bash
# Optional: cap inflight D2H offloads per offload group to N (omit or None in most setups).
# Required as a non-None non-negative integer when fine-grained activation offloading is used with
-# local full-iteration CUDA graphs (full_iteration in cuda_graph_scope); see prose below.
+# full-iteration CUDA graphs (--cuda-graph-impl full_iteration); see prose below.
--fine-grained-offloading-max-inflight-offloads
```
TransformerConfig.fine_grained_offloading_max_inflight_offloads caps, per offload group (for example `moe_act`, `qkv_linear`), how many D2H copies may be in flight before a main-stream wait_event. 0 waits after each offload; larger values allow more overlap; None skips these joins.
-With full-iteration CUDA graphs (local graph impl, full_iteration in cuda_graph_scope) and fine-grained activation offloading enabled, set it to a non-None integer: that path does not rely on record_stream, so explicit joins are required.
+With full-iteration CUDA graphs (`--cuda-graph-impl full_iteration`) and fine-grained activation offloading enabled, set it to a non-None integer: that path does not rely on record_stream, so explicit joins are required.
## Compatible With Fine-Grained Recomputation
diff --git a/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py
index 2a2b1d63a21..a5afb7e5c60 100644
--- a/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py
+++ b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py
@@ -6,24 +6,27 @@
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)))
+
+sys.path.append(
+ os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
+)
+from typing import Union
+
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.models.gpt.gpt_layer_specs import (
+ get_gpt_layer_local_spec,
+ get_gpt_layer_with_transformer_engine_spec,
+)
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
+from megatron.inference.text_generation import generate_and_post_process
+from megatron.training import get_args, get_model, get_tokenizer, print_rank_0
+from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
+from megatron.training.checkpointing import load_checkpoint
+from megatron.training.initialize import initialize_megatron
+
def model_provider(pre_process=True, post_process=True) -> GPTModel:
"""Builds the model.
@@ -44,20 +47,17 @@ def model_provider(pre_process=True, post_process=True) -> GPTModel:
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
+ 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
+ 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
+ num_experts=args.num_experts, moe_grouped_gemm=args.moe_grouped_gemm
)
else:
transformer_layer_spec = import_module(args.spec)
@@ -73,37 +73,46 @@ def model_provider(pre_process=True, post_process=True) -> GPTModel:
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
+ 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')
+ 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()
@@ -111,6 +120,7 @@ def generate_samples_unconditional(model):
cnt = 0
num_samples = args.num_samples
from tqdm import tqdm
+
pbar = tqdm(total=num_samples)
while True:
@@ -118,16 +128,23 @@ def generate_samples_unconditional(model):
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)
+ 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}
+ datum = {
+ 'text': generation[len(prompt) :],
+ 'all_text': generation,
+ 'prompt': prompt,
+ 'id': cnt,
+ }
yield datum
cnt += 1
pbar.update()
@@ -148,6 +165,7 @@ def generate_samples_conditional(model):
num_samples = args.num_samples
cnt = 0
from tqdm import tqdm
+
pbar = tqdm(total=num_samples)
fname = open(args.sample_input_file, "r")
@@ -171,16 +189,23 @@ def generate_samples_conditional(model):
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)
+ 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}
+ datum = {
+ 'text': generation[len(prompt) :],
+ 'all_text': generation,
+ 'prompt': prompt,
+ 'id': cnt,
+ }
yield datum
cnt += 1
pbar.update()
@@ -207,8 +232,7 @@ 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))
+ 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:
@@ -220,11 +244,15 @@ def generate_and_write_samples_conditional(model):
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})
+ 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
diff --git a/examples/bert/pretrain_bert.py b/examples/bert/pretrain_bert.py
index 3eb95ecf396..4dd6160f795 100644
--- a/examples/bert/pretrain_bert.py
+++ b/examples/bert/pretrain_bert.py
@@ -7,26 +7,30 @@
import torch
import torch.nn.functional as F
-from megatron.training import get_args
-from megatron.training import print_rank_0
-from megatron.training import get_timers
-from megatron.core import tensor_parallel
+from megatron.core import mpu, tensor_parallel
+from megatron.core.datasets.bert_dataset import (
+ BERTMaskedWordPieceDataset,
+ BERTMaskedWordPieceDatasetConfig,
+)
+from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
+from megatron.core.datasets.utils import get_blend_from_list
from megatron.core.enums import ModelType
+from megatron.core.models.bert.bert_layer_specs import (
+ bert_layer_local_spec,
+ bert_layer_with_transformer_engine_spec,
+)
from megatron.core.models.bert.bert_model import BertModel
-from megatron.training import pretrain
-from megatron.training.utils import average_losses_across_data_parallel_group
-from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
-from megatron.training.argument_utils import pretrain_cfg_container_from_args
-from megatron.core.transformer.spec_utils import import_module
-from megatron.core.models.bert.bert_layer_specs import bert_layer_with_transformer_engine_spec, bert_layer_local_spec
from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
-from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
-from megatron.core.datasets.bert_dataset import BERTMaskedWordPieceDataset, BERTMaskedWordPieceDatasetConfig
-from megatron.core.datasets.utils import get_blend_from_list
-from megatron.core import mpu, tensor_parallel
+from megatron.core.transformer.spec_utils import import_module
+from megatron.training import get_args, get_timers, pretrain, print_rank_0
+from megatron.training.argument_utils import pretrain_cfg_container_from_args
+from megatron.training.arguments import core_transformer_config_from_args, parse_and_validate_args
+from megatron.training.utils import average_losses_across_data_parallel_group
-def model_provider(pre_process=True, post_process=True, vp_stage=None, config=None, pg_collection=None):
+def model_provider(
+ pre_process=True, post_process=True, vp_stage=None, config=None, pg_collection=None
+):
"""Build the model."""
print_rank_0('building BERT model ...')
@@ -37,11 +41,11 @@ def model_provider(pre_process=True, post_process=True, vp_stage=None, config=No
num_tokentypes = 2 if args.bert_binary_head else 0
if args.spec is None:
- transformer_layer_spec = bert_layer_with_transformer_engine_spec #default spec
+ transformer_layer_spec = bert_layer_with_transformer_engine_spec # default spec
elif args.spec[0] == 'local':
print_rank_0('Using Local spec for transformer layers')
transformer_layer_spec = bert_layer_local_spec
- else :
+ else:
transformer_layer_spec = import_module(args.spec)
model = BertModel(
@@ -55,7 +59,8 @@ def model_provider(pre_process=True, post_process=True, vp_stage=None, config=No
parallel_output=True,
pre_process=pre_process,
post_process=post_process,
- vp_stage=vp_stage)
+ vp_stage=vp_stage,
+ )
return model
@@ -64,8 +69,7 @@ def get_batch(data_iterator):
"""Build the batch."""
# Items and their type.
- keys = ['text', 'types', 'labels',
- 'is_random', 'loss_mask', 'padding_mask']
+ keys = ['text', 'types', 'labels', 'is_random', 'loss_mask', 'padding_mask']
datatype = torch.int64
# Broadcast data.
@@ -91,23 +95,19 @@ def loss_func(loss_mask, sentence_order, output_tensor):
lm_loss_ = lm_loss_.float()
loss_mask = loss_mask.float()
- lm_loss = torch.sum(
- lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum()
+ lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum()
if sop_logits is not None:
- sop_loss = F.cross_entropy(sop_logits.view(-1, 2).float(),
- sentence_order.view(-1),
- ignore_index=-1)
+ sop_loss = F.cross_entropy(
+ sop_logits.view(-1, 2).float(), sentence_order.view(-1), ignore_index=-1
+ )
sop_loss = sop_loss.float()
loss = lm_loss + sop_loss
- averaged_losses = average_losses_across_data_parallel_group(
- [lm_loss, sop_loss])
- return loss, {'lm loss': averaged_losses[0],
- 'sop loss': averaged_losses[1]}
+ averaged_losses = average_losses_across_data_parallel_group([lm_loss, sop_loss])
+ return loss, {'lm loss': averaged_losses[0], 'sop loss': averaged_losses[1]}
else:
loss = lm_loss
- averaged_losses = average_losses_across_data_parallel_group(
- [lm_loss])
+ averaged_losses = average_losses_across_data_parallel_group([lm_loss])
return loss, {'lm loss': averaged_losses[0]}
@@ -118,16 +118,14 @@ def forward_step(data_iterator, model):
# Get the batch.
timers('batch-generator', log_level=2).start()
- tokens, types, sentence_order, loss_mask, lm_labels, padding_mask = get_batch(
- data_iterator)
+ tokens, types, sentence_order, loss_mask, lm_labels, padding_mask = get_batch(data_iterator)
timers('batch-generator').stop()
if not args.bert_binary_head:
types = None
# Forward pass through the model.
- output_tensor = model(tokens, padding_mask,
- tokentype_ids=types, lm_labels=lm_labels)
+ output_tensor = model(tokens, padding_mask, tokentype_ids=types, lm_labels=lm_labels)
return output_tensor, partial(loss_func, loss_mask, sentence_order)
@@ -145,7 +143,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None
blend_per_split=[
get_blend_from_list(args.train_data_path),
get_blend_from_list(args.valid_data_path),
- get_blend_from_list(args.test_data_path)
+ get_blend_from_list(args.test_data_path),
],
split=args.split,
path_to_cache=args.data_cache_path,
@@ -162,8 +160,7 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None
allow_ambiguous_pad_tokens=args.allow_ambiguous_pad_tokens,
)
- print_rank_0('> building train, validation, and test datasets '
- 'for BERT ...')
+ print_rank_0('> building train, validation, and test datasets ' 'for BERT ...')
train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
BERTMaskedWordPieceDataset,
@@ -184,6 +181,10 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None
args = parse_and_validate_args(args_defaults={'tokenizer_type': 'BertWordPieceLowerCase'})
full_config = pretrain_cfg_container_from_args(args)
- pretrain(full_config, train_valid_test_datasets_provider, model_provider,
- ModelType.encoder_or_decoder,
- forward_step)
+ pretrain(
+ full_config,
+ train_valid_test_datasets_provider,
+ model_provider,
+ ModelType.encoder_or_decoder,
+ forward_step,
+ )
diff --git a/examples/dynamic_context_parallel/README.md b/examples/dynamic_context_parallel/README.md
new file mode 100644
index 00000000000..70864e42293
--- /dev/null
+++ b/examples/dynamic_context_parallel/README.md
@@ -0,0 +1,161 @@
+# Dynamic Context Parallel Benchmark
+
+This example compares regular DP-balanced packed-sequence training against
+Dynamic Context Parallelism (DCP) on the same variable-length mock workload.
+
+The script reuses the normal Megatron-LM training stack:
+
+- `pretrain_gpt.py` builds and trains the GPT model.
+- `MockVarlenDataset` creates THD-format variable-length samples.
+- `DefaultDynamicCPScheduler` is enabled only for the DCP run.
+
+No model class or custom dataset class is introduced by this example.
+
+## Run
+
+From the Megatron-LM repository root:
+
+```bash
+GPUS_PER_NODE=8 bash examples/dynamic_context_parallel/benchmark_dcp.sh
+```
+
+The default topology is `TP=1, CP=4, PP=1`, so the default run expects eight
+GPUs and uses data parallel size two. By default the script sets
+`NUM_MICROBATCHES=8` and
+`GLOBAL_BATCH_SIZE=MICRO_BATCH_SIZE * DP_SIZE * NUM_MICROBATCHES`. Keeping more
+than one microbatch per global batch gives the DCP scheduler enough
+variable-length samples to assign smaller local CP groups instead of expanding
+every sample back to the full fixed CP size. This mock THD workload also keeps
+`MICRO_BATCH_SIZE=1`; increase effective batch size with `NUM_MICROBATCHES` and
+data-parallel ranks.
+
+The script runs two jobs:
+
+1. Baseline packed sequence training:
+
+ ```text
+ --sequence-packing-scheduler dp_balanced
+ --context-parallel-size 4
+ ```
+
+2. Dynamic CP training:
+
+ ```text
+ --dynamic-context-parallel
+ --sequence-packing-scheduler default_dynamic_cp
+ --context-parallel-size 4
+ ```
+
+Both runs use the same model, batch size, sequence distribution, and
+`--max-seqlen-per-dp-cp-rank`. The script also sets
+`--moe-token-dispatcher-type alltoall`, which Megatron-Core currently requires
+when sequence packing is enabled. It sets `--num-workers 0` by default to keep
+mock collation in the main process for this variable-length THD workload.
+
+## What Makes DCP Useful Here
+
+The mock dataset draws sequence lengths from a lognormal distribution:
+
+```json
+{"mode":"distribution","type":"lognormal","format":"thd","min_seq_len":128,"max_seq_len":8192,"mean_seq_len":1024,"lognormal_sigma":1.5}
+```
+
+With the default `--max-seqlen-per-dp-cp-rank 2048`, DCP can assign different
+local CP sizes:
+
+- Short samples up to 2048 tokens can use one rank.
+- Medium samples up to 4096 tokens can use two ranks.
+- Long samples up to 8192 tokens can use four ranks.
+
+The baseline keeps the full fixed CP size for the packed workload. DCP can
+spread short samples over the DPxCP domain instead of making every sample occupy
+the full CP group.
+
+## Output
+
+At the end, the script prints:
+
+```text
+=== Dynamic CP benchmark summary ===
+Iteration-time statistics exclude the first 10 logged iterations.
+Baseline dp_balanced average: ...
+Dynamic CP average: ...
+Average speedup: ...
+Baseline dp_balanced median: ...
+Dynamic CP median: ...
+Median speedup: ...
+Baseline dp_balanced 10% trimmed avg: ...
+Dynamic CP 10% trimmed avg: ...
+10% trimmed mean speedup: ...
+```
+
+It parses Megatron-LM's regular training log line:
+
+```text
+elapsed time per iteration (ms): ...
+```
+
+Logs and TensorBoard output are written under `dcp_benchmark_output/` by
+default.
+
+## Slurm Benchmark Results
+
+The following single-run measurements were collected on June 9, 2026 on a
+Slurm cluster using four GPUs per node and
+`/home/tolong/nvidian+nemo+26.02.rc5.sqsh`.
+The run used the default benchmark shape (`TP=1`, `CP=4`, `PP=1`,
+`TRAIN_ITERS=30`, `WARMUP_ITERS=10`, `NUM_MICROBATCHES=8`,
+`MAX_SEQLEN_PER_DP_CP_RANK=2048`) and the default lognormal mock VarlenDataset
+distribution shown above. Statistics exclude the first 10 logged iterations.
+
+The 10% trimmed mean is the primary comparison because short DCP runs can have
+large first-use spikes when dynamic groups are exercised. Arithmetic means are
+included to show that variance.
+
+| Nodes | GPUs | Slurm job | Baseline trimmed mean (ms) | DCP trimmed mean (ms) | Trimmed speedup | Baseline avg (ms) | DCP avg (ms) |
+| ----- | ---- | --------- | -------------------------- | --------------------- | --------------- | ----------------- | ------------ |
+| 1 | 4 | 3239751 | 195.269 | 151.906 | 1.285x | 216.125 | 155.150 |
+| 2 | 8 | 3239752 | 220.119 | 155.850 | 1.412x | 226.515 | 157.905 |
+| 4 | 16 | 3239753 | 226.775 | 187.125 | 1.212x | 231.620 | 189.470 |
+| 8 | 32 | 3239754 | 286.206 | 208.088 | 1.375x | 287.780 | 225.875 |
+| 16 | 64 | 3239955 | 271.988 | 181.912 | 1.495x | 281.480 | 217.620 |
+
+## Useful Overrides
+
+Use a larger model:
+
+```bash
+NUM_LAYERS=32 HIDDEN_SIZE=4096 FFN_HIDDEN_SIZE=16384 NUM_ATTENTION_HEADS=32 \
+GPUS_PER_NODE=8 bash examples/dynamic_context_parallel/benchmark_dcp.sh
+```
+
+Use a more skewed long/short distribution:
+
+```bash
+VARLEN_DATASET_JSON='{"mode":"distribution","type":"lognormal","format":"thd","min_seq_len":128,"max_seq_len":8192,"mean_seq_len":768,"lognormal_sigma":1.8}' \
+GPUS_PER_NODE=8 bash examples/dynamic_context_parallel/benchmark_dcp.sh
+```
+
+Use exact sequence lengths from a CSV:
+
+```bash
+VARLEN_DATASET_JSON='{"mode":"file","format":"thd","path":"/path/to/lengths.csv"}' \
+GPUS_PER_NODE=8 bash examples/dynamic_context_parallel/benchmark_dcp.sh
+```
+
+Reduce the runtime:
+
+```bash
+TRAIN_ITERS=12 WARMUP_ITERS=3 GPUS_PER_NODE=8 \
+bash examples/dynamic_context_parallel/benchmark_dcp.sh
+```
+
+Use a specific Python interpreter:
+
+```bash
+PYTHON=/opt/venv/bin/python GPUS_PER_NODE=8 \
+bash examples/dynamic_context_parallel/benchmark_dcp.sh
+```
+
+For stable numbers, keep the same GPU allocation, run more iterations, and
+avoid checkpointing or evaluation during the measured window.
diff --git a/examples/dynamic_context_parallel/benchmark_dcp.sh b/examples/dynamic_context_parallel/benchmark_dcp.sh
new file mode 100755
index 00000000000..1643c4799f9
--- /dev/null
+++ b/examples/dynamic_context_parallel/benchmark_dcp.sh
@@ -0,0 +1,268 @@
+#!/bin/bash
+
+set -euo pipefail
+
+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}
+export NVTE_ALLOW_NONDETERMINISTIC_ALGO=${NVTE_ALLOW_NONDETERMINISTIC_ALGO:-1}
+
+GPUS_PER_NODE=${GPUS_PER_NODE:-8}
+NUM_NODES=${NUM_NODES:-1}
+NODE_RANK=${NODE_RANK:-0}
+MASTER_ADDR=${MASTER_ADDR:-localhost}
+MASTER_PORT=${MASTER_PORT:-6000}
+PYTHON=${PYTHON:-python}
+
+TP_SIZE=${TP_SIZE:-1}
+CP_SIZE=${CP_SIZE:-4}
+PP_SIZE=${PP_SIZE:-1}
+MOE_TOKEN_DISPATCHER_TYPE=${MOE_TOKEN_DISPATCHER_TYPE:-alltoall}
+WORLD_SIZE=$((GPUS_PER_NODE * NUM_NODES))
+MODEL_PARALLEL_SIZE=$((TP_SIZE * CP_SIZE * PP_SIZE))
+
+NUM_LAYERS=${NUM_LAYERS:-12}
+HIDDEN_SIZE=${HIDDEN_SIZE:-2048}
+FFN_HIDDEN_SIZE=${FFN_HIDDEN_SIZE:-8192}
+NUM_ATTENTION_HEADS=${NUM_ATTENTION_HEADS:-16}
+SEQ_LENGTH=${SEQ_LENGTH:-8192}
+MAX_POSITION_EMBEDDINGS=${MAX_POSITION_EMBEDDINGS:-$SEQ_LENGTH}
+
+MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-1}
+GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-}
+NUM_MICROBATCHES=${NUM_MICROBATCHES:-8}
+TRAIN_ITERS=${TRAIN_ITERS:-30}
+WARMUP_ITERS=${WARMUP_ITERS:-10}
+LOG_INTERVAL=${LOG_INTERVAL:-1}
+NUM_WORKERS=${NUM_WORKERS:-0}
+
+MAX_SEQLEN_PER_DP_CP_RANK=${MAX_SEQLEN_PER_DP_CP_RANK:-2048}
+MIN_DYNAMIC_CONTEXT_PARALLEL_SIZE=${MIN_DYNAMIC_CONTEXT_PARALLEL_SIZE:-1}
+
+VOCAB_SIZE=${VOCAB_SIZE:-131072}
+NULL_TOKENIZER_PAD_ID=${NULL_TOKENIZER_PAD_ID:-0}
+
+OUTPUT_DIR=${OUTPUT_DIR:-"${PWD}/dcp_benchmark_output"}
+DATA_CACHE_PATH=${DATA_CACHE_PATH:-"${OUTPUT_DIR}/data_cache"}
+DEFAULT_VARLEN_DATASET_JSON='{"mode":"distribution","type":"lognormal","format":"thd","min_seq_len":128,"max_seq_len":8192,"mean_seq_len":1024,"lognormal_sigma":1.5}'
+VARLEN_DATASET_JSON=${VARLEN_DATASET_JSON:-$DEFAULT_VARLEN_DATASET_JSON}
+
+PRETRAIN_SCRIPT_PATH="pretrain_gpt.py"
+
+if [[ ! -f "$PRETRAIN_SCRIPT_PATH" ]]; then
+ echo "Error: $PRETRAIN_SCRIPT_PATH not found. Run this script from the Megatron-LM repo root."
+ exit 1
+fi
+
+if (( WORLD_SIZE < MODEL_PARALLEL_SIZE )); then
+ echo "Error: need at least TP_SIZE * CP_SIZE * PP_SIZE GPUs."
+ echo "Got GPUS_PER_NODE=${GPUS_PER_NODE}, NUM_NODES=${NUM_NODES}, TP_SIZE=${TP_SIZE}, CP_SIZE=${CP_SIZE}, PP_SIZE=${PP_SIZE}."
+ exit 1
+fi
+
+if (( WORLD_SIZE % MODEL_PARALLEL_SIZE != 0 )); then
+ echo "Error: total GPUs must be divisible by TP_SIZE * CP_SIZE * PP_SIZE."
+ echo "Got WORLD_SIZE=${WORLD_SIZE}, TP_SIZE=${TP_SIZE}, CP_SIZE=${CP_SIZE}, PP_SIZE=${PP_SIZE}."
+ exit 1
+fi
+
+DP_SIZE=$((WORLD_SIZE / MODEL_PARALLEL_SIZE))
+MICRO_BATCH_TIMES_DP=$((MICRO_BATCH_SIZE * DP_SIZE))
+GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-$((MICRO_BATCH_TIMES_DP * NUM_MICROBATCHES))}
+
+if (( MICRO_BATCH_SIZE != 1 )); then
+ echo "Error: this variable-length THD mock benchmark expects MICRO_BATCH_SIZE=1."
+ echo "Increase effective batch size with NUM_MICROBATCHES and data-parallel ranks."
+ exit 1
+fi
+
+if (( GLOBAL_BATCH_SIZE % MICRO_BATCH_TIMES_DP != 0 )); then
+ echo "Error: GLOBAL_BATCH_SIZE must be divisible by MICRO_BATCH_SIZE * DP_SIZE."
+ echo "Got GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE}, MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE}, DP_SIZE=${DP_SIZE}."
+ exit 1
+fi
+EFFECTIVE_NUM_MICROBATCHES=$((GLOBAL_BATCH_SIZE / MICRO_BATCH_TIMES_DP))
+
+if (( TRAIN_ITERS <= WARMUP_ITERS )); then
+ echo "Error: TRAIN_ITERS must be greater than WARMUP_ITERS."
+ exit 1
+fi
+
+mkdir -p "$OUTPUT_DIR" "$DATA_CACHE_PATH"
+
+echo "WORLD_SIZE=${WORLD_SIZE} DP_SIZE=${DP_SIZE} CP_SIZE=${CP_SIZE} TP_SIZE=${TP_SIZE} PP_SIZE=${PP_SIZE}"
+echo "MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE} GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE} NUM_MICROBATCHES=${EFFECTIVE_NUM_MICROBATCHES}"
+
+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 "$HIDDEN_SIZE"
+ --ffn-hidden-size "$FFN_HIDDEN_SIZE"
+ --num-attention-heads "$NUM_ATTENTION_HEADS"
+ --seq-length "$SEQ_LENGTH"
+ --max-position-embeddings "$MAX_POSITION_EMBEDDINGS"
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --transformer-impl transformer_engine
+ --attention-backend flash
+ --moe-token-dispatcher-type "$MOE_TOKEN_DISPATCHER_TYPE"
+)
+
+TRAINING_ARGS=(
+ --micro-batch-size "$MICRO_BATCH_SIZE"
+ --global-batch-size "$GLOBAL_BATCH_SIZE"
+ --train-iters "$TRAIN_ITERS"
+ --lr-decay-iters "$TRAIN_ITERS"
+ --lr 1.5e-4
+ --min-lr 1.0e-5
+ --lr-decay-style cosine
+ --lr-warmup-iters 0
+ --weight-decay 1.0e-2
+ --clip-grad 1.0
+ --bf16
+ --calculate-per-token-loss
+ --no-gradient-accumulation-fusion
+)
+
+PARALLEL_ARGS=(
+ --tensor-model-parallel-size "$TP_SIZE"
+ --pipeline-model-parallel-size "$PP_SIZE"
+ --context-parallel-size "$CP_SIZE"
+)
+
+if (( TP_SIZE > 1 )); then
+ PARALLEL_ARGS+=(--sequence-parallel)
+fi
+
+DATA_ARGS=(
+ --use-varlen-dataset
+ --mock-data
+ --varlen-mock-dataset-config-json "$VARLEN_DATASET_JSON"
+ --tokenizer-type NullTokenizer
+ --vocab-size "$VOCAB_SIZE"
+ --null-tokenizer-pad-id "$NULL_TOKENIZER_PAD_ID"
+ --split 99,1,0
+ --data-cache-path "$DATA_CACHE_PATH"
+ --dataloader-type single
+ --num-workers "$NUM_WORKERS"
+)
+
+LOGGING_ARGS=(
+ --log-interval "$LOG_INTERVAL"
+ --log-throughput
+ --timing-log-level 0
+ --eval-interval 1000000
+ --eval-iters 1
+ --save-interval 1000000
+ --distributed-backend nccl
+ --distributed-timeout-minutes 60
+)
+
+extract_iteration_stats_ms() {
+ local log_file=$1
+ "$PYTHON" - "$log_file" "$WARMUP_ITERS" <<'PY'
+import re
+import statistics
+import sys
+
+log_file = sys.argv[1]
+warmup = int(sys.argv[2])
+pattern = re.compile(r"elapsed time per iteration \(ms\): ([0-9.]+)")
+
+values = []
+with open(log_file, errors="ignore") as f:
+ for line in f:
+ match = pattern.search(line)
+ if match:
+ values.append(float(match.group(1)))
+
+values = values[warmup:]
+if not values:
+ print("nan nan nan")
+ raise SystemExit
+
+average = sum(values) / len(values)
+median = statistics.median(values)
+trim_count = int(len(values) * 0.1)
+trimmed = sorted(values)
+if len(trimmed) - 2 * trim_count > 0:
+ trimmed = trimmed[trim_count : len(trimmed) - trim_count]
+trimmed_average = sum(trimmed) / len(trimmed)
+
+print(f"{average:.3f} {median:.3f} {trimmed_average:.3f}")
+PY
+}
+
+calc_speedup() {
+ local base=$1
+ local dcp=$2
+ awk -v base="$base" -v dcp="$dcp" 'BEGIN {
+ if (base > 0 && dcp > 0) {
+ printf "%.2f", base / dcp
+ } else {
+ printf "nan"
+ }
+ }'
+}
+
+run_case() {
+ local name=$1
+ shift
+
+ local tensorboard_dir="${OUTPUT_DIR}/tensorboard_${name}"
+ local log_file="${OUTPUT_DIR}/${name}.log"
+ rm -rf "$tensorboard_dir"
+ mkdir -p "$tensorboard_dir"
+
+ echo
+ echo "=== Running ${name} ==="
+ echo "Log: ${log_file}"
+
+ "$PYTHON" -m torch.distributed.run "${DISTRIBUTED_ARGS[@]}" \
+ "$PRETRAIN_SCRIPT_PATH" \
+ "${MODEL_ARGS[@]}" \
+ "${TRAINING_ARGS[@]}" \
+ "${PARALLEL_ARGS[@]}" \
+ "${DATA_ARGS[@]}" \
+ "${LOGGING_ARGS[@]}" \
+ --tensorboard-dir "$tensorboard_dir" \
+ "$@" 2>&1 | tee "$log_file"
+}
+
+run_case baseline \
+ --sequence-packing-scheduler dp_balanced \
+ --max-seqlen-per-dp-cp-rank "$MAX_SEQLEN_PER_DP_CP_RANK"
+
+run_case dcp \
+ --dynamic-context-parallel \
+ --sequence-packing-scheduler default_dynamic_cp \
+ --min-dynamic-context-parallel-size "$MIN_DYNAMIC_CONTEXT_PARALLEL_SIZE" \
+ --max-seqlen-per-dp-cp-rank "$MAX_SEQLEN_PER_DP_CP_RANK"
+
+read -r baseline_avg_ms baseline_median_ms baseline_trimmed_ms < <(extract_iteration_stats_ms "${OUTPUT_DIR}/baseline.log")
+read -r dcp_avg_ms dcp_median_ms dcp_trimmed_ms < <(extract_iteration_stats_ms "${OUTPUT_DIR}/dcp.log")
+
+avg_speedup=$(calc_speedup "$baseline_avg_ms" "$dcp_avg_ms")
+median_speedup=$(calc_speedup "$baseline_median_ms" "$dcp_median_ms")
+trimmed_speedup=$(calc_speedup "$baseline_trimmed_ms" "$dcp_trimmed_ms")
+
+echo
+echo "=== Dynamic CP benchmark summary ==="
+echo "Iteration-time statistics exclude the first ${WARMUP_ITERS} logged iterations."
+echo "Baseline dp_balanced average: ${baseline_avg_ms} ms"
+echo "Dynamic CP average: ${dcp_avg_ms} ms"
+echo "Average speedup: ${avg_speedup}x"
+echo "Baseline dp_balanced median: ${baseline_median_ms} ms"
+echo "Dynamic CP median: ${dcp_median_ms} ms"
+echo "Median speedup: ${median_speedup}x"
+echo "Baseline dp_balanced 10% trimmed avg: ${baseline_trimmed_ms} ms"
+echo "Dynamic CP 10% trimmed avg: ${dcp_trimmed_ms} ms"
+echo "10% trimmed mean speedup: ${trimmed_speedup}x"
+echo
+echo "Logs and TensorBoard data are under ${OUTPUT_DIR}"
diff --git a/examples/inference/advanced/gpt_dynamic_inference.py b/examples/inference/advanced/gpt_dynamic_inference.py
index 53f201207ab..9bb5b0cf08f 100644
--- a/examples/inference/advanced/gpt_dynamic_inference.py
+++ b/examples/inference/advanced/gpt_dynamic_inference.py
@@ -11,10 +11,11 @@
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
+from megatron.training.arguments import parse_and_validate_args
+
sys.path.append(
os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
)
@@ -241,7 +242,10 @@ def _process_step_result(result):
# Suspend.
if attempted_step_count % args.suspend_resume_interval == 0:
- print("**** step %d/%d ... suspend." % (engine.context.step_count, attempted_step_count))
+ print(
+ "**** step %d/%d ... suspend."
+ % (engine.context.step_count, attempted_step_count)
+ )
engine.suspend()
# Resume, 0+ attempted steps later.
@@ -251,7 +255,10 @@ def _process_step_result(result):
% args.suspend_resume_interval
== 0
):
- print("**** step %d/%d ... resume." % (engine.context.step_count, attempted_step_count))
+ print(
+ "**** step %d/%d ... resume."
+ % (engine.context.step_count, attempted_step_count)
+ )
engine.resume()
# If engine suspended, continue to next iter.
@@ -469,7 +476,9 @@ def escape_str(s):
# Attach peak memory metrics; the functional test only validates these
# if the fields exist in the golden values.
json_results.update(peak_mem_stats)
- json_results["lifetime_prefill_token_count"] = engine.context.lifetime_prefill_token_count
+ 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:
diff --git a/examples/inference/advanced/gpt_dynamic_inference_with_coordinator.py b/examples/inference/advanced/gpt_dynamic_inference_with_coordinator.py
index f5191e980a9..34380e86c6f 100644
--- a/examples/inference/advanced/gpt_dynamic_inference_with_coordinator.py
+++ b/examples/inference/advanced/gpt_dynamic_inference_with_coordinator.py
@@ -9,7 +9,6 @@
from collections import defaultdict
from typing import List
-from megatron.training.arguments import parse_and_validate_args
import torch
import torch.distributed as dist
@@ -19,13 +18,14 @@
from megatron.core.inference.inference_client import InferenceClient
from megatron.core.inference.inference_request import DynamicInferenceRequestRecord
from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.utils import configure_nvtx_profiling
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
+from megatron.training.arguments import parse_and_validate_args
# pylint: disable=line-too-long
@@ -74,11 +74,15 @@ async def main(
)
# 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
+ 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 = 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:
@@ -104,7 +108,10 @@ async def main(
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:
+ 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
@@ -121,7 +128,10 @@ async def main(
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:
+ 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
@@ -160,7 +170,7 @@ async def main(
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:
diff --git a/examples/inference/advanced/gpt_static_inference.py b/examples/inference/advanced/gpt_static_inference.py
index 89cc0d5d8b8..b616eefc723 100644
--- a/examples/inference/advanced/gpt_static_inference.py
+++ b/examples/inference/advanced/gpt_static_inference.py
@@ -5,7 +5,6 @@
import time
from argparse import Namespace
-from megatron.training.arguments import parse_and_validate_args
import torch
from megatron.core.inference.contexts import StaticInferenceContext
@@ -20,6 +19,7 @@
)
from megatron.core.tokenizers.utils.build_tokenizer import build_tokenizer
from megatron.core.transformer.module import MegatronModule
+from megatron.training.arguments import parse_and_validate_args
sys.path.append(
os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
diff --git a/examples/inference/advanced/simple_t5_batch_inference.py b/examples/inference/advanced/simple_t5_batch_inference.py
index 1aca74b3176..3591936a2ae 100644
--- a/examples/inference/advanced/simple_t5_batch_inference.py
+++ b/examples/inference/advanced/simple_t5_batch_inference.py
@@ -4,9 +4,10 @@
import sys
from argparse import Namespace
+import pretrain_t5
import torch
+from pretrain_t5 import model_provider
-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 (
@@ -21,7 +22,6 @@
)
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))
diff --git a/examples/inference/launch_inference_server.py b/examples/inference/launch_inference_server.py
index c5e3289d277..c980e142c6d 100644
--- a/examples/inference/launch_inference_server.py
+++ b/examples/inference/launch_inference_server.py
@@ -39,14 +39,14 @@ def add_serve_args(parser: ArgumentParser) -> ArgumentParser:
group.add_argument("--coordinator-port", type=int, default=None)
group.add_argument("--host", type=str, default="0.0.0.0", help="HTTP bind host")
group.add_argument("--port", type=int, default=5000, help="HTTP bind port")
- group.add_argument(
- "--parsers", type=str, nargs="+", default=[], help="Response parser names"
- )
+ group.add_argument("--parsers", type=str, nargs="+", default=[], help="Response parser names")
group.add_argument(
"--verbose", action="store_true", default=False, help="Per-request HTTP logging"
)
group.add_argument(
- "--frontend-replicas", type=int, default=4,
+ "--frontend-replicas",
+ type=int,
+ default=4,
help="Number of HTTP frontend processes spawned on the primary rank.",
)
return parser
diff --git a/examples/inference/offline_inference.py b/examples/inference/offline_inference.py
index b39cd19903a..167ad1084d1 100644
--- a/examples/inference/offline_inference.py
+++ b/examples/inference/offline_inference.py
@@ -97,9 +97,8 @@ def _validate_prompt_lengths(args, llm, requests):
for idx, r in enumerate(requests)
if len(r.prompt_tokens) > llm.context.max_tokens
}
- assert not invalid, (
- "request idxs with prompts longer than context.max_tokens: "
- ", ".join(f"{k}({v})" for k, v in invalid.items())
+ assert not invalid, "request idxs with prompts longer than context.max_tokens: " ", ".join(
+ f"{k}({v})" for k, v in invalid.items()
)
@@ -118,9 +117,7 @@ def _print_setup_prefix(setup_prefix: str) -> None:
print("~~~")
-def _report_results(
- args, setup_prefix, results, throughputs, total_time, peak_mem_stats, captured
-):
+def _report_results(args, setup_prefix, results, throughputs, total_time, peak_mem_stats, captured):
if dist.get_rank() != 0:
return
@@ -139,9 +136,7 @@ def _report_results(
peak_resvd_gb = stats["reserved_bytes.all.peak"] / 1024**3
throughput = throughputs[-1] if throughputs else 0.0
capture_str = (
- f"{captured['capture_stats']['time']:.2f} sec"
- if captured["capture_stats"]
- else "--"
+ f"{captured['capture_stats']['time']:.2f} sec" if captured["capture_stats"] else "--"
)
print("~~~")
print(
@@ -276,7 +271,15 @@ def main():
prompts_list = [r.prompt_text for r in requests]
- runner_args = (args, model, tokenizer, inference_config, requests, prompts_list, sampling_params)
+ runner_args = (
+ args,
+ model,
+ tokenizer,
+ inference_config,
+ requests,
+ prompts_list,
+ sampling_params,
+ )
if args.async_mode:
asyncio.run(_run_async(*runner_args))
else:
diff --git a/examples/mimo/train.py b/examples/mimo/train.py
index 594170faa7e..52be3f7ec58 100644
--- a/examples/mimo/train.py
+++ b/examples/mimo/train.py
@@ -9,27 +9,25 @@
from functools import partial
from typing import Any, Dict, Iterator
-from megatron.training.argument_utils import pretrain_cfg_container_from_args
-from megatron.training.arguments import parse_and_validate_args
import torch
-from megatron.training import get_args, pretrain, print_rank_0
from megatron.core.parallel_state import (
+ get_context_parallel_group,
+ get_data_parallel_group,
get_tensor_model_parallel_group,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_src_rank,
- get_context_parallel_group,
- get_data_parallel_group,
)
+from megatron.training import get_args, pretrain, print_rank_0
+from megatron.training.argument_utils import pretrain_cfg_container_from_args
+from megatron.training.arguments import parse_and_validate_args
sys.path.append(
os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
)
from data.energon_avlm_task_encoder import llava_avlm_dataloader_provider
from data.energon_vlm_task_encoder import llava_vlm_dataloader_provider
-from data.mock import (
- train_valid_test_datasets_provider as mock_train_valid_test_datasets_provider,
-)
+from data.mock import train_valid_test_datasets_provider as mock_train_valid_test_datasets_provider
from model_providers.llava_avlm import model_provider_llava_avlm
from model_providers.llava_vlm import model_provider_llava_vlm
from model_providers.mock import model_provider_mock_vlm_single_encoder
@@ -51,13 +49,24 @@
"llava_avlm": llava_avlm_dataloader_provider,
}
+
def add_mimo_args(parser):
"""Add MIMO-specific arguments to the parser."""
group = parser.add_argument_group('MIMO', 'MIMO specific arguments')
# MIMO-specific parameters
- group.add_argument('--dataset-provider', type=str, default='mock', help='Dataset provider to choose from [mock, llava_vlm, video_llava_vlm, llava_avlm]')
- group.add_argument('--model-provider', type=str, default='mock', help='Model provider to choose from [mock, llava_vlm, video_llava_vlm, llava_avlm]')
+ group.add_argument(
+ '--dataset-provider',
+ type=str,
+ default='mock',
+ help='Dataset provider to choose from [mock, llava_vlm, video_llava_vlm, llava_avlm]',
+ )
+ group.add_argument(
+ '--model-provider',
+ type=str,
+ default='mock',
+ help='Model provider to choose from [mock, llava_vlm, video_llava_vlm, llava_avlm]',
+ )
# mock dataloader related args
# can control mock samples with total seq length and image seq length
@@ -72,17 +81,29 @@ def add_mimo_args(parser):
'--audio-encoder-model', type=str, default=None, help='Audio encoder model name'
)
group.add_argument(
- '--hf-assign-unused-tokens', type=str, nargs='+', default=None,
- help='Assigning unused tokens to special tokens. Example: '
- '--hf-assign-unused-tokens "