update to core_v0.9

CHANGELOG.md

+# Changelog
+
+## NVIDIA Megatron Core 0.9.0
+
+- Uneven pipeline parallelism
+  - Enable pipeline parallelism where first and last ranks have fewer transformer layers than the intermediate ranks
+- Per layer CUDAGraph support for GPT training with Transformer Engine modules
+- Enable different TP sizes for the vision encoder
+- Enable pipeline parallelism for T5 & Llava models
+- Support multi-tile multi-image input in Llava models
+- MoE
+  - FP8 support
+  - Runtime upcycling support
+  - Dispatcher implementation optimizations
+  - Shared expert support with overlapping optimizations
+    - Qwen Model support
+- Known Issues
+  - When using sequence parallel, during the transformer block forward pass, dropout is not using the appropriate rng context.
+
+
+## NVIDIA Megatron Core 0.8.0
+
+- Multimodal
+  - Added initial support for training vision language models using the LLaVA architecture
+  - Added initial support for inference with multimodal inputs
+  - End-to-end multimodal example from data collection to training to evaluation is provided in examples/multimodal
+- MoE
+  - Context Parallel support.
+  - Distributed checkpoint support for grouped GEMM.
+- Mamba
+
+## NVIDIA Megatron Core 0.7.0
+
+- MoE
+  - Token drop support
+  - Several efficiency optimizations
+  - Improved model parallelism
+  - Memory optimizations
+- Distributed checkpointing
+  - Enabled for Retro
+  - Asynchronous checkpoint saving
+- Several minor bug fixes, speed improvements, and memory optimizations
+
+## NVIDIA Megatron Core 0.6.0
+
+- MoE (Mixture of Experts)
+  - Performance optimization
+    - Communication optimization for multi GPU and Single GPU 
+    - 23% improvement (323 TFLOPS/GPU) over MCore 0.5.0 on Mixtral with Hopper BF16
+    - GroupedMLP enhancement for Hopper
+    - DP Overlapping. Support overlapping computation with gradient reduction and parameter gathering.
+  - All-to-All based Token Dispatcher
+  - Layer-wise logging for load balancing loss.
+  - Improved expert parallel support including distributed optimizer.
+- Distributed optimizer
+- RETRO
+  - Data processing
+- BERT
+  - Distributed checkpointing
+- Dist checkpointing
+  - PyTorch native distributed backend
+  - Improved saving/loading speed
+- TensorRT-LLM Export
+  - Integration with TensorRT Model Optimizer Post-training quantization (PTQ)
+  - Text generation driver to perform PTQ in Megatron-LM
+  - Llama2 and Nemotron3-8b examples to use TensorRT-LLM unified build API to build engine after training.
+- Several minor enhancements, bug fixes, and documentation updates
+
+## NVIDIA Megatron Core 0.5.0
+
+### Key Features and Enhancements
+
+Megatron core documentation is now [live!](https://docs.nvidia.com/megatron-core/developer-guide/latest/user-guide/index.html#quick-start)
+
+### Model Features
+
+- MoE (Mixture of Experts)
+  - Support for Z-loss, Load balancing and Sinkhorn
+  - Layer and communications refactor
+  - Richer parallelism mappings and EP can be combined with other model parallel techniques for larger MoE variants, e.g. EP + TP + DP + SP + PP
+  - Token dropless architecture with Top-K routing
+  - Performance optimization with with GroupedGEMM when number of local experts is > 1
+  - Distributed checkpointing
+- Interleaved rotary embedding
+
+### Datasets
+
+- Masked WordPiece datasets for BERT and T5
+- Raw and mock datasets
+
+### Parallelism
+
+### Performance
+
+- Activation offloading to CPU
+- Rope and Swiglu fusion
+- Sliding window attention (via Transformer Engine)
+
+### General Improvements
+
+- Timers
+
+## NVIDIA Megatron Core 0.4.0
+
+### Key Features and Enhancements
+
+#### Models
+
+- BERT
+- RETRO
+- T5
+
+#### Parallelism
+
+- Mixture of Experts support for GPT
+- Model parallel efficient Distributed Data Parallel (DDP)
+- Context Parallel (2D Tensor Parallel) support
+
+#### Datasets
+
+- GPT Dataset
+- Blended Dataset

CODEOWNERS

+[Core-ADLR] @mcore-reviewers/core-adlr
+megatron/core/ 
+
+[Core-NeMo] @mcore-reviewers/core-nemo
+megatron/core/ 
+
+^[Core-MLPerf] @mcore-reviewers/mlperf
+megatron/core/
+
+[MoE-ADLR] @mcore-reviewers/moe-adlr
+megatron/core/transformer/moe/
+
+[MoE-Moe] @mcore-reviewers/moe-moe
+megatron/core/transformer/moe/
+
+[Datasets] @mcore-reviewers/datasets
+megatron/core/datasets/
+
+[BERT] @mcore-reviewers/bert
+megatron/core/models/bert/
+
+[GPT] @mcore-reviewers/gpt
+megatron/core/models/gpt/
+
+[Retro] @mcore-reviewers/retro
+megatron/core/models/retro/
+
+[Distributed Checkpointing] @mcore-reviewers/dist-checkpointing
+megatron/core/dist_checkpointing/
+
+[Distributed Optimizer] @mcore-reviewers/dist-optimizer
+megatron/core/optimizer/distrib_optimizer/ 
+
+[Inference] @mcore-reviewers/inference
+megatron/core/inference/
+
+[Quantization and Inference (QAT)] @mcore-reviewers/quantization-and-inference
+megatron/core/inference/
+
+; [Context Parallelism] @mcore-reviewers/context-parallelism
+; 
+
+[CI] @mcore-reviewers/ci
+.gitlab-ci.yml
+Dockerfile.ci
+jet-tests.yml
+tests/

Dockerfile.ci

+# syntax=docker/dockerfile:experimental
+
+ARG FROM_IMAGE_NAME
+FROM $FROM_IMAGE_NAME as main
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN sed -i -e 's/^APT/# APT/' -e 's/^DPkg/# DPkg/' \
+      /etc/apt/apt.conf.d/docker-clean
+
+RUN apt-get update && \
+      apt-get install -y --no-install-recommends gettext && \
+      apt-get clean
+
+RUN 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 Mamba begin #####
+RUN pip uninstall -y triton && \
+    pip install triton==2.1.0
+
+# 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 && \
+    pip uninstall -y causal-conv1d && \
+    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 && \
+    pip uninstall -y mamba-ssm && \
+    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
+##### For Mamba end #####
+
+##### For JET-API start #####
+RUN apt-get update && \ 
+    apt-get install -y python3-venv && \
+    apt-get clean -y && \
+    python -m venv /opt/jet
+##### For JET-API end #####
+
+RUN pip3 install --no-cache-dir \
+      einops \
+      flask-restful \
+      nltk \
+      pytest \
+      pytest-cov \
+      pytest_mock \
+      pytest-random-order \
+      sentencepiece \
+      wrapt \
+      git+https://github.com/fanshiqing/grouped_gemm@v1.1.2 \
+      zarr \
+      tensorstore==0.1.45 \
+      wandb
+
+COPY . /workspace/megatron-lm
+
+COPY . /workspace/megatron-lm
+RUN cp -r /workspace/megatron-lm /opt && \
+    pip install /opt/megatron-lm  
+
+
+##### For NVIDIANS only #####
+FROM main as jet
+ARG CACHEBUST=0
+RUN --mount=type=secret,id=JET_INDEX_URLS \
+    JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) && \
+    /opt/jet/bin/pip install jet-api --upgrade $JET_INDEX_URLS 
+ENV PATH="$PATH:/opt/jet/bin"
+###
\ No newline at end of file

Dockerfile.linting

+# syntax=docker/dockerfile:experimental
+
+ARG FROM_IMAGE_NAME
+FROM $FROM_IMAGE_NAME as main
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN sed -i -e 's/^APT/# APT/' -e 's/^DPkg/# DPkg/' \
+      /etc/apt/apt.conf.d/docker-clean
+
+
+RUN pip3 install --no-cache-dir \
+      black==24.4.2 \
+      isort==5.13.2 \
+      flake8==7.1.0 \
+      pylint==3.2.6
+
+COPY . /opt/megatron-lm
+
+WORKDIR /opt/megatron-lm
+
+FROM main as jet
\ No newline at end of file

LICENSE

@@ -29,13 +29,14 @@ The following applies to all files unless otherwise noted:
 --
 
 This repository also contains code from Hugging Face Inc., Google Research,
-Facebook (from their Fairseq and Dino projects), Microsoft(from their 
-Swin-Transformer project)and Philip Popien. Files from these 
-organizations have notices at  the top of each file. Below are 
-licenses used in those files, as indicated.
+Facebook (from their Fairseq, Dino, and ParlAI projects), Microsoft (from their
+Swin-Transformer project), Philip Popien, and the Mamba project (Tri Dao and
+Albert Gu). Files from these organizations have notices at the top of each file.
+Below are licenses used in those files, as indicated.
 
 
-------------- LICENSE FOR Facebook, huggingface and Google Research code  --------------
+--------------------------------------------------------------------------------
+-- LICENSE FOR Facebook, huggingface, Google Research, LLaVA, and Mamba code  --
 
 
                                  Apache License
@@ -240,12 +241,10 @@ licenses used in those files, as indicated.
    See the License for the specific language governing permissions and
    limitations under the License.
 
-------------- LICENSE FOR Facebook Fairseq code --------------
+------------- LICENSE FOR Facebook, Inc. and its affiliates, Meta Platforms, Inc. and its affiliates, Microsoft Corporation, and OpenGVLab/InternVL  --------------
 
 MIT License
 
-Copyright (c) Facebook, Inc. and its affiliates.
-
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
 in the Software without restriction, including without limitation the rights
@@ -264,28 +263,3 @@ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 
-------------- LICENSE FOR Mircrosoft Swin transformer code --------------
-
-MIT License
-
-Copyright (c) Microsoft Corporation.
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE
-
-

MANIFEST.in

+include megatron/core/requirements.txt
+include megatron/core/README.md
\ No newline at end of file

README.md

@@ -6,7 +6,6 @@
 - [训练](#训练)
   - [数据预处理](#数据预处理)
   - [GPT预训练](#gpt预训练)
-    - [单卡训练](#单卡训练)
     - [分布式多卡训练](#分布式多卡训练)
 - [GPT文本生成](#gpt文本生成)
 - [参考](#参考)
@@ -16,23 +15,23 @@
 <pre>
 pip install -r requirements.txt
 </pre>
-
 2. 安装DCU相关whl包
+
 DCU相关包下载目录：[https://cancon.hpccube.com:65024/4/main](https://cancon.hpccube.com:65024/4/main)
 
-pytorch whl包：pytorch ---> dtk-24.04
+pytorch whl包：pytorch ---> dtk-24.04.1
 根据python版本,下载对应pytorch的whl包
 
 <pre>
 pip install torch* (下载的torch的whl包)
 </pre>
-torchvision whl包：vision ---> dtk-24.04
+torchvision whl包：vision ---> dtk-24.04.1
 根据python版本,下载对应torchvision的whl包
 
 <pre>
 pip install torchvision* (下载的torchvision的whl包)
 </pre>
-apex whl包：apex ---> dtk-24.04
+apex whl包：apex ---> dtk-24.04.1
 根据python版本,下载对应apex的whl包
 
 <pre>
@@ -40,6 +39,13 @@ pip install apex* (下载的apex的whl包)
 </pre>
 若使用 pip install 下载安装过慢，可添加源：-i https://pypi.tuna.tsinghua.edu.cn/simple/
 
+3. 安装unsloth
+
+<pre>
+cd ./unsloth
+pip3 install -e .
+</pre>
+
 # 下载词汇文件
 
 <pre>
@@ -63,8 +69,7 @@ xz -d oscar-1GB.jsonl.xz
 python tools/preprocess_data.py \
     --input oscar-1GB.jsonl \ 
     --output-prefix ./dataset/my-gpt2 \
-    --vocab gpt2-vocab.json \
-    --dataset-impl mmap \
+    --vocab-file gpt2-vocab.json \
     --tokenizer-type GPT2BPETokenizer \
     --merge-file gpt2-merges.txt \
     --append-eod \
@@ -74,8 +79,7 @@ python tools/preprocess_data.py \
 参数说明
 --input				输入数据集路径，即oscar-1GB.jsonl.xz解压后的文件路径
 --output-prefix		输出数据路径，处理后会自动加上_text_document后缀
---vocab				下载的gpt2-vocab.json词表文件路径
---dataset-impl		dataset类型
+--vocab-file				下载的gpt2-vocab.json词表文件路径
 --tokenizer-type 	tokenizer类型
 --merge-file		下载的gpt2-merges.txt文件路径		
 --append-eod		添加结束标志符		
@@ -89,16 +93,16 @@ python tools/preprocess_data.py \
   ```bash
   VOCAB_FILE=gpt2-vocab.json
   MERGE_FILE=gpt2-merges.txt
-  DATA_PATH=./dataset/my-gpt2_text_document
+  DATA_PATH="./dataset/my-gpt2_text_document"
   ```
 
 - 执行多卡训练
 
   ```
   #np为起的进程数，np\hostfile均需按实际填写
-  mpirun -np 4 --hostfile hostfile single.sh（基于单节点四卡）
+  mpirun -np 4 --hostfile hostfile single.sh localhost（基于单节点四卡）
   ```
 
 # 参考
 
-- [README_ORIGIN](README_ORIGIN.md)
\ No newline at end of file
+- [README_ORIGIN](README_ORIGIN.md)

docs/distrib_optimizer.md

-# 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
-
-![Data flow](images/distrib_optimizer/data_flow.png)
-
-## Sharding scheme
-
-![Sharding scheme](images/distrib_optimizer/sharding_scheme.png)
-
-## 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)

docs/llama_mistral.md

+# 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 Huggingface.
+
+The [Llama-2](https://ai.meta.com/llama/) and [Llama-3](https://llama.meta.com/) family of models are an open-source set of pretrained & finetuned (for chat) models that have achieved strong results across a wide set of benchmarks. At their times of release, both Llama-2 and Llama-3 models achieved among the best results for open-source models, and were competitive with leading closed-source models (see https://arxiv.org/pdf/2307.09288.pdf and https://ai.meta.com/blog/meta-llama-3/).
+
+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.
+
+# 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 Meta/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.
+
+## Contents
+  * [Download Meta or Huggingface checkpoints](#download-meta-or-huggingface-checkpoints)
+  * [Convert checkpoint format](#convert-checkpoint-format)
+    * [Meta format](#meta-format)
+    * [Huggingface format](#huggingface-format)
+  * [Launch model](#launch-model)
+    * [Megatron](#launch-megatron)
+    * [Meta](#launch-meta)
+    * [Huggingface](#launch-hf)
+  * [Benchmark results](#benchmark-results)
+
+## Download Meta or Huggingface checkpoints
+
+Users must first apply for access to download the Llama-2 checkpoints either directly from [Meta](https://ai.meta.com/resources/models-and-libraries/llama-downloads/) or through [Huggingface](https://huggingface.co/docs/transformers/main/model_doc/llama2) (HF). The checkpoints are available in two formats, Meta's native format (available from both the Meta and HF links), and HF's format (available only from HF). Either format can be converted to Megatron, as detailed next.
+
+## Convert checkpoint format
+
+We recommend passing `--dtype bf16` for training or finetuning. Inference can be done in bfloat16 or float16.
+
+### Meta format
+
+The Meta format checkpoints are converted to HF format as an intermediate step before converting to Megatron format. The `transformers` package is required, and must have version >=4.31.0 (e.g., `pip install transformers>=4.31.0`). (**Note**: we have specifically tested with versions `4.31.0` and `4.32.0`; your experience may vary with newer versions.) Assuming the downloaded checkpoints are in `$CHECKPOINT_DIR` (with separate sub-directories for 7B, 13B, 70B, etc.), the following example command can be used to convert from Llama-2 format to HF format in bfloat16:
+
+```
+python tools/checkpoint/convert.py --model-type GPT \
+>   --loader llama_mistral \
+>   --saver megatron \
+>   --checkpoint-type meta \
+>   --model-size llama2-7B \
+>   --load-dir $LLAMA_META_FORMAT_DIR \
+>   --save-dir ${MEGATRON_FORMAT_DIR} \
+>   --tokenizer-model ${TOKENIZER_MODEL} \
+>   --target-tensor-parallel-size ${TP} \
+>   --target-pipeline-parallel-size ${PP} \
+>   --bf16
+```
+
+Valid values for `--model-size` are `llama2-7B`, `llama2-13B`, and `llama2-70B` (for pretrained-only models), and `llama2-7Bf`, `llama2-13Bf`, and `llama2-70Bf` (for chat-finetuned models).
+
+### Huggingface format
+
+The HF checkpoints can be converted to Megatron format by using Megatron's own Llama-2 checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`). One important argument that must be set correctly is the tensor parallel size (`TP`) for each model. The following table shows these values:
+
+| Model size | Tensor parallel size (`TP`) |
+| ---------- | --------------------------- |
+|  7B        | 1                           |
+| 13B        | 2                           |
+| 70B        | 8                           |
+
+Using these values for `TP`, along with the path to the Llama-2 tokenizer model (automatically downloaded with original checkpoint download; see `${TOKENIZER_MODEL}` below), run the following command from the root of your Megatron source code to convert from HF format to Megatron format:
+
+```
+$>: python tools/checkpoint/convert.py \
+ >    --model-type GPT \
+ >    --loader llama_mistral \
+ >    --saver megatron \
+ >    --target-tensor-parallel-size ${TP} \
+ >    --checkpoint-type hf
+ >    --load-dir ${HF_FORMAT_DIR} \
+ >    --save-dir ${MEGATRON_FORMAT_DIR} \
+ >    --tokenizer-model ${TOKENIZER_MODEL}
+```
+
+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 Meta
+
+Meta checkpoints can be launched with: https://github.com/facebookresearch/llama
+
+### Launch Huggingface
+
+Huggingface checkpoints can be launched with: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+
+## 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: `|<llama_score> - <megatron_score>| / <llama_score>`, 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
+
+Llama-3 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.
+
+## Contents
+  * [Download Huggingface checkpoints](#download-huggingface-checkpoints)
+  * [Convert checkpoint format](#convert-checkpoint-format)
+    * [Huggingface format](#huggingface-format)
+  * [Validate checkpoint](#optional-validate-checkpoint)
+  * [Launch model](#launch-model)
+
+## Download Huggingface checkpoints
+
+Users must first apply for access to download the Llama-3 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's own Llama-3 checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`). One important argument that must be set correctly is the tensor parallel size (`TP`) for each model. The following table shows these values:
+
+| Model size | Tensor parallel size (`TP`) |
+| ---------- | --------------------------- |
+|  8B        | 1                           |
+| 70B        | 8                           |
+
+Using these values for `TP`, along with the path to the Llama-3 tokenizer model (automatically downloaded with original checkpoint download; see `${TOKENIZER_MODEL}` below), run the following command from the root of your Megatron source code to convert from HF format to Megatron format:
+
+```
+$>: python tools/checkpoint/convert.py \
+ >    --bf16 \
+ >    --model-type GPT \
+ >    --loader llama_mistral \
+ >    --saver mcore \
+ >    --target-tensor-parallel-size ${TP} \
+ >    --checkpoint-type hf
+ >    --load-dir ${HF_FORMAT_DIR} \
+ >    --save-dir ${MEGATRON_FORMAT_DIR} \
+ >    --tokenizer-model ${TOKENIZER_MODEL}
+ >    --model-size llama3-8B \
+```
+
+Valid values for `--model-size` are `llama3-8B` and `llama3-70B` (for pretrained-only models), and `llama3-8Bf` and `llama3-70Bf` (for chat-finetuned models).
+
+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/llama_mistral/run_text_generation_llama3.sh <PATH_TO_CONVERTED_MCORE_CHECKPOINT> <PATH_TO_DOWNLOADED_HUGGINGFACE_CHECKPOINT>`.
+
+Once running, query the server with `curl 'http://<TEXT_GENERATION_SERVER_IP>:5000/api' -X 'PUT' -H 'Content-Type: application/json; charset=UTF-8'  -d '{"prompts":["<SOME_PROMPT>"], "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 <PATH_TO_DOWNLOADED_HUGGINGFACE_CHECKPOINT> --prompt <SOME_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 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 \
+```
+
+# 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.
+
+## Contents
+  * [Download Huggingface checkpoints](#download-huggingface-checkpoints)
+  * [Convert checkpoint format](#convert-checkpoint-format)
+  * [(Optional) Validate checkpoint](#optional-validate-checkpoint)
+  * [Launch model](#launch-model)
+
+## Download Huggingface checkpoints
+
+Users must first apply for access to download the Mistral-7b checkpoints through [Huggingface](https://huggingface.co/mistralai/Mistral-7B-v0.3) (HF).
+
+## Convert checkpoint format
+
+The HF checkpoints can be converted to Megatron format by using Megatron's own Mistral checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`).
+
+Using the path to the Mistral tokenizer model (downloaded alongside the HF checkpoint), run the following command from the root of your Megatron source code to convert from HF format to mcore format:
+
+```
+$>: python tools/checkpoint/convert.py \
+ >    --bf16 \
+ >    --model-type GPT \
+ >    --loader llama_mistral \
+ >    --saver mcore \
+ >    --target-tensor-parallel-size ${TP} \
+ >    --checkpoint-type hf \
+ >    --load-dir ${HF_FORMAT_DIR} \
+ >    --save-dir ${MEGATRON_FORMAT_DIR} \
+ >    --tokenizer-model ${TOKENIZER_MODEL} \
+ >    --model-size mistral-7B \
+```
+
+Valid values for `--model-size` are mistral-7B for the pretrained model or mistral-7Bf for the chat fine-tuned model.
+
+After this conversion, we are ready to load the checkpoints into an mcore GPT model.
+
+## (Optional) Validate checkpoints
+
+A Megatron-LM text generation server for Mistral-7B can be launched using the script `examples/llama_mistral/run_text_generation_mistral.sh <PATH_TO_CONVERTED_MCORE_CHECKPOINT> <PATH_TO_DOWNLOADED_HUGGINGFACE_CHECKPOINT>`.
+
+Once running, query the server with `curl 'http://<TEXT_GENERATION_SERVER_IP>:5000/api' -X 'PUT' -H 'Content-Type: application/json; charset=UTF-8'  -d '{"prompts":["<SOME_PROMPT>"], "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 <PATH_TO_DOWNLOADED_HUGGINGFACE_CHECKPOINT> --prompt <SOME_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 use the Llama architecture and may be converted from HuggingFace to Megatron using the commands in [Llama3](#llama-3).

docs/source/api-guide/context_parallel.rst

+context\_parallel package
+=========================
+
+Context parallelism overview 
+----------------------------
+
+.. figure:: ../images/context_parallel/CP_overview.png
+   :alt: cp_overview
+   :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 dimension of sequence length. Unlike prior SP (sequence parallelism) which only splits the sequence of Dropout and LayerNorm activations, CP partitions the network inputs and all activations along sequence dimension. With CP, all modules except attention (e.g., Linear, LayerNorm, etc.) can work as usual without any changes, because they do not have inter-token operations. As for attention, the Q (query) of each token needs to compute with the KV (key and value) of all tokens in the same sequence. Hence, CP requires additional all-gather across GPUs to collect the full sequence of KV. Correspondingly, reduce-scatter should be applied to the activation gradients of KV in backward propagation. To reduce activation memory footprint, each GPU only stores the KV of a sequence chunk in forward and gathers KV again in backward. KV communication happens between a GPU and its counterparts in other TP groups. The all-gather and reduce-scatter are transformed to point-to-point communications in ring topology under the hood. Exchanging KV also can leverage MQA/GQA to reduce communication volumes, as they only have one or few attention heads for KV.
+
+For example, in Figure 1, assuming sequence length is 8K, each GPU processes 4K tokens. GPU0 and GPU2 compose a CP group, they exchange KV with each other. Same thing also happens between GPU1 and GPU3. CP is similar to `Ring Attention <https://arxiv.org/abs/2310.01889>`_ but provides better performance by (1) leveraging the latest OSS and cuDNN flash attention kernels; (2) removing unnecessary computation resulted from low-triangle causal masking and achieving optimal load balance among GPUs.
+
+Context parallelism benefits 
+----------------------------
+
+.. figure:: ../images/context_parallel/CP_results.png
+   :alt: cp_results
+   :align: center
+   
+   Figure 2: Speedup of 175B GPT with various TP+CP combinations vs. full recompute (i.e., TP8CP1).
+
+LLM encounters OOM (out of memory) issue with long context (i.e., long sequence length) because of linearly increasing memory footprint of activations. Recomputing activations in backward can avoid OOM but also introduce significant overheads (~30% with full recompute). Enlarging TP (tensor model parallelism) can fix the OOM issue as well, but it potentially makes compute (e.g., Linear) too short to overlap communication latencies. To be clear, scaling out to more GPUs with bigger TP can hit the overlapping problem no matter if OOM happens.
+
+CP can better address the issues. With CP, each GPU only computes on a part of the sequence, which reduces both computation and communication by CP times. Therefore, there are no concerns about the overlapping between them. The activation memory footprint per GPU is also CP times smaller, hence no OOM issue any more. As Figure 2 shows, the combinations of TP and CP can achieve optimal performance by eliminating recompute overheads and making the best tradeoff between computation and communications.
+
+Enabling context parallelism
+----------------------------
+
+CP support has been added to GPT. All models that share GPT code path also should be able to benefit from CP, such as Llama. CP can work with TP (tensor model parallelism), PP (pipeline model parallelism), and DP (data parallelism), where the total number of GPUs equals TPxCPxPPxDP. CP also can work with different attention variants, including MHA/MQA/GQA, uni-directional and bi-directional masking.
+
+CP is enabled by simply setting context_parallel_size=<CP_SIZE> in command line. Default context_parallel_size is 1, which means CP is disabled. Running with CP requires Megatron-Core (>=0.5.0) and Transformer Engine (>=1.1).

docs/source/api-guide/datasets.rst

+datasets package
+================
+
+.. mdinclude :: ../../../megatron/core/datasets/readme.md
+
+Submodules
+----------
+
+datasets.blended\_megatron\_dataset\_config module
+---------------------------------------------------
+
+.. automodule:: core.datasets.blended_megatron_dataset_config
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.blended\_megatron\_dataset\_builder module
+---------------------------------------------------
+
+.. automodule:: core.datasets.blended_megatron_dataset_builder
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.megatron\_tokenizer module
+-----------------------------------
+
+.. automodule:: core.datasets.megatron_tokenizer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.indexed\_dataset module
+--------------------------------
+
+.. automodule:: core.datasets.indexed_dataset
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.megatron\_dataset module
+---------------------------------
+
+.. automodule:: core.datasets.megatron_dataset
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.gpt\_dataset module
+----------------------------
+
+.. automodule:: core.datasets.gpt_dataset
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.masked\_dataset module
+-------------------------------
+
+.. automodule:: core.datasets.masked_dataset
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.bert\_dataset module
+-----------------------------
+
+.. automodule:: core.datasets.bert_dataset
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.t5\_dataset module
+---------------------------
+
+.. automodule:: core.datasets.t5_dataset
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.blended\_dataset module
+----------------------------------
+
+.. automodule:: core.datasets.blended_dataset
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+datasets.utils module
+---------------------
+
+.. automodule:: core.datasets.utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.datasets
+   :members:
+   :undoc-members:
+   :show-inheritance:
+

docs/source/api-guide/dist_checkpointing.rst

+dist\_checkpointing package
+===========================
+
+A library for saving and loading the distributed checkpoints.
+A "distributed checkpoint" can have various underlying formats (current default format is based on Zarr)
+but has a distinctive property - the checkpoint saved in one parallel configuration (tensor/pipeline/data parallelism)
+can be loaded in a different parallel configuration.
+
+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.
+
+
+Subpackages
+-----------
+
+.. toctree::
+   :maxdepth: 4
+
+   dist_checkpointing.strategies
+
+Submodules
+----------
+
+dist\_checkpointing.serialization module
+----------------------------------------
+
+.. automodule:: core.dist_checkpointing.serialization
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+dist\_checkpointing.mapping module
+----------------------------------
+
+.. automodule:: core.dist_checkpointing.mapping
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+dist\_checkpointing.optimizer module
+------------------------------------
+
+.. automodule:: core.dist_checkpointing.optimizer
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+dist\_checkpointing.core module
+-------------------------------
+
+.. automodule:: core.dist_checkpointing.core
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+dist\_checkpointing.dict\_utils module
+--------------------------------------
+
+.. automodule:: core.dist_checkpointing.dict_utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+
+dist\_checkpointing.utils module
+--------------------------------
+
+.. automodule:: core.dist_checkpointing.utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.dist_checkpointing
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/source/api-guide/dist_checkpointing.strategies.rst

+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.
+
+Submodules
+----------
+
+dist\_checkpointing.strategies.base module
+------------------------------------------
+
+.. automodule:: core.dist_checkpointing.strategies.base
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+dist\_checkpointing.strategies.tensorstore module
+-------------------------------------------------
+
+.. automodule:: core.dist_checkpointing.strategies.tensorstore
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+dist\_checkpointing.strategies.two\_stage module
+------------------------------------------------
+
+.. automodule:: core.dist_checkpointing.strategies.two_stage
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+dist\_checkpointing.strategies.zarr module
+------------------------------------------
+
+.. automodule:: core.dist_checkpointing.strategies.zarr
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.dist_checkpointing.strategies
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/source/api-guide/dist_optimizer.md

+# Distributed Optimizer
+
+The motivation for the distributed optimizer is to save memory by distributing the optimizer state evenly across data parallel ranks (https://arxiv.org/abs/1910.02054), versus the naive method of replicating the optimizer state across data parallel ranks.
+
+Theoretical memory savings vary depending on the combination of the datatype of the model's parameters (`param_dtype`) and main gradients accumulated across data-parallel replicas (`grad_dtype`). We always use `fp32` main parameters for optimizer steps. In the current implementation, the theoretical number of bytes per parameter is (where d is the data parallel size):
+
+|        | 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  |
+
+Our implementation of the distributed optimizer uses contiguous buffers for parameters and main gradients; model gradients are copied over to the main gradients as soon as they are fully computed.
+
+The figures below illustrate the distributed optimizer's sharding scheme, and the key steps of the distributed optimizer's parameter update:
+
+## Data flow
+
+![Data flow](../images/distrib_optimizer/data_flow.png)
+
+## Sharding scheme
+
+![Sharding scheme](../images/distrib_optimizer/sharding_scheme.png)
+
+## Key steps
+
+_(note: using illustrations above, assuming `bf16` model weights, `bf16` model gradients that are computed by the backward pass and `fp32` main gradients that are also used for optimizer steps; we always use `fp32` main weights for optimizer steps)_
+
+- Backward pass finishes (gradient buffer holds 16 `fp32` gradient elements).
+- Call reduce-scatter on each DP rank.
+- Each DP rank now has 4 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 4 `fp32` main parameter elements into the corresponding `bf16` parameter buffer (each element is cast from fp32 to fp16).
+- Call all-gather on each DP rank.
+- The parameter buffer now contains all 16, fully updated, `bf16` model parameter elements. Parameters in PyTorch modules already point to the appropriate locations in this parameter buffer, and thus forward passes are ready to run after the all-gather completes.
+- At this point, the gradient buffer is also ready to be zero'd for the next iteration.

docs/source/api-guide/distributed.rst

+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).
+
+Submodules
+----------
+
+distributed.distributed\_data\_parallel
+---------------------------------------
+
+Model wrapper for distributed data parallelism. Stores gradients in a
+contiguous buffer, and supports the option of overlapping communication
+(all-reduce or reduce-scatter) with backprop computation by breaking up
+full model's gradients into smaller buckets and running all-reduce /
+reduce-scatter on each bucket asynchronously. 
+
+.. automodule:: core.distributed.distributed_data_parallel
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+distributed.finalize\_model\_grads
+----------------------------------
+
+Finalize model gradients for optimizer step across all used parallelism modes.
+Synchronizes the all-reduce / reduce-scatter of model gradients across DP replicas,
+all-reduces the layernorm gradients for sequence parallelism, embedding gradients
+across first and last pipeline stages (if not tied), and expert gradients for expert
+parallelism.
+
+.. automodule:: core.distributed.finalize_model_grads
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+
+Module contents
+---------------
+
+Contains functionality to synchronize gradients across different ranks before
+optimizer step.
+
+.. automodule:: core.distributed
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/source/api-guide/fusions.rst

+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.
+
+Submodules
+----------
+
+fusions.fused\_bias\_dropout module
+-----------------------------------
+
+This module uses PyTorch JIT to fuse the bias add and dropout operations. Since dropout is not used during inference, different functions are used when in train mode and when in inference mode.
+
+.. automodule:: core.fusions.fused_bias_dropout
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+fusions.fused\_bias\_gelu module
+--------------------------------
+
+This module uses PyTorch JIT to fuse the bias add and GeLU nonlinearity operations.
+
+.. automodule:: core.fusions.fused_bias_gelu
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+fusions.fused\_layer\_norm module
+---------------------------------
+
+This module provides a wrapper around various fused LayerNorm implementation in Apex.
+
+.. automodule:: core.fusions.fused_layer_norm
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+fusions.fused\_softmax module
+-----------------------------
+
+This module provides wrappers around variations of Softmax in Apex.
+
+.. automodule:: core.fusions.fused_softmax
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+fusions.fused\_cross\_entropy\_loss module
+------------------------------------------
+
+This module uses PyTorch JIT to fuse the cross entropy loss calculation and batches communication calls.
+
+.. automodule:: core.fusions.fused_cross_entropy
+   :members:
+   :undoc-members:
+   :show-inheritance:
+