-Below are some of the projects where we have directly used Megatron:
-* [BERT and GPT Studies Using Megatron](https://arxiv.org/pdf/1909.08053.pdf)
-* [BioMegatron: Larger Biomedical Domain Language Model](https://www.aclweb.org/anthology/2020.emnlp-main.379.pdf)
-* [End-to-End Training of Neural Retrievers for Open-Domain Question Answering](https://arxiv.org/abs/2101.00408)
-* [Large Scale Multi-Actor Generative Dialog Modeling](https://www.aclweb.org/anthology/2020.acl-main.8.pdf)
-* [Local Knowledge Powered Conversational Agents](https://arxiv.org/abs/2010.10150)
-* [MEGATRON-CNTRL: Controllable Story Generation with External Knowledge Using Large-Scale Language Models](https://www.aclweb.org/anthology/2020.emnlp-main.226.pdf)
-* [RACE Reading Comprehension Dataset Leaderboard](http://www.qizhexie.com/data/RACE_leaderboard.html)
-* [Training Question Answering Models From Synthetic Data](https://www.aclweb.org/anthology/2020.emnlp-main.468.pdf)
-* [Few-shot Instruction Prompts for Pretrained Language Models to Detect Social Biases](https://arxiv.org/abs/2112.07868)
-* [Exploring the Limits of Domain-Adaptive Training for Detoxifying Large-Scale Language Models](https://arxiv.org/abs/2202.04173)
-* [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://arxiv.org/abs/2201.11990)
-* [Multi-Stage Prompting for Knowledgeable Dialogue Generation](https://arxiv.org/abs/2203.08745)
+Megatron-LM & Megatron-Core
+===========================
+
+
+[](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html)
+[](./setup.py)
+[](./LICENSE)
-Megatron is also used in [NeMo Megatron](https://developer.nvidia.com/nvidia-nemo#nemo-megatron), a framework to help enterprises overcome the challenges of building and training sophisticated natural language processing models with billions and trillions of parameters.
+
-Our codebase is capable of efficiently training very large (hundreds of billions of parameters) language models with both model and data parallelism. To demonstrate how the code scales with multiple GPUs and model sizes, we consider GPT models from 1 billion all the way to 1 trillion parameters. All models use a vocabulary size of 51,200 and a sequence length of 2048. We vary hidden size, number of attention heads, and number of layers to arrive at a specifc model size. As the model size increases, we also modestly increase the batch size. We leverage [NVIDIA's Selene supercomputer](https://www.top500.org/system/179842/) to perform scaling studies and use up to 3072 [A100](https://www.nvidia.com/en-us/data-center/a100/) GPUs for the largest model. Each cluster node has 8 NVIDIA 80GB A100 GPUs. The graph below shows that we scale nearly linear up to 1 trillion parameter models running on 3072 GPUs. Note that these results are from benchmark runs and these models were not trained to convergence; however, the FLOPs are measured for end-to-end training, i.e., includes all operations including data loading, optimization, and even logging.
+# Latest News
-
+- **[2024/7]** Megatron-Core v0.7 improves scalability and training resiliency and adds support for multimodal training ([blog](https://developer.nvidia.com/blog/train-generative-ai-models-more-efficiently-with-new-nvidia-megatron-core-functionalities/)).
+- **[2024/6]** Megatron-Core added supports for Mamba-based models. Check out our paper [An Empirical Study of Mamba-based Language Models](https://arxiv.org/pdf/2406.07887) and [code example](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba).
+- **[2024/1 Announcement]** NVIDIA has released the core capabilities in **Megatron-LM** into [**Megatron-Core**](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core) in this repository. Megatron-Core expands upon Megatron-LM's GPU-optimized techniques with more cutting-edge innovations on system-level optimizations, featuring composable and modular APIs. Explore the [Megatron-Core intro](#megatron-core) for more details.
-The following table shows both model (MFU) and hardware (HFU) FLOPs utilization for select configurations up to 1T parameters (see [our paper](https://arxiv.org/pdf/2205.05198) for a description of how these are calculated). As the model size increases, we achieve better GPU utilization and for the one trillion parameter model, we reach a MFU and HFU of 56.3% and 57.0%, respectively. Note that these numbers are also measured on benchmark runs and in this case are measured using a data parallel size of one. Data parallelism introduces some overhead due to the gradient all-reduce required between the data parallel groups. However, for large transformer models, this overhead is not large and can almost entirely eliminted by overlapping the gradient all-reduce with backpropagation.
-| Model Size | Model FLOPs Utilization | Hardware FLOPs Utilization |
-| :---: | :---: | :---: |
-| 22B | 41.5% | 43.7% |
-| 175B | 51.4% | 52.8% |
-| 530B | 56.0% | 57.0% |
-| 1T | 56.3% | 57.0% |
-# Contents
- * [Contents](#contents)
+# Table of Contents
+ * [Megatron Overview](#megatron-overview)
+ * [Megatron-LM](#megatron-lm)
+ * [Megatron-Core](#megatron-core)
+ * [Training Speed and Scalability](#training-speed-and-scalability)
* [Setup](#setup)
* [Downloading Checkpoints](#downloading-checkpoints)
* [Usage](#usage)
@@ -44,7 +36,7 @@ The following table shows both model (MFU) and hardware (HFU) FLOPs utilization
* [Distributed Optimizer](#distributed-optimizer)
* [FlashAttention](#flashattention)
* [GPT-3 Example](#gpt-3-example)
- * [Retro](#retro)
+ * [Retro and InstructRetro](#retro-and-instructretro)
* [Evaluation and Tasks](#evaluation-and-tasks)
* [GPT Text Generation](#gpt-text-generation)
* [GPT Evaluation](#gpt-evaluation)
@@ -53,9 +45,40 @@ The following table shows both model (MFU) and hardware (HFU) FLOPs utilization
* [BERT Task Evaluation](#bert-task-evaluation)
* [RACE Evaluation](#race-evaluation)
* [MNLI Evaluation](#mnli-evaluation)
+ * [Llama-2 Inference and Finetuning](#llama-2-inference-and-finetuning)
* [Datasets](#datasets)
* [Collecting Wikipedia Training Data](#collecting-wikipedia-training-data)
* [Collecting GPT Webtext Data](#collecting-gpt-webtext-data)
+ * [Reproducibility](#reproducibility)
+ * [Projects using Megatron](#projects-using-megatron)
+
+# Megatron Overview
+This repository comprises two essential components: **Megatron-LM** and **Megatron-Core**. Megatron-LM serves as a ressearch-oriented framework leveraging Megatron-Core for large language model (LLM) training. Megatron-Core, on the other hand, is a library of GPU optimized training techniques that comes with formal product support including versioned APIs and regular releases. You can use Megatron-Core alongside Megatron-LM or [Nvidia NeMo Framework](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/nlp/nemo_megatron/mcore_customization.html) for an end-to-end and cloud-native solution. Alternatively, you can integrate Megatron-Core's building blocks into your preferred training framework.
+
+## Megatron-LM
+First introduced in 2019, Megatron ([1](https://arxiv.org/pdf/1909.08053.pdf), [2](https://arxiv.org/pdf/2104.04473.pdf), and [3](https://arxiv.org/pdf/2205.05198)) sparked a wave of innovation in the AI community, enabling researchers and developers to utilize the underpinnings of this library to further LLM advancements. Today, many of the most popular LLM developer frameworks have been inspired by and built directly leveraging the open-source Megatron-LM library, spurring a wave of foundation models and AI startups. Some of the most popular LLM frameworks built on top of Megatron-LM include [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [HuggingFace Accelerate](https://github.com/huggingface/accelerate), and [NVIDIA NeMo Framework](https://www.nvidia.com/en-us/ai-data-science/generative-ai/nemo-framework/). A list of projects that have directly used Megatron can be found [here](#projects-using-megatron).
+
+## Megatron-Core
+Megatron-Core is an open-source PyTorch-based library that contains GPU-optimized techniques and cutting-edge system-level optimizations. It abstracts them into composable and modular APIs, allowing full flexibility for developers and model researchers to train custom transformers at-scale on NVIDIA accelerated computing infrastructure. This library is compatible with all NVIDIA Tensor Core GPUs, including FP8 acceleration support for [NVIDIA Hopper architectures](https://www.nvidia.com/en-us/data-center/technologies/hopper-architecture/).
+
+Megatron-Core offers core building blocks such as attention mechanisms, transformer blocks and layers, normalization layers, and embedding techniques. Additional functionality like activation recomputation, distributed checkpointing is also natively built-in to the library. The building blocks and functionality are all GPU optimized, and can be built with advanced parallelization strategies for optimal training speed and stability on NVIDIA Accelerated Computing Infrastructure. Another key component of the Megatron-Core library includes advanced model parallelism techniques (tensor, sequence, pipeline, context, and MoE expert parallelism).
+
+Megatron-Core can be used with [NVIDIA NeMo](https://www.nvidia.com/en-us/ai-data-science/products/nemo/), an enterprise-grade AI platform. Alternatively, you can explore Megatron-Core with the native PyTorch training loop [here](https://github.com/NVIDIA/Megatron-LM/tree/main/examples). Visit [Megatron-Core documentation](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html) to learn more.
+
+
+# Training Speed and Scalability
+Our codebase is capable of efficiently training large language models (i.e., models with hundreds of billions of parameters) with both model and data parallelism. To demonstrate how our software scales with multiple GPUs and model sizes, we consider GPT models ranging from 2 billion parameters to 462 billion parameters. All models use a vocabulary size of 131,072 and a sequence length of 4096. We vary hidden size, number of attention heads, and number of layers to arrive at a specific model size. As the model size increases, we also modestly increase batch size. Our experiments use up to 6144 [H100](https://www.nvidia.com/en-us/data-center/h100/) GPUs. We perform fine-grained overlapping of data-parallel (`--overlap-grad-reduce --overlap-param-gather`), tensor-parallel (`--tp-comm-overlap`) and pipeline-parallel communication (enabled by default) with computation to improve scalability. The reported throughputs are measured for end-to-end training and include all operations including data loading, optimizer steps, communication, and even logging. Note that we did not train these models to convergence.
+
+
+
+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.
+
+
+
+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%.
+
+
+
# Setup
We strongly recommend using the latest release of [NGC's PyTorch container](https://ngc.nvidia.com/catalog/containers/nvidia:pytorch) with DGX nodes. If you can't use this for some reason, use the latest pytorch, cuda, nccl, and NVIDIA [APEX](https://github.com/NVIDIA/apex#quick-start) releases. Data preprocessing requires [NLTK](https://www.nltk.org/install.html), though this is not required for training, evaluation, or downstream tasks.
@@ -67,7 +90,7 @@ docker run --gpus all -it --rm -v /path/to/megatron:/workspace/megatron -v /path
```
## Downloading Checkpoints
-We have provided pretrained [BERT-345M](https://ngc.nvidia.com/catalog/models/nvidia:megatron_bert_345m) and [GPT-345M](https://ngc.nvidia.com/catalog/models/nvidia:megatron_lm_345m) checkpoints for use to evaluate or finetuning downstream tasks. To access these checkpoints, first [sign up](https://ngc.nvidia.com/signup) for and [setup](https://ngc.nvidia.com/setup/installers/cli) the NVIDIA GPU Cloud (NGC) Registry CLI. Further documentation for downloading models can be found in the [NGC documentation](https://docs.nvidia.com/dgx/ngc-registry-cli-user-guide/index.html#topic_6_4_1).
+We have provided pretrained [BERT-345M](https://ngc.nvidia.com/catalog/models/nvidia:megatron_bert_345m) and [GPT-345M](https://ngc.nvidia.com/catalog/models/nvidia:megatron_lm_345m) checkpoints to evaluate or for finetuning downstream tasks. To access these checkpoints, first [sign up](https://ngc.nvidia.com/signup) for and [setup](https://ngc.nvidia.com/setup/installers/cli) the NVIDIA GPU Cloud (NGC) Registry CLI. Further documentation for downloading models can be found in the [NGC documentation](https://docs.nvidia.com/dgx/ngc-registry-cli-user-guide/index.html#topic_6_4_1).
Alternatively, you can directly download the checkpoints using:
@@ -89,7 +112,7 @@ After installation, there are several possible workflows. The most comprehensive
However, steps 1 and 2 can be replaced by using one of the pretrained models mentioned above.
-We've provided several scripts for pretraining both BERT and GPT in [`examples`](./examples) directory, as well as scripts for both zero-shot and fine-tuned downstream tasks including MNLI, RACE, WikiText103, and LAMBADA evaluation. There is also a script for GPT interactive text generation.
+We've provided several scripts for pretraining both BERT and GPT in the [`examples`](./examples) directory, as well as scripts for both zero-shot and fine-tuned downstream tasks including MNLI, RACE, WikiText103, and LAMBADA evaluation. There is also a script for GPT interactive text generation.
# Training
## Data Preprocessing
@@ -101,13 +124,12 @@ The training data requires preprocessing. First, place your training data in a l
The name of the `text` field of the json can be changed by using the `--json-key` flag in [`preprocess_data.py`](./tools/preprocess_data.py) The other metadata are optional and are not used in training.
-The loose json is then processed into a binary format for training. To convert the json into mmap, cached index file, or the lazy loader format use `preprocess_data.py`. Set the `--dataset-impl` flag to `mmap`, `cached`, or `lazy`, respectively (default is `mmap`). An example script to prepare data for BERT training is:
+The loose json is then processed into a binary format for training. To convert the json into mmap format use `preprocess_data.py`. An example script to prepare data for BERT training is:
python tools/preprocess_data.py \
--input my-corpus.json \
--output-prefix my-bert \
- --vocab bert-vocab.txt \
- --dataset-impl mmap \
+ --vocab-file bert-vocab.txt \
--tokenizer-type BertWordPieceLowerCase \
--split-sentences
@@ -124,8 +146,7 @@ Some minor modifications are required for GPT data preprocessing, namely, the ad
python tools/preprocess_data.py \
--input my-corpus.json \
--output-prefix my-gpt2 \
- --vocab gpt2-vocab.json \
- --dataset-impl mmap \
+ --vocab-file gpt2-vocab.json \
--tokenizer-type GPT2BPETokenizer \
--merge-file gpt2-merges.txt \
--append-eod
@@ -138,27 +159,28 @@ Further command line arguments are described in the source file [`preprocess_dat
## BERT Pretraining
-The [`examples/pretrain_bert.sh`](./examples/pretrain_bert.sh) script runs single GPU 345M parameter BERT pretraining. Debugging is the primary use for single GPU training, as the code base and command line arguments are optimized for highly distributed training. Most of the arguments are fairly self-explanatory. By default, the learning rate decays linearly over the training iterations starting at `--lr` to a minimum set by `--min-lr` over `--lr-decay-iters` iterations. The fraction of training iterations used for warmup is set by `--lr-warmup-fraction`. While this is single GPU training, the batch size specified by `--micro-batch-size` is a single forward-backward path batch-size and the code will perform gradient accumulation steps until it reaches `global-batch-size` which is the batch size per iteration. The data is partitioned into a 949:50:1 ratio for training/validation/test sets (default is 969:30:1). This partitioning happens on the fly, but is consistent across runs with the same random seed (1234 by default, or specified manually with `--seed`). We use `train-iters` as the training iterations requested. Alternatively, one can provide `--train-samples` which is total number of samples to train on. If this option is present, then instead of providing `--lr-decay-iters`, one will need to provide `--lr-decay-samples`.
+The [`examples/bert/train_bert_340m_distributed.sh`](examples/bert/train_bert_340m_distributed.sh) script runs single GPU 345M parameter BERT pretraining. Debugging is the primary use for single GPU training, as the code base and command line arguments are optimized for highly distributed training. Most of the arguments are fairly self-explanatory. By default, the learning rate decays linearly over the training iterations starting at `--lr` to a minimum set by `--min-lr` over `--lr-decay-iters` iterations. The fraction of training iterations used for warmup is set by `--lr-warmup-fraction`. While this is single GPU training, the batch size specified by `--micro-batch-size` is a single forward-backward path batch-size and the code will perform gradient accumulation steps until it reaches `global-batch-size` which is the batch size per iteration. The data is partitioned into a 949:50:1 ratio for training/validation/test sets (default is 969:30:1). This partitioning happens on the fly, but is consistent across runs with the same random seed (1234 by default, or specified manually with `--seed`). We use `train-iters` as the training iterations requested. Alternatively, one can provide `--train-samples` which is total number of samples to train on. If this option is present, then instead of providing `--lr-decay-iters`, one will need to provide `--lr-decay-samples`.
-The logging, checkpoint-saving, and evaluation intervals are specified. Checkpointing the activations facilitates the training of larger models and/or batches. Note that the `--data-path` now includes the additional `_text_sentence` suffix added in preprocessing, but does not include the file extensions.
+The logging, checkpoint-saving, and evaluation interval options are specified. Note that the `--data-path` now includes the additional `_text_sentence` suffix added in preprocessing, but does not include the file extensions.
-Further command line arguments are described in the source file [`arguments.py`](./megatron/arguments.py).
+Further command line arguments are described in the source file [`arguments.py`](./megatron/training/arguments.py).
-To run `examples/pretrain_bert.sh`, make any desired modifications including setting the environment variables for `CHECKPOINT_PATH`, `VOCAB_FILE`, and `DATA_PATH`. Make sure to set these variables to their paths in the container. Then launch the container with Megatron and necessary paths mounted (as explained in [Setup](#setup)) and run the example script.
+To run `train_bert_340m_distributed.sh`, make any desired modifications including setting the environment variables for `CHECKPOINT_PATH`, `VOCAB_FILE`, and `DATA_PATH`. Make sure to set these variables to their paths in the container. Then launch the container with Megatron and necessary paths mounted (as explained in [Setup](#setup)) and run the example script.
## GPT Pretraining
-The `examples/pretrain_gpt.sh` script runs single GPU 345M parameter GPT pretraining. As mentioned above, single GPU training is primarily intended for debugging purposes, as the code is optimized for distributed training.
+The `examples/gpt3/train_gpt3_175b_distributed.sh` script runs single GPU 345M parameter GPT pretraining. As mentioned above, single GPU training is primarily intended for debugging purposes, as the code is optimized for distributed training.
It follows largely the same format as the previous BERT script with a few notable differences: the tokenization scheme used is BPE (which requires a merge table and a `json` vocabulary file) instead of WordPiece, the model architecture allows for longer sequences (note that the max position embedding must be greater than or equal to the maximum sequence length), and the `--lr-decay-style` has been set to cosine decay. Note that the `--data-path` now includes the additional `_text_document` suffix added in preprocessing, but does not include the file extensions.
-Further command line arguments are described in the source file [`arguments.py`](./megatron/arguments.py).
+Further command line arguments are described in the source file [`arguments.py`](./megatron/training/arguments.py).
-`examples/pretrain_gpt.sh` can be launched the same way as described for BERT. Set the env vars and make any other modifications, launch the container with appropriate mounts, and run the script.
+`train_gpt3_175b_distributed.sh` can be launched the same way as described for BERT. Set the env vars and make any other modifications, launch the container with appropriate mounts, and run the script.
+More details in [`examples/gpt3/README.md`](./examples/gpt3/README.md)
## T5 Pretraining
-Very similar to BERT and GPT, the `examples/pretrain_t5.sh` script runs single GPU "base" (~220M parameter) T5 pretraining. The primary difference from BERT and GPT is the addition of the following arguments to accommodate the T5 architecture:
+Very similar to BERT and GPT, the `examples/t5/train_t5_220m_distributed.sh` script runs single GPU "base" (~220M parameter) T5 pretraining. The primary difference from BERT and GPT is the addition of the following arguments to accommodate the T5 architecture:
* `--kv-channels` sets the inner dimension of the "key" and "value" matrices of all attention mechanisms in the model. For BERT and GPT this defaults to the hidden size divided by the number of attention heads, but can be configured for T5.
@@ -168,19 +190,19 @@ Very similar to BERT and GPT, the `examples/pretrain_t5.sh` script runs single G
All of the other arguments remain as they were for BERT and GPT pretraining. Run this example with the same steps described above for the other scripts.
+More details in [`examples/t5/README.md`](./examples/t5/README.md)
+
## Distributed Pretraining
-The `examples/pretrain_{bert,gpt,t5}_distributed.sh` scripts use the PyTorch distributed launcher for distributed training. As such, multi-node training can be achieved by properly setting environment variables. See the official PyTorch [documentation](https://pytorch.org/docs/stable/elastic/run.html#launcher-api) for further description of these [environment variables](https://pytorch.org/docs/stable/distributed.html#environment-variable-initialization). By default, multi-node training uses the [nccl](https://developer.nvidia.com/nccl) distributed backend. A simple set of additional arguments and the use of the PyTorch distributed module with the `torchrun` elastic launcher (equivalent to `python -m torch.distributed.run`) are the only additional requirements to adopt distributed training. See any of `examples/pretrain_{bert,gpt,t5}_distributed.sh` for more details.
+The `pretrain_{bert,gpt,t5}_distributed.sh` scripts use the PyTorch distributed launcher for distributed training. As such, multi-node training can be achieved by properly setting environment variables. See the official PyTorch [documentation](https://pytorch.org/docs/stable/elastic/run.html#launcher-api) for further description of these [environment variables](https://pytorch.org/docs/stable/distributed.html#environment-variable-initialization). By default, multi-node training uses the [nccl](https://developer.nvidia.com/nccl) distributed backend. A simple set of additional arguments and the use of the PyTorch distributed module with the `torchrun` elastic launcher (equivalent to `python -m torch.distributed.run`) are the only additional requirements to adopt distributed training. See any of `pretrain_{bert,gpt,t5}_distributed.sh` for more details.
-We use two types of parallelism: data and model parallelism. We facilitate two distributed data parallel implementations: a simple one of our own that performs gradient all-reduce at the end of back propagation step, and Torch's distributed data parallel wrapper that overlaps gradient reduction with back propagation computation. To switch between these two options use `--DDP-impl local` or `--DDP-impl torch`, respectively. As expected, Torch distributed data parallelism is more efficient at larger model sizes. For example, for the 8.3 billion parameters model running on 512 GPUs, the scaling increases from 60% to 76% when Torch's distributed data parallel is used. However, the overlapping method requires more memory and for some configurations (e.g., 2.5 billion parameters using 2-way model parallel and 1.2 billion parameters with no model parallel) can make the overall training slower as a result. We empirically found that using a smaller model in those cases improves the training time.
+We use two types of parallelism: data and model parallelism. Our data parallelism implementation is in `megatron/core/distributed`, and supports overlapping of the gradient reduction with the backward pass when the `--overlap-grad-reduce` command-line option is used.
-Second, we developed a simple and efficient two-dimensional model-parallel approach. To use tensor model parallelism (splitting execution of a single transformer module over multiple GPUs, see Section 3 of [our paper](https://arxiv.org/pdf/1909.08053.pdf)), add the `--tensor-model-parallel-size` flag to specify the number of GPUs among which to split the model, along with the arguments passed to the distributed launcher as mentioned above. To use sequence parallelism specify `--sequence-parallel`, which requires tensor model parallel as it split among the same GPUs (more details in Section 4.2.2 of [our paper](https://arxiv.org/pdf/2205.05198.pdf)).
+Second, we developed a simple and efficient two-dimensional model-parallel approach. To use the first dimension, tensor model parallelism (splitting execution of a single transformer module over multiple GPUs, see Section 3 of [our paper](https://arxiv.org/pdf/1909.08053.pdf)), add the `--tensor-model-parallel-size` flag to specify the number of GPUs among which to split the model, along with the arguments passed to the distributed launcher as mentioned above. To use the second dimension, sequence parallelism, specify `--sequence-parallel`, which also requires tensor model parallelism to be enabled because it splits across the same GPUs (more details in Section 4.2.2 of [our paper](https://arxiv.org/pdf/2205.05198.pdf)).
To use pipeline model parallelism (sharding the transformer modules into stages with an equal number of transformer modules on each stage, and then pipelining execution by breaking the batch into smaller microbatches, see Section 2.2 of [our paper](https://arxiv.org/pdf/2104.04473.pdf)), use the `--pipeline-model-parallel-size` flag to specify the number of stages to split the model into (e.g., splitting a model with 24 transformer layers across 4 stages would mean each stage gets 6 transformer layers each).
-
-
-We have examples of how to use these two different forms of model parallelism the example scripts ending in `distributed_with_mp.sh`:
+We have examples of how to use these two different forms of model parallelism the example scripts ending in `distributed_with_mp.sh`.
Other than these minor changes, the distributed training is identical to the training on a single GPU.
@@ -188,13 +210,15 @@ The interleaved pipelining schedule (more details in Section 2.2.2 of [our paper
## Activation Checkpointing and Recomputation
-To reduce GPU memory usage so deploy a large model to a training system, we support activation checkpointing and recomputation. We support two levels of recompute granularity: `selective` and `full`. Selective recomputation is the default and recommended in almost all cases. It saves the activations that take less space and are expensive to recompute and recomputes activations that take a lot of space but are relatively cheap to recompute (see [our paper](https://arxiv.org/pdf/2205.05198) for details). To enable selective activation recompute simply use `--recompute-activations`.
+To reduce GPU memory usage when training a large model, we support various forms of activation checkpointing and recomputation. Instead of all activations being stored in memory to be used during backprop, as was traditionally the case in deep learning models, only activations at certain "checkpoints" in the model are retained (or stored) in memory, and the other activations are recomputed on-the-fly when needed for backprop. Note that this kind of checkpointing, *activation* checkpointing, is very different from the checkpointing of model parameters and optimizer state, which is mentioned elsewhere.
+
+We support two levels of recompute granularity: `selective` and `full`. Selective recomputation is the default and is recommended in almost all cases. This mode retains in memory the activations that take less memory storage space and are more expensive to recompute and recomputes the activations that take more memory storage space but are relatively inexpensive to recompute. See [our paper](https://arxiv.org/pdf/2205.05198) for details. You should find that this mode maximizes performance while minimizing the memory required to store activations. To enable selective activation recompute simply use `--recompute-activations`.
-For cases where memory is very tight, `full` checkpointing saves just the inputs to a transformer layer, or a block of transformer layers, and recomputes everything else. To turn on full activation recompute use `--recompute-granularity full`. When using full activation recomputation, there are two methods: `uniform` and `block`, chosen using the `--recompute-method` argument.
+For cases where memory is very limited, `full` recompute saves just the inputs to a transformer layer, or a group, or block, of transformer layers, and recomputes everything else. To enable full activation recompute use `--recompute-granularity full`. When using `full` activation recompute, there are two methods: `uniform` and `block`, chosen using the `--recompute-method` argument.
-* Uniform method uniformly divides the Transformer layers into groups of layers and stores the input activations of each group in the memory. The baseline group size is 1 and, in this case, the input activation of each Transformer layer is checkpointed. When the GPU memory is insufficient, increasing the number of layers per group reduces the memory usage thus enables running a bigger model. For example, when using the number of layers per group of 4, the input activation of each group of 4 Transformer layers is checkpointed.
+* The `uniform` method uniformly divides the transformer layers into groups of layers (each group of size `--recompute-num-layers`) and stores the input activations of each group in memory. The baseline group size is 1 and, in this case, the input activation of each transformer layer is stored. When the GPU memory is insufficient, increasing the number of layers per group reduces the memory usage, enabling a bigger model to be trained. For example, when `--recompute-num-layers` is set to 4, only the input activation of each group of 4 transformer layers is stored.
-* Block method checkpoints the input activations of a set number of individual Transformer layers per pipeline stage and do the rest of layers without any checkpointing. This method can be used to skip checkpointing some Transformer layers until the GPU memory is fully used, which is applicable only when there is unused GPU memory. Checkpointing fewer transformer layers avoids unnecessary activation recomputation in the backprop thus improves training performance. For example, when we specify 5 layers to checkpoint of 8 layers per pipeline stage, the input activations of only the first 5 Transformer layers are checkpointed and activation recomputation for the rest 3 layers is not needed in the backprop.
+* The `block` method recomputes the input activations of a specific number (given by `--recompute-num-layers`) of individual transformer layers per pipeline stage and stores the input activations of the remaining layers in the pipeline stage. Reducing `--recompute-num-layers` results in storing the input activations to more transformer layers, which reduces the activation recomputation required in the backprop, thus improving training performance while increasing memory usage. For example, when we specify 5 layers to recompute of 8 layers per pipeline stage, the input activations of only the first 5 transformer layers are recomputed in the backprop step while the input activations for the final 3 layers are stored. `--recompute-num-layers` can be incrementally increased until the amount of memory storage space required is just small enough to fit in the available memory, thereby both maximally utilizing memory and maximizing performance.
## Distributed Optimizer
@@ -211,6 +235,8 @@ Theoretical memory savings vary depending on the combination of the model's para
| bf16 param, fp32 grads | 18 | 6 + 12/d |
| fp32 param, fp32 grads | 16 | 8 + 8/d |
+As with regular data parallelism, overlapping of the gradient reduction (in this case, a reduce-scatter) with the backward pass can be facilitated using the `--overlap-grad-reduce` flag. Additionally, overlapping of the parameter all-gather can be overlapped with the forward pass using `--overlap-param-gather`.
+
## FlashAttention
Usage: `--use-flash-attn`. Support attention head dimensions at most 128.
@@ -226,23 +252,35 @@ pip install flash-attn
## GPT-3 Example
-In `examples/pretrain_gpt3_175B.sh` we have provided an example of how to configure Megatron to run [GPT-3](https://arxiv.org/abs/2005.14165) with 175 billion parameters on 1024 GPUs. The script is designed for [slurm](https://slurm.schedmd.com/documentation.html) with [pyxis](https://github.com/NVIDIA/pyxis) plugin but can be easily adopted to any other scheduler. It uses 8-way and 16-way tensor and pipeline parallelism, respectively. With options `global-batch-size 1536` and `rampup-batch-size 16 16 5859375`, the training will start with global batch size 16 and linearly increase the global batch size to 1536 over 5,859,375 samples with incrmeental steps 16. The training dataset can be either a single set or a multiple datasets combined with a set of weights.
+In `examples/gpt3/train_gpt3_175b_distributed.sh` we have provided an example of how to configure Megatron to train [GPT-3](https://arxiv.org/abs/2005.14165) with 175 billion parameters on 1024 GPUs. The script is designed for [slurm](https://slurm.schedmd.com/documentation.html) with [pyxis](https://github.com/NVIDIA/pyxis) plugin but can be easily adopted to any other scheduler. It uses 8-way tensor parallelism and 16-way pipeline parallelism. With options `global-batch-size 1536` and `rampup-batch-size 16 16 5859375`, the training will start with global batch size 16 and linearly increase the global batch size to 1536 over 5,859,375 samples with incremental steps 16. The training dataset can be either a single set or a multiple datasets combined with a set of weights.
With full global batch size of 1536 on 1024 A100 GPUs, each iteration takes around 32 seconds resulting in 138 teraFLOPs per GPU which is 44% of the theoretical peak FLOPs.
+## Retro and InstructRetro
+
-## Retro
+Retro [(Borgeaud et al., 2022)](https://arxiv.org/abs/2112.04426) is an autoregressive decoder-only language model (LM) pretrained with retrieval-augmentation.
+Retro features practical scalability to support large-scale pretraining from scratch by retrieving from trillions of tokens.
+Pretraining with retrieval provides a more efficient storage mechanism of factual knowledge, when compared to storing factual knowledge implicitly within the network's parameters, thus largely reducing model parameters while achieving lower perplexity than standard GPT.
+Retro also provides the flexibility to update the
+knowledge stored in LMs [(Wang et al., 2023a)](https://arxiv.org/abs/2304.06762)
+by updating the retrieval database without training LMs again.
-See:
+InstructRetro [(Wang et al., 2023b)](https://arxiv.org/abs/2310.07713) further scales up the size of Retro to 48B, featuring the largest LLM pretrained with retrieval (as of December 2023).
+The obtained foundation model, Retro 48B, largely outperforms the GPT counterpart in terms of perplexity.
+With instruction tuning on Retro, InstructRetro demonstrates significant improvement over the instruction tuned GPT on downstream tasks in the zero-shot setting. Specifically, the average improvement of InstructRetro is 7% over its GPT counterpart across 8 short-form QA tasks, and 10% over GPT across 4 challenging long-form QA tasks. We also find that one can ablate the encoder from InstructRetro architecture and directly use the InstructRetro decoder backbone as GPT, while achieving comparable results.
-- `tools/retro/README.md` for an overview.
-- `tools/retro/examples/get_preprocess_cmd.sh` for an example of common preprocessing arguments.
-- `tools/retro/examples/preprocess_data.sh` for an example of how to preprocess data.
-- `tools/retro/examples/pretrain_model.sh` for an example of how to pretrain a model.
+In this repo, we provide an end-to-end reproduction guide to implement Retro and InstructRetro, covering
+- **Retrieval database construction**, which supports billions or even trillions of tokens as a large-scale retrieval database.
+- **Pretraining with retrieval**, which supports pretraining from scratch and pretraining from a pretrained GPT model (Retro-fitting).
+- **Instruction tuning**, where we provide an open-source instruction tuning dataset and the training recipe for instruction tuning on Retro.
+- **Downstream task evaluation**, where we provide the text generation and evaluation scripts for zero-shot question answering tasks.
-Retro is a retrieval-enhanced model that is based on GPT. As described in [Improving language models by retrieving from trillions of tokens](https://arxiv.org/abs/2112.04426), Retro retrieves from a database of document chunks by performing locality search using a sample's tokens. The retrieval database can be large -- often billions or even trillions of tokens -- and provides a more efficient storage mechanism of factual knowledge, when compared to storing factual knowledge implicitly within the network's parameters.
+See [tools/retro/README.md](tools/retro/README.md) for a detailed overview.
-Using Retro requires two steps: 1) preprocessing the retrieval database and pretraining neighbors, and 2) pretraining a model using this data. Please see `tools/retro/README.md` for a detailed overview.
+## Mamba-based Language Models
+
+See [examples/mamba](./examples/mamba) for details.
(view) [num_splits, np, hn, h]
- -->(tranpose) [np, num_splits, hn, h]
- -->(view) [np * num_splits * hn, h] """
-
- intermediate_shape = \
- (num_splits, num_attention_heads_per_partition,
- hidden_size_per_attention_head) + input_shape[1:]
-
- t = t.view(*intermediate_shape)
- t = t.transpose(0, 1).contiguous()
- else:
- """[np * hn * num_splits, h]
- -->(view) [np, hn, num_splits, h]
- -->(tranpose) [np, num_splits, hn, h]
- -->(view) [np * num_splits * hn, h] """
-
- intermediate_shape = \
- (num_attention_heads_per_partition,
- hidden_size_per_attention_head, num_splits) +\
- input_shape[1:]
-
- t = t.view(*intermediate_shape)
- t = t.transpose(1, 2).contiguous()
- t = t.view(*input_shape)
-
- return t
-
-def fix_query_key_value_ordering(model, checkpoint_version):
- """Fix up query/key/value matrix ordering if checkpoint
- version is smaller than 2.0
- """
- if checkpoint_version < 2.0:
- if isinstance(model, list):
- assert len(model)==1
- model = model[0]
- for name, param in model.named_parameters():
- if name.endswith(('.query_key_value.weight', '.query_key_value.bias')):
- if checkpoint_version == 0:
- fixed_param = _transpose_first_dim(param.data, 3, True, model)
- elif checkpoint_version == 1.0:
- fixed_param = _transpose_first_dim(param.data, 3, False, model)
- else:
- print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")
- sys.exit()
- param.data.copy_(fixed_param)
- if name.endswith(('.key_value.weight', '.key_value.bias')):
- if checkpoint_version == 0:
- fixed_param = _transpose_first_dim(param.data, 2, True, model)
- elif checkpoint_version == 1.0:
- fixed_param = _transpose_first_dim(param.data, 2, False, model)
- else:
- print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")
- sys.exit()
- param.data.copy_(fixed_param)
- print_rank_0(" succesfully fixed query-key-values ordering for"
- " checkpoint version {}".format(checkpoint_version))
-
-
-def _load_base_checkpoint(load_dir, rank0=False):
- """ Load the base state_dict from the given directory
-
- If rank0 is true, just loads rank 0 checkpoint, ignoring arguments.
- """
-
- # Read the tracker file and set the iteration.
- tracker_filename = get_checkpoint_tracker_filename(load_dir)
-
- # If no tracker file, return nothing
- if not os.path.isfile(tracker_filename):
- if not rank0:
- print_rank_0('WARNING: could not find the metadata file {} '.format(
- tracker_filename))
- print_rank_0(' will not load any checkpoints and will start from '
- 'random')
- return None, False
-
- # Otherwise, read the tracker file and either set the iteration or
- # mark it as a release checkpoint.
- iteration, release = read_metadata(tracker_filename)
-
- # Checkpoint.
- if rank0:
- checkpoint_name = find_checkpoint_rank_0(load_dir, iteration, release)
- else:
- checkpoint_name = get_checkpoint_name(load_dir, iteration, release)
- if release:
- print_rank_0(f' loading release checkpoint from {load_dir}')
- else:
- print_rank_0(f' loading checkpoint from {load_dir} at iteration {iteration}')
-
- # Load the checkpoint.
- try:
- state_dict = torch.load(checkpoint_name, map_location='cpu')
- except ModuleNotFoundError:
- from megatron.fp16_deprecated import loss_scaler
- # For backward compatibility.
- if not rank0:
- print_rank_0(' > deserializing using the old code structure ...')
- sys.modules['fp16.loss_scaler'] = sys.modules[
- 'megatron.fp16_deprecated.loss_scaler']
- sys.modules['megatron.fp16.loss_scaler'] = sys.modules[
- 'megatron.fp16_deprecated.loss_scaler']
- state_dict = torch.load(checkpoint_name, map_location='cpu')
- sys.modules.pop('fp16.loss_scaler', None)
- sys.modules.pop('megatron.fp16.loss_scaler', None)
- except BaseException as e:
- print_rank_0('could not load the checkpoint')
- print_rank_0(e)
- sys.exit()
-
- return state_dict, release
-
-
-def load_args_from_checkpoint(args, load_arg='load'):
- """Set required arguments from the checkpoint specified in the
- arguments.
-
- Will overwrite arguments that have a non-None default value, but
- will leave any arguments that default to None as set.
-
- Returns the same args NameSpace with the new values added/updated.
-
- If no checkpoint is specified in args, or if the checkpoint is
- there but invalid, the arguments will not be modified
-
- """
- load_dir = getattr(args, load_arg)
-
- if load_dir is None:
- print_rank_0('No load directory specified, using provided arguments.')
- return args
-
- state_dict, release = _load_base_checkpoint(load_dir, rank0=True)
-
- # Args.
- if not state_dict:
- print_rank_0('Checkpoint not found to provide arguments, using provided arguments.')
- return args
-
- if 'args' not in state_dict:
- print_rank_0('Checkpoint provided does not have arguments saved, using provided arguments.')
- return args
-
- checkpoint_args = state_dict['args']
- checkpoint_version = state_dict.get('checkpoint_version', 0)
- args.iteration = state_dict['iteration']
-
- # One-off conversion for foundation models
- if hasattr(checkpoint_args, 'disable_bias_linear'):
- setattr(checkpoint_args, 'add_bias_linear', not getattr(checkpoint_args, 'disable_bias_linear'))
-
- def _set_arg(arg_name, old_arg_name=None, force=False):
- if not force and getattr(args, arg_name, None) is not None:
- return
-
- if old_arg_name is not None:
- checkpoint_value = getattr(checkpoint_args, old_arg_name, None)
- else:
- checkpoint_value = getattr(checkpoint_args, arg_name, None)
-
- if checkpoint_value is not None:
- print_rank_0(f"Setting {arg_name} to {checkpoint_value} from checkpoint")
- setattr(args, arg_name, checkpoint_value)
- else:
- print_rank_0(f"Checkpoint did not provide arguments {arg_name}")
-
- _set_arg('num_layers')
- _set_arg('hidden_size')
- _set_arg('ffn_hidden_size')
- _set_arg('seq_length')
- _set_arg('num_attention_heads')
- _set_arg('kv_channels')
- _set_arg('max_position_embeddings')
- _set_arg('add_position_embedding', force=True)
- _set_arg('use_rotary_position_embeddings', force=True)
- _set_arg('rotary_percent', force=True)
- _set_arg('add_bias_linear', force=True)
- _set_arg('swiglu', force=True)
- _set_arg('untie_embeddings_and_output_weights', force=True)
- _set_arg('apply_layernorm_1p', force=True)
- _set_arg('tokenizer_type')
- _set_arg('padded_vocab_size')
- if checkpoint_version < 3.0:
- _set_arg('tensor_model_parallel_size',
- 'model_parallel_size')
- else:
- _set_arg('tensor_model_parallel_size', force=True)
- _set_arg('pipeline_model_parallel_size', force=True)
- _set_arg('virtual_pipeline_model_parallel_size', force=True)
- _set_arg('num_layers_per_virtual_pipeline_stage')
- return args, checkpoint_args
-
-
-def load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', strict=True):
- """Load a model checkpoint and return the iteration.
- strict (bool): whether to strictly enforce that the keys in
- :attr:`state_dict` of the checkpoint match the names of
- parameters and buffers in model.
- """
- args = get_args()
- load_dir = getattr(args, load_arg)
-
- model = unwrap_model(model)
-
- state_dict, release = _load_base_checkpoint(load_dir, rank0=False)
-
- # Checkpoint not loaded.
- if state_dict is None:
-
- # Conditionally exit at this point.
- if args.exit_on_missing_checkpoint:
- print_rank_0(">> '--exit-on-missing-checkpoint' set ... exiting. <<")
- torch.distributed.barrier()
- sys.exit()
-
- # Iteration defaults to 0.
- return 0
-
- # Set checkpoint version.
- set_checkpoint_version(state_dict.get('checkpoint_version', 0))
-
- # Set iteration.
- if args.finetune or release:
- iteration = 0
- else:
- try:
- iteration = state_dict['iteration']
- except KeyError:
- try: # Backward compatible with older checkpoints
- iteration = state_dict['total_iters']
- except KeyError:
- print_rank_0('A metadata file exists but unable to load '
- 'iteration from checkpoint {}, exiting'.format(
- checkpoint_name))
- sys.exit()
-
- # Check arguments.
- assert args.consumed_train_samples == 0
- assert args.consumed_valid_samples == 0
- if 'args' in state_dict and not args.finetune:
- checkpoint_args = state_dict['args']
- check_checkpoint_args(checkpoint_args)
- args.consumed_train_samples = getattr(checkpoint_args,
- 'consumed_train_samples', 0)
- update_num_microbatches(consumed_samples=args.consumed_train_samples)
- args.consumed_valid_samples = getattr(checkpoint_args,
- 'consumed_valid_samples', 0)
- else:
- print_rank_0('could not find arguments in the checkpoint ...')
-
- # Model.
- if len(model) == 1:
- model[0].load_state_dict(state_dict['model'], strict=strict)
- else:
- for i in range(len(model)):
- mpu.set_virtual_pipeline_model_parallel_rank(i)
- model[i].load_state_dict(state_dict['model%d' % i], strict=strict)
-
- # Fix up query/key/value matrix ordering if needed.
- checkpoint_version = get_checkpoint_version()
- print_rank_0(f' checkpoint version {checkpoint_version}')
- fix_query_key_value_ordering(model, checkpoint_version)
-
- # Optimizer.
- if not release and not args.finetune and not args.no_load_optim:
- try:
- # Load state dict.
- if optimizer is not None:
- optimizer.load_state_dict(state_dict['optimizer'])
-
- # Load distributed optimizer's custom parameter state.
- if args.use_distributed_optimizer:
- tracker_filename = get_checkpoint_tracker_filename(load_dir)
- iteration, release = read_metadata(tracker_filename)
- model_checkpoint_name = \
- get_checkpoint_name(load_dir, iteration, release)
- optim_checkpoint_name = \
- get_distributed_optimizer_checkpoint_name(
- model_checkpoint_name)
- optimizer.load_parameter_state(optim_checkpoint_name)
-
- # Load scheduler.
- if opt_param_scheduler is not None:
- if 'lr_scheduler' in state_dict: # backward compatbility
- opt_param_scheduler.load_state_dict(state_dict['lr_scheduler'])
- else:
- opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler'])
- except KeyError:
- print_rank_0('Unable to load optimizer from checkpoint {}. '
- 'Specify --no-load-optim or --finetune to prevent '
- 'attempting to load the optimizer state, '
- 'exiting ...'.format(checkpoint_name))
- sys.exit()
- else:
- if (args.fp16 or args.bf16) and optimizer is not None:
- optimizer.reload_model_params()
-
- # rng states.
- if not release and not args.finetune and not args.no_load_rng:
- try:
- if 'rng_state' in state_dict:
- # access rng_state for data parallel rank
- if args.data_parallel_random_init:
-
- rng_state = state_dict['rng_state'][mpu.get_data_parallel_rank()]
- else:
- rng_state = state_dict['rng_state'][0]
- random.setstate(rng_state['random_rng_state'])
- np.random.set_state(rng_state['np_rng_state'])
- torch.set_rng_state(rng_state['torch_rng_state'])
- torch.cuda.set_rng_state(rng_state['cuda_rng_state'])
- # Check for empty states array
- if not rng_state['rng_tracker_states']:
- raise KeyError
- tensor_parallel.get_cuda_rng_tracker().set_states(
- rng_state['rng_tracker_states'])
- else: # backward compatability
- random.setstate(state_dict['random_rng_state'])
- np.random.set_state(state_dict['np_rng_state'])
- torch.set_rng_state(state_dict['torch_rng_state'])
- torch.cuda.set_rng_state(state_dict['cuda_rng_state'])
- # Check for empty states array
- if not state_dict['rng_tracker_states']:
- raise KeyError
- tensor_parallel.get_cuda_rng_tracker().set_states(
- state_dict['rng_tracker_states'])
- except KeyError:
- print_rank_0('Unable to load rng state from checkpoint {}. '
- 'Specify --no-load-rng or --finetune to prevent '
- 'attempting to load the rng state, '
- 'exiting ...'.format(checkpoint_name))
- sys.exit()
-
- # Some utilities want to load a checkpoint without distributed being initialized
- if torch.distributed.is_initialized():
- torch.distributed.barrier()
-
- print_rank_0(f' successfully loaded checkpoint from {args.load} '
- f'at iteration {iteration}')
-
- return iteration
-
-
-def load_biencoder_checkpoint(model, only_query_model=False,
- only_context_model=False, custom_load_path=None):
- """
- selectively load retrieval models for indexing/retrieving
- from saved checkpoints
- """
-
- args = get_args()
-
- model = unwrap_model(model)
-
- load_path = custom_load_path if custom_load_path is not None else args.load
-
- tracker_filename = get_checkpoint_tracker_filename(load_path)
- with open(tracker_filename, 'r') as f:
- iteration = int(f.read().strip())
-
- checkpoint_name = get_checkpoint_name(load_path, iteration,
- args.use_distributed_optimizer,
- release=False)
-
- if mpu.get_data_parallel_rank() == 0:
- print('global rank {} is loading checkpoint {}'.format(
- torch.distributed.get_rank(), checkpoint_name))
-
- state_dict = torch.load(model_checkpoint_name, map_location='cpu')
- ret_state_dict = state_dict['model']
-
- if only_query_model:
- ret_state_dict.pop('context_model')
- if only_context_model:
- ret_state_dict.pop('query_model')
-
- assert len(model) == 1
- model[0].load_state_dict(ret_state_dict)
- torch.distributed.barrier()
-
- if mpu.get_data_parallel_rank() == 0:
- print(' successfully loaded {}'.format(checkpoint_name))
-
- return model
diff --git a/megatron/core/QuickStart.md b/megatron/core/QuickStart.md
new file mode 100644
index 0000000000..6deb1a5f76
--- /dev/null
+++ b/megatron/core/QuickStart.md
@@ -0,0 +1,250 @@
+## Quick Start
+
+The following guide is a short getting started guide for Megatron Core. In it you:
+
+* Initialize Megatron Core on 2 GPUS.
+* Build a GPT model with tensor model parallel size 2, pipeline parallel size 1
+* Train it for a five iterations using Megatron Core schedules
+* Save the model using the distributed checkpointing format
+* Load the model saved above.
+
+**NOTE:** The following sample was tested using Megatron Core version 0.8.0 and NGC PyTorch Container version 24.02.
+
+### Environment Setup
+
+```
+docker run --ipc=host --shm-size=512m --gpus 2 -it nvcr.io/nvidia/pytorch:24.02-py3
+
+git clone https://github.com/NVIDIA/Megatron-LM.git && cd Megatron-LM
+```
+
+
+### Writing Your First Training Loop
+
+In the following steps you create a sample GPT model split across tensors (Tensor model parallel) on 2 GPUS, and run a forward pass through it using a MockGPT dataset helper class that we created in Megatron Core.
+
+
+
+**NOTE:** All of the following steps are in the [run_simple_mcore_train_loop.py](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/run_simple_mcore_train_loop.py) script.
+
+To run the ``run_simple_mcore_train_loop.py`` script:
+
+```
+PYTHONPATH=$PYTHON_PATH:./megatron torchrun --nproc-per-node 2 examples/run_simple_mcore_train_loop.py
+```
+
+
+
+**STEP 1 - Initialize Distributed Training and Model Parallel Setup**
+
+The following utility, when called, initializes your distributed setup.
+
+```python
+import os
+import torch
+from megatron.core import parallel_state
+
+def initialize_distributed(tensor_model_parallel_size = 1, pipeline_model_parallel_size = 1):
+ # Torch setup for distributed training
+ rank = int(os.environ['LOCAL_RANK'])
+ world_size = torch.cuda.device_count()
+ torch.cuda.set_device(rank)
+ torch.distributed.init_process_group(world_size=world_size, rank=rank)
+
+ # Megatron core distributed training initialization
+ parallel_state.initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size)
+```
+
+
+**STEP 2 - GPT Model Setup**
+
+In this step, you create a GPT model. For a list of other configurations that you can pass into the model open and review [transformer_config.py](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/transformer/transformer_config.py).
+
+```
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
+
+def model_provider():
+ """Build the model."""
+
+ transformer_config = TransformerConfig(
+ num_layers=2,
+ hidden_size=12,
+ num_attention_heads=4,
+ use_cpu_initialization=True,
+ pipeline_dtype=torch.float32)
+
+ gpt_model = GPTModel(
+ config=transformer_config,
+ transformer_layer_spec=get_gpt_layer_local_spec(),
+ vocab_size=100,
+ max_sequence_length=64)
+
+ return gpt_model
+```
+
+
+**STEP 3 - GPT Mock Dataset Setup**
+
+In the following step, you explore the mock dataset utility.
+
+* To train the model using your data, use the GPTDataset class in [gpt_dataset.py](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/datasets/gpt_dataset.py).
+
+* To find more information about Megatron Core data pipeline, see the [data pipeline readme.md](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/datasets/readme.md?ref_type=heads).
+
+```
+import torch
+from torch.utils.data import DataLoader
+
+from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
+from megatron.core.datasets.gpt_dataset import GPTDatasetConfig, MockGPTDataset
+from megatron.training.tokenizer.tokenizer import _NullTokenizer
+from megatron.core.datasets.utils import compile_helpers
+
+_SEQUENCE_LENGTH = 64
+
+def get_train_data_iterator():
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ if torch.distributed.get_rank() == 0:
+ compile_helpers()
+ torch.distributed.barrier()
+ else:
+ compile_helpers()
+
+ config = GPTDatasetConfig(
+ random_seed=0,
+ sequence_length=_SEQUENCE_LENGTH,
+ reset_position_ids=False,
+ reset_attention_mask=False,
+ eod_mask_loss=False,
+ tokenizer=_NullTokenizer(vocab_size=_SEQUENCE_LENGTH),
+ )
+
+ datasets = BlendedMegatronDatasetBuilder(
+ MockGPTDataset, [1000, None, None], lambda: True, config
+ ).build()
+
+ train_dataloader = DataLoader(datasets[0], batch_size=8, shuffle=True)
+
+ train_iterator = iter(train_dataloader)
+
+ return train_iterator
+
+```
+
+
+**STEP 4 - Forward Step Function**
+
+Megatron Core uses [schedules.py](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/pipeline_parallel/schedules.py) to run the model. It is sufficient to define a forward step function, which takes as input the data iterator and the model and produces as output the output tensor and a loss function.
+
+```python
+from functools import partial
+
+def forward_step_func(data_iterator, model):
+
+ def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor):
+
+ losses = output_tensor.float()
+ loss_mask = loss_mask.view(-1).float()
+ loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
+ # If you have data parallel reduce loss across data parallel groups.
+ # If pipeline parallel, loss computation is done only in last stage.
+
+ return loss, {'lm loss': loss}
+
+ data = next(data_iterator)
+ tokens = data['tokens'].to(device)
+ attention_mask = data['attention_mask'].to(device)
+ position_ids = data['position_ids'].to(device)
+ labels = data['labels'].to(device)
+ loss_mask = data['loss_mask'].to(device)
+
+ output_tensor = model(tokens, position_ids, attention_mask,
+ labels=labels)
+
+ return output_tensor, partial(loss_func, loss_mask)
+```
+
+
+**STEP 5 - Load and Save Distributed Checkpoint**
+
+Megatron Core uses distributed checkpoints for loading and saving models. This gives you the flexibility to convert the model from one model parallel setting to another when you load a model. For example, a model trained with tensor parallel size 2, can be loaded again as tensor model parallel size 4, and so forth.
+
+```python
+from megatron.core import dist_checkpointing
+
+def save_distributed_checkpoint(checkpoint_path, gpt_model):
+ sharded_state_dict = gpt_model.sharded_state_dict(prefix='')
+ dist_checkpointing.save(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)
+
+def load_distributed_checkpoint(checkpoint_path, gpt_model):
+ sharded_state_dict=gpt_model.sharded_state_dict(prefix='')
+ checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)
+ gpt_model.load_state_dict(checkpoint)
+ return gpt_model
+```
+
+
+**STEP 6 - Main Function**
+
+The following code snippet is the main function that needs to go into your script. It runs the model for 5 iterations, saves the model, and loads the data model.
+
+```python
+from pathlib import Path
+from torch.optim import Adam
+from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+
+if __name__ == "__main__":
+ initialize_distributed(tensor_model_parallel_size=2, pipeline_model_parallel_size=1)
+ model_parallel_cuda_manual_seed(123)
+
+ gpt_model = model_provider()
+ device = torch.device("cuda")
+ gpt_model.to(device)
+
+ optim = Adam(gpt_model.parameters())
+
+ train_iterator = get_train_data_iterator()
+
+ forward_backward_func = get_forward_backward_func()
+
+ # Running the model for 5 iterations
+ for _ in range(5):
+ optim.zero_grad()
+
+ losses_reduced = forward_backward_func(
+ forward_step_func=forward_step_func,
+ data_iterator=train_iterator,
+ model=gpt_model,
+ num_microbatches=1,
+ seq_length=64,
+ micro_batch_size=8,
+ decoder_seq_length=64,
+ forward_only=False)
+
+ optim.step()
+
+ print(f'Losses reduced : {losses_reduced}')
+
+ # Saving the model
+ save_distributed_checkpoint(gpt_model=gpt_model, checkpoint_path='/workspace/ckpt')
+
+ # Loading the model
+ gpt_model = load_distributed_checkpoint(gpt_model=gpt_model, checkpoint_path='/workspace/ckpt')
+ gpt_model.to(device)
+ print('Successfully loaded the model')
+```
+
+
+
+
+### Extending Further
+
+The example you explored here is a basic training loop in Megatron Core. To review more advanced examples, explore [pretrain_gpt.py]. ``pretrain_gpt.py`` has more complex training loops that includes the following and other Megatron Core features:
+
+* pipeline parallel
+* context parallel
+* rope embeddings
+* mixture of experts
diff --git a/megatron/core/README.md b/megatron/core/README.md
index 0c8c61738d..38970b0c47 100644
--- a/megatron/core/README.md
+++ b/megatron/core/README.md
@@ -1 +1,14 @@
-Megatron Core is a library for efficient and scalable training of transformer based models.
+# Megatron-Core
+
+Megatron-Core is an open-source PyTorch-based library that contains GPU-optimized techniques and cutting-edge system-level optimizations. It abstracts them into composable and modular APIs, allowing full flexibility for developers and model researchers to train custom transformers at-scale on NVIDIA accelerated computing infrastructure. This library is compatible with all NVIDIA Tensor Core GPUs, including FP8 acceleration support for [NVIDIA Hopper architectures](https://www.nvidia.com/en-us/data-center/technologies/hopper-architecture/).
+
+Megatron-Core offers core building blocks such as attention mechanisms, transformer blocks and layers, normalization layers, and embedding techniques. Additional functionality like activation re-computation, distributed checkpointing is also natively built-in to the library. The building blocks and functionality are all GPU optimized, and can be built with advanced parallelization strategies for optimal training speed and stability on NVIDIA Accelerated Computing Infrastructure. Another key component of the Megatron-Core library includes advanced model parallelism techniques (tensor, sequence, pipeline, context, and MoE expert parallelism).
+
+Megatron-Core can be used with [NVIDIA NeMo](https://www.nvidia.com/en-us/ai-data-science/products/nemo/), an enterprise-grade AI platform. Alternatively, you can explore Megatron-Core with the native PyTorch training loop [here](https://github.com/NVIDIA/Megatron-LM/tree/main/examples). Visit [Megatron-Core documentation](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html) to learn more.
+
+## Quick links
+
+- [Benchmark using NVIDIA NeMo](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html#performance-benchmarks)
+- [Multimodal example (LLaVA training pipeline)](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/multimodal)
+- [Mixture-of-Experts](https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/transformer/moe)
+- [Training Mamba-based Language Models](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/mamba)
diff --git a/megatron/core/README_STRAGGLER.md b/megatron/core/README_STRAGGLER.md
new file mode 100644
index 0000000000..fe9062c851
--- /dev/null
+++ b/megatron/core/README_STRAGGLER.md
@@ -0,0 +1,93 @@
+## StragglerDetector for a TP Group
+
+The file `megatron/core/utils.py` has a class named `StragglerDetector` which supports Python Contexts.
+It can be used to find straggling TP group based on the RTT of the ranks in the TP Group. It also collects
+Power/Temp/Utilization for GPUs, which can additionally be used to narrow down to the exact GPU in the TP Group,
+assuming the straggling was caused by hardware anomaly in a given GPU.
+This class supports collecting timing events for various steps of a given iteration. It
+keeps collecting such timing events on a per rank basis, and when the reporter is invoked
+during a logging interval, it computes the min and max of certain metric across all
+ranks and logs the observed metric and the rank as follows
+
+```
+ 0: INFO:megatron.core.utils:[2024-03-14 23:07:56] | MnRtt/Rnk: 3453.08ms/8 | MxRtt/Rnk: 3468.20ms/0 | MnPwr/Rnk: 601796W/8 | MxPwr/Rnk: 683801W/18 | MnTmp/Rnk: 52C/0 | MxTmp/Rnk: 65C/21 | MnUtl/Rnk: 97%/8 | MxUtl/Rnk: 100%/6 | MnClk/Rnk: 1950MHz/28 | MxClk/Rnk: 1980MHz/0 | MnDRtt/Rnk: 14.27ms/23 | MxDRtt/Rnk: 34.65ms/3 | MnEtpt/Rnk: 296.02TF/0 | MxEtpt/Rnk: 297.32TF/8
+```
+
+
+### Description of the metrics
+
+Each metric is prefixed with `Mn` or `Mx` to represent `Minimum` or `Maximum`. Each metric is also suffixed with the rank where the metric was measured. The metrics are averaged over the logging interval. Between the prefix and the rank is the name of the metric as follows
+
+- Rtt : RoundTrip Time (time spent in all the traced ops per iteration)
+- Pwr : GPU Power
+- Tmp : GPU Temperature
+- Utl : GPU Utilization
+- Clk : GPU Clock
+- DRtt: get_batch latency
+- Etpt: Estimated throughput. This is derived from actual computed throughput dividied by Rtt. Since we do not collect timing for backward pass, the value is further divided by three to come up with estimated throughput.
+
+
+### Command Line activation
+To start using the StragglerDetector, need to pass the following argument `--log-straggler`. It optionally also takes two additional parameters. Default disabled
+- `--disable-straggler-on-startup` - whether to keept the StragglerDetector disabled on startup and enable later. Default enabled
+- `--straggler-ctrlr-port` - The StragglerDetector can toggle between on/off just by sending `curl Rank0Host:port`. Default port is 65535. Every time it is turned
+- `--straggler-minmax-count` - If set to > 1 (N), it prints N Top and Bottom Etpt/Rank pairs as shown below
+```
+ 0: INFO:megatron.core.utils:^^^^ Bottom 4 Ranks with lowest Etpt(TF): 296.02/0, 296.17/2, 296.23/1, 296.23/4,
+ 0: INFO:megatron.core.utils:^^^^ Top 4 Ranks with highest Etpt(TF): 297.28/15, 297.28/11, 297.32/12, 297.32/8,
+```
+
+
+### Programming the StragglerDetector
+The StragglerDetector class supports context, and its implementation is a Singleton.
+- Initialization
+
+```
+ # initialization, where StragglerDetector will be used
+ from megatron.core.utils import StragglerDetector
+ stimer = StragglerDetector()
+```
+
+- One time for each rank
+
+```
+ # one time before the training loop starts
+ stimer.configure(world, rank, enabled=True, port=65545)
+
+ # Arguments to configure
+ # world : World Size
+ # rank : The rank of this trainer
+ # mmcnt : (Optional) Number of ranks to print for showing Min/Max Etpt
+ # amp : (Optional) Set to 3.0 if we only use timers in fwd pass
+ # port : (Optional) control port, useful only for rank-0
+ # prefill : (Optional) howmany Events to pre-populate
+ # enabled : (Optional) whether or not collection is enabled on startup
+```
+
+- To Capture time
+
+```
+ # whereever timing need to be captured
+ with stimer:
+ do_operation()
+
+ # special case for get_batch
+ with stimer(bdata=True):
+ input,... = get_batch(iterator,...)
+```
+
+- Logging in main training loop
+
+```
+ # logging
+ total_flops = 0.0
+ iteration = 0
+ # inside the main training loop
+ while training:
+ iteration += 1
+ do_step()
+ total_flops += get_computed_flops()
+ if iteration % log_interval:
+ stimer.report(total_flops, log_interval)
+ total_flops = 0.0
+```
diff --git a/megatron/core/__init__.py b/megatron/core/__init__.py
index cb437d5dae..0eccb1d02e 100644
--- a/megatron/core/__init__.py
+++ b/megatron/core/__init__.py
@@ -1,6 +1,24 @@
-import megatron.core.parallel_state
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
import megatron.core.tensor_parallel
import megatron.core.utils
+from megatron.core import parallel_state
+from megatron.core.distributed import DistributedDataParallel
+from megatron.core.inference_params import InferenceParams
+from megatron.core.model_parallel_config import ModelParallelConfig
+from megatron.core.package_info import (
+ __contact_emails__,
+ __contact_names__,
+ __description__,
+ __download_url__,
+ __homepage__,
+ __keywords__,
+ __license__,
+ __package_name__,
+ __repository_url__,
+ __shortversion__,
+ __version__,
+)
+from megatron.core.timers import Timers
# Alias parallel_state as mpu, its legacy name
mpu = parallel_state
@@ -9,4 +27,8 @@
"parallel_state",
"tensor_parallel",
"utils",
+ "DistributedDataParallel",
+ "InferenceParams",
+ "ModelParallelConfig",
+ "Timers",
]
diff --git a/megatron/core/config_logger.py b/megatron/core/config_logger.py
new file mode 100644
index 0000000000..231a0226be
--- /dev/null
+++ b/megatron/core/config_logger.py
@@ -0,0 +1,104 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import dataclasses
+import json
+import os
+
+import torch
+import torch.nn as nn
+
+from megatron.core import parallel_state
+
+
+def get_config_logger_path(config):
+ return getattr(config, 'config_logger_dir', '')
+
+
+def has_config_logger_enabled(config):
+ return get_config_logger_path(config) != ''
+
+
+# For each prefix, holds a counter and increases it every time we dump with this
+# prefix.
+__config_logger_path_counts = {}
+
+
+def get_path_count(path):
+ """
+ keeps tracks of number of times we've seen the input `path` and return count-1
+ """
+ global __config_logger_path_counts
+ if not path in __config_logger_path_counts:
+ __config_logger_path_counts[path] = 0
+ count = __config_logger_path_counts[path]
+ __config_logger_path_counts[path] += 1
+ return count
+
+
+def get_path_with_count(path):
+ """
+ calls get_path_count and appends returned value to path
+ """
+ return f'{path}.iter{get_path_count(path)}'
+
+
+class JSONEncoderWithMcoreTypes(json.JSONEncoder):
+ def default(self, o):
+ if type(o).__name__ in ['function', 'ProcessGroup']:
+ return str(o)
+ if type(o).__name__ in ['dict', 'OrderedDict']:
+ return {k: self.default(v) for k, v in o.items()}
+ if type(o).__name__ in ['list', 'ModuleList']:
+ return [self.default(val) for val in o]
+ if type(o).__name__ == 'UniqueDescriptor':
+ return {
+ attr: self.default(getattr(o, attr))
+ for attr in filter(lambda x: not x.startswith('__'), dir(o))
+ }
+ if type(o) is torch.dtype:
+ return str(o)
+ # if it's a Float16Module, add "Float16Module" to the output dict
+ if type(o).__name__ == 'Float16Module':
+ return {'Float16Module': {'module': self.default(o.module)}}
+ # If it's a nn.Module subchild, either print its children or itself if leaf.
+ if issubclass(type(o), nn.Module):
+ if len(getattr(o, '_modules', {})) > 0:
+ return {key: self.default(val) for key, val in o._modules.items()}
+ else:
+ return str(o)
+ if type(o).__name__ in ['ABCMeta', 'type', 'AttnMaskType']:
+ return str(o)
+ if dataclasses.is_dataclass(o) or type(o).__name__ in ['ModuleSpec', 'TransformerConfig']:
+ return dataclasses.asdict(o)
+ try:
+ return super().default(o)
+ except:
+ return str(o)
+
+
+def log_config_to_disk(config, dict_data, prefix=''):
+ """
+ Encodes the input dict (dict_data) using the JSONEncoderWithMcoreTypes
+ and dumps to disk, as specified via path
+ """
+ path = get_config_logger_path(config)
+ assert path is not None, 'Expected config_logger_dir to be non-empty in config.'
+
+ if 'self' in dict_data:
+ if prefix == '':
+ prefix = type(dict_data['self']).__name__
+ del dict_data['self']
+
+ if not os.path.exists(path):
+ os.makedirs(path, exist_ok=True)
+
+ rank = parallel_state.get_all_ranks()
+ path = get_path_with_count(os.path.join(path, f'{prefix}.rank_{rank}'))
+ if type(dict_data).__name__ == 'OrderedDict':
+ torch.save(dict_data, f'{path}.pth')
+ else:
+ with open(f'{path}.json', 'w') as fp:
+ json.dump(dict_data, fp, cls=JSONEncoderWithMcoreTypes)
+
+
+__all__ = ['has_config_logger_enabled', 'log_config_to_disk']
diff --git a/megatron/data/Makefile b/megatron/core/datasets/Makefile
similarity index 100%
rename from megatron/data/Makefile
rename to megatron/core/datasets/Makefile
diff --git a/megatron/fused_kernels/tests/__init__.py b/megatron/core/datasets/__init__.py
similarity index 100%
rename from megatron/fused_kernels/tests/__init__.py
rename to megatron/core/datasets/__init__.py
diff --git a/megatron/core/datasets/bert_dataset.py b/megatron/core/datasets/bert_dataset.py
new file mode 100644
index 0000000000..78ae2edf62
--- /dev/null
+++ b/megatron/core/datasets/bert_dataset.py
@@ -0,0 +1,192 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+import numpy
+
+from megatron.core.datasets.indexed_dataset import IndexedDataset
+from megatron.core.datasets.masked_dataset import (
+ MaskedWordPieceDataset,
+ MaskedWordPieceDatasetConfig,
+)
+from megatron.core.datasets.utils import Split
+
+
+@dataclass
+class BERTMaskedWordPieceDatasetConfig(MaskedWordPieceDatasetConfig):
+ """Configuration object for Megatron Core BERT WordPiece datasets"""
+
+ classification_head: bool = None
+ """Option to perform the next sequence prediction during sampling"""
+
+ def __post_init__(self) -> None:
+ """Do asserts and set fields post init"""
+ super().__post_init__()
+
+ assert self.classification_head is not None
+
+
+class BERTMaskedWordPieceDataset(MaskedWordPieceDataset):
+ """The BERT dataset that assumes WordPiece tokenization
+
+ Args:
+ indexed_dataset (IndexedDataset): The IndexedDataset around which to build the MegatronDataset
+
+ dataset_path (str): The real path on disk to the dataset, for bookkeeping
+
+ indexed_indices (numpy.ndarray): The set of the documents indices to expose
+
+ num_samples (Optional[int]): The number of samples to draw from the indexed dataset. When None, build as many samples as correspond to one epoch.
+
+ index_split (Split): The indexed_indices Split
+
+ config (BERTMaskedWordPieceDatasetConfig): The config
+ """
+
+ def __init__(
+ self,
+ indexed_dataset: IndexedDataset,
+ dataset_path: str,
+ indexed_indices: numpy.ndarray,
+ num_samples: Optional[int],
+ index_split: Split,
+ config: BERTMaskedWordPieceDatasetConfig,
+ ) -> None:
+ super().__init__(
+ indexed_dataset, dataset_path, indexed_indices, num_samples, index_split, config
+ )
+
+ self.token_lookup = list(self.config.tokenizer.inv_vocab.keys())
+ # Account for the single
and two token ids
+ self.sample_index = self._build_sample_index(
+ self.config.sequence_length - 3, 2 if self.config.classification_head else 1
+ )
+
+ @staticmethod
+ def _key_config_attributes() -> List[str]:
+ """Inherited method implementation
+
+ Returns:
+ List[str]: The key config attributes
+ """
+ return super(
+ BERTMaskedWordPieceDataset, BERTMaskedWordPieceDataset
+ )._key_config_attributes() + ["classification_head"]
+
+ def __getitem__(self, idx: int) -> Dict[str, Union[int, numpy.ndarray]]:
+ """Abstract method implementation
+
+ Args:
+ idx (int): The index into the dataset
+
+ Returns:
+ Dict[str, Union[int, numpy.ndarray]]: The
+ """
+ idx_beg, idx_end, target_sequence_length = self.sample_index[idx]
+ sample = [self.dataset[i] for i in range(idx_beg, idx_end)]
+ numpy_random_state = numpy.random.RandomState(seed=(self.config.random_seed + idx) % 2**32)
+
+ assert target_sequence_length <= self.config.sequence_length
+
+ # Split the sample into contiguous subsegments A and B
+ pivot = len(sample)
+ is_next_random = False
+ if self.config.classification_head:
+ assert len(sample) > 1, "the sample must contain at least two sentences"
+ pivot = 1
+ if len(sample) >= 3:
+ pivot = numpy_random_state.randint(low=1, high=len(sample))
+ is_next_random = numpy_random_state.random() < 0.5
+ split_A = []
+ for sample_a in sample[:pivot]:
+ split_A.extend(sample_a)
+ split_B = []
+ for sample_b in sample[pivot:]:
+ split_B.extend(sample_b)
+ if is_next_random:
+ split_A, split_B = split_B, split_A
+
+ # Trim the subsegments from either end to a desired joint length
+ length_A = len(split_A)
+ length_B = len(split_B)
+ if length_A + length_B <= target_sequence_length:
+ truncated = False
+ else:
+ while length_A + length_B > target_sequence_length:
+ split = split_A if length_A > length_B else split_B
+ if numpy_random_state.random() < 0.5:
+ del split[0]
+ else:
+ del split[-1]
+ length_A = len(split_A)
+ length_B = len(split_B)
+ truncated = True
+
+ # Merge the subsegments and create the token assignment labels
+ tokens = [self.config.tokenizer.cls, *split_A, self.config.tokenizer.sep]
+ assignments = [0 for _ in range(1 + len(split_A) + 1)]
+ if split_B:
+ tokens += [*split_B, self.config.tokenizer.sep]
+ assignments += [1 for _ in range(len(split_B) + 1)]
+
+ # Masking
+ tokens, masked_positions, masked_labels, _, _ = self._create_masked_lm_predictions(
+ tokens, target_sequence_length, numpy_random_state
+ )
+
+ # Pad the sequences and convert to NumPy
+ length_toks = len(tokens)
+ length_pads = self.config.sequence_length - length_toks
+ assert length_pads >= 0
+
+ tokens = numpy.array(tokens, dtype=numpy.int64)
+ tokens = numpy.pad(tokens, (0, length_pads), constant_values=self.config.tokenizer.pad)
+
+ assignments = numpy.array(assignments, dtype=numpy.int64)
+ assignments = numpy.pad(
+ assignments, (0, length_pads), constant_values=self.config.tokenizer.pad
+ )
+
+ # Get the padding mask
+ mask_pads = numpy.ones(length_toks, dtype=numpy.int64)
+ mask_pads = numpy.pad(
+ mask_pads, (0, length_pads), constant_values=self.config.tokenizer.pad
+ )
+
+ # Mask the labels
+ labels = numpy.zeros(self.config.sequence_length, dtype=numpy.int64) - 1
+ labels[masked_positions] = masked_labels
+
+ # Get the loss mask
+ mask_loss = numpy.zeros(self.config.sequence_length, dtype=numpy.int64)
+ mask_loss[masked_positions] = 1
+
+ return {
+ "text": tokens,
+ "types": assignments,
+ "labels": labels,
+ "is_random": int(is_next_random),
+ "padding_mask": mask_pads,
+ "loss_mask": mask_loss,
+ "truncated": int(truncated),
+ }
+
+ def _get_token_mask(self, numpy_random_state: numpy.random.RandomState) -> Optional[int]:
+ """Abstract method implementation
+
+ 80% of the time, replace the token id with mask token id. 10% of the time, replace token id
+ with a random token id from the vocabulary. 10% of the time, do nothing.
+
+ Args:
+ numpy_random_state (RandomState): The NumPy random state
+
+ Returns:
+ Optional[int]: The replacement token id or None
+ """
+ if numpy_random_state.random() < 0.8:
+ return self.config.tokenizer.mask
+ else:
+ if numpy_random_state.random() >= 0.5:
+ return self.token_lookup[numpy_random_state.randint(0, len(self.token_lookup))]
+ return None
diff --git a/megatron/core/datasets/blended_dataset.py b/megatron/core/datasets/blended_dataset.py
new file mode 100644
index 0000000000..be0b7a4a08
--- /dev/null
+++ b/megatron/core/datasets/blended_dataset.py
@@ -0,0 +1,201 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import hashlib
+import json
+import logging
+import os
+import time
+from collections import OrderedDict
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy
+import torch
+
+from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
+from megatron.core.datasets.megatron_dataset import MegatronDataset
+from megatron.core.datasets.utils import normalize
+from megatron.core.utils import log_single_rank
+
+logger = logging.getLogger(__name__)
+
+_VERBOSE = False
+
+
+class BlendedDataset(torch.utils.data.Dataset):
+ """Conjugating class for a set of MegatronDataset instances
+
+ Args:
+ datasets (List[MegatronDataset]): The MegatronDataset instances to blend
+
+ weights (List[Union[int, float]]): The weights that determine the dataset blend ratios
+
+ size (Optional[int]): The number of samples to draw from the blend. If None, for each dataset index idx draw exactly weights[idx] samples from datasets[idx].
+
+ config (BlendedMegatronDatasetConfig): The config
+
+ Raises:
+ RuntimeError: When the dataset has fewer or more samples than 'size' post-initialization
+ """
+
+ def __init__(
+ self,
+ datasets: List[MegatronDataset],
+ weights: List[Union[int, float]],
+ size: Optional[int],
+ config: BlendedMegatronDatasetConfig,
+ ) -> None:
+ assert len(datasets) == len(weights)
+ assert len(datasets) < 32767
+ assert all(map(lambda _: type(_) == type(datasets[0]), datasets))
+ assert all(map(lambda _: _.index_split == datasets[0].index_split, datasets))
+ assert all(map(lambda _: _ > 0, weights))
+ assert all(map(lambda _: type(_) == type(weights[0]), weights))
+ if size is None and isinstance(weights[0], float):
+ assert all(map(lambda _: _ == int(_), weights))
+
+ # Alert user to unnecessary blending
+ if len(datasets) == 1:
+ log_single_rank(
+ logger, logging.WARNING, f"Building a BlendedDataset for a single MegatronDataset"
+ )
+
+ if size is not None:
+ weights = normalize(weights)
+
+ self.datasets = datasets
+ self.split = self.datasets[0].index_split
+ self.weights = weights
+ self.size = size
+ self.config = config
+
+ unique_identifiers = OrderedDict()
+ unique_identifiers["class"] = type(self).__name__
+ unique_identifiers["datasets"] = [dataset.unique_identifiers for dataset in self.datasets]
+ unique_identifiers["split"] = self.split.name
+ unique_identifiers["weights"] = self.weights
+ unique_identifiers["size"] = self.size
+ unique_identifiers["renormalize_blend_weights"] = self.config.renormalize_blend_weights
+
+ self.unique_description = json.dumps(
+ unique_identifiers, indent=4, default=lambda obj: obj.unique_identifiers
+ )
+ self.unique_description_hash = hashlib.md5(
+ self.unique_description.encode("utf-8")
+ ).hexdigest()
+
+ self.built_anew_on_cache_miss = False
+
+ self.dataset_index, self.dataset_sample_index = self._build_indices()
+
+ def __len__(self) -> int:
+ return self.dataset_index.shape[0]
+
+ def __getitem__(self, idx: int) -> Dict[str, Union[int, numpy.ndarray]]:
+ dataset_id = self.dataset_index[idx]
+ dataset_sample_id = self.dataset_sample_index[idx]
+ return {"dataset_id": dataset_id, **self.datasets[dataset_id][dataset_sample_id]}
+
+ def _build_indices(self) -> Tuple[numpy.ndarray, numpy.ndarray]:
+ """Build and optionally cache the dataset index and the dataset sample index
+
+ The dataset index is a 1-D mapping which determines the dataset to query. The dataset
+ sample index is a 1-D mapping which determines the sample to request from the queried
+ dataset.
+
+ Returns:
+ Tuple[numpy.ndarray, numpy.ndarray]: The dataset index and the dataset sample index
+ """
+ path_to_cache = self.config.path_to_cache
+
+ if path_to_cache:
+ get_path_to = lambda suffix: os.path.join(
+ path_to_cache,
+ f"{self.unique_description_hash}-{type(self).__name__}-{self.split.name}-{suffix}",
+ )
+ path_to_description = get_path_to("description.txt")
+ path_to_dataset_index = get_path_to("dataset_index.npy")
+ path_to_dataset_sample_index = get_path_to("dataset_sample_index.npy")
+ cache_hit = all(
+ map(
+ os.path.isfile,
+ [path_to_description, path_to_dataset_index, path_to_dataset_sample_index],
+ )
+ )
+ else:
+ cache_hit = False
+
+ if not path_to_cache or (not cache_hit and torch.distributed.get_rank() == 0):
+ log_single_rank(
+ logger, logging.INFO, f"Build and save the {type(self).__name__} indices"
+ )
+ self.built_anew_on_cache_miss = True
+
+ # Build the dataset and dataset sample indexes
+ log_single_rank(
+ logger, logging.INFO, f"\tBuild and save the dataset and dataset sample indexes"
+ )
+ t_beg = time.time()
+ from megatron.core.datasets import helpers
+
+ if self.size is not None:
+ dataset_index = numpy.zeros(self.size, dtype=numpy.int16)
+ dataset_sample_index = numpy.zeros(self.size, dtype=numpy.int64)
+ helpers.build_blending_indices(
+ dataset_index,
+ dataset_sample_index,
+ self.weights,
+ len(self.datasets),
+ self.size,
+ _VERBOSE,
+ )
+ else:
+ size = sum(self.weights)
+ dataset_index = numpy.zeros(size, dtype=numpy.int16)
+ dataset_sample_index = numpy.zeros(size, dtype=numpy.int64)
+ helpers.build_exhaustive_blending_indices(
+ dataset_index, dataset_sample_index, self.weights, len(self.datasets)
+ )
+
+ if path_to_cache:
+ os.makedirs(path_to_cache, exist_ok=True)
+ # Write the description
+ with open(path_to_description, "wt") as writer:
+ writer.write(self.unique_description)
+ # Save the indexes
+ numpy.save(path_to_dataset_index, dataset_index, allow_pickle=True)
+ numpy.save(path_to_dataset_sample_index, dataset_sample_index, allow_pickle=True)
+ else:
+ log_single_rank(
+ logger,
+ logging.WARNING,
+ f"Unable to save the {type(self).__name__} indexes because path_to_cache is None",
+ )
+
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ return dataset_index, dataset_sample_index
+
+ log_single_rank(logger, logging.INFO, f"Load the {type(self).__name__} indices")
+
+ log_single_rank(
+ logger, logging.INFO, f"\tLoad the dataset index from {path_to_dataset_index}"
+ )
+ t_beg = time.time()
+ dataset_index = numpy.load(path_to_dataset_index, allow_pickle=True, mmap_mode='r')
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"\tLoad the dataset sample index from {path_to_dataset_sample_index}",
+ )
+ t_beg = time.time()
+ dataset_sample_index = numpy.load(
+ path_to_dataset_sample_index, allow_pickle=True, mmap_mode='r'
+ )
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ return dataset_index, dataset_sample_index
diff --git a/megatron/core/datasets/blended_megatron_dataset_builder.py b/megatron/core/datasets/blended_megatron_dataset_builder.py
new file mode 100644
index 0000000000..c9cf4abf63
--- /dev/null
+++ b/megatron/core/datasets/blended_megatron_dataset_builder.py
@@ -0,0 +1,528 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+import math
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Callable, Iterable, List, Optional, Type, Union
+
+import numpy
+import torch
+
+from megatron.core.datasets.blended_dataset import BlendedDataset
+from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
+from megatron.core.datasets.megatron_dataset import LowLevelDataset, MegatronDataset
+from megatron.core.datasets.utils import Split, normalize
+from megatron.core.parallel_state import get_virtual_pipeline_model_parallel_rank
+from megatron.core.utils import log_single_rank
+
+logger = logging.getLogger(__name__)
+
+MidLevelDataset = MegatronDataset
+
+TopLevelDataset = Union[BlendedDataset, MidLevelDataset]
+
+DistributedDataset = Union[
+ TopLevelDataset, MidLevelDataset, LowLevelDataset, torch.utils.data.Dataset
+]
+
+
+class BlendedMegatronDatasetBuilder(object):
+ """Builder class for the BlendedDataset and MegatronDataset classes
+
+ Args:
+ cls (Type[MegatronDataset]): The class to instantiate, must inherit from MegatronDataset
+
+ sizes (List[Optional[int]]): The minimum total number of samples to draw, or None, per split
+
+ is_built_on_rank (Callable): A callable which returns True if the dataset should be built on the current rank and False otherwise. It should be Megatron Core parallelism aware i.e. global rank, local group rank, and virtual rank may inform its return value.
+
+ config (BlendedMegatronDatasetConfig): The config object which informs dataset creation
+ """
+
+ def __init__(
+ self,
+ cls: Type[MidLevelDataset],
+ sizes: List[int],
+ is_built_on_rank: Callable,
+ config: BlendedMegatronDatasetConfig,
+ ):
+ self.cls = cls
+ self.sizes = sizes
+ self.is_built_on_rank = is_built_on_rank
+ self.config = config
+
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Building dataset splits with cls={cls.__name__}, sizes={self.sizes}, and config={self.config}",
+ )
+
+ if not self.config.mock:
+ for split in Split:
+ size_is_none = self.sizes[split.value] is None
+ if self.config.blend_per_split is None:
+ weights_are_none = self.config.blend[1] is None
+ else:
+ if self.config.blend_per_split[split.value] is None:
+ continue
+ weights_are_none = self.config.blend_per_split[split.value][1] is None
+ if size_is_none:
+ assert (
+ weights_are_none
+ ), f"size_is_none => weights_are_none fails for {split.name} split"
+
+ if torch.distributed.is_initialized():
+ gb_rank = torch.distributed.get_rank()
+ vp_rank = get_virtual_pipeline_model_parallel_rank()
+ if gb_rank == 0 and (vp_rank == 0 or vp_rank is None):
+ assert (
+ self.is_built_on_rank()
+ ), "is_built_on_rank must return True when global rank = 0 and vp rank = 0"
+
+ def build(self) -> List[Optional[TopLevelDataset]]:
+ """Build all dataset splits according to the provided blend(s)
+
+ This method is distributed-aware and must be called on all ranks.
+
+ The dataset splits returned can vary according to the config. Supply config.blend and
+ config.split to build BlendedDataset and/or MegatronDataset splits from the same
+ distribution. Supply config.blend_per_split to build BlendedDataset and/or MegatronDataset
+ splits from separate distributions. In either case, for each split, handle the following
+ cases:
+
+ (1) The split is None
+ - do nothing
+
+ (2) The split has one contributing dataset, and...
+
+ (a) 'size' is not None
+ - Build a mid-level dataset with low-level dataset sampling in proportion to the size
+
+ (b) 'size' is None
+ - Build mid-level datasets with no excess low-level dataset sampling
+
+ (3) The split has multiple contributing datasets, and...
+
+ (a) 'weights' is not None and 'size' is not None
+ - Build mid-level datasets with low-level dataset sampling in proportion to their weights and the size
+ - Build a top-level dataset of length marginally greater than 'size' with mid-level dataset sampling in proportion to their weights and the size
+
+ (b) 'weights' is not None and 'size' is None
+ - Error
+
+ (c) 'weights' is None and 'size' is not None
+ - Build mid-level datasets with no excess low-level dataset sampling
+ - Build a top-level dataset of length 'size' with mid-level dataset sampling in proportion to their lengths and the size
+
+ - The 'size' of the top-level dataset is capped at the sum of the mid-level dataset lengths
+
+ (d) 'weights' is None and 'size' is None
+ - Build mid-level datasets with no excess low-level dataset sampling
+ - Build a top-level dataset with no excess mid-level dataset sampling
+
+ Returns:
+ List[Optional[TopLevelDataset]]: A list containing a dataset instance (or None) per split
+ """
+ datasets = self._build_blended_dataset_splits()
+
+ for dataset in datasets:
+ if dataset is not None and len(dataset) > 0:
+ if isinstance(dataset, BlendedDataset):
+ if dataset.built_anew_on_cache_miss or any(
+ x.built_anew_on_cache_miss for x in dataset.datasets
+ ):
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Verifying NumPy indices for {type(dataset).__name__} {dataset.split.name} split",
+ )
+ else:
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"NumPy indices for {type(dataset).__name__} {dataset.split.name} split are fully cached, skipping verification",
+ )
+ continue
+ # Check blend size
+ assert dataset.size is None or dataset.size == dataset.dataset_index.shape[0]
+ # Check blend access of mid-level datasets
+ _, sizes = numpy.unique(dataset.dataset_index, return_counts=True)
+ for i, dataset_and_size in enumerate(zip(dataset.datasets, sizes)):
+ if len(dataset_and_size[0]) < dataset_and_size[1]:
+ raise IndexError(
+ f"The {dataset.split.name} blend oversamples (N = {dataset_and_size[1]}) {type(dataset_and_size[0]).__name__} {i} (len = {len(dataset_and_size[0])}). "
+ f"Set renormalize_blend_weights to True and re-run. File an issue if the problem is not resolved."
+ )
+
+ return datasets
+
+ def _build_blended_dataset_splits(self) -> List[Optional[TopLevelDataset]]:
+ """Build all dataset splits according to the provided blend(s)
+
+ See the BlendedMegatronDatasetBuilder.build alias for more information.
+
+ Returns:
+ List[Optional[TopLevelDataset]]: A list containing a dataset instance (or None) per split
+ """
+ ##
+ # Return fake "mock" datasets
+ ##
+ if self.config.mock:
+ split = self.config.split_matrix
+ try:
+ return self._build_megatron_dataset_splits(None, split, self.sizes)
+ except Exception as error:
+ raise Exception(
+ f"{self.cls.__name__} failed to build as a mock data generator"
+ ) from error
+
+ ##
+ # All splits come from the same distribution
+ ##
+ elif self.config.blend:
+ prefixes, weights = self.config.blend
+ if weights is not None:
+ weights = normalize(weights)
+
+ split = self.config.split_matrix
+
+ # Blend consists of a single prefix
+ if len(prefixes) == 1 and weights is None:
+ return self._build_megatron_dataset_splits(prefixes[0], split, self.sizes)
+
+ # Build the mid-level datasets
+ if weights is None:
+ sizes_per_dataset = [[None for split in Split] for prefix in prefixes]
+ else:
+ sizes_per_dataset = _get_size_per_split_per_dataset(weights, self.sizes)
+
+ # build each dataset in parallel
+ megatron_datasets = self._build_megatron_datasets_parallel(
+ prefixes, split, sizes_per_dataset
+ )
+
+ # Build the top-level datasets
+ blended_datasets = [None] * len(Split)
+ for i in range(len(Split)):
+ if split[i] is not None:
+ weights_i = weights
+ if weights_i is not None and self.sizes[i] is not None:
+ size_per_dataset = list(zip(*sizes_per_dataset))[i]
+ size_i = sum(size_per_dataset)
+ if self.config.renormalize_blend_weights:
+ weights_i = list(map(lambda _size: _size / size_i, size_per_dataset))
+ elif weights_i is None:
+ try:
+ weights_i = [
+ len(megatron_dataset) for megatron_dataset in megatron_datasets[i]
+ ]
+ except TypeError:
+ weights_i = [0 for _ in prefixes]
+ if self.sizes[i] is not None:
+ size_i = min(self.sizes[i], sum(weights_i))
+ else:
+ size_i = None # => the size will be sum(weights_i)
+ else:
+ raise RuntimeError
+ blended_datasets[i] = self.build_generic_dataset(
+ BlendedDataset,
+ self.is_built_on_rank,
+ True, # synchronize_ranks, default behavior to build on rank-0 first
+ megatron_datasets[i],
+ weights_i,
+ size_i,
+ self.config,
+ )
+
+ return blended_datasets
+
+ ##
+ # Each split comes from a separate distribution
+ ##
+ else:
+ blended_datasets = [None] * len(Split)
+ for i in range(len(Split)):
+ split_spoof = [None] * len(Split)
+ split_spoof[i] = (0.0, 1.0)
+ sizes_spoof = [0] * len(Split)
+ sizes_spoof[i] = self.sizes[i]
+
+ # Blend is provided for the split
+ blend = self.config.blend_per_split[i]
+ if blend is not None:
+ prefixes, weights = blend
+ if weights is not None:
+ weights = normalize(weights)
+
+ # Blend consists of a sigle prefix
+ if len(prefixes) == 1:
+ blended_datasets[i] = self._build_megatron_dataset_splits(
+ prefixes[0], split_spoof, sizes_spoof
+ )[i]
+ continue
+
+ # Build mid-level datasets
+ if weights is None:
+ sizes_per_dataset = [[None for split in Split] for prefix in prefixes]
+ else:
+ sizes_per_dataset = _get_size_per_split_per_dataset(weights, sizes_spoof)
+
+ # build each dataset in parallel
+ megatron_datasets = self._build_megatron_datasets_parallel(
+ prefixes, split_spoof, sizes_per_dataset
+ )[i]
+
+ # Build top-level dataset
+ if weights is not None and self.sizes[i] is not None:
+ size_per_dataset = list(zip(*sizes_per_dataset))[i]
+ size = sum(size_per_dataset)
+ if self.config.renormalize_blend_weights:
+ weights = list(map(lambda _size: _size / size, size_per_dataset))
+ elif weights is None:
+ try:
+ weights = [
+ len(megatron_dataset) for megatron_dataset in megatron_datasets
+ ]
+ except TypeError:
+ weights = [0 for _ in prefixes]
+ if self.sizes[i] is not None:
+ size = min(self.sizes[i], sum(weights))
+ else:
+ size = None # => the size will be sum(weights)
+ else:
+ raise RuntimeError
+ blended_datasets[i] = self.build_generic_dataset(
+ BlendedDataset,
+ self.is_built_on_rank,
+ True, # synchronize_ranks, default behavior to build on rank-0 first
+ megatron_datasets,
+ weights,
+ size,
+ self.config,
+ )
+
+ return blended_datasets
+
+ def _build_megatron_datasets_parallel(
+ self, prefixes: List[str], split: List[float], sizes_per_dataset: List[List[int]]
+ ) -> List[List[Optional[MegatronDataset]]]:
+ """Build the megatron datasets for a list of prefixes in parallel
+
+ Args:
+ prefixes (List[str]): The list of prefix strings
+
+ split (List[float]): The dataset split ratios (must sum to 1.00)
+
+ sizes_per_dataset (List[List[int]]): The number of samples to request
+ per MegatronDataset per spilt
+
+ Returns:
+ List[List[Optional[MegatronDataset]]]: For each split, have a list of
+ MegatronDataset per prefix
+ """
+
+ # Helper function to wrap the threading logic
+ def _threading_helper(
+ megatron_datasets: List[List[Optional[MegatronDataset]]],
+ num_workers: int,
+ prefixes: List[str],
+ split: List[float],
+ sizes_per_dataset: List[List[int]],
+ ) -> None:
+ with ThreadPoolExecutor(max_workers=num_workers) as executor:
+ all_futures = []
+ for i in range(len(prefixes)):
+ all_futures.append(
+ executor.submit(
+ self._build_megatron_dataset_splits,
+ prefixes[i],
+ split,
+ sizes_per_dataset[i],
+ False, # synchronize_ranks, barrier is called in this function
+ )
+ )
+ for future in all_futures:
+ try:
+ megatron_datasets_split = future.result()
+ for j in range(len(megatron_datasets_split)):
+ megatron_datasets[j].append(megatron_datasets_split[j])
+ except Exception as err:
+ raise err
+
+ megatron_datasets = [[] for _ in range(len(Split))]
+ num_dataset_builder_threads = self.config.num_dataset_builder_threads
+
+ if torch.distributed.is_initialized():
+ rank = torch.distributed.get_rank()
+ # First, build on rank 0
+ if rank == 0:
+ num_workers = num_dataset_builder_threads
+ if num_workers > 1:
+ # since only rank 0 is running, scale up the thread count
+ # but not too much to avoid overloading storage on miss path.
+ # if user set num_dataset_builder_threads to 1,
+ # i.e. meant for serial build, do not scale up.
+ num_workers *= min(2, max(1, torch.cuda.device_count()))
+ _threading_helper(
+ megatron_datasets, num_workers, prefixes, split, sizes_per_dataset
+ )
+
+ torch.distributed.barrier()
+
+ # Then, build on other ranks; guaranteed to be data_cache hit
+ if rank != 0:
+ _threading_helper(
+ megatron_datasets,
+ num_dataset_builder_threads,
+ prefixes,
+ split,
+ sizes_per_dataset,
+ )
+ else:
+ _threading_helper(
+ megatron_datasets, num_dataset_builder_threads, prefixes, split, sizes_per_dataset
+ )
+
+ return megatron_datasets
+
+ def _build_megatron_dataset_splits(
+ self,
+ dataset_path: Optional[str],
+ split: List[float],
+ sizes: List[int],
+ synchronize_ranks: bool = True,
+ ) -> List[Optional[MidLevelDataset]]:
+ """Build each MidLevelDataset split from a single LowLevelDataset
+
+ Args:
+ dataset_path (Optional[str]): The path on disk which defines the underlying LowLevelDataset, or None for mock dataset classes
+
+ split (List[Tuple[float, float]]): The dataset split matrix
+
+ sizes (List[int]): The number of total samples to draw from each split
+
+ synchronize_ranks (bool): Whether to call barrier for rank-0 / barrier / other-ranks behavior. Set to False when we enforce this behavior at higher level.
+
+ Returns:
+ List[Optional[MidLevelDataset]]: The MidLevelDataset (or None) per split
+ """
+ # short-cut if we are not building on this rank
+ if torch.distributed.is_initialized() and not self.is_built_on_rank():
+ for i in range(len(Split)):
+ if split[i] is not None and synchronize_ranks:
+ torch.distributed.barrier()
+ return [None] * len(Split)
+
+ # Build the low level dataset
+ low_level_dataset = self.cls.build_low_level_dataset(dataset_path, self.config)
+
+ # Build the split indices for the low level dataset
+ num_elements = self.cls.numel_low_level_dataset(low_level_dataset)
+ split_indices = []
+ for i, _ in enumerate(Split):
+ if split[i] is not None:
+ beg = int(round(split[i][0] * float(num_elements)))
+ end = int(round(split[i][1] * float(num_elements)))
+ split_indices.append(numpy.arange(start=beg, stop=end, step=1, dtype=numpy.int32))
+ else:
+ split_indices.append(None)
+
+ # Build the mid level dataset
+ mid_level_datasets = []
+ for i, _split in enumerate(Split):
+ if split[i] is None:
+ mid_level_datasets.append(None)
+ else:
+ mid_level_datasets.append(
+ self.build_generic_dataset(
+ self.cls,
+ self.is_built_on_rank,
+ synchronize_ranks,
+ low_level_dataset,
+ dataset_path,
+ split_indices[i],
+ sizes[i],
+ _split,
+ self.config,
+ )
+ )
+
+ return mid_level_datasets
+
+ @staticmethod
+ def build_generic_dataset(
+ cls: Union[Type[DistributedDataset], Callable],
+ is_built_on_rank: Callable,
+ synchronize_ranks: bool,
+ *args: Any,
+ ) -> Optional[Union[DistributedDataset, Iterable]]:
+ """Build the DistributedDataset
+
+ Return None if and only if the underlying dataset class is not built on the current rank
+ and torch.distributed is initialized.
+
+ Args:
+ cls (Union[Type[DistributedDataset], Callable]): The DistributedDataset class to be built. In special cases, e.g. when we are building the low level dataset for a RawMegatronDataset instance, we can accept a Callable which returns an Iterable.
+
+ synchronize_ranks (bool): Whether to call barrier for rank-0 / barrier / other-ranks behavior. Set to False when we enforce this behavior at higher level.
+
+ args (Tuple[Any]): The positional arguments used to build the provided DistributedDataset class
+
+ Raises:
+ Exception: When the dataset constructor raises an OSError
+
+ Returns:
+ Optional[Union[DistributedDataset, Iterable]]: The DistributedDataset instantion, the Iterable instantiation, or None
+ """
+ if torch.distributed.is_initialized():
+ rank = torch.distributed.get_rank()
+
+ dataset = None
+
+ # First, build on rank 0
+ if rank == 0 and is_built_on_rank():
+ try:
+ dataset = cls(*args)
+ except OSError as err:
+ log = (
+ f"Failed to write dataset materials to the data cache directory. "
+ + f"Please supply a directory to which you have write access via "
+ + f"the path_to_cache attribute in BlendedMegatronDatasetConfig and "
+ + f"retry. Refer to the preserved traceback above for more information."
+ )
+ raise Exception(log) from err
+
+ if synchronize_ranks:
+ torch.distributed.barrier()
+
+ # After, build on other ranks
+ if rank != 0 and is_built_on_rank():
+ dataset = cls(*args)
+
+ return dataset
+
+ return cls(*args)
+
+
+def _get_size_per_split_per_dataset(
+ normalized_weights: List[float], target_size_per_split: List[int]
+) -> List[List[int]]:
+ """Determine the contribution of the MegatronDataset splits to the BlendedDataset splits
+
+ Args:
+ normalized_weights (List[float]): e.g. [0.3, 0.7]
+
+ target_size_per_split (List[int]): The number of samples to target for each BlendedDataset split
+
+ Returns:
+ List[List[int]]: The number of samples to request per MegatronDataset per split
+ """
+ assert numpy.isclose(sum(normalized_weights), 1.0)
+
+ # Use 0.5% target margin to ensure we satiate the request
+ sizes_per_dataset = [
+ [int(math.ceil(target_size * weight * 1.005)) for target_size in target_size_per_split]
+ for weight in normalized_weights
+ ]
+
+ return sizes_per_dataset
diff --git a/megatron/core/datasets/blended_megatron_dataset_config.py b/megatron/core/datasets/blended_megatron_dataset_config.py
new file mode 100644
index 0000000000..52bc31f62e
--- /dev/null
+++ b/megatron/core/datasets/blended_megatron_dataset_config.py
@@ -0,0 +1,177 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import functools
+import logging
+import re
+from dataclasses import dataclass, field
+from typing import List, Optional, Tuple
+
+from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer
+from megatron.core.datasets.utils import Split, log_single_rank, normalize
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BlendedMegatronDatasetConfig:
+ """Configuration object for Megatron Core datasets"""
+
+ random_seed: int
+ """The seed for all RNG during dataset creation."""
+
+ sequence_length: int
+ """The sequence length."""
+
+ blend: Optional[Tuple[List[str], Optional[List[float]]]] = None
+ """The blend, consisting of a list of dataset prefixes and optionally a list of dataset
+ weights. For example, [["dataset-path1", "dataset-path2"], [0.3, 0.7]]. When the weights are
+ None, they are inferred from the lengths of the contributing datasets. Not to be used with
+ 'blend_per_split'. Defaults to None.
+ """
+
+ blend_per_split: Optional[List[Optional[Tuple[List[str], Optional[List[float]]]]]] = None
+ """A set of blends, as defined above, one for each split distribution. Not to be used with
+ 'blend'. Defauls to None.
+ """
+
+ renormalize_blend_weights: bool = False
+ """Renormalize the blend weights to account for mid-level dataset oversampling done to ensure
+ fulfillmenet of the of the requested number of samples. Defaults to False for backward
+ comparability in the data sample order.
+ """
+
+ split: Optional[str] = None
+ """The split string, a comma separated weighting for the dataset splits when drawing samples
+ from a single distribution. Not to be used with 'blend_per_split'. Defaults to None.
+ """
+
+ split_matrix: Optional[List[Tuple[float, float]]] = field(init=False, default=None)
+ """The split matrix consisting of non-overlapping book-ends of each split in order. For more
+ information, refer to 'convert_split_vector_to_split_matrix'. Created automatically from
+ 'split'. Not to be passed in to the constructor.
+ """
+
+ num_dataset_builder_threads: int = 1
+ """The number of threads to use for dataset building."""
+
+ path_to_cache: Optional[str] = None
+ """Where all re-useable dataset indices are to be cached."""
+
+ mmap_bin_files: bool = True
+ """Whether to mmap the .bin files or use file pointers."""
+
+ mock: bool = field(init=False, default=False)
+ """Whether to bypass real data loading and validation in favor of mock data generation.
+ Created automatically from 'blend' and 'blend_per_split'. Not to be passed in to the
+ constructor.
+ """
+
+ tokenizer: Optional[MegatronTokenizer] = None
+ """The MegatronTokenizer instance or None. Required for datasets which do online tokenization."""
+
+ def __post_init__(self) -> None:
+ """Do asserts and set fields post init"""
+ if self.blend_per_split is not None and any(self.blend_per_split):
+ assert self.blend is None, "blend and blend_per_split are incompatible"
+ assert self.split is None, "split and blend_per_split are incompatible"
+ assert len(self.blend_per_split) == len(
+ Split
+ ), f"blend_per_split must contain {len(Split)} blends"
+ for split in Split:
+ if self.blend_per_split[split.value] is None:
+ log_single_rank(
+ logger, logging.INFO, f"blend not provided for {split.name} split"
+ )
+ else:
+ assert self.blend_per_split[split.value][1] is None or len(
+ self.blend_per_split[split.value][0]
+ ) == len(
+ self.blend_per_split[split.value][1]
+ ), "blend per split prefixes and weights must be equal in number"
+ else:
+ if self.blend is not None:
+ assert self.blend[1] is None or len(self.blend[0]) == len(
+ self.blend[1]
+ ), "blend prefixes and weights must be equal in number"
+ assert self.split is not None, "split must be provided when blend is not None"
+ else:
+ self.mock = True
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Let mock = True, as both blend and blend_per_split are None",
+ )
+ self.split = "1,1,1"
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Let split = {self.split}, an arbitrarily even split, as mock is True",
+ )
+ split_vector = parse_and_normalize_split(self.split)
+ self.split_matrix = convert_split_vector_to_split_matrix(split_vector)
+ log_single_rank(logger, logging.INFO, f"Let split_matrix = {self.split_matrix}")
+
+
+def parse_and_normalize_split(split: str) -> List[float]:
+ """Parse the dataset split ratios from a string
+
+ Args:
+ split (str): The train valid test split string e.g. "99,1,0"
+
+ Returns:
+ List[float]: The trian valid test split ratios e.g. [0.99, 0.01, 0.0]
+ """
+ split = list(map(float, re.findall(r"[.0-9]+", split)))
+ split = split + [0.0 for _ in range(len(Split) - len(split))]
+
+ assert len(split) == len(Split)
+ assert all(map(lambda _: _ >= 0.0, split))
+
+ split = normalize(split)
+
+ return split
+
+
+def convert_split_vector_to_split_matrix(
+ vector_a: List[float], vector_b: Optional[List[float]] = None
+) -> List[Optional[Tuple[float, float]]]:
+ """Build the split matrix from one or optionally two contributing split vectors.
+
+ Ex. a standard conversion:
+
+ [0.99, 0.01, 0.0] -> [(0, 0.99), (0.99, 1.0), None]
+
+ Ex. a conversion for Retro when Retro pretraining uses a [0.99, 0.01, 0.0] split and Retro
+ preprocessing used a [0.98, 0.02, 0.0] split:
+
+ [0.99, 0.01, 0.0], [0.98, 0.02, 0.0] -> [(0, 0.98), (0.99, 1.0), None]
+
+ Args:
+ vector_a (List[float]): The primary split vector
+
+ vector_b (Optional[List[float]]): An optional secondary split vector which constrains the primary split vector. Defaults to None.
+
+ Returns:
+ List[Tuple[float, float]]: The split matrix consisting of book-ends of each split in order
+ """
+ if vector_b is None:
+ vector_b = vector_a
+
+ # [.900, .090, .010] -> [0.00, .900, .990, 100]
+ expansion_a = functools.reduce(lambda a, b: a + [a[len(a) - 1] + b], [[0], *vector_a])
+ expansion_b = functools.reduce(lambda a, b: a + [a[len(a) - 1] + b], [[0], *vector_b])
+
+ # [0.00, .900, .990, 100.0] -> [(0.00, .900), (.900, .990), (.990, 100)]
+ bookends_a = list(zip(expansion_a[:-1], expansion_a[1:]))
+ bookends_b = list(zip(expansion_b[:-1], expansion_b[1:]))
+
+ # gather per-split overlap or None
+ matrix = []
+ for bookend_a, bookend_b in zip(bookends_a, bookends_b):
+ if min(bookend_a[1], bookend_b[1]) <= max(bookend_a[0], bookend_b[0]):
+ overlap = None
+ else:
+ overlap = (max(bookend_a[0], bookend_b[0]), min(bookend_a[1], bookend_b[1]))
+ matrix.append(overlap)
+
+ return matrix
diff --git a/megatron/core/datasets/gpt_dataset.py b/megatron/core/datasets/gpt_dataset.py
new file mode 100644
index 0000000000..115727de92
--- /dev/null
+++ b/megatron/core/datasets/gpt_dataset.py
@@ -0,0 +1,778 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+import os
+import time
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple
+
+import numpy
+import torch
+
+from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
+from megatron.core.datasets.indexed_dataset import IndexedDataset
+from megatron.core.datasets.megatron_dataset import MegatronDataset
+from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer
+from megatron.core.datasets.utils import Split
+from megatron.core.datasets.utils_s3 import S3Config, is_s3_path
+from megatron.core.utils import log_single_rank
+
+logger = logging.getLogger(__name__)
+
+_PAD_TOKEN_ID = -1
+
+
+@dataclass
+class GPTDatasetConfig(BlendedMegatronDatasetConfig):
+ """Configuration object for Megatron Core GPT datasets"""
+
+ reset_position_ids: bool = None
+ """Option to reset the position IDs in the dataset at an interval"""
+
+ reset_attention_mask: bool = None
+ """Option to reset the attention mask from the dataset"""
+
+ eod_mask_loss: bool = None
+ """Option to enable the EOD mask loss"""
+
+ create_attention_mask: bool = True
+ """Option to enable the attention masks generation. Can be disabled if attention kernel
+ generates masks by itself.
+ """
+
+ drop_last_partial_validation_sequence: bool = True
+ """Option to drop the last partial validation sequence"""
+
+ add_extra_token_to_sequence: bool = True
+ """Option to draw sequences with one extra token to ensure the sample input tokens and sample
+ output tokens are both of the desired sequence length
+ """
+
+ s3_cache_path: str = None
+ """Path for caching indices for s3 dataloading."""
+
+ def __post_init__(self) -> None:
+ """Do asserts and set fields post init"""
+ super().__post_init__()
+
+ assert self.tokenizer is not None
+
+ assert self.reset_position_ids is not None
+ assert self.reset_attention_mask is not None
+ assert self.eod_mask_loss is not None
+
+
+class GPTDataset(MegatronDataset):
+ """The base GPT dataset
+
+ Args:
+ indexed_dataset (IndexedDataset): The IndexedDataset around which to build the GPTDataset
+
+ dataset_path (Optional[str]): The real path on disk to the dataset, for bookkeeping
+
+ indexed_indices (numpy.ndarray): The set of the documents indices to expose
+
+ num_samples (Optional[int]): The number of samples to draw from the indexed dataset. When None, build as many samples as correspond to one epoch.
+
+ index_split (Split): The indexed_indices Split
+
+ config (GPTDatasetConfig): The config
+ """
+
+ def __init__(
+ self,
+ indexed_dataset: IndexedDataset,
+ dataset_path: Optional[str],
+ indexed_indices: numpy.ndarray,
+ num_samples: Optional[int],
+ index_split: Split,
+ config: GPTDatasetConfig,
+ ) -> None:
+ super().__init__(
+ indexed_dataset, dataset_path, indexed_indices, num_samples, index_split, config
+ )
+ self.masks_and_position_ids_are_cacheable = not any(
+ [
+ self.config.reset_position_ids,
+ self.config.reset_attention_mask,
+ self.config.eod_mask_loss,
+ ]
+ )
+ self.masks_and_position_ids_are_cached = False
+ self.cached_attention_mask = None
+ self.cached_loss_mask = None
+ self.cached_position_ids = None
+
+ try:
+ self._pad_token_id = self.config.tokenizer.pad
+ except Exception:
+ self._pad_token_id = _PAD_TOKEN_ID
+
+ (self.document_index, self.sample_index, self.shuffle_index) = (
+ self._build_document_sample_shuffle_indices()
+ )
+
+ @staticmethod
+ def numel_low_level_dataset(low_level_dataset: IndexedDataset) -> int:
+ """Abstract method implementation
+
+ For GPT, the underlying IndexedDataset should be split by sequence, as opposed to, say,
+ BERT, which should be split by document
+
+ Args:
+ low_level_dataset (IndexedDataset): The underlying IndexedDataset
+
+ Returns:
+ int: The number of unique elements in the underlying IndexedDataset
+ """
+ return low_level_dataset.sequence_lengths.shape[0]
+
+ @staticmethod
+ def build_low_level_dataset(dataset_path: str, config: GPTDatasetConfig) -> IndexedDataset:
+ """Abstract method implementation
+
+ Args:
+ dataset_path (str): The real path prefix to the IndexedDataset .bin and .idx files
+
+ config (GPTDatasetConfig): The config
+
+ Returns:
+ IndexedDataset: The underlying IndexedDataset
+ """
+ if is_s3_path(dataset_path):
+ return IndexedDataset(
+ dataset_path,
+ multimodal=False,
+ mmap=config.mmap_bin_files,
+ s3_config=S3Config(path_to_idx_cache=config.s3_cache_path),
+ )
+ return IndexedDataset(dataset_path, multimodal=False, mmap=config.mmap_bin_files)
+
+ def __len__(self) -> int:
+ """Abstract method implementation
+
+ Returns:
+ int: The length of the dataset
+ """
+ return self.sample_index.shape[0] - 1
+
+ def __getitem__(self, idx: Optional[int]) -> Dict[str, torch.Tensor]:
+ """Abstract method implementation
+
+ Args:
+ idx (Optioal[int]): The index into the dataset
+
+ Returns:
+ Dict[str, torch.Tensor]: The sample information wrapped in a dictionary
+ """
+ if idx is None:
+ # Batch padding sequence so the index does not matter
+ text, _ = self._query_document_sample_shuffle_indices(0)
+ else:
+ text, _ = self._query_document_sample_shuffle_indices(idx)
+
+ text = torch.from_numpy(text).long()
+ if self.config.add_extra_token_to_sequence:
+ tokens = text[:-1].contiguous()
+ labels = text[1:].contiguous()
+ else:
+ tokens = text
+ labels = torch.roll(text, shifts=-1, dims=0)
+ labels[-1] = self._pad_token_id
+
+ if (
+ not self.masks_and_position_ids_are_cacheable
+ or not self.masks_and_position_ids_are_cached
+ ):
+ attention_mask, loss_mask, position_ids = _get_ltor_masks_and_position_ids(
+ tokens,
+ self.config.tokenizer.eod,
+ self.config.reset_position_ids,
+ self.config.reset_attention_mask,
+ self.config.eod_mask_loss,
+ self.config.create_attention_mask,
+ )
+ if self.masks_and_position_ids_are_cacheable:
+ self.cached_attention_mask = attention_mask
+ self.cached_loss_mask = loss_mask
+ self.cached_position_ids = position_ids
+ self.masks_and_position_ids_are_cached = True
+ else:
+ attention_mask = self.cached_attention_mask
+ loss_mask = self.cached_loss_mask
+ position_ids = self.cached_position_ids
+
+ # For padded sequences, mask the loss
+ loss_mask[labels == self._pad_token_id] = 0.0
+
+ # For padded sequences, ensure the embedding layer can map the token ID
+ tokens[tokens == self._pad_token_id] = 0
+ labels[labels == self._pad_token_id] = 0
+
+ # Batch padding sequence so we mask the loss
+ if idx is None:
+ loss_mask = torch.zeros_like(loss_mask)
+
+ if self.config.create_attention_mask:
+ return {
+ "tokens": tokens,
+ "labels": labels,
+ "attention_mask": attention_mask,
+ "loss_mask": loss_mask,
+ "position_ids": position_ids,
+ }
+ else:
+ return {
+ "tokens": tokens,
+ "labels": labels,
+ "loss_mask": loss_mask,
+ "position_ids": position_ids,
+ }
+
+ def _query_document_sample_shuffle_indices(
+ self, idx: int
+ ) -> Tuple[numpy.ndarray, numpy.ndarray]:
+ """Get the text (token ids) and document ids for a given index
+
+ Args:
+ idx (int): The index into the dataset
+
+ Returns:
+ Tuple[numpy.ndarray, numpy.ndarray]: The text ids and document ids
+ """
+ # Do the shuffle mapping
+ idx = self.shuffle_index[idx]
+
+ # Get the beginning and end documents and offsets
+ doc_index_beg, doc_index_beg_offset = self.sample_index[idx]
+ doc_index_end, doc_index_end_offset = self.sample_index[idx + 1]
+
+ document_ids = []
+ sample_parts = []
+
+ # Sample spans a single document
+ if doc_index_beg == doc_index_end:
+ # Add the document id
+ document_ids.append(self.document_index[doc_index_beg])
+
+ # Add the entire sample
+ sample_parts.append(
+ self.dataset.get(
+ self.document_index[doc_index_beg],
+ offset=doc_index_beg_offset,
+ length=doc_index_end_offset
+ - doc_index_beg_offset
+ + self.config.add_extra_token_to_sequence,
+ )
+ )
+
+ # Sample spans multiple documents
+ else:
+ for i in range(doc_index_beg, doc_index_end + 1):
+ # Add the document id
+ document_ids.append(self.document_index[i])
+
+ # Add the sample part
+ offset = 0 if i > doc_index_beg else doc_index_beg_offset
+ length = (
+ None
+ if i < doc_index_end
+ else doc_index_end_offset + self.config.add_extra_token_to_sequence
+ )
+ sample_parts.append(
+ self.dataset.get(self.document_index[i], offset=offset, length=length)
+ )
+ assert len(document_ids) == len(
+ sample_parts
+ ), f"len(document_ids) ({len(document_ids)}) != len(sample_parts) ({len(sample_parts)})"
+
+ length = sum(map(len, sample_parts))
+
+ # Pad the sample if necessary
+ if length < (self.config.sequence_length + self.config.add_extra_token_to_sequence):
+ sample_parts.append(
+ [self._pad_token_id]
+ * (self.config.sequence_length + self.config.add_extra_token_to_sequence - length)
+ )
+
+ return (
+ numpy.concatenate(sample_parts, dtype=numpy.int64),
+ numpy.array(document_ids, dtype=numpy.int64),
+ )
+
+ def _build_document_sample_shuffle_indices(
+ self,
+ ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
+ """Build the document index, the sample index, and the shuffle index
+
+ The document index:
+ -- 1-D
+ -- An ordered array of document ids
+
+ The sample index:
+ -- 2-D
+ -- The document indices and offsets which mark the start of every sample
+
+ The shuffle index:
+ -- 1-D
+ -- A random permutation of index range of the sample index
+
+ Returns:
+ Tuple[numpy.ndarray, numpy.ndarray]: The document index, the sample index, and the shuffle index
+ """
+ path_to_cache = self.config.path_to_cache
+ if path_to_cache is None and not self.config.mock:
+ path_to_cache = os.path.join(
+ self.dataset.path_prefix, "cache", f"{type(self).__name__}_indices"
+ )
+
+ if path_to_cache:
+ get_path_to = lambda suffix: os.path.join(
+ path_to_cache,
+ f"{self.unique_description_hash}-{type(self).__name__}-{self.index_split.name}-{suffix}",
+ )
+ path_to_description = get_path_to("description.txt")
+ path_to_document_index = get_path_to("document_index.npy")
+ path_to_sample_index = get_path_to("sample_index.npy")
+ path_to_shuffle_index = get_path_to("shuffle_index.npy")
+ cache_hit = all(
+ map(
+ os.path.isfile,
+ [
+ path_to_description,
+ path_to_document_index,
+ path_to_sample_index,
+ path_to_shuffle_index,
+ ],
+ )
+ )
+ else:
+ cache_hit = False
+
+ if not path_to_cache or (
+ not cache_hit
+ and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0)
+ ):
+
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Build and save the {type(self).__name__} {self.index_split.name} indices",
+ )
+ self.built_anew_on_cache_miss = True
+ t_beg = time.time()
+
+ sequence_length = self.config.sequence_length
+ num_tokens_per_epoch = self._get_num_tokens_per_epoch()
+ num_epochs = self._get_num_epochs(num_tokens_per_epoch)
+
+ if num_epochs == 1:
+ separate_final_epoch = False
+ else:
+ # Get the number of samples for the last epoch
+ num_samples_sans_final_epoch = (
+ (num_epochs - 1) * num_tokens_per_epoch
+ - self.config.add_extra_token_to_sequence
+ ) // sequence_length
+ num_samples_from_final_epoch = self.num_samples - num_samples_sans_final_epoch
+ num_samples_per_epoch = (
+ num_tokens_per_epoch - self.config.add_extra_token_to_sequence
+ ) // sequence_length
+
+ # num_samples_from_final_epoch should be non-negative
+ assert num_samples_from_final_epoch >= 0
+
+ # num_samples_from_final_epoch should not exceed max value
+ assert num_samples_from_final_epoch <= num_samples_per_epoch + 1
+
+ # Separate the final epoch if it falls below the threshold
+ threshold = 0.80
+ separate_final_epoch = num_samples_from_final_epoch < int(
+ threshold * num_samples_per_epoch
+ )
+
+ log_single_rank(
+ logger,
+ logging.DEBUG,
+ f"> num_samples_from_final_epoch: {num_samples_from_final_epoch}",
+ )
+ log_single_rank(logger, logging.DEBUG, f"> threshold: {threshold}")
+ log_single_rank(
+ logger, logging.DEBUG, f"> num_samples_per_epoch: {num_samples_per_epoch}"
+ )
+
+ log_single_rank(
+ logger, logging.DEBUG, f"> separate_final_epoch: {separate_final_epoch}"
+ )
+
+ numpy_random_state = numpy.random.RandomState(self.config.random_seed)
+
+ # Build the document index
+ document_index = _build_document_index(
+ self.indices, num_epochs, numpy_random_state, separate_final_epoch
+ )
+
+ drop_last_partial_sequence = True
+ if self.index_split == Split.valid:
+ drop_last_partial_sequence = self.config.drop_last_partial_validation_sequence
+
+ # Build the sample index
+ from megatron.core.datasets import helpers
+
+ if self.index_split == Split.valid:
+ drop_last_partial_sequence = self.config.drop_last_partial_validation_sequence
+ else:
+ drop_last_partial_sequence = True
+
+ assert document_index.dtype == numpy.int32
+ assert self.dataset.sequence_lengths.dtype == numpy.int32
+ if len(document_index) * 2 > len(self.dataset.sequence_lengths):
+ # Heuristic: if "access density" of sequence_lengths is relatively high,
+ # force loading the mmap-ed array into memory by taking a copy.
+ # System performance benefits come from two aspects:
+ # 1. **sequentially** pre-loading the whole file if we're gonna read a large fraction anyways.
+ # 2. GIL is held when calling into c++ code; making the c++ func faster improves parallelism.
+ sequence_lengths_for_cpp = self.dataset.sequence_lengths.copy()
+ else:
+ sequence_lengths_for_cpp = self.dataset.sequence_lengths
+ sample_index = helpers.build_sample_idx(
+ sequence_lengths_for_cpp,
+ document_index,
+ sequence_length,
+ num_epochs,
+ num_tokens_per_epoch,
+ drop_last_partial_sequence,
+ self.config.add_extra_token_to_sequence,
+ )
+
+ # Build the shuffle index
+ if separate_final_epoch:
+ shuffle_index = _build_shuffle_index(
+ num_samples_sans_final_epoch, sample_index.shape[0] - 1, numpy_random_state
+ )
+ else:
+ shuffle_index = _build_shuffle_index(
+ sample_index.shape[0] - 1, sample_index.shape[0] - 1, numpy_random_state
+ )
+
+ if path_to_cache:
+ os.makedirs(path_to_cache, exist_ok=True)
+ # Write the description
+ with open(path_to_description, "wt") as writer:
+ writer.write(self.unique_description)
+ numpy.save(path_to_document_index, document_index, allow_pickle=True)
+ numpy.save(path_to_sample_index, sample_index, allow_pickle=True)
+ numpy.save(path_to_shuffle_index, shuffle_index, allow_pickle=True)
+ else:
+ log_single_rank(
+ logger,
+ logging.WARNING,
+ f"Unable to save the {type(self).__name__} indexes because path_to_cache is None",
+ )
+
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ log_single_rank(
+ logger, logging.INFO, f"> total number of samples: {sample_index.shape[0] - 1}"
+ )
+ log_single_rank(logger, logging.INFO, f"> total number of epochs: {num_epochs}")
+
+ return document_index, sample_index, shuffle_index
+
+ log_single_rank(
+ logger, logging.INFO, f"Load the {type(self).__name__} {self.index_split.name} indices"
+ )
+
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"\tLoad the document index from {os.path.basename(path_to_document_index)}",
+ )
+ t_beg = time.time()
+ document_index = numpy.load(path_to_document_index, allow_pickle=True, mmap_mode='r')
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"\tLoad the sample index from {os.path.basename(path_to_sample_index)}",
+ )
+ t_beg = time.time()
+ sample_index = numpy.load(path_to_sample_index, allow_pickle=True, mmap_mode='r')
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"\tLoad the shuffle index from {os.path.basename(path_to_shuffle_index)}",
+ )
+ t_beg = time.time()
+ shuffle_index = numpy.load(path_to_shuffle_index, allow_pickle=True, mmap_mode='r')
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ log_single_rank(
+ logger, logging.INFO, f"> total number of samples: {sample_index.shape[0] - 1}"
+ )
+
+ return document_index, sample_index, shuffle_index
+
+ def _get_num_tokens_per_epoch(self) -> int:
+ """Calculate the number of tokens in a single epoch
+
+ Returns:
+ int: The number of tokens in a single epoch
+ """
+ return int(numpy.sum(self.dataset.sequence_lengths[self.indices]))
+
+ def _get_num_epochs(self, num_tokens_per_epoch: int) -> int:
+ """Calculate the number of epochs
+
+ Args:
+ num_tokens_per_epoch (int): The number of tokens in a single epoch
+
+ Returns:
+ int: The number of epochs
+ """
+ num_epochs = 1
+ num_tokens = num_tokens_per_epoch
+ if self.num_samples is None:
+ return num_epochs
+ else:
+ num_tokens_requested = (
+ self.num_samples * self.config.sequence_length
+ ) + self.config.add_extra_token_to_sequence
+ while num_tokens < num_tokens_requested:
+ num_epochs += 1
+ num_tokens += num_tokens_per_epoch
+ return num_epochs
+
+
+def _build_document_index(
+ documents: numpy.ndarray,
+ num_epochs: int,
+ numpy_random_state: numpy.random.RandomState,
+ separate_final_epoch: bool,
+) -> numpy.ndarray:
+ """Build an array with length = num epochs * num documents
+
+ Args:
+ documents (numpy.ndarray): the subset of exposed document indices
+
+ num_epochs (int): The number of epochs
+
+ numpy_random_state (numpy.random.RandomState): The NumPy random state
+
+ separate_final_epoch (bool): Whether to exclude the last epoch from the global shuffle
+
+ Returns:
+ numpy.ndarray: The document index
+ """
+ if not separate_final_epoch or num_epochs == 1:
+ document_index = numpy.mgrid[0:num_epochs, 0 : len(documents)][1]
+ document_index[:] = documents
+ document_index = document_index.reshape(-1)
+ document_index = document_index.astype(numpy.int32)
+ numpy_random_state.shuffle(document_index)
+ return document_index
+
+ doc_idx_first = _build_document_index(documents, num_epochs - 1, numpy_random_state, False)
+ doc_idx_last = _build_document_index(documents, 1, numpy_random_state, False)
+ return numpy.concatenate((doc_idx_first, doc_idx_last))
+
+
+def _build_shuffle_index(
+ num_samples: int, total_size: int, numpy_random_state: numpy.random.RandomState
+) -> numpy.ndarray:
+ """Build the range [0, size) and shuffle
+
+ Args:
+ num_samples (int): The size of the first shuffle range [0, num_samples)
+
+ total_size (int): The size of the entire index. If larger than 'num_samples', it defines the second shuffle range [num_samples, total_size)
+
+ numpy_random_state (numpy.random.RandomState): The NumPy random state
+
+ Returns:
+ numpy.ndarray: The shuffle index
+ """
+ dtype_ = numpy.uint32
+ if total_size >= (numpy.iinfo(numpy.uint32).max - 1):
+ dtype_ = numpy.int64
+
+ shuffle_idx_first = numpy.arange(start=0, stop=num_samples, step=1, dtype=dtype_)
+ numpy_random_state.shuffle(shuffle_idx_first)
+ if num_samples == total_size:
+ return shuffle_idx_first
+
+ shuffle_idx_last = numpy.arange(start=num_samples, stop=total_size, step=1, dtype=dtype_)
+ numpy_random_state.shuffle(shuffle_idx_last)
+
+ return numpy.concatenate((shuffle_idx_first, shuffle_idx_last))
+
+
+def _get_ltor_masks_and_position_ids(
+ data: torch.Tensor,
+ eod_token: int,
+ reset_position_ids: bool,
+ reset_attention_mask: bool,
+ eod_mask_loss: bool,
+ create_attention_mask: bool,
+):
+ """Build masks and position id for left to right model.
+
+ Args:
+ data (torch.Tensor): The data tenor that holds the tokens from the dataset
+
+ eod_token (int): ID of the token to that is considered the EOD
+
+ reset_position_ids (bool): Switch to reset the document position ID's
+
+ reset_attention_mask (bool): Switch to reset the attention mask
+
+ eod_mask_loss (bool): Switch to enable the EOD mask loss
+
+ create_attention_mask (bool): Switch to enable the attention masks generation. Can be disabled if attention kernel generates masks by itself.
+
+ Returns:
+ torch.Tensor: Attention mask needed to be used for Attention
+
+ torch.Tensor: The mask used for loss value during training
+
+ torch.Tensor: The position ID's of the token
+ """
+ seq_length = data.numel()
+
+ if create_attention_mask:
+ attention_mask = torch.tril(
+ torch.ones((seq_length, seq_length), device=data.device)
+ ).unsqueeze(0)
+ else:
+ attention_mask = None
+
+ # Loss mask.
+ loss_mask = torch.ones(seq_length, dtype=torch.float, device=data.device)
+ if eod_mask_loss:
+ loss_mask[data == eod_token] = 0.0
+
+ # Position ids.
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device)
+ # We need to clone as the ids will be modifed based on batch index.
+ if reset_position_ids:
+ position_ids = position_ids.clone()
+
+ if reset_position_ids or reset_attention_mask:
+ # Find indices where EOD token is.
+ eod_index = position_ids[data == eod_token]
+ # Detach indices from positions if going to modify positions.
+ if reset_position_ids:
+ eod_index = eod_index.clone()
+
+ # Loop through EOD indices:
+ prev_index = 0
+ for j in range(eod_index.numel()):
+ i = eod_index[j]
+ # Mask attention loss.
+ if reset_attention_mask and attention_mask is not None:
+ attention_mask[0, (i + 1) :, : (i + 1)] = 0
+ # Reset positions.
+ if reset_position_ids:
+ position_ids[(i + 1) :] -= i + 1 - prev_index
+ prev_index = i + 1
+
+ if attention_mask is not None:
+ # Convert attention mask to binary:
+ attention_mask = attention_mask < 0.5
+
+ return attention_mask, loss_mask, position_ids
+
+
+class MockGPTLowLevelDataset:
+
+ seed: int = 0
+ size: int = 100000
+ max_sequence_length: int = 4096
+
+ def __init__(self, tokenizer: MegatronTokenizer) -> None:
+ self.tokenizer = tokenizer
+ rng = numpy.random.default_rng(seed=self.seed)
+ self.sequence_lengths = rng.integers(
+ low=1, high=self.max_sequence_length, size=self.size, dtype=numpy.int32
+ )
+
+ def __len__(self) -> int:
+ return self.size
+
+ def __getitem__(self, idx: int) -> numpy.number:
+ length = self.sequence_lengths[idx]
+ sample = numpy.int64(
+ numpy.concatenate([numpy.arange(length - 1) + 1, [self.tokenizer.eod]])
+ )
+ return sample
+
+ def get(self, idx: int, offset: int = 0, length: Optional[int] = None) -> numpy.ndarray:
+ if length is None:
+ length = self.sequence_lengths[idx] - offset
+ return self[idx][offset : offset + length]
+
+
+class MockGPTDataset(GPTDataset):
+ """The mock GPT dataset
+
+ Args:
+ indexed_dataset (MockGPTLowLevelDataset): The MockGPTLowLevelDataset around which to build the MockGPTDataset
+
+ dataset_path (Optional[str]): This argument is of no consequence for the MockGPTDataset
+
+ indices (numpy.ndarray): The set of the dataset indices to expose
+
+ num_samples (int): The number of samples to draw from the dataset
+
+ index_split (Split): The indices Split
+
+ config (GPTDatasetConfig): The config
+ """
+
+ def __init__(
+ self,
+ dataset: MockGPTLowLevelDataset,
+ dataset_path: Optional[str],
+ indices: numpy.ndarray,
+ num_samples: int,
+ index_split: Split,
+ config: GPTDatasetConfig,
+ ) -> None:
+ assert config.mock
+
+ super().__init__(dataset, dataset_path, indices, num_samples, index_split, config)
+
+ @staticmethod
+ def numel_low_level_dataset(low_level_dataset: MockGPTLowLevelDataset) -> int:
+ """Abstract method implementation
+
+ Args:
+ low_level_dataset (MockGPTLowLevelDataset): The underlying MockGPTLowLevelDataset
+
+ Returns:
+ int: The number of unique elements in the underlying MockGPTLowLevelDataset
+ """
+ return len(low_level_dataset)
+
+ @staticmethod
+ def build_low_level_dataset(
+ dataset_path: Optional[str], config: GPTDatasetConfig
+ ) -> MockGPTLowLevelDataset:
+ """Abstract method implementation
+
+ Args:
+ dataset_path (Optional[str]): This argument is of no consequence for the MockGPTLowLevelDataset
+
+ config (GPTDatasetConfig): The config
+
+ Returns:
+ MockGPTLowLevelDataset: The underlying MockGPTLowLevelDataset
+ """
+ return MockGPTLowLevelDataset(config.tokenizer)
diff --git a/megatron/core/datasets/helpers.cpp b/megatron/core/datasets/helpers.cpp
new file mode 100644
index 0000000000..0b05f09d7a
--- /dev/null
+++ b/megatron/core/datasets/helpers.cpp
@@ -0,0 +1,839 @@
+/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. */
+
+/* Helper methods for fast index mapping builds */
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+namespace py = pybind11;
+using namespace std;
+
+const int32_t LONG_SENTENCE_LEN = 512;
+
+
+void build_exhaustive_blending_indices(py::array_t &dataset_index, py::array_t &dataset_sample_index, const py::array_t &sizes, const int32_t num_datasets) {
+ /*
+ Build blending indices by sampling exactly as many samples from dataset[i]
+ as is requested by sizes[i] for all i in the range [0, num_datasets).
+ */
+ auto dataset_index_ptr = dataset_index.mutable_unchecked<1>();
+ auto dataset_sample_index_ptr = dataset_sample_index.mutable_unchecked<1>();
+ auto sizes_ptr = sizes.unchecked<1>();
+
+ int64_t total_size = 0;
+ int64_t dataset_sample_counts[num_datasets];
+ std::set dataset_unspent_indices;
+ for (int32_t i = 0; i < num_datasets; ++i) {
+ total_size += sizes_ptr[i];
+ dataset_sample_counts[i] = 0;
+ dataset_unspent_indices.insert(i);
+ }
+
+ // still need fractional weights to sample in proportion to sizes
+ double weights[num_datasets];
+ for (int32_t i = 0; i < num_datasets; ++i) {
+ weights[i] = sizes_ptr[i] / static_cast(total_size);
+ }
+
+ int64_t index_sample = 0;
+ while (dataset_unspent_indices.size() > 0) {
+ double index_sample_double = std::max(static_cast(index_sample), 1.0);
+
+ int64_t error_argmax;
+ double error_max = std::numeric_limits::lowest();
+
+ for (int32_t index_dataset : dataset_unspent_indices) {
+ double error = weights[index_dataset] * index_sample_double - static_cast(dataset_sample_counts[index_dataset]);
+ if (error > error_max) {
+ error_argmax = index_dataset;
+ error_max = error;
+ }
+ }
+
+ // Populate the indices.
+ dataset_index_ptr[index_sample] = static_cast(error_argmax);
+ dataset_sample_index_ptr[index_sample] = dataset_sample_counts[error_argmax];
+
+ // Update the total samples.
+ dataset_sample_counts[error_argmax] += 1;
+
+ if (sizes_ptr[error_argmax] - static_cast(dataset_sample_counts[error_argmax]) == 0) {
+ dataset_unspent_indices.erase(error_argmax);
+ }
+
+ index_sample += 1;
+ }
+}
+
+void build_blending_indices(py::array_t &dataset_index,
+ py::array_t &dataset_sample_index,
+ const py::array_t &weights,
+ const int32_t num_datasets,
+ const int64_t size, const bool verbose)
+{
+ /* Given multiple datasets and a weighting array, build samples
+ such that it follows those wieghts.*/
+
+ if (verbose)
+ {
+ std::cout << "> building indices for blended datasets ..." << std::endl;
+ }
+
+ // Get the pointer access without the checks.
+ auto dataset_index_ptr = dataset_index.mutable_unchecked<1>();
+ auto dataset_sample_index_ptr = dataset_sample_index.mutable_unchecked<1>();
+ auto weights_ptr = weights.unchecked<1>();
+
+ // Initialize buffer for number of samples used for each dataset.
+ int64_t current_samples[num_datasets];
+ for (int64_t i = 0; i < num_datasets; ++i)
+ {
+ current_samples[i] = 0;
+ }
+
+ // For each sample:
+ for (int64_t sample_idx = 0; sample_idx < size; ++sample_idx)
+ {
+
+ // Determine where the max error in sampling is happening.
+ auto sample_idx_double = std::max(static_cast(sample_idx), 1.0);
+ int64_t max_error_index = 0;
+ double max_error = weights_ptr[0] * sample_idx_double -
+ static_cast(current_samples[0]);
+ for (int64_t dataset_idx = 1; dataset_idx < num_datasets; ++dataset_idx)
+ {
+ double error = weights_ptr[dataset_idx] * sample_idx_double -
+ static_cast(current_samples[dataset_idx]);
+ if (error > max_error)
+ {
+ max_error = error;
+ max_error_index = dataset_idx;
+ }
+ }
+
+ // Populate the indices.
+ dataset_index_ptr[sample_idx] = static_cast(max_error_index);
+ dataset_sample_index_ptr[sample_idx] = current_samples[max_error_index];
+
+ // Update the total samples.
+ current_samples[max_error_index] += 1;
+ }
+
+ // print info
+ if (verbose)
+ {
+ std::cout << " > sample ratios:" << std::endl;
+ for (int64_t dataset_idx = 0; dataset_idx < num_datasets; ++dataset_idx)
+ {
+ auto ratio = static_cast(current_samples[dataset_idx]) /
+ static_cast(size);
+ std::cout << " dataset " << dataset_idx << ", input: " << weights_ptr[dataset_idx] << ", achieved: " << ratio << std::endl;
+ }
+ }
+}
+
+py::array build_sample_idx(const py::array_t &sizes_,
+ const py::array_t &doc_idx_,
+ const int32_t seq_length,
+ const int32_t num_epochs,
+ const int64_t tokens_per_epoch,
+ const bool drop_last_partial_sequence = true,
+ const int add_extra_token_to_sequence = 1)
+{
+ /* Sample index (sample_idx) is used for gpt2 like dataset for which
+ the documents are flattened and the samples are built based on this
+ 1-D flatten array. It is a 2D array with sizes [number-of-samples + 1, 2]
+ where [..., 0] contains the index into `doc_idx` and [..., 1] is the
+ starting offset in that document.*/
+
+ // Consistency checks.
+ assert(seq_length > 1);
+ assert(num_epochs > 0);
+ assert(tokens_per_epoch > 1);
+
+ // Remove bound checks.
+ auto sizes = sizes_.unchecked<1>();
+ auto doc_idx = doc_idx_.unchecked<1>();
+
+ // Mapping and it's length (1D).
+ int64_t num_samples = 0;
+ if (drop_last_partial_sequence == true)
+ {
+ num_samples = (num_epochs * tokens_per_epoch - add_extra_token_to_sequence) / seq_length;
+ }
+ else
+ {
+ num_samples = ceil(float(num_epochs * tokens_per_epoch - add_extra_token_to_sequence) / seq_length);
+ }
+ int64_t *sample_idx = new int64_t[2 * (num_samples + 1)];
+
+ // Index into sample_idx.
+ int64_t sample_index = 0;
+ // Index into doc_idx.
+ int64_t doc_idx_index = 0;
+ // Begining offset for each document.
+ int32_t doc_offset = 0;
+ // Start with first document and no offset.
+ sample_idx[2 * sample_index] = doc_idx_index;
+ sample_idx[2 * sample_index + 1] = doc_offset;
+ ++sample_index;
+
+ while (sample_index <= num_samples)
+ {
+ // Start with a fresh sequence.
+ int32_t remaining_seq_length = seq_length + add_extra_token_to_sequence;
+ while (remaining_seq_length != 0)
+ {
+ // Get the document length.
+ auto doc_id = doc_idx[doc_idx_index];
+ auto doc_length = sizes[doc_id] - doc_offset;
+ // And add it to the current sequence.
+ remaining_seq_length -= doc_length;
+ // If we have more than a full sequence, adjust offset and set
+ // remaining length to zero so we return from the while loop.
+ // Note that -1 here is for the same reason we have -1 in
+ // `_num_epochs` calculations.
+ if (remaining_seq_length <= 0)
+ {
+ doc_offset += (remaining_seq_length + doc_length - add_extra_token_to_sequence);
+ remaining_seq_length = 0;
+ }
+ else
+ {
+ // Otherwise, start from the begining of the next document.
+ if (doc_idx_index == (doc_idx_.shape(0) - 1))
+ {
+ // If we have reached the end of the documents, break.
+ assert(sample_index == num_samples);
+ doc_offset = sizes[doc_idx[doc_idx_index]] - add_extra_token_to_sequence;
+ break;
+ }
+ ++doc_idx_index;
+ doc_offset = 0;
+ }
+ }
+ // Record the sequence.
+ sample_idx[2 * sample_index] = doc_idx_index;
+ sample_idx[2 * sample_index + 1] = doc_offset;
+ ++sample_index;
+ }
+
+ // Method to deallocate memory.
+ py::capsule free_when_done(sample_idx, [](void *mem_)
+ {
+ int64_t *mem = reinterpret_cast(mem_);
+ delete[] mem; });
+
+ // Return the numpy array.
+ const auto byte_size = sizeof(int64_t);
+ return py::array(std::vector{num_samples + 1, 2}, // shape
+ {2 * byte_size, byte_size}, // C-style contiguous strides
+ sample_idx, // the data pointer
+ free_when_done); // numpy array references
+}
+
+inline int32_t get_target_sample_len(const int32_t short_seq_ratio,
+ const int32_t max_length,
+ std::mt19937 &rand32_gen)
+{
+ /* Training sample length. */
+ if (short_seq_ratio == 0)
+ {
+ return max_length;
+ }
+ const auto random_number = rand32_gen();
+ if ((random_number % short_seq_ratio) == 0)
+ {
+ return 2 + random_number % (max_length - 1);
+ }
+ return max_length;
+}
+
+template
+py::array build_mapping_impl(const py::array_t &docs_,
+ const py::array_t &sizes_,
+ const int32_t num_epochs,
+ const uint64_t max_num_samples,
+ const int32_t max_seq_length,
+ const double short_seq_prob,
+ const int32_t seed,
+ const bool verbose,
+ const int32_t min_num_sent)
+{
+ /* Build a mapping of (start-index, end-index, sequence-length) where
+ start and end index are the indices of the sentences in the sample
+ and sequence-length is the target sequence length.
+ */
+
+ // Consistency checks.
+ assert(num_epochs > 0);
+ assert(max_seq_length > 1);
+ assert(short_seq_prob >= 0.0);
+ assert(short_seq_prob <= 1.0);
+ assert(seed > 0);
+
+ // Remove bound checks.
+ auto docs = docs_.unchecked<1>();
+ auto sizes = sizes_.unchecked<1>();
+
+ // For efficiency, convert probability to ratio. Note: rand() generates int.
+ int32_t short_seq_ratio = 0;
+ if (short_seq_prob > 0)
+ {
+ short_seq_ratio = static_cast(round(1.0 / short_seq_prob));
+ }
+
+ if (verbose)
+ {
+ const auto sent_start_index = docs[0];
+ const auto sent_end_index = docs[docs_.shape(0) - 1];
+ const auto num_sentences = sent_end_index - sent_start_index;
+ cout << " using:" << endl
+ << std::flush;
+ cout << " number of documents: " << docs_.shape(0) - 1 << endl
+ << std::flush;
+ cout << " sentences range: [" << sent_start_index << ", " << sent_end_index << ")" << endl
+ << std::flush;
+ cout << " total number of sentences: " << num_sentences << endl
+ << std::flush;
+ cout << " number of epochs: " << num_epochs << endl
+ << std::flush;
+ cout << " maximum number of samples: " << max_num_samples << endl
+ << std::flush;
+ cout << " maximum sequence length: " << max_seq_length << endl
+ << std::flush;
+ cout << " short sequence probability: " << short_seq_prob << endl
+ << std::flush;
+ cout << " short sequence ration (1/prob): " << short_seq_ratio << endl
+ << std::flush;
+ cout << " seed: " << seed << endl
+ << std::flush;
+ }
+
+ // Mapping and it's length (1D).
+ int64_t num_samples = -1;
+ DocIdx *maps = NULL;
+
+ // Perform two iterations, in the first iteration get the size
+ // and allocate memory and in the second iteration populate the map.
+ bool second = false;
+ for (int32_t iteration = 0; iteration < 2; ++iteration)
+ {
+
+ // Set the seed so both iterations produce the same results.
+ std::mt19937 rand32_gen(seed);
+
+ // Set the flag on second iteration.
+ second = (iteration == 1);
+
+ // Counters:
+ uint64_t empty_docs = 0;
+ uint64_t one_sent_docs = 0;
+ uint64_t long_sent_docs = 0;
+
+ // Current map index.
+ uint64_t map_index = 0;
+
+ // For each epoch:
+ for (int32_t epoch = 0; epoch < num_epochs; ++epoch)
+ {
+ if (map_index >= max_num_samples)
+ {
+ if (verbose && (!second))
+ {
+ cout << " reached " << max_num_samples << " samples after "
+ << epoch << " epochs ..." << endl
+ << std::flush;
+ }
+ break;
+ }
+ // For each document:
+ for (int32_t doc = 0; doc < (docs.shape(0) - 1); ++doc)
+ {
+
+ // Document sentences are in [sent_index_first, sent_index_last)
+ const auto sent_index_first = docs[doc];
+ const auto sent_index_last = docs[doc + 1];
+
+ // At the begining of the document previous index is the
+ // start index.
+ auto prev_start_index = sent_index_first;
+
+ // Remaining documents.
+ auto num_remain_sent = sent_index_last - sent_index_first;
+
+ // Some bookkeeping
+ if ((epoch == 0) && (!second))
+ {
+ if (num_remain_sent == 0)
+ {
+ ++empty_docs;
+ }
+ if (num_remain_sent == 1)
+ {
+ ++one_sent_docs;
+ }
+ }
+
+ // Detect documents with long sentences.
+ bool contains_long_sentence = false;
+ if (num_remain_sent > 1)
+ {
+ for (auto sent_index = sent_index_first;
+ sent_index < sent_index_last; ++sent_index)
+ {
+ if (sizes[sent_index] > LONG_SENTENCE_LEN)
+ {
+ if ((epoch == 0) && (!second))
+ {
+ ++long_sent_docs;
+ }
+ contains_long_sentence = true;
+ break;
+ }
+ }
+ }
+
+ // If we have more than two sentences.
+ if ((num_remain_sent >= min_num_sent) && (!contains_long_sentence))
+ {
+
+ // Set values.
+ auto seq_len = int32_t{0};
+ auto num_sent = int32_t{0};
+ auto target_seq_len = get_target_sample_len(short_seq_ratio,
+ max_seq_length,
+ rand32_gen);
+
+ // Loop through sentences.
+ for (auto sent_index = sent_index_first;
+ sent_index < sent_index_last; ++sent_index)
+ {
+
+ // Add the size and number of sentences.
+ seq_len += sizes[sent_index];
+ ++num_sent;
+ --num_remain_sent;
+
+ // If we have reached the target length.
+ // and if not only one sentence is left in the document.
+ // and if we have at least two sentneces.
+ // and if we have reached end of the document.
+ if (((seq_len >= target_seq_len) &&
+ (num_remain_sent > 1) &&
+ (num_sent >= min_num_sent)) ||
+ (num_remain_sent == 0))
+ {
+
+ // Check for overflow.
+ if ((3 * map_index + 2) >
+ std::numeric_limits::max())
+ {
+ cout << "number of samples exceeded maximum "
+ << "allowed by type int64: "
+ << std::numeric_limits::max()
+ << endl;
+ throw std::overflow_error("Number of samples");
+ }
+
+ // Populate the map.
+ if (second)
+ {
+ const auto map_index_0 = 3 * map_index;
+ maps[map_index_0] = static_cast(prev_start_index);
+ maps[map_index_0 + 1] = static_cast(sent_index + 1);
+ maps[map_index_0 + 2] = static_cast(target_seq_len);
+ }
+
+ // Update indices / counters.
+ ++map_index;
+ prev_start_index = sent_index + 1;
+ target_seq_len = get_target_sample_len(short_seq_ratio,
+ max_seq_length,
+ rand32_gen);
+ seq_len = 0;
+ num_sent = 0;
+ }
+
+ } // for (auto sent_index=sent_index_first; ...
+ } // if (num_remain_sent > 1) {
+ } // for (int doc=0; doc < num_docs; ++doc) {
+ } // for (int epoch=0; epoch < num_epochs; ++epoch) {
+
+ if (!second)
+ {
+ if (verbose)
+ {
+ cout << " number of empty documents: " << empty_docs << endl
+ << std::flush;
+ cout << " number of documents with one sentence: " << one_sent_docs << endl
+ << std::flush;
+ cout << " number of documents with long sentences: " << long_sent_docs << endl
+ << std::flush;
+ cout << " will create mapping for " << map_index << " samples" << endl
+ << std::flush;
+ }
+ assert(maps == NULL);
+ assert(num_samples < 0);
+ maps = new DocIdx[3 * map_index];
+ num_samples = static_cast(map_index);
+ }
+
+ } // for (int iteration=0; iteration < 2; ++iteration) {
+
+ // Shuffle.
+ // We need a 64 bit random number generator as we might have more
+ // than 2 billion samples.
+ std::mt19937_64 rand64_gen(seed + 1);
+ for (auto i = (num_samples - 1); i > 0; --i)
+ {
+ const auto j = static_cast(rand64_gen() % (i + 1));
+ const auto i0 = 3 * i;
+ const auto j0 = 3 * j;
+ // Swap values.
+ swap(maps[i0], maps[j0]);
+ swap(maps[i0 + 1], maps[j0 + 1]);
+ swap(maps[i0 + 2], maps[j0 + 2]);
+ }
+
+ // Method to deallocate memory.
+ py::capsule free_when_done(maps, [](void *mem_)
+ {
+ DocIdx *mem = reinterpret_cast(mem_);
+ delete[] mem; });
+
+ // Return the numpy array.
+ const auto byte_size = sizeof(DocIdx);
+ return py::array(std::vector{num_samples, 3}, // shape
+ {3 * byte_size, byte_size}, // C-style contiguous strides
+ maps, // the data pointer
+ free_when_done); // numpy array references
+}
+
+py::array build_mapping(const py::array_t &docs_,
+ const py::array_t &sizes_,
+ const int num_epochs,
+ const uint64_t max_num_samples,
+ const int max_seq_length,
+ const double short_seq_prob,
+ const int seed,
+ const bool verbose,
+ const int32_t min_num_sent)
+{
+
+ if (sizes_.size() > std::numeric_limits::max())
+ {
+ if (verbose)
+ {
+ cout << " using uint64 for data mapping..." << endl
+ << std::flush;
+ }
+ return build_mapping_impl(docs_, sizes_, num_epochs,
+ max_num_samples, max_seq_length,
+ short_seq_prob, seed, verbose,
+ min_num_sent);
+ }
+ else
+ {
+ if (verbose)
+ {
+ cout << " using uint32 for data mapping..." << endl
+ << std::flush;
+ }
+ return build_mapping_impl(docs_, sizes_, num_epochs,
+ max_num_samples, max_seq_length,
+ short_seq_prob, seed, verbose,
+ min_num_sent);
+ }
+}
+
+template
+py::array build_blocks_mapping_impl(const py::array_t &docs_,
+ const py::array_t &sizes_,
+ const py::array_t &titles_sizes_,
+ const int32_t num_epochs,
+ const uint64_t max_num_samples,
+ const int32_t max_seq_length,
+ const int32_t seed,
+ const bool verbose,
+ const bool use_one_sent_blocks)
+{
+ /* Build a mapping of (start-index, end-index, sequence-length) where
+ start and end index are the indices of the sentences in the sample
+ and sequence-length is the target sequence length.
+ */
+
+ // Consistency checks.
+ assert(num_epochs > 0);
+ assert(max_seq_length > 1);
+ assert(seed > 0);
+
+ // Remove bound checks.
+ auto docs = docs_.unchecked<1>();
+ auto sizes = sizes_.unchecked<1>();
+ auto titles_sizes = titles_sizes_.unchecked<1>();
+
+ if (verbose)
+ {
+ const auto sent_start_index = docs[0];
+ const auto sent_end_index = docs[docs_.shape(0) - 1];
+ const auto num_sentences = sent_end_index - sent_start_index;
+ cout << " using:" << endl
+ << std::flush;
+ cout << " number of documents: " << docs_.shape(0) - 1 << endl
+ << std::flush;
+ cout << " sentences range: [" << sent_start_index << ", " << sent_end_index << ")" << endl
+ << std::flush;
+ cout << " total number of sentences: " << num_sentences << endl
+ << std::flush;
+ cout << " number of epochs: " << num_epochs << endl
+ << std::flush;
+ cout << " maximum number of samples: " << max_num_samples << endl
+ << std::flush;
+ cout << " maximum sequence length: " << max_seq_length << endl
+ << std::flush;
+ cout << " seed: " << seed << endl
+ << std::flush;
+ }
+
+ // Mapping and its length (1D).
+ int64_t num_samples = -1;
+ DocIdx *maps = NULL;
+
+ // Acceptable number of sentences per block.
+ int min_num_sent = 2;
+ if (use_one_sent_blocks)
+ {
+ min_num_sent = 1;
+ }
+
+ // Perform two iterations, in the first iteration get the size
+ // and allocate memory and in the second iteration populate the map.
+ bool second = false;
+ for (int32_t iteration = 0; iteration < 2; ++iteration)
+ {
+
+ // Set the flag on second iteration.
+ second = (iteration == 1);
+
+ // Current map index.
+ uint64_t map_index = 0;
+
+ uint64_t empty_docs = 0;
+ uint64_t one_sent_docs = 0;
+ uint64_t long_sent_docs = 0;
+ // For each epoch:
+ for (int32_t epoch = 0; epoch < num_epochs; ++epoch)
+ {
+ // assign every block a unique id
+ int32_t block_id = 0;
+
+ if (map_index >= max_num_samples)
+ {
+ if (verbose && (!second))
+ {
+ cout << " reached " << max_num_samples << " samples after "
+ << epoch << " epochs ..." << endl
+ << std::flush;
+ }
+ break;
+ }
+ // For each document:
+ for (int32_t doc = 0; doc < (docs.shape(0) - 1); ++doc)
+ {
+
+ // Document sentences are in [sent_index_first, sent_index_last)
+ const auto sent_index_first = docs[doc];
+ const auto sent_index_last = docs[doc + 1];
+ const auto target_seq_len = max_seq_length - titles_sizes[doc];
+
+ // At the begining of the document previous index is the
+ // start index.
+ auto prev_start_index = sent_index_first;
+
+ // Remaining documents.
+ auto num_remain_sent = sent_index_last - sent_index_first;
+
+ // Some bookkeeping
+ if ((epoch == 0) && (!second))
+ {
+ if (num_remain_sent == 0)
+ {
+ ++empty_docs;
+ }
+ if (num_remain_sent == 1)
+ {
+ ++one_sent_docs;
+ }
+ }
+ // Detect documents with long sentences.
+ bool contains_long_sentence = false;
+ if (num_remain_sent >= min_num_sent)
+ {
+ for (auto sent_index = sent_index_first;
+ sent_index < sent_index_last; ++sent_index)
+ {
+ if (sizes[sent_index] > LONG_SENTENCE_LEN)
+ {
+ if ((epoch == 0) && (!second))
+ {
+ ++long_sent_docs;
+ }
+ contains_long_sentence = true;
+ break;
+ }
+ }
+ }
+ // If we have enough sentences and no long sentences.
+ if ((num_remain_sent >= min_num_sent) && (!contains_long_sentence))
+ {
+
+ // Set values.
+ auto seq_len = int32_t{0};
+ auto num_sent = int32_t{0};
+
+ // Loop through sentences.
+ for (auto sent_index = sent_index_first;
+ sent_index < sent_index_last; ++sent_index)
+ {
+
+ // Add the size and number of sentences.
+ seq_len += sizes[sent_index];
+ ++num_sent;
+ --num_remain_sent;
+
+ // If we have reached the target length.
+ // and there are an acceptable number of sentences left
+ // and if we have at least the minimum number of sentences.
+ // or if we have reached end of the document.
+ if (((seq_len >= target_seq_len) &&
+ (num_remain_sent >= min_num_sent) &&
+ (num_sent >= min_num_sent)) ||
+ (num_remain_sent == 0))
+ {
+
+ // Populate the map.
+ if (second)
+ {
+ const auto map_index_0 = 4 * map_index;
+ // Each sample has 4 items: the starting sentence index, ending sentence index,
+ // the index of the document from which the block comes (used for fetching titles)
+ // and the unique id of the block (used for creating block indexes)
+
+ maps[map_index_0] = static_cast(prev_start_index);
+ maps[map_index_0 + 1] = static_cast(sent_index + 1);
+ maps[map_index_0 + 2] = static_cast(doc);
+ maps[map_index_0 + 3] = static_cast(block_id);
+ }
+
+ // Update indices / counters.
+ ++map_index;
+ ++block_id;
+ prev_start_index = sent_index + 1;
+ seq_len = 0;
+ num_sent = 0;
+ }
+ } // for (auto sent_index=sent_index_first; ...
+ } // if (num_remain_sent > 1) {
+ } // for (int doc=0; doc < num_docs; ++doc) {
+ } // for (int epoch=0; epoch < num_epochs; ++epoch) {
+
+ if (!second)
+ {
+ if (verbose)
+ {
+ cout << " number of empty documents: " << empty_docs << endl
+ << std::flush;
+ cout << " number of documents with one sentence: " << one_sent_docs << endl
+ << std::flush;
+ cout << " number of documents with long sentences: " << long_sent_docs << endl
+ << std::flush;
+ cout << " will create mapping for " << map_index << " samples" << endl
+ << std::flush;
+ }
+ assert(maps == NULL);
+ assert(num_samples < 0);
+ maps = new DocIdx[4 * map_index];
+ num_samples = static_cast(map_index);
+ }
+
+ } // for (int iteration=0; iteration < 2; ++iteration) {
+
+ // Shuffle.
+ // We need a 64 bit random number generator as we might have more
+ // than 2 billion samples.
+ std::mt19937_64 rand64_gen(seed + 1);
+ for (auto i = (num_samples - 1); i > 0; --i)
+ {
+ const auto j = static_cast(rand64_gen() % (i + 1));
+ const auto i0 = 4 * i;
+ const auto j0 = 4 * j;
+ // Swap values.
+ swap(maps[i0], maps[j0]);
+ swap(maps[i0 + 1], maps[j0 + 1]);
+ swap(maps[i0 + 2], maps[j0 + 2]);
+ swap(maps[i0 + 3], maps[j0 + 3]);
+ }
+
+ // Method to deallocate memory.
+ py::capsule free_when_done(maps, [](void *mem_)
+ {
+ DocIdx *mem = reinterpret_cast(mem_);
+ delete[] mem; });
+
+ // Return the numpy array.
+ const auto byte_size = sizeof(DocIdx);
+ return py::array(std::vector{num_samples, 4}, // shape
+ {4 * byte_size, byte_size}, // C-style contiguous strides
+ maps, // the data pointer
+ free_when_done); // numpy array references
+}
+
+py::array build_blocks_mapping(const py::array_t &docs_,
+ const py::array_t &sizes_,
+ const py::array_t &titles_sizes_,
+ const int num_epochs,
+ const uint64_t max_num_samples,
+ const int max_seq_length,
+ const int seed,
+ const bool verbose,
+ const bool use_one_sent_blocks)
+{
+
+ if (sizes_.size() > std::numeric_limits::max())
+ {
+ if (verbose)
+ {
+ cout << " using uint64 for data mapping..." << endl
+ << std::flush;
+ }
+ return build_blocks_mapping_impl(docs_, sizes_, titles_sizes_,
+ num_epochs, max_num_samples, max_seq_length, seed, verbose, use_one_sent_blocks);
+ }
+ else
+ {
+ if (verbose)
+ {
+ cout << " using uint32 for data mapping..." << endl
+ << std::flush;
+ }
+ return build_blocks_mapping_impl(docs_, sizes_, titles_sizes_,
+ num_epochs, max_num_samples, max_seq_length, seed, verbose, use_one_sent_blocks);
+ }
+}
+
+PYBIND11_MODULE(helpers, m)
+{
+ m.def("build_mapping", &build_mapping);
+ m.def("build_blocks_mapping", &build_blocks_mapping);
+ m.def("build_sample_idx", &build_sample_idx);
+ m.def("build_blending_indices", &build_blending_indices);
+ m.def("build_exhaustive_blending_indices", &build_exhaustive_blending_indices);
+}
diff --git a/megatron/core/datasets/indexed_dataset.py b/megatron/core/datasets/indexed_dataset.py
new file mode 100644
index 0000000000..29975336f1
--- /dev/null
+++ b/megatron/core/datasets/indexed_dataset.py
@@ -0,0 +1,857 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Essentially re-written in entirety
+
+import logging
+import os
+import shutil
+import struct
+import time
+from abc import ABC, abstractmethod
+from enum import Enum
+from functools import lru_cache
+from itertools import accumulate
+from types import TracebackType
+from typing import List, Optional, Tuple, Type, Union
+
+try:
+ import boto3
+except ModuleNotFoundError:
+ pass
+import numpy
+import torch
+
+from megatron.core.datasets.utils_s3 import (
+ S3Config,
+ is_s3_path,
+ maybe_download_file,
+ object_exists,
+ parse_s3_path,
+)
+from megatron.core.utils import log_single_rank
+
+logger = logging.getLogger(__name__)
+
+_INDEX_HEADER = b"MMIDIDX\x00\x00"
+
+
+class DType(Enum):
+ """The NumPy data type Enum for writing/reading the IndexedDataset indices"""
+
+ uint8 = 1
+ int8 = 2
+ int16 = 3
+ int32 = 4
+ int64 = 5
+ float64 = 6
+ float32 = 7
+ uint16 = 8
+
+ @classmethod
+ def code_from_dtype(cls, value: Type[numpy.number]) -> int:
+ """Get the code from the dtype
+
+ Args:
+ value (Type[numpy.number]): The dtype
+
+ Returns:
+ int: The code
+ """
+ return cls[value.__name__].value
+
+ @classmethod
+ def dtype_from_code(cls, value: int) -> Type[numpy.number]:
+ """Get the dtype from the code
+
+ Args:
+ value (int): The code
+
+ Returns:
+ Type[numpy.number]: The dtype
+ """
+ return getattr(numpy, cls(value).name)
+
+ @staticmethod
+ def size(key: Union[int, Type[numpy.number]]) -> int:
+ """Get the size of the dtype/code in bytes
+
+ Args:
+ key (Union[int, Type[numpy.number]]): The dtype or code
+
+ Raises:
+ ValueError: If the key is neither dtype nor integer code
+
+ Returns:
+ int: The size of the dtype/code in in bytes
+ """
+ if isinstance(key, int):
+ return DType.dtype_from_code(key)().itemsize
+ elif numpy.number in key.__mro__:
+ return key().itemsize
+ else:
+ raise ValueError
+
+ @staticmethod
+ def optimal_dtype(cardinality: Optional[int]) -> Type[numpy.number]:
+ """Get the dtype to use for an index of a certain cardinality
+
+ Args:
+ cardinality (Optional[int]): The number of elements to be indexed
+
+ Returns:
+ Type[numpy.number]: The dtype to use for the index
+ """
+ if cardinality is not None and cardinality < 65500:
+ return numpy.uint16
+ else:
+ return numpy.int32
+
+
+class _IndexWriter(object):
+ """Object class to write the index (.idx) file
+
+ Args:
+ idx_path (str): The path to the index file
+
+ dtype (Type[numpy.number]): The dtype of the index file
+ """
+
+ def __init__(self, idx_path: str, dtype: Type[numpy.number]) -> None:
+ self.idx_path = idx_path
+ self.dtype = dtype
+
+ def __enter__(self) -> "_IndexWriter":
+ """Enter the context introduced by the 'with' keyword
+
+ Returns:
+ _IndexWriter: The instance
+ """
+ self.idx_writer = open(self.idx_path, "wb")
+ # fixed, vestigial practice
+ self.idx_writer.write(_INDEX_HEADER)
+ # fixed, vestigial practice
+ self.idx_writer.write(struct.pack(" Optional[bool]:
+ """Exit the context introduced by the 'with' keyword
+
+ Args:
+ exc_type (Optional[Type[BaseException]]): Exception type
+
+ exc_val (Optional[BaseException]): Exception value
+
+ exc_tb (Optional[TracebackType]): Exception traceback object
+
+ Returns:
+ Optional[bool]: Whether to silence the exception
+ """
+ self.idx_writer.close()
+
+ def write(
+ self,
+ sequence_lengths: List[int],
+ sequence_modes: Optional[List[int]],
+ document_indices: List[int],
+ ) -> None:
+ """Write the index (.idx) file
+
+ Args:
+ sequence_lengths (List[int]): The length of each sequence
+
+ sequence_modes (Optional[List[int]]): The mode of each sequences
+
+ document_indices (List[int]): The seqyebce indices demarcating the end of each document
+ """
+ sequence_pointers = self._sequence_pointers(sequence_lengths)
+
+ # the number of sequences in the dataset
+ sequence_count = len(sequence_lengths)
+ self.idx_writer.write(struct.pack(" List[int]:
+ """Build the sequence pointers per the sequence lengths and dtype size
+
+ Args:
+ sequence_lengths (List[int]): The length of each sequence
+
+ Returns:
+ List[int]: The pointer to the beginning of each sequence
+ """
+ itemsize = DType.size(self.dtype)
+ curr_ptr = 0
+ list_ptr = []
+ for length in sequence_lengths:
+ list_ptr.append(curr_ptr)
+ curr_ptr += length * itemsize
+ return list_ptr
+
+
+class _IndexReader(object):
+ """Object class to read the index (.idx) file
+
+ Args:
+ idx_path (str): The path to the index file
+
+ multimodal (bool): Whether the dataset is multimodal
+ """
+
+ def __init__(self, idx_path: str, multimodal: bool) -> None:
+
+ log_single_rank(logger, logging.INFO, f"Load the {type(self).__name__} from {idx_path}")
+
+ with open(idx_path, "rb") as stream:
+ header = stream.read(9)
+ assert header == _INDEX_HEADER, f"bad header, cannot read: {idx_path}"
+
+ version = struct.unpack(" time elapsed: {t_end - t_beg:4f} seconds")
+
+ log_single_rank(logger, logging.INFO, f"\tExtract the sequence pointers")
+ t_beg = time.time()
+ self.sequence_pointers = numpy.frombuffer(
+ self.bin_buffer,
+ dtype=numpy.int64,
+ count=self.sequence_count,
+ offset=offset + self.sequence_lengths.nbytes,
+ )
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ log_single_rank(logger, logging.INFO, f"\tExtract the document indices")
+ t_beg = time.time()
+ self.document_indices = numpy.frombuffer(
+ self.bin_buffer,
+ dtype=numpy.int64,
+ count=self.document_count,
+ offset=offset + self.sequence_lengths.nbytes + self.sequence_pointers.nbytes,
+ )
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ self.sequence_modes = None
+ if multimodal:
+ log_single_rank(logger, logging.INFO, f"\tExtract the sequence modes")
+ t_beg = time.time()
+ self.sequence_modes = numpy.frombuffer(
+ self.bin_buffer,
+ dtype=numpy.int8,
+ count=self.sequence_count,
+ offset=offset
+ + self.sequence_lengths.nbytes
+ + self.sequence_pointers.nbytes
+ + self.document_indices.nbytes,
+ )
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ assert self.sequence_lengths.shape[0] == len(self)
+ assert self.sequence_lengths.shape[0] == self.sequence_count
+ assert self.sequence_lengths.shape[0] == self.document_indices[-1]
+
+ log_single_rank(logger, logging.INFO, f"> total number of sequences: {len(self)}")
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"> total number of documents: {self.document_indices.shape[0] - 1}",
+ )
+
+ def __del__(self) -> None:
+ """Clean up the object"""
+ if hasattr(self, "bin_buffer_mmap"):
+ self.bin_buffer_mmap._mmap.close()
+ del self.bin_buffer_mmap
+
+ def __len__(self) -> int:
+ """Return the length of the dataset
+
+ Returns:
+ int: The length of the dataset
+ """
+ return self.sequence_count
+
+ @lru_cache(maxsize=8)
+ def __getitem__(self, idx: int) -> Tuple[numpy.int32, numpy.int64, Optional[numpy.int8]]:
+ """Return the pointer, length, and mode at the index
+
+ Args:
+ idx (int): The index into the dataset
+
+ Returns:
+ Tuple[numpy.int32, numpy.int64, Optional[numpy.int8]]: The pointer, length and mode at the index
+ """
+ return (
+ self.sequence_pointers[idx],
+ self.sequence_lengths[idx],
+ self.sequence_modes[idx] if self.sequence_modes is not None else None,
+ )
+
+
+class _BinReader(ABC):
+ """Abstract class to read the data (.bin) file"""
+
+ @abstractmethod
+ def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray:
+ """Read bytes into a numpy array.
+
+ Args:
+ dtype (Type[numpy.number]): Data-type of the returned array.
+
+ count (int): Number of items to read.
+
+ offset (int): Start reading from this offset (in bytes).
+
+ Returns:
+ numpy.ndarray: An array with `count` items and data-type `dtype` constructed from reading bytes from the data file starting at `offset`.
+ """
+ pass
+
+
+class _MMapBinReader(_BinReader):
+ """A _BinReader that memory maps the data (.bin) file
+
+ Args:
+ bin_path (str): bin_path (str): The path to the data (.bin) file.
+ """
+
+ def __init__(self, bin_path: str) -> None:
+ self._bin_buffer_mmap = numpy.memmap(bin_path, mode="r", order="C")
+ self._bin_buffer = memoryview(self._bin_buffer_mmap)
+
+ def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray:
+ """Read bytes into a numpy array.
+
+ Args:
+ dtype (Type[numpy.number]): Data-type of the returned array.
+
+ count (int): Number of items to read.
+
+ offset (int): Start reading from this offset (in bytes).
+
+ Returns:
+ numpy.ndarray: An array with `count` items and data-type `dtype` constructed from reading bytes from the data file starting at `offset`.
+ """
+ return numpy.frombuffer(self._bin_buffer, dtype=dtype, count=count, offset=offset)
+
+ def __del__(self) -> None:
+ """Clean up the object."""
+ if self._bin_buffer_mmap is not None:
+ self._bin_buffer_mmap._mmap.close()
+ del self._bin_buffer_mmap
+
+
+class _FileBinReader(_BinReader):
+ """A _BinReader that reads from the data (.bin) file using a file pointer
+
+ Args:
+ bin_path (str): bin_path (str): The path to the data (.bin) file.
+ """
+
+ def __init__(self, bin_path: str) -> None:
+ self._bin_path = bin_path
+
+ def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray:
+ """Read bytes into a numpy array.
+
+ Args:
+ dtype (Type[numpy.number]): Data-type of the returned array.
+
+ count (int): Number of items to read.
+
+ offset (int): Start reading from this offset (in bytes).
+
+ Returns:
+ numpy.ndarray: An array with `count` items and data-type `dtype` constructed from reading bytes from the data file starting at `offset`.
+ """
+ sequence = numpy.empty(count, dtype=dtype)
+ with open(self._bin_path, mode='rb', buffering=0) as bin_buffer_file:
+ bin_buffer_file.seek(offset)
+ bin_buffer_file.readinto(sequence)
+ return sequence
+
+
+class _S3BinReader(_BinReader):
+ """A _BinReader that reads from the data (.bin) file from S3
+
+ Args:
+ bin_path (str): bin_path (str): The path to the data (.bin) file.
+
+ bin_chunk_nbytes (int, optional): If not None, then maintain an in-memory cache to speed up calls to the `read` method. Furthermore, on a cache miss, download this number of bytes to refresh the cache. Otherwise (None), do not maintain an in-memory cache. A class that inherits from _BinReader may not implement caching in which case it should assert that `bin_chunk_nbytes` is None at initialization.
+ """
+
+ def __init__(self, bin_path: str, bin_chunk_nbytes: int) -> None:
+ assert bin_chunk_nbytes > 0
+ self._client = boto3.client("s3")
+ self._s3_bucket, self._s3_key = parse_s3_path(bin_path)
+ self._cache = None
+ self._cache_bytes_start = None
+ self._cache_bytes_end = None
+ self._cache_nbytes = bin_chunk_nbytes
+
+ def _extract_from_cache(self, offset: int, size: int) -> bytes:
+ """Extract `size` bytes starting at `offset` bytes into the cache"""
+ start = offset - self._cache_bytes_start
+ assert start >= 0
+ end = start + size
+ assert end <= len(self._cache)
+ return self._cache[start:end]
+
+ def read(self, dtype: Type[numpy.number], count: int, offset: int) -> numpy.ndarray:
+ """Read bytes into a numpy array.
+
+ Let `size` be the `count` * `DType.size(dtype)`. If the requested span of bytes [`offset`,
+ `offset` + `size`) is covered by the in-memory cache maintained by this class, then this
+ function extracts the requested span from that cache and returns it. Otherwise, this
+ function first refreshes the cache and then extracts the requested span from the refreshed
+ cache and returns it.
+
+ The cache is refreshed based on `offset` and `size`. In particular, we divide all the bytes
+ in an S3 object into blocks, where each block contains `bin_chunk_nbytes` bytes. We assign
+ each block an index starting from 0. We take the block with index (`offset` //
+ `bin_chunk_nbytes`) to refresh the cache. If this new block still does not cover the
+ requested span, we extend it just enough to include `offset` + `size`.
+
+ Args:
+ dtype (Type[numpy.number]): Data-type of the returned array.
+
+ count (int): Number of items to read.
+
+ offset (int): Start reading from this offset (in bytes).
+
+ Returns:
+ numpy.ndarray: An array with `count` items and data-type `dtype` constructed from reading bytes from the data file starting at `offset`.
+ """
+ size = count * DType.size(dtype)
+ if (
+ self._cache is not None
+ and offset >= self._cache_bytes_start
+ and offset + size <= self._cache_bytes_end
+ ):
+ return numpy.frombuffer(self._extract_from_cache(offset, size), dtype=dtype)
+
+ bytes_start = (offset // self._cache_nbytes) * self._cache_nbytes
+ assert bytes_start >= 0
+ assert offset >= bytes_start
+ bytes_end = max(bytes_start + self._cache_nbytes, offset + size)
+ assert bytes_end >= 1
+ self._cache = self._client.get_object(
+ Bucket=self._s3_bucket,
+ Key=self._s3_key,
+ # Subtract 1, because the end of Range is inclusive.
+ Range=f'bytes={bytes_start}-{bytes_end-1}',
+ )['Body'].read()
+ self._cache_bytes_start = bytes_start
+ self._cache_bytes_end = bytes_end
+ return numpy.frombuffer(self._extract_from_cache(offset, size), dtype=dtype)
+
+ def __del__(self) -> None:
+ """Clean up the object"""
+ self._client.close()
+
+
+class IndexedDataset(torch.utils.data.Dataset):
+ """The low-level interface dataset class
+
+ Args:
+ path_prefix (str): The index (.idx) and data (.bin) prefix
+
+ multimodal (bool): Whether the dataset is multimodal. Defaults to False.
+
+ mmap (bool): Whether to mmap the .bin files. Defaults to True.
+
+ s3_config (Optional[S3Config]): Supplied only for data stored on S3. IndexedDataset downloads the index (.idx) file to `s3_config.path_to_idx_cache` and streams data from the data (.bin) file in `s3_config.bin_chunk_nbytes` blocks. Note that `mmap` must be disabled for S3 data loading. Defaults to None.
+ """
+
+ def __init__(
+ self,
+ path_prefix: str,
+ multimodal: bool = False,
+ mmap: bool = True,
+ s3_config: Optional[S3Config] = None,
+ ) -> None:
+ super().__init__()
+ self.path_prefix = None
+ self.multimodal = None
+ self.mmap = None
+ self.s3_config = None
+
+ self.index = None
+ self.bin_reader = None
+
+ if is_s3_path(path_prefix) and s3_config is not None:
+ idx_path = get_idx_path(path_prefix)
+ cache_idx_path = os.path.join(s3_config.path_to_idx_cache, os.path.basename(idx_path))
+ maybe_download_file(idx_path, cache_idx_path)
+
+ self.initialize(path_prefix, multimodal, mmap, s3_config)
+
+ def initialize(
+ self, path_prefix: str, multimodal: bool, mmap: bool, s3_config: Optional[S3Config]
+ ) -> None:
+ """Initialize the dataset
+
+ This method is called by IndexedDataset.__init__ during object creation and by
+ IndexedDataset.__setstate__ during un-pickling
+
+ Args:
+ path_prefix (str): The index (.idx) and data (.bin) prefix
+
+ multimodal (bool): Whether the dataset is multimodal
+
+ mmap (bool): Whether to mmap the .bin file
+
+ s3_config (Optional[S3Config]): See IndexedDataset docstring for details.
+ """
+ idx_path = get_idx_path(path_prefix)
+ bin_path = get_bin_path(path_prefix)
+ if s3_config is None:
+ assert os.path.exists(idx_path) and os.path.exists(
+ bin_path
+ ), f"One or both of the .idx and .bin files cannot be found at the path prefix {path_prefix}"
+ self.path_prefix = path_prefix
+ self.multimodal = multimodal
+ self.mmap = mmap
+ self.s3_config = s3_config
+ if mmap:
+ assert not s3_config
+ self.bin_reader = _MMapBinReader(bin_path)
+ elif s3_config:
+ assert not mmap
+ self.bin_reader = _S3BinReader(bin_path, s3_config.bin_chunk_nbytes)
+ idx_path = os.path.join(
+ s3_config.path_to_idx_cache, os.path.basename(get_idx_path(path_prefix))
+ )
+ else:
+ self.bin_reader = _FileBinReader(bin_path)
+ self.index = _IndexReader(idx_path, self.multimodal)
+
+ def __getstate__(self) -> Tuple[str, bool, bool, Optional[S3Config]]:
+ """Get the state during pickling
+
+ Returns:
+ Tuple[str, bool, bool, Optional[S3Config]]: The state tuple
+ """
+ return self.path_prefix, self.multimodal, self.mmap, self.s3_config
+
+ def __setstate__(self, state: Tuple[str, bool, bool, Optional[S3Config]]) -> None:
+ """Set the state during un-pickling
+
+ Args:
+ state (Tuple[str, bool, bool, Optional[S3Config]]): The state tuple
+ """
+ path_prefix, multimodal, mmap, s3_config = state
+ self.initialize(path_prefix, multimodal, mmap, s3_config)
+
+ def __del__(self) -> None:
+ """Clean up the object"""
+ del self.bin_reader
+ del self.index
+
+ def __len__(self) -> int:
+ """Return the length of the dataset i.e. the number of sequences in the index
+
+ Returns:
+ int: The length of the dataset
+ """
+ return len(self.index)
+
+ def __getitem__(
+ self, idx: Union[int, numpy.integer, slice]
+ ) -> Union[numpy.ndarray, Tuple[numpy.ndarray, numpy.ndarray]]:
+ """Return from the dataset
+
+ Args:
+ idx (Union[int, numpy.integer, slice]): The index or index slice into the dataset
+
+ Raises:
+ ValueError: When the index slice is non-contiguous
+
+ TypeError: When the index is of an unexpected type
+
+ Returns:
+ Union[numpy.ndarray, Tuple[numpy.ndarray, numpy.ndarray]]: The sequence tokens and modes at the index or index slice
+ """
+ if isinstance(idx, (int, numpy.integer)):
+ sequence_pointer, sequence_length, sequence_mode = self.index[idx]
+ sequence = self.bin_reader.read(
+ dtype=self.index.dtype, count=sequence_length, offset=sequence_pointer
+ )
+ return (sequence, sequence_mode) if sequence_mode is not None else sequence
+ elif isinstance(idx, slice):
+ start, stop, step = idx.indices(len(self))
+ if step != 1:
+ raise ValueError("Slices into indexed_dataset must be contiguous")
+ sequence_lengths = self.index.sequence_lengths[idx]
+ sequence_modes = self.index.sequence_modes[idx] if self.multimodal else None
+ sequence_offsets = list(accumulate(sequence_lengths))
+ sequences = numpy.split(
+ self.bin_reader.read(
+ dtype=self.index.dtype,
+ count=sum(sequence_lengths),
+ offset=self.index.sequence_pointers[start],
+ ),
+ sequence_offsets[:-1],
+ )
+ return (sequences, sequence_modes) if sequence_modes is not None else sequences
+ else:
+ raise TypeError("Unexpected type received for idx: {}".format(type(idx)))
+
+ def get(self, idx: int, offset: int = 0, length: Optional[int] = None) -> numpy.ndarray:
+ """Retrieve a single item from the dataset with the option to only
+ return a portion of the item.
+
+ get(idx) is the same as [idx] but get() does not support slicing.
+
+ Args:
+ idx (Union[int, numpy.integer]): The index into the dataset
+
+ offset (int): The integer token offset in the sequence
+
+ length (int): The number of tokens to grab from the sequence
+
+ Returns:
+ Union[numpy.ndarray, Tuple[numpy.ndarray, numpy.ndarray]]: The sequence tokens and modes at the index
+ """
+ sequence_pointer, sequence_length, sequence_mode = self.index[idx]
+ if length is None:
+ length = sequence_length - offset
+ sequence_pointer += offset * DType.size(self.index.dtype)
+ sequence = self.bin_reader.read(
+ dtype=self.index.dtype, count=length, offset=sequence_pointer
+ )
+ return (sequence, sequence_mode) if sequence_mode is not None else sequence
+
+ @property
+ def sequence_lengths(self) -> numpy.ndarray:
+ """Get the sequence lengths
+
+ Returns:
+ numpy.ndarray: The sequence lengths
+ """
+ return self.index.sequence_lengths
+
+ @property
+ def document_indices(self) -> numpy.ndarray:
+ """Get the document indices
+
+ Returns:
+ numpy.ndarray: The document indices
+ """
+ return self.index.document_indices
+
+ def get_document_indices(self) -> numpy.ndarray:
+ """Get the document indices
+
+ This method is slated for deprecation.
+
+ Returns:
+ numpy.ndarray: The document indices
+ """
+ return self.index.document_indices
+
+ def set_document_indices(self, document_indices: numpy.ndarray) -> None:
+ """Set the document indices
+
+ This method is slated for deprecation.
+
+ Args:
+ document_indices (numpy.ndarray): The document indices
+ """
+ self.index.document_indices = document_indices
+
+ @property
+ def sequence_modes(self) -> numpy.ndarray:
+ """Get the sequence modes
+
+ Returns:
+ numpy.ndarray: The sequence modes
+ """
+ return self.index.sequence_modes
+
+ @staticmethod
+ def exists(path_prefix: str) -> bool:
+ """Return whether the IndexedDataset exists on disk at the prefix
+
+ Args:
+ path_prefix (str): The prefix to the index (.idx) and data (.bin) files
+
+ Returns:
+ bool: Whether the IndexedDataset exists on disk at the prefix
+ """
+ if is_s3_path(path_prefix):
+ s3_client = boto3.client("s3")
+ return object_exists(s3_client, get_idx_path(path_prefix)) and object_exists(
+ s3_client, get_bin_path(path_prefix)
+ )
+ return os.path.exists(get_idx_path(path_prefix)) and os.path.exists(
+ get_bin_path(path_prefix)
+ )
+
+
+class IndexedDatasetBuilder(object):
+ """Builder class for the IndexedDataset class
+
+ Args:
+ bin_path (str): The path to the data (.bin) file
+
+ dtype (Type[numpy.number], optional): The dtype of the index file. Defaults to numpy.int32.
+
+ multimodal (bool, optional): Whether the dataset is multimodal. Defaults to False.
+ """
+
+ def __init__(
+ self, bin_path: str, dtype: Type[numpy.number] = numpy.int32, multimodal: bool = False
+ ) -> None:
+ self.data_file = open(bin_path, "wb")
+ self.dtype = dtype
+ self.multimodal = multimodal
+
+ self.sequence_lengths = []
+ self.document_indices = [0]
+ self.sequence_modes = [] if self.multimodal else None
+
+ def add_item(self, tensor: torch.Tensor, mode: int = 0) -> None:
+ """Add a single item to the dataset
+
+ Args:
+ tensor (torch.Tensor): The item to add to the data file
+
+ mode (int, optional): The mode for the item. Defaults to 0.
+ """
+ np_array = numpy.array(tensor.numpy(), dtype=self.dtype)
+ self.data_file.write(np_array.tobytes(order="C"))
+ self.sequence_lengths.append(np_array.size)
+ if self.multimodal:
+ self.sequence_modes.append(mode)
+
+ def add_document(
+ self, tensor: torch.Tensor, lengths: List[int], modes: Optional[List[int]] = None
+ ) -> None:
+ """Add an entire document to the dataset
+
+ Args:
+ tensor (torch.Tensor): The document to add
+
+ lengths (List[int]): The lengths of each item in the document
+
+ modes (Optional[List[int]], optional): The modes for each item in the document. Defaults to None.
+ """
+ np_array = numpy.array(tensor, dtype=self.dtype)
+ self.data_file.write(np_array.tobytes(order="C"))
+ self.sequence_lengths.extend(lengths)
+ self.document_indices.append(len(self.sequence_lengths))
+ if self.multimodal:
+ self.sequence_modes.extend(modes if modes is not None else [0] * lengths)
+
+ def end_document(self) -> None:
+ """Finalize the document, for use with IndexedDatasetBuilder.add_item"""
+ self.document_indices.append(len(self.sequence_lengths))
+
+ def add_index(self, path_prefix: str) -> None:
+ """Add an entire IndexedDataset to the dataset
+
+ Args:
+ path_prefix (str): The index (.idx) and data (.bin) prefix
+ """
+ # Concatenate index
+ index = _IndexReader(get_idx_path(path_prefix), multimodal=self.multimodal)
+ assert index.dtype == self.dtype
+
+ offset = len(self.sequence_lengths)
+ self.sequence_lengths.extend(index.sequence_lengths)
+ self.document_indices.extend((offset + index.document_indices)[1:])
+
+ if self.multimodal:
+ self.sequence_modes.extend(index.sequence_modes)
+
+ # Concatenate data
+ with open(get_bin_path(path_prefix), "rb") as f:
+ shutil.copyfileobj(f, self.data_file)
+
+ def finalize(self, idx_path: str) -> None:
+ """Clean up and write the index (.idx) file
+
+ Args:
+ idx_path (str): The path to the index file
+ """
+ self.data_file.close()
+ with _IndexWriter(idx_path, self.dtype) as writer:
+ writer.write(self.sequence_lengths, self.sequence_modes, self.document_indices)
+
+
+def get_idx_path(path_prefix: str) -> str:
+ """Get the path to the index file from the prefix
+
+ Args:
+ path_prefix (str): The prefix
+
+ Returns:
+ str: The path to the index file
+ """
+ return path_prefix + ".idx"
+
+
+def get_bin_path(path_prefix: str) -> str:
+ """Get the path to the data file from the prefix
+
+ Args:
+ path_prefix (str): The prefix
+
+ Returns:
+ str: The path to the data file
+ """
+ return path_prefix + ".bin"
diff --git a/megatron/core/datasets/masked_dataset.py b/megatron/core/datasets/masked_dataset.py
new file mode 100644
index 0000000000..9db6c67eb1
--- /dev/null
+++ b/megatron/core/datasets/masked_dataset.py
@@ -0,0 +1,423 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+import os
+import time
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import numpy
+import torch
+
+from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
+from megatron.core.datasets.indexed_dataset import IndexedDataset
+from megatron.core.datasets.megatron_dataset import MegatronDataset
+from megatron.core.datasets.utils import Split
+from megatron.core.utils import log_single_rank
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MaskedWordPieceDatasetConfig(BlendedMegatronDatasetConfig):
+ """Configuration object for Megatron Core Masked WordPiece datasets"""
+
+ masking_probability: float = None
+ """The probability we mask a candidate N-gram"""
+
+ short_sequence_probability: float = None
+ """The probability we return a sequence shorter than the target sequence length"""
+
+ masking_max_ngram: int = None
+ """The maximum length N-gram to consider masking or permuting"""
+
+ masking_do_full_word: bool = None
+ """Whether we mask the the whole word or its component parts"""
+
+ masking_do_permutation: bool = None
+ """Whether we shuffle a subset of candidate N-grams in addition"""
+
+ masking_use_longer_ngrams: bool = None
+ """Whether to favor longer N-grams over shorter N-grams"""
+
+ masking_use_geometric_distribution: bool = None
+ """Whether to draw the size of the N-gram from a geometric distribution according to SpanBERT
+ https://arxiv.org/abs/1907.10529 (Section 3.1)
+ """
+
+ def __post_init__(self) -> None:
+ """Do asserts and set fields post init"""
+ super().__post_init__()
+
+ assert self.tokenizer is not None
+
+ assert self.masking_probability is not None
+ assert self.short_sequence_probability is not None
+ assert self.masking_max_ngram is not None
+ assert self.masking_do_full_word is not None
+ assert self.masking_do_permutation is not None
+ assert self.masking_use_longer_ngrams is not None
+ assert self.masking_use_geometric_distribution is not None
+
+ assert self.masking_probability > 0 and self.masking_probability < 1.0
+ assert self.short_sequence_probability >= 0 and self.short_sequence_probability <= 1.0
+ assert self.masking_max_ngram > 0
+ assert not (self.masking_use_geometric_distribution and self.masking_do_permutation)
+
+ if self.masking_use_geometric_distribution and self.masking_use_longer_ngrams:
+ log_single_rank(
+ logger,
+ logging.WARNING,
+ "The use of a geometric distribution overrides the default distribution",
+ )
+
+
+class MaskedWordPieceDataset(MegatronDataset):
+ """The semi-abstract base class for masked WordPiece datasets
+
+ This implementation makes the rigid assumption that all inheritor datasets are built upon the
+ IndexedDataset class. This assumption may be pushed down to the inheritors in future if
+ necessary.
+
+ NB: WordPiece tokenization prepends a double hash "##" to all tokens/pieces in a word, save the
+ first token/piece.
+
+ Args:
+ indexed_dataset (IndexedDataset): The IndexedDataset around which to build the MegatronDataset
+
+ dataset_path (str): The real path on disk to the dataset, for bookkeeping
+
+ indexed_indices (numpy.ndarray): The set of the documents indices to expose
+
+ num_samples (Optional[int]): The number of samples to draw from the indexed dataset. When None, build as many samples as correspond to one epoch.
+
+ index_split (Split): The indexed_indices Split
+
+ config (MaskedWordPieceDatasetConfig): The config
+ """
+
+ def __init__(
+ self,
+ indexed_dataset: IndexedDataset,
+ dataset_path: str,
+ indexed_indices: numpy.ndarray,
+ num_samples: Optional[int],
+ index_split: Split,
+ config: MaskedWordPieceDatasetConfig,
+ ) -> None:
+ super().__init__(
+ indexed_dataset, dataset_path, indexed_indices, num_samples, index_split, config
+ )
+
+ @staticmethod
+ def numel_low_level_dataset(low_level_dataset: IndexedDataset) -> int:
+ return low_level_dataset.document_indices.shape[0] - 1
+
+ @staticmethod
+ def build_low_level_dataset(
+ dataset_path: str, config: MaskedWordPieceDatasetConfig
+ ) -> IndexedDataset:
+ return IndexedDataset(dataset_path)
+
+ @staticmethod
+ def _key_config_attributes() -> List[str]:
+ """Inherited method implementation
+
+ Returns:
+ List[str]: The key config attributes
+ """
+ return super(MaskedWordPieceDataset, MaskedWordPieceDataset)._key_config_attributes() + [
+ "masking_probability",
+ "short_sequence_probability",
+ "masking_max_ngram",
+ "masking_do_full_word",
+ "masking_do_permutation",
+ "masking_use_longer_ngrams",
+ "masking_use_geometric_distribution",
+ ]
+
+ def __len__(self) -> int:
+ return self.sample_index.shape[0]
+
+ def _build_sample_index(
+ self, sequence_length: int, min_sentences_per_sample: int
+ ) -> numpy.ndarray:
+ path_to_cache = self.config.path_to_cache
+ if path_to_cache is None:
+ path_to_cache = os.path.join(
+ self.dataset.path_prefix, "cache", f"{type(self).__name__}_indices"
+ )
+
+ get_path_to = lambda suffix: os.path.join(
+ path_to_cache, f"{self.unique_description_hash}-{type(self).__name__}-{suffix}"
+ )
+ path_to_description = get_path_to("description.txt")
+ path_to_sample_index = get_path_to("sample_index.npy")
+ cache_hit = all(map(os.path.isfile, [path_to_description, path_to_sample_index]))
+
+ if self.num_samples is not None:
+ num_epochs = numpy.iinfo(numpy.int32).max - 1
+ else:
+ num_epochs = 1
+
+ if not cache_hit and torch.distributed.get_rank() == 0:
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Build and save the {type(self).__name__} {self.index_split.name} indices",
+ )
+ self.built_anew_on_cache_miss = True
+
+ os.makedirs(path_to_cache, exist_ok=True)
+
+ # Write the description
+ with open(path_to_description, "wt") as writer:
+ writer.write(self.unique_description)
+
+ # Build the sample index
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"\tBuild and save the sample index to {os.path.basename(path_to_sample_index)}",
+ )
+ t_beg = time.time()
+ from megatron.core.datasets import helpers
+
+ # Add +1 for access to document upper bound
+ indices = numpy.append(self.indices, self.indices[-1] + 1)
+
+ sample_index = helpers.build_mapping(
+ self.dataset.document_indices[indices],
+ self.dataset.sequence_lengths,
+ num_epochs,
+ self.num_samples,
+ sequence_length,
+ self.config.short_sequence_probability,
+ self.config.random_seed,
+ False,
+ min_sentences_per_sample,
+ )
+ numpy.save(path_to_sample_index, sample_index, allow_pickle=True)
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ log_single_rank(
+ logger, logging.INFO, f"> total number of samples: {sample_index.shape[0]}"
+ )
+ log_single_rank(logger, logging.INFO, f"> total number of epochs: {num_epochs}")
+
+ return sample_index
+
+ log_single_rank(
+ logger, logging.INFO, f"Load the {type(self).__name__} {self.index_split.name} indices"
+ )
+
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"\tLoad the sample index from {os.path.basename(path_to_sample_index)}",
+ )
+ t_beg = time.time()
+ sample_index = numpy.load(path_to_sample_index, allow_pickle=True, mmap_mode="r")
+ t_end = time.time()
+ log_single_rank(logger, logging.DEBUG, f"\t> time elapsed: {t_end - t_beg:4f} seconds")
+
+ return sample_index
+
+ def _create_masked_lm_predictions(
+ self,
+ token_ids: List[int],
+ target_sequence_length: int,
+ numpy_random_state: numpy.random.RandomState,
+ ) -> Tuple[List[int], List[int], List[int], List[int], List[Tuple[List[int], List[int]]]]:
+ """Creates the predictions for the masked LM objective
+
+ Args:
+ token_ids (List[int]): The token ids
+ target_sequence_length (int): The target sequence length
+ numpy_random_state (numpy.random.RandomState): The NumPy random state
+
+ Returns:
+ Tuple[List[int], List[int], List[int], List[int], List[Tuple[List[int], List[int]]]]:
+ 1. masked_token_ids -> The masked sequence
+ 2. masked_positions -> The indices for the masked token ids
+ 3. masked_labels -> The original token ids for the masked token ids
+ 4. boundaries -> The sentence and word boundaries for the sequence
+ 4. masked_spans -> The masked positions and labels with N-gram info intact
+ """
+ # Build the token sentence and word boundaries and the masking candidates
+ # e.g. [cls, id, ##id, ##id, id, ##id, sep, id, ##id, sep]
+ # -> boundaries: [1, 1, 0, 0, 1, 0, 1, 1, 0, 1]
+ # -> candidates with whole word masking: [[1, 2, 3], [4, 5], [7, 8]]
+ # -> candidates sans whole word masking: [[1], [2], [3], [4], [5], [7], [8]]
+ boundaries = []
+ candidates = []
+ for i, token_id in enumerate(token_ids):
+ if token_id == self.config.tokenizer.cls or token_id == self.config.tokenizer.sep:
+ boundaries.append(1)
+ else:
+ if not self.config.tokenizer.inv_vocab[token_id].startswith("##"):
+ boundaries.append(1)
+ candidates.append([i])
+ else:
+ boundaries.append(0)
+ if self.config.masking_do_full_word and len(candidates) > 0:
+ candidates[-1].append(i)
+ else:
+ candidates.append([i])
+
+ n_maskings = min(
+ self.config.masking_probability * target_sequence_length,
+ max(1, int(round(len(token_ids) * self.config.masking_probability))),
+ )
+
+ ngram_nvals = numpy.arange(self.config.masking_max_ngram, dtype=numpy.int64) + 1
+
+ # By default, the N-gram probabilites are inversely proportional to N
+ # e.g. N = 3
+ # -> P = array([0.54545455, 0.27272727, 0.18181818])
+ nprobs = 1.0 / ngram_nvals
+ nprobs = nprobs / nprobs.sum(keepdims=True)
+ if self.config.masking_use_longer_ngrams:
+ nprobs = nprobs[::-1]
+
+ # Create a nested list of depth 3
+ # layer 1: the candidate dimension
+ # layer 2: the N-gram dimension
+ # layer 3: the token dimension
+ candidate_ngrams = [
+ [candidates[idx : idx + n] for n in ngram_nvals] for idx in range(len(candidates))
+ ]
+ numpy_random_state.shuffle(candidate_ngrams)
+
+ masked_token_ids = list(token_ids)
+ masked_positions_and_labels = []
+ masked_spans = []
+ masked_indices = set()
+ for candidate_idx in range(len(candidate_ngrams)):
+ n_ngrams = len(candidate_ngrams[candidate_idx])
+
+ # Stop when we hit our desired number of maskings
+ if len(masked_positions_and_labels) >= n_maskings:
+ break
+
+ # Do nothing for candidates with no ngrams
+ if not candidate_ngrams[candidate_idx]:
+ continue
+
+ # Choose the initial value of N
+ if self.config.masking_use_geometric_distribution:
+ # Sample N from a geometric distribution with p = 0.2 and clip
+ # i.e. SpanBERT
+ # -> https://arxiv.org/abs/1907.10529 (Section 3.1)
+ p = 0.2
+ n = min(numpy_random_state.geometric(p), self.config.masking_max_ngram)
+ else:
+ p = nprobs[:n_ngrams] / nprobs[:n_ngrams].sum(keepdims=True)
+ n = numpy_random_state.choice(ngram_nvals[:n_ngrams], p=p)
+
+ while True:
+ ngram_indices = sum(candidate_ngrams[candidate_idx][n - 1], [])
+ n = n - 1
+ # Success: masking this N-gram puts us below the desired number of maskings
+ if n_maskings >= len(masked_positions_and_labels) + len(ngram_indices):
+ skip_candidate = False
+ break
+ # Failure: no N-grams remain for this candidate
+ if n == 0:
+ skip_candidate = True
+ break
+
+ # Do nothing for candidates whose 1-gram is too long
+ if skip_candidate:
+ continue
+
+ # Do nothing for candidate indices which have already been masked
+ if any(map(lambda idx: idx in masked_indices, ngram_indices)):
+ continue
+
+ # Mask the tokens and record their original positions and values
+ for index in ngram_indices:
+ masked_indices.add(index)
+ mask = self._get_token_mask(numpy_random_state)
+ if mask is None:
+ masked_token_ids[index] = token_ids[index]
+ else:
+ masked_token_ids[index] = mask
+ masked_positions_and_labels.append((index, token_ids[index]))
+
+ masked_spans.append((ngram_indices, [token_ids[index] for index in ngram_indices]))
+
+ assert len(masked_positions_and_labels) <= n_maskings
+
+ numpy_random_state.shuffle(candidate_ngrams)
+
+ if self.config.masking_do_permutation:
+
+ n_swappings = n_maskings
+
+ permuted_indices = set()
+ for candidate_idx in range(len(candidate_ngrams)):
+ n_ngrams = len(candidate_ngrams[candidate_idx])
+
+ if len(permuted_indices) >= n_swappings:
+ break
+
+ # Do nothing for candidates with no ngrams
+ if not candidate_ngrams[candidate_idx]:
+ continue
+
+ p = nprobs[:n_ngrams] / nprobs[:n_ngrams].sum(keepdims=True)
+ n = numpy.random.choice(ngram_nvals[:n_ngrams], p=p)
+
+ while True:
+ ngram_indices = sum(candidate_ngrams[candidate_idx][n - 1], [])
+ n = n - 1
+ # Success: swapping this N-gram puts us below the desired number of swappings
+ if n_swappings >= len(permuted_indices) + len(ngram_indices):
+ skip_candidate = False
+ break
+ # Failure: no N-grams remain for this candidate
+ if n == 0:
+ skip_candidate = True
+ break
+
+ # Do nothing for candidates whose 1-gram is too long
+ if skip_candidate:
+ continue
+
+ # Do nothing for candidate indices which have already been masked or permuted
+ if any(
+ map(lambda idx: idx in masked_indices or idx in permuted_indices, ngram_indices)
+ ):
+ continue
+
+ for index in ngram_indices:
+ permuted_indices.add(index)
+
+ assert len(permuted_indices) <= n_swappings
+
+ permuted_indices = sorted(permuted_indices)
+ permuted_indices_copy = list(permuted_indices)
+ numpy_random_state.shuffle(permuted_indices_copy)
+ masked_token_ids_copy = list(masked_token_ids)
+
+ for idx, idx_copy in zip(permuted_indices, permuted_indices_copy):
+ masked_token_ids[idx] = masked_token_ids_copy[idx_copy]
+ masked_positions_and_labels.append((idx, masked_token_ids_copy[idx]))
+
+ masked_positions_and_labels = sorted(masked_positions_and_labels, key=lambda x: x[0])
+ masked_positions = []
+ masked_labels = []
+ for position, label in masked_positions_and_labels:
+ masked_positions.append(position)
+ masked_labels.append(label)
+
+ masked_spans = sorted(masked_spans, key=lambda x: x[0][0])
+
+ return masked_token_ids, masked_positions, masked_labels, boundaries, masked_spans
+
+ @abstractmethod
+ def _get_token_mask(self, numpy_random_state: numpy.random.RandomState) -> Optional[int]:
+ pass
diff --git a/megatron/core/datasets/megatron_dataset.py b/megatron/core/datasets/megatron_dataset.py
new file mode 100644
index 0000000000..15a9a53328
--- /dev/null
+++ b/megatron/core/datasets/megatron_dataset.py
@@ -0,0 +1,139 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import hashlib
+import json
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+import numpy
+import torch
+
+from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
+from megatron.core.datasets.indexed_dataset import IndexedDataset
+from megatron.core.datasets.utils import Split
+
+LowLevelDataset = Union[IndexedDataset, Iterable]
+
+
+class MegatronDataset(ABC, torch.utils.data.Dataset):
+ """The highest level wrapper class from which all dataset classes should inherit
+
+ Args:
+ dataset (LowLevelDataset): The dataset around which to build the MegatronDataset
+
+ dataset_path (Optional[str]): The real path on disk to the dataset, for bookkeeping
+
+ indices (numpy.ndarray): The set of the documents indices to expose
+
+ num_samples (Optional[int]): The minimum number of samples to build from the indexed dataset. When None, build as many samples as correspond to one epoch.
+
+ index_split (Split): The indices Split
+
+ config (BlendedMegatronDatasetConfig): The config
+ """
+
+ def __init__(
+ self,
+ dataset: LowLevelDataset,
+ dataset_path: Optional[str],
+ indices: numpy.ndarray,
+ num_samples: Optional[int],
+ index_split: Split,
+ config: BlendedMegatronDatasetConfig,
+ ) -> None:
+ self.dataset = dataset
+ self.dataset_path = dataset_path
+ self.indices = indices
+ self.num_samples = num_samples
+ self.index_split = index_split
+ self.config = config
+
+ self.unique_identifiers = OrderedDict()
+
+ self.unique_identifiers["class"] = type(self).__name__
+ self.unique_identifiers["dataset_path"] = self.dataset_path
+ self.unique_identifiers["num_samples"] = self.num_samples
+ self.unique_identifiers["index_split"] = self.index_split.name
+ for attr in self._key_config_attributes():
+ self.unique_identifiers[attr] = getattr(self.config, attr)
+
+ self.unique_description = json.dumps(
+ self.unique_identifiers, indent=4, default=lambda obj: obj.unique_identifiers
+ )
+ self.unique_description_hash = hashlib.md5(
+ self.unique_description.encode("utf-8")
+ ).hexdigest()
+
+ self.built_anew_on_cache_miss = False
+
+ @staticmethod
+ def numel_low_level_dataset(low_level_dataset: LowLevelDataset) -> int:
+ """Return the number of elements in the underlying low level dataset for the purpose of
+ segregating the train/valid/test split indices
+
+ It may be that the low level dataset can be split any number of ways, depending on the mid
+ level dataset it supports, which is why we define the "number of elements" function
+ separately from the __len__ function here in the mid level dataset class
+
+ Args:
+ low_level_dataset (LowLevelDataset): The underlying low level dataset
+
+ Returns:
+ int: The number of elements in the underlying low level dataset
+ """
+ raise NotImplementedError
+
+ @staticmethod
+ def build_low_level_dataset(
+ dataset_path: str, config: BlendedMegatronDatasetConfig
+ ) -> LowLevelDataset:
+ """Build the low level dataset via a function to be called from within
+ BlendedMegatronDatasetBuilder.build_generic_dataset
+
+ It may be that the low level dataset spans any subset of train/valid/test splits, which is
+ why we define a static "build" function separately from the constructor in the mid level
+ dataset class
+
+ Args:
+ dataset_path (str): The real path on disk to the dataset
+
+ config (BlendedMegatronDatasetConfig): The dataset config
+
+ Returns:
+ LowLevelDataset: The low level dataset
+ """
+ raise NotImplementedError
+
+ @staticmethod
+ def _key_config_attributes() -> List[str]:
+ """Return all config attributes which contribute to uniquely identifying the dataset.
+
+ These attributes will be used to build a uniquely identifying string and MD5 hash which
+ will be used to cache/load dataset resources from run to run.
+
+ Returns:
+ List[str]: The key config attributes
+ """
+ return ["random_seed", "sequence_length", "split", "split_matrix", "tokenizer"]
+
+ @abstractmethod
+ def __len__(self) -> int:
+ """Return the length of the dataset
+
+ Returns:
+ int: See abstract implementation
+ """
+ pass
+
+ @abstractmethod
+ def __getitem__(self, idx: int) -> Dict[str, Union[torch.Tensor, numpy.ndarray]]:
+ """Return from the dataset
+
+ Args:
+ idx (int): The index into the dataset
+
+ Returns:
+ Dict[str, Union[torch.Tensor, numpy.ndarray]]: See abstract implementation
+ """
+ pass
diff --git a/megatron/core/datasets/megatron_tokenizer.py b/megatron/core/datasets/megatron_tokenizer.py
new file mode 100644
index 0000000000..84f3546cf3
--- /dev/null
+++ b/megatron/core/datasets/megatron_tokenizer.py
@@ -0,0 +1,154 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import json
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+from typing import Any
+
+import numpy
+
+
+class MegatronTokenizer(ABC):
+ """Abstract class for tokenizer
+
+ Absent a config or class-specific tracking of which objects are uniquely identifying, we must
+ include all key word arguments as unique identifiers
+
+ Args:
+ tokenizer_paths (Tuple[str]): All tokenizer source paths or prefixes
+
+ tokenizer_options (Dict[str, Any]): All tokenizer options
+ """
+
+ def __init__(self, *tokenizer_paths: str, **tokenizer_options: Any):
+
+ self.unique_identifiers = OrderedDict()
+ self.unique_identifiers["class"] = type(self).__name__
+ self.unique_identifiers["tokenizer_path"] = list(tokenizer_paths)
+ for option in tokenizer_options:
+ self.unique_identifiers[option] = str(tokenizer_options[option])
+
+ self.unique_description = json.dumps(self.unique_identifiers, indent=4)
+
+ super().__init__()
+
+ @abstractmethod
+ def tokenize(self, text: str) -> numpy.ndarray:
+ """Convert text to embedding ids
+
+ Args:
+ text (str): The text to convert
+
+ Returns:
+ numpy.ndarray: The converted embedding ids
+ """
+ pass
+
+ def detokenize(self, ids: numpy.ndarray) -> str:
+ """Convert embedding ids to text
+
+ Args:
+ ids (numpy.ndarray): The ids to convert
+
+ Returns:
+ str: The converted text
+
+ Raises:
+ NotImplementedError: Non-abstract, optional method
+ """
+ raise NotImplementedError("{} has no method 'detokenize'".format(type(self).__name__))
+
+ def offsets(self, ids: list[int], text: str) -> list[int]:
+ """Convert embedding ids to text offsets
+
+ Args:
+ ids (list[int]): The ids to convert
+ text (str): The text to convert
+
+ Returns:
+ list[int]: The converted offsets
+
+ Raises:
+ NotImplementedError: Non-abstract, optional method
+ """
+ raise NotImplementedError("{} has no method 'offsets'".format(type(self).__name__))
+
+ @property
+ @abstractmethod
+ def vocab(self):
+ """Dictionary from vocab text token to id token"""
+ pass
+
+ @property
+ @abstractmethod
+ def inv_vocab(self):
+ """Dictionary from vocab id token to text token"""
+ pass
+
+ @property
+ @abstractmethod
+ def vocab_size(self):
+ """The vocabulary size"""
+ pass
+
+ @property
+ def cls(self):
+ """The CLS token id
+
+ Raises:
+ NotImplementedError: Non-abstract, optional attribute
+ """
+ raise NotImplementedError("{} has no attribute 'cls'".format(type(self).__name__))
+
+ @property
+ def sep(self):
+ """The SEP token id
+
+ Raises:
+ NotImplementedError: Non-abstract, optional attribute
+ """
+ raise NotImplementedError("{} has no attribute 'sep'".format(type(self).__name__))
+
+ @property
+ def pad(self):
+ """The PAD token id
+
+ Raises:
+ NotImplementedError: Non-abstract, optional attribute
+ """
+ raise NotImplementedError("{} has no attribute 'pad'".format(type(self).__name__))
+
+ @property
+ def eod(self):
+ """The EOD token id
+
+ Raises:
+ NotImplementedError: Non-abstract, optional attribute
+ """
+ raise NotImplementedError("{} has no attribute 'eod'".format(type(self).__name__))
+
+ @property
+ def bos(self):
+ """The BOS token id
+
+ Raises:
+ NotImplementedError: Non-abstract, optional attribute
+ """
+ raise NotImplementedError("{} has no attribute 'bos'".format(type(self).__name__))
+
+ @property
+ def eos(self):
+ """The EOS token id
+
+ Raises:
+ NotImplementedError: Non-abstract, optional attribute
+ """
+ raise NotImplementedError("{} has no attribute 'eos'".format(type(self).__name__))
+
+ @property
+ def mask(self):
+ """The MASK token id
+
+ Raises:
+ NotImplementedError: Non-abstract, optional attribute
+ """
+ raise NotImplementedError("{} has no attribute 'mask'".format(type(self).__name__))
diff --git a/megatron/core/datasets/multimodal_dataset.py b/megatron/core/datasets/multimodal_dataset.py
new file mode 100644
index 0000000000..0a3e93a15b
--- /dev/null
+++ b/megatron/core/datasets/multimodal_dataset.py
@@ -0,0 +1,62 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Callable, Dict
+
+import torch
+
+from megatron.core.datasets.gpt_dataset import GPTDatasetConfig, MockGPTDataset
+
+
+@dataclass
+class MultimodalDatasetConfig(GPTDatasetConfig):
+ """Configuration object for Megatron Core Multimodal datasets.
+
+ Note: This is unused at the moment and may be missing features. Follow-up changes will use this.
+ """
+
+ image_h: int = None
+ """Image height."""
+
+ image_w: int = None
+ """Image width."""
+
+ # Function to preprocess the data sample to a format expected by a specific model. By default, do nothing.
+ preprocess_func: Callable[[Dict[str, torch.Tensor]], Dict[str, torch.Tensor]] = lambda x: x
+ """Optional function to preprocess data samples for a specific model."""
+
+ def __post_init__(self) -> None:
+ super().__post_init__()
+
+ assert self.image_h is not None
+ assert self.image_w is not None
+
+
+class MockMultimodalDataset(MockGPTDataset):
+ """Mock multimodal dataset.
+
+
+ This is unused at the moment and may be missing features. Follow-up changes will use this.
+ """
+
+ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+ """Return a sample that contains a dummy image, text sequence and the associated labels and cost and attention masks.
+
+ Args:
+ idx (int): The integer seed for mock data generation.
+
+ Returns:
+ Dict[str, torch.Tensor]: The mock data.
+ """
+ # Get a text sample.
+ sample = super().__getitem__(idx)
+
+ # Add mock input image.
+ sample["image"] = torch.zeros(
+ (3, self.config.image_h, self.config.image_w), dtype=torch.float32
+ )
+
+ # Run optional data preprocessing.
+ preprocess_func = self.config.preprocess_func
+
+ return preprocess_func(sample)
diff --git a/megatron/core/datasets/readme.md b/megatron/core/datasets/readme.md
new file mode 100644
index 0000000000..12ade943b5
--- /dev/null
+++ b/megatron/core/datasets/readme.md
@@ -0,0 +1,193 @@
+# Data Pipeline
+
+## Data pre-processing
+
+Data preprocessing is built around the following classes:
+
+1. `IndexedDatasetBuilder`
+2. `IndexedDataset`
+
+At the moment, an end-to-end data preprocessing implementation is left to the user. See the class docstring(s) for more details.
+
+#### IndexedDatasetBuilder
+
+The `IndexedDatasetBuilder` is capable of building and merging `IndexedDataset` instances.
+
+#### IndexedDataset
+
+The `IndexedDataset` class is the lowest-level data interface in Megatron Core. Internally, an `IndexedDataset` instance references two binaries: the data file (`.bin`) contains document/sequence data and the index file (`.idx`) contains document/sequence metadata.
+
+The index file stores dataset-level metadata first:
+- The index header, for backward compatibility
+- The index version, for backward compatibility
+- A numeric code corresponding to the data type used to write data to the data file
+- The number of sequences in the dataset
+- The number of documents in the dataset
+
+The index file stores document-level and sequence-level metadata second:
+- In order, the number of elements per sequence
+- In order, the byte offset (pointer) per sequence
+- In order, the consecutive sequence index range `[...)` per document
+- In order, the mode per sequence (in the multimodal case)
+
+## Data loading: construction
+
+Building the data loaders is a distributed-aware process built around the following classes:
+
+1. `BlendedMegatronDatasetConfig`
+2. `BlendedMegatronDatasetBuilder`
+3. `IndexedDataset`
+3. `MegatronDataset`
+4. `BlendedDataset`
+
+See the class docstrings for more details.
+
+#### BlendedMegatronDatasetConfig (extendable)
+
+The `BlendedMegatronDatasetConfig` class parameterizes the `BlendedMegatronDatasetBuilder` and in turn the `MegatronDataset` and `BlendedDataset`.
+
+Different training/inference regimes will require different extensions e.g. the `GPTDatasetConfig`
+
+#### BlendedMegatronDatasetBuilder
+
+The `BlendedMegatronDatasetBuilder` class builds the highest-level data interfaces in Megatron Core.
+
+**NB:** All ranks should attempt to build the dataset via the `BlendedMegatronDatasetBuilder` or the program will hang. Which ranks follow through on their attempts can be controlled via the `BlendedMegatronDatasetConfig`.
+
+#### IndexedDataset
+
+The `IndexedDataset` class is the lowest-level data interface in Megatron Core.
+
+The `IndexedDataset` should already exist on disk before attempting to build any of the high-level data interfaces.
+
+
+#### MegatronDataset (extendable)
+
+The `MegatronDataset` abstract class is a high-level data interface in Megatron Core. It is an abstraction built upon the `IndexedDataset`.
+
+Different training/inference regimes will require different extensions e.g. the `GPTDataset`
+
+#### BlendedDataset
+
+The `BlendedDataset` class is a high-level data interface in Megatron Core. It is an abstraction built upon the `MegatronDataset`.
+
+The `BlendedDataset` is only necessary when a blend multiple data distributions, i.e. multiple `MegatronDataset` instances, should contribute to a certain dataset split. The blend can be controlled via the `BlendedMegatronDatasetConfig`.
+
+## Data loading: implementation
+
+### GPTDataset
+
+The `GPTDataset` is parameterized by the following variables: the underlying `IndexedDataset` instance `indexed_dataset`, the split indices `indexed_indices` (the congituous subset of document or sequence indices used for training, validation, and testing), the number of samples `N`, the sequence length `S`, and the random seed `R`.
+
+The `GPTDataset` creates three index mappings to facilitate lookup: (1) the document index, (2) the sample index, and (3) the shuffle index.
+
+1. The document index _Do_idx_ is a 1-D array mapping from _i_ to document index of length `E * |indexed_indices|` where `E` corresponds to the minimum number of epochs such that `E * |indexed_indices| >= N`. The document index is shuffled according to `R`.
+
+ ```
+ Given:
+
+ N = 15
+ indexed_indices = [5, 6, 7, 8, 9]
+ E = 3
+
+ Then, for example:
+
+ Do_idx = [8, 8, 9, 6, 7, 5, 8, 5, 6, 6, 5, 9, 7, 7, 9]
+ ```
+
+2. The sample index _Sa_idx_ is a 2-D array mapping from _j_ to pairs of (_i_, _Do_idx_[ _i_ ] offset) of shape `[N + 1, 2]`. The rows _j_ and _j_ + 1 serve as the left and right bounds for the _j_-th sample.
+
+ ```
+ Given:
+
+ S = 1024
+
+ Then, for example:
+
+ Sa_idx[0] = (0, 0)
+ Sa_idx[1] = (0, 1024) => Do_idx[0] has length greater than S
+ Sa_idx[2] = (1, 512) => Do_idx[0] has length 1536
+ Sa_idx[3] = (2, 0) => Do_idx[1] has length 1536
+ Sa_idx[4] = (5, 300) => Do_idx[2:5] are shorter documents relative to Do_idx[0:2]
+ Sa_idx[5] = (6, 24) => Do_idx[5] has length 1300
+ ```
+
+3. The shuffle index _Sh_idx_ is a 1-D array mapping from _k_ to _j_ of length `N`. The shuffle index is shuffled according to `R`.
+
+ ```
+ Given
+
+ N = 10
+
+ Then, for example:
+
+ Sh_idx = [4, 0, 2, 6, 1, 9, 5, 8, 7, 3]
+ ```
+
+To query the `GPTDataset` for the _k_-th sample we do the following
+
+- Use the shuffle index to get the index _j_ into the sample index.
+
+ ```
+ j = Sh_idx[k]
+ ```
+- Use the sample index to get the left and right sample-bounding indices into the document index and the starting token offset for each document.
+
+ ```
+ i, offset = Sa_idx[j]
+ i_next, offset_next = Sa_idx[j + 1]
+ ```
+- Use the document index to retrieve `S` tokens from consecutive (in the document index) documents.
+
+ ```
+ sample = []
+ sample += indexed_dataset[Do_idx[i]][offset:]
+ if i != i_next:
+ sample += indexed_dataset[Do_idx[i + 1:i_next]]
+ sample += indexed_dataset[Do_idx[i_next]][:offset_next]
+ ```
+
+To save time during initialization, each index is built/cached sequentially on one process rank and subsequently loaded in parallel on other process ranks. The cached indices are unique to a hash generated in the `MegatronDataset.__init__` function.
+
+### BlendedDataset
+
+The `BlendedDataset` is parameterized by the following variables: the underlying `MegatronDataset` instances `D`, the weights `W` (one per dataset), and the size `S`. The `BlendedDataset` will draw samples from contributing datasets in proportion to the weights until achieving a composite dataset of the desired size. During each sampling step, we draw a single sample from the dataset which has the greatest sampling error.
+
+The `BlendedDataset` creates two "blending" indices to facilitate lookup: (1) the dataset index and (2) the dataset sample index.
+
+1. The dataset index _Da_idx_ is a 1-D array mapping from _i_ to dataset index of length `S`.
+
+ ```
+ Given
+
+ D = [d0, d1, d2]
+ W = [1/2, 1/4, 1/4]
+ S = 4
+
+ Then, for example:
+
+ Da_idx = [0, 1, 2, 0]
+
+ ```
+
+2. The dataset sample index _Sa_idx_ is a 1-D mapping from _i_ to the sample index for dataset _Da_idx[i]_ of length `S`.
+
+ ```
+ Given
+
+ Da_idx = [0, 1, 2, 0]
+
+ Then, for example:
+
+ Sa_idx = [0, 0, 0, 1]
+ ```
+
+To query the `BlendedDataset` for the _k_-th sample we do the following
+
+- Use the dataset index to retrieve the corresponding dataset from `D` and the dataset sample index to retrieve the corresponding sample from that dataset.
+
+ ```
+ sample = D[Da_idx[k]][Sa_idx[k]]
+ ```
+
+To save time during initialization, each index is built/cached sequentially on one process rank and subsequently loaded in parallel on other process ranks. The cached indices are unique to a hash generated in the `BlendedDataset.__init__` function.
diff --git a/megatron/core/datasets/retro/__init__.py b/megatron/core/datasets/retro/__init__.py
new file mode 100644
index 0000000000..7ce970c6e9
--- /dev/null
+++ b/megatron/core/datasets/retro/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from .config import RetroGPTChunkDatasets
+from .query.multi_split_gpt_dataset import MultiSplitGPTDataset, MultiSplitGPTDatasetConfig
+from .query.retro_dataset import get_retro_datasets
diff --git a/megatron/core/datasets/retro/config/__init__.py b/megatron/core/datasets/retro/config/__init__.py
new file mode 100644
index 0000000000..3635bedb3f
--- /dev/null
+++ b/megatron/core/datasets/retro/config/__init__.py
@@ -0,0 +1,16 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""
+Exports:
+
+ - Embedder: Base class for all Bert embedders.
+ - RetroBertEmbedders: Container class for in-memory and on-disk embedders.
+ - RetroPreprocessingConfig: Configuration class for all of Retro preprocessing.
+ - RetroGPTChunkDatasets: Container class for train, valid, and test datasets.
+ - RetroTokenizers: Container class for GPT and Bert tokenizers.
+"""
+
+from .bert_embedders import Embedder, RetroBertEmbedders
+from .config import RetroPreprocessingConfig
+from .gpt_chunk_datasets import RetroGPTChunkDatasets
+from .tokenizers import RetroTokenizers
diff --git a/megatron/core/datasets/retro/config/bert_embedders.py b/megatron/core/datasets/retro/config/bert_embedders.py
new file mode 100644
index 0000000000..8f3fe85c4a
--- /dev/null
+++ b/megatron/core/datasets/retro/config/bert_embedders.py
@@ -0,0 +1,48 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Container dataclass for holding both in-memory and on-disk Bert embedders."""
+
+import abc
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+
+
+class Embedder(abc.ABC):
+ """Base class for all Bert embedders.
+
+ All embedders should be able to embed either an entire text dataset (to a 2D
+ numpy array), or a single text string (to a 1D numpy array).
+ """
+
+ @abc.abstractmethod
+ def embed_text_dataset(self, text_dataset: torch.utils.data.Dataset) -> np.ndarray:
+ """Embed a text dataset.
+
+ Args:
+ text_dataset (torch.utils.data.Dataset): Text dataset to embed. Each sample of the text dataset should output a dict with a key 'text' and a string value.
+
+ Returns:
+ A 2D ndarray with shape (len(text_dataset), dimension(embedder)).
+ """
+
+ @abc.abstractmethod
+ def embed_text(self, text: str) -> np.ndarray:
+ """Embed a simple string of text.
+
+ Args:
+ text (str): A single text sample.
+
+ Returns:
+ A 1D ndarray with shape (dimensions(embedder),).
+ """
+
+
+@dataclass
+class RetroBertEmbedders:
+ """Container dataclass for in-memory and on-disk Bert embedders."""
+
+ disk: Embedder
+ mem: Embedder
diff --git a/megatron/core/datasets/retro/config/config.py b/megatron/core/datasets/retro/config/config.py
new file mode 100644
index 0000000000..ac9ca84124
--- /dev/null
+++ b/megatron/core/datasets/retro/config/config.py
@@ -0,0 +1,135 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Retro preprocessing config."""
+
+from dataclasses import dataclass
+
+from megatron.core.transformer import TransformerConfig
+
+from .bert_embedders import RetroBertEmbedders
+from .gpt_chunk_datasets import RetroGPTChunkDatasets
+from .tokenizers import RetroTokenizers
+
+
+@dataclass
+class RetroPreprocessingConfig(TransformerConfig):
+ """Configuration object for Retro preprocessing.
+
+ *Note* : Arguments prefixed with '--retro-gpt-*' or '--retro-bert-*' are
+ included and named as such to more easily handle managing both models
+ running at the same time. Megatron is not optimized to run two models at
+ once, so this naming convention makes it clearer.
+
+ Args:
+
+ retro_project_dir (str): Retro project directory, which contains the preprocessed data for for pretraining. This directory is built during preprocessing (see tools/retro/README.md), and contains subdirectories for the chunk database and pretraining neighbors.
+ retro_tasks (str): Comma-separated list of tasks to run. Run entire preprocesing pipeline by using '--retro-tasks build'. Alternatively, run individual stages with tasks (in this order) 'db-build', 'index-build', or 'query-pretraining-neighbors'. For example, '--retro-tasks db-build,index-build,query-pretraining-neighbors' is equivalent to '--retro-tasks build'; or the argument can contain a subset of these tasks. Stages must always be run in the correct order (listed above).
+ retro_task_validate (float): If defined, validate a randomly sampled subset of the existing results of the given task. Each task implements a 'validate' method that is responsible for sampling a `retro_task_validate` fraction of the existing results, and then checking for bitwise equality with the current code base. (E.g., `--retro-task-validate 0.01`.)
+ retro_block_size (int): Number of chunks to process at a time when generating Bert embeddings and querying the search index. Partial results for each block are generally saved to disk in separate files.
+ retro_doc_block_size (int): Number of documents to processe at time when processing token datasets into chunk databases. The partial chunk database for each block is saved into a separate file.
+ retro_gpt_seed (int): Random seed used for python, numpy, pytorch, and cuda.
+ retro_gpt_data_path (str): Path to the training dataset. Accepted format: 1) a single data path, 2) multiple datasets in the form: dataset1-weight dataset1-path dataset2-weight dataset2-path ... It is used with --split when a single dataset used for all three: train, valid and test. It is exclusive to the other --*-data-path args.
+ retro_gpt_data_cache_path (str): Path to a directory to hold cached index files.
+ retro_gpt_split (str): Comma-separated list of proportions for training, validation, and test split. For example the split `90,5,5` will use 90%% of data for training, 5%% for validation and 5%% for test.
+ retro_gpt_train_samples (int): Total number of samples to train over all training runs.
+ retro_gpt_eval_interval (int): GPT evaluation interval.
+ retro_gpt_eval_iters (int): GPT evaluation iterations.
+ retro_gpt_tokenizer_type (str): GPT tokenizer type.
+ retro_gpt_tokenizer_model (str): GPT tokenizer model file.
+ retro_gpt_vocab_file (str): GPT vocab file.
+ retro_gpt_merge_file (str): GPT merge file.
+ retro_gpt_seq_length (int): GPT sequence length.
+ retro_gpt_global_batch_size (int): GPT global batch size.
+ retro_gpt_chunk_length (int): GPT chunk length.
+ retro_bert_tokenizer_type (str): Bert tokenizer type (for when using '--bert-embedder-type megatron').
+ retro_bert_vocab_file (str): Bert vocab file.
+ retro_bert_batch_size (int): Micro-batch size for processing Bert embeddings.
+ retro_bert_max_chunk_length (int): Maximum sequence length for Bert embeddings. (Named 'chunk' here in reference to these Bert sequences being converted from GPT chunks.)
+ retro_index_type (str): A 'faiss-base' index is a simple, un-optimized wrapper around a Faiss index. A 'faiss-par-add' index optimizes the 'add()' method by making it multi-node and multi-process, but with bit-wise equivalent results.
+ retro_index_str (str): Index string used for calling faiss.index_factory(). For example, 'IVF262144_HNSW32,Flat' or 'OPQ32_256,IVF4194304_HNSW32,PQ32'.
+ retro_index_ntrain (int): Number of database chunks to use for training the index. This value must be less or equal to the total number of chunks in the database.
+ retro_index_train_load_fraction (float): Fraction of sampled chunks to use for training the index. Useful when our total sampled embeddings use too much memory; lowering the load fraction is less costly than re-embedding a new sampled dataset from scratch.
+ retro_index_add_load_fraction (float): Fraction of database chunks to use for adding to the index. Useful when our total index size would use too much memory; lowering the load fraction is less costly than re-designing our token datasets.
+ retro_index_delete_training_embeddings (bool): Delete training embeddings for the search index. Useful for debugging.
+ retro_index_delete_added_codes (bool): Delete added codes for the search index. Useful for debugging.
+ retro_query_ef_search (int): Index ef-search parameter for Hierarchical Navigable Small Worlds (HNSW) during querying.
+ retro_query_nprobe (int): Index nprobe parameter for Inverted File (IVF) during querying.
+ retro_query_num_neighbors_query (int): Number of neighbors to retrieve when calling index.search().
+ retro_query_num_neighbors_save (int): Number of neighbors to save to disk after the index's returned neighbors. If longer than target value, neighbors truncated; and if shorter than target value, neighbors are padded with -1's.
+ retro_bert_embedders (RetroBertEmbedders): Set of Bert embedders used for embedding chunks. Contains entries: 1) 'mem' for an in-memory embedder, and 2) 'disk' for an embedder that saves results in blocks to disk.
+ retro_gpt_chunk_datasets (RetroGPTChunkDatasets): GPT datasets for 'train', 'valid', and 'test'.
+ retro_tokenizers (RetroTokenizers): GPT ('gpt') and Bert ('bert') tokenizers.
+ """
+
+ # Basic.
+ retro_project_dir: str = None
+ retro_tasks: str = 'build'
+ retro_task_validate: float = None
+ retro_block_size: int = 100000
+ retro_doc_block_size: int = 100000
+
+ # GPT.
+ retro_gpt_seed: int = 1234
+ retro_gpt_data_path: list = None # basic list here, for parsing purposes
+ retro_gpt_data_cache_path: str = None
+ retro_gpt_split: str = '969,30,1'
+ retro_gpt_train_samples: int = None
+ retro_gpt_eval_interval: int = None
+ retro_gpt_eval_iters: int = None
+ retro_gpt_tokenizer_type: str = None
+ retro_gpt_tokenizer_model: str = None
+ retro_gpt_vocab_file: str = None
+ retro_gpt_merge_file: str = None
+ retro_gpt_seq_length: int = None
+ retro_gpt_global_batch_size: int = None
+ retro_gpt_chunk_length: int = 64
+
+ # Bert.
+ retro_bert_tokenizer_type: str = None
+ retro_bert_vocab_file: str = None
+ retro_bert_batch_size: int = 128
+ retro_bert_max_chunk_length: int = 256
+
+ # Index.
+ retro_index_type: str = 'faiss-par-add'
+ retro_index_str: str = None
+ retro_index_ntrain: int = None
+ retro_index_train_load_fraction: float = 1.0
+ retro_index_add_load_fraction: float = 1.0
+ retro_index_delete_training_embeddings: bool = True
+ retro_index_delete_added_codes: bool = True
+
+ # Query.
+ retro_query_ef_search: int = 256
+ retro_query_nprobe: int = 65536
+ retro_query_num_neighbors_query: int = 200
+ retro_query_num_neighbors_save: int = 20
+
+ # Tools.
+ retro_bert_embedders: RetroBertEmbedders = None
+ retro_gpt_chunk_datasets: RetroGPTChunkDatasets = None
+ retro_tokenizers: RetroTokenizers = None
+
+ def __post_init__(self) -> None:
+ """Validate Retro config."""
+
+ # Validate required attributes.
+ assert self.retro_project_dir is not None
+ assert self.retro_tasks is not None
+ assert self.retro_gpt_data_path is not None or self.retro_gpt_data_cache_path is not None
+ assert self.retro_gpt_train_samples is not None
+ assert self.retro_gpt_eval_interval is not None
+ assert self.retro_gpt_eval_iters is not None
+ assert self.retro_gpt_tokenizer_type is not None
+ assert self.retro_gpt_tokenizer_model is not None or (
+ self.retro_gpt_vocab_file is not None and self.retro_gpt_merge_file is not None
+ )
+ assert self.retro_gpt_seq_length is not None
+ assert self.retro_gpt_global_batch_size is not None
+ assert self.retro_bert_tokenizer_type is not None
+ assert self.retro_bert_vocab_file is not None
+ assert self.retro_index_str is not None
+ assert self.retro_index_ntrain is not None
+
+ # Split retro tasks.
+ self.retro_tasks = self.retro_tasks.split(",")
diff --git a/megatron/core/datasets/retro/config/gpt_chunk_datasets.py b/megatron/core/datasets/retro/config/gpt_chunk_datasets.py
new file mode 100644
index 0000000000..831b1d812b
--- /dev/null
+++ b/megatron/core/datasets/retro/config/gpt_chunk_datasets.py
@@ -0,0 +1,15 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Container dataclass for GPT chunk datasets (train, valid, and test)."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class RetroGPTChunkDatasets:
+ """Container dataclass for GPT chunk datasets."""
+
+ # Each dict contains 'dataset', 'neighbor_dir', and 'num_active_chunks'.
+ train: dict = None
+ valid: dict = None
+ test: dict = None
diff --git a/megatron/core/datasets/retro/config/tokenizers.py b/megatron/core/datasets/retro/config/tokenizers.py
new file mode 100644
index 0000000000..2e731c83b9
--- /dev/null
+++ b/megatron/core/datasets/retro/config/tokenizers.py
@@ -0,0 +1,15 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Container class for GPT and Bert tokenizers."""
+
+from dataclasses import dataclass
+
+from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer
+
+
+@dataclass
+class RetroTokenizers:
+ """Container class for GPT and Bert tokenizers."""
+
+ gpt: MegatronTokenizer = None
+ bert: MegatronTokenizer = None
diff --git a/megatron/core/datasets/retro/db/__init__.py b/megatron/core/datasets/retro/db/__init__.py
new file mode 100644
index 0000000000..f1f460b3b0
--- /dev/null
+++ b/megatron/core/datasets/retro/db/__init__.py
@@ -0,0 +1,9 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""
+Exports:
+
+ - build_db: Build a chunk database from a list of indexed datasets.
+"""
+
+from .build import build_db
diff --git a/megatron/core/datasets/retro/db/build.py b/megatron/core/datasets/retro/db/build.py
new file mode 100644
index 0000000000..44b9038230
--- /dev/null
+++ b/megatron/core/datasets/retro/db/build.py
@@ -0,0 +1,633 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Build a chunk database from a list of indexed datasets.
+
+Building a chunk database consists of.
+
+ - Breaking each document of each indexed dataset into consecutive
+ retro_gpt_chunk_length chunks.
+ - Re-tokenize each chunk into Bert, and discard any chunks with empty Bert
+ tokens.
+ - Save chunk offsets to disk for each indexed dataset.
+"""
+
+import glob
+import os
+import types
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from megatron.core.datasets.indexed_dataset import IndexedDataset
+from megatron.core.datasets.retro.config import RetroPreprocessingConfig
+from megatron.core.datasets.retro.external_libs import h5py
+from megatron.core.datasets.retro.utils import (
+ extract_data_config,
+ get_blocks_by_rank,
+ log_retro_rank_0,
+ retro_makedir,
+)
+
+from .utils import (
+ get_indexed_dataset_infos,
+ get_indexed_dataset_infos_path,
+ get_individual_chunk_db,
+ get_individual_db_dir,
+ get_individual_db_paths,
+ get_individual_doc_offsets,
+ get_merged_db_path_map,
+ init_indexed_dataset_infos,
+ load_indexed_datasets,
+ save_indexed_dataset_infos,
+)
+
+
+def build_partial_db(
+ config: types.SimpleNamespace,
+ dataset_idx: int,
+ n_datasets: int,
+ indexed_dataset: IndexedDataset,
+ block_id: int,
+ n_blocks: int,
+ block: dict,
+ proc_id: int,
+ n_procs: int,
+) -> Tuple[int, list, list, dict]:
+ """Process a document index range of the indexed dataset.
+
+ The chunk database is built in parallel blocks, since de-tokenizing &
+ re-tokenizing for Bert-length computation is expensive. This method
+ iterates each document and extracts sequential 'chunk-length' sequences
+ from each document.
+
+ Args:
+ config (types.SimpleNamespace): Subset of Retro config, containing 'chunk_length', 'gpt_eod', 'gpt_detokenize', 'bert_tokenize', and 'task_validate'.
+ dataset_idx (int): Index of this dataset out of all blended datasets.
+ n_datasets (int): Total number of blended datasets.
+ indexed_dataset (IndexedDataset): Indexed dataset to be chunked.
+ block_id (int): Block index out of all blocks to be processed.
+ n_blocks (int): Total number of blocks to be processed.
+ block (dict): Range information such as start/end points for chunking idnexed dataset.
+ proc_id (int): Process ID for tracking parallel process order.
+ n_procs (int): Total number of parallel processes.
+
+ Returns:
+ A tuple containing:
+
+ - Process ID.
+ - List of valid chunks.
+ - List of invalid chunks (i.e., chunks that converted to empty Bert embeddings.).
+ - Dict mapping document ID to number of valid chunks.
+ """
+
+ # Document start/end indexes.
+ doc_range = block["range"]
+ n_docs = doc_range[1] - doc_range[0]
+ n_docs_per_proc = int(np.ceil(n_docs / n_procs))
+ doc_start_id = doc_range[0] + proc_id * n_docs_per_proc
+ doc_end_id = min(doc_range[1], doc_start_id + n_docs_per_proc)
+
+ # Print progress.
+ progress_proc_ids = set(range(n_procs)) if torch.distributed.get_rank() == 0 else set()
+ if proc_id in progress_proc_ids:
+ log_retro_rank_0(
+ " > building partial chunk db, proc %d / %d, docs %d:%d / %d."
+ % (proc_id, n_procs, doc_start_id, doc_end_id, n_docs)
+ )
+
+ # Progress bars (snapshot of overall progress).
+ doc_id_iter = range(doc_start_id, doc_end_id)
+ pbar = (
+ tqdm(doc_id_iter, "parse doc chunks", miniters=len(doc_id_iter) // 20)
+ if proc_id in progress_proc_ids
+ else doc_id_iter
+ )
+
+ # Iterate documents & parse chunks.
+ chunk_db_valid: List[Tuple] = []
+ chunk_db_invalid: List[Tuple] = []
+ doc_size_map = {}
+ for doc_id in pbar:
+
+ # Progress description.
+ try:
+ pbar.set_description(
+ "%sds %d / %d, block %d / %d, proc %d / %d."
+ % (
+ "" if config.task_validate is None else "[validate] ",
+ dataset_idx,
+ n_datasets,
+ block_id,
+ n_blocks,
+ proc_id,
+ n_procs,
+ )
+ )
+ except Exception:
+ pass
+
+ # Remove EOD token.
+ doc = indexed_dataset.get(doc_id)
+ if doc[-1].item() == config.gpt_eod:
+ doc = doc[:-1]
+ doc_len = len(doc)
+
+ # Chunk start/end indexes.
+ chunk_start_idxs = list(range(0, doc_len, config.chunk_length))
+ chunk_end_idxs = [min(doc_len, s + config.chunk_length) for s in chunk_start_idxs]
+
+ # Re-tokenize each chunk to Bert/Wordpiece (empty bert -> 'invalid').
+ doc_size_map[doc_id] = 0
+ for i, chunk_start_idx in enumerate(chunk_start_idxs):
+
+ # Re-tokenize.
+ chunk_end_idx = chunk_end_idxs[i]
+ gpt_token_ids = indexed_dataset.get(
+ idx=doc_id, offset=chunk_start_idx, length=chunk_end_idx - chunk_start_idx
+ )
+ text = config.gpt_detokenize(gpt_token_ids.tolist())
+ bert_token_ids = config.bert_tokenize(text)
+
+ # 'Valid' for non-empty Bert chunks; 'invalid' otherwise.
+ if len(bert_token_ids) == 0:
+ _chunk_db = chunk_db_invalid
+ else:
+ _chunk_db = chunk_db_valid
+ doc_size_map[doc_id] += 1
+ _chunk_db.append((doc_id, chunk_start_idx, chunk_end_idx, len(bert_token_ids)))
+
+ return proc_id, chunk_db_valid, chunk_db_invalid, doc_size_map
+
+
+def build_block_db(
+ config: RetroPreprocessingConfig,
+ dataset_idx: int,
+ n_datasets: int,
+ indexed_dataset: IndexedDataset,
+ n_procs: int,
+ executor: ProcessPoolExecutor,
+ n_missing_blocks: int,
+ block_idx: int,
+ block: dict,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+ """Split each document within block into consecutive retro_gpt_chunk_length size chunks.
+
+ Args:
+ config (RetroPreprocessingConfig): For DB building, we make use of attributes 'chunk_length', 'gpt_eod', 'gpt_detokenize', 'bert_tokenize', and 'task_validate'.
+ dataset_idx (int): Index of this dataset out of all blended datasets.
+ n_datasets (int): Total number of blended datasets.
+ indexed_dataset (IndexedDataset): Indexed dataset to be chunked.
+ n_procs (int): Total number of parallel processes.
+ executor (ProcessPoolExecutor): Executor for launching parallel processes.
+ n_missing_blocks (int): Total number of blocks to be processed.
+ block_idx (int): Block index out of all blocks to be processed.
+ block (dict): Range information such as start/end points for chunking idnexed dataset.
+
+ Returns:
+ A tuple containing:
+
+ - List of valid chunks.
+ - List of invalid chunks (i.e., chunks that converted to empty Bert embeddings.).
+ - Dict mapping document ID to number of valid chunks.
+ """
+
+ # Build partial dbs.
+ log_retro_rank_0(' > build partial dbs.')
+ futures = []
+ for proc_id in range(n_procs): # not true process id
+ futures.append(
+ executor.submit(
+ build_partial_db,
+ types.SimpleNamespace(
+ chunk_length=config.retro_gpt_chunk_length,
+ gpt_eod=config.retro_tokenizers.gpt.eod,
+ gpt_detokenize=config.retro_tokenizers.gpt.detokenize,
+ bert_tokenize=config.retro_tokenizers.bert.tokenize,
+ task_validate=config.retro_task_validate,
+ ),
+ dataset_idx,
+ n_datasets,
+ indexed_dataset,
+ block_idx,
+ n_missing_blocks,
+ block,
+ proc_id,
+ n_procs,
+ )
+ )
+ partial_chunk_dbs = []
+ for future in as_completed(futures):
+ partial_chunk_dbs.append(future.result())
+
+ # Concatenate chunks.
+ partial_chunk_dbs.sort(key=lambda item: item[0]) # sort by proc_id
+ chunk_db_valid = [
+ item for partial_chunk_db in partial_chunk_dbs for item in partial_chunk_db[1]
+ ]
+ chunk_db_invalid = [
+ item for partial_chunk_db in partial_chunk_dbs for item in partial_chunk_db[2]
+ ]
+
+ # Convert to numpy.
+ log_retro_rank_0(' > converting chunk db to numpy.')
+ chunk_db_valid = np.array(chunk_db_valid, dtype="uint32")
+ chunk_db_invalid = np.array(chunk_db_invalid, dtype="uint32")
+
+ # Document offsets.
+ doc_sizes = [
+ (d, s) for partial_chunk_db in partial_chunk_dbs for d, s in partial_chunk_db[3].items()
+ ]
+ doc_sizes.sort(key=lambda item: item[0])
+ doc_offsets = np.cumsum([item[1] for item in doc_sizes]).astype("uint64")
+ doc_offsets = np.stack(
+ (np.array([item[0] for item in doc_sizes], dtype="uint64"), doc_offsets), axis=1
+ )
+
+ return chunk_db_valid, chunk_db_invalid, doc_offsets
+
+
+def save_block_db(
+ block: dict, chunk_db_valid: np.ndarray, chunk_db_invalid: np.ndarray, doc_offsets: np.ndarray
+) -> None:
+ """Save block of chunked tokens to disk. These blocks are later used for
+ training and adding to the vector index.
+
+ Args:
+ block (dict): Range information such as start/end points for chunking idnexed dataset.
+ chunk_db_valid (np.ndarray): Array of valid chunk indexes.
+ chunk_db_invalid (np.ndarray): Array of invalid chunk indexes.
+ doc_offsets (np.ndarray): Array of document offsets by chunks.
+ """
+ log_retro_rank_0(" > saving individual db.")
+ with h5py.File(block["path"], "w") as f:
+ dset = f.create_dataset("chunks_valid", data=chunk_db_valid)
+ dset = f.create_dataset("chunks_invalid", data=chunk_db_invalid)
+ dset = f.create_dataset("doc_offsets", data=doc_offsets)
+
+
+def build_individual_db(
+ config: RetroPreprocessingConfig, dataset_idx: int, n_datasets: int, dataset_info: dict
+) -> None:
+ """Process a single indexed dataset & extract chunks.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ dataset_idx (int): Dataset index within blended dataset.
+ n_datasets (int): Total number of datasets within blended dataset.
+ dataset_info (dict): Metadata for dataset (see `save_indexed_dataset_infos()` in `utils.py` for more detail).
+ """
+
+ # Make directory.
+ db_dir = get_individual_db_dir(config.retro_project_dir, dataset_info["prefix"])
+ retro_makedir(config, db_dir)
+
+ # Indexed dataset.
+ indexed_dataset = dataset_info["dataset"]
+
+ # Missing DB blocks (split by documents).
+ blocks = get_blocks_by_rank(
+ db_dir,
+ len(indexed_dataset),
+ config.retro_doc_block_size,
+ validate=lambda f: f["chunks_valid"].shape == (0,) or f["chunks_valid"].shape[1] == 4,
+ sample=config.retro_task_validate,
+ )
+ if config.retro_task_validate is None:
+ active_blocks = blocks.missing
+ else:
+ assert blocks.n_missing_world == 0
+ active_blocks = blocks.existing
+
+ # Prevent missing-path-write race condition.
+ torch.distributed.barrier()
+
+ # Nothing to do?
+ if config.retro_task_validate is None and not active_blocks:
+ return
+
+ # Num processes.
+ if blocks.n_missing_world == 1:
+ n_procs = 128
+ elif blocks.n_missing_world <= 2:
+ n_procs = 64
+ elif blocks.n_missing_world <= 4:
+ n_procs = 32
+ elif blocks.n_missing_world <= 8:
+ n_procs = 16
+ else:
+ n_procs = 8
+
+ # Process documents in parallel.
+ with ProcessPoolExecutor(max_workers=n_procs) as executor:
+ for block_idx, block in enumerate(active_blocks):
+
+ if block is not None:
+
+ # Build block DB.
+ chunk_db_valid, chunk_db_invalid, doc_offsets = build_block_db(
+ config=config,
+ dataset_idx=dataset_idx,
+ n_datasets=n_datasets,
+ indexed_dataset=indexed_dataset,
+ n_procs=n_procs,
+ executor=executor,
+ n_missing_blocks=len(active_blocks),
+ block_idx=block_idx,
+ block=block,
+ )
+
+ if config.retro_task_validate is None:
+ # Save block DB.
+ save_block_db(
+ block=block,
+ chunk_db_valid=chunk_db_valid,
+ chunk_db_invalid=chunk_db_invalid,
+ doc_offsets=doc_offsets,
+ )
+
+ else:
+
+ # Load existing block DB.
+ with h5py.File(block["path"]) as f:
+ existing_chunks_valid = np.copy(f["chunks_valid"])
+ existing_chunks_invalid = np.copy(f["chunks_invalid"])
+ existing_doc_offsets = np.copy(f["doc_offsets"])
+
+ # Check equality.
+ log_retro_rank_0(" > validate.")
+ assert np.array_equal(existing_chunks_valid, chunk_db_valid)
+ assert np.array_equal(existing_chunks_invalid, chunk_db_invalid)
+ assert np.array_equal(existing_doc_offsets, doc_offsets)
+
+ # Wait for all ranks to finish block.
+ log_retro_rank_0(" > waiting for all ranks to finish block.")
+ torch.distributed.barrier()
+
+ log_retro_rank_0(" > finished saving individual db.")
+
+
+def build_individual_dbs(
+ config: RetroPreprocessingConfig, indexed_dataset_infos: List[Dict]
+) -> None:
+ """Iterate each indexed dataset & process its chunks.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ indexed_dataset_infos (List[Dict]): Preprocessing metadata for each dataset.
+ """
+
+ # Build individual DBs.
+ log_retro_rank_0(" > build individual chunk dbs.")
+ for ds_idx, ds_info in enumerate(indexed_dataset_infos):
+
+ # Progress.
+ log_retro_rank_0(
+ " > building individual db, dataset %d / %d ... '%s'."
+ % (ds_idx, len(indexed_dataset_infos), ds_info["prefix"])
+ )
+
+ # Process single dataset.
+ build_individual_db(config, ds_idx, len(indexed_dataset_infos), ds_info)
+
+
+def update_chunk_counts(
+ config: RetroPreprocessingConfig, indexed_dataset_infos: List[Dict]
+) -> None:
+ """Set n_chunks_train & n_chunks sampled for each individual DB.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ indexed_dataset_infos (List[Dict]): Preprocessing metadata for each dataset (i.e., 'prefix', 'ratio', 'n_chunks', etc.).
+ """
+
+ if torch.distributed.get_rank() != 0:
+ return
+
+ # Data ratio sum (for setting index training chunks).
+ data_ratio_sum = sum([d["ratio"] for d in indexed_dataset_infos])
+
+ # Training split size (split at document level).
+ train_fraction = float(extract_data_config(config).split.split(",")[0]) / 100
+ assert train_fraction > 0 and train_fraction <= 1
+
+ # Set n_chunks (including n_chunks_sampled for unambiguity).
+ log_retro_rank_0(" > compute n_chunks.")
+ for ds_index, ds_info in enumerate(indexed_dataset_infos):
+
+ db_paths = get_individual_db_paths(config.retro_project_dir, ds_info["prefix"])
+
+ # Update counts.
+ ds_info["n_docs"] = len(ds_info["dataset"].document_indices) - 1
+ ds_info["n_docs_train"] = int(train_fraction * ds_info["n_docs"])
+ ds_info["n_chunks"] = 0 # previously, 'n_chunks_valid'
+ ds_info["n_chunks_train"] = 0
+ ds_info["n_chunks_invalid"] = 0
+ for db_path in tqdm(
+ db_paths, "%d/%d, %s" % (ds_index, len(indexed_dataset_infos), ds_info["prefix"])
+ ):
+ with h5py.File(db_path, "r") as f:
+ ds_info["n_chunks"] += len(f["chunks_valid"])
+ ds_info["n_chunks_invalid"] += len(f["chunks_invalid"])
+ ds_info["n_chunks_train"] += (
+ (np.copy(f["chunks_valid"][:, 0]) < ds_info["n_docs_train"]).sum().item()
+ )
+
+ ds_info["n_chunks_sampled"] = int(
+ config.retro_index_ntrain * ds_info["ratio"] / data_ratio_sum
+ )
+
+ # Verify counts.
+ assert ds_info["n_chunks_train"] <= ds_info["n_chunks"], "n_train (%d) > n_total (%d)." % (
+ ds_info["n_chunks_train"],
+ ds_info["n_chunks"],
+ )
+ assert (
+ ds_info["n_chunks_sampled"] <= ds_info["n_chunks_train"]
+ ), "n_sampled (%d) > n_train (%d)." % (
+ ds_info["n_chunks_sampled"],
+ ds_info["n_chunks_train"],
+ )
+
+
+def merge_dbs(project_dir: str, indexed_dataset_infos: List[Dict], db_type: str) -> None:
+ """Merge individual DBs into single DB.
+
+ Args:
+ project_dir (str): Retro project dir.
+ indexed_dataset_infos (List[Dict]): Preprocessing metadata for each dataset (i.e., 'prefix', 'ratio', 'n_chunks', etc.).
+ db_type (str): DB type (e.g., 'sampled', 'train', or 'valid').
+ """
+
+ if torch.distributed.get_rank() != 0:
+ return
+
+ log_retro_rank_0(" > build %s chunk db." % db_type)
+
+ # Count chunks.
+ if db_type == "sampled":
+ n_chunks_key = "n_chunks_sampled"
+ n_docs_key = None
+ elif db_type == "train":
+ n_chunks_key = "n_chunks_train"
+ n_docs_key = "n_docs_train"
+ elif db_type == "valid":
+ n_docs_key = None
+ else:
+ raise Exception("handle db_type '%s'." % db_type)
+
+ if db_type == "valid":
+ n_chunks = sum(m["n_chunks"] - m["n_chunks_train"] for m in indexed_dataset_infos)
+ else:
+ n_chunks = sum(m[n_chunks_key] for m in indexed_dataset_infos)
+ n_docs = None if n_docs_key is None else sum(m[n_docs_key] for m in indexed_dataset_infos)
+
+ # DB path.
+ db_path = get_merged_db_path_map(project_dir)[db_type]
+
+ # Delete existing chunk db if incorrect size.
+ if os.path.exists(db_path):
+
+ try:
+
+ f = h5py.File(db_path)
+ n_alloc = len(f["chunks"]) # total allocated
+ n_written = f["n_written"][0].item() # total written
+ f.close()
+
+ if n_chunks != n_alloc or n_chunks != n_written:
+ os.remove(db_path)
+
+ except Exception as e:
+ if isinstance(e, OSError):
+ os.remove(db_path)
+ elif isinstance(e, KeyError):
+ f.close()
+ os.remove(db_path)
+ else:
+ raise e
+
+ # Build merged chunk db.
+ if not os.path.exists(db_path):
+
+ os.makedirs(os.path.dirname(db_path), exist_ok=True)
+ f = h5py.File(db_path, "w")
+
+ # Initialize output arrays.
+ merged_chunk_db: np.ndarray = f.create_dataset("chunks", (n_chunks, 5), dtype="uint32")
+ merged_doc_offsets: np.ndarray = (
+ None
+ if n_docs_key is None
+ else f.create_dataset("doc_offsets", (n_docs, 3), dtype="uint64")
+ )
+ n_written = f.create_dataset("n_written", (1,), dtype="uint64")
+ n_written[0] = 0
+
+ # Iterate indexed datasets & collect chunks.
+ chunk_start_index = 0
+ doc_start_index = 0
+ doc_start_offset = 0
+ for ds_idx, ds_info in enumerate(indexed_dataset_infos):
+ log_retro_rank_0(
+ " > merging dbs; '%s', dataset %d / %d ... '%s'."
+ % (db_type, ds_idx, len(indexed_dataset_infos), ds_info["prefix"])
+ )
+ individual_chunk_db: np.ndarray = get_individual_chunk_db(project_dir, ds_idx, ds_info)
+ individual_doc_offsets: np.ndarray = (
+ None
+ if n_docs_key is None
+ else get_individual_doc_offsets(project_dir, ds_idx, ds_info)
+ )
+
+ if db_type == "valid":
+ individual_chunk_db = individual_chunk_db[ds_info["n_chunks_train"] :]
+ if n_docs_key is None:
+ individual_doc_offsets = None
+ else:
+ train_doc_offset = individual_doc_offsets[ds_info["n_docs_train"] - 1, 2]
+ individual_doc_offsets = np.copy(
+ individual_doc_offsets[ds_info["n_docs_train"] :]
+ )
+ individual_doc_offsets[:, 2] -= train_doc_offset
+
+ log_retro_rank_0("~~~")
+ log_retro_rank_0(individual_doc_offsets)
+ log_retro_rank_0(train_doc_offset)
+ raise Exception("test me.")
+ else:
+ individual_chunk_db = individual_chunk_db[: ds_info[n_chunks_key]]
+ individual_doc_offsets = (
+ None
+ if n_docs_key is None
+ else np.copy(individual_doc_offsets[: ds_info[n_docs_key]])
+ )
+
+ merged_chunk_db[chunk_start_index : chunk_start_index + len(individual_chunk_db)] = (
+ individual_chunk_db
+ )
+ chunk_start_index += len(individual_chunk_db)
+ n_written[0] = chunk_start_index
+ if n_docs_key is not None:
+ individual_doc_offsets[:, 2] += doc_start_offset
+ doc_end_index = doc_start_index + individual_doc_offsets.shape[0]
+ merged_doc_offsets[doc_start_index:doc_end_index] = individual_doc_offsets
+ doc_start_index = doc_end_index
+ doc_start_offset = individual_doc_offsets[-1, 2].item()
+
+ f.close()
+
+
+def build_merged_dbs(project_dir: str, indexed_dataset_infos: List[Dict]) -> None:
+ """Merge individual dataset components into single database.
+
+ This method merges databases for DB types:
+ - 'sampled': used for training the vector index.
+ - 'train': used for adding to the trained vector index.
+ - 'valid': can be used for validating/testing the vector index.
+
+ Args:
+ project_dir (str): Retro project dir.
+ indexed_dataset_infos (List[Dict]): Preprocessing metadata for each dataset (i.e., 'prefix', 'ratio', 'n_chunks', etc.).
+ """
+ merge_dbs(project_dir, indexed_dataset_infos, "sampled")
+ merge_dbs(project_dir, indexed_dataset_infos, "train")
+ merge_dbs(project_dir, indexed_dataset_infos, "valid")
+
+
+def build_db(config: RetroPreprocessingConfig) -> None:
+ """Extract token chunks from each indexed dataset.
+
+ Iterate each document of each indexed dataset, extract that document's chunks, and save to a 'DB' (hdf5 file).
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ project_dir = config.retro_project_dir
+
+ # Indexed dataset info.
+ if config.retro_task_validate is None:
+ indexed_dataset_infos = init_indexed_dataset_infos(config)
+ else:
+ indexed_dataset_infos = get_indexed_dataset_infos(config.retro_project_dir)
+ # Build individual dbs.
+ build_individual_dbs(config, indexed_dataset_infos)
+
+ # If validating, return here.
+ if config.retro_task_validate is not None:
+ return
+
+ # Single-process going forward.
+ if torch.distributed.get_rank() != 0:
+ return
+
+ # Update n_chunks & save indexed dataset infos.
+ if not os.path.exists(get_indexed_dataset_infos_path(project_dir)):
+ update_chunk_counts(config, indexed_dataset_infos)
+ save_indexed_dataset_infos(project_dir, indexed_dataset_infos)
+ indexed_dataset_infos = get_indexed_dataset_infos(project_dir)
+
+ # Builded merged dbs.
+ build_merged_dbs(project_dir, indexed_dataset_infos)
diff --git a/megatron/core/datasets/retro/db/dataset.py b/megatron/core/datasets/retro/db/dataset.py
new file mode 100644
index 0000000000..f9053622ab
--- /dev/null
+++ b/megatron/core/datasets/retro/db/dataset.py
@@ -0,0 +1,105 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""A DBDataset is for iterating the chunks of the chunk database.
+
+This dataset is used for both training a vector index, and adding vectors to a
+trained index.
+"""
+
+from typing import List
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from megatron.core.datasets.indexed_dataset import IndexedDataset
+
+
+class DBDataset(torch.utils.data.Dataset):
+ """Dataset for iterating chunks.
+
+ Args:
+ db_path (str): Path of HDF5-format chunk database.
+ indexed_datasets (List[IndexedDataset]): Indexed datasets used to build database.
+ chunks (np.ndarray): Array of chunk indexes, for indexing into indexed datasets. Format [dataset_idx, doc_id, start_idx, end_idx, bert_length].
+ chunk_length (int): Max GPT chunk length (e.g., 64).
+ eod_token_id (int): EOD token ID.
+ """
+
+ def __init__(
+ self,
+ db_path: str,
+ indexed_datasets: List[IndexedDataset],
+ chunks: np.ndarray,
+ chunk_length: int,
+ eod_token_id: int,
+ ):
+
+ assert chunks.shape[1] == 5, (
+ "expected 5 columns (dataset_idx, "
+ "doc_idx, token_start_idx, token_end_idx, bert_chunk_length); "
+ "found %d columns." % chunks.shape[1]
+ )
+
+ self.db_path = db_path
+ self.indexed_datasets = indexed_datasets
+ self.chunks = chunks
+ self.doc_chunk_map = None
+
+ self.max_chunk_length = chunk_length
+ self.eod_token_id = eod_token_id
+
+ def __len__(self) -> int:
+ """Length of DB dataset.
+
+ Returns:
+ Number of chunks contained in the dataset.
+ """
+ return self.chunks.shape[0]
+
+ def __getitem__(self, chunk_id: int) -> dict:
+ """DB dataset sample.
+
+ Args:
+ chunk_id (int): Index of chunk within dataset.
+
+ Returns:
+ A dict containing:
+ - 'doc_id': Document index within indexed dataset.
+ - 'text': GPT token IDs.
+ """
+
+ # Chunk start/end indexes.
+ indexed_dataset_id, doc_id, token_start_idx, token_end_idx, _ = [
+ value.item() for value in self.chunks[chunk_id]
+ ]
+ chunk_length = token_end_idx - token_start_idx
+ indexed_dataset = self.indexed_datasets[indexed_dataset_id]
+
+ # Chunk token ids.
+ token_ids = indexed_dataset.get(doc_id, offset=token_start_idx, length=chunk_length)
+
+ # Extend chunks to max_chunk_length by padding with EOD tokens.
+ if chunk_length != self.max_chunk_length:
+ assert chunk_length < self.max_chunk_length, "invalid chunk len."
+ token_ids = token_ids.tolist()
+ token_ids += [self.eod_token_id] * (self.max_chunk_length - chunk_length)
+
+ return {"doc_id": doc_id, "text": np.array(token_ids, dtype=np.int64)}
+
+ def load_doc_tuples(self) -> None:
+ """Load the dataset & document ids.
+
+ Load the dataset id & document id of each chunk in the database, to
+ be used for causality filtering during querying.
+ """
+ self.doc_tuples = np.zeros(shape=(len(self), 2), dtype="uint32")
+ block_size = int(1e6)
+ for start_idx in tqdm(
+ range(0, len(self), block_size),
+ "load doc tuples",
+ miniters=(len(self) // block_size) // 10,
+ disable=torch.distributed.get_rank() != 0,
+ ):
+ end_idx = min(len(self), start_idx + block_size)
+ self.doc_tuples[start_idx:end_idx] = self.chunks[start_idx:end_idx, :2]
diff --git a/megatron/core/datasets/retro/db/utils.py b/megatron/core/datasets/retro/db/utils.py
new file mode 100644
index 0000000000..e8578a09d5
--- /dev/null
+++ b/megatron/core/datasets/retro/db/utils.py
@@ -0,0 +1,367 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Utilities for building a chunk database."""
+
+import glob
+import json
+import os
+from typing import Dict, List, Optional
+
+import numpy as np
+
+from megatron.core.datasets.indexed_dataset import IndexedDataset
+from megatron.core.datasets.retro.config import RetroPreprocessingConfig
+from megatron.core.datasets.retro.external_libs import h5py
+from megatron.core.models.retro.utils import get_gpt_data_dir
+
+from .dataset import DBDataset
+
+
+def get_db_dir(project_dir: str) -> str:
+ """Sub-directory for DB data.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+
+ Returns:
+ Path of the DB sub-directory within the project.
+ """
+ return os.path.join(project_dir, "db")
+
+
+def init_indexed_dataset_infos(config: RetroPreprocessingConfig) -> List[Dict]:
+ """Gather meta-info about each indexed dataset.
+
+ The returned info array allows for easy access to the configuration, and
+ helps remove ambiguity.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ List of processing metadata for each dataset, including:
+ - ratio: Data split weight.
+ - prefix: Relative path to dataset under DB sub-directory.
+ """
+
+ data_dir = get_gpt_data_dir(config.retro_project_dir)
+ data_blend: List[str] = config.retro_gpt_data_path
+ assert len(data_blend) % 2 == 0, "currently, only blended dataset is supported."
+
+ # Dataset infos.
+ infos = []
+ for i in range(0, len(data_blend), 2):
+ ratio = float(data_blend[i])
+ prefix = data_blend[i + 1]
+ path = os.path.join(data_dir, prefix + ".bin")
+ assert os.path.exists(path), "couldn't find '%s'." % path
+ infos.append({"ratio": ratio, "prefix": prefix})
+
+ # Load indexed datasets.
+ load_indexed_datasets(config.retro_project_dir, infos)
+
+ return infos
+
+
+def get_indexed_dataset_infos_path(project_dir: str) -> str:
+ """Path to indexed dataset meta-infos.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+
+ Returns:
+ Path to the `indexed_dataset_infos.json` file.
+ """
+ return os.path.join(get_db_dir(project_dir), "indexed_dataset_infos.json")
+
+
+def save_indexed_dataset_infos(project_dir: str, indexed_dataset_infos: List[Dict]) -> None:
+ """Save dataset order & meta-info.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ indexed_dataset_infos (List[Dict]): List of metadata for each dataset, with each entry containing:
+
+ - ratio: Data split weight.
+ - prefix: Relative path to dataset under DB sub-directory.
+ - n_docs: Number of documents.
+ - n_docs_train: Number of documents used for pretraining.
+ - n_chunks: Number of valid chunks.
+ - n_chunks_train: Number of valid chunks used for pretraining.
+ - n_chunks_invalid: Number of invalid chunks.
+ - n_chunks_sampled: Number of valid chunks used for vector index training.
+ """
+
+ # Remove 'dataset' field.
+ clean_infos = []
+ for info in indexed_dataset_infos:
+ info = dict(info)
+ del info["dataset"]
+ clean_infos.append(info)
+
+ # Save.
+ with open(get_indexed_dataset_infos_path(project_dir), "w") as f:
+ json.dump(clean_infos, f, indent=4)
+
+
+def load_indexed_datasets(project_dir: str, indexed_dataset_infos: List[Dict]) -> None:
+ """Loaded indexed datasets into memory-mapped datasets.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ indexed_dataset_infos (List[Dict]): List of metadata for each dataset (see `save_indexed_dataset_infos()` for more details.
+ """
+ data_dir = get_gpt_data_dir(project_dir)
+ for info in indexed_dataset_infos:
+ info["dataset"] = IndexedDataset(os.path.join(data_dir, info["prefix"]), mmap=True)
+
+
+def get_indexed_dataset_infos(project_dir: str) -> List[Dict]:
+ """Load indexed dataset meta-infos.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+
+ Returns:
+ List of metadata for each dataset (see `save_indexed_dataset_infos()` for more details.
+ """
+
+ # Load json.
+ path = get_indexed_dataset_infos_path(project_dir)
+ with open(path) as f:
+ infos = json.load(f)
+
+ # Load indexed datasets.
+ load_indexed_datasets(project_dir, infos)
+
+ return infos
+
+
+def get_individual_db_dir(project_dir: str, prefix: str) -> str:
+ """Individual DB's directory.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ prefix (str): Unique relative path to dataset within project dir.
+
+ Returns:
+ Path to the given datasets's chunk database.
+ """
+ return os.path.join(get_db_dir(project_dir), "individual", prefix)
+
+
+def get_individual_db_paths(project_dir: str, prefix: str) -> List[str]:
+ """Get paths of all database blocks of an individual dataset.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ prefix (str): Unique relative path to dataset within project dir.
+
+ Returns:
+ Paths to each HDF5 chunk database files that comprises this datasets full chunk database.
+ """
+ return sorted(glob.glob(get_individual_db_dir(project_dir, prefix) + "/*hdf5"))
+
+
+def get_individual_chunk_db(project_dir: str, ds_id: int, ds_info: dict) -> np.ndarray:
+ """Load individual dataset's chunk DB.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ ds_id (int): Index of dataset within blended dataset.
+ ds_info (dict): Preprocessing metadata for dataset (see `save_indexed_dataset_infos()` for more detail).
+
+ Returns:
+ Array of chunk start/end indexes for this dataset, where the chunk indexes can be used for indexing into the corresponding indexed dataset.
+ """
+ paths = get_individual_db_paths(project_dir, ds_info["prefix"])
+ # *Note*: convert to dataset, rather than copying to memory.
+ db = np.zeros((ds_info["n_chunks"], 5), dtype="uint32")
+ db[:, 0] = ds_id
+ start_idx = 0
+ for path in paths:
+ f = h5py.File(path, "r")
+ n_chunks_current = f["chunks_valid"].shape[0]
+ db[start_idx : (start_idx + n_chunks_current), 1:] = f["chunks_valid"]
+ start_idx += n_chunks_current
+ f.close()
+
+ assert start_idx == ds_info["n_chunks"]
+
+ return db
+
+
+def get_individual_doc_offsets(project_dir: str, ds_id: int, ds_info: dict) -> np.ndarray:
+ """Load individual dataset's document offsets.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ ds_id (int): Index of dataset within blended dataset.
+ ds_info (dict): Preprocessing metadata for dataset (see `save_indexed_dataset_infos()` for more detail).
+
+ Returns:
+ Array of document offsets by chunk index for this dataset.
+ """
+ paths = get_individual_db_paths(project_dir, ds_info["prefix"])
+ # *Note*: convert to dataset, rather than copying to memory.
+ doc_offsets = np.zeros((ds_info["n_docs"], 3), dtype="uint64")
+ doc_offsets[:, 0] = ds_id
+ start_idx = 0
+ start_offset = 0
+ for path in paths:
+ with h5py.File(path) as f:
+ current_doc_offsets = np.copy(f["doc_offsets"])
+ current_doc_offsets[:, 1] += start_offset
+ current_ndocs = current_doc_offsets.shape[0]
+ doc_offsets[start_idx : (start_idx + current_ndocs), 1:] = current_doc_offsets
+ start_idx += current_ndocs
+ start_offset = current_doc_offsets[-1, 1].item()
+
+ return doc_offsets
+
+
+def get_merged_db_path_map(project_dir: str) -> dict:
+ """Paths to merged datasets.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+
+ Returns:
+ A dict of chunk databases, one for each of:
+ - sampled: Chunks used for training the vector index.
+ - train: Chunks used for pretraining 'train' dataset.
+ - valid: Chunks used for pretraining 'valid' dataset.
+ """
+ base_dir = get_db_dir(project_dir)
+ return {
+ "sampled": os.path.join(base_dir, "merged", "sampled.hdf5"),
+ "train": os.path.join(base_dir, "merged", "train.hdf5"),
+ "valid": os.path.join(base_dir, "merged", "valid.hdf5"),
+ }
+
+
+def get_merged_dataset(
+ project_dir: str,
+ chunk_length: int,
+ eod_token_id: int,
+ db_type: str,
+ indexed_dataset_infos: Optional[List[Dict]] = None,
+) -> DBDataset:
+ """Get merged dataset.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ chunk_length (int): GPT chunk length (e.g., 64).
+ eod_token_id (int): EOD token ID.
+ db_type (str): DB type (e.g., 'sampled', 'train', or 'valid').
+ indexed_dataset_infos (Optional[List[Dict]]): Optionally, pre-loaded list of dataset metadata (see `save_indexed_dataset_infos()` for more detail). If not provided, the indexed dataset infos will be loaded from disk.
+
+ Returns:
+ A DBDataset, which is a dataset that wraps the HDF5 chunk index array.
+ """
+
+ if not indexed_dataset_infos:
+ indexed_dataset_infos = get_indexed_dataset_infos(project_dir)
+
+ # Load chunks.
+ db_path = get_merged_db_path_map(project_dir)[db_type]
+ f = h5py.File(db_path, "r")
+ chunks = f["chunks"]
+
+ # DB dataset.
+ indexed_datasets = [info["dataset"] for info in indexed_dataset_infos]
+ dataset = DBDataset(
+ db_path=db_path,
+ indexed_datasets=indexed_datasets,
+ chunks=chunks,
+ chunk_length=chunk_length,
+ eod_token_id=eod_token_id,
+ )
+
+ return dataset
+
+
+def get_merged_sampled_dataset(
+ project_dir: str,
+ chunk_length: int,
+ eod_token_id: int,
+ indexed_dataset_infos: Optional[List[Dict]] = None,
+) -> DBDataset:
+ """Get sampled dataset (for training the vector index).
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ chunk_length (int): GPT chunk length (e.g., 64).
+ eod_token_id (int): EOD token ID.
+ indexed_dataset_infos (Optional[List[Dict]]): Optionally, pre-loaded list of dataset metadata (see `save_indexed_dataset_infos()` for more detail). If not provided, the indexed dataset infos will be loaded from disk.
+
+ Returns:
+ A DBDataset, which is a dataset that wraps the HDF5 chunk index array.
+ """
+ return get_merged_dataset(
+ project_dir, chunk_length, eod_token_id, "sampled", indexed_dataset_infos
+ )
+
+
+def get_merged_train_dataset(
+ project_dir: str,
+ chunk_length: int,
+ eod_token_id: int,
+ indexed_dataset_infos: Optional[List[Dict]] = None,
+) -> DBDataset:
+ """Get training dataset (for adding to the vector index).
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ chunk_length (int): GPT chunk length (e.g., 64).
+ eod_token_id (int): EOD token ID.
+ indexed_dataset_infos (Optional[List[Dict]]): Optionally, pre-loaded list of dataset metadata (see `save_indexed_dataset_infos()` for more detail). If not provided, the indexed dataset infos will be loaded from disk.
+
+ Returns:
+ A DBDataset, which is a dataset that wraps the HDF5 chunk index array.
+ """
+ return get_merged_dataset(
+ project_dir, chunk_length, eod_token_id, "train", indexed_dataset_infos
+ )
+
+
+def get_merged_valid_dataset(
+ project_dir: str,
+ chunk_length: int,
+ eod_token_id: int,
+ indexed_dataset_infos: Optional[List[Dict]] = None,
+) -> DBDataset:
+ """Get validation dataset (for testing the vector index).
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ chunk_length (int): GPT chunk length (e.g., 64).
+ eod_token_id (int): EOD token ID.
+ indexed_dataset_infos (Optional[List[Dict]]): Optionally, pre-loaded list of dataset metadata (see `save_indexed_dataset_infos()` for more detail). If not provided, the indexed dataset infos will be loaded from disk.
+
+ Returns:
+ A DBDataset, which is a dataset that wraps the HDF5 chunk index array.
+ """
+ return get_merged_dataset(
+ project_dir, chunk_length, eod_token_id, "valid", indexed_dataset_infos
+ )
+
+
+def get_merged_datasets(project_dir: str, chunk_length: int, eod_token_id: int) -> dict:
+ """Get all merged datasets.
+
+ Args:
+ project_dir (str): Path to Retro project dir.
+ chunk_length (int): GPT chunk length (e.g., 64).
+ eod_token_id (int): EOD token ID.
+
+ Returns:
+ A dict mapping DB type ('sampled', 'train', or 'valid') to the corresponding DBDataset, which is a dataset that wraps the HDF5 chunk index array.
+ """
+ fns = {
+ "sampled": get_merged_sampled_dataset,
+ "train": get_merged_train_dataset,
+ "valid": get_merged_valid_dataset,
+ }
+ datasets = {key: fn(project_dir, chunk_length, eod_token_id) for key, fn in fns.items()}
+ return datasets
diff --git a/megatron/core/datasets/retro/external_libs.py b/megatron/core/datasets/retro/external_libs.py
new file mode 100644
index 0000000000..c057eba25c
--- /dev/null
+++ b/megatron/core/datasets/retro/external_libs.py
@@ -0,0 +1,15 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Required external libraries for Retro preprocessing."""
+
+import importlib
+
+required_libs = ["faiss", "h5py", "transformers"] # for huggingface bert
+
+for lib in required_libs:
+ try:
+ globals()[lib] = importlib.import_module(lib)
+ except ImportError as e:
+ raise Exception(
+ f"Missing one or more packages required for Retro preprocessing: {required_libs}. Tried importing '{lib}'."
+ )
diff --git a/megatron/core/datasets/retro/index/__init__.py b/megatron/core/datasets/retro/index/__init__.py
new file mode 100644
index 0000000000..d069f55f22
--- /dev/null
+++ b/megatron/core/datasets/retro/index/__init__.py
@@ -0,0 +1,11 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""
+Exports:
+
+ - train_index: Train an index on representative vectors.
+ - add_to_index: Add vectors to a trained index.
+ - build_index: Wrapper function that calls above two functions.
+"""
+
+from .build import add_to_index, build_index, train_index
diff --git a/megatron/core/datasets/retro/index/build.py b/megatron/core/datasets/retro/index/build.py
new file mode 100644
index 0000000000..1f310d89c3
--- /dev/null
+++ b/megatron/core/datasets/retro/index/build.py
@@ -0,0 +1,313 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Construct an index.
+
+Constructing an index generally happens in two phases:
+
+ - index.train(): Train an index on a representative set of vectors.
+ - index.add(): Add vectors to an index, to be available for retrieval.
+"""
+
+import os
+import shutil
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from megatron.core.datasets.retro.config import RetroPreprocessingConfig
+from megatron.core.datasets.retro.db.utils import (
+ get_merged_sampled_dataset,
+ get_merged_train_dataset,
+)
+from megatron.core.datasets.retro.external_libs import h5py
+from megatron.core.datasets.retro.utils import GPTToTextDataset
+
+from .factory import IndexFactory
+from .utils import (
+ get_training_data_block_dir,
+ get_training_data_block_paths,
+ get_training_data_merged_path,
+ get_training_data_root_dir,
+)
+
+##################################################
+# Train index.
+##################################################
+
+
+def get_empty_index_path(config: RetroPreprocessingConfig) -> str:
+ """Path of empty index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Path to the empty (trained, but without added samples) vector index.
+ """
+ index = IndexFactory.get_index(config.retro_index_type)
+ empty_index_path = index.get_empty_index_path(config)
+ return empty_index_path
+
+
+def get_block_nload(block_path: str, load_fraction: float) -> int:
+ """Compute number of blocks to load.
+
+ This is computed by multiplying the total number of available blocks with the
+ fraction of blocks to load.
+
+ Args:
+ block_path (str): Path to HDF5 file containing block of data. File must contain key 'data'.
+ load_fraction (float): Fraction (0 < load_fraction <= 1) of block samples to load.
+
+ Returns:
+ Number of block samples to load.
+ """
+ with h5py.File(block_path) as fi:
+ return int(load_fraction * fi["data"].shape[0])
+
+
+def merge_embedding_blocks(config: RetroPreprocessingConfig) -> None:
+ """Merge individual embedding blocks into a single binary mmap file.
+
+ The embeddings are initially stored in block-sized (e.g., ~100k embeddings per
+ block) HDF5 files. These individual block files must be merged into a single
+ file before training, to be based as a numpy mmap array to the index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ if torch.distributed.get_rank() != 0:
+ return
+
+ # Get block, merged paths.
+ load_fraction = config.retro_index_train_load_fraction
+ block_paths = get_training_data_block_paths(config)
+ bin_path = get_training_data_merged_path(config)
+
+ # Skip, if already built.
+ if os.path.exists(bin_path):
+ return
+
+ # Merge blocks.
+ with open(bin_path, "wb") as fo:
+ byte_offset = 0
+ for block_idx, block_path in enumerate(
+ tqdm(
+ block_paths,
+ "merge train embeddings",
+ miniters=len(block_paths) // 10,
+ disable=torch.distributed.get_rank() != 0,
+ )
+ ):
+ with h5py.File(block_path) as fi:
+
+ nload = get_block_nload(block_path, load_fraction)
+ block = np.array(fi["data"][:nload], copy=False)
+
+ fo.write(block.tobytes())
+
+ byte_offset += block.size * block.itemsize
+ fo.seek(byte_offset)
+
+
+def get_text_dataset_for_training(config: RetroPreprocessingConfig) -> GPTToTextDataset:
+ """Convert GPT token chunk dataset to a text dataset for passing to the
+ embedder.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ The text dataset consisting of tokens converted from sampled chunk database.
+ """
+ gpt_dataset = get_merged_sampled_dataset(
+ project_dir=config.retro_project_dir,
+ chunk_length=config.retro_gpt_chunk_length,
+ eod_token_id=config.retro_tokenizers.gpt.eod,
+ )
+ text_dataset = GPTToTextDataset(gpt_dataset, config.retro_tokenizers.gpt)
+ return text_dataset
+
+
+def embed_training_chunks(config: RetroPreprocessingConfig) -> None:
+ """Embed DB chunks.
+
+ Store chunks in blocks on disk. These blocks will later be merged into
+ a single dataset for training the index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ merged_train_data_path = get_training_data_merged_path(config)
+ if os.path.exists(merged_train_data_path):
+ return
+
+ # Get training text dataset.
+ text_dataset = get_text_dataset_for_training(config)
+
+ # Embed dataset.
+ embedder = config.retro_bert_embedders.disk
+ embedder.embed_text_dataset("index", get_training_data_block_dir(config), text_dataset)
+
+ # Merge embeddings.
+ merge_embedding_blocks(config)
+
+
+def train_on_embeddings(config: RetroPreprocessingConfig) -> None:
+ """Train index on embedded DB chunks.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+ index = IndexFactory.get_index(config.retro_index_type)
+ index.train(config)
+
+
+def remove_embeddings(config: RetroPreprocessingConfig) -> None:
+ """Remove embeddings after training.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+ torch.distributed.barrier()
+ if torch.distributed.get_rank() != 0:
+ return
+ empty_index_path = get_empty_index_path(config)
+ assert os.path.isfile(empty_index_path)
+ shutil.rmtree(get_training_data_root_dir(config), ignore_errors=True)
+
+
+def _train_index(config: RetroPreprocessingConfig) -> None:
+ """Train index on DB chunks.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ # Check if trained index already exists.
+ if not os.path.isfile(get_empty_index_path(config)):
+
+ # Embed training chunks.
+ embed_training_chunks(config)
+
+ # Train index on embeddings.
+ train_on_embeddings(config)
+
+ # Wait for (single-process) training to complete.
+ torch.distributed.barrier()
+
+ # Remove embeddings.
+ if config.retro_index_delete_training_embeddings:
+ remove_embeddings(config)
+
+
+def train_index(config: RetroPreprocessingConfig) -> None:
+ """Entry point for training the index.
+
+ We select whether to train a new index, or validate an existing index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ # Train new index.
+ if config.retro_task_validate is None:
+ _train_index(config)
+
+ # Validate existing trained index.
+ else:
+ from .validate import validate_training_embeddings
+
+ validate_training_embeddings(config)
+
+
+##################################################
+# Add to index.
+##################################################
+
+
+def get_text_dataset_for_adding(config: RetroPreprocessingConfig) -> GPTToTextDataset:
+ """Convert GPT token chunk dataset to a text dataset for passing to the
+ embedder.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ The text dataset that consists of tokens converted from the 'train' chunk database. These are the chunks used for retrieval by the pretraining 'train' dataset.
+ """
+ gpt_dataset = get_merged_train_dataset(
+ project_dir=config.retro_project_dir,
+ chunk_length=config.retro_gpt_chunk_length,
+ eod_token_id=config.retro_tokenizers.gpt.eod,
+ )
+ text_dataset = GPTToTextDataset(gpt_dataset, config.retro_tokenizers.gpt)
+ return text_dataset
+
+
+def _add_to_index(config: RetroPreprocessingConfig) -> str:
+ """Add DB chunks to index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Path to the populated index.
+ """
+
+ # Get index.
+ index = IndexFactory.get_index(config.retro_index_type)
+
+ # Get text dataset.
+ text_dataset = get_text_dataset_for_adding(config)
+
+ # Add to index.
+ output_index_path = index.add(config, text_dataset)
+
+ return output_index_path
+
+
+def add_to_index(config: RetroPreprocessingConfig) -> None:
+ """Entry point for adding to the index.
+
+ We select whether to add to a new index, or validate an existing index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ # Add to new index.
+ if config.retro_task_validate is None:
+ _add_to_index(config)
+
+ # Validate existing encodings.
+ else:
+ from .validate import validate_added_encodings
+
+ validate_added_encodings(config)
+
+
+##################################################
+# Build index (train + add).
+##################################################
+
+
+def build_index(config: RetroPreprocessingConfig) -> None:
+ """Build index.
+
+ Building index involves sequentially running stages above:
+ - Train index (on sampled training chunks).
+ - Add to index (on all training chunks).
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ # Train index.
+ train_index(config)
+
+ # Add to index.
+ add_to_index(config)
diff --git a/megatron/core/datasets/retro/index/factory.py b/megatron/core/datasets/retro/index/factory.py
new file mode 100644
index 0000000000..f88084ddb1
--- /dev/null
+++ b/megatron/core/datasets/retro/index/factory.py
@@ -0,0 +1,40 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""The IndexFactory constructs an index from an index type string."""
+
+from megatron.core.datasets.retro.index.index import Index
+
+from .indexes import FaissBaseIndex, FaissParallelAddIndex
+
+
+class IndexFactory:
+ """Get index.
+
+ Index type generally read from argument '--retro-index-ty'.
+ """
+
+ @classmethod
+ def get_index_class(cls, index_type: str) -> type:
+ """Get an index class, given a type string.
+
+ Args:
+ index_type (str): One of 'faiss-base' (naive Faiss index wrapper) or 'faiss-par-add' (Faiss index wrapper with near embarrassingly parallel index.add().
+
+ Returns:
+ An `Index` sub-type corresponding to the `index_type`.
+ """
+ return {"faiss-base": FaissBaseIndex, "faiss-par-add": FaissParallelAddIndex}[index_type]
+
+ @classmethod
+ def get_index(cls, index_type: str) -> Index:
+ """Construct an index from an index type string.
+
+ Args:
+ index_type (str): One of 'faiss-base' (naive Faiss index wrapper) or 'faiss-par-add' (Faiss index wrapper with near embarrassingly parallel index.add().
+
+ Returns:
+ An `Index` instance corresponding to the `index_type`.
+ """
+ index_class = cls.get_index_class(index_type)
+ index = index_class()
+ return index
diff --git a/megatron/core/datasets/retro/index/index.py b/megatron/core/datasets/retro/index/index.py
new file mode 100644
index 0000000000..c6bd13fbee
--- /dev/null
+++ b/megatron/core/datasets/retro/index/index.py
@@ -0,0 +1,133 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Base class for all vector indexes.
+
+A vector index is a type of retrieval database that is queried using vectors,
+and returns vectors that are 'similar' (e.g., by cosine distance) to the query
+vector. The construction and usage of an index generally has the following
+pattern:
+
+ - Train the index on representative vectors.
+ - Add vectors to the index (i.e., vectors available for retrieval)
+ - Query index with new vector, to retrieve similar vector indexes.
+"""
+
+import abc
+import os
+from typing import List, Tuple
+
+import numpy as np
+import torch
+
+from megatron.core.datasets.retro.config import Embedder, RetroPreprocessingConfig
+from megatron.core.datasets.retro.external_libs import faiss
+from megatron.core.datasets.retro.utils import GPTToTextDataset
+
+from .utils import get_index_dir
+
+
+class Index(abc.ABC):
+ """Abstract base class for indexes.
+
+ *Note* : While currently only Faiss-based classes are implemented, in the
+ future, this class will be extended with other types of indexes that have
+ different performance-accuracy trade-offs.
+
+ The primary methods to override are:
+ - train() : Train index on the sampled training chunks.
+ - add() : Add all training chunks to index.
+ """
+
+ @classmethod
+ def make_object_verbose(cls, index: faiss.Index, verbose: bool) -> None:
+ """Make index object verbose.
+
+ Args:
+ index (faiss.Index): Faiss object to set verbose.
+ verbose (bool): Sets whether index should log status updates during training and adding.
+ """
+ assert isinstance(verbose, bool)
+ faiss.ParameterSpace().set_index_parameter(index, "verbose", verbose)
+
+ def get_empty_index_path(self, config: RetroPreprocessingConfig) -> str:
+ """Get file path to empty index (i.e., trained, but unpopulated).
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ File path to empty index (i.e., this index has had index.train() called, but not yet index.add()).
+ """
+ return os.path.join(
+ get_index_dir(config), "empty_%.3f.faissindex" % config.retro_index_train_load_fraction
+ )
+
+ def get_empty_index(self, config: RetroPreprocessingConfig) -> faiss.Index:
+ """Get empty index (i.e., trained, but unpopulated).
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Empty Faiss index, loaded from storage.
+ """
+ return faiss.read_index(self.get_empty_index_path(config))
+
+ def get_added_index_path(self, config: RetroPreprocessingConfig) -> str:
+ """Get file path to index that has been populated with vectors.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ File path to added index (i.e., this index has had both index.train() and index.add() called).
+ """
+ return os.path.join(
+ get_index_dir(config),
+ "added_%.3f_%.3f.faissindex"
+ % (config.retro_index_train_load_fraction, config.retro_index_add_load_fraction),
+ )
+
+ def get_added_index(self, config: RetroPreprocessingConfig) -> faiss.Index:
+ """Get index that has been populated with vectors.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ 'Added' (i.e., populated) Faiss index, loaded from storage.
+ """
+ return faiss.read_index(self.get_added_index_path(config))
+
+ @abc.abstractmethod
+ def train(self, config: RetroPreprocessingConfig) -> None:
+ """Train index on a representative set of vectors.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ @abc.abstractmethod
+ def add(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> None:
+ """Add vectors to index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ text_dataset (GPTToTextDataset): Text dataset that will be embedded and added to the index.
+ """
+
+ def embed_text_dataset_block(
+ self, embedder: Embedder, text_dataset: GPTToTextDataset, _range: Tuple[int, int]
+ ) -> np.ndarray:
+ """Embed a range of a text dataset.
+
+ Args:
+ embedder (Embedder): Embedder used for embedding a text dataset.
+ text_dataset (GPTToTextDataset): Text dataset that will be embedded.
+ _range (Tuple[int, int]): Start/end sample indices within text dataset used for embedding.
+
+ Returns:
+ An array of embeddings, with shape (len(text_dataset), dimension(embedder)).
+ """
+ sub_dataset = torch.utils.data.Subset(text_dataset, range(*_range))
+ return embedder.embed_text_dataset(sub_dataset)
diff --git a/megatron/core/datasets/retro/index/indexes/__init__.py b/megatron/core/datasets/retro/index/indexes/__init__.py
new file mode 100644
index 0000000000..c445909fea
--- /dev/null
+++ b/megatron/core/datasets/retro/index/indexes/__init__.py
@@ -0,0 +1,10 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""
+Exports:
+- FaissBaseIndex: Unoptimized Faiss index wrapper
+- FaissParallelAddIndex: Optimized index.add() for Faiss index.
+"""
+
+from .faiss_base import FaissBaseIndex
+from .faiss_par_add import FaissParallelAddIndex
diff --git a/megatron/core/datasets/retro/index/indexes/faiss_base.py b/megatron/core/datasets/retro/index/indexes/faiss_base.py
new file mode 100644
index 0000000000..c1daf3f533
--- /dev/null
+++ b/megatron/core/datasets/retro/index/indexes/faiss_base.py
@@ -0,0 +1,150 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""
+This class implements a simple, un-optimized wrapper around a Faiss index, that
+implements the Index interface (see ..index.py). While this class is
+instantiable, it is meant to be extended with optimizations in classes that
+inherit from this class (see FaissParAddIndex, for an example).
+"""
+
+import os
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from megatron.core.datasets.retro.config import RetroPreprocessingConfig
+from megatron.core.datasets.retro.external_libs import faiss
+from megatron.core.datasets.retro.index.index import Index
+from megatron.core.datasets.retro.index.utils import (
+ get_training_data_merged_path,
+ num_samples_to_block_ranges,
+)
+from megatron.core.datasets.retro.utils import GPTToTextDataset, log_retro_rank_0
+
+
+class FaissBaseIndex(Index):
+ """Base class for Faiss-base indexes.
+
+ This class wraps a Faiss index, and adds additional functionality for training
+ and adding codes. This base class performs a naive sequential code adding,
+ while the optimized FaissParallelAddIndex class performs a parallel
+ index.add().
+ """
+
+ def _train(self, config: RetroPreprocessingConfig) -> None:
+ """Train index (rank 0's method).
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ assert torch.distributed.get_rank() == 0
+
+ # Set num threads (torch.distributed reset it to 1).
+ faiss.omp_set_num_threads(64)
+
+ empty_index_path = self.get_empty_index_path(config)
+
+ # Index already exists? -> return.
+ if os.path.isfile(empty_index_path):
+ return
+
+ # Load data.
+ merged_path = get_training_data_merged_path(config)
+ inp = np.memmap(merged_path, dtype="f4", mode="r").reshape((-1, config.hidden_size))
+
+ # Init index.
+ index = faiss.index_factory(config.hidden_size, config.retro_index_str)
+
+ # Move to GPU.
+ log_retro_rank_0("> move faiss index to gpu.")
+ index_ivf = faiss.extract_index_ivf(index)
+ clustering_index = faiss.index_cpu_to_all_gpus(faiss.IndexFlatL2(index_ivf.d))
+ index_ivf.clustering_index = clustering_index
+ log_retro_rank_0("> finished moving to gpu.")
+ self.make_object_verbose(index, True)
+ self.make_object_verbose(index_ivf, True)
+ self.make_object_verbose(index_ivf.quantizer, True)
+ self.make_object_verbose(index_ivf.clustering_index, True)
+
+ # Train index.
+ index.train(inp)
+
+ # Save index.
+ faiss.write_index(index, empty_index_path)
+
+ def train(self, config: RetroPreprocessingConfig) -> None:
+ """Train index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ # Single process only.
+ if torch.distributed.get_rank() == 0:
+ self._train(config)
+
+ torch.distributed.barrier()
+
+ def _add(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> None:
+ """Add to index (rank 0's method).
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ text_dataset (GPTToTextDataset): Text dataset that will be embedded and added to the index.
+ """
+
+ assert torch.distributed.get_rank() == 0
+
+ dataset_sample_ranges = num_samples_to_block_ranges(len(text_dataset))
+
+ # Set num threads (torch.distributed reset it to 1).
+ faiss.omp_set_num_threads(64)
+
+ # Bert embedder.
+ embedder = config.bert_embedders.mem
+
+ # Empty/added index paths.
+ empty_index_path = self.get_empty_index_path()
+ added_index_path = self.get_added_index_path()
+
+ # Skip adding, if index exists.
+ if os.path.isfile(added_index_path):
+ return
+
+ # Read trained index.
+ index = faiss.read_index(empty_index_path)
+
+ # Iterate data blocks & add.
+ for sample_range in tqdm(dataset_sample_ranges, "faiss_base.add"):
+
+ # Embed text.
+ embeds = self.embed_text_dataset_block(embedder, text_dataset, sample_range)
+
+ # Add to index.
+ index.add(embeds)
+
+ # Write index.
+ faiss.write_index(index, added_index_path)
+
+ def add(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> str:
+ """Add to index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ text_dataset (GPTToTextDataset): Text dataset that will be embedded and added to the index.
+
+ Returns:
+ File path to the populated index.
+ """
+
+ # Single process only.
+ if torch.distributed.get_rank() == 0:
+ self._add(config, text_dataset)
+
+ # Wait for rank 0.
+ torch.distributed.barrier()
+
+ # Get output index path, for return.
+ return self.get_added_index_path(config)
diff --git a/megatron/core/datasets/retro/index/indexes/faiss_par_add.py b/megatron/core/datasets/retro/index/indexes/faiss_par_add.py
new file mode 100644
index 0000000000..e014217262
--- /dev/null
+++ b/megatron/core/datasets/retro/index/indexes/faiss_par_add.py
@@ -0,0 +1,208 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Multi-process & multi-node version of Faiss's index.add().
+
+This class inherits from FaissBaseIndex, and optimizes the 'add()' method by
+making it multi-node and multi-process, with bit-wise equivalence to
+FaissBaseIndex. This allows 'add()' to scale out to very large datasets, since
+the vast majority of the computational effort is embarrassingly parallel.
+"""
+
+import os
+import shutil
+from typing import Tuple
+
+import numpy as np
+import psutil
+import torch
+from tqdm import tqdm
+
+from megatron.core.datasets.retro.config import Embedder, RetroPreprocessingConfig
+from megatron.core.datasets.retro.external_libs import faiss, h5py
+from megatron.core.datasets.retro.index.utils import get_added_code_paths, get_added_codes_dir
+from megatron.core.datasets.retro.utils import (
+ GPTToTextDataset,
+ get_blocks_by_rank,
+ log_retro_rank_0,
+ retro_makedir,
+)
+
+from .faiss_base import FaissBaseIndex
+
+
+class FaissParallelAddIndex(FaissBaseIndex):
+ """
+ This class parallelizes both 1) encoding vectors, and 2) adding codes to the
+ index. This class is more performant than naive use of Faiss, because most
+ of the computational work is in encoding the vectors, which is an
+ embarassingly parallel operation.
+ """
+
+ def encode_block(
+ self, index: faiss.Index, embedder: Embedder, text_dataset: GPTToTextDataset, block: dict
+ ) -> Tuple[np.ndarray, np.ndarray]:
+ """Encode sub-dataset block, to be later added to index.
+
+ Encode the data subset, generally in blocks of 1M vectors each. For
+ each block, the empty/trained index is loaded, codes are computed
+ via index.sa_encode(), and the resulting codes are saved to disk.
+
+ Args:
+ index (faiss.Index): Faiss index object.
+ embedder (Embedder): Embedder used to embed text dataset.
+ text_dataset (GPTToTextDataset): Text dataset to be embedded and encoded.
+ block (dict): Range information specifying start/end indices within text dataset.
+
+ Returns:
+ A tuple of (embeddings, encodings) for the given block subset of the text dataset.
+ """
+
+ # Embed block.
+ embeddings = self.embed_text_dataset_block(embedder, text_dataset, block["range"])
+
+ # Encode block.
+ log_retro_rank_0("encode.")
+ codes = index.sa_encode(embeddings)
+
+ # Return embeddings for validation purposes.
+ return embeddings, codes
+
+ def save_block(self, config: RetroPreprocessingConfig, block: dict, codes: np.ndarray) -> None:
+ """Save block of codes to disk.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ block (dict): Range information specifying the start/end indices within the encoded text dataset. Here, the 'path' item is used for writing the encodings to storage.
+ codes (np.ndarray): Block of encodings to be saved to storage.
+ """
+ # Save neighbors.
+ log_retro_rank_0("save codes.")
+ retro_makedir(config, os.path.dirname(block["path"]))
+ with h5py.File(block["path"], "w") as f:
+ f.create_dataset("data", data=codes)
+
+ def encode(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> None:
+ """Encode text dataset, to be later added to index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ text_dataset (GPTToTextDataset): Text dataset to be encoded by the index.
+ """
+
+ codes_dir = get_added_codes_dir(config)
+ retro_makedir(config, codes_dir)
+
+ # Index.
+ index = self.get_empty_index(config)
+
+ # Bert embedder.
+ embedder = config.retro_bert_embedders.mem
+
+ # Missing code blocks.
+ def validate(f: h5py.File) -> None:
+ """Validation method for validating loaded encodings.
+
+ Args:
+ f (h5py.File): File that contains encodings.
+ """
+ assert len(f["data"].shape) == 2
+
+ blocks = get_blocks_by_rank(
+ codes_dir, len(text_dataset), config.retro_block_size, validate=validate
+ )
+
+ # Encode each block.
+ for block_index, block in enumerate(blocks.missing):
+
+ if block is not None:
+
+ # Progress.
+ log_retro_rank_0(
+ "encode block %d / %d ... %s."
+ % (block_index, len(blocks.missing), block["path"])
+ )
+
+ # Encode and save.
+ _, codes = self.encode_block(index, embedder, text_dataset, block)
+ self.save_block(config, block, codes)
+
+ # Synchronize progress across all ranks. (for easier observation)
+ log_retro_rank_0(" > waiting for other ranks to finish block.")
+ torch.distributed.barrier()
+
+ def add_codes(self, config: RetroPreprocessingConfig) -> None:
+ """Read codes from disk, and add them to the index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ if torch.distributed.get_rank() != 0:
+ return
+
+ added_index_path = self.get_added_index_path(config)
+ if os.path.exists(added_index_path):
+ return
+
+ # Index.
+ log_retro_rank_0("read empty index.")
+ index = self.get_empty_index(config)
+ index_ivf = faiss.extract_index_ivf(index)
+
+ # Add codes.
+ log_retro_rank_0("add codes.")
+ code_paths = get_added_code_paths(config)
+ pbar = tqdm(code_paths)
+ for code_path in pbar:
+ pbar.set_description(
+ "add codes, mem %.3f gb, %.1f%%"
+ % (psutil.virtual_memory()[3] / 1024**3, psutil.virtual_memory()[2])
+ )
+ with h5py.File(code_path) as f:
+
+ nload = int(config.retro_index_add_load_fraction * f["data"].shape[0])
+ offset = int(os.path.basename(code_path).split("-")[0])
+ xids = np.arange(offset, offset + nload)
+ codes = np.copy(f["data"][:nload])
+ index_ivf.add_sa_codes(codes, xids)
+
+ # Update index's ntotal.
+ index.ntotal = index_ivf.ntotal
+
+ # Write index.
+ log_retro_rank_0("write added index.")
+ faiss.write_index(index, added_index_path)
+
+ def remove_codes(self, config: RetroPreprocessingConfig) -> None:
+ """Remove added codes after adding to index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+ if torch.distributed.get_rank() != 0:
+ return
+ assert os.path.isfile(self.get_added_index_path(config))
+
+ if config.retro_index_delete_added_codes:
+ raise Exception("remove?")
+ shutil.rmtree(get_added_codes_dir(config), ignore_errors=True)
+
+ def add(self, config: RetroPreprocessingConfig, text_dataset: GPTToTextDataset) -> None:
+ """Add vectors to index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ text_dataset (GPTToTextDataset): Text dataset that will be embedded and added to the index.
+ """
+
+ # Encode chunks.
+ self.encode(config, text_dataset)
+
+ # Add codes to index.
+ self.add_codes(config)
+
+ # Wait for (single-process) adding to complete.
+ torch.distributed.barrier()
+
+ # Remove codes.
+ self.remove_codes(config)
diff --git a/megatron/core/datasets/retro/index/utils.py b/megatron/core/datasets/retro/index/utils.py
new file mode 100644
index 0000000000..58229439ae
--- /dev/null
+++ b/megatron/core/datasets/retro/index/utils.py
@@ -0,0 +1,126 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Utilities for building an index."""
+
+import glob
+import os
+from typing import List, Tuple
+
+from megatron.core.datasets.retro.config import RetroPreprocessingConfig
+from megatron.core.datasets.retro.utils import retro_makedir
+
+
+def get_index_dir(config: RetroPreprocessingConfig) -> str:
+ """Create sub-directory for this index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Path to index sub-directory within Retro project.
+ """
+
+ # Directory path.
+ index_dir_path = os.path.join(
+ config.retro_project_dir, "index", config.retro_index_type, config.retro_index_str
+ )
+
+ # Make directory.
+ retro_makedir(config, index_dir_path)
+
+ return index_dir_path
+
+
+def num_samples_to_block_ranges(
+ config: RetroPreprocessingConfig, num_samples: int
+) -> List[Tuple[int, int]]:
+ """Split a range (length num_samples) into sequence of block ranges
+ of size block_size.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ num_samples (int): Split `num_samples` into consecutive block ranges, where each block is size `config.retro_block_size`.
+
+ Returns:
+ A list of tuples where each item is the (start, end) index for a given block.
+ """
+ block_size = config.retro_block_size
+ start_idxs = list(range(0, num_samples, block_size))
+ end_idxs = [min(num_samples, s + block_size) for s in start_idxs]
+ ranges = list(zip(start_idxs, end_idxs))
+ return ranges
+
+
+def get_training_data_root_dir(config: RetroPreprocessingConfig) -> str:
+ """Get root directory for embeddings (blocks and merged data).
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Path to the training data directory, which contains both training embedding blocks and the final merged training embeddings.
+ """
+ return os.path.join(config.retro_project_dir, "index", "train_emb")
+
+
+def get_training_data_block_dir(config: RetroPreprocessingConfig) -> str:
+ """Get directory for of saved embedding blocks.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Path to the directory containing the training embedding blocks, which will be later merged into a single embedding array.
+ """
+ return os.path.join(get_training_data_root_dir(config), "blocks")
+
+
+def get_training_data_block_paths(config: RetroPreprocessingConfig) -> List[str]:
+ """Get paths to saved embedding blocks.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Paths of all training embedding blocks.
+ """
+ return sorted(glob.glob(get_training_data_block_dir(config) + "/*.hdf5"))
+
+
+def get_training_data_merged_path(config: RetroPreprocessingConfig) -> str:
+ """Get path to merged training embeddings.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Path to the merged training embedding binary file.
+ """
+ return os.path.join(
+ get_training_data_root_dir(config),
+ "train_%.3f.bin" % config.retro_index_train_load_fraction,
+ )
+
+
+def get_added_codes_dir(config: RetroPreprocessingConfig) -> str:
+ """Get directory of saved encodings.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Path to the directory containing the vector encodings for adding to the index.
+ """
+ return os.path.join(get_index_dir(config), "add_codes")
+
+
+def get_added_code_paths(config: RetroPreprocessingConfig) -> List[str]:
+ """Get paths to all saved encodings.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ Paths of all vector encoding blocks, for adding to the index.
+ """
+ return sorted(glob.glob(get_added_codes_dir(config) + "/*.hdf5"))
diff --git a/megatron/core/datasets/retro/index/validate.py b/megatron/core/datasets/retro/index/validate.py
new file mode 100644
index 0000000000..57306707c4
--- /dev/null
+++ b/megatron/core/datasets/retro/index/validate.py
@@ -0,0 +1,191 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Validate an index's data.
+
+This module contains functionality for checking for bitwise equality across code
+changes. The training and adding steps of index construction can be validated
+separately. The following high-level checks are supported:
+
+ - Training: Validate that saved training embeddings are bitwise equal with a
+ sample set of freshly computed embeddings. (*Note*:
+ `--no-retro-index-delete-training-embeddings` must be used.)
+ - Adding: Validate that the saved encodings are bitwise equal with a sample of
+ sample set of freshly computed encodings. (*Note*:
+ `--no-retro-index-delete-added-codes` must be used.)
+"""
+
+import typing
+
+import numpy as np
+import torch
+from torch.utils.data import Subset
+
+from megatron.core.datasets.retro.config import RetroPreprocessingConfig
+from megatron.core.datasets.retro.external_libs import h5py
+from megatron.core.datasets.retro.utils import (
+ GPTToTextDataset,
+ get_blocks_by_rank,
+ log_retro_rank_0,
+)
+
+from .build import get_text_dataset_for_adding, get_text_dataset_for_training
+from .factory import IndexFactory
+from .utils import get_added_codes_dir, get_training_data_block_dir
+
+##################################################
+# Validate trained index.
+##################################################
+
+
+def validate_training_embeddings(config: RetroPreprocessingConfig) -> None:
+ """Validate training embeddings.
+
+ Steps:
+ - Randomly sample subset of text dataset blocks.
+ - Embed each block.
+ - Compare against saved embeddings.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ # Training text dataset.
+ text_dataset = get_text_dataset_for_training(config)
+
+ # Sample existing blocks.
+ blocks = get_blocks_by_rank(
+ dirname=get_training_data_block_dir(config),
+ n_samples=len(text_dataset),
+ block_size=config.retro_block_size,
+ validate=None,
+ sample=config.retro_task_validate,
+ )
+
+ assert blocks.n_missing_world == 0
+
+ # Embed & validate blocks.
+ embedder = config.retro_bert_embedders.mem
+ for block_idx, block in enumerate(blocks.existing):
+
+ # Missing block lists are extended with None to have equal-length
+ # lists. Skip the Nones.
+ if block is not None:
+
+ # Progress. (*note*: move world progress to here.)
+ log_retro_rank_0(
+ "embed training block %d / %d ... %s."
+ % (block_idx, len(blocks.existing), block["path"])
+ )
+
+ # Load existing block embeddings.
+ with h5py.File(block["path"]) as f:
+ existing_embeddings = np.copy(f["data"])
+
+ # Embed block.
+ sub_dataset = Subset(text_dataset, range(*block["range"]))
+ embeddings = embedder.embed_text_dataset(sub_dataset, "train")
+
+ # Check equality.
+ log_retro_rank_0(" > validate.")
+ assert np.array_equal(existing_embeddings, embeddings)
+
+ # Synchronize progress across all ranks. (for easier observation)
+ log_retro_rank_0(" > waiting for other ranks to finish block.")
+ torch.distributed.barrier()
+
+ log_retro_rank_0(" > finished validating training embeddings.")
+
+
+##################################################
+# Validate filled index.
+##################################################
+
+
+def validate_added_encodings(config: RetroPreprocessingConfig) -> None:
+ """Validate added encodings.
+
+ Steps:
+ - Randomly sample subset of text dataset blocks.
+ - Encode each block.
+ - Compare against saved encodings.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ # Index.
+ index = IndexFactory.get_index(config.retro_index_type)
+ inner_index = index.get_empty_index(config)
+
+ # Text dataset.
+ text_dataset = get_text_dataset_for_adding(config)
+
+ # Sample existing blocks.
+ def validate(f: h5py.File) -> None:
+ """Validation method for validating encoding blocks.
+
+ Args:
+ f (h5py.File): File with block of encodings.
+ """
+ assert len(f["data"].shape) == 2
+
+ blocks = get_blocks_by_rank(
+ dirname=get_added_codes_dir(config),
+ n_samples=len(text_dataset),
+ block_size=config.retro_block_size,
+ validate=validate,
+ sample=config.retro_task_validate,
+ )
+
+ assert blocks.n_missing_world == 0
+
+ # Encode and validate blocks.
+ embedder = config.retro_bert_embedders.mem
+ for block_idx, block in enumerate(blocks.existing):
+
+ if block is not None:
+
+ # Progress.
+ log_retro_rank_0(
+ "encode block %d / %d ... %s." % (block_idx, len(blocks.existing), block["path"])
+ )
+
+ # Load existing codes.
+ with h5py.File(block["path"]) as f:
+ existing_codes = np.copy(f["data"])
+
+ # Encode block.
+ embeddings, codes = index.encode_block(inner_index, embedder, text_dataset, block)
+
+ # Check equality.
+ log_retro_rank_0(" > validate.")
+ assert np.array_equal(existing_codes, codes)
+
+ # Synchronize progress across all ranks. (for easier observation)
+ log_retro_rank_0(" > waiting for other ranks to finish block.")
+ torch.distributed.barrier()
+
+ log_retro_rank_0(" > finished validating added encodings.")
+
+
+##################################################
+# Validate index (trained + filled).
+##################################################
+
+
+def validate_index(config: RetroPreprocessingConfig) -> None:
+ """Validate index.
+
+ Validating index involves sequentially running stages above:
+ - Validate trained index.
+ - Validate filled index.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ # Validate training embeddings.
+ validate_training_embeddings(config)
+
+ # Validate added codes.
+ validate_added_encodings(config)
diff --git a/megatron/core/datasets/retro/query/__init__.py b/megatron/core/datasets/retro/query/__init__.py
new file mode 100644
index 0000000000..ac9483373c
--- /dev/null
+++ b/megatron/core/datasets/retro/query/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/datasets/retro/query/gpt_chunk_dataset.py b/megatron/core/datasets/retro/query/gpt_chunk_dataset.py
new file mode 100644
index 0000000000..6191a30a31
--- /dev/null
+++ b/megatron/core/datasets/retro/query/gpt_chunk_dataset.py
@@ -0,0 +1,109 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""
+A GPTChunkDataset is a wrapper around a regular GPTDataset, that sequentially
+chunks the sample tokens into `retro_chunk_length` sized smaller samples.
+
+For example, if the GPTDataset has 100 samples and a sequence length of 2048, and
+retro_chunk_length is 64, then the GPTChunkDataset will contain 100*(2048/64) =
+3200 samples, each with length 64.
+"""
+
+import torch
+
+from megatron.core.datasets.gpt_dataset import GPTDataset
+from megatron.core.datasets.retro.utils import get_num_chunks_per_sample
+
+from .utils import get_neighbor_dir
+
+
+class GPTChunkDataset(torch.utils.data.Dataset):
+ """Pretraining chunk dataset wraps a standard GPT dataset.
+
+ This dataset conceptually divides each sample (e.g., length 2048)
+ into chunks (e.g., length 64) and restructures them into a list of
+ chunks (e.g., length num_samples * num_chunks_per_sample).
+
+ Args:
+ sample_dataset (GPTDataset): Original GPT dataset, with `sequence_length` size samples.
+ sample_length (int): Alias for `sequence_length`.
+ chunk_length (int): Retro chunk length (e.g., 64).
+ """
+
+ def __init__(self, sample_dataset: GPTDataset, sample_length: int, chunk_length: int):
+
+ super().__init__()
+
+ self.sample_dataset = sample_dataset
+ self.chunk_length = chunk_length
+ self.n_chunks_per_sample = get_num_chunks_per_sample(sample_length, chunk_length)
+ self.n_samples = len(sample_dataset)
+ self.n_chunks = self.n_samples * self.n_chunks_per_sample
+
+ def __len__(self) -> int:
+ """Get dataset length.
+
+ Returns:
+ Dataset length.
+ """
+ return self.n_chunks
+
+ def __getitem__(self, idx: int) -> dict:
+ """Get sample, including represented document IDs.
+
+ Args:
+ idx (int): Sample index.
+
+ Returns:
+ A sample, which contains both the chunk-length token sample ('text') along with all document_ids ('doc_ids') contained withing the full `sequence_length` sample.
+ """
+
+ # Convert global chunk index to global sample index & local chunk index.
+ sample_idx = idx // self.n_chunks_per_sample
+ chunk_idx = idx % self.n_chunks_per_sample
+
+ # Extract sample data.
+ sample = self.sample_dataset[sample_idx]
+ sample_token_ids = sample["text"]
+ sample_doc_ids = sample["document_ids"]
+
+ # Chunk start/end token idxs.
+ token_start_idx = chunk_idx * self.chunk_length
+ token_end_idx = token_start_idx + self.chunk_length
+ chunk_token_ids = sample_token_ids[token_start_idx:token_end_idx]
+
+ # Sample.
+ return {"doc_ids": sample_doc_ids, "text": chunk_token_ids}
+
+
+def build_gpt_chunk_datasets_from_gpt_datasets(
+ project_dir: str, gpt_datasets: dict, sample_length: int, chunk_length: int
+) -> dict:
+ """Get train, valid, test GPT chunk datasets.
+
+ Args:
+ project_dir (str): Retro project dir.
+ gpt_datasets (dict): Mapping of 'train', 'valid', and 'test' GPT datasets (original, unchunked datasets).
+ sample_length (int): Alias of `sequence_length`.
+ chunk_length (int): Retro chunk length (e.g., 64).
+
+ Returns:
+ A ?
+ """
+
+ # GPT chunk datasets.
+ chunk_datasets = {
+ key: (
+ {
+ "dataset": GPTChunkDataset(sample_ds, sample_length, chunk_length),
+ "neighbor_dir": get_neighbor_dir(project_dir, key, sample_ds),
+ "num_active_chunks": num_active_samples
+ * get_num_chunks_per_sample(sample_length, chunk_length),
+ }
+ if sample_ds
+ else None
+ )
+ for key, (sample_ds, num_active_samples) in gpt_datasets.items()
+ }
+
+ return chunk_datasets
diff --git a/megatron/core/datasets/retro/query/multi_split_gpt_dataset.py b/megatron/core/datasets/retro/query/multi_split_gpt_dataset.py
new file mode 100644
index 0000000000..97a891fd14
--- /dev/null
+++ b/megatron/core/datasets/retro/query/multi_split_gpt_dataset.py
@@ -0,0 +1,107 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""A MultiSplitGPTDataset can handle multiple intersecting split strings, as well
+as returning all of the document IDs of a sample."""
+
+import logging
+from dataclasses import dataclass
+from typing import Dict, List
+
+import numpy
+
+from megatron.core.datasets.blended_megatron_dataset_config import (
+ convert_split_vector_to_split_matrix,
+ parse_and_normalize_split,
+)
+from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig
+from megatron.core.datasets.indexed_dataset import IndexedDataset
+from megatron.core.datasets.utils import Split
+from megatron.core.utils import log_single_rank
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MultiSplitGPTDatasetConfig(GPTDatasetConfig):
+ """Configuration object for Megatron Core blended and Retro datasets.
+
+ Args:
+ return_document_ids (bool): Whether to return the document ids when querying the dataset. Turn this option on during preprocessing.
+ split_preprocessing (str): The Retro preprocessing split string. It follows the same pattern convention as 'split'. Not to be used with 'blend_per_split'.
+ """
+
+ return_document_ids: bool = None
+
+ split_preprocessing: str = None
+
+ def __post_init__(self) -> None:
+ """Validate config attributes."""
+ super().__post_init__()
+ assert self.split is not None, "the Retro data pipeline does not support 'blend_per_split'"
+ assert self.return_document_ids is not None, "this attribute must be user defined"
+ assert self.split_preprocessing is not None, "this attribute must be user defined"
+ split_vector = parse_and_normalize_split(self.split)
+ split_preprocessing_vector = parse_and_normalize_split(self.split_preprocessing)
+ if not numpy.allclose(split_vector, split_preprocessing_vector):
+ self.split_matrix = convert_split_vector_to_split_matrix(
+ split_vector, split_preprocessing_vector
+ )
+ log_single_rank(
+ logger,
+ logging.WARNING,
+ f"split =/= split_preprocessing. Let split_matrix = {self.split_matrix}",
+ )
+
+
+class MultiSplitGPTDataset(GPTDataset):
+ """Retro's customized GPT dataset.
+
+ Args:
+ indexed_dataset (IndexedDataset): The IndexedDataset around which to build the MegatronDataset.
+ dataset_path (str): The real path on disk to the dataset, for bookkeeping.
+ indexed_indices (numpy.ndarray): The set of the documents indices to expose.
+ num_samples (int): The number of samples to draw from the indexed dataset.
+ index_split (Split): The indexed_indices Split.
+ config (MultiSplitGPTDatasetConfig): The Retro-specific container for all config sourced parameters.
+ """
+
+ def __init__(
+ self,
+ indexed_dataset: IndexedDataset,
+ dataset_path: str,
+ indexed_indices: numpy.ndarray,
+ num_samples: int,
+ index_split: Split,
+ config: MultiSplitGPTDatasetConfig,
+ ) -> None:
+ super().__init__(
+ indexed_dataset, dataset_path, indexed_indices, num_samples, index_split, config
+ )
+
+ def __getitem__(self, idx: int) -> Dict[str, numpy.ndarray]:
+ """Get dataset sample.
+
+ Args:
+ idx (int): The index into the dataset.
+
+ Returns:
+ Dict[str, numpy.ndarray]: The text ids and (optionally) the document ids wrapped in a dictionary.
+ """
+ text, document_ids = self._query_document_sample_shuffle_indices(idx)
+ if self.config.return_document_ids:
+ return {"text": text, "document_ids": document_ids}
+ else:
+ return {"text": text}
+
+ @staticmethod
+ def _key_config_attributes() -> List[str]:
+ """Add custom attributes for building unique dataset hash.
+
+ The preprocessing split used for preprocessing will constrain the samples available for pretraining.
+
+ Returns:
+ List[str]: The key config attributes.
+ """
+ return super(MultiSplitGPTDataset, MultiSplitGPTDataset)._key_config_attributes() + [
+ "split_preprocessing"
+ ]
diff --git a/megatron/core/datasets/retro/query/query.py b/megatron/core/datasets/retro/query/query.py
new file mode 100644
index 0000000000..9da3381712
--- /dev/null
+++ b/megatron/core/datasets/retro/query/query.py
@@ -0,0 +1,393 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Entry point for querying an index using a GPTChunkDataset.
+
+Querying involves:
+
+ - Iterate all chunks in the GPTChunkDataset.
+ - Query index for neighbor chunk IDs (i.e., chunks from the chunk database).
+ - Save neighbor chunk IDs to disk, for use in building a RetroDataset sample
+ during pretraining.
+"""
+
+import os
+import time
+import typing
+
+import numpy as np
+import psutil
+import torch
+from tqdm import tqdm
+
+from megatron.core.datasets.retro.config import RetroPreprocessingConfig
+from megatron.core.datasets.retro.db.dataset import DBDataset
+from megatron.core.datasets.retro.db.utils import (
+ get_merged_train_dataset as get_db_merged_train_dataset,
+)
+from megatron.core.datasets.retro.external_libs import faiss, h5py
+from megatron.core.datasets.retro.index.factory import IndexFactory
+from megatron.core.datasets.retro.index.index import Index
+from megatron.core.datasets.retro.index.utils import get_index_dir
+from megatron.core.datasets.retro.query.gpt_chunk_dataset import GPTChunkDataset
+from megatron.core.datasets.retro.utils import (
+ GPTToTextDataset,
+ get_blocks_by_rank,
+ log_retro_rank_0,
+ retro_makedir,
+)
+
+from .gpt_chunk_dataset import build_gpt_chunk_datasets_from_gpt_datasets
+
+
+def get_index(config: RetroPreprocessingConfig, ondisk: bool = False) -> faiss.Index:
+ """Read index from disk.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ ondisk (bool): If `ondisk = True`, memory map the index. (For debugging purposes only; very non-performant.)
+
+ Returns:
+ A Faiss index, loaded from storage.
+ """
+
+ # Load index.
+ index_wrapper = IndexFactory.get_index(config.retro_index_type)
+ index_dir = get_index_dir(config)
+ added_index_path = index_wrapper.get_added_index_path(config)
+ if ondisk:
+ index = faiss.read_index(added_index_path, faiss.IO_FLAG_MMAP)
+ else:
+ index = faiss.read_index(added_index_path)
+
+ # Search parameters.
+ faiss.ParameterSpace().set_index_parameter(index, "efSearch", config.retro_query_ef_search)
+ faiss.ParameterSpace().set_index_parameter(index, "nprobe", config.retro_query_nprobe)
+
+ return index
+
+
+def embed_block(
+ config: RetroPreprocessingConfig, gpt_dataset: GPTChunkDataset, block: dict
+) -> np.ndarray:
+ """Embed block of chunks.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ gpt_dataset (GPTChunkDataset): Chunk dataset to be embedded.
+ block (dict): Range information containing start/end indices of subset of chunk dataset.
+
+ Returns:
+ Embeddings array, with shape (len(block["range"]), dimension(embedder)).
+ """
+ text_block_dataset = torch.utils.data.Subset(
+ GPTToTextDataset(gpt_dataset, config.retro_tokenizers.gpt), range(*block["range"])
+ )
+ return config.retro_bert_embedders.mem.embed_text_dataset(text_block_dataset)
+
+
+def query_embeddings(
+ config: RetroPreprocessingConfig,
+ db_dataset: DBDataset,
+ index: Index,
+ embeddings: np.ndarray,
+ chunk_id_range: range,
+ sample_map: dict,
+ n_chunks_per_sample: int,
+ verbose: bool = True,
+) -> typing.Tuple[np.ndarray, np.ndarray]:
+ """Query neighbors of a block of embeddings.
+
+ Querying includes:
+ - Query index for neighbor chunk IDs.
+ - Filter chunk IDs that have the same document ID as the queried embedding.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ db_dataset (DBDataset): Dataset containing chunk database entries.
+ index (Index): Vector index populated with chunk database indices.
+ embeddings (np.ndarray): Embeddings from GPT chunk dataset.
+ chunk_id_range (range): Chunk ID range from GPT chunk dataset.
+ sample_map (dict): Mapping of sample_idx to dataset_idx and document_ids. Used for document filtering.
+ n_chunks_per_sample (int): Number of chunks per sample (e.g., sequence_length / chunk_length).
+ verbose (bool): Log querying progress.
+
+ Returns:
+ A tuple of original (unfiltered) neighbor IDs, and filtered (by document ID) neighbor IDs.
+ """
+
+ # Query neighbor ids.
+ if verbose:
+ log_retro_rank_0("search.")
+ t = time.time()
+ assert index.ntotal > 0, "check we don't accidentally have an empty index."
+ _, query_neighbor_ids = index.search(embeddings, config.retro_query_num_neighbors_query)
+ if verbose:
+ log_retro_rank_0(" time : %.3f sec." % (time.time() - t))
+
+ # Filter banned neighbor ids.
+ if verbose:
+ log_retro_rank_0("filter banned neighbor ids.")
+ filtered_neighbor_ids = np.full(
+ shape=(len(query_neighbor_ids), config.retro_query_num_neighbors_save),
+ fill_value=-1,
+ dtype="int64",
+ )
+ min_chunk_id, max_chunk_id = chunk_id_range
+ for chunk_id in range(min_chunk_id, max_chunk_id):
+
+ sample_id = chunk_id // n_chunks_per_sample
+ sample = sample_map[sample_id]
+ sample_dataset_idx = sample["dataset_idx"].item()
+ sample_doc_ids = sample["doc_ids"].tolist()
+ sample_doc_tuples = [(sample_dataset_idx, d) for d in sample_doc_ids]
+
+ # Get valid neighbors (!= -1).
+ query_row = [i for i in query_neighbor_ids[chunk_id - min_chunk_id] if i >= 0]
+
+ # Filter row.
+ filtered_row = [
+ i
+ for i in query_row
+ if tuple(db_dataset.doc_tuples[i].tolist()) not in sample_doc_tuples
+ ]
+ filtered_row = filtered_row[: config.retro_query_num_neighbors_save]
+ filtered_row += [-1] * (config.retro_query_num_neighbors_save - len(filtered_row))
+ filtered_neighbor_ids[chunk_id - min_chunk_id] = filtered_row
+
+ return query_neighbor_ids, filtered_neighbor_ids
+
+
+def query_embedding_block(
+ config: RetroPreprocessingConfig,
+ db_dataset: DBDataset,
+ index: Index,
+ embeddings: np.ndarray,
+ chunk_id_range: range,
+ sample_map: dict,
+ n_chunks_per_sample: int,
+) -> typing.Tuple[np.ndarray, np.ndarray]:
+ """Query a block of embeddings.
+
+ The block is broken into smaller sub-blocks, for easier tracking of progress.
+ Both the raw neighbor IDs and the filtered neighbor IDs (i.e., chunks with the
+ same document ID are removed) are collected.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ db_dataset (DBDataset): Dataset containing chunk database entries.
+ index (Index): Vector index populated with chunk database indices.
+ embeddings (np.ndarray): Embeddings from GPT chunk dataset.
+ chunk_id_range (range): Chunk ID range from GPT chunk dataset.
+ sample_map (dict): Mapping of sample_idx to dataset_idx and document_ids. Used for document filtering.
+ n_chunks_per_sample (int): Number of chunks per sample (e.g., sequence_length / chunk_length).
+
+ Returns:
+ A tuple of original (unfiltered) neighbor IDs, and filtered (by document ID) neighbor IDs.
+ """
+
+ query_neighbor_ids = []
+ filtered_neighbor_ids = []
+
+ # Query in sub-blocks.
+ partial_block_size = 1000
+ for partial_start_idx in tqdm(
+ range(0, len(embeddings), partial_block_size),
+ " search",
+ miniters=(len(embeddings) // partial_block_size) // 10,
+ disable=torch.distributed.get_rank() != 0,
+ ):
+ partial_end_idx = min(len(embeddings), partial_start_idx + partial_block_size)
+ partial_embeddings = embeddings[partial_start_idx:partial_end_idx]
+ partial_chunk_id_range = (
+ chunk_id_range[0] + partial_start_idx,
+ chunk_id_range[0] + partial_end_idx,
+ )
+ partial_query_neighbor_ids, partial_filtered_neighbor_ids = query_embeddings(
+ config,
+ db_dataset,
+ index,
+ partial_embeddings,
+ partial_chunk_id_range,
+ sample_map,
+ n_chunks_per_sample,
+ verbose=False,
+ )
+ query_neighbor_ids.append(partial_query_neighbor_ids)
+ filtered_neighbor_ids.append(partial_filtered_neighbor_ids)
+
+ # Concatenate.
+ query_neighbor_ids = np.concatenate(query_neighbor_ids, axis=0)
+ filtered_neighbor_ids = np.concatenate(filtered_neighbor_ids, axis=0)
+
+ return query_neighbor_ids, filtered_neighbor_ids
+
+
+def query_block_neighbors(
+ config: RetroPreprocessingConfig,
+ db_dataset: DBDataset,
+ query_dataset: GPTChunkDataset,
+ index: Index,
+ block: dict,
+) -> None:
+ """Query neighbors of a dataset block (i.e., range).
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ db_dataset (DBDataset): Dataset containing chunk database entries.
+ query_dataset (GPTChunkDataset): GPT chunk dataset to be queried.
+ index (Index): Vector index populated with chunk database indices.
+ block (dict): Range information containing start/end indices for querying GPT chunk dataset.
+ """
+
+ n_chunks_per_sample = query_dataset.n_chunks_per_sample
+
+ # Sample map.
+ sample_ids = sorted(
+ list(set(chunk_id // n_chunks_per_sample for chunk_id in range(*block["range"])))
+ )
+ sample_map = {}
+ for i in sample_ids:
+ sample = query_dataset.sample_dataset[i]
+ sample_map[i] = {"dataset_idx": sample["dataset_id"], "doc_ids": sample["document_ids"]}
+
+ # Embed block.
+ embeddings = embed_block(config, query_dataset, block)
+
+ # Query embeddings.
+ _, filtered_neighbor_ids = query_embedding_block(
+ config, db_dataset, index, embeddings, block["range"], sample_map, n_chunks_per_sample
+ )
+
+ if config.retro_task_validate is None:
+ # Save neighbors.
+ log_retro_rank_0("save neighbors.")
+ retro_makedir(config, os.path.dirname(block["path"]))
+ f = h5py.File(block["path"], "w")
+ f.create_dataset("neighbors", data=filtered_neighbor_ids)
+ f.close()
+
+ else:
+ # Validate neighbors.
+ with h5py.File(block["path"]) as f:
+ existing_neighbor_ids = np.copy(f["neighbors"])
+ assert np.array_equal(existing_neighbor_ids, filtered_neighbor_ids)
+
+
+def query_dataset_neighbors(
+ config: RetroPreprocessingConfig,
+ db_dataset: DBDataset,
+ query_dataset: GPTChunkDataset,
+ num_active_chunks: int,
+ prefix: str,
+ neighbor_dir: str,
+ index: Index,
+) -> None:
+ """Query neighbors of each chunk within a dataset.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ db_dataset (DBDataset): Dataset containing chunk database entries.
+ query_dataset (GPTChunkDataset): GPT chunk dataset to be queried.
+ num_active_chunks (int): The 'active' chunks are the subset of the GPT chunk dataset that aren't being queried. This argument is used when validating the correctness of a subset of the GPT chunk dataset.
+ prefix (str): Extra string for logging progress.
+ neighbor_dir (str): File path to directory for saving neighbor IDs.
+ index (Index): Vector index populated with chunk database indices.
+ """
+
+ def validate(f: h5py.File) -> None:
+ """Validation method for validating saved neighbor IDs.
+
+ Args:
+ f (h5py.File): File containing save neighbor IDs.
+ """
+ assert (
+ f["neighbors"].shape[1] == config.retro_query_num_neighbors_save
+ ), "neighbors.shape == %s; num_neighbors_target == %d." % (
+ str(f["neighbors"].shape),
+ config.retro_num_neighbors_target,
+ )
+
+ if config.retro_task_validate is None:
+ retro_makedir(config, neighbor_dir)
+ blocks = get_blocks_by_rank(
+ neighbor_dir, num_active_chunks, config.retro_block_size, validate=validate
+ )
+ active_blocks = blocks.missing
+ else:
+ blocks = get_blocks_by_rank(
+ neighbor_dir,
+ num_active_chunks,
+ config.retro_block_size,
+ validate=validate,
+ sample=config.retro_task_validate,
+ )
+ assert blocks.n_missing_world == 0
+ active_blocks = blocks.existing
+
+ # Query each block.
+ for block_index, block in enumerate(active_blocks):
+
+ if block is not None:
+
+ # Progress.
+ log_retro_rank_0(
+ "%squery '%s' block %d / %d ... %s ... mem %.3f gb, %.1f%%."
+ % (
+ "" if config.retro_task_validate is None else "[validate] ",
+ prefix,
+ block_index,
+ len(active_blocks),
+ os.path.basename(block["path"]),
+ psutil.virtual_memory()[3] / 1024**3,
+ psutil.virtual_memory()[2],
+ )
+ )
+
+ # Query block neighbors.
+ query_block_neighbors(config, db_dataset, query_dataset, index, block)
+
+ # Synchronize progress across all ranks. (for easier observation)
+ log_retro_rank_0(" > waiting for other ranks to finish block.")
+ torch.distributed.barrier()
+
+
+def query_neighbors(config: RetroPreprocessingConfig) -> None:
+ """Query pretraining datasets (train & valid).
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ """
+
+ # Num threads.
+ faiss.omp_set_num_threads(64)
+
+ # Load chunk db dataset.
+ log_retro_rank_0("load chunk db dataset.")
+ db_dataset = get_db_merged_train_dataset(
+ project_dir=config.retro_project_dir,
+ chunk_length=config.retro_gpt_chunk_length,
+ eod_token_id=config.retro_tokenizers.gpt.eod,
+ )
+ db_dataset.load_doc_tuples()
+
+ # Load index.
+ log_retro_rank_0(" > get index.")
+ index = get_index(config)
+
+ # Query each (i.e., train, valid, test) dataset.
+ log_retro_rank_0(" > query.")
+ for prefix, info in vars(config.retro_gpt_chunk_datasets).items():
+ if info is None:
+ continue
+ log_retro_rank_0(
+ " > query '%s' dataset ... %d samples." % (prefix, info["num_active_chunks"])
+ )
+ query_dataset_neighbors(
+ config,
+ db_dataset,
+ info["dataset"],
+ info["num_active_chunks"],
+ prefix,
+ info["neighbor_dir"],
+ index,
+ )
diff --git a/megatron/core/datasets/retro/query/retro_dataset.py b/megatron/core/datasets/retro/query/retro_dataset.py
new file mode 100644
index 0000000000..6c3b9ae60c
--- /dev/null
+++ b/megatron/core/datasets/retro/query/retro_dataset.py
@@ -0,0 +1,238 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""
+A RetroDataset wraps both:
+
+ - A GPTDataset (which is nested as GPTChunkDataset -> MultiSplitGPTDataset ->
+ GPTDataset).
+ - Neighbor IDs of chunks in the chunk database, that were saved during
+ preprocessing.
+
+Both the GPT sample data and the neighbor IDs are returned within a sample from
+this dataset.
+"""
+
+import os
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+import torch
+
+from megatron.core.datasets.retro.db.dataset import DBDataset
+from megatron.core.datasets.retro.db.utils import get_merged_train_dataset as get_db_dataset
+from megatron.core.datasets.retro.external_libs import h5py
+from megatron.core.datasets.retro.utils import BlockPathMap, log_retro_rank_0
+from megatron.core.models.retro import RetroConfig
+
+from .gpt_chunk_dataset import GPTChunkDataset, build_gpt_chunk_datasets_from_gpt_datasets
+from .utils import get_query_dir
+
+
+class RetroDataset(torch.utils.data.Dataset):
+ """Dataset of retro samples.
+
+ Each sample contains the original GPT sample, along with the token IDs
+ of each neighbor of each chunk within the sequence. Neighbor array has
+ shape (num_chunks_per_sample, num_neighbors, num_retrieved_tokens).
+
+ ** Note: chunk dataset wraps original GPT dataset (see gpt_chunk_dataset.py).
+
+ Args:
+ num_queried_samples (int): Total number of queried samples.
+ num_neighbors (int): Total number of saved neighbors.
+ num_retrieved_chunks (int): Number of retrieved chunks (e.g., 2 for neighbor + continuation).
+ block_size (int): Number of neighbor entries per file.
+ db_dataset (DBDataset): Chunk database used for retrieval.
+ chunk_dataset (GPTChunkDataset): GPT chunk dataset, which is a wrapper around a standard GPT dataset that breaks each sample into chunks.
+ neighbor_path_map (BlockPathMap): Mapping of neighbor ID to file path.
+ """
+
+ def __init__(
+ self,
+ num_queried_samples: int,
+ num_neighbors: int,
+ num_retrieved_chunks: int,
+ block_size: int,
+ db_dataset: DBDataset,
+ chunk_dataset: GPTChunkDataset,
+ neighbor_path_map: BlockPathMap,
+ ):
+ super().__init__()
+
+ self.num_queried_samples = num_queried_samples
+ self.num_neighbors = num_neighbors
+ self.num_retrieved_chunks = num_retrieved_chunks
+ self.block_size = block_size
+ self.db_dataset = db_dataset
+ self.chunk_dataset = chunk_dataset
+ self.neighbor_path_map = neighbor_path_map
+
+ def __len__(self) -> int:
+ """Dataset length.
+
+ Returns:
+ Number of samples in dataset.
+ """
+ return len(self.chunk_dataset.sample_dataset)
+
+ def __getitem__(self, sample_idx: int) -> dict:
+ """Get dataset sample.
+
+ Args:
+ sample_idx (int): Index of sample in dataset.
+
+ Returns:
+ A dict consisting of GPT sample (attribute 'text') and corresponding neighbor chunk IDs ('neighbor_chunks', for indexing chunk database) and neighbor token IDs (corresponding chunk database GPT tokens).
+ """
+ n_chunks_per_sample = self.chunk_dataset.n_chunks_per_sample
+
+ # Wrap sample idx around number of queried samples.
+ sample_idx = sample_idx % self.num_queried_samples
+
+ # Get standard sample.
+ sample = self.chunk_dataset.sample_dataset[sample_idx]
+
+ # Sample idx to chunk idxs.
+ chunk_idxs = list(
+ range(sample_idx * n_chunks_per_sample, (sample_idx + 1) * n_chunks_per_sample)
+ )
+
+ # Collect retrieved tokens.
+ all_retrieved_chunk_ids = []
+ all_retrieved_token_ids = []
+ for chunk_idx in chunk_idxs:
+
+ # Neighbor chunk ids.
+ neighbor_path = self.neighbor_path_map[chunk_idx]
+ with h5py.File(neighbor_path, "r") as f:
+ neighbor_chunk_ids = f["neighbors"][
+ chunk_idx % self.block_size, : self.num_neighbors
+ ].tolist()
+
+ # Retrieved (neighbor + continuation) token ids.
+ retrieved_chunk_ids = []
+ retrieved_token_ids = []
+ for neighbor_chunk_id in neighbor_chunk_ids:
+ current_chunk_ids = [
+ i % len(self.db_dataset)
+ for i in range(neighbor_chunk_id, neighbor_chunk_id + self.num_retrieved_chunks)
+ ]
+ current_token_ids = [self.db_dataset[ci]["text"] for ci in current_chunk_ids]
+ retrieved_chunk_ids.append(current_chunk_ids)
+ retrieved_token_ids.append(current_token_ids)
+
+ # Collect retrieved tokens.
+ all_retrieved_chunk_ids.append(retrieved_chunk_ids)
+ all_retrieved_token_ids.append(retrieved_token_ids)
+
+ # Reshape retrieved tokens.
+ all_retrieved_chunk_ids = np.array(all_retrieved_chunk_ids).reshape(
+ (n_chunks_per_sample, self.num_neighbors, -1)
+ )
+ all_retrieved_token_ids = np.array(all_retrieved_token_ids).reshape(
+ (n_chunks_per_sample, self.num_neighbors, -1)
+ )
+
+ # Sample.
+ sample: Dict[str, np.ndarray] = {
+ **sample,
+ "neighbor_chunks": all_retrieved_chunk_ids,
+ "neighbor_tokens": all_retrieved_token_ids,
+ }
+
+ return sample
+
+
+def get_retro_datasets(
+ config: RetroConfig, gpt_datasets: dict, sample_length: int, eod_token_id: int
+) -> Tuple[Optional[RetroDataset], Optional[RetroDataset], Optional[RetroDataset]]:
+ """Get train, valid, test retro datasets.
+
+ Args:
+ config (RetroConfig): Retro preprocessing config.
+ gpt_datasets (dict): Mapping of data split key ('train', 'valid', or 'test') to the original sequence-length GPT dataset (i.e., not the chunk dataset).
+ sample_length (int): Alias to `sequence_length`.
+ eod_token_id (int): GPT EOD token ID.
+
+ Returns:
+ A tuple of 'train', 'valid', and 'test' `RetroDataset`s.
+ """
+
+ # DB dataset.
+ db_dataset = get_db_dataset(
+ project_dir=config.retro_project_dir,
+ chunk_length=config.retro_chunk_length,
+ eod_token_id=eod_token_id,
+ )
+
+ # GPT chunk datasets.
+ chunk_ds_info_map = build_gpt_chunk_datasets_from_gpt_datasets(
+ project_dir=config.retro_project_dir,
+ gpt_datasets=gpt_datasets,
+ sample_length=sample_length,
+ chunk_length=config.retro_chunk_length,
+ )
+
+ # Retro datasets.
+ retro_dataset_map: Dict[str, Optional[RetroDataset]] = {}
+ query_dir = get_query_dir(config.retro_project_dir)
+ for data_key, chunk_ds_info in chunk_ds_info_map.items():
+
+ # Skip unused datasets.
+ if chunk_ds_info is None:
+ retro_dataset_map[data_key] = None
+ continue
+
+ # For consistency with preprocessing, the neighbor_dir is overwritten
+ # (from its setting in `build_gpt_chunk_datasets_from_gpt_datasets()`
+ # above). This is one piece -- along with setting data_path and
+ # train_samples from config.json -- of ensuring consistency between
+ # preprocessing and pretraining.
+ chunk_dataset = chunk_ds_info["dataset"]
+ chunk_ds_info["neighbor_dir"] = os.path.join(
+ query_dir, config.retro_neighbor_dirs[data_key]
+ )
+ neighbor_dir = chunk_ds_info["neighbor_dir"]
+ neighbor_path_map = BlockPathMap.from_dir(
+ dir=neighbor_dir, block_size=config.retro_block_size
+ )
+
+ # Verify num chunks.
+ n_active_chunks = chunk_ds_info["num_active_chunks"]
+ n_neighbor_chunks = neighbor_path_map.max_idx
+
+ if not os.path.isdir(neighbor_dir):
+ if torch.distributed.get_rank() == 0:
+ raise Exception(
+ "neighbor directory '%s' not found; please "
+ "compare --train-samples, --seq-length, --seed, "
+ "--eval-iters, and --eval-interval, with "
+ "retro preprocessing args." % neighbor_dir
+ )
+ torch.distributed.barrier()
+ exit()
+
+ if config.retro_verify_neighbor_count and n_active_chunks != n_neighbor_chunks:
+ if torch.distributed.get_rank() == 0:
+ log_retro_rank_0("neighbor_dir : %s" % neighbor_dir)
+ log_retro_rank_0("neighbor_path_map : %s" % neighbor_path_map)
+ raise Exception(
+ "num sampled chunks (%d) != num neighbor chunks "
+ "(%d); did you complete querying the entire "
+ "pretraining dataset?" % (n_active_chunks, n_neighbor_chunks)
+ )
+ torch.distributed.barrier()
+ exit()
+
+ # Retro dataset.
+ retro_dataset_map[data_key] = RetroDataset(
+ num_queried_samples=gpt_datasets[data_key][1],
+ num_neighbors=config.retro_num_neighbors,
+ num_retrieved_chunks=config.retro_num_retrieved_chunks,
+ block_size=config.retro_block_size,
+ db_dataset=db_dataset,
+ chunk_dataset=chunk_dataset,
+ neighbor_path_map=neighbor_path_map,
+ )
+
+ return (retro_dataset_map["train"], retro_dataset_map["valid"], retro_dataset_map["test"])
diff --git a/megatron/core/datasets/retro/query/utils.py b/megatron/core/datasets/retro/query/utils.py
new file mode 100644
index 0000000000..b4e0c67009
--- /dev/null
+++ b/megatron/core/datasets/retro/query/utils.py
@@ -0,0 +1,35 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Utilities for querying the pretraining dataset."""
+
+import os
+
+from megatron.core.datasets.megatron_dataset import MegatronDataset
+
+
+def get_query_dir(project_dir: str) -> str:
+ """Get root directory of all saved query data.
+
+ Args:
+ project_dir (str): Retro project dir.
+
+ Returns:
+ Path to query sub-directory in Retro project.
+ """
+ return os.path.join(project_dir, "query")
+
+
+def get_neighbor_dir(project_dir: str, key: str, dataset: MegatronDataset) -> str:
+ """Get directory containing neighbor IDs for a dataset (i.e., train, valid, or test).
+
+ Args:
+ project_dir (str): Retro project dir.
+ key (str): Dataset split key; 'train', 'valid', or 'test'.
+ dataset (MegatronDataset): Dataset containing unique hash for finding corresponding neighbors.
+
+ Returns:
+ Path to directory containing this dataset's neighbors within Retro project.
+ """
+ return os.path.join(
+ get_query_dir(project_dir), os.path.basename(f"{key}_{dataset.unique_description_hash}")
+ )
diff --git a/megatron/core/datasets/retro/utils.py b/megatron/core/datasets/retro/utils.py
new file mode 100644
index 0000000000..31c0be14c8
--- /dev/null
+++ b/megatron/core/datasets/retro/utils.py
@@ -0,0 +1,349 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Utilities for Retro preprocessing."""
+
+import glob
+import logging
+import os
+from collections import defaultdict
+from types import SimpleNamespace
+from typing import Any, Callable, Dict, List, Optional
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from megatron.core import parallel_state
+from megatron.core.datasets.retro.config import RetroPreprocessingConfig
+from megatron.core.datasets.retro.query.multi_split_gpt_dataset import (
+ MultiSplitGPTDataset,
+ MultiSplitGPTDatasetConfig,
+)
+from megatron.core.utils import log_single_rank
+
+from .external_libs import h5py
+
+logger = logging.getLogger(__name__)
+
+
+def log_retro_rank_0(message: str) -> None:
+ """Log on rank 0.
+
+ Args:
+ message (str): Message to log.
+ """
+ log_single_rank(logger, logging.INFO, "[RETRO] " + message)
+
+
+def retro_makedir(config: RetroPreprocessingConfig, path: str) -> None:
+ """Make a directory, conditional on not being in validation mode.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+ path (str): Path to directory.
+ """
+ if config.retro_task_validate is None:
+ os.makedirs(path, exist_ok=True)
+
+
+def extract_data_config(config: RetroPreprocessingConfig) -> MultiSplitGPTDatasetConfig:
+ """Extract data config from dataset.
+
+ Args:
+ config (RetroPreprocessingConfig): Retro preprocessing config.
+
+ Returns:
+ The config object used to build the dataset.
+ """
+ return config.retro_gpt_chunk_datasets.train["dataset"].sample_dataset.config
+
+
+def get_num_chunks_per_sample(sample_length: int, chunk_length: int) -> int:
+ """Compute seq_length // chunk_length.
+
+ Args:
+ sample_length (int): Alias of `sequence_length`.
+ chunk_length (int): Retro chunk length (e.g., 64).
+
+ Returns:
+ Number of chunks per sample (i.e., `sequence_length` / `chunk_length`).
+ """
+ assert sample_length % chunk_length == 0
+ return sample_length // chunk_length
+
+
+class GPTToTextDataset(torch.utils.data.Dataset):
+ """Dataset to convert GPT tokens to text.
+
+ Args:
+ gpt_dataset (MultiSplitGPTDataset): GPT dataset, which outputs GPT token samples.
+ gpt_tokenizer (Any): GPT tokenizer.
+ """
+
+ def __init__(self, gpt_dataset: MultiSplitGPTDataset, gpt_tokenizer: Any):
+
+ super().__init__()
+
+ self.gpt_dataset = gpt_dataset
+ self.gpt_tokenizer = gpt_tokenizer
+
+ def __len__(self) -> int:
+ """Dataset length.
+
+ Returns:
+ Number of samples in the dataset.
+ """
+ return len(self.gpt_dataset)
+
+ def __getitem__(self, idx: int) -> dict:
+ """Get dataset sample.
+
+ Args:
+ idx (int): Index of sample.
+
+ Returns:
+ A dict containing attribute 'text' of type string.
+ """
+ gpt_token_ids = self.gpt_dataset[idx]["text"].tolist()
+ text = self.gpt_tokenizer.detokenize(gpt_token_ids)
+ return {"text": text}
+
+
+def get_blocks(
+ dirname: str, n_samples: int, block_size: int, validate: Callable = None
+) -> SimpleNamespace:
+ """Divide range [0, num_samples) to sequence of block ranges.
+
+ This is a core method within the concept of block processing. The idea
+ is to divide a range (size n_samples) into a sequence of blocks. Each
+ block corresponds to a file within 'dirname' with name
+ '{start_idx}-{end_idx}.hdf5'. This method checks for the existence of
+ these files, and returns two lists, one for existing blocks and one for
+ missing blocks.
+
+ Args:
+ dirname (str): Path to directory containing block files.
+ n_samples (int): Ideal number of samples. The total number of saved block data is <=n_samples.
+ block_size (int): Max number of samples per block file (e.g., 100000).
+ validate (Callable): Method for validating each block file during load.
+
+ Returns:
+ A namespace consisting of 2 lists: existing blocks, and missing blocks. The total number of samples between the existing and missing blocks should equal n_samples above.
+ """
+
+ assert os.path.isdir(dirname), "missing directory '%s.'" % dirname
+
+ # Block ranges.
+ block_start_idxs = list(range(0, n_samples, block_size))
+ block_end_idxs = [min(n_samples, i + block_size) for i in block_start_idxs]
+ block_ranges = list(zip(block_start_idxs, block_end_idxs))
+
+ # All block files (existing + missing).
+ n_digits = int(np.ceil(np.log(n_samples) / np.log(10)) + 1)
+ all_blocks = [
+ {
+ "range": r,
+ "path": os.path.join(
+ dirname, "%s-%s.hdf5" % tuple([str(i).zfill(n_digits) for i in r])
+ ),
+ }
+ for r in block_ranges
+ ]
+ all_block_path_set = set(block["path"] for block in all_blocks)
+
+ # Validate function.
+ validate = (lambda f: None) if validate is None else validate
+
+ # Delete corrupt files.
+ if torch.distributed.get_rank() == 0:
+ existing_block_paths = [
+ block["path"] for block in all_blocks if os.path.exists(block["path"])
+ ]
+ for index, path in enumerate(tqdm(existing_block_paths, "validating block.")):
+
+ assert path in all_block_path_set, "unexpected filename, '%s'." % path
+
+ try:
+ f = h5py.File(path, "r")
+ except Exception:
+ os.remove(path)
+ continue
+
+ try:
+ validate(f)
+ except Exception:
+ os.remove(path)
+ finally:
+ f.close()
+
+ # Wait for files to be deleted.
+ torch.distributed.barrier()
+
+ # Collect blocks.
+ blocks = SimpleNamespace(
+ existing=[b for b in all_blocks if os.path.exists(b["path"])],
+ missing=[b for b in all_blocks if not os.path.exists(b["path"])],
+ )
+
+ return blocks
+
+
+def get_blocks_by_rank(
+ dirname: str,
+ n_samples: int,
+ block_size: int,
+ validate: Callable = None,
+ sample: Optional[float] = None,
+) -> SimpleNamespace:
+ """Divide existing and missing blocks evenly across all ranks.
+
+ See 'get_blocks()' above for description. The returned lists of existing and
+ missing blocks are split evenly across ranks via interleaving. This way,
+ each rank has a roughly equal number of blocks to process for a
+ downstream operation.
+
+ Args:
+ dirname (str): Path to directory containing block files.
+ n_samples (int): Ideal number of samples. The total number of saved block data is <=n_samples.
+ block_size (int): Max number of samples per block file (e.g., 100000).
+ validate (Callable): Method for validating each block file during load.
+ sample (Optional[float]): If provided, sample a random subset of the blocks. Used for validating preprocessing correctness.
+
+ Returns:
+ A namespace consisting of 2 lists: existing blocks, and missing blocks. Each of these two lists is potentially a sub-sample of the total set of existing and missing blocks, depending on whether sampling is used. Additionally, the attributes n_existing_world and n_missing_world are the total number of existing and missing blocks, independent of samples. Therefore, (n_existing_world + n_missing_world) * block_size == n_samples.
+ """
+
+ # Get world blocks.
+ blocks = get_blocks(dirname, n_samples, block_size, validate)
+
+ # This rank's existing and missing files.
+ data_parallel_rank = parallel_state.get_data_parallel_rank()
+ data_parallel_world_size = parallel_state.get_data_parallel_world_size()
+ rank_existing_blocks = blocks.existing[
+ data_parallel_rank : len(blocks.existing) : data_parallel_world_size
+ ]
+ rank_missing_blocks = blocks.missing[
+ data_parallel_rank : len(blocks.missing) : data_parallel_world_size
+ ]
+
+ # Extend rank's existing and missing blocks (with None) such that all ranks
+ # have equal length lists. This allows for easier tracking of global progress.
+ def get_world_max(n: int) -> int:
+ """Get max value across ranks.
+
+ Args:
+ n (int): Value on this rank.
+
+ Returns:
+ Max value across all ranks.
+ """
+ n_tensor = torch.cuda.LongTensor([n])
+ torch.distributed.all_reduce(n_tensor, op=torch.distributed.ReduceOp.MAX)
+ return n_tensor.item()
+
+ max_n_existing = get_world_max(len(rank_existing_blocks))
+ max_n_missing = get_world_max(len(rank_missing_blocks))
+
+ rank_existing_blocks += [None] * (max_n_existing - len(rank_existing_blocks))
+ rank_missing_blocks += [None] * (max_n_missing - len(rank_missing_blocks))
+
+ # Collect blocks.
+ blocks = SimpleNamespace(
+ n_existing_world=len(blocks.existing),
+ n_missing_world=len(blocks.missing),
+ existing=rank_existing_blocks,
+ missing=rank_missing_blocks,
+ )
+
+ if sample is not None:
+ # Sample existing and missing blocks evenly across all ranks. The
+ # returned lists of blocks are randomly sampled (without replacement)
+ # to yield `sample * len(blocks)` number of blocks.
+
+ # Randomly sample blocks.
+ def sample_blocks(_blocks: List[Optional[Dict]]) -> List[Optional[Dict]]:
+ """Sample a random subset of all blocks.
+
+ Args:
+ _blocks (List[Optional[Dict]]): List of all blocks.
+
+ Returns:
+ A random subset of the blocks.
+ """
+ n_blocks_sample = int(np.ceil(sample * len(_blocks)))
+ sampled_blocks: List[Optional[Dict]] = [b for b in _blocks if b is not None]
+
+ np.random.seed(None)
+ np.random.shuffle(sampled_blocks)
+
+ sampled_blocks = sampled_blocks[:n_blocks_sample]
+ sampled_blocks += [None] * (n_blocks_sample - len(sampled_blocks))
+
+ return sampled_blocks
+
+ blocks.existing = sample_blocks(blocks.existing)
+ blocks.missing = sample_blocks(blocks.missing)
+
+ return blocks
+
+
+class BlockPathMap:
+ """Map an index to its containing block path.
+
+ The common use for this class is to have a directory of files containing
+ blocks of processed data, of uniform block size (e.g., 100k samples per
+ file). Each file must follow a naming convention of 'startIdx-endIdx.[ext]',
+ where 'endIdx' minus 'startIdx' must equal the block size, with the possible
+ exception of the final block. Given an input index, this class maps the
+ index to the containing block file.
+
+ Args:
+ block_paths (List[str]): List of paths to saved block files.
+ block_size (int): Max number of samples per block file (e.g., 100000).
+ """
+
+ @classmethod
+ def from_dir(cls, dir: str, block_size: int, ext: str = "hdf5") -> Any:
+ """Get list of block files, and create map.
+
+ Args:
+ dir (str): Path to directory containing saved block files.
+ block_size (int): Max number of samples per block file (e.g., 100000).
+ ext (str): Block file extension (e.g., 'hdf5').
+
+ Returns:
+ A mapping of sample index to block file path.
+ """
+ assert os.path.isdir(dir), f"directory not found, '{dir}'."
+ return cls(sorted(glob.glob(dir + f"/*.{ext}")), block_size)
+
+ def __init__(self, block_paths: List[str], block_size: int):
+ self.max_idx = 0
+ self.block_path_map = {}
+ for block_path in block_paths:
+ name = os.path.splitext(os.path.basename(block_path))[0]
+ start_idx, end_idx = [int(i) for i in name.split("-")]
+ self.block_path_map[start_idx] = block_path
+ self.max_idx = max(self.max_idx, end_idx)
+ self.block_size = block_size
+
+ def __str__(self) -> str:
+ """Stringify the mapping.
+
+ Returns:
+ A string representation of this block path map.
+ """
+ return "%d paths" % len(self.block_path_map)
+
+ def __getitem__(self, idx: int) -> str:
+ """Get block path from index.
+
+ Args:
+ idx (int): Index of sample.
+
+ Returns:
+ The path to the block file containing the sample index.
+ """
+ block_start_idx = self.block_size * (idx // self.block_size)
+ block_path = self.block_path_map[block_start_idx]
+ return block_path
diff --git a/megatron/core/datasets/t5_dataset.py b/megatron/core/datasets/t5_dataset.py
new file mode 100644
index 0000000000..b54e4f5315
--- /dev/null
+++ b/megatron/core/datasets/t5_dataset.py
@@ -0,0 +1,228 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from collections import deque
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+import numpy
+
+from megatron.core.datasets.indexed_dataset import IndexedDataset
+from megatron.core.datasets.masked_dataset import (
+ MaskedWordPieceDataset,
+ MaskedWordPieceDatasetConfig,
+)
+from megatron.core.datasets.utils import Split
+
+
+@dataclass
+class T5MaskedWordPieceDatasetConfig(MaskedWordPieceDatasetConfig):
+ """Configuration object for Megatron Core T5 WordPiece datasets
+
+ NB: As a temporary holdover from Megatron-LM. The T5 tokenizer has an attribute which defines
+ a number of special sentinel tokens used during sampling. The assert in __post_init__ serves to
+ preserve compatibility with Megatron-LM until the T5 tokenizer is in Megatron Core.
+ """
+
+ sequence_length_encoder: Optional[int] = field(init=False, default=None)
+ """A sequence_length alias and the sequence length for the encoder"""
+
+ sequence_length_decoder: int = None
+ """The sequence length for the decoder"""
+
+ def __post_init__(self) -> None:
+ """Do asserts and set fields post init"""
+ super().__post_init__()
+
+ self.sequence_length_encoder = self.sequence_length
+
+ assert self.sequence_length_encoder is not None
+ assert self.sequence_length_decoder is not None
+
+ assert len(self.tokenizer.additional_special_tokens_ids) > 0
+
+
+class T5MaskedWordPieceDataset(MaskedWordPieceDataset):
+ """The T5 dataset that assumes WordPiece tokenization
+
+ Args:
+ indexed_dataset (IndexedDataset): The IndexedDataset around which to build the MegatronDataset
+
+ dataset_path (str): The real path on disk to the dataset, for bookkeeping
+
+ indexed_indices (numpy.ndarray): The set of the documents indices to expose
+
+ num_samples (Optional[int]): The number of samples to draw from the indexed dataset. When None, build as many samples as correspond to one epoch.
+
+ index_split (Split): The indexed_indices Split
+
+ config (T5MaskedWordPieceDatasetConfig): The config
+ """
+
+ def __init__(
+ self,
+ indexed_dataset: IndexedDataset,
+ dataset_path: str,
+ indexed_indices: numpy.ndarray,
+ num_samples: Optional[int],
+ index_split: Split,
+ config: T5MaskedWordPieceDatasetConfig,
+ ) -> None:
+ super().__init__(
+ indexed_dataset, dataset_path, indexed_indices, num_samples, index_split, config
+ )
+
+ self.token_lookup = list(self.config.tokenizer.inv_vocab.keys())
+ # Account for the single and single token ids
+ self.sample_index = self._build_sample_index(self.config.sequence_length - 2, 1)
+
+ @staticmethod
+ def _key_config_attributes() -> List[str]:
+ """Inherited method implementation
+
+ Returns:
+ List[str]: The key config attributes
+ """
+ return super(
+ T5MaskedWordPieceDataset, T5MaskedWordPieceDataset
+ )._key_config_attributes() + ["sequence_length_decoder"]
+
+ def __getitem__(self, idx: int) -> Dict[str, Union[int, numpy.ndarray]]:
+ """Abstract method implementation
+
+ Args:
+ idx (int): The index into the dataset
+
+ Returns:
+ Dict[str, Union[int, numpy.ndarray]]: The
+ """
+ idx_beg, idx_end, target_sequence_length = self.sample_index[idx]
+ sample = [self.dataset[i] for i in range(idx_beg, idx_end)]
+
+ numpy_random_state = numpy.random.RandomState(seed=(self.config.random_seed + idx) % 2**32)
+
+ assert target_sequence_length <= self.config.sequence_length
+
+ # Flatten the sample into a list of tokens
+ tokens = [token for sentence in sample for token in sentence]
+
+ # Truncate the list of tokens to a desired length
+ truncated = len(tokens) > target_sequence_length
+ tokens = tokens[:target_sequence_length]
+
+ # Masking
+ (tokens, _, _, _, masked_spans) = self._create_masked_lm_predictions(
+ tokens, target_sequence_length, numpy_random_state
+ )
+
+ # Prepare the encoder input and decoder input and output
+ sentinels = deque(self.config.tokenizer.additional_special_tokens_ids)
+ encoder_input = []
+ decoder_input = [self.config.tokenizer.bos]
+ decoder_output = []
+ idx_beg = 0
+ for indices, labels in masked_spans:
+ sentinel = sentinels.popleft()
+
+ # set the end index
+ idx_end = indices[0]
+
+ encoder_input.extend(tokens[idx_beg:idx_end])
+ encoder_input.append(sentinel)
+
+ decoder_input.append(sentinel)
+ decoder_input.extend(labels)
+
+ decoder_output.append(sentinel)
+ decoder_output.extend(labels)
+
+ # set the start index
+ idx_beg = indices[-1] + 1
+
+ encoder_input.extend(tokens[idx_beg:])
+ decoder_output.append(self.config.tokenizer.eos)
+
+ # Pad the sequences and convert to NumPy
+ length_toks_encoder = len(encoder_input)
+ length_toks_decoder = len(decoder_input)
+ length_pads_encoder = self.config.sequence_length_encoder - length_toks_encoder
+ length_pads_decoder = self.config.sequence_length_decoder - length_toks_decoder
+ assert length_pads_encoder >= 0
+ assert length_pads_decoder >= 0
+
+ encoder_input = numpy.array(encoder_input, dtype=numpy.int64)
+ encoder_input = numpy.pad(
+ encoder_input, (0, length_pads_encoder), constant_values=self.config.tokenizer.pad
+ )
+
+ decoder_input = numpy.array(decoder_input, dtype=numpy.int64)
+ decoder_input = numpy.pad(
+ decoder_input, (0, length_pads_decoder), constant_values=self.config.tokenizer.pad
+ )
+
+ # Create attention and history masks
+ mask_encoder = self._make_attention_mask(encoder_input, encoder_input)
+ mask_encoder_decoder = self._make_attention_mask(decoder_input, encoder_input)
+ mask_decoder = self._make_attention_mask(decoder_input, decoder_input)
+ mask_decoder = mask_decoder * self._make_history_mask(decoder_input)
+
+ # Mask the labels
+ decoder_output = numpy.array(decoder_output, dtype=numpy.int64)
+ decoder_output = numpy.pad(decoder_output, (0, length_pads_decoder), constant_values=-1)
+
+ # Get the loss mask
+ loss_mask = numpy.zeros(self.config.sequence_length_decoder, dtype=numpy.int64)
+ loss_mask[:length_toks_decoder] = 1
+
+ return {
+ "text_enc": encoder_input,
+ "text_dec": decoder_input,
+ "labels": decoder_output,
+ "loss_mask": loss_mask,
+ "truncated": int(truncated),
+ "enc_mask": mask_encoder,
+ "dec_mask": mask_decoder,
+ "enc_dec_mask": mask_encoder_decoder,
+ }
+
+ @staticmethod
+ def _make_attention_mask(
+ source_block: numpy.ndarray, target_block: numpy.ndarray
+ ) -> numpy.ndarray:
+ """Return a 2-D attention mask
+
+ Args:
+ source_block (numpy.ndarray): A 1-D array
+ target_block (numpy.ndarray): A 1-D array
+
+ Returns:
+ numpy.ndarray: The 2-D attention mask
+ """
+ mask = (target_block[None, :] >= 1) * (source_block[:, None] >= 1)
+ return mask.astype(numpy.int64)
+
+ @staticmethod
+ def _make_history_mask(block: numpy.ndarray) -> numpy.ndarray:
+ """Return a 2-D history (lower-left-triangular) mask
+
+ Args:
+ block (numpy.ndarray): A 1-D array
+
+ Returns:
+ numpy.ndarray: The 2-D history (lower-left-triangular) mask
+ """
+ arange = numpy.arange(block.shape[0])
+ mask = arange[None,] <= arange[:, None]
+ return mask.astype(numpy.int64)
+
+ def _get_token_mask(self, numpy_random_state: numpy.random.RandomState) -> int:
+ """Abstract method implementation
+
+ 100% of the time, replace the token id with mask token id.
+
+ Args:
+ numpy_random_state (RandomState): The NumPy random state
+
+ Returns:
+ int: The mask token id
+ """
+ return self.config.tokenizer.mask
diff --git a/megatron/core/datasets/utils.py b/megatron/core/datasets/utils.py
new file mode 100644
index 0000000000..8d887d4a4a
--- /dev/null
+++ b/megatron/core/datasets/utils.py
@@ -0,0 +1,87 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+from enum import Enum
+from typing import List, Optional, Tuple
+
+import numpy
+import torch
+
+from ..utils import log_single_rank
+
+logger = logging.getLogger(__name__)
+
+
+class Split(Enum):
+ train = 0
+ valid = 1
+ test = 2
+
+
+def compile_helpers():
+ """Compile C++ helper functions at runtime. Make sure this is invoked on a single process."""
+ import os
+ import subprocess
+
+ command = ["make", "-C", os.path.abspath(os.path.dirname(__file__))]
+ if subprocess.run(command).returncode != 0:
+ import sys
+
+ log_single_rank(logger, logging.ERROR, "Failed to compile the C++ dataset helper functions")
+ sys.exit(1)
+
+
+def normalize(weights: List[float]) -> List[float]:
+ """Do non-exponentiated normalization
+
+ Args:
+ weights (List[float]): The weights
+
+ Returns:
+ List[float]: The normalized weights
+ """
+ w = numpy.array(weights, dtype=numpy.float64)
+ w_sum = numpy.sum(w)
+ w = (w / w_sum).tolist()
+ return w
+
+
+def get_blend_from_list(
+ blend: Optional[List[str]],
+) -> Optional[Tuple[List[str], Optional[List[float]]]]:
+ """Get the megatron.core.datasets.blended_megatron_dataset_config.BlendedMegatronDatasetConfig blend from the blend list
+
+ Args:
+ blend (Optional[List[str]]): The blend list, which can be either (1) a list of prefixes, e.g. ["path/to/dataset_1_prefix", "path/to/dataset_2_prefix"], or (2) a flattened, zipped list of weights and prefixes, e.g. ["30", "path/to/dataset_1_prefix", "70", "path/to/dataset_2_prefix"]
+
+ Returns:
+ Optional[Tuple[List[str], Optional[List[float]]]]: The blend, consisting of a list of dataset prefixes and optionally a list of dataset weights, e.g. [["path/to/dataset_1_prefix", "path/to/dataset_2_prefix"], [30.0, 70.0]].
+ """
+ if blend is None:
+ return None
+
+ if len(blend) % 2 == 1:
+ weight_per_dataset = None
+ raw_prefix_per_dataset = blend
+ else:
+ raw_weight_per_dataset, raw_prefix_per_dataset = zip(
+ *[(blend[i], blend[i + 1]) for i in range(0, len(blend), 2)]
+ )
+
+ weight_per_dataset = []
+ for rwpd in raw_weight_per_dataset:
+ try:
+ weight = float(rwpd)
+ except ValueError:
+ weight = None
+ weight_per_dataset.append(weight)
+
+ is_none = map(lambda _: _ is None, weight_per_dataset)
+ if any(is_none):
+ assert all(is_none)
+ weight_per_dataset = None
+ raw_prefix_per_dataset = blend
+
+ prefix_per_dataset = [rppd.strip() for rppd in raw_prefix_per_dataset]
+
+ return prefix_per_dataset, weight_per_dataset
diff --git a/megatron/core/datasets/utils_s3.py b/megatron/core/datasets/utils_s3.py
new file mode 100644
index 0000000000..61103b429d
--- /dev/null
+++ b/megatron/core/datasets/utils_s3.py
@@ -0,0 +1,164 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import os
+from typing import Any, Dict, NamedTuple, Protocol, Tuple
+
+import torch
+
+try:
+ import boto3
+ import botocore.exceptions as exceptions
+except ModuleNotFoundError:
+ pass
+
+S3_PREFIX = "s3://"
+
+
+class S3Config(NamedTuple):
+ """Config when the data (.bin) file and the index (.idx) file are in S3
+
+ TODO: These parameters are few and can be consolidated with parameters specific to bin reader
+ classes - @jkamalu
+
+ Attributes:
+
+ path_to_idx_cache (str): The local directory where we will store the index (.idx) file
+
+ bin_chunk_nbytes (int): If the number of bytes is too small, then we send a request to S3 at each call of the `read` method in _S3BinReader, which is slow, because each request has a fixed cost independent of the size of the byte range requested. If the number of bytes is too large, then we only rarely have to send requests to S3, but it takes a lot of time to complete the request when we do, which can block training. We've found that 256 * 1024 * 1024 (i.e., 256 MiB) has worked well (though we have not put that much effort into tuning it), so we default to it.
+ """
+
+ path_to_idx_cache: str
+
+ bin_chunk_nbytes: int = 256 * 1024 * 1024
+
+
+class S3Client(Protocol):
+ """The protocol which all s3 clients should abide by"""
+
+ def download_file(self, Bucket: str, Key: str, Filename: str) -> None: ...
+
+ def upload_file(self, Filename: str, Bucket: str, Key: str) -> None: ...
+
+ def head_object(self, Bucket: str, Key: str) -> Dict[str, Any]: ...
+
+ def get_object(self, Bucket: str, Key: str, Range: str) -> Dict[str, Any]: ...
+
+ def close(self) -> None: ...
+
+
+def is_s3_path(path: str) -> bool:
+ """Ascertain whether a path is in S3
+
+ Args:
+ path (str): The path
+
+ Returns:
+ bool: True if the path is in S3, False otherwise
+ """
+ return path.startswith(S3_PREFIX)
+
+
+def parse_s3_path(path: str) -> Tuple[str, str]:
+ """Parses the given S3 path returning correspsonding bucket and key.
+
+ Args:
+ path (str): The S3 path
+
+ Returns:
+ Tuple[str, str]: A (bucket, key) tuple
+ """
+ assert is_s3_path(path)
+ parts = path.replace(S3_PREFIX, "").split("/")
+ bucket = parts[0]
+ if len(parts) > 1:
+ key = "/".join(parts[1:])
+ assert S3_PREFIX + bucket + "/" + key == path
+ else:
+ key = ""
+ return bucket, key
+
+
+def object_exists(client: S3Client, path: str) -> bool:
+ """Ascertain whether the object at the given S3 path exists in S3
+
+ Args:
+ client (S3Client): The S3 client
+
+ path (str): The S3 path
+
+ Raises:
+ botocore.exceptions.ClientError: The error code is 404
+
+ Returns:
+ bool: True if the object exists in S3, False otherwise
+ """
+ parsed_s3_path = parse_s3_path(path)
+ try:
+ response = client.head_object(bucket=parsed_s3_path[0], key=parsed_s3_path[1])
+ except exceptions.ClientError as e:
+ if e.response["Error"]["Code"] != "404":
+ raise e
+ return True
+
+
+def _download_file(client: S3Client, s3_path: str, local_path: str) -> None:
+ """Download the object at the given S3 path to the given local file system path
+
+ Args:
+ client (S3Client): The S3 client
+
+ s3_path (str): The S3 source path
+
+ local_path (str): The local destination path
+ """
+ dirname = os.path.dirname(local_path)
+ os.makedirs(dirname, exist_ok=True)
+ parsed_s3_path = parse_s3_path(s3_path)
+ client.download_file(parsed_s3_path[0], parsed_s3_path[1], local_path)
+
+
+def maybe_download_file(s3_path: str, local_path: str) -> None:
+ """Download the object at the given S3 path to the given local file system path
+
+ In a distributed setting, downloading the S3 object proceeds in stages in order
+ to try to have the minimum number of processes download the object in order for
+ all the ranks to have access to the downloaded object.
+
+ Args:
+ s3_path (str): The S3 source path
+
+ local_path (str): The local destination path
+ """
+
+ if torch.distributed.is_initialized():
+ rank = torch.distributed.get_rank()
+ local_rank = rank % torch.cuda.device_count()
+ else:
+ rank = 0
+ local_rank = 0
+
+ s3_client = boto3.client("s3")
+
+ if (not os.path.exists(local_path)) and (rank == 0):
+ _download_file(s3_client, s3_path, local_path)
+
+ if torch.distributed.is_initialized():
+ torch.distributed.barrier()
+
+ # If the `local_path` is in a file system that is not
+ # shared across all the ranks, then we assume it's in the
+ # host file system and each host needs to download the file.
+ if (not os.path.exists(local_path)) and (local_rank == 0):
+ _download_file(s3_client, s3_path, local_path)
+
+ if torch.distributed.is_initialized():
+ torch.distributed.barrier()
+
+ # If the `local_path` still does not exist, then we assume
+ # each rank is saving to a separate location.
+ if not os.path.exists(local_path):
+ _download_file(s3_client, s3_path, local_path)
+
+ if torch.distributed.is_initialized():
+ torch.distributed.barrier()
+
+ assert os.path.exists(local_path)
diff --git a/megatron/core/dist_checkpointing/__init__.py b/megatron/core/dist_checkpointing/__init__.py
new file mode 100644
index 0000000000..a065b5f36a
--- /dev/null
+++ b/megatron/core/dist_checkpointing/__init__.py
@@ -0,0 +1,11 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+from .core import check_is_distributed_checkpoint
+from .mapping import LocalNonpersistentObject, LocalNonpersitentObject, ShardedTensor
+from .serialization import (
+ load,
+ load_common_state_dict,
+ load_plain_tensors,
+ load_tensors_metadata,
+ save,
+)
diff --git a/megatron/core/dist_checkpointing/core.py b/megatron/core/dist_checkpointing/core.py
new file mode 100644
index 0000000000..af6ebff6ec
--- /dev/null
+++ b/megatron/core/dist_checkpointing/core.py
@@ -0,0 +1,77 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Module for managing distributed checkpoints metadata. """
+
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Optional
+
+CONFIG_FNAME = 'metadata.json'
+
+
+class CheckpointingException(Exception):
+ """Base checkpointing related exception"""
+
+ pass
+
+
+@dataclass
+class CheckpointingConfig:
+ """Documents backends used in the checkpoint.
+
+ Checkpoint config keeps track of formats used for storing the sharded tensors
+ (sharded_backend) and other objects (common_backend).
+
+ Note that versioning is not for the checkpoint content (which is application specific),
+ but for the checkpoint format itself.
+ """
+
+ sharded_backend: str
+ sharded_backend_version: int = 1
+ common_backend: str = 'torch'
+ common_backend_version: int = 1
+
+
+def check_is_distributed_checkpoint(checkpoint_dir):
+ """Checks if `metadata.json` exists in the checkpoint and is a valid config.
+
+ Args:
+ checkpoint_dir: checkpoint directory
+
+ Returns:
+ bool: True if `metadata.json` exists in the checkpoint and is a valid config.
+ """
+ return maybe_load_config(checkpoint_dir) is not None
+
+
+def maybe_load_config(checkpoint_dir: str) -> Optional[CheckpointingConfig]:
+ """Returns checkpoint config if `checkpoint_dir` is a distributed checkpoint and None otherwise
+
+ Args:
+ checkpoint_dir: checkpoint directory
+
+ Returns:
+ CheckpointingConfig (optional): None if checkpoint is not a valid distributed checkpoint
+ """
+ config_path = Path(checkpoint_dir, CONFIG_FNAME)
+ if not config_path.exists():
+ return None
+ with config_path.open() as f:
+ config_dict = json.load(f)
+ return CheckpointingConfig(**config_dict)
+
+
+def save_config(config: CheckpointingConfig, checkpoint_dir: str):
+ """Save given config to checkpoint directory.
+
+ Args:
+ config: checkpoint config
+ checkpoint_dir: checkpoint directory
+
+ Returns:
+ None
+ """
+ config_path = Path(checkpoint_dir, CONFIG_FNAME)
+ with config_path.open('w') as f:
+ json.dump(asdict(config), f)
diff --git a/megatron/core/dist_checkpointing/dict_utils.py b/megatron/core/dist_checkpointing/dict_utils.py
new file mode 100644
index 0000000000..438925112c
--- /dev/null
+++ b/megatron/core/dist_checkpointing/dict_utils.py
@@ -0,0 +1,245 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Utilities for operating with dicts and lists.
+
+All functions in this module handle nesting of dicts and lists.
+Other objects (e.g. tuples) are treated as atomic leaf types that cannot be traversed.
+"""
+
+from collections import defaultdict
+from typing import Any, Callable, Dict, Iterable, List, Tuple, TypeVar, Union
+
+import numpy as np
+import torch
+
+U, V = TypeVar("U"), TypeVar("V")
+
+
+def extract_matching_values(
+ x: Union[dict, list], predicate: Callable[[Any], bool], return_lists_as_dicts: bool = False
+) -> Tuple[Union[dict, list], Union[dict, list]]:
+ """Return matching and nonmatching values. Keeps hierarchy.
+
+ Args:
+ x (Union[dict, list]) : state dict to process. Top-level argument must be a dict or list
+ predicate (object -> bool): determines matching values
+ return_lists_as_dicts (bool): if True, matching lists will be turned
+ into dicts, with keys indicating the indices of original elements.
+ Useful for reconstructing the original hierarchy.
+ """
+
+ def _set_elem(target, k, v):
+ if return_lists_as_dicts:
+ target[k] = v
+ else:
+ target.append(v)
+
+ if isinstance(x, dict):
+ matching_vals = {}
+ nonmatching_vals = {}
+ for k, v in x.items():
+ if isinstance(v, (list, dict)):
+ match, nonmatch = extract_matching_values(v, predicate, return_lists_as_dicts)
+ if match:
+ matching_vals[k] = match
+ if nonmatch or not v:
+ nonmatching_vals[k] = nonmatch
+ elif predicate(v):
+ matching_vals[k] = v
+ else:
+ nonmatching_vals[k] = v
+ elif isinstance(x, list): # type: ignore
+ matching_vals = {} if return_lists_as_dicts else []
+ nonmatching_vals = {} if return_lists_as_dicts else []
+ for ind, v in enumerate(x):
+ if isinstance(v, (list, dict)) and v:
+ match, nonmatch = extract_matching_values(v, predicate, return_lists_as_dicts)
+ if match:
+ _set_elem(matching_vals, ind, match)
+ if nonmatch or not v:
+ _set_elem(nonmatching_vals, ind, nonmatch)
+ else:
+ target = matching_vals if predicate(v) else nonmatching_vals
+ _set_elem(target, ind, v)
+ else:
+ raise ValueError(f'Unexpected top-level object type: {type(x)}')
+ return matching_vals, nonmatching_vals
+
+
+def diff(x1: Any, x2: Any, prefix: Tuple = ()) -> Tuple[list, list, list]:
+ """Recursive diff of dicts.
+
+ Args:
+ x1 (object): left dict
+ x2 (object): right dict
+ prefix (tuple): tracks recursive calls. Used for reporting differing keys.
+
+ Returns:
+ Tuple[list, list, list]: tuple of:
+ - only_left: Prefixes present only in left dict
+ - only_right: Prefixes present only in right dict
+ - mismatch: values present in both dicts but not equal across dicts.
+ For tensors equality of all elems is checked.
+ Each element is a tuple (prefix, type of left value, type of right value).
+ """
+ mismatch = []
+ if isinstance(x1, dict) and isinstance(x2, dict):
+ only_left = [prefix + (k,) for k in x1.keys() - x2.keys()]
+ only_right = [prefix + (k,) for k in x2.keys() - x1.keys()]
+ for k in x2.keys() & x1.keys():
+ _left, _right, _mismatch = diff(x1[k], x2[k], prefix + (k,))
+ only_left.extend(_left)
+ only_right.extend(_right)
+ mismatch.extend(_mismatch)
+ elif isinstance(x1, list) or isinstance(x1, tuple) or isinstance(x1, np.ndarray):
+ assert type(x1) == type(x2)
+ only_left = list(range(len(x1) - 1, len(x2) - 1, -1))
+ only_right = list(range(len(x1) - 1, len(x2) - 1, -1))
+ for i, (v1, v2) in enumerate(zip(x1, x2)):
+ _left, _right, _mismatch = diff(v1, v2, prefix + (i,))
+ only_left.extend(_left)
+ only_right.extend(_right)
+ mismatch.extend(_mismatch)
+ else:
+ only_left = []
+ only_right = []
+ if isinstance(x1, torch.Tensor) and isinstance(x2, torch.Tensor):
+ _is_mismatch = not torch.all(x1 == x2)
+ # TODO: change with concrete type that has both replica_id and data attrs
+ elif hasattr(x1, 'replica_id') and hasattr(x2, 'replica_id'):
+ assert type(x1) == type(x2)
+ only_left, only_right, mismatch = diff(
+ x1.data, x2.data, prefix + (type(x1),)
+ ) # type: ignore
+ _is_mismatch = False
+ else:
+ try:
+ _is_mismatch = bool(x1 != x2)
+ except RuntimeError:
+ _is_mismatch = True
+
+ if _is_mismatch:
+ mismatch.append((prefix, type(x1), type(x2)))
+
+ return only_left, only_right, mismatch
+
+
+def inspect_types(x: Any, prefix: Tuple = (), indent: int = 4):
+ """Helper to print types of (nested) dict values."""
+ print_indent = lambda: print(' ' * indent * len(prefix), end='')
+ if isinstance(x, dict):
+ print()
+ for k, v in x.items():
+ print_indent()
+ print(f'> {k}: ', end='')
+ inspect_types(v, prefix + (k,), indent)
+ elif isinstance(x, list):
+ print()
+ for i, v in enumerate(x):
+ print_indent()
+ print(f'- {i}: ', end='')
+ inspect_types(v, prefix + (i,), indent)
+ else:
+ if isinstance(x, torch.Tensor):
+ print(f'Tensor of shape {x.shape}')
+ else:
+ try:
+ x_str = str(x)
+ except:
+ x_str = ''
+ if len(x_str) > 30:
+ x_str = x_str[:30] + '... (truncated)'
+ print(f'[{type(x)}]: {x_str}')
+
+
+def nested_values(x: Union[dict, list]):
+ """Returns iterator over (nested) values of a given dict or list."""
+ x_iter = x.values() if isinstance(x, dict) else x
+ for v in x_iter:
+ if isinstance(v, (dict, list)):
+ yield from nested_values(v)
+ else:
+ yield v
+
+
+def nested_items_iter(x: Union[dict, list]):
+ """Returns iterator over (nested) tuples (container, key, value) of a given dict or list."""
+ x_iter = x.items() if isinstance(x, dict) else enumerate(x)
+ for k, v in x_iter:
+ if isinstance(v, (dict, list)):
+ yield from nested_items_iter(v)
+ else:
+ yield x, k, v
+
+
+def dict_map(f: Callable, d: dict):
+ """`map` equivalent for dicts."""
+ for sub_d, k, v in nested_items_iter(d):
+ sub_d[k] = f(v)
+
+
+def dict_map_with_key(f: Callable, d: dict):
+ """`map` equivalent for dicts with a function that accepts tuple (key, value)."""
+ for sub_d, k, v in nested_items_iter(d):
+ sub_d[k] = f(k, v)
+
+
+def dict_list_map_inplace(f: Callable[[U], V], x: Union[Dict, List, U]):
+ """Maps dicts and lists *in-place* with a given function."""
+ if isinstance(x, dict):
+ for k, v in x.items():
+ x[k] = dict_list_map_inplace(f, v)
+ elif isinstance(x, list):
+ x[:] = (dict_list_map_inplace(f, v) for v in x)
+ else:
+ return f(x)
+ return x
+
+
+def dict_list_map_outplace(f: Callable[[U], V], x: Union[Dict, List, U]) -> Union[Dict, List, V]:
+ """Maps dicts and lists *out-of-place* with a given function."""
+ if isinstance(x, dict):
+ return {k: dict_list_map_outplace(f, v) for k, v in x.items()}
+ elif isinstance(x, list):
+ return [dict_list_map_outplace(f, v) for v in x]
+ else:
+ return f(x)
+
+
+def merge(x1: Union[dict, list], x2: Union[dict, list], key: Tuple[Union[str, int], ...] = ()):
+ """Merges dicts and lists recursively."""
+ if isinstance(x1, dict) and isinstance(x2, dict):
+ for k, v2 in x2.items():
+ if k not in x1:
+ x1[k] = v2
+ else:
+ x1[k] = merge(x1[k], v2, key=key + (k,))
+ elif isinstance(x1, list) and isinstance(x2, list):
+ if len(x1) != len(x2):
+ raise ValueError(
+ f'Cannot merge two lists with different lengths ({len(x1)} and {len(x2)}, '
+ f'encountered at level {key})'
+ )
+ for i, v2 in enumerate(x2):
+ x1[i] = merge(x1[i], v2, key=key + (i,))
+ else:
+ raise ValueError(
+ f'Duplicate non-dict and non-list values encountered: `{x1}` and `{x2}` '
+ f'(at level {key})'
+ )
+ return x1
+
+
+def map_reduce(
+ xs: Iterable,
+ key_fn: Callable = lambda x: x,
+ value_fn: Callable = lambda x: x,
+ reduce_fn: Callable = lambda x: x,
+) -> dict:
+ """Simple map-reduce implementation following `more_itertools.map_reduce` interface."""
+ res = defaultdict(list)
+ for x in xs:
+ res[key_fn(x)].append(value_fn(x))
+ for k in res:
+ res[k] = reduce_fn(res[k])
+ return dict(res)
diff --git a/megatron/core/dist_checkpointing/exchange_utils.py b/megatron/core/dist_checkpointing/exchange_utils.py
new file mode 100644
index 0000000000..2106fe574c
--- /dev/null
+++ b/megatron/core/dist_checkpointing/exchange_utils.py
@@ -0,0 +1,519 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+"""Utilities for exchanging data between ranks."""
+
+import logging
+from collections import defaultdict
+from functools import reduce
+from itertools import zip_longest
+from time import time
+from typing import Dict, List, NamedTuple, Optional, Set, Tuple, TypeVar, cast
+
+import numpy as np
+import torch
+
+from .core import CheckpointingException
+from .dict_utils import nested_values
+from .mapping import ShardedStateDict, ShardedTensor, is_main_replica
+from .utils import _sharded_tensor_shard_id, _ShardId
+
+# TODO: remove TE references once the TE bug is fixed
+# Check if Transformer Engine has Float8Tensor class
+HAVE_TE_FLOAT8TENSOR = False
+try:
+ from transformer_engine.pytorch.float8_tensor import Float8Tensor
+
+ HAVE_TE_FLOAT8TENSOR = True
+except (ImportError, ModuleNotFoundError):
+ # Float8Tensor not found
+ pass
+
+
+def is_float8tensor(tensor: torch.Tensor) -> bool:
+ """Check if a tensor is a Transformer Engine Float8Tensor"""
+ return HAVE_TE_FLOAT8TENSOR and isinstance(tensor, Float8Tensor)
+
+
+logger = logging.getLogger(__name__)
+
+
+class ShardDistribution(NamedTuple):
+ """Represents a distribution of ShardedTensors.
+
+ Given distribution is valid only for a specific parallelization group,
+ which is implicit here (not referenced by this class).
+
+ Args:
+ main_rank_for_shard (Dict[_ShardId, int]): specifies which rank should hold
+ the main replica for a given shard
+ shards_in_this_group (Set[_ShardId]): which shards have a main replica
+ in this parallelization group
+ shard_to_metadata (Dict[_ShardId, ShardedTensor]): maps ShardedTensor
+ identifier to the original ShardedTensor
+ all_ranks_for_shard (Dict[_ShardId, List[int]]): specifies which ranks
+ need a given shard in a given parallelization group
+
+ """
+
+ main_rank_for_shard: Dict[_ShardId, int]
+ shards_in_this_group: Set[_ShardId]
+ shard_to_metadata: Dict[_ShardId, ShardedTensor]
+ all_ranks_for_shard: Dict[_ShardId, List[int]]
+
+
+def _shard_size(sh_ten: ShardedTensor):
+ """Returns size in bytes of a given sharded tensor."""
+ if sh_ten.flattened_range is None:
+ numel = np.product(sh_ten.local_shape)
+ else:
+ numel = sh_ten.flattened_range.stop - sh_ten.flattened_range.start
+ return numel * torch._utils._element_size(sh_ten.dtype)
+
+
+def _get_empty_tensor_for_exchange(
+ shard_id: _ShardId,
+ needed_shards: Dict[_ShardId, ShardedTensor],
+ unneeded_shards: Dict[_ShardId, ShardedTensor],
+ loaded_tensors: Dict[_ShardId, torch.Tensor],
+) -> Tuple[torch.Tensor, Optional[torch.device]]:
+ """Determines the empty tensor to use for exchange.
+
+ If shard_id is needed by this rank, it will be in the `unloaded_shards`.
+ Otherwise, the metadata for this tensor can be found in `shard_to_metadata`
+
+ Args:
+ shard_id (_ShardId): shard_id that will be exchanged
+ needed_shards (Dict[_ShardId, ShardedTensor]): mapping from shard ids
+ to metadata for shards needed by this rank
+ unneeded_shards (Dict[_ShardId, ShardedTensor]): mapping from shard ids
+ to metadata for shards that can be discarded after exchange
+ loaded_tensors (Dict[_ShardId, torch.Tensor]): mapping where useful tensors
+ are placed in
+
+ Returns:
+ Tuple[torch.Tensor, Optional[torch.device]]: empty CUDA tensor to be exchanged,
+ and the device of the original state dict tensor (if there was any)
+ """
+ local_unloaded_sh_ten = needed_shards.get(shard_id)
+ if local_unloaded_sh_ten is None:
+ orig_device = None # this tensor will be discarded anyway
+ sh_ten = unneeded_shards[shard_id]
+ if sh_ten.data is None:
+ sh_ten.init_data('cuda')
+ tensor = sh_ten.data
+ sh_ten.data = None # won't be used. free memory
+ else:
+ tensor = sh_ten.data
+ if tensor.device.type == 'cpu':
+ tensor = torch.empty_like(tensor, device='cuda')
+ else:
+ local_unloaded_sh_ten.init_data('cuda')
+ orig_device = local_unloaded_sh_ten.data.device
+ tensor = local_unloaded_sh_ten.data
+ if tensor.device.type == 'cpu':
+ tensor = torch.empty_like(tensor, device='cuda')
+ loaded_tensors[shard_id] = tensor
+ return tensor, orig_device
+
+
+T = TypeVar('T')
+
+
+def distribute_shards_to_ranks(
+ shard_to_ranks: Dict[T, List[int]], shard_to_size: Dict[T, int], num_ranks: int
+) -> Dict[T, int]:
+ """Computes uniform distribution of workload across ranks, based on sizes.
+
+ Currently, the assignment is greedy, based on:
+ 1. Firstly, the coverage of each shard
+ (how many ranks the shard is available on; lower coverage is assigned first)
+ 2. Secondly, the size of each shard (larger size is assigned first)
+ 3. Finally, shard id for differentiation.
+
+ Third step is added because we rely on the fact that
+ the assignment is deterministic on all ranks.
+
+ Args:
+ shard_to_ranks (Dict[T, List[int]]): mapping of rank access to shards
+ shard_to_size (Dict[T, int]): sizes of each shard
+ num_ranks (int): number of ranks in the parallelization group
+
+ Returns (Dict[T, int]): assignment of shard to rank (which rank should do the work
+ to achieve maximal uniformity)
+ """
+ shard_to_ranks = {k: tuple(v) for k, v in shard_to_ranks.items()}
+ shard_to_saving_rank = {}
+ rank_sizes = [(0, rank) for rank in range(num_ranks)]
+
+ # start from tensors of lowest coverage, then go by tensor size from largest (hence minus size)
+ for shard_id, shard_ranks in sorted(
+ shard_to_ranks.items(),
+ key=lambda sh_id_ranks: (
+ len(sh_id_ranks[1]),
+ -shard_to_size[sh_id_ranks[0]],
+ sh_id_ranks[0],
+ ),
+ ):
+ # assign greedily to the least occupied rank
+ size, rank = min((size, rank) for size, rank in rank_sizes if rank in shard_ranks)
+
+ shard_to_saving_rank[shard_id] = rank
+ rank_sizes[rank] = (size + shard_to_size[shard_id], rank)
+
+ logger.debug(f'distribute_shards_to_ranks distribution: {rank_sizes}')
+
+ return shard_to_saving_rank
+
+
+def determine_main_replica_uniform_distribution(
+ sharded_state_dict: ShardedStateDict,
+ parallelization_group: torch.distributed.ProcessGroup,
+ ignore_groups: bool = False,
+) -> Optional[ShardDistribution]:
+ """Computes the save distribution.
+
+ Should be used in conjunction with `distribute_main_replicas_with_precomputed_distribution`
+ which applies the computed save distribution.
+
+ We rely on the fact that the assignment algorithm is deterministic on all ranks,
+ so there is no extra communication needed after metadata exchange.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): state dict to compute the distribution of
+ parallelization_group (ProcessGroup): distribution will be computed
+ within this process group
+ ignore_groups (bool, optional): whether the distribution defines groups.
+ This option is primarily used during loading, as it ensures that all replicas,
+ including non-main ones, are loaded by this parallelization group
+ Defaults to False.
+
+ Returns (ShardDistribution, optional): distribution that can be used to apply the
+ parallelization. Returns None if the process_group is trivial (1 rank)
+
+ """
+ group_size = torch.distributed.get_world_size(group=parallelization_group)
+ if group_size <= 1:
+ return
+ local_shards = list(
+ sh_base
+ for sh_base in nested_values(sharded_state_dict)
+ if isinstance(sh_base, ShardedTensor)
+ )
+ local_shards_no_data = [ten.without_data() for ten in local_shards]
+
+ all_shards = [None] * torch.distributed.get_world_size(group=parallelization_group)
+ torch.distributed.all_gather_object(
+ all_shards, local_shards_no_data, group=parallelization_group
+ )
+
+ shard_to_ranks = defaultdict(list)
+ shard_to_size = {}
+ shard_to_metadata = {}
+ shards_in_this_parallelization_group: Set[_ShardId] = set()
+ for rank, rank_shards in enumerate(all_shards):
+ for sh_ten in rank_shards:
+ shard_id = _sharded_tensor_shard_id(sh_ten)
+ shard_to_ranks[shard_id].append(rank)
+ if shard_id not in shard_to_size:
+ shard_to_size[shard_id] = _shard_size(sh_ten)
+ shard_to_metadata[shard_id] = sh_ten
+ if is_main_replica(sh_ten.replica_id) or ignore_groups:
+ shards_in_this_parallelization_group.add(shard_id)
+
+ shard_to_ranks = {
+ k: v for k, v in shard_to_ranks.items() if k in shards_in_this_parallelization_group
+ }
+
+ shard_to_saving_rank = distribute_shards_to_ranks(
+ shard_to_ranks, shard_to_size, len(all_shards)
+ )
+
+ return ShardDistribution(
+ shard_to_saving_rank,
+ shards_in_this_parallelization_group,
+ shard_to_metadata,
+ shard_to_ranks,
+ )
+
+
+@torch.no_grad()
+def exchange_loaded_tensors_gather_rounds(
+ loaded_tensors: Dict[_ShardId, torch.Tensor],
+ unloaded_shards: Dict[_ShardId, ShardedTensor],
+ shard_distribution: ShardDistribution = None,
+ parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> Dict[_ShardId, torch.Tensor]:
+ """Exchange the tensors loaded by different ranks with several all_gather calls.
+
+ Groups tensors by dtype, divide tensors that will be exchanged into rounds
+ and execute all_gather for tensors from each round.
+
+ Note: the loading is distributed across ranks based on total loaded size
+ in bytes, so there is no guarantee that number of rounds needed for each
+ rank will be similar, which might result in a lot of almost empty
+ all_gathers. The solution would be to group all tensors into a one
+ bytes tensor and do a single all_gather (with similarly sized messages).
+
+ Args:
+ loaded_tensors (Dict[_ShardId, torch.Tensor]): mapping from ShardedTensor
+ shard ids to tensors already loaded by this rank.
+ unloaded_shards (Dict[_ShardId, torch.Tensor]): mapping from ShardedTensor
+ shard ids to ShardedTensors that aren't loaded yet.
+ shard_distribution (ShardDistribution): distribution of all shards
+ parallelization_group (ProcessGroup, optional): process group used for load
+ distribution. Tensors will be exchanged within this group
+
+ Returns:
+ Dict[_ShardId, torch.Tensor]: dictionary mapping shard ids to tensors
+ needed by this rank to load a given state dict. Includes
+ previously loaded tensors (from `loaded_tensors` input)
+ """
+ main_rank_for_shard, _, shard_to_metadata, all_ranks_for_shard = shard_distribution
+ local_rank = torch.distributed.get_rank(group=parallelization_group)
+
+ all_loaded_tensors = dict(loaded_tensors)
+
+ # Group by dtype so that we all_gather tensors of the same dtype
+ for dtype in sorted(set(map(lambda sh_ten: sh_ten.dtype, shard_to_metadata.values())), key=str):
+
+ start = time()
+ # shards_by_rank maps rank to tensors loaded by this rank
+ shards_by_rank: List[List[torch.Tensor]] = [
+ [] for _ in range(torch.distributed.get_world_size(group=parallelization_group))
+ ]
+ for shard_id, rank in main_rank_for_shard.items():
+ if len(all_ranks_for_shard[shard_id]) == 1:
+ assert all_ranks_for_shard[shard_id][0] == main_rank_for_shard[shard_id], (
+ f'When there is only 1 ranks that needs a given shard,'
+ f' it should be the loading rank.'
+ f' Got: needs [{all_ranks_for_shard[shard_id][0]}]'
+ f' vs loads [{main_rank_for_shard[shard_id]}]'
+ )
+ # Skipping the exchange since only the loading rank needs this tensor
+ # TODO: we can employ some optimizations even for `len(shard_to_ranks) > 1`
+ # case, e.g. P2P exchange. Currently handling this case saves most of the
+ # work though.
+ continue
+ if shard_to_metadata[shard_id].dtype == dtype:
+ shards_by_rank[rank].append(shard_id)
+
+ # Transpose `shards_by_rank` to form exchange rounds
+ shards_by_round = zip_longest(*shards_by_rank, fillvalue=None)
+ for round_idx, round_shard_ids in enumerate(shards_by_round):
+ round_tensors = []
+ orig_devices = {}
+ for rank, shard_id in enumerate(round_shard_ids):
+ if shard_id is None:
+ # if no more useful data, the given rank will exchange empty tensor
+ local_ten = torch.empty(0, dtype=dtype, device='cuda')
+ orig_device = None
+ else:
+ assert isinstance(shard_id, tuple), type(shard_id)
+ if rank == local_rank:
+ assert shard_id in all_loaded_tensors, (shard_id, all_loaded_tensors.keys())
+ orig_device = all_loaded_tensors[shard_id]
+ all_loaded_tensors[shard_id] = all_loaded_tensors[shard_id].cuda()
+ local_ten = all_loaded_tensors[shard_id]
+ else:
+ local_ten, orig_device = _get_empty_tensor_for_exchange(
+ shard_id, unloaded_shards, shard_to_metadata, all_loaded_tensors
+ )
+ # Because of a TE bug, we have to exchange a nominal dtype instead of FP8
+ # It's ok to keep the nominal dtype after exchange, because TE will handle
+ # this during state dict load.
+ # TODO: remove it once the bug is fixed
+ if is_float8tensor(local_ten):
+ local_ten = local_ten.from_float8()
+ all_loaded_tensors[shard_id] = local_ten
+
+ round_tensors.append(local_ten)
+ if orig_device is not None:
+ orig_devices[shard_id] = orig_device
+
+ torch.distributed.all_gather(
+ list(round_tensors),
+ round_tensors[local_rank],
+ group=parallelization_group,
+ async_op=False,
+ )
+
+ # Move tensors back to CPU if originally was on CPU
+ for shard_id, orig_device in orig_devices.items():
+ all_loaded_tensors[shard_id] = all_loaded_tensors[shard_id].to(orig_device)
+
+ del round_tensors # remove tensor references
+
+ end = time()
+ if torch.distributed.get_rank() == 0:
+ logger.debug(f'{dtype} exchange rounds all_gather schedule took {end - start}s')
+
+ return all_loaded_tensors
+
+
+def exchange_loaded_tensors_gather_object(
+ loaded_tensors: Dict[_ShardId, torch.Tensor],
+ unloaded_shards: Dict[_ShardId, ShardedTensor],
+ shard_distribution: ShardDistribution,
+ parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> Dict[_ShardId, torch.Tensor]:
+ """Exchange the tensors loaded by different ranks with a simple all_gather_object call.
+
+ This version can be used for debugging purposes do to its simplistic
+ implementation. Shouldn't be used if performance is important.
+
+ Args:
+ loaded_tensors (Dict[_ShardId, torch.Tensor]): mapping from ShardedTensor
+ shard ids to tensors already loaded by this rank.
+ unloaded_shards (Dict[_ShardId, torch.Tensor]): mapping from ShardedTensor
+ shard ids to ShardedTensors that aren't loaded yet.
+ shard_distribution (ShardDistribution): distribution of all shards
+ parallelization_group (ProcessGroup, optional): process group used for load
+ distribution. Tensors will be exchanged within this group
+
+ Returns:
+ Dict[_ShardId, torch.Tensor]: dictionary mapping shard ids to tensors
+ needed by this rank to load a given state dict. Includes
+ previously loaded tensors (from `loaded_tensors` input)
+
+ """
+ all_loaded_tensors_list = [None] * torch.distributed.get_world_size(group=parallelization_group)
+ torch.distributed.all_gather_object(
+ all_loaded_tensors_list, loaded_tensors, group=parallelization_group
+ )
+ all_loaded_tensors_list = cast(List[Dict[_ShardId, torch.Tensor]], all_loaded_tensors_list)
+ all_loaded_tensors = reduce(lambda x, y: {**x, **y}, all_loaded_tensors_list)
+
+ # Error checks
+ if len(all_loaded_tensors) != sum(map(len, all_loaded_tensors_list)):
+ err_msg = 'Duplicate shard ids loaded by different ranks'
+ if torch.distributed.get_rank() == 0:
+ logger.error(
+ f'{err_msg}. Shards ids by rank:'
+ f' {[lt.keys() for lt in all_loaded_tensors_list]}'
+ )
+ raise CheckpointingException(err_msg)
+
+ return all_loaded_tensors
+
+
+@torch.no_grad()
+def exchange_loaded_tensors_broadcast(
+ loaded_tensors: Dict[_ShardId, torch.Tensor],
+ unloaded_shards: Dict[_ShardId, ShardedTensor],
+ shard_distribution: ShardDistribution,
+ parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> Dict[_ShardId, torch.Tensor]:
+ """Exchange the tensors loaded by different ranks by a series of broadcasts.
+
+ For each rank for each loaded tensor do a broadcast to the whole group.
+ A reasonable tradeoff in terms of performance and simplicity.
+
+ Args:
+ loaded_tensors (Dict[_ShardId, torch.Tensor]): mapping from ShardedTensor
+ shard ids to tensors already loaded by this rank.
+ unloaded_shards (Dict[_ShardId, ShardedTensor]): mapping from ShardedTensor
+ shard ids to ShardedTensors that aren't loaded yet.
+ shard_distribution (ShardDistribution): distribution of all shards
+ parallelization_group (ProcessGroup, optional): process group used for load
+ distribution. Tensors will be exchanged within this group
+
+ Returns:
+ Dict[_ShardId, torch.Tensor]: dictionary mapping shard ids to tensors
+ needed by this rank to load a given state dict. Includes
+ previously loaded tensors (from `loaded_tensors` input)
+ """
+ main_rank_for_shard, _, shard_to_metadata, all_ranks_for_shard = shard_distribution
+ local_rank = torch.distributed.get_rank(group=parallelization_group)
+
+ all_loaded_tensors = dict(loaded_tensors)
+
+ start = time()
+
+ for idx, (shard_id, rank) in enumerate(main_rank_for_shard.items()):
+ if len(all_ranks_for_shard[shard_id]) == 1:
+ assert all_ranks_for_shard[shard_id][0] == main_rank_for_shard[shard_id], (
+ f'When there is only 1 ranks that needs a given shard,'
+ f' it should be the loading rank.'
+ f'Got: needs [{all_ranks_for_shard[shard_id][0]}]'
+ f' vs loads [{main_rank_for_shard[shard_id]}]'
+ )
+ # Skipping the exchange since only the loading rank needs this tensor
+ # TODO: we can employ some optimizations even for `len(shard_to_ranks) > 1` case,
+ # e.g. P2P exchange. Currently handling this case saves most of the work though.
+ continue
+ if rank == local_rank:
+ assert shard_id in all_loaded_tensors, (shard_id, all_loaded_tensors.keys())
+ orig_device = all_loaded_tensors[shard_id].device
+ local_ten = all_loaded_tensors[shard_id].cuda()
+ else:
+ local_ten, orig_device = _get_empty_tensor_for_exchange(
+ shard_id, unloaded_shards, shard_to_metadata, all_loaded_tensors
+ )
+
+ # Because of a TE bug, we have to exchange a nominal dtype instead of FP8
+ # It's ok to keep the nominal dtype after exchange, because TE will handle
+ # this during state dict load.
+ # TODO: remove it once the bug is fixed
+ if is_float8tensor(local_ten):
+ local_ten = local_ten.from_float8()
+ all_loaded_tensors[shard_id] = local_ten
+
+ global_src_rank = (
+ rank
+ if parallelization_group == None
+ else torch.distributed.get_global_rank(parallelization_group, rank)
+ )
+ # We can do async_op=True only if there is no CPU-copy follow-up
+ torch.distributed.broadcast(
+ local_ten,
+ src=global_src_rank,
+ group=parallelization_group,
+ async_op=orig_device is None,
+ )
+ # Move tensor back to CPU if originally was on CPU
+ if orig_device is not None:
+ all_loaded_tensors[shard_id] = local_ten.to(orig_device)
+ del local_ten
+
+ end = time()
+ if torch.distributed.get_rank() == 0:
+ logger.debug(f'exchange broadcast schedule took {end - start}s')
+
+ return all_loaded_tensors
+
+
+def exchange_by_distribution(
+ loaded_tensors: Dict[_ShardId, torch.Tensor],
+ unloaded_shards: Dict[_ShardId, ShardedTensor],
+ shard_distribution: ShardDistribution = None,
+ parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
+ exchange_algo='broadcast',
+) -> Dict[_ShardId, torch.Tensor]:
+ """Exchange tensors loaded by different ranks using the specified exchange_algo.
+
+ Args:
+ loaded_tensors (Dict[_ShardId, torch.Tensor]): mapping from ShardedTensor
+ shard ids to tensors already loaded by this rank.
+ unloaded_shards (Dict[_ShardId, ShardedTensor]): mapping from ShardedTensor
+ shard ids to ShardedTensors that aren't loaded yet.
+ shard_distribution (ShardDistribution): distribution of all shards
+ parallelization_group (ProcessGroup, optional): process group used for load
+ distribution. Tensors will be exchanged within this group
+ exchange_algo (str): The algorithm used for performing exchanges.
+ Defaults to 'broadcast'.
+
+ Returns:
+ Dict[_ShardId, torch.Tensor]: dictionary mapping shard ids to tensors
+ needed by this rank to load a given state dict. Includes
+ previously loaded tensors (from `loaded_tensors` input)
+ """
+
+ if exchange_algo == 'gather_object':
+ exchange_fn = exchange_loaded_tensors_gather_object
+ elif exchange_algo == 'gather_rounds':
+ exchange_fn = exchange_loaded_tensors_gather_rounds
+ elif exchange_algo == 'broadcast':
+ exchange_fn = exchange_loaded_tensors_broadcast
+ else:
+ raise NotImplementedError(f'Unrecognized gather algorithm: {exchange_algo}')
+ return exchange_fn(loaded_tensors, unloaded_shards, shard_distribution, parallelization_group)
diff --git a/megatron/core/dist_checkpointing/mapping.py b/megatron/core/dist_checkpointing/mapping.py
new file mode 100644
index 0000000000..90d4fcdc22
--- /dev/null
+++ b/megatron/core/dist_checkpointing/mapping.py
@@ -0,0 +1,729 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Core library classes for representing sharding of tensors and objects.
+
+The main expected usage is wrapping torch.Tensors in state dicts with
+ShardedTensor class (mostly with the ShardedTensor.from_rank_offsets classmethod).
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field, replace
+from itertools import chain
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from .core import CheckpointingException
+from .dict_utils import dict_list_map_inplace
+
+logger = logging.getLogger(__name__)
+
+# These type definitions are just hints to differentiate a plain model state
+# dict (StateDict) from a state dict with tensors replaced with ShardedTensors
+# (ShardedStateDict).
+StateDict = Dict[str, Any]
+ShardedStateDict = Dict[str, Any]
+ReplicaId = Union[int, Tuple[int, ...]]
+
+
+class ShardedBase(ABC):
+ """Base class for ShardedTensor and ShardedStateDict."""
+
+ key: str
+ data: object
+ replica_id: ReplicaId
+
+ @abstractmethod
+ def validate_metadata_integrity(self):
+ """Codifies the constraints on metadata attributes."""
+
+ @abstractmethod
+ def without_data(self) -> 'ShardedBase':
+ """Returns a new ShardedBase instance with data=None."""
+ raise NotImplementedError
+
+
+@dataclass
+class ShardedTensor(ShardedBase):
+ """Represents a mapping between a local tensor and a global tensor.
+
+ Global tensor is assumed to consist of many local tensors distributed
+ between different processes.
+
+ Args:
+ key: unique identifier of a global tensor
+ data: local tensor data. Can be None only for consistency validation
+ dtype: tensor dtype
+ local_shape: local tensor shape
+ global_shape: global tensor shape
+ global_offset: offset of a local tensor in a global tensor,
+ specified in number of tensor elements
+ axis_fragmentations: global tensor fragmentation of each axis
+ replica_id: indicates given local tensor's replication wrt.
+ local tensors in different processes
+ prepend_axis_num: number of axes prepended to the local tensor to
+ reflect global tensor shape. The behavior is similar to
+ unsqueezing the local tensor.
+ allow_shape_mismatch: if True, during loading, the global shape of
+ a stored tensor does not have to match the expected global shape.
+ Useful for representing tensors with flexible shape,
+ e.g. padded.
+ flattened_range: specifies a slice that should be applied to a
+ flattened tensor with `local_shape` in order to get
+ the tensor stored as `data`
+ """
+
+ key: str
+ data: Optional[torch.Tensor] = field(repr=False)
+ dtype: torch.dtype
+ local_shape: Tuple[int, ...]
+ global_shape: Tuple[int, ...]
+ global_offset: Tuple[int, ...]
+ axis_fragmentations: Optional[Tuple[int, ...]]
+ replica_id: ReplicaId = 0
+ prepend_axis_num: int = 0
+ allow_shape_mismatch: bool = False
+ flattened_range: Optional[slice] = None
+
+ def __post_init__(self):
+ self.validate_metadata_integrity()
+
+ def validate_metadata_integrity(self) -> None:
+ """Codifies the constraints on metadata attributes.
+
+ Meeting those constraints is guaranteed when instantiating a ShardedTensor
+ class with `from_rank_offsets` or `from_rank_offsets_flat` constructors.
+
+ Returns:
+ None
+ """
+ has_flattened_range = self.flattened_range is not None
+ if self.data is not None:
+ if self.data.dtype != self.dtype:
+ raise CheckpointingException(
+ f'Data dtype should match `dtype` attribute for {self}'
+ )
+ if not has_flattened_range and self.data.shape != self.local_shape:
+ raise CheckpointingException(
+ f'Data shape should match `local_shape` attribute for {self}'
+ )
+ if has_flattened_range:
+ if self.data.ndim != 1:
+ raise CheckpointingException(f'Data should be 1D for a flattened {self}')
+ real_data = self.data
+ try:
+ self.data = None
+ self.init_data(device='meta')
+ if self.data.shape != real_data.shape:
+ raise CheckpointingException(
+ f'Data shape doesnt match expected {self.data.shape} for {self}'
+ )
+ finally:
+ self.data = real_data
+
+ if len(self.global_shape) != len(self.global_offset):
+ raise CheckpointingException(
+ f'Global offset dimensions should be equal to global shape dimensions for {self}'
+ )
+ if len(self.local_shape) + self.prepend_axis_num != len(self.global_shape):
+ raise CheckpointingException(
+ f'Local shape together with `prepend_axis_num` dimensions should be '
+ f'equal to global shape dimensions for {self}'
+ )
+
+ for off, sh in zip(self.global_offset[self.prepend_axis_num :], self.local_shape):
+ if off % sh != 0:
+ raise CheckpointingException(
+ f'Global offset ({off}) must be divisible by local shape ({sh}) for {self}.'
+ )
+
+ if has_flattened_range and self.flattened_range.step is not None:
+ raise CheckpointingException(
+ f'`step` argument in the flattened range of a ShardedTensor is not supported.'
+ )
+
+ def global_slice(self) -> Tuple[Union[int, slice], ...]:
+ """
+ Returns a tuple of int and slice objects representing a slice of the
+ global tensor that this ShardedTensor corresponds to.
+ """
+ assert len(self.global_offset) == len(self.local_shape) + self.prepend_axis_num
+ return tuple(
+ chain(
+ (off for off in self.global_offset[: self.prepend_axis_num]),
+ (
+ slice(off, off + sh)
+ for off, sh in zip(
+ self.global_offset[self.prepend_axis_num :], self.local_shape
+ )
+ ),
+ )
+ )
+
+ def global_coordinates(self) -> Tuple[np.ndarray, ...]:
+ """
+ Returns a tuple of np.ndarrays representing the coordinates of the global tensor
+ that this ShardedTensor corresponds to.
+ """
+ if self.flattened_range is None:
+ raise CheckpointingException(
+ f'`global_coordinates` is undefined for'
+ f' {self.__class__.__name__} without `flattened_range`'
+ )
+
+ local_coords = self.local_coordinates()
+ assert len(local_coords) + self.prepend_axis_num == len(self.global_offset), (
+ len(local_coords),
+ self,
+ )
+ global_coords = tuple(
+ c + off
+ for c, off in zip((0,) * self.prepend_axis_num + local_coords, self.global_offset)
+ )
+ return global_coords
+
+ def local_coordinates(self) -> Tuple[np.ndarray, ...]:
+ """
+ Returns a tuple of np.ndarrays representing the coordinates of the local tensor
+ that this ShardedTensor corresponds to.
+ """
+ if self.flattened_range is None:
+ raise CheckpointingException(
+ f'`local_coordinates` is undefined for'
+ f' {self.__class__.__name__} without `flattened_range`'
+ )
+
+ # TODO: np.unravel_index?
+ mask = np.zeros(np.product(self.local_shape), dtype=bool)
+ mask[self.flattened_range] = True
+ return np.nonzero(mask.reshape(self.local_shape))
+
+ def local_chunk_offset_in_global(self) -> Tuple[int, ...]:
+ """Offset of a local chunk in a global array of chunks.
+
+ Returns:
+ Tuple[int, ...]: the offset of the whole local chunk in a global array of chunks.
+ """
+ assert len(self.global_offset) == len(self.local_shape) + self.prepend_axis_num
+ chunk_offset = list(self.global_offset[: self.prepend_axis_num])
+ for off, sh in zip(self.global_offset[self.prepend_axis_num :], self.local_shape):
+ assert off % sh == 0, str(self)
+ chunk_offset.append(off // sh)
+ return tuple(chunk_offset)
+
+ def max_allowed_chunks(self) -> Tuple[int, ...]:
+ """
+ Returns the maximum allowed chunks for this ShardedTensor.
+ """
+ chunks = []
+ for axis_sh, axis_fragm in zip(self.global_shape, self.axis_fragmentations):
+ if not self.allow_shape_mismatch and axis_sh % axis_fragm != 0:
+ raise CheckpointingException(
+ f'Axis shape ({axis_sh}) not divisible by axis fragmentation ({axis_fragm}'
+ )
+ axis_chunk_size = axis_sh // axis_fragm
+ chunks.append(axis_chunk_size)
+ return tuple(chunks)
+
+ def without_data(self):
+ return replace(self, data=None)
+
+ @classmethod
+ def from_rank_offsets(
+ cls,
+ key: str,
+ data: torch.Tensor,
+ *rank_offsets: Tuple[int, int, int],
+ replica_id: ReplicaId = 0,
+ prepend_axis_num: int = 0,
+ flattened_range: None = None,
+ **init_kwargs,
+ ):
+ """Allows to construct the ShardedTensor given offset specified in process ranks.
+
+ Args:
+ key (str): unique key
+ data (torch.Tensor): local tensor data
+ rank_offsets (Tuple[int, int, int]): each tuple
+ (axis, axis_rank_offset, axis_fragm) says that if
+ global tensor is divided into `axis_fragm` fragment along `axis`
+ axis, then local tensor data corresponds to the `axis_rank_offset` chunk.
+ replica_id (ReplicaId): see ShardedTensor
+ prepend_axis_num (int): see ShardedTensor
+ flattened_range (None): must be None when using this constructor
+ init_kwargs: passed to ShardedTensor.__init__
+ """
+ if flattened_range is not None:
+ raise ValueError(
+ 'Cannot instantiate a flat ShardedTensor with `from_rank_offsets` method.'
+ ' Use `from_rank_offsets_flat` instead'
+ )
+ global_offset = [0] * (data.ndim + prepend_axis_num)
+ global_shape = ([1] * prepend_axis_num) + list(data.shape)
+ axis_fragmentations = [1] * (data.ndim + prepend_axis_num)
+ _seen_axis = set()
+ for axis, axis_rank_offset, axis_fragm in rank_offsets:
+ assert axis >= 0 and axis_rank_offset >= 0 and axis_fragm >= 0, (
+ axis,
+ axis_rank_offset,
+ axis_fragm,
+ )
+ assert (
+ axis_rank_offset < axis_fragm
+ ), 'Rank offset must be lower than axis fragmentation'
+ if axis in _seen_axis:
+ raise CheckpointingException('Duplicated axis specified')
+ _seen_axis.add(axis)
+
+ local_axis_shape = 1 if axis < prepend_axis_num else data.shape[axis - prepend_axis_num]
+ global_shape[axis] = axis_fragm * local_axis_shape
+ global_offset[axis] = axis_rank_offset * local_axis_shape
+ axis_fragmentations[axis] = axis_fragm
+
+ return cls(
+ key,
+ data,
+ data.dtype,
+ tuple(data.shape),
+ tuple(global_shape),
+ tuple(global_offset),
+ tuple(axis_fragmentations),
+ replica_id,
+ prepend_axis_num,
+ flattened_range=flattened_range,
+ **init_kwargs,
+ )
+
+ @classmethod
+ def from_rank_offsets_flat(
+ cls,
+ key: str,
+ data: torch.Tensor,
+ non_flat_local_shape: Tuple[int, ...],
+ *args,
+ flattened_range: Optional[slice] = None,
+ **kwargs,
+ ):
+ """Allows to construct a *flattened* ShardedTensor given offset specified in process ranks.
+
+ Args:
+ key (str):
+ data (torch.Tensor): this should be a flattened data tensor
+ non_flat_local_shape (Tuple[int, ...]): expected local shape of a non-flat chunk
+ *args: passed unchanged to the `from_rank_offsets` constructor
+ flattened_range (slice): see ShardedTensor. Defaults to None, but must be set to
+ a non-None slice.
+ **kwargs:
+
+ Returns:
+ ShardedTensor: constructed ShardedTensor instance
+ """
+ if flattened_range is None:
+ raise CheckpointingException(
+ 'Cannot instantiate a non-flat ShardedTensor with `from_rank_offsets_flat` method.'
+ ' Use `from_rank_offsets` instead'
+ )
+ if data.ndim != 1:
+ raise CheckpointingException(
+ f'Flattened ShardedTensor requires 1D data, got shape: {data.shape}'
+ )
+ if flattened_range.stop - flattened_range.start != data.numel():
+ raise CheckpointingException(
+ f'Flattened ShardedTensor data length ({data.numel()}) must meet the '
+ f'slice length: {flattened_range.stop - flattened_range.start}'
+ )
+
+ non_flat_data_meta = torch.empty(*non_flat_local_shape, dtype=data.dtype, device='meta')
+ sh_ten = cls.from_rank_offsets(key, non_flat_data_meta, *args, **kwargs)
+ instance = replace(sh_ten, data=data, flattened_range=flattened_range)
+ instance.validate_metadata_integrity()
+ return instance
+
+ def init_data(self, device: Union[str, torch.device], init_fn=torch.empty):
+ """
+ Initialize the tensor data of this ShardedTensor.
+
+ Only called if `data` attribute is None.
+
+ Args:
+ device (Union[str, torch.device]): device to place the tensor on
+ init_fn (Callable, optional): function to use to initialize the tensor.
+ Defaults to `torch.empty`.
+ """
+ if self.data is not None:
+ return
+ self.data = init_fn(self.local_shape, dtype=self.dtype, device=device)
+ if self.flattened_range is not None:
+ self.data = self.data.flatten()[self.flattened_range.start : self.flattened_range.stop]
+
+ def narrow(self, dim: int, start: int, length: int) -> List['ShardedTensor']:
+ """This is an analogue of torch.narrow for ShardedTensors.
+
+ Narrowing assumes that we narrow a local tensor on each rank.
+ This has consequences on local_shape, global_shape, global_offset, etc.
+
+ Args:
+ dim (int): dimension to narrow. Doesn't include prepended axes.
+ start (int): start element
+ length (int): length of the slice
+
+ Returns:
+ List[ShardedTensor]: narrowed ShardedTensors. For non-flat tensors,
+ the list will always have 1 element. For flat ShardedTensors the number of
+ elements varies depending on `dim` and on overlap, because flat
+ tensors must be contiguous. In particular the list can be empty.
+ """
+ prepended_dim = dim + self.prepend_axis_num
+ local_length_along_dim = self.local_shape[dim]
+
+ def _update_tuple(x, ind, val):
+ x = list(x)
+ x[ind] = val
+ return tuple(x)
+
+ def _safe_div(x, y):
+ assert x % y == 0, (x, y)
+ return x // y
+
+ # Decrease global shape and global offset by `length / local_length_along_dim`
+ assert (
+ self.global_shape[prepended_dim] % local_length_along_dim == 0
+ ), f'Only regular grid of local tensors is supported for narrowing, got: {self}'
+ assert (
+ self.global_offset[prepended_dim] % local_length_along_dim == 0
+ ), f'Only regular grid of local tensors is supported for narrowing, got: {self}'
+ global_shape = _update_tuple(
+ self.global_shape,
+ prepended_dim,
+ _safe_div(self.global_shape[prepended_dim] * length, local_length_along_dim),
+ )
+ global_offset = _update_tuple(
+ self.global_offset,
+ prepended_dim,
+ _safe_div(self.global_offset[prepended_dim] * length, local_length_along_dim),
+ )
+
+ if self.flattened_range is None:
+ new_data = self.data.narrow(dim, start, length)
+ # always a single result tensor
+ return [
+ replace(
+ self,
+ data=new_data,
+ local_shape=new_data.shape,
+ global_shape=global_shape,
+ global_offset=global_offset,
+ )
+ ]
+ else:
+ if dim != 0:
+ raise CheckpointingException(
+ f'Narrowing along the first axis is supported for now only, got dim={dim}'
+ )
+
+ # If dim=0, we will always get 0 or 1 resulting tensor.
+ # If dim>1, in general there can be more result tensors (e.g. max 3 for dim=1)
+
+ # For on original flat ShardedTensor of local shape [3, 4] and
+ # flattened_range=slice(5, 10),
+ # the X signs mark the actual (flat) data in `self.data`
+ # notice 12 (3*4) total "virtual" elements, out of which 5 is actual data.
+ # flat original: [.....XXXXX..]
+
+ # If we narrow to start=1, length=1 in the original local shape dimensions,
+ # the overlapping flat slice would be:
+ # narrow to: [....XXXX....]
+ # flat overlap: [.....XXX....]
+
+ # Now `data` is flattened and sliced, so we must compute local_shape manually
+ local_shape = _update_tuple(self.local_shape, dim, length)
+ other_dims_volume = np.prod(
+ _update_tuple(local_shape, dim, 1)
+ ) # 4 in the example above
+ volume_before_split = other_dims_volume * start # 4 in the example above
+ volume_of_split = other_dims_volume * length # 4 in the example above
+
+ flat_slice_start_shifted = (
+ self.flattened_range.start - volume_before_split
+ ) # 5 - 4 = 1 in the example above
+ flat_slice_stop_shifted = (
+ self.flattened_range.stop - volume_before_split
+ ) # 10 - 4 = 6 in the example above
+
+ # Find an intersection of
+ # (flat_slice_start_shifted, flat_slice_stop_shifted) vs (0, volume_of_split)
+
+ if flat_slice_stop_shifted <= 0 or flat_slice_start_shifted >= volume_of_split:
+ return [] # no intersection
+
+ # new_flattened_range = slice(1, 4) in the example above
+ new_flattened_range = slice(
+ max(flat_slice_start_shifted, 0), min(flat_slice_stop_shifted, volume_of_split)
+ )
+ # Apply the intersection to the flattened data tensor.
+ # Compute start and slice appropriate length
+ intersection_slice_start = (
+ new_flattened_range.start - flat_slice_start_shifted
+ ) # 0 in the example above
+ new_data = self.data[
+ intersection_slice_start : intersection_slice_start
+ + new_flattened_range.stop
+ - new_flattened_range.start
+ ]
+
+ return [
+ replace(
+ self,
+ data=new_data,
+ local_shape=local_shape,
+ global_shape=global_shape,
+ global_offset=global_offset,
+ flattened_range=new_flattened_range,
+ )
+ ]
+
+
+def is_main_replica(replica_id: ReplicaId):
+ """Checks if given `replica_id` is considered as main.
+
+ "Main" replica is:
+ - integer 0
+ - or an iterable with all 0 elements
+
+ It is the application responsibility to set correct replicas for sharded tensors.
+
+ Args:
+ replica_id (Union[int, Tuple[int, ...]]): replica id
+
+ Returns:
+ (bool): True for a "main" replica
+ """
+ if isinstance(replica_id, int):
+ return replica_id == 0
+ return all(r == 0 for r in replica_id)
+
+
+class LocalNonpersistentObject:
+ """Object that should not be stored in a checkpoint, but restored locally.
+
+ Wrapping any object inside the state dict with LocalNonpersistentObject
+ will result in:
+ - during saving, this object will *not* be stored in the checkpoint
+ - during loading, a local version of this object will be placed in a state dict
+ """
+
+ def __init__(self, obj):
+ self.obj = obj
+
+ def unwrap(self):
+ """Returns the original object."""
+ return self.obj
+
+
+# TODO: Delete once NeMo fixes typo.
+LocalNonpersitentObject = LocalNonpersistentObject
+
+
+@dataclass
+class ShardedObject(ShardedBase):
+ """Represents a mapping between a local object and a global object.
+
+ Global object is assumed to consist of many local objects distributed
+ between different processes.
+
+ NOTE: Contrary to ShardedTensor, it's impossible to change global object
+ sharding. Conceptually, ShardedObject is a fully-sharded ShardedTensor
+ with atomic arbitrary typed elements.
+
+ Args:
+ key: unique identifier of a global tensor
+ data: local object data. Can be None only for consistency validation
+ global_shape: global object shape
+ global_offset: offset of a local object in a global object, specified in number of shards
+ replica_id: indicates local object replication wrt. local objects in different processes
+ """
+
+ key: str
+ data: object
+ global_shape: Tuple[int, ...]
+ global_offset: Tuple[int, ...]
+ replica_id: ReplicaId = 0
+
+ def __post_init__(self):
+ self.validate_metadata_integrity()
+
+ def validate_metadata_integrity(self):
+ if len(self.global_shape) != len(self.global_offset):
+ raise CheckpointingException(
+ f'Global offset dimensions should be equal to global shape dimensions for {self}'
+ )
+
+ def without_data(self):
+ return replace(self, data=None)
+
+ @property
+ def unique_key(self):
+ """returns a unique key for this object"""
+ return (
+ f'{self.key}/shard_'
+ f'{".".join(map(str, self.global_offset))}_'
+ f'{".".join(map(str, self.global_shape))}'
+ )
+
+ def __str__(self):
+ return f'{self.__class__.__name__}(key=\'{self.key}\')'
+
+ @classmethod
+ def empty_from_unique_key(cls, unique_key, replica_id: ReplicaId = 0) -> 'ShardedObject':
+ """Instantiates a ShardedObject from a unique key.
+
+ Args:
+ unique_key: a string of the form
+ /shard__
+ replica_id: indicates local object replication wrt.
+ local objects in different processes
+
+ Returns:
+ a ShardedObject with data=None
+ """
+ key, shard_key = unique_key.split('/')
+ shard_str, offset, shape = shard_key.split('_')
+ assert shard_str == 'shard'
+ offset = tuple(map(int, offset.split('.')))
+ shape = tuple(map(int, shape.split('.')))
+ if len(shape) + 1 == len(offset):
+ # This is a backward-compatible fix. We don't know the last
+ # element of global shape so set it to -1.
+ shape += (-1,)
+ return cls(key, None, shape, offset, replica_id)
+
+
+FactoryBuildFn = Callable[[str, torch.Tensor, ReplicaId, Optional[slice]], ShardedStateDict]
+FactoryMergeFn = Callable[[StateDict], torch.Tensor]
+
+
+@dataclass
+class ShardedTensorFactory(ShardedBase):
+ """Allows to apply transformations to tensors before/after serialization.
+
+ The essence of those transformations is that they can be applied to
+ optimizer states the same way they are applied to the model params.
+ The ultimate state dict with sharded tensors must depend functionally on
+ `build_fn` arguments (key, data, replica_id, flattened_range),
+ which will be provided by the optimizer.
+
+ Builder creates a sub-state-dict out of a tensor before saving, and merger
+ merges the corresponding state dict after loading.
+
+ Args:
+ key (str): unique identifier of the factory
+ data (torch.Tensor): original model parameter that will be further
+ transformed by this factory
+ build_fn (callable): function that transforms the original tensor
+ to a sharded state dict
+ merge_fn (callable): function that transforms loaded subtree back
+ into a single tensor (inverse of `build_fn`)
+ replica_id (ReplicaId): indicates factory replication wrt.
+ factories in different processes
+ flattened_range (slice, optional): indicates additional flattening
+ applied to the ShardedTensors produced by the factory
+ """
+
+ key: str
+ data: torch.Tensor
+ build_fn: FactoryBuildFn
+ merge_fn: FactoryMergeFn
+ replica_id: ReplicaId = 0
+ flattened_range: Optional[slice] = None
+
+ def build(self):
+ """Builds a ShardedStateDict from the original tensor"""
+ return self.build_fn(self.key, self.data, self.replica_id, self.flattened_range)
+
+ def validate_metadata_integrity(self):
+ """No reasonable checks can be applied"""
+ pass
+
+ def without_data(self):
+ return replace(self, data=None)
+
+
+def apply_factories(sharded_state_dict: ShardedStateDict):
+ """Turn ShardedTensorFactories into ShardedTensors *in-place*.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): state dict possibly
+ containing ShardedTensorFactory objects
+
+ Returns:
+ None: state dict is modified in place
+ """
+
+ def apply(x):
+ if isinstance(x, ShardedTensorFactory):
+ x = x.build()
+ return x
+
+ dict_list_map_inplace(apply, sharded_state_dict)
+
+
+def apply_factory_merges(
+ x1: StateDict, x2: ShardedStateDict, key: Tuple[str, ...] = ()
+) -> StateDict:
+ """Apply merges defined by ShardedTensorFactories *in-place*.
+
+ Args:
+ x1 (StateDict): state dict loaded from the checkpoint
+ x2 (ShardedStateDict): subset of `x1` (in terms of dict keys)
+ with ShardedTensorFactory
+ as (possibly nested) values that define how to
+ merge objects from the `x1` state dict
+ key (Tuple[str, ...]): current key in a recursive call.
+ Used only for reporting meaningful errors
+
+ Returns:
+ StateDict: `x1` modified in-place
+ """
+ if isinstance(x2, ShardedTensorFactory):
+ return x2.merge_fn(x1)
+
+ # There rest is almost the same as the `merge` function from `dict_utils`
+ if isinstance(x1, dict) and isinstance(x2, dict):
+ for k, v2 in x2.items():
+ if k not in x1:
+ raise ValueError(
+ f'Different dict keys encountered in `apply_factory_merges` '
+ f'({x1.keys()} vs {x2.keys()})'
+ )
+ else:
+ x1[k] = apply_factory_merges(x1[k], v2, key=key + (k,))
+ elif isinstance(x1, list) and isinstance(x2, list):
+ if len(x1) != len(x2):
+ err_msg = (
+ f'Cannot merge two lists with different lengths '
+ f'({len(x1)} and {len(x2)}, encountered at key {key})'
+ )
+ logger.error(err_msg + f'\nx1: {x1}\nx2: {x2}')
+ raise ValueError(err_msg)
+ for i, v2 in enumerate(x2):
+ x1[i] = apply_factory_merges(x1[i], v2, key=key + (i,))
+ elif isinstance(x1, list) and isinstance(x2, dict):
+ for k, v2 in x2.items():
+ if not isinstance(k, int):
+ raise ValueError(
+ f'Invalid dict key {k} non-integer type encountered '
+ f'in a list-dict merge at level {key}'
+ )
+ if k >= len(x1):
+ raise ValueError(
+ f'Dict key {k} out of bound for list of length'
+ f'{len(x1)} (encountered at level {key})'
+ )
+ x1[k] = apply_factory_merges(x1[k], v2, key=key + (k,))
+ else:
+ raise ValueError(
+ f'Duplicate non-dict and non-list values encountered: `{x1}` and `{x2} (at key {key})`'
+ )
+ return x1
diff --git a/megatron/core/dist_checkpointing/optimizer.py b/megatron/core/dist_checkpointing/optimizer.py
new file mode 100644
index 0000000000..2d231a24ff
--- /dev/null
+++ b/megatron/core/dist_checkpointing/optimizer.py
@@ -0,0 +1,129 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Helpers for defining sharding for optimizer states based on existing sharding for model parameters. """
+
+import logging
+from copy import deepcopy
+from dataclasses import replace
+from itertools import chain
+from typing import Dict, Iterable, List, Tuple, Union
+
+logger = logging.getLogger(__name__)
+
+import torch
+
+from .dict_utils import nested_values
+from .mapping import (
+ LocalNonpersistentObject,
+ ShardedStateDict,
+ ShardedTensor,
+ ShardedTensorFactory,
+ StateDict,
+)
+from .utils import extract_sharded_tensors_and_factories
+
+
+def get_optim_param_to_id_map(optim_params_iter: Iterable[torch.nn.Parameter]) -> Dict[int, int]:
+ param_mappings = {}
+ for i, param in enumerate(optim_params_iter):
+ if id(param) not in param_mappings:
+ param_mappings[id(param)] = i
+ return param_mappings
+
+
+def get_param_id_to_sharded_param_map(
+ model_sharded_state_dict: ShardedStateDict, optim_params_iter: Iterable[torch.nn.Parameter]
+) -> Dict[int, Union[ShardedTensor, ShardedTensorFactory]]:
+ """Generate mapping from optimizer state ids to model sharded parameters.
+
+ Args:
+ model_sharded_state_dict: sharded state dict with all model sharded tensors (can have any structure)
+ optim_params_iter: iterable which iterates over model parameters tracked by the optimizer.
+ The iteration must be in the same order as in the optimizer parameters.
+
+ Returns:
+ Dict[int, Union[ShardedTensor, ShardedTensorFactory]]: mapping from optimizer state ids
+ to model sharded parameters.
+ """
+ model_sharded_state_dict, _ = extract_sharded_tensors_and_factories(model_sharded_state_dict)
+ id_to_sharded_param_map = {}
+ param_to_id_map = get_optim_param_to_id_map(optim_params_iter)
+ for ten in nested_values(model_sharded_state_dict):
+ if id(ten.data) in param_to_id_map:
+ id_to_sharded_param_map[param_to_id_map[id(ten.data)]] = ten
+ else:
+ logger.debug(f'{ten} is not tracked by the optimizer')
+
+ if not id_to_sharded_param_map:
+ logger.warning(
+ "Sharded parameters mapping is empty. It means tensors in model state dict"
+ " do not correspond to tensors in optimizer parameters map."
+ " Make sure to call state_dict with `keep_vars=True`."
+ )
+ return id_to_sharded_param_map
+
+
+def make_sharded_optimizer_tensor(
+ model_param: Union[ShardedTensor, ShardedTensorFactory], optim_param: torch.Tensor, prefix: str
+) -> Union[ShardedTensor, ShardedTensorFactory]:
+ """Build a ShardedTensor or ShardedTensorFactory for optimizer param based on model param
+
+ Args:
+ model_param (Union[ShardedTensor, ShardedTensorFactory]): model param
+ optim_param (torch.Tensor): corresponding optimizer param
+ prefix (str): optimizer prefix for the ShardedTensor or ShardedTensorFactory
+
+ Returns:
+ Union[ShardedTensor, ShardedTensorFactory]: wrapped optimizer parameter
+ """
+ if isinstance(model_param, ShardedTensorFactory):
+ return replace(model_param, key=f'{prefix}.{model_param.key}', data=optim_param)
+
+ assert (
+ tuple(optim_param.shape) == model_param.local_shape
+ ), f'Optimizer shape ({tuple(optim_param.shape)} does not match model shape ({model_param.local_shape})'
+ sh_ten = replace(
+ model_param, key=f'{prefix}.{model_param.key}', data=optim_param, dtype=optim_param.dtype
+ )
+ sh_ten.validate_metadata_integrity()
+ return sh_ten
+
+
+def optim_state_to_sharding_state(
+ optim_state_dict: StateDict,
+ id_to_sharded_param_map: Dict[int, ShardedTensor],
+ exclude_keys: Tuple[str] = (),
+):
+ """Turn optimizer state dict to sharded state dict based on model state dict *in-place*.
+
+ Can be used to add sharding information to most common optimizer state dict.
+ Creates separate ShardedTensors for each key in `optim_state_dict['state']`
+ (e.g. for torch.optim.Adam there will be separate tensors for `exp_avg` and `exp_avg_sq`)
+
+ Args:
+ optim_state_dict (StateDict): optimizer state dict with
+ state parameters under `state` key and group hyperparameters under `param_groups` -> `params` key.
+ id_to_sharded_param_map (Dict[int, ShardedTensor]): mapping from optimizer param ids to model sharded tensors.
+ Can be generated with `get_param_id_to_sharded_param_map` function
+ exclude_keys (Tuple[str]): optimizer state keys to exclude from the final state dict.
+
+ Returns:
+ None: state dict is modified in place
+ """
+ sharded_state = {}
+ for param_id, param_state in optim_state_dict['state'].items():
+ sharded_state[param_id] = {}
+ for state_key, param in param_state.items():
+ if state_key in exclude_keys:
+ continue
+ if param_id in id_to_sharded_param_map:
+ sharded_state[param_id][state_key] = make_sharded_optimizer_tensor(
+ id_to_sharded_param_map[param_id], param, prefix=f'optimizer.state.{state_key}'
+ )
+ else:
+ raise ValueError(f'Param id {param_id} does not match any model sharded param')
+
+ optim_state_dict['param_groups'] = deepcopy(optim_state_dict['param_groups'])
+ for group in optim_state_dict['param_groups']:
+ group['params'] = LocalNonpersistentObject(group['params'])
+ optim_state_dict['state'] = sharded_state
diff --git a/megatron/core/dist_checkpointing/serialization.py b/megatron/core/dist_checkpointing/serialization.py
new file mode 100644
index 0000000000..5493c96bbd
--- /dev/null
+++ b/megatron/core/dist_checkpointing/serialization.py
@@ -0,0 +1,411 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Entrypoints for saving and loading the distributed checkpoints.
+
+Functions `load` and `save` are equivalents of `torch.load` and `torch.save`
+but expect torch.Tensors to be wrapped with classes from the `mapping module`.
+Additionally, `load` expects the sharded state dict argument as a guidance for
+loading the sharded tensors.
+"""
+
+import logging
+from pathlib import Path
+from typing import Dict, Optional, Set, Tuple, Union
+
+import torch
+
+from . import ShardedTensor
+from .core import CheckpointingConfig, save_config
+from .dict_utils import extract_matching_values, merge
+from .mapping import (
+ CheckpointingException,
+ ShardedObject,
+ ShardedStateDict,
+ StateDict,
+ apply_factory_merges,
+)
+from .state_dict_transformation import load_preprocess, save_preprocess
+from .strategies.async_utils import AsyncRequest
+from .strategies.base import (
+ AsyncSaveShardedStrategy,
+ LoadCommonStrategy,
+ LoadShardedStrategy,
+ SaveCommonStrategy,
+ SaveShardedStrategy,
+ StrategyAction,
+ get_default_strategy,
+)
+from .utils import extract_sharded_base
+from .validation import (
+ StrictHandling,
+ determine_global_metadata,
+ parse_strict_flag,
+ validate_integrity_and_strict_load,
+ validate_sharded_objects_handling,
+ verify_checkpoint_and_load_strategy,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# flat state dict with sharded objects without any data
+CkptShardedMetadata = Dict[str, Union[ShardedTensor, ShardedObject]]
+
+
+def load(
+ sharded_state_dict: ShardedStateDict,
+ checkpoint_dir: str,
+ sharded_strategy: Union[LoadShardedStrategy, Tuple[str, int], None] = None,
+ common_strategy: Union[LoadCommonStrategy, Tuple[str, int], None] = None,
+ validate_access_integrity: bool = True,
+ strict: Union[str, StrictHandling] = StrictHandling.ASSUME_OK_UNEXPECTED,
+) -> Union[StateDict, Tuple[StateDict, Set[str], Set[str]]]:
+ """Loading entrypoint.
+
+ In the steps below, the following verbs refer to corresponding objects:
+ - load = load from checkpoint
+ - extract = extract from sharded_state_dict
+ - add = add to the final state dict
+ Steps:
+ 1. Load common state dict and form the base of the result state dict
+ 2. Apply factories to sharded_state_dict
+ 3. Extract LocalNonPersistentObject and add
+ 4. (optional) Extract ShardedObjects, load and add
+ 5. Extract ShardedBase, load, apply factory merges and add
+
+ Args:
+ sharded_state_dict (ShardedStateDict): state dict of the existing model
+ populated with ShardedTensors. Used as a mapping to determine which
+ parts of global tensors stored in the checkpoint should be loaded.
+ checkpoint_dir (str): directory with the checkpoint
+ sharded_strategy (LoadShardedStrategy, Tuple[str, int], optional):
+ configures loading behavior for sharded tensors
+ common_strategy (LoadCommonStrategy, Tuple[str, int], optional):
+ configures loading behavior for common data
+ validate_access_integrity (bool default = True): checks if each tensor shard is accessed
+ exactly once (as main replica) by some process
+ strict (StrictHandling, str, optional): determines the behavior in case of a mismatch
+ between the requested sharded state dict and the checkpoint. See `StrictHandling` docs
+ for more details. Some values affect the return value of this function
+ (missing and unexpected keys are returned).
+ Defaults to `True` (StrictHandling.ASSUME_OK_UNEXPECTED) which doesn't
+ incur any performance overhead. Other recommended values
+ are: `False` (StrictHandling.LOG_UNEXPECTED) which logs only unexpected keys
+ or `StrictHandling.RETURN_ALL` which returns all mismatch keys.
+
+ Returns:
+ StateDict or Tuple[StateDict, Set[str], Set[str]]: in most cases only
+ the loaded state dict is returned. If `strict` flag was set to
+ """
+ sharded_strategy, common_strategy = verify_checkpoint_and_load_strategy(
+ checkpoint_dir, sharded_strategy, common_strategy
+ )
+
+ checkpoint_dir = Path(checkpoint_dir)
+ common_state_dict = common_strategy.load_common(checkpoint_dir)
+ if not sharded_state_dict:
+ return common_state_dict
+
+ sharded_state_dict, nonpersistent_state_dict, sh_ten_factories = load_preprocess(
+ sharded_state_dict
+ )
+ merge(common_state_dict, nonpersistent_state_dict)
+
+ # At this point we are only dealing with ShardedBase objects
+ sharded_state_dict, _ = extract_sharded_base(sharded_state_dict)
+
+ # Validation
+ ckpt_sharded_metadata = None
+ local_metadata, global_metadata = None, None
+ strict = parse_strict_flag(strict)
+ if StrictHandling.requires_explicit_ckpt_mismatch_check(strict):
+ ckpt_sharded_metadata = load_sharded_metadata(
+ str(checkpoint_dir), sharded_strategy, common_strategy
+ )
+ if validate_access_integrity or StrictHandling.requires_global_app_metadata(strict):
+ local_metadata, global_metadata = determine_global_metadata(sharded_state_dict)
+
+ sharded_state_dict, missing_keys, unexpected_keys = validate_integrity_and_strict_load(
+ sharded_state_dict,
+ strict,
+ validate_access_integrity,
+ local_metadata,
+ global_metadata,
+ ckpt_sharded_metadata,
+ )
+
+ # ShardedBase loading
+ if not sharded_strategy.can_handle_sharded_objects:
+ validate_sharded_objects_handling(sharded_strategy, common_strategy)
+ sharded_objects_state_dict, sharded_state_dict = extract_matching_values(
+ sharded_state_dict, lambda v: isinstance(v, ShardedObject)
+ )
+ sharded_objects = common_strategy.load_sharded_objects(
+ sharded_objects_state_dict, checkpoint_dir
+ )
+ merge(common_state_dict, sharded_objects)
+
+ loaded_state_dict = sharded_strategy.load(sharded_state_dict, checkpoint_dir)
+
+ merge(common_state_dict, loaded_state_dict)
+
+ loaded_state_dict = apply_factory_merges(common_state_dict, sh_ten_factories)
+
+ if StrictHandling.requires_returning_mismatch_keys(strict):
+ return common_state_dict, missing_keys, unexpected_keys
+ else:
+ return common_state_dict
+
+
+def load_common_state_dict(checkpoint_dir: Path) -> StateDict:
+ """Load common (non-sharded) objects state dict from the checkpoint.
+
+ Args:
+ checkpoint_dir (Path): checkpoint directory
+
+ Returns:
+ StateDict: state dict with non-sharded objects from the checkpoint
+ """
+ sharded_strategy, common_strategy = verify_checkpoint_and_load_strategy(str(checkpoint_dir))
+ return common_strategy.load_common(checkpoint_dir)
+
+
+def load_tensors_metadata(
+ checkpoint_dir: str, sharded_strategy: Union[LoadShardedStrategy, None] = None
+) -> CkptShardedMetadata:
+ """Load tensors metadata from the checkpoint.
+
+ Returns a dictionary similar to a sharded state dict, but note that
+ the dictionary keys are simply ShardedTensor keys (contrary to the
+ actual sharded state dicts where keys correspond to state dict keys).
+
+ Dict values are ShardedTensors without any sharding (so, the only useful
+ information is tensors global shape and dtype).
+
+ Concrete implementation depends on the loading strategy. If no strategy is
+ given, a default for a given backend is used.
+
+ Args:
+ checkpoint_dir (str): checkpoint directory to load from
+ sharded_strategy (LoadShardedStrategy, optional): sharded strategy to load metadata.
+ Defaults to None - in this case a default load strategy for a given checkpoint type
+ is used.
+
+ Returns:
+ CkptShardedMetadata: flat state dict without data describing ShardedTensors
+ in the checkpoint
+ """
+ sharded_strategy, common_strategy = verify_checkpoint_and_load_strategy(
+ checkpoint_dir, sharded_strategy
+ )
+ return sharded_strategy.load_tensors_metadata(Path(checkpoint_dir))
+
+
+def load_sharded_metadata(
+ checkpoint_dir: str,
+ sharded_strategy: Union[LoadShardedStrategy, None] = None,
+ common_strategy: Union[LoadCommonStrategy, None] = None,
+) -> CkptShardedMetadata:
+ """Load sharded metadata from the checkpoint.
+
+ Similar to `load_tensors_metadata`, but includes also ShardedObjects.
+
+ Returns a dictionary similar to a sharded state dict, but note that
+ the dictionary keys are simply ShardedTensor keys (contrary to the
+ actual sharded state dicts where keys correspond to state dict keys).
+
+ Dict values are ShardedTensors without any sharding (so, the only useful
+ information is tensors global shape and dtype).
+
+ Concrete implementation depends on the loading strategy. If no strategy is
+ given, a default for a given backend is used.
+
+ Args:
+ checkpoint_dir (str): checkpoint directory to load from
+ sharded_strategy (LoadShardedStrategy, optional): sharded strategy to load metadata.
+ Defaults to None - in this case a default load strategy for a given checkpoint type
+ is used.
+ common_strategy (LoadCommonStrategy, optional): common strategy to load metadata.
+ Defaults to None - in this case a default load strategy for a given checkpoint type is
+ used. This strategy won't be used unless `sharded_strategy` can't handle ShardedObjects
+
+ Returns:
+ CkptShardedMetadata: flat state dict without data describing ShardedTensors
+ and ShardedObjects in the checkpoint
+ """
+ sharded_strategy, common_strategy = verify_checkpoint_and_load_strategy(
+ checkpoint_dir, sharded_strategy, common_strategy
+ )
+ sharded_metadata = sharded_strategy.load_sharded_metadata(Path(checkpoint_dir))
+ if not sharded_strategy.can_handle_sharded_objects:
+ validate_sharded_objects_handling(sharded_strategy, common_strategy)
+ common_metadata = common_strategy.load_sharded_metadata(Path(checkpoint_dir))
+ sharded_metadata = merge(sharded_metadata, common_metadata)
+ return sharded_metadata
+
+
+def load_plain_tensors(checkpoint_dir: str) -> StateDict:
+ """Load checkpoint tensors without any sharding and plain structure.
+
+ NOTE: common state dict is NOT included.
+
+ Args:
+ checkpoint_dir (str): checkpoint directory to load the tensors from.
+
+ Returns:
+ StateDict: checkpoint state dict containing only torch.Tensors.
+ """
+ sharded_state_dict = load_tensors_metadata(checkpoint_dir)
+ # Don't validate integrity because shards will be overlapped
+ # if world_size > 1 (all processes load whole tensors)
+ return load(sharded_state_dict, checkpoint_dir, validate_access_integrity=False)
+
+
+#
+# def load_plain_tensors_and_objects(checkpoint_dir: str) -> StateDict:
+# """Load checkpoint tensors and objects without any sharding and plain structure.
+#
+# NOTE: state dict structure might be different than the one used for checkpoint saving.
+# NOTE: common state dict is NOT included.
+#
+# Args:
+# checkpoint_dir (str): checkpoint directory to load the state dict from.
+#
+# Returns:
+# StateDict: complete checkpoint state dict without any sharding.
+# """
+# sharded_state_dict = load_tensors_metadata(checkpoint_dir)
+# # Don't validate integrity because shards will be overlapped
+# # if world_size > 1 (all processes load whole tensors)
+# return load(sharded_state_dict, checkpoint_dir, validate_access_integrity=False)
+
+
+def save(
+ sharded_state_dict: ShardedStateDict,
+ checkpoint_dir: str,
+ sharded_strategy: Union[SaveShardedStrategy, Tuple[str, int], None] = None,
+ common_strategy: Union[SaveCommonStrategy, Tuple[str, int], None] = None,
+ validate_access_integrity: bool = True,
+ async_sharded_save: bool = False,
+) -> Optional[AsyncRequest]:
+ """Saving entrypoint.
+
+ Extracts ShardedTensors from the given state dict. Rank 0 saves the
+ "regular" part of the checkpoint to common torch file.
+ The ShardedTensors are saved according to a strategy specified by the
+ config.
+
+ Steps:
+ 1. Apply factories
+ 2. Extract and discard LocalNonPersistentObject
+ 3. Extract all ShardedBase object
+ 4. Save all other objects to common.pt
+ 5. (optional) Extract and save ShardedObjects
+ 6. Save all ShardedBase objects
+ 7. Write metadata.json file with backend and version metadata.
+
+ Step (6) can be performed asynchronously (see `async_sharded_save`), in this
+ case the actual save is embodied in the returned async request and can be
+ scheduled by the external caller. For async request, step (7) is added as
+ one of the finalization functions, so that metadata.json is written only
+ if the checkpoint is complete.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): state dict of the populated with
+ ShardedTensors. Used as a mapping to determine how local tensors
+ should be saved as global tensors in the checkpoint.
+ checkpoint_dir (str): directory to save the checkpoint to
+ sharded_strategy (SaveShardedStrategy, Tuple[str, int], optional):
+ configures sharded tensors saving behavior and backend
+ common_strategy (SaveCommonStrategy, Tuple[str, int], optional):
+ configures common data saving behavior and backend
+ validate_access_integrity (bool default = True): checks if each tensor shard is accessed
+ exactly once (as main replica) by some process
+ async_sharded_save (bool, optional): if True, for the sharded state dict part
+ an async save implementation will be called, with the AsyncRequest
+ being returned to the caller. Note that it is the caller responsibility to
+ actually schedule the async save. Defaults to False.
+
+ Returns:
+ AsyncRequest (optional): if `async_sharded_save` is True, returns
+ async request that should be scheduled by the caller of this function.
+ None otherwise.
+ """
+ checkpoint_dir = Path(checkpoint_dir)
+
+ if torch.distributed.get_rank() == 0:
+ if not checkpoint_dir.exists():
+ raise CheckpointingException(
+ f'Checkpoint destination directory does not exist: {checkpoint_dir}'
+ )
+
+ if next(checkpoint_dir.iterdir(), None) is not None:
+ raise CheckpointingException(
+ f'Checkpoint destination directory ({checkpoint_dir}) is not empty'
+ )
+
+ if common_strategy is not None:
+ raise NotImplementedError('The only supported common strategy is torch')
+
+ if sharded_strategy is None:
+ sharded_strategy = get_default_save_sharded_strategy()
+ if not isinstance(sharded_strategy, SaveShardedStrategy):
+ assert isinstance(sharded_strategy, tuple), type(sharded_strategy)
+ sharded_strategy = get_default_strategy(StrategyAction.SAVE_SHARDED, *sharded_strategy)
+
+ if common_strategy is None:
+ common_strategy = get_default_save_common_strategy()
+ if not isinstance(common_strategy, SaveCommonStrategy):
+ assert isinstance(common_strategy, tuple), type(common_strategy)
+ common_strategy = get_default_strategy(StrategyAction.SAVE_COMMON, *common_strategy)
+
+ sharded_state_dict, state_dict = save_preprocess(sharded_state_dict, validate_access_integrity)
+
+ common_strategy.save_common(state_dict, checkpoint_dir)
+
+ if not sharded_strategy.can_handle_sharded_objects:
+ validate_sharded_objects_handling(sharded_strategy, common_strategy)
+ sharded_objects_state_dict, sharded_state_dict = extract_matching_values(
+ sharded_state_dict, lambda v: isinstance(v, ShardedObject)
+ )
+ common_strategy.save_sharded_objects(sharded_objects_state_dict, checkpoint_dir)
+
+ def metadata_finalize_fn():
+ if torch.distributed.get_rank() == 0:
+ save_config(
+ CheckpointingConfig(sharded_strategy.backend, sharded_strategy.version),
+ checkpoint_dir,
+ )
+ torch.distributed.barrier()
+
+ if not async_sharded_save:
+ sharded_strategy.save(sharded_state_dict, checkpoint_dir)
+ metadata_finalize_fn()
+ return
+
+ if not isinstance(sharded_strategy, AsyncSaveShardedStrategy):
+ raise CheckpointingException(
+ f'Cannot apply async_save to non-async strategy {sharded_strategy}'
+ )
+ async_request = sharded_strategy.async_save(sharded_state_dict, checkpoint_dir)
+ async_request.finalize_fns.append(metadata_finalize_fn)
+ return async_request
+
+
+def get_default_save_sharded_strategy(
+ backend: str = 'torch_dist', version: int = 1
+) -> SaveShardedStrategy:
+ """Get default save sharded strategy."""
+ return get_default_strategy(StrategyAction.SAVE_SHARDED, backend, version)
+
+
+def get_default_save_common_strategy(
+ backend: str = 'torch', version: int = 1
+) -> SaveCommonStrategy:
+ """Get default save common strategy."""
+ return get_default_strategy(StrategyAction.SAVE_COMMON, backend, version)
+
+
+def get_default_load_sharded_strategy(checkpoint_dir: str) -> LoadShardedStrategy:
+ """Get default load sharded strategy."""
+ return verify_checkpoint_and_load_strategy(checkpoint_dir)[0]
diff --git a/megatron/core/dist_checkpointing/state_dict_transformation.py b/megatron/core/dist_checkpointing/state_dict_transformation.py
new file mode 100644
index 0000000000..ebb960e384
--- /dev/null
+++ b/megatron/core/dist_checkpointing/state_dict_transformation.py
@@ -0,0 +1,253 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Utilities for transforming state_dict, including a tensor-aware implementation."""
+
+import logging
+from time import time
+from typing import Any, Optional
+
+import torch
+
+from .dict_utils import dict_list_map_inplace, extract_matching_values, merge, nested_values
+from .exchange_utils import determine_main_replica_uniform_distribution, exchange_by_distribution
+from .mapping import (
+ ShardedObject,
+ ShardedStateDict,
+ ShardedTensor,
+ ShardedTensorFactory,
+ apply_factories,
+ apply_factory_merges,
+)
+from .utils import (
+ _sharded_object_id,
+ _sharded_tensor_shard_id,
+ extract_nonpersistent,
+ extract_sharded_base,
+)
+from .validation import determine_global_metadata, validate_sharding_integrity
+
+logger = logging.getLogger(__name__)
+
+
+def save_preprocess(sharded_state_dict: ShardedStateDict, validate_access_integrity: bool = True):
+ """Preprocesses the given state dictionary by applying factories,
+ discarding non-persistent data and extracting the common state dictionary.
+ Optionally, it can validate sharding integrity.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): The initial state dictionary to be preprocessed.
+ validate_access_integrity (bool): If True, triggers validation of sharding integrity.
+
+ Returns:
+ Tuple[ShardedStateDict, dict]:
+ The preprocessed sharded state dictionary and the common state dictionary.
+ """
+ apply_factories(sharded_state_dict)
+ _, sharded_state_dict = extract_nonpersistent(sharded_state_dict)
+ sharded_part, common_state_dict = extract_sharded_base(sharded_state_dict)
+ if validate_access_integrity:
+ validate_sharding_integrity(determine_global_metadata(sharded_part)[1])
+ return sharded_part, common_state_dict
+
+
+def load_preprocess(sharded_state_dict: ShardedStateDict):
+ """Preprocesses the given state dictionary by applying factories
+ and extracting non-persistent data, without modifying the original dictionary.
+
+ Args:
+ sharded_state_dict (ShardedStateDict):
+ The initial state dictionary to be processed (remains unchanged).
+
+ Returns:
+ Tuple[ShardedStateDict, dict, dict]:
+ - A preprocessed copy of the sharded state dictionary.
+ - A dictionary containing non-persistent state data.
+ - A dictionary of `ShardedTensorFactory` instances.
+ """
+ # Create a copy of sharded_state_dict as the passed in state dict may have
+ # references that prevent tensors from being deallocated
+ sharded_state_dict, _ = extract_matching_values(sharded_state_dict, lambda x: True)
+
+ sh_ten_factories, _ = extract_matching_values(
+ sharded_state_dict,
+ lambda x: isinstance(x, ShardedTensorFactory),
+ return_lists_as_dicts=True,
+ )
+ apply_factories(sharded_state_dict)
+
+ # Data inside sh_ten_factories no longer needed so delete them to reduce memory usage
+ dict_list_map_inplace(ShardedTensorFactory.without_data, sh_ten_factories)
+ # Non-persistent objects
+ nonpersistent_state_dict, sharded_state_dict = extract_nonpersistent(sharded_state_dict)
+ dict_list_map_inplace(lambda o: o.unwrap(), nonpersistent_state_dict)
+ return sharded_state_dict, nonpersistent_state_dict, sh_ten_factories
+
+
+def prepare_state_dict_for_save(
+ sharded_state_dict: ShardedStateDict,
+ async_prepare: bool = False,
+ algo: str = 'atomic',
+ validate_access_integrity: bool = True,
+ parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
+ to_cpu: bool = True,
+):
+ """Creates a tensor-aware state dictionary that can be saved using the Local Checkpoint Manager.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): The initial state dictionary.
+ async_prepare (bool): If True, enables asynchronous preparation.
+ algo (str): The algorithm used to create the tensor-aware state dictionary.
+ validate_access_integrity (bool): If True, validates sharding integrity.
+ parallelization_group (torch.distributed.ProcessGroup):
+ The process group used for exchanges to avoid duplications.
+ to_cpu (bool): If True, moves all tensors from device to CPU.
+
+ Returns:
+ ShardedStateDict: The tensor-aware state dictionary.
+ """
+
+ _start = time()
+
+ if async_prepare:
+ raise NotImplementedError('Async state_dict preparation is not yet implemented')
+ if algo != 'atomic' and algo != 'fully_parallel':
+ raise NotImplementedError(
+ 'Only "atomic" and "fully_parallel" sharding algorithms are supported.'
+ )
+ fully_parallel = algo == 'fully_parallel'
+
+ sharded_part, common_state_dict = save_preprocess(sharded_state_dict, validate_access_integrity)
+ sharded_tensors = []
+ sharded_objects = []
+ for sh_base in nested_values(sharded_part):
+ if isinstance(sh_base, ShardedTensor):
+ sharded_tensors.append(sh_base)
+ else:
+ assert isinstance(sh_base, ShardedObject)
+ sharded_objects.append(sh_base)
+ if fully_parallel:
+ shard_to_saving_rank, _, shard_to_metadata = determine_main_replica_uniform_distribution(
+ sharded_part, parallelization_group, True
+ )
+
+ raw_tensors, raw_objects = {}, {}
+ for ten in sharded_tensors:
+ shard_id = _sharded_tensor_shard_id(ten)
+ if not fully_parallel or shard_to_saving_rank[shard_id] == torch.distributed.get_rank():
+ # TODO cover creating copies on host in CheckpointManager.save()
+ if to_cpu:
+ raw_tensors[shard_id] = ten.data.to("cpu", non_blocking=True)
+ else:
+ raw_tensors[shard_id] = ten.data
+ ten.data = None
+ for obj in sharded_objects:
+ raw_objects[_sharded_object_id(obj)] = obj.data
+ obj.data = None
+
+ logger.debug(f'prepare_state_dict_for_save took {time() - _start}')
+
+ state_dict_for_save = {
+ 'raw_tensors': raw_tensors,
+ 'raw_objects': raw_objects,
+ 'common': common_state_dict,
+ 'sharded_state_dict': sharded_part,
+ }
+ if fully_parallel:
+ state_dict_for_save['shard_to_rank'] = shard_to_saving_rank
+ state_dict_for_save['shard_to_metadata'] = shard_to_metadata
+ return state_dict_for_save
+
+
+def recreate_state_dict_after_load(
+ sharded_state_dict: ShardedStateDict,
+ loaded_state_dict: ShardedStateDict,
+ algo: str = 'atomic',
+ exchange_algo: str = 'broadcast',
+ validate_access_integrity: bool = True,
+ parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
+):
+ """Creates a final sharded state dictionary from a tensor-aware state dictionary.
+
+ Args:
+ sharded_state_dict (ShardedStateDict):
+ The initial sharded state dictionary generated from the model.
+ loaded_state_dict (ShardedStateDict):
+ Tensor-aware state dictionary used to fill in missing data in the sharded state.
+ algo (str): The algorithm used to reconstruct the state dictionary
+ from the tensor-aware state dictionary.
+ exchange_algo (str): The algorithm used for tensor exchanges during retrieval.
+ validate_access_integrity (bool): If True, performs validation of sharding integrity.
+ parallelization_group (torch.distributed.ProcessGroup):
+ The process group used for efficient exchanges during retrieval.
+
+ Returns:
+ ShardedStateDict: The finalized sharded state dictionary.
+ """
+
+ if algo != 'atomic' and algo != 'fully_parallel':
+ raise NotImplementedError(
+ 'Only "atomic" and "fully_parallel" sharding algorithms are supported.'
+ )
+ fully_parallel = algo == 'fully_parallel'
+
+ # __adding__ common part
+ recreated_state_dict, _ = extract_matching_values(loaded_state_dict["common"], lambda x: True)
+
+ if not sharded_state_dict:
+ return recreated_state_dict
+ # TODO validate laoded_state_dict["sharded_state_dict"] and sharded_state_dict are compatible
+
+ sharded_state_dict, nonpersistent_state_dict, sh_ten_factories = load_preprocess(
+ sharded_state_dict
+ )
+ # __adding__ nonpersistent part
+ merge(recreated_state_dict, nonpersistent_state_dict)
+
+ sharded_part, _ = extract_sharded_base(sharded_state_dict)
+ if validate_access_integrity:
+ validate_sharding_integrity(determine_global_metadata(sharded_part)[1])
+
+ # load sharded tensors and sharded objects to sharded_part
+ loaded_tensors = loaded_state_dict['raw_tensors']
+ # TODO cover restoring the original device (H2D) in CheckpointManager.load()
+ for k, v in loaded_tensors.items():
+ loaded_tensors[k] = v.cuda() # H2D
+ if fully_parallel:
+ distribution = (
+ loaded_state_dict['shard_to_rank'],
+ None,
+ loaded_state_dict['shard_to_metadata'],
+ )
+ unloaded_shards = {}
+ for sh_base in nested_values(sharded_part):
+ if isinstance(sh_base, ShardedTensor):
+ shard_id = _sharded_tensor_shard_id(sh_base)
+ if shard_id not in loaded_tensors:
+ unloaded_shards[shard_id] = sh_base
+ loaded_tensors = exchange_by_distribution(
+ loaded_tensors, unloaded_shards, distribution, parallelization_group, exchange_algo
+ )
+ loaded_objects = loaded_state_dict['raw_objects']
+
+ def load_sharded_base(x: Any):
+ if isinstance(x, ShardedTensor):
+ shard_id = _sharded_tensor_shard_id(x)
+ if shard_id not in loaded_tensors:
+ raise Exception(
+ 'The current local checkpoint implementation assumes'
+ 'consistent tensor sharding during load and save operations.'
+ f'However, the expected shard {x} (ID: {shard_id})'
+ f'was not found in the checkpoint. (IDs: {loaded_tensors.keys()})'
+ )
+ x = loaded_tensors[shard_id]
+ if isinstance(x, ShardedObject):
+ object_id = _sharded_object_id(x)
+ assert object_id in loaded_objects, (x, object_id, loaded_objects.keys())
+ x = loaded_objects[object_id]
+ return x
+
+ dict_list_map_inplace(load_sharded_base, sharded_part)
+ sharded_part = apply_factory_merges(sharded_part, sh_ten_factories)
+ # __adding__ sharded_part
+ merge(recreated_state_dict, sharded_part)
+ return recreated_state_dict
diff --git a/megatron/core/dist_checkpointing/strategies/__init__.py b/megatron/core/dist_checkpointing/strategies/__init__.py
new file mode 100644
index 0000000000..a786b8e84a
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Various loading and saving strategies """
+from megatron.core.dist_checkpointing.strategies.common import register_default_common_strategies
+
+# We load "common" strategies by default to be always available
+register_default_common_strategies()
diff --git a/megatron/core/dist_checkpointing/strategies/async_utils.py b/megatron/core/dist_checkpointing/strategies/async_utils.py
new file mode 100644
index 0000000000..7cdda8ac32
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/async_utils.py
@@ -0,0 +1,224 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""
+This module provides an async utilities which allow to start
+a checkpoint save process in the background.
+"""
+import logging
+from collections import deque
+from time import time
+from typing import Callable, List, NamedTuple, Optional, Tuple
+
+import torch
+from torch import multiprocessing as mp
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncRequest(NamedTuple):
+ """Represents an async request that needs to be scheduled for execution.
+
+ Args:
+ async_fn (Callable, optional): async function to call. None represents noop.
+ async_fn_args (Tuple): args to pass to `async_fn`.
+ finalize_fns (List[Callable]): list of functions to call to finalize the request.
+ These functions will be called synchronously after `async_fn` is done
+ *on all ranks*.
+ """
+
+ async_fn: Optional[Callable]
+ async_fn_args: Tuple
+ finalize_fns: List[Callable]
+ is_frozen: bool = False
+
+ def add_finalize_fn(self, fn: Callable) -> None:
+ """Adds a new finalize function to the request.
+
+ Args:
+ fn (Callable): function to add to the async request. This function
+ will be called *after* existing finalization functions.
+
+ Returns:
+ None
+ """
+ if self.is_frozen:
+ raise RuntimeError('Cannot add finalization functions to a frozen AsyncRequest')
+ self.finalize_fns.append(fn)
+
+ def execute_sync(self) -> None:
+ """Helper to synchronously execute the request.
+
+ This logic is equivalent to what should happen in case of the async call.
+ """
+ if self.async_fn is not None:
+ self.async_fn(*self.async_fn_args)
+ torch.distributed.barrier()
+ for finalize_fn in self.finalize_fns:
+ finalize_fn()
+
+ def freeze(self) -> 'AsyncRequest':
+ """Freezes the async request, disallowing adding new finalization functions.
+
+ Returns:
+ AsyncRequest: new async request with all same fields except for the
+ `is_frozen` flag.
+ """
+ return self._replace(is_frozen=True)
+
+
+class DistributedAsyncCaller:
+ """Wrapper around mp.Process that ensures correct semantic of distributed finalization.
+
+ Starts process asynchronously and allows checking if all processes on all ranks are done.
+ """
+
+ def __init__(self):
+ self.process: Optional[mp.Process] = None
+ self.start_time: Optional[float] = None
+
+ def schedule_async_call(self, async_fn: Optional[Callable], save_args: Tuple) -> None:
+ """Spawn a process with `async_fn` as the target.
+
+ This method must be called on all ranks.
+
+ Args:
+ async_fn (Callable, optional): async function to call. If None,
+ no process will be started.
+ save_args (Tuple): async function args.
+ """
+ if async_fn is None:
+ return # nothing to do
+ start_sync = time()
+ torch.cuda.synchronize()
+ end_sync = time()
+ logger.debug(
+ f"rank: {torch.distributed.get_rank()}, takes {end_sync - start_sync} to finish D2H "
+ )
+
+ ctx = mp.get_context('fork')
+ self.start_time = time()
+ self.process = ctx.Process(target=async_fn, args=save_args)
+ self.process.start()
+ init_time = time()
+ logger.debug(
+ f"rank: {torch.distributed.get_rank()}, takes {init_time - self.start_time} to schedule async ckpt "
+ )
+
+ def is_current_async_call_done(self, blocking=False) -> bool:
+ """Check if async save is finished on all ranks.
+
+ For semantic correctness, requires rank synchronization in each check.
+ This method must be called on all ranks.
+
+ Args:
+ blocking (bool, optional): if True, will wait until the call is done
+ on all ranks. Otherwise, returns immediately if at least one rank
+ is still active. Defaults to False.
+
+ Returns:
+ bool: True if all ranks are done (immediately of after active wait
+ if `blocking` is True), False if at least one rank is still active.
+ """
+ # The following takes the same overhead as torch.distributed.barrier (single integer all-reduce)
+ is_alive = int(self.process.is_alive()) if self.process is not None else 0
+ ten = torch.tensor([is_alive], dtype=torch.int, device=torch.cuda.current_device())
+ logger.debug(
+ f"rank: {torch.distributed.get_rank()}, DistributedAsyncCaller is_alive: {is_alive}"
+ )
+ torch.distributed.all_reduce(ten)
+ if ten[0] > 0 and not blocking:
+ return False
+ else:
+ if self.process is not None:
+ logger.debug(f"rank: {torch.distributed.get_rank()}, joining self.process")
+ self.process.join()
+ self.process = None
+
+ logger.debug(
+ f"DistributedAsyncCaller: Async process join finished after {time() - self.start_time:.2f}s from forking"
+ )
+ self.start_time = None
+ return True
+
+
+class _ActiveAsyncRequest(NamedTuple):
+ """Helper to represent an active async call.
+
+ Args:
+ idx (int): index of the call (starting from 0)
+ async_caller (DistributedAsyncCaller): async caller instance that represents
+ the async process handling the async request
+ async_request (AsyncRequest): async request that is being called
+ """
+
+ idx: int
+ async_caller: DistributedAsyncCaller
+ async_request: AsyncRequest
+
+
+class AsyncCallsQueue:
+ """Manages a queue of async calls.
+
+ Allows adding a new async call with `schedule_async_request` and finalizing
+ active calls with `maybe_finalize_async_calls`.
+ """
+
+ def __init__(self):
+ self.async_calls: deque[_ActiveAsyncRequest] = deque([])
+ self.call_idx: int = -1
+
+ def schedule_async_request(self, async_request: AsyncRequest) -> int:
+ """Start a new async call and add it to a queue of active async calls.
+
+ This method must be called on all ranks.
+
+ Args:
+ async_request (AsyncRequest): async request to start.
+
+ Returns:
+ int: index of the async call that was started.
+ This can help the user keep track of the async calls.
+ """
+ self.call_idx += 1
+ async_caller = DistributedAsyncCaller()
+ async_request = async_request.freeze()
+ async_caller.schedule_async_call(async_request.async_fn, async_request.async_fn_args)
+ self.async_calls.append(_ActiveAsyncRequest(self.call_idx, async_caller, async_request))
+ return self.call_idx
+
+ def maybe_finalize_async_calls(self, blocking=False) -> List[int]:
+ """Finalizes all available calls.
+
+ This method must be called on all ranks.
+
+ Args:
+ blocking (bool, optional): if True, will wait until all active requests
+ are done. Otherwise, finalizes only the async request that already
+ finished. Defaults to False.
+ Returns:
+ List[int]: list of indices (as returned by `schedule_async_request`)
+ of async calls that have been successfully finalized.
+ """
+ call_idx_finalized = []
+ while self.async_calls:
+ next_async_done = self.async_calls[0].async_caller.is_current_async_call_done(blocking)
+ if not next_async_done:
+ break
+ call_idx, _, async_request = self.async_calls.popleft()
+ for finalize_fn in async_request.finalize_fns:
+ finalize_fn()
+ ten = torch.tensor([call_idx], dtype=torch.int, device=torch.cuda.current_device())
+ torch.distributed.all_reduce(ten, op=torch.distributed.ReduceOp.MAX)
+ assert (
+ ten.item() == call_idx
+ ), 'Unmatched async calls. That probably means not all ranks are participating in async finalization'
+ call_idx_finalized.append(call_idx)
+ return call_idx_finalized
+
+ def get_num_unfinalized_calls(self):
+ """Get the number of active async calls."""
+ return len(self.async_calls)
+
+ def close(self):
+ """Finalize all calls upon closing."""
+ self.maybe_finalize_async_calls(blocking=True)
diff --git a/megatron/core/dist_checkpointing/strategies/base.py b/megatron/core/dist_checkpointing/strategies/base.py
new file mode 100644
index 0000000000..35fca1f350
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/base.py
@@ -0,0 +1,223 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Strategies base interfaces. """
+
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from enum import Enum
+from pathlib import Path
+from typing import Any, DefaultDict, Union
+
+from ..mapping import CheckpointingException, ShardedStateDict, StateDict
+from .async_utils import AsyncCallsQueue, AsyncRequest
+
+
+class StrategyAction(Enum):
+ """Specifies save vs load and sharded vs common action."""
+
+ LOAD_COMMON = 'load_common'
+ LOAD_SHARDED = 'load_sharded'
+ SAVE_COMMON = 'save_common'
+ SAVE_SHARDED = 'save_sharded'
+
+
+default_strategies: DefaultDict[str, dict[tuple, Any]] = defaultdict(dict)
+
+async_calls = AsyncCallsQueue()
+
+
+def get_default_strategy(action: StrategyAction, backend: str, version: int):
+ """Retrieves a default strategy for a given action, backend and version."""
+ try:
+ if backend == 'zarr':
+ error_hint = ' Please install `zarr` and `tensorstore<=0.1.45` packages'
+ from .tensorstore import register_default_tensorstore_strategies
+
+ register_default_tensorstore_strategies()
+ from .zarr import register_default_zarr_strategies
+
+ register_default_zarr_strategies()
+ elif backend == 'torch_dist':
+ error_hint = ' Please use PyTorch version >=2.1'
+ from .torch import register_default_torch_strategies
+
+ register_default_torch_strategies()
+ except ImportError as e:
+ raise CheckpointingException(
+ f'Cannot import a default strategy for: {(action.value, backend, version)}. '
+ f'Error: {e}. Hint: {error_hint}'
+ ) from e
+ try:
+ return default_strategies[action.value][(backend, version)]
+ except KeyError as e:
+ raise CheckpointingException(
+ f'Cannot find a default strategy for: {(action.value, backend, version)}'
+ ) from e
+
+
+def register_default_strategy(
+ action: StrategyAction,
+ backend: str,
+ version: int,
+ strategy: Union['SaveStrategyBase', 'LoadStrategyBase'],
+):
+ """Adds a given strategy to the registry of default strategies.
+
+ Args:
+ action (StrategyAction): specifies save/load and sharded/common
+ backend (str): backend that the strategy becomes a default for
+ version (int): version that the strategy becomes a default for
+ strategy (SaveStrategyBase, LoadStrategyBase): strategy to register
+ """
+ default_strategies[action.value][(backend, version)] = strategy
+
+
+class LoadStrategyBase(ABC):
+ """Base class for a load strategy. Requires implementing checks for compatibility with a
+ given checkpoint version."""
+
+ @abstractmethod
+ def check_backend_compatibility(self, loaded_backend):
+ """Verifies if this strategy is compatible with `loaded_backend`."""
+ raise NotImplementedError
+
+ @abstractmethod
+ def check_version_compatibility(self, loaded_version):
+ """Verifies if this strategy is compatible with `loaded_version`."""
+ raise NotImplementedError
+
+ @property
+ def can_handle_sharded_objects(self):
+ """Returns whether or not this strategy can handle loading ShardedObjects."""
+ return False
+
+
+class SaveStrategyBase(ABC):
+ """Base class for a save strategy. Requires defining a backend type and
+ version of the saved format."""
+
+ def __init__(self, backend: str, version: int):
+ self.backend = backend
+ self.version = version
+
+ @property
+ def can_handle_sharded_objects(self):
+ """Returns whether or not this strategy can handle saving ShardedObjects."""
+ return False
+
+ def __str__(self):
+ return f'{self.__class__.__name__}({self.backend}, {self.version})'
+
+
+class LoadCommonStrategy(LoadStrategyBase):
+ """Load strategy for common (non-sharded) objects"""
+
+ @abstractmethod
+ def load_common(self, checkpoint_dir: Path):
+ """Load common part of the checkpoint."""
+ raise NotImplementedError
+
+ @abstractmethod
+ def load_sharded_objects(
+ self, sharded_objects_state_dict: ShardedStateDict, checkpoint_dir: Path
+ ):
+ """Load sharded objects from the checkpoint."""
+ raise NotImplementedError
+
+ def load_sharded_metadata(self, checkpoint_dir: Path) -> ShardedStateDict:
+ """Load just the metadata from the checkpoint."""
+ if not self.can_handle_sharded_objects:
+ return {}
+ raise NotImplementedError
+
+
+class LoadShardedStrategy(LoadStrategyBase):
+ """Load strategy for sharded tensors"""
+
+ @abstractmethod
+ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
+ """Load the sharded part of the checkpoint."""
+ raise NotImplementedError
+
+ @abstractmethod
+ def load_tensors_metadata(self, checkpoint_dir: Path):
+ """Load tensors metadata from the checkpoint for ShardedTensors.
+
+ Returns a dictionary similar to a sharded state dict, but note that
+ the dictionary keys are simply ShardedTensor keys (contrary to the
+ actual sharded state dicts where keys correspond to state dict keys).
+
+ Dict values are ShardedTensors without any data and sharding (so, the
+ only useful information is tensors global shape and dtype).
+ """
+ raise NotImplementedError(
+ f'Loading only tensors metadata not implemented for {self.__class__.__name__}'
+ )
+
+ def load_sharded_metadata(self, checkpoint_dir: Path):
+ """Load sharded metadata from the checkpoint for ShardedTensors and ShardedObjects.
+
+ Returns a dictionary similar to a sharded state dict, but note that
+ the dictionary keys are simply sharded keys (contrary to the
+ actual sharded state dicts where keys correspond to state dict keys).
+
+ Dict values are ShardedTensors or ShardedObjects without any data and sharding.
+ """
+ if not self.can_handle_sharded_objects:
+ return self.load_tensors_metadata(checkpoint_dir)
+ raise NotImplementedError(
+ f'Loading only sharded metadata not implemented for {self.__class__.__name__}'
+ )
+
+
+class SaveCommonStrategy(SaveStrategyBase):
+ """Save strategy for common (non-sharded) objects"""
+
+ @abstractmethod
+ def save_common(self, common_state_dict: StateDict, checkpoint_dir: Path):
+ """Save common part of the state dict."""
+ raise NotImplementedError
+
+ def save_sharded_objects(
+ self, sharded_objects_state_dict: ShardedStateDict, checkpoint_dir: Path
+ ):
+ """Save sharded objects from the state dict."""
+ raise NotImplementedError
+
+
+class SaveShardedStrategy(SaveStrategyBase):
+ """Save strategy for sharded tensors"""
+
+ @abstractmethod
+ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
+ """Save the sharded part of the state dict."""
+ raise NotImplementedError
+
+
+class AsyncSaveShardedStrategy(SaveShardedStrategy):
+ """Save strategy suitable for async save."""
+
+ @abstractmethod
+ def async_save(
+ self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path
+ ) -> AsyncRequest:
+ """Perform preparation and return an AsyncRequest to the external caller.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict to save
+ checkpoint_dir (Path): checkpoint target directory
+
+ Returns:
+ AsyncRequest: represents the async save function and finalization function.
+ It is the caller responsibility to actually schedule the async save.
+ """
+ raise NotImplementedError
+
+ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
+ """Each async strategy can be trivially used as a sync strategy."""
+ async_request = self.async_save(sharded_state_dict, checkpoint_dir)
+ # multiprocessing routines may cause issue when called on parent process
+ # We keep this verbose call for now
+ global async_calls
+ async_calls.schedule_async_request(async_request)
+ async_calls.maybe_finalize_async_calls(blocking=True)
diff --git a/megatron/core/dist_checkpointing/strategies/common.py b/megatron/core/dist_checkpointing/strategies/common.py
new file mode 100644
index 0000000000..f2c87b4d60
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/common.py
@@ -0,0 +1,157 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+""" Common strategies. """
+
+import logging
+import os
+from pathlib import Path
+
+import torch
+
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict, StateDict
+from megatron.core.dist_checkpointing.strategies.base import (
+ SaveCommonStrategy,
+ StrategyAction,
+ register_default_strategy,
+)
+
+from ..dict_utils import dict_list_map_inplace, nested_values
+from ..mapping import CheckpointingException, ShardedObject, is_main_replica
+from ..strategies.base import LoadCommonStrategy
+
+COMMON_STATE_FNAME = 'common.pt'
+
+logger = logging.getLogger(__name__)
+
+
+def register_default_common_strategies():
+ """Register default common strategies."""
+ register_default_strategy(StrategyAction.LOAD_COMMON, 'torch', 1, TorchCommonLoadStrategy())
+ register_default_strategy(
+ StrategyAction.SAVE_COMMON, 'torch', 1, TorchCommonSaveStrategy('torch', 1)
+ )
+
+
+class TorchCommonSaveStrategy(SaveCommonStrategy):
+ """Common save strategy leveraging native torch save/load."""
+
+ def save_common(self, common_state_dict: StateDict, checkpoint_dir: Path):
+ """Save common part of the state dict."""
+ if torch.distributed.get_rank() == 0:
+ torch.save(common_state_dict, checkpoint_dir / COMMON_STATE_FNAME)
+
+ def save_sharded_objects(
+ self, sharded_objects_state_dict: ShardedStateDict, checkpoint_dir: Path
+ ):
+ """Save sharded objects from the state dict."""
+ for sh_obj in nested_values(sharded_objects_state_dict):
+ if is_main_replica(sh_obj.replica_id):
+ save_path = checkpoint_dir / f'{sh_obj.unique_key}.pt'
+ os.makedirs(save_path.parent, exist_ok=True)
+ torch.save(sh_obj.data, save_path)
+
+ def can_handle_sharded_objects(self):
+ """This strategy can handle ShardedObjects."""
+ return True
+
+
+class TorchCommonLoadStrategy(LoadCommonStrategy):
+ """Common load strategy leveraging native torch save/load."""
+
+ def load_common(self, checkpoint_dir: Path):
+ """Load common (non-sharded) objects state dict from the checkpoint.
+
+ Args:
+ checkpoint_dir (Path): checkpoint directory
+
+ Returns:
+ StateDict: state dict with non-sharded objects from the checkpoint
+ """
+ load_path = Path(checkpoint_dir) / COMMON_STATE_FNAME
+ try:
+ return torch.load(load_path, map_location='cpu')
+ except FileNotFoundError as e:
+ err_msg = f'Common file {load_path} does not exist'
+ ckpt_files = [f.name for f in checkpoint_dir.iterdir()]
+ logger.debug(f'{err_msg}. Checkpoint directory content: {ckpt_files}')
+ raise CheckpointingException(err_msg) from e
+
+ def load_sharded_objects(
+ self, sharded_objects_state_dict: ShardedStateDict, checkpoint_dir: Path
+ ):
+ """Replaces all ShardedObject from a given state dict with values loaded from the
+ checkpoint.
+
+ Args:
+ sharded_objects_state_dict (ShardedStateDict):
+ sharded state dict defining what objects should be loaded.
+ checkpoint_dir (Path): checkpoint directory
+
+ Returns:
+ None: sharded state dict is modified in place
+ """
+
+ def load_sharded_object(sh_obj: ShardedObject):
+ sh_obj.data = None
+ load_path = checkpoint_dir / f'{sh_obj.unique_key}.pt'
+ try:
+ loaded_obj = torch.load(load_path)
+ except FileNotFoundError as e:
+ # Backward compatible logic: previously the save format was incorrect
+ old_load_path = (checkpoint_dir / sh_obj.unique_key).with_suffix('.pt')
+ try:
+ loaded_obj = torch.load(old_load_path)
+ except FileNotFoundError:
+ err_msg = f'Object shard {load_path} not found'
+ obj_subdir = checkpoint_dir / sh_obj.key
+ if obj_subdir.exists():
+ obj_files = [f.name for f in obj_subdir.iterdir()]
+ logger.debug(
+ f'{err_msg}. Object {sh_obj.key} directory content: {obj_files}'
+ )
+ else:
+ ckpt_files = [f.name for f in checkpoint_dir.iterdir()]
+ logger.debug(
+ f'{err_msg}. Object {sh_obj.key} directory does not exist. Checkpoint'
+ f' directory content: {ckpt_files}'
+ )
+ raise CheckpointingException(err_msg) from e
+ return loaded_obj
+
+ return dict_list_map_inplace(load_sharded_object, sharded_objects_state_dict)
+
+ def load_sharded_metadata(self, checkpoint_dir: Path) -> ShardedStateDict:
+ sharded_metadata = {}
+ for subdir in checkpoint_dir.iterdir():
+ if not subdir.is_dir():
+ continue
+ shard_files = list(subdir.glob('shard_*.pt'))
+ if not shard_files:
+ continue
+ sh_objs = []
+ for shard_file in shard_files:
+ full_key = f'{subdir.name}/{shard_file.stem}'
+ sh_objs.append(ShardedObject.empty_from_unique_key(full_key))
+
+ # This is a backward-compatibility fix, where the last global shape is missing in the
+ # name
+ if sh_objs[0].global_shape[-1] < 0:
+ max_last_offset = max(map(lambda sh_obj: sh_obj.global_offset[-1], sh_objs))
+ for sh_obj in sh_objs:
+ sh_obj.global_shape = (*sh_obj.global_shape[:-1], max_last_offset + 1)
+
+ # Update the sharded state dict
+ for sh_obj in sh_objs:
+ sharded_metadata[sh_obj.unique_key] = sh_obj
+ return sharded_metadata
+
+ @property
+ def can_handle_sharded_objects(self):
+ """This strategy can handle ShardedObjects."""
+ return True
+
+ def check_backend_compatibility(self, loaded_version):
+ pass
+
+ def check_version_compatibility(self, loaded_version):
+ pass
diff --git a/megatron/core/dist_checkpointing/strategies/filesystem_async.py b/megatron/core/dist_checkpointing/strategies/filesystem_async.py
new file mode 100644
index 0000000000..9d0be4d6e7
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/filesystem_async.py
@@ -0,0 +1,386 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Storage writer for PyT Distributed format allowing asynchronous save. """
+import gc
+import logging
+import os
+import queue
+from contextlib import contextmanager
+from itertools import chain
+from pathlib import Path
+from time import time
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import psutil
+import torch
+from torch import multiprocessing as mp
+from torch.distributed.checkpoint import FileSystemWriter
+from torch.distributed.checkpoint.filesystem import DEFAULT_SUFFIX, _StoragePrefix, _write_item
+from torch.distributed.checkpoint.planner import SavePlan, SavePlanner, WriteItem, WriteItemType
+from torch.distributed.checkpoint.storage import WriteResult
+from torch.futures import Future
+
+logger = logging.getLogger(__name__)
+
+WriteBucket = Tuple[Path, str, Tuple[list, list]] # represents writes to a single file
+
+_results_queue = None
+
+
+def _get_write_results_queue():
+ global _results_queue
+ if _results_queue is None:
+ ctx = mp.get_context('spawn')
+ _results_queue = ctx.Manager().Queue()
+ return _results_queue
+
+
+@contextmanager
+def _disable_gc():
+ """Temporarily disables GC."""
+ gc_enabled = gc.isenabled()
+ try:
+ if gc_enabled:
+ gc.disable()
+ yield
+ finally:
+ if gc_enabled:
+ gc.enable()
+
+
+class FileSystemWriterAsync(FileSystemWriter):
+ """
+ Async-enabled implementation of FileSystemWriter using file IO.
+
+ This class doesn't spawn the async process itself, relies on the external async mechanism.
+
+ Flow:
+ 1. Call `write_data`
+ 2. Externally start async process with `get_save_function_and_args` function and args
+ 3. The async function to call is `writer_proxy_func` which calls
+ `write_preloaded_data` in multiple processes
+
+ After saving is finalized on all ranks:
+ 4. Call `super().finish` with the results gathered in `self.writer_result`
+
+ Note that step (3) above can also be called synchronously.
+
+ Currently, it's assumed that a separate writer is created for each ckpt save
+ (intermediate state is stored as writer attributes).
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ if not self.single_file_per_rank:
+ raise NotImplementedError(
+ 'single_file_per_rank flag not supported for FileSystemWriterAsync'
+ )
+
+ # Intermediate state between preparation and finalization
+ self.write_buckets: Optional[List[WriteBucket]] = None
+ self.results_queue: Optional[mp.Queue] = None
+
+ def prepare_write_data(self, plan: SavePlan, planner: SavePlanner) -> None:
+ """
+ First stage of async saving. Copy data to CPU and plan the local saving.
+
+ Args:
+ plan (SavePlan): save plan generated by the PyT Distributed compatible planner
+ planner (SavePlanner): save planner used to resolve the bytes and tensor data
+
+ Returns: None, but stores the save plan in `self.write_buckets`
+ """
+ storage_plan: _StoragePrefix = plan.storage_data
+ start = time()
+ logger.debug(f"thread_count: {self.thread_count}, time: {start}")
+ item_buckets = _split_by_size_and_type(self.thread_count, plan.items)
+ logger.debug(f"bucket_prep, time: {time() - start}")
+
+ start = time()
+ # move tensors from GPU to CPU before starting async writing
+ # We do D2H synchronously for now
+ file_count = 0
+
+ def gen_file():
+ nonlocal file_count
+ file_name = f"{storage_plan.prefix}{file_count}{DEFAULT_SUFFIX}"
+ file_count += 1
+ return file_name
+
+ # Prepare bytes / tensor data in each bucket, which will be assigned to each writer process
+ self.write_buckets = []
+ for bucket in item_buckets:
+ bytes_data = [
+ (item, planner.resolve_data(item))
+ for item in bucket
+ if item.type == WriteItemType.BYTE_IO
+ ]
+ tensor_data = [
+ (item, planner.resolve_data(item).detach().to("cpu", non_blocking=True))
+ for item in bucket
+ if item.type != WriteItemType.BYTE_IO
+ ]
+ if len(bytes_data) > 0 or len(tensor_data) > 0:
+ file_name = gen_file()
+ self.write_buckets.append(
+ (self.path / file_name, file_name, (bytes_data, tensor_data))
+ )
+
+ # Check if there is anything to write on this rank
+ if len(self.write_buckets) > 0:
+ assert len(self.write_buckets) <= self.thread_count, (
+ len(self.write_buckets),
+ self.thread_count,
+ )
+ self.results_queue = _get_write_results_queue()
+ else:
+ self.results_queue = None
+ end = time()
+ logger.debug(f"D2H and push, time: {end - start}")
+
+ def get_save_function_and_args(self) -> Tuple[Optional[Callable], Tuple]:
+ """
+ Get function that saves the data to storage along with its arguments.
+ Allows the external caller to apply the save function synchronously or asynchronously.
+
+ Returns: None (if there is nothing to write on this rank) or a tuple of:
+ - the function that saves the data
+ - arguments to that function
+ """
+ if not self.write_buckets:
+ return None, ()
+ return (self.write_preloaded_data_multiproc, (self.write_buckets, self.results_queue))
+
+ @staticmethod
+ @_disable_gc()
+ def write_preloaded_data_multiproc(
+ write_buckets: List[WriteBucket], global_results_queue: mp.Queue
+ ) -> None:
+ """
+ Performs saving data to storage with multiple processes.
+
+ Starts predefined number of processes and uses 2 queues to make sure the results
+ are complete:
+ - local_results_queue - to send the actual results
+ - count_queue - small queue to mark worker as completed
+
+ Using just one queue disallowed proper exception handling.
+
+ This method is meant to be run in a forked subprocess.
+ Triggering GC during execution leads to CUDA errors
+ (cleaning up tensors owned by the parent process).
+ To prevent this, we disable the GC explicitly for this function with _disable_gc.
+
+ Args:
+ write_buckets (List[WriteBucket]): write plan
+ global_results_queue (mp.Queue): mp.Queue to collect Dict[List[WriteResults]] (or an Exception)
+ from parallel write processes to the main training process
+ Returns: None
+ """
+ w_start = time()
+ write_results_or_exc: Union[dict, Exception] = dict()
+ ctx = mp.get_context('fork')
+ local_results_queue = ctx.Queue()
+ count_queue = ctx.JoinableQueue()
+ p_list = []
+ for i, write_bucket in enumerate(write_buckets):
+ try:
+ count_queue.put(i)
+ p_list.append(
+ ctx.Process(
+ target=FileSystemWriterAsync.write_preloaded_data,
+ args=(i, write_bucket, local_results_queue, count_queue, True),
+ )
+ )
+ except Exception as e:
+ err_msg = f'An error is caught while a proc {i} is created, error: {e}'
+ logger.error(err_msg)
+ write_results_or_exc = RuntimeError(err_msg)
+
+ if not isinstance(write_results_or_exc, Exception):
+ for p in p_list:
+ p.start()
+
+ logger.debug('FileSystemWriterAsync: collecting worker results...')
+
+ # To make sure all nodes are completed
+ count_queue.join()
+ # At this point, all workers completed, so the queue should have exactly `len(write_buckets)` items
+ for proc_idx in range(len(write_buckets)):
+ try:
+ local_proc_idx, local_results_or_exc = local_results_queue.get()
+ except queue.Empty:
+ write_results_or_exc = RuntimeError(
+ f'Unexpected empty `local_results_queue` (got only {proc_idx}/{len(write_buckets)} items)'
+ )
+ break
+ else:
+ if isinstance(local_results_or_exc, Exception):
+ err_msg = f"Local process {local_proc_idx} encountered an error: {local_results_or_exc}"
+ logger.error(err_msg)
+ write_results_or_exc = local_results_or_exc
+ break
+ else:
+ assert isinstance(local_results_or_exc, list), type(local_results_or_exc)
+ write_results_or_exc[local_proc_idx] = local_results_or_exc
+ p_list[local_proc_idx].join()
+
+ logger.debug('FileSystemWriterAsync: collected worker results successfully')
+
+ global_results_queue.put(write_results_or_exc)
+
+ w_end = time()
+ logger.debug(
+ f"{w_end}, rank: {torch.distributed.get_rank()}, write(sync,parallel): {w_end - w_start}"
+ )
+
+ @staticmethod
+ @_disable_gc()
+ def write_preloaded_data(
+ local_proc_idx: int,
+ write_bucket: WriteBucket,
+ results_queue: mp.SimpleQueue,
+ count_queue: mp.JoinableQueue,
+ use_fsync: bool,
+ ) -> None:
+ """
+ Performs actual data saving to storage.
+
+ Args:
+ local_proc_idx (int): index of a local process that performs writing
+ write_bucket (WriteBucket): data to write to storage
+ results_queue (mp.Queue): queue to return the write results to the proxy checkpoint process.
+ count_queue (mp.JoinableQueue): queue to marks worker task as completed
+ use_fsync (bool): if True, calls os.fsync at the end of saving
+
+ Returns: None, the write result are put into the `queue`
+ """
+ mem_before = _process_memory()
+
+ local_results = []
+ try:
+ file_name, storage_key, (bytes_data, tensor_data) = write_bucket
+ with open(file_name, "wb") as stream:
+ for write_item, data in bytes_data:
+ local_results.append(_write_item(stream, data, write_item, storage_key))
+
+ for write_item, tensor in tensor_data:
+ assert tensor.is_cpu
+ local_results.append(_write_item(stream, tensor, write_item, storage_key))
+
+ if use_fsync:
+ os.fsync(stream.fileno())
+ local_output = (local_proc_idx, local_results)
+ except Exception as e:
+ local_output = (local_proc_idx, e)
+
+ results_queue.put(local_output)
+ # Signal this process is done.
+ count_queue.get()
+ count_queue.task_done()
+
+ mem_after = _process_memory()
+ logger.debug(
+ f"{local_proc_idx} consumed: {mem_after - mem_before}, before: {mem_before}, after: {mem_after}"
+ )
+
+ def write_data(self, plan: SavePlan, planner: SavePlanner) -> Future[List[WriteResult]]:
+ raise NotImplementedError('write_data not implemented for FileSystemWriterAsync')
+
+ def retrieve_write_results(self) -> List[WriteResult]:
+ """
+ Turn the latest dict including write results from `self.results_queue` into a single results lists. Includes error check.
+
+ Returns (List[WriteResult]): the list of write results from all local processes performing the save.
+
+ """
+ assert self.write_buckets is not None
+
+ if self.results_queue is None:
+ write_results_or_exc = {}
+ else:
+ try:
+ write_results_or_exc = self.results_queue.get_nowait()
+ except queue.Empty:
+ raise RuntimeError(f'results_queue should not be empty')
+
+ if isinstance(write_results_or_exc, Exception):
+ raise RuntimeError(f'Worker failure: {write_results_or_exc}') from write_results_or_exc
+ write_results: dict = write_results_or_exc
+ if len(write_results) != len(self.write_buckets):
+ raise RuntimeError(
+ f'Incomplete worker results (expected {len(self.write_buckets)}, got {len(write_results)}.'
+ f' This probably indicates a worker failure.'
+ )
+ return list(chain.from_iterable(write_results.values()))
+
+
+def _split_by_size_and_type(bins: int, items: List[WriteItem]) -> List[List[WriteItem]]:
+ """
+ Splits write items according to item size into close to uniform bins.
+
+ Same as torch.distributed.checkpoint.filesystem._split_by_size_and_type,
+ but with a fixed _item_size function.
+
+ Args:
+ bins (int): numbers of bins to split to
+ items (List[WriteItem]): list of write items
+
+ Returns (List[List[WriteItem]]): write items split to bins
+ """
+ if bins == 1:
+ return [items]
+
+ bytes_items = [wi for wi in items if wi.type == WriteItemType.BYTE_IO]
+ tensor_items = [wi for wi in items if wi.type != WriteItemType.BYTE_IO]
+
+ buckets: List[List[WriteItem]] = [[] for _ in range(bins)]
+ bucket_sizes = [0 for _ in range(bins)]
+
+ tensor_items.sort(key=_item_size, reverse=True)
+
+ # Assign bytes with a simple round-robin
+ for i, item in enumerate(bytes_items):
+ buckets[i % bins].append(item)
+
+ # Then, assign tensors according to their sizes
+ for item in tensor_items:
+ # TODO replace with headq
+ idx = min(enumerate(bucket_sizes), key=lambda x: x[1])[0]
+ buckets[idx].append(item)
+ bucket_sizes[idx] += _item_size(item)
+
+ return buckets
+
+
+def _item_size(item: WriteItem) -> int:
+ """
+ Calculates size (in bytes) of a single write item.
+
+ Same as torch.distributed.checkpoint.filesystem._item_size,
+ but fixes computing chunk size (with item.tensor_data.chunk.sizes)
+
+ Args:
+ item (WriteItem): write item to compute the size of
+
+ Returns (int): size of an item in bytes
+ """
+ size = 1
+ assert item.tensor_data is not None
+ # can't use math.prod as PT needs to support older python
+ for s in item.tensor_data.chunk.sizes:
+ size *= s
+
+ dtype = item.tensor_data.properties.dtype
+ return size * torch._utils._element_size(dtype)
+
+
+def _process_memory() -> int:
+ """
+ Get memory used by current process.
+
+ Returns (int): memory used by current process
+ """
+ process = psutil.Process(os.getpid())
+ mem_info = process.memory_info()
+ return mem_info.rss
diff --git a/megatron/core/dist_checkpointing/strategies/fully_parallel.py b/megatron/core/dist_checkpointing/strategies/fully_parallel.py
new file mode 100644
index 0000000000..56523daf1c
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/fully_parallel.py
@@ -0,0 +1,439 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import logging
+from pathlib import Path
+from time import time
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+
+from megatron.core.dist_checkpointing import ShardedTensor
+from megatron.core.dist_checkpointing.core import CheckpointingException
+from megatron.core.dist_checkpointing.dict_utils import (
+ dict_list_map_inplace,
+ extract_matching_values,
+ merge,
+ nested_values,
+)
+from megatron.core.dist_checkpointing.exchange_utils import (
+ ShardDistribution,
+ determine_main_replica_uniform_distribution,
+ exchange_by_distribution,
+)
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict, StateDict, is_main_replica
+from megatron.core.dist_checkpointing.strategies.base import (
+ AsyncSaveShardedStrategy,
+ LoadShardedStrategy,
+ SaveShardedStrategy,
+)
+from megatron.core.dist_checkpointing.utils import _sharded_tensor_shard_id, _ShardId
+from megatron.core.dist_checkpointing.validation import (
+ determine_global_metadata,
+ validate_sharding_integrity,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FullyParallelSaveStrategyWrapper(AsyncSaveShardedStrategy):
+ """Wraps arbitrary strategy and distributes the save during `save`.
+
+ The save distribution happens without any *data* communication.
+ Only the *metadata* is exchanged and based on data replication on different
+ ranks, we try to distribute the save as uniformly as possible.
+
+ This wrapper assumes, that setting `replica_id` to 0 will make the
+ underlying strategy do the saving on current rank. All the other `replica_id`s
+ are set to 1.
+
+ Currently, the save distribution is realized with a greedy algorithm
+ described in `distribute_shards_to_ranks`.
+
+ Args:
+ strategy (SaveShardedStrategy): base strategy to wrap
+ parallelization_group (ProcessGroup, optional): process group to use for save
+ distribution. Note that this doesn't have to match exactly the
+ data distribution, but should cover the replication pattern
+ to maximize performance. Defaults to the whole world.
+ do_cache_distribution (bool, optional): whether to cache the save distribution
+ from previous calls. Should be set to True only if the state dict
+ structure between the calls is always the same. Defaults to True.
+ """
+
+ def __init__(
+ self,
+ strategy: SaveShardedStrategy,
+ parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
+ do_cache_distribution: bool = False,
+ ):
+ super().__init__(strategy.backend, strategy.version)
+ self.base_strategy = strategy
+ self.parallelization_group = parallelization_group
+ self.do_cache_distribution = do_cache_distribution
+
+ self.cached_distribution: Optional[ShardDistribution] = None
+
+ def async_save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
+ if not isinstance(self.base_strategy, AsyncSaveShardedStrategy):
+ raise CheckpointingException(
+ f'Cannot apply async_save to non-async base strategy {self.base_strategy}'
+ )
+ self.apply_saving_parallelization(sharded_state_dict)
+ return self.base_strategy.async_save(sharded_state_dict, checkpoint_dir)
+
+ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
+ self.apply_saving_parallelization(sharded_state_dict)
+ return self.base_strategy.save(sharded_state_dict, checkpoint_dir)
+
+ def apply_saving_parallelization(self, sharded_state_dict: ShardedStateDict) -> None:
+ """Distributes the save across ranks by exchanging metadata.
+
+ Exchanges metadata from the state dict and computes the uniform
+ (as close as possible) distribution of saves among the ranks.
+
+ If `self.do_cache_distribution` is True, caches the distribution between
+ the calls and subsequent distributions happen without any inter-rank
+ communication.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): state dict to distribute the saving
+
+ Returns: None
+ """
+ start = time()
+ if self.do_cache_distribution and self.cached_distribution is not None:
+ logger.debug(f'Apply *cached* save parallelization')
+ precomputed_distribution = self.cached_distribution
+ else:
+ logger.debug(f'Apply save parallelization')
+ precomputed_distribution = determine_main_replica_uniform_distribution(
+ sharded_state_dict, self.parallelization_group
+ )
+
+ distribute_main_replicas_with_precomputed_distribution(
+ sharded_state_dict, self.parallelization_group, precomputed_distribution
+ )
+ if self.cached_distribution is None:
+ # First time applying the parallelization
+ validate_sharding_integrity(determine_global_metadata(sharded_state_dict)[1])
+ if self.do_cache_distribution:
+ self.cached_distribution = precomputed_distribution
+ end = time()
+ logger.debug(f"parallel save sharding, time: {end - start}")
+
+ @property
+ def can_handle_sharded_objects(self):
+ return self.base_strategy.can_handle_sharded_objects
+
+
+class FullyParallelLoadStrategyWrapper(LoadShardedStrategy):
+ """Wraps arbitrary load strategy and distributes the load during `load`.
+
+ See `load` method docs for details.
+
+ Args:
+ strategy (LoadShardedStrategy): base strategy to wrap
+ parallelization_group (ProcessGroup, optional): process group to use for load
+ distribution. Note that this doesn't have to match exactly the
+ data distribution, but should cover the replication pattern
+ to maximize performance. Defaults to the whole world.
+ In most cases, it's recommended to set it to the DP group.
+ do_cache_distribution (bool, optional): whether to cache the load distribution
+ from previous calls. Should be set to True only if the state dict
+ structure between the calls is always the same. Defaults to False,
+ since the loading in general happens only once during training.
+ Note that the load distribution *cannot* be reused as a save distribution,
+ because save/load is not fully symmetrical.
+ exchange_algo (str): algorithm to use for exchanging the data.
+ Options:
+ - broadcast - each rank broadcasts individual tensors to others
+ - gather_object (default) - ranks all_gather_object the whole loaded state dicts
+ - gather_rounds (default) - ranks all gather individual tensors in rounds
+ See method docs for more details.
+ """
+
+ def __init__(
+ self,
+ strategy: LoadShardedStrategy,
+ parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
+ do_cache_distribution: bool = False,
+ exchange_algo: str = 'broadcast',
+ ):
+ super().__init__()
+ self.base_strategy = strategy
+ if parallelization_group is None:
+ parallelization_group = (
+ dist.GroupMember.WORLD
+ ) # explicit group needed for torch.distributed.get_global_rank call
+ self.parallelization_group = parallelization_group
+ self.do_cache_distribution = do_cache_distribution
+ self.exchange_algo = exchange_algo
+
+ self.cached_distribution: Optional[ShardDistribution] = None
+
+ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> StateDict:
+ """Distributes the load and calls underlying strategy only for parts of the state dict.
+
+ Steps:
+ 1. Load metadata is exchanged between the ranks in the parallelization group.
+ 2. Each rank deterministically plans the load for the whole workload
+ so that the loads are as uniform as possible.
+ 3. Each ranks loads its planned shard of the checkpoint.
+ 4. All ranks exchange the loaded shards.
+
+ Internode communication is involved in steps (1) (with metadata)
+ and (4) (with actual data). Storage interaction is involved in step (3).
+
+ Currently, the load distribution (step 2) is realized with a greedy algorithm
+ described in `distribute_shards_to_ranks` (same as for saving distribution).
+
+ Currently, the shards are all gathered between all ranks in the parallelization
+ group. This might not be optimal (some ranks do not need all tensors),
+ but it's a reasonable approximation for an optimal exchange in most scenarios.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict to load
+ checkpoint_dir (Path): checkpoint directory to load from
+
+ Returns:
+ StateDict: loaded state dict. The state dict should be equivalent to
+ a state dict that would be loaded with the underlying strategy
+ without this wrapper.
+ """
+ if torch.distributed.get_world_size(self.parallelization_group) <= 1:
+ return self.base_strategy.load(sharded_state_dict, checkpoint_dir)
+
+ # Step 1 and 2: exchange load metadata and distribute the load
+ start = time()
+ precomputed_distribution = self.apply_loading_parallelization(sharded_state_dict)
+ assert (
+ precomputed_distribution is not None
+ ), 'Expecting non-trivial distribution for non-trivial parallelization group'
+ end = time()
+ logger.debug(f'self.apply_loading_parallelization took {end - start}s')
+ start = end
+
+ # Step 3: load part of the checkpoint.
+ # Load only sharded objects first. ShardedTensors will be loaded separately
+ # so that we can keep track of sharded tensors loaded by this rank
+ (sharded_tensors, sharded_state_dict, to_load_shards, unloaded_shards) = (
+ self._defer_loading_sharded_tensors(sharded_state_dict)
+ )
+ loaded_state_dict = self.base_strategy.load(sharded_state_dict, checkpoint_dir)
+
+ end = time()
+ logger.debug(f'Base load of ShardedObjects took {end - start}s')
+ start = end
+
+ # Load sharded tensors separately
+ loaded_tensors = self.base_strategy.load(to_load_shards, checkpoint_dir)
+
+ end = time()
+ logger.debug(f'Base load of ShardedTensors took {end - start}s')
+ start = end
+
+ # Step 4: exchange data between ranks
+ logger.debug(f'Applying parallel load with algo {self.exchange_algo}')
+ all_loaded_tensors = exchange_by_distribution(
+ loaded_tensors,
+ unloaded_shards,
+ precomputed_distribution,
+ self.parallelization_group,
+ self.exchange_algo,
+ )
+ if not set(unloaded_shards.keys()).issubset(all_loaded_tensors.keys()):
+ missing_shards = set(unloaded_shards.keys()) - all_loaded_tensors.keys()
+ raise CheckpointingException(
+ f'Missing shards after fully parallel loading: {missing_shards}'
+ )
+
+ sync_start = time()
+ torch.cuda.synchronize()
+ end = time()
+ logger.debug(f'torch.cuda.synchronize took {end - sync_start}s')
+ logger.debug(f'self.exchange_loaded_tensors took {end - start}s')
+
+ self.fill_in_deferred_sharded_tensors(sharded_tensors, all_loaded_tensors)
+ merge(loaded_state_dict, sharded_tensors)
+ return loaded_state_dict
+
+ def _defer_loading_sharded_tensors(
+ self, sharded_state_dict: ShardedStateDict
+ ) -> Tuple[
+ ShardedStateDict,
+ ShardedStateDict,
+ Dict[_ShardId, ShardedTensor],
+ Dict[_ShardId, ShardedTensor],
+ ]:
+ """Divides state dict into parts loaded by this vs other ranks.
+
+ ShardedTensors with main replica_id will be loaded by this rank,
+ others will be received by other ranks (after loading from storage).
+
+ Args:
+ sharded_state_dict (ShardedStateDict): state dict with ShardedTensor
+ that will be divided.
+
+ Returns: a tuple of:
+ - ShardedStateDict: sub-state dict only with ShardedTensors
+ - ShardedStateDict: sub-state dict with non-ShardedTensors
+ - Dict[_ShardId, ShardedTensor]: ShardedTensor are uniquely identified
+ by shard ids. This is a mapping from shard id to a corresponding
+ ShardedTensor for tensors loaded by *this* rank
+ - Dict[_ShardId, ShardedTensor]: mapping from shard id to a corresponding
+ ShardedTensor for tensors loaded by *other* ranks
+ """
+ to_load_shards = {}
+ unloaded_shards = {}
+
+ sharded_tensors, sharded_state_dict = extract_matching_values(
+ sharded_state_dict, lambda v: isinstance(v, ShardedTensor)
+ )
+
+ def wrap_non_main_replicas(x):
+ if isinstance(x, ShardedTensor):
+ # Assign shard to be loaded or not
+ if is_main_replica(x.replica_id):
+ to_load_shards[_sharded_tensor_shard_id(x)] = x
+ else:
+ unloaded_shards[_sharded_tensor_shard_id(x)] = x
+ return x
+
+ dict_list_map_inplace(wrap_non_main_replicas, sharded_tensors)
+ return sharded_tensors, sharded_state_dict, to_load_shards, unloaded_shards
+
+ def apply_loading_parallelization(
+ self, sharded_state_dict: ShardedStateDict
+ ) -> Optional[ShardDistribution]:
+ """Distributes the load across ranks by exchanging metadata.
+
+ Exchanges metadata from the state dict and computes the uniform
+ (as close as possible) distribution of loads among the ranks.
+ Marks ShardedTensors to be loaded by the current rank with replica_id 0
+ (and others with non 0 values).
+
+ If `self.do_cache_distribution` is True, caches the distribution between
+ the calls and subsequent distributions happen without any inter-rank
+ communication.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): state dict to distribute the loading
+
+ Returns:
+ ShardDistribution (optional): the computed loading distribution
+ """
+ if self.do_cache_distribution and self.cached_distribution is not None:
+ logger.debug(f'Apply *cached* load parallelization')
+ precomputed_distribution = self.cached_distribution
+ else:
+ logger.debug(f'Apply load parallelization')
+ precomputed_distribution = determine_main_replica_uniform_distribution(
+ sharded_state_dict, self.parallelization_group, True
+ )
+
+ distribute_main_replicas_with_precomputed_distribution(
+ sharded_state_dict, self.parallelization_group, precomputed_distribution
+ )
+ if self.do_cache_distribution:
+ self.cached_distribution = precomputed_distribution
+
+ return precomputed_distribution
+
+ def fill_in_deferred_sharded_tensors(
+ self, sharded_state_dict: ShardedStateDict, loaded_tensors: Dict[_ShardId, torch.Tensor]
+ ) -> None:
+ """Fill in tensors not loaded by current rank with tensors from `loaded_tensors` map.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict to fill in.
+ ShardedTensors are completely replaced with corresponding torch.Tensors.
+ loaded_tensors (Dict[_ShardId, torch.Tensor]): dict allowing to map
+ ShardedTensor from the sharded_state_dict to loaded tensors.
+
+ Returns:
+
+ """
+
+ def fill_in_sharded_tensor(x):
+ if isinstance(x, ShardedTensor):
+ try:
+ x = loaded_tensors[_sharded_tensor_shard_id(x)]
+ except KeyError as e:
+ raise CheckpointingException(
+ f'Missing loaded tensor shard: {_sharded_tensor_shard_id(x)}'
+ ) from e
+
+ return x
+
+ dict_list_map_inplace(fill_in_sharded_tensor, sharded_state_dict)
+
+ @property
+ def can_handle_sharded_objects(self):
+ return self.base_strategy.can_handle_sharded_objects
+
+ def load_tensors_metadata(self, checkpoint_dir: Path):
+ return self.base_strategy.load_tensors_metadata(checkpoint_dir)
+
+ def load_sharded_metadata(self, checkpoint_dir: Path):
+ return self.base_strategy.load_sharded_metadata(checkpoint_dir)
+
+ def check_backend_compatibility(self, loaded_version):
+ return self.base_strategy.check_backend_compatibility(loaded_version)
+
+ def check_version_compatibility(self, loaded_version):
+ return self.base_strategy.check_version_compatibility(loaded_version)
+
+
+def distribute_main_replicas_with_precomputed_distribution(
+ sharded_state_dict: ShardedStateDict,
+ parallelization_group: torch.distributed.ProcessGroup,
+ precomputed_distribution: Optional[ShardDistribution],
+):
+ """Applies the save distribution computed with `determine_main_replica_uniform_distribution`.
+
+ Based on rank assignment, sets replica ids of the shards saved by current rank to 0
+ and all the other replica ids to 1.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): state dict to apply the save distribution to
+ parallelization_group (ProcessGroup): distribution will be applied within this
+ process group. Must match with the process group passed to
+ `determine_main_replica_uniform_distribution`.
+ precomputed_distribution (ShardDistribution): distribution computed with
+ `determine_main_replica_uniform_distribution`
+
+ Returns: None
+
+ Example replica ids of tensors A, B, C before distribution:
+ rank0: A: (0, 0, 0), B: (0, 0, 0), C: (0, 0, 0)
+ rank1: A: (0, 0, 1), B: (0, 0, 1), C: (0, 0, 1)
+ rank2: A: (0, 0, 2), B: (0, 0, 2), C: (0, 0, 2)
+
+ Replicas after distribution for the example above:
+ rank0: A: 0, B: 1, C: 1
+ rank1: A: 1, B: 0, C: 1
+ rank2: A: 1, B: 1, C: 0
+ """
+ if torch.distributed.get_world_size(group=parallelization_group) <= 1:
+ return
+ if precomputed_distribution is None:
+ raise ValueError(
+ 'precomputed_distribution must be not None for non-trivial parallelization group'
+ )
+
+ local_shards = list(
+ sh_base
+ for sh_base in nested_values(sharded_state_dict)
+ if isinstance(sh_base, ShardedTensor)
+ )
+
+ rank_within_dp_group = torch.distributed.get_rank(parallelization_group)
+ for sh_ten in local_shards:
+ shard_id = _sharded_tensor_shard_id(sh_ten)
+ if (
+ shard_id in precomputed_distribution.shards_in_this_group
+ and rank_within_dp_group == precomputed_distribution.main_rank_for_shard[shard_id]
+ ):
+ sh_ten.replica_id = 0
+ else:
+ sh_ten.replica_id = 1
diff --git a/megatron/core/dist_checkpointing/strategies/resharding.py b/megatron/core/dist_checkpointing/strategies/resharding.py
new file mode 100644
index 0000000000..c1c2bcec84
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/resharding.py
@@ -0,0 +1,315 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+""" Performant resharding of flattened tensors.
+
+Tensors that are first sharded (e.g. across TP) and then flattened cause
+very irregular access patterns during loading. The idea for performant save/load
+is to store tensors with global shape [X, Y, Z] and local shape [x, y, z]
+as tensors with global shape [X // x, Y // y, Z // z, x * y * z] and
+local shape [1, 1, 1, x * y * z]. This allows parallel save of tensors along the
+last (flattened) dimension. During loading, some additional resharding is needed.
+"""
+import logging
+import math
+from dataclasses import dataclass
+from itertools import product
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from torch.distributed.checkpoint import ChunkStorageMetadata
+from torch.distributed.checkpoint.resharding import _shards_get_overlap_region_wrt_saved_tensor
+
+from megatron.core.dist_checkpointing import ShardedTensor
+from megatron.core.dist_checkpointing.core import CheckpointingException
+from megatron.core.dist_checkpointing.dict_utils import (
+ dict_list_map_inplace,
+ extract_matching_values,
+)
+from megatron.core.dist_checkpointing.mapping import (
+ ReplicaId,
+ ShardedStateDict,
+ ShardedTensorFactory,
+ StateDict,
+ apply_factories,
+ apply_factory_merges,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TensorReformulationMetadata:
+ """Metadata needed to restore the original tensor shape.
+
+ Args:
+ ckpt_orig_global_shape (Tuple[int, ...]): original global shape of the tensor
+ saved in the checkpoint. This is the global shape of the application,
+ further reformulated into `ckpt_reform_global_shape` while saving.
+ ckpt_reform_global_shape (Tuple[int, ...]): reformulated global shape of the tensor
+ saved in the checkpoint. This is the actual saved shape.
+ """
+
+ ckpt_orig_global_shape: Tuple[int, ...]
+ ckpt_reform_global_shape: Tuple[int, ...]
+
+ def __post_init__(self):
+ assert self.ckpt_orig_global_shape
+
+
+def nd_flattened_tensor_reformulated_global_shape(sh_ten: ShardedTensor) -> Tuple[int, ...]:
+ """Reformulated global shape of the flattened N-D ShardedTensor.
+
+ N-D tensor global shape [X, Y, Z] and local shape [x, y, z]
+ is reformulated into global shape [X // x, Y // y, Z // z, x * y * z] and
+ local shape [1, 1, 1, x * y * z], to allow parallel save of tensors along the
+ last (flattened) dimension.
+
+ Args:
+ sh_ten (ShardedTensor): flattened N-D ShardedTensor (N > 1)
+
+ Returns:
+ Tuple[int, ...]: reformulated tensor shape
+ """
+ assert is_nd_flattened_tensor(sh_ten), sh_ten
+ return sh_ten.axis_fragmentations + (int(np.prod(sh_ten.local_shape)),)
+
+
+def is_nd_flattened_tensor(sh_ten: Any) -> bool:
+ """Checks if ShardedTensor is flattened and more than 1-dimensional
+
+ Args:
+ sh_ten (Any): any object
+
+ Returns:
+ bool: whether the given object is a flattened ShardedTensor and is N-dimensional (N > 1)
+ """
+ return (
+ isinstance(sh_ten, ShardedTensor)
+ and sh_ten.flattened_range is not None
+ and len(sh_ten.global_shape) > 1
+ )
+
+
+# information needed to restore. With current implementation, this is a nested state dict
+# with ShardedTensorFactories which is basically a ShardedStateDict type
+ReformulationRestoreMetadata = ShardedStateDict
+
+
+def apply_nd_flattened_tensors_reformulation(
+ sharded_state_dict: ShardedStateDict,
+ reformulation_metadata: Dict[str, TensorReformulationMetadata],
+) -> Tuple[ShardedStateDict, ReformulationRestoreMetadata]:
+ """Applies N-D reformulation to a given sharded state dict.
+
+ After applying the method and loading the reformulated state dict,
+ the `restore_nd_flattened_tensors_formulation` needs to be applied.
+
+ Current implementation uses ShardedTensorFactories for convenience of
+ restoring the original structure, but it's just an implementation detail.
+ Turns N-D ShardedTensors into factories and immediately applies them,
+ keeping the data needed to restore the original structure.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict potentially
+ with tensors to reformulate.
+ reformulation_metadata (Dict[str, TensorReformulationMetadata]): dict
+ containing all metadata needed for reformulating tensors in `sharded_state_dict`.
+ for each N-D flattened tensor `sh_ten` in `sharded_state_dict` there must be an
+ entry with `sh_ten.key`.
+
+ Returns:
+ tuple:
+ ShardedStateDict - reformulated sharded state dict
+ ReformulationRestoreMetadata - data needed to restore the original formulation
+ with `restore_nd_flattened_tensors_formulation`
+ """
+
+ def maybe_reformulate_nd_flattened_tensor(sh_ten: Any):
+ if not isinstance(sh_ten, ShardedTensor) or not is_nd_flattened_tensor(sh_ten):
+ return sh_ten
+ # N-D flattened ShardedTensor
+ try:
+ sh_ten_reformulation_metadata = reformulation_metadata[sh_ten.key]
+ except KeyError as e:
+ raise CheckpointingException(
+ f'Missing reformulation metadata for tensor {sh_ten}. Existing keys: {reformulation_metadata.keys()}'
+ ) from e
+
+ ckpt_actual_saved_shape = sh_ten_reformulation_metadata.ckpt_reform_global_shape
+ app_actual_load_shape = nd_flattened_tensor_reformulated_global_shape(sh_ten)
+ if ckpt_actual_saved_shape == app_actual_load_shape:
+ # Same shape - no need to reshard
+ return sh_ten
+
+ return reformulate_single_nd_flattened_tensor(sh_ten, sh_ten_reformulation_metadata)
+
+ # Turn N-D tensors into factories and immediately apply them
+ dict_list_map_inplace(maybe_reformulate_nd_flattened_tensor, sharded_state_dict)
+ sh_ten_factories, _ = extract_matching_values(
+ sharded_state_dict,
+ lambda x: isinstance(x, ShardedTensorFactory),
+ return_lists_as_dicts=True,
+ )
+ apply_factories(sharded_state_dict)
+
+ # Unlink `data` pointers to free memory
+ def unlink_data(x):
+ x.data = None
+ return x
+
+ dict_list_map_inplace(unlink_data, sh_ten_factories)
+ return sharded_state_dict, sh_ten_factories
+
+
+def restore_nd_flattened_tensors_formulation(
+ state_dict: StateDict, formulation_restore_metadata: ReformulationRestoreMetadata
+) -> StateDict:
+ """Restores the original state dict from a reformulated form.
+
+ Inverse of `apply_nd_flattened_tensors_reformulation`.
+
+ Args:
+ state_dict (StateDict): state dict obtained by loading a reformulated
+ sharded state dict.
+ formulation_restore_metadata (ReformulationRestoreMetadata): metadata returned by
+ `apply_nd_flattened_tensors_reformulation` function
+
+ Returns:
+ StateDict: state dict with the original tensors formulation restored
+ """
+ return apply_factory_merges(state_dict, formulation_restore_metadata)
+
+
+def reformulate_single_nd_flattened_tensor(
+ sh_ten: ShardedTensor, reformulation_metadata: TensorReformulationMetadata
+) -> Union[Any, ShardedTensorFactory]:
+ """Reformulates shapes of a single N-D flattened ShardedTensor.
+
+ We need to define a pair of transformations:
+ - turn N-D ShardedTensor with original formulation into multiple reformulated ShardedTensors
+ - merge multiple reformulated loaded torch.Tensors into a single original tensor
+ Current implementation uses ShardedTensorFactories as a convenient mechanism
+ for specifying and keeping track of those transformations.
+
+ Args:
+ sh_ten (ShardedTensor): sharded tensor to reformulate.
+ reformulation_metadata (TensorReformulationMetadata): metadata needed to
+ perform the reformulation
+
+ Returns:
+ ShardedTensorFactory: factory that keeps information how to reformulate
+ (build) the ShardedTensor and then restore original formulation (merge)
+ after loading.
+ """
+ rmd = reformulation_metadata
+ # Data won't be needed - remove unnecessary tensor references
+ sh_ten = sh_ten.without_data()
+
+ # Based on reformulation_metadata, determine other tensor shapes and metadata
+ ckpt_axis_fragmentation = rmd.ckpt_reform_global_shape[:-1]
+ for sh, fragm in zip(rmd.ckpt_orig_global_shape, ckpt_axis_fragmentation):
+ assert sh % fragm == 0, (sh_ten, rmd.ckpt_reform_global_shape)
+ ckpt_local_shape_with_prepended_axis = tuple(
+ sh // fragm for sh, fragm in zip(rmd.ckpt_orig_global_shape, ckpt_axis_fragmentation)
+ )
+ assert (
+ ckpt_local_shape_with_prepended_axis[: sh_ten.prepend_axis_num]
+ == (1,) * sh_ten.prepend_axis_num
+ ), (ckpt_local_shape_with_prepended_axis, sh_ten)
+ ckpt_local_shape = ckpt_local_shape_with_prepended_axis[sh_ten.prepend_axis_num :]
+
+ # Iterate over reformulated shapes needed by the application and from checkpoint,
+ # and generate new ShardedTensors that match the checkpoint sharding.
+ overlap_dim_offsets = []
+ assert len(ckpt_axis_fragmentation) == len(sh_ten.axis_fragmentations), (
+ ckpt_axis_fragmentation,
+ sh_ten,
+ )
+ for dim, (app_chunk_dim_offset, ckpt_fragm, app_fragm) in enumerate(
+ zip(
+ sh_ten.local_chunk_offset_in_global(),
+ ckpt_axis_fragmentation,
+ sh_ten.axis_fragmentations,
+ )
+ ):
+ # without `int`, it's an exact offset of the app shard expressed in ckpt_local_shape units
+ first_overlap_dim_offset = int(ckpt_fragm / app_fragm * app_chunk_dim_offset)
+ # `math.ceil` argument is an exact offset of the app next shard expressed in ckpt_local_shape units
+ next_overlap_dim_offset = math.ceil(ckpt_fragm / app_fragm * (app_chunk_dim_offset + 1))
+ overlap_dim_offsets.append(range(first_overlap_dim_offset, next_overlap_dim_offset))
+
+ logger.debug(
+ f'Generated the following number of overlap shards for each dimension: {list(map(len, overlap_dim_offsets))}'
+ f' for fragmentation ckpt {ckpt_axis_fragmentation} vs app {sh_ten.axis_fragmentations} and chunk offset {sh_ten.local_chunk_offset_in_global()}'
+ )
+ reformulated_sh_tens = {}
+ for chunk_offset in product(*overlap_dim_offsets):
+ global_offset = tuple(
+ chunk_off * chunk_shape
+ for chunk_off, chunk_shape in zip(chunk_offset, ckpt_local_shape_with_prepended_axis)
+ )
+ reformulated_sh_tens[(global_offset, ckpt_local_shape)] = ShardedTensor(
+ sh_ten.key,
+ None,
+ sh_ten.dtype,
+ ckpt_local_shape,
+ rmd.ckpt_orig_global_shape,
+ global_offset,
+ ckpt_axis_fragmentation,
+ sh_ten.replica_id,
+ sh_ten.prepend_axis_num,
+ sh_ten.allow_shape_mismatch,
+ flattened_range=slice(0, rmd.ckpt_reform_global_shape[-1]), # whole ckpt shard
+ )
+
+ # Now, we have to define the transformations from application sharding
+ # to checkpoint sharding.
+
+ @torch.no_grad()
+ def sh_ten_build_fn(*args, **kwargs):
+ # Here we simply return the precomputed tensors.
+ return reformulated_sh_tens
+
+ @torch.no_grad()
+ def sh_ten_merge_fn(sub_state_dict):
+ # This is the non-flattened local tensor with original formulation
+ # that we are going to fill with shards loaded from the checkpoint.
+ app_non_flat_ten = torch.empty(
+ sh_ten.local_shape,
+ dtype=sh_ten.dtype,
+ device=sh_ten.data.device if sh_ten.data is not None else None,
+ )
+
+ assert len(sub_state_dict) > 0
+ for (ckpt_global_offset, ckpt_local_shape), ckpt_ten in sub_state_dict.items():
+ # For each ckpt shard, we fill the appropriate application shard part
+ dest_ten = app_non_flat_ten
+ src_ten = ckpt_ten.view(ckpt_local_shape)
+ # We don't need narrowing over `prepend_axis_num` axes so we take the [sh_ten.prepend_axis_num:] offsets slice
+ for (
+ dim,
+ offset_for_saved_tensor,
+ offset_for_current_tensor,
+ length,
+ ) in _shards_get_overlap_region_wrt_saved_tensor(
+ saved_shard=ChunkStorageMetadata(
+ ckpt_global_offset[sh_ten.prepend_axis_num :], ckpt_local_shape
+ ),
+ current_shard=ChunkStorageMetadata(
+ sh_ten.global_offset[sh_ten.prepend_axis_num :], sh_ten.local_shape
+ ),
+ ):
+ src_ten = src_ten.narrow(dim, offset_for_saved_tensor, length)
+ dest_ten = dest_ten.narrow(dim, offset_for_current_tensor, length)
+ dest_ten.copy_(src_ten)
+ return app_non_flat_ten.flatten()[sh_ten.flattened_range]
+
+ return ShardedTensorFactory(
+ sh_ten.key,
+ sh_ten.data,
+ sh_ten_build_fn,
+ sh_ten_merge_fn,
+ sh_ten.replica_id,
+ sh_ten.flattened_range,
+ )
diff --git a/megatron/core/dist_checkpointing/strategies/state_dict_saver.py b/megatron/core/dist_checkpointing/strategies/state_dict_saver.py
new file mode 100644
index 0000000000..7b35209f21
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/state_dict_saver.py
@@ -0,0 +1,162 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+""" State dict saver for PyT Distributed format allowing asynchronous save. """
+
+from logging import getLogger
+from time import time
+from typing import TYPE_CHECKING, Optional, Tuple, cast
+
+import torch
+import torch.distributed as dist
+from torch.distributed.checkpoint import CheckpointException
+from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
+from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE, Metadata
+from torch.distributed.checkpoint.planner import SavePlan, SavePlanner
+from torch.distributed.checkpoint.utils import _DistWrapper, _get_failure_dict
+
+if TYPE_CHECKING:
+ from .filesystem_async import FileSystemWriterAsync
+
+
+logger = getLogger(__name__)
+
+
+def save_state_dict_async_plan(
+ state_dict: STATE_DICT_TYPE,
+ storage_writer: 'FileSystemWriterAsync',
+ process_group: Optional[dist.ProcessGroup] = None,
+ coordinator_rank: int = 0,
+ planner: Optional[SavePlanner] = None,
+ cached_ckpt_structure: Optional[Tuple[SavePlan, SavePlan, bool]] = None,
+) -> Tuple[Tuple['FileSystemWriterAsync', Metadata, _DistWrapper], SavePlan, bool]:
+ """
+ First stage of saving a state dict to storage.
+
+ This is an async adjustment of torch.distributed.checkpoint.state_dict_saver.
+ In order to support async save, saving should be split into three parts:
+ 1. Planning
+ 2. Actual saving
+ 3. Finalization
+
+ Out of these, step (2) *must* happen asynchronously.
+ The first step is realized with this function.
+
+ The planning part consists of several steps, described here:
+ https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner
+
+ Args:
+ state_dict (STATE_DICT_TYPE): state dict to save
+ storage_writer (FileSystemWriterAsync): in current version only an instance of
+ FileSystemWriterAsync
+ process_group (dist.ProcessGroup, optional): process group used for save planning
+ coordinator_rank (int, optional): coordinator rank for planning. Defaults to 0.
+ planner (SavePlanner, optional): save planner for torch.distributed.checkpoint format
+ cached_ckpt_structure (Tuple[SavePlan, SavePlan, bool], Optional):
+ Each object of this tuple will be used in the order as following
+ cached_central_plan (SavePlan): a globally coordinated save plan
+ cached in the previous iteration
+ cached_local_plan (SavePlan): a local plan
+ cached in the previous iteration
+ validated_cache_reuse (bool): boolean value to tell global_metadata and planning dict
+ is consistent over iterations
+
+ Returns: Tuple of:
+ - storage writer (the one passed as input)
+ - metadata from planning
+ - distributed wrapper used for planning
+ The return value of this function should be passed as an input to
+ `save_state_dict_async_finalize` and cached_plan to skip `reduce_scatter` at planning.
+ """
+ cached_central_plan, cached_local_plan, validated_cache_reuse = (None, None, False)
+ if cached_ckpt_structure:
+ cached_central_plan, cached_local_plan, validated_cache_reuse = cached_ckpt_structure
+
+ rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+ dist_wrapper = _DistWrapper(process_group, True, coordinator_rank)
+ if planner is None:
+ planner = DefaultSavePlanner()
+ assert planner is not None
+
+ global_metadata = None
+ logger.debug(f"rank: {rank}, starting state dict save")
+ local_plan = cached_local_plan
+
+ def local_step():
+ nonlocal local_plan
+ assert planner is not None
+ # PyTorch 2.4 introduced additional `metadata` argument,
+ # we have to reference `is_coordinator` args by name
+ planner.set_up_planner(state_dict, is_coordinator=dist_wrapper.is_coordinator)
+ storage_writer.set_up_storage_writer(dist_wrapper.is_coordinator)
+ if not validated_cache_reuse and local_plan is None:
+ local_plan = planner.create_local_plan()
+ local_plan = storage_writer.prepare_local_plan(local_plan)
+ return local_plan
+
+ def global_step(all_local_plans):
+ nonlocal global_metadata
+ assert planner is not None
+ all_local_plans, global_metadata = planner.create_global_plan(all_local_plans)
+ all_local_plans = storage_writer.prepare_global_plan(all_local_plans)
+ return all_local_plans
+
+ # Execute local and global planning
+ start_plan = time()
+ if validated_cache_reuse and cached_central_plan:
+ logger.debug(f"rank: {rank}, Passed cache reusable")
+ local_step()
+ central_plan = cached_central_plan
+ else:
+ central_plan = dist_wrapper.reduce_scatter("plan", local_step, global_step)
+ central_plan = planner.finish_plan(central_plan)
+ end_plan = time()
+ logger.debug(f"rank: {rank}, plan time: {end_plan - start_plan}")
+ # Prepare async writing of tensors.
+ # The `storage_writer` will store the information about tensors it needs to save
+ start = time()
+ storage_writer.prepare_write_data(central_plan, planner)
+ end = time()
+ logger.debug(f"{time()} rank: {rank}, write(async) time: {end - start}")
+ return (
+ (storage_writer, cast(Metadata, global_metadata), dist_wrapper),
+ central_plan,
+ local_plan,
+ cached_central_plan == central_plan,
+ )
+
+
+def save_state_dict_async_finalize(
+ storage_writer: 'FileSystemWriterAsync', global_metadata: Metadata, dist_wrapper: _DistWrapper
+) -> None:
+ """
+ Finalization of save_state_dict_async_plan.
+
+ The input arguments are the same as the save_state_dict_async_plan output,
+ the `write_results` are retrieved from the storage_writer.
+
+ Args:
+ storage_writer (FileSystemWriterAsync): storage writer used for planning
+ global_metadata (Metadata): metadata created during planning
+ dist_wrapper (_DistWrapper): distributed wrapper created during planning
+
+ Returns: None
+ """
+ write_results = storage_writer.retrieve_write_results()
+
+ # Gather the write results that will be saved to the metadata file.
+ gather_start = time()
+ all_results = dist_wrapper.gather_object(write_results)
+ gather_end = time()
+ logger.debug(f"{gather_end}, {torch.distributed.get_rank()}, gather: {gather_end-gather_start}")
+
+ # Store the metadata on coordinator rank
+ if dist_wrapper.is_coordinator:
+ node_failures = _get_failure_dict(all_results)
+ if len(node_failures) == 0:
+ assert global_metadata is not None
+ write_start = time()
+ storage_writer.finish(global_metadata, all_results)
+ write_end = time()
+ logger.debug(f"{write_end}, metadata_write: {write_end - write_start}")
+ else:
+ raise CheckpointException("write", node_failures)
diff --git a/megatron/core/dist_checkpointing/strategies/tensorstore.py b/megatron/core/dist_checkpointing/strategies/tensorstore.py
new file mode 100644
index 0000000000..0b20bf5e77
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/tensorstore.py
@@ -0,0 +1,128 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Strategies using TensorStore to load and save Zarr arrays. """
+
+from functools import partial
+from itertools import starmap
+from logging import getLogger
+from pathlib import Path
+
+import tensorstore as ts
+import torch
+
+from ..core import CheckpointingException
+from ..dict_utils import dict_list_map_inplace
+from ..mapping import ShardedStateDict, ShardedTensor
+from .base import LoadShardedStrategy, StrategyAction, register_default_strategy
+from .zarr import load_zarr_based_sharded_metadata, postprocess_numpy_array
+
+logger = getLogger(__name__)
+
+
+def register_default_tensorstore_strategies():
+ """Register default strategies leveraging tensorstore."""
+ register_default_strategy(
+ StrategyAction.LOAD_SHARDED, 'zarr', 1, TensorStoreLoadShardedStrategy()
+ )
+
+
+class TensorStoreLoadShardedStrategy(LoadShardedStrategy):
+ """Load strategy for Zarr backend using `tensorstore` for loading."""
+
+ def __init__(self, load_directly_on_device: bool = False):
+ super().__init__()
+ self.load_directly_on_device = load_directly_on_device
+
+ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
+ if torch.distributed.get_rank() == 0:
+ print(f'Loading distributed checkpoint with {self.__class__.__name__}')
+ if self.load_directly_on_device:
+ print(f'Loading distributed checkpoint directly on the GPU')
+ load_fn = partial(
+ _load_from_array,
+ checkpoint_dir=checkpoint_dir,
+ load_directly_on_device=self.load_directly_on_device,
+ )
+ dict_list_map_inplace(load_fn, sharded_state_dict)
+ return sharded_state_dict
+
+ def load_tensors_metadata(self, checkpoint_dir: Path):
+ def get_ts_shape_dtype(path):
+ arr = open_ts_array(path)
+ return arr.shape, arr.dtype.numpy_dtype
+
+ return load_zarr_based_sharded_metadata(checkpoint_dir, get_ts_shape_dtype)
+
+ def check_backend_compatibility(self, loaded_version):
+ pass # TODO
+
+ def check_version_compatibility(self, loaded_version):
+ pass # TODO
+
+
+def merge_global_slice_with_shape(global_slice, actual_shape, key):
+ """Intersects the global slice with the actual shape (prevent overflow)."""
+
+ def _merge_slice(dim_slice, dim_size):
+ if isinstance(dim_slice, slice):
+ assert (
+ dim_slice.start < dim_size
+ ), f'Got empty slice for ShardedTensor {key} ({dim_slice}, {dim_size})'
+ if dim_slice.stop > dim_size:
+ dim_slice = slice(dim_slice.start, dim_size, dim_slice.step)
+ return dim_slice
+
+ assert len(global_slice) == len(actual_shape), (global_slice, actual_shape, key)
+ return tuple(starmap(_merge_slice, zip(global_slice, actual_shape)))
+
+
+def _load_from_array(
+ sharded_tensor: ShardedTensor,
+ checkpoint_dir: Path,
+ load_directly_on_device: bool = False,
+ apply_flattened_range: bool = True,
+):
+ x = _load_regular_chunk(sharded_tensor, checkpoint_dir)
+ ten = postprocess_numpy_array(x, sharded_tensor, apply_flattened_range)
+ if load_directly_on_device:
+ sharded_tensor.data.data.copy_(ten)
+ return sharded_tensor.data
+ else:
+ return ten
+
+
+def _load_regular_chunk(sharded_tensor: ShardedTensor, checkpoint_dir: Path):
+ assert isinstance(sharded_tensor, ShardedTensor), type(sharded_tensor)
+ arr = open_ts_array(checkpoint_dir / sharded_tensor.key)
+ if sharded_tensor.global_shape == arr.shape:
+ x = (
+ arr[sharded_tensor.global_slice()].read().result()
+ ) # flattened tensors loading is delayed
+ elif sharded_tensor.allow_shape_mismatch:
+ global_slice = merge_global_slice_with_shape(
+ sharded_tensor.global_slice(), arr.shape, sharded_tensor.key
+ )
+ x = arr[global_slice].read().result() # flattened tensors loading is delayed
+ else:
+ _msg = (
+ f'Global shape mismatch for loaded ({arr.shape})'
+ f' and expected ({sharded_tensor.global_shape}) tensor'
+ f' for key {sharded_tensor.key}'
+ )
+ raise CheckpointingException(_msg)
+ return x
+
+
+def open_ts_array(arr_path: Path):
+ """Opens a Zarr file array with Tensorstore with basic setting.
+
+ Args:
+ arr_path (Path): path to a Zarr (Tensorstore) array
+ """
+ spec = {'driver': 'zarr', 'metadata_key': '.zarray', 'kvstore': {}}
+ spec['kvstore'] = {'driver': 'file', 'path': str(arr_path)}
+ try:
+ arr = ts.open(ts.Spec(spec), open=True).result()
+ except Exception as e:
+ raise CheckpointingException(f'Array {arr_path} could not be loaded. Error: {e}') from e
+ return arr
diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py
new file mode 100644
index 0000000000..e1bd843417
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/torch.py
@@ -0,0 +1,840 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Strategies using PyTorch distributed.checkpoint as an underlying format. """
+import io
+from collections import ChainMap, defaultdict
+from dataclasses import dataclass
+from itertools import product
+from logging import getLogger
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, cast
+
+import torch
+from packaging.version import Version as PkgVersion
+from torch.distributed import checkpoint
+from torch.distributed._shard.metadata import ShardMetadata
+from torch.distributed._shard.sharded_tensor import Shard, ShardedTensorMetadata, TensorProperties
+from torch.distributed._sharded_tensor import ShardedTensor as TorchShardedTensor
+from torch.distributed._tensor import DTensor
+from torch.distributed.checkpoint import (
+ BytesStorageMetadata,
+ DefaultLoadPlanner,
+ DefaultSavePlanner,
+ FileSystemReader,
+ LoadPlan,
+ Metadata,
+ ReadItem,
+ SavePlan,
+ TensorStorageMetadata,
+ WriteItem,
+)
+from torch.distributed.checkpoint._nested_dict import FLATTEN_MAPPING, unflatten_state_dict
+from torch.distributed.checkpoint._traverse import OBJ_PATH, traverse_state_dict
+from torch.distributed.checkpoint.metadata import Metadata
+from torch.distributed.checkpoint.planner_helpers import _create_write_items
+
+from ..core import CheckpointingException
+from ..dict_utils import nested_values
+from ..mapping import (
+ ShardedBase,
+ ShardedObject,
+ ShardedStateDict,
+ ShardedTensor,
+ StateDict,
+ is_main_replica,
+)
+from .async_utils import AsyncRequest
+from .base import (
+ AsyncSaveShardedStrategy,
+ LoadShardedStrategy,
+ StrategyAction,
+ register_default_strategy,
+)
+from .filesystem_async import FileSystemWriterAsync
+from .resharding import (
+ TensorReformulationMetadata,
+ apply_nd_flattened_tensors_reformulation,
+ is_nd_flattened_tensor,
+ nd_flattened_tensor_reformulated_global_shape,
+ restore_nd_flattened_tensors_formulation,
+)
+from .state_dict_saver import save_state_dict_async_finalize, save_state_dict_async_plan
+
+try:
+ if not torch.cuda.is_available():
+ raise ImportError
+ from transformer_engine.pytorch.float8_tensor import Float8Tensor
+
+ HAVE_TE = True
+except ImportError:
+ HAVE_TE = False
+
+
+def register_default_torch_strategies():
+ """Register default strategies related to PyT Distributed backend."""
+ register_default_strategy(
+ StrategyAction.LOAD_SHARDED, 'torch_dist', 1, TorchDistLoadShardedStrategy()
+ )
+ register_default_strategy(
+ StrategyAction.SAVE_SHARDED, 'torch_dist', 1, TorchDistSaveShardedStrategy('torch_dist', 1)
+ )
+
+
+logger = getLogger(__name__)
+
+
+def flatten_state_dict(
+ state_dict: ShardedStateDict,
+) -> Tuple[ShardedStateDict, Dict[str, OBJ_PATH]]:
+ """Flattens state dict into a single level dict.
+
+ It's a copy of torch.distributed.checkpoint._nested_dict.flatten_state_dict
+ which also accepts ShardedBase tensors as terminal objects
+
+ Args:
+ state_dict (ShardedStateDict): state dict to be flattened
+
+ Returns (tuple): flattened state dict and a mapping allowing to recreate the original one
+
+ """
+ flattened = {}
+ mappings = {}
+
+ def flat_copy(path: OBJ_PATH, value: Any) -> None:
+ new_fqn = ".".join(map(str, path))
+ if new_fqn in flattened:
+ raise ValueError(f"duplicated flatten key {new_fqn}")
+ flattened[new_fqn] = value
+ mappings[new_fqn] = path
+
+ traverse_state_dict(state_dict, flat_copy, lambda x: isinstance(x, (torch.Tensor, ShardedBase)))
+ return flattened, mappings
+
+
+def sharded_tensor_to_torch_sharded_tensor(
+ sh_tens: List[ShardedTensor], rank: Optional[int] = None
+) -> TorchShardedTensor:
+ """Convert MCore ShardedTensor to PyT ShardedTensor. PyT requires information about all chunks.
+
+ On high-level, this function follows the logic of
+ torch.distributed.fsdp._shard_utils._create_chunk_sharded_tensor.
+ Additionally, it saves `prepend_axis_num` and `has_flattened_range` (specific to MCore)
+ as attributes for further restoration in `_unwrap_pyt_sharded_tensor`.
+
+ NOTE: this function assumes regular (grid) sharding of the MCore ShardedTensor.
+ The only local irregularities could be introduced with a `flattened_range` attribute.
+
+ This function handles 3 different type of ShardedTensors:
+ 1. Non-flat regular ShardedTensors (`not has_flattened_range`)
+ 2. 1D flattened ShardedTensors (`is_flattened_range_1d`)
+ 3. N-D flattened ShardedTensors (`has_flattened_range`)
+
+ (1) and (2) type are saved according to their original shape.
+ Type (3) however requires global shape adjustment for efficiency:
+ we treat [X, Y, Z] global shape tensor with local shape [x, y, z]
+ as a [X // x, Y // y, Z // z, x * y * z] tensor with last axis
+ partitioned according to `flattened_range` slices.
+ This will need special handling while resharding.
+
+ Args:
+ sh_tens (List[ShardedTensor]): list of sharded tensors to convert
+ rank (int, optional): current process rank passed to PyT ShardedTensor.
+ If None, assumes rank in the default pg.
+
+ Returns (TorchShardedTensor): PyT ShardedTensor containing all passed shards.
+
+ """
+ if rank is None:
+ rank = torch.distributed.get_rank()
+
+ some_sh_ten = sh_tens[0]
+ has_flattened_range = some_sh_ten.flattened_range is not None
+ is_flattened_range_1d = has_flattened_range and len(some_sh_ten.global_shape) == 1
+
+ for sh_ten in sh_tens:
+ assert (sh_ten.flattened_range is not None) == has_flattened_range, sh_tens
+ if not sh_ten.data.is_contiguous():
+ sh_ten.data = sh_ten.data.contiguous()
+
+ local_global_offsets = {}
+
+ prepend_axis_num = sh_tens[0].prepend_axis_num
+ # Determine local shards according to tensor type (see docs)
+ if is_flattened_range_1d:
+ # Type (2) case: 1D flattened ShardedTensors
+ for sh_ten in sh_tens:
+ assert len(sh_ten.global_offset) == 1, sh_ten
+ assert sh_ten.prepend_axis_num == 0, sh_ten
+ local_global_offsets.setdefault(sh_ten.global_offset, []).append(sh_ten)
+
+ global_shape = some_sh_ten.global_shape
+ offsets_shape = (
+ some_sh_ten.local_shape
+ ) # local shape is not flattened, we need it for chunk offsets
+
+ local_shards = [
+ Shard.from_tensor_and_offsets(
+ sh_ten.data,
+ [
+ sh_ten.global_offset[0] + sh_ten.flattened_range.start
+ ], # additional flattened offset
+ rank,
+ )
+ for sh_ten in sh_tens
+ ]
+
+ elif has_flattened_range:
+ # Type (3) case: N-D flattened ShardedTensors
+ for sh_ten in sh_tens:
+ local_global_offsets.setdefault(sh_ten.local_chunk_offset_in_global(), []).append(
+ sh_ten
+ )
+ assert sh_ten.data.ndim == 1, sh_ten
+ sh_ten.data = sh_ten.data.view((1,) * len(sh_ten.global_shape) + (-1,))
+
+ # Global shape reformulation:
+ global_shape = nd_flattened_tensor_reformulated_global_shape(some_sh_ten)
+ offsets_shape = (1,) * len(
+ some_sh_ten.global_shape
+ ) # reformulated global shape has shape equal ti number of local chunks
+
+ local_shards = [
+ Shard.from_tensor_and_offsets(
+ sh_ten.data,
+ list(
+ sh_ten.local_chunk_offset_in_global() + (sh_ten.flattened_range.start,)
+ ), # additional flattened offset
+ rank,
+ )
+ for sh_ten in sh_tens
+ ]
+ else:
+ # Type (1) case: non-flat regular ShardedTensors
+ for sh_ten in sh_tens:
+ local_global_offsets.setdefault(sh_ten.global_offset, []).append(sh_ten)
+ sh_ten.data = sh_ten.data.view(
+ (1,) * prepend_axis_num + sh_ten.local_shape
+ ) # adjust to prepended_axis_num
+
+ global_shape = some_sh_ten.global_shape
+ offsets_shape = some_sh_ten.data.shape # includes prepended axes
+
+ local_shards = [
+ Shard.from_tensor_and_offsets(
+ sh_ten.data, list(sh_ten.global_offset), rank # simple case
+ )
+ for sh_ten in sh_tens
+ ]
+
+ # Create a ShardedTensor without invoking communication. Determine global shards
+ world_size = torch.distributed.get_world_size()
+ shard_metadata = []
+ # NOTE: here we assume a regular grid of shards
+ for fragment_offsets in product(*map(range, some_sh_ten.axis_fragmentations)):
+ offset = tuple(map(lambda x: x[0] * x[1], zip(fragment_offsets, offsets_shape)))
+ if offset in local_global_offsets:
+ # local shard
+ placement = f"rank:{rank}/cuda"
+ for sh_ten in local_global_offsets[offset]:
+ if is_flattened_range_1d:
+ offset = (sh_ten.global_offset[0] + sh_ten.flattened_range.start,)
+ size = sh_ten.data.shape
+ elif has_flattened_range:
+ assert offset == sh_ten.local_chunk_offset_in_global()
+ # This is not an actual offset, but an offset of the whole shard
+ # This is needed for a PyT Dist internal integrity check
+ offset = sh_ten.local_chunk_offset_in_global() + (0,)
+ size = (1,) * len(offsets_shape) + global_shape[-1:]
+ else:
+ size = sh_ten.data.shape
+ shard_metadata.append(ShardMetadata(offset, size, placement))
+
+ else:
+ # pylint: disable=line-too-long
+ # for shards from other ranks we provide simplistic data - this information will be discarded
+ # during TorchShardedTensor._init_from_local_shards_and_global_metadata call.
+ # Due to a bug in PyT 24.05 container we must specify some concrete rank within a world size.
+ # The exact rank doesn't matter as long as it's different than my rank - hence (rank + 1) % WS.
+ placement = f"rank:{(rank + 1) % world_size}/cuda"
+ if has_flattened_range and not is_flattened_range_1d:
+ offset = offset + (0,)
+ size = (1,) * len(offsets_shape) + global_shape[-1:]
+ else:
+ size = offsets_shape
+ shard_metadata.append(ShardMetadata(offset, size, placement))
+
+ tensor = some_sh_ten.data
+ sharded_tensor_metadata = ShardedTensorMetadata(
+ shards_metadata=shard_metadata,
+ size=torch.Size(global_shape),
+ tensor_properties=TensorProperties(
+ dtype=tensor.dtype,
+ layout=tensor.layout,
+ requires_grad=tensor.requires_grad,
+ memory_format=torch.contiguous_format,
+ pin_memory=tensor.is_pinned(),
+ ),
+ )
+ pyt_sh_ten = TorchShardedTensor._init_from_local_shards_and_global_metadata(
+ local_shards, sharded_tensor_metadata=sharded_tensor_metadata, process_group=None
+ )
+ # Store MCore related data as PyTShardedTensor attribute.
+ # This won't be stored in the checkpoint, only for runtime purposes
+ pyt_sh_ten.mcore_sh_ten = sh_ten.without_data()
+ pyt_sh_ten.mcore_metadata = {}
+ if has_flattened_range and not is_flattened_range_1d:
+ pyt_sh_ten.mcore_metadata['nd_reformulated_orig_global_shape'] = sh_ten.global_shape
+ return pyt_sh_ten
+
+
+def mcore_to_pyt_state_dict(
+ state_dict: Dict[str, List[ShardedBase]],
+ is_loading: bool = False,
+ init_device: torch.device = torch.device("cpu"),
+) -> Dict[str, Union[TorchShardedTensor, io.BytesIO]]:
+ """Convert state dict with ShardedTensors and ShardedObjects
+ to state dict compatible with PyT Dist format.
+
+ Operates in-place and returns the original state dict.
+
+ Args:
+ state_dict (Dict[str, List[ShardedBase]]): flattened state dict, where values
+ are lists of either ShardedTensor or ShardedObjects.
+ is_loading (bool, optional): flag indicating if loading or saving. Defaults to False.
+ init_device (torch.device, optional): device to initialize potentially missing tensors
+ during loading. Defaults to 'cpu'.
+
+ Returns (Dict[str, Union[TorchShardedTensor, io.BytesIO]]): original dictionary with values
+ converted either into PyT ShardedTensors or io.BytesIO.
+
+ """
+ rank = torch.distributed.get_rank()
+ pyt_state_dict = {}
+
+ def _mcore_to_torch_sharded_tensor(sh_tens: List[ShardedTensor]) -> TorchShardedTensor:
+ """Build a PyT ShardedTensor from given shards.
+
+ During loading:
+ - if data is None, initialize it with an empty tensor (will be used to copy the data into)
+ - if `allow_shape_mismatch` is True, the data is initialized with zeros
+ prior to loading (not all parts of the tensor will be read from the checkpoint)
+ """
+ assert all(isinstance(sh_ten, ShardedTensor) for sh_ten in sh_tens), sh_tens
+ for sh_ten in sh_tens:
+ if sh_ten.data is None:
+ if is_loading:
+ sh_ten.init_data(
+ init_device,
+ init_fn=torch.zeros if sh_ten.allow_shape_mismatch else torch.empty,
+ )
+ else:
+ raise CheckpointingException(f'`data` attr is None for {sh_ten}')
+ else:
+ sh_ten.data = sh_ten.data.detach()
+ if sh_ten.allow_shape_mismatch and is_loading:
+ sh_ten.data.zero_()
+
+ torch_sh_ten = sharded_tensor_to_torch_sharded_tensor(sh_tens, rank)
+ torch_sh_ten.key = sh_tens[0].key
+ return torch_sh_ten
+
+ def _mcore_to_torch_sharded_object(sh_objs: List[ShardedObject]) -> io.BytesIO:
+ """Build io.BytesIO from given sharded objects data."""
+ assert all(isinstance(sh_obj, ShardedObject) for sh_obj in sh_objs), sh_objs
+ serialized_data = io.BytesIO()
+ torch.save([sh_obj.data for sh_obj in sh_objs], serialized_data)
+ return serialized_data
+
+ for k, v in state_dict.items():
+ if isinstance(v[0], ShardedTensor):
+ v = cast(List[ShardedTensor], v)
+ pyt_state_dict[k] = _mcore_to_torch_sharded_tensor(v)
+ else:
+ v = cast(List[ShardedObject], v)
+ pyt_state_dict[k] = _mcore_to_torch_sharded_object(v)
+
+ return pyt_state_dict
+
+
+def _unwrap_pyt_sharded_tensor(sh_ten: TorchShardedTensor) -> List[torch.Tensor]:
+ """Unwrap tensor from PyT ShardedTensor instance.
+
+ If `prepend_axis_num` was non-zero (which is specific to MCore ShardedTensor)
+ then the tensor has additional singleton dimensions which should be squeezed.
+ """
+ mcore_sh_ten = sh_ten.mcore_sh_ten
+ ret_tensors = []
+ for sh in sh_ten.local_shards():
+ ten = sh.tensor
+ if mcore_sh_ten.flattened_range is not None:
+ assert ten.shape[:-1] == (1,) * (len(ten.shape) - 1), ten.shape
+ ten = ten.view(-1)
+ else:
+ for _ in range(mcore_sh_ten.prepend_axis_num):
+ ten = ten.squeeze(0)
+ ret_tensors.append(ten)
+ return ret_tensors
+
+
+def _replace_state_dict_keys_with_sharded_keys(
+ sharded_state_dict: ShardedStateDict, keep_only_main_replica: bool = False
+) -> Tuple[Dict[str, List[ShardedBase]], FLATTEN_MAPPING, Dict[str, List[str]]]:
+ """Group ShardedBase objects by keys and
+ return mappings required for recreating the original dict."""
+ flat_sd, flat_mapping = flatten_state_dict(sharded_state_dict)
+ rename_mapping = defaultdict(list)
+ new_flat_sd = defaultdict(list)
+ for k, sh_base in flat_sd.items():
+ assert isinstance(sh_base, ShardedBase), type(sh_base)
+ key = sh_base.unique_key if isinstance(sh_base, ShardedObject) else sh_base.key
+ if is_main_replica(sh_base.replica_id) or not keep_only_main_replica:
+ rename_mapping[key].append(k)
+ new_flat_sd[key].append(sh_base)
+ return new_flat_sd, flat_mapping, rename_mapping
+
+
+def _replace_sharded_keys_with_state_dict_keys(
+ state_dict: Dict[str, List[Union[torch.Tensor, io.BytesIO]]],
+ flat_mapping: FLATTEN_MAPPING,
+ rename_mapping: Dict[str, List[str]],
+):
+ """Inverse of _replace_state_dict_keys_with_sharded_keys."""
+ recovered_sd = {}
+ for k, tensors in state_dict.items():
+ assert len(tensors) == len(rename_mapping[k])
+ for ten, recovered_k in zip(tensors, rename_mapping[k]):
+ recovered_sd[recovered_k] = ten
+
+ return unflatten_state_dict(recovered_sd, flat_mapping)
+
+
+def _restore_dict_types(x: Union[dict, list, Any], keys_template: Union[dict, list, Any]):
+ """Recursively update `x` keys, based on `keys_template`."""
+ if isinstance(keys_template, dict):
+ assert isinstance(x, dict), type(x)
+ for k, v in keys_template.items():
+ if not isinstance(k, str):
+ assert str(k) in x, (k, x.keys)
+ x[k] = x.pop(str(k))
+ _restore_dict_types(x[k], v)
+ elif isinstance(keys_template, list):
+ assert isinstance(x, list), type(x)
+ for x_val, templ_val in zip(x, keys_template):
+ _restore_dict_types(x_val, templ_val)
+
+
+@dataclass(frozen=True)
+class MCoreSavePlan(SavePlan):
+ """SavePlan with MCore specific data."""
+
+ mcore_data: Dict[str, Dict[str, Any]] = None # Mcore related data about each tensor
+
+
+class MCoreSavePlanner(DefaultSavePlanner):
+ """Differs with the default planner by saving BytesIO objects on all ranks.
+
+ In the integration of MCore with PyT Distributed format, BytesIO objects
+ come from ShardedObjects, which should be treated as separate objects on each rank
+ (not common on all ranks).
+
+ Also, the objects are already packed in io.BytesIO, so no need to redo it
+ in transform_object.
+ """
+
+ def __init__(
+ self,
+ *args,
+ dedup_replicated_tensors: Optional[bool] = None,
+ nd_flattened_global_shapes: Optional[Dict[str, Tuple[int, ...]]] = None,
+ **kwargs,
+ ) -> None:
+ # `dedup_replicated_tensors` was deprecated in 2.3; this check avoids warnings
+ # during saving.
+ if PkgVersion(torch.__version__) <= PkgVersion("2.2"):
+ kwargs['dedup_replicated_tensors'] = dedup_replicated_tensors
+ super().__init__(*args, **kwargs)
+ self.nd_flattened_global_shapes = nd_flattened_global_shapes or {}
+
+ def create_local_plan(self) -> SavePlan:
+ """Adds IOBytes write request on non-coordinator ranks."""
+
+ # NOTE: for PyT 2.4.0a0 we can't rely on `create_default_local_save_plan` because
+ # some alpha versions (specifically 2.4.0a0+f70bd71a48 in 24.06 NGC PyTorch container)
+ # add iobytes request only on coordinator ranks and some alpha versions
+ # (specifically 2.4.0a0+3bcc3cddb5 in 24.07 NGC PyTorch container)
+ # add those requests on all ranks. We inline a simplified version of this method below.
+ write_items = []
+ for fqn, obj in self.state_dict.items():
+ assert not isinstance(
+ obj, DTensor
+ ) # translation from MCore ShardedTensors shouldn't result in DTensors
+ # Create write requests for tensor and bytes values.
+ # For MCore, these should be already non-duplicates.
+ write_items += _create_write_items(fqn, obj)
+
+ self.plan = MCoreSavePlan(
+ items=write_items,
+ planner_data=self.mappings,
+ mcore_data={
+ k: sh_ten.mcore_metadata
+ for k, sh_ten in self.state_dict.items()
+ if isinstance(sh_ten, TorchShardedTensor)
+ },
+ )
+ return self.plan
+
+ def create_global_plan(self, all_plans: List[MCoreSavePlan]) -> Tuple[List[SavePlan], Metadata]:
+ """Merges MCore data for all plans."""
+ global_plan, metadata = super().create_global_plan(all_plans)
+ metadata.mcore_data = dict(ChainMap(*(plan.mcore_data for plan in all_plans)))
+ return global_plan, metadata
+
+ def transform_object(self, write_item: WriteItem, object: Any):
+ """Make no transformations - bytes objects are already serialized."""
+ return object
+
+
+class MCoreLoadPlanner(DefaultLoadPlanner):
+ """Adds global shape validation to the default planner.
+
+ If global shape validation can be ignored (shouldn't!), the default
+ load planner can be used.
+ """
+
+ def __init__(
+ self, *args, shapes_validation_sharded_tensors: Iterable[ShardedTensor] = (), **kwargs
+ ) -> None:
+ super().__init__(*args, **kwargs)
+ self.shapes_validation_sharded_tensors = shapes_validation_sharded_tensors
+ self._intermediate_read_item_and_target: Optional[Tuple[ReadItem, torch.Tensor]] = None
+
+ def _validate_global_shapes(self, metadata, sharded_tensors):
+ for sh_ten in sharded_tensors:
+ if sh_ten.key not in metadata.state_dict_metadata:
+ raise KeyError(
+ f"{sh_ten.key} from model not in state dict:"
+ f" {sorted(metadata.state_dict_metadata.keys())}"
+ )
+ loaded_shape = metadata.state_dict_metadata[sh_ten.key].size
+ if not is_nd_flattened_tensor(sh_ten):
+ expected_shape = sh_ten.global_shape
+ else:
+ expected_shape = nd_flattened_tensor_reformulated_global_shape(sh_ten)
+ if loaded_shape != expected_shape:
+ _msg = (
+ f'Global shape mismatch for loaded ({loaded_shape})'
+ f' and expected ({expected_shape}) tensor'
+ f' for key {sh_ten.key}'
+ )
+ raise CheckpointingException(_msg)
+
+ def create_local_plan(self) -> LoadPlan:
+ """Runs additional shapes validation."""
+ self._validate_global_shapes(self.metadata, self.shapes_validation_sharded_tensors)
+ return super().create_local_plan()
+
+ def resolve_tensor(self, read_item: ReadItem):
+ """Override to add FP8 support.
+
+ Narrowing the Float8Tensor can create incontiguous tensors and there are
+ no `copy` kernels for such cases. This method creates a contiguous FP8
+ tensors so that the subsequent `copy_` in FileSystemReader succeeds.
+ Note that this requires tracking the original tensor
+ (as `self._intermediate_read_item_and_target` attribute)
+ and restoring it in `commit_tensor` method.
+ """
+ target_tensor = super().resolve_tensor(read_item)
+ if (
+ not target_tensor.is_contiguous()
+ and HAVE_TE
+ and isinstance(target_tensor, Float8Tensor)
+ ):
+ self._intermediate_read_item_and_target = (read_item, target_tensor)
+ target_tensor = Float8Tensor.make_like(
+ target_tensor, data=target_tensor._data.contiguous()
+ )
+ return target_tensor
+
+ def commit_tensor(self, read_item: ReadItem, tensor: torch.Tensor) -> None:
+ """Restores the original FP8 tensor saved in `resolve_tensor`."""
+ if self._intermediate_read_item_and_target is not None:
+ interm_read_item, target_tensor = self._intermediate_read_item_and_target
+ assert (
+ interm_read_item is read_item
+ ), '`commit_tensor` method should be called right after `resolve_tensor`'
+ target_tensor.copy_(tensor)
+ tensor = target_tensor
+ self._intermediate_read_item_and_target = None
+ return super().commit_tensor(read_item, tensor)
+
+
+class TorchDistSaveShardedStrategy(AsyncSaveShardedStrategy):
+ """Async save strategy for the PyT Distributed format.
+
+ The idea is to translate MCore ShardedTensors into PyT ShardedTensors
+ and use the async-adjusted torch.distributed.checkpoint saving mechanism
+ provided by the FileSystemWriterAsync writer.
+ """
+
+ def __init__(
+ self,
+ backend: str,
+ version: int,
+ keep_only_main_replica: bool = True,
+ thread_count: int = 2,
+ cached_metadata: bool = False,
+ ):
+ """Adds parameters specific to PyT Distributed format
+ Args:
+ backend (str): format backend string
+ version (int): format version
+ keep_only_main_replica (bool, optional): PyT Distributed has a mechanism
+ for deduplication, but replica_id aware deduplication is more coherent.
+ Default is True (recommended to keep it).
+ thread_count (int, optional): threads to use during saving.
+ Affects the number of files in the checkpoint (saving ranks * num_threads).
+ cached_metadata (bool, optional): Enables using cached global metadata to avoid
+ gathering local metadata every checkpointing invocation
+ """
+ super().__init__(backend, version)
+ self.keep_only_main_replica = keep_only_main_replica
+ self.thread_count = thread_count
+
+ # Cached SavePlans to skip plan in `save_state_dict_async_plan`
+ # cached outcome of `SavePlan.prepare_global_plan`,
+ # which aggregates local plans from all ranks
+ self.cached_central_plan: SavePlan = None
+ # cached outcome of `SavePlan.prepare_local_plan` describes how local state_dict is written
+ self.cached_local_plan: SavePlan = None
+ # Cached global metadata, only `coordinator` for dist-ckpt holds
+ # if central plans are consistent over iters
+ self.cached_global_metadata: Metadata = None
+ # This variable records if the ckpt structures are consistent
+ # so the following checkpoint savings reuse `cached_global_metadata`
+ self.validated_cache_reuse: bool = False
+ # The knob to enable cached metadata communication in saving
+ self.use_cached_ckpt_structure: bool = cached_metadata
+
+ def async_save(
+ self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path
+ ) -> AsyncRequest:
+ """Translates MCore ShardedTensors to PyT ShardedTensors & saves in PyT Distributed format.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict to save
+ checkpoint_dir (Path): checkpoint directory
+
+ Returns: None
+ """
+ # Translate the state dict
+ (sharded_state_dict, flat_mapping, rename_mapping) = (
+ _replace_state_dict_keys_with_sharded_keys(
+ sharded_state_dict, self.keep_only_main_replica
+ )
+ )
+ pyt_state_dict = mcore_to_pyt_state_dict(sharded_state_dict, False)
+ # Use PyT saving mechanism
+ writer = FileSystemWriterAsync(checkpoint_dir, thread_count=self.thread_count)
+ # This should be set differently if we run in a smaller process group than the default
+ coordinator = 0
+ # Try twice to validate the generated `central_plan` is the same across iterations
+ # If so, reuse `cached_central_plan` and `cached_global_metadata`
+ # From the 3rd iteration, `save_state_dict_async_plan` will not generate `global_metadata`
+ # (return None) so `self.cached_global_metadata` is reused
+ args_cached_plans = None
+ if self.use_cached_ckpt_structure:
+ args_cached_plans = (
+ self.cached_central_plan,
+ self.cached_local_plan,
+ self.validated_cache_reuse,
+ )
+
+ (
+ save_state_dict_ret,
+ self.cached_central_plan,
+ self.cached_local_plan,
+ self.validated_cache_reuse,
+ ) = save_state_dict_async_plan(
+ pyt_state_dict,
+ writer,
+ None,
+ coordinator,
+ planner=MCoreSavePlanner(dedup_replicated_tensors=not self.keep_only_main_replica),
+ cached_ckpt_structure=args_cached_plans,
+ )
+ rank = torch.distributed.get_rank()
+ if self.use_cached_ckpt_structure:
+ if self.validated_cache_reuse:
+ logger.debug(f"rank: {rank}, cache validated")
+ if save_state_dict_ret[1]: # when global_metadata is not cached
+ self.cached_global_metadata = save_state_dict_ret[1] # Cache Metadata
+ # Only Coordinator rank holds cached global_metadata
+ # (None is returned for global_metadata)
+ elif coordinator == rank:
+ logger.debug(f"rank: {rank}, reuse metadata, {save_state_dict_ret[1]}")
+ save_state_dict_ret = list(save_state_dict_ret)
+ save_state_dict_ret[1] = self.cached_global_metadata
+
+ return self._get_save_and_finalize_callbacks(writer, save_state_dict_ret)
+
+ def _get_save_and_finalize_callbacks(self, writer, save_state_dict_ret) -> AsyncRequest:
+ save_fn_args = writer.get_save_function_and_args()
+ save_fn, save_args = save_fn_args
+
+ def finalize_fn():
+ save_state_dict_async_finalize(*save_state_dict_ret)
+ torch.distributed.barrier()
+
+ return AsyncRequest(save_fn, save_args, [finalize_fn])
+
+ def can_handle_sharded_objects(self):
+ return True
+
+
+def get_reformulation_metadata(
+ sharded_state_dict: ShardedStateDict, checkpoint_dir: Path
+) -> Dict[str, TensorReformulationMetadata]:
+ """Reads MCore data for N-D flattened tensors from checkpoint metadata during ckpt load.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict to load
+ checkpoint_dir (Path): checkpoint directory
+
+ Returns:
+ Dict[str, TensorReformulationMetadata] - dictionary that maps keys of every
+ N-D flattened tensor from the sharded_state_dict to its original global shape
+ as stored in `mcore_data` in the checkpoint.
+ """
+ ckpt_metadata = FileSystemReader(checkpoint_dir).read_metadata()
+ reformulation_metadata = {}
+ for sh_ten in nested_values(sharded_state_dict):
+ if not is_nd_flattened_tensor(sh_ten):
+ continue
+ try:
+ ckpt_global_shape = ckpt_metadata.mcore_data[sh_ten.key][
+ 'nd_reformulated_orig_global_shape'
+ ]
+ except KeyError as e:
+ raise CheckpointingException(
+ f'Cannot find global shape metadata for N-D flattened tensor {sh_ten} '
+ f'in checkpoint metadata: {ckpt_metadata.mcore_data}'
+ ) from e
+
+ reformulation_metadata[sh_ten.key] = TensorReformulationMetadata(
+ ckpt_global_shape, ckpt_metadata.state_dict_metadata[sh_ten.key].size
+ )
+ return reformulation_metadata
+
+
+class TorchDistLoadShardedStrategy(LoadShardedStrategy):
+ """Basic load strategy for the PyT Distributed format."""
+
+ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> StateDict:
+ """Translates MCore ShardedTensors to PyT ShardedTensors & loads from PyT Distributed fmt.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict with mapping
+ information to instruct loading
+ checkpoint_dir (Path): checkpoint directory
+
+ Returns: loaded state dict
+ """
+ # Apply N-D tensors resharding
+ sharded_state_dict, formulation_restore_data = apply_nd_flattened_tensors_reformulation(
+ sharded_state_dict, get_reformulation_metadata(sharded_state_dict, checkpoint_dir)
+ )
+
+ flexible_shape_sharded_tensors = [
+ sh_ten
+ for sh_ten in nested_values(sharded_state_dict)
+ if isinstance(sh_ten, ShardedTensor) and not sh_ten.allow_shape_mismatch
+ ]
+
+ orig_sharded_state_dict = sharded_state_dict
+ # MCore state dict to PyT Distributed compatible
+ (sharded_state_dict, flat_mapping, rename_mapping) = (
+ _replace_state_dict_keys_with_sharded_keys(sharded_state_dict)
+ )
+ pyt_state_dict = mcore_to_pyt_state_dict(sharded_state_dict, True)
+ # Load PyT Distributed format
+ checkpoint.load_state_dict(
+ pyt_state_dict,
+ FileSystemReader(checkpoint_dir),
+ planner=MCoreLoadPlanner(
+ shapes_validation_sharded_tensors=flexible_shape_sharded_tensors
+ ),
+ )
+ pyt_state_dict = cast(
+ Dict[str, Union[TorchShardedTensor, List[io.BytesIO]]], pyt_state_dict
+ )
+ # Unwrap ShardedTensors and return to original state dict
+ mcore_state_dict = {
+ k: v if not isinstance(v, TorchShardedTensor) else _unwrap_pyt_sharded_tensor(v)
+ for k, v in pyt_state_dict.items()
+ }
+ mcore_state_dict = _replace_sharded_keys_with_state_dict_keys(
+ mcore_state_dict, flat_mapping, rename_mapping
+ )
+ _restore_dict_types(mcore_state_dict, orig_sharded_state_dict)
+ # Apply N-D tensors resharding postprocessing
+ mcore_state_dict = restore_nd_flattened_tensors_formulation(
+ mcore_state_dict, formulation_restore_data
+ )
+ return mcore_state_dict
+
+ def load_tensors_metadata(self, checkpoint_dir: Path, metadata: Metadata = None):
+ """Uses tensors metadata stored in the metadata file."""
+ if metadata is None:
+ fs_reader = FileSystemReader(checkpoint_dir)
+ metadata = fs_reader.read_metadata()
+
+ mcore_data = getattr(metadata, 'mcore_data', {})
+ sharded_metadata = {}
+ for k, tp in metadata.state_dict_metadata.items():
+ if not isinstance(tp, TensorStorageMetadata):
+ continue # load only tensors
+
+ nd_orig_global_shape = mcore_data.get(k, {}).get('nd_reformulated_orig_global_shape')
+ if nd_orig_global_shape is None:
+ # Regular tensor
+ sharded_metadata[k] = ShardedTensor.from_rank_offsets(
+ k, torch.empty(tp.size, **tp.properties.__dict__, device='meta')
+ ).without_data()
+ else:
+ # N-D flattened tensor
+ unflat_ten = torch.empty(
+ nd_orig_global_shape, **tp.properties.__dict__, device='meta'
+ )
+ flat_ten = unflat_ten.flatten()
+ sharded_metadata[k] = ShardedTensor.from_rank_offsets_flat(
+ k,
+ flat_ten,
+ unflat_ten.shape,
+ flattened_range=slice(0, unflat_ten.numel()), # whole slice
+ ).without_data()
+
+ return sharded_metadata
+
+ def load_sharded_metadata(self, checkpoint_dir: Path) -> ShardedStateDict:
+ """Uses tensors and objects metadata stored in the metadata file."""
+ fs_reader = FileSystemReader(checkpoint_dir)
+ metadata = fs_reader.read_metadata()
+
+ sharded_metadata = {}
+ for metadata_key, storage_metadata in metadata.state_dict_metadata.items():
+ if not isinstance(storage_metadata, BytesStorageMetadata):
+ continue
+ sh_obj = ShardedObject.empty_from_unique_key(metadata_key)
+ sharded_metadata[sh_obj.unique_key] = sh_obj
+
+ sharded_metadata.update(self.load_tensors_metadata(checkpoint_dir, metadata))
+ return sharded_metadata
+
+ def can_handle_sharded_objects(self):
+ return True
+
+ def check_backend_compatibility(self, loaded_version):
+ pass # TODO
+
+ def check_version_compatibility(self, loaded_version):
+ pass # TODO
diff --git a/megatron/core/dist_checkpointing/strategies/two_stage.py b/megatron/core/dist_checkpointing/strategies/two_stage.py
new file mode 100644
index 0000000000..72e60bc79b
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/two_stage.py
@@ -0,0 +1,254 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" 2-stage checkpoint loading. """
+import os
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from functools import partial, wraps
+from itertools import chain
+from logging import DEBUG, INFO, StreamHandler, getLogger
+from operator import attrgetter, itemgetter
+from pathlib import Path
+from typing import Iterable, List, NamedTuple, Optional, Tuple, Union
+
+import torch
+
+from ..dict_utils import dict_list_map_inplace, map_reduce, nested_values
+from ..mapping import ShardedStateDict, ShardedTensor, StateDict
+from .base import LoadShardedStrategy
+from .tensorstore import TensorStoreLoadShardedStrategy, _load_from_array, open_ts_array
+from .zarr import flatten_range, load_zarr_based_sharded_metadata
+
+_import_trigger = None
+
+
+timers = defaultdict(list)
+
+logger = getLogger(__name__)
+
+
+def timed(verbose=True):
+ def timed_dec(fn):
+ name = fn.__name__
+
+ @wraps(fn)
+ def wrapped(*args, **kwargs):
+ if verbose:
+ logger.debug(f'{name} init')
+ start = time.time()
+ ret = fn(*args, **kwargs)
+ took = time.time() - start
+ if verbose:
+ logger.debug(f'{name} took {took}s')
+ timers[name].append(took)
+ return ret
+
+ return wrapped
+
+ return timed_dec
+
+
+@dataclass
+class _ShardedTensorMetadata:
+ global_rank: int
+ sharded_tensor_no_data: ShardedTensor
+ dist_group_rank: Tuple[int] # id of distributed group
+ dist_group_ranks: Tuple[int] # id of distributed group
+ data_size: Optional[int] = None # bytes
+
+
+def sharded_tensor_chunk_id(sharded_tensor: ShardedTensor):
+ return (sharded_tensor.key, sharded_tensor.global_offset)
+
+
+class TwoStageDataParallelLoadShardedStrategy(LoadShardedStrategy):
+ """Loads one checkpoint replica from storage and broadcasts to other nodes.
+
+ This strategy loads checkpoint from storage on minimal set of nodes
+ and distributes the checkpoint to other nodes with torch.distributed.
+ Loading is performed with tensorstore.
+
+ Steps:
+ 0. (optional) create Gloo distributed groups
+ 1. Exchange ShardedTensors metadata between all nodes
+ 2. Align needed tensors within DP groups
+ 3. For each globally unique tensor:
+ 3.a) on one of the ranks load it from storage to CPU and move to CUDA
+ 3.b) allocate CUDA tensor on other ranks
+ 3.c) broadcast within DP group
+ 3.d) copy tensor content to the model param location
+ 3.e) free tensor buffers from a) and b)
+
+ Notes:
+ 1. Loading and broadcasting is done sequentially to avoid both host and device OOMs
+ 2. There is a lot of overlap potential between all three steps done for each tensor:
+ 2.a) loading from storage to numpy
+ 2.b) moving CPU tensors to CUDA
+ 2.c) broadcast
+ """
+
+ def __init__(self, data_parallel_group, cpu_transfer=True):
+ super().__init__()
+
+ self.cpu_transfer = cpu_transfer
+ self.data_parallel_group_orig = data_parallel_group
+ self.data_parallel_group = None if cpu_transfer else data_parallel_group
+ self.dp_group_ranks = tuple(
+ sorted(torch.distributed.get_process_group_ranks(data_parallel_group))
+ )
+ self.dp_group_rank = torch.distributed.get_rank(self.data_parallel_group_orig)
+ self.global_rank = torch.distributed.get_rank()
+
+ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
+ self.maybe_init_gloo_group()
+ all_tensors_sorted = self._build_load_plan(sharded_state_dict)
+ self._exchange_loaded_tensors(all_tensors_sorted, sharded_state_dict, checkpoint_dir)
+ # TODO: fix hang in summarize_load_times
+ # self.summarize_load_times()
+ return sharded_state_dict
+
+ def summarize_load_times(self):
+ torch.distributed.barrier()
+ logger.info('Checkpoint loading finished. Summary:')
+ # TODO: `timers` keys are not guaranteed to be the same across ranks which causes hangs
+ for key, times in sorted(timers.items()):
+ times_sum = sum(times)
+ max_times = torch.tensor([times_sum], device='cuda')
+ avg_times = torch.tensor([times_sum], device='cuda')
+ torch.distributed.all_reduce(max_times, op=torch.distributed.ReduceOp.MAX)
+ torch.distributed.all_reduce(avg_times, op=torch.distributed.ReduceOp.SUM)
+ avg_times /= torch.distributed.get_world_size()
+ if torch.distributed.get_rank() == 0:
+ logger.info(f'{key}: max {max_times[0]}, avg {avg_times[0]}')
+
+ @timed(verbose=False)
+ def load_tensor_from_storage(self, checkpoint_dir, ten_meta: _ShardedTensorMetadata):
+ logger.debug(f'_load_from_array({ten_meta.sharded_tensor_no_data.key}) init')
+ ret = _load_from_array(
+ ten_meta.sharded_tensor_no_data,
+ checkpoint_dir,
+ load_directly_on_device=False,
+ apply_flattened_range=False,
+ )
+ logger.debug(f'_load_from_array({ten_meta.sharded_tensor_no_data.key}) DONE')
+ return ret
+
+ @timed()
+ def maybe_init_gloo_group(self):
+ if not self.cpu_transfer:
+ return
+ all_groups = [None] * torch.distributed.get_world_size()
+ torch.distributed.all_gather_object(all_groups, self.dp_group_ranks)
+ all_groups = set(tuple(sorted(gr)) for gr in all_groups)
+ for group_ranks in sorted(all_groups):
+ gloo_pg = torch.distributed.new_group(ranks=group_ranks, backend='gloo')
+ if self.global_rank in group_ranks:
+ self.data_parallel_group = gloo_pg
+ assert self.dp_group_rank == torch.distributed.get_rank(self.data_parallel_group)
+
+ def check_backend_compatibility(self, loaded_version):
+ pass # TODO
+
+ def check_version_compatibility(self, loaded_version):
+ pass # TODO
+
+ @timed()
+ def _build_load_plan(
+ self, sharded_state_dict: ShardedStateDict
+ ) -> List[_ShardedTensorMetadata]:
+ local_meta = [
+ _ShardedTensorMetadata(
+ self.global_rank,
+ sharded_ten.without_data(),
+ self.dp_group_rank,
+ self.dp_group_ranks,
+ )
+ for sharded_ten in nested_values(sharded_state_dict)
+ ]
+ all_meta = [None] * torch.distributed.get_world_size(group=self.data_parallel_group)
+ torch.distributed.all_gather_object(all_meta, local_meta, group=self.data_parallel_group)
+ all_meta = list(chain.from_iterable(all_meta))
+ all_tensors_sorted = self.deduplicate_chunks(all_meta)
+ return all_tensors_sorted
+
+ @timed()
+ def deduplicate_chunks(self, ten_metas: List[_ShardedTensorMetadata]):
+ """Group tensors by chunk and then pick the tensor with the lowest rank.
+
+ NOTE: with proper loading overlap, loading from randomized ranks
+ (instead of the smallest one) could be beneficial here.
+ """
+ ten_metas = map_reduce(
+ ten_metas,
+ key_fn=lambda meta: sharded_tensor_chunk_id(meta.sharded_tensor_no_data),
+ reduce_fn=partial(min, key=attrgetter('dist_group_rank')),
+ )
+ all_metas_sorted = list(map(itemgetter(1), sorted(ten_metas.items())))
+ return all_metas_sorted
+
+ @timed()
+ def _exchange_loaded_tensors(
+ self, ten_metas: List[_ShardedTensorMetadata], sharded_state_dict, checkpoint_dir
+ ):
+ logger.debug(f'_exchange_loaded_tensors, num ten_metas: {len(ten_metas)}')
+ for ten_meta in ten_metas:
+
+ src_rank = torch.distributed.get_global_rank(
+ self.data_parallel_group, ten_meta.dist_group_rank
+ )
+
+ if self.dp_group_rank == ten_meta.dist_group_rank:
+ exchange_tensor = self.load_tensor_from_storage(checkpoint_dir, ten_meta)
+ if not self.cpu_transfer:
+ exchange_tensor = exchange_tensor.cuda()
+ else:
+ # TODO: for non-flattened ranges we could reuse the buffer from the start here
+ exchange_tensor = torch.empty(
+ ten_meta.sharded_tensor_no_data.local_shape,
+ device='cpu' if self.cpu_transfer else 'cuda',
+ dtype=ten_meta.sharded_tensor_no_data.dtype,
+ )
+
+ logger.debug(
+ f'exchange {ten_meta.sharded_tensor_no_data.key}, {exchange_tensor.shape}({exchange_tensor.numel()}), broadcast({src_rank} -> {self.dp_group_ranks})'
+ )
+ torch.distributed.broadcast(
+ exchange_tensor, group=self.data_parallel_group, src=src_rank
+ )
+ self._distribute_data_to_state_dict(ten_meta, exchange_tensor, sharded_state_dict)
+ logger.debug(f'exchange {ten_meta.sharded_tensor_no_data.key} done')
+
+ # free buffer memory
+ exchange_tensor = None
+
+ @timed(verbose=False)
+ def _distribute_data_to_state_dict(
+ self,
+ ten_meta: _ShardedTensorMetadata,
+ loaded_ten: torch.Tensor,
+ sharded_state_dict: ShardedStateDict,
+ ):
+ tensor_key = sharded_tensor_chunk_id(ten_meta.sharded_tensor_no_data)
+
+ def _fill_in_data(t: Union[ShardedTensor, torch.Tensor]):
+ if not isinstance(t, ShardedTensor) or sharded_tensor_chunk_id(t) != tensor_key:
+ # already filled-in or key not matching
+ return t
+ sharded_tensor: ShardedTensor = t
+ x = loaded_ten
+ if sharded_tensor.flattened_range is not None:
+ x = flatten_range(sharded_tensor, x)
+
+ # Reuse existing buffer
+ sharded_tensor.data.data.copy_(x)
+ return sharded_tensor.data
+
+ dict_list_map_inplace(_fill_in_data, sharded_state_dict)
+
+ def load_tensors_metadata(self, checkpoint_dir: Path):
+ def get_ts_shape_dtype(path):
+ arr = open_ts_array(path)
+ return arr.shape, arr.dtype.numpy_dtype
+
+ return load_zarr_based_sharded_metadata(checkpoint_dir, get_ts_shape_dtype)
diff --git a/megatron/core/dist_checkpointing/strategies/zarr.py b/megatron/core/dist_checkpointing/strategies/zarr.py
new file mode 100644
index 0000000000..6a52b12257
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/zarr.py
@@ -0,0 +1,321 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Strategies using Zarr as an underlying format. """
+import logging
+import os
+from functools import partial
+from logging import getLogger
+from pathlib import Path
+from typing import Callable, List, Optional, Tuple
+
+import numpy as np
+import torch
+import zarr
+
+from ..core import CheckpointingException
+from ..dict_utils import dict_list_map_inplace, nested_values
+from ..mapping import ShardedStateDict, ShardedTensor, is_main_replica
+from .base import (
+ LoadShardedStrategy,
+ SaveShardedStrategy,
+ StrategyAction,
+ register_default_strategy,
+)
+
+logger = logging.getLogger(__name__)
+
+numpy_to_torch_dtype_dict = {
+ np.dtype('bool'): torch.bool,
+ np.dtype('uint8'): torch.uint8,
+ np.dtype('int8'): torch.int8,
+ np.dtype('int16'): torch.int16,
+ np.dtype('int32'): torch.int32,
+ np.dtype('int64'): torch.int64,
+ np.dtype('float16'): torch.float16,
+ np.dtype('float32'): torch.float32,
+ np.dtype('float64'): torch.float64,
+ np.dtype('complex64'): torch.complex64,
+ np.dtype('complex128'): torch.complex128,
+}
+
+torch_to_numpy_dtype_dict = {v: k for k, v in numpy_to_torch_dtype_dict.items()}
+
+
+try:
+ # Register a bfloat16 type with this import
+ import tensorstore # pylint: disable=unused-import
+
+ HAS_BFLOAT16 = True
+ numpy_to_torch_dtype_dict[np.dtype('bfloat16')] = torch.bfloat16
+ torch_to_numpy_dtype_dict[torch.bfloat16] = np.dtype('bfloat16')
+except ImportError:
+ HAS_BFLOAT16 = False
+
+logger = getLogger(__name__)
+
+
+def register_default_zarr_strategies():
+ """Register default strategies related to Zarr backend."""
+ register_default_strategy(
+ StrategyAction.SAVE_SHARDED, 'zarr', 1, ZarrSaveShardedStrategy('zarr', 1)
+ )
+
+
+class ZarrSaveShardedStrategy(SaveShardedStrategy):
+ """Save strategy for Zarr backend."""
+
+ def __init__(self, backend: str, version: int):
+ super().__init__(backend, version)
+ logger.warning(
+ f'`zarr` distributed checkpoint backend is deprecated.'
+ ' Please switch to PyTorch Distributed format (`torch_dist`).'
+ )
+
+ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
+ sharded_tensors = list(nested_values(sharded_state_dict))
+ arrays = _create_or_open_zarr_arrays(sharded_tensors, checkpoint_dir)
+ for ten, arr in zip(sharded_tensors, arrays):
+ _save_to_existing_array(ten, arr)
+ torch.distributed.barrier()
+
+
+def _create_or_open_zarr_arrays(
+ sharded_tensors: List[ShardedTensor], checkpoint_dir: Path
+) -> List[Optional[zarr.Array]]:
+ """Returns list of zarr arrays corresponding to given tensors.
+
+ For a sharded tensors that:
+ a) is main replica and represents the first chunk (all offsets 0), creates the Zarr array
+ b) is main replica but not the first chunk,
+ opens the arrays created in (a) (possibly by other process)
+ c) otherwise, sets the corresponding array to None since it won't be used
+
+ Args:
+ sharded_tensors (List[ShardedTensor]): sharded tensors from a given rank
+ that will be saved to checkpoint
+ checkpoint_dir (Path): checkpoint in which the arrays will be created
+ """
+ arrays = []
+ for ten in sharded_tensors:
+ arr = _create_zarr_array(ten, checkpoint_dir) if _should_create_array(ten) else None
+ arrays.append(arr)
+
+ torch.distributed.barrier()
+ # Open arrays created above by other processes
+ for arr_idx, ten in enumerate(sharded_tensors):
+ if arrays[arr_idx] is not None:
+ # array created by this process
+ assert _should_create_array(ten), ten
+ continue
+ if not is_main_replica(ten.replica_id):
+ # this array won't be needed for saving and can stay None
+ continue
+ open_kwargs = {}
+ if ten.flattened_range is not None:
+ open_kwargs['synchronizer'] = zarr.ProcessSynchronizer(
+ str(checkpoint_dir / f'{ten.key}.sync')
+ )
+ arrays[arr_idx] = _open_zarr_array_verbose(checkpoint_dir / ten.key, 'r+', **open_kwargs)
+ return arrays
+
+
+def _should_create_array(ten: ShardedTensor):
+ return (
+ is_main_replica(ten.replica_id)
+ and set(ten.global_offset) == {0}
+ and (ten.flattened_range is None or ten.flattened_range.start == 0)
+ )
+
+
+def _save_to_existing_array(sharded_tensor: ShardedTensor, arr: Optional[zarr.Array]):
+ if not is_main_replica(sharded_tensor.replica_id):
+ return
+ assert arr is not None
+ x = sharded_tensor.data
+ x = x.detach().cpu()
+ torch.cuda.synchronize()
+ if x.dtype == torch.bfloat16:
+ x = x.float()
+ x = x.numpy()
+ x = x.astype('bfloat16')
+ else:
+ x = x.numpy()
+
+ if sharded_tensor.flattened_range is None:
+ arr[sharded_tensor.global_slice()] = x
+ else:
+ arr.set_coordinate_selection(sharded_tensor.global_coordinates(), x)
+
+
+def _create_zarr_array(sharded_tensor: ShardedTensor, checkpoint_dir: Path):
+ np_dtype = torch_to_numpy_dtype_dict[sharded_tensor.dtype]
+ try:
+ arr = zarr.create(
+ sharded_tensor.global_shape,
+ dtype=np_dtype,
+ store=checkpoint_dir / sharded_tensor.key,
+ chunks=sharded_tensor.max_allowed_chunks(),
+ compressor=None,
+ fill_value=None,
+ write_empty_chunks=True,
+ )
+ logger.debug(f'Created a new Zarr array at {checkpoint_dir / sharded_tensor.key}')
+ except zarr.errors.ContainsArrayError as e:
+ raise CheckpointingException(
+ f'Array {checkpoint_dir / sharded_tensor.key} already exists'
+ ) from e
+
+ if HAS_BFLOAT16 and np_dtype == np.dtype('bfloat16'):
+ arr._dtype = np_dtype
+ zarray = arr.store['.zarray']
+ arr.store['.zarray'] = zarray.replace(b' exp_sh:
+ assert False, (
+ f'Expected shape ({exp_sh}) smaller than actual ({x_sh})'
+ f' for {repr(expected_sharded_ten)}'
+ )
+ else:
+ pad_args.extend((0, exp_sh - x_sh))
+ # TODO: behavior control with envvar is for testing purposes only, remove it
+ if not int(os.environ.get('DIST_CKPT_PAD_REPLICATE', 0)):
+ return torch.nn.functional.pad(x, pad_args)
+
+ # unsqueeze and squeeze to get shapes supported by cudnn
+ print(f'Replicating last row for {expected_sharded_ten.key}')
+ if x.dtype == torch.bfloat16:
+ return (
+ torch.nn.functional.pad(x.float().unsqueeze(0), pad_args, mode='replicate')
+ .squeeze(0)
+ .bfloat16()
+ )
+ return torch.nn.functional.pad(x.unsqueeze(0), pad_args, mode='replicate').squeeze(0)
+
+
+def load_zarr_based_sharded_metadata(
+ checkpoint_dir: Path, get_shape_dtype_fn: Callable[[str], Tuple[Tuple[int], np.dtype]]
+) -> ShardedStateDict:
+ """Load metadata of Zarr arrays.
+
+ Args:
+ checkpoint_dir (str): checkpoint root directory
+ get_shape_dtype_fn (str -> ((int, ...), np.dtype)): a function returning
+ an array shape and dtype for a given Zarr array path
+ """
+ sharded_state_dict = {}
+ for subdir in checkpoint_dir.iterdir():
+ if not subdir.is_dir() or not (subdir / '.zarray').exists():
+ continue
+ key = subdir.name
+ arr_shape, arr_dtype = get_shape_dtype_fn(str(subdir))
+
+ sharded_state_dict[key] = ShardedTensor(
+ key,
+ None,
+ numpy_to_torch_dtype_dict[arr_dtype],
+ arr_shape,
+ arr_shape,
+ tuple(0 for _ in arr_shape),
+ tuple(1 for _ in arr_shape),
+ )
+ return sharded_state_dict
diff --git a/megatron/core/dist_checkpointing/utils.py b/megatron/core/dist_checkpointing/utils.py
new file mode 100644
index 0000000000..9186e4790a
--- /dev/null
+++ b/megatron/core/dist_checkpointing/utils.py
@@ -0,0 +1,219 @@
+# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+
+""" Helpers for manipulating sharded tensors and sharded state dicts. """
+
+from typing import Dict, Optional, Tuple
+
+from .dict_utils import dict_list_map_inplace, extract_matching_values
+from .mapping import (
+ LocalNonpersistentObject,
+ ShardedBase,
+ ShardedObject,
+ ShardedStateDict,
+ ShardedTensor,
+ ShardedTensorFactory,
+ StateDict,
+)
+
+# _ShardId uniquely identifies a ShardedTensor. This is a subset of ShardedTensor
+# attributes: key (str), global_offset (tuple) and flattened_range (optional tuple)
+_ShardId = Tuple[str, tuple, Optional[tuple]]
+
+
+def _sharded_tensor_shard_id(sharded_tensor: ShardedTensor) -> _ShardId:
+ """Unique id of the sharded tensor data.
+
+ Should yield the same value for same data replicated on different ranks.
+
+ Args:
+ sharded_tensor (ShardedTensor): sharded tensor representing the data shard
+
+ Returns (tuple): unique id of a data shard
+ """
+ f_range = sharded_tensor.flattened_range
+ return (
+ sharded_tensor.key,
+ sharded_tensor.global_offset,
+ None if f_range is None else (f_range.start, f_range.stop),
+ )
+
+
+def _sharded_object_id(sharded_object: ShardedObject) -> _ShardId:
+ """Unique id of the sharded object data.
+
+ Should yield the same value for same data replicated on different ranks.
+
+ Args:
+ sharded_object (ShardedObject): sharded object representing the data shard
+
+ Returns (tuple): unique id of a data shard
+ """
+ return (sharded_object.key, sharded_object.global_offset, sharded_object.global_shape)
+
+
+def extract_sharded_tensors(
+ sharded_state_dict: ShardedStateDict,
+) -> Tuple[ShardedStateDict, StateDict]:
+ """Extract a dict consisting of only ShardedTensor objects
+ from a given state dict with any objects.
+
+ Args:
+ sharded_state_dict: state dict possibly containing ShardedTensor objects
+
+ Returns:
+ Tuple[ShardedStateDict, StateDict]: tuple of:
+ - state dict with all ShardedTensor (keeping the original state dict structure)
+ - state dict with all objects other than ShardedTensor
+ (keeping the original state dict structure)
+ """
+ return extract_matching_values(sharded_state_dict, lambda v: isinstance(v, ShardedTensor))
+
+
+def extract_sharded_tensors_and_factories(
+ sharded_state_dict: ShardedStateDict,
+) -> Tuple[ShardedStateDict, StateDict]:
+ """Extract a dict consisting of only ShardedTensor and ShardedTensorFactory objects
+ from a given state dict with any objects.
+
+ Args:
+ sharded_state_dict:
+ state dict possibly containing ShardedTensor and ShardedTensorFactory objects
+
+ Returns:
+ Tuple[ShardedStateDict, StateDict]: tuple of:
+ - state dict with all ShardedTensor and ShardedTensorFactory
+ (keeping the original state dict structure)
+ - state dict with all other objects (keeping the original state dict structure)
+ """
+ return extract_matching_values(
+ sharded_state_dict, lambda v: isinstance(v, (ShardedTensor, ShardedTensorFactory))
+ )
+
+
+def extract_sharded_tensors_or_nonpersistent(
+ sharded_state_dict: ShardedStateDict,
+) -> Tuple[ShardedStateDict, StateDict]:
+ """Extract a dict consisting of only ShardedTensor, ShardedTensorFactory
+ and LocalNonpersistentObject objects from a given state dict with any objects.
+
+ Args:
+ sharded_state_dict: state dict possibly containing ShardedTensor, ShardedTensorFactory
+ and LocalNonpersistentObject objects
+
+ Returns:
+ Tuple[ShardedStateDict, StateDict]: tuple of:
+ - state dict with all ShardedTensor, ShardedTensorFactory and LocalNonpersistentObject
+ (keeping the original state dict structure)
+ - state dict with all other objects (keeping the original state dict structure)
+ """
+ return extract_matching_values(
+ sharded_state_dict,
+ lambda v: isinstance(v, (ShardedTensor, LocalNonpersistentObject, ShardedTensorFactory)),
+ )
+
+
+def extract_sharded_base(
+ sharded_state_dict: ShardedStateDict,
+) -> Tuple[ShardedStateDict, StateDict]:
+ """Extract a dict consisting of only ShardedBase from a given state dict with any objects.
+
+ Args:
+ sharded_state_dict: state dict possibly containing ShardedBase objects
+
+ Returns:
+ Tuple[ShardedStateDict, StateDict]: tuple of:
+ - state dict with all ShardedBase objects (keeping the original state dict structure)
+ - state dict with all other objects (keeping the original state dict structure)
+ """
+ return extract_matching_values(sharded_state_dict, lambda v: isinstance(v, ShardedBase))
+
+
+def extract_nonpersistent(
+ sharded_state_dict: ShardedStateDict,
+) -> Tuple[ShardedStateDict, StateDict]:
+ """Extract a dict consisting of only LocalNonpersistentObjects from a given state dict.
+
+ Args:
+ sharded_state_dict: state dict possibly containing LocalNonpersistentObjects
+
+ Returns:
+ Tuple[ShardedStateDict, StateDict]: tuple of:
+ - state dict with all LocalNonpersistentObjects
+ (keeping the original state dict structure)
+ - state dict with all other objects (keeping the original state dict structure)
+ """
+
+ return extract_matching_values(
+ sharded_state_dict, lambda v: isinstance(v, LocalNonpersistentObject)
+ )
+
+
+def add_prefix_for_sharding(sharded_state_dict: ShardedStateDict, prefix: str):
+ """Prepend a given prefix to all ShardedBase objects in a given state dict *in-place*.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict
+ prefix (str): prefix to be prepended
+
+ Returns:
+ None: state dict is modified in-place
+ """
+
+ def add_prefix(t):
+ if isinstance(t, ShardedBase):
+ t.key = f'{prefix}{t.key}'
+ return t
+
+ dict_list_map_inplace(add_prefix, sharded_state_dict)
+
+
+def replace_prefix_for_sharding(
+ sharded_state_dict: ShardedStateDict, old_prefix: str, new_prefix: str
+):
+ """Replaces the given prefix in *all* sharded keys in a given state dict.
+
+ Errors out if some key does not begin with a given prefix.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict to replace keys in
+ old_prefix (str): prefix to be replaced in each key
+ new_prefix (str): new prefix
+
+ Returns:
+ None: state dict is modified in place
+ """
+
+ def _replace_prefix(x):
+ if isinstance(x, (ShardedTensor, ShardedTensorFactory, ShardedObject)):
+ if not x.key.startswith(old_prefix):
+ raise ValueError(f'Expected {x.key} to begin with prefix {old_prefix}')
+ x.key = f'{new_prefix}{x.key[len(old_prefix):]}' # str.removeprefix in Python >= 3.9
+ return x
+
+ dict_list_map_inplace(_replace_prefix, sharded_state_dict)
+
+
+def apply_prefix_mapping(sharded_state_dict: ShardedStateDict, prefix_map: Dict[str, str]):
+ """Replaces prefixes *only in keys matching* with one of prefixes in the map.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict to replace keys in
+ prefix_map (Dict[str, str]):
+ map of old->new prefixes. The first matching prefix for each key is used
+
+ Returns:
+ None: state dict is modified in place
+ """
+
+ def _replace_prefixes(x):
+ if not isinstance(x, (ShardedTensor, ShardedTensorFactory, ShardedObject)):
+ return x
+ for old_prefix, new_prefix in prefix_map.items():
+ if x.key.startswith(old_prefix):
+ x.key = (
+ f'{new_prefix}{x.key[len(old_prefix):]}' # str.removeprefix in Python >= 3.9
+ )
+ break
+ return x
+
+ dict_list_map_inplace(_replace_prefixes, sharded_state_dict)
diff --git a/megatron/core/dist_checkpointing/validation.py b/megatron/core/dist_checkpointing/validation.py
new file mode 100644
index 0000000000..cd11b82ed6
--- /dev/null
+++ b/megatron/core/dist_checkpointing/validation.py
@@ -0,0 +1,525 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import logging
+from collections import Counter, defaultdict
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union
+
+import numpy as np
+import torch
+
+from megatron.core.dist_checkpointing import ShardedTensor
+from megatron.core.dist_checkpointing.core import CheckpointingException, maybe_load_config
+from megatron.core.dist_checkpointing.dict_utils import (
+ extract_matching_values,
+ map_reduce,
+ nested_values,
+)
+from megatron.core.dist_checkpointing.mapping import (
+ ShardedBase,
+ ShardedObject,
+ ShardedStateDict,
+ is_main_replica,
+)
+from megatron.core.dist_checkpointing.strategies.base import (
+ LoadCommonStrategy,
+ LoadShardedStrategy,
+ SaveCommonStrategy,
+ SaveShardedStrategy,
+ StrategyAction,
+ get_default_strategy,
+)
+
+if TYPE_CHECKING:
+ from megatron.core.dist_checkpointing.serialization import CkptShardedMetadata
+
+logger = logging.getLogger(__name__)
+
+# list of local saved/loaded ShardedBase objects
+_LocalMetadata = List[Union[ShardedTensor, ShardedObject]]
+# list of lists of global saved/loaded ShardedBase objects (each list element corresponds to global rank)
+_GlobalMetadata = List[_LocalMetadata]
+
+
+class StrictHandling(Enum):
+ """Determines handling of load mismatch (non-empty "unexpected" or "missing" keys).
+
+ Different flags carry different implications on performance and behaviour and
+ are divided into two groups:
+ - *_UNEXPECTED
+ - *_ALL
+ The first group ignores missing keys (present in the checkpoint but missing
+ in the sharded state dict) which is created in order to avoid inter-rank
+ metadata exchange. Note that the metadata exchange will happen anyway
+ with `load(..., validate_access_integrity=True)` flag in which case using the
+ `*_ALL` option is recommended as it provides a more thorough check with no
+ performance penalty wrt. `*_UNEXPECTED` group.
+
+ All options except for the first one (`ASSUME_OK_UNEXPECTED`) require
+ extra disk access before the load in order to remove unexpected keys
+ from the sharded state dict requested to load.
+ """
+
+ # Relies on the underlying strategy to raise error on unexpected keys
+ ASSUME_OK_UNEXPECTED = 'assume_ok_unexpected'
+ # Logs (with WARNING level) "unexpected" keys. Missing keys are ignored.
+ # This is treated as a reasonable default for a "non-strict" load
+ LOG_UNEXPECTED = 'log_unexpected'
+ # Logs (with WARNING level) all mismatched keys.
+ LOG_ALL = 'log_all'
+ # Raise error on unexpected keys before load attempt.
+ # Gives cleaner error message than `ASSUME_OK_UNEXPECTED` but requires
+ # extra disk access.
+ RAISE_UNEXPECTED = 'raise_unexpected'
+ # Raise error on any mismatch. Similar to `RAISE_UNEXPECTED` but requires
+ # metadata exchange.
+ RAISE_ALL = 'raise_all'
+ # "Unexpected" mismatches are not reported, but returned by the `load`
+ # function along with the loaded state dict. Missing keys are ignored.
+ RETURN_UNEXPECTED = 'return_unexpected'
+ # All mismatches are returned along with the loaded state dict.
+ RETURN_ALL = 'return_all'
+ # Simply ignores mismatches (not recommended)
+ IGNORE_ALL = 'ignore_all'
+
+ @staticmethod
+ def requires_explicit_ckpt_mismatch_check(val: 'StrictHandling') -> bool:
+ """Whether a given strict flag involves mismatch check against the checkpoint."""
+ return val != StrictHandling.ASSUME_OK_UNEXPECTED
+
+ @staticmethod
+ def requires_global_app_metadata(val: 'StrictHandling') -> bool:
+ """Whether a given strict option requires global metadata for validation."""
+ return val in (
+ StrictHandling.IGNORE_ALL,
+ StrictHandling.RAISE_ALL,
+ StrictHandling.RETURN_ALL,
+ StrictHandling.LOG_ALL,
+ )
+
+ @staticmethod
+ def requires_returning_mismatch_keys(val: 'StrictHandling') -> bool:
+ """Whether a given strict option results in extra return value from the `load` function."""
+ return val in (StrictHandling.RETURN_UNEXPECTED, StrictHandling.RETURN_ALL)
+
+
+def parse_strict_flag(strict: Union[str, StrictHandling]) -> StrictHandling:
+ """Parse user passed strict flag from a string to StrictHandling instance.
+
+ Args:
+ strict (str, StrictHandling): strict flag to parse. If already an instance
+ of StrictHandling, this function is a noop.
+
+ Returns:
+ StrictHandling: enum instance
+ """
+ if isinstance(strict, StrictHandling):
+ return strict
+ try:
+ return StrictHandling(strict)
+ except (ValueError, TypeError) as e:
+ raise ValueError(f'Invalid strict flag: {e}') from e
+
+
+def validate_integrity_and_strict_load(
+ sharded_state_dict: ShardedStateDict,
+ strict: StrictHandling,
+ validate_access_integrity: bool,
+ local_metadata: Optional[_LocalMetadata] = None,
+ global_metadata: Optional[_GlobalMetadata] = None,
+ ckpt_sharded_metadata: Optional['CkptShardedMetadata'] = None,
+) -> Tuple[ShardedStateDict, Set[str], Set[str]]:
+ """Validates sharding integrity and potential mismatches with the checkpoint.
+
+ `validate_access_integrity` controls sharding integrity check (orthogonal
+ to strictness checking) which verifies `sharded_state_dict` runtime completeness
+ (in isolation from the actual checkpoint).
+
+ `strict` flag controls handling of mismatches between the requested
+ sharded state dict to load and the actual checkpoint. See `StrictHandling`
+ docs for details regarding flag behavior and performance implications
+ (disk interactions or inter-rank communication).
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict to verify.
+ strict (StrictHandling): flag determining how to handle sharded keys mismatch.
+ validate_access_integrity (bool): whether to perform sharding validation.
+ local_metadata (_LocalMetadata, optional): local sharded state dict metadata.
+ Defaults to None, in which case it's determined based on `sharded_state_dict`.
+ global_metadata (_GlobalMetadata, optional): global sharded state dict metadata
+ (exchanged between ranks). Defaults to None, in which case "missing"
+ keys are not determined.
+ ckpt_sharded_metadata (CkptShardedMetadata, optional): sharded metadata
+ from the checkpoint. Defaults to None, which only makes sense
+ for the `StrictHandling.ASSUME_OK_UNEXPECTED` strict value.
+
+ Returns:
+ Tuple[ShardedStateDict, Set[str], Set[str]]: tuple of: sharded state dict
+ without unexpected keys, missing and unexpected keys. Missing keys are equal
+ on all ranks, unexpected keys might differ across ranks. Additionally,
+ missing keys might be erroneously empty (depending on `strict` value).
+ """
+ missing_keys, unexpected_keys = [], []
+ if StrictHandling.requires_explicit_ckpt_mismatch_check(strict):
+ if ckpt_sharded_metadata is None:
+ raise CheckpointingException(
+ 'Cannot verify checkpoint mismatch with ckpt_sharded_metadata=None.'
+ )
+ if local_metadata is None:
+ local_metadata = [
+ sh_base.without_data() for sh_base in nested_values(sharded_state_dict)
+ ]
+ # We don't want to check for missing keys even if we could
+ _skip_missing_keys = strict in (
+ StrictHandling.ASSUME_OK_UNEXPECTED,
+ StrictHandling.LOG_UNEXPECTED,
+ StrictHandling.RAISE_UNEXPECTED,
+ StrictHandling.RETURN_UNEXPECTED,
+ )
+ missing_keys, unexpected_keys = _determine_missing_and_unexpected_keys(
+ ckpt_sharded_metadata, local_metadata, None if _skip_missing_keys else global_metadata
+ )
+
+ sharded_state_dict = adjust_non_strict_load(sharded_state_dict, unexpected_keys)
+
+ if strict == StrictHandling.IGNORE_ALL:
+ missing_keys, unexpected_keys = [], []
+ elif strict in (StrictHandling.RAISE_UNEXPECTED, StrictHandling.RAISE_ALL):
+ maybe_report_missing_and_unexpected_keys(missing_keys, unexpected_keys, True)
+ elif strict in (StrictHandling.LOG_UNEXPECTED, StrictHandling.LOG_ALL):
+ maybe_report_missing_and_unexpected_keys(missing_keys, unexpected_keys, False)
+
+ if validate_access_integrity:
+ if global_metadata is None:
+ raise CheckpointingException(
+ 'Cannot check sharding intergrity without global_metadata (None).'
+ )
+ validate_sharding_integrity(global_metadata)
+
+ return sharded_state_dict, missing_keys, unexpected_keys
+
+
+def verify_checkpoint_and_load_strategy(
+ checkpoint_dir: str,
+ sharded_strategy: Union[LoadShardedStrategy, Tuple[str, int], None] = None,
+ common_strategy: Union[LoadCommonStrategy, Tuple[str, int], None] = None,
+) -> Tuple[LoadShardedStrategy, LoadCommonStrategy]:
+ """Verifies if checkpoint metadata exists and matches given strategies.
+
+ If no strategies are passed, they are determined based on the checkpoint metadata.
+
+ Args:
+ checkpoint_dir (str): checkpoint directory
+ sharded_strategy (LoadShardedStrategy, Tuple[str, int], optional): sharded load strategy to be verified
+ if compatible with the checkpoint content. If None, the default sharded load strategy
+ for the checkpoint backend will be returned.
+ common_strategy (LoadCommonStrategy, Tuple[str, int], optional): common load strategy to be verified
+ if compatible with the checkpoint content. If None, the default common load strategy
+ for the checkpoint backend will be returned.
+ """
+ if not Path(checkpoint_dir).exists():
+ raise CheckpointingException(f'Checkpoint directory {checkpoint_dir} does not exist')
+
+ saved_config = maybe_load_config(checkpoint_dir)
+ if saved_config is None:
+ raise CheckpointingException(f'{checkpoint_dir} is not a distributed checkpoint')
+
+ if sharded_strategy is None:
+ sharded_strategy = get_default_strategy(
+ StrategyAction.LOAD_SHARDED,
+ saved_config.sharded_backend,
+ saved_config.sharded_backend_version,
+ )
+ elif isinstance(sharded_strategy, tuple):
+ sharded_strategy = get_default_strategy(StrategyAction.LOAD_SHARDED, *sharded_strategy)
+
+ if common_strategy is None:
+ common_strategy = get_default_strategy(
+ StrategyAction.LOAD_COMMON,
+ saved_config.common_backend,
+ saved_config.common_backend_version,
+ )
+ elif isinstance(common_strategy, tuple):
+ sharded_strategy = get_default_strategy(StrategyAction.LOAD_COMMON, *common_strategy)
+
+ sharded_strategy.check_backend_compatibility(saved_config.sharded_backend)
+ sharded_strategy.check_version_compatibility(saved_config.sharded_backend_version)
+ common_strategy.check_backend_compatibility(saved_config.common_backend)
+ common_strategy.check_version_compatibility(saved_config.common_backend_version)
+ return sharded_strategy, common_strategy
+
+
+def adjust_non_strict_load(
+ sharded_state_dict: ShardedStateDict, sharded_keys_to_remove: Set[str]
+) -> ShardedStateDict:
+ """Adjusts sharded state dict removing keys not existing in the checkpoint.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): sharded state dict to modify
+ sharded_keys_to_remove (Set[str]): keys to remove from the state dict
+
+ Returns:
+ ShardedStateDict: state dict without ShardedBase objects with specified keys
+ """
+
+ def is_unexpected_key(x: ShardedBase):
+ assert isinstance(x, ShardedBase), f'Unexpected type {type(x)}'
+ return x.key in sharded_keys_to_remove
+
+ _, sharded_state_dict = extract_matching_values(sharded_state_dict, is_unexpected_key)
+ return sharded_state_dict
+
+
+def _determine_missing_and_unexpected_keys(
+ ckpt_sharded_metadata: 'CkptShardedMetadata',
+ local_metadata: _LocalMetadata,
+ global_metadata: Optional[_GlobalMetadata] = None,
+) -> Tuple[Set[str], Set[str]]:
+ """Determines load mismatches based on metadata.
+
+ There is an asymmetry between "unexpected" and "missing" keys.
+ Unexpected keys can be determined based only on local metadata.
+ Missing keys must be based on global metadata, since other ranks might access
+ different keys than the current rank.
+ In consequence, the return value of this function is different on each rank:
+ "missing_keys" are equal, but "unexpected_keys" might differ across ranks.
+
+ Args:
+ ckpt_sharded_metadata (CkptShardedMetadata): sharded state dict (without data)
+ constructed based on the checkpoint content
+ local_metadata (_LocalMetadata): list of local ShardedBase objects
+ requested to be loaded by this rank
+ global_metadata (_GlobalMetadata, optional): list of global ShardedBase objects
+ requested to be loaded by all ranks. Defaults to None, in which case
+ returned "missing" keys are empty.
+
+ Returns:
+ Tuple[Set[str], Set[str]]: missing and unexpected keys. Missing keys are equal
+ on all ranks, unexpected keys might differ across ranks. If passed
+ `global_metadata` is empty, returned missing keys are empty as well.
+
+ """
+ local_accessed_keys = set(sh_base.key for sh_base in local_metadata)
+ ckpt_keys = set(sh_base.key for sh_base in ckpt_sharded_metadata.values())
+ unexpected_keys = local_accessed_keys - ckpt_keys
+ if global_metadata is not None:
+ global_accessed_keys = set(
+ sh_base.key for rank_metadata in global_metadata for sh_base in rank_metadata
+ )
+ missing_keys = ckpt_keys - global_accessed_keys
+ else:
+ missing_keys = set()
+
+ if missing_keys:
+ logger.debug(f'Dist ckpt load missing keys: {missing_keys}')
+ if unexpected_keys:
+ logger.debug(f'Dist ckpt load unexpected keys: {unexpected_keys}')
+
+ return missing_keys, unexpected_keys
+
+
+def maybe_report_missing_and_unexpected_keys(
+ missing_keys: Set[str], unexpected_keys: Set[str], raise_error: bool = True
+) -> None:
+ """Raises or logs an error in case missing or unexpected keys are non-empty.
+
+ Args:
+ missing_keys (Set[str]): missing keys in the state dict
+ unexpected_keys (Set[str]): unexpected keys in the state dict
+ raise_error: If True, raises error on mismatch. Otherwise, logs mismatch
+ with WARNING level.
+
+ Returns:
+ None
+
+ Raises:
+ CheckpointingException: if `raise_error` is True and at least one of
+ `missing_keys` or `unexpected_keys` are non-empty.
+ """
+ if not missing_keys and not unexpected_keys:
+ return
+ missing_title_msg = (
+ f'Some keys found in the checkpoint are missing in the provided sharded state dict. '
+ )
+ missing_body_msg = f'Missing keys (for all ranks): {missing_keys}. '
+ unexpected_title_msg = f'Unexpected keys (not found in the checkpoint) encountered in the provided sharded state dict. '
+ unexpected_body_msg = f'Unexpected keys (for this rank): {unexpected_keys}. '
+ error_msg = ''
+ if missing_keys:
+ error_msg += missing_title_msg
+ if unexpected_keys:
+ error_msg += unexpected_title_msg
+
+ error_msg += '\n'
+ if missing_keys:
+ error_msg += missing_body_msg
+ if unexpected_keys:
+ error_msg += unexpected_body_msg
+
+ if raise_error:
+ raise CheckpointingException(error_msg)
+ else:
+ logger.warning(error_msg)
+
+
+def validate_sharding_integrity(global_metadata: _GlobalMetadata) -> None:
+ """Validate if the ShardedTensors and ShardedObjects from multiple processes define correct sharding.
+
+ Local ShardedTensors and ShardedObject metadata is exchanged with `torch.distributed.all_gather_object`
+ and then process with global rank 0 checks if main replicas of the shards:
+ - cover the whole global tensors
+ - don't overlap
+
+ Args:
+ global_metadata (_GlobalMetadata): ShardedTensor and ShardedObject objects from all ranks.
+
+ Returns:
+ None
+
+ Raises:
+ CheckpointingException for invalid access pattern
+ """
+ if torch.distributed.get_rank() != 0:
+ return
+
+ key_shardings = defaultdict(list)
+ for rank, rank_shardings in enumerate(global_metadata):
+ for sharding in rank_shardings:
+ key_shardings[sharding.key].append((rank, sharding))
+ for key, shardings in key_shardings.items():
+ if isinstance(shardings[0][1], ShardedObject):
+ _validate_objects_for_key(shardings)
+ else:
+ _validate_sharding_for_key(shardings)
+
+
+def _validate_sharding_for_key(rank_sharding: List[Tuple[int, ShardedTensor]]):
+ some_rank_shard = rank_sharding[0][1]
+ global_shape = some_rank_shard.global_shape
+ local_shape = some_rank_shard.local_shape
+ dtype = some_rank_shard.dtype
+ has_flattened_range = some_rank_shard.flattened_range is not None
+ for rank, sharding in rank_sharding:
+ assert sharding.dtype == dtype, (sharding.dtype, dtype, some_rank_shard)
+ assert sharding.global_shape == global_shape, (
+ sharding.global_shape,
+ global_shape,
+ some_rank_shard,
+ )
+ assert sharding.local_shape == local_shape, (
+ sharding.local_shape,
+ local_shape,
+ some_rank_shard,
+ )
+ assert (sharding.flattened_range is not None) == has_flattened_range, (
+ (sharding.flattened_range is not None),
+ has_flattened_range,
+ some_rank_shard,
+ )
+
+ shard_access_cnt = _compute_shards_access(rank_sharding)
+ if has_flattened_range:
+ map_reduce(
+ rank_sharding,
+ lambda x: x[1].global_offset,
+ lambda x: x[1],
+ _validate_sharding_for_key_flattened,
+ )
+ else:
+ if not torch.all(shard_access_cnt == 1):
+ logger.error(f'Invalid access pattern for {rank_sharding[0][1]}: {shard_access_cnt}')
+ raise CheckpointingException(f'Invalid access pattern for {rank_sharding[0][1]}')
+
+
+def _compute_shards_access(rank_sharding):
+ shard_access_cnt = torch.zeros(
+ rank_sharding[0][1].axis_fragmentations, dtype=torch.int, device='cpu'
+ )
+ for rank, sharding in rank_sharding:
+ if is_main_replica(sharding.replica_id):
+ shard_access_cnt[sharding.local_chunk_offset_in_global()] += 1
+ return shard_access_cnt
+
+
+def _validate_sharding_for_key_flattened(tensors_by_shard):
+ all_slices = []
+ local_shape = tensors_by_shard[0].local_shape
+ for sharding in tensors_by_shard:
+ assert sharding.local_shape == local_shape
+ sharding: ShardedTensor
+ if not is_main_replica(sharding.replica_id):
+ continue
+
+ all_slices.append((sharding.flattened_range.start, sharding.flattened_range.stop))
+
+ starts, stops = map(np.asarray, zip(*sorted(all_slices)))
+ if (
+ starts[0] != 0
+ or stops[-1] != np.product(local_shape)
+ or not np.all(starts[1:] == stops[:-1])
+ ):
+ logger.error(
+ f'Flattened ranges dont cover the whole shard {tensors_by_shard[0]}. Ranges: {(starts, stops)}'
+ )
+ raise CheckpointingException(
+ f'Flattened ranges dont cover the whole shard {tensors_by_shard[0]}. Ranges: {(starts, stops)}'
+ )
+
+
+def _validate_objects_for_key(sharded_objects: List[ShardedObject]):
+ """Ensure uniqueness of saved objects."""
+ unique_keys = [
+ sh_obj.unique_key for _, sh_obj in sharded_objects if is_main_replica(sh_obj.replica_id)
+ ]
+ if len(unique_keys) != len(set(unique_keys)):
+ duplicates = {k: cnt for k, cnt in Counter(unique_keys).items() if cnt > 1}
+ logger.error(f'Duplicate ShardedObject keys and counts: {duplicates}')
+ raise CheckpointingException(f'Duplicate ShardedObject keys: {list(duplicates.keys())}')
+ expected_shard_num = np.prod(sharded_objects[0][1].global_shape)
+ if len(unique_keys) != expected_shard_num:
+ err_msg = f'Invalid access pattern: {expected_shard_num - len(unique_keys)} ShardedObject are missing.'
+ logger.error(f'{err_msg} Existing shards: {unique_keys}')
+ raise CheckpointingException(err_msg)
+
+
+def determine_global_metadata(
+ sharded_state_dict: ShardedStateDict,
+) -> Tuple[_LocalMetadata, _GlobalMetadata]:
+ """Exchanges local metadata with `all_gather_object` to determine global metadata.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): local sharded state dict
+
+ Returns:
+ Tuple[_LocalMetadata, _GlobalMetadata]: local and global ShardedBase objects with stripped data
+ """
+ local_metadata = [ten.without_data() for ten in nested_values(sharded_state_dict)]
+ global_metadata = [None] * torch.distributed.get_world_size()
+ torch.distributed.all_gather_object(global_metadata, local_metadata)
+ return local_metadata, global_metadata
+
+
+def validate_sharded_objects_handling(
+ sharded_strategy: Union[SaveShardedStrategy, LoadShardedStrategy],
+ common_strategy: Union[SaveCommonStrategy, LoadCommonStrategy],
+) -> None:
+ """Checks if either of the passed strategies can handle sharded objects.
+
+ Args:
+ sharded_strategy (Union[SaveShardedStrategy, LoadShardedStrategy]): sharded strategy used for saving/loading
+ common_strategy (Union[SaveCommonStrategy, LoadCommonStrategy]): common strategy used for saving/loading
+
+ Returns:
+ None
+
+ Raises:
+ CheckpointingException: if both strategies can't handle ShardedObjects
+ """
+ if (
+ not sharded_strategy.can_handle_sharded_objects
+ and not common_strategy.can_handle_sharded_objects
+ ):
+ raise CheckpointingException(
+ f'Either sharded strategy or common strategy must implement ShardedObjects handling.'
+ f' Both {sharded_strategy} and {common_strategy} specify can_handle_sharded_objects=False'
+ )
diff --git a/megatron/core/distributed/__init__.py b/megatron/core/distributed/__init__.py
new file mode 100644
index 0000000000..e43ae115ae
--- /dev/null
+++ b/megatron/core/distributed/__init__.py
@@ -0,0 +1,10 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from .distributed_data_parallel import DistributedDataParallel
+from .distributed_data_parallel_config import DistributedDataParallelConfig
+from .finalize_model_grads import finalize_model_grads
+
+# For backwards compatibility. ParamAndGradBuffer will be deprecated in future release.
+# ParamAndGradBuffer (which is an alias of _ParamAndGradBuffer) is not intended to be
+# consumed directly by external code.
+from .param_and_grad_buffer import ParamAndGradBuffer
diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py
new file mode 100644
index 0000000000..6e5bbd96d7
--- /dev/null
+++ b/megatron/core/distributed/distributed_data_parallel.py
@@ -0,0 +1,485 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+from contextlib import contextmanager
+
+import torch
+
+from .. import parallel_state
+from ..config_logger import has_config_logger_enabled, log_config_to_disk
+from ..transformer.module import MegatronModule
+from ..transformer.transformer_config import TransformerConfig
+from ..utils import is_float8tensor, log_single_rank
+from .distributed_data_parallel_config import DistributedDataParallelConfig
+from .param_and_grad_buffer import _ParamAndGradBuffer, partition_buckets
+
+logger = logging.getLogger(__name__)
+
+
+class DistributedDataParallel(MegatronModule):
+ """
+ DDP wrapper which stores grads in contiguous buffers. Also has option of overlapping
+ communication with backprop computation by breaking up full model's gradients into smaller
+ buckets and running all-reduce / reduce-scatter on each bucket asynchronously. This class
+ also provides the option to do the gradient accumulation in a type other than the param type
+ (e.g., fp32 for a bf16 model).
+
+ Args:
+ config: Transformer config object.
+ ddp_config: DistributedDataParallel config object.
+ module: Underlying model.
+ disable_bucketing: If true, force assign all parameters to a single bucket. If false,
+ use standard bucketing policy: assign parameters to smaller buckets and all-reduce
+ per bucket _if_ overlap_grad_reduce is True and pp_rank is 0.
+
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ ddp_config: DistributedDataParallelConfig,
+ module: torch.nn.Module,
+ disable_bucketing: bool = False,
+ ):
+ super().__init__(config=config)
+ if has_config_logger_enabled(config):
+ log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+ self.module = module
+
+ # If bucket_size is not provided as an input, use sane default.
+ # If using very large dp_sizes, make buckets larger to ensure that chunks used in NCCL
+ # ring-reduce implementations are large enough to remain bandwidth-bound rather than
+ # latency-bound.
+ if ddp_config.bucket_size is None:
+ ddp_config.bucket_size = max(
+ 40000000, 1000000 * parallel_state.get_data_parallel_world_size()
+ )
+ # Set bucket_size to infinity if overlap_grad_reduce is False.
+ if not ddp_config.overlap_grad_reduce:
+ ddp_config.bucket_size = None
+
+ self.ddp_config = ddp_config
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f'Setting up DistributedDataParallel with config {self.ddp_config}',
+ )
+
+ # Turn off bucketing if we are on a pipeline stage that is not the first (since
+ # data-parallel communication on these stages is not on the critical path), or if
+ # disable_bucketing is True (e.g., we might not want to break up model parameters
+ # into buckets for model chunks after the first in the interleaved schedule).
+ self.bucket_size = self.ddp_config.bucket_size
+ if parallel_state.get_pipeline_model_parallel_rank() > 0:
+ self.bucket_size = None
+ if disable_bucketing:
+ self.bucket_size = None
+
+ self.param_to_bucket_group = {}
+
+ # Group parameters by their gradient type.
+ param_to_name = {}
+ dense_params = []
+ expert_parallel_params = []
+ self.params_with_grad = []
+ for name, param in self.module.named_parameters():
+ if not param.requires_grad:
+ continue
+
+ # Track params with grad to enable direct setting
+ # of param.grad_added_to_main_grad
+ self.params_with_grad.append(param)
+
+ param.grad_added_to_main_grad = False
+ param_to_name[param] = name
+
+ if getattr(param, 'allreduce', True):
+ dense_params.append(param)
+ else:
+ expert_parallel_params.append(param)
+
+ def _allocate_buffers_for_parameters(
+ input_params, data_parallel_group, gradient_scaling_factor
+ ):
+ param_and_grad_dtype_to_params = {}
+ param_and_grad_dtype_to_offsets = {}
+ param_and_grad_dtype_to_indices = {}
+
+ # Group parameters by their gradient type.
+ for param in input_params:
+ assert param.requires_grad
+
+ param_dtype = param.dtype
+ if is_float8tensor(param):
+ # Currently TE's Float8Tensor is a wrapper of torch.Tensor. It has a "fake"
+ # dtype (usually a higher precision dtype such as bfloat16), but its actual
+ # data is stored in the form of a torch uint8 tensor within the Float8Tensor's
+ # ".data" attribute. Therefore, when creating the param buffer for fp8 params,
+ # it is necessary to use torch.uint8, not the "fake" dtype got from
+ # "param.dtype".
+ param_dtype = torch.uint8
+ grad_dtype = torch.float if self.ddp_config.grad_reduce_in_fp32 else param.dtype
+
+ params = param_and_grad_dtype_to_params.get((param_dtype, grad_dtype), [])
+ params.append(param)
+ param_and_grad_dtype_to_params[(param_dtype, grad_dtype)] = params
+
+ # Get the index of each param among the params with same dtype, if a param is fp8,
+ # use its "fake" high precision dtype to find which params have same dtype with it.
+ # For example:
+ # Case 1:
+ # params = [p1(bf16), p2(bf16), p3(bf16), p4(bf16)]
+ # param_and_grad_dtype_to_indices = {
+ # (torch.bfloat16, torch.float32): [0, 1, 2, 3],
+ # }
+ # Case 2:
+ # params = [p1(bf16), p2(fp8), p3(fp8), p4(bf16)]
+ # param_and_grad_dtype_to_indices = {
+ # (torch.bfloat16, torch.float32): [0, 3],
+ # (torch.uint8, torch.float32): [1, 2],
+ # }
+ # We need these indices to load a non-native-fp8 checkpoint in native-fp8 mode.
+ offset = param_and_grad_dtype_to_offsets.get((param.dtype, grad_dtype), 0)
+ param_and_grad_dtype_to_offsets[(param.dtype, grad_dtype)] = offset + 1
+ indices = param_and_grad_dtype_to_indices.get((param_dtype, grad_dtype), [])
+ indices.append(offset)
+ param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)] = indices
+
+ if not config.calculate_per_token_loss:
+ target_gradient_scaling_factor = 1.0 / parallel_state.get_data_parallel_world_size(
+ with_context_parallel=True
+ )
+ if self.ddp_config.average_in_collective:
+ # Collective is averaging gradients in collective with data_parallel_group.
+ assert (
+ gradient_scaling_factor
+ / torch.distributed.get_world_size(group=data_parallel_group)
+ == target_gradient_scaling_factor
+ )
+ else:
+ assert gradient_scaling_factor == target_gradient_scaling_factor
+
+ # Allocate the grad buffers and map the grads.
+ buffers = []
+ for (param_dtype, grad_dtype), params in param_and_grad_dtype_to_params.items():
+ buffers.append(
+ _ParamAndGradBuffer(
+ self.ddp_config,
+ param_dtype,
+ grad_dtype,
+ params,
+ data_parallel_group,
+ self.bucket_size,
+ param_to_name,
+ gradient_scaling_factor,
+ param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)],
+ )
+ )
+
+ # In some scenarios, we want to put buckets from different buffers into a group so that
+ # their communication can be aggregated. For example, when there are both fp8 buffers
+ # and bf16 buffers in the model and vpp is enabled, each model chunk will have an fp8
+ # bucket and a bf16 bucket, which doubles the number of communication kernels, and
+ # because of the use of CUDA_DEVICE_MAX_CONNECTIONS=1, having multiple back-to-back
+ # communications will prevent the overlap of the communication kernels with computation
+ # kernels.
+ # If bucketing is explicitly disabled, then put all buckets in a buffer into a single
+ # bucket group.
+ bucket_groups = partition_buckets(buffers, force_single_bucket_group=disable_bucketing)
+
+ # Set `next_param_gather_bucket_group` for different bucket groups by iterating through
+ # buckets in reverse order (since all-gathers happen in reverse order of buckets).
+ if self.ddp_config.use_distributed_optimizer and self.ddp_config.overlap_param_gather:
+ num_bucket_groups = len(bucket_groups)
+ for i in range(1, num_bucket_groups):
+ bucket_groups[num_bucket_groups - i].next_param_gather_bucket_group = (
+ bucket_groups[num_bucket_groups - i - 1]
+ )
+
+ # Create map from param to bucket group, used in pre_hook.
+ for bucket_group in bucket_groups:
+ for bucket in bucket_group.buckets:
+ for param in bucket.params_list:
+ self.param_to_bucket_group[param] = bucket_group
+
+ return buffers, bucket_groups
+
+ if config.calculate_per_token_loss:
+ gradient_scaling_factor = 1.0
+ expert_gradient_scaling_factor = 1.0
+ else:
+ if self.ddp_config.average_in_collective:
+ gradient_scaling_factor = 1.0
+ expert_gradient_scaling_factor = (
+ 1.0 / parallel_state.get_expert_model_parallel_world_size()
+ )
+ else:
+ data_parallel_world_size = parallel_state.get_data_parallel_world_size(
+ with_context_parallel=True
+ )
+ gradient_scaling_factor = 1.0 / data_parallel_world_size
+ expert_gradient_scaling_factor = 1.0 / data_parallel_world_size
+
+ # Allocate the param+grad buffers for dense params' grads.
+ self.buffers, self.bucket_groups = _allocate_buffers_for_parameters(
+ dense_params,
+ parallel_state.get_data_parallel_group(with_context_parallel=True),
+ gradient_scaling_factor=gradient_scaling_factor,
+ )
+
+ # Allocate separate param+grad buffers for expert parallel params' grads.
+ self.expert_parallel_buffers, self.expert_parallel_bucket_groups = (
+ _allocate_buffers_for_parameters(
+ expert_parallel_params,
+ parallel_state.get_data_modulo_expert_parallel_group(with_context_parallel=True),
+ gradient_scaling_factor=expert_gradient_scaling_factor,
+ )
+ )
+
+ # Delete references to weight_tensor if they exist since we don't want two parameter copies
+ # if we re-mapped parameters (which happens when we use the distributed optimizer).
+ # This is a temporary workaround around a TE bug that is fixed with
+ # https://github.com/NVIDIA/TransformerEngine/pull/719.
+ if self.ddp_config.use_distributed_optimizer:
+
+ @torch.no_grad()
+ def unmap_weight_tensor(m):
+ if hasattr(m, 'weight_tensor'):
+ m.weight_tensor = None
+
+ self.module.apply(unmap_weight_tensor)
+
+ # Register backward hook.
+ # Accumulation function for the gradients need to be stored so they
+ # don't go out of scope.
+ self.grad_accs = []
+ for param in self.module.parameters():
+ if param.requires_grad:
+ # Expand so we get access to grad_fn.
+ param_tmp = param.expand_as(param)
+ # Get the gradient accumulator function.
+ grad_acc = param_tmp.grad_fn.next_functions[0][0]
+ grad_acc.register_hook(self._make_backward_post_hook(param))
+ self.grad_accs.append(grad_acc)
+
+ self.use_forward_hook = (
+ self.ddp_config.use_distributed_optimizer and self.ddp_config.overlap_param_gather
+ )
+ self.remove_forward_pre_hook_handles = {}
+ if self.use_forward_hook:
+ self.enable_forward_pre_hook()
+ self.overlap_param_gather_with_optimizer_step = False
+
+ def enable_forward_pre_hook(self):
+ """
+ Enable forward pre-hooks needed for param all-gather overlap with forward compute.
+ """
+ assert self.use_forward_hook
+ assert len(self.remove_forward_pre_hook_handles) == 0
+ # Register forward pre-hook for all sub-modules.
+ for module in self.module.modules():
+ self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook(
+ self._make_forward_pre_hook()
+ )
+
+ def disable_forward_pre_hook(self):
+ """
+ Disable forward pre-hooks needed for param all-gather overlap with forward compute.
+ """
+ assert self.use_forward_hook
+ # De-register forward pre-hook for all sub-modules.
+ for module in self.module.modules():
+ assert self.remove_forward_pre_hook_handles[module] is not None
+ self.remove_forward_pre_hook_handles[module].remove()
+ del self.remove_forward_pre_hook_handles[module]
+ assert len(self.remove_forward_pre_hook_handles) == 0
+
+ # Force synchronize parameters.
+ self.start_param_sync(force_sync=True)
+
+ def forward(self, *inputs, **kwargs):
+ """
+ Calls the wrapped module's forward() method.
+ """
+ return self.module(*inputs, **kwargs)
+
+ def _make_forward_pre_hook(self):
+ """
+ Create a forward pre-hook to wait on all-gather handles when necessary (i.e.,
+ when a module uses a parameter in a bucket with a still incomplete all-gather).
+ """
+
+ def hook(module, *unused):
+ assert (
+ self.use_forward_hook
+ ), "Should use pre-hook only when overlap_param_gather is True"
+
+ # Make sure all parameters in this module have been all-gathered as necessary.
+ for param in module.parameters(recurse=False):
+ # Skip parameters without an associated buffer (such parameters have a
+ # .requires_grad field equal to False).
+ if param not in self.param_to_bucket_group:
+ continue
+ assert param.requires_grad
+
+ # If aligning param all-gather across pipeline stages, all-gather is dispatched
+ # by start_param_sync calls in core/pipeline_parallelism/schedules.py.
+ # If overlapping param all-gather with optimizer step, then all-gather has
+ # already been dispatched in optimizer step.
+ skip_next_bucket_dispatch = (
+ self.ddp_config.align_param_gather
+ or self.overlap_param_gather_with_optimizer_step
+ )
+ self.param_to_bucket_group[param].finish_param_sync(
+ skip_next_bucket_dispatch=skip_next_bucket_dispatch
+ )
+
+ return hook
+
+ def _make_backward_post_hook(self, param: torch.nn.Parameter):
+ """
+ Creates a backward post-hook to dispatch an all-reduce / reduce-scatter when
+ ready (i.e., when all grads in a bucket have been computed in all microbatches
+ in a batch).
+ """
+
+ def hook(*unused):
+ if param in self.param_to_bucket_group:
+ assert param.requires_grad
+ if self.ddp_config.overlap_grad_reduce:
+ assert (
+ param.grad is not None
+ ), 'param.grad being None is not safe when overlap_grad_reduce is True'
+ if param.grad is not None and (
+ not param.grad_added_to_main_grad or getattr(param, 'zero_out_wgrad', False)
+ ):
+ param.main_grad.add_(param.grad.data)
+ param.grad = None
+
+ if self.ddp_config.overlap_grad_reduce:
+ self.param_to_bucket_group[param].register_grad_ready(param)
+
+ return hook
+
+ @contextmanager
+ def no_sync(self):
+ """
+ Context manager that turns off gradient synchronization.
+ """
+ for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
+ bucket_group.is_last_microbatch = False
+ try:
+ yield
+ finally:
+ for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
+ bucket_group.is_last_microbatch = True
+
+ def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bool = False):
+ """
+ Initiates param sync (all-gather) communication operations for all model parameters.
+
+ By default, when overlap_param_gather is set to True, dispatches asynchronous communication
+ calls; when overlap_param_gather is set to False, calls synchronous communication
+ ops. Can override this default behavior using flags below.
+
+ Args:
+ force_sync (bool, optional): force synchronous collective regardless of
+ other settings.
+ force_dispatch (bool, optional): force dispatch regardless of other settings.
+ """
+ if not force_sync:
+ # If overlapping param AG with optimizer step, AG should not be dispatched again
+ # in forward_backward_step.
+ if self.overlap_param_gather_with_optimizer_step and not force_dispatch:
+ return
+
+ for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
+ bucket_group.start_param_sync(force_sync=force_sync)
+
+ def start_grad_sync(self, *unused):
+ """
+ Initiates grad sync (all-reduce or reduce-scatter) communication operations
+ for all model gradients.
+
+ When overlap_grad_reduce is set to True, dispatches asynchronous communication
+ calls. When overlap_grad_reduce is set to False, calls synchronous
+ communication ops.
+ """
+ for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
+ bucket_group.start_grad_sync()
+
+ def finish_grad_sync(self):
+ """
+ Finishes grad sync (all-reduce or reduce-scatter) communication operations
+ for all model gradients.
+
+ When overlap_grad_reduce is set to True, waits for asynchronous communication
+ calls to complete. When overlap_grad_reduce is set to False, calls synchronous
+ communication ops.
+ """
+ for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
+ bucket_group.finish_grad_sync()
+
+ def scale_gradients(self, scaling_factor: float):
+ """Scale all gradients inside the buffers by `scaling_factor`."""
+ for buffer in self.buffers + self.expert_parallel_buffers:
+ buffer.scale_gradients(scaling_factor)
+
+ def zero_grad_buffer(self):
+ """
+ Zeros out all grad buffers. Needs to be called at the beginning of each
+ training iteration.
+ """
+ for param in self.params_with_grad:
+ param.grad_added_to_main_grad = False
+ for buffer in self.buffers + self.expert_parallel_buffers:
+ buffer.reset()
+ for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
+ bucket_group.reset()
+
+ def broadcast_params(self):
+ """
+ Syncs parameters across all DP ranks.
+ """
+ for param in self.module.parameters():
+ is_expert_parallel = not getattr(param, 'allreduce', True)
+
+ if is_expert_parallel:
+ data_parallel_group = parallel_state.get_data_modulo_expert_parallel_group(
+ with_context_parallel=True
+ )
+ else:
+ data_parallel_group = parallel_state.get_data_parallel_group(
+ with_context_parallel=True
+ )
+ torch.distributed.broadcast(
+ param.data,
+ src=torch.distributed.get_global_rank(data_parallel_group, 0),
+ group=data_parallel_group,
+ )
+
+ def state_dict(self, prefix='', keep_vars=False):
+ """
+ Returns a dictionary containing references to the whole state of the
+ wrapped module.
+
+ Both parameters and persistent buffers (e.g. running averages) are included.
+ Keys are corresponding parameter and buffer names. Parameters and buffers
+ set to None are not included.
+ """
+ return self.module.state_dict(prefix=prefix, keep_vars=keep_vars)
+
+ def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False):
+ """
+ Returns wrapped module's state_dict for checkpoint saving.
+ """
+ return self.module.state_dict_for_save_checkpoint(prefix=prefix, keep_vars=keep_vars)
+
+ def load_state_dict(self, state_dict, strict=True):
+ """
+ Copies parameters and buffers from state_dict into the wrapped module and its
+ descendants. If strict is True, then the keys of state_dict must exactly match
+ the keys returned by this module’s state_dict() function.
+ """
+ self.module.load_state_dict(state_dict, strict=strict)
diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py
new file mode 100644
index 0000000000..14068ea367
--- /dev/null
+++ b/megatron/core/distributed/distributed_data_parallel_config.py
@@ -0,0 +1,44 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class DistributedDataParallelConfig:
+ """Configuration for DistributedDataParallel."""
+
+ grad_reduce_in_fp32: bool = False
+ """If true, reduce grads in fp32."""
+
+ overlap_grad_reduce: bool = False
+ """If true, overlap grad all-reduce / reduce-scatter with backward compute."""
+
+ overlap_param_gather: bool = False
+ """If true, overlap param all-gather with forward compute."""
+
+ align_param_gather: bool = False
+ """If true, all PP stages will launch param all-gathers simultaneously. Otherwise, each
+ PP stage will independently launch as needed.
+ """
+
+ use_distributed_optimizer: bool = False
+ """If true, issue reduce-scatter collectives to aggregate gradients and clean up
+ originally allocated model parameters, otherwise issue all-reduce collectives.
+ """
+
+ check_for_nan_in_grad: bool = False
+ """ If true, check for NaNs in gradients _before_ communication collective."""
+
+ bucket_size: Optional[int] = None
+ """Maximum number of parameters in each bucket. If unspecified, MCore uses a default
+ value of max(40000000, 1000000 * dp_size) parameters (larger DP sizes need larger
+ buckets to ensure collectives do not become latency-bound)."""
+
+ average_in_collective: bool = False
+ """If true, compute average in collective directly, as opposed to dividing by the
+ dp_size first and then computing sum in the collective."""
+
+ fp8_param_gather: bool = False
+ """If true, keep the compute param in fp8 (do not use any other intermediate dtype) and
+ perform the param all-gather in fp8."""
diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py
new file mode 100644
index 0000000000..ff5046afa5
--- /dev/null
+++ b/megatron/core/distributed/finalize_model_grads.py
@@ -0,0 +1,162 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from typing import List, Optional
+
+import torch
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+from .. import parallel_state
+from ..transformer.transformer_config import TransformerConfig
+from ..utils import get_attr_wrapped_model, get_model_config
+
+
+def _allreduce_word_embedding_grads(model: List[torch.nn.Module], config: TransformerConfig):
+ """
+ All-reduce word embedding grads.
+
+ Reduce grads across first and last stages to ensure that word_embeddings parameters stay in
+ sync.
+ """
+
+ if (
+ parallel_state.is_rank_in_embedding_group(ignore_virtual=True)
+ and torch.distributed.get_world_size(parallel_state.get_embedding_group()) > 1
+ ):
+ if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
+ model_module = model[0]
+ elif parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+ model_module = model[-1]
+ else: # We do not support an interleaved schedule for models with encoders yet.
+ model_module = model[0]
+
+ model_module = get_attr_wrapped_model(model_module, 'pre_process', return_model_obj=True)
+ if model_module.share_embeddings_and_output_weights:
+ weight = model_module.shared_embedding_or_output_weight()
+ grad = weight.main_grad
+ torch.distributed.all_reduce(grad, group=parallel_state.get_embedding_group())
+
+
+def _allreduce_position_embedding_grads(model: List[torch.nn.Module], config: TransformerConfig):
+ """
+ All-reduce position_embeddings grad across encoder and decoder stages to ensure that position
+ embeddings parameters stay in sync.
+ """
+ if (
+ parallel_state.is_rank_in_position_embedding_group()
+ and torch.distributed.get_world_size(parallel_state.get_position_embedding_group()) > 1
+ ):
+ if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
+ model_module = model[0]
+ elif parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+ model_module = model[-1]
+ else: # We do not support an interleaved schedule for models with encoders yet.
+ model_module = model[0]
+
+ model_module = get_attr_wrapped_model(model_module, 'pre_process', return_model_obj=True)
+ assert hasattr(model_module, 'position_embeddings')
+ grad = model_module.position_embeddings.weight.main_grad
+ torch.distributed.all_reduce(grad, group=parallel_state.get_position_embedding_group())
+
+
+def _allreduce_embedding_grads(model: List[torch.nn.Module], config: TransformerConfig):
+ """
+ All-reduce both word and position embeddings.
+ """
+ _allreduce_word_embedding_grads(model, config)
+ _allreduce_position_embedding_grads(model, config)
+
+
+def _allreduce_layernorm_grads(model: List[torch.nn.Module], config: TransformerConfig):
+ """
+ All-reduce layernorm grads (for sequence parallelism).
+ """
+
+ # All-reduce layernorm parameters across model parallel nodes
+ # when sequence parallelism is used
+ if parallel_state.get_tensor_model_parallel_world_size() > 1 and (
+ config.sequence_parallel or config.qk_layernorm
+ ):
+ grads = []
+ for model_chunk in model:
+ for name, param in get_attr_wrapped_model(model_chunk, 'named_parameters')():
+ if (
+ param.requires_grad
+ and getattr(param, 'sequence_parallel', False)
+ or 'q_layernorm' in name
+ or 'k_layernorm' in name
+ ):
+ grad = param.main_grad
+ grads.append(grad.data)
+ if grads:
+ coalesced = _flatten_dense_tensors(grads)
+ torch.distributed.all_reduce(
+ coalesced, group=parallel_state.get_tensor_model_parallel_group()
+ )
+ for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):
+ buf.copy_(synced)
+
+
+def finalize_model_grads(model: List[torch.nn.Module], num_tokens: Optional[torch.Tensor] = None):
+ """
+ All-reduce all model grads across DP replicas, layernorm grads for sequence parallelism,
+ embedding grads across first and last pipeline stages (if not tied),
+ scale gradients by `num_tokens`.
+ """
+
+ config = get_model_config(model[0])
+
+ # All-reduce / reduce-scatter across DP replicas.
+ if config.timers is not None:
+ config.timers('all-grads-sync', log_level=1).start(barrier=config.barrier_with_L1_time)
+ for model_chunk in model:
+ model_chunk.finish_grad_sync()
+ if config.timers is not None:
+ config.timers('all-grads-sync').stop()
+
+ # All-reduce layer-norm grads (for sequence parallelism).
+ if config.timers is not None:
+ config.timers('layernorm-grads-all-reduce', log_level=1).start(
+ barrier=config.barrier_with_L1_time
+ )
+ _allreduce_layernorm_grads(model, config)
+ if config.timers is not None:
+ config.timers('layernorm-grads-all-reduce').stop()
+
+ # All-reduce embedding grads (for pipeline parallelism).
+ if config.timers is not None:
+ config.timers('embedding-grads-all-reduce', log_level=1).start(
+ barrier=config.barrier_with_L1_time
+ )
+ _allreduce_embedding_grads(model, config)
+ if config.timers is not None:
+ config.timers('embedding-grads-all-reduce').stop()
+
+ # normalize gradients for per-token loss normalization.
+ # if we are using by the number of tokens, then we use that as a divisor. this number
+ # will be the total number of non-padded tokens in the global batch.
+ if num_tokens is not None:
+
+ # the number of tokens is only present on the last stage, so broadcast it
+ # to the other ranks in the pipeline parallel group.
+ last_rank = parallel_state.get_pipeline_model_parallel_last_rank()
+ pp_group = parallel_state.get_pipeline_model_parallel_group()
+
+ if not isinstance(last_rank, list):
+ assert not isinstance(last_rank, list)
+ last_rank = [last_rank]
+ assert not isinstance(pp_group, list)
+ pp_group = [pp_group]
+
+ # need to do a broadcast for every pp group, even though num_tokens should be the same.
+ num_tokens_list = []
+ for lr, group in zip(last_rank, pp_group):
+ torch.distributed.broadcast(num_tokens, src=lr, group=group)
+ num_tokens_list.append(torch.clone(num_tokens))
+ assert all(x.item() == num_tokens_list[0] for x in num_tokens_list)
+
+ # all-reduce across DP ranks.
+ torch.distributed.all_reduce(num_tokens, group=parallel_state.get_data_parallel_group())
+ for model_chunk in model:
+ if num_tokens > 0:
+ scaling = 1.0 / num_tokens
+ model_chunk.scale_gradients(scaling)
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
new file mode 100644
index 0000000000..351ff9e0bf
--- /dev/null
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -0,0 +1,769 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+import math
+import os
+import warnings
+from enum import Enum
+from typing import Dict, List, Optional
+
+import torch
+from torch.distributed import _coalescing_manager
+
+from ..utils import is_float8tensor, log_on_each_pipeline_stage
+from .distributed_data_parallel_config import DistributedDataParallelConfig
+
+logger = logging.getLogger(__name__)
+
+
+class BufferType(Enum):
+ """
+ Enumeration for buffer type.
+ """
+
+ PARAM = 1
+ GRAD = 2
+
+
+def shard_buffer(buffer: torch.Tensor, data_parallel_world_size: int):
+ """
+ Shard buffer into data_parallel_world_size chunks of equal size.
+ """
+ assert buffer.numel() % data_parallel_world_size == 0
+ shard_size = buffer.numel() // data_parallel_world_size
+ sharded_buffer = [
+ buffer[(r * shard_size) : ((r + 1) * shard_size)] for r in range(data_parallel_world_size)
+ ]
+ return sharded_buffer
+
+
+class _ParamAndGradBucket:
+ """
+ Bucket to keep track of a subset of the model's parameters and gradients.
+
+ Args:
+ params: List of parameters whose gradients are collated in this bucket.
+ param_data: View in ParamAndGradBuffer.param_data that this bucket is responsible for.
+ grad_data: View in ParamAndGradBuffer.grad_data that this bucket is responsible for.
+ offset: Offset of this bucket's view in the larger ParamAndGradBuffer.
+ numel_unpadded: Number of unpadded elements in bucket.
+ gradient_scaling_factor: This factor is utilized to scale gradients prior to their
+ communication. Its application is twofold: it facilitates the averaging of gradients
+ and the scaling of gradients in the context of the Mixture of Experts (MoE) model.
+ bucket_id: Index of bucket in buffer.
+ """
+
+ def __init__(
+ self,
+ params: List[torch.nn.Parameter],
+ param_data: Optional[torch.Tensor],
+ grad_data: torch.Tensor,
+ offset: int,
+ numel_unpadded: int,
+ gradient_scaling_factor: float,
+ bucket_id: int,
+ ):
+ self.params_list = params
+ self.params = set(params)
+ # Make sure there are no duplicate params.
+ assert len(self.params_list) == len(self.params)
+ self.param_data = param_data
+ self.grad_data = grad_data
+ # The distributed optimizer needs to keep track of this bucket's offset
+ # within the full grad_buffer.
+ self.offset = offset
+ self.numel_unpadded = numel_unpadded
+ self.gradient_scaling_factor = gradient_scaling_factor
+ self.bucket_id = bucket_id
+
+
+class _ParamAndGradBucketGroup:
+ """
+ Put multiple buckets into a group so that their communications can be aggregated together.
+ Provides functionality to register when params in the bucket group have grads ready to be
+ synced; an asynchronous communication call is automatically launched when _all_ params in
+ the bucket group have grads ready.
+
+ Args:
+ buckets: A list of buckets.
+ ddp_config: DistributedDataParallel config object.
+ data_parallel_group: Data-parallel process group.
+ data_parallel_world_size: World size using the data-parallel group group.
+ """
+
+ def __init__(
+ self,
+ buckets: List[_ParamAndGradBucket],
+ ddp_config: DistributedDataParallelConfig,
+ data_parallel_group: torch.distributed.ProcessGroup,
+ data_parallel_world_size: int,
+ ):
+ self.buckets = buckets
+ self.ddp_config = ddp_config
+ self.data_parallel_group = data_parallel_group
+ self.data_parallel_world_size = data_parallel_world_size
+ self.data_parallel_rank = torch.distributed.get_rank(group=data_parallel_group)
+
+ # State for bookkeeping: params is the set of parameters this bucket group is
+ # responsible for, params_with_grad is the set of parameters with grads
+ # available. When overlap_grad_reduce is True, communication (all-reduce
+ # or reduce-scatter) is issued when params_with_grad equals params.
+ self.param_to_bucket = {}
+ self.params = set()
+ for bucket in self.buckets:
+ for param in bucket.params_list:
+ self.param_to_bucket[param] = bucket
+ self.params.add(param)
+
+ self.next_param_gather_bucket_group = None
+
+ self.reset()
+ self.param_gather_handle = None
+ self.param_gather_dispatched = False
+ self.grad_reduce_handle = None
+
+ def reset(self):
+ """
+ Reset metadata in bucket group in preparation for the next iteration of training.
+ """
+ self.params_with_grad = set()
+ self.is_last_microbatch = True
+
+ def check_for_nan_in_grad(self):
+ """
+ Make sure norm of grads in bucket are not NaN prior to data-parallel
+ all-reduce / reduce-scatter.
+ """
+ global_rank = torch.distributed.get_rank()
+ norm_is_nan = self.buckets[0].grad_data.norm(p=2).isnan()
+ for i in range(1, len(self.buckets)):
+ norm_is_nan.logical_or_(self.buckets[i].grad_data.norm(p=2).isnan())
+ assert not norm_is_nan, (
+ f'Rank {global_rank}: found NaN in local grad norm in '
+ f'backward pass before data-parallel communication collective. '
+ f'Device: {torch.cuda.current_device()}, node: {os.uname()[1]}'
+ )
+
+ def start_param_sync(self, force_sync: bool = False):
+ """
+ Initiates all necessary param all-gathers for this bucket.
+
+ When ddp_config.overlap_param_gather is set to True, dispatches an asynchronous
+ communication call (unless force_sync is True). When ddp_config.overlap_param_gather
+ is set to False, makes synchronous call.
+
+ Args:
+ force_sync (bool, optional): force synchronous collective regardless of
+ other settings if true.
+ """
+ assert self.ddp_config.use_distributed_optimizer
+
+ if force_sync:
+ if self.param_gather_handle is not None:
+ self.param_gather_handle.wait()
+ self.param_gather_handle = None
+ return
+ else:
+ assert self.param_gather_handle is None
+
+ async_op = self.ddp_config.overlap_param_gather and not force_sync
+ # Coalesce communication kernels across buckets in the bucket group.
+ with _coalescing_manager(self.data_parallel_group, async_ops=async_op) as cm:
+ for bucket in self.buckets:
+ local_data_view = shard_buffer(bucket.param_data, self.data_parallel_world_size)[
+ self.data_parallel_rank
+ ]
+ torch.distributed._all_gather_base(
+ bucket.param_data,
+ local_data_view,
+ group=self.data_parallel_group,
+ async_op=async_op,
+ )
+ if async_op:
+ self.param_gather_handle = cm
+ else:
+ # When using `_coalescing_manager`, even if a synchronous op (async_op=False) is used,
+ # `cm` is not None, which is different from when `_coalescing_manager` is not used in
+ # which case the torch.distributed._all_gather_base() will return None. In order to
+ # maintain consistency with prior code, we need to manually set communication handle to
+ # None.
+ self.param_gather_handle = None
+ self.param_gather_dispatched = True
+
+ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
+ """
+ Finishes param sync communication operation for this bucket. Dispatches
+ next bucket's param sync if available, unless skip_next_bucket_dispatch
+ is True.
+
+ When ddp_config.overlap_param_gather is set to True, waits for asynchronous
+ communication call to complete (and dispatches one if one is not already
+ outstanding). Throws assertion error if ddp_config.overlap_param_gather is set to
+ False.
+
+ Args:
+ skip_next_bucket_dispatch (bool, optional): if true, dispatch next
+ bucket's communication if available.
+ """
+ assert self.ddp_config.use_distributed_optimizer
+ assert self.ddp_config.overlap_param_gather
+
+ # If current bucket's param AG has not been dispatched, dispatch it now (e.g., first
+ # AG bucket in first model chunk if ddp_config.align_param_gather is False).
+ if not self.param_gather_dispatched:
+ self.start_param_sync()
+
+ if self.param_gather_handle is not None:
+ self.param_gather_handle.wait()
+ self.param_gather_handle = None
+ # Dispatch next bucket's asynchronous param AG.
+ if self.next_param_gather_bucket_group is not None and not skip_next_bucket_dispatch:
+ self.next_param_gather_bucket_group.start_param_sync()
+
+ def start_grad_sync(self):
+ """
+ Initiates grad sync (all-reduce or reduce-scatter) communication operations
+ for all buckets in the bucket group.
+
+ When ddp_config.overlap_grad_reduce is set to True, dispatches an asynchronous
+ communication call. When ddp_config.overlap_grad_reduce is set to False, makes
+ synchronous call.
+ """
+ assert (
+ self.grad_reduce_handle is None
+ ), 'Should not have multiple communication calls outstanding at once'
+
+ if self.ddp_config.check_for_nan_in_grad:
+ self.check_for_nan_in_grad()
+
+ # gradient_scaling_factor already takes into account whether we are computing
+ # an average or sum in the data-parallel collective.
+ for bucket in self.buckets:
+ if bucket.gradient_scaling_factor != 1.0:
+ bucket.grad_data *= bucket.gradient_scaling_factor
+
+ # Decide reduce_op.
+ reduce_op = torch.distributed.ReduceOp.SUM
+ if self.ddp_config.average_in_collective:
+ reduce_op = torch.distributed.ReduceOp.AVG
+
+ # Use async communications only when overlap_grad_reduce is True.
+ async_op = self.ddp_config.overlap_grad_reduce
+ # Coalesce communication kernels across buckets in the bucket group.
+ with _coalescing_manager(self.data_parallel_group, async_ops=async_op) as cm:
+ for bucket in self.buckets:
+ if self.ddp_config.use_distributed_optimizer:
+ local_data_view = shard_buffer(bucket.grad_data, self.data_parallel_world_size)[
+ self.data_parallel_rank
+ ]
+ torch.distributed._reduce_scatter_base(
+ local_data_view,
+ bucket.grad_data,
+ op=reduce_op,
+ group=self.data_parallel_group,
+ async_op=async_op,
+ )
+ else:
+ torch.distributed.all_reduce(
+ bucket.grad_data,
+ op=reduce_op,
+ group=self.data_parallel_group,
+ async_op=async_op,
+ )
+ if async_op:
+ self.grad_reduce_handle = cm
+ else:
+ # When using `_coalescing_manager`, even if a synchronous op (async_op=False) is used,
+ # `cm` is not None, which is different from when `_coalescing_manager` is not used in
+ # which case the torch.distributed._reduce_scatter_base() will return None. In order to
+ # maintain consistency with prior code, we need to manually set communication handle to
+ # None.
+ self.grad_reduce_handle = None
+
+ def finish_grad_sync(self):
+ """
+ Finishes grad sync (all-reduce or reduce-scatter) communication operations
+ for all buckets in the bucket group.
+
+ When ddp_config.overlap_grad_reduce is set to True, waits for asynchronous
+ communication call to complete. When ddp_config.overlap_grad_reduce is set to False,
+ makes synchronous call.
+ """
+ # If overlap_grad_reduce is False, start (and finish) synchronous communication call here.
+ self.param_gather_dispatched = False
+ if not self.ddp_config.overlap_grad_reduce:
+ self.start_grad_sync()
+ return
+ assert self.grad_reduce_handle is not None, (
+ f'Communication call has not been issued for this bucket '
+ f'({len(self.params_with_grad)}/{len(self.params)} params have grad available)'
+ )
+ self.grad_reduce_handle.wait()
+ self.grad_reduce_handle = None
+
+ def register_grad_ready(self, param: torch.nn.Parameter):
+ """
+ Registers grads for the passed-in param to be "ready" for grad sync.
+
+ When the number of microbatches is greater than 1, we only want to register
+ grads as ready when processing the last microbatch and ddp_config.overlap_grad_reduce
+ is True.
+ """
+ assert (
+ self.ddp_config.overlap_grad_reduce
+ ), 'register_grad_ready() should only be called when overlap_grad_reduce is True'
+ if self.is_last_microbatch:
+ assert param in self.param_to_bucket, 'Param is not in the bucket group'
+ assert param not in self.params_with_grad, 'Cannot set grad twice'
+ self.params_with_grad.add(param)
+ # If all params in bucket group have grads available, issue communication call.
+ if len(self.params_with_grad) == len(self.params):
+ self.start_grad_sync()
+
+
+class _ParamAndGradBuffer:
+ """
+ Groups parameters and gradients into a contiguous buffer, and then breaks the buffer into
+ buckets with roughly `bucket_size` parameters each.
+
+ Args:
+ ddp_config: DistributedDataParallel config object.
+ param_dtype: Type of param tensor.
+ grad_dtype: Type of grad tensor.
+ params: List of parameters whose parameters and gradients are collated in the underlying
+ tensor.
+ data_parallel_group: Data-parallel process group.
+ bucket_size: The rough size of each bucket in terms of number of parameters.
+ param_to_name: Mapping from `torch.nn.Parameter` to name (for logging purposes).
+ gradient_scaling_factor: This factor is utilized to scale gradients prior to their
+ communication. Its application is twofold: it facilitates the averaging of gradients
+ and the scaling of gradients in the context of the Mixture of Experts (MoE) model.
+ param_indices: The index of each param among the params with same dtype, if a param is fp8,
+ use its "fake" high precision dtype to determine which params have same dtype with it.
+ These indices are needed when loading a non-native-fp8 checkpoint in native-fp8 mode.
+ """
+
+ def __init__(
+ self,
+ ddp_config: DistributedDataParallelConfig,
+ param_dtype: torch.dtype,
+ grad_dtype: torch.dtype,
+ params: List[torch.nn.Parameter],
+ data_parallel_group: torch.distributed.ProcessGroup,
+ bucket_size: int,
+ param_to_name: Dict[torch.nn.Parameter, str],
+ gradient_scaling_factor: float,
+ param_indices: List[int],
+ ):
+ self.ddp_config = ddp_config
+ self.params = params
+ self.param_indices = param_indices
+
+ # Check that params are unique.
+ unique_params = set()
+ for param in params:
+ assert param not in unique_params
+ unique_params.add(param)
+ del unique_params
+
+ # Store attributes that will be needed later.
+ self.param_dtype = param_dtype
+ self.grad_dtype = grad_dtype
+ self.data_parallel_group = data_parallel_group
+ self.data_parallel_world_size = torch.distributed.get_world_size(
+ group=self.data_parallel_group
+ )
+ self.gradient_scaling_factor = gradient_scaling_factor
+
+ # Data structures to store underlying buckets and relevant indexing data.
+ self.buckets = []
+ self.param_to_bucket = {} # Param -> bucket mapping.
+ self.param_index_map = {} # Param -> location in buffer mapping (used in dist. optimizer).
+
+ def _pad(number_to_be_padded: int, divisor: int) -> int:
+ return int(math.ceil(number_to_be_padded / divisor) * divisor)
+
+ def _pad_end_of_bucket_if_needed(bucket_end_index: int) -> int:
+ """
+ Pads end index of bucket if using distributed optimizer (to ensure uniform sharding).
+ """
+ if self.ddp_config.use_distributed_optimizer:
+ # Workaround for TE bug causing cuBLAS to pick an incompatible algorithm.
+ # This also helps cuBLAS pick more efficient algorithms for GEMMs.
+ # We now ensure that all buckets start at a memory address that is 256-byte
+ # aligned (128 values since params and grads use >= 16-bit precision).
+ return _pad(bucket_end_index, math.lcm(self.data_parallel_world_size, 128))
+ return bucket_end_index
+
+ def _pad_start_of_param_if_needed(param_start_index: int) -> int:
+ """
+ Pads start index of param if using distributed optimizer (to ensure "good" alignment).
+ """
+ if self.ddp_config.use_distributed_optimizer:
+ # Ensure that params start at 128-byte aligned addresses (64 values
+ # since params are >= 16-bit precision).
+ return _pad(param_start_index, 64)
+ return param_start_index
+
+ # First, figure out how many elements should be in the underlying buffer storage.
+ # Note that if we need to split the buffer into smaller buckets, each of these
+ # might need to be padded as well (if using the distributed optimizer).
+ param_start_index = 0
+ bucket_start_index = param_start_index
+ bucket_params = set()
+ self.bucket_indices = []
+ per_bucket_numel_unpadded = []
+ bucket_id = 0
+
+ def _update_bucket_metadata(param_end_index: int) -> int:
+ """
+ Record metadata for the bucket starting at bucket_start_index and ending with the
+ passed-in param_end_index. Returns the bucket's end_index.
+ """
+ nonlocal bucket_start_index, bucket_params, bucket_id
+ per_bucket_numel_unpadded.append(param_end_index - bucket_start_index)
+ bucket_end_index = _pad_end_of_bucket_if_needed(param_end_index)
+
+ # Record metadata of new bucket.
+ self.bucket_indices.append((bucket_start_index, bucket_end_index))
+ bucket_start_index = bucket_end_index
+
+ # Prepare for next bucket.
+ bucket_params = set()
+ bucket_id += 1
+
+ # Return the potentially padded bucket_end_index.
+ return bucket_end_index
+
+ def _does_param_require_new_bucket(param):
+ """
+ Split shared embedding parameters into separate bucket if using distributed
+ optimizer that makes use of reduce-scatters instead of all-reduces.
+ This ensures that the first and last pipeline stage partition optimizer state
+ for the shared embedding parameters the same way across DP replicas, allowing
+ the DP reduce-scatter to be before the embedding all-reduce.
+ """
+ return (
+ getattr(param, "shared_embedding", False)
+ and self.ddp_config.use_distributed_optimizer
+ )
+
+ for param in params[::-1]:
+ # Iterate through parameters in reverse order to roughly follow backprop order.
+
+ this_numel = param.data.nelement()
+ param_start_index = _pad_start_of_param_if_needed(param_start_index)
+
+ # Create bucket with collected parameters if current param needs its own bucket.
+ if _does_param_require_new_bucket(param):
+ # We are creating a bucket for the already accumulated parameters, whose params
+ # end at the current param_start_index.
+ if self.ddp_config.use_distributed_optimizer:
+ # Make sure new bucket is appropriately padded.
+ if param_start_index % self.data_parallel_world_size != 0:
+ param_start_index = _pad_end_of_bucket_if_needed(param_start_index)
+ if len(bucket_params) > 0:
+ bucket_end_index = _update_bucket_metadata(param_start_index)
+
+ param_end_index = param_start_index + this_numel
+ self.param_index_map[param] = (param_start_index, param_end_index, bucket_id)
+ bucket_params.add(param)
+
+ # If we have enough elements already or the current param is part of the shared
+ # embedding layer and needs a separate bucket, form a new bucket.
+ if (
+ bucket_size is not None and (param_end_index - bucket_start_index) >= bucket_size
+ ) or _does_param_require_new_bucket(param):
+ bucket_end_index = _update_bucket_metadata(param_end_index)
+ param_start_index = bucket_end_index
+ else:
+ param_start_index = param_end_index
+
+ # Add remaining params to a new bucket.
+ if len(bucket_params) > 0:
+ bucket_end_index = _update_bucket_metadata(param_end_index)
+
+ # Next, create underlying storage for buffer (with numel elements that includes
+ # padding as necessary).
+ self.numel = bucket_end_index
+ self.numel_unpadded = sum(per_bucket_numel_unpadded)
+ assert self.numel_unpadded <= self.numel
+ if self.ddp_config.use_distributed_optimizer:
+ assert self.numel % self.data_parallel_world_size == 0
+ else:
+ assert self.numel == self.numel_unpadded
+
+ self.param_data = None
+ # Only re-map param tensors if using distributed optimizer.
+ if self.ddp_config.use_distributed_optimizer:
+ self.param_data = torch.zeros(
+ self.numel,
+ dtype=self.param_dtype,
+ device=torch.cuda.current_device(),
+ requires_grad=False,
+ )
+ self.grad_data = torch.zeros(
+ self.numel,
+ dtype=self.grad_dtype,
+ device=torch.cuda.current_device(),
+ requires_grad=False,
+ )
+
+ # Finally, map param.data and param.main_grad fields to buffers.
+ bucket_params = []
+ bucket_start_index = 0
+ cur_bucket_id = 0
+ for param in params[::-1]:
+ param_start_index, param_end_index, bucket_id = self.param_index_map[param]
+
+ # Assign param.data to appropriate segment of self.param_data.
+ if self.param_data is not None:
+ old_param_data = param.data
+ new_param_data = self._get(
+ param.data.shape, param_start_index, buffer_type=BufferType.PARAM
+ )
+ if is_float8tensor(param):
+ param._data = new_param_data
+ else:
+ param.data = new_param_data
+ assert old_param_data._base is None
+ # Copy tensor values (from initialization or checkpoint).
+ param.data.detach().copy_(old_param_data)
+ del old_param_data
+
+ param.main_grad = self._get(
+ param.data.shape, param_start_index, buffer_type=BufferType.GRAD
+ )
+ if bucket_id != cur_bucket_id:
+ bucket_end_index = _pad_end_of_bucket_if_needed(param_start_index)
+ self.buckets.append(
+ self._new_bucket(
+ bucket_params=bucket_params,
+ start_index=bucket_start_index,
+ end_index=bucket_end_index,
+ numel_unpadded=per_bucket_numel_unpadded[cur_bucket_id],
+ bucket_id=cur_bucket_id,
+ )
+ )
+ bucket_start_index = bucket_end_index
+ bucket_params = []
+ assert cur_bucket_id + 1 == len(self.buckets)
+ assert bucket_id == cur_bucket_id + 1
+ cur_bucket_id = bucket_id
+ bucket_params.append(param)
+
+ # Add remaining params to a new bucket.
+ if len(bucket_params) > 0:
+ bucket_end_index = _pad_end_of_bucket_if_needed(param_end_index)
+ self.buckets.append(
+ self._new_bucket(
+ bucket_params=bucket_params,
+ start_index=bucket_start_index,
+ end_index=bucket_end_index,
+ numel_unpadded=per_bucket_numel_unpadded[cur_bucket_id],
+ bucket_id=cur_bucket_id,
+ )
+ )
+
+ # Log buckets for all PP stages.
+ log_strs = []
+ log_strs.append(
+ f'Number of buckets for gradient all-reduce / reduce-scatter: {len(self.buckets)}'
+ )
+ for index, bucket in enumerate(self.buckets):
+ numel = 0
+ for param in bucket.params:
+ numel += param.data.nelement()
+ log_strs.append(f'Params for bucket {index+1} ({numel} elements):')
+ for param in bucket.params:
+ log_strs.append(f'\t{param_to_name[param]}')
+ log_on_each_pipeline_stage(logger, logging.INFO, '\n'.join(log_strs))
+
+ def scale_gradients(self, scaling_factor: float) -> None:
+ """Scale the gradient data by `scaling_factor`."""
+ self.grad_data *= scaling_factor
+
+ def _get(self, shape: torch.Size, start_index: int, buffer_type: BufferType) -> torch.Tensor:
+ """
+ Return a tensor with the input `shape` as a view into the 1-D data starting at
+ `start_index`.
+ """
+ end_index = start_index + shape.numel()
+ assert end_index <= self.numel, 'Requested tensor is out of buffer range'
+ if buffer_type == BufferType.PARAM:
+ assert self.param_data is not None
+ buffer_tensor = self.param_data[start_index:end_index]
+ elif buffer_type == BufferType.GRAD:
+ buffer_tensor = self.grad_data[start_index:end_index]
+ else:
+ raise Exception("Illegal buffer type provided to GradBuffer._get() function")
+ buffer_tensor = buffer_tensor.view(shape)
+ return buffer_tensor
+
+ def _new_bucket(
+ self,
+ bucket_params: List[torch.nn.Parameter],
+ start_index: int,
+ end_index: int,
+ numel_unpadded: int,
+ bucket_id: int,
+ ) -> _ParamAndGradBucket:
+ """
+ Helper function that creates a new bucket. Also updates param->bucket mapping.
+ """
+
+ # Assert that indices are correctly padded (if needed), and that bucket
+ # position is same as originally computed.
+ if self.ddp_config.use_distributed_optimizer:
+ assert start_index % self.data_parallel_world_size == 0
+ assert end_index % self.data_parallel_world_size == 0
+ assert (start_index, end_index) == self.bucket_indices[bucket_id]
+
+ # Get appropriate view into global ParamAndGradBuffer.
+ bucketed_param_data = None
+ if self.param_data is not None:
+ bucketed_param_data = self._get(
+ torch.Size([end_index - start_index]), start_index, buffer_type=BufferType.PARAM
+ )
+ bucketed_grad_data = self._get(
+ torch.Size([end_index - start_index]), start_index, buffer_type=BufferType.GRAD
+ )
+ bucket = _ParamAndGradBucket(
+ params=bucket_params,
+ param_data=bucketed_param_data,
+ grad_data=bucketed_grad_data,
+ offset=start_index,
+ numel_unpadded=numel_unpadded,
+ gradient_scaling_factor=self.gradient_scaling_factor,
+ bucket_id=bucket_id,
+ )
+ for bucket_param in bucket_params:
+ assert bucket_param not in self.param_to_bucket
+ self.param_to_bucket[bucket_param] = bucket
+
+ return bucket
+
+ def reset(self):
+ """
+ Zero out the underlying grad_buffer.
+ """
+ self.grad_data.zero_()
+
+
+def partition_buckets(
+ buffers: List[_ParamAndGradBuffer], force_single_bucket_group: bool = False
+) -> List[_ParamAndGradBucketGroup]:
+ """
+ Automatically regroup the buckets of input buffers and return a list of bucket groups.
+
+ In some scenarios, we need to put buckets from different buffers into a group so that their
+ communication can be aggregated.
+
+ For example, when there are both fp8 weights and bf16 biases in the model and virtual
+ pipeline parallelism is enabled, each model chunk will have an fp8 bucket and a bf16 bucket,
+ which doubles the number of communication kernels, and because of the use of
+ CUDA_DEVICE_MAX_CONNECTIONS=1, having multiple back-to-back communications will prevent the
+ overlap of communication kernels with computation kernels.
+
+ The grouping strategy is:
+ 1. If force_single_bucket_group is True, put all buckets across all buffers into a single
+ bucket group.
+ 2. If force_single_bucket_group is False, when there is no fp8 buffer in the input buffers,
+ let each bucket group have only one bucket.
+ 3. If force_single_bucket_group is False, when using fp8 params, merge all non-fp8 buckets
+ into the last fp8 bucket group.
+ - Since the non-fp8 parameters (typically the biases of various layers) are relatively
+ small, they are likely to be grouped into a single non-fp8 bucket.
+ - The fp8 buckets start from the end of the model, i.e., the first bucket corresponds to
+ the end of the model, while the last bucket corresponds to the beginning.
+ - If we combine the non-fp8 bucket with the first fp8 bucket, we cannot initiate the
+ reduce-scatter to synchronize gradients after the backward pass at the end of the model
+ has completed. This is because we need to wait for the non-fp8 params from the beginning
+ layers to obtain their gradients.
+ - Combining the non-fp8 bucket with the last fp8 bucket can help avoid this issue.
+
+ Args:
+ buffers (list): list of input buffers.
+ single_bucket_group_per_buffer (bool, optional): force group all buckets in each buffer
+ into a single bucket group.
+ """
+
+ if len(buffers) == 0:
+ return []
+
+ dtype_to_buffer_map = {}
+ for buffer in buffers:
+ dtype = buffer.param_dtype
+ # Make sure that the param_dtype of any two buffers is different.
+ assert dtype not in dtype_to_buffer_map
+ dtype_to_buffer_map[dtype] = buffer
+
+ # Case 1: Put all buckets into a single bucket group if force_single_bucket_group is True.
+ if force_single_bucket_group:
+ buckets = []
+ ddp_config = buffers[0].ddp_config
+ data_parallel_group = buffers[0].data_parallel_group
+ data_parallel_world_size = buffers[0].data_parallel_world_size
+ for buffer in buffers:
+ assert ddp_config == buffer.ddp_config
+ assert data_parallel_group == buffer.data_parallel_group
+ assert data_parallel_world_size == buffer.data_parallel_world_size
+ buckets.extend(buffer.buckets)
+
+ bucket_group = _ParamAndGradBucketGroup(
+ buckets, ddp_config, data_parallel_group, data_parallel_world_size
+ )
+ return [bucket_group]
+
+ if torch.uint8 not in dtype_to_buffer_map:
+ # Case 2: When there is no fp8 buffer in the input buffers, let each bucket group have
+ # only one bucket.
+ bucket_groups = []
+ for buffer in buffers:
+ for bucket in buffer.buckets:
+ bucket_groups.append(
+ _ParamAndGradBucketGroup(
+ [bucket],
+ buffer.ddp_config,
+ buffer.data_parallel_group,
+ buffer.data_parallel_world_size,
+ )
+ )
+ return bucket_groups
+ else:
+ # Case 3: When using fp8 params, merge all non-fp8 buckets into the last fp8 bucket group.
+ non_fp8_buckets = []
+ for buffer in buffers:
+ if buffer.param_dtype != torch.uint8:
+ for bucket in buffer.buckets:
+ non_fp8_buckets.append(bucket)
+
+ bucket_groups = []
+ fp8_buffer = dtype_to_buffer_map[torch.uint8]
+ for bucket in fp8_buffer.buckets:
+ if len(bucket_groups) == len(fp8_buffer.buckets) - 1:
+ # The last bucket group.
+ group_buckets = [bucket] + non_fp8_buckets
+ else:
+ # The first N-1 bucket groups.
+ group_buckets = [bucket]
+ bucket_groups.append(
+ _ParamAndGradBucketGroup(
+ group_buckets,
+ buffer.ddp_config,
+ buffer.data_parallel_group,
+ buffer.data_parallel_world_size,
+ )
+ )
+ return bucket_groups
+
+
+# For backwards compatibility. ParamAndGradBuffer will be deprecated in future release.
+# _ParamAndGradBuffer is not intended to be consumed directly by external code.
+class ParamAndGradBuffer(_ParamAndGradBuffer):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ warnings.warn(
+ "`ParamAndGradBuffer` will be deprecated in a future release, and is not "
+ "intended to be used by external code."
+ )
diff --git a/megatron/core/enums.py b/megatron/core/enums.py
index cf1452b23e..46e7d3b766 100644
--- a/megatron/core/enums.py
+++ b/megatron/core/enums.py
@@ -2,6 +2,7 @@
import enum
+
class ModelType(enum.Enum):
encoder_or_decoder = 1
encoder_and_decoder = 2
diff --git a/megatron/core/export/__init__.py b/megatron/core/export/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/export/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/export/data_type.py b/megatron/core/export/data_type.py
new file mode 100644
index 0000000000..38fbdea8f6
--- /dev/null
+++ b/megatron/core/export/data_type.py
@@ -0,0 +1,5 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from enum import Enum
+
+DataType = Enum('DataType', ["bfloat16", "float16", "float32"])
diff --git a/megatron/core/export/export_config.py b/megatron/core/export/export_config.py
new file mode 100644
index 0000000000..2cc1e208be
--- /dev/null
+++ b/megatron/core/export/export_config.py
@@ -0,0 +1,19 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass
+
+
+@dataclass
+class ExportConfig:
+ """Base configuration for Megatron Core Export
+
+ These parameters control the export setting for trtllm
+ """
+
+ inference_tp_size: int = 1
+
+ inference_pp_size: int = 1
+
+ use_parallel_embedding: bool = False
+
+ use_embedding_sharing: bool = False
diff --git a/megatron/core/export/model_type.py b/megatron/core/export/model_type.py
new file mode 100644
index 0000000000..6a33d6440e
--- /dev/null
+++ b/megatron/core/export/model_type.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from enum import Enum
+
+ModelType = Enum(
+ 'ModelType', ["gpt", "gptnext", "llama", "falcon", "starcoder", "mixtral", "gemma"]
+)
diff --git a/megatron/core/export/trtllm/__init__.py b/megatron/core/export/trtllm/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/export/trtllm/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/export/trtllm/engine_builder/__init__.py b/megatron/core/export/trtllm/engine_builder/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/export/trtllm/engine_builder/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py b/megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py
new file mode 100644
index 0000000000..e729fec410
--- /dev/null
+++ b/megatron/core/export/trtllm/engine_builder/trtllm_engine_builder.py
@@ -0,0 +1,148 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import tensorrt_llm
+from tensorrt_llm._common import check_max_num_tokens
+from tensorrt_llm.builder import BuildConfig
+from tensorrt_llm.commands.build import build as build_trtllm
+from tensorrt_llm.logger import logger
+from tensorrt_llm.lora_manager import LoraConfig
+from tensorrt_llm.models.modeling_utils import optimize_model, preprocess_weights
+from tensorrt_llm.plugin import PluginConfig
+
+
+class TRTLLMEngineBuilder:
+ """A utility class to build TRTLLM engine"""
+
+ @staticmethod
+ def build_and_save_engine(
+ engine_dir: str,
+ trtllm_model_weights: dict,
+ trtllm_model_config,
+ max_input_len: int = 1024,
+ max_output_len: int = 1024,
+ max_batch_size: int = 4,
+ lora_ckpt_list=None,
+ use_lora_plugin=None,
+ max_lora_rank: int = 64,
+ lora_target_modules=None,
+ max_prompt_embedding_table_size: int = 0,
+ paged_kv_cache: bool = True,
+ remove_input_padding: bool = True,
+ paged_context_fmha: bool = False,
+ use_refit: bool = False,
+ max_num_tokens: int = None,
+ max_seq_len: int = None,
+ opt_num_tokens: int = None,
+ max_beam_width: int = 1,
+ tokens_per_block: int = 128,
+ multiple_profiles: bool = False,
+ gpt_attention_plugin: str = "auto",
+ gemm_plugin: str = "auto",
+ ):
+ """Method to build the TRTLLM Engine
+
+ This method uses the TRTLLMEngineBuilder to build and save the engine to engine dir
+
+ Args:
+ engine_dir (str): The file path to save the engine
+ trtllm_model_weights (dict): The TRTLLM converted model weights dict
+ trtllm_model_config : The TRTLLM Config
+ max_input_len (int, optional): Max input length. Defaults to 1024.
+ max_output_len (int, optional): Max output length. Defaults to 1024.
+ max_batch_size (int, optional): Max batch size. Defaults to 4.
+ model_type (ModelType, optional): ModelType enum. Defaults to ModelType.gpt.
+ lora_ckpt_list (_type_, optional): Lora checkpoint list. Defaults to None.
+ use_lora_plugin (_type_, optional): Use lora plugin. Defaults to None.
+ max_lora_rank (int, optional): Max lora rank. Defaults to 64.
+ lora_target_modules (_type_, optional): Lora target modules. Defaults to None.
+ max_prompt_embedding_table_size (int, optional): Defaults to 0.
+ paged_kv_cache (bool, optional): Use Paged KV cache. Defaults to True.
+ remove_input_padding (bool, optional): Remove input padding. Defaults to True.
+ paged_context_fmha (bool, optional): Paged context fmha. Defaults to False.
+ use_refit (bool, optional): Use refit. Defaults to False.
+ max_num_tokens (int, optional): Max num of tokens. Defaults to None.
+ max_seq_len (int, optional): Max seq length. Defaults to None.
+ opt_num_tokens (int, optional): Opt number of tokens. Defaults to None.
+ max_beam_width (int, optional): Max beam width. Defaults to 1.
+ tokens_per_block (int, optional): Nmber of tokens per block. Defaults to 128.
+ multiple_profiles (bool, optional): Use multiple profiles. Defaults to False.
+ gpt_attention_plugin (str, optional): Gpt attention plugin to use. Defaults to "auto".
+ gemm_plugin (str, optional): Gemma plugin to use. Defaults to "auto".
+ """
+ architecture = (
+ "LLaMAForCausalLM"
+ if trtllm_model_config.architecture == "LlamaForCausalLM"
+ else trtllm_model_config.architecture
+ )
+ try:
+ model_cls = getattr(tensorrt_llm.models, architecture)
+ except:
+ raise AttributeError(f"Could not find TRTLLM model for architecture: {architecture}!")
+
+ logger.set_level("info")
+ plugin_config = PluginConfig()
+ plugin_config.gpt_attention_plugin = gpt_attention_plugin
+ plugin_config.gemm_plugin = gemm_plugin
+ if paged_kv_cache:
+ plugin_config.enable_paged_kv_cache(tokens_per_block=tokens_per_block)
+ else:
+ plugin_config.paged_kv_cache = False
+ plugin_config.remove_input_padding = remove_input_padding
+ plugin_config.use_paged_context_fmha = paged_context_fmha
+ plugin_config.multiple_profiles = multiple_profiles
+
+ if max_seq_len is None:
+ max_seq_len = max_input_len + max_output_len
+
+ max_num_tokens, opt_num_tokens = check_max_num_tokens(
+ max_num_tokens=max_num_tokens,
+ opt_num_tokens=opt_num_tokens,
+ max_seq_len=max_seq_len,
+ max_batch_size=max_batch_size,
+ max_input_len=max_input_len,
+ max_beam_width=max_beam_width,
+ remove_input_padding=remove_input_padding,
+ enable_context_fmha=plugin_config.context_fmha,
+ tokens_per_block=tokens_per_block,
+ multiple_profiles=multiple_profiles,
+ )
+
+ build_dict = {
+ 'max_input_len': max_input_len,
+ 'max_output_len': max_output_len,
+ 'max_batch_size': max_batch_size,
+ 'max_beam_width': max_beam_width,
+ 'max_seq_len': max_seq_len,
+ 'max_num_tokens': max_num_tokens,
+ 'opt_num_tokens': opt_num_tokens,
+ 'max_prompt_embedding_table_size': max_prompt_embedding_table_size,
+ 'gather_context_logits': False,
+ 'gather_generation_logits': False,
+ 'strongly_typed': False,
+ 'builder_opt': None,
+ 'use_refit': use_refit,
+ 'multiple_profiles': multiple_profiles,
+ }
+ build_config = BuildConfig.from_dict(build_dict, plugin_config=plugin_config)
+
+ if use_lora_plugin is not None:
+ # build_config.plugin_config.set_lora_plugin(use_lora_plugin)
+ # build_config.plugin_config._lora_plugin = use_lora_plugin
+ lora_config = LoraConfig(
+ lora_dir=lora_ckpt_list,
+ lora_ckpt_source='nemo', # TODO : NEED TO SEE HOW TO HANDLE THIS FOR MCORE
+ max_lora_rank=max_lora_rank,
+ lora_target_modules=lora_target_modules,
+ )
+ build_config.lora_config = lora_config
+
+ model = model_cls.from_config(trtllm_model_config)
+ model = optimize_model(
+ model,
+ use_parallel_embedding=trtllm_model_config.use_parallel_embedding,
+ share_embedding_table=trtllm_model_config.share_embedding_table,
+ )
+ preprocess_weights(trtllm_model_weights, trtllm_model_config)
+ model.load(trtllm_model_weights)
+ engine = build_trtllm(model, build_config)
+ engine.save(engine_dir)
diff --git a/megatron/core/export/trtllm/model_to_trllm_mapping/__init__.py b/megatron/core/export/trtllm/model_to_trllm_mapping/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/export/trtllm/model_to_trllm_mapping/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/export/trtllm/model_to_trllm_mapping/default_conversion_dict.py b/megatron/core/export/trtllm/model_to_trllm_mapping/default_conversion_dict.py
new file mode 100644
index 0000000000..cad9315034
--- /dev/null
+++ b/megatron/core/export/trtllm/model_to_trllm_mapping/default_conversion_dict.py
@@ -0,0 +1,18 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.export.model_type import ModelType
+from megatron.core.export.trtllm.model_to_trllm_mapping.falcon_model import FALCON_DICT
+from megatron.core.export.trtllm.model_to_trllm_mapping.gemma_model import GEMMA_DICT
+from megatron.core.export.trtllm.model_to_trllm_mapping.gpt_model import GPT_DICT
+from megatron.core.export.trtllm.model_to_trllm_mapping.gpt_next_model import GPT_NEXT_DICT
+from megatron.core.export.trtllm.model_to_trllm_mapping.llama_model import LLAMA_DICT
+from megatron.core.export.trtllm.model_to_trllm_mapping.starcoder_model import STARCODER_DICT
+
+DEFAULT_CONVERSION_DICT = {
+ ModelType.llama: LLAMA_DICT,
+ ModelType.falcon: FALCON_DICT,
+ ModelType.gemma: GEMMA_DICT,
+ ModelType.starcoder: STARCODER_DICT,
+ ModelType.gpt: GPT_DICT,
+ ModelType.gptnext: GPT_NEXT_DICT,
+}
diff --git a/megatron/core/export/trtllm/model_to_trllm_mapping/falcon_model.py b/megatron/core/export/trtllm/model_to_trllm_mapping/falcon_model.py
new file mode 100644
index 0000000000..d1469d02ba
--- /dev/null
+++ b/megatron/core/export/trtllm/model_to_trllm_mapping/falcon_model.py
@@ -0,0 +1,26 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.export.trtllm.trtllm_layers import TRTLLMLayers
+
+# pylint: disable=line-too-long
+FALCON_DICT = {
+ # INPUT
+ 'embedding.word_embeddings.weight': TRTLLMLayers.vocab_embedding,
+ 'embedding.position_embeddings.weight': TRTLLMLayers.position_embedding,
+ # ATTENTION
+ 'decoder.layers.input_layernorm.weight': TRTLLMLayers.input_layernorm_weight,
+ 'decoder.layers.input_layernorm.bias': TRTLLMLayers.input_layernorm_bias,
+ 'decoder.layers.self_attention.linear_qkv.weight': TRTLLMLayers.attention_qkv_weight,
+ 'decoder.layers.self_attention.linear_qkv.layer_norm_weight': TRTLLMLayers.input_layernorm_weight,
+ 'decoder.layers.self_attention.linear_proj.weight': TRTLLMLayers.attention_dense_weight,
+ # MLP
+ 'decoder.layers.pre_mlp_layernorm.weight': TRTLLMLayers.post_layernorm_weight,
+ 'decoder.layers.pre_mlp_layernorm.bias': TRTLLMLayers.post_layernorm_bias,
+ 'decoder.layers.mlp.linear_fc1.weight': TRTLLMLayers.mlp_fc_weight,
+ 'decoder.layers.mlp.linear_fc2.weight': TRTLLMLayers.mlp_projection_weight,
+ # FINAL LAYER NORM
+ 'decoder.final_layernorm.weight': TRTLLMLayers.final_layernorm_weight,
+ 'decoder.final_layernorm.bias': TRTLLMLayers.final_layernorm_bias,
+ # OUTPUT LAYER
+ 'output_layer.weight': TRTLLMLayers.lm_head,
+}
diff --git a/megatron/core/export/trtllm/model_to_trllm_mapping/gemma_model.py b/megatron/core/export/trtllm/model_to_trllm_mapping/gemma_model.py
new file mode 100644
index 0000000000..47a0211706
--- /dev/null
+++ b/megatron/core/export/trtllm/model_to_trllm_mapping/gemma_model.py
@@ -0,0 +1,21 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.export.trtllm.trtllm_layers import TRTLLMLayers
+
+# pylint: disable=line-too-long
+GEMMA_DICT = {
+ # INPUT
+ 'embedding.word_embeddings.weight': TRTLLMLayers.vocab_embedding,
+ # ATTENTION
+ 'decoder.layers.self_attention.linear_qkv.weight': TRTLLMLayers.attention_qkv_weight,
+ 'decoder.layers.self_attention.linear_qkv.layer_norm_weight': TRTLLMLayers.input_layernorm_weight,
+ 'decoder.layers.self_attention.linear_proj.weight': TRTLLMLayers.attention_dense_weight,
+ # MLP
+ 'decoder.layers.mlp.linear_fc1.weight': TRTLLMLayers.mlp_fc_weight,
+ 'decoder.layers.mlp.linear_fc2.weight': TRTLLMLayers.mlp_projection_weight,
+ 'decoder.layers.mlp.linear_fc1.layer_norm_weight': TRTLLMLayers.post_layernorm_weight,
+ # FINAL LAYER NORM
+ 'decoder.final_layernorm.weight': TRTLLMLayers.final_layernorm_weight,
+ # OUTPUT LAYER
+ 'output_layer.weight': TRTLLMLayers.lm_head,
+}
diff --git a/megatron/core/export/trtllm/model_to_trllm_mapping/gpt_model.py b/megatron/core/export/trtllm/model_to_trllm_mapping/gpt_model.py
new file mode 100644
index 0000000000..eda27600c6
--- /dev/null
+++ b/megatron/core/export/trtllm/model_to_trllm_mapping/gpt_model.py
@@ -0,0 +1,28 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.export.trtllm.trtllm_layers import TRTLLMLayers
+
+GPT_DICT = {
+ # INPUT
+ 'embedding.word_embeddings.weight': TRTLLMLayers.vocab_embedding,
+ 'embedding.position_embeddings.weight': TRTLLMLayers.position_embedding,
+ # ATTENTION
+ 'decoder.layers.input_layernorm.weight': TRTLLMLayers.input_layernorm_weight,
+ 'decoder.layers.input_layernorm.bias': TRTLLMLayers.input_layernorm_bias,
+ 'decoder.layers.self_attention.linear_qkv.weight': TRTLLMLayers.attention_qkv_weight,
+ 'decoder.layers.self_attention.linear_qkv.bias': TRTLLMLayers.attention_qkv_bias,
+ 'decoder.layers.self_attention.linear_proj.weight': TRTLLMLayers.attention_dense_weight,
+ 'decoder.layers.self_attention.linear_proj.bias': TRTLLMLayers.attention_dense_bias,
+ # MLP
+ 'decoder.layers.pre_mlp_layernorm.weight': TRTLLMLayers.post_layernorm_weight,
+ 'decoder.layers.pre_mlp_layernorm.bias': TRTLLMLayers.post_layernorm_bias,
+ 'decoder.layers.mlp.linear_fc1.weight': TRTLLMLayers.mlp_fc_weight,
+ 'decoder.layers.mlp.linear_fc1.bias': TRTLLMLayers.mlp_fc_bias,
+ 'decoder.layers.mlp.linear_fc2.weight': TRTLLMLayers.mlp_projection_weight,
+ 'decoder.layers.mlp.linear_fc2.bias': TRTLLMLayers.mlp_projection_bias,
+ # FINAL LAYER NORM
+ 'decoder.final_layernorm.weight': TRTLLMLayers.final_layernorm_weight,
+ 'decoder.final_layernorm.bias': TRTLLMLayers.final_layernorm_bias,
+ # OUTPUT LAYER
+ 'output_layer.weight': TRTLLMLayers.lm_head,
+}
diff --git a/megatron/core/export/trtllm/model_to_trllm_mapping/gpt_next_model.py b/megatron/core/export/trtllm/model_to_trllm_mapping/gpt_next_model.py
new file mode 100644
index 0000000000..ac5f84ef1b
--- /dev/null
+++ b/megatron/core/export/trtllm/model_to_trllm_mapping/gpt_next_model.py
@@ -0,0 +1,24 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.export.trtllm.trtllm_layers import TRTLLMLayers
+
+# pylint: disable=line-too-long
+GPT_NEXT_DICT = {
+ # INPUT
+ 'embedding.word_embeddings.weight': TRTLLMLayers.vocab_embedding,
+ # ATTENTION
+ 'decoder.layers.self_attention.linear_proj.weight': TRTLLMLayers.attention_dense_weight,
+ 'decoder.layers.self_attention.linear_qkv.layer_norm_weight': TRTLLMLayers.input_layernorm_weight,
+ 'decoder.layers.self_attention.linear_qkv.layer_norm_bias': TRTLLMLayers.input_layernorm_bias,
+ 'decoder.layers.self_attention.linear_qkv.weight': TRTLLMLayers.attention_qkv_weight,
+ # MLP
+ 'decoder.layers.mlp.linear_fc1.layer_norm_weight': TRTLLMLayers.post_layernorm_weight,
+ 'decoder.layers.mlp.linear_fc1.layer_norm_bias': TRTLLMLayers.post_layernorm_bias,
+ 'decoder.layers.mlp.linear_fc1.weight': TRTLLMLayers.mlp_fc_weight,
+ 'decoder.layers.mlp.linear_fc2.weight': TRTLLMLayers.mlp_projection_weight,
+ # FINAL LAYER NORM
+ 'decoder.final_layernorm.weight': TRTLLMLayers.final_layernorm_weight,
+ 'decoder.final_layernorm.bias': TRTLLMLayers.final_layernorm_bias,
+ # OUTPUT LAYER
+ 'output_layer.weight': TRTLLMLayers.lm_head,
+}
diff --git a/megatron/core/export/trtllm/model_to_trllm_mapping/llama_model.py b/megatron/core/export/trtllm/model_to_trllm_mapping/llama_model.py
new file mode 100644
index 0000000000..5fd2067081
--- /dev/null
+++ b/megatron/core/export/trtllm/model_to_trllm_mapping/llama_model.py
@@ -0,0 +1,22 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.export.trtllm.trtllm_layers import TRTLLMLayers
+
+# pylint: disable=line-too-long
+LLAMA_DICT = {
+ # INPUT
+ 'embedding.word_embeddings.weight': TRTLLMLayers.vocab_embedding,
+ 'embedding.position_embeddings.weight': TRTLLMLayers.position_embedding,
+ # ATTENTION
+ 'decoder.layers.self_attention.linear_qkv.weight': TRTLLMLayers.attention_qkv_weight,
+ 'decoder.layers.self_attention.linear_qkv.layer_norm_weight': TRTLLMLayers.input_layernorm_weight,
+ 'decoder.layers.self_attention.linear_proj.weight': TRTLLMLayers.attention_dense_weight,
+ # MLP
+ 'decoder.layers.mlp.linear_fc1.weight': TRTLLMLayers.mlp_fc_weight,
+ 'decoder.layers.mlp.linear_fc2.weight': TRTLLMLayers.mlp_projection_weight,
+ 'decoder.layers.mlp.linear_fc1.layer_norm_weight': TRTLLMLayers.post_layernorm_weight,
+ # FINAL LAYER NORM
+ 'decoder.final_layernorm.weight': TRTLLMLayers.final_layernorm_weight,
+ # OUTPUT LAYER
+ 'output_layer.weight': TRTLLMLayers.lm_head,
+}
diff --git a/megatron/core/export/trtllm/model_to_trllm_mapping/starcoder_model.py b/megatron/core/export/trtllm/model_to_trllm_mapping/starcoder_model.py
new file mode 100644
index 0000000000..dce61d26c5
--- /dev/null
+++ b/megatron/core/export/trtllm/model_to_trllm_mapping/starcoder_model.py
@@ -0,0 +1,30 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.export.trtllm.trtllm_layers import TRTLLMLayers
+
+# pylint: disable=line-too-long
+STARCODER_DICT = {
+ # INPUT
+ 'embedding.word_embeddings.weight': TRTLLMLayers.vocab_embedding,
+ # ATTENTION
+ 'decoder.layers.input_layernorm.weight': TRTLLMLayers.input_layernorm_weight,
+ 'decoder.layers.input_layernorm.bias': TRTLLMLayers.input_layernorm_bias,
+ 'decoder.layers.self_attention.linear_qkv.weight': TRTLLMLayers.attention_qkv_weight,
+ 'decoder.layers.self_attention.linear_qkv.bias': TRTLLMLayers.attention_qkv_bias,
+ 'decoder.layers.self_attention.linear_qkv.layer_norm_weight': TRTLLMLayers.input_layernorm_weight,
+ 'decoder.layers.self_attention.linear_qkv.layer_norm_bias': TRTLLMLayers.input_layernorm_bias,
+ 'decoder.layers.self_attention.linear_proj.weight': TRTLLMLayers.attention_dense_weight,
+ 'decoder.layers.self_attention.linear_proj.bias': TRTLLMLayers.attention_dense_bias,
+ # MLP
+ 'decoder.layers.mlp.linear_fc1.weight': TRTLLMLayers.mlp_fc_weight,
+ 'decoder.layers.mlp.linear_fc1.bias': TRTLLMLayers.mlp_fc_bias,
+ 'decoder.layers.mlp.linear_fc2.weight': TRTLLMLayers.mlp_projection_weight,
+ 'decoder.layers.mlp.linear_fc2.bias': TRTLLMLayers.mlp_projection_bias,
+ 'decoder.layers.mlp.linear_fc1.layer_norm_weight': TRTLLMLayers.post_layernorm_weight,
+ 'decoder.layers.mlp.linear_fc1.layer_norm_bias': TRTLLMLayers.post_layernorm_bias,
+ # FINAL LAYER NORM
+ 'decoder.final_layernorm.weight': TRTLLMLayers.final_layernorm_weight,
+ 'decoder.final_layernorm.bias': TRTLLMLayers.final_layernorm_bias,
+ # OUTPUT LAYER
+ 'output_layer.weight': TRTLLMLayers.lm_head,
+}
diff --git a/megatron/core/export/trtllm/trt_model_config.py b/megatron/core/export/trtllm/trt_model_config.py
new file mode 100644
index 0000000000..2ed09398c2
--- /dev/null
+++ b/megatron/core/export/trtllm/trt_model_config.py
@@ -0,0 +1,15 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import tensorrt_llm
+
+from megatron.core.export.model_type import ModelType
+
+TRT_MODEL_CONFIG = {
+ ModelType.gpt: tensorrt_llm.models.gpt.config.GPTConfig,
+ ModelType.gptnext: tensorrt_llm.models.gpt.config.GPTConfig,
+ ModelType.starcoder: tensorrt_llm.models.gpt.config.GPTConfig,
+ ModelType.mixtral: tensorrt_llm.models.llama.config.LLaMAConfig,
+ ModelType.llama: tensorrt_llm.models.llama.config.LLaMAConfig,
+ ModelType.gemma: tensorrt_llm.models.GemmaConfig,
+ ModelType.falcon: tensorrt_llm.models.falcon.config.FalconConfig,
+}
diff --git a/megatron/core/export/trtllm/trt_model_type.py b/megatron/core/export/trtllm/trt_model_type.py
new file mode 100644
index 0000000000..f45ff1786e
--- /dev/null
+++ b/megatron/core/export/trtllm/trt_model_type.py
@@ -0,0 +1,13 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.export.model_type import ModelType
+
+TRT_MODEL_TYPE_STRING = {
+ ModelType.gpt: 'GPTForCausalLM',
+ ModelType.gptnext: 'GPTForCausalLM',
+ ModelType.starcoder: 'GPTForCausalLM',
+ ModelType.mixtral: 'LlamaForCausalLM',
+ ModelType.llama: 'LlamaForCausalLM',
+ ModelType.gemma: 'GemmaForCausalLM',
+ ModelType.falcon: 'FalconForCausalLM',
+}
diff --git a/megatron/core/export/trtllm/trtllm_helper.py b/megatron/core/export/trtllm/trtllm_helper.py
new file mode 100644
index 0000000000..d8bef18b33
--- /dev/null
+++ b/megatron/core/export/trtllm/trtllm_helper.py
@@ -0,0 +1,461 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import tensorrt_llm
+from tensorrt_llm.functional import non_gated_version
+from tensorrt_llm.layers import MoeConfig
+
+from megatron.core.export.data_type import DataType
+from megatron.core.export.export_config import ExportConfig
+from megatron.core.export.model_type import ModelType
+from megatron.core.export.trtllm.engine_builder.trtllm_engine_builder import TRTLLMEngineBuilder
+from megatron.core.export.trtllm.model_to_trllm_mapping.default_conversion_dict import (
+ DEFAULT_CONVERSION_DICT,
+)
+from megatron.core.export.trtllm.trt_model_config import TRT_MODEL_CONFIG
+from megatron.core.export.trtllm.trt_model_type import TRT_MODEL_TYPE_STRING
+
+# pylint: disable=line-too-long
+from megatron.core.export.trtllm.trtllm_weights_converter.distributed_trtllm_model_weights_converter import (
+ DistributedTRTLLMModelWeightsConverter,
+)
+from megatron.core.export.trtllm.trtllm_weights_converter.single_device_trtllm_model_weights_converter import (
+ SingleDeviceTRTLLMModelWeightsConverter,
+)
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class TRTLLMHelper:
+ """TRTLLM Helper class to convert export and build TRTLLM model."""
+
+ def __init__(
+ self,
+ transformer_config: TransformerConfig,
+ model_type: ModelType,
+ trtllm_conversion_dict: dict = {},
+ position_embedding_type: str = 'learned_absolute',
+ max_position_embeddings: int = None,
+ rotary_percentage: int = 1.0,
+ rotary_base: int = 10000,
+ moe_tp_mode: int = 2,
+ multi_query_mode: bool = False,
+ activation: str = "gelu",
+ seq_len_interpolation_factor: float = None,
+ moe_renorm_mode=None,
+ share_embeddings_and_output_weights=False,
+ ):
+ """Constructor for the TRTLLMHelper
+
+ There are two public API's supported by this helper.
+ a) get_trtllm_pretrained_config_and_model_weights
+ b) build_and_save_engine
+
+ Args:
+ transformer_config (TransformerConfig): The transformer config
+ model_type (ModelType): The type of the input model. Enum (megatron.core.export.model_type.ModelType)
+ conversion_dict (dict, optional): A conversion dictionary that will map your model layer names to trtllm equivalent layer names. Sample dictionaries are given megatron/core/export/model_mapping. NOTE: Ingore layer numbers in the model layer names. (e.g) decoder.layers.0.attention_qkv.weight will be decoder.layers.attention_qkv.weight in the mapping dictionary. Defaults to {}.
+ position_embedding_type (str, optional): The position embedding type. Defaults to None.
+ max_position_embeddings (int, optional): Max posistion embeddings value. Defaults to None.
+ rotary_percentage (int, optional): The rotary percentage if using rope embedding. Defaults to 1.0.
+ rotary_base (int, optional): The rotary base (theta value) if using rope embeddings. Defaults to 10000.
+ moe_tp_mode (int, optional): TRTLLM Config. Defaults to 2.
+ multi_query_mode (bool, optional): Defaults to False.
+ activation (str, optional): Defaults to "gelu".
+ seq_len_interpolation_factor (float, optional): The sequence length interpolation factor if using rope embeddings. Defaults to None.
+ moe_renorm_mode (optional) : Renormalization mode if using mixture of experts. Defaults to None.
+ share_embeddings_and_output_weights (bool, optional): True if input and output layers share weights. Defaults to False.
+ """
+
+ self.transformer_config = transformer_config
+ self.model_type = model_type
+ self.trtllm_conversion_dict = DEFAULT_CONVERSION_DICT[model_type]
+ self.trtllm_conversion_dict.update(trtllm_conversion_dict)
+ assert position_embedding_type in [
+ 'learned_absolute',
+ 'rope',
+ ], f"Position embedding type should be one of learned_absolute, rope. You entered {position_embedding_type}"
+ self.position_embedding_type = position_embedding_type
+ self.max_position_embeddings = max_position_embeddings
+ self.rotary_percentage = rotary_percentage
+ self.rotary_base = rotary_base
+ self.moe_tp_mode = moe_tp_mode
+ self.multi_query_mode = multi_query_mode
+ self.activation = activation
+ self.seq_len_interpolation_factor = seq_len_interpolation_factor
+ self.moe_renorm_mode = moe_renorm_mode
+ self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+
+ def _get_trtllm_config(
+ self,
+ export_config: ExportConfig,
+ world_size: int,
+ gpus_per_node: int,
+ vocab_size_padded: int,
+ dtype: DataType,
+ ):
+ """Get TRTLLM Config
+
+ Returns appropriate TRTLLM PretrainedConfig used by TRTLLM for building engine
+
+ Args:
+ export_config (ExportConfig): The export config that defines inference tp , pp size etc.
+ world_size (int): The number of gpus (Mostly TP * PP)
+ gpus_per_node (int): Num gpus per node
+ vocab_size_padded (int): Padded vocab size
+ dtype (DataType): The datatype or model precision
+
+ Returns:
+ GPTConfig or the LLamaConfig or the PretrainedConfig constructed from your model config
+ """
+ hidden_act = self.activation
+ hidden_act = (
+ hidden_act.split("-")[-1]
+ if self.transformer_config.num_moe_experts
+ else non_gated_version(hidden_act)
+ )
+
+ config = {
+ 'architecture': TRT_MODEL_TYPE_STRING[self.model_type],
+ 'dtype': dtype.name,
+ 'num_hidden_layers': self.transformer_config.num_layers,
+ 'num_attention_heads': self.transformer_config.num_attention_heads,
+ 'num_key_value_heads': (
+ self.transformer_config.num_query_groups
+ if self.transformer_config.num_query_groups
+ else self.transformer_config.num_attention_heads
+ ),
+ 'head_size': self.transformer_config.kv_channels,
+ 'hidden_size': self.transformer_config.hidden_size,
+ 'intermediate_size': self.transformer_config.ffn_hidden_size,
+ 'norm_epsilon': self.transformer_config.layernorm_epsilon,
+ 'vocab_size': vocab_size_padded,
+ 'position_embedding_type': (
+ "rope_gpt_neox" if self.position_embedding_type == "rope" else "learned_absolute"
+ ),
+ 'max_position_embeddings': self.max_position_embeddings,
+ 'hidden_act': hidden_act,
+ 'use_parallel_embedding': export_config.use_parallel_embedding,
+ 'embedding_sharding_dim': 0,
+ 'share_embedding_table': export_config.use_embedding_sharing,
+ 'quantization': {'quant_algo': None, 'kv_cache_quant_algo': None},
+ 'bias': self.transformer_config.add_bias_linear,
+ 'apply_query_key_layer_scaling': False,
+ 'rotary_pct': self.rotary_percentage,
+ 'rotary_base': self.rotary_base,
+ 'moe_num_experts': (
+ 0
+ if self.transformer_config.moe_router_topk == 0
+ else (self.transformer_config.num_moe_experts or 1)
+ ),
+ 'moe_top_k': self.transformer_config.moe_router_topk,
+ 'moe_normalization_mode': self.moe_renorm_mode
+ or MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE,
+ 'moe_tp_mode': self.moe_tp_mode,
+ 'logits_dtype': 'float32',
+ 'world_size': world_size,
+ 'tp_size': export_config.inference_tp_size,
+ 'pp_size': export_config.inference_pp_size,
+ 'gpus_per_node': gpus_per_node,
+ }
+
+ if self.model_type == ModelType.falcon:
+ config["new_decoder_architecture"] = (
+ False if self.transformer_config.num_layers == 32 else True
+ )
+ config["parallel_attention"] = True
+
+ if self.seq_len_interpolation_factor is not None:
+ config["rotary_scaling"] = {
+ "type": "linear",
+ "factor": float(self.seq_len_interpolation_factor),
+ }
+
+ config_cls = TRT_MODEL_CONFIG[self.model_type]
+ return config_cls(**config)
+
+ # pylint: disable=line-too-long
+ def get_trtllm_pretrained_config_and_model_weights(
+ self,
+ model_state_dict,
+ dtype: DataType,
+ export_config: ExportConfig = None,
+ on_device_distributed_conversion: bool = False,
+ vocab_size: int = None,
+ gpus_per_node: int = None,
+ state_dict_split_by_layer_numbers: bool = True,
+ ):
+ """Get TRTLLM Config and Converted Model Weights
+
+ This function returns the trtllm model weights as a list.
+ There are two modes for conversion. The default is to use a single device cpu/gpu for conversion.
+ NOTE: For faster performance, if your entire model will fit in memory, pre transfer the model state dict to cuda device and then call this function.
+ For on device conversion it returns weights which will be used on the device itself.
+ Same thing happens with the pretrained config
+
+ Args:
+ model_state_dict (dict, optional): The input model state dictionary (Entire model state loaded on CPU). Used only when on device conversion is set to False. Defaults to None.
+ False, or the model state dict of each GPU in the case of on_device conversion)
+ export_config (ExportConfig): The export config used to define inference tp size, pp size etc. Used only for on device conversion.
+ dtype (DataType): The data type of model precision
+ on_device_distributed_conversion (bool, optional): Convert on gpus in distributed setting. This assumes that the model state dict is sharded according to required inference model parallelism and that each gpu gets its part of the model state dict . Defaults to False.
+ vocab_size (int, optional): The vocabulary size. Defaults to None.
+ gpus_per_node (int, optional): The number of gpus per node. Used for on device conversion.
+ state_dict_split_by_layer_numbers (bool, optional): Are the model layers split by layer numbers in state dict. For example : mlp.fc1.weight can be represented like mlp.fc1.weight of shape [num_layers, hidden_dim, ffn_hidden_dim]} or it can be like mlp.fc1.layers.0.weight of shape [hidden_dim, ffn_hidden_dim], then mlp.fc1.layers.1.weight ... for all layers. If you use represenation 2 set this to True. Defaults to True
+
+ Returns:
+ Two lists . First list of trtllm converted model weights(Either on device, or a list of weights for each gpu) and the trtllm_model_configs.
+ """
+ if on_device_distributed_conversion:
+ assert (vocab_size is not None, "Need to pass in vocab_size for on device")
+ assert (
+ self.model_type in [ModelType.gpt, ModelType.gptnext, ModelType.llama],
+ "On device conversion only supported for model types gptnext and llama",
+ )
+ assert (
+ export_config is None,
+ "Export config is inferred based on the parallel state. If you want to set inference tp 2, then load the model with this TP2 setting and just pass in the model state dict. ",
+ )
+ assert (
+ gpus_per_node is not None
+ ), "Need to pass in gpus_per_node for on device conversion"
+ trtllm_model_weights_on_device, trtllm_model_config = (
+ self._get_trtllm_pretrained_config_and_model_weights_in_distributed_setting(
+ model_state_dict, dtype, vocab_size, gpus_per_node
+ )
+ )
+ return [trtllm_model_weights_on_device], [trtllm_model_config]
+
+ else:
+ assert not (
+ self.share_embeddings_and_output_weights and not export_config.use_embedding_sharing
+ ), "Found share_embeddings_and_output_weights is True in the model. So set export_config.use_embedding_sharing to True"
+ assert (
+ vocab_size is None
+ ), "Vocab size is inferred from the input layer for cpu conversion. So leave it as None"
+ trtllm_model_weights_list, trtllm_model_config_list = (
+ self._get_trtllm_pretrained_config_and_model_weights_list_on_single_device(
+ export_config,
+ model_state_dict,
+ dtype,
+ gpus_per_node,
+ state_dict_split_by_layer_numbers,
+ )
+ )
+
+ return trtllm_model_weights_list, trtllm_model_config_list
+
+ def _get_trtllm_pretrained_config_and_model_weights_in_distributed_setting(
+ self, model_state_dict: dict, dtype: DataType, vocab_size: int, gpus_per_node: int
+ ):
+ """Get the TRTLLM Pretrained config and model weights list in a distributed setting
+
+ This function assumes the model state dict is distributed according to model parallelism .
+ Each device gets its own model state dict
+
+ Args:
+ export_config (ExportConfig): The export config to set inference tp, pp size etc.
+ model_state_dict (dict): The model state dictionary (All collected on cpu)
+ dtype (DataType): The data type or model precision
+ vocab_size (int): Tokenizer vocab size
+ gpus_per_node (int): The number of gpus per node
+
+ Returns:
+ Two lists . List of trtllm converted model weights and trtllm model configs (One for each gpu).
+ """
+
+ distributed_trtllm_model_weights_converter = DistributedTRTLLMModelWeightsConverter(
+ transformer_config=self.transformer_config,
+ dtype=dtype,
+ multi_query_mode=self.multi_query_mode,
+ activation=self.activation,
+ )
+ distributed_trtllm_model_weights_converter.convert(
+ model_state_dict=model_state_dict,
+ trtllm_conversion_dict=self.trtllm_conversion_dict,
+ tokenizer_vocab_size=vocab_size,
+ )
+
+ export_config = ExportConfig(
+ inference_pp_size=distributed_trtllm_model_weights_converter.inference_pp_size,
+ inference_tp_size=distributed_trtllm_model_weights_converter.inference_tp_size,
+ use_parallel_embedding=True,
+ use_embedding_sharing=self.share_embeddings_and_output_weights,
+ )
+
+ world_size = export_config.inference_tp_size * export_config.inference_pp_size
+
+ trtllm_model_config = self._get_trtllm_config(
+ export_config=export_config,
+ world_size=world_size,
+ gpus_per_node=gpus_per_node,
+ vocab_size_padded=vocab_size,
+ dtype=dtype,
+ )
+
+ model_parallel_rank = (
+ distributed_trtllm_model_weights_converter.pp_rank
+ * distributed_trtllm_model_weights_converter.inference_tp_size
+ + distributed_trtllm_model_weights_converter.tp_rank
+ )
+
+ trtllm_model_config.mapping = tensorrt_llm.Mapping(
+ world_size=world_size,
+ rank=model_parallel_rank,
+ tp_size=export_config.inference_tp_size,
+ pp_size=export_config.inference_pp_size,
+ )
+
+ return distributed_trtllm_model_weights_converter.trtllm_model_weights, trtllm_model_config
+
+ def _get_trtllm_pretrained_config_and_model_weights_list_on_single_device(
+ self,
+ export_config: ExportConfig,
+ model_state_dict: dict,
+ dtype: DataType,
+ gpus_per_node=None,
+ state_dict_split_by_layer_numbers=True,
+ ):
+ """Get the TRTLLM Pretrained config and model weights list (one per gpu rank) on single device (CPU/GPU)
+
+ This function assumes the entire model state dict is present in CPU or on one GPU
+
+ Args:
+ export_config (ExportConfig): The export config to set inference tp, pp size etc.
+ model_state_dict (dict): The model state dictionary (All collected on cpu)
+ dtype (DataType): The data type or model precision
+ gpus_per_node (int, optional): Number of gpus per node
+ state_dict_split_by_layer_numbers (bool, optional): Are the model layers split by layer numbers in state dict. For example : mlp.fc1.weight can be represented like mlp.fc1.weight of shape [num_layers, hidden_dim, ffn_hidden_dim]} or it can be like mlp.fc1.layers.0.weight of shape [hidden_dim, ffn_hidden_dim], then mlp.fc1.layers.1.weight ... for all layers. If you use represenation 2 set this to True. Defaults to True
+
+ Returns:
+ Two lists . List of trtllm converted model weights and trtllm model configs (One for each gpu).
+ """
+ trtllm_model_configs_list = []
+ trtllm_model_weights_list = []
+
+ single_device_trtllm_model_weights_converter = SingleDeviceTRTLLMModelWeightsConverter(
+ export_config=export_config,
+ transformer_config=self.transformer_config,
+ dtype=dtype,
+ activation=self.activation,
+ multi_query_mode=self.multi_query_mode,
+ )
+ # Convert the input model state dict to trtllm model weights dictionary
+ single_device_trtllm_model_weights_converter.convert(
+ model_state_dict=model_state_dict,
+ trtllm_conversion_dict=self.trtllm_conversion_dict,
+ state_dict_split_by_layer_numbers=state_dict_split_by_layer_numbers,
+ )
+
+ vocab_size_padded = single_device_trtllm_model_weights_converter.get_padded_vocab_size()
+ world_size = export_config.inference_tp_size * export_config.inference_pp_size
+ gpus_per_node = gpus_per_node or export_config.inference_tp_size
+
+ for gpu_rank in range(world_size):
+ mapping = tensorrt_llm.Mapping(
+ world_size=world_size,
+ rank=gpu_rank,
+ tp_size=export_config.inference_tp_size,
+ pp_size=export_config.inference_pp_size,
+ )
+
+ # Important to create a new instance everytime so that the list elements have differnt rank values in the mapping object
+ trtllm_model_config = self._get_trtllm_config(
+ export_config=export_config,
+ world_size=world_size,
+ gpus_per_node=gpus_per_node,
+ vocab_size_padded=vocab_size_padded,
+ dtype=dtype,
+ )
+ trtllm_model_config.mapping = mapping
+ trtllm_model_configs_list.append(trtllm_model_config)
+
+ # Get the model weights for each rank and append it to the trtllm_model_weights_list
+ trtllm_model_weights_per_gpu = (
+ single_device_trtllm_model_weights_converter.get_local_model_weights_per_gpu(
+ mapping, trtllm_model_config
+ )
+ )
+ trtllm_model_weights_list.append(trtllm_model_weights_per_gpu)
+
+ return trtllm_model_weights_list, trtllm_model_configs_list
+
+ def build_and_save_engine(
+ self,
+ engine_dir: str,
+ trtllm_model_weights: dict,
+ trtllm_model_config,
+ max_input_len: int = 1024,
+ max_output_len: int = 1024,
+ max_batch_size: int = 4,
+ lora_ckpt_list=None,
+ use_lora_plugin=None,
+ max_lora_rank: int = 64,
+ lora_target_modules=None,
+ max_prompt_embedding_table_size: int = 0,
+ paged_kv_cache: bool = True,
+ remove_input_padding: bool = True,
+ paged_context_fmha: bool = False,
+ use_refit: bool = False,
+ max_num_tokens: int = None,
+ max_seq_len: int = None,
+ opt_num_tokens: int = None,
+ max_beam_width: int = 1,
+ tokens_per_block: int = 128,
+ multiple_profiles: bool = False,
+ gpt_attention_plugin: str = "auto",
+ gemm_plugin: str = "auto",
+ ):
+ """Method to build the TRTLLM Engine
+
+ This method uses the TRTLLMEngineBuilder to build and save the engine to engine dir
+
+ Args:
+ engine_dir (str): The file path to save the engine
+ trtllm_model_weights (dict): The TRTLLM converted model weights dict
+ trtllm_model_config : The TRTLLM Config
+ max_input_len (int, optional): Max input length. Defaults to 1024.
+ max_output_len (int, optional): Max output length. Defaults to 1024.
+ max_batch_size (int, optional): Max batch size. Defaults to 4.
+ lora_ckpt_list (_type_, optional): Lora checkpoint list. Defaults to None.
+ use_lora_plugin (_type_, optional): Use lora plugin. Defaults to None.
+ max_lora_rank (int, optional): Max lora rank. Defaults to 64.
+ lora_target_modules (_type_, optional): Lora target modules. Defaults to None.
+ max_prompt_embedding_table_size (int, optional): Max size of prompt embedding table. Defaults to 0.
+ paged_kv_cache (bool, optional): Use Paged KV cache. Defaults to True.
+ remove_input_padding (bool, optional): Remove input padding. Defaults to True.
+ paged_context_fmha (bool, optional): Paged context fmha. Defaults to False.
+ use_refit (bool, optional): Use refit. Defaults to False.
+ max_num_tokens (int, optional): Max num of tokens. Defaults to None.
+ max_seq_len (int, optional): Max seq length. Defaults to None.
+ opt_num_tokens (int, optional): Opt number of tokens. Defaults to None.
+ max_beam_width (int, optional): Max beam width. Defaults to 1.
+ tokens_per_block (int, optional): Nmber of tokens per block. Defaults to 128.
+ multiple_profiles (bool, optional): Use multiple profiles. Defaults to False.
+ gpt_attention_plugin (str, optional): Gpt attention plugin to use. Defaults to "auto".
+ gemm_plugin (str, optional): Gemma plugin to use. Defaults to "auto".
+ """
+
+ TRTLLMEngineBuilder.build_and_save_engine(
+ engine_dir,
+ trtllm_model_weights,
+ trtllm_model_config,
+ max_input_len,
+ max_output_len,
+ max_batch_size,
+ lora_ckpt_list,
+ use_lora_plugin,
+ max_lora_rank,
+ lora_target_modules,
+ max_prompt_embedding_table_size,
+ paged_kv_cache,
+ remove_input_padding,
+ paged_context_fmha,
+ use_refit,
+ max_num_tokens,
+ max_seq_len,
+ opt_num_tokens,
+ max_beam_width,
+ tokens_per_block,
+ multiple_profiles,
+ gpt_attention_plugin,
+ gemm_plugin,
+ )
diff --git a/megatron/core/export/trtllm/trtllm_layers.py b/megatron/core/export/trtllm/trtllm_layers.py
new file mode 100644
index 0000000000..0cf805dcb6
--- /dev/null
+++ b/megatron/core/export/trtllm/trtllm_layers.py
@@ -0,0 +1,157 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import re
+from enum import Enum
+from typing import Tuple
+
+
+class TRTLLMLayers(Enum):
+ """TRTLLM Layer names
+
+ This Enum will be used to map input model layer names to TRTLLM Layer names
+ """
+
+ # ONE TIME LAYERS (NOT ASSOCIATED TO TRANSFORMER BLOCK)
+ # Input layers
+ position_embedding = 'transformer.position_embedding.weight'
+ vocab_embedding = 'transformer.vocab_embedding.weight'
+ lm_head = 'lm_head.weight'
+
+ # Output layers
+ final_layernorm_weight = 'transformer.ln_f.weight'
+ final_layernorm_bias = 'transformer.ln_f.bias'
+
+ # TRANSFORMER LAYERS
+ # Attention block related layers
+ input_layernorm_weight = 'transformer.layers.input_layernorm.weight'
+ input_layernorm_bias = 'transformer.layers.input_layernorm.bias'
+ attention_qkv_weight = 'transformer.layers.attention.qkv.weight'
+ attention_qkv_bias = 'transformer.layers.attention.qkv.bias'
+ attention_dense_weight = 'transformer.layers.attention.dense.weight'
+ attention_dense_bias = 'transformer.layers.attention.dense.bias'
+
+ # mlp layers
+ mlp_fc_weight = 'transformer.layers.mlp.fc.weight'
+ mlp_fc_bias = 'transformer.layers.mlp.fc.bias'
+ post_layernorm_weight = 'transformer.layers.post_layernorm.weight'
+ post_layernorm_bias = 'transformer.layers.post_layernorm.bias'
+ mlp_projection_weight = 'transformer.layers.mlp.proj.weight'
+ mlp_projection_bias = 'transformer.layers.mlp.proj.bias'
+
+ # mixture of expert layers
+ mlp_router_weight = 'transformer.layers.mlp.router.weight'
+ mlp_fc_weight_mixture_of_experts = 'transformer.layers.mlp.fc.weight.expert'
+ mlp_projection_weight_mixture_of_experts = 'transformer.layers.mlp.proj.weight.expert'
+
+ @staticmethod
+ def return_layer_name_and_number(layer_name: str) -> Tuple[str, int]:
+ """Helper function to return layer name and number
+ Given an input layer e.g decoder.layers.2.self_attention.linear_qkv.weight,
+ this function returns decoder.layers.self_attention.linear_qkv.weight and layernumber 2.
+ In case no layer number is present, it returns None for the layer number
+ Args:
+ layer_name (dict): The input layer name
+
+ Returns:
+ Tuple[str, int]: The layer name , layer number (layer number could be None)
+ """
+ # Use regular expression to find the number specifically after 'layers.'
+ match = re.search(r'(?<=layers\.)\d+(?=\.)', layer_name)
+ if match:
+ # Extract the number and remove it from the layer name
+ number = match.group(0)
+ layer_name_without_number = re.sub(r'\.{}\.'.format(number), '.', layer_name)
+ return layer_name_without_number, int(number)
+ else:
+ # Return the original name if no number is found
+ return layer_name, None
+
+ # pylint: disable=line-too-long
+ @staticmethod
+ def rename_input_layer_names_to_trtllm_layer_names(
+ model_state_dict: dict,
+ trtllm_conversion_dict: dict,
+ state_dict_split_by_layer_numbers: bool = True,
+ ) -> dict:
+ """Helper function to rename model layer names to TRTLLM Layer names
+
+ We go through each layer (keys) in the model state dict,
+ and map it to the equivalent TRTLLMLayer name (megatron/core/export/trtllm/trtllm).
+ If we have a layer number associated with layer, we extract it out,
+ map the original layer name to equivalent trtllm layer name and add layer number back.
+ CPU Conversion will pass in model state dict without layer numbers
+ (i.e decoder.layers.mlp.linear_fc1.weight of shape [num_layers, hidden_dim, 4 * hidden_dim]) .
+ GPU conversion will pass model state dict with each layer seperated
+ (i.e decoder.layers.2.mlp.linear_fc1.weight of shape [hidden_dim, 4 * hidden_dim]).
+
+ Args:
+ model_state_dict (dict): The original model state dict
+ trtllm_conversion_dict (dict): The conversion dictionary mapping input model layer names to trtllm layer names
+ state_dict_split_by_layer_numbers (bool, optional): Are the model layers split by layer numbers in state dict. For example : mlp.fc1.weight can be represented like mlp.fc1.weight of shape [num_layers, hidden_dim, ffn_hidden_dim]} or it can be like mlp.fc1.layers.0.weight of shape [hidden_dim, ffn_hidden_dim], then mlp.fc1.layers.1.weight ... for all layers. If you use represenation 2 set this to True. Defaults to True
+
+ Raises:
+ ValueError: In case the keys dont match to trtllm keys or if all model layers are not mapped to equivalent trtllm keys
+
+ Returns:
+ dict: The model state dict with the key (i.e original model layer name) replaced by trtllm layer names
+ """
+ for original_model_layer_name in list(model_state_dict.keys()):
+ if "_extra_state" in original_model_layer_name:
+ del model_state_dict[original_model_layer_name]
+ continue
+
+ original_layer_name_without_number, layer_number = (
+ TRTLLMLayers.return_layer_name_and_number(original_model_layer_name)
+ )
+ if 'layers' in original_layer_name_without_number and state_dict_split_by_layer_numbers:
+ assert (
+ layer_number is not None
+ ), f"Layer number is None for {original_model_layer_name} and state_dict_split_by_layer_numbers is set to True. Consider setting it False"
+
+ if original_layer_name_without_number not in trtllm_conversion_dict:
+ raise ValueError(
+ f'Unable to rename key {original_layer_name_without_number}. Provide an appropriate mapping in the trtllm_conversion_dict when you initialize TRTLLMHelper'
+ )
+
+ trtllm_layer = trtllm_conversion_dict[original_layer_name_without_number]
+ assert isinstance(
+ trtllm_layer, TRTLLMLayers
+ ), f"{trtllm_layer} is not supported for conversion. Please use one of the TRTLLMLayerNames we provided in megatron/core/export/trtllm/trtllm_layer_names"
+
+ value = model_state_dict.pop(original_model_layer_name)
+
+ if layer_number is not None:
+ trtllm_layer_name_with_number = re.sub(
+ r'(?<=layers\.)', f'{layer_number}.', trtllm_layer.value
+ )
+ model_state_dict[trtllm_layer_name_with_number] = value
+ else:
+ model_state_dict[trtllm_layer.value] = value
+
+ return model_state_dict
+
+
+# These layers are not associated within the transformer block.
+# So they dont have a layer number (i.e independant of number of layers in the model)
+NON_TRANSFORMER_LAYERS_NAMES = [
+ TRTLLMLayers.vocab_embedding.value,
+ TRTLLMLayers.position_embedding.value,
+ TRTLLMLayers.lm_head.value,
+ TRTLLMLayers.final_layernorm_weight.value,
+ TRTLLMLayers.final_layernorm_bias.value,
+]
+
+
+def get_layer_name_without_prefix(layer: TRTLLMLayers) -> str:
+ """Get TRTLayer name without prefix
+
+ Given a layer e.g TRTLLMLayers.attention_qkv_weight it returns 'attention.qkv.weight'
+
+ Args:
+ layer (TRTLLMLayers): The TRTLLMLayer
+
+ Returns:
+ str: The TRTLLMLayers suffix (i.e Removing transformer.layers. fromt he layer name)
+ """
+ layer_name_without_prefix = layer.value.replace("transformer.layers.", "")
+ return layer_name_without_prefix
diff --git a/megatron/core/export/trtllm/trtllm_weights_converter/__init__.py b/megatron/core/export/trtllm/trtllm_weights_converter/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/export/trtllm/trtllm_weights_converter/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/export/trtllm/trtllm_weights_converter/distributed_trtllm_model_weights_converter.py b/megatron/core/export/trtllm/trtllm_weights_converter/distributed_trtllm_model_weights_converter.py
new file mode 100644
index 0000000000..035e23a16c
--- /dev/null
+++ b/megatron/core/export/trtllm/trtllm_weights_converter/distributed_trtllm_model_weights_converter.py
@@ -0,0 +1,258 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import torch
+from tqdm import tqdm
+
+from megatron.core import parallel_state
+from megatron.core.export.data_type import DataType
+from megatron.core.export.trtllm.trtllm_layers import NON_TRANSFORMER_LAYERS_NAMES, TRTLLMLayers
+from megatron.core.export.trtllm.trtllm_layers import get_layer_name_without_prefix as suffix
+from megatron.core.tensor_parallel.utils import VocabUtility
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+def str_dtype_to_torch(dtype: DataType):
+ """Get torch datatype from input datatype"""
+ from tensorrt_llm._utils import str_dtype_to_torch
+
+ return str_dtype_to_torch(dtype.name)
+
+
+# pylint: disable=line-too-long
+class DistributedTRTLLMModelWeightsConverter:
+ """The TRTLLM Converter class used for GPU (on device) conversion
+
+ This class is used to convert models sharded and on gpus. (It assumes that the model is already sharded appropriate to how you want to export it). (i.e) If you want to export to tp2pp2, then load the model in tp2pp2 setting and pass in their respective state dictionaries
+ """
+
+ def __init__(
+ self,
+ transformer_config: TransformerConfig,
+ dtype: DataType,
+ multi_query_mode: bool = False,
+ activation: str = "gelu",
+ ):
+ """Constructor for the TRTLLMModelWeightsConverterGPU class
+
+ This class is responsible to convert the model weights to TRTLLM equivalent weights.
+
+ Args:
+ transformer_config (TransformerConfig): The transformer config
+ dtype (DataType): The data type or model precision
+ multi_query_mode (bool, optional): Defaults to False.
+ activation (str, optional): Defaults to "gelu".
+ """
+ self.transformer_config = transformer_config
+ self.trtllm_model_weights = {}
+ self.storage_type = str_dtype_to_torch(dtype)
+ self.activation = activation
+ num_kv_heads = self.transformer_config.num_query_groups
+ if num_kv_heads == 0:
+ if multi_query_mode:
+ num_kv_heads = 1
+ else:
+ num_kv_heads = self.transformer_config.num_attention_heads
+ self.num_kv_heads = num_kv_heads
+
+ self.inference_pp_size = parallel_state.get_pipeline_model_parallel_world_size()
+ self.inference_tp_size = parallel_state.get_tensor_model_parallel_world_size()
+ self.tp_rank = parallel_state.get_tensor_model_parallel_rank()
+ self.pp_rank = parallel_state.get_pipeline_model_parallel_rank()
+ self.tp_group = parallel_state.get_tensor_model_parallel_group()
+ vp_size = parallel_state.get_virtual_pipeline_model_parallel_world_size()
+
+ assert (
+ vp_size is None or vp_size == 1
+ ), "Virtual parallelism is not supported in GPU Converter. Gather the VP chunks and use PP config."
+
+ def _add_to_trtllm_model_weights(self, val: torch.Tensor, layer_name: str):
+ assert torch.is_tensor(val), f"Expected a tensor for {layer_name} but got {type(val)}"
+ val = val.to(self.storage_type)
+ val = val.detach().contiguous()
+ if val.ndim >= 2:
+ val = torch.transpose(val.reshape(val.shape[0], -1), 0, 1)
+ if layer_name not in self.trtllm_model_weights:
+ self.trtllm_model_weights[layer_name] = torch.empty(
+ val.size(), dtype=val.dtype, layout=val.layout, device="cpu", pin_memory=True
+ )
+ self.trtllm_model_weights[layer_name] = val
+
+ def _convert_transformer_layer(self, layer_name: str, val: torch.Tensor):
+ """Convert Transformer layers to TRTLLM weights
+
+ Transformer layers referes to layers within the transformber block. They have a layer number associated with them. Depending on the layer we either directly save it to trtllm_model_weights, or split it across some dimension and save the splits
+
+ Args:
+ model_state_dict (dict): The input model state dictionary (All collected on CPU)
+ layer (TRTLLMLayerNames): The TRTLLM Layer that we want to change
+ """
+ if val.ndim == 2:
+ val = val.T
+
+ if (
+ layer_name.endswith(suffix(TRTLLMLayers.input_layernorm_weight))
+ or layer_name.endswith(suffix(TRTLLMLayers.input_layernorm_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.post_layernorm_weight))
+ or layer_name.endswith(suffix(TRTLLMLayers.post_layernorm_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.attention_dense_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.attention_dense_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.mlp_projection_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.mlp_router_weight))
+ or layer_name.endswith(suffix(TRTLLMLayers.attention_dense_weight))
+ or layer_name.endswith(suffix(TRTLLMLayers.mlp_projection_weight))
+ ):
+ # Same as layernorm1p in NeMo
+ if (
+ self.transformer_config.layernorm_zero_centered_gamma
+ and self.transformer_config.normalization == "LayerNorm"
+ and 'layernorm.weight' in layer_name
+ ):
+ val = val + 1.0
+
+ self._add_to_trtllm_model_weights(val=val, layer_name=layer_name)
+
+ elif layer_name.endswith(suffix(TRTLLMLayers.mlp_fc_weight)) or layer_name.endswith(
+ suffix(TRTLLMLayers.mlp_fc_bias)
+ ):
+
+ split_gated_activation = self.activation in [
+ "swiglu",
+ "geglu",
+ "fast-swiglu",
+ "fast-geglu",
+ ]
+ if split_gated_activation:
+ vals, gates = [[n] for n in torch.chunk(val, 2, axis=-1)]
+ gate_layer_name = layer_name.replace("fc", "gate")
+ self._add_to_trtllm_model_weights(val=gates[0], layer_name=gate_layer_name)
+ val = vals[0]
+
+ self._add_to_trtllm_model_weights(val=val, layer_name=layer_name)
+
+ elif layer_name.endswith(suffix(TRTLLMLayers.attention_qkv_bias)):
+ qkv_hidden_dim = val.shape[0]
+ size_per_head = (
+ qkv_hidden_dim
+ // (self.transformer_config.num_attention_heads + 2 * self.num_kv_heads)
+ * self.inference_tp_size
+ )
+ q_num = self.transformer_config.num_attention_heads // self.num_kv_heads
+
+ # We first concat all sub weights per tp rank together.
+ val = val.reshape(self.num_kv_heads // self.inference_tp_size, q_num + 2, size_per_head)
+ qkv = torch.split(val, [q_num, 1, 1], dim=1)
+ split_vals = torch.concatenate(
+ [qkv[0].reshape(-1), qkv[1].reshape(-1), qkv[2].reshape(-1)], dim=0
+ )
+ self._add_to_trtllm_model_weights(val=split_vals, layer_name=layer_name)
+
+ # TODO : Should add a atten layer dimension "qkvqkv, qqkkvv etc to see how to reshape here"
+ elif layer_name.endswith(suffix(TRTLLMLayers.attention_qkv_weight)):
+ hidden_dim = val.shape[0]
+ size_per_head = self.transformer_config.kv_channels
+ if size_per_head is None:
+ size_per_head = hidden_dim // self.transformer_config.num_attention_heads
+ q_num = self.transformer_config.num_attention_heads // self.num_kv_heads
+
+ val = val.reshape(
+ hidden_dim, self.num_kv_heads // self.inference_tp_size, q_num + 2, size_per_head
+ )
+ qkv = torch.split(val, [q_num, 1, 1], dim=2)
+ split_vals = torch.concatenate(
+ [
+ qkv[0].reshape(hidden_dim, -1),
+ qkv[1].reshape(hidden_dim, -1),
+ qkv[2].reshape(hidden_dim, -1),
+ ],
+ dim=1,
+ )
+ self._add_to_trtllm_model_weights(val=split_vals, layer_name=layer_name)
+
+ else:
+ raise ValueError(f"{layer_name} cannot be handled by GPU converter")
+
+ def _convert_non_transformer_layer(self, model_state_dict: dict, layer_name: str):
+ """Convert Non Transformer layers to TRTLLM weights
+
+ Non transformer layers referes to layers that occur only once in the model (e.g Embedding , final output layer etc. ) They dont have any layer number associated with them. We remove this layer from the original state dict and cast it to storage type and convert to numpy and add it to trtllm_model_weights
+
+ Args:
+ model_state_dict (dict): The input model state dictionary (All collected on CPU)
+ layer (TRTLLMLayerNames): The TRTLLM Layer that we want to change
+ """
+ if layer_name in model_state_dict:
+ val = model_state_dict.pop(layer_name)
+ self._add_to_trtllm_model_weights(val=val, layer_name=layer_name)
+
+ # ----------------Convert Embeddings----------------
+ def _get_remove_vocab_padding(self, layer_name, model_state_dict, tokenizer_vocab_size):
+ val = model_state_dict.get(layer_name, None)
+ if val is None:
+ return None
+
+ if self.inference_tp_size > 1: # Gather padded tensor chunks
+ vocab_size_padded = val.shape[0] * self.inference_tp_size
+ vocab_start_index, vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
+ vocab_size_padded, self.tp_rank, self.inference_tp_size
+ )
+ dim_size = list(val.size())
+ dim_size[0] = vocab_size_padded
+ gathered_val = torch.zeros(
+ dim_size, dtype=val.dtype, device=torch.cuda.current_device()
+ )
+ gathered_val[vocab_start_index:vocab_end_index] = val
+ torch.distributed.all_reduce(gathered_val, group=self.tp_group)
+ val = gathered_val
+ unpadded = val[:tokenizer_vocab_size]
+ if self.inference_tp_size > 1: # Split gathered val for val parallel embedding
+ vocab_start_index, vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
+ tokenizer_vocab_size, self.tp_rank, self.inference_tp_size
+ )
+ unpadded = unpadded[vocab_start_index:vocab_end_index]
+ return unpadded.T # TRTLLM expects (vocab_size, hidden_size) so need extra transpose
+
+ @torch.no_grad()
+ def convert(
+ self, model_state_dict: dict, trtllm_conversion_dict: dict, tokenizer_vocab_size: int
+ ):
+ """Convert model weights to trtllm model weights
+
+ This method goes through each layer in the model state dict and converts to equivalent trtllm model weights. It also handles splitting across TP dimension , expert split etc.
+
+ Args:
+ model_state_dict (dict): The full model state dict (all on CPU)
+ trtllm_conversion_dict (dict): The conversion dictionary used to convert model layer names to trtllm layer names
+ tokenizer_vocab_size (int): The vocab size of the tokenizer
+ """
+
+ # First step is to convert input model layer names to equivalent trtllm layer names
+ model_state_dict = TRTLLMLayers.rename_input_layer_names_to_trtllm_layer_names(
+ model_state_dict=model_state_dict, trtllm_conversion_dict=trtllm_conversion_dict
+ )
+
+ # Convert the non transformer layers
+ for layer_name in NON_TRANSFORMER_LAYERS_NAMES:
+ if (
+ layer_name in TRTLLMLayers.vocab_embedding.value
+ or layer_name in TRTLLMLayers.lm_head.value
+ ):
+ # For embedding layers alone we do some pre processing
+ embed_val = self._get_remove_vocab_padding(
+ layer_name, model_state_dict, tokenizer_vocab_size
+ )
+ model_state_dict[layer_name] = embed_val
+ # TODO : Check if this handling of position embedding is right.
+ if layer_name == TRTLLMLayers.position_embedding.value:
+ position_embedding = model_state_dict[layer_name]
+ req_position_embedding = position_embedding.chunk(self.inference_tp_size)[
+ self.tp_rank
+ ]
+ model_state_dict[layer_name] = req_position_embedding.T
+ self._convert_non_transformer_layer(
+ model_state_dict=model_state_dict, layer_name=layer_name
+ )
+
+ for layer_name, value in tqdm(
+ model_state_dict.items(), desc="Converting to TRTLLM Weights"
+ ):
+ self._convert_transformer_layer(layer_name, value)
diff --git a/megatron/core/export/trtllm/trtllm_weights_converter/single_device_trtllm_model_weights_converter.py b/megatron/core/export/trtllm/trtllm_weights_converter/single_device_trtllm_model_weights_converter.py
new file mode 100644
index 0000000000..c7a98972d2
--- /dev/null
+++ b/megatron/core/export/trtllm/trtllm_weights_converter/single_device_trtllm_model_weights_converter.py
@@ -0,0 +1,437 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import re
+
+import torch
+from tqdm import tqdm
+
+from megatron.core.export.data_type import DataType
+from megatron.core.export.export_config import ExportConfig
+from megatron.core.export.trtllm.trtllm_layers import NON_TRANSFORMER_LAYERS_NAMES, TRTLLMLayers
+from megatron.core.export.trtllm.trtllm_layers import get_layer_name_without_prefix as suffix
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+# pylint: disable=line-too-long
+# TODO: Writing TRT imports this way so that it can be mocked in the test_trtllm_cpu_converter.py unit test
+# TODO: Figure out how to patch it directly from the trtllm library
+def pad_vocab_size(vocab_size: int, tp_size: int):
+ """Pad vocab size based on inference size"""
+ from tensorrt_llm._utils import pad_vocab_size
+
+ return pad_vocab_size(vocab_size, tp_size)
+
+
+def str_dtype_to_torch(dtype: DataType):
+ """Get torch datatype from input datatype"""
+ from tensorrt_llm._utils import str_dtype_to_torch
+
+ return str_dtype_to_torch(dtype.name)
+
+
+class SingleDeviceTRTLLMModelWeightsConverter:
+ """Class to convert Model weights to TRTLLM weights on CPU"""
+
+ def __init__(
+ self,
+ export_config: ExportConfig,
+ transformer_config: TransformerConfig,
+ dtype: DataType,
+ multi_query_mode: bool = False,
+ activation: str = "gelu",
+ ):
+ """Constructor for the TRTLLMModelWeightsConverterCPU class
+
+ This class is responsible to convert the model weights to TRTLLM equivalent weights and also split them for each GPU rank and return as a list.
+
+ Args:
+ export_config (ExportConfig): The export config with inference tp size, pp size etc.
+ transformer_config (TransformerConfig): The transformer config
+ dtype (DataType): The data type or model precision
+ multi_query_mode (bool, optional): Defaults to False.
+ activation (str, optional): Defaults to "gelu".
+ """
+ self.export_config = export_config
+ self.transformer_config = transformer_config
+ self.trtllm_model_weights = {}
+ self.storage_type = str_dtype_to_torch(dtype)
+ self.activation = activation
+ num_kv_heads = self.transformer_config.num_query_groups
+ if num_kv_heads == 0:
+ if multi_query_mode:
+ num_kv_heads = 1
+ else:
+ num_kv_heads = self.transformer_config.num_attention_heads
+ self.num_kv_heads = num_kv_heads
+
+ def _convert_non_transformer_layer(self, model_state_dict: dict, layer_name: str):
+ """Convert Non Transformer layers to TRTLLM weights
+
+ Non transformer layers referes to layers that occur only once in the model (e.g Embedding , final output layer etc. ) They dont have any layer number associated with them. We remove this layer from the original state dict and cast it to storage type and convert to numpy and add it to trtllm_model_weights
+
+ Args:
+ model_state_dict (dict): The input model state dictionary (All collected on CPU)
+ layer_name (str): The TRTLLM Layer name that we want to convert
+ """
+ if layer_name in model_state_dict:
+ val = model_state_dict.pop(layer_name)
+ val = val.to(self.storage_type).detach().contiguous()
+ self.trtllm_model_weights[layer_name] = val
+
+ def _convert_transformer_layer(self, layer_name: str, val: torch.Tensor):
+ """Convert Transformer layers to TRTLLM weights
+
+ Transformer layers referes to layers within the transformber block. They have a layer number associated with them. Depending on the layer we either directly save it to trtllm_model_weights, or split it across some dimension and save the splits
+
+ Args:
+ model_state_dict (dict): The input model state dictionary (All collected on CPU)
+ layer (TRTLLMLayerNames): The TRTLLM Layer that we want to change
+ """
+
+ def _add_to_trtllm_model_weights(val: torch.Tensor, layer_name: str, split_type=None):
+ """Add the input weight to trtllm_model_weights
+
+ Depending on split (Expert split/Tensor split/None) we split the input data and add accordingly
+
+ Args:
+ val (torch.Tensor): The model weight to be added
+ layer_name (str): The TRTLLMlayername as a string
+ split_type (str, optional): The split type. Defaults to None.
+ """
+ if split_type == 'expert_split':
+ for split_num, split_val in enumerate(val):
+ self.trtllm_model_weights[f'{layer_name}.{split_num}.bin'] = (
+ split_val.to(self.storage_type).detach().contiguous()
+ )
+ elif split_type == 'tensor_split':
+ for split_num, split_val in enumerate(val):
+ if split_val.ndim >= 2:
+ split_val = torch.transpose(split_val.reshape(split_val.shape[0], -1), 1, 0)
+
+ self.trtllm_model_weights[f'{layer_name}.{split_num}.bin'] = (
+ split_val.to(self.storage_type).detach().contiguous()
+ )
+ else:
+ if val.ndim >= 2:
+ val = torch.transpose(val.reshape(val.shape[0], -1), 1, 0)
+ self.trtllm_model_weights[layer_name] = (
+ val.to(self.storage_type).detach().contiguous()
+ )
+
+ if val.ndim == 2:
+ val = val.T
+
+ if (
+ layer_name.endswith(suffix(TRTLLMLayers.input_layernorm_weight))
+ or layer_name.endswith(suffix(TRTLLMLayers.input_layernorm_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.post_layernorm_weight))
+ or layer_name.endswith(suffix(TRTLLMLayers.post_layernorm_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.attention_dense_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.attention_dense_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.mlp_projection_bias))
+ or layer_name.endswith(suffix(TRTLLMLayers.mlp_router_weight))
+ ):
+ # Same as layernorm1p in NeMo
+ if (
+ self.transformer_config.layernorm_zero_centered_gamma
+ and self.transformer_config.normalization == "LayerNorm"
+ and 'layernorm.weight' in layer_name
+ ):
+ val = val + 1.0
+
+ _add_to_trtllm_model_weights(val=val, layer_name=layer_name, split_type=None)
+
+ elif layer_name.endswith(
+ suffix(TRTLLMLayers.attention_dense_weight)
+ ) or layer_name.endswith(suffix(TRTLLMLayers.mlp_projection_weight)):
+ split_vals = torch.chunk(val, self.export_config.inference_tp_size, axis=0)
+ _add_to_trtllm_model_weights(
+ val=split_vals, layer_name=layer_name, split_type='tensor_split'
+ )
+
+ elif layer_name.endswith(suffix(TRTLLMLayers.mlp_fc_weight)) or layer_name.endswith(
+ suffix(TRTLLMLayers.mlp_fc_bias)
+ ):
+ split_gated_activation = self.activation in [
+ "swiglu",
+ "geglu",
+ "fast-swiglu",
+ "fast-geglu",
+ ]
+ if split_gated_activation:
+ val, gate = torch.chunk(val, 2, axis=-1)
+ gate_layer_name = layer_name.replace("fc", "gate")
+ split_vals = torch.chunk(gate, self.export_config.inference_tp_size, axis=-1)
+ _add_to_trtllm_model_weights(
+ val=split_vals, layer_name=gate_layer_name, split_type='tensor_split'
+ )
+
+ split_vals = torch.chunk(val, self.export_config.inference_tp_size, axis=-1)
+ _add_to_trtllm_model_weights(
+ val=split_vals, layer_name=layer_name, split_type='tensor_split'
+ )
+
+ elif layer_name.endswith(suffix(TRTLLMLayers.attention_qkv_bias)):
+ qkv_hidden_dim = val.shape[0]
+ size_per_head = qkv_hidden_dim // (
+ self.transformer_config.num_attention_heads + 2 * self.num_kv_heads
+ )
+ q_num = self.transformer_config.num_attention_heads // self.num_kv_heads
+
+ # We first concat all sub weights per tp rank together.
+ val = val.reshape(self.num_kv_heads, q_num + 2, size_per_head)
+
+ qkv = torch.split(val, [q_num, 1, 1], dim=1)
+ q_split = torch.chunk(qkv[0], self.export_config.inference_tp_size, axis=0)
+ k_split = torch.chunk(qkv[1], self.export_config.inference_tp_size, axis=0)
+ v_split = torch.chunk(qkv[2], self.export_config.inference_tp_size, axis=0)
+
+ # Concatenate Q, K, and V together
+ split_vals = [
+ torch.concatenate(
+ [q_split[i].reshape(-1), k_split[i].reshape(-1), v_split[i].reshape(-1)], dim=0
+ )
+ for i in range(self.export_config.inference_tp_size)
+ ]
+ _add_to_trtllm_model_weights(
+ val=split_vals, layer_name=layer_name, split_type='tensor_split'
+ )
+
+ # TODO : Should add a atten layer dimension "qkvqkv, qqkkvv etc to see how to reshape here"
+ elif layer_name.endswith(suffix(TRTLLMLayers.attention_qkv_weight)):
+ hidden_dim = val.shape[0]
+ size_per_head = self.transformer_config.kv_channels
+ if size_per_head is None:
+ size_per_head = hidden_dim // self.transformer_config.num_attention_heads
+ q_num = self.transformer_config.num_attention_heads // self.num_kv_heads
+
+ # When the merge factor exceeds 1, the 'vals' list will have multiple entries.
+ # Depending on the format, 'vals' can look like either [QQQQ..KV, QQQQ..KV, ...](for GQA) or [QKV, QKV, ...](for MHA).
+ # We first concat all sub weights per tp rank together.
+ val = val.reshape(hidden_dim, self.num_kv_heads, q_num + 2, size_per_head)
+
+ # Split the QKV to separate variables.
+ qkv = torch.split(val, [q_num, 1, 1], dim=2)
+
+ query_groups_shape = qkv[0].shape
+ if len(query_groups_shape) > 1:
+ if (query_groups_shape[1] % self.export_config.inference_tp_size) != 0:
+ raise Exception(
+ "Number of query groups of the models is {0}. Please select tensor parallelism size "
+ "that can split the number of query groups to equal number of query matrices in the "
+ "each GPU.".format(query_groups_shape[1])
+ )
+
+ q_split = torch.chunk(qkv[0], self.export_config.inference_tp_size, axis=1)
+ k_split = torch.chunk(qkv[1], self.export_config.inference_tp_size, axis=1)
+ v_split = torch.chunk(qkv[2], self.export_config.inference_tp_size, axis=1)
+
+ # Concatenate Q, K, and V together
+ split_vals = [
+ torch.concatenate(
+ [
+ q_split[i].reshape(hidden_dim, -1),
+ k_split[i].reshape(hidden_dim, -1),
+ v_split[i].reshape(hidden_dim, -1),
+ ],
+ dim=1,
+ )
+ for i in range(self.export_config.inference_tp_size)
+ ]
+ _add_to_trtllm_model_weights(
+ val=split_vals, layer_name=layer_name, split_type='tensor_split'
+ )
+
+ elif layer_name.endswith(suffix(TRTLLMLayers.mlp_fc_weight_mixture_of_experts)):
+ w1, w3 = torch.chunk(val, 2, axis=1)
+ # w1 splits
+ split_w1s = torch.chunk(w1, self.export_config.inference_tp_size, axis=1)
+ # w3 splits
+ split_w3s = torch.chunk(w3, self.export_config.inference_tp_size, axis=1)
+
+ split_vals = [torch.concatenate(item, dim=1) for item in zip(split_w3s, split_w1s)]
+ layer_name = layer_name.replace(".expert", "") # Remove suffix .expert from key
+ _add_to_trtllm_model_weights(
+ val=split_vals, layer_name=layer_name, split_type='expert_split'
+ )
+
+ elif layer_name.endswith(suffix(TRTLLMLayers.mlp_projection_weight_mixture_of_experts)):
+ split_vals = torch.chunk(val, self.export_config.inference_tp_size, axis=-1)
+ layer_name = layer_name.replace(".expert", "") # Remove suffix .expert from key
+ _add_to_trtllm_model_weights(
+ val=split_vals, layer_name=layer_name, split_type='expert_split'
+ )
+ else:
+ raise ValueError(f"{layer_name} cannot be handled by converter")
+
+ @torch.no_grad()
+ def convert(
+ self, model_state_dict: dict, trtllm_conversion_dict, state_dict_split_by_layer_numbers=True
+ ):
+ """Convert model weights to trtllm model weights
+
+ This method goes through each layer in the model state dict and converts to equivalent trtllm model weights. It also handles splitting across TP dimension , expert split etc.
+
+ Args:
+ model_state_dict (dict): The full model state dict (all on CPU)
+ trtllm_conversion_dict (dict): The conversion dictionary used to convert model layer names to trtllm layer names
+ state_dict_split_by_layer_numbers (bool, optional): Are the model layers split by layer numbers in state dict. For example : mlp.fc1.weight can be represented like mlp.fc1.weight of shape [num_layers, hidden_dim, ffn_hidden_dim]} or it can be like mlp.fc1.layers.0.weight of shape [hidden_dim, ffn_hidden_dim], then mlp.fc1.layers.1.weight ... for all layers. If you use represenation 2 set this to True. Defaults to True
+ """
+
+ # First step is to convert input model layer names to equivalent trtllm layer names
+ model_state_dict = TRTLLMLayers.rename_input_layer_names_to_trtllm_layer_names(
+ model_state_dict=model_state_dict,
+ trtllm_conversion_dict=trtllm_conversion_dict,
+ state_dict_split_by_layer_numbers=state_dict_split_by_layer_numbers,
+ )
+
+ # Convert the non transformer layers
+ for layer_name in NON_TRANSFORMER_LAYERS_NAMES:
+ # For vocab embedding layer alone we pad the weights to be divisible by inference tp size
+ if (
+ layer_name == TRTLLMLayers.vocab_embedding.value
+ and self.export_config.use_parallel_embedding
+ ):
+ val = model_state_dict[TRTLLMLayers.vocab_embedding.value]
+ vocab_size = val.shape[0]
+ if vocab_size % self.export_config.inference_tp_size != 0:
+ vocab_size_padded = pad_vocab_size(
+ vocab_size, self.export_config.inference_tp_size
+ )
+ pad_width = vocab_size_padded - vocab_size
+ val = torch.nn.functional.pad(val, (0, 0, 0, pad_width), value=0)
+ model_state_dict[layer_name] = val
+
+ self._convert_non_transformer_layer(
+ model_state_dict=model_state_dict, layer_name=layer_name
+ )
+
+ transformer_layers_dict = {}
+ # Convert the transformer layers
+ if state_dict_split_by_layer_numbers:
+ # Already model dict is split by layer numbers
+ transformer_layers_dict = model_state_dict
+ else:
+ # Here we split the model state dict into individual layers
+ for layer_name in list(model_state_dict.keys()):
+ value = model_state_dict.pop(layer_name)
+ for layer_number in range(self.transformer_config.num_layers):
+ # e.g transformer.layers.mlp.fc.bias => transformer.layers.2.mlp.fc.bias
+ layer_name_with_layer_number = re.sub(
+ r'(?<=layers\.)', f'{layer_number}.', layer_name
+ )
+ transformer_layers_dict[layer_name_with_layer_number] = value[layer_number]
+
+ for layer_name, value in tqdm(
+ transformer_layers_dict.items(), desc="Converting to TRTLLM Weights"
+ ):
+ self._convert_transformer_layer(layer_name, value)
+
+ def get_padded_vocab_size(self) -> int:
+ """Return the paded vocab size
+
+ We extract the lm head and vocab embedding and use that to determine padded_vocab_size
+
+ Returns:
+ int: Padded vocab size
+ """
+ lm_head_weight = self.trtllm_model_weights.get(TRTLLMLayers.lm_head.value, None)
+ vocab_size = self.trtllm_model_weights[TRTLLMLayers.vocab_embedding.value].shape[0]
+ vocab_size_padded = (
+ vocab_size
+ if lm_head_weight is None
+ else pad_vocab_size(vocab_size, self.export_config.inference_tp_size)
+ )
+ return vocab_size_padded
+
+ def get_local_model_weights_per_gpu(self, mapping, trtllm_model_config: dict):
+ """Get the trtllm model weights split per gpu
+
+ Given the trtllm mapping information (tp, pp rank etc) we split the model weights in a list, with each element of the list corresponding to the weights of each gpu rank
+
+ Args:
+ mapping : The trtllm mapping information
+ trtllm_model_config (dict): The trtllm model config
+ """
+
+ def _split(torch_tensor, tp_size, idx, dim=0):
+ """Splits the np tensor v on dim and return the idx's slice."""
+ if tp_size == 1:
+ return torch_tensor
+ if len(torch_tensor.shape) == 1:
+ return torch.chunk(torch_tensor, tp_size)[idx].contiguous()
+ else:
+ return torch.chunk(torch_tensor, tp_size, axis=dim)[idx].contiguous()
+
+ pp_layer_range = mapping.pp_layers(self.transformer_config.num_layers)
+
+ trtllm_model_weights_per_gpu = {}
+ for layer_name, value in self.trtllm_model_weights.items():
+ if layer_name in NON_TRANSFORMER_LAYERS_NAMES:
+ continue
+
+ # Happens in the case of TP split or expert split
+ if layer_name.endswith(".bin"):
+ if layer_name.endswith(f"{mapping.tp_rank}.bin"):
+ layer_name = layer_name.replace(f".{mapping.tp_rank}.bin", "")
+ else:
+ continue
+
+ layer_num = int(layer_name.split(".")[2])
+ if layer_num in pp_layer_range:
+ layer_name = layer_name.replace(
+ f"layers.{layer_num}", f"layers.{layer_num - pp_layer_range[0]}"
+ )
+ else:
+ continue
+ if (
+ hasattr(trtllm_model_config, 'new_decoder_architecture')
+ and trtllm_model_config.new_decoder_architecture
+ and "post_layernorm" in layer_name
+ ):
+ layer_name = layer_name.replace("post_layernorm", "mlp_layernorm")
+
+ trtllm_model_weights_per_gpu[layer_name] = value
+
+ if mapping.is_first_pp_rank():
+ embedding_weight = (
+ _split(
+ self.trtllm_model_weights[TRTLLMLayers.vocab_embedding.value],
+ mapping.tp_size,
+ mapping.tp_rank,
+ )
+ if self.export_config.use_parallel_embedding
+ else self.trtllm_model_weights[TRTLLMLayers.vocab_embedding.value]
+ )
+
+ trtllm_model_weights_per_gpu[TRTLLMLayers.vocab_embedding.value] = embedding_weight
+
+ pos_embedding_weight = self.trtllm_model_weights.get(
+ TRTLLMLayers.position_embedding.value
+ )
+ if pos_embedding_weight is not None:
+ if self.export_config.use_parallel_embedding:
+ pos_embedding_weight = _split(
+ pos_embedding_weight, mapping.tp_size, mapping.tp_rank
+ )
+
+ trtllm_model_weights_per_gpu[TRTLLMLayers.position_embedding.value] = (
+ pos_embedding_weight
+ )
+
+ if mapping.is_last_pp_rank():
+ lm_head_weight = self.trtllm_model_weights.get(TRTLLMLayers.lm_head.value, None)
+ if lm_head_weight is not None:
+ trtllm_model_weights_per_gpu[TRTLLMLayers.lm_head.value] = _split(
+ lm_head_weight, mapping.tp_size, mapping.tp_rank
+ )
+
+ trtllm_model_weights_per_gpu[TRTLLMLayers.final_layernorm_weight.value] = (
+ self.trtllm_model_weights[TRTLLMLayers.final_layernorm_weight.value]
+ )
+
+ ln_f_bias = self.trtllm_model_weights.get(TRTLLMLayers.final_layernorm_bias.value)
+ if ln_f_bias is not None:
+ trtllm_model_weights_per_gpu[TRTLLMLayers.final_layernorm_bias.value] = ln_f_bias
+
+ return trtllm_model_weights_per_gpu
diff --git a/megatron/mpu/tests/__init__.py b/megatron/core/extensions/__init__.py
similarity index 100%
rename from megatron/mpu/tests/__init__.py
rename to megatron/core/extensions/__init__.py
diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
new file mode 100644
index 0000000000..bf5159c759
--- /dev/null
+++ b/megatron/core/extensions/transformer_engine.py
@@ -0,0 +1,1087 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import dataclasses
+import os
+import warnings
+from typing import Callable
+
+import torch
+import transformer_engine as te
+from packaging.version import Version as PkgVersion
+from torch import Tensor
+from torch.nn.parameter import Parameter
+
+from megatron.core import ModelParallelConfig, parallel_state
+from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.parallel_state import (
+ get_context_parallel_global_ranks,
+ get_context_parallel_group,
+ get_tensor_and_expert_parallel_world_size,
+ get_tensor_model_parallel_group,
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+)
+from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
+from megatron.core.tensor_parallel.layers import (
+ _initialize_affine_weight_cpu,
+ set_tensor_model_parallel_attributes,
+)
+from megatron.core.tensor_parallel.utils import divide
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+from megatron.core.utils import get_te_version, is_te_min_version
+
+
+def _get_extra_te_kwargs(config: TransformerConfig):
+ extra_transformer_engine_kwargs = {"params_dtype": config.params_dtype}
+
+ if is_te_min_version("0.12.0"):
+ if config.use_cpu_initialization:
+ extra_transformer_engine_kwargs["device"] = 'cpu'
+ else:
+ extra_transformer_engine_kwargs["device"] = torch.cuda.current_device()
+ return extra_transformer_engine_kwargs
+
+
+def condition_init_method(config, init_method):
+ """Condition TE init_method on config.perform_initialization."""
+ return init_method if config.perform_initialization else (lambda w: None)
+
+
+class TENorm:
+ """
+ A conditional wrapper to initialize an instance of Transformer-Engine's
+ `LayerNorm` or `RMSNorm` based on input
+ """
+
+ # TODO should we ditch normalization config and just use spec to choose LayerNorm vs RMSNorm?
+ def __new__(cls, config: TransformerConfig, hidden_size: int, eps: float = 1e-5):
+ if config.normalization == "LayerNorm":
+ instance = te.pytorch.LayerNorm(
+ hidden_size=hidden_size,
+ eps=eps,
+ sequence_parallel=config.sequence_parallel,
+ zero_centered_gamma=config.layernorm_zero_centered_gamma,
+ **_get_extra_te_kwargs(config),
+ )
+ elif config.normalization == "RMSNorm":
+ assert hasattr(
+ te.pytorch, "RMSNorm"
+ ), "Transformer-Engine >= v0.11 required to use this feature"
+ instance = te.pytorch.RMSNorm(
+ hidden_size=hidden_size,
+ eps=eps,
+ sequence_parallel=config.sequence_parallel,
+ zero_centered_gamma=config.layernorm_zero_centered_gamma,
+ **_get_extra_te_kwargs(config),
+ )
+ else:
+ raise Exception('Only LayerNorm and RMSNorm are curently supported')
+
+ return instance
+
+
+class TELinear(te.pytorch.Linear):
+ """
+ Wrapper for the Transformer-Engine's `Linear` layer.
+
+ Note that if Megatron's parallel_state has not been initialized
+ yet, the tp_group passed to TE will be None and must be set later
+ via set_tensor_parallel_group().
+ """
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ parallel_mode: str,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ skip_bias_add: bool,
+ skip_weight_param_allocation: bool,
+ tp_comm_buffer_name: str = None,
+ is_expert: bool = False,
+ ):
+ self.config = config
+
+ # TE returns a zero length Tensor when bias=False and
+ # return_bias=True, but we prefer None. So in that case we
+ # tell TE to not return the bias, and return None
+ # ourselves. This way our forward always returns two values
+ # and we don't have to deal with the zero length Tensor.
+ self.te_return_bias = skip_bias_add and bias
+ self.is_first_microbatch = True
+ self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
+ if skip_weight_param_allocation:
+ raise ValueError(
+ 'Transformer Engine linear layers do not support skip_weight_param_allocation'
+ )
+
+ extra_kwargs = _get_extra_te_kwargs(config)
+
+ if is_te_min_version("0.8.0"):
+ if self.config.tp_comm_overlap:
+ if is_te_min_version("1.5.0"):
+ # Use old overlap flags if they were supplied instead
+ extra_kwargs["ub_overlap_ag"] = (
+ self.config.tp_comm_overlap_ag
+ if hasattr(self.config, "tp_comm_overlap_ag")
+ else self.config.tp_comm_split_ag or self.config.tp_comm_atomic_ag
+ )
+ extra_kwargs["ub_overlap_rs"] = (
+ self.config.tp_comm_overlap_rs
+ if hasattr(self.config, "tp_comm_overlap_rs")
+ else self.config.tp_comm_split_rs or self.config.tp_comm_atomic_rs
+ )
+ # Disable ub overlap for experts.
+ if is_expert:
+ extra_kwargs["ub_overlap_ag"] = False
+ extra_kwargs["ub_overlap_rs"] = False
+ else:
+ extra_kwargs["ub_split_ag"] = self.config.tp_comm_split_ag
+ extra_kwargs["ub_atomic_gemm_ag"] = self.config.tp_comm_atomic_ag
+ extra_kwargs["ub_split_rs"] = self.config.tp_comm_split_rs
+ extra_kwargs["ub_atomic_gemm_rs"] = self.config.tp_comm_atomic_rs
+ # Disable ub overlap for experts.
+ if is_expert:
+ extra_kwargs["ub_split_ag"] = False
+ extra_kwargs["ub_atomic_gemm_ag"] = False
+ extra_kwargs["ub_split_rs"] = False
+ extra_kwargs["ub_atomic_gemm_rs"] = False
+ if is_te_min_version("1.0.0", check_equality=False):
+ assert (
+ tp_comm_buffer_name is not None
+ ), "Buffer name should be set to configure communication overlap settings"
+ extra_kwargs["ub_name"] = tp_comm_buffer_name
+
+ self.expert_parallel = self.config.expert_model_parallel_size > 1
+ if is_expert and self.expert_parallel:
+ rng_tracker_name = get_expert_parallel_rng_tracker_name()
+ else:
+ rng_tracker_name = None
+ if is_te_min_version("1.7.0"):
+ extra_kwargs["rng_tracker_name"] = rng_tracker_name
+
+ # Disable communications in TE when using SP or EP by making TE agnostic of model parallel.
+ tp_size = self.config.tensor_model_parallel_size
+ tp_group = get_tensor_model_parallel_group(check_initialized=False)
+ if is_expert and (self.config.sequence_parallel or self.expert_parallel):
+ if self.config.moe_extended_tp:
+ tp_size = get_tensor_and_expert_parallel_world_size()
+ if parallel_mode == "column":
+ output_size = divide(output_size, tp_size)
+ elif parallel_mode == "row":
+ input_size = divide(input_size, tp_size)
+ parallel_mode = None
+ tp_size = 1
+ tp_group = None
+
+ super().__init__(
+ in_features=input_size,
+ out_features=output_size,
+ sequence_parallel=self.config.sequence_parallel,
+ fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
+ tp_group=tp_group,
+ tp_size=tp_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ return_bias=self.te_return_bias,
+ parallel_mode=parallel_mode,
+ **extra_kwargs,
+ )
+
+ for param in self.parameters():
+ setattr(param, 'allreduce', not (is_expert and self.expert_parallel))
+
+ def forward(self, x):
+ """Forward."""
+ _is_first_microbatch = (
+ None if self.disable_parameter_transpose_cache else self.is_first_microbatch
+ )
+ out = super().forward(x, is_first_microbatch=_is_first_microbatch)
+ self.is_first_microbatch = False
+
+ # TE only returns a tuple when return_bias is True, otherwise
+ # it returns a single Tensor, we always want to return two
+ # values regardless of the arguments.
+ if self.te_return_bias:
+ return out
+ return out, None
+
+
+class TELayerNormColumnParallelLinear(te.pytorch.LayerNormLinear):
+ """
+ Wrapper for the Transformer-Engine's `LayerNormLinear` layer that combines
+ layernorm and linear layers
+ """
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ config: TransformerConfig,
+ init_method: Callable,
+ gather_output: bool,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ skip_weight_param_allocation: bool = False,
+ tp_comm_buffer_name: str = None,
+ ):
+ self.config = config
+
+ if gather_output:
+ raise ValueError('Transformer Engine linear layers do not support gather_output = True')
+
+ if is_expert:
+ raise ValueError('Transformer Engine linear layers do not yet support MoE')
+
+ if skip_weight_param_allocation:
+ raise ValueError(
+ 'Transformer Engine linear layers do not support skip_weight_param_allocation'
+ )
+
+ # TE returns a zero length Tensor when bias=False and
+ # return_bias=True, but we prefer None. So in that case we
+ # tell TE to not return the bias, and return None
+ # ourselves. This way our forward always returns two values
+ # and we don't have to deal with the zero length Tensor.
+ self.te_return_bias = skip_bias_add and bias
+ self.is_first_microbatch = True
+ self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
+ extra_kwargs = _get_extra_te_kwargs(config)
+
+ # Only Transformer-Engine version >= 0.11.0 supports `RMSNorm`
+ if is_te_min_version("0.11.0"):
+ extra_kwargs["normalization"] = self.config.normalization
+ elif self.config.normalization != "LayerNorm":
+ te_version = get_te_version()
+ raise ValueError(
+ f"Transformer Engine v{te_version} does not support {self.config.normalization}."
+ )
+
+ if is_te_min_version("0.8.0"):
+ if self.config.tp_comm_overlap:
+ extra_kwargs["ub_bulk_wgrad"] = self.config.tp_comm_bulk_wgrad
+ extra_kwargs["ub_bulk_dgrad"] = self.config.tp_comm_bulk_dgrad
+ if is_te_min_version("1.5.0", check_equality=False):
+ # Use old overlap flags if they were supplied instead
+ extra_kwargs["ub_overlap_ag"] = (
+ self.config.tp_comm_overlap_ag
+ if hasattr(self.config, "tp_comm_overlap_ag")
+ else self.config.tp_comm_split_ag or self.config.tp_comm_atomic_ag
+ )
+ if is_te_min_version("1.6.0.dev0", check_equality=False):
+ extra_kwargs["ub_overlap_rs_dgrad"] = (
+ self.config.tp_comm_overlap_rs_dgrad
+ if hasattr(self.config, "tp_comm_overlap_rs_dgrad")
+ else False
+ )
+ if tp_comm_buffer_name == 'qkv' and self.config.tp_comm_overlap_disable_qkv:
+ extra_kwargs["ub_overlap_ag"] = False
+ extra_kwargs["ub_overlap_rs_dgrad"] = False
+
+ if tp_comm_buffer_name == 'fc1' and self.config.tp_comm_overlap_disable_fc1:
+ extra_kwargs["ub_overlap_ag"] = False
+ extra_kwargs["ub_overlap_rs_dgrad"] = False
+ else:
+ extra_kwargs["ub_atomic_gemm_ag"] = self.config.tp_comm_atomic_ag
+ extra_kwargs["ub_split_ag"] = self.config.tp_comm_split_ag
+ if is_te_min_version("1.0.0", check_equality=False):
+ assert (
+ tp_comm_buffer_name is not None
+ ), "Buffer name should be set to configure communication overlap settings"
+ extra_kwargs["ub_name"] = tp_comm_buffer_name
+
+ super().__init__(
+ in_features=input_size,
+ out_features=output_size,
+ eps=self.config.layernorm_epsilon,
+ sequence_parallel=self.config.sequence_parallel,
+ fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
+ tp_group=get_tensor_model_parallel_group(check_initialized=False),
+ tp_size=self.config.tensor_model_parallel_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ init_method=(
+ condition_init_method(config, init_method)
+ if not config.use_cpu_initialization
+ else lambda w: None
+ ),
+ bias=bias,
+ return_bias=self.te_return_bias,
+ parallel_mode="column",
+ return_layernorm_output=False,
+ zero_centered_gamma=self.config.layernorm_zero_centered_gamma,
+ **extra_kwargs,
+ )
+
+ world_size = get_tensor_model_parallel_world_size()
+ rank = get_tensor_model_parallel_rank()
+
+ if config.use_cpu_initialization:
+ output_size_per_partition = divide(output_size, world_size)
+ _ = _initialize_affine_weight_cpu(
+ self.weight,
+ output_size,
+ input_size,
+ output_size_per_partition,
+ 0,
+ init_method,
+ stride=1,
+ return_master_weight=False,
+ rank=rank,
+ world_size=world_size,
+ skip_set_tensor_parallel_attributes=True,
+ )
+ if bias:
+ self.bias = Parameter(
+ torch.empty(output_size_per_partition, dtype=config.params_dtype)
+ )
+ set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ with torch.no_grad():
+ self.bias.zero_()
+ setattr(self.bias, 'allreduce', True)
+
+ def forward(self, x):
+ """Forward."""
+ _is_first_microbatch = (
+ None if self.disable_parameter_transpose_cache else self.is_first_microbatch
+ )
+ out = super().forward(x, is_first_microbatch=_is_first_microbatch)
+ self.is_first_microbatch = False
+
+ # TE only returns a tuple when return_bias is True, otherwise
+ # it returns a single Tensor, we always want to return two
+ # values regardless of the arguments.
+ if self.te_return_bias:
+ return out
+ return out, None
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Sharding along axis 0, bias sharded"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(
+ state_dict, prefix, {'weight': 0, 'bias': 0}, sharded_offsets
+ )
+
+
+class TEColumnParallelLinear(TELinear):
+ """
+ Wrapper for the Transformer-Engine's `Linear` layer but specialized similar
+ to megatron's `ColumnParallelLinear` layer.
+ """
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ gather_output: bool,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ skip_weight_param_allocation: bool = False,
+ tp_comm_buffer_name: str = None,
+ ):
+ if gather_output:
+ raise ValueError('Transformer Engine linear layers do not support gather_output = True')
+
+ super().__init__(
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="column",
+ config=config,
+ init_method=(
+ condition_init_method(config, init_method)
+ if not config.use_cpu_initialization
+ else lambda w: None
+ ),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ is_expert=is_expert,
+ skip_weight_param_allocation=skip_weight_param_allocation,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ )
+
+ world_size = get_tensor_model_parallel_world_size()
+ rank = get_tensor_model_parallel_rank()
+ if config.use_cpu_initialization:
+ output_size_per_partition = divide(output_size, world_size)
+ _ = _initialize_affine_weight_cpu(
+ self.weight,
+ output_size,
+ input_size,
+ output_size_per_partition,
+ 0,
+ init_method,
+ stride=1,
+ return_master_weight=False,
+ rank=rank,
+ world_size=world_size,
+ skip_set_tensor_parallel_attributes=True,
+ )
+ if bias:
+ self.bias = Parameter(
+ torch.empty(output_size_per_partition, dtype=config.params_dtype)
+ )
+ set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
+ with torch.no_grad():
+ self.bias.zero_()
+ setattr(self.bias, 'allreduce', True)
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Sharding along axis 0, bias sharded"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(
+ state_dict, prefix, {'weight': 0, 'bias': 0}, sharded_offsets
+ )
+
+
+class TERowParallelLinear(TELinear):
+ """
+ Wrapper for the Transformer-Engine's `Linear` layer but specialized similar
+ to megatron's `RowParallelLinear` layer.
+ """
+
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ input_is_parallel: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ tp_comm_buffer_name: str = None,
+ ):
+ if not input_is_parallel:
+ raise ValueError(
+ "Transformer Engine linear layers do not support input_is_parallel = False"
+ )
+
+ super().__init__(
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="row",
+ config=config,
+ init_method=(
+ condition_init_method(config, init_method)
+ if not config.use_cpu_initialization
+ else lambda w: None
+ ),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ skip_weight_param_allocation=False, # We don't currently use this for row parallel layers # pylint: disable=line-too-long
+ is_expert=is_expert,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ )
+ world_size = get_tensor_model_parallel_world_size()
+ rank = get_tensor_model_parallel_rank()
+ if config.use_cpu_initialization:
+ input_size_per_partition = divide(input_size, world_size)
+ self.master_weight = _initialize_affine_weight_cpu(
+ self.weight,
+ output_size,
+ input_size,
+ input_size_per_partition,
+ 1,
+ init_method,
+ stride=1,
+ return_master_weight=False,
+ params_dtype=config.params_dtype,
+ rank=rank,
+ world_size=world_size,
+ skip_set_tensor_parallel_attributes=True,
+ )
+ if bias:
+ self.bias = Parameter(torch.empty(output_size, dtype=config.params_dtype))
+ # Always initialize bias to zero.
+ with torch.no_grad():
+ self.bias.zero_()
+ setattr(self.bias, 'allreduce', True)
+ setattr(self.bias, 'sequence_parallel', config.sequence_parallel)
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Sharding along axis 1, bias not sharded"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(
+ state_dict, prefix, {'weight': 1}, sharded_offsets
+ )
+
+
+class TEDotProductAttention(te.pytorch.DotProductAttention):
+ """
+ Wrapper for the Transformer-Engine's `DotProductAttention` layer that also
+ has "flash attention" enabled.
+
+ Note that if Megatron's parallel_state has not been initialized yet, the
+ tp_group and cp_group passed to TE will be None and must be set later
+ via set_tensor_parallel_group() and set_context_parallel_group().
+ """
+
+ cp_stream: torch.cuda.Stream = None
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ layer_number: int,
+ attn_mask_type: AttnMaskType,
+ attention_type: str,
+ attention_dropout: float = None,
+ softmax_scale: float = None,
+ k_channels: int = None,
+ v_channels: int = None,
+ ):
+ self.config = config
+ self.te_forward_mask_type = False
+ self.qkv_format: str = 'sbhd'
+
+ if self.config.apply_query_key_layer_scaling != bool(
+ int(os.getenv('NVTE_APPLY_QK_LAYER_SCALING', '0'))
+ ):
+ raise ValueError(
+ f"apply_query_key_layer_scaling is {self.config.apply_query_key_layer_scaling} "
+ f"but environment variable NVTE_APPLY_QK_LAYER_SCALING is "
+ f"{os.getenv('NVTE_APPLY_QK_LAYER_SCALING')}. Transformer Engine does not support "
+ f"setting query key layer scaling via argument, so these two must match."
+ )
+
+ extra_kwargs = {}
+ if is_te_min_version("0.11.0"):
+ extra_kwargs["num_gqa_groups"] = self.config.num_query_groups
+ elif self.config.num_query_groups != self.config.num_attention_heads:
+ raise ValueError(
+ f"Transformer Engine v{get_te_version()} does not support Grouped Query Attention, "
+ f"use a newer version of Transformer Engine. "
+ f"(num_query_groups ({self.config.num_query_groups}) != "
+ f"num_attention_heads ({self.config.num_attention_heads}))"
+ )
+
+ if is_te_min_version("0.10.0"):
+ extra_kwargs["attention_type"] = attention_type
+ # older version don't need attention_type
+
+ if is_te_min_version("0.12.0", check_equality=False):
+ self.te_forward_mask_type = True
+
+ # Only Transformer-Engine version >= 1.0.0 supports context parallelism
+ if is_te_min_version("1.0.0"):
+ if getattr(TEDotProductAttention, "cp_stream") is None:
+ TEDotProductAttention.cp_stream = torch.cuda.Stream()
+ extra_kwargs["cp_group"] = get_context_parallel_group(check_initialized=False)
+ extra_kwargs["cp_global_ranks"] = get_context_parallel_global_ranks(
+ check_initialized=False
+ )
+ extra_kwargs["cp_stream"] = TEDotProductAttention.cp_stream
+ else:
+ assert (
+ self.config.context_parallel_size == 1
+ ), "Only Transformer-Engine version >= 1.0.0 supports context parallelism!"
+
+ if self.config.deterministic_mode:
+ if int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")) != 0:
+ raise RuntimeError(
+ "deterministic_mode is on and we are using DotProductAttention from "
+ "Transformer Engine, but NVTE_ALLOW_NONDETERMINISTIC_ALGO is not 0. "
+ f"Currently set to: {os.getenv('NVTE_ALLOW_NONDETERMINISTIC_ALGO', 'not set')}."
+ )
+
+ if config.window_size is not None:
+ # Check version
+ assert is_te_min_version("1.2.0"), (
+ f"Transformer-Engine v{get_te_version()} must be >= 1.2.0 to support"
+ "sliding window attention."
+ )
+ extra_kwargs['window_size'] = config.window_size
+
+ if is_te_min_version("1.10.0"):
+ # TE 1.10.0 introduces the ability to set the different k and v channels
+ kv_channels = (
+ (k_channels, v_channels)
+ if k_channels is not None and v_channels is not None
+ else self.config.kv_channels
+ )
+ extra_kwargs['softmax_scale'] = softmax_scale
+ else:
+ kv_channels = self.config.kv_channels
+
+ super().__init__(
+ num_attention_heads=self.config.num_attention_heads,
+ kv_channels=kv_channels,
+ attention_dropout=(
+ self.config.attention_dropout if attention_dropout is None else attention_dropout
+ ),
+ attn_mask_type=attn_mask_type.name,
+ sequence_parallel=self.config.sequence_parallel,
+ tp_size=self.config.tensor_model_parallel_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ tp_group=get_tensor_model_parallel_group(check_initialized=False),
+ layer_number=layer_number,
+ **extra_kwargs,
+ )
+
+ def forward(
+ self,
+ query: Tensor,
+ key: Tensor,
+ value: Tensor,
+ attention_mask: Tensor,
+ attn_mask_type: AttnMaskType,
+ packed_seq_params: PackedSeqParams = None,
+ ):
+ """Forward."""
+ packed_seq_kwargs = (
+ dataclasses.asdict(packed_seq_params) if packed_seq_params is not None else {}
+ )
+ # overwrite self.qkv_format depending on self.config.apply_rope_fusion, which can be set
+ # after init
+ if self.config.apply_rope_fusion and is_te_min_version("0.13.0", check_equality=False):
+ self.qkv_format = 'bshd'
+
+ qkv_format = packed_seq_kwargs.get('qkv_format', self.qkv_format)
+
+ if get_te_version() < PkgVersion("1.3.0"):
+ # TE 1.3.0 introduces precomputing max_seqlen to remove unnecessary kernels and D2H
+ # copies (#555)
+ # These two arguments did not exist prior to 1.3.0
+ packed_seq_kwargs.pop("max_seqlen_q", None)
+ packed_seq_kwargs.pop("max_seqlen_kv", None)
+
+ if get_te_version() < PkgVersion("1.10.0"):
+ # TE 1.8.0 introduces cu_seqlens_padded which is the cu_seqlens with paddings counted
+ # in each individual sequence in THD format dataset
+ # These two arguments did not exist prior to 1.8.0.Full support added in 1.10.0 (#1012)
+ packed_seq_kwargs.pop("cu_seqlens_q_padded", None)
+ packed_seq_kwargs.pop("cu_seqlens_kv_padded", None)
+
+ if self.config.apply_rope_fusion and qkv_format == 'bshd':
+ query, key, value = [x.transpose(0, 1).contiguous() for x in (query, key, value)]
+ # In PyTorch, the following two tensors are in fact the same:
+ # Tensor with shape (1, S, H, D) and stride (S*H*D, H*D, D, 1)
+ # Tensor with shape (1, S, H, D) and stride (H*D, H*D, D, 1)
+ # Stride for a dimension that is 1 has no meaning, so tensors created two different ways
+ # can have same shape but different strides.
+ # We unify them to the first one to pass the stride check in TE
+ if value.shape == key.shape and value.shape[0] == 1 and value.stride() != key.stride():
+ value = value.as_strided(value.shape, key.stride())
+
+ if self.te_forward_mask_type:
+ if qkv_format == 'thd' and is_te_min_version("1.7.0"):
+ # thd format uses flash attention with cuDNN kernel which requires is_padding=True,
+ # so the only acceptable mask types are `padding_causal` and `padding`. These do not
+ # necessarily indicate there are padded tokens in the sequence.
+ if attn_mask_type == AttnMaskType.causal:
+ attn_mask_type = AttnMaskType.padding_causal
+ elif attn_mask_type == AttnMaskType.no_mask:
+ attn_mask_type = AttnMaskType.padding
+ core_attn_out = super().forward(
+ query,
+ key,
+ value,
+ attention_mask,
+ attn_mask_type=attn_mask_type.name,
+ **packed_seq_kwargs,
+ )
+ else:
+ core_attn_out = super().forward(query, key, value, attention_mask, **packed_seq_kwargs)
+
+ if self.config.apply_rope_fusion and qkv_format == 'bshd':
+ return core_attn_out.transpose(0, 1)
+ else:
+ return core_attn_out
+
+
+if is_te_min_version("1.9.0.dev0"):
+
+ class TEGroupedLinear(te.pytorch.GroupedLinear):
+ """
+ Wrapper for the Transformer-Engine's `GroupedLinear` layer.
+
+ Note that if Megatron's parallel_state has not been initialized
+ yet, the tp_group passed to TE will be None and must be set later
+ via set_tensor_parallel_group().
+ """
+
+ def __init__(
+ self,
+ num_gemms: int,
+ input_size: int,
+ output_size: int,
+ *,
+ parallel_mode: str,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool = False,
+ tp_comm_buffer_name: str = None,
+ ):
+ self.config = config
+
+ # TE returns a zero length Tensor when bias=False and
+ # return_bias=True, but we prefer None. So in that case we
+ # tell TE to not return the bias, and return None
+ # ourselves. This way our forward always returns two values
+ # and we don't have to deal with the zero length Tensor.
+ self.te_return_bias = skip_bias_add and bias
+ self.is_first_microbatch = True
+ self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache
+
+ extra_kwargs = _get_extra_te_kwargs(config)
+ extra_kwargs["ub_name"] = tp_comm_buffer_name
+
+ self.expert_parallel = self.config.expert_model_parallel_size > 1
+ if self.expert_parallel:
+ extra_kwargs["rng_tracker_name"] = get_expert_parallel_rng_tracker_name()
+
+ # For MoE models, the comms between TP and EP group is explicitly handled by
+ # MoE token dispatcher. So we disable comms by making TE agnostic of model parallel.
+ self.explicit_expert_comm = is_expert and (
+ config.tensor_model_parallel_size > 1 or self.expert_parallel
+ )
+ tp_group = get_tensor_model_parallel_group(check_initialized=False)
+ if self.explicit_expert_comm and config.moe_extended_tp:
+ tp_size = parallel_state.get_tensor_and_expert_parallel_world_size()
+ else:
+ tp_size = parallel_state.get_tensor_model_parallel_world_size()
+ if self.explicit_expert_comm:
+ if parallel_mode == "column":
+ output_size = divide(output_size, tp_size)
+ elif parallel_mode == "row":
+ input_size = divide(input_size, tp_size)
+ parallel_mode = None
+ tp_size = 1
+ tp_group = None
+
+ super().__init__(
+ num_gemms=num_gemms,
+ in_features=input_size,
+ out_features=output_size,
+ sequence_parallel=self.config.sequence_parallel,
+ fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
+ tp_group=tp_group,
+ tp_size=tp_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ return_bias=self.te_return_bias,
+ parallel_mode=parallel_mode,
+ **extra_kwargs,
+ )
+
+ for param in self.parameters():
+ setattr(param, 'allreduce', not (is_expert and self.expert_parallel))
+
+ def forward(self, x, m_splits):
+ """Forward."""
+ _is_first_microbatch = (
+ None if self.disable_parameter_transpose_cache else self.is_first_microbatch
+ )
+ out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch)
+ self.is_first_microbatch = False
+
+ # TE only returns a tuple when return_bias is True, otherwise
+ # it returns a single Tensor, we always want to return two
+ # values regardless of the arguments.
+ if self.te_return_bias:
+ return out
+ return out, None
+
+ def _sharded_state_dict_grouped(
+ self, tp_axis_map, prefix='', sharded_offsets=(), metadata=None
+ ):
+ """
+ prefix should be module_name to make keys identical to sequetial ones.
+ """
+ sharded_state_dict = {}
+ full_state_dict = self.state_dict(prefix='', keep_vars=True)
+ num_global_experts = (
+ parallel_state.get_expert_model_parallel_world_size() * self.num_gemms
+ )
+ local_expert_indices_offset = (
+ parallel_state.get_expert_model_parallel_rank() * self.num_gemms
+ )
+ ep_axis = len(sharded_offsets)
+ for gemm_idx in range(self.num_gemms):
+ state_dict = {
+ f'{gemm_idx}.weight': full_state_dict[f'weight{gemm_idx}'],
+ f'{gemm_idx}._extra_state': full_state_dict['_extra_state'],
+ }
+ if self.use_bias:
+ state_dict[f'{gemm_idx}.bias'] = full_state_dict[f'bias{gemm_idx}']
+ sub_sd = make_sharded_tensors_for_checkpoint(
+ state_dict,
+ '',
+ tp_axis_map,
+ (
+ *sharded_offsets,
+ (ep_axis, local_expert_indices_offset + gemm_idx, num_global_experts),
+ ),
+ )
+ # Remove expert layers indexing from sharded keys
+ replace_prefix_for_sharding(sub_sd, f'{gemm_idx}.', prefix)
+ sharded_state_dict.update(
+ {
+ f'{prefix}weight{gemm_idx}': sub_sd[f'{gemm_idx}.weight'],
+ # TODO: TE's GroupedLinear only has one _extra_state for all experts.
+ # We need sharding or build/merge fn to handle _extra_state correctly.
+ f'{prefix}_extra_state{"" if gemm_idx == 0 else gemm_idx}': sub_sd[
+ f'{gemm_idx}._extra_state'
+ ],
+ }
+ )
+ if self.use_bias:
+ sharded_state_dict[f'{prefix}bias{gemm_idx}'] = sub_sd[f'{gemm_idx}.bias']
+ # Adjust replica ids - replication along DP modulo EP
+ for k, sh_ten in sharded_state_dict.items():
+ replica_id = sh_ten.replica_id
+ assert (
+ len(replica_id) == 3
+ ), f'Expected replica_id for {k} to be in (PP, TP, DP) format, got: {replica_id}'
+ sh_ten.replica_id = (
+ *replica_id[:2],
+ parallel_state.get_data_modulo_expert_parallel_rank(),
+ )
+ return sharded_state_dict
+
+ class TEColumnParallelGroupedLinear(TEGroupedLinear):
+ """
+ Wrapper for the Transformer-Engine's `GroupedLinear` layer but specialized
+ to column-parallel style.
+ """
+
+ def __init__(
+ self,
+ num_gemms: int,
+ input_size: int,
+ output_size: int,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ tp_comm_buffer_name: str = None,
+ ):
+
+ super().__init__(
+ num_gemms=num_gemms,
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="column",
+ config=config,
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ is_expert=is_expert,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ )
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """
+ For each gemm, sharding along axis 0, bias sharded.
+ Assume sharded_offsets[-1] is the expert parallel offset.
+ """
+ tp_axis_map = {}
+ for gemm_idx in range(self.num_gemms):
+ tp_axis_map.update({f'{gemm_idx}.weight': 0, f'{gemm_idx}.bias': 0})
+ return super()._sharded_state_dict_grouped(
+ tp_axis_map, prefix, sharded_offsets, metadata
+ )
+
+ class TERowParallelGroupedLinear(TEGroupedLinear):
+ """
+ Wrapper for the Transformer-Engine's `GroupedLinear` layer but specialized
+ to row-parallel style.
+ """
+
+ def __init__(
+ self,
+ num_gemms: int,
+ input_size: int,
+ output_size: int,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ skip_bias_add: bool,
+ is_expert: bool,
+ tp_comm_buffer_name: str = None,
+ ):
+
+ super().__init__(
+ num_gemms=num_gemms,
+ input_size=input_size,
+ output_size=output_size,
+ parallel_mode="row",
+ config=config,
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ is_expert=is_expert,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ )
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """
+ For each gemm, sharding along axis 1, bias not sharded.
+ Assume sharded_offsets[-1] is the expert parallel offset.
+ """
+ tp_axis_map = {f'{gemm_idx}.weight': 1 for gemm_idx in range(self.num_gemms)}
+ return super()._sharded_state_dict_grouped(
+ tp_axis_map, prefix, sharded_offsets, metadata
+ )
+
+else:
+
+ TEGroupedLinear = None
+ TEColumnParallelGroupedLinear = None
+ TERowParallelGroupedLinear = None
+
+
+class TEDelayedScaling(te.common.recipe.DelayedScaling):
+ """
+ Wrapper for the Transformer-Engine's `DelayedScaling` layer.
+ """
+
+ def __init__(
+ self,
+ config: ModelParallelConfig,
+ fp8_format: int,
+ override_linear_precision: tuple = (False, False, False),
+ ):
+ extra_kwargs = _get_extra_te_kwargs(config)
+ if is_te_min_version("1.6.0.dev0"):
+ extra_kwargs["fp8_dpa"] = config.fp8_dot_product_attention
+ extra_kwargs["fp8_mha"] = config.fp8_multi_head_attention
+ if get_te_version() < PkgVersion("1.8.0"):
+ extra_kwargs["interval"] = config.fp8_interval
+ elif config.fp8_interval != 1:
+ warnings.warn("fp8_interval is deprecated and ignored from Transformer-Engine v1.8.0.")
+
+ super().__init__(
+ margin=config.fp8_margin,
+ fp8_format=fp8_format,
+ amax_compute_algo=config.fp8_amax_compute_algo,
+ amax_history_len=config.fp8_amax_history_len,
+ override_linear_precision=override_linear_precision,
+ **extra_kwargs,
+ )
+
+
+class TECudaRNGStatesTracker(te.pytorch.distributed.CudaRNGStatesTracker):
+ """Wraps TransformerEngine's CudaRNGStatesTracker so that it is
+ interchangeable with Megatron's RNG tracker"""
+
+ def is_initialized(self):
+ """Checks if the internal RNG state has been set wirth set_states()."""
+ return self._is_initialized
+
+ def reset(self):
+ """Reset the internal RNG state."""
+ super().reset()
+ self._is_initialized = False
+
+ def set_states(self, states):
+ """Set the internal RNG state."""
+ super().set_states(states)
+ self._is_initialized = True
+
+ def add(self, name, seed):
+ """Track the rng state."""
+ super().add(name, seed)
+ self._is_initialized = True
+
+
+def te_checkpoint(
+ forward_func,
+ distribute_saved_activations,
+ get_rng_state_tracker,
+ tp_group,
+ hidden_states,
+ attention_mask,
+ context,
+ context_mask,
+ rotary_pos_emb,
+):
+ """Checkpointing with Transformer-Engine."""
+ from transformer_engine.pytorch.distributed import checkpoint
+
+ if is_te_min_version("1.5.0"):
+ return checkpoint(
+ forward_func,
+ hidden_states,
+ attention_mask,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ distribute_saved_activations=distribute_saved_activations,
+ get_rng_state_tracker=get_rng_state_tracker,
+ tp_group=tp_group,
+ )
+ else:
+ return checkpoint(
+ forward_func,
+ distribute_saved_activations,
+ get_rng_state_tracker,
+ tp_group,
+ hidden_states,
+ attention_mask,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ )
+
+
+try:
+
+ from transformer_engine.pytorch.attention import _SplitAlongDim
+
+ SplitAlongDim = _SplitAlongDim.apply
+
+except ImportError:
+
+ SplitAlongDim = None
+
+try:
+
+ from transformer_engine.pytorch.cpu_offload import (
+ get_cpu_offload_context as _get_cpu_offload_context,
+ )
+
+ def get_cpu_offload_context(
+ enabled, num_layers, model_layers, activation_offloading, weight_offloading
+ ):
+ """Get CPU offload context and sync function."""
+ if is_te_min_version("1.10.0.dev0"):
+ context, sync_func = _get_cpu_offload_context(
+ enabled, num_layers, model_layers, activation_offloading, weight_offloading
+ )
+ else:
+ context, sync_func = _get_cpu_offload_context(
+ enabled, num_layers, activation_offloading, weight_offloading
+ )
+
+ return context, sync_func
+
+except ImportError:
+
+ get_cpu_offload_context = None
diff --git a/tests/pipeline_parallel/__init__.py b/megatron/core/fusions/__init__.py
similarity index 100%
rename from tests/pipeline_parallel/__init__.py
rename to megatron/core/fusions/__init__.py
diff --git a/megatron/core/fusions/fused_bias_dropout.py b/megatron/core/fusions/fused_bias_dropout.py
new file mode 100644
index 0000000000..c7fa8419a0
--- /dev/null
+++ b/megatron/core/fusions/fused_bias_dropout.py
@@ -0,0 +1,73 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+from typing import Optional, Tuple
+
+import torch
+
+from megatron.core.jit import jit_fuser
+
+
+def _bias_dropout_add_func(x_with_bias, residual, prob, training):
+ # type: (Tuple[Tensor, Optional[Tensor]], Tensor, float, bool) -> Tensor
+ # NOTE: Previously, the argument `bias` used to be passed as
+ # `bias.expand_as(residual)` when the `bias_dropout_func` is called from the
+ # transformer layer but broadcasting should automatically take care of that.
+ # Also, looking at broadcasting semantics, `expand_as` and broadcasting
+ # seem to be identical performance-wise (both just change the view).
+
+ x, bias = x_with_bias # unpack
+
+ # If we want to train mixed precision, then the output of this function
+ # should be half precision. However, in AMP O1, the input (residual) is
+ # in fp32, and it will up-cast the result to fp32, causing pipeline parallel
+ # GPU communication to hang. Therefore, we need to cast residual to the same
+ # dtype as x.
+ residual = residual if residual.dtype == x.dtype else residual.to(x.dtype)
+
+ # The Dropout operation, Residual Addition and the tensor returning can be
+ # done generically outside the if statement, but that stops fusing of Bias
+ # Addition-Dropout-Residual Addition operation. So doing it together inside
+ # the conditional branch to improve performance
+ if bias is not None:
+ x = x + bias
+ out = torch.nn.functional.dropout(x, p=prob, training=training)
+ out = residual + out
+ return out
+ else:
+ out = torch.nn.functional.dropout(x, p=prob, training=training)
+ out = residual + out
+ return out
+
+
+def bias_dropout_add_unfused(training):
+ def _bias_dropout_add(x_with_bias, residual, prob):
+ return _bias_dropout_add_func(x_with_bias, residual, prob, training)
+
+ return _bias_dropout_add
+
+
+@jit_fuser
+def bias_dropout_add_fused_train(
+ x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]], residual: torch.Tensor, prob: float
+) -> torch.Tensor:
+ return _bias_dropout_add_func(x_with_bias, residual, prob, True)
+
+
+@jit_fuser
+def bias_dropout_add_fused_inference(
+ x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]], residual: torch.Tensor, prob: float
+) -> torch.Tensor:
+ return _bias_dropout_add_func(x_with_bias, residual, prob, False)
+
+
+def get_bias_dropout_add(training, fused):
+ if fused:
+ # jit scripting for a nn.module (with dropout) is not
+ # triggering the fusion kernel. For now, we use two
+ # different nn.functional routines to account for varying
+ # dropout semantics during training and inference phases.
+ if training:
+ return bias_dropout_add_fused_train
+ else:
+ return bias_dropout_add_fused_inference
+ else:
+ return bias_dropout_add_unfused(training)
diff --git a/megatron/core/fusions/fused_bias_geglu.py b/megatron/core/fusions/fused_bias_geglu.py
new file mode 100644
index 0000000000..70ef348828
--- /dev/null
+++ b/megatron/core/fusions/fused_bias_geglu.py
@@ -0,0 +1,85 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import torch
+
+from megatron.core.jit import jit_fuser
+
+###### BIAS GELU FUSION/ NO AUTOGRAD ################
+# 1/sqrt(2*pi)-> 0.3989423
+# 1/sqrt(2) -> 0.70710678
+# sqrt(2/pi) -> 0.79788456
+# this function is tanh approximation of gelu
+# actual gelu is:
+# x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
+
+
+@jit_fuser
+def geglu(y):
+ y_1, y_2 = torch.chunk(y, 2, -1)
+ return (y_1 * 0.5 * (1.0 + torch.tanh(0.79788456 * y_1 * (1 + 0.044715 * y_1 * y_1)))) * y_2
+
+
+@jit_fuser
+def bias_geglu(bias, y):
+ y = y + bias
+ return geglu(y)
+
+
+# gradient of tanh approximation of gelu
+# gradient of actual gelu is:
+# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
+@jit_fuser
+def geglu_back(g, y):
+ y_1, y_2 = torch.chunk(y, 2, -1)
+ tanh_out = torch.tanh(0.79788456 * y_1 * (1 + 0.044715 * y_1 * y_1))
+ # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
+ ff = 0.5 * y_1 * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * y_1 * y_1)) + 0.5 * (
+ 1 + tanh_out
+ )
+ return torch.cat(((g * y_2) * ff, g * (y_1 * 0.5 * (1.0 + tanh_out))), -1)
+
+
+@jit_fuser
+def bias_geglu_back(g, y, bias):
+ y = y + bias
+ return geglu_back(g, y)
+
+
+class BiasGeGLUFunction(torch.autograd.Function):
+ @staticmethod
+ # bias is an optional argument
+ def forward(ctx, input, bias):
+ ctx.save_for_backward(input, bias)
+ return bias_geglu(input, bias)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ input, bias = ctx.saved_tensors
+ tmp = bias_geglu_back(grad_output, input, bias)
+ return tmp, tmp
+
+
+class GeGLUFunction(torch.autograd.Function):
+ @staticmethod
+ # bias is an optional argument
+ def forward(ctx, input):
+ ctx.save_for_backward(input)
+ return geglu(input)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ input = ctx.saved_tensors
+ tmp = geglu_back(grad_output, input[0])
+ return tmp
+
+
+def bias_geglu_impl(input, bias):
+ ori_shape = input.shape
+ assert len(ori_shape) in [2, 3]
+ input = input.view(-1, ori_shape[-1])
+ if bias is not None:
+ output = BiasGeGLUFunction.apply(input, bias)
+ else:
+ output = GeGLUFunction.apply(input)
+
+ return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1)
diff --git a/megatron/core/fusions/fused_bias_gelu.py b/megatron/core/fusions/fused_bias_gelu.py
new file mode 100644
index 0000000000..8cc90f6174
--- /dev/null
+++ b/megatron/core/fusions/fused_bias_gelu.py
@@ -0,0 +1,55 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import torch
+
+from megatron.core.jit import jit_fuser
+
+# BIAS GELU FUSION/ NO AUTOGRAD ################
+# 1/sqrt(2*pi)-> 0.3989423
+# 1/sqrt(2) -> 0.70710678
+# sqrt(2/pi) -> 0.79788456
+# this function is tanh approximation of gelu
+# actual gelu is:
+# x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
+
+
+@jit_fuser
+def bias_gelu(bias, y):
+ x = bias + y
+ return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+
+
+# gradient of tanh approximation of gelu
+# gradient of actual gelu is:
+# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
+@jit_fuser
+def bias_gelu_back(g, bias, y):
+ x = bias + y
+ tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+ # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
+ ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (
+ 1 + tanh_out
+ )
+ return ff * g
+
+
+class GeLUFunction(torch.autograd.Function):
+ @staticmethod
+ # bias is an optional argument
+ def forward(ctx, input, bias):
+ ctx.save_for_backward(input, bias)
+ return bias_gelu(bias, input)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ input, bias = ctx.saved_tensors
+ tmp = bias_gelu_back(grad_output, bias, input)
+ return tmp, tmp
+
+ # This is required to make Sphinx happy :-(
+ @classmethod
+ def apply(cls, *args, **kwargs):
+ return super().apply(*args, **kwargs)
+
+
+bias_gelu_impl = GeLUFunction.apply
diff --git a/megatron/core/fusions/fused_bias_swiglu.py b/megatron/core/fusions/fused_bias_swiglu.py
new file mode 100644
index 0000000000..fd3ac3ec6f
--- /dev/null
+++ b/megatron/core/fusions/fused_bias_swiglu.py
@@ -0,0 +1,89 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import torch
+import torch.nn.functional as F
+
+from megatron.core.jit import jit_fuser
+
+###### BIAS SWIGLU FUSION/ NO AUTOGRAD ################
+
+
+@jit_fuser
+def swiglu(y):
+ y_1, y_2 = torch.chunk(y, 2, -1)
+ return F.silu(y_1) * y_2
+
+
+@jit_fuser
+def bias_swiglu(y, bias):
+ y = y + bias
+ return swiglu(y)
+
+
+# gradient of tanh approximation of gelu
+# gradient of actual gelu is:
+# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
+@jit_fuser
+def swiglu_back(g, y):
+ y_1, y_2 = torch.chunk(y, 2, -1)
+ return torch.cat(
+ (g * torch.sigmoid(y_1) * (1 + y_1 * (1 - torch.sigmoid(y_1))) * y_2, g * F.silu(y_1)), -1
+ )
+
+
+@jit_fuser
+def bias_swiglu_back(g, y, bias):
+ y = y + bias
+ return swiglu_back(g, y)
+
+
+class BiasSwiGLUFunction(torch.autograd.Function):
+ @staticmethod
+ # bias is an optional argument
+ def forward(ctx, input, bias, fp8_input_store):
+ input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+ ctx.save_for_backward(input_for_backward, bias)
+ ctx.ori_input_dtype = input.dtype
+ ctx.fp8_input_store = fp8_input_store
+ return bias_swiglu(input, bias)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ input, bias = ctx.saved_tensors
+ input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
+ tmp = bias_swiglu_back(grad_output, input, bias)
+ return tmp, tmp, None
+
+
+class SwiGLUFunction(torch.autograd.Function):
+ @staticmethod
+ # bias is an optional argument
+ def forward(ctx, input, fp8_input_store):
+ input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+ ctx.save_for_backward(input_for_backward)
+ ctx.ori_input_dtype = input.dtype
+ ctx.fp8_input_store = fp8_input_store
+ return swiglu(input)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ input = ctx.saved_tensors[0]
+ input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input
+ tmp = swiglu_back(grad_output, input)
+ return tmp, None
+
+
+def bias_swiglu_impl(input, bias, fp8_input_store=False):
+ ori_shape = input.shape
+ assert len(ori_shape) in [2, 3]
+ input = input.view(-1, ori_shape[-1])
+ if bias is not None:
+ output = BiasSwiGLUFunction.apply(input, bias, fp8_input_store)
+ else:
+ output = SwiGLUFunction.apply(input, fp8_input_store)
+
+ return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1)
+
+
+# bias_swiglu_impl = BiasSwiGLUFunction.apply
+# swiglu_impl = SwiGLUFunction.apply
diff --git a/megatron/core/fusions/fused_cross_entropy.py b/megatron/core/fusions/fused_cross_entropy.py
new file mode 100644
index 0000000000..909cc403cf
--- /dev/null
+++ b/megatron/core/fusions/fused_cross_entropy.py
@@ -0,0 +1,143 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Tuple
+
+import torch
+
+from megatron.core.jit import jit_fuser
+from megatron.core.parallel_state import (
+ get_tensor_model_parallel_group,
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+)
+from megatron.core.tensor_parallel.cross_entropy import VocabParallelCrossEntropy
+from megatron.core.tensor_parallel.utils import VocabUtility
+
+
+@jit_fuser
+def calculate_logits_max(vocab_parallel_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+
+ vocab_parallel_logits, logits_max = VocabParallelCrossEntropy.calculate_logits_max(
+ vocab_parallel_logits
+ )
+
+ return vocab_parallel_logits, logits_max
+
+
+@jit_fuser
+def calculate_predicted_logits(
+ vocab_parallel_logits: torch.Tensor,
+ target: torch.Tensor,
+ logits_max: torch.Tensor,
+ vocab_start_index: int,
+ vocab_end_index: int,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+
+ (target_mask, masked_target_1d, predicted_logits, sum_exp_logits, exp_logits) = (
+ VocabParallelCrossEntropy.calculate_predicted_logits(
+ vocab_parallel_logits, target, logits_max, vocab_start_index, vocab_end_index
+ )
+ )
+
+ predicted_logits_sum_exp_logits = torch.cat((predicted_logits, sum_exp_logits))
+
+ return target_mask, masked_target_1d, predicted_logits_sum_exp_logits, exp_logits
+
+
+@jit_fuser
+def calculate_cross_entropy_loss(
+ exp_logits: torch.Tensor, predicted_logits_sum_exp_logits: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+ split_val = predicted_logits_sum_exp_logits.size()[0] // 2
+ predicted_logits, sum_exp_logits = torch.split(predicted_logits_sum_exp_logits, split_val)
+
+ exp_logits, loss = VocabParallelCrossEntropy.calculate_cross_entropy_loss(
+ exp_logits, predicted_logits, sum_exp_logits
+ )
+
+ return exp_logits, loss
+
+
+@jit_fuser
+def calculate_gradients(
+ softmax: torch.Tensor,
+ grad_output: torch.Tensor,
+ target_mask: torch.Tensor,
+ masked_target_1d: torch.Tensor,
+) -> torch.Tensor:
+
+ (grad_2d, arange_1d, softmax_update, grad_input) = (
+ VocabParallelCrossEntropy.prepare_gradient_calculation_operands(softmax, target_mask)
+ )
+
+ grad_input = VocabParallelCrossEntropy.calculate_gradients(
+ grad_2d, arange_1d, masked_target_1d, softmax_update, grad_input, grad_output
+ )
+
+ grad_input = grad_input.to(torch.bfloat16)
+
+ return grad_input
+
+
+class _VocabParallelCrossEntropy(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, vocab_parallel_logits, target):
+
+ vocab_parallel_logits, logits_max = calculate_logits_max(vocab_parallel_logits)
+ torch.distributed.all_reduce(
+ logits_max, op=torch.distributed.ReduceOp.MAX, group=get_tensor_model_parallel_group()
+ )
+
+ # Get the partition's vocab indices
+ get_vocab_range = VocabUtility.vocab_range_from_per_partition_vocab_size
+ partition_vocab_size = vocab_parallel_logits.size()[-1]
+ rank = get_tensor_model_parallel_rank()
+ world_size = get_tensor_model_parallel_world_size()
+ vocab_start_index, vocab_end_index = get_vocab_range(partition_vocab_size, rank, world_size)
+
+ (target_mask, masked_target_1d, predicted_logits_sum_exp_logits, exp_logits) = (
+ calculate_predicted_logits(
+ vocab_parallel_logits, target, logits_max, vocab_start_index, vocab_end_index
+ )
+ )
+
+ # All reduce is needed to get the chunks from other GPUs.
+ # In the fused case, tensors are batches to invoke a single
+ # AllReduce call
+ torch.distributed.all_reduce(
+ predicted_logits_sum_exp_logits,
+ op=torch.distributed.ReduceOp.SUM,
+ group=get_tensor_model_parallel_group(),
+ )
+
+ exp_logits, loss = calculate_cross_entropy_loss(exp_logits, predicted_logits_sum_exp_logits)
+
+ # Store softmax, target-mask and masked-target for backward pass.
+ ctx.save_for_backward(exp_logits, target_mask, masked_target_1d)
+
+ return loss
+
+ @staticmethod
+ def backward(ctx, grad_output):
+
+ # Retreive tensors from the forward path.
+ softmax, target_mask, masked_target_1d = ctx.saved_tensors
+
+ grad_input = calculate_gradients(softmax, grad_output, target_mask, masked_target_1d)
+
+ return grad_input, None
+
+
+def fused_vocab_parallel_cross_entropy(vocab_parallel_logits, target):
+ """
+ Performs cross entropy loss when logits are split across tensor parallel ranks
+
+ Args:
+ vocab_parallel_logits: logits split across tensor parallel ranks
+ dimension is [sequence_length, batch_size, hidden_size]
+
+ target: correct vocab ids of dimseion [sequence_length, micro_batch_size]
+
+ """
+ return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target)
diff --git a/megatron/core/fusions/fused_layer_norm.py b/megatron/core/fusions/fused_layer_norm.py
new file mode 100644
index 0000000000..d02ae7aa4d
--- /dev/null
+++ b/megatron/core/fusions/fused_layer_norm.py
@@ -0,0 +1,169 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import importlib
+import inspect
+import numbers
+
+import torch
+from torch import Tensor
+from torch.nn import init
+from torch.nn.parameter import Parameter
+
+from megatron.core.transformer import TransformerConfig
+from megatron.core.utils import make_viewless_tensor
+
+try:
+ from apex.contrib.layer_norm.layer_norm import FastLayerNormFN
+
+ HAVE_PERSIST_LAYER_NORM = True
+except ImportError:
+ HAVE_PERSIST_LAYER_NORM = False
+
+try:
+ from apex.normalization.fused_layer_norm import FusedLayerNormAffineFunction
+
+ HAVE_FUSED_LAYER_NORM = True
+except ImportError:
+ HAVE_FUSED_LAYER_NORM = False
+
+
+class FusedLayerNorm(torch.nn.Module):
+ """Layer Norm, fused into a single CUDA kernel.
+
+ Args:
+ hidden_size (int): Transformer hidden dimension.
+
+ eps (float): Epsilon added to denominator, for numerical stability.
+
+ persist_layer_norm (bool): Use persistent fused layer norm kernel.
+ This kernel supports only a set of hidden sizes. Please
+ check persist_ln_hidden_sizes if your hidden size is supported.
+
+ zero_centered_gamma (bool): Adjust LayerNorm weights such that they are
+ centered around zero. This improves numerical stability.
+
+ config (TransformerConfig): Transformer config. Include to match custom
+ layer norm interfaces.
+
+ normalization (str): Normalization type, used for Transformer Engine.
+ Must equal 'LayerNorm' here.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ hidden_size: int,
+ eps: float = 1e-5,
+ persist_layer_norm: bool = True,
+ zero_centered_gamma: bool = False,
+ normalization: str = "LayerNorm", # included to match TE interface
+ ):
+ super().__init__()
+
+ self.config = config
+
+ self.zero_centered_gamma = self.config.layernorm_zero_centered_gamma
+ assert (
+ self.config.normalization == "LayerNorm"
+ ), f'({self.config.normalization}) is not supported in FusedLayerNorm'
+
+ # List of hiddens sizes supported in the persistent layer norm kernel
+ # If the hidden size is not supported, fall back to the non-persistent
+ # kernel.
+ persist_ln_hidden_sizes = [
+ 1024,
+ 1536,
+ 2048,
+ 2304,
+ 3072,
+ 3840,
+ 4096,
+ 5120,
+ 6144,
+ 8192,
+ 10240,
+ 12288,
+ 12800,
+ 15360,
+ 16384,
+ 18432,
+ 20480,
+ 24576,
+ 25600,
+ 30720,
+ 32768,
+ 40960,
+ 49152,
+ 65536,
+ ]
+ persist_layer_norm = self.config.persist_layer_norm
+ if hidden_size not in persist_ln_hidden_sizes or not HAVE_PERSIST_LAYER_NORM:
+ persist_layer_norm = False
+
+ if not persist_layer_norm and not HAVE_FUSED_LAYER_NORM:
+ # TODO: Add pytorch only layer norm
+ raise ValueError(f'Apex must be installed to use FusedLayerNorm.')
+
+ if isinstance(hidden_size, numbers.Integral):
+ hidden_size = (hidden_size,)
+ self.hidden_size = torch.Size(hidden_size)
+ self.eps = eps
+ # Parameters need to be initialized with torch.empty rather than torch.Tensor for correct device placement with nemo2.
+ self.weight = Parameter(torch.empty(*hidden_size))
+ self.bias = Parameter(torch.empty(*hidden_size))
+ self.reset_parameters()
+ self.persist_layer_norm = persist_layer_norm
+ self.sequence_parallel = self.config.sequence_parallel
+
+ # set sequence parallelism flag on weight and bias parameters
+ setattr(self.weight, 'sequence_parallel', self.sequence_parallel)
+ setattr(self.bias, 'sequence_parallel', self.sequence_parallel)
+
+ def reset_parameters(self):
+
+ if self.zero_centered_gamma:
+ init.zeros_(self.weight)
+ init.zeros_(self.bias)
+ else:
+ init.ones_(self.weight)
+ init.zeros_(self.bias)
+
+ def forward(self, input: Tensor) -> Tensor:
+
+ weight = self.weight + 1 if self.zero_centered_gamma else self.weight
+
+ if self.persist_layer_norm:
+ if 'memory_efficient' in inspect.getfullargspec(FastLayerNormFN.forward).args:
+ output = FastLayerNormFN.apply(
+ input, weight, self.bias, self.eps, self.config.memory_efficient_layer_norm
+ )
+ else:
+ output = FastLayerNormFN.apply(input, weight, self.bias, self.eps)
+
+ # Apex's fast layer norm function outputs a 'view' tensor (i.e., has
+ # a populated '_base' field). This will result in schedule.py's
+ # deallocate_output_tensor() throwing an error, so a viewless tensor is
+ # created to prevent this.
+ output = make_viewless_tensor(
+ inp=output, requires_grad=input.requires_grad, keep_graph=True
+ )
+
+ else:
+ if (
+ 'memory_efficient'
+ in inspect.getfullargspec(FusedLayerNormAffineFunction.forward).args
+ ):
+ return FusedLayerNormAffineFunction.apply(
+ input,
+ weight,
+ self.bias,
+ self.hidden_size,
+ self.eps,
+ self.config.memory_efficient_layer_norm,
+ )
+ else:
+ return FusedLayerNormAffineFunction.apply(
+ input, weight, self.bias, self.hidden_size, self.eps
+ )
+
+ return output
diff --git a/megatron/model/fused_softmax.py b/megatron/core/fusions/fused_softmax.py
similarity index 81%
rename from megatron/model/fused_softmax.py
rename to megatron/core/fusions/fused_softmax.py
index ed29262acd..c7bfbb768b 100644
--- a/megatron/model/fused_softmax.py
+++ b/megatron/core/fusions/fused_softmax.py
@@ -1,9 +1,12 @@
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+from typing import Optional
import torch
import torch.nn as nn
-from megatron.model.enums import AttnMaskType
+
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.utils import get_default_causal_mask
class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
@@ -19,9 +22,7 @@ def forward(ctx, inputs, scale):
import scaled_upper_triang_masked_softmax_cuda
scale_t = torch.tensor([scale])
- softmax_results = scaled_upper_triang_masked_softmax_cuda.forward(
- inputs, scale_t[0]
- )
+ softmax_results = scaled_upper_triang_masked_softmax_cuda.forward(inputs, scale_t[0])
ctx.save_for_backward(softmax_results, scale_t)
return softmax_results
@@ -62,9 +63,7 @@ def backward(ctx, output_grads):
softmax_results, scale_t = ctx.saved_tensors
- input_grads = scaled_masked_softmax_cuda.backward(
- output_grads, softmax_results, scale_t[0]
- )
+ input_grads = scaled_masked_softmax_cuda.backward(output_grads, softmax_results, scale_t[0])
return input_grads, None, None
@@ -81,9 +80,7 @@ def forward(ctx, inputs, scale):
scale_t = torch.tensor([scale])
- softmax_results = scaled_softmax_cuda.forward(
- inputs, scale_t[0]
- )
+ softmax_results = scaled_softmax_cuda.forward(inputs, scale_t[0])
ctx.save_for_backward(softmax_results, scale_t)
return softmax_results
@@ -93,9 +90,7 @@ def backward(ctx, output_grads):
softmax_results, scale_t = ctx.saved_tensors
- input_grads = scaled_softmax_cuda.backward(
- output_grads, softmax_results, scale_t[0]
- )
+ input_grads = scaled_softmax_cuda.backward(output_grads, softmax_results, scale_t[0])
return input_grads, None, None
@@ -103,7 +98,7 @@ class FusedScaleMaskSoftmax(nn.Module):
"""
fused operation: scaling + mask + softmax
- Arguments:
+ Args:
input_in_fp16: flag to indicate if input in fp16 data format.
input_in_bf16: flag to indicate if input in bf16 data format.
attn_mask_type: attention mask type (pad or causal)
@@ -136,11 +131,14 @@ def __init__(
self.softmax_in_fp32 = softmax_in_fp32
self.scale = scale
- assert (
- self.scale is None or softmax_in_fp32
- ), "softmax should be in fp32 when scaled"
+ assert self.scale is None or softmax_in_fp32, "softmax should be in fp32 when scaled"
+
+ def forward(self, input: torch.Tensor, mask: Optional[torch.Tensor]):
+ """Forward pass of softmax with masked input.
- def forward(self, input, mask):
+ In case attn_mask_type is causal the mask is generated and None can be passed.
+ A user-defined mask is only needed when attn_mask_type is not causal.
+ """
# [b, np, sq, sk]
assert input.dim() == 4
@@ -157,7 +155,7 @@ def is_kernel_available(self, mask, b, np, sq, sk):
and self.input_in_float16 # input must be fp16
and 16 < sk <= 4096 # sk must be 16 ~ 2048
and sq % 4 == 0 # sq must be divisor of 4
- and sk % 4 == 0 # sk must be divisor of 4
+ and sk % 4 == 0 # sk must be divisor of 4
and attn_batches % 4 == 0 # np * b must be divisor of 4
):
if 0 <= sk <= 4096:
@@ -195,6 +193,15 @@ def forward_torch_softmax(self, input, mask):
if self.scale is not None:
input = input * self.scale
+
+ # Generate causal mask if not given
+ sq, sk = input.size(2), input.size(3)
+ if self.attn_mask_type == AttnMaskType.causal and mask is None and sq > 1:
+ # If sq == 1 then either KV cache is used or one-element context is passed
+ # so keeping mask=None in this case; subsequent code should handle it
+ assert sq == sk, "causal mask is only for self attention"
+ mask = get_default_causal_mask(sq)
+
mask_output = self.mask_func(input, mask) if mask is not None else input
probs = torch.nn.Softmax(dim=-1)(mask_output)
diff --git a/megatron/core/inference/__init__.py b/megatron/core/inference/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/inference/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/inference/ammo_support/__init__.py b/megatron/core/inference/ammo_support/__init__.py
new file mode 100644
index 0000000000..12be50cefe
--- /dev/null
+++ b/megatron/core/inference/ammo_support/__init__.py
@@ -0,0 +1,8 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import warnings
+
+warnings.warn(
+ "The 'megatron.core.inference.ammo_support' module is deprecated and will be removed in a future release. "
+ "Please use megatron.core.inference.modelopt_support instead",
+ DeprecationWarning,
+)
diff --git a/megatron/core/inference/ammo_support/gpt/model_specs.py b/megatron/core/inference/ammo_support/gpt/model_specs.py
new file mode 100644
index 0000000000..ba3bd9fa0f
--- /dev/null
+++ b/megatron/core/inference/ammo_support/gpt/model_specs.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from megatron.core.inference.modelopt_support.gpt.model_specs import get_gpt_layer_modelopt_spec
diff --git a/megatron/core/inference/ammo_support/gpt/state_dict_hooks.py b/megatron/core/inference/ammo_support/gpt/state_dict_hooks.py
new file mode 100644
index 0000000000..8532366222
--- /dev/null
+++ b/megatron/core/inference/ammo_support/gpt/state_dict_hooks.py
@@ -0,0 +1,5 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from megatron.core.inference.modelopt_support.gpt.state_dict_hooks import (
+ mcore_gpt_load_legacy_state_dict_pre_hook,
+ mcore_gpt_load_te_state_dict_pre_hook,
+)
diff --git a/megatron/core/inference/common_inference_params.py b/megatron/core/inference/common_inference_params.py
new file mode 100644
index 0000000000..22353088f8
--- /dev/null
+++ b/megatron/core/inference/common_inference_params.py
@@ -0,0 +1,29 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from dataclasses import dataclass
+
+
+@dataclass
+class CommonInferenceParams:
+ """Inference parameters sent along with the prompts
+
+ For an explanation of these parameters refer to this blog https://ivibudh.medium.com/a-guide-to-controlling-llm-model-output-exploring-top-k-top-p-and-temperature-parameters-ed6a31313910
+ """
+
+ temperature: float = 1.0
+ top_k: int = 0
+ top_p: float = 0.0
+ return_log_probs: bool = False
+ num_tokens_to_generate: int = 30
+
+ def add_attributes(self, attribute_value_pair: dict):
+ """Utility to add more attributes to inference params
+
+ Use this method to pass in a custom dictonary to add more inference parameter attributes to the instance you created. Use as follows
+ c = CommonInferenceParams
+ c.add_attributes({'min_length':4, 'eod_id':153})
+
+ Args:
+ attribute_value_pair (dict): A dictionary containing attributes as the key names and their values as the values.
+ """
+ for key, value in attribute_value_pair.items():
+ setattr(self, key, value)
diff --git a/megatron/core/inference/communication_utils.py b/megatron/core/inference/communication_utils.py
new file mode 100644
index 0000000000..0c23a583de
--- /dev/null
+++ b/megatron/core/inference/communication_utils.py
@@ -0,0 +1,50 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import torch
+
+from megatron.core import parallel_state
+
+
+def _is_cuda(tensor):
+ """Check if a tensor is not none and is cuda."""
+ assert tensor is not None
+ assert tensor.is_cuda
+
+
+def broadcast_from_last_pipeline_stage(size, dtype, tensor=None):
+ """Broadcast a tensor from last pipeline stage to all ranks."""
+
+ if parallel_state.is_pipeline_last_stage():
+ _is_cuda(tensor)
+ assert tensor.is_contiguous()
+ else:
+ tensor = torch.empty(size, dtype=dtype, device=torch.cuda.current_device())
+ # Get the group and corresponding source rank.
+ src = parallel_state.get_pipeline_model_parallel_last_rank()
+ group = parallel_state.get_pipeline_model_parallel_group()
+ torch.distributed.broadcast(tensor, src, group)
+ return tensor
+
+
+def recv_from_prev_pipeline_rank_(recv_buffer=None):
+ """Receive from previous pipeline stage and update the
+ input buffer inplace."""
+ recv_prev_op = torch.distributed.P2POp(
+ torch.distributed.irecv, recv_buffer, parallel_state.get_pipeline_model_parallel_prev_rank()
+ )
+ reqs = torch.distributed.batch_isend_irecv([recv_prev_op])
+ for req in reqs:
+ req.wait()
+ # To protect against race condition when using batch_isend_irecv().
+ torch.cuda.synchronize()
+
+
+def send_to_next_pipeline_rank(tensor=None):
+ """Send output to the next pipeline stage."""
+ send_next_op = torch.distributed.P2POp(
+ torch.distributed.isend, tensor, parallel_state.get_pipeline_model_parallel_next_rank()
+ )
+ reqs = torch.distributed.batch_isend_irecv([send_next_op])
+ for req in reqs:
+ req.wait()
+ # To protect against race condition when using batch_isend_irecv().
+ torch.cuda.synchronize()
diff --git a/megatron/core/inference/engines/__init__.py b/megatron/core/inference/engines/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/inference/engines/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/inference/engines/abstract_engine.py b/megatron/core/inference/engines/abstract_engine.py
new file mode 100644
index 0000000000..6893f6a905
--- /dev/null
+++ b/megatron/core/inference/engines/abstract_engine.py
@@ -0,0 +1,17 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from abc import ABC, abstractmethod
+from typing import List
+
+
+class AbstractEngine(ABC):
+ @staticmethod
+ @abstractmethod
+ def generate(self) -> dict:
+ """The abstract backend's generate function.
+
+ To define a new backend, implement this and return the outputs as a dictionary.
+
+ Returns:
+ dict: The output dictionary containing keys for `input_prompt`, `generated_text`, `generated_tokens`.
+ """
+ pass
diff --git a/megatron/core/inference/engines/mcore_engine.py b/megatron/core/inference/engines/mcore_engine.py
new file mode 100644
index 0000000000..fe8160228b
--- /dev/null
+++ b/megatron/core/inference/engines/mcore_engine.py
@@ -0,0 +1,113 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from typing import Dict, List
+
+import torch
+
+from megatron.core.inference.common_inference_params import CommonInferenceParams
+from megatron.core.inference.engines.abstract_engine import AbstractEngine
+from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.inference.scheduler import Scheduler
+from megatron.core.inference.text_generation_controllers.simple_text_generation_controller import (
+ SimpleTextGenerationController,
+)
+
+
+class MCoreEngine(AbstractEngine):
+ """The Megatron core backend constructor
+
+ This is the backend that does a simple forward pass on the model.
+ Supports any model that is callable (Accepts the inputs and outputs the tensor)
+
+ Args:
+ text_generation_controller (SimpleTextGenerationController): A text generation
+ controller that will be used to define how to preprocess prompts, generate
+ outputs and detokenizer the output tokens.
+ max_batch_size : The maxinum number of requests to process at once
+ random_seed (int, optional): Use a random seed if you want deterministic
+ results. Defaults to None.
+ """
+
+ def __init__(
+ self,
+ text_generation_controller: SimpleTextGenerationController,
+ max_batch_size,
+ random_seed: int = None,
+ ):
+ self.text_generation_controller = text_generation_controller
+ self.random_seed = random_seed
+ self.scheduler = Scheduler(max_batch_size=max_batch_size)
+
+ def generate(
+ self,
+ prompts: List[str],
+ add_BOS: bool = False,
+ encoder_prompts: List[str] = None,
+ common_inference_params: CommonInferenceParams = None,
+ ) -> dict:
+ """The megatron core inference backend generate function
+
+ This backend returns the output generations as a dictionary.
+ It returns the prompt tokens along with the generated tokens, the prompt
+ plus the generated string and the output log probabilities if requested
+
+ Args:
+ prompts (List[str]): All the prompts as a list of strings
+ add_BOS (bool): Whether to add BOS token to beginning of prompts
+ encoder_prompts (List[dict]): All the encoder prompts as a list of strings
+ common_inference_params (CommonInferenceParams): The inference parameters
+
+ Returns:
+ List[InferenceRequest]: The output is list of inference requests containing the
+ generated tokens, texts and log probs if required
+ """
+ # TODO :M core- get rng state tracker
+ if self.random_seed:
+ torch.random.manual_seed(self.random_seed)
+
+ for i in range(len(prompts)):
+ prompt = prompts[i]
+ encoder_prompt = encoder_prompts[i] if encoder_prompts is not None else None
+ prompt_tokens = self.text_generation_controller.tokenize_prompt(prompt, add_BOS)
+
+ self.scheduler.add_request(
+ prompt=prompt,
+ prompt_tokens=prompt_tokens,
+ encoder_prompt=encoder_prompt,
+ inference_parameters=common_inference_params,
+ )
+
+ self.run_engine()
+
+ result: List[InferenceRequest] = self.scheduler.completed_request_pool.values()
+ return result
+
+ def run_engine(self):
+ """Main functionality to run inference
+
+ Runs the engine until there are no requests in the queue.
+
+ Args:
+ dynamic_generation (bool, optional): Set this to True, if you want
+ to enable dynamic batching. Mainly used with an inference server.
+ Defaults to False.
+ """
+ while self.scheduler.have_requests_pending():
+ active_requests: Dict[int, InferenceRequest] = self.scheduler.active_request_pool.copy()
+ result_dict: Dict[int, InferenceRequest] = (
+ self.text_generation_controller.generate_all_output_tokens_static_batch(
+ active_requests
+ )
+ )
+
+ self.scheduler.update_requests_pools(result_dict=result_dict)
+
+ # TODO: Later for dynamic batching we will do something like this
+ """
+ if dynamic_batching:
+ result_dict: Dict[
+ int, InferenceRequest
+ ] = self.text_generation_controller.generate_output_tokens_one_step_dynamic_batch(
+ active_requests
+ )
+ self.scheduler.update_requests_pools(result_dict=result_dict)
+ """
diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py
new file mode 100644
index 0000000000..4825dfd366
--- /dev/null
+++ b/megatron/core/inference/inference_request.py
@@ -0,0 +1,39 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from dataclasses import dataclass
+from enum import Enum
+from typing import List
+
+import torch
+
+from megatron.core.inference.common_inference_params import CommonInferenceParams
+
+
+# class syntax
+class Status(Enum):
+ """Enum for status"""
+
+ WAITING_IN_QUEUE = 1
+ ACTIVE_AND_GENERATING_TOKENS = 2
+ ACTIVE_BUT_NOT_GENERATING_TOKENS = 3
+ COMPLETED = 4
+
+
+@dataclass
+class InferenceRequest:
+ """Class for one inference request
+
+ Containing relevant data for an inference request
+
+ """
+
+ request_id: str
+ prompt: str
+ inference_parameters: CommonInferenceParams
+ prompt_tokens: List[int]
+ arrival_time: float
+ status: Status
+ encoder_prompt: str = None
+ generated_text: str = None
+ generated_tokens: torch.Tensor = None
+ generated_log_probs: torch.Tensor = None
+ generated_length: int = 0
diff --git a/megatron/core/inference/model_inference_wrappers/__init__.py b/megatron/core/inference/model_inference_wrappers/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/inference/model_inference_wrappers/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
new file mode 100644
index 0000000000..b7f58efcfe
--- /dev/null
+++ b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
@@ -0,0 +1,234 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import abc
+import math
+from argparse import Namespace
+from typing import Iterable, List, Union
+
+import torch
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.inference.common_inference_params import CommonInferenceParams
+from megatron.core.inference.communication_utils import (
+ recv_from_prev_pipeline_rank_,
+ send_to_next_pipeline_rank,
+)
+from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
+ InferenceWrapperConfig,
+)
+from megatron.core.inference_params import InferenceParams
+from megatron.core.models.gpt.gpt_model import GPTModel
+
+
+class AbstractModelInferenceWrapper(abc.ABC):
+ def __init__(
+ self,
+ model: Union['LegacyGPTModel', GPTModel],
+ inference_wrapper_config: InferenceWrapperConfig,
+ ):
+ """Constructor for the model inference wrapper
+
+ The wrapper prepares the model for inference, provides the required input data and runs the forward pass.
+
+ Args:
+ model (Union[GPTModel, LegacyGPTModel]): The actual GPT model (MCore or MLM)
+ args (Namespace): The commadline arguments that were passed
+ """
+ assert not isinstance(
+ model, Iterable
+ ), 'interleaving schedule is not supported for inference'
+ self.model = model
+ self.inference_wrapper_config = inference_wrapper_config
+ self.pipeline_communication_dtype = (
+ torch.float
+ if self.inference_wrapper_config.fp32_residual_connection
+ else self.inference_wrapper_config.params_dtype
+ )
+
+ def prep_model_for_inference(self, prompts_tokens: torch.Tensor):
+ """A utility function for preparing model for inference
+
+ The function gets called once before the auto regressive inference loop. It puts the model in eval mode , and gets some model and inference data parameters. Extend this to build position ids ,attention mask etc, so that required slices can be extracted during the forward pass.
+
+ Args:
+ prompts_tokens (torch.Tensor): A tensor of shape [batch_size, max_seq_len]
+
+ """
+ self.model.eval()
+
+ # For TP only model both is_pp_first_stage and _is_pp_last_stage returns True
+ self.model_is_pipeline_parallel = not (
+ parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage()
+ )
+ self.prompts_tokens = prompts_tokens
+ batch_size, max_sequence_length = self.prompts_tokens.shape
+ self.inference_params = InferenceParams(batch_size, max_sequence_length)
+
+ @abc.abstractmethod
+ def get_batch_for_context_window(self) -> List:
+ """Returns the input data for inference
+
+ This function gets called iteratively in the inference loop . It can be used to extract relevant input from the prompt tokens, attention mask etc. required for each step in inference.
+
+ """
+ pass
+
+ def forward_pass_without_pipeline_parallel(self, inference_input: List) -> torch.Tensor:
+ """Utility to carry out simple forward pass for TP or no model parallel models
+
+ Runs a very simple forward pass for model. Used in the case of models without any parallelism or only tensor parallelism.
+
+ Args:
+ inference_input (List): A list containg the inputs for the gpt model [tokens, position ids, attention mask]
+
+ Returns:
+ torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size]
+ """
+ tokens, position_ids, attention_mask = inference_input
+ logits = self.model(
+ tokens, position_ids, attention_mask, inference_params=self.inference_params
+ )
+ logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits)
+ self.inference_params.sequence_len_offset += tokens.size(1)
+
+ return logits
+
+ def _allocate_recv_buffer(self, batch_size, seq_len):
+ """Receive happens between the layers with size [seq_len, batch_size, hidden_size]."""
+ recv_size = (seq_len, batch_size, self.inference_wrapper_config.hidden_size)
+ return torch.empty(
+ recv_size, dtype=self.pipeline_communication_dtype, device=torch.cuda.current_device()
+ )
+
+ def forward_pass_with_pipeline_parallel_small_input_batch(
+ self, inference_input: List
+ ) -> torch.Tensor:
+ """Utility to carry out forward pass for PP models with very small inputs
+
+ If a model is pipeline parallel, yet, the input global batch is very small, we compute a foward pass on the entire global batch, rather than splitting it up into micro batches and doing something more complex as in the forward_pass_with_pipeline_parallel_large_input_batch method
+
+ Args:
+ inference_input (List): A list containg the inputs for the gpt model [tokens, position ids, attention mask]
+
+ Returns:
+ torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size]
+ """
+ tokens, position_ids, attention_mask = inference_input
+ batch_size, seq_len = tokens.shape
+ recv_buffer = None
+ if not parallel_state.is_pipeline_first_stage():
+ recv_buffer = self._allocate_recv_buffer(batch_size, seq_len)
+ recv_from_prev_pipeline_rank_(recv_buffer)
+
+ self.model.set_input_tensor(recv_buffer)
+ output_tensor = self.model(
+ tokens, position_ids, attention_mask, inference_params=self.inference_params
+ )
+
+ if not parallel_state.is_pipeline_last_stage():
+ send_to_next_pipeline_rank(output_tensor.type(dtype=self.pipeline_communication_dtype))
+
+ self.inference_params.sequence_len_offset += seq_len
+
+ logits = None
+ if parallel_state.is_pipeline_last_stage():
+ logits = output_tensor
+ logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits)
+
+ return logits
+
+ def forward_pass_with_pipeline_parallel_large_input_batch(
+ self, inference_input: List
+ ) -> torch.Tensor:
+ """Utility to carry out forward pass PP models.
+
+ Runs the forward pass for models which are pipeline parallel. This is more complex than forward_pass_with_pipeline_parallel_small_input_batch coz this splits the global batch into small micro batches and runs them through the model.
+
+ Args:
+ inference_input (List): A list containg the inputs for the gpt model [tokens, position ids, attention mask]
+
+ Returns:
+ torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size]
+ """
+ tokens, position_ids, attention_mask = inference_input
+ micro_batch_size = max(
+ 1,
+ self.inference_wrapper_config.inference_batch_times_seqlen_threshold // tokens.size(1),
+ )
+ batch_size, seq_len = tokens.shape
+ # Round up to account for the last partial micro batch if present
+ num_micro_batches = math.ceil(batch_size / micro_batch_size)
+
+ logits = None
+ # Preallocate memory for output logits.
+ if parallel_state.is_pipeline_last_stage():
+ logits = torch.empty(
+ (batch_size, seq_len, self.inference_wrapper_config.padded_vocab_size),
+ dtype=torch.float32,
+ device=torch.cuda.current_device(),
+ )
+
+ recv_buffer = None
+ if not parallel_state.is_pipeline_first_stage():
+ recv_buffer = self._allocate_recv_buffer(micro_batch_size, seq_len)
+ for micro_batch_index in range(num_micro_batches):
+ start = micro_batch_index * micro_batch_size
+ end = min(start + micro_batch_size, batch_size)
+ tokens2use = tokens[start:end, ...]
+ position_ids2use = position_ids[start:end, ...]
+ current_micro_batch_size = end - start
+
+ # Need to change recv buffer shape for the last partial microbatch (if exists)
+ if current_micro_batch_size != micro_batch_size:
+ recv_buffer = self._allocate_recv_buffer(current_micro_batch_size, seq_len)
+
+ if not parallel_state.is_pipeline_first_stage():
+ recv_from_prev_pipeline_rank_(recv_buffer)
+
+ self.model.set_input_tensor(recv_buffer)
+ output_tensor = self.model(
+ tokens2use, position_ids2use, attention_mask, inference_params=self.inference_params
+ )
+
+ if not parallel_state.is_pipeline_last_stage():
+ send_to_next_pipeline_rank(output_tensor)
+
+ self.inference_params.batch_size_offset += current_micro_batch_size
+
+ if parallel_state.is_pipeline_last_stage():
+ output_tensor = tensor_parallel.gather_from_tensor_model_parallel_region(
+ output_tensor
+ )
+ logits[start:end, ...] = output_tensor
+
+ # Once done with all micro batches, we reset batch size offset and seq len offset
+ self.inference_params.sequence_len_offset += seq_len
+ self.inference_params.batch_size_offset = 0
+
+ # NOTE: Only returns the logits on the last pipeline stage
+ return logits
+
+ def run_one_forward_step(self, inference_input: List) -> torch.Tensor:
+ """The forward pass of the model for inference
+
+ Appropriate utility is called for the forward pass depending on the type of model parallelism used
+
+ Args:
+ inference_input (List): A list containg the inputs for the gpt model [tokens, position ids, attention mask]
+
+ Returns:
+ torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size]. The logits are returned only in the last pipeline stage for PP models.
+ """
+ if self.model_is_pipeline_parallel:
+ tokens = inference_input[0]
+ current_batch_size, seq_len = tokens.shape
+ # If input batch is large, we need to split into micro batches and run the forward pass
+ if (
+ current_batch_size * seq_len
+ > self.inference_wrapper_config.inference_batch_times_seqlen_threshold
+ ):
+ return self.forward_pass_with_pipeline_parallel_large_input_batch(inference_input)
+ else:
+ # If input batch is very small we can do a simple forward pass on the entire global batch
+ return self.forward_pass_with_pipeline_parallel_small_input_batch(inference_input)
+ else:
+ return self.forward_pass_without_pipeline_parallel(inference_input)
diff --git a/megatron/core/inference/model_inference_wrappers/gpt/__init__.py b/megatron/core/inference/model_inference_wrappers/gpt/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/inference/model_inference_wrappers/gpt/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py
new file mode 100644
index 0000000000..87b1d2df77
--- /dev/null
+++ b/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py
@@ -0,0 +1,85 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from argparse import Namespace
+from typing import List, Tuple
+
+import torch
+
+from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
+ AbstractModelInferenceWrapper,
+)
+from megatron.core.models.gpt import GPTModel
+
+
+class GPTInferenceWrapper(AbstractModelInferenceWrapper):
+ def __init__(self, model: GPTModel, args: Namespace):
+ """Constructor for the model inference wrapper
+
+ The wrapper prepares the model for inference, provides the required input data, and runs the forward pass
+
+ Args:
+ model (GPTModel): The GPT model (MCore or legacy)
+ args (Namespace): The command line arguments that were passed
+ """
+ super().__init__(model, args)
+
+ def prep_model_for_inference(self, prompts_tokens: torch.Tensor):
+ """A utility function for preparing model for inference
+
+ This function is called before the forward pass. It puts the model in eval mode, builds position ids, and creates attention masks so that required slices can be extracted during the forward pass.
+
+ Args:
+ prompts_tokens (torch.Tensor): A tensor of shape [batch_size, max_seq_len]
+ """
+
+ super().prep_model_for_inference(prompts_tokens=prompts_tokens)
+ self.attention_mask, self.position_ids = self._build_attention_mask_and_position_ids(
+ prompts_tokens
+ )
+
+ def _build_attention_mask_and_position_ids(
+ self, prompts_tokens: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Builds the full attention mask and position ids for the input tokens
+
+ Args:
+ prompts_tokens (torch.Tensor): A tensor of shape [batch_size, max_seq_len]
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]: The attention mask of shape [1, 1, max_seq_len, max_seq_len] and position ids of shape [batch_size, max_seq_len]
+ """
+ seq_length = prompts_tokens.size(1)
+ attention_mask = torch.tril(
+ torch.ones((1, seq_length, seq_length), device=prompts_tokens.device)
+ ).view(1, 1, seq_length, seq_length)
+ # Convert to boolean
+ attention_mask = attention_mask < 0.5
+
+ position_ids = (
+ torch.arange(seq_length, dtype=torch.long, device=prompts_tokens.device)
+ .unsqueeze(0)
+ .expand_as(prompts_tokens)
+ )
+
+ return attention_mask, position_ids
+
+ def get_batch_for_context_window(
+ self, context_start_position: int, context_end_position: int
+ ) -> List:
+ """Returns the inference data given context window
+
+ This function gets called iteratively in a loop . Given the start and end context positions , it extracts the appropriate data.
+
+ Args:
+ context_start_position (int): Start of the context window. During the first inference step it is mostly 0
+ context_end_position (int): End of the context window. During the last inference step it will mostly be the max generated sequence length.
+
+ Returns:
+ List: A list of inputs that will be used by your model in the forward step
+ """
+ tokens2use = self.prompts_tokens[:, context_start_position:context_end_position]
+ positions2use = self.position_ids[:, context_start_position:context_end_position]
+ attention_mask2use = self.attention_mask[
+ ..., context_start_position:context_end_position, :context_end_position
+ ]
+ data_at_step_idx = [tokens2use, positions2use, attention_mask2use]
+ return data_at_step_idx
diff --git a/megatron/core/inference/model_inference_wrappers/inference_wrapper_config.py b/megatron/core/inference/model_inference_wrappers/inference_wrapper_config.py
new file mode 100644
index 0000000000..e22550e7e3
--- /dev/null
+++ b/megatron/core/inference/model_inference_wrappers/inference_wrapper_config.py
@@ -0,0 +1,40 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class InferenceWrapperConfig:
+ """Config for the model inference wrapper
+
+ NOTE : All the arguments here are obtained from arguments.py file
+ """
+
+ hidden_size: int
+ """Receive happens between the layers during PP with size [seq_len, batch_size, hidden_size]"""
+
+ params_dtype: torch.dtype
+ """Can be torch.float or torch.half if --fp16 is used, or torch.bfloat16 if --bf16 is used"""
+
+ inference_batch_times_seqlen_threshold: int
+ """if batch-size times sequence-length is smaller than this threshold then we will not use pipelining, otherwise we will."""
+
+ padded_vocab_size: int
+ """The final padded vocab size (Padded to make it divisible by --make-vocab-size-divisible-by value)"""
+
+ fp32_residual_connection: bool = False
+ """Move residual connections to fp32. Obtained from arguments.py"""
+
+ def add_attributes(self, attribute_value_pair: dict):
+ """Utility to add more attributes to inference params
+
+ Use this method to pass in a custom dictonary to add more config to the instance you created. Use as follows
+ c = InferenceWrapperConfig
+ c.add_attributes({'precision':'fp32'})
+
+ Args:
+ attribute_value_pair (dict): A dictionary containing attributes as the key names and their values as the values.
+ """
+ for key, value in attribute_value_pair.items():
+ setattr(self, key, value)
diff --git a/megatron/core/inference/model_inference_wrappers/t5/__init__.py b/megatron/core/inference/model_inference_wrappers/t5/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/inference/model_inference_wrappers/t5/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/inference/model_inference_wrappers/t5/t5_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/t5/t5_inference_wrapper.py
new file mode 100644
index 0000000000..10e1da4812
--- /dev/null
+++ b/megatron/core/inference/model_inference_wrappers/t5/t5_inference_wrapper.py
@@ -0,0 +1,205 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from argparse import Namespace
+from collections import deque
+from typing import Any, List, Tuple
+
+import numpy
+import torch
+
+from megatron.core import tensor_parallel
+from megatron.core.datasets.t5_dataset import T5MaskedWordPieceDataset
+from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
+ AbstractModelInferenceWrapper,
+)
+from megatron.core.models.T5 import T5Model
+
+
+class T5InferenceWrapper(AbstractModelInferenceWrapper):
+ """Constructor for the model inference wrapper
+
+ The wrapper prepares the model for inference, provides the required input
+ data, and runs the forward pass
+
+ Args:
+ model (T5Model): The T5 model (MCore or legacy)
+ args (Namespace): The command line arguments that were passed
+ """
+
+ def __init__(self, model: T5Model, args: Namespace):
+ super().__init__(model, args)
+
+ def prep_model_for_inference(
+ self, prompts_tokens: torch.Tensor, encoder_prompts: List[str] = None, tokenizer: Any = None
+ ):
+ """A utility function for preparing model for inference
+
+ This function is called before the forward pass. It puts the model in eval mode, builds
+ position ids, and creates attention masks so that required slices can be extracted during
+ the forward pass.
+
+ Args:
+ prompts_tokens (torch.Tensor): A tensor of shape [batch_size, max_sequence_length]
+ encoder_prompts (dict): List of string of encoder input prompts
+ tokenizer (_type_): Tokenizer used for tokenizing and detokenizing text
+ """
+
+ super().prep_model_for_inference(prompts_tokens=prompts_tokens)
+
+ encoder_prompts_tokens_list = [
+ self.tokenize_encoder_prompt(encoder_prompt, tokenizer)
+ for encoder_prompt in encoder_prompts
+ ]
+ self.batch_encoder_prompts_tokens = self.pad_encoder_prompts_tokens(
+ encoder_prompts_tokens_list, self.model.max_sequence_length, tokenizer
+ )
+
+ # create batch mask for encoder_prompt (self.batch_input_tokens) and
+ # decoder_input (self.prompts_tokens), similar to megatron/core/datasets/t5_dataset.py
+ decoder_prompts_tokens = self.prompts_tokens.cpu().numpy()
+ encoder_prompts_tokens = self.batch_encoder_prompts_tokens.cpu().numpy()
+ self.batch_mask_encoder = []
+ self.batch_mask_decoder = []
+ self.batch_mask_encoder_decoder = []
+ for i in range(len(self.prompts_tokens)):
+ self.batch_mask_encoder.append(
+ T5MaskedWordPieceDataset._make_attention_mask(
+ encoder_prompts_tokens[i], encoder_prompts_tokens[i]
+ )
+ )
+ self.batch_mask_decoder.append(
+ T5MaskedWordPieceDataset._make_attention_mask(
+ decoder_prompts_tokens[i], decoder_prompts_tokens[i]
+ )
+ * T5MaskedWordPieceDataset._make_history_mask(decoder_prompts_tokens[i])
+ )
+ self.batch_mask_encoder_decoder.append(
+ T5MaskedWordPieceDataset._make_attention_mask(
+ decoder_prompts_tokens[i], encoder_prompts_tokens[i]
+ )
+ )
+ self.batch_mask_encoder = torch.tensor(numpy.array(self.batch_mask_encoder)).cuda()
+ self.batch_mask_decoder = torch.tensor(numpy.array(self.batch_mask_decoder)).cuda()
+ self.batch_mask_encoder_decoder = torch.tensor(
+ numpy.array(self.batch_mask_encoder_decoder)
+ ).cuda()
+ self.batch_mask_encoder = self.batch_mask_encoder < 0.5
+ self.batch_mask_decoder = self.batch_mask_decoder < 0.5
+ self.batch_mask_encoder_decoder = self.batch_mask_encoder_decoder < 0.5
+
+ def tokenize_encoder_prompt(
+ self, encoder_prompt: str, tokenizer
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Utility to tokenize the encoder_prompt
+
+ Args:
+ encoder_prompt (str): The encoder_prompt
+ tokenizer (_type_): Tokenizer used for tokenizing and detokenizing string
+
+ Returns:
+ torch.Tensor: Returns the tokenized prompt
+ """
+
+ # if there is the word "" in prompt, replacing it with special_additional_token,
+ # similar to processing step in megatron/core/datasets/t5_dataset.py
+ divided_encoder_prompt_list = encoder_prompt.split("")
+ masks_count = len(divided_encoder_prompt_list) - 1
+ sentinels = deque(tokenizer.additional_special_tokens_ids)
+
+ encoder_prompt_tokens = []
+ for divided_encoder_prompt in divided_encoder_prompt_list:
+ divided_encoder_prompt_tokens = tokenizer.tokenize(divided_encoder_prompt)
+ encoder_prompt_tokens.extend(divided_encoder_prompt_tokens)
+ if masks_count > 0:
+ sentinel = sentinels.popleft()
+ encoder_prompt_tokens.extend([sentinel])
+
+ return encoder_prompt_tokens
+
+ def pad_encoder_prompts_tokens(
+ self, encoder_prompts_tokens_list: List[List[int]], max_sequence_length: int, tokenizer
+ ) -> torch.Tensor:
+ """Method to pad input prompts
+
+ Given a list of prompts, pad them all to uniform length
+
+ Args:
+ encoder_prompts_tokens_list (List[List[int]]): A list containing the
+ encoder_input_tokens
+ max_sequence_length (int): Maximum of the length of the encoder inputs tokens
+ tokenizer (_type_): Tokenizer used for tokenizing and detokenizing text
+
+ Returns:
+ torch.Tensor: A torch tensor of shape [bs, max_sequence_length]
+ """
+
+ for encoder_prompt_tokens in encoder_prompts_tokens_list:
+ padding_size = max_sequence_length - len(encoder_prompt_tokens)
+ encoder_prompt_tokens.extend([tokenizer.pad] * padding_size)
+
+ return torch.tensor(encoder_prompts_tokens_list).cuda()
+
+ def get_batch_for_context_window(
+ self, context_start_position: int, context_end_position: int
+ ) -> List:
+ """Returns the inference data given context window
+
+ This function gets called iteratively in a loop . Given the start and end context
+ positions , it extracts the appropriate data.
+
+ Args:
+ context_start_position (int): Start of the context window. During
+ the first inference step it is mostly 0
+ context_end_position (int): End of the context window. During the
+ last inference step it will mostly be the max generated sequence length.
+
+ Returns:
+ List: A list of inputs that will be used by your model in the forward step
+ """
+
+ # rerun encoder every step
+ # T5 inference not yet support kv_cache
+ encoder_tokens2use = self.batch_encoder_prompts_tokens
+ decoder_tokens2use = self.prompts_tokens[:, :context_end_position]
+ encoder_mask2use = self.batch_mask_encoder
+ decoder_mask2use = self.batch_mask_decoder[:, :context_end_position, :context_end_position]
+ encoder_decoder_mask2use = self.batch_mask_encoder_decoder[:, :context_end_position, :]
+ data_at_step_idx = [
+ encoder_tokens2use,
+ decoder_tokens2use,
+ encoder_mask2use,
+ decoder_mask2use,
+ encoder_decoder_mask2use,
+ ]
+
+ return data_at_step_idx
+
+ def forward_pass_without_pipeline_parallel(self, inference_input: List) -> torch.Tensor:
+ """Utility to carry out simple forward pass for TP or no model parallel models
+
+ Runs a very simple forward pass for model. Used in the case of models without
+ any parallelism or only tensor parallelism.
+
+ Args:
+ inference_input (List): A list containg the inputs for the gpt
+ model [tokens, position ids, attention mask]
+
+ Returns:
+ torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size]
+ """
+ [encoder_tokens, decoder_tokens, encoder_mask, decoder_mask, encoder_decoder_mask] = (
+ inference_input
+ )
+ tokens = decoder_tokens
+
+ # T5 inference not yet support kv_cache
+ logits = self.model(
+ encoder_tokens,
+ decoder_tokens,
+ encoder_mask,
+ decoder_mask,
+ encoder_decoder_mask,
+ inference_params=None,
+ )
+ logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits)
+
+ return logits
diff --git a/megatron/core/inference/modelopt_support/__init__.py b/megatron/core/inference/modelopt_support/__init__.py
new file mode 100644
index 0000000000..f8eb8f3d9f
--- /dev/null
+++ b/megatron/core/inference/modelopt_support/__init__.py
@@ -0,0 +1,8 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+"""Integrations with NVIDIA TensorRT Model Optimizer (referred as ModelOpt).
+
+ModelOpt is a library comprising state-of-the-art model optimization techniques including quantization and sparsity to
+compress model for efficient inference on NVIDIA GPUs. ModelOpt is integrated with Megatron-core to provide a seamless
+experience for users to optimize their Megatron-core models for inference. More details on ModelOpt including
+installation and usage can be found at https://github.com/NVIDIA/TensorRT-Model-Optimizer.
+"""
diff --git a/megatron/core/inference/modelopt_support/gpt/__init__.py b/megatron/core/inference/modelopt_support/gpt/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/inference/modelopt_support/gpt/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/inference/modelopt_support/gpt/model_specs.py b/megatron/core/inference/modelopt_support/gpt/model_specs.py
new file mode 100644
index 0000000000..ba1ab8993d
--- /dev/null
+++ b/megatron/core/inference/modelopt_support/gpt/model_specs.py
@@ -0,0 +1,57 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.extensions.transformer_engine import TEDotProductAttention, TENorm
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+
+
+# Use this spec for ModelOpt PTQ and TensorRT-LLM export
+def get_gpt_layer_modelopt_spec(
+ remap_te_layernorm: bool = False, qk_layernorm: bool = False
+) -> ModuleSpec:
+ """Mix the native spec with TENorm.
+
+ This is essentially the native local spec except for the layernorm implementation
+ is using TENorm from Transformer-Engine. The issue is that FusedLayerNorm from apex
+ has stopped supporting RMSNorm needed by llama.
+ """
+ sharded_state_dict_keys_map = {}
+ if remap_te_layernorm:
+ sharded_state_dict_keys_map = {
+ 'input_layernorm.': 'self_attention.linear_qkv.layer_norm_',
+ 'pre_mlp_layernorm.': 'mlp.linear_fc1.layer_norm_',
+ }
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=TENorm,
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=ColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=RowParallelLinear,
+ q_layernorm=TENorm if qk_layernorm else IdentityOp,
+ k_layernorm=TENorm if qk_layernorm else IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=TENorm,
+ mlp=ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear
+ ),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ # Map TE-layernorm-fusion keys back
+ sharded_state_dict_keys_map=sharded_state_dict_keys_map,
+ ),
+ )
diff --git a/megatron/core/inference/modelopt_support/gpt/state_dict_hooks.py b/megatron/core/inference/modelopt_support/gpt/state_dict_hooks.py
new file mode 100644
index 0000000000..15c3527c94
--- /dev/null
+++ b/megatron/core/inference/modelopt_support/gpt/state_dict_hooks.py
@@ -0,0 +1,133 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from logging import getLogger
+
+import torch
+
+logger = getLogger(__name__)
+
+
+def mcore_gpt_load_legacy_state_dict_pre_hook(
+ state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+):
+ """Register a pre-hook to fix the state_dict key difference.
+
+ This prehook is used when trying to load the legacy Megatron-LM GPTModel into its
+ megatron/core variant that uses native ParallelLinear and Transformer-Engine Norm.
+ Only this particular spec supports post-training quantization and TensorRT-LLM
+ config export through `nvidia-modelopt` package.
+
+ Args:
+ state_dict: state dictionary
+ prefix: module name prefix
+ local_metadata: local metatdata
+ strict: whether is in strict mode
+ missing_keys: missing state dict keys
+ unexpected_keys: unexpected state dict keys
+ error_msgs: error messages
+ """
+ if "modelopt_state" in state_dict:
+ state_dict.pop("modelopt_state")
+
+ if "language_model" in state_dict:
+ language_model_state_dict = state_dict.pop("language_model")
+ if "embedding" in language_model_state_dict:
+ if "word_embeddings" in language_model_state_dict["embedding"]:
+ for key, param in language_model_state_dict["embedding"]["word_embeddings"].items():
+ state_dict.update({"embedding.word_embeddings." + key: param})
+ if "position_embeddings" in language_model_state_dict["embedding"]:
+ for key, param in language_model_state_dict["embedding"][
+ "position_embeddings"
+ ].items():
+ state_dict.update({"embedding.position_embeddings." + key: param})
+ if "transformer" in language_model_state_dict:
+ for key, param in language_model_state_dict["transformer"].items():
+ state_dict.update({"decoder." + key: param})
+ else:
+ for key, param in language_model_state_dict["encoder"].items():
+ state_dict.update({"decoder." + key: param})
+ if "output_layer" in language_model_state_dict:
+ for key, param in language_model_state_dict["output_layer"].items():
+ state_dict.update({"output_layer." + key: param})
+
+ if torch.distributed.get_rank() == 0:
+ logger.info("ModelOptGPTModel {}".format(state_dict.keys()))
+
+ module_name_rewrite_list = [
+ ("input_norm", "input_layernorm"),
+ (".attention.query_key_value", ".self_attention.linear_qkv"),
+ (".attention.dense", ".self_attention.linear_proj"),
+ ("self_attention.query_key_value", "self_attention.linear_qkv"),
+ ("self_attention.dense", "self_attention.linear_proj"),
+ ("post_attention_layernorm", "pre_mlp_layernorm"),
+ ("post_attention_norm", "pre_mlp_layernorm"),
+ ("dense_h_to_4h", "linear_fc1"),
+ ("dense_4h_to_h", "linear_fc2"),
+ ("final_norm", "final_layernorm"),
+ ]
+
+ key_rewrite_list = []
+
+ for key, _ in state_dict.items():
+ for old_name, new_name in module_name_rewrite_list:
+ if old_name in key:
+ key_rewrite_list += [(key, key.replace(old_name, new_name))]
+
+ for old_key, new_key in key_rewrite_list:
+ if torch.distributed.get_rank() == 0:
+ logger.info("replace {} with {}".format(old_key, new_key))
+ state_dict[new_key] = state_dict[old_key]
+ state_dict.pop(old_key)
+
+
+def mcore_gpt_load_te_state_dict_pre_hook(
+ state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+):
+ """Register a pre-hook to fix the state_dict key difference of.
+
+ This prehook is used when trying to load the megatron/core GPTModel that uses a
+ fused Transformer-Engine ParallelLinear into the variant that uses native ParallelLinear
+ and Transformer-Engine Norm (effectively to restore the fusion).
+ Only this particular spec supports post-training quantization and TensorRT-LLM
+ config export through `nvidia-modelopt` package.
+
+ Args:
+ state_dict: state dictionary
+ prefix: module name prefix
+ local_metadata: local metatdata
+ strict: whether is in strict mode
+ missing_keys: missing state dict keys
+ unexpected_keys: unexpected state dict keys
+ error_msgs: error messages
+ """
+ if "modelopt_state" in state_dict:
+ state_dict.pop("modelopt_state")
+
+ key_with_te_extra_state_to_pop = []
+
+ for key, _ in state_dict.items():
+ if "_extra_state" in key:
+ key_with_te_extra_state_to_pop += [key]
+
+ for key in key_with_te_extra_state_to_pop:
+ state_dict.pop(key)
+
+ module_name_rewrite_list = [
+ ("self_attention.linear_qkv.layer_norm_weight", "input_layernorm.weight"),
+ ("self_attention.linear_qkv.layer_norm_bias", "input_layernorm.bias"),
+ ("mlp.linear_fc1.layer_norm_weight", "pre_mlp_layernorm.weight"),
+ ("mlp.linear_fc1.layer_norm_bias", "pre_mlp_layernorm.bias"),
+ ]
+
+ key_rewrite_list = []
+
+ for key, _ in state_dict.items():
+ for old_name, new_name in module_name_rewrite_list:
+ if old_name in key:
+ key_rewrite_list += [(key, key.replace(old_name, new_name))]
+
+ for old_key, new_key in key_rewrite_list:
+ if torch.distributed.get_rank() == 0:
+ logger.info("replace {} with {}".format(old_key, new_key))
+ state_dict[new_key] = state_dict[old_key]
+ state_dict.pop(old_key)
diff --git a/megatron/core/inference/scheduler.py b/megatron/core/inference/scheduler.py
new file mode 100644
index 0000000000..00ab81b4ab
--- /dev/null
+++ b/megatron/core/inference/scheduler.py
@@ -0,0 +1,127 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import time
+import typing
+from collections import OrderedDict
+from typing import Dict
+
+import torch
+
+from megatron.core.inference.common_inference_params import CommonInferenceParams
+from megatron.core.inference.inference_request import InferenceRequest, Status
+from megatron.core.inference.utils import Counter
+
+
+class Scheduler:
+ """Scheduler for handling requests to inference engine
+
+ This class is responsible for handing of all the incomign requests
+
+ Args:
+ max_batch_size (int): The max batch size that we can pass to the
+ inference engine at a time.
+ """
+
+ def __init__(self, max_batch_size: int):
+ self.max_batch_size = max_batch_size
+ self.active_request_pool: Dict[int, InferenceRequest] = OrderedDict()
+ self.waiting_request_pool: Dict[int, InferenceRequest] = OrderedDict()
+ self.completed_request_pool: Dict[int, InferenceRequest] = OrderedDict()
+ self.request_counter = Counter()
+
+ def add_request(
+ self,
+ prompt: str,
+ prompt_tokens: torch.Tensor,
+ encoder_prompt: str = None,
+ inference_parameters: CommonInferenceParams = None,
+ arrival_time: float = None,
+ ):
+ """Add an incoming request
+
+ This method will add the request to either the active pool or the waiting pool
+ depending on the batch size.
+
+ Args:
+ prompt (str): Input prompt string
+ prompt_tokens (torch.Tensor): A torch tensor having the input prompts tokenized
+ encoder_prompt (str): Encoder input string
+ inference_parameters (CommonInferenceParams): The inference parameters
+ arrival_time (float, optional): The incoming request time. Defaults to None.
+ """
+ request_id = str(next(self.request_counter))
+
+ if arrival_time is None:
+ arrival_time = time.time()
+
+ status = (
+ Status.ACTIVE_BUT_NOT_GENERATING_TOKENS
+ if len(self.active_request_pool) < self.max_batch_size
+ else Status.WAITING_IN_QUEUE
+ )
+
+ inference_request = InferenceRequest(
+ request_id=request_id,
+ prompt=prompt,
+ inference_parameters=inference_parameters,
+ arrival_time=arrival_time,
+ prompt_tokens=prompt_tokens,
+ status=status,
+ encoder_prompt=encoder_prompt,
+ )
+
+ if status == status.ACTIVE_BUT_NOT_GENERATING_TOKENS:
+ self.active_request_pool[request_id] = inference_request
+ else:
+ self.waiting_request_pool[request_id] = inference_request
+
+ def have_requests_pending(self) -> bool:
+ """Method to check if there are requests pending
+
+ This method returns False only when there are no active requests or waiting requests.
+ """
+ num_requests_pending = len(self.active_request_pool) + len(self.waiting_request_pool)
+ return num_requests_pending > 0
+
+ def add_earliest_waiting_request_to_active_pool(self):
+ """Utility to add the waiting request to active pool
+
+ This method will add the earliest request (FIFO) that is in the waiting request
+ pool to the active request pool.
+ """
+ assert (
+ len(self.active_request_pool) < self.max_batch_size
+ ), "Active request pool is already full. Cant add any more requests"
+ if len(self.waiting_request_pool) > 0:
+ (earliest_waiting_request_request_id, earliest_waiting_request) = (
+ self.waiting_request_pool.popitem(last=False)
+ )
+ earliest_waiting_request.status = Status.ACTIVE_BUT_NOT_GENERATING_TOKENS
+ self.active_request_pool[earliest_waiting_request_request_id] = earliest_waiting_request
+
+ def update_requests_pools(self, result_dict: typing.OrderedDict[int, InferenceRequest] = None):
+ """Update request pool status
+
+ This method will full up the active request pool, if it has less than max batch size
+ elements from the waiting request pool.
+ If provided with a request dict, it will put the completed requests into the completed
+ request pool and add waiting request into active pool.
+
+ Args:
+ result (typing.OrderedDict[int, InferenceRequest], optional): The result returned
+ by the engine. A dictionary with keys as the request ids, and values as the
+ requests. Defaults to None
+ """
+ for result_request_id in list(result_dict.keys()):
+ active_request = self.active_request_pool[result_request_id]
+
+ # If a request has completed put it into the completed request pool.
+ if active_request.status == Status.COMPLETED:
+ completed_request = self.active_request_pool.pop(result_request_id)
+ self.completed_request_pool[result_request_id] = completed_request
+
+ # If the active request pool is not full, add waiting requests in FIFO order
+ while (
+ len(self.active_request_pool) < self.max_batch_size
+ and len(self.waiting_request_pool) > 0
+ ):
+ self.add_earliest_waiting_request_to_active_pool()
diff --git a/megatron/core/inference/text_generation_controllers/__init__.py b/megatron/core/inference/text_generation_controllers/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/inference/text_generation_controllers/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/inference/text_generation_controllers/encoder_decoder_text_generation_controller.py b/megatron/core/inference/text_generation_controllers/encoder_decoder_text_generation_controller.py
new file mode 100644
index 0000000000..61beff0211
--- /dev/null
+++ b/megatron/core/inference/text_generation_controllers/encoder_decoder_text_generation_controller.py
@@ -0,0 +1,35 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from typing import OrderedDict
+
+import torch
+
+from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.inference.text_generation_controllers.simple_text_generation_controller import (
+ SimpleTextGenerationController,
+)
+
+
+class EncoderDecoderTextGenerationController(SimpleTextGenerationController):
+ """The text generation controller for encoder-decoder architecture
+
+ This class ingherits from SimpleTextGenerationController, adding features
+ relating to encoder input encoder_prompt
+
+ """
+
+ def prep_model_for_inference(
+ self, prompts_tokens: torch.Tensor, active_requests: OrderedDict[int, InferenceRequest]
+ ):
+ """Preparing batch for inference, using respective wrapper's prep_model_for_inference method
+
+ Args:
+ prompts_tokens (torch.Tensor): A tensor of shape [batch_size, max_sequence_length]
+ active_requests (OrderedDict[int, InferenceRequest]): The input active requests
+ """
+ encoder_prompts = list(
+ map(lambda request: request.encoder_prompt, active_requests.values())
+ )
+
+ self.inference_wrapped_model.prep_model_for_inference(
+ prompts_tokens=prompts_tokens, encoder_prompts=encoder_prompts, tokenizer=self.tokenizer
+ )
diff --git a/megatron/core/inference/text_generation_controllers/simple_text_generation_controller.py b/megatron/core/inference/text_generation_controllers/simple_text_generation_controller.py
new file mode 100644
index 0000000000..0667af8373
--- /dev/null
+++ b/megatron/core/inference/text_generation_controllers/simple_text_generation_controller.py
@@ -0,0 +1,400 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from typing import List, OrderedDict, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from megatron.core import parallel_state
+from megatron.core.inference.common_inference_params import CommonInferenceParams
+from megatron.core.inference.communication_utils import broadcast_from_last_pipeline_stage
+from megatron.core.inference.inference_request import InferenceRequest, Status
+from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
+ AbstractModelInferenceWrapper,
+)
+
+
+class SimpleTextGenerationController:
+ """The basic text generation controller
+
+ This class is responsible for tokenizing the input , running the inference, sampling
+ and also detokenizing the output
+
+ Args:
+ inference_wrapped_model (AbstractModelInferenceWrapper): A model that
+ is wrapped using the specs given in the abstract_model_inference_wrapper.py
+ tokenizer (_type_): Tokenizer used for tokenizing and detokenizing the prompts
+ """
+
+ def __init__(self, inference_wrapped_model: AbstractModelInferenceWrapper, tokenizer):
+ self.inference_wrapped_model = inference_wrapped_model
+ self.tokenizer = tokenizer
+
+ # For models without pipeline parallelism, is_first_stage and is_last_stage returns True
+ self.model_is_pipeline_parallel = not (
+ parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage()
+ )
+
+ def tokenize_prompt(
+ self, prompt: str, add_BOS: bool = False
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Utility to tokenize the input prompts
+
+ Args:
+ prompt (str): The input prompt
+
+ Returns:
+ torch.Tensor: Returns the tokenized prompt
+ """
+ prompt_tokens = self.tokenizer.tokenize(prompt)
+
+ if add_BOS:
+ prompt_tokens = [self.tokenizer.bos] + prompt_tokens
+
+ return prompt_tokens
+
+ def detokenize_generations(self, prompt_tokens_with_generated_tokens: torch.Tensor) -> str:
+ """Detokenize the output generations
+
+ Args:
+ prompt_tokens_with_generated_tokens (torch.Tensor): The input prompt
+ tokens plus the generated tokens
+
+ Returns:
+ str: The detokenized output
+ """
+ tokens = prompt_tokens_with_generated_tokens.cpu().numpy().tolist()
+ return self.tokenizer.detokenize(tokens)
+
+ def sample_from_logits(
+ self,
+ last_token_logits: torch.Tensor,
+ common_inference_params: CommonInferenceParams,
+ vocab_size: int = None,
+ ) -> torch.Tensor:
+ """Samples the logits to generate outputs
+
+ Given the logits of the last token, this function samples it
+ according to the parameters defined in common_inference_params
+ and returns the samples
+
+ Args:
+ last_token_logits (torch.Tensor): The last token logits. A tensor of
+ size [batch_size, vocab_size]
+ common_inference_params (CommonInferenceParams): The paramters to use
+ for inference
+ vocab_size (int): Obtained from the tokenizer. Defaults to None
+
+ Returns:
+ torch.Tensor: 1D tensor of the sampled logits with [batch_size] elements
+ """
+
+ top_p = common_inference_params.top_p
+ top_k = common_inference_params.top_k
+ temperature = common_inference_params.temperature
+
+ assert not (top_k > 0 and top_p > 0), 'Cannot have top-p and top-k both greater than zero'
+ assert top_p <= 1.0, 'top-p should be in (0,1]'
+
+ def modify_logits_for_top_k_filtering(logits, top_k):
+ """Set the logits for none top-k values to -inf."""
+ filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None]
+ logits.masked_fill_(filter_, float('-Inf'))
+
+ def modify_logits_for_top_p_filtering(logits, top_p):
+ """Set the logits for none top-p values to -inf."""
+ # First sort and calculate cumulative sum of probabilities.
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+ cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+
+ # Filteration based on the cumulative sum.
+ filter_ = cumulative_probs > top_p
+ # This shift by 1 is weird and I cannot justify it. This existed
+ # in the original implementation:
+ # https://github.com/ari-holtzman/degen/blob/master/gen.py
+ # and I guess it is needed so keeping it for now.
+ filter_[:, 1:] = filter_[:, :-1].clone()
+ # Make sure we at least have one token to select from.
+ filter_[..., 0] = 0
+
+ # Fill in the filtered part
+ filter_ = filter_.scatter(1, sorted_indices, filter_)
+ logits.masked_fill_(filter_, float('-Inf'))
+
+ # Greedy sampling
+ if top_k == 1:
+ sampled_logits = torch.argmax(last_token_logits, dim=-1)
+ else:
+ last_token_logits = last_token_logits.clone()
+ if temperature != 1.0:
+ last_token_logits.div_(temperature)
+
+ if top_k > 1:
+ assert top_k <= last_token_logits.size(1), 'top-k is larger than logit size.'
+ if vocab_size:
+ assert top_k < vocab_size, 'top-k is larger than vocab size.'
+ modify_logits_for_top_k_filtering(last_token_logits, top_k)
+
+ elif top_p > 0.0:
+ modify_logits_for_top_p_filtering(last_token_logits, top_p)
+
+ # After filtering, we need to recalculate the distribution.
+ probabilities = last_token_logits.softmax(dim=-1)
+ sampled_logits = torch.multinomial(probabilities, num_samples=1).view(-1)
+
+ # If vocab size is provided, make sure the samples are in in the range [0, vocab-size).
+ if vocab_size:
+ sampled_logits = torch.clamp(sampled_logits, min=0, max=(vocab_size - 1))
+ return sampled_logits
+
+ def update_generation_status(
+ self,
+ updated_prompts_tokens: torch.Tensor,
+ generation_started: torch.Tensor,
+ current_context_end_position: int,
+ is_generation_done_tensor: torch.Tensor,
+ generated_sequence_lengths: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Checks which prompts have reached an end condition
+
+ We check which prompts have reached an end condition and set the corresponding
+ flags of the is_generation_done_tensor to True. The generated sequence lengths
+ increase as we keep generating, until that prompts hits an end condition. The
+ generation_started tensor determines which prompts have started generating.
+
+ Args:
+ updated_prompts_tokens (torch.Tensor): The prompts tokens updated with the latest
+ generated tokens. A tensor of shape [batch_size, max_seq_len]
+ (i.e max_seq_len = max_prompt_len + tokens_to_generate)
+ generation_started (torch.Tensor): A boolean tensor of shape [batch_size]. True
+ indicates the prompt at that index has started generating tokens.
+ current_context_end_position (int): An integer indicating which position to
+ extract from the prompts tokens to get the latest generated tokens.
+ is_generation_done_tensor (torch.Tensor): A boolean tensor of shape [batch_size].
+ True indicates the prompt at that index has reached end condition.
+ generated_sequence_lengths (torch.Tensor): A int tensor of shape [batch_size].
+ Each value represents the generated sequence lengths for that prompt.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Returns the boolean
+ is_generation_done_tensor and the generated_sequence_lengths after updating it
+ """
+ latest_samples = updated_prompts_tokens[:, current_context_end_position]
+ # Make sure we are checking eod criterion only for prompts that have started generating
+ # (i.e) We only look at the generated tokenns and not the input tokens.
+ reached_eod = (latest_samples == self.tokenizer.eod) & generation_started
+ is_generation_done_tensor = is_generation_done_tensor | reached_eod
+ # We increment generated sequence lengths when that prompt has not hit the
+ # EOD and generation has started
+ generated_sequence_lengths += ~is_generation_done_tensor & generation_started
+
+ return is_generation_done_tensor, generated_sequence_lengths
+
+ def pad_input_prompt_tokens(
+ self,
+ batch_prompt_tokens_list: List[List[int]],
+ max_prompt_length_in_batch: int,
+ num_tokens_to_generate: int,
+ ) -> torch.Tensor:
+ """Method to pad input prompts
+
+ Given a list of prompts, pad them all to uniform length
+
+ Args:
+ batch_prompt_tokens_list (List[List[int]]): A list containing the prompt tokens
+ max_prompt_length_in_batch (int): Maximum of the length of the input prompt tokens
+ num_tokens_togenerate (int): The number of tokens to generate for each prompt
+
+ Returns:
+ torch.Tensor: A torch tensor of shape [bs, max_seq_len] (i.e)
+ max_seq_len = max_prompt_length_in_batch + num_tokens_to_generate,
+ with extra indices for each tensor padded with mask id.
+ """
+ max_seq_len = max_prompt_length_in_batch + num_tokens_to_generate
+
+ for prompt_tokens in batch_prompt_tokens_list:
+ padding_size = max_seq_len - len(prompt_tokens)
+ prompt_tokens.extend([self.tokenizer.eod] * padding_size)
+
+ return torch.tensor(batch_prompt_tokens_list).cuda()
+
+ def generate_output_tokens_dynamic_batch(
+ self, active_requests: OrderedDict[int, InferenceRequest]
+ ) -> OrderedDict[int, InferenceRequest]:
+ """Utility to generate the output tokens and probabilities for the prompts
+
+ This utility generates the output tokens for a dynamic batch. It will run one forward step
+ at a time, and pass control back to the engine, which will update the request pool and call
+ this method again.
+
+ Args:
+ active_requests (OrderedDict[int, InferenceRequest]): The input active requests.
+
+ Returns:
+ OrderedDict[int, InferenceRequest]: The result for each of the incoming requests
+ after running one forward step.
+ """
+ raise Exception("Not implemented yet")
+
+ def generate_all_output_tokens_static_batch(
+ self, active_requests: OrderedDict[int, InferenceRequest]
+ ) -> OrderedDict[int, InferenceRequest]:
+ """Utility to generate the all the output tokens and probabilities for the prompts .
+
+ This utility generates the output tokens for a static batch. It runs the forward steps till
+ all prompts complete generation, updates the status of these requests to completed, adds
+ the generated result and returns these requests
+
+ Args:
+ active_requests (OrderedDict[int, InferenceRequest]): The input active requests.
+
+ Returns:
+ OrderedDict[int, InferenceRequest]: The result for each of the incoming requests
+ """
+ batch_prompt_tokens_list = list(
+ map(lambda request: request.prompt_tokens, active_requests.values())
+ )
+ prompt_lengths_in_batch = torch.tensor(
+ [len(prompt_tokens) for prompt_tokens in batch_prompt_tokens_list]
+ ).cuda()
+ max_prompt_length_in_batch = max(prompt_lengths_in_batch)
+ min_prompt_length_in_batch = min(prompt_lengths_in_batch)
+
+ # For batch inference the inference params are the same for all request
+ common_inference_params: CommonInferenceParams = list(active_requests.values())[
+ 0
+ ].inference_parameters
+
+ # max_seq_len = max_prompt_length_in_batch + num_tokens_to_generate
+ batch_prompt_tokens = self.pad_input_prompt_tokens(
+ batch_prompt_tokens_list,
+ max_prompt_length_in_batch=max_prompt_length_in_batch,
+ num_tokens_to_generate=common_inference_params.num_tokens_to_generate,
+ )
+ batch_size, max_sequence_length = batch_prompt_tokens.shape
+
+ # Pre allocate log probs tensor
+ output_log_probs = None
+ if common_inference_params.return_log_probs:
+ output_log_probs = torch.empty(
+ (batch_size, max_sequence_length - 1), dtype=torch.float32
+ ).cuda()
+
+ # An array to check which of the prompts have reached end of generation condition
+ is_generation_done_tensor = torch.zeros(batch_size, dtype=torch.bool).cuda()
+
+ # An array to act as a counter to keep track of generated sequence lengths
+ generated_sequence_lengths = torch.zeros(batch_size).cuda()
+
+ with torch.no_grad():
+
+ self.prep_model_for_inference(
+ prompts_tokens=batch_prompt_tokens, active_requests=active_requests
+ )
+
+ context_start_position = 0
+ # Pick the context window that we need to pass through the network.
+ for context_end_position in range(min_prompt_length_in_batch, max_sequence_length):
+
+ inference_input = self.inference_wrapped_model.get_batch_for_context_window(
+ context_start_position, context_end_position
+ )
+
+ # Returns the final logits of shape [batch_size, context_length, vocab_size]
+ # Note: This is returned in all TP ranks or last PP stage in PP models
+ logits = self.inference_wrapped_model.run_one_forward_step(inference_input)
+ if self.model_is_pipeline_parallel:
+ context_length = context_end_position - context_start_position
+ logits = broadcast_from_last_pipeline_stage(
+ [batch_size, context_length, self.tokenizer.vocab_size],
+ dtype=torch.float32,
+ tensor=logits,
+ )
+
+ # Indicates which of the input prompts have started generating tokens.
+ # A 1D boolean tensor with [batch_size] elements (i.e) The shortest
+ # prompts will start generating first and so on
+ generation_started = prompt_lengths_in_batch <= context_end_position
+ last_token_logits = logits[:, -1, :]
+ sampled_logits = self.sample_from_logits(
+ last_token_logits, common_inference_params, self.tokenizer.vocab_size
+ )
+
+ # Substitute the sampled logits only for only the prompts that
+ # have started generating tokens
+ batch_prompt_tokens[generation_started, context_end_position] = sampled_logits[
+ generation_started
+ ]
+
+ if common_inference_params.return_log_probs:
+ log_probs = F.log_softmax(logits, dim=2)
+ indices = torch.unsqueeze(
+ batch_prompt_tokens[
+ :, (context_start_position + 1) : (context_end_position + 1)
+ ],
+ 2,
+ )
+ # Get the log probabilities for only the prompt tokens
+ output_log_probs[:, context_start_position:context_end_position] = torch.gather(
+ log_probs, 2, indices
+ ).squeeze(2)
+
+ context_start_position = context_end_position
+
+ # Check end of generation status for each tensor
+ # and update generated sequence lengths
+ (is_generation_done_tensor, generated_sequence_lengths) = (
+ self.update_generation_status(
+ updated_prompts_tokens=batch_prompt_tokens,
+ generation_started=generation_started,
+ current_context_end_position=context_end_position,
+ is_generation_done_tensor=is_generation_done_tensor,
+ generated_sequence_lengths=generated_sequence_lengths,
+ )
+ )
+ # Boolean flag indicating if all prompts are finished
+ all_prompts_done = torch.all(is_generation_done_tensor)
+ if all_prompts_done:
+ break
+
+ # Include all the generated tokens
+ batch_prompt_tokens_with_generations = batch_prompt_tokens[:, : (context_end_position + 1)]
+ if common_inference_params.return_log_probs:
+ output_log_probs = output_log_probs[:, :context_end_position]
+
+ generated_sequence_lengths[
+ generated_sequence_lengths > common_inference_params.num_tokens_to_generate
+ ] = common_inference_params.num_tokens_to_generate
+
+ for idx, request in enumerate(active_requests.values()):
+ input_prompt_length = int(prompt_lengths_in_batch[idx])
+ # Shorter prompts might have generated more than required tokens. So we trim them down
+ required_sequence_length = int(
+ min(generated_sequence_lengths[idx], common_inference_params.num_tokens_to_generate)
+ )
+ # Extract only the generated tokens
+ required_result_tokens = batch_prompt_tokens_with_generations[
+ idx, input_prompt_length : (input_prompt_length + required_sequence_length)
+ ]
+
+ request.generated_length = required_sequence_length
+ request.generated_tokens = required_result_tokens
+ request.generated_log_probs = (
+ None
+ if output_log_probs is None
+ else output_log_probs[idx, input_prompt_length:required_sequence_length]
+ )
+ request.status = Status.COMPLETED
+ request.generated_text = self.detokenize_generations(required_result_tokens)
+
+ return active_requests
+
+ def prep_model_for_inference(
+ self, prompts_tokens: torch.Tensor, active_requests: OrderedDict[int, InferenceRequest]
+ ):
+ """Preparing batch for inference, using respective wrapper's prep_model_for_inference method
+
+ Args:
+ prompts_tokens (torch.Tensor): A tensor of shape [batch_size, max_sequence_length]
+ active_requests (OrderedDict[int, InferenceRequest]): The input active requests
+ """
+ self.inference_wrapped_model.prep_model_for_inference(prompts_tokens=prompts_tokens)
diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
new file mode 100644
index 0000000000..bdb1021ef5
--- /dev/null
+++ b/megatron/core/inference/utils.py
@@ -0,0 +1,17 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+class Counter:
+ """A simple counter class
+
+ This class is responsible for assigning request ids to incoming requests
+ """
+
+ def __init__(self, start: int = 0) -> None:
+ self.counter = start
+
+ def __next__(self) -> int:
+ i = self.counter
+ self.counter += 1
+ return i
+
+ def reset(self) -> None:
+ self.counter = 0
diff --git a/megatron/core/inference_params.py b/megatron/core/inference_params.py
new file mode 100644
index 0000000000..0db49e3115
--- /dev/null
+++ b/megatron/core/inference_params.py
@@ -0,0 +1,31 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+class InferenceParams:
+ """Inference parameters that are passed to the main model in order
+ to efficienly calculate and store the context during inference."""
+
+ def __init__(self, max_batch_size, max_sequence_length):
+ self.max_sequence_length = max_sequence_length
+ self.max_batch_size = max_batch_size
+ self.sequence_len_offset = 0
+ self.batch_size_offset = 0
+ self.key_value_memory_dict = {}
+
+ def swap_key_value_dict(self, batch_idx):
+ "swap between batches"
+ if len(self.key_value_memory_dict) == 0:
+ raise ValueError("should not swap when dict in empty")
+
+ for layer_number in self.key_value_memory_dict.keys():
+ inference_key_memory, inference_value_memory = self.key_value_memory_dict[layer_number]
+ assert (
+ len(batch_idx) == inference_key_memory.shape[1]
+ ) # make sure batch size is the same
+ new_inference_key_memory = inference_key_memory[:, batch_idx]
+ new_inference_value_memory = inference_value_memory[:, batch_idx]
+ self.key_value_memory_dict[layer_number] = (
+ new_inference_key_memory,
+ new_inference_value_memory,
+ )
+
+ def __str__(self):
+ return f"InferenceParams(max_seq_len = {self.max_sequence_length}, max_batch_size = {self.max_batch_size}, sequence_len_offset = {self.sequence_len_offset}, batch_size_offset = {self.batch_size_offset}, key_value_memory_dict = {self.key_value_memory_dict.keys()})"
diff --git a/megatron/core/jit.py b/megatron/core/jit.py
new file mode 100644
index 0000000000..8bb18d393c
--- /dev/null
+++ b/megatron/core/jit.py
@@ -0,0 +1,11 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import torch
+
+TORCH_MAJOR = int(torch.__version__.split(".")[0])
+TORCH_MINOR = int(torch.__version__.split(".")[1])
+
+jit_fuser = torch.jit.script
+# nvFuser is deprecated in PyTorch JIT starting from 2.2
+if (TORCH_MAJOR > 2) or (TORCH_MAJOR == 2 and TORCH_MINOR >= 2):
+ jit_fuser = torch.compile
diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py
new file mode 100644
index 0000000000..f2751673e4
--- /dev/null
+++ b/megatron/core/model_parallel_config.py
@@ -0,0 +1,344 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Callable, ContextManager, Optional
+
+import torch
+
+
+@dataclass
+class ModelParallelConfig:
+ """Base configuration for Megatron Core
+
+ The initialization function has an argument for each parameter.
+ """
+
+ ###################
+ # Model parallelism
+ ###################
+ tensor_model_parallel_size: int = 1
+ """Intra-layer model parallelism. Splits tensors across GPU ranks."""
+
+ pipeline_model_parallel_size: int = 1
+ """Inter-layer model parallelism. Splits transformer layers across GPU ranks."""
+
+ virtual_pipeline_model_parallel_size: Optional[int] = None
+ """Interleaved pipeline parallelism is used to improve performance by reducing the pipeline
+ bubble. Considers a transformer block as a list of smaller transformer (virtual) blocks.
+ The number of virtual blocks per pipeline model parallel rank is the virtual model parallel
+ size. See Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM:
+ arxiv.org/pdf/2104.04473.pdf for more details.
+ """
+
+ sequence_parallel: bool = False
+ """Makes tensor parallelism more memory efficient for LLMs (20B+) by parallelizing layer norms
+ and dropout sequentially. See Reducing Activation Recomputation in Large Transformer Models
+ (https://arxiv.org/abs/2205.05198) for more details.
+ """
+
+ context_parallel_size: int = 1
+ """Splits network input along sequence dimension across GPU ranks."""
+
+ expert_model_parallel_size: int = 1
+ """Distributes Moe Experts across sub data parallel dimension."""
+
+ moe_extended_tp: bool = False
+ """Alternative parallelization strategy for expert parallelism. Instead of distributing experts
+ across expert_model_parallel_size, each expert is sharded along extendended tensor parallel
+ domain (tensor_model_paralle_size * expert_model_parallel_size). It avoids the load balancing
+ problem with MOE training.
+ """
+
+ ###################
+ # Initialization
+ ###################
+ perform_initialization: bool = True
+ """If true, weights are initialized. This option can be useful when you know you are going to
+ load values from a checkpoint.
+ """
+
+ use_cpu_initialization: bool = False
+ """When set to False, we initialize the weights directly on the GPU. CPU initialization is the
+ same regardless of tensor model parallelism, but GPU initialization is not. Transferring
+ weights from CPU to GPU can take a significant amount of time for large models.
+ """
+
+ ###################
+ # Training
+ ###################
+ fp16: bool = False
+ """If true, train with fp16 mixed precision training."""
+
+ bf16: bool = False
+ """If true, train with bf16 mixed precision training."""
+
+ params_dtype: torch.dtype = torch.float32
+ """dtype used when intializing the weights."""
+
+ timers: Callable = None
+ """Timers object to call for various timing functions. See megatron.core.timers.Timers"""
+
+ finalize_model_grads_func: Callable = None
+ """Function that finalizes gradients on all workers. Could include ensuring that grads are
+ all-reduced across data parallelism, pipeline parallelism, and sequence parallelism
+ dimensions.
+ """
+
+ grad_scale_func: Callable = None
+ """If using loss scaling, this function should take the loss and return the scaled loss. If
+ None, no function is called on the loss.
+ """
+
+ no_sync_func: Callable = None
+ """Function that creates a context that suppresses asynchronous data-parallel communication. If
+ the model is an instance of core.distributed.DistributedDataParallel, the default is to use
+ core.distributed.DistributedDataParallel.no_sync.
+ """
+
+ grad_sync_func: Callable = None
+ """Function that launches asynchronous gradient reductions (e.g. distributed optimizer gradient
+ reduce-scatters). The function should take one argument: an iterable of parameters whose
+ gradients are to be synchronized.
+ """
+
+ param_sync_func: Callable = None
+ """Function that launches asynchronous parameter synchronizations (e.g. distributed optimizer
+ parameter all-gathers). The function should take one argument: an iterable of parameters to
+ be synchronized.
+ """
+
+ deterministic_mode: bool = False
+ """If true, code that has deterministic execution will be chosen. This usually
+ means slower execution, but is good for debugging and testing. Defaults to False."""
+
+ enable_autocast: bool = False
+ """If true runs the forward step function inside torch.autocast context."""
+
+ autocast_dtype: torch.dtype = None
+ """dtype to pass to torch.amp.autocast when enabled. If None, is set to pipeline_dtype."""
+
+ num_microbatches_with_partial_activation_checkpoints: Optional[int] = None
+ """If int, set the number of microbatches where not all of the layers will be checkpointed and
+ recomputed. The rest of the microbatches within the window of maximum outstanding
+ microbatches will recompute all layers (either full recompute or selective recompute). If
+ None, the checkpoint and recompute will be left up to the forward_step function.
+
+ """
+
+ ###################
+ # Optimizations
+ ###################
+ gradient_accumulation_fusion: bool = False
+ """If true, fuses weight gradient accumulation to GEMMs. Requires the custom CUDA extension
+ fused_weight_gradient_mlp_cuda module. To use gradient_accumulation_fusion you must install
+ APEX with --cpp_ext and --cuda_ext. For example: "pip install --global-option=\"--cpp_ext\"
+ --global-option=\"--cuda_ext\" ". Note that the extension requires CUDA>=11. Otherwise, you
+ must turn off gradient accumulation fusion.
+ """
+
+ async_tensor_model_parallel_allreduce: bool = False
+ """NOTE: Deprecated. This flag is ignored."""
+
+ use_te_rng_tracker: bool = False
+ """If true, uses RNG state tracker in TransformerEngine if exists.
+ """
+
+ tp_comm_overlap: bool = False
+ """If true, allows overlapping of Linear layer execution with tensor parallel communication
+ collectives like AllGather/ReduceScatter. Overlapping is done for the linear layers wherever
+ possible during the forward and the backward pass.
+ """
+
+ tp_comm_bulk_wgrad: bool = True
+ """If true, allows All-Gather overlap with Bprop activation gradient GEMM. Don't care if
+ tp_comm_overlap is False.
+ """
+
+ tp_comm_bulk_dgrad: bool = True
+ """If true, allows Reduce-Scatter overlap with Bprop weight gradient GEMM. Don't care if
+ tp_comm_overlap is False.
+ """
+
+ tp_comm_overlap_ag: bool = True
+ """If true, allows All-Gather overlap with GEMM by pipelining the GEMM and All-Gather.
+ Don't care if tp_comm_overlap is False.
+ """
+
+ tp_comm_overlap_rs: bool = True
+ """If true, allows Reduce-Scatter overlap with GEMM by pipelining the GEMM and Reduce-Scatter.
+ Don't care if tp_comm_overlap is False.
+ """
+
+ tp_comm_overlap_rs_dgrad: bool = False
+ """If true, allows Reduce-Scatter overlap with DGRAD GEMM by pipelining the
+ GEMM and Reduce-Scatter splits. Don't care if tp_comm_overlap is False.
+ """
+
+ tp_comm_split_ag: bool = True
+ """Deprecated from TransformerEngine v1.6.0.
+ If true, allows All-Gather overlap with Fprop GEMM by pipelining the GEMM and All-Gather
+ splits. Don't care if tp_comm_overlap is False.
+ """
+
+ tp_comm_atomic_ag: bool = False
+ """Deprecated from TransformerEngine v1.6.0.
+ If true, allows All-Gather overlap with Fprop GEMM by pipelining the GEMM and All-Gather
+ both done atomically. Don't care if tp_comm_overlap is False.
+ """
+
+ tp_comm_split_rs: bool = True
+ """Deprecated from TransformerEngine v1.6.0.
+ If true, allows Reduce-Scatter overlap with Fprop GEMM by pipelining the GEMM and
+ Reduce-Scatter splits. Don't care if tp_comm_overlap is False.
+ """
+
+ tp_comm_atomic_rs: bool = False
+ """Deprecated from TransformerEngine v1.6.0.
+ If true, allows Reduce-Scatter overlap with Fprop GEMM by pipelining the GEMM and
+ Reduce-Scatter both done atomically. Don't care if tp_comm_overlap is False.
+ """
+
+ cross_entropy_loss_fusion: bool = False
+ """If this is enabled, the fused cross entropy implementation would be used.
+ Defaults to False.
+ """
+
+ tp_comm_overlap_disable_qkv: bool = False
+ """
+ If true, the AllGather -> Gemm overlap for QKV gets disabled
+ """
+
+ tp_comm_overlap_disable_fc1: bool = False
+ """
+ If true, the AllGather -> Gemm overlap for FC1 layer of MLP gets disabled
+ """
+
+ tp_comm_bootstrap_backend: str = 'nccl'
+ """
+ Set the bootstrapping backend out of 'nccl', 'mpi', and 'gloo'
+ """
+
+ ###################
+ # Pipeline Parallel
+ ###################
+ pipeline_dtype: torch.dtype = None
+ """dtype used in p2p communication, usually params_dtype"""
+
+ variable_seq_lengths: bool = False
+ """Support for variable sequence lengths across microbatches. Setting this communicates the size
+ of tensors during pipeline parallelism communication, because of this extra overhead it
+ should only be set if the sequence length varies by microbatch within a global batch.
+ """
+
+ overlap_p2p_comm: bool = False
+ """When True some of the peer to peer communication for pipeline parallelism will overlap with
+ computation. Must be False if batch_p2p_comm is true.
+ """
+
+ batch_p2p_comm: bool = True
+ """Use batch_isend_irecv instead of individual isend/irecv calls. Must be False if
+ overlap_p2p_comm is True.
+ """
+
+ batch_p2p_sync: bool = True
+ """When using batch_isend_irecv, do a cuda.device.synchronize afterward to work around a bug in
+ older version of PyTorch.
+ """
+
+ use_ring_exchange_p2p: bool = False
+ """Use custom ring_exchange kernel instead of torch.distributed.batch_isend_irecv(). Requires
+ custom built torch with torch.distributed.ring_exchange.
+ """
+
+ deallocate_pipeline_outputs: bool = False
+ """If True, output data is deallocated after the tensor is sent to the next pipeline stage.
+ Helps with saving memory, does nothing when pipeline parallel is not used.
+ """
+
+ defer_embedding_wgrad_compute: bool = False
+ """If true, defers the embedding WGRAD GEMMs while pipeline flush is
+ taking place enabling us to hide pipeline flush latency. Defaults to False.
+ """
+
+ wgrad_deferral_limit: int = 0
+ """This value tunes the number of micro-batches for which the embedding weight gradient compute
+ needs to be deferred to pipeline flush, this argument is invalid if
+ `defer_embedding_wgrad_compute` is False.
+ Defaults to 0, which means all micro-batches are deferred.
+ """
+
+ pipeline_model_parallel_split_rank: Optional[int] = None
+ """If int, rank where encoder and decoder should be split in cases where the model has both an
+ encoder and decoder (e.g., T5). Ignored if None.
+ """
+
+ ###################
+ # CPU Offloading
+ ###################
+ cpu_offloading: bool = False
+ """When set to True, all the activations are offloaded to the CPU asynchronously."""
+
+ cpu_offloading_num_layers: int = 0
+ """Tells the number of transformer layers for which activations has to be offloaded."""
+
+ _cpu_offloading_context: ContextManager = (
+ None
+ # Used for internal use only, not to be set by a user.
+ # TODO: Need to move to the 'right' place when possible.
+ )
+ """For internal use only, do not set."""
+
+ cpu_offloading_activations: bool = True
+ """If True, offloads the activations to CPU."""
+
+ cpu_offloading_weights: bool = True
+ """If True, offloads the weights to CPU."""
+
+ ###################
+ # Timing
+ ###################
+ barrier_with_L1_time: bool = True
+ """If true, use barrier with level 1 time measurements. It is up to the user to make sure
+ calling barrier with their timers will not result in hangs. This can happen if for example
+ the user adds a level 1 timer that is not called by all ranks.
+ """
+
+ def __post_init__(self):
+ """Python dataclass method that is used to modify attributes after initialization.
+ See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more
+ details.
+ """
+ if self.sequence_parallel:
+ if self.tensor_model_parallel_size <= 1:
+ raise ValueError("Can not use sequence paralllelism without tensor parallelism")
+
+ if self.pipeline_model_parallel_size > 1:
+ if self.pipeline_dtype is None:
+ raise ValueError(
+ "When using pipeline parallelism, pipeline_dtype must be specified"
+ )
+
+ if self.autocast_dtype is None:
+ self.autocast_dtype = self.params_dtype
+
+ if self.defer_embedding_wgrad_compute and self.pipeline_model_parallel_size == 1:
+ raise ValueError(
+ "Cannot defer embedding wgrad compute when pipeline model parallel is not used"
+ )
+
+ if self.defer_embedding_wgrad_compute and not self.gradient_accumulation_fusion:
+ raise ValueError(
+ "Cannot defer embedding wgrad compute when gradient accumulation fusion is not used"
+ )
+
+ if self.defer_embedding_wgrad_compute and self.wgrad_deferral_limit < 0:
+ raise ValueError(
+ "Wgrad deferral limit should be greater than or equal to 0 when it is enabled!"
+ )
+
+ if self.expert_model_parallel_size > 1 and self.tensor_model_parallel_size > 1:
+ if self.sequence_parallel is False:
+ raise ValueError(
+ "When using expert parallelism and tensor parallelism, sequence parallelism "
+ "must be used"
+ )
diff --git a/megatron/core/models/T5/__init__.py b/megatron/core/models/T5/__init__.py
new file mode 100644
index 0000000000..2551f81e65
--- /dev/null
+++ b/megatron/core/models/T5/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from .t5_model import T5Model
diff --git a/megatron/core/models/T5/t5_model.py b/megatron/core/models/T5/t5_model.py
new file mode 100644
index 0000000000..bce998c6e8
--- /dev/null
+++ b/megatron/core/models/T5/t5_model.py
@@ -0,0 +1,422 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from typing import List, Literal, Optional
+
+import torch
+from torch import Tensor
+
+from megatron.core import InferenceParams, tensor_parallel
+from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
+from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
+from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.transformer.enums import ModelType
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class T5LMHead(MegatronModule):
+ """Masked LM head for T5
+
+ Args:
+ config (TransformerConfig): transformer config
+ parallel_output (bool): wether output logits being distributed or not.
+ vocab_size (int): vocabulary size
+ pre_process (bool): Include embedding layer
+ share_embeddings_and_output_weights (bool): When True, input
+ embeddings and output logit weights are shared.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ parallel_output: bool,
+ vocab_size: int,
+ pre_process: bool = True,
+ share_embeddings_and_output_weights: bool = False,
+ ):
+ super(T5LMHead, self).__init__(config=config)
+
+ if has_config_logger_enabled(config):
+ log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+ self.parallel_output = parallel_output
+
+ self.output_layer = tensor_parallel.ColumnParallelLinear(
+ config.hidden_size,
+ vocab_size,
+ config=config,
+ init_method=config.init_method,
+ bias=share_embeddings_and_output_weights,
+ skip_bias_add=not share_embeddings_and_output_weights,
+ gather_output=not self.parallel_output,
+ skip_weight_param_allocation=pre_process and share_embeddings_and_output_weights,
+ )
+
+ def forward(self, hidden_states: Tensor, word_embeddings_weight: Tensor) -> Tensor:
+ """Forward pass.
+
+ Args:
+ hidden_states (Tensor): output hidden states from decoder
+ word_embeddings_weight (Tensor): word embedding weight
+
+ Returns:
+ Tensor: logits tensor
+ """
+
+ logits, _ = self.output_layer(hidden_states, weight=word_embeddings_weight)
+ return logits
+
+
+class T5Model(LanguageModule):
+ """T5 Language model.
+
+ Args:
+ config (TransformerConfig): transformer config
+
+ encoder_config (TransformerConfig): encoder transformer config
+
+ transformer_encoder_layer_spec (ModuleSpec): transformer layer
+ customization specs for encoder
+
+ transformer_decoder_layer_spec (ModuleSpec): transformer layer
+ customization specs for decoder
+
+ vocab_size (int): vocabulary size
+
+ max_sequence_length (int): maximum size of sequence. This is used for positional embedding
+
+ pre_process (bool): Include embedding layer (used with pipeline parallelism)
+
+ post_process (bool): Include an output layer (used with pipeline parallelism)
+
+ fp16_lm_cross_entropy (bool, optional): Defaults to False
+
+ parallel_output (bool): Do not gather the outputs,
+ keep them split across tensor parallel ranks
+
+ share_embeddings_and_output_weights (bool): When True,
+ input embeddings and output logit weights are shared. Defaults to False.
+
+ position_embedding_type (string): Position embedding type.
+ Options ['learned_absolute', 'rope'].
+ Defaults is 'learned_absolute'.
+
+ rotary_percent (float): Percent of rotary dimension to use for rotary position embeddings.
+ Defaults to 1.0 (100%). Ignored unless position_embedding_type is 'rope'.
+
+ seq_len_interpolation_factor (float): scale of linearly interpolating
+ RoPE for longer sequences. The value must be a float larger than 1.0.
+ Defaults to None.
+
+ add_encoder (bool): Create the encoder (used with pipeline parallelism).
+ When using pipelining, the encoder will only be created on a subset
+ of the pipeline ranks.
+
+ add_decoder (bool): Include an output layer (used with pipeline parallelism).
+ As with `add_encoder`, when using this model and pipelining,
+ the decoder will only be created on a subset of the pipeline ranks.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ encoder_config: TransformerConfig,
+ transformer_encoder_layer_spec: ModuleSpec,
+ transformer_decoder_layer_spec: ModuleSpec,
+ vocab_size: int,
+ max_sequence_length: int,
+ pre_process: bool = True,
+ post_process: bool = True,
+ fp16_lm_cross_entropy: bool = False,
+ parallel_output: bool = True,
+ share_embeddings_and_output_weights: bool = False,
+ position_embedding_type: Literal['learned_absolute', 'rope'] = 'learned_absolute',
+ rotary_percent: float = 1.0,
+ seq_len_interpolation_factor: Optional[float] = None,
+ add_encoder: bool = True,
+ add_decoder: bool = True,
+ ):
+
+ super(T5Model, self).__init__(config=config)
+
+ self.config: TransformerConfig = config
+ self.encoder_config: TransformerConfig = encoder_config
+ self.transformer_encoder_layer_spec: ModuleSpec = transformer_encoder_layer_spec
+ self.transformer_decoder_layer_spec: ModuleSpec = transformer_decoder_layer_spec
+ self.vocab_size = vocab_size
+ self.max_sequence_length = max_sequence_length
+ self.pre_process = pre_process
+ self.post_process = post_process
+ self.add_encoder = add_encoder
+ self.add_decoder = add_decoder
+ self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
+ self.parallel_output = parallel_output
+ self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+ self.position_embedding_type = position_embedding_type
+ self.encoder_hidden_state = None
+
+ self.model_type = ModelType.encoder_and_decoder
+
+ # Tells schedules.py that this model has a skip connection
+ # between the encoder's output and the decoder
+ # (and hence both the encoder and decoder's tensors are required for correct backprop).
+ self.xattn_needed = True
+
+ # specify the position embeddings as a member
+ # variable in the T5 class so that they are easy to
+ # find for `finalize_model_grads._allreduce_position_embedding_grads`
+ self.position_embeddings = None
+ if self.pre_process:
+ self.embedding = LanguageModelEmbedding(
+ config=self.config,
+ vocab_size=self.vocab_size,
+ max_sequence_length=self.max_sequence_length,
+ position_embedding_type=self.position_embedding_type,
+ )
+ self.position_embeddings = self.embedding.position_embeddings
+
+ # Rotary Position Embeddings
+ if self.position_embedding_type == 'rope':
+ self.rotary_pos_emb = RotaryEmbedding(
+ kv_channels=self.config.kv_channels,
+ rotary_percent=rotary_percent,
+ rotary_interleaved=self.config.rotary_interleaved,
+ seq_len_interpolation_factor=seq_len_interpolation_factor,
+ use_cpu_initialization=self.config.use_cpu_initialization,
+ )
+
+ # Transformer encoder
+ encoder_spec, decoder_spec = (
+ self.transformer_encoder_layer_spec,
+ self.transformer_decoder_layer_spec,
+ )
+ if self.add_encoder:
+ self.encoder = TransformerBlock(
+ config=self.encoder_config,
+ spec=encoder_spec,
+ pre_process=self.pre_process,
+ post_process=self.post_process,
+ )
+ else:
+ self.encoder = None
+
+ if self.add_decoder:
+ # Transformer decoder
+ self.decoder = TransformerBlock(
+ config=self.config,
+ spec=decoder_spec,
+ pre_process=self.pre_process,
+ post_process=self.post_process,
+ )
+ else:
+ self.decoder = None
+
+ # Output
+ if post_process:
+ self.lm_head = T5LMHead(
+ config,
+ parallel_output,
+ self.vocab_size,
+ self.pre_process,
+ self.share_embeddings_and_output_weights,
+ )
+ self.output_layer = self.lm_head.output_layer
+
+ if self.pre_process or self.post_process:
+ self.setup_embeddings_and_output_layer()
+
+ def forward(
+ self,
+ encoder_input_ids: Tensor,
+ decoder_input_ids: Tensor,
+ encoder_attn_mask: Tensor,
+ decoder_attn_mask: Tensor,
+ encoder_decoder_attn_mask: Tensor,
+ lm_labels: Tensor = None,
+ encoder_hidden_states: Tensor = None,
+ output_encoder_hidden_only: bool = False,
+ inference_params: InferenceParams = None,
+ ) -> Tensor:
+ """Forward pass.
+
+ Args:
+ encoder_input_ids (Tensor): input ids for encoder
+ decoder_input_ids (Tensor): input ids for decoder
+ encoder_attn_mask (Tensor): self-attention mask for encoder
+ decoder_attn_mask (Tensor): self-attention mask for decoder
+ encoder_decoder_attn_mask (Tensor): cross-attention mask between encoder and decoder
+ lm_labels (Tensor): labels for decoder output
+ inference_params (InferenceParams): relevant arguments for inferencing
+
+ Returns:
+ Tensor: loss tensor
+ """
+
+ (encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask) = (
+ t5_extended_attention_mask(
+ [encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask]
+ )
+ )
+
+ ## Encoder forward
+ if encoder_hidden_states is None:
+
+ # Encoder position ids
+ encoder_position_ids = t5_position_ids(encoder_input_ids)
+
+ # Encoder embedding.
+ if self.pre_process:
+ encoder_input = self.embedding(
+ input_ids=encoder_input_ids, position_ids=encoder_position_ids
+ )
+ else:
+ # intermediate stage of pipeline
+ encoder_input = None
+
+ # Rotary positional embeddings
+ rotary_pos_emb = None
+ if self.position_embedding_type == 'rope':
+ rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+ inference_params, self.encoder, encoder_input, self.config
+ )
+ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)
+
+ # Run encoder.
+ if self.add_encoder:
+ encoder_hidden_states = self.encoder(
+ hidden_states=encoder_input,
+ attention_mask=encoder_attn_mask,
+ inference_params=inference_params,
+ rotary_pos_emb=rotary_pos_emb,
+ )
+ else:
+ encoder_hidden_states = self.encoder_hidden_state
+
+ if not self.add_decoder or output_encoder_hidden_only:
+ return encoder_hidden_states
+
+ ## Decoder forward
+ # Decoder position ids
+ decoder_position_ids = t5_position_ids(decoder_input_ids)
+
+ # Decoder embedding.
+ if self.pre_process:
+ decoder_input = self.embedding(
+ input_ids=decoder_input_ids, position_ids=decoder_position_ids
+ )
+ else:
+ # intermediate stage of pipeline
+ decoder_input = None ### should it take encoder_hidden_states
+
+ # Rotary positional embeddings
+ rotary_pos_emb = None
+ if self.position_embedding_type == 'rope':
+ rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+ inference_params, self.decoder, decoder_input, self.config
+ )
+ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)
+
+ # Run decoder.
+ decoder_hidden_states = self.decoder(
+ hidden_states=decoder_input,
+ attention_mask=decoder_attn_mask,
+ context=encoder_hidden_states,
+ context_mask=encoder_decoder_attn_mask,
+ inference_params=inference_params,
+ rotary_pos_emb=rotary_pos_emb,
+ )
+
+ if self.post_process:
+ lm_logits = self.lm_head(
+ decoder_hidden_states, self.shared_embedding_or_output_weight()
+ )
+ if lm_labels is None:
+ # [s b h] => [b s h]
+ return lm_logits.transpose(0, 1).contiguous()
+ else:
+ # [b s] => [s b]
+ lm_loss = self.compute_language_model_loss(lm_labels, lm_logits)
+ return lm_loss
+ else:
+ return decoder_hidden_states
+
+ def set_input_tensor(self, input_tensor):
+ """See megatron.model.transformer.set_input_tensor()"""
+
+ # This is usually handled in schedules.py but some inference code still
+ # gives us non-lists or None
+ if not isinstance(input_tensor, list):
+ input_tensor = [input_tensor]
+
+ if self.add_encoder and self.add_decoder:
+ assert (
+ len(input_tensor) == 1
+ ), 'input_tensor should only be length 1 for stage with both encoder and decoder'
+ self.encoder.set_input_tensor(input_tensor[0])
+ elif self.add_encoder:
+ assert (
+ len(input_tensor) == 1
+ ), 'input_tensor should only be length 1 for stage with only encoder'
+ self.encoder.set_input_tensor(input_tensor[0])
+ elif self.add_decoder:
+ if len(input_tensor) == 2:
+ self.decoder.set_input_tensor(input_tensor[0])
+ self.encoder_hidden_state = input_tensor[1]
+ elif len(input_tensor) == 1:
+ self.decoder.set_input_tensor(None)
+ self.encoder_hidden_state = input_tensor[0]
+ else:
+ raise Exception('input_tensor must have either length 1 or 2')
+ else:
+ raise Exception('Stage must have at least either encoder or decoder')
+
+ def shared_embedding_or_output_weight(self) -> Tensor:
+ """Function to share the input embeddings and output logit weights."""
+
+ if self.pre_process:
+ return self.embedding.word_embeddings.weight
+ elif self.post_process:
+ return self.lm_head.output_layer.weight
+ return None
+
+
+def t5_extended_attention_mask(attention_mask_list: List[Tensor]) -> List[Tensor]:
+ """Creates the extended attention mask
+
+ Converts the attention mask of dimension [batch size, seq_len, seq_len]
+ to [batch size, 1, seq_len, seq_len]
+
+ Args:
+ attention_mask (Tensor): The input attention mask
+
+ Returns:
+ Tensor: The extended binary attention mask
+ """
+
+ def attn_mask_postprocess(attn_mask):
+ # [b, 1, s, s]
+ extended_attention_mask = attn_mask.unsqueeze(1)
+ return extended_attention_mask
+
+ return [
+ (attn_mask_postprocess(attn_mask) if attn_mask is not None else None)
+ for attn_mask in attention_mask_list
+ ]
+
+
+def t5_position_ids(token_ids: Tensor) -> Tensor:
+ """Calculate position ids from token ids
+ Args:
+ token_ids (Tensor): input tokens
+
+ Returns:
+ Tensor: position ids
+ """
+ seq_length = token_ids.size(1)
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=token_ids.device)
+ position_ids = position_ids.unsqueeze(0).expand_as(token_ids)
+
+ return position_ids
diff --git a/megatron/core/models/T5/t5_spec.py b/megatron/core/models/T5/t5_spec.py
new file mode 100644
index 0000000000..ecdcdbc260
--- /dev/null
+++ b/megatron/core/models/T5/t5_spec.py
@@ -0,0 +1,248 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.attention import (
+ CrossAttention,
+ CrossAttentionSubmodules,
+ SelfAttention,
+ SelfAttentionSubmodules,
+)
+from megatron.core.transformer.dot_product_attention import DotProductAttention
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
+from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+
+try:
+ from megatron.core.extensions.transformer_engine import (
+ TEColumnParallelLinear,
+ TEDotProductAttention,
+ TELayerNormColumnParallelLinear,
+ TENorm,
+ TERowParallelLinear,
+ )
+
+ HAVE_TE = True
+except ImportError:
+ HAVE_TE = False
+
+try:
+ import apex # pylint: disable=unused-import
+
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+
+ HAVE_APEX = True
+ LNImpl = FusedLayerNorm
+except ImportError:
+ import warnings
+
+ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm
+
+ warnings.warn(f'Apex is not installed. Falling back to Torch LayerNorm')
+ LNImpl = WrappedTorchLayerNorm
+
+
+def encoder_model_with_transformer_engine_default_spec() -> ModuleSpec:
+ """T5 encoder TE spec (uses Transformer Engine components)."""
+
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.arbitrary},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=TELayerNormColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ q_layernorm=IdentityOp,
+ k_layernorm=IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ mlp=ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear
+ ),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+
+
+def decoder_model_with_transformer_engine_default_spec() -> ModuleSpec:
+ """T5 decoder TE spec (uses Transformer Engine components)."""
+
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=TELayerNormColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ q_layernorm=IdentityOp,
+ k_layernorm=IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_cross_attn_layernorm=TENorm,
+ cross_attention=ModuleSpec(
+ module=CrossAttention,
+ params={"attn_mask_type": AttnMaskType.arbitrary},
+ submodules=CrossAttentionSubmodules(
+ linear_q=TEColumnParallelLinear,
+ linear_kv=TEColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ ),
+ ),
+ cross_attn_bda=get_bias_dropout_add,
+ mlp=ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear
+ ),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+
+
+def encoder_model_with_local_spec() -> ModuleSpec:
+ """T5 encoder local spec (uses Megatron-Core components)."""
+
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=LNImpl,
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.arbitrary},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ q_layernorm=IdentityOp,
+ k_layernorm=IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=LNImpl,
+ mlp=ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear
+ ),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ sharded_state_dict_keys_map={
+ 'input_layernorm.': 'self_attention.linear_qkv.layer_norm_',
+ 'pre_mlp_layernorm.': 'mlp.linear_fc1.layer_norm_',
+ },
+ ),
+ )
+
+
+def decoder_model_with_local_spec() -> ModuleSpec:
+ """T5 decoder local spec (uses Megatron-Core components)."""
+
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=LNImpl,
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ q_layernorm=IdentityOp,
+ k_layernorm=IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_cross_attn_layernorm=LNImpl,
+ cross_attention=ModuleSpec(
+ module=CrossAttention,
+ params={"attn_mask_type": AttnMaskType.arbitrary},
+ submodules=CrossAttentionSubmodules(
+ linear_q=ColumnParallelLinear,
+ linear_kv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ ),
+ ),
+ cross_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=LNImpl,
+ mlp=ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear
+ ),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ sharded_state_dict_keys_map={
+ 'input_layernorm.': 'self_attention.linear_qkv.layer_norm_',
+ 'pre_mlp_layernorm.': 'mlp.linear_fc1.layer_norm_',
+ },
+ ),
+ )
+
+
+def get_t5_encoder_with_transformer_engine_block_spec(
+ num_layers: int,
+) -> TransformerBlockSubmodules:
+ """T5 encoder block spec for Transformer Engine
+
+ Args:
+ config (TransformerConfig): config, containing number of layers for encoder
+ """
+
+ layer_spec = encoder_model_with_transformer_engine_default_spec()
+ block_spec = TransformerBlockSubmodules([layer_spec] * num_layers, layer_norm=TENorm)
+ return block_spec
+
+
+def get_t5_decoder_with_transformer_engine_block_spec(
+ num_layers: int,
+) -> TransformerBlockSubmodules:
+ """T5 decoder block spec for Transformer Engine
+
+ Args:
+ config (TransformerConfig): config, containing number of layers for decoder
+ """
+
+ layer_spec = decoder_model_with_transformer_engine_default_spec()
+ block_spec = TransformerBlockSubmodules([layer_spec] * num_layers, layer_norm=TENorm)
+ return block_spec
+
+
+def get_t5_encoder_with_local_block_spec(num_layers: int) -> TransformerBlockSubmodules:
+ """T5 encoder block spec for local (uses Megatron-Core components)
+
+ Args:
+ num_layers (int): number of encoder layers
+ """
+
+ layer_spec = encoder_model_with_local_spec()
+ block_spec = TransformerBlockSubmodules([layer_spec] * num_layers, layer_norm=TENorm)
+ return block_spec
+
+
+def get_t5_decoder_with_local_block_spec(num_layers: int) -> TransformerBlockSubmodules:
+ """T5 decoder block spec for local (uses Megatron-Core components)
+
+ Args:
+ num_layers (int): number of decoder layers
+ """
+
+ layer_spec = decoder_model_with_local_spec()
+ block_spec = TransformerBlockSubmodules([layer_spec] * num_layers, layer_norm=TENorm)
+ return block_spec
diff --git a/megatron/core/models/__init__.py b/megatron/core/models/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/megatron/core/models/bert/__init__.py b/megatron/core/models/bert/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/megatron/core/models/bert/bert_layer_specs.py b/megatron/core/models/bert/bert_layer_specs.py
new file mode 100644
index 0000000000..cd51c124c9
--- /dev/null
+++ b/megatron/core/models/bert/bert_layer_specs.py
@@ -0,0 +1,92 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
+from megatron.core.transformer.dot_product_attention import DotProductAttention
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+
+try:
+ from megatron.core.extensions.transformer_engine import (
+ TEDotProductAttention,
+ TELayerNormColumnParallelLinear,
+ TERowParallelLinear,
+ )
+
+ HAVE_TE = True
+except ImportError:
+ HAVE_TE = False
+
+try:
+ import apex # pylint: disable=unused-import
+
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+
+ HAVE_APEX = True
+ LNImpl = FusedLayerNorm
+except ImportError:
+ import warnings
+
+ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm
+
+ warnings.warn(f'Apex is not installed. Falling back to Torch LayerNorm')
+ LNImpl = WrappedTorchLayerNorm
+
+# Use this spec to use lower level Transformer Engine modules (required for fp8 training)
+bert_layer_with_transformer_engine_spec = ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.padding},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=TELayerNormColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ q_layernorm=IdentityOp,
+ k_layernorm=IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ mlp=ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear
+ ),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ ),
+)
+
+# Use this spec for an implementation using only modules in megatron core
+bert_layer_local_spec = ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=LNImpl,
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.padding},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ q_layernorm=IdentityOp,
+ k_layernorm=IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=LNImpl,
+ mlp=ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ sharded_state_dict_keys_map={
+ 'input_layernorm.': 'self_attention.linear_qkv.layer_norm_',
+ 'pre_mlp_layernorm.': 'mlp.linear_fc1.layer_norm_',
+ },
+ ),
+)
diff --git a/megatron/core/models/bert/bert_lm_head.py b/megatron/core/models/bert/bert_lm_head.py
new file mode 100644
index 0000000000..fd26ebd16f
--- /dev/null
+++ b/megatron/core/models/bert/bert_lm_head.py
@@ -0,0 +1,54 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import torch
+from torch import Tensor
+
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import get_linear_layer
+
+try:
+ import apex
+
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+
+ HAVE_APEX = True
+ LNImpl = FusedLayerNorm
+except ImportError:
+ import warnings
+
+ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm
+
+ warnings.warn(f'Apex is not installed. Falling back to Torch LayerNorm')
+ LNImpl = WrappedTorchLayerNorm
+
+
+class BertLMHead(MegatronModule):
+ """Masked LM head for Bert.
+
+ Args:
+ hidden_size: hidden size
+ config (TransformerConfig): TransformerConfig object
+ """
+
+ def __init__(self, hidden_size: int, config: TransformerConfig):
+ super().__init__(config=config)
+
+ # TODO: Should switch this to TE ?
+ self.dense = get_linear_layer(
+ hidden_size, hidden_size, config.init_method, config.perform_initialization
+ )
+
+ setattr(self.dense.weight, 'sequence_parallel', config.sequence_parallel)
+ setattr(self.dense.bias, 'sequence_parallel', config.sequence_parallel)
+
+ self.layer_norm = LNImpl(
+ config=config, hidden_size=hidden_size, eps=config.layernorm_epsilon
+ )
+
+ self.gelu = torch.nn.functional.gelu
+
+ def forward(self, hidden_states: Tensor) -> Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.gelu(hidden_states)
+ hidden_states = self.layer_norm(hidden_states)
+ return hidden_states
diff --git a/megatron/core/models/bert/bert_model.py b/megatron/core/models/bert/bert_model.py
new file mode 100644
index 0000000000..eb08d4cfd6
--- /dev/null
+++ b/megatron/core/models/bert/bert_model.py
@@ -0,0 +1,366 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+import os
+import warnings
+from typing import Literal, Optional
+
+import torch
+from torch import Tensor
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
+from megatron.core.models.bert.bert_layer_specs import bert_layer_local_spec
+from megatron.core.models.bert.bert_lm_head import BertLMHead
+from megatron.core.models.bert.pooler import Pooler
+from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
+from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.transformer.enums import AttnMaskType, ModelType
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import get_linear_layer
+from megatron.core.utils import get_te_version as _get_te_version
+from megatron.core.utils import is_te_min_version
+
+
+def get_te_version():
+ """Included for backwards compatibility."""
+ warnings.warn("`get_te_version` will be deprecated in a future release")
+ return _get_te_version()
+
+
+class BertModel(LanguageModule):
+ """Transformer language model.
+
+ Args:
+ config (TransformerConfig): transformer config
+ num_tokentypes (int) : Set to 2 when args.bert_binary_head is True, and 0 otherwise.
+ Defaults to 0.
+ transformer_layer_spec (ModuleSpec): Specifies module to use for transformer layers
+ vocab_size (int): vocabulary size
+ max_sequence_length (int): maximum size of sequence. This is used for positional embedding
+ pre_process (bool): Include embedding layer (used with pipeline parallelism)
+ post_process (bool): Include an output layer (used with pipeline parallelism)
+ parallel_output (bool): Do not gather the outputs, keep them split across tensor parallel
+ ranks
+ share_embeddings_and_output_weights (bool): When True, input embeddings and output logit
+ weights are shared. Defaults to False.
+ position_embedding_type (string): Position embedding type.
+ Options ['learned_absolute', 'rope']. Defaults is 'learned_absolute'.
+ rotary_percent (float): Percent of rotary dimension to use for rotary position embeddings.
+ Defaults to 1.0 (100%). Ignored unless position_embedding_type is 'rope'.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ num_tokentypes: int,
+ transformer_layer_spec: ModuleSpec,
+ vocab_size: int,
+ max_sequence_length: int,
+ pre_process: bool = True,
+ post_process: bool = True,
+ fp16_lm_cross_entropy: bool = False,
+ parallel_output: bool = True,
+ share_embeddings_and_output_weights: bool = False,
+ position_embedding_type: Literal['learned_absolute', 'rope'] = 'learned_absolute',
+ rotary_percent: float = 1.0,
+ seq_len_interpolation_factor: Optional[float] = None,
+ add_binary_head=True,
+ return_embeddings=False,
+ ):
+ super(BertModel, self).__init__(config=config)
+
+ if has_config_logger_enabled(config):
+ log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+ if return_embeddings:
+ assert self.post_process and self.add_binary_head
+
+ self.config: TransformerConfig = config
+ self.transformer_layer_spec: ModuleSpec = transformer_layer_spec
+ self.vocab_size = vocab_size
+ self.max_sequence_length = max_sequence_length
+ self.pre_process = pre_process
+ self.post_process = post_process
+ self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
+ self.parallel_output = parallel_output
+ self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+ self.position_embedding_type = position_embedding_type
+ self.add_binary_head = add_binary_head
+ self.return_embeddings = return_embeddings
+
+ # megatron core pipelining currently depends on model type
+ self.model_type = ModelType.encoder_or_decoder
+
+ self.attn_mask_dimensions = self._sanity_check_attention_and_get_attn_mask_dimension()
+
+ # Embeddings.
+ if self.pre_process:
+ self.embedding = LanguageModelEmbedding(
+ config=self.config,
+ vocab_size=self.vocab_size,
+ max_sequence_length=self.max_sequence_length,
+ position_embedding_type=position_embedding_type,
+ num_tokentypes=num_tokentypes,
+ )
+
+ if self.position_embedding_type == 'rope':
+ self.rotary_pos_emb = RotaryEmbedding(
+ kv_channels=self.config.kv_channels,
+ rotary_percent=rotary_percent,
+ rotary_interleaved=self.config.rotary_interleaved,
+ seq_len_interpolation_factor=seq_len_interpolation_factor,
+ use_cpu_initialization=self.config.use_cpu_initialization,
+ )
+
+ # Transformer.
+ self.encoder = TransformerBlock(
+ config=self.config,
+ spec=self.transformer_layer_spec,
+ pre_process=self.pre_process,
+ post_process=self.post_process,
+ )
+
+ # Output
+ if post_process:
+ # TODO: Make sure you are passing in the mpu_vocab_size properly
+ self.lm_head = BertLMHead(config.hidden_size, config)
+
+ self.output_layer = tensor_parallel.ColumnParallelLinear(
+ config.hidden_size,
+ self.vocab_size,
+ config=config,
+ init_method=config.init_method,
+ bias=True,
+ skip_bias_add=False,
+ gather_output=not self.parallel_output,
+ skip_weight_param_allocation=pre_process and share_embeddings_and_output_weights,
+ )
+
+ self.binary_head = None
+ if self.add_binary_head:
+ # TODO: Shoudl switch this to TE ?
+ self.binary_head = get_linear_layer(
+ config.hidden_size, 2, config.init_method, config.perform_initialization
+ )
+
+ self.pooler = Pooler(
+ config.hidden_size, config.init_method, config, config.sequence_parallel
+ )
+
+ if self.pre_process or self.post_process:
+ self.setup_embeddings_and_output_layer()
+
+ # pylint: disable=line-too-long
+ def _sanity_check_attention_and_get_attn_mask_dimension(self) -> str:
+ """We do some checks and return attention mask dimensions for self attention
+
+ Transformer engine library underwent a lot of change. So we need to change dimensions of
+ the attention mask depending on the TE version. We also santiy check some arguments.
+
+ 1. If we use local version of attention dimension of the mask is [b,1,s,s]
+ 2. If we use transformer engine > 1.10 we support all 3 backends with padding mask and [b,1,s,s]
+ 3. If we use transformer engine >= 1.7 but less than 1.10
+ a ) Flash and Fused attention uses padding mask with [b,1,1,s]
+ b ) Unfused attention works with arbitrary mask with [b,1,s,s]
+ 4. If we use transformer engine < 1.7
+ Flash and fused attention is not supported. Unfused attention will work with padding mask [b,1,s,s]
+
+ Default if you dont set any NVTE_ATTN flag will it will just use the fused path for transformer engine version >= 1.7 and unfused path for other
+
+ Args:
+ transformer_layer_spec (ModuleSpec): The transformer layer spec
+
+ Returns:
+ str: A string showing the format of the attn mask dimensions
+ """
+ attn_mask_dimensions = None
+ # For local layer spec we just use b1ss
+ if self.transformer_layer_spec == bert_layer_local_spec:
+ attn_mask_dimensions = "b1ss"
+ else:
+ attn_mask_type = self.transformer_layer_spec.submodules.self_attention.params[
+ 'attn_mask_type'
+ ]
+ flash_attention_enabled = os.getenv('NVTE_FLASH_ATTN') == '1'
+ fused_attention_enabled = os.getenv('NVTE_FUSED_ATTN') == '1'
+ # For TE >= 1.10 (We always use padding mask and use b11s)
+ if is_te_min_version("1.10.0"):
+ attn_mask_dimensions = "b11s"
+ if attn_mask_type != AttnMaskType.padding:
+ warnings.warn(
+ f'For TE versions >= 1.10 , flash/fused/unfused support padding mask. Setting attention mask from {attn_mask_type} to padding'
+ )
+ self.transformer_layer_spec.submodules.self_attention.params[
+ 'attn_mask_type'
+ ] = AttnMaskType.padding
+ # For 1.7 >= TE < 1.10 flash and fused path use padding mask with b11s and unfused path uses arbitrary mask with b1ss
+ elif is_te_min_version("1.7.0"):
+ if flash_attention_enabled or fused_attention_enabled:
+ attn_mask_dimensions = "b11s"
+ else:
+ if attn_mask_type != AttnMaskType.arbitrary:
+ warnings.warn(
+ f'For TE versions >= 1.7 but < 1.10 , unfused path supports only arbitrary mask. Setting attention mask from {attn_mask_type} to arbitray'
+ )
+ self.transformer_layer_spec.submodules.self_attention.params[
+ 'attn_mask_type'
+ ] = AttnMaskType.arbitrary
+ attn_mask_dimensions = "b1ss"
+ # For TE < 1.7 we only support unfused attention with b1ss and padding mask
+ else:
+ attn_mask_dimensions = "b1ss"
+ assert not flash_attention_enabled and not fused_attention_enabled, (
+ "Flash and fused attention is not supported with transformer engine version "
+ "< 1.7. Set NVTE_FLASH_ATTN=0 and NVTE_FUSED_ATTN=0 or upgrade transformer "
+ "engine >= 1.7"
+ )
+
+ return attn_mask_dimensions
+
+ def bert_extended_attention_mask(self, attention_mask: Tensor) -> Tensor:
+ """Creates the extended attention mask
+
+ Converts the attention mask of dimension
+ [batch size, 1, seq len] to [batch size, 1, seq len, seq len]
+ or [batch size, 1, 1, seq_len] and makes it binary
+
+ Args:
+ attention_mask (Tensor): The input attention mask
+
+ Returns:
+ Tensor: The extended binary attention mask
+ """
+ # We create a 3D attention mask from a 2D tensor mask.
+ if self.attn_mask_dimensions == "b1ss":
+ # [b, 1, s]
+ attention_mask_b1s = attention_mask.unsqueeze(1)
+ # [b, s, 1]
+ attention_mask_bs1 = attention_mask.unsqueeze(2)
+ # [b, s, s]
+ attention_mask_bss = attention_mask_b1s * attention_mask_bs1
+ # [b, 1, s, s]
+ extended_attention_mask = attention_mask_bss.unsqueeze(1)
+ else:
+ # [b, 1, 1, s]
+ extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)
+
+ # Convert attention mask to binary:
+ extended_attention_mask = extended_attention_mask < 0.5
+
+ return extended_attention_mask
+
+ def bert_position_ids(self, token_ids):
+ """Position ids for bert model"""
+ # Create position ids
+ seq_length = token_ids.size(1)
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=token_ids.device)
+ position_ids = position_ids.unsqueeze(0).expand_as(token_ids)
+
+ return position_ids
+
+ def set_input_tensor(self, input_tensor: Tensor) -> None:
+ """Sets input tensor to the model.
+
+ See megatron.model.transformer.set_input_tensor()
+
+ Args:
+ input_tensor (Tensor): Sets the input tensor for the model.
+ """
+ # This is usually handled in schedules.py but some inference code still
+ # gives us non-lists or None
+ if not isinstance(input_tensor, list):
+ input_tensor = [input_tensor]
+
+ assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert'
+ self.encoder.set_input_tensor(input_tensor[0])
+
+ def forward(
+ self,
+ input_ids: Tensor,
+ attention_mask: Tensor,
+ tokentype_ids: Tensor = None,
+ lm_labels: Tensor = None,
+ inference_params=None,
+ ):
+ """Forward function of BERT model
+
+ Forward function of the BERT Model This function passes the input tensors
+ through the embedding layer, and then the encoder and finally into the post
+ processing layer (optional).
+
+ It either returns the Loss values if labels are given or the final hidden units
+ """
+ extended_attention_mask = self.bert_extended_attention_mask(attention_mask)
+
+ if parallel_state.is_pipeline_first_stage():
+ input_ids = input_ids
+ position_ids = self.bert_position_ids(input_ids)
+ else:
+ position_ids = None
+ input_ids = None
+
+ # Encoder embedding.
+ if self.pre_process:
+ encoder_input = self.embedding(
+ input_ids=input_ids, position_ids=position_ids, tokentype_ids=tokentype_ids
+ )
+ else:
+ # intermediate stage of pipeline
+ # encoder will get hidden_states from encoder.input_tensor
+ encoder_input = None
+
+ # Rotary positional embeddings (Why not move this into BERT/GPTEmberdding ?)
+ rotary_pos_emb = None
+ if self.position_embedding_type == 'rope':
+ rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+ inference_params, self.encoder, encoder_input, self.config
+ )
+ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)
+
+ # Run encoder.
+ hidden_states = self.encoder(
+ hidden_states=encoder_input,
+ attention_mask=extended_attention_mask,
+ inference_params=inference_params,
+ rotary_pos_emb=rotary_pos_emb,
+ )
+ if not self.post_process:
+ return hidden_states
+
+ if self.add_binary_head:
+ pooled_output = self.pooler(hidden_states, 0)
+
+ if self.return_embeddings:
+ embeddings = torch.transpose(hidden_states, 0, 1)
+ masks = torch.sum(attention_mask, dim=1)
+ # Collect masked embeddings.
+ output = torch.zeros(
+ size=(embeddings.shape[0], embeddings.shape[2]),
+ dtype=torch.float32,
+ device=torch.cuda.current_device(),
+ )
+ for i, (embedding, mask) in enumerate(zip(embeddings, masks)):
+ output[i, :] = torch.mean(embedding[1 : mask - 1], dim=0)
+ return output
+
+ # logits and loss
+ output_weight = None
+ if self.share_embeddings_and_output_weights:
+ output_weight = self.shared_embedding_or_output_weight()
+
+ hidden_states_after_lm_head = self.lm_head(hidden_states=hidden_states)
+ logits, _ = self.output_layer(hidden_states_after_lm_head, weight=output_weight)
+
+ binary_logits = None
+ if self.binary_head is not None:
+ binary_logits = self.binary_head(pooled_output)
+
+ if lm_labels is None:
+ # [s b h] => [b s h]
+ return logits.transpose(0, 1).contiguous(), binary_logits
+
+ loss = self.compute_language_model_loss(lm_labels, logits)
+
+ return loss, binary_logits
diff --git a/megatron/core/models/bert/pooler.py b/megatron/core/models/bert/pooler.py
new file mode 100644
index 0000000000..e0de1a845a
--- /dev/null
+++ b/megatron/core/models/bert/pooler.py
@@ -0,0 +1,52 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import torch
+from torch import Tensor
+
+from megatron.core import tensor_parallel
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import get_linear_layer
+
+
+class Pooler(MegatronModule):
+ """Pooler layer.
+
+ Pool hidden states of a specific token (for example start of the
+ sequence) and add a linear transformation followed by a tanh.
+
+ Args:
+ hidden_size (int): The hidden size_
+ init_method (callable): weight initialization method for the linear layer. bias is set to zero.
+ config (TransformerConfig): The transformer configuration
+ sequence_parallel (bool): Using squence parallel ? Defaults to False
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ init_method: callable,
+ config: TransformerConfig,
+ sequence_parallel: bool = False,
+ ):
+ super(Pooler, self).__init__(config)
+ # TODO: Shoudl switch this to TE ?
+ self.dense = get_linear_layer(
+ hidden_size, hidden_size, init_method, config.perform_initialization
+ )
+ self.sequence_parallel = sequence_parallel
+
+ def forward(self, hidden_states: Tensor, sequence_index=0):
+ # hidden_states: [s, b, h]
+ # sequence_index: index of the token to pool.
+
+ # gather data along sequence dimensions
+ # same pooler is run on all tensor parallel nodes
+ if self.sequence_parallel:
+ hidden_states = tensor_parallel.gather_from_sequence_parallel_region(
+ hidden_states, tensor_parallel_output_grad=False
+ )
+
+ pooled = hidden_states[sequence_index, :, :]
+ pooled = self.dense(pooled)
+ pooled = torch.tanh(pooled)
+ return pooled
diff --git a/megatron/core/models/common/__init__.py b/megatron/core/models/common/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/megatron/core/models/common/embeddings/__init__.py b/megatron/core/models/common/embeddings/__init__.py
new file mode 100644
index 0000000000..865f96da5d
--- /dev/null
+++ b/megatron/core/models/common/embeddings/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from .rope_utils import apply_rotary_pos_emb
+from .rotary_pos_embedding import RotaryEmbedding
+from .yarn_rotary_pos_embedding import YarnRotaryEmbedding, _yarn_get_mscale
diff --git a/megatron/core/models/common/embeddings/language_model_embedding.py b/megatron/core/models/common/embeddings/language_model_embedding.py
new file mode 100644
index 0000000000..bc1a2de9cb
--- /dev/null
+++ b/megatron/core/models/common/embeddings/language_model_embedding.py
@@ -0,0 +1,137 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+from megatron.core import tensor_parallel
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class LanguageModelEmbedding(MegatronModule):
+ """Language model embeddings.
+
+ Args:
+ config (TransformerConfig): config object with all necessary configs for TransformerBlock
+ vocab_size (int): vocabulary size
+ max_sequence_length (int): maximum size of sequence. This
+ is used for positional embedding
+ add_position_embedding (bool): Add a position embedding.
+ embedding_dropout_prob (float): dropout probability for embeddings
+ num_tokentypes (int): Set to 0 without binary head, and 2 with a binary head . Defaults to 0.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ vocab_size: int,
+ max_sequence_length: int,
+ position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'learned_absolute',
+ num_tokentypes: int = 0,
+ ):
+ super().__init__(config=config)
+
+ self.config: TransformerConfig = config
+ self.vocab_size: int = vocab_size
+ self.max_sequence_length: int = max_sequence_length
+ self.add_position_embedding: bool = position_embedding_type == 'learned_absolute'
+ self.num_tokentypes = num_tokentypes
+ self.reduce_scatter_embeddings = (
+ (not self.add_position_embedding)
+ and self.num_tokentypes <= 0
+ and self.config.sequence_parallel
+ )
+
+ # Word embeddings (parallel).
+ self.word_embeddings = tensor_parallel.VocabParallelEmbedding(
+ num_embeddings=self.vocab_size,
+ embedding_dim=self.config.hidden_size,
+ init_method=self.config.init_method,
+ reduce_scatter_embeddings=self.reduce_scatter_embeddings,
+ config=self.config,
+ )
+
+ # Position embedding (serial).
+ if self.add_position_embedding:
+ self.position_embeddings = torch.nn.Embedding(
+ self.max_sequence_length, self.config.hidden_size
+ )
+
+ # Initialize the position embeddings.
+ if self.config.perform_initialization:
+ self.config.init_method(self.position_embeddings.weight)
+
+ if self.num_tokentypes > 0:
+ self.tokentype_embeddings = torch.nn.Embedding(
+ self.num_tokentypes, self.config.hidden_size
+ )
+ # Initialize the token-type embeddings.
+ if self.config.perform_initialization:
+ self.config.init_method(self.tokentype_embeddings.weight)
+ else:
+ self.tokentype_embeddings = None
+
+ # Embeddings dropout
+ self.embedding_dropout = torch.nn.Dropout(self.config.hidden_dropout)
+
+ def zero_parameters(self):
+ """Zero out all parameters in embedding."""
+ self.word_embeddings.weight.data.fill_(0)
+ self.word_embeddings.weight.shared = True
+ self.position_embeddings.weight.data.fill_(0)
+ self.position_embeddings.weight.shared = True
+ if self.num_tokentypes > 0:
+ self.tokentype_embeddings.weight.data.fill_(0)
+ self.tokentype_embeddings.weight.shared = True
+
+ def forward(self, input_ids: Tensor, position_ids: Tensor, tokentype_ids: int = None) -> Tensor:
+ """Forward pass of the embedding module.
+
+ Args:
+ input_ids (Tensor): The input tokens
+ position_ids (Tensor): The position id's used to calculate position embeddings
+ tokentype_ids (int): The token type ids. Used when args.bert_binary_head is set to True. Defaults to None
+
+ Returns:
+ Tensor: The output embeddings
+ """
+ word_embeddings = self.word_embeddings(input_ids)
+ if self.add_position_embedding:
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings = word_embeddings + position_embeddings
+ else:
+ embeddings = word_embeddings
+
+ if not self.reduce_scatter_embeddings:
+ # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
+ embeddings = embeddings.transpose(0, 1).contiguous()
+
+ if tokentype_ids is not None:
+ assert self.tokentype_embeddings is not None
+ # [b s h] -> [s b h] (So that it can be added with embeddings)
+ tokentype_embedding = self.tokentype_embeddings(tokentype_ids).permute(1, 0, 2)
+ embeddings = embeddings + tokentype_embedding
+ else:
+ assert self.tokentype_embeddings is None
+
+ # If the input flag for fp32 residual connection is set, convert for float.
+ if self.config.fp32_residual_connection:
+ embeddings = embeddings.float()
+
+ # Dropout.
+ if self.config.sequence_parallel:
+ if not self.reduce_scatter_embeddings:
+ embeddings = tensor_parallel.scatter_to_sequence_parallel_region(embeddings)
+ # `scatter_to_sequence_parallel_region` returns a view, which prevents
+ # the original tensor from being garbage collected. Clone to facilitate GC.
+ # Has a small runtime cost (~0.5%).
+ if self.config.clone_scatter_output_in_embedding:
+ embeddings = embeddings.clone()
+ with tensor_parallel.get_cuda_rng_tracker().fork():
+ embeddings = self.embedding_dropout(embeddings)
+ else:
+ embeddings = self.embedding_dropout(embeddings)
+
+ return embeddings
diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py
new file mode 100644
index 0000000000..accb251961
--- /dev/null
+++ b/megatron/core/models/common/embeddings/rope_utils.py
@@ -0,0 +1,191 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+ from megatron.core.transformer.transformer_config import TransformerConfig
+
+import logging
+
+import torch
+from torch import Tensor
+
+from megatron.core import parallel_state
+
+logger = logging.getLogger(__name__)
+
+try:
+ from apex.transformer.functional import (
+ fused_apply_rotary_pos_emb,
+ fused_apply_rotary_pos_emb_thd,
+ )
+
+ HAVE_APPLY_ROPE_FUSION = True
+except ImportError:
+ HAVE_APPLY_ROPE_FUSION = False
+
+
+def get_pos_emb_on_this_cp_rank(pos_emb: Tensor, seq_dim: int) -> Tensor:
+ """Get the position embedding on the current context parallel rank.
+
+ Args:
+ pos_emb (Tensor): Positional embedding tensor
+ seq_dim (int): Sequence dimension
+ """
+ cp_size = parallel_state.get_context_parallel_world_size()
+ cp_rank = parallel_state.get_context_parallel_rank()
+ cp_idx = torch.tensor(
+ [cp_rank, (2 * cp_size - cp_rank - 1)], device="cpu", pin_memory=True
+ ).cuda(non_blocking=True)
+ pos_emb = pos_emb.view(
+ *pos_emb.shape[:seq_dim], 2 * cp_size, -1, *pos_emb.shape[(seq_dim + 1) :]
+ )
+ pos_emb = pos_emb.index_select(seq_dim, cp_idx)
+ pos_emb = pos_emb.view(*pos_emb.shape[:seq_dim], -1, *pos_emb.shape[(seq_dim + 2) :])
+ return pos_emb
+
+
+def _rotate_half(x: Tensor, rotary_interleaved: bool) -> Tensor:
+ """Change sign so the last dimension becomes [-odd, +even]
+
+ Args:
+ x (Tensor): Input tensor
+
+ Returns:
+ Tensor: Tensor rotated half
+ """
+ if not rotary_interleaved:
+ x1, x2 = torch.chunk(x, 2, dim=-1)
+ return torch.cat((-x2, x1), dim=-1)
+ else:
+ x1 = x[:, :, :, ::2]
+ x2 = x[:, :, :, 1::2]
+ x_new = torch.stack((-x2, x1), dim=-1)
+ return x_new.view(x_new.shape[0], x_new.shape[1], x_new.shape[2], -1)
+
+
+def _apply_rotary_pos_emb_bshd(
+ t: Tensor,
+ freqs: Tensor,
+ rotary_interleaved: bool = False,
+ multi_latent_attention: bool = False,
+ mscale: float = 1.0,
+) -> Tensor:
+ """Apply rotary positional embedding to input tensor T.
+
+ check https://kexue.fm/archives/8265 for detailed formulas
+
+ Args:
+ t (Tensor): Input tensor T is of shape [seq_length, ... , dim]
+ freqs (Tensor): Rotary Positional embedding tensor freq is of shape [seq_length, ..., dim]
+
+ Returns:
+ Tensor: The input tensor after applying RoPE
+ """
+ rot_dim = freqs.shape[-1]
+
+ # ideally t_pass is empty so rotary pos embedding is applied to all tensor t
+ t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+
+ if multi_latent_attention:
+ x1 = t[..., 0::2]
+ x2 = t[..., 1::2]
+ t = torch.cat((x1, x2), dim=-1)
+
+ # first part is cosine component
+ # second part is sine component, need to change signs with _rotate_half method
+ cos_ = (torch.cos(freqs) * mscale).to(t.dtype)
+ sin_ = (torch.sin(freqs) * mscale).to(t.dtype)
+
+ t = (t * cos_) + (_rotate_half(t, rotary_interleaved) * sin_)
+ return torch.cat((t, t_pass), dim=-1)
+
+
+def _apply_rotary_pos_emb_thd(
+ t: Tensor,
+ cu_seqlens: Tensor,
+ freqs: Tensor,
+ rotary_interleaved: bool = False,
+ multi_latent_attention: bool = False,
+ mscale: float = 1.0,
+) -> Tensor:
+ """A baseline implementation of applying RoPE for `thd` format.
+
+ Args:
+ t (Tensor): Input tensor T is of shape [t, h, d]
+ cu_seqlens(Tensor): Cumulative sum of sequence lengths in a batch for `t`,
+ with shape [b + 1] and dtype torch.int32.
+ freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d]
+
+ Returns:
+ Tensor: Shape [t, h, d]. The input tensor after applying RoPE.
+ """
+
+ seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+ return torch.cat(
+ [
+ _apply_rotary_pos_emb_bshd(
+ x.unsqueeze(1),
+ freqs[: x.size(0)],
+ rotary_interleaved=rotary_interleaved,
+ multi_latent_attention=multi_latent_attention,
+ mscale=mscale,
+ )
+ for x in torch.split(t, seqlens)
+ ]
+ ).squeeze(1)
+
+
+def apply_rotary_pos_emb(
+ t: Tensor,
+ freqs: Tensor,
+ config: TransformerConfig,
+ cu_seqlens: Optional[Tensor] = None,
+ mscale: float = 1.0,
+):
+ """
+ Reroute to the appropriate apply_rotary_pos_emb function depending on
+ fused/unfused kernels, or bshd (conventional) / thd (packed seq) format
+ """
+ if config.apply_rope_fusion and not HAVE_APPLY_ROPE_FUSION:
+ # setting apply_rope_fusion in config to False
+ # so that subsequent queries to this config also return False
+ config.apply_rope_fusion = False
+ if not getattr(apply_rotary_pos_emb, "printed_fused_warning", False):
+ logger.warning(
+ "Setting apply_rope_fusion to false because its implementation"
+ " is not included in Apex. Try upgrading to the latest version"
+ )
+ apply_rotary_pos_emb.printed_fused_warning = True
+
+ if getattr(config, "multi_latent_attention", False) and config.rotary_interleaved:
+ logger.warning(
+ "rotary_interleaved is not supported with multi_latent_attention, setting it to False"
+ )
+ config.rotary_interleaved = False
+
+ if config.apply_rope_fusion:
+ if cu_seqlens is None:
+ return fused_apply_rotary_pos_emb(t, freqs, transpose_output_memory=True)
+ else:
+ return fused_apply_rotary_pos_emb_thd(t, cu_seqlens, freqs)
+ else:
+ if cu_seqlens is None:
+ return _apply_rotary_pos_emb_bshd(
+ t,
+ freqs,
+ rotary_interleaved=config.rotary_interleaved,
+ multi_latent_attention=config.multi_latent_attention,
+ mscale=mscale,
+ )
+ else:
+ return _apply_rotary_pos_emb_thd(
+ t,
+ cu_seqlens,
+ freqs,
+ rotary_interleaved=config.rotary_interleaved,
+ multi_latent_attention=config.multi_latent_attention,
+ mscale=mscale,
+ )
diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py
new file mode 100644
index 0000000000..5232faec60
--- /dev/null
+++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py
@@ -0,0 +1,186 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from megatron.core.transformer.transformer_config import TransformerConfig
+ from megatron.core.transformer.transformer_block import TransformerBlock
+
+import logging
+import math
+
+import torch
+from torch import Tensor, nn
+
+from megatron.core import parallel_state
+from megatron.core.models.common.embeddings.rope_utils import ( # for backward compatibility; pylint: disable=unused-import
+ _apply_rotary_pos_emb_bshd,
+ _apply_rotary_pos_emb_thd,
+ _rotate_half,
+ apply_rotary_pos_emb,
+ get_pos_emb_on_this_cp_rank,
+)
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = ['RotaryEmbedding']
+
+
+class RotaryEmbedding(nn.Module):
+ """Rotary Embedding for language model.
+
+ Args:
+ kv_channels (int): Projection weights dimension in multi-head attention. Obtained
+ from transformer config
+ rotary_percent (float): Percent of rotary dimension to use for rotary position
+ embeddings.
+ rotary_interleaved (bool, optional): If True, interleaved rotary position embeddings.
+ Defaults to False.
+ seq_len_interpolation_factor (float, optional): scale of linearly interpolating RoPE
+ for longer sequences. The value must be a float larger than 1.0. Defaults to None
+ rotary_base (int, optional): Base period for rotary position embeddings. Defaults to
+ 10000.
+ rope_scaling (bool, optional): Apply rope scaling as used in llama 3.1
+ use_cpu_initialization (bool, optional): If False, initialize the inv_freq directly
+ on the GPU. Defaults to False
+ """
+
+ def __init__(
+ self,
+ kv_channels: int,
+ rotary_percent: float,
+ rotary_interleaved: bool = False,
+ seq_len_interpolation_factor: float = None,
+ rotary_base: int = 10000,
+ rope_scaling: bool = False,
+ use_cpu_initialization: bool = False,
+ ) -> None:
+ super().__init__()
+
+ dim = kv_channels
+ if rotary_percent < 1.0:
+ dim = int(dim * rotary_percent)
+ self.rotary_interleaved = rotary_interleaved
+
+ self.seq_len_interpolation_factor = seq_len_interpolation_factor
+ device = 'cpu' if use_cpu_initialization else torch.cuda.current_device()
+ self.inv_freq = 1.0 / (
+ rotary_base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
+ )
+
+ if rope_scaling:
+ self.inv_freq = self._apply_scaling(self.inv_freq)
+
+ def _apply_scaling(
+ self,
+ freqs,
+ factor=8,
+ low_freq_factor=1,
+ high_freq_factor=4,
+ original_max_position_embeddings=8192,
+ ):
+ # This implementation is adapted from:
+ # https://github.com/huggingface/transformers/blob/2a5a6ad18aa22e98429bb5ecb880660328030ea0/src/transformers/modeling_rope_utils.py#L303-L343
+
+ factor = factor # `8` in the original implementation
+ low_freq_factor = low_freq_factor # `1` in the original implementation
+ high_freq_factor = high_freq_factor # `4` in the original implementation
+ old_context_len = original_max_position_embeddings # `8192` in the original implementation
+
+ low_freq_wavelen = old_context_len / low_freq_factor
+ high_freq_wavelen = old_context_len / high_freq_factor
+
+ wavelen = 2 * math.pi / freqs
+ # wavelen < high_freq_wavelen: do nothing
+ # wavelen > low_freq_wavelen: divide by factor
+ inv_freq_llama = torch.where(wavelen > low_freq_wavelen, freqs / factor, freqs)
+ # otherwise: interpolate between the two, using a smooth factor
+ smooth_factor = (old_context_len / wavelen - low_freq_factor) / (
+ high_freq_factor - low_freq_factor
+ )
+ smoothed_inv_freq = (
+ 1 - smooth_factor
+ ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
+ is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
+ inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
+
+ return inv_freq_llama
+
+ def forward(self, max_seq_len: int, offset: int = 0) -> Tensor:
+ """Forward pass of RoPE embedding.
+
+ Args:
+ max_seq_len (int): Maximum size of sequence
+ offset (int, optional): _description_. Defaults to 0.
+
+ Returns:
+ Tensor: Embeddings after applying RoPE.
+ """
+ if self.inv_freq.device.type == 'cpu':
+ # move `inv_freq` to GPU once at the first micro-batch forward pass
+ self.inv_freq = self.inv_freq.to(device=torch.cuda.current_device())
+ seq = (
+ torch.arange(max_seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+ + offset
+ )
+
+ if self.seq_len_interpolation_factor is not None:
+ seq *= 1 / self.seq_len_interpolation_factor
+
+ freqs = torch.outer(seq, self.inv_freq)
+ # first part even vector components, second part odd vector components,
+ # 2 * dim in dimension size
+ if not self.rotary_interleaved:
+ emb = torch.cat((freqs, freqs), dim=-1)
+ else:
+ emb = torch.stack((freqs.view(-1, 1), freqs.view(-1, 1)), dim=-1).view(
+ freqs.shape[0], -1
+ )
+ # emb [seq_length, .., dim]
+ emb = emb[:, None, None, :]
+ if parallel_state.get_context_parallel_world_size() > 1:
+ # slice rotary_pos_emb along sequence dimension and select the parition of the current
+ # CP rank
+ emb = get_pos_emb_on_this_cp_rank(emb, 0)
+ return emb
+
+ def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+ state_dict.pop(f'{prefix}inv_freq', None)
+ return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+ def get_rotary_seq_len(
+ self,
+ inference_params,
+ transformer: TransformerBlock,
+ transformer_input: Tensor,
+ transformer_config: TransformerConfig,
+ ) -> float:
+ """Function to get the rotary sequence length.
+
+ Args:
+ inference_params : Used during Inference time
+ transformer (TransformerBlock): The transformer block (decoder/encoder) used
+ by the model
+ transformer_input (Tensor): Input tensor to the transformer
+ transformer_config (TransformerConfig): Transformer config used by the model
+
+ Returns:
+ float: The rotary sequence length
+ """
+ if inference_params is not None:
+ rotary_seq_len = inference_params.max_sequence_length
+ else:
+ if transformer.input_tensor is not None:
+ rotary_seq_len = transformer.input_tensor.size(0)
+ else:
+ rotary_seq_len = transformer_input.size(0)
+
+ if transformer_config.sequence_parallel:
+ rotary_seq_len *= transformer_config.tensor_model_parallel_size
+
+ rotary_seq_len *= transformer_config.context_parallel_size
+
+ return rotary_seq_len
diff --git a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py
new file mode 100644
index 0000000000..14d147ea34
--- /dev/null
+++ b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py
@@ -0,0 +1,169 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+import math
+
+import torch
+from torch import Tensor
+
+from megatron.core import parallel_state
+from megatron.core.models.common.embeddings.rope_utils import get_pos_emb_on_this_cp_rank
+from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+
+logger = logging.getLogger(__name__)
+
+
+class YarnRotaryEmbedding(RotaryEmbedding):
+ """Yarn Rotary Embedding for language model.
+
+ Args:
+ kv_channels (int): Projection weights dimension in multi-head attention. Obtained from
+ transformer config
+ rotary_percent (float): Percent of rotary dimension to use for rotary position embeddings.
+ rotary_interleaved (bool, optional): If True, interleaved rotary position embeddings.
+ Defaults to False.
+ seq_len_interpolation_factor (float, optional): scale of linearly interpolating RoPE for
+ longer sequences. The value must be a float larger than 1.0. Defaults to None
+ rotary_base (float, optional): Base period for rotary position embeddings. Defaults to
+ 10000.
+ use_cpu_initialization (bool, optional): If False, initialize the inv_freq directly on
+ the GPU. Defaults to False
+ scaling_factor (float, optional): Scaling factor for Yarn RoPE. Defaults to 1.0.
+ original_max_position_embeddings (int, optional): Original maximum position embeddings
+ length. Defaults to 4096.
+ beta_fast (float, optional): Fast beta value for Yarn RoPE. Defaults to 32.
+ beta_slow (float, optional): Slow beta value for Yarn RoPE. Defaults to 1.
+ mscale (float, optional): Mscale value for Yarn RoPE. Defaults to 1.
+ mscale_all_dim (float, optional): Mscale all dim value for Yarn RoPE. Defaults to 0.
+ """
+
+ def __init__(
+ self,
+ kv_channels: int,
+ rotary_percent: float = 1.0,
+ rotary_interleaved: bool = False,
+ seq_len_interpolation_factor: float = None,
+ rotary_base: float = 10000.0,
+ use_cpu_initialization: bool = False,
+ scaling_factor: float = 1.0,
+ original_max_position_embeddings: int = 4096,
+ beta_fast: float = 32.0,
+ beta_slow: float = 1.0,
+ mscale: float = 1.0,
+ mscale_all_dim: float = 0.0,
+ ):
+ self.dim = kv_channels
+ self.rotary_base = rotary_base
+ self.scaling_factor = scaling_factor
+ self.original_max_position_embeddings = original_max_position_embeddings
+ self.beta_fast = beta_fast
+ self.beta_slow = beta_slow
+ self.mscale = mscale
+ self.mscale_all_dim = mscale_all_dim
+
+ device = 'cpu' if use_cpu_initialization else torch.cuda.current_device()
+ self.inv_freq_extra = 1.0 / (
+ self.rotary_base
+ ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
+ )
+ self.inv_freq_inter = 1.0 / (
+ self.scaling_factor
+ * self.rotary_base
+ ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
+ )
+ super().__init__(
+ kv_channels,
+ rotary_percent,
+ rotary_interleaved,
+ seq_len_interpolation_factor,
+ rotary_base,
+ use_cpu_initialization,
+ )
+
+ def forward(self, max_seq_len: int, offset: int = 0) -> Tensor:
+
+ assert (
+ not self.rotary_interleaved
+ ), "Yarn RoPE does not support interleaved rotary embeddings"
+
+ if self.inv_freq_extra.device.type == 'cpu':
+ # move `inv_freq_extra` to GPU once at the first micro-batch forward pass
+ self.inv_freq_extra = self.inv_freq_extra.to(device=torch.cuda.current_device())
+
+ if self.inv_freq_inter.device.type == 'cpu':
+ # move `inv_freq_inter` to GPU once at the first micro-batch forward pass
+ self.inv_freq_inter = self.inv_freq_inter.to(device=torch.cuda.current_device())
+
+ low, high = _yarn_find_correction_range(
+ self.beta_fast,
+ self.beta_slow,
+ self.dim,
+ self.rotary_base,
+ self.original_max_position_embeddings,
+ )
+ inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, self.dim // 2).to(
+ device=self.inv_freq_extra.device, dtype=torch.float32
+ )
+ inv_freq = self.inv_freq_inter * (1 - inv_freq_mask) + self.inv_freq_extra * inv_freq_mask
+
+ seq = (
+ torch.arange(
+ max_seq_len, device=self.inv_freq_extra.device, dtype=self.inv_freq_extra.dtype
+ )
+ + offset
+ )
+
+ freqs = torch.outer(seq, inv_freq)
+
+ _mscale = float(
+ _yarn_get_mscale(self.scaling_factor, self.mscale)
+ / _yarn_get_mscale(self.scaling_factor, self.mscale_all_dim)
+ )
+
+ emb = torch.cat((freqs, freqs), dim=-1)
+ # emb [seq_length, .., dim]
+ emb = emb[:, None, None, :]
+ if parallel_state.get_context_parallel_world_size() > 1:
+ # slice rotary_pos_emb along sequence dimension
+ # and select the parition of the current CP rank
+ emb = get_pos_emb_on_this_cp_rank(emb, 0)
+ return emb, _mscale
+
+
+# Inverse dim formula to find dim based on number of rotations
+def _yarn_find_correction_dim(
+ num_rotations: float, dim: int, rotary_base: float = 10000, max_position_embeddings: int = 2048
+) -> float:
+ return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
+ 2 * math.log(rotary_base)
+ )
+
+
+# Find dim range bounds based on rotations
+def _yarn_find_correction_range(
+ low_rot: float,
+ high_rot: float,
+ dim: int,
+ rotary_base: float = 10000,
+ max_position_embeddings: int = 2048,
+) -> tuple[int, int]:
+ low = math.floor(_yarn_find_correction_dim(low_rot, dim, rotary_base, max_position_embeddings))
+ high = math.ceil(_yarn_find_correction_dim(high_rot, dim, rotary_base, max_position_embeddings))
+ return max(low, 0), min(high, dim - 1) # Clamp values just in case
+
+
+def _yarn_linear_ramp_mask(min: float, max: float, dim: int) -> Tensor:
+ if min == max:
+ max += 0.001 # Prevent singularity
+
+ linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
+ ramp_func = torch.clamp(linear_func, 0, 1)
+ return ramp_func
+
+
+def _yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
+ if scale <= 1:
+ return 1.0
+ return 0.1 * mscale * math.log(scale) + 1.0
diff --git a/megatron/core/models/common/language_module/__init__.py b/megatron/core/models/common/language_module/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
new file mode 100644
index 0000000000..7075e57f98
--- /dev/null
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -0,0 +1,204 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import logging
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.fusions.fused_cross_entropy import fused_vocab_parallel_cross_entropy
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import make_tp_sharded_tensor_for_checkpoint
+
+
+class LanguageModule(MegatronModule):
+ """Base language module that has common helper functions used across GPT, BERT etc.
+
+ Args:
+ config (TransformerConfig): Input transformer config for the model
+ """
+
+ def __init__(self, config: TransformerConfig) -> None:
+ super().__init__(config=config)
+
+ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor:
+ """Computes the language model loss (Cross entropy across vocabulary)
+
+ Args:
+ labels (Tensor): The labels of dimension [batch size, seq length]
+ logits (Tensor): The final logits returned by the output layer of the transformer model
+
+ Returns:
+ Tensor: Loss tensor of dimensions [batch size, sequence_length]
+ """
+ # [b s] => [s b]
+ labels = labels.transpose(0, 1).contiguous()
+ if self.config.cross_entropy_loss_fusion:
+ loss = fused_vocab_parallel_cross_entropy(logits, labels)
+ else:
+ loss = tensor_parallel.vocab_parallel_cross_entropy(logits, labels)
+
+ # [s b] => [b, s]
+ loss = loss.transpose(0, 1).contiguous()
+ return loss
+
+ def setup_embeddings_and_output_layer(self) -> None:
+ """Sets up embedding layer in first stage and output layer in last stage.
+
+ This function initalizes word embeddings in the final stage when we are
+ using pipeline parallelism and sharing word embeddings, and sets up param
+ attributes on the embedding and output layers.
+ """
+
+ # Set `is_embedding_or_output_parameter` attribute.
+ if self.pre_process:
+ self.embedding.word_embeddings.weight.is_embedding_or_output_parameter = True
+ if self.post_process and self.output_layer.weight is not None:
+ self.output_layer.weight.is_embedding_or_output_parameter = True
+
+ if not self.share_embeddings_and_output_weights:
+ return
+
+ if parallel_state.get_pipeline_model_parallel_world_size() == 1:
+ # Zero out wgrad if sharing embeddings between two layers on same
+ # pipeline stage to make sure grad accumulation into main_grad is
+ # correct and does not include garbage values (e.g., from torch.empty).
+ self.shared_embedding_or_output_weight().zero_out_wgrad = True
+ return
+
+ if parallel_state.is_pipeline_first_stage() and self.pre_process and not self.post_process:
+ self.shared_embedding_or_output_weight().shared_embedding = True
+
+ if self.post_process and not self.pre_process:
+ assert not parallel_state.is_pipeline_first_stage()
+ # set word_embeddings weights to 0 here, then copy first
+ # stage's weights using all_reduce below.
+ self.output_layer.weight.data.fill_(0)
+ self.output_layer.weight.shared = True
+ self.output_layer.weight.shared_embedding = True
+
+ # Parameters are shared between the word embeddings layers, and the
+ # heads at the end of the model. In a pipelined setup with more than
+ # one stage, the initial embedding layer and the head are on different
+ # workers, so we do the following:
+ # 1. Create a second copy of word_embeddings on the last stage, with
+ # initial parameters of 0.0.
+ # 2. Do an all-reduce between the first and last stage to ensure that
+ # the two copies of word_embeddings start off with the same
+ # parameter values.
+ # 3. In the training loop, before an all-reduce between the grads of
+ # the two word_embeddings layers to ensure that every applied weight
+ # update is the same on both stages.
+
+ # Ensure that first and last stages have the same initial parameter
+ # values.
+ if torch.distributed.is_initialized():
+ if parallel_state.is_rank_in_embedding_group():
+ weight = self.shared_embedding_or_output_weight()
+ weight.data = weight.data.cuda()
+ torch.distributed.all_reduce(
+ weight.data, group=parallel_state.get_embedding_group()
+ )
+
+ elif not getattr(LanguageModule, "embedding_warning_printed", False):
+ logging.getLogger(__name__).warning(
+ "Distributed processes aren't initialized, so the output layer "
+ "is not initialized with weights from the word embeddings. "
+ "If you are just manipulating a model this is fine, but "
+ "this needs to be handled manually. If you are training "
+ "something is definitely wrong."
+ )
+ LanguageModule.embedding_warning_printed = True
+
+ def shared_embedding_or_output_weight(self) -> Tensor:
+ """Gets the emedding weight or output logit weights when share embedding and output weights set to True.
+
+ Returns:
+ Tensor: During pre processing it returns the input embeddings weight while during post processing it returns the final output layers weight
+ """
+ if self.pre_process:
+ return self.embedding.word_embeddings.weight
+ elif self.post_process:
+ return self.output_layer.weight
+ return None
+
+ def sharded_state_dict(
+ self,
+ prefix: str = '',
+ sharded_offsets: Tuple[Tuple[int, int, int]] = (),
+ metadata: Optional[dict] = None,
+ ) -> ShardedStateDict:
+ """Sharded state dict implementation that handles the output layer weights tying.
+
+ Args:
+ prefix (str): Module name prefix.
+ sharded_offsets (tuple): PP related offsets, expected to be empty at this module level.
+ metadata (Optional[Dict]): metadata controlling sharded state dict creation.
+
+ Returns:
+ ShardedStateDict: sharded state dict for the LanguageModel
+ """
+ assert not sharded_offsets, "Unexpected sharded offsets"
+ sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+
+ first_stage_word_emb_key = f'{prefix}embedding.word_embeddings.weight'
+ output_layer_weight_key = f'{prefix}output_layer.weight'
+ output_layer_bias_key = f'{prefix}output_layer.bias'
+
+ if self.share_embeddings_and_output_weights:
+ self.tie_embeddings_and_output_weights_state_dict(
+ sharded_state_dict, output_layer_weight_key, first_stage_word_emb_key
+ )
+ elif self.post_process:
+ # Make sure the output layer follows the embeddings padding logic
+ sharded_state_dict[output_layer_weight_key].allow_shape_mismatch = True
+
+ # Regardless of sharing the output weights with embeddings, we must handle the bias padding
+ if self.post_process and output_layer_bias_key in sharded_state_dict:
+ sharded_state_dict[output_layer_bias_key].allow_shape_mismatch = True
+
+ return sharded_state_dict
+
+ def tie_embeddings_and_output_weights_state_dict(
+ self,
+ sharded_state_dict: ShardedStateDict,
+ output_layer_weight_key: str,
+ first_stage_word_emb_key: str,
+ ) -> None:
+ """Ties the embedding and output weights in a given sharded state dict.
+
+ Args:
+ sharded_state_dict (ShardedStateDict): state dict with the weight to tie
+ output_layer_weight_key (str): key of the output layer weight in the state dict.
+ This entry will be replaced with a tied version
+ first_stage_word_emb_key (str): this must be the same as the
+ ShardedTensor.key of the first stage word embeddings.
+
+ Returns: None, acts in-place
+ """
+ if not self.post_process:
+ # No output layer
+ assert output_layer_weight_key not in sharded_state_dict, sharded_state_dict.keys()
+ return
+
+ if self.pre_process:
+ # Output layer is equivalent to the embedding already
+ return
+
+ # Replace the default output layer with a one sharing the weights with the embedding
+ del sharded_state_dict[output_layer_weight_key]
+ tensor = self.shared_embedding_or_output_weight()
+ last_stage_word_emb_replica_id = (
+ 1, # copy of first stage embedding
+ 0,
+ parallel_state.get_data_parallel_rank(with_context_parallel=True),
+ )
+
+ sharded_state_dict[output_layer_weight_key] = make_tp_sharded_tensor_for_checkpoint(
+ tensor=tensor,
+ key=first_stage_word_emb_key,
+ replica_id=last_stage_word_emb_replica_id,
+ allow_shape_mismatch=True,
+ )
diff --git a/megatron/core/models/common/vision_module/__init__.py b/megatron/core/models/common/vision_module/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/megatron/core/models/common/vision_module/vision_module.py b/megatron/core/models/common/vision_module/vision_module.py
new file mode 100644
index 0000000000..5dc51873a4
--- /dev/null
+++ b/megatron/core/models/common/vision_module/vision_module.py
@@ -0,0 +1,17 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+"""Megatron Vision Module."""
+
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+# Note: This is only a stub at the moment. This will be expanded in follow-up changes.
+class VisionModule(MegatronModule):
+ """Base vision module that has common helper functions used across CLIP, ViT, etc.
+
+ Args:
+ config (TransformerConfig): Input transformer config for the model
+ """
+
+ def __init__(self, config: TransformerConfig) -> None:
+ super().__init__(config=config)
diff --git a/megatron/core/models/gpt/__init__.py b/megatron/core/models/gpt/__init__.py
new file mode 100644
index 0000000000..8bbecfcb09
--- /dev/null
+++ b/megatron/core/models/gpt/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from .gpt_model import GPTModel
diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
new file mode 100755
index 0000000000..1db68dc886
--- /dev/null
+++ b/megatron/core/models/gpt/gpt_layer_specs.py
@@ -0,0 +1,247 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Optional
+
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
+from megatron.core.transformer.dot_product_attention import DotProductAttention
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules
+from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
+from megatron.core.transformer.multi_latent_attention import (
+ MLASelfAttention,
+ MLASelfAttentionSubmodules,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+
+try:
+ from megatron.core.extensions.transformer_engine import (
+ TEColumnParallelGroupedLinear,
+ TEColumnParallelLinear,
+ TEDotProductAttention,
+ TELayerNormColumnParallelLinear,
+ TENorm,
+ TERowParallelGroupedLinear,
+ TERowParallelLinear,
+ )
+
+ HAVE_TE = True
+except ImportError:
+ HAVE_TE = False
+
+try:
+ import apex # pylint: disable=unused-import
+
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+
+ HAVE_APEX = True
+ LNImpl = FusedLayerNorm
+except ImportError:
+ import warnings
+
+ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm
+
+ warnings.warn('Apex is not installed. Falling back to Torch LayerNorm')
+ LNImpl = WrappedTorchLayerNorm
+
+
+def get_gpt_layer_with_transformer_engine_spec(
+ num_experts: Optional[int] = None,
+ moe_grouped_gemm: Optional[bool] = False,
+ qk_layernorm: Optional[bool] = False,
+ multi_latent_attention: Optional[bool] = False,
+ fp8: Optional[str] = None,
+) -> ModuleSpec:
+ """Use this spec to use lower-level Transformer Engine modules (required for fp8 training).
+
+
+ Args:
+ num_experts (int, optional): Number of experts. Defaults to None.
+ moe_grouped_gemm (bool, optional): To use Grouped GEMM. Defaults to False.
+ qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False.
+ fp8 (str, optional): Flag to decide the linear layer spec for MoE. Defaults to None.
+
+ Returns:
+ ModuleSpec: Module specification with TE modules
+ """
+ mlp = _get_mlp_module_spec(
+ use_te=True, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm, fp8=fp8
+ )
+
+ if multi_latent_attention:
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=MLASelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=MLASelfAttentionSubmodules(
+ linear_q_proj=TEColumnParallelLinear,
+ linear_q_down_proj=TEColumnParallelLinear,
+ linear_q_up_proj=TEColumnParallelLinear,
+ linear_kv_down_proj=TEColumnParallelLinear,
+ linear_kv_up_proj=TEColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ q_layernorm=TENorm if qk_layernorm else IdentityOp,
+ kv_layernorm=TENorm if qk_layernorm else IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=TENorm if num_experts else IdentityOp,
+ input_layernorm=TENorm if num_experts else IdentityOp,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+ else:
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=TELayerNormColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ # TENorm significantly harms convergence when used
+ # for QKLayerNorm; we instead use the Apex implementation.
+ q_layernorm=FusedLayerNorm if qk_layernorm else IdentityOp,
+ k_layernorm=FusedLayerNorm if qk_layernorm else IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=TENorm if num_experts else IdentityOp,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+
+
+def get_gpt_layer_local_spec(
+ num_experts: Optional[int] = None,
+ moe_grouped_gemm: Optional[bool] = False,
+ qk_layernorm: Optional[bool] = False,
+ multi_latent_attention: Optional[bool] = False,
+) -> ModuleSpec:
+ """Use this spec for an implementation using only modules in Megatron-Core.
+
+
+ Args:
+ num_experts (int, optional): Number of experts. Defaults to None.
+ moe_grouped_gemm (bool, optional): To use Grouped GEMM. Defaults to False.
+ qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False.
+
+ Returns:
+ ModuleSpec: Module specification with Megatron-Core modules
+ """
+ mlp = _get_mlp_module_spec(
+ use_te=False, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm
+ )
+ if multi_latent_attention:
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=MLASelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=MLASelfAttentionSubmodules(
+ linear_q_proj=ColumnParallelLinear,
+ linear_q_down_proj=ColumnParallelLinear,
+ linear_q_up_proj=ColumnParallelLinear,
+ linear_kv_down_proj=ColumnParallelLinear,
+ linear_kv_up_proj=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ q_layernorm=LNImpl if qk_layernorm else IdentityOp,
+ kv_layernorm=LNImpl if qk_layernorm else IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=LNImpl if num_experts else IdentityOp,
+ input_layernorm=LNImpl if num_experts else IdentityOp,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+ else:
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=LNImpl,
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ q_layernorm=LNImpl if qk_layernorm else IdentityOp,
+ k_layernorm=LNImpl if qk_layernorm else IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=LNImpl,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ sharded_state_dict_keys_map={
+ 'input_layernorm.': 'self_attention.linear_qkv.layer_norm_',
+ 'pre_mlp_layernorm.': 'mlp.linear_fc1.layer_norm_',
+ },
+ ),
+ )
+
+
+def _get_mlp_module_spec(
+ use_te: Optional[bool] = True,
+ num_experts: Optional[int] = None,
+ moe_grouped_gemm: Optional[bool] = False,
+ fp8: Optional[str] = None,
+) -> ModuleSpec:
+ """Helper function to get module spec for MLP/MoE"""
+ if num_experts is None:
+ # Dense MLP w/ or w/o TE modules.
+ return ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=TELayerNormColumnParallelLinear if use_te else ColumnParallelLinear,
+ linear_fc2=TERowParallelLinear if use_te else RowParallelLinear,
+ ),
+ )
+ else:
+ # Mixture of experts with modules in megatron core.
+ if use_te and moe_grouped_gemm:
+ linear_fc1 = TEColumnParallelGroupedLinear
+ linear_fc2 = TERowParallelGroupedLinear
+ elif use_te and fp8:
+ linear_fc1 = TEColumnParallelLinear
+ linear_fc2 = TERowParallelLinear
+ else:
+ linear_fc1 = ColumnParallelLinear
+ linear_fc2 = RowParallelLinear
+
+ use_te_grouped_gemm = use_te and TEColumnParallelGroupedLinear is not None
+
+ return ModuleSpec(
+ module=MoELayer,
+ submodules=MoESubmodules(
+ experts=(
+ MLPSubmodules(linear_fc1=linear_fc1, linear_fc2=linear_fc2)
+ if not moe_grouped_gemm or use_te_grouped_gemm
+ else None
+ ),
+ shared_experts=ModuleSpec(
+ module=SharedExpertMLP,
+ params={"gate": False},
+ submodules=MLPSubmodules(
+ linear_fc1=TEColumnParallelLinear if use_te else ColumnParallelLinear,
+ linear_fc2=TERowParallelLinear if use_te else RowParallelLinear,
+ ),
+ ),
+ ),
+ )
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
new file mode 100644
index 0000000000..bd52f89680
--- /dev/null
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -0,0 +1,290 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from collections import OrderedDict
+from typing import Dict, Literal, Optional
+
+from torch import Tensor
+
+from megatron.core import InferenceParams, tensor_parallel
+from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
+from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.enums import ModelType
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class GPTModel(LanguageModule):
+ """GPT Transformer language model.
+
+ Args:
+ config (TransformerConfig):
+ Transformer config
+ transformer_layer_spec (ModuleSpec):
+ Specifies module to use for transformer layers
+ vocab_size (int):
+ Vocabulary size
+ max_sequence_length (int):
+ maximum size of sequence. This is used for positional embedding
+ pre_process (bool, optional):
+ Include embedding layer (used with pipeline parallelism). Defaults to True.
+ post_process (bool, optional):
+ Include an output layer (used with pipeline parallelism). Defaults to True.
+ fp16_lm_cross_entropy (bool, optional):
+ Defaults to False.
+ parallel_output (bool, optional):
+ Do not gather the outputs, keep them split across tensor
+ parallel ranks. Defaults to True.
+ share_embeddings_and_output_weights (bool, optional):
+ When True, input embeddings and output logit weights are shared. Defaults to False.
+ position_embedding_type (Literal[learned_absolute,rope], optional):
+ Position embedding type.. Defaults to 'learned_absolute'.
+ rotary_percent (float, optional):
+ Percent of rotary dimension to use for rotary position embeddings.
+ Ignored unless position_embedding_type is 'rope'. Defaults to 1.0.
+ rotary_base (int, optional):
+ Base period for rotary position embeddings. Ignored unless
+ position_embedding_type is 'rope'.
+ Defaults to 10000.
+ seq_len_interpolation_factor (Optional[float], optional):
+ scale of linearly interpolating RoPE for longer sequences.
+ The value must be a float larger than 1.0. Defaults to None.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ transformer_layer_spec: ModuleSpec,
+ vocab_size: int,
+ max_sequence_length: int,
+ pre_process: bool = True,
+ post_process: bool = True,
+ fp16_lm_cross_entropy: bool = False,
+ parallel_output: bool = True,
+ share_embeddings_and_output_weights: bool = False,
+ position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'learned_absolute',
+ rotary_percent: float = 1.0,
+ rotary_base: int = 10000,
+ rope_scaling: bool = False,
+ seq_len_interpolation_factor: Optional[float] = None,
+ ) -> None:
+ super().__init__(config=config)
+
+ if has_config_logger_enabled(config):
+ log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+ self.transformer_layer_spec: ModuleSpec = transformer_layer_spec
+ self.vocab_size = vocab_size
+ self.max_sequence_length = max_sequence_length
+ self.pre_process = pre_process
+ self.post_process = post_process
+ self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
+ self.parallel_output = parallel_output
+ self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+ self.position_embedding_type = position_embedding_type
+
+ # megatron core pipelining currently depends on model type
+ # TODO: remove this dependency ?
+ self.model_type = ModelType.encoder_or_decoder
+
+ # These 4 attributes are needed for TensorRT-LLM export.
+ self.max_position_embeddings = max_sequence_length
+ self.rotary_percent = rotary_percent
+ self.rotary_base = rotary_base
+ self.rotary_scaling = rope_scaling
+
+ if self.pre_process:
+ self.embedding = LanguageModelEmbedding(
+ config=self.config,
+ vocab_size=self.vocab_size,
+ max_sequence_length=self.max_sequence_length,
+ position_embedding_type=position_embedding_type,
+ )
+
+ if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention:
+ self.rotary_pos_emb = RotaryEmbedding(
+ kv_channels=self.config.kv_channels,
+ rotary_percent=rotary_percent,
+ rotary_interleaved=self.config.rotary_interleaved,
+ seq_len_interpolation_factor=seq_len_interpolation_factor,
+ rotary_base=rotary_base,
+ rope_scaling=rope_scaling,
+ use_cpu_initialization=self.config.use_cpu_initialization,
+ )
+
+ # Transformer.
+ self.decoder = TransformerBlock(
+ config=self.config,
+ spec=transformer_layer_spec,
+ pre_process=self.pre_process,
+ post_process=self.post_process,
+ )
+
+ # Output
+ if post_process:
+ if self.config.defer_embedding_wgrad_compute:
+ # The embedding activation buffer preserves a reference to the input activations
+ # of the final embedding projection layer GEMM. It will hold the activations for
+ # all the micro-batches of a global batch for the last pipeline stage. Once we are
+ # done with all the back props for all the microbatches for the last pipeline stage,
+ # it will be in the pipeline flush stage. During this pipeline flush we use the
+ # input activations stored in embedding activation buffer and gradient outputs
+ # stored in gradient buffer to calculate the weight gradients for the embedding
+ # final linear layer.
+ self.embedding_activation_buffer = []
+ self.grad_output_buffer = []
+ else:
+ self.embedding_activation_buffer = None
+ self.grad_output_buffer = None
+
+ self.output_layer = tensor_parallel.ColumnParallelLinear(
+ config.hidden_size,
+ self.vocab_size,
+ config=config,
+ init_method=config.init_method,
+ bias=False,
+ skip_bias_add=False,
+ gather_output=not self.parallel_output,
+ skip_weight_param_allocation=self.pre_process
+ and self.share_embeddings_and_output_weights,
+ embedding_activation_buffer=self.embedding_activation_buffer,
+ grad_output_buffer=self.grad_output_buffer,
+ )
+
+ if self.pre_process or self.post_process:
+ self.setup_embeddings_and_output_layer()
+
+ if has_config_logger_enabled(self.config):
+ log_config_to_disk(
+ self.config, self.state_dict(), prefix=f'{type(self).__name__}_init_ckpt'
+ )
+
+ def set_input_tensor(self, input_tensor: Tensor) -> None:
+ """Sets input tensor to the model.
+
+ See megatron.model.transformer.set_input_tensor()
+
+ Args:
+ input_tensor (Tensor): Sets the input tensor for the model.
+ """
+ # This is usually handled in schedules.py but some inference code still
+ # gives us non-lists or None
+ if not isinstance(input_tensor, list):
+ input_tensor = [input_tensor]
+
+ assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert'
+ self.decoder.set_input_tensor(input_tensor[0])
+
+ def forward(
+ self,
+ input_ids: Tensor,
+ position_ids: Tensor,
+ attention_mask: Tensor,
+ decoder_input: Tensor = None,
+ labels: Tensor = None,
+ inference_params: InferenceParams = None,
+ packed_seq_params: PackedSeqParams = None,
+ extra_block_kwargs: dict = None,
+ runtime_gather_output: Optional[bool] = None,
+ ) -> Tensor:
+ """Forward function of the GPT Model This function passes the input tensors
+ through the embedding layer, and then the decoeder and finally into the post
+ processing layer (optional).
+
+ It either returns the Loss values if labels are given or the final hidden units
+
+ Args:
+ runtime_gather_output (bool): Gather output at runtime. Default None means
+ `parallel_output` arg in the constructor will be used.
+ """
+ # If decoder_input is provided (not None), then input_ids and position_ids are ignored.
+ # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
+
+ # Decoder embedding.
+ if decoder_input is not None:
+ pass
+ elif self.pre_process:
+ decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+ else:
+ # intermediate stage of pipeline
+ # decoder will get hidden_states from encoder.input_tensor
+ decoder_input = None
+
+ # Rotary positional embeddings (embedding is None for PP intermediate devices)
+ rotary_pos_emb = None
+ if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention:
+ rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+ inference_params, self.decoder, decoder_input, self.config
+ )
+ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)
+
+ # Run decoder.
+ hidden_states = self.decoder(
+ hidden_states=decoder_input,
+ attention_mask=attention_mask,
+ inference_params=inference_params,
+ rotary_pos_emb=rotary_pos_emb,
+ packed_seq_params=packed_seq_params,
+ **(extra_block_kwargs or {}),
+ )
+
+ if not self.post_process:
+ return hidden_states
+
+ # logits and loss
+ output_weight = None
+ if self.share_embeddings_and_output_weights:
+ output_weight = self.shared_embedding_or_output_weight()
+ logits, _ = self.output_layer(
+ hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output
+ )
+
+ if has_config_logger_enabled(self.config):
+ payload = OrderedDict(
+ {
+ 'input_ids': input_ids,
+ 'position_ids': position_ids,
+ 'attention_mask': attention_mask,
+ 'decoder_input': decoder_input,
+ 'logits': logits,
+ }
+ )
+ log_config_to_disk(self.config, payload, prefix='input_and_logits')
+
+ if labels is None:
+ # [s b h] => [b s h]
+ return logits.transpose(0, 1).contiguous()
+
+ loss = self.compute_language_model_loss(labels, logits)
+
+ return loss
+
+ def sharded_state_dict(
+ self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[Dict] = None
+ ) -> ShardedStateDict:
+ """Sharded state dict implementation for GPTModel backward-compatibility
+ (removing extra state).
+
+ Args:
+ prefix (str): Module name prefix.
+ sharded_offsets (tuple): PP related offsets, expected to be empty at this module level.
+ metadata (Optional[Dict]): metadata controlling sharded state dict creation.
+
+ Returns:
+ ShardedStateDict: sharded state dict for the GPTModel
+ """
+ sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+ output_layer_extra_state_key = f'{prefix}output_layer._extra_state'
+
+ # Old GPT checkpoints only stored the output layer weight key. So we remove the
+ # _extra_state key but check that it doesn't contain any data anyway
+ output_extra_state = sharded_state_dict.pop(output_layer_extra_state_key, None)
+ assert not (
+ output_extra_state and output_extra_state.data
+ ), f'Expected output layer extra state to be empty, got: {output_extra_state}'
+
+ return sharded_state_dict
diff --git a/megatron/core/models/mamba/__init__.py b/megatron/core/models/mamba/__init__.py
new file mode 100644
index 0000000000..5aaf852401
--- /dev/null
+++ b/megatron/core/models/mamba/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from .mamba_model import MambaModel
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
new file mode 100755
index 0000000000..e5fa9efa72
--- /dev/null
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -0,0 +1,67 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.extensions.transformer_engine import (
+ TEDotProductAttention,
+ TELayerNormColumnParallelLinear,
+ TERowParallelLinear,
+)
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules
+from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules
+from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules
+from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+
+mamba_stack_spec = ModuleSpec(
+ module=MambaStack,
+ submodules=MambaStackSubmodules(
+ mamba_layer=ModuleSpec(
+ module=MambaLayer,
+ submodules=MambaLayerSubmodules(
+ mixer=ModuleSpec(
+ module=MambaMixer,
+ submodules=MambaMixerSubmodules(
+ in_proj=TELayerNormColumnParallelLinear, out_proj=TERowParallelLinear
+ ),
+ ),
+ mamba_bda=get_bias_dropout_add,
+ ),
+ ),
+ # Started with spec from gpt_layer_specs.py (with MLP removed)
+ # Using the TE spec because we had problems getting the non-TE spec
+ # working
+ attention_layer=ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=TELayerNormColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ ),
+ ),
+ # Started with spec from gpt_layer_specs.py
+ # Using the TE spec because we had problems getting the non-TE spec
+ # working
+ mlp_layer=ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ mlp=ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear
+ ),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ ),
+ ),
+ ),
+)
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
new file mode 100644
index 0000000000..5794b1b41a
--- /dev/null
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -0,0 +1,228 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Literal, Optional
+
+from torch import Tensor
+
+from megatron.core import InferenceParams, tensor_parallel
+from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
+from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
+from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.transformer.enums import ModelType
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class MambaModel(LanguageModule):
+ """Mamba language model.
+
+ Args:
+ config (TransformerConfig): Transformer config
+ mamba_stack_spec (ModuleSpec): Specifies the modules to use for the various layer types
+ vocab_size (int): Vocabulary size
+ max_sequence_length (int): maximum size of sequence.
+ This is used for positional embedding
+ pre_process (bool, optional): Include embedding layer
+ (used with pipeline parallelism). Defaults to True.
+ mamba_ssm_ngroups (int, optional): Specifies the number of groups to use.
+ The default value is 8, as in the NVIDIA Mamba2 (pure and hybrid) 8b.
+ However, in the original Mamba2 paper, the checkpoints use a setting of 1.
+ Defaults to 8.
+ hybrid_attention_ratio (float, optional): The target ratio of attention
+ layers to total layers
+ hybrid_mlp_ratio (float, optional): The target ratio of mlp layers to total layers
+ hybrid_override_pattern (str, optional): The hybrid layer pattern to override with
+ post_process (bool, optional): Include an output layer (used with pipeline parallelism).
+ Defaults to True.
+ fp16_lm_cross_entropy (bool, optional): Defaults to False.
+ parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor
+ parallel ranks. Defaults to True.
+ share_embeddings_and_output_weights (bool, optional): When True, input embeddings and
+ output logit weights are shared. Defaults to False.
+ position_embedding_type (Literal[learned_absolute,rope,none], optional): Position
+ embedding type. Defaults to 'none'.
+ rotary_percent (float, optional): Percent of rotary dimension to use for rotary position
+ embeddings. Ignored unless position_embedding_type is 'rope'. Defaults to 1.0.
+ rotary_base (int, optional): Base period for rotary position embeddings. Ignored unless
+ position_embedding_type is 'rope'. Defaults to 10000.
+ seq_len_interpolation_factor (Optional[float], optional): scale of linearly
+ interpolating RoPE for longer sequences. The value must be a float larger than 1.0.
+ Defaults to None.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ mamba_stack_spec: ModuleSpec,
+ vocab_size: int,
+ max_sequence_length: int,
+ mamba_ssm_ngroups: int = 8,
+ pre_process: bool = True,
+ hybrid_attention_ratio: float = 0.0,
+ hybrid_mlp_ratio: float = 0.0,
+ hybrid_override_pattern: str = None,
+ post_process: bool = True,
+ fp16_lm_cross_entropy: bool = False,
+ parallel_output: bool = True,
+ share_embeddings_and_output_weights: bool = False,
+ # Mamba with no attention has no need for position embeddings, so none is default
+ position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'none',
+ rotary_percent: float = 1.0,
+ rotary_base: int = 10000,
+ seq_len_interpolation_factor: Optional[float] = None,
+ ) -> None:
+ super().__init__(config=config)
+
+ if has_config_logger_enabled(config):
+ log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+ self.mamba_stack_spec: ModuleSpec = mamba_stack_spec
+ self.vocab_size = vocab_size
+ self.max_sequence_length = max_sequence_length
+ self.mamba_ssm_ngroups = mamba_ssm_ngroups
+ self.pre_process = pre_process
+ self.hybrid_attention_ratio = hybrid_attention_ratio
+ self.hybrid_mlp_ratio = hybrid_mlp_ratio
+ self.hybrid_override_pattern = hybrid_override_pattern
+ self.post_process = post_process
+ self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
+ self.parallel_output = parallel_output
+ self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+ self.position_embedding_type = position_embedding_type
+
+ # megatron core pipelining currently depends on model type
+ # TODO: remove this dependency ?
+ self.model_type = ModelType.encoder_or_decoder
+
+ if self.pre_process:
+ self.embedding = LanguageModelEmbedding(
+ config=self.config,
+ vocab_size=self.vocab_size,
+ max_sequence_length=self.max_sequence_length,
+ position_embedding_type=position_embedding_type,
+ )
+
+ if self.position_embedding_type == 'rope':
+ self.rotary_pos_emb = RotaryEmbedding(
+ kv_channels=self.config.kv_channels,
+ rotary_percent=rotary_percent,
+ seq_len_interpolation_factor=seq_len_interpolation_factor,
+ rotary_base=rotary_base,
+ use_cpu_initialization=self.config.use_cpu_initialization,
+ )
+
+ self.decoder = build_module(
+ mamba_stack_spec,
+ self.config,
+ mamba_ssm_ngroups=self.mamba_ssm_ngroups,
+ pre_process=self.pre_process,
+ hybrid_attention_ratio=self.hybrid_attention_ratio,
+ hybrid_mlp_ratio=self.hybrid_mlp_ratio,
+ hybrid_override_pattern=self.hybrid_override_pattern,
+ post_process=self.post_process,
+ dtype=config.params_dtype,
+ )
+
+ # Output
+ if post_process:
+ self.output_layer = tensor_parallel.ColumnParallelLinear(
+ config.hidden_size,
+ self.vocab_size,
+ config=config,
+ init_method=config.init_method,
+ bias=False,
+ skip_bias_add=False,
+ gather_output=not self.parallel_output,
+ skip_weight_param_allocation=self.pre_process
+ and self.share_embeddings_and_output_weights,
+ )
+
+ if self.pre_process or self.post_process:
+ self.setup_embeddings_and_output_layer()
+
+ def set_input_tensor(self, input_tensor: Tensor) -> None:
+ """Sets input tensor to the model.
+
+ See megatron.model.transformer.set_input_tensor()
+
+ Args:
+ input_tensor (Tensor): Sets the input tensor for the model.
+ """
+ # This is usually handled in schedules.py but some inference code still
+ # gives us non-lists or None
+ if not isinstance(input_tensor, list):
+ input_tensor = [input_tensor]
+
+ assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert'
+ self.decoder.set_input_tensor(input_tensor[0])
+
+ def forward(
+ self,
+ input_ids: Tensor,
+ position_ids: Tensor,
+ attention_mask: Tensor,
+ decoder_input: Tensor = None,
+ labels: Tensor = None,
+ inference_params: InferenceParams = None,
+ ) -> Tensor:
+ """Forward function of the Mamba model. This function passes the input tensors
+ through the embedding layer, and then the decoder and finally into the post
+ processing layer (optional).
+
+ It either returns the Loss values if labels are given or the final hidden units
+ """
+ # If decoder_input is provided (not None), then input_ids and position_ids are ignored.
+ # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
+
+ # Decoder embedding.
+ if decoder_input is not None:
+ pass
+ elif self.pre_process:
+ decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+ else:
+ # intermediate stage of pipeline
+ # decoder will get hidden_states from encoder.input_tensor
+ decoder_input = None
+
+ rotary_pos_emb = None
+ if self.position_embedding_type == 'rope':
+ rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+ inference_params, self.decoder, decoder_input, self.config
+ )
+ rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)
+
+ # The following assert will currently fail when running inference.
+ # Commented out for now.
+ # TODO (duncan/rwaleffe): (1) confirm that the externally-generated
+ # attention mask is not needed and is ignored by the model in
+ # inference mode, (2) reduce the size of the externally-generated
+ # attention mask to prevent CPU OOM (as we did for training), (3)
+ # force the attention mask passed to the model in inference mode to
+ # be None, so this assert will succeed.
+ # assert attention_mask is None, "The attention mask is ignored and should be set to None"
+
+ # Run decoder.
+ hidden_states = self.decoder(
+ hidden_states=decoder_input,
+ attention_mask=attention_mask,
+ inference_params=inference_params,
+ rotary_pos_emb=rotary_pos_emb,
+ )
+
+ if not self.post_process:
+ return hidden_states
+
+ # logits and loss
+ output_weight = None
+ if self.share_embeddings_and_output_weights:
+ output_weight = self.shared_embedding_or_output_weight()
+ logits, _ = self.output_layer(hidden_states, weight=output_weight)
+
+ if labels is None:
+ # [s b h] => [b s h]
+ return logits.transpose(0, 1).contiguous()
+
+ loss = self.compute_language_model_loss(labels, logits)
+
+ return loss
diff --git a/megatron/core/models/multimodal/__init__.py b/megatron/core/models/multimodal/__init__.py
new file mode 100644
index 0000000000..f8011007a5
--- /dev/null
+++ b/megatron/core/models/multimodal/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/models/multimodal/llava_model.py b/megatron/core/models/multimodal/llava_model.py
new file mode 100644
index 0000000000..29f18ee725
--- /dev/null
+++ b/megatron/core/models/multimodal/llava_model.py
@@ -0,0 +1,660 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import logging
+from collections import namedtuple
+from functools import partial
+from typing import List, Optional
+
+import torch
+
+from megatron.core import InferenceParams, tensor_parallel
+from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
+from megatron.core.extensions.transformer_engine import TEDotProductAttention
+from megatron.core.models.gpt import GPTModel
+from megatron.core.models.vision.clip_vit_model import CLIPViTModel, get_num_image_embeddings
+from megatron.core.models.vision.multimodal_projector import MultimodalProjector
+from megatron.core.parallel_state import get_tensor_model_parallel_world_size
+from megatron.core.transformer import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import is_te_min_version
+
+IMAGE_TOKEN_INDEX = -200 # ID for images in the input sequence.
+IGNORE_INDEX = -100 # ID for labels that should be ignored.
+
+
+# Note: This is under development and may be missing features.
+class LLaVAModel(MegatronModule):
+ """LLaVA multi-modal model.
+
+ Args:
+ language_transformer_config (TransformerConfig): Transformer config for the language model.
+ language_transformer_layer_spec (ModuleSpec): Language model spec.
+ language_vocab_size (int): Language model vocabulary size.
+ language_max_sequence_length (int): Language model maximum sequence length.
+ vision_transformer_config (TransformerConfig): Transformer config for the vision model.
+ vision_transformer_layer_spec (ModuleSpec): Vision model spec.
+ drop_vision_class_token (bool): Drop vision class token(s) before the language model.
+ vision_projection_config (TransformerConfig): Vision projection config.
+ vision_projection_layer_spec (ModuleSpec): Vision projection spec.
+ vision_projection_type (str): Type of the vision projection. Default: 2-layer MLP.
+ allow_missing_vision_projection_checkpoint (bool): Allow vision projection weights to be
+ missing when loading a checkpoint. Default False.
+ parallel_output (bool): Keep outputs split across tensor parallel ranks.
+ This is typically True for training and False for inference.
+ language_position_embedding_type (str): Language model position embedding type.
+ language_rotary_percent (float): RoPE percent. Defaults to 1.0.
+ pre_process (bool): Include embedding layer in the decoder (used with pipeline parallel).
+ post_process (bool): Include output layer in the decoder (used with pipeline parallel).
+ add_encoder (bool): Construct the encoder (used with pipeline parallel).
+ When we use pipelining, the encoder will live on only the first stage
+ add_decoder (bool): Construct the decoder (used with pipeline parallel).
+ When we use pipelining, the decoder will live on every stage after the first one.
+ img_h (int): Input image height.
+ img_w (int): Input image width.
+ patch_dim (int): The size of each image patch side.
+ language_rotary_base (int): RoPE base.
+ """
+
+ def __init__(
+ self,
+ language_transformer_config: TransformerConfig,
+ language_transformer_layer_spec: ModuleSpec,
+ language_vocab_size: int,
+ language_max_sequence_length: int,
+ vision_transformer_config: TransformerConfig,
+ vision_transformer_layer_spec: ModuleSpec,
+ drop_vision_class_token: bool,
+ vision_projection_config: TransformerConfig,
+ vision_projection_layer_spec: ModuleSpec,
+ vision_projection_type: str = "mlp",
+ allow_missing_vision_projection_checkpoint: bool = False,
+ parallel_output: bool = True,
+ language_position_embedding_type: str = 'learned_absolute',
+ language_rotary_percent: float = 1.0,
+ pre_process: bool = True,
+ post_process: bool = True,
+ add_encoder: bool = True,
+ add_decoder: bool = True,
+ img_h: int = 336,
+ img_w: int = 336,
+ patch_dim: int = 14,
+ language_rotary_base: int = 10000,
+ language_rope_scaling: bool = False,
+ ) -> None:
+ super().__init__(config=language_transformer_config)
+
+ if has_config_logger_enabled(language_transformer_config):
+ log_config_to_disk(language_transformer_config, locals(), prefix=type(self).__name__)
+
+ logging.getLogger(__name__).warning(
+ "LLaVA model is under active development. "
+ "It may be missing features and its methods may change."
+ )
+
+ self.pre_process = pre_process
+ self.post_process = post_process
+ self.add_encoder = add_encoder
+ self.add_decoder = add_decoder
+
+ self.encoder_hidden_state = None
+ self.vision_model = None
+ self.vision_projection = None
+ self.language_model = None
+
+ self.sequence_parallel_lm = language_transformer_config.sequence_parallel
+ if self.sequence_parallel_lm:
+ assert (
+ language_transformer_layer_spec.submodules.self_attention.submodules.core_attention
+ == TEDotProductAttention
+ ), "Sequence Parallelism is supported only with Transformer Engine DotProductAttention."
+ self.tp_comm_overlap_lm = language_transformer_config.tp_comm_overlap
+
+ # This attribute is needed to check if an all-reduce is required
+ # on the word embeddings inside `finalize_model_grads._allreduce_word_embedding_grads`.
+ self.share_embeddings_and_output_weights = False
+ if self.add_decoder:
+ self.language_model = GPTModel(
+ config=language_transformer_config,
+ transformer_layer_spec=language_transformer_layer_spec,
+ vocab_size=language_vocab_size,
+ max_sequence_length=language_max_sequence_length,
+ parallel_output=parallel_output,
+ position_embedding_type=language_position_embedding_type,
+ rotary_percent=language_rotary_percent,
+ pre_process=self.pre_process,
+ post_process=self.post_process,
+ rotary_base=language_rotary_base,
+ rope_scaling=language_rope_scaling,
+ )
+ self.share_embeddings_and_output_weights = (
+ self.language_model.share_embeddings_and_output_weights
+ )
+ self._language_max_sequence_length = language_max_sequence_length
+ self._language_is_pipeline_parallel = (
+ language_transformer_config.pipeline_model_parallel_size > 1
+ )
+
+ class_token_len = 1
+ if self.add_encoder:
+ self._drop_vision_class_token = drop_vision_class_token
+ add_class_token = True
+ if vision_transformer_config.vision_model_type == "siglip":
+ class_token_len = 0
+ add_class_token = False
+ error_msg = (
+ "Siglip does not support vision class token, "
+ "set disable-vision-class-token to False."
+ )
+ assert not self._drop_vision_class_token, error_msg
+ self.vision_model = CLIPViTModel(
+ vision_transformer_config,
+ vision_transformer_layer_spec,
+ img_h=img_h,
+ img_w=img_w,
+ class_token_len=class_token_len,
+ patch_dim=patch_dim,
+ model_subtype=vision_transformer_config.vision_model_type,
+ add_class_token=add_class_token,
+ )
+ # Map (intermediate) vision model outputs to the language model input dimension.
+ self.vision_projection = MultimodalProjector(
+ vision_projection_config,
+ vision_projection_layer_spec,
+ vision_projection_type,
+ vision_transformer_config.hidden_size, # input size to the projection.
+ )
+ # Ignore missing weights for the vision projection during checkpoint loading.
+ # This should be disabled by default but can be enabled if your checkpoint contains
+ # pretrained vision and language models but not the projection from vision model
+ # outputs to language model inputs.
+ if allow_missing_vision_projection_checkpoint:
+ vision_projection_param_names = [
+ f"vision_projection.{name}"
+ for name in self.vision_projection.state_dict().keys()
+ ]
+ self.vision_projection.register_load_state_dict_post_hook(
+ partial(_load_state_dict_hook_ignore_param_names, vision_projection_param_names)
+ )
+
+ self._img_seq_len = get_num_image_embeddings(
+ img_h,
+ img_w,
+ patch_dim,
+ vision_transformer_config.vision_model_type,
+ drop_vision_class_token,
+ class_token_len,
+ )
+
+ def shared_embedding_or_output_weight(self):
+ """This is a convenience method to surface the language model's word embeddings, which is
+ necessary for `finalize_model_grads._allreduce_word_embedding_grads`."""
+ if self.add_decoder:
+ return self.language_model.shared_embedding_or_output_weight()
+ return None
+
+ def set_input_tensor(self, input_tensor) -> None:
+ """Set model chunk input tensor."""
+ # This is usually handled in schedules.py but some inference code still
+ # gives us non-lists or None
+ if not isinstance(input_tensor, list):
+ input_tensor = [input_tensor]
+ assert len(input_tensor) == 1, 'input_tensor should only be length 1 for llava'
+
+ if self.add_encoder and self.add_decoder:
+ self.vision_model.set_input_tensor(input_tensor[0])
+ elif self.add_encoder:
+ self.vision_model.set_input_tensor(input_tensor[0])
+ elif self.pre_process:
+ self.encoder_hidden_state = input_tensor[0]
+ else:
+ self.language_model.set_input_tensor(input_tensor[0])
+
+ def freeze(
+ self, freeze_language_model: bool, freeze_vision_model: bool, freeze_vision_projection: bool
+ ):
+ """Freeze model modules.
+
+ Make specific modules non-trainable by setting requires_grad to False.
+
+ Args:
+ freeze_language_model (bool): Freeze the language model module.
+ freeze_vision_model (bool): Freeze the vision model module.
+ freeze_vision_projection (bool): Freeze the vision projection module.
+ """
+ modules = []
+ if freeze_language_model and self.language_model is not None:
+ modules.append(self.language_model)
+ if freeze_vision_model and self.vision_model is not None:
+ modules.append(self.vision_model)
+ if freeze_vision_projection and self.vision_projection is not None:
+ modules.append(self.vision_projection)
+
+ for module in modules:
+ for param in module.parameters():
+ param.requires_grad = False
+
+ def _preprocess_data(
+ self,
+ image_embeddings,
+ language_embeddings,
+ input_ids,
+ loss_mask,
+ labels,
+ use_inference_kv_cache,
+ image_token_index,
+ num_image_tiles,
+ attention_mask,
+ ):
+ """Preprocess input data before input to language model.
+
+ This function is adopted from
+ https://github.com/huggingface/transformers/blob/85817d98fb60977c97e3014196a462b732d2ed1a/src/transformers/models/llava_next/modeling_llava_next.py#L409
+ for our input data conventions.
+
+ image_token_index = -200 indicates the image position in the input_ids = [0, 1, -200, 2, 3]
+ and labels = [1, -200, 2, 3, 4], for example.
+ We want to replace the image position (-200) with image_embeddings and return the following:
+ - final_embeddings = [0, 1, image_embeddings, 2, 3],
+ - final_labels = [1, -100, 2, 3, 4]
+ - final_loss_mask = [1, 0, 0, 1, 1]
+
+ This function handles samples without images (text-only sample). It also handles samples
+ with images that are split into multiples tiles.
+
+ If pipeline parallelism is not used, then self.pre_process and self.post_process
+ are both True and we update both input embeddings, labels and loss masks (if available).
+
+ If pipeline parallelism is used, then we do the following
+ - the first language model chunk has self.pre_process = True and
+ self.post_process = False. We update input embeddings.
+ - the middle language model chunk(s) has self.pre_process = False and
+ self.post_process = False. We don't need to update anything.
+ - the last language model chunk has self.pre_process = False and
+ self.post_process = True. We update labels and loss mask.
+
+ TODO: This function should adjust the attention mask too.
+ Currently, we assume the language model uses a causal mask.
+
+ Returns:
+ final_embedding (torch.Tensor): image and text embeddings [combined_seq_len, b, h].
+ final_labels (torch.Tensor): labels for image and text positions [b, combined_seq_len].
+ final_loss_mask (torch.Tensor): loss mask [b, combined_seq_len].
+ """
+ assert self.add_decoder, "input text preprocessing is only needed for the language model"
+
+ # No pre- or postprocessing needed.
+ # With pipeline parallel > 2, this means a chunk in the middle of the model.
+ if not self.pre_process and not self.post_process:
+ return language_embeddings, loss_mask, labels, attention_mask
+
+ # If using the inference KV cache, the image tokens are already computed.
+ if use_inference_kv_cache:
+ return language_embeddings, loss_mask, labels, attention_mask
+
+ img_seq_len = self._img_seq_len
+ batch_size, text_seq_len = input_ids.shape
+
+ has_labels = labels is not None
+ if has_labels:
+ assert (
+ labels.shape == loss_mask.shape
+ ), f"mismatching labels shape {labels.shape} and loss mask shape {loss_mask.shape}"
+
+ # Create indices for new text and label positions.
+ with torch.no_grad():
+ image_token_mask = input_ids == image_token_index
+ num_images_per_sample = torch.sum(image_token_mask, dim=-1)
+
+ # Number of tiles per sample.
+ num_image_tiles_batch = num_image_tiles.split(num_images_per_sample.tolist(), dim=0)
+ num_image_tiles_batch = torch.tensor(
+ [x.sum() for x in num_image_tiles_batch], device=input_ids.device
+ )
+
+ # Sequence length for each sample is the image sequence length multiplied by
+ # the number of tiles for that image, minus image token indices,
+ # plus text sequence length.
+ seq_lens = num_image_tiles_batch * img_seq_len - num_images_per_sample + text_seq_len
+ max_seq_len = seq_lens.max()
+ # Pipeline parallel expects fixed input size. Check if we need to pad.
+ if (
+ self._language_is_pipeline_parallel
+ and max_seq_len < self._language_max_sequence_length
+ ):
+ max_seq_len = self._language_max_sequence_length
+
+ if self.sequence_parallel_lm:
+ if self.tp_comm_overlap_lm:
+ # If shorter: Pad to language_max_sequence_length to use TP Comm overlap.
+ # If longer: Gets truncated later.
+ if max_seq_len < self._language_max_sequence_length:
+ padded_seq_len = self._language_max_sequence_length
+ else:
+ # Pad to multiple of tp size for sequence parallelism
+ tp_world_size = get_tensor_model_parallel_world_size()
+ padded_seq_len = int(
+ (max_seq_len + (tp_world_size - 1)) // tp_world_size * tp_world_size
+ )
+ sp_padding_needed = padded_seq_len - max_seq_len
+ max_seq_len = padded_seq_len
+ batch_indices, non_image_indices = torch.where(input_ids != image_token_index)
+
+ # New position ids for the text tokens, shifted by the image sequence length.
+ # E.g. for input_ids = [-200, 1, 2, 3] and img_seq_len = 576, we get
+ # new_position_ids = [576, 577, 578, 579]. text_position_ids are then [577, 578, 579].
+ image_token_mask_lens = image_token_mask.int().clone()
+ # -1 is for the removed image token index.
+ image_token_mask_lens[image_token_mask] = num_image_tiles * img_seq_len - 1
+ # +1 is needed here for the cumulative sum. -1 is adjusting for zero-based indexing.
+ new_position_ids = torch.cumsum((image_token_mask_lens + 1), dim=-1) - 1
+ text_position_ids = new_position_ids[batch_indices, non_image_indices]
+
+ # Labels are shifted to left by one.
+ # So, shift text position ids and non-image indices to left by one.
+ if has_labels:
+ label_text_position_ids = text_position_ids - 1
+ valid_label_text_position_ids = label_text_position_ids >= 0
+ label_text_position_ids = label_text_position_ids[valid_label_text_position_ids]
+
+ label_batch_indices = batch_indices[valid_label_text_position_ids]
+
+ label_non_image_indices = non_image_indices - 1
+ valid_label_non_image_indices = label_non_image_indices >= 0
+ label_non_image_indices = label_non_image_indices[valid_label_non_image_indices]
+
+ # Create a mask for the image embedding positions.
+ images_mask = torch.full(
+ (batch_size, max_seq_len), True, dtype=torch.bool, device=input_ids.device
+ )
+ # No images in the text positions.
+ images_mask[batch_indices, text_position_ids] = False
+ # Samples can have different amount of images tokens.
+ # new_position_ids[:, -1] gives the last text position id for each sample.
+ # Padding is needed when the number of image tokens differs.
+ first_padding_idx = new_position_ids[:, -1] + 1
+ images_mask[
+ torch.arange(max_seq_len, device=first_padding_idx.device).repeat(batch_size, 1)
+ >= first_padding_idx.unsqueeze(1)
+ ] = False
+
+ # Create the final input embedding (if this is the first language model stage).
+ final_embedding = None
+ if self.pre_process:
+ embed_dim = language_embeddings.shape[-1]
+ final_embedding = torch.zeros(
+ batch_size,
+ max_seq_len,
+ embed_dim,
+ dtype=language_embeddings.dtype,
+ device=language_embeddings.device,
+ )
+
+ # Put text embeddings to the text positions in the result tensor.
+ final_embedding[batch_indices, text_position_ids] = language_embeddings[
+ batch_indices, non_image_indices
+ ]
+
+ # Put image embeddings to image positions.
+ final_embedding[images_mask] = (
+ image_embeddings.permute(1, 0, 2).reshape(-1, embed_dim).contiguous()
+ )
+
+ # Create the final labels and loss mask (if this is the last language model stage).
+ final_labels, final_loss_mask = None, None
+ if has_labels:
+ final_labels = torch.full(
+ (batch_size, max_seq_len), IGNORE_INDEX, dtype=labels.dtype, device=labels.device
+ )
+ final_loss_mask = torch.full(
+ (batch_size, max_seq_len), 0, dtype=loss_mask.dtype, device=loss_mask.device
+ )
+
+ # Put text labels and loss mask to the text positions.
+ final_labels[label_batch_indices, label_text_position_ids] = labels[
+ label_batch_indices, label_non_image_indices
+ ]
+
+ final_loss_mask[batch_indices, text_position_ids] = loss_mask[
+ batch_indices, non_image_indices
+ ]
+
+ # For labels, pick the last label index that got dropped by the shift to left.
+ label_extra_text_position_ids = seq_lens - 1
+ batch_range = torch.arange(len(label_extra_text_position_ids))
+ final_labels[batch_range, label_extra_text_position_ids] = labels[batch_range, -1]
+
+ # Loss mask the image positions.
+ final_loss_mask[images_mask] = 0
+
+ # Loss mask last text position just before an image
+ # so that text token does not need to predict the first image token.
+ batch_image_indices, image_indices = torch.where(image_token_mask)
+ # Indices just before image tokens. If it's -1, skip it.
+ before_image_indices = image_indices - 1
+ valid = before_image_indices >= 0
+ valid_batch_image_indices = batch_image_indices[valid]
+ valid_before_image_indices = before_image_indices[valid]
+ # Map those indices those position ids.
+ valid_before_image_indices = new_position_ids[
+ valid_batch_image_indices, valid_before_image_indices
+ ]
+
+ final_loss_mask[valid_batch_image_indices, valid_before_image_indices] = 0
+
+ if final_embedding is not None and has_labels:
+ assert (
+ final_embedding.shape[:2] == final_labels.shape == final_loss_mask.shape
+ ), "unexpected shapes after data preprocessing"
+
+ truncate_labels = has_labels and final_labels.shape[1] > self._language_max_sequence_length
+ if truncate_labels:
+ final_labels = final_labels[:, : self._language_max_sequence_length]
+ final_loss_mask = final_loss_mask[:, : self._language_max_sequence_length]
+
+ if final_embedding is not None:
+ final_embedding = final_embedding.transpose(1, 0).contiguous()
+ # Truncate if exceeding the language model's max sequence length.
+ if final_embedding.shape[0] > self._language_max_sequence_length:
+ final_embedding = final_embedding[: self._language_max_sequence_length]
+ if self.sequence_parallel_lm:
+ # Create an attention mask. This ensures correct computation.
+ # This is done even when no padding was done as we set mask_type to
+ # 'padding' or 'padding_causal' when using SP.
+ if attention_mask is None:
+ # Create base attention mask with original seq len to indicate valid tokens
+ attention_mask = (
+ torch.ones(
+ (
+ final_embedding.shape[1],
+ final_embedding.shape[0] - sp_padding_needed,
+ ),
+ device=final_embedding.device,
+ )
+ .unsqueeze(1)
+ .unsqueeze(1)
+ ) # [b, 1, 1, final seq len - sp_padding_needed]
+ if sp_padding_needed > 0:
+ # Add the padding portion of the mask
+ attention_mask = torch.nn.functional.pad(attention_mask, (0, sp_padding_needed))
+ if is_te_min_version("1.7.0"):
+ # Attention mask True/False meaning flipped in 1.7.0
+ attention_mask = attention_mask < 0.5
+ final_embedding = tensor_parallel.scatter_to_sequence_parallel_region(
+ final_embedding
+ )
+
+ return final_embedding, final_labels, final_loss_mask, attention_mask
+
+ def forward(
+ self,
+ images: torch.Tensor,
+ input_ids: torch.Tensor,
+ position_ids: torch.Tensor,
+ attention_mask: torch.Tensor,
+ labels: Optional[torch.Tensor] = None,
+ loss_mask: Optional[torch.Tensor] = None,
+ inference_params: Optional[InferenceParams] = None,
+ num_image_tiles: Optional[List[int]] = None,
+ image_token_index: Optional[int] = IMAGE_TOKEN_INDEX,
+ runtime_gather_output: Optional[bool] = None,
+ ) -> torch.Tensor:
+ """Forward function of the LLaVA model.
+
+ Args:
+ images (torch.Tensor): input images of shape [num_tiles, img_h, img_w].
+ num_tiles means the number of image tiles in this batch.
+ num_tiles = 0 if the batch doesn't contain images.
+ input_ids (torch.Tensor): input text ids [batch, text_seq_len].
+ position_ids (torch.Tensor): input text position ids [batch, text_seq_len].
+ attention_mask (torch.Tensor): Language model attention mask
+ [batch, 1, 1, combined_seq_len].
+ labels (torch.Tensor): Optional target text labels [batch, combined_seq_len].
+ loss_mask (torch.Tensor): Text loss mask [batch, text_seq_len].
+ inference_params (InferenceParams): Inference-time parameters including KV cache.
+ num_image_tiles (list of int): Number of tiles per image. Default 1 tile per image.
+ image_token_index (int): ID for input images.
+ runtime_gather_output (bool): Gather output at runtime. Default None means
+ `parallel_output` arg in the constructor will be used.
+
+ Returns:
+ output (torch.Tensor): Loss of shape [b, s] if labels are provided,
+ otherwise logits of shape [b, s, vocab_size].
+ loss_mask (torch.Tensor): Loss mask expanded to combined sequence length. Shape [b, s].
+ """
+ use_inference_kv_cache = (
+ inference_params is not None
+ and "image_tokens_count" in inference_params.key_value_memory_dict
+ )
+ has_images = images.shape[0] > 0
+
+ # If running inference, we can skip image token computation
+ # if they were computed already earlier for this sample.
+ if use_inference_kv_cache:
+ image_embeddings = None
+ elif self.add_encoder and not has_images:
+ # If no images provided, use an empty image embeddings tensor.
+ image_embeddings = torch.tensor([], dtype=images.dtype, device=images.device).reshape(
+ 0, 0, 0
+ )
+ elif self.add_encoder and has_images:
+ image_embeddings = self.vision_model(images) # [num_tiles, img_seq_len, h_vision]
+ if self._drop_vision_class_token:
+ image_embeddings = image_embeddings[:, self.vision_model.class_token_len :, :]
+ # contiguous() required as `permute` can sparsify the tensor and this breaks pipelining
+ image_embeddings = image_embeddings.permute(
+ 1, 0, 2
+ ).contiguous() # [img_seq_len, num_tiles, h_vision]
+
+ # map vision model output size to language model input size.
+ image_embeddings = self.vision_projection(
+ image_embeddings
+ ) # [img_seq_len, num_tiles, h_language]
+
+ # TODO: Support batched inference.
+ # In inference, the language model KV cache will be updated for image token positions.
+ # Store the image tokens sequence length to be used as an offset to the KV cache later.
+ if inference_params is not None:
+ inference_params.key_value_memory_dict["image_tokens_count"] = (
+ image_embeddings.shape[0] * image_embeddings.shape[1]
+ )
+ else:
+ image_embeddings = self.encoder_hidden_state
+
+ if not self.add_decoder:
+ return image_embeddings, loss_mask
+
+ language_embeddings = None
+ if self.pre_process:
+ input_ids_text = input_ids.clone()
+ input_ids_text[input_ids_text == image_token_index] = 0
+ # Note: This adds absolute position embedding but not RoPE.
+ # Each image is counted as one position.
+ # RoPE is added in language_model forward. Each image embedding is one position.
+ if self.sequence_parallel_lm:
+ # Pad to nearest multiple of TP world size for embedding.
+ tp_world_size = get_tensor_model_parallel_world_size()
+ padded_seq_len = (
+ int(
+ (input_ids_text.shape[1] + tp_world_size - 1)
+ // tp_world_size
+ * tp_world_size
+ )
+ - input_ids_text.shape[1]
+ )
+ if padded_seq_len != 0:
+ input_ids_text = torch.nn.functional.pad(input_ids_text, (0, padded_seq_len))
+ if position_ids is not None:
+ position_ids = torch.nn.functional.pad(position_ids, (0, padded_seq_len))
+ language_embeddings = self.language_model.embedding(
+ input_ids=input_ids_text, position_ids=position_ids
+ ) # [text_seq_len, b, h_language]
+ if self.sequence_parallel_lm:
+ # Gather the language embeddings back.
+ # We use the full embedding to insert image embeddings
+ # and then scatter to avoid load imbalance.
+ language_embeddings = tensor_parallel.gather_from_sequence_parallel_region(
+ language_embeddings, tensor_parallel_output_grad=False
+ )
+ # Remove the padding done for SP as we'll need new padding calculation
+ # after image embeddings are inserted.
+ if padded_seq_len != 0:
+ language_embeddings = language_embeddings[:-padded_seq_len]
+ language_embeddings = language_embeddings.transpose(
+ 1, 0
+ ).contiguous() # [b, text_seq_len, h_language]
+
+ # Assume 1 tile per image if the number of tiles is not provided.
+ if num_image_tiles is None:
+ num_image_tiles = torch.ones(images.shape[0], dtype=torch.int, device=input_ids.device)
+
+ # Preprocess input, labels and loss mask.
+ combined_embeddings, new_labels, new_loss_mask, attention_mask = self._preprocess_data(
+ image_embeddings,
+ language_embeddings,
+ input_ids,
+ loss_mask,
+ labels,
+ use_inference_kv_cache,
+ image_token_index,
+ num_image_tiles,
+ attention_mask,
+ ) # [combined_seq_len, b, h_language], [b, combined_seq_len], [b, combined_seq_len]
+
+ output = self.language_model(
+ input_ids=None,
+ position_ids=None,
+ attention_mask=attention_mask,
+ decoder_input=combined_embeddings,
+ labels=new_labels,
+ inference_params=inference_params,
+ runtime_gather_output=runtime_gather_output,
+ )
+
+ if labels is None or loss_mask is None:
+ return output
+
+ return output, new_loss_mask
+
+
+def _load_state_dict_hook_ignore_param_names(
+ param_names: List[str], module: torch.nn.Module, incompatible_keys: namedtuple
+):
+ """Hook to ignore missing keys during checkpoint loading.
+
+ By default, this should not be used to avoid accidentally missing weights in checkpoint loading.
+
+ Example use case: Use this if you want to load a checkpoint that contains vision and language
+ model weights but not the vision projection weights.
+
+ Args:
+ param_names (list str): Parameter names allowed to be missing when calling load_state_dict.
+ module (torch.nn.Module): The torch module this hook applies to. Required by the torch API.
+ incompatible_keys (namedtuple): Namedtuple with fields missing_keys and unexpected_keys,
+ which collect the missing and unexpected keys, respectively.
+ """
+ for param_name in param_names:
+ if param_name in incompatible_keys.missing_keys:
+ logging.getLogger(__name__).warning(
+ f"{param_name} being removed from incompatible_keys.missing_keys in LlavaModel"
+ )
+ incompatible_keys.missing_keys.remove(param_name)
diff --git a/megatron/core/models/multimodal/llava_spec.py b/megatron/core/models/multimodal/llava_spec.py
new file mode 100644
index 0000000000..40e58d0bfc
--- /dev/null
+++ b/megatron/core/models/multimodal/llava_spec.py
@@ -0,0 +1,87 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from megatron.core.extensions.transformer_engine import (
+ TEDotProductAttention,
+ TELayerNormColumnParallelLinear,
+ TENorm,
+ TERowParallelLinear,
+)
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.models.gpt.gpt_layer_specs import _get_mlp_module_spec
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
+from megatron.core.transformer.dot_product_attention import DotProductAttention
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+
+try:
+ import apex # pylint: disable=unused-import
+
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+
+ HAVE_APEX = True
+ LNImpl = FusedLayerNorm
+except ImportError:
+ import warnings
+
+ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm
+
+ warnings.warn(f'Apex is not installed. Falling back to Torch LayerNorm')
+ LNImpl = WrappedTorchLayerNorm
+
+
+def decoder_model_with_transformer_engine_default_spec(
+ num_experts: int = None, moe_grouped_gemm: bool = False, qk_layernorm: bool = False
+) -> ModuleSpec:
+ """LLava decoder TE spec (uses Transformer Engine components)."""
+ mlp = _get_mlp_module_spec(
+ use_te=True, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm
+ )
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=TELayerNormColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ q_layernorm=TENorm if qk_layernorm else IdentityOp,
+ k_layernorm=TENorm if qk_layernorm else IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+
+
+def decoder_model_with_local_default_spec(
+ num_experts: int = None, moe_grouped_gemm: bool = False, qk_layernorm: bool = False
+) -> ModuleSpec:
+ """LLava decoder local spec."""
+ mlp = _get_mlp_module_spec(
+ use_te=False, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm
+ )
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=LNImpl,
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=LNImpl,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
diff --git a/megatron/core/models/retro/__init__.py b/megatron/core/models/retro/__init__.py
new file mode 100644
index 0000000000..ea7cea6d8f
--- /dev/null
+++ b/megatron/core/models/retro/__init__.py
@@ -0,0 +1,13 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""
+Exports:
+
+ - RetroConfig: configuration dataclass for RetroModel.
+ - RetroModel: The Retro model.
+ - get_retro_decoder_block_spec: Get spec for Retro decoder transformer block.
+"""
+
+from .config import RetroConfig
+from .decoder_spec import get_retro_decoder_block_spec
+from .model import RetroModel
diff --git a/megatron/core/models/retro/base_attention.py b/megatron/core/models/retro/base_attention.py
new file mode 100644
index 0000000000..ee8656d96a
--- /dev/null
+++ b/megatron/core/models/retro/base_attention.py
@@ -0,0 +1,43 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Base class for decoder and encoder attention modules."""
+
+from megatron.core.models.retro.config import RetroConfig
+from megatron.core.transformer.attention import CrossAttention, CrossAttentionSubmodules
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.module import MegatronModule
+
+
+class BaseRetroCrossAttention(MegatronModule):
+ """Base class for Retro cross attention, for both encoder & decoder layers.
+
+ This class collects the retro arguments below (i.e., num neighbors, chunk
+ length, and retrieve length) for use in Retro's custom cross attention
+ operators.
+
+ Args:
+ config (RetroConfig): Retro config.
+ submodules (CrossAttentionSubmodules): Cross attention submodules.
+ layer_number (int): Layer number within transformer block.
+ attn_mask_type (AttnMaskType): Mask type ('causal' or 'padding').
+ """
+
+ def __init__(
+ self,
+ config: RetroConfig,
+ submodules: CrossAttentionSubmodules,
+ layer_number: int = 1,
+ attn_mask_type: AttnMaskType = AttnMaskType.padding,
+ ):
+ super().__init__(config=config)
+
+ self.attn = CrossAttention(
+ config=config,
+ submodules=submodules,
+ layer_number=layer_number,
+ attn_mask_type=attn_mask_type,
+ )
+
+ self.retro_num_neighbors = config.retro_num_neighbors
+ self.retro_chunk_length = config.retro_chunk_length
+ self.retro_retrieved_length = config.retro_retrieved_length
diff --git a/megatron/core/models/retro/config.py b/megatron/core/models/retro/config.py
new file mode 100644
index 0000000000..d4b5c9684b
--- /dev/null
+++ b/megatron/core/models/retro/config.py
@@ -0,0 +1,85 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Configuration dataclass for a RetroModel."""
+
+import os
+from dataclasses import dataclass
+
+from megatron.core.transformer import TransformerConfig
+from megatron.core.utils import is_te_min_version
+
+
+@dataclass
+class RetroConfig(TransformerConfig):
+ """Configuration object for Retro models."""
+
+ # Retro.
+ retro_project_dir: str = None
+ """Retro project directory, which contains the preprocessed data for for pretraining. This
+ directory is built during preprocessing (see tools/retro/README.md), and contains
+ subdirectories for the chunk database and pretraining neighbors.
+ """
+
+ retro_block_size: int = None
+ """Number of records to load per data file, as saved during preprocessing. Block processing is
+ used for efficient data preprocessing.
+ """
+
+ retro_chunk_length: int = None
+ """Chunk length used for performing chunked- cross-attention (CCA)."""
+
+ retro_encoder_num_layers: int = 2
+ """Number of layers to use for the retrieval encoder."""
+
+ retro_encoder_hidden_dropout: float = 0.1
+ """Hidden dropout for retrieval encoder."""
+
+ retro_encoder_attention_dropout: float = 0.1
+ """Attention dropout for retrieval encoder."""
+
+ retro_neighbor_dirs: dict = None
+ """Directory names of saved neighbor id files for train, valid, and test datasets."""
+
+ retro_num_neighbors: int = 2
+ """Number of neighbors to retrieve during pretraining."""
+
+ retro_num_retrieved_chunks: int = 2
+ """Number of chunks to retrieve from the retrieval database."""
+
+ retro_retrieved_length: int = None
+ """Cached value of retro_num_retrieved_chunks * retro_chunk_length (i.e., the total number of
+ retrieved tokens; neighbor + continuation).
+ """
+
+ retro_split_preprocessing: str = None
+ """Data split used during data preprocessing."""
+
+ retro_verify_neighbor_count: bool = True
+ """Verify that len(GPT dataset) == len(saved neighbors)."""
+
+ def __post_init__(self) -> None:
+ """Validate Retro config."""
+
+ super().__post_init__()
+
+ # Validate Transformer Engine version.
+ if is_te_min_version("1.3"):
+ try:
+ assert os.getenv("NVTE_FLASH_ATTN") == "0"
+ assert os.getenv("NVTE_FUSED_ATTN") == "0"
+ except Exception as e:
+ raise Exception(
+ "When using Transformer Engine >= 1.3, environment vars NVTE_FLASH_ATTN "
+ "and NVTE_FUSED_ATTN most both be defined and set to '0'. "
+ "Currently, NVTE_FLASH_ATTN == %s, NVTE_FUSED_ATTN == %s."
+ % (
+ os.getenv("NVTE_FLASH_ATTN", "[unset]"),
+ os.getenv("NVTE_FUSED_ATTN", "[unset]"),
+ )
+ )
+
+ # Preprocessing split should be defined.
+ assert self.retro_split_preprocessing is not None
+
+ # Pre-compute retrieved length.
+ self.retro_retrieved_length = self.retro_num_retrieved_chunks * self.retro_chunk_length
diff --git a/megatron/core/models/retro/decoder_attention.py b/megatron/core/models/retro/decoder_attention.py
new file mode 100644
index 0000000000..6b7a04d884
--- /dev/null
+++ b/megatron/core/models/retro/decoder_attention.py
@@ -0,0 +1,305 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Retro's cross attention modules for the decoder block."""
+
+from functools import partial
+from typing import Callable
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from megatron.core import InferenceParams
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.models.retro.base_attention import BaseRetroCrossAttention
+from megatron.core.models.retro.config import RetroConfig
+from megatron.core.models.retro.utils import get_all_true_mask
+from megatron.core.transformer import ModuleSpec
+from megatron.core.transformer.attention import CrossAttentionSubmodules
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_block import TransformerBlock
+
+
+class RetroDecoderCrossAttention(BaseRetroCrossAttention):
+ """Retro decoder's chunked cross attention operator.
+
+ See this paper for more details: https://arxiv.org/abs/2112.04426.
+ Neighboring chunks retrieved from the chunk database are used here for
+ chunked-cross attention.
+
+ ** Note about 'encoder_block_spec' **
+
+ Retro is an encoder-decoder model that uses its encoder for encoding
+ neighboring chunks that are retrieved from a chunk database. These
+ encoded neighbors are then used in the decoder stack for performing
+ chunked-cross attention (see paper link above).
+
+ In contrast to the T5 model, the encoder and decoder are computationally
+ intertwined, since the input to the encoder is the output of the self-
+ attention of the first decoder layer. As such, the encoder block itself
+ is instantiated within the first Retro decoder layer, in order to receive
+ the self-attention's output. (Note, that only the first decoder layer
+ instantiates an encoder block, and the remaining decoder layers use the
+ encoder output from the first decoder layer.)
+
+ Args:
+ config (RetroConfig): Retro config.
+ submodules (CrossAttentionSubmodules): Cross attention submodules.
+ layer_number (int): Layer number within transformer block.
+ attn_mask_type (AttnMaskType): Mask type ('causal' or 'padding').
+ encoder_block_spec (ModuleSpec): The first Retro decoder layer is provided with a transformer block spec to construct the neighbor encoder.
+ """
+
+ def __init__(
+ self,
+ config: RetroConfig,
+ submodules: CrossAttentionSubmodules,
+ layer_number: int = 1,
+ attn_mask_type: AttnMaskType = AttnMaskType.padding,
+ encoder_block_spec: ModuleSpec = None,
+ ):
+ super().__init__(
+ config=config,
+ submodules=submodules,
+ layer_number=layer_number,
+ attn_mask_type=attn_mask_type,
+ )
+
+ if encoder_block_spec:
+ self.encoder = TransformerBlock(
+ config=config, spec=encoder_block_spec, pre_process=True, post_process=False
+ )
+ # self._encoder_key = 'encoder' # ... necessary?
+ else:
+ self.encoder = None
+
+ def forward(
+ self,
+ hidden_states: Tensor,
+ attention_mask: Tensor,
+ key_value_states: Tensor = None,
+ inference_params: InferenceParams = None,
+ # rotary_pos_emb: Tensor = None, # ... unsupported for retro.
+ ) -> dict:
+ """Cross attention for Retro decoder.
+
+ Notation:
+ ns : Sequence length.
+ bs : Batch size.
+ d : Hidden size.
+ l : Number of chunks per sample (i.e., seq_length/chunk_length).
+ m : Number of tokens per chunk.
+ k : Number of neighbors.
+ r : Number of retrieved tokens (neighbors + continuation).
+
+ Args:
+ hidden_states (Tensor): Transformer layer hidden states.
+ attention_mask (Tensor): Attention mask.
+ key_value_states (Tensor): Neighbor embeddings if first decoder layer, else encoder output.
+ inference_params (InferenceParams): Inference params.
+
+ Returns:
+ A dict consisting of the attention output and context, along with other scalars necessary for performing the downstream bias-dropout-add.
+ """
+
+ # hidden_states: [ ns, bs, d ]
+ # key_value_states: [ r, k*bs*l, d ]
+
+ ns, bs, d = hidden_states.shape
+ l = int(np.ceil(ns / self.retro_chunk_length))
+
+ # Retrieve neighbors.
+ if self.encoder:
+
+ # Sequence length remainder.
+ first_ns = ns % self.retro_chunk_length
+
+ # Case 1: Sequence length not divisible by chunk length.
+ if first_ns > 0:
+
+ # Split sequence into first partial chunk & remaining chunks.
+ first_chunk, rest_chunk = hidden_states[:first_ns], hidden_states[first_ns:]
+
+ # Pad partial chunk with zeros.
+ first_chunk = torch.nn.functional.pad(
+ first_chunk, (0, 0, 0, 0, 0, self.retro_chunk_length - first_ns), 'constant', 0
+ )
+
+ # Concatenate padded chunk with remaining chunks.
+ chunked_output = torch.cat((first_chunk, rest_chunk), dim=0) # [ l*m, bs, d ]
+
+ # Case 2: Sequence length is divisible by chunk length.
+ else:
+ chunked_output = hidden_states # [ l*m, bs, d ]
+
+ # Chunk & permute hidden states.
+ # - hidden_states: [ l*m, bs, d ]
+ # - chunked_output: [ m, bs*l, d ]
+ chunked_output = (
+ chunked_output.reshape(l, self.retro_chunk_length, bs, d)
+ .permute(1, 2, 0, 3)
+ .reshape(self.retro_chunk_length, bs * l, d)
+ .contiguous()
+ )
+
+ # flash attn: [ b, h, sq, sk ]
+ # fused attn: [ b, 1, 1, sq ]
+ chunked_output_mask = get_all_true_mask(
+ size=(1, 1, chunked_output.shape[0], key_value_states.shape[0]),
+ device=chunked_output.device,
+ )
+
+ # Encode neighbors. (Note: 'key_value_states' re-assigned here.)
+ key_value_states = self.encoder(
+ hidden_states=key_value_states,
+ attention_mask=attention_mask,
+ context=chunked_output,
+ context_mask=chunked_output_mask,
+ inference_params=inference_params,
+ ) # [ r, k*bs*l, d ]
+ key_value_states = key_value_states.reshape(
+ self.retro_retrieved_length * self.retro_num_neighbors, bs * l, d
+ ) # [ r*k, bs*l, d ]
+
+ # Attend starting at last token of first chunk.
+ pad = (ns - 1) % self.retro_chunk_length
+ attending_chunks = hidden_states[pad:]
+
+ # Pad attending tokens to sequence length.
+ padded_chunks = torch.nn.functional.pad(
+ attending_chunks, (0, 0, 0, 0, 0, self.retro_chunk_length - 1), 'constant', 0
+ )
+
+ # Permute attending chunks.
+ # - padded_chunks: [ l*m, bs, d ]
+ # - padded_chunked_output: [ m, bs*l, d ] (matches 'chunked_output' above)
+ padded_chunked_output = padded_chunks.reshape(l, self.retro_chunk_length, bs, d).permute(
+ 1, 2, 0, 3
+ )
+ padded_chunked_output = padded_chunked_output.reshape(
+ self.retro_chunk_length, bs * l, d
+ ).contiguous()
+
+ # flash attn: [ b, h, sq, sk ]
+ # fused attn: [ b, 1, 1, sq ]
+ padded_chunked_output_mask = get_all_true_mask(
+ size=(1, 1, padded_chunked_output.shape[0], key_value_states.shape[0]),
+ device=padded_chunked_output.device,
+ )
+
+ # Attend to encoded neighbors.
+ attention_output, attention_bias = self.attn(
+ hidden_states=padded_chunked_output,
+ attention_mask=padded_chunked_output_mask,
+ key_value_states=key_value_states,
+ )
+
+ # Return dimensions for bias-dropout step.
+ return {
+ "ns": ns,
+ "bs": bs,
+ "d": d,
+ "l": l,
+ "pad": pad,
+ "attention_output": attention_output, # [ m, bs*l, d ]
+ "attention_bias": attention_bias, # [ d ]
+ "context": key_value_states, # [ r*k, bs*l, d ]
+ }
+
+
+class RetroDecoderBiasDropoutAdd(MegatronModule):
+ """Retro decoder's bias-dropout-add operator.
+
+ This operator takes care of reshaping and permuting the output from the
+ chunk dimension to the sequence dimension.
+
+ Args:
+ config (RetroConfig): Retro config.
+ """
+
+ def __init__(self, config: RetroConfig):
+ super().__init__(config=config)
+ self.retro_chunk_length = config.retro_chunk_length
+
+ @classmethod
+ def _forward(
+ cls,
+ x_with_bias: dict,
+ residual: Tensor,
+ prob: float,
+ retro_chunk_length: int,
+ bias_dropout_add: Callable,
+ ) -> Tensor:
+ """Per-chunk bias-dropout-add.
+
+ Args:
+ x_with_bias (dict): Attention output and bias, along with other Retro relevant parameters.
+ residual (Tensor): Transformer layer residual.
+ prob (float): Dropout probability.
+ retro_chunk_length (int): Retro chunk length (e.g., 64).
+ bias_dropout_add (Callable): Bias-dropout-add function.
+
+ Returns:
+ Output of bias-dropout-add.
+ """
+
+ # Extract input dict.
+ ns = x_with_bias["ns"]
+ bs = x_with_bias["bs"]
+ d = x_with_bias["d"]
+ l = x_with_bias["l"]
+ pad = x_with_bias["pad"]
+ attention_output = x_with_bias["attention_output"] # [ m, bs*l, d ]
+ attention_bias = x_with_bias["attention_bias"] # [ d ]
+
+ # Re-enable torch grad to enable fused optimization.
+ with torch.enable_grad():
+
+ # Bias-dropout-add.
+ x = bias_dropout_add(
+ (
+ attention_output,
+ None if attention_bias is None else attention_bias.expand_as(attention_output),
+ ),
+ torch.zeros_like(attention_output),
+ prob,
+ )
+
+ # Permute chunks back to sequence dimension.
+ # 1. [ m, bs*l, d ]
+ # 2. [ m, bs, l, d ]
+ # 3. [ l, m, bs, d ]
+ # 4. [ m*l, bs, d ] == [ ns, bs, d ]
+ x = (
+ x.reshape(retro_chunk_length, bs, l, d)
+ .permute(2, 0, 1, 3)
+ .reshape(retro_chunk_length * l, bs, d)
+ )
+
+ # Prepend zeros for non-attending tokens.
+ x = torch.nn.functional.pad(x, (0, 0, 0, 0, pad, 0), 'constant', 0)[
+ :ns
+ ] # [ ns, bs, d ]
+
+ # Add residual. [ ns, bs, d ]
+ x = x + residual
+
+ # Output. [ ns, bs, d ]
+ return x
+
+ def forward(self, training: bool, fused: bool) -> partial:
+ """Retro decoder bias-dropout-add.
+
+ Args:
+ training (bool): If training, then apply dropout.
+ fused (bool): Fuse bias-dropout-add.
+
+ Returns:
+ The partial function for performing bias-dropout-add.
+ """
+ return partial(
+ self._forward,
+ retro_chunk_length=self.retro_chunk_length,
+ bias_dropout_add=get_bias_dropout_add(training, fused),
+ )
diff --git a/megatron/core/models/retro/decoder_spec.py b/megatron/core/models/retro/decoder_spec.py
new file mode 100644
index 0000000000..2ad234b96b
--- /dev/null
+++ b/megatron/core/models/retro/decoder_spec.py
@@ -0,0 +1,185 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Specs for Retro decoder."""
+
+import typing
+
+from megatron.core import parallel_state
+from megatron.core.models.gpt.gpt_layer_specs import (
+ get_gpt_layer_local_spec,
+ get_gpt_layer_with_transformer_engine_spec,
+)
+from megatron.core.models.retro.config import RetroConfig
+from megatron.core.models.retro.decoder_attention import (
+ RetroDecoderBiasDropoutAdd,
+ RetroDecoderCrossAttention,
+)
+from megatron.core.models.retro.encoder_spec import get_retro_encoder_block_spec
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer import ModuleSpec
+from megatron.core.transformer.attention import CrossAttentionSubmodules
+from megatron.core.transformer.dot_product_attention import DotProductAttention
+from megatron.core.transformer.transformer_block import (
+ TransformerBlockSubmodules,
+ get_num_layers_to_build,
+)
+
+try:
+ import apex # pylint: disable=unused-import
+
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+
+ HAVE_APEX = True
+ LNImpl = FusedLayerNorm
+except ImportError:
+ import warnings
+
+ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm
+
+ warnings.warn(f'Apex is not installed. Falling back to Torch LayerNorm')
+ LNImpl = WrappedTorchLayerNorm
+
+try:
+ from megatron.core.extensions.transformer_engine import (
+ TEColumnParallelLinear,
+ TEDotProductAttention,
+ TENorm,
+ TERowParallelLinear,
+ )
+
+ HAVE_TE = True
+except ImportError:
+ HAVE_TE = False
+
+
+def get_retro_decoder_layer_te_spec(
+ encoder_block_spec: typing.Union[ModuleSpec, TransformerBlockSubmodules, None] = None
+) -> ModuleSpec:
+ """Retro decoder TE spec (uses Transformer Engine components).
+
+ A Retro decoder layer uses custom attention and bias-dropout-add operators
+ to perform chunked-cross attention. Additionally, the first Retro decoder
+ layer instantiates an entire encoder transformer block. As such, the decoder
+ cross attention module takes an optional encoder block spec, which is only
+ provided for the first Retro decoder layer.
+
+ Args:
+ encoder_block_spec (ModuleSpec): Retro encoder block spec, to be provided for
+ the first Retro decoder layer.
+
+ Returns:
+ A module spec with Transformer Engine modules.
+ """
+ spec = get_gpt_layer_with_transformer_engine_spec()
+ spec.submodules.pre_cross_attn_layernorm = TENorm
+ spec.submodules.cross_attention = ModuleSpec(
+ module=RetroDecoderCrossAttention,
+ params={"encoder_block_spec": encoder_block_spec},
+ submodules=CrossAttentionSubmodules(
+ linear_q=TEColumnParallelLinear,
+ linear_kv=TEColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ ),
+ )
+ spec.submodules.cross_attn_bda = ModuleSpec(module=RetroDecoderBiasDropoutAdd)
+ return spec
+
+
+def get_retro_decoder_layer_local_spec(
+ encoder_block_spec: typing.Optional[ModuleSpec] = None,
+) -> ModuleSpec:
+ """Retro decoder local spec (uses Megatron-Core components).
+
+ A Retro decoder layer uses custom attention and bias-dropout-add operators
+ to perform chunked-cross attention. Additionally, the first Retro decoder
+ layer instantiates an entire encoder transformer block. As such, the decoder
+ cross attention module takes an optional encoder block spec, which is only
+ provided for the first Retro decoder layer.
+
+ Args:
+ encoder_block_spec (ModuleSpec): Retro encoder block spec, to be provided
+ for the first Retro decoder layer.
+
+ Returns:
+ A module spec with local modules.
+ """
+ spec = get_gpt_layer_local_spec()
+ spec.submodules.pre_cross_attn_layernorm = LNImpl
+ spec.submodules.cross_attention = ModuleSpec(
+ module=RetroDecoderCrossAttention,
+ params={"encoder_block_spec": encoder_block_spec},
+ submodules=CrossAttentionSubmodules(
+ linear_q=ColumnParallelLinear,
+ linear_kv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ ),
+ )
+ spec.submodules.cross_attn_bda = ModuleSpec(module=RetroDecoderBiasDropoutAdd)
+ return spec
+
+
+def get_retro_decoder_block_spec(
+ config: RetroConfig, use_transformer_engine: bool
+) -> TransformerBlockSubmodules:
+ """Retro decoder block spec.
+
+ Retro decoder block implementation details:
+ - The retro decoder block consists of interleaved GPT layers
+ and customized Retro decoder layers.
+ - The Retro decoder layers are spaced three layers apart,
+ and start on layer 6 or 9 (depending on the total number of layers).
+ - The first decoder layer instantiates an encoder block,
+ and it therefore passes in an encoder_block_spec.
+
+ Args:
+ config (RetroConfig): Retro config.
+ use_transformer_engine (bool): If True, use Transformer Engine (instead of local modules.
+
+ Returns:
+ Transformer block submodules for the given spec.
+ """
+
+ # Num layers.
+ assert (
+ parallel_state.get_pipeline_model_parallel_world_size() == 1
+ ), "retro does not currently support pipeline parallelism."
+ assert (
+ parallel_state.get_virtual_pipeline_model_parallel_world_size() is None
+ ), "retro does not currently support virtual pipeline parallelism."
+ num_layers = get_num_layers_to_build(config)
+
+ # Retro layer numbers.
+ retro_layer_start = 6 if num_layers <= 15 else 9
+ retro_layer_numbers = list(range(retro_layer_start, num_layers + 1, 3))
+
+ # Layer specs.
+ gpt_layer_spec = (
+ get_gpt_layer_with_transformer_engine_spec()
+ if use_transformer_engine
+ else get_gpt_layer_local_spec()
+ )
+ get_retro_decoder_layer_spec = (
+ get_retro_decoder_layer_te_spec
+ if use_transformer_engine
+ else get_retro_decoder_layer_local_spec
+ )
+ retro_layer_spec = get_retro_decoder_layer_spec()
+ retro_layer_spec_with_retriever = get_retro_decoder_layer_spec(
+ get_retro_encoder_block_spec(config, use_transformer_engine)
+ )
+
+ layer_specs = []
+ for layer_number in range(1, num_layers + 1):
+ if layer_number == retro_layer_numbers[0]:
+ layer_specs.append(retro_layer_spec_with_retriever)
+ elif layer_number in retro_layer_numbers:
+ layer_specs.append(retro_layer_spec)
+ else:
+ layer_specs.append(gpt_layer_spec)
+
+ # Block spec.
+ block_spec = TransformerBlockSubmodules(layer_specs=layer_specs)
+
+ return block_spec
diff --git a/megatron/core/models/retro/encoder_attention.py b/megatron/core/models/retro/encoder_attention.py
new file mode 100644
index 0000000000..76625abe33
--- /dev/null
+++ b/megatron/core/models/retro/encoder_attention.py
@@ -0,0 +1,226 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Retro's cross attention modules for the encoder block."""
+
+from functools import partial
+from typing import Callable, List, Optional, Tuple, Type
+
+import torch
+from torch import Tensor
+
+from megatron.core import InferenceParams
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.models.retro.base_attention import BaseRetroCrossAttention
+from megatron.core.models.retro.config import RetroConfig
+from megatron.core.models.retro.utils import get_all_true_mask
+from megatron.core.transformer.module import MegatronModule
+
+
+class RetroEncoderCrossAttention(BaseRetroCrossAttention):
+ """Retro encoder's cross attention operator.
+
+ See this paper for more details: https://arxiv.org/abs/2112.04426.
+ Neighboring chunks are retrieved from the chunk database, encoded, and
+ used by the decoder layers for chunked cross attention.
+
+ Args:
+ config (RetroConfig): Retro config.
+ submodules (CrossAttentionSubmodules): Cross attention submodules.
+ layer_number (int): Layer number within transformer block.
+ attn_mask_type (AttnMaskType): Mask type ('causal' or 'padding').
+ """
+
+ def forward(
+ self,
+ hidden_states: Tensor,
+ attention_mask: Tensor,
+ key_value_states: Tensor = None,
+ inference_params: InferenceParams = None,
+ # rotary_pos_emb: Tensor = None, # unsupported for retro.
+ ) -> List[Tuple[Tensor, Optional[Tensor], Tensor]]:
+ """Cross attention for Retro encoder.
+
+ Notation:
+ ns : Sequence length.
+ bs : Batch size.
+ d : Hidden size.
+ l : Number of chunks per sample (i.e., seq_length/chunk_length).
+ k : Number of neighbors.
+ r : Number of retrieved tokens (neighbors + continuation).
+
+ Args:
+ hidden_states (Tensor): Transformer layer hidden states.
+ attention_mask (Tensor): Attention mask.
+ key_value_states (Tensor): Neighbor embeddings.
+ inference_params (InferenceParams): Inference params.
+
+ Returns:
+ List of tuples, where each tuple is (attention_output, attention_bias, residual).
+ """
+
+ # Input shape. [ r, bs*l*k, d ]
+ ns, bs, d = hidden_states.shape
+
+ # Reshape sequence into neighboring chunks.
+ # - hidden_states: [ r, bs*l*k, d ]
+ # - chunked_outputs: [ r, bs*l, k, d ]
+ chunked_outputs = hidden_states.reshape(
+ self.retro_retrieved_length, -1, self.retro_num_neighbors, d
+ )
+
+ # flash attn: [ b, h, sq, sk ]
+ # fused attn: [ b, 1, 1, sq ]
+ chunked_output_mask = get_all_true_mask(
+ size=(1, 1, chunked_outputs.shape[0], key_value_states.shape[0]),
+ device=chunked_outputs.device,
+ )
+
+ # Per-chunk attention.
+ attention_output_tuples = []
+ for k in range(self.retro_num_neighbors):
+
+ # Attend to current neighboring chunks.
+ # - chunked_output: [ r, bs*l, d ]
+ # - key_value_states: [ m, bs*l, d ]
+ # - attention_output: [ r, bs*l, d ]
+ # - attention_bias: [ d ]
+ chunked_output = chunked_outputs[:, :, k].contiguous()
+ attention_output, attention_bias = self.attn(
+ hidden_states=chunked_output, # Q (neighbor embedding)
+ attention_mask=chunked_output_mask,
+ key_value_states=key_value_states, # K, V (hidden act)
+ )
+
+ # Residual connection. [ r, bs*l, d ]
+ residual = chunked_output
+
+ # Collect tensors.
+ attention_output_tuples.append((attention_output, attention_bias, residual))
+
+ # Output. (List[Tuple[( [ r, bs*l, d ], [ d ] )]])
+ return attention_output_tuples
+
+
+class RetroEncoderBiasDropoutAdd(MegatronModule):
+ """Retro encoder's bias-dropout-add operator.
+
+ This operator applies bias-dropout-add individually on each neighboring
+ chunk that is retrieved from the chunk database.
+
+ Args:
+ config (RetroConfig): Retro config.
+ """
+
+ def __init__(self, config: RetroConfig):
+ super().__init__(config=config)
+ self.retro_num_neighbors = config.retro_num_neighbors
+
+ @classmethod
+ def _forward(
+ cls,
+ x_with_bias: List[Tuple[Tensor, Optional[Tensor], Tensor]],
+ residual: Tensor,
+ prob: float,
+ retro_num_neighbors: int,
+ bias_dropout_add: Callable,
+ ) -> Tensor:
+ """Per-chunk bias-dropout-add.
+
+ Args:
+ x_with_bias (dict): Attention output and bias tuple.
+ residual (Tensor): Transformer layer residual.
+ prob (float): Dropout probability.
+ retro_num_neighbors (int): Number of retrieved neighbor chunks (e.g., 2).
+ bias_dropout_add (Callable): Bias-dropout-add function.
+
+ Returns:
+ Output of bias-dropout-add.
+ """
+
+ # Re-enable torch grad to enable fused optimization.
+ with torch.enable_grad():
+
+ # Per-neighbor bias-dropout-add.
+ # - attention_output: [ r, bs*l, d ]
+ # - attention_bias: [ d ]
+ # - residual: [ r, bs*l, d ]
+ # - output: [ r, bs*l, d ]
+ outputs = [
+ bias_dropout_add(
+ (
+ attention_output,
+ None if attention_bias is None else attention_bias.expand_as(residual),
+ ),
+ residual,
+ prob,
+ )
+ for attention_output, attention_bias, residual in x_with_bias
+ ]
+
+ # Concatenate outputs (to shape [r, k*bs*l, d]; see notation above).
+ r, _, d = outputs[0].shape
+ output = torch.stack(outputs, dim=1).reshape(r, -1, d)
+
+ # Output. [ r, k*bs*l, d ]
+ return output
+
+ def forward(self, training: bool, fused: bool) -> partial:
+ """Retro decoder bias-dropout-add.
+
+ Args:
+ training (bool): If training, then apply dropout.
+ fused (bool): Fuse bias-dropout-add.
+
+ Returns:
+ A partial function for performing bias-dropout-add.
+ """
+ return partial(
+ self._forward,
+ retro_num_neighbors=self.retro_num_neighbors,
+ bias_dropout_add=get_bias_dropout_add(training, fused),
+ )
+
+
+class RetroEncoderLayerNorm(MegatronModule):
+ """Retro encoder's layernorm operator.
+
+ This operator applies layernorm individually on each neighboring chunk that
+ is retrieved from the chunk database, and then concatenates the chunks into
+ a single tensor.
+
+ Args:
+ config (RetroConfig): Retro config.
+ submodules (Type): Layer norm class. (Named 'submodules' to fit external interface.)
+ """
+
+ def __init__(self, config: RetroConfig, submodules: Type, **kwargs: dict):
+ super().__init__(config=config)
+ norm_class = submodules
+ self.norm = norm_class(config=config, **kwargs)
+ self.retro_num_neighbors = config.retro_num_neighbors
+
+ def forward(self, input: Tensor) -> Tensor:
+ """Per-chunk layer norm.
+
+ Args:
+ input (Tensor): Input chunks, concatenated into a single tensor.
+
+ Returns:
+ Output of the layer norm.
+ """
+
+ # Input shape: [ r, k*bs*l, d ]. (see notation above in attention module)
+
+ # Split input into 'num_neighbors' tensors.
+ chunk_size = input.shape[1] // self.retro_num_neighbors
+ inputs = torch.split(input, chunk_size, dim=1)
+
+ # Norm.
+ outputs = [self.norm(inp.contiguous()) for inp in inputs]
+
+ # Concatenate layer norms (to shape [r, k*bs*l, d]; see notation above).
+ r, _, d = inputs[0].shape
+ output = torch.stack(outputs, dim=1).reshape(r, -1, d)
+
+ # Output. [ r, k*bs*l, d ]
+ return output
diff --git a/megatron/core/models/retro/encoder_spec.py b/megatron/core/models/retro/encoder_spec.py
new file mode 100644
index 0000000000..b8a969bd84
--- /dev/null
+++ b/megatron/core/models/retro/encoder_spec.py
@@ -0,0 +1,168 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Specs for Retro encoder."""
+
+from megatron.core.models.gpt.gpt_layer_specs import (
+ get_gpt_layer_local_spec,
+ get_gpt_layer_with_transformer_engine_spec,
+)
+from megatron.core.models.retro.config import RetroConfig
+from megatron.core.models.retro.encoder_attention import (
+ RetroEncoderBiasDropoutAdd,
+ RetroEncoderCrossAttention,
+ RetroEncoderLayerNorm,
+)
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer import ModuleSpec
+from megatron.core.transformer.attention import CrossAttentionSubmodules
+from megatron.core.transformer.dot_product_attention import DotProductAttention
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
+
+try:
+ from megatron.core.extensions.transformer_engine import (
+ TEColumnParallelLinear,
+ TEDotProductAttention,
+ TENorm,
+ TERowParallelLinear,
+ )
+
+ HAVE_TE = True
+except ImportError:
+ HAVE_TE = False
+
+try:
+ import apex # pylint: disable=unused-import
+
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+
+ HAVE_APEX = True
+ LNImpl = FusedLayerNorm
+except ImportError:
+ import warnings
+
+ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm
+
+ warnings.warn(f'Apex is not installed. Falling back to Torch LayerNorm')
+ LNImpl = WrappedTorchLayerNorm
+
+
+def get_retro_encoder_layer_te_spec() -> ModuleSpec:
+ """Retro encoder TE spec (uses Transformer Engine components).
+
+ A Retro encoder layer uses custom attention, bias-dropout-add, and layernorm
+ operators to encode neighboring chunks that are retrieved from the chunk
+ database. Each operator is responsible for iterating the retrieved chunks
+ and processing them individually.
+
+ Returns:
+ A module spec if Transformer Engine modules.
+ """
+ spec = get_gpt_layer_with_transformer_engine_spec()
+ spec.submodules.pre_cross_attn_layernorm = TENorm
+ spec.submodules.cross_attention = ModuleSpec(
+ module=RetroEncoderCrossAttention,
+ params={"attn_mask_type": AttnMaskType.padding},
+ submodules=CrossAttentionSubmodules(
+ linear_q=TEColumnParallelLinear,
+ linear_kv=TEColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ ),
+ )
+ spec.submodules.cross_attn_bda = ModuleSpec(module=RetroEncoderBiasDropoutAdd)
+ spec.submodules.pre_mlp_layernorm = ModuleSpec(module=RetroEncoderLayerNorm, submodules=TENorm)
+ spec.submodules.mlp = ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(linear_fc1=TEColumnParallelLinear, linear_fc2=TERowParallelLinear),
+ )
+ return spec
+
+
+def get_retro_encoder_layer_local_spec() -> ModuleSpec:
+ """Retro encoder local spec (uses Megatron-Core components).
+
+ A Retro encoder layer uses custom attention, bias-dropout-add, and layernorm
+ operators to encode neighboring chunks that are retrieved from the chunk
+ database. Each operator is responsible for iterating the retrieved chunks
+ and processing them individually.
+
+ Returns:
+ A module spec if local modules.
+ """
+ spec = get_gpt_layer_local_spec()
+ spec.submodules.pre_cross_attn_layernorm = LNImpl
+ spec.submodules.cross_attention = ModuleSpec(
+ module=RetroEncoderCrossAttention,
+ params={"attn_mask_type": AttnMaskType.padding},
+ submodules=CrossAttentionSubmodules(
+ linear_q=ColumnParallelLinear,
+ linear_kv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ ),
+ )
+ spec.submodules.cross_attn_bda = ModuleSpec(module=RetroEncoderBiasDropoutAdd)
+ spec.submodules.pre_mlp_layernorm = ModuleSpec(module=RetroEncoderLayerNorm, submodules=LNImpl)
+ spec.submodules.mlp = ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear),
+ )
+ spec.submodules.sharded_state_dict_keys_map = {
+ 'input_layernorm.': 'self_attention.linear_qkv.layer_norm_'
+ } # pre_mlp_layernorm doesn't need remapping
+ return spec
+
+
+def get_retro_encoder_block_spec(
+ config: RetroConfig, use_transformer_engine: bool
+) -> TransformerBlockSubmodules:
+ """Retro encoder block spec.
+
+ The retro encoder block consists of one customized Retro encoder layer
+ (layer 1), and all of the following layers are standard GPT layers.
+
+ Args:
+ config (RetroConfig): Retro config.
+ use_transformer_engine (bool): If True, use Transformer Engine (instead of local modules).
+
+ Returns:
+ Transformer block submodules for the given spec.
+ """
+
+ # Num layers.
+ num_layers = config.retro_encoder_num_layers
+ retro_layer_numbers = [1]
+
+ # Layer specs.
+ gpt_layer_spec = (
+ get_gpt_layer_with_transformer_engine_spec()
+ if use_transformer_engine
+ else get_gpt_layer_local_spec()
+ )
+ get_retro_encoder_layer_spec = (
+ get_retro_encoder_layer_te_spec
+ if use_transformer_engine
+ else get_retro_encoder_layer_local_spec
+ )
+ retro_layer_spec = get_retro_encoder_layer_spec()
+ for spec in (gpt_layer_spec, retro_layer_spec):
+ spec.params["hidden_dropout"] = config.retro_encoder_hidden_dropout
+ spec.submodules.self_attention.params["attn_mask_type"] = AttnMaskType.padding
+ spec.submodules.self_attention.submodules.core_attention = ModuleSpec(
+ module=TEDotProductAttention if use_transformer_engine else DotProductAttention,
+ params={"attention_dropout": config.retro_encoder_attention_dropout},
+ )
+
+ layer_specs = []
+ for layer_number in range(1, num_layers + 1):
+ if layer_number in retro_layer_numbers:
+ layer_specs.append(retro_layer_spec)
+ else:
+ layer_specs.append(gpt_layer_spec)
+
+ # Block spec.
+ block_spec = TransformerBlockSubmodules(layer_specs=layer_specs)
+
+ return block_spec
diff --git a/megatron/core/models/retro/model.py b/megatron/core/models/retro/model.py
new file mode 100644
index 0000000000..8142c91f7a
--- /dev/null
+++ b/megatron/core/models/retro/model.py
@@ -0,0 +1,99 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Retro Model."""
+from typing import Dict, Optional
+
+from torch import Tensor
+
+from megatron.core import InferenceParams
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.models.gpt import GPTModel
+
+
+class RetroModel(GPTModel):
+ """Retro Model.
+
+ A Retro model mostly re-uses the GPTModel interface, with the only difference
+ being the embedding of the 'context' this is used by Retro for processing
+ neighbor tokens. This embedded context is then forwarded to the Transformer
+ Block.
+ """
+
+ def forward(
+ self,
+ input_ids: Tensor,
+ position_ids: Tensor,
+ attention_mask: Tensor,
+ context_input_ids: Tensor = None,
+ context_position_ids: Tensor = None,
+ context_mask: Tensor = None,
+ decoder_input: Tensor = None,
+ labels: Tensor = None,
+ inference_params: InferenceParams = None,
+ ) -> Tensor:
+ """RetroModel forward method.
+
+ Foward input tokens & mask, along with neighbor tokens & mask, through
+ the Retro model..
+
+ Args:
+ input_ids (Tensor): Input token IDs.
+ position_ids (Tensor): Input position IDs.
+ attention_mask (Tensor): Input attention mask.
+ context_input_ids (Tensor): Context (i.e., neighbor) token IDs.
+ context_position_ids (Tensor): Context (i.e., neighbor) position IDs.
+ context_mask (Tensor): Context (i.e., neighbor) attention mask.
+ decoder_input (Tensor): When using pipeline parallelism, input_ids and position_ids will only be used on the first stage, and for all other stages decoder_input will be provided via communication from the previous stage.
+ labels (Tensor): The labels of dimension [batch size, seq length].
+ inference_params (InferenceParams): Parameters for inference.
+
+ Returns:
+ Output tensor of forward pass.
+ """
+
+ # Argument shapes:
+ # Notation:
+ # ns : Sequence length.
+ # bs : Batch size.
+ # d : Hidden size.
+ # l : Number of chunks per sample (i.e., seq_length/chunk_length).
+ # k : Number of neighbors.
+ # r : Number of retrieved tokens (neighbors + continuation).
+ # - input_ids: [ bs, ns ]
+ # - context_ids: [ k*bs*l, r ]
+ # - context: [ r, k*bs*l, d ]
+ # - output: [ ns, bs, d ]
+
+ # Context embedding (e.g., for Retro neighbor tokens).
+ if context_input_ids is not None:
+ context = self.embedding(context_input_ids, context_position_ids)
+ else:
+ context = None
+
+ # Call GPTModel.forward, and pass in embedded context.
+ return super().forward(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ decoder_input=decoder_input,
+ labels=labels,
+ inference_params=inference_params,
+ extra_block_kwargs={"context": context, "context_mask": context_mask},
+ )
+
+ def sharded_state_dict(
+ self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[Dict] = None
+ ) -> ShardedStateDict:
+ """Get sharded state dict.
+
+ Args:
+ prefix (str): Module name prefix.
+ sharded_offsets (tuple): Offsets of local shard within global tensor.
+ metadata (Optional[Dict]): Shard metadata.
+
+ Returns:
+ A ?
+ """
+ metadata = metadata or {}
+ metadata['non_homogeneous_layers'] = True
+ return super().sharded_state_dict(prefix, sharded_offsets, metadata)
diff --git a/megatron/core/models/retro/utils.py b/megatron/core/models/retro/utils.py
new file mode 100644
index 0000000000..7d83c5d306
--- /dev/null
+++ b/megatron/core/models/retro/utils.py
@@ -0,0 +1,24 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import os
+
+import torch
+
+
+def get_config_path(project_dir: str) -> str:
+ """Config copy stored within retro project dir."""
+ return os.path.join(project_dir, "config.json")
+
+
+def get_gpt_data_dir(project_dir: str) -> str:
+ """Get project-relative directory of GPT bin/idx datasets."""
+ return os.path.join(project_dir, "data")
+
+
+# ** Note ** : Retro's compatibility between cross attention and Flash/Fused
+# Attention is currently a work in progress. We default to returning None for
+# now.
+# def get_all_true_mask(size, device):
+# return torch.full(size=size, fill_value=True, dtype=torch.bool, device=device)
+def get_all_true_mask(size, device):
+ return None
diff --git a/megatron/core/models/vision/__init__.py b/megatron/core/models/vision/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/megatron/core/models/vision/clip_vit_model.py b/megatron/core/models/vision/clip_vit_model.py
new file mode 100644
index 0000000000..53c3feddee
--- /dev/null
+++ b/megatron/core/models/vision/clip_vit_model.py
@@ -0,0 +1,190 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Optional, Union
+
+import torch
+
+from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
+from megatron.core.extensions.transformer_engine import TENorm
+from megatron.core.models.common.vision_module.vision_module import VisionModule
+from megatron.core.transformer.enums import ModelType
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+# Note: This is under development and is missing features like position embedding interpolation.
+class CLIPViTModel(VisionModule):
+ """CLIP ViT vision model.
+
+ Args:
+ transformer_config (TransformerConfig): Transformer config.
+ transformer_layer_spec (ModuleSpec): Specifies module to use for transformer layers.
+ ln_pre_impl (ModuleSpec or type): Specifies the layer norm type to use for ln_pre.
+ add_class_token (bool, optional): Include a class token. Defaults to True.
+ class_token_len (int): Class token length. Defaults to 1 but 8 may be faster.
+ patch_dim (int): Image patch size.
+ img_h (int): Input image height.
+ img_w (int): Input image width.
+ """
+
+ def __init__(
+ self,
+ transformer_config: TransformerConfig,
+ transformer_layer_spec: ModuleSpec,
+ ln_pre_impl: Union[ModuleSpec, type] = TENorm,
+ ln_post_impl: Union[ModuleSpec, type] = TENorm,
+ add_class_token: bool = True,
+ class_token_len: int = 1,
+ patch_dim: int = 14,
+ img_h: int = 336,
+ img_w: int = 336,
+ model_subtype: str = "clip",
+ ) -> None:
+
+ error_msg = f"CLIPViTModel model subtype {model_subtype} is not supported."
+ assert model_subtype in ["clip", "siglip"], error_msg
+
+ if model_subtype == "siglip":
+ assert class_token_len == 0, "SigLIP does not support class tokens."
+ assert not add_class_token, "SigLIP does not support class tokens."
+
+ super().__init__(config=transformer_config)
+
+ if has_config_logger_enabled(transformer_config):
+ log_config_to_disk(transformer_config, locals(), prefix=type(self).__name__)
+
+ self.class_token_len = class_token_len
+ self.visual_hidden_size = transformer_config.hidden_size
+ self.patch_dim = patch_dim
+ self.img_h = img_h
+ self.img_w = img_w
+
+ assert self.img_h % self.patch_dim == 0
+ assert self.img_w % self.patch_dim == 0
+ self.num_patches_per_dim_h = self.img_h // self.patch_dim
+ self.num_patches_per_dim_w = self.img_w // self.patch_dim
+ self.num_patches = self.num_patches_per_dim_h * self.num_patches_per_dim_w
+
+ self.add_class_token = add_class_token
+ self.class_token_len = class_token_len
+
+ self.seq_length = self.num_patches + (self.class_token_len if self.add_class_token else 0)
+
+ self.ln_pre = None
+ self.ln_post = None
+ if model_subtype == "clip":
+ self.ln_pre = build_module(
+ ln_pre_impl,
+ config=transformer_config,
+ hidden_size=self.visual_hidden_size,
+ eps=transformer_config.layernorm_epsilon,
+ )
+ conv_bias = False
+ padding = 0
+ if model_subtype == "siglip":
+ self.ln_post = build_module(
+ ln_post_impl,
+ config=transformer_config,
+ hidden_size=self.visual_hidden_size,
+ eps=transformer_config.layernorm_epsilon,
+ )
+ conv_bias = True
+ padding = "valid"
+
+ self.conv1 = torch.nn.Conv2d(
+ in_channels=3,
+ out_channels=self.visual_hidden_size,
+ kernel_size=self.patch_dim,
+ stride=self.patch_dim,
+ bias=conv_bias,
+ padding=padding,
+ )
+
+ self.position_ids = torch.arange(self.seq_length).expand(1, -1).cuda()
+
+ self.position_embeddings = torch.nn.Embedding(self.seq_length, self.visual_hidden_size)
+
+ self.add_class_token = add_class_token
+ if self.add_class_token:
+ self.class_token = torch.nn.Parameter(
+ torch.randn(1, self.class_token_len, self.visual_hidden_size)
+ )
+
+ self.model_type = ModelType.encoder_or_decoder
+
+ # Transformer layers.
+ # TODO: Make pre_process and post_process configurable.
+ # NOTE: a final layer norm and/or linear layer in some implementations are omitted here.
+ # They can be added separately where needed.
+ self.decoder = TransformerBlock(
+ config=transformer_config,
+ spec=transformer_layer_spec,
+ pre_process=True,
+ post_process=False,
+ )
+
+ def set_input_tensor(self, input_tensor: torch.Tensor) -> None:
+ """Sets input tensor to the model.
+
+ Args:
+ input_tensor (Tensor): Sets the input tensor for the model.
+ """
+ self.decoder.set_input_tensor(input_tensor)
+
+ def forward(
+ self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+ ) -> torch.Tensor:
+ """Forward function of the CLIP ViT Model. This function passes the input tensors
+ through the embedding layer and then the transformer.
+
+ Args:
+ x (torch.Tensor): input data of shape [batch, img_h, img_w]
+ attention_mask (torch.Tensor with dtype=bool): Attention mask to use.
+
+ Returns:
+ x (torch.Tensor): output after final transformer block of shape [b, s, h].
+ """
+ x = self.conv1(x) # shape = [batch, hidden_size, grid, grid]
+ x = x.reshape(x.shape[0], x.shape[1], -1) # [batch, hidden_size, grid ** 2]
+ x = x.permute(0, 2, 1) # [batch, grid ** 2, hidden_size]
+
+ if self.add_class_token:
+ class_token = self.class_token.expand(
+ x.shape[0], -1, -1
+ ) # [batch, class_token_len, hidden_size]
+ x = torch.cat(
+ [class_token, x], dim=1
+ ) # [batch, grid ** 2 + class_token_len, hidden_size]
+
+ assert x.shape[1] == self.seq_length, f"{x.shape[1]} != {self.seq_length}"
+ x = x + self.position_embeddings(self.position_ids)
+ if self.ln_pre:
+ x = self.ln_pre(x)
+ x = x.permute(1, 0, 2) # [b, s, h] -> [s, b, h]
+ # `permute` can make the tensor non-contiguous, breaking pipelining.
+ x = x.contiguous()
+
+ x = self.decoder(x, attention_mask)
+ x = x.permute(1, 0, 2) # [s, b, h] -> [b, s, h]
+ x = x.contiguous()
+ if self.ln_post:
+ x = self.ln_post(x)
+ return x
+
+
+def get_num_image_embeddings(
+ img_h, img_w, patch_dim, vision_model_type, disable_vision_class_token, class_token_len
+):
+ """Get the number of image embeddings per image tile."""
+ if vision_model_type == "siglip":
+ keep_class_token = False
+ elif vision_model_type == "clip":
+ keep_class_token = not disable_vision_class_token
+
+ num_patches_per_dim_h = img_h // patch_dim
+ num_patches_per_dim_w = img_w // patch_dim
+ num_patches = num_patches_per_dim_h * num_patches_per_dim_w
+ num_image_embeddings_per_tile = num_patches + (class_token_len if keep_class_token else 0)
+
+ return num_image_embeddings_per_tile
diff --git a/megatron/core/models/vision/multimodal_projector.py b/megatron/core/models/vision/multimodal_projector.py
new file mode 100644
index 0000000000..18e62c68a5
--- /dev/null
+++ b/megatron/core/models/vision/multimodal_projector.py
@@ -0,0 +1,67 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from megatron.core import tensor_parallel
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import make_viewless_tensor
+
+
+class MultimodalProjector(MegatronModule):
+ """
+ MultimodalProjector will take the encoded input with input_size hidden state and project
+ it into the hidden size of the language model for multimodal training. When projector is
+ type affine linear_fc1 from submodules is used.
+
+ Args:
+ transformer_config (TransformerConfig): Transformer config
+ submodules (MLPSubmodules): Specifies MLP submodules for mlp type projector
+ projector_type (str): Projector type
+ input_size (int): Input size from feature encoder
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: MLPSubmodules,
+ projector_type: str,
+ input_size: int,
+ ):
+ super().__init__(config=config)
+ self.projector_type = projector_type
+
+ assert submodules is not None, "MLPSubmodules must be provided"
+
+ if self.projector_type == "mlp":
+ self.encoder = MLP(config=config, submodules=submodules, input_size=input_size)
+ elif self.projector_type == "affine":
+ self.encoder = build_module(
+ submodules.linear_fc1,
+ input_size,
+ config.hidden_size,
+ config=config,
+ init_method=config.init_method,
+ gather_output=True,
+ bias=config.add_bias_linear,
+ skip_bias_add=True,
+ is_expert=False,
+ tp_comm_buffer_name=None,
+ )
+ else:
+ raise Exception(f"Unsupported multimodal projection type {self.projector_type}")
+
+ def forward(self, hidden_states):
+ # Run encoder.
+ encoder_output, encoder_output_bias = self.encoder(hidden_states)
+
+ if encoder_output_bias is not None:
+ encoder_output = encoder_output + encoder_output_bias
+
+ # the encoder produces "viewed" tensor. This will result in schedule.py's
+ # deallocate_output_tensor() throwing an error, so a viewless tensor is
+ # created to prevent this.
+ encoder_output = make_viewless_tensor(
+ inp=encoder_output, requires_grad=True, keep_graph=True
+ )
+
+ return encoder_output
diff --git a/megatron/core/models/vision/vit_layer_specs.py b/megatron/core/models/vision/vit_layer_specs.py
new file mode 100644
index 0000000000..da9066b007
--- /dev/null
+++ b/megatron/core/models/vision/vit_layer_specs.py
@@ -0,0 +1,95 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.extensions.transformer_engine import (
+ TEDotProductAttention,
+ TELayerNormColumnParallelLinear,
+ TERowParallelLinear,
+)
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
+from megatron.core.transformer.dot_product_attention import DotProductAttention
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+
+try:
+ import apex # pylint: disable=unused-import
+
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+
+ HAVE_APEX = True
+ LNImpl = FusedLayerNorm
+except ImportError:
+ import warnings
+
+ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm
+
+ warnings.warn(f'Apex is not installed. Falling back to Torch LayerNorm')
+ LNImpl = WrappedTorchLayerNorm
+
+
+# Use this spec to use lower level Transformer Engine modules (required for fp8 training)
+def get_vit_layer_with_transformer_engine_spec() -> ModuleSpec:
+ '''
+ Returns ViT layer spec with Transformer Engine layers
+ '''
+ mlp = _get_mlp_module_spec(use_te=True)
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.no_mask},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=TELayerNormColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=IdentityOp,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+
+
+def get_vit_layer_with_local_spec() -> ModuleSpec:
+ '''
+ Returns ViT layer spec with Mcore local layers
+ '''
+ mlp = _get_mlp_module_spec(use_te=False)
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=LNImpl,
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=ColumnParallelLinear,
+ core_attention=DotProductAttention,
+ linear_proj=RowParallelLinear,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ pre_mlp_layernorm=LNImpl,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+
+
+# Helper function to get module spec for MLP/MoE
+def _get_mlp_module_spec(use_te: bool = True) -> ModuleSpec:
+ # Dense MLP w/ or w/o TE modules.
+ return ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=TELayerNormColumnParallelLinear if use_te else ColumnParallelLinear,
+ linear_fc2=TERowParallelLinear if use_te else RowParallelLinear,
+ ),
+ )
diff --git a/megatron/core/num_microbatches_calculator.py b/megatron/core/num_microbatches_calculator.py
new file mode 100644
index 0000000000..5850e512ca
--- /dev/null
+++ b/megatron/core/num_microbatches_calculator.py
@@ -0,0 +1,498 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""Megatron Core number of microbatches calculators."""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import List, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+# TODO: global_var merge into mcore?
+_GLOBAL_NUM_MICROBATCHES_CALCULATOR: Union[
+ 'ConstantNumMicroBatchesCalculator', 'RampupBatchsizeNumMicroBatchesCalculator'
+] = None
+
+
+def get_num_microbatches() -> int:
+ """Get number of microbatches."""
+ return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get()
+
+
+def get_current_global_batch_size() -> int:
+ """Get current global batch size."""
+ return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get_current_global_batch_size()
+
+
+def get_micro_batch_size() -> int:
+ """Get micro batch size."""
+ return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get_micro_batch_size()
+
+
+def get_current_running_global_batch_size() -> int:
+ """Get current running global batch size, taking into account number of DP replicas might be
+ incompatible with true global batch size if `decrease_batch_size_if_needed` is True."""
+ return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get_current_running_global_batch_size()
+
+
+def update_num_microbatches(
+ consumed_samples: int, consistency_check: bool = True, verbose: bool = False
+) -> None:
+ """Update number of microbatches.
+
+ Args:
+ consumed_samples (int):
+ Number of samples consumed.
+ consistency_check (bool, optional):
+ Option to check current schedule's consistency. Defaults to True.
+ verbose (bool, optional):
+ Option to control logging. Defaults to False.
+ """
+ _GLOBAL_NUM_MICROBATCHES_CALCULATOR.update(consumed_samples, consistency_check, verbose)
+
+
+def init_num_microbatches_calculator(
+ rank: int,
+ rampup_batch_size: Optional[List[int]],
+ global_batch_size: int,
+ micro_batch_size: int,
+ data_parallel_size: int,
+ decrease_batch_size_if_needed: bool = False,
+) -> None:
+ """Initialize number of microbatches calculator. Supporting backward compatibility.
+
+ Args:
+ rank (int):
+ Rank of the GPU, only rank 0 will log the information.
+ rampup_batch_size (Optional[List[int]]):
+ Rampup batch size, should be in format of [start_global_batch_size,
+ batch_size_increment, ramup_samples].
+ global_batch_size (int):
+ Global batch size for the model.
+ micro_batch_size (int):
+ Micro batch size at initialization.
+ data_parallel_size (int):
+ Data parallel size.
+ decrease_batch_size_if_needed (bool, optional):
+ If true, scale down batch size to ensure divisibility by DP size * microbatch size.
+ Defaults to False.
+ """
+ _configure_global_num_microbatches_calculator(
+ rank,
+ rampup_batch_size,
+ global_batch_size,
+ micro_batch_size,
+ data_parallel_size,
+ decrease_batch_size_if_needed,
+ init=True,
+ )
+
+
+def destroy_num_microbatches_calculator():
+ """Destroy number of microbatches calculator."""
+ global _GLOBAL_NUM_MICROBATCHES_CALCULATOR
+ _GLOBAL_NUM_MICROBATCHES_CALCULATOR = None
+
+
+def reconfigure_num_microbatches_calculator(
+ rank: int,
+ rampup_batch_size: Optional[List[int]],
+ global_batch_size: int,
+ micro_batch_size: int,
+ data_parallel_size: int,
+ decrease_batch_size_if_needed: bool = False,
+) -> None:
+ """Reconfigure number of microbatches calculator. Supporting backward compatibility.
+
+ Args:
+ rank (int):
+ Rank of the GPU, only rank 0 will log the information.
+ rampup_batch_size (Optional[List[int]]):
+ Rampup batch size, should be in format of
+ [start_global_batch_size, batch_size_increment, ramup_samples].
+ global_batch_size (int):
+ Global batch size for the model.
+ micro_batch_size (int):
+ Micro batch size at initialization.
+ data_parallel_size (int):
+ Data parallel size.
+ decrease_batch_size_if_needed (bool, optional):
+ If true, scale down batch size to ensure divisibility by DP size * microbatch size.
+ Defaults to False.
+ """
+ _configure_global_num_microbatches_calculator(
+ rank,
+ rampup_batch_size,
+ global_batch_size,
+ micro_batch_size,
+ data_parallel_size,
+ decrease_batch_size_if_needed,
+ init=False,
+ )
+
+
+def _configure_global_num_microbatches_calculator(
+ rank: int,
+ rampup_batch_size: Optional[List[int]],
+ global_batch_size: int,
+ micro_batch_size: int,
+ data_parallel_size: int,
+ decrease_batch_size_if_needed: bool = False,
+ init: bool = False,
+) -> None:
+ """Configure number of microbatches calculator. Can be used for initialization and
+ reconfiguration.
+
+ Args:
+ rank (int):
+ Rank of the GPU, only rank 0 will log the information.
+ rampup_batch_size (Optional[List[int]]):
+ Rampup batch size, should be in format of
+ [start_global_batch_size, batch_size_increment, ramup_samples].
+ global_batch_size (int):
+ Global batch size for the model.
+ micro_batch_size (int):
+ Micro batch size at initialization.
+ data_parallel_size (int):
+ Data parallel size.
+ decrease_batch_size_if_needed (bool, optional):
+ If true, scale down batch size to ensure divisibility by DP size * microbatch size.
+ Defaults to False.
+ init (bool, optional):
+ If true, initialize the calculator. Defaults to False.
+ """
+ global _GLOBAL_NUM_MICROBATCHES_CALCULATOR
+
+ if init:
+ assert (
+ _GLOBAL_NUM_MICROBATCHES_CALCULATOR is None
+ ), 'num microbatches calculator is already initialized.'
+
+ _GLOBAL_NUM_MICROBATCHES_CALCULATOR = _build_num_microbatches_calculator(
+ rank,
+ rampup_batch_size,
+ global_batch_size,
+ micro_batch_size,
+ data_parallel_size,
+ decrease_batch_size_if_needed,
+ )
+
+
+def _build_num_microbatches_calculator(
+ rank: int,
+ rampup_batch_size: Optional[List[int]],
+ global_batch_size: int,
+ micro_batch_size: int,
+ data_parallel_size: int,
+ decrease_batch_size_if_needed: bool,
+) -> Union['ConstantNumMicroBatchesCalculator', 'RampupBatchsizeNumMicroBatchesCalculator']:
+ """Build number of microbatches calculator. Internal helper method.
+
+ Args:
+ rank (int):
+ Rank of the GPU, only rank 0 will log the information.
+ rampup_batch_size (Optional[List[int]]):
+ Rampup batch size, should be in format of
+ [start_global_batch_size, batch_size_increment, ramup_samples].
+ global_batch_size (int):
+ Global batch size for the model.
+ micro_batch_size (int):
+ Micro batch size at initialization.
+ data_parallel_size (int):
+ Data parallel size.
+ decrease_batch_size_if_needed (bool):
+ If true, scale down batch size to ensure divisibility by DP size * microbatch size.
+
+ """
+
+ # Constant batch size.
+ if rampup_batch_size is None:
+ num_microbatches_calculator = ConstantNumMicroBatchesCalculator(
+ global_batch_size,
+ micro_batch_size,
+ data_parallel_size,
+ decrease_batch_size_if_needed,
+ rank,
+ )
+ if rank == 0:
+ logger.info(
+ f'setting number of microbatches to constant {num_microbatches_calculator.get()}'
+ )
+ # Batch size ramp up.
+ else:
+ assert len(rampup_batch_size) == 3, (
+ 'expected the following '
+ 'format: --rampup-batch-size '
+ ' '
+ )
+ start_global_batch_size = int(rampup_batch_size[0])
+ batch_size_increment = int(rampup_batch_size[1])
+ ramup_samples = int(rampup_batch_size[2])
+ if rank == 0:
+ logger.info(
+ f'will use batch size rampup starting from global batch size '
+ f'{start_global_batch_size} to global batch size {global_batch_size} with batch'
+ f'size increments {batch_size_increment} over {ramup_samples} samples.'
+ )
+ num_microbatches_calculator = RampupBatchsizeNumMicroBatchesCalculator(
+ global_batch_size,
+ micro_batch_size,
+ data_parallel_size,
+ decrease_batch_size_if_needed,
+ rank,
+ start_global_batch_size,
+ batch_size_increment,
+ ramup_samples,
+ )
+
+ return num_microbatches_calculator
+
+
+def _round(batch_size: int, divisor: int) -> int:
+ """Round `batch_size` down to nearest batch size divisible by `divisor`."""
+ return (batch_size // divisor) * divisor
+
+
+class NumMicroBatchesCalculator(ABC):
+ """Base class for number of microbatches calculator."""
+
+ def __init__(self) -> None:
+ self.num_micro_batches = None
+ self.current_global_batch_size = None
+ self.micro_batch_size = None
+ self.current_running_global_batch_size = None
+
+ def get(self) -> int:
+ """Get number of microbatches."""
+ return self.num_micro_batches
+
+ def get_current_global_batch_size(self) -> int:
+ """Get current global batch size."""
+ return self.current_global_batch_size
+
+ def get_micro_batch_size(self) -> int:
+ """Get current global batch size."""
+ return self.micro_batch_size
+
+ def get_current_running_global_batch_size(self) -> int:
+ """Get current running global batch size. If decrease_batch_size_if_needed is False,
+ this just equals global batch size."""
+ return self.current_running_global_batch_size
+
+ @abstractmethod
+ def update(self, consumed_samples, consistency_check, verbose=False) -> None:
+ """Update number of microbatches depending on batch size rampup."""
+ pass
+
+
+class ConstantNumMicroBatchesCalculator(NumMicroBatchesCalculator):
+ """Calculator of number of microbatches with constant global batch size.
+
+ Args:
+ global_batch_size (int):
+ Global batch size.
+ micro_batch_size (int):
+ Micro batch size.
+ data_parallel_size (int):
+ Data parallel size.
+ decrease_batch_size_if_needed (bool):
+ If true, decrease batch size to ensure divisibility by DP size * microbatch size
+ (if needed).
+ rank (int):
+ Rank (to determine whether logging should be performed).
+ """
+
+ def __init__(
+ self,
+ global_batch_size: int,
+ micro_batch_size: int,
+ data_parallel_size: int,
+ decrease_batch_size_if_needed: bool,
+ rank: int,
+ ) -> None:
+
+ micro_batch_times_data_parallel_size = micro_batch_size * data_parallel_size
+ if decrease_batch_size_if_needed:
+ running_global_batch_size = _round(
+ global_batch_size, micro_batch_times_data_parallel_size
+ )
+ assert running_global_batch_size % micro_batch_times_data_parallel_size == 0
+ if rank == 0:
+ logger.info(
+ f'decreasing batch size from {global_batch_size} to {running_global_batch_size}'
+ f'to keep divisiblity by micro_batch_size={micro_batch_size} * '
+ f'data_parallel_size={data_parallel_size}'
+ )
+ self.num_micro_batches = (
+ running_global_batch_size // micro_batch_times_data_parallel_size
+ )
+ else:
+ assert global_batch_size % micro_batch_times_data_parallel_size == 0, (
+ 'global batch size ({}) is not divisible by micro batch size ({})'
+ ' times data parallel size ({})'.format(
+ global_batch_size, micro_batch_size, data_parallel_size
+ )
+ )
+ running_global_batch_size = global_batch_size
+ self.num_micro_batches = global_batch_size // micro_batch_times_data_parallel_size
+ assert (
+ self.num_micro_batches >= 1
+ ), 'number of microbatches should be at least 1, got {}.'.format(self.num_micro_batches)
+
+ self.current_global_batch_size = global_batch_size
+ self.current_running_global_batch_size = running_global_batch_size
+ self.micro_batch_size = micro_batch_size
+
+ def update(self, consumed_samples, consistency_check, verbose=False) -> None:
+ pass
+
+
+class RampupBatchsizeNumMicroBatchesCalculator(NumMicroBatchesCalculator):
+ """Calculator of number of microbatches with batch size rampup.
+ Over `steps = (global-batch-size - start-batch-size) / batch_size_increment` increment batch
+ size from start-batch-size to global-batch-size using rampup-samples / steps
+ samples.
+
+ Args:
+ global_batch_size (int):
+ Global batch size post rampup.
+ micro_batch_size (int):
+ Micro batch size.
+ data_parallel_size (int):
+ Data parallel size.
+ decrease_batch_size_if_needed (bool):
+ If true, decrease batch size to ensure divisibility by DP size * microbatch size
+ (if needed).
+ rank (int):
+ Rank (to determine whether logging should be performed).
+ start_global_batch_size (int):
+ Global batch size to start with.
+ batch_size_increment (int):
+ Global batch size increments.
+ ramup_samples (int):
+ Number of samples to use ramp up global
+ batch size from `start_global_batch_size` to `global_batch_size`.
+ """
+
+ def __init__(
+ self,
+ global_batch_size: int,
+ micro_batch_size: int,
+ data_parallel_size: int,
+ decrease_batch_size_if_needed: bool,
+ rank: int,
+ start_global_batch_size: int,
+ batch_size_increment: int,
+ ramup_samples: int,
+ ) -> None:
+ assert global_batch_size > 0, 'global batch size should be positive, got {}.'.format(
+ global_batch_size
+ )
+ assert start_global_batch_size > 0, 'start batch size should be positive, got {}.'.format(
+ start_global_batch_size
+ )
+ assert batch_size_increment > 0, 'batch size increment should be positive, got {}.'.format(
+ batch_size_increment
+ )
+ assert ramup_samples >= 0, 'ramp-up samples should be non-negative, got {}.'.format(
+ ramup_samples
+ )
+
+ self.global_batch_size = global_batch_size
+ self.micro_batch_size = micro_batch_size
+ self.data_parallel_size = data_parallel_size
+ self.decrease_batch_size_if_needed = decrease_batch_size_if_needed
+ self.rank = rank
+ self.start_global_batch_size = start_global_batch_size
+ self.batch_size_increment = batch_size_increment
+ self.ramup_samples = ramup_samples
+
+ self.micro_batch_times_data_parallel_size = self.micro_batch_size * self.data_parallel_size
+ assert self.micro_batch_times_data_parallel_size > 0
+ self.current_global_batch_size = None
+
+ diff_batch_size = self.global_batch_size - self.start_global_batch_size
+ assert diff_batch_size >= 0, (
+ 'expected global batch size to be greater than or equal to start batch size, '
+ f'got {self.global_batch_size} and {self.start_global_batch_size}'
+ )
+ assert diff_batch_size % batch_size_increment == 0, (
+ 'expected '
+ f'global batch size interval ({diff_batch_size}) to be divisible by global batch '
+ f'size increment ({batch_size_increment})'
+ )
+
+ num_increments = diff_batch_size // self.batch_size_increment
+ self.rampup_samples_per_increment = self.ramup_samples / num_increments
+
+ # Initialize number of microbatches.
+ self.update(0, consistency_check=False, verbose=True)
+
+ def update(self, consumed_samples: int, consistency_check: bool, verbose: bool = False) -> None:
+ """Update number of microbatches.
+
+ Args:
+ consumed_samples (int): Number of samples consumed.
+ consistency_check (bool): Option to check current schedule's consistency.
+ verbose (bool, optional): Option to control logging. Defaults to False.
+ """
+
+ # Update current global batch size.
+ global_batch_size_changed = False
+ old_current_global_batch_size = self.current_global_batch_size
+ if consumed_samples > self.ramup_samples:
+ self.current_global_batch_size = self.global_batch_size
+ else:
+ steps = int(consumed_samples / self.rampup_samples_per_increment)
+ self.current_global_batch_size = (
+ self.start_global_batch_size + steps * self.batch_size_increment
+ )
+ assert self.current_global_batch_size <= self.global_batch_size
+
+ if old_current_global_batch_size != self.current_global_batch_size:
+ global_batch_size_changed = True
+ if self.rank == 0 and global_batch_size_changed and verbose:
+ if old_current_global_batch_size is None:
+ logger.info(f'setting initial batch size to {self.current_global_batch_size}')
+ else:
+ logger.info(
+ f'ramping up batch size from {old_current_global_batch_size} to '
+ f'{self.current_global_batch_size}'
+ )
+
+ # Check consistency of the current global batch size.
+ if consistency_check and not self.decrease_batch_size_if_needed:
+ assert (
+ self.current_global_batch_size % self.micro_batch_times_data_parallel_size == 0
+ ), (
+ 'current global '
+ 'batch size ({}) is not divisible by micro-batch-size ({}) times'
+ 'data parallel size ({})'.format(
+ self.current_global_batch_size, self.micro_batch_size, self.data_parallel_size
+ )
+ )
+
+ if (
+ self.decrease_batch_size_if_needed
+ and self.current_global_batch_size % self.micro_batch_times_data_parallel_size != 0
+ ):
+ self.current_running_global_batch_size = _round(
+ self.current_global_batch_size, self.micro_batch_times_data_parallel_size
+ )
+ if self.rank == 0 and global_batch_size_changed and verbose:
+ logger.info(
+ f'decreasing batch size from {self.current_global_batch_size} to '
+ f'{self.current_running_global_batch_size} to keep divisiblity by '
+ f'micro_batch_size={self.micro_batch_size} * '
+ f'data_parallel_size={self.data_parallel_size}'
+ )
+ assert (
+ self.current_running_global_batch_size % self.micro_batch_times_data_parallel_size
+ == 0
+ )
+ else:
+ self.current_running_global_batch_size = self.current_global_batch_size
+
+ self.num_micro_batches = (
+ self.current_running_global_batch_size // self.micro_batch_times_data_parallel_size
+ )
diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
new file mode 100644
index 0000000000..4a83564ce7
--- /dev/null
+++ b/megatron/core/optimizer/__init__.py
@@ -0,0 +1,445 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import logging
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+
+try:
+ from transformer_engine.pytorch.optimizers import FusedAdam as Adam
+ from transformer_engine.pytorch.optimizers import FusedSGD as SGD
+except ImportError:
+ try:
+ from apex.optimizers import FusedAdam as Adam
+ from apex.optimizers import FusedSGD as SGD
+ except ImportError:
+ import warnings
+
+ warnings.warn(
+ f'Transformer Engine and Apex are not installed. Falling back to Torch optimizers.'
+ )
+
+ # Apex's FusedAdam is a drop-in replacement for torch's AdamW.
+ # pylint: disable-next=line-too-long.
+ # See https://github.com/NVIDIA/apex/blob/7b73b12361068a10b0f44844534613f252a5ea75/apex/optimizers/fused_adam.py#L16.
+ from torch.optim import AdamW as Adam, SGD
+
+from megatron.core import mpu
+
+from ..distributed.param_and_grad_buffer import _ParamAndGradBuffer
+from ..transformer.module import MegatronModule
+from ..utils import log_single_rank
+from .distrib_optimizer import DistributedOptimizer
+from .grad_scaler import ConstantGradScaler, DynamicGradScaler
+from .optimizer import (
+ ChainedOptimizer,
+ Float16OptimizerWithFloat16Params,
+ FP32Optimizer,
+ MegatronOptimizer,
+)
+from .optimizer_config import OptimizerConfig
+
+logger = logging.getLogger(__name__)
+
+
+def _get_param_groups(
+ model_chunks: List[MegatronModule],
+ no_weight_decay_cond: Optional[Callable],
+ scale_lr_cond: Optional[Callable],
+ lr_mult: float,
+ lr: float,
+ min_lr: float,
+ decoupled_lr: Optional[float],
+ decoupled_min_lr: Optional[float],
+) -> List[Dict]:
+ """Create parameter groups for optimizer.
+
+ Creates parameter groups based on weight decay condition (regularized vs
+ non regularized), learning rate scale condition (lr vs lr_mult * lr),
+ and whether it is expert parameters. scale_lr_cond is used during finetuning
+ where head of the network requires a scaled version of the base learning rate.
+
+ Args:
+ model_chunks (List[MegatronModule]): model chunks to create parameter
+ groups for.
+ no_weight_decay_cond (func, optional): function to determine whether a
+ parameter should not perform weight decay.
+ scale_lr_cond (func, optional): function to determine whether a parameter
+ should have a scaled learning rate.
+ lr_mult (float): learning rate multiplier for parameters that
+ satisfy scale_lr_cond.
+ lr (float): learning rate.
+ min_lr (float): minimum learning rate.
+ decoupled_lr (Optional[float]): optional decoupled learning rate.
+ decoupled_min_lr (Optional[float]): optional decoupled minimum learning rate.
+
+ Returns:
+ List of parameter groups.
+ """
+
+ use_decoupled_learning_rate = decoupled_lr is not None
+
+ # Map (wd_mult, lr_mult, is_expert_parallel, is_decoupled_lr) to params.
+ params_map = {}
+ for model_chunk in model_chunks:
+ for name, param in model_chunk.named_parameters():
+ if not param.requires_grad:
+ continue
+
+ is_expert_parallel = not getattr(param, 'allreduce', True)
+
+ if no_weight_decay_cond is not None:
+ no_wd = no_weight_decay_cond(name, param)
+ else:
+ # Do not regularize biases and norm parameters.
+ no_wd = name.endswith(".bias") or len(param.shape) == 1
+
+ if scale_lr_cond is not None:
+ scale_lr = scale_lr_cond(name, param)
+ else:
+ scale_lr = False
+
+ if not no_wd and not scale_lr:
+ wd_mult, _lr_mult = 1.0, 1.0
+ elif not no_wd and scale_lr:
+ wd_mult, _lr_mult = 1.0, lr_mult
+ elif no_wd and not scale_lr:
+ wd_mult, _lr_mult = 0.0, 1.0
+ else:
+ wd_mult, _lr_mult = 0.0, lr_mult
+
+ is_decoupled_lr = False
+ # For input/embedding and output layer: embedding.word_embeddings.weight /
+ # output_layer.weight.
+ if use_decoupled_learning_rate and getattr(
+ param, 'is_embedding_or_output_parameter', False
+ ):
+ is_decoupled_lr = True
+
+ key = (wd_mult, _lr_mult, is_expert_parallel, is_decoupled_lr)
+ if key not in params_map:
+ params_map[key] = []
+ params_map[key].append(param)
+
+ param_groups = []
+ for (wd_mult, _lr_mult, is_expert_parallel, is_decoupled_lr), params in params_map.items():
+ assert len(params) > 0
+ param_group = {
+ 'params': params,
+ 'wd_mult': wd_mult,
+ 'lr_mult': _lr_mult,
+ 'is_expert_parallel': is_expert_parallel,
+ 'is_decoupled_lr': is_decoupled_lr,
+ }
+ param_groups.append(param_group)
+
+ param_groups = _update_min_and_max_lr_in_param_groups(
+ param_groups,
+ lr=lr,
+ min_lr=min_lr,
+ decoupled_lr=decoupled_lr,
+ decoupled_min_lr=decoupled_min_lr,
+ )
+
+ return param_groups
+
+
+def _update_min_and_max_lr_in_param_groups(
+ param_groups: List[Dict],
+ lr: float,
+ min_lr: float,
+ decoupled_lr: Optional[float],
+ decoupled_min_lr: Optional[float],
+) -> List[Dict]:
+ """
+ Updates `max_lr` and `min_lr` values in each parameter group, and returns new list.
+ By default, each group will use `lr` / `min_lr` as `max_lr` / `min_lr`.
+ If `decoupled_lr` is provided, then `decoupled_lr` / `decoupled_min_lr` will be used
+ as `max_lr` / `min_lr` for the input and output layer.
+
+ Args:
+ param_groups (List): parameter groups whose 'max_lr' and `min_lr` fields need to
+ be adjusted.
+ lr (float): learning rate.
+ min_lr (float): minimum learning rate.
+ decoupled_lr (Optional[float]): optional decoupled learning rate.
+ decoupled_min_lr (Optional[float]): optional decoupled minimum learning rate.
+
+ Returns:
+ List of adjusted parameter groups.
+ """
+
+ if decoupled_min_lr is None:
+ decoupled_min_lr = min_lr
+
+ for param_group in param_groups:
+ if param_group['is_decoupled_lr']:
+ assert decoupled_lr is not None
+ param_group['max_lr'] = decoupled_lr
+ param_group['min_lr'] = decoupled_min_lr
+ else:
+ param_group['max_lr'] = lr
+ param_group['min_lr'] = min_lr
+ return param_groups
+
+
+def _get_param_groups_and_buffers(
+ model_chunks: List[MegatronModule],
+ model_chunk_offset: int,
+ config: OptimizerConfig,
+ no_weight_decay_cond: Optional[Callable],
+ scale_lr_cond: Optional[Callable],
+ lr_mult: float,
+ filter_fn: Callable,
+ buffer_name: str,
+) -> Tuple[List[Dict], Dict[int, List[_ParamAndGradBuffer]]]:
+ """Returns parameter groups and buffer for optimizer.
+
+ Args:
+ model_chunks (List[MegatronModule]): model chunks to create parameter
+ groups for.
+ model_chunk_offset (int): offset of model_chunks in global model_chunks list.
+ config (OptimizerConfig): optimizer configuration object.
+ no_weight_decay_cond (func, optional): function to determine whether a
+ parameter should not perform weight decay.
+ scale_lr_cond (func, optional): function to determine whether a parameter
+ should have a scaled learning rate.
+ lr_mult (float): learning rate multiplier for parameters that
+ satisfy scale_lr_cond.
+ lr (float): learning rate.
+ min_lr (float): minimum learning rate.
+ filter_fn (callable): filtering function for param_groups.
+ buffer_name (str): name of buffer.
+
+ Returns:
+ List of parameter groups and dictionary of model chunk IDs to buffers.
+ """
+ param_groups = _get_param_groups(
+ model_chunks,
+ no_weight_decay_cond,
+ scale_lr_cond,
+ lr_mult,
+ lr=config.lr,
+ min_lr=config.min_lr,
+ decoupled_lr=config.decoupled_lr,
+ decoupled_min_lr=config.decoupled_min_lr,
+ )
+ param_groups = list(filter(filter_fn, param_groups))
+ buffers = {}
+ for model_chunk_idx, model_chunk in enumerate(model_chunks):
+ if hasattr(model_chunk, buffer_name):
+ buffers[model_chunk_idx + model_chunk_offset] = getattr(model_chunk, buffer_name)
+
+ return param_groups, buffers
+
+
+def _get_megatron_optimizer_based_on_param_groups(
+ config: OptimizerConfig,
+ model_chunks: List[MegatronModule],
+ param_groups: List,
+ per_model_buffers: Optional[Dict[int, List[_ParamAndGradBuffer]]] = None,
+ model_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+ data_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+ data_parallel_group_gloo: Optional[torch.distributed.ProcessGroup] = None,
+ data_parallel_group_idx: Optional[int] = None,
+) -> MegatronOptimizer:
+ """Get Megatron optimizer based on parameter groups.
+
+ Args:
+ config (OptimizerConfig): optimizer configuration object.
+ model_chunks (list): list of model chunks.
+ param_groups (list): list of parameter groups.
+ per_model_buffers (dict, optional): buffers for distributed optimizer. Defaults to None.
+ data_parallel_group (torch.distributed.ProcessGroup, optional): data-parallel group for
+ distributed optimizer. Defaults to None.
+ data_parallel_group_gloo (torch.distributed.ProcessGroup, optional): gloo data-parallel
+ group for distributed optimizer. Defaults to None.
+ data_parallel_group_idx (int, optional): data-parallel group index for distributed
+ optimizer. Defaults to None.
+
+ Returns:
+ Instance of MegatronOptimizer.
+ """
+ if config.optimizer == 'adam':
+ optimizer = Adam(
+ param_groups,
+ lr=config.lr,
+ weight_decay=config.weight_decay,
+ betas=(config.adam_beta1, config.adam_beta2),
+ eps=config.adam_eps,
+ )
+
+ def init_state_fn(opt):
+ for group in opt.param_groups:
+ for p in group['params']:
+ if len(opt.state[p]) == 0:
+ opt.state[p]['exp_avg'] = torch.zeros_like(p.data)
+ opt.state[p]['exp_avg_sq'] = torch.zeros_like(p.data)
+
+ elif config.optimizer == 'sgd':
+ optimizer = SGD(
+ param_groups,
+ lr=config.lr,
+ weight_decay=config.weight_decay,
+ momentum=config.sgd_momentum,
+ )
+ init_state_fn = None
+ else:
+ raise Exception('{} optimizer is not supported.'.format(config.optimizer))
+
+ # Mixed precision optimizer.
+ # - Note: both the Float16Optimizer and the DistributedOptimizer inherit
+ # from the MixedPrecisionOptimizer, which manages any optimizer where
+ # the model params and main params are distinct.
+ if config.fp16 or config.bf16 or config.use_distributed_optimizer:
+
+ # Grad scaler:
+ # if loss-scale is provided, instantiate the constant scaler.
+ # if we are using fp16 and loss-scale is not present, use a
+ # dynamic scaler.
+ # otherwise we are running in bf16 with no loss-scale so
+ # leave it as None.
+ grad_scaler = None
+
+ # Constant loss scale.
+ if config.loss_scale:
+ grad_scaler = ConstantGradScaler(config.loss_scale)
+
+ # Dynamic loss scale.
+ else:
+ if config.fp16:
+ grad_scaler = DynamicGradScaler(
+ initial_scale=config.initial_loss_scale,
+ min_scale=config.min_loss_scale,
+ growth_factor=2.0,
+ backoff_factor=0.5,
+ growth_interval=config.loss_scale_window,
+ hysteresis=config.hysteresis,
+ )
+
+ optimizer_args = [optimizer, config, grad_scaler, init_state_fn]
+ if config.use_distributed_optimizer:
+ optimizer = DistributedOptimizer(
+ *optimizer_args,
+ model_chunks=model_chunks,
+ per_model_buffers=per_model_buffers,
+ data_parallel_group=data_parallel_group,
+ data_parallel_group_gloo=data_parallel_group_gloo,
+ data_parallel_group_idx=data_parallel_group_idx,
+ )
+ else:
+ optimizer = Float16OptimizerWithFloat16Params(*optimizer_args)
+ setattr(optimizer, 'model_parallel_group', model_parallel_group)
+ else:
+ # FP32 optimizer.
+ optimizer = FP32Optimizer(optimizer, config, init_state_fn)
+ setattr(optimizer, 'model_parallel_group', model_parallel_group)
+
+ return optimizer
+
+
+def get_megatron_optimizer(
+ config: OptimizerConfig,
+ model_chunks: List[MegatronModule],
+ no_weight_decay_cond: Optional[Callable] = None,
+ scale_lr_cond: Optional[Callable] = None,
+ lr_mult: float = 1.0,
+) -> MegatronOptimizer:
+ """Retrieve the Megatron optimizer for model chunks.
+
+ We use separate optimizers for expert parameters and non-expert parameters.
+
+ Args:
+ config (OptimizerConfig): optimizer configuration object.
+ model_chunks (List[MegatronModule]): model chunks to get optimizer for.
+ no_weight_decay_cond (func, optional): function to determine whether a parameter
+ should not perform weight decay. Defaults to None.
+ scale_lr_cond (func, optional): function to determine whether a parameter
+ should have a scaled learning rate. Defaults to None.
+ lr_mult (float, optional): learning rate multiplier for parameters that
+ satisfy scale_lr_cond. Defaults to 1.0.
+
+ Returns:
+ Instance of MegatronOptimizer.
+ """
+
+ log_single_rank(logger, logging.INFO, f'Setting up optimizer with config {config}')
+
+ # Separate out first model chunk if overlapping param AG with optimizer step.
+ if config.overlap_param_gather_with_optimizer_step:
+ all_dense_model_chunks = [[model_chunks[0]], model_chunks[1:]]
+ overlap_param_gather_with_optimizer_step_flags = [True, False]
+ else:
+ all_dense_model_chunks = [model_chunks]
+ overlap_param_gather_with_optimizer_step_flags = [False]
+ model_parallel_rank = torch.distributed.get_rank(mpu.get_model_parallel_group())
+
+ optimizers = []
+ model_chunk_offset = 0
+ for dense_model_chunks, overlap_param_gather_with_optimizer_step in zip(
+ all_dense_model_chunks, overlap_param_gather_with_optimizer_step_flags
+ ):
+ param_groups, buffers = _get_param_groups_and_buffers(
+ dense_model_chunks,
+ model_chunk_offset=model_chunk_offset,
+ config=config,
+ no_weight_decay_cond=no_weight_decay_cond,
+ scale_lr_cond=scale_lr_cond,
+ lr_mult=lr_mult,
+ filter_fn=lambda g: not g['is_expert_parallel'],
+ buffer_name='buffers',
+ )
+ for model_chunk in dense_model_chunks:
+ model_chunk.overlap_param_gather_with_optimizer_step = (
+ overlap_param_gather_with_optimizer_step
+ )
+ optimizers.append(
+ _get_megatron_optimizer_based_on_param_groups(
+ config,
+ model_chunks=dense_model_chunks,
+ param_groups=param_groups,
+ per_model_buffers=buffers,
+ model_parallel_group=mpu.get_model_parallel_group(),
+ data_parallel_group=mpu.get_data_parallel_group(with_context_parallel=True),
+ data_parallel_group_gloo=mpu.get_data_parallel_group_gloo(
+ with_context_parallel=True
+ ),
+ data_parallel_group_idx=model_parallel_rank,
+ )
+ )
+ model_chunk_offset += 1
+
+ moe_param_groups, moe_buffers = _get_param_groups_and_buffers(
+ model_chunks,
+ model_chunk_offset=0,
+ config=config,
+ no_weight_decay_cond=no_weight_decay_cond,
+ scale_lr_cond=scale_lr_cond,
+ lr_mult=lr_mult,
+ filter_fn=lambda g: g['is_expert_parallel'],
+ buffer_name='expert_parallel_buffers',
+ )
+ if len(moe_param_groups) > 0:
+ model_parallel_world_size = torch.distributed.get_world_size(mpu.get_model_parallel_group())
+ expert_parallel_rank = mpu.get_expert_model_parallel_rank()
+ optimizers.append(
+ _get_megatron_optimizer_based_on_param_groups(
+ config,
+ model_chunks=model_chunks,
+ param_groups=moe_param_groups,
+ per_model_buffers=moe_buffers,
+ model_parallel_group=mpu.get_model_parallel_group(with_expert_parallel=True),
+ data_parallel_group=mpu.get_data_modulo_expert_parallel_group(
+ with_context_parallel=True
+ ),
+ data_parallel_group_gloo=mpu.get_data_modulo_expert_parallel_group_gloo(
+ with_context_parallel=True
+ ),
+ data_parallel_group_idx=expert_parallel_rank * model_parallel_world_size
+ + model_parallel_rank,
+ )
+ )
+
+ if len(optimizers) == 1:
+ return optimizers[0]
+
+ return ChainedOptimizer(optimizers)
diff --git a/megatron/core/optimizer/clip_grads.py b/megatron/core/optimizer/clip_grads.py
new file mode 100644
index 0000000000..708ccd019e
--- /dev/null
+++ b/megatron/core/optimizer/clip_grads.py
@@ -0,0 +1,193 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Gradient clipping."""
+
+import os
+from typing import List, Optional, Union
+
+import torch
+from torch import inf
+
+try:
+ from transformer_engine.pytorch.optimizers import (
+ multi_tensor_applier,
+ multi_tensor_l2norm,
+ multi_tensor_scale,
+ )
+
+ l2_norm_impl = multi_tensor_l2norm
+ multi_tensor_scale_impl = multi_tensor_scale
+except ImportError:
+ try:
+ import amp_C
+ from apex.multi_tensor_apply import multi_tensor_applier
+
+ l2_norm_impl = amp_C.multi_tensor_l2norm
+ multi_tensor_scale_impl = amp_C.multi_tensor_scale
+ except ImportError:
+ import warnings
+
+ warnings.warn(
+ f'Transformer Engine and Apex are not installed. '
+ 'Falling back to local implementations of multi_tensor_applier, '
+ 'multi_tensor_l2norm, and multi_tensor_scale'
+ )
+
+ from megatron.core.utils import (
+ local_multi_tensor_applier,
+ local_multi_tensor_l2_norm,
+ local_multi_tensor_scale,
+ )
+
+ multi_tensor_applier = local_multi_tensor_applier
+ l2_norm_impl = local_multi_tensor_l2_norm
+ multi_tensor_scale_impl = local_multi_tensor_scale
+
+
+from ..tensor_parallel import param_is_not_tensor_parallel_duplicate
+from ..transformer.module import param_is_not_shared
+
+
+def get_grad_norm_fp32(
+ grads_for_norm: Union[List[torch.Tensor], torch.Tensor],
+ norm_type: Union[int, float] = 2,
+ model_parallel_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> float:
+ """Calculate the norm of gradients in fp32.
+
+ This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and
+ added functionality to handle model parallel parameters.
+
+ Arguments:
+ grads_for_norm (Iterable[Tensor] or Tensor): an iterable of Tensors or a single
+ Tensor that will be used for calculating the grad norm.
+ norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+ infinity norm.
+ model_parallel_group (group): given the nature of the distributed
+ optimizer, this is passed as an argument.
+
+ Returns:
+ Total norm of the parameters (viewed as a single vector).
+ """
+
+ if isinstance(grads_for_norm, torch.Tensor):
+ grads_for_norm = [grads_for_norm]
+
+ # Norm parameters.
+ norm_type = float(norm_type)
+ total_norm = 0.0
+
+ # Calculate norm.
+ if norm_type == inf:
+ total_norm = max(grad.abs().max() for grad in grads_for_norm)
+ total_norm_cuda = torch.tensor([float(total_norm)], dtype=torch.float, device='cuda')
+ # Take max across all model-parallel GPUs.
+ torch.distributed.all_reduce(
+ total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=model_parallel_group
+ )
+ total_norm = total_norm_cuda[0].item()
+
+ else:
+ if norm_type == 2.0:
+ dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='cuda')
+ # Use apex's multi-tensor applier for efficiency reasons.
+ # Multi-tensor applier takes a function and a list of list
+ # and performs the operation on that list all in one kernel.
+ if grads_for_norm:
+ grad_norm, _ = multi_tensor_applier(
+ l2_norm_impl,
+ dummy_overflow_buf,
+ [grads_for_norm],
+ False, # no per-parameter norm
+ )
+ else:
+ grad_norm = torch.tensor([0], dtype=torch.float, device='cuda')
+ # Since we will be summing across data parallel groups,
+ # we need the pow(norm-type).
+ total_norm = grad_norm**norm_type
+
+ else:
+ for grad in grads_for_norm:
+ grad_norm = torch.norm(grad, norm_type)
+ total_norm += grad_norm**norm_type
+
+ # Sum across all model-parallel GPUs.
+ torch.distributed.all_reduce(
+ total_norm, op=torch.distributed.ReduceOp.SUM, group=model_parallel_group
+ )
+ total_norm = total_norm.item() ** (1.0 / norm_type)
+
+ return total_norm
+
+
+def clip_grad_by_total_norm_fp32(
+ parameters: Union[List[torch.Tensor], torch.Tensor],
+ max_norm: Union[int, float],
+ total_norm: float,
+):
+ """Clips gradient of an iterable of parameters in fp32 by total norm.
+
+ Note that the gradients are modified in place.
+
+ Args:
+ parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+ single Tensor that will have gradients normalized.
+ max_norm (float or int): max norm of the gradients.
+ total_norm (float): total norm of the gradients.
+ """
+ # Grads.
+ grads = []
+ for param in parameters:
+ if param.grad is not None:
+ assert param.grad.type() == 'torch.cuda.FloatTensor'
+ grads.append(param.grad.detach())
+
+ # Scale.
+ clip_coeff = max_norm / (total_norm + 1.0e-6)
+ if clip_coeff < 1.0:
+ dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='cuda')
+ multi_tensor_applier(
+ multi_tensor_scale_impl, dummy_overflow_buf, [grads, grads], clip_coeff
+ )
+
+
+def count_zeros_fp32(
+ parameters: Union[List[torch.Tensor], torch.Tensor],
+ model_parallel_group: torch.distributed.ProcessGroup,
+) -> float:
+ """Counts the number of zeros in gradients associated with the passed-in list of
+ parameters.
+
+ Args:
+ parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+ single Tensor that will have the number of zeros in its corresponding
+ gradient counted.
+ model_parallel_group (torch.distributed.ProcessGroup, optional): model-parallel
+ group over which grad norm needs to be aggregated.
+ """
+
+ if isinstance(parameters, torch.Tensor):
+ parameters = [parameters]
+
+ # Filter parameters based on:
+ # - grad should not be none
+ # - parameter should not be shared
+ # - should not be a replica due to tensor model parallelism
+ total_num_zeros = torch.tensor([0.0], dtype=torch.float, device='cuda')
+ for param in parameters:
+ grad_not_none = param.grad is not None
+ is_not_shared = param_is_not_shared(param)
+ is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param)
+ if grad_not_none and is_not_shared and is_not_tp_duplicate:
+ grad = param.grad.detach()
+ num_zeros = grad.numel() - torch.count_nonzero(grad)
+ total_num_zeros = num_zeros + total_num_zeros
+
+ # Sum across all model-parallel GPUs.
+ torch.distributed.all_reduce(
+ total_num_zeros, op=torch.distributed.ReduceOp.SUM, group=model_parallel_group
+ )
+
+ total_num_zeros = total_num_zeros.item()
+
+ return total_num_zeros
diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
new file mode 100644
index 0000000000..dfa8d51979
--- /dev/null
+++ b/megatron/core/optimizer/distrib_optimizer.py
@@ -0,0 +1,1839 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Megatron distributed optimizer."""
+
+
+import itertools
+import warnings
+from dataclasses import replace
+from logging import getLogger
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+
+HAVE_APEX_OR_TE = True
+try:
+ from transformer_engine.pytorch.optimizers import FusedAdam as Adam
+except ImportError:
+ try:
+ from apex.optimizers import FusedAdam as Adam
+ except ImportError:
+ from torch.optim import Adam
+
+ HAVE_APEX_OR_TE = False
+
+from .. import tensor_parallel
+from ..config_logger import has_config_logger_enabled, log_config_to_disk
+from ..dist_checkpointing import ShardedTensor
+from ..dist_checkpointing.dict_utils import nested_values
+from ..dist_checkpointing.mapping import (
+ LocalNonpersistentObject,
+ ShardedObject,
+ ShardedStateDict,
+ ShardedTensorFactory,
+)
+from ..dist_checkpointing.utils import extract_sharded_tensors_and_factories
+from ..distributed.param_and_grad_buffer import _ParamAndGradBuffer, partition_buckets
+from ..transformer.module import MegatronModule
+from ..utils import is_float8tensor
+from .grad_scaler import MegatronGradScaler
+from .optimizer import (
+ MixedPrecisionOptimizer,
+ _multi_tensor_copy_this_to_that,
+ _zero_grad_group_helper,
+)
+from .optimizer_config import OptimizerConfig
+
+try:
+ # This will be used when "--fp8-param-gather" is enabled.
+ # When BF16/FP16 parameters don't exist, we need to cast the FP32 main parameters to
+ # FP8 directly in the optimizer.
+ from transformer_engine.pytorch.cpp_extensions import cast_to_fp8
+except:
+ pass
+
+logger = getLogger(__name__)
+
+
+class Range:
+ """
+ A range represents a start and end points for indexing a shard
+ from a full tensor.
+
+ Args:
+ start (int): Start index.
+ end (int): End index.
+ """
+
+ def __init__(self, start: int, end: int):
+ self.start = start
+ self.end = end
+ self.size = end - start
+
+ def normalize(self, start: int = 0):
+ """Shift start/end indexes to start at new start index.
+
+ Both start and end indexes will be shifted by [new start] - [old start].
+
+ Args:
+ start (int): New start index.
+ """
+ return Range(start, start + self.size)
+
+ def __str__(self):
+ return "%d,%d [%d]" % (self.start, self.end, self.size)
+
+ def __len__(self):
+ return self.end - self.start
+
+
+class DistributedOptimizer(MixedPrecisionOptimizer):
+ """Distributed optimizer, for all data types (fp16, bf16, and fp32).
+
+ See __init__() below for argument details.
+ """
+
+ @classmethod
+ def _build_model_gbuf_param_range_map(
+ cls,
+ param_world_index_map: Dict[torch.nn.Parameter, Tuple],
+ gbuf_world_range: Range,
+ bucket_offset: int,
+ ):
+ """
+ Build mapping from param reference to grad buffer shard ranges.
+
+ This method builds a mapping from parameter references to grad
+ buffer shard ranges, specific to each data-parallel (DP) rank's
+ set of 'owned' parameters. Each grad buffer (padded to be an even
+ multiple of DP-world-size) is conceptually divided into DP-world-size
+ contiguous regions, where each DP rank 'owns' a contiguous region.
+ Ownership in this sense means DP rank is responsible for reducing
+ the relevant subset of grads, and updating the relevant subset of
+ params.
+
+ This conceptual partitioning of the grad buffer does NOT respect
+ parameter boundaries, and as such it is assumed that each created
+ range references a shard (or subset) of the full parameter. It is
+ easiest to think of each DP rank as operating (i.e., reducing,
+ gathering) purely on views into the grad buffer, for all model-to-
+ main & main-to-model operations.
+
+ This method creates four ranges:
+ - The param's range within the entire grad buffer (i.e., world index).
+ - The param's range within the relevant grad bucket's buffer.
+ - The param's range within the DP rank's local view of the grad buffer.
+ - The param's range within itself (i.e., its shard).
+ """
+
+ # Param range map.
+ param_range_map = {}
+ for param, param_world_indexes in param_world_index_map.items():
+
+ # Param range.
+ param_world_start, param_world_end, _ = param_world_indexes
+ param_local_start = max(0, param_world_start - gbuf_world_range.start)
+ param_local_end = min(gbuf_world_range.size, param_world_end - gbuf_world_range.start)
+
+ # Add param, if within local gbuf range.
+ if param_local_end > param_local_start:
+ param_local_range = Range(param_local_start, param_local_end)
+ param_world_range = param_local_range.normalize(
+ param_local_start + gbuf_world_range.start
+ )
+ param_world_range_in_bucket = Range(
+ param_world_range.start - bucket_offset, param_world_range.end - bucket_offset
+ )
+ sub_param_start = max(0, gbuf_world_range.start - param_world_start)
+ sub_param_range = param_local_range.normalize(sub_param_start)
+ param_range_map[param] = {
+ "gbuf_world": param_world_range,
+ "gbuf_world_in_bucket": param_world_range_in_bucket,
+ "gbuf_local": param_local_range,
+ "param": sub_param_range,
+ }
+
+ return param_range_map
+
+ @classmethod
+ def _build_model_gbuf_range(cls, param_and_grad_buffer: _ParamAndGradBuffer, bucket_index: int):
+ """
+ Build mapping between params and their grad buffers.
+
+ This method does the initial setup for the method above. This setup
+ includes determining the shard ranges into the param_and_grad_buffer
+ for each data-parallel (DP) rank. Each DP rank keeps range info for
+ all other DP ranks, for the purpose of creating args for
+ reduce-scatter and all-gather.
+ """
+
+ data_parallel_rank = torch.distributed.get_rank(param_and_grad_buffer.data_parallel_group)
+ data_parallel_world_size = param_and_grad_buffer.data_parallel_group.size()
+
+ bucket = param_and_grad_buffer.buckets[bucket_index]
+ gbuf_size = bucket.grad_data.numel()
+ assert (
+ gbuf_size % data_parallel_world_size == 0
+ ), f"Each bucket's buffer size should be divisible by {data_parallel_world_size}"
+ max_gbuf_range_size = gbuf_size // data_parallel_world_size
+
+ # All world ranges (i.e., across all data parallel ranks).
+ gbuf_world_all_ranges = []
+ for r in range(data_parallel_world_size):
+ # Compute start of chunk in this bucket.
+ gbuf_world_start = r * max_gbuf_range_size
+ gbuf_world_end = min(gbuf_size, gbuf_world_start + max_gbuf_range_size)
+ # Add bucket's offset in grad buffer.
+ gbuf_world_range = Range(
+ gbuf_world_start + bucket.offset, gbuf_world_end + bucket.offset
+ )
+ gbuf_world_all_ranges.append(gbuf_world_range)
+
+ # Local DP's ranges.
+ gbuf_world_range = gbuf_world_all_ranges[data_parallel_rank]
+
+ # Get each param's ranges.
+ param_range_map = cls._build_model_gbuf_param_range_map(
+ param_and_grad_buffer.param_index_map, gbuf_world_range, bucket.offset
+ )
+
+ # Group into dict.
+ data = {"param_map": param_range_map}
+
+ return data
+
+ @classmethod
+ def _build_gbuf_range_map(cls, param_and_grad_buffer: _ParamAndGradBuffer):
+ """
+ Build mapping between params and their grad buffers. These mappings are
+ partitioned according to data type.
+
+ Iterate through all buckets of grad buffer to construct param ranges
+ that this rank "owns" (the dp_rank'th shard of each bucket, where each
+ shard is 1/dp_world_size of the bucket).
+
+ Args:
+ param_and_grad_buffer (_ParamAndGradBuffer): buffer to build mapping for.
+ """
+ return {
+ (param_and_grad_buffer.param_dtype, param_and_grad_buffer.grad_dtype): [
+ cls._build_model_gbuf_range(param_and_grad_buffer, bucket_index)
+ for bucket_index in range(len(param_and_grad_buffer.buckets))
+ ]
+ }
+
+ @classmethod
+ def _build_model_param_gbuf_map(
+ cls, gbuf_ranges: List[Dict]
+ ) -> Dict[torch.nn.Parameter, Tuple]:
+ """
+ Create a reverse of the gbuf_ranges, for referencing in opposite direction.
+ """
+ param_gbuf_map = {}
+ for gbuf_index, gbuf_range_map in enumerate(gbuf_ranges):
+ for dtype, gbuf_range_map_for_all_buckets in gbuf_range_map.items():
+ for bucket_index, gbuf_range_map in enumerate(gbuf_range_map_for_all_buckets):
+ for param, _ in gbuf_range_map["param_map"].items():
+ assert param not in param_gbuf_map, (
+ "Param should not be in param_gbuf_map; each param only belongs "
+ "to a single bucket."
+ )
+ param_gbuf_map[param] = (gbuf_index, dtype, bucket_index)
+ return param_gbuf_map
+
+ @classmethod
+ def _build_optimizer_group_ranges(cls, param_groups: List[Dict], gbuf_ranges: List[Dict]):
+ """
+ Create optimizer groups.
+
+ Given the set of parameter shard ranges that are owned by the current
+ data-parallel (DP) rank, gather the set of parameters that will be
+ used (in the method below) to create the current DP's optimizer
+ groups.
+ """
+
+ # Param group map.
+ # World param group map.
+ # - Store a mapping of for all parameters
+ # across all DP ranks. This is necessary because it is our first
+ # cross reference between the DDP mappings and the optimizer group
+ # parameters. This mapping only for use in the next step of building
+ # the local mapping over this DP rank's parameters.
+ world_param_group_map = {}
+ for group_index, group in enumerate(param_groups):
+ for param in group["params"]:
+ assert param.requires_grad
+ world_param_group_map[param] = group_index
+
+ # Optimizer group ranges & param-group mapping.
+ # - Build a mapping from groups to their contained parameters, and also
+ # from parameters to their containing group index and order within
+ # the group. The group index and order are particularly important for
+ # saving and loading checkpoints.
+ local_param_group_map = {}
+ group_ranges = [{"params": []} for _ in param_groups]
+ for gbuf_range_map in gbuf_ranges:
+ for dtype, gbuf_range_map_for_all_buckets in gbuf_range_map.items():
+ for gbuf_range_map in gbuf_range_map_for_all_buckets:
+ for param in gbuf_range_map["param_map"]:
+ group_index = world_param_group_map[param]
+ group_range = group_ranges[group_index]
+ group_range["params"].append(param)
+ local_param_group_map[param] = (group_index, len(group_range["params"]) - 1)
+
+ # Squeeze zero-size group ranges.
+ for group_index, group_range in enumerate(group_ranges):
+ group_range["orig_group"] = param_groups[group_index]
+ group_range["orig_group_idx"] = param_groups[group_index]
+
+ return local_param_group_map, group_ranges
+
+ @classmethod
+ def _build_model_and_main_param_groups(
+ cls,
+ gbuf_ranges: List[Dict],
+ param_gbuf_map: Dict[torch.nn.Parameter, Tuple],
+ opt_group_ranges: List,
+ ):
+ """
+ Create main parameter groups needed for the optimizer step.
+
+ These groups encompass both: 1) groups used by this class, for
+ reducing/gather, and 2) groups used by the inner optimizer for the
+ parameter update. Given that the conceptual grad buffer partitioning
+ (created in earlier method) doesn't respect parameter boundaries,
+ the optimizer operates on shards of the model parameters, rather than
+ the full parameters.
+ """
+
+ # Parameter groups:
+ # model_float16_groups: original float16 parameters
+ # model_fp32_groups: original fp32 parameters
+ # shard_float16_groups: shards of original float16 parameters
+ # shard_fp32_groups: shards of original fp32 parameters
+ # shard_fp32_from_float16_groups: fp32 copy of float16 parameters
+ model_float16_groups = []
+ model_fp32_groups = []
+ shard_float16_groups = []
+ shard_fp32_groups = []
+ shard_fp32_from_float16_groups = []
+
+ # Allocate (or slice) each group's param shard.
+ for group_range in opt_group_ranges:
+
+ # Params of this group.
+ model_float16_params_this_group = []
+ model_fp32_params_this_group = []
+ shard_float16_params_this_group = []
+ shard_fp32_params_this_group = []
+ shard_fp32_from_float16_params_this_group = []
+ model_float16_groups.append(model_float16_params_this_group)
+ model_fp32_groups.append(model_fp32_params_this_group)
+ shard_float16_groups.append(shard_float16_params_this_group)
+ shard_fp32_groups.append(shard_fp32_params_this_group)
+ shard_fp32_from_float16_groups.append(shard_fp32_from_float16_params_this_group)
+
+ for model_param in group_range["params"]:
+
+ assert model_param.requires_grad
+
+ gbuf_index, dtype, bucket_index = param_gbuf_map[model_param]
+ gbuf_range = gbuf_ranges[gbuf_index][dtype][bucket_index]
+ param_range = gbuf_range["param_map"][model_param]["param"]
+
+ # fp16, bf16 params.
+ if model_param.type() in ['torch.cuda.HalfTensor', 'torch.cuda.BFloat16Tensor']:
+
+ # Clone model -> main.
+ shard_model_param = model_param.detach().view(-1)[
+ param_range.start : param_range.end
+ ]
+
+ # If we use FP8 params to initialize FP32 main params (compared to using the
+ # bf16/fp16 params to initialize the main params), there will be a loss of
+ # precision at the beginning of training (this problem will not occur if the
+ # training is long enough or if the main params are loaded from a checkpoint).
+ if is_float8tensor(model_param) and hasattr(
+ model_param, 'get_high_precision_init_val'
+ ):
+ shard_main_param = (
+ model_param.get_high_precision_init_val()
+ .view(-1)[param_range.start : param_range.end]
+ .clone()
+ .to(shard_model_param.device)
+ .float()
+ )
+ model_param.clear_high_precision_init_val()
+ else:
+ shard_main_param = shard_model_param.clone().float()
+
+ tensor_parallel.copy_tensor_model_parallel_attributes(
+ shard_model_param, model_param
+ )
+ tensor_parallel.copy_tensor_model_parallel_attributes(
+ shard_main_param, model_param
+ )
+ if hasattr(model_param, 'shared'):
+ shard_model_param.shared = model_param.shared
+ shard_main_param.shared = model_param.shared
+
+ # Add to group.
+ model_float16_params_this_group.append(model_param)
+ shard_float16_params_this_group.append(shard_model_param)
+ shard_fp32_from_float16_params_this_group.append(shard_main_param)
+
+ # fp32 params.
+ elif model_param.type() == 'torch.cuda.FloatTensor':
+ shard_model_param = model_param.view(-1)[param_range.start : param_range.end]
+ model_fp32_params_this_group.append(model_param)
+ shard_fp32_params_this_group.append(shard_model_param)
+ tensor_parallel.copy_tensor_model_parallel_attributes(
+ shard_model_param, model_param
+ )
+ if hasattr(model_param, 'shared'):
+ shard_model_param.shared = model_param.shared
+
+ else:
+ raise TypeError(
+ 'Wrapped parameters must be one of '
+ 'torch.cuda.FloatTensor, '
+ 'torch.cuda.HalfTensor, or '
+ 'torch.cuda.BFloat16Tensor. '
+ 'Received {}'.format(model_param.type())
+ )
+
+ # Update optimizer's params.
+ group_range["orig_group"]["params"] = [
+ *shard_fp32_params_this_group,
+ *shard_fp32_from_float16_params_this_group,
+ ]
+
+ return (
+ model_float16_groups,
+ model_fp32_groups,
+ shard_float16_groups,
+ shard_fp32_groups,
+ shard_fp32_from_float16_groups,
+ )
+
+ def __init__(
+ self,
+ optimizer: torch.optim.Optimizer,
+ config: OptimizerConfig,
+ grad_scaler: MegatronGradScaler,
+ init_state_fn: Optional[Callable],
+ model_chunks: List[MegatronModule],
+ per_model_buffers: Dict[int, List[_ParamAndGradBuffer]],
+ data_parallel_group: torch.distributed.ProcessGroup,
+ data_parallel_group_gloo: torch.distributed.ProcessGroup,
+ data_parallel_group_idx: int,
+ ):
+ """
+ Distributed optimizer, for all data types (fp16, bf16, and fp32).
+
+ The steps in this method create the core mapping between param and grad buffers,
+ parameters, and parameter shard ranges, that is needed for converting between model
+ param indexes and main parameter shard indexes. This method also updates the optimizer
+ parameter groups with the newly created shards.
+
+ Args:
+ optimizer (torch.optim.Optimizer): base optimizer such as Adam or SGD.
+ config (OptimizerConfig): configuration object for optimizer.
+ grad_scaler (MegatronGradScaler): used for scaling gradients. Note that
+ this can be None. This case happens when `bf16 = True` and we don't
+ use any loss scale. Note that for `bf16 = True`, we can have
+ a constant gradient scaler. Also for `bf16 = False`, we
+ always require a grad scaler.
+ init_state_fn (Callable, optional): function to initialize state in the optimizer.
+ model_chunks (List[MegatronModule]): list of model chunks.
+ per_model_buffers (Dict[int, List[ParamAndGradBuffer]]): the implementation of the
+ distributed optimizer is centered on using a contiguous buffer for
+ communicating grads & params between the model state and the optimizer state.
+ You can find a more detailed description in
+ https://github.com/NVIDIA/Megatron-LM/blob/main/docs/source/distrib_optimizer.md.
+ data_parallel_group (torch.distributed.ProcessGroup): data-parallel group to use to
+ all-gather params after optimizer.step().
+ data_parallel_group_gloo (torch.distributed.ProcessGroup): gloo data-parallel group
+ (used in checkpoint loading and saving).
+ data_parallel_group_idx (int): index in data-parallel group (used by
+ distributed checkpointing logic).
+ """
+
+ if has_config_logger_enabled(config):
+ log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+ assert (
+ HAVE_APEX_OR_TE
+ ), f'Please install Apex or Transformer Engine to use DistributedOptimizer.'
+
+ super().__init__(optimizer, config, grad_scaler, init_state_fn)
+ self.model_chunks = model_chunks
+ self.ddp_config = self.model_chunks[0].ddp_config
+ for model_chunk in self.model_chunks:
+ assert self.ddp_config == model_chunk.ddp_config
+
+ assert isinstance(
+ optimizer, Adam
+ ), "Only Adam currently supported, due to checkpointing requirements."
+
+ # Model grad buffer ranges.
+ assert per_model_buffers is not None, "per_model_buffers must be provided"
+ self.buffers = list(itertools.chain(*per_model_buffers.values()))
+ self.per_model_buffers = per_model_buffers
+ self.data_parallel_group = data_parallel_group
+ self.data_parallel_group_gloo = data_parallel_group_gloo
+ self.data_parallel_group_idx = data_parallel_group_idx
+
+ self.gbuf_idx_to_model_idx_map = {}
+ gbuf_idx = 0
+ for model_idx, buffers in self.per_model_buffers.items():
+ for _ in buffers:
+ self.gbuf_idx_to_model_idx_map[gbuf_idx] = model_idx
+ gbuf_idx += 1
+
+ self.per_model_bucket_groups = {}
+ for model_idx, buffers in self.per_model_buffers.items():
+ self.per_model_bucket_groups[model_idx] = partition_buckets(buffers)
+
+ self.gbuf_ranges = []
+ self.per_bucket_numel = []
+ self.per_bucket_numel_unpadded = []
+ for buffer in self.buffers:
+
+ self.per_bucket_numel.append(
+ {
+ (buffer.param_dtype, buffer.grad_dtype): [
+ bucket.grad_data.numel() for bucket in buffer.buckets
+ ]
+ }
+ )
+ self.per_bucket_numel_unpadded.append(
+ {
+ (buffer.param_dtype, buffer.grad_dtype): [
+ bucket.numel_unpadded for bucket in buffer.buckets
+ ]
+ }
+ )
+ self.gbuf_ranges.append(self._build_gbuf_range_map(buffer))
+ self.model_param_gbuf_map = self._build_model_param_gbuf_map(self.gbuf_ranges)
+
+ # Optimizer ranges.
+ (self.model_param_group_index_map, self.opt_group_ranges) = (
+ self._build_optimizer_group_ranges(self.optimizer.param_groups, self.gbuf_ranges)
+ )
+
+ # Allocate main param shards.
+ (
+ self.model_float16_groups,
+ self.model_fp32_groups,
+ self.shard_float16_groups,
+ self.shard_fp32_groups,
+ self.shard_fp32_from_float16_groups,
+ ) = self._build_model_and_main_param_groups(
+ self.gbuf_ranges, self.model_param_gbuf_map, self.opt_group_ranges
+ )
+
+ # Update optimizer groups.
+ # - Also, leverage state_dict() and load_state_dict() to
+ # recast preexisting per-param state tensors.
+ self.optimizer.param_groups = [g["orig_group"] for g in self.opt_group_ranges]
+ self.optimizer.load_state_dict(self.optimizer.state_dict())
+
+ def enable_pre_hook(self):
+ """
+ Enable forward pre-hook needed for param all-gather overlap with forward compute.
+ """
+ warnings.warn(
+ "`DistributedOptimizer.enable_pre_hook` will be deprecated in a future release. "
+ "Use `DistributedDataParallel.enable_forward_pre_hook` directly."
+ )
+ for model_chunk in self.model_chunks:
+ model_chunk.enable_forward_pre_hook()
+
+ def disable_pre_hook(self):
+ """
+ Disable forward pre-hook needed for param all-gather overlap with forward compute.
+ """
+ warnings.warn(
+ "`DistributedOptimizer.disable_pre_hook` will be deprecated in a future release. "
+ "Use `DistributedDataParallel.disable_forward_pre_hook` directly."
+ )
+ for model_chunk in self.model_chunks:
+ model_chunk.disable_forward_pre_hook()
+
+ def _get_model_param_range_map(self, param: torch.nn.Parameter):
+ """
+ Given a model param, get the index sub-range of the param that this
+ data-parallel rank owns.
+ """
+ gbuf_index, dtype, bucket_index = self.model_param_gbuf_map[param]
+ gbuf_range_map = self.gbuf_ranges[gbuf_index][dtype][bucket_index]
+ param_range_map = gbuf_range_map["param_map"][param]
+ return param_range_map
+
+ def get_model_parallel_group(self) -> torch.distributed.ProcessGroup:
+ """
+ With the distributed optimizer, the model parallel group is the
+ entire world.
+ """
+ return None
+
+ def state_dict(self):
+ """
+ The state dict contains all non-DP-rank-dependent (i.e., non-parameter-
+ related) optimizer variables. The returned state dict can be stored in
+ the standard model/RNG checkpoint file. The parameter and dependent
+ optimizer state (e.g., exp_avg, exp_avg_sq) are stored in a separate
+ checkpoint file by calling 'save_parameter_state()'.
+ """
+
+ inner_state_dict = self.optimizer.state_dict()
+ state_dict = {}
+
+ # Extract 'step', for non-Apex/TE support.
+ if not HAVE_APEX_OR_TE:
+ steps = list(set([s["step"].item() for s in inner_state_dict["state"].values()]))
+ assert len(steps) == 1
+ step = steps[0]
+
+ # Optimizer state (do not store parameter state here).
+ state_dict['optimizer'] = {k: v for k, v in inner_state_dict.items() if k != "state"}
+ for param_group in state_dict["optimizer"]["param_groups"]:
+ del param_group["params"]
+ if not HAVE_APEX_OR_TE:
+ # Native PyTorch param group requires step (i.e., iteration).
+ param_group["step"] = step
+
+ # Grad scaler state.
+ if self.grad_scaler:
+ state_dict['grad_scaler'] = self.grad_scaler.state_dict()
+
+ return state_dict
+
+ def load_state_dict(self, state_dict):
+ """Load the state dict.
+
+ As detailed in state_dict(), the state dict contains all non-
+ parameter-related variables. This method is notably longer than
+ state_dict(), because the Torch optimizers state has yet to be
+ allocated at this point, and so we must do a cross referencing between
+ the optimizers state (and the ordering it expects for parameter state)
+ and this DP rank's shards. The optimizer at this point does not contain
+ any tensor dimension information, so we must get these dimensions from
+ the DP shards mapped during DistributedOptimizer.__init__().
+
+ The tensor parameter state is loaded via load_parameter_state(), and
+ so this method also must populate the loaded state dict with dummy
+ tensor data (i.e., via torch.empty() below). This will be overwritten
+ during load_parameter_state().
+
+ ** Note: Torch optimizer's state structure. **
+ The Torch optimizer stores its state in two levels. The top level is a
+ list of groups, where each group contains a list of integer indexes
+ (corresponding to parameters) that index into a master parameter list
+ that is shared by all groups. As such, three values are necessary for
+ maintaining this ordering:
+
+ - group_index : The group to which a parameter belongs.
+ - group_order : The index of a parameter within its group.
+ - state_order : The index of a parameter within the shared parameter
+ list.
+ """
+
+ # Get the Torch optimizer's state dict.
+ # - This 'inner' optimizer at this point is unallocated, and only
+ # contains an integer ordering of parameters within each group, and
+ # the ordering of parameters within its flattened parameter state
+ # list.
+ inner_state_dict = self.optimizer.state_dict()
+ state_dict_param_groups = [
+ {**group, "params": list(inner_state_dict["param_groups"][idx]["params"])}
+ for idx, group in enumerate(state_dict["optimizer"]["param_groups"])
+ ]
+
+ # Allocate or retrieve optimizer state (i.e., tensors).
+ if len(self.optimizer.state) == 0:
+ # Allocate empty optimizer state if not previously initialized.
+ # - If len(self.optimizer.state) == 0, this means that the optimizer
+ # state has not been previously initialized. Once it has been
+ # initialized, we skip this code block to avoid reallocating
+ # empty tensors (i.e., torch.empty), which in turn reduces memory
+ # fragmentation.
+ # - Real data is overwritten during load_parameter_state().
+ state_dict_state = []
+ for gbuf_range_maps in self.gbuf_ranges:
+ for gbuf_range_map_for_all_buckets in gbuf_range_maps.values():
+ for gbuf_range_map in gbuf_range_map_for_all_buckets:
+ for model_param, param_range_map in gbuf_range_map["param_map"].items():
+
+ # Get parameter ordering information (see method docstring
+ # for details).
+ group_index, group_order = self.model_param_group_index_map[model_param]
+ state_order = inner_state_dict["param_groups"][group_index]["params"][
+ group_order
+ ]
+
+ # Allocate dummy tensors.
+ numel = len(param_range_map["gbuf_world"])
+ init_shard = lambda: torch.empty(
+ (numel,), dtype=torch.float32, device=torch.cuda.current_device()
+ )
+
+ state_dict_state.append(
+ (state_order, {"exp_avg": init_shard(), "exp_avg_sq": init_shard()})
+ )
+
+ # Sort by state order (see method docstring for details).
+ state_dict_state.sort(key=lambda s: s[0])
+ state_dict_state = {s[0]: s[1] for s in state_dict_state}
+
+ else:
+ # Retrieve existing optimizer state.
+ state_dict_state = inner_state_dict["state"]
+
+ # Extract 'step', for non-Apex/TE support.
+ if not HAVE_APEX_OR_TE:
+ steps = list(set([g["step"] for g in state_dict["optimizer"]["param_groups"]]))
+ assert len(steps) == 1
+ step = torch.tensor(steps[0], dtype=torch.float)
+
+ for s in state_dict_state.values():
+ # Native PyTorch state dict requires step (i.e., iteration).
+ s["step"] = step
+
+ # Optimizer.
+ self.optimizer.load_state_dict(
+ {"state": state_dict_state, "param_groups": state_dict_param_groups}
+ )
+
+ # Grad scaler.
+ if 'grad_scaler' not in state_dict:
+ if self.config.fp16:
+ logger.info(
+ '***WARNING*** found an old checkpoint, will not ' 'load grad scaler ...'
+ )
+ else:
+ if self.grad_scaler:
+ self.grad_scaler.load_state_dict(state_dict['grad_scaler'])
+ else:
+ logger.info(
+ '***WARNING*** fould the grad scaler in the '
+ 'checkpoint but it is None in the class. '
+ 'Skipping loading grad scaler ...'
+ )
+
+ if 'param_state' in state_dict:
+ assert 'param_state_sharding_type' in state_dict, state_dict.keys()
+ param_state = state_dict['param_state']
+ sharding_type = state_dict['param_state_sharding_type']
+ logger.info(f'Loading distributed optimizer sharded state of type {sharding_type}')
+ if sharding_type == 'dp_zero_gather_scatter':
+ self.load_parameter_state_from_dp_zero(param_state)
+ elif sharding_type == 'fully_sharded_bucket_space':
+ self.load_parameter_state_from_fs_bucket_space(param_state)
+ elif sharding_type == 'fully_sharded_model_space':
+ self.load_parameter_state_from_fs_model_space(param_state)
+ else:
+ raise NotImplementedError(f'Unknown sharding_type: {sharding_type}')
+
+ def get_parameter_state_fs_bucket_space(self):
+ """Get internal representation of parameter state without any copies and modifications.
+
+ This is referred to as "fully sharded bucket space" because the optimizer state is
+ fully sharded (e.g. no gather involved) and bucket-centric (the state
+ follows the internal structure of the Distributed Optimizer buckets)
+ as opposed to model-centric (typical structure of PyT optimizers)
+ """
+ state = {
+ "per_bucket_numel": self.per_bucket_numel,
+ "per_bucket_numel_unpadded": self.per_bucket_numel_unpadded,
+ }
+ for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges):
+
+ # Iterate grad buffers (by data type).
+ dtype_state = {}
+ assert len(gbuf_range_maps) == 1, "single dtype supported, for now."
+ for dtype, gbuf_range_map_for_all_buckets in gbuf_range_maps.items():
+ buckets_state = []
+ for bucket_idx, gbuf_range_map in enumerate(gbuf_range_map_for_all_buckets):
+ bucket_state = []
+ for model_param, param_range_map in gbuf_range_map["param_map"].items():
+
+ # Main param & optimizer states.
+ group_index, group_order = self.model_param_group_index_map[model_param]
+ main_param = self.optimizer.param_groups[group_index]["params"][group_order]
+ optim_state = self.optimizer.state[main_param]
+
+ tensors = {
+ "param": main_param,
+ **optim_state,
+ "gbuf_local_start": param_range_map["gbuf_local"].start,
+ "gbuf_local_end": param_range_map["gbuf_local"].end,
+ }
+ bucket_state.append(tensors)
+ buckets_state.append(bucket_state)
+ dtype_state[dtype] = buckets_state
+ state[gbuf_idx] = dtype_state
+ return state
+
+ def get_parameter_state_dp_zero(self):
+ """Get parameter state (i.e., parameter & optimizer tensors).
+
+ This method performs two steps:
+ - For each DP rank, copy param & optimizer shards to contiguous CPU
+ buffers (e.g., one buffer each for main_param, exp_avg, and
+ exp_avg_sq).
+ - Gather contiguous buffers on DP rank 0 and concatenate to world
+ buffers.
+ """
+
+ # Data parallelism variables.
+ data_parallel_world_size = self.data_parallel_group_gloo.size()
+ data_parallel_rank = torch.distributed.get_rank(self.data_parallel_group_gloo)
+ data_parallel_group_gloo = self.data_parallel_group_gloo
+ data_parallel_global_ranks = torch.distributed.get_process_group_ranks(
+ self.data_parallel_group_gloo
+ )
+
+ # Collect param states.
+ state = {"buckets_coalesced": True}
+ for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges):
+
+ # Iterate grad buffers (by data type).
+ dtype_state = {}
+ assert len(gbuf_range_maps) == 1, "single dtype supported, for now."
+ for dtype, gbuf_range_map_for_all_buckets in gbuf_range_maps.items():
+ buffer_numel_unpadded = self.buffers[gbuf_idx].numel_unpadded
+ # Create coalesced tensors for all state related to parameters in this buffer.
+ world_tensors = {}
+ if data_parallel_rank == 0:
+ world_tensors = {
+ key: torch.zeros(
+ (buffer_numel_unpadded,), dtype=torch.float32, device="cpu"
+ )
+ for key in ("param", "exp_avg", "exp_avg_sq")
+ }
+ world_tensors["numel_unpadded"] = buffer_numel_unpadded
+ offset_in_world_tensors = 0
+ for bucket_idx, gbuf_range_map in enumerate(gbuf_range_map_for_all_buckets):
+
+ # Compute local DP contiguous shard's size.
+ gbuf_world_numel = self.buffers[gbuf_idx].buckets[bucket_idx].grad_data.numel()
+ assert gbuf_world_numel % data_parallel_world_size == 0
+ gbuf_local_numel = gbuf_world_numel // data_parallel_world_size
+
+ gbuf_world_numel_unpadded = (
+ self.buffers[gbuf_idx].buckets[bucket_idx].numel_unpadded
+ )
+ assert gbuf_world_numel_unpadded <= gbuf_world_numel
+
+ local_shards = {
+ key: torch.zeros((gbuf_local_numel,), dtype=torch.float32, device="cpu")
+ for key in ("param", "exp_avg", "exp_avg_sq")
+ }
+
+ # Build contiguous DP rank shards (for param + optim states).
+ for model_param, param_range_map in gbuf_range_map["param_map"].items():
+
+ # Main param & optimizer states.
+ group_index, group_order = self.model_param_group_index_map[model_param]
+ main_param = self.optimizer.param_groups[group_index]["params"][group_order]
+ optim_state = self.optimizer.state[main_param]
+
+ tensors = {"param": main_param, **optim_state}
+
+ # Copy states into contiguous shard.
+ gbuf_local_start = param_range_map["gbuf_local"].start
+ gbuf_local_end = param_range_map["gbuf_local"].end
+ for key in local_shards:
+ local_shards[key][gbuf_local_start:gbuf_local_end].data.copy_(
+ tensors[key].detach().cpu()
+ )
+
+ # Gather contiguous shards on DP rank 0.
+ for key, send_tensor in local_shards.items():
+
+ # Gather tensor list.
+ if data_parallel_rank == 0:
+ recv_tensors = [
+ torch.zeros((gbuf_local_numel,), dtype=torch.float32, device="cpu")
+ for _ in range(data_parallel_world_size)
+ ]
+ else:
+ recv_tensors = None
+
+ # Gather.
+ torch.distributed.gather(
+ send_tensor,
+ recv_tensors,
+ data_parallel_global_ranks[0],
+ data_parallel_group_gloo,
+ )
+
+ # Concatenate.
+ if data_parallel_rank == 0:
+ recv_tensors_concatenated = torch.cat(recv_tensors)
+ # Copy this bucket's collected all-gather tensors into the right place
+ # in the tensor for the buffer. The tensor for the buffer gets rid of
+ # the padding between buckets.
+ start = offset_in_world_tensors
+ end = offset_in_world_tensors + gbuf_world_numel_unpadded
+ world_tensors[key][start:end].copy_(
+ recv_tensors_concatenated[:gbuf_world_numel_unpadded]
+ )
+
+ offset_in_world_tensors += gbuf_world_numel_unpadded
+
+ # Collect world state.
+ dtype_state[dtype] = world_tensors
+ state[gbuf_idx] = dtype_state
+
+ return state
+
+ def save_parameter_state(self, filename: str):
+ """Save the distributed parameter state on DP rank 0.
+
+ Args:
+ filename (str): path to save parameter state to.
+ """
+
+ state_dict = self.get_parameter_state_dp_zero()
+ if torch.distributed.get_rank(self.data_parallel_group) == 0:
+ torch.save(state_dict, filename)
+
+ def sharded_state_dict(
+ self,
+ model_sharded_state_dict: ShardedStateDict,
+ is_loading: bool = False,
+ sharding_type: str = 'fully_sharded_model_space',
+ ):
+ """
+ Chooses between 3 param state sharding implementations as requested by `sharding_type`.
+
+ Regular state dict parameters are saved on DP rank 0 and loaded on all ranks.
+ """
+ if not is_loading and sharding_type == 'fully_sharded_bucket_space':
+ logger.warning(
+ '`fully_sharded_bucket_space` sharding for DistributedOptimizer'
+ ' checkpoint is deprecated and will be removed in the future.'
+ ' Please switch to `full_sharded_model_space`.'
+ )
+
+ state_dict = self.state_dict()
+ if sharding_type != 'fully_sharded_model_space':
+ # State dict differs between different model parallel groups
+ state_dict = {
+ k: ShardedObject(
+ f'optimizer.distributed.dp_group_idx_{self.data_parallel_group_idx}.{k}',
+ v,
+ (1,),
+ (0,),
+ replica_id=torch.distributed.get_rank(self.data_parallel_group),
+ )
+ for k, v in state_dict.items()
+ }
+
+ if is_loading:
+ # Call the distributed optimizer's specialized load_state_dict(),
+ # which conditionally skips re-allocating the optimizer's state if
+ # already initialized, which in turn reduces memory fragmentation.
+ self.load_state_dict(self.state_dict())
+
+ if sharding_type == 'fully_sharded_bucket_space':
+ param_state = self.sharded_param_state_fs_bucket_space(
+ model_sharded_state_dict, is_loading
+ )
+
+ elif sharding_type == 'dp_zero_gather_scatter':
+ param_state = self.sharded_param_state_dp_zero(model_sharded_state_dict, is_loading)
+ elif sharding_type == 'fully_sharded_model_space':
+ param_state = self.sharded_param_state_fs_model_space(
+ model_sharded_state_dict, is_loading
+ )
+ else:
+ raise NotImplementedError(f'Unknown sharding_type: {sharding_type}')
+
+ state_dict['param_state'] = param_state
+ state_dict['param_state_sharding_type'] = sharding_type
+ return state_dict
+
+ def sharded_param_state_dp_zero(
+ self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False
+ ):
+ """Naive implementation which reuses gather/scatter from the legacy ckpt format.
+
+ During saving, gathers the parameters state on DP rank 0 and saves a ShardedObject
+ with fixed TPxPP structure. During loading, loads the saved data on DP rank 0
+ (None on other ranks). Relies on the parameters scatter done in load_state_dict.
+ """
+ if is_loading:
+ param_state_data = None
+ else:
+ # Gather on rank 0
+ param_state_data = self.get_parameter_state_dp_zero()
+
+ if torch.distributed.get_rank(self.data_parallel_group) == 0:
+ # Fixed TPxPP. Save on DP rank 0 only
+ param_state = ShardedObject(
+ f'optimizer.distributed.dp_group_idx_{self.data_parallel_group_idx}.param_state',
+ param_state_data,
+ (1,),
+ (0,),
+ )
+ else:
+ # DP ranks > 0 don't save. During loading, the param_state needs to be None.
+ param_state = LocalNonpersistentObject(None)
+
+ return param_state
+
+ def sharded_param_state_fs_bucket_space(
+ self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False
+ ):
+ """Sharded state dict where each noncontiguous buffer is a separate ShardedTensor.
+
+ Results in fully parallel save and load without any inter-process
+ communication or intermediate buffers/copies.
+ """
+ data_parallel_rank = torch.distributed.get_rank(self.data_parallel_group)
+ data_parallel_world_size = torch.distributed.get_world_size(self.data_parallel_group)
+
+ state = self.get_parameter_state_fs_bucket_space()
+ # per_bucket_numel metadata is saved separately for each TPxPP domain.
+ for per_bucket_key in ('per_bucket_numel', 'per_bucket_numel_unpadded'):
+ key = (
+ f'optimizer.distributed.dp_group_idx_{self.data_parallel_group_idx}'
+ f'.{per_bucket_key}'
+ )
+ state[per_bucket_key] = ShardedObject(
+ key, state[per_bucket_key], (1,), (0,), replica_id=data_parallel_rank
+ )
+
+ for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges):
+ for dtype, gbuf_range_map_for_all_buckets in state[gbuf_idx].items():
+ for bucket_idx, bucket_state in enumerate(gbuf_range_map_for_all_buckets):
+ # Compute local DP contiguous shard's size.
+ gbuf_world_numel = self.buffers[gbuf_idx].buckets[bucket_idx].grad_data.numel()
+ assert gbuf_world_numel % data_parallel_world_size == 0
+ gbuf_local_numel = gbuf_world_numel // data_parallel_world_size
+
+ sharded_bucket_key = (
+ f'optimizer.distributed.dp_group_idx_{self.data_parallel_group_idx}'
+ f'.gbuf_idx_{gbuf_idx}.dtype_{dtype}.bucket_idx_{bucket_idx}'
+ )
+
+ # The global ckpt tensors must be fully covered.
+ # We add extra empty padding if necessary
+ assert bucket_state, 'empty bucket encountered'
+
+ # Insert padding between parameter tensors to ensure full coverage as needed.
+ all_pad_tensors = {}
+ for i in range(len(bucket_state) - 1):
+ next_param_start = bucket_state[i + 1]['gbuf_local_start']
+ cur_param_end = bucket_state[i]['gbuf_local_end']
+ if next_param_start != cur_param_end:
+ pad_tensors = {
+ k: torch.empty(
+ next_param_start - cur_param_end, dtype=v.dtype, device=v.device
+ )
+ for k, v in bucket_state[i].items()
+ if isinstance(v, torch.Tensor)
+ }
+ all_pad_tensors[i + 1] = {
+ **pad_tensors,
+ 'gbuf_local_start': cur_param_end,
+ 'gbuf_local_end': next_param_start,
+ 'padding': True,
+ }
+
+ # Insert from end so that insertion positions are still correct.
+ indices_to_insert = sorted(list(all_pad_tensors.keys()))
+ for index_to_insert in reversed(indices_to_insert):
+ bucket_state.insert(index_to_insert, all_pad_tensors[index_to_insert])
+
+ if bucket_state[-1]['gbuf_local_end'] != gbuf_local_numel:
+ pad_tensors = {
+ k: torch.empty(
+ gbuf_local_numel - bucket_state[-1]['gbuf_local_end'],
+ dtype=v.dtype,
+ device=v.device,
+ )
+ for k, v in bucket_state[-1].items()
+ if isinstance(v, torch.Tensor)
+ }
+ bucket_state.append(
+ {
+ **pad_tensors,
+ 'gbuf_local_start': bucket_state[-1]['gbuf_local_end'],
+ 'gbuf_local_end': gbuf_local_numel,
+ 'padding': True,
+ }
+ )
+
+ # Each tensor is mapped to a slice (`flattened_range`)
+ # of a DP-local shard of size `gbuf_local_numel`.
+ for bucket_params_idx in range(len(bucket_state)):
+ tensors = bucket_state[bucket_params_idx]
+ gbuf_local_start = tensors.pop('gbuf_local_start')
+ gbuf_local_end = tensors.pop('gbuf_local_end')
+ if 'padding' not in tensors:
+ tensors['padding'] = False
+
+ for key in tensors:
+ if key == 'padding':
+ tensors[key] = LocalNonpersistentObject(tensors[key])
+ continue
+ assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), (
+ tensors[key].shape,
+ gbuf_local_start,
+ gbuf_local_end,
+ )
+
+ tensors[key] = ShardedTensor(
+ f'{sharded_bucket_key}.{key}',
+ tensors[key],
+ tensors[key].dtype,
+ (gbuf_local_numel,),
+ (data_parallel_world_size * gbuf_local_numel,),
+ (data_parallel_rank * gbuf_local_numel,),
+ axis_fragmentations=(data_parallel_world_size,),
+ flattened_range=slice(gbuf_local_start, gbuf_local_end),
+ allow_shape_mismatch=True,
+ )
+ return state
+
+ def sharded_param_state_fs_model_space(
+ self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False
+ ):
+ """Sharded state dict where each buffer is mapped to corresponding model param.
+
+ In this approach the optimizer state tensors are directly related to model parameters
+ by linking them with metadata from `model_sharded_state_dict`.
+ This will allow changing TP and PP while using DistOpt (as with other optimizers).
+ """
+
+ param_to_sharded_metadata = {}
+ model_sharded_state_dict, _ = extract_sharded_tensors_and_factories(
+ model_sharded_state_dict
+ )
+ for sh_base in nested_values(model_sharded_state_dict):
+ param_to_sharded_metadata[sh_base.data] = sh_base
+
+ prefix = 'optimizer.state'
+ state = {}
+
+ # Not stored in the checkpoint, used only to identify params in
+ # `sharded_param_state_fs_model_space`.
+ param_idx = 0
+ for gbuf_range_maps in self.gbuf_ranges:
+ for gbuf_range_map_for_all_buckets in gbuf_range_maps.values():
+ for gbuf_range_map in gbuf_range_map_for_all_buckets:
+ for model_param, param_range_map in gbuf_range_map["param_map"].items():
+ group_index, group_order = self.model_param_group_index_map[model_param]
+ param_range = param_range_map['param']
+
+ main_param = self.optimizer.param_groups[group_index]["params"][group_order]
+ optim_state = self.optimizer.state[main_param]
+
+ tensors = {"fp32_param": main_param, **optim_state}
+ # Match optimizer parameter with model ShardedTensor (or
+ # ShardedTensorFactory).
+ try:
+ sharded_metadata = param_to_sharded_metadata[model_param]
+ except KeyError as e:
+ raise ValueError(
+ f'Model param {model_param} not in model_sharded_state_dict'
+ ) from e
+
+ # Set DP corresponding replica_id coordinate to 0.
+ assert (
+ len(sharded_metadata.replica_id) == 3
+ ), f'Expected replica_id format (PP, TP, DP), got: {sharded_metadata}'
+ replica_id = (*sharded_metadata.replica_id[:2], 0)
+
+ # Instantiate ShardedTensor (or ShardedTensorFactory) for optimizer
+ # params.
+ for state_key, state_ten in tensors.items():
+ replace_kwargs = dict(
+ key=f'{prefix}.{state_key}.{sharded_metadata.key}',
+ data=state_ten,
+ dtype=state_ten.dtype,
+ flattened_range=slice(param_range.start, param_range.end),
+ replica_id=replica_id,
+ )
+ if isinstance(sharded_metadata, ShardedTensorFactory):
+ replace_kwargs.pop('dtype')
+ tensors[state_key] = replace(sharded_metadata, **replace_kwargs)
+ tensors[state_key].validate_metadata_integrity()
+ state[param_idx] = tensors
+ param_idx += 1
+ return state
+
+ def load_parameter_state_from_fs_bucket_space(self, state_dict):
+ """Loads the parameter state from an internal representation.
+
+ Inverse of the `get_parameter_state_fs_bucket_space` method.
+ """
+ logger.warning(
+ '`fully_sharded_bucket_space` sharding for DistributedOptimizer'
+ 'checkpoint is deprecated. Please switch to `full_sharded_model_space`'
+ )
+
+ if state_dict is not None and "per_bucket_numel_unpadded" in state_dict:
+ per_bucket_numel_unpadded_in_checkpoint = state_dict["per_bucket_numel_unpadded"]
+ assert self.per_bucket_numel_unpadded == per_bucket_numel_unpadded_in_checkpoint, (
+ f"Number of unpadded elements in each bucket need to be the same in current run "
+ f"({self.per_bucket_numel_unpadded}) and checkpoint "
+ f"({per_bucket_numel_unpadded_in_checkpoint})"
+ )
+
+ for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges):
+ assert len(gbuf_range_maps) == 1, "single dtype supported, for now."
+ for dtype, gbuf_range_map_for_all_buckets in gbuf_range_maps.items():
+ for bucket_idx, gbuf_range_map in enumerate(gbuf_range_map_for_all_buckets):
+ bucket_state = state_dict[gbuf_idx][dtype][bucket_idx]
+ bucket_state = [
+ bucket_state_elem
+ for bucket_state_elem in bucket_state
+ if not bucket_state_elem['padding']
+ ]
+
+ assert len(bucket_state) == len(gbuf_range_map["param_map"]), (
+ len(bucket_state),
+ len(gbuf_range_map["param_map"]),
+ )
+ for src_tensors, (model_param, param_range_map) in zip(
+ bucket_state, gbuf_range_map["param_map"].items()
+ ):
+ # Main param & optimizer states.
+ group_index, group_order = self.model_param_group_index_map[model_param]
+ main_param = self.optimizer.param_groups[group_index]["params"][group_order]
+ optim_state = self.optimizer.state[main_param]
+
+ dst_tensors = {"param": main_param, **optim_state}
+ for key in dst_tensors:
+ dst_tensors[key].copy_(src_tensors[key])
+
+ @torch.no_grad()
+ def load_parameter_state_from_fs_model_space(self, state_dict):
+ """Loads the parameter state from a "model space" representation.
+
+ Inverse of the `sharded_param_state_fs_model_space` method.
+ """
+ param_idx = 0 # matching order with `sharded_param_state_fs_model_space`
+ for gbuf_range_maps in self.gbuf_ranges:
+ for gbuf_range_map_for_all_buckets in gbuf_range_maps.values():
+ for gbuf_range_map in gbuf_range_map_for_all_buckets:
+ for model_param, param_range_map in gbuf_range_map["param_map"].items():
+ group_index, group_order = self.model_param_group_index_map[model_param]
+ main_param = self.optimizer.param_groups[group_index]["params"][group_order]
+ optim_state = self.optimizer.state[main_param]
+
+ src_tensors = state_dict[param_idx]
+ dst_tensors = {"fp32_param": main_param, **optim_state}
+ for key in dst_tensors:
+ dst_tensors[key].copy_(src_tensors[key])
+
+ param_idx += 1
+
+ @classmethod
+ def _update_legacy_world_tensors(cls, old_tensors, new_numels):
+ '''Reshard buckets (where each bucket is a tensor) to new target
+ numels, where the total numel remains the same.'''
+
+ old_total = sum([t.numel() for t in old_tensors])
+ new_total = sum(new_numels)
+
+ assert old_total == new_total
+
+ unified_tensor = torch.cat(old_tensors, dim=0)
+
+ new_tensors = []
+ start_idx = 0
+ for new_numel in new_numels:
+ new_tensors.append(unified_tensor[start_idx : (start_idx + new_numel)])
+ start_idx += new_numel
+
+ return new_tensors
+
+ def load_parameter_state_from_dp_zero_legacy(self, state_dict):
+ """Load parameter state (i.e., parameter & optimizer tensors) from DP 0 rank,
+ using the legacy checkpoint format as described below.
+
+ The difference between this method and `load_parameter_state_from_dp_zero_modern()`
+ is that this method is used for updating the format of checkpoints that
+ were saved using code from before Feb 13, 2024. Starting on this date, a
+ new format was used (i.e., different format for the parameter mapping and
+ bucket sharding).
+
+ Use arg `--ckpt-convert-update-legacy-dist-opt-format` to call this
+ method, along with `--ckpt-convert-format` and `--ckpt-convert-save` to
+ update a legacy-format checkpoint to the modern format.
+ """
+
+ # Data parallelism variables.
+ data_parallel_world_size = self.data_parallel_group_gloo.size()
+ data_parallel_rank = torch.distributed.get_rank(self.data_parallel_group_gloo)
+ data_parallel_group_gloo = self.data_parallel_group_gloo
+ data_parallel_global_ranks = torch.distributed.get_process_group_ranks(
+ self.data_parallel_group_gloo
+ )
+
+ # Scatter tensors to all DP ranks.
+ for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges):
+ for dtype, gbuf_range_map_for_all_buckets in gbuf_range_maps.items():
+ if data_parallel_rank == 0:
+ buffer_numel_unpadded = self.buffers[gbuf_idx].numel_unpadded
+ model_numels = [b.numel_unpadded for b in self.buffers[gbuf_idx].buckets]
+ checkpoint_numels = [
+ t.numel() for t in state_dict[gbuf_idx][torch.float32]["param"]
+ ]
+ assert sum(model_numels) == sum(checkpoint_numels)
+ for key in ("param", "exp_avg", "exp_avg_sq"):
+ legacy_world_tensors = self._update_legacy_world_tensors(
+ state_dict[gbuf_idx][torch.float32][key],
+ [
+ self.buffers[gbuf_idx].buckets[bi].numel_unpadded
+ for bi in range(len(gbuf_range_map_for_all_buckets))
+ ],
+ )
+ offset_in_world_tensors = 0
+ for bucket_idx, gbuf_range_map in enumerate(gbuf_range_map_for_all_buckets):
+ # Compute local DP contiguous shard's size.
+ gbuf_world_numel = (
+ self.buffers[gbuf_idx].buckets[bucket_idx].grad_data.numel()
+ )
+ assert gbuf_world_numel % data_parallel_world_size == 0
+ gbuf_local_numel = gbuf_world_numel // data_parallel_world_size
+ gbuf_world_numel_unpadded = (
+ self.buffers[gbuf_idx].buckets[bucket_idx].numel_unpadded
+ )
+ assert gbuf_world_numel_unpadded <= gbuf_world_numel
+
+ # Contiguous local shards (received from DP rank 0).
+ recv_tensor = torch.empty(
+ (gbuf_local_numel,), dtype=torch.float32, device="cpu"
+ )
+
+ # Scatter tensor list.
+ if data_parallel_rank == 0:
+
+ start = offset_in_world_tensors
+ end = offset_in_world_tensors + gbuf_world_numel_unpadded
+
+ world_tensor = legacy_world_tensors[bucket_idx]
+ assert (
+ world_tensor.numel() == gbuf_world_numel_unpadded
+ ), "%d vs. %d." % (world_tensor.numel(), gbuf_world_numel_unpadded)
+ offset_in_world_tensors += gbuf_world_numel_unpadded
+
+ # Pad world_tensor to gbuf_world_numel. Don't pad at the front,
+ # pad at the back.
+ world_tensor = torch.nn.functional.pad(
+ world_tensor, (0, gbuf_world_numel - gbuf_world_numel_unpadded)
+ )
+ assert world_tensor.numel() == gbuf_world_numel
+ gbuf_start_idxs = list(range(0, gbuf_world_numel, gbuf_local_numel))
+ send_tensors = [
+ world_tensor[i : (i + gbuf_local_numel)] for i in gbuf_start_idxs
+ ]
+ else:
+ send_tensors = None
+
+ # Scatter.
+ torch.distributed.scatter(
+ recv_tensor,
+ send_tensors,
+ data_parallel_global_ranks[0],
+ data_parallel_group_gloo,
+ )
+
+ # Copy local contiguous shards to param/optim shards.
+ for model_param, param_range_map in gbuf_range_map["param_map"].items():
+
+ # Main param & optimizer states.
+ group_index, group_order = self.model_param_group_index_map[model_param]
+ main_param = self.optimizer.param_groups[group_index]["params"][
+ group_order
+ ]
+ if key == "param":
+ tensor_to_copy_into = main_param
+ else:
+ optim_state = self.optimizer.state[main_param]
+ tensor_to_copy_into = optim_state[key]
+
+ # Copy states into contiguous shard.
+ gbuf_local_start = param_range_map["gbuf_local"].start
+ gbuf_local_end = param_range_map["gbuf_local"].end
+ tensor_to_copy_into.data.copy_(
+ recv_tensor[gbuf_local_start:gbuf_local_end]
+ )
+
+ def load_parameter_state_from_dp_zero(self, state_dict, *, update_legacy_format=False):
+ """Load parameter state (i.e., parameter & optimizer tensors) from DP 0 rank,
+ using the new checkpoint format with coalesced state across buckets.
+
+ This method performs the reverse of get_parameter_state_dp_zero():
+ - Scatter contiguous buffers from DP rank 0 to each DP rank (each DP
+ rank receives its relevant subset of the world buffers).
+ - For each DP rank, copy param & optimizer shards from contiguous CPU
+ buffers. (e.g., one buffer each for main_param, exp_avg, and
+ exp_avg_sq).
+ """
+
+ # Selectively load from a legacy checkpoint. The legacy format was used
+ # prior to Feb 13, 2024.
+ if update_legacy_format:
+ return self.load_parameter_state_from_dp_zero_legacy(state_dict)
+
+ # Data parallelism variables.
+ data_parallel_world_size = self.data_parallel_group_gloo.size()
+ data_parallel_rank = torch.distributed.get_rank(self.data_parallel_group_gloo)
+ data_parallel_group_gloo = self.data_parallel_group_gloo
+ data_parallel_global_ranks = torch.distributed.get_process_group_ranks(
+ self.data_parallel_group_gloo
+ )
+
+ if data_parallel_rank == 0:
+ # Do nothing if "--fp8-param-gather" is not used.
+ self.split_state_dict_if_needed(state_dict)
+
+ # Scatter tensors to all DP ranks.
+ for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges):
+ for dtype, gbuf_range_map_for_all_buckets in gbuf_range_maps.items():
+ if data_parallel_rank == 0:
+ buffer_numel_unpadded = self.buffers[gbuf_idx].numel_unpadded
+ checkpoint_numel_unpadded = state_dict[gbuf_idx][dtype]["numel_unpadded"]
+ assert buffer_numel_unpadded == checkpoint_numel_unpadded, (
+ f"Number of unpadded elements must be same in current run "
+ f"({buffer_numel_unpadded}) and checkpoint ({checkpoint_numel_unpadded})"
+ )
+ for key in ("param", "exp_avg", "exp_avg_sq"):
+ offset_in_world_tensors = 0
+ for bucket_idx, gbuf_range_map in enumerate(gbuf_range_map_for_all_buckets):
+ # Compute local DP contiguous shard's size.
+ gbuf_world_numel = (
+ self.buffers[gbuf_idx].buckets[bucket_idx].grad_data.numel()
+ )
+ assert gbuf_world_numel % data_parallel_world_size == 0
+ gbuf_local_numel = gbuf_world_numel // data_parallel_world_size
+ gbuf_world_numel_unpadded = (
+ self.buffers[gbuf_idx].buckets[bucket_idx].numel_unpadded
+ )
+ assert gbuf_world_numel_unpadded <= gbuf_world_numel
+
+ # Contiguous local shards (received from DP rank 0).
+ recv_tensor = torch.zeros(
+ (gbuf_local_numel,), dtype=torch.float32, device="cpu"
+ )
+
+ # Scatter tensor list.
+ if data_parallel_rank == 0:
+ world_tensors = state_dict[gbuf_idx][dtype][key]
+
+ start = offset_in_world_tensors
+ end = offset_in_world_tensors + gbuf_world_numel_unpadded
+ assert 0 <= start < end <= world_tensors.numel()
+ world_tensor = world_tensors[start:end]
+ offset_in_world_tensors += gbuf_world_numel_unpadded
+
+ # Pad world_tensor to gbuf_world_numel. Don't pad at the front,
+ # pad at the back.
+ world_tensor = torch.nn.functional.pad(
+ world_tensor, (0, gbuf_world_numel - gbuf_world_numel_unpadded)
+ )
+ assert world_tensor.numel() == gbuf_world_numel
+ gbuf_start_idxs = list(range(0, gbuf_world_numel, gbuf_local_numel))
+ send_tensors = [
+ world_tensor[i : (i + gbuf_local_numel)] for i in gbuf_start_idxs
+ ]
+ else:
+ send_tensors = None
+
+ # Scatter.
+ torch.distributed.scatter(
+ recv_tensor,
+ send_tensors,
+ data_parallel_global_ranks[0],
+ data_parallel_group_gloo,
+ )
+
+ # Copy local contiguous shards to param/optim shards.
+ for model_param, param_range_map in gbuf_range_map["param_map"].items():
+
+ # Main param & optimizer states.
+ group_index, group_order = self.model_param_group_index_map[model_param]
+ main_param = self.optimizer.param_groups[group_index]["params"][
+ group_order
+ ]
+ if key == "param":
+ tensor_to_copy_into = main_param
+ else:
+ optim_state = self.optimizer.state[main_param]
+ tensor_to_copy_into = optim_state[key]
+
+ # Copy states into contiguous shard.
+ gbuf_local_start = param_range_map["gbuf_local"].start
+ gbuf_local_end = param_range_map["gbuf_local"].end
+ tensor_to_copy_into.data.copy_(
+ recv_tensor[gbuf_local_start:gbuf_local_end]
+ )
+
+ def split_state_dict_if_needed(self, state_dict):
+ """
+ When "--fp8-param-gather" is disabled, weights and biases are stored in the same
+ `ParamAndGradBuffer`. So, when saving a checkpoint, the optimizer's main parameters are
+ saved in a single continuous tensor (this also applies to "exp_avg" and "exp_avg_sq").
+
+ However, when "--fp8-param-gather" is enabled, weights(in fp8 dtype) and biases(in bf16/fp16
+ dtype) are stored in separate `ParamAndGradBuffer`. Therefore, when we enabled
+ "--fp8-param-gather", and want to load a checkpoint saved without "--fp8-param-gather", we
+ need to split the weights(fp8) and biases(bf16/fp16) in the static_dict into two separate
+ tensors.
+ """
+ # Skip if there is no fp8 buffers.
+ fp8_gbuf_indices = []
+ for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges):
+ for dtype, _ in gbuf_range_maps.items():
+ if is_float8tensor(self.buffers[gbuf_idx].params[0]):
+ fp8_gbuf_indices.append(gbuf_idx)
+ if len(fp8_gbuf_indices) == 0:
+ return
+
+ dtype_to_gbuf_idx = {}
+ for key in state_dict.keys():
+ if key != 'buckets_coalesced':
+ for dtype in state_dict[key].keys():
+ assert dtype not in dtype_to_gbuf_idx
+ if dtype[0] == torch.uint8:
+ # If the `state_dict`` already contains a torch.uint8 buffer, we assumed
+ # that the fp8 weights and fp16/bf16 biases in the checkpoint are already
+ # separated. In this case, no action is required, so we can return directly.
+ return
+ dtype_to_gbuf_idx[dtype] = key
+
+ # 1. Replace the gbuf_idx in the checkpoint with the new gbuf_idx.
+ # 2. Copy the non-tensor data (i.e., the "buckets_coalesced") to `new_state_dict`.
+ new_state_dict = {'buckets_coalesced': state_dict['buckets_coalesced']}
+ for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges):
+ for dtype, _ in gbuf_range_maps.items():
+ if not is_float8tensor(self.buffers[gbuf_idx].params[0]):
+ new_state_dict[gbuf_idx] = state_dict[dtype_to_gbuf_idx[dtype]]
+
+ for fp8_gbuf_idx in fp8_gbuf_indices:
+ # Note that `self.buffers[fp8_gbuf_idx].params[0].dtype` is the dummy dtype of
+ # `Float8Tensor`, not torch.uint8.
+ non_fp8_param_and_grad_dtype = (
+ self.buffers[fp8_gbuf_idx].params[0].dtype,
+ self.buffers[fp8_gbuf_idx].grad_dtype,
+ )
+
+ # Iterate through all buffers to find the one that needs to be split.
+ non_fp8_gbuf_idx = None
+ for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges):
+ for dtype, _ in gbuf_range_maps.items():
+ if dtype == non_fp8_param_and_grad_dtype:
+ non_fp8_gbuf_idx = gbuf_idx
+ assert non_fp8_gbuf_idx is not None
+
+ # We need the fp8_flags to determine the order of weight (fp8) and bias (fp16/bf16) in
+ # the buffer.
+ index_to_fp8_map = {}
+ for index in self.buffers[fp8_gbuf_idx].param_indices:
+ assert index not in index_to_fp8_map
+ index_to_fp8_map[index] = True
+ for index in self.buffers[non_fp8_gbuf_idx].param_indices:
+ assert index not in index_to_fp8_map
+ index_to_fp8_map[index] = False
+ param_indices = (
+ self.buffers[fp8_gbuf_idx].param_indices
+ + self.buffers[non_fp8_gbuf_idx].param_indices
+ )
+ assert min(param_indices) == 0
+ assert max(param_indices) == len(param_indices) - 1
+ fp8_flags = []
+ for i in range(len(param_indices)):
+ fp8_flag.append(index_to_fp8_map[i])
+
+ fp8_buffer = self.buffers[fp8_gbuf_idx]
+ non_fp8_buffer = self.buffers[non_fp8_gbuf_idx]
+
+ fp8_idx = len(fp8_buffer.params) - 1
+ non_fp8_idx = len(non_fp8_buffer.params) - 1
+ offsets, fp8_offsets, non_fp8_offsets = [0], [0], [0]
+
+ # Because the parameters in `ParamAndGradBuffer` are traversed in reverse order, the
+ # flag here also needs to be traversed in reverse order.
+ for fp8_flag in fp8_flags[::-1]:
+ if fp8_flag:
+ numel = fp8_buffer.params[fp8_idx].nelement()
+ fp8_idx -= 1
+ offsets.append(offsets[-1] + numel)
+ fp8_offsets.append(fp8_offsets[-1] + numel)
+ else:
+ numel = non_fp8_buffer.params[non_fp8_idx].nelement()
+ non_fp8_idx -= 1
+ offsets.append(offsets[-1] + numel)
+ non_fp8_offsets.append(non_fp8_offsets[-1] + numel)
+
+ # Split the target buffer into two separate buffers.
+ fp8_state_dict, non_fp8_state_dict = {}, {}
+ for key in ['param', 'exp_avg', 'exp_avg_sq']:
+ tensor = state_dict[non_fp8_gbuf_idx][non_fp8_param_and_grad_dtype][key]
+ fp8_tensor = torch.empty([fp8_offsets[-1]], dtype=tensor.dtype)
+ non_fp8_tensor = torch.empty([non_fp8_offsets[-1]], dtype=tensor.dtype)
+
+ fp8_idx, non_fp8_idx = 0, 0
+ for i in range(len(offsets) - 1):
+ if fp8_flags[-(i + 1)]:
+ fp8_tensor[fp8_offsets[fp8_idx] : fp8_offsets[fp8_idx + 1]].copy_(
+ tensor[offsets[i] : offsets[i + 1]]
+ )
+ fp8_idx += 1
+ else:
+ non_fp8_tensor[
+ non_fp8_offsets[non_fp8_idx] : non_fp8_offsets[non_fp8_idx + 1]
+ ].copy_(tensor[offsets[i] : offsets[i + 1]])
+ non_fp8_idx += 1
+
+ fp8_state_dict[key] = fp8_tensor
+ non_fp8_state_dict[key] = non_fp8_tensor
+
+ fp8_state_dict['numel_unpadded'] = fp8_offsets[-1]
+ non_fp8_state_dict['numel_unpadded'] = non_fp8_offsets[-1]
+
+ # Add the two separate buffers into `new_state_dict`.
+ new_state_dict[fp8_gbuf_idx] = {}
+ new_state_dict[fp8_gbuf_idx][(torch.uint8, fp8_buffer.grad_dtype)] = fp8_state_dict
+ new_state_dict[non_fp8_gbuf_idx][non_fp8_param_and_grad_dtype] = non_fp8_state_dict
+
+ # Inplace update state_dict
+ state_dict.clear()
+ for key, value in new_state_dict.items():
+ state_dict[key] = value
+
+ def load_parameter_state(self, filename: str, *, update_legacy_format=False):
+ """Load the distributed parameter state from disk.
+
+ Args:
+ filename (str): path to load parameter state from.
+ """
+ state_dict = None
+ if torch.distributed.get_rank(self.data_parallel_group) == 0:
+ state_dict = torch.load(filename)
+
+ self.load_parameter_state_from_dp_zero(
+ state_dict, update_legacy_format=update_legacy_format
+ )
+
+ def zero_grad(self, set_to_none: bool = True):
+ """
+ Zeroes grads for the model related parameters, i.e., model_float16_groups
+ and model_fp32_groups. We additionally zero the remaining groups as a
+ memory optimization to reduce fragmentation; in the case of
+ set_to_none==True, the space used by this field can be safely deallocated.
+
+ Args:
+ set_to_none (bool): if true, set grads to None.
+ """
+ for groups in (
+ self.model_float16_groups,
+ self.model_fp32_groups,
+ self.shard_float16_groups, # grad empty/unused here?
+ self.shard_fp32_groups, # throws grad-access warning
+ self.shard_fp32_from_float16_groups,
+ ):
+ for group in groups:
+ _zero_grad_group_helper(group, set_to_none)
+
+ def _collect_main_grad_data_for_unscaling(self):
+ """
+ Note: this should be equivalent to the float-16 optimizer's method,
+ but written differently, so the two should be combined.
+ """
+ return [
+ param.grad.data for group in self.optimizer.param_groups for param in group["params"]
+ ]
+
+ def _get_model_and_main_params_data_float16(self):
+ """
+ Get aligned list of model and main params.
+ """
+ model_data = []
+ main_data = []
+ for model_group, main_group in zip(
+ self.shard_float16_groups, self.shard_fp32_from_float16_groups
+ ):
+ for model_param, main_param in zip(model_group, main_group):
+ model_data.append(model_param.data)
+ main_data.append(main_param.data)
+ return model_data, main_data
+
+ def _copy_model_grads_to_main_grads(self):
+ """
+ Copy model grads to main grads.
+
+ Since this step follows a reduce-scatter through the DDP's grad
+ buffer, this method is responsible for copying the updated grads
+ from the grad buffer to the main shard's grad field.
+ """
+
+ # Utility method for copying group grads.
+ def copy_group_grads(model_groups, shard_main_groups):
+ for model_group, shard_main_group in zip(model_groups, shard_main_groups):
+ for model_param, shard_main_param in zip(model_group, shard_main_group):
+
+ param_range_map = self._get_model_param_range_map(model_param)
+ param_range = param_range_map["param"]
+ assert param_range.size == shard_main_param.nelement()
+
+ model_grad = model_param.main_grad
+ shard_model_grad = model_grad.view(-1)[param_range.start : param_range.end]
+ shard_main_param.grad = shard_model_grad.float()
+
+ # Copy model groups to shard groups.
+ copy_group_grads(self.model_float16_groups, self.shard_fp32_from_float16_groups)
+ copy_group_grads(self.model_fp32_groups, self.shard_fp32_groups)
+
+ def _copy_main_params_to_model_params(self):
+ """
+ Copy main params to model params.
+
+ Since this step is followed by an all-gather through the DDP's grad
+ buffer, this method is responsible for copying the updated params
+ from the main shards into the correct position in the grad buffer.
+ """
+
+ # Utility method for copying group params.
+ def copy_group_params(shard_main_groups, model_groups):
+ for shard_main_group, model_group in zip(shard_main_groups, model_groups):
+ for shard_main_param, model_param in zip(shard_main_group, model_group):
+
+ param_range_map = self._get_model_param_range_map(model_param)
+ world_range = param_range_map["gbuf_world_in_bucket"]
+
+ assert world_range.size == shard_main_param.nelement()
+
+ gbuf_index, _, bucket_id = self.model_param_gbuf_map[model_param]
+ model_param_buffer = self.buffers[gbuf_index].buckets[bucket_id].param_data
+
+ shard_model_param = model_param_buffer.view(-1)[
+ world_range.start : world_range.end
+ ]
+
+ if is_float8tensor(model_param):
+ # 1. When "--fp8-param-gather" is disabled, the main param is first cast to
+ # BF16/FP16, and then cast to FP8, so the amax_history is calculated
+ # using BF16/FP16 param.
+ # 2. When "--fp8-param-gather" is enabled, we can cast the FP32 main param
+ # to FP8 directly, which results in slightly different results with
+ # higher speed. In theory, this does not affect convergence.
+ # TODO: The following code maintains the logic of the point-1 above. It can
+ # be deleted if it is not necessary.
+ shard_main_param = shard_main_param.to(model_param.dtype)
+
+ cast_to_fp8(
+ shard_main_param.view(1, -1),
+ model_param._fp8_meta['scaling_fwd'],
+ model_param._fp8_meta_index,
+ model_param._fp8_dtype,
+ out=shard_model_param.view(1, -1),
+ )
+ else:
+ shard_model_param.data.copy_(shard_main_param)
+
+ # Copy shard groups to model groups.
+ copy_group_params(self.shard_fp32_from_float16_groups, self.model_float16_groups)
+ copy_group_params(self.shard_fp32_groups, self.model_fp32_groups)
+
+ def _copy_model_params_to_main_params(self):
+ """
+ Copy model params to main params.
+
+ During finetuning, this method is used to reload the main params from
+ the model params. This copy does not make use of the grad buffer as
+ an intermediary.
+ """
+
+ # Utility method for copying group params.
+ def copy_group_params(model_groups, shard_main_groups):
+ for model_group, shard_main_group in zip(model_groups, shard_main_groups):
+ for model_param, shard_main_param in zip(model_group, shard_main_group):
+
+ param_range_map = self._get_model_param_range_map(model_param)
+ param_range = param_range_map["param"]
+ assert param_range.size == shard_main_param.nelement()
+
+ shard_model_param = model_param.view(-1)[param_range.start : param_range.end]
+ shard_main_param.data.copy_(shard_model_param)
+
+ # Copy model groups to shard groups.
+ copy_group_params(self.model_float16_groups, self.shard_fp32_from_float16_groups)
+ copy_group_params(self.model_fp32_groups, self.shard_fp32_groups)
+
+ def _update_fp8_scale_inv_and_amax(self):
+ """
+ If detect FP8 parameters, update their `_scale_inv` and do reduce-max for their
+ `amax_history`.
+ """
+ amaxes = []
+ scales = []
+ scale_invs = []
+ # Iterate over all parameters inside this optimizer to find FP8 parameters.
+ for buffer in self.buffers:
+ for bucket in buffer.buckets:
+ for param in bucket.params_list:
+ if is_float8tensor(param):
+ fp8_meta = param._fp8_meta['scaling_fwd']
+ fp8_meta_index = param._fp8_meta_index
+ amaxes.append(fp8_meta.amax_history[0][fp8_meta_index].view(1))
+ scales.append(fp8_meta.scale[fp8_meta_index].view(1))
+ scale_invs.append(param._scale_inv.view(1))
+ # Reset transpose cache
+ param._reset_caches()
+
+ # If there is no FP8 parameters, skip all operations.
+ if len(scales) > 0:
+ dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='cuda')
+
+ # Update scaling factors.
+ packed_scales = torch.empty(len(scales), dtype=torch.float32, device=scales[0].device)
+ packed_scale_views = [packed_scales[i].view(1) for i in range(len(scales))]
+ _multi_tensor_copy_this_to_that(scales, packed_scale_views, dummy_overflow_buf)
+ torch.reciprocal(packed_scales, out=packed_scales)
+ _multi_tensor_copy_this_to_that(packed_scale_views, scale_invs, dummy_overflow_buf)
+
+ # Reduce amaxes.
+ # Note: Assume each param has a separate amax.
+ packed_amaxes = torch.empty(len(amaxes), dtype=torch.float32, device=amaxes[0].device)
+ packed_amax_views = [packed_amaxes[i].view(1) for i in range(len(amaxes))]
+ _multi_tensor_copy_this_to_that(amaxes, packed_amax_views, dummy_overflow_buf)
+ torch.distributed.all_reduce(
+ packed_amaxes, op=torch.distributed.ReduceOp.MAX, group=self.data_parallel_group
+ )
+ _multi_tensor_copy_this_to_that(packed_amax_views, amaxes, dummy_overflow_buf)
+
+ @torch.no_grad()
+ def step_with_ready_grads(self) -> bool:
+ """Step the optimizer with ready gradients, return successful.
+ Under the hood, either launch synchronous param all-gathers or get ready to launch
+ asynchorous all-gathers that get overlapped with the next forward pass.
+ """
+ update_successful = super().step_with_ready_grads()
+
+ # If there is no FP8 parameters, this will do nothing.
+ self._update_fp8_scale_inv_and_amax()
+
+ timers = self.config.timers
+ if timers is not None:
+ timers('params-all-gather', log_level=1).start(barrier=self.config.barrier_with_L1_time)
+ # If not overlapping all-gather for parameters, launch synchronous all-gather
+ # communication calls here. If overlapping all-gather for parameters, the following
+ # the first all-gather is launched asynchronously in the next optimizer.zero_grad()
+ # call and subsequent all-gathers are launched in the forward pre-hook.
+ if not self.ddp_config.overlap_param_gather:
+ for model_chunk in self.model_chunks:
+ model_chunk.start_param_sync()
+ if timers is not None:
+ timers('params-all-gather').stop()
+
+ return update_successful
diff --git a/megatron/optimizer/grad_scaler.py b/megatron/core/optimizer/grad_scaler.py
similarity index 56%
rename from megatron/optimizer/grad_scaler.py
rename to megatron/core/optimizer/grad_scaler.py
index 66f7c907a4..abdd1e7b60 100644
--- a/megatron/optimizer/grad_scaler.py
+++ b/megatron/core/optimizer/grad_scaler.py
@@ -1,19 +1,18 @@
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
"""Megatron grad scaler."""
-from abc import ABC
-from abc import abstractmethod
+from abc import ABC, abstractmethod
+from typing import Dict
import torch
class MegatronGradScaler(ABC):
-
- def __init__(self, initial_scale):
+ def __init__(self, initial_scale: float):
"""Initialize scale value with the input initial scale."""
assert initial_scale > 0.0
- self._scale = torch.cuda.FloatTensor([initial_scale])
+ self._scale = torch.tensor([initial_scale], dtype=torch.float, device='cuda')
@property
def scale(self):
@@ -24,7 +23,7 @@ def inv_scale(self):
return self._scale.double().reciprocal().float()
@abstractmethod
- def update(self, found_inf):
+ def update(self, found_inf: bool):
pass
@abstractmethod
@@ -32,14 +31,16 @@ def state_dict(self):
pass
@abstractmethod
- def load_state_dict(self, state_dict):
+ def load_state_dict(self, state_dict: Dict):
pass
-
class ConstantGradScaler(MegatronGradScaler):
+ """
+ Constant grad scaler (loss scale is never adjusted regardless of NaNs seen in gradients).
+ """
- def update(self, found_inf):
+ def update(self, found_inf: bool):
pass
def state_dict(self):
@@ -49,26 +50,48 @@ def load_state_dict(self, state_dict):
pass
-
class DynamicGradScaler(MegatronGradScaler):
-
- def __init__(self, initial_scale, min_scale,
- growth_factor, backoff_factor,
- growth_interval, hysteresis):
- """"Grad scaler with dynamic scale that gets adjusted
- during training."""
+ """
+ Grad scaler with dynamic scale that gets adjusted during training.
+
+ Reduces loss scale by `backoff_factor` if `hysteresis` number of NaNs are seen in a row. Increases
+ loss scale by `growth_factor` if NaNs are not seen for `growth_interval` iterations.
+ """
+
+ def __init__(
+ self,
+ initial_scale: float,
+ min_scale: float,
+ growth_factor: float,
+ backoff_factor: float,
+ growth_interval: int,
+ hysteresis: int,
+ ):
+ """
+ Grad scaler with dynamic scale that gets adjusted during training.
+
+ Args:
+ initial_scale (float): Initial loss scale value.
+ min_scale (float): Minimum loss scale value.
+ growth_factor (float): Factor to grow loss scale by if NaNs are not seen in `growth_interval`
+ training iterations. Must be greater than 1.
+ backoff_factor (float): Factor to decrease loss scale by if NaNs are seen in `hysteresis`
+ consecutive training iterations. Must be between 0 and 1.
+ growth_interval (int): Number of training iterations of no NaNs before loss scale is increased.
+ hysteresis (int): Number of training iterations of consecutive NaNs before loss scale is decreased.
+ """
super(DynamicGradScaler, self).__init__(initial_scale)
# Lower bound on the scale.
assert min_scale > 0.0
assert min_scale <= initial_scale
- self.min_scale = torch.cuda.FloatTensor([min_scale])
+ self.min_scale = torch.tensor([min_scale], dtype=torch.float, device='cuda')
# Growth and backoff factors for the scale.
assert growth_factor > 1.0
- self.growth_factor = torch.cuda.FloatTensor([growth_factor])
+ self.growth_factor = torch.tensor([growth_factor], dtype=torch.float, device='cuda')
assert backoff_factor < 1.0
assert backoff_factor > 0.0
- self.backoff_factor = torch.cuda.FloatTensor([backoff_factor])
+ self.backoff_factor = torch.tensor([backoff_factor], dtype=torch.float, device='cuda')
# Interval over which if we don't see any inf/nan,
# we will scale the grad scale by the growth factor.
assert growth_interval > 0
@@ -82,8 +105,10 @@ def __init__(self, initial_scale, min_scale,
self._growth_tracker = 0
self._hysteresis_tracker = self.hysteresis
-
- def update(self, found_inf):
+ def update(self, found_inf: bool):
+ """
+ Updates internal state in grad scaler based on whether NaNs are seen in grads or not.
+ """
# If we have an inf/nan, growth tracker is set to 0
# and hysterisis tracker is reduced by 1.
@@ -92,8 +117,7 @@ def update(self, found_inf):
self._hysteresis_tracker -= 1
# Now if we are out of hysteresis count, scale down the loss.
if self._hysteresis_tracker <= 0:
- self._scale = torch.max(self._scale * self.backoff_factor,
- self.min_scale)
+ self._scale = torch.max(self._scale * self.backoff_factor, self.min_scale)
else:
# If there is no nan/inf, increment the growth tracker.
self._growth_tracker += 1
@@ -105,7 +129,6 @@ def update(self, found_inf):
# and scale up the loss scale.
self._scale = self._scale * self.growth_factor
-
def state_dict(self):
state_dict = {}
state_dict['scale'] = self._scale
@@ -113,8 +136,7 @@ def state_dict(self):
state_dict['hysteresis_tracker'] = self._hysteresis_tracker
return state_dict
-
- def load_state_dict(self, state_dict):
+ def load_state_dict(self, state_dict: Dict):
self._scale = state_dict['scale'].cuda(torch.cuda.current_device())
self._growth_tracker = state_dict['growth_tracker']
self._hysteresis_tracker = state_dict['hysteresis_tracker']
diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py
new file mode 100644
index 0000000000..7f2bbc0832
--- /dev/null
+++ b/megatron/core/optimizer/optimizer.py
@@ -0,0 +1,1069 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Megatron optimizer."""
+
+import copy
+import math
+import warnings
+from abc import ABC, abstractmethod
+from itertools import chain
+from logging import getLogger
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+
+try:
+ from transformer_engine.pytorch.optimizers import multi_tensor_applier, multi_tensor_scale
+
+ multi_tensor_scale_impl = multi_tensor_scale
+except ImportError:
+ try:
+ from apex.multi_tensor_apply import multi_tensor_applier
+ except ImportError:
+ from megatron.core.utils import local_multi_tensor_applier
+
+ multi_tensor_applier = local_multi_tensor_applier
+ try:
+ import amp_C
+
+ l2_norm_impl = amp_C.multi_tensor_l2norm
+ multi_tensor_scale_impl = amp_C.multi_tensor_scale
+ except ImportError:
+ from megatron.core.utils import local_multi_tensor_l2_norm, local_multi_tensor_scale
+
+ l2_norm_impl = local_multi_tensor_l2_norm
+ multi_tensor_scale_impl = local_multi_tensor_scale
+
+from .. import parallel_state, tensor_parallel
+from ..config_logger import has_config_logger_enabled, log_config_to_disk
+from ..dist_checkpointing.mapping import ShardedStateDict
+from ..dist_checkpointing.optimizer import (
+ get_param_id_to_sharded_param_map,
+ make_sharded_optimizer_tensor,
+ optim_state_to_sharding_state,
+)
+from ..dist_checkpointing.utils import add_prefix_for_sharding
+from ..transformer.module import param_is_not_shared
+from .clip_grads import clip_grad_by_total_norm_fp32, count_zeros_fp32, get_grad_norm_fp32
+from .grad_scaler import MegatronGradScaler
+from .optimizer_config import OptimizerConfig
+
+logger = getLogger(__name__)
+
+
+def _zero_grad_group_helper(group: List[torch.nn.Parameter], set_to_none: bool):
+ """
+ Zero out the gradient for a group of parameters.
+ Note: copied from torch.optim.optimizer.
+ """
+ for param in group:
+ if param.grad is not None:
+ if set_to_none:
+ param.grad = None
+ else:
+ if param.grad.grad_fn is not None:
+ param.grad.detach_()
+ else:
+ param.grad.requires_grad_(False)
+ param.grad.zero_()
+
+
+def _multi_tensor_copy_this_to_that(
+ this: List[torch.Tensor], that: List[torch.Tensor], overflow_buf: Optional[torch.Tensor] = None
+):
+ """
+ Use multi-tensor-applier to copy values from one list to another.
+ We don't have a bfloat16 implementation so for now if the overflow_buf
+ is not provided, we default back to simple loop copy to be compatible
+ with bfloat16.
+ """
+ if overflow_buf:
+ overflow_buf.fill_(0)
+ # Scaling with factor `1.0` is equivalent to copy.
+ multi_tensor_applier(multi_tensor_scale_impl, overflow_buf, [this, that], 1.0)
+ else:
+ for this_, that_ in zip(this, that):
+ that_.copy_(this_)
+
+
+class MegatronOptimizer(ABC):
+ """
+ Base class for all Megatron optimizers.
+
+ Args:
+ optimizer (torch.optim.Optimizer): base optimizer such as Adam or SGD.
+ config (OptimizerConfig): configuration object for optimizer.
+ init_state_fn (Callable, optional): function to initialize state in the optimizer.
+ """
+
+ def __init__(
+ self,
+ optimizer: torch.optim.Optimizer,
+ config: OptimizerConfig,
+ init_state_fn: Callable = lambda x: None,
+ ):
+ """Input optimizer is the base optimizer (e.g., Adam)."""
+ self.optimizer = optimizer
+ assert self.optimizer, 'no optimizer is provided.'
+ self.config = config
+ self.init_state_fn = init_state_fn
+
+ def get_parameters(self) -> List[torch.nn.Parameter]:
+ """
+ Get list of parameters wrapped in optimizer.
+ """
+ params = []
+ for param_group in self.optimizer.param_groups:
+ for param in param_group['params']:
+ params.append(param)
+ return params
+
+ def get_main_grads_for_grad_norm(self) -> List[torch.Tensor]:
+ """
+ Get main_grads that should be taken into account to compute the grad norm.
+ Filter parameters based on:
+ - grad should not be None.
+ - parameter should not be shared (i.e., grads shouldn't be double counted while
+ computing norms).
+ - should not be a replica due to tensor model parallelism.
+ """
+ params = self.get_parameters()
+ grads_for_norm = []
+ for param in params:
+ grad = param.grad
+ grad_not_none = grad is not None
+ is_not_shared = param_is_not_shared(param)
+ is_not_tp_duplicate = tensor_parallel.param_is_not_tensor_parallel_duplicate(param)
+ if grad_not_none and is_not_shared and is_not_tp_duplicate:
+ grads_for_norm.append(grad)
+
+ return grads_for_norm
+
+ def get_model_parallel_group(self) -> torch.distributed.ProcessGroup:
+ """Default returned here, but the distributed optimizer overrides this."""
+ if hasattr(self, 'model_parallel_group'):
+ return self.model_parallel_group
+ return parallel_state.get_model_parallel_group()
+
+ @abstractmethod
+ def prepare_grads(self) -> bool:
+ """Pre-processing gradients before the optimizer step, returns whether inf/nan is found."""
+ return False
+
+ @abstractmethod
+ def step_with_ready_grads(self) -> bool:
+ """Step the optimizer with ready gradients, return successful."""
+ return True
+
+ @torch.no_grad()
+ def get_grad_norm(self):
+ """Compute and return grad norm."""
+ grads_for_norm = self.get_main_grads_for_grad_norm()
+ total_norm = get_grad_norm_fp32(
+ grads_for_norm, model_parallel_group=self.get_model_parallel_group()
+ )
+ return total_norm
+
+ def clip_grad_norm(self, clip_grad: float) -> float:
+ """Compute and return grad norm, also clip grads."""
+ params = self.get_parameters()
+ grads_for_norm = self.get_main_grads_for_grad_norm()
+ grad_norm = get_grad_norm_fp32(
+ grads_for_norm, model_parallel_group=self.get_model_parallel_group()
+ )
+ clip_grad_by_total_norm_fp32(params, clip_grad, grad_norm)
+ return grad_norm
+
+ def count_zeros(self) -> float:
+ """Count number of zeros in model's gradients."""
+ params = self.get_parameters()
+ return count_zeros_fp32(params, model_parallel_group=self.get_model_parallel_group())
+
+ @abstractmethod
+ def zero_grad(self, set_to_none: bool = True):
+ """Zero gradients and prepare for next forward pass."""
+ pass
+
+ @abstractmethod
+ def get_loss_scale(self) -> torch.Tensor:
+ """
+ Get current loss scale factor.
+ NOTE: The output should be a CUDA tensor of size 1.
+ """
+ pass
+
+ def scale_loss(self, loss: torch.Tensor) -> torch.Tensor:
+ """Simple scaling."""
+ return self.get_loss_scale() * loss
+
+ def start_param_sync(self, model_index: int, *unused):
+ """
+ Start parameter synchronization for all optimizers.
+ This is a no-op for all non-distributed optimizers.
+ """
+ pass
+
+ @abstractmethod
+ def reload_model_params(self):
+ """Refreshes any internal state from the current model parameters.
+ Call whenever the parameters are changed outside of the optimizer.
+ For example, when we load a model from a checkpoint without loading
+ the optimizer, the model parameters are updated but for fp16 optimizer
+ with main parameters, the main parameters need to also be updated."""
+ pass
+
+ @abstractmethod
+ def state_dict(self):
+ """Return state_dict."""
+ pass
+
+ @abstractmethod
+ def load_state_dict(self, state_dict):
+ """Load pass-in `state_dict`."""
+ pass
+
+ # Promote state so it can be retrieved or set via
+ # "optimizer_instance.state"
+ def _get_state(self):
+ return self.optimizer.state
+
+ def _set_state(self, value):
+ self.optimizer.state = value
+
+ state = property(_get_state, _set_state)
+
+ # Promote param_groups so it can be retrieved or set via
+ # "optimizer_instance.param_groups"
+ # (for example, to adjust the learning rate)
+ def _get_param_groups(self):
+ return self.optimizer.param_groups
+
+ def _set_param_groups(self, value):
+ self.optimizer.param_groups = value
+
+ param_groups = property(_get_param_groups, _set_param_groups)
+
+ @abstractmethod
+ def step(self):
+ """Step the optimizer."""
+ pass
+
+ @abstractmethod
+ def sharded_state_dict(
+ self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False
+ ) -> ShardedStateDict:
+ """Builds sharded state dict for the optimizer, based on model's sharded state dict.
+
+ Args:
+ model_sharded_state_dict (ShardedStateDict): sharded state dict of the model
+ is_loading (bool, optional): flag indicating whether the state dict will be
+ used to save or load the optimizer state. Defaults to False.
+
+ Returns: optimizer sharded state dict
+ """
+
+ @staticmethod
+ def _extract_common_per_param_step(state_dict) -> Union[int, torch.Tensor]:
+ common_step = None
+ for param_idx, param_state in state_dict['state'].items():
+ param_step = param_state.get('step', None)
+ if param_step is not None:
+ if common_step is None:
+ common_step = param_step
+ elif common_step != param_step:
+ raise ValueError(
+ "The optimizer step differs per parameter. Mcore only supports "
+ "optimizers whose step is shared across all parameters."
+ )
+ return common_step
+
+ @staticmethod
+ def _restore_common_per_param_step(state_dict: Dict, step: Union[int, torch.Tensor]):
+ for param_idx, param_state in state_dict['state'].items():
+ param_state['step'] = copy.deepcopy(step)
+
+
+class MixedPrecisionOptimizer(MegatronOptimizer):
+ """Base class for both the float-16 and the distributed optimizer.
+
+ Args:
+ optimizer (torch.optim.Optimizer): base optimizer such as Adam or SGD.
+ config (OptimizerConfig): configuration object for optimizer.
+ grad_scaler (MegatronGradScaler): used for scaling gradients. Note that
+ this can be None. This case happens when `bf16 = True` and we don't
+ use any loss scale. Note that for `bf16 = True`, we can have
+ a constant gradient scaler. Also for `bf16 = False`, we
+ always require a grad scaler.
+ init_state_fn (Callable, optional): function to initialize state in the optimizer.
+ """
+
+ def __init__(
+ self,
+ optimizer: torch.optim.Optimizer,
+ config: OptimizerConfig,
+ grad_scaler: Optional[MegatronGradScaler],
+ init_state_fn: Callable,
+ ):
+ if has_config_logger_enabled(config):
+ log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+ super().__init__(optimizer, config, init_state_fn)
+ self.grad_scaler = grad_scaler
+
+ # None grad scaler is only supported for bf16.
+ if self.grad_scaler is None:
+ assert not self.config.fp16, 'fp16 expects a grad scaler.'
+
+ # Tensor used to determine if a nan/if has happend.
+ # Any non-zero value indicates inf/nan.
+ # Note that we keep this for the cases that grad scaler is none.
+ # We still record nan/inf if we have a bfloat16 with a grad scaler.
+ if self.grad_scaler:
+ self.found_inf = torch.tensor([0.0], dtype=torch.float, device='cuda')
+
+ # Dummy tensor needed for apex multi-apply tensor.
+ # For bfloat, we don't have multi-tensor apply and for now
+ # we set it to none so the multi-tensor apply gets ignored.
+ if self.config.bf16:
+ self._dummy_overflow_buf = None
+ else:
+ self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device='cuda')
+
+ # In case grad scaler is not passed, define the unity scale.
+ if self.grad_scaler is None:
+ self._scale_one = torch.tensor([1.0], dtype=torch.float, device='cuda')
+
+ def get_loss_scale(self):
+ if self.grad_scaler is None:
+ return self._scale_one
+ return self.grad_scaler.scale
+
+ def reload_model_params(self):
+ self._copy_model_params_to_main_params()
+
+ def _unscale_main_grads_and_check_for_nan(self):
+
+ # Collect main grads.
+ main_grads = self._collect_main_grad_data_for_unscaling()
+
+ # Reset found inf.
+ self.found_inf.fill_(0.0)
+
+ # Unscale and set found inf/nan
+ torch._amp_foreach_non_finite_check_and_unscale_(
+ main_grads, self.found_inf, self.grad_scaler.inv_scale
+ )
+
+ # Update across all model parallel instances.
+ torch.distributed.all_reduce(
+ self.found_inf, op=torch.distributed.ReduceOp.MAX, group=self.get_model_parallel_group()
+ )
+
+ # Check for nan.
+ found_inf_flag = self.found_inf.item() > 0
+
+ return found_inf_flag
+
+ @torch.no_grad()
+ def prepare_grads(self) -> bool:
+ """Pre-processing gradients before the optimizer step, returns whether inf/nan is found."""
+ timers = self.config.timers
+
+ # Copy gradients from model params to main params.
+ if timers is not None:
+ timers('optimizer-copy-to-main-grad', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ self._copy_model_grads_to_main_grads()
+ if timers is not None:
+ timers('optimizer-copy-to-main-grad').stop()
+
+ # Do unscale, check for inf, and update grad scaler only for
+ # the case that grad scaler is provided.
+ if self.grad_scaler:
+
+ # Unscale and check for inf/nan.
+ if timers is not None:
+ timers('optimizer-unscale-and-check-inf', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ found_inf_flag = self._unscale_main_grads_and_check_for_nan()
+ if timers is not None:
+ timers('optimizer-unscale-and-check-inf').stop()
+
+ # We are done with scaling gradients
+ # so we can update the loss scale.
+ self.grad_scaler.update(found_inf_flag)
+
+ return found_inf_flag
+
+ return False
+
+ @torch.no_grad()
+ def step_with_ready_grads(self) -> bool:
+ """Step the optimizer with ready gradients, return successful."""
+ timers = self.config.timers
+ # Step the optimizer.
+ if timers is not None:
+ timers('optimizer-inner-step', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ self.optimizer.step()
+ if timers is not None:
+ timers('optimizer-inner-step').stop()
+
+ # Update params from main params.
+ if timers is not None:
+ timers('optimizer-copy-main-to-model-params', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ self._copy_main_params_to_model_params()
+ if timers is not None:
+ timers('optimizer-copy-main-to-model-params').stop()
+
+ return True
+
+ @torch.no_grad()
+ def step(self):
+ timers = self.config.timers
+
+ found_inf_flag = self.prepare_grads()
+ if found_inf_flag:
+ return False, None, None
+
+ # Clip the main gradients.
+ if timers is not None:
+ timers('optimizer-clip-main-grad', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ grad_norm = None
+ if self.config.clip_grad > 0.0:
+ grad_norm = self.clip_grad_norm(self.config.clip_grad)
+ if timers is not None:
+ timers('optimizer-clip-main-grad').stop()
+
+ # Count the zeros in the grads.
+ if timers is not None:
+ timers('optimizer-count-zeros', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ num_zeros_in_grad = self.count_zeros() if self.config.log_num_zeros_in_grad else None
+ if timers is not None:
+ timers('optimizer-count-zeros').stop()
+
+ success = self.step_with_ready_grads()
+
+ # Successful update.
+ return success, grad_norm, num_zeros_in_grad
+
+
+class Float16OptimizerWithFloat16Params(MixedPrecisionOptimizer):
+ """Float16 optimizer for fp16 and bf16 data types.
+
+ Args:
+ optimizer (torch.optim.Optimizer): base optimizer such as Adam or SGD.
+ config (OptimizerConfig): configuration object for optimizer.
+ grad_scaler (MegatronGradScaler): used for scaling gradients. Note that
+ this can be None. This case happens when `bf16 = True` and we don't
+ use any loss scale. Note that for `bf16 = True`, we can have
+ a constant gradient scaler. Also for `bf16 = False`, we
+ always require a grad scaler.
+ init_state_fn (Callable, optional): function to initialize state in the optimizer.
+ """
+
+ def __init__(
+ self,
+ optimizer: torch.optim.Optimizer,
+ config: OptimizerConfig,
+ grad_scaler: MegatronGradScaler,
+ init_state_fn: Callable,
+ ):
+
+ super().__init__(optimizer, config, grad_scaler, init_state_fn)
+
+ # Handle main parameters.
+
+ # Three groups of parameters:
+ # float16_groups: original float16 parameters
+ # fp32_from_float16_groups: fp32 copy of float16 parameters
+ # fp32_from_fp32_groups: original fp32 parameters
+ self.float16_groups = []
+ self.fp32_from_float16_groups = []
+ self.fp32_from_fp32_groups = []
+
+ # For all the groups in the original optimizer:
+ for param_group in self.optimizer.param_groups:
+ float16_params_this_group = []
+ fp32_params_this_group = []
+ fp32_from_float16_params_this_group = []
+ # For all the parameters in this group:
+ for i, param in enumerate(param_group['params']):
+ if param.requires_grad:
+
+ # float16 params:
+ if param.type() in ['torch.cuda.HalfTensor', 'torch.cuda.BFloat16Tensor']:
+ float16_params_this_group.append(param)
+ # Create a copy
+ main_param = param.detach().clone().float()
+ # Copy tensor model parallel attributes.
+ tensor_parallel.copy_tensor_model_parallel_attributes(main_param, param)
+ if hasattr(param, 'shared'):
+ main_param.shared = param.shared
+ # Replace the optimizer params with the new fp32 copy.
+ param_group['params'][i] = main_param
+
+ fp32_from_float16_params_this_group.append(main_param)
+ # Reset existing state dict key to the new main param.
+ if param in self.optimizer.state:
+ self.optimizer.state[main_param] = self.optimizer.state.pop(param)
+ # fp32 params.
+ elif param.type() == 'torch.cuda.FloatTensor':
+ fp32_params_this_group.append(param)
+ param_group['params'][i] = param
+
+ else:
+ raise TypeError(
+ 'Wrapped parameters must be one of '
+ 'torch.cuda.FloatTensor, '
+ 'torch.cuda.HalfTensor, or '
+ 'torch.cuda.BFloat16Tensor. '
+ 'Received {}'.format(param.type())
+ )
+
+ self.float16_groups.append(float16_params_this_group)
+ self.fp32_from_float16_groups.append(fp32_from_float16_params_this_group)
+ self.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+ def zero_grad(self, set_to_none=True):
+ """We only need to zero the model related parameters, i.e.,
+ float16_groups & fp32_from_fp32_groups. We additionally zero
+ fp32_from_float16_groups as a memory optimization to reduce
+ fragmentation; in the case of set_to_none==True, the space
+ used by this field can be safely deallocated at this point."""
+ for group in self.float16_groups:
+ _zero_grad_group_helper(group, set_to_none)
+ for group in self.fp32_from_float16_groups:
+ _zero_grad_group_helper(group, set_to_none)
+ for group in self.fp32_from_fp32_groups:
+ _zero_grad_group_helper(group, set_to_none)
+
+ def _collect_main_grad_data_for_unscaling(self):
+
+ main_grads = []
+
+ # fp32 params from float16 ones.
+ for main_group in self.fp32_from_float16_groups:
+ for main_param in main_group:
+ if main_param.grad is not None:
+ main_grads.append(main_param.grad.data)
+
+ # Append fp32 parameters.
+ for main_group in self.fp32_from_fp32_groups:
+ for main_param in main_group:
+ if main_param.grad is not None:
+ main_grads.append(main_param.grad.data)
+
+ return main_grads
+
+ def _get_model_and_main_params_data_float16(self):
+ model_data = []
+ main_data = []
+ for model_group, main_group in zip(self.float16_groups, self.fp32_from_float16_groups):
+ for model_param, main_param in zip(model_group, main_group):
+ model_data.append(model_param.data)
+ main_data.append(main_param.data)
+ return model_data, main_data
+
+ def _copy_model_grads_to_main_grads(self):
+ # This only needs to be done for the float16 group.
+ for model_group, main_group in zip(self.float16_groups, self.fp32_from_float16_groups):
+ for model_param, main_param in zip(model_group, main_group):
+ if hasattr(model_param, 'main_grad'):
+ main_param.grad = model_param.main_grad.float()
+ else:
+ if model_param.grad is not None:
+ main_param.grad = model_param.grad.float()
+
+ # Safe to deallocate model's grad/main_grad after copying.
+ # (If using contiguous buffers, main_grad's memory should
+ # persist and therefore should not be deallocated.)
+ model_param.grad = None
+
+ # For fp32 grads, we need to reset the grads to main grad.
+ for model_group in self.fp32_from_fp32_groups:
+ for model_param in model_group:
+ model_param.grad = model_param.main_grad
+
+ def _copy_main_params_to_model_params(self):
+ # Only needed for the float16 params.
+ model_data, main_data = self._get_model_and_main_params_data_float16()
+ _multi_tensor_copy_this_to_that(
+ this=main_data, that=model_data, overflow_buf=self._dummy_overflow_buf
+ )
+
+ def _copy_model_params_to_main_params(self):
+ # Only needed for the float16 params.
+ model_data, main_data = self._get_model_and_main_params_data_float16()
+ _multi_tensor_copy_this_to_that(
+ this=model_data, that=main_data, overflow_buf=self._dummy_overflow_buf
+ )
+
+ def state_dict(self):
+ state_dict = {}
+ state_dict['optimizer'] = self.optimizer.state_dict()
+ if self.grad_scaler:
+ state_dict['grad_scaler'] = self.grad_scaler.state_dict()
+ state_dict['fp32_from_fp16_params'] = self.fp32_from_float16_groups
+ return state_dict
+
+ def sharded_state_dict(
+ self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False
+ ):
+
+ if is_loading:
+ self.init_state_fn(self.optimizer)
+
+ state_dict = self.state_dict()
+
+ id_to_sharded_param_map = get_param_id_to_sharded_param_map(
+ model_sharded_state_dict, chain.from_iterable(g for g in self.float16_groups)
+ )
+
+ # Convert fp32_from_fp16_params
+ assert len(state_dict['fp32_from_fp16_params']) == len(
+ state_dict['optimizer']['param_groups']
+ )
+ state_dict['fp32_from_fp16_params'] = [
+ [
+ make_sharded_optimizer_tensor(
+ id_to_sharded_param_map[param_id],
+ fp32_param,
+ prefix=f'optimizer.state.fp32_param',
+ )
+ for param_id, fp32_param in zip(state_group['params'], fp32_group)
+ ]
+ for fp32_group, state_group in zip(
+ state_dict['fp32_from_fp16_params'], state_dict['optimizer']['param_groups']
+ )
+ ]
+
+ step = self._extract_common_per_param_step(state_dict['optimizer'])
+
+ # Convert regular optimizer state
+ # all optimizer parameters passed to optim_state_to_sharding_state are
+ # expected to have the same shape as the model parameters,
+ # so we save the step separately and ignore it here
+ optim_state_to_sharding_state(
+ state_dict['optimizer'], id_to_sharded_param_map, exclude_keys="step"
+ )
+ # save step as a shared step among all parameters. Separate per-parameter
+ # steps are not supported
+ state_dict['optimizer']['state']['common_step'] = step
+ return state_dict
+
+ def load_state_dict(self, state_dict):
+ pipeline_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
+ # Optimizer.
+ optimizer_key = 'optimizer'
+ if optimizer_key not in state_dict:
+ optimizer_key = 'optimizer_state_dict'
+ logger.info('***WARNING*** loading optimizer from ' 'an old checkpoint ...')
+ if 'common_step' in state_dict[optimizer_key]['state']:
+ common_step = state_dict[optimizer_key]['state'].pop('common_step')
+ self._restore_common_per_param_step(state_dict[optimizer_key], common_step)
+ self.optimizer.load_state_dict(state_dict[optimizer_key])
+
+ # Grad scaler.
+ if 'grad_scaler' not in state_dict:
+ if self.config.fp16:
+ logger.info(
+ '***WARNING*** found an old checkpoint, will not ' 'load grad scaler ...'
+ )
+ else:
+ if self.grad_scaler:
+ self.grad_scaler.load_state_dict(state_dict['grad_scaler'])
+ else:
+ logger.info(
+ '***WARNING*** fould the grad scaler in the '
+ 'checkpoint but it is None in the class. '
+ 'Skipping loading grad scaler ...'
+ )
+
+ # Copy data for the main params.
+ fp32_from_float16_params_key = 'fp32_from_fp16_params'
+ if fp32_from_float16_params_key not in state_dict:
+ fp32_from_float16_params_key = 'fp32_from_fp16'
+ for current_group, saved_group in zip(
+ self.fp32_from_float16_groups, state_dict[fp32_from_float16_params_key]
+ ):
+ for current_param, saved_param in zip(current_group, saved_group):
+ current_param.data.copy_(saved_param.data)
+
+
+class FP32Optimizer(MegatronOptimizer):
+ """Float32 optimizer.
+
+ Args:
+ optimizer (torch.optim.Optimizer): base optimizer such as Adam or SGD.
+ config (OptimizerConfig): configuration object for optimizer.
+ init_state_fn (Callable, optional): function to initialize state in the optimizer.
+ """
+
+ def __init__(
+ self, optimizer: torch.optim.Optimizer, config: OptimizerConfig, init_state_fn: Callable
+ ):
+ if has_config_logger_enabled(config):
+ log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+ super(FP32Optimizer, self).__init__(optimizer, config, init_state_fn)
+
+ self._scale = torch.tensor([1.0], dtype=torch.float, device='cuda')
+
+ def zero_grad(self, set_to_none=True):
+ """Copied from torch.optim.optimizer"""
+ for group in self.optimizer.param_groups:
+ _zero_grad_group_helper(group['params'], set_to_none)
+
+ def get_loss_scale(self):
+ """FP32 optimizer does not do any scaling."""
+ return self._scale
+
+ @torch.no_grad()
+ def prepare_grads(self) -> bool:
+ """Pre-processing gradients before the optimizer step, returns whether inf/nan is found."""
+ timers = self.config.timers
+
+ # Copy main_grads to grads.
+ if timers is not None:
+ timers('optimizer-copy-to-main-grad', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ for param_group in self.optimizer.param_groups:
+ for param in param_group['params']:
+ param.grad = param.main_grad
+ if timers is not None:
+ timers('optimizer-copy-to-main-grad').stop()
+
+ return False
+
+ @torch.no_grad()
+ def step_with_ready_grads(self) -> bool:
+ """Step the optimizer with ready gradients, return successful."""
+ timers = self.config.timers
+
+ # Update parameters.
+ if timers is not None:
+ timers('optimizer-inner-step', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ self.optimizer.step()
+ if timers is not None:
+ timers('optimizer-inner-step').stop()
+
+ return True
+
+ @torch.no_grad()
+ def step(self):
+ """Clip gradients (if needed) and step the base optimizer.
+ Always return successful since there is no overflow."""
+ timers = self.config.timers
+
+ found_inf_flag = self.prepare_grads()
+ if found_inf_flag:
+ return False, None, None
+
+ # Clip gradients.
+ if timers is not None:
+ timers('optimizer-clip-main-grad', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ grad_norm = None
+ if self.config.clip_grad > 0.0:
+ grad_norm = self.clip_grad_norm(self.config.clip_grad)
+ if timers is not None:
+ timers('optimizer-clip-main-grad').stop()
+
+ # Count the zeros in the grads.
+ if timers is not None:
+ timers('optimizer-count-zeros', log_level=1).start(
+ barrier=self.config.barrier_with_L1_time
+ )
+ num_zeros_in_grad = self.count_zeros() if self.config.log_num_zeros_in_grad else None
+ if timers is not None:
+ timers('optimizer-count-zeros').stop()
+
+ success = self.step_with_ready_grads()
+
+ # No overflow for FP32 optimizer.
+ return success, grad_norm, num_zeros_in_grad
+
+ def reload_model_params(self):
+ pass
+
+ def state_dict(self):
+ return self.optimizer.state_dict()
+
+ def load_state_dict(self, state_dict):
+ pipeline_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
+ if 'common_step' in state_dict['state']:
+ common_step = state_dict['state'].pop('common_step')
+ self._restore_common_per_param_step(state_dict, common_step)
+ self.optimizer.load_state_dict(state_dict)
+
+ def sharded_state_dict(
+ self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False
+ ):
+ if is_loading:
+ self.init_state_fn(self.optimizer)
+
+ state_dict = self.state_dict()
+ id_to_sharded_param_map = get_param_id_to_sharded_param_map(
+ model_sharded_state_dict, self.get_parameters()
+ )
+ step = self._extract_common_per_param_step(state_dict)
+
+ # all optimizer parameters passed to optim_state_to_sharding_state are
+ # expected to have the same shape as the model parameters,
+ # so we save the step separately and ignore it here
+ optim_state_to_sharding_state(state_dict, id_to_sharded_param_map, exclude_keys="step")
+ # save step as a shared step among all parameters. Separate per-parameter
+ # steps are not supported
+ state_dict['state']['common_step'] = step
+ return state_dict
+
+
+class ProxyDict:
+ """
+ A dictionary-like object that proxies to a list of dictionaries.
+
+ e.g., ProxyDict([{'a': 1}, {'b': 2}]) behaves like:
+ {
+ (0, 'a'): 1,
+ (1, 'b'): 2,
+ }
+ We use tuples as keys to avoid ambiguity with the keys of the inner dicts.
+ """
+
+ def __init__(self, inner_dicts: List[dict]):
+ self._inner_dicts = inner_dicts
+
+ def __getitem__(self, key: Tuple[int, str]):
+ idx, inner_key = key
+ return self._inner_dicts[idx].get(inner_key)
+
+ def __setitem__(self, key: Tuple[int, str], value: Any):
+ idx, inner_key = key
+ self._inner_dicts[idx][inner_key] = value
+
+ def __len__(self) -> int:
+ return sum([len(inner_dict) for inner_dict in self._inner_dicts])
+
+ def __iter__(self):
+ for idx, inner_dict in enumerate(self._inner_dicts):
+ for inner_key in inner_dict:
+ yield (idx, inner_key)
+
+ def items(self):
+ """Return generator over underlying items."""
+ for idx, inner_dict in enumerate(self._inner_dicts):
+ for inner_key, value in inner_dict.items():
+ yield (idx, inner_key), value
+
+
+class ChainedOptimizer(MegatronOptimizer):
+ """ChainedOptimizer is designed for a collection of optimizers.
+
+ These optimizers are responsible for different parts of multiple models for
+ a training task and will be executed one-by-one when the model is updated.
+
+ Args:
+ chained_optimizers: a list of optimizers.
+ """
+
+ def __init__(self, chained_optimizers: List[MegatronOptimizer]):
+ self.model_chunks = []
+ self.config = getattr(chained_optimizers[0], 'config', None)
+ for optimizer in chained_optimizers:
+ if hasattr(optimizer, 'model_chunks'):
+ for model_chunk in optimizer.model_chunks:
+ if model_chunk not in self.model_chunks:
+ self.model_chunks.append(model_chunk)
+ assert self.config == getattr(optimizer, 'config', None)
+ self.chained_optimizers = chained_optimizers
+
+ @property
+ def param_groups(self) -> List[dict]:
+ """Get param_groups aggregated over underlying optimizers."""
+ param_groups = []
+ for optimizer in self.chained_optimizers:
+ param_groups += optimizer.param_groups
+ return param_groups
+
+ @property
+ def state(self) -> ProxyDict:
+ """
+ Return optimizer state with tuple keys, where the first element is the
+ index of the optimizer in the list of chained optimizers.
+ """
+ return ProxyDict([opt.state for opt in self.chained_optimizers])
+
+ def zero_grad(self, set_to_none=True):
+ for optimizer in self.chained_optimizers:
+ optimizer.zero_grad(set_to_none)
+
+ def get_loss_scale(self):
+ return self.chained_optimizers[0].get_loss_scale()
+
+ def reload_model_params(self):
+ for optimizer in self.chained_optimizers:
+ optimizer.reload_model_params()
+
+ def state_dict(self):
+ return [optimizer.state_dict() for optimizer in self.chained_optimizers]
+
+ def sharded_state_dict(
+ self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False, **kwargs
+ ):
+ sharded_state_dict = {}
+ for optimizer_idx, optimizer in enumerate(self.chained_optimizers):
+ optim_state_dict = optimizer.sharded_state_dict(
+ model_sharded_state_dict, is_loading, **kwargs
+ )
+ add_prefix_for_sharding(optim_state_dict, f'chained_{optimizer_idx}.')
+ sharded_state_dict[optimizer_idx] = optim_state_dict
+ return sharded_state_dict
+
+ def load_state_dict(self, state_dict):
+ if len(self.chained_optimizers) != len(state_dict):
+ raise RuntimeError(
+ f'Expected {len(self.chained_optimizers)} entries'
+ f' in state dict, but got {len(state_dict)}.'
+ )
+ if isinstance(state_dict, dict):
+ state_dict = (v for k, v in sorted(state_dict.items()))
+ for optimizer, state in zip(self.chained_optimizers, state_dict):
+ optimizer.load_state_dict(state)
+
+ @torch.no_grad()
+ def prepare_grads(self) -> bool:
+ """Pre-processing gradients before the optimizer step, returns whether inf/nan is found."""
+ found_inf_flag = False
+ for optimizer in self.chained_optimizers:
+ found_inf_flag |= optimizer.prepare_grads()
+
+ return found_inf_flag
+
+ @torch.no_grad()
+ def step_with_ready_grads(self) -> bool:
+ """Step the optimizer with ready gradients, return successful."""
+ success = True
+ for optimizer_idx, optimizer in enumerate(self.chained_optimizers):
+ success &= optimizer.step_with_ready_grads()
+ if self.config.overlap_param_gather_with_optimizer_step and optimizer_idx == 0:
+ assert success
+ assert len(optimizer.model_chunks) == 1
+ optimizer.model_chunks[0].start_param_sync(force_dispatch=True)
+
+ return success
+
+ def disable_pre_hook(self):
+ """Disable pre-hooks for underlying distributed optimizers."""
+ warnings.warn(
+ "`ChainedOptimizer.disable_pre_hook` will be deprecated in a future release. "
+ "Use `DistributedDataParallel.disable_forward_pre_hook` directly."
+ )
+ for model_chunk in self.model_chunks:
+ model_chunk.disable_forward_pre_hook()
+
+ def enable_pre_hook(self):
+ """Enable pre-hooks for underlying distributed optimizers."""
+ warnings.warn(
+ "`ChainedOptimizer.enable_pre_hook` will be deprecated in a future release. "
+ "Use `DistributedDataParallel.enable_forward_pre_hook` directly."
+ )
+ for model_chunk in self.model_chunks:
+ model_chunk.enable_forward_pre_hook()
+
+ @torch.no_grad()
+ def step(self):
+ """ChainedOptimizer will step all optimizers one by one."""
+ found_inf_flag = self.prepare_grads()
+ if found_inf_flag:
+ return False, None, None
+
+ # Get grad norm.
+ grad_norms = []
+ for optimizer in self.chained_optimizers:
+ _grad_norm = optimizer.get_grad_norm()
+ grad_norms += [_grad_norm if _grad_norm else 0.0]
+ grad_norm = math.sqrt(sum([x**2 for x in grad_norms]))
+
+ # Clip gradients.
+ for optimizer in self.chained_optimizers:
+ if optimizer.config.clip_grad > 0.0:
+ clip_grad_by_total_norm_fp32(
+ optimizer.get_parameters(),
+ max_norm=optimizer.config.clip_grad,
+ total_norm=grad_norm,
+ )
+
+ # Count the zeros in the grads.
+ num_zeros_in_grad = 0
+ for optimizer in self.chained_optimizers:
+ num_zeros_in_grad += (
+ optimizer.count_zeros() if optimizer.config.log_num_zeros_in_grad else 0
+ )
+
+ update_successful = self.step_with_ready_grads()
+
+ return update_successful, grad_norm, num_zeros_in_grad
+
+ def save_parameter_state(self, filename: str):
+ """Save the distributed parameter states of all optimizers to a file.
+
+ Args:
+ filename (str): path to save parameter state to.
+ """
+ save_states = False
+ states = []
+ for optimizer in self.chained_optimizers:
+ if hasattr(optimizer, 'get_parameter_state_dp_zero'):
+ state_dict = optimizer.get_parameter_state_dp_zero()
+
+ # Save checkpoint economically, only when DP rank = 0, state dict
+ # needs to be saved.
+ if torch.distributed.get_rank(optimizer.data_parallel_group) == 0:
+ states.append(state_dict)
+ save_states = True
+ else:
+ states.append(None)
+ else:
+ states.append(None)
+
+ if save_states:
+ torch.save(states, filename)
+
+ def load_parameter_state(self, filename: str, *, update_legacy_format: bool = False):
+ """Load the distributed parameter states of all optimizers from a file.
+
+ Args:
+ filename (str): path to load parameter state from.
+ """
+ states = None
+ for idx, optimizer in enumerate(self.chained_optimizers):
+ if not hasattr(optimizer, 'load_parameter_state_from_dp_zero'):
+ continue
+
+ # Lazy loading checkpoint, state dict is needed only when DP rank = 0.
+ if torch.distributed.get_rank(optimizer.data_parallel_group) == 0 and states is None:
+ states = torch.load(filename)
+
+ state_dict = states[idx] if states else None
+ optimizer.load_parameter_state_from_dp_zero(
+ state_dict, update_legacy_format=update_legacy_format
+ )
+
+ def start_param_sync(self, model_index: int, *unused):
+ """Start parameter synchronization for all optimizers."""
+ for optimizer in self.chained_optimizers:
+ optimizer.start_param_sync(model_index, *unused)
diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py
new file mode 100644
index 0000000000..8876d925cb
--- /dev/null
+++ b/megatron/core/optimizer/optimizer_config.py
@@ -0,0 +1,116 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import torch
+
+
+@dataclass
+class OptimizerConfig:
+ """Configuration for optimizer."""
+
+ ##############
+ # General
+ ##############
+ optimizer: str = 'adam'
+ """Optimizer to use (one of Adam or SGD)."""
+
+ lr: Optional[float] = None
+ """Initial learning rate. Depending on decay style and initial warmup, the learning rate at each
+ iteration would be different.
+ """
+
+ min_lr: Optional[float] = None
+ """Minumum value for learning rate. The scheduler clip values below this threshold."""
+
+ decoupled_lr: Optional[float] = None
+ """Separate learning rate for the input and output layer."""
+
+ decoupled_min_lr: Optional[float] = None
+ """Minimum value for learning rate for the input and output layer. The scheduler clip values
+ below this threshold.
+ """
+
+ weight_decay: float = 0.01
+ """Weight decay coefficient for L2 regularization."""
+
+ ##############
+ # Precision
+ ##############
+ fp16: bool = False
+ """If true, train with fp16 mixed precision training. Defaults to False."""
+
+ bf16: bool = False
+ """If true, train with bf16 mixed precision training. Defaults to False."""
+
+ params_dtype: torch.dtype = torch.float32
+ """dtype used when intializing the weights. Defaults to torch.float32."""
+
+ ###############
+ # Loss scaling
+ ###############
+ loss_scale: Optional[float] = None
+ """Static loss scaling, positive power of 2 values can improve fp16 convergence. If None,
+ dynamic loss scaling is used.
+ """
+
+ initial_loss_scale: float = 2**32
+ """Initial loss-scale for dynamic loss scaling."""
+
+ min_loss_scale: float = 1.0
+ """Minimum loss scale for dynamic loss scaling."""
+
+ loss_scale_window: float = 1000
+ """Window over which to raise/lower dynamic scale."""
+
+ hysteresis: int = 2
+ """Hysteresis for dynamic loss scaling."""
+
+ ##############
+ # Optimizer
+ ##############
+ # Adam
+ adam_beta1: float = 0.9
+ """First coefficient for computing running averages of gradient and its square in Adam
+ optimizer.
+ """
+
+ adam_beta2: float = 0.999
+ """Second coefficient for computing running averages of gradient and its square in Adam
+ optimizer.
+ """
+
+ adam_eps: float = 1e-08
+ """Term added to the denominator to improve numerical stability in Adam optimizer."""
+
+ # SGD.
+ sgd_momentum: float = 0.9
+ """Momentum factor for SGD optimizer."""
+
+ #######################
+ # Distributed optimizer
+ #######################
+ use_distributed_optimizer: bool = False
+ """Distribute optimizer state over data-parallel replicas."""
+
+ overlap_param_gather_with_optimizer_step: bool = False
+ """If true, overlap param all-gather of first bucket with optimizer step."""
+
+ ################
+ # Miscellaneous
+ ################
+ clip_grad: float = 1.0
+ """Gradient clipping based on global L2 norm."""
+
+ log_num_zeros_in_grad: bool = False
+ """If true, calculate and log the number of zeros in gradient."""
+
+ barrier_with_L1_time: bool = False
+ """If true, use barrier with level 1 time measurements."""
+
+ timers: Callable = None
+ """Function to get timers."""
+
+ config_logger_dir: str = ""
+ """When non-empty, dumps entry-point configs to config_logger_dir"""
diff --git a/megatron/core/optimizer_param_scheduler.py b/megatron/core/optimizer_param_scheduler.py
new file mode 100644
index 0000000000..43c106f4f5
--- /dev/null
+++ b/megatron/core/optimizer_param_scheduler.py
@@ -0,0 +1,297 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""Learning rate decay and weight decay incr functions."""
+import logging
+import math
+from typing import Optional
+
+from megatron.core.optimizer import MegatronOptimizer
+from megatron.core.utils import log_single_rank
+
+logger = logging.getLogger(__name__)
+
+
+class OptimizerParamScheduler:
+ """Anneals learning rate and weight decay
+
+ Args:
+ optimizer (MegatronOptimizer): the optimizer to be used
+ init_lr (float): initial learning rate
+ max_lr (float): maximum learning rate
+ min_lr (float): minimum learning rate
+ lr_warmup_steps (int): number of warmup steps
+ lr_decay_steps (int): number of decay steps
+ lr_decay_style (str): decay style for learning rate
+ start_wd (float): initial weight decay
+ end_wd (float): final weight decay
+ wd_incr_steps (int): number of weight decay increment steps
+ wd_incr_style (str): weight decay increment style
+ use_checkpoint_opt_param_scheduler (bool, optional): whether to use the checkpoint values
+ for the optimizer param scheduler
+ override_opt_param_scheduler (bool, optional): whether to override the optimizer param
+ scheduler values with the class values
+ wsd_decay_steps (int, optional): number of weight decay decay steps
+ lr_wsd_decay_style (str, optional): decay style for learning rate during weight decay decay
+ steps
+
+ """
+
+ def __init__(
+ self,
+ optimizer: MegatronOptimizer,
+ init_lr: float,
+ max_lr: float,
+ min_lr: float,
+ lr_warmup_steps: int,
+ lr_decay_steps: int,
+ lr_decay_style: str,
+ start_wd: float,
+ end_wd: float,
+ wd_incr_steps: int,
+ wd_incr_style: str,
+ use_checkpoint_opt_param_scheduler: Optional[bool] = True,
+ override_opt_param_scheduler: Optional[bool] = False,
+ wsd_decay_steps: Optional[int] = None,
+ lr_wsd_decay_style: Optional[str] = None,
+ ) -> None:
+
+ # Class values.
+ self.optimizer = optimizer
+
+ self.init_lr = init_lr
+ self.max_lr = float(max_lr)
+ self.min_lr = min_lr
+ assert self.min_lr >= 0.0
+ assert self.max_lr >= self.min_lr
+ assert self.init_lr <= self.max_lr
+
+ self.lr_warmup_steps = lr_warmup_steps
+ self.num_steps = 0
+ self.lr_decay_steps = lr_decay_steps
+ self.wsd_decay_steps = wsd_decay_steps
+ self.lr_wsd_decay_style = lr_wsd_decay_style
+ assert self.lr_decay_steps > 0
+ assert self.lr_warmup_steps < self.lr_decay_steps
+
+ self.lr_decay_style = lr_decay_style
+ if self.lr_decay_style == "WSD":
+ assert self.wsd_decay_steps is not None
+
+ self.start_wd = start_wd
+ self.end_wd = end_wd
+ assert self.start_wd >= 0.0
+ assert self.end_wd >= self.start_wd
+ self.wd_incr_steps = wd_incr_steps
+ self.wd_incr_style = wd_incr_style
+
+ self.override_opt_param_scheduler = override_opt_param_scheduler
+ self.use_checkpoint_opt_param_scheduler = use_checkpoint_opt_param_scheduler
+ if self.override_opt_param_scheduler:
+ assert not self.use_checkpoint_opt_param_scheduler, (
+ 'both override and ' 'use-checkpoint are set.'
+ )
+
+ # Set the learning rate
+ self.step(0)
+ log_single_rank(logger, logging.INFO, f"> learning rate decay style: {self.lr_decay_style}")
+
+ def get_wd(self) -> float:
+ """Weight decay incr functions"""
+ if self.num_steps > self.wd_incr_steps:
+ return self.end_wd
+
+ if self.wd_incr_style == 'constant':
+ assert self.start_wd == self.end_wd
+ return self.end_wd
+
+ incr_ratio = float(self.num_steps) / float(self.wd_incr_steps)
+ assert incr_ratio >= 0.0
+ assert incr_ratio <= 1.0
+ delta_wd = self.end_wd - self.start_wd
+
+ if self.wd_incr_style == 'linear':
+ coeff = incr_ratio
+ elif self.wd_incr_style == 'cosine':
+ coeff = 0.5 * (math.cos(math.pi * (1 - incr_ratio)) + 1.0)
+ else:
+ raise Exception(f'{self.wd_incr_style} weight decay increment style is not supported.')
+
+ return self.start_wd + coeff * delta_wd
+
+ def get_lr(self, param_group: dict) -> float:
+ """Learning rate decay functions from:
+ https://openreview.net/pdf?id=BJYwwY9ll pg. 4
+
+ Args:
+ param_group (dict): parameter group from the optimizer.
+ """
+
+ max_lr = param_group.get('max_lr', self.max_lr)
+ min_lr = param_group.get('min_lr', self.min_lr)
+
+ # Use linear warmup for the initial part.
+ if self.lr_warmup_steps > 0 and self.num_steps <= self.lr_warmup_steps:
+ return self.init_lr + (
+ (max_lr - self.init_lr) * float(self.num_steps) / float(self.lr_warmup_steps)
+ )
+
+ # If the learning rate is constant, just return the initial value.
+ if self.lr_decay_style == 'constant':
+ return max_lr
+
+ # For any steps larger than `self.lr_decay_steps`, use `min_lr`.
+ if self.num_steps > self.lr_decay_steps:
+ return min_lr
+
+ # If we are done with the warmup period, use the decay style.
+ if self.lr_decay_style == 'inverse-square-root':
+ warmup_steps = max(self.lr_warmup_steps, 1)
+ num_steps = max(self.num_steps, 1)
+ lr = max_lr * warmup_steps**0.5 / (num_steps**0.5)
+ return max(min_lr, lr)
+
+ num_steps_ = self.num_steps - self.lr_warmup_steps
+ decay_steps_ = self.lr_decay_steps - self.lr_warmup_steps
+ decay_ratio = float(num_steps_) / float(decay_steps_)
+ assert decay_ratio >= 0.0
+ assert decay_ratio <= 1.0
+ delta_lr = max_lr - min_lr
+
+ if self.lr_decay_style == 'linear':
+ coeff = 1.0 - decay_ratio
+ elif self.lr_decay_style == 'cosine':
+ coeff = 0.5 * (math.cos(math.pi * decay_ratio) + 1.0)
+ elif self.lr_decay_style == 'WSD':
+ wsd_anneal_start_ = self.lr_decay_steps - self.wsd_decay_steps
+ if self.num_steps <= wsd_anneal_start_:
+ coeff = 1.0
+ else:
+ wsd_steps = self.num_steps - wsd_anneal_start_
+ wsd_decay_ratio = float(wsd_steps) / float(self.wsd_decay_steps)
+ if self.lr_wsd_decay_style == "linear":
+ coeff = 1.0 - wsd_decay_ratio
+ elif self.lr_wsd_decay_style == "cosine":
+ coeff = 0.5 * (math.cos(math.pi * wsd_decay_ratio) + 1.0)
+ elif self.lr_wsd_decay_style == "exponential":
+ coeff = (2.0 * math.pow(0.5, wsd_decay_ratio)) - 1.0
+ else:
+ raise Exception(f'{self.lr_decay_style} decay style is not supported.')
+
+ return min_lr + coeff * delta_lr
+
+ def step(self, increment: int) -> None:
+ """Set lr for all parameters groups.
+
+ Args:
+ increment (int): number of steps to increment
+ """
+ self.num_steps += increment
+ new_wd = self.get_wd()
+ for param_group in self.optimizer.param_groups:
+ new_lr = self.get_lr(param_group)
+ param_group['lr'] = new_lr * param_group.get('lr_mult', 1.0)
+ param_group['weight_decay'] = new_wd * param_group.get('wd_mult', 1.0)
+
+ def state_dict(self) -> dict:
+ """Return the state dict."""
+ state_dict = {
+ 'max_lr': self.max_lr,
+ 'lr_warmup_steps': self.lr_warmup_steps,
+ 'num_steps': self.num_steps,
+ 'lr_decay_style': self.lr_decay_style,
+ 'lr_decay_steps': self.lr_decay_steps,
+ 'min_lr': self.min_lr,
+ 'start_wd': self.start_wd,
+ 'end_wd': self.end_wd,
+ 'wd_incr_style': self.wd_incr_style,
+ 'wd_incr_steps': self.wd_incr_steps,
+ }
+ return state_dict
+
+ def _check_and_set(self, cls_value: float, sd_value: float, name: str) -> float:
+ """Auxiliary function for checking the values in the checkpoint and
+ setting them.
+
+ Args:
+ cls_value (float): class value
+ sd_value (float): checkpoint value
+ name (str): name of the parameter
+ """
+
+ if self.override_opt_param_scheduler:
+ log_single_rank(logger, logging.INFO, f" > overriding {name} value to {cls_value}")
+ return cls_value
+
+ if not self.use_checkpoint_opt_param_scheduler:
+ assert cls_value == sd_value, (
+ f'OptimizerParamScheduler: class input value {cls_value} and checkpoint'
+ f'value {sd_value} for {name} do not match'
+ )
+
+ log_single_rank(logger, logging.INFO, f" > using checkpoint value {sd_value} for {name}")
+ return sd_value
+
+ def load_state_dict(self, state_dict: dict) -> None:
+ """Load the state dict.
+
+ Args:
+ state_dict (dict): state dict to be load
+ """
+
+ if 'start_lr' in state_dict:
+ max_lr_ = state_dict['start_lr']
+ else:
+ max_lr_ = state_dict['max_lr']
+ self.max_lr = self._check_and_set(self.max_lr, max_lr_, 'learning rate')
+
+ self.min_lr = self._check_and_set(
+ self.min_lr, state_dict['min_lr'], 'minimum learning rate'
+ )
+
+ if 'warmup_iter' in state_dict:
+ lr_warmup_steps_ = state_dict['warmup_iter']
+ elif 'warmup_steps' in state_dict:
+ lr_warmup_steps_ = state_dict['warmup_steps']
+ else:
+ lr_warmup_steps_ = state_dict['lr_warmup_steps']
+ self.lr_warmup_steps = self._check_and_set(
+ self.lr_warmup_steps, lr_warmup_steps_, 'warmup iterations'
+ )
+
+ if 'end_iter' in state_dict:
+ lr_decay_steps_ = state_dict['end_iter']
+ elif 'decay_steps' in state_dict:
+ lr_decay_steps_ = state_dict['decay_steps']
+ else:
+ lr_decay_steps_ = state_dict['lr_decay_steps']
+ self.lr_decay_steps = self._check_and_set(
+ self.lr_decay_steps, lr_decay_steps_, 'total number of iterations'
+ )
+
+ if 'decay_style' in state_dict:
+ lr_decay_style_ = state_dict['decay_style']
+ else:
+ lr_decay_style_ = state_dict['lr_decay_style']
+ self.lr_decay_style = self._check_and_set(
+ self.lr_decay_style, lr_decay_style_, 'learning rate decay style'
+ )
+
+ if 'num_iters' in state_dict:
+ num_steps = state_dict['num_iters']
+ else:
+ num_steps = state_dict['num_steps']
+ self.step(increment=num_steps)
+
+ if 'start_wd' in state_dict:
+ self.start_wd = self._check_and_set(
+ self.start_wd, state_dict['start_wd'], "start weight decay"
+ )
+ self.end_wd = self._check_and_set(self.end_wd, state_dict['end_wd'], "end weight decay")
+ self.wd_incr_steps = self._check_and_set(
+ self.wd_incr_steps,
+ state_dict['wd_incr_steps'],
+ "total number of weight decay iterations",
+ )
+ self.wd_incr_style = self._check_and_set(
+ self.wd_incr_style, state_dict['wd_incr_style'], "weight decay incr style"
+ )
diff --git a/megatron/core/package_info.py b/megatron/core/package_info.py
index 6f53034623..6135dc52c8 100644
--- a/megatron/core/package_info.py
+++ b/megatron/core/package_info.py
@@ -2,9 +2,9 @@
MAJOR = 0
-MINOR = 1
+MINOR = 10
PATCH = 0
-PRE_RELEASE = ''
+PRE_RELEASE = 'rc0'
# Use the following formatting: (major, minor, patch, pre-release)
VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)
@@ -14,10 +14,16 @@
__package_name__ = 'megatron_core'
__contact_names__ = 'NVIDIA'
-__contact_emails__ = 'nemo-toolkit@nvidia.com' # use NeMo Email
-__homepage__ = 'https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/' # use NeMo homepage
+__contact_emails__ = 'nemo-toolkit@nvidia.com' # use NeMo Email
+__homepage__ = (
+ 'https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/' # use NeMo homepage
+)
__repository_url__ = 'https://github.com/NVIDIA/Megatron-LM/megatron/core'
__download_url__ = 'https://github.com/NVIDIA/Megatron-LM/releases'
-__description__ = 'Megatron Core - a library for efficient and scalable training of transformer based models'
+__description__ = (
+ 'Megatron Core - a library for efficient and scalable training of transformer based models'
+)
__license__ = 'BSD-3'
-__keywords__ = 'deep learning, machine learning, gpu, NLP, NLU, language, transformer, nvidia, pytorch, torch'
+__keywords__ = (
+ 'deep learning, machine learning, gpu, NLP, NLU, language, transformer, nvidia, pytorch, torch'
+)
diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py
new file mode 100644
index 0000000000..dff0cc5992
--- /dev/null
+++ b/megatron/core/packed_seq_params.py
@@ -0,0 +1,20 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from dataclasses import dataclass
+
+from torch import Tensor
+
+
+@dataclass
+class PackedSeqParams:
+ '''
+ parameters to TEDotProductAttention and fused rope kernels for the
+ `thd` (packed) sequence format
+ '''
+
+ qkv_format: str = None
+ cu_seqlens_q: Tensor = None
+ cu_seqlens_kv: Tensor = None
+ cu_seqlens_q_padded: Tensor = None
+ cu_seqlens_kv_padded: Tensor = None
+ max_seqlen_q: Tensor = None
+ max_seqlen_kv: Tensor = None
diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py
index fcc1b6b9cf..e9043b647c 100644
--- a/megatron/core/parallel_state.py
+++ b/megatron/core/parallel_state.py
@@ -2,8 +2,14 @@
"""Model and data parallel groups."""
+import os
+import warnings
+from datetime import timedelta
+from functools import partial
+from itertools import cycle
+from typing import Callable, List, Optional
+
import torch
-from typing import Optional
from .utils import GlobalMemoryBuffer
@@ -13,6 +19,8 @@
_PIPELINE_MODEL_PARALLEL_GROUP = None
# Model parallel group (both intra- and pipeline) that the current rank belongs to.
_MODEL_PARALLEL_GROUP = None
+# Model parallel group (both intra-, pipeline, and expert) that the current rank belongs to.
+_MODEL_AND_EXPERT_PARALLEL_GROUP = None
# Embedding group.
_EMBEDDING_GROUP = None
# Position embedding group.
@@ -20,18 +28,33 @@
# Data parallel group that the current rank belongs to.
_DATA_PARALLEL_GROUP = None
_DATA_PARALLEL_GROUP_GLOO = None
-# FP8 amax reduction group.
-_AMAX_REDUCTION_GROUP = None
+# tensor model parallel group and data parallel group combined
+# used for fp8 and moe training
+_TENSOR_AND_DATA_PARALLEL_GROUP = None
+# Expert parallel group that the current rank belongs to.
+_EXPERT_MODEL_PARALLEL_GROUP = None
+_TENSOR_AND_EXPERT_PARALLEL_GROUP = None
+_DATA_MODULO_EXPERT_PARALLEL_GROUP = None
+_DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO = None
+_DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP = None
+_DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP_GLOO = None
+
_VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = None
_VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
_PIPELINE_MODEL_PARALLEL_SPLIT_RANK = None
+_PIPELINE_MODEL_PARALLEL_DECODER_START = None
+
# These values enable us to change the mpu sizes on the fly.
_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None
_MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
+_MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE = None
+_MPU_DATA_PARALLEL_WORLD_SIZE = None
+_MPU_DATA_PARALLEL_RANK = None
_MPU_TENSOR_MODEL_PARALLEL_RANK = None
_MPU_PIPELINE_MODEL_PARALLEL_RANK = None
+_MPU_EXPERT_MODEL_PARALLEL_RANK = None
# A list of ranks that have a copy of the embedding.
_EMBEDDING_GLOBAL_RANKS = None
@@ -47,20 +70,305 @@
# rank when broadcasting weights from src to all other data parallel ranks
_DATA_PARALLEL_GLOBAL_RANKS = None
+# A list of global ranks for each tensor model parallel group to ease calculation of
+# the first local rank in the tensor model parallel group
+_TENSOR_MODEL_PARALLEL_GLOBAL_RANKS = None
+
+# Context parallel group that the current rank belongs to
+_CONTEXT_PARALLEL_GROUP = None
+# A list of global ranks for each context parallel group to ease calculation of the
+# destination rank when exchanging KV/dKV between context parallel_ranks
+_CONTEXT_PARALLEL_GLOBAL_RANKS = None
+
+# Data parallel group information with context parallel combined.
+_DATA_PARALLEL_GROUP_WITH_CP = None
+_DATA_PARALLEL_GROUP_WITH_CP_GLOO = None
+_DATA_PARALLEL_GLOBAL_RANKS_WITH_CP = None
+
+# combined parallel group of TP and CP
+_TENSOR_AND_CONTEXT_PARALLEL_GROUP = None
+
+# combined parallel group of TP, DP, and CP used for fp8
+_TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP = None
+
# Memory buffers to avoid dynamic memory allocation
_GLOBAL_MEMORY_BUFFER = None
+# MOE logging
+_MOE_LAYER_WISE_LOGGING_TRACKER = {}
+
+
+def get_nccl_options(pg_name, nccl_comm_cfgs):
+ """Set the NCCL process group options.
+
+ Args:
+ pg_name (str): process group name
+ nccl_comm_cfgs (dict): nccl communicator configurations
+
+ When an option (e.g., max_ctas) is not found in the config, use the NCCL default setting.
+ """
+ if pg_name in nccl_comm_cfgs:
+ nccl_options = torch.distributed.ProcessGroupNCCL.Options()
+ nccl_options.config.cga_cluster_size = nccl_comm_cfgs[pg_name].get('cga_cluster_size', 4)
+ nccl_options.config.max_ctas = nccl_comm_cfgs[pg_name].get('max_ctas', 32)
+ nccl_options.config.min_ctas = nccl_comm_cfgs[pg_name].get('min_ctas', 1)
+ return nccl_options
+ else:
+ return None
+
+
+def generate_masked_orthogonal_rank_groups(
+ world_size: int, parallel_size: List[int], mask: List[bool]
+) -> List[List[int]]:
+ """Generate orthogonal parallel groups based on the parallel size and mask.
+
+ Arguments:
+ world_size (int): world size
+
+ parallel_size (List[int]):
+ The parallel size of each orthogonal parallel type. For example, if
+ tensor_parallel_size = 2, pipeline_model_parallel_group = 3, data_parallel_size = 4,
+ and the parallel mapping order is tp-pp-dp, then the parallel_size = [2, 3, 4].
+
+ mask (List[bool]):
+ The mask controls which parallel methods the generated groups represent. If mask[i] is
+ True, it means the generated group contains the i-th parallelism method. For example,
+ if parallel_size = [tp_size, pp_size, dp_size], and mask = [True, False , True], then
+ the generated group is the `tp-dp` group, if the mask = [False, True, False], then the
+ generated group is the `pp` group.
+
+ Algorithm:
+ For orthogonal parallelism, such as tp/dp/pp/cp, the global_rank and
+ local_rank satisfy the following equation:
+ global_rank = tp_rank + dp_rank * tp_size + pp_rank * tp_size * dp_size (1)
+ tp_rank \in [0, tp_size)
+ dp_rank \in [0, dp_size)
+ pp_rank \in [0, pp_size)
+
+ If we want to get the `dp_group` (tp_size * pp_size groups of dp_size ranks each.
+ For example, if the gpu size is 8 and order is 'tp-pp-dp', size is '2-2-2', and the
+ dp_group here is [[0, 4], [1, 5], [2, 6], [3, 7]].)
+ The tp_rank and pp_rank will be combined to form the `dp_group_index`.
+ dp_group_index = tp_rank + pp_rank * tp_size (2)
+
+ So, Given that tp_rank and pp_rank satisfy equation (2), and dp_rank in
+ range(0, dp_size), the ranks in dp_group[dp_group_index] satisfies the
+ equation (1).
+
+ This function solve this math problem.
+
+ For example, if the parallel_size = [tp_size, dp_size, pp_size] = [2, 3, 4],
+ and the mask = [False, True, False]. Then,
+ dp_group_index(0) = tp_rank(0) + pp_rank(0) * 2
+ dp_group_index(1) = tp_rank(1) + pp_rank(0) * 2
+ ...
+ dp_group_index(7) = tp_rank(1) + pp_rank(3) * 2
+
+ dp_group[0] = 0 + range(0, 3) * 2 + 0 = [0, 2, 4]
+ dp_group[1] = 1 + range(0, 3) * 2 + 0 = [1, 3, 5]
+ ...
+ dp_group[7] = 1 + range(0, 3) * 2 + 3 * 2 * 3 = [19, 21, 23]
+ """
+
+ def prefix_product(a: List[int], init=1) -> List[int]:
+ r = [init]
+ for v in a:
+ init = init * v
+ r.append(init)
+ return r
+
+ def inner_product(a: List[int], b: List[int]) -> int:
+ return sum([x * y for x, y in zip(a, b)])
+
+ def decompose(index, shape, stride=None):
+ '''
+ This function solve the math problem below:
+ There is an equation:
+ index = sum(idx[i] * stride[i])
+ And given the value of index, stride.
+ Return the idx.
+ This function will used to get the pp/dp/pp_rank
+ from group_index and rank_in_group.
+ '''
+ if stride is None:
+ stride = prefix_product(shape)
+ idx = [(index // d) % s for s, d in zip(shape, stride)]
+ # stride is a prefix_product result. And the value of stride[-1]
+ # is not used.
+ assert (
+ sum([x * y for x, y in zip(idx, stride[:-1])]) == index
+ ), "idx {} with shape {} mismatch the return idx {}".format(index, shape, idx)
+ return idx
+
+ masked_shape = [s for s, m in zip(parallel_size, mask) if m]
+ unmasked_shape = [s for s, m in zip(parallel_size, mask) if not m]
+
+ global_stride = prefix_product(parallel_size)
+ masked_stride = [d for d, m in zip(global_stride, mask) if m]
+ unmasked_stride = [d for d, m in zip(global_stride, mask) if not m]
+
+ group_size = prefix_product(masked_shape)[-1]
+ num_of_group = world_size // group_size
+
+ ranks = []
+ for group_index in range(num_of_group):
+ # get indices from unmaksed for group_index.
+ decomposed_group_idx = decompose(group_index, unmasked_shape)
+ rank = []
+ for rank_in_group in range(group_size):
+ # get indices from masked for rank_in_group.
+ decomposed_rank_idx = decompose(rank_in_group, masked_shape)
+ rank.append(
+ inner_product(decomposed_rank_idx, masked_stride)
+ + inner_product(decomposed_group_idx, unmasked_stride)
+ )
+ ranks.append(rank)
+ return ranks
+
+
+class RankGenerator(object):
+ """A class for generating rank groups for different modes of parallelism."""
+
+ def __init__(
+ self, tp: int, ep: int, dp: int, pp: int, cp: int, order: str, rank_offset: int = 0
+ ) -> None:
+ self.tp = tp
+ self.ep = ep
+ self.dp = dp
+ self.pp = pp
+ self.cp = cp
+ self.rank_offset = rank_offset
+ self.world_size = tp * dp * pp * cp
+
+ self.name_to_size = {
+ "tp": self.tp,
+ "pp": self.pp,
+ "dp": self.dp,
+ "ep": self.ep,
+ "cp": self.cp,
+ }
+ self.order = order
+ order = order.lower()
+
+ if 'ep' in order:
+ if 'ep-dp' not in order and 'dp-ep' not in order:
+ raise RuntimeError(f"The ep and dp must be adjacent in order ({self.order}).")
+
+ for name in self.name_to_size.keys():
+ if name not in order and self.name_to_size[name] != 1:
+ raise RuntimeError(
+ f"The size of ({name}) is ({self.name_to_size[name]}), but you haven't"
+ f"specified the order ({self.order})."
+ )
+ elif name not in order:
+ order = order + '-' + name
+
+ self.order_w_ep = order
+ self.order_wo_ep = '-'.join([token for token in order.split('-') if token != 'ep'])
+ self.ordered_size_wo_ep = []
+ self.ordered_size_w_ep = []
+
+ for token in order.split('-'):
+ if token == 'dp':
+ self.ordered_size_w_ep.append(self.dp // self.ep)
+ self.ordered_size_wo_ep.append(self.dp)
+ elif token == 'ep':
+ self.ordered_size_w_ep.append(self.ep)
+ else:
+ self.ordered_size_w_ep.append(self.name_to_size[token])
+ self.ordered_size_wo_ep.append(self.name_to_size[token])
+
+ def get_mask(self, order: str, token: str):
+ """Create a mask for the specified tokens based on the given order.
+
+ Args:
+ order (str): The order of parallelism types (e.g., 'tp-dp-pp').
+ token (str): The specific parallelism types to include in the mask,
+ separated by hyphens (e.g., 'tp-dp').
+ """
+ ordered_token = order.split('-')
+ token = token.split('-')
+ mask = [False] * len(ordered_token)
+ for t in token:
+ mask[ordered_token.index(t)] = True
+ return mask
+
+ def get_ranks(self, token, independent_ep=False):
+ """Get rank group by input token.
+
+ Args:
+ token (str):
+ Specify the ranks type that want to get. If we want
+ to obtain multiple parallel types, we can use a hyphen
+ '-' to separate them. For example, if we want to obtain
+ the TP_DP group, the token should be 'tp-dp'.
+
+ independent_ep (bool: True):
+ This flag controls whether we treat EP and DP independently.
+ EP shares ranks with DP, if we want to get ranks related to
+ EP, we should set the flag. For example, get_ranks('dp', True)
+ will get DP modulo EP group, and get_ranks('dp', False) will
+ get full DP group.
+ """
+ if independent_ep:
+ parallel_size = self.ordered_size_w_ep
+ order = self.order_w_ep
+ else:
+ parallel_size = self.ordered_size_wo_ep
+ order = self.order_wo_ep
+ mask = self.get_mask(order, token)
+ ranks = generate_masked_orthogonal_rank_groups(self.world_size, parallel_size, mask)
+ if self.rank_offset > 0:
+ for rank_group in ranks:
+ for i in range(len(rank_group)):
+ rank_group[i] += self.rank_offset
+ return ranks
+
+
+def default_embedding_ranks(pp_ranks, split_rank=None):
+ """Return the default ranks that constitute the stages on which the word embeddings live.
+ For most models, these are the first and last pipeline stages.
+
+ We also support the deprecated split rank argument for backwards compatibility."""
+ if len(pp_ranks) == 1:
+ return [pp_ranks[0]]
+ elif split_rank is not None and pp_ranks[split_rank] not in (pp_ranks[0], pp_ranks[-1]):
+ return [pp_ranks[0], pp_ranks[split_rank], pp_ranks[-1]]
+ else:
+ return [pp_ranks[0], pp_ranks[-1]]
+
+
+def default_position_embedding_ranks(pp_ranks, split_rank=None):
+ """Return the default ranks that constitute the stages on which the position embeddings live.
+ For most models, this is only the first pipeline stage.
+
+ We also support the deprecated split rank argument for backwards compatibility."""
+ if split_rank is not None and pp_ranks[0] != pp_ranks[split_rank]:
+ return [pp_ranks[0], pp_ranks[split_rank]]
+ else:
+ return [pp_ranks[0]]
+
def initialize_model_parallel(
tensor_model_parallel_size: int = 1,
pipeline_model_parallel_size: int = 1,
virtual_pipeline_model_parallel_size: Optional[int] = None,
pipeline_model_parallel_split_rank: Optional[int] = None,
- use_fp8: bool = False,
+ use_sharp: bool = False,
+ context_parallel_size: int = 1,
+ expert_model_parallel_size: int = 1,
+ nccl_communicator_config_path: Optional[str] = None,
+ distributed_timeout_minutes: int = 30,
+ order: str = "tp-cp-ep-dp-pp",
+ encoder_tensor_model_parallel_size: Optional[int] = 0,
+ encoder_pipeline_model_parallel_size: Optional[int] = 0,
+ get_embedding_ranks: Optional[Callable[[List[int], Optional[int]], List[int]]] = None,
+ get_position_embedding_ranks: Optional[Callable[[List[int], Optional[int]], List[int]]] = None,
) -> None:
+ # pylint: disable=line-too-long
"""Initialize model data parallel groups.
- Arguments:
+ Args:
tensor_model_parallel_size (int, default = 1):
The number of GPUs to split individual tensors across.
@@ -87,7 +395,7 @@ def initialize_model_parallel(
GPU 3: [7, 8] [15, 16]
pipeline_model_parallel_split_rank (int, optional):
- For models with both an encoder and decoder, the rank in
+ DEPRECATED. For models with both an encoder and decoder, the rank in
pipeline to switch between encoder and decoder (i.e. the
first rank of the decoder). This allows the user to set
the pipeline parallel size of the encoder and decoder
@@ -96,10 +404,72 @@ def initialize_model_parallel(
pipeline_model_parallel_split_rank is 3, then ranks 0-2
will be the encoder and ranks 3-7 will be the decoder.
- use_fp8 (bool, default = False):
- Construct GPU groups needed for FP8 training, namely for
- amax reduction across the product of the data-parallel and
- tensor-parallel groups.
+ use_sharp (bool, default = False):
+ Set the use of SHARP for the collective communications of
+ data-parallel process groups. When `True`, run barrier
+ within each data-parallel process group, which specifies
+ the SHARP application target groups.
+
+ context_parallel_size (int, default = 1):
+ The number of tensor parallel GPU groups to split the
+ network input sequence length across. Compute of attention
+ module requires tokens of full sequence length, so GPUs
+ in a context parallel group need to communicate with each
+ other to exchange information of other sequence chunks.
+ Each GPU and its counterparts in other tensor parallel
+ groups compose a context parallel group.
+
+ For example, assume we have 8 GPUs, if tensor model parallel
+ size is 4 and context parallel size is 2, the network input
+ will be split into two sequence chunks, which are processed
+ by 2 different groups of 4 GPUs. One chunk is processed by
+ GPU0-3, the other chunk is processed by GPU4-7. Four groups
+ are build to do context parallel communications: [GPU0, GPU4],
+ [GPU1, GPU5], [GPU2, GPU6], and [GPU3, GPU7].
+
+ Context parallelism partitions sequence length, so it has no
+ impact on weights, which means weights are duplicated among
+ GPUs in a context parallel group. Hence, weight gradients
+ all-reduce is required in backward. For simplicity, we piggyback
+ GPUs of context parallelism on data parallel group for
+ weight gradient all-reduce.
+
+ expert_model_parallel_size (int, default = 1):
+ The number of Mixture of Experts parallel GPUs in each expert
+ parallel group.
+
+ nccl_communicator_config_path (str, default = None):
+ Path to the yaml file of NCCL communicator configurations.
+ `min_ctas`, `max_ctas`, and `cga_cluster_size` can be set
+ for each communicator.
+
+ distributed_timeout_minutes (int, default = 30): Timeout, in
+ minutes,for operations executed against distributed
+ process groups. See PyTorch documentation at
+ https://pytorch.org/docs/stable/distributed.html for
+ caveats.
+
+ order (str, default=tp-dp-pp):
+ The rank initialization order of parallelism. Now we support
+ tp-dp-pp and tp-pp-dp orders.
+
+ encoder_tensor_model_parallel_size (int, default = 0):
+ The number of GPUs to split individual tensors across in the encoder. If 0,
+ then we use the default, decoder's tensor model parallel size.
+
+ encoder_pipeline_model_parallel_size (int, default = 0):
+ The number of tensor parallel GPU groups to allocate to the encoder. As an example,
+ if pipeline_model_parallel_size is 4 and encoder_pipeline_model_parallel_size is 2,
+ then the encoder will use the first two pipeline stages for its layers, and the total
+ amount of pipelineing is 6.
+
+ get_embedding_ranks (Callable[[List[int], Optional[int]], List[int]], optional, default=None):
+ A function that takes in a list of ranks for a pipeline group and returns
+ those ranks that should have embeddings.
+
+ get_position_embedding_ranks (Callable[[List[int], Optional[int]], List[int]], optional, default=None):
+ A function that takes in a list of ranks for a pipeline group, and returns
+ those ranks that should have position embeddings.
Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
@@ -118,27 +488,69 @@ def initialize_model_parallel(
ranks 8 to 15 belong to the second box.
"""
+ if encoder_pipeline_model_parallel_size is None:
+ encoder_pipeline_model_parallel_size = 0
+
+ if encoder_tensor_model_parallel_size == 0 and encoder_pipeline_model_parallel_size > 0:
+ encoder_tensor_model_parallel_size = tensor_model_parallel_size
+
+ if get_embedding_ranks is None:
+ get_embedding_ranks = partial(
+ default_embedding_ranks, split_rank=pipeline_model_parallel_split_rank
+ )
+
+ if get_position_embedding_ranks is None:
+ get_position_embedding_ranks = partial(
+ default_position_embedding_ranks, split_rank=pipeline_model_parallel_split_rank
+ )
+
+ if encoder_pipeline_model_parallel_size > 0:
+ global _PIPELINE_MODEL_PARALLEL_DECODER_START
+ _PIPELINE_MODEL_PARALLEL_DECODER_START = encoder_pipeline_model_parallel_size
+
# Get world size and rank. Ensure some consistencies.
assert torch.distributed.is_initialized()
world_size: int = torch.distributed.get_world_size()
- if world_size % (tensor_model_parallel_size * pipeline_model_parallel_size) != 0:
+ if encoder_tensor_model_parallel_size > 0:
+ assert encoder_pipeline_model_parallel_size > 0
+ assert (
+ encoder_tensor_model_parallel_size <= tensor_model_parallel_size
+ ), "We do not support encoders with more TP than the decoder."
+
+ encoder_model_size = (
+ encoder_tensor_model_parallel_size
+ * encoder_pipeline_model_parallel_size
+ * context_parallel_size
+ )
+ decoder_model_size = (
+ tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size
+ )
+ total_model_size = encoder_model_size + decoder_model_size
+
+ if world_size % total_model_size != 0:
+ raise RuntimeError(f"world_size ({world_size}) is not divisible by {total_model_size}")
+
+ data_parallel_size: int = world_size // total_model_size
+
+ if data_parallel_size % expert_model_parallel_size != 0:
raise RuntimeError(
- f"world_size ({world_size}) is not divisible by tensor_model_parallel_size "
- f"({tensor_model_parallel_size}) x pipeline_model_parallel_size ({pipeline_model_parallel_size})"
+ f"data_parallel_size ({data_parallel_size}) is not divisible by "
+ "expert_model_parallel_size "
)
- data_parallel_size: int = world_size // (tensor_model_parallel_size *
- pipeline_model_parallel_size)
+ encoder_world_size = encoder_model_size * data_parallel_size
+ decoder_world_size = decoder_model_size * data_parallel_size
- num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size
- num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size
- num_data_parallel_groups: int = world_size // data_parallel_size
+ assert (
+ encoder_world_size + decoder_world_size == world_size
+ ), f"{encoder_world_size=} + {decoder_world_size=} != {world_size=}"
if virtual_pipeline_model_parallel_size is not None:
- if not pipeline_model_parallel_size > 2:
- raise RuntimeError("pipeline-model-parallel size should be greater than 2 with "
- "interleaved schedule")
+ if not pipeline_model_parallel_size > 1:
+ raise RuntimeError(
+ "pipeline-model-parallel size should be greater than 1 with interleaved schedule"
+ )
global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
_VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = 0
@@ -150,108 +562,304 @@ def initialize_model_parallel(
rank = torch.distributed.get_rank()
+ nccl_comm_cfgs = {}
+ if nccl_communicator_config_path is not None:
+ try:
+ import yaml
+ except ImportError:
+ raise RuntimeError(
+ "Cannot import `yaml`. Setting custom nccl communicator configs "
+ "requires the yaml package."
+ )
+
+ with open(nccl_communicator_config_path, "r") as stream:
+ nccl_comm_cfgs = yaml.safe_load(stream)
+
+ if encoder_world_size > 0:
+ encoder_rank_generator = RankGenerator(
+ tp=encoder_tensor_model_parallel_size,
+ ep=1,
+ dp=data_parallel_size,
+ pp=encoder_pipeline_model_parallel_size,
+ cp=context_parallel_size,
+ order=order,
+ rank_offset=0,
+ )
+ else:
+ encoder_rank_generator = None
+
+ decoder_rank_generator = RankGenerator(
+ tp=tensor_model_parallel_size,
+ ep=expert_model_parallel_size,
+ dp=data_parallel_size,
+ pp=pipeline_model_parallel_size,
+ cp=context_parallel_size,
+ order=order,
+ rank_offset=encoder_world_size,
+ )
+
+ def generator_wrapper(group_type, **kwargs):
+ """The `RankGenerator` class produces a hyper-rectangle for a given set of
+ tensor, pipeline, data, expert, and context parallelism. If we have an encoder,
+ in addition to the default decoder, we essentially instantiate two `RankGenerator`
+ classes to construct the parallelism for each module separately, and we then have
+ to stitch them together for the right groups. For now, this means pp and tp-pp."""
+ d_ranks = decoder_rank_generator.get_ranks(group_type, **kwargs)
+ if encoder_rank_generator is None:
+ for x in d_ranks:
+ yield x
+ return
+ e_ranks = encoder_rank_generator.get_ranks(group_type, **kwargs)
+ if group_type == 'pp':
+ # Map 1 encoder tp rank to several decoder tp ranks, because
+ # these won't be the same size.
+ for x, y in zip(cycle(e_ranks), d_ranks):
+ yield x + y
+ elif group_type == 'tp-pp':
+ # For this group, we can just return the concatenated
+ # groups together, because their sizes are the same.
+ assert len(e_ranks) == len(d_ranks)
+ for x, y in zip(e_ranks, d_ranks):
+ yield x + y
+ else:
+ for x in e_ranks:
+ yield x
+ for x in d_ranks:
+ yield x
+
+ timeout = timedelta(minutes=distributed_timeout_minutes)
+
# Build the data-parallel groups.
global _DATA_PARALLEL_GROUP
global _DATA_PARALLEL_GROUP_GLOO
global _DATA_PARALLEL_GLOBAL_RANKS
+ global _DATA_PARALLEL_GROUP_WITH_CP
+ global _DATA_PARALLEL_GROUP_WITH_CP_GLOO
+ global _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP
assert _DATA_PARALLEL_GROUP is None, 'data parallel group is already initialized'
- all_data_parallel_group_ranks = []
- for i in range(pipeline_model_parallel_size):
- start_rank = i * num_pipeline_model_parallel_groups
- end_rank = (i + 1) * num_pipeline_model_parallel_groups
- for j in range(tensor_model_parallel_size):
- ranks = range(start_rank + j, end_rank, tensor_model_parallel_size)
- all_data_parallel_group_ranks.append(list(ranks))
- group = torch.distributed.new_group(ranks)
- group_gloo = torch.distributed.new_group(ranks, backend="gloo")
- if rank in ranks:
- _DATA_PARALLEL_GROUP = group
- _DATA_PARALLEL_GROUP_GLOO = group_gloo
- _DATA_PARALLEL_GLOBAL_RANKS = ranks
+
+ for ranks in generator_wrapper('dp'):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('dp', nccl_comm_cfgs)
+ )
+ group_gloo = torch.distributed.new_group(ranks, timeout=timeout, backend="gloo")
+ if rank in ranks:
+ _DATA_PARALLEL_GROUP = group
+ _DATA_PARALLEL_GROUP_GLOO = group_gloo
+ _DATA_PARALLEL_GLOBAL_RANKS = ranks
+
+ for ranks_with_cp in generator_wrapper('dp-cp'):
+ group_with_cp = torch.distributed.new_group(
+ ranks_with_cp, timeout=timeout, pg_options=get_nccl_options('dp_cp', nccl_comm_cfgs)
+ )
+ group_with_cp_gloo = torch.distributed.new_group(
+ ranks_with_cp, timeout=timeout, backend="gloo"
+ )
+ if rank in ranks_with_cp:
+ _DATA_PARALLEL_GROUP_WITH_CP = group_with_cp
+ _DATA_PARALLEL_GROUP_WITH_CP_GLOO = group_with_cp_gloo
+ _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP = ranks_with_cp
+
+ # Apply SHARP to DP process groups
+ if use_sharp:
+ if rank == 0:
+ print(
+ "The number of process groups to use SHARP with depends on the type "
+ "of the network switch. Nvidia QM1 switch supports SAHRP up to 8 "
+ "process groups and QM2 supports up to 256 process groups. We apply "
+ "SHARP to the communications of the data-parallel domain. If the "
+ "number of data-parallel process groups is larger than the max "
+ "process groups that the network switch supports, the communication "
+ "will fall back to non-SHARP operators. To enable SHARP, "
+ "`#SBATCH_NETWORK=sharp` should be set in the sbatch script."
+ )
+ torch.distributed.barrier(
+ group=get_data_parallel_group(with_context_parallel=True),
+ device_ids=[torch.cuda.current_device()],
+ )
+ # Set `NCCL_COLLNET_ENABLE=0` to restrict SHARP application to DP process groups
+ os.environ["NCCL_COLLNET_ENABLE"] = "0"
+
+ # Build the context-parallel groups.
+ global _CONTEXT_PARALLEL_GROUP
+ global _CONTEXT_PARALLEL_GLOBAL_RANKS
+ assert _CONTEXT_PARALLEL_GROUP is None, 'context parallel group is already initialized'
+ for ranks in generator_wrapper('cp'):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('cp', nccl_comm_cfgs)
+ )
+ if rank in ranks:
+ _CONTEXT_PARALLEL_GROUP = group
+ _CONTEXT_PARALLEL_GLOBAL_RANKS = ranks
# Build the model-parallel groups.
global _MODEL_PARALLEL_GROUP
assert _MODEL_PARALLEL_GROUP is None, 'model parallel group is already initialized'
- for i in range(data_parallel_size):
- ranks = [data_parallel_group_ranks[i]
- for data_parallel_group_ranks in all_data_parallel_group_ranks]
- group = torch.distributed.new_group(ranks)
+ for ranks in generator_wrapper('tp-pp'):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('mp', nccl_comm_cfgs)
+ )
if rank in ranks:
_MODEL_PARALLEL_GROUP = group
+ # Build the model-parallel groups with expert parallel
+ global _MODEL_AND_EXPERT_PARALLEL_GROUP
+ assert (
+ _MODEL_AND_EXPERT_PARALLEL_GROUP is None
+ ), 'model and expert parallel group is already initialized'
+ for ranks in generator_wrapper('tp-ep-pp', independent_ep=True):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('mp_exp', nccl_comm_cfgs)
+ )
+ if rank in ranks:
+ _MODEL_AND_EXPERT_PARALLEL_GROUP = group
+
# Build the tensor model-parallel groups.
global _TENSOR_MODEL_PARALLEL_GROUP
- assert _TENSOR_MODEL_PARALLEL_GROUP is None, \
- 'tensor model parallel group is already initialized'
- for i in range(num_tensor_model_parallel_groups):
- ranks = range(i * tensor_model_parallel_size,
- (i + 1) * tensor_model_parallel_size)
- group = torch.distributed.new_group(ranks)
+ global _TENSOR_MODEL_PARALLEL_GLOBAL_RANKS
+ assert (
+ _TENSOR_MODEL_PARALLEL_GROUP is None
+ ), 'tensor model parallel group is already initialized'
+ for ranks in generator_wrapper('tp'):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('tp', nccl_comm_cfgs)
+ )
if rank in ranks:
_TENSOR_MODEL_PARALLEL_GROUP = group
+ _TENSOR_MODEL_PARALLEL_GLOBAL_RANKS = ranks
# Build the pipeline model-parallel groups and embedding groups
# (first and last rank in each pipeline model-parallel group).
global _PIPELINE_MODEL_PARALLEL_GROUP
global _PIPELINE_GLOBAL_RANKS
- assert _PIPELINE_MODEL_PARALLEL_GROUP is None, \
- 'pipeline model parallel group is already initialized'
+ assert (
+ _PIPELINE_MODEL_PARALLEL_GROUP is None
+ ), 'pipeline model parallel group is already initialized'
global _EMBEDDING_GROUP
global _EMBEDDING_GLOBAL_RANKS
assert _EMBEDDING_GROUP is None, 'embedding group is already initialized'
global _POSITION_EMBEDDING_GROUP
global _POSITION_EMBEDDING_GLOBAL_RANKS
- assert _POSITION_EMBEDDING_GROUP is None, \
- 'position embedding group is already initialized'
- for i in range(num_pipeline_model_parallel_groups):
- ranks = range(i, world_size, num_pipeline_model_parallel_groups)
- group = torch.distributed.new_group(ranks)
+ assert _POSITION_EMBEDDING_GROUP is None, 'position embedding group is already initialized'
+ for ranks in generator_wrapper('pp'):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('pp', nccl_comm_cfgs)
+ )
if rank in ranks:
- _PIPELINE_MODEL_PARALLEL_GROUP = group
- _PIPELINE_GLOBAL_RANKS = ranks
- # Setup embedding group (to exchange gradients between
- # first and last stages).
- if len(ranks) > 1:
- embedding_ranks = [ranks[0], ranks[-1]]
- position_embedding_ranks = [ranks[0]]
- if pipeline_model_parallel_split_rank is not None:
- if ranks[pipeline_model_parallel_split_rank] not in embedding_ranks:
- embedding_ranks = [ranks[0],
- ranks[pipeline_model_parallel_split_rank],
- ranks[-1]]
- if ranks[pipeline_model_parallel_split_rank] not in position_embedding_ranks:
- position_embedding_ranks = [ranks[0],
- ranks[pipeline_model_parallel_split_rank]]
- else:
- embedding_ranks = ranks
- position_embedding_ranks = ranks
-
- group = torch.distributed.new_group(embedding_ranks)
+ if _PIPELINE_MODEL_PARALLEL_GROUP is None:
+ _PIPELINE_MODEL_PARALLEL_GROUP = group
+ _PIPELINE_GLOBAL_RANKS = ranks
+ elif isinstance(_PIPELINE_GLOBAL_RANKS[0], list):
+ _PIPELINE_MODEL_PARALLEL_GROUP.append(group)
+ _PIPELINE_GLOBAL_RANKS.append(ranks)
+ else:
+ _PIPELINE_MODEL_PARALLEL_GROUP = [_PIPELINE_MODEL_PARALLEL_GROUP, group]
+ _PIPELINE_GLOBAL_RANKS = [_PIPELINE_GLOBAL_RANKS, ranks]
+
+ embedding_ranks = get_embedding_ranks(ranks)
+ group = torch.distributed.new_group(
+ embedding_ranks, timeout=timeout, pg_options=get_nccl_options('embd', nccl_comm_cfgs)
+ )
if rank in embedding_ranks:
_EMBEDDING_GROUP = group
- if rank in ranks:
_EMBEDDING_GLOBAL_RANKS = embedding_ranks
- group = torch.distributed.new_group(position_embedding_ranks)
+ position_embedding_ranks = get_position_embedding_ranks(ranks)
+ group = torch.distributed.new_group(
+ position_embedding_ranks,
+ timeout=timeout,
+ pg_options=get_nccl_options('embd', nccl_comm_cfgs),
+ )
if rank in position_embedding_ranks:
_POSITION_EMBEDDING_GROUP = group
- if rank in ranks:
_POSITION_EMBEDDING_GLOBAL_RANKS = position_embedding_ranks
- # Build the FP8 groups.
- global _AMAX_REDUCTION_GROUP
- assert _AMAX_REDUCTION_GROUP is None, \
- 'FP8 amax reduction group is already initialized'
- if use_fp8:
- amax_group_size: int = tensor_model_parallel_size * data_parallel_size
- num_amax_groups: int = world_size // amax_group_size
- for i in range(num_amax_groups):
- start_rank = i * amax_group_size
- end_rank = (i + 1) * amax_group_size
- ranks = range(start_rank, end_rank)
- group = torch.distributed.new_group(ranks)
- if rank in ranks:
- _AMAX_REDUCTION_GROUP = group
+ # Build the tensor + data parallel groups.
+ global _TENSOR_AND_DATA_PARALLEL_GROUP
+ global _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP
+ assert (
+ _TENSOR_AND_DATA_PARALLEL_GROUP is None
+ ), 'Tensor + data parallel group is already initialized'
+ for ranks in generator_wrapper('tp-dp-cp'):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('tp_dp_cp', nccl_comm_cfgs)
+ )
+ if rank in ranks:
+ _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP = group
+ for ranks in generator_wrapper('tp-dp'):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('tp_dp', nccl_comm_cfgs)
+ )
+ if rank in ranks:
+ _TENSOR_AND_DATA_PARALLEL_GROUP = group
+
+ global _TENSOR_AND_CONTEXT_PARALLEL_GROUP
+ assert (
+ _TENSOR_AND_CONTEXT_PARALLEL_GROUP is None
+ ), 'Tensor + context parallel group is already initialized'
+ for ranks in generator_wrapper('tp-cp'):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('tp_cp', nccl_comm_cfgs)
+ )
+ if rank in ranks:
+ _TENSOR_AND_CONTEXT_PARALLEL_GROUP = group
+
+ # Build the tensor + expert parallel groups
+ global _EXPERT_MODEL_PARALLEL_GROUP
+ assert _EXPERT_MODEL_PARALLEL_GROUP is None, 'Expert parallel group is already initialized'
+ global _TENSOR_AND_EXPERT_PARALLEL_GROUP
+ assert (
+ _TENSOR_AND_EXPERT_PARALLEL_GROUP is None
+ ), 'Tensor + expert parallel group is already initialized'
+ global _DATA_MODULO_EXPERT_PARALLEL_GROUP
+ assert (
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP is None
+ ), 'Data modulo expert group is already initialized'
+ global _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP
+ assert (
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP is None
+ ), 'Data modulo expert group with context parallel is already initialized'
+ global _DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO
+ global _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP_GLOO
+
+ for ranks in generator_wrapper('tp-ep', independent_ep=True):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('tp_exp', nccl_comm_cfgs)
+ )
+ if rank in ranks:
+ _TENSOR_AND_EXPERT_PARALLEL_GROUP = group
+
+ for ranks in generator_wrapper('ep', independent_ep=True):
+ group = torch.distributed.new_group(
+ ranks, pg_options=get_nccl_options('exp', nccl_comm_cfgs)
+ )
+ if rank in ranks:
+ _EXPERT_MODEL_PARALLEL_GROUP = group
+
+ for ranks in generator_wrapper('dp', independent_ep=True):
+ group = torch.distributed.new_group(
+ ranks, timeout=timeout, pg_options=get_nccl_options('dp_modulo_exp', nccl_comm_cfgs)
+ )
+ group_gloo = torch.distributed.new_group(ranks, backend="gloo")
+ if rank in ranks:
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP = group
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO = group_gloo
+
+ for ranks in generator_wrapper('dp-cp', independent_ep=True):
+ # Lazy initialization of the group
+ if get_context_parallel_world_size() > 1:
+ group = torch.distributed.new_group(
+ ranks,
+ timeout=timeout,
+ pg_options=get_nccl_options('dp_modulo_exp_cp', nccl_comm_cfgs),
+ )
+ group_gloo = torch.distributed.new_group(ranks, backend="gloo")
+ else:
+ group = _DATA_MODULO_EXPERT_PARALLEL_GROUP
+ group_gloo = _DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO
+ if rank in ranks:
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP = group
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP_GLOO = group_gloo
# Initialize global memory buffer
# This isn't really "parallel state" but there isn't another good place to
@@ -260,94 +868,230 @@ def initialize_model_parallel(
_set_global_memory_buffer()
-def is_unitialized():
+def is_initialized():
"""Useful for code segments that may be accessed with or without mpu initialization"""
- return _DATA_PARALLEL_GROUP is None
+ return _DATA_PARALLEL_GROUP is not None
+
+
+def is_unitialized() -> bool:
+ """Check if parallel state has been initialized
+
+ Deprecated. Use is_initialized instead.
+
+ """
+ warnings.warn("is_unitialized is deprecated, use is_initialized instead", DeprecationWarning)
+ return not is_initialized()
def model_parallel_is_initialized():
- """Check if model and data parallel groups are initialized."""
- if _TENSOR_MODEL_PARALLEL_GROUP is None or \
- _PIPELINE_MODEL_PARALLEL_GROUP is None or \
- _DATA_PARALLEL_GROUP is None:
+ """Check if model- and data-parallel groups are initialized."""
+ if (
+ _TENSOR_MODEL_PARALLEL_GROUP is None
+ or _PIPELINE_MODEL_PARALLEL_GROUP is None
+ or _DATA_PARALLEL_GROUP is None
+ ):
return False
return True
-def get_model_parallel_group():
- """Get the model parallel group the caller rank belongs to."""
- assert _MODEL_PARALLEL_GROUP is not None, \
- 'model parallel group is not initialized'
+def get_model_parallel_group(with_expert_parallel=False):
+ """Get the model-parallel group the caller rank belongs to."""
+ if with_expert_parallel:
+ assert (
+ _MODEL_AND_EXPERT_PARALLEL_GROUP is not None
+ ), 'model parallel group is not initialized'
+ return _MODEL_AND_EXPERT_PARALLEL_GROUP
+ assert _MODEL_PARALLEL_GROUP is not None, 'model parallel group is not initialized'
return _MODEL_PARALLEL_GROUP
-def get_tensor_model_parallel_group():
- """Get the tensor model parallel group the caller rank belongs to."""
- assert _TENSOR_MODEL_PARALLEL_GROUP is not None, \
- 'intra_layer_model parallel group is not initialized'
+def get_tensor_model_parallel_group(check_initialized=True):
+ """Get the tensor-model-parallel group the caller rank belongs to."""
+ if check_initialized:
+ assert (
+ _TENSOR_MODEL_PARALLEL_GROUP is not None
+ ), 'tensor model parallel group is not initialized'
return _TENSOR_MODEL_PARALLEL_GROUP
def get_pipeline_model_parallel_group():
- """Get the pipeline model parallel group the caller rank belongs to."""
- assert _PIPELINE_MODEL_PARALLEL_GROUP is not None, \
- 'pipeline_model parallel group is not initialized'
+ """Get the pipeline-model-parallel group the caller rank belongs to."""
+ assert (
+ _PIPELINE_MODEL_PARALLEL_GROUP is not None
+ ), 'pipeline_model parallel group is not initialized'
return _PIPELINE_MODEL_PARALLEL_GROUP
-def get_data_parallel_group():
- """Get the data parallel group the caller rank belongs to."""
- assert _DATA_PARALLEL_GROUP is not None, \
- 'data parallel group is not initialized'
- return _DATA_PARALLEL_GROUP
+def get_data_parallel_group(with_context_parallel=False):
+ """Get the data-parallel group the caller rank belongs to."""
+ if with_context_parallel:
+ assert (
+ _DATA_PARALLEL_GROUP_WITH_CP is not None
+ ), 'data parallel group with context parallel combined is not initialized'
+ return _DATA_PARALLEL_GROUP_WITH_CP
+ else:
+ assert _DATA_PARALLEL_GROUP is not None, 'data parallel group is not initialized'
+ return _DATA_PARALLEL_GROUP
+
+
+def get_data_parallel_group_gloo(with_context_parallel=False):
+ """Get the Gloo data-parallel group the caller rank belongs to."""
+ if with_context_parallel:
+ assert (
+ _DATA_PARALLEL_GROUP_WITH_CP_GLOO is not None
+ ), 'data parallel group-gloo with context parallel combined is not initialized'
+ return _DATA_PARALLEL_GROUP_WITH_CP_GLOO
+ else:
+ assert _DATA_PARALLEL_GROUP_GLOO is not None, 'data parallel group-gloo is not initialized'
+ return _DATA_PARALLEL_GROUP_GLOO
-def get_data_parallel_group_gloo():
- """Get the data parallel group-gloo the caller rank belongs to."""
- assert _DATA_PARALLEL_GROUP_GLOO is not None, \
- 'data parallel group-gloo is not initialized'
- return _DATA_PARALLEL_GROUP_GLOO
+def get_context_parallel_group(check_initialized=True):
+ """Get the context-parallel group the caller rank belongs to."""
+ if check_initialized:
+ assert _CONTEXT_PARALLEL_GROUP is not None, 'context parallel group is not initialized'
+ return _CONTEXT_PARALLEL_GROUP
+
+
+def get_context_parallel_global_ranks(check_initialized=True):
+ """Get all global ranks of the context-parallel group that the caller rank belongs to."""
+ if check_initialized:
+ assert (
+ _CONTEXT_PARALLEL_GLOBAL_RANKS is not None
+ ), 'context parallel group is not initialized'
+ return _CONTEXT_PARALLEL_GLOBAL_RANKS
def get_embedding_group():
"""Get the embedding group the caller rank belongs to."""
- assert _EMBEDDING_GROUP is not None, \
- 'embedding group is not initialized'
+ assert _EMBEDDING_GROUP is not None, 'embedding group is not initialized'
return _EMBEDDING_GROUP
def get_position_embedding_group():
"""Get the position embedding group the caller rank belongs to."""
- assert _POSITION_EMBEDDING_GROUP is not None, \
- 'position embedding group is not initialized'
+ assert _POSITION_EMBEDDING_GROUP is not None, 'position embedding group is not initialized'
return _POSITION_EMBEDDING_GROUP
-def get_amax_reduction_group():
+def get_amax_reduction_group(with_context_parallel=False, tp_only_amax_red=False):
"""Get the FP8 amax reduction group the caller rank belongs to."""
- assert _AMAX_REDUCTION_GROUP is not None, \
- 'FP8 amax reduction group is not initialized'
- return _AMAX_REDUCTION_GROUP
+ if with_context_parallel:
+ if not tp_only_amax_red:
+ assert (
+ _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP is not None
+ ), 'FP8 amax reduction group is not initialized'
+ return _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP
+ else:
+ assert (
+ _TENSOR_AND_CONTEXT_PARALLEL_GROUP is not None
+ ), 'FP8 amax reduction group is not initialized'
+ return _TENSOR_AND_CONTEXT_PARALLEL_GROUP
+ else:
+ if not tp_only_amax_red:
+ assert (
+ _TENSOR_AND_DATA_PARALLEL_GROUP is not None
+ ), 'FP8 amax reduction group is not initialized'
+ return _TENSOR_AND_DATA_PARALLEL_GROUP
+ else:
+ assert (
+ _TENSOR_MODEL_PARALLEL_GROUP is not None
+ ), 'FP8 amax reduction group is not initialized'
+ return _TENSOR_MODEL_PARALLEL_GROUP
+
+
+def get_tensor_and_data_parallel_group(with_context_parallel=False):
+ """Get the tensor- and data-parallel group the caller rank belongs to."""
+ if with_context_parallel:
+ assert (
+ _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP is not None
+ ), 'tensor and data parallel group is not initialized'
+ return _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP
+ else:
+ assert (
+ _TENSOR_AND_DATA_PARALLEL_GROUP is not None
+ ), 'tensor and data parallel group is not initialized'
+ return _TENSOR_AND_DATA_PARALLEL_GROUP
+
+
+def get_tensor_and_context_parallel_group():
+ """Get the tensor- and context-parallel group the caller rank belongs to."""
+ assert (
+ _TENSOR_AND_CONTEXT_PARALLEL_GROUP is not None
+ ), 'tensor and context parallel group is not initialized'
+ return _TENSOR_AND_CONTEXT_PARALLEL_GROUP
+
+
+def get_expert_model_parallel_group():
+ """Get the expert-model-parallel group the caller rank belongs to."""
+ assert (
+ _EXPERT_MODEL_PARALLEL_GROUP is not None
+ ), 'expert model parallel group is not initialized'
+ return _EXPERT_MODEL_PARALLEL_GROUP
+
+
+def get_tensor_and_expert_parallel_group():
+ """Get the tensor- and expert-parallel group the caller rank belongs to."""
+ assert (
+ _TENSOR_AND_EXPERT_PARALLEL_GROUP is not None
+ ), 'tensor and expert parallel group is not initialized'
+ return _TENSOR_AND_EXPERT_PARALLEL_GROUP
+
+
+def get_data_modulo_expert_parallel_group(with_context_parallel=False):
+ """Get the data-modulo-expert-parallel group the caller rank belongs to."""
+ if with_context_parallel:
+ assert (
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP is not None
+ ), 'data modulo expert parallel group with context parallel is not initialized'
+ return _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP
+ else:
+ assert (
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP is not None
+ ), 'data modulo expert parallel group is not initialized'
+ return _DATA_MODULO_EXPERT_PARALLEL_GROUP
+
+
+def get_data_modulo_expert_parallel_group_gloo(with_context_parallel=False):
+ """Get the Gloo data-modulo-expert-parallel group the caller rank belongs to."""
+ if with_context_parallel:
+ assert (
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP_GLOO is not None
+ ), 'data modulo expert parallel group-gloo with context parallel is not initialized'
+ return _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP_GLOO
+ else:
+ assert (
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO is not None
+ ), 'data modulo expert parallel group-gloo is not initialized'
+ return _DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO
+
+
+def set_expert_model_parallel_world_size(world_size):
+ """Sets the expert-model-parallel world size."""
+ global _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE
+ _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE = world_size
def set_tensor_model_parallel_world_size(world_size):
- """Set the tensor model parallel size"""
+ """Set the tensor-model-parallel size"""
global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size
def set_pipeline_model_parallel_world_size(world_size):
- """Set the pipeline model parallel size"""
+ """Set the pipeline-model-parallel size"""
global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
_MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = world_size
+
def set_virtual_pipeline_model_parallel_world_size(world_size):
- """Set the pipeline model parallel size"""
+ """Set the pipeline-model-parallel size"""
global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
_VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = world_size
+
def get_tensor_model_parallel_world_size():
- """Return world size for the tensor model parallel group."""
+ """Return world size for the tensor-model-parallel group."""
global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None:
return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
@@ -355,33 +1099,49 @@ def get_tensor_model_parallel_world_size():
def get_pipeline_model_parallel_world_size():
- """Return world size for the pipeline model parallel group."""
+ """Return world size for the pipeline-model-parallel group."""
global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
if _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE is not None:
return _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
- return torch.distributed.get_world_size(group=get_pipeline_model_parallel_group())
+
+ pp_group = get_pipeline_model_parallel_group()
+ if isinstance(pp_group, list):
+ # Implicit assumption that each PP group is the same size.
+ sizes = []
+ for group in _PIPELINE_GLOBAL_RANKS:
+ sizes.append(len(group))
+ assert all(x == sizes[0] for x in sizes)
+ return torch.distributed.get_world_size(group=pp_group[0])
+ else:
+ return torch.distributed.get_world_size(group=pp_group)
+
+
+def set_expert_model_parallel_rank(rank):
+ """Set expert-model-parallel rank."""
+ global _MPU_EXPERT_MODEL_PARALLEL_RANK
+ _MPU_EXPERT_MODEL_PARALLEL_RANK = rank
def set_tensor_model_parallel_rank(rank):
- """Set tensor model parallel rank."""
+ """Set tensor-model-parallel rank."""
global _MPU_TENSOR_MODEL_PARALLEL_RANK
_MPU_TENSOR_MODEL_PARALLEL_RANK = rank
def set_pipeline_model_parallel_rank(rank):
- """Set pipeline model parallel rank."""
+ """Set pipeline-model-parallel rank."""
global _MPU_PIPELINE_MODEL_PARALLEL_RANK
_MPU_PIPELINE_MODEL_PARALLEL_RANK = rank
def set_pipeline_model_parallel_split_rank(rank):
- """Set pipeline model parallel split rank."""
+ """Set pipeline-model-parallel split rank. DEPRECATED."""
global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
_PIPELINE_MODEL_PARALLEL_SPLIT_RANK = rank
def get_tensor_model_parallel_rank():
- """Return my rank for the tensor model parallel group."""
+ """Return caller's rank for the tensor-model-parallel group."""
global _MPU_TENSOR_MODEL_PARALLEL_RANK
if _MPU_TENSOR_MODEL_PARALLEL_RANK is not None:
return _MPU_TENSOR_MODEL_PARALLEL_RANK
@@ -389,15 +1149,27 @@ def get_tensor_model_parallel_rank():
def get_pipeline_model_parallel_rank():
- """Return my rank for the pipeline model parallel group."""
+ """Return caller's rank for the pipeline-model-parallel group."""
global _MPU_PIPELINE_MODEL_PARALLEL_RANK
if _MPU_PIPELINE_MODEL_PARALLEL_RANK is not None:
return _MPU_PIPELINE_MODEL_PARALLEL_RANK
- return torch.distributed.get_rank(group=get_pipeline_model_parallel_group())
+ rank = torch.distributed.get_rank()
+ pp_group = get_pipeline_model_parallel_group()
+ if isinstance(pp_group, list):
+ # Assume that if the caller exist in multiple PP groups, then it has the same index.
+ indices = []
+ for group in _PIPELINE_GLOBAL_RANKS:
+ for i, r in enumerate(group):
+ if r == rank:
+ indices.append(i)
+ assert all(x == indices[0] for x in indices)
+ return torch.distributed.get_rank(group=pp_group[0])
+ else:
+ return torch.distributed.get_rank(group=pp_group)
def get_pipeline_model_parallel_split_rank():
- """Return pipeline model parallel split rank."""
+ """Return pipeline-model-parallel split rank."""
global _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
return _PIPELINE_MODEL_PARALLEL_SPLIT_RANK
@@ -405,29 +1177,35 @@ def get_pipeline_model_parallel_split_rank():
def is_pipeline_first_stage(ignore_virtual=False):
"""Return True if in the first pipeline model-parallel stage, False otherwise."""
if not ignore_virtual:
- if get_virtual_pipeline_model_parallel_world_size() is not None and \
- get_virtual_pipeline_model_parallel_rank() != 0:
+ if (
+ get_virtual_pipeline_model_parallel_world_size() is not None
+ and get_virtual_pipeline_model_parallel_rank() != 0
+ ):
return False
return get_pipeline_model_parallel_rank() == 0
def is_pipeline_last_stage(ignore_virtual=False):
- """Return True if in the last pipeline model-parallel stage, False otherwise."""
+ """Return True if in the last pipeline-model-parallel stage, False otherwise."""
if not ignore_virtual:
- virtual_pipeline_model_parallel_world_size = \
+ virtual_pipeline_model_parallel_world_size = (
get_virtual_pipeline_model_parallel_world_size()
- if virtual_pipeline_model_parallel_world_size is not None and \
- get_virtual_pipeline_model_parallel_rank() != (
- virtual_pipeline_model_parallel_world_size - 1):
+ )
+ if (
+ virtual_pipeline_model_parallel_world_size is not None
+ and get_virtual_pipeline_model_parallel_rank()
+ != (virtual_pipeline_model_parallel_world_size - 1)
+ ):
return False
- return get_pipeline_model_parallel_rank() == (
- get_pipeline_model_parallel_world_size() - 1)
+ return get_pipeline_model_parallel_rank() == (get_pipeline_model_parallel_world_size() - 1)
def is_rank_in_embedding_group(ignore_virtual=False):
"""Return true if current rank is in embedding group, False otherwise."""
rank = torch.distributed.get_rank()
global _EMBEDDING_GLOBAL_RANKS
+ if _EMBEDDING_GLOBAL_RANKS is None:
+ return False
if ignore_virtual:
return rank in _EMBEDDING_GLOBAL_RANKS
if rank in _EMBEDDING_GLOBAL_RANKS:
@@ -444,7 +1222,7 @@ def is_rank_in_position_embedding_group():
"""Return true if current rank is in position embedding group, False otherwise."""
rank = torch.distributed.get_rank()
global _POSITION_EMBEDDING_GLOBAL_RANKS
- return rank in _POSITION_EMBEDDING_GLOBAL_RANKS
+ return _POSITION_EMBEDDING_GLOBAL_RANKS is not None and rank in _POSITION_EMBEDDING_GLOBAL_RANKS
def is_pipeline_stage_before_split(rank=None):
@@ -477,13 +1255,42 @@ def is_pipeline_stage_after_split(rank=None):
return False
+def is_inside_encoder(rank=None):
+ """Return True if pipeline stage executes encoder block for a model
+ with both encoder and decoder."""
+ if get_pipeline_model_parallel_world_size() == 1:
+ return True
+ if rank is None:
+ rank = get_pipeline_model_parallel_rank()
+ global _PIPELINE_MODEL_PARALLEL_DECODER_START
+ if _PIPELINE_MODEL_PARALLEL_DECODER_START is None:
+ return True
+ if rank < _PIPELINE_MODEL_PARALLEL_DECODER_START:
+ return True
+ return False
+
+
+def is_inside_decoder(rank=None):
+ """Return True if pipeline stage executes decoder block for a model
+ with both encoder and decoder."""
+ if get_pipeline_model_parallel_world_size() == 1:
+ return True
+ if rank is None:
+ rank = get_pipeline_model_parallel_rank()
+ global _PIPELINE_MODEL_PARALLEL_DECODER_START
+ if _PIPELINE_MODEL_PARALLEL_DECODER_START is None:
+ return True
+ if rank >= _PIPELINE_MODEL_PARALLEL_DECODER_START:
+ return True
+ return False
+
+
def is_pipeline_stage_at_split():
"""Return true if pipeline stage executes decoder block and next
stage executes encoder block for a model with both encoder and
decoder."""
rank = get_pipeline_model_parallel_rank()
- return is_pipeline_stage_before_split(rank) and \
- is_pipeline_stage_after_split(rank+1)
+ return is_pipeline_stage_before_split(rank) and is_pipeline_stage_after_split(rank + 1)
def get_virtual_pipeline_model_parallel_rank():
@@ -504,110 +1311,337 @@ def get_virtual_pipeline_model_parallel_world_size():
return _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
-def set_virtual_pipeline_model_parallel_world_size(world_size):
- """Set the virtual pipeline-parallel world size"""
- global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
- _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = world_size
-
-
def get_tensor_model_parallel_src_rank():
"""Calculate the global rank corresponding to the first local rank
in the tensor model parallel group."""
- global_rank = torch.distributed.get_rank()
- local_world_size = get_tensor_model_parallel_world_size()
- return (global_rank // local_world_size) * local_world_size
+ assert (
+ _TENSOR_MODEL_PARALLEL_GLOBAL_RANKS is not None
+ ), "Tensor model parallel group is not initialized"
+ return _TENSOR_MODEL_PARALLEL_GLOBAL_RANKS[0]
-def get_data_parallel_src_rank():
+def get_data_parallel_src_rank(with_context_parallel=False):
"""Calculate the global rank corresponding to the first local rank
in the data parallel group."""
- assert _DATA_PARALLEL_GLOBAL_RANKS is not None, \
- "Data parallel group is not initialized"
- return _DATA_PARALLEL_GLOBAL_RANKS[0]
+ if with_context_parallel:
+ assert (
+ _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP is not None
+ ), "Data parallel group with context parallel combined is not initialized"
+ return _DATA_PARALLEL_GLOBAL_RANKS_WITH_CP[0]
+ else:
+ assert _DATA_PARALLEL_GLOBAL_RANKS is not None, "Data parallel group is not initialized"
+ return _DATA_PARALLEL_GLOBAL_RANKS[0]
def get_pipeline_model_parallel_first_rank():
- """Return the global rank of the first process in the pipeline for the
- current tensor parallel group"""
- assert _PIPELINE_GLOBAL_RANKS is not None, \
- "Pipeline parallel group is not initialized"
- return _PIPELINE_GLOBAL_RANKS[0]
+ """Return the global rank of the first stage in the current rank's pipeline."""
+ assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
+ if isinstance(_PIPELINE_GLOBAL_RANKS[0], list):
+ # I assume the first rank is the same for all pp groups right now.
+ for rank_group in _PIPELINE_GLOBAL_RANKS:
+ assert rank_group[0] == _PIPELINE_GLOBAL_RANKS[0][0]
+ return _PIPELINE_GLOBAL_RANKS[0][0]
+ else:
+ return _PIPELINE_GLOBAL_RANKS[0]
def get_pipeline_model_parallel_last_rank():
- """Return the global rank of the last process in the pipeline for the
- current tensor parallel group"""
- assert _PIPELINE_GLOBAL_RANKS is not None, \
- "Pipeline parallel group is not initialized"
+ """Return the global rank of the last stage in the current rank's pipeline."""
+ assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
last_rank_local = get_pipeline_model_parallel_world_size() - 1
return _PIPELINE_GLOBAL_RANKS[last_rank_local]
+
def get_pipeline_model_parallel_next_rank():
- """Return the global rank that follows the caller in the pipeline"""
- assert _PIPELINE_GLOBAL_RANKS is not None, \
- "Pipeline parallel group is not initialized"
+ """Return the global rank that follows the caller in the pipeline, for each
+ pipeline-parallel group that the rank is part of.
+
+ If it is just part of one group, an int is returned, otherwise a list of ints.
+ """
+ assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
rank_in_pipeline = get_pipeline_model_parallel_rank()
world_size = get_pipeline_model_parallel_world_size()
- return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline + 1) % world_size]
+ if isinstance(_PIPELINE_GLOBAL_RANKS[0], list):
+ to_return = []
+ for group in _PIPELINE_GLOBAL_RANKS:
+ to_return.append(group[(rank_in_pipeline + 1) % world_size])
+ return to_return
+ else:
+ return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline + 1) % world_size]
def get_pipeline_model_parallel_prev_rank():
- """Return the global rank that preceeds the caller in the pipeline"""
- assert _PIPELINE_GLOBAL_RANKS is not None, \
- "Pipeline parallel group is not initialized"
+ """Return the global rank that precedes the caller in the pipeline, for each
+ pipeline-parallel group that the rank is part of.
+
+ If it is just part of one group, an int is returned, otherwise a list of ints.
+ """
+ assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
rank_in_pipeline = get_pipeline_model_parallel_rank()
world_size = get_pipeline_model_parallel_world_size()
- return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline - 1) % world_size]
+ if isinstance(_PIPELINE_GLOBAL_RANKS[0], list):
+ to_return = []
+ for group in _PIPELINE_GLOBAL_RANKS:
+ to_return.append(group[(rank_in_pipeline - 1) % world_size])
+ return to_return
+ else:
+ return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline - 1) % world_size]
-def get_data_parallel_world_size():
+def get_data_parallel_world_size(with_context_parallel=False):
"""Return world size for the data parallel group."""
- return torch.distributed.get_world_size(group=get_data_parallel_group())
+ global _MPU_DATA_PARALLEL_WORLD_SIZE
+ if _MPU_DATA_PARALLEL_WORLD_SIZE is not None:
+ return _MPU_DATA_PARALLEL_WORLD_SIZE
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ return torch.distributed.get_world_size(
+ group=get_data_parallel_group(with_context_parallel=with_context_parallel)
+ )
+ else:
+ return 0
+
+
+def set_data_parallel_rank(rank):
+ """Return world size for the data parallel group."""
+ global _MPU_DATA_PARALLEL_RANK
+ _MPU_DATA_PARALLEL_RANK = rank
+
+
+def get_data_parallel_rank(with_context_parallel=False):
+ """Return caller's rank in the data-parallel group."""
+ global _MPU_DATA_PARALLEL_RANK
+ if _MPU_DATA_PARALLEL_RANK is not None:
+ return _MPU_DATA_PARALLEL_RANK
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ return torch.distributed.get_rank(
+ group=get_data_parallel_group(with_context_parallel=with_context_parallel)
+ )
+ else:
+ return 0
+
+
+def get_context_parallel_world_size():
+ """Return world size for the context parallel group."""
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ return torch.distributed.get_world_size(group=get_context_parallel_group())
+ else:
+ return 0
+
+
+def get_context_parallel_rank():
+ """Return caller's rank in the context-parallel group."""
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ return torch.distributed.get_rank(group=get_context_parallel_group())
+ else:
+ return 0
+
+
+def get_tensor_and_context_parallel_world_size():
+ """Return world size for the tensor and context-parallel group."""
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ return torch.distributed.get_world_size(group=get_tensor_and_context_parallel_group())
+ else:
+ return 0
+
+
+def get_tensor_and_context_parallel_rank():
+ """Return caller's rank in the joint tensor-model-parallel and context-parallel group."""
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ return torch.distributed.get_rank(group=get_tensor_and_context_parallel_group())
+ else:
+ return 0
+
+
+def get_expert_model_parallel_world_size():
+ """Return world size for the expert-model-parallel group."""
+ if _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE is not None:
+ return _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ tensor_and_expert_parallel_world_size = torch.distributed.get_world_size(
+ group=get_tensor_and_expert_parallel_group()
+ )
+ return tensor_and_expert_parallel_world_size // get_tensor_model_parallel_world_size()
+ else:
+ return 0
-def get_data_parallel_rank():
- """Return my rank for the data parallel group."""
- return torch.distributed.get_rank(group=get_data_parallel_group())
+def get_tensor_and_expert_parallel_world_size():
+ """Return world size for the expert model parallel group times model parallel group.
+ Currently, each expert will also be distributed across TP group by default.
+ """
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ tensor_and_expert_parallel_world_size = torch.distributed.get_world_size(
+ group=get_tensor_and_expert_parallel_group()
+ )
+ return tensor_and_expert_parallel_world_size
+ else:
+ return 0
+
+
+def get_expert_model_parallel_rank():
+ """Return caller's rank in the expert-model-parallel group."""
+ if _MPU_EXPERT_MODEL_PARALLEL_RANK is not None:
+ return _MPU_EXPERT_MODEL_PARALLEL_RANK
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ tensor_and_expert_parallel_rank = torch.distributed.get_rank(
+ group=get_tensor_and_expert_parallel_group()
+ )
+ return tensor_and_expert_parallel_rank // get_tensor_model_parallel_world_size()
+ else:
+ return 0
+
+
+def get_data_modulo_expert_parallel_rank(with_context_parallel=False):
+ """Return caller's rank in the context-parallel group."""
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ return torch.distributed.get_rank(
+ group=get_data_modulo_expert_parallel_group(with_context_parallel=with_context_parallel)
+ )
+ else:
+ return 0
+
+
+def get_tensor_and_expert_parallel_rank():
+ """Return caller's rank in the joint tensor- and expert-model-parallel group."""
+ if torch.distributed.is_available() and torch.distributed.is_initialized():
+ return torch.distributed.get_rank(group=get_tensor_and_expert_parallel_group())
+ else:
+ return 0
+
def _set_global_memory_buffer():
- """Initialize global buffer"""
+ """Initialize global buffer."""
global _GLOBAL_MEMORY_BUFFER
assert _GLOBAL_MEMORY_BUFFER is None, 'global memory buffer is already initialized'
_GLOBAL_MEMORY_BUFFER = GlobalMemoryBuffer()
+
def get_global_memory_buffer():
"""Return the global GlobalMemoryBuffer object"""
assert _GLOBAL_MEMORY_BUFFER is not None, 'global memory buffer is not initialized'
return _GLOBAL_MEMORY_BUFFER
+def destroy_global_memory_buffer():
+ """Sets the global memory buffer to None"""
+ global _GLOBAL_MEMORY_BUFFER
+ _GLOBAL_MEMORY_BUFFER = None
+
+
+def get_all_ranks():
+ """Get caller's rank in tensor-model-parallel, data-parallel, context-parallel,
+ pipeline-model-parallel and expert-model-parallel groups."""
+ ranks = [
+ get_tensor_model_parallel_rank(),
+ get_data_parallel_rank(),
+ get_context_parallel_rank(),
+ get_pipeline_model_parallel_rank(),
+ get_expert_model_parallel_rank(),
+ ]
+ return '_'.join(map(lambda x: str(x or 0), ranks))
+
+
+def get_moe_layer_wise_logging_tracker():
+ """Return the moe layer wise tracker."""
+ global _MOE_LAYER_WISE_LOGGING_TRACKER
+ return _MOE_LAYER_WISE_LOGGING_TRACKER
+
+
def destroy_model_parallel():
"""Set the groups to none."""
global _MODEL_PARALLEL_GROUP
_MODEL_PARALLEL_GROUP = None
+
+ global _MODEL_AND_EXPERT_PARALLEL_GROUP
+ _MODEL_AND_EXPERT_PARALLEL_GROUP = None
+
global _TENSOR_MODEL_PARALLEL_GROUP
_TENSOR_MODEL_PARALLEL_GROUP = None
+
global _PIPELINE_MODEL_PARALLEL_GROUP
_PIPELINE_MODEL_PARALLEL_GROUP = None
+
global _DATA_PARALLEL_GROUP
_DATA_PARALLEL_GROUP = None
+
+ global _DATA_PARALLEL_GROUP_WITH_CP
+ _DATA_PARALLEL_GROUP_WITH_CP = None
+
+ global _CONTEXT_PARALLEL_GROUP
+ _CONTEXT_PARALLEL_GROUP = None
+
+ global _CONTEXT_PARALLEL_GLOBAL_RANKS
+ _CONTEXT_PARALLEL_GLOBAL_RANKS = None
+
global _EMBEDDING_GROUP
_EMBEDDING_GROUP = None
+
global _POSITION_EMBEDDING_GROUP
_POSITION_EMBEDDING_GROUP = None
- global _AMAX_REDUCTION_GROUP
- _AMAX_REDUCTION_GROUP = None
+
+ global _TENSOR_AND_DATA_PARALLEL_GROUP
+ _TENSOR_AND_DATA_PARALLEL_GROUP = None
+
+ global _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP
+ _TENSOR_AND_DATA_PARALLEL_GROUP_WITH_CP = None
+
+ global _TENSOR_AND_CONTEXT_PARALLEL_GROUP
+ _TENSOR_AND_CONTEXT_PARALLEL_GROUP = None
+
+ global _EXPERT_MODEL_PARALLEL_GROUP
+ _EXPERT_MODEL_PARALLEL_GROUP = None
+
+ global _TENSOR_AND_EXPERT_PARALLEL_GROUP
+ _TENSOR_AND_EXPERT_PARALLEL_GROUP = None
+
+ global _DATA_MODULO_EXPERT_PARALLEL_GROUP
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP = None
+
+ global _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP = None
+
global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
_VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = None
+
global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
_VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
+
global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None
+
global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
_MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
+
global _MPU_TENSOR_MODEL_PARALLEL_RANK
_MPU_TENSOR_MODEL_PARALLEL_RANK = None
+
global _MPU_PIPELINE_MODEL_PARALLEL_RANK
_MPU_PIPELINE_MODEL_PARALLEL_RANK = None
+
global _GLOBAL_MEMORY_BUFFER
_GLOBAL_MEMORY_BUFFER = None
+
+ global _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE
+ _MPU_EXPERT_MODEL_PARALLEL_WORLD_SIZE = None
+
+ global _MPU_EXPERT_MODEL_PARALLEL_RANK
+ _MPU_EXPERT_MODEL_PARALLEL_RANK = None
+
+ global _DATA_PARALLEL_GROUP_GLOO
+ if _DATA_PARALLEL_GROUP_GLOO is not None:
+ torch.distributed.destroy_process_group(_DATA_PARALLEL_GROUP_GLOO)
+ _DATA_PARALLEL_GROUP_GLOO = None
+
+ global _DATA_PARALLEL_GROUP_WITH_CP_GLOO
+ _DATA_PARALLEL_GROUP_WITH_CP_GLOO = None
+
+ global _DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO
+ if _DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO is not None:
+ torch.distributed.destroy_process_group(_DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO)
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_GLOO = None
+
+ global _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP_GLOO
+ _DATA_MODULO_EXPERT_PARALLEL_GROUP_WITH_CP_GLOO = None
+
+ global _MOE_LAYER_WISE_LOGGING_TRACKER
+ _MOE_LAYER_WISE_LOGGING_TRACKER = {}
diff --git a/megatron/core/pipeline_parallel/__init__.py b/megatron/core/pipeline_parallel/__init__.py
index 00cd1ff382..37b3a5a972 100644
--- a/megatron/core/pipeline_parallel/__init__.py
+++ b/megatron/core/pipeline_parallel/__init__.py
@@ -1 +1,2 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
from .schedules import get_forward_backward_func
diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py
index 301583132a..3e33e7c2f8 100644
--- a/megatron/core/pipeline_parallel/p2p_communication.py
+++ b/megatron/core/pipeline_parallel/p2p_communication.py
@@ -1,30 +1,32 @@
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
-from functools import reduce
import operator
-from typing import Optional, List, Union, Callable, Tuple
+from functools import reduce
+from typing import Callable, List, Optional, Tuple, Union
import torch
from megatron import core
+from megatron.core import ModelParallelConfig
from megatron.core.parallel_state import (
get_pipeline_model_parallel_group,
- get_pipeline_model_parallel_prev_rank,
get_pipeline_model_parallel_next_rank,
+ get_pipeline_model_parallel_prev_rank,
+ get_pipeline_model_parallel_rank,
+ get_pipeline_model_parallel_world_size,
)
# Types
Shape = Union[List[int], torch.Size]
-def _communicate_shapes(tensor_send_next, tensor_send_prev,
- recv_prev, recv_next,
- use_ring_exchange_p2p):
+
+def _communicate_shapes(tensor_send_next, tensor_send_prev, recv_prev, recv_next, config):
"""Communicate tensor shapes between stages. Used to communicate
tensor shapes before the actual tensor communication happens.
This is required when the sequence lengths across micro batches
are not uniform.
- Takes the following arguments:
+ Args:
tensor_send_next: tensor to send to next rank (no tensor sent if
set to None).
tensor_send_prev: tensor to send to prev rank (no tensor sent if
@@ -42,49 +44,59 @@ def _communicate_shapes(tensor_send_next, tensor_send_prev,
send_prev_shape_tensor = None
send_next_shape_tensor = None
if recv_prev:
- recv_prev_shape_tensor = torch.empty((3),
- device=torch.cuda.current_device(),
- dtype=torch.int64)
+ recv_prev_shape_tensor = torch.empty(
+ (3), device=torch.cuda.current_device(), dtype=torch.int64
+ )
if recv_next:
- recv_next_shape_tensor = torch.empty((3),
- device=torch.cuda.current_device(),
- dtype=torch.int64)
+ recv_next_shape_tensor = torch.empty(
+ (3), device=torch.cuda.current_device(), dtype=torch.int64
+ )
if tensor_send_prev is not None:
- send_prev_shape_tensor = torch.tensor(tensor_send_prev.size(),
- device=torch.cuda.current_device(),
- dtype=torch.int64)
+ send_prev_shape_tensor = torch.tensor(
+ tensor_send_prev.size(), device=torch.cuda.current_device(), dtype=torch.int64
+ )
if tensor_send_next is not None:
- send_next_shape_tensor = torch.tensor(tensor_send_next.size(),
- device=torch.cuda.current_device(),
- dtype=torch.int64)
-
- if use_ring_exchange_p2p:
- torch.distributed.ring_exchange(tensor_send_prev=send_prev_shape_tensor,
- tensor_recv_prev=recv_prev_shape_tensor,
- tensor_send_next=send_next_shape_tensor,
- tensor_recv_next=recv_next_shape_tensor,
- group=mpu.get_pipeline_model_parallel_group())
+ send_next_shape_tensor = torch.tensor(
+ tensor_send_next.size(), device=torch.cuda.current_device(), dtype=torch.int64
+ )
+
+ if config.use_ring_exchange_p2p:
+ torch.distributed.ring_exchange(
+ tensor_send_prev=send_prev_shape_tensor,
+ tensor_recv_prev=recv_prev_shape_tensor,
+ tensor_send_next=send_next_shape_tensor,
+ tensor_recv_next=recv_next_shape_tensor,
+ group=get_pipeline_model_parallel_group(),
+ )
else:
ops = []
if send_prev_shape_tensor is not None:
send_prev_op = torch.distributed.P2POp(
- torch.distributed.isend, send_prev_shape_tensor,
- mpu.get_pipeline_model_parallel_prev_rank())
+ torch.distributed.isend,
+ send_prev_shape_tensor,
+ get_pipeline_model_parallel_prev_rank(),
+ )
ops.append(send_prev_op)
if recv_prev_shape_tensor is not None:
recv_prev_op = torch.distributed.P2POp(
- torch.distributed.irecv, recv_prev_shape_tensor,
- mpu.get_pipeline_model_parallel_prev_rank())
+ torch.distributed.irecv,
+ recv_prev_shape_tensor,
+ get_pipeline_model_parallel_prev_rank(),
+ )
ops.append(recv_prev_op)
if send_next_shape_tensor is not None:
send_next_op = torch.distributed.P2POp(
- torch.distributed.isend, send_next_shape_tensor,
- mpu.get_pipeline_model_parallel_next_rank())
+ torch.distributed.isend,
+ send_next_shape_tensor,
+ get_pipeline_model_parallel_next_rank(),
+ )
ops.append(send_next_op)
if recv_next_shape_tensor is not None:
recv_next_op = torch.distributed.P2POp(
- torch.distributed.irecv, recv_next_shape_tensor,
- mpu.get_pipeline_model_parallel_next_rank())
+ torch.distributed.irecv,
+ recv_next_shape_tensor,
+ get_pipeline_model_parallel_next_rank(),
+ )
ops.append(recv_next_op)
if len(ops) > 0:
reqs = torch.distributed.batch_isend_irecv(ops)
@@ -106,19 +118,132 @@ def _communicate_shapes(tensor_send_next, tensor_send_prev,
return recv_prev_shape, recv_next_shape
-def _communicate(*, tensor_send_next: Optional[torch.Tensor],
- tensor_send_prev: Optional[torch.Tensor],
- recv_prev: bool,
- recv_next: bool,
- tensor_shape: Shape,
- dtype: Optional[torch.dtype],
- variable_seq_lengths: bool = False,
- use_ring_exchange_p2p: bool = False,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+def _batched_p2p_ops(
+ *,
+ tensor_send_prev: Optional[torch.Tensor],
+ tensor_recv_prev: Optional[torch.Tensor],
+ tensor_send_next: Optional[torch.Tensor],
+ tensor_recv_next: Optional[torch.Tensor],
+ group: torch.distributed.ProcessGroup,
+ prev_pipeline_rank: int,
+ next_pipeline_rank: int,
+):
+ ops = []
+ if tensor_send_prev is not None:
+ send_prev_op = torch.distributed.P2POp(
+ torch.distributed.isend, tensor_send_prev, prev_pipeline_rank, group
+ )
+ ops.append(send_prev_op)
+ if tensor_recv_prev is not None:
+ recv_prev_op = torch.distributed.P2POp(
+ torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank, group
+ )
+ ops.append(recv_prev_op)
+ if tensor_send_next is not None:
+ send_next_op = torch.distributed.P2POp(
+ torch.distributed.isend, tensor_send_next, next_pipeline_rank, group
+ )
+ ops.append(send_next_op)
+ if tensor_recv_next is not None:
+ recv_next_op = torch.distributed.P2POp(
+ torch.distributed.irecv, tensor_recv_next, next_pipeline_rank, group
+ )
+ ops.append(recv_next_op)
+ if len(ops) > 0:
+ reqs = torch.distributed.batch_isend_irecv(ops)
+ else:
+ reqs = []
+ return reqs
+
+
+def _p2p_ops(
+ *,
+ tensor_send_prev: Optional[torch.Tensor],
+ tensor_recv_prev: Optional[torch.Tensor],
+ tensor_send_next: Optional[torch.Tensor],
+ tensor_recv_next: Optional[torch.Tensor],
+ group: torch.distributed.ProcessGroup,
+ prev_pipeline_rank: int,
+ next_pipeline_rank: int,
+):
+ reqs = []
+ rank = get_pipeline_model_parallel_rank()
+ even_send_odd_recv_group = group
+ if get_pipeline_model_parallel_world_size() == 2:
+ # Use the global process group for one of the two p2p communications
+ # to allow the overlap of the independent communications.
+ # Using the global process group is compatible because the pipeline-parallel
+ # communications set the source and destination by global rank.
+ even_recv_odd_send_group = torch.distributed.group.WORLD
+ else:
+ even_recv_odd_send_group = group
+
+ if get_pipeline_model_parallel_rank() % 2 == 0:
+ if tensor_send_next is not None:
+ send_next_req = torch.distributed.isend(
+ tensor=tensor_send_next, dst=next_pipeline_rank, group=even_send_odd_recv_group
+ )
+ reqs.append(send_next_req)
+
+ if tensor_recv_prev is not None:
+ recv_prev_req = torch.distributed.irecv(
+ tensor=tensor_recv_prev, src=prev_pipeline_rank, group=even_recv_odd_send_group
+ )
+ reqs.append(recv_prev_req)
+
+ if tensor_send_prev is not None:
+ send_prev_req = torch.distributed.isend(
+ tensor=tensor_send_prev, dst=prev_pipeline_rank, group=even_send_odd_recv_group
+ )
+ reqs.append(send_prev_req)
+
+ if tensor_recv_next is not None:
+ recv_next_req = torch.distributed.irecv(
+ tensor=tensor_recv_next, src=next_pipeline_rank, group=even_recv_odd_send_group
+ )
+ reqs.append(recv_next_req)
+
+ else:
+ if tensor_recv_prev is not None:
+ recv_prev_req = torch.distributed.irecv(
+ tensor=tensor_recv_prev, src=prev_pipeline_rank, group=even_send_odd_recv_group
+ )
+ reqs.append(recv_prev_req)
+
+ if tensor_send_next is not None:
+ send_next_req = torch.distributed.isend(
+ tensor=tensor_send_next, dst=next_pipeline_rank, group=even_recv_odd_send_group
+ )
+ reqs.append(send_next_req)
+
+ if tensor_recv_next is not None:
+ recv_next_req = torch.distributed.irecv(
+ tensor=tensor_recv_next, src=next_pipeline_rank, group=even_send_odd_recv_group
+ )
+ reqs.append(recv_next_req)
+
+ if tensor_send_prev is not None:
+ send_prev_req = torch.distributed.isend(
+ tensor=tensor_send_prev, dst=prev_pipeline_rank, group=even_recv_odd_send_group
+ )
+ reqs.append(send_prev_req)
+ return reqs
+
+
+def _communicate(
+ *,
+ tensor_send_next: Optional[torch.Tensor],
+ tensor_send_prev: Optional[torch.Tensor],
+ recv_prev: bool,
+ recv_next: bool,
+ tensor_shape: Shape,
+ config: ModelParallelConfig,
+ wait_on_reqs: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor]:
"""Communicate tensors between stages. Used as helper method in other
communication methods that are used in megatron/schedules.py.
- Arguments:
+ Args:
tensor_send_next (torch.Tensor, optional):
Tensor to send to next rank (no tensor sent if None)
@@ -136,23 +261,9 @@ def _communicate(*, tensor_send_next: Optional[torch.Tensor],
tensors sent and received in a single function call are
the same shape).
- dtype (torch.dtype, required if either recv_{prev,next} is True):
- this must be the type of the tensors that will be
- received, will typically be params_dtype, but in the case
- of fp32 residual connections might be torch.float.
-
- variable_seq_lengths (bool, optional, default=False):
- Support for variable sequence lengths across
- microbatches. Setting this communicates the size of
- tensors during pipeline parallelism communication, because
- of this extra overhead it should only be set if the
- sequence length is not constant during training.
-
- use_ring_exchange_p2p (bool, optional, default = False):
- Use custom ring_exchange kernel instead of
- torch.distributed.batch_isend_irecv(). Requires custom
- built torch with torch.distributed.ring_exchange.
-
+ wait_on_reqs (boolean, optional, default=False):
+ For non-batched p2p communication, wait on each request
+ before returning.
Returns:
tuple containing
@@ -162,91 +273,142 @@ def _communicate(*, tensor_send_next: Optional[torch.Tensor],
"""
- # Create placeholder tensors for receive in forward and backward directions
- # if needed.
- tensor_recv_prev = None
- tensor_recv_next = None
+ tensor_recv_prev_func = None
+ tensor_recv_next_func = None
- if not variable_seq_lengths:
+ if not config.variable_seq_lengths:
recv_prev_shape = tensor_shape
recv_next_shape = tensor_shape
else:
- recv_prev_shape, recv_next_shape = \
- _communicate_shapes(tensor_send_next,
- tensor_send_prev,
- recv_prev,
- recv_next)
+ recv_prev_shape, recv_next_shape = _communicate_shapes(
+ tensor_send_next, tensor_send_prev, recv_prev, recv_next, config
+ )
+
+ def create_tensor_recv_prev():
+ return torch.empty(
+ recv_prev_shape,
+ requires_grad=True,
+ device=torch.cuda.current_device(),
+ dtype=config.pipeline_dtype,
+ )
+
+ def create_tensor_recv_next():
+ return torch.empty(
+ recv_next_shape,
+ requires_grad=True,
+ device=torch.cuda.current_device(),
+ dtype=config.pipeline_dtype,
+ )
if recv_prev:
- if dtype is None:
- raise RuntimeError("dtype must be provided if recv_prev is True")
+ if config.pipeline_dtype is None:
+ raise RuntimeError("pipeline_dtype must be provided if recv_prev is True")
if tensor_shape is None:
raise RuntimeError(
"tensor_shape must be specified if recv_prev is True. "
"Common tensor_shape is (seq_length, micro_batch_size, hidden_size)"
)
- tensor_recv_prev = torch.empty(recv_prev_shape,
- requires_grad=True,
- device=torch.cuda.current_device(),
- dtype=dtype)
+ tensor_recv_prev_func = create_tensor_recv_prev
+
if recv_next:
- if dtype is None:
+ if config.pipeline_dtype is None:
raise RuntimeError("dtype must be provided if recv_next is True")
if tensor_shape is None:
raise RuntimeError(
"tensor_shape must be specified if recv_next is True. "
"Common tensor_shape is (seq_length, micro_batch_size, hidden_size)"
)
- tensor_recv_next = torch.empty(recv_next_shape,
- requires_grad=True,
- device=torch.cuda.current_device(),
- dtype=dtype)
+ tensor_recv_next_func = create_tensor_recv_next
# Send tensors in both the forward and backward directions as appropriate.
- if use_ring_exchange_p2p:
- torch.distributed.ring_exchange(tensor_send_prev=tensor_send_prev,
- tensor_recv_prev=tensor_recv_prev,
- tensor_send_next=tensor_send_next,
- tensor_recv_next=tensor_recv_next,
- group=get_pipeline_model_parallel_group())
+ if config.use_ring_exchange_p2p:
+
+ def _ring_exchange_wrapper(**kwargs):
+ torch.distributed.ring_exchange(**kwargs)
+ return []
+
+ p2p_func = _ring_exchange_wrapper
+ elif config.batch_p2p_comm:
+ assert wait_on_reqs
+ p2p_func = _batched_p2p_ops
else:
- ops = []
- if tensor_send_prev is not None:
- send_prev_op = torch.distributed.P2POp(
- torch.distributed.isend, tensor_send_prev,
- get_pipeline_model_parallel_prev_rank())
- ops.append(send_prev_op)
- if tensor_recv_prev is not None:
- recv_prev_op = torch.distributed.P2POp(
- torch.distributed.irecv, tensor_recv_prev,
- get_pipeline_model_parallel_prev_rank())
- ops.append(recv_prev_op)
- if tensor_send_next is not None:
- send_next_op = torch.distributed.P2POp(
- torch.distributed.isend, tensor_send_next,
- get_pipeline_model_parallel_next_rank())
- ops.append(send_next_op)
- if tensor_recv_next is not None:
- recv_next_op = torch.distributed.P2POp(
- torch.distributed.irecv, tensor_recv_next,
- get_pipeline_model_parallel_next_rank())
- ops.append(recv_next_op)
- if len(ops) > 0:
- reqs = torch.distributed.batch_isend_irecv(ops)
- for req in reqs:
- req.wait()
+ p2p_func = _p2p_ops
+
+ # Each rank can now be part of several different pipeline parallel groups
+ # (specifically, this can occur when encoder tensor parallelism != decoder
+ # tensor parallelism, and hence a rank in the encoder is going to feed
+ # several different decoder ranks. We therefore have to receive or send tensors
+ # from several groups. For convenience, I wrap everything into lists.
+ pp_group = get_pipeline_model_parallel_group()
+ next_rank = get_pipeline_model_parallel_next_rank()
+ prev_rank = get_pipeline_model_parallel_prev_rank()
+ if not isinstance(pp_group, list):
+ pp_group = [pp_group]
+ assert not isinstance(next_rank, list)
+ next_rank = [next_rank]
+ assert not isinstance(prev_rank, list)
+ prev_rank = [prev_rank]
+
+ reqs = []
+ tensor_recv_prev_list = []
+ tensor_recv_next_list = []
+
+ for group, nr, pr in zip(pp_group, next_rank, prev_rank):
+ if tensor_recv_prev_func is not None:
+ tensor_recv_prev = tensor_recv_prev_func()
+ tensor_recv_prev_list.append(tensor_recv_prev)
+ else:
+ tensor_recv_prev = None
+
+ if tensor_recv_next_func is not None:
+ tensor_recv_next = tensor_recv_next_func()
+ tensor_recv_next_list.append(tensor_recv_next)
+ else:
+ tensor_recv_next = None
+
+ reqs.extend(
+ p2p_func(
+ tensor_send_prev=tensor_send_prev,
+ tensor_recv_prev=tensor_recv_prev,
+ tensor_send_next=tensor_send_next,
+ tensor_recv_next=tensor_recv_next,
+ group=group,
+ prev_pipeline_rank=pr,
+ next_pipeline_rank=nr,
+ )
+ )
+
+ if wait_on_reqs and len(reqs) > 0:
+ for req in reqs:
+ req.wait()
+ reqs = None
+
+ if config.batch_p2p_comm and config.batch_p2p_sync:
# To protect against race condition when using batch_isend_irecv().
# User should assert that we have a modern enough PyTorch to not need this
torch.cuda.synchronize()
- return tensor_recv_prev, tensor_recv_next
+ def _handle_tensor_list(x):
+ """This basically handles all the cases that we expect to see. Either the list None,
+ or it's a singleton (the usual cases, since most ranks only belong to one pipeline group),
+ or everything returned is None, or everything returned is not None, and it has to be summed
+ together."""
+ if len(x) == 0:
+ return None
+ if len(x) == 1:
+ return x[0]
+ if all(xx is None for xx in x):
+ return None
+ return torch.stack(x, dim=0).sum(dim=0, dtype=torch.float32).to(x[0].dtype)
+
+ tensor_recv_prev = _handle_tensor_list(tensor_recv_prev_list)
+ tensor_recv_next = _handle_tensor_list(tensor_recv_next_list)
+ return tensor_recv_prev, tensor_recv_next, reqs
-def recv_forward(tensor_shape: Shape,
- dtype: torch.dtype,
- timers: Callable = None) -> torch.Tensor:
- """ Receive tensor from previous rank in pipeline (forward receive).
+def recv_forward(tensor_shape: Shape, config: ModelParallelConfig) -> torch.Tensor:
+ """Receive tensor from previous rank in pipeline (forward receive).
See _communicate for argument details.
"""
@@ -254,23 +416,22 @@ def recv_forward(tensor_shape: Shape,
if core.parallel_state.is_pipeline_first_stage():
input_tensor = None
else:
- if timers is not None:
- timers('forward-recv', log_level=2).start()
- input_tensor, _ = _communicate(
+ if config.timers is not None:
+ config.timers('forward-recv', log_level=2).start()
+ input_tensor, _, _ = _communicate(
tensor_send_next=None,
tensor_send_prev=None,
recv_prev=True,
recv_next=False,
tensor_shape=tensor_shape,
- dtype=dtype)
- if timers is not None:
- timers('forward-recv').stop()
+ config=config,
+ )
+ if config.timers is not None:
+ config.timers('forward-recv').stop()
return input_tensor
-def recv_backward(tensor_shape: Shape,
- dtype: torch.dtype,
- timers: Callable = None) -> torch.Tensor:
+def recv_backward(tensor_shape: Shape, config: ModelParallelConfig) -> torch.Tensor:
"""Receive tensor from next rank in pipeline (backward receive).
See _communicate for argument details.
@@ -278,65 +439,65 @@ def recv_backward(tensor_shape: Shape,
if core.parallel_state.is_pipeline_last_stage():
output_tensor_grad = None
else:
- if timers is not None:
- timers('backward-recv', log_level=2).start()
- _, output_tensor_grad = _communicate(
+ if config.timers is not None:
+ config.timers('backward-recv', log_level=2).start()
+ _, output_tensor_grad, _ = _communicate(
tensor_send_next=None,
tensor_send_prev=None,
recv_prev=False,
recv_next=True,
tensor_shape=tensor_shape,
- dtype=dtype)
- if timers is not None:
- timers('backward-recv').stop()
+ config=config,
+ )
+ if config.timers is not None:
+ config.timers('backward-recv').stop()
return output_tensor_grad
-def send_forward(output_tensor: torch.Tensor,
- timers: Callable = None) -> None:
+def send_forward(output_tensor: torch.Tensor, config: ModelParallelConfig) -> None:
"""Send tensor to next rank in pipeline (forward send).
See _communicate for argument details.
"""
if not core.parallel_state.is_pipeline_last_stage():
- if timers is not None:
- timers('forward-send', log_level=2).start()
+ if config.timers is not None:
+ config.timers('forward-send', log_level=2).start()
_communicate(
tensor_send_next=output_tensor,
tensor_send_prev=None,
recv_prev=False,
recv_next=False,
tensor_shape=None,
- dtype=None)
- if timers is not None:
- timers('forward-send').stop()
+ config=config,
+ )
+ if config.timers is not None:
+ config.timers('forward-send').stop()
-def send_backward(input_tensor_grad: torch.Tensor,
- timers: Callable = None) -> None:
+def send_backward(input_tensor_grad: torch.Tensor, config: ModelParallelConfig) -> None:
"""Send tensor to previous rank in pipeline (backward send).
See _communicate for argument details.
"""
if not core.parallel_state.is_pipeline_first_stage():
- if timers is not None:
- timers('backward-send', log_level=2).start()
+ if config.timers is not None:
+ config.timers('backward-send', log_level=2).start()
_communicate(
tensor_send_next=None,
tensor_send_prev=input_tensor_grad,
recv_prev=False,
recv_next=False,
tensor_shape=None,
- dtype=None)
- if timers is not None:
- timers('backward-send').stop()
+ config=config,
+ )
+ if config.timers is not None:
+ config.timers('backward-send').stop()
-def send_forward_recv_backward(output_tensor: torch.Tensor,
- tensor_shape: Shape,
- dtype: torch.dtype,
- timers: Callable = None) -> torch.Tensor:
+def send_forward_recv_backward(
+ output_tensor: torch.Tensor, tensor_shape: Shape, config: ModelParallelConfig
+) -> torch.Tensor:
"""Batched send and recv with next rank in pipeline.
See _communicate for argument details.
@@ -344,24 +505,24 @@ def send_forward_recv_backward(output_tensor: torch.Tensor,
if core.parallel_state.is_pipeline_last_stage():
output_tensor_grad = None
else:
- if timers is not None:
- timers('forward-send-backward-recv', log_level=2).start()
- _, output_tensor_grad = _communicate(
+ if config.timers is not None:
+ config.timers('forward-send-backward-recv', log_level=2).start()
+ _, output_tensor_grad, _ = _communicate(
tensor_send_next=output_tensor,
tensor_send_prev=None,
recv_prev=False,
recv_next=True,
tensor_shape=tensor_shape,
- dtype=dtype)
- if timers is not None:
- timers('forward-send-backward-recv').stop()
+ config=config,
+ )
+ if config.timers is not None:
+ config.timers('forward-send-backward-recv').stop()
return output_tensor_grad
-def send_backward_recv_forward(input_tensor_grad: torch.Tensor,
- tensor_shape: Shape,
- dtype: torch.dtype,
- timers: Callable = None) -> torch.Tensor:
+def send_backward_recv_forward(
+ input_tensor_grad: torch.Tensor, tensor_shape: Shape, config: ModelParallelConfig
+) -> torch.Tensor:
"""Batched send and recv with previous rank in pipeline.
See _communicate for argument details.
@@ -369,88 +530,101 @@ def send_backward_recv_forward(input_tensor_grad: torch.Tensor,
if core.parallel_state.is_pipeline_first_stage():
input_tensor = None
else:
- if timers is not None:
- timers('backward-send-forward-recv', log_level=2).start()
- input_tensor, _ = _communicate(
+ if config.timers is not None:
+ config.timers('backward-send-forward-recv', log_level=2).start()
+ input_tensor, _, _ = _communicate(
tensor_send_next=None,
tensor_send_prev=input_tensor_grad,
recv_prev=True,
recv_next=False,
tensor_shape=tensor_shape,
- dtype=dtype)
- if timers is not None:
- timers('backward-send-forward-recv').stop()
+ config=config,
+ )
+ if config.timers is not None:
+ config.timers('backward-send-forward-recv').stop()
return input_tensor
-def send_forward_recv_forward(output_tensor: torch.Tensor,
- recv_prev: bool,
- tensor_shape: Shape,
- dtype: torch.dtype,
- timers: Callable = None) -> torch.Tensor:
+def send_forward_recv_forward(
+ output_tensor: torch.Tensor,
+ recv_prev: bool,
+ tensor_shape: Shape,
+ config: ModelParallelConfig,
+ overlap_p2p_comm: bool = False,
+) -> torch.Tensor:
"""Batched recv from previous rank and send to next rank in pipeline.
See _communicate for argument details.
"""
- if timers is not None:
- timers('forward-send-forward-recv', log_level=2).start()
- input_tensor, _ = _communicate(
+ if config.timers is not None:
+ config.timers('forward-send-forward-recv', log_level=2).start()
+ input_tensor, _, wait_handles = _communicate(
tensor_send_next=output_tensor,
tensor_send_prev=None,
recv_prev=recv_prev,
recv_next=False,
tensor_shape=tensor_shape,
- dtype=dtype)
- if timers is not None:
- timers('forward-send-forward-recv').stop()
+ wait_on_reqs=(not overlap_p2p_comm),
+ config=config,
+ )
+ if config.timers is not None:
+ config.timers('forward-send-forward-recv').stop()
+ if overlap_p2p_comm:
+ return input_tensor, wait_handles
return input_tensor
-def send_backward_recv_backward(input_tensor_grad: torch.Tensor,
- recv_next: bool,
- tensor_shape: Shape,
- dtype: torch.dtype,
- timers: Callable = None) -> torch.Tensor:
+def send_backward_recv_backward(
+ input_tensor_grad: torch.Tensor,
+ recv_next: bool,
+ tensor_shape: Shape,
+ config: ModelParallelConfig,
+ overlap_p2p_comm: bool = False,
+) -> torch.Tensor:
"""Batched recv from next rank and send to previous rank in pipeline.
See _communicate for argument details.
"""
- if timers is not None:
- timers('backward-send-backward-recv', log_level=2).start()
- _, output_tensor_grad = _communicate(
+ if config.timers is not None:
+ config.timers('backward-send-backward-recv', log_level=2).start()
+ _, output_tensor_grad, wait_handles = _communicate(
tensor_send_next=None,
tensor_send_prev=input_tensor_grad,
recv_prev=False,
recv_next=recv_next,
tensor_shape=tensor_shape,
- dtype=dtype)
- if timers is not None:
- timers('backward-send-backward-recv').stop()
+ wait_on_reqs=(not overlap_p2p_comm),
+ config=config,
+ )
+ if config.timers is not None:
+ config.timers('backward-send-backward-recv').stop()
+ if overlap_p2p_comm:
+ return output_tensor_grad, wait_handles
return output_tensor_grad
def send_forward_backward_recv_forward_backward(
- output_tensor: torch.Tensor,
- input_tensor_grad: torch.Tensor,
- recv_prev: bool,
- recv_next: bool,
- tensor_shape: Shape,
- dtype: torch.dtype,
- timers: Callable = None) -> Tuple[torch.Tensor, torch.Tensor]:
+ output_tensor: torch.Tensor,
+ input_tensor_grad: torch.Tensor,
+ recv_prev: bool,
+ recv_next: bool,
+ tensor_shape: Shape,
+ config: ModelParallelConfig,
+) -> torch.Tensor:
"""Batched send and recv with previous and next ranks in pipeline.
See _communicate for argument details.
"""
- if timers is not None:
- timers('forward-backward-send-forward-backward-recv',
- log_level=2).start()
- input_tensor, output_tensor_grad = _communicate(
+ if config.timers is not None:
+ config.timers('forward-backward-send-forward-backward-recv', log_level=2).start()
+ input_tensor, output_tensor_grad, _ = _communicate(
tensor_send_next=output_tensor,
tensor_send_prev=input_tensor_grad,
recv_prev=recv_prev,
recv_next=recv_next,
tensor_shape=tensor_shape,
- dtype=dtype)
- if timers is not None:
- timers('forward-backward-send-forward-backward-recv').stop()
+ config=config,
+ )
+ if config.timers is not None:
+ config.timers('forward-backward-send-forward-backward-recv').stop()
return input_tensor, output_tensor_grad
diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py
index 5007a44cd2..f082dbc6df 100644
--- a/megatron/core/pipeline_parallel/schedules.py
+++ b/megatron/core/pipeline_parallel/schedules.py
@@ -1,20 +1,27 @@
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
import contextlib
-from typing import Callable, Iterator, List, Optional, Union
+from typing import Iterator, List, Union
import torch
from torch.autograd.variable import Variable
-from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP
from megatron.core import parallel_state
-from megatron.core.pipeline_parallel import p2p_communication
from megatron.core.enums import ModelType
-from megatron.core.utils import get_attr_wrapped_model, get_model_type
+from megatron.core.pipeline_parallel import p2p_communication
+from megatron.core.transformer.moe.router import MoEAuxLossAutoScaler
+from megatron.core.utils import (
+ drain_embedding_wgrad_compute,
+ get_attr_wrapped_model,
+ get_model_config,
+ get_model_type,
+ get_model_xattn,
+)
# Types
Shape = Union[List[int], torch.Size]
+
def get_forward_backward_func():
"""Retrieves the appropriate forward_backward function given the
configuration of parallel_state.
@@ -24,6 +31,10 @@ def get_forward_backward_func():
world size and virtual pipeline model parallel world size in the
global parallel_state.
+ Note that if using sequence parallelism, the sequence length component of
+ the tensor shape is updated to original_sequence_length /
+ tensor_model_parallel_world_size.
+
The function returned takes the following arguments:
forward_step_func (required): A function that takes a data
@@ -32,6 +43,13 @@ def get_forward_backward_func():
take one torch.Tensor and return a torch.Tensor of loss and a
dictionary of string -> torch.Tensor.
+ A third argument, checkpoint_activations_microbatch, indicates
+ that the activations for this microbatch should be
+ checkpointed. A None value for this argument indicates that
+ the default from the configuration should be used. This is
+ used when the
+ num_microbatches_with_partial_activation_checkpoints is used.
+
For example:
def loss_func(loss_mask, output_tensor):
@@ -57,62 +75,29 @@ def forward_step(data_iterator, model):
passed as is to forward_step_func. Expected to be a list of
iterators in the case of interleaved pipeline parallelism.
- model (required): the actual model. Expected to be a list of
- modules in the case of interleaved pipeline parallelism.
+ model (required): the actual model. Expected to be a list of modules in the case of interleaved
+ pipeline parallelism. Must be a (potentially wrapped) megatron.core.models.MegatronModule.
num_microbatches (int, required):
The number of microbatches to go through
- dtype (required when using pipeline parallelism): dtype used in
- p2p communication, usually params_dtype
-
- tensor_shape (required when using pipeline parallelism): Shape of
- tensor. The tensor is expected to be 3D and its order of
- dimension is supposed to be ``(sequence, batch, hidden)``.
-
- decoder_seq_length (int, required for ModelType.encoder_and_decoder models):
- Sequence length of the decoder portion, used to determine tensor shapes.
-
- grad_scaler (optional, default=None): If using loss scaling,
- this function should take the loss and return the scaled
- loss. If None, no function is called on the loss.
-
- sequence_parallel (optional, default=False):
- Set to :obj:`True` for this function to handle sequence
- length. When :obj:`True`, the sequence length on each tensor
- model parallel rank is updated to
- :math:`original\_sequence\_length /
- tensor\_model\_parallel\_world\_size`.
- TODO: Do we need this? Just roll into tensor_shape arg?
-
- forward_only (optional, default=False): Perform only the forward step
+ seq_length (int, required): Sequence length of the current global batch. If this is a dual-stack
+ transformer, this is the encoder's sequence length. This is ignored if variable_seq_lengths
+ in the config is True. Otherwise, each microbatch in the current global batch size must use
+ this sequence length.
- timers (optional, default=None): TODO
+ micro_batch_size (int, required): The number of sequences in a microbatch.
- collect_non_loss_data: TODO
+ decoder_seq_length (int, optional): The sequence length for the decoder in a dual-stack
+ transformer. This is ignored for a single-stack transformer.
- enable_autocast (optional, default=False): If True, runs the
- forward_step_func call inside torch.autocast context
+ forward_only (optional, default = False): Perform only the forward step
- deallocate_pipeline_outputs (optional, default=False): If True, output data
- is deallocated after the tensor is sent to the next pipeline stage.
- Helps with saving memory, does nothing when pipeline parallel is
- not used.
-
- no_sync_func (optional): Function that creates a context that
- suppresses asynchronous data-parallel communication. If the
- model is an instance of torch.nn.DistributedDataParallel, the
- default is to use torch.nn.DistributedDataParallel.no_sync.
+ collect_non_loss_data (optional, bool, default=False): TODO
- grad_sync_func (optional): Function that launches asynchronous
- gradient reductions (e.g. distributed optimizer gradient
- reduce-scatters). The function should take one argument: an
- iterable of parameters whose gradients are to be synchronized.
-
- param_sync_func (optional): Function that launches asynchronous
- parameter synchronizations (e.g. distributed optimizer
- parameter all-gathers). The function should take one argument:
- an iterable of parameters to be synchronized.
+ first_val_step (bool, optional): Is the first step of the validation phase. Used by
+ Transformer Engine modules to only update their fp8 weights only on the first validation
+ step.
"""
pipeline_model_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
@@ -125,6 +110,7 @@ def forward_step(data_iterator, model):
forward_backward_func = forward_backward_no_pipelining
return forward_backward_func
+
def deallocate_output_tensor(out, deallocate_pipeline_outputs=False):
'''Pseudo-deallocate (i.e., set to scalar) the output tensor's '.data' field.
@@ -134,15 +120,10 @@ def deallocate_output_tensor(out, deallocate_pipeline_outputs=False):
'''
if (out is None) or (not deallocate_pipeline_outputs):
return
- assert isinstance(out, torch.Tensor), \
- "expected Tensor, found %s." % type(out).__name__
- assert out._base is None, \
- "counter-productive to free a view of another tensor."
- out.data = torch.empty(
- (1,),
- device = out.device,
- dtype = out.dtype,
- )
+ assert isinstance(out, torch.Tensor), "expected Tensor, found %s." % type(out).__name__
+ assert out._base is None, "counter-productive to free a view of another tensor."
+ out.data = torch.empty((1,), device=out.device, dtype=out.dtype)
+
def custom_backward(output, grad_output):
'''Directly call C++ autograd engine.
@@ -153,54 +134,127 @@ def custom_backward(output, grad_output):
grad have the same shape, while C++'s 'backward' does not.
'''
- assert output.numel() == 1, \
- "output should be pseudo-'freed' in schedule, to optimize memory"
- assert isinstance(output, torch.Tensor), \
- "output == '%s'." % type(output).__name__
- assert isinstance(grad_output, (torch.Tensor, type(None))), \
+ assert output.numel() == 1, "output should be pseudo-'freed' in schedule, to optimize memory"
+ assert isinstance(output, torch.Tensor), "output == '%s'." % type(output).__name__
+ assert isinstance(grad_output, (torch.Tensor, type(None))), (
"grad_output == '%s'." % type(grad_output).__name__
+ )
# Handle scalar output
if grad_output is None:
assert output.numel() == 1, "implicit grad requires scalar output."
- grad_output = torch.ones_like(
- output,
- memory_format = torch.preserve_format,
- )
+ grad_output = torch.ones_like(output, memory_format=torch.preserve_format)
# Call c++ engine [ see torch/csrc/autograd/python_engine.cpp ]
Variable._execution_engine.run_backward(
- tensors = (output,),
- grad_tensors = (grad_output,),
- keep_graph = False,
- create_graph = False,
- inputs = tuple(),
+ tensors=(output,),
+ grad_tensors=(grad_output,),
+ keep_graph=False,
+ create_graph=False,
+ inputs=tuple(),
allow_unreachable=True,
accumulate_grad=True,
)
-
-
-
-def forward_step(forward_step_func,
- data_iterator,
- model,
- num_microbatches,
- input_tensor,
- forward_data_store,
- timers,
- collect_non_loss_data=False,
- autocast_dtype=torch.float,
- enable_autocast=False):
+def set_current_microbatch(model, microbatch_id):
+ decoder_exists = True
+ decoder = None
+ try:
+ decoder = get_attr_wrapped_model(model, "decoder")
+ except RuntimeError:
+ decoder_exists = False
+ if decoder_exists and decoder is not None:
+ decoder.current_microbatch = microbatch_id
+
+
+def forward_step(
+ forward_step_func,
+ data_iterator,
+ model,
+ num_microbatches,
+ input_tensor,
+ forward_data_store,
+ config,
+ collect_non_loss_data=False,
+ checkpoint_activations_microbatch=None,
+ is_first_microbatch=False,
+ current_microbatch=None,
+ encoder_decoder_xattn=False,
+):
"""Forward step for passed-in model.
- If first stage, input tensor is obtained from data_iterator, otherwise
- passed-in input_tensor is used.
+ If it is the first stage, the input tensor is obtained from the data_iterator.
+ Otherwise, the passed-in input_tensor is used.
+
+ Args:
+ forward_step_func (callable):
+ The forward step function for the model that takes the
+ data iterator as the first argument, and model as the second.
+ This user's forward step is expected to output a tuple of two elements:
+
+ 1. The output object from the forward step. This output object needs to be a
+ tensor or some kind of collection of tensors. The only hard requirement
+ for this object is that it needs to be acceptible as input into the second
+ function.
+ 2. A function to reduce (optionally) the output from the forward step. This
+ could be a reduction over the loss from the model, it could be a function that
+ grabs the output from the model and reformats, it could be a function that just
+ passes through the model output. This function must have one of the following
+ patterns, and depending on the pattern different things happen internally:
+
+ a. A tuple of reduced loss and some other data. Note that in this case
+ the first argument is divided by the number of global microbatches,
+ assuming it is a loss, so that the loss is stable as a function of
+ the number of devices the step is split across.
+ b. A triple of reduced loss, number of tokens, and some other data. This
+ is similar to case (a), but the loss is further averaged across the
+ number of tokens in the batch. If the user is not already averaging
+ across the number of tokens, this pattern is useful to use.
+ c. Any arbitrary data the user wants (eg a dictionary of tensors, a list
+ of tensors, etc in the case of inference). To trigger case 3 you need
+ to specify `collect_non_loss_data=True` and you may also want to
+ specify `forward_only=True` in the call to the parent forward_backward
+ function.
+ data_iterator (iterator):
+ The data iterator.
+ model (nn.Module):
+ The model to perform the forward step on.
+ num_microbatches (int):
+ The number of microbatches.
+ input_tensor (Tensor or list[Tensor]):
+ The input tensor(s) for the forward step.
+ forward_data_store (list):
+ The list to store the forward data. If you go down path 2.a or
+ 2.b for the return of your forward reduction function then this will store only the
+ final dimension of the output, for example the metadata output by the loss function.
+ If you go down the path of 2.c then this will store the entire output of the forward
+ reduction function applied to the model output.
+ config (object):
+ The configuration object.
+ collect_non_loss_data (bool, optional):
+ Whether to collect non-loss data. Defaults to False.
+ This is the path to use if you want to collect arbitrary output from the model forward,
+ such as with inference use cases. Defaults to False.
+ checkpoint_activations_microbatch (int, optional):
+ The microbatch to checkpoint activations.
+ Defaults to None.
+ is_first_microbatch (bool, optional):
+ Whether it is the first microbatch. Defaults to False.
+ current_microbatch (int, optional):
+ The current microbatch. Defaults to None.
+
+ Returns:
+ Tensor or list[Tensor]: The output object(s) from the forward step.
+ Tensor: The number of tokens.
+ """
+ if config.timers is not None:
+ config.timers('forward-compute', log_level=2).start()
- Returns output tensor."""
- if timers is not None:
- timers('forward-compute', log_level=2).start()
+ if is_first_microbatch and hasattr(model, 'set_is_first_microbatch'):
+ model.set_is_first_microbatch()
+ if current_microbatch is not None:
+ set_current_microbatch(model, current_microbatch)
unwrap_output_tensor = False
if not isinstance(input_tensor, list):
@@ -210,41 +264,69 @@ def forward_step(forward_step_func,
set_input_tensor = get_attr_wrapped_model(model, "set_input_tensor")
set_input_tensor(input_tensor)
- if enable_autocast:
- context_manager = torch.autocast("cuda", dtype=autocast_dtype)
+ if config.enable_autocast:
+ context_manager = torch.autocast("cuda", dtype=config.autocast_dtype)
else:
context_manager = contextlib.nullcontext()
with context_manager:
- output_tensor, loss_func = forward_step_func(data_iterator, model)
+ if checkpoint_activations_microbatch is None:
+ output_tensor, loss_func = forward_step_func(data_iterator, model)
+ else:
+ output_tensor, loss_func = forward_step_func(
+ data_iterator, model, checkpoint_activations_microbatch
+ )
+ num_tokens = torch.tensor(0, dtype=torch.int)
if parallel_state.is_pipeline_last_stage():
if not collect_non_loss_data:
- output_tensor = loss_func(output_tensor)
- loss, loss_reduced = output_tensor
- output_tensor = loss / num_microbatches
+ outputs = loss_func(output_tensor)
+ if len(outputs) == 3:
+ output_tensor, num_tokens, loss_reduced = outputs
+ if not config.calculate_per_token_loss:
+ output_tensor /= num_tokens
+ output_tensor /= num_microbatches
+ else:
+ # preserve legacy loss averaging behavior (ie, over the number of microbatches)
+ assert len(outputs) == 2
+ output_tensor, loss_reduced = outputs
+ output_tensor /= num_microbatches
forward_data_store.append(loss_reduced)
else:
data = loss_func(output_tensor, non_loss_data=True)
forward_data_store.append(data)
- if timers is not None:
- timers('forward-compute').stop()
+ if config.timers is not None:
+ config.timers('forward-compute').stop()
+
+ # Set the loss scale for the auxiliary loss of the MoE layer.
+ # Since we use a trick to do backward on the auxiliary loss, we need to set the scale
+ # explicitly.
+ if hasattr(config, 'num_moe_experts') and config.num_moe_experts is not None:
+ # Calculate the loss scale based on the grad_scale_func if available, else default to 1.
+ loss_scale = (
+ config.grad_scale_func(torch.ones(1, device=output_tensor.device))
+ if config.grad_scale_func is not None
+ else torch.tensor(1.0)
+ )
+ # Set the loss scale
+ MoEAuxLossAutoScaler.set_loss_scale(loss_scale / num_microbatches)
- # If T5 model (or other model with encoder and decoder)
- # and in decoder stack, then send encoder_hidden_state
+ # If T5 model and in decoder stack, then send encoder_hidden_state
# downstream as well.
model_type = get_model_type(model)
+ if (
+ model_type == ModelType.encoder_and_decoder
+ and encoder_decoder_xattn
+ and parallel_state.is_inside_decoder()
+ ):
+ return [output_tensor, input_tensor[-1]], num_tokens
- if parallel_state.is_pipeline_stage_after_split() and \
- model_type == ModelType.encoder_and_decoder:
- return [output_tensor, input_tensor[-1]]
if unwrap_output_tensor:
- return output_tensor
- return [output_tensor]
+ return output_tensor, num_tokens
+ return [output_tensor], num_tokens
-def backward_step(grad_scaler, input_tensor, output_tensor,
- output_tensor_grad, model_type, timers, deallocate_pipeline_outputs=False):
+def backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config):
"""Backward step through passed-in output tensor.
If last stage, output_tensor_grad is None, otherwise gradient of loss
@@ -257,8 +339,8 @@ def backward_step(grad_scaler, input_tensor, output_tensor,
# needs to be modified slightly to support arbitrary numbers of skip
# connections.
- if timers is not None:
- timers('backward-compute', log_level=2).start()
+ if config.timers is not None:
+ config.timers('backward-compute', log_level=2).start()
# Retain the grad on the input_tensor.
unwrap_input_tensor_grad = False
@@ -275,10 +357,10 @@ def backward_step(grad_scaler, input_tensor, output_tensor,
output_tensor_grad = [output_tensor_grad]
# Backward pass.
- if output_tensor_grad[0] is None and grad_scaler is not None:
- output_tensor = grad_scaler(output_tensor[0])
-
- if deallocate_pipeline_outputs:
+ if output_tensor_grad[0] is None and config.grad_scale_func is not None:
+ output_tensor[0] = config.grad_scale_func(output_tensor[0])
+
+ if config.deallocate_pipeline_outputs:
custom_backward(output_tensor[0], output_tensor_grad[0])
else:
torch.autograd.backward(output_tensor[0], grad_tensors=output_tensor_grad[0])
@@ -295,39 +377,43 @@ def backward_step(grad_scaler, input_tensor, output_tensor,
# Handle single skip connection if it exists (encoder_hidden_state in
# model with encoder and decoder).
- if parallel_state.get_pipeline_model_parallel_world_size() > 1 and \
- parallel_state.is_pipeline_stage_after_split() and \
- model_type == ModelType.encoder_and_decoder:
+ if (
+ parallel_state.get_pipeline_model_parallel_world_size() > 1
+ and model_type == ModelType.encoder_and_decoder
+ and len(output_tensor_grad) > 1 # excludes models that lack a skip connection.
+ ):
if output_tensor_grad[1] is not None:
+ assert input_tensor_grad[-1] is not None
input_tensor_grad[-1].add_(output_tensor_grad[1])
if unwrap_input_tensor_grad:
input_tensor_grad = input_tensor_grad[0]
- if timers is not None:
- timers('backward-compute').stop()
+ if config.timers is not None:
+ config.timers('backward-compute').stop()
return input_tensor_grad
-def forward_backward_no_pipelining(*,
- forward_step_func,
- data_iterator: Union[Iterator, List[Iterator]],
- model: Union[torch.nn.Module, List[torch.nn.Module]],
- num_microbatches: int,
- dtype: Optional[torch.dtype] = None,
- tensor_shape: Optional[Shape] = None, # unused
- decoder_seq_length: Optional[int] = None, # unused
- grad_scaler: Callable = None,
- sequence_parallel: bool = False, # unused
- forward_only: bool = False,
- timers: Callable = None,
- collect_non_loss_data: bool = False,
- enable_autocast: bool = False,
- deallocate_pipeline_outputs: bool = False,
- no_sync_func: Optional[Callable] = None,
- grad_sync_func: Optional[Callable] = None, # unused
- param_sync_func: Optional[Callable] = None, # unused
- ):
+def check_first_val_step(first_val_step, forward_only, cond):
+ if (first_val_step is not None) and forward_only:
+ return first_val_step and cond
+ else:
+ return cond
+
+
+def forward_backward_no_pipelining(
+ *,
+ forward_step_func,
+ data_iterator: Union[Iterator, List[Iterator]],
+ model: Union[torch.nn.Module, List[torch.nn.Module]],
+ num_microbatches: int,
+ seq_length: int, # unused
+ micro_batch_size: int, # unused
+ decoder_seq_length: int = None, # unused
+ forward_only: bool = False,
+ collect_non_loss_data: bool = False,
+ first_val_step: bool = None,
+):
"""Run forward and backward passes with no pipeline parallelism
(no inter-stage communication).
@@ -338,16 +424,19 @@ def forward_backward_no_pipelining(*,
"""
if isinstance(model, list):
- assert len(model) == 1, \
- "non-pipeline-parallel schedule does not support model chunking"
+ assert len(model) == 1, "non-pipeline-parallel schedule does not support model chunking"
model = model[0]
if isinstance(data_iterator, list):
- assert len(data_iterator) == 1, \
- "non-pipeline-parallel schedule does not support model chunking"
+ assert (
+ len(data_iterator) == 1
+ ), "non-pipeline-parallel schedule does not support model chunking"
data_iterator = data_iterator[0]
- if no_sync_func is None and isinstance(model, torchDDP):
- no_sync_func = model.no_sync
+ config = get_model_config(model)
+ if config.timers is not None:
+ config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)
+
+ no_sync_func = config.no_sync_func
if no_sync_func is None:
no_sync_func = contextlib.nullcontext
@@ -355,81 +444,174 @@ def forward_backward_no_pipelining(*,
forward_data_store = []
input_tensor, output_tensor_grad = None, None
+ total_num_tokens = torch.zeros([], dtype=torch.int, device="cuda")
with no_sync_func():
for i in range(num_microbatches - 1):
- output_tensor = forward_step(forward_step_func, data_iterator,
- model, num_microbatches, input_tensor, forward_data_store,
- timers, collect_non_loss_data, dtype, enable_autocast)
+ output_tensor, num_tokens = forward_step(
+ forward_step_func,
+ data_iterator,
+ model,
+ num_microbatches,
+ input_tensor,
+ forward_data_store,
+ config,
+ collect_non_loss_data,
+ is_first_microbatch=check_first_val_step(first_val_step, forward_only, i == 0),
+ current_microbatch=i,
+ )
+ total_num_tokens += num_tokens.item()
if not forward_only:
- backward_step(grad_scaler, input_tensor, output_tensor,
- output_tensor_grad, model_type, timers, deallocate_pipeline_outputs)
+ backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
# Run computation for last microbatch out of context handler (want to
# synchronize gradients).
- output_tensor = forward_step(forward_step_func, data_iterator,
- model, num_microbatches, input_tensor, forward_data_store,
- timers, collect_non_loss_data, dtype, enable_autocast)
+ output_tensor, num_tokens = forward_step(
+ forward_step_func,
+ data_iterator,
+ model,
+ num_microbatches,
+ input_tensor,
+ forward_data_store,
+ config,
+ collect_non_loss_data,
+ is_first_microbatch=check_first_val_step(
+ first_val_step, forward_only, num_microbatches == 1
+ ),
+ current_microbatch=num_microbatches - 1,
+ )
+ total_num_tokens += num_tokens.item()
if not forward_only:
- backward_step(grad_scaler, input_tensor, output_tensor,
- output_tensor_grad, model_type, timers, deallocate_pipeline_outputs)
+ backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
+
+ if config.finalize_model_grads_func is not None and not forward_only:
+ # Finalize model grads (perform full grad all-reduce / reduce-scatter for
+ # data parallelism and layernorm all-reduce for sequence parallelism).
+ config.finalize_model_grads_func(
+ [model], total_num_tokens if config.calculate_per_token_loss else None
+ )
+
+ if config.timers is not None:
+ config.timers('forward-backward').stop()
return forward_data_store
-def forward_backward_pipelining_with_interleaving(*,
- forward_step_func,
- data_iterator: Union[Iterator, List[Iterator]],
- model: Union[torch.nn.Module, List[torch.nn.Module]],
- num_microbatches: int,
- dtype: torch.dtype,
- tensor_shape: Shape,
- decoder_seq_length: Optional[int] = None,
- grad_scaler: Callable = None,
- sequence_parallel: bool = False,
- forward_only: bool = False,
- timers: Callable = None,
- collect_non_loss_data: bool = False,
- enable_autocast: bool = False,
- deallocate_pipeline_outputs: bool = False,
- no_sync_func: Optional[Callable] = None,
- grad_sync_func: Optional[Callable] = None,
- param_sync_func: Optional[Callable] = None,
- ):
+def clear_embedding_activation_buffer(config, model):
+
+ if (
+ parallel_state.is_pipeline_last_stage(ignore_virtual=True)
+ and config.defer_embedding_wgrad_compute
+ ):
+ if isinstance(model, list):
+ embedding_module = get_attr_wrapped_model(
+ model[-1], 'post_process', return_model_obj=True
+ )
+ else:
+ embedding_module = get_attr_wrapped_model(model, 'post_process', return_model_obj=True)
+
+ # Need to ensure no stray activations exists in this buffer
+ embedding_module.embedding_activation_buffer.clear()
+
+ return embedding_module
+ else:
+ return None
+
+
+def finish_embedding_wgrad_compute(config, embedding_module):
+ if (
+ parallel_state.is_pipeline_last_stage(ignore_virtual=True)
+ and config.defer_embedding_wgrad_compute
+ ):
+ embedding_activation_buffer = embedding_module.embedding_activation_buffer
+ grad_output_buffer = embedding_module.grad_output_buffer
+ weight = (
+ embedding_module.output_layer.weight
+ if embedding_module.share_embeddings_and_output_weights
+ else embedding_module.shared_embedding_or_output_weight()
+ )
+
+ drain_embedding_wgrad_compute(
+ config, embedding_activation_buffer, grad_output_buffer, weight
+ )
+
+
+def forward_backward_pipelining_with_interleaving(
+ *,
+ forward_step_func,
+ data_iterator: Union[Iterator, List[Iterator]],
+ model: Union[torch.nn.Module, List[torch.nn.Module]],
+ num_microbatches: int,
+ seq_length: int,
+ micro_batch_size: int,
+ decoder_seq_length: int = None,
+ forward_only: bool = False,
+ collect_non_loss_data: bool = False,
+ first_val_step: bool = None,
+):
"""Run interleaved 1F1B schedule (model split into model chunks), with
communication between pipeline stages as needed.
Returns dictionary with losses if the last stage, empty dict otherwise."""
- assert isinstance(model, list), \
- "interleaved pipeline parallelism expected model chunking"
- assert all(isinstance(chunk, torch.nn.Module) for chunk in model), \
- "invalid model chunking"
- assert isinstance(data_iterator, list), \
- "interleaved pipeline parallelism expected each model chunk to have a data iterator"
+ assert isinstance(model, list), "interleaved pipeline parallelism expected model chunking"
+ assert all(isinstance(chunk, torch.nn.Module) for chunk in model), "invalid model chunking"
+ assert isinstance(
+ data_iterator, list
+ ), "interleaved pipeline parallelism expected each model chunk to have a data iterator"
+
+ config = get_model_config(model[0])
+ if config.overlap_p2p_comm and config.batch_p2p_comm:
+ raise ValueError("Can not use both overlap_p2p_comm and batch_p2p_comm")
+
+ # Needed only when gradients are finalized in M-Core
+ if config.finalize_model_grads_func is not None and not forward_only:
+ embedding_module = clear_embedding_activation_buffer(config, model)
+
+ if config.timers is not None:
+ config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)
# Disable async grad reductions
- if no_sync_func is None and all(isinstance(chunk, torchDDP) for chunk in model):
+ no_sync_func = config.no_sync_func
+ if isinstance(no_sync_func, list):
+
def multi_no_sync():
stack = contextlib.ExitStack()
- for chunk in model:
- stack.enter_context(chunk.no_sync())
+ for model_chunk_no_sync_func in config.no_sync_func:
+ stack.enter_context(model_chunk_no_sync_func())
return stack
+
no_sync_func = multi_no_sync
if no_sync_func is None:
no_sync_func = contextlib.nullcontext
no_sync_context = None
+
+ if config.grad_sync_func is not None and not isinstance(config.grad_sync_func, list):
+ config.grad_sync_func = [config.grad_sync_func for _ in model]
+
+ if config.param_sync_func is not None and not isinstance(config.param_sync_func, list):
+ config.param_sync_func = [config.param_sync_func for _ in model]
+
+ # Disable config.grad_sync_func and config.param_sync_func if only running forward passes.
+ # They will be re-enabled at the end of this function.
+ grad_sync_func, param_sync_func = None, None
+ if forward_only:
+ grad_sync_func, param_sync_func = config.grad_sync_func, config.param_sync_func
+ config.grad_sync_func, config.param_sync_func = None, None
+
def disable_grad_sync():
"""Disable asynchronous grad reductions"""
nonlocal no_sync_context
if no_sync_context is None:
no_sync_context = no_sync_func()
no_sync_context.__enter__()
+
def enable_grad_sync():
"""Enable asynchronous grad reductions"""
nonlocal no_sync_context
if no_sync_context is not None:
no_sync_context.__exit__(None, None, None)
no_sync_context = None
+
disable_grad_sync()
# Model chunk IDs with synchronized grads
@@ -437,6 +619,8 @@ def enable_grad_sync():
input_tensors = [[] for _ in range(len(model))]
output_tensors = [[] for _ in range(len(model))]
+ total_num_tokens = torch.tensor(0, dtype=torch.int).cuda()
+
forward_data_store = []
if not forward_only:
output_tensor_grads = [[] for _ in range(len(model))]
@@ -454,17 +638,16 @@ def enable_grad_sync():
if model_type == ModelType.encoder_and_decoder:
raise RuntimeError("Interleaving is not supported with an encoder and decoder model.")
- if decoder_seq_length is not None and decoder_seq_length != tensor_shape[0]:
- raise RuntimeError("Interleaving is not supported with a different decoder sequence length.")
-
- if sequence_parallel:
- seq_length, batch_size, hidden = tensor_shape
- tensor_shape = (
- seq_length // parallel_state.get_tensor_model_parallel_world_size(),
- batch_size,
- hidden,
+ if decoder_seq_length is not None and decoder_seq_length != seq_length:
+ raise RuntimeError(
+ "Interleaving is not supported with a different decoder sequence length."
)
+ tensor_shape = [seq_length, micro_batch_size, config.hidden_size]
+ tensor_shape[0] = tensor_shape[0] // parallel_state.get_context_parallel_world_size()
+ if config.sequence_parallel:
+ tensor_shape[0] = tensor_shape[0] // parallel_state.get_tensor_model_parallel_world_size()
+
# Compute number of warmup and remaining microbatches.
num_model_chunks = len(model)
total_num_microbatches = num_microbatches * num_model_chunks
@@ -482,32 +665,48 @@ def enable_grad_sync():
num_warmup_microbatches = total_num_microbatches
all_warmup_microbatches = True
else:
- num_warmup_microbatches = \
- (pipeline_parallel_size - pipeline_parallel_rank - 1) * 2
- num_warmup_microbatches += (
- num_model_chunks - 1) * pipeline_parallel_size
- num_warmup_microbatches = min(num_warmup_microbatches,
- total_num_microbatches)
- num_microbatches_remaining = \
- total_num_microbatches - num_warmup_microbatches
+ num_warmup_microbatches = (pipeline_parallel_size - pipeline_parallel_rank - 1) * 2
+ num_warmup_microbatches += (num_model_chunks - 1) * pipeline_parallel_size
+ num_warmup_microbatches = min(num_warmup_microbatches, total_num_microbatches)
+ num_microbatches_remaining = total_num_microbatches - num_warmup_microbatches
+
+ # Checkpoint the activations of partial Transformer layers in a number of micro-batches
+ # within the maximum outstanding micro-batch backpropagations.
+ # Micro-batches with the ids less than 'num_microbatches_with_partial_activation_checkpoints'
+ # checkpoint partial Transformer layers (or skip checkpointing) and
+ # the rest of micro-batches within a window of micro-batches checkpoint
+ # all Transformer layers. The window of micro-batches is set by the maximum
+ # outstanding backpropagations and becomes smaller at later pipeline stages.
+ # Please refer the appendix C in https://arxiv.org/pdf/2205.05198.pdf
+ max_outstanding_backprops = None
+ if config.num_microbatches_with_partial_activation_checkpoints is not None:
+ max_outstanding_backprops = num_warmup_microbatches + 1
# Synchronize params for first two model chunks
- if param_sync_func is not None:
- param_sync_func(model[0].parameters())
- param_sync_func(model[1].parameters())
+ if config.param_sync_func is not None:
+ config.param_sync_func[0](model[0].parameters())
+ config.param_sync_func[1](model[1].parameters())
def get_model_chunk_id(microbatch_id, forward):
"""Helper method to get the model chunk ID given the iteration number."""
microbatch_id_in_group = microbatch_id % (pipeline_parallel_size * num_model_chunks)
model_chunk_id = microbatch_id_in_group // pipeline_parallel_size
if not forward:
- model_chunk_id = (num_model_chunks - model_chunk_id - 1)
+ model_chunk_id = num_model_chunks - model_chunk_id - 1
return model_chunk_id
+ def get_microbatch_id_in_model_chunk(iteration_id, forward):
+ """Helper method to get the microbatch_id within model chunk given the iteration number."""
+ assert forward
+ iteration_group_id = iteration_id // (pipeline_parallel_size * num_model_chunks)
+ microbatch_id_in_model_chunk = (iteration_group_id * pipeline_parallel_size) + (
+ iteration_id % pipeline_parallel_size
+ )
+ return microbatch_id_in_model_chunk
+
def is_first_microbatch_for_model_chunk(microbatch_id: int) -> bool:
"""Check if an iteration is the first for a model chunk."""
microbatch_group_size = pipeline_parallel_size * num_model_chunks
- num_microbatch_groups = num_microbatches // microbatch_group_size
microbatch_group_id = microbatch_id // microbatch_group_size
microbatch_id_in_group = microbatch_id % microbatch_group_size
if microbatch_group_id == 0:
@@ -518,7 +717,7 @@ def is_first_microbatch_for_model_chunk(microbatch_id: int) -> bool:
def is_last_microbatch_for_model_chunk(microbatch_id: int) -> bool:
"""Check if an iteration is the last for a model chunk."""
microbatch_group_size = pipeline_parallel_size * num_model_chunks
- num_microbatch_groups = num_microbatches // microbatch_group_size
+ num_microbatch_groups = total_num_microbatches // microbatch_group_size
microbatch_group_id = microbatch_id // microbatch_group_size
microbatch_id_in_group = microbatch_id % microbatch_group_size
if microbatch_group_id == num_microbatch_groups - 1:
@@ -526,8 +725,7 @@ def is_last_microbatch_for_model_chunk(microbatch_id: int) -> bool:
else:
return False
-
- def forward_step_helper(microbatch_id):
+ def forward_step_helper(microbatch_id, current_microbatch, checkpoint_activations_microbatch):
"""Helper method to run forward step with model split into chunks
(run set_virtual_pipeline_model_parallel_rank() before calling
forward_step())."""
@@ -539,31 +737,44 @@ def forward_step_helper(microbatch_id):
# To reduce idling from mismatched microbatch times, we launch
# asynchronous communication at the same time across the
# pipeline-parallel group.
- if param_sync_func is not None:
+ if config.param_sync_func is not None:
param_sync_microbatch_id = microbatch_id + pipeline_parallel_rank
- if param_sync_microbatch_id < num_microbatches and is_first_microbatch_for_model_chunk(param_sync_microbatch_id):
+ if (
+ param_sync_microbatch_id < total_num_microbatches
+ and is_first_microbatch_for_model_chunk(param_sync_microbatch_id)
+ ):
param_sync_chunk_id = get_model_chunk_id(param_sync_microbatch_id, forward=True) + 1
if 1 < param_sync_chunk_id < num_model_chunks:
- param_sync_func(model[param_sync_chunk_id].parameters())
+ config.param_sync_func[param_sync_chunk_id](
+ model[param_sync_chunk_id].parameters()
+ )
# forward step
if parallel_state.is_pipeline_first_stage():
- if len(input_tensors[model_chunk_id]) == \
- len(output_tensors[model_chunk_id]):
+ if len(input_tensors[model_chunk_id]) == len(output_tensors[model_chunk_id]):
input_tensors[model_chunk_id].append(None)
input_tensor = input_tensors[model_chunk_id][-1]
- output_tensor = forward_step(forward_step_func,
- data_iterator[model_chunk_id],
- model[model_chunk_id],
- num_microbatches,
- input_tensor,
- forward_data_store,
- timers,
- collect_non_loss_data,
- dtype,
- enable_autocast)
+
+ output_tensor, num_tokens = forward_step(
+ forward_step_func,
+ data_iterator[model_chunk_id],
+ model[model_chunk_id],
+ num_microbatches,
+ input_tensor,
+ forward_data_store,
+ config,
+ collect_non_loss_data,
+ checkpoint_activations_microbatch,
+ check_first_val_step(
+ first_val_step, forward_only, is_first_microbatch_for_model_chunk(microbatch_id)
+ ),
+ current_microbatch=current_microbatch,
+ )
output_tensors[model_chunk_id].append(output_tensor)
+ nonlocal total_num_tokens
+ total_num_tokens += num_tokens.item()
+
# if forward-only, no need to save tensors for a backward pass
if forward_only:
input_tensors[model_chunk_id].pop()
@@ -579,7 +790,7 @@ def backward_step_helper(microbatch_id):
parallel_state.set_virtual_pipeline_model_parallel_rank(model_chunk_id)
# launch grad synchronization (default)
- if grad_sync_func is None and is_last_microbatch_for_model_chunk(microbatch_id):
+ if config.grad_sync_func is None and is_last_microbatch_for_model_chunk(microbatch_id):
enable_grad_sync()
synchronized_model_chunks.add(model_chunk_id)
@@ -589,26 +800,23 @@ def backward_step_helper(microbatch_id):
input_tensor = input_tensors[model_chunk_id].pop(0)
output_tensor = output_tensors[model_chunk_id].pop(0)
output_tensor_grad = output_tensor_grads[model_chunk_id].pop(0)
- input_tensor_grad = \
- backward_step(grad_scaler,
- input_tensor,
- output_tensor,
- output_tensor_grad,
- model_type,
- timers,
- deallocate_pipeline_outputs)
+ input_tensor_grad = backward_step(
+ input_tensor, output_tensor, output_tensor_grad, model_type, config
+ )
# launch grad synchronization (custom grad sync)
# Note: Asynchronous communication tends to slow down compute.
# To reduce idling from mismatched microbatch times, we launch
# asynchronous communication at the same time across the
# pipeline-parallel group.
- if grad_sync_func is not None:
+ if config.grad_sync_func is not None:
grad_sync_microbatch_id = microbatch_id - pipeline_parallel_rank
- if grad_sync_microbatch_id >= 0 and is_last_microbatch_for_model_chunk(grad_sync_microbatch_id):
+ if grad_sync_microbatch_id >= 0 and is_last_microbatch_for_model_chunk(
+ grad_sync_microbatch_id
+ ):
grad_sync_chunk_id = get_model_chunk_id(grad_sync_microbatch_id, forward=False)
enable_grad_sync()
- grad_sync_func(model[grad_sync_chunk_id].parameters())
+ config.grad_sync_func[grad_sync_chunk_id](model[grad_sync_chunk_id].parameters())
synchronized_model_chunks.add(grad_sync_chunk_id)
disable_grad_sync()
@@ -616,13 +824,33 @@ def backward_step_helper(microbatch_id):
# Run warmup forward passes.
parallel_state.set_virtual_pipeline_model_parallel_rank(0)
- input_tensors[0].append(
- p2p_communication.recv_forward(tensor_shape, dtype, timers=timers))
+ input_tensors[0].append(p2p_communication.recv_forward(tensor_shape, config))
+
+ fwd_wait_handles = None
+ bwd_wait_handles = None
+
for k in range(num_warmup_microbatches):
- output_tensor = forward_step_helper(k)
+
+ if fwd_wait_handles is not None:
+ for req in fwd_wait_handles:
+ req.wait()
+
+ # Decide to checkpoint all layers' activations of the current micro-batch
+ if max_outstanding_backprops is not None:
+ checkpoint_activations_microbatch = (
+ k % max_outstanding_backprops
+ >= config.num_microbatches_with_partial_activation_checkpoints
+ )
+ else:
+ checkpoint_activations_microbatch = None
+
+ current_microbatch = get_microbatch_id_in_model_chunk(k, forward=True)
+ output_tensor = forward_step_helper(
+ k, current_microbatch, checkpoint_activations_microbatch
+ )
# Determine if tensor should be received from previous stage.
- next_forward_model_chunk_id = get_model_chunk_id(k+1, forward=True)
+ next_forward_model_chunk_id = get_model_chunk_id(k + 1, forward=True)
recv_prev = True
if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
if next_forward_model_chunk_id == 0:
@@ -636,108 +864,257 @@ def backward_step_helper(microbatch_id):
# Send and receive tensors as appropriate (send tensors computed
# in this iteration; receive tensors for next iteration).
- if k == (num_warmup_microbatches - 1) and not forward_only and \
- not all_warmup_microbatches:
- input_tensor_grad = None
- recv_next = True
- if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
- recv_next = False
- input_tensor, output_tensor_grad = \
- p2p_communication.send_forward_backward_recv_forward_backward(
- output_tensor, input_tensor_grad,
- recv_prev=recv_prev, recv_next=recv_next,
- tensor_shape=tensor_shape, dtype=dtype,
- timers=timers)
- output_tensor_grads[num_model_chunks-1].append(output_tensor_grad)
+ if not config.overlap_p2p_comm:
+ if (
+ k == (num_warmup_microbatches - 1)
+ and not forward_only
+ and not all_warmup_microbatches
+ ):
+ input_tensor_grad = None
+ recv_next = True
+ if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+ recv_next = False
+ (input_tensor, output_tensor_grad) = (
+ p2p_communication.send_forward_backward_recv_forward_backward(
+ output_tensor,
+ input_tensor_grad,
+ recv_prev=recv_prev,
+ recv_next=recv_next,
+ tensor_shape=tensor_shape,
+ config=config,
+ )
+ )
+ output_tensor_grads[num_model_chunks - 1].append(output_tensor_grad)
+ else:
+ input_tensor = p2p_communication.send_forward_recv_forward(
+ output_tensor, recv_prev=recv_prev, tensor_shape=tensor_shape, config=config
+ )
+ input_tensors[next_forward_model_chunk_id].append(input_tensor)
else:
- input_tensor = \
- p2p_communication.send_forward_recv_forward(
- output_tensor, recv_prev=recv_prev,
- tensor_shape=tensor_shape, dtype=dtype,
- timers=timers)
- input_tensors[next_forward_model_chunk_id].append(input_tensor)
- deallocate_output_tensor(output_tensor, deallocate_pipeline_outputs)
+ input_tensor, fwd_wait_handles = p2p_communication.send_forward_recv_forward(
+ output_tensor,
+ recv_prev=recv_prev,
+ tensor_shape=tensor_shape,
+ config=config,
+ overlap_p2p_comm=True,
+ )
+
+ if (
+ k == (num_warmup_microbatches - 1)
+ and not forward_only
+ and not all_warmup_microbatches
+ ):
+ input_tensor_grad = None
+ recv_next = True
+ if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+ recv_next = False
+
+ (output_tensor_grad, bwd_wait_handles) = (
+ p2p_communication.send_backward_recv_backward(
+ input_tensor_grad,
+ recv_next=recv_next,
+ tensor_shape=tensor_shape,
+ config=config,
+ overlap_p2p_comm=True,
+ )
+ )
+
+ output_tensor_grads[num_model_chunks - 1].append(output_tensor_grad)
+ input_tensors[next_forward_model_chunk_id].append(input_tensor)
+
+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
# Run 1F1B in steady state.
for k in range(num_microbatches_remaining):
# Forward pass.
forward_k = k + num_warmup_microbatches
- output_tensor = forward_step_helper(forward_k)
- # Backward pass.
- backward_k = k
- input_tensor_grad = backward_step_helper(backward_k)
+ # Decide to checkpoint all layers' activations of the current micro-batch
+ if max_outstanding_backprops is not None:
+ checkpoint_activations_microbatch = (
+ forward_k % max_outstanding_backprops
+ >= config.num_microbatches_with_partial_activation_checkpoints
+ )
+ else:
+ checkpoint_activations_microbatch = None
+
+ current_microbatch = get_microbatch_id_in_model_chunk(forward_k, forward=True)
+ if config.overlap_p2p_comm:
+ if fwd_wait_handles is not None:
+ for req in fwd_wait_handles:
+ req.wait()
+
+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
+
+ output_tensor = forward_step_helper(
+ forward_k, current_microbatch, checkpoint_activations_microbatch
+ )
+
+ # Determine if current stage has anything to send in either direction,
+ # otherwise set tensor to None.
+ forward_model_chunk_id = get_model_chunk_id(forward_k, forward=True)
+ parallel_state.set_virtual_pipeline_model_parallel_rank(forward_model_chunk_id)
+
+ # Last virtual stage no activation tensor to send
+ if parallel_state.is_pipeline_last_stage():
+ output_tensor = None
+
+ # Determine if peers are sending, and where in data structure to put
+ # received tensors.
+ recv_prev = True
+ if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
+ # First stage is ahead of last stage by (pipeline_parallel_size - 1).
+ next_forward_model_chunk_id = get_model_chunk_id(
+ forward_k - (pipeline_parallel_size - 1), forward=True
+ )
+ if next_forward_model_chunk_id == (num_model_chunks - 1):
+ recv_prev = False
+ next_forward_model_chunk_id += 1
+ else:
+ next_forward_model_chunk_id = get_model_chunk_id(forward_k + 1, forward=True)
- # Send output_tensor and input_tensor_grad, receive input_tensor
- # and output_tensor_grad.
+ # If last iteration, don't receive; we already received one extra
+ # before the start of the for loop.
+ if k == (num_microbatches_remaining - 1):
+ recv_prev = False
- # Determine if current stage has anything to send in either direction,
- # otherwise set tensor to None.
- forward_model_chunk_id = get_model_chunk_id(forward_k, forward=True)
- parallel_state.set_virtual_pipeline_model_parallel_rank(forward_model_chunk_id)
- if parallel_state.is_pipeline_last_stage():
- output_tensor = None
+ # Send activation tensor to the next stage and receive activation tensor from the
+ # previous stage
+ input_tensor, fwd_wait_handles = p2p_communication.send_forward_recv_forward(
+ output_tensor,
+ recv_prev=recv_prev,
+ tensor_shape=tensor_shape,
+ config=config,
+ overlap_p2p_comm=True,
+ )
+ # assert fwd_wait_handles is not None
+
+ if bwd_wait_handles is not None:
+ for req in bwd_wait_handles:
+ req.wait()
+
+ # Backward pass.
+ backward_k = k
+ input_tensor_grad = backward_step_helper(backward_k)
+
+ backward_model_chunk_id = get_model_chunk_id(backward_k, forward=False)
+ parallel_state.set_virtual_pipeline_model_parallel_rank(backward_model_chunk_id)
+
+ # First virtual stage no activation gradient tensor to send
+ if parallel_state.is_pipeline_first_stage():
+ input_tensor_grad = None
+
+ # Determine if the current virtual stage has an activation gradient tensor to receive
+ recv_next = True
+ if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+ # Last stage is ahead of first stage by (pipeline_parallel_size - 1).
+ next_backward_model_chunk_id = get_model_chunk_id(
+ backward_k - (pipeline_parallel_size - 1), forward=False
+ )
+ if next_backward_model_chunk_id == 0:
+ recv_next = False
+ next_backward_model_chunk_id -= 1
+ else:
+ next_backward_model_chunk_id = get_model_chunk_id(backward_k + 1, forward=False)
+
+ output_tensor_grad, bwd_wait_handles = p2p_communication.send_backward_recv_backward(
+ input_tensor_grad,
+ recv_next=recv_next,
+ tensor_shape=tensor_shape,
+ config=config,
+ overlap_p2p_comm=True,
+ )
+
+ else: # no p2p overlap
+ output_tensor = forward_step_helper(
+ forward_k, current_microbatch, checkpoint_activations_microbatch
+ )
+
+ # Backward pass.
+ backward_k = k
+ input_tensor_grad = backward_step_helper(backward_k)
+
+ # Send output_tensor and input_tensor_grad, receive input_tensor
+ # and output_tensor_grad.
+
+ # Determine if current stage has anything to send in either direction,
+ # otherwise set tensor to None.
+ forward_model_chunk_id = get_model_chunk_id(forward_k, forward=True)
+ parallel_state.set_virtual_pipeline_model_parallel_rank(forward_model_chunk_id)
+ if parallel_state.is_pipeline_last_stage():
+ output_tensor = None
+
+ backward_model_chunk_id = get_model_chunk_id(backward_k, forward=False)
+ parallel_state.set_virtual_pipeline_model_parallel_rank(backward_model_chunk_id)
+ if parallel_state.is_pipeline_first_stage():
+ input_tensor_grad = None
+
+ # Determine if peers are sending, and where in data structure to put
+ # received tensors.
+ recv_prev = True
+ if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
+ # First stage is ahead of last stage by (pipeline_parallel_size - 1).
+ next_forward_model_chunk_id = get_model_chunk_id(
+ forward_k - (pipeline_parallel_size - 1), forward=True
+ )
+ if next_forward_model_chunk_id == (num_model_chunks - 1):
+ recv_prev = False
+ next_forward_model_chunk_id += 1
+ else:
+ next_forward_model_chunk_id = get_model_chunk_id(forward_k + 1, forward=True)
- backward_model_chunk_id = get_model_chunk_id(backward_k, forward=False)
- parallel_state.set_virtual_pipeline_model_parallel_rank(backward_model_chunk_id)
- if parallel_state.is_pipeline_first_stage():
- input_tensor_grad = None
+ recv_next = True
+ if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+ # Last stage is ahead of first stage by (pipeline_parallel_size - 1).
+ next_backward_model_chunk_id = get_model_chunk_id(
+ backward_k - (pipeline_parallel_size - 1), forward=False
+ )
+ if next_backward_model_chunk_id == 0:
+ recv_next = False
+ next_backward_model_chunk_id -= 1
+ else:
+ next_backward_model_chunk_id = get_model_chunk_id(backward_k + 1, forward=False)
- # Determine if peers are sending, and where in data structure to put
- # received tensors.
- recv_prev = True
- if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
- # First stage is ahead of last stage by (pipeline_parallel_size - 1).
- next_forward_model_chunk_id = get_model_chunk_id(
- forward_k - (pipeline_parallel_size - 1), forward=True)
- if next_forward_model_chunk_id == (num_model_chunks - 1):
+ # If last iteration, don't receive; we already received one extra
+ # before the start of the for loop.
+ if k == (num_microbatches_remaining - 1):
recv_prev = False
- next_forward_model_chunk_id += 1
- else:
- next_forward_model_chunk_id = get_model_chunk_id(forward_k + 1,
- forward=True)
-
- recv_next = True
- if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
- # Last stage is ahead of first stage by (pipeline_parallel_size - 1).
- next_backward_model_chunk_id = get_model_chunk_id(
- backward_k - (pipeline_parallel_size - 1), forward=False)
- if next_backward_model_chunk_id == 0:
- recv_next = False
- next_backward_model_chunk_id -= 1
- else:
- next_backward_model_chunk_id = get_model_chunk_id(backward_k + 1,
- forward=False)
-
- # If last iteration, don't receive; we already received one extra
- # before the start of the for loop.
- if k == (num_microbatches_remaining - 1):
- recv_prev = False
- # Communicate tensors.
- input_tensor, output_tensor_grad = \
- p2p_communication.send_forward_backward_recv_forward_backward(
- output_tensor, input_tensor_grad,
- recv_prev=recv_prev, recv_next=recv_next,
- tensor_shape=tensor_shape, dtype=dtype, timers=timers)
- deallocate_output_tensor(output_tensor, deallocate_pipeline_outputs)
+ # Communicate tensors.
+ (input_tensor, output_tensor_grad) = (
+ p2p_communication.send_forward_backward_recv_forward_backward(
+ output_tensor,
+ input_tensor_grad,
+ recv_prev=recv_prev,
+ recv_next=recv_next,
+ tensor_shape=tensor_shape,
+ config=config,
+ )
+ )
+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
# Put input_tensor and output_tensor_grad in data structures in the
# right location.
if recv_prev:
input_tensors[next_forward_model_chunk_id].append(input_tensor)
if recv_next:
- output_tensor_grads[next_backward_model_chunk_id].append(
- output_tensor_grad)
+ output_tensor_grads[next_backward_model_chunk_id].append(output_tensor_grad)
+
+ deallocate_output_tensor(output_tensor, config.deallocate_pipeline_outputs)
# Run cooldown backward passes (flush out pipeline).
if not forward_only:
+ if config.overlap_p2p_comm and bwd_wait_handles is not None:
+ for wait_handle in bwd_wait_handles:
+ wait_handle.wait()
+
if all_warmup_microbatches:
- output_tensor_grads[num_model_chunks-1].append(
- p2p_communication.recv_backward(tensor_shape, dtype=dtype, timers=timers))
+ output_tensor_grads[num_model_chunks - 1].append(
+ p2p_communication.recv_backward(tensor_shape, config=config)
+ )
for k in range(num_microbatches_remaining, total_num_microbatches):
input_tensor_grad = backward_step_helper(k)
- next_backward_model_chunk_id = get_model_chunk_id(k+1, forward=False)
+ next_backward_model_chunk_id = get_model_chunk_id(k + 1, forward=False)
recv_next = True
if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
if next_backward_model_chunk_id == (num_model_chunks - 1):
@@ -746,211 +1123,265 @@ def backward_step_helper(microbatch_id):
recv_next = False
output_tensor_grads[next_backward_model_chunk_id].append(
p2p_communication.send_backward_recv_backward(
- input_tensor_grad, recv_next=recv_next,
- tensor_shape=tensor_shape, dtype=dtype,
- timers=timers))
-
- # Launch any remaining grad reductions
- enable_grad_sync()
- if grad_sync_func is not None:
- params = []
- for model_chunk_id in range(num_model_chunks):
- if model_chunk_id not in synchronized_model_chunks:
- params.extend(model[model_chunk_id].parameters())
- synchronized_model_chunks.add(model_chunk_id)
- if params:
- grad_sync_func(params)
+ input_tensor_grad, recv_next=recv_next, tensor_shape=tensor_shape, config=config
+ )
+ )
+
+ # Launch any remaining grad reductions.
+ enable_grad_sync()
+ if config.grad_sync_func is not None:
+ for model_chunk_id in range(num_model_chunks):
+ if model_chunk_id not in synchronized_model_chunks:
+ config.grad_sync_func[model_chunk_id](model[model_chunk_id].parameters())
+ synchronized_model_chunks.add(model_chunk_id)
+
+ if config.finalize_model_grads_func is not None and not forward_only:
+
+ # If defer_embedding_wgrad_compute is enabled we need to do the
+ # weight gradient GEMM's here.
+ finish_embedding_wgrad_compute(config, embedding_module)
+
+ # Finalize model grads (perform full grad all-reduce / reduce-scatter for
+ # data parallelism, layernorm all-reduce for sequence parallelism, and
+ # embedding all-reduce for pipeline parallelism).
+ config.finalize_model_grads_func(
+ model, total_num_tokens if config.calculate_per_token_loss else None
+ )
+
+ # Restore config.grad_sync_func and config.param_sync_func.
+ if forward_only:
+ config.grad_sync_func, config.param_sync_func = grad_sync_func, param_sync_func
+
+ if config.timers is not None:
+ config.timers('forward-backward').stop()
return forward_data_store
-def get_tensor_shapes(*,
- rank: int,
- model_type: ModelType,
- tensor_shape: Shape,
- decoder_seq_length: int,
- sequence_parallel: bool):
- # Determine right tensor sizes (based on position of rank with respect to split
- # rank) and model size.
- # Send two tensors if model is T5 and rank is in decoder stage:
- # first tensor is decoder (pre-transpose),
- # second tensor is encoder (post-transpose).
- # If model is T5 and rank is at the boundary:
- # send one tensor (post-transpose from encoder).
- # Otherwise, send one tensor (pre-transpose).
- tensor_shapes = []
- assert (
- len(tensor_shape) == 3
- ), f"`tensor_shape` should be [sequence_length, micro_batch_size, hidden_size] but {tensor_shape}"
+def get_tensor_shapes(
+ *,
+ rank: int,
+ model_type: ModelType,
+ seq_length: int,
+ micro_batch_size: int,
+ decoder_seq_length: int,
+ config,
+ encoder_decoder_xattn: bool,
+):
+ # Determine right tensor sizes (based on position of rank with
+ # respect to split rank) and model size.
+ # Send two tensors if model decoder requires the encoder's output
+ # (via cross-attention) and rank is in decoder stage.
+ # first tensor is decoder.
+ # second tensor is encoder.
+ # If model has an encoder & decoder and rank is at the boundary:
+ # send one tensor.
+ # Otherwise, send one tensor.
+ tensor_shapes = []
- seq_length, micro_batch_size, hidden_size = tensor_shape
+ seq_length = seq_length // parallel_state.get_context_parallel_world_size()
+ if model_type == ModelType.encoder_and_decoder:
+ decoder_seq_length = decoder_seq_length // parallel_state.get_context_parallel_world_size()
- if sequence_parallel:
+ if config.sequence_parallel:
seq_length = seq_length // parallel_state.get_tensor_model_parallel_world_size()
+ if model_type == ModelType.encoder_and_decoder:
+ decoder_seq_length = (
+ decoder_seq_length // parallel_state.get_tensor_model_parallel_world_size()
+ )
if model_type == ModelType.encoder_and_decoder:
- if sequence_parallel:
- decoder_seq_length = decoder_seq_length // parallel_state.get_tensor_model_parallel_world_size()
-
- if parallel_state.is_pipeline_stage_before_split(rank):
- tensor_shapes.append((seq_length, micro_batch_size, hidden_size))
+ if parallel_state.is_inside_encoder(rank):
+ tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
+ elif encoder_decoder_xattn:
+ tensor_shapes.append((decoder_seq_length, micro_batch_size, config.hidden_size))
+ tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
else:
- tensor_shapes.append((decoder_seq_length, micro_batch_size, hidden_size))
- tensor_shapes.append((seq_length, micro_batch_size, hidden_size))
- else:
- tensor_shapes.append((seq_length, micro_batch_size, hidden_size))
+ tensor_shapes.append((decoder_seq_length, micro_batch_size, config.hidden_size))
+ else: # model_type == ModelType.encoder_or_decoder
+ tensor_shapes.append((seq_length, micro_batch_size, config.hidden_size))
return tensor_shapes
-
-def recv_forward(tensor_shapes, dtype, timers):
+def recv_forward(tensor_shapes, config):
input_tensors = []
for tensor_shape in tensor_shapes:
if tensor_shape is None:
input_tensors.append(None)
else:
- input_tensors.append(p2p_communication.recv_forward(tensor_shape, dtype,
- timers=timers))
+ input_tensors.append(p2p_communication.recv_forward(tensor_shape, config))
return input_tensors
-def recv_backward(tensor_shapes, dtype, timers):
+def recv_backward(tensor_shapes, config):
output_tensor_grads = []
for tensor_shape in tensor_shapes:
if tensor_shape is None:
output_tensor_grads.append(None)
else:
- output_tensor_grads.append(p2p_communication.recv_backward(tensor_shape, dtype,
- timers=timers))
+ output_tensor_grads.append(p2p_communication.recv_backward(tensor_shape, config))
return output_tensor_grads
-def send_forward(output_tensors, tensor_shapes, timers):
+def send_forward(output_tensors, tensor_shapes, config):
if not isinstance(output_tensors, list):
output_tensors = [output_tensors]
- for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
+ for output_tensor, tensor_shape in zip(output_tensors, tensor_shapes):
if tensor_shape is None:
continue
- p2p_communication.send_forward(output_tensor, timers=timers)
+ p2p_communication.send_forward(output_tensor, config)
-def send_backward(input_tensor_grads, tensor_shapes, timers):
+def send_backward(input_tensor_grads, tensor_shapes, config):
if not isinstance(input_tensor_grads, list):
input_tensor_grads = [input_tensor_grads]
- for (input_tensor_grad, tensor_shape) in zip(input_tensor_grads, tensor_shapes):
+ for input_tensor_grad, tensor_shape in zip(input_tensor_grads, tensor_shapes):
if tensor_shape is None:
continue
- p2p_communication.send_backward(input_tensor_grad, timers=timers)
+ p2p_communication.send_backward(input_tensor_grad, config)
-def send_forward_recv_backward(output_tensors, tensor_shapes, dtype, timers):
+def send_forward_recv_backward(output_tensors, tensor_shapes, config):
if not isinstance(output_tensors, list):
output_tensors = [output_tensors]
output_tensor_grads = []
- for (output_tensor, tensor_shape) in zip(output_tensors, tensor_shapes):
+ for output_tensor, tensor_shape in zip(output_tensors, tensor_shapes):
if tensor_shape is None:
output_tensor_grads.append(None)
continue
output_tensor_grad = p2p_communication.send_forward_recv_backward(
- output_tensor, tensor_shape, dtype, timers=timers)
+ output_tensor, tensor_shape, config
+ )
output_tensor_grads.append(output_tensor_grad)
return output_tensor_grads
-def send_backward_recv_forward(input_tensor_grads, tensor_shapes, dtype, timers):
+def send_backward_recv_forward(input_tensor_grads, tensor_shapes, config):
if not isinstance(input_tensor_grads, list):
input_tensor_grads = [input_tensor_grads]
input_tensors = []
- for (input_tensor_grad, tensor_shape) in zip(input_tensor_grads, tensor_shapes):
+ for input_tensor_grad, tensor_shape in zip(input_tensor_grads, tensor_shapes):
if tensor_shape is None:
input_tensors.append(None)
continue
input_tensor = p2p_communication.send_backward_recv_forward(
- input_tensor_grad, tensor_shape, dtype, timers=timers)
+ input_tensor_grad, tensor_shape, config
+ )
input_tensors.append(input_tensor)
return input_tensors
-def forward_backward_pipelining_without_interleaving(*,
- forward_step_func,
- data_iterator: Union[Iterator, List[Iterator]],
- model: Union[torch.nn.Module, List[torch.nn.Module]],
- num_microbatches: int,
- dtype: torch.dtype,
- tensor_shape: Shape,
- decoder_seq_length: Optional[int] = None,
- grad_scaler: Callable = None,
- sequence_parallel: bool = False,
- forward_only: bool = False,
- timers: Callable = None,
- collect_non_loss_data: bool = False,
- enable_autocast: bool = False,
- deallocate_pipeline_outputs: bool = False,
- no_sync_func: Optional[Callable] = None,
- grad_sync_func: Optional[Callable] = None,
- param_sync_func: Optional[Callable] = None, # unused
- ):
+def forward_backward_pipelining_without_interleaving(
+ *,
+ forward_step_func,
+ data_iterator: Union[Iterator, List[Iterator]],
+ model: Union[torch.nn.Module, List[torch.nn.Module]],
+ num_microbatches: int,
+ seq_length: int,
+ micro_batch_size: int,
+ decoder_seq_length: int = None,
+ forward_only: bool = False,
+ collect_non_loss_data: bool = False,
+ first_val_step: bool = None,
+):
"""Run non-interleaved 1F1B schedule, with communication between pipeline
- stages.
-
- Returns dictionary with losses if the last stage, empty dict otherwise."""
+ stages. Returns dictionary with losses if the last stage, empty dict otherwise."""
if isinstance(model, list):
- assert len(model) == 1, \
- "non-interleaved pipeline parallelism does not support model chunking"
+ assert (
+ len(model) == 1
+ ), "non-interleaved pipeline parallelism does not support model chunking"
model = model[0]
if isinstance(data_iterator, list):
- assert len(data_iterator) == 1, \
- "non-pipeline-parallel schedule does not support model chunking"
+ assert (
+ len(data_iterator) == 1
+ ), "non-pipeline-parallel schedule does not support model chunking"
data_iterator = data_iterator[0]
+ config = get_model_config(model)
+ if config.overlap_p2p_comm:
+ raise ValueError(
+ "Non-interleaved pipeline parallelism does not support overlapping p2p communication"
+ )
+
+ # Needed only when gradients are finalized in M-Core
+ if config.finalize_model_grads_func is not None and not forward_only:
+ embedding_module = clear_embedding_activation_buffer(config, model)
+
+ if config.timers is not None:
+ config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)
+
# Disable async grad reductions
- if no_sync_func is None and isinstance(model, torchDDP):
- no_sync_func = model.no_sync
+ no_sync_func = config.no_sync_func
if no_sync_func is None:
no_sync_func = contextlib.nullcontext
no_sync_context = None
+
def disable_grad_sync():
"""Disable asynchronous grad reductions"""
nonlocal no_sync_context
if no_sync_context is None:
no_sync_context = no_sync_func()
no_sync_context.__enter__()
+
def enable_grad_sync():
"""Enable asynchronous grad reductions"""
nonlocal no_sync_context
if no_sync_context is not None:
no_sync_context.__exit__(None, None, None)
no_sync_context = None
+
disable_grad_sync()
# Compute number of warmup microbatches.
- num_warmup_microbatches = \
- (parallel_state.get_pipeline_model_parallel_world_size() -
- parallel_state.get_pipeline_model_parallel_rank() - 1)
- num_warmup_microbatches = min(
- num_warmup_microbatches,
- num_microbatches)
- num_microbatches_remaining = \
- num_microbatches - num_warmup_microbatches
+ num_warmup_microbatches = (
+ parallel_state.get_pipeline_model_parallel_world_size()
+ - parallel_state.get_pipeline_model_parallel_rank()
+ - 1
+ )
+ num_warmup_microbatches = min(num_warmup_microbatches, num_microbatches)
+ num_microbatches_remaining = num_microbatches - num_warmup_microbatches
+
+ # Checkpoint the activations of partial Transformer layers in a number of micro-batches
+ # within the maximum outstanding micro-batch backpropagations.
+ # Micro-batches with the ids less than 'num_microbatches_with_partial_activation_checkpoints'
+ # checkpoint partial Transformer layers (or skip checkpointing) and
+ # the rest of micro-batches within a window of micro-batches checkpoint
+ # all Transformer layers. The window of micro-batches is set by the maximum
+ # outstanding backpropagations and becomes smaller at later pipeline stages.
+ # Please refer the appendix C in https://arxiv.org/pdf/2205.05198.pdf
+ max_outstanding_backprops = None
+ if config.num_microbatches_with_partial_activation_checkpoints is not None:
+ max_outstanding_backprops = num_warmup_microbatches + 1
model_type = get_model_type(model)
+ encoder_decoder_xattn = get_model_xattn(model)
rank = parallel_state.get_pipeline_model_parallel_rank()
- recv_tensor_shapes = get_tensor_shapes(rank=rank-1,
- model_type=model_type,
- tensor_shape=tensor_shape,
- decoder_seq_length=decoder_seq_length,
- sequence_parallel=sequence_parallel)
- send_tensor_shapes = get_tensor_shapes(rank=rank,
- model_type=model_type,
- tensor_shape=tensor_shape,
- decoder_seq_length=decoder_seq_length,
- sequence_parallel=sequence_parallel)
+ recv_tensor_shapes = get_tensor_shapes(
+ rank=rank - 1,
+ model_type=model_type,
+ seq_length=seq_length,
+ micro_batch_size=micro_batch_size,
+ decoder_seq_length=decoder_seq_length,
+ config=config,
+ encoder_decoder_xattn=encoder_decoder_xattn,
+ )
+ send_tensor_shapes = get_tensor_shapes(
+ rank=rank,
+ model_type=model_type,
+ seq_length=seq_length,
+ micro_batch_size=micro_batch_size,
+ decoder_seq_length=decoder_seq_length,
+ config=config,
+ encoder_decoder_xattn=encoder_decoder_xattn,
+ )
# Input, output tensors only need to be saved when doing backward passes
input_tensors = None
output_tensors = None
+ total_num_tokens = torch.tensor(0, dtype=torch.int).cuda()
+
if not forward_only:
input_tensors = []
output_tensors = []
@@ -958,64 +1389,112 @@ def enable_grad_sync():
# Run warmup forward passes.
for i in range(num_warmup_microbatches):
- input_tensor = recv_forward(recv_tensor_shapes, dtype, timers=timers)
- output_tensor = forward_step(forward_step_func, data_iterator, model, num_microbatches,
- input_tensor, forward_data_store,
- timers, collect_non_loss_data, dtype, enable_autocast)
- send_forward(output_tensor, send_tensor_shapes, timers=timers)
+ # Decide to checkpoint all layers' activations of the current micro-batch
+ if max_outstanding_backprops is not None:
+ checkpoint_activations_microbatch = (
+ i % max_outstanding_backprops
+ >= config.num_microbatches_with_partial_activation_checkpoints
+ )
+ else:
+ checkpoint_activations_microbatch = None
+
+ input_tensor = recv_forward(recv_tensor_shapes, config)
+ output_tensor, num_tokens = forward_step(
+ forward_step_func,
+ data_iterator,
+ model,
+ num_microbatches,
+ input_tensor,
+ forward_data_store,
+ config,
+ collect_non_loss_data,
+ checkpoint_activations_microbatch,
+ check_first_val_step(first_val_step, forward_only, i == 0),
+ current_microbatch=i,
+ encoder_decoder_xattn=encoder_decoder_xattn,
+ )
+ send_forward(output_tensor, send_tensor_shapes, config)
+ total_num_tokens += num_tokens.item()
if not forward_only:
input_tensors.append(input_tensor)
output_tensors.append(output_tensor)
- deallocate_output_tensor(output_tensor[0], deallocate_pipeline_outputs)
+ deallocate_output_tensor(output_tensor[0], config.deallocate_pipeline_outputs)
# Before running 1F1B, need to receive first forward tensor.
# If all microbatches are run in warmup / cooldown phase, then no need to
# receive this tensor here.
if num_microbatches_remaining > 0:
- input_tensor = recv_forward(recv_tensor_shapes, dtype, timers=timers)
+ input_tensor = recv_forward(recv_tensor_shapes, config)
# Run 1F1B in steady state.
for i in range(num_microbatches_remaining):
- last_iteration = (i == (num_microbatches_remaining - 1))
+ last_iteration = i == (num_microbatches_remaining - 1)
- output_tensor = forward_step(forward_step_func, data_iterator, model, num_microbatches,
- input_tensor, forward_data_store,
- timers, collect_non_loss_data, dtype, enable_autocast)
+ # Decide to checkpoint all layers' activations of the current micro-batch
+ if max_outstanding_backprops is not None:
+ checkpoint_activations_microbatch = (
+ (i + num_warmup_microbatches) % max_outstanding_backprops
+ ) >= config.num_microbatches_with_partial_activation_checkpoints
+ else:
+ checkpoint_activations_microbatch = None
+
+ output_tensor, num_tokens = forward_step(
+ forward_step_func,
+ data_iterator,
+ model,
+ num_microbatches,
+ input_tensor,
+ forward_data_store,
+ config,
+ collect_non_loss_data,
+ checkpoint_activations_microbatch,
+ check_first_val_step(
+ first_val_step, forward_only, (i == 0) and (num_warmup_microbatches == 0)
+ ),
+ current_microbatch=i + num_warmup_microbatches,
+ encoder_decoder_xattn=encoder_decoder_xattn,
+ )
+ total_num_tokens += num_tokens.item()
if forward_only:
- send_forward(output_tensor, send_tensor_shapes, timers=timers)
+ send_forward(output_tensor, send_tensor_shapes, config)
if not last_iteration:
- input_tensor = recv_forward(recv_tensor_shapes, dtype, timers=timers)
+ input_tensor = recv_forward(recv_tensor_shapes, config)
else:
- output_tensor_grad = \
- send_forward_recv_backward(output_tensor,
- send_tensor_shapes, dtype,
- timers=timers)
+ output_tensor_grad = send_forward_recv_backward(
+ output_tensor, send_tensor_shapes, config
+ )
# Add input_tensor and output_tensor to end of list.
input_tensors.append(input_tensor)
output_tensors.append(output_tensor)
- deallocate_output_tensor(output_tensor[0], deallocate_pipeline_outputs)
+ deallocate_output_tensor(output_tensor[0], config.deallocate_pipeline_outputs)
# Pop input_tensor and output_tensor from the start of the list for
# the backward pass.
input_tensor = input_tensors.pop(0)
output_tensor = output_tensors.pop(0)
- input_tensor_grad = \
- backward_step(grad_scaler, input_tensor, output_tensor,
- output_tensor_grad, model_type, timers, deallocate_pipeline_outputs)
+ # Enable grad sync for the last microbatch in the batch if the full
+ # backward pass completes in the 1F1B stage.
+ if num_warmup_microbatches == 0 and last_iteration:
+ if config.grad_sync_func is None or rank == 0:
+ enable_grad_sync()
+
+ input_tensor_grad = backward_step(
+ input_tensor, output_tensor, output_tensor_grad, model_type, config
+ )
if last_iteration:
input_tensor = None
- send_backward(input_tensor_grad, recv_tensor_shapes, timers=timers)
+ send_backward(input_tensor_grad, recv_tensor_shapes, config)
else:
- input_tensor = \
- send_backward_recv_forward(
- input_tensor_grad, recv_tensor_shapes, dtype, timers=timers)
+ input_tensor = send_backward_recv_forward(
+ input_tensor_grad, recv_tensor_shapes, config
+ )
# Run cooldown backward passes.
if not forward_only:
@@ -1026,25 +1505,41 @@ def enable_grad_sync():
# async grad reduction in first pipeline stage. Other
# pipeline stages do grad reduction during pipeline
# bubble.
- if i == num_warmup_microbatches-1:
- if grad_sync_func is None or rank == 0:
+ if i == num_warmup_microbatches - 1:
+ if config.grad_sync_func is None or rank == 0:
enable_grad_sync()
input_tensor = input_tensors.pop(0)
output_tensor = output_tensors.pop(0)
- output_tensor_grad = recv_backward(send_tensor_shapes, dtype, timers=timers)
+ output_tensor_grad = recv_backward(send_tensor_shapes, config)
- input_tensor_grad = \
- backward_step(grad_scaler, input_tensor, output_tensor,
- output_tensor_grad, model_type, timers, deallocate_pipeline_outputs)
+ input_tensor_grad = backward_step(
+ input_tensor, output_tensor, output_tensor_grad, model_type, config
+ )
- send_backward(input_tensor_grad, recv_tensor_shapes, timers=timers)
+ send_backward(input_tensor_grad, recv_tensor_shapes, config)
- # Launch any remaining grad reductions
- if no_sync_context is not None:
- enable_grad_sync()
- if grad_sync_func is not None:
- grad_sync_func(model.parameters())
+ # Launch any remaining grad reductions.
+ if no_sync_context is not None:
+ enable_grad_sync()
+ if config.grad_sync_func is not None:
+ config.grad_sync_func(model.parameters())
+
+ if config.finalize_model_grads_func is not None and not forward_only:
+
+ # If defer_embedding_wgrad_compute is enabled we need to do the
+ # weight gradient GEMM's here.
+ finish_embedding_wgrad_compute(config, embedding_module)
+
+ # Finalize model grads (perform full grad all-reduce / reduce-scatter for
+ # data parallelism, layernorm all-reduce for sequence parallelism, and
+ # embedding all-reduce for pipeline parallelism).
+ config.finalize_model_grads_func(
+ [model], total_num_tokens if config.calculate_per_token_loss else None
+ )
+
+ if config.timers is not None:
+ config.timers('forward-backward').stop()
return forward_data_store
diff --git a/megatron/core/requirements.txt b/megatron/core/requirements.txt
index 08ed5eeb4b..a03ef133e7 100644
--- a/megatron/core/requirements.txt
+++ b/megatron/core/requirements.txt
@@ -1 +1,2 @@
-torch
\ No newline at end of file
+torch
+packaging
diff --git a/megatron/core/ssm/__init__.py b/megatron/core/ssm/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py
new file mode 100644
index 0000000000..20754b5c25
--- /dev/null
+++ b/megatron/core/ssm/mamba_block.py
@@ -0,0 +1,336 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2024, Tri Dao, Albert Gu.
+
+# Some of this code was adopted from https://github.com/state-spaces/mamba/
+# This source code is licensed under the Apache license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import Union
+
+from torch import Tensor, nn
+
+from megatron.core import parallel_state
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
+from megatron.core.extensions.transformer_engine import TENorm
+from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols as LayerSymbols
+from megatron.core.ssm.mamba_hybrid_layer_allocation import allocate_layers
+from megatron.core.tensor_parallel import get_cuda_rng_tracker
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import sharded_state_dict_default
+from megatron.core.utils import make_viewless_tensor
+
+
+# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
+def _init_weights(
+ module,
+ n_layer,
+ initializer_range=0.02, # Now only used for embedding layer.
+ rescale_prenorm_residual=True,
+ n_residuals_per_layer=1, # Change to 2 if we have MLP
+):
+ with get_cuda_rng_tracker().fork():
+ if isinstance(module, nn.Linear):
+ if not getattr(module.weight, "_no_reinit", False):
+ nn.init.normal_(module.weight, std=initializer_range)
+ if module.bias is not None:
+ if not getattr(module.bias, "_no_reinit", False):
+ nn.init.zeros_(module.bias)
+ elif isinstance(module, nn.Embedding):
+ nn.init.normal_(module.weight, std=initializer_range)
+
+ for name, p in module.named_parameters():
+ if name in ["in_proj.weight", "x_proj.weight", "conv1d.weight", "out_proj.weight"]:
+ nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+
+ if rescale_prenorm_residual:
+ # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+ # > A modified initialization which accounts for the accumulation on the
+ # > residual path with model depth. Scale
+ # > the weights of residual layers at initialization by a factor of
+ # > 1/√N where N is the # of residual layers.
+ # > -- GPT-2 :: https://openai.com/blog/better-language-models/
+ #
+ # Reference (Megatron-LM):
+ # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+ for name, p in module.named_parameters():
+ if name in ["out_proj.weight", "fc2.weight"]:
+ # Special Scaled Initialization
+ nn.init.normal_(
+ p,
+ mean=0.0,
+ std=initializer_range / math.sqrt(n_residuals_per_layer * n_layer),
+ )
+
+
+@dataclass
+class MambaStackSubmodules:
+ """
+ A class for the module specs for the MambaStack.
+ """
+
+ mamba_layer: Union[ModuleSpec, type] = IdentityOp
+ attention_layer: Union[ModuleSpec, type] = IdentityOp
+ mlp_layer: Union[ModuleSpec, type] = IdentityOp
+
+
+class MambaStack(MegatronModule):
+ """
+ Constructor for the MambaStack class.
+
+ Args:
+ config (TransformerConfig): the transformer configuration
+ submodules (MambaStackSubmodules): the submodules for the stack
+ mamba_ssm_ngroups (int, optional): the number of groups for the
+ MAMBA SSM. Defaults to 8.
+ residual_in_fp32 (bool, optional): whether to do residual connections
+ in fp32. Defaults to False.
+ pre_process (bool, optional): whether to include an embedding layer.
+ Defaults to True.
+ hybrid_attention_ratio (float, optional): the target ratio of attention layers to
+ total layers. Defaults to 0.0.
+ hybrid_mlp_ratio (float, optional): the target ratio of mlp layers to total
+ layers. Defaults to 0.0.
+ hybrid_override_pattern (str, optional): the hybrid layer pattern to override
+ with. Defaults to None.
+ post_layer_norm (bool, optional): whether to include a final layer norm.
+ Defaults to True.
+ post_process (bool, optional): whether to include an output layer.
+ Defaults to True.
+ device (optional): the device to use. Defaults to None.
+ dtype (optional): the data type to use. Defaults to None.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: MambaStackSubmodules,
+ mamba_ssm_ngroups: int = 8,
+ residual_in_fp32=False,
+ pre_process: bool = True,
+ hybrid_attention_ratio: float = 0.0,
+ hybrid_mlp_ratio: float = 0.0,
+ hybrid_override_pattern: str = None,
+ post_layer_norm: bool = True,
+ post_process: bool = True,
+ device=None,
+ dtype=None,
+ ) -> None:
+ super().__init__(config=config)
+ self.residual_in_fp32 = residual_in_fp32
+ self.pre_process = pre_process
+ self.post_layer_norm = post_layer_norm
+ self.post_process = post_process
+
+ # Required for pipeline parallel schedules
+ self.input_tensor = None
+
+ self.hybrid_attention_ratio = hybrid_attention_ratio
+ self.hybrid_mlp_ratio = hybrid_mlp_ratio
+ self.hybrid_override_pattern = hybrid_override_pattern
+
+ layer_type_list = allocate_layers(
+ self.config.num_layers,
+ self.hybrid_attention_ratio,
+ self.hybrid_mlp_ratio,
+ self.hybrid_override_pattern,
+ )
+
+ pp_layer_offset = 0
+ if parallel_state.get_pipeline_model_parallel_world_size() > 1:
+ pp_layer_offset, layer_type_list = self._select_layers_for_pipeline_parallel(
+ layer_type_list
+ )
+
+ self.layers = nn.ModuleList()
+ for i, layer_type in enumerate(layer_type_list):
+ if layer_type == LayerSymbols.MAMBA:
+ layer = build_module(
+ submodules.mamba_layer,
+ config=self.config,
+ mamba_ssm_ngroups=mamba_ssm_ngroups,
+ residual_in_fp32=residual_in_fp32,
+ layer_number=i + 1 + pp_layer_offset,
+ )
+ elif layer_type == LayerSymbols.ATTENTION:
+ # Transformer layers apply their own pp_layer_offset
+ layer = build_module(
+ submodules.attention_layer, config=self.config, layer_number=i + 1
+ )
+ elif layer_type == LayerSymbols.MLP:
+ # Transformer layers apply their own pp_layer_offset
+ layer = build_module(submodules.mlp_layer, config=self.config, layer_number=i + 1)
+ else:
+ assert True, "unexpected layer_type"
+ self.layers.append(layer)
+
+ # Required for activation recomputation
+ self.num_layers_per_pipeline_rank = len(self.layers)
+
+ if self.post_process and self.post_layer_norm:
+ # Final layer norm before output.
+ self.final_norm = TENorm(
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon,
+ )
+
+ self.apply(partial(_init_weights, n_layer=self.config.num_layers))
+
+ def _select_layers_for_pipeline_parallel(self, layer_type_list):
+ pipeline_rank = parallel_state.get_pipeline_model_parallel_rank()
+ num_layers_per_pipeline_rank = (
+ self.config.num_layers // parallel_state.get_pipeline_model_parallel_world_size()
+ )
+
+ assert parallel_state.get_virtual_pipeline_model_parallel_world_size() is None, (
+ "The Mamba hybrid model does not currently support "
+ "virtual/interleaved pipeline parallelism"
+ )
+
+ offset = pipeline_rank * num_layers_per_pipeline_rank
+ selected_list = layer_type_list[offset : offset + num_layers_per_pipeline_rank]
+
+ return offset, selected_list
+
+ def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None):
+ """
+ Allocate inference cache for each layer.
+
+ Args:
+ batch_size (int): The batch size to use for inference.
+ max_seqlen (int): The maximum sequence length to use
+ for inference.
+ dtype (optional): The data type to use for allocation.
+ Defaults to the data type of the model.
+ """
+ return {
+ i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype)
+ for i, layer in enumerate(self.layers)
+ }
+
+ def set_input_tensor(self, input_tensor: Tensor):
+ """Set input tensor to be used instead of forward()'s input.
+
+ When doing pipeline parallelism the input from the previous
+ stage comes from communication, not from the input, so the
+ model's forward_step_func won't have it. This function is thus
+ used by internal code to bypass the input provided by the
+ forward_step_func"""
+ self.input_tensor = input_tensor
+
+ def forward(
+ self,
+ hidden_states: Tensor,
+ attention_mask: Tensor,
+ inference_params=None,
+ rotary_pos_emb: Tensor = None,
+ ):
+ """
+ Forward function of the MambaStack class.
+
+ It either returns the Loss values if labels are given or the
+ final hidden units
+
+ Args:
+ hidden_states (Tensor): the input tensor.
+ attention_mask (Tensor): the attention mask.
+ inference_params (InferenceParams): the inference parameters.
+ rotary_pos_emb (Tensor, optional): the rotary positional embeddings.
+ Defaults to None.
+ Returns:
+ Tensor: the output tensor.
+ """
+ if not self.pre_process:
+ # See set_input_tensor()
+ hidden_states = self.input_tensor
+
+ if inference_params:
+ # NOTE(bnorick): match InferenceParams attributes for
+ # mamba_ssm.utils.generation.InferenceParams,
+ # this hack supports eval
+ inference_params.max_seqlen = inference_params.max_sequence_length
+ inference_params.seqlen_offset = inference_params.sequence_len_offset
+
+ for layer in self.layers:
+ hidden_states = layer(
+ hidden_states,
+ attention_mask,
+ inference_params=inference_params,
+ rotary_pos_emb=rotary_pos_emb,
+ )
+
+ # The attention layer (currently a simplified transformer layer)
+ # outputs a tuple of (hidden_states, context). Context is intended
+ # for cross-attention, and is not needed in our model.
+ if isinstance(hidden_states, tuple):
+ hidden_states = hidden_states[0]
+
+ # Final layer norm.
+ if self.post_process and self.post_layer_norm:
+ hidden_states = self.final_norm(hidden_states)
+
+ # Ensure that the tensor passed between pipeline parallel stages is
+ # viewless. See related notes in TransformerBlock and TransformerLayer
+ output = make_viewless_tensor(
+ inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True
+ )
+
+ return hidden_states
+
+ def sharded_state_dict(
+ self, prefix: str = '', sharded_offsets: tuple = (), metadata: dict = None
+ ) -> ShardedStateDict:
+ """
+ Returns a sharded state dictionary for the current object.
+
+ This function constructs a sharded state dictionary by iterating over the layers
+ in the current object, computing the sharded state dictionary for each layer,
+ and combining the results into a single dictionary.
+
+ Parameters:
+ prefix (str): The prefix to use for the state dictionary keys.
+ sharded_offsets (tuple): The sharded offsets to use for the state dictionary.
+ metadata (dict): Additional metadata to use when computing the sharded state dictionary.
+
+ Returns:
+ dict: The sharded state dictionary for the current object.
+ """
+
+ sharded_state_dict = {}
+ layer_prefix = f'{prefix}layers.'
+
+ for local_layer_idx, layer in enumerate(self.layers):
+
+ global_layer_offset = layer.layer_number - 1 # self.layer_number starts at 1
+ state_dict_prefix = (
+ f'{layer_prefix}{local_layer_idx}.' # module list index in MambaBlock
+ )
+
+ sharded_prefix = f'{layer_prefix}{global_layer_offset}.'
+ sharded_pp_offset = []
+
+ layer_sharded_state_dict = layer.sharded_state_dict(
+ state_dict_prefix, sharded_pp_offset, metadata
+ )
+
+ replace_prefix_for_sharding(layer_sharded_state_dict, state_dict_prefix, sharded_prefix)
+
+ sharded_state_dict.update(layer_sharded_state_dict)
+
+ # Add modules other than self.layers
+ for name, module in self.named_children():
+ if not module is self.layers:
+ sharded_state_dict.update(
+ sharded_state_dict_default(
+ module, f'{prefix}{name}.', sharded_offsets, metadata
+ )
+ )
+
+ return sharded_state_dict
diff --git a/megatron/core/ssm/mamba_hybrid_layer_allocation.py b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
new file mode 100644
index 0000000000..abfa2ae305
--- /dev/null
+++ b/megatron/core/ssm/mamba_hybrid_layer_allocation.py
@@ -0,0 +1,191 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+
+if __name__ != "__main__":
+ from megatron.core.utils import log_single_rank
+else:
+ from typing import Any
+
+ def log_single_rank(logger: logging.Logger, *args: Any, rank: int = 0, **kwargs: Any):
+ print(*args[1:], **kwargs)
+
+
+logger = logging.getLogger(__name__)
+
+
+class Symbols:
+ MAMBA = 'M'
+ ATTENTION = '*'
+ MLP = '-'
+ VALID = {MAMBA, ATTENTION, MLP}
+
+
+def _allocate_auto(
+ total_layers_count: int, target_attention_ratio: float, target_mlp_ratio: float
+) -> list:
+ # First, allocate attention (evenly spaced, starting and ending with mamba)
+ attention_layers_count: int = round(total_layers_count * target_attention_ratio)
+ mamba_layers_count: int = total_layers_count - attention_layers_count
+ mamba_sections_count: int = attention_layers_count + 1
+ mamba_section_length: float = mamba_layers_count / mamba_sections_count
+
+ layer_type_list = [Symbols.MAMBA] * total_layers_count
+ x: float = mamba_section_length
+ for l in range(total_layers_count):
+ if x < 0.5:
+ layer_type_list[l] = Symbols.ATTENTION
+ x += mamba_section_length
+ else:
+ x -= 1
+
+ # Next, allocate mlp
+ # (evenly distributed, but right-justified, not replacing attention)
+ mlp_layers_count: int = round(total_layers_count * target_mlp_ratio)
+ if mlp_layers_count > 0:
+ mamba_layers_count -= mlp_layers_count
+ mamba_to_mlp_ratio: float = mamba_layers_count / mlp_layers_count
+
+ x: float = mamba_to_mlp_ratio
+ for l in range(total_layers_count):
+ if layer_type_list[l] == Symbols.MAMBA:
+ if x < 0.5:
+ layer_type_list[l] = Symbols.MLP
+ x += mamba_to_mlp_ratio
+ else:
+ x -= 1
+
+ return layer_type_list
+
+
+def _allocate_override(total_layers_count: int, override_pattern: str) -> list:
+ layer_type_list = list(override_pattern)
+ override_pattern_length = len(layer_type_list)
+ if override_pattern_length != total_layers_count:
+ raise ValueError(
+ "The hybrid override pattern is the wrong "
+ f"length: got {override_pattern_length}, expected "
+ f"{total_layers_count}"
+ )
+ for l in layer_type_list:
+ if l not in Symbols.VALID:
+ raise ValueError(f"In hybrid override pattern, '{l}' is not " f"one of {Symbols.VALID}")
+
+ return layer_type_list
+
+
+def _layer_counts_match(a: list, b: list) -> bool:
+ for s in Symbols.VALID:
+ if a.count(s) != b.count(s):
+ return False
+ return True
+
+
+def allocate_layers(
+ total_layers_count: int,
+ target_attention_ratio: float,
+ target_mlp_ratio: float,
+ override_pattern: str = None,
+) -> list:
+ assert total_layers_count > 0
+ assert target_attention_ratio >= 0.0 and target_attention_ratio <= 1.0
+ assert target_mlp_ratio >= 0.0 and target_mlp_ratio <= 1.0
+ assert target_attention_ratio + target_mlp_ratio <= 1.0
+ # Note: target_mamba_ratio = 1.0 - target_attention_ratio - target_mlp_ratio
+
+ layer_type_list = _allocate_auto(total_layers_count, target_attention_ratio, target_mlp_ratio)
+
+ if override_pattern is not None:
+ layer_type_list_override = _allocate_override(total_layers_count, override_pattern)
+ log_single_rank(logger, logging.INFO, "Using hybrid override pattern")
+ if (target_attention_ratio > 0.0 or target_mlp_ratio > 0.0) and not _layer_counts_match(
+ layer_type_list_override, layer_type_list
+ ):
+ raise ValueError(
+ "The number of each type of layer in the override "
+ "pattern must match the number in the overridden "
+ "pattern."
+ )
+ if layer_type_list_override == layer_type_list:
+ log_single_rank(
+ logger, logging.INFO, "The override pattern matches the overridden pattern"
+ )
+ else:
+ log_single_rank(logger, logging.INFO, "Warning: overriding pattern A with pattern B")
+ log_single_rank(logger, logging.INFO, f"A: {''.join(layer_type_list)}")
+ log_single_rank(logger, logging.INFO, f"B: {''.join(layer_type_list_override)}")
+ layer_type_list = layer_type_list_override
+
+ if target_attention_ratio > 0.0 or target_mlp_ratio > 0.0 or override_pattern is not None:
+ actual_attention_layers_count = layer_type_list.count(Symbols.ATTENTION)
+ actual_attention_ratio = actual_attention_layers_count / total_layers_count
+ actual_mlp_layers_count = layer_type_list.count(Symbols.MLP)
+ actual_mlp_ratio = actual_mlp_layers_count / total_layers_count
+ allocation_string = ''.join(layer_type_list)
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Hybrid allocation ({Symbols.MAMBA} is mamba, "
+ f"{Symbols.ATTENTION} is attention, "
+ f"{Symbols.MLP} is mlp):",
+ )
+ log_single_rank(logger, logging.INFO, allocation_string)
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"{actual_attention_layers_count} attention layers in "
+ f"{total_layers_count} total layers.",
+ )
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Target attention ratio: {target_attention_ratio:.2f}. "
+ f"Actual attention ratio: {actual_attention_ratio:.2f}.",
+ )
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"{actual_mlp_layers_count} mlp layers in " f"{total_layers_count} total layers.",
+ )
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Target mlp ratio: {target_mlp_ratio:.2f}. "
+ f"Actual mlp ratio: {actual_mlp_ratio:.2f}.",
+ )
+ return layer_type_list
+
+
+if __name__ == "__main__":
+ test_cases = [
+ # (10, 0.2, 0.0),
+ # (48, 0.0, 0.0), # will not print anything
+ # (48, 0.1, 0.0),
+ # 48, 0.3, 0.0),
+ # (48, 0.5, 0.0),
+ # (48, 0.6, 0.0),
+ # (48, 0.7, 0.0),
+ # (10, 0.0, 0.1),
+ # (10, 0.0, 0.3),
+ # (10, 0.0, 0.5),
+ # (10, 0.1, 0.1),
+ # (10, 0.2, 0.2),
+ # (10, 0.3, 0.3),
+ # (10, 0.5, 0.5),
+ # (48, 0.2, 0.3),
+ # (48, 0.5, 0.2),
+ # (48, 0.5, 0.2, "MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-"),
+ # (48, 0.25, 0.25, "MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-"),
+ # (48, 0.25, 0.25, "MM-*MM-*MM*-MM*-MM*-MM*-M*M-M*M-M*M-M*M-*MM-*MM-"),
+ # (48, 0.0, 0.2, "MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-"),
+ # (48, 0.2, 0.0, "MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-"),
+ # (48, 0.0, 0.0, "MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-MM*-"),
+ # (48, 0.5, 0.5),
+ # (10, 0.3, 0.2, "MMM*-*M*M-"),
+ # (10, 0.3, 0.2, "MM*M-*M*M-"),
+ (9, 0.0, 0.0, "M*-M*-M*-"),
+ (9, 0.0, 0.0, "MMMMMMMMM"),
+ ]
+ for t in test_cases:
+ print("")
+ allocate_layers(*t)
diff --git a/megatron/core/ssm/mamba_layer.py b/megatron/core/ssm/mamba_layer.py
new file mode 100644
index 0000000000..f0776746dd
--- /dev/null
+++ b/megatron/core/ssm/mamba_layer.py
@@ -0,0 +1,116 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2024, Tri Dao, Albert Gu.
+
+# Some of this code was adopted from https://github.com/state-spaces/mamba/
+# This source code is licensed under the Apache license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+from torch import Tensor
+
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+@dataclass
+class MambaLayerSubmodules:
+ """
+ Configuration class for specifying the submodules of a Mamba layer.
+
+ This class defines the structure and default implementations for various
+ components of a Mamba layer, allowing for flexible customization of the
+ layer's architecture.
+
+ Args:
+ norm (Union[ModuleSpec, type]): Specification for the input layer normalization.
+ mixer (Union[ModuleSpec, type]): Specification for the along-sequence mixing mechanism.
+ mamba_bda (Union[ModuleSpec, type]): Specification for the bias-dropout-add operation
+ after the mixer.
+ """
+
+ norm: Union[ModuleSpec, type] = IdentityOp
+ mixer: Union[ModuleSpec, type] = IdentityOp
+ mamba_bda: Union[ModuleSpec, type] = IdentityOp
+
+
+class MambaLayer(MegatronModule):
+ """
+ A single Mamba layer.
+
+ Mamba layer takes input with size [s, b, h] and returns an
+ output of the same size.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: MambaLayerSubmodules,
+ mamba_ssm_ngroups=8,
+ layer_number: int = 1,
+ residual_in_fp32=False,
+ ):
+ """Initialize Mamba Layer."""
+ super().__init__(config)
+ self.config = config
+ self.layer_number = layer_number
+ self.residual_in_fp32 = residual_in_fp32
+ self.hidden_dropout = config.hidden_dropout
+ self.mixer = build_module(
+ submodules.mixer,
+ self.config,
+ d_model=self.config.hidden_size,
+ ngroups=mamba_ssm_ngroups,
+ layer_number=layer_number,
+ )
+ self.norm = build_module(submodules.norm, self.config, self.config.hidden_size)
+ self.mamba_bda = build_module(submodules.mamba_bda)
+ self.bias_dropout_add_exec_handler = torch.enable_grad
+
+ def forward(
+ self,
+ hidden_states: Tensor,
+ attention_mask: Tensor, # Not used in MambaLayer
+ inference_params=None,
+ rotary_pos_emb: Tensor = None, # Not used in MambaLayer
+ ):
+ """
+ Perform a forward pass through the Mamba layer.
+
+ This method implements the core computation of a Mamba layer, including
+ the convolution and the selective SSM/SSD.
+
+ Args:
+ hidden_states (Tensor): Input tensor of shape [s, b, h] where s is sequence length,
+ b is batch size, and h is hidden size.
+ attention_mask (Tensor): Mask tensor for self-attention. Not used by this layer.
+ inference_params (object, optional): Parameters for inference-time optimizations.
+ rotary_pos_emb (Tensor, optional): Rotary positional embeddings.
+
+ Returns:
+ output (Tensor): Transformed hidden states of shape [s, b, h].
+ """
+
+ residual = hidden_states
+ if self.residual_in_fp32:
+ residual = residual.to(torch.float32)
+
+ hidden_states = hidden_states.to(dtype=self.config.params_dtype)
+ hidden_states = self.norm(hidden_states)
+
+ mixer_out_with_bias = self.mixer(hidden_states, inference_params=inference_params)
+
+ with self.bias_dropout_add_exec_handler():
+ hidden_states = self.mamba_bda(self.training, self.config.bias_dropout_fusion)(
+ mixer_out_with_bias, residual, self.hidden_dropout
+ )
+
+ return hidden_states
+
+ def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None):
+ """Allocate the inference cache."""
+ return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype)
diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py
new file mode 100644
index 0000000000..6448f30d9c
--- /dev/null
+++ b/megatron/core/ssm/mamba_mixer.py
@@ -0,0 +1,718 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2024, Tri Dao, Albert Gu.
+
+# Some of this code was adopted from https://github.com/state-spaces/mamba/
+# This source code is licensed under the Apache license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from dataclasses import dataclass, replace
+from typing import List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from megatron.core.dist_checkpointing import ShardedTensor
+from megatron.core.dist_checkpointing.mapping import ReplicaId, ShardedTensorFactory
+from megatron.core.parallel_state import get_tensor_model_parallel_world_size
+from megatron.core.tensor_parallel import get_cuda_rng_tracker
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import (
+ make_sharded_tensors_for_checkpoint,
+ sharded_state_dict_default,
+)
+
+try:
+ from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+except ImportError:
+ selective_state_update = None
+
+try:
+ from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+except ImportError:
+ causal_conv1d_fn = None
+ causal_conv1d_update = None
+
+try:
+ from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated
+ from mamba_ssm.ops.triton.ssd_combined import (
+ mamba_chunk_scan_combined,
+ mamba_split_conv1d_scan_combined,
+ )
+except ImportError:
+ raise ImportError("mamba-ssm is required by the Mamba model but cannot be imported")
+
+try:
+ from einops import rearrange, repeat
+except ImportError:
+ raise ImportError("einops is required by the Mamba model but cannot be imported")
+
+
+class ExtendedRMSNorm(RMSNormGated):
+ """
+ RMSNormGated with sharded state dict.
+ """
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Sharding along axis 0, bias not sharded"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(
+ state_dict, prefix, {'weight': 0}, sharded_offsets
+ )
+
+
+@dataclass
+class MambaMixerSubmodules:
+ """
+ Contains the module specs for the input and output linear layers.
+ """
+
+ in_proj: Union[ModuleSpec, type] = None
+ out_proj: Union[ModuleSpec, type] = None
+
+
+class MambaMixer(MegatronModule):
+ """
+ Args:
+ config: The config of the model.
+ submodules: Contains the module specs for the input and output linear layers.
+ d_model: The hidden size of the model.
+ d_state: The state size of the SSM.
+ d_conv: The number of channels in the causal convolution.
+ conv_init: The initialization range for the causal convolution weights.
+ expand: The expansion factor for the SSM.
+ headdim: The hidden size of each attention head.
+ ngroups: The number of attention heads.
+ A_init_range: The initialization range for the attention weights.
+ D_has_hdim: Whether the D parameter has the same number of dimensions as the hidden
+ state.
+ rmsnorm: Whether to use root mean square normalization.
+ norm_before_gate: Whether to apply normalization before the gating mechanism.
+ dt_min: The minimum value of the dt parameter.
+ dt_max: The maximum value of the dt parameter.
+ dt_init: The initialization value of the dt parameter.
+ dt_scale: The scaling factor for the dt parameter.
+ dt_init_floor: The minimum value of the dt parameter after initialization.
+ bias: Whether to use bias in the linear layers.
+ conv_bias: Whether to use bias in the causal convolution.
+ chunk_size: The chunk size for the fused kernel.
+ use_mem_eff_path: Whether to use the memory-efficient path for the Mamba model.
+ layer_number: The layer number of this Mamba layer.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: MambaMixerSubmodules,
+ d_model,
+ d_state=128,
+ d_conv=4,
+ conv_init=None,
+ expand=2,
+ headdim=64,
+ ngroups=8,
+ A_init_range=(1, 16),
+ D_has_hdim=False,
+ rmsnorm=True,
+ norm_before_gate=False,
+ dt_min=0.001,
+ dt_max=0.1,
+ dt_init="random",
+ dt_scale=1.0,
+ dt_init_floor=1e-4,
+ bias=False,
+ conv_bias=True,
+ # Fused kernel and sharding options
+ chunk_size=128,
+ use_mem_eff_path=True,
+ layer_number=None,
+ ):
+ super().__init__(config)
+ self.config = config
+ self.d_model = d_model
+ self.d_state = d_state
+ self.d_conv = d_conv
+ self.conv_init = conv_init
+ self.expand = expand
+ self.d_inner = int(self.expand * self.d_model)
+ self.headdim = headdim
+ self.ngroups = ngroups
+ assert self.d_inner % self.headdim == 0
+ self.nheads = self.d_inner // self.headdim
+ self.D_has_hdim = D_has_hdim
+ self.rmsnorm = rmsnorm
+ self.norm_before_gate = norm_before_gate
+ self.chunk_size = chunk_size
+ self.use_mem_eff_path = use_mem_eff_path
+ self.layer_number = layer_number
+
+ self.tensor_model_parallel_size = get_tensor_model_parallel_world_size()
+ assert self.d_inner % self.tensor_model_parallel_size == 0
+ assert self.ngroups % self.tensor_model_parallel_size == 0
+ assert self.nheads % self.tensor_model_parallel_size == 0
+ assert not bias
+ assert not self.norm_before_gate
+
+ self.d_inner_local = self.d_inner // self.tensor_model_parallel_size
+ self.ngroups_local = self.ngroups // self.tensor_model_parallel_size
+ self.nheads_local = self.nheads // self.tensor_model_parallel_size
+
+ assert self.d_inner_local % self.ngroups_local == 0
+
+ # Assume sequence parallelism: input is already partitioned along the
+ # sequence dimension
+ self.in_proj = build_module(
+ submodules.in_proj,
+ self.d_model,
+ self.d_inner * 2 + 2 * self.ngroups * self.d_state + self.nheads, # AB CD E
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=bias,
+ skip_bias_add=False,
+ is_expert=False,
+ tp_comm_buffer_name='fc1',
+ )
+
+ conv_dim = self.d_inner_local + 2 * self.ngroups_local * self.d_state # A CD
+ with get_cuda_rng_tracker().fork():
+ # weight dim: [conv_dim, conv_dim, d_conv]
+ self.conv1d = nn.Conv1d(
+ in_channels=conv_dim,
+ out_channels=conv_dim,
+ bias=conv_bias,
+ kernel_size=d_conv,
+ groups=conv_dim,
+ padding=d_conv - 1,
+ device=torch.cuda.current_device(),
+ dtype=config.params_dtype,
+ )
+ setattr(self.conv1d.weight, 'tensor_model_parallel', True)
+ setattr(self.conv1d.bias, 'tensor_model_parallel', True)
+
+ if self.conv_init is not None:
+ nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init)
+
+ self.activation = "silu"
+ self.act = nn.SiLU()
+
+ with get_cuda_rng_tracker().fork():
+ # Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max
+ dt = torch.exp(
+ torch.rand(
+ self.nheads_local, device=torch.cuda.current_device(), dtype=config.params_dtype
+ )
+ * (math.log(dt_max) - math.log(dt_min))
+ + math.log(dt_min)
+ ).clamp(min=dt_init_floor)
+ # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+ inv_dt = dt + torch.log(-torch.expm1(-dt))
+ with torch.no_grad():
+ self.dt_bias = nn.Parameter(inv_dt)
+ # Our initialization would set all Linear.bias to zero,
+ # need to mark this one as _no_reinit
+ self.dt_bias._no_reinit = True
+ # Just to be explicit. Without this we already don't
+ # put wd on dt_bias because of the check
+
+ # name.endswith("bias") in param_grouping.py
+ self.dt_bias._no_weight_decay = True
+
+ assert A_init_range[0] > 0 and A_init_range[1] >= A_init_range[0]
+ A = torch.empty(
+ self.nheads_local, dtype=torch.float32, device=torch.cuda.current_device()
+ ).uniform_(*A_init_range)
+ A_log = torch.log(A) # Keep A_log in fp32
+ self.A_log = nn.Parameter(A_log)
+ self.A_log._no_weight_decay = True
+ setattr(self.A_log, 'tensor_model_parallel', True)
+
+ # D "skip" parameter
+ self.D = nn.Parameter(
+ torch.ones(
+ self.d_inner_local if self.D_has_hdim else self.nheads_local,
+ device=torch.cuda.current_device(),
+ )
+ ) # Keep in fp32
+ self.D._no_weight_decay = True
+ setattr(self.D, 'tensor_model_parallel', True)
+
+ if self.rmsnorm:
+ assert RMSNormGated is not None
+ self.norm = ExtendedRMSNorm(
+ self.d_inner_local,
+ eps=1e-5,
+ group_size=self.d_inner_local // self.ngroups_local,
+ norm_before_gate=self.norm_before_gate,
+ device=torch.cuda.current_device(),
+ dtype=config.params_dtype,
+ )
+
+ # Assume sequence parallelism: input is partitioned along d_inner and
+ # output is partitioned along the sequence dimension
+ self.out_proj = build_module(
+ submodules.out_proj,
+ self.d_inner,
+ self.d_model,
+ config=self.config,
+ init_method=self.config.output_layer_init_method,
+ bias=bias,
+ input_is_parallel=True,
+ skip_bias_add=True,
+ is_expert=False,
+ tp_comm_buffer_name='fc2',
+ )
+
+ def forward(self, hidden_states, inference_params=None):
+ """
+ hidden_states: (nL, B, D) / (L B D)
+ Returns: same shape as hidden_states
+ """
+ _, batch, dim = hidden_states.shape
+
+ conv_state, ssm_state = None, None
+ if inference_params is not None:
+ assert not self.config.sequence_parallel
+ conv_state, ssm_state = self._get_states_from_cache(inference_params, batch)
+ if inference_params.seqlen_offset > 0:
+ # The states are updated inplace
+ out, out_bias, _, _ = self.step(hidden_states, conv_state, ssm_state)
+ return out, out_bias
+
+ # (nheads_local)
+ A = -torch.exp(self.A_log.float())
+
+ xz, _ = self.in_proj(hidden_states)
+
+ # transpose: l b pd --> b l pd
+ xz = rearrange(xz, "l b d -> b l d").contiguous()
+
+ if self.use_mem_eff_path and inference_params is None:
+ assert ssm_state is None
+
+ if self.conv1d.bias is not None:
+ self.conv1d.bias.data_ptr()
+
+ y = mamba_split_conv1d_scan_combined(
+ xz,
+ rearrange(self.conv1d.weight, "d 1 w -> d w"),
+ self.conv1d.bias,
+ self.dt_bias.float(),
+ A,
+ D=(
+ rearrange(self.D.float(), "(h p) -> h p", p=self.headdim)
+ if self.D_has_hdim
+ else self.D
+ ),
+ chunk_size=self.chunk_size,
+ activation=self.activation,
+ headdim=None if self.D_has_hdim else self.headdim,
+ ngroups=self.ngroups_local,
+ norm_before_gate=self.norm_before_gate,
+ )
+
+ if self.rmsnorm:
+ y = self.norm(y)
+ else:
+ z, xBC, dt = torch.split(
+ xz,
+ [
+ self.d_inner_local,
+ self.d_inner_local + 2 * self.ngroups_local * self.d_state,
+ self.nheads_local,
+ ],
+ dim=-1,
+ )
+
+ # transpose: b l pd --> b pd l
+ xBC = rearrange(xBC, "b l d -> b d l").contiguous()
+
+ # Compute short convolution
+ if conv_state is not None:
+ # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv
+ # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise.
+ conv_state.copy_(
+ F.pad(xBC, (self.d_conv - xBC.shape[-1], 0))
+ ) # Update state (B D W)
+
+ seqlen = xBC.size(2)
+ if causal_conv1d_fn is None:
+ xBC = self.act(self.conv1d(xBC)[..., :seqlen])
+ else:
+ assert self.activation in ["silu", "swish"]
+ xBC = causal_conv1d_fn(
+ x=xBC,
+ weight=rearrange(self.conv1d.weight, "d 1 w -> d w"),
+ bias=self.conv1d.bias,
+ activation=self.activation,
+ )
+
+ # transpose b pd l --> b l pd
+ xBC = rearrange(xBC, "b d l -> b l d").contiguous()
+
+ x, B, C = torch.split(
+ xBC,
+ [
+ self.d_inner_local,
+ self.ngroups_local * self.d_state,
+ self.ngroups_local * self.d_state,
+ ],
+ dim=-1,
+ )
+
+ # TODO Vijay: fuse most of the transposes with the GEMMS
+ x = rearrange(x, "b l (h p) -> b l h p", p=self.headdim).contiguous()
+ dt = dt.contiguous()
+ B = rearrange(B, "b l (g n) -> b l g n", n=self.d_state).contiguous()
+ C = rearrange(C, "b l (g n) -> b l g n", n=self.d_state).contiguous()
+ z = rearrange(z, "b l (h p) -> b l h p", p=self.headdim).contiguous()
+ y = mamba_chunk_scan_combined(
+ x,
+ dt,
+ A,
+ B,
+ C,
+ self.chunk_size,
+ D=(
+ rearrange(self.D.float(), "(h p) -> h p", p=self.headdim)
+ if self.D_has_hdim
+ else self.D
+ ),
+ z=z if not self.rmsnorm else None,
+ dt_bias=self.dt_bias.float(),
+ dt_softplus=True,
+ return_final_states=ssm_state is not None,
+ )
+
+ if ssm_state is not None:
+ y, last_state = y
+ ssm_state.copy_(last_state)
+
+ if self.rmsnorm:
+ y = rearrange(y, "b l h p -> b l (h p)").contiguous()
+ z = rearrange(z, "b l h p -> b l (h p)").contiguous()
+ y = self.norm(y, z)
+ else:
+ y = rearrange(y, "b l h p -> b l (h p)").contiguous()
+
+ y = rearrange(y, "b l d -> l b d").contiguous()
+ out, out_bias = self.out_proj(y)
+
+ return out, out_bias
+
+ def step(self, hidden_states, conv_state, ssm_state):
+ """
+ Performs inference step for decoding
+ """
+ # assert self.ngroups_local == 1, "Only support ngroups=1 for inference for now"
+ dtype = hidden_states.dtype
+ assert hidden_states.shape[0] == 1, "Only support decoding with 1 token at a time for now"
+
+ # l b d --> b d
+ hidden_states = hidden_states.squeeze(0)
+
+ # b d_model --> b p(2d)
+ xz, _ = self.in_proj(hidden_states)
+
+ z, xBC, dt = torch.split(
+ xz,
+ [
+ self.d_inner_local,
+ self.d_inner_local + 2 * self.ngroups_local * self.d_state,
+ self.nheads_local,
+ ],
+ dim=-1,
+ )
+
+ # Conv step
+ if causal_conv1d_update is None:
+ conv_state.copy_(torch.roll(conv_state, shifts=-1, dims=-1)) # Update state (B D W)
+ conv_state[:, :, -1] = xBC
+ xBC = torch.sum(
+ conv_state * rearrange(self.conv1d.weight, "d 1 w -> d w"), dim=-1
+ ) # (B D)
+ if self.conv1d.bias is not None:
+ xBC = xBC + self.conv1d.bias
+ xBC = self.act(xBC).to(dtype=dtype)
+ else:
+ xBC = causal_conv1d_update(
+ xBC,
+ conv_state,
+ rearrange(self.conv1d.weight, "d 1 w -> d w"),
+ self.conv1d.bias,
+ self.activation,
+ )
+
+ x, B, C = torch.split(
+ xBC,
+ [
+ self.d_inner_local,
+ self.ngroups_local * self.d_state,
+ self.ngroups_local * self.d_state,
+ ],
+ dim=-1,
+ )
+ A = -torch.exp(self.A_log.float())
+
+ # SSM step
+ if selective_state_update is None:
+ if self.ngroups_local > 1:
+ B = rearrange(B, "b (g n) -> b g n", n=self.d_state)
+ C = rearrange(C, "b (g n) -> b g n", n=self.d_state)
+ B = repeat(B, "b g n -> b (g h) n", h=self.d_inner_local // self.ngroups_local)
+ C = repeat(C, "b g n -> b (g h) n", h=self.d_inner_local // self.ngroups_local)
+
+ dt = repeat(dt, "b h -> b (h p)", p=self.headdim)
+ dt_bias = repeat(self.dt_bias, "h -> (h p)", p=self.headdim)
+ A = repeat(A, "h -> (h p) n", p=self.headdim, n=self.d_state)
+ D = repeat(self.D, "h -> (h p)", p=self.headdim)
+
+ dt = F.softplus(dt + dt_bias.to(dtype=dt.dtype))
+ dA = torch.exp(torch.einsum("bd,dn->bdn", dt, A))
+
+ dB_x = torch.einsum('bd,bdn,bd->bdn', dt, B, x)
+ ssm_state.copy_(
+ ssm_state * rearrange(dA, "b (h p) n -> b h p n", p=self.headdim)
+ + rearrange(dB_x, "b (h p) n -> b h p n", p=self.headdim)
+ )
+
+ y = torch.einsum(
+ "bdn,bdn->bd",
+ rearrange(ssm_state.to(dtype), "b h p n -> b (h p) n", p=self.headdim),
+ C,
+ )
+ y = y + D.to(dtype) * x
+ if not self.rmsnorm:
+ y = y * self.act(z) # (B D)
+ else:
+ # Discretize A and B (b (g n))
+ dt = F.softplus(dt + self.dt_bias.to(dtype=dt.dtype)) # (batch, nheads)
+ dA = torch.exp(dt * A)
+ x = rearrange(x, "b (h p) -> b h p", p=self.headdim)
+ dBx = torch.einsum("bh,bn,bhp->bhpn", dt, B, x)
+ ssm_state.copy_(ssm_state * rearrange(dA, "b h -> b h 1 1") + dBx)
+ y = torch.einsum("bhpn,bn->bhp", ssm_state.to(dtype), C)
+ y = y + rearrange(self.D.to(dtype), "h -> h 1") * x
+ y = rearrange(y, "b h p -> b (h p)")
+ if not self.rmsnorm:
+ y = y * self.act(z) # (B D)
+ else:
+ A = repeat(A, "h -> h p n", p=self.headdim, n=self.d_state).to(dtype=torch.float32)
+ dt = repeat(dt, "b h -> b h p", p=self.headdim)
+ dt_bias = repeat(self.dt_bias, "h -> h p", p=self.headdim)
+ D = repeat(self.D, "h -> h p", p=self.headdim)
+ B = rearrange(B, "b (g n) -> b g n", g=self.ngroups_local)
+ C = rearrange(C, "b (g n) -> b g n", g=self.ngroups_local)
+ x_reshaped = rearrange(x, "b (h p) -> b h p", p=self.headdim)
+ if not self.rmsnorm:
+ z = rearrange(z, "b (h p) -> b h p", p=self.headdim)
+ y = selective_state_update(
+ ssm_state,
+ x_reshaped,
+ dt,
+ A,
+ B,
+ C,
+ D,
+ z=z if not self.rmsnorm else None,
+ dt_bias=dt_bias,
+ dt_softplus=True,
+ )
+ y = rearrange(y, "b h p -> b (h p)")
+
+ if self.rmsnorm:
+ y = self.norm(y, z)
+
+ # b pd --> b d
+ out, out_bias = self.out_proj(y)
+ return out.unsqueeze(0), out_bias, conv_state, ssm_state
+
+ def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None):
+ """
+ allocate inference cache
+ """
+ device = self.out_proj.weight.device
+ conv_dtype = self.conv1d.weight.dtype if dtype is None else dtype
+ conv_state = torch.zeros(
+ batch_size, self.conv1d.weight.shape[0], self.d_conv, device=device, dtype=conv_dtype
+ )
+ ssm_dtype = self.in_proj.weight.dtype if dtype is None else dtype
+ # ssm_dtype = torch.float32
+ ssm_state = torch.zeros(
+ batch_size,
+ self.nheads_local,
+ self.headdim,
+ self.d_state,
+ device=device,
+ dtype=ssm_dtype,
+ )
+ return conv_state, ssm_state
+
+ def _get_states_from_cache(self, inference_params, batch_size, initialize_states=False):
+ assert self.layer_number is not None
+ if self.layer_number not in inference_params.key_value_memory_dict:
+ conv_state = torch.zeros(
+ batch_size,
+ self.conv1d.weight.shape[0],
+ self.d_conv,
+ device=self.conv1d.weight.device,
+ dtype=self.conv1d.weight.dtype,
+ )
+ ssm_state = torch.zeros(
+ batch_size,
+ self.nheads_local,
+ self.headdim,
+ self.d_state,
+ device=self.in_proj.weight.device,
+ dtype=self.in_proj.weight.dtype,
+ )
+ inference_params.key_value_memory_dict[self.layer_number] = (conv_state, ssm_state)
+ else:
+ conv_state, ssm_state = inference_params.key_value_memory_dict[self.layer_number]
+ # TODO: What if batch size changes between generation, and we reuse the same states?
+ if initialize_states:
+ conv_state.zero_()
+ ssm_state.zero_()
+ return conv_state, ssm_state
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Provide a sharded state dictionary for distributed checkpointing."""
+ sharded_state_dict = {}
+ # Parameters
+ self._save_to_state_dict(sharded_state_dict, '', keep_vars=True)
+ sharded_state_dict = make_sharded_tensors_for_checkpoint(
+ sharded_state_dict,
+ prefix,
+ tensor_parallel_layers_axis_map={
+ 'A_log': 0,
+ 'dt_bias': 0,
+ 'D': 0,
+ }, # parameters sharded across TP
+ sharded_offsets=sharded_offsets,
+ )
+ # Submodules
+ for name, module in self.named_children():
+ if name == 'conv1d':
+ # Add TP sharding for Conv1d
+ module_sd = module.state_dict(prefix='', keep_vars=True)
+ module_sharded_sd = make_sharded_tensors_for_checkpoint(
+ module_sd, f'{prefix}{name}.', {f'weight': 0, f'bias': 0}, sharded_offsets
+ )
+
+ else:
+ module_sharded_sd = sharded_state_dict_default(
+ module, f'{prefix}{name}.', sharded_offsets, metadata
+ )
+
+ sharded_state_dict.update(module_sharded_sd)
+
+ # At this point the TP sharding is correctly defined fo each tensor, but some of the tensors
+ # must be additionally split into separate parts
+ # in_proj
+ in_proj_dim = (
+ self.d_inner_local * 2 + 2 * self.ngroups_local * self.d_state + self.nheads_local
+ )
+ assert sharded_state_dict[f'{prefix}in_proj.weight'].data.size(0) == in_proj_dim, (
+ in_proj_dim,
+ sharded_state_dict[f'{prefix}in_proj.weight'],
+ )
+
+ sharded_state_dict[f'{prefix}in_proj.weight'] = _split_tensor_factory(
+ sharded_state_dict[f'{prefix}in_proj.weight'],
+ [
+ self.d_inner_local,
+ self.d_inner_local,
+ self.ngroups_local * self.d_state,
+ self.ngroups_local * self.d_state,
+ self.nheads_local,
+ ],
+ ['z', 'x', 'B', 'C', 'dt'],
+ 0,
+ )
+
+ conv_dim = self.d_inner_local + 2 * self.ngroups_local * self.d_state
+ assert sharded_state_dict[f'{prefix}conv1d.weight'].data.size(0) == conv_dim, (
+ conv_dim,
+ sharded_state_dict[f'{prefix}conv1d.weight'],
+ )
+ assert sharded_state_dict[f'{prefix}conv1d.bias'].data.size(0) == conv_dim, (
+ conv_dim,
+ sharded_state_dict[f'{prefix}conv1d.bias'],
+ )
+
+ for conv_layer_name in ['conv1d.weight', 'conv1d.bias']:
+ sharded_state_dict[f'{prefix}{conv_layer_name}'] = _split_tensor_factory(
+ sharded_state_dict[f'{prefix}{conv_layer_name}'],
+ [
+ self.d_inner_local,
+ self.ngroups_local * self.d_state,
+ self.ngroups_local * self.d_state,
+ ],
+ ['x', 'B', 'C'],
+ 0,
+ )
+
+ return sharded_state_dict
+
+
+def _split_tensor_factory(
+ orig_sh_ten: ShardedTensor, split_sections: List[int], split_names: List[str], split_dim: int
+) -> ShardedTensorFactory:
+ """Builds a factory that splits a given ShardedTensor into several independent chunks."""
+ assert isinstance(orig_sh_ten, ShardedTensor), type(orig_sh_ten)
+ orig_sh_ten_no_data = orig_sh_ten.without_data() # remove `data` reference
+
+ if sum(split_sections) != orig_sh_ten_no_data.local_shape[split_dim]:
+ raise ValueError(
+ f'Split sections must cover the whole dimension size, '
+ f'got {split_sections=} vs dimensions size '
+ f'{orig_sh_ten_no_data.local_shape[split_dim]}'
+ )
+
+ assert not isinstance(
+ split_sections, int
+ ), 'Splitting into predefined section sizes is supported (`split_sections` must be a list)'
+ assert len(split_sections) == len(split_names), (len(split_sections), len(split_names))
+
+ @torch.no_grad()
+ def sh_ten_build_fn(
+ key: str, t: torch.Tensor, replica_id: ReplicaId, flattened_range: Optional[slice]
+ ):
+ factory_sh_ten = replace(
+ orig_sh_ten_no_data,
+ key=key,
+ data=t,
+ dtype=t.dtype,
+ replica_id=replica_id,
+ flattened_range=flattened_range,
+ )
+
+ chunk_sh_tens = []
+ split_start = 0
+ for split_size, split_name in zip(split_sections, split_names):
+ split_chunks = factory_sh_ten.narrow(split_dim, split_start, split_size)
+ for sh_ten in split_chunks:
+ sh_ten.key = f'{sh_ten.key}.{split_name}'
+ chunk_sh_tens.extend(split_chunks)
+ split_start += split_size
+
+ assert split_start == orig_sh_ten_no_data.local_shape[split_dim], (
+ split_start,
+ orig_sh_ten_no_data.local_shape[split_dim],
+ )
+ assert sum(sh_ten.data.numel() for sh_ten in chunk_sh_tens) == t.numel(), (
+ chunk_sh_tens,
+ t.shape,
+ )
+ return chunk_sh_tens
+
+ @torch.no_grad()
+ def sh_ten_merge_fn(sub_state_dict):
+ return torch.cat(sub_state_dict)
+
+ return ShardedTensorFactory(
+ orig_sh_ten.key, orig_sh_ten.data, sh_ten_build_fn, sh_ten_merge_fn, orig_sh_ten.replica_id
+ )
diff --git a/megatron/core/ssm/triton_cache_manager.py b/megatron/core/ssm/triton_cache_manager.py
new file mode 100644
index 0000000000..781f17d32c
--- /dev/null
+++ b/megatron/core/ssm/triton_cache_manager.py
@@ -0,0 +1,81 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright 2018-2020 Philippe Tillet
+# Copyright 2020-2022 OpenAI
+
+# Some of this code was adopted from https://github.com/triton-lang/triton
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import uuid
+from pathlib import Path
+
+try:
+ from triton import __version__ as triton_version
+ from triton.runtime.cache import FileCacheManager
+except ImportError:
+ raise ImportError("triton is required by the Mamba model but cannot be imported")
+
+
+def _version_no_greater_than(version, version_limit):
+ major, minor, _ = map(int, version.split('.'))
+ limit_major, limit_minor = map(int, version_limit.split('.'))
+ return major < limit_major or (major == limit_major and minor <= limit_minor)
+
+
+def default_cache_dir():
+ """Provides a default path for the Triton cache directory."""
+ return os.path.join(Path.home(), ".triton", "cache")
+
+
+class ParallelFileCacheManager(FileCacheManager):
+ """
+ This patched version of ParallelFileCacheManager prevents errors related
+ to the builing of the Triton compiler cache when the number of model
+ parallel ranks is greater than one, including when certain types of file
+ system are used (such as Lustre).
+
+ Usage:
+ export TRITON_CACHE_DIR=
+ export TRITON_CACHE_MANAGER=megatron.core.ssm.triton_cache_manager:ParallelFileCacheManager
+
+ This patch implements the changes in the following two Triton project pull
+ requests:
+ 1. https://github.com/triton-lang/triton/pull/3544
+ 2. https://github.com/triton-lang/triton/pull/4295
+
+ The above changes will probably be included in Triton release version 3.1,
+ making this patch no longer necessary.
+ """
+
+ def put(self, data, filename, binary=True) -> str:
+ """A patched version of put, implementing PR 3544 and PR 4295."""
+ patch_limit = '3.0'
+ assert _version_no_greater_than(triton_version, patch_limit), (
+ "Assertion failed: ParallelFileCacheManager patch should not be "
+ f"used beyond Triton version {patch_limit}."
+ )
+ if not self.cache_dir:
+ raise RuntimeError("Could not create or locate cache dir")
+ binary = isinstance(data, bytes)
+ if not binary:
+ data = str(data)
+ assert self.lock_path is not None
+ filepath = self._make_path(filename)
+ # Random ID to avoid any collisions
+ rnd_id = str(uuid.uuid4())
+ # we use the PID in case a bunch of these around so we can see what PID made it
+ pid = os.getpid()
+ # use temp dir to be robust against program interruptions
+ temp_dir = os.path.join(self.cache_dir, f"tmp.pid_{pid}_{rnd_id}")
+ os.makedirs(temp_dir, exist_ok=True)
+ temp_path = os.path.join(temp_dir, filename)
+
+ mode = "wb" if binary else "w"
+ with open(temp_path, mode) as f:
+ f.write(data)
+ # Replace is guaranteed to be atomic on POSIX systems if it succeeds
+ # so filepath cannot see a partial write
+ os.replace(temp_path, filepath)
+ os.removedirs(temp_dir)
+ return filepath
diff --git a/megatron/core/tensor_parallel/__init__.py b/megatron/core/tensor_parallel/__init__.py
index 4abec79c16..41d87431fe 100644
--- a/megatron/core/tensor_parallel/__init__.py
+++ b/megatron/core/tensor_parallel/__init__.py
@@ -1,36 +1,43 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
from .cross_entropy import vocab_parallel_cross_entropy
from .data import broadcast_data
-
from .layers import (
ColumnParallelLinear,
RowParallelLinear,
VocabParallelEmbedding,
- set_tensor_model_parallel_attributes,
- set_defaults_if_not_set_tensor_model_parallel_attributes,
copy_tensor_model_parallel_attributes,
+ linear_with_grad_accumulation_and_async_allreduce,
param_is_not_tensor_parallel_duplicate,
- linear_with_grad_accumulation_and_async_allreduce
-
+ set_defaults_if_not_set_tensor_model_parallel_attributes,
+ set_tensor_model_parallel_attributes,
)
-
from .mappings import (
+ all_gather_last_dim_from_tensor_parallel_region,
+ all_to_all,
+ all_to_all_hp2sp,
+ all_to_all_sp2hp,
copy_to_tensor_model_parallel_region,
- gather_from_tensor_model_parallel_region,
gather_from_sequence_parallel_region,
- scatter_to_tensor_model_parallel_region,
+ gather_from_sequence_parallel_region_to_moe,
+ gather_from_tensor_model_parallel_region,
+ reduce_from_tensor_model_parallel_region,
+ reduce_scatter_last_dim_to_tensor_parallel_region,
+ reduce_scatter_to_sequence_parallel_region,
+ reduce_scatter_to_sequence_parallel_region_from_moe,
scatter_to_sequence_parallel_region,
+ scatter_to_tensor_model_parallel_region,
)
-
from .random import (
checkpoint,
get_cuda_rng_tracker,
+ get_data_parallel_rng_tracker_name,
+ get_expert_parallel_rng_tracker_name,
model_parallel_cuda_manual_seed,
)
-
from .utils import (
+ gather_split_1d_tensor,
split_tensor_along_last_dim,
split_tensor_into_1d_equal_chunks,
- gather_split_1d_tensor,
)
__all__ = [
@@ -38,7 +45,7 @@
"vocab_parallel_cross_entropy",
# data.py
"broadcast_data",
- #layers.py
+ # layers.py
"ColumnParallelLinear",
"RowParallelLinear",
"VocabParallelEmbedding",
@@ -51,15 +58,19 @@
"copy_to_tensor_model_parallel_region",
"gather_from_tensor_model_parallel_region",
"gather_from_sequence_parallel_region",
-# "reduce_from_tensor_model_parallel_region",
+ "reduce_from_tensor_model_parallel_region",
+ "reduce_scatter_to_sequence_parallel_region",
"scatter_to_tensor_model_parallel_region",
"scatter_to_sequence_parallel_region",
# random.py
"checkpoint",
"get_cuda_rng_tracker",
"model_parallel_cuda_manual_seed",
+ "get_expert_parallel_rng_tracker_name",
# utils.py
"split_tensor_along_last_dim",
"split_tensor_into_1d_equal_chunks",
"gather_split_1d_tensor",
+ "gather_from_sequence_parallel_region_to_moe",
+ "reduce_scatter_to_sequence_parallel_region_from_moe",
]
diff --git a/megatron/core/tensor_parallel/cross_entropy.py b/megatron/core/tensor_parallel/cross_entropy.py
index 9147dbbadd..0066d126fd 100644
--- a/megatron/core/tensor_parallel/cross_entropy.py
+++ b/megatron/core/tensor_parallel/cross_entropy.py
@@ -1,36 +1,46 @@
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+from typing import Tuple
+
import torch
from megatron.core.parallel_state import (
get_tensor_model_parallel_group,
get_tensor_model_parallel_rank,
- get_tensor_model_parallel_world_size
+ get_tensor_model_parallel_world_size,
)
from .utils import VocabUtility
-class _VocabParallelCrossEntropy(torch.autograd.Function):
+class VocabParallelCrossEntropy:
+ """
+ Computes the Cross Entropy Loss splitting the Vocab size across tensor parallel
+ ranks. This implementation is used in both fused and unfused cross entropy implementations
+ """
@staticmethod
- def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0):
+ def calculate_logits_max(
+ vocab_parallel_logits: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ vocab_parallel_logits = vocab_parallel_logits.float()
# Maximum value along vocab dimension across all GPUs.
logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]
- torch.distributed.all_reduce(logits_max,
- op=torch.distributed.ReduceOp.MAX,
- group=get_tensor_model_parallel_group())
- # Subtract the maximum value.
- vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1)
- # Get the partition's vocab indecies
- get_vocab_range = VocabUtility.vocab_range_from_per_partition_vocab_size
- partition_vocab_size = vocab_parallel_logits.size()[-1]
- rank = get_tensor_model_parallel_rank()
- world_size = get_tensor_model_parallel_world_size()
- vocab_start_index, vocab_end_index = get_vocab_range(
- partition_vocab_size, rank, world_size)
+ return vocab_parallel_logits, logits_max
+
+ @staticmethod
+ def calculate_predicted_logits(
+ vocab_parallel_logits: torch.Tensor,
+ target: torch.Tensor,
+ logits_max: torch.Tensor,
+ vocab_start_index: int,
+ vocab_end_index: int,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+
+ # In-place subtraction reduces memory pressure.
+ vocab_parallel_logits -= logits_max.unsqueeze(dim=-1)
# Create a mask of valid vocab ids (1 means it needs to be masked).
target_mask = (target < vocab_start_index) | (target >= vocab_end_index)
@@ -40,26 +50,25 @@ def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0):
# Get predicted-logits = logits[target].
# For Simplicity, we convert logits to a 2-D tensor with size
# [*, partition-vocab-size] and target to a 1-D tensor of size [*].
+ partition_vocab_size = vocab_parallel_logits.size()[-1]
logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
masked_target_1d = masked_target.view(-1)
- arange_1d = torch.arange(start=0, end=logits_2d.size()[0],
- device=logits_2d.device)
+ arange_1d = torch.arange(start=0, end=logits_2d.size()[0], device=logits_2d.device)
predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
predicted_logits_1d = predicted_logits_1d.clone().contiguous()
predicted_logits = predicted_logits_1d.view_as(target)
predicted_logits[target_mask] = 0.0
- # All reduce is needed to get the chunks from other GPUs.
- torch.distributed.all_reduce(predicted_logits,
- op=torch.distributed.ReduceOp.SUM,
- group=get_tensor_model_parallel_group())
- # Sum of exponential of logits along vocab dimension across all GPUs.
exp_logits = vocab_parallel_logits
torch.exp(vocab_parallel_logits, out=exp_logits)
sum_exp_logits = exp_logits.sum(dim=-1)
- torch.distributed.all_reduce(sum_exp_logits,
- op=torch.distributed.ReduceOp.SUM,
- group=get_tensor_model_parallel_group())
+
+ return target_mask, masked_target_1d, predicted_logits, sum_exp_logits, exp_logits
+
+ @staticmethod
+ def calculate_cross_entropy_loss(
+ exp_logits: torch.Tensor, predicted_logits: torch.Tensor, sum_exp_logits: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
# Loss = log(sum(exp(logits))) - predicted-logit.
loss = torch.log(sum_exp_logits) - predicted_logits
@@ -67,6 +76,85 @@ def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0):
# Normalize and optionally smooth logits
exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))
+ return exp_logits, loss
+
+ @staticmethod
+ def prepare_gradient_calculation_operands(
+ softmax: torch.Tensor, target_mask: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+
+ # All the inputs have softmax as thier gradient.
+ grad_input = softmax
+ # For simplicity, work with the 2D gradient.
+ partition_vocab_size = softmax.size()[-1]
+ grad_2d = grad_input.view(-1, partition_vocab_size)
+
+ # Add the gradient from matching classes.
+ arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device)
+
+ softmax_update = 1.0 - target_mask.view(-1).float()
+
+ return grad_2d, arange_1d, softmax_update, grad_input
+
+ @staticmethod
+ def calculate_gradients(
+ grad_2d: torch.Tensor,
+ arange_1d: torch.Tensor,
+ masked_target_1d: torch.Tensor,
+ softmax_update: torch.Tensor,
+ grad_input: torch.Tensor,
+ grad_output: torch.Tensor,
+ ) -> torch.Tensor:
+
+ grad_2d[arange_1d, masked_target_1d] -= softmax_update
+
+ # Finally elementwise multiplication with the output gradients.
+ grad_input.mul_(grad_output.unsqueeze(dim=-1))
+
+ return grad_input
+
+
+class _VocabParallelCrossEntropy(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0):
+
+ vocab_parallel_logits, logits_max = VocabParallelCrossEntropy.calculate_logits_max(
+ vocab_parallel_logits
+ )
+ torch.distributed.all_reduce(
+ logits_max, op=torch.distributed.ReduceOp.MAX, group=get_tensor_model_parallel_group()
+ )
+
+ # Get the partition's vocab indices
+ get_vocab_range = VocabUtility.vocab_range_from_per_partition_vocab_size
+ partition_vocab_size = vocab_parallel_logits.size()[-1]
+ rank = get_tensor_model_parallel_rank()
+ world_size = get_tensor_model_parallel_world_size()
+ vocab_start_index, vocab_end_index = get_vocab_range(partition_vocab_size, rank, world_size)
+
+ (target_mask, masked_target_1d, predicted_logits, sum_exp_logits, exp_logits) = (
+ VocabParallelCrossEntropy.calculate_predicted_logits(
+ vocab_parallel_logits, target, logits_max, vocab_start_index, vocab_end_index
+ )
+ )
+
+ # All reduce is needed to get the chunks from other GPUs.
+ torch.distributed.all_reduce(
+ predicted_logits,
+ op=torch.distributed.ReduceOp.SUM,
+ group=get_tensor_model_parallel_group(),
+ )
+
+ torch.distributed.all_reduce(
+ sum_exp_logits,
+ op=torch.distributed.ReduceOp.SUM,
+ group=get_tensor_model_parallel_group(),
+ )
+
+ exp_logits, loss = VocabParallelCrossEntropy.calculate_cross_entropy_loss(
+ exp_logits, predicted_logits, sum_exp_logits
+ )
+
vocab_size = exp_logits.size(-1)
if label_smoothing > 0:
"""
@@ -87,7 +175,6 @@ def forward(ctx, vocab_parallel_logits, target, label_smoothing=0.0):
loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs
ctx.label_smoothing, ctx.vocab_size = label_smoothing, vocab_size
- ctx.save_for_backward(exp_logits, target_mask, masked_target_1d)
# Store softmax, target-mask and masked-target for backward pass.
ctx.save_for_backward(exp_logits, target_mask, masked_target_1d)
@@ -101,28 +188,22 @@ def backward(ctx, grad_output):
softmax, target_mask, masked_target_1d = ctx.saved_tensors
label_smoothing, vocab_size = ctx.label_smoothing, ctx.vocab_size
- # All the inputs have softmax as thier gradient.
- grad_input = softmax
- # For simplicity, work with the 2D gradient.
- partition_vocab_size = softmax.size()[-1]
- grad_2d = grad_input.view(-1, partition_vocab_size)
-
- # Add the gradient from matching classes.
- arange_1d = torch.arange(start=0, end=grad_2d.size()[0],
- device=grad_2d.device)
-
- softmax_update = 1.0 - target_mask.view(-1).float()
+ (grad_2d, arange_1d, softmax_update, grad_input) = (
+ VocabParallelCrossEntropy.prepare_gradient_calculation_operands(softmax, target_mask)
+ )
if label_smoothing > 0:
smoothing = label_smoothing * vocab_size / (vocab_size - 1)
grad_2d[arange_1d, masked_target_1d] -= (1.0 - smoothing) * softmax_update
average_grad = 1 / vocab_size
grad_2d[arange_1d, :] -= smoothing * average_grad
- else:
- grad_2d[arange_1d, masked_target_1d] -= softmax_update
- # Finally elementwise multiplication with the output gradients.
- grad_input.mul_(grad_output.unsqueeze(dim=-1))
+ # Finally elementwise multiplication with the output gradients.
+ grad_input.mul_(grad_output.unsqueeze(dim=-1))
+ else:
+ grad_input = VocabParallelCrossEntropy.calculate_gradients(
+ grad_2d, arange_1d, masked_target_1d, softmax_update, grad_input, grad_output
+ )
return grad_input, None, None
@@ -131,9 +212,9 @@ def vocab_parallel_cross_entropy(vocab_parallel_logits, target, label_smoothing=
"""
Performs cross entropy loss when logits are split across tensor parallel ranks
- Arguments:
+ Args:
vocab_parallel_logits: logits split across tensor parallel ranks
- dimension is [sequence_length, batch_size, hidden_size]
+ dimension is [sequence_length, batch_size, vocab_size/num_parallel_ranks]
target: correct vocab ids of dimseion [sequence_length, micro_batch_size]
diff --git a/megatron/core/tensor_parallel/data.py b/megatron/core/tensor_parallel/data.py
index b911790dae..c549f74d73 100644
--- a/megatron/core/tensor_parallel/data.py
+++ b/megatron/core/tensor_parallel/data.py
@@ -8,15 +8,17 @@
get_tensor_model_parallel_src_rank,
)
-
_MAX_DATA_DIM = 5
def _check_data_types(keys, data, target_dtype):
"""Check that all the keys have the same target data type."""
for key in keys:
- assert data[key].dtype == target_dtype, '{} has data type {} which '\
- 'is different than {}'.format(key, data[key].dtype, target_dtype)
+ assert (
+ data[key].dtype == target_dtype
+ ), '{} has data type {} which ' 'is different than {}'.format(
+ key, data[key].dtype, target_dtype
+ )
def _build_key_size_numel_dictionaries(keys, data):
@@ -35,9 +37,10 @@ def _build_key_size_numel_dictionaries(keys, data):
offset += max_dim
# Move to GPU and broadcast.
- sizes_cuda = torch.cuda.LongTensor(sizes)
- torch.distributed.broadcast(sizes_cuda, get_tensor_model_parallel_src_rank(),
- group=get_tensor_model_parallel_group())
+ sizes_cuda = torch.tensor(sizes, dtype=torch.long, device='cuda')
+ torch.distributed.broadcast(
+ sizes_cuda, get_tensor_model_parallel_src_rank(), group=get_tensor_model_parallel_group()
+ )
# Move back to cpu and unpack.
sizes_cpu = sizes_cuda.cpu()
@@ -66,7 +69,7 @@ def broadcast_data(keys, data, datatype):
"""Broadcast data from rank zero of each model parallel group to the
members of the same model parallel group.
- Arguments:
+ Args:
keys: list of keys in the data disctionary to be broadcasted
data: data dictionary of string keys and cpu tensor values.
datatype: torch data type of all tensors in data associated
@@ -74,24 +77,21 @@ def broadcast_data(keys, data, datatype):
"""
# Build (key, size) and (key, number of elements) dictionaries along
# with the total number of elements on all ranks.
- key_size, key_numel, total_numel = _build_key_size_numel_dictionaries(keys,
- data)
+ key_size, key_numel, total_numel = _build_key_size_numel_dictionaries(keys, data)
# Pack on rank zero.
if get_tensor_model_parallel_rank() == 0:
# Check that all keys have the same data type.
_check_data_types(keys, data, datatype)
# Flatten the data associated with the keys
- flatten_data = torch.cat(
- [data[key].contiguous().view(-1) for key in keys], dim=0).cuda()
+ flatten_data = torch.cat([data[key].contiguous().view(-1) for key in keys], dim=0).cuda()
else:
- flatten_data = torch.empty(total_numel,
- device=torch.cuda.current_device(),
- dtype=datatype)
+ flatten_data = torch.empty(total_numel, device=torch.cuda.current_device(), dtype=datatype)
# Broadcast
- torch.distributed.broadcast(flatten_data, get_tensor_model_parallel_src_rank(),
- group=get_tensor_model_parallel_group())
+ torch.distributed.broadcast(
+ flatten_data, get_tensor_model_parallel_src_rank(), group=get_tensor_model_parallel_group()
+ )
# Unpack
output = {}
diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py
index 15e0fbb025..903b4ed873 100644
--- a/megatron/core/tensor_parallel/layers.py
+++ b/megatron/core/tensor_parallel/layers.py
@@ -1,41 +1,40 @@
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
# Parts of the code here are adapted from PyTorch
# repo: https://github.com/pytorch/pytorch
-import math
import os
-from typing import Optional
import warnings
+from typing import Any, Callable, List, Optional, Tuple
import torch
import torch.nn.functional as F
-import torch.nn.init as init
+from torch.cuda.amp import custom_bwd, custom_fwd
from torch.nn.parameter import Parameter
-from torch.cuda.amp import custom_fwd, custom_bwd
-
+from megatron.core.model_parallel_config import ModelParallelConfig
from megatron.core.parallel_state import (
+ get_global_memory_buffer,
+ get_tensor_and_expert_parallel_rank,
+ get_tensor_and_expert_parallel_world_size,
+ get_tensor_model_parallel_group,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
- get_tensor_model_parallel_group,
- get_global_memory_buffer,
)
+
+from ..dist_checkpointing.mapping import ShardedStateDict
+from ..transformer.utils import make_sharded_tensors_for_checkpoint
+from ..utils import make_tp_sharded_tensor_for_checkpoint, prepare_input_tensors_for_wgrad_compute
from .mappings import (
copy_to_tensor_model_parallel_region,
- gather_from_tensor_model_parallel_region,
gather_from_sequence_parallel_region,
+ gather_from_tensor_model_parallel_region,
reduce_from_tensor_model_parallel_region,
- scatter_to_tensor_model_parallel_region,
reduce_scatter_to_sequence_parallel_region,
+ scatter_to_tensor_model_parallel_region,
)
-
-from .random import get_cuda_rng_tracker
-from .utils import (
- divide,
- split_tensor_along_last_dim,
- VocabUtility,
-)
+from .random import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
+from .utils import VocabUtility, divide
_grad_accum_fusion_available = True
try:
@@ -43,17 +42,23 @@
except ImportError:
_grad_accum_fusion_available = False
-_MODEL_PARALLEL_ATTRIBUTE_DEFAULTS = {'tensor_model_parallel': False,
- 'partition_dim': -1,
- 'partition_stride': 1}
+_MODEL_PARALLEL_ATTRIBUTE_DEFAULTS = {
+ 'tensor_model_parallel': False,
+ 'partition_dim': -1,
+ 'partition_stride': 1,
+}
+
def param_is_not_tensor_parallel_duplicate(param):
- return (hasattr(param, 'tensor_model_parallel') and
- param.tensor_model_parallel) or (
- get_tensor_model_parallel_rank() == 0)
+ """Returns true if the passed-in parameter is not a duplicate parameter
+ on another TP rank."""
+ return (hasattr(param, 'tensor_model_parallel') and param.tensor_model_parallel) or (
+ get_tensor_model_parallel_rank() == 0
+ )
def set_tensor_model_parallel_attributes(tensor, is_parallel, dim, stride):
+ """Sets tp attributes to tensor"""
# Make sure the attributes are not set.
for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
assert not hasattr(tensor, attribute)
@@ -64,67 +69,85 @@ def set_tensor_model_parallel_attributes(tensor, is_parallel, dim, stride):
def set_defaults_if_not_set_tensor_model_parallel_attributes(tensor):
+ """Set default model parallel attributes if not set explicitly already."""
+
def maybe_set(attribute, value):
if not hasattr(tensor, attribute):
setattr(tensor, attribute, value)
+
for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
maybe_set(attribute, _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS[attribute])
def copy_tensor_model_parallel_attributes(destination_tensor, source_tensor):
+ """Copy model parallel attributes from one tensor to another."""
+
def maybe_copy(attribute):
if hasattr(source_tensor, attribute):
- setattr(destination_tensor, attribute,
- getattr(source_tensor, attribute))
+ setattr(destination_tensor, attribute, getattr(source_tensor, attribute))
+
for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
maybe_copy(attribute)
-def _initialize_affine_weight_gpu(weight, init_method,
- partition_dim, stride=1):
+def _initialize_affine_weight_gpu(
+ weight, init_method, partition_dim, stride=1, expert_parallel=False
+):
"""Initialize affine weight for model parallel on GPU."""
- set_tensor_model_parallel_attributes(tensor=weight,
- is_parallel=True,
- dim=partition_dim,
- stride=stride)
-
- with get_cuda_rng_tracker().fork():
- init_method(weight)
-
-
-def _initialize_affine_weight_cpu(weight, output_size, input_size,
- per_partition_size, partition_dim,
- init_method, stride=1,
- return_master_weight=False,
- *, params_dtype=torch.float32):
+ set_tensor_model_parallel_attributes(
+ tensor=weight, is_parallel=True, dim=partition_dim, stride=stride
+ )
+
+ if not expert_parallel:
+ with get_cuda_rng_tracker().fork():
+ init_method(weight)
+ else:
+ with get_cuda_rng_tracker().fork(get_expert_parallel_rng_tracker_name()):
+ init_method(weight)
+
+
+def _initialize_affine_weight_cpu(
+ weight,
+ output_size,
+ input_size,
+ per_partition_size,
+ partition_dim,
+ init_method,
+ stride=1,
+ return_master_weight=False,
+ *,
+ params_dtype=torch.float32,
+ rank=None,
+ world_size=None,
+ skip_set_tensor_parallel_attributes=False,
+):
"""Initialize affine weight for model parallel.
Build the master weight on all processes and scatter
the relevant chunk."""
- set_tensor_model_parallel_attributes(tensor=weight,
- is_parallel=True,
- dim=partition_dim,
- stride=stride)
+ if not skip_set_tensor_parallel_attributes:
+ set_tensor_model_parallel_attributes(
+ tensor=weight, is_parallel=True, dim=partition_dim, stride=stride
+ )
# Initialize master weight
- master_weight = torch.empty(output_size, input_size,
- dtype=torch.float,
- requires_grad=False)
+ master_weight = torch.empty(output_size, input_size, dtype=torch.float, requires_grad=False)
init_method(master_weight)
master_weight = master_weight.to(dtype=params_dtype)
-
# Split and copy
per_partition_per_stride_size = divide(per_partition_size, stride)
- weight_list = torch.split(master_weight, per_partition_per_stride_size,
- dim=partition_dim)
- rank = get_tensor_model_parallel_rank()
- world_size = get_tensor_model_parallel_world_size()
+ weight_list = torch.split(master_weight, per_partition_per_stride_size, dim=partition_dim)
+ if rank is None:
+ rank = get_tensor_model_parallel_rank()
+ world_size = get_tensor_model_parallel_world_size()
my_weight_list = weight_list[rank::world_size]
with torch.no_grad():
- torch.cat(my_weight_list, dim=partition_dim, out=weight)
+ # all tensors must live on the same device
+ cpu_weight = torch.cat(my_weight_list, dim=partition_dim).to_dense()
+ weight.data.copy_(cpu_weight)
if return_master_weight:
return master_weight
return None
@@ -135,107 +158,267 @@ class VocabParallelEmbedding(torch.nn.Module):
This is mainly adapted from torch.nn.Embedding and all the default
values are kept.
- Arguments:
+
+ Args:
num_embeddings: vocabulary size.
embedding_dim: size of hidden state.
+ reduce_scatter_embeddings: Decides whether to perform ReduceScatter after embedding lookup
- Keyword Arguments:
- init_method: method to initialize weights.
- params_dtype
- use_cpu_initialization
- perform_initialization
+ Keyword Args:
+ config: A megatron.core.ModelParallelConfig object
"""
- def __init__(self, num_embeddings: int, embedding_dim: int, *,
- init_method=init.xavier_normal_,
- params_dtype: torch.dtype=torch.float32,
- use_cpu_initialization: bool=False,
- perform_initialization: bool=True):
+ def __init__(
+ self,
+ num_embeddings: int,
+ embedding_dim: int,
+ *,
+ init_method: Callable,
+ reduce_scatter_embeddings: bool = False,
+ config: ModelParallelConfig,
+ ):
super(VocabParallelEmbedding, self).__init__()
# Keep the input dimensions.
self.num_embeddings = num_embeddings
self.embedding_dim = embedding_dim
- # Set the detauls for compatibility.
- self.padding_idx = None
- self.max_norm = None
- self.norm_type = 2.
- self.scale_grad_by_freq = False
- self.sparse = False
- self._weight = None
+ self.reduce_scatter_embeddings = reduce_scatter_embeddings
self.tensor_model_parallel_size = get_tensor_model_parallel_world_size()
# Divide the weight matrix along the vocaburaly dimension.
- self.vocab_start_index, self.vocab_end_index = \
+ (self.vocab_start_index, self.vocab_end_index) = (
VocabUtility.vocab_range_from_global_vocab_size(
- self.num_embeddings, get_tensor_model_parallel_rank(),
- self.tensor_model_parallel_size)
- self.num_embeddings_per_partition = self.vocab_end_index - \
- self.vocab_start_index
+ self.num_embeddings,
+ get_tensor_model_parallel_rank(),
+ self.tensor_model_parallel_size,
+ )
+ )
+ self.num_embeddings_per_partition = self.vocab_end_index - self.vocab_start_index
+ self.deterministic_mode = config.deterministic_mode
# Allocate weights and initialize.
- if use_cpu_initialization:
- self.weight = Parameter(torch.empty(
- self.num_embeddings_per_partition, self.embedding_dim,
- dtype=params_dtype))
- if perform_initialization:
+ if config.use_cpu_initialization:
+ self.weight = Parameter(
+ torch.empty(
+ self.num_embeddings_per_partition, self.embedding_dim, dtype=config.params_dtype
+ )
+ )
+ if config.perform_initialization:
_initialize_affine_weight_cpu(
- self.weight, self.num_embeddings, self.embedding_dim,
- self.num_embeddings_per_partition, 0, init_method,
- params_dtype=params_dtype)
+ self.weight,
+ self.num_embeddings,
+ self.embedding_dim,
+ self.num_embeddings_per_partition,
+ 0,
+ init_method,
+ params_dtype=config.params_dtype,
+ )
else:
- self.weight = Parameter(torch.empty(
- self.num_embeddings_per_partition, self.embedding_dim,
- device=torch.cuda.current_device(), dtype=params_dtype))
- if perform_initialization:
- _initialize_affine_weight_gpu(self.weight, init_method,
- partition_dim=0, stride=1)
+ self.weight = Parameter(
+ torch.empty(
+ self.num_embeddings_per_partition,
+ self.embedding_dim,
+ device=torch.cuda.current_device(),
+ dtype=config.params_dtype,
+ )
+ )
+ if config.perform_initialization:
+ _initialize_affine_weight_gpu(self.weight, init_method, partition_dim=0, stride=1)
def forward(self, input_):
+ """Forward.
+
+ Args:
+ input_ (torch.Tensor): Input tensor.
+ """
if self.tensor_model_parallel_size > 1:
# Build the mask.
- input_mask = (input_ < self.vocab_start_index) | \
- (input_ >= self.vocab_end_index)
+ input_mask = (input_ < self.vocab_start_index) | (input_ >= self.vocab_end_index)
# Mask the input.
masked_input = input_.clone() - self.vocab_start_index
masked_input[input_mask] = 0
else:
masked_input = input_
- # Get the embeddings.
- output_parallel = F.embedding(masked_input, self.weight,
- self.padding_idx, self.max_norm,
- self.norm_type, self.scale_grad_by_freq,
- self.sparse)
+ # Get the embeddings.
+ if self.deterministic_mode:
+ output_parallel = self.weight[masked_input]
+ else:
+ # F.embedding currently has a non-deterministic backward function
+ output_parallel = F.embedding(masked_input, self.weight)
# Mask the output embedding.
if self.tensor_model_parallel_size > 1:
output_parallel[input_mask, :] = 0.0
- # Reduce across all the model parallel GPUs.
- output = reduce_from_tensor_model_parallel_region(output_parallel)
+
+ if self.reduce_scatter_embeddings:
+ # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
+ output_parallel = output_parallel.transpose(0, 1).contiguous()
+ output = reduce_scatter_to_sequence_parallel_region(output_parallel)
+ else:
+ # Reduce across all the model parallel GPUs.
+ output = reduce_from_tensor_model_parallel_region(output_parallel)
return output
+ def sharded_state_dict(
+ self,
+ prefix: str = '',
+ sharded_offsets: Tuple[Tuple[int, int, int]] = (),
+ metadata: Optional[dict] = None,
+ ) -> ShardedStateDict:
+ """Non-default implementation for embeddings due to `allow_shape_mismatch` param"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+
+ weight_prefix = f'{prefix}weight'
+ return {
+ weight_prefix: make_tp_sharded_tensor_for_checkpoint(
+ tensor=state_dict['weight'],
+ key=weight_prefix,
+ allow_shape_mismatch=True,
+ prepend_offsets=sharded_offsets,
+ )
+ }
+
+
+class LinearWithFrozenWeight(torch.autograd.Function):
+ """Linear operator that does not calculate gradient for weight.
+ This op and LinearWithGradAccumulationAndAsyncCommunication performs
+ mathematically-identical forward and DGRAD.
+
+ Conceptually this op is the same as torch.nn.functional.linear with
+ weight.requires_grad==False, but in experiments they are not identical
+ mathematically."""
+
+ @staticmethod
+ @custom_fwd
+ def forward(ctx, input, weight, bias, allreduce_dgrad):
+ """Forward with frozen weight."""
+ ctx.save_for_backward(weight)
+ ctx.allreduce_dgrad = allreduce_dgrad
+ output = torch.matmul(input, weight.t())
+ if bias is not None:
+ output = output + bias
+ return output
+
+ @staticmethod
+ @custom_bwd
+ def backward(ctx, grad_output):
+ """Backward with frozen weight."""
+ (weight,) = ctx.saved_tensors
+ grad_input = grad_output.matmul(weight)
+
+ if ctx.allreduce_dgrad:
+ # All-reduce. Note: here async and sync are effectively the same.
+ torch.distributed.all_reduce(grad_input, group=get_tensor_model_parallel_group())
+
+ return grad_input, None, None, None
+
+
+def linear_with_frozen_weight(
+ input: torch.Tensor,
+ weight: torch.Tensor,
+ bias: Optional[torch.Tensor],
+ gradient_accumulation_fusion: bool,
+ async_grad_allreduce: bool,
+ sequence_parallel: bool,
+ grad_output_buffer: Optional[List[torch.Tensor]] = None,
+ wgrad_deferral_limit: None = None,
+ allreduce_dgrad: bool = None,
+) -> torch.Tensor:
+ """Linear layer execution with weight.requires_grad == False.
+
+ This function handles linear layers with weight frozen (untrainable).
+ In the forward, it only saves weight and does not save input activations.
+ In the backward, it does not perform weight gradient calculation, or
+ weight gradient allreduce.
+
+ Args:
+
+ input (torch.Tensor required): input like torch.nn.functional.linear
+
+ weight (torch.Tensor required): weight like torch.nn.functional.linear
+
+ bias (torch.Tensor optional): bias like torch.nn.functional.linear
+
+ gradient_accumulation_fusion (bool required): dummy argument, used to
+ keep the API unified between all forward implementation functions.
+
+ async_grad_allreduce (bool required): dummy argument, used to
+ keep the API unified between all forward implementation functions.
+
+ sequence_parallel (bool required): Indicates that sequence
+ parallelism is used and thus in the forward pass the input is
+ all gathered, and the backward pass the input gradients are
+ reduce scattered.
+
+ grad_output_buffer (List[torch.Tensor] optional): dummy argument, used to
+ keep the API unified between all forward implementation functions.
+
+ wgrad_deferral_limit (int optional): dummy argument, used to
+ keep the API unified between all forward implementation functions.
+
+ allreduce_dgrad (bool): Do the allreduce of input gradients.
+ Here, async and sync allreduce are the same. If sequence_parallel is
+ True, this must be False, as no all reduce is performed.
+
+ """
+
+ assert grad_output_buffer is None, (
+ "grad_output_buffer kwarg is only supported with "
+ "linear_with_grad_accumulation_and_async_allreduce"
+ )
+
+ assert wgrad_deferral_limit is None, (
+ "This arg is only supported with " "linear_with_grad_accumulation_and_async_allreduce"
+ )
+
+ if sequence_parallel:
+ input = gather_from_sequence_parallel_region(input, tensor_parallel_output_grad=True)
+ else:
+ input = input
+
+ if allreduce_dgrad is None:
+ warnings.warn(
+ "`async_grad_allreduce` is deprecated and will be removed in a future release. "
+ "Please ue `allreduce_dgrad` instead."
+ )
+ allreduce_dgrad = async_grad_allreduce
+
+ args = [input, weight, bias, allreduce_dgrad]
+
+ return LinearWithFrozenWeight.apply(*args)
+
class LinearWithGradAccumulationAndAsyncCommunication(torch.autograd.Function):
"""See linear_with_grad_accumulation_and_async_allreduce"""
@staticmethod
@custom_fwd
- def forward(ctx, input, weight, bias, gradient_accumulation_fusion,
- async_grad_allreduce, sequence_parallel):
+ def forward(
+ ctx,
+ input,
+ weight,
+ bias,
+ gradient_accumulation_fusion,
+ allreduce_dgrad,
+ sequence_parallel,
+ grad_output_buffer,
+ wgrad_deferral_limit,
+ ):
+ """Forward."""
ctx.save_for_backward(input, weight)
ctx.use_bias = bias is not None
ctx.gradient_accumulation_fusion = gradient_accumulation_fusion
- ctx.async_grad_allreduce = async_grad_allreduce
+ ctx.allreduce_dgrad = allreduce_dgrad
ctx.sequence_parallel = sequence_parallel
+ ctx.wgrad_deferral_limit = wgrad_deferral_limit
+ ctx.grad_output_buffer = grad_output_buffer
if sequence_parallel:
world_size = get_tensor_model_parallel_world_size()
dim_size = list(input.size())
dim_size[0] = dim_size[0] * world_size
- all_gather_buffer = \
- get_global_memory_buffer().get_tensor(dim_size, input.dtype, "mpu")
+ all_gather_buffer = get_global_memory_buffer().get_tensor(dim_size, input.dtype, "mpu")
torch.distributed._all_gather_base(
- all_gather_buffer,
- input,
- group=get_tensor_model_parallel_group())
+ all_gather_buffer, input, group=get_tensor_model_parallel_group()
+ )
total_input = all_gather_buffer
else:
total_input = input
@@ -248,91 +431,128 @@ def forward(ctx, input, weight, bias, gradient_accumulation_fusion,
@staticmethod
@custom_bwd
def backward(ctx, grad_output):
+ """Backward."""
input, weight = ctx.saved_tensors
use_bias = ctx.use_bias
+ grad_output_buffer = ctx.grad_output_buffer
+ wgrad_deferral_limit = ctx.wgrad_deferral_limit
+
+ wgrad_compute = True
+ if grad_output_buffer is not None:
+ if wgrad_deferral_limit == 0 or len(grad_output_buffer) < wgrad_deferral_limit:
+ grad_output_buffer.append(grad_output)
+ wgrad_compute = False
+
+ if wgrad_compute:
+ if ctx.sequence_parallel:
+ world_size = get_tensor_model_parallel_world_size()
+ dim_size = list(input.size())
+ dim_size[0] = dim_size[0] * world_size
+
+ all_gather_buffer = get_global_memory_buffer().get_tensor(
+ dim_size, input.dtype, "mpu"
+ )
+ handle = torch.distributed._all_gather_base(
+ all_gather_buffer, input, group=get_tensor_model_parallel_group(), async_op=True
+ )
- if ctx.sequence_parallel:
- world_size = get_tensor_model_parallel_world_size()
- dim_size = list(input.size())
- dim_size[0] = dim_size[0] * world_size
-
- all_gather_buffer = \
- get_global_memory_buffer().get_tensor(dim_size, input.dtype, "mpu")
- handle = torch.distributed._all_gather_base(
- all_gather_buffer,
- input,
- group=get_tensor_model_parallel_group(), async_op=True)
-
- # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
- # gather is scheduled before the input gradient computation
- total_input = all_gather_buffer
- else:
- total_input = input
+ # Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
+ # gather is scheduled before the input gradient computation
+ total_input = all_gather_buffer
+ else:
+ total_input = input
grad_input = grad_output.matmul(weight)
- if ctx.sequence_parallel:
+ if ctx.sequence_parallel and wgrad_compute:
handle.wait()
- # Doing gather + slicing during the NeMo forward pass can make this tensor
- # not be contiguous. PyTorch only checks if the tensor is contiguous, and only
- # clones it if it's not contiguous:
- # https://github.com/pytorch/pytorch/blob/c47cf9bc7f9e02f649ab4ed53fe4d35732c92ab6/torch/_refs/__init__.py#L2761
- grad_output = grad_output.contiguous()
- # Convert the tensor shapes to 2D for execution compatibility
- grad_output = grad_output.view(grad_output.shape[0] * grad_output.shape[1],
- grad_output.shape[2])
- total_input = total_input.view(total_input.shape[0] * total_input.shape[1],
- total_input.shape[2])
-
- if ctx.async_grad_allreduce:
+ if wgrad_compute:
+ grad_output, total_input = prepare_input_tensors_for_wgrad_compute(
+ grad_output, total_input
+ )
+
+ if ctx.allreduce_dgrad:
# Asynchronous all-reduce
handle = torch.distributed.all_reduce(
- grad_input, group=get_tensor_model_parallel_group(), async_op=True)
+ grad_input, group=get_tensor_model_parallel_group(), async_op=True
+ )
# Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
# all-reduce is scheduled before the weight gradient computation
if ctx.sequence_parallel:
- assert not ctx.async_grad_allreduce
+ assert not ctx.allreduce_dgrad
dim_size = list(input.size())
- sub_grad_input = torch.empty(dim_size, dtype=input.dtype,
- device=torch.cuda.current_device(),
- requires_grad=False)
+ sub_grad_input = torch.empty(
+ dim_size, dtype=input.dtype, device=torch.cuda.current_device(), requires_grad=False
+ )
# reduce_scatter
- handle = torch.distributed._reduce_scatter_base(sub_grad_input, grad_input,
- group=get_tensor_model_parallel_group(),
- async_op=True)
+ handle = torch.distributed._reduce_scatter_base(
+ sub_grad_input, grad_input, group=get_tensor_model_parallel_group(), async_op=True
+ )
# Here we rely on CUDA_DEVICE_MAX_CONNECTIONS=1 to ensure that the
# reduce scatter is scheduled before the weight gradient computation
-
if ctx.gradient_accumulation_fusion:
- if weight.main_grad.dtype == torch.float32:
- fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(total_input, grad_output, weight.main_grad)
- elif weight.main_grad.dtype in (torch.float16, torch.bfloat16):
- fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16(total_input, grad_output, weight.main_grad)
+ if wgrad_compute:
+ if weight.main_grad.dtype == torch.float32:
+ fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(
+ total_input, grad_output, weight.main_grad
+ )
+ elif weight.main_grad.dtype in (torch.float16, torch.bfloat16):
+ fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16(
+ total_input, grad_output, weight.main_grad
+ )
+ else:
+ raise RuntimeError("Unsupported gradient type for gradient accumulation fusion")
+
+ if hasattr(weight, 'grad_added_to_main_grad'):
+ # When overlap_grad_reduce is True, need to ensure that backward hooks
+ # are all run on the main backprop thread to prevent deadlocks. Setup
+ # dummy grad_weight tensor to prevent backward hooks from being run
+ # in a background thread.
+ if getattr(weight, 'zero_out_wgrad', False):
+ grad_weight = torch.zeros(
+ weight.main_grad.shape,
+ dtype=input.dtype,
+ device=torch.cuda.current_device(),
+ requires_grad=False,
+ )
+ else:
+ grad_weight = torch.empty(
+ weight.main_grad.shape,
+ dtype=input.dtype,
+ device=torch.cuda.current_device(),
+ requires_grad=False,
+ )
+ weight.grad_added_to_main_grad = True
else:
- raise RuntimeError("Unsupported gradient type for gradient accumulation fusion")
- grad_weight = None
+ grad_weight = None
else:
grad_weight = grad_output.t().matmul(total_input)
grad_bias = grad_output.sum(dim=0) if use_bias else None
if ctx.sequence_parallel:
handle.wait()
- return sub_grad_input, grad_weight, grad_bias, None, None, None
+ # Need to return None's as gradient has to flow for all the input arguments
+ # provided during forward
+ return sub_grad_input, grad_weight, grad_bias, None, None, None, None, None
- if ctx.async_grad_allreduce:
+ if ctx.allreduce_dgrad:
handle.wait()
- return grad_input, grad_weight, grad_bias, None, None, None
+ return grad_input, grad_weight, grad_bias, None, None, None, None, None
+
def linear_with_grad_accumulation_and_async_allreduce(
input: torch.Tensor,
weight: torch.Tensor,
bias: Optional[torch.Tensor],
gradient_accumulation_fusion: bool,
- async_grad_allreduce: bool,
- sequence_parallel_enabled: bool,
+ sequence_parallel: bool,
+ allreduce_dgrad: bool,
+ async_grad_allreduce: Optional[bool] = None,
+ grad_output_buffer: Optional[List[torch.Tensor]] = None,
+ wgrad_deferral_limit: Optional[int] = 0,
) -> torch.Tensor:
"""Linear layer execution with asynchronous communication and
gradient accumulation fusion in backprop.
@@ -359,105 +579,155 @@ def linear_with_grad_accumulation_and_async_allreduce(
CUDA_DEVICE_MAX_CONNECTIONS=1 forces the kernels to be scheduled
in the order they are called.
- Arguments:
+ Args:
+ input (torch.Tensor required): input like torch.nn.functional.linear
- input (torch.Tensor required): input like torch.nn.functional.linear
+ weight (torch.Tensor required): weight like torch.nn.functional.linear
- weight (torch.Tensor required): weight like torch.nn.functional.linear
+ bias (torch.Tensor optional): bias like torch.nn.functional.linear
- bias (torch.Tensor optional): bias like torch.nn.functional.linear
+ gradient_accumulation_fusion (bool required): Perform the gradient
+ accumulation fusion, requires the custom CUDA extension
+ fused_weight_gradient_mlp_cuda module. To use
+ gradient_accumulation_fusion you must install APEX with
+ --cpp_ext and --cuda_ext. For example: "pip install
+ --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext .\"
+ " Note that the extension requires CUDA>=11. Otherwise, you
+ must turn off gradient accumulation fusion."
+
+ allreduce_dgrad (bool required): Do the allreduce of input gradients.
+ The allreduce is done asynchronously with the computation of weight
+ gradients. If sequence_parallel is True, this must be
+ False, as no all reduce is performed.
+
+ async_grad_allreduce (bool optional): Do the allreduce of input
+ gradients asyncronously with the computation of weight
+ gradients. If sequence_parallel is True, this must be
+ False, as no all reduce is performed. Will be deprecated with 0.10.0
+
+ sequence_parallel (bool required): Indicates that sequence
+ parallelism is used and thus in the forward pass the input is
+ all gathered, and the backward pass the input gradients are
+ reduce scattered.
+
+ grad_output_buffer (List[torch.Tensor] optional): Buffer used to save
+ output gradients when embedding table wgrad compute is deferred.
+ Defaults to None.
+
+ wgrad_deferral_limit (int optional): Limit on the number of
+ micro-batches for which embedding weight gradient GEMM should be
+ deferred. Disable by setting this to 0. Defaults to 0.
- gradient_accumulation_fusion (bool required): Perform the gradient
- accumulation fusion, requires the custom CUDA extension
- fused_weight_gradient_mlp_cuda module. To use
- gradient_accumulation_fusion you must install APEX with
- --cpp_ext and --cuda_ext. For example: "pip install
- --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext .\"
- " Note that the extension requires CUDA>=11. Otherwise, you
- must turn off gradient accumulation fusion."
-
- async_grad_allreduce (bool required): Do the allreduce of input
- gradients asyncronously with the computation of weight
- gradients. If sequence_parallel_enabled is True, this must be
- False, as no all reduce is performed.
-
- sequence_parallel_enabled (bool required): Indicates that sequence
- parallelism is used and thus in the forward pass the input is
- all gathered, and the backward pass the input gradients are
- reduce scattered.
"""
+ if async_grad_allreduce is not None:
+ warnings.warn(
+ "async_grad_allreduce is deprecated, not in use anymore and will"
+ " be fully removed with 0.10.0. Please use allreduce_dgrad instead."
+ )
+
args = [
input,
weight,
bias,
gradient_accumulation_fusion,
- async_grad_allreduce,
- sequence_parallel_enabled,
+ allreduce_dgrad,
+ sequence_parallel,
+ grad_output_buffer,
+ wgrad_deferral_limit,
]
if not linear_with_grad_accumulation_and_async_allreduce.warned:
if os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1":
- if sequence_parallel_enabled:
+ if sequence_parallel:
warnings.warn(
"When using sequence parallelism it is recommended to set the "
"environment variable CUDA_DEVICE_MAX_CONNECTIONS to 1 for "
- "maximum speedup")
+ "maximum speedup"
+ )
linear_with_grad_accumulation_and_async_allreduce.warned = True
- if async_grad_allreduce:
+ if allreduce_dgrad:
warnings.warn(
"When using async grad allreduce it is recommended to set the "
"environment variable CUDA_DEVICE_MAX_CONNECTIONS to 1 for "
- "maximum speedup")
+ "maximum speedup"
+ )
linear_with_grad_accumulation_and_async_allreduce.warned = True
return LinearWithGradAccumulationAndAsyncCommunication.apply(*args)
+
linear_with_grad_accumulation_and_async_allreduce.warned = False
+
class ColumnParallelLinear(torch.nn.Module):
"""Linear layer with column parallelism.
The linear layer is defined as Y = XA + b. A is parallelized along
its second dimension as A = [A_1, ..., A_p].
- Arguments:
- input_size: first dimension of matrix A.
- output_size: second dimension of matrix A.
-
- Keyword Arguments
- bias: If true, add bias
- gather_output: If true, call all-gather on output and make Y available
- to all GPUs, otherwise, every GPU will have its output
- which is Y_i = XA_i
- init_method: method to initialize weights. Note that bias is always set
- to zero.
- stride: For the strided linear layers.
- keep_master_weight_for_test: This was added for testing and should be
- set to False. It returns the master weights
- used for initialization.
- skip_bias_add: This was added to enable performance optimations where bias
- can be fused with other elementwise operations. we skip
- adding bias but instead return it.
- async_tensor_model_parallel_allreduce:
- params_dtype:
- use_cpu_initialization:
- gradient_accumulation_fusion:
- sequence_parallel_enabled:
+ Args:
+ input_size:
+ first dimension of matrix A.
+ output_size:
+ second dimension of matrix A.
+ bias:
+ If true, add bias
+ gather_output:
+ If true, call all-gather on output and make Y available to all GPUs,
+ otherwise, every GPU will have its output which is Y_i = XA_i
+ init_method:
+ method to initialize weights. Note that bias is always set to zero.
+ stride:
+ For the strided linear layers.
+ keep_master_weight_for_test:
+ This was added for testing and should be set to False. It
+ returns the master weights used for initialization.
+ skip_bias_add:
+ If True, do not add the bias term, instead return it to be added by the
+ caller. This enables performance optimations where bias can be fused with other
+ elementwise operations.
+ skip_weight_param_allocation:
+ If True, weight parameter is not allocated and must be passed
+ as a keyword argument `weight` during the forward pass. Note that this does not
+ affect bias, which will be allocated if bias is True. Defaults to False.
+ embedding_activation_buffer:
+ This buffer holds the input activations of the final embedding
+ linear layer on the last pipeline stage when defer_embedding_wgrad_compute is enabled.
+ grad_output_buffer:
+ This buffer holds the gradient outputs of the final embedding linear
+ layer on the last pipeline stage when defer_embedding_wgrad_compute is enabled.
+ is_expert:
+ If True, the layer is treated as an MoE expert layer.
+ config:
+ ModelParallelConfig object
+ tp_comm_buffer_name:
+ Communication buffer name is not used in non-Transformer-Engine modules.
+ disable_grad_reduce:
+ If True, reduction of output gradients across tensor-parallel ranks
+ will be disabled. Defaults to False. This feature is used by Lora Adapter in Nemo to
+ delay and fuse reduction along with other gradients for performance optimization.
"""
- def __init__(self, input_size, output_size, *,
- bias=True, gather_output=True,
- init_method=init.xavier_normal_, stride=1,
- keep_master_weight_for_test=False,
- skip_bias_add=False,
- async_tensor_model_parallel_allreduce=True,
- params_dtype=torch.float32,
- use_cpu_initialization=False,
- perform_initialization=True,
- gradient_accumulation_fusion=False,
- sequence_parallel_enabled: bool = False,
- ):
+ def __init__(
+ self,
+ input_size,
+ output_size,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias=True,
+ gather_output=False,
+ stride=1,
+ keep_master_weight_for_test=False,
+ skip_bias_add=False,
+ skip_weight_param_allocation: bool = False,
+ embedding_activation_buffer: Optional[List[torch.Tensor]] = None,
+ grad_output_buffer: Optional[List[torch.Tensor]] = None,
+ is_expert: bool = False,
+ tp_comm_buffer_name: str = None, # Not used
+ disable_grad_reduce: bool = False,
+ ):
super(ColumnParallelLinear, self).__init__()
# Keep input parameters
@@ -465,216 +735,401 @@ def __init__(self, input_size, output_size, *,
self.output_size = output_size
self.gather_output = gather_output
# Divide the weight matrix along the last dimension.
- world_size = get_tensor_model_parallel_world_size()
- self.output_size_per_partition = divide(output_size, world_size)
self.skip_bias_add = skip_bias_add
+ self.is_expert = is_expert
+ self.expert_parallel = config.expert_model_parallel_size > 1
+ self.embedding_activation_buffer = embedding_activation_buffer
+ self.grad_output_buffer = grad_output_buffer
+ self.config = config
+ self.disable_grad_reduce = disable_grad_reduce
+
+ self.explicit_expert_comm = self.is_expert and (
+ config.tensor_model_parallel_size > 1 or self.expert_parallel
+ )
+ if self.explicit_expert_comm and config.moe_extended_tp:
+ world_size = get_tensor_and_expert_parallel_world_size()
+ rank = get_tensor_and_expert_parallel_rank()
+ else:
+ world_size = get_tensor_model_parallel_world_size()
+ rank = get_tensor_model_parallel_rank()
+
+ self.output_size_per_partition = divide(output_size, world_size)
# Parameters.
# Note: torch.nn.functional.linear performs XA^T + b and as a result
# we allocate the transpose.
# Initialize weight.
- if use_cpu_initialization:
- self.weight = Parameter(torch.empty(self.output_size_per_partition,
- self.input_size,
- dtype=params_dtype))
- if perform_initialization:
- self.master_weight = _initialize_affine_weight_cpu(
- self.weight, self.output_size, self.input_size,
- self.output_size_per_partition, 0, init_method,
- stride=stride, return_master_weight=keep_master_weight_for_test)
+ if not skip_weight_param_allocation:
+ if config.use_cpu_initialization:
+ self.weight = Parameter(
+ torch.empty(
+ self.output_size_per_partition, self.input_size, dtype=config.params_dtype
+ )
+ )
+ if config.perform_initialization:
+ self.master_weight = _initialize_affine_weight_cpu(
+ self.weight,
+ self.output_size,
+ self.input_size,
+ self.output_size_per_partition,
+ 0,
+ init_method,
+ stride=stride,
+ return_master_weight=keep_master_weight_for_test,
+ rank=rank,
+ world_size=world_size,
+ )
+ else:
+ self.weight = Parameter(
+ torch.empty(
+ self.output_size_per_partition,
+ self.input_size,
+ device=torch.cuda.current_device(),
+ dtype=config.params_dtype,
+ )
+ )
+ if config.perform_initialization:
+ _initialize_affine_weight_gpu(
+ self.weight,
+ init_method,
+ partition_dim=0,
+ stride=stride,
+ expert_parallel=(self.is_expert and self.expert_parallel),
+ )
+
+ setattr(self.weight, 'allreduce', not (self.is_expert and self.expert_parallel))
else:
- self.weight = Parameter(torch.empty(
- self.output_size_per_partition, self.input_size,
- device=torch.cuda.current_device(), dtype=params_dtype))
- if perform_initialization:
- _initialize_affine_weight_gpu(self.weight, init_method,
- partition_dim=0, stride=stride)
+ self.weight = None
if bias:
- if use_cpu_initialization:
- self.bias = Parameter(torch.empty(
- self.output_size_per_partition, dtype=params_dtype))
+ if config.use_cpu_initialization:
+ self.bias = Parameter(
+ torch.empty(self.output_size_per_partition, dtype=config.params_dtype)
+ )
else:
- self.bias = Parameter(torch.empty(
- self.output_size_per_partition,
- device=torch.cuda.current_device(),
- dtype=params_dtype))
+ self.bias = Parameter(
+ torch.empty(
+ self.output_size_per_partition,
+ device=torch.cuda.current_device(),
+ dtype=config.params_dtype,
+ )
+ )
set_tensor_model_parallel_attributes(self.bias, True, 0, stride)
- # Always initialize bias to zero.
- with torch.no_grad():
- self.bias.zero_()
+ if config.perform_initialization:
+ # Always initialize bias to zero.
+ with torch.no_grad():
+ self.bias.zero_()
+ setattr(self.bias, 'allreduce', not (self.is_expert and self.expert_parallel))
else:
self.register_parameter('bias', None)
- self.async_tensor_model_parallel_allreduce = (
- async_tensor_model_parallel_allreduce and
- world_size > 1)
- if sequence_parallel_enabled:
- if world_size <= 1:
- warnings.warn(
- f"`sequence_parallel_enabled` is set to `True`, but tensor model parallel size is {world_size}. "
- f"Disabling sequence parallel."
- )
- sequence_parallel_enabled = False
- self.sequence_parallel_enabled = sequence_parallel_enabled
+ self.sequence_parallel = config.sequence_parallel
+ if self.sequence_parallel and world_size <= 1:
+ warnings.warn(
+ "`sequence_parallel` is set to `True`, but tensor model parallel size "
+ f"is {world_size}. Disabling sequence parallel."
+ )
+ self.sequence_parallel = False
- if gradient_accumulation_fusion:
- if not _grad_accum_fusion_available:
- raise RuntimeError(
- "ColumnParallelLinear was called with gradient_accumulation_fusion set "
- "to True but the custom CUDA extension fused_weight_gradient_mlp_cuda "
- "module is not found. To use gradient_accumulation_fusion you must "
- "install APEX with --cpp_ext and --cuda_ext. For example: "
- "pip install --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext .\" "
- "Note that the extension requires CUDA>=11. Otherwise, you must turn off "
- "gradient accumulation fusion."
- )
- self.gradient_accumulation_fusion = gradient_accumulation_fusion
+ self.allreduce_dgrad = (
+ world_size > 1 and not self.sequence_parallel and not self.disable_grad_reduce
+ )
- if self.async_tensor_model_parallel_allreduce and self.sequence_parallel_enabled:
+ if config.gradient_accumulation_fusion and not _grad_accum_fusion_available:
raise RuntimeError(
- "`async_tensor_model_parallel_allreduce` and `sequence_parallel_enabled` "
- "cannot be enabled at the same time."
+ "ColumnParallelLinear was called with gradient_accumulation_fusion set "
+ "to True but the custom CUDA extension fused_weight_gradient_mlp_cuda "
+ "module is not found. To use gradient_accumulation_fusion you must "
+ "install APEX with --cpp_ext and --cuda_ext. For example: "
+ "pip install --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext .\" "
+ "Note that the extension requires CUDA>=11. Otherwise, you must turn off "
+ "gradient accumulation fusion."
)
+ self.gradient_accumulation_fusion = config.gradient_accumulation_fusion
+ if self.allreduce_dgrad and self.sequence_parallel:
+ raise RuntimeError(
+ "`allreduce_dgrad` and `sequence_parallel` cannot be enabled at the same time."
+ )
- def forward(self, input_):
+ self._forward_impl = linear_with_grad_accumulation_and_async_allreduce
+
+ # Hook adding a default empty _extra_state for state dict
+ self._register_load_state_dict_pre_hook(
+ lambda state_dict, prefix, *args, **kwargs: state_dict.setdefault(
+ f'{prefix}_extra_state'
+ )
+ )
+
+ def forward(
+ self,
+ input_: torch.Tensor,
+ weight: Optional[torch.Tensor] = None,
+ runtime_gather_output: Optional[bool] = None,
+ ):
"""Forward of ColumnParallelLinear
Args:
- input_: 3D tensor whose order of dimension is [sequence, batch, hidden]
+ input_:
+ 3D tensor whose order of dimension is [sequence, batch, hidden]
+ weight (optional):
+ weight tensor to use, compulsory when skip_weight_param_allocation is True.
+ runtime_gather_output (bool): Gather output at runtime. Default None means
+ `gather_output` arg in the constructor will be used.
Returns:
- output
- bias
+
"""
+ if weight is None:
+ if self.weight is None:
+ raise RuntimeError(
+ "weight was not supplied to ColumnParallelLinear forward pass "
+ "and skip_weight_param_allocation is True."
+ )
+ weight = self.weight
+ else:
+ # Check the weight passed in is the correct shape
+ expected_shape = (self.output_size_per_partition, self.input_size)
+ if weight.shape != expected_shape:
+ raise RuntimeError(
+ f"supplied weight's shape is {tuple(weight.shape)}, "
+ f"not {expected_shape} as expected"
+ )
+
+ if self.config._cpu_offloading_context is not None:
+ if self.config._cpu_offloading_context.inside_context is True:
+ assert (
+ self.config.cpu_offloading is False
+ ), "CPU Offloading cannot be enabled while using non-TE modules"
+
bias = self.bias if not self.skip_bias_add else None
- if self.async_tensor_model_parallel_allreduce or \
- self.sequence_parallel_enabled:
+ if (
+ self.allreduce_dgrad
+ or self.sequence_parallel
+ or self.explicit_expert_comm
+ or self.disable_grad_reduce
+ ):
input_parallel = input_
else:
input_parallel = copy_to_tensor_model_parallel_region(input_)
+
+ if self.config.defer_embedding_wgrad_compute:
+ if (
+ self.config.wgrad_deferral_limit == 0
+ or len(self.embedding_activation_buffer) < self.config.wgrad_deferral_limit
+ ):
+ self.embedding_activation_buffer.append(input_parallel)
+
# Matrix multiply.
- output_parallel = linear_with_grad_accumulation_and_async_allreduce(
+ if not weight.requires_grad:
+ self._forward_impl = linear_with_frozen_weight
+ else:
+ self._forward_impl = linear_with_grad_accumulation_and_async_allreduce
+
+ allreduce_dgrad = False if self.explicit_expert_comm else self.allreduce_dgrad
+
+ output_parallel = self._forward_impl(
input=input_parallel,
- weight=self.weight,
+ weight=weight,
bias=bias,
gradient_accumulation_fusion=self.gradient_accumulation_fusion,
- async_grad_allreduce=self.async_tensor_model_parallel_allreduce,
- sequence_parallel_enabled=self.sequence_parallel_enabled,
+ async_grad_allreduce=allreduce_dgrad,
+ sequence_parallel=False if self.explicit_expert_comm else self.sequence_parallel,
+ grad_output_buffer=(
+ self.grad_output_buffer if self.config.defer_embedding_wgrad_compute else None
+ ),
+ wgrad_deferral_limit=(
+ self.config.wgrad_deferral_limit
+ if self.config.defer_embedding_wgrad_compute
+ else None
+ ),
+ allreduce_dgrad=allreduce_dgrad,
)
- if self.gather_output:
+
+ gather_output = self.gather_output
+ # Use the runtime gather output if it's set explicitly.
+ if runtime_gather_output is not None:
+ gather_output = runtime_gather_output
+
+ if gather_output:
# All-gather across the partitions.
- assert not self.sequence_parallel_enabled
+ assert not self.sequence_parallel
output = gather_from_tensor_model_parallel_region(output_parallel)
else:
output = output_parallel
output_bias = self.bias if self.skip_bias_add else None
return output, output_bias
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Sharding along axis 0, bias sharded"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(
+ state_dict, prefix, {'weight': 0, 'bias': 0}, sharded_offsets
+ )
+
+ def set_extra_state(self, state: Any):
+ """Extra state is ignored"""
+
+ def get_extra_state(self) -> None:
+ """Keep compatibility with TE state dict."""
+ return None
+
class RowParallelLinear(torch.nn.Module):
"""Linear layer with row parallelism.
- The linear layer is defined as Y = XA + b. A is parallelized along
- its first dimension and X along its second dimension as:
- - -
- | A_1 |
- | . |
- A = | . | X = [X_1, ..., X_p]
- | . |
- | A_p |
- - -
- Arguments:
- input_size: first dimension of matrix A.
- output_size: second dimension of matrix A.
-
- Keyword Arguments:
- bias: If true, add bias. Note that bias is not parallelized.
- input_is_parallel: If true, we assume that the input is already
- split across the GPUs and we do not split
- again.
- init_method: method to initialize weights. Note that bias is always set
- to zero.
- stride: For the strided linear layers.
- keep_master_weight_for_test: This was added for testing and should be
- set to False. It returns the master weights
- used for initialization.
- skip_bias_add: This was added to enable performance optimization where bias
- can be fused with other elementwise operations. We skip
- adding bias but instead return it.
- params_dtype:
- use_cpu_initialization:
- perform_initialization:
- gradient_accumulation_fusion:
- sequence_parallel_enabled:
+ The linear layer is defined as Y = XA + b. A is parallelized along its first dimension and X
+ along its second dimension. A = transpose([A_1 .. A_p]) X = [X_1, ..., X_p]
+
+ Args:
+ input_size:
+ first dimension of matrix A.
+ output_size:
+ second dimension of matrix A.
+ bias:
+ If true, add bias. Note that bias is not parallelized.
+ input_is_parallel:
+ If true, we assume that the input is already split across the GPUs
+ and we do not split again.
+ init_method:
+ method to initialize weights. Note that bias is always set to zero.
+ stride:
+ For the strided linear layers.
+ keep_master_weight_for_test:
+ This was added for testing and should be set to False. It returns the master weights
+ used for initialization.
+ skip_bias_add:
+ If True, do not add the bias term, instead return it to be added by the
+ caller. This enables performance optimations where bias can be fused with other
+ elementwise operations.
+ is_expert:
+ If True, the layer is treated as an MoE expert layer
+ tp_comm_buffer_name:
+ Communication buffer name. Not used in non-Transformer-Engine modules.
+ config:
+ ModelParallelConfig object
+
"""
- def __init__(self, input_size, output_size, *,
- bias=True, input_is_parallel=False,
- init_method=init.xavier_normal_, stride=1,
- keep_master_weight_for_test=False,
- skip_bias_add=False,
- params_dtype=torch.float32,
- use_cpu_initialization=False,
- perform_initialization=True,
- gradient_accumulation_fusion=False,
- sequence_parallel_enabled: bool = False,
- ):
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ config: ModelParallelConfig,
+ init_method: Callable,
+ bias: bool,
+ input_is_parallel: bool,
+ skip_bias_add: bool,
+ stride: int = 1,
+ keep_master_weight_for_test: bool = False,
+ is_expert: bool = False,
+ tp_comm_buffer_name: str = None, # Not used
+ ):
super(RowParallelLinear, self).__init__()
# Keep input parameters
self.input_size = input_size
self.output_size = output_size
self.input_is_parallel = input_is_parallel
+ self.skip_bias_add = skip_bias_add
+ self.config = config
+ self.is_expert = is_expert
+ self.expert_parallel = config.expert_model_parallel_size > 1
+ self.gradient_accumulation_fusion = config.gradient_accumulation_fusion
+ self.sequence_parallel = config.sequence_parallel
+ if self.sequence_parallel and not self.input_is_parallel:
+ raise RuntimeError("To enable `sequence_parallel`, `input_is_parallel` must be `True`")
+
+ self.explicit_expert_comm = self.is_expert and (
+ config.tensor_model_parallel_size > 1 or self.expert_parallel
+ )
+
# Divide the weight matrix along the last dimension.
- world_size = get_tensor_model_parallel_world_size()
+ if self.explicit_expert_comm and config.moe_extended_tp:
+ world_size = get_tensor_and_expert_parallel_world_size()
+ rank = get_tensor_and_expert_parallel_rank()
+ else:
+ world_size = get_tensor_model_parallel_world_size()
+ rank = get_tensor_model_parallel_rank()
+
self.input_size_per_partition = divide(input_size, world_size)
- self.skip_bias_add = skip_bias_add
- self.gradient_accumulation_fusion = gradient_accumulation_fusion
- self.sequence_parallel_enabled = sequence_parallel_enabled
- if self.sequence_parallel_enabled and not self.input_is_parallel:
- raise RuntimeError("To enable `sequence_parallel_enabled`, `input_is_parallel` must be `True`")
# Parameters.
# Note: torch.nn.functional.linear performs XA^T + b and as a result
# we allocate the transpose.
# Initialize weight.
- if use_cpu_initialization:
- self.weight = Parameter(torch.empty(self.output_size,
- self.input_size_per_partition,
- dtype=params_dtype))
- if perform_initialization:
+ if config.use_cpu_initialization:
+ self.weight = Parameter(
+ torch.empty(
+ self.output_size, self.input_size_per_partition, dtype=config.params_dtype
+ )
+ )
+ if config.perform_initialization:
self.master_weight = _initialize_affine_weight_cpu(
- self.weight, self.output_size, self.input_size,
- self.input_size_per_partition, 1, init_method,
- stride=stride, return_master_weight=keep_master_weight_for_test,
- params_dtype=params_dtype)
+ self.weight,
+ self.output_size,
+ self.input_size,
+ self.input_size_per_partition,
+ 1,
+ init_method,
+ stride=stride,
+ return_master_weight=keep_master_weight_for_test,
+ params_dtype=config.params_dtype,
+ rank=rank,
+ world_size=world_size,
+ )
else:
- self.weight = Parameter(torch.empty(
- self.output_size, self.input_size_per_partition,
- device=torch.cuda.current_device(), dtype=params_dtype))
- if perform_initialization:
- _initialize_affine_weight_gpu(self.weight, init_method,
- partition_dim=1, stride=stride)
+ self.weight = Parameter(
+ torch.empty(
+ self.output_size,
+ self.input_size_per_partition,
+ device=torch.cuda.current_device(),
+ dtype=config.params_dtype,
+ )
+ )
+ if config.perform_initialization:
+ _initialize_affine_weight_gpu(
+ self.weight,
+ init_method,
+ partition_dim=1,
+ stride=stride,
+ expert_parallel=(self.is_expert and self.expert_parallel),
+ )
+ setattr(self.weight, 'allreduce', not (self.is_expert and self.expert_parallel))
+
if bias:
- if use_cpu_initialization:
- self.bias = Parameter(torch.empty(self.output_size,
- dtype=params_dtype))
+ if config.use_cpu_initialization:
+ self.bias = Parameter(torch.empty(self.output_size, dtype=config.params_dtype))
else:
- self.bias = Parameter(torch.empty(
- self.output_size, device=torch.cuda.current_device(),
- dtype=params_dtype))
- setattr(self.bias, 'sequence_parallel', sequence_parallel_enabled)
-
- # Always initialize bias to zero.
- with torch.no_grad():
- self.bias.zero_()
+ self.bias = Parameter(
+ torch.empty(
+ self.output_size,
+ device=torch.cuda.current_device(),
+ dtype=config.params_dtype,
+ )
+ )
+
+ if config.perform_initialization:
+ # Always initialize bias to zero.
+ with torch.no_grad():
+ self.bias.zero_()
+ setattr(self.bias, 'allreduce', not (self.is_expert and self.expert_parallel))
+ setattr(self.bias, 'sequence_parallel', self.sequence_parallel)
else:
self.register_parameter('bias', None)
+ self._forward_impl = linear_with_grad_accumulation_and_async_allreduce
+ # Hook adding a default empty _extra_state for state dict
+ self._register_load_state_dict_pre_hook(
+ lambda state_dict, prefix, *args, **kwargs: state_dict.setdefault(
+ f'{prefix}_extra_state'
+ )
+ )
def forward(self, input_):
"""Forward of RowParallelLinear
@@ -686,31 +1141,64 @@ def forward(self, input_):
- output
- bias
"""
+
+ if self.config._cpu_offloading_context is not None:
+ if self.config._cpu_offloading_context.inside_context is True:
+ assert (
+ self.config.cpu_offloading is False
+ ), "CPU Offloading cannot be enabled while using non-TE modules"
+
# Set up backprop all-reduce.
if self.input_is_parallel:
input_parallel = input_
else:
- assert not self.sequence_parallel_enabled
+ assert not self.sequence_parallel
input_parallel = scatter_to_tensor_model_parallel_region(input_)
# Matrix multiply.
- output_parallel = linear_with_grad_accumulation_and_async_allreduce(
+ if not self.weight.requires_grad:
+ self._forward_impl = linear_with_frozen_weight
+ else:
+ self._forward_impl = linear_with_grad_accumulation_and_async_allreduce
+
+ allreduce_dgrad = False
+
+ output_parallel = self._forward_impl(
input=input_parallel,
weight=self.weight,
bias=None,
gradient_accumulation_fusion=self.gradient_accumulation_fusion,
- async_grad_allreduce=False,
- sequence_parallel_enabled=False,
+ async_grad_allreduce=allreduce_dgrad,
+ sequence_parallel=False,
+ grad_output_buffer=None,
+ allreduce_dgrad=allreduce_dgrad,
)
# All-reduce across all the partitions.
- if self.sequence_parallel_enabled:
+ if self.explicit_expert_comm:
+ assert self.skip_bias_add
+ output_ = output_parallel
+ elif self.sequence_parallel:
output_ = reduce_scatter_to_sequence_parallel_region(output_parallel)
else:
output_ = reduce_from_tensor_model_parallel_region(output_parallel)
if not self.skip_bias_add:
- output = output_ + self.bias if self.bias is not None else output_
+ output = (output_ + self.bias) if self.bias is not None else output_
output_bias = None
else:
output = output_
output_bias = self.bias
return output, output_bias
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Sharding along axis 1, bias not sharded"""
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ return make_sharded_tensors_for_checkpoint(
+ state_dict, prefix, {'weight': 1}, sharded_offsets
+ )
+
+ def set_extra_state(self, state: Any):
+ """Extra state is ignored"""
+
+ def get_extra_state(self) -> None:
+ """Keep compatibility with TE state dict."""
+ return None
diff --git a/megatron/core/tensor_parallel/mappings.py b/megatron/core/tensor_parallel/mappings.py
index 624be8054e..3addd8d2ee 100644
--- a/megatron/core/tensor_parallel/mappings.py
+++ b/megatron/core/tensor_parallel/mappings.py
@@ -3,10 +3,14 @@
import torch
from megatron.core.parallel_state import (
+ get_expert_model_parallel_group,
+ get_global_memory_buffer,
+ get_tensor_and_expert_parallel_group,
+ get_tensor_model_parallel_group,
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
- get_tensor_model_parallel_group,
)
+
from .utils import split_tensor_along_last_dim
@@ -14,11 +18,11 @@ def _reduce(input_):
"""All-reduce the input tensor across model parallel group."""
# Bypass the function if we are using only 1 GPU.
- if get_tensor_model_parallel_world_size()==1:
+ if get_tensor_model_parallel_world_size() == 1:
return input_
# All-reduce.
- torch.distributed.all_reduce(input_, group=get_tensor_model_parallel_group())
+ torch.distributed.all_reduce(input_.contiguous(), group=get_tensor_model_parallel_group())
return input_
@@ -53,13 +57,14 @@ def _split_along_first_dim(input_):
# Split along first dimension.
dim_size = input_.size()[0]
- assert dim_size % world_size == 0, \
- "First dimension of the tensor should be divisible by tensor parallel size"
+ assert (
+ dim_size % world_size == 0
+ ), "First dimension of the tensor should be divisible by tensor parallel size"
local_dim_size = dim_size // world_size
rank = get_tensor_model_parallel_rank()
dim_offset = rank * local_dim_size
- output = input_[dim_offset:dim_offset+local_dim_size].contiguous()
+ output = input_[dim_offset : dim_offset + local_dim_size].contiguous()
return output
@@ -72,55 +77,161 @@ def _gather_along_last_dim(input_):
if world_size == 1:
return input_
- # Size and dimension.
- last_dim = input_.dim() - 1
- rank = get_tensor_model_parallel_rank()
+ dim_size = list(input_.size())
+ dim_size[0] = dim_size[0] * world_size
- tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
- tensor_list[rank] = input_
- torch.distributed.all_gather(tensor_list, input_, group=get_tensor_model_parallel_group())
+ output = torch.empty(dim_size, dtype=input_.dtype, device=torch.cuda.current_device())
+ torch.distributed.all_gather_into_tensor(
+ output, input_.contiguous(), group=get_tensor_model_parallel_group()
+ )
+ tensor_list = output.chunk(world_size, dim=0)
+ output = torch.cat(tensor_list, dim=-1).contiguous()
+
+ return output
- # Note: torch.cat already creates a contiguous tensor.
- output = torch.cat(tensor_list, dim=last_dim).contiguous()
+def _reduce_scatter_along_last_dim(input_):
+ """Reduce-scatter tensors on the last dimension."""
+ world_size = get_tensor_model_parallel_world_size()
+ target_shape = list(input_.size())
+ target_shape[-1] = target_shape[-1] // world_size
+ input_ = input_.reshape(-1, input_.shape[-1])
+ split_tensors = torch.split(
+ input_, split_size_or_sections=input_.shape[-1] // world_size, dim=1
+ )
+ concat_tensor = torch.cat(split_tensors, dim=0)
+ output = _reduce_scatter_along_first_dim(concat_tensor).reshape(target_shape)
return output
-def _gather_along_first_dim(input_):
- """Gather tensors and concatinate along the first dimension."""
+def _gather_along_first_dim(input_, output_split_sizes=None):
+ """Gather tensors and concatenate along the first dimension.
+
+ Args:
+ input_tensor (torch.Tensor):
+ A tensor to be gathered.
+ output_split_sizes (List[int], optional):
+ A list specifying the sizes of the output splits along the first dimension.
+ If None, equal splitting is assumed. Default: None.
+
+ Returns:
+ torch.Tensor: Gathered tensor.
+ """
world_size = get_tensor_model_parallel_world_size()
# Bypass the function if we are using only 1 GPU.
if world_size == 1:
return input_
+ dim_size = list(input_.size())
+ if output_split_sizes is None:
+ dim_size[0] = dim_size[0] * world_size
+
+ output = torch.empty(dim_size, dtype=input_.dtype, device=torch.cuda.current_device())
+ torch.distributed._all_gather_base(
+ output, input_.contiguous(), group=get_tensor_model_parallel_group()
+ )
+ else:
+ dim_size[0] = sum(output_split_sizes)
+ output = torch.empty(dim_size, dtype=input_.dtype, device=torch.cuda.current_device())
+ output_tensor_list = list(torch.split(output, output_split_sizes, dim=0))
+ torch.distributed.all_gather(
+ output_tensor_list, input_, group=get_tensor_model_parallel_group()
+ )
+
+ return output
+
+
+def _reduce_scatter_along_first_dim(input_, input_split_sizes=None):
+ """Reduce-scatter the input tensor across model parallel group.
+
+ Args:
+ input_ (torch.Tensor): The input tensor to be reduce-scattered.
+ input_split_sizes (List[int], optional): A list specifying the sizes of
+ the input splits along the first dimension for each rank. If None,
+ equal splitting is assumed. Default: None.
+ """
+ world_size = get_tensor_model_parallel_world_size()
+ # Bypass the function if we are using only 1 GPU.
+ if world_size == 1:
+ return input_
+
+ if input_split_sizes is None:
+ dim_size = list(input_.size())
+ assert (
+ dim_size[0] % world_size == 0
+ ), "First dimension of the tensor should be divisible by tensor parallel size"
+
+ dim_size[0] = dim_size[0] // world_size
+
+ output = torch.empty(dim_size, dtype=input_.dtype, device=torch.cuda.current_device())
+ torch.distributed._reduce_scatter_base(
+ output, input_.contiguous(), group=get_tensor_model_parallel_group()
+ )
+ else:
+ rank = torch.distributed.get_rank(get_tensor_model_parallel_group())
+ input_tensor_list = list(torch.split(input_, input_split_sizes, dim=0))
+ output = torch.empty_like(input_tensor_list[rank])
+ torch.distributed.reduce_scatter(
+ output, input_tensor_list, group=get_tensor_model_parallel_group()
+ )
+ return output
+
+
+def _gather_along_first_dim_moe(input_, use_global_buffer=False):
+ """Gather tensors and concatenate along the first dimension."""
+ group = get_tensor_and_expert_parallel_group()
+ world_size = torch.distributed.get_world_size(group=group)
+ # Bypass the function if we are using only 1 GPU.
+ if world_size == 1:
+ return input_
+
dim_size = list(input_.size())
dim_size[0] = dim_size[0] * world_size
- output = torch.empty(dim_size, dtype=input_.dtype,
- device=torch.cuda.current_device())
- torch.distributed._all_gather_base(output, input_.contiguous(),
- group=get_tensor_model_parallel_group())
+ if use_global_buffer:
+ output = get_global_memory_buffer().get_tensor(dim_size, input_.dtype, "mpu")
+ else:
+ output = torch.empty(dim_size, dtype=input_.dtype, device=torch.cuda.current_device())
+ torch.distributed._all_gather_base(output, input_.contiguous(), group=group)
return output
-def _reduce_scatter_along_first_dim(input_):
+
+def _reduce_scatter_along_first_dim_moe(input_, use_global_buffer=False):
"""Reduce-scatter the input tensor across model parallel group."""
- world_size = get_tensor_model_parallel_world_size()
+ group = get_tensor_and_expert_parallel_group()
+ world_size = torch.distributed.get_world_size(group=group)
# Bypass the function if we are using only 1 GPU.
if world_size == 1:
return input_
dim_size = list(input_.size())
- assert dim_size[0] % world_size == 0, \
- "First dimension of the tensor should be divisible by tensor parallel size"
-
+ assert dim_size[0] % world_size == 0
dim_size[0] = dim_size[0] // world_size
-
- output = torch.empty(dim_size, dtype=input_.dtype,
- device=torch.cuda.current_device())
- torch.distributed._reduce_scatter_base(output, input_.contiguous(),
- group=get_tensor_model_parallel_group())
+
+ if use_global_buffer:
+ output = get_global_memory_buffer().get_tensor(dim_size, input_.dtype, "mpu")
+ else:
+ output = torch.empty(dim_size, dtype=input_.dtype, device=torch.cuda.current_device())
+ torch.distributed._reduce_scatter_base(output, input_.contiguous(), group=group)
+ return output
+
+
+def _gather_along_first_dim_expert_parallel(input_):
+ """Gather tensors and concatenate along the first dimension."""
+ group = get_expert_model_parallel_group()
+ world_size = torch.distributed.get_world_size(group=group)
+ # Bypass the function if we are using only 1 GPU.
+ if world_size == 1:
+ return input_
+
+ dim_size = list(input_.size())
+ dim_size[0] = dim_size[0] * world_size
+
+ output = torch.empty(dim_size, dtype=input_.dtype, device=torch.cuda.current_device())
+ torch.distributed._all_gather_base(output, input_.contiguous(), group=group)
+
return output
@@ -129,14 +240,17 @@ class _CopyToModelParallelRegion(torch.autograd.Function):
@staticmethod
def symbolic(graph, input_):
+ """Symbolic function for tracing."""
return input_
-
+
@staticmethod
def forward(ctx, input_):
+ """Forward function."""
return input_
@staticmethod
def backward(ctx, grad_output):
+ """Backward function."""
return _reduce(grad_output)
@@ -145,14 +259,17 @@ class _ReduceFromModelParallelRegion(torch.autograd.Function):
@staticmethod
def symbolic(graph, input_):
+ """Symbolic function for tracing."""
return _reduce(input_)
-
+
@staticmethod
def forward(ctx, input_):
+ """Forward function."""
return _reduce(input_)
@staticmethod
def backward(ctx, grad_output):
+ """Backward function."""
return grad_output
@@ -161,14 +278,17 @@ class _ScatterToModelParallelRegion(torch.autograd.Function):
@staticmethod
def symbolic(graph, input_):
+ """Symbolic function for tracing."""
return _split_along_last_dim(input_)
@staticmethod
def forward(ctx, input_):
+ """Forward function."""
return _split_along_last_dim(input_)
@staticmethod
def backward(ctx, grad_output):
+ """Backward function."""
return _gather_along_last_dim(grad_output)
@@ -177,14 +297,17 @@ class _GatherFromModelParallelRegion(torch.autograd.Function):
@staticmethod
def symbolic(graph, input_):
+ """Symbolic function for tracing."""
return _gather_along_last_dim(input_)
-
+
@staticmethod
def forward(ctx, input_):
+ """Forward function."""
return _gather_along_last_dim(input_)
@staticmethod
def backward(ctx, grad_output):
+ """Backward function."""
return _split_along_last_dim(grad_output)
@@ -193,87 +316,314 @@ class _ScatterToSequenceParallelRegion(torch.autograd.Function):
@staticmethod
def symbolic(graph, input_):
+ """Symbolic function for tracing."""
return _split_along_first_dim(input_)
@staticmethod
def forward(ctx, input_):
+ """Forward function."""
return _split_along_first_dim(input_)
@staticmethod
def backward(ctx, grad_output):
+ """Backward function."""
return _gather_along_first_dim(grad_output)
class _GatherFromSequenceParallelRegion(torch.autograd.Function):
- """Gather the input from sequence parallel region and concatinate."""
+ """Gather the input from sequence parallel region and concatinate."""
@staticmethod
- def symbolic(graph, input_, tensor_parallel_output_grad=True):
- return _gather_along_first_dim(input_)
-
+ def symbolic(graph, input_, tensor_parallel_output_grad=True, output_split_sizes=None):
+ """Symbolic function for tracing."""
+ return _gather_along_first_dim(input_, output_split_sizes)
+
@staticmethod
- def forward(ctx, input_, tensor_parallel_output_grad=True):
+ def forward(ctx, input_, tensor_parallel_output_grad=True, output_split_sizes=None):
+ """Forward function."""
ctx.tensor_parallel_output_grad = tensor_parallel_output_grad
- return _gather_along_first_dim(input_)
+ ctx.output_split_sizes = output_split_sizes
+ return _gather_along_first_dim(input_, ctx.output_split_sizes)
@staticmethod
def backward(ctx, grad_output):
+ """Backward function."""
tensor_parallel_output_grad = ctx.tensor_parallel_output_grad
# If the computation graph after the gather operation is
- # in the tensor parallel mode, output gradients need to reduce
- # scattered and whereas if the computation is duplicated,
+ # in the tensor parallel mode, output gradients need to reduce
+ # scattered and whereas if the computation is duplicated,
# output gradients need to be scattered.
if tensor_parallel_output_grad:
- return _reduce_scatter_along_first_dim(grad_output), None
+ return (
+ _reduce_scatter_along_first_dim(grad_output, ctx.output_split_sizes),
+ None,
+ None,
+ )
else:
- return _split_along_first_dim(grad_output), None
+ assert ctx.output_split_sizes is None
+ return _split_along_first_dim(grad_output), None, None
class _ReduceScatterToSequenceParallelRegion(torch.autograd.Function):
"""Reduce scatter the input from the model parallel region."""
+ @staticmethod
+ def symbolic(graph, input_, input_split_sizes=None):
+ """Symbolic function for tracing."""
+ return _reduce_scatter_along_first_dim(input_, input_split_sizes)
+
+ @staticmethod
+ def forward(ctx, input_, input_split_sizes=None):
+ """Forward function."""
+ ctx.input_split_sizes = input_split_sizes
+ return _reduce_scatter_along_first_dim(input_, input_split_sizes)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ """Backward function."""
+ input_split_sizes = ctx.input_split_sizes
+ return _gather_along_first_dim(grad_output, input_split_sizes), None
+
+
+class _GatherFromSequenceParallelRegionToMOE(torch.autograd.Function):
+ """Gather the input from model parallel region and concatenate.""" # TODO
+
+ @staticmethod
+ def symbolic(graph, input_, use_global_buffer=False):
+ """Symbolic function for tracing."""
+ return _gather_along_first_dim_moe(input_, use_global_buffer)
+
+ @staticmethod
+ def forward(ctx, input_, use_global_buffer=False):
+ """Forward function."""
+ ctx.use_global_buffer = use_global_buffer
+ return _gather_along_first_dim_moe(input_, use_global_buffer)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ """Backward function."""
+ use_global_buffer = ctx.use_global_buffer
+ return _reduce_scatter_along_first_dim_moe(grad_output, use_global_buffer), None
+
+
+class _ReduceScatterToSequenceParallelRegionFromMOE(torch.autograd.Function):
+ """Reduce scatter the input from the model parallel region."""
+
+ @staticmethod
+ def symbolic(graph, input_, use_global_buffer=False):
+ """Symbolic function for tracing."""
+ return _reduce_scatter_along_first_dim_moe(input_, use_global_buffer)
+
+ @staticmethod
+ def forward(ctx, input_, use_global_buffer=False):
+ """Forward function."""
+ ctx.use_global_buffer = use_global_buffer
+ return _reduce_scatter_along_first_dim_moe(input_, use_global_buffer)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ """Backward function."""
+ use_global_buffer = ctx.use_global_buffer
+ return _gather_along_first_dim_moe(grad_output, use_global_buffer), None
+
+
+class _AllGatherFromTensorParallelRegion(torch.autograd.Function):
+ """Gather the input from model parallel region and concatenate."""
+
+ @staticmethod
+ def symbolic(graph, input_):
+ """Symbolic function for tracing."""
+ return _gather_along_last_dim(input_)
+
+ @staticmethod
+ def forward(ctx, input_):
+ """Forward function."""
+ return _gather_along_last_dim(input_)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ """Backward function."""
+ return _reduce_scatter_along_last_dim(grad_output)
+
+
+class _ReduceScatterToTensorParallelRegion(torch.autograd.Function):
+ """Reduce scatter the input from the model parallel region."""
+
@staticmethod
def symbolic(graph, input_):
- return _reduce_scatter_along_first_dim(input_)
-
+ """Symbolic function for tracing."""
+ return _reduce_scatter_along_last_dim(input_)
+
@staticmethod
def forward(ctx, input_):
- return _reduce_scatter_along_first_dim(input_)
+ """Forward function."""
+ return _reduce_scatter_along_last_dim(input_)
@staticmethod
def backward(ctx, grad_output):
- return _gather_along_first_dim(grad_output)
+ """Backward function."""
+ return _gather_along_last_dim(grad_output)
+
+
+class _AllToAll(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, group, input, output_split_sizes, input_split_sizes):
+ """Forward function."""
+ ctx.group = group
+ ctx.output_split_sizes = output_split_sizes
+ ctx.input_split_sizes = input_split_sizes
+
+ world_size = torch.distributed.get_world_size(group=group)
+ # Bypass the function if we are using only 1 GPU.
+ if world_size == 1:
+ return input
+
+ input = input.contiguous()
+ if output_split_sizes is None:
+ # Equal split (all2all)
+ output = torch.empty_like(input)
+ else:
+ # Unequal split (all2all-v)
+ output = input.new_empty(
+ size=[sum(output_split_sizes)] + list(input.size()[1:]),
+ dtype=input.dtype,
+ device=torch.cuda.current_device(),
+ )
+ torch.distributed.all_to_all_single(
+ output,
+ input,
+ output_split_sizes=output_split_sizes,
+ input_split_sizes=input_split_sizes,
+ group=group,
+ )
+ return output
+
+ @staticmethod
+ def backward(ctx, *grad_output):
+ """Backward function."""
+ return (
+ None,
+ _AllToAll.apply(ctx.group, *grad_output, ctx.input_split_sizes, ctx.output_split_sizes),
+ None,
+ None,
+ )
# -----------------
# Helper functions.
# -----------------
+
def copy_to_tensor_model_parallel_region(input_):
+ """Wrapper for autograd function"""
return _CopyToModelParallelRegion.apply(input_)
def reduce_from_tensor_model_parallel_region(input_):
+ """Wrapper for autograd function"""
return _ReduceFromModelParallelRegion.apply(input_)
def scatter_to_tensor_model_parallel_region(input_):
+ """Wrapper for autograd function"""
return _ScatterToModelParallelRegion.apply(input_)
def gather_from_tensor_model_parallel_region(input_):
+ """Wrapper for autograd function"""
return _GatherFromModelParallelRegion.apply(input_)
def scatter_to_sequence_parallel_region(input_):
+ """Wrapper for autograd function"""
return _ScatterToSequenceParallelRegion.apply(input_)
-def gather_from_sequence_parallel_region(input_, tensor_parallel_output_grad=True):
- return _GatherFromSequenceParallelRegion.apply(input_, tensor_parallel_output_grad)
+def gather_from_sequence_parallel_region(
+ input_, tensor_parallel_output_grad=True, output_split_sizes=None
+):
+ """Wrapper for autograd function"""
+ return _GatherFromSequenceParallelRegion.apply(
+ input_, tensor_parallel_output_grad, output_split_sizes
+ )
+
+
+def reduce_scatter_to_sequence_parallel_region(input_, input_split_sizes=None):
+ """Wrapper for autograd function"""
+ return _ReduceScatterToSequenceParallelRegion.apply(input_, input_split_sizes)
+
+
+def gather_from_sequence_parallel_region_to_moe(input_, use_global_buffer=False):
+ """Wrapper for autograd function"""
+ return _GatherFromSequenceParallelRegionToMOE.apply(input_, use_global_buffer)
+
+
+def reduce_scatter_to_sequence_parallel_region_from_moe(input_, use_global_buffer=False):
+ """Wrapper for autograd function"""
+ return _ReduceScatterToSequenceParallelRegionFromMOE.apply(input_, use_global_buffer)
+
+def all_gather_last_dim_from_tensor_parallel_region(input_):
+ """Wrapper for autograd function"""
+ return _AllGatherFromTensorParallelRegion.apply(input_)
-def reduce_scatter_to_sequence_parallel_region(input_):
- return _ReduceScatterToSequenceParallelRegion.apply(input_)
+def reduce_scatter_last_dim_to_tensor_parallel_region(input_):
+ """Wrapper for autograd function"""
+ return _ReduceScatterToTensorParallelRegion.apply(input_)
+
+
+def all_to_all(group, input_, output_split_sizes_=None, input_split_sizes=None):
+ """Wrapper for autograd function"""
+ return _AllToAll.apply(group, input_, output_split_sizes_, input_split_sizes)
+
+
+def all_to_all_sp2hp(input_):
+ """
+ Perform AlltoAll communication on tensor parallel group, transform the input tensor from shape
+ [num_tokens/TP, H] to [num_tokens, H/TP].
+
+ Args:
+ input_ (torch.Tensor):
+ The input tensor which has been distributed along the sequence
+ dimension.
+
+ Returns:
+ torch.Tensor: The output tensor with shape [num_tokens, H/TP].
+
+ """
+ world_size = get_tensor_model_parallel_world_size()
+ tp_group = get_tensor_model_parallel_group()
+ input_ = input_.reshape(-1, input_.shape[-1])
+ split_tensors = torch.split(
+ input_, split_size_or_sections=input_.shape[-1] // world_size, dim=1
+ )
+ concat_tensor = torch.cat(split_tensors, dim=0)
+ output = all_to_all(tp_group, concat_tensor)
+ return output
+
+
+def all_to_all_hp2sp(input_):
+ """
+ Perform AlltoAll communication on tensor parallel group, transform the input tensor from shape
+ [num_tokens, H/TP] to [num_tokens/TP, H].
+
+ Args:
+ input_ (torch.Tensor):
+ The input tensor which has been distributed along the hidden
+ dimension.
+
+ Returns:
+ torch.Tensor: The output tensor with shape [num_tokens/TP, H].
+ """
+ world_size = get_tensor_model_parallel_world_size()
+ input_ = input_.reshape(-1, input_.shape[-1])
+ tp_group = get_tensor_model_parallel_group()
+ input_exchanged = all_to_all(tp_group, input_)
+ input_reshaped = input_exchanged.reshape(-1, input_exchanged.shape[-1])
+ split_tensors = torch.split(
+ input_reshaped, split_size_or_sections=input_reshaped.shape[0] // world_size, dim=0
+ )
+ output = torch.cat(split_tensors, dim=-1)
+ return output
diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py
index 23059fc1f5..4b144d4163 100644
--- a/megatron/core/tensor_parallel/random.py
+++ b/megatron/core/tensor_parallel/random.py
@@ -4,28 +4,26 @@
# repo: https://github.com/pytorch/pytorch
import contextlib
+import logging
import torch
from torch import _C
-from torch.cuda import _lazy_call, device as device_ctx_manager
+from torch.cuda import _lazy_call
+from torch.cuda import device as device_ctx_manager
from torch.utils.checkpoint import detach_variable
from megatron.core.parallel_state import (
- get_data_parallel_rank,
- get_tensor_model_parallel_group,
+ get_expert_model_parallel_rank,
get_tensor_model_parallel_rank,
- get_tensor_model_parallel_world_size,
)
+from megatron.core.utils import is_te_min_version, safely_set_viewless_tensor_data
-from .utils import (
- split_tensor_into_1d_equal_chunks,
- gather_split_1d_tensor,
-)
-
-from megatron.core.utils import safely_set_viewless_tensor_data
+from .utils import gather_split_1d_tensor, split_tensor_into_1d_equal_chunks
# Default name for the model parallel rng tracker.
_MODEL_PARALLEL_RNG_TRACKER_NAME = 'model-parallel-rng'
+_EXPERT_PARALLEL_RNG_TRACKER_NAME = 'expert-parallel-rng'
+_DATA_PARALLEL_RNG_TRACKER_NAME = 'data-parallel-rng'
def _set_cuda_rng_state(new_state, device=-1):
@@ -42,6 +40,7 @@ def _set_cuda_rng_state(new_state, device=-1):
def cb():
with device_ctx_manager(device):
_C._cuda_setRNGState(new_state)
+
else:
# newer PyTorch
if device == -1:
@@ -61,6 +60,17 @@ def cb():
_lazy_call(cb)
+def get_expert_parallel_rng_tracker_name():
+ """Get the expert parallel rng tracker name"""
+ global _EXPERT_PARALLEL_RNG_TRACKER_NAME
+ return _EXPERT_PARALLEL_RNG_TRACKER_NAME
+
+
+def get_data_parallel_rng_tracker_name():
+ """Get the data parallel rng tracker name"""
+ global _DATA_PARALLEL_RNG_TRACKER_NAME
+ return _DATA_PARALLEL_RNG_TRACKER_NAME
+
class CudaRNGStatesTracker:
"""Tracker for the cuda RNG states.
@@ -72,14 +82,22 @@ class CudaRNGStatesTracker:
"""
def __init__(self):
- # Map from a string name to the cuda rng state.
- self.states_ = {}
- # Seeds are just for book keeping and ensure no seed is set twice.
- self.seeds_ = set()
+ self.reset()
+
+ def is_initialized(self):
+ """Checks if the internal RNG state has been set wirth set_states()."""
+ return self._is_initialized
def reset(self):
"""Set to the initial state (no tracker)."""
+
+ # Track if initialized.
+ self._is_initialized = False
+
+ # Map from a string name to the cuda rng state.
self.states_ = {}
+
+ # Seeds are just for book keeping and ensure no seed is set twice.
self.seeds_ = set()
def get_states(self):
@@ -93,10 +111,12 @@ def get_states(self):
def set_states(self, states):
"""Set the rng states. For efficiency purposes, we do not check
the size of seed for compatibility."""
+ self._is_initialized = True
self.states_ = states
def add(self, name, seed):
"""Track the rng state."""
+ self._is_initialized = True
# Check seed is not already used.
if seed in self.seeds_:
raise Exception('seed {} already exists'.format(seed))
@@ -123,10 +143,15 @@ def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME):
orig_cuda_rng_state = torch.cuda.get_rng_state()
# Set rng state to the desired one
_set_cuda_rng_state(self.states_[name])
+ # Record cpu RNG state
+ cpu_rng_state = torch.get_rng_state()
# Do the stuff we wanted to do.
try:
yield
finally:
+ # Throw a warning if cpu RNG state changed
+ if not torch.all(cpu_rng_state == torch.get_rng_state()).item():
+ logging.getLogger(__name__).warning('CPU RNG state changed within GPU RNG context')
# Update the current rng state for later use.
self.states_[name] = torch.cuda.get_rng_state()
# And set the state to the original state we started with.
@@ -134,11 +159,35 @@ def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME):
# RNG tracker object.
-_CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker()
+_CUDA_RNG_STATE_TRACKER = None
+_CUDA_RNG_STATE_TRACKER_INITIALIZED = False
+
+
+def initialize_rng_tracker(use_te_rng_tracker: bool = False):
+ """Create the RNG tracker. 'use_te_rng_tracker' determines whether to use
+ Megatron or TransformerEngine's implementation.
+ In particular, TransformerEngine's implementation is cudagraphable and supports FP8.
+ """
+
+ global _CUDA_RNG_STATE_TRACKER
+ global _CUDA_RNG_STATE_TRACKER_INITIALIZED
+ if _CUDA_RNG_STATE_TRACKER_INITIALIZED:
+ return
+
+ if use_te_rng_tracker:
+ if not is_te_min_version("1.5.0"):
+ raise RuntimeError("use_te_rng_tracker requires TransformerEngine version >= 1.5")
+ from megatron.core.extensions.transformer_engine import TECudaRNGStatesTracker
+
+ _CUDA_RNG_STATE_TRACKER = TECudaRNGStatesTracker()
+ else:
+ _CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker()
+ _CUDA_RNG_STATE_TRACKER_INITIALIZED = True
-def get_cuda_rng_tracker():
+def get_cuda_rng_tracker(use_te_rng_tracker=False):
"""Get cuda rng tracker."""
+ initialize_rng_tracker(use_te_rng_tracker)
return _CUDA_RNG_STATE_TRACKER
@@ -150,14 +199,12 @@ def model_parallel_cuda_manual_seed(seed):
after this function. Basically, this is replacement for that
function.
Two set of RNG states are tracked:
- default state: This is for data parallelism and is the same among a
- set of model parallel GPUs but different across
- different model paralle groups. This is used for
- example for dropout in the non-tensor-model-parallel regions.
- tensor-model-parallel state: This state is different among a set of model
- parallel GPUs, but the same across data parallel
- groups. This is used for example for dropout in
- model parallel regions.
+ default state: This is for data parallelism and is the same among a set of model parallel GPUs
+ but different across different model parallel groups. This is used for example for dropout
+ in the non-tensor-model-parallel regions.
+ tensor-model-parallel state: This state is different among a set of model parallel GPUs,
+ but the same across data parallel groups. This is used for example for dropout
+ in model parallel regions.
"""
# 2718 is just for fun and any POSITIVE value will work.
offset = seed + 2718
@@ -165,26 +212,33 @@ def model_parallel_cuda_manual_seed(seed):
# Data parallel gets the original seed.
data_parallel_seed = seed
+ initialize_rng_tracker()
_CUDA_RNG_STATE_TRACKER.reset()
# Set the default state.
torch.cuda.manual_seed(data_parallel_seed)
+ _CUDA_RNG_STATE_TRACKER.add(_DATA_PARALLEL_RNG_TRACKER_NAME, data_parallel_seed)
+
# and model parallel state.
- _CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME,
- tensor_model_parallel_seed)
+ _CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME, tensor_model_parallel_seed)
+
+ expert_parallel_seed = (
+ seed + 1024 + 100 * get_expert_model_parallel_rank() + get_tensor_model_parallel_rank()
+ )
+ _CUDA_RNG_STATE_TRACKER.add(_EXPERT_PARALLEL_RNG_TRACKER_NAME, expert_parallel_seed)
class CheckpointFunction(torch.autograd.Function):
- """This function is adapted from torch.utils.checkpoint with
- two main changes:
- 1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state`
- 2) the states in the model parallel tracker are also properly
- tracked/set/reset.
+ """Checkpoint Function
+
+ This function is adapted from torch.utils.checkpoint with two main changes:
+ 1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state`
+ 2) the states in the model parallel tracker are also properly tracked/set/reset.
"""
+
@staticmethod
def forward(ctx, run_function, distribute_saved_activations, *args):
ctx.run_function = run_function
- ctx.distribute_saved_activations \
- = distribute_saved_activations
+ ctx.distribute_saved_activations = distribute_saved_activations
# Copy the rng states.
ctx.fwd_cpu_rng_state = torch.get_rng_state()
@@ -199,8 +253,8 @@ def forward(ctx, run_function, distribute_saved_activations, *args):
if distribute_saved_activations:
ctx.input_0_shape = args[0].data.shape
safely_set_viewless_tensor_data(
- args[0],
- split_tensor_into_1d_equal_chunks(args[0].data, new_buffer=True))
+ args[0], split_tensor_into_1d_equal_chunks(args[0].data, new_buffer=True)
+ )
# Store everything.
ctx.save_for_backward(*args)
@@ -210,13 +264,15 @@ def forward(ctx, run_function, distribute_saved_activations, *args):
@staticmethod
def backward(ctx, *args):
if not torch.autograd._is_checkpoint_valid():
- raise RuntimeError("Checkpointing is not compatible with .grad(), "
- "please use .backward() if possible")
+ raise RuntimeError(
+ "Checkpointing is not compatible with .grad(), "
+ "please use .backward() if possible"
+ )
inputs = ctx.saved_tensors
if ctx.distribute_saved_activations:
safely_set_viewless_tensor_data(
- inputs[0],
- gather_split_1d_tensor(inputs[0].data).view(ctx.input_0_shape))
+ inputs[0], gather_split_1d_tensor(inputs[0].data).view(ctx.input_0_shape)
+ )
# Store the current states.
bwd_cpu_rng_state = torch.get_rng_state()
@@ -240,14 +296,15 @@ def backward(ctx, *args):
if isinstance(outputs, torch.Tensor):
outputs = (outputs,)
+
+ # filter out non tensor outputs for backward pass
+ outputs, args = zip(*filter(lambda x: torch.is_tensor(x[0]), zip(outputs, args)))
torch.autograd.backward(outputs, args)
- grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp
- for inp in detached_inputs)
+ grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs)
return (None, None) + grads
def checkpoint(function, distribute_saved_activations, *args):
"""Checkpoint a model or part of the model.
This has been directly copied from torch.utils.checkpoint."""
- return CheckpointFunction.apply(function,
- distribute_saved_activations, *args)
+ return CheckpointFunction.apply(function, distribute_saved_activations, *args)
diff --git a/megatron/core/tensor_parallel/utils.py b/megatron/core/tensor_parallel/utils.py
index a4c7cb77cc..d7c191b411 100644
--- a/megatron/core/tensor_parallel/utils.py
+++ b/megatron/core/tensor_parallel/utils.py
@@ -1,26 +1,31 @@
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
-import torch
from typing import List, Sequence
-from megatron.core.utils import divide
+import torch
+
from megatron.core import parallel_state
+from megatron.core.parallel_state import (
+ get_tensor_model_parallel_group,
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+)
+from megatron.core.utils import divide
+
def split_tensor_along_last_dim(
- tensor: torch.Tensor,
- num_partitions: int,
- contiguous_split_chunks: bool = False,
+ tensor: torch.Tensor, num_partitions: int, contiguous_split_chunks: bool = False
) -> List[torch.Tensor]:
- """ Split a tensor along its last dimension.
+ """Split a tensor along its last dimension.
- Arguments:
- tensor: input tensor.
- num_partitions: number of partitions to split the tensor
- contiguous_split_chunks: If True, make each chunk contiguous
- in memory.
+ Args:
+ tensor: input tensor.
+ num_partitions: number of partitions to split the tensor
+ contiguous_split_chunks: If True, make each chunk contiguous
+ in memory.
- Returns:
- A list of Tensors
+ Returns:
+ A list of Tensors
"""
# Get the size and dimension.
last_dim = tensor.dim() - 1
@@ -33,28 +38,31 @@ def split_tensor_along_last_dim(
return tensor_list
+
def split_tensor_into_1d_equal_chunks(tensor, new_buffer=False):
- """ Break a tensor into equal 1D chunks across tensor parallel ranks.
+ """Break a tensor into equal 1D chunks across tensor parallel ranks.
- Returns a Tensor or View with this rank's portion of the data.
+ Returns a Tensor or View with this rank's portion of the data.
- Arguments:
- tensor: The tensor to split
+ Args:
+ tensor: The tensor to split
- Keyword Arguments:
- new_buffer (bool): If True, returns a new Tensor.
- If False, returns a view into the existing Tensor.
- Default is False
+ Keyword Args:
+ new_buffer (bool): If True, returns a new Tensor.
+ If False, returns a view into the existing Tensor.
+ Default is False
"""
- partition_size = torch.numel(tensor) // \
- parallel_state.get_tensor_model_parallel_world_size()
+ partition_size = torch.numel(tensor) // parallel_state.get_tensor_model_parallel_world_size()
start_index = partition_size * parallel_state.get_tensor_model_parallel_rank()
end_index = start_index + partition_size
if new_buffer:
- data = torch.empty(partition_size, dtype=tensor.dtype,
- device=torch.cuda.current_device(),
- requires_grad=False)
+ data = torch.empty(
+ partition_size,
+ dtype=tensor.dtype,
+ device=torch.cuda.current_device(),
+ requires_grad=False,
+ )
data.copy_(tensor.view(-1)[start_index:end_index])
else:
data = tensor.view(-1)[start_index:end_index]
@@ -62,33 +70,33 @@ def split_tensor_into_1d_equal_chunks(tensor, new_buffer=False):
def gather_split_1d_tensor(tensor):
- """ Opposite of split_tensor_into_1d_equal_chunks. Gather values from tensor
- model parallel ranks.
+ """Opposite of split_tensor_into_1d_equal_chunks. Gather values from tensor
+ model parallel ranks.
- Returns a new Tensor with the gathered data.
+ Returns a new Tensor with the gathered data.
- Arguments:
- tensor: A Tensor or view of this rank's portion of the data.
+ Args:
+ tensor: A Tensor or view of this rank's portion of the data.
"""
- numel_gathered = torch.numel(tensor) * \
- parallel_state.get_tensor_model_parallel_world_size()
- gathered = torch.empty(numel_gathered, dtype=tensor.dtype,
- device=torch.cuda.current_device(),
- requires_grad=False)
+ numel_gathered = torch.numel(tensor) * parallel_state.get_tensor_model_parallel_world_size()
+ gathered = torch.empty(
+ numel_gathered, dtype=tensor.dtype, device=torch.cuda.current_device(), requires_grad=False
+ )
# TODO: This API is experimental in pytorch (as of Feb 2022) and
# this might break in future pytorch releases. We chose this API
# as opposed to torch.distributed.all_gather for efficiency reasons.
# This API calls directly NCCL all-gather versus the former does
# internal copies and can potentially cause slow down.
- torch.distributed._all_gather_base(gathered, tensor,
- group=parallel_state.get_tensor_model_parallel_group())
+ torch.distributed._all_gather_base(
+ gathered, tensor, group=parallel_state.get_tensor_model_parallel_group()
+ )
return gathered
class VocabUtility:
- """ Split the vocabulary into `world_size` chunks and return the first
- and last index of the vocabulary belonging to the `rank`
- partition: Note that indices in [fist, last)
+ """Split the vocabulary into `world_size` chunks and return the first
+ and last index of the vocabulary belonging to the `rank`
+ partition: Note that indices in [fist, last)
"""
@@ -101,7 +109,9 @@ def vocab_range_from_per_partition_vocab_size(
return index_f, index_l
@staticmethod
- def vocab_range_from_global_vocab_size(global_vocab_size: int, rank: int, world_size: int) -> Sequence[int]:
+ def vocab_range_from_global_vocab_size(
+ global_vocab_size: int, rank: int, world_size: int
+ ) -> Sequence[int]:
per_partition_vocab_size = divide(global_vocab_size, world_size)
return VocabUtility.vocab_range_from_per_partition_vocab_size(
per_partition_vocab_size, rank, world_size
diff --git a/megatron/timers.py b/megatron/core/timers.py
similarity index 54%
rename from megatron/timers.py
rename to megatron/core/timers.py
index a9478fa014..e7070e37d8 100644
--- a/megatron/timers.py
+++ b/megatron/core/timers.py
@@ -2,16 +2,14 @@
"""Megatron timers."""
-from abc import ABC
-from abc import abstractmethod
import time
+from abc import ABC, abstractmethod
+from typing import List
import torch
-
class TimerBase(ABC):
-
def __init__(self, name):
self.name = name
@@ -32,9 +30,7 @@ def elapsed(self, reset=True, barrier=False):
pass
-
class DummyTimer(TimerBase):
-
def __init__(self):
super().__init__('dummy timer')
@@ -48,13 +44,13 @@ def reset(self):
return
def elapsed(self, reset=True, barrier=False):
- raise Exception('dummy timer should not be used to '
- 'calculate elapsed time')
-
+ raise Exception('dummy timer should not be used to calculate elapsed time')
class Timer(TimerBase):
"""
+ Timer class with ability to start/stop.
+
Comment on using `barrier`: If this flag is passed, then all
the caller processes will wait till all reach the timing routine.
It is up to the user to make sure all the ranks in `barrier_group`
@@ -64,20 +60,33 @@ class Timer(TimerBase):
"""
def __init__(self, name):
+ """Initialize Timer.
+
+ Args:
+ name (str): Name of the timer.
+ """
super().__init__(name)
self._elapsed = 0.0
+ self._active_time = 0.0
self._started = False
# Note that None will default to the global process group
self._barrier_group = None
self._start_time = time.time()
-
def set_barrier_group(self, barrier_group):
- self._barrier_group = barrier_group
+ """Sets barrier group.
+ Args:
+ barrier_group (ProcessGroup): Torch ProcessGroup for barrier.
+ """
+ self._barrier_group = barrier_group
def start(self, barrier=False):
- """Start the timer."""
+ """Start the timer.
+
+ Args:
+ barrier (bool, optional): Synchronizes ranks before starting. Defaults to False.
+ """
assert not self._started, 'timer has already been started'
if barrier:
torch.distributed.barrier(group=self._barrier_group)
@@ -85,25 +94,37 @@ def start(self, barrier=False):
self._start_time = time.time()
self._started = True
-
def stop(self, barrier=False):
- """Stop the timer."""
+ """Stop the timer.
+
+ Args:
+ barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False.
+ """
assert self._started, 'timer is not started'
if barrier:
torch.distributed.barrier(group=self._barrier_group)
torch.cuda.synchronize()
- self._elapsed += (time.time() - self._start_time)
+ elapsed = time.time() - self._start_time
+ self._elapsed += elapsed
+ self._active_time += elapsed
self._started = False
-
def reset(self):
"""Reset timer."""
+ # Don't reset _active_time
self._elapsed = 0.0
self._started = False
-
def elapsed(self, reset=True, barrier=False):
- """Calculate the elapsed time."""
+ """Calculates the elapsed time and restarts timer.
+
+ Args:
+ reset (bool, optional): Resets timer before restarting. Defaults to True.
+ barrier (bool, optional): Synchronizes ranks before stopping. Defaults to False.
+
+ Returns:
+ float: Elapsed time.
+ """
_started = self._started
# If the timing in progress, end it first.
if self._started:
@@ -118,37 +139,53 @@ def elapsed(self, reset=True, barrier=False):
self.start(barrier=barrier)
return _elapsed
+ def active_time(self):
+ return self._active_time
class Timers:
- """Group of timers."""
+ """Class for a group of Timers."""
def __init__(self, log_level, log_option):
+ """Initialize group of timers.
+
+ Args:
+ log_level (int): Log level to control what timers are enabled.
+ log_option (str): Setting for logging statistics over ranks for all the timers. Allowed: ['max', 'minmax', 'all'].
+ """
self._log_level = log_level
+ allowed_log_options = set(['max', 'minmax', 'all'])
+ assert (
+ log_option in allowed_log_options
+ ), 'input log option {} is invalid. It must be one of {}'.format(
+ log_option, allowed_log_options
+ )
self._log_option = log_option
self._timers = {}
self._log_levels = {}
self._dummy_timer = DummyTimer()
self._max_log_level = 2
-
def __call__(self, name, log_level=None):
+ """Call timer with name and log level."""
# If the timer has already been set, then check if the log-level
# is provided, it matches the one that the timer was created with.
if name in self._timers:
if log_level is not None:
- assert log_level == self._log_levels[name], \
- 'input log level {} does not match already existing '\
- 'log level {} for {} timer'.format(
- log_level, self._log_levels[name], name)
+ assert log_level == self._log_levels[name], (
+ 'input log level {} does not match already existing '
+ 'log level {} for {} timer'.format(log_level, self._log_levels[name], name)
+ )
return self._timers[name]
# If timer does not exist and no log level is provided,
# set it to the max log level which is 2.
if log_level is None:
log_level = self._max_log_level
- assert log_level <= self._max_log_level, \
- 'log level {} is larger than max supported log level {}'.format(
- log_level, self._max_log_level)
+ assert (
+ log_level <= self._max_log_level
+ ), 'log level {} is larger than max supported log level {}'.format(
+ log_level, self._max_log_level
+ )
# Now if the input log level is larger than the one set for
# the timers class, just ignore it and return a dummy timer.
if log_level > self._log_level:
@@ -158,18 +195,21 @@ def __call__(self, name, log_level=None):
self._log_levels[name] = log_level
return self._timers[name]
-
def _get_elapsed_time_all_ranks(self, names, reset, barrier):
- """
+ """Returns elapsed times of timers in names.
Assumptions:
- All the ranks call this function.
- `names` are identical on all ranks.
If the above assumptions are not met, calling this function will
result in hang.
- Arguments:
- - names: list of timer names
- - reset: reset the timer after recording the elapsed time
- - barrier: if set, do a global barrier before time measurments
+
+ Args:
+ names (List[str]): list of timer names
+ reset (bool): reset the timer after recording the elapsed time
+ barrier (bool): if set, do a global barrier before time measurments
+
+ Returns:
+ torch.tensor: Tensor of size [world_size, len(names)] with times in float.
"""
# First make sure all the callers are in sync.
@@ -184,30 +224,28 @@ def _get_elapsed_time_all_ranks(self, names, reset, barrier):
# pytorch yet. It is simpler to deal with a single tensor
# and since we are only gathering a small amount of data,
# it should be ok to use all-gather instead of gather.
- rank_name_to_time = torch.zeros((world_size, len(names)),
- dtype=torch.float,
- device=torch.cuda.current_device())
+ rank_name_to_time = torch.zeros(
+ (world_size, len(names)), dtype=torch.float, device=torch.cuda.current_device()
+ )
for i, name in enumerate(names):
if name in self._timers:
# Here we don't need to pass the barrier flag as all
# the processes are already in sync. This avoids the
# issue of different timers having different barrier
# groups inside their class.
- rank_name_to_time[rank, i] = self._timers[name].elapsed(
- reset=reset)
+ rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset)
# See the note above for why we are not using gather.
- torch.distributed._all_gather_base(rank_name_to_time.view(-1),
- rank_name_to_time[rank, :].view(-1))
+ torch.distributed._all_gather_base(
+ rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1)
+ )
return rank_name_to_time
-
def _get_global_min_max_time(self, names, reset, barrier, normalizer):
"""Report only min and max times across all ranks."""
- rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset,
- barrier)
+ rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
name_to_min_max_time = {}
for i, name in enumerate(names):
rank_to_time = rank_name_to_time[:, i]
@@ -217,32 +255,32 @@ def _get_global_min_max_time(self, names, reset, barrier, normalizer):
if rank_to_time.numel() > 0:
name_to_min_max_time[name] = (
rank_to_time.min().item() / normalizer,
- rank_to_time.max().item() / normalizer)
+ rank_to_time.max().item() / normalizer,
+ )
return name_to_min_max_time
-
- def _get_global_min_max_time_string(self, names, reset, barrier,
- normalizer, max_only):
- name_to_min_max_time = self._get_global_min_max_time(
- names, reset, barrier, normalizer)
+ def _get_global_min_max_time_string(self, names, reset, barrier, normalizer, max_only):
+ """Report strings for max/minmax times across all ranks."""
+ name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer)
if not name_to_min_max_time:
return None
- output_string = '(min, max) time across ranks (ms):'
+ if max_only:
+ output_string = 'max time across ranks (ms):'
+ else:
+ output_string = '(min, max) time across ranks (ms):'
for name in name_to_min_max_time:
min_time, max_time = name_to_min_max_time[name]
if max_only:
- output_string += '\n {}: {:.2f}'.format(
- (name+' ').ljust(48, '.'), max_time)
+ output_string += '\n {}: {:.2f}'.format((name + ' ').ljust(48, '.'), max_time)
else:
output_string += '\n {}: ({:.2f}, {:.2f})'.format(
- (name+' ').ljust(48, '.'), min_time, max_time)
+ (name + ' ').ljust(48, '.'), min_time, max_time
+ )
return output_string
-
def _get_all_ranks_time_string(self, names, reset, barrier, normalizer):
"""Report times across all ranks."""
- rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset,
- barrier)
+ rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
output_string = 'times across ranks (ms):'
no_reported_timing = True
@@ -255,49 +293,103 @@ def _get_all_ranks_time_string(self, names, reset, barrier, normalizer):
not_yet_found = False
output_string += '\n {}:'.format(name)
output_string += '\n rank {:2d}: {:.2f}'.format(
- rank, rank_name_to_time[rank, i] / normalizer)
+ rank, rank_name_to_time[rank, i] / normalizer
+ )
if no_reported_timing:
return None
return output_string
+ def get_all_timers_string(
+ self,
+ names: List[str] = None,
+ normalizer: float = 1.0,
+ reset: bool = True,
+ barrier: bool = False,
+ ):
+ """Returns the output string with logged timer values according to configured options.
+
+ Args:
+ names (List[str]): Names of the timers to log. If None, all registered timers are fetched. Defaults to None.
+ normalizer (float, optional): Normalizes the timer values by the factor. Defaults to 1.0.
+ reset (bool, optional): Whether to reset timer values after logging. Defaults to True.
+ barrier (bool, optional): Whether to do a global barrier before time measurments. Defaults to False.
+
+ Raises:
+ Exception: Raises if log option is invalid.
+
+ Returns:
+ str: Formatted string with the timer values.
+ """
- def log(self, names, rank=None, normalizer=1.0, reset=True, barrier=False):
- """Log a group of timers."""
+ if names == None: # get all registered timers
+ names = self._timers.keys()
- # Print.
assert normalizer > 0.0
if self._log_option in ['max', 'minmax']:
max_only = False
if self._log_option == 'max':
max_only = True
output_string = self._get_global_min_max_time_string(
- names, reset, barrier, normalizer/1000.0, max_only)
+ names, reset, barrier, normalizer / 1000.0, max_only
+ )
elif self._log_option == 'all':
- output_string = self._get_all_ranks_time_string(names,
- reset, barrier,
- normalizer/1000.0)
+ output_string = self._get_all_ranks_time_string(
+ names, reset, barrier, normalizer / 1000.0
+ )
else:
- raise Exception('unknown timing log option {}'.format(
- self._log_option))
+ raise Exception('unknown timing log option {}'.format(self._log_option))
+ return output_string
+
+ def log(
+ self,
+ names: List[str],
+ rank: int = None,
+ normalizer: float = 1.0,
+ reset: bool = True,
+ barrier: bool = False,
+ ):
+ """logs the timers passed in names to stdout. Example usage is to log average per step value for timer 'foo',
+ this function can be called with normalizer factor set to logging interval.
+
+ Args:
+ names (List[str]): Names of the timers to log.
+ rank (int, optional): logs the timers to a specific rank. If set to None, logs to the last rank. Defaults to None.
+ normalizer (float, optional): Normalizes the timer values by the factor. Defaults to 1.0.
+ reset (bool, optional): Whether to reset timer values after logging. Defaults to True.
+ barrier (bool, optional): Whether to do a global barrier before time measurments. Defaults to False.
+ """
+ output_string = self.get_all_timers_string(names, normalizer, reset, barrier)
# If no input rank is provided, log on last rank.
if rank is None:
rank = torch.distributed.get_world_size() - 1
if rank == torch.distributed.get_rank() and output_string is not None:
print(output_string, flush=True)
-
- def write(self, names, writer, iteration, normalizer=1.0,
- reset=False, barrier=False):
- """Write timers to a tensorboard writer
- Note that we only report maximum time across ranks to tensorboard.
+ def write(
+ self,
+ names: List[str],
+ writer,
+ iteration: int,
+ normalizer: float = 1.0,
+ reset: bool = True,
+ barrier: bool = False,
+ ):
+ """Write timers to a tensorboard writer. Note that we only report maximum time across ranks to tensorboard.
+
+ Args:
+ names (List[str]): Names of the timers to log.
+ writer (SummaryWriter): Tensorboard SummaryWriter object
+ iteration (int): Current iteration.
+ normalizer (float, optional): Normalizes the timer values by the factor. Defaults to 1.0.
+ reset (bool, optional): Whether to reset timer values after logging. Defaults to True.
+ barrier (bool, optional): Whether to do a global barrier before time measurments. Defaults to False.
"""
# currently when using add_scalars,
# torch.utils.add_scalars makes each timer its own run, which
# polutes the runs list, so we just add each as a scalar
assert normalizer > 0.0
- name_to_min_max_time = self._get_global_min_max_time(
- names, reset, barrier, normalizer)
+ name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer)
if writer is not None:
for name in name_to_min_max_time:
_, max_time = name_to_min_max_time[name]
diff --git a/megatron/core/transformer/__init__.py b/megatron/core/transformer/__init__.py
new file mode 100644
index 0000000000..0e3cdcfa57
--- /dev/null
+++ b/megatron/core/transformer/__init__.py
@@ -0,0 +1,6 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from .module import MegatronModule
+from .spec_utils import ModuleSpec, build_module
+from .transformer_config import MLATransformerConfig, TransformerConfig
+from .transformer_layer import TransformerLayer, TransformerLayerSubmodules
diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
new file mode 100644
index 0000000000..850dec88e1
--- /dev/null
+++ b/megatron/core/transformer/attention.py
@@ -0,0 +1,582 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.models.common.embeddings import apply_rotary_pos_emb
+from megatron.core.parallel_state import (
+ get_data_parallel_group,
+ get_data_parallel_rank,
+ get_data_parallel_world_size,
+ get_tensor_model_parallel_group,
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+)
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.utils import divide
+
+from .enums import AttnMaskType
+from .transformer_config import TransformerConfig
+
+try:
+ import transformer_engine # pylint: disable=unused-import
+
+ HAVE_TE = True
+ from megatron.core.transformer.custom_layers.transformer_engine import SplitAlongDim
+except ImportError:
+ HAVE_TE = False
+ SplitAlongDim = None
+
+
+@dataclass
+class SelfAttentionSubmodules:
+ linear_qkv: Union[ModuleSpec, type] = None
+ core_attention: Union[ModuleSpec, type] = None
+ linear_proj: Union[ModuleSpec, type] = None
+ q_layernorm: Union[ModuleSpec, type] = None
+ k_layernorm: Union[ModuleSpec, type] = None
+
+
+@dataclass
+class CrossAttentionSubmodules:
+ linear_q: Union[ModuleSpec, type] = None
+ linear_kv: Union[ModuleSpec, type] = None
+ core_attention: Union[ModuleSpec, type] = None
+ linear_proj: Union[ModuleSpec, type] = None
+
+
+class Attention(MegatronModule, ABC):
+ """Attention layer abstract class.
+
+ This layer only contains common modules required for the "self attn" and
+ "cross attn" specializations.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: Union[SelfAttentionSubmodules, CrossAttentionSubmodules],
+ layer_number: int,
+ attn_mask_type: AttnMaskType,
+ attention_type: str,
+ ):
+ super().__init__(config=config)
+
+ self.config = config
+ self.layer_number = layer_number
+ self.attn_mask_type = attn_mask_type
+ self.attention_type = attention_type
+
+ # For normal attention without groups, num_query_groups == num_attention_heads,
+ # so these two will be the same
+ self.query_projection_size = self.config.kv_channels * self.config.num_attention_heads
+ self.kv_projection_size = self.config.kv_channels * self.config.num_query_groups
+
+ # Per attention head and per partition values.
+ world_size = parallel_state.get_tensor_model_parallel_world_size()
+ self.hidden_size_per_attention_head = divide(
+ self.query_projection_size, self.config.num_attention_heads
+ )
+ self.num_attention_heads_per_partition = divide(self.config.num_attention_heads, world_size)
+ self.num_query_groups_per_partition = divide(self.config.num_query_groups, world_size)
+
+ self.core_attention = build_module(
+ submodules.core_attention,
+ config=self.config,
+ layer_number=self.layer_number,
+ attn_mask_type=self.attn_mask_type,
+ attention_type=self.attention_type,
+ )
+
+ self.checkpoint_core_attention = self.config.recompute_granularity == 'selective'
+
+ # Output.
+ self.linear_proj = build_module(
+ submodules.linear_proj,
+ self.query_projection_size,
+ self.config.hidden_size,
+ config=self.config,
+ init_method=self.config.output_layer_init_method,
+ bias=self.config.add_bias_linear,
+ input_is_parallel=True,
+ skip_bias_add=True,
+ is_expert=False,
+ tp_comm_buffer_name='proj',
+ )
+
+ def _checkpointed_attention_forward(
+ self,
+ query,
+ key,
+ value,
+ attention_mask,
+ rotary_pos_emb=None,
+ attn_mask_type=None,
+ packed_seq_params=None,
+ ):
+ """Forward method with selective activation checkpointing."""
+
+ def custom_forward(*inputs):
+ query = inputs[0]
+ key = inputs[1]
+ value = inputs[2]
+ attention_mask = inputs[3]
+ attn_mask_type = inputs[5]
+ attn_mask_type = AttnMaskType(attn_mask_type.item())
+ output_ = self.core_attention(
+ query,
+ key,
+ value,
+ attention_mask,
+ attn_mask_type=attn_mask_type,
+ packed_seq_params=packed_seq_params,
+ )
+ return output_
+
+ if attn_mask_type is None:
+ attn_mask_type = self.attn_mask_type
+ attn_mask_type = torch.tensor([attn_mask_type.value], dtype=torch.int)
+ hidden_states = tensor_parallel.checkpoint(
+ custom_forward, False, query, key, value, attention_mask, rotary_pos_emb, attn_mask_type
+ )
+
+ return hidden_states
+
+ def _allocate_memory(self, inference_max_sequence_length, batch_size, dim, dtype):
+ """Allocate memory to store kv cache during inference."""
+
+ return torch.empty(
+ inference_max_sequence_length,
+ batch_size,
+ self.num_query_groups_per_partition,
+ dim,
+ dtype=dtype,
+ device=torch.cuda.current_device(),
+ )
+
+ def _adjust_key_value_for_inference(self, inference_params, key, value, rotary_pos_emb):
+ """
+ Saves the generated key and value tensors to the end of the buffers in inference_params.
+ Returns the full size keys and values from the provided inference_params, as well as
+ adjusted rotary_pos_emb.
+
+ Returns a tuple: (key, value, rotary_pos_emb)
+
+ """
+ attn_mask_type = self.attn_mask_type
+ if inference_params is None:
+ return key, value, rotary_pos_emb, attn_mask_type
+
+ # =================================================
+ # Pre-allocate memory for key-values for inference.
+ # =================================================
+ if self.layer_number not in inference_params.key_value_memory_dict:
+ inf_max_seq_length = inference_params.max_sequence_length
+ inf_max_batch_size = inference_params.max_batch_size
+ inference_key_memory = self._allocate_memory(
+ inf_max_seq_length, inf_max_batch_size, key.shape[-1], key.dtype
+ )
+ inference_value_memory = self._allocate_memory(
+ inf_max_seq_length, inf_max_batch_size, value.shape[-1], value.dtype
+ )
+ inference_params.key_value_memory_dict[self.layer_number] = (
+ inference_key_memory,
+ inference_value_memory,
+ )
+ else:
+ # Get the pre-allocated buffers for this layer
+ inference_key_memory, inference_value_memory = inference_params.key_value_memory_dict[
+ self.layer_number
+ ]
+
+ if inference_params.sequence_len_offset > 0:
+ # This should mean that we are past the prompt forward_step
+ # and so we need to turn off masking
+ attn_mask_type = AttnMaskType.no_mask
+
+ batch_start = inference_params.batch_size_offset
+ batch_end = batch_start + key.size(1)
+ assert batch_end <= inference_key_memory.size(1)
+ sequence_start = inference_params.sequence_len_offset
+ sequence_end = sequence_start + key.size(0)
+ assert sequence_end <= inference_key_memory.size(0)
+ # Copy key and values.
+ inference_key_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = key
+ inference_value_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = value
+ key = inference_key_memory[:sequence_end, batch_start:batch_end, ...]
+ value = inference_value_memory[:sequence_end, batch_start:batch_end, ...]
+
+ # adjust the key rotary positional embedding
+ if rotary_pos_emb is None:
+ return key, value, rotary_pos_emb, attn_mask_type
+
+ q_pos_emb, k_pos_emb = rotary_pos_emb
+ q_pos_emb = q_pos_emb[sequence_start:sequence_end, :, :, :]
+ k_pos_emb = k_pos_emb[:sequence_end, :, :, :]
+ rotary_pos_emb = (q_pos_emb, k_pos_emb)
+
+ return key, value, rotary_pos_emb, attn_mask_type
+
+ @abstractmethod
+ def get_query_key_value_tensors(self, hidden_states, key_value_states):
+ """
+ This method needs to be implemented based on whether the derived class
+ is "self-attn" or "cross-attn".
+ """
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask,
+ key_value_states=None,
+ inference_params=None,
+ rotary_pos_emb=None,
+ packed_seq_params=None,
+ ):
+ # hidden_states: [sq, b, h]
+
+ # For self attention we just duplicate the rotary_pos_emb if it isn't already
+ if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple):
+ rotary_pos_emb = (rotary_pos_emb,) * 2
+
+ # =====================
+ # Query, Key, and Value
+ # =====================
+ # Get the query, key and value tensors based on the type of attention -
+ # self or cross attn.
+ query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states)
+
+ # ===================================================
+ # Adjust key, value, and rotary_pos_emb for inference
+ # ===================================================
+ key, value, rotary_pos_emb, attn_mask_type = self._adjust_key_value_for_inference(
+ inference_params, key, value, rotary_pos_emb
+ )
+
+ if packed_seq_params is not None:
+ query = query.squeeze(1)
+ key = key.squeeze(1)
+ value = value.squeeze(1)
+
+ # ================================================
+ # relative positional embedding (rotary embedding)
+ # ================================================
+ if rotary_pos_emb is not None:
+ q_pos_emb, k_pos_emb = rotary_pos_emb
+
+ if packed_seq_params is not None:
+ cu_seqlens_q = packed_seq_params.cu_seqlens_q
+ cu_seqlens_kv = packed_seq_params.cu_seqlens_kv
+ else:
+ cu_seqlens_q = cu_seqlens_kv = None
+ query = apply_rotary_pos_emb(
+ query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q
+ )
+ key = apply_rotary_pos_emb(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv)
+
+ # TODO, can apply positional embedding to value_layer so it has
+ # absolute positional embedding.
+ # otherwise, only relative positional embedding takes effect
+ # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb)
+
+ # ==================================
+ # core attention computation
+ # ==================================
+
+ if self.checkpoint_core_attention and self.training:
+ core_attn_out = self._checkpointed_attention_forward(
+ query,
+ key,
+ value,
+ attention_mask,
+ attn_mask_type=attn_mask_type,
+ packed_seq_params=packed_seq_params,
+ )
+ else:
+ core_attn_out = self.core_attention(
+ query,
+ key,
+ value,
+ attention_mask,
+ attn_mask_type=attn_mask_type,
+ packed_seq_params=packed_seq_params,
+ )
+
+ if packed_seq_params is not None:
+ # reshape to same output shape as unpacked case
+ # (t, np, hn) -> (t, b=1, h=np*hn)
+ # t is the pack size = sum (sq_i)
+ # note that batch is a dummy dimension in the packed case
+ core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1)
+
+ # =================
+ # Output. [sq, b, h]
+ # =================
+
+ output, bias = self.linear_proj(core_attn_out)
+
+ return output, bias
+
+
+class SelfAttention(Attention):
+ """Self-attention layer class
+
+ Self-attention layer takes input with size [s, b, h]
+ and returns output of the same size.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: SelfAttentionSubmodules,
+ layer_number: int,
+ attn_mask_type=AttnMaskType.padding,
+ ):
+ super().__init__(
+ config=config,
+ submodules=submodules,
+ layer_number=layer_number,
+ attn_mask_type=attn_mask_type,
+ attention_type="self",
+ )
+
+ self.linear_qkv = build_module(
+ submodules.linear_qkv,
+ self.config.hidden_size,
+ self.query_projection_size + 2 * self.kv_projection_size,
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=self.config.add_bias_linear or self.config.add_qkv_bias,
+ skip_bias_add=False,
+ is_expert=False,
+ tp_comm_buffer_name='qkv',
+ )
+
+ if submodules.q_layernorm is not None:
+ self.q_layernorm = build_module(
+ submodules.q_layernorm,
+ hidden_size=self.hidden_size_per_attention_head,
+ config=self.config,
+ eps=self.config.layernorm_epsilon,
+ )
+ else:
+ self.q_layernorm = None
+
+ if submodules.k_layernorm is not None:
+ self.k_layernorm = build_module(
+ submodules.k_layernorm,
+ hidden_size=self.hidden_size_per_attention_head,
+ config=self.config,
+ eps=self.config.layernorm_epsilon,
+ )
+ else:
+ self.k_layernorm = None
+
+ def run_realtime_tests(self):
+ """Performs a consistency check.
+
+ This function makes sure that tensors across devices are the same during an experiment.
+ This is often not guaranteed to be so because of silent hardware failures (eg, memory
+ corruption loading a checkpoint, network traffic corruption encountered during
+ data transmission).
+
+ (TODO) In the future, more tensors should be checked across the training run and
+ checked every X iterations. This is left for future work. Equality of tensors is probably
+ not required; transmitting hashes is sufficient."""
+
+ if not self.config.qk_layernorm:
+ return
+
+ # check that all tensor parallel and data parallel ranks have the same
+ # Q & K layernorm parameters.
+ rank = get_data_parallel_rank()
+ inputs = torch.stack(
+ [
+ self.q_layernorm.weight.data,
+ self.q_layernorm.bias.data,
+ self.k_layernorm.weight.data,
+ self.k_layernorm.bias.data,
+ ]
+ )
+ dp_list = [torch.empty_like(inputs) for _ in range(get_data_parallel_world_size())]
+ dp_list[rank] = inputs
+ torch.distributed.all_gather(dp_list, inputs, group=get_data_parallel_group())
+
+ def _compare(srcs, tgts, names, parallelism):
+ assert len(srcs) == len(tgts) == len(names)
+ for src, tgt, name in zip(srcs, tgts, names):
+ assert torch.all(src == tgt), (
+ f"Discrepancy between {name} in {parallelism} ranks {i} and {rank}. "
+ f"Diff: {torch.norm(src - tgt)}"
+ )
+
+ for i, dp in enumerate(dp_list):
+ q_w, q_b, k_w, k_b = torch.unbind(dp)
+ _compare(
+ [q_w, q_b, k_w, k_b],
+ [
+ self.q_layernorm.weight.data,
+ self.q_layernorm.bias.data,
+ self.k_layernorm.weight.data,
+ self.k_layernorm.bias.data,
+ ],
+ ["q_w", "q_b", "k_w", "k_b"],
+ "DP",
+ )
+
+ rank = get_tensor_model_parallel_rank()
+ tp_list = [torch.empty_like(inputs) for _ in range(get_tensor_model_parallel_world_size())]
+ tp_list[rank] = inputs
+ torch.distributed.all_gather(tp_list, inputs, group=get_tensor_model_parallel_group())
+
+ for i, tp in enumerate(tp_list):
+ q_w, q_b, k_w, k_b = torch.unbind(tp)
+ _compare(
+ [q_w, q_b, k_w, k_b],
+ [
+ self.q_layernorm.weight.data,
+ self.q_layernorm.bias.data,
+ self.k_layernorm.weight.data,
+ self.k_layernorm.bias.data,
+ ],
+ ["q_w", "q_b", "k_w", "k_b"],
+ "TP",
+ )
+
+ def get_query_key_value_tensors(self, hidden_states, key_value_states=None):
+ """
+ Derives `query`, `key` and `value` tensors from `hidden_states`.
+ """
+ # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)]
+ mixed_qkv, _ = self.linear_qkv(hidden_states)
+
+ # [sq, b, hp] --> [sq, b, ng, (np/ng + 2) * hn]
+ new_tensor_shape = mixed_qkv.size()[:-1] + (
+ self.num_query_groups_per_partition,
+ (
+ (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2)
+ * self.hidden_size_per_attention_head
+ ),
+ )
+ mixed_qkv = mixed_qkv.view(*new_tensor_shape)
+
+ split_arg_list = [
+ (
+ self.num_attention_heads_per_partition
+ // self.num_query_groups_per_partition
+ * self.hidden_size_per_attention_head
+ ),
+ self.hidden_size_per_attention_head,
+ self.hidden_size_per_attention_head,
+ ]
+
+ if SplitAlongDim is not None:
+
+ # [sq, b, ng, (np/ng + 2) * hn]
+ # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn]
+ (query, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list)
+ else:
+
+ # [sq, b, ng, (np/ng + 2) * hn]
+ # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn]
+ (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3)
+
+ # [sq, b, ng, np/ng * hn] -> [sq, b, np, hn]
+ query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head)
+
+ if self.q_layernorm is not None:
+ query = self.q_layernorm(query)
+
+ if self.k_layernorm is not None:
+ key = self.k_layernorm(key)
+
+ if self.config.test_mode:
+ self.run_realtime_tests()
+
+ return query, key, value
+
+
+class CrossAttention(Attention):
+ """Cross-attention layer class
+
+ Cross-attention layer takes input with size [s, b, h] and context with size
+ [s, b, h] and returns output of the same size.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: CrossAttentionSubmodules,
+ layer_number: int,
+ attn_mask_type=AttnMaskType.padding,
+ ):
+ super().__init__(
+ config=config,
+ submodules=submodules,
+ layer_number=layer_number,
+ attn_mask_type=attn_mask_type,
+ attention_type="cross",
+ )
+
+ if self.config.num_query_groups != self.config.num_attention_heads:
+ raise ValueError("Group query attention is not currently supported in cross attention.")
+ assert self.query_projection_size == self.kv_projection_size
+
+ self.linear_q = build_module(
+ submodules.linear_q,
+ self.config.hidden_size,
+ self.query_projection_size,
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=self.config.add_bias_linear,
+ skip_bias_add=False,
+ is_expert=False,
+ )
+
+ self.linear_kv = build_module(
+ submodules.linear_kv,
+ self.config.hidden_size,
+ 2 * self.kv_projection_size,
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=self.config.add_bias_linear,
+ skip_bias_add=False,
+ is_expert=False,
+ )
+
+ def get_query_key_value_tensors(self, hidden_states, key_value_states):
+ """
+ Derives `query` tensor from `hidden_states`, and `key`/`value` tensors
+ from `key_value_states`.
+ """
+ # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)]
+ mixed_kv, _ = self.linear_kv(key_value_states)
+
+ # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn]
+ new_tensor_shape = mixed_kv.size()[:-1] + (
+ self.num_attention_heads_per_partition,
+ 2 * self.hidden_size_per_attention_head,
+ )
+ mixed_kv = mixed_kv.view(*new_tensor_shape)
+
+ # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn]
+ (key, value) = tensor_parallel.split_tensor_along_last_dim(mixed_kv, 2)
+
+ # Attention head [sq, b, h] --> [sq, b, hp]
+ query, _ = self.linear_q(hidden_states)
+
+ # [sq, b, hp] --> [sq, b, np, hn]
+ new_tensor_shape = query.size()[:-1] + (
+ self.num_attention_heads_per_partition,
+ self.hidden_size_per_attention_head,
+ )
+ query = query.view(*new_tensor_shape)
+
+ return query, key, value
diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py
new file mode 100644
index 0000000000..2588980b5b
--- /dev/null
+++ b/megatron/core/transformer/cuda_graphs.py
@@ -0,0 +1,313 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+import time
+from enum import Enum
+
+import torch
+
+from megatron.core.transformer.module import MegatronModule
+
+try:
+ from transformer_engine.pytorch import make_graphed_callables
+ from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
+
+ HAVE_TE_GRAPHS = True
+except:
+ HAVE_TE_GRAPHS = False
+
+
+class GraphStatus(Enum):
+ """An Enum to track if a cudagraph is ready to perform a forward or backward pass."""
+
+ FWD_READY = 0
+ BWD_READY = 1
+
+
+class GraphStatusFunc(torch.autograd.Function):
+ """Inserts a node into the autograd graph that tracks whether an object has an outstanding
+ backward pass by toggling the value of GraphStatus. This is mainly used to detect when to create
+ multiple graphs per transformer layer for pipeline parallelism.
+ We don't use backward module hooks as they change forward output tensors to views, see:
+ https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_hook
+ """
+
+ @staticmethod
+ def forward(ctx, runner, obj):
+ """Occurs immediately before the graph's forward pass.
+ Marks the graph's backward pass as ready."""
+ ctx.runner = runner
+ runner.status = GraphStatus.BWD_READY
+ return obj
+
+ @staticmethod
+ def backward(ctx, grad):
+ """Occurs immediately after the graph's backward pass.
+ Marks the graph's forward pass as ready."""
+ assert ctx.runner.status == GraphStatus.BWD_READY
+ ctx.runner.status = GraphStatus.FWD_READY
+ return None, grad
+
+
+class TensorDescription:
+ """Records the attributes of a tensor. Used to check if a
+ tensor argument matches the tensor with which the module
+ was graph captured with."""
+
+ def __init__(self, tensor):
+ self.shape = tuple(tensor.shape)
+ self.dtype = tensor.dtype
+ self.device = tensor.device
+
+ def matches_tensor(self, tensor):
+ """Check if 'tensor' matches the attributes of this TensorDescription."""
+
+ assert torch.is_tensor(tensor)
+ return (
+ tensor.shape == self.shape
+ and tensor.dtype == self.dtype
+ and tensor.device == self.device
+ )
+
+
+class CudaGraphCallable(torch.nn.Module):
+ """Wraps a module to be cudagraphable, records the output of the cudagraph.
+ Reinserts non-tensor args, kwargs that were previously filtered out by 'get_tensor_args'.
+ """
+
+ def __init__(self, module, groundtruth_args, groundtruth_kwargs):
+ super().__init__()
+ self.add_module('base_module', module)
+
+ # The Pytorch cudagraph API requires only tensor inputs, so we strip
+ # non-tensor arguments and reinsert them in forward() using these groundtruth attributes.
+ # We will also check future calls to the cudagraph against these to ensure the cudagraph
+ # is called with the same inputs as it was captured with.
+ self.groundtruth_outputs = []
+ self.groundtruth_args = tuple(
+ TensorDescription(a) if torch.is_tensor(a) else a for a in groundtruth_args
+ )
+ self.groundtruth_kwargs = {
+ k: TensorDescription(v) if torch.is_tensor(v) else v
+ for k, v in groundtruth_kwargs.items()
+ }
+
+ def forward(self, *arg_tensors, **kwarg_tensors):
+ """Call the forward pass of the cudagraph. Also checks the outputs
+ of the cudagraph matches what the graph was traced with."""
+
+ args = list(self.groundtruth_args)
+ arg_tensors = list(arg_tensors)
+ for idx, groundtruth_arg in enumerate(self.groundtruth_args):
+ if isinstance(groundtruth_arg, TensorDescription):
+ args[idx] = arg_tensors.pop(0)
+
+ kwargs = dict(self.groundtruth_kwargs)
+ for k, v in self.groundtruth_kwargs.items():
+ if isinstance(v, TensorDescription):
+ kwargs[k] = kwarg_tensors[k]
+
+ # Use forward() instead of __call__ to avoid triggering hooks
+ out = self.base_module.forward(*args, **kwargs)
+ if torch.is_tensor(out):
+ out = tuple(out)
+
+ self.groundtruth_outputs = [TensorDescription(o) if torch.is_tensor(o) else o for o in out]
+
+ out = tuple(o for o in out if torch.is_tensor(o))
+ assert (
+ len(out) > 0
+ ), """A graphed module returned no tensors in training mode, however the graphed module
+ must output at least one tensor, so that a corresponding backward node
+ may be registered in the autograd graph."""
+
+ if len(out) == 1:
+ return out[0]
+ return out
+
+
+class CudaGraphRunner(torch.nn.Module):
+ """Wraps a single cudagraph and its expected arguments. Checks that
+ the provided args are the same as what the graph was traced with.
+ """
+
+ def __init__(self, graphed_module, wrapped_module):
+ super().__init__()
+
+ self.graphed_module = graphed_module
+ self.groundtruth_args = wrapped_module.groundtruth_args
+ self.groundtruth_kwargs = wrapped_module.groundtruth_kwargs
+ self.groundtruth_outputs = wrapped_module.groundtruth_outputs
+ self.status = GraphStatus.FWD_READY
+
+ def static_args_match(self, args, kwargs):
+ """Check the the passed args, kwargs match with the arg, kwargs
+ the graph was created with."""
+
+ def check(val, ref):
+ if isinstance(ref, TensorDescription):
+ return ref.matches_tensor(val)
+ return ref == val
+
+ if len(args) != len(self.groundtruth_args):
+ return False
+ for idx, groundtruth_arg in enumerate(self.groundtruth_args):
+ if not check(args[idx], groundtruth_arg):
+ return False
+
+ if kwargs.keys() != self.groundtruth_kwargs.keys():
+ return False
+ for k, v in self.groundtruth_kwargs.items():
+ if not check(kwargs[k], v):
+ return False
+ return True
+
+ def forward(self, args, kwargs, is_first_microbatch=None):
+ """Call the forward pass of the cuda graph."""
+ if self.training and torch.is_grad_enabled():
+ args = list(args)
+ for pos in range(len(args)):
+ if torch.is_tensor(args[pos]):
+ args[pos] = GraphStatusFunc.apply(self, args[pos])
+ for k, v in kwargs.items():
+ if torch.is_tensor(v):
+ kwargs[k] = GraphStatusFunc.apply(self, v)
+
+ ret_tensors = self.graphed_module(is_first_microbatch=is_first_microbatch, *args, **kwargs)
+ ret_tensors = [ret_tensors] if torch.is_tensor(ret_tensors) else list(ret_tensors)
+ out = tuple(
+ ret_tensors.pop(0) if isinstance(o, TensorDescription) else o
+ for o in self.groundtruth_outputs
+ )
+
+ # Check that the static graph matches what was recorded during graph capture
+ assert len(out) == len(self.groundtruth_outputs)
+ for idx, o in enumerate(self.groundtruth_outputs):
+ if isinstance(o, TensorDescription):
+ assert o.matches_tensor(out[idx])
+ else:
+ assert o == out[idx]
+
+ if len(out) == 1:
+ return out[0]
+ return out
+
+
+class CudaGraphManager(torch.nn.Module):
+ """Creates and runs cudagraphs for a megatron module."""
+
+ def __init__(self):
+ super().__init__()
+ self.cudagraph_runners = []
+ self.is_first_microbatch = True
+ assert HAVE_TE_GRAPHS, "CudaGraphManager currently requires TransformerEngine"
+
+ # Cudagraph stream capture requires no operations on the default stream prior to the
+ # capture, so change to a side stream. At graph capture change it back.
+ self.stream = torch.cuda.current_stream()
+ torch.cuda.set_stream(torch.cuda.Stream())
+
+ def __call__(self, megatron_module, args, kwargs):
+ """Calls the forward pass of the cudagraphed module.
+
+ Args:
+ megatron_module (torch.nn.module): The megatron module to be graphed and run
+
+ args (tuple): The positional args to be passed to the module.
+
+ kwargs (dict): The keyword args to be passed to the module.
+
+ """
+
+ # param.data_ptr() below is used to trigger any hooks that have attached to the parameter.
+ # Specifically, this is trying to trigger the param sync hook for the APEX optimizer, which
+ # triggers param syncs by hooking into any param references.
+ # However cudagraphs disables this, so we workaround by manually referencing params here.
+ # For more information see:
+ # https://github.com/NVIDIA/apex/blob/7001836/apex/contrib/optimizers/distributed_fused_adam.py#L885C9
+ for param in megatron_module.parameters():
+ param.data_ptr()
+
+ runner = None
+ for _runner in self.cudagraph_runners:
+ if _runner.static_args_match(args, kwargs) and _runner.status == GraphStatus.FWD_READY:
+ runner = _runner
+ break
+
+ if runner is None:
+ if self.training and torch.is_grad_enabled():
+ runner = self.create_cudagraph_module(megatron_module, args, kwargs)
+ self.cudagraph_runners.append(runner)
+ logging.getLogger(__name__).info(
+ f"Creating cudagraph; now have {len(self.cudagraph_runners)}"
+ )
+ else:
+ # No cudagraphs were found in inference mode, so fallback to eager since
+ # tensor.requires_grad is needed to correctly trace the backward graph.
+ return super(MegatronModule, megatron_module).__call__(*args, **kwargs)
+
+ tensor_args, tensor_kwargs = self.get_tensor_args(args, kwargs)
+ out = runner(tensor_args, tensor_kwargs, is_first_microbatch=self.is_first_microbatch)
+ self.is_first_microbatch = False
+ return out
+
+ def get_tensor_args(self, args, kwargs):
+ """Filter out non-tensor arguments from args and kwargs.
+ Needed since 'make_graphed_callables' expects Torch.tensor arg, kwargs."""
+ tensor_kwargs = {}
+ for k, v in kwargs.items():
+ if torch.is_tensor(v):
+ tensor_kwargs[k] = v
+ tensor_args = tuple(arg for arg in args if torch.is_tensor(arg))
+ return tensor_args, tensor_kwargs
+
+ def create_cudagraph_module(self, megatron_module, args, kwargs):
+ """Record the graph capture stream. Runs warmup iterations of
+ megatron_module, and creates a autograd function, where the
+ forward, backward functions are the cudagraphs of module's forward,
+ backward passes. Finally wraps this cudagraph function with a CudaGraphRunner.
+ """
+
+ torch.cuda.synchronize()
+ torch.cuda.set_stream(self.stream)
+ start = time.time()
+
+ wrapped_module = CudaGraphCallable(megatron_module, args, kwargs)
+ sample_args, sample_kwargs = self.get_tensor_args(args, kwargs)
+
+ # Cudagraphs require no autograd history recorded on sample inputs
+ sample_args_detached = tuple(n.detach() for n in sample_args)
+ sample_kwargs_detached = {k: v.detach() for k, v in sample_kwargs.items()}
+ sample_args_copy = tuple(torch.clone(n) for n in sample_args_detached)
+ sample_kwargs_copy = {k: torch.clone(v) for k, v in sample_kwargs_detached.items()}
+
+ # Zero out input args inplace so cudagraph warmup doesnt affect grads
+ for orig, detach in zip(sample_args, sample_args_detached):
+ detach.zero_()
+ detach.requires_grad = orig.requires_grad
+ for k, detach in sample_kwargs_detached.items():
+ detach.zero_()
+ detach.requires_grad = sample_kwargs[k].requires_grad
+
+ fp8_enabled = megatron_module.config.fp8 is not None
+ fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_enabled else None
+ graphed_module = make_graphed_callables(
+ modules=wrapped_module,
+ sample_args=sample_args_detached,
+ sample_kwargs=sample_kwargs_detached,
+ _order=[1, -1],
+ allow_unused_input=True,
+ fp8_enabled=fp8_enabled,
+ fp8_recipe=fp8_recipe,
+ fp8_weight_caching=True,
+ )
+
+ # Restore zeroed out sample args
+ # Detach again since pytorch prohibits inplace ops on leaf nodes
+ for orig, copy in zip(sample_args, sample_args_copy):
+ orig.detach().copy_(copy)
+ for k, orig in sample_kwargs.items():
+ orig.detach().copy_(sample_kwargs_copy[k])
+
+ logging.getLogger(__name__).info(f'Time spent in cudagraph capture: {time.time() - start}s')
+ return CudaGraphRunner(graphed_module, wrapped_module)
diff --git a/megatron/core/transformer/custom_layers/__init__.py b/megatron/core/transformer/custom_layers/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/megatron/core/transformer/custom_layers/transformer_engine.py b/megatron/core/transformer/custom_layers/transformer_engine.py
new file mode 100644
index 0000000000..02ce9ad5a7
--- /dev/null
+++ b/megatron/core/transformer/custom_layers/transformer_engine.py
@@ -0,0 +1,12 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import warnings
+
+warnings.warn(
+ """The 'megatron.core.transformer.custom_layers.transformer_engine'
+ module is deprecated and will be removed in 0.10.0. Please use
+ 'megatron.core.extensions.transformer_engine' instead.""",
+ DeprecationWarning,
+ stacklevel=2,
+)
+from megatron.core.extensions.transformer_engine import *
diff --git a/megatron/core/transformer/dot_product_attention.py b/megatron/core/transformer/dot_product_attention.py
new file mode 100644
index 0000000000..d5c014cabf
--- /dev/null
+++ b/megatron/core/transformer/dot_product_attention.py
@@ -0,0 +1,202 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+
+import math
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.fusions.fused_softmax import FusedScaleMaskSoftmax
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import attention_mask_func
+from megatron.core.utils import divide
+
+
+class DotProductAttention(MegatronModule):
+ """
+ Region where selective activation recomputation is applied.
+ This region is memory intensive but less compute intensive which
+ makes activation checkpointing more efficient for LLMs (20B+).
+ See Reducing Activation Recomputation in Large Transformer Models:
+ https://arxiv.org/abs/2205.05198 for more details.
+
+ We use the following notation:
+ h: hidden size
+ n: number of attention heads
+ p: number of tensor model parallel partitions
+ b: batch size
+ s: sequence length
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ layer_number: int,
+ attn_mask_type: AttnMaskType,
+ attention_type: str,
+ attention_dropout: float = None,
+ softmax_scale: float = None,
+ ):
+ super().__init__(config=config)
+
+ self.config: TransformerConfig = config
+
+ assert (
+ self.config.context_parallel_size == 1
+ ), "Context parallelism is only supported by TEDotProductAttention!"
+
+ assert (
+ self.config.window_size is None
+ ), "Sliding Window Attention is only supported by TEDotProductAttention!"
+
+ self.layer_number = max(1, layer_number)
+ self.attn_mask_type = attn_mask_type
+ self.attention_type = attention_type # unused for now
+
+ projection_size = self.config.kv_channels * self.config.num_attention_heads
+
+ # Per attention head and per partition values.
+ world_size = parallel_state.get_tensor_model_parallel_world_size()
+ self.hidden_size_per_partition = divide(projection_size, world_size)
+ self.hidden_size_per_attention_head = divide(projection_size, config.num_attention_heads)
+ self.num_attention_heads_per_partition = divide(self.config.num_attention_heads, world_size)
+ self.num_query_groups_per_partition = divide(self.config.num_query_groups, world_size)
+
+ coeff = None
+ if softmax_scale is None:
+ self.softmax_scale = 1.0 / math.sqrt(self.hidden_size_per_attention_head)
+ else:
+ self.softmax_scale = softmax_scale
+
+ if self.config.apply_query_key_layer_scaling:
+ coeff = self.layer_number
+ self.softmax_scale /= coeff
+
+ self.scale_mask_softmax = FusedScaleMaskSoftmax(
+ input_in_fp16=self.config.fp16,
+ input_in_bf16=self.config.bf16,
+ attn_mask_type=self.attn_mask_type,
+ scaled_masked_softmax_fusion=self.config.masked_softmax_fusion,
+ mask_func=attention_mask_func,
+ softmax_in_fp32=self.config.attention_softmax_in_fp32,
+ scale=coeff,
+ )
+
+ # Dropout. Note that for a single iteration, this layer will generate
+ # different outputs on different number of parallel partitions but
+ # on average it should not be partition dependent.
+ self.attention_dropout = torch.nn.Dropout(
+ self.config.attention_dropout if attention_dropout is None else attention_dropout
+ )
+
+ def forward(
+ self,
+ query: Tensor,
+ key: Tensor,
+ value: Tensor,
+ attention_mask: Tensor,
+ attn_mask_type: AttnMaskType = None,
+ packed_seq_params: Optional[PackedSeqParams] = None,
+ ):
+ assert packed_seq_params is None, (
+ "Packed sequence is not supported by DotProductAttention."
+ "Please use TEDotProductAttention instead."
+ )
+
+ # ===================================
+ # Raw attention scores. [b, n/p, s, s]
+ # ===================================
+
+ # expand the key and value [sk, b, ng, hn] -> [sk, b, np, hn]
+ # This is a noop for normal attention where ng == np. When using group query attention this
+ # creates a view that has the keys and values virtually repeated along their dimension to
+ # match the number of queries.
+
+ # attn_mask_type is not used.
+ if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1:
+ key = key.repeat_interleave(
+ self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
+ )
+ value = value.repeat_interleave(
+ self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2
+ )
+
+ # [b, np, sq, sk]
+ output_size = (query.size(1), query.size(2), query.size(0), key.size(0))
+
+ # [sq, b, np, hn] -> [sq, b * np, hn]
+ # This will be a simple view when doing normal attention, but in group query attention
+ # the key and value tensors are repeated to match the queries so you can't use
+ # simple strides to extract the queries.
+ query = query.reshape(output_size[2], output_size[0] * output_size[1], -1)
+ # [sk, b, np, hn] -> [sk, b * np, hn]
+ key = key.view(output_size[3], output_size[0] * output_size[1], -1)
+
+ # preallocting input tensor: [b * np, sq, sk]
+ matmul_input_buffer = parallel_state.get_global_memory_buffer().get_tensor(
+ (output_size[0] * output_size[1], output_size[2], output_size[3]), query.dtype, "mpu"
+ )
+
+ # Raw attention scores. [b * np, sq, sk]
+ matmul_result = torch.baddbmm(
+ matmul_input_buffer,
+ query.transpose(0, 1), # [b * np, sq, hn]
+ key.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk]
+ beta=0.0,
+ alpha=self.softmax_scale,
+ )
+
+ # change view to [b, np, sq, sk]
+ attention_scores = matmul_result.view(*output_size)
+
+ # ===========================
+ # Attention probs and dropout
+ # ===========================
+
+ # attention scores and attention mask [b, np, sq, sk]
+ attention_probs: Tensor = self.scale_mask_softmax(attention_scores, attention_mask)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+
+ if not self.config.sequence_parallel:
+ with tensor_parallel.get_cuda_rng_tracker().fork():
+ attention_probs = self.attention_dropout(attention_probs)
+ else:
+ attention_probs = self.attention_dropout(attention_probs)
+
+ # =========================
+ # Context layer. [sq, b, hp]
+ # =========================
+
+ # value -> context layer.
+ # [sk, b, np, hn] --> [b, np, sq, hn]
+
+ # context layer shape: [b, np, sq, hn]
+ output_size = (value.size(1), value.size(2), query.size(0), value.size(3))
+
+ # change view [sk, b * np, hn]
+ value = value.view(value.size(0), output_size[0] * output_size[1], -1)
+
+ # change view [b * np, sq, sk]
+ attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
+
+ # matmul: [b * np, sq, hn]
+ context = torch.bmm(attention_probs, value.transpose(0, 1))
+
+ # change view [b, np, sq, hn]
+ context = context.view(*output_size)
+
+ # [b, np, sq, hn] --> [sq, b, np, hn]
+ context = context.permute(2, 0, 1, 3).contiguous()
+
+ # [sq, b, np, hn] --> [sq, b, hp]
+ new_context_shape = context.size()[:-2] + (self.hidden_size_per_partition,)
+ context = context.view(*new_context_shape)
+
+ return context
diff --git a/megatron/core/transformer/enums.py b/megatron/core/transformer/enums.py
new file mode 100644
index 0000000000..99d0ddefbd
--- /dev/null
+++ b/megatron/core/transformer/enums.py
@@ -0,0 +1,28 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+import enum
+
+
+# can we get rid of this?
+# it's being used in pipeline schedules
+class ModelType(enum.Enum):
+ encoder_or_decoder = 1
+ encoder_and_decoder = 2
+
+
+# class LayerType(enum.Enum):
+# encoder = 1
+# decoder = 2
+
+
+class AttnType(enum.Enum):
+ self_attn = 1
+ cross_attn = 2
+
+
+class AttnMaskType(enum.Enum):
+ padding = 1
+ causal = 2
+ no_mask = 3 # only used for TE
+ padding_causal = 4 # only used for thd attention
+ arbitrary = 5
diff --git a/megatron/core/transformer/identity_op.py b/megatron/core/transformer/identity_op.py
new file mode 100644
index 0000000000..5d9388ffcc
--- /dev/null
+++ b/megatron/core/transformer/identity_op.py
@@ -0,0 +1,28 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+import torch
+
+
+class IdentityOp(torch.nn.Module):
+ """
+ This is a placeholder for IdentityOp(x) -> x
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__()
+
+ def forward(self, x, *args, **kwargs):
+ return x
+
+
+class IdentityFuncOp(IdentityOp):
+ """
+ This is a placeholder for IdentityFuncOp(...)(x) -> IdentityOp(x) -> x.
+ Such a func is handy for ops like `bias_dropout_fusion` which themselves
+ return a function at runtime based on passed arguments
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__()
+
+ def forward(self, *args, **kwargs):
+ return super().forward
diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py
new file mode 100644
index 0000000000..e82d6ecd20
--- /dev/null
+++ b/megatron/core/transformer/mlp.py
@@ -0,0 +1,255 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from megatron.core import parallel_state
+from megatron.core.dist_checkpointing import ShardedTensor
+from megatron.core.dist_checkpointing.mapping import (
+ ReplicaId,
+ ShardedStateDict,
+ ShardedTensorFactory,
+)
+from megatron.core.fusions.fused_bias_geglu import bias_geglu_impl
+from megatron.core.fusions.fused_bias_gelu import bias_gelu_impl
+from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+
+
+@dataclass
+class MLPSubmodules:
+ linear_fc1: Union[ModuleSpec, type] = None
+ linear_fc2: Union[ModuleSpec, type] = None
+
+
+class MLP(MegatronModule):
+ """
+ MLP will take the input with h hidden state, project it to 4*h
+ hidden dimension, perform nonlinear transformation, and project the
+ state back into h hidden dimension.
+
+
+ Returns an output and a bias to be added to the output.
+ If config.add_bias_linear is False, the bias returned is None.
+
+ We use the following notation:
+ h: hidden size
+ p: number of tensor model parallel partitions
+ b: batch size
+ s: sequence length
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: MLPSubmodules,
+ is_expert: bool = False,
+ input_size: int = None,
+ ):
+ super().__init__(config=config)
+
+ self.config: TransformerConfig = config
+
+ self.input_size = input_size if input_size != None else self.config.hidden_size
+
+ # If this is a gated linear unit we double the output width, see https://arxiv.org/pdf/2002.05202.pdf
+ ffn_hidden_size = self.config.ffn_hidden_size
+ if self.config.gated_linear_unit:
+ ffn_hidden_size *= 2
+
+ self.linear_fc1 = build_module(
+ submodules.linear_fc1,
+ self.input_size,
+ ffn_hidden_size,
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=self.config.add_bias_linear,
+ skip_bias_add=True,
+ is_expert=is_expert,
+ tp_comm_buffer_name='fc1',
+ )
+
+ self.activation_func = self.config.activation_func
+
+ self.linear_fc2 = build_module(
+ submodules.linear_fc2,
+ self.config.ffn_hidden_size,
+ self.config.hidden_size,
+ config=self.config,
+ init_method=self.config.output_layer_init_method,
+ bias=self.config.add_bias_linear,
+ input_is_parallel=True,
+ skip_bias_add=True,
+ is_expert=is_expert,
+ tp_comm_buffer_name='fc2',
+ )
+
+ def forward(self, hidden_states):
+
+ # [s, b, 4 * h/p]
+ intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states)
+
+ if self.config.bias_activation_fusion:
+ if self.activation_func == F.gelu:
+ if self.config.gated_linear_unit:
+ intermediate_parallel = bias_geglu_impl(intermediate_parallel, bias_parallel)
+ else:
+ assert self.config.add_bias_linear is True
+ intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel)
+ elif self.activation_func == F.silu and self.config.gated_linear_unit:
+ intermediate_parallel = bias_swiglu_impl(
+ intermediate_parallel,
+ bias_parallel,
+ self.config.activation_func_fp8_input_store,
+ )
+ else:
+ raise ValueError("Only support fusion of gelu and swiglu")
+ else:
+ if bias_parallel is not None:
+ intermediate_parallel = intermediate_parallel + bias_parallel
+ if self.config.gated_linear_unit:
+
+ def glu(x):
+ x = torch.chunk(x, 2, dim=-1)
+ return self.config.activation_func(x[0]) * x[1]
+
+ intermediate_parallel = glu(intermediate_parallel)
+ else:
+ intermediate_parallel = self.activation_func(intermediate_parallel)
+
+ # [s, b, h]
+ output, output_bias = self.linear_fc2(intermediate_parallel)
+
+ return output, output_bias
+
+ def sharded_state_dict(
+ self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None
+ ) -> ShardedStateDict:
+ sharded_state_dict = {}
+ for name, module in self._modules.items():
+ sub_sd = module.sharded_state_dict(f'{prefix}{name}.', sharded_offsets, metadata)
+ if self.config.gated_linear_unit and name == 'linear_fc1':
+ assert f'{prefix}{name}.weight' in sub_sd, sub_sd.keys()
+ for k, v in sub_sd.items():
+ if k in (f'{prefix}{name}.weight', f'{prefix}{name}.bias'):
+ sub_sd[k] = apply_swiglu_sharded_factory(v, sharded_offsets)
+ sharded_state_dict.update(sub_sd)
+ return sharded_state_dict
+
+
+def apply_swiglu_sharded_factory(original_sh_ten, sharded_offsets):
+ # We must split the tensor into 2 parts, each sharded separately.
+ # This requires a ShardedTensorFactory which `chunk`s during saving
+ # and `cat`s during loading
+ tp_rank = parallel_state.get_tensor_model_parallel_rank()
+ tp_size = parallel_state.get_tensor_model_parallel_world_size()
+ swiglu_shard_axis = 0
+ prepend_axis_num = len(sharded_offsets)
+ original_shape = original_sh_ten.local_shape
+ original_numel = int(np.prod(original_shape))
+
+ @torch.no_grad()
+ def sh_ten_build_fn(
+ key: str, t: torch.Tensor, replica_id: ReplicaId, flattened_range: Optional[slice]
+ ):
+ offset_w = (swiglu_shard_axis + prepend_axis_num, tp_rank, tp_size * 2)
+ offset_v = (swiglu_shard_axis + prepend_axis_num, tp_size + tp_rank, tp_size * 2)
+ if flattened_range is None:
+ tensor_w, tensor_v = torch.chunk(t, 2, dim=swiglu_shard_axis)
+ return [
+ ShardedTensor.from_rank_offsets(
+ key,
+ tensor_w,
+ *sharded_offsets,
+ offset_w,
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ ),
+ ShardedTensor.from_rank_offsets(
+ key,
+ tensor_v,
+ *sharded_offsets,
+ offset_v,
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ ),
+ ]
+ else:
+ # Here we need to map a slice `t` (`flattened_range` specifies slice start and stop)
+ # of the *original* flattened tensor into slices `w` and `v` of chunked
+ # and flattened tensor.
+ # Example:
+ # If original tensor has (16, 5) shape and flattened_range is `slice(8, 64)`,
+ # then `t` has shape `(56,)` and we need to create 2 tensors:
+ # w: first 32 elements of `t` with flattened_range slice(8, 40)
+ # v: last 24 elements of `t` with flattened_range slice(0, 24)
+ # Global offsets are the same as in the non-flattened case
+ assert t.ndim == 1, (key, t.shape)
+ non_flat_local_shape = (original_shape[0] // 2, *original_shape[1:])
+ chunk_numel = original_numel // 2
+ result = []
+ if flattened_range.start < chunk_numel:
+ # Non-empty `w` chunk
+ tensor_w = t[: chunk_numel - flattened_range.start]
+ flattened_range_w = slice(
+ flattened_range.start, min(chunk_numel, flattened_range.stop)
+ )
+ assert len(tensor_w) == flattened_range_w.stop - flattened_range_w.start
+ result.append(
+ ShardedTensor.from_rank_offsets_flat(
+ key,
+ tensor_w,
+ non_flat_local_shape,
+ *sharded_offsets,
+ offset_w,
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ flattened_range=flattened_range_w,
+ )
+ )
+ if flattened_range.stop > chunk_numel:
+ # Non-empty `v` chunk
+ tensor_v = t[-(flattened_range.stop - chunk_numel) :]
+ flattened_range_v = slice(
+ max(chunk_numel, flattened_range.start) - chunk_numel,
+ flattened_range.stop - chunk_numel,
+ )
+ assert len(tensor_v) == flattened_range_v.stop - flattened_range_v.start, (
+ len(tensor_v),
+ flattened_range_v,
+ )
+
+ result.append(
+ ShardedTensor.from_rank_offsets_flat(
+ key,
+ tensor_v,
+ non_flat_local_shape,
+ *sharded_offsets,
+ offset_v,
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ flattened_range=flattened_range_v,
+ )
+ )
+ assert sum(sh_ten.data.numel() for sh_ten in result) == t.numel(), (result, t.shape)
+ return result
+
+ def sh_ten_merge_fn(sub_state_dict):
+ with torch.no_grad():
+ return torch.cat(sub_state_dict)
+
+ return ShardedTensorFactory(
+ original_sh_ten.key,
+ original_sh_ten.data,
+ sh_ten_build_fn,
+ sh_ten_merge_fn,
+ original_sh_ten.replica_id,
+ )
diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py
new file mode 100644
index 0000000000..c89acec400
--- /dev/null
+++ b/megatron/core/transformer/module.py
@@ -0,0 +1,195 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Megatron Module."""
+from typing import Optional, Tuple
+
+import torch
+from torch.autograd import Variable
+from torch.nn.parameter import Parameter
+
+from megatron.core import parallel_state
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import (
+ make_sharded_tensors_for_checkpoint,
+ sharded_state_dict_default,
+)
+
+_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor)
+_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor)
+_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor)
+
+
+def param_is_not_shared(param):
+ return not hasattr(param, 'shared') or not param.shared
+
+
+class MegatronModule(torch.nn.Module):
+ """Base Megatron module inhertied by all Models.
+
+ Megatron specific extensions of torch Module with support
+ for pipelining
+
+ Args:
+ config (TransformerConfig): Transformer config
+ """
+
+ # def __init__(self, config: TransformerConfig, share_word_embeddings=True):
+ def __init__(self, config: TransformerConfig):
+ super().__init__()
+ self.config = config
+
+ def state_dict_for_save_checkpoint(self, prefix: str = '', keep_vars: bool = False):
+ """Override state dict for saving checkpoints Use this function to override the
+ state dict for saving checkpoints.
+
+ Args:
+ prefix (str, optional): _description_. Defaults to ''.
+ keep_vars (bool, optional): _description_. Defaults to False.
+
+ Returns:
+ _type_: _description_
+ """
+
+ return self.state_dict(prefix=prefix, keep_vars=keep_vars)
+
+ def sharded_state_dict(
+ self,
+ prefix: str = '',
+ sharded_offsets: Tuple[Tuple[int, int, int]] = (),
+ metadata: Optional[dict] = None,
+ ) -> ShardedStateDict:
+ """Default implementation for sharded state dict for distributed checkpointing.
+
+ General definition of sharded_state_dict simply calls `sharded_state_dict_default`
+ (which call sharded_state_dict method if possible or a default implementation otherwise)
+ recursively on all submodules.
+
+ Args:
+ prefix (str): prefix for the state dict keys
+ sharded_offsets (Tuple[Tuple[int, int, int]], optional): sharding already
+ applied (e.g. PP related) by sup-modules. Passed along to ShardedTensor
+ metadata (dict, optional): metadata passed recursively to sharded_state_dict methods
+
+ Returns:
+ dict: dictionary of state dict keys mapped to ShardedTensors
+ """
+ sharded_state_dict = {}
+ # Save parameters
+ self._save_to_state_dict(sharded_state_dict, '', keep_vars=True)
+ sharded_state_dict = make_sharded_tensors_for_checkpoint(
+ sharded_state_dict, prefix, sharded_offsets=sharded_offsets
+ )
+ # Recurse into submodules
+ for name, module in self.named_children():
+ sharded_state_dict.update(
+ sharded_state_dict_default(module, f'{prefix}{name}.', sharded_offsets, metadata)
+ )
+ return sharded_state_dict
+
+ def set_is_first_microbatch(self):
+ """Sets the is_first_microbatch flag if it exists and config.fp8==True.
+ When this flag is set, TE modules will update their fp8 parameter cache.
+ """
+ if self.config.fp8 is not None:
+ if not hasattr(self, "modules_with_is_first_microbatch"):
+ self.modules_with_is_first_microbatch = []
+ for m in self.modules():
+ if hasattr(m, "is_first_microbatch"):
+ self.modules_with_is_first_microbatch.append(m)
+ for m in self.modules_with_is_first_microbatch:
+ m.is_first_microbatch = True
+
+
+def conversion_helper(val, conversion):
+ if not isinstance(val, (tuple, list)):
+ return conversion(val)
+ rtn = [conversion_helper(v, conversion) for v in val]
+ if isinstance(val, tuple):
+ rtn = tuple(rtn)
+ return rtn
+
+
+def fp32_to_float16(val, float16_convertor):
+ def half_conversion(val):
+ val_typecheck = val
+ if isinstance(val_typecheck, (Parameter, Variable)):
+ val_typecheck = val.data
+ if isinstance(val_typecheck, _FLOAT_TYPES):
+ val = float16_convertor(val)
+ return val
+
+ return conversion_helper(val, half_conversion)
+
+
+def float16_to_fp32(val):
+ def float_conversion(val):
+ val_typecheck = val
+ if isinstance(val_typecheck, (Parameter, Variable)):
+ val_typecheck = val.data
+ if isinstance(val_typecheck, (_BF16_TYPES, _HALF_TYPES)):
+ val = val.float()
+ return val
+
+ return conversion_helper(val, float_conversion)
+
+
+class Float16Module(MegatronModule):
+ """Float 16 Module.
+
+ Attributes:
+ config (TransformerConfig): Transformer config
+ fp16 (bool) : Specifies if the model runs in fp16 mode
+ bf16 (bool) : Specifies if the model runs in bf16 mode
+
+ Args:
+ config (TransformerConfig): The transformer config used to initalize the model
+ """
+
+ def __init__(self, config: TransformerConfig, module: torch.nn.Module):
+ super(Float16Module, self).__init__(config)
+ self.config = config
+ self.fp16 = config.fp16
+ self.bf16 = config.bf16
+
+ if self.fp16:
+ self.add_module('module', module.half())
+
+ def float16_convertor(val):
+ return val.half()
+
+ elif self.bf16:
+ self.add_module('module', module.bfloat16())
+
+ def float16_convertor(val):
+ return val.bfloat16()
+
+ else:
+ raise Exception('Either config.fp16 or config.bf16 should be True.')
+
+ self.float16_convertor = float16_convertor
+
+ def set_input_tensor(self, input_tensor):
+ return self.module.set_input_tensor(input_tensor)
+
+ def forward(self, *inputs, **kwargs):
+ if parallel_state.is_pipeline_first_stage():
+ inputs = fp32_to_float16(inputs, self.float16_convertor)
+ outputs = self.module(*inputs, **kwargs)
+ if parallel_state.is_pipeline_last_stage():
+ outputs = float16_to_fp32(outputs)
+ return outputs
+
+ def state_dict(self, destination=None, prefix='', keep_vars=False):
+ return self.module.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars)
+
+ def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False):
+ """Retrieve state_dict from the module being wrapped."""
+ return self.module.state_dict_for_save_checkpoint(prefix=prefix, keep_vars=keep_vars)
+
+ def sharded_state_dict(self, prefix='', *args, **kwargs):
+ """Retrieve sharded_state_dict from the module being wrapped."""
+ return self.module.sharded_state_dict(prefix, *args, **kwargs)
+
+ def load_state_dict(self, state_dict, strict=True):
+ self.module.load_state_dict(state_dict, strict=strict)
diff --git a/megatron/core/transformer/moe/README.md b/megatron/core/transformer/moe/README.md
new file mode 100644
index 0000000000..a7ee75bcbf
--- /dev/null
+++ b/megatron/core/transformer/moe/README.md
@@ -0,0 +1,366 @@
+# Megatron Core MoE Key Features
+
+Megatron-Core offers rich parallelism mappings, combining Expert Parallelism with tensor, data, sequence, and pipeline parallelism. This boosts Mixtral 8X7B bf16 training to achieve **438 TFLOPS** as of MCore v0.8.
+
+
+### Parallelism
+- **Expert Parallelism**
+ - A specific method of parallelism for MoE models, where experts are partitioned onto different workers and each worker processes a different batch of training samples, each worker process one or more experts for each MoE layer.
+- **3D Parallelism**: Data Parallelism, Tensor Parallelism, Pipeline Parallelism
+ - Note: When using MoE with expert parallelism and tensor parallelism, sequence parallelism must be enabled.
+- **Context Parallelism**:
+ - Split the sequence dimension to support long context training.
+- **Richer parallel mappings**: EP can be combined with DP/TP/PP/CP for handling larger MoE variants.
+- **Full distributed optimizer support.**
+
+### Router and Load Balancing
+- Router type:
+ - Top-K MLP router
+- Load Balancing algorithms:
+ - Sinkhorn (S-BASE)
+ - Aux loss / Load balancing loss
+
+### Performance Optimizations
+- GroupedGEMM when num local experts > 1
+ - Supported dtype: bf16
+ - Performance improvements for larger MoE models
+- Enable `--tp-comm-overlap` for MoE
+
+### Token Dispatch Mechanism
+- Dropless / No token drop
+- Token drop, with or without padding to capacity
+
+### Ease of use
+- Checkpoint converter for Mixtral models, see the [example](https://github.com/NVIDIA/Megatron-LM/tree/main/examples/mixtral) for details.
+- Distributed checkpoining
+- Per-layer logging
+
+## Upcoming features
+- Token permutation / unpermutation fusion
+- Fused Sinkhorn Kernel
+- FP8 training support
+
+# User Guide
+
+### MoE Related Arguments
+
+| Item | Description |
+| --- | --- |
+| --num-experts | Number of Experts in MoE (None means no MoE) |
+| --expert-model-parallel-size | Degree of expert model parallelism. Default is 1. |
+| --moe-grouped-gemm | When there are multiple experts per rank, launch multiple local GEMM kernels in multiple streams to improve the utilization and performance with GroupedLinear in TransformerEngine. |
+| --moe-router-load-balancing-type | Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer, "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss". |
+| --moe-router-topk | Number of experts to route to for each token. The default is 2. |
+| --moe-aux-loss-coeff | Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended. Default is 0.0. |
+| --moe-z-loss-coeff | Scaling coefficient for the z-loss: a starting value of 1e-3 is recommended. Default is None. |
+| --moe-input-jitter-eps | Add noise to the input tensor by applying jitter with a specified epsilon value. Default is None. |
+| --moe-token-dispatcher-type | Determines the token dispatcher type. Choices are "allgather", "alltoall" and "alltoall_seq". Default is "allgather". We recommend using 'alltoall' if expert parallelism is applied. We have upgraded the "alltoall" dispatcher in place during MCore v0.9, while retaining the original implementation, renamed as "alltoall_seq".|
+| --moe-per-layer-logging | Enable per-layer logging for MoE, currently supports auxiliary loss and z loss. |
+| --moe-expert-capacity-factor | The capacity factor for each expert, None means no token will be dropped. Default is None. |
+| --moe-pad-expert-input-to-capacity | Pads the input for each expert to match the expert capacity length, effective only after the --moe-expert-capacity-factor is set. |
+| --moe-token-drop-policy | The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with the lowest probabilities will be dropped. If "position", tokens at the end of each batch will be dropped. |
+| --moe-layer-recompute | Enable activation checkpointing for moe_layer, should be used when memory is not sufficient. |
+| --moe-extended-tp | (Experimental) Alternative parallelization strategy for expert parallelism. Instead of distributing experts across *expert_model_parallel_size*, each expert is sharded along extendended tensor parallel domain (tensor_model_paralle_size * expert_model_parallel_size). It avoids the load balancing problem with MOE training. Only available with `--moe-token-dispatcher-type allgather`. |
+| --moe-shared-expert-intermediate-size | Set shared expert total ffn hidden size. It should be equal to `num_shared_experts * ffn_size_of_each_shared_expert` if there are multiple shared experts. None means no shared expert. |
+| --moe-shared-expert-overlap | (Experimental, may changed) If this is set, the communications/computations in the shared experts and the dispatcher will overlap (The `alltoall` dispatcher is needed.) Otherwise, the shared expert runs after the routed experts. |
+| --moe-use-upcycling | Load the dense model checkpoint, convert it into an MoE model at runtime and start training. The converted model will be saved to the path specified by `--save` before training begins. Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.|
+
+
+## Usage
+
+### Quick Start
+To train a top-2 MoE model with 8 experts and auxiliary loss, include the following arguments:
+
+```bash
+--num-experts 8
+--expert-model-parallel-size 8
+--moe-grouped-gemm
+--moe-router-load-balancing-type aux_loss # options: aux_loss, sinkhorn, none. Default is aux_loss.
+--moe-router-topk 2
+--moe-aux-loss-coeff 1e-2
+--use-distributed-optimizer
+--moe-token-dispatcher-type alltoall
+```
+
+To enable the token drop mechanism, such as GShard and SwitchTransformer, include the following arguments:
+
+```bash
+--moe-expert-capacity-factor 1.0
+--moe-pad-expert-input-to-capacity # Optional
+```
+
+The following figure illustrates differenting dropping strategies in MCore:
+
+
+
+1. The default dropless strategy will not drop or pad any token.
+2. By setting `--moe-expert-capacity-factor`, the tokens exceed the capacity of expert will be dropped based on their selected probabilities.
+ The dropping is performed before the token exchange operation between EP ranks when EP > 1.
+ The formula of capacity is `capacity = num_tokens_per_rank * topk * capacity_factor / num_experts`.
+3. By setting `--moe-pad-expert-input-to-capacity`, the experts with tokens less than capacity will be padded to the capacity.
+
+### Fine-tuning Mixtral Models
+Megatron-Core has full support for Mixtral MoE models, and we provide the checkpoint converter for Mixtral models from huggingface format to MCore format.
+
+
+### Distributed Checkpointing
+MCore v0.7 introduced fully parallel and asynchronous saving capabilities to distributed checkpointing,
+which addresses the issues of low efficiency in the traditional checkpoint saving methods.
+It also solved the problem of incompatibility between checkpoints of different parallel mappings in the traditional format.
+With the new distributed checkpointing solution, MCore can achieve flexible parallelism configurations by saving and loading the unified format checkpoints.
+Compared to native PyTorch solution, MCore achieves up to 50x reduction in checkpointing overhead.
+
+From MCore v0.8, MoE supports Distributed Checkpointing, which means users can save and load with any combination of parallelism and it is currently available, including expert parallel.
+1. Loading weight and distributed optimizer states with TPxCPxEPxPP resharding with SequentialMLP is supported in version 0.8.
+2. GroupedMLP weight resharding is supported in version 0.8.0 and optimizer state resharding is supported in version 0.10.0. Switching between GroupedMLP/SequentialMLP when loading and saving is partially supported.
+3. TEGroupedMLP has fully support on distributed checkpointing and is fully exchangable with SequentialMLP in version 0.9.0.
+4. Optimizer state resharding cannot do across EP=1 with EP>1 due to the different optimizer type.
+
+Usage
+- `--ckpt-format torch_dist` The main argument, it will attempt to save and load using distributed checkpointing.
+- `--auto-detect-ckpt-format` With this, it can load both distributed checkpointing and legacy checkpointing.
+
+Checkpoint compatibility across SequentialMLP, GroupedMLP, and TEGroupedMLP:
+```text
+ ┌───────────────┐ ┌───────────────┐ ┌───────────────┐
+ │ GroupedMLP │ │ SequentialMLP │ │ TEGroupedMLP │
+ │ │ │ │ │ │
+ │ │ │ │ │ │
+ │ ┌───────────┐ │ │ ┌───────────┐ │ │ ┌───────────┐ │
+ │ │legacy ckpt│ │ │ │legacy ckpt│ │ │ │legacy ckpt│ │
+ │ └─────┬─────┘ │ │ └─────┬─────┘ │ │ └─────┬─────┘ │
+ │ ▼ │ │ ▼ │ │ ▼ │
+ │ ┌─────────┐ │ │ ┌─────────┐ │ │ ┌─────────┐ │
+ │ │dist ckpt│ │ │ │dist ckpt│ │ │ │dist ckpt│ │
+┌──►│ │ weight │ │◄────────►│ │ weight │ │◄────────►│ │ weight │ │◄──┐
+│ │ └─────────┘ │ │ └─────────┘ │ │ └─────────┘ │ │
+└───┼───────────────┼──────────┼───────────────┼──────────┼───────────────┼───┘
+ │┌─────────────┐│ │┌─────────────┐│ │┌─────────────┐│
+ ││ dist ckpt ││ ││ dist ckpt ││ ││ dist ckpt ││
+ ││optim states ││ ││optim states ││◄────────►││optim states ││
+ │└─────────────┘│ │└─────────────┘│ │└─────────────┘│
+ └───────────────┘ └───────────────┘ └───────────────┘
+```
+
+Best practices for distributed checkpointing:
+1. Convert a legacy checkpoint to a distributed checkpoint. To achieve this, we can add both `--ckpt-format torch_dist --auto-detect-ckpt-format`, then it will load the legacy one and save as the distributed checkpoint format later when the training progress tries to save checkpoints.
+2. Convert checkpoint of the legacy GroupedMLP to TEGroupedMLP. This is only supported for the weight parts. To achieve this, we can use the above method to convert the legacy checkpoint to a distributed checkpoint of the legacy GroupedMLP. After updating the libraries and using TEGroupedMLP, we can directly load the previously saved checkpoint by adding argument `--no-load-optim`.
+
+### Shared Experts
+MCore v0.9 introduced the shared expert feature. We can enable this feature by setting suitable `--moe-shared-expert-intermediate-size`.
+
+The parallelism patterns of the shared experts follow the settings of the dense part, i.e., the attention module. The shared experts are not distributed but replicated in EP ranks.
+
+We also have an experimental feature that tries to overlap the communications and computations in the shared experts and the dispatcher.
+We can set `--moe-shared-expert-overlap` and use `alltoall` dispatcher to enable it.
+The overlapping relies on the envirionment setting `CUDA_DEVICE_MAX_CONNECTIONS=1`.
+The `AllGather` and `ReduceScatter` communications in the shared experts are overlapped with `permute`/`unpermute` in the dispatcher.
+The `MLP` computation part in the shared experts are overlapped with the `AlltoAll` communications in the dispatcher.
+Both the forward and the backward pass can overlap. But to get the overlapping in the backward pass, the PyTorch version should `>= 2.2.0`.
+
+### Upcycling
+Use `--moe-use-upcycling` to enable the upcycling feature, which will load the dense model from the directory specified by `--load`, convert it into an MoE model at runtime and start training. The converted model will be saved to the path specified by `--save` before training begins. Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.
+
+The MoE model structure is defined through script arguments. All MoE-related arguments (such as `--num-experts`) can be customized; however, other model structure arguments must be consistent with those of the dense model.
+
+## MoE training example:
+
+Click here.
+
+```bash
+#!/bin/bash
+
+# Runs Mixtral 8x7B model on 32 H100/A100 GPUs
+# The Dropless MoE suffers from an imbalanced token distribution at the early stage of training (the first few hundred iterations), which may lead to poor performance and out-of-memory (OOM) issues.
+# To check the performance of a Dropless MoE model, we should run the model for at least 500 iterations or resume from trained checkpoints.
+
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+GPUS_PER_NODE=8
+# Change for multinode config
+MASTER_ADDR=${MASTER_ADDR:-"localhost"}
+MASTER_PORT=${MASTER_PORT:-"6000"}
+NNODES=${NNODES:-"1"}
+NODE_RANK=${RANK:-"0"}
+WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))
+
+CHECKPOINT_PATH=$1
+TOKENIZER_MODEL=$2
+DATA_PATH=$3
+
+DISTRIBUTED_ARGS=(
+ --nproc_per_node $GPUS_PER_NODE
+ --nnodes $NNODES
+ --node_rank $NODE_RANK
+ --master_addr $MASTER_ADDR
+ --master_port $MASTER_PORT
+)
+
+MODEL_ARGS=(
+ --disable-bias-linear
+ --seq-length 4096
+ --max-position-embeddings 32768
+ --num-layers 32
+ --hidden-size 4096
+ --ffn-hidden-size 14336
+ --num-attention-heads 32
+ --init-method-std 0.01
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --normalization RMSNorm
+ --position-embedding-type rope
+ --swiglu
+ --untie-embeddings-and-output-weights
+ --group-query-attention
+ --num-query-groups 8
+ --no-masked-softmax-fusion
+ --no-position-embedding
+)
+
+MOE_ARGS=(
+ --num-experts 8
+ --expert-model-parallel-size 8
+ --moe-router-load-balancing-type aux_loss # options: aux_loss, sinkhorn, None. Default is aux_loss.
+ --moe-router-topk 2
+ --moe-aux-loss-coeff 1e-2
+ --moe-grouped-gemm
+)
+
+DATA_ARGS=(
+ --tokenizer-type Llama2Tokenizer
+ --tokenizer-model ${TOKENIZER_MODEL}
+ --data-path $DATA_PATH
+ --split 99990,8,2
+)
+
+TRAINING_ARGS=(
+ --micro-batch-size 1
+ --global-batch-size 128
+ --lr 1e-4
+ --train-iters 500000
+ --lr-decay-iters 320000
+ --lr-decay-style cosine
+ --min-lr 1.0e-5
+ --weight-decay 0.1
+ --lr-warmup-iters 500
+ --clip-grad 1.0
+ --bf16
+ --overlap-grad-reduce
+ --overlap-param-gather
+)
+
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size 1
+ --pipeline-model-parallel-size 4
+ --num-layers-per-virtual-pipeline-stage 8
+ --sequence-parallel
+ --use-distributed-optimizer
+)
+
+LOGGING_ARGS=(
+ --log-interval 1 \
+ --save-interval 10000 \
+ --eval-interval 1000 \
+ --eval-iters 10 \
+ --save $CHECKPOINT_PATH \
+ --load $CHECKPOINT_PATH \
+ --tensorboard-dir "${CHECKPOINT_PATH}/tensorboard" \
+ --no-load-optim \
+ --no-load-rng
+)
+
+if [ -n "${WANDB_API_KEY}" ]; then
+ LOGGING_ARGS+=(
+ --wandb-project ${WANDB_PROJECT:-"Mixtral-Finetuning"}
+ --wandb-exp-name ${WANDB_NAME:-"Mixtral_8x7B"}
+ )
+fi
+
+torchrun ${DISTRIBUTED_ARGS[@]} pretrain_gpt.py \
+ ${MODEL_ARGS[@]} \
+ ${MOE_ARGS[@]} \
+ ${DATA_ARGS[@]} \
+ ${TRAINING_ARGS[@]} \
+ ${MODEL_PARALLEL_ARGS[@]} \
+ ${LOGGING_ARGS[@]}
+```
+
+
+# Performance Best Practice
+
+### Tuning Guide of Parallel Mappings
+
+To find a good parallel mapping that help you achieve a high throughput of a new model, there are some general rule that could help. Here is an overview of properties in different aspects for each parallel strategy.
+
+| Parallel Strategy | Peak Activation Memory | Weight Memory | Optimizer states | Communication (Per-Layer) |
+|:-----------------:|:-------------------------------:|:--------------:|:---------------------------------:|:-------------------------:|
+| TP | 1/N (with SP on) | 1/N | 1/N | High |
+| EP | 1 | 1/N in MoELayer| 1/N | Medium |
+| PP | 1 (>1 with virtual pipeline) | 1/N | 1/N | Medium |
+| CP | 1/N | 1 | 1/N (with distributed optimizer) | Medium |
+| DP | 1 | 1 | 1/N (with distributed optimizer) | Low |
+
+For a specific model, the best parallel mapping varies based on the model architecture, trained sequence length and the hardware platform.
+Here we provide some general rules to get better performance:
+1. Keep the model parallism size as small as possible.
+ - For the large language models, model parallism is often required to prevent OOM, but it will bring communication overhead and hurt performance.
+ - With distributed optimizer, master weights and optimizer states will be sharded across all DP ranks with slight communication overhead.
+ So try to reduce the model parallism size and increase data parallism size when there are lots of free GPU memory during training.
+2. Ensure the EPxTP communication winthin the NVLink domain.
+ - Communications of EP and TP should remain within the NVLink domain as much as possible, as both are communication-intensive.
+ - If the model is too large and requires scaling across multiple nodes, consider PP before TP and EP. See item 3 for details.
+3. Use Pipeline Parallelism to scale the model further.
+ - Enable Virtual Pipeline Parallelism(VPP) to reduce pp bubbles when PP_size >= 2 by setting `num_layers_per_virtual_pipeline_stage`.
+ - VPP_size tuning: the legal values of vpp_size are all common divisors of num_layers/pp_size, E.g., num_layers=24, pp_size=4, then we can pick vpp_size from {1, 2, 3, 6}. The larger the vpp_size, the lower the pipeline bubbles, while the larger number of P2P communications between each PP stages. Empirically a value in the middle often gives the best trade-off. `VPP_size=num_layers / PP_size / num_layers_per_virtual_pipeline_stage`
+4. Prefer EP over TP for the expert layer when possible:
+ - TP saves more memory than EP, but EP can achieve better GEMM efficiency and less communication overhead than TP.
+ - If EP size increased to the number of expert, the local token permutation/un-permutation for experts computation are omitted.
+ - Simplify the computation graph of MoE layers, more convenient for performing potential comm-computation overlapping.
+ - In practice, EP8TP1 is better than EP4TP2 for 8x7B.
+5. Enable Context Parallelism for long context training.
+ - The efficiency of CP largely depends on whether its communication can be overlapped with computation.
+ - Emperically, use CP when sequence length >= 8K.
+
+
+### End-to-End Training Practice
+**Use the latest NVIDIA PyTorch or NeMo Docker Image**
+- [NGC PyTorch Image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)
+- [NGC NeMo Image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo)
+
+**Token Dispatcher Choices**
+- Token Dispatcher sends tokens to the designated expert, involves tensor rearangement and communications.
+- Dispatcher `allgather` is the default option. It achieves better performance and efficiency when only tensor parallelism is used or when the Top-k value is very large.
+- Dispatcher `alltoall` is recommended if expert parallelism is applied.
+- Dispatcher `alltoall_seq` is the original implementation of `alltoall` and is retained for potential compatibility risk.
+
+**Enable Communication Overlap**
+- Enable `--overlap-param-gather` and `--overlap-grad-reduce` with distributed optimizer.
+- Enable `--tp-comm-overlap` when TP>1.
+- Enable p2p comm overlap when PP > 1 by setting `num_layers_per_virtual_pipeline_stage`.
+
+**Enable GroupedGEMM when num_local_experts>1 with `--moe-grouped-gemm`**
+- GroupedGEMM has higher efficiency than vanilla sequential GEMMs for each expert.
+- Recommend to use the TE version of Grouped GEMM (by upgrading to MCore v0.8 and TE v1.9), which support Gradient Accumulation Fusion and FP8 Training.
+
+**OOM Caused by Token Distribution Imbalance when Training From Scratch**
+MoE suffers from a severe load imbalance issue when the router is under-trained, leading to the model easily running out of memory (OOM), which typically occurs in the first 100~300 steps when training from scratch.
+Therefore, there are two recommended ways during the first 200 steps to avoid the OOM problem, which can be removed after the token distribution is more stable:
+1. Use Extended-TP(`-moe-extended-tp`) to replace EP with TP in MoELayer, this can prevent the load imbalancing between EP ranks. Since current ETP implementation has some memeory overhead, you can further enable activation recomputation only for MoE Layer by adding `--moe-layer-recompute`.
+2. Setting capacity factor to a relatively small number like 1.0 by adding `--moe-token-capacity-factor 1.0`.
+
+### Reference Best Parallel Mapping
+
+Here are the reference parallel mappings of MCore v0.8 for Mixtral 8x7B and 8x22B models:
+| Model | Vocab Size| Dispatcher | Precision | #GPUs | SEQ LEN | TP | EP | PP | VP | MBS | GBS |
+|:-----------------------:|:---------:|:----------:|:---------:|:-----:|:-------:|:--:|:--:|:--:|:--:|:---:|:---:|
+| Mixtral 8x7B(Dropless) | 32K | All-to-All | BF16 | 64 | 4096 | 1 | 8 | 4 | 8 | 1 | 256 |
+| Mixtral 8x22B(Dropless) | 32K | All-to-All | BF16 | 128 | 4096 | 4 | 2 | 8 | 7 | 1 | 256 |
+
+Detailed Benchmark Information:
+Server:
+- 8xH100 80GB HBM3
+- NVLink 4th Generation
+- InfiniBand 8x400 Gbit/s
+
+Docker Image:
+- PyTorch 24.04 with TransformerEngine v1.9
\ No newline at end of file
diff --git a/megatron/core/transformer/moe/__init__.py b/megatron/core/transformer/moe/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
new file mode 100644
index 0000000000..1bb5da588b
--- /dev/null
+++ b/megatron/core/transformer/moe/experts.py
@@ -0,0 +1,808 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import itertools
+from copy import deepcopy
+from functools import partial
+from math import ceil
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter
+
+from megatron.core import parallel_state
+from megatron.core.dist_checkpointing import ShardedTensor
+from megatron.core.dist_checkpointing.mapping import (
+ LocalNonpersistentObject,
+ ReplicaId,
+ ShardedStateDict,
+ ShardedTensorFactory,
+)
+from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
+from megatron.core.fusions.fused_bias_geglu import bias_geglu_impl
+from megatron.core.fusions.fused_bias_gelu import bias_gelu_impl
+from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl
+from megatron.core.jit import jit_fuser
+from megatron.core.tensor_parallel.layers import (
+ _initialize_affine_weight_cpu,
+ _initialize_affine_weight_gpu,
+)
+from megatron.core.tensor_parallel.utils import divide
+from megatron.core.transformer.mlp import MLP, MLPSubmodules, apply_swiglu_sharded_factory
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.moe import grouped_gemm_util as gg
+from megatron.core.transformer.spec_utils import build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.utils import make_sharded_object_for_checkpoint
+
+
+class GroupedMLP(MegatronModule):
+ """An efficient implementation of the Experts layer using GroupedGEMM.
+
+ Executes multiple experts in parallel to maximize computational efficiency.
+ """
+
+ def __init__(self, num_local_experts: int, config: TransformerConfig):
+ super().__init__(config=config)
+ self.config: TransformerConfig = config
+ self.num_local_experts = num_local_experts
+ gg.assert_grouped_gemm_is_available()
+ assert (
+ config.add_bias_linear == False
+ ), "bias not supported in Grouped GEMM yet, please set '--disable-bias-linear' instead."
+
+ self.expert_parallel = config.expert_model_parallel_size > 1
+ if self.config.gated_linear_unit:
+ if self.config.activation_func not in (F.silu, F.gelu):
+ raise ValueError("Activation function must be silu or gelu when using GroupedMLP.")
+
+ @jit_fuser
+ def glu(x):
+ x = torch.chunk(x, 2, dim=-1)
+ return self.config.activation_func(x[0]) * x[1]
+
+ self.activation_func = glu
+ else:
+ self.activation_func = self.config.activation_func
+
+ # How many feature each rank holds for fc1 and fc2, respectively.
+ self.moe_extended_tp = config.moe_extended_tp
+ if config.moe_extended_tp:
+ tp_size = parallel_state.get_tensor_and_expert_parallel_world_size()
+ else:
+ tp_size = parallel_state.get_tensor_model_parallel_world_size()
+
+ fc1_output_size = self.config.ffn_hidden_size * self.num_local_experts
+ if config.gated_linear_unit:
+ # Project to 4h. If using swiglu double the output width,
+ # see https://arxiv.org/pdf/2002.05202.pdf
+ fc1_output_size *= 2
+ fc1_output_size_per_partition = divide(fc1_output_size, tp_size)
+
+ fc2_input_size = self.config.ffn_hidden_size * self.num_local_experts
+ fc2_input_size_per_partition = divide(fc2_input_size, tp_size)
+
+ # Note: The current kernel implementations of grouped_gemm
+ # does not support transposition with CUTLASS grouped GEMM
+ # (https://github.com/fanshiqing/grouped_gemm/blob/main/csrc/grouped_gemm.cu#L355-L358)
+ # and as a result we avoid allocate the transpose of weights.
+ # Initialize weight.
+ if config.use_cpu_initialization:
+ self.weight1 = Parameter(
+ torch.empty(
+ self.config.hidden_size,
+ fc1_output_size_per_partition,
+ dtype=config.params_dtype,
+ )
+ )
+ self.weight2 = Parameter(
+ torch.empty(
+ fc2_input_size_per_partition, self.config.hidden_size, dtype=config.params_dtype
+ )
+ )
+ if config.perform_initialization:
+ _initialize_affine_weight_cpu(
+ self.weight1,
+ self.config.hidden_size,
+ fc1_output_size,
+ fc1_output_size_per_partition,
+ partition_dim=1,
+ init_method=config.init_method,
+ params_dtype=config.params_dtype,
+ )
+ _initialize_affine_weight_cpu(
+ self.weight2,
+ fc2_input_size,
+ self.config.hidden_size,
+ fc2_input_size_per_partition,
+ partition_dim=0,
+ init_method=config.output_layer_init_method,
+ params_dtype=config.params_dtype,
+ )
+ else:
+ self.weight1 = Parameter(
+ torch.empty(
+ self.config.hidden_size,
+ fc1_output_size_per_partition,
+ device=torch.cuda.current_device(),
+ dtype=config.params_dtype,
+ )
+ )
+ self.weight2 = Parameter(
+ torch.empty(
+ fc2_input_size_per_partition,
+ self.config.hidden_size,
+ device=torch.cuda.current_device(),
+ dtype=config.params_dtype,
+ )
+ )
+ if config.perform_initialization:
+ _initialize_affine_weight_gpu(
+ self.weight1,
+ config.init_method,
+ partition_dim=1,
+ expert_parallel=self.expert_parallel,
+ )
+ _initialize_affine_weight_gpu(
+ self.weight2,
+ config.output_layer_init_method,
+ partition_dim=0,
+ expert_parallel=self.expert_parallel,
+ )
+ setattr(self.weight1, 'allreduce', not self.expert_parallel)
+ setattr(self.weight2, 'allreduce', not self.expert_parallel)
+
+ def remove_extra_states_check(self, incompatible_keys):
+ """
+ Remove _extra_state from unexpected keys.
+ These keys are for dist ckpt compatibility with SequentialMLP.
+ """
+ keys = deepcopy(incompatible_keys.unexpected_keys)
+ for key in keys:
+ if '_extra_state' in key:
+ incompatible_keys.unexpected_keys.remove(key)
+
+ self.register_load_state_dict_post_hook(remove_extra_states_check)
+
+ def forward(self, permuted_local_hidden_states: torch.Tensor, tokens_per_expert: torch.Tensor):
+ """Forward step of the GroupedMLP."""
+ if permuted_local_hidden_states.nelement() != 0:
+ # Reshape the weights for the grouped GEMMs.
+ w1 = self.weight1.view(self.num_local_experts, self.config.hidden_size, -1)
+ w2 = self.weight2.view(self.num_local_experts, -1, self.config.hidden_size)
+
+ fc1_output = gg.ops.gmm(
+ permuted_local_hidden_states, w1, tokens_per_expert, trans_b=False
+ )
+
+ intermediate_parallel = self.activation_func(fc1_output)
+
+ fc2_output = gg.ops.gmm(intermediate_parallel, w2, tokens_per_expert, trans_b=False)
+ else:
+ # No token is allocated for local experts.
+ assert torch.count_nonzero(tokens_per_expert) == 0
+
+ # Make sure params of experts still have gradients even given zero tokens.
+ w1 = self.weight1.view(self.config.hidden_size, -1)
+ w2 = self.weight2.view(-1, self.config.hidden_size)
+ h = torch.matmul(permuted_local_hidden_states, w1)
+ h = self.activation_func(h)
+ h = torch.matmul(h, w2)
+
+ fc2_output = h
+
+ return fc2_output, None
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """
+ Maps local expert to global experts.
+ The sharded_state_dict for the weight parts are compatible with the SequentialMLP,
+ whereas the optimizer states are not due to the limitation from weight transposing.
+ That is, for finetuning scenario, the checkpoint is compatible with the SequentialMLP.
+ """
+ if self.moe_extended_tp:
+ raise NotImplementedError(
+ 'Currently distributed checkpointing is not supported for moe_extended_tp'
+ )
+
+ sharded_state_dict = {}
+ num_global_experts = (
+ parallel_state.get_expert_model_parallel_world_size() * self.num_local_experts
+ )
+ local_expert_indices_offset = (
+ parallel_state.get_expert_model_parallel_rank() * self.num_local_experts
+ )
+ tp_size = parallel_state.get_tensor_model_parallel_world_size()
+ tp_rank = parallel_state.get_tensor_model_parallel_rank()
+
+ prepend_axis_num = len(sharded_offsets)
+ replica_id = (
+ 0,
+ 0,
+ parallel_state.get_data_modulo_expert_parallel_rank(with_context_parallel=True),
+ )
+
+ local_ffn_dim_size = (
+ self.weight2.numel() // self.num_local_experts // self.config.hidden_size
+ )
+
+ @torch.no_grad()
+ def sh_ten_build_fn(
+ key: str,
+ t: torch.Tensor,
+ replica_id: ReplicaId,
+ flattened_range: Optional[slice],
+ tp_axis: int,
+ with_glu: bool,
+ ):
+ # TODO: write a generic implementation to cover both cases with and without GLU
+ if tp_axis == 1:
+ # weight1
+ if with_glu:
+ last_dim_size = local_ffn_dim_size * 2
+ else:
+ last_dim_size = local_ffn_dim_size
+ real_shape = (self.num_local_experts, self.config.hidden_size, last_dim_size)
+ elif tp_axis == 0:
+ # weight2
+ real_shape = (self.num_local_experts, local_ffn_dim_size, self.config.hidden_size)
+ assert with_glu == False
+ else:
+ raise ValueError("tp_axis should be 0 or 1.")
+ if flattened_range is None:
+ # weights
+ t = t.view(real_shape).transpose(-1, -2)
+ # change tp_axis due to the transposing
+ tp_axis = 1 - tp_axis
+ if with_glu:
+ local_tensors = torch.chunk(t, 2, -2)
+ sub_states = [
+ ShardedTensor.from_rank_offsets(
+ key,
+ local_tensors[0].contiguous(),
+ *sharded_offsets,
+ (
+ prepend_axis_num,
+ parallel_state.get_expert_model_parallel_rank(),
+ parallel_state.get_expert_model_parallel_world_size(),
+ ),
+ (prepend_axis_num + 1, tp_rank, tp_size * 2),
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ ),
+ ShardedTensor.from_rank_offsets(
+ key,
+ local_tensors[1].contiguous(),
+ *sharded_offsets,
+ (
+ prepend_axis_num,
+ parallel_state.get_expert_model_parallel_rank(),
+ parallel_state.get_expert_model_parallel_world_size(),
+ ),
+ (prepend_axis_num + 1, tp_size + tp_rank, tp_size * 2),
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ ),
+ ]
+ else:
+ sub_states = ShardedTensor.from_rank_offsets(
+ key,
+ t.contiguous(),
+ *sharded_offsets,
+ (
+ prepend_axis_num,
+ parallel_state.get_expert_model_parallel_rank(),
+ parallel_state.get_expert_model_parallel_world_size(),
+ ),
+ (prepend_axis_num + 1 + tp_axis, tp_rank, tp_size),
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ )
+ else:
+ # flattened optmizer states
+ # the non-flattened weight shape is [local_expert_num, hidden_size, ffn_size]
+ #
+ # For the case without GLU, it is straightforward, we just need to split each
+ # expert along the dim-0.
+ #
+ # For the case with GLU, we need to split the experts along dim-0 and split the
+ # two tensors for GLU along dim-2.
+ # To split along the non-first dim, we need to chunk the tensor into small pieces,
+ # since they belong to different tenors and are interleaved in the flattened space.
+ # Refer to the below sketch graph.
+ # |................| |........|........|
+ # |............FFFF| |........|....BBBB|
+ # |FFFFFFFFFFFFFFFF| -> |AAAAAAAA|BBBBBBBB|
+ # |FFFFFFFFFFFFFFFF| |AAAAAAAA|BBBBBBBB|
+ # |FF..............| |AA......|........|
+ # |................| |........|........|
+ #
+ # But too many chunks have severe performance issues. We merge these chunks during
+ # the save process along with some length information and recover them during the
+ # load process.
+ assert t.ndim == 1, (key, t.shape)
+ if with_glu:
+ non_flat_local_shape = (1, self.config.hidden_size, local_ffn_dim_size)
+ chunk_numel = local_ffn_dim_size
+ sub_states = []
+ start_pos = 0
+ for local_expert_idx in range(self.num_local_experts):
+ first_glu_idx = -1
+ w_start_range = -1
+ v_start_range = -1
+ w_tensors = []
+ v_tensors = []
+ w_lens = []
+ v_lens = []
+ for input_dim_idx in range(self.config.hidden_size):
+ for glu_idx in range(2):
+ local_idx = (
+ local_expert_idx * self.config.hidden_size * 2
+ + input_dim_idx * 2
+ + glu_idx
+ )
+ if (
+ flattened_range.start < chunk_numel * (local_idx + 1)
+ and flattened_range.stop > chunk_numel * local_idx
+ ):
+ if first_glu_idx == -1:
+ first_glu_idx = glu_idx
+ end_pos = min(
+ flattened_range.stop,
+ chunk_numel * (local_idx + 1) - flattened_range.start,
+ )
+ local_tensor = t[start_pos:end_pos]
+ local_flattened_range = slice(
+ max(0, flattened_range.start - chunk_numel * local_idx),
+ min(
+ chunk_numel,
+ flattened_range.stop - chunk_numel * local_idx,
+ ),
+ )
+ assert (
+ len(local_tensor)
+ == local_flattened_range.stop - local_flattened_range.start
+ )
+ start_pos += len(local_tensor)
+ expert_global_idx = (
+ local_expert_indices_offset + local_expert_idx
+ )
+ if glu_idx == 0:
+ w_tensors.append(local_tensor)
+ w_lens.append(len(local_tensor))
+ if w_start_range == -1:
+ w_start_range = max(
+ 0, flattened_range.start - chunk_numel * local_idx
+ )
+ else:
+ v_tensors.append(local_tensor)
+ v_lens.append(len(local_tensor))
+ if v_start_range == -1:
+ v_start_range = max(
+ 0, flattened_range.start - chunk_numel * local_idx
+ )
+ sub_states.append(
+ {
+ 'w_tensors': ShardedTensor.from_rank_offsets_flat(
+ key,
+ (
+ torch.cat(w_tensors, -1)
+ if len(w_tensors) > 0
+ else torch.Tensor()
+ ),
+ non_flat_local_shape,
+ *sharded_offsets,
+ (prepend_axis_num, expert_global_idx, num_global_experts),
+ (prepend_axis_num + 1 + tp_axis, tp_rank, tp_size * 2),
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ flattened_range=slice(
+ w_start_range, w_start_range + sum(w_lens)
+ ),
+ ),
+ 'w_lens': LocalNonpersistentObject(w_lens),
+ 'v_tensors': ShardedTensor.from_rank_offsets_flat(
+ key,
+ (
+ torch.cat(v_tensors, -1)
+ if len(v_tensors) > 0
+ else torch.Tensor()
+ ),
+ non_flat_local_shape,
+ *sharded_offsets,
+ (prepend_axis_num, expert_global_idx, num_global_experts),
+ (
+ prepend_axis_num + 1 + tp_axis,
+ tp_rank + tp_size,
+ tp_size * 2,
+ ),
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ flattened_range=slice(
+ v_start_range, v_start_range + sum(v_lens)
+ ),
+ ),
+ 'v_lens': LocalNonpersistentObject(v_lens),
+ 'first_glu_idx': LocalNonpersistentObject(first_glu_idx),
+ }
+ )
+ else:
+ non_flat_local_shape = (
+ real_shape[0] // self.num_local_experts,
+ *real_shape[1:],
+ )
+ chunk_numel = local_ffn_dim_size * self.config.hidden_size
+ sub_states = []
+ start_pos = 0
+ for local_expert_idx in range(self.num_local_experts):
+ if (
+ flattened_range.start < chunk_numel * (local_expert_idx + 1)
+ and flattened_range.stop > chunk_numel * local_expert_idx
+ ):
+ end_pos = min(
+ flattened_range.stop,
+ chunk_numel * (local_expert_idx + 1) - flattened_range.start,
+ )
+ local_tensor = t[start_pos:end_pos]
+ local_flattened_range = slice(
+ max(0, flattened_range.start - chunk_numel * local_expert_idx),
+ min(
+ chunk_numel,
+ flattened_range.stop - chunk_numel * local_expert_idx,
+ ),
+ )
+ assert (
+ len(local_tensor)
+ == local_flattened_range.stop - local_flattened_range.start
+ )
+ start_pos += len(local_tensor)
+ expert_global_idx = local_expert_indices_offset + local_expert_idx
+ sub_states.append(
+ ShardedTensor.from_rank_offsets_flat(
+ key,
+ local_tensor,
+ non_flat_local_shape,
+ *sharded_offsets,
+ (prepend_axis_num, expert_global_idx, num_global_experts),
+ (prepend_axis_num + 1 + tp_axis, tp_rank, tp_size),
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ flattened_range=local_flattened_range,
+ )
+ )
+ return sub_states
+
+ @torch.no_grad()
+ def sh_ten_merge_fn(sub_state_dict, tp_axis: int, with_glu: bool):
+ if tp_axis == 1:
+ # weight1
+ weight_shape = (self.config.hidden_size, -1)
+ elif tp_axis == 0:
+ # weight2
+ weight_shape = (-1, self.config.hidden_size)
+ assert with_glu == False
+ else:
+ raise ValueError("tp_axis should be 0 or 1.")
+ if isinstance(sub_state_dict, list) and isinstance(sub_state_dict[0], dict):
+ # flattened tensor with glu
+ res = []
+ for local_expert_dict in sub_state_dict:
+ w_tensors = torch.split(
+ local_expert_dict['w_tensors'], local_expert_dict['w_lens']
+ )
+ v_tensors = torch.split(
+ local_expert_dict['v_tensors'], local_expert_dict['v_lens']
+ )
+ first_glu_idx = local_expert_dict['first_glu_idx']
+ if first_glu_idx == 0:
+ res += [
+ x for x in itertools.chain(*itertools.zip_longest(w_tensors, v_tensors))
+ ]
+ else:
+ res += [
+ x for x in itertools.chain(*itertools.zip_longest(v_tensors, w_tensors))
+ ]
+ return torch.cat(res)
+ elif isinstance(sub_state_dict, list) and sub_state_dict[0].ndim == 1:
+ # flattened tensor without glu
+ return torch.cat(sub_state_dict)
+ else:
+ if with_glu:
+ sub_state_dict = torch.cat(sub_state_dict, -2)
+ return sub_state_dict.transpose(-1, -2).reshape(weight_shape)
+
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ for name, tensor in state_dict.items():
+ if name == 'weight1':
+ tp_axis = 1
+ with_glu = self.config.gated_linear_unit
+ wkey = f'{prefix}experts.linear_fc1.weight'
+ else:
+ tp_axis = 0
+ with_glu = False
+ wkey = f'{prefix}experts.linear_fc2.weight'
+ sharded_state_dict[f'{prefix}{name}'] = ShardedTensorFactory(
+ wkey,
+ tensor,
+ partial(sh_ten_build_fn, tp_axis=tp_axis, with_glu=with_glu),
+ partial(sh_ten_merge_fn, tp_axis=tp_axis, with_glu=with_glu),
+ replica_id,
+ )
+
+ replica_id = (
+ 0,
+ parallel_state.get_tensor_model_parallel_rank(),
+ parallel_state.get_data_modulo_expert_parallel_rank(with_context_parallel=True),
+ )
+ # Add fake _extra_state to be compatible with SequentialMLP
+ for expert_local_idx in range(self.num_local_experts):
+ expert_global_idx = local_expert_indices_offset + expert_local_idx
+ expert_sharded_offsets = (
+ *sharded_offsets,
+ (len(sharded_offsets), expert_global_idx, num_global_experts),
+ )
+ for mod in ['linear_fc1', 'linear_fc2']:
+ sharded_state_dict[f'{prefix}expert{expert_global_idx}.{mod}._extra_state'] = (
+ make_sharded_object_for_checkpoint(
+ None,
+ f'{prefix}experts.{mod}._extra_state',
+ expert_sharded_offsets,
+ replica_id,
+ )
+ )
+
+ return sharded_state_dict
+
+
+class TEGroupedMLP(MegatronModule):
+ """An efficient implementation of the Experts layer using TE's GroupedLinear.
+
+ Executes multiple experts in parallel to maximize computational efficiency.
+ """
+
+ def __init__(self, num_local_experts, config: TransformerConfig, submodules: MLPSubmodules):
+ super().__init__(config=config)
+ self.moe_extended_tp = config.moe_extended_tp
+ self.num_local_experts = num_local_experts
+ self.input_size = self.config.hidden_size
+
+ # Double the output width with gated linear unit, see https://arxiv.org/pdf/2002.05202.pdf
+ ffn_hidden_size = self.config.ffn_hidden_size
+ if self.config.gated_linear_unit:
+ ffn_hidden_size *= 2
+
+ self.linear_fc1 = build_module(
+ submodules.linear_fc1,
+ self.num_local_experts,
+ self.input_size,
+ ffn_hidden_size,
+ config=self.config,
+ init_method=self.config.init_method,
+ bias=self.config.add_bias_linear,
+ skip_bias_add=True,
+ is_expert=True,
+ tp_comm_buffer_name='fc1',
+ )
+
+ self.activation_func = self.config.activation_func
+
+ self.linear_fc2 = build_module(
+ submodules.linear_fc2,
+ self.num_local_experts,
+ self.config.ffn_hidden_size,
+ self.config.hidden_size,
+ config=self.config,
+ init_method=self.config.output_layer_init_method,
+ bias=self.config.add_bias_linear,
+ skip_bias_add=True,
+ is_expert=True,
+ tp_comm_buffer_name='fc2',
+ )
+
+ def remove_extra_states_check(self, incompatible_keys):
+ """
+ Remove extra _extra_state from unexpected keys.
+ These keys are for dist ckpt compatibility with SequentialMLP.
+ """
+ keys = deepcopy(incompatible_keys.unexpected_keys)
+ for key in keys:
+ if '_extra_state' in key:
+ incompatible_keys.unexpected_keys.remove(key)
+
+ self.register_load_state_dict_post_hook(remove_extra_states_check)
+
+ def forward(
+ self, permuted_local_hidden_states: torch.Tensor, tokens_per_expert: torch.Tensor
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """Forward of TEGroupedMLP
+
+ Args:
+ permuted_local_hidden_states (torch.Tensor): The permuted input hidden states of the
+ local experts.
+ tokens_per_expert (torch.Tensor): The number of tokens per expert.
+
+ Return:
+ output (torch.Tensor): The output of the local experts.
+ """
+ tokens_per_expert = tokens_per_expert.tolist()
+ intermediate_parallel, bias_parallel = self.linear_fc1(
+ permuted_local_hidden_states, tokens_per_expert
+ )
+
+ if self.config.bias_activation_fusion:
+ if self.activation_func == F.gelu:
+ if self.config.gated_linear_unit:
+ intermediate_parallel = bias_geglu_impl(intermediate_parallel, bias_parallel)
+ else:
+ assert self.config.add_bias_linear is True
+ intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel)
+ elif self.activation_func == F.silu and self.config.gated_linear_unit:
+ intermediate_parallel = bias_swiglu_impl(
+ intermediate_parallel,
+ bias_parallel,
+ self.config.activation_func_fp8_input_store,
+ )
+ else:
+ raise ValueError("Only support fusion of gelu and swiglu")
+ else:
+ if bias_parallel is not None:
+ intermediate_parallel = intermediate_parallel + bias_parallel
+ if self.config.gated_linear_unit:
+
+ def glu(x):
+ x = torch.chunk(x, 2, dim=-1)
+ return self.config.activation_func(x[0]) * x[1]
+
+ intermediate_parallel = glu(intermediate_parallel)
+ else:
+ intermediate_parallel = self.activation_func(intermediate_parallel)
+
+ output, output_bias = self.linear_fc2(intermediate_parallel, tokens_per_expert)
+
+ return output, output_bias
+
+ def sharded_state_dict(
+ self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None
+ ) -> ShardedStateDict:
+ """
+ Maps local expert to global experts.
+ The sharded state dict is interchangable with SequentialMLP's.
+ """
+ if self.moe_extended_tp:
+ raise NotImplementedError(
+ 'Currently distributed checkpointing is not supported for moe_extended_tp'
+ )
+ sharded_state_dict = {}
+ for name, module in self._modules.items():
+ sub_sd = module.sharded_state_dict(f'{name}.', sharded_offsets, metadata)
+ if name == 'linear_fc1' and self.config.gated_linear_unit:
+ num_global_experts = (
+ parallel_state.get_expert_model_parallel_world_size() * self.num_local_experts
+ )
+ local_expert_indices_offset = (
+ parallel_state.get_expert_model_parallel_rank() * self.num_local_experts
+ )
+ ep_axis = len(sharded_offsets)
+ for i in range(self.num_local_experts):
+ new_sharded_offsets = (
+ *sharded_offsets,
+ (ep_axis, local_expert_indices_offset + i, num_global_experts),
+ )
+ for k in (f'{name}.weight{i}', f'{name}.bias{i}'):
+ if k in sub_sd:
+ sub_sd[k] = apply_swiglu_sharded_factory(sub_sd[k], new_sharded_offsets)
+ # Add prefix here to match sequential's keys
+ replace_prefix_for_sharding(sub_sd, f'{name}.', f'{prefix}experts.{name}.')
+ sharded_state_dict.update({f"{prefix}{k}": v for k, v in sub_sd.items()})
+ return sharded_state_dict
+
+
+class SequentialMLP(MegatronModule):
+ """An implementation of the Experts layer using a sequence of MLP layers.
+
+ This class executes each expert sequentially.
+ """
+
+ def __init__(self, num_local_experts, config: TransformerConfig, submodules: MLPSubmodules):
+ super().__init__(config=config)
+ self.add_bias = config.add_bias_linear
+ self.moe_extended_tp = config.moe_extended_tp
+ self.num_local_experts = num_local_experts
+ self.local_experts = torch.nn.ModuleList()
+ for _ in range(self.num_local_experts):
+ expert = MLP(self.config, submodules, is_expert=True)
+ self.local_experts.append(expert)
+
+ def _pad_tensor_for_fp8(self, hidden):
+ """Padding tensor shape to multiples of 16."""
+ actual_num_tokens = hidden.shape[0]
+ divisor = 16
+ padded_num_tokens = ceil(actual_num_tokens / divisor) * divisor - actual_num_tokens
+ if padded_num_tokens > 0:
+ pad_tensor = torch.zeros(
+ padded_num_tokens, hidden.shape[1], dtype=hidden.dtype, device=hidden.device
+ )
+ hidden = torch.cat((hidden, pad_tensor), dim=0)
+ return hidden
+
+ def forward(self, permuted_local_hidden_states: torch.Tensor, tokens_per_expert: torch.Tensor):
+ """Forward step of the SequentialMLP."""
+ if self.num_local_experts == 1:
+ if self.config.fp8:
+ hidden = self._pad_tensor_for_fp8(permuted_local_hidden_states)
+ output, output_bias = self.local_experts[0](hidden)
+ output = output[: permuted_local_hidden_states.shape[0]]
+ else:
+ output, output_bias = self.local_experts[0](permuted_local_hidden_states)
+
+ return output, output_bias
+ else:
+ tokens_per_expert = tokens_per_expert.tolist()
+ tokens_list = torch.split(permuted_local_hidden_states, tokens_per_expert)
+
+ output_local_list = []
+ output_bias_list = []
+
+ for expert, tokens in zip(self.local_experts, tokens_list):
+ if self.config.fp8:
+ hidden = self._pad_tensor_for_fp8(tokens)
+ output, output_bias = expert(hidden)
+ output = output[: tokens.shape[0]]
+ else:
+ output, output_bias = expert(tokens)
+ output_local_list.append(output)
+ if self.add_bias:
+ output_bias_list.append(output_bias.expand_as(output))
+
+ output_local = torch.cat(output_local_list, dim=0)
+ if self.add_bias:
+ output_bias_local = torch.cat(output_bias_list, dim=0)
+ else:
+ output_bias_local = None
+
+ return output_local, output_bias_local
+
+ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+ """Maps local expert to global experts."""
+ if self.moe_extended_tp:
+ raise NotImplementedError(
+ 'Currently distributed checkpointing is not supported for moe_extended_tp'
+ )
+
+ sharded_state_dict = {}
+ num_global_experts = (
+ parallel_state.get_expert_model_parallel_world_size() * self.num_local_experts
+ )
+ local_expert_indices_offset = (
+ parallel_state.get_expert_model_parallel_rank() * self.num_local_experts
+ )
+
+ expert_sharded_prefix = f'{prefix}experts.'
+ for expert_local_idx, expert in enumerate(self.local_experts):
+ expert_global_idx = local_expert_indices_offset + expert_local_idx
+ expert_state_dict_prefix = f'{prefix}local_experts.{expert_local_idx}.'
+ expert_sharded_offsets = (
+ *sharded_offsets,
+ (len(sharded_offsets), expert_global_idx, num_global_experts),
+ )
+
+ expert_state_dict = expert.sharded_state_dict(
+ expert_state_dict_prefix, expert_sharded_offsets, metadata
+ )
+ # Remove expert layers indexing from sharded keys
+ replace_prefix_for_sharding(
+ expert_state_dict, expert_state_dict_prefix, expert_sharded_prefix
+ )
+ # Adjust replica ids - replication along DP modulo EP
+ for k, sh_ten in expert_state_dict.items():
+ replica_id = sh_ten.replica_id
+ assert (
+ len(replica_id) == 3
+ ), f'Expected replica_id for {k} to be in (PP, TP, DP) format, got: {replica_id}'
+ sh_ten.replica_id = (
+ *replica_id[:2],
+ parallel_state.get_data_modulo_expert_parallel_rank(with_context_parallel=True),
+ )
+
+ sharded_state_dict.update(expert_state_dict)
+ return sharded_state_dict
diff --git a/megatron/core/transformer/moe/grouped_gemm_util.py b/megatron/core/transformer/moe/grouped_gemm_util.py
new file mode 100644
index 0000000000..5dd344816b
--- /dev/null
+++ b/megatron/core/transformer/moe/grouped_gemm_util.py
@@ -0,0 +1,22 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+try:
+ import grouped_gemm
+except ImportError:
+ grouped_gemm = None
+
+
+def grouped_gemm_is_available():
+ """Check if grouped_gemm is available."""
+ return grouped_gemm is not None
+
+
+def assert_grouped_gemm_is_available():
+ """Assert that grouped_gemm is available."""
+ assert grouped_gemm_is_available(), (
+ "Grouped GEMM is not available. Please run "
+ "`pip install git+https://github.com/fanshiqing/grouped_gemm@v1.1.4`."
+ )
+
+
+ops = grouped_gemm.ops if grouped_gemm_is_available() else None
diff --git a/megatron/core/transformer/moe/legacy_a2a_token_dispatcher.py b/megatron/core/transformer/moe/legacy_a2a_token_dispatcher.py
new file mode 100644
index 0000000000..872c36aaa9
--- /dev/null
+++ b/megatron/core/transformer/moe/legacy_a2a_token_dispatcher.py
@@ -0,0 +1,304 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from typing import List, Optional, Tuple
+
+import torch
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.tensor_parallel.mappings import _gather_along_first_dim_expert_parallel
+from megatron.core.transformer.moe.moe_utils import permute, unpermute
+from megatron.core.transformer.moe.token_dispatcher import MoETokenDispatcher
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class MoEAlltoAllSEQTokenDispatcher(MoETokenDispatcher):
+ """
+ The legacy implementation of the AlltoAll-based token dispatcher, which handles token dispatching on the sequence level instead of token level. The core of this implementation lies each device dispatching on the entire sequence, with the hidden state being partitioned.
+ Note: This class is a replica of the MoEAlltoAllTokenDispatcher from version 0.8.
+ """
+
+ def __init__(
+ self, num_local_experts: int, local_expert_indices: List[int], config: TransformerConfig
+ ) -> None:
+ """
+ Initialize the AlltoAll token dispatcher.
+
+ Args:
+ num_local_experts (int): Number of local experts on the current device.
+ local_expert_indices (List[int]): Indices of local experts on the current device.
+ config (TransformerConfig): Configuration for the transformer model.
+ """
+ super().__init__(config=config)
+ self.hidden_shape = None
+ self.num_input_tokens = None
+ self.num_local_experts = num_local_experts
+ self.num_experts = config.num_moe_experts
+ assert self.num_local_experts > 0, "Expected at least one expert"
+ if self.num_local_experts > 1:
+ self.expert_ids_per_ep_rank = torch.tensor(
+ [i % self.num_local_experts for i in range(self.num_experts)],
+ dtype=torch.int32,
+ device=torch.cuda.current_device(),
+ )
+ self.local_expert_indices = local_expert_indices
+ assert (
+ len(self.local_expert_indices) == self.num_local_experts
+ ), "Invalid local expert indices"
+ for i in range(len(self.local_expert_indices) - 1):
+ assert (
+ self.local_expert_indices[i] == self.local_expert_indices[i + 1] - 1
+ ), "local_expert_indices must be continous"
+ self.router_topk = config.moe_router_topk
+ self.add_bias = config.add_bias_linear
+ self.ep_size = config.expert_model_parallel_size
+ self.probs = None
+ self.input_splits = None
+ self.output_splits = None
+ self.num_global_tokens_per_local_expert = None
+
+ # Token drop and padding.
+ # We need to keep track of the token num if we drop tokens without padding them.
+ self.num_out_tokens = None
+ # Drop and pad the input to capacity.
+ self.drop_and_pad = self.config.moe_pad_expert_input_to_capacity
+ if self.drop_and_pad:
+ assert self.config.moe_expert_capacity_factor is not None
+ self.capacity = None
+
+ # A cuda stream synchronization is needed in self.token_permutation() in some cases,
+ # because there are several non-blocking DtoH data transfers called in self.preprocess().
+ # The synchronization happens at different points based on MoE settings as late as possible.
+ # Valid sync points are "before_permutation_1", "before_ep_alltoall", "before_finish", and "no_sync".
+ self.cuda_sync_point = "no_sync"
+
+ def preprocess(self, indices: torch.Tensor) -> torch.Tensor:
+ """
+ Preprocess token indices for AlltoAll communication and token permutation. This method computes the number of tokens assigned to each expert based on the input indices.
+ It also initializes the necessary data structures for AlltoAll communication, such as input
+ and output splits, and the mapping between global tokens and local experts.
+
+ Args:
+ indices (torch.Tensor): Tensor of indices mapping tokens to experts.
+
+ Returns:
+ torch.Tensor: Tensor containing the number of tokens assigned to local expert.
+ """
+ num_local_tokens_per_expert = torch.histc(
+ indices, bins=self.num_experts, min=0, max=self.num_experts
+ )
+ # num_local_tokens_per_expert: [num_experts]
+
+ ep_size = self.config.expert_model_parallel_size
+ if self.drop_and_pad:
+ # probs: [num_experts, capacity]
+ self.capacity = self.probs.size(1)
+ num_tokens_per_local_expert = torch.full(
+ (self.num_local_experts,), self.capacity * self.ep_size, dtype=torch.long
+ )
+ return num_tokens_per_local_expert
+ elif self.config.moe_expert_capacity_factor is not None:
+ # Token drop but no pad. A synchronization is needed before the first
+ # permutation to get the `num_out_tokens` CPU value.
+ self.num_out_tokens = num_local_tokens_per_expert.sum().to(
+ torch.device("cpu"), non_blocking=True
+ )
+ self.cuda_sync_point = "before_permutation_1"
+ elif ep_size > 1:
+ # Token dropless and enable ep. A synchronization is needed before expert parallel
+ # AlltoAll communication to get the `input_splits` and `output_splits` CPU values.
+ self.cuda_sync_point = "before_ep_alltoall"
+ else:
+ # Token dropless and no ep. A synchronization is needed before the token_permutation()
+ # function returns to get the `tokens_per_expert` CPU value.
+ self.cuda_sync_point = "before_finish"
+
+ if ep_size > 1:
+ # ===================================================
+ # Calculate input_splits, output_splits for alltoall-v.
+ # ===================================================
+ self.input_splits = (
+ num_local_tokens_per_expert.reshape(ep_size, self.num_local_experts)
+ .sum(axis=1)
+ .to(torch.device("cpu"), non_blocking=True)
+ .numpy()
+ )
+ num_global_tokens_per_expert = _gather_along_first_dim_expert_parallel(
+ num_local_tokens_per_expert
+ ).reshape(ep_size, self.num_experts)
+ self.num_global_tokens_per_local_expert = num_global_tokens_per_expert[
+ :, self.local_expert_indices[0] : self.local_expert_indices[-1] + 1
+ ]
+ self.output_splits = (
+ self.num_global_tokens_per_local_expert.sum(axis=-1)
+ .to(torch.device("cpu"), non_blocking=True)
+ .numpy()
+ )
+ num_tokens_per_local_expert = self.num_global_tokens_per_local_expert.sum(axis=0).to(
+ torch.device("cpu"), non_blocking=True
+ )
+ # ===================================================
+ # num_global_tokens_per_expert: [ep_size, num_experts]
+ # num_global_tokens_per_local_expert: [ep_size, num_local_experts]
+ # num_tokens_per_local_expert: [num_local_experts]
+ # ===================================================
+ else:
+ self.num_global_tokens_per_local_expert = num_local_tokens_per_expert.reshape(
+ -1, self.num_experts
+ )
+ num_tokens_per_local_expert = num_local_tokens_per_expert.to(
+ torch.device("cpu"), non_blocking=True
+ )
+
+ if self.num_local_experts > 1:
+ # No further synchronization is needed because torch.repeat_interleave() calls stream
+ # synchronization internally when the `output_size` parameter is not provided.
+ self.cuda_sync_point = "no_sync"
+ self.global_input_tokens_local_experts_indices = torch.repeat_interleave(
+ self.expert_ids_per_ep_rank, self.num_global_tokens_per_local_expert.ravel()
+ )
+
+ return num_tokens_per_local_expert
+
+ def token_permutation(
+ self, hidden_states: torch.Tensor, probs: torch.Tensor, indices: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Dispatch tokens to local experts using AlltoAll communication.
+
+ Args:
+ hidden_states (torch.Tensor): Input token embeddings.
+ probs (torch.Tensor): Probs of tokens assigned to experts.
+ indices (torch.Tensor): Indices of tokens assigned to experts.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]:
+ - Permuted token embeddings for local experts.
+ - Number of tokens per expert.
+ """
+ # Preprocess: Get the metadata for communication, permutation and computation operations.
+ self.hidden_shape = hidden_states.shape
+ self.probs = probs
+ assert probs.dim() == 2, "Expected 2D tensor for probs"
+ assert indices.dim() == 2, "Expected 2D tensor for indices"
+ hidden_states = hidden_states.view(-1, self.hidden_shape[-1])
+ tokens_per_expert = self.preprocess(indices)
+
+ # Perform tensor parallel AlltoAll communication
+ # hidden_states: [S*B/TP, H] -> [S*B, H/TP]
+ if parallel_state.get_tensor_model_parallel_world_size() > 1:
+ hidden_states = tensor_parallel.all_to_all_sp2hp(hidden_states)
+
+ # Permutation 1: input to AlltoAll input
+ self.hidden_shape_before_permute = hidden_states.shape
+ if self.cuda_sync_point == "before_permutation_1":
+ torch.cuda.current_stream().synchronize()
+ permutated_local_input_tokens, self.reversed_local_input_permutation_mapping = permute(
+ hidden_states,
+ indices,
+ num_out_tokens=self.num_out_tokens,
+ padded_mode=self.drop_and_pad,
+ )
+
+ # Perform expert parallel AlltoAll communication
+ if self.cuda_sync_point == "before_ep_alltoall":
+ torch.cuda.current_stream().synchronize()
+ global_input_tokens = tensor_parallel.all_to_all(
+ parallel_state.get_expert_model_parallel_group(),
+ permutated_local_input_tokens,
+ self.output_splits,
+ self.input_splits,
+ )
+
+ # Permutation 2: Sort alltoall output by local experts when num_local_experts > 1.
+ if self.num_local_experts > 1:
+ if not self.drop_and_pad:
+ global_input_tokens, self.reversed_global_input_permutation_mapping = permute(
+ global_input_tokens, self.global_input_tokens_local_experts_indices
+ )
+ else:
+ global_input_tokens = global_input_tokens.reshape(
+ self.ep_size, self.num_local_experts, self.capacity, -1
+ )
+ global_input_tokens = (
+ global_input_tokens.transpose(0, 1)
+ .reshape(self.num_local_experts * self.ep_size * self.capacity, -1)
+ .contiguous()
+ )
+
+ # Perform tensor parallel AllGather on the hidden dimension to obtain the input tokens.
+ # global_input_tokens: [SEQL, H/TP] -> [SEQL, H]
+ if parallel_state.get_tensor_model_parallel_world_size() > 1:
+ global_input_tokens = tensor_parallel.all_gather_last_dim_from_tensor_parallel_region(
+ global_input_tokens
+ )
+ if self.cuda_sync_point == "before_finish":
+ torch.cuda.current_stream().synchronize()
+
+ return global_input_tokens, tokens_per_expert
+
+ def token_unpermutation(
+ self, hidden_states: torch.Tensor, bias: torch.Tensor = None
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """
+ Reverse the token permutation to restore the original order.
+
+ Args:
+ hidden_states (torch.Tensor): Output from local experts.
+ bias (torch.Tensor, optional): Bias tensor (not supported).
+
+ Returns:
+ Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ - Unpermuted token embeddings in the original order.
+ - None (bias is not supported).
+ """
+ assert bias is None, "Bias is not supported in MoEAlltoAllTokenDispatcher"
+
+ # Perform tensor parallel Reduce-Scatter
+ # hidden_states: [SEQL, H] -> [SEQL, H/TP]
+ if parallel_state.get_tensor_model_parallel_world_size() > 1:
+ hidden_states = tensor_parallel.reduce_scatter_last_dim_to_tensor_parallel_region(
+ hidden_states
+ )
+
+ # Unpermutation 2: expert output to AlltoAll input
+ if self.num_local_experts > 1:
+ if not self.drop_and_pad:
+ hidden_states = unpermute(
+ hidden_states, self.reversed_global_input_permutation_mapping
+ )
+ else:
+ hidden_states = hidden_states.reshape(
+ self.num_local_experts, self.ep_size, self.capacity, -1
+ )
+ hidden_states = (
+ hidden_states.transpose(0, 1)
+ .reshape(self.ep_size * self.num_local_experts * self.capacity, -1)
+ .contiguous()
+ )
+
+ # Perform expert parallel AlltoAll communication
+ # hidden_states: [SEQL, H] -> [SEQL, H/TP]
+ permutated_local_input_tokens = tensor_parallel.all_to_all(
+ parallel_state.get_expert_model_parallel_group(),
+ hidden_states,
+ self.input_splits,
+ self.output_splits,
+ )
+
+ # Unpermutation 1: AlltoAll output to output
+ output = unpermute(
+ permutated_local_input_tokens,
+ self.reversed_local_input_permutation_mapping,
+ probs=self.probs,
+ padded_mode=self.drop_and_pad,
+ restore_shape=self.hidden_shape_before_permute,
+ )
+
+ # Perform tensor parallel AlltoAll communication
+ # output: [S*B, H/TP] -> [S*B/TP, H]
+ if parallel_state.get_tensor_model_parallel_world_size() > 1:
+ output = tensor_parallel.all_to_all_hp2sp(output)
+
+ # Reshape the output tensor
+ output = output.view(self.hidden_shape)
+ return output, None
diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
new file mode 100644
index 0000000000..8b393abc77
--- /dev/null
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -0,0 +1,164 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.transformer.mlp import MLPSubmodules
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.moe.experts import GroupedMLP, SequentialMLP, TEGroupedMLP
+from megatron.core.transformer.moe.legacy_a2a_token_dispatcher import MoEAlltoAllSEQTokenDispatcher
+from megatron.core.transformer.moe.router import TopKRouter
+from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
+from megatron.core.transformer.moe.token_dispatcher import (
+ MoEAllGatherTokenDispatcher,
+ MoEAlltoAllTokenDispatcher,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+@dataclass
+class MoESubmodules:
+ """MoE Layer Submodule spec"""
+
+ experts: Union[ModuleSpec, type] = None
+ shared_experts: Union[ModuleSpec, type] = None
+
+
+class BaseMoELayer(MegatronModule, ABC):
+ """Base class for a mixture of experts layer.
+
+ Args:
+ config (TransformerConfig): Configuration object for the transformer model.
+ """
+
+ def __init__(self, config: TransformerConfig, layer_number: int = None):
+ super(BaseMoELayer, self).__init__(config)
+ self.config = config
+ self.expert_parallel_size = parallel_state.get_expert_model_parallel_world_size()
+ assert self.expert_parallel_size > 0, "Expected non-negative expert parallel size"
+
+ if self.config.moe_extended_tp:
+ self.num_local_experts = self.config.num_moe_experts
+ local_expert_indices_offset = 0
+ else:
+ assert self.config.num_moe_experts % self.expert_parallel_size == 0
+ self.num_local_experts = self.config.num_moe_experts // self.expert_parallel_size
+ local_expert_indices_offset = (
+ parallel_state.get_expert_model_parallel_rank() * self.num_local_experts
+ )
+
+ self.use_shared_expert = self.config.moe_shared_expert_intermediate_size is not None
+ self.shared_expert_overlap = self.config.moe_shared_expert_overlap
+
+ self.local_expert_indices = [
+ local_expert_indices_offset + i for i in range(self.num_local_experts)
+ ]
+ assert all(map(lambda x: x < self.config.num_moe_experts, self.local_expert_indices))
+ self.router = None
+ self.experts = None
+ self.shared_experts = None
+ self.token_dispatcher = None
+ self.layer_number = layer_number
+
+ @abstractmethod
+ def forward(self, hidden_states):
+ """Forward method for the MoE layer."""
+ pass
+
+ def set_layer_number(self, layer_number: int):
+ """Set the layer number for the MoE layer."""
+ self.layer_number = layer_number
+ self.router.set_layer_number(layer_number)
+
+
+class MoELayer(BaseMoELayer):
+ """Mixture of experts Layer **currently only supports no token dropping**.
+
+ Args:
+ BaseMoELayer (MegatronModule): Base class for MoE layers
+ """
+
+ def __init__(
+ self, config: TransformerConfig, submodules: MLPSubmodules = None, layer_number: int = None
+ ):
+ self.submodules = submodules
+ super(MoELayer, self).__init__(config=config, layer_number=layer_number)
+ self.moe_layer_recompute = config.moe_layer_recompute
+
+ # Initialize router
+ self.router = TopKRouter(config=self.config)
+
+ # Initialize experts
+ if self.config.moe_grouped_gemm:
+ if isinstance(self.submodules.experts, MLPSubmodules):
+ self.experts = TEGroupedMLP(
+ self.num_local_experts, self.config, self.submodules.experts
+ )
+ else:
+ self.experts = GroupedMLP(self.num_local_experts, self.config)
+ else:
+ assert isinstance(self.submodules.experts, MLPSubmodules)
+ self.experts = SequentialMLP(
+ self.num_local_experts, self.config, self.submodules.experts
+ )
+
+ # Initialize token dispatcher
+ if config.moe_token_dispatcher_type == "allgather":
+ self.token_dispatcher = MoEAllGatherTokenDispatcher(
+ self.num_local_experts, self.local_expert_indices, config=self.config
+ )
+ elif config.moe_token_dispatcher_type == "alltoall":
+ self.token_dispatcher = MoEAlltoAllTokenDispatcher(
+ self.num_local_experts, self.local_expert_indices, config=self.config
+ )
+ elif config.moe_token_dispatcher_type == "alltoall_seq":
+ self.token_dispatcher = MoEAlltoAllSEQTokenDispatcher(
+ self.num_local_experts, self.local_expert_indices, config=self.config
+ )
+ else:
+ raise ValueError(
+ f"Unsupported token dispatcher type: {config.moe_token_dispatcher_type}"
+ )
+
+ # Initialize shared experts
+ if self.use_shared_expert:
+ self.shared_experts = SharedExpertMLP(self.config, self.submodules.shared_experts)
+ if self.shared_expert_overlap:
+ self.token_dispatcher.set_shared_experts(self.shared_experts)
+
+ def forward(self, hidden_states: torch.Tensor):
+ if (
+ self.training
+ and self.config.tensor_model_parallel_size > 1
+ and not self.config.sequence_parallel
+ ):
+ raise ValueError(
+ "During training, performance may degrade if MoE and tensor parallelism"
+ "are enabled without also enabling sequence parallelism."
+ )
+
+ # process MoE
+ def custom_forward(hidden_states):
+ probs, indices = self.router(hidden_states)
+ (dispatched_input, tokens_per_expert) = self.token_dispatcher.token_permutation(
+ hidden_states, probs, indices
+ )
+ expert_output, mlp_bias = self.experts(dispatched_input, tokens_per_expert)
+ output, mlp_bias = self.token_dispatcher.token_unpermutation(expert_output, mlp_bias)
+ if self.use_shared_expert and not self.shared_expert_overlap:
+ # if shared_expert_overlap is True, the expert calculation happens in
+ # the token_dispatcher to overlap communications and computations
+ output += self.shared_experts(hidden_states)
+ return output, mlp_bias
+
+ if self.moe_layer_recompute:
+ output, mlp_bias = tensor_parallel.checkpoint(custom_forward, False, hidden_states)
+ else:
+ output, mlp_bias = custom_forward(hidden_states)
+
+ return output, mlp_bias
diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
new file mode 100644
index 0000000000..02a2cccca5
--- /dev/null
+++ b/megatron/core/transformer/moe/moe_utils.py
@@ -0,0 +1,558 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import math
+
+import torch
+
+from megatron.core import parallel_state
+
+
+def switch_load_balancing_loss_func(
+ probs: torch.Tensor,
+ tokens_per_expert: torch.Tensor,
+ topk: int,
+ moe_aux_loss_coeff: float,
+ sequence_partition_group=None,
+):
+ """Calculate the auxiliary loss for load balancing.
+ Refer to the Switch Transformer paper (https://arxiv.org/abs/2101.03961) for details.
+
+ Args:
+ probs (torch.Tensor): Softmax probabilities output by the router for each token.
+ Shape in [num_tokens, num_experts].
+ tokens_per_expert (torch.Tensor): Number of tokens assigned to each expert.
+ Shape in [num_experts]
+ topk (int): The number of experts selected for each token.
+ moe_aux_loss_coeff (float): The coefficient for the auxiliary loss.
+ sequence_partition_group (optional): The parallel group over which the sequence is
+ partitioned. If None, no partitioning is applied.
+ Defaults to None.
+
+ Returns:
+ torch.Tensor: The auxiliary loss for load balancing.
+ """
+ num_sub_sequence = 1
+
+ # If the sequence is partitioned by certain parallelism strategies like Sequence Parallelism
+ # or Context Parallelism, compute the gradient of the auxiliary loss with respect to the full
+ # sequence.
+ if sequence_partition_group is not None:
+ # We can keep `aggregated_probs_per_expert` local since we don't need the gradient for
+ # `tokens_per_expert`, saving one allreduce operation for `aggregated_probs_per_expert`.
+ num_sub_sequence = torch.distributed.get_world_size(sequence_partition_group)
+ torch.distributed.all_reduce(tokens_per_expert, group=sequence_partition_group)
+
+ num_tokens = probs.shape[0] * num_sub_sequence
+ num_experts = probs.shape[1]
+
+ # The formula of aux_loss: aux_loss = sum((probs_per_expert/num_tokens) *
+ # (tokens_per_expert/(num_tokens*topk))) * num_experts * moe_aux_loss_coeff.
+ # This can be simplified to fuse the division and multiplication operations.
+ aggregated_probs_per_expert = probs.sum(dim=0)
+ aux_loss = torch.sum(aggregated_probs_per_expert * tokens_per_expert) * (
+ num_experts * moe_aux_loss_coeff / (num_tokens * num_tokens * topk)
+ )
+ return aux_loss
+
+
+def z_loss_func(logits, z_loss_coeff):
+ """Encourages the router's logits to remain small to enhance stability.
+ Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
+
+ Args:
+ logits (torch.Tensor): The logits of the router.
+
+ Returns:
+ torch.Tensor: The logits after applying the z-loss.
+ """
+
+ z_loss = torch.mean(torch.square(torch.logsumexp(logits, dim=-1))) * z_loss_coeff
+ return z_loss
+
+
+def sinkhorn(cost: torch.Tensor, tol: float = 0.0001):
+ """Sinkhorn based MoE routing function"""
+ cost = torch.exp(cost)
+ d0 = torch.ones(cost.size(0), device=cost.device, dtype=cost.dtype)
+ d1 = torch.ones(cost.size(1), device=cost.device, dtype=cost.dtype)
+
+ eps = 0.00000001
+ error = 1e9
+ d1_old = d1
+ while error > tol:
+ d0 = (1 / d0.size(0)) * 1 / (torch.sum(d1 * cost, 1) + eps)
+ d1 = (1 / d1.size(0)) * 1 / (torch.sum(d0.unsqueeze(1) * cost, 0) + eps)
+ error = torch.mean(torch.abs(d1_old - d1))
+ d1_old = d1
+ return d1 * cost * d0.unsqueeze(1)
+
+
+def get_capacity(num_tokens: int, num_experts: int, capacity_factor: float, min_capacity=None):
+ """
+ Calculate the capacity of each expert.
+
+ Args:
+ num_tokens (int): num of the input tokens.
+ num_experts (int): num of the experts.
+ capacity_factor (float): Capacity factor.
+ min_capacity (int, optional): Minimum capacity. Defaults to None.
+
+ Returns:
+ Tensor: Capacity of each expert.
+ """
+ capacity = math.ceil((num_tokens / num_experts) * capacity_factor)
+ if min_capacity is not None and capacity < min_capacity:
+ capacity = min_capacity
+ return capacity
+
+
+class MoEAuxLossAutoScaler(torch.autograd.Function):
+ """An AutoScaler that compute and scales the grad for auxiliary loss."""
+
+ main_loss_backward_scale: torch.Tensor = torch.tensor(1.0)
+
+ @staticmethod
+ def forward(ctx, output: torch.Tensor, aux_loss: torch.Tensor):
+ """Preserve the aux_loss by storing it in the context to avoid garbage collection.
+
+ Args:
+ output (torch.Tensor): The output tensor.
+ aux_loss (torch.Tensor): The auxiliary loss tensor.
+
+ Returns:
+ torch.Tensor: The output tensor.
+ """
+ ctx.save_for_backward(aux_loss)
+ return output
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor):
+ """Compute and scale the gradient for auxiliary loss..
+
+ Args:
+ grad_output (torch.Tensor): The gradient of the output.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]: The gradient of the output, scaled auxiliary loss
+ gradient.
+ """
+ (aux_loss,) = ctx.saved_tensors
+ aux_loss_backward_scale = MoEAuxLossAutoScaler.main_loss_backward_scale
+ scaled_aux_loss_grad = torch.ones_like(aux_loss) * aux_loss_backward_scale
+ return grad_output, scaled_aux_loss_grad
+
+ @staticmethod
+ def set_loss_scale(scale: torch.Tensor):
+ """set the scale of the aux loss.
+
+ Args:
+ scale (torch.Tensor): The scale value to set. Please ensure that the scale passed in
+ matches the scale of the main_loss.
+ """
+ MoEAuxLossAutoScaler.main_loss_backward_scale = scale
+
+
+def permute(tokens, indices, num_out_tokens: int = None, padded_mode: bool = False):
+ """Permute the tokens based on the indices. Token with the same index will be grouped together.
+ The input indices shape is [tokens, top_k], it indicates which experts were selected by each
+ token separately.
+ Args:
+ tokens (torch.Tensor): The input token tensor.
+ indices (torch.Tensor): The token to expert indices tensor, should have a shape of
+ [num_tokens] or [num_tokens, topk].
+ num_out_tokens (int, optional): The effective output token count, when enabling the
+ capacity factor, should equal the number of tokens not
+ dropped. By default, set to None, meaning no tokens are
+ dropped.
+ padded_mode (bool, optional): If True, indicating the indices are padded to
+ [num_expert, capacity] to denote selected tokens per expert.
+ Defaults to False.
+
+ Returns:
+ torch.Tensor: The permuted tensor.
+ torch.Tensor: The sorted_indices corresponding permuted tensor.
+ """
+ if padded_mode:
+ return permute_with_padded_tokens(tokens, indices)
+
+ if indices.dim() == 1:
+ indices = indices.unsqueeze(1)
+
+ topk = indices.size(1)
+ flatten_indices = indices.view(-1)
+ sorted_indices = torch.argsort(flatten_indices, stable=True)
+ if num_out_tokens is not None:
+ sorted_indices = sorted_indices[:num_out_tokens]
+ moe_gather_indices = (sorted_indices // topk).unsqueeze(1).expand(-1, tokens.size(-1))
+ permuted_tokens = moe_gather.apply(tokens, moe_gather_indices)
+
+ return permuted_tokens, sorted_indices
+
+
+def unpermute(
+ permuted_tokens: torch.Tensor,
+ sorted_indices: torch.Tensor,
+ probs: torch.Tensor = None,
+ padded_mode: bool = False,
+ restore_shape: torch.Size = None,
+):
+ """Unpermute a tensor of permuted tokens based on sorted indices, and optionally merge the
+ tokens with their corresponding probabilities.
+
+ Args:
+ permuted_tokens (torch.Tensor): 2D tensor [num_tokens*topk, hidden]. The tensor of permuted
+ tokens to be unpermuted.
+ sorted_indices (torch.Tensor): 1D tensor [num_tokens*topk]. The tensor of sorted indices
+ used to unpermute the tokens.
+ probs (torch.Tensor, optional): 2D tensor [num_tokens, topk]. The tensor of probabilities
+ corresponding to the permuted tokens. If provided,
+ the unpermuted tokens will be merged with their respective
+ probabilities.
+ padded_mode (bool, optional): If True, indicating the indices are padded to
+ [num_expert, capacity] to denote selected tokens per expert.
+ Defaults to False.
+ restore_shape (torch.Size, optional): The input shape before permutation, only used in
+ padding mode. Defaults to None.
+
+ Returns:
+ torch.Tensor: The unpermuted tokens, optionally merged with probabilities.
+ """
+ if padded_mode:
+ return unpermute_with_padded_tokens(
+ permuted_tokens, sorted_indices, probs, restore_shape=restore_shape
+ )
+
+ assert sorted_indices.numel() == permuted_tokens.size(
+ 0
+ ), f"Got {sorted_indices.numel()} != {permuted_tokens.size(0)}."
+ if probs is not None:
+ # Unpermute and merge the tokens with their probabilities
+ num_unpermuted_tokens = probs.numel()
+ assert probs.dim() == 2, f"Expected 2D tensor for probs, got {probs.dim()} dims."
+ topk = probs.size(1)
+ else:
+ # Unpermute the tokens without merge
+ num_unpermuted_tokens = permuted_tokens.size(0)
+ topk = 1
+
+ output_size = [num_unpermuted_tokens, permuted_tokens.shape[-1]]
+ moe_scatter_indices = sorted_indices.unsqueeze(1).expand(-1, permuted_tokens.size(-1))
+ unpermuted_tokens = moe_scatter.apply(permuted_tokens, moe_scatter_indices, output_size)
+ unpermuted_tokens = unpermuted_tokens.reshape(-1, topk, permuted_tokens.size(-1))
+ if probs is not None:
+ unpermuted_tokens = unpermuted_tokens * probs.unsqueeze(-1)
+ unpermuted_tokens = unpermuted_tokens.sum(dim=1)
+
+ return unpermuted_tokens
+
+
+def permute_with_padded_tokens(tokens, indices):
+ """Permute the tokens based on the indices, only used in padding mode.
+ The input indices shape is [num_expert, capacity], it indicates which tokens were selected
+ by each expert separately.
+ Args:
+ tokens (torch.Tensor): The input token tensor.
+ indices (torch.Tensor): A tensor with shape [num_expert, capacity], indicating the selected
+ tokens for each expert.
+
+ Returns:
+ torch.Tensor: The permuted tensor.
+ torch.Tensor: The sorted_indices corresponding permuted tensor.
+ """
+ permuted_tokens = tokens.index_select(dim=0, index=indices.view(-1))
+
+ return permuted_tokens, indices
+
+
+def unpermute_with_padded_tokens(
+ permuted_tokens: torch.Tensor,
+ indices: torch.Tensor,
+ probs: torch.Tensor,
+ restore_shape: torch.Size,
+) -> torch.Tensor:
+ """
+ Unpermutes a padded permuted tokens based on sorted indices and merges the tokens with their
+ corresponding probabilities.
+
+ This function takes a tensor of permuted tokens and reorders them according to the provided
+ indices. It also combines the tokens with their associated probabilities.
+
+ Parameters:
+ permuted_tokens (torch.Tensor): A 2D tensor containing permuted tokens.
+ indices (torch.Tensor): A tensor with shape [num_expert, capacity], indicating the selected
+ tokens for each expert.
+ probs (torch.Tensor): A tensor with the same shape as indices, containing probabilities
+ corresponding to each token.
+ restore_shape (torch.Size): The target shape for the unpermuted tokens tensor.
+
+ Returns:
+ torch.Tensor: A tensor of unpermuted tokens, merged with their probabilities.
+
+ """
+ # Ensure permuted_tokens is 2D
+ assert permuted_tokens.dim() == 2, f"Got {permuted_tokens.dim()}D."
+
+ # Reshape and expand probabilities and indices to match permuted_tokens
+ probs = probs.view(-1).unsqueeze(-1)
+ indices = indices.view(-1, 1).expand(-1, permuted_tokens.shape[1])
+ assert (
+ permuted_tokens.shape == indices.shape
+ ), "Shape mismatch between permuted_tokens and indices."
+
+ # Combine tokens with their probabilities
+ combined_output = probs * permuted_tokens
+
+ # Prepare a tensor of zeros with the desired output shape
+ empty_tokens = torch.zeros(
+ restore_shape, dtype=combined_output.dtype, device=combined_output.device
+ )
+
+ # Scatter the combined tokens back to their original positions
+ unpermuted_tokens = torch.scatter_add(empty_tokens, 0, indices, combined_output)
+
+ return unpermuted_tokens
+
+
+def sort_chunks_by_idxs(input: torch.Tensor, split_sizes: torch.Tensor, sorted_idxs: torch.Tensor):
+ """Split and sort the input tensor based on the split_sizes and sorted indices."""
+ input = torch.split(input, split_sizes.tolist(), dim=0)
+ output = torch.cat([input[i] for i in sorted_idxs], dim=0)
+ return output
+
+
+def topk_softmax_with_capacity(
+ logits: torch.Tensor,
+ topk: int,
+ capacity_factor: float = None,
+ pad_to_capacity: bool = False,
+ drop_policy: str = "probs",
+ use_pre_softmax: bool = False,
+ deterministic_mode: bool = False,
+):
+ """Apply capacity and padding to the top-k selection.
+ Args:
+ logits (torch.Tensor): Logits tensor.
+ topk (int): The number of experts to select for each token.
+ capacity_factor (int): The capacity factor of each expert. Will drop tokens if the number
+ of tokens exceeds the capacity.
+ pad_to_capacity (bool): Whether to need padding in token drop mode.
+ drop_policy (str): The policy to drop tokens. Can be either "prob" or "position".
+ If "prob", the tokens with the lowest probabilities will be dropped.
+ If "position", tokens at the end of each batch will be dropped.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Probs, indices and tokens_per_expert
+ tensor.
+
+ (1) If there's no token padding, the shape of probs and indices is [tokens, top_k],
+ indicating the selected experts for each token.
+ (2) If there's token padding, the shape of probs and indices is [num_expert, capacity],
+ indicating the tokens selected for each expert.
+ """
+ assert logits.dim() == 2, f"Expected 2D logits [num_tokens, num_experts], got {logits.dim()}."
+ num_tokens = logits.shape[0]
+ num_experts = logits.shape[1]
+ if use_pre_softmax:
+ # Pre softmax
+ scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)
+ probs, top_indices = torch.topk(scores, k=topk, dim=1)
+ else:
+ # Post softmax
+ if topk == 1:
+ # Requires applying softmax before selecting the top-k when k is 1,
+ # since softmax on a [num_tokens, 1] would yield a zero gradient.
+ raise ValueError("Please use --moe-router-pre-softmax when topk is 1.")
+ scores, top_indices = torch.topk(logits, k=topk, dim=1)
+ probs = torch.softmax(scores, dim=-1, dtype=torch.float32).type_as(logits)
+
+ if capacity_factor is None:
+ # TopK without capacity
+ if deterministic_mode:
+ tokens_per_expert = torch.bincount(top_indices.view(-1), minlength=num_experts)
+ else:
+ tokens_per_expert = torch.histc(top_indices, bins=num_experts, min=0, max=num_experts)
+ return probs, top_indices, tokens_per_expert
+ else:
+ # TopK with capacity
+ expert_capacity = get_capacity(
+ num_tokens=num_tokens * topk, num_experts=num_experts, capacity_factor=capacity_factor
+ )
+ # TopK selection, Maskout unused experts
+ topk_masked_gates = torch.zeros_like(logits).scatter(1, top_indices, probs)
+ topk_mask = torch.zeros_like(logits).scatter(1, top_indices, 1)
+
+ # Maskout exceeded tokens
+ if drop_policy == "probs":
+ capacity_probs, capacity_indices = torch.topk(
+ topk_masked_gates, k=expert_capacity, dim=0, sorted=False
+ )
+ capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1)
+ elif drop_policy == "position":
+ _, capacity_indices = torch.topk(topk_mask, k=expert_capacity, dim=0, sorted=False)
+ capacity_mask = torch.zeros_like(logits).scatter(0, capacity_indices, 1)
+ capacity_probs = torch.gather(topk_masked_gates, 0, capacity_indices)
+ else:
+ raise ValueError(f"Invalid drop_policy: {drop_policy}")
+
+ if pad_to_capacity:
+ final_probs, final_indices = (
+ capacity_probs.T.contiguous(),
+ capacity_indices.T.contiguous(),
+ )
+ tokens_per_expert_before_capacity = topk_mask.sum(dim=0)
+ else:
+ # Get exceed mask and maskout exceeded probs and indices
+ final_mask = torch.logical_and(topk_mask, capacity_mask)
+ drop_mask = torch.logical_not(final_mask)
+ exceed_mask = torch.gather(drop_mask, 1, top_indices)
+ final_probs = probs * torch.logical_not(exceed_mask)
+ final_indices = top_indices.clone().masked_fill_(
+ exceed_mask, torch.iinfo(torch.long).max
+ )
+ tokens_per_expert_before_capacity = topk_mask.sum(dim=0)
+ return final_probs, final_indices, tokens_per_expert_before_capacity
+
+
+def save_to_aux_losses_tracker(
+ name: str,
+ loss: torch.Tensor,
+ layer_number: int,
+ num_layers: int,
+ reduce_group: torch.distributed.ProcessGroup = None,
+ avg_group: torch.distributed.ProcessGroup = None,
+):
+ """Save the auxiliary loss for logging.
+ Args:
+ name (str): The name of the loss.
+ loss (torch.Tensor): The loss tensor.
+ layer_number (int): Layer index of the loss.
+ num_layers (int): The number of total layers.
+ reduce_group (torch.distributed.ProcessGroup): The group for reducing the loss.
+ mean_group (torch.distributed.ProcessGroup): The group for averaging the loss.
+ """
+ # Skip aux loss logging if layer_number is None.
+ if layer_number is None:
+ return
+
+ tracker = parallel_state.get_moe_layer_wise_logging_tracker()
+ if name not in tracker:
+ tracker[name] = {}
+ tracker[name]["values"] = torch.zeros(num_layers, device=loss.device)
+ tracker[name]["values"][layer_number - 1] += loss.detach() # Aggregate the loss for the layer.
+ tracker[name]["reduce_group"] = reduce_group
+ tracker[name]["avg_group"] = avg_group
+
+
+def clear_aux_losses_tracker():
+ """Clear the auxiliary losses."""
+ tracker = parallel_state.get_moe_layer_wise_logging_tracker()
+ for name in tracker:
+ tracker[name]["values"].zero_()
+ tracker[name]["reduce_group"] = None
+ tracker[name]["avg_group"] = None
+
+
+def reduce_aux_losses_tracker_across_ranks():
+ """Collect and reduce the auxiliary losses across ranks."""
+ tracker = parallel_state.get_moe_layer_wise_logging_tracker()
+ for name in tracker:
+ values = tracker[name]["values"]
+ # Collect aux losses across PP.
+ torch.distributed.all_reduce(
+ values, group=parallel_state.get_pipeline_model_parallel_group()
+ )
+ # Reduce aux losses across ranks.
+ if tracker[name].get('reduce_group') is not None:
+ torch.distributed.all_reduce(values, group=tracker[name].get('reduce_group'))
+ if tracker[name].get('avg_group') is not None:
+ torch.distributed.all_reduce(
+ values, group=tracker[name]['avg_group'], op=torch.distributed.ReduceOp.AVG
+ )
+
+
+def track_moe_metrics(
+ loss_scale, iteration, writer, wandb_writer=None, total_loss_dict=None, per_layer_logging=False
+):
+ """Track the MoE metrics for logging."""
+ # Aux loss logging
+ reduce_aux_losses_tracker_across_ranks()
+ tracker = parallel_state.get_moe_layer_wise_logging_tracker()
+ if writer is not None:
+ aux_losses = {k: v['values'].float() * loss_scale for k, v in tracker.items()}
+ for name, loss_list in aux_losses.items():
+ if total_loss_dict is not None:
+ if name not in total_loss_dict:
+ total_loss_dict[name] = loss_list.mean()
+ else:
+ total_loss_dict[name] += loss_list.mean()
+
+ # currently when using add_scalars,
+ # torch.utils.add_scalars makes each timer its own run, which
+ # polutes the runs list, so we just add each as a scalar
+ writer.add_scalar(name, loss_list.mean(), iteration)
+ if per_layer_logging:
+ for i, loss in enumerate(loss_list.tolist()):
+ writer.add_scalar(f"moe/{name}_layer_{i}", loss, iteration)
+
+ # W&B logging lacks support for logging multiple scalars simultaneously.
+ # As a workaround, we log each scalar individually first, then we can create
+ # a custom panel to manually group them to a single plot.
+ if wandb_writer:
+ wandb_writer.log({f"{name}": loss_list.mean()}, iteration)
+ if per_layer_logging:
+ wandb_writer.log(
+ {
+ f"moe/{name}_layer_{i}": loss
+ for i, loss in enumerate(loss_list.tolist())
+ },
+ iteration,
+ )
+
+ clear_aux_losses_tracker()
+
+
+class moe_gather(torch.autograd.Function):
+ """Gather the input tensor based on the map tensor."""
+
+ @staticmethod
+ def forward(ctx, input_, map_):
+ """Gather the input tensor based on the map tensor."""
+ ctx.input_size = input_.size()
+ ctx.map = map_
+ return torch.gather(input_, 0, map_)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ """Scatter the grad_output tensor based on the map tensor."""
+ input_size = ctx.input_size
+ map_ = ctx.map
+
+ output = torch.zeros(
+ input_size, dtype=grad_output.dtype, device=torch.cuda.current_device()
+ )
+ output.scatter_add_(0, map_, grad_output)
+ return output, None, None
+
+
+class moe_scatter(torch.autograd.Function):
+ """Scatter the input tensor based on the map tensor."""
+
+ @staticmethod
+ def forward(ctx, input_, map_, output_size=None):
+ """Scatter the input tensor based on the map tensor."""
+ ctx.map = map_
+
+ if output_size is not None:
+ output = torch.zeros(output_size, dtype=input_.dtype, device=input_.device)
+ else:
+ output = torch.zeros_like(input_)
+
+ output.scatter_add_(0, map_, input_)
+ return output
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ """Gather the grad_output tensor based on the map tensor."""
+ map_ = ctx.map
+ grad_input = torch.gather(grad_output, 0, map_)
+ return grad_input, None, None, None
diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
new file mode 100644
index 0000000000..3e85ec53c5
--- /dev/null
+++ b/megatron/core/transformer/moe/router.py
@@ -0,0 +1,309 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+from abc import ABC, abstractmethod
+
+import torch
+
+from megatron.core import parallel_state
+from megatron.core.tensor_parallel import (
+ gather_from_sequence_parallel_region,
+ get_cuda_rng_tracker,
+ get_data_parallel_rng_tracker_name,
+)
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.moe.moe_utils import (
+ MoEAuxLossAutoScaler,
+ save_to_aux_losses_tracker,
+ sinkhorn,
+ switch_load_balancing_loss_func,
+ topk_softmax_with_capacity,
+ z_loss_func,
+)
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class Router(ABC, MegatronModule):
+ """Base Router class"""
+
+ def __init__(self, config: TransformerConfig) -> None:
+ """
+ Initialize the Router module.
+
+ Args:
+ config (TransformerConfig): Configuration object for the Transformer model.
+ """
+ super().__init__(config)
+ self.config = config
+ self.num_experts = self.config.num_moe_experts
+ self.moe_aux_loss_func = None
+ self.layer_number = None
+
+ # Initialize the gate weights.
+ self.weight = torch.nn.Parameter(
+ torch.empty((self.config.num_moe_experts, self.config.hidden_size), dtype=torch.float32)
+ )
+ if config.perform_initialization:
+ if get_cuda_rng_tracker().is_initialized():
+ with get_cuda_rng_tracker().fork(get_data_parallel_rng_tracker_name()):
+ config.init_method(self.weight)
+ else:
+ config.init_method(self.weight)
+ self.weight.data = self.weight.data.to(dtype=config.params_dtype)
+ setattr(self.weight, 'sequence_parallel', config.sequence_parallel)
+
+ def gating(self, input: torch.Tensor):
+ """Forward pass of the router gate.
+
+ Args:
+ input (torch.Tensor): Input tensor.
+
+ Returns:
+ torch.Tensor: Logits tensor.
+ """
+ if self.weight.device.type == 'cpu':
+ # move weights to GPU
+ self.weight.data = self.weight.data.to(device=torch.cuda.current_device())
+ logits = torch.nn.functional.linear(input, self.weight)
+ return logits
+
+ @abstractmethod
+ def routing(self, logits: torch.Tensor):
+ """Routing function.
+
+ Args:
+ logits (torch.Tensor): Logits tensor.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]:
+ Tuple of tensors representing max probs and the indices.
+ """
+ raise NotImplementedError("Routing function not implemented.")
+
+ @abstractmethod
+ def forward(self, input: torch.Tensor):
+ """
+ Forward pass of the router.
+
+ Args:
+ input (torch.Tensor): Input tensor.
+ """
+ raise NotImplementedError("Forward function not implemented.")
+
+ def set_layer_number(self, layer_number: int):
+ """Set the layer number for the router."""
+ self.layer_number = layer_number
+
+
+class TopKRouter(Router):
+ """Route each token to the top-k experts."""
+
+ def __init__(self, config: TransformerConfig) -> None:
+ """Initialize the zero token dropping router.
+
+ Args:
+ config (TransformerConfig): The configuration for the transformer model.
+ """
+ super().__init__(config=config)
+ self.topk = self.config.moe_router_topk
+ self.routing_type = self.config.moe_router_load_balancing_type
+ self.input_jitter = None
+
+ def sinkhorn_load_balancing(self, logits: torch.Tensor):
+ """Apply sinkhorn routing to the logits tensor.
+
+ Args:
+ logits (torch.Tensor): The logits tensor.
+
+ Returns:
+ torch.Tensor: The logits tensor after applying sinkhorn routing.
+ """
+
+ def _sinkhorn_activation(logits):
+ if self.topk == 1:
+ logits = torch.sigmoid(logits)
+ else: # k > 1
+ logits = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)
+ return logits
+
+ assert self.config.moe_aux_loss_coeff == 0, "Sinkhorn routing does not support aux loss."
+ if self.training:
+ with torch.no_grad():
+ norm_logits = sinkhorn(
+ logits.to(dtype=torch.float32)
+ ) # explicit fp32 conversion for stability
+ _, indices = torch.topk(norm_logits, k=self.topk, dim=1)
+ logits = _sinkhorn_activation(logits)
+ scores = torch.gather(logits, 1, indices)
+ else:
+ logits = _sinkhorn_activation(logits)
+ scores, indices = torch.topk(logits, k=self.topk, dim=1)
+ return scores, indices
+
+ def aux_loss_load_balancing(self, logits: torch.Tensor):
+ """Apply loss-based load balancing to the logits tensor.
+
+ Args:
+ logits (torch.Tensor): the logits tensor after gating, shape: [num_tokens, num_experts].
+
+ Returns:
+ probs (torch.Tensor): the probabilities tensor after load balancing.
+ indices (torch.Tensor): the indices tensor after top-k selection.
+ """
+ probs, indices, tokens_per_expert = topk_softmax_with_capacity(
+ logits,
+ self.topk,
+ capacity_factor=self.config.moe_expert_capacity_factor,
+ pad_to_capacity=self.config.moe_pad_expert_input_to_capacity,
+ drop_policy=self.config.moe_token_drop_policy,
+ use_pre_softmax=self.config.moe_router_pre_softmax,
+ deterministic_mode=self.config.deterministic_mode,
+ )
+
+ if self.training:
+ # Apply load balancing loss
+ scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
+ probs = self.apply_load_balancing_loss(scores, tokens_per_expert, activation=probs)
+ return probs, indices
+
+ def apply_load_balancing_loss(
+ self,
+ probs: torch.Tensor,
+ num_local_tokens_per_expert: torch.Tensor,
+ activation: torch.Tensor,
+ ):
+ """Applies auxiliary loss to the MoE layer.
+
+ Args:
+ probs (torch.Tensor):
+ The probs output by the router for each token. [num_tokens, num_experts]
+ num_local_tokens_per_expert (torch.Tensor):
+ The number of tokens per expert. [num_experts]
+ activation (torch.Tensor): The activation tensor to attach the gradient function to.
+
+ Returns:
+ torch.Tensor: The activation tensor with the attached gradient function.
+ """
+ moe_aux_loss_coeff = self.config.moe_aux_loss_coeff
+ sequence_partition_group = None
+ if self.config.moe_token_dispatcher_type == "alltoall_seq":
+ sequence_partition_group = parallel_state.get_context_parallel_group()
+ moe_aux_loss_coeff /= parallel_state.get_tensor_model_parallel_world_size()
+ else:
+ sequence_partition_group = parallel_state.get_tensor_and_context_parallel_group()
+
+ aux_loss = switch_load_balancing_loss_func(
+ probs,
+ num_local_tokens_per_expert,
+ self.topk,
+ moe_aux_loss_coeff,
+ sequence_partition_group=sequence_partition_group,
+ )
+ save_to_aux_losses_tracker(
+ "load_balancing_loss",
+ aux_loss / moe_aux_loss_coeff,
+ self.layer_number,
+ self.config.num_layers,
+ reduce_group=sequence_partition_group,
+ )
+ activation = MoEAuxLossAutoScaler.apply(activation, aux_loss)
+ return activation
+
+ def apply_z_loss(self, logits):
+ """Encourages the router's logits to remain small to enhance stability.
+ Please refer to the ST-MoE paper (https://arxiv.org/pdf/2202.08906.pdf) for details.
+
+ Args:
+ logits (torch.Tensor): The logits of the router.
+
+ Returns:
+ torch.Tensor: The logits after applying the z-loss.
+ """
+ if self.config.moe_z_loss_coeff is not None and self.training:
+ moe_z_loss_coeff = (
+ self.config.moe_z_loss_coeff
+ / parallel_state.get_tensor_and_context_parallel_world_size()
+ )
+ z_loss = z_loss_func(logits, moe_z_loss_coeff)
+ logits = MoEAuxLossAutoScaler.apply(logits, z_loss)
+ save_to_aux_losses_tracker(
+ "z_loss", z_loss / moe_z_loss_coeff, self.layer_number, self.config.num_layers
+ )
+ return logits
+
+ def apply_input_jitter(self, input: torch.Tensor):
+ """Add noise to the input tensor.
+ Refer to https://arxiv.org/abs/2101.03961.
+
+ Args:
+ input (Tensor): Input tensor.
+
+ Returns:
+ Tensor: Jittered input.
+ """
+ if self.config.moe_input_jitter_eps is not None:
+ eps = self.config.moe_input_jitter_eps
+ if self.input_jitter is None:
+ self.input_jitter = torch.distributions.uniform.Uniform(
+ torch.tensor(1.0 - eps, device=input.device),
+ torch.tensor(1.0 + eps, device=input.device),
+ ).rsample
+ return input * self.input_jitter(input.shape)
+ else:
+ return input
+
+ def routing(self, logits: torch.Tensor):
+ """Top-k routing function
+
+ Args:
+ logits (torch.Tensor): Logits tensor after gating.
+
+ Returns:
+ probs (torch.Tensor): the probabilities tensor after load balancing.
+ indices (torch.Tensor): the indices tensor after top-k selection.
+ """
+ logits = logits.view(-1, self.config.num_moe_experts)
+
+ # Apply Z-Loss
+ logits = self.apply_z_loss(logits)
+
+ if self.config.moe_token_dispatcher_type == "alltoall_seq":
+ # Gather the logits from the TP region
+ logits = gather_from_sequence_parallel_region(logits)
+
+ if self.routing_type == "sinkhorn":
+ scores, indices = self.sinkhorn_load_balancing(logits)
+ elif self.routing_type == "aux_loss":
+ scores, indices = self.aux_loss_load_balancing(logits)
+ elif self.routing_type == "none":
+ # A naive top-k routing without load balancing
+ scores, indices, _ = topk_softmax_with_capacity(
+ logits,
+ self.topk,
+ capacity_factor=self.config.moe_expert_capacity_factor,
+ pad_to_capacity=self.config.moe_pad_expert_input_to_capacity,
+ drop_policy=self.config.moe_token_drop_policy,
+ use_pre_softmax=self.config.moe_router_pre_softmax,
+ deterministic_mode=self.config.deterministic_mode,
+ )
+ else:
+ raise ValueError(f"Unsupported MoE routing type: {self.routing_type}")
+
+ return scores, indices
+
+ def forward(self, input: torch.Tensor):
+ """
+ Forward pass of the router.
+
+ Args:
+ input (torch.Tensor): Input tensor.
+ """
+ self.hidden = input.shape[-1]
+
+ # Apply input jitter
+ input = self.apply_input_jitter(input)
+ logits = self.gating(input)
+ logits = logits.view(-1, self.config.num_moe_experts)
+
+ scores, indices = self.routing(logits)
+
+ return scores, indices
diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py
new file mode 100644
index 0000000000..c2d9c188e3
--- /dev/null
+++ b/megatron/core/transformer/moe/shared_experts.py
@@ -0,0 +1,262 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import warnings
+from copy import deepcopy
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.fusions.fused_bias_geglu import bias_geglu_impl
+from megatron.core.fusions.fused_bias_gelu import bias_gelu_impl
+from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl
+from megatron.core.tensor_parallel.mappings import (
+ copy_to_tensor_model_parallel_region,
+ gather_from_sequence_parallel_region,
+ reduce_from_tensor_model_parallel_region,
+ reduce_scatter_to_sequence_parallel_region,
+)
+from megatron.core.tensor_parallel.random import (
+ get_cuda_rng_tracker,
+ get_data_parallel_rng_tracker_name,
+)
+from megatron.core.transformer.mlp import MLP
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import make_sharded_tensor_for_checkpoint
+
+
+class SharedExpertMLP(MLP):
+ """
+ MLP layer for Shared Experts.
+ """
+
+ # This stream is used when '--moe-shared-expert-overlap' is set.
+ # The shared experts are scheduled into this stream to be overlapped with the dispatcher.
+ stream = None
+
+ def __init__(self, config: TransformerConfig, spec: ModuleSpec):
+ config = deepcopy(config)
+ assert config.add_bias_linear == False, "bias is not supported in the shared experts, "
+ "please set '--disable-bias-linear' instead."
+
+ config.ffn_hidden_size = config.moe_shared_expert_intermediate_size
+ super().__init__(config=config, submodules=spec.submodules)
+
+ self.use_shared_expert_gate = spec.params.get("gate", False)
+ if self.use_shared_expert_gate:
+ self.gate_weight = torch.nn.Parameter(torch.empty((1, self.config.hidden_size)))
+ if config.perform_initialization:
+ if get_cuda_rng_tracker().is_initialized():
+ with get_cuda_rng_tracker().fork(get_data_parallel_rng_tracker_name()):
+ config.init_method(self.gate_weight)
+ else:
+ config.init_method(self.gate_weight)
+ self.gate_weight.data = self.gate_weight.data.to(dtype=config.params_dtype)
+ setattr(self.gate_weight, 'sequence_parallel', self.config.sequence_parallel)
+ else:
+ self.gate_weight = None
+
+ if self.config.moe_shared_expert_overlap:
+ # disable TP related AG/RS communications in the linear module
+ for linear in [self.linear_fc1, self.linear_fc2]:
+ if hasattr(linear, 'parallel_mode'):
+ # TELinear
+ linear.parallel_mode = None
+ else:
+ # MCore legacy Linear
+ linear.explicit_expert_comm = True
+
+ # The overlapped version is splitted into some separated functions and is put inside
+ # the token dispatcher. These functions should be called in this order and no one can
+ # be skipped:
+ # pre_forward_comm(input)
+ # linear_fc1_forward_and_act()
+ # linear_fc2_forward()
+ # post_forward_comm()
+ # output = get_output()
+ #
+ # We use cached intermediate results to avoid messy arg passing in the dispatcher.
+ self.cached_fc1_input = None
+ self.cached_fc2_input = None
+ self.cached_fc2_output = None
+ self.cached_output = None
+ self.gate_score = None
+
+ if self.stream is None:
+ self.stream = torch.cuda.Stream()
+
+ def forward(self, hidden_states):
+ """Forward function"""
+ output, _ = super().forward(hidden_states)
+ if self.use_shared_expert_gate:
+ logits = torch.nn.functional.linear(hidden_states, self.gate_weight)
+ gate_score = torch.nn.functional.sigmoid(logits)
+ output = output * gate_score
+ return output
+
+ def sharded_state_dict(
+ self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None
+ ) -> ShardedStateDict:
+ """Gets sharded state dict."""
+ sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+ if self.use_shared_expert_gate:
+ name = 'gate_weight'
+ state_dict = self.state_dict(prefix='', keep_vars=True)
+ sub_sd = {
+ f'{prefix}{name}': make_sharded_tensor_for_checkpoint(
+ state_dict[name], f'{prefix}{name}', prepend_offsets=sharded_offsets
+ )
+ }
+ sharded_state_dict.update(sub_sd)
+ return sharded_state_dict
+
+ def pre_forward_comm(self, input):
+ """
+ All Gather for SP before forward.
+ This function is used to overlap shared experts with the dispatcher.
+ It is only useful when --moe-shared-expert-overlap is set and may be changed.
+ """
+ assert self.config.moe_shared_expert_overlap
+ assert self.cached_output is None
+ self.stream.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(self.stream):
+ if self.use_shared_expert_gate:
+ logits = torch.nn.functional.linear(input, self.gate_weight)
+ self.gate_score = torch.nn.functional.sigmoid(logits)
+ if self.config.sequence_parallel:
+ self.cached_fc1_input = gather_from_sequence_parallel_region(
+ input, tensor_parallel_output_grad=True
+ )
+ else:
+ self.cached_fc1_input = copy_to_tensor_model_parallel_region(input)
+ set_tensor_grad_fn_sequence_sr(self.cached_fc1_input, torch.iinfo(torch.int).max)
+
+ def linear_fc1_forward_and_act(self, overlapped_comm_output=None):
+ """
+ Do Linear FC1 and activation function forward.
+ This function is used to overlap shared experts with the dispatcher.
+ It is only useful when --moe-shared-expert-overlap is set and may be changed.
+ """
+ assert self.config.moe_shared_expert_overlap
+ assert self.cached_fc1_input is not None
+ if overlapped_comm_output is not None:
+ set_tensor_grad_fn_sequence_sr(overlapped_comm_output, torch.iinfo(torch.int).max)
+ with torch.cuda.stream(self.stream):
+ # [s, b, 4 * h/p]
+ intermediate_parallel, bias_parallel = self.linear_fc1(self.cached_fc1_input)
+ self.cached_fc1_input = None
+
+ if self.config.bias_activation_fusion:
+ if self.activation_func == F.gelu:
+ if self.config.gated_linear_unit:
+ intermediate_parallel = bias_geglu_impl(
+ intermediate_parallel, bias_parallel
+ )
+ else:
+ assert self.config.add_bias_linear is True
+ intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel)
+ elif self.activation_func == F.silu and self.config.gated_linear_unit:
+ intermediate_parallel = bias_swiglu_impl(
+ intermediate_parallel,
+ bias_parallel,
+ self.config.activation_func_fp8_input_store,
+ )
+ else:
+ raise ValueError("Only support fusion of gelu and swiglu")
+ else:
+ if bias_parallel is not None:
+ intermediate_parallel = intermediate_parallel + bias_parallel
+ if self.config.gated_linear_unit:
+
+ def glu(x):
+ x = torch.chunk(x, 2, dim=-1)
+ return self.config.activation_func(x[0]) * x[1]
+
+ intermediate_parallel = glu(intermediate_parallel)
+ else:
+ intermediate_parallel = self.activation_func(intermediate_parallel)
+
+ self.cached_fc2_input = intermediate_parallel
+
+ def linear_fc2_forward(self, overlapped_comm_output=None):
+ """
+ Do Linear FC2 forward.
+ This function is used to overlap shared experts with the dispatcher.
+ It is only useful when --moe-shared-expert-overlap is set and may be changed.
+ """
+ assert self.config.moe_shared_expert_overlap
+ assert self.cached_fc2_input is not None
+ if overlapped_comm_output is not None:
+ set_tensor_grad_fn_sequence_sr(overlapped_comm_output, torch.iinfo(torch.int).max)
+ with torch.cuda.stream(self.stream):
+ # [s, b, h]
+ self.cached_fc2_output, _ = self.linear_fc2(self.cached_fc2_input)
+ self.cached_fc2_input = None
+
+ def post_forward_comm(self):
+ """
+ Reduce scatter for SP after forward.
+ This function is used to overlap shared experts with the dispatcher.
+ It is only useful when --moe-shared-expert-overlap is set and may be changed.
+ """
+ assert self.config.moe_shared_expert_overlap
+ assert self.cached_fc2_output is not None
+ with torch.cuda.stream(self.stream):
+ if self.config.sequence_parallel:
+ self.cached_output = reduce_scatter_to_sequence_parallel_region(
+ self.cached_fc2_output
+ )
+ else:
+ self.cached_output = reduce_from_tensor_model_parallel_region(
+ self.cached_fc2_output
+ )
+ self.cached_fc2_output = None
+ set_tensor_grad_fn_sequence_sr(self.cached_output, torch.iinfo(torch.int).max)
+
+ def get_output(self):
+ """
+ Gets the module forward output.
+ This function is used to overlap shared experts with the dispatcher.
+ It is only useful when --moe-shared-expert-overlap is set and may be changed.
+ """
+ assert self.config.moe_shared_expert_overlap
+ assert self.cached_output is not None
+ with torch.cuda.stream(self.stream):
+ if self.use_shared_expert_gate:
+ assert self.gate_score is not None
+ output = self.cached_output * self.gate_score
+ self.gate_score = None
+ else:
+ output = self.cached_output
+ self.cached_output = None
+ torch.cuda.current_stream().wait_stream(self.stream)
+ return output
+
+
+TORCH_MAJOR = int(torch.__version__.split(".")[0])
+TORCH_MINOR = int(torch.__version__.split(".")[1])
+TORCH_LAST = torch.__version__.split(".")[2]
+
+
+def set_tensor_grad_fn_sequence_sr(tensor, value):
+ """
+ Set sequence_sr for the grad_fn of a tensor to control the backward order.
+ For older PyTorch version, do nothing (backward order is not changed).
+ The bigger the value is, the earlier the grad_fn is scheduled.
+ """
+ if (
+ (TORCH_MAJOR > 2)
+ or (TORCH_MAJOR == 2 and TORCH_MINOR > 2)
+ or (TORCH_MAJOR == 2 and TORCH_MINOR == 2 and '+' not in TORCH_LAST)
+ ):
+ # In NVIDIA PyTorch container 24.01, the PyTorch version is 2.2.0a0+81ea7a4,
+ # which does not contian the set_sequence_nr commit.
+ if tensor is not None and tensor.grad_fn is not None:
+ tensor.grad_fn._set_sequence_nr(value)
+ else:
+ warnings.warn(
+ "WARNING : PyTorch is too old to set sequence_sr and the performance may not "
+ "optimal. Please use PyTorch >= 2.2.0 for better performance."
+ )
diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
new file mode 100644
index 0000000000..db1b1920fa
--- /dev/null
+++ b/megatron/core/transformer/moe/token_dispatcher.py
@@ -0,0 +1,627 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from abc import abstractmethod
+from typing import List, Optional, Tuple
+
+import torch
+
+from megatron.core import parallel_state, tensor_parallel
+from megatron.core.tensor_parallel.mappings import (
+ _gather_along_first_dim_moe,
+ gather_from_sequence_parallel_region,
+ reduce_scatter_to_sequence_parallel_region,
+)
+from megatron.core.transformer.moe.moe_utils import (
+ moe_gather,
+ moe_scatter,
+ permute,
+ sort_chunks_by_idxs,
+ unpermute,
+)
+from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+""" We use the following notation throughout this file:
+ H: hidden size
+ B: micro batch size
+ S: sequence length
+ TP: tensor model parallel size
+ EP: expert model parallel size
+ num_local_tokens: S/TP*B
+ num_global_tokens: num_local_tokens*TP*EP
+"""
+
+
+class MoETokenDispatcher:
+ """
+ MoE Token Dispatcher
+ """
+
+ def __init__(self, config: TransformerConfig) -> None:
+ """
+ Initialize the MoE Token Dispatcher.
+ """
+ self.config = config
+ self.shared_experts: Optional[SharedExpertMLP] = None
+
+ @abstractmethod
+ def token_permutation(self, tokens: torch.Tensor, indices: torch.Tensor):
+ """Dispatch tokens to experts.
+
+ Args:
+ tokens (torch.Tensor): Input tokens.
+ indices (torch.Tensor): indices tensor.
+
+ Returns:
+ torch.Tensor: Tokens tensor.
+ """
+ raise NotImplementedError("Dispatch function not implemented.")
+
+ @abstractmethod
+ def token_unpermutation(
+ self, expert_output: torch.Tensor, probs: torch.Tensor, indices: torch.Tensor
+ ):
+ """Restores the expert output to its original ordering.
+
+ Args:
+ expert_output (torch.Tensor): The output tensor from the expert models.
+ probs (torch.Tensor): Each token's score with each expert.
+ indices (torch.Tensor): The indices used to reorder the expert output.
+
+ Returns:
+ (torch.Tensor, torch.Tensor): Unpermuted activation and optional bias.
+ """
+ raise NotImplementedError("Restore function not implemented.")
+
+ def set_shared_experts(self, shared_experts):
+ """Set shared expert to the dispatcher."""
+ self.shared_experts = shared_experts
+
+
+class MoEAllGatherTokenDispatcher(MoETokenDispatcher):
+ """
+ AllGather Based Token dispatcher.
+ Note that this allgather spans the communication domain of TP*EP:
+ """
+
+ def __init__(
+ self, num_local_experts: int, local_expert_indices: List[int], config: TransformerConfig
+ ) -> None:
+ """
+ Initialize the zero token dropping router.
+ """
+ super().__init__(config=config)
+ self.num_local_experts = num_local_experts
+ assert self.num_local_experts > 0, "Expected at least one expert"
+ self.local_expert_indices = local_expert_indices
+ assert len(self.local_expert_indices) > 0, "Expected at least one local expert index"
+ self.router_topk = config.moe_router_topk
+ self.add_bias = config.add_bias_linear
+
+ # self.local_probs: probs of global token assignment to local experts.
+ self.local_probs = None
+
+ # self.global_local_map: 2D tensor. A mask of mapping between global and local tokens where
+ # each element is True if it's between the local_expert_indices. Only useful when cross
+ # device token permutation is enabled and **AllGahter** is performed.
+ self.global_local_map = None
+
+ def token_permutation(
+ self, hidden_states: torch.Tensor, max_prob: torch.Tensor, max_ind: torch.Tensor
+ ):
+ """Dispatch tokens to local experts. It's composed of two stages:
+ (1) Permute the tokens across the expert parallel devices. After this stage,
+ each device receives all of the tokens assigned to its local set of experts
+ in its local HBM.
+ (2) Permute the tokens locally so that they are grouped by their expert
+ assignment. After the stage (1), the tokens are grouped by which device
+ they came from. We re-order them locally for subsequent efficient computation.
+
+ Args:
+ hidden_states: 3D tensor [S/TP, B, H]. Input tokens.
+ max_prob: 2D tensor [S/TP*B, topk]. Each row of max_prob contains
+ the probility distribution across `topk` experts for one local token.
+ For 'aux_loss' load balancing, the sum of the values in each row is 1,
+ thus for `top1` gating, it degenerates into a full 1 tensor.
+ max_ind: 2D tensor [num_local_tokens, topk], where
+ `num_local_tokens=S/TP*B`. Token assignment to global experts.
+
+ Returns:
+ permuted_local_hidden_states: Permutation of tokens to local experts group.
+ tokens_per_expert: the number of tokens each local expert to process.
+ """
+ self.hidden_shape = hidden_states.shape
+ # [S/TP, B, H] -> [S*B/TP, H]
+ hidden_states = hidden_states.view(-1, self.hidden_shape[-1])
+
+ # Permute the tokens across the expert parallel devices.
+ if (self.config.tensor_model_parallel_size > 1) or (
+ self.config.expert_model_parallel_size > 1
+ ):
+ ## local_indices calculation
+ with torch.no_grad():
+ # [num_local_tokens, topk] -> [num_global_tokens, topk], where:
+ # num_local_tokens=(S/TP)*B, num_global_tokens=S*B*EP
+ global_indices = tensor_parallel.gather_from_sequence_parallel_region_to_moe(
+ max_ind
+ )
+ # Create a mask of mapping between global and local tokens where each
+ # element is True if it's between the local_expert_indices
+ global_local_mask = (global_indices >= self.local_expert_indices[0]) & (
+ global_indices <= self.local_expert_indices[-1]
+ )
+ local_indices = global_indices.masked_select(global_local_mask)
+
+ ## local_probs calculation
+ # max_prob: [S/TP*B, topk] -> global_probs: [S*B*EP, topk]
+ global_probs = tensor_parallel.gather_from_sequence_parallel_region_to_moe(max_prob)
+ self.local_probs = global_probs.masked_select(global_local_mask)
+ self.local_probs = self.local_probs.view(-1, 1)
+ # Note that this allgather spans the communication domain of TP*EP.
+ # [(S/TP)*B, H] -> [((S/TP)*B)*(TP*EP), H] = [S*B*EP, H]
+ global_hidden_states = tensor_parallel.gather_from_sequence_parallel_region_to_moe(
+ hidden_states, use_global_buffer=True
+ )
+ # Reshape global_local_mask to be compatible with Tensor.gather
+ global_local_map = global_local_mask.nonzero()[:, 0]
+ self.global_local_map = global_local_map.view(-1, 1).expand(-1, hidden_states.shape[-1])
+ local_hidden_states = moe_gather.apply(global_hidden_states, self.global_local_map)
+ else:
+ if self.router_topk > 1:
+ global_local_mask = torch.ones_like(max_ind).bool()
+ local_indices = max_ind.masked_select(global_local_mask)
+ self.local_probs = max_prob.masked_select(global_local_mask)
+ self.local_probs = self.local_probs.view(-1, 1)
+ global_local_map = global_local_mask.nonzero()[:, 0]
+ self.global_local_map = global_local_map.view(-1, 1).expand(
+ -1, hidden_states.shape[-1]
+ )
+ local_hidden_states = torch.gather(hidden_states, 0, self.global_local_map)
+ else:
+ local_indices = max_ind
+ self.local_probs = max_prob.view(-1, 1)
+ local_hidden_states = hidden_states
+ self.global_local_map = None
+
+ with torch.no_grad():
+ # The indices of local_indices that give its sorted order along dim 0.
+ self.indices = torch.argsort(local_indices, dim=0)
+ if self.config.deterministic_mode:
+ tokens_per_expert = torch.bincount(
+ local_indices.view(-1), minlength=self.config.num_moe_experts
+ )
+ if self.num_local_experts < self.config.num_moe_experts:
+ tokens_per_expert = tokens_per_expert[
+ self.local_expert_indices[0] : self.local_expert_indices[-1] + 1
+ ]
+ else:
+ tokens_per_expert = torch.histc(
+ local_indices,
+ bins=self.num_local_experts,
+ min=self.local_expert_indices[0],
+ max=self.local_expert_indices[-1],
+ )
+ tokens_per_expert = tokens_per_expert.cpu().to(torch.long)
+
+ # Stage2: permute the tokens locally so that they are grouped by their expert assignment
+ # Reshape indices to be compatible with Tensor.gather
+
+ permuted_local_hidden_states, self.reversed_local_input_permutation_mapping = permute(
+ local_hidden_states, local_indices
+ )
+
+ return permuted_local_hidden_states, tokens_per_expert
+
+ def token_unpermutation(self, hidden_states: torch.Tensor, bias: torch.Tensor = None):
+ """
+ Reverse process of `dispatch()` which permutes the output of local
+ experts locallay and across expert parallel rank into the original order to
+ produce the final output.
+
+ Args:
+ hidden_states: 2D tensor [num_permuted_tokens_for_local_experts, H],
+ output of local experts.
+ bias (optional): The bias tensor.
+
+ Returns:
+ output_total: un-permuted updated hidden states output from all local experts
+ with shape of [S/TP, B, H]
+ """
+ # Stage1: unpermute the tokens and bias locally respectively.
+ # Scale the expert output prior to reduction and subsequent to local unpermutation if k > 1.
+
+ unpermuted_local_hidden = unpermute(
+ hidden_states, self.reversed_local_input_permutation_mapping
+ )
+ unpermuted_local_hidden = unpermuted_local_hidden * self.local_probs
+
+ unpermuted_local_bias = None
+ if self.add_bias:
+ assert bias is not None
+ unpermuted_local_bias = torch.zeros_like(hidden_states)
+ unpermuted_local_bias = unpermute(bias, self.reversed_local_input_permutation_mapping)
+ unpermuted_local_bias = unpermuted_local_bias * self.local_probs
+
+ output_total = unpermuted_local_hidden
+ output_bias_total = unpermuted_local_bias
+
+ # Unpermute the tokens across expert parallel devices.
+ if (self.config.tensor_model_parallel_size > 1) or (
+ self.config.expert_model_parallel_size > 1
+ ):
+ assert (
+ self.global_local_map is not None
+ ), "global_local_map is necessary for `AllGather`."
+ ep_group_size = parallel_state.get_tensor_and_expert_parallel_world_size()
+ # hidden_shape: [S/TP, B, H], gloal_num_tokens = S/TP*B*(TP*EP)
+ global_num_tokens = self.hidden_shape[0] * self.hidden_shape[1] * ep_group_size
+ global_hidden_shape = [global_num_tokens, hidden_states.shape[-1]]
+ assert self.global_local_map.shape == unpermuted_local_hidden.shape
+ unpermuted_global_hidden = moe_scatter.apply(
+ unpermuted_local_hidden, self.global_local_map, global_hidden_shape
+ )
+ output_total = tensor_parallel.reduce_scatter_to_sequence_parallel_region_from_moe(
+ unpermuted_global_hidden
+ )
+ if self.add_bias:
+ # Unpermute the bias across expert parallel devices.
+ unpermuted_global_bias = torch.zeros_like(unpermuted_global_hidden)
+ unpermuted_global_bias = unpermuted_global_bias.scatter_add(
+ 0, self.global_local_map, unpermuted_local_bias
+ )
+ output_bias_total = (
+ tensor_parallel.reduce_scatter_to_sequence_parallel_region_from_moe(
+ unpermuted_global_bias
+ )
+ )
+ # bias is duplicated across tensor parallelism ranks;
+ # reduce scatter reduces bias across tensor parallel_ranks
+ output_bias_total = (
+ output_bias_total / parallel_state.get_tensor_model_parallel_world_size()
+ )
+ else:
+ if self.router_topk > 1:
+ global_num_tokens = self.hidden_shape[0] * self.hidden_shape[1]
+ global_hidden_shape = [global_num_tokens, hidden_states.shape[-1]]
+ unpermuted_global_hidden = torch.zeros(
+ global_hidden_shape,
+ dtype=hidden_states.dtype,
+ device=torch.cuda.current_device(),
+ )
+ output_total = unpermuted_global_hidden.scatter_add(
+ 0, self.global_local_map, unpermuted_local_hidden
+ )
+ if self.add_bias:
+ unpermuted_global_bias = torch.zeros_like(unpermuted_global_hidden)
+ output_bias_total = unpermuted_global_bias.scatter_add(
+ 0, self.global_local_map, unpermuted_local_bias
+ )
+
+ output_total = output_total.view(self.hidden_shape)
+ if self.add_bias:
+ output_bias_total = output_bias_total.view(self.hidden_shape)
+ else:
+ output_bias_total = None
+
+ return output_total, output_bias_total
+
+
+class MoEAlltoAllTokenDispatcher(MoETokenDispatcher):
+ """
+ AlltoAll-based token dispatcher.
+ """
+
+ def __init__(
+ self, num_local_experts: int, local_expert_indices: List[int], config: TransformerConfig
+ ) -> None:
+ """
+ Initialize the AlltoAll token dispatcher.
+
+ Args:
+ num_local_experts (int): Number of local experts on the current device.
+ local_expert_indices (List[int]): Indices of local experts on the current device.
+ config (TransformerConfig): Configuration for the transformer model.
+ """
+ super().__init__(config=config)
+ self.hidden_shape = None
+ self.num_local_experts = num_local_experts
+ self.num_experts = config.num_moe_experts
+ assert self.num_local_experts > 0, "Expected at least one expert"
+ self.local_expert_indices = local_expert_indices
+ assert (
+ len(self.local_expert_indices) == self.num_local_experts
+ ), "Invalid local expert indices"
+ for i in range(len(self.local_expert_indices) - 1):
+ assert (
+ self.local_expert_indices[i] == self.local_expert_indices[i + 1] - 1
+ ), "local_expert_indices must be continous"
+ self.ep_size = config.expert_model_parallel_size
+ self.tp_size = config.tensor_model_parallel_size
+ self.probs = None
+
+ # [ep_size]. Represents the number of tokens sent by the current rank to other
+ # EP ranks.
+ self.input_splits = None
+ # [ep_size]. Represents the number of tokens received by the current rank from
+ # other EP ranks.
+ self.output_splits = None
+ # [tp_size]. Represents the number of tokens received by the current rank from
+ # other TP ranks.
+ self.output_splits_tp = None
+ # [tp_size * ep_size, num_local_experts]. Represents the number of tokens sent
+ # to each local expert by all ranks.
+ self.num_global_tokens_per_local_expert_cpu = None
+ input_chunk_idxs = torch.arange(self.num_experts * self.tp_size)
+ # [num_local_experts, tp_size * ep_size]. Sort the input chunks by local experts.
+ self.sort_input_by_local_experts = (
+ input_chunk_idxs.reshape(-1, self.num_local_experts).T.ravel().tolist()
+ )
+ # [tp_size * ep_size, num_local_experts]. Restore the output chunks by local experts.
+ self.restore_output_by_local_experts = (
+ input_chunk_idxs.reshape(self.num_local_experts, -1).T.ravel().tolist()
+ )
+
+ # Token drop and padding.
+ # We need to keep track of the token num if we drop tokens without padding them.
+ self.num_out_tokens = None
+ # Drop and pad the input to capacity.
+ self.drop_and_pad = self.config.moe_pad_expert_input_to_capacity
+ if self.drop_and_pad:
+ assert self.config.moe_expert_capacity_factor is not None
+ self.capacity = None
+
+ # A cuda stream synchronization is needed in self.token_permutation() in some cases,
+ # because there are several non-blocking DtoH data transfers called in self.preprocess().
+ # The synchronization happens at different points based on MoE settings as late as possible.
+ # Valid sync points are "before_permutation_1", "before_ep_alltoall", "before_finish",
+ # and "no_sync".
+ self.cuda_sync_point = "no_sync"
+
+ self.shared_experts = None
+
+ def preprocess(self, indices: torch.Tensor) -> torch.Tensor:
+ """
+ Preprocess token indices for AlltoAll communication and token permutation. This method
+ computes the number of tokens assigned to each expert based on the input indices.
+ It also initializes the necessary data structures for AlltoAll communication, such as input
+ and output splits, and the mapping between global tokens and local experts.
+
+ Args:
+ indices (torch.Tensor): Tensor of indices mapping tokens to experts.
+
+ Returns:
+ torch.Tensor: Tensor containing the number of tokens assigned to local expert.
+ """
+ if self.config.deterministic_mode:
+ num_local_tokens_per_expert = torch.bincount(
+ indices.view(-1), minlength=self.num_experts
+ )
+ else:
+ num_local_tokens_per_expert = torch.histc(
+ indices, bins=self.num_experts, min=0, max=self.num_experts
+ )
+ # num_local_tokens_per_expert: [num_experts]
+
+ tp_rank = parallel_state.get_tensor_model_parallel_rank()
+ if self.drop_and_pad:
+ # probs: [num_experts, local_capacity]
+ self.capacity = self.probs.size(1)
+ num_tokens_per_local_expert = torch.full(
+ (self.num_local_experts,),
+ self.capacity * self.tp_size * self.ep_size,
+ dtype=torch.long,
+ )
+ # [tp_size * ep_size, num_local_experts].
+ self.num_global_tokens_per_local_expert_cpu = torch.full(
+ (self.num_experts * self.tp_size,), self.capacity, dtype=torch.long
+ )
+ return num_tokens_per_local_expert
+ elif self.config.moe_expert_capacity_factor is not None:
+ # Token drop but no pad. A synchronization is needed before the first
+ # permutation to get the `num_out_tokens` CPU value.
+ self.num_out_tokens = num_local_tokens_per_expert.sum().to(
+ torch.device("cpu"), non_blocking=True
+ )
+ self.cuda_sync_point = "before_permutation_1"
+ elif self.ep_size > 1 or self.num_local_experts > 1:
+ # Token dropless and enable ep. A synchronization is needed before expert parallel
+ # AlltoAll communication to get the `input_splits` and `output_splits` CPU values.
+ self.cuda_sync_point = "before_ep_alltoall"
+ else:
+ # Token dropless and no ep. A synchronization is needed before the token_permutation()
+ # function returns to get the `tokens_per_expert` CPU value.
+ self.cuda_sync_point = "before_finish"
+
+ if self.ep_size > 1 or self.tp_size > 1:
+ # ===================================================
+ # Calculate input_splits, output_splits for alltoall/allgather in variable size.
+ # ===================================================
+ self.input_splits = (
+ num_local_tokens_per_expert.reshape(self.ep_size, self.num_local_experts)
+ .sum(axis=1)
+ .to(torch.device("cpu"), non_blocking=True)
+ .numpy()
+ )
+ # Gather the global distribution of tokens across ranks.
+ # num_global_tokens_per_expert represents the number of tokens sent to each
+ # expert by all ranks.
+ # [tp_size, ep_size, num_experts]
+ num_global_tokens_per_expert = (
+ _gather_along_first_dim_moe(num_local_tokens_per_expert)
+ .reshape(self.ep_size, self.tp_size, self.num_experts)
+ .transpose(0, 1)
+ )
+ # [tp_size, ep_size, num_experts] -> [tp_size, ep_size, num_local_experts]
+ num_global_tokens_per_local_expert = num_global_tokens_per_expert[
+ :, :, self.local_expert_indices[0] : self.local_expert_indices[-1] + 1
+ ].contiguous()
+ # [tp_size, ep_size, num_local_experts] -> [tp_size, ep_size]
+ num_global_tokens_per_rank = num_global_tokens_per_local_expert.sum(axis=2)
+ # [tp_size, ep_size] -> [ep_size]
+ # self.output_splits represents the number of tokens received by the current rank
+ # from other EP rank.
+ self.output_splits = (
+ num_global_tokens_per_rank[tp_rank]
+ .to(torch.device("cpu"), non_blocking=True)
+ .numpy()
+ )
+ # [tp_size, ep_size] -> [tp_size]
+ # self.output_splits_tp represents the number of tokens received by the current
+ # rank from other TP rank.
+ self.output_splits_tp = (
+ num_global_tokens_per_rank.sum(axis=1)
+ .to(torch.device("cpu"), non_blocking=True)
+ .numpy()
+ )
+ # [tp_size, ep_size, num_local_experts] -> [num_local_experts]
+ num_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=(0, 1)).to(
+ torch.device("cpu"), non_blocking=True
+ )
+ else:
+ num_global_tokens_per_local_expert = num_local_tokens_per_expert.reshape(
+ self.num_experts
+ )
+ num_tokens_per_local_expert = num_local_tokens_per_expert.to(
+ torch.device("cpu"), non_blocking=True
+ )
+
+ if self.num_local_experts > 1:
+ self.num_global_tokens_per_local_expert_cpu = num_global_tokens_per_local_expert.view(
+ -1, self.num_local_experts
+ ).to(torch.device("cpu"), non_blocking=True)
+
+ return num_tokens_per_local_expert
+
+ def token_permutation(
+ self, hidden_states: torch.Tensor, probs: torch.Tensor, indices: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Dispatch tokens to local experts using AlltoAll communication.
+
+ Args:
+ hidden_states (torch.Tensor): Input token embeddings.
+ probs (torch.Tensor): Probs of tokens assigned to experts.
+ indices (torch.Tensor): Indices of tokens assigned to experts.
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]:
+ - Permuted token embeddings for local experts.
+ - Number of tokens per expert.
+ """
+ # Preprocess: Get the metadata for communication, permutation and computation operations.
+ self.hidden_shape = hidden_states.shape
+ self.probs = probs
+ assert probs.dim() == 2, "Expected 2D tensor for probs"
+ assert indices.dim() == 2, "Expected 2D tensor for indices"
+ hidden_states = hidden_states.view(-1, self.hidden_shape[-1])
+ tokens_per_expert = self.preprocess(indices)
+
+ if self.shared_experts is not None:
+ self.shared_experts.pre_forward_comm(hidden_states.view(self.hidden_shape))
+
+ # Permutation 1: input to AlltoAll input
+ self.hidden_shape_before_permute = hidden_states.shape
+ if self.cuda_sync_point == "before_permutation_1":
+ torch.cuda.current_stream().synchronize()
+ permutated_local_input_tokens, self.reversed_local_input_permutation_mapping = permute(
+ hidden_states,
+ indices,
+ num_out_tokens=self.num_out_tokens,
+ padded_mode=self.drop_and_pad,
+ )
+
+ # Perform expert parallel AlltoAll communication
+ if self.cuda_sync_point == "before_ep_alltoall":
+ torch.cuda.current_stream().synchronize()
+ global_input_tokens = tensor_parallel.all_to_all(
+ parallel_state.get_expert_model_parallel_group(),
+ permutated_local_input_tokens,
+ self.output_splits,
+ self.input_splits,
+ )
+ if self.shared_experts is not None:
+ self.shared_experts.linear_fc1_forward_and_act(global_input_tokens)
+
+ if parallel_state.get_tensor_model_parallel_world_size() > 1:
+ global_input_tokens = gather_from_sequence_parallel_region(
+ global_input_tokens,
+ output_split_sizes=(
+ self.output_splits_tp.tolist() if self.output_splits_tp is not None else None
+ ),
+ )
+
+ # Permutation 2: Sort tokens by local expert.
+ if self.num_local_experts > 1:
+ global_input_tokens = sort_chunks_by_idxs(
+ global_input_tokens,
+ self.num_global_tokens_per_local_expert_cpu.ravel(),
+ self.sort_input_by_local_experts,
+ )
+
+ if self.cuda_sync_point == "before_finish":
+ torch.cuda.current_stream().synchronize()
+
+ return global_input_tokens, tokens_per_expert
+
+ def token_unpermutation(
+ self, hidden_states: torch.Tensor, bias: torch.Tensor = None
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """
+ Reverse the token permutation to restore the original order.
+
+ Args:
+ hidden_states (torch.Tensor): Output from local experts.
+ bias (torch.Tensor, optional): Bias tensor (not supported).
+
+ Returns:
+ Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ - Unpermuted token embeddings in the original order.
+ - None (bias is not supported).
+ """
+ assert bias is None, "Bias is not supported in MoEAlltoAllTokenDispatcher"
+
+ # Unpermutation 2: Unsort tokens by local expert.
+ if self.num_local_experts > 1:
+ hidden_states = sort_chunks_by_idxs(
+ hidden_states,
+ self.num_global_tokens_per_local_expert_cpu.T.ravel(),
+ self.restore_output_by_local_experts,
+ )
+
+ if parallel_state.get_tensor_model_parallel_world_size() > 1:
+ hidden_states = reduce_scatter_to_sequence_parallel_region(
+ hidden_states,
+ input_split_sizes=(
+ self.output_splits_tp.tolist() if self.output_splits_tp is not None else None
+ ),
+ )
+
+ # Perform expert parallel AlltoAll communication
+ # hidden_states: [SEQL, H] -> [SEQL, H/TP]
+ permutated_local_input_tokens = tensor_parallel.all_to_all(
+ parallel_state.get_expert_model_parallel_group(),
+ hidden_states,
+ self.input_splits,
+ self.output_splits,
+ )
+ if self.shared_experts is not None:
+ self.shared_experts.linear_fc2_forward(permutated_local_input_tokens)
+ self.shared_experts.post_forward_comm()
+
+ # Unpermutation 1: Unsort input tokens to restore the original order.
+ output = unpermute(
+ permutated_local_input_tokens,
+ self.reversed_local_input_permutation_mapping,
+ probs=self.probs,
+ padded_mode=self.drop_and_pad,
+ restore_shape=self.hidden_shape_before_permute,
+ )
+
+ # Reshape the output tensor
+ output = output.view(self.hidden_shape)
+
+ # Add shared experts output
+ if self.shared_experts is not None:
+ shared_expert_output = self.shared_experts.get_output()
+ output += shared_expert_output
+ return output, None
diff --git a/megatron/core/transformer/moe/upcycling_utils.py b/megatron/core/transformer/moe/upcycling_utils.py
new file mode 100644
index 0000000000..b905fc99be
--- /dev/null
+++ b/megatron/core/transformer/moe/upcycling_utils.py
@@ -0,0 +1,196 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved.
+""" Helpers for converting a dense model to a MoE model in runtime """
+from megatron.core import mpu
+
+
+def _get_keys_endswith(model, suffix):
+ """
+ Retrieve keys from the model that end with a specified suffix.
+ """
+ return [k for k in model if k.endswith(suffix)]
+
+
+def _covert_to_moe_state_dict(state_dict, moe_model):
+ """
+ Convert a dense model's state_dict to a MoE model's state_dict.
+
+ This function takes the state dictionary of a dense model and modifies it to fit the
+ structure required by a Mixture of Experts model. It handles the necessary
+ transformations for weights and biases specific to the MoE architecture.
+
+ Args:
+ state_dict (dict): The dense model's state_dict.
+ moe_model (nn.Module): The MoE model instance from which to get the submodule
+ and state_dict, must be a model without FP16 and/or
+ DDP wrapper.
+
+ Returns:
+ dict: The converted MoE model state_dict, ready for use in the MoE architecture.
+ """
+
+ mlp = moe_model.get_submodule('decoder.layers.0.mlp')
+
+ moe_state_dict = moe_model.state_dict()
+ new_state_dict = state_dict
+
+ mlp_lm_weight_keys = _get_keys_endswith(new_state_dict, 'mlp.linear_fc1.layer_norm_weight')
+ mlp_lm_bias_keys = _get_keys_endswith(new_state_dict, 'mlp.linear_fc1.layer_norm_bias')
+ mlp_fc1_weight_keys = _get_keys_endswith(new_state_dict, 'mlp.linear_fc1.weight')
+ mlp_fc2_weight_keys = _get_keys_endswith(new_state_dict, 'mlp.linear_fc2.weight')
+ mlp_fc1_bias_keys = _get_keys_endswith(new_state_dict, 'mlp.linear_fc1.bias')
+ mlp_fc2_bias_keys = _get_keys_endswith(new_state_dict, 'mlp.linear_fc2.bias')
+ mlp_fc1_extra_state_keys = _get_keys_endswith(new_state_dict, 'mlp.linear_fc1._extra_state')
+ mlp_fc2_extra_state_keys = _get_keys_endswith(new_state_dict, 'mlp.linear_fc2._extra_state')
+
+ for key in mlp_lm_weight_keys:
+ params = new_state_dict.pop(key)
+ new_key = key.replace('mlp.linear_fc1.layer_norm_weight', 'pre_mlp_layernorm.weight')
+ new_state_dict[new_key] = params
+
+ for key in mlp_lm_bias_keys:
+ params = new_state_dict.pop(key)
+ new_key = key.replace('mlp.linear_fc1.layer_norm_bias', 'pre_mlp_layernorm.bias')
+ new_state_dict[new_key] = params
+
+ for mlp_weight_key in mlp_fc1_weight_keys:
+ router_key = mlp_weight_key.replace('mlp.linear_fc1.weight', 'mlp.router.weight')
+ new_state_dict[router_key] = moe_state_dict[router_key].data.data.clone()
+
+ use_te_grouped_gemm = 'decoder.layers.0.mlp.experts.linear_fc1.weight0' in moe_state_dict
+
+ if mlp.config.moe_grouped_gemm and use_te_grouped_gemm:
+ for mlp_weight_key in mlp_fc1_weight_keys:
+ weight_tensor = new_state_dict.pop(mlp_weight_key)
+ for expert_i in range(mlp.num_local_experts):
+ new_key = mlp_weight_key.replace(
+ 'mlp.linear_fc1.weight', f'mlp.experts.linear_fc1.weight{expert_i}'
+ )
+ new_state_dict[new_key] = weight_tensor.clone()
+
+ for mlp_weight_key in mlp_fc2_weight_keys:
+ weight_tensor = new_state_dict.pop(mlp_weight_key)
+ for expert_i in range(mlp.num_local_experts):
+ new_key = mlp_weight_key.replace(
+ 'mlp.linear_fc2.weight', f'mlp.experts.linear_fc2.weight{expert_i}'
+ )
+ new_state_dict[new_key] = weight_tensor.clone()
+
+ for extra_state_key in mlp_fc1_extra_state_keys:
+ new_state_dict.pop(extra_state_key)
+ new_key = extra_state_key.replace(
+ 'mlp.linear_fc1._extra_state', 'mlp.experts.linear_fc1._extra_state'
+ )
+ new_state_dict[new_key] = None
+
+ for extra_state_key in mlp_fc2_extra_state_keys:
+ new_state_dict.pop(extra_state_key)
+ new_key = extra_state_key.replace(
+ 'mlp.linear_fc2._extra_state', 'mlp.experts.linear_fc2._extra_state'
+ )
+ new_state_dict[new_key] = None
+
+ elif mlp.config.moe_grouped_gemm:
+ for mlp_weight_key in mlp_fc1_weight_keys:
+ weight_tensor = new_state_dict.pop(mlp_weight_key)
+ shape = weight_tensor.shape
+ weight_tensor = weight_tensor.repeat(mlp.num_local_experts, 1, 1)
+ weight_tensor = weight_tensor.permute(0, 2, 1).reshape(
+ shape[1], mlp.num_local_experts * shape[0]
+ )
+ new_key = mlp_weight_key.replace('mlp.linear_fc1.weight', 'mlp.experts.weight1')
+ new_state_dict[new_key] = weight_tensor
+
+ for mlp_weight_key in mlp_fc2_weight_keys:
+ weight_tensor = new_state_dict.pop(mlp_weight_key)
+ shape = weight_tensor.shape
+ weight_tensor = weight_tensor.repeat(mlp.num_local_experts, 1, 1)
+ weight_tensor = weight_tensor.permute(0, 2, 1).reshape(
+ mlp.num_local_experts * shape[1], shape[0]
+ )
+ new_key = mlp_weight_key.replace('mlp.linear_fc2.weight', 'mlp.experts.weight2')
+ new_state_dict[new_key] = weight_tensor
+
+ else:
+
+ def covert_to_experts(keys):
+ for key in keys:
+ params = new_state_dict.pop(key)
+ new_key_format_str = key.replace('mlp', 'mlp.experts.local_experts.{}')
+ for expert_i in range(mlp.num_local_experts):
+ new_key = new_key_format_str.format(expert_i)
+ if hasattr(params, 'clone'):
+ new_state_dict[new_key] = params.clone()
+ else:
+ # set extra_state to None for now
+ new_state_dict[new_key] = None
+
+ covert_to_experts(mlp_fc1_weight_keys)
+ covert_to_experts(mlp_fc2_weight_keys)
+ covert_to_experts(mlp_fc1_bias_keys)
+ covert_to_experts(mlp_fc2_bias_keys)
+ covert_to_experts(mlp_fc1_extra_state_keys)
+ covert_to_experts(mlp_fc2_extra_state_keys)
+
+ return new_state_dict
+
+
+def upcycle_state_dict(moe_model, dense_model):
+ """
+ Convert a dense model's state_dict to a MoE model's state_dict.
+
+ This function facilitates the conversion of the state_dict from a dense model to
+ a MoE model, ensuring that the parameters are correctly mapped for each model.
+
+ Args:
+ moe_model (nn.Module): The MoE model, must be a model without FP16 and/or DDP wrapper.
+ dense_model (nn.Module): The dense model instance.
+
+ Returns:
+ dict: A dictionary containing the converted state_dict for the MoE model.
+ """
+
+ state_dict = {}
+ if len(moe_model) == 1:
+ assert len(dense_model) == 1
+ state_dict['model'] = _covert_to_moe_state_dict(dense_model[0].state_dict(), moe_model[0])
+ else:
+ assert len(moe_model) == len(dense_model)
+ for i in range(len(moe_model)):
+ mpu.set_virtual_pipeline_model_parallel_rank(i)
+ state_dict['model%d' % i] = _covert_to_moe_state_dict(
+ dense_model[i].state_dict(), moe_model[i]
+ )
+ return state_dict
+
+
+def load_and_upcycle_model(
+ load_dense_ckpt_func, moe_model, dense_model, strict=True, load_args=(), load_kwargs={}
+):
+ """
+ Load a dense model checkpoint and convert it to a MoE model.
+
+ This function loads a checkpoint for a dense model and converts it to the MoE model format,
+ allowing for the integration of the dense model's parameters into the MoE architecture.
+
+ Args:
+ load_dense_ckpt_func (callable): The function to load the dense model checkpoint.
+ moe_model (nn.Module): The MoE model instance.
+ dense_model (nn.Module): The dense model instance.
+ strict (bool): Whether to strictly load the state dictionary (default is True).
+ load_args (tuple): Positional arguments to pass to the loading function.
+ load_kwargs (dict): Keyword arguments to pass to the loading function.
+ """
+
+ iteration, num_floating_point_operations_so_far = load_dense_ckpt_func(
+ *load_args, **load_kwargs
+ )
+ state_dict = upcycle_state_dict(moe_model, dense_model)
+
+ if len(moe_model) == 1:
+ moe_model[0].load_state_dict(state_dict['model'], strict=strict)
+ else:
+ for i in range(len(moe_model)):
+ mpu.set_virtual_pipeline_model_parallel_rank(i)
+ moe_model[i].load_state_dict(state_dict['model%d' % i], strict=strict)
+
+ return iteration, num_floating_point_operations_so_far
diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py
new file mode 100644
index 0000000000..d637e2b448
--- /dev/null
+++ b/megatron/core/transformer/multi_latent_attention.py
@@ -0,0 +1,375 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+
+import math
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+
+from megatron.core import parallel_state
+from megatron.core.models.common.embeddings import (
+ YarnRotaryEmbedding,
+ _yarn_get_mscale,
+ apply_rotary_pos_emb,
+)
+from megatron.core.transformer.attention import Attention
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import MLATransformerConfig
+
+
+@dataclass
+class MLASelfAttentionSubmodules:
+ """Submodules for the MLA self-attention layer."""
+
+ linear_q_proj: Union[ModuleSpec, type] = None
+ linear_q_down_proj: Union[ModuleSpec, type] = None
+ linear_q_up_proj: Union[ModuleSpec, type] = None
+ linear_kv_down_proj: Union[ModuleSpec, type] = None
+ linear_kv_up_proj: Union[ModuleSpec, type] = None
+ core_attention: Union[ModuleSpec, type] = None
+ linear_proj: Union[ModuleSpec, type] = None
+ q_layernorm: Union[ModuleSpec, type] = None
+ kv_layernorm: Union[ModuleSpec, type] = None
+
+
+class MultiLatentAttention(Attention):
+ """Multi-Latent Attention layer abstract class.
+
+ This layer only contains common modules required for the "self attn" and
+ "cross attn" specializations.
+ """
+
+ def __init__(
+ self,
+ config: MLATransformerConfig,
+ submodules: Union[MLASelfAttentionSubmodules],
+ layer_number: int,
+ attn_mask_type: AttnMaskType,
+ attention_type: str,
+ ) -> None:
+ world_size = parallel_state.get_tensor_model_parallel_world_size()
+ assert (
+ world_size == 1
+ ), "MLA is not supported with Tensor Parallelism yet, \
+ use Expert Parallelism and Pipeline Parallelism for better performance."
+
+ super().__init__(
+ config=config,
+ submodules=submodules,
+ layer_number=layer_number,
+ attention_type=attention_type,
+ attn_mask_type=attn_mask_type,
+ )
+
+ self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads
+
+ self.q_head_dim = self.config.qk_head_dim + self.config.qk_pos_emb_head_dim
+
+ mscale = _yarn_get_mscale(self.config.rotary_scaling_factor, self.config.mscale)
+ self.softmax_scale = mscale * mscale / math.sqrt(self.q_head_dim)
+
+ self.rotary_pos_emb = YarnRotaryEmbedding(
+ self.config.qk_pos_emb_head_dim,
+ rotary_base=self.config.rotary_base,
+ scaling_factor=self.config.rotary_scaling_factor,
+ original_max_position_embeddings=self.config.max_position_embeddings,
+ beta_fast=self.config.beta_fast,
+ beta_slow=self.config.beta_slow,
+ mscale=self.config.mscale,
+ mscale_all_dim=self.config.mscale_all_dim,
+ )
+
+ self.core_attention = build_module(
+ submodules.core_attention,
+ config=self.config,
+ layer_number=self.layer_number,
+ attn_mask_type=self.attn_mask_type,
+ attention_type=self.attention_type,
+ softmax_scale=self.softmax_scale,
+ k_channels=self.q_head_dim,
+ v_channels=self.config.v_head_dim,
+ )
+
+ # Output.
+ self.linear_proj = build_module(
+ submodules.linear_proj,
+ self.query_projection_size,
+ self.config.hidden_size,
+ config=self.config,
+ init_method=self.config.output_layer_init_method,
+ bias=self.config.add_bias_linear,
+ input_is_parallel=True,
+ skip_bias_add=True,
+ is_expert=False,
+ tp_comm_buffer_name='proj',
+ )
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask,
+ key_value_states=None,
+ inference_params=None,
+ rotary_pos_emb=None,
+ packed_seq_params=None,
+ position_ids=None,
+ ):
+ assert rotary_pos_emb is None, "Rotary position embeddings should not be passed into MLA."
+
+ # hidden_states: [sq, b, h]
+
+ # =====================
+ # Query, Key, and Value
+ # =====================
+ # Get the query, key and value tensors based on the type of attention -
+ # self or cross attn.
+ # query: [96, 1, 16, 128], key:[96, 1, 16, 128], value:[96, 1, 16, 128]
+ query, key, value = self.get_query_key_value_tensors(
+ hidden_states,
+ key_value_states,
+ position_ids,
+ packed_seq_params,
+ inference_params=inference_params,
+ )
+
+ # ===================================================
+ # Adjust key, value for inference
+ # ===================================================
+ # rotary_pos_emb = None
+ key, value, _, attn_mask_type = self._adjust_key_value_for_inference(
+ inference_params, key, value, rotary_pos_emb=None
+ )
+
+ # ==================================
+ # core attention computation
+ # ==================================
+ # Need corresponding TE change
+ if self.checkpoint_core_attention and self.training:
+ core_attn_out = self._checkpointed_attention_forward(
+ query, key, value, attention_mask, packed_seq_params=packed_seq_params
+ )
+ else:
+ core_attn_out = self.core_attention(
+ query,
+ key,
+ value,
+ attention_mask,
+ packed_seq_params=packed_seq_params,
+ attn_mask_type=attn_mask_type,
+ )
+
+ if packed_seq_params is not None:
+ # reshape to same output shape as unpacked case
+ # (t, np, hn) -> (t, b=1, h=np*hn)
+ # t is the pack size = sum (sq_i)
+ # note that batch is a dummy dimension in the packed case
+ core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1)
+
+ # =================
+ # Output. [sq, b, h]
+ # =================
+ output, bias = self.linear_proj(core_attn_out)
+
+ return output, bias
+
+
+class MLASelfAttention(MultiLatentAttention):
+ """MLA Self-attention layer class
+
+ Self-attention layer takes input with size [s, b, h]
+ and returns output of the same size.
+ """
+
+ def __init__(
+ self,
+ config: MLATransformerConfig,
+ submodules: MLASelfAttentionSubmodules,
+ layer_number: int,
+ attn_mask_type=AttnMaskType.padding,
+ ):
+ super().__init__(
+ config=config,
+ submodules=submodules,
+ layer_number=layer_number,
+ attn_mask_type=attn_mask_type,
+ attention_type="self",
+ )
+
+ if self.config.q_lora_rank is None:
+ # Not projectiing query
+ self.linear_q_proj = build_module(
+ submodules.linear_q_proj,
+ self.config.hidden_size,
+ self.config.num_attention_heads * self.q_head_dim,
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=False,
+ skip_bias_add=False,
+ is_expert=False,
+ )
+
+ else:
+
+ self.linear_q_down_proj = build_module(
+ submodules.linear_q_down_proj,
+ self.config.hidden_size,
+ self.config.q_lora_rank,
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=False,
+ skip_bias_add=False,
+ is_expert=False,
+ )
+
+ self.linear_q_up_proj = build_module(
+ submodules.linear_q_up_proj,
+ self.config.q_lora_rank,
+ self.config.num_attention_heads * self.q_head_dim,
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=False,
+ skip_bias_add=False,
+ is_expert=False,
+ )
+
+ self.linear_kv_down_proj = build_module(
+ submodules.linear_kv_down_proj,
+ self.config.hidden_size,
+ self.config.kv_lora_rank + self.config.qk_pos_emb_head_dim,
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=False,
+ skip_bias_add=False,
+ is_expert=False,
+ )
+
+ self.linear_kv_up_proj = build_module(
+ submodules.linear_kv_up_proj,
+ self.config.kv_lora_rank,
+ self.config.num_attention_heads * (self.config.qk_head_dim + self.config.v_head_dim),
+ config=self.config,
+ init_method=self.config.init_method,
+ gather_output=False,
+ bias=False,
+ skip_bias_add=False,
+ is_expert=False,
+ )
+
+ if self.config.q_lora_rank is not None:
+ self.q_layernorm = build_module(
+ submodules.q_layernorm,
+ hidden_size=self.config.q_lora_rank,
+ config=self.config,
+ eps=self.config.layernorm_epsilon,
+ )
+
+ self.kv_layernorm = build_module(
+ submodules.kv_layernorm,
+ hidden_size=self.config.kv_lora_rank,
+ config=self.config,
+ eps=self.config.layernorm_epsilon,
+ )
+
+ def get_query_key_value_tensors(
+ self,
+ hidden_states,
+ key_value_states=None,
+ position_ids=None,
+ packed_seq_params=None,
+ inference_params=None,
+ ):
+ """
+ Derives `query`, `key` and `value` tensors from `hidden_states`.
+ """
+ # s = sequence length, b = batch size, h = hidden size, n = num attention heads
+ # Attention heads [s, b, n*h]
+ assert (
+ hidden_states.ndim == 3
+ ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D"
+ q_len, bsz, _ = hidden_states.size()
+
+ if self.config.q_lora_rank is not None:
+ q_compressed, _ = self.linear_q_down_proj(hidden_states)
+ q_compressed = self.q_layernorm(q_compressed)
+ q, _ = self.linear_q_up_proj(q_compressed)
+ else:
+ # hidden_states:[s, b, 2048], q: [s, b, n * 192]
+ q, _ = self.linear_q_proj(hidden_states)
+
+ # q: [s, b, n, 192]
+ q = q.view(q_len, bsz, self.num_attention_heads_per_partition, self.q_head_dim)
+
+ # q: [s, b, n, 128], q_pos_emb: [s, b, n, 64]
+ q_no_pe, q_pos_emb = torch.split(
+ q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1
+ )
+
+ # kv_combined: [s, b, 576]
+ kv_combined, _ = self.linear_kv_down_proj(hidden_states)
+
+ # kv_compressed:[s, b, 512], k_pos_emb: [s, b, 64]
+ kv_compressed, k_pos_emb = torch.split(
+ kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1
+ )
+
+ # kv: [s, b, 2048]
+ kv, _ = self.linear_kv_up_proj(self.kv_layernorm(kv_compressed))
+
+ # kv: [s, b, n, 256]
+ kv = kv.view(
+ q_len,
+ bsz,
+ self.num_attention_heads_per_partition,
+ self.config.qk_head_dim + self.config.v_head_dim,
+ )
+
+ # k_no_pe: [s, b, n, 128], value: [s, b, n, 128]
+ k_no_pe, value = torch.split(kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1)
+
+ # rotary_pos_emb:[s, b, 1, 64]
+ rotary_pos_emb = self.rotary_pos_emb(max_seq_len=self.config.max_position_embeddings)
+
+ if len(rotary_pos_emb) == 2:
+ mscale = rotary_pos_emb[1]
+ rotary_pos_emb = rotary_pos_emb[0]
+
+ if inference_params is not None:
+ # add offset to the sequence start for inference
+ sequence_start = inference_params.sequence_len_offset
+ sequence_end = sequence_start + q_len
+ rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end]
+
+ # [s, b, 64] -> [s, b, 1, 64]
+ k_pos_emb = torch.unsqueeze(k_pos_emb, 2)
+
+ if packed_seq_params is not None:
+ cu_seqlens_q = packed_seq_params.cu_seqlens_q
+ cu_seqlens_kv = packed_seq_params.cu_seqlens_kv
+ else:
+ cu_seqlens_q = cu_seqlens_kv = None
+
+ # q_pos_emb: [s, b, n, 64], k_pos_emb:[s, b, 1, 64]
+ q_pos_emb = apply_rotary_pos_emb(
+ q_pos_emb, rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q, mscale=mscale
+ )
+ k_pos_emb = apply_rotary_pos_emb(
+ k_pos_emb, rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv, mscale=mscale
+ )
+
+ # query: [s, b, n, 192]
+ query = torch.cat([q_no_pe, q_pos_emb], dim=-1)
+
+ # key: [s, b, n, 192]
+ key = torch.cat([k_no_pe, k_pos_emb], dim=-1)
+
+ query = query.contiguous()
+ key = key.contiguous()
+ value = value.contiguous()
+
+ return query, key, value
diff --git a/megatron/core/transformer/spec_utils.py b/megatron/core/transformer/spec_utils.py
new file mode 100644
index 0000000000..b3de854173
--- /dev/null
+++ b/megatron/core/transformer/spec_utils.py
@@ -0,0 +1,106 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+import types
+from dataclasses import dataclass, field
+from typing import Tuple, Union
+
+
+@dataclass
+class ModuleSpec:
+ """This is a Module Specification dataclass.
+
+ Specification defines the location of the module (to import dynamically)
+ or the imported module itself. It also defines the params that need to be
+ passed to initialize the module.
+
+ Args:
+ module (Union[Tuple, type]): A tuple describing the location of the
+ module class e.g. `(module.location, ModuleClass)` or the imported
+ module class itself e.g. `ModuleClass` (which is already imported
+ using `from module.location import ModuleClass`).
+ params (dict): A dictionary of params that need to be passed while init.
+
+ """
+
+ module: Union[Tuple, type]
+ params: dict = field(default_factory=lambda: {})
+ submodules: type = None
+
+
+def import_module(module_path: Tuple[str]):
+ """Import a named object from a module in the context of this function.
+
+ TODO: make this importer module more robust, at least make sure there
+ are no side effects of using this as is
+ """
+ base_path, name = module_path
+ try:
+ module = __import__(base_path, globals(), locals(), [name])
+ except ImportError as e:
+ print(f"couldn't import module due to {e}")
+ return None
+ return vars(module)[name]
+
+
+def get_module(spec_or_module: Union[ModuleSpec, type], **additional_kwargs):
+ # If a module clas is already provided return it as is
+ if isinstance(spec_or_module, (type, types.FunctionType)):
+ return spec_or_module
+
+ # If the module is provided instead of module path, then return it as is
+ if isinstance(spec_or_module.module, (type, types.FunctionType)):
+ return spec_or_module.module
+
+ # Otherwise, return the dynamically imported module from the module path
+ return import_module(spec_or_module.module)
+
+
+def build_module(spec_or_module: Union[ModuleSpec, type], *args, **kwargs):
+ # If the passed `spec_or_module` is
+ # a `Function`, then return it as it is
+ # NOTE: to support an already initialized module add the following condition
+ # `or isinstance(spec_or_module, torch.nn.Module)` to the following if check
+ if isinstance(spec_or_module, types.FunctionType):
+ return spec_or_module
+
+ # If the passed `spec_or_module` is actually a spec (instance of
+ # `ModuleSpec`) and it specifies a `Function` using its `module`
+ # field, return the `Function` as it is
+ if isinstance(spec_or_module, ModuleSpec) and isinstance(
+ spec_or_module.module, types.FunctionType
+ ):
+ return spec_or_module.module
+
+ # Check if a module class is provided as a spec or if the module path
+ # itself is a class
+ if isinstance(spec_or_module, type):
+ module = spec_or_module
+ elif hasattr(spec_or_module, "module") and isinstance(spec_or_module.module, type):
+ module = spec_or_module.module
+ else:
+ # Otherwise, dynamically import the module from the module path
+ module = import_module(spec_or_module.module)
+
+ # If the imported module is actually a `Function` return it as it is
+ if isinstance(module, types.FunctionType):
+ return module
+
+ # Finally return the initialized module with params from the spec as well
+ # as those passed as **kwargs from the code
+
+ # Add the `submodules` argument to the module init call if it exists in the
+ # spec.
+ if hasattr(spec_or_module, "submodules") and spec_or_module.submodules is not None:
+ kwargs["submodules"] = spec_or_module.submodules
+
+ try:
+ return module(
+ *args, **spec_or_module.params if hasattr(spec_or_module, "params") else {}, **kwargs
+ )
+ except Exception as e:
+ # improve the error message since we hide the module name in the line above
+ import sys
+
+ raise type(e)(f"{str(e)} when instantiating {module.__name__}").with_traceback(
+ sys.exc_info()[2]
+ )
diff --git a/megatron/core/transformer/torch_layer_norm.py b/megatron/core/transformer/torch_layer_norm.py
new file mode 100644
index 0000000000..11cf406f04
--- /dev/null
+++ b/megatron/core/transformer/torch_layer_norm.py
@@ -0,0 +1,44 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import warnings
+
+import torch
+
+from megatron.core.transformer import TransformerConfig
+
+
+class WrappedTorchLayerNorm(torch.nn.LayerNorm):
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ hidden_size: int,
+ eps: float = 1e-5,
+ persist_layer_norm: bool = False, ## TODO: unused arguments. See https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/issues/223
+ zero_centered_gamma: bool = False,
+ normalization: str = "LayerNorm", # included to match TE interface
+ ):
+ self.config = config
+ assert (
+ not self.config.layernorm_zero_centered_gamma
+ ), f"zero_centered_gamma not supported by torch LayerNorm"
+
+ assert (
+ self.config.normalization == "LayerNorm"
+ ), f'({self.config.normalization}) is not supported in by torch Layernorm'
+
+ assert (
+ not self.config.persist_layer_norm
+ ), f"persist_layer_norm not supported by torch LayerNorm"
+
+ assert (
+ not self.config.sequence_parallel
+ ), f"sequence parallel not supported by torch LayerNorm"
+
+ assert (
+ not self.config.memory_efficient_layer_norm
+ ), f"memory_efficient_layer_norm not supported by torch LayerNorm"
+
+ super().__init__(
+ normalized_shape=hidden_size, ## applied to last len(normalized_shape.size) dimensions
+ eps=eps,
+ )
diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py
new file mode 100755
index 0000000000..3a88f1ab22
--- /dev/null
+++ b/megatron/core/transformer/transformer_block.py
@@ -0,0 +1,597 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import torch
+from torch import Tensor
+
+from megatron.core import InferenceParams, parallel_state, tensor_parallel
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
+from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.transformer_layer import BaseTransformerLayer
+from megatron.core.transformer.utils import sharded_state_dict_default
+from megatron.core.utils import is_te_min_version, make_viewless_tensor
+
+try:
+ from megatron.core.extensions.transformer_engine import (
+ TEDelayedScaling,
+ TENorm,
+ get_cpu_offload_context,
+ te_checkpoint,
+ )
+
+ HAVE_TE = True
+ LayerNormImpl = TENorm
+except ImportError:
+ HAVE_TE = False
+ get_cpu_offload_context = None
+
+ try:
+ import apex # pylint: disable=unused-import
+
+ LayerNormImpl = FusedLayerNorm
+
+ except ImportError:
+ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm
+
+ LayerNormImpl = WrappedTorchLayerNorm
+
+
+def get_num_layers_to_build(config: TransformerConfig) -> int:
+ """
+ Determine the number of transformer layers to build for the current pipeline stage.
+ Args:
+ config (TransformerConfig): Configuration object containing transformer model parameters.
+
+ Returns:
+ int: The number of layers to be built for the current pipeline stage.
+ """
+ if config.first_pipeline_num_layers is not None or config.last_pipeline_num_layers is not None:
+ assert (
+ parallel_state.get_virtual_pipeline_model_parallel_world_size() is None
+ ), "Uneven number of layer not compatible with interleaved pipeline schedule"
+
+ # Number of layers to distribute over rest of pipeline stages
+ layers_to_distribute = config.num_layers
+ # Number of pipeline stages left for distributing transformer layers
+ pipeline_stages_left = parallel_state.get_pipeline_model_parallel_world_size()
+
+ if config.first_pipeline_num_layers is not None:
+ layers_to_distribute -= config.first_pipeline_num_layers
+ pipeline_stages_left -= 1
+ if parallel_state.is_pipeline_first_stage():
+ return config.first_pipeline_num_layers
+
+ if config.last_pipeline_num_layers is not None:
+ layers_to_distribute -= config.last_pipeline_num_layers
+ pipeline_stages_left -= 1
+ if parallel_state.is_pipeline_last_stage():
+ return config.last_pipeline_num_layers
+
+ assert (
+ layers_to_distribute % pipeline_stages_left == 0
+ ), "With uneven pipelineing the left over layers must be divisible by left over stages"
+ num_layers_per_pipeline_rank = layers_to_distribute // pipeline_stages_left
+ else:
+ pipeline_ranks = config.pipeline_model_parallel_size
+ num_layers_per_pipeline_rank = config.num_layers // pipeline_ranks
+
+ if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None:
+ # Interleaved pipeline parallelism:
+ # Number of layers in each model chunk is the number of layers in the stage,
+ # divided by the number of model chunks in a stage.
+ # With 8 layers, 2 stages, and 4 model chunks, we want an assignment of
+ # layers to stages like (each list is a model chunk):
+ # Stage 0: [0] [2] [4] [6]
+ # Stage 1: [1] [3] [5] [7]
+ # With 8 layers, 2 stages, and 2 virtual stages, we want an assignment of
+ # layers to stages like (each list is a model chunk):
+ # Stage 0: [0, 1] [4, 5]
+ # Stage 1: [2, 3] [6, 7]
+
+ vp_size = parallel_state.get_virtual_pipeline_model_parallel_world_size()
+
+ num_layers_per_virtual_rank = num_layers_per_pipeline_rank // vp_size
+
+ num_layers_to_build = num_layers_per_virtual_rank
+
+ else:
+ # Non-interleaved pipeline parallelism:
+ # Each stage gets a contiguous set of layers.
+
+ num_layers_to_build = num_layers_per_pipeline_rank
+
+ return num_layers_to_build
+
+
+@dataclass
+class TransformerBlockSubmodules:
+ """
+ Dataclass for specifying the submodules of a transformer block.
+
+ This class defines the structure for configuring the layers and normalization
+ within a transformer block, allowing for flexible and customizable architecture designs.
+
+ Args:
+ layer_specs (List[ModuleSpec], optional): A list of module specifications for
+ the layers within the transformer block. Each specification typically
+ defines a complete transformer layer (e.g., self-attention, feed-forward network).
+ layer_norm (Optional[Union[ModuleSpec, torch.nn.Module]], optional): Specification
+ or instance of the layer normalization to be applied.
+ """
+
+ layer_specs: List[ModuleSpec] = None
+ layer_norm: Optional[Union[ModuleSpec, torch.nn.Module]] = None
+
+
+def _get_block_submodules(
+ config: TransformerConfig, spec: Union[TransformerBlockSubmodules, ModuleSpec]
+) -> TransformerBlockSubmodules:
+ """
+ Retrieve or construct TransformerBlockSubmodules based on the provided specification.
+
+ Args:
+ config (TransformerConfig): Configuration object for the transformer model.
+ spec (Union[TransformerBlockSubmodules, ModuleSpec]): Specification for the
+ transformer block submodules. Can be either a TransformerBlockSubmodules
+ instance or a ModuleSpec.
+
+ Returns:
+ TransformerBlockSubmodules: The submodules for the transformer block.
+ """
+
+ # Transformer block submodules.
+ if isinstance(spec, TransformerBlockSubmodules):
+ return spec
+
+ # ModuleSpec here is generally assumed to be for a transformer layer that
+ # is implemented in `transformer_layer.py` or if it subclasses
+ # `BaseTransformerLayer` from the `transformer_layer.py` file.
+ elif isinstance(spec, ModuleSpec):
+ if issubclass(spec.module, TransformerBlock):
+ return spec.submodules
+ elif issubclass(spec.module, BaseTransformerLayer):
+ num_layers = get_num_layers_to_build(config)
+ return TransformerBlockSubmodules(
+ layer_specs=[spec] * num_layers, layer_norm=LayerNormImpl
+ )
+ else:
+ raise Exception(f"specialize for {spec.module.__name__}.")
+ else:
+ raise Exception(f"specialize for {type(spec).__name__}.")
+
+
+class TransformerBlock(MegatronModule):
+ """Transformer class."""
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ spec: Union[TransformerBlockSubmodules, ModuleSpec],
+ post_layer_norm: bool = True,
+ pre_process: bool = True,
+ post_process: bool = True,
+ ):
+ super().__init__(config=config)
+
+ self.submodules = _get_block_submodules(config, spec)
+ self.post_layer_norm = post_layer_norm
+ self.pre_process = pre_process
+ self.post_process = post_process
+ # Dictionary to store CUDA graphs. Number of items in the dictionary = len(self.layers).
+ # Item `i` in the dictionary is a list of `N` CUDA graphs for layer 'i' where N is the
+ # number of microbatches. Multiple CUDA graphs per layer is required to support
+ # pipelining which requires running FWD graph of multiple microbatches before BWD graph.
+ # To enable CUDA graph, this dictionary should be populated in the model training script
+ # with the graphs returned by make_graphed_callables API before the first trainng step.
+ self.cuda_graphs = {}
+ self.current_microbatch = -1
+
+ # required for pipeline parallel schedules
+ self.input_tensor = None
+
+ self.checkpoint_core_attention = self.config.recompute_granularity == 'selective'
+
+ if get_cpu_offload_context is not None:
+ (self.offload_context, self.group_prefetch_offload_commit_async) = (
+ get_cpu_offload_context(
+ self.config.cpu_offloading,
+ self.config.cpu_offloading_num_layers,
+ self.config.num_layers,
+ self.config.cpu_offloading_activations,
+ self.config.cpu_offloading_weights,
+ )
+ )
+ self.config._cpu_offloading_context = (
+ self.offload_context if self.config.cpu_offloading else None
+ )
+ else:
+ assert (
+ self.config.cpu_offloading is False
+ ), "CPU Offloading is enabled when TE is not present"
+
+ self.offload_context, self.group_prefetch_offload_commit_async = nullcontext(), None
+ self.config._cpu_offloading_context = None
+
+ self._build_layers()
+ self.num_layers_per_pipeline_rank = len(self.layers)
+ self.tp_only_amax_red = config.tp_only_amax_red
+
+ def _build_layers(self):
+ # Transformer layers.
+ # @jcasper can we improve how we deal with layer_number?
+ # currently it's only used in CoreAttention?
+ # if self.apply_query_key_layer_scaling:
+ # coeff = self.layer_number
+ # self.norm_factor *= coeff
+ def build_layer(layer_spec, layer_number):
+ return build_module(layer_spec, config=self.config, layer_number=layer_number)
+
+ # offset is implicit in TransformerLayer
+ self.layers = torch.nn.ModuleList(
+ [
+ build_layer(layer_spec, i + 1)
+ for i, layer_spec in enumerate(self.submodules.layer_specs)
+ ]
+ )
+
+ # @TODO: add back standalone_embedding_stage (see issue #293)
+ # In pipeline parallelism, we want to add this LN only to the last stage of the pipeline
+ # self.post_process and self.post_layer_norm guide this behavior
+ if self.submodules.layer_norm and self.post_process and self.post_layer_norm:
+ self.final_layernorm = build_module(
+ self.submodules.layer_norm,
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon,
+ )
+ else:
+ self.final_layernorm = None # Either this or nn.Identity
+
+ def _get_layer(self, layer_number: int):
+ return self.layers[layer_number]
+
+ def _checkpointed_forward(
+ self,
+ hidden_states: Tensor,
+ attention_mask: Tensor,
+ context: Tensor,
+ context_mask: Tensor,
+ rotary_pos_emb: Tensor,
+ packed_seq_params: PackedSeqParams,
+ ):
+ """Forward method with activation checkpointing."""
+
+ def custom(start: int, end: int):
+ def custom_forward(
+ hidden_states, attention_mask, context, context_mask, rotary_pos_emb
+ ):
+ for index in range(start, end):
+ layer = self._get_layer(index)
+ hidden_states, context = layer(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ context=context,
+ context_mask=context_mask,
+ rotary_pos_emb=rotary_pos_emb,
+ inference_params=None,
+ packed_seq_params=packed_seq_params,
+ )
+ return hidden_states, context
+
+ return custom_forward
+
+ def checkpoint_handler(forward_func):
+ if self.config.fp8:
+ return te_checkpoint(
+ forward_func,
+ self.config.distribute_saved_activations,
+ tensor_parallel.random.get_cuda_rng_tracker,
+ parallel_state.get_tensor_model_parallel_group(),
+ hidden_states,
+ attention_mask,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ )
+ else:
+ return tensor_parallel.checkpoint(
+ forward_func,
+ self.config.distribute_saved_activations,
+ hidden_states,
+ attention_mask,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ )
+
+ if self.config.recompute_method == 'uniform':
+ # Uniformly divide the total number of Transformer layers and checkpoint
+ # the input activation of each divided chunk.
+ # A method to further reduce memory usage reducing checkpoints.
+ layer_idx = 0
+ while layer_idx < self.num_layers_per_pipeline_rank:
+ hidden_states, context = checkpoint_handler(
+ custom(layer_idx, layer_idx + self.config.recompute_num_layers)
+ )
+
+ layer_idx += self.config.recompute_num_layers
+
+ elif self.config.recompute_method == 'block':
+ # Checkpoint the input activation of only a set number of individual
+ # Transformer layers and skip the rest.
+ # A method fully use the device memory removing redundant re-computation.
+ recompute_skip_num_layers = 0
+ for layer_idx in range(self.num_layers_per_pipeline_rank):
+ # Skip recomputation when input grad computation is not needed.
+ # Need to have at least one input tensor with gradient computation
+ # for re-enterant autograd engine.
+ if self.config.fp8 and not hidden_states.requires_grad:
+ recompute_skip_num_layers += 1
+ if (
+ layer_idx >= recompute_skip_num_layers
+ and layer_idx < self.config.recompute_num_layers + recompute_skip_num_layers
+ ):
+ hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1))
+ else:
+ hidden_states, context = custom(layer_idx, layer_idx + 1)(
+ hidden_states, attention_mask, context, context_mask, rotary_pos_emb
+ )
+ else:
+ raise ValueError("Invalid activation recompute method.")
+
+ return hidden_states
+
+ def set_input_tensor(self, input_tensor: Tensor):
+ """Set input tensor to be used instead of forward()'s input.
+
+ When doing pipeline parallelism the input from the previous
+ stage comes from communication, not from the input, so the
+ model's forward_step_func won't have it. This function is thus
+ used by internal code to bypass the input provided by the
+ forward_step_func"""
+ self.input_tensor = input_tensor
+
+ def get_cuda_graph_optional_args(
+ self,
+ attention_mask: Tensor,
+ context: Tensor,
+ context_mask: Tensor,
+ rotary_pos_emb: Tensor,
+ inference_params: InferenceParams,
+ packed_seq_params: PackedSeqParams,
+ ):
+ """Get optional tensor arguments for CUDA graph."""
+
+ optional_inputs = {}
+ optional_inputs['is_first_microbatch'] = self.current_microbatch == 0
+ try:
+ import transformer_engine.pytorch as te # pylint: disable=unused-import
+
+ if is_te_min_version("1.10.0", check_equality=False):
+ assert not any(
+ [attention_mask, context, context_mask, rotary_pos_emb]
+ ), "Keyword Arguments not supported with CUDA graph."
+ else:
+ optional_inputs['attention_mask'] = attention_mask
+ optional_inputs['context'] = context
+ optional_inputs['context_mask'] = context_mask
+ optional_inputs['rotary_pos_emb'] = rotary_pos_emb
+ except ImportError:
+ raise RuntimeError("CUDAGraph requires TransformerEngine, but not installed")
+ return optional_inputs
+
+ def forward(
+ self,
+ hidden_states: Tensor,
+ attention_mask: Tensor,
+ context: Tensor = None,
+ context_mask: Tensor = None,
+ rotary_pos_emb: Tensor = None,
+ inference_params: InferenceParams = None,
+ packed_seq_params: PackedSeqParams = None,
+ ):
+ """
+ Perform the forward pass through the transformer block.
+
+ This method handles the core computation of the transformer, including
+ self-attention, optional cross-attention, and feed-forward operations.
+
+ Args:
+ hidden_states (Tensor): Input tensor of shape [s, b, h] where s is the
+ sequence length, b is the batch size, and h is the hidden size.
+ attention_mask (Tensor): Boolean tensor of shape [1, 1, s, s] for masking
+ self-attention.
+ context (Tensor, optional): Context tensor for cross-attention.
+ context_mask (Tensor, optional): Mask for cross-attention context
+ rotary_pos_emb (Tensor, optional): Rotary positional embeddings.
+ inference_params (InferenceParams, optional): Parameters for inference-time
+ optimizations.
+ packed_seq_params (PackedSeqParams, optional): Parameters for packed sequence
+ processing.
+
+ Returns:
+ Union[Tensor, Tuple[Tensor, Tensor]]: The output hidden states tensor of shape
+ [s, b, h], and optionally the updated context tensor if cross-attention is used.
+ """
+
+ if not self.pre_process:
+ # See set_input_tensor()
+ hidden_states = self.input_tensor
+
+ # Viewless tensor.
+ # - We only need to create a viewless tensor in the case of micro batch
+ # size (mbs) == 1, since in this case, 'hidden_states.transpose()'
+ # above creates a view tensor, and '.contiguous()' is a pass-through.
+ # For mbs >= 2, '.contiguous()' creates a new tensor, eliminating
+ # the need to make it viewless.
+ #
+ # However, we don't explicitly check mbs == 1 here because
+ # make_viewless_tensor() has negligible overhead when its input
+ # is already viewless.
+ #
+ # - For the 'else' case above, calling make_viewless_tensor() here is
+ # likely redundant, since p2p_communication.py (likely originator)
+ # already creates viewless tensors. That said, make_viewless_tensor()
+ # is called here to be future-proof and corner-case-proof.
+ hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
+
+ if self.config.sequence_parallel:
+ rng_context = tensor_parallel.get_cuda_rng_tracker().fork()
+ else:
+ rng_context = nullcontext()
+
+ if self.config.fp8:
+ import transformer_engine # To keep out TE dependency when not training in fp8
+
+ if self.config.fp8 == "e4m3":
+ fp8_format = transformer_engine.common.recipe.Format.E4M3
+ elif self.config.fp8 == "hybrid":
+ fp8_format = transformer_engine.common.recipe.Format.HYBRID
+ else:
+ raise ValueError("E4M3 and HYBRID are the only supported FP8 formats.")
+
+ fp8_recipe = TEDelayedScaling(
+ config=self.config,
+ fp8_format=fp8_format,
+ override_linear_precision=(False, False, not self.config.fp8_wgrad),
+ )
+ fp8_group = None
+ if parallel_state.model_parallel_is_initialized():
+ fp8_group = parallel_state.get_amax_reduction_group(
+ with_context_parallel=True, tp_only_amax_red=self.tp_only_amax_red
+ )
+ fp8_context = transformer_engine.pytorch.fp8_autocast(
+ enabled=True, fp8_recipe=fp8_recipe, fp8_group=fp8_group
+ )
+ else:
+ fp8_context = nullcontext()
+
+ with rng_context and fp8_context:
+ # Forward pass.
+ if self.config.recompute_granularity == 'full' and self.training:
+ hidden_states = self._checkpointed_forward(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ context=context,
+ context_mask=context_mask,
+ rotary_pos_emb=rotary_pos_emb,
+ packed_seq_params=packed_seq_params,
+ )
+ else:
+ for l_no, layer in enumerate(self.layers):
+ with self.offload_context:
+ layer.use_cudagraph = True
+ if (len(self.cuda_graphs) == 0) or (not self.training):
+ hidden_states, context = layer(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ context=context,
+ context_mask=context_mask,
+ rotary_pos_emb=rotary_pos_emb,
+ inference_params=inference_params,
+ packed_seq_params=packed_seq_params,
+ )
+ else:
+ # CUDA graph replay for layer `l_no` and microbatch
+ # `self.current_microbatch`. TransformerEngine versions>=1.10
+ # allow keyword arguments with CUDA graph. However, CUDA graph
+ # acccepts only Tensor inputs and Tensor outputs. Hence,
+ # `inference_params` and `packed_seq_params` are excluded from
+ # input list while output is limited to `hidden_states`.
+ cg_index = self.current_microbatch % len(self.cuda_graphs[l_no])
+ assert not any(
+ [inference_params, packed_seq_params]
+ ), "CUDA graph accepts only Tensor inputs."
+ optional_inputs = self.get_cuda_graph_optional_args(
+ attention_mask,
+ context,
+ context_mask,
+ rotary_pos_emb,
+ inference_params,
+ packed_seq_params,
+ )
+ hidden_states = self.cuda_graphs[l_no][cg_index](
+ hidden_states, **optional_inputs
+ )
+
+ if (
+ torch.is_grad_enabled()
+ and self.config.cpu_offloading
+ and self.group_prefetch_offload_commit_async is not None
+ ):
+ hidden_states = self.group_prefetch_offload_commit_async(hidden_states)
+
+ # Final layer norm.
+ if self.final_layernorm is not None:
+ hidden_states = self.final_layernorm(hidden_states)
+ # TENorm produces a "viewed" tensor. This will result in schedule.py's
+ # deallocate_output_tensor() throwing an error, so a viewless tensor is
+ # created to prevent this.
+ hidden_states = make_viewless_tensor(
+ inp=hidden_states, requires_grad=True, keep_graph=True
+ )
+
+ return hidden_states
+
+ def sharded_state_dict(
+ self, prefix: str = '', sharded_offsets: tuple = (), metadata: dict = None
+ ) -> ShardedStateDict:
+ """
+ Generate a sharded state dictionary for the transformer block.
+
+ Args:
+ prefix (str, optional): Prefix to be added to all keys in the state dict.
+ Defaults to an empty string.
+ sharded_offsets (tuple, optional): Tuple of sharding offsets.
+ metadata (dict, optional): Additional metadata for sharding.
+ Can specify if layers are non-homogeneous. Defaults to None.
+
+ Returns:
+ ShardedStateDict: A dictionary containing the sharded state of the model.
+ """
+ assert not sharded_offsets, "Unexpected sharded offsets"
+ non_homogeneous_layers = metadata is not None and metadata.get(
+ 'non_homogeneous_layers', False
+ )
+ sharded_state_dict = {}
+
+ layer_prefix = f'{prefix}layers.'
+ num_layers = self.config.num_layers
+ for layer in self.layers:
+ offset = layer._get_layer_offset()
+
+ global_layer_offset = layer.layer_number - 1 # self.layer_number starts at 1
+ state_dict_prefix = f'{layer_prefix}{global_layer_offset - offset}.' # module list index in TransformerBlock # pylint: disable=line-too-long
+ if non_homogeneous_layers:
+ sharded_prefix = f'{layer_prefix}{global_layer_offset}.'
+ sharded_pp_offset = []
+ else:
+ sharded_prefix = layer_prefix
+ sharded_pp_offset = [
+ (0, global_layer_offset, num_layers)
+ ] # PP sharding offset for ShardedTensors
+ layer_sharded_state_dict = layer.sharded_state_dict(
+ state_dict_prefix, sharded_pp_offset, metadata
+ )
+ replace_prefix_for_sharding(layer_sharded_state_dict, state_dict_prefix, sharded_prefix)
+
+ sharded_state_dict.update(layer_sharded_state_dict)
+
+ # Add modules other than self.layers
+ for name, module in self.named_children():
+ if not module is self.layers:
+ sharded_state_dict.update(
+ sharded_state_dict_default(
+ module, f'{prefix}{name}.', sharded_offsets, metadata
+ )
+ )
+
+ return sharded_state_dict
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
new file mode 100644
index 0000000000..a63171686a
--- /dev/null
+++ b/megatron/core/transformer/transformer_config.py
@@ -0,0 +1,566 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Callable, Optional, Tuple
+
+import torch.nn.functional as F
+
+from ..model_parallel_config import ModelParallelConfig
+from ..utils import get_te_version, init_method_normal, is_te_min_version, scaled_init_method_normal
+
+
+@dataclass
+class TransformerConfig(ModelParallelConfig):
+ """Configuration object for megatron-core transformers.
+
+ The initialization function has an argument for each parameter,
+ including those in ModelParallelConfig.
+ """
+
+ ####################
+ # model architecture
+ ####################
+ num_layers: int = 0
+ """Number of transformer layers in a transformer block."""
+
+ first_pipeline_num_layers: int = None
+ """Number of transformer layers on first pipeline stage.
+ None implies equal layer division across PP ranks."""
+
+ last_pipeline_num_layers: int = None
+ """Number of transformer layers on last pipeline stage.
+ None implies equal layer division across PP ranks."""
+
+ hidden_size: int = 0
+ """Transformer hidden size."""
+
+ num_attention_heads: int = 0
+ """Number of transformer attention heads."""
+
+ num_query_groups: int = None
+ """Number of query groups for group query attention. If None, normal attention is used."""
+
+ ffn_hidden_size: int = None
+ """Transformer Feed-Forward Network hidden size. This is set to 4*hidden_size
+ if not provided."""
+
+ kv_channels: int = None
+ """Projection weights dimension in multi-head attention. This is set to hidden_size //
+ num_attention_heads if not provided."""
+
+ hidden_dropout: float = 0.1
+ """Dropout probability for transformer hidden state."""
+
+ attention_dropout: float = 0.1
+ """Post attention dropout probability."""
+
+ fp32_residual_connection: bool = False
+ """If true, move residual connections to fp32."""
+
+ # @jcasper should we keep this option?
+ apply_residual_connection_post_layernorm: bool = False
+ """If True, uses the original BERT residule connection ordering."""
+
+ layernorm_epsilon: float = 1e-5
+ """Epsilon value for any LayerNorm operations."""
+
+ layernorm_zero_centered_gamma: bool = False
+ """If set to True, the LayerNorm is adjusted to center the gamma values around 0. This improves
+ numerical stability."""
+
+ add_bias_linear: bool = True
+ """Include a bias term in all linear layers (QKV projections, after core attention, and two in
+ MLP layer)."""
+
+ add_qkv_bias: bool = False
+ """Add a bias term only for QKV projections."""
+
+ gated_linear_unit: bool = False
+ """Use a gated linear unit for the first linear layer in the MLP."""
+
+ activation_func: Callable = F.gelu
+ """Activation function to use for the non-linearity in the MLP."""
+
+ activation_func_fp8_input_store: bool = False
+ """Store the input of MLP activation function in FP8 for backprop to save memory.
+ The stored input is casted back to the original precision before backprop compuatation."""
+
+ num_moe_experts: int = None
+ """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None
+ for no MoE."""
+
+ rotary_interleaved: bool = False
+ """True is rotate pairs of even and odd dimensions (RoFormer style), False is rotate pairs of
+ first half and second half (LLaMa style). Default to False."""
+
+ window_size: Optional[Tuple[int, int]] = None
+ """If not None, then will use sliding window attention. The size of the window is specified by
+ the numbers inside the tuple; -1 is special value meaning "infinite window size"."""
+
+ normalization: bool = "LayerNorm"
+ """Which norm to use for normalization layers, valid options are `LayerNorm` and `RMSNorm`."""
+
+ qk_layernorm: bool = False
+ """Whether to apply LayerNorm to the query and key embeddings."""
+
+ test_mode: bool = False
+ """Whether to run real-time tests."""
+
+ calculate_per_token_loss: bool = False
+ """Whether cross entropy loss is calculated over the actual number of non-padded tokens in the
+ global batch, versus the default behavior of assuming all tokens are non-padded."""
+
+ multi_latent_attention: bool = False
+ """Whether to use multi-latent attention."""
+
+ ####################
+ # initialization
+ ####################
+ init_method: Callable = None
+ """Method to initialize weights. Note that bias is always set to zero. Should be a function that
+ takes a single Tensor and initializes it. If None, will be set to
+ megatron.core.utils.init_method_normal(init_method_std) which is torch nn init normal with
+ mean=0.0 and std=init_method_std."""
+
+ output_layer_init_method: Callable = None
+ """Method to initialize weights of the output layer of both attention and MLP blocks. If None,
+ will be set to megatron.core.utils.scaled_init_method_normal(init_method_std) which is torch nn
+ init normal with mean=0.0 and std=init_method_std / math.sqrt(2.0 * num_layers)."""
+
+ init_method_std: float = 0.02
+ """Standard deviation of the zero mean normal for the default initialization method, not used if
+ init_method and output_layer_init_method are provided."""
+
+ ####################
+ # mixed-precision
+ ####################
+ apply_query_key_layer_scaling: bool = False
+ """If true, scale Q * K^T by 1 / layer-number. This improve numeric stability when training with
+ fp16."""
+
+ attention_softmax_in_fp32: bool = True
+ """If True, run attention masking and softmax in fp32. This should be True if
+ apply_query_key_layer_scaling is True."""
+
+ ####################
+ # fusion
+ ####################
+ bias_activation_fusion: bool = False
+ """If True, fuses bias addition and the activation function when possible."""
+
+ masked_softmax_fusion: bool = False
+ """If True, uses softmax fusion."""
+
+ persist_layer_norm: bool = False
+ """If True, uses the persistent fused layer norm kernel. This kernel only supports a fixed set
+ of hidden sizes."""
+
+ memory_efficient_layer_norm: bool = False
+ """If True, and using local layers (not from TransformerEngine), tells Apex to use the memory
+ efficient fused LayerNorm kernel. Ignored if not using LayerNorm."""
+
+ bias_dropout_fusion: bool = False # TODO: this should be bias_dropout_add_fusion?
+ """If True, uses bias dropout fusion."""
+
+ apply_rope_fusion: bool = False
+ """If True, use fused RoPE kernel."""
+
+ ####################
+ # activation recomputation
+ ####################
+ recompute_granularity: str = None
+ """Determines which type of activation recompute to use. Megatron-core supports 'selective'
+ activation checkpointing where only the memory intensive part of attention is checkpointed.
+ These memory intensive activations are also less compute intensive which makes activation
+ checkpointing more efficient for LLMs (20B+). See Reducing Activation Recomputation in Large
+ Transformer Models (https://arxiv.org/abs/2205.05198) for more details. 'full' will checkpoint
+ the entire transformer layer. If None, no recompute is performed and all activations are saved.
+ If set, must be 'selective' or 'full'. 'selective' always uses all layers.
+ """
+
+ recompute_method: str = None
+ """Determines which transformer layers will be recomputed. uniform will uniformly divide the
+ total number of transformer layers in a transformer block and recompute the input activation of
+ each divided chunk at the specified granularity. block will recompute the input activations for
+ only a set number of transformer layers per pipeline stage. The rest of the layers in the
+ pipeline stage will not have any activations recomputed. If None, and recompute is enabled, all
+ layers will do recomputation. If set, must be 'uniform' or 'block'."""
+
+ recompute_num_layers: int = None
+ """When recompute_method is uniform, recompute_num_layers is the number of transformer layers in
+ each uniformly divided recompute unit. When recompute_method is block, recompute_num_layers is
+ the number of transformer layers to recompute within each pipeline stage. Must be None for
+ 'selective' activation checkpointing."""
+
+ distribute_saved_activations: bool = None
+ """If True, distribute recomputed activations across the model parallel group."""
+
+ ####################
+ # fp8 related
+ ####################
+ fp8: str = None
+ """If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined
+ choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8
+ activation and weight tensors and e5m2 for all FP8 output activation gradient tensors."""
+
+ fp8_margin: int = 0
+ """Margin for the scaling factor computation."""
+
+ fp8_interval: int = 1
+ """DEPRECATED from TransformerEngine v1.8.0. This flag is ignored.
+ Controls how often the scaling factor is recomputed.
+ """
+
+ fp8_amax_history_len: int = 1
+ """The length of the amax history window used for scaling factor computation."""
+
+ fp8_amax_compute_algo: str = "most_recent"
+ """Algorithm used for choosing the `amax` value for the scaling factor computation. There are 2
+ predefined choices: `max` chooses the largest `amax` in the history window, while `most_recent`
+ always chooses the most recently seen value.
+
+ """
+
+ fp8_wgrad: bool = True
+ """When set to False, override FP8 config options and do the wgrad computation
+ in higher precision."""
+
+ fp8_dot_product_attention: bool = False
+ """When set to True, use the FP8 implementation of Dot Product Attention."""
+
+ fp8_multi_head_attention: bool = False
+ """When set to True, use the FP8 implementation of Multi Head Attention."""
+
+ tp_only_amax_red: bool = False
+ """When set to True, reduce the FP8 AMAX only in the TP or TP-CP domain"""
+
+ ####################
+ # MoE related
+ ####################
+ moe_shared_expert_intermediate_size: int = None
+ """Shared expert total ffn hidden size.
+ It should be equal to 'num_shared_experts * ffn_size_of_each_shared_expert' if
+ there are multiple shared experts.
+ None means no shared expert."""
+
+ moe_shared_expert_overlap: bool = False
+ """Enable overlapping between shared expert computations and dispatcher communications.
+ Without this, the shared epxerts execute after the routed experts."""
+
+ moe_router_load_balancing_type: str = "aux_loss"
+ """Determines the load balancing strategy for the router. "aux_loss" corresponds to the load
+ balancing loss used in GShard and SwitchTransformer, "sinkhorn" corresponds to the balancing
+ algorithm used in S-BASE, and "none" implies no load balancing."""
+
+ moe_router_topk: int = 2
+ """Number of experts to route to for each token."""
+
+ moe_router_pre_softmax: bool = False
+ """Enable pre-softmax routing for MoE, which means softmax is before the top-k selection.
+ By default, softmax is done after top-k."""
+
+ moe_grouped_gemm: bool = False
+ """When there are multiple experts per rank, compress multiple local (potentially small) gemms
+ in a single kernel launch to improve the utilization and performance by leveraging the Grouped
+ GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm).
+ """
+
+ moe_aux_loss_coeff: float = 0 # 1e-2 would be a good start value for load balance loss.
+ """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended."""
+
+ moe_z_loss_coeff: float = None # 1e-3 would be a good start value for z-loss
+ """Scaling coefficient for the z-loss. A starting value of 1e-3 is recommended."""
+
+ moe_input_jitter_eps: float = None
+ """Add noise to the input tensor by applying jitter with a specified epsilon value."""
+
+ moe_token_dropping: bool = False # TODO: Support token dropping.
+ """This feature involves selectively dropping and padding tokens for each expert to achieve a
+ specified capacity, similar to GShard, Switch-Transformer, and DeepSpeed-MoE. Note that this is
+ currently unsupported so should remain False."""
+
+ moe_token_dispatcher_type: str = "allgather"
+ """The type of token dispatcher to use. The default is 'allgather'.
+ Options are 'allgather' and 'alltoall'."""
+
+ moe_per_layer_logging: bool = False
+ """Enable per-layer logging for MoE, currently supports auxiliary loss and z loss."""
+
+ moe_expert_capacity_factor: float = None
+ """moe_expert_capacity_factor (float): The capacity factor for each expert, None means no token
+ will be dropped. The default is None."""
+
+ moe_pad_expert_input_to_capacity: bool = False
+ """moe_pad_expert_input_to_capacity (bool): If True, pads the input for each expert to match
+ the expert capacity length, effective only after the moe_expert_capacity_factor is set. The
+ default setting is False."""
+
+ moe_token_drop_policy: str = 'probs'
+ """The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with
+ the lowest probabilities will be dropped. If "position", tokens at the end of each batch will
+ be dropped.
+ """
+
+ moe_layer_recompute: bool = False
+ """Memory optimization: checkpointing moe_layer to save actiavtion memory."""
+
+ ####################
+ # miscellaneous
+ ####################
+ clone_scatter_output_in_embedding: bool = True
+ """When set to True, clone the output of scatter_to_sequence_parallel_region in embedding layer
+ to facilitate garbage collection of input."""
+
+ disable_parameter_transpose_cache: bool = False
+ """When set to true, the parameter transposes are not cached for subsequent iterations."""
+
+ enable_cuda_graph: bool = False
+ """When set to true, TransformerLayer layers are swapped with a CUDA graphed version."""
+
+ external_cuda_graph: bool = False
+ """When set to true, TransformerLayer layers are swapped with user provided CUDA graphs."""
+
+ config_logger_dir: str = ""
+ """When non-empty, dumps entry-point configs to config_logger_dir"""
+
+ def __post_init__(self):
+ """Python dataclass method that is used to modify attributes after initialization.
+ See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more
+ details.
+ """
+ super().__post_init__()
+ if self.fp16 and self.bf16:
+ raise ValueError(
+ f'Only one of self.fp16: {self.fp16} and self.bf16 {self.bf16} should be True.'
+ )
+
+ if self.num_attention_heads % self.tensor_model_parallel_size != 0:
+ raise ValueError(
+ f"num_attention_heads ({self.num_attention_heads}) must be a multiple of "
+ f"tensor_model_parallel_size ({self.tensor_model_parallel_size})."
+ )
+
+ if self.ffn_hidden_size is None:
+ self.ffn_hidden_size = 4 * self.hidden_size
+
+ if self.kv_channels is None:
+ self.kv_channels = self.hidden_size // self.num_attention_heads
+
+ if self.num_query_groups is None:
+ self.num_query_groups = self.num_attention_heads
+
+ if self.num_query_groups % self.tensor_model_parallel_size != 0:
+ raise ValueError(
+ f"num_query_groups ({self.num_query_groups}) must be a multiple of "
+ f"tensor_model_parallel_size ({self.tensor_model_parallel_size})."
+ )
+
+ if self.apply_query_key_layer_scaling:
+ self.attention_softmax_in_fp32 = True
+
+ if self.expert_model_parallel_size > 1 and self.num_moe_experts is None:
+ raise ValueError('num_moe_experts must be non None to use expert-parallel.')
+
+ if self.num_moe_experts is not None and self.num_moe_experts <= 0:
+ raise ValueError('num_moe_experts must be non-negative.')
+
+ if self.moe_shared_expert_intermediate_size is not None:
+ if self.moe_shared_expert_intermediate_size <= 0:
+ raise ValueError(
+ f'moe_shared_expert_intermediate_size must be '
+ f'num_shared_experts * ffn_size_of_each_shared_expert, '
+ f'but got {self.moe_shared_expert_intermediate_size}'
+ )
+ if self.moe_shared_expert_overlap and self.moe_token_dispatcher_type not in [
+ "alltoall"
+ ]:
+ raise ValueError(
+ f'moe_shared_expert_overlap only works with alltoall token dispatcher.'
+ )
+
+ if self.moe_expert_capacity_factor is not None:
+ if self.moe_token_dispatcher_type not in ["alltoall", "alltoall_seq"]:
+ raise ValueError(
+ 'moe_expert_capacity_factor only works with alltoall token dispatcher'
+ )
+ if self.moe_expert_capacity_factor < 0:
+ self.moe_expert_capacity_factor = None
+ if self.moe_router_load_balancing_type not in ["aux_loss", "none"]:
+ raise ValueError(
+ 'moe_expert_capacity_factor only works with aux_loss or none load balancing'
+ )
+
+ if self.moe_pad_expert_input_to_capacity:
+ if self.moe_expert_capacity_factor is None:
+ raise ValueError(
+ 'moe_expert_capacity_factor must be set to use moe_pad_expert_input_to_capacity'
+ )
+
+ if self.cpu_offloading and (
+ self.cpu_offloading_num_layers < 0 or self.cpu_offloading_num_layers >= self.num_layers
+ ):
+ raise ValueError(
+ f'CPU offloading can be done only for layers less than {self.num_layers}'
+ )
+
+ if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
+ raise ValueError(
+ 'Currently there is no support for Pipeline parallelism with CPU offloading'
+ )
+
+ if self.cpu_offloading and self.recompute_granularity is not None:
+ raise ValueError(
+ 'CPU offloading does not work when activation recomputation is enabled'
+ )
+
+ if self.recompute_granularity is not None:
+ if self.recompute_granularity not in ['full', 'selective']:
+ raise ValueError(
+ f'When using recompute_granuarlity: {self.recompute_granularity} must be "full"'
+ 'or "selective".'
+ )
+
+ if self.recompute_method is not None:
+ if self.recompute_method not in ['block', 'uniform']:
+ raise ValueError(
+ f'recompute_method: {self.recompute_method} must be "block" or "uniform".'
+ )
+ elif self.recompute_granularity != 'selective':
+ raise ValueError(
+ f'Using recompute_granularity: {self.recompute_granularity} so '
+ 'recompute_method must be "block" or "uniform"'
+ )
+
+ if self.recompute_granularity != 'selective' and self.recompute_num_layers is None:
+ raise ValueError(
+ f'When using recompute_granularity: {self.recompute_granularity} '
+ 'recompute_num_layers must be between '
+ '1 and num_layers_per_pipeline_rank: '
+ f'{self.num_layers // self.pipeline_model_parallel_size}'
+ )
+ elif (
+ self.recompute_granularity == 'selective' and self.recompute_num_layers is not None
+ ):
+ raise ValueError(
+ f'When using recompute_granularity: {self.recompute_granularity} '
+ 'recompute_num_layers must be None.'
+ )
+
+ if self.distribute_saved_activations and self.sequence_parallel:
+ raise ValueError(
+ f'distribute_saved_activations: {self.distribute_saved_activations} must be '
+ f'false when sequence parallel is enabled: {self.sequence_parallel}'
+ )
+
+ if self.virtual_pipeline_model_parallel_size is not None:
+ if not self.num_layers % self.virtual_pipeline_model_parallel_size == 0:
+ raise ValueError(
+ f'num_layers: {self.num_layers} must be divisible by '
+ f'virtual_model_parallel_size {self.virtual_pipeline_model_parallel_size}'
+ )
+
+ if self.apply_query_key_layer_scaling:
+ self.attention_softmax_in_fp32 = True
+
+ if self.bias_activation_fusion:
+ if self.activation_func not in [F.gelu, F.silu]:
+ raise ValueError(
+ "When bias_activation_fusion is True, activation function should be either "
+ "gelu or swiglu"
+ )
+ if (
+ self.activation_func == F.gelu
+ and not self.gated_linear_unit
+ and not self.add_bias_linear
+ ):
+ raise ValueError(
+ "When bias_activation_fusion is True, gated_linear_unit is False, "
+ "and activation function is gelu, add_bias_linear must also be True."
+ )
+ if self.activation_func_fp8_input_store:
+ if self.activation_func != F.silu or not self.gated_linear_unit:
+ raise ValueError("Storing activation input in FP8 is supported only for SwiGLU.")
+ if self.apply_rope_fusion and self.rotary_interleaved:
+ raise ValueError('rotary_interleaved does not work with apply_rope_fusion.')
+
+ if self.init_method is None:
+ self.init_method = init_method_normal(self.init_method_std)
+
+ if self.output_layer_init_method is None:
+ self.output_layer_init_method = scaled_init_method_normal(
+ self.init_method_std, self.num_layers
+ )
+
+ if self.moe_extended_tp:
+ if self.moe_token_dispatcher_type != 'allgather':
+ raise ValueError(
+ "Moe extended TP parallelism only applies to allgather based token dispatcher."
+ )
+ extended_tp_size = self.tensor_model_parallel_size * self.expert_model_parallel_size
+ if self.ffn_hidden_size % extended_tp_size != 0:
+ raise ValueError(
+ f'ffn_hidden_size: {self.ffn_hidden_size} must be divisible by '
+ f'extended_tp_size {extended_tp_size}'
+ )
+
+ if self.num_moe_experts and self.fp8:
+ # TE version below 1.7.0 will raise Error when handle zeros tokens for expert
+ if not is_te_min_version("1.7.0.dev0"):
+ raise ValueError(
+ "Only transformer-engine>=1.7.0 supports MoE FP8 training, "
+ f"but your version is {get_te_version()}."
+ )
+
+ if self.moe_grouped_gemm:
+ raise ValueError("Grouped GEMM of MoE not support fp8 for now.")
+
+
+@dataclass
+class MLATransformerConfig(TransformerConfig):
+ """Configuration object for megatron-core Multi-Latent Attention (MLA) transformers.
+
+ The initialization function has an argument for each parameter, including those in
+ ModelParallelConfig. Included YaRN RoPE parameters that is fused in MLA.
+ """
+
+ multi_latent_attention: bool = True
+ """Whether to use Multi-Latent Attention."""
+
+ q_lora_rank: int = 512
+ """Rank of Query tensor's low rank representation."""
+
+ kv_lora_rank: int = 512
+ """Rank of Key and Value tensors' low rank representation."""
+
+ qk_head_dim: int = 128
+ """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim"""
+
+ qk_pos_emb_head_dim: int = 64
+ """Dimension of the position embedding in the QK projection."""
+
+ v_head_dim: int = 128
+ """Dimension of the head in the V projection."""
+
+ rotary_base: float = 10000
+ """Rotary base for the rotary embeddings."""
+
+ rotary_scaling_factor: float = 40
+ """Rotary scaling factor for the rotary embeddings."""
+
+ normalization: str = "RMSNorm"
+ """Default normalization layer for MLA models is RMSNorm."""
+
+ max_position_embeddings: int = 163840
+ """Maximum position embeddings for the original model."""
+
+ beta_fast: float = 32
+ """Beta fast for YaRN RoPE."""
+
+ beta_slow: float = 1
+ """Beta slow for YaRN RoPE."""
+
+ mscale: float = 0.707
+ """Mscale for YaRN RoPE in Multi-Latent Attention."""
+
+ mscale_all_dim: float = 0.707
+ """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention."""
diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
new file mode 100644
index 0000000000..7f5f14944e
--- /dev/null
+++ b/megatron/core/transformer/transformer_layer.py
@@ -0,0 +1,377 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Union
+
+import torch
+
+from megatron.core import parallel_state
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.dist_checkpointing.utils import apply_prefix_mapping
+from megatron.core.transformer.cuda_graphs import CudaGraphManager
+from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import make_viewless_tensor
+
+
+@dataclass
+class TransformerLayerSubmodules:
+ """
+ Configuration class for specifying the submodules of a transformer layer.
+
+ This class defines the structure and default implementations for various
+ components of a transformer layer, allowing for flexible customization
+ of the layer's architecture.
+
+ Args:
+ input_layernorm (Union[ModuleSpec, type]): Specification for the input layer normalization.
+ self_attention (Union[ModuleSpec, type]): Specification for the self-attention mechanism.
+ self_attn_bda (Union[ModuleSpec, type]): Specification for the bias-dropout-add operation
+ after self-attention.
+ pre_cross_attn_layernorm (Union[ModuleSpec, type]): Specification for the layer
+ normalization before cross-attention.
+ cross_attention (Union[ModuleSpec, type]): Specification for the cross-attention mechanism.
+ cross_attn_bda (Union[ModuleSpec, type]): Specification for the bias-dropout-add operation
+ after cross-attention.
+ pre_mlp_layernorm (Union[ModuleSpec, type]): Specification for the layer normalization
+ before the MLP.
+ mlp (Union[ModuleSpec, type]): Specification for the MLP.
+ mlp_bda (Union[ModuleSpec, type]): Specification for the bias-dropout-add operation
+ after the MLP.
+ sharded_state_dict_keys_map (Dict[str, str]): Mapping for sharded tensor keys to be applied
+ in the `sharded_state_dict` method.
+ """
+
+ input_layernorm: Union[ModuleSpec, type] = IdentityOp
+ self_attention: Union[ModuleSpec, type] = IdentityOp
+ self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp
+
+ pre_cross_attn_layernorm: Union[ModuleSpec, type] = IdentityOp
+ cross_attention: Union[ModuleSpec, type] = IdentityOp
+ cross_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp
+
+ pre_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp
+ mlp: Union[ModuleSpec, type] = IdentityOp
+ mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp
+
+ # Mapping for sharded tensor keys to be applied in `sharded_state_dict` method
+ sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict)
+
+
+class BaseTransformerLayer(ABC):
+ """A common parent class for `TransformerLayer` like implementations.
+
+ A dummy class that is subclassed by similar `TransformerLayer`s e.g. the
+ `TransformerLayer` in this file and possibly other `TransformerLayer`
+ implementations that aim to use `TransformerBlock` as the base module.
+ The main purpose is to check if any layer (or module) provided in the spec
+ is a subclass of this class to allow fanning-out of that spec for all the
+ layers in the `TransformerBlock`. See `_get_block_submodules` method
+ implementation in `transformer_block.py` file for more details.
+ """
+
+ def __init__(self):
+ pass
+
+
+class TransformerLayer(MegatronModule, BaseTransformerLayer):
+ """A single transformer layer.
+
+ Transformer layer takes input with size [s, b, h] and returns an
+ output of the same size.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: TransformerLayerSubmodules,
+ layer_number: int = 1,
+ hidden_dropout: float = None,
+ ):
+ super().__init__(config=config)
+
+ if config.enable_cuda_graph and self.training:
+ assert (
+ not config.cpu_offloading and config.recompute_granularity is None
+ ), "Cudagraphs not supported"
+ self.cudagraph_manager = CudaGraphManager()
+
+ self.submodules_config = submodules
+ self.layer_number = layer_number + self._get_layer_offset()
+ self.hidden_dropout = config.hidden_dropout if hidden_dropout is None else hidden_dropout
+
+ # [Module 1: Input Layernorm] Optional Layernorm on the input data
+ # TODO: add pytorch only layernorm
+ self.input_layernorm = build_module(
+ submodules.input_layernorm,
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon,
+ )
+
+ # [Module 2: SelfAttention]
+ self.self_attention = build_module(
+ submodules.self_attention, config=self.config, layer_number=layer_number
+ )
+
+ # [Module 3: BiasDropoutFusion]
+ self.self_attn_bda = build_module(submodules.self_attn_bda)
+
+ # [Module 4: Post SelfAttention] Optional Layernorm after self-attn
+ self.pre_cross_attn_layernorm = build_module(
+ submodules.pre_cross_attn_layernorm,
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon,
+ )
+
+ # [Module 5: CrossAttention]
+ self.cross_attention = build_module(
+ submodules.cross_attention, config=self.config, layer_number=layer_number
+ )
+
+ # [Module 6: BiasDropoutFusion]
+ self.cross_attn_bda = build_module(submodules.cross_attn_bda, config=self.config)
+
+ # [Module 7: Pre MLP] Optional Layernorm before MLP
+ self.pre_mlp_layernorm = build_module(
+ submodules.pre_mlp_layernorm,
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon,
+ )
+
+ # [Module 8: MLP block]
+ # TODO how to set the gpt_layer_spec.py when we have moe_frequency > 1,
+ # where MLP and MoE layer both appear alternately?
+ self.mlp = build_module(submodules.mlp, config=self.config)
+ if hasattr(self.mlp, 'set_layer_number'):
+ self.mlp.set_layer_number(self.layer_number)
+
+ # [Module 9: BiasDropoutFusion]
+ self.mlp_bda = build_module(submodules.mlp_bda)
+
+ # @jcasper how should we handle nvfuser?
+ # Set bias+dropout+add fusion grad_enable execution handler.
+ # TORCH_MAJOR = int(torch.__version__.split('.')[0])
+ # TORCH_MINOR = int(torch.__version__.split('.')[1])
+ # use_nvfuser = TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 10)
+ # self.bias_dropout_add_exec_handler = nullcontext if use_nvfuser else torch.enable_grad
+ self.bias_dropout_add_exec_handler = torch.enable_grad
+
+ def _get_layer_offset(self):
+ """Get the index number of this layer, given the level of pipelining."""
+ pipeline_rank = parallel_state.get_pipeline_model_parallel_rank()
+
+ num_layers_per_pipeline_rank = (
+ self.config.num_layers // self.config.pipeline_model_parallel_size
+ )
+
+ if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None:
+ vp_rank = parallel_state.get_virtual_pipeline_model_parallel_rank()
+ vp_size = parallel_state.get_virtual_pipeline_model_parallel_world_size()
+
+ total_num_layers = self.config.num_layers
+ num_layers_per_virtual_rank = num_layers_per_pipeline_rank // vp_size
+ total_virtual_chunks = total_num_layers // vp_size
+ offset = vp_rank * total_virtual_chunks + (pipeline_rank * num_layers_per_virtual_rank)
+
+ else:
+ # Each stage gets a contiguous set of layers.
+ if parallel_state.get_pipeline_model_parallel_world_size() > 1:
+ if (
+ self.config.first_pipeline_num_layers is not None
+ or self.config.last_pipeline_num_layers is not None
+ ):
+ # Calculate number of pipelines for distributing layers
+ middle_pipeline_stages = parallel_state.get_pipeline_model_parallel_world_size()
+ middle_pipeline_stages -= sum(
+ [
+ 1 if x is not None else 0
+ for x in (
+ self.config.first_pipeline_num_layers,
+ self.config.last_pipeline_num_layers,
+ )
+ ]
+ )
+
+ # Calculate layers to distribute
+ first_pipeline_offset = (
+ 0
+ if self.config.first_pipeline_num_layers is None
+ else self.config.first_pipeline_num_layers
+ )
+ last_pipeline_offset = (
+ 0
+ if self.config.first_pipeline_num_layers is None
+ else self.config.last_pipeline_num_layers
+ )
+
+ middle_num_layers = (
+ self.config.num_layers - first_pipeline_offset - last_pipeline_offset
+ )
+
+ if middle_pipeline_stages > 0:
+ num_layers_per_pipeline_rank = middle_num_layers // middle_pipeline_stages
+ else:
+ num_layers_per_pipeline_rank = 0
+
+ middle_pipeline_rank = (
+ pipeline_rank
+ if self.config.first_pipeline_num_layers is None
+ else pipeline_rank - 1
+ )
+
+ if pipeline_rank == 0:
+ offset = 0
+ else:
+ offset = (
+ middle_pipeline_rank * num_layers_per_pipeline_rank
+ ) + first_pipeline_offset
+ else:
+ offset = pipeline_rank * num_layers_per_pipeline_rank
+ else:
+ offset = 0
+
+ return offset
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ context=None,
+ context_mask=None,
+ rotary_pos_emb=None,
+ inference_params=None,
+ packed_seq_params=None,
+ ):
+ """
+ Perform a forward pass through the transformer layer.
+
+ This method implements the core computation of a transformer layer, including
+ self-attention, cross-attention (if applicable), and feed-forward operations.
+
+ Args:
+ hidden_states (Tensor): Input tensor of shape [s, b, h] where s is sequence length,
+ b is batch size, and h is hidden size.
+ attention_mask (Tensor): Mask tensor for self-attention.
+ context (Tensor, optional): Context tensor for cross-attention.
+ context_mask (Tensor, optional): Mask tensor for cross-attention.
+ rotary_pos_emb (Tensor, optional): Rotary positional embeddings.
+ inference_params (object, optional): Parameters for inference-time optimizations.
+ packed_seq_params (object, optional): Parameters for packed sequence processing.
+
+ Returns:
+ Tuple[Tensor, Tensor]: A tuple containing:
+ output (Tensor): Transformed hidden states of shape [s, b, h].
+ context (Tensor): Updated context tensor if cross-attention is used,
+ otherwise None.
+ """
+
+ # Residual connection.
+ residual = hidden_states
+
+ # Optional Input Layer norm
+ input_layernorm_output = self.input_layernorm(hidden_states)
+
+ # Self attention.
+ attention_output_with_bias = self.self_attention(
+ input_layernorm_output,
+ attention_mask=attention_mask,
+ inference_params=inference_params,
+ rotary_pos_emb=rotary_pos_emb,
+ packed_seq_params=packed_seq_params,
+ )
+
+ # TODO: could we move `bias_dropout_add_exec_handler` itself
+ # inside the module provided in the `bias_dropout_add_spec` module?
+ with self.bias_dropout_add_exec_handler():
+ hidden_states = self.self_attn_bda(self.training, self.config.bias_dropout_fusion)(
+ attention_output_with_bias, residual, self.hidden_dropout
+ )
+
+ # Residual connection.
+ residual = hidden_states
+
+ # Optional Layer norm after self-attention
+ pre_cross_attn_layernorm_output = self.pre_cross_attn_layernorm(hidden_states)
+
+ # Cross attention.
+ attention_output_with_bias = self.cross_attention(
+ pre_cross_attn_layernorm_output,
+ attention_mask=context_mask,
+ key_value_states=context,
+ inference_params=inference_params,
+ )
+
+ if isinstance(attention_output_with_bias, dict) and "context" in attention_output_with_bias:
+ context = attention_output_with_bias["context"]
+
+ # TODO: could we move `bias_dropout_add_exec_handler` itself
+ # inside the module provided in the `bias_dropout_add_spec` module?
+ with self.bias_dropout_add_exec_handler():
+ hidden_states = self.cross_attn_bda(self.training, self.config.bias_dropout_fusion)(
+ attention_output_with_bias, residual, self.hidden_dropout
+ )
+
+ # Residual connection.
+ residual = hidden_states
+
+ # Optional Layer norm post the cross-attention.
+ pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states)
+
+ # MLP.
+ mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
+
+ # TODO: could we move `bias_dropout_add_exec_handler` itself
+ # inside the module provided in the `bias_dropout_add_spec` module?
+ with self.bias_dropout_add_exec_handler():
+ hidden_states = self.mlp_bda(self.training, self.config.bias_dropout_fusion)(
+ mlp_output_with_bias, residual, self.hidden_dropout
+ )
+
+ # Jit compiled function creates 'view' tensor. This tensor
+ # potentially gets saved in the MPU checkpoint function context,
+ # which rejects view tensors. While making a viewless tensor here
+ # won't result in memory savings (like the data loader, or
+ # p2p_communication), it serves to document the origin of this
+ # 'view' tensor.
+ output = make_viewless_tensor(
+ inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True
+ )
+
+ # CUDA graph requires returned values to be Tensors
+ if self.config.external_cuda_graph and self.training:
+ return output
+ return output, context
+
+ def sharded_state_dict(
+ self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None
+ ) -> ShardedStateDict:
+ """
+ Generate a sharded state dictionary for the transformer layer.
+
+ Args:
+ prefix (str, optional): Prefix to be added to all keys in the state dict.
+ sharded_offsets (tuple, optional): Tuple of sharding offsets.
+ metadata (Optional[dict], optional): Additional metadata for sharding.
+
+ Returns:
+ ShardedStateDict: A dictionary containing the sharded state of the transformer layer.
+ """
+ sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+ prefixed_map = {
+ f'{prefix}{k}': f'{prefix}{v}'
+ for k, v in self.submodules_config.sharded_state_dict_keys_map.items()
+ }
+ if prefixed_map:
+ apply_prefix_mapping(sharded_state_dict, prefixed_map)
+ return sharded_state_dict
+
+ def __call__(self, *args, **kwargs):
+ if hasattr(self, 'cudagraph_manager'):
+ return self.cudagraph_manager(self, args, kwargs)
+ return super(MegatronModule, self).__call__(*args, **kwargs)
diff --git a/megatron/core/transformer/utils.py b/megatron/core/transformer/utils.py
new file mode 100644
index 0000000000..4781b68d2a
--- /dev/null
+++ b/megatron/core/transformer/utils.py
@@ -0,0 +1,188 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Utilities for transformer layers."""
+from functools import lru_cache
+from operator import itemgetter
+from typing import Any, Dict, Iterable, Iterator, Optional, Tuple, Union
+
+import torch
+
+from megatron.core import parallel_state
+from megatron.core.dist_checkpointing.mapping import ShardedObject, ShardedStateDict, StateDict
+from megatron.core.jit import jit_fuser
+from megatron.core.utils import (
+ make_sharded_tensor_for_checkpoint,
+ make_tp_sharded_tensor_for_checkpoint,
+)
+
+
+def get_linear_layer(rows, columns, init_method, perform_initialization=True):
+ """Simple linear layer with weight initialization."""
+ layer = torch.nn.Linear(rows, columns)
+ if perform_initialization: # Take from modelparallel config
+ init_method(layer.weight)
+ with torch.no_grad():
+ layer.bias.zero_()
+ return layer
+
+
+@lru_cache(maxsize=32)
+def get_default_causal_mask(sq: int) -> torch.Tensor:
+ """Return the causal upper triangular mask for softmax input."""
+ return torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool()
+
+
+def attention_mask_func(attention_scores, attention_mask):
+ attention_scores.masked_fill_(attention_mask, -10000.0)
+ return attention_scores
+
+
+@jit_fuser
+def gelu_impl(x):
+ """OpenAI's gelu implementation."""
+ return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x)))
+
+
+def openai_gelu(x):
+ return gelu_impl(x)
+
+
+# This is actually Python equivalent of torch.nn.functional.gelu(), also with type hints for ONNX exporter
+@jit_fuser
+def erf_gelu(x):
+ return (
+ x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype) + torch.ones_like(x).to(dtype=x.dtype))
+ )
+
+
+def make_sharded_tensors_for_checkpoint(
+ state_dict: StateDict,
+ prefix: str,
+ tensor_parallel_layers_axis_map: Optional[Dict[str, int]] = None,
+ sharded_offsets: Iterable[Tuple[int, int, int]] = (),
+ extra_state_suffix: str = '_extra_state',
+):
+ """Wraps tensors from transformer layers with ShardedTensor or ShardedObject.
+
+ For a given `state_dict`, wraps:
+ - all _extra_states with ShardedObject
+ - all tensors specified in tensor_parallel_layers_axis_map with TP and DP sharded ShardedTensor
+ - other values with DP sharded ShardedTensor
+
+ Args:
+ state_dict (StateDict): state_dict to convert
+ prefix (str): prefix appended to keys in final state dict
+ tensor_parallel_layers_axis_map (Dict[str, int], optional): dict mapping layer
+ names to the axis for TP sharding
+ sharded_offsets (Iterable[Tuple[int, int, int]], optional): sharding already
+ applied (e.g. PP related), passed along to ShardedTensor
+ extra_state_suffix (str, default = '_extra_state'): layers with this
+ suffix will be wrapped with ShardedObject instead of ShardedTensor.
+
+ """
+
+ if tensor_parallel_layers_axis_map is None:
+ tensor_parallel_layers_axis_map = {}
+
+ sharded_state_dict = {}
+ for layer_name in state_dict.keys():
+ tensor = state_dict[layer_name]
+ layer_key = f'{prefix}{layer_name}'
+
+ if layer_name.endswith(extra_state_suffix):
+ sharded_state_dict[layer_key] = make_sharded_object_for_checkpoint(
+ tensor, layer_key, sharded_offsets
+ )
+
+ elif layer_name in tensor_parallel_layers_axis_map:
+ tp_axis = tensor_parallel_layers_axis_map[layer_name]
+ sharded_state_dict[layer_key] = make_tp_sharded_tensor_for_checkpoint(
+ tensor, layer_key, tp_axis, prepend_offsets=sharded_offsets
+ )
+
+ else:
+ sharded_state_dict[layer_key] = make_sharded_tensor_for_checkpoint(
+ tensor, layer_key, prepend_offsets=sharded_offsets
+ )
+
+ return sharded_state_dict
+
+
+def make_sharded_object_for_checkpoint(
+ obj: Any,
+ key: str,
+ sharded_offsets: Iterable[Tuple[int, int, int]] = (),
+ replica_id: Union[None, int, Tuple[int, ...]] = None,
+ **kwargs,
+):
+ """Helper for instantiating a non-sharded ShardedObject (replicated across TP and DP group).
+
+ Args:
+ obj (object): any object to be sharded
+ key (str): unique identifier of the object
+ sharded_offsets (Iterable[Tuple[int, int, int]]): offsets normally
+ prepended to ShardedTensors, will be used as global offsets for
+ ShardedObject
+ replica_id (Union[None, int, Tuple[int, ...]]): replica id
+ """
+ if replica_id is None:
+ replica_id = (
+ 0,
+ parallel_state.get_tensor_model_parallel_rank(),
+ parallel_state.get_data_parallel_rank(with_context_parallel=True),
+ )
+
+ return ShardedObject(key, obj, *_get_extra_state_offsets(sharded_offsets), replica_id, **kwargs)
+
+
+def _get_extra_state_offsets(
+ sharded_offsets: Iterable[Tuple[int, int, int]]
+) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
+ """Turns ShardedTensor offsets into offsets suitable for ShardedObject."""
+ if sharded_offsets:
+ sharded_offsets = sorted(sharded_offsets, key=itemgetter(0)) # sort by axis
+ axis, extra_state_offset, extra_state_shape = zip(*sharded_offsets)
+ assert list(axis) == list(
+ range(len(axis))
+ ), f'Expected contiguous axis for offsets: {sharded_offsets}'
+ else:
+ extra_state_shape = (1,)
+ extra_state_offset = (0,)
+ return extra_state_shape, extra_state_offset
+
+
+def sharded_state_dict_default(
+ module: torch.nn.Module,
+ prefix: str = '',
+ sharded_offsets: Tuple[Tuple[int, int, int]] = (),
+ metadata: Optional[dict] = None,
+) -> ShardedStateDict:
+ """Provides implementation for sharded_state_dict method for non-MegatronModules.
+
+ Tries to call `module.sharded_state_dict` when possible,
+ otherwise uses regular state dict and assumes tensors are replicated across TP and DP.
+
+ `keep_vars=True` is passed to module.state_dict so that optimizer states
+ can be sharded later on.
+
+ Args:
+ module (torch.nn.Module): module which sharded state dict we want to obtain
+ prefix (str): prefix for the state dict keys
+ sharded_offsets (Tuple[Tuple[int, int, int]], optional): sharding already
+ applied (e.g. PP related) by sup-modules. Passed along to ShardedTensor
+ metadata (dict, optional): metadata passed to module sharded_state_dict method
+
+ Returns:
+ dict: dictionary of state dict keys mapped to ShardedTensors
+ """
+
+ if hasattr(module, 'sharded_state_dict'):
+ module_sharded_sd = module.sharded_state_dict(
+ prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata
+ )
+ else:
+ module_sd = module.state_dict(prefix='', keep_vars=True)
+ module_sharded_sd = make_sharded_tensors_for_checkpoint(
+ module_sd, prefix, {}, sharded_offsets
+ )
+ return module_sharded_sd
diff --git a/megatron/core/utils.py b/megatron/core/utils.py
index 7214b0c271..f3910926ab 100644
--- a/megatron/core/utils.py
+++ b/megatron/core/utils.py
@@ -1,19 +1,63 @@
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
"""Utility functions used throughout Megatron core"""
-from functools import reduce
+import array
+import hashlib
+import logging
+import math
import operator
+import queue
+import socket
+import sys
+import threading
+import time
+import traceback
+from dataclasses import dataclass
+from datetime import datetime
+from functools import reduce
+from importlib.metadata import version
+from types import TracebackType
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
import torch
+from packaging.version import Version as PkgVersion
from megatron.core import parallel_state
+from megatron.core.dist_checkpointing.mapping import ShardedTensor
+
+logger = logging.getLogger(__name__)
+
+
+_te_version = None
+
+
+def get_te_version():
+ """Get TE version from __version__; if not available use pip's. Use caching."""
+
+ def get_te_version_str():
+ import transformer_engine as te
+
+ if hasattr(te, '__version__'):
+ return str(te.__version__)
+ else:
+ return version("transformer-engine")
+
+ global _te_version
+ if _te_version is None:
+ _te_version = PkgVersion(get_te_version_str())
+ return _te_version
+
+
+def is_te_min_version(version, check_equality=True):
+ """Check if minimum version of `transformer-engine` is installed."""
+ if check_equality:
+ return get_te_version() >= PkgVersion(version)
+ return get_te_version() > PkgVersion(version)
def ensure_divisibility(numerator, denominator):
"""Ensure that numerator is divisible by the denominator."""
- assert numerator % denominator == 0, "{} is not divisible by {}".format(
- numerator, denominator
- )
+ assert numerator % denominator == 0, "{} is not divisible by {}".format(numerator, denominator)
def divide(numerator, denominator):
@@ -22,22 +66,53 @@ def divide(numerator, denominator):
ensure_divisibility(numerator, denominator)
return numerator // denominator
-def get_attr_wrapped_model(model, attr):
- """Get an attribute from a wrapped model"""
+
+def get_attr_wrapped_model(model, attr, allow_none=True, return_model_obj=False):
+ """Get an attribute from a wrapped model.
+ If return_model_obj is true, return the object that has the 'attr' attribute;
+ otherwise, return the attribute directly."""
if isinstance(model, list):
raise RuntimeError("_get_attr_wrapped_model given a list of models")
- while not hasattr(model, attr):
+ if allow_none:
+
+ def condition(model, attr):
+ return not hasattr(model, attr)
+
+ else:
+
+ def condition(model, attr):
+ return getattr(model, attr, None) is None
+
+ while condition(model, attr):
if not hasattr(model, "module"):
raise RuntimeError(f"_get_attr_wrapped_model couldn't find attribute {attr}")
model = model.module
+
+ if return_model_obj:
+ return model
return getattr(model, attr)
+
def get_model_type(model):
+ """Returns model_type attribute"""
return get_attr_wrapped_model(model, 'model_type')
+def get_model_xattn(model):
+ """Returns whether the model has the xattn_needed attribute"""
+ try:
+ return get_attr_wrapped_model(model, 'xattn_needed')
+ except RuntimeError:
+ return False
+
+
+def get_model_config(model):
+ """Returns the config attribute, allowed to return None"""
+ return get_attr_wrapped_model(model, 'config', allow_none=False)
+
+
class GlobalMemoryBuffer:
"""Global buffer to avoid dynamic memory allocations.
Caller should ensure that buffers of the same name
@@ -47,60 +122,65 @@ def __init__(self):
self.buffer = {}
def get_tensor(self, tensor_shape, dtype, name):
+ """
+ Returns (potentially) a sub-tensor from the self.buffer for the given shape.
+ """
required_len = reduce(operator.mul, tensor_shape, 1)
- if self.buffer.get((name, dtype), None) is None or \
- self.buffer[(name, dtype)].numel() < required_len:
- self.buffer[(name, dtype)] = \
- torch.empty(required_len,
- dtype=dtype,
- device=torch.cuda.current_device(),
- requires_grad=False)
+ if (
+ self.buffer.get((name, dtype), None) is None
+ or self.buffer[(name, dtype)].numel() < required_len
+ ):
+ self.buffer[(name, dtype)] = torch.empty(
+ required_len, dtype=dtype, device=torch.cuda.current_device(), requires_grad=False
+ )
return self.buffer[(name, dtype)][0:required_len].view(*tensor_shape)
+
def _kernel_make_viewless_tensor(inp, requires_grad):
- '''Make a viewless tensor.
+ """Make a viewless tensor.
View tensors have the undesirable side-affect of retaining a reference
to the originally-viewed tensor, even after manually setting the '.data'
field. This method creates a new tensor that links to the old tensor's
data, without linking the viewed tensor, referenced via the '._base'
field.
- '''
- out = torch.empty(
- (1,),
- dtype = inp.dtype,
- device = inp.device,
- requires_grad = requires_grad,
- )
+ """
+ out = torch.empty((1,), dtype=inp.dtype, device=inp.device, requires_grad=requires_grad)
out.data = inp.data
return out
+
class MakeViewlessTensor(torch.autograd.Function):
- '''
+ """
Autograd function to make a viewless tensor.
This function should be used in cases where the computation graph needs
to be propagated, but we only want a viewless tensor (e.g.,
ParallelTransformer's hidden_states). Call this function by passing
'keep_graph = True' to 'make_viewless_tensor()'.
- '''
+ """
+
@staticmethod
def forward(ctx, inp, requires_grad):
+ """Runs the fwd pass of _kernel_make_viewless_tensor"""
return _kernel_make_viewless_tensor(inp, requires_grad)
+
@staticmethod
def backward(ctx, grad_output):
+ """No-op"""
return grad_output, None
+
def make_viewless_tensor(inp, requires_grad, keep_graph):
- '''
+ """
Entry-point for creating viewless tensors.
This method should be used, rather than calling 'MakeViewlessTensor'
or '_kernel_make_viewless_tensor' directly. This method acts as a
switch for determining if an autograd function or a regular method
should be used to create the tensor.
- '''
+ """
# return tensor as-is, if not a 'view'
if inp._base is None:
@@ -112,11 +192,12 @@ def make_viewless_tensor(inp, requires_grad, keep_graph):
else:
return _kernel_make_viewless_tensor(inp, requires_grad)
-def assert_viewless_tensor(tensor, extra_msg = None):
- '''Assert that a tensor is not a view (i.e., its '._base' field is
- not set).'''
+
+def assert_viewless_tensor(tensor, extra_msg=None):
+ """Assert that a tensor is not a view (i.e., its '._base' field is
+ not set)."""
if isinstance(tensor, list):
- [ assert_viewless_tensor(t) for t in tensor ]
+ [assert_viewless_tensor(t) for t in tensor]
return tensor
if not isinstance(tensor, torch.Tensor):
return tensor
@@ -127,11 +208,1100 @@ def assert_viewless_tensor(tensor, extra_msg = None):
) % extra_msg
return tensor
+
def safely_set_viewless_tensor_data(tensor, new_data_tensor):
- '''Safely set tensor's '.data' field.
+ """Safely set tensor's '.data' field.
Check first that the tensor is viewless (i.e., '._base' not set). If not,
raise an exception.
- '''
- assert_viewless_tensor(tensor, extra_msg = "FYI, tensor._base has shape %s, and new_data_tensor has shape %s." % ("--" if tensor._base is None else tensor._base.shape, new_data_tensor.shape))
+ """
+ assert_viewless_tensor(
+ tensor,
+ extra_msg="FYI, tensor._base has shape %s, and new_data_tensor has shape %s."
+ % ("--" if tensor._base is None else tensor._base.shape, new_data_tensor.shape),
+ )
tensor.data = new_data_tensor
+
+
+def init_method_normal(sigma):
+ """Init method based on N(0, sigma)."""
+
+ def init_(tensor):
+ return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)
+
+ return init_
+
+
+def scaled_init_method_normal(sigma, num_layers):
+ """Init method based on N(0, sigma/sqrt(2*num_layers)."""
+ std = sigma / math.sqrt(2.0 * num_layers)
+
+ def init_(tensor):
+ return torch.nn.init.normal_(tensor, mean=0.0, std=std)
+
+ return init_
+
+
+def log_single_rank(logger: logging.Logger, *args: Any, rank: int = 0, **kwargs: Any):
+ """If torch distributed is initialized, log only on rank
+
+ Args:
+ logger (logging.Logger): The logger to write the logs
+
+ args (Tuple[Any]): All logging.Logger.log positional arguments
+
+ rank (int, optional): The rank to write on. Defaults to 0.
+
+ kwargs (Dict[str, Any]): All logging.Logger.log keyword arguments
+ """
+ if torch.distributed.is_initialized():
+ if torch.distributed.get_rank() == rank:
+ logger.log(*args, **kwargs)
+ else:
+ logger.log(*args, **kwargs)
+
+
+def log_on_each_pipeline_stage(logger: logging.Logger, *args: Any, **kwargs: Any):
+ """Log on first rank in each pipeline stage
+
+ Args:
+ logger (logging.Logger): The logger to write the logs
+
+ args (Tuple[Any]): All logging.Logger.log positional arguments
+
+ kwargs (Dict[str, Any]): All logging.Logger.log keyword arguments
+ """
+ assert torch.distributed.is_initialized()
+
+ if (
+ parallel_state.get_data_parallel_rank(with_context_parallel=True) == 0
+ and parallel_state.get_tensor_model_parallel_rank() == 0
+ ):
+ logger.log(*args, **kwargs)
+
+
+def check_param_hashes_across_dp_replicas(
+ model: List[torch.nn.Module], cross_check: bool = False
+) -> bool:
+ """Computes hashes of all parameters in model, all-gathers hashes across DP replicas,
+ and then checks for equality between the locally-computed hashes and those of other ranks.
+
+ NOTE: This function computes SHA-1 hashes on the CPU and thus needs to move all param
+ tensors from GPU to CPU first; as a result, this function is not intended to be called
+ very frequently in the main training loop.
+
+ Args:
+ model (List[torch.nn.Module]): List of model chunks whose parameter hashes need to
+ be checked.
+ cross_check (bool): If true, will check whether hashes match across all DP replicas.
+
+ Returns:
+ True if all param hashes match with corresponding hash on DP replica 0 or
+ across all replicas if cross_check is enabled, False otherwise.
+ """
+
+ # Compute per-parameter hashes on this rank.
+ params = []
+ local_param_hashes = []
+ for model_chunk_id, model_chunk in enumerate(model):
+ for param_name, param in model_chunk.named_parameters():
+ param_hash = torch.frombuffer(
+ array.array(
+ 'B', hashlib.sha1(param.data.to("cpu").float().numpy(force=True)).digest()
+ ),
+ dtype=torch.uint8,
+ )
+ params.append((model_chunk_id, param_name, param))
+ local_param_hashes.append(param_hash)
+ local_param_hashes = torch.stack(local_param_hashes)
+
+ # Collect per-parameter hashes across all ranks in DP group.
+ all_param_hashes = [
+ torch.zeros_like(local_param_hashes)
+ for _ in range(parallel_state.get_data_parallel_world_size())
+ ]
+ torch.distributed.all_gather(
+ all_param_hashes, local_param_hashes, group=parallel_state.get_data_parallel_group_gloo()
+ )
+
+ # Make sure local per-parameter hash matches DP rank 0.
+ param_hashes_match = torch.equal(local_param_hashes, all_param_hashes[0])
+ if not param_hashes_match:
+ for i, (model_chunk_id, param_name, param) in enumerate(params):
+ if not torch.equal(local_param_hashes[i], all_param_hashes[0][i]):
+ rank = torch.distributed.get_rank()
+ logger.info(
+ f"[Rank {rank}] Hash not matching for {param_name} in model chunk"
+ f"{model_chunk_id}"
+ )
+ if cross_check:
+ # Make sure all ranks have the same hash.
+ return all(map(lambda x: torch.equal(local_param_hashes, x), all_param_hashes))
+ else:
+ return param_hashes_match
+
+
+def make_tp_sharded_tensor_for_checkpoint(
+ tensor, key, tp_axis=0, replica_id=None, prepend_offsets=(), **kwargs
+):
+ """Helper for instantiating a ShardedTensor where the `tp_axis` dimension
+ is sharded across TP group.
+
+ Optionally, can provide offsets which prepend new dimensions to the tensor.
+ """
+
+ prepend_axis_num = len(prepend_offsets)
+
+ if replica_id is None:
+ replica_id = (0, 0, parallel_state.get_data_parallel_rank(with_context_parallel=True))
+
+ return ShardedTensor.from_rank_offsets(
+ key,
+ tensor,
+ *prepend_offsets,
+ (
+ tp_axis + prepend_axis_num,
+ parallel_state.get_tensor_model_parallel_rank(),
+ parallel_state.get_tensor_model_parallel_world_size(),
+ ),
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ **kwargs,
+ )
+
+
+def make_sharded_tensor_for_checkpoint(tensor, key, prepend_offsets=(), replica_id=None, **kwargs):
+ """Helper for instantiating a non-sharded ShardedTensor (replicated across TP and DP group).
+
+ Optionally, can provide offsets which prepend new dimensions to the tensor.
+ """
+
+ prepend_axis_num = len(prepend_offsets)
+
+ if replica_id is None:
+ replica_id = (
+ 0,
+ parallel_state.get_tensor_model_parallel_rank(),
+ parallel_state.get_data_parallel_rank(with_context_parallel=True),
+ )
+
+ return ShardedTensor.from_rank_offsets(
+ key,
+ tensor,
+ *prepend_offsets,
+ replica_id=replica_id,
+ prepend_axis_num=prepend_axis_num,
+ **kwargs,
+ )
+
+
+def prepare_input_tensors_for_wgrad_compute(grad_output, all_gathered_input):
+ """Ensure grad_output is stored in a contiguous buffer."""
+ # Doing gather + slicing during the NeMo forward pass can make this tensor
+ # not be contiguous. PyTorch only checks if the tensor is contiguous, and only
+ # clones it if it's not contiguous:
+ # https://github.com/pytorch/pytorch/blob/c47cf9bc7f9e02f649ab4ed53fe4d35732c92ab6/torch/_refs/__init__.py#L2761
+ grad_output = grad_output.contiguous()
+ # Convert the tensor shapes to 2D for execution compatibility
+ if grad_output.dim() == 3:
+ grad_output = grad_output.view(
+ grad_output.shape[0] * grad_output.shape[1], grad_output.shape[2]
+ )
+ all_gathered_input = all_gathered_input.view(
+ all_gathered_input.shape[0] * all_gathered_input.shape[1], all_gathered_input.shape[2]
+ )
+
+ return grad_output, all_gathered_input
+
+
+def drain_embedding_wgrad_compute(config, embedding_activation_buffer, grad_output_buffer, weight):
+ """Helper for performing embedding wgrad GEMM's during the pipeline drain phase, pipelines the
+ AllGather and GEMM's.
+
+ Should only be used when pipeline model parallelism and gradient accumulation
+ fusion are enabled.
+ """
+
+ assert len(embedding_activation_buffer) == len(
+ grad_output_buffer
+ ), "Length of activation and gradient buffers need to be equal!"
+
+ import fused_weight_gradient_mlp_cuda
+
+ from megatron.core.parallel_state import (
+ get_global_memory_buffer,
+ get_tensor_model_parallel_group,
+ get_tensor_model_parallel_world_size,
+ )
+
+ input = embedding_activation_buffer.pop(0)
+ world_size = get_tensor_model_parallel_world_size()
+ dim_size = list(input.size())
+ dim_size[0] = dim_size[0] * world_size
+
+ all_gathered_input = [None, None]
+ if config.sequence_parallel:
+ all_gather_buffer = get_global_memory_buffer().get_tensor(dim_size, input.dtype, "mpu_0")
+ handle = torch.distributed._all_gather_base(
+ all_gather_buffer, input, group=get_tensor_model_parallel_group(), async_op=False
+ )
+
+ all_gathered_input[0] = all_gather_buffer
+ all_gather_buffer = None
+ else:
+ all_gathered_input[0] = input
+
+ input = None
+
+ def wgrad_compute(all_gathered_input, grad_output, weight):
+
+ grad_output, all_gathered_input = prepare_input_tensors_for_wgrad_compute(
+ grad_output, all_gathered_input
+ )
+
+ if config.gradient_accumulation_fusion:
+ if weight.main_grad.dtype == torch.float32:
+ fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp32(
+ all_gathered_input, grad_output, weight.main_grad
+ )
+ elif weight.main_grad.dtype in (torch.float16, torch.bfloat16):
+ fused_weight_gradient_mlp_cuda.wgrad_gemm_accum_fp16(
+ all_gathered_input, grad_output, weight.main_grad
+ )
+ else:
+ raise RuntimeError("Unsupported gradient type for gradient accumulation fusion")
+
+ # We have all_gathered_input list acting as a double buffer here,
+ # since we are pipelining the AllGather and GEMM,one buffer all gathers
+ # the input while the other buffer reads from it for the GEMM. We use i
+ # and (i+1) for indexing to enable this double buffering.
+ for i in range(len(embedding_activation_buffer)):
+ input = embedding_activation_buffer.pop(0)
+ if config.sequence_parallel:
+ name = "mpu_" + str((i + 1) % 2)
+ all_gather_buffer = get_global_memory_buffer().get_tensor(dim_size, input.dtype, name)
+ handle = torch.distributed._all_gather_base(
+ all_gather_buffer, input, group=get_tensor_model_parallel_group(), async_op=True
+ )
+
+ all_gathered_input[(i + 1) % 2] = all_gather_buffer
+ all_gather_buffer = None
+ else:
+ all_gathered_input[(i + 1) % 2] = input
+
+ grad_output = grad_output_buffer.pop(0)
+ wgrad_compute(all_gathered_input[i % 2], grad_output, weight)
+ drain_idx = (i + 1) % 2
+ input, all_gathered_input[i % 2], grad_output = None, None, None
+
+ if config.sequence_parallel:
+ handle.wait()
+
+ grad_output = grad_output_buffer.pop(0)
+ wgrad_compute(all_gathered_input[drain_idx], grad_output, weight)
+ input, all_gathered_input[drain_idx], grad_output = None, None, None
+
+
+def local_multi_tensor_applier(op, noop_flag_buffer, tensor_lists, *args):
+ """Multi tensor op applier"""
+ return op(2048 * 32, noop_flag_buffer, tensor_lists, *args)
+
+
+# computes l2 norm for a list of contiguous tensors
+# works as a drop-in replacement for amp_C.multi_tensor_l2norm
+def local_multi_tensor_l2_norm(chunk_size, noop_flag, tensor_lists, per_tensor, *args):
+ """
+ Computes l2 norm for a list of contiguous tensors
+ works as a drop-in replacement for amp_C.multi_tensor_l2norm
+ """
+ l2 = [[(torch.norm(tensor)) for tensor in tensor_list] for tensor_list in tensor_lists]
+ l2_reduced = torch.norm(torch.tensor(l2))
+ l2_cuda = torch.tensor([float(l2_reduced)], dtype=torch.float, device='cuda')
+ return l2_cuda, None
+
+
+# works as a drop-in replacement for amp_C.multi_tensor_scale
+def local_multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale):
+ """Works as a drop-in replacement for amp_C.multi_tensor_scale."""
+ inputs, targets = tensor_lists[0], tensor_lists[1]
+ if inputs == targets:
+ for i in range(len(targets)):
+ # for parity with apex implementation
+ targets[i] *= scale
+ else:
+ for i in range(len(targets)):
+ targets[i] = inputs[i] * scale
+
+
+class _ValueWithRank:
+ """This is an internal class, not for use outside this module
+
+ Attributes:
+ _rank (int): rank for the value
+ _value (float) : the value it stores, eg elapsed time
+ _unit (str) : unit for the value
+ """
+
+ def __init__(self, value: float, rank: int, unit: str = "") -> None:
+ """Initializer
+
+ Args:
+ _value (float): the initial value with which it is inited
+ _rank (int): the rank number
+ _unit (str) : the unit of the value, eg ms or flops
+ """
+ self._rank = rank
+ self._value = value
+ self._unit = unit
+
+ def __lt__(self, other) -> bool:
+ """Check if value of self is smaller than other's value
+
+ Args:
+ other (_ValueWithRank): The other object to compare with
+
+ Returns:
+ bool: True if lhs._value of operand is less than rhs._value, else False
+ """
+ return self._value < other._value
+
+ def __gt__(self, other) -> bool:
+ """Check if value of self is larger than other's value
+
+ Args:
+ other (_ValueWithRank): The other object to compare with
+
+ Returns:
+ bool: True if lhs._value of operand is greater than rhs._value, else False
+ """
+ return self._value > other._value
+
+ def __call__(self) -> Tuple[float, int, str]:
+ """Returns the value, the rank, and unit as a Tuple
+
+ Returns:
+ Tuple[float, int, str]: value, rank, unit
+ """
+ return self._value, self._rank, self._unit
+
+ def __str__(self) -> str:
+ """String representation of the object
+
+ Returns:
+ str: strigified object
+ """
+
+ return f"{self._value:.2f}{self._unit}/{self._rank}"
+
+
+@dataclass
+class _StragglerData:
+ """This is an internal dataclass, not for use outside this module
+
+ Attributes:
+ min_elapsed (_ValueWithRank) min iteration time across all ranks
+ max_elapsed (_ValueWithRank) max iteration time across all ranks
+ min_btime (_ValueWithRank) min cpu time across all ranks
+ max_btime (_ValueWithRank) max cpu time across all ranks
+ min_temp (_ValueWithRank): min gpu temp across all ranks
+ max_temp (_ValueWithRank): max gpu temp across all ranks
+ min_power (_ValueWithRank) min gpu power across all ranks
+ max_power (_ValueWithRank) max gpu power across all ranks
+ min_util (_ValueWithRank): min gpu util across all ranks
+ max_util (_ValueWithRank): max gpu util across all ranks
+ min_clock (_ValueWithRank): min gpu clock across all ranks
+ max_clock (_ValueWithRank) max gpu clock across all ranks
+ aflops (List[_ValueWithRank]): sorted array of (_ValueWithRank)
+ """
+
+ # gemm time
+ min_elapsed = _ValueWithRank(sys.float_info.max, 0, "ms")
+ max_elapsed = _ValueWithRank(sys.float_info.min, 0, "ms")
+ # get_batch time
+ min_btime = _ValueWithRank(sys.float_info.max, 0, "us")
+ max_btime = _ValueWithRank(sys.float_info.min, 0, "us")
+ # temp
+ min_temp = _ValueWithRank(sys.float_info.max, 0, "C")
+ max_temp = _ValueWithRank(sys.float_info.min, 0, "C")
+ # power
+ min_power = _ValueWithRank(sys.float_info.max, 0, "W")
+ max_power = _ValueWithRank(sys.float_info.min, 0, "W")
+ # util
+ min_util = _ValueWithRank(sys.float_info.max, 0, "%")
+ max_util = _ValueWithRank(sys.float_info.min, 0, "%")
+ # clock
+ min_clock = _ValueWithRank(sys.float_info.max, 0, "MHz")
+ max_clock = _ValueWithRank(sys.float_info.min, 0, "MHz")
+ aflops: Union[List[_ValueWithRank], None] = None
+
+
+class StragglerDetector:
+ """Singleton Class implementing per rank Straggler Detector
+
+ It use cuda events to time operation of choice using the
+ start and stop methods which can be directly invoked using
+ the class instance or can be used like a python context.
+ After collection, a report() method is available to display
+ the collected metrics. It is only supported if CUDA is
+ available. megatron/core/README_STRAGGLER.md for more info
+
+ Note:
+ The instance and class attributes mentioned below are all
+ private to the class and has no use outside the class
+
+ Attributes:
+ _off (bool): current state of the toggle
+ start (FunctionType): start method
+ stop (FunctionType): stop method
+ world (int): world size
+ rank (int): rank for this instance
+ mmcnt (int): number of ranks to report
+ port (int): control port
+ amp (float): amplification factor for TFLOPs, default 3.0
+ toggle (bool): whether to start/stop detector collection
+ bdata (bool): when true, just collect get_batch
+ dev (int): cuda device
+ evt_q (LifoQueue): cuda event queue
+ start_gemm_ev (list[torch.cuda.Event]): cuda start event
+ stop_gemm_ev (list[torch.cuda.Event]): cuda stop event
+ start_data_ev (list[torch.cuda.Event]): cuda start event
+ stop_data_ev (list[torch.cuda.Event]): cuda stop event
+ start_gemm_tm (list[int]): start time (wallclock)
+ stop_gemm_tm (list[int]): stop time (wallclock)
+ start_data_tm (list[int]): start time for get_batch
+ stop_data_tm (list[int]): stop time for get_batch
+ sock (socket): the controller socket
+ ctrlr (Thread): the controller thread
+ """
+
+ _configured = False
+ """Indicates if the singleton instance is configured or not
+ """
+
+ def __new__(cls: Type["StragglerDetector"]) -> "StragglerDetector":
+ """Constructor
+ Creates an instance of the class if not created
+
+ Args:
+ cls (Type['StragglerDetector']): The class type
+
+ Returns:
+ StragglerDetector: the class instance
+ """
+
+ if not hasattr(cls, "_instance"):
+ cls._instance = super(StragglerDetector, cls).__new__(cls)
+ return cls._instance
+
+ def __init__(self) -> None:
+ """Initializer
+
+ The inital state of the StragglerDetector instance is disabled.
+ The enabled state is indicated using self._off member variable
+ and the proerty enabled.
+ """
+ self._off: bool = True
+ self.start = self.null_method
+ self.stop = self.null_method
+ self.world: int = 0
+ self.rank: int = 0
+ self.mmcnt: int = 1
+ self.port: int = 0
+ self.amp: float = 3.0
+ self.toggle: bool = False
+ self.bdata: bool = False
+ self.dev: Union[torch.device, int, None] = None
+ self.evt_q: Union[queue.LifoQueue, None] = None
+ self.start_gemm_ev: List[torch.cuda.Event] = []
+ self.stop_gemm_ev: List[torch.cuda.Event] = []
+ self.start_data_ev: List[torch.cuda.Event] = []
+ self.stop_data_ev: List[torch.cuda.Event] = []
+ self.start_gemm_tm: List[int] = []
+ self.stop_gemm_tm: List[int] = []
+ self.start_data_tm: List[int] = []
+ self.stop_data_tm: List[int] = []
+ self.sock: Union[socket.socket, None] = None
+ self.ctrlr: Union[threading.Thread, None] = None
+
+ def configure(
+ self,
+ world: int,
+ rank: int,
+ mmcnt: int = 1,
+ amp: float = 3.0,
+ port: int = 65535,
+ prefill: int = 1024,
+ enabled: bool = False,
+ ) -> None:
+ """This method is called to configure the Singleton instance
+
+ It should be called once per instantiation per process.
+
+ Note:
+ The constructor keeps the state of instance disabled
+ i.e no collection will happen even when start/stop methods are
+ called. Only when enabled is True (self._off is True), the
+ start/stop method pointers get assigned the real collection
+ methods, otherwise they are initialized with null_method
+
+ Args:
+ world (int): World Size
+ rank (int): The rank of this trainer
+ mmcnt (int, optional): Number of ranks to print for showing Min/Max Etpt.
+ Defaults to 1.
+ amp (float, optional): Set to 3.0 if we only use timers in fwd pass.
+ Defaults to 3.0.
+ port (int, optional): Control port, useful only for rank-0. Defaults to 65535.
+ prefill (int, optional): Howmany Events to pre-populate. Defaults to 1024.
+ enabled (bool, optional): Whether or not collection is enabled on startup.
+ Defaults to False.
+ """
+ if StragglerDetector._configured:
+ # don't throw
+ return
+ StragglerDetector._configured = True
+ self.bdata = False
+ self.start = self.null_method
+ self.stop = self.null_method
+ self._off = True
+ # No CUDA, No Support
+ if torch.cuda.is_available():
+ self._off = not enabled
+ self.world = world
+ self.rank = rank
+ self.mmcnt = mmcnt if mmcnt > 1 else 1
+ self.amp = amp
+ self.port = port
+ self.toggle = False
+ self.bdata = False
+ self.evt_q = queue.LifoQueue()
+ self.start_gemm_ev = []
+ self.stop_gemm_ev = []
+ self.start_data_ev = []
+ self.stop_data_ev = []
+ self.start_gemm_tm = []
+ self.stop_gemm_tm = []
+ self.start_data_tm = []
+ self.stop_data_tm = []
+ backend = torch.distributed.get_backend()
+ if backend == "nccl":
+ self.dev = torch.cuda.current_device()
+ else:
+ self.dev = torch.device("cpu")
+ # cache some events
+ for _ in range(prefill):
+ self.evt_q.put(torch.cuda.Event(enable_timing=True))
+ if self.rank == 0:
+ # Start the controller
+ self._controller()
+ if not self._off:
+ self.start = self.start_method
+ self.stop = self.stop_method
+
+ def reset(self) -> None:
+ """This method is called to reset the metrics state of the instance
+
+ It is generally called from within elapsed() after extracting per rank metrics.
+ """
+ if self._off:
+ return
+ # Pool them
+ if self.evt_q is not None:
+ _ = [self.evt_q.put(ev) for ev in self.start_gemm_ev]
+ _ = [self.evt_q.put(ev) for ev in self.stop_gemm_ev]
+ _ = [self.evt_q.put(ev) for ev in self.start_data_ev]
+ _ = [self.evt_q.put(ev) for ev in self.stop_data_ev]
+ self.start_gemm_ev = []
+ self.stop_gemm_ev = []
+ self.start_data_ev = []
+ self.stop_data_ev = []
+ # Use regular timers
+ self.start_gemm_tm = []
+ self.stop_gemm_tm = []
+ self.start_data_tm = []
+ self.stop_data_tm = []
+ self.bdata = False
+
+ def start_method(self) -> None:
+ """This method adds the start timers.
+
+ Both cuda event and perf_counter are added. If bdata is set to
+ true from __call__, this method skips inserting cuda
+ timer. This way it can be used to measure time spent on
+ CPU - generally useful for timing get_batch()
+ """
+ # Not reentrant
+ if self.evt_q is not None and self.evt_q.qsize() > 1:
+ sev = self.evt_q.get() # no try-catch
+ eev = self.evt_q.get() # no try-catch
+ else:
+ sev = torch.cuda.Event(enable_timing=True)
+ eev = torch.cuda.Event(enable_timing=True)
+ # First check if this start is for data
+ if self.bdata:
+ self.start_data_ev.append(sev)
+ self.stop_data_ev.append(eev)
+ self.start_data_tm.append(0)
+ self.stop_data_tm.append(0)
+ idx = len(self.stop_data_tm) - 1
+ self.start_data_tm[idx] = time.perf_counter_ns()
+ self.start_data_ev[idx].record()
+ self.bdata = False
+ return
+ self.start_gemm_ev.append(sev)
+ self.stop_gemm_ev.append(eev)
+ self.start_gemm_tm.append(0)
+ self.stop_gemm_tm.append(0)
+ idx = len(self.stop_gemm_tm) - 1
+ self.start_gemm_tm[idx] = time.perf_counter_ns()
+ self.start_gemm_ev[idx].record()
+
+ def stop_method(self) -> None:
+ """This method adds the stop timers.
+
+ Both cuda event and perf_counter are added. If bdata is set to
+ true from __call__, this method skips inserting cuda
+ timer. Also see start_method()
+ """
+ # Not reentrant
+ # First check if this stop is for data
+ idx = len(self.stop_data_tm) - 1
+ if idx >= 0 and self.stop_data_tm[idx] == 0:
+ self.stop_data_tm[idx] = time.perf_counter_ns()
+ self.stop_data_ev[idx].record()
+ return
+ idx = len(self.stop_gemm_tm) - 1
+ if idx >= 0 and self.stop_gemm_tm[idx] == 0:
+ self.stop_gemm_tm[idx] = time.perf_counter_ns()
+ self.stop_gemm_ev[idx].record()
+
+ def elapsed(self) -> Tuple[float, float, int, int, int, int]:
+ """This method is called from report(), or can be called directly
+
+ It is called to collect all the elapsed time since last reset().
+ It finally calls reset()
+
+ Returns:
+ Tuple[float, float, int, int, int, int]: see below for returns
+ delta : time spent in kernel
+ batch_delta : time spent in get_batch
+ temp : observed gpu temp
+ power : observed gpu power
+ util : observed gpu utilization
+ clock : observed gpu clock
+ """
+ if self._off:
+ # match with return below
+ return 0, 0, 0, 0, 0, 0
+ ls_ev = len(self.start_gemm_ev)
+ le_ev = len(self.stop_gemm_ev)
+ ls_bs = len(self.start_data_ev)
+ ls_be = len(self.stop_data_ev)
+ delta = 0.0
+ batch_delta = 0.0
+ temp = 0
+ power = 0
+ clock = 0
+ if ls_ev != le_ev:
+ logger.warning(f"Event Start/Stop out of sync {ls_ev}/{le_ev}")
+ elif ls_bs != ls_be:
+ logger.warning(f"get_batch Start/Stop out of sync {ls_bs}/{ls_be}")
+ else:
+ temp = torch.cuda.temperature()
+ power = torch.cuda.power_draw()
+ util = torch.cuda.utilization()
+ clock = torch.cuda.clock_rate()
+ torch.cuda.synchronize()
+ # Process Events
+ for i in range(ls_ev):
+ e_ev = self.start_gemm_ev[i].elapsed_time(self.stop_gemm_ev[i])
+ e_tm = (self.stop_gemm_tm[i] - self.start_gemm_tm[i]) / 1e6 # ns to ms
+ # Pick the larger of Event and perf_counter time?
+ delta += max(e_ev, e_tm)
+ # Process get_batch
+ for i in range(ls_bs):
+ b_ev = self.start_data_ev[i].elapsed_time(self.stop_data_ev[i])
+ b_tm = (self.stop_data_tm[i] - self.start_data_tm[i]) / 1e6 # ns to ms
+ # data fetching has prefetch, hence take the max, instead of avg
+ batch_delta = max(batch_delta, max(b_ev, b_tm))
+ self.reset() # Prepare for next round
+ # time in ms, batch_delta in ms, check return above
+ return delta, batch_delta, temp, power, util, clock
+
+ def report(self, total_flops: float = 0.0, log_interval: int = 0) -> bool:
+ """Function to log the min/max metircs and the associated rank over a time period
+
+ It finds the slowest and fastest rank among all ranks. It should be
+ called by all ranks, but only rank-0 prints the analysis
+ At the end it checks, if the straggler detector should
+ remain active or if it should be deactivated.
+
+ Args:
+ total_flops (float, optional): The theoretical flops over the period. Defaults to 0.0.
+ log_interval (int, optional): The training interval over which reporting is called(ms)
+ Defaults to 0.
+
+ Returns:
+ bool: True if reported, else False
+ """
+ ret = False
+ if not self._off and total_flops > 0.0 and log_interval > 0:
+ elapsed, btime, temp, power, util, clock = self.elapsed() # get raw time
+ # btime (get_batch time is max in the iteration)
+ ptime = elapsed / (log_interval * 1.0) # avg per iteration elapsed time, ms
+ api_flops = total_flops / (log_interval * 1.0) # avg per iteration flops, ms
+ apir_flops = api_flops / (
+ ptime * 10**9 * self.world
+ ) # this is avg per iteration this rank's thruput, TFLOP/s (note 10**9),
+ et_flops = apir_flops / self.amp # Estimated TFLOPs, not tracing backward
+
+ o_dt = self._min_max(
+ ptime, btime, float(temp), float(power), float(util), float(clock), et_flops
+ )
+ if self.rank == 0 and o_dt is not None and o_dt.aflops is not None:
+ now = f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]"
+ min_flops, min_frank, _ = o_dt.aflops[0]()
+ max_flops, max_frank, _ = o_dt.aflops[-1]()
+ logger.info(
+ f"{now} | "
+ f"MnRtt/Rnk: {o_dt.min_elapsed} | "
+ f"MxRtt/Rnk: {o_dt.max_elapsed} | "
+ f"MnPwr/Rnk: {o_dt.min_power} | "
+ f"MxPwr/Rnk: {o_dt.max_power} | "
+ f"MnTmp/Rnk: {o_dt.min_temp} | "
+ f"MxTmp/Rnk: {o_dt.max_temp} | "
+ f"MnUtl/Rnk: {o_dt.min_util} | "
+ f"MxUtl/Rnk: {o_dt.max_util} | "
+ f"MnClk/Rnk: {o_dt.min_clock} | "
+ f"MxClk/Rnk: {o_dt.max_clock} | "
+ f"MnDRtt/Rnk: {o_dt.min_btime} | "
+ f"MxDRtt/Rnk: {o_dt.max_btime} | "
+ f"MnEtpt/Rnk: {min_flops:.2f}TF/{min_frank} | "
+ f"MxEtpt/Rnk: {max_flops:.2f}TF/{max_frank}"
+ )
+ if self.mmcnt > 1 and self.mmcnt < self.world:
+ line = f"^^^^ Bottom {self.mmcnt} Ranks with lowest Etpt(TF):"
+ for i in range(self.mmcnt):
+ line += f" {o_dt.aflops[i]},"
+ logger.info(line)
+ line = f"^^^^ Top {self.mmcnt} Ranks with highest Etpt(TF):"
+ shift = self.world - self.mmcnt
+ for i in range(self.mmcnt):
+ line += f" {o_dt.aflops[i+shift]},"
+ logger.info(line)
+ ret = True
+
+ # Check/Communicate if tracking is turned off or on
+ self._check_toggle()
+ return ret
+
+ def _check_toggle(self) -> None:
+ """Helper method to check if a request to toggle the collection state was made
+
+ It checks iof collection state toggle req was made via the server listening on
+ rank-0 since last call to report(). Called by report(). Calling this method
+ indirectly from report() is the only way to activate the change that is made
+ via rank-0
+ """
+ # If no change just commnunicate the current
+ off = self._off
+ if self.rank == 0 and self.toggle:
+ off = not self._off
+ self.toggle = False
+ st = torch.tensor(off, dtype=torch.bool, device=self.dev)
+ torch.distributed.broadcast(st, 0) # Blocking
+ # save old switch
+ off = self._off
+ self._off = bool(st.item())
+ if off != self._off:
+ if not self._off:
+ self.start = self.start_method
+ self.stop = self.stop_method
+ state = "ON"
+ else:
+ self.start = self.null_method
+ self.stop = self.null_method
+ state = "OFF"
+ if self.rank == 0:
+ logger.info(f"Toggling StragglerDetector State {state}")
+
+ def _handler(self) -> None:
+ """Thread function for the controller.
+
+ It is a tcp-server that listens on a port. Uses HTTP protocol.
+ If connected to it using curl, it indicates a toggle of the
+ collection state. The actual toggling happens at the end of
+ calling report() when _check_toggle() is called.
+ """
+ resp = r"HTTP/1.0 200 OK\r\nConnection: Close\r\nContent-length: "
+
+ if self.rank == 0:
+ state = "OFF" if self._off else "ON"
+ logger.info(
+ f"Controller ready to recv " f"commands on port {self.port}. Current state {state}"
+ )
+ while True and self.sock is not None:
+ try:
+ conn, _ = self.sock.accept()
+ _ = conn.recv(1024)
+ self.toggle = True
+ state = "ON" if self._off else "OFF"
+ msg = f"Will turn StragglerDetector {state} at next logging interval"
+ msg_len = len(msg)
+ final_resp = f"{resp}{msg_len}\r\n\r\n{msg}"
+ conn.send(final_resp.encode())
+ conn.close()
+ logger.info(msg)
+ except Exception as err:
+ logger.error(f"Error in stragler handler.. {str(err)}")
+ return
+
+ def _controller(self):
+ """Installs a controller listener that is used to toggle collection state.
+
+ Called from configure(). Ignored for all ranks other than rank-0
+ """
+ try:
+ if self.rank == 0:
+ neth = "0.0.0.0"
+ netp = self.port
+ self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+ self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+ self.sock.bind((neth, netp))
+ self.sock.listen(128)
+ self.ctrlr = threading.Thread(
+ target=self._handler, args=(), name="straggler", daemon=True
+ )
+ self.ctrlr.start()
+ except Exception as err:
+ logger.warning(f"StragglerDetector cannot be controlled.. {str(err)}")
+
+ def _min_max(
+ self,
+ ptime: float,
+ btime: float,
+ temp: float,
+ power: float,
+ util: float,
+ clock: float,
+ flops: float,
+ ) -> Union[_StragglerData, None]:
+ """Helper function to find the min/max values
+
+ Args:
+ ptime (float): avg per iteration gpu time
+ btime (float): avg per iteration cpu time
+ temp (float): gpu temp at the time of reporting
+ power (float): gpu power at the time of reporting
+ util (float): gpu util at the time of reporting
+ clock (float): gpu clock at the time of reporting
+ flops (float): estimated flops for the rank
+
+ Returns:
+ Union[_StragglerData, None]: It contains the min/max of few metrics and the
+ corresponding rank it also has sorted list of
+ all (flops, rank) sorted by flops (aflops)
+ or returns None if collecton is disabled
+ """
+ if self._off:
+ return None
+ # initialize output data object
+ o_dt = _StragglerData()
+
+ prof_data: Dict[str, Union[int, float]] = {}
+ data_list: List[Dict[str, Union[int, float]]] = []
+ prof_data["rank"] = self.rank
+ prof_data["time"] = ptime
+ prof_data["btime"] = btime
+ prof_data["temp"] = temp
+ prof_data["power"] = power
+ prof_data["util"] = util
+ prof_data["clock"] = clock
+ prof_data["flops"] = flops
+
+ if self.rank == 0:
+ data_list = [prof_data] * self.world
+
+ # this is blocking by default
+ torch.distributed.gather_object(prof_data, object_gather_list=data_list, dst=0)
+
+ if self.rank == 0:
+ min_ctime = min(data_list, key=lambda k: k["time"]) # elapsed
+ max_ctime = max(data_list, key=lambda k: k["time"]) # elapsed
+
+ min_cbatch = min(data_list, key=lambda k: k["btime"]) # batch time
+ max_cbatch = max(data_list, key=lambda k: k["btime"]) # batch time
+
+ min_ctemp = min(data_list, key=lambda k: k["temp"]) # temp
+ max_ctemp = max(data_list, key=lambda k: k["temp"]) # temp
+
+ min_cpower = min(data_list, key=lambda k: k["power"]) # power
+ max_cpower = max(data_list, key=lambda k: k["power"]) # power
+
+ min_cutil = min(data_list, key=lambda k: k["util"]) # gpu util
+ max_cutil = max(data_list, key=lambda k: k["util"]) # gpu util
+
+ min_cclock = min(data_list, key=lambda k: k["clock"]) # gpu clock
+ max_cclock = max(data_list, key=lambda k: k["clock"]) # gpu clock
+
+ min_val = min_ctime["time"]
+ min_rank = min_ctime["rank"]
+ max_val = max_ctime["time"]
+ max_rank = max_ctime["rank"]
+ o_dt.min_elapsed = _ValueWithRank(min_val, int(min_rank), "ms")
+ o_dt.max_elapsed = _ValueWithRank(max_val, int(max_rank), "ms")
+
+ min_val = min_cbatch["btime"]
+ min_rank = min_cbatch["rank"]
+ max_val = max_cbatch["btime"]
+ max_rank = max_cbatch["rank"]
+ o_dt.min_btime = _ValueWithRank(min_val, int(min_rank), "ms")
+ o_dt.max_btime = _ValueWithRank(max_val, int(max_rank), "ms")
+
+ min_val = min_ctemp["temp"]
+ min_rank = min_ctemp["rank"]
+ max_val = max_ctemp["temp"]
+ max_rank = max_ctemp["rank"]
+ o_dt.min_temp = _ValueWithRank(min_val, int(min_rank), "C")
+ o_dt.max_temp = _ValueWithRank(max_val, int(max_rank), "C")
+
+ min_val = min_cpower["power"]
+ min_rank = min_cpower["rank"]
+ max_val = max_cpower["power"]
+ max_rank = max_cpower["rank"]
+ o_dt.min_power = _ValueWithRank(min_val, int(min_rank), "W")
+ o_dt.max_power = _ValueWithRank(max_val, int(max_rank), "W")
+
+ min_val = min_cutil["util"]
+ min_rank = min_cutil["rank"]
+ max_val = max_cutil["util"]
+ max_rank = max_cutil["rank"]
+ o_dt.min_util = _ValueWithRank(min_val, int(min_rank), "%")
+ o_dt.max_util = _ValueWithRank(max_val, int(max_rank), "%")
+
+ min_val = min_cclock["clock"]
+ min_rank = min_cclock["rank"]
+ max_val = max_cclock["clock"]
+ max_rank = max_cclock["rank"]
+ o_dt.min_clock = _ValueWithRank(min_val, int(min_rank), "MHz")
+ o_dt.max_clock = _ValueWithRank(max_val, int(max_rank), "MHz")
+
+ o_dt.aflops = [
+ _ValueWithRank(d.get("flops", 0.0), int(d.get("rank", -1)))
+ for _, d in enumerate(data_list)
+ ]
+ o_dt.aflops.sort(key=lambda val_with_rank: val_with_rank()[0])
+ # wait for everyone here
+ torch.distributed.barrier()
+
+ return o_dt
+
+ @property
+ def enabled(self) -> bool:
+ """Can be called to check the enabled state of the instance
+
+ Note:
+ After the request to toggle the state, the
+ actual state change happens at end of call
+ to report()
+ """
+ return not self._off
+
+ @property
+ def configured(self) -> bool:
+ """Can be called to check if the the instance is already configured
+
+ Returns:
+ bool: returns True if configure was called and was a success, else False
+ """
+ return StragglerDetector._configured
+
+ @property
+ def my_rank(self):
+ """Can be called to get configured rank of this instance
+
+ Returns:
+ int: Configured rank for this instance
+ """
+ return self.rank
+
+ @property
+ def world_size(self) -> int:
+ """Can be called to get configured world of this instance
+
+ Returns:
+ int: World size configured for this instance
+ """
+ return self.world
+
+ def null_method(self) -> None:
+ """Default method to initialize start/stop method ptrs"""
+ pass
+
+ def __enter__(self) -> "StragglerDetector":
+ """Define context/instance entry
+
+ Returns:
+ StragglerDetector: the instance
+ """
+ self.start()
+ return self
+
+ def __call__(self, bdata: bool = False) -> "StragglerDetector":
+ """Callable for the instance. Set context state,
+
+ Useful when the context is used for cpu timers only when bdata=True
+
+ Args:
+ bdata (bool, optional): when true, only enables cpu timers. Defaults to False.
+
+ Returns:
+ StragglerDetector: the instance
+ """
+ self.bdata = bdata
+ return self
+
+ def __exit__(
+ self,
+ ex_type: Optional[Type[BaseException]],
+ ex_val: Optional[BaseException],
+ ex_tb: Optional[TracebackType],
+ ) -> bool:
+ """Define context/instance exit, calls the stop method
+
+ Args:
+ ex_type (Optional[Type[BaseException]]): Exception type
+ ex_val (Optional[BaseException]): _description_
+ ex_tb (Optional[TracebackType]): _description_
+
+ Returns:
+ bool: True if the exception was handled
+ """
+ # Should not suppress errors even if turned off
+ if ex_type is not None:
+ err = traceback.format_exception(ex_type, ex_val, ex_tb)
+ logger.warning(f"{str(ex_val)}\n{err}")
+ self.stop()
+ return False
+
+
+# Singleton, global visibility
+__straggler__ = StragglerDetector()
+"""StragglerDetector: private module variable, not be directly accessed
+"""
+
+
+# Check if Transformer Engine has Float8Tensor class
+HAVE_TE_FLOAT8TENSOR = False
+try:
+ from transformer_engine.pytorch.float8_tensor import Float8Tensor
+
+ HAVE_TE_FLOAT8TENSOR = True
+except (ImportError, ModuleNotFoundError):
+ # Float8Tensor not found
+ pass
+
+
+def is_float8tensor(tensor: torch.Tensor) -> bool:
+ """Check if a tensor is a Transformer Engine Float8Tensor"""
+ return HAVE_TE_FLOAT8TENSOR and isinstance(tensor, Float8Tensor)
diff --git a/megatron/data/__init__.py b/megatron/data/__init__.py
deleted file mode 100644
index cd5f898c6b..0000000000
--- a/megatron/data/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from . import indexed_dataset
diff --git a/megatron/data/bert_dataset.py b/megatron/data/bert_dataset.py
deleted file mode 100644
index 036e6bccc9..0000000000
--- a/megatron/data/bert_dataset.py
+++ /dev/null
@@ -1,183 +0,0 @@
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
-
-"""BERT Style dataset."""
-
-import numpy as np
-import torch
-
-from megatron import (
- get_args,
- get_tokenizer,
- mpu,
- print_rank_0
-)
-from megatron.data.dataset_utils import (
- get_samples_mapping,
- get_a_and_b_segments,
- truncate_segments,
- create_tokens_and_tokentypes,
- create_masked_lm_predictions
-)
-
-class BertDataset(torch.utils.data.Dataset):
-
- def __init__(self, name, indexed_dataset, data_prefix,
- num_epochs, max_num_samples, masked_lm_prob,
- max_seq_length, short_seq_prob, seed, binary_head):
-
- # Params to store.
- self.name = name
- self.seed = seed
- self.masked_lm_prob = masked_lm_prob
- self.max_seq_length = max_seq_length
- self.binary_head = binary_head
-
- # Dataset.
- self.indexed_dataset = indexed_dataset
-
- # Build the samples mapping.
- self.samples_mapping = get_samples_mapping(self.indexed_dataset,
- data_prefix,
- num_epochs,
- max_num_samples,
- self.max_seq_length - 3, # account for added tokens
- short_seq_prob,
- self.seed,
- self.name,
- self.binary_head)
-
- # Vocab stuff.
- tokenizer = get_tokenizer()
- self.vocab_id_list = list(tokenizer.inv_vocab.keys())
- self.vocab_id_to_token_dict = tokenizer.inv_vocab
- self.cls_id = tokenizer.cls
- self.sep_id = tokenizer.sep
- self.mask_id = tokenizer.mask
- self.pad_id = tokenizer.pad
-
- def __len__(self):
- return self.samples_mapping.shape[0]
-
- def __getitem__(self, idx):
- start_idx, end_idx, seq_length = self.samples_mapping[idx]
- sample = [self.indexed_dataset[i] for i in range(start_idx, end_idx)]
- # Note that this rng state should be numpy and not python since
- # python randint is inclusive whereas the numpy one is exclusive.
- # We % 2**32 since numpy requres the seed to be between 0 and 2**32 - 1
- np_rng = np.random.RandomState(seed=((self.seed + idx) % 2**32))
- return build_training_sample(sample, seq_length,
- self.max_seq_length, # needed for padding
- self.vocab_id_list,
- self.vocab_id_to_token_dict,
- self.cls_id, self.sep_id,
- self.mask_id, self.pad_id,
- self.masked_lm_prob, np_rng,
- self.binary_head)
-
-
-
-
-def build_training_sample(sample,
- target_seq_length, max_seq_length,
- vocab_id_list, vocab_id_to_token_dict,
- cls_id, sep_id, mask_id, pad_id,
- masked_lm_prob, np_rng, binary_head):
- """Biuld training sample.
-
- Arguments:
- sample: A list of sentences in which each sentence is a list token ids.
- target_seq_length: Desired sequence length.
- max_seq_length: Maximum length of the sequence. All values are padded to
- this length.
- vocab_id_list: List of vocabulary ids. Used to pick a random id.
- vocab_id_to_token_dict: A dictionary from vocab ids to text tokens.
- cls_id: Start of example id.
- sep_id: Separator id.
- mask_id: Mask token id.
- pad_id: Padding token id.
- masked_lm_prob: Probability to mask tokens.
- np_rng: Random number genenrator. Note that this rng state should be
- numpy and not python since python randint is inclusive for
- the opper bound whereas the numpy one is exclusive.
- """
-
- if binary_head:
- # We assume that we have at least two sentences in the sample
- assert len(sample) > 1
- assert target_seq_length <= max_seq_length
-
- # Divide sample into two segments (A and B).
- if binary_head:
- tokens_a, tokens_b, is_next_random = get_a_and_b_segments(sample,
- np_rng)
- else:
- tokens_a = []
- for j in range(len(sample)):
- tokens_a.extend(sample[j])
- tokens_b = []
- is_next_random = False
-
- # Truncate to `target_sequence_length`.
- max_num_tokens = target_seq_length
- truncated = truncate_segments(tokens_a, tokens_b, len(tokens_a),
- len(tokens_b), max_num_tokens, np_rng)
-
- # Build tokens and toketypes.
- tokens, tokentypes = create_tokens_and_tokentypes(tokens_a, tokens_b,
- cls_id, sep_id)
-
- # Masking.
- max_predictions_per_seq = masked_lm_prob * max_num_tokens
- (tokens, masked_positions, masked_labels, _, _) = create_masked_lm_predictions(
- tokens, vocab_id_list, vocab_id_to_token_dict, masked_lm_prob,
- cls_id, sep_id, mask_id, max_predictions_per_seq, np_rng)
-
- # Padding.
- tokens_np, tokentypes_np, labels_np, padding_mask_np, loss_mask_np \
- = pad_and_convert_to_numpy(tokens, tokentypes, masked_positions,
- masked_labels, pad_id, max_seq_length)
-
- train_sample = {
- 'text': tokens_np,
- 'types': tokentypes_np,
- 'labels': labels_np,
- 'is_random': int(is_next_random),
- 'loss_mask': loss_mask_np,
- 'padding_mask': padding_mask_np,
- 'truncated': int(truncated)}
- return train_sample
-
-
-def pad_and_convert_to_numpy(tokens, tokentypes, masked_positions,
- masked_labels, pad_id, max_seq_length):
- """Pad sequences and convert them to numpy."""
-
- # Some checks.
- num_tokens = len(tokens)
- padding_length = max_seq_length - num_tokens
- assert padding_length >= 0, \
- f"num_tokens ({num_tokens}) is greater than " \
- "max_seq_length ({max_seq_length})."
- assert len(tokentypes) == num_tokens
- assert len(masked_positions) == len(masked_labels)
-
- # Tokens and token types.
- filler = [pad_id] * padding_length
- tokens_np = np.array(tokens + filler, dtype=np.int64)
- tokentypes_np = np.array(tokentypes + filler, dtype=np.int64)
-
- # Padding mask.
- padding_mask_np = np.array([1] * num_tokens + [0] * padding_length,
- dtype=np.int64)
-
- # Lables and loss mask.
- labels = [-1] * max_seq_length
- loss_mask = [0] * max_seq_length
- for i in range(len(masked_positions)):
- assert masked_positions[i] < num_tokens
- labels[masked_positions[i]] = masked_labels[i]
- loss_mask[masked_positions[i]] = 1
- labels_np = np.array(labels, dtype=np.int64)
- loss_mask_np = np.array(loss_mask, dtype=np.int64)
-
- return tokens_np, tokentypes_np, labels_np, padding_mask_np, loss_mask_np
diff --git a/megatron/data/blendable_dataset.py b/megatron/data/blendable_dataset.py
deleted file mode 100644
index 453b362f3e..0000000000
--- a/megatron/data/blendable_dataset.py
+++ /dev/null
@@ -1,64 +0,0 @@
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
-
-"""Blendable dataset."""
-
-import time
-
-import numpy as np
-import torch
-
-from megatron import print_rank_0
-
-class BlendableDataset(torch.utils.data.Dataset):
-
-
- def __init__(self, datasets, weights, size):
-
- self.datasets = datasets
- num_datasets = len(datasets)
- assert num_datasets == len(weights)
-
- self.size = size
-
- # Normalize weights.
- weights = np.array(weights, dtype=np.float64)
- sum_weights = np.sum(weights)
- assert sum_weights > 0.0
- weights /= sum_weights
-
- # Build indicies.
- start_time = time.time()
- assert num_datasets < 255
- self.dataset_index = np.zeros(self.size, dtype=np.uint8)
- self.dataset_sample_index = np.zeros(self.size, dtype=np.int64)
-
- from megatron.data import helpers
- helpers.build_blending_indices(self.dataset_index,
- self.dataset_sample_index,
- weights, num_datasets, self.size,
- torch.distributed.get_rank() == 0)
- print_rank_0('> elapsed time for building blendable dataset indices: '
- '{:.2f} (sec)'.format(time.time() - start_time))
-
- # Check size
- _ = self.__getitem__(self.size - 1)
- try:
- _ = self.__getitem__(self.size)
- raise RuntimeError('BlendedDataset size is improperly bounded')
- except IndexError:
- pass
- print_rank_0('> size of blendable dataset: '
- '{} samples'.format(self.size))
-
-
- def __len__(self):
- return self.size
-
-
- def __getitem__(self, idx):
- dataset_idx = self.dataset_index[idx]
- sample_idx = self.dataset_sample_index[idx]
- return {
- "dataset_idx" : dataset_idx,
- **self.datasets[dataset_idx][sample_idx],
- }
diff --git a/megatron/data/gpt_dataset.py b/megatron/data/gpt_dataset.py
deleted file mode 100644
index 602e511678..0000000000
--- a/megatron/data/gpt_dataset.py
+++ /dev/null
@@ -1,524 +0,0 @@
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
-
-"""GPT style dataset."""
-
-import os
-import time
-
-import numpy as np
-import torch
-
-from megatron import print_rank_0
-from megatron.core import mpu
-from megatron.data.blendable_dataset import BlendableDataset
-from megatron.data.dataset_utils import get_datasets_weights_and_num_samples
-from megatron.data.dataset_utils import get_train_valid_test_split_
-from megatron.data.indexed_dataset import make_dataset as make_indexed_dataset
-
-
-def build_train_valid_test_datasets(data_prefix, data_impl, splits_string,
- train_valid_test_num_samples,
- seq_length, seed, skip_warmup,
- train_data_prefix=None,
- valid_data_prefix=None,
- test_data_prefix=None,
- return_doc_ids=False):
- """Build train, valid, and test datasets."""
-
- if data_prefix:
- print_rank_0("Single data path provided for train, valid & test")
-
- # Single dataset.
- if len(data_prefix) == 1:
- return _build_train_valid_test_datasets(data_prefix[0],
- data_impl, splits_string,
- train_valid_test_num_samples,
- seq_length, seed, skip_warmup)
-
- # Blending dataset.
- # Parse the values.
- output = get_datasets_weights_and_num_samples(data_prefix,
- train_valid_test_num_samples)
- prefixes, weights, datasets_train_valid_test_num_samples = output
- train_num_samples, valid_num_samples, test_num_samples = map(
- sum,
- zip(*datasets_train_valid_test_num_samples)
- )
-
- # Build individual datasets.
- train_datasets = []
- valid_datasets = []
- test_datasets = []
- for i in range(len(prefixes)):
- train_ds, valid_ds, test_ds = _build_train_valid_test_datasets(
- prefixes[i], data_impl, splits_string,
- datasets_train_valid_test_num_samples[i],
- seq_length, seed, skip_warmup,
- return_doc_ids)
- if train_ds:
- train_datasets.append(train_ds)
- if valid_ds:
- valid_datasets.append(valid_ds)
- if test_ds:
- test_datasets.append(test_ds)
-
- # Blend.
- blending_train_dataset = None
- if train_datasets:
- blending_train_dataset = BlendableDataset(train_datasets, weights, train_num_samples)
- blending_valid_dataset = None
- if valid_datasets:
- blending_valid_dataset = BlendableDataset(valid_datasets, weights, valid_num_samples)
- blending_test_dataset = None
- if test_datasets:
- blending_test_dataset = BlendableDataset(test_datasets, weights, test_num_samples)
-
- return (blending_train_dataset, blending_valid_dataset,
- blending_test_dataset)
-
- else:
- print_rank_0("Separate data paths provided for train, valid & test. Split string will be ignored.")
-
- train_dataset, valid_dataset, test_dataset = None, None, None
- # Single dataset.
- if train_data_prefix is not None:
- train_dataset = build_dataset("train", train_data_prefix, data_impl,
- train_valid_test_num_samples[0],
- seq_length, seed, skip_warmup)
-
- if valid_data_prefix is not None:
- valid_dataset = build_dataset("valid", valid_data_prefix, data_impl,
- train_valid_test_num_samples[1],
- seq_length, seed, False)
-
- if test_data_prefix is not None:
- test_dataset = build_dataset("test", test_data_prefix, data_impl,
- train_valid_test_num_samples[2],
- seq_length, seed, False)
-
- return (train_dataset, valid_dataset, test_dataset)
-
-
-def _build_train_valid_test_datasets(data_prefix, data_impl, splits_string,
- train_valid_test_num_samples,
- seq_length, seed, skip_warmup,
- return_doc_ids=False):
- """Build train, valid, and test datasets."""
-
- # Indexed dataset.
- indexed_dataset = get_indexed_dataset_(data_prefix,
- data_impl,
- skip_warmup)
-
- total_num_of_documents = indexed_dataset.sizes.shape[0]
- splits = get_train_valid_test_split_(splits_string, total_num_of_documents)
-
- # Print stats about the splits.
- print_rank_0(' > dataset split:')
-
- def print_split_stats(name, index):
- print_rank_0(' {}:'.format(name))
- print_rank_0(' document indices in [{}, {}) total of {} '
- 'documents'.format(splits[index], splits[index + 1],
- splits[index + 1] - splits[index]))
- print_split_stats('train', 0)
- print_split_stats('validation', 1)
- print_split_stats('test', 2)
-
- def build_dataset(index, name):
- dataset = None
- if splits[index + 1] > splits[index]:
- documents = np.arange(start=splits[index], stop=splits[index + 1],
- step=1, dtype=np.int32)
- dataset = GPTDataset(name, data_prefix,
- documents, indexed_dataset,
- train_valid_test_num_samples[index],
- seq_length, seed,
- return_doc_ids)
- return dataset
-
- train_dataset = build_dataset(0, 'train')
- valid_dataset = build_dataset(1, 'valid')
- test_dataset = build_dataset(2, 'test')
-
- return (train_dataset, valid_dataset, test_dataset)
-
-
-def build_dataset(dataset_name, data_prefix, data_impl, num_samples,
- seq_length, seed, skip_warmup):
- dataset = None
- if len(data_prefix) == 1:
- dataset = _build_dataset(dataset_name,
- data_prefix[0], data_impl,
- num_samples, seq_length,
- seed, skip_warmup)
- else:
- # Blending dataset.
- # Parse the values.
- output = get_datasets_weights_and_num_samples(data_prefix, num_samples)
- prefixes, weights, dataset_num_samples = output
- num_samples = sum(dataset_num_samples)
-
- # Build individual datasets.
- datasets = []
- for i in range(len(prefixes)):
- ds = _build_dataset(dataset_name, prefixes[i],
- data_impl, dataset_num_samples[i],
- seq_length, seed, skip_warmup)
- if ds:
- datasets.append(ds)
-
- if datasets:
- dataset = BlendableDataset(datasets, weights, num_samples)
-
- return dataset
-
-
-def _build_dataset(dataset_name, data_prefix, data_impl,
- num_samples, seq_length, seed, skip_warmup):
- """
- Build dataset. This method is called when individual
- train, valid, test datasets are provided
- """
-
- # Indexed dataset.
- indexed_dataset = get_indexed_dataset_(data_prefix,
- data_impl,
- skip_warmup)
-
- total_num_of_documents = indexed_dataset.sizes.shape[0]
-
- print_rank_0(' {}:'.format(dataset_name))
- print_rank_0(' document indices in [0, {}) total of {} '
- 'documents'.format(total_num_of_documents, total_num_of_documents))
-
- documents = np.arange(start=0, stop=total_num_of_documents,
- step=1, dtype=np.int32)
-
- dataset = GPTDataset(dataset_name, data_prefix,
- documents, indexed_dataset,
- num_samples, seq_length, seed)
-
- return dataset
-
-
-def get_indexed_dataset_(data_prefix, data_impl, skip_warmup):
- """Build indexed dataset."""
- print_rank_0(' > building dataset index ...')
-
- start_time = time.time()
- indexed_dataset = make_indexed_dataset(data_prefix,
- data_impl,
- skip_warmup)
- print_rank_0(' > finished creating indexed dataset in {:4f} '
- 'seconds'.format(time.time() - start_time))
- print_rank_0(' number of documents: {}'.format(
- indexed_dataset.sizes.shape[0]))
-
- return indexed_dataset
-
-
-class GPTDataset(torch.utils.data.Dataset):
-
- def __init__(self, name, data_prefix, documents, indexed_dataset,
- num_samples, seq_length, seed,
- return_doc_ids=False):
-
- self.name = name
- self.indexed_dataset = indexed_dataset
- self.return_doc_ids = return_doc_ids
-
- # Checks
- assert np.min(documents) >= 0
- assert np.max(documents) < indexed_dataset.sizes.shape[0]
-
- # Build index mappings.
- self.doc_idx, self.sample_idx, self.shuffle_idx, self.index_prefix = \
- _build_index_mappings(self.name, data_prefix,
- documents, self.indexed_dataset.sizes,
- num_samples, seq_length, seed)
-
-
- def __len__(self):
- # -1 is due to data structure used to retieve the index:
- # sample i --> [sample_idx[i], sample_idx[i+1])
- return self.sample_idx.shape[0] - 1
-
- def __getitem__(self, idx):
- # Get the shuffled index.
- idx = self.shuffle_idx[idx]
- # Start and end documents and offsets.
- doc_index_f = self.sample_idx[idx][0]
- doc_index_l = self.sample_idx[idx + 1][0]
- offset_f = self.sample_idx[idx][1]
- offset_l = self.sample_idx[idx + 1][1]
- # If we are within the same document, just extract the chunk.
- doc_ids = []
- if doc_index_f == doc_index_l:
- doc_ids.append(self.doc_idx[doc_index_f])
- sample = self.indexed_dataset.get(self.doc_idx[doc_index_f],
- offset=offset_f,
- length=offset_l - offset_f + 1)
- else:
- # Otherwise, get the rest of the initial document.
- doc_ids.append(self.doc_idx[doc_index_f])
- sample_list = [self.indexed_dataset.get(self.doc_idx[doc_index_f],
- offset=offset_f)]
- # Loop over all in between documents and add the entire document.
- for i in range(doc_index_f + 1, doc_index_l):
- doc_ids.append(self.doc_idx[i])
- sample_list.append(self.indexed_dataset.get(self.doc_idx[i]))
- # And finally add the relevant portion of last document.
- doc_ids.append(self.doc_idx[doc_index_l])
- sample_list.append(self.indexed_dataset.get(
- self.doc_idx[doc_index_l],
- length=offset_l + 1))
- sample = np.concatenate(sample_list)
-
- if self.return_doc_ids: # for retro preprocessing
- return {'text': np.array(sample, dtype=np.int64),
- 'doc_ids': np.array(doc_ids, dtype=np.int64)}
- else:
- return {'text': np.array(sample, dtype=np.int64)}
-
-
-def _build_index_mappings(name, data_prefix, documents, sizes,
- num_samples, seq_length, seed):
- """Build doc-idx, sample-idx, and shuffle-idx.
- doc-idx: is an array (ordered) of documents to be used in training.
- sample-idx: is the start document index and document offset for each
- training sample.
- shuffle-idx: maps the sample index into a random index into sample-idx.
- """
- # Number of tokens in each epoch and number of required epochs.
- tokens_per_epoch = _num_tokens(documents, sizes)
- num_epochs = _num_epochs(tokens_per_epoch, seq_length, num_samples)
-
- # rng state
- np_rng = np.random.RandomState(seed=seed)
-
- # Filename of the index mappings.
- index_prefix = '{}_indexmap'.format(name)
- index_prefix += '_{}ns'.format(num_samples)
- index_prefix += '_{}sl'.format(seq_length)
- index_prefix += '_{}s'.format(seed)
- _filename = data_prefix + '_' + index_prefix
- doc_idx_filename = _filename + '_doc_idx.npy'
- sample_idx_filename = _filename + '_sample_idx.npy'
- shuffle_idx_filename = _filename + '_shuffle_idx.npy'
-
- # Build the indexed mapping if not exist.
- if torch.distributed.get_rank() == 0 and \
- (not os.path.isfile(doc_idx_filename) or
- not os.path.isfile(sample_idx_filename) or
- not os.path.isfile(shuffle_idx_filename)):
-
- print_rank_0(' > WARNING: could not find index map files, building '
- 'the indices on rank 0 ...')
-
- # For the last epoch, decide whether include the entire epoch
- # in the global shuffle or not.
-
- # If we need only one epoch, then separating last epoch does
- # not mean anything.
- if num_epochs == 1:
- separate_last_epoch = False
- print(' > only one epoch required, setting '
- 'separate_last_epoch to False', flush=True)
-
- else:
- # Get the number of samples for the last epoch
- num_samples_from_epochs_minus_one = (
- (num_epochs - 1) * tokens_per_epoch - 1) // seq_length
- last_epoch_num_samples = num_samples - \
- num_samples_from_epochs_minus_one
- assert last_epoch_num_samples >= 0, \
- 'last epoch number of samples should be non-negative.'
- num_samples_per_epoch = (tokens_per_epoch - 1) // seq_length
- assert last_epoch_num_samples < (num_samples_per_epoch + 1), \
- 'last epoch number of samples exceeded max value.'
- # If we have less than 80% of the samples for the last epoch,
- # seperate out the epoch and treat it differently.
- # Note: the 80% number is just based on common sense and can
- # be adjusted if needed.
- separate_last_epoch = (last_epoch_num_samples <
- int(0.80 * num_samples_per_epoch))
- if separate_last_epoch:
- string = ' > last epoch number of samples ({}) is smaller '\
- 'than 80% of number of samples per epoch ({}), '\
- 'setting separate_last_epoch to True'
- else:
- string = ' > last epoch number of samples ({}) is larger '\
- 'than 80% of number of samples per epoch ({}), '\
- 'setting separate_last_epoch to False'
- print(string.format(last_epoch_num_samples,
- num_samples_per_epoch), flush=True)
-
- # doc-idx.
- start_time = time.time()
- doc_idx = _build_doc_idx(documents, num_epochs, np_rng,
- separate_last_epoch)
- np.save(doc_idx_filename, doc_idx, allow_pickle=True)
- print_rank_0(' > elasped time to build and save doc-idx mapping '
- '(seconds): {:4f}'.format(time.time() - start_time))
- # sample-idx.
- start_time = time.time()
- # Use C++ implementation for speed.
- # First compile and then import.
- from megatron.data import helpers
- assert doc_idx.dtype == np.int32
- assert sizes.dtype == np.int32
- sample_idx = helpers.build_sample_idx(sizes, doc_idx, seq_length,
- num_epochs, tokens_per_epoch)
- np.save(sample_idx_filename, sample_idx, allow_pickle=True)
- print_rank_0(' > elasped time to build and save sample-idx mapping '
- '(seconds): {:4f}'.format(time.time() - start_time))
- # shuffle-idx.
- start_time = time.time()
- # -1 is due to data structure used to retieve the index:
- # sample i --> [sample_idx[i], sample_idx[i+1])
- if separate_last_epoch:
- num_samples_ = num_samples_from_epochs_minus_one
- else:
- num_samples_ = sample_idx.shape[0] - 1
- shuffle_idx = _build_shuffle_idx(num_samples_,
- sample_idx.shape[0] - 1, np_rng)
- np.save(shuffle_idx_filename, shuffle_idx, allow_pickle=True)
- print_rank_0(' > elasped time to build and save shuffle-idx mapping'
- ' (seconds): {:4f}'.format(time.time() - start_time))
-
- # This should be a barrier but nccl barrier assumes
- # device_index=rank which is not the case for model
- # parallel case
- counts = torch.cuda.LongTensor([1])
- torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
- torch.distributed.all_reduce(counts, group=mpu.get_pipeline_model_parallel_group())
- assert counts[0].item() == (
- torch.distributed.get_world_size() //
- torch.distributed.get_world_size(group=mpu.get_tensor_model_parallel_group()))
-
- # Load mappings.
- start_time = time.time()
- print_rank_0(' > loading doc-idx mapping from {}'.format(
- doc_idx_filename))
- doc_idx = np.load(doc_idx_filename, allow_pickle=True, mmap_mode='r')
- print_rank_0(' > loading sample-idx mapping from {}'.format(
- sample_idx_filename))
- sample_idx = np.load(sample_idx_filename, allow_pickle=True, mmap_mode='r')
- print_rank_0(' > loading shuffle-idx mapping from {}'.format(
- shuffle_idx_filename))
- shuffle_idx = np.load(shuffle_idx_filename, allow_pickle=True, mmap_mode='r')
- print_rank_0(' loaded indexed file in {:3.3f} seconds'.format(
- time.time() - start_time))
- print_rank_0(' total number of samples: {}'.format(
- sample_idx.shape[0]))
- print_rank_0(' total number of epochs: {}'.format(num_epochs))
-
- return doc_idx, sample_idx, shuffle_idx, index_prefix
-
-
-def _num_tokens(documents, sizes):
- """Total number of tokens in the dataset."""
- return np.sum(sizes[documents])
-
-
-def _num_epochs(tokens_per_epoch, seq_length, num_samples):
- """Based on number of samples and sequence lenght, calculate how many
- epochs will be needed."""
- num_epochs = 0
- total_tokens = 0
- while True:
- num_epochs += 1
- total_tokens += tokens_per_epoch
- # -1 is because we need to retrieve seq_length + 1 token each time
- # but the last token will overlap with the first token of the next
- # sample except for the last sample.
- if ((total_tokens - 1) // seq_length) >= num_samples:
- return num_epochs
-
-
-def _build_doc_idx(documents, num_epochs, np_rng, separate_last_epoch):
- """Build an array with length = number-of-epochs * number-of-dcuments.
- Each index is mapped to a corresponding document."""
- if not separate_last_epoch or num_epochs == 1:
- doc_idx = np.mgrid[0:num_epochs, 0:len(documents)][1]
- doc_idx[:] = documents
- doc_idx = doc_idx.reshape(-1)
- doc_idx = doc_idx.astype(np.int32)
- np_rng.shuffle(doc_idx)
- return doc_idx
-
- doc_idx_first = _build_doc_idx(documents, num_epochs-1, np_rng, False)
- doc_idx_last = _build_doc_idx(documents, 1, np_rng, False)
- return np.concatenate((doc_idx_first, doc_idx_last))
-
-
-def _build_sample_idx(sizes, doc_idx, seq_length,
- num_epochs, tokens_per_epoch):
- """Sample index mapping is a 2D array with sizes
- [number-of-samples + 1, 2] where [..., 0] contains
- the index into `doc_idx` and [..., 1] is the
- starting offset in that document."""
-
- # Total number of samples. For -1 see comments in `_num_epochs`.
- num_samples = (num_epochs * tokens_per_epoch - 1) // seq_length
- sample_idx = np.zeros([num_samples + 1, 2], dtype=np.int32)
-
- # Index into sample_idx.
- sample_index = 0
- # Index into doc_idx.
- doc_idx_index = 0
- # Begining offset for each document.
- doc_offset = 0
- # Start with first document and no offset.
- sample_idx[sample_index][0] = doc_idx_index
- sample_idx[sample_index][1] = doc_offset
- sample_index += 1
- while sample_index <= num_samples:
- # Start with a fresh sequence.
- remaining_seq_length = seq_length + 1
- while remaining_seq_length != 0:
- # Get the document length.
- doc_id = doc_idx[doc_idx_index]
- doc_length = sizes[doc_id] - doc_offset
- # And add it to the current sequence.
- remaining_seq_length -= doc_length
- # If we have more than a full sequence, adjust offset and set
- # remaining length to zero so we return from the while loop.
- # Note that -1 here is for the same reason we have -1 in
- # `_num_epochs` calculations.
- if remaining_seq_length <= 0:
- doc_offset += (remaining_seq_length + doc_length - 1)
- remaining_seq_length = 0
- else:
- # Otherwise, start from the begining of the next document.
- doc_idx_index += 1
- doc_offset = 0
- # Record the sequence.
- sample_idx[sample_index][0] = doc_idx_index
- sample_idx[sample_index][1] = doc_offset
- sample_index += 1
-
- return sample_idx
-
-
-def _build_shuffle_idx(num_samples, total_size, np_rng):
- """Build the range [0, size) and shuffle."""
- print(' > building shuffle index with split [0, {}) and [{}, {}) '
- '...'.format(num_samples, num_samples, total_size), flush=True)
-
- dtype_ = np.uint32
- if total_size >= (np.iinfo(np.uint32).max - 1):
- dtype_ = np.int64
-
- shuffle_idx_first = np.arange(start=0, stop=num_samples,
- step=1, dtype=dtype_)
- np_rng.shuffle(shuffle_idx_first)
- if num_samples == total_size:
- return shuffle_idx_first
-
- shuffle_idx_last = np.arange(start=num_samples, stop=total_size,
- step=1, dtype=dtype_)
- np_rng.shuffle(shuffle_idx_last)
-
- return np.concatenate((shuffle_idx_first, shuffle_idx_last))
diff --git a/megatron/data/helpers.cpp b/megatron/data/helpers.cpp
deleted file mode 100644
index 09f5f97626..0000000000
--- a/megatron/data/helpers.cpp
+++ /dev/null
@@ -1,701 +0,0 @@
-/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. */
-
-/* Helper methods for fast index mapping builds */
-
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-
-namespace py = pybind11;
-using namespace std;
-
-const int32_t LONG_SENTENCE_LEN = 512;
-
-
-void build_blending_indices(py::array_t& dataset_index,
- py::array_t& dataset_sample_index,
- const py::array_t& weights,
- const int32_t num_datasets,
- const int64_t size, const bool verbose) {
- /* Given multiple datasets and a weighting array, build samples
- such that it follows those wieghts.*/
-
- if (verbose) {
- std::cout << "> building indices for blendable datasets ..." << std::endl;
- }
-
- // Get the pointer access without the checks.
- auto dataset_index_ptr = dataset_index.mutable_unchecked<1>();
- auto dataset_sample_index_ptr = dataset_sample_index.mutable_unchecked<1>();
- auto weights_ptr = weights.unchecked<1>();
-
- // Initialize buffer for number of samples used for each dataset.
- int64_t current_samples[num_datasets];
- for(int64_t i = 0; i < num_datasets; ++i) {
- current_samples[i] = 0;
- }
-
- // For each sample:
- for(int64_t sample_idx = 0; sample_idx < size; ++sample_idx) {
-
- // Determine where the max error in sampling is happening.
- auto sample_idx_double = std::max(static_cast(sample_idx), 1.0);
- int64_t max_error_index = 0;
- double max_error = weights_ptr[0] * sample_idx_double -
- static_cast(current_samples[0]);
- for (int64_t dataset_idx = 1; dataset_idx < num_datasets; ++dataset_idx) {
- double error = weights_ptr[dataset_idx] * sample_idx_double -
- static_cast(current_samples[dataset_idx]);
- if (error > max_error) {
- max_error = error;
- max_error_index = dataset_idx;
- }
- }
-
- // Populate the indices.
- dataset_index_ptr[sample_idx] = static_cast(max_error_index);
- dataset_sample_index_ptr[sample_idx] = current_samples[max_error_index];
-
- // Update the total samples.
- current_samples[max_error_index] += 1;
-
- }
-
- // print info
- if (verbose) {
- std::cout << " > sample ratios:" << std::endl;
- for (int64_t dataset_idx = 0; dataset_idx < num_datasets; ++dataset_idx) {
- auto ratio = static_cast(current_samples[dataset_idx]) /
- static_cast(size);
- std::cout << " dataset " << dataset_idx << ", input: " <<
- weights_ptr[dataset_idx] << ", achieved: " << ratio << std::endl;
- }
- }
-
-}
-
-
-py::array build_sample_idx(const py::array_t& sizes_,
- const py::array_t& doc_idx_,
- const int32_t seq_length,
- const int32_t num_epochs,
- const int64_t tokens_per_epoch) {
- /* Sample index (sample_idx) is used for gpt2 like dataset for which
- the documents are flattened and the samples are built based on this
- 1-D flatten array. It is a 2D array with sizes [number-of-samples + 1, 2]
- where [..., 0] contains the index into `doc_idx` and [..., 1] is the
- starting offset in that document.*/
-
- // Consistency checks.
- assert(seq_length > 1);
- assert(num_epochs > 0);
- assert(tokens_per_epoch > 1);
-
- // Remove bound checks.
- auto sizes = sizes_.unchecked<1>();
- auto doc_idx = doc_idx_.unchecked<1>();
-
- // Mapping and it's length (1D).
- int64_t num_samples = (num_epochs * tokens_per_epoch - 1) / seq_length;
- int32_t* sample_idx = new int32_t[2*(num_samples+1)];
-
- cout << " using:" << endl << std::flush;
- cout << " number of documents: " <<
- doc_idx_.shape(0) / num_epochs << endl << std::flush;
- cout << " number of epochs: " << num_epochs <<
- endl << std::flush;
- cout << " sequence length: " << seq_length <<
- endl << std::flush;
- cout << " total number of samples: " << num_samples <<
- endl << std::flush;
-
- // Index into sample_idx.
- int64_t sample_index = 0;
- // Index into doc_idx.
- int64_t doc_idx_index = 0;
- // Begining offset for each document.
- int32_t doc_offset = 0;
- // Start with first document and no offset.
- sample_idx[2 * sample_index] = doc_idx_index;
- sample_idx[2 * sample_index + 1] = doc_offset;
- ++sample_index;
-
- while (sample_index <= num_samples) {
- // Start with a fresh sequence.
- int32_t remaining_seq_length = seq_length + 1;
- while (remaining_seq_length != 0) {
- // Get the document length.
- auto doc_id = doc_idx[doc_idx_index];
- auto doc_length = sizes[doc_id] - doc_offset;
- // And add it to the current sequence.
- remaining_seq_length -= doc_length;
- // If we have more than a full sequence, adjust offset and set
- // remaining length to zero so we return from the while loop.
- // Note that -1 here is for the same reason we have -1 in
- // `_num_epochs` calculations.
- if (remaining_seq_length <= 0) {
- doc_offset += (remaining_seq_length + doc_length - 1);
- remaining_seq_length = 0;
- } else {
- // Otherwise, start from the begining of the next document.
- ++doc_idx_index;
- doc_offset = 0;
- }
- }
- // Record the sequence.
- sample_idx[2 * sample_index] = doc_idx_index;
- sample_idx[2 * sample_index + 1] = doc_offset;
- ++sample_index;
- }
-
- // Method to deallocate memory.
- py::capsule free_when_done(sample_idx, [](void *mem_) {
- int32_t *mem = reinterpret_cast(mem_);
- delete[] mem;
- });
-
- // Return the numpy array.
- const auto byte_size = sizeof(int32_t);
- return py::array(std::vector{num_samples+1, 2}, // shape
- {2*byte_size, byte_size}, // C-style contiguous strides
- sample_idx, // the data pointer
- free_when_done); // numpy array references
-
-}
-
-
-inline int32_t get_target_sample_len(const int32_t short_seq_ratio,
- const int32_t max_length,
- std::mt19937& rand32_gen) {
- /* Training sample length. */
- if (short_seq_ratio == 0) {
- return max_length;
- }
- const auto random_number = rand32_gen();
- if ((random_number % short_seq_ratio) == 0) {
- return 2 + random_number % (max_length - 1);
- }
- return max_length;
-}
-
-
-template
-py::array build_mapping_impl(const py::array_t& docs_,
- const py::array_t& sizes_,
- const int32_t num_epochs,
- const uint64_t max_num_samples,
- const int32_t max_seq_length,
- const double short_seq_prob,
- const int32_t seed,
- const bool verbose,
- const int32_t min_num_sent) {
- /* Build a mapping of (start-index, end-index, sequence-length) where
- start and end index are the indices of the sentences in the sample
- and sequence-length is the target sequence length.
- */
-
- // Consistency checks.
- assert(num_epochs > 0);
- assert(max_seq_length > 1);
- assert(short_seq_prob >= 0.0);
- assert(short_seq_prob <= 1.0);
- assert(seed > 0);
-
- // Remove bound checks.
- auto docs = docs_.unchecked<1>();
- auto sizes = sizes_.unchecked<1>();
-
- // For efficiency, convert probability to ratio. Note: rand() generates int.
- int32_t short_seq_ratio = 0;
- if (short_seq_prob > 0) {
- short_seq_ratio = static_cast(round(1.0 / short_seq_prob));
- }
-
- if (verbose) {
- const auto sent_start_index = docs[0];
- const auto sent_end_index = docs[docs_.shape(0) - 1];
- const auto num_sentences = sent_end_index - sent_start_index;
- cout << " using:" << endl << std::flush;
- cout << " number of documents: " << docs_.shape(0) - 1 <<
- endl << std::flush;
- cout << " sentences range: [" << sent_start_index <<
- ", " << sent_end_index << ")" << endl << std::flush;
- cout << " total number of sentences: " << num_sentences <<
- endl << std::flush;
- cout << " number of epochs: " << num_epochs <<
- endl << std::flush;
- cout << " maximum number of samples: " << max_num_samples <<
- endl << std::flush;
- cout << " maximum sequence length: " << max_seq_length <<
- endl << std::flush;
- cout << " short sequence probability: " << short_seq_prob <<
- endl << std::flush;
- cout << " short sequence ration (1/prob): " << short_seq_ratio <<
- endl << std::flush;
- cout << " seed: " << seed << endl <<
- std::flush;
- }
-
- // Mapping and it's length (1D).
- int64_t num_samples = -1;
- DocIdx* maps = NULL;
-
- // Perform two iterations, in the first iteration get the size
- // and allocate memory and in the second iteration populate the map.
- bool second = false;
- for (int32_t iteration=0; iteration<2; ++iteration) {
-
- // Set the seed so both iterations produce the same results.
- std::mt19937 rand32_gen(seed);
-
- // Set the flag on second iteration.
- second = (iteration == 1);
-
- // Counters:
- uint64_t empty_docs = 0;
- uint64_t one_sent_docs = 0;
- uint64_t long_sent_docs = 0;
-
- // Current map index.
- uint64_t map_index = 0;
-
- // For each epoch:
- for (int32_t epoch=0; epoch= max_num_samples) {
- if (verbose && (!second)) {
- cout << " reached " << max_num_samples << " samples after "
- << epoch << " epochs ..." << endl << std::flush;
- }
- break;
- }
- // For each document:
- for (int32_t doc=0; doc<(docs.shape(0) - 1); ++doc) {
-
- // Document sentences are in [sent_index_first, sent_index_last)
- const auto sent_index_first = docs[doc];
- const auto sent_index_last = docs[doc + 1];
-
- // At the begining of the document previous index is the
- // start index.
- auto prev_start_index = sent_index_first;
-
- // Remaining documents.
- auto num_remain_sent = sent_index_last - sent_index_first;
-
- // Some bookkeeping
- if ((epoch == 0) && (!second)) {
- if (num_remain_sent == 0) {
- ++empty_docs;
- }
- if (num_remain_sent == 1) {
- ++one_sent_docs;
- }
- }
-
- // Detect documents with long sentences.
- bool contains_long_sentence = false;
- if (num_remain_sent > 1) {
- for (auto sent_index=sent_index_first;
- sent_index < sent_index_last; ++sent_index) {
- if (sizes[sent_index] > LONG_SENTENCE_LEN){
- if ((epoch == 0) && (!second)) {
- ++long_sent_docs;
- }
- contains_long_sentence = true;
- break;
- }
- }
- }
-
- // If we have more than two sentences.
- if ((num_remain_sent >= min_num_sent) && (!contains_long_sentence)) {
-
- // Set values.
- auto seq_len = int32_t{0};
- auto num_sent = int32_t{0};
- auto target_seq_len = get_target_sample_len(short_seq_ratio,
- max_seq_length,
- rand32_gen);
-
- // Loop through sentences.
- for (auto sent_index=sent_index_first;
- sent_index < sent_index_last; ++sent_index) {
-
- // Add the size and number of sentences.
- seq_len += sizes[sent_index];
- ++num_sent;
- --num_remain_sent;
-
- // If we have reached the target length.
- // and if not only one sentence is left in the document.
- // and if we have at least two sentneces.
- // and if we have reached end of the document.
- if (((seq_len >= target_seq_len) &&
- (num_remain_sent > 1) &&
- (num_sent >= min_num_sent) ) || (num_remain_sent == 0)) {
-
- // Check for overflow.
- if ((3 * map_index + 2) >
- std::numeric_limits::max()) {
- cout << "number of samples exceeded maximum "
- << "allowed by type int64: "
- << std::numeric_limits::max()
- << endl;
- throw std::overflow_error("Number of samples");
- }
-
- // Populate the map.
- if (second) {
- const auto map_index_0 = 3 * map_index;
- maps[map_index_0] = static_cast(prev_start_index);
- maps[map_index_0 + 1] = static_cast(sent_index + 1);
- maps[map_index_0 + 2] = static_cast(target_seq_len);
- }
-
- // Update indices / counters.
- ++map_index;
- prev_start_index = sent_index + 1;
- target_seq_len = get_target_sample_len(short_seq_ratio,
- max_seq_length,
- rand32_gen);
- seq_len = 0;
- num_sent = 0;
- }
-
- } // for (auto sent_index=sent_index_first; ...
- } // if (num_remain_sent > 1) {
- } // for (int doc=0; doc < num_docs; ++doc) {
- } // for (int epoch=0; epoch < num_epochs; ++epoch) {
-
- if (!second) {
- if (verbose) {
- cout << " number of empty documents: " << empty_docs <<
- endl << std::flush;
- cout << " number of documents with one sentence: " <<
- one_sent_docs << endl << std::flush;
- cout << " number of documents with long sentences: " <<
- long_sent_docs << endl << std::flush;
- cout << " will create mapping for " << map_index <<
- " samples" << endl << std::flush;
- }
- assert(maps == NULL);
- assert(num_samples < 0);
- maps = new DocIdx[3*map_index];
- num_samples = static_cast(map_index);
- }
-
- } // for (int iteration=0; iteration < 2; ++iteration) {
-
- // Shuffle.
- // We need a 64 bit random number generator as we might have more
- // than 2 billion samples.
- std::mt19937_64 rand64_gen(seed + 1);
- for (auto i=(num_samples - 1); i > 0; --i) {
- const auto j = static_cast(rand64_gen() % (i + 1));
- const auto i0 = 3 * i;
- const auto j0 = 3 * j;
- // Swap values.
- swap(maps[i0], maps[j0]);
- swap(maps[i0 + 1], maps[j0 + 1]);
- swap(maps[i0 + 2], maps[j0 + 2]);
- }
-
- // Method to deallocate memory.
- py::capsule free_when_done(maps, [](void *mem_) {
- DocIdx *mem = reinterpret_cast(mem_);
- delete[] mem;
- });
-
- // Return the numpy array.
- const auto byte_size = sizeof(DocIdx);
- return py::array(std::vector{num_samples, 3}, // shape
- {3*byte_size, byte_size}, // C-style contiguous strides
- maps, // the data pointer
- free_when_done); // numpy array references
-
-}
-
-
-py::array build_mapping(const py::array_t& docs_,
- const py::array_t& sizes_,
- const int num_epochs,
- const uint64_t max_num_samples,
- const int max_seq_length,
- const double short_seq_prob,
- const int seed,
- const bool verbose,
- const int32_t min_num_sent) {
-
- if (sizes_.size() > std::numeric_limits::max()) {
- if (verbose) {
- cout << " using uint64 for data mapping..." << endl << std::flush;
- }
- return build_mapping_impl(docs_, sizes_, num_epochs,
- max_num_samples, max_seq_length,
- short_seq_prob, seed, verbose,
- min_num_sent);
- } else {
- if (verbose) {
- cout << " using uint32 for data mapping..." << endl << std::flush;
- }
- return build_mapping_impl(docs_, sizes_, num_epochs,
- max_num_samples, max_seq_length,
- short_seq_prob, seed, verbose,
- min_num_sent);
- }
-}
-
-template
-py::array build_blocks_mapping_impl(const py::array_t& docs_,
- const py::array_t& sizes_,
- const py::array_t& titles_sizes_,
- const int32_t num_epochs,
- const uint64_t max_num_samples,
- const int32_t max_seq_length,
- const int32_t seed,
- const bool verbose,
- const bool use_one_sent_blocks) {
- /* Build a mapping of (start-index, end-index, sequence-length) where
- start and end index are the indices of the sentences in the sample
- and sequence-length is the target sequence length.
- */
-
- // Consistency checks.
- assert(num_epochs > 0);
- assert(max_seq_length > 1);
- assert(seed > 0);
-
- // Remove bound checks.
- auto docs = docs_.unchecked<1>();
- auto sizes = sizes_.unchecked<1>();
- auto titles_sizes = titles_sizes_.unchecked<1>();
-
- if (verbose) {
- const auto sent_start_index = docs[0];
- const auto sent_end_index = docs[docs_.shape(0) - 1];
- const auto num_sentences = sent_end_index - sent_start_index;
- cout << " using:" << endl << std::flush;
- cout << " number of documents: " << docs_.shape(0) - 1 <<
- endl << std::flush;
- cout << " sentences range: [" << sent_start_index <<
- ", " << sent_end_index << ")" << endl << std::flush;
- cout << " total number of sentences: " << num_sentences <<
- endl << std::flush;
- cout << " number of epochs: " << num_epochs <<
- endl << std::flush;
- cout << " maximum number of samples: " << max_num_samples <<
- endl << std::flush;
- cout << " maximum sequence length: " << max_seq_length <<
- endl << std::flush;
- cout << " seed: " << seed << endl <<
- std::flush;
- }
-
- // Mapping and its length (1D).
- int64_t num_samples = -1;
- DocIdx* maps = NULL;
-
- // Acceptable number of sentences per block.
- int min_num_sent = 2;
- if (use_one_sent_blocks) {
- min_num_sent = 1;
- }
-
- // Perform two iterations, in the first iteration get the size
- // and allocate memory and in the second iteration populate the map.
- bool second = false;
- for (int32_t iteration=0; iteration<2; ++iteration) {
-
- // Set the flag on second iteration.
- second = (iteration == 1);
-
- // Current map index.
- uint64_t map_index = 0;
-
- uint64_t empty_docs = 0;
- uint64_t one_sent_docs = 0;
- uint64_t long_sent_docs = 0;
- // For each epoch:
- for (int32_t epoch=0; epoch= max_num_samples) {
- if (verbose && (!second)) {
- cout << " reached " << max_num_samples << " samples after "
- << epoch << " epochs ..." << endl << std::flush;
- }
- break;
- }
- // For each document:
- for (int32_t doc=0; doc<(docs.shape(0) - 1); ++doc) {
-
- // Document sentences are in [sent_index_first, sent_index_last)
- const auto sent_index_first = docs[doc];
- const auto sent_index_last = docs[doc + 1];
- const auto target_seq_len = max_seq_length - titles_sizes[doc];
-
- // At the begining of the document previous index is the
- // start index.
- auto prev_start_index = sent_index_first;
-
- // Remaining documents.
- auto num_remain_sent = sent_index_last - sent_index_first;
-
- // Some bookkeeping
- if ((epoch == 0) && (!second)) {
- if (num_remain_sent == 0) {
- ++empty_docs;
- }
- if (num_remain_sent == 1) {
- ++one_sent_docs;
- }
- }
- // Detect documents with long sentences.
- bool contains_long_sentence = false;
- if (num_remain_sent >= min_num_sent) {
- for (auto sent_index=sent_index_first;
- sent_index < sent_index_last; ++sent_index) {
- if (sizes[sent_index] > LONG_SENTENCE_LEN){
- if ((epoch == 0) && (!second)) {
- ++long_sent_docs;
- }
- contains_long_sentence = true;
- break;
- }
- }
- }
- // If we have enough sentences and no long sentences.
- if ((num_remain_sent >= min_num_sent) && (!contains_long_sentence)) {
-
- // Set values.
- auto seq_len = int32_t{0};
- auto num_sent = int32_t{0};
-
- // Loop through sentences.
- for (auto sent_index=sent_index_first;
- sent_index < sent_index_last; ++sent_index) {
-
- // Add the size and number of sentences.
- seq_len += sizes[sent_index];
- ++num_sent;
- --num_remain_sent;
-
- // If we have reached the target length.
- // and there are an acceptable number of sentences left
- // and if we have at least the minimum number of sentences.
- // or if we have reached end of the document.
- if (((seq_len >= target_seq_len) &&
- (num_remain_sent >= min_num_sent) &&
- (num_sent >= min_num_sent) ) || (num_remain_sent == 0)) {
-
- // Populate the map.
- if (second) {
- const auto map_index_0 = 4 * map_index;
- // Each sample has 4 items: the starting sentence index, ending sentence index,
- // the index of the document from which the block comes (used for fetching titles)
- // and the unique id of the block (used for creating block indexes)
-
- maps[map_index_0] = static_cast(prev_start_index);
- maps[map_index_0 + 1] = static_cast(sent_index + 1);
- maps[map_index_0 + 2] = static_cast(doc);
- maps[map_index_0 + 3] = static_cast(block_id);
- }
-
- // Update indices / counters.
- ++map_index;
- ++block_id;
- prev_start_index = sent_index + 1;
- seq_len = 0;
- num_sent = 0;
- }
- } // for (auto sent_index=sent_index_first; ...
- } // if (num_remain_sent > 1) {
- } // for (int doc=0; doc < num_docs; ++doc) {
- } // for (int epoch=0; epoch < num_epochs; ++epoch) {
-
- if (!second) {
- if (verbose) {
- cout << " number of empty documents: " << empty_docs <<
- endl << std::flush;
- cout << " number of documents with one sentence: " <<
- one_sent_docs << endl << std::flush;
- cout << " number of documents with long sentences: " <<
- long_sent_docs << endl << std::flush;
- cout << " will create mapping for " << map_index <<
- " samples" << endl << std::flush;
- }
- assert(maps == NULL);
- assert(num_samples < 0);
- maps = new DocIdx[4*map_index];
- num_samples = static_cast(map_index);
- }
-
- } // for (int iteration=0; iteration < 2; ++iteration) {
-
- // Shuffle.
- // We need a 64 bit random number generator as we might have more
- // than 2 billion samples.
- std::mt19937_64 rand64_gen(seed + 1);
- for (auto i=(num_samples - 1); i > 0; --i) {
- const auto j = static_cast(rand64_gen() % (i + 1));
- const auto i0 = 4 * i;
- const auto j0 = 4 * j;
- // Swap values.
- swap(maps[i0], maps[j0]);
- swap(maps[i0 + 1], maps[j0 + 1]);
- swap(maps[i0 + 2], maps[j0 + 2]);
- swap(maps[i0 + 3], maps[j0 + 3]);
- }
-
- // Method to deallocate memory.
- py::capsule free_when_done(maps, [](void *mem_) {
- DocIdx *mem = reinterpret_cast(mem_);
- delete[] mem;
- });
-
- // Return the numpy array.
- const auto byte_size = sizeof(DocIdx);
- return py::array(std::vector{num_samples, 4}, // shape
- {4*byte_size, byte_size}, // C-style contiguous strides
- maps, // the data pointer
- free_when_done); // numpy array references
-
-}
-
-py::array build_blocks_mapping(const py::array_t& docs_,
- const py::array_t& sizes_,
- const py::array_t& titles_sizes_,
- const int num_epochs,
- const uint64_t max_num_samples,
- const int max_seq_length,
- const int seed,
- const bool verbose,
- const bool use_one_sent_blocks) {
-
- if (sizes_.size() > std::numeric_limits::max()) {
- if (verbose) {
- cout << " using uint64 for data mapping..." << endl << std::flush;
- }
- return build_blocks_mapping_impl(docs_, sizes_, titles_sizes_,
- num_epochs, max_num_samples, max_seq_length, seed, verbose, use_one_sent_blocks);
- } else {
- if (verbose) {
- cout << " using uint32 for data mapping..." << endl << std::flush;
- }
- return build_blocks_mapping_impl(docs_, sizes_, titles_sizes_,
- num_epochs, max_num_samples, max_seq_length, seed, verbose, use_one_sent_blocks);
- }
-}
-
-PYBIND11_MODULE(helpers, m) {
- m.def("build_mapping", &build_mapping);
- m.def("build_blocks_mapping", &build_blocks_mapping);
- m.def("build_sample_idx", &build_sample_idx);
- m.def("build_blending_indices", &build_blending_indices);
-}
diff --git a/megatron/data/indexed_dataset.py b/megatron/data/indexed_dataset.py
deleted file mode 100644
index 4286e69b45..0000000000
--- a/megatron/data/indexed_dataset.py
+++ /dev/null
@@ -1,584 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-#
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-
-
-# copied from fairseq/fairseq/data/indexed_dataset.py
-# Removed IndexedRawTextDataset since it relied on Fairseq dictionary
-# other slight modifications to remove fairseq dependencies
-# Added document index to index file and made it accessible.
-# An empty sentence no longer separates documents.
-
-from functools import lru_cache
-import os
-import shutil
-import struct
-from itertools import accumulate
-
-import numpy as np
-import torch
-from megatron import print_rank_0
-
-
-def __best_fitting_dtype(vocab_size=None):
- if vocab_size is not None and vocab_size < 65500:
- return np.uint16
- else:
- return np.int32
-
-
-def get_available_dataset_impl():
- return ['lazy', 'cached', 'mmap']
-
-
-def infer_dataset_impl(path):
- if IndexedDataset.exists(path):
- with open(index_file_path(path), 'rb') as f:
- magic = f.read(8)
- if magic == IndexedDataset._HDR_MAGIC:
- return 'cached'
- elif magic == MMapIndexedDataset.Index._HDR_MAGIC[:8]:
- return 'mmap'
- else:
- return None
- else:
- print(f"Dataset does not exist: {path}")
- print("Path should be a basename that both .idx and .bin can be appended to get full filenames.")
- return None
-
-
-def make_builder(out_file, impl, vocab_size=None):
- if impl == 'mmap':
- return MMapIndexedDatasetBuilder(out_file, dtype=__best_fitting_dtype(vocab_size))
- else:
- return IndexedDatasetBuilder(out_file)
-
-
-def make_dataset(path, impl, skip_warmup=False):
- if not IndexedDataset.exists(path):
- print(f"Dataset does not exist: {path}")
- print("Path should be a basename that both .idx and .bin can be appended to get full filenames.")
- return None
- if impl == 'infer':
- impl = infer_dataset_impl(path)
- if impl == 'lazy' and IndexedDataset.exists(path):
- return IndexedDataset(path)
- elif impl == 'cached' and IndexedDataset.exists(path):
- return IndexedCachedDataset(path)
- elif impl == 'mmap' and MMapIndexedDataset.exists(path):
- return MMapIndexedDataset(path, skip_warmup)
- print(f"Unknown dataset implementation: {impl}")
- return None
-
-
-def dataset_exists(path, impl):
- if impl == 'mmap':
- return MMapIndexedDataset.exists(path)
- else:
- return IndexedDataset.exists(path)
-
-
-def read_longs(f, n):
- a = np.empty(n, dtype=np.int64)
- f.readinto(a)
- return a
-
-
-def write_longs(f, a):
- f.write(np.array(a, dtype=np.int64))
-
-
-dtypes = {
- 1: np.uint8,
- 2: np.int8,
- 3: np.int16,
- 4: np.int32,
- 5: np.int64,
- 6: np.float32,
- 7: np.double,
- 8: np.uint16
-}
-
-
-def code(dtype):
- for k in dtypes.keys():
- if dtypes[k] == dtype:
- return k
- raise ValueError(dtype)
-
-
-def index_file_path(prefix_path):
- return prefix_path + '.idx'
-
-
-def data_file_path(prefix_path):
- return prefix_path + '.bin'
-
-
-def create_doc_idx(sizes):
- doc_idx = [0]
- for i, s in enumerate(sizes):
- if s == 0:
- doc_idx.append(i + 1)
- return doc_idx
-
-
-class IndexedDataset(torch.utils.data.Dataset):
- """Loader for IndexedDataset"""
- _HDR_MAGIC = b'TNTIDX\x00\x00'
-
- def __init__(self, path):
- super().__init__()
- self.path = path
- self.data_file = None
- self.read_index(path)
-
- def read_index(self, path):
- with open(index_file_path(path), 'rb') as f:
- magic = f.read(8)
- assert magic == self._HDR_MAGIC, (
- 'Index file doesn\'t match expected format. '
- 'Make sure that --dataset-impl is configured properly.'
- )
- version = f.read(8)
- assert struct.unpack('= self._len:
- raise IndexError('index out of range')
-
- def __del__(self):
- if self.data_file:
- self.data_file.close()
-
- # @lru_cache(maxsize=8)
- def __getitem__(self, idx):
- if not self.data_file:
- self.read_data(self.path)
- if isinstance(idx, int):
- i = idx
- self.check_index(i)
- tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]]
- a = np.empty(tensor_size, dtype=self.dtype)
- self.data_file.seek(self.data_offsets[i] * self.element_size)
- self.data_file.readinto(a)
- return a
- elif isinstance(idx, slice):
- start, stop, step = idx.indices(len(self))
- if step != 1:
- raise ValueError("Slices into indexed_dataset must be contiguous")
- sizes = self.sizes[self.dim_offsets[start]:self.dim_offsets[stop]]
- size = sum(sizes)
- a = np.empty(size, dtype=self.dtype)
- self.data_file.seek(self.data_offsets[start] * self.element_size)
- self.data_file.readinto(a)
- offsets = list(accumulate(sizes))
- sents = np.split(a, offsets[:-1])
- return sents
-
- def __len__(self):
- return self._len
-
- def num_tokens(self, index):
- return self.sizes[index]
-
- def size(self, index):
- return self.sizes[index]
-
- @staticmethod
- def exists(path):
- return (
- os.path.exists(index_file_path(path)) and os.path.exists(data_file_path(path))
- )
-
- @property
- def supports_prefetch(self):
- return False # avoid prefetching to save memory
-
-
-class IndexedCachedDataset(IndexedDataset):
-
- def __init__(self, path):
- super().__init__(path)
- self.cache = None
- self.cache_index = {}
-
- @property
- def supports_prefetch(self):
- return True
-
- def prefetch(self, indices):
- if all(i in self.cache_index for i in indices):
- return
- if not self.data_file:
- self.read_data(self.path)
- indices = sorted(set(indices))
- total_size = 0
- for i in indices:
- total_size += self.data_offsets[i + 1] - self.data_offsets[i]
- self.cache = np.empty(total_size, dtype=self.dtype)
- ptx = 0
- self.cache_index.clear()
- for i in indices:
- self.cache_index[i] = ptx
- size = self.data_offsets[i + 1] - self.data_offsets[i]
- a = self.cache[ptx: ptx + size]
- self.data_file.seek(self.data_offsets[i] * self.element_size)
- self.data_file.readinto(a)
- ptx += size
- if self.data_file:
- # close and delete data file after prefetch so we can pickle
- self.data_file.close()
- self.data_file = None
-
- # @lru_cache(maxsize=8)
- def __getitem__(self, idx):
- if isinstance(idx, int):
- i = idx
- self.check_index(i)
- tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]]
- a = np.empty(tensor_size, dtype=self.dtype)
- ptx = self.cache_index[i]
- np.copyto(a, self.cache[ptx: ptx + a.size])
- return a
- elif isinstance(idx, slice):
- # Hack just to make this work, can optimizer later if necessary
- sents = []
- for i in range(*idx.indices(len(self))):
- sents.append(self[i])
- return sents
-
-
-class IndexedDatasetBuilder(object):
- element_sizes = {
- np.uint8: 1,
- np.int8: 1,
- np.int16: 2,
- np.int32: 4,
- np.int64: 8,
- np.float: 4,
- np.double: 8
- }
-
- def __init__(self, out_file, dtype=np.int32):
- self.out_file = open(out_file, 'wb')
- self.dtype = dtype
- self.data_offsets = [0]
- self.dim_offsets = [0]
- self.sizes = []
- self.element_size = self.element_sizes[self.dtype]
- self.doc_idx = [0]
-
- def add_item(self, tensor):
- bytes = self.out_file.write(np.array(tensor.numpy(), dtype=self.dtype))
- self.data_offsets.append(self.data_offsets[-1] + bytes / self.element_size)
- for s in tensor.size():
- self.sizes.append(s)
- self.dim_offsets.append(self.dim_offsets[-1] + len(tensor.size()))
-
- def end_document(self):
- self.doc_idx.append(len(self.sizes))
-
- def merge_file_(self, another_file):
- index = IndexedDataset(another_file)
- assert index.dtype == self.dtype
-
- doc_offset = len(self.sizes)
-
- begin = self.data_offsets[-1]
- for data_offset in index.data_offsets[1:]:
- self.data_offsets.append(begin + data_offset)
- self.sizes.extend(index.sizes)
-
- begin = self.dim_offsets[-1]
- for dim_offset in index.dim_offsets[1:]:
- self.dim_offsets.append(begin + dim_offset)
-
- self.doc_idx.extend((doc_offset + index.doc_idx)[1:])
-
- with open(data_file_path(another_file), 'rb') as f:
- while True:
- data = f.read(1024)
- if data:
- self.out_file.write(data)
- else:
- break
-
- def finalize(self, index_file):
- self.out_file.close()
- index = open(index_file, 'wb')
- index.write(b'TNTIDX\x00\x00')
- index.write(struct.pack(' 0, "Provide the argument --vocab-extra-ids 100 to the script"
-
- def __len__(self):
- return self.samples_mapping.shape[0]
-
- def __getitem__(self, idx):
-
- start_index, end_index, seq_length = self.samples_mapping[idx]
- sample = []
- for index in range(start_index, end_index):
- sample.append(self.indexed_dataset[index])
- # Note that this rng state should be numpy and not python since
- # python randint is inclusive whereas the numpy one is exclusive.
- np_rng = np.random.RandomState(seed=(self.seed + idx))
- return build_training_sample(sample, seq_length,
- self.max_seq_length, # needed for padding
- self.max_seq_length_dec,
- self.vocab_id_list,
- self.vocab_id_to_token_dict,
- self.cls_id, self.sep_id,
- self.mask_id, self.pad_id,
- self.masked_lm_prob, np_rng,
- self.bos_id, self.eos_id,
- self.sentinel_tokens)
-
-
-def build_training_sample(sample, target_seq_length,
- max_seq_length, max_seq_length_dec,
- vocab_id_list, vocab_id_to_token_dict,
- cls_id, sep_id, mask_id, pad_id,
- masked_lm_prob, np_rng, bos_id=None,
- eos_id=None, sentinel_tokens=None):
- """Build training sample.
-
- Arguments:
- sample: A list of sentences in which each sentence is a list token ids.
- target_seq_length: Desired sequence length.
- max_seq_length: Maximum length of the sequence. All values are padded to
- this length.
- vocab_id_list: List of vocabulary ids. Used to pick a random id.
- vocab_id_to_token_dict: A dictionary from vocab ids to text tokens.
- cls_id: Start of example id.
- sep_id: Separator id.
- mask_id: Mask token id.
- pad_id: Padding token id.
- masked_lm_prob: Probability to mask tokens.
- np_rng: Random number genenrator. Note that this rng state should be
- numpy and not python since python randint is inclusive for
- the opper bound whereas the numpy one is exclusive.
- bos_id: start of decoder example id
- eos_id: end of generation id
- sentinel_tokens: unique value to be substituted for every replaced span
- """
-
- assert target_seq_length <= max_seq_length
-
- # flatten sentences into one list
- tokens = [token for sentence in sample for token in sentence]
-
- # Truncate to `target_sequence_length`.
- max_num_tokens = target_seq_length
- truncated = len(tokens) > max_num_tokens
- tokens = tokens[:max_num_tokens]
-
- # Masking.
- max_predictions_per_seq = masked_lm_prob * max_num_tokens
- (tokens, masked_positions, masked_labels, _, masked_spans) = create_masked_lm_predictions(
- tokens, vocab_id_list, vocab_id_to_token_dict, masked_lm_prob,
- cls_id, sep_id, mask_id, max_predictions_per_seq, np_rng,
- max_ngrams=10, geometric_dist=True, masking_style="t5")
-
- # Padding.
- tokens_enc, tokens_dec_in, labels, enc_mask, \
- dec_mask, enc_dec_mask, loss_mask \
- = pad_and_convert_to_numpy(tokens, masked_positions,
- masked_labels, pad_id, max_seq_length,
- max_seq_length_dec, masked_spans,
- bos_id, eos_id, sentinel_tokens)
-
- train_sample = {
- 'text_enc': tokens_enc,
- 'text_dec': tokens_dec_in,
- 'labels': labels,
- 'loss_mask': loss_mask,
- 'truncated': int(truncated),
- 'enc_mask': enc_mask,
- 'dec_mask': dec_mask,
- 'enc_dec_mask': enc_dec_mask,
- }
- return train_sample
-
-
-def pad_and_convert_to_numpy(tokens, masked_positions,
- masked_labels, pad_id,
- max_seq_length, max_seq_length_dec,
- masked_spans=None, bos_id=None,
- eos_id=None, sentinel_tokens=None):
- """Pad sequences and convert them to numpy."""
-
- sentinel_tokens = collections.deque(sentinel_tokens)
- t5_input = []
- (t5_decoder_in, t5_decoder_out) = ([bos_id], [])
- (start_index, end_index) = (0, None)
- for span in masked_spans:
- flag = sentinel_tokens.popleft()
-
- # Append the same tokens in decoder input and output
- t5_decoder_in.append(flag)
- t5_decoder_in.extend(span.label)
- t5_decoder_out.append(flag)
- t5_decoder_out.extend(span.label)
-
- end_index = span.index[0]
- t5_input.extend(tokens[start_index: end_index])
- t5_input.append(flag)
-
- # the next start index is the token after the last span token
- start_index = span.index[-1] + 1
-
- # Add token to the t5_decoder_out
- t5_decoder_out.append(eos_id)
-
- # Add the remaining tokens to the t5 input
- t5_input.extend(tokens[start_index:])
-
- # assert (len(t5_input) - len(masked_spans)) + \
- # (len(t5_decoder_in) - (len(masked_spans) + 1)) == len(tokens)
-
- # Some checks.
-
- # Encoder-side padding mask.
- num_tokens = len(t5_input)
- padding_length = max_seq_length - num_tokens
- assert padding_length >= 0
- assert len(masked_positions) == len(masked_labels)
-
- # Tokens..
- filler = [pad_id] * padding_length
- tokens_enc = np.array(t5_input + filler, dtype=np.int64)
-
- # Decoder-side padding mask.
- num_tokens_dec = len(t5_decoder_in)
- padding_length_dec = max_seq_length_dec - num_tokens_dec
- assert padding_length_dec >= 0
- filler_dec = [pad_id] * padding_length_dec
- tokens_dec_in = np.array(t5_decoder_in + filler_dec, dtype=np.int64)
-
- # Create attention masks
- enc_mask = make_attention_mask(tokens_enc, tokens_enc)
- enc_dec_mask = make_attention_mask(tokens_dec_in, tokens_enc)
- dec_mask = make_attention_mask(tokens_dec_in, tokens_dec_in)
- dec_mask = dec_mask * make_history_mask(tokens_dec_in)
-
- # Labels mask.
- labels = t5_decoder_out + ([-1] * padding_length_dec)
- labels = np.array(labels, dtype=np.int64)
-
- # Loss mask
- loss_mask = ([1] * num_tokens_dec) + ([0] * padding_length_dec)
- loss_mask = np.array(loss_mask, dtype=np.int64)
-
- return tokens_enc, tokens_dec_in, labels, enc_mask, \
- dec_mask, enc_dec_mask, loss_mask
-
-
-def make_attention_mask(source_block, target_block):
- """
- Returns a 2-dimensional (2-D) attention mask
- :param source_block: 1-D array
- :param target_block: 1-D array
- """
- mask = (target_block[None, :] >= 1) * (source_block[:, None] >= 1)
- mask = mask.astype(np.int64)
- # (source_length, target_length)
- return mask
-
-
-def make_attention_mask_3d(source_block, target_block):
- """
- Returns a 3-dimensional (3-D) attention mask
- :param source_block: 1-D array
- :param target_block: 1-D array
- """
- mask = (target_block[:, None, :] >= 1) * (source_block[:, :, None] >= 1)
- # (batch, source_length, target_length)
- # mask = mask.astype(np.int64)
- return mask
-
-
-def make_history_mask(block):
- length = block.shape[0]
- arange = np.arange(length)
- history_mask = (arange[None, ] <= arange[:, None])
- history_mask = history_mask.astype(np.int64)
- return history_mask
-
-
-def make_history_mask_3d(block):
- batch, length = block.shape
- arange = torch.arange(length, device=block.device)
- history_mask = (arange[None, ] <= arange[:, None])[None, ]
- history_mask = history_mask.expand(batch, length, length)
- return history_mask
diff --git a/megatron/data/test/test_indexed_dataset.py b/megatron/data/test/test_indexed_dataset.py
deleted file mode 100644
index 12fec8d819..0000000000
--- a/megatron/data/test/test_indexed_dataset.py
+++ /dev/null
@@ -1,125 +0,0 @@
-# This file isn't really a formal automated test, it's just a place to
-# put some code used during development and manual testing of
-# indexed_dataset.
-
-from megatron.data import indexed_dataset
-from megatron.tokenizer import build_tokenizer
-import argparse
-import os
-import sys
-
-import torch
-
-script_dir = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(os.path.join(script_dir, "../../../"))
-
-
-def test_indexed_dataset(args):
- ds = indexed_dataset.make_dataset(args.data, args.dataset_impl)
- tokenizer = build_tokenizer(args)
- print(len(ds.doc_idx))
- print(len(ds))
- print(ds.doc_idx[-1])
- if ds.supports_prefetch:
- # just prefetch the whole thing in test (so assume it is small)
- ds.prefetch(range(len(ds)))
- if args.count > len(ds.doc_idx) - 1:
- args.count = len(ds.doc_idx) - 1
-
- for i in range(args.count):
- start = ds.doc_idx[i]
- end = ds.doc_idx[i + 1]
- ids = ds[start:end]
- print(f"Document {i}:")
- print("--------------")
- for s in ids:
- assert len(s) > 0
- l = s.data.tolist()
- text = tokenizer.detokenize(l)
- print(text)
- print("---")
-
-
-def test_indexed_dataset_get(args):
- ds = indexed_dataset.make_dataset(args.data, args.dataset_impl)
- tokenizer = build_tokenizer(args)
- size = ds.sizes[0]
- print(f"size: {size}")
- full = ds.get(0)
- print(full)
- # print(tokenizer.detokenize(full.data.tolist()))
- print("---")
- end = ds.get(0, offset=size - 10)
- print(end)
- # print(tokenizer.detokenize(end.data.tolist()))
-
- start = ds.get(0, length=10)
- print(start)
- # print(tokenizer.detokenize(start.data.tolist()))
-
- part = ds.get(0, offset=2, length=8)
- print(part)
- # print(tokenizer.detokenize(part.data.tolist()))
-
-# def test_albert_dataset(args):
-# # tokenizer = FullBertTokenizer(args.vocab, do_lower_case=True)
-# # idataset = indexed_dataset.make_dataset(args.data, args.dataset_impl)
-# # ds = AlbertDataset(idataset, tokenizer)
-# ds = AlbertDataset.from_paths(args.vocab, args.data, args.dataset_impl,
-# args.epochs, args.max_num_samples,
-# args.masked_lm_prob, args.seq_length,
-# args.short_seq_prob, args.seed)
-# truncated = 0
-# total = 0
-# for i, s in enumerate(ds):
-# ids = s['text']
-# tokens = ds.tokenizer.convert_ids_to_tokens(ids)
-# print(tokens)
-# if i >= args.count-1:
-# exit()
-
-
-def main():
- parser = argparse.ArgumentParser()
- parser.add_argument('--data', type=str, help='prefix to data files')
- parser.add_argument('--dataset-impl', type=str, default='infer',
- choices=['lazy', 'cached', 'mmap', 'infer'])
- parser.add_argument('--count', type=int, default=10,
- help='Number of samples/documents to print')
-
- group = parser.add_argument_group(title='tokenizer')
- group.add_argument('--tokenizer-type', type=str, required=True,
- choices=['BertWordPieceLowerCase',
- 'GPT2BPETokenizer'],
- help='What type of tokenizer to use.')
- group.add_argument('--vocab-file', type=str, default=None,
- help='Path to the vocab file')
- group.add_argument('--merge-file', type=str, default=None,
- help='Path to the BPE merge file (if necessary).')
-
- parser.add_argument('--epochs', type=int, default=5,
- help='Number of epochs to plan for')
- parser.add_argument('--max-num-samples', type=int, default=None,
- help='Maximum number of samples to plan for')
- parser.add_argument('--masked-lm-prob', type=float, default=0.15,
- help='probability of masking tokens')
- parser.add_argument('--seq-length', type=int, default=512,
- help='maximum sequence length')
- parser.add_argument('--short-seq-prob', type=float, default=0.1,
- help='probability of creating a short sequence')
- parser.add_argument('--seed', type=int, default=1234,
- help='random seed')
- args = parser.parse_args()
- args.rank = 0
- args.make_vocab_size_divisible_by = 128
- args.tensor_model_parallel_size = 1
-
- if args.dataset_impl == "infer":
- args.dataset_impl = indexed_dataset.infer_dataset_impl(args.data)
-
-# test_albert_dataset(args)
- test_indexed_dataset_get(args)
-
-
-if __name__ == "__main__":
- main()
diff --git a/megatron/data/test/test_preprocess_data.sh b/megatron/data/test/test_preprocess_data.sh
deleted file mode 100755
index d121c85958..0000000000
--- a/megatron/data/test/test_preprocess_data.sh
+++ /dev/null
@@ -1,10 +0,0 @@
-#!/bin/bash
-
-IMPL=cached
-python ../preprocess_data.py \
- --input test_samples.json \
- --vocab vocab.txt \
- --dataset-impl ${IMPL} \
- --output-prefix test_samples_${IMPL} \
- --workers 1 \
- --log-interval 2
diff --git a/megatron/fused_kernels/scaled_masked_softmax.cpp b/megatron/fused_kernels/scaled_masked_softmax.cpp
deleted file mode 100644
index 4c8a8c2ee3..0000000000
--- a/megatron/fused_kernels/scaled_masked_softmax.cpp
+++ /dev/null
@@ -1,83 +0,0 @@
-/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. */
-
-#include
-#include
-#include
-
-namespace multihead_attn {
-namespace fused_softmax {
-namespace scaled_masked_softmax {
-
-torch::Tensor fwd_cuda(
- torch::Tensor const& input,
- torch::Tensor const& mask,
- float scale_factor);
-
-torch::Tensor bwd_cuda(
- torch::Tensor const& output_grads,
- torch::Tensor const& softmax_results,
- float scale_factor);
-
-int get_batch_per_block_cuda(
- int query_seq_len,
- int key_seq_len,
- int batches,
- int attn_heads);
-
-torch::Tensor fwd(
- torch::Tensor const& input,
- torch::Tensor const& mask,
- float scale_factor) {
- AT_ASSERTM(input.dim() == 4, "expected 4D tensor");
- AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) ||
- (input.scalar_type() == at::ScalarType::BFloat16),
- "Only fp16 and bf16 are supported");
- AT_ASSERTM(mask.dim() == 4, "expected 4D tensor");
-
- return fwd_cuda(input, mask, scale_factor);
-}
-
-torch::Tensor bwd(
- torch::Tensor const& output_grads,
- torch::Tensor const& softmax_results,
- float scale_factor) {
-
- AT_ASSERTM(output_grads.dim() == 4, "expected 3D tensor");
- AT_ASSERTM(softmax_results.dim() == 4, "expected 3D tensor");
-
- AT_ASSERTM((output_grads.scalar_type() == at::ScalarType::Half) ||
- (output_grads.scalar_type() == at::ScalarType::BFloat16),
- "Only fp16 and bf16 are supported");
- AT_ASSERTM((softmax_results.scalar_type() == at::ScalarType::Half) ||
- (softmax_results.scalar_type() == at::ScalarType::BFloat16),
- "Only fp16 and bf16 are supported");
-
- return bwd_cuda(output_grads, softmax_results, scale_factor);
-}
-
-int get_batch_per_block(
- int query_seq_len,
- int key_seq_len,
- int batches,
- int attn_heads) {
- return get_batch_per_block_cuda(query_seq_len, key_seq_len, batches, attn_heads);
-}
-
-} // end namespace scaled_masked_softmax
-} // end namespace fused_softmax
-} // end namespace multihead_attn
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
- m.def("forward",
- &multihead_attn::fused_softmax::scaled_masked_softmax::fwd,
- "Self Multihead Attention scaled, time masked softmax -- Forward.");
-
- m.def("backward",
- &multihead_attn::fused_softmax::scaled_masked_softmax::bwd,
- "Self Multihead Attention scaled, time masked softmax -- Backward.");
-
- m.def("get_batch_per_block",
- &multihead_attn::fused_softmax::scaled_masked_softmax::get_batch_per_block,
- "Return Batch per block size."
- );
-}
diff --git a/megatron/fused_kernels/scaled_masked_softmax.h b/megatron/fused_kernels/scaled_masked_softmax.h
deleted file mode 100644
index 21ebbd5228..0000000000
--- a/megatron/fused_kernels/scaled_masked_softmax.h
+++ /dev/null
@@ -1,710 +0,0 @@
-/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. */
-
-#pragma once
-
-#include
-#include
-#include
-#include
-#include
-#include
-#include
-
-namespace {
-
-template
-__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);
-
-template <>
-__device__ __inline__ void copy_vector(c10::BFloat16 *dst, const c10::BFloat16 *src) { *dst = *src; }
-
-template <>
-__device__ __inline__ void copy_vector(c10::BFloat16 *dst, const c10::BFloat16 *src) { *((float2*) dst) = *((float2*) src); }
-
-template <>
-__device__ __inline__ void copy_vector(c10::Half *dst, const c10::Half *src) { *dst = *src; }
-
-template <>
-__device__ __inline__ void copy_vector(c10::Half *dst, const c10::Half *src) { *((float2*) dst) = *((float2*) src); }
-
-template <>
-__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) { *dst = *src; }
-
-template <>
-__device__ __inline__ void copy_vector(uint8_t *dst, const uint8_t *src) {*((half2*) dst) = *((half2*) src); }
-
-int log2_ceil(int value) {
- int log2_value = 0;
- while ((1 << log2_value) < value) ++log2_value;
- return log2_value;
-}
-
-template
-struct Add {
- __device__ __forceinline__ T operator()(T a, T b) const {
- return a + b;
- }
-};
-
-template
-struct Max {
- __device__ __forceinline__ T operator()(T a, T b) const {
- return a < b ? b : a;
- }
-};
-
-template
-__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff)
-{
-#if CUDA_VERSION >= 9000
- return __shfl_xor_sync(mask, value, laneMask, width);
-#else
- return __shfl_xor(value, laneMask, width);
-#endif
-}
-
-template class ReduceOp>
-__device__ __forceinline__ void warp_reduce(acc_t* sum) {
- ReduceOp r;
- #pragma unroll
- for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);
- sum[i] = r(sum[i], b);
- }
- }
-}
-
-
-/*
- * Extended softmax (from native aten pytorch) with following additional features
- * 1) input scaling
- */
-template
-__global__ void scaled_softmax_warp_forward(
- output_t *dst,
- const input_t *src,
- const acc_t scale,
- int micro_batch_size,
- int element_count)
-{
- // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
- // warp_size of method warp_softmax_forward_kernel.
- constexpr int next_power_of_two = 1 << log2_elements;
- constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
- constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
- constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
- constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;
-
- // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )
- // gridDim/blockIdx = (seq_len, attn_heads, batches)
- int first_batch = (blockDim.y * (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z))+ threadIdx.y) * WARP_BATCH;
-
- // micro_batch_size might not be a multiple of WARP_BATCH. Check how
- // many batches have to computed within this WARP.
- int local_batches = micro_batch_size - first_batch;
- if (local_batches > WARP_BATCH)
- local_batches = WARP_BATCH;
-
- // there might be multiple batches per warp. compute the index within the batch
- int local_idx = threadIdx.x;
-
- src += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
- dst += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
-
- // load data from global memory
- acc_t elements[WARP_BATCH][WARP_ITERATIONS];
- input_t temp_data[ELEMENTS_PER_LDG_STG];
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- int batch_element_count = (i >= local_batches) ? 0 : element_count;
-
- #pragma unroll
- for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) {
- int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
-
- if (element_index < batch_element_count) {
- int itr_idx = i*element_count+it*WARP_SIZE;
- copy_vector(temp_data, src + itr_idx);
-
- #pragma unroll
- for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
- elements[i][it + element] = (acc_t)temp_data[element] * scale;
- }
- } else {
- #pragma unroll
- for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
- elements[i][it + element] = -std::numeric_limits::infinity();
- }
- }
- }
- }
-
- // compute max_value
- acc_t max_value[WARP_BATCH];
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- max_value[i] = elements[i][0];
- #pragma unroll
- for (int it = 1; it < WARP_ITERATIONS; ++it) {
- max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
- }
- }
- warp_reduce(max_value);
-
- acc_t sum[WARP_BATCH] { 0.0f };
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- #pragma unroll
- for (int it = 0; it < WARP_ITERATIONS; ++it) {
- elements[i][it] = std::exp((elements[i][it] - max_value[i]));
- sum[i] += elements[i][it];
- }
- }
- warp_reduce(sum);
-
- // store result
- output_t out[ELEMENTS_PER_LDG_STG];
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- if (i >= local_batches)
- break;
- #pragma unroll
- for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) {
- int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
- if (element_index < element_count) {
- #pragma unroll
- for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
- out[element] = elements[i][it + element] / sum[i];
- }
- copy_vector(dst + i * element_count + it * WARP_SIZE, out);
- } else {
- break;
- }
- }
- }
-}
-
-
-/*
- * Extended softmax (from native aten pytorch) with following additional features
- * 1) input scaling
- * 2) Explicit masking
- */
-template
-__global__ void scaled_masked_softmax_warp_forward(
- output_t *dst,
- const input_t *src,
- const uint8_t *mask,
- const acc_t scale,
- int micro_batch_size,
- int element_count,
- int pad_batches)
-{
- // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
- // warp_size of method warp_softmax_forward_kernel.
- constexpr int next_power_of_two = 1 << log2_elements;
- constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
- constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
- constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
- constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;
-
- // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )
- // gridDim/blockIdx = (seq_len, attn_heads, batches)
- int first_batch = (blockDim.y * (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z))+ threadIdx.y) * WARP_BATCH;
- int pad_first_batch = 0;
- if (pad_batches != 1) { // bert style
- pad_first_batch = (blockDim.y * (blockIdx.x + gridDim.x * blockIdx.z) + threadIdx.y) * WARP_BATCH;
- } else { // gpt2 style
- pad_first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
- }
-
- // micro_batch_size might not be a multiple of WARP_BATCH. Check how
- // many batches have to computed within this WARP.
- int local_batches = micro_batch_size - first_batch;
- if (local_batches > WARP_BATCH)
- local_batches = WARP_BATCH;
-
- // there might be multiple batches per warp. compute the index within the batch
- int local_idx = threadIdx.x;
-
- src += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
- dst += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
- mask += pad_first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
-
- // load data from global memory
- acc_t elements[WARP_BATCH][WARP_ITERATIONS];
- input_t temp_data[ELEMENTS_PER_LDG_STG];
- uint8_t temp_mask[ELEMENTS_PER_LDG_STG];
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- int batch_element_count = (i >= local_batches) ? 0 : element_count;
-
- #pragma unroll
- for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) {
- int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
-
- if (element_index < batch_element_count) {
- int itr_idx = i*element_count+it*WARP_SIZE;
- copy_vector(temp_data, src + itr_idx);
- copy_vector(temp_mask, mask + itr_idx);
-
- #pragma unroll
- for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
- if (temp_mask[element] != 1) {
- elements[i][it + element] = (acc_t)temp_data[element] * scale;
- } else {
- elements[i][it + element] = -10000.0;
- }
- }
- } else {
- #pragma unroll
- for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
- elements[i][it + element] = -std::numeric_limits::infinity();
- }
- }
- }
- }
-
- // compute max_value
- acc_t max_value[WARP_BATCH];
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- max_value[i] = elements[i][0];
- #pragma unroll
- for (int it = 1; it < WARP_ITERATIONS; ++it) {
- max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
- }
- }
- warp_reduce(max_value);
-
- // compute scale value to account for full mask
- acc_t scale_value[WARP_BATCH];
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- scale_value[i] = (max_value[i] == -10000.0) ? 0.0 : 1.0;
- }
-
- acc_t sum[WARP_BATCH] { 0.0f };
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- #pragma unroll
- for (int it = 0; it < WARP_ITERATIONS; ++it) {
- elements[i][it] = std::exp((elements[i][it] - max_value[i]));
- sum[i] += elements[i][it];
- }
- }
- warp_reduce(sum);
-
- // store result
- output_t out[ELEMENTS_PER_LDG_STG];
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- if (i >= local_batches)
- break;
- #pragma unroll
- for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) {
- int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
- if (element_index < element_count) {
- #pragma unroll
- for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
- out[element] = elements[i][it + element] * scale_value[i] / sum[i];
- }
- copy_vector(dst + i * element_count + it * WARP_SIZE, out);
- } else {
- break;
- }
- }
- }
-}
-
-template
-__global__ void scaled_masked_softmax_warp_backward(
- output_t *gradInput,
- input_t *grad,
- const input_t *output,
- acc_t scale,
- int micro_batch_size,
- int element_count)
-{
- // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
- // warp_size of method warp_softmax_backward_kernel.
- constexpr int next_power_of_two = 1 << log2_elements;
- constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
- constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
- constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
- constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;
-
- // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )
- // gridDim/blockIdx = (seq_len, attn_heads, batches)
- int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
-
- // micro_batch_size might not be a multiple of WARP_BATCH. Check how
- // many batches have to computed within this WARP.
- int local_batches = micro_batch_size - first_batch;
- if (local_batches > WARP_BATCH)
- local_batches = WARP_BATCH;
-
- // there might be multiple batches per warp. compute the index within the batch
- int local_idx = threadIdx.x;
-
- // the first element to process by the current thread
- int thread_offset = first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
- grad += thread_offset;
- output += thread_offset;
- gradInput += thread_offset;
-
- // load data from global memory
- acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS] { 0.0f };
- acc_t output_reg[WARP_BATCH][WARP_ITERATIONS] { 0.0f };
- input_t temp_grad[ELEMENTS_PER_LDG_STG];
- input_t temp_output[ELEMENTS_PER_LDG_STG];
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- int batch_element_count = (i >= local_batches) ? 0 : element_count;
-
- #pragma unroll
- for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) {
- int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
- if (element_index < batch_element_count) {
- copy_vector(temp_grad, grad + i * element_count + it * WARP_SIZE);
- copy_vector(temp_output, output + i * element_count + it * WARP_SIZE);
-
- #pragma unroll
- for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
- output_reg[i][it + element] = (acc_t)temp_output[element];
- }
- #pragma unroll
- for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
- grad_reg[i][it + element] = (acc_t)temp_grad[element] * output_reg[i][it + element];
- }
- }
- }
- }
-
- acc_t sum[WARP_BATCH];
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- sum[i] = grad_reg[i][0];
- #pragma unroll
- for (int it = 1; it < WARP_ITERATIONS; ++it) {
- sum[i] += grad_reg[i][it];
- }
- }
- warp_reduce(sum);
-
- // store result
- #pragma unroll
- for (int i = 0; i < WARP_BATCH; ++i) {
- if (i >= local_batches)
- break;
- #pragma unroll
- for (int it = 0; it < WARP_ITERATIONS; it+=ELEMENTS_PER_LDG_STG) {
- int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
- if (element_index < element_count) {
- // compute gradients
- output_t out[ELEMENTS_PER_LDG_STG];
- #pragma unroll
- for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
- out[element] = (output_t)(scale * (grad_reg[i][it + element] - output_reg[i][it + element] * sum[i]));
- }
- copy_vector(gradInput + i * element_count + it * WARP_SIZE, out);
- }
- }
- }
-}
-} // end of anonymous namespace
-
-int get_batch_per_block(int query_seq_len, int key_seq_len, int batches, int attn_heads){
- int log2_elements = log2_ceil(key_seq_len);
- const int next_power_of_two = 1 << log2_elements;
-
- int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
- int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
-
- constexpr int threads_per_block = 128;
- int warps_per_block = (threads_per_block / warp_size);
- int batches_per_block = warps_per_block * batches_per_warp;
-
- return batches_per_block;
-}
-
-template
-void dispatch_scaled_softmax_forward(
- output_t *dst,
- const input_t *src,
- const input_t scale,
- int query_seq_len,
- int key_seq_len,
- int batches,
- int attn_heads)
-{
- TORCH_INTERNAL_ASSERT(key_seq_len >= 0 && key_seq_len <= 4096 );
- if (key_seq_len == 0) {
- return;
- } else {
- int log2_elements = log2_ceil(key_seq_len);
- const int next_power_of_two = 1 << log2_elements;
- int batch_count = batches * attn_heads * query_seq_len;
-
- // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_forward.
- int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
-
- // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_forward.
- int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
-
- // use 128 threads per block to maximimize gpu utilization
- constexpr int threads_per_block = 128;
-
- int warps_per_block = (threads_per_block / warp_size);
- int batches_per_block = warps_per_block * batches_per_warp;
- TORCH_INTERNAL_ASSERT(query_seq_len%batches_per_block == 0);
- dim3 blocks(query_seq_len/batches_per_block, attn_heads, batches);
- dim3 threads(warp_size, warps_per_block, 1);
- // Launch code would be more elegant if C++ supported FOR CONSTEXPR
- switch (log2_elements) {
- case 0: // 1
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 1: // 2
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 2: // 4
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 3: // 8
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 4: // 16
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 5: // 32
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 6: // 64
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 7: // 128
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 8: // 256
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 9: // 512
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 10: // 1024
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 11: // 2048
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- case 12: // 4096
- scaled_softmax_warp_forward
- <<>>(dst, src, scale, batch_count, key_seq_len);
- break;
- default:
- break;
- }
- }
-}
-
-template
-void dispatch_scaled_masked_softmax_forward(
- output_t *dst,
- const input_t *src,
- const uint8_t *mask,
- const input_t scale,
- int query_seq_len,
- int key_seq_len,
- int batches,
- int attn_heads,
- int pad_batches)
-{
- TORCH_INTERNAL_ASSERT(key_seq_len >= 0 && key_seq_len <= 4096 );
- if (key_seq_len == 0) {
- return;
- } else {
- int log2_elements = log2_ceil(key_seq_len);
- const int next_power_of_two = 1 << log2_elements;
- int batch_count = batches * attn_heads * query_seq_len;
-
- // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_forward.
- int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
-
- // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_forward.
- int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
-
- // use 128 threads per block to maximimize gpu utilization
- constexpr int threads_per_block = 128;
-
- int warps_per_block = (threads_per_block / warp_size);
- int batches_per_block = warps_per_block * batches_per_warp;
- TORCH_INTERNAL_ASSERT(query_seq_len%batches_per_block == 0);
- dim3 blocks(query_seq_len/batches_per_block, attn_heads, batches);
- dim3 threads(warp_size, warps_per_block, 1);
- // Launch code would be more elegant if C++ supported FOR CONSTEXPR
- switch (log2_elements) {
- case 0: // 1
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 1: // 2
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 2: // 4
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 3: // 8
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 4: // 16
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 5: // 32
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 6: // 64
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 7: // 128
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 8: // 256
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 9: // 512
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 10: // 1024
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 11: // 2048
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- case 12: // 4096
- scaled_masked_softmax_warp_forward
- <<>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
- break;
- default:
- break;
- }
- }
-}
-
-template
-void dispatch_scaled_masked_softmax_backward(
- output_t *grad_input,
- input_t *grad,
- const input_t *output,
- const acc_t scale,
- int query_seq_len,
- int key_seq_len,
- int batches,
- int attn_heads)
-{
- TORCH_INTERNAL_ASSERT( key_seq_len >= 0 && key_seq_len <= 4096 );
- if (key_seq_len == 0) {
- return;
- } else {
- int log2_elements = log2_ceil(key_seq_len);
- const int next_power_of_two = 1 << log2_elements;
- int batch_count = batches * attn_heads * query_seq_len;
-
- // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_backward.
- int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
-
- // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_backward.
- int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
-
- // use 128 threads per block to maximimize gpu utilization
- constexpr int threads_per_block = 128;
-
- int warps_per_block = (threads_per_block / warp_size);
- int batches_per_block = warps_per_block * batches_per_warp;
- int blocks = batch_count/batches_per_block;
- dim3 threads(warp_size, warps_per_block, 1);
- // Launch code would be more elegant if C++ supported FOR CONSTEXPR
- switch (log2_elements) {
- case 0: // 1
- scaled_masked_softmax_warp_backward
- <<>>(grad_input, grad, output, scale, batch_count, key_seq_len);
- break;
- case 1: // 2
- scaled_masked_softmax_warp_backward
- <<>>(grad_input, grad, output, scale, batch_count, key_seq_len);
- break;
- case 2: // 4
- scaled_masked_softmax_warp_backward
- <<>>(grad_input, grad, output, scale, batch_count, key_seq_len);
- break;
- case 3: // 8
- scaled_masked_softmax_warp_backward