GPU optimized techniques for training transformer models at-scale
+ +[](https://docs.nvidia.com/megatron-core/developer-guide/latest/index.html) +[](./setup.py) +[](./LICENSE) + +
+
+# 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.
+
+# Table of Contents
+
+- [Megatron-LM \& Megatron-Core](#megatron-lm--megatron-core)
+- [Latest News](#latest-news)
+- [Table of 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)
+- [Training](#training)
+ - [Data Preprocessing](#data-preprocessing)
+ - [BERT Pretraining](#bert-pretraining)
+ - [GPT Pretraining](#gpt-pretraining)
+ - [T5 Pretraining](#t5-pretraining)
+ - [Distributed Pretraining](#distributed-pretraining)
+ - [Activation Checkpointing and Recomputation](#activation-checkpointing-and-recomputation)
+ - [Distributed Optimizer](#distributed-optimizer)
+ - [FlashAttention](#flashattention)
+ - [GPT-3 Example](#gpt-3-example)
+ - [Retro and InstructRetro](#retro-and-instructretro)
+ - [Mamba-based Language Models](#mamba-based-language-models)
+ - [Mixture of Experts](#mixture-of-experts)
+- [Evaluation and Tasks](#evaluation-and-tasks)
+ - [GPT Text Generation](#gpt-text-generation)
+ - [Detoxify GPT via Self-generation](#detoxify-gpt-via-self-generation)
+ - [GPT Evaluation](#gpt-evaluation)
+ - [WikiText Perplexity Evaluation](#wikitext-perplexity-evaluation)
+ - [LAMBADA Cloze Accuracy](#lambada-cloze-accuracy)
+ - [BERT Task Evaluation](#bert-task-evaluation)
+ - [RACE Evaluation](#race-evaluation)
+ - [MNLI Evaluation](#mnli-evaluation)
+ - [Llama-2 Inference and Finetuning](#llama-2-inference-and-finetuning)
+- [Model Optimization and Deployment](#model-optimization-and-deployment)
+ - [Quantization and TensorRT-LLM Deployment](#quantization-and-tensorrt-llm-deployment)
+- [Datasets](#datasets)
+ - [Collecting Wikipedia Training Data](#collecting-wikipedia-training-data)
+ - [Collecting GPT Webtext Data](#collecting-gpt-webtext-data)
+- [Reproducibility](#reproducibility)
+- [Checkpoint conversion](#checkpoint-conversion)
+ - [Model class conversion](#model-class-conversion)
+ - [Checkpoint format conversion](#checkpoint-format-conversion)
+- [Projects Using Megatron](#projects-using-megatron)
+
+# Megatron Overview
+This repository comprises two essential components: **Megatron-LM** and **Megatron-Core**. Megatron-LM serves as a research-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%.
+
+
-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)
-* [Evaluating Parameter Efficient Learning for Generation](https://aclanthology.org/2022.emnlp-main.319.pdf)
-
-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 specific 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.
-
-
-
-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 eliminated 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)
- * [Setup](#setup)
- * [Downloading Checkpoints](#downloading-checkpoints)
- * [Usage](#usage)
- * [Training](#training)
- * [Data Preprocessing](#data-preprocessing)
- * [BERT Pretraining](#bert-pretraining)
- * [GPT Pretraining](#gpt-pretraining)
- * [T5 Pretraining](#t5-pretraining)
- * [Distributed Pretraining](#distributed-pretraining)
- * [Activation Checkpointing and Recomputation](#activation-checkpointing-and-recomputation)
- * [Distributed Optimizer](#distributed-optimizer)
- * [FlashAttention](#flashattention)
- * [GPT-3 Example](#gpt-3-example)
- * [Retro](#retro)
- * [Evaluation and Tasks](#evaluation-and-tasks)
- * [GPT Text Generation](#gpt-text-generation)
- * [GPT Evaluation](#gpt-evaluation)
- * [WikiText Perplexity Evaluation](#wikitext-perplexity-evaluation)
- * [LAMBADA Cloze Accuracy](#lambada-cloze-accuracy)
- * [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)
# 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.
@@ -70,7 +100,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:
@@ -92,7 +122,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
@@ -139,27 +169,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.
@@ -169,19 +200,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.
@@ -189,13 +220,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
@@ -212,6 +245,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.
@@ -227,23 +262,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 incremental 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.
+## Mixture of Experts
+MoE (Mixture of Experts) is a powerful LLM architecture implemented in the Megatron-Core framework, designed to enhance the efficiency and scalability of large language models. It leverages **Expert Parallelism**, allowing multiple experts to be distributed across different workers, where each worker processes distinct batches of training samples. This method significantly increases computational throughput, enabling models to achieve high performance metrics, such as 47% MFU during BF16 training for 8x7B on H100.
+
+Key Features of MoE:
+- **Parallelism Techniques**: MoE combines various parallelism strategies, including Expert Parallelism, Data Parallelism, Tensor Parallelism, Sequence Paralleism, Pipeline Parallelism, and Context Parallelism. This combination allows for handling larger model variants effectively.
+- **Router and Load Balancing**: The system employs advanced routing mechanisms like the Top-K router and utilizes load balancing algorithms to optimize token distribution among experts.
+- **Performance Optimizations**: Techniques such as GroupedGEMM and FP8 training enhance the efficiency of MoE models, particularly when multiple experts are involved.
+- **Token Dispatch Mechanism**: MoE supports both dropless and token drop strategies to manage token distribution effectively across experts.
+
+For a comprehensive overview of MoE training configurations and optimizations, please refer to the detailed README located at [megatron/core/transformer/moe/README.md](./megatron/core/transformer/moe/README.md).
+
# Evaluation and Tasks
We provide several command line arguments, detailed in the scripts listed below, to handle various zero-shot and fine-tuned downstream tasks. However, you can also finetune your model from a pretrained checkpoint on other corpora as desired. To do so, simply add the `--finetune` flag and adjust the input files and training parameters within the original training script. The iteration count will be reset to zero, and the optimizer and internal state will be reinitialized. If the fine-tuning is interrupted for any reason, be sure to remove the `--finetune` flag before continuing, otherwise the training will start again from the beginning.
@@ -332,7 +390,7 @@ We provide several command line arguments, detailed in the scripts listed below,
Because evaluation requires substantially less memory than training, it may be advantageous to merge a model trained in parallel for use on fewer GPUs in downstream tasks. The following script accomplishes this. This example reads in a GPT model with 4-way tensor and 4-way pipeline model parallelism and writes out a model with 2-way tensor and 2-way pipeline model parallelism.
-python tools/checkpoint/util.py \
+python tools/checkpoint/convert.py \
--model-type GPT \
--load-dir checkpoints/gpt3_tp4_pp4 \
--save-dir checkpoints/gpt3_tp2_pp2 \
@@ -345,7 +403,7 @@ Several downstream tasks are described for both GPT and BERT models below. They
## GPT Text Generation
-We have included a simple REST server to use for text generation in `tools/run_text_generation_server.py`. You run it much like you would start a pretraining job, specifying an appropriate pretrained checkpoint. There are also few optional parameters: `temperature`, `top-k`and `top-p`. See `--help` or the source file for more information. See [examples/run_text_generation_server_345M.sh](examples/run_text_generation_server_345M.sh) for an example of how to run the server.
+We have included a simple REST server to use for text generation in `tools/run_text_generation_server.py`. You run it much like you would start a pretraining job, specifying an appropriate pretrained checkpoint. There are also few optional parameters: `temperature`, `top-k`and `top-p`. See `--help` or the source file for more information. See [examples/inference/run_text_generation_server_345M.sh](examples/inference/run_text_generation_server_345M.sh) for an example of how to run the server.
Once the server is running you can use `tools/text_generation_cli.py` to query it, it takes one argument which is the host the server is running on.
@@ -359,12 +417,12 @@ You can also use CURL or any other tools to query the server directly:
curl 'http://localhost:5000/api' -X 'PUT' -H 'Content-Type: application/json; charset=UTF-8' -d '{"prompts":["Hello world"], "tokens_to_generate":1}'
-See [megatron/text_generation_server.py](megatron/text_generation_server.py) for more API options.
+See [megatron/inference/text_generation_server.py](megatron/inference/text_generation_server.py) for more API options.
### Detoxify GPT via Self-generation
-We include an example in `examples/detxoify_lm/` to detoxify language models by leveraging the generative power of language models.
+We include an example in `examples/academic_paper_scripts/detxoify_lm/` to detoxify language models by leveraging the generative power of language models.
-See [examples/detxoify_lm/README.md](examples/detxoify_lm/README.md) for step-by-step tutorials on how to perform domain-adaptive training and detoxify LM using self-generated corpus.
+See [examples/academic_paper_scripts/detxoify_lm/README.md](examples/academic_paper_scripts/detxoify_lm/README.md) for step-by-step tutorials on how to perform domain-adaptive training and detoxify LM using self-generated corpus.
## GPT Evaluation
@@ -407,7 +465,7 @@ python tasks/main.py \
### LAMBADA Cloze Accuracy
To compute LAMBADA cloze accuracy (the accuracy of predicting the last token given the preceding tokens) we utilize a detokenized, processed version of the [LAMBADA dataset](https://github.com/cybertronai/bflm/blob/master/lambada_test.jsonl).
-We use the following command to run LAMBADA evaluation on a 345M parameter model. Note that the `--strict-lambada` flag should be used to require whole word matching. Make that `lambada` is part of the file path.
+We use the following command to run LAMBADA evaluation on a 345M parameter model. Note that the `--strict-lambada` flag should be used to require whole word matching. Ensure that `lambada` is part of the file path.
TASK="LAMBADA"
@@ -503,7 +561,13 @@ python tasks/main.py \
The Llama-2 [family of models](https://ai.meta.com/llama/) are an open-source set of pretrained & finetuned (for chat) models that have achieved strong results across a wide set of benchmarks. At the time of release, Llama-2 models achieved among the best results for open-source models, and were competitive with the closed-source GPT-3.5 model (see https://arxiv.org/pdf/2307.09288.pdf).
-The Llama-2 checkpoints can be loaded into Megatron for inference and finetuning. See documentation [here](docs/llama2.md).
+The Llama-2 checkpoints can be loaded into Megatron for inference and finetuning. See documentation [here](docs/llama_mistral.md).
+
+# Model Optimization and Deployment
+Megatron-Core (MCore) `GPTModel` family supports advanced quantization algorithms and high-performance inference through TensorRT-LLM.
+
+## Quantization and TensorRT-LLM Deployment
+See [Megatron Model Optimization and Deployment](examples/inference/quantization/README.md) for `llama2` and `nemotron3` examples.
# Datasets
We do not host any datasets for GPT or BERT training, however, we detail their collection so that our results may be reproduced.
@@ -511,16 +575,93 @@ We do not host any datasets for GPT or BERT training, however, we detail their c
## Collecting Wikipedia Training Data
We recommend following the Wikipedia data extraction process specified by Google research: "the recommended pre-processing is to download [the latest dump](https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2), extract the text with [WikiExtractor.py](https://github.com/attardi/wikiextractor), and then apply any necessary cleanup to convert it into plain text."
-We recommend using the `--json` argument when using WikiExtractor, which will dump the Wikipedia data into loose json format (one json per line), making it more manageable on the file system and also readily consumable by our codebase. We recommend further preprocessing this json dataset by nltk punctuation standardization. For BERT training, use the `--split-sentences` flag to `preprocess_data.py` as described [above](#data-preprocessing) to include sentence breaks in the produced index. If you'd like to use Wikipedia data for GPT training you should still clean it with nltk/spacy/ftfy, but do not use the `--split-sentences` flag.
+We recommend using the `--json` argument when using WikiExtractor, which will dump the Wikipedia data into loose json format (one json object per line), making it more manageable on the file system and also readily consumable by our codebase. We recommend further preprocessing this json dataset with nltk punctuation standardization. For BERT training, use the `--split-sentences` flag to `preprocess_data.py` as described [above](#data-preprocessing) to include sentence breaks in the produced index. If you'd like to use Wikipedia data for GPT training you should still clean it with nltk/spacy/ftfy, but do not use the `--split-sentences` flag.
## Collecting GPT Webtext Data
-We utilize the publicly available [OpenWebText](https://github.com/eukaryote31/openwebtext) library from [jcpeterson](https://github.com/jcpeterson/openwebtext) and [eukaryote31's](https://github.com/eukaryote31/openwebtext) work to download urls. We then filtered, cleaned, and deduplicated all downloaded content according to the procedure described in our [openwebtext](./tools/openwebtext) directory. For reddit URLs corresponding to content up to October 2018 we arrived at approximately 37GB of content.
+We utilize the publicly available [OpenWebText](https://github.com/eukaryote31/openwebtext) library from [jcpeterson](https://github.com/jcpeterson/openwebtext) and [eukaryote31's](https://github.com/eukaryote31/openwebtext) work to download urls. We then filter, clean, and deduplicate all downloaded content according to the procedure described in our [openwebtext](./tools/openwebtext) directory. For reddit URLs corresponding to content up to October 2018 we arrived at approximately 37GB of content.
# Reproducibility
-Megatron training is intended to be bitwise reproducible. This means that the same training config run twice in the same HW and SW environment should produce identical model checkpoints, losses and accuracy metric values (iteration time metrics may vary).
+Megatron training can be bitwise reproducible; to enable this mode use `--deterministic-mode`. This means that the same training config run twice in the same HW and SW environment should produce identical model checkpoints, losses and accuracy metric values (iteration time metrics may vary).
+
+There are currently three known Megatron optimizations that break reproducibility whilst still producing almost identical training runs:
+1. The specific NCCL algorithm that is used during an all-reduce (as specified by the environment variable `NCCL_ALGO`) is important. We have tested the following: `^NVLS`, `Tree`, `Ring`, `CollnetDirect`, `CollnetChain`. The code admits the use of `^NVLS`, which allows NCCL the choice of non-NVLS algorithms; its choice seems to be stable.
+2. Flash attention is non-deterministic; do not use `--use-flash-attn`.
+3. If using Transformer Engine, you must also set the environment variable `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0`.
+
+In addition, determinisim has only been verified in NGC PyTorch containers up to and newer than 23.12. If you observe nondeterminism in Megatron training under other circumstances please open an issue.
+
+# Checkpoint conversion
+
+We support two forms of model conversion:
+
+1. Model class conversion (i.e., the `GPTModel` in `model.legacy` vs. `model.core`)
+2. Checkpoint format conversion (i.e., distributed vs. non-distributed checkpoint)
-There are currently two known Megatron optimizations that break reproducibility whilst still producing almost identical training runs. The following workarounds should be applied in cases where reproducibility is required:
-1. When training using `--bf16`, reproducbility is only obtained when the checkpointing and resume schedule of training is identical. If the checkpointing schedule will change, i.e. checkpointing and resume will occur at different iterations, the option `--no-bias-gelu-fusion` should be used.
-2. Flash attention is non-deterministic. If reproducibility is required do not use `--use-flash-attn`.
+## Model class conversion
-These sources of non-determinism are under active investigation. If you observe non-determinism in Megatron training under other circumstances please open an issue.
+Megatron supports converting between different model classes, including internal model classes (we currently have the older `legacy` models, and the newer `core` models) and external model classes (such as Meta, Huggingface, Mistral, and Mixtral models). Additionally, during this conversion, one can update the parallel state of the model (i.e., changing tensor and pipeline model parallelism).
+
+ We provide the tool `tools/checkpoint/convert.py` to convert between model classes. Some important arguments include:
+
+- `--model-type`: `GPT` or `BERT`
+- `--loader`: format of the existing checkpoint. Supported formats include:
+ - `legacy`: our older model classes (under `megatron.legacy.model`)
+ - `core`: our newer model classes (under `megatron.core.models`)
+ - `llama_mistral`: for loading Llama and Mistral models (supports Meta and Huggingface formats)
+ - `mixtral_hf`: for loading Mixtral models (Huggingface only)
+- `--load-dir`: directory for loading the existing checkpoint
+- `--saver`: `legacy` or `core` (see descriptions under `--loader`)
+- `--save-dir`: directory for saving the new checkpoint
+- `--target-tensor-parallel-size`: new tensor model parallel size
+- `--target-pipeline-parallel-size`: new pipeline model parallel size
+
+For more argument details, please see the main script (`convert.py`), loader scripts (`loader_core.py`, `loader_legacy.py`, `loader_llama_mistral.py`, `loader_mixtral_hf.py`), or saver scripts (`saver_core.py`, `saver_legacy.py`).
+
+An example command for converting a GPT model from the old format (`legacy`) to the new format (`core`) would look as follows:
+
+```
+python tools/checkpoint/convert.py \
+> --model-type GPT \
+> --loader legacy \
+> --load-dir ${LEGACY_FORMAT_DIR} \
+> --saver core \
+> --save-dir ${CORE_FORMAT_DIR} \
+> --target-tensor-parallel-size ${TP} \
+> --target-pipeline-parallel-size ${PP} \
+```
+
+For examples of converting Llama/Mistral models into Megatron, please see [here](docs/llama_mistral.md).
+
+## Checkpoint format conversion
+
+Megatron offers multiple checkpoint formats, including:
+
+- `torch`: Basic checkpoint format with sequential read & writes, and is tied to a specific tensor/pipeline model parallel state (TP/PP states, respectively). (While a specific checkpoint is tied to a specific TP/PP state, a checkpoint can still be manually converted via the model class converter described above).
+- `torch_dist`: Distributed checkpoint format, for fast parallel reads & writes, and also is parallel state agnostic (i.e., one can load the same checkpoint to different TP/PP setups).
+
+Generally speaking, `torch_dist` is the more modern and recommended checkpoint format due to its speed. However, depending on the use case, it may be desirable to convert between these two formats. To do so, launch your *training* script (e.g., via `pretrain_gpt.py`) as you normally would, but with two additional arguments:
+
+- `--ckpt-convert-format ${FORMAT}`: `${FORMAT}` can be one of `torch` or `torch_dist`, as described above.
+- `--ckpt-convert-save ${PATH_TO_SAVE_NEW_FORMAT}`: this path should be different than your existing `--load`/`--save` paths, to avoid overwriting the existing checkpoint. After converting, use this new path for your `--load`/`--save` paths.
+
+The general idea of this checkpoint format converter is that it launches the model just as one normally would for training, but before running any training iterations, it saves to the new checkpoint format, and then exits. It is important to note that all other launch args should remain the same, in order for the system to understand the previous checkpoint format.
+
+# Projects Using Megatron
+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)
+* [Evaluating Parameter Efficient Learning for Generation](https://aclanthology.org/2022.emnlp-main.319.pdf)
+* [Exploring the Limits of Domain-Adaptive Training for Detoxifying Large-Scale Language Models](https://arxiv.org/abs/2202.04173)
+* [Shall We Pretrain Autoregressive Language Models with Retrieval? A Comprehensive Study](https://arxiv.org/abs/2304.06762)
+* [InstructRetro: Instruction Tuning post Retrieval-Augmented Pretraining](https://arxiv.org/abs/2310.07713)
+* [An Empirical Study of Mamba-based Language Models](https://arxiv.org/abs/2406.07887)
diff --git a/compute_memory_usage.py b/compute_memory_usage.py
deleted file mode 100644
index ca6e3aacde..0000000000
--- a/compute_memory_usage.py
+++ /dev/null
@@ -1,79 +0,0 @@
-from megatron.initialize import initialize_megatron
-from megatron import get_args
-
-
-def compute_weight_and_optimizer_memory(args):
- assert args.sequence_parallel
- num_parameters_in_transformer_layers = (
- 10
- * args.num_layers
- * args.hidden_size
- * args.hidden_size
- * (
- 1
- + (args.num_query_groups / (5.0 * args.num_attention_heads))
- + (2 / (5 * args.hidden_size))
- + (1 / (5 * args.num_layers * args.hidden_size))
- )
- )
- embedding_size = args.hidden_size * args.padded_vocab_size
- if args.untie_embeddings_and_output_weights:
- num_parameters_with_embeddings = num_parameters_in_transformer_layers + (2 * embedding_size)
- else:
- num_parameters_with_embeddings = num_parameters_in_transformer_layers + embedding_size
- print(f"Number of parameters in billions: {num_parameters_with_embeddings / 10**9:.2f}")
-
- # Most loaded model shard has (1/pp_size transformer layers + 1 embedding layer) / tp_size.
- num_parameters_on_most_loaded_model_shard = (
- (num_parameters_in_transformer_layers / args.pipeline_model_parallel_size) + embedding_size
- ) / args.tensor_model_parallel_size
- # Other shards just have (1/pp_size transformer layers) / tp_size.
- num_parameters_on_other_model_shards = num_parameters_in_transformer_layers / (
- args.pipeline_model_parallel_size * args.tensor_model_parallel_size
- )
-
- print(
- f"Number of parameters in most loaded shard in billions: {num_parameters_on_most_loaded_model_shard / 10**9:.4f}"
- )
- print(
- f"Number of parameters in other shards in billions: {num_parameters_on_other_model_shards / 10**9:.4f}"
- )
-
- num_bytes_per_parameter = (
- 18 if not args.use_distributed_optimizer else 6 + (12 / args.data_parallel_size)
- )
- return num_parameters_on_most_loaded_model_shard * num_bytes_per_parameter
-
-
-def compute_activation_memory(args):
- # Using formula in Table 2 of https://arxiv.org/pdf/2205.05198.pdf.
- assert args.recompute_granularity == 'selective'
- activation_memory = (
- args.seq_length * args.micro_batch_size * args.hidden_size * args.num_layers
- ) * 34
-
- # Multiply by interleaved PP memory factor.
- activation_memory *= 1 + (
- (args.pipeline_model_parallel_size - 2)
- / (args.pipeline_model_parallel_size * args.virtual_pipeline_model_parallel_size)
- )
- return activation_memory / args.tensor_model_parallel_size
-
-
-def compute_total_memory(args):
- weight_and_optimizer_memory = compute_weight_and_optimizer_memory(args)
- activation_memory = compute_activation_memory(args)
- total_memory = weight_and_optimizer_memory + activation_memory
- print(
- f"(DP size, PP size, TP size) = {(args.data_parallel_size, args.pipeline_model_parallel_size, args.tensor_model_parallel_size)}, "
- f"Weight and optimizer memory: {weight_and_optimizer_memory / (1024 * 1024):.2f} MB, "
- f"Activation memory: {activation_memory / (1024 * 1024):.2f} MB, "
- f"Total memory: {total_memory / (1024 * 1024):.2f} MB\n"
- )
-
-
-if __name__ == "__main__":
- initialize_megatron(allow_no_cuda=True, skip_mpu_initialization=True)
- args = get_args()
-
- compute_total_memory(args)
diff --git a/docs/distrib_optimizer.md b/docs/distrib_optimizer.md
deleted file mode 100644
index def23b20eb..0000000000
--- a/docs/distrib_optimizer.md
+++ /dev/null
@@ -1,54 +0,0 @@
-# Distributed Optimizer
-
-The motivation for the distributed optimizer is to save memory by distributing the optimizer state evenly across data parallel ranks, versus the current method of replicating the optimizer state across data parallel ranks. As described in https://arxiv.org/abs/1910.02054, this branch specifically implements the following:
-
-- [yes] distribute all 'non-overlapping' optimizer state (i.e., model params already in fp32 are NOT distributed)
-- [no] distribute model gradients
-- [no] distribute model parameters
-
-Theoretical memory savings vary depending on the combination of the model's param dtype and grad dtype. In the current implementation, the theoretical number of bytes per parameter is (where 'd' is the data parallel size):
-
-| | Non-distributed optim | Distributed optim |
-| ------ | ------ | ------ |
-| float16 param, float16 grads | 20 | 4 + 16/d |
-| float16 param, fp32 grads | 18 | 6 + 12/d |
-| fp32 param, fp32 grads | 16 | 8 + 8/d |
-
-The implementation of the distributed optimizer is centered on using the contiguous grad buffer for communicating grads & params between the model state and the optimizer state. The grad buffer at any given moment either holds:
-
-1. all model grads
-2. a 1/d size _copy_ of the main grads (before copying to the optimizer state)
-3. a 1/d size _copy_ of the main params (after copying from the optimizer state)
-4. all model params
-5. zeros (or None), between iterations
-
-The grad buffer is used for performing reduce-scatter and all-gather operations, for passing grads & params between the model state and optimizer state. With this implementation, no dynamic buffers are allocated.
-
-The figures below illustrate the grad buffer's sharding scheme, and the key steps of the distributed optimizer's param update:
-
-## Data flow
-
-
-
-## Sharding scheme
-
-
-
-## Key steps
-
-_(note: using illustrations above, and assuming fp16 grads)_
-
-- Backward pass finishes (grad buffer holds 16 fp16 grad elements)
-- Call reduce-scatter on each DP rank
-- Each DP rank now has 4 elements within the grad buffer that are fully reduced (remaining 12 elements are garbage)
-- Each DP rank copies its relevant 4 fp16 grad elements from the grad buffer into 4 fp32 main grad elements (separate buffer, owned by the optimizer); i.e.
- - DP rank 0 copies elements [0:4]
- - DP rank 1 copies elements [4:8]
- - DP rank 2 copies elements [8:12]
- - DP rank 3 copies elements [12:16]
-- Optimizer.step()
-- Each DP rank copies its 4 fp32 main (/optimizer) param elements into the corresponding 4 fp16 elements in the grad buffer
-- Call all-gather on each DP rank
-- Grad buffer now contains all 16, fully updated, fp16 model param elements
-- Copy updated model params from grad buffer into their respective param tensors
-- (At this point, grad buffer is ready to be zero'd for the next iteration)
diff --git a/docs/images/distrib_optimizer/data_flow.png b/docs/images/distrib_optimizer/data_flow.png
deleted file mode 100644
index d48fc134c4..0000000000
Binary files a/docs/images/distrib_optimizer/data_flow.png and /dev/null differ
diff --git a/docs/images/distrib_optimizer/sharding_scheme.png b/docs/images/distrib_optimizer/sharding_scheme.png
deleted file mode 100644
index b07c25b05f..0000000000
Binary files a/docs/images/distrib_optimizer/sharding_scheme.png and /dev/null differ
diff --git a/docs/llama2.md b/docs/llama2.md
deleted file mode 100644
index 9043a2b95d..0000000000
--- a/docs/llama2.md
+++ /dev/null
@@ -1,171 +0,0 @@
-# Llama-2 Inference and Finetuning
-
-The Llama-2 [family of models](https://ai.meta.com/llama/) are an open-source set of pretrained & finetuned (for chat) models that have achieved strong results across a wide set of benchmarks. At the time of release, Llama-2 models achieved among the best results for open-source models, and were competitive with the closed-source GPT-3.5 model (see https://arxiv.org/pdf/2307.09288.pdf).
-
-Llama-2 checkpoints can be loaded into Megatron for inference and for finetuning. Loading these checkpoints consists of three steps:
-
-1. Get access to download the checkpoints.
-2. Convert the checkpoints from Meta/Huggingface format to Megatron format.
-3. Setup arguments for launching the model.
-
-The following sections detail these steps. The final section lists benchmark result comparisons between: 1) Llama-2 inference code running the Meta-format checkpoints, and 2) Megatron inference code running the converted checkpoints.
-
-# Contents
- * [Download Meta or Huggingface checkpoints](#download-meta-or-huggingface-checkpoints)
- * [Convert checkpoint format](#convert-checkpoint-format)
- * [Meta format](#meta-format)
- * [Huggingface format](#huggingface-format)
- * [Launch model](#launch-model)
- * [Megatron](#launch-megatron)
- * [Meta](#launch-meta)
- * [Huggingface](#launch-hf)
- * [Benchmark results](#benchmark-results)
-
-# Download Meta or Huggingface checkpoints
-
-Users must first apply for access to download the Llama-2 checkpoints either directly from [Meta](https://ai.meta.com/resources/models-and-libraries/llama-downloads/) or through [Huggingface](https://huggingface.co/docs/transformers/main/model_doc/llama2) (HF). The checkpoints are available in two formats, Meta's native format (available from both the Meta and HF links), and HF's format (available only from HF). Either format can be converted to Megatron, as detailed next.
-
-# Convert checkpoint format
-
-Depending on which checkpoint format is downloaded (Meta or HF), one or two steps must be taken to convert to Megatron format.
-
-### Meta format
-
-The Meta format checkpoints must first be converted to HF format before converting to Megatron format. The `transformers` package is required for the first step, and must have version >=4.31.0 (e.g., `pip install transformers>=4.31.0`). (**Note**: we have specifically tested with versions `4.31.0` and `4.32.0`; your experience may vary with newer versions.) Assuming the downloaded checkpoints are in `$CHECKPOINT_DIR` (with separate sub-directories for 7B, 13B, 70B, etc.), the following example command can be used to convert from Llama-2 format to HF format:
-
-```
-$>: python $LIB_DIR/transformers/models/llama/convert_llama_weights_to_hf.py \
- > --input_dir $LLAMA_FORMAT_DIR \
- > --output_dir $HF_FORMAT_DIR \
- > --model_size 7B`
-```
-
-Valid values for `--model_size` include `7B`, `13B`, and `70B` (for pretrained-only models), and `7Bf`, `13Bf`, and `70Bf` (for chat-finetuned models). Use `python convert_llama_weights_to_hf.py --help` for additional argument details. Once the checkpoints have been converted to HF format, proceed to the Huggingface format section below.
-
-### Huggingface format
-
-The HF checkpoints can be converted to Megatron format by using Megatron's own Llama-2 checkpoint converter for HF format (see script `tools/checkpoint/loader_llama2_hf.py`). One important argument that must be set correctly is the tensor parallel size (`TP`) for each model. The following table shows these values:
-
-| Model size | Tensor parallel size (`TP`) |
-| ---------- | --------------------------- |
-| 7B | 1 |
-| 13B | 2 |
-| 70B | 8 |
-
-Using these values for `TP`, along with the path to the Llama-2 tokenizer model (automatically downloaded with original checkpoint download; see `${TOKENIZER_MODEL}` below), run the following command from the root of your Megatron source code to convert from HF format to Megatron format:
-
-```
-$>: python tools/checkpoint/util.py \
- > --model-type GPT \
- > --loader llama2_hf \
- > --saver megatron \
- > --target-tensor-parallel-size ${TP} \
- > --load-dir ${HF_FORMAT_DIR} \
- > --save-dir ${MEGATRON_FORMAT_DIR} \
- > --tokenizer-model ${TOKENIZER_MODEL}
-```
-
-After this conversion, we are ready to load the checkpoints into a Megatron GPT model.
-
-# Launch model
-
-### Launch Megatron
-
-If loading for either inference or finetuning, use the following arguments:
-
-```
---tensor-model-parallel-size ${TP} \
---pipeline-model-parallel-size 1 \
---seq-length 4096 \
---max-position-embeddings 4096 \
---tokenizer-type Llama2Tokenizer \
---tokenizer-model ${TOKENIZER_MODEL} \
---load ${CHECKPOINT_DIR} \
---exit-on-missing-checkpoint \
---use-checkpoint-args \
---no-load-optim \
---no-load-rng \
---fp16 \
---untie-embeddings-and-output-weights \
---use-rotary-position-embeddings \
---normalization RMSNorm \
---no-position-embedding \
---no-masked-softmax-fusion \
---no-query-key-layer-scaling \
-```
-
-### Launch Meta
-
-Meta checkpoints can be launched with: https://github.com/facebookresearch/llama
-
-### Launch Huggingface
-
-Huggingface checkpoints can be launched with: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
-
-# Benchmark results
-
-The tables below list the benchmark comparisons between native Llama-2 (using Meta's checkpoint and Meta's inference code) and Megatron (using a converted HF checkpoint and Megatron's inference code).
-
-The values are the percent error between Megatron and Llama-2, calculated using the formula: `| - | / `, where the type of score is detailed before each table. Across all tests (80 total per model size), the mean error is 0.15%. The small difference in benchmark scores between the two models is due to minor arithmetic differences in implementation that alter the numerics slightly. Some of the factors that influence this difference include:
-
-- Megatron performs batch matrix multiplications in a couple places, such as within self attention and in SwiGLU, that Llama performs separately.
-- Megatron uses `torch.baddbmm` within self attention, versus Llama using `torch.matmul`.
-- Megatron uses a `sin`/`cos` implementation for rotary position embeddings, versus Llama using a `polar`/`complex` implementation.
-- Llama calls `torch.set_default_dtype(torch.float16)` during initialization, which Megatron does not.
-
-### Big Bench
-
-Score type: multiple choice grade.
-
-| bigbench / standard | 7b | 13b | 70b |
-| -- | -- | -- | -- |
-| date_understanding | 0.29% | 0.13% | 0.12% |
-| general_knowledge | 0.00% | 0.00% | 0.00% |
-| human_organs_senses | 0.00% | 0.00% | 0.00% |
-| intent_recognition | 0.00% | 0.11% | 0.00% |
-| riddle_sense | 0.00% | 0.00% | 0.00% |
-| similarities_abstraction | 0.00% | 0.58% | 0.00% |
-| simple_arithmetic_json_multiple_choice | 0.00% | 0.00% | 0.00% |
-| undo_permutation | 0.19% | 0.19% | 0.18% |
-
-### Multilingual
-
-Score type: multiple choice grade.
-
-| multilingual / xcopa | 7b | 13b | 70b |
-| -- | -- | -- | -- |
-| en-template-mGPT-remove-punctuation | 0.08% | 0.00% | 0.00% |
-| et-template-mGPT-remove-punctuation | 0.00% | 0.13% | 0.25% |
-| ht-template-mGPT-remove-punctuation | 0.26% | 0.13% | 0.26% |
-| id-template-mGPT-remove-punctuation | 0.11% | 0.00% | 0.19% |
-| it-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% |
-| qu-template-mGPT-remove-punctuation | 0.00% | 0.00% | 0.27% |
-| sw-template-mGPT-remove-punctuation | 0.14% | 0.13% | 0.13% |
-| th-template-mGPT-remove-punctuation | 0.25% | 0.13% | 0.13% |
-| tr-template-mGPT-remove-punctuation | 0.26% | 0.00% | 0.34% |
-| vi-template-mGPT-remove-punctuation | 0.00% | 0.11% | 0.00% |
-| zh-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% |
-
-### LM Evaluation Harness
-
-Score type: multiple choice grade.
-
-| lm-eval | 7b | 13b | 70b |
-| -- | -- | -- | -- |
-| boolq | 0.04% | 0.04% | 0.07% |
-| hellaswag | 0.02% | 0.03% | 0.03% |
-| piqa | 0.00% | 0.00% | 0.07% |
-| winogrande | 0.00% | 0.11% | 0.20% |
-
-### MMLU
-
-Score type: multiple choice grade.
-
-Note: the number in brackets is the number of sub-tasks for each supercategory.
-
-| mmlu | 7b | 13b | 70b |
-| -- | -- | -- | -- |
-| stem [18] | 0.79% | 0.05% | 0.01% |
-| humanities [13] | 0.19% | 0.01% | 0.02% |
-| other (business, health, misc.) [14] | 0.08% | 0.06% | 0.12% |
-| social sciences [12] | 0.37% | 0.21% | 0.01% |
diff --git a/docs/llama_mistral.md b/docs/llama_mistral.md
new file mode 100644
index 0000000000..5dd61866e8
--- /dev/null
+++ b/docs/llama_mistral.md
@@ -0,0 +1,444 @@
+# Llama, Mistral and other Llama-like model support in Megatron-LM
+
+NOTE: In order to simplify code we now only support converting llama-3.x and mistral checkpoints downloaded from Huggingface.
+
+The [Llama-2](https://ai.meta.com/llama/) and [Llama-3.x](https://llama.meta.com/) family of models are an open-source set of pretrained & finetuned (for chat) models that have achieved strong results across a wide set of benchmarks. At their times of release, both Llama-2 and Llama-3 models achieved among the best results for open-source models, and were competitive with leading closed-source models (see https://arxiv.org/pdf/2307.09288.pdf and https://ai.meta.com/blog/meta-llama-3/).
+
+Similarly, [Mistral-7b](https://mistral.ai/news/announcing-mistral-7b/) is an open-source model with pretrained and finetuned (for chat) variants that achieve strong benchmark results.
+
+Architecturally Llama-2, Llama-3 and Mistral-7b are very similar. As such Megatron can support loading checkpoints from all three for inference and finetuning. Converting the checkpoints and loading them is slightly different for each model and is detailed for each below.
+
+# Contents
+
+- [Llama, Mistral and other Llama-like model support in Megatron-LM](#llama-mistral-and-other-llama-like-model-support-in-megatron-lm)
+- [Contents](#contents)
+- [Llama-2](#llama-2)
+ - [Download Meta or Huggingface checkpoints](#download-meta-or-huggingface-checkpoints)
+ - [Convert checkpoint format](#convert-checkpoint-format)
+ - [Meta format](#meta-format)
+ - [Huggingface format](#huggingface-format)
+ - [Launch model](#launch-model)
+ - [Launch Megatron](#launch-megatron)
+ - [Launch Meta](#launch-meta)
+ - [Launch Huggingface](#launch-huggingface)
+ - [Benchmark results](#benchmark-results)
+ - [Big Bench](#big-bench)
+ - [Multilingual](#multilingual)
+ - [LM Evaluation Harness](#lm-evaluation-harness)
+ - [MMLU](#mmlu)
+- [Llama-3.x](#llama-3x)
+ - [Download Huggingface checkpoints](#download-huggingface-checkpoints)
+ - [Convert checkpoint format](#convert-checkpoint-format-1)
+ - [Huggingface format](#huggingface-format-1)
+ - [(Optional) Validate checkpoints](#optional-validate-checkpoints)
+ - [Launch model](#launch-model-1)
+- [Mistral-7b](#mistral-7b)
+ - [Download Huggingface checkpoints](#download-huggingface-checkpoints-2)
+ - [Convert checkpoint format](#convert-checkpoint-format-3)
+ - [(Optional) Validate checkpoints](#optional-validate-checkpoints-2)
+ - [Launch model](#launch-model-3)
+- [Other Llama-like model support](#other-llama-like-model-support)
+- [Known numerical differences](#known-numerical-differences)
+- [Using legacy model format](#using-legacy-model-format)
+
+
+# Llama-2
+
+Llama-2 checkpoints can be loaded into Megatron for inference and for finetuning. Loading these checkpoints consists of three steps:
+
+1. Get access to download the checkpoints.
+2. Convert the checkpoints from Meta/Huggingface format to Megatron format.
+3. Setup arguments for launching the model.
+
+The following sections detail these steps. The final section lists benchmark result comparisons between: 1) Llama-2 inference code running the Meta-format checkpoints, and 2) Megatron inference code running the converted checkpoints.
+
+## Download Meta or Huggingface checkpoints
+
+Users must first apply for access to download the Llama-2 checkpoints either directly from [Meta](https://ai.meta.com/resources/models-and-libraries/llama-downloads/) or through [Huggingface](https://huggingface.co/docs/transformers/main/model_doc/llama2) (HF). The checkpoints are available in two formats, Meta's native format (available from both the Meta and HF links), and HF's format (available only from HF). Either format can be converted to Megatron, as detailed next.
+
+## Convert checkpoint format
+
+We recommend passing `--dtype bf16` for training or finetuning. Inference can be done in bfloat16 or float16.
+
+### Meta format
+
+The Meta format checkpoints are converted to HF format as an intermediate step before converting to Megatron format. The `transformers` package is required, and must have version >=4.31.0 (e.g., `pip install transformers>=4.31.0`). (**Note**: we have specifically tested with versions `4.31.0` and `4.32.0`; your experience may vary with newer versions.) Assuming the downloaded checkpoints are in `$CHECKPOINT_DIR` (with separate sub-directories for 7B, 13B, 70B, etc.), the following example command can be used to convert from Llama-2 format to HF format in bfloat16:
+
+```
+python tools/checkpoint/convert.py \
+> --model-type GPT \
+> --loader llama_mistral \
+> --load-dir ${META_FORMAT_DIR} \
+> --model-size ${MODEL_SIZE} \
+> --checkpoint-type meta \
+> --tokenizer-model ${TOKENIZER_MODEL} \
+> --saver core \
+> --save-dir ${MEGATRON_FORMAT_DIR} \
+> --target-tensor-parallel-size ${TP} \
+> --target-pipeline-parallel-size ${PP} \
+> --bf16
+```
+
+Valid values for `--model-size` are `llama2-7B`, `llama2-13B`, and `llama2-70B` (for pretrained-only models), and `llama2-7Bf`, `llama2-13Bf`, and `llama2-70Bf` (for chat-finetuned models).
+
+### Huggingface format
+
+The HF checkpoints can be converted to Megatron format by using Megatron's own Llama-2 checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`). One important argument that must be set correctly is the tensor parallel size (`TP`) for each model. The following table shows these values:
+
+| Model size | Tensor parallel size (`TP`) |
+| ---------- | --------------------------- |
+| 7B | 1 |
+| 13B | 2 |
+| 70B | 8 |
+
+Using these values for `TP`, along with the path to the Llama-2 tokenizer model (automatically downloaded with original checkpoint download; see `${TOKENIZER_MODEL}` below), run the following command from the root of your Megatron source code to convert from HF format to Megatron format:
+
+```
+python tools/checkpoint/convert.py \
+> --model-type GPT \
+> --loader llama_mistral \
+> --load-dir ${HF_FORMAT_DIR} \
+> --model-size ${MODEL_SIZE} \
+> --checkpoint-type hf \
+> --tokenizer-model ${TOKENIZER_MODEL} \
+> --saver core \
+> --save-dir ${MEGATRON_FORMAT_DIR} \
+> --target-tensor-parallel-size ${TP} \
+> --target-pipeline-parallel-size ${PP} \
+> --bf16
+```
+
+After this conversion, we are ready to load the checkpoints into a Megatron GPT model.
+
+## Launch model
+
+### Launch Megatron
+
+If loading for either inference or finetuning, use the following arguments:
+
+```
+--tensor-model-parallel-size ${TP} \
+--pipeline-model-parallel-size 1 \
+--seq-length 4096 \
+--max-position-embeddings 4096 \
+--tokenizer-type Llama2Tokenizer \
+--tokenizer-model ${TOKENIZER_MODEL} \
+--load ${CHECKPOINT_DIR} \
+--exit-on-missing-checkpoint \
+--use-checkpoint-args \
+--no-load-optim \
+--no-load-rng \
+--untie-embeddings-and-output-weights \
+--use-rotary-position-embeddings \
+--normalization RMSNorm \
+--no-position-embedding \
+--no-masked-softmax-fusion \
+--attention-softmax-in-fp32
+```
+
+**Note:** If you converted to the legacy model format (i.e., `--saver legacy`), please see [here](#using-legacy-model-format).
+
+### Launch Meta
+
+Meta checkpoints can be launched with: https://github.com/facebookresearch/llama
+
+### Launch Huggingface
+
+Huggingface checkpoints can be launched with: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+
+## Benchmark results
+
+The tables below list the benchmark comparisons between native Llama-2 (using Meta's checkpoint and Meta's inference code) and Megatron (using a converted HF checkpoint and Megatron's inference code).
+
+The values are the percent error between Megatron and Llama-2, calculated using the formula: `| - | / `, where the type of score is detailed before each table. Across all tests (80 total per model size), the mean error is 0.15%. The small difference in benchmark scores between the two models is due to minor arithmetic differences in implementation that alter the numerics slightly. Some of the factors that influence this difference include:
+
+- Megatron performs batch matrix multiplications in a couple places, such as within self attention and in SwiGLU, that Llama performs separately.
+- Megatron uses `torch.baddbmm` within self attention, versus Llama using `torch.matmul`.
+- Megatron uses a `sin`/`cos` implementation for rotary position embeddings, versus Llama using a `polar`/`complex` implementation.
+- Llama calls `torch.set_default_dtype(torch.float16)` during initialization, which Megatron does not.
+
+### Big Bench
+
+Score type: multiple choice grade.
+
+| bigbench / standard | 7b | 13b | 70b |
+| -- | -- | -- | -- |
+| date_understanding | 0.29% | 0.13% | 0.12% |
+| general_knowledge | 0.00% | 0.00% | 0.00% |
+| human_organs_senses | 0.00% | 0.00% | 0.00% |
+| intent_recognition | 0.00% | 0.11% | 0.00% |
+| riddle_sense | 0.00% | 0.00% | 0.00% |
+| similarities_abstraction | 0.00% | 0.58% | 0.00% |
+| simple_arithmetic_json_multiple_choice | 0.00% | 0.00% | 0.00% |
+| undo_permutation | 0.19% | 0.19% | 0.18% |
+
+### Multilingual
+
+Score type: multiple choice grade.
+
+| multilingual / xcopa | 7b | 13b | 70b |
+| -- | -- | -- | -- |
+| en-template-mGPT-remove-punctuation | 0.08% | 0.00% | 0.00% |
+| et-template-mGPT-remove-punctuation | 0.00% | 0.13% | 0.25% |
+| ht-template-mGPT-remove-punctuation | 0.26% | 0.13% | 0.26% |
+| id-template-mGPT-remove-punctuation | 0.11% | 0.00% | 0.19% |
+| it-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% |
+| qu-template-mGPT-remove-punctuation | 0.00% | 0.00% | 0.27% |
+| sw-template-mGPT-remove-punctuation | 0.14% | 0.13% | 0.13% |
+| th-template-mGPT-remove-punctuation | 0.25% | 0.13% | 0.13% |
+| tr-template-mGPT-remove-punctuation | 0.26% | 0.00% | 0.34% |
+| vi-template-mGPT-remove-punctuation | 0.00% | 0.11% | 0.00% |
+| zh-template-mGPT-remove-punctuation | 0.00% | 0.10% | 0.09% |
+
+### LM Evaluation Harness
+
+Score type: multiple choice grade.
+
+| lm-eval | 7b | 13b | 70b |
+| -- | -- | -- | -- |
+| boolq | 0.04% | 0.04% | 0.07% |
+| hellaswag | 0.02% | 0.03% | 0.03% |
+| piqa | 0.00% | 0.00% | 0.07% |
+| winogrande | 0.00% | 0.11% | 0.20% |
+
+### MMLU
+
+Score type: multiple choice grade.
+
+Note: the number in brackets is the number of sub-tasks for each supercategory.
+
+| mmlu | 7b | 13b | 70b |
+| -- | -- | -- | -- |
+| stem [18] | 0.79% | 0.05% | 0.01% |
+| humanities [13] | 0.19% | 0.01% | 0.02% |
+| other (business, health, misc.) [14] | 0.08% | 0.06% | 0.12% |
+| social sciences [12] | 0.37% | 0.21% | 0.01% |
+
+# Llama-3.x
+
+Llama-3.x checkpoints can be loaded into Megatron for inference and for finetuning. Loading these checkpoints consists of several steps:
+
+1. Get access to download the checkpoints (weights and tokenizer).
+2. Convert the checkpoints from Huggingface format to Megatron format.
+3. (Optional) Validate converted checkpoints
+4. Setup arguments for launching the model.
+
+The following sections detail these steps.
+
+## Download Huggingface checkpoints
+
+Users must first apply for access to download the Llama-3.x checkpoints from [Huggingface](https://huggingface.co/meta-llama).
+
+## Convert checkpoint format
+
+We recommend passing `--dtype bf16` for training or finetuning. Inference can be done in bfloat16 or float16.
+
+### Huggingface format
+
+The HF checkpoints can be converted to Megatron format by using Megatron's own Llama-3.x checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`). One important argument that must be set correctly is the tensor parallel size (`TP`) for each model. The following table shows these values:
+
+| Model size | Tensor parallel size (`TP`) |
+| ---------- | --------------------------- |
+| 1B | 1 |
+| 3B | 1 |
+| 8B | 1 |
+| 70B | 8 |
+
+Using these values for `TP`, along with the path to the Llama-3.x tokenizer model (automatically downloaded with original checkpoint download; see `${TOKENIZER_MODEL}` below), run the following command from the root of your Megatron source code to convert from HF format to Megatron format:
+
+```
+$>: python tools/checkpoint/convert.py \
+ > --bf16 \
+ > --model-type GPT \
+ > --loader llama_mistral \
+ > --saver core \
+ > --target-tensor-parallel-size ${TP} \
+ > --checkpoint-type hf \
+ > --load-dir ${HF_FORMAT_DIR} \
+ > --save-dir ${MEGATRON_FORMAT_DIR} \
+ > --tokenizer-model ${TOKENIZER_MODEL} \
+ > --model-size llama3 \
+```
+
+After this conversion, we are ready to load the checkpoints into a Megatron GPT model.
+
+## (Optional) Validate checkpoints
+
+A Megatron-LM text generation server for Llama3 can be launched using the script `examples/inference/llama_mistral/run_text_generation_llama3.sh `. For Llama3.1, please use `examples/inference/llama_mistral/run_text_generation_llama3.1.sh`.
+
+Once running, query the server with `curl 'http://:5000/api' -X 'PUT' -H 'Content-Type: application/json; charset=UTF-8' -d '{"prompts":[""], "tokens_to_generate":100, "top_k":1}'`.
+
+A reference generation for comparison can be obtained from the Huggingface transformers library by running `python examples/llama_mistral/huggingface_reference.py --model_path --prompt `.
+
+## Launch model
+
+If loading for either inference or finetuning, use the following arguments for Llama 3.0:
+
+```
+--tensor-model-parallel-size ${TP} \
+--pipeline-model-parallel-size 1 \
+--seq-length 8192 \
+--max-position-embeddings 8192 \
+--tokenizer-type HuggingFaceTokenizer \
+--tokenizer-model ${TOKENIZER_MODEL} \
+--load ${CHECKPOINT_DIR} \
+--exit-on-missing-checkpoint \
+--use-checkpoint-args \
+--no-load-optim \
+--no-load-rng \
+--untie-embeddings-and-output-weights \
+--normalization RMSNorm \
+--position-embedding-type rope \
+--no-masked-softmax-fusion \
+--attention-softmax-in-fp32 \
+--disable-bias-linear \
+--transformer-impl transformer_engine \
+--group-query-attention 8 \
+--attention-dropout 0.0 \
+--hidden-dropout 0.0 \
+--rotary-base 500000 \
+--rotary-percent 1.0 \
+--ffn-hidden-size 14336 \
+--num-attention-heads 32 \
+--swiglu \
+--bf16 \
+```
+
+For Llama3.1 please use the following arguments:
+
+```
+--tensor-model-parallel-size ${TP} \
+--pipeline-model-parallel-size 1 \
+--seq-length 8192 \
+--max-position-embeddings 131072 \
+--tokenizer-type HuggingFaceTokenizer \
+--tokenizer-model ${TOKENIZER_MODEL} \
+--load ${CHECKPOINT_DIR} \
+--exit-on-missing-checkpoint \
+--use-checkpoint-args \
+--no-load-optim \
+--no-load-rng \
+--untie-embeddings-and-output-weights \
+--normalization RMSNorm \
+--position-embedding-type rope \
+--no-masked-softmax-fusion \
+--attention-softmax-in-fp32 \
+--disable-bias-linear \
+--transformer-impl transformer_engine \
+--group-query-attention 8 \
+--attention-dropout 0.0 \
+--hidden-dropout 0.0 \
+--rotary-base 500000 \
+--rotary-percent 1.0 \
+--use-rope-scaling \
+--ffn-hidden-size 14336 \
+--num-attention-heads 32 \
+--swiglu \
+--bf16 \
+```
+
+**Note:** If you converted to the legacy model format (i.e., `--saver legacy`), please see [here](#using-legacy-model-format).
+
+# Mistral-7b
+
+Megatron currently supports loading the v0.3 release of Mistral-7b (which does not use sliding window attention and offers a larger 32768 vocabulary) for inference and finetuning. Loading these checkpoints consists of several steps:
+
+1. Get access to download the checkpoints (weights and tokenizer).
+2. Convert the checkpoints from HuggingFace format to Megatron format.
+3. (Optional) Validate converted checkpoints
+4. Setup arguments for launching the model.
+
+The following sections detail these steps.
+
+## Download Huggingface checkpoints
+
+Users must first apply for access to download the Mistral-7b checkpoints through [Huggingface](https://huggingface.co/mistralai/Mistral-7B-v0.3) (HF).
+
+## Convert checkpoint format
+
+The HF checkpoints can be converted to Megatron format by using Megatron's own Mistral checkpoint converter for HF format (see script `tools/checkpoint/loader_llama_mistral.py`).
+
+Using the path to the Mistral tokenizer model (downloaded alongside the HF checkpoint), run the following command from the root of your Megatron source code to convert from HF format to the Megatron core format:
+
+```
+$>: python tools/checkpoint/convert.py \
+ > --bf16 \
+ > --model-type GPT \
+ > --loader llama_mistral \
+ > --saver core \
+ > --target-tensor-parallel-size ${TP} \
+ > --checkpoint-type hf \
+ > --load-dir ${HF_FORMAT_DIR} \
+ > --save-dir ${MEGATRON_FORMAT_DIR} \
+ > --tokenizer-model ${TOKENIZER_MODEL} \
+ > --model-size mistral \
+```
+
+After this conversion, we are ready to load the checkpoints into a Megatron core GPT model.
+
+## (Optional) Validate checkpoints
+
+A Megatron-LM text generation server for Mistral-7B can be launched using the script `examples/inference/llama_mistral/run_text_generation_mistral.sh `.
+
+Once running, query the server with `curl 'http://:5000/api' -X 'PUT' -H 'Content-Type: application/json; charset=UTF-8' -d '{"prompts":[""], "tokens_to_generate":100, "top_k":1}'`.
+
+A reference generation for comparison can be obtained from the Huggingface transformers library by running `python examples/inference/llama_mistral/huggingface_reference.py --model_path --prompt `.
+
+## Launch model
+
+If loading for either inference or finetuning, use the following arguments:
+
+```
+--tensor-model-parallel-size ${TP} \
+--pipeline-model-parallel-size 1 \
+--seq-length 4096 \
+--max-position-embeddings 4096 \
+--tokenizer-type HuggingFaceTokenizer \
+--tokenizer-model ${TOKENIZER_MODEL} \
+--load ${CHECKPOINT_DIR} \
+--exit-on-missing-checkpoint \
+--use-checkpoint-args \
+--no-load-optim \
+--no-load-rng \
+--untie-embeddings-and-output-weights \
+--normalization RMSNorm \
+--position-embedding-type rope \
+--no-masked-softmax-fusion \
+--attention-softmax-in-fp32
+--apply-layernorm-1p \
+--transformer-impl transformer_engine \
+--group-query-attention 8 \
+--disable-bia-linear \
+--rotary-base 1000000 \
+--rotary-percent 1.0 \
+--swiglu \
+--ffn-hidden-size 14336 \
+--num-attention-heads 32
+```
+
+**Note:** If you converted to the legacy model format (i.e., `--saver legacy`), please see [here](#using-legacy-model-format).
+
+# Other Llama-like model support
+
+*Note: Experimental*
+
+Many models such as Yi-34B and Qwen2.x use the Llama architecture and may be converted from HuggingFace to Megatron using the commands in [Llama-3.x](#llama-3x).
+
+# Known numerical differences
+
+It is not expected that the megatron and Huggingface implementations of llama3.x and mistral models will produce numerically identical results. There are multiple points where small numerical differences are expected. This is a non-exhaustive list:
+
+1. TransformerEngine (TE) uses the model params_dtype inside RMSNorm whereas the Huggingface implementation uses fp32. See for details: https://github.com/NVIDIA/TransformerEngine/issues/1132
+2. Huggingface `transformers` implements the q, k and v projections in self-attention as separate GEMMs whereas Megatron core combines them into a single GEMM for efficiency. This leads to small numerical differences.
+
+# Using legacy model format
+
+In all the checkpoint conversion examples used in this document, the saver format `--saver core` is used, signifying that the newer (and recommended) Megatron GPT model class will be used. I.e.:
+
+- old class: `megatron.legacy.model.gpt_model.GPTModel`
+- new class: `megatron.core.models.gpt.gpt_model.GPTModel`
+
+Using this new format is the recommended approach. However, if your use case requires using the older class (i.e., convert using `--saver legacy`), then when launching training or finetuning, the following args must be added:
+
+- `--use-legacy-models`: use the older model class
+- `--ckpt-format torch`: use the `torch` checkpoint format, which is the only checkpoint format that is compatible with the legacy model format
diff --git a/docs/source/api-guide/context_parallel.rst b/docs/source/api-guide/context_parallel.rst
new file mode 100644
index 0000000000..c08defd210
--- /dev/null
+++ b/docs/source/api-guide/context_parallel.rst
@@ -0,0 +1,35 @@
+context\_parallel package
+=========================
+
+Context parallelism overview
+----------------------------
+
+.. figure:: ../images/context_parallel/CP_overview.png
+ :alt: cp_overview
+ :align: center
+
+ Figure 1: A transformer layer running with TP2CP2. Communications next to Attention are for CP, others are for TP. (AG/RS: all-gather in forward and reduce-scatter in backward, RS/AG: reduce-scatter in forward and all-gather in backward, /AG: no-op in forward and all-gather in backward).
+
+Context Parallelism ("CP") is a parallelization scheme on the dimension of sequence length. Unlike prior SP (sequence parallelism) which only splits the sequence of Dropout and LayerNorm activations, CP partitions the network inputs and all activations along sequence dimension. With CP, all modules except attention (e.g., Linear, LayerNorm, etc.) can work as usual without any changes, because they do not have inter-token operations. As for attention, the Q (query) of each token needs to compute with the KV (key and value) of all tokens in the same sequence. Hence, CP requires additional all-gather across GPUs to collect the full sequence of KV. Correspondingly, reduce-scatter should be applied to the activation gradients of KV in backward propagation. To reduce activation memory footprint, each GPU only stores the KV of a sequence chunk in forward and gathers KV again in backward. KV communication happens between a GPU and its counterparts in other TP groups. The all-gather and reduce-scatter are transformed to point-to-point communications in ring topology under the hood. Exchanging KV also can leverage MQA/GQA to reduce communication volumes, as they only have one or few attention heads for KV.
+
+For example, in Figure 1, assuming sequence length is 8K, each GPU processes 4K tokens. GPU0 and GPU2 compose a CP group, they exchange KV with each other. Same thing also happens between GPU1 and GPU3. CP is similar to `Ring Attention `_ but provides better performance by (1) leveraging the latest OSS and cuDNN flash attention kernels; (2) removing unnecessary computation resulted from low-triangle causal masking and achieving optimal load balance among GPUs.
+
+Context parallelism benefits
+----------------------------
+
+.. figure:: ../images/context_parallel/CP_results.png
+ :alt: cp_results
+ :align: center
+
+ Figure 2: Speedup of 175B GPT with various TP+CP combinations vs. full recompute (i.e., TP8CP1).
+
+LLM encounters OOM (out of memory) issue with long context (i.e., long sequence length) because of linearly increasing memory footprint of activations. Recomputing activations in backward can avoid OOM but also introduce significant overheads (~30% with full recompute). Enlarging TP (tensor model parallelism) can fix the OOM issue as well, but it potentially makes compute (e.g., Linear) too short to overlap communication latencies. To be clear, scaling out to more GPUs with bigger TP can hit the overlapping problem no matter if OOM happens.
+
+CP can better address the issues. With CP, each GPU only computes on a part of the sequence, which reduces both computation and communication by CP times. Therefore, there are no concerns about the overlapping between them. The activation memory footprint per GPU is also CP times smaller, hence no OOM issue anymore. As Figure 2 shows, the combinations of TP and CP can achieve optimal performance by eliminating recompute overheads and making the best tradeoff between computation and communications.
+
+Enabling context parallelism
+----------------------------
+
+CP support has been added to GPT. All models that share GPT code path also should be able to benefit from CP, such as Llama. CP can work with TP (tensor model parallelism), PP (pipeline model parallelism), and DP (data parallelism), where the total number of GPUs equals TPxCPxPPxDP. CP also can work with different attention variants, including MHA/MQA/GQA, uni-directional and bi-directional masking.
+
+CP is enabled by simply setting context_parallel_size= in command line. Default context_parallel_size is 1, which means CP is disabled. Running with CP requires Megatron-Core (>=0.5.0) and Transformer Engine (>=1.1).
diff --git a/docs/source/api-guide/datasets.rst b/docs/source/api-guide/datasets.rst
new file mode 100644
index 0000000000..247a3f07d3
--- /dev/null
+++ b/docs/source/api-guide/datasets.rst
@@ -0,0 +1,104 @@
+datasets package
+================
+
+.. mdinclude :: ../../../megatron/core/datasets/readme.md
+
+Submodules
+----------
+
+datasets.blended\_megatron\_dataset\_config module
+---------------------------------------------------
+
+.. automodule:: core.datasets.blended_megatron_dataset_config
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.blended\_megatron\_dataset\_builder module
+---------------------------------------------------
+
+.. automodule:: core.datasets.blended_megatron_dataset_builder
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.megatron\_tokenizer module
+-----------------------------------
+
+.. automodule:: core.datasets.megatron_tokenizer
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.indexed\_dataset module
+--------------------------------
+
+.. automodule:: core.datasets.indexed_dataset
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.megatron\_dataset module
+---------------------------------
+
+.. automodule:: core.datasets.megatron_dataset
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.gpt\_dataset module
+----------------------------
+
+.. automodule:: core.datasets.gpt_dataset
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.masked\_dataset module
+-------------------------------
+
+.. automodule:: core.datasets.masked_dataset
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.bert\_dataset module
+-----------------------------
+
+.. automodule:: core.datasets.bert_dataset
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.t5\_dataset module
+---------------------------
+
+.. automodule:: core.datasets.t5_dataset
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.blended\_dataset module
+----------------------------------
+
+.. automodule:: core.datasets.blended_dataset
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+datasets.utils module
+---------------------
+
+.. automodule:: core.datasets.utils
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.datasets
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
diff --git a/docs/source/api-guide/dist_checkpointing.rst b/docs/source/api-guide/dist_checkpointing.rst
new file mode 100644
index 0000000000..7e384a08a3
--- /dev/null
+++ b/docs/source/api-guide/dist_checkpointing.rst
@@ -0,0 +1,79 @@
+dist\_checkpointing package
+===========================
+
+A library for saving and loading the distributed checkpoints.
+A "distributed checkpoint" can have various underlying formats (current default format is based on Zarr)
+but has a distinctive property - the checkpoint saved in one parallel configuration (tensor/pipeline/data parallelism)
+can be loaded in a different parallel configuration.
+
+Using the library requires defining sharded state_dict dictionaries with functions from *mapping* and *optimizer* modules.
+Those state dicts can be saved or loaded with a *serialization* module using strategies from *strategies* module.
+
+
+Subpackages
+-----------
+
+.. toctree::
+ :maxdepth: 4
+
+ dist_checkpointing.strategies
+
+Submodules
+----------
+
+dist\_checkpointing.serialization module
+----------------------------------------
+
+.. automodule:: core.dist_checkpointing.serialization
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+dist\_checkpointing.mapping module
+----------------------------------
+
+.. automodule:: core.dist_checkpointing.mapping
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+dist\_checkpointing.optimizer module
+------------------------------------
+
+.. automodule:: core.dist_checkpointing.optimizer
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+dist\_checkpointing.core module
+-------------------------------
+
+.. automodule:: core.dist_checkpointing.core
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+dist\_checkpointing.dict\_utils module
+--------------------------------------
+
+.. automodule:: core.dist_checkpointing.dict_utils
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+
+dist\_checkpointing.utils module
+--------------------------------
+
+.. automodule:: core.dist_checkpointing.utils
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.dist_checkpointing
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/dist_checkpointing.strategies.rst b/docs/source/api-guide/dist_checkpointing.strategies.rst
new file mode 100644
index 0000000000..41e674c761
--- /dev/null
+++ b/docs/source/api-guide/dist_checkpointing.strategies.rst
@@ -0,0 +1,50 @@
+dist\_checkpointing.strategies package
+======================================
+
+Package defining different checkpoint formats (backends) and saving/loading algorithms (strategies).
+
+Strategies can be used for implementing new checkpoint formats or implementing new (more optimal for a given use case) ways of saving/loading of existing formats.
+Strategies are passed to `dist_checkpointing.load` and `dist_checkpointing.save` functions and control the actual saving/loading procedure.
+
+Submodules
+----------
+
+dist\_checkpointing.strategies.base module
+------------------------------------------
+
+.. automodule:: core.dist_checkpointing.strategies.base
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+dist\_checkpointing.strategies.tensorstore module
+-------------------------------------------------
+
+.. automodule:: core.dist_checkpointing.strategies.tensorstore
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+dist\_checkpointing.strategies.two\_stage module
+------------------------------------------------
+
+.. automodule:: core.dist_checkpointing.strategies.two_stage
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+dist\_checkpointing.strategies.zarr module
+------------------------------------------
+
+.. automodule:: core.dist_checkpointing.strategies.zarr
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.dist_checkpointing.strategies
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/dist_optimizer.md b/docs/source/api-guide/dist_optimizer.md
new file mode 100644
index 0000000000..34f42d5343
--- /dev/null
+++ b/docs/source/api-guide/dist_optimizer.md
@@ -0,0 +1,40 @@
+# Distributed Optimizer
+
+The motivation for the distributed optimizer is to save memory by distributing the optimizer state evenly across data parallel ranks (https://arxiv.org/abs/1910.02054), versus the naive method of replicating the optimizer state across data parallel ranks.
+
+Theoretical memory savings vary depending on the combination of the datatype of the model's parameters (`param_dtype`) and main gradients accumulated across data-parallel replicas (`grad_dtype`). We always use `fp32` main parameters for optimizer steps. In the current implementation, the theoretical number of bytes per parameter is (where d is the data parallel size):
+
+| | Non-distributed optim | Distributed optim |
+| ------ | ------ | ------ |
+| `fp16` parameters, `fp16` gradients | 20 | 4 + 16/d |
+| `bf16` parameters, `fp32` gradients | 18 | 6 + 12/d |
+| `fp32` parameters, `fp32` gradients | 16 | 8 + 8/d |
+
+Our implementation of the distributed optimizer uses contiguous buffers for parameters and main gradients; model gradients are copied over to the main gradients as soon as they are fully computed.
+
+The figures below illustrate the distributed optimizer's sharding scheme, and the key steps of the distributed optimizer's parameter update:
+
+## Data flow
+
+
+
+## Sharding scheme
+
+
+
+## Key steps
+
+_(note: using illustrations above, assuming `bf16` model weights, `bf16` model gradients that are computed by the backward pass and `fp32` main gradients that are also used for optimizer steps; we always use `fp32` main weights for optimizer steps)_
+
+- Backward pass finishes (gradient buffer holds 16 `fp32` gradient elements).
+- Call reduce-scatter on each DP rank.
+- Each DP rank now has 4 elements within the gradient buffer that are fully reduced (remaining 12 elements are garbage).
+ - DP rank 0 has gradient values for elements [0:4].
+ - DP rank 1 has gradient values for elements [4:8].
+ - DP rank 2 has gradient values for elements [8:12].
+ - DP rank 3 has gradient values for elements [12:16].
+- Optimizer.step().
+- Each DP rank copies its 4 `fp32` main parameter elements into the corresponding `bf16` parameter buffer (each element is cast from fp32 to fp16).
+- Call all-gather on each DP rank.
+- The parameter buffer now contains all 16, fully updated, `bf16` model parameter elements. Parameters in PyTorch modules already point to the appropriate locations in this parameter buffer, and thus forward passes are ready to run after the all-gather completes.
+- At this point, the gradient buffer is also ready to be zero'd for the next iteration.
diff --git a/docs/source/api-guide/distributed.rst b/docs/source/api-guide/distributed.rst
new file mode 100644
index 0000000000..737820331c
--- /dev/null
+++ b/docs/source/api-guide/distributed.rst
@@ -0,0 +1,53 @@
+distributed package
+===================
+
+This package contains various utilities to finalize model weight gradients
+on each rank before the optimizer step. This includes a distributed data
+parallelism wrapper to all-reduce or reduce-scatter the gradients across
+data-parallel replicas, and a `finalize\_model\_grads` method to
+synchronize gradients across different parallelism modes (e.g., 'tied'
+layers on different pipeline stages, or gradients for experts in a MoE on
+different ranks due to expert parallelism).
+
+Submodules
+----------
+
+distributed.distributed\_data\_parallel
+---------------------------------------
+
+Model wrapper for distributed data parallelism. Stores gradients in a
+contiguous buffer, and supports the option of overlapping communication
+(all-reduce or reduce-scatter) with backprop computation by breaking up
+full model's gradients into smaller buckets and running all-reduce /
+reduce-scatter on each bucket asynchronously.
+
+.. automodule:: core.distributed.distributed_data_parallel
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+distributed.finalize\_model\_grads
+----------------------------------
+
+Finalize model gradients for optimizer step across all used parallelism modes.
+Synchronizes the all-reduce / reduce-scatter of model gradients across DP replicas,
+all-reduces the layernorm gradients for sequence parallelism, embedding gradients
+across first and last pipeline stages (if not tied), and expert gradients for expert
+parallelism.
+
+.. automodule:: core.distributed.finalize_model_grads
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+
+Module contents
+---------------
+
+Contains functionality to synchronize gradients across different ranks before
+optimizer step.
+
+.. automodule:: core.distributed
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/encoder_decoder_parallelism.rst b/docs/source/api-guide/encoder_decoder_parallelism.rst
new file mode 100644
index 0000000000..7cdff941de
--- /dev/null
+++ b/docs/source/api-guide/encoder_decoder_parallelism.rst
@@ -0,0 +1,54 @@
+encoder-decoder-parallelism package
+===================================
+
+Mcore (as of 0.9) supports heterogeneous parallelism for encoder-decoder models.
+In particular, the user is now able to specify the amount of tensor and pipeline parallelism and have it be
+distinct from that in the decoder.
+
+Submodules
+----------
+
+Encoder Pipeline Parallelism
+----------------------------
+
+Supported in: T5, LLaVa.
+
+The new argument for encoder parallelism is `--encoder-pipeline-model-parallel-size`. This argument is completely distinct
+from the usual argument that controls pipelining: `--pipeline-model-parallel-size`, which controls the amount of pipelining in the decoder
+in the context of encoder-decoder models.
+
+The total amount of pipelining in an encoder-decoder model is the sum of these two arguments. By default, the amount of
+encoder pipelining is 0, and the amount of decoder pipelining is 1, meaning that the encoder & decoder share the single pipeline rank.
+If `--pipeline-model-parallel-size` > 1,then the amount of encoder parallelism has to be specified and has to be greater than 0.
+This is because we are not able to share pipeline ranks between the encoder and decoder anymore.
+
+Encoder Tensor Parallelism
+--------------------------
+
+Supported in: LLaVa.
+
+Since we expect encoders to be much smaller than decoders, we also give users the ability to set a different amount of tensor
+parallelism than the decoder. This is achieved with the argument `--encoder-tensor-model-parallel-size`. To use this option, you must
+be using encoder pipeline parallelism (ie, `--encoder-pipeline-model-parallel-size` > 0).
+
+Unlike with encoder pipeline parallelism, which was unrestricted by the amount of decoder pipeline parallelism, we only allow encoders to have
+less than or the same amount of tensor parallelism as the decoder. The summary of how we do this is that within p2p_communication.py, we have
+to send the activations of one encoder rank to several decoder ranks; correspondingly, we have to add support for summing gradients from several
+(downstream) decoder ranks for the encoder rank. We have not seen a quantization-related degradation from summing these gradient tensors
+together yet; it could happen in very large models.
+
+
+Number of GPUs Required
+-----------------------
+
+The total amount of GPUs required to train a model when these options enabled is:
+
+dp * etp * epp * cp + dp * tp * pp * cp
+
+where:
+dp: amount of data parallelism (this is the same for the encoder & decoder)
+[e]tp: amount of tensor parallelism
+[e]pp: amount of pipeline parallelism
+cp: amount of context parallelism (as with dp, this is the same for the encoder & decoder)
+
+The default value of this argument is 0; in practice, we will use the amount of tensor parallelism in the decoder to construct the encoder.
diff --git a/docs/source/api-guide/fusions.rst b/docs/source/api-guide/fusions.rst
new file mode 100644
index 0000000000..22782ca84e
--- /dev/null
+++ b/docs/source/api-guide/fusions.rst
@@ -0,0 +1,65 @@
+fusions package
+===============
+
+This package provides modules that provide commonly fused
+operations. Fusing operations improves compute efficiency by
+increasing the amount of work done each time a tensor is read from
+memory. To perform the fusion, modules in this either rely on PyTorch
+functionality for doing just-in-time compilation
+(i.e. `torch.jit.script` in older PyTorch versions of `torch.compile`
+in recent versions), or call into custom kernels in external libraries
+such as Apex or TransformerEngine.
+
+Submodules
+----------
+
+fusions.fused\_bias\_dropout module
+-----------------------------------
+
+This module uses PyTorch JIT to fuse the bias add and dropout operations. Since dropout is not used during inference, different functions are used when in train mode and when in inference mode.
+
+.. automodule:: core.fusions.fused_bias_dropout
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+fusions.fused\_bias\_gelu module
+--------------------------------
+
+This module uses PyTorch JIT to fuse the bias add and GeLU nonlinearity operations.
+
+.. automodule:: core.fusions.fused_bias_gelu
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+fusions.fused\_layer\_norm module
+---------------------------------
+
+This module provides a wrapper around various fused LayerNorm implementation in Apex.
+
+.. automodule:: core.fusions.fused_layer_norm
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+fusions.fused\_softmax module
+-----------------------------
+
+This module provides wrappers around variations of Softmax in Apex.
+
+.. automodule:: core.fusions.fused_softmax
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+fusions.fused\_cross\_entropy\_loss module
+------------------------------------------
+
+This module uses PyTorch JIT to fuse the cross entropy loss calculation and batches communication calls.
+
+.. automodule:: core.fusions.fused_cross_entropy
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
diff --git a/docs/source/api-guide/index.rst b/docs/source/api-guide/index.rst
new file mode 100644
index 0000000000..dac785af04
--- /dev/null
+++ b/docs/source/api-guide/index.rst
@@ -0,0 +1,20 @@
+API Guide
+=========
+
+.. toctree::
+ :maxdepth: 4
+
+ models
+ tensor_parallel
+ context_parallel
+ pipeline_parallel
+ fusions
+ transformer
+ moe
+ dist_checkpointing
+ dist_optimizer
+ distributed
+ datasets
+ num_microbatches_calculator
+ optimizer_param_scheduler
+ encoder_decoder_parallelism
\ No newline at end of file
diff --git a/docs/source/api-guide/models.bert.rst b/docs/source/api-guide/models.bert.rst
new file mode 100644
index 0000000000..1b562ce72c
--- /dev/null
+++ b/docs/source/api-guide/models.bert.rst
@@ -0,0 +1,22 @@
+models.bert package
+===================
+Useful package for training bert and bert like encoder only models. It optionally comes with a binary head that can be used for classification tasks .
+
+Submodules
+----------
+
+models.bert.bert\_model module
+------------------------------
+
+.. automodule:: core.models.bert.bert_model
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.models.bert
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/models.gpt.rst b/docs/source/api-guide/models.gpt.rst
new file mode 100644
index 0000000000..31c4da6a9c
--- /dev/null
+++ b/docs/source/api-guide/models.gpt.rst
@@ -0,0 +1,22 @@
+models.gpt package
+==================
+This is the implementation of the popular GPT model. It supports several features like model parallelization (Tensor Parallel, Pipeline Parallel, Data Parallel) , mixture of experts, FP8 , Distributed optimizer etc. We are constantly adding new features. So be on the lookout or raise an issue if you want to have something added.
+
+Submodules
+----------
+
+models.gpt.gpt\_model module
+----------------------------
+
+.. automodule:: core.models.gpt.gpt_model
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.models.gpt
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/models.rst b/docs/source/api-guide/models.rst
new file mode 100644
index 0000000000..12c40e4f35
--- /dev/null
+++ b/docs/source/api-guide/models.rst
@@ -0,0 +1,21 @@
+models package
+==============
+This package contains most of the popular LLMs . Currently we have support for GPT, Bert, T5 and Retro . This is an ever growing list so keep an eye out.
+
+Subpackages
+-----------
+
+.. toctree::
+ :maxdepth: 4
+
+ models.gpt
+ models.t5
+ models.bert
+
+Module contents
+---------------
+
+.. automodule:: core.models
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/models.t5.rst b/docs/source/api-guide/models.t5.rst
new file mode 100644
index 0000000000..1cc3315682
--- /dev/null
+++ b/docs/source/api-guide/models.t5.rst
@@ -0,0 +1,21 @@
+models.t5 package
+=================
+
+Submodules
+----------
+
+models.t5.t5\_model module
+--------------------------
+
+.. automodule:: core.models.T5.t5_model
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.models.T5
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/moe.rst b/docs/source/api-guide/moe.rst
new file mode 100644
index 0000000000..9afc01e080
--- /dev/null
+++ b/docs/source/api-guide/moe.rst
@@ -0,0 +1,4 @@
+Mixture of Experts package
+==========================
+
+.. mdinclude :: ../../../megatron/core/transformer/moe/README.md
diff --git a/docs/source/api-guide/num_microbatches_calculator.rst b/docs/source/api-guide/num_microbatches_calculator.rst
new file mode 100644
index 0000000000..4790b31749
--- /dev/null
+++ b/docs/source/api-guide/num_microbatches_calculator.rst
@@ -0,0 +1,12 @@
+Microbatches Calculator
+=======================
+This api is used to calculate the number of microbatches required to fit a given model on a given batch size.
+
+
+Module contents
+---------------
+
+.. automodule:: core.num_microbatches_calculator
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/optimizer_param_scheduler.rst b/docs/source/api-guide/optimizer_param_scheduler.rst
new file mode 100644
index 0000000000..caf5d8abfb
--- /dev/null
+++ b/docs/source/api-guide/optimizer_param_scheduler.rst
@@ -0,0 +1,12 @@
+Optimizer Parameters Scheduler
+==============================
+This api is used to calculate the learning rate and weight decay for the optimizer.
+
+
+Module contents
+---------------
+
+.. automodule:: core.optimizer_param_scheduler
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/pipeline_parallel.rst b/docs/source/api-guide/pipeline_parallel.rst
new file mode 100644
index 0000000000..5c67079a70
--- /dev/null
+++ b/docs/source/api-guide/pipeline_parallel.rst
@@ -0,0 +1,47 @@
+pipeline\_parallel package
+==========================
+
+This package contains implementations for two different pipeline parallelism
+schedules (one without interleaving and one with interleaving, see `Efficient
+Large-Scale Language Model Training on GPU Clusters Using Megatron-LM `_
+for details), and a default no-pipelining schedule. It also contains methods
+for the point-to-point communication that is needed between pipeline stages.
+
+Submodules
+----------
+
+pipeline\_parallel.p2p\_communication module
+--------------------------------------------
+
+Contains implementations for the various point-to-point communication needed
+(e.g., `recv_forward` and `recv_backward`) in the different pipeline parallelism
+schedules.
+
+.. automodule:: core.pipeline_parallel.p2p_communication
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+pipeline\_parallel.schedules module
+-----------------------------------
+
+Contains implementations for two pipeline parallelism schedules
+(`forward_backward_pipelining_with_interleaving`for pipeline parallelism with
+interleaving, `forward_backward_pipelining_without_interleaving` for pipeline
+parallelism without interleaving) and a default no-pipelining schedule
+(`forward_backward_no_pipelining`). `get_forward_backward_func` returns the right
+scheduling function to use based on the configuration being trained
+(e.g., if pipeline-parallel size is 1, use `forward_backward_no_pipelining`).
+
+.. automodule:: core.pipeline_parallel.schedules
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.pipeline_parallel
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/tensor_parallel.rst b/docs/source/api-guide/tensor_parallel.rst
new file mode 100644
index 0000000000..d8ae9dea22
--- /dev/null
+++ b/docs/source/api-guide/tensor_parallel.rst
@@ -0,0 +1,67 @@
+tensor\_parallel package
+========================
+
+This package contains an implementation for tensor parallelism in transformer
+models (see `Megatron-LM: Training Multi-Billion Parameter Language Models
+Using Model Parallelism `_ and `Reducing
+Activation Recomputation in Large Transformer Models `_
+for details).
+
+Submodules
+----------
+
+tensor\_parallel.cross\_entropy module
+--------------------------------------
+
+.. automodule:: core.tensor_parallel.cross_entropy
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+tensor\_parallel.data module
+----------------------------
+
+.. automodule:: core.tensor_parallel.data
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+tensor\_parallel.layers module
+------------------------------
+
+.. automodule:: core.tensor_parallel.layers
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+tensor\_parallel.mappings module
+--------------------------------
+
+.. automodule:: core.tensor_parallel.mappings
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+tensor\_parallel.random module
+------------------------------
+
+.. automodule:: core.tensor_parallel.random
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+tensor\_parallel.utils module
+-----------------------------
+
+.. automodule:: core.tensor_parallel.utils
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.tensor_parallel
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/api-guide/transformer.rst b/docs/source/api-guide/transformer.rst
new file mode 100644
index 0000000000..6e2e894d54
--- /dev/null
+++ b/docs/source/api-guide/transformer.rst
@@ -0,0 +1,136 @@
+transformer package
+===================
+
+The `transformer` package provides a customizable and configurable
+implementation of the transformer model architecture. Each component
+of a transformer stack, from entire layers down to individual linear
+layers, can be customized by swapping in different PyTorch modules
+using the "spec" parameters (see `here
+`_). The
+configuration of the transformer (hidden size, number of layers,
+number of attention heads, etc.) is provided via a `TransformerConfig`
+object.
+
+Submodules
+----------
+
+transformer.attention module
+----------------------------
+
+This is the entire attention portion, either self or cross attention,
+of a transformer layer including the query, key, and value
+projections, a "core" attention calculation (e.g. dot product
+attention), and final output linear projection.
+
+.. automodule:: core.transformer.attention
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+transformer.dot\_product\_attention module
+------------------------------------------
+
+This is a PyTorch-only implementation of dot product attention. A more
+efficient implementation, like those provided by FlashAttention or
+CUDNN's FusedAttention, are typically used when training speed is
+important.
+
+.. automodule:: core.transformer.dot_product_attention
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+transformer.enums module
+------------------------
+
+.. automodule:: core.transformer.enums
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+transformer.identity\_op module
+-------------------------------
+
+This provides a pass-through module that can be used in specs to
+indicate that the operation should not be performed. For example, when
+using LayerNorm with the subsequent linear layer, an IdentityOp can be
+passed in as the LayerNorm module to use.
+
+.. automodule:: core.transformer.identity_op
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+transformer.mlp module
+----------------------
+
+This is the entire MLP portion of the transformer layer with an input
+projection, non-linearity, and output projection.
+
+.. automodule:: core.transformer.mlp
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+transformer.module module
+-------------------------
+
+This provides a common base class for all modules used in the
+transformer that contains some common functionality.
+
+.. automodule:: core.transformer.module
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+transformer.transformer\_block module
+-------------------------------------
+
+A block, or stack, of several transformer layers. The layers can all
+be the same or each can be unique.
+
+.. automodule:: core.transformer.transformer_block
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+transformer.transformer\_config module
+--------------------------------------
+
+This contains all of the configuration options for the
+transformer. Using a dataclass reduces code bloat by keeping all
+arguments together in a dataclass instead of passing several arguments
+through multiple layers of function calls.
+
+.. automodule:: core.transformer.transformer_config
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+transformer.transformer\_layer module
+-------------------------------------
+
+A single standard transformer layer including attention and MLP blocks.
+
+.. automodule:: core.transformer.transformer_layer
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+transformer.utils module
+------------------------
+
+Various utilities used in the transformer implementation.
+
+.. automodule:: core.transformer.utils
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: core.transformer
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/docs/source/images/context_parallel/CP_overview.png b/docs/source/images/context_parallel/CP_overview.png
new file mode 100644
index 0000000000..38c55b371a
Binary files /dev/null and b/docs/source/images/context_parallel/CP_overview.png differ
diff --git a/docs/source/images/context_parallel/CP_results.png b/docs/source/images/context_parallel/CP_results.png
new file mode 100644
index 0000000000..e0415ce86e
Binary files /dev/null and b/docs/source/images/context_parallel/CP_results.png differ
diff --git a/docs/source/images/distrib_optimizer/data_flow.png b/docs/source/images/distrib_optimizer/data_flow.png
new file mode 100644
index 0000000000..01f5cfb2e7
Binary files /dev/null and b/docs/source/images/distrib_optimizer/data_flow.png differ
diff --git a/docs/source/images/distrib_optimizer/sharding_scheme.png b/docs/source/images/distrib_optimizer/sharding_scheme.png
new file mode 100644
index 0000000000..e48dd95024
Binary files /dev/null and b/docs/source/images/distrib_optimizer/sharding_scheme.png differ
diff --git a/docs/source/images/moe/token_drop.png b/docs/source/images/moe/token_drop.png
new file mode 100644
index 0000000000..1c335ee7aa
Binary files /dev/null and b/docs/source/images/moe/token_drop.png differ
diff --git a/docs/source/index.rst b/docs/source/index.rst
new file mode 100644
index 0000000000..f2a89b8ac7
--- /dev/null
+++ b/docs/source/index.rst
@@ -0,0 +1,23 @@
+.. Lumache documentation master file, created by
+ sphinx-quickstart on Tue Aug 15 13:44:10 2023.
+ You can adapt this file completely to your liking, but it should at least
+ contain the root `toctree` directive.
+
+Megatron Core User Guide
+===================================
+
+**Megatron Core** is a Python library that has the core components required to build your language models.
+A reference implementation of Megatron Core can be found in `NeMo `_ It offers a *simple* and
+*intuitive* API.
+
+.. toctree::
+ :maxdepth: 2
+ :caption: User Guide
+
+ user-guide/index
+
+.. toctree::
+ :maxdepth: 3
+ :caption: API Guide
+
+ api-guide/index
diff --git a/docs/source/user-guide/index.rst b/docs/source/user-guide/index.rst
new file mode 100644
index 0000000000..0fb996a4f0
--- /dev/null
+++ b/docs/source/user-guide/index.rst
@@ -0,0 +1,4 @@
+User Guide
+============
+
+.. mdinclude:: ../../../megatron/core/QuickStart.md
\ No newline at end of file
diff --git a/examples/detxoify_lm/README.md b/examples/academic_paper_scripts/detxoify_lm/README.md
similarity index 100%
rename from examples/detxoify_lm/README.md
rename to examples/academic_paper_scripts/detxoify_lm/README.md
diff --git a/examples/detxoify_lm/annotations/filter-selfgeneration.py b/examples/academic_paper_scripts/detxoify_lm/annotations/filter-selfgeneration.py
similarity index 100%
rename from examples/detxoify_lm/annotations/filter-selfgeneration.py
rename to examples/academic_paper_scripts/detxoify_lm/annotations/filter-selfgeneration.py
diff --git a/examples/detxoify_lm/annotations/perspective_api_annotate.py b/examples/academic_paper_scripts/detxoify_lm/annotations/perspective_api_annotate.py
similarity index 98%
rename from examples/detxoify_lm/annotations/perspective_api_annotate.py
rename to examples/academic_paper_scripts/detxoify_lm/annotations/perspective_api_annotate.py
index fd82c2a2ae..9736db099a 100644
--- a/examples/detxoify_lm/annotations/perspective_api_annotate.py
+++ b/examples/academic_paper_scripts/detxoify_lm/annotations/perspective_api_annotate.py
@@ -107,7 +107,7 @@ def get_score(line):
except UnicodeDecodeError:
try:
decoded_text = encoded_text[:20476].decode('utf8')
- except:
+ except Exception:
print("Error occurred")
data['score'] = None
return json.dumps(data)
@@ -138,7 +138,7 @@ def get_scores(lines):
except UnicodeDecodeError:
try:
decoded_text = encoded_text[:20476].decode('utf8')
- except:
+ except Exception:
print("Error occurred")
data['score'] = None
all_data.append(json.dumps(data))
diff --git a/examples/detxoify_lm/annotations/preprocess.sh b/examples/academic_paper_scripts/detxoify_lm/annotations/preprocess.sh
similarity index 100%
rename from examples/detxoify_lm/annotations/preprocess.sh
rename to examples/academic_paper_scripts/detxoify_lm/annotations/preprocess.sh
diff --git a/examples/detxoify_lm/finetune_gpt.py b/examples/academic_paper_scripts/detxoify_lm/finetune_gpt.py
similarity index 88%
rename from examples/detxoify_lm/finetune_gpt.py
rename to examples/academic_paper_scripts/detxoify_lm/finetune_gpt.py
index f1bbba5bda..6a3696d388 100644
--- a/examples/detxoify_lm/finetune_gpt.py
+++ b/examples/academic_paper_scripts/detxoify_lm/finetune_gpt.py
@@ -10,19 +10,20 @@
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
os.path.pardir, os.path.pardir)))
-from megatron import get_args
-from megatron import get_timers
-from megatron import get_tokenizer
-from megatron import print_rank_0
+from megatron.training import get_args
+from megatron.training import get_timers
+from megatron.training import get_tokenizer
+from megatron.training import print_rank_0
from megatron.core import mpu
from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
from megatron.core.datasets.blended_megatron_dataset_config import GPTDatasetConfig
from megatron.core.datasets.gpt_dataset import GPTDataset
-from megatron.model import GPTModel
+from megatron.core.datasets.utils import get_blend_from_list
+from megatron.legacy.model import GPTModel
from megatron.core.enums import ModelType
from megatron.training import pretrain
-from megatron.utils import get_ltor_masks_and_position_ids
-from megatron.utils import average_losses_across_data_parallel_group
+from megatron.training.utils import get_ltor_masks_and_position_ids
+from megatron.training.utils import average_losses_across_data_parallel_group
def model_provider(pre_process=True, post_process=True):
"""Build the model."""
@@ -105,8 +106,9 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
train_ds, _, test_ds = BlendedMegatronDatasetBuilder(
GPTDataset,
train_val_test_num_samples,
+ lambda: True,
GPTDatasetConfig(
- blend=args.data_path,
+ blend=get_blend_from_list(args.data_path),
split=args.split,
random_seed=args.seed,
sequence_length=args.seq_length,
@@ -119,8 +121,9 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
_, valid_ds, _ = BlendedMegatronDatasetBuilder(
GPTDataset,
train_val_test_num_samples,
+ lambda: True,
GPTDatasetConfig(
- blend=args.data_path2,
+ blend=get_blend_from_list(args.data_path2),
split="98,2,0",
random_seed=1234,
sequence_length=2048,
diff --git a/examples/detxoify_lm/finetune_gpt_distributed-1.3b.sh b/examples/academic_paper_scripts/detxoify_lm/finetune_gpt_distributed-1.3b.sh
similarity index 100%
rename from examples/detxoify_lm/finetune_gpt_distributed-1.3b.sh
rename to examples/academic_paper_scripts/detxoify_lm/finetune_gpt_distributed-1.3b.sh
diff --git a/examples/detxoify_lm/generate-1.3b.sh b/examples/academic_paper_scripts/detxoify_lm/generate-1.3b.sh
similarity index 100%
rename from examples/detxoify_lm/generate-1.3b.sh
rename to examples/academic_paper_scripts/detxoify_lm/generate-1.3b.sh
diff --git a/examples/detxoify_lm/generate_samples_gpt.py b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py
similarity index 68%
rename from examples/detxoify_lm/generate_samples_gpt.py
rename to examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py
index 47e1590ea5..895a45d024 100644
--- a/examples/detxoify_lm/generate_samples_gpt.py
+++ b/examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py
@@ -9,23 +9,84 @@
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
os.path.pardir, os.path.pardir)))
import torch
-from megatron import get_args
-from megatron import get_tokenizer
-from megatron import print_rank_0
-from megatron.checkpointing import load_checkpoint
+from megatron.training import get_args
+from megatron.training import get_tokenizer
+from megatron.training import print_rank_0
+from megatron.training.checkpointing import load_checkpoint
from megatron.core import mpu
-from megatron.initialize import initialize_megatron
-from megatron.model import GPTModel
+from megatron.training.initialize import initialize_megatron
+from megatron.legacy.model import GPTModel
from megatron.training import get_model
-from megatron.text_generation import generate_and_post_process
+from megatron.inference.text_generation import generate_and_post_process
+from megatron.training.arguments import core_transformer_config_from_args
+from megatron.core.models.gpt import GPTModel
+from typing import Union
+import megatron.legacy.model
+from megatron.core.transformer.spec_utils import import_module
+from megatron.training.arguments import core_transformer_config_from_args
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_local_spec
+def model_provider(pre_process=True, post_process=True) -> Union[GPTModel, megatron.legacy.model.GPTModel]:
+ """Builds the model.
-def model_provider(pre_process=True, post_process=True):
- """Build the model."""
+ If you set the use_legacy_models to True, it will return the legacy GPT model and if not the core GPT model.
+
+ Args:
+ pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True.
+ post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True.
+
+
+ Returns:
+ Union[GPTModel, megatron.legacy.model.GPTModel]: The returned model
+ """
+ args = get_args()
print_rank_0('building GPT model ...')
- model = GPTModel(num_tokentypes=0, parallel_output=False,
- pre_process=pre_process, post_process=post_process)
+ config = core_transformer_config_from_args(args)
+
+ if args.use_legacy_models:
+ model = megatron.legacy.model.GPTModel(
+ config,
+ num_tokentypes=0,
+ parallel_output=False,
+ pre_process=pre_process,
+ post_process=post_process
+ )
+ else:
+ if args.spec is None:
+ if args.transformer_impl == 'local':
+ transformer_layer_spec = get_gpt_layer_local_spec(
+ num_experts=args.num_experts,
+ moe_grouped_gemm=args.moe_grouped_gemm
+ )
+ elif args.transformer_impl == 'transformer_engine':
+ transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+ num_experts=args.num_experts,
+ moe_grouped_gemm=args.moe_grouped_gemm
+ )
+ else:
+ raise ValueError(f"Invalid transformer_impl {args.transformer_impl}")
+ elif args.spec[0] == 'local':
+ transformer_layer_spec = get_gpt_layer_local_spec(
+ num_experts=args.num_experts,
+ moe_grouped_gemm=args.moe_grouped_gemm
+ )
+ else:
+ transformer_layer_spec = import_module(args.spec)
+
+ model = GPTModel(
+ config=config,
+ transformer_layer_spec=transformer_layer_spec,
+ vocab_size=args.padded_vocab_size,
+ max_sequence_length=args.max_position_embeddings,
+ pre_process=pre_process,
+ post_process=post_process,
+ fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+ parallel_output=False,
+ share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+ position_embedding_type=args.position_embedding_type,
+ rotary_percent=args.rotary_percent
+ )
return model
diff --git a/examples/detxoify_lm/perspective_api.py b/examples/academic_paper_scripts/detxoify_lm/perspective_api.py
similarity index 100%
rename from examples/detxoify_lm/perspective_api.py
rename to examples/academic_paper_scripts/detxoify_lm/perspective_api.py
diff --git a/examples/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh b/examples/academic_paper_scripts/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh
similarity index 100%
rename from examples/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh
rename to examples/academic_paper_scripts/detxoify_lm/self_generation/selfgenerate-1.3b-unconditional.sh
diff --git a/examples/msdp/README.md b/examples/academic_paper_scripts/msdp/README.md
similarity index 100%
rename from examples/msdp/README.md
rename to examples/academic_paper_scripts/msdp/README.md
diff --git a/examples/msdp/data_processing.sh b/examples/academic_paper_scripts/msdp/data_processing.sh
similarity index 100%
rename from examples/msdp/data_processing.sh
rename to examples/academic_paper_scripts/msdp/data_processing.sh
diff --git a/examples/msdp/eval_knwl_generation.sh b/examples/academic_paper_scripts/msdp/eval_knwl_generation.sh
similarity index 100%
rename from examples/msdp/eval_knwl_generation.sh
rename to examples/academic_paper_scripts/msdp/eval_knwl_generation.sh
diff --git a/examples/msdp/eval_resp_generation.sh b/examples/academic_paper_scripts/msdp/eval_resp_generation.sh
similarity index 100%
rename from examples/msdp/eval_resp_generation.sh
rename to examples/academic_paper_scripts/msdp/eval_resp_generation.sh
diff --git a/examples/msdp/prep_resp_gen.sh b/examples/academic_paper_scripts/msdp/prep_resp_gen.sh
similarity index 100%
rename from examples/msdp/prep_resp_gen.sh
rename to examples/academic_paper_scripts/msdp/prep_resp_gen.sh
diff --git a/examples/msdp/prompt_knwl_gen.sh b/examples/academic_paper_scripts/msdp/prompt_knwl_gen.sh
similarity index 100%
rename from examples/msdp/prompt_knwl_gen.sh
rename to examples/academic_paper_scripts/msdp/prompt_knwl_gen.sh
diff --git a/examples/msdp/prompt_resp_gen.sh b/examples/academic_paper_scripts/msdp/prompt_resp_gen.sh
similarity index 100%
rename from examples/msdp/prompt_resp_gen.sh
rename to examples/academic_paper_scripts/msdp/prompt_resp_gen.sh
diff --git a/examples/sc21/CONFIG.sh b/examples/academic_paper_scripts/sc21/CONFIG.sh
similarity index 100%
rename from examples/sc21/CONFIG.sh
rename to examples/academic_paper_scripts/sc21/CONFIG.sh
diff --git a/examples/sc21/README.md b/examples/academic_paper_scripts/sc21/README.md
similarity index 100%
rename from examples/sc21/README.md
rename to examples/academic_paper_scripts/sc21/README.md
diff --git a/examples/sc21/SBATCH.sh b/examples/academic_paper_scripts/sc21/SBATCH.sh
similarity index 100%
rename from examples/sc21/SBATCH.sh
rename to examples/academic_paper_scripts/sc21/SBATCH.sh
diff --git a/examples/sc21/SRUN.sh b/examples/academic_paper_scripts/sc21/SRUN.sh
similarity index 100%
rename from examples/sc21/SRUN.sh
rename to examples/academic_paper_scripts/sc21/SRUN.sh
diff --git a/examples/sc21/run_figure_11.sh b/examples/academic_paper_scripts/sc21/run_figure_11.sh
similarity index 100%
rename from examples/sc21/run_figure_11.sh
rename to examples/academic_paper_scripts/sc21/run_figure_11.sh
diff --git a/examples/sc21/run_figure_12.sh b/examples/academic_paper_scripts/sc21/run_figure_12.sh
similarity index 100%
rename from examples/sc21/run_figure_12.sh
rename to examples/academic_paper_scripts/sc21/run_figure_12.sh
diff --git a/examples/sc21/run_figure_13.sh b/examples/academic_paper_scripts/sc21/run_figure_13.sh
similarity index 100%
rename from examples/sc21/run_figure_13.sh
rename to examples/academic_paper_scripts/sc21/run_figure_13.sh
diff --git a/examples/sc21/run_figure_14.sh b/examples/academic_paper_scripts/sc21/run_figure_14.sh
similarity index 100%
rename from examples/sc21/run_figure_14.sh
rename to examples/academic_paper_scripts/sc21/run_figure_14.sh
diff --git a/examples/sc21/run_figure_15.sh b/examples/academic_paper_scripts/sc21/run_figure_15.sh
similarity index 100%
rename from examples/sc21/run_figure_15.sh
rename to examples/academic_paper_scripts/sc21/run_figure_15.sh
diff --git a/examples/sc21/run_figure_16.sh b/examples/academic_paper_scripts/sc21/run_figure_16.sh
similarity index 100%
rename from examples/sc21/run_figure_16.sh
rename to examples/academic_paper_scripts/sc21/run_figure_16.sh
diff --git a/examples/sc21/run_figure_17.sh b/examples/academic_paper_scripts/sc21/run_figure_17.sh
similarity index 100%
rename from examples/sc21/run_figure_17.sh
rename to examples/academic_paper_scripts/sc21/run_figure_17.sh
diff --git a/examples/sc21/run_figure_18.sh b/examples/academic_paper_scripts/sc21/run_figure_18.sh
similarity index 100%
rename from examples/sc21/run_figure_18.sh
rename to examples/academic_paper_scripts/sc21/run_figure_18.sh
diff --git a/examples/sc21/run_table_1.sh b/examples/academic_paper_scripts/sc21/run_table_1.sh
similarity index 100%
rename from examples/sc21/run_table_1.sh
rename to examples/academic_paper_scripts/sc21/run_table_1.sh
diff --git a/examples/bert/README.md b/examples/bert/README.md
new file mode 100644
index 0000000000..6c1fe95bf0
--- /dev/null
+++ b/examples/bert/README.md
@@ -0,0 +1,53 @@
+# BERT MODEL
+
+## Table of contents
+- [1. Training Setup](#1-training-setup)
+- [2. Configurations](#2-configurations)
+
+## 1. Training setup
+
+
+To run the model using a docker container run it as follows
+```
+PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:24.01-py3
+CHECKPOINT_PATH="" #
+TENSORBOARD_LOGS_PATH=""#
+VOCAB_FILE="" #//bert-vocab.txt
+DATA_PATH="" #_text_document
+
+docker run \
+ --gpus=all \
+ --ipc=host \
+ --workdir /workspace/megatron-lm \
+ -v /path/to/data:/path/to/data \
+ -v /path/to/megatron-lm:/workspace/megatron-lm \
+ megatron-lm nvcr.io/nvidia/pytorch:24.01-py3 \
+ bash examples/bert/train_bert_340m_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH $VOCAB_FILE $DATA_PATH "
+
+```
+NOTE: Depending on the environment you are running it the above command might like slightly different.
+
+
+## 2. Configurations
+
+The example in this folder shows you how to run 340m large model. There are other configs you could run as well
+
+### 4B
+```
+ --num-layers 48 \
+ --hidden-size 2560 \
+ --num-attention-heads 32 \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+
+```
+
+### 20B
+```
+ --num-layers 48 \
+ --hidden-size 6144 \
+ --num-attention-heads 96 \
+ --tensor-model-parallel-size 4 \
+ --pipeline-model-parallel-size 4 \
+
+```
\ No newline at end of file
diff --git a/examples/bert/train_bert_340m_distributed.sh b/examples/bert/train_bert_340m_distributed.sh
new file mode 100644
index 0000000000..f0d9c87c8b
--- /dev/null
+++ b/examples/bert/train_bert_340m_distributed.sh
@@ -0,0 +1,79 @@
+#!/bin/bash
+
+# Runs the "340M" parameter model (Bert - Large)
+
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+GPUS_PER_NODE=8
+# Change for multinode config
+MASTER_ADDR=localhost
+MASTER_PORT=6000
+NUM_NODES=1
+NODE_RANK=0
+WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
+
+CHECKPOINT_PATH=$1 #
+TENSORBOARD_LOGS_PATH=$2 #
+VOCAB_FILE=$3 #/bert-vocab.json
+DATA_PATH=$4 #_text_document
+
+DISTRIBUTED_ARGS=(
+ --nproc_per_node $GPUS_PER_NODE
+ --nnodes $NUM_NODES
+ --master_addr $MASTER_ADDR
+ --master_port $MASTER_PORT
+)
+
+BERT_MODEL_ARGS=(
+ --num-layers 24
+ --hidden-size 1024
+ --num-attention-heads 16
+ --seq-length 512
+ --max-position-embeddings 512
+ --attention-backend auto # Can use (flash/fused/unfused/local)
+)
+
+TRAINING_ARGS=(
+ --micro-batch-size 4
+ --global-batch-size 32
+ --train-iters 1000000
+ --weight-decay 1e-2
+ --clip-grad 1.0
+ --fp16
+ --lr 0.0001
+ --lr-decay-iters 990000
+ --lr-decay-style linear
+ --min-lr 1.0e-5
+ --weight-decay 1e-2
+ --lr-warmup-fraction .01
+ --clip-grad 1.0
+)
+
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size 8
+ --pipeline-model-parallel-size 16
+)
+
+DATA_ARGS=(
+ --data-path $DATA_PATH
+ --vocab-file $VOCAB_FILE
+ --split 949,50,1
+)
+
+EVAL_AND_LOGGING_ARGS=(
+ --log-interval 100
+ --save-interval 10000
+ --eval-interval 1000
+ --save $CHECKPOINT_PATH
+ --load $CHECKPOINT_PATH
+ --eval-iters 10
+ --tensorboard-dir $TENSORBOARD_LOGS_PATH
+)
+
+torchrun ${DISTRIBUTED_ARGS[@]} pretrain_bert.py \
+ ${BERT_MODEL_ARGS[@]} \
+ ${TRAINING_ARGS[@]} \
+ ${MODEL_PARALLEL_ARGS[@]} \
+ ${DATA_ARGS[@]} \
+ ${EVAL_AND_LOGGING_ARGS[@]}
+
\ No newline at end of file
diff --git a/examples/evaluate_retriever_nq.sh b/examples/evaluate_retriever_nq.sh
deleted file mode 100644
index a579b5fd94..0000000000
--- a/examples/evaluate_retriever_nq.sh
+++ /dev/null
@@ -1,37 +0,0 @@
-#!/bin/bash
-
-# Evaluate natural question test data given Wikipedia embeddings and pretrained
-# ICT model or a finetuned model for Natural Question task
-
-# Datasets can be downloaded from the following link:
-# https://github.com/facebookresearch/DPR/blob/master/data/download_data.py
-
-EVIDENCE_DATA_DIR=
-EMBEDDING_PATH=
-CHECKPOINT_PATH=
-
-QA_FILE=
-
-python tasks/main.py \
- --task RETRIEVER-EVAL \
- --tokenizer-type BertWordPieceLowerCase \
- --num-layers 12 \
- --hidden-size 768 \
- --num-attention-heads 12 \
- --tensor-model-parallel-size 1 \
- --micro-batch-size 128 \
- --seq-length 512 \
- --max-position-embeddings 512 \
- --load ${CHECKPOINT_PATH} \
- --evidence-data-path ${EVIDENCE_DATA_DIR} \
- --embedding-path ${EMBEDDING_PATH} \
- --retriever-seq-length 256 \
- --vocab-file bert-vocab.txt\
- --qa-data-test ${QA_FILE} \
- --faiss-use-gpu \
- --retriever-report-topk-accuracies 1 5 20 100 \
- --fp16 \
- --indexer-log-interval 1000 \
- --indexer-batch-size 128
-
-
diff --git a/examples/evaluate_zeroshot_gpt.sh b/examples/evaluate_zeroshot_gpt.sh
deleted file mode 100755
index 2cc1c5a760..0000000000
--- a/examples/evaluate_zeroshot_gpt.sh
+++ /dev/null
@@ -1,37 +0,0 @@
-#!/bin/bash
-
-WORLD_SIZE=8
-
-DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
- --nnodes 1 \
- --node_rank 0 \
- --master_addr localhost \
- --master_port 6000"
-
-TASK="LAMBADA"
-
-VALID_DATA=
-VOCAB_FILE=gpt2-vocab.json
-MERGE_FILE=gpt2-merges.txt
-CHECKPOINT=checkpoints/gpt2_345m
-
-
-python -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/main.py \
- --task $TASK \
- --valid-data $VALID_DATA \
- --tokenizer-type GPT2BPETokenizer \
- --strict-lambada \
- --vocab-file $VOCAB_FILE \
- --merge-file $MERGE_FILE \
- --load $CHECKPOINT \
- --tensor-model-parallel-size 1 \
- --num-layers 24 \
- --hidden-size 1024 \
- --num-attention-heads 16 \
- --batch-size 8 \
- --seq-length 1024 \
- --max-position-embeddings 1024 \
- --log-interval 10 \
- --fp16 \
- --no-load-optim \
- --no-load-rng
diff --git a/examples/export/README.md b/examples/export/README.md
new file mode 100644
index 0000000000..ddb8216f94
--- /dev/null
+++ b/examples/export/README.md
@@ -0,0 +1,10 @@
+# Megatron Core Export
+
+This module is used to export megatron core models to different inference frameworks.
+Currently we support TRTLLM export . In the future we will be adding support for VLLM etc.
+
+## PTQ AND EXPORT
+Follow the instructions in [ptq_and_trtllm_export](./ptq_and_trtllm_export) to do post training quantization, followed by an export to TRTLLM format.
+
+# TRTLLM EXPORT
+Follow the instructions in [trtllm_export](./trtllm_export/) to do export to TRTLLM checkpoint format alone.
\ No newline at end of file
diff --git a/examples/export/knowledge_distillation/pretrain_gpt_modelopt.py b/examples/export/knowledge_distillation/pretrain_gpt_modelopt.py
new file mode 100644
index 0000000000..65d0727d8c
--- /dev/null
+++ b/examples/export/knowledge_distillation/pretrain_gpt_modelopt.py
@@ -0,0 +1,136 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Pretrain GPT."""
+import os
+import sys
+from functools import partial
+
+# This file isn't located in project root, but to import, it should pretend to be.
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../")))
+
+from megatron.core import mpu
+from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
+from megatron.core.datasets.gpt_dataset import GPTDataset, GPTDatasetConfig, MockGPTDataset
+from megatron.core.datasets.utils import get_blend_from_list
+from megatron.core.enums import ModelType
+from megatron.core.models.gpt import GPTModel
+from megatron.core.utils import StragglerDetector
+from megatron.inference.arguments import add_modelopt_args
+from megatron.inference.gpt import loss_func, model_provider
+from megatron.training import get_args, get_timers, get_tokenizer, pretrain
+from megatron.training.utils import (
+ get_batch_on_this_cp_rank,
+ get_batch_on_this_tp_rank,
+ print_rank_0,
+)
+
+stimer = StragglerDetector()
+
+
+def get_batch(data_iterator):
+ """Generate a batch."""
+
+ # TODO: this is pretty hacky, find a better way
+ if (not mpu.is_pipeline_first_stage()) and (not mpu.is_pipeline_last_stage()):
+ return None, None, None, None, None
+
+ # get batches based on the TP rank you are on
+ batch = get_batch_on_this_tp_rank(data_iterator)
+
+ # slice batch along sequence dimension for context parallelism
+ batch = get_batch_on_this_cp_rank(batch)
+
+ return batch.values()
+
+
+def forward_step(data_iterator, model: GPTModel):
+ """Forward training step.
+
+ Args:
+ data_iterator : Input data iterator
+ model (GPTModel): The GPT Model
+ """
+ timers = get_timers()
+
+ # Get the batch.
+ timers('batch-generator', log_level=2).start()
+ global stimer
+ with stimer(bdata=True):
+ tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data_iterator)
+ timers('batch-generator').stop()
+
+ with stimer:
+ output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
+
+ # [ModelOpt]: model is needed to access ModelOpt distillation losses
+ return output_tensor, partial(loss_func, loss_mask, model)
+
+
+def is_dataset_built_on_rank():
+ return (
+ mpu.is_pipeline_first_stage() or mpu.is_pipeline_last_stage()
+ ) and mpu.get_tensor_model_parallel_rank() == 0
+
+
+def core_gpt_dataset_config_from_args(args):
+ tokenizer = get_tokenizer()
+
+ return GPTDatasetConfig(
+ random_seed=args.seed,
+ sequence_length=args.seq_length,
+ blend=get_blend_from_list(args.data_path),
+ blend_per_split=[
+ get_blend_from_list(args.train_data_path),
+ get_blend_from_list(args.valid_data_path),
+ get_blend_from_list(args.test_data_path),
+ ],
+ split=args.split,
+ num_dataset_builder_threads=args.num_dataset_builder_threads,
+ path_to_cache=args.data_cache_path,
+ mmap_bin_files=args.mmap_bin_files,
+ tokenizer=tokenizer,
+ reset_position_ids=args.reset_position_ids,
+ reset_attention_mask=args.reset_attention_mask,
+ eod_mask_loss=args.eod_mask_loss,
+ create_attention_mask=args.create_attention_mask_in_dataloader,
+ )
+
+
+def train_valid_test_datasets_provider(train_val_test_num_samples):
+ """Build the train test and validation datasets.
+
+ Args:
+ train_val_test_num_samples : A list containing the number of samples in train test and validation.
+ """
+ args = get_args()
+
+ config = core_gpt_dataset_config_from_args(args)
+
+ if args.mock_data:
+ dataset_type = MockGPTDataset
+ else:
+ dataset_type = GPTDataset
+
+ print_rank_0("> building train, validation, and test datasets for GPT ...")
+
+ train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(
+ dataset_type, train_val_test_num_samples, is_dataset_built_on_rank, config
+ ).build()
+
+ print_rank_0("> finished creating GPT datasets ...")
+
+ return train_ds, valid_ds, test_ds
+
+
+if __name__ == "__main__":
+ # Temporary for transition to core datasets
+ train_valid_test_datasets_provider.is_distributed = True
+
+ pretrain(
+ train_valid_test_datasets_provider,
+ model_provider,
+ ModelType.encoder_or_decoder,
+ forward_step,
+ args_defaults={"tokenizer_type": "GPT2BPETokenizer"},
+ extra_args_provider=add_modelopt_args,
+ )
diff --git a/examples/export/ptq_and_trtllm_export/README.md b/examples/export/ptq_and_trtllm_export/README.md
new file mode 100644
index 0000000000..2605910869
--- /dev/null
+++ b/examples/export/ptq_and_trtllm_export/README.md
@@ -0,0 +1,295 @@
+# Megatron Model Optimization and Deployment
+
+## Installation
+We recommend that users follow TensorRT-LLM's official installation guide to build it from source
+and proceed with a containerized environment (`docker.io/tensorrt_llm/release:latest`):
+
+```sh
+git clone https://github.com/NVIDIA/TensorRT-LLM.git
+cd TensorRT-LLM
+git checkout v0.10.0
+make -C docker release_build
+```
+
+> **TROUBLE SHOOTING:** rather than copying each folder separately in `docker/Dockerfile.multi`,
+> you may need to copy the entire dir as `COPY ./ /src/tensorrt_llm` since a `git submodule` is
+> called later which requires `.git` to continue.
+
+Once the container is built, install `nvidia-modelopt` and additional dependencies for sharded checkpoint support:
+```sh
+pip install "nvidia-modelopt[all]~=0.13.0" --extra-index-url https://pypi.nvidia.com
+pip install zarr tensorstore!=0.1.46
+```
+TensorRT-LLM quantization functionalities are currently packaged in `nvidia-modelopt`.
+You can find more documentation about `nvidia-modelopt` [here](https://nvidia.github.io/TensorRT-Model-Optimizer/).
+
+## Support Matrix
+
+The following matrix shows the current support for the PTQ + TensorRT-LLM export flow.
+
+| model | fp16 | int8_sq | fp8 | int4_awq |
+|-----------------------------|------|---------| ----| -------- |
+| nextllm-2b | x | x | x | |
+| nemotron3-8b | x | | x | |
+| nemotron3-15b | x | | x | |
+| llama2-text-7b | x | x | x | TP2 |
+| llama2-chat-70b | x | x | x | TP4 |
+
+Our PTQ + TensorRT-LLM flow has native support on MCore `GPTModel` with a mixed layer spec (native ParallelLinear
+and Transformer-Engine Norm (`TENorm`). Note that this is not the default mcore gpt spec. You can still load the
+following checkpoint formats with some remedy:
+
+| GPTModel | sharded | remedy arguments |
+|-----------------------------------|---------|---------------------------------------------|
+| megatron.legacy.model | | `--export-legacy-megatron` |
+| TE-Fused (default mcore gpt spec) | | `--export-te-mcore-model` |
+| TE-Fused (default mcore gpt spec) | x | |
+
+> **TROUBLE SHOOTING:** If you are trying to load an unpacked `.nemo` sharded checkpoint, then typically you will
+> need to adding `additional_sharded_prefix="model."` to `modelopt_load_checkpoint()` since NeMo has an additional
+> `model.` wrapper on top of the `GPTModel`.
+
+> **NOTE:** flag `--export-legacy-megatron` may not work on all legacy checkpoint versions.
+
+## Examples
+
+> **NOTE:** we only provide a simple text generation script to test the generated TensorRT-LLM engines. For
+> a production-level API server or enterprise support, see [NeMo](https://github.com/NVIDIA/NeMo) and TensorRT-LLM's
+> backend for [NVIDIA Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server).
+
+### Minitron-8B FP8 Quantization and TensorRT-LLM Deployment
+First download the nemotron checkpoint from https://huggingface.co/nvidia/Minitron-8B-Base, extract the
+sharded checkpoint from the `.nemo` tarbal and fix the tokenizer file name.
+
+> **NOTE:** The following cloning method uses `ssh`, and assume you have registered the `ssh-key` in Hugging Face.
+> If you are want to clone with `https`, then `git clone https://huggingface.co/nvidia/Minitron-8B-Base` with an access token.
+
+```sh
+git lfs install
+git clone git@hf.co:nvidia/Minitron-8B-Base
+cd Minitron-8B-Base/nemo
+tar -xvf minitron-8b-base.nemo
+cd ../..
+```
+
+Now launch the PTQ + TensorRT-LLM export script,
+```sh
+bash examples/export/ptq_and_trtllm_export/ptq_trtllm_minitron_8b ./Minitron-8B-Base None
+```
+By default, `cnn_dailymail` is used for calibration. The `GPTModel` will have quantizers for simulating the
+quantization effect. The checkpoint will be saved optionally (with quantizers as additional states) and can
+be restored for further evaluation or quantization-aware training. TensorRT-LLM checkpoint and engine are exported to `/tmp/trtllm_ckpt` and
+built in `/tmp/trtllm_engine` by default.
+
+The script expects `${CHECKPOINT_DIR}` (`./Minitron-8B-Base/nemo`) to have the following structure:
+
+> **NOTE:** The .nemo checkpoint after extraction (including examples below) should all have the following strucure.
+
+```
+├── model_weights
+│ ├── common.pt
+│ ...
+│
+├── model_config.yaml
+│...
+```
+
+> **NOTE:** The script is using `TP=8`. Change `$TP` in the script if your checkpoint has a different tensor
+> model parallelism.
+
+Then build TensorRT engine and run text generation example using the newly built TensorRT engine
+
+```sh
+export trtllm_options=" \
+ --checkpoint_dir /tmp/trtllm_ckpt \
+ --output_dir /tmp/trtllm_engine \
+ --max_input_len 2048 \
+ --max_seq_len 512 \
+ --max_batch_size 8 "
+
+trtllm-build ${trtllm_options}
+
+python examples/export/ptq_and_trtllm_export/trtllm_text_generation.py --tokenizer nvidia/Minitron-8B-Base
+```
+
+### mistral-12B FP8 Quantization and TensorRT-LLM Deployment
+First download the nemotron checkpoint from https://huggingface.co/nvidia/Mistral-NeMo-12B-Base, extract the
+sharded checkpoint from the `.nemo` tarbal.
+
+> **NOTE:** The following cloning method uses `ssh`, and assume you have registered the `ssh-key` in Hugging Face.
+> If you are want to clone with `https`, then `git clone https://huggingface.co/nvidia/Mistral-NeMo-12B-Base` with an access token.
+
+```sh
+git lfs install
+git clone git@hf.co:nvidia/Mistral-NeMo-12B-Base
+cd Mistral-NeMo-12B-Base
+tar -xvf Mistral-NeMo-12B-Base.nemo
+cd ..
+```
+
+Then log in to huggingface so that you can access to model
+
+> **NOTE:** You need a token generated from huggingface.co/settings/tokens and access to mistralai/Mistral-Nemo-Base-2407 on huggingface
+
+```sh
+pip install -U "huggingface_hub[cli]"
+huggingface-cli login
+```
+
+Now launch the PTQ + TensorRT-LLM checkpoint export script,
+
+```sh
+bash examples/export/ptq_and_trtllm_export/ptq_trtllm_mistral_12b.sh ./Mistral-NeMo-12B-Base None
+```
+
+Then build TensorRT engine and run text generation example using the newly built TensorRT engine
+
+```sh
+export trtllm_options=" \
+ --checkpoint_dir /tmp/trtllm_ckpt \
+ --output_dir /tmp/trtllm_engine \
+ --max_input_len 2048 \
+ --max_seq_len 512 \
+ --max_batch_size 8 "
+
+trtllm-build ${trtllm_options}
+
+python examples/export/ptq_and_trtllm_export/trtllm_text_generation.py --tokenizer mistralai/Mistral-Nemo-Base-2407
+```
+
+
+### llama2-text-7b INT8 SmoothQuant and TensorRT-LLM Deployment
+> **NOTE:** Due to the LICENSE issue, we do not provide a MCore checkpoint to download. Users can follow
+> the instruction in `docs/llama2.md` to convert the checkpoint to megatron legacy `GPTModel` format and
+> use `--export-legacy-megatron` flag which will remap the checkpoint to the MCore `GPTModel` spec
+> that we support.
+
+```sh
+bash examples/export/ptq_and_trtllm_export/ptq_trtllm_llama_7b.sh ${CHECKPOINT_DIR}
+```
+
+The script expect `${CHECKPOINT_DIR}` to have the following structure:
+```
+├── hf
+│ ├── tokenizer.config
+│ ├── tokenizer.model
+│ ...
+│
+├── iter_0000001
+│ ├── mp_rank_00
+│ ...
+│
+├── latest_checkpointed_iteration.txt
+```
+In short, other than the converted llama megatron checkpoint, also put the Hugging Face checkpoint inside as
+the source of the tokenizer.
+
+Then build TensorRT engine and run text generation example using the newly built TensorRT engine
+
+```sh
+export trtllm_options=" \
+ --checkpoint_dir /tmp/trtllm_ckpt \
+ --output_dir /tmp/trtllm_engine \
+ --max_input_len 2048 \
+ --max_seq_len 512 \
+ --max_batch_size 8 "
+
+trtllm-build ${trtllm_options}
+
+python examples/export/ptq_and_trtllm_export/trtllm_text_generation.py --tokenizer meta-llama/Llama-2-7b
+```
+
+### llama3-8b / llama3.1-8b INT8 SmoothQuant and TensorRT-LLM Deployment
+> **NOTE:** For llama3.1, the missing rope_scaling parameter will be fixed in modelopt-0.19 and trtllm-0.13.
+
+> **NOTE:** There are two ways to acquire the checkpoint. Users can follow
+> the instruction in `docs/llama2.md` to convert the checkpoint to megatron legacy `GPTModel` format and
+> use `--export-legacy-megatron` flag which will remap the checkpoint to the MCore `GPTModel` spec
+> that we support.
+> Or Users can download [nemo model](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/llama38bnemo) from NGC and extract the sharded checkpoint from the .nemo tarbal.
+
+If users choose to download the model from NGC, first extract the sharded checkpoint from the .nemo tarbal.
+
+```sh
+tar -xvf 8b_pre_trained_bf16.nemo
+```
+
+> **NOTE:** You need a token generated from huggingface.co/settings/tokens and access to meta-llama/Llama-3.1-8B or meta-llama/Llama-3-8B on huggingface
+
+```sh
+pip install -U "huggingface_hub[cli]"
+huggingface-cli login
+```
+
+Now launch the PTQ + TensorRT-LLM checkpoint export script for llama-3,
+
+```sh
+bash examples/export/ptq_and_trtllm_export/ptq_trtllm_llama3_8b.sh ./llama-3-8b-nemo_v1.0 None
+```
+
+or llama-3.1
+
+```sh
+bash examples/export/ptq_and_trtllm_export/ptq_trtllm_llama3_1_8b.sh ./llama-3_1-8b-nemo_v1.0 None
+```
+
+Then build TensorRT engine and run text generation example using the newly built TensorRT engine
+
+```sh
+export trtllm_options=" \
+ --checkpoint_dir /tmp/trtllm_ckpt \
+ --output_dir /tmp/trtllm_engine \
+ --max_input_len 2048 \
+ --max_seq_len 512 \
+ --max_batch_size 8 "
+
+trtllm-build ${trtllm_options}
+
+python examples/export/ptq_and_trtllm_export/trtllm_text_generation.py --tokenizer meta-llama/Meta-Llama-3-8B
+# For llama-3
+
+python examples/export/ptq_and_trtllm_export/trtllm_text_generation.py --tokenizer meta-llama/Meta-Llama-3.1-8B
+#For llama-3.1
+```
+
+
+### Mixtral-8x7B FP8 Quantization and TensorRT-LLM Deployment
+First download the nemotron checkpoint from https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/mixtral-8x7b-v01, extract the
+sharded checkpoint from the `.nemo` tarbal.
+
+```sh
+ngc registry model download-version "nvidia/nemo/mixtral-8x7b-v01:1.0"
+cd mixtral-8x7b-v01_v1.0
+tar -xvf mixtral.nemo
+cd ..
+```
+
+Then log in to huggingface so that you can access to model
+
+> **NOTE:** You need a token generated from huggingface.co/settings/tokens and access to mistralai/Mixtral-8x7B-v0.1 on huggingface
+
+```sh
+pip install -U "huggingface_hub[cli]"
+huggingface-cli login
+```
+
+Now launch the PTQ + TensorRT-LLM checkpoint export script,
+
+```sh
+bash examples/export/ptq_and_trtllm_export/ptq_trtllm_mixtral_8x7b.sh ./mixtral-8x7b-v01_v1.0/
+```
+
+Then build TensorRT engine and run text generation example using the newly built TensorRT engine
+
+```sh
+export trtllm_options=" \
+ --checkpoint_dir /tmp/trtllm_ckpt \
+ --output_dir /tmp/trtllm_engine \
+ --max_input_len 2048 \
+ --max_seq_len 512 \
+ --max_batch_size 8 "
+
+trtllm-build ${trtllm_options}
+
+python examples/export/ptq_and_trtllm_export/trtllm_text_generation.py --tokenizer mistralai/Mixtral-8x7B-v0.1
+```
diff --git a/examples/export/ptq_and_trtllm_export/ptq_trtllm_llama2_7b.sh b/examples/export/ptq_and_trtllm_export/ptq_trtllm_llama2_7b.sh
new file mode 100644
index 0000000000..ebcc448955
--- /dev/null
+++ b/examples/export/ptq_and_trtllm_export/ptq_trtllm_llama2_7b.sh
@@ -0,0 +1,80 @@
+#!/bin/bash
+set -e
+
+DEFAULT_NAME="/checkpoints/llama2-text-7b_v0.2.0"
+NAME="${1:-$DEFAULT_NAME}"
+
+DEFAULT_QUANT_CFG="int8_sq"
+QUANT_CFG="${2:-$DEFAULT_QUANT_CFG}"
+
+# CHANGE THE FOLLOWING IF YOU MOUNT YOUR DATA AND CHECKPOINTS DIFFERENTLY IN THE CONTAINER.
+TP="8"
+INFERENCE_TP=${TP}
+DECODER_TYPE="llama"
+CHECKPOINT_LOAD_DIR="${NAME}"
+TOKENIZER_MODEL="${CHECKPOINT_LOAD_DIR}/hf/tokenizer.model"
+
+# LLaMA2 text 7b has ffn_hidden_size 11008. int4_awq requires a block_size of 128 as a result the TP can at most be 2
+if [ "$QUANT_CFG" = "int4_awq" ]; then
+ INFERENCE_TP="2"
+fi
+
+additional_options=" \
+ --export-quant-cfg ${QUANT_CFG} \
+ --export-legacy-megatron \
+ --export-te-mcore-model \
+ --calib-batch-size 8 \
+ --decoder ${DECODER_TYPE} \
+ --export-dir /tmp/trtllm_ckpt \
+ --inference-tensor-parallel ${INFERENCE_TP} "
+
+trtllm_options=" \
+ --tensorrt-llm-checkpoint-dir /tmp/trtllm_ckpt \
+ --engine-dir /tmp/trtllm_engine \
+ --tokenizer ${CHECKPOINT_LOAD_DIR}/hf \
+ --max-input-len 2048 \
+ --max-output-len 512 \
+ --max-batch-size 8 "
+
+# DO NOT CHANGE THE SETTING BELOW UNLESS YOU KNOW WHAT YOU ARE DOING!!!
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+options=" \
+ --disable-bias-linear \
+ --swiglu \
+ --no-rope-fusion \
+ --untie-embeddings-and-output-weights \
+ --use-rotary-position-embeddings \
+ --normalization RMSNorm \
+ --rotary-percent 1.0 \
+ --no-position-embedding \
+ --no-masked-softmax-fusion \
+ --no-bias-gelu-fusion \
+ --no-bias-dropout-fusion \
+ --no-async-tensor-model-parallel-allreduce \
+ --tensor-model-parallel-size ${TP} \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --ffn-hidden-size 11008 \
+ --num-attention-heads 32 \
+ --seq-length 4096 \
+ --max-position-embeddings 4096 \
+ --micro-batch-size 1 \
+ --make-vocab-size-divisible-by 1 \
+ --tokenizer-type Llama2Tokenizer \
+ --tokenizer-model ${TOKENIZER_MODEL} \
+ --save-interval 1000000 \
+ --use-dist-ckpt \
+ --load ${CHECKPOINT_LOAD_DIR} \
+ --fp16"
+
+# Precompile CUDA extentions
+python -c "import modelopt.torch.quantization.extensions as ext; print(ext.cuda_ext); print(ext.cuda_ext_fp8)"
+
+# Acquire launch configuration where variable launch_config will be set
+launch_config="--nproc_per_node=${TP}"
+
+# Launch multi-process with torchrun
+torchrun ${launch_config} examples/export/ptq_and_trtllm_export/text_generation_ptq.py ${options} ${additional_options}
+
diff --git a/examples/export/ptq_and_trtllm_export/ptq_trtllm_llama3_1_8b.sh b/examples/export/ptq_and_trtllm_export/ptq_trtllm_llama3_1_8b.sh
new file mode 100644
index 0000000000..94ee12db41
--- /dev/null
+++ b/examples/export/ptq_and_trtllm_export/ptq_trtllm_llama3_1_8b.sh
@@ -0,0 +1,75 @@
+#!/bin/bash
+set -e
+
+DEFAULT_NAME="/checkpoints/llama-3_1-8b-nemo_v1.0"
+NAME="${1:-$DEFAULT_NAME}"
+
+DEFAULT_QUANT_CFG="int8_sq"
+QUANT_CFG="${2:-$DEFAULT_QUANT_CFG}"
+
+# CHANGE THE FOLLOWING IF YOU MOUNT YOUR DATA AND CHECKPOINTS DIFFERENTLY IN THE CONTAINER.
+TP="1"
+INFERENCE_TP=${TP}
+DECODER_TYPE="llama"
+CHECKPOINT_LOAD_DIR="${NAME}"
+
+# LLaMA2 text 7b has ffn_hidden_size 11008. int4_awq requires a block_size of 128 as a result the TP can at most be 2
+if [ "$QUANT_CFG" = "int4_awq" ]; then
+ INFERENCE_TP="2"
+fi
+
+additional_options=" \
+ --export-quant-cfg ${QUANT_CFG} \
+ --export-legacy-megatron \
+ --export-te-mcore-model \
+ --calib-batch-size 8 \
+ --decoder ${DECODER_TYPE} \
+ --export-dir /tmp/trtllm_ckpt \
+ --inference-tensor-parallel ${INFERENCE_TP} "
+
+# DO NOT CHANGE THE SETTING BELOW UNLESS YOU KNOW WHAT YOU ARE DOING!!!
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+options=" \
+ --disable-bias-linear \
+ --attention-backend unfused \
+ --swiglu \
+ --no-rope-fusion \
+ --untie-embeddings-and-output-weights \
+ --use-rotary-position-embeddings \
+ --normalization RMSNorm \
+ --rotary-percent 1.0 \
+ --hidden-dropout 0.0 \
+ --attention-dropout 0.0 \
+ --no-bias-gelu-fusion \
+ --no-bias-dropout-fusion \
+ --no-async-tensor-model-parallel-allreduce \
+ --tensor-model-parallel-size ${TP} \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --ffn-hidden-size 14336 \
+ --num-attention-heads 32 \
+ --seq-length 131072 \
+ --max-position-embeddings 131072 \
+ --micro-batch-size 4 \
+ --make-vocab-size-divisible-by 128 \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model meta-llama/Meta-Llama-3.1-8B \
+ --save-interval 1000000 \
+ --use-rope-scaling \
+ --use-dist-ckpt \
+ --load ${CHECKPOINT_LOAD_DIR} \
+ --rotary-base 500000 \
+ --fp16"
+
+# Precompile CUDA extentions
+python -c "import modelopt.torch.quantization.extensions as ext; print(ext.cuda_ext); print(ext.cuda_ext_fp8)"
+
+# Acquire launch configuration where variable launch_config will be set
+launch_config="--nproc_per_node=${TP}"
+
+# Launch multi-process with torchrun
+torchrun ${launch_config} examples/export/ptq_and_trtllm_export/text_generation_ptq.py ${options} ${additional_options}
diff --git a/examples/export/ptq_and_trtllm_export/ptq_trtllm_llama3_8b.sh b/examples/export/ptq_and_trtllm_export/ptq_trtllm_llama3_8b.sh
new file mode 100644
index 0000000000..dfa5a80c26
--- /dev/null
+++ b/examples/export/ptq_and_trtllm_export/ptq_trtllm_llama3_8b.sh
@@ -0,0 +1,75 @@
+#!/bin/bash
+set -e
+
+DEFAULT_NAME="/checkpoints/llama-3_1-8b-nemo_v1.0"
+NAME="${1:-$DEFAULT_NAME}"
+
+DEFAULT_QUANT_CFG="int8_sq"
+QUANT_CFG="${2:-$DEFAULT_QUANT_CFG}"
+
+
+# CHANGE THE FOLLOWING IF YOU MOUNT YOUR DATA AND CHECKPOINTS DIFFERENTLY IN THE CONTAINER.
+TP="1"
+INFERENCE_TP=${TP}
+DECODER_TYPE="llama"
+CHECKPOINT_LOAD_DIR="${NAME}"
+
+# LLaMA2 text 7b has ffn_hidden_size 11008. int4_awq requires a block_size of 128 as a result the TP can at most be 2
+if [ "$QUANT_CFG" = "int4_awq" ]; then
+ INFERENCE_TP="2"
+fi
+
+additional_options=" \
+ --export-quant-cfg ${QUANT_CFG} \
+ --export-legacy-megatron \
+ --export-te-mcore-model \
+ --calib-batch-size 8 \
+ --decoder ${DECODER_TYPE} \
+ --export-dir /tmp/trtllm_ckpt \
+ --inference-tensor-parallel ${INFERENCE_TP} "
+
+# DO NOT CHANGE THE SETTING BELOW UNLESS YOU KNOW WHAT YOU ARE DOING!!!
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+options=" \
+ --disable-bias-linear \
+ --attention-backend unfused \
+ --swiglu \
+ --no-rope-fusion \
+ --untie-embeddings-and-output-weights \
+ --use-rotary-position-embeddings \
+ --normalization RMSNorm \
+ --rotary-percent 1.0 \
+ --hidden-dropout 0.0 \
+ --attention-dropout 0.0 \
+ --no-bias-gelu-fusion \
+ --no-bias-dropout-fusion \
+ --no-async-tensor-model-parallel-allreduce \
+ --tensor-model-parallel-size ${TP} \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --ffn-hidden-size 14336 \
+ --num-attention-heads 32 \
+ --seq-length 8192 \
+ --max-position-embeddings 8192 \
+ --micro-batch-size 4 \
+ --make-vocab-size-divisible-by 128 \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model meta-llama/Meta-Llama-3-8B \
+ --save-interval 1000000 \
+ --use-dist-ckpt \
+ --load ${CHECKPOINT_LOAD_DIR} \
+ --rotary-base 500000 \
+ --fp16"
+
+# Precompile CUDA extentions
+python -c "import modelopt.torch.quantization.extensions as ext; print(ext.cuda_ext); print(ext.cuda_ext_fp8)"
+
+# Acquire launch configuration where variable launch_config will be set
+launch_config="--nproc_per_node=${TP}"
+
+# Launch multi-process with torchrun
+torchrun ${launch_config} examples/export/ptq_and_trtllm_export/text_generation_ptq.py ${options} ${additional_options}
diff --git a/examples/export/ptq_and_trtllm_export/ptq_trtllm_minitron_8b.sh b/examples/export/ptq_and_trtllm_export/ptq_trtllm_minitron_8b.sh
new file mode 100644
index 0000000000..6e57972e30
--- /dev/null
+++ b/examples/export/ptq_and_trtllm_export/ptq_trtllm_minitron_8b.sh
@@ -0,0 +1,70 @@
+#!/bin/bash
+set -e
+
+DEFAULT_NAME="/checkpoints/nemotron3-8b_v0.3.0"
+NAME="${1:-$DEFAULT_NAME}"
+
+DEFAULT_QUANT_CFG="fp8"
+QUANT_CFG="${2:-$DEFAULT_QUANT_CFG}"
+
+# CHANGE THE FOLLOWING IF YOU MOUNT YOUR DATA AND CHECKPOINTS DIFFERENTLY IN THE CONTAINER.
+TP="8"
+INFERENCE_TP=${TP}
+DECODER_TYPE="gptnext"
+CHECKPOINT_LOAD_DIR="${NAME}/nemo"
+
+if [ "$QUANT_CFG" = "int4_awq" ]; then
+ INFERENCE_TP="1"
+fi
+
+additional_options=" \
+ --export-quant-cfg ${QUANT_CFG} \
+ --export-legacy-megatron \
+ --export-te-mcore-model \
+ --calib-batch-size 8 \
+ --decoder ${DECODER_TYPE} \
+ --export-dir /tmp/trtllm_ckpt \
+ --inference-tensor-parallel ${INFERENCE_TP} "
+
+# DO NOT CHANGE THE SETTING BELOW UNLESS YOU KNOW WHAT YOU ARE DOING!!!
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+options=" \
+ --apply-layernorm-1p \
+ --attn-attention unfused \
+ --untie-embeddings-and-output-weights \
+ --disable-bias-linear \
+ --no-rope-fusion \
+ --no-position-embedding \
+ --use-rotary-position-embeddings \
+ --rotary-percent 0.5 \
+ --squared-relu \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --tensor-model-parallel-size ${TP} \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --ffn-hidden-size 16384 \
+ --group-query-attention \
+ --num-attention-heads 48 \
+ --kv-channels 128 \
+ --seq-length 4096 \
+ --num-query-groups 8 \
+ --max-position-embeddings 4096 \
+ --micro-batch-size 4 \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model nvidia/Minitron-8B-Base \
+ --save-interval 1000000 \
+ --load ${CHECKPOINT_LOAD_DIR} \
+ --bf16 \
+ --use-dist-ckpt"
+
+# Precompile CUDA extentions
+python -c "import modelopt.torch.quantization.extensions as ext; print(ext.cuda_ext); print(ext.cuda_ext_fp8)"
+
+# Acquire launch configuration where variable launch_config will be set
+launch_config="--nproc_per_node=${TP}"
+
+# Launch multi-process with torchrun
+torchrun ${launch_config} examples/export/ptq_and_trtllm_export/text_generation_ptq.py ${options} ${additional_options}
diff --git a/examples/export/ptq_and_trtllm_export/ptq_trtllm_mistral_12b.sh b/examples/export/ptq_and_trtllm_export/ptq_trtllm_mistral_12b.sh
new file mode 100644
index 0000000000..8469945f08
--- /dev/null
+++ b/examples/export/ptq_and_trtllm_export/ptq_trtllm_mistral_12b.sh
@@ -0,0 +1,71 @@
+#!/bin/bash
+set -e
+
+DEFAULT_NAME="/checkpoints/Mistral-NeMo-12B-Base"
+NAME="${1:-$DEFAULT_NAME}"
+
+DEFAULT_QUANT_CFG="fp8"
+QUANT_CFG="${2:-$DEFAULT_QUANT_CFG}"
+
+# CHANGE THE FOLLOWING IF YOU MOUNT YOUR DATA AND CHECKPOINTS DIFFERENTLY IN THE CONTAINER.
+TP="8"
+INFERENCE_TP=${TP}
+DECODER_TYPE="llama"
+CHECKPOINT_LOAD_DIR="${NAME}"
+
+if [ "$QUANT_CFG" = "int4_awq" ]; then
+ INFERENCE_TP="1"
+fi
+
+additional_options=" \
+ --export-quant-cfg ${QUANT_CFG} \
+ --export-legacy-megatron \
+ --export-te-mcore-model \
+ --calib-batch-size 8 \
+ --decoder ${DECODER_TYPE} \
+ --export-dir /tmp/trtllm_ckpt \
+ --inference-tensor-parallel ${INFERENCE_TP} "
+
+# DO NOT CHANGE THE SETTING BELOW UNLESS YOU KNOW WHAT YOU ARE DOING!!!
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+options=" \
+ --untie-embeddings-and-output-weights \
+ --attention-backend unfused \
+ --disable-bias-linear \
+ --use-rotary-position-embeddings \
+ --rotary-percent 1.0 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --tensor-model-parallel-size ${TP} \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 40 \
+ --hidden-size 5120 \
+ --ffn-hidden-size 14336 \
+ --num-attention-heads 32 \
+ --seq-length 8192 \
+ --kv-channels 128 \
+ --normalization RMSNorm \
+ --swiglu \
+ --num-query-groups 8 \
+ --group-query-attention \
+ --position-embedding-type rope \
+ --max-position-embeddings 8192 \
+ --micro-batch-size 1 \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tiktoken-pattern v2 \
+ --tokenizer-model mistralai/Mistral-Nemo-Base-2407 \
+ --save-interval 1000000 \
+ --load ${CHECKPOINT_LOAD_DIR} \
+ --fp16 \
+ --rotary-base 1000000 \
+ --use-dist-ckpt"
+
+# Precompile CUDA extentions
+python -c "import modelopt.torch.quantization.extensions as ext; print(ext.cuda_ext); print(ext.cuda_ext_fp8)"
+
+# Acquire launch configuration where variable launch_config will be set
+launch_config="--nproc_per_node=${TP}"
+
+# Launch multi-process with torchrun
+torchrun ${launch_config} examples/export/ptq_and_trtllm_export/text_generation_ptq.py ${options} ${additional_options}
diff --git a/examples/export/ptq_and_trtllm_export/ptq_trtllm_mixtral_8x7b.sh b/examples/export/ptq_and_trtllm_export/ptq_trtllm_mixtral_8x7b.sh
new file mode 100644
index 0000000000..d2a4edee47
--- /dev/null
+++ b/examples/export/ptq_and_trtllm_export/ptq_trtllm_mixtral_8x7b.sh
@@ -0,0 +1,84 @@
+#!/bin/bash
+set -e
+
+DEFAULT_NAME="/checkpoints/Mistral-NeMo-12B-Base"
+NAME="${1:-$DEFAULT_NAME}"
+
+DEFAULT_QUANT_CFG="fp8"
+QUANT_CFG="${2:-$DEFAULT_QUANT_CFG}"
+
+# NOTE: UNFUSED ATTENTION MUST BE USED TO AVOID ADDITIONAL STATE_DICT KEY MISMATCH.
+export NVTE_FLASH_ATTN=0
+export NVTE_FUSED_ATTN=0
+export NVTE_UNFUSED_ATTN=1
+
+# CHANGE THE FOLLOWING IF YOU MOUNT YOUR DATA AND CHECKPOINTS DIFFERENTLY IN THE CONTAINER.
+TP="8"
+INFERENCE_TP=${TP}
+DECODER_TYPE="llama"
+CHECKPOINT_LOAD_DIR="${NAME}"
+
+if [ "$QUANT_CFG" = "int4_awq" ]; then
+ INFERENCE_TP="1"
+fi
+
+additional_options=" \
+ --export-quant-cfg ${QUANT_CFG} \
+ --export-legacy-megatron \
+ --export-te-mcore-model \
+ --calib-batch-size 8 \
+ --decoder ${DECODER_TYPE} \
+ --export-dir /tmp/trtllm_ckpt \
+ --inference-tensor-parallel ${INFERENCE_TP} "
+
+# DO NOT CHANGE THE SETTING BELOW UNLESS YOU KNOW WHAT YOU ARE DOING!!!
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+options=" \
+ --untie-embeddings-and-output-weights \
+ --no-masked-softmax-fusion \
+ --no-position-embedding \
+ --use-mcore-models \
+ --disable-bias-linear \
+ --rotary-percent 1.0 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --tensor-model-parallel-size ${TP} \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --ffn-hidden-size 14336 \
+ --num-attention-heads 32 \
+ --seq-length 4096 \
+ --kv-channels 128 \
+ --normalization RMSNorm \
+ --swiglu \
+ --num-query-groups 8 \
+ --num-experts 8 \
+ --moe-router-topk 2 \
+ --moe-aux-loss-coeff 1e-2 \
+ --moe-router-load-balancing-type aux_loss \
+ --group-query-attention \
+ --position-embedding-type rope \
+ --no-rope-fusion \
+ --max-position-embeddings 32768 \
+ --micro-batch-size 1 \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tiktoken-pattern v2 \
+ --tokenizer-model mistralai/Mixtral-8x7B-Instruct-v0.1 \
+ --save-interval 1000000 \
+ --load ${CHECKPOINT_LOAD_DIR} \
+ --bf16 \
+ --rotary-base 1000000 \
+ --use-dist-ckpt"
+
+# Precompile CUDA extentions
+python -c "import modelopt.torch.quantization.extensions as ext; print(ext.cuda_ext); print(ext.cuda_ext_fp8)"
+
+# Acquire launch configuration where variable launch_config will be set
+launch_config="--nproc_per_node=${TP}"
+
+# Launch multi-process with torchrun
+torchrun ${launch_config} examples/export/ptq_and_trtllm_export/text_generation_ptq.py ${options} ${additional_options}
+
+
diff --git a/examples/export/ptq_and_trtllm_export/text_generation_ptq.py b/examples/export/ptq_and_trtllm_export/text_generation_ptq.py
new file mode 100644
index 0000000000..c915cec790
--- /dev/null
+++ b/examples/export/ptq_and_trtllm_export/text_generation_ptq.py
@@ -0,0 +1,222 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""Sample Generate GPT."""
+import functools
+import os
+import sys
+from pathlib import Path
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../")))
+
+import modelopt.torch.quantization as mtq
+import torch
+from datasets import load_dataset
+from tqdm import tqdm
+
+# [ModelOpt]: changing the default model provider to the ModelOpt version
+from megatron.core import mpu
+from megatron.inference.arguments import add_modelopt_args
+from megatron.inference.checkpointing import load_modelopt_checkpoint
+from megatron.inference.gpt.model_provider import model_provider
+from megatron.inference.text_generation import generate_and_post_process
+from megatron.training import get_args, get_model, initialize_megatron
+from megatron.training.checkpointing import save_checkpoint
+from megatron.training.utils import print_rank_0, unwrap_model
+
+QUANT_CFG_CHOICES = {
+ "int8": mtq.INT8_DEFAULT_CFG,
+ "int8_sq": mtq.INT8_SMOOTHQUANT_CFG,
+ "fp8": mtq.FP8_DEFAULT_CFG,
+ "int4_awq": mtq.INT4_AWQ_CFG,
+ "w4a8_awq": mtq.W4A8_AWQ_BETA_CFG,
+ "int4": mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
+}
+
+
+def add_trtllm_ckpt_export_args(parser):
+ """Add additional arguments for TensorRT-LLM."""
+ group = parser.add_argument_group(title="trtllm")
+
+ group.add_argument(
+ "--export-dir", type=str, help="The output TensorRT-LLM checkpoint.",
+ )
+ group.add_argument(
+ "--decoder", type=str, choices=["gptnext", 'llama'], help="The decoder type of the model.",
+ )
+ group.add_argument(
+ "--inference-tensor-parallel",
+ type=int,
+ help="Tensor parallel for the inference time, can be different from the training config.",
+ default=1,
+ )
+
+
+def add_text_generate_ptq_args(parser):
+ """Add additional arguments for ModelOpt text generation PTQ."""
+ group = parser.add_argument_group(title='ModelOpt text generation ptq')
+ group.add_argument(
+ "--calib-dataset",
+ type=str,
+ default="cnn_dailymail",
+ help="Calibration datasets from HuggingFace datasets.",
+ )
+ group.add_argument(
+ "--calib-batch-size", type=int, default=4, help="Batch size to use for ptq calibration."
+ )
+ group.add_argument(
+ "--calib-size", type=int, default=512, help="Samples to use for ptq calibration."
+ )
+ parser.add_argument(
+ "--prompts",
+ type=str,
+ default=(
+ "Born in north-east France, Soyer trained as a|Born in California, Soyer trained as a"
+ ),
+ help="Input texts. Please use | to separate different batches.",
+ )
+ add_modelopt_args(parser)
+ add_trtllm_ckpt_export_args(parser)
+ return parser
+
+
+def get_calib_dataloader(
+ data="cnn_dailymail", batch_size=4, calib_size=512, max_sequence_length=512
+):
+ if data == "pileval":
+ dataset = load_dataset(
+ "json", data_files="https://the-eye.eu/public/AI/pile/val.jsonl.zst", split="train"
+ )
+ text_column = "text"
+ elif data == "wikitext":
+ dataset = load_dataset("wikitext", "wikitext-103-v1", split="train")
+ text_column = "text"
+ elif data == "cnn_dailymail":
+ dataset = load_dataset("cnn_dailymail", name="3.0.0", split="train")
+ text_column = "article"
+
+ calib_size = max(min(len(dataset), calib_size), batch_size)
+ for i in range(calib_size // batch_size):
+ batch = dataset[i * batch_size : (i + 1) * batch_size][text_column]
+ for j in range(len(batch)):
+ batch[j] = batch[j][:max_sequence_length]
+ yield batch
+
+
+
+if __name__ == "__main__":
+ initialize_megatron(
+ extra_args_provider=add_text_generate_ptq_args,
+ args_defaults={
+ 'tokenizer_type': 'GPT2BPETokenizer',
+ 'no_load_rng': True,
+ 'no_load_optim': True,
+ },
+ )
+
+ args = get_args()
+ if args.num_layers_per_virtual_pipeline_stage is not None:
+ print_rank_0("Interleaved pipeline schedule is not yet supported for text generation.")
+ exit()
+
+ print_rank_0("WARNING: Forcing exit_on_missing_checkpoint to True for text generation.")
+ args.exit_on_missing_checkpoint = True
+ if hasattr(args, 'moe_grouped_gemm') and args.moe_grouped_gemm == True:
+ print_rank_0("WARNING: Forcing moe_grouped_gemm to False for PTQ and export.")
+ args.moe_grouped_gemm = False
+
+ # Set up model and load checkpoint
+ # [ModelOpt]: make sure that output logits are allgathered.
+ text_generation_model_provider = functools.partial(model_provider, parallel_output=False)
+ model = get_model(text_generation_model_provider, wrap_with_ddp=False)
+
+ if args.load is not None:
+ load_modelopt_checkpoint(model, strict=not args.untie_embeddings_and_output_weights)
+ print_rank_0("Done loading checkpoint")
+
+ # Removing virtual pipeline parallel and other wrapper
+ assert len(model) == 1, "Above condition should have caught this"
+ unwrapped_model = unwrap_model(model)
+
+ all_prompts = args.prompts.split("|")
+
+ def custom_prompt_forward_loop_func(model):
+ for prompt in tqdm(all_prompts):
+ if mpu.is_pipeline_first_stage() and mpu.get_tensor_model_parallel_rank() == 0:
+ (
+ prompts_plus_generations,
+ prompts_plus_generations_segments,
+ logprobs,
+ _,
+ ) = generate_and_post_process(
+ model,
+ prompts=[prompt],
+ tokens_to_generate=128,
+ return_output_log_probs=True,
+ temperature=1.0,
+ )
+ print_rank_0(prompts_plus_generations)
+ else:
+ generate_and_post_process(model)
+
+ def hf_dataset_forword_loop_func(model):
+ dataloader = get_calib_dataloader(args.calib_dataset, args.calib_batch_size, args.calib_size)
+ for prompts in tqdm(dataloader, total=args.calib_size//args.calib_batch_size):
+ if mpu.is_pipeline_first_stage() and mpu.get_tensor_model_parallel_rank() == 0:
+ (
+ prompts_plus_generations,
+ prompts_plus_generations_segments,
+ logprobs,
+ _,
+ ) = generate_and_post_process(
+ model,
+ prompts=prompts,
+ tokens_to_generate=0,
+ return_output_log_probs=False,
+ temperature=1.0,
+ )
+ else:
+ generate_and_post_process(model)
+
+ ptq_forward_loop_func = custom_prompt_forward_loop_func
+ if args.calib_dataset is not None:
+ ptq_forward_loop_func = hf_dataset_forword_loop_func
+
+ if args.export_quant_cfg in QUANT_CFG_CHOICES:
+ mtq_config = QUANT_CFG_CHOICES[args.export_quant_cfg]
+ if "*output_layer*" not in mtq_config["quant_cfg"]:
+ mtq_config["quant_cfg"]["*output_layer*"] = {"enable": False}
+ if "awq" in args.export_quant_cfg:
+ weight_quantizer = mtq_config["quant_cfg"]["*weight_quantizer"] # type: ignore
+ if isinstance(weight_quantizer, list):
+ weight_quantizer = weight_quantizer[0]
+ weight_quantizer["block_sizes"][-1] = 128
+ print_rank_0("Quantizing the model...")
+ mtq.quantize(unwrapped_model[0], mtq_config, ptq_forward_loop_func)
+
+ custom_prompt_forward_loop_func(model[0])
+
+ if args.save is not None and args.export_quant_cfg in QUANT_CFG_CHOICES:
+ save_checkpoint(1, unwrapped_model, None, None, 0)
+
+ print_rank_0(f"Fake Quantized Model:\n {unwrapped_model[0]}")
+
+ if args.export_dir:
+ assert args.decoder in ["gptnext", "llama"], f"Decoder type {args.decoder} not supported."
+ Path(args.export_dir).mkdir(parents=True, exist_ok=True)
+ print_rank_0("Exporting TensorRT-LLM checkpoints.")
+
+ from modelopt.torch.export import export_tensorrt_llm_checkpoint
+
+ # In TRT LLM, squared relu activation does not support bf16. So we use fp16 by default.
+ export_tensorrt_llm_checkpoint(
+ unwrapped_model[0],
+ args.decoder,
+ torch.bfloat16 if args.bf16 else torch.float16,
+ export_dir=args.export_dir,
+ inference_tensor_parallel=args.inference_tensor_parallel,
+ inference_pipeline_parallel=1,
+ use_nfs_workspace=True,
+ )
+
+ print_rank_0(f"TensorRT-LLM checkpoints saved to {args.export_dir}")
+ torch.distributed.barrier()
diff --git a/examples/export/ptq_and_trtllm_export/trtllm_text_generation.py b/examples/export/ptq_and_trtllm_export/trtllm_text_generation.py
new file mode 100644
index 0000000000..ab8aa25a96
--- /dev/null
+++ b/examples/export/ptq_and_trtllm_export/trtllm_text_generation.py
@@ -0,0 +1,64 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+"""An example script to run the tensorrt_llm engine."""
+
+import argparse
+from pathlib import Path
+import subprocess
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from modelopt.deploy.llm import LLM
+from tensorrt_llm.models import PretrainedConfig
+from transformers import AutoTokenizer, T5Tokenizer
+import tensorrt_llm
+
+
+def parse_arguments():
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--tokenizer", type=str, default="")
+ parser.add_argument("--engine-dir", type=str, default="/tmp/trtllm_engine")
+ parser.add_argument(
+ "--input-texts",
+ type=str,
+ default=(
+ "Born in north-east France, Soyer trained as a|Born in California, Soyer trained as a"
+ ),
+ help="Input texts. Please use | to separate different batches.",
+ )
+ return parser.parse_args()
+
+
+def run(args):
+ try:
+ tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True)
+ except Exception as e:
+ raise Exception(f"Failed to load tokenizer: {e}")
+
+ print(tokenizer, tokenizer.vocab_size)
+
+ input_texts = args.input_texts.split("|")
+ assert input_texts, "input_text not specified"
+ print(input_texts)
+
+ free_memory_before = torch.cuda.mem_get_info()
+
+ # This is a ModelOpt wrapper on top of tensorrt_llm.hlapi.llm.LLM
+ llm_engine = LLM(args.engine_dir, tokenizer)
+
+ torch.cuda.cudart().cudaProfilerStart()
+ # outputs = llm_engine.generate_text(input_texts, args.max_output_len, args.max_beam_width)
+ outputs = llm_engine.generate(input_texts)
+ torch.cuda.cudart().cudaProfilerStop()
+
+ free_memory_after = torch.cuda.mem_get_info()
+ print(
+ f"Used GPU memory: {(free_memory_before[0] - free_memory_after[0]) / 1024 / 1024 / 1024} GB"
+ )
+ print(outputs)
+
+
+if __name__ == "__main__":
+ args = parse_arguments()
+ run(args)
diff --git a/examples/export/trtllm_export/README.md b/examples/export/trtllm_export/README.md
new file mode 100644
index 0000000000..52cad78583
--- /dev/null
+++ b/examples/export/trtllm_export/README.md
@@ -0,0 +1,161 @@
+# Megatron Core To TRTLLM Export Documentation
+This guide will walk you through how you can use the megatron core export for exporting models to trtllm format
+
+### Contents
+- [Megatron Core To TRTLLM Export Documentation](#megatron-core-to-trtllm-export-documentation)
+- [Contents](#contents)
+ - [1. Quick Start](#1-quick-start)
+ - [1.1 Understanding The Code](#11-understanding-the-code)
+ - [1.2 Running The Code](#12-running-the-code)
+ - [2. GPU Export](#2-gpu-export)
+ - [3. Future work](#4-future-work)
+
+#### 1. Quick Start
+This will walk you through the flow of converting an mcore gpt model to trtllm format using single device mode. The file can be found at [gpt_single_device_cpu_export.py](./single_device_export/gpt_single_device_cpu_export.py)
+
+NOTE: For faster performance, if your entire model will fit into gpu memory, pre transfer the model state dict to gpu and then call the get_trtllm_pretrained_config_and_model_weights function.
+
+
+
+##### 1.1 Understanding The Code
+***STEP 1 - We initialize model parallel and other default arguments***
+We initalize tp and pp to 1 so that we can get the full model state dict on cpu
+```python
+ initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
+```
+
+***STEP 2 - We load the model using the model_provider_function***
+NOTE: We create a simple gpt model
+
+```python
+ transformer_config = TransformerConfig(
+ num_layers=2,
+ hidden_size=64, # Needs to be atleast 32 times num_attn_heads
+ num_attention_heads=2,
+ use_cpu_initialization=True,
+ pipeline_dtype=torch.float32,
+ )
+
+ gpt_model = GPTModel(
+ config=transformer_config,
+ transformer_layer_spec=get_gpt_layer_local_spec(),
+ vocab_size=100,
+ max_sequence_length=_SEQUENCE_LENGTH,
+ )
+
+ # Optionally you can also load a model using this code
+ # sharded_state_dict=gpt_model.sharded_state_dict(prefix='')
+ # checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)
+ # gpt_model.load_state_dict(checkpoint)
+
+```
+
+***STEP 3 - Instantiate the TRTLLM Helper***
+We instantiate the [TRTLLM Helper](../../../megatron/core/export/trtllm/trtllm_helper.py) For the GPT model we instantiate trtllm_helper as shown below.
+```python
+ if hasattr(gpt_model, "rotary_pos_emb"):
+ seq_len_interpolation_factor = gpt_model.rotary_pos_emb.seq_len_interpolation_factor
+
+ trtllm_helper = TRTLLMHelper(
+ transformer_config=gpt_model.config,
+ model_type=ModelType.gpt,
+ position_embedding_type = gpt_model.position_embedding_type,
+ max_position_embeddings = gpt_model.max_position_embeddings,
+ rotary_percentage = gpt_model.rotary_percent,
+ rotary_base = gpt_model.rotary_base,
+ moe_tp_mode = 2,
+ multi_query_mode = False,
+ activation = "gelu",
+ seq_len_interpolation_factor = seq_len_interpolation_factor,
+ share_embeddings_and_output_weights=gpt_model.share_embeddings_and_output_weights
+ )
+```
+
+***STEP 4 - Get the TRTLLM Weights and configs***
+To convert model weights to trtllm weights and configs, we use the [single_device_converter](../../../megatron/core/export/trtllm/trtllm_weights_converter/single_device_trtllm_model_weights_converter.py). We pass as inputs the model state dict, and export config. In this example we use inference tp size as 2 for the export.
+
+```python
+ model_state_dict={}
+ for key , val in gpt_model.state_dict().items():
+ # val is non for _extra_state layers . We filter it out
+ if val is not None:
+ model_state_dict[key] = val
+
+ export_config = ExportConfig(inference_tp_size = 2)
+ weight_list, config_list = trtllm_helper.get_trtllm_pretrained_config_and_model_weights(
+ model_state_dict= model_state_dict,
+ dtype = DataType.bfloat16,
+ export_config=export_config
+ )
+```
+
+***STEP 5 - Build the TRTLLM Engine***
+Following code is used to build the TRTLLM Engine.
+
+```python
+ for trtllm_model_weights, trtllm_model_config in zip(weight_list, config_list):
+ trtllm_helper.build_and_save_engine(
+ max_input_len=256,
+ max_output_len=256,
+ max_batch_size=8,
+ engine_dir='/opt/megatron-lm/engine',
+ trtllm_model_weights=trtllm_model_weights,
+ trtllm_model_config=trtllm_model_config,
+ lora_ckpt_list=None,
+ use_lora_plugin=None,
+ max_lora_rank=64,
+ lora_target_modules=None,
+ max_prompt_embedding_table_size=0,
+ paged_kv_cache=True,
+ remove_input_padding=True,
+ paged_context_fmha=False,
+ use_refit=False,
+ max_num_tokens=None,
+ max_seq_len=512,
+ opt_num_tokens=None,
+ max_beam_width=1,
+ tokens_per_block=128,
+ multiple_profiles=False,
+ gpt_attention_plugin="auto",
+ gemm_plugin="auto",
+ )
+```
+
+
+##### 1.2 Running The Code
+An example run script is shown below.
+
+```
+# In a workstation
+MLM_PATH=/path/to/megatron-lm
+CONTAINER_IMAGE=gitlab-master.nvidia.com:5005/dl/joc/nemo-ci/trtllm_0.12/train:pipe.17669124-x86
+
+docker run -it --gpus=all --ipc=host -v $MLM_PATH/:/opt/megatron-lm $CONTAINER_IMAGE bash
+
+# Inside the container run the following.
+
+cd /opt/megatron-lm/
+
+CUDA_VISIBLE_DEVICES=0 torchrun --nproc-per-node 1 examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py
+```
+
+
+
+#### 2. GPU Export
+You can use the [gpt_distributed_gpu_export.py](./distributed_export/gpt_distributed_gpu_export.py) to run a more optimized on device distributed. version of trtllm export. Internally this uses the [distributed_converter](../../../megatron/core/export/trtllm/trtllm_weights_converter/distributed_trtllm_model_weights_converter.py) to convert model weights on device.
+In the single device version you collect all the model weights on CPU/GPU, convert it to trtllm format, and then store the engine back on disk. In the GPU version you load each individual state dict on the gpus, convert it on the device itself and store the engine on disk.
+
+To run the gpu version
+
+```
+CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc-per-node 2 examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py
+```
+
+
+
+#### 3. Future work
+The following are planned for the future releases .
+* Pipeline parallellism for export (Work in progress)
+* GPU Export for more models (Work in progress for some models)
+* Refit functionality
+* VLLM Support
\ No newline at end of file
diff --git a/examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py b/examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py
new file mode 100644
index 0000000000..57d44f9f62
--- /dev/null
+++ b/examples/export/trtllm_export/distributed_export/gpt_distributed_gpu_export.py
@@ -0,0 +1,117 @@
+import os
+import torch
+from megatron.core import parallel_state
+from megatron.core import dist_checkpointing
+from megatron.core.export.model_type import ModelType
+from megatron.core.export.data_type import DataType
+from megatron.core.export.trtllm.trtllm_helper import TRTLLMHelper
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
+
+
+_SEQUENCE_LENGTH = 64
+_VOCAB_SIZE = 256
+
+def initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1):
+ parallel_state.destroy_model_parallel()
+
+ # Torch setup for distributed training
+ rank = int(os.environ['LOCAL_RANK'])
+ world_size = torch.cuda.device_count()
+ torch.cuda.set_device(rank)
+ torch.distributed.init_process_group(world_size=world_size, rank=rank)
+
+ # Megatron core distributed training initialization
+ parallel_state.initialize_model_parallel(tensor_model_parallel_size = tensor_model_parallel_size, pipeline_model_parallel_size=pipeline_model_parallel_size)
+
+def model_provider():
+ """Build the model."""
+
+ transformer_config = TransformerConfig(
+ num_layers=2,
+ hidden_size=64,
+ num_attention_heads=2,
+ use_cpu_initialization=True,
+ pipeline_dtype=torch.float32
+ )
+
+ gpt_model = GPTModel(
+ config=transformer_config,
+ transformer_layer_spec=get_gpt_layer_local_spec(),
+ vocab_size=_VOCAB_SIZE,
+ max_sequence_length=_SEQUENCE_LENGTH,
+ )
+
+ return gpt_model
+
+def load_distributed_checkpoint(checkpoint_path, gpt_model):
+ sharded_state_dict=gpt_model.sharded_state_dict(prefix='')
+ checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)
+ gpt_model.load_state_dict(checkpoint)
+ return gpt_model
+
+if __name__ == "__main__":
+ initialize_distributed(tensor_model_parallel_size=2, pipeline_model_parallel_size=1)
+ model_parallel_cuda_manual_seed(123)
+
+ gpt_model = model_provider()
+ device = torch.device("cuda")
+ gpt_model.to(device)
+
+ # Optionally you can also load a gpt model from ckpt_path using this code below
+ # gpt_model = load_distributed_checkpoint(gpt_model=gpt_model, checkpoint_path=ckpt_path)
+
+ seq_len_interpolation_factor = None
+ if hasattr(gpt_model, "rotary_pos_emb"):
+ seq_len_interpolation_factor = gpt_model.rotary_pos_emb.seq_len_interpolation_factor
+
+ trtllm_helper = TRTLLMHelper(
+ transformer_config=gpt_model.config,
+ model_type=ModelType.gpt,
+ position_embedding_type = gpt_model.position_embedding_type,
+ max_position_embeddings = gpt_model.max_position_embeddings,
+ rotary_percentage = gpt_model.rotary_percent,
+ rotary_base = gpt_model.rotary_base,
+ moe_tp_mode = 2,
+ multi_query_mode = False,
+ activation = "gelu",
+ seq_len_interpolation_factor = seq_len_interpolation_factor,
+ share_embeddings_and_output_weights=gpt_model.share_embeddings_and_output_weights
+ )
+
+
+ trtllm_model_weights, trtllm_model_config = trtllm_helper.get_trtllm_pretrained_config_and_model_weights(
+ model_state_dict= gpt_model.state_dict(),
+ dtype = DataType.bfloat16,
+ on_device_distributed_conversion=True,
+ vocab_size=_VOCAB_SIZE,
+ gpus_per_node=2,
+ )
+
+ trtllm_helper.build_and_save_engine(
+ max_input_len=256,
+ max_output_len=256,
+ max_batch_size=8,
+ engine_dir='/opt/megatron-lm/engine',
+ trtllm_model_weights=trtllm_model_weights[0],
+ trtllm_model_config=trtllm_model_config[0],
+ lora_ckpt_list=None,
+ use_lora_plugin=None,
+ max_lora_rank=64,
+ lora_target_modules=None,
+ max_prompt_embedding_table_size=0,
+ paged_kv_cache=True,
+ remove_input_padding=True,
+ paged_context_fmha=False,
+ use_refit=False,
+ max_num_tokens=None,
+ max_seq_len=512,
+ opt_num_tokens=None,
+ max_beam_width=1,
+ tokens_per_block=128,
+ multiple_profiles=False,
+ gpt_attention_plugin="auto",
+ gemm_plugin="auto",
+ )
diff --git a/examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py b/examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py
new file mode 100644
index 0000000000..587e7cfdd3
--- /dev/null
+++ b/examples/export/trtllm_export/single_device_export/gpt_single_device_cpu_export.py
@@ -0,0 +1,118 @@
+import os
+import torch
+from megatron.core import parallel_state
+from megatron.core import dist_checkpointing
+from megatron.core.export.model_type import ModelType
+from megatron.core.export.data_type import DataType
+from megatron.core.export.export_config import ExportConfig
+from megatron.core.export.trtllm.trtllm_helper import TRTLLMHelper
+from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.models.gpt.gpt_model import GPTModel
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
+
+
+_SEQUENCE_LENGTH = 64
+
+
+def initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1):
+ parallel_state.destroy_model_parallel()
+
+ # Torch setup for distributed training
+ rank = int(os.environ['LOCAL_RANK'])
+ world_size = torch.cuda.device_count()
+ torch.cuda.set_device(rank)
+ torch.distributed.init_process_group(world_size=world_size, rank=rank)
+
+ # Megatron core distributed training initialization
+ parallel_state.initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size)
+
+def model_provider():
+ """Build the model."""
+
+ transformer_config = TransformerConfig(
+ num_layers=2,
+ hidden_size=64, # Needs to be atleast 32 times num_attn_heads
+ num_attention_heads=2,
+ use_cpu_initialization=True,
+ pipeline_dtype=torch.float32,
+ )
+
+ gpt_model = GPTModel(
+ config=transformer_config,
+ transformer_layer_spec=get_gpt_layer_local_spec(),
+ vocab_size=100,
+ max_sequence_length=_SEQUENCE_LENGTH,
+ )
+
+ return gpt_model
+
+def load_distributed_checkpoint(checkpoint_path, gpt_model):
+ sharded_state_dict=gpt_model.sharded_state_dict(prefix='')
+ checkpoint = dist_checkpointing.load(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path)
+ gpt_model.load_state_dict(checkpoint)
+ return gpt_model
+
+if __name__ == "__main__":
+ # Need to use TP1 PP1 for export on single device
+ initialize_distributed(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
+ model_parallel_cuda_manual_seed(123)
+
+ gpt_model = model_provider()
+
+ # Optionally you can also load a gpt model from ckpt_path using this code below
+ # gpt_model = load_distributed_checkpoint(gpt_model=gpt_model, checkpoint_path=ckpt_path)
+
+ seq_len_interpolation_factor = None
+ if hasattr(gpt_model, "rotary_pos_emb"):
+ seq_len_interpolation_factor = gpt_model.rotary_pos_emb.seq_len_interpolation_factor
+
+ trtllm_helper = TRTLLMHelper(
+ transformer_config=gpt_model.config,
+ model_type=ModelType.gpt,
+ position_embedding_type = gpt_model.position_embedding_type,
+ max_position_embeddings = gpt_model.max_position_embeddings,
+ rotary_percentage = gpt_model.rotary_percent,
+ rotary_base = gpt_model.rotary_base,
+ moe_tp_mode = 2,
+ multi_query_mode = False,
+ activation = "gelu",
+ seq_len_interpolation_factor = seq_len_interpolation_factor,
+ share_embeddings_and_output_weights=gpt_model.share_embeddings_and_output_weights
+ )
+
+
+ export_config = ExportConfig(inference_tp_size = 2)
+ # NOTE : For faster performance, if your entire model will fit in gpu memory, transfer model state dict to GPU and then call this api
+ weight_list, config_list = trtllm_helper.get_trtllm_pretrained_config_and_model_weights(
+ model_state_dict= gpt_model.state_dict(),
+ dtype = DataType.bfloat16,
+ export_config=export_config
+ )
+
+ for trtllm_model_weights, trtllm_model_config in zip(weight_list, config_list):
+ trtllm_helper.build_and_save_engine(
+ max_input_len=256,
+ max_output_len=256,
+ max_batch_size=8,
+ engine_dir='/opt/megatron-lm/engine',
+ trtllm_model_weights=trtllm_model_weights,
+ trtllm_model_config=trtllm_model_config,
+ lora_ckpt_list=None,
+ use_lora_plugin=None,
+ max_lora_rank=64,
+ lora_target_modules=None,
+ max_prompt_embedding_table_size=0,
+ paged_kv_cache=True,
+ remove_input_padding=True,
+ paged_context_fmha=False,
+ use_refit=False,
+ max_num_tokens=None,
+ max_seq_len=512,
+ opt_num_tokens=None,
+ max_beam_width=1,
+ tokens_per_block=128,
+ multiple_profiles=False,
+ gpt_attention_plugin="auto",
+ gemm_plugin="auto",
+ )
\ No newline at end of file
diff --git a/examples/finetune_mnli_distributed.sh b/examples/finetune_mnli_distributed.sh
deleted file mode 100755
index a3f9accbcc..0000000000
--- a/examples/finetune_mnli_distributed.sh
+++ /dev/null
@@ -1,43 +0,0 @@
-#!/bin/bash
-
-WORLD_SIZE=8
-
-DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
- --nnodes 1 \
- --node_rank 0 \
- --master_addr localhost \
- --master_port 6000"
-
-TRAIN_DATA="data/glue_data/MNLI/train.tsv"
-VALID_DATA="data/glue_data/MNLI/dev_matched.tsv \
- data/glue_data/MNLI/dev_mismatched.tsv"
-PRETRAINED_CHECKPOINT=checkpoints/bert_345m
-VOCAB_FILE=bert-vocab.txt
-CHECKPOINT_PATH=checkpoints/bert_345m_mnli
-
-python -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/main.py \
- --task MNLI \
- --seed 1234 \
- --train-data $TRAIN_DATA \
- --valid-data $VALID_DATA \
- --tokenizer-type BertWordPieceLowerCase \
- --vocab-file $VOCAB_FILE \
- --epochs 5 \
- --pretrained-checkpoint $PRETRAINED_CHECKPOINT \
- --tensor-model-parallel-size 1 \
- --num-layers 24 \
- --hidden-size 1024 \
- --num-attention-heads 16 \
- --micro-batch-size 8 \
- --lr 5.0e-5 \
- --lr-decay-style linear \
- --lr-warmup-fraction 0.065 \
- --seq-length 512 \
- --max-position-embeddings 512 \
- --save-interval 500000 \
- --save $CHECKPOINT_PATH \
- --log-interval 10 \
- --eval-interval 100 \
- --eval-iters 50 \
- --weight-decay 1.0e-1 \
- --fp16
diff --git a/examples/finetune_race_distributed.sh b/examples/finetune_race_distributed.sh
deleted file mode 100755
index 3d92253388..0000000000
--- a/examples/finetune_race_distributed.sh
+++ /dev/null
@@ -1,46 +0,0 @@
-#!/bin/bash
-
-WORLD_SIZE=8
-
-DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
- --nnodes 1 \
- --node_rank 0 \
- --master_addr localhost \
- --master_port 6000"
-
-TRAIN_DATA="data/RACE/train/middle"
-VALID_DATA="data/RACE/dev/middle \
- data/RACE/dev/high"
-VOCAB_FILE=bert-vocab.txt
-PRETRAINED_CHECKPOINT=checkpoints/bert_345m
-CHECKPOINT_PATH=checkpoints/bert_345m_race
-
-python -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/main.py \
- --task RACE \
- --seed 1234 \
- --train-data $TRAIN_DATA \
- --valid-data $VALID_DATA \
- --tokenizer-type BertWordPieceLowerCase \
- --vocab-file $VOCAB_FILE \
- --epochs 3 \
- --pretrained-checkpoint $PRETRAINED_CHECKPOINT \
- --tensor-model-parallel-size 1 \
- --num-layers 24 \
- --hidden-size 1024 \
- --num-attention-heads 16 \
- --micro-batch-size 4 \
- --lr 1.0e-5 \
- --lr-decay-style linear \
- --lr-warmup-fraction 0.06 \
- --seq-length 512 \
- --max-position-embeddings 512 \
- --save-interval 100000 \
- --save $CHECKPOINT_PATH \
- --log-interval 10 \
- --eval-interval 100 \
- --eval-iters 50 \
- --weight-decay 1.0e-1 \
- --clip-grad 1.0 \
- --hidden-dropout 0.1 \
- --attention-dropout 0.1 \
- --fp16
diff --git a/examples/finetune_retriever_distributed.sh b/examples/finetune_retriever_distributed.sh
deleted file mode 100755
index 535a2e053d..0000000000
--- a/examples/finetune_retriever_distributed.sh
+++ /dev/null
@@ -1,56 +0,0 @@
-#!/bin/bash
-
-# Finetune a BERT or pretrained ICT model using Google natural question data
-# Datasets can be downloaded from the following link:
-# https://github.com/facebookresearch/DPR/blob/master/data/download_data.py
-
-WORLD_SIZE=8
-
-DISTRIBUTED_ARGS="--nproc_per_node $WORLD_SIZE \
- --nnodes 1 \
- --node_rank 0 \
- --master_addr localhost \
- --master_port 6000"
-
-CHECKPOINT_PATH=
-
-# Load either of the below
-BERT_LOAD_PATH=
-PRETRAINED_CHECKPOINT=
-
-python -m torch.distributed.launch $DISTRIBUTED_ARGS ./tasks/main.py \
- --task RET-FINETUNE-NQ \
- --train-with-neg \
- --train-hard-neg 1 \
- --pretrained-checkpoint ${PRETRAINED_CHECKPOINT} \
- --num-layers 12 \
- --hidden-size 768 \
- --num-attention-heads 12 \
- --tensor-model-parallel-size 1 \
- --tokenizer-type BertWordPieceLowerCase \
- --train-data nq-train.json \
- --valid-data nq-dev.json \
- --save ${CHECKPOINT_PATH} \
- --load ${CHECKPOINT_PATH} \
- --vocab-file bert-vocab.txt \
- --bert-load ${BERT_LOAD_PATH} \
- --save-interval 5000 \
- --log-interval 10 \
- --eval-interval 20000 \
- --eval-iters 100 \
- --indexer-log-interval 1000 \
- --faiss-use-gpu \
- --DDP-impl torch \
- --fp16 \
- --retriever-report-topk-accuracies 1 5 10 20 100 \
- --seq-length 512 \
- --retriever-seq-length 256 \
- --max-position-embeddings 512 \
- --retriever-score-scaling \
- --epochs 80 \
- --micro-batch-size 8 \
- --eval-micro-batch-size 16 \
- --indexer-batch-size 128 \
- --lr 2e-5 \
- --lr-warmup-fraction 0.01 \
- --weight-decay 1e-1
diff --git a/examples/gpt3/README.md b/examples/gpt3/README.md
index fec51e1fea..8d6f267416 100644
--- a/examples/gpt3/README.md
+++ b/examples/gpt3/README.md
@@ -10,7 +10,7 @@
To run the model using a docker container run it as follows
```
-PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:23.09-py3
+PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:24.01-py3
CHECKPOINT_PATH="" #
TENSORBOARD_LOGS_PATH=""#
VOCAB_FILE="" #/gpt2-vocab.json
@@ -23,8 +23,8 @@ docker run \
--workdir /workspace/megatron-lm \
-v /path/to/data:/path/to/data \
-v /path/to/megatron-lm:/workspace/megatron-lm \
- megatron-lm nvcr.io/nvidia/pytorch:23.04-py3 \
- bash /examples/gpt3/train_gpt3_175b_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH $VOCAB_FILE $MERGE_FILE $DATA_PATH "
+ megatron-lm nvcr.io/nvidia/pytorch:24.01-py3 \
+ bash examples/gpt3/train_gpt3_175b_distributed.sh $CHECKPOINT_PATH $TENSORBOARD_LOGS_PATH $VOCAB_FILE $MERGE_FILE $DATA_PATH "
```
NOTE: Depending on the environment you are running it the above command might like slightly different.
@@ -34,7 +34,7 @@ NOTE: Depending on the environment you are running it the above command might li
The example in this folder shows you how to run 175B model. There are other configs you could run as well
-### 345M
+### 345M
```
--num-layers 12 \
--hidden-size 512 \
@@ -45,7 +45,7 @@ The example in this folder shows you how to run 175B model. There are other conf
```
-### 857M
+### 857M
```
--num-layers 24 \
--hidden-size 1024 \
diff --git a/examples/gpt3/gpt_config.yaml b/examples/gpt3/gpt_config.yaml
new file mode 100644
index 0000000000..8ef4f41e96
--- /dev/null
+++ b/examples/gpt3/gpt_config.yaml
@@ -0,0 +1,301 @@
+# WARNING: Yaml configs is currently an experimental feature
+language_model:
+ # model architecture
+ num_layers: 24
+ hidden_size: 1024
+ num_attention_heads: 16
+ num_query_groups: null
+
+ ffn_hidden_size: null
+ kv_channels: null
+ hidden_dropout: 0.0
+ attention_dropout: 0.0
+ fp32_residual_connection: False
+
+ apply_residual_connection_post_layernorm: False
+ layernorm_epsilon: 1.e-5
+ layernorm_zero_centered_gamma: True
+ add_bias_linear: False
+ bias_activation_fusion: False
+ add_qkv_bias: False
+ gated_linear_unit: False
+ activation_func: swiglu
+ num_moe_experts: null
+ rotary_interleaved: False
+ window_size: null
+
+ # initialization
+ init_method: null
+ init_method_std: 0.02
+ output_layer_init_method: null
+
+ # mixed-precision
+ apply_query_key_layer_scaling: False
+ attention_softmax_in_fp32: False
+
+ # fusion
+ bias_swiglu_fusion: True
+ masked_softmax_fusion: True
+ persist_layer_norm: False
+ memory_efficient_layer_norm: False
+ bias_dropout_fusion: True
+ apply_rope_fusion: True
+
+ # activation recomputation
+ recompute_granularity: null
+ recompute_method: null
+ recompute_num_layers: null
+ distribute_saved_activations: null
+
+ # fp8 related
+ fp8: null
+ fp8_margin: 0
+ fp8_interval: 1
+ fp8_amax_history_len: 1
+ fp8_amax_compute_algo: "most_recent"
+ fp8_wgrad: True
+
+ # miscellaneous
+ clone_scatter_output_in_embedding: True
+
+ normalization: "LayerNorm" # alt value supported by TE: "RMSNorm"
+
+ # MoE related
+ moe_router_load_balancing_type: "aux_loss"
+ moe_router_topk: 2
+ moe_router_group_topk: null
+ moe_router_num_groups: null
+ moe_grouped_gemm: False
+ moe_aux_loss_coeff: 0 # 1e-2 would be a good start value for load balance loss.
+ moe_z_loss_coeff: null # 1e-3 would be a good start value for z-loss
+ moe_input_jitter_eps: null
+ moe_token_dropping: False
+
+model_parallel:
+ # Model parallelism
+ tensor_model_parallel_size: 1
+ context_parallel_size: 1
+ pipeline_model_parallel_size: 1
+ virtual_pipeline_model_parallel_size: null
+ sequence_parallel: True
+ expert_model_parallel_size: 1
+
+ # Initialization
+ perform_initialization: True
+ use_cpu_initialization: null
+
+ # Training
+ fp16: False
+ bf16: True
+ params_dtype: null # Set from above arguments for core
+ timers: null
+
+ # Optimizations
+ gradient_accumulation_fusion: True
+ async_tensor_model_parallel_allreduce: True
+ tp_comm_overlap: False
+
+ # Debug Options
+ tp_comm_split_ag: True
+ tp_comm_atomic_ag: True
+ tp_comm_split_rs: True
+ tp_comm_atomic_rs: True
+ tp_comm_bulk_wgrad: True
+ tp_comm_bulk_dgrad: True
+
+ # Parallelism
+ finalize_model_grads_func: null
+
+ # Pipeline Parallel
+ pipeline_dtype: null
+ grad_scale_func: null
+ enable_autocast: False
+ autocast_dtype: null
+ variable_seq_lengths: False
+ num_microbatches_with_partial_activation_checkpoints: null
+ overlap_p2p_comm: False
+ batch_p2p_comm: True
+ batch_p2p_sync: True
+ use_ring_exchange_p2p: False
+ deallocate_pipeline_outputs: False
+ no_sync_func: null
+ grad_sync_func: null
+ param_sync_func: null
+ pipeline_model_parallel_split_rank: null
+
+ # CPU Offloading
+ cpu_offloading: False
+ cpu_offloading_num_layers: 0
+ _cpu_offloading_context: null
+ cpu_offloading_weights: False
+ cpu_offloading_activations: True
+
+ # Timing
+ barrier_with_L1_time: True
+
+# training:
+use_legacy_models: False
+spec: null
+micro_batch_size: 2
+global_batch_size: 128
+rampup_batch_size: [32, 32, 65324160]
+check_for_nan_in_loss_and_grad: True
+num_layers_per_virtual_pipeline_stage: null
+
+encoder_num_layers: null
+decoder_num_layers: null
+rotary_seq_len_interpolation_factor: null
+add_position_embedding: False
+make_vocab_size_divisible_by: 128
+group_query_attention: False
+
+
+exit_signal_handler: False
+exit_duration_in_mins: null
+exit_interval: null
+
+untie_embeddings_and_output_weights: True
+position_embedding_type: rope
+rotary_percent: 0.5
+openai_gelu: False
+squared_relu: False
+swiglu: True
+onnx_safe: null
+bert_binary_head: True
+max_position_embeddings: 4096
+
+transformer_impl: local
+use_flash_attn: False
+seed: 1234
+data_parallel_random_init: False
+
+# Optimizer
+optimizer: adam
+lr: 2.5e-4
+lr_decay_style: cosine
+lr_decay_iters: null
+lr_decay_samples: 255126953
+lr_warmup_fraction: null
+lr_warmup_iters: 0
+lr_warmup_samples: 81381
+lr_warmup_init: 0.0
+min_lr: 2.5e-5
+weight_decay: 0.1
+start_weight_decay: null
+end_weight_decay: null
+weight_decay_incr_style: constant
+clip_grad: 1.0
+adam_beta1: 0.9
+adam_beta2: 0.95
+adam_eps: 1.e-08
+sgd_momentum: 0.9
+override_opt_param_scheduler: False
+use_checkpoint_opt_param_scheduler: False
+
+# checkpointing arguments
+save: null
+save_interval: 20000
+no_save_optim: null
+no_save_rng: null
+load: null
+no_load_optim: null
+no_load_rng: null
+finetune: False
+use_checkpoint_args: False
+exit_on_missing_checkpoint: False
+
+# loss arguments
+loss_scale: null
+initial_loss_scale: 4294967296
+min_loss_scale: 1.0
+loss_scale_window: 1000
+hysteresis: 2
+accumulate_allreduce_grads_in_fp32: False
+fp16_lm_cross_entropy: False
+
+# distributed arguments
+distributed_backend: nccl
+distributed_timeout_minutes: 10
+overlap_grad_reduce: False
+align_grad_reduce: True
+overlap_param_gather: False
+align_param_gather: False
+scatter_gather_tensors_in_pipeline: True
+local_rank: null
+lazy_mpu_init: null
+empty_unused_memory_level: 0
+standalone_embedding_stage: False
+use_distributed_optimizer: False
+nccl_communicator_config_path: null
+
+train_iters: null
+eval_iters: 32
+eval_interval: 2000
+skip_train: False
+
+adlr_autoresume: False
+adlr_autoresume_interval: 1000
+
+# garbage collection
+manual_gc: False
+manual_gc_interval: 0
+manual_gc_eval: True
+
+tp_comm_overlap_cfg: null
+
+#data
+data_path: null
+split: '99,1,0'
+train_data_path: null
+valid_data_path: null
+test_data_path: null
+data_cache_path: null
+mock_data: False
+vocab_size: null
+vocab_file: null
+merge_file: null
+vocab_extra_ids: 0
+seq_length: 4096
+encoder_seq_length: null
+decoder_seq_length: null
+retriever_seq_length: 256
+sample_rate: 1.0
+mask_prob: 0.15
+short_seq_prob: 0.1
+num_workers: 2
+tokenizer_type: GPTSentencePieceTokenizer
+tokenizer_model: null
+reset_position_ids: False
+reset_attention_mask: False
+eod_mask_loss: False
+train_samples: 268554688
+dataloader_type: null
+
+#profile:
+profile: False
+profile_ranks: [0]
+profile_step_end: 12
+profile_step_start: 10
+
+#logging:
+log_params_norm: True
+log_num_zeros_in_grad: True
+log_throughput: False
+log_progress: False
+timing_log_level: 0
+timing_log_option: minmax
+tensorboard_log_interval: 1
+tensorboard_queue_size: 1000
+log_timers_to_tensorboard: False
+log_validation_ppl_to_tensorboard: False
+log_memory_to_tensorboard: False
+log_world_size_to_tensorboard: False
+log_loss_scale_to_tensorboard: True
+wandb_project: ''
+wandb_exp_name: ''
+wandb_save_dir: ''
+enable_one_logger: True
+one_logger_project: megatron-lm
+one_logger_run_name: null
+log_interval: 100
+tensorboard_dir: null
diff --git a/examples/gpt3/train_gpt3_175b_distributed.sh b/examples/gpt3/train_gpt3_175b_distributed.sh
index 01ca2e0309..7d2c01b315 100755
--- a/examples/gpt3/train_gpt3_175b_distributed.sh
+++ b/examples/gpt3/train_gpt3_175b_distributed.sh
@@ -12,11 +12,11 @@ NUM_NODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
-CHECKPOINT_PATH=$0 #
-TENSORBOARD_LOGS_PATH=$1 #
-VOCAB_FILE=$2 #/gpt2-vocab.json
-MERGE_FILE=$3 #/gpt2-merges.txt
-DATA_PATH=$4 #_text_document
+CHECKPOINT_PATH=$1 #
+TENSORBOARD_LOGS_PATH=$2 #
+VOCAB_FILE=$3 #/gpt2-vocab.json
+MERGE_FILE=$4 #/gpt2-merges.txt
+DATA_PATH=$5 #_text_document
DISTRIBUTED_ARGS=(
--nproc_per_node $GPUS_PER_NODE
@@ -31,6 +31,7 @@ GPT_MODEL_ARGS=(
--num-attention-heads 96
--seq-length 2048
--max-position-embeddings 2048
+ --attention-backend auto # Can use (flash/fused/unfused/local)
)
TRAINING_ARGS=(
@@ -49,7 +50,6 @@ TRAINING_ARGS=(
--min-lr 6.0e-6
--lr-warmup-fraction .001
--lr-decay-iters 430000
- --use-mcore-models
)
MODEL_PARALLEL_ARGS=(
diff --git a/examples/inference/README.md b/examples/inference/README.md
new file mode 100644
index 0000000000..b4b07cbc6a
--- /dev/null
+++ b/examples/inference/README.md
@@ -0,0 +1,279 @@
+### Megatron Core Inference Documentation
+This guide provides an example for Megatron Core for running model inference.
+
+### Contents
+- [Megatron Core Inference Documentation](#megatron-core-inference-documentation)
+- [Contents](#contents)
+ - [1. Quick Start](#1-quick-start)
+ - [1.1 Understanding The Code](#11-understanding-the-code)
+ - [1.2 Running The Code](#12-running-the-code)
+ - [2. Flow of Control In MCore Backend](#2-flow-of-control-in-mcore-backend)
+ - [3. Customizing The Inference Pipeline](#3-customizing-the-inference-pipeline)
+ - [3.1. Create Your Own Inference Backend](#31-create-your-own-inference-backend)
+ - [3.2. Create Your Own Text Generation Controller](#32-create-your-own-text-generation-controller)
+ - [3.3. Support Other Models](#33-support-other-models)
+ - [3.3. Modify Inference Parameters](#33-modify-inference-parameters)
+ - [4. Future work](#4-future-work)
+
+
+
+#### 1. Quick Start
+This example runs batch inference on a GPT model trained using Megatron Core. The entrypoint is [simple_gpt_batch_inference.py](./gpt/gpt_batch_inference.py)
+
+
+
+##### 1.1 Code Walkthrough
+***STEP 1 - Initialize model parallel and other default arguments***
+The micro batch size is set as 1 as it is not used in tensor-parallelism only, and for pipeline-parallel models it is calculated at runtime.
+```python
+ initialize_megatron(
+ args_defaults={'no_load_rng': True, 'no_load_optim': True, 'micro_batch_size': 1}
+ )
+```
+
+***STEP 2 - Load the model using the model_provider_function***
+NOTE: The model provider function supports both MCore and Legacy models.
+
+```python
+ model = get_model(model_provider, wrap_with_ddp=False)
+ load_checkpoint(model, None, None)
+ model = model[0]
+```
+
+***STEP 3 - Choose an engine***
+Text generation requires an inference engine, which includes a scheduler. The default engine is the [Megatron Core engine](../../megatron/core/inference/engine/mcore_engine.py) with a simple [text generation controller](../../megatron/core/inference/text_generation_controllers/text_generation_controller.py). TRTLLMEngine will be supported in the future.
+```python
+ inference_wrapped_model = GPTInferenceWrapper(model, args)
+ text_generation_controller = TextGenerationController(
+ inference_wrapped_model=inference_wrapped_model,
+ tokenizer=tokenizer
+ )
+ inference_backend = MCoreEngine(
+ text_generation_controller=text_generation_controller, max_batch_size=args.max_batch_size
+ )
+```
+
+***STEP 4 - Run text generation***
+The [SamplingParams](../../megatron/core/inference/sampling_params.py) contains suggested defaults. Customize this to change top_p, top_k, number of tokens to generate etc.
+*Note: The result is returned as a list of [InferenceRequests](../../megatron/core/inference/inference_request.py)*
+```python
+ results: List[InferenceRequest] = inference_engine.generate(
+ prompts=args.prompts, sampling_params=sampling_params
+ )
+
+ if torch.distributed.get_rank() == 0:
+ for idx, result in enumerate(results):
+ print(f' ------------- RESULT FOR PROMPT {idx} --------------- ')
+ result = {
+ 'id': result.request_id,
+ 'input_prompt': result.prompt,
+ 'generated_text': result.generated_text,
+ 'generated_tokens' : result.generated_tokens
+ }
+ print(result)
+```
+
+
+
+##### 1.2 Running The Code
+An example run script is shown below. Set the tokenizer paths, inference params, and other settings appropriately.
+
+For a quick recap on sampling parameters, refer to [this blog](https://ivibudh.medium.com/a-guide-to-controlling-llm-model-output-exploring-top-k-top-p-and-temperature-parameters-ed6a31313910).
+
+```
+# In a slurm cluster (You could also use docker)
+ACCOUNT=
+MLM_PATH=/path/to/megatron-lm
+GPT_CKPT=/path/to/gpt/ckpt
+VOCAB_MERGE_FILE_PATH=/path/to/vocab/and/merge/file
+CONTAINER_IMAGE=nvcr.io/ea-bignlp/ga-participants/nemofw-training:23.11
+
+srun --account $ACCOUNT \
+--job-name=$ACCOUNT:inference \
+--partition=batch \
+--time=01:00:00 \
+--container-image $CONTAINER_IMAGE \
+--container-mounts $MLM_PATH:/workspace/megatron-lm/,$GPT_CKPT:/workspace/mcore_gpt_ckpt,$VOCAB_MERGE_FILE_PATH:/workspace/tokenizer \
+--no-container-mount-home \
+--pty /bin/bash \
+
+# Inside the container run the following.
+
+cd megatron-lm/
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+TOKENIZER_ARGS=(
+ --vocab-file /workspace/tokenizer/gpt2-vocab.json
+ --merge-file /workspace/tokenizer/gpt2-merges.txt
+ --tokenizer-type GPT2BPETokenizer
+)
+
+MODEL_ARGS=(
+ --use-checkpoint-args
+ --use-mcore-models
+ --load /workspace/mcore_gpt_ckpt
+)
+
+INFERENCE_SPECIFIC_ARGS=(
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --num-tokens-to-generate 20
+ --max-batch-size 4
+)
+
+torchrun --nproc-per-node=4 examples/inference/gpt/simple_gpt_batch_inference.py \
+ ${TOKENIZER_ARGS[@]} \
+ ${MODEL_ARGS[@]} \
+ ${INFERENCE_SPECIFIC_ARGS[@]} \
+ --prompts "prompt one " "sample prompt two" "sample prompt 3"
+
+NOTE: Other parameters which can be customized for inference are :-
+--temperature (Sampling temperature)
+--top_k (top_k sampling)
+--top_p (top_p sampling)
+--num-tokens-to-generate (Number of tokens to generate for each prompt)
+--inference-batch-times-seqlen-threshold (During inference, if batch-size times sequence-length is smaller than this threshold then we will not use pipelining, otherwise we will.')
+--use-dist-ckpt (If using dist checkpoint format for the model)
+--use-legacy-models (If using legacy gpt model instead of mcore gpt model)
+
+```
+
+
+
+
+
+#### 2. Control Flow in the MCore Backend
+An example of inference with static batching is provided in [gpt_batch_inference.py](./gpt/gpt_batch_inference.py).
+* [mcore_engine](../../megatron/core/inference/engines/mcore_engine.py) **generate()** function is called with the input prompts.
+* The `Scheduler` in the engine will add these prompts to the [active requests] pool (../../megatron/core/inference/inference_request.py) until max batch size is hit. Remaining requests will be added to the waiting requests pool.
+* The engine will run until all requests (waiting + active) are completed.
+ * The active requests are passed into **generate_all_output_tokens_static_batch()** of the text generation controller .
+ * This function uses the **prep_model_for_inference()** method of the [model_inference_wrappers](../../megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py) and runs an autoregressive sampling loop
+ * In the autoregressive loop, the **get_batch_for_context_window()** method of the inference wrapper is called to slice out the input tokens and masks
+ * Input tokens and masks are passed it into the **run_one_forward_step()** method, which calls the model `.forward()` method to get the output logits
+ * Output logits are synchronized across all pipeline parallel ranks
+ * The text generation controller obtains the log probabilities and samples tokens based on the strategy defined in the sampling parameters.
+ * The sampled tokens are then appended to the input prompt tokens for the next iteration
+ * The **update_generation_status()** method of the text generation controller checks which prompts have finished generating or hit a stop condition
+ * After the inference loop, the result is detokenized and stored as an attribute of the InferenceRequest. These requests are marked as completed.
+ * The **update_requests_pool()** method of the scheduler moves completed requests into the completed request pool and waiting requests into the active request pool
+
+
+
+#### 3. Customizing The Inference Pipeline
+
+The inference pipeline supports three levels of customization:
+
+* **Inference engine** - The MCore Engine is currently supported. Change this to add a new backend.
+* **Text generation controller** - The main sampling loop. This can be customized to support alternative tokenization, detokenization, or to implement a new sampling strategy.
+* **Inference Wrapped Model** - Change this to support a new model.
+* **Modify Inference Parameters** - Change this to update top_p, top_k, number of tokens to be generated, temperature, or other sampling parameters.
+
+
+
+##### 3.1. Create Your Own Inference Backend
+The [abstract_engine.py](./../../megatron/core/inference/engine/abstract_engine.py) file contains a `generate` method that can be extended to support a new backend.
+
+```python
+class AbstractEngine(ABC):
+ @staticmethod
+ def generate(self) -> dict:
+ """The abstract backend's generate function.
+
+ To define a new backend, implement this method and return the outputs as a dictionary.
+```
+
+
+
+##### 3.2. Implement a new Sampling Loop
+
+The [TextGenerationController](../../megatron/core/inference/text_generation_controllers/text_generation_controller.py) contains the main sampling loop and can be modified to support new tokenization, detokenization, or sampling strategies.
+
+``` python
+class TextGenerationController:
+
+ def tokenize_prompt(self, prompt: str) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Utility to tokenize the input prompts"""
+
+ def sample_from_logits(
+ self,
+ last_token_logits: torch.Tensor,
+ sampling_params: SamplingParams,
+ vocab_size: int,
+ ) -> torch.Tensor:
+ """Samples the logits to generate outputs
+
+ Given the logits of the last token, this function samples according to the parameters defined in sampling_params and returns the sampled tokens.
+ """
+
+ def update_generation_status(
+ self,
+ updated_prompts_tokens: torch.Tensor,
+ generation_started: torch.Tensor,
+ current_context_end_position: int,
+ is_generation_done_tensor: torch.Tensor,
+ generated_sequence_lengths: torch.Tensor,
+ ) -> torch.Tensor:
+ """Function to check which prompts have reached an end condition
+
+ We check which prompts have reached an end condition and set the corresponding flags of the is_generation_done_tensor to True . The generated sequence lengths increases as we keep generating, until that prompts hits an eod condition. The generation started status tensor helps us determine which prompts have started generating
+ """
+
+ def generate_all_output_tokens_static_batch(
+ self, active_requests: OrderedDict[int, InferenceRequest],
+ ) -> OrderedDict[int, InferenceRequest]:
+ """Utility to generate all the output tokens and probabilities for the prompts .
+
+ This utility generates the output tokens for a static batch. It runs the forward steps till all prompts complete generation, updates the status of these requests to completed, adds the generated result and returns these requests
+ """
+
+ def detokenize_generations(self, prompt_tokens_with_generated_tokens: torch.Tensor) -> str:
+ """Detokenize the output generations"""
+```
+
+
+
+##### 3.3. Support Other Models
+Extend [abstract_model_inference_wrapper.py](./../../megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py) to support other models. The abstract model wrapper implements:
+* Forward method which calls the model `forward` method depending on model parallel settings
+* Initializes the model and puts it in `.eval()` mode
+* Setup for the input parameters (max batch size, max seq length)
+
+The following methods should be implemented:
+```python
+class AbstractModelInferenceWrapper:
+ def prep_model_for_inference(self, prompts_tokens: torch.Tensor):
+ """A utility function for preparing model for inference
+
+ The function gets called once before the auto regressive inference loop. It puts the model in eval mode , and gets some model and inference data parameters. Extend this to build position ids ,attention mask etc, so that required slices can be extracted during the forward pass
+ """
+
+ @abc.abstractclassmethod
+ def get_batch_for_context_window(self) -> List:
+ """Returns the input data for inference
+
+ This function gets called iteratively in the inference loop. It can be used to extract relevant input from the prompt tokens, attention mask etc. required for each step in inference.
+```
+
+Refer to [gpt_inference_wrapper.py](../../megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py) for an example of implementing this for GPTModel.
+
+
+
+##### 3.3. Modify Inference Parameters
+We use [common inference params](../../megatron/core/inference/sampling_params.py) for text generation. Customize this if you want to change top_p, top_k, number of tokens to generate etc. If you want to add other attributes that you would use in the inference loop, you can do that as shown below
+
+```
+from megatron.core.inference.sampling_params import SamplingParams
+
+c = SamplingParams(temperature=0.5)
+c.add_attributes({'min_length':4, 'eod_id':153})
+```
+
+
+
+#### 4. Future work
+The following features are planned for the future releases.
+* Dynamic batching
+* Paged Attention
+* TRTLLM Engine support
+* Support for multimodal inference
\ No newline at end of file
diff --git a/examples/inference/gpt/gpt_batch_inference.py b/examples/inference/gpt/gpt_batch_inference.py
new file mode 100644
index 0000000000..9c2bada4b6
--- /dev/null
+++ b/examples/inference/gpt/gpt_batch_inference.py
@@ -0,0 +1,200 @@
+import os
+from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
+ InferenceWrapperConfig,
+)
+from pretrain_gpt import model_provider
+import torch
+import sys
+import time
+import tqdm
+import warnings
+from argparse import Namespace
+from megatron.core.inference.engines.abstract_engine import AbstractEngine
+from megatron.core.inference.engines.mcore_engine import MCoreEngine
+from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
+ GPTInferenceWrapper,
+)
+from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.inference.text_generation_controllers.text_generation_controller import (
+ TextGenerationController,
+)
+from megatron.core.transformer.module import MegatronModule
+
+sys.path.append(
+ os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
+)
+
+from megatron.training import get_args
+from megatron.training import get_tokenizer
+from megatron.training.checkpointing import load_checkpoint
+from megatron.core import mpu
+from megatron.training.initialize import initialize_megatron
+from megatron.training import get_model
+import asyncio
+from typing import AsyncIterator, List
+
+
+
+def add_text_generate_args(parser):
+ """Text generation arguments."""
+ group = parser.add_argument_group(title='text generation')
+
+ group.add_argument("--temperature", type=float, default=1.0, help='Sampling temperature.')
+ group.add_argument("--top_k", type=int, default=1, help='Top k sampling.')
+ group.add_argument("--top_p", type=float, default=0.0, help='Top p sampling.')
+ group.add_argument(
+ "--return-log-probs",
+ action='store_true',
+ default=False,
+ help='Return the log probabilities of the final output tokens',
+ )
+ group.add_argument(
+ "--num-tokens-to-generate",
+ type=int,
+ default=30,
+ help='Number of tokens to generate for each prompt',
+ )
+ group.add_argument(
+ "--prompts",
+ metavar='N',
+ type=str,
+ nargs='+',
+ help='Input prompts with each prompt within quotes and seperated by space',
+ )
+ group.add_argument(
+ "--max-batch-size", type=int, default=8, dest="inference_max_requests",
+ help='Max number of prompts to process at once'
+ )
+ group.add_argument("--stream", action="store_true", default=False, help="Stream output tokens")
+ return parser
+
+
+def get_inference_engine(args: Namespace, model: MegatronModule) -> AbstractEngine:
+ """Utility to get the relevant backend for running inference
+
+ This function will automatically chose the TRTLLMBackend when possible, and if not revert to Mcore backend if the user does not specify any backends. TRT LLM Backend is not implmented yet.
+
+ Args:
+ args (Namespace): The user arguments parsed from command line
+ model (MegatronModule): The megatron model .
+
+ Returns:
+ AbstractBackend: The chosen backend
+ """
+ tokenizer = get_tokenizer()
+
+ inference_wrapper_config = InferenceWrapperConfig(
+ hidden_size=args.hidden_size,
+ inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold,
+ fp32_residual_connection=args.fp32_residual_connection,
+ params_dtype=args.params_dtype,
+ padded_vocab_size=args.padded_vocab_size,
+ inference_max_requests=args.inference_max_requests,
+ inference_max_seq_length=args.inference_max_seq_length,
+ )
+
+ inference_wrapped_model = GPTInferenceWrapper(model, inference_wrapper_config)
+ text_generation_controller = TextGenerationController(inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer)
+ return MCoreEngine(text_generation_controller=text_generation_controller)
+
+
+async def generate(
+ inference_engine: MCoreEngine,
+ sampling_params: SamplingParams,
+ prompts: List[str],
+) -> List[InferenceRequest]:
+ async def collect_stream(prompt, request_id, stream_generator):
+ print(f"Request {request_id}: {prompt}", end="", flush=True)
+ prev_idx = 0
+ async for output in stream_generator:
+ print(output.generated_text[prev_idx:], end="", flush=True)
+ prev_idx = len(output.generated_text)
+ print()
+
+ request_ids: List[str] = [
+ inference_engine.add_request(
+ prompt=prompt, inference_parameters=sampling_params, streaming=True
+ )
+ for prompt in prompts
+ ]
+ stream_generators = [inference_engine.get_stream_generator(request_id) for request_id in request_ids]
+
+ tasks = [
+ asyncio.create_task(collect_stream(prompt, request_id, stream_generator))
+ for (prompt, request_id, stream_generator) in zip(prompts, request_ids, stream_generators)
+ ]
+
+ await inference_engine.run_engine_async()
+ await asyncio.gather(*tasks)
+
+ results: List[InferenceRequest] = [
+ inference_engine.scheduler.completed_request_pool[request_id] for request_id in request_ids
+ ]
+
+ return results
+
+def main():
+ """Main program."""
+
+ # Note: The default args passed here can be overwritten by using appropriate params (check arguments.py file)
+ # Micro batch size is not needed to be set by user. (It is calculated based on inference-batch-times-seqlen-threshold argument)
+ initialize_megatron(
+ extra_args_provider=add_text_generate_args,
+ args_defaults={
+ 'no_load_rng': True,
+ 'no_load_optim': True,
+ 'micro_batch_size': 1,
+ 'exit_on_missing_checkpoint': True,
+ },
+ )
+
+ # Set up model and load checkpoint
+ model = get_model(model_provider, wrap_with_ddp=False)
+ load_checkpoint(model, None, None)
+ model = model[0]
+
+ args = get_args()
+
+ inference_engine = get_inference_engine(args, model)
+
+ sampling_params = SamplingParams(
+ temperature=args.temperature,
+ top_k=args.top_k,
+ top_p=args.top_p,
+ return_log_probs=args.return_log_probs,
+ num_tokens_to_generate=args.num_tokens_to_generate,
+ )
+
+ if args.enable_cuda_graph:
+ print(f"Running warmup for CUDA graphs...")
+ inference_engine.generate(
+ prompts=args.prompts, sampling_params=sampling_params
+ )
+
+ start_time = time.perf_counter()
+ if args.stream:
+ results: List[InferenceRequest] = asyncio.run(generate(inference_engine, sampling_params, args.prompts))
+ else:
+ results: List[InferenceRequest] = inference_engine.generate(
+ prompts=args.prompts, sampling_params=sampling_params,
+ )
+ end_time = time.perf_counter()
+ latency = end_time - start_time
+
+ if torch.distributed.get_rank() == 0:
+ for idx, result in enumerate(results):
+ print(f' \n------------- RESULT FOR PROMPT {idx} --------------- ')
+ result = {
+ 'id': result.request_id,
+ 'input_prompt': result.prompt,
+ 'generated_text': result.generated_text,
+ 'generated_tokens': result.generated_tokens,
+ 'latency': latency,
+ }
+ print(result)
+
+ torch.distributed.destroy_process_group()
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/inference/llama_mistral/huggingface_reference.py b/examples/inference/llama_mistral/huggingface_reference.py
new file mode 100644
index 0000000000..9d8f4465f6
--- /dev/null
+++ b/examples/inference/llama_mistral/huggingface_reference.py
@@ -0,0 +1,25 @@
+import argparse
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+# Set up argument parsing
+parser = argparse.ArgumentParser(description="Script for text generation with a specific model and prompt.")
+parser.add_argument('--prompt', type=str, required=True, help="Prompt text to use for text generation")
+parser.add_argument('--model-path', type=str, required=True, help="Path to the Huggingface model checkpoint")
+
+# Parse command-line arguments
+args = parser.parse_args()
+
+model_path = args.model_path
+prompt = args.prompt
+
+config = AutoConfig.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path, config=config)
+model = AutoModelForCausalLM.from_pretrained(model_path, config=config).cuda()
+
+inputs = tokenizer(prompt, return_tensors="pt")
+for key in inputs:
+ inputs[key] = inputs[key].cuda()
+# top_k, top_p and do_sample are set for greedy argmax based sampling
+
+outputs = model.generate(**inputs, max_length=100, do_sample=False, top_p=0, top_k=0, temperature=1.0)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
\ No newline at end of file
diff --git a/examples/inference/llama_mistral/run_text_generation_llama3.1.sh b/examples/inference/llama_mistral/run_text_generation_llama3.1.sh
new file mode 100755
index 0000000000..06584f0917
--- /dev/null
+++ b/examples/inference/llama_mistral/run_text_generation_llama3.1.sh
@@ -0,0 +1,56 @@
+#!/bin/bash
+# This example will start serving the Llama3.1-8B model
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NVTE_APPLY_QK_LAYER_SCALING=0
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr 0.0.0.0 \
+ --master_port 6000"
+
+# Ensure CHECKPOINT and TOKENIZER_MODEL are provided
+if [ -z "$1" ] || [ -z "$2" ]; then
+ echo "Error: You must provide CHECKPOINT and TOKENIZER_MODEL as command-line arguments."
+ echo "Usage: $0 /path/to/checkpoint /path/to/tokenizer_model"
+ exit 1
+fi
+
+# Assign command-line arguments to variables
+CHECKPOINT=$1
+TOKENIZER_MODEL=$2
+
+pip install flask-restful
+
+torchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \
+ --use-checkpoint-args \
+ --disable-bias-linear \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model ${TOKENIZER_MODEL} \
+ --transformer-impl transformer_engine \
+ --normalization RMSNorm \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --no-masked-softmax-fusion \
+ --attention-softmax-in-fp32 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --untie-embeddings-and-output-weights \
+ --position-embedding-type rope \
+ --rotary-percent 1.0 \
+ --rotary-base 500000 \
+ --use-rope-scaling \
+ --use-rotary-position-embeddings \
+ --swiglu \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --ffn-hidden-size 14336 \
+ --load ${CHECKPOINT} \
+ --num-attention-heads 32 \
+ --max-position-embeddings 131072 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --seq-length 8192
diff --git a/examples/inference/llama_mistral/run_text_generation_llama3.sh b/examples/inference/llama_mistral/run_text_generation_llama3.sh
new file mode 100755
index 0000000000..c5fc4103ab
--- /dev/null
+++ b/examples/inference/llama_mistral/run_text_generation_llama3.sh
@@ -0,0 +1,55 @@
+#!/bin/bash
+# This example will start serving the Llama3-8B model
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NVTE_APPLY_QK_LAYER_SCALING=0
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr 0.0.0.0 \
+ --master_port 6000"
+
+# Ensure CHECKPOINT and TOKENIZER_MODEL are provided
+if [ -z "$1" ] || [ -z "$2" ]; then
+ echo "Error: You must provide CHECKPOINT and TOKENIZER_MODEL as command-line arguments."
+ echo "Usage: $0 /path/to/checkpoint /path/to/tokenizer_model"
+ exit 1
+fi
+
+# Assign command-line arguments to variables
+CHECKPOINT=$1
+TOKENIZER_MODEL=$2
+
+pip install flask-restful
+
+torchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \
+ --use-checkpoint-args \
+ --disable-bias-linear \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model ${TOKENIZER_MODEL} \
+ --transformer-impl transformer_engine \
+ --normalization RMSNorm \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --no-masked-softmax-fusion \
+ --attention-softmax-in-fp32 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --untie-embeddings-and-output-weights \
+ --position-embedding-type rope \
+ --rotary-percent 1.0 \
+ --rotary-base 500000 \
+ --use-rotary-position-embeddings \
+ --swiglu \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --ffn-hidden-size 14336 \
+ --load ${CHECKPOINT} \
+ --num-attention-heads 32 \
+ --max-position-embeddings 8192 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --seq-length 8192
diff --git a/examples/inference/llama_mistral/run_text_generation_mistral.sh b/examples/inference/llama_mistral/run_text_generation_mistral.sh
new file mode 100755
index 0000000000..4358fd494c
--- /dev/null
+++ b/examples/inference/llama_mistral/run_text_generation_mistral.sh
@@ -0,0 +1,53 @@
+#!/bin/bash
+# This example will start serving the Mistral-7B-v0.3 model
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr 0.0.0.0 \
+ --master_port 6000"
+
+# Ensure CHECKPOINT and TOKENIZER_MODEL are provided
+if [ -z "$1" ] || [ -z "$2" ]; then
+ echo "Error: You must provide CHECKPOINT and TOKENIZER_MODEL as command-line arguments."
+ echo "Usage: $0 /path/to/checkpoint /path/to/tokenizer_model"
+ exit 1
+fi
+
+# Assign command-line arguments to variables
+CHECKPOINT=$1
+TOKENIZER_MODEL=$2
+
+pip install flask-restful
+
+torchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \
+ --tokenizer-type HuggingFaceTokenizer \
+ --tokenizer-model ${TOKENIZER_MODEL} \
+ --use-checkpoint-args \
+ --apply-layernorm-1p \
+ --transformer-impl transformer_engine \
+ --normalization RMSNorm \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --no-masked-softmax-fusion \
+ --use-flash-attn \
+ --untie-embeddings-and-output-weights \
+ --disable-bias-linear \
+ --position-embedding-type rope \
+ --rotary-percent 1.0 \
+ --rotary-base 1000000 \
+ --swiglu \
+ --ffn-hidden-size 14336 \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --load ${CHECKPOINT} \
+ --num-attention-heads 32 \
+ --max-position-embeddings 4096 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --seq-length 4096 \
+ --seed 101
diff --git a/examples/run_text_generation_server_345M.sh b/examples/inference/run_text_generation_server_345M.sh
similarity index 92%
rename from examples/run_text_generation_server_345M.sh
rename to examples/inference/run_text_generation_server_345M.sh
index a151b98467..e8e61adb16 100755
--- a/examples/run_text_generation_server_345M.sh
+++ b/examples/inference/run_text_generation_server_345M.sh
@@ -26,9 +26,6 @@ torchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \
--fp16 \
--micro-batch-size 1 \
--seq-length 1024 \
- --out-seq-length 1024 \
- --temperature 1.0 \
--vocab-file $VOCAB_FILE \
--merge-file $MERGE_FILE \
- --top_p 0.9 \
--seed 42
diff --git a/examples/run_text_generation_server_345M_8_tensor_parallel.sh b/examples/inference/run_text_generation_server_345M_8_tensor_parallel.sh
similarity index 92%
rename from examples/run_text_generation_server_345M_8_tensor_parallel.sh
rename to examples/inference/run_text_generation_server_345M_8_tensor_parallel.sh
index 027ab42172..368cec3b31 100755
--- a/examples/run_text_generation_server_345M_8_tensor_parallel.sh
+++ b/examples/inference/run_text_generation_server_345M_8_tensor_parallel.sh
@@ -24,9 +24,6 @@ python -m torch.distributed.launch $DISTRIBUTED_ARGS tools/run_text_generation_s
--fp16 \
--micro-batch-size 1 \
--seq-length 1024 \
- --out-seq-length 1024 \
- --temperature 1.0 \
--vocab-file $VOCAB_FILE \
--merge-file $MERGE_FILE \
- --top_p 0.9 \
--seed 42
diff --git a/examples/inference/t5/simple_t5_batch_inference.py b/examples/inference/t5/simple_t5_batch_inference.py
new file mode 100644
index 0000000000..b4226d7de0
--- /dev/null
+++ b/examples/inference/t5/simple_t5_batch_inference.py
@@ -0,0 +1,157 @@
+import os
+import sys
+from argparse import Namespace
+
+import torch
+
+import pretrain_t5
+from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.inference.engines.abstract_engine import AbstractEngine
+from megatron.core.inference.engines.mcore_engine import MCoreEngine
+from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
+ InferenceWrapperConfig,
+)
+from megatron.core.inference.model_inference_wrappers.t5.t5_inference_wrapper import (
+ T5InferenceWrapper,
+)
+from megatron.core.inference.text_generation_controllers.encoder_decoder_text_generation_controller import (
+ EncoderDecoderTextGenerationController,
+)
+from megatron.core.transformer.module import MegatronModule
+from pretrain_t5 import model_provider
+
+sys.path.append(
+ os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
+)
+
+from typing import List
+
+from megatron.core import mpu
+from megatron.training import get_args, get_model, get_tokenizer
+from megatron.training.checkpointing import load_checkpoint
+from megatron.training.initialize import initialize_megatron
+
+
+def add_text_generate_args(parser):
+ """Text generation arguments."""
+ group = parser.add_argument_group(title='text generation')
+
+ group.add_argument("--temperature", type=float, default=1.0, help='Sampling temperature.')
+ group.add_argument("--top_k", type=int, default=1, help='Top k sampling.')
+ group.add_argument("--top_p", type=float, default=0.0, help='Top p sampling.')
+ group.add_argument(
+ "--return-log-probs",
+ action='store_true',
+ default=False,
+ help='Return the log probabilities of the final output tokens',
+ )
+ group.add_argument(
+ "--num-tokens-to-generate",
+ type=int,
+ default=30,
+ help='Number of tokens to generate for each prompt',
+ )
+ group.add_argument(
+ "--encoder-prompts",
+ metavar='N',
+ type=str,
+ nargs='+',
+ help='Encoder input prompts with each prompt within quotes and seperated by space',
+ )
+ group.add_argument(
+ "--max-batch-size", type=int, default=1, help='Max number of prompts to process at once'
+ )
+ return parser
+
+
+def get_inference_engine(args: Namespace, model: MegatronModule) -> AbstractEngine:
+ """Utility to get the relevant backend for running inference
+
+ This function will automatically chose the TRTLLMBackend when possible, and if not revert to Mcore backend if the user does not specify any backends. TRT LLM Backend is not implmented yet.
+
+ Args:
+ args (Namespace): The user arguments parsed from command line
+ model (MegatronModule): The megatron model .
+
+ Returns:
+ AbstractBackend: The chosen backend
+ """
+ tokenizer = get_tokenizer()
+
+ inference_wrapper_config = InferenceWrapperConfig(
+ hidden_size=args.hidden_size,
+ inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold,
+ fp32_residual_connection=args.fp32_residual_connection,
+ params_dtype=args.params_dtype,
+ padded_vocab_size=args.padded_vocab_size,
+ )
+
+ inference_wrapped_model = T5InferenceWrapper(model, inference_wrapper_config)
+ text_generation_controller = EncoderDecoderTextGenerationController(
+ inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer
+ )
+ return MCoreEngine(
+ text_generation_controller=text_generation_controller, max_batch_size=args.max_batch_size
+ )
+
+
+def main():
+ """Main program."""
+
+ # Note: The default args passed here can be overwritten by using appropriate params (check arguments.py file)
+ # Micro batch size is not needed to be set by user. (It is calculated based on inference-batch-times-seqlen-threshold argument)
+ initialize_megatron(
+ extra_args_provider=add_text_generate_args,
+ args_defaults={
+ 'no_load_rng': True,
+ 'no_load_optim': True,
+ 'micro_batch_size': 1,
+ 'exit_on_missing_checkpoint': True,
+ },
+ )
+
+ # Set up model and load checkpoint
+ model = get_model(model_provider, wrap_with_ddp=False)
+ load_checkpoint(model, None, None)
+ model = model[0]
+
+ args = get_args()
+
+ inference_engine = get_inference_engine(args, model)
+
+ sampling_params = SamplingParams(
+ temperature=args.temperature,
+ top_k=args.top_k,
+ top_p=args.top_p,
+ return_log_probs=args.return_log_probs,
+ num_tokens_to_generate=args.num_tokens_to_generate,
+ )
+
+ tokenizer = get_tokenizer()
+ decoder_prompts = [""] * len(
+ args.encoder_prompts
+ ) # for T5, the prompt is provided as encoder input, hence decoder_prompts is empty
+ args.prompts = decoder_prompts
+
+ results: List[InferenceRequest] = inference_engine.generate(
+ prompts=args.prompts,
+ add_BOS=True,
+ encoder_prompts=args.encoder_prompts,
+ sampling_params=sampling_params,
+ )
+
+ if torch.distributed.get_rank() == 0:
+ for idx, result in enumerate(results):
+ print(f' \n------------- RESULT FOR PROMPT {idx} --------------- ')
+ result = {
+ 'id': result.request_id,
+ 'input_prompt': result.prompt,
+ 'generated_text': result.generated_text,
+ 'generated_tokens': result.generated_tokens,
+ }
+ print(result)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/mamba/.gitignore b/examples/mamba/.gitignore
new file mode 100644
index 0000000000..940f4797e4
--- /dev/null
+++ b/examples/mamba/.gitignore
@@ -0,0 +1,4 @@
+checkpoints/
+data-cache/
+tensorboard/
+triton-cache/
diff --git a/examples/mamba/Dockerfile b/examples/mamba/Dockerfile
new file mode 100644
index 0000000000..2e194095b7
--- /dev/null
+++ b/examples/mamba/Dockerfile
@@ -0,0 +1,32 @@
+FROM nvcr.io/nvidia/pytorch:24.01-py3
+
+RUN pip uninstall -y triton && \
+ pip install triton==2.1.0 sentencepiece==0.1.99 flask-restful
+
+# The causal-conv1d and mamba-ssm packages below are built from scratch here
+# (which takes significant time) because there are no wheels available on PyPI
+# for these relatively newer versions of the packages that are compatible with
+# the older NGC-variant PyTorch version (e.g. version 2.2.0.dev231106) that we
+# are using (in the NGC base container). Generally, if the package is not
+# compatible with the PyTorch version, then it will generate a Python import
+# error. The package authors tend to only release wheels for new versions of
+# these pacakges which are compatible with the versions of regular PyTorch and
+# NGC-variant PyTorch that are newer at the time of release. So, to use newer
+# versions of these packages with relatively older versions of the NGC PyTorch
+# container, we tend to have to build the packages from scratch.
+
+RUN cd /tmp && \
+ git clone https://github.com/Dao-AILab/causal-conv1d.git && \
+ cd causal-conv1d && \
+ git checkout v1.2.2.post1 && \
+ CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install . && \
+ cd .. && \
+ rm -rf causal-conv1d
+
+RUN cd /tmp && \
+ git clone https://github.com/state-spaces/mamba.git && \
+ cd mamba && \
+ git checkout v2.0.3 && \
+ MAMBA_FORCE_BUILD=TRUE pip install . && \
+ cd .. && \
+ rm -rf mamba
diff --git a/examples/mamba/README.md b/examples/mamba/README.md
new file mode 100644
index 0000000000..f8f6d79683
--- /dev/null
+++ b/examples/mamba/README.md
@@ -0,0 +1,94 @@
+# Mamba-based Language Models
+
+## Introduction
+
+This document is an entrypoint into the code used for
+[An Empirical Study of Mamba-based Language Models](https://arxiv.org/abs/2406.07887).
+
+We are releasing the parameters for some of the models described in that
+technical report via
+[HuggingFace](https://huggingface.co/collections/nvidia/ssms-666a362c5c3bb7e4a6bcfb9c).
+The code in the `main` branch is no longer compatible with the `Mamba2-*`
+checkpoints. You can load them using the
+[fixed snapshot of the code used for the technical report](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba).
+
+## Installation
+
+Create and run a Docker container using the [Dockerfile](./Dockerfile).
+
+```
+docker build -t your_image_name:your_tag .
+docker run --gpus all -it --rm \
+ -v /path/to/megatron:/workspace/megatron \
+ -v /path/to/dataset:/workspace/dataset \
+ -v /path/to/checkpoints:/workspace/checkpoints \
+ -w /workspace/megatron/examples/mamba \
+ your_image_name:your_tag
+```
+
+## Train
+
+[`train.sh`](./train.sh) is an example pretraining script, showing how to run on
+a single node. Select between 800M-scale and 8B-scale models by setting the
+`MODEL_SCALE` variable. The 8B-scale hybrid model architecture is the same as
+the one described in the technical report.
+
+## Text Generation
+
+Use [`run_text_gen_server_8b.sh`](./run_text_gen_server_8b.sh) to start a text
+generation server using an 8B hybrid checkpoint. This is configured to run the
+8B hybrid model described in the technical report, with tensor model parallel
+set to 1.
+
+The arguments in the script will need to be changed if using a checkpoint with a
+different model parallel configuration or other differences, such as model
+architecture. For example, to run the 8B pure Mamba-2 model, change
+`--hybrid-attention-ratio` and `--hybrid-mlp-ratio` to 0.0, or remove them.
+
+Use [`run_text_gen_server_8b_gpt3.sh`](./run_text_gen_server_8b_gpt3.sh) to start
+a text generation server using the 8B reference Transformer checkpoint.
+
+## Checkpoint Formats
+
+For inference, the model must be configured to match the checkpoint file used,
+including the hybrid layer configuration and model parallel configuration.
+
+If you need to convert a hybrid checkpoint file to a different tensor parallel
+or pipeline parallel size, use
+[the hybrid conversion script](../../tools/checkpoint/hybrid_conversion.py).
+There is an example run command at the end of that file.
+
+Before running that script, you will need to set `PYTHONPATH` to include the
+root directory of your Megatron-LM repository clone.
+
+```
+export PYTHONPATH=:PYTHONPATH
+```
+
+## Hybrid Options
+
+`--hybrid-attention-ratio ATT` specifies a target ratio of attention layers
+to total layers. For example, 4 attention layers out of 48 total layers is
+specified by `--hybrid-attention-ratio 0.08`.
+
+`--hybrid-mlp-ratio MLP` specifies a target ratio of MLP layers to total
+layers. For example, 24 MLP layers out of 48 total layers is specified by
+`--hybrid-mlp-ratio 0.5`.
+
+* (`ATT` + `MLP`) must be less than or equal to 1.0.
+* (1.0 - `ATT` - `MLP`) is the hybrid mamba ratio, the ratio of mamba layers to
+total layers.
+* `ATT` = `MLP` = 0 is a pure Mamba model.
+* `ATT` = `MLP` = 0.5 is a transfomer model.
+
+If either `ATT` or `MLP` is greater than 0.0 or if `--hybrid-override-pattern`
+is specified, the logfile will include information about the hybrid layer
+pattern used. `--hybrid-override-pattern` can be used to specify a different
+pattern than the default, algorithmically-generated one.
+
+## Mamba vs Mamba-2
+
+This codebase currently only supports Mamba-2, and not the original version of
+Mamba. However, the
+[fixed snapshot of the code used for the technical report](https://github.com/NVIDIA/Megatron-LM/tree/ssm/examples/mamba)
+can be configured to run the original version of Mamba.
diff --git a/examples/mamba/run_text_gen_server_8b.sh b/examples/mamba/run_text_gen_server_8b.sh
new file mode 100755
index 0000000000..8d3137f244
--- /dev/null
+++ b/examples/mamba/run_text_gen_server_8b.sh
@@ -0,0 +1,50 @@
+#!/bin/bash
+
+# Use: ./run_text_gen_server_8b.sh
+# To launch the client: python ../../tools/text_generation_cli.py
+
+CHECKPOINT_PATH=$1
+TOKENIZER_PATH=$2
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr localhost \
+ --master_port 6000"
+
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NCCL_IB_TIMEOUT=19
+export NCCL_IB_QPS_PER_CONNECTION=4
+
+export TRITON_CACHE_DIR="./triton-cache/"
+export TRITON_CACHE_MANAGER="megatron.core.ssm.triton_cache_manager:ParallelFileCacheManager"
+
+torchrun $DISTRIBUTED_ARGS ../../tools/run_mamba_text_generation_server.py \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --untie-embeddings-and-output-weights \
+ --num-layers 56 \
+ --hidden-size 4096 \
+ --load ${CHECKPOINT_PATH} \
+ --num-attention-heads 32 \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --hybrid-attention-ratio 0.08 \
+ --hybrid-mlp-ratio 0.5 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --disable-bias-linear \
+ --normalization RMSNorm \
+ --seq-length 4096 \
+ --max-position-embeddings 4096 \
+ --position-embedding-type none \
+ --tokenizer-type GPTSentencePieceTokenizer \
+ --tokenizer-model ${TOKENIZER_PATH} \
+ --distributed-backend nccl \
+ --distributed-timeout-minutes 1440 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --use-mcore-models \
+ --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \
+ --seed 42
diff --git a/examples/mamba/run_text_gen_server_8b_gpt3.sh b/examples/mamba/run_text_gen_server_8b_gpt3.sh
new file mode 100644
index 0000000000..5413b245ed
--- /dev/null
+++ b/examples/mamba/run_text_gen_server_8b_gpt3.sh
@@ -0,0 +1,46 @@
+#!/bin/bash
+
+# Use: ./run_text_gen_server_8b_gpt3.sh
+# To launch the client: python ../../tools/text_generation_cli.py
+
+CHECKPOINT_PATH=$1
+TOKENIZER_PATH=$2
+
+DISTRIBUTED_ARGS="--nproc_per_node 1 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr localhost \
+ --master_port 6000"
+
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NCCL_IB_TIMEOUT=19
+export NCCL_IB_QPS_PER_CONNECTION=4
+
+torchrun $DISTRIBUTED_ARGS ../../tools/run_text_generation_server.py \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 1 \
+ --use-flash-attn \
+ --apply-layernorm-1p \
+ --untie-embeddings-and-output-weights \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --load ${CHECKPOINT_PATH} \
+ --num-attention-heads 32 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --disable-bias-linear \
+ --seq-length 4096 \
+ --max-position-embeddings 4096 \
+ --position-embedding-type rope \
+ --rotary-percent 0.5 \
+ --squared-relu \
+ --tokenizer-type GPTSentencePieceTokenizer \
+ --tokenizer-model ${TOKENIZER_PATH} \
+ --distributed-backend nccl \
+ --distributed-timeout-minutes 1440 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --use-mcore-models \
+ --transformer-impl local \
+ --seed 42
diff --git a/examples/mamba/train.sh b/examples/mamba/train.sh
new file mode 100755
index 0000000000..3952a997d4
--- /dev/null
+++ b/examples/mamba/train.sh
@@ -0,0 +1,105 @@
+#!/bin/bash
+
+# Use: ./train.sh
+
+MODEL_SCALE="800M" # or "8B"
+
+case "${MODEL_SCALE}" in
+ "800M")
+ TENSOR_MODEL_PARALLEL_SIZE=1
+ NUM_LAYERS=48
+ HIDDEN_SIZE=1024
+ NUM_ATTENTION_HEADS=16
+ GLOBAL_BATCH_SIZE=32
+ ;;
+ "8B")
+ TENSOR_MODEL_PARALLEL_SIZE=4
+ NUM_LAYERS=56
+ HIDDEN_SIZE=4096
+ NUM_ATTENTION_HEADS=32
+ GLOBAL_BATCH_SIZE=8
+ ;;
+ *)
+ echo "Invalid version specified"
+ exit 1
+ ;;
+esac
+
+DATA_PATH=$1
+TOKENIZER_PATH=$2
+
+export NCCL_IB_SL=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NCCL_IB_TIMEOUT=19
+export NCCL_IB_QPS_PER_CONNECTION=4
+
+CHECKPOINT_DIR="./checkpoints"
+DATACACHE_DIR="./data-cache"
+TENSORBOARD_DIR="./tensorboard"
+
+mkdir -p ${CHECKPOINT_DIR}
+mkdir -p ${DATACACHE_DIR}
+mkdir -p ${TENSORBOARD_DIR}
+
+export TRITON_CACHE_DIR="./triton-cache/"
+export TRITON_CACHE_MANAGER="megatron.core.ssm.triton_cache_manager:ParallelFileCacheManager"
+
+SEQ_LEN=4096
+TRAIN_SAMPLES=73242188 # 300B tokens / 4096
+LR_WARMUP_SAMPLES=50000
+LR_DECAY_SAMPLES=73192188 # TRAIN_SAMPLES - LR_WARMUP_SAMPLES
+
+options=" \
+ --tensor-model-parallel-size ${TENSOR_MODEL_PARALLEL_SIZE} \
+ --sequence-parallel \
+ --pipeline-model-parallel-size 1 \
+ --use-distributed-optimizer \
+ --overlap-param-gather \
+ --overlap-grad-reduce \
+ --untie-embeddings-and-output-weights \
+ --init-method-std 0.02 \
+ --position-embedding-type none \
+ --num-layers ${NUM_LAYERS} \
+ --hidden-size ${HIDDEN_SIZE} \
+ --num-attention-heads ${NUM_ATTENTION_HEADS} \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --hybrid-attention-ratio 0.08 \
+ --hybrid-mlp-ratio 0.5 \
+ --seq-length ${SEQ_LEN} \
+ --max-position-embeddings ${SEQ_LEN} \
+ --train-samples ${TRAIN_SAMPLES} \
+ --lr-warmup-samples ${LR_WARMUP_SAMPLES} \
+ --lr-decay-samples ${LR_DECAY_SAMPLES} \
+ --save ${CHECKPOINT_DIR} \
+ --load ${CHECKPOINT_DIR} \
+ --data-path ${DATA_PATH} \
+ --data-cache-path ${DATACACHE_DIR} \
+ --split 99,1,0 \
+ --tokenizer-type GPTSentencePieceTokenizer \
+ --tokenizer-model ${TOKENIZER_PATH} \
+ --distributed-backend nccl \
+ --micro-batch-size 4 \
+ --global-batch-size ${GLOBAL_BATCH_SIZE} \
+ --lr 2.5e-4 \
+ --min-lr 2.5e-5 \
+ --lr-decay-style cosine \
+ --weight-decay 0.1 \
+ --clip-grad 1.0 \
+ --attention-dropout 0.0 \
+ --hidden-dropout 0.0 \
+ --disable-bias-linear \
+ --normalization RMSNorm \
+ --adam-beta1 0.9 \
+ --adam-beta2 0.95 \
+ --log-interval 10 \
+ --save-interval 2000 \
+ --eval-interval 2000 \
+ --eval-iters 32 \
+ --bf16 \
+ --use-mcore-models \
+ --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \
+ --no-create-attention-mask-in-dataloader \
+ --tensorboard-dir ${TENSORBOARD_DIR}"
+
+torchrun --nproc_per_node 8 ../../pretrain_mamba.py ${options}
diff --git a/examples/merge_mp_bert.sh b/examples/merge_mp_bert.sh
deleted file mode 100755
index 1383433284..0000000000
--- a/examples/merge_mp_bert.sh
+++ /dev/null
@@ -1,18 +0,0 @@
-#!/bin/bash
-
-TENSOR_MODEL_PARALLEL_SIZE=2
-
-VOCAB_FILE=bert-vocab.txt
-CHECKPOINT_PATH=checkpoints/bert_345m
-
-WORLD_SIZE=$TENSOR_MODEL_PARALLEL_SIZE python tools/merge_mp_partitions.py \
- --model-type BERT \
- --tensor-model-parallel-size $TENSOR_MODEL_PARALLEL_SIZE \
- --tokenizer-type BertWordPieceLowerCase \
- --vocab-file $VOCAB_FILE \
- --num-layers 24 \
- --hidden-size 1024 \
- --num-attention-heads 16 \
- --seq-length 512 \
- --max-position-embeddings 512 \
- --load $CHECKPOINT_PATH
diff --git a/examples/mixtral/README.md b/examples/mixtral/README.md
new file mode 100644
index 0000000000..e85eccd6ef
--- /dev/null
+++ b/examples/mixtral/README.md
@@ -0,0 +1,132 @@
+# Mixtral 8x7B Model Inference and Finetuning
+
+## Download Mixtral 8x7B Checkpoints
+Download Mixtral 8x7B HF format checkpoint from [HF-hub](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1/)
+
+Or you can simply run this following script to download Mixtral 8x7B into a specific folder.
+```python
+from huggingface_hub import snapshot_download
+SAVED_DIR = "" # Specify the saved directory
+# Download HF checkpoints
+snapshot_download(repo_id="mistralai/Mixtral-8x7B-v0.1", ignore_patterns=["*.pt"], local_dir=SAVED_DIR, local_dir_use_symlinks=False)
+```
+
+## Convert Mixtral 8x7B checkpoints from HF to MCore
+The HF checkpoints can be converted to Megatron format by using the provided checkpoint converter for HF format.
+The target model parallel size(e.g. TP,PP,EP) should be specified.
+
+Currently the converter doesn't support distributed checkpointing yet, so each different parallel config requires a specific checkpoint.
+- For training, the recommended model parallel config is TP1EP8PP4
+- For inference, the recommended model parallel config is TP1EP1PP2
+
+```
+TOKENIZER_MODEL=/workspace/checkpoints/mixtral-hf/tokenizer.model
+MEGATRON_PATH="/workspace/megatron-lm"
+export PYTHONPATH=$MEGATRON_PATH:$PYTHONPATH
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+TARGET_TP_SIZE=""
+TARGET_EP_SIZE=""
+TARGET_PP_SIZE=""
+
+HF_FORMAT_DIR=/workspace/checkpoints/mixtral-hf
+MEGATRON_FORMAT_DIR=/workspace/checkpoints/mixtral-mcore-TP${TARGET_TP_SIZE}PP${TARGET_PP_SIZE}EP${TARGET_EP_SIZE}
+
+python tools/checkpoint/convert.py \
+--model-type GPT \
+--loader loader_mixtral_hf \
+--saver mcore \
+--target-tensor-parallel-size ${TARGET_TP_SIZE} \
+--target-pipeline-parallel-size ${TARGET_PP_SIZE} \
+--target-expert-parallel-size ${TARGET_EP_SIZE} \
+--load-dir ${HF_FORMAT_DIR} \
+--save-dir ${MEGATRON_FORMAT_DIR} \
+--tokenizer-model ${TOKENIZER_MODEL}
+```
+
+## Text generation with Mixtral 8x7B
+Inference with Mixtral 8x7B requires at least 2 GPUS, such that a distributed checkpoint with EP>=2 or PP>=2 converted with above script is needed.
+
+The Megatron-LM have included a simple REST server to use for text generation in `tools/run_text_generation_server.py`, launch it with the following script:
+```
+#!/bin/bash
+# This example will start serving the Mixtral 8x7B model.
+DISTRIBUTED_ARGS="--nproc_per_node 2 \
+ --nnodes 1 \
+ --node_rank 0 \
+ --master_addr localhost \
+ --master_port 6000"
+
+CHECKPOINT=
+TOKENIZER_MODEL=
+
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+
+pip install flask-restful
+
+torchrun $DISTRIBUTED_ARGS tools/run_text_generation_server.py \
+ --tensor-model-parallel-size 1 \
+ --pipeline-model-parallel-size 2 \
+ --expert-model-parallel-size 1 \
+ --load ${CHECKPOINT} \
+ --tokenizer-type Llama2Tokenizer \
+ --tokenizer-model $TOKENIZER_MODEL \
+ --use-mcore-models \
+ --max-position-embeddings 32768 \
+ --num-layers 32 \
+ --hidden-size 4096 \
+ --ffn-hidden-size 14336 \
+ --num-attention-heads 32 \
+ --normalization RMSNorm \
+ --disable-bias-linear \
+ --position-embedding-type rope \
+ --no-position-embedding \
+ --swiglu \
+ --untie-embeddings-and-output-weights \
+ --group-query-attention \
+ --num-query-groups 8 \
+ --bf16 \
+ --micro-batch-size 1 \
+ --seq-length 1024 \
+ --seed 42 \
+ --num-experts 8 \
+ --moe-router-topk 2 \
+ --moe-token-dispatcher-type alltoall \
+ --moe-grouped-gemm \
+ --mock-data \
+ --rotary-base 1000000
+```
+
+Once the server is running you can use `tools/text_generation_cli.py` to query it, it takes one argument which is the host the server is running on.
+
+```
+python tools/text_generation_cli.py localhost:5000
+```
+
+
+## Finetuning from pretrained Mixtral 8x7B
+To finetuning pretrained Mixtral 8x7B, use the following scripts:
+
+
+```bash
+PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:24.04-py3
+CHECKPOINT_PATH="" # Speicfy path to checkpoint dir
+TOKENIZER_MODEL="" # Specify path to tokenizer.model
+DATA_PATH="" # Specify path to data
+
+docker run \
+ --gpus=all \
+ --ipc=host \
+ --workdir /workspace/megatron-lm \
+ -v /path/to/data:/path/to/data \
+ -v /path/to/megatron-lm:/workspace/megatron-lm \
+ $PYTORCH_IMAGE \
+ bash examples/mixtral/train_mixtral_8x7b_distributed.sh $CHECKPOINT_PATH $TOKENIZER_MODEL $DATA_PATH
+```
+
+The above functionality also applys to Mixtral 8x22B actually, you should set the model config (including hidden_size/head_num/num_layers/ffn_hidden_size) properly according to the original [config](https://huggingface.co/mistralai/Mixtral-8x22B-v0.1/blob/main/config.json).
+
+## Acknowledgements
+Contributors outside NVIDIA for the huggingface converter and example of Mixtral models in Megatron-Core:
+- Peng Li
+- Jun Huang
diff --git a/examples/mixtral/train_mixtral_8x7b_distributed.sh b/examples/mixtral/train_mixtral_8x7b_distributed.sh
new file mode 100644
index 0000000000..ed44d60f5c
--- /dev/null
+++ b/examples/mixtral/train_mixtral_8x7b_distributed.sh
@@ -0,0 +1,116 @@
+#!/bin/bash
+
+# Runs Mixtral 8x7B model
+
+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=${SLURM_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=(
+ --use-mcore-models
+ --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
+ --rotary-base 1000000
+)
+
+MOE_ARGS=(
+ --num-experts 8
+ --moe-router-topk 2
+ --moe-router-load-balancing-type aux_loss
+ --moe-aux-loss-coeff 1e-2
+ --moe-grouped-gemm
+ --moe-token-dispatcher-type alltoall
+ --overlap-param-gather
+ --overlap-grad-reduce
+)
+
+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 256
+ --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
+)
+
+MODEL_PARALLEL_ARGS=(
+ --tensor-model-parallel-size 1
+ --pipeline-model-parallel-size 4
+ --expert-model-parallel-size 8
+ --use-distributed-optimizer
+ --sequence-parallel
+)
+
+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"}
+ --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[@]}
diff --git a/examples/multimodal/Dockerfile b/examples/multimodal/Dockerfile
new file mode 100644
index 0000000000..7b54091ae6
--- /dev/null
+++ b/examples/multimodal/Dockerfile
@@ -0,0 +1,26 @@
+FROM nvcr.io/nvidia/pytorch:24.02-py3
+
+RUN apt update && \
+ apt -y upgrade && \
+ apt install -y --no-install-recommends \
+ software-properties-common \
+ build-essential \
+ python3-pip \
+ python3-dev \
+ bash \
+ git \
+ vim \
+ tmux \
+ python-is-python3 \
+ default-jre
+
+RUN pip install --upgrade pip
+RUN pip install einops einops-exts sentencepiece braceexpand webdataset packaging
+RUN pip install transformers datasets accelerate timm
+RUN pip install pytest-cov pytest_mock nltk wrapt
+RUN pip install zarr "tensorstore==0.1.45"
+RUN pip install black isort click==8.0.2
+RUN pip install pycocoevalcap megatron-energon mistral-common tiktoken
+RUN pip install git+https://github.com/openai/CLIP.git
+# Use --no-deps for the following to avoid outdated and unnecessary dependencies.
+RUN pip install open_clip_torch open-flamingo[eval] --no-deps
diff --git a/examples/multimodal/README.md b/examples/multimodal/README.md
new file mode 100644
index 0000000000..a65839f8f1
--- /dev/null
+++ b/examples/multimodal/README.md
@@ -0,0 +1,157 @@
+# Multimodal Example
+
+*NOTE: This example is under active development and is expected change.*
+
+The following walks through all the steps required to pretrain and instruction tune a llava architecture vision-language model (VLM). It is important to precisely follow all steps to obtain the benchmark scores at the end.
+
+This example has been tested on an A100 based DGX cluster. Pretraining and instruction tuning took approximately 1 day and 11 hours respectively on 64 GPUs using four way tensor parallelism (tp=4). Training speed will scale approximately linearly with number of GPUs available.
+
+Multimodal support in megatron is still under active development. This example is not intended to produce state-of-the-art model quality (that would require more data and model refinements), it is merely intended to demonstrate the multimodal functionality in megatron. If you hit any problems, please open a github issue.
+
+## Setup
+
+### Docker container
+
+You can build a docker container using `examples/multimodal/Dockerfile` to run this example.
+
+### Language model
+
+Follow the instructions in [Mistral](../../docs/llama_mistral.md#mistral-7b) to download weights for Mistral-7B-Instruct-v0.3 from HuggingFace and convert to mcore format with tensor parallel size 4.
+Please use the tokenizer from HuggingFace.
+
+### Vision model
+
+This example uses the OpenAI CLIP `ViT-L/14@336px` Vision model. To download the weights from OpenAI and convert them to a format that can be loaded in megatron, please run the following:
+
+```
+python examples/multimodal/model_converter/clip_converter.py --download-root /some/download/folder --output /some/output/folder --tensor-parallel-size 4 --use-te
+```
+
+### Combined model checkpoint
+
+Update the paths to point to the mcore converted CLIP and Mistral models and run the following script to combine the Mistral and CLIP models into a single multimodal checkpoint folder:
+
+```
+examples/multimodal/combine_lm_vision_checkpoints.sh /path/to/mistral/model /path/to/clip/model /output/dir
+```
+
+## Training
+
+### Pretraining
+
+1. Download the LLavA-Pretrain dataset from Hugging Face and unzip the images folder (NOTE: 79GB of disk space required):
+
+ ```
+ git clone https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain
+ cd LLaVA-Pretrain
+ unzip images.zip
+ ```
+
+3. Run the following script to convert the data to webdataset format:
+
+ ```
+ cd
+ python examples/multimodal/convert_llava_pretrain_to_wds.py
+ ```
+
+4. Run the following command to convert to megatron-energon format:
+
+ ```
+ cd /wds
+ energon prepare ./
+ ```
+
+ select the following values for the presented options:
+
+ ```
+ > Please enter a desired train/val/test split like "0.5, 0.2, 0.3" or "8,1,1": 9,1,0
+ > Do you want to create a dataset.yaml interactively? [Y/n]: Y
+ > Please enter a number to choose a class: 10 (VQAWebdataset)
+ > Do you want to set a simple field_map[Y] (or write your own sample_loader [n])? [Y/n]: Y
+ > Please enter a webdataset field name for 'image' (): jpg
+ > Please enter a webdataset field name for 'context' (): json[0][value]
+ > Please enter a webdataset field name for 'answers' (typing.Optional[typing.List[str]], default: None): json[1][value]
+ > Please enter a webdataset field name for 'answer_weights' (typing.Optional[torch.Tensor], default: None):
+ ```
+
+5. Update `pretrain_dataset.yaml` so that both `path` variables point to `LLaVA-Pretrain/wds`
+
+6. Run the following script to pretrain a llava model for image captioning:
+
+ ```
+ cd
+ examples/multimodal/pretrain_mistral_clip.sh
+ ```
+
+All being well you should observe training and validation loss curves similar to the following:
+
+
+
+These curves were obtained with global batch size of 256. Changing this value will likely change the curves. For pretraining and instruction tuning llava models we have found that loss curves are an unreliable predictor of downstream task performance. Therefore it is necessary to run test generation and evaluation on a range of metrics to understand model quality. We intend to add training time zero-shot evaluation in a future update.
+
+You can execute the pretraining script multiple times to resume training. On resuming, the latest model, optimizer, and dataloader state are loaded.
+
+### SFT
+
+1. Prepare an instruction tuning dataset such in [megatron-energon format](https://nvidia.github.io/Megatron-Energon/data_prep.html#). NOTE: we do not provide instructions for this.
+
+2. Update `sft_dataset.yaml` so that both `path` variables point to the train and val splits of your instruction tuning dataset.
+
+Run the following script to instruction tune the pre-trained llava model:
+
+ ```
+ examples/multimodal/sft_mistral_clip.sh
+ ```
+
+You can execute the SFT script multiple times to resume training. On resuming, the latest model, optimizer, and dataloader state are loaded.
+
+## Evaluation
+
+### Generation
+
+Run the following script:
+
+```
+examples/multimodal/text_generation_mistral_clip.sh --input-image-path /path/to/input/images --output-path /some/output/directory \
+ --model-path /path/to/model.pt --gt-path /path/to/groundtruth/file --task generation-task-name
+```
+
+where `--task generation-task-name` is the name of the evaluation benchmark such as `captioning` or `MMMU`.
+
+### After pretraining
+
+#### COCO captioning
+
+1. Download the COCO 2014 test image set:
+
+ ```wget http://images.cocodataset.org/zips/test2014.zip```
+
+2. Download COCO test image annotations:
+
+ ```https://storage.googleapis.com/sfr-vision-language-research/datasets/coco_karpathy_test.json```
+
+3. First, run text generation using `--task captioning`.
+
+4. Run the following command:
+
+ ```
+ python examples/multimodal/evaluate_coco.py --input-path /output/directory/from/generation --groundtruth-path /path/to/groundtruth/file
+ ```
+
+For the mistral-7b-instruct plus clip llava model you should obtain a COCO CIDer score of approximately 94.
+
+### After SFT
+
+#### MMMU
+
+The official MMMU repository is not pip installable currently so please clone their code in `examples/multimodal` by running `git clone https://github.com/MMMU-Benchmark/MMMU.git`.
+
+The MMMU dataset is loaded from HuggingFace automatically as part of the code.
+
+Run text generation using `--task MMMU`. Then, run the following command:
+
+```
+python examples/multimodal/evaluate_mmmu.py --input-path /output/directory/from/generation
+```
+
+For the mistral-7b-instruct plus clip instruction tuned llava model you should obtain a MMMU score of approximately 38.
diff --git a/examples/multimodal/assets/pretrain_curves.png b/examples/multimodal/assets/pretrain_curves.png
new file mode 100644
index 0000000000..7981a73ba1
Binary files /dev/null and b/examples/multimodal/assets/pretrain_curves.png differ
diff --git a/examples/multimodal/combine_lm_vision_checkpoints.sh b/examples/multimodal/combine_lm_vision_checkpoints.sh
new file mode 100755
index 0000000000..52de16ecd2
--- /dev/null
+++ b/examples/multimodal/combine_lm_vision_checkpoints.sh
@@ -0,0 +1,57 @@
+#/bin/bash
+MCORE_LM=$1 #
+MCORE_VISION=$2 #
+OUTPUT_DIR=$3 #
+MODEL_TYPE=$4 # Model type. Default: Mistral CLIP example.
+
+if [[ $MODEL_TYPE == "nvlm" ]]; then
+ # NVLM TP=8
+ python examples/multimodal/combine_state_dicts.py \
+ --input \
+ ${MCORE_LM}/iter_0000001/mp_rank_00/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_00/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_01/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_01/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_02/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_02/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_03/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_03/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_04/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_04/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_05/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_05/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_06/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_06/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_07/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_07/model_optim_rng.pt \
+ --prefixes language_model vision_model language_model vision_model language_model vision_model language_model vision_model language_model vision_model language_model vision_model language_model vision_model language_model vision_model \
+ --output \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_00/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_01/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_02/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_03/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_04/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_05/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_06/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_07/model_optim_rng.pt
+else
+ # Mistral CLIP example TP=4.
+ python examples/multimodal/combine_state_dicts.py \
+ --input \
+ ${MCORE_LM}/iter_0000001/mp_rank_00/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_00/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_01/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_01/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_02/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_02/model_optim_rng.pt \
+ ${MCORE_LM}/iter_0000001/mp_rank_03/model_optim_rng.pt \
+ ${MCORE_VISION}/iter_0000001/mp_rank_03/model_optim_rng.pt \
+ --prefixes language_model vision_model language_model vision_model language_model vision_model language_model vision_model \
+ --output \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_00/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_01/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_02/model_optim_rng.pt \
+ ${OUTPUT_DIR}/iter_0000001/mp_rank_03/model_optim_rng.pt
+fi
+
+echo 1 > ${OUTPUT_DIR}/latest_checkpointed_iteration.txt
diff --git a/examples/multimodal/combine_state_dicts.py b/examples/multimodal/combine_state_dicts.py
new file mode 100644
index 0000000000..2f7028474c
--- /dev/null
+++ b/examples/multimodal/combine_state_dicts.py
@@ -0,0 +1,81 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+
+import argparse
+import os
+import sys
+
+import torch
+
+# Add megatron to the path.
+sys.path.append(
+ os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
+)
+
+
+def combine(input_files, module_prefixes, output_files):
+ num_inputs_per_output = int(len(input_files) / len(output_files))
+
+ for output_idx, output_file in enumerate(output_files):
+ combined_state_dict = None
+
+ lb = output_idx * num_inputs_per_output
+ ub = (output_idx + 1) * num_inputs_per_output
+ current_input_files = input_files[lb:ub]
+ current_module_prefixes = module_prefixes[lb:ub]
+
+ for i, (input_file, module_prefix) in enumerate(
+ zip(current_input_files, current_module_prefixes)
+ ):
+ # initialize the combined state dict using the first provided input file
+ current_state_dict = torch.load(input_file)
+ if i == 0:
+ combined_state_dict = current_state_dict.copy()
+ combined_state_dict["model"] = dict()
+
+ # copy model state dict and prefix names with the given module keys.
+ for k, v in current_state_dict["model"].items():
+ combined_state_dict["model"]["%s.%s" % (module_prefix, k)] = v
+
+ output_dir = os.path.dirname(output_file)
+ if not os.path.exists(output_dir):
+ os.makedirs(output_dir, exist_ok=True)
+ torch.save(combined_state_dict, output_file)
+ print("saved:", output_file)
+
+ print("done.")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description="""
+ Combine multiple state dicts into a single state dict.
+ The combined state dict is first initialized by taking a copy of the first provided input state dict.
+ To avoid conflicts in model parameter names, a prefix must be provided for each input file.
+ Model parameter names will be renamed from to ..
+
+
+ Example usage:
+ python combine_state_dicts.py --input language_model.pt vision_model.pt --prefixes language_model vision_model --output multimodal.pt
+ """,
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ )
+ parser.add_argument("--input", nargs="*", required=True, help="paths to input state dict files")
+ parser.add_argument(
+ "--prefixes",
+ nargs="*",
+ required=True,
+ help="prefixes to use with each input model's parameters",
+ )
+ parser.add_argument(
+ "--output", nargs="*", required=True, help="path(s) to output state dict file"
+ )
+
+ args = parser.parse_args()
+
+ assert len(args.input) > 1, "must provide more than 1 input model to combine"
+ assert len(args.input) == len(args.prefixes), "each input model must have a corresponding key"
+ assert (
+ len(args.input) % len(args.output) == 0
+ ), "each output file must use the same number of input files"
+
+ combine(args.input, args.prefixes, args.output)
diff --git a/examples/multimodal/config.py b/examples/multimodal/config.py
new file mode 100644
index 0000000000..e0de36f7a2
--- /dev/null
+++ b/examples/multimodal/config.py
@@ -0,0 +1,280 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from dataclasses import dataclass
+
+import torch
+
+from megatron.training.activations import fast_gelu, quick_gelu, squared_relu
+
+
+def get_language_model_config(config):
+ if config.language_model_type == "llama3_8b":
+ config.activation_func = torch.nn.functional.silu
+ config.add_bias_linear = False
+ config.bias_activation_fusion = False
+ config.gated_linear_unit = True
+ config.apply_query_key_layer_scaling = False
+ config.layernorm_zero_centered_gamma = (
+ False # Zero centered gamma not supported for RMSNorm
+ )
+ config.bias_dropout_fusion = False
+ config.apply_rope_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.ffn_hidden_size = 14336
+ elif config.language_model_type == "llama3.1_8b":
+ config.activation_func = torch.nn.functional.silu
+ config.add_bias_linear = False
+ config.bias_activation_fusion = False
+ config.gated_linear_unit = True
+ config.apply_query_key_layer_scaling = False
+ config.layernorm_zero_centered_gamma = (
+ False # Zero centered gamma not supported for RMSNorm
+ )
+ config.bias_dropout_fusion = False
+ config.apply_rope_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.ffn_hidden_size = 14336
+ elif config.language_model_type == "llama3.1_70B":
+ config.activation_func = torch.nn.functional.silu
+ config.add_bias_linear = False
+ config.bias_activation_fusion = False
+ config.gated_linear_unit = True
+ config.apply_query_key_layer_scaling = False
+ config.layernorm_zero_centered_gamma = (
+ False # Zero centered gamma not supported for RMSNorm
+ )
+ config.bias_dropout_fusion = False
+ config.apply_rope_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.ffn_hidden_size = 28672
+ elif config.language_model_type == "mistral_7b":
+ config.activation_func = torch.nn.functional.silu
+ config.add_bias_linear = False
+ config.bias_activation_fusion = False
+ config.gated_linear_unit = True
+ config.apply_query_key_layer_scaling = False
+ config.layernorm_zero_centered_gamma = (
+ False # Zero centered gamma not supported for RMSNorm
+ )
+ config.bias_dropout_fusion = False
+ config.apply_rope_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.ffn_hidden_size = 14336
+ elif config.language_model_type == "yi-34b":
+ config.activation_func = torch.nn.functional.silu
+ config.add_bias_linear = False
+ config.bias_activation_fusion = False
+ config.gated_linear_unit = True
+ config.apply_query_key_layer_scaling = False
+ config.layernorm_zero_centered_gamma = (
+ False # Zero centered gamma not supported for RMSNorm
+ )
+ config.bias_dropout_fusion = False
+ config.apply_rope_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.ffn_hidden_size = 20480
+ elif config.language_model_type == "qwen2.5_7B":
+ config.activation_func = torch.nn.functional.silu
+ config.add_bias_linear = False
+ config.add_qkv_bias = True
+ config.bias_activation_fusion = False
+ config.gated_linear_unit = True
+ config.apply_query_key_layer_scaling = False
+ config.layernorm_zero_centered_gamma = (
+ False # Zero centered gamma not supported for RMSNorm
+ )
+ config.bias_dropout_fusion = False
+ config.apply_rope_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.ffn_hidden_size = 18944
+ elif config.language_model_type == "qwen2.0_72B":
+ config.activation_func = torch.nn.functional.silu
+ config.add_bias_linear = False
+ config.add_qkv_bias = True
+ config.bias_activation_fusion = False
+ config.gated_linear_unit = True
+ config.apply_query_key_layer_scaling = False
+ config.layernorm_zero_centered_gamma = (
+ False # Zero centered gamma not supported for RMSNorm
+ )
+ config.bias_dropout_fusion = False
+ config.apply_rope_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.ffn_hidden_size = 29568
+ elif config.language_model_type == "llama3.2_1b":
+ config.activation_func = torch.nn.functional.silu
+ config.add_bias_linear = False
+ config.bias_activation_fusion = False
+ config.gated_linear_unit = True
+ config.apply_query_key_layer_scaling = False
+ config.layernorm_zero_centered_gamma = (
+ False # Zero centered gamma not supported for RMSNorm
+ )
+ config.bias_dropout_fusion = False
+ config.apply_rope_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.ffn_hidden_size = 8192
+ elif config.language_model_type.startswith("huggingface"):
+ # Loaded from HuggingFace config file.
+ pass
+ else:
+ raise ValueError(f"unknown language model type {config.language_model_type}")
+
+ return config
+
+
+def get_vision_model_config(config, apply_query_key_layer_scaling):
+ if config.vision_model_type == "clip":
+ config.num_layers = 24
+ config.num_attention_heads = 16
+ config.add_bias_linear = True
+ config.add_qkv_bias = True
+ config.hidden_size = 1024
+ config.hidden_dropout = 0.0
+ config.attention_dropout = 0.0
+ config.ffn_hidden_size = 4096
+ config.gated_linear_unit = False
+ config.activation_func = quick_gelu
+ config.kv_channels = 64
+ config.num_query_groups = 16
+ config.layernorm_zero_centered_gamma = False
+ config.apply_query_key_layer_scaling = apply_query_key_layer_scaling
+ config.bias_activation_fusion = False
+ config.bias_dropout_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.normalization = 'LayerNorm'
+ config.apply_rope_fusion = False
+ elif config.vision_model_type == "siglip":
+ config.num_layers = 27
+ config.num_attention_heads = 16
+ config.add_bias_linear = True
+ config.add_qkv_bias = True
+ config.hidden_size = 1152
+ config.hidden_dropout = 0.0
+ config.attention_dropout = 0.0
+ config.ffn_hidden_size = 4304
+ config.gated_linear_unit = False
+ config.activation_func = fast_gelu
+ config.kv_channels = 72
+ config.num_query_groups = 16
+ config.layernorm_zero_centered_gamma = False
+ config.apply_query_key_layer_scaling = apply_query_key_layer_scaling
+ config.bias_activation_fusion = False
+ config.bias_dropout_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.normalization = 'LayerNorm'
+ config.apply_rope_fusion = False
+ config.qk_layernorm = False
+ config.layernorm_epsilon = 1e-6
+ elif config.vision_model_type == "internvit":
+ config.num_layers = 45
+ config.num_attention_heads = ((24 // config.tensor_model_parallel_size) + 1) * config.tensor_model_parallel_size
+ config.num_query_groups = config.num_attention_heads
+ config.add_bias_linear = True
+ config.add_qkv_bias = False
+ config.hidden_size = 3200
+ config.hidden_dropout = 0.0
+ config.attention_dropout = 0.0
+ config.ffn_hidden_size = 12800
+ config.gated_linear_unit = False
+ config.activation_func = torch.nn.functional.gelu
+ config.layernorm_zero_centered_gamma = False
+ config.apply_query_key_layer_scaling = apply_query_key_layer_scaling
+ config.bias_activation_fusion = False
+ config.bias_dropout_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.normalization = 'RMSNorm'
+ config.layernorm_epsilon = 1e-6
+ config.apply_rope_fusion = False
+ elif config.vision_model_type == "radio":
+ config.num_layers = 32
+ config.num_attention_heads = 16
+ config.add_bias_linear = True
+ config.add_qkv_bias = True
+ config.hidden_size = 1280
+ config.ffn_hidden_size = 5120
+ config.gated_linear_unit = False
+ config.activation_func = fast_gelu
+ config.kv_channels = 80
+ config.num_query_groups = 16
+ config.layernorm_zero_centered_gamma = False
+ config.apply_query_key_layer_scaling = apply_query_key_layer_scaling
+ config.bias_activation_fusion = False
+ config.bias_dropout_fusion = False
+ config.attention_softmax_in_fp32 = True
+ config.normalization = 'LayerNorm'
+ config.apply_rope_fusion = False
+ config.qk_layernorm = False
+ config.layernorm_epsilon = 1e-6
+ elif config.vision_model_type.startswith("huggingface"):
+ # Loaded from HuggingFace config file.
+ pass
+ else:
+ raise ValueError(f"unknown vision model type {config.vision_model_type}")
+
+ return config
+
+
+def get_vision_projection_config(config, hidden_size):
+ config.gated_linear_unit = False
+ config.bias_activation_fusion = False
+ config.add_bias_linear = False
+ config.hidden_size = hidden_size # Used as the vision projection output size, i.e., the input to the language model.
+ if config.language_model_type == "llama3_8b":
+ config.ffn_hidden_size = 14336
+ config.activation_func = torch.nn.functional.gelu
+ elif config.language_model_type == "llama3.1_8b":
+ config.ffn_hidden_size = 4096
+ config.activation_func = torch.nn.functional.gelu
+ config.layernorm_epsilon = 1e-5
+ config.add_bias_linear = True
+ config.normalization = "LayerNorm"
+ elif config.language_model_type == "mistral_7b":
+ config.ffn_hidden_size = 14336
+ config.activation_func = torch.nn.functional.gelu
+ config.normalization = None
+ elif config.language_model_type == "yi-34b":
+ config.ffn_hidden_size = 20480
+ config.normalization = "LayerNorm"
+ config.activation_func = torch.nn.functional.gelu
+ elif config.language_model_type == "qwen2.5_7B":
+ config.ffn_hidden_size = 3584
+ config.activation_func = torch.nn.functional.gelu
+ elif config.language_model_type == "qwen2.0_72B":
+ config.ffn_hidden_size = 29568
+ config.normalization = "LayerNorm"
+ config.activation_func = torch.nn.functional.gelu
+ elif config.language_model_type == "llama3.2_1b":
+ config.ffn_hidden_size = 2048
+ config.activation_func = torch.nn.functional.gelu
+ config.normalization = "LayerNorm"
+ elif config.language_model_type.startswith("huggingface"):
+ config.activation_func = torch.nn.functional.gelu
+ from transformers import AutoConfig
+ hf_config = AutoConfig.from_pretrained(config.huggingface_model_name_or_path)
+ if "qwen" in hf_config.model_type:
+ config.ffn_hidden_size = 1536
+ else:
+ raise ValueError(f"unknown language model type {config.language_model_type}")
+
+ return config
+
+
+@dataclass
+class EvaluationConfig:
+ """Evaluation related configuration."""
+ task: str
+
+ temperature: float = 1.0
+ top_p: float = 0.0
+ top_k: int = 0
+
+ out_seq_length: int = 32
+
+ output_path: str = ""
+
+ input_image_path: str = ""
+ gt_path: str = ""
+
+ num_partitions: int = 1
+ partition_id: int = 0
+ num_samples_per_partition: int = 0
diff --git a/examples/multimodal/convert_llava_pretrain_to_wds.py b/examples/multimodal/convert_llava_pretrain_to_wds.py
new file mode 100644
index 0000000000..0092aef246
--- /dev/null
+++ b/examples/multimodal/convert_llava_pretrain_to_wds.py
@@ -0,0 +1,31 @@
+import json
+import os
+import webdataset as wds
+
+from tqdm import tqdm
+
+llava_pretrain_dir = ''
+
+# Paths to the dataset files
+json_file = os.path.join(llava_pretrain_dir, 'blip_laion_cc_sbu_558k.json')
+output = os.path.join(llava_pretrain_dir, 'wds')
+
+if not os.path.exists(output):
+ os.mkdir(output)
+
+# Load data
+with open(json_file, 'r') as f:
+ data = json.load(f)
+
+with wds.ShardWriter(os.path.join(output, 'pretrain-%d.tar'), maxcount=10000) as shard_writer:
+ for entry in tqdm(data):
+ with open(os.path.join(llava_pretrain_dir, entry['image']), "rb") as img_file:
+ image_data = img_file.read()
+ sample = {
+ "__key__": entry['id'],
+ "jpg": image_data,
+ "json": json.dumps(entry['conversations']).encode("utf-8"),
+ }
+ shard_writer.write(sample)
+
+print(f"Dataset successfully converted to wds")
diff --git a/examples/multimodal/dataloader_provider.py b/examples/multimodal/dataloader_provider.py
new file mode 100644
index 0000000000..aef2186834
--- /dev/null
+++ b/examples/multimodal/dataloader_provider.py
@@ -0,0 +1,169 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import os
+
+import torch
+from dataset_helpers import TaskEncoder, print_error_handler
+
+from megatron.core import parallel_state
+from megatron.energon import (
+ LimitDataset,
+ RepeatDataset,
+ WorkerConfig,
+ get_loader,
+ get_savable_loader,
+ get_train_dataset,
+ get_val_datasets,
+)
+from megatron.core.num_microbatches_calculator import get_num_microbatches
+from megatron.core.parallel_state import get_tensor_model_parallel_rank, get_pipeline_model_parallel_world_size, get_pipeline_model_parallel_rank
+from megatron.training import get_args
+from megatron.training.checkpointing import get_checkpoint_name
+
+
+def datasets_provider(worker_config=None):
+ """Create multimodal train, validation and test datasets."""
+ args = get_args()
+
+ dname = args.data_path[0] if type(args.data_path) is list else args.data_path
+ train_dataset = get_train_dataset(
+ dname,
+ batch_size=args.micro_batch_size,
+ task_encoder=TaskEncoder(),
+ worker_config=worker_config,
+ max_samples_per_sequence=None,
+ shuffle_buffer_size=None,
+ packing_buffer_size=args.packing_buffer_size,
+ handler=print_error_handler,
+ image_decode="pil",
+ )
+
+ val_datasets = get_val_datasets(
+ dname,
+ batch_size=args.micro_batch_size,
+ # This is the total number over all workers
+ # limit=args.eval_iters * get_num_microbatches(),
+ task_encoder=TaskEncoder(),
+ worker_config=worker_config,
+ packing_buffer_size=args.packing_buffer_size,
+ handler=print_error_handler,
+ image_decode="pil",
+ )
+ val_datasets_without_source_datasets = [
+ # Limit the dataset to eval_iters * num_microbatches
+ LimitDataset(
+ # Repeat the inner dataset in case it's too short
+ RepeatDataset(val_ds, worker_config=worker_config),
+ length=args.eval_iters * get_num_microbatches(),
+ worker_config=worker_config,
+ reset_after_epoch=True,
+ )
+ for val_ds, _src_ds in val_datasets
+ ]
+
+ return train_dataset, val_datasets_without_source_datasets, None
+
+
+def is_first_or_last_stage(pp_size, encoder_pipeline_model_parallel_size):
+ """Check if the current pipeline parallel stage is the first or last stage."""
+ if pp_size == 1: # No pipeline parallelism.
+ return True
+
+ is_valid_rank = False
+ pp_rank = get_pipeline_model_parallel_rank()
+ if encoder_pipeline_model_parallel_size == 0:
+ # No separate pipeline stage for the vision model. Run the dataloader on the first and last pipeline stage.
+ is_valid_rank = pp_rank in (0, pp_size-1)
+ elif encoder_pipeline_model_parallel_size == 1:
+ # Separate pipeline stage for the vision model. Run the dataloader on the first vision and LM stage and last LM stage.
+ is_valid_rank = pp_rank in (0, 1, pp_size-1)
+ else:
+ raise NotImplementedError("encoder-pipeline-model-parallel-size > 1 is not supported yet")
+
+ return is_valid_rank
+
+
+def is_dataloader_rank(encoder_pipeline_model_parallel_size):
+ """Check if we should have the dataloader on this tensor and pipeline parallel rank."""
+ # Run dataloader only on the first tensor parallel rank (will be broadcasted to others).
+ is_first_rank = get_tensor_model_parallel_rank() == 0
+
+ pp_size = get_pipeline_model_parallel_world_size()
+ is_first_rank = is_first_rank and is_first_or_last_stage(pp_size, encoder_pipeline_model_parallel_size)
+
+ return is_first_rank
+
+
+def train_valid_test_dataloaders_provider(train_val_test_num_samples):
+ """Build multimodal train, validation and test dataloaders."""
+ args = get_args()
+
+ # Dataloader is only on specific ranks.
+ if not is_dataloader_rank(args.encoder_pipeline_model_parallel_size):
+ return None, None, None
+
+ worker_debug_path = None
+ worker_log_level = 0
+
+ rank = parallel_state.get_data_parallel_rank()
+ world_size = parallel_state.get_data_parallel_world_size()
+ data_parallel_group = parallel_state.get_data_parallel_group()
+
+ worker_config = WorkerConfig(
+ rank=rank,
+ world_size=world_size,
+ num_workers=args.num_workers,
+ data_parallel_group=data_parallel_group,
+ worker_debug_path=worker_debug_path,
+ worker_log_level=worker_log_level,
+ )
+ train_ds, valid_ds1, test_ds = datasets_provider(worker_config)
+
+ train_dataloader = get_savable_loader(train_ds, worker_config=worker_config)
+ if args.load is not None:
+ if getattr(args, "dataloader_save", None):
+ dp_rank = parallel_state.get_data_parallel_rank()
+ data_save_name = get_checkpoint_name(
+ args.dataloader_save,
+ args.iteration,
+ pipeline_rank=0, # Only the first pipeline parallel rank stores the dataloader checkpoint.
+ basename=f"train_dataloader_dprank{dp_rank:03d}.pt",
+ )
+ if os.path.exists(data_save_name):
+ try:
+ dataset_state_dict = torch.load(data_save_name, map_location="cpu")
+ train_dataloader.restore_state_rank(dataset_state_dict["dataloader_state_dict"])
+ print(f"restored dataset state from {data_save_name}")
+ except Exception as e:
+ print("loading dataset state failed. Skipping. " + str(e))
+ else:
+ print(f"dataset state {data_save_name} does not exist")
+
+ valid_dataloader = [
+ EnergonDataloader(get_loader(valid_ds, worker_config=worker_config))
+ for valid_ds in valid_ds1
+ ]
+ test_dataloader = None
+
+ return EnergonDataloader(train_dataloader), valid_dataloader, EnergonDataloader(test_dataloader)
+
+
+class EnergonDataloader:
+ """A wrapper to use Megatron Energon dataloader with the Megatron-LM training loop."""
+ def __init__(self, dataloader):
+ self._dataloader = dataloader
+ self._iter = iter(cyclic_iter(dataloader))
+
+ def __next__(self):
+ return self._iter.__next__()
+
+ def __iter__(self):
+ return self._iter.__iter__()
+
+ def save_state(self):
+ return self._dataloader.save_state_rank()
+
+
+def cyclic_iter(iter):
+ while True:
+ for x in iter:
+ yield x
diff --git a/examples/multimodal/dataset_helpers.py b/examples/multimodal/dataset_helpers.py
new file mode 100644
index 0000000000..39fa44b338
--- /dev/null
+++ b/examples/multimodal/dataset_helpers.py
@@ -0,0 +1,906 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import bisect
+import dataclasses
+import json
+import re
+import sys
+import traceback
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union
+
+from image_processing import find_closest_aspect_ratio, find_closest_area_weighted_aspect_ratio, get_visual_transform
+from PIL import Image
+from torchvision.transforms import ToPILImage
+import numpy as np
+import torch
+
+from energon_util import OfflineTargetAspectRatioSample, SampleListSample
+from megatron.core.models.multimodal.llava_model import IGNORE_INDEX, IMAGE_TOKEN, VIDEO_TOKEN
+from megatron.core.models.vision.clip_vit_model import get_num_image_embeddings
+from megatron.energon import (
+ Batch,
+ CaptioningSample,
+ DefaultTaskEncoder,
+ OCRSample,
+ Sample,
+ SimilarityInterleavedSample,
+ VQASample,
+ MultiChoiceVQASample
+)
+from megatron.energon.task_encoder.base import stateless
+from megatron.training import get_args, get_tokenizer
+
+
+@dataclass
+class ImageTaskSample(Sample):
+ __key__: str
+ __restore_key__: Tuple[Union[str, int, tuple], ...]
+ __subflavor__: Dict
+ __subflavors__: Dict
+ # (c, h, w)
+ imgs: List[torch.Tensor]
+ num_tiles: List[int]
+ tokens: torch.Tensor
+ total_len: int # Total token count in the sample, including text and image tokens
+ labels: torch.Tensor = None
+
+
+@dataclass
+class ImageTaskSamplePacked(Sample):
+ """Dataclass to store a single packed sample (not a batch).
+
+ P = Number of sub-samples in the packed sample
+ seq_len = Total sequence length
+ num_imgs = Number of images across all samples in the packed sample
+ """
+
+ __key__: str # Sample name
+ __restore_key__: Tuple[Union[str, int, tuple], ...]
+ __subflavor__: Dict # Sample metadata. Deprecated.
+ __subflavors__: Dict # Sample metadata.
+ tokens: torch.Tensor # Input tokens packed into a single tensor (seq_len,)
+ labels: torch.Tensor # Target tokens packed into a single tensor (seq_len,)
+ imgs: List[torch.Tensor] # Input images
+ num_tiles: List[int] # Number of tiles for each image of each sample (num_imgs)
+ max_length: int # Maximum length across sub-samples.
+ cu_lengths: List[int] # Cumulative length of each sub-sample in this packed sample incl. text and image tokens (P,)
+
+
+# Typing for the resulting batch data after encode_batch()
+@dataclass
+class ImageTaskBatchPacked(Batch):
+ """Dataclass to store a batch of packed samples.
+
+ N = Batch size
+ P = Number of samples in the packed sample
+ seq_len = Maximum sequence length
+ num_imgs = Number of images across all samples in the packed sample
+ """
+
+ __key__: List[str] # Sample names
+ __restore_key__: Tuple[Union[str, int, tuple], ...]
+ __subflavor__: Dict # Sample metadata. Deprecated.
+ __subflavors__: List[Dict] # Sample metadatas.
+ tokens: torch.Tensor # Input tokens packed and padded (N, seq_len)
+ labels: torch.Tensor # Target tokens packed and padded (N, seq_len)
+ imgs: torch.Tensor # All image tiles stacked into a single tensor (num_tiles, C, H, W)
+ num_tiles: List[List[int]] # Number of tiles per image (N, num_imgs)
+ max_lengths: List[int] # Maximum length across sub-samples (N,)
+ cu_lengths: List[List[int]] # Cumulative length of each sub-sample in each packed sample of the batch (N, P)
+
+
+# Based on https://github.com/hiyouga/LLaMA-Factory/blob/641d0dab08d96a93c34657742213d8994d9ed476/src/llamafactory/data/processors/processor_utils.py#L19
+# Copyright (c) 2024 LLaMA-Factory. Apache license 2.0.
+def search_for_fit(numbers: List[int], capacity: int) -> int:
+ """Finds the index of largest number that fits into the knapsack with the given capacity."""
+ index = bisect.bisect(numbers, capacity)
+ return -1 if index == 0 else (index - 1)
+
+
+# Based on https://github.com/hiyouga/LLaMA-Factory/blob/641d0dab08d96a93c34657742213d8994d9ed476/src/llamafactory/data/processors/processor_utils.py#L27
+# Copyright (c) 2024 LLaMA-Factory. Apache license 2.0.
+def greedy_knapsack(item_sizes: List[int], samples: List, max_capacity: int) -> List:
+ """Greedy algorithm with binary search for the knapsack problem.
+
+ Pack as many samples as possible given a maximum capacity and capacities of individual samples.
+ Used if sequence packing is enabled.
+ """
+ assert len(item_sizes) == len(samples), "sample lengths and samples must have the same length."
+
+ knapsacks = []
+
+ if len(item_sizes) == 0:
+ return knapsacks
+
+ # Sort sample lengths and samples together.
+ sorted_item_sizes, sorted_samples = zip(*sorted(zip(item_sizes, samples), key=lambda x: x[0]))
+ sorted_item_sizes = list(sorted_item_sizes)
+ sorted_samples = list(sorted_samples)
+
+ # Check if all samples fit in the knapsack capacity.
+ if sorted_item_sizes[-1] > max_capacity:
+ raise ValueError(f"knapsack: A sample is larger {sorted_item_sizes[-1]} than the max_sequence_length {max_capacity}.")
+
+ while sorted_item_sizes:
+ current_knapsack = []
+ remaining_capacity = max_capacity
+
+ while True:
+ idx = search_for_fit(sorted_item_sizes, remaining_capacity)
+ if idx == -1:
+ break # Can't fit more samples.
+
+ remaining_capacity -= sorted_item_sizes[idx]
+
+ sorted_item_sizes.pop(idx)
+ sample = sorted_samples.pop(idx)
+ current_knapsack.append(sample)
+
+ knapsacks.append(current_knapsack)
+
+ return knapsacks
+
+
+class TaskEncoder(DefaultTaskEncoder[OCRSample, OCRSample, ImageTaskBatchPacked, dict]):
+ """A simple task encoder for VLMs."""
+
+ def __init__(
+ self
+ ):
+ super().__init__()
+
+ self.args = get_args()
+
+ self.tokenizer = get_tokenizer()
+ with open(self.args.prompt_path, "r") as f:
+ self.manual_prompts = json.load(f)
+ self.dataloader_seq_length = self.args.dataloader_seq_length # Always return samples of this length.
+ self.packing_seq_length = self.args.packing_seq_length # Packing sequence length, if packing is enabled.
+ self.is_packing_enabled = self.args.packing_buffer_size is not None and self.args.packing_buffer_size > 0
+
+ if self.dataloader_seq_length and self.packing_seq_length:
+ assert self.dataloader_seq_length >= self.packing_seq_length, "dataloader sequence length must be greater than or equal to the packing sequence length"
+
+ if self.is_packing_enabled:
+ assert self.packing_seq_length > 0, "packing sequence length must be set"
+
+ self.num_image_embeddings_per_tile = get_num_image_embeddings(
+ self.args.img_h,
+ self.args.img_w,
+ self.args.patch_dim,
+ self.args.vision_model_type,
+ self.args.disable_vision_class_token,
+ 1,
+ self.args.pixel_shuffle,
+ self.args.use_tile_tags,
+ )
+
+ self.txt_to_token_dict = {}
+
+ self.img_h, self.img_w = self.args.img_h, self.args.img_w
+ self.img_token_id = self.tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
+ # This map is used to reduce the number of tiles used per image if the number of tokens is
+ # larger than the decoder_seq_length.
+ self.num_tiles_degradation_map = {12:8, 8:6, 6:4, 4:2, 2:1, 1:1}
+
+ self.find_closest_aspect_ratio_fn = (
+ find_closest_area_weighted_aspect_ratio if self.args.use_area_weighted_aspect_ratio
+ else find_closest_aspect_ratio)
+
+ def _get_total_seq_length(self, input_ids, num_tiles):
+ """Calculate expected sequence length given text tokens length and number of tiles."""
+ total_num_images = len(num_tiles)
+ total_num_tiles = sum(num_tiles)
+ total_len = len(input_ids) + total_num_tiles * self.num_image_embeddings_per_tile - total_num_images
+ return total_len
+
+ def _truncate_for_packing(self, input_ids, target, num_tiles):
+ """Truncate tokens and labels if they exceed packing sequence length."""
+ total_num_images = len(num_tiles)
+ total_num_tiles = sum(num_tiles)
+ total_img_embeddings_len = total_num_tiles * self.num_image_embeddings_per_tile
+ max_text_tokens = self.packing_seq_length - total_img_embeddings_len + total_num_images
+
+ input_ids = input_ids[:max_text_tokens]
+ target = target[:max_text_tokens]
+
+ # If truncate causes all labels to be ignored, then skip the sample
+ if (target == IGNORE_INDEX).all():
+ raise ValueError(f"all targets will be ignored after truncation: {input_ids}")
+
+ return input_ids, target
+
+ @stateless(restore_seeds=True)
+ def encode_sample(self, sample: Union[CaptioningSample, OCRSample, VQASample, SimilarityInterleavedSample]):
+ if isinstance(sample, OCRSample):
+ if "pdfa" in sample.__key__:
+ yield self.combined_ocr_encoder(sample, task_type='encode_pdf')
+ elif "multi" in sample.__key__:
+ yield self.combined_ocr_encoder(sample, task_type='_encode_ocr')
+ else:
+ yield self.combined_ocr_encoder(sample, task_type='encode_ocr_ref')
+ elif isinstance(sample, CaptioningSample):
+ yield self.encode_captioning(sample)
+ elif isinstance(sample, VQASample):
+ is_llava_training = sample.__subflavors__["is_llava_training"] if "is_llava_training" in sample.__subflavors__ else False
+
+ if "llava" in sample.__key__ or is_llava_training:
+ yield self.encode_llava_pretrain(sample)
+ else:
+ yield self.encode_any_single_turn_vqa(sample)
+ elif isinstance(sample, SimilarityInterleavedSample):
+ yield self.encode_llava_sft(sample)
+ elif isinstance(sample, MultiChoiceVQASample):
+ yield self.encode_any_single_turn_vqa(sample)
+ # Because the SampleListSample is defined in the Megatron module but loaded by the Energon
+ # library, we need to resort to the more brittle check:
+ elif type(sample).__name__ == "SampleListSample":
+ yield self.encode_sample_list(sample)
+ else:
+ raise NotImplementedError("Sample format not supported", sample)
+
+ def encode_captioning(self, sample: CaptioningSample):
+ """Encode CaptioningSample."""
+ augment = sample.__subflavors__.get("augmentation")
+
+ imgs = get_visual_transform(
+ sample.image, self.img_h, self.img_w, self.args.use_tiling, self.args.max_num_tiles, self.args.use_thumbnail, augment,
+ self.args.vision_model_type, find_closest_aspect_ratio_fn=self.find_closest_aspect_ratio_fn
+ )
+ num_tiles = [len(imgs)]
+
+ prompt_list = self.manual_prompts["CaptioningPretraining"]["raw"]
+
+ prompt_idx = np.random.randint(len(prompt_list))
+ cur_prompt = prompt_list[prompt_idx]
+ cur_prompt = IMAGE_TOKEN + "\n" + cur_prompt + "\n"
+
+ caption = sample.caption.strip()
+
+ split_by_line_flag = sample.__subflavors__.get("SplitByLine")
+ if split_by_line_flag:
+ caption_list = caption.split('\n')
+ caption = np.random.choice(caption_list)
+
+ conv = [
+ # Note: no system message.
+ {"role": "user", "content": cur_prompt},
+ {"role": "assistant", "content": caption},
+ ]
+
+ input_ids, target = self.tokenizer.tokenize_conversation(conv, True, False)
+
+ if self.is_packing_enabled:
+ input_ids, target = self._truncate_for_packing(input_ids, target, num_tiles)
+
+ return ImageTaskSample(
+ __key__=sample.__key__,
+ __restore_key__=sample.__restore_key__,
+ __subflavor__=None,
+ __subflavors__=sample.__subflavors__,
+ imgs=imgs,
+ num_tiles=num_tiles,
+ tokens=torch.tensor(input_ids),
+ labels=torch.tensor(target),
+ total_len=self._get_total_seq_length(input_ids, num_tiles),
+ )
+
+ def encode_llava_pretrain(self, sample: VQASample):
+ """Encode pretrain sample in LLAVA style."""
+ augment = sample.__subflavors__.get("augmentation", False)
+
+ imgs = get_visual_transform(
+ sample.image, self.img_h, self.img_w, self.args.use_tiling, self.args.max_num_tiles, self.args.use_thumbnail, augment,
+ self.args.vision_model_type, find_closest_aspect_ratio_fn=self.find_closest_aspect_ratio_fn
+ )
+ num_tiles = [len(imgs)]
+
+ # LLAVA training: override text-prompt with just the image.
+ conv = [
+ # Note: no system message.
+ {"role": "user", "content": IMAGE_TOKEN + "\n"},
+ {"role": "assistant", "content": sample.answers},
+ ]
+
+ input_ids, target = self.tokenizer.tokenize_conversation(conv, True, False)
+
+ if self.is_packing_enabled:
+ input_ids, target = self._truncate_for_packing(input_ids, target, num_tiles)
+
+ return ImageTaskSample(
+ __key__=sample.__key__,
+ __restore_key__=sample.__restore_key__,
+ __subflavor__=None,
+ __subflavors__=sample.__subflavors__,
+ imgs=imgs,
+ num_tiles=num_tiles,
+ tokens=torch.tensor(input_ids),
+ labels=torch.tensor(target),
+ total_len=self._get_total_seq_length(input_ids, num_tiles),
+ )
+
+ def encode_sample_list(self, samples: SampleListSample):
+ """We encode the list of samples using encode_llava_sft on each sample."""
+ error_msg = ("You probably don't want to use online packing since SampleListSample is "
+ "usually used along offline packing.")
+ assert not self.is_packing_enabled, error_msg
+ encoded_samples = []
+ current_length = 0
+ for sample in samples.samples:
+ encoded_sample = self.encode_llava_sft(sample, truncate_for_sample_list_packing=True)
+ if current_length + encoded_sample.total_len > self.packing_seq_length:
+ break
+ else:
+ encoded_samples.append(encoded_sample)
+ current_length += encoded_sample.total_len
+ return self.pack_selected_samples(encoded_samples)
+
+ def encode_llava_sft(self, sample: Union[SimilarityInterleavedSample, OfflineTargetAspectRatioSample], truncate_for_sample_list_packing=False):
+ """Encode SFT sample."""
+ augment = sample.__subflavors__['augmentation'] if 'augmentation' in sample.__subflavors__ else False
+ has_video = sample.__subflavors__['has_video'] if 'has_video' in sample.__subflavors__ else False
+
+ # If the target aspect ratio are provided by the dataset, we use them instead of computing
+ # them with the self.find_closest_aspect_ratio_fn function.
+ local_find_closest_aspect_ratio_fn = self.find_closest_aspect_ratio_fn
+ if type(sample).__name__ == "OfflineTargetAspectRatioSample":
+ target_aspect_ratio = tuple(sample.target_aspect_ratio[0])
+ assert target_aspect_ratio is not None, "Sample of type OfflineTargetAspectRatioSample needs to define the target aspect ratio."
+ local_find_closest_aspect_ratio_fn = lambda *args, **kwargs: target_aspect_ratio
+
+ has_image = False
+ # We infer whether the sample has image or not.
+ if hasattr(sample, "images") and not has_video:
+ # If this is a text-only sample and we are freezing the LM,
+ # then use a dummy input image.
+ if len(sample.images) == 0 and self.args.freeze_LM:
+ empty_img = Image.new('RGB', (self.args.img_w, self.args.img_h), (255, 255, 255))
+ sample.images.append(empty_img)
+ if len(sample.images) > 0:
+ has_image = True
+
+ # Note: Some tokenizers may ignore the system prompt.
+ conversation = [{"role": "system", "content": "Answer the questions."}]
+ # Format the conversation as a list of "user" / "assistant" turns.
+ for text in sample.texts:
+ error_msg = f"unexpected role {text['from']} in {sample.texts}"
+ assert text["from"] in ["human", "gpt"], error_msg
+ conversation.append({
+ "role": "user" if text["from"] == "human" else "assistant",
+ "content": text["value"]})
+
+ # Replace the image tags with IMAGE_TOKEN and count the number of image tags
+ number_image_tags = 0
+ image_tag_ids_list = []
+ for turn in conversation:
+ if turn["role"] == "user":
+ image_tag_ids = [int(x) - 1 for x in re.findall(r"", turn["content"])]
+ image_tag_ids_list.extend(image_tag_ids)
+ turn["content"] = re.sub(r"", IMAGE_TOKEN, turn["content"])
+ # For videos, we use the image token to locate where to put the frames.
+ if has_video:
+ turn["content"] = turn["content"].replace(VIDEO_TOKEN, IMAGE_TOKEN)
+ number_image_tags += turn["content"].count(IMAGE_TOKEN)
+
+ # We re-order the images in sample.images according to how they appear in the conversation.
+ if len(image_tag_ids_list) > 0:
+ sample.images = [sample.images[idx] for idx in image_tag_ids_list]
+
+ # If there is only one image, but several image tags, we assume all the tags refer to the
+ # same image and duplicate the image:
+ if not has_video and len(sample.images) == 1 and number_image_tags > 1:
+ sample.images = sample.images * number_image_tags
+
+ # We currently only support one video per sample.
+ number_of_images = 1 if has_video else len(sample.images)
+ # Fail if there are more image or video tags than image or videos:
+ error_msg = (
+ f"Found {number_image_tags} image tags for {number_of_images} images. {sample.texts}")
+ assert number_image_tags <= number_of_images, error_msg
+
+ # If there are less image of video tags than image or videos, prepend the tags to the first
+ # user message:
+ if number_image_tags < number_of_images:
+ for turn in conversation:
+ if turn["role"] == "user":
+ turn["content"] = IMAGE_TOKEN*(number_of_images-number_image_tags) + "\n" + turn["content"]
+ break
+
+ input_ids, target = self.tokenizer.tokenize_conversation(conversation, True, False)
+
+ if has_image:
+ imgs = []
+ num_tiles = []
+ max_num_tiles = self.args.max_num_tiles
+ # We keep a buffer of 4 tokens for the question,
+ # the rest can be used for image tokens.
+ max_image_token_allowed = self.args.decoder_seq_length - len(input_ids) - 4
+ # We start by extracting as many tiles per image as possible, and decrease the max
+ # number of tiles if there are too many image tokens.
+ while True:
+ imgs = []
+ num_tiles = []
+ for img in sample.images:
+ img_tiles = get_visual_transform(
+ img, self.img_h, self.img_w, self.args.use_tiling, max_num_tiles,
+ self.args.use_thumbnail, augment, self.args.vision_model_type,
+ find_closest_aspect_ratio_fn=local_find_closest_aspect_ratio_fn)
+ imgs += img_tiles
+ num_tiles += [len(img_tiles)]
+ if max_num_tiles == 1:
+ break
+ if sum(num_tiles) * self.num_image_embeddings_per_tile > max_image_token_allowed:
+ if max_num_tiles in self.num_tiles_degradation_map:
+ max_num_tiles = self.num_tiles_degradation_map[max_num_tiles]
+ else:
+ raise RuntimeError((
+ f"Tried to decrease the number of tiles {max_num_tiles} but it's not ",
+ f"defined in the degradation map {self.num_tiles_degradation_map}"))
+ else:
+ break
+ elif has_video:
+ # We don't use tiling for videos to limit the number of tokens.
+ use_tiling=False
+ # Grab the selected frames of the video as a tensor with shape
+ # fhwc: (num_frames, num_channels, height, width).
+ video_fchw = sample.images.frames
+ if video_fchw.shape[0] == 0:
+ raise ValueError(f"Video {sample.__key__} {sample.__restore_key__} {sample.texts} has no frames.")
+ selected_frames = torch.linspace(
+ 0, video_fchw.shape[0] - 1, self.args.num_frames).long()
+ video_fchw = video_fchw[selected_frames]
+ imgs = []
+ for video_chw in video_fchw:
+ to_pil = ToPILImage()
+ video_chw = to_pil(video_chw)
+ imgs += get_visual_transform(
+ video_chw, self.img_h, self.img_w, use_tiling, self.args.max_num_tiles,
+ self.args.use_thumbnail, augment, self.args.vision_model_type,
+ find_closest_aspect_ratio_fn=local_find_closest_aspect_ratio_fn)
+ num_tiles = [len(imgs)]
+ else:
+ imgs = num_tiles = []
+
+ if self.is_packing_enabled or truncate_for_sample_list_packing:
+ input_ids, target = self._truncate_for_packing(input_ids, target, num_tiles)
+
+ # Some final checks with respect to the number of image tokens and images on the tokenized
+ # conversation. There can still be errors, for instance if a non-video sample happens to
+ # have our pre-defined video token, or if the packing truncation removed a necessary image
+ # tag.
+ number_image_token = np.sum(input_ids == self.img_token_id)
+ error_msg = (
+ f"Found {number_image_token} image tokens for len({num_tiles}) = {len(num_tiles)} image tiles in {conversation}.")
+ assert number_image_token == len(num_tiles), error_msg
+ error_msg = (
+ f"Found sum({num_tiles}) = {np.sum(num_tiles)} tiles for {len(imgs)} images in {conversation}.")
+ assert np.sum(num_tiles) == len(imgs), error_msg
+
+ # We need to ensure that there are at least some trainable tokens in the sample.
+ assert self.target_has_trainable_tokens(input_ids, num_tiles, target), "Sample has no trainable tokens."
+
+ return ImageTaskSample(
+ __key__=sample.__key__,
+ __restore_key__=sample.__restore_key__,
+ __subflavor__=None,
+ __subflavors__=sample.__subflavors__,
+ imgs=imgs,
+ num_tiles=num_tiles,
+ tokens=torch.tensor(input_ids),
+ labels=torch.tensor(target),
+ total_len=self._get_total_seq_length(input_ids, num_tiles),
+ )
+
+ def target_has_trainable_tokens(self, input_ids, num_tiles, target):
+ # Compute the loss mask based on extending the image tags with the proper
+ # number of image tokens, extracting the first self.args.decoder_seq_length tokens, and
+ # ensuring that some of these tokens have a loss mask > 0.
+ # Note that this is a bit hacky because we reproduce here parts of the logics which are in
+ # the model itself. Ideally, the data sampler would return the already processed inputs
+ # and targets to avoid this duplication.
+ expanded_target = target.copy()
+ expanded_target[input_ids==self.img_token_id] = self.img_token_id
+ expanded_target = self.replace_value_with_repetition(
+ expanded_target, self.img_token_id,
+ self.num_image_embeddings_per_tile * np.array(num_tiles), IGNORE_INDEX)
+ loss_mask = torch.ones(torch.tensor(expanded_target).size(), dtype=torch.float)
+ loss_mask[expanded_target == self.tokenizer.pad] = 0.0 # mask paddings
+ loss_mask[expanded_target == IGNORE_INDEX] = 0.0 # mask prompts
+ loss_mask = torch.cat((loss_mask[1:], torch.zeros((1,))))
+ loss_mask = loss_mask[:self.args.decoder_seq_length]
+ return torch.sum(loss_mask) > 0
+
+ def replace_value_with_repetition(self, arr, token_to_replace, num_repetition, new_token):
+ """
+ Replace every occurrence of value V in the input array with R repetitions of W.
+
+ Args:
+ arr (Array): Input array to be modified
+ token_to_replace: token to be replaced
+ new_token: new token
+ num_repetition (Array): number of repetition of new token.
+
+ Returns:
+ Array: New array with token_to_replace replaced by num_repetition repetitions of
+ new_token
+ """
+ error_msg = "The number of image tokens must match the length of the tile tensor."
+ assert np.sum(arr==token_to_replace) == len(num_repetition), error_msg
+ result = []
+ idx = 0
+ for item in arr:
+ if item == token_to_replace:
+ # If the current item matches token_to_replace, add R copies of W
+ result.extend([new_token] * num_repetition[idx])
+ idx += 1
+ else:
+ # Otherwise, keep the original item
+ result.append(item)
+
+ return np.array(result)
+
+ def encode_any_single_turn_vqa(self, sample):
+ """Encode MultiChoiceVQA or VQA sample."""
+ augment = sample.__subflavors__['augmentation'] if 'augmentation' in sample.__subflavors__ else False
+ has_video = sample.__subflavors__['has_video'] if 'has_video' in sample.__subflavors__ else False
+
+ if has_video:
+ # Grab the selected frames of the video as a tensor with shape
+ # fhwc: (num_frames, height, width, num_channels).
+ video_fhwc = sample.image.permute(0, 2, 3, 1)
+ selected_frames = torch.linspace(
+ 0, video_fhwc.shape[0] - 1, self.args.num_frames).long()
+ video_frame_fhwc = video_fhwc[selected_frames]
+ imgs = []
+ for video_frame_hwc in video_frame_fhwc:
+ imgs += get_visual_transform(
+ video_frame_hwc, self.img_h, self.img_w,
+ self.args.use_tiling, self.args.max_num_tiles,
+ self.args.use_thumbnail, augment, self.args.vision_model_type,
+ find_closest_aspect_ratio_fn=self.find_closest_aspect_ratio_fn)
+ else:
+ imgs = get_visual_transform(
+ sample.image, self.img_h, self.img_w, self.args.use_tiling, self.args.max_num_tiles,
+ self.args.use_thumbnail, augment, self.args.vision_model_type,
+ find_closest_aspect_ratio_fn=self.find_closest_aspect_ratio_fn
+ )
+
+ num_tiles = [len(imgs)]
+
+ if isinstance(sample, MultiChoiceVQASample):
+ cur_prompt = format_multichoice_question(sample.context, sample.choices)
+ if IMAGE_TOKEN not in cur_prompt:
+ cur_prompt = IMAGE_TOKEN + "\n" + cur_prompt
+ cur_answer = format_multichoice_answer(sample.correct_choice_idx)
+ elif isinstance(sample, VQASample):
+ if 'docvqa' in sample.__key__:
+ prompt_list = self.manual_prompts["VQASFT"]["docvqa"]
+ elif sample.__subflavors__.get("VQASFT"):
+ prompt_list = self.manual_prompts["VQASFT"]["raw"]
+ else:
+ prompt_list = ["{}"]
+
+ prompt_idx = np.random.randint(len(prompt_list))
+ cur_prompt = prompt_list[prompt_idx]
+
+ cur_prompt = cur_prompt.format(sample.context)
+
+ if IMAGE_TOKEN not in cur_prompt:
+ cur_prompt = IMAGE_TOKEN + "\n" + cur_prompt
+
+ if isinstance(sample.answers, list):
+ answer_list = sample.answers
+ weight_list = np.array(sample.answer_weights).astype(np.float32)
+ weight_list = weight_list / np.sum(weight_list)
+ answer_idx = np.random.choice(weight_list.shape[0], 1, p=weight_list)[0]
+ cur_answer = answer_list[answer_idx]
+ else:
+ cur_answer = sample.answers
+ else:
+ raise NotImplementedError("Unsupported data type provided", sample)
+
+ conversation = [
+ {"role": "system", "content": "Answer the questions."},
+ {"role": "user", "content": cur_prompt},
+ {"role": "assistant", "content": str(cur_answer)},
+ ]
+
+ input_ids, target = self.tokenizer.tokenize_conversation(conversation, True, False)
+
+ if self.is_packing_enabled:
+ input_ids, target = self._truncate_for_packing(input_ids, target, num_tiles)
+
+ return ImageTaskSample(
+ __key__=sample.__key__,
+ __restore_key__=sample.__restore_key__,
+ __subflavor__=None,
+ __subflavors__=sample.__subflavors__,
+ imgs=imgs,
+ num_tiles=num_tiles,
+ tokens=torch.tensor(input_ids),
+ labels=torch.tensor(target),
+ total_len=self._get_total_seq_length(input_ids, num_tiles),
+ )
+
+ def combined_ocr_encoder(self, sample, task_type):
+ """Encode OCR samples."""
+ augment = sample.__subflavors__['augmentation'] if 'augmentation' in sample.__subflavors__ else False
+
+ if task_type == "encode_pdf":
+ sample, cur_prompt, cur_answer = self.encode_pdf_prompt(sample)
+ elif task_type == "encode_ocr_ref":
+ sample, cur_prompt, cur_answer = self.encode_ocr_ref_prompt(sample)
+ elif task_type == "_encode_ocr":
+ sample, cur_prompt, cur_answer = self.encode_ocr_prompt(sample)
+
+ imgs = get_visual_transform(
+ sample.image, self.img_h, self.img_w, self.args.use_tiling, self.args.max_num_tiles,
+ self.args.use_thumbnail, augment, self.args.vision_model_type,
+ find_closest_aspect_ratio_fn=self.find_closest_aspect_ratio_fn
+ )
+ num_tiles = [len(imgs)]
+
+ conversation = [
+ {"role": "system", "content": "Answer the questions."},
+ {"role": "user", "content": cur_prompt},
+ {"role": "assistant", "content": str(cur_answer)},
+ ]
+
+ input_ids, target = self.tokenizer.tokenize_conversation(conversation, True, False)
+
+ if self.is_packing_enabled:
+ input_ids, target = self._truncate_for_packing(input_ids, target, num_tiles)
+
+ return ImageTaskSample(
+ __key__=sample.__key__,
+ __restore_key__=sample.__restore_key__,
+ __subflavor__=None,
+ __subflavors__=sample.__subflavors__,
+ imgs=imgs,
+ num_tiles=num_tiles,
+ tokens=torch.tensor(input_ids),
+ labels=torch.tensor(target),
+ total_len=self._get_total_seq_length(input_ids, num_tiles),
+ )
+
+ def encode_pdf_prompt(self, sample: OCRSample) -> ImageTaskSample:
+ """Encode OCR sample."""
+ prompt_list = self.manual_prompts["DocPretraining"]["raw"]
+ prompt_idx = np.random.randint(len(prompt_list))
+ cur_prompt = prompt_list[prompt_idx]
+ if IMAGE_TOKEN not in cur_prompt:
+ cur_prompt = IMAGE_TOKEN + "\n" + cur_prompt
+
+ # Make sure there is no extra IMAGE_TOKEN tag.
+ sample.text = sample.text.replace(IMAGE_TOKEN, "")
+
+ caption = sample.text.strip()
+
+ split_by_line_flag = sample.__subflavors__.get("SplitByLine")
+ if split_by_line_flag:
+ caption_list = caption.split('\n')
+ caption = np.random.choice(caption_list)
+ cur_answer = caption
+
+ return sample, cur_prompt, cur_answer
+
+ def encode_ocr_ref_prompt(self, sample: OCRSample) -> ImageTaskSample:
+ """Encode OCR sample."""
+ ref = sample.text
+ region = sample.words_boxes
+
+ # Make sure there is no extra IMAGE_TOKEN tag
+ ref = ref.replace(IMAGE_TOKEN, "")
+
+ if len(region) == 4:
+ region = f"({region[0]},{region[1]}),({region[2]},{region[3]}) "
+ else:
+ region = f"({region[0]},{region[1]}),({region[2]},{region[3]}),({region[4]},{region[5]}),({region[6]},{region[7]}) "
+
+ # Randomly choose between two tasks
+ task_idx = np.random.randint(2)
+ if task_idx == 0:
+ # Referring Grounding
+ prompt_list = self.manual_prompts["DocPretraining"]["referring_grounding"]
+ prompt_content = ref
+ answer = region
+ else:
+ # Grounded OCR
+ prompt_list = self.manual_prompts["DocPretraining"]["grounded_ocr"]
+ prompt_content = region
+ answer = ref
+
+ prompt_idx = np.random.randint(len(prompt_list))
+ cur_prompt = prompt_list[prompt_idx]
+ cur_prompt = cur_prompt.format(prompt_content)
+ if IMAGE_TOKEN not in cur_prompt:
+ cur_prompt = IMAGE_TOKEN + "\n" + cur_prompt
+
+ return sample, cur_prompt, answer
+
+ def bbox_coord_to_label(self, text, bbox):
+ """Format bbox coordinates as text."""
+ assert len(bbox) == 4 or len(bbox) == 8
+
+ # Make sure there is no extra IMAGE_TOKEN tag
+ text = text.replace(IMAGE_TOKEN, "")
+
+ if len(bbox) == 4:
+ label_str = f"{text}({bbox[0]},{bbox[1]}),({bbox[2]},{bbox[3]}) "
+ else:
+ label_str = f"{text}({bbox[0]},{bbox[1]}),({bbox[2]},{bbox[3]}),({bbox[4]},{bbox[5]}),({bbox[6]},{bbox[7]}) "
+
+ return label_str
+
+ def encode_ocr_prompt(self, sample: OCRSample) -> ImageTaskSample:
+ """Encode OCR sample."""
+ if isinstance(sample.words_boxes[0], int):
+ answer = self.bbox_coord_to_label(sample.text, sample.words_boxes)
+ elif isinstance(sample.words_boxes[0], list):
+ answer = ""
+ for i, bbox in enumerate(sample.words_boxes):
+ answer += self.bbox_coord_to_label(sample.words_text[i], bbox)
+
+ prompt_list = self.manual_prompts["DocPretraining"]["ocr_multi"]
+ prompt_idx = np.random.randint(len(prompt_list))
+ cur_prompt = prompt_list[prompt_idx]
+
+ if IMAGE_TOKEN not in cur_prompt:
+ cur_prompt = IMAGE_TOKEN + "\n" + cur_prompt
+ cur_answer = answer
+
+ return sample, cur_prompt, cur_answer
+
+ def batch(self, samples: List[Union[ImageTaskSample, ImageTaskSamplePacked]]) -> ImageTaskBatchPacked:
+ # Stack images to [num_tiles, c, h, w]. If there are no images (text-only), then use a dummy image.
+ imgs = [img for s in samples for img in s.imgs]
+ if len(imgs) > 0:
+ imgs = torch.stack(imgs)
+ else:
+ imgs = torch.tensor([[0]], dtype=torch.float32)
+
+ # If the user hasn't defined a target dataloader sequence length, then use the max along the sample lengths.
+ max_seq_len = self.dataloader_seq_length
+ if not max_seq_len:
+ max_seq_len = max(len(s.tokens) for s in samples)
+
+ tokens = np.full((len(samples), max_seq_len), self.tokenizer.pad, dtype=np.int64)
+ # +1 to accommodate shift to left by one later.
+ labels = np.full((len(samples), max_seq_len + 1), self.tokenizer.pad, dtype=np.int64)
+
+ for i, s in enumerate(samples):
+ # If the sample/target length exceeds the target sequence length, then truncate.
+ text_len = min(max_seq_len, len(s.tokens))
+ target_len = min(max_seq_len+1, len(s.labels))
+
+ tokens[i, :text_len] = s.tokens[:text_len]
+ labels[i, :target_len] = s.labels[:target_len]
+
+ num_tiles = torch.tensor([n for s in samples for n in s.num_tiles], dtype=torch.int32)
+ if len(num_tiles) == 0:
+ num_tiles = torch.tensor([[0]], dtype=torch.int32)
+
+ # Cumulative sample lengths are needed for packing, otherwise use dummy values.
+ cu_lengths = torch.tensor([[0]], dtype=torch.int32)
+ max_lengths = torch.tensor([[0]], dtype=torch.int32)
+
+ if self.is_packing_enabled:
+ cu_lengths = torch.stack([s.cu_lengths for s in samples])
+ max_lengths = torch.tensor([s.max_length for s in samples], dtype=torch.int32)
+
+ return ImageTaskBatchPacked(
+ __key__=[s.__key__ for s in samples],
+ __restore_key__=[s.__restore_key__ for s in samples],
+ __subflavor__=None,
+ __subflavors__=samples[0].__subflavors__,
+ tokens=tokens,
+ labels=labels,
+ imgs=imgs,
+ num_tiles=num_tiles,
+ cu_lengths=cu_lengths,
+ max_lengths=max_lengths,
+ )
+
+ def encode_batch(self, batch: ImageTaskBatchPacked) -> dict:
+ raw = dataclasses.asdict(batch)
+ del raw["__subflavors__"]
+ return raw
+
+ def select_samples_to_pack(self, samples: List[ImageTaskSample]) -> List[List[ImageTaskSample]]:
+ """Selects which samples will be packed together.
+
+ NOTE: Energon dataloader calls this method internally if packing is used.
+ Please see https://nvidia.github.io/Megatron-Energon/packing.html
+ """
+ lengths = [sample.total_len for sample in samples]
+
+ packed_samples = greedy_knapsack(lengths, samples, self.packing_seq_length)
+
+ return packed_samples
+
+ @stateless
+ def pack_selected_samples(self, samples: List[ImageTaskSample]) -> List[ImageTaskSamplePacked]:
+ """
+ Function to pack a list of ImageTaskSample into a single ImageTaskSamplePacked.
+
+ NOTE: Energon dataloader calls this method internally if packing is used.
+ Please see https://nvidia.github.io/Megatron-Energon/packing.html
+
+ Args:
+ samples: List of ImageTaskSample instances to pack into one sample.
+
+ Returns:
+ ImageTaskSamplePacked instance.
+ """
+ packing_seq_len = self.packing_seq_length
+
+ packed_tokens = []
+ packed_labels = []
+ packed_imgs = []
+
+ current_length = 0
+ max_length = 0
+ cu_lengths = [0]
+
+ # Process each sample and build lists that we will concatenate to create the packed sample.
+ for _, sample in enumerate(samples):
+ sample_len = sample.total_len
+
+ if sample_len > max_length:
+ max_length = sample_len
+
+ # If adding this sample exceeds the max length, stop.
+ # This should not happen. The select_samples_to_pack method should have already ensured that the samples fit.
+ if current_length + sample_len > packing_seq_len:
+ raise ValueError(f"Packed sample exceeds the maximum sequence length of {packing_seq_len}: {samples}")
+
+ # Add the sample's tokens and labels
+ packed_tokens.append(sample.tokens)
+ packed_labels.append(sample.labels)
+
+ # Add the images
+ packed_imgs += sample.imgs
+
+ current_length += sample_len
+ cu_lengths.append(current_length)
+
+ # Concatenate packed tokens and labels.
+ packed_tokens = torch.cat(packed_tokens, dim=0)
+ packed_labels = torch.cat(packed_labels, dim=0)
+
+ return ImageTaskSamplePacked(
+ __key__=",".join([s.__key__ for s in samples]),
+ __restore_key__=(), # Will be set by energon based on `samples`
+ __subflavor__=None,
+ __subflavors__=samples[0].__subflavors__,
+ tokens=packed_tokens,
+ labels=packed_labels,
+ imgs=packed_imgs,
+ cu_lengths=torch.tensor(cu_lengths, dtype=torch.int32),
+ max_length=max_length,
+ num_tiles=[n for s in samples for n in s.num_tiles],
+ )
+
+
+def print_error_handler(exc: Exception, key: Optional[str]):
+ print(
+ f"The following exception occurred in the dataloader for sample {key} and is skipped",
+ file=sys.stderr,
+ )
+ traceback.print_exc()
+
+
+def format_multichoice_question(question, multichoice_options):
+ """Format multi-choice question."""
+ options_text = ["{}. {}\n".format(chr(ord('A') + i), option) for i, option in
+ zip(range(len(multichoice_options)), multichoice_options)]
+ options_text = "".join(options_text)
+
+ options_text = f"{options_text}Answer with the option's letter from the given choices directly."
+
+ return "{}\n{}".format(question, options_text)
+
+
+def format_multichoice_answer(idx):
+ """Format multi-choice answer."""
+ return chr(ord('A') + idx)
diff --git a/examples/multimodal/energon_util.py b/examples/multimodal/energon_util.py
new file mode 100644
index 0000000000..661e691c53
--- /dev/null
+++ b/examples/multimodal/energon_util.py
@@ -0,0 +1,48 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import torch
+import warnings
+from dataclasses import dataclass
+from typing import Any, List
+
+from megatron.energon import Sample
+from megatron.energon.epathlib.epath import EPath
+from megatron.energon.flavors.webdataset import DefaultDecoderWebdatasetFactory
+
+
+@dataclass
+class SampleListSample(Sample):
+ """Sample type for a list of samples of any type which needs to be packed together.
+
+ This is useful for datasets which are packed offline.
+ """
+
+ #: The images of the sequence
+ samples: List[Any]
+
+
+class SampleListWebdataset(DefaultDecoderWebdatasetFactory[SampleListSample]):
+ __sample_type__ = SampleListSample
+
+ def __init__(self, path: EPath, **kwargs):
+ warnings.warn(
+ f"{type(self)} is deprecated, use the default instead and set the sample_type:\n"
+ f"To convert, update your {path}/.nv-meta/dataset.yaml to:\n"
+ f"# remove top-level __module__ and __class__\n"
+ f"sample_type:\n"
+ f" __module__: megatron.energon\n"
+ f" __class__: {self.__sample_type__.__name__}\n"
+ f"# Keep the remaining content",
+ DeprecationWarning,
+ )
+ super().__init__(path, **kwargs)
+
+
+@dataclass
+class OfflineTargetAspectRatioSample(Sample):
+ """Sample type for image + text samples with target aspect ratio computed offline."""
+
+ #: The images of the sequence
+ images: List[torch.Tensor]
+ #: The texts of the sequence
+ texts: List[str]
+ target_aspect_ratio: List[List]
diff --git a/examples/multimodal/evaluation/evaluate_ai2d.py b/examples/multimodal/evaluation/evaluate_ai2d.py
new file mode 100644
index 0000000000..39b866ae4a
--- /dev/null
+++ b/examples/multimodal/evaluation/evaluate_ai2d.py
@@ -0,0 +1,52 @@
+import argparse
+import json
+
+from evaluate_mmmu import get_input_output_paths
+from evaluate_vqav2 import compute_vqa_accuracy
+
+
+def merge_input_files(input_path):
+ """Merge input files to a format compatible with the evaluator."""
+ input_file_paths, output_file_path = get_input_output_paths(input_path, task="AI2D")
+
+ results = dict()
+
+ for input_file_path in input_file_paths:
+ with open(input_file_path, "r") as input_file:
+ for line in input_file:
+ res = json.loads(line)
+ sample_id = res["sample_id"]
+
+ # Ignore possible duplicates.
+ if sample_id in results:
+ continue
+
+ results[sample_id] = {
+ "question_id": sample_id,
+ "answer": res["answer"],
+ "gt_answer": res["gt_answer"],
+ }
+
+ results = list(results.values())
+
+ with open(output_file_path, "w") as output_file:
+ json.dump(results, output_file)
+
+ return output_file_path
+
+
+def ai2d_eval(input_path):
+ """Run AI2D evaluation."""
+ result_file_path = merge_input_files(input_path)
+ avg_acc = compute_vqa_accuracy(result_file_path, task="AI2D")
+ return avg_acc
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--input-path', type=str, help="Path to input file(s)")
+ args = parser.parse_args()
+
+ avg_acc = ai2d_eval(args.input_path)
+
+ print(f"===== AI2D Accuracy {avg_acc:.2f}% =====")
diff --git a/examples/multimodal/evaluation/evaluate_chartqa.py b/examples/multimodal/evaluation/evaluate_chartqa.py
new file mode 100644
index 0000000000..53d4944f46
--- /dev/null
+++ b/examples/multimodal/evaluation/evaluate_chartqa.py
@@ -0,0 +1,48 @@
+import argparse
+import json
+
+from evaluate_mmmu import get_input_output_paths
+from evaluate_vqav2 import compute_vqa_accuracy
+
+
+def merge_input_files(input_path):
+ """Merge input files to a format compatible with the evaluator."""
+ input_file_paths, output_file_path = get_input_output_paths(input_path, task="ChartQA")
+
+ results = dict()
+
+ for input_file_path in input_file_paths:
+ with open(input_file_path, "r") as input_file:
+ for line in input_file:
+ res = json.loads(line)
+ sample_id = res["sample_id"]
+
+ # Ignore possible duplicates.
+ if sample_id in results:
+ continue
+
+ res["question_id"] = sample_id
+ results[sample_id] = res
+
+ results = list(results.values())
+
+ with open(output_file_path, "w") as output_file:
+ json.dump(results, output_file)
+
+ return output_file_path
+
+
+def chartqa_eval(input_path):
+ """Run ChartQA evaluation."""
+ result_file_path = merge_input_files(input_path)
+ return compute_vqa_accuracy(result_file_path, task="ChartQA")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--input-path', type=str, help="Path to input file(s)")
+ args = parser.parse_args()
+
+ avg_acc = chartqa_eval(args.input_path)
+
+ print(f"ChartQA accuracy: {avg_acc:.2f}")
diff --git a/examples/multimodal/evaluation/evaluate_coco.py b/examples/multimodal/evaluation/evaluate_coco.py
new file mode 100644
index 0000000000..8eeb367e8f
--- /dev/null
+++ b/examples/multimodal/evaluation/evaluate_coco.py
@@ -0,0 +1,66 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import argparse
+import json
+
+from evaluate_mmmu import get_input_output_paths
+from pycocoevalcap.eval import COCOEvalCap
+from pycocotools.coco import COCO
+
+
+def convert_to_coco_format(input_path):
+ """Convert input files to COCO compatible format."""
+ input_file_paths, output_file_path = get_input_output_paths(input_path, task="captioning")
+
+ results = dict()
+
+ for input_file_path in input_file_paths:
+ with open(input_file_path, "r") as input_file:
+ for line in input_file:
+ res = json.loads(line)
+ sample_id = res["sample_id"]
+
+ # Ignore possible duplicates.
+ if sample_id in results:
+ continue
+
+ caption = res["caption"].rstrip(".").lower()
+ results[sample_id] = {
+ "image_id": sample_id,
+ "caption": caption,
+ }
+
+ results = list(results.values())
+
+ with open(output_file_path, "w") as output_file:
+ json.dump(results, output_file, indent=4)
+
+ return output_file_path
+
+
+def coco_captioning_eval(input_path, groundtruth_file):
+ """Run COCO captioning evaluation."""
+ coco = COCO(groundtruth_file)
+ input_file = convert_to_coco_format(input_path)
+ coco_result = coco.loadRes(input_file)
+
+ coco_eval = COCOEvalCap(coco, coco_result)
+
+ # Evaluate on the input subset of images.
+ coco_eval.params["image_id"] = coco_result.getImgIds()
+
+ coco_eval.evaluate()
+
+ print("========== COCO captioning scores ==========")
+ for metric, score in coco_eval.eval.items():
+ print(f"{metric} {score * 100:.3f}")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--input-path", type=str, required=True, help="Path to input file(s)")
+ parser.add_argument(
+ "--groundtruth-path", type=str, required=True, help="Path to groundtruth file"
+ )
+ args = parser.parse_args()
+
+ coco_captioning_eval(args.input_path, args.groundtruth_path)
diff --git a/examples/multimodal/evaluation/evaluate_mathvista.py b/examples/multimodal/evaluation/evaluate_mathvista.py
new file mode 100644
index 0000000000..a55f312f21
--- /dev/null
+++ b/examples/multimodal/evaluation/evaluate_mathvista.py
@@ -0,0 +1,122 @@
+import argparse
+import json
+import re
+
+from evaluate_mmmu import get_input_output_paths
+from MMMU.mmmu.utils.eval_utils import parse_multi_choice_response
+from open_flamingo.eval.vqa_metric import VQAEval
+
+
+def merge_input_files(input_path):
+ """Merge input files to a format compatible with the evaluator."""
+ input_file_paths, output_file_path = get_input_output_paths(input_path, task="MathVista")
+
+ results = dict()
+
+ for input_file_path in input_file_paths:
+ with open(input_file_path, "r") as input_file:
+ for line in input_file:
+ res = json.loads(line)
+ sample_id = res["sample_id"]
+
+ # Remove possible duplicates.
+ if sample_id in results:
+ continue
+
+ results[sample_id] = res
+
+ results = list(results.values())
+
+ with open(output_file_path, "w") as output_file:
+ json.dump(results, output_file)
+
+ return output_file_path
+
+
+def extra_processing(text):
+ """Extra processing."""
+ # Max decimal point capped to 2 decimal point
+ regex = re.compile(r'^\d+\.\d+$')
+ decimal = regex.findall(text)
+
+ if len(decimal) > 0:
+ non_decimal = len(decimal[0].split(".")[0])
+
+ # if decimal values are all 0, trim them
+ decimal_digits = [int(d) for d in decimal[0].split(".")[1]]
+ if sum(decimal_digits) == 0:
+ text = decimal[0][:non_decimal]
+ else:
+ text = decimal[0][: non_decimal + 3]
+
+ # remove % and trailing .
+ text = text.replace("%", "")
+ if text[-1] == ".":
+ text = text[:-1]
+
+ return text
+
+
+def extract_answer(text):
+ """Extract answer."""
+ alphabet = re.findall(r'[a-zA-Z]+', text)
+ if len(alphabet) > 0 and "e+" not in text:
+ template = re.findall(r'answer is -*\d+\.*\d*', text)
+ if len(template) > 0:
+ text = template[0]
+
+ numbers = re.findall(r'-*\d+\.*\d*', text)
+ text = numbers[0] if len(numbers) > 0 else text
+
+ return text
+
+
+def compute_mathvista_accuracy(result_file):
+ """Compute MathVista accuracy."""
+ merged_results = json.load(open(result_file))
+
+ vqa = VQAEval(vqa=None, vqaRes=None)
+ acc = 0
+ for res in merged_results:
+ pred_ans = res["answer"]
+ if res["question_type"] == "multi_choice":
+ pred_ans = parse_multi_choice_response(pred_ans, res["all_choices"], res["index2ans"])
+ else:
+ pred_ans = vqa.processPunctuation(pred_ans)
+ pred_ans = vqa.processDigitArticle(pred_ans)
+ # Extra processing and extraction.
+ pred_ans = extra_processing(pred_ans)
+ pred_ans = extract_answer(pred_ans)
+
+ gt_ans = res["gt_answer"]
+ if isinstance(gt_ans, list):
+ assert len(gt_ans) == 1, f"Expected 1 groundtruth, got {gt_ans}"
+ gt_ans = gt_ans[0]
+
+ if res["question_type"] != "multi_choice":
+ gt_ans = vqa.processPunctuation(gt_ans)
+ gt_ans = vqa.processDigitArticle(gt_ans)
+
+ gt_ans = extra_processing(gt_ans)
+
+ if pred_ans == gt_ans:
+ acc += 1
+ acc = acc / len(merged_results) * 100
+ return acc
+
+
+def mathvista_eval(input_path):
+ """Run MathVista evaluation."""
+ result_file_path = merge_input_files(input_path)
+ acc = compute_mathvista_accuracy(result_file_path)
+ return acc
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--input-path', type=str, help="Path to input file(s)")
+ args = parser.parse_args()
+
+ acc = mathvista_eval(args.input_path)
+
+ print(f"===== MathVista accuracy: {acc} =====")
diff --git a/examples/multimodal/evaluation/evaluate_mmmu.py b/examples/multimodal/evaluation/evaluate_mmmu.py
new file mode 100644
index 0000000000..798c42bfa7
--- /dev/null
+++ b/examples/multimodal/evaluation/evaluate_mmmu.py
@@ -0,0 +1,116 @@
+import argparse
+import glob
+import json
+import os
+import sys
+import re
+import subprocess
+
+# Get the absolute path of the parent directory
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))
+# Add the parent directory to sys.path
+sys.path.insert(0, parent_dir)
+
+from run_text_generation import get_output_path
+from config import EvaluationConfig
+
+
+def get_input_output_paths(input_path, task):
+ """Get all input files and an output path for a merged file."""
+ # Single input file.
+ if os.path.exists(input_path):
+ input_file_paths = [input_path]
+ output_file_path = input_path.replace(".jsonl", "-merged.json")
+ # Select multiple partitions and dp ranks.
+ else:
+ cfg = EvaluationConfig(task=task, output_path=input_path, partition_id="*")
+ pattern = get_output_path(cfg, dp_rank="*")
+ input_file_paths = glob.glob(pattern)
+
+ output_file_path = input_path + f"-{task}-merged.json"
+
+ return input_file_paths, output_file_path
+
+
+def convert_to_mmmu_format(input_path):
+ """Convert input files to MMMU compatible format."""
+ input_file_paths, output_file_path = get_input_output_paths(input_path, "MMMU")
+
+ output = dict()
+
+ for input_file_path in input_file_paths:
+ with open(input_file_path, "r") as input_file:
+ for line in input_file:
+ res = json.loads(line)
+
+ sample_id = res["sample_id"]
+ prediction = res["prediction"]
+
+ if res["question_type"] == "multiple-choice":
+ from MMMU.mmmu.utils.eval_utils import parse_multi_choice_response
+
+ prediction = parse_multi_choice_response(
+ prediction, res["all_choices"], res["index2ans"]
+ )
+
+ # MMMU eval script expects just a sample_id to prediction mapping.
+ # Skip possible duplicates.
+ if sample_id in output:
+ continue
+
+ output[sample_id] = prediction
+
+ with open(output_file_path, "w") as output_file:
+ json.dump(output, output_file)
+
+ return output_file_path
+
+
+def mmmu_eval(input_path, groundtruth_path):
+ """Run MMMU evaluation."""
+ result_file = convert_to_mmmu_format(input_path)
+
+ # The MMMU repo has a script for running the actual evaluation but no API. So launching the script here.
+ output = subprocess.run(
+ [
+ "python",
+ "examples/multimodal/MMMU/mmmu/main_eval_only.py",
+ "--output_path",
+ result_file,
+ "--answer_path",
+ groundtruth_path,
+ ],
+ capture_output=True,
+ text=True,
+ )
+
+ print(output.stderr)
+ print(output.stdout)
+
+ m = re.search("'Overall': {'num': \d+, 'acc': (\d.\d+)}", output.stdout)
+
+ return float(m.group(1)) * 100.0
+
+
+def main():
+ """Run MMMU evaluation."""
+ # Using the validation groundtruth file from the MMMU repo by default. This assumes you have cloned the MMMU github repo here.
+ default_groundtruth_path = "examples/multimodal/MMMU/mmmu/answer_dict_val.json"
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--input-path", type=str, required=True, help="Path to input file(s)")
+ parser.add_argument(
+ "--groundtruth-path",
+ type=str,
+ default=default_groundtruth_path,
+ help="Path to groundtruth file. Defaults to the validation file in the MMMU repo.",
+ )
+ args = parser.parse_args()
+
+ avg_acc = mmmu_eval(args.input_path, args.groundtruth_path)
+
+ print(f"MMMU average accuracy: {avg_acc:.2f}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/multimodal/evaluation/evaluate_ocrbench.py b/examples/multimodal/evaluation/evaluate_ocrbench.py
new file mode 100644
index 0000000000..b37473a67d
--- /dev/null
+++ b/examples/multimodal/evaluation/evaluate_ocrbench.py
@@ -0,0 +1,137 @@
+import argparse
+import json
+
+from evaluate_mmmu import get_input_output_paths
+
+
+def merge_input_files(input_path):
+ """Merge input files to a format compatible with the evaluator."""
+ input_file_paths, output_file_path = get_input_output_paths(input_path, task="OCRBench")
+
+ results = dict()
+
+ for input_file_path in input_file_paths:
+ with open(input_file_path, "r") as input_file:
+ for line in input_file:
+ res = json.loads(line)
+ sample_id = res["sample_id"]
+
+ # Remove possible duplicates.
+ if sample_id in results:
+ continue
+
+ results[sample_id] = res
+
+ results = list(results.values())
+
+ with open(output_file_path, "w") as output_file:
+ json.dump(results, output_file)
+
+ return output_file_path
+
+
+def compute_ocrbench_score(result_file):
+ """Compute OCRBench score."""
+ merged_results = json.load(open(result_file))
+
+ # OCRBench score calculation is adopted from https://github.com/Yuliang-Liu/MultimodalOCR/blob/1b7713f44c91f30f64efb6d3e494c416861ef15f/example.py#L1
+ # MIT License. Copyright (c) 2023 Yuliang Liu
+ score = {
+ "Regular Text Recognition": 0,
+ "Irregular Text Recognition": 0,
+ "Artistic Text Recognition": 0,
+ "Handwriting Recognition": 0,
+ "Digit String Recognition": 0,
+ "Non-Semantic Text Recognition": 0,
+ "Scene Text-centric VQA": 0,
+ "Doc-oriented VQA": 0,
+ "Doc-oriented VQA": 0,
+ "Key Information Extraction": 0,
+ "Handwritten Mathematical Expression Recognition": 0,
+ }
+
+ for res in merged_results:
+ predict = res["answer"]
+ answers = res["gt_answer"]
+
+ dataset_name = res["dataset_name"]
+ ocr_type = res["data_type"]
+
+ if dataset_name == "HME100k":
+ if isinstance(answers, list):
+ for j in range(len(answers)):
+ answer = answers[j].strip().replace("\n", " ").replace(" ", "")
+ predict = predict.strip().replace("\n", " ").replace(" ", "")
+ if answer in predict:
+ score[ocr_type] += 1
+ else:
+ answers = answers.strip().replace("\n", " ").replace(" ", "")
+ predict = predict.strip().replace("\n", " ").replace(" ", "")
+ if answers in predict:
+ score[ocr_type] += 1
+ else:
+ if isinstance(answers, list):
+ for j in range(len(answers)):
+ answer = answers[j].lower().strip().replace("\n", " ")
+ predict = predict.lower().strip().replace("\n", " ")
+ if answer in predict:
+ score[ocr_type] += 1
+ else:
+ answers = answers.lower().strip().replace("\n", " ")
+ predict = predict.lower().strip().replace("\n", " ")
+ if answers in predict:
+ score[ocr_type] += 1
+
+ recognition_score = (
+ score['Regular Text Recognition']
+ + score['Irregular Text Recognition']
+ + score['Artistic Text Recognition']
+ + score['Handwriting Recognition']
+ + score['Digit String Recognition']
+ + score['Non-Semantic Text Recognition']
+ )
+ final_score = (
+ recognition_score
+ + score['Scene Text-centric VQA']
+ + score['Doc-oriented VQA']
+ + score['Key Information Extraction']
+ + score['Handwritten Mathematical Expression Recognition']
+ )
+ result_log = f"""###########################OCRBench##############################
+Text Recognition(Total 300): {recognition_score}
+------------------Details of Recognition Score-------------------
+Regular Text Recognition(Total 50): {score['Regular Text Recognition']}
+Irregular Text Recognition(Total 50): {score['Irregular Text Recognition']}
+Artistic Text Recognition(Total 50): {score['Artistic Text Recognition']}
+Handwriting Recognition(Total 50): {score['Handwriting Recognition']}
+Digit String Recognition(Total 50): {score['Digit String Recognition']}
+Non-Semantic Text Recognition(Total 50): {score['Non-Semantic Text Recognition']}
+----------------------------------------------------------------
+Scene Text-centric VQA(Total 200): {score['Scene Text-centric VQA']}
+----------------------------------------------------------------
+Doc-oriented VQA(Total 200): {score['Doc-oriented VQA']}
+----------------------------------------------------------------
+Key Information Extraction(Total 200): {score['Key Information Extraction']}
+----------------------------------------------------------------
+Handwritten Mathematical Expression Recognition(Total 100): {score['Handwritten Mathematical Expression Recognition']}
+----------------------Final Score-------------------------------
+Final Score(Total 1000): {final_score}"""
+
+ return result_log, final_score
+
+
+def ocrbench_eval(input_path):
+ """Run OCRBench evaluation."""
+ result_file_path = merge_input_files(input_path)
+ result_log, score = compute_ocrbench_score(result_file_path)
+ return result_log, score
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--input-path', type=str, help="Path to input file(s)")
+ args = parser.parse_args()
+
+ result_log, _ = ocrbench_eval(args.input_path)
+
+ print(result_log)
diff --git a/examples/multimodal/evaluation/evaluate_textvqa.py b/examples/multimodal/evaluation/evaluate_textvqa.py
new file mode 100644
index 0000000000..af782bdf03
--- /dev/null
+++ b/examples/multimodal/evaluation/evaluate_textvqa.py
@@ -0,0 +1,52 @@
+import argparse
+import json
+
+from evaluate_mmmu import get_input_output_paths
+from evaluate_vqav2 import compute_vqa_accuracy
+
+
+def merge_input_files(input_path):
+ """Merge input files to a format compatible with the evaluator."""
+ input_file_paths, output_file_path = get_input_output_paths(input_path, task="TextVQA")
+
+ results = dict()
+
+ for input_file_path in input_file_paths:
+ with open(input_file_path, "r") as input_file:
+ for line in input_file:
+ res = json.loads(line)
+ sample_id = res["sample_id"]
+
+ # Remove possible duplicates.
+ if sample_id in results:
+ continue
+
+ results[sample_id] = {
+ "question_id": sample_id,
+ "answer": res["answer"],
+ "gt_answer": res["gt_answer"],
+ }
+
+ results = list(results.values())
+
+ with open(output_file_path, "w") as output_file:
+ json.dump(results, output_file)
+
+ return output_file_path
+
+
+def textvqa_eval(input_path):
+ """Run TextVQA evaluation."""
+ result_file_path = merge_input_files(input_path)
+ avg_acc = compute_vqa_accuracy(result_file_path, task="TextVQA")
+ return avg_acc
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--input-path', type=str, help="Path to input file(s)")
+ args = parser.parse_args()
+
+ avg_acc = textvqa_eval(args.input_path)
+
+ print(f"===== TextVQA Accuracy {avg_acc:.2f}% =====")
diff --git a/examples/multimodal/evaluation/evaluate_vqav2.py b/examples/multimodal/evaluation/evaluate_vqav2.py
new file mode 100644
index 0000000000..7807d80723
--- /dev/null
+++ b/examples/multimodal/evaluation/evaluate_vqav2.py
@@ -0,0 +1,109 @@
+import argparse
+import json
+
+from evaluate_mmmu import get_input_output_paths
+from open_flamingo.eval.vqa_metric import VQAEval
+
+
+def merge_input_files(input_path):
+ """Merge input files to a format compatible with the evaluator."""
+ input_file_paths, output_file_path = get_input_output_paths(input_path, task="VQAv2")
+
+ results = dict()
+
+ for input_file_path in input_file_paths:
+ with open(input_file_path, "r") as input_file:
+ for line in input_file:
+ res = json.loads(line)
+ sample_id = res["sample_id"]
+
+ # Skip possible duplicates.
+ if sample_id in results:
+ continue
+
+ res["question_id"] = sample_id
+ results[sample_id] = res
+
+ results = list(results.values())
+
+ with open(output_file_path, "w") as output_file:
+ json.dump(results, output_file)
+
+ return output_file_path
+
+
+def is_number(n: str):
+ """Check if input is a number."""
+ try:
+ float(n)
+ return True
+ except ValueError:
+ return False
+
+
+def compute_vqa_accuracy(result_file, task):
+ """Compute VQA accuracy."""
+ merged_results = json.load(open(result_file))
+
+ vqa = VQAEval(vqa=None, vqaRes=None)
+ all_acc = []
+ for res in merged_results:
+ pred = res["answer"]
+ pred = vqa.processPunctuation(pred)
+ pred = vqa.processDigitArticle(pred)
+
+ gt = res["gt_answer"]
+ gt = [vqa.processPunctuation(ans) for ans in gt]
+ gt = [vqa.processDigitArticle(ans) for ans in gt]
+
+ # ChartQA uses relaxed accuracy:
+ # "We consider an answer to be correct if it is within 5% of the gold answer.
+ # For non-numeric answers, we still need an exact match to consider an answer to be correct."
+ if task == "ChartQA":
+ acc = 0.0
+ assert len(gt) == 1, "expected exactly one groundtruth answer."
+ gt = gt[0]
+
+ pred = pred.rstrip("%")
+ gt = gt.rstrip("%")
+
+ if is_number(pred) and is_number(gt):
+ pred = float(pred)
+ gt = float(gt)
+ if pred >= (gt * 0.95) and pred <= (gt * 1.05):
+ acc = 1.0
+ elif pred == gt:
+ acc = 1.0
+
+ all_acc.append(acc)
+ elif task in ("VQAv2", "TextVQA"):
+ num_match = sum([pred == ans for ans in gt])
+ acc = min(1.0, num_match / 3.0)
+ all_acc.append(acc)
+ elif task == "AI2D":
+ assert len(gt) == 1, f"Expected exactly 1 GT, got {gt}"
+ acc = pred == gt[0]
+ all_acc.append(acc)
+ else:
+ raise NotImplementedError(f"unknown task {task}")
+
+ acc_avg = sum(all_acc) / len(all_acc) * 100
+
+ return acc_avg
+
+
+def vqav2_eval(input_path):
+ """Run VQAv2 evaluation."""
+ result_file = merge_input_files(input_path)
+ avg_acc = compute_vqa_accuracy(result_file, task="VQAv2")
+ return avg_acc
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--input-path', type=str, help="Path to input file(s)")
+ args = parser.parse_args()
+
+ avg_acc = vqav2_eval(args.input_path)
+
+ print(f"===== VQAv2 Accuracy {avg_acc:.2f}% =====")
diff --git a/examples/multimodal/evaluation/evaluation_datasets.py b/examples/multimodal/evaluation/evaluation_datasets.py
new file mode 100644
index 0000000000..dc7ebf3c30
--- /dev/null
+++ b/examples/multimodal/evaluation/evaluation_datasets.py
@@ -0,0 +1,914 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+"""Evaluation datasets."""
+import glob
+import itertools
+import json
+import os
+import re
+from collections import defaultdict
+
+import numpy as np
+import torch
+from image_processing import get_visual_transform
+from PIL import Image
+
+from megatron.training import print_rank_0
+
+
+def _get_partition_bounds(
+ total_num_samples, num_samples_per_partition, num_partitions, partition_id
+):
+ if num_samples_per_partition == 0:
+ samples_per_partition = [
+ int(x) for x in np.linspace(0, total_num_samples, num_partitions + 1)
+ ]
+ return samples_per_partition[partition_id], samples_per_partition[partition_id + 1]
+ return num_samples_per_partition * partition_id, num_samples_per_partition * (partition_id + 1)
+
+
+class VQADataset(torch.utils.data.Dataset):
+ """VQA evaluation dataset."""
+
+ def __init__(
+ self,
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ keys,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ ):
+ samples = json.load(open(gt_path, encoding='utf-8'))
+ if "data" in samples:
+ samples = samples["data"]
+
+ # Optionally, process only a subset of the input files.
+ if num_partitions > 0:
+ lb, ub = _get_partition_bounds(
+ len(samples), num_samples_per_partition, num_partitions, partition_id
+ )
+ samples = samples[lb:ub]
+
+ self._keys = keys
+ self._samples = samples
+ self._input_image_path = input_image_path
+ self._img_h = img_h
+ self._img_w = img_w
+ self._use_tiling = use_tiling
+ self._max_num_tiles = max_num_tiles
+ self._use_thumbnail = use_thumbnail
+ self._vision_model_type = vision_model_type
+
+ def __len__(self):
+ return len(self._samples)
+
+ def __getitem__(self, idx):
+ sample = self._samples[idx]
+
+ img_file = "{}/{}".format(self._input_image_path, sample[self._keys["image_id"]])
+ if not os.path.exists(img_file):
+ img_file += ".jpg"
+
+ if not os.path.exists(img_file):
+ img_file = img_file.replace('.jpg', '.png')
+
+ img = Image.open(img_file)
+ imgs = get_visual_transform(
+ img,
+ self._img_h,
+ self._img_w,
+ self._use_tiling,
+ self._max_num_tiles,
+ self._use_thumbnail,
+ augment=False,
+ vision_model_type=self._vision_model_type,
+ )
+ tile_count = torch.tensor([len(imgs)], dtype=torch.int)
+
+ sample_id = idx
+ if "sample_id" in self._keys:
+ sample_id = sample[self._keys["sample_id"]]
+
+ metadata = "" # Not used.
+
+ return (
+ torch.stack(imgs),
+ tile_count,
+ sample_id,
+ sample[self._keys["question"]],
+ sample[self._keys["answer"]],
+ metadata,
+ )
+
+
+class CaptioningDataset(torch.utils.data.Dataset):
+ """Captioning evaluation dataset."""
+
+ def __init__(
+ self,
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ ):
+ image_files = sorted(glob.glob(input_image_path + "/*"))
+
+ # Optionally, process only a subset of the input files.
+ if num_partitions > 0:
+ lb, ub = _get_partition_bounds(
+ len(image_files), num_samples_per_partition, num_partitions, partition_id
+ )
+ image_files = image_files[lb:ub]
+
+ gts = json.load(open(gt_path))
+ answers = defaultdict(list)
+ for gt in gts["annotations"]:
+ answers[gt["image_id"]].append(gt['caption'])
+
+ self._image_files = image_files
+ self._answers = answers
+ self._img_h = img_h
+ self._img_w = img_w
+ self._use_tiling = use_tiling
+ self._max_num_tiles = max_num_tiles
+ self._use_thumbnail = use_thumbnail
+ self._vision_model_type = vision_model_type
+
+ def __len__(self):
+ return len(self._image_files)
+
+ def __getitem__(self, idx):
+ img_file = self._image_files[idx]
+ image_id = int(img_file.split("_")[-1].split(".")[0])
+
+ img = Image.open(img_file)
+ imgs = get_visual_transform(
+ img,
+ self._img_h,
+ self._img_w,
+ self._use_tiling,
+ self._max_num_tiles,
+ self._use_thumbnail,
+ augment=False,
+ vision_model_type=self._vision_model_type,
+ )
+
+ tile_count = torch.tensor([len(imgs)], dtype=torch.int)
+
+ question = "" # Fixed for all samples.
+ metadata = "" # Not used.
+
+ return torch.stack(imgs), tile_count, image_id, question, self._answers[image_id], metadata
+
+
+class MMMUDataset(torch.utils.data.Dataset):
+ """MMMU evaluation dataset."""
+
+ def __init__(
+ self,
+ input_image_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ prompt_style,
+ vision_model_type,
+ ):
+ import datasets
+ from MMMU.mmmu.utils.data_utils import CAT_SHORT2LONG, load_yaml
+
+ # The following downloads the MMMU dataset from HuggingFace and uses the API from the MMMU github repo to run MMMU evaluation.
+ all_mmmu_datasets = []
+
+ hf_datasets_cache = os.environ["HF_DATASETS_CACHE"]
+ assert hf_datasets_cache != "", "Please set the environment variable HF_DATASETS_CACHE."
+
+ for subject in CAT_SHORT2LONG.values():
+ # Use a local copy of the dataset if exists (can be faster) or the HF one.
+ if os.path.exists(input_image_path):
+ subject_dataset = datasets.load_dataset(
+ os.path.join(input_image_path, subject),
+ split=datasets.Split.VALIDATION,
+ cache_dir=hf_datasets_cache,
+ verification_mode="no_checks",
+ )
+ else:
+ subject_dataset = datasets.load_dataset(
+ "MMMU/MMMU",
+ subject,
+ split=datasets.Split.VALIDATION,
+ cache_dir=hf_datasets_cache,
+ )
+
+ all_mmmu_datasets.append(subject_dataset)
+
+ dataset = datasets.concatenate_datasets(all_mmmu_datasets)
+
+ dataset = [s for s in dataset if s['id'].startswith("val")]
+
+ # Optionally, process only a subset of the input files.
+ if num_partitions > 0:
+ lb, ub = _get_partition_bounds(
+ len(dataset), num_samples_per_partition, num_partitions, partition_id
+ )
+ dataset = dataset[lb:ub]
+
+ # Using the LLaVA config from the MMMU repo.
+ config = load_yaml("examples/multimodal/MMMU/mmmu/configs/llava1.5.yaml")
+ for k, v in config.items():
+ if isinstance(v, list):
+ assert len(v) == 1, "only one value supported."
+ config[k] = v[0]
+
+ self._config = config
+
+ self._dataset = dataset
+
+ self._img_h = img_h
+ self._img_w = img_w
+ self._use_tiling = use_tiling
+ self._max_num_tiles = max_num_tiles
+ self._use_thumbnail = use_thumbnail
+ self._prompt_style = prompt_style
+ self._vision_model_type = vision_model_type
+
+ def __len__(self):
+ return len(self._dataset)
+
+ def __getitem__(self, idx):
+ from MMMU.mmmu.utils.data_utils import construct_prompt, process_single_sample
+
+ sample = self._dataset[idx]
+
+ # Use the single image approach from the MMMU repo.
+ if self._prompt_style == "single_image":
+ sample = process_single_sample(sample)
+ sample = construct_prompt(sample, self._config)
+
+ img = sample["image"]
+ sample_imgs = get_visual_transform(
+ img,
+ self._img_h,
+ self._img_w,
+ self._use_tiling,
+ self._max_num_tiles,
+ self._use_thumbnail,
+ augment=False,
+ vision_model_type=self._vision_model_type,
+ )
+ sample_num_tiles = [len(sample_imgs)]
+
+ prompt = sample["final_input_prompt"]
+ for i in range(8):
+ prompt = prompt.replace(f"", "")
+ sample["final_input_prompt"] = f"\n{prompt}"
+ elif self._prompt_style == "vlmevalkit":
+ sample = construct_prompt(sample, self._config)
+
+ if sample["question_type"] == "multiple-choice":
+ question = sample["question"]
+
+ options = ""
+ for k, v in sample["index2ans"].items():
+ options += f"{k}. {v}\n"
+
+ final_prompt = f"{question}\n"
+ if "hint" in sample:
+ final_prompt += f"Hint: {sample['hint']}\n"
+
+ if "task_instructions" in sample:
+ final_prompt += f"Task instructions: {sample['task_instructions']}\n"
+
+ final_prompt += options
+ final_prompt += "Answer with the option's letter from the given choices directly."
+
+ sample["final_input_prompt"] = final_prompt.rstrip()
+ else:
+ question = sample["question"]
+ final_prompt = f"{question}\n"
+ final_prompt += "Answer the question directly."
+ sample["final_input_prompt"] = final_prompt.rstrip()
+
+ sample_imgs = []
+ sample_num_tiles = []
+
+ img_indices = sorted(list(set(re.findall(r""
+
+ img = sample[img_key]
+ assert img is not None, f"{img_str} is in prompt but not in sample images"
+
+ imgs = get_visual_transform(
+ img,
+ self._img_h,
+ self._img_w,
+ self._use_tiling,
+ adjusted_max_num_tiles,
+ self._use_thumbnail,
+ augment=False,
+ vision_model_type=self._vision_model_type,
+ ) # List of tiles.
+
+ sample_imgs.extend(imgs)
+ sample_num_tiles.append(len(imgs))
+
+ sample["final_input_prompt"] = " ".join([f'' for i in range(len(img_indices))]) + "\n" + sample["final_input_prompt"]
+ elif self._prompt_style == "multi_image":
+ sample = construct_prompt(sample, self._config)
+
+ sample_imgs = []
+ sample_num_tiles = []
+
+ img_indices = re.findall(r""
+
+ img = sample[img_key]
+ assert img is not None, f"{img_str} is in prompt but not in sample images"
+
+ # Note: Only replace the current image tag.
+ sample["final_input_prompt"] = sample["final_input_prompt"].replace(
+ img_str, "", 1
+ )
+
+ imgs = get_visual_transform(
+ img,
+ self._img_h,
+ self._img_w,
+ self._use_tiling,
+ adjusted_max_num_tiles,
+ self._use_thumbnail,
+ augment=False,
+ vision_model_type=self._vision_model_type,
+ ) # List of tiles.
+
+ sample_imgs.extend(imgs)
+ sample_num_tiles.append(len(imgs))
+
+ # Sanity check.
+ for i in range(1, 8):
+ assert (
+ f"" not in sample["final_input_prompt"]
+ ), "prompt contains unhandled image tags"
+ else:
+ raise ValueError(f"unknown prompt style {self._prompt_style}")
+
+ # MMMU specific metadata.
+ metadata = {"question_type": sample["question_type"]}
+ if sample["question_type"] == "multiple-choice":
+ metadata["index2ans"] = sample["index2ans"]
+ metadata["all_choices"] = sample["all_choices"]
+
+ prompt = sample['final_input_prompt']
+
+ tile_count = torch.tensor(sample_num_tiles, dtype=torch.int)
+
+ return (
+ torch.stack(sample_imgs),
+ tile_count,
+ sample["id"],
+ prompt,
+ sample["answer"],
+ metadata,
+ )
+
+
+class VideoMMEDataset(torch.utils.data.Dataset):
+ "Video MME evaluation dataset."
+
+ def __init__(
+ self,
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ num_frames,
+ vision_model_type,
+ ):
+ ground_truth_original = json.load(open(gt_path))
+ ground_truth = []
+ for gt in ground_truth_original:
+ video_path = gt["url"]
+ video_path = video_path.replace("https://www.youtube.com/watch?v=", "")
+ video_path = video_path.replace("https://m.youtube.com/watch?v=", "")
+ video_path = os.path.join(input_image_path, video_path + ".mp4")
+ if not os.path.exists(video_path):
+ continue
+ gt["video_path"] = video_path
+ ground_truth.append(gt)
+
+ ground_truth = sorted(ground_truth, key=lambda gt: gt["video_path"])
+ print_rank_0(f"Found {len(ground_truth)} videos to process.")
+
+ if num_partitions > 0:
+ start_idx, end_idx = _get_partition_bounds(
+ len(ground_truth), num_samples_per_partition, num_partitions, partition_id
+ )
+ ground_truth = ground_truth[start_idx:end_idx]
+
+ self._ground_truth = ground_truth
+ self._img_h = img_h
+ self._img_w = img_w
+ self._use_tiling = False
+ self._max_num_tiles = max_num_tiles
+ self._use_thumbnail = use_thumbnail
+ self._num_frames = num_frames
+ self._vision_model_type = vision_model_type
+
+ def __len__(self):
+ return len(self._ground_truth)
+
+ def __getitem__(self, idx):
+ from torchvision.io import read_video
+
+ gt = self._ground_truth[idx]
+
+ video, _, _ = read_video(gt["video_path"], start_pts=0, end_pts=None, pts_unit='sec')
+ video = video.numpy()
+ selected_frames = torch.linspace(0, video.shape[0] - 1, self._num_frames).long()
+ video_frames = video[selected_frames]
+ if self._num_frames == 1:
+ video_frames = video_frames[None]
+
+ imgs = []
+ for img in video_frames:
+ from torchvision.transforms import ToPILImage
+ to_pil = ToPILImage()
+ img = to_pil(img)
+ imgs += get_visual_transform(
+ img, self._img_h, self._img_w, self._use_tiling, self._max_num_tiles,
+ self._use_thumbnail, augment=False, vision_model_type=self._vision_model_type
+ )
+
+ for question in gt["questions"]:
+ # Very hacky, but we essentially re-create gt holding only the
+ # question of interest. This is the make this generation script
+ # compatible with the Video MME evaluation script.
+ question_dict = {
+ "video_id": gt["video_id"],
+ "duration_category": gt["duration_category"],
+ "video_category": gt["video_category"],
+ "video_subcategory": gt["video_subcategory"],
+ "url": gt["url"],
+ "questions": [question],
+ }
+
+ num_tiles = torch.tensor([len(imgs)], dtype=torch.int)
+
+ answer = ""
+ metadata = ""
+
+ return (
+ torch.stack(imgs),
+ num_tiles,
+ question["question_id"],
+ question_dict,
+ answer,
+ metadata,
+ )
+
+
+class OCRBenchDataset(torch.utils.data.Dataset):
+ """OCRBench evaluation dataset."""
+
+ def __init__(
+ self,
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ ):
+ gt = json.load(open(gt_path, encoding='utf-8'))
+
+ if num_partitions > 0:
+ start_idx, end_idx = _get_partition_bounds(
+ len(gt), num_samples_per_partition, num_partitions, partition_id
+ )
+ gt = gt[start_idx:end_idx]
+
+ self._input_image_path = input_image_path
+ self._gt = gt
+ self._img_h = img_h
+ self._img_w = img_w
+ self._use_tiling = use_tiling
+ self._max_num_tiles = max_num_tiles
+ self._use_thumbnail = use_thumbnail
+ self._vision_model_type = vision_model_type
+
+ def __len__(self):
+ return len(self._gt)
+
+ def __getitem__(self, idx):
+ img_path = os.path.join(self._input_image_path, self._gt[idx]['image_path'])
+
+ img = Image.open(img_path)
+ imgs = get_visual_transform(
+ img,
+ self._img_h,
+ self._img_w,
+ self._use_tiling,
+ self._max_num_tiles,
+ self._use_thumbnail,
+ augment=False,
+ vision_model_type=self._vision_model_type,
+ )
+
+ tile_count = torch.tensor([len(imgs)], dtype=torch.int)
+
+ metadata = {
+ "dataset_name": self._gt[idx]["dataset_name"],
+ "data_type": self._gt[idx]["type"],
+ }
+
+ return (
+ torch.stack(imgs),
+ tile_count,
+ idx,
+ self._gt[idx]["question"],
+ self._gt[idx]["answers"],
+ metadata,
+ )
+
+
+class MathVistaDataset(torch.utils.data.Dataset):
+ """MathVista evaluation dataset."""
+
+ def __init__(
+ self,
+ input_image_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ ):
+ import datasets
+
+ hf_datasets_cache = os.environ["HF_DATASETS_CACHE"]
+ assert hf_datasets_cache != "", "Please set the environment variable HF_DATASETS_CACHE."
+
+ if os.path.exists(input_image_path):
+ dataset = datasets.load_dataset(
+ input_image_path, cache_dir=hf_datasets_cache, verification_mode="no_checks"
+ )
+ else:
+ dataset = datasets.load_dataset(
+ "AI4Math/MathVista", split="testmini", cache_dir=hf_datasets_cache
+ )
+
+ if num_partitions > 0:
+ start_idx, end_idx = _get_partition_bounds(
+ len(dataset), num_samples_per_partition, num_partitions, partition_id
+ )
+ dataset = dataset[start_idx:end_idx]
+
+ self._dataset = dataset
+ self._img_h = img_h
+ self._img_w = img_w
+ self._use_tiling = use_tiling
+ self._max_num_tiles = max_num_tiles
+ self._use_thumbnail = use_thumbnail
+ self._vision_model_type = vision_model_type
+
+ def __len__(self):
+ return len(self._dataset["pid"])
+
+ def __getitem__(self, idx):
+ # Already a PIL object.
+ img = self._dataset['decoded_image'][idx]
+
+ imgs = get_visual_transform(
+ img,
+ self._img_h,
+ self._img_w,
+ self._use_tiling,
+ self._max_num_tiles,
+ self._use_thumbnail,
+ augment=False,
+ vision_model_type=self._vision_model_type,
+ )
+
+ tile_count = torch.tensor([len(imgs)], dtype=torch.int)
+
+ question_id = self._dataset["pid"][idx]
+ question = self._dataset["question"][idx]
+ question_type = self._dataset["question_type"][idx] # free_form or multi_choice
+ query = self._dataset["query"][idx]
+ choices = self._dataset["choices"][idx]
+ answer = self._dataset["answer"][idx]
+
+ if question_type == 'multi_choice':
+ start_chr = 'A'
+ choices_str = ''
+ index2ans = {}
+ all_choices = []
+ for choice in choices:
+ all_choices.append(start_chr)
+ index2ans[start_chr] = choice
+ choices_str += f"{start_chr}. {choice}\n"
+ start_chr = chr(ord(start_chr) + 1)
+
+ question = question + '\n' + choices_str
+ question = question + "Answer with the option's letter from the given choices directly."
+ answer = chr(ord('A') + choices.index(answer))
+ else:
+ question = query.replace("Hint: ", "")
+ index2ans = {}
+ all_choices = []
+
+ metadata = {
+ "question_type": question_type,
+ "index2ans": index2ans,
+ "all_choices": all_choices,
+ }
+
+ return torch.stack(imgs), tile_count, question_id, question, answer, metadata
+
+
+class AI2DDataset(torch.utils.data.Dataset):
+ """AI2D evaluation dataset."""
+
+ def __init__(
+ self,
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ no_mask,
+ vision_model_type,
+ ):
+ with open(gt_path, 'r') as f:
+ jsonl = list(f)
+
+ gt = [json.loads(json_str) for json_str in jsonl]
+
+ if num_partitions > 0:
+ start_idx, end_idx = _get_partition_bounds(
+ len(gt), num_samples_per_partition, num_partitions, partition_id
+ )
+ gt = gt[start_idx:end_idx]
+
+ self._gt = gt
+ self._input_image_path = input_image_path
+ self._img_h = img_h
+ self._img_w = img_w
+ self._use_tiling = use_tiling
+ self._max_num_tiles = max_num_tiles
+ self._use_thumbnail = use_thumbnail
+ self._no_mask = no_mask
+ self._vision_model_type = vision_model_type
+
+ def __len__(self):
+ return len(self._gt)
+
+ def __getitem__(self, idx):
+ img_path = os.path.join(self._input_image_path, self._gt[idx]['image'])
+ if self._no_mask:
+ img_path.replace("AI2D_TEST", "AI2D_TEST_NO_MASK_IMAGES")
+
+ img = Image.open(img_path)
+ imgs = get_visual_transform(
+ img,
+ self._img_h,
+ self._img_w,
+ self._use_tiling,
+ self._max_num_tiles,
+ self._use_thumbnail,
+ augment=False,
+ vision_model_type=self._vision_model_type,
+ )
+
+ tile_count = torch.tensor([len(imgs)], dtype=torch.int)
+
+ metadata = "" # Not used.
+
+ return (
+ torch.stack(imgs),
+ tile_count,
+ self._gt[idx]["question_id"],
+ self._gt[idx]["question"],
+ self._gt[idx]["answer"],
+ metadata,
+ )
+
+
+def get_evaluation_dataset(
+ task,
+ input_image_path,
+ gt_path,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ num_frames,
+ vision_model_type,
+):
+ """Get an evaluation dataset."""
+ if task == "TextVQA":
+ keys = {
+ "image_id": "image_id",
+ "sample_id": "question_id",
+ "question": "question",
+ "answer": "answers",
+ }
+
+ dataset = VQADataset(
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ keys,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ )
+ elif task == "VQAv2":
+ keys = {
+ "image_id": "image",
+ "sample_id": "question_id",
+ "question": "question",
+ "answer": "answer",
+ }
+
+ dataset = VQADataset(
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ keys,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ )
+ elif task == "ChartQA":
+ keys = {"image_id": "imgname", "question": "query", "answer": "label"}
+
+ dataset = VQADataset(
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ keys,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ )
+ elif task == "captioning":
+ dataset = CaptioningDataset(
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ )
+ elif task == 'MMMU':
+ # Note:
+ # - prompt_style="single_image" uses only one image like in the MMMU repo example.
+ # - prompt_style="multi_image" uses multiple input images.
+ # - prompt_style="vlmevalkit" is similar to https://github.com/open-compass/VLMEvalKit/blob/5d3cebcf18ef4bfbadc3bd3ef80bdc7aad2c6557/vlmeval/vlm/internvl_chat.py#L499
+ dataset = MMMUDataset(
+ input_image_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ prompt_style="single_image",
+ vision_model_type=vision_model_type,
+ )
+ elif task == "VideoMME":
+ dataset = VideoMMEDataset(
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ num_frames,
+ vision_model_type,
+ )
+ elif task == "OCRBench":
+ dataset = OCRBenchDataset(
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ )
+ elif task == "MathVista":
+ dataset = MathVistaDataset(
+ input_image_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ vision_model_type,
+ )
+ elif task == "AI2D":
+ dataset = AI2DDataset(
+ input_image_path,
+ gt_path,
+ num_samples_per_partition,
+ num_partitions,
+ partition_id,
+ img_h,
+ img_w,
+ use_tiling,
+ max_num_tiles,
+ use_thumbnail,
+ no_mask=False,
+ vision_model_type=vision_model_type,
+ )
+ else:
+ raise NotImplementedError(f"unsupported task {task}")
+
+ return dataset
diff --git a/examples/multimodal/image_processing.py b/examples/multimodal/image_processing.py
new file mode 100644
index 0000000000..15175a634e
--- /dev/null
+++ b/examples/multimodal/image_processing.py
@@ -0,0 +1,143 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. Except portions as noted which are Copyright (c) 2023 OpenGVLab and licensed under the MIT license found in LICENSE.
+from torchvision import transforms as T
+from torchvision.transforms import Compose
+from torchvision.transforms.functional import InterpolationMode
+
+
+IMAGENET_PIXEL_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_PIXEL_STD = [0.229, 0.224, 0.225]
+SIGLIP_PIXEL_MEAN = [0.5, 0.5, 0.5]
+SIGLIP_PIXEL_STD = [0.5, 0.5, 0.5]
+CLIP_PIXEL_MEAN = [0.48145466, 0.4578275, 0.40821073]
+CLIP_PIXEL_STD = [0.26862954, 0.26130258, 0.27577711]
+
+
+pixel_statistics = {
+ "clip": (CLIP_PIXEL_MEAN, CLIP_PIXEL_STD),
+ "siglip": (SIGLIP_PIXEL_MEAN, SIGLIP_PIXEL_STD),
+ "internvit": (IMAGENET_PIXEL_MEAN, IMAGENET_PIXEL_STD),
+ "radio": (CLIP_PIXEL_MEAN, CLIP_PIXEL_STD),
+ "huggingface": (SIGLIP_PIXEL_MEAN, SIGLIP_PIXEL_STD),
+}
+
+
+# From https://github.com/OpenGVLab/InternVL/blob/c62fa4f7c850165d7386bdc48ac6bc5a6fab0864/internvl_chat/internvl/train/dataset.py#L685
+# Copyright (c) 2023 OpenGVLab.
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+ best_ratio_diff = float('inf')
+ best_ratio = (1, 1)
+ area = width * height
+ for ratio in target_ratios:
+ target_aspect_ratio = ratio[0] / ratio[1]
+ ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+ if ratio_diff < best_ratio_diff:
+ best_ratio_diff = ratio_diff
+ best_ratio = ratio
+ elif ratio_diff == best_ratio_diff:
+ if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+ best_ratio = ratio
+ return best_ratio
+
+
+def find_closest_area_weighted_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+ """
+ Find the best number of tiles based on the aspect ratio and the area covered by the tiles.
+ """
+ best_factor = float('-inf')
+ best_ratio = (1, 1)
+ area = width * height
+ for ratio in target_ratios:
+ target_aspect_ratio = ratio[0] / ratio[1]
+ factor_based_on_area_n_ratio = (
+ min((ratio[0]*ratio[1]*image_size*image_size)/ area, 0.6) *
+ min(target_aspect_ratio/aspect_ratio, aspect_ratio/target_aspect_ratio))
+ if factor_based_on_area_n_ratio > best_factor:
+ best_factor = factor_based_on_area_n_ratio
+ best_ratio = ratio
+ return best_ratio
+
+
+def get_visual_transform(
+ img, img_h, img_w, use_tiling=False, max_num_tiles=1, use_thumbnail=False, augment=False,
+ vision_model_type="clip", find_closest_aspect_ratio_fn=find_closest_aspect_ratio):
+ pixel_mean, pixel_std = pixel_statistics[vision_model_type]
+
+ assert not augment, "Image augmentation not implemented."
+ transform = build_transform(img_h, pixel_mean, pixel_std, vision_model_type)
+
+ if use_tiling:
+ assert img_h == img_w, "dynamic tiling expects equal tile height and width"
+ imgs = dynamic_preprocess(
+ img, min_num=1, max_num=max_num_tiles, image_size=img_h, use_thumbnail=use_thumbnail,
+ find_closest_aspect_ratio_fn=find_closest_aspect_ratio_fn)
+ imgs = [transform(img) for img in imgs]
+ else:
+ imgs = [transform(img)]
+
+ return imgs
+
+
+# From https://github.com/OpenGVLab/InternVL/blob/c62fa4f7c850165d7386bdc48ac6bc5a6fab0864/internvl_chat/internvl/train/dataset.py#L702
+# Copyright (c) 2023 OpenGVLab.
+def dynamic_preprocess(
+ image, min_num=1, max_num=6, image_size=448, use_thumbnail=False,
+ find_closest_aspect_ratio_fn=find_closest_aspect_ratio):
+ orig_width, orig_height = image.size
+ aspect_ratio = orig_width / orig_height
+
+ # calculate the existing image aspect ratio
+ target_ratios = set(
+ (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
+ i * j <= max_num and i * j >= min_num)
+ target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+ # find the closest aspect ratio to the target
+ target_aspect_ratio = find_closest_aspect_ratio_fn(
+ aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+ # calculate the target width and height
+ target_width = image_size * target_aspect_ratio[0]
+ target_height = image_size * target_aspect_ratio[1]
+ blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+ # resize the image
+ resized_img = image.resize((target_width, target_height))
+ processed_images = []
+ for i in range(blocks):
+ box = (
+ (i % (target_width // image_size)) * image_size,
+ (i // (target_width // image_size)) * image_size,
+ ((i % (target_width // image_size)) + 1) * image_size,
+ ((i // (target_width // image_size)) + 1) * image_size
+ )
+ # split the image
+ split_img = resized_img.crop(box)
+ processed_images.append(split_img)
+ assert len(processed_images) == blocks
+ if use_thumbnail and len(processed_images) != 1:
+ thumbnail_img = image.resize((image_size, image_size))
+ processed_images.append(thumbnail_img)
+ return processed_images
+
+
+# Based on https://github.com/openai/CLIP/blob/dcba3cb2e2827b402d2701e7e1c7d9fed8a20ef1/clip/clip.py#L79
+# and https://github.com/OpenGVLab/InternVL/blob/aa521e6eb1df4cf153aa4118fcf13e673c055d46/internvl_chat/internvl/train/dataset.py#L276
+def build_transform(input_size, pixel_mean, pixel_std, vision_model_type):
+ if vision_model_type in ("siglip", "internvit", "radio", "huggingface"):
+ transform = T.Compose([
+ T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+ T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+ T.ToTensor(),
+ T.Normalize(mean=pixel_mean, std=pixel_std)
+ ])
+ elif vision_model_type == "clip":
+ transform = Compose([
+ T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+ T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+ T.ToTensor(),
+ T.Normalize(mean=pixel_mean, std=pixel_std),
+ ])
+ else:
+ raise NotImplementedError(f"image processing not defined for vision model {vision_model_type}")
+
+ return transform
diff --git a/examples/multimodal/layer_specs.py b/examples/multimodal/layer_specs.py
new file mode 100644
index 0000000000..dff1181d94
--- /dev/null
+++ b/examples/multimodal/layer_specs.py
@@ -0,0 +1,139 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import torch
+
+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 (
+ TEColumnParallelLinear,
+ TEDotProductAttention,
+ TELayerNormColumnParallelLinear,
+ TENorm,
+ TERowParallelLinear,
+ )
+
+ HAVE_TE = True
+except ImportError:
+ HAVE_TE = False
+
+try:
+ import apex
+
+ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
+ from megatron.core.transformer.torch_norm import WrappedTorchNorm
+
+ HAVE_APEX = True
+ LNImpl = FusedLayerNorm
+except ImportError:
+ import warnings
+
+ from megatron.core.transformer.torch_norm import WrappedTorchNorm
+
+ warnings.warn(f'Apex is not installed. Falling back to Torch Norm')
+ LNImpl = WrappedTorchNorm
+
+
+def get_layer_spec(is_vit, normalization) -> ModuleSpec:
+ attn_mask_type = AttnMaskType.no_mask if is_vit else AttnMaskType.causal
+ if normalization == "LayerNorm":
+ norm = LNImpl
+ elif normalization == "RMSNorm":
+ if HAVE_TE:
+ norm = TENorm
+ else:
+ version = torch.__version__.split('.')
+ version_geq_2_4 = (
+ int(TORCH_VERSION[0]) > 2
+ or (
+ int(TORCH_VERSION[0]) == 2
+ and int(TORCH_VERSION[1]) >= 4
+ )
+ )
+ assert version_geq_2_4, "Torch version >= 2.4.0 is required for RMSNorm"
+ if HAVE_APEX:
+ warnings.warn(f'Apex does not support RMSNorm. Falling back to Torch Norm')
+ norm = WrappedTorchNorm
+ else:
+ raise RuntimeError("unknown normalization", normalization)
+
+ mlp = get_mlp_module_spec(use_te=False) # doesn't include norm.
+
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=norm,
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": attn_mask_type},
+ 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=norm,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+
+
+def get_layer_spec_te(is_vit=False, padding=False) -> ModuleSpec:
+ attn_mask_type = AttnMaskType.no_mask if is_vit else AttnMaskType.causal
+ # Padding mask is needed for e.g. Context Parallel.
+ if padding:
+ assert not is_vit, "padding_causal mask not used with ViT"
+ attn_mask_type = AttnMaskType.padding_causal
+
+ mlp = get_norm_mlp_module_spec_te()
+ return ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": attn_mask_type},
+ 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_mlp_layernorm=IdentityOp,
+ mlp=mlp,
+ mlp_bda=get_bias_dropout_add,
+ ),
+ )
+
+
+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=TEColumnParallelLinear if use_te else ColumnParallelLinear,
+ linear_fc2=TERowParallelLinear if use_te else RowParallelLinear,
+ ),
+ )
+
+
+def get_norm_mlp_module_spec_te() -> ModuleSpec:
+ return ModuleSpec(
+ module=MLP,
+ submodules=MLPSubmodules(
+ linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear
+ ),
+ )
diff --git a/examples/multimodal/manual_prompts.json b/examples/multimodal/manual_prompts.json
new file mode 100644
index 0000000000..b0dfd84801
--- /dev/null
+++ b/examples/multimodal/manual_prompts.json
@@ -0,0 +1,48 @@
+{
+ "COMMENT": "Sources for these prompts include https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain/viewer and https://huggingface.co/datasets/HuggingFaceM4/M3IT",
+ "Captioning": {
+ "raw": [
+ "Can you briefly explain what you see in the image?",
+ "Describe what's happening in this image in one short sentence.",
+ "Write a short caption that accurately represents the content of this image.",
+ "Please generate a descriptive caption for the image provided.",
+ "How would you summarize the scene depicted in the picture in short?",
+ "Describe the image briefly.",
+ "Write a succinct description of the image, capturing its main components, the relationships between them, and any notable details.",
+ "Create a concise caption that accurately describes the main elements in the image provided.",
+ "Write a brief, yet comprehensive, description of the image.",
+ "Describe the image in a clear and concise manner.",
+ "For the given image, provide a one-sentence summary that captures the most important details.",
+ "Generate a short caption for the picture.",
+ "Write a short and informative description that highlights the primary subjects and actions occurring in the given image.",
+ "Provide a concise and informative caption for the image, focusing on the primary subjects.",
+ "Write a clear description of the image, make sure the key features are well covered.",
+ "Offer a succinct explanation of the picture presented."
+ ]
+ },
+ "CaptioningPretraining": {
+ "raw": [
+ "Generate a short caption of the image.",
+ "Describe the image concisely.",
+ "Provide a brief description of the given image."
+ ],
+ "llava": [
+ "Give a brief description of image.",
+ "Give a brief description of the image.",
+ "Provide a brief description of the given image.",
+ "Provide a one-sentence caption for the provided image.",
+ "Write a terse but informative summary of the picture.",
+ "Describe the image concisely.",
+ "Generate a clear and concise summary of the photo."
+ ]
+ },
+ "OCR": {
+ "raw": [
+ "Can you read the text from image and output here?",
+ "Extract and document the text from the provided image.",
+ "Converting the text embedded in this image into a readable document.",
+ "Transcribe all the text you find.",
+ "Can you extract all visible text from the image here?"
+ ]
+ }
+}
diff --git a/examples/multimodal/model.py b/examples/multimodal/model.py
new file mode 100644
index 0000000000..1a82b12037
--- /dev/null
+++ b/examples/multimodal/model.py
@@ -0,0 +1,254 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import warnings
+from copy import deepcopy
+
+import torch
+from config import get_language_model_config, get_vision_model_config, get_vision_projection_config
+from layer_specs import get_layer_spec, get_layer_spec_te, get_mlp_module_spec, get_norm_mlp_module_spec_te
+
+from megatron.core.models.multimodal.llava_model import IMAGE_TOKEN, LLaVAModel
+from megatron.core.models.vision.clip_vit_model import get_num_image_embeddings
+from megatron.training import get_args, get_tokenizer, print_rank_0
+from megatron.training.arguments import core_transformer_config_from_args
+
+
+def model_provider(
+ pre_process=True, post_process=True, add_encoder=True, add_decoder=True, parallel_output=True
+) -> LLaVAModel:
+ """Builds the model.
+
+ Args:
+ pre_process (bool): Include the embedding layer in the gpt decoder (used with pipeline parallelism). Defaults to True.
+ post_process (bool): Include an output layer and a layernorm in the gpt decoder (used with pipeline parallelism). Defaults to True.
+ add_encoder (bool): Construct the encoder module (used with pipeline parallelism). Defaults to True. When we use pipelining, the encoder
+ will live on only a subset of the pipeline stages (specifically, only the first stage).
+ add_decoder (bool): Construct the decoder module (used with pipeline parallelism). Defaults to True. When we use pipelining, the decoder
+ will live on only a subset of the pipeline stages (specifically, every stage after the first one).
+ parallel_output (bool): Enable parallel model output.
+
+ Returns:
+ model: A multimodal model.
+ """
+ args = get_args()
+ assert args.encoder_pipeline_model_parallel_size <= 1, "LLaVA does not support pp>1 for encoder on it's own pipeline rank"
+
+ use_te = args.use_te
+
+ print_rank_0('building a multimodal model ...')
+
+ num_image_embeddings = get_num_image_embeddings(
+ args.img_h,
+ args.img_w,
+ args.patch_dim,
+ args.vision_model_type,
+ args.disable_vision_class_token,
+ 1,
+ args.pixel_shuffle,
+ args.use_tile_tags,
+ )
+ old_seq_length = args.seq_length
+ args.seq_length = args.encoder_seq_length = num_image_embeddings
+ if torch.distributed.get_rank() == 0 and old_seq_length != args.seq_length:
+ warnings.warn(
+ f"Changed seq_length and encoder_seq_length (vision model sequence length) from {old_seq_length} to num_image_tokens ({num_image_embeddings})"
+ )
+
+ max_num_image_embeddings = (args.max_num_tiles + int(args.use_thumbnail)) * num_image_embeddings
+
+ assert (
+ args.decoder_seq_length is not None
+ ), "Please provide --decoder-seq-length to set the language model sequence length"
+ assert (
+ args.decoder_seq_length > max_num_image_embeddings
+ ), "Language model sequence length must be greater than the maximum number of image embeddings"
+ if args.decoder_seq_length > args.max_position_embeddings:
+ args.max_position_embeddings = args.decoder_seq_length
+ warnings.warn(
+ f"Expanded max_position_embeddings to {args.max_position_embeddings} to accommodate the maximum language model sequence length"
+ )
+
+ base_config = core_transformer_config_from_args(get_args())
+ base_config.language_model_type = args.language_model_type
+ base_config.vision_model_type = args.vision_model_type
+ base_config.calculate_per_token_loss = True
+
+ language_config = deepcopy(base_config)
+ language_config = get_language_model_config(language_config)
+
+ if use_te:
+ # Padding mask needed for SP/CP.
+ padding = args.context_parallel_size > 1 and args.sequence_parallel
+ language_transformer_layer_spec = get_layer_spec_te(
+ is_vit=False, padding=padding
+ ) # TENorm detects LayerNorm/RMS automatically.
+ else:
+ language_transformer_layer_spec = get_layer_spec(
+ is_vit=False, normalization=language_config.normalization
+ )
+
+ vision_model_type = args.vision_model_type
+ vision_config = deepcopy(base_config)
+ vision_config = get_vision_model_config(
+ vision_config, apply_query_key_layer_scaling=args.apply_query_key_layer_scaling
+ )
+ if vision_model_type.startswith("huggingface"):
+ assert args.encoder_tensor_model_parallel_size < 2, "Huggingface vision encoders do not support --encoder-tensor-model-parallel-size > 1"
+ assert args.encoder_pipeline_model_parallel_size == 0, "Huggingface vision encoders do not support --encoder-pipeline-model-parallel-size > 0"
+ assert not args.sequence_parallel, "Huggingface models do not support --sequence-parallel"
+ assert args.context_parallel_size < 2, "Huggingface models do not support --context-parallel-size > 1"
+ assert args.vision_huggingface_model_name_or_path is not None, "Providing --vision-huggingface-model-name-or-path is necessary when using huggingface vision model"
+
+ vision_config.huggingface_model_name_or_path = args.vision_huggingface_model_name_or_path
+
+ from transformers import AutoConfig
+ huggingface_config = AutoConfig.from_pretrained(vision_config.huggingface_model_name_or_path)
+ vision_config.hidden_size = huggingface_config.hidden_size
+
+ vision_model_type = args.vision_model_type
+ if vision_model_type in ["clip", "siglip", "radio"]:
+ if use_te:
+ vision_transformer_layer_spec = get_layer_spec_te(
+ is_vit=True
+ ) # TENorm detects LayerNorm/RMS automatically.
+ else:
+ vision_transformer_layer_spec = get_layer_spec(
+ is_vit=True, normalization=vision_config.normalization
+ )
+ elif vision_model_type == "internvit":
+ from nvlm.internvit import get_internvit_layer_spec
+ vision_transformer_layer_spec = get_internvit_layer_spec(use_te=use_te)
+ elif vision_model_type.startswith("huggingface"):
+ vision_transformer_layer_spec = None
+ else:
+ raise RuntimeError("unsupported vision model type", vision_model_type)
+
+ vision_projection_config = deepcopy(base_config)
+
+ if base_config.language_model_type.startswith("huggingface"):
+ assert args.tensor_model_parallel_size == 1, "Huggingface models do not support --tensor-model-parallel-size > 1"
+ assert args.pipeline_model_parallel_size < 2, "Huggingface models do not support --pipeline-model-parallel-size > 1"
+ assert not args.sequence_parallel, "Huggingface models do not support --sequence-parallel"
+ assert args.context_parallel_size < 2, "Huggingface models do not support --context-parallel-size > 1"
+ assert args.language_huggingface_model_name_or_path is not None, "Providing --language-huggingface-model-name-or-path is necessary when using huggingface language model"
+
+ language_config.huggingface_model_name_or_path = args.language_huggingface_model_name_or_path
+ # Pass to vision projection config so can choose the correct ffn hidden size
+ vision_projection_config.huggingface_model_name_or_path = args.language_huggingface_model_name_or_path
+
+ vision_projection_config = get_vision_projection_config(
+ vision_projection_config, language_config.hidden_size
+ )
+
+ # --encoder-pipeline-model-parallel-size 1 will enable a separate pipeline stage for the vision model.
+ if args.encoder_pipeline_model_parallel_size > 0:
+ assert (
+ args.encoder_pipeline_model_parallel_size == 1
+ ), "vision model and projection can only live on 1 pipeline stage."
+
+ if args.encoder_tensor_model_parallel_size > 0:
+ vision_config.tensor_model_parallel_size = args.encoder_tensor_model_parallel_size
+ vision_projection_config.tensor_model_parallel_size = (
+ args.encoder_tensor_model_parallel_size
+ )
+
+ # Make sure vision model pipeline parallel size is not inherited from the language model pipeline parallel size.
+ # 0 is not a valid for the config value, hence max(1, ).
+ vision_config.pipeline_model_parallel_size = max(1, args.encoder_pipeline_model_parallel_size)
+ vision_projection_config.pipeline_model_parallel_size = vision_config.pipeline_model_parallel_size
+
+ # Make sure the vision model does not inherit first and last pipeline num layers from the language model.
+ vision_config.first_pipeline_num_layers = vision_config.last_pipeline_num_layers = None
+
+ if vision_projection_config.normalization:
+ vision_projection_layer_spec = get_norm_mlp_module_spec_te().submodules
+ else:
+ vision_projection_layer_spec = get_mlp_module_spec(use_te=use_te).submodules
+
+ # Toggle --recompute* for the vision and language model separately.
+ if args.recompute_vision:
+ if vision_config.recompute_method is not None and vision_config.recompute_granularity is not None:
+ vision_config.recompute_num_layers = vision_config.num_layers
+ else:
+ vision_config.recompute_granularity = None
+ vision_config.recompute_method = None
+ vision_config.recompute_num_layers = None
+
+ vision_projection_config.recompute_granularity = None
+ vision_projection_config.recompute_method = None
+ vision_projection_config.recompute_num_layers = None
+
+ # TODO: Vision model and projection do not use SP/CP yet.
+ vision_config.sequence_parallel = False
+ vision_config.context_parallel_size = 1
+ vision_config.tp_comm_overlap = False
+
+ vision_projection_config.sequence_parallel = False
+ vision_projection_config.context_parallel_size = 1
+ vision_projection_config.tp_comm_overlap = False
+
+ tokenizer = get_tokenizer()
+ image_token_index = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
+ assert image_token_index is not None, f"IMAGE_TOKEN={IMAGE_TOKEN} needs to be added using the --special-tokens arg."
+
+ tile_tags = _get_tile_tags(args, tokenizer)
+
+ model = LLaVAModel(
+ language_transformer_config=language_config,
+ language_transformer_layer_spec=language_transformer_layer_spec,
+ language_vocab_size=args.padded_vocab_size,
+ language_max_sequence_length=args.decoder_seq_length,
+ vision_transformer_config=vision_config,
+ vision_transformer_layer_spec=vision_transformer_layer_spec,
+ drop_vision_class_token=args.disable_vision_class_token,
+ vision_projection_config=vision_projection_config,
+ vision_projection_layer_spec=vision_projection_layer_spec,
+ vision_projection_type="mlp",
+ allow_missing_vision_projection_checkpoint=args.allow_missing_vision_projection_checkpoint,
+ parallel_output=parallel_output,
+ share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+ language_position_embedding_type=args.position_embedding_type,
+ language_rotary_percent=args.rotary_percent,
+ pre_process=pre_process,
+ post_process=post_process,
+ add_encoder=add_encoder,
+ add_decoder=add_decoder,
+ img_h=args.img_h,
+ img_w=args.img_w,
+ patch_dim=args.patch_dim,
+ language_rotary_base=args.rotary_base,
+ language_rope_scaling=args.use_rope_scaling,
+ image_token_index=image_token_index,
+ pixel_shuffle=args.pixel_shuffle,
+ tile_tags=tile_tags,
+ )
+
+ model.freeze(
+ freeze_language_model=args.freeze_LM,
+ freeze_vision_model=args.freeze_ViT,
+ freeze_vision_projection=False,
+ )
+
+ return model
+
+
+def _get_tile_tags(args, tokenizer):
+ """Tile tags are used in NVLM to surround image tiles with text tags."""
+ if not args.use_tile_tags:
+ return None
+
+ # We expect the tokenized length of the tags is same.
+ thumbnail_tag_text = ""
+ if args.tokenizer_prompt_format == "nvlm-yi-34b":
+ thumbnail_tag_text = ""
+
+ assert args.max_num_tiles <= 6, "Up to 6 tile tags used"
+ tile_tags_text = [f"" for i in range(1, args.max_num_tiles + 1)] + [thumbnail_tag_text]
+
+ start_idx = 0
+ if tokenizer._prompt_config.has_bos:
+ start_idx = 1
+
+ # Convert to tokens [num_tiles, tile_seq_len].
+ tile_tags = [tokenizer.tokenize(t)[start_idx:] for t in tile_tags_text]
+
+ return tile_tags
diff --git a/examples/multimodal/model_converter/clip_converter.py b/examples/multimodal/model_converter/clip_converter.py
new file mode 100644
index 0000000000..696c810890
--- /dev/null
+++ b/examples/multimodal/model_converter/clip_converter.py
@@ -0,0 +1,163 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import argparse
+import os
+
+import torch
+
+import clip
+
+
+def convert(download_root, output_path, tensor_parallel_size, use_te):
+ device = "cuda"
+
+ model, _ = clip.load("ViT-L/14@336px", device=device, download_root=download_root)
+
+ state_dict = model.state_dict()
+ new_state_dicts = [{"model": dict()} for _ in range(tensor_parallel_size)]
+
+ # Indices from mapping pytorch multihead attention to megatron.
+ kv_channels = 64
+ hidden_dim = 1024
+ num_heads = 16
+ indices = []
+ for i in range(num_heads):
+ lb = i * kv_channels
+ ub = (i + 1) * kv_channels
+ indices.append(torch.arange(lb, ub, dtype=torch.int))
+ indices.append(torch.arange(hidden_dim + lb, hidden_dim + ub, dtype=torch.int))
+ indices.append(torch.arange(2 * hidden_dim + lb, 2 * hidden_dim + ub, dtype=torch.int))
+
+ indices = torch.cat(indices)
+
+ for name, tensor in state_dict.items():
+ # Skip text model.
+ if "visual" not in name:
+ continue
+
+ # Skip final layers not used in our model.
+ if name == "visual.proj" or "ln_post" in name:
+ continue
+
+ # Map parameter names to ones used in megatron.
+ new_name = ""
+ new_tensor = tensor
+ if new_tensor.dtype == torch.float16:
+ new_tensor = new_tensor.to(torch.float32)
+
+ # This is used for chunking some tensors to target tensor parallel size.
+ chunk_dim = None
+
+ if "class_embedding" in name:
+ new_name = "class_token"
+ # Our model uses class token that is expanded to input dimensions already.
+ new_tensor = new_tensor.expand(1, 1, -1)
+ elif "positional_embedding" in name:
+ new_name = "position_embeddings.weight"
+ elif "conv1" in name:
+ new_name = "conv1.weight"
+ elif "ln_pre.weight" in name:
+ new_name = "ln_pre.weight"
+ elif "ln_pre.bias" in name:
+ new_name = "ln_pre.bias"
+ elif "transformer.resblocks" in name:
+ layer_idx = name.split(".")[3]
+ base = f"decoder.layers.{layer_idx}"
+
+ if "attn.in_proj_weight" in name:
+ new_name = f"{base}.self_attention.linear_qkv.weight"
+ new_tensor = new_tensor[indices]
+ chunk_dim = 0
+ elif "attn.in_proj_bias" in name:
+ new_name = f"{base}.self_attention.linear_qkv.bias"
+ new_tensor = new_tensor[indices]
+ chunk_dim = 0
+ elif "attn.out_proj.weight" in name:
+ new_name = f"{base}.self_attention.linear_proj.weight"
+ chunk_dim = 1
+ elif "attn.out_proj.bias" in name:
+ new_name = f"{base}.self_attention.linear_proj.bias"
+ elif "ln_1.weight" in name:
+ new_name = f"{base}.input_layernorm.weight"
+ if use_te:
+ new_name = f"{base}.self_attention.linear_qkv.layer_norm_weight"
+ elif "ln_1.bias" in name:
+ new_name = f"{base}.input_layernorm.bias"
+ if use_te:
+ new_name = f"{base}.self_attention.linear_qkv.layer_norm_bias"
+ elif "mlp.c_fc.weight" in name:
+ new_name = f"{base}.mlp.linear_fc1.weight"
+ chunk_dim = 0
+ elif "mlp.c_fc.bias" in name:
+ new_name = f"{base}.mlp.linear_fc1.bias"
+ chunk_dim = 0
+ elif "mlp.c_proj.weight" in name:
+ new_name = f"{base}.mlp.linear_fc2.weight"
+ chunk_dim = 1
+ elif "mlp.c_proj.bias" in name:
+ new_name = f"{base}.mlp.linear_fc2.bias"
+ elif "ln_2.weight" in name:
+ new_name = f"{base}.pre_mlp_layernorm.weight"
+ if use_te:
+ new_name = f"{base}.mlp.linear_fc1.layer_norm_weight"
+ elif "ln_2.bias" in name:
+ new_name = f"{base}.pre_mlp_layernorm.bias"
+ if use_te:
+ new_name = f"{base}.mlp.linear_fc1.layer_norm_bias"
+
+ assert new_name != "", f"unexpected layer name {name}"
+
+ if chunk_dim is None:
+ new_tensors = [new_tensor for _ in range(tensor_parallel_size)]
+ else:
+ new_tensors = torch.chunk(new_tensor, tensor_parallel_size, dim=chunk_dim)
+
+ for i in range(tensor_parallel_size):
+ # chunk() creates a view of a bigger tensor. clone() is used here to avoid excessive storage.
+ new_state_dicts[i]["model"][new_name] = new_tensors[i].clone()
+
+ # TE sets _extra_state (for FP8 purposes), so set an empty one here for compatibility.
+ extra_state_layers = ("linear_qkv", "linear_proj", "linear_fc1", "linear_fc2")
+ is_extra_state_layer = any([l in new_name for l in extra_state_layers])
+ if use_te and is_extra_state_layer:
+ layer = new_name.split(".")[-2]
+ if layer in extra_state_layers:
+ extra_state_name = (
+ new_name[: new_name.rfind(".") + 1] + "_extra_state"
+ ) # Replace the weight name.
+ new_state_dicts[i]["model"][extra_state_name] = None
+
+ for i in range(tensor_parallel_size):
+ output_dir_tp = os.path.join(output_path, "iter_0000001", f"mp_rank_0{i}")
+ os.makedirs(output_dir_tp)
+ output_path_tp = os.path.join(output_dir_tp, "model_optim_rng.pt")
+ torch.save(new_state_dicts[i], output_path_tp)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description="""
+Convert OpenAI CLIP VIT weights to megatron format.
+
+
+Example usage:
+python clip_converter.py --download-root /some/download/folder --output /some/output/folder --tensor-parallel-size 4
+""",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ )
+
+ parser.add_argument(
+ "--download-root", type=str, required=True, help="Download folder for OpenAI CLIP weights"
+ )
+ parser.add_argument(
+ "--output", type=str, required=True, help="output directory for megatron state dict file(s)"
+ )
+ parser.add_argument(
+ "--tensor-parallel-size", type=int, default=1, help="model tensor parallel size"
+ )
+ parser.add_argument("--use-te", action="store_true", help="Use Transformer Engine")
+
+ args = parser.parse_args()
+
+ convert(args.download_root, args.output, args.tensor_parallel_size, args.use_te)
+
+ print("done.")
diff --git a/examples/multimodal/model_converter/internvit_converter.py b/examples/multimodal/model_converter/internvit_converter.py
new file mode 100755
index 0000000000..48404c2084
--- /dev/null
+++ b/examples/multimodal/model_converter/internvit_converter.py
@@ -0,0 +1,162 @@
+import argparse
+import os
+
+import torch
+from transformers import AutoModel
+
+
+def convert(model_name, output_path, tensor_parallel_size, use_te):
+ """Convert InternViT HF checkpoint to mcore."""
+ hf_model = AutoModel.from_pretrained(
+ model_name,
+ trust_remote_code=True
+ )
+
+ hf_state_dict = hf_model.state_dict()
+ new_state_dicts = [{"model": dict()} for _ in range(tensor_parallel_size)]
+
+ hidden_size = 3200
+ num_heads = 25
+ dim = 128
+
+ order = torch.ones(3 * hidden_size).long()
+
+ for j in range(num_heads):
+ for i in range(dim):
+ order[i + dim*3*j] = j*dim+i
+ order[dim + i + dim*3*j] = j*dim+i+num_heads*dim
+ order[dim*2 + i + dim*3*j] = j*dim+i+num_heads*dim*2
+
+ for name, tensor in hf_state_dict.items():
+ # Map parameter names to ones used in megatron.
+ new_name = ""
+ new_tensor = tensor
+
+ # This is used for chunking some tensors to target tensor parallel size.
+ chunk_dim = None
+
+ if "embeddings.class_embedding" in name:
+ new_name = "class_token"
+ elif "embeddings.patch_embedding.weight" in name:
+ new_name = "conv1.weight"
+ elif "embeddings.patch_embedding.bias" in name:
+ new_name = "conv1.bias"
+ elif "embeddings.position_embedding" in name:
+ new_name = "position_embeddings.weight"
+ new_tensor = new_tensor.squeeze(0)
+ elif "encoder.layers" in name:
+ layer_idx = name.split(".")[2]
+
+ base = f"decoder.layers.{layer_idx}"
+
+ head_dim = 128
+
+ if tensor_parallel_size == 1:
+ num_padded_heads = 25
+ elif tensor_parallel_size == 8:
+ # Note: 25 is not divisible by 8 and we don't currently support uneven heads split with tensor parallelism.
+ # So we pad with dummy all-zero heads. Please use a nice even number of attention heads in your model.
+ num_padded_heads = 32
+ else:
+ raise NotImplementedError("invalid tensor parallel size value:", tensor_parallel_size)
+
+ if "ls1" in name:
+ new_name = f"{base}.ls1"
+ elif "ls2" in name:
+ new_name = f"{base}.ls2"
+ elif "attn.qkv.weight" in name:
+ new_name = f"{base}.self_attention.linear_qkv.weight"
+ num_tensors = 3
+ padded_dim = head_dim * num_padded_heads * num_tensors
+ padded_tensor = torch.zeros((padded_dim, new_tensor.shape[-1]), dtype=new_tensor.dtype, device=new_tensor.device)
+ padded_tensor[:new_tensor.shape[0], :] = new_tensor[order]
+ new_tensor = padded_tensor
+ chunk_dim = 0
+ elif "attn.q_norm.weight" in name:
+ new_name = f"{base}.self_attention.q_layernorm.weight"
+ num_tensors = 1
+ padded_dim = head_dim * num_padded_heads * num_tensors
+ padded_tensor = torch.zeros(padded_dim, dtype=new_tensor.dtype, device=new_tensor.device)
+ padded_tensor[:new_tensor.shape[0]] = new_tensor
+ new_tensor = padded_tensor
+ chunk_dim = 0
+ elif "attn.k_norm.weight" in name:
+ new_name = f"{base}.self_attention.k_layernorm.weight"
+ num_tensors = 1
+ padded_dim = head_dim * num_padded_heads * num_tensors
+ padded_tensor = torch.zeros(padded_dim, dtype=new_tensor.dtype, device=new_tensor.device)
+ padded_tensor[:new_tensor.shape[0]] = new_tensor
+ new_tensor = padded_tensor
+ chunk_dim = 0
+ elif "attn.proj.weight" in name:
+ new_name = f"{base}.self_attention.linear_proj.weight"
+ num_tensors = 1
+ padded_dim = head_dim * num_padded_heads * num_tensors
+ padded_tensor = torch.zeros((new_tensor.shape[0], padded_dim), dtype=new_tensor.dtype, device=new_tensor.device)
+ padded_tensor[:, :new_tensor.shape[-1]] = new_tensor
+ new_tensor = padded_tensor
+ chunk_dim = 1
+ elif "attn.proj.bias" in name:
+ new_name = f"{base}.self_attention.linear_proj.bias"
+ elif "mlp.fc1.weight" in name:
+ new_name = f"{base}.mlp.linear_fc1.weight"
+ chunk_dim = 0
+ elif "mlp.fc1.bias" in name:
+ new_name = f"{base}.mlp.linear_fc1.bias"
+ chunk_dim = 0
+ elif "mlp.fc2.weight" in name:
+ new_name = f"{base}.mlp.linear_fc2.weight"
+ chunk_dim = 1
+ elif "mlp.fc2.bias" in name:
+ new_name = f"{base}.mlp.linear_fc2.bias"
+ elif "norm1" in name:
+ new_name = f"{base}.input_layernorm.weight"
+ elif "norm2" in name:
+ new_name = f"{base}.pre_mlp_layernorm.weight"
+ else:
+ raise RuntimeError("unexpected transformer layer name", name)
+ else:
+ raise RuntimeError("unexpected layer name", name)
+
+ assert new_name != "", f"unexpected layer name {name}"
+
+ # TE sets _extra_state (for FP8 purposes), so set an empty one here for compatibility.
+ extra_state_layers = ("linear_qkv", "linear_proj", "linear_fc1", "linear_fc2")
+ is_extra_state_layer = any([l in new_name for l in extra_state_layers])
+ if use_te and is_extra_state_layer:
+ layer = new_name.split(".")[-2]
+ if layer in extra_state_layers:
+ extra_state_name = (
+ new_name[: new_name.rfind(".") + 1] + "_extra_state"
+ ) # Replace the weight name.
+ for i in range(tensor_parallel_size):
+ new_state_dicts[i]["model"][extra_state_name] = None
+
+ if chunk_dim is None:
+ new_tensors = [new_tensor for _ in range(tensor_parallel_size)]
+ else:
+ new_tensors = torch.chunk(new_tensor, tensor_parallel_size, dim=chunk_dim)
+
+ for i in range(tensor_parallel_size):
+ new_state_dicts[i]["model"][new_name] = new_tensors[i].clone()
+
+ for i in range(tensor_parallel_size):
+ output_dir_tp = os.path.join(output_path, f"iter_0000001/mp_rank_0{i}")
+ os.makedirs(output_dir_tp, exist_ok=True)
+ output_path_tp = os.path.join(output_dir_tp, "model_optim_rng.pt")
+ torch.save(new_state_dicts[i], output_path_tp)
+ print("saved file", output_path_tp)
+
+ print("done")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description="InternVIT HuggingFace to Mcore converter")
+ parser.add_argument("--model-name", type=str, default="OpenGVLab/InternViT-6B-448px-V1-5", help="Model name in HuggingFace")
+ parser.add_argument("--output-dir", type=str, required=True, help="Output directory for the mcore model.")
+ parser.add_argument("--use-te", action="store_true", default=True)
+ parser.add_argument("--tensor-parallel-size", type=int, required=True)
+
+ args = parser.parse_args()
+
+ convert(args.model_name, args.output_dir, args.tensor_parallel_size, args.use_te)
diff --git a/examples/multimodal/model_converter/radio_converter.py b/examples/multimodal/model_converter/radio_converter.py
new file mode 100644
index 0000000000..05750a66eb
--- /dev/null
+++ b/examples/multimodal/model_converter/radio_converter.py
@@ -0,0 +1,152 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import argparse
+import os
+
+import torch
+
+def convert(output_path, tensor_parallel_size, use_te, version):
+ device = "cuda"
+
+ model = torch.hub.load('NVlabs/RADIO', 'radio_model', version=version, progress=True)
+
+ state_dict = model.state_dict()
+ new_state_dicts = [{"model": dict()} for _ in range(tensor_parallel_size)]
+
+ # Indices from mapping pytorch multihead attention to megatron.
+ kv_channels = 80
+ hidden_dim = 1280
+ num_heads = 16
+ indices = []
+ for i in range(num_heads):
+ lb = i * kv_channels
+ ub = (i + 1) * kv_channels
+ indices.append(torch.arange(lb, ub, dtype=torch.int))
+ indices.append(torch.arange(hidden_dim + lb, hidden_dim + ub, dtype=torch.int))
+ indices.append(torch.arange(2 * hidden_dim + lb, 2 * hidden_dim + ub, dtype=torch.int))
+
+ indices = torch.cat(indices)
+
+ for name, tensor in state_dict.items():
+ # Map parameter names to ones used in megatron.
+ new_name = ""
+ new_tensor = tensor
+ if new_tensor.dtype == torch.float16:
+ new_tensor = new_tensor.to(torch.float32)
+
+ # This is used for chunking some tensors to target tensor parallel size.
+ chunk_dim = None
+
+ if "summary_idxs" in name:
+ continue
+ elif "patch_generator" in name:
+ if "embedder" in name:
+ new_name = "embedder.weight"
+ chunk_dim = 0
+ elif "cls_token" in name:
+ new_name = "class_token"
+ elif "pos_embed" in name:
+ new_name = "position_embeddings"
+ elif "input_conditioner" in name:
+ continue
+ elif "blocks" in name:
+ layer_idx = name.split(".")[2]
+ base = f"decoder.layers.{layer_idx}"
+
+ if "attn.qkv.weight" in name:
+ new_name = f"{base}.self_attention.linear_qkv.weight"
+ new_tensor = new_tensor[indices]
+ chunk_dim = 0
+ elif "attn.qkv.bias" in name:
+ new_name = f"{base}.self_attention.linear_qkv.bias"
+ new_tensor = new_tensor[indices]
+ chunk_dim = 0
+ elif "attn.proj.weight" in name:
+ new_name = f"{base}.self_attention.linear_proj.weight"
+ chunk_dim = 1
+ elif "attn.proj.bias" in name:
+ new_name = f"{base}.self_attention.linear_proj.bias"
+ elif "norm1.weight" in name:
+ new_name = f"{base}.input_layernorm.weight"
+ if use_te:
+ new_name = f"{base}.self_attention.linear_qkv.layer_norm_weight"
+ elif "norm1.bias" in name:
+ new_name = f"{base}.input_layernorm.bias"
+ if use_te:
+ new_name = f"{base}.self_attention.linear_qkv.layer_norm_bias"
+ elif "mlp.fc1.weight" in name:
+ new_name = f"{base}.mlp.linear_fc1.weight"
+ chunk_dim = 0
+ elif "mlp.fc1.bias" in name:
+ new_name = f"{base}.mlp.linear_fc1.bias"
+ chunk_dim = 0
+ elif "mlp.fc2.weight" in name:
+ new_name = f"{base}.mlp.linear_fc2.weight"
+ chunk_dim = 1
+ elif "mlp.fc2.bias" in name:
+ new_name = f"{base}.mlp.linear_fc2.bias"
+ elif "norm2.weight" in name:
+ new_name = f"{base}.pre_mlp_layernorm.weight"
+ if use_te:
+ new_name = f"{base}.mlp.linear_fc1.layer_norm_weight"
+ elif "norm2.bias" in name:
+ new_name = f"{base}.pre_mlp_layernorm.bias"
+ if use_te:
+ new_name = f"{base}.mlp.linear_fc1.layer_norm_bias"
+
+ assert new_name != "", f"unexpected layer name {name}"
+
+ if chunk_dim is None:
+ new_tensors = [new_tensor for _ in range(tensor_parallel_size)]
+ else:
+ new_tensors = torch.chunk(new_tensor, tensor_parallel_size, dim=chunk_dim)
+
+ for i in range(tensor_parallel_size):
+ # chunk() creates a view of a bigger tensor. clone() is used here to avoid excessive storage.
+ new_state_dicts[i]["model"][new_name] = new_tensors[i].clone()
+
+ # TE sets _extra_state (for FP8 purposes), so set an empty one here for compatibility.
+ extra_state_layers = ("linear_qkv", "linear_proj", "linear_fc1", "linear_fc2")
+ is_extra_state_layer = any([l in new_name for l in extra_state_layers])
+ if use_te and is_extra_state_layer:
+ layer = new_name.split(".")[-2]
+ if layer in extra_state_layers:
+ extra_state_name = (
+ new_name[: new_name.rfind(".") + 1] + "_extra_state"
+ ) # Replace the weight name.
+ new_state_dicts[i]["model"][extra_state_name] = None
+
+ for i in range(tensor_parallel_size):
+ output_dir_tp = os.path.join(output_path, "iter_0000001", f"mp_rank_0{i}")
+ os.makedirs(output_dir_tp)
+ output_path_tp = os.path.join(output_dir_tp, "model_optim_rng.pt")
+ torch.save(new_state_dicts[i], output_path_tp)
+ with open(os.path.join(output_path, "latest_checkpointed_iteration.txt"), "w") as f:
+ f.write("1")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description="""
+Convert RADIO weights to megatron format.
+
+
+Example usage:
+python radio_converter.py --output /some/output/folder --tensor-parallel-size 4
+""",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ )
+
+ parser.add_argument(
+ "--output", type=str, required=True, help="output directory for megatron state dict file(s)"
+ )
+ parser.add_argument(
+ "--tensor-parallel-size", type=int, default=1, help="model tensor parallel size"
+ )
+ parser.add_argument("--use-te", action="store_true", help="Use Transformer Engine")
+ parser.add_argument("--version", type=str, default="radio_v2.5-h", help="Version of radio to load for conversion")
+
+ args = parser.parse_args()
+
+ convert(args.output, args.tensor_parallel_size, args.use_te, args.version)
+
+ print("done.")
diff --git a/examples/multimodal/model_converter/siglip_converter.py b/examples/multimodal/model_converter/siglip_converter.py
new file mode 100644
index 0000000000..666cda15eb
--- /dev/null
+++ b/examples/multimodal/model_converter/siglip_converter.py
@@ -0,0 +1,154 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import argparse
+import os
+from transformers import PaliGemmaForConditionalGeneration
+import torch
+
+
+def convert(output_path, tensor_parallel_size, use_te):
+ device = "cuda"
+
+ model_id = "google/paligemma-3b-pt-448"
+ model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).eval()
+
+ model = model.to(device)
+
+ print(model.config)
+ for name, tensor in model.state_dict().items():
+ if "vision_model" not in name:
+ continue
+ shape_str = "(" + ", ".join([str(x) for x in tensor.shape]) + ")"
+ print(f"{name:<75} {shape_str:>20}")
+
+ state_dict = model.state_dict()
+ new_state_dicts = [{"model": dict()} for _ in range(tensor_parallel_size)]
+
+ def add_chunck_tensor(new_tensor, new_name, chunk_dim=None):
+ if chunk_dim is None:
+ new_tensors = [new_tensor for _ in range(tensor_parallel_size)]
+ else:
+ new_tensors = torch.chunk(new_tensor, tensor_parallel_size, dim=chunk_dim)
+
+ for i in range(tensor_parallel_size):
+ # chunk() creates a view of a bigger tensor. clone() is used here to avoid excessive storage.
+ new_state_dicts[i]["model"][new_name] = new_tensors[i].clone()
+
+ # TE sets _extra_state (for FP8 purposes), so set an empty one here for compatibility.
+ extra_state_layers = ("linear_qkv", "linear_proj", "linear_fc1", "linear_fc2")
+ is_extra_state_layer = any([l in new_name for l in extra_state_layers])
+ if use_te and is_extra_state_layer:
+ layer = new_name.split(".")[-2]
+ if layer in extra_state_layers:
+ extra_state_name = (
+ new_name[: new_name.rfind(".") + 1] + "_extra_state"
+ ) # Replace the weight name.
+ new_state_dicts[i]["model"][extra_state_name] = None
+
+ for name, tensor in state_dict.items():
+ if tensor.dtype == torch.float16:
+ state_dict[name] = tensor.to(torch.float32)
+
+ add_chunck_tensor(
+ state_dict["vision_tower.vision_model.embeddings.position_embedding.weight"],
+ "position_embeddings.weight")
+ add_chunck_tensor(
+ state_dict["vision_tower.vision_model.embeddings.patch_embedding.weight"],
+ "conv1.weight")
+ add_chunck_tensor(
+ state_dict["vision_tower.vision_model.embeddings.patch_embedding.bias"],
+ "conv1.bias")
+
+ head_dim = 72
+ num_head = 16
+ for layer_idx in range(27):
+ origin_base = f"vision_tower.vision_model.encoder.layers.{layer_idx}"
+ target_base = f"decoder.layers.{layer_idx}"
+
+ for param_type in ["weight", "bias"]:
+ # QKV
+ q_proj_params = state_dict[f"{origin_base}.self_attn.q_proj.{param_type}"]
+ k_proj_params = state_dict[f"{origin_base}.self_attn.k_proj.{param_type}"]
+ v_proj_params = state_dict[f"{origin_base}.self_attn.v_proj.{param_type}"]
+ # Do some tensor manipulation because megatron expect one tensor
+ # projection for the QKV in the order
+ # [(Q1, K1, V1), (Q2, K2, V2), ...] where Qi is the query of the
+ # i-th head with dimension num_head.
+ new_tensor = torch.concatenate([
+ q_proj_params.view(num_head, head_dim, -1),
+ k_proj_params.view(num_head, head_dim, -1),
+ v_proj_params.view(num_head, head_dim, -1)], axis=1).view(
+ 3*head_dim*num_head, -1)
+ if param_type == "bias":
+ new_tensor = new_tensor[:, 0]
+ new_name = f"{target_base}.self_attention.linear_qkv.{param_type}"
+ add_chunck_tensor(new_tensor, new_name, chunk_dim=0)
+ # linear_proj
+ add_chunck_tensor(
+ state_dict[f"{origin_base}.self_attn.out_proj.{param_type}"],
+ f"{target_base}.self_attention.linear_proj.{param_type}",
+ chunk_dim=1 if param_type == "weight" else None)
+ # layer_norm
+ new_name = f"{target_base}.input_layernorm.{param_type}"
+ if use_te:
+ new_name = f"{target_base}.self_attention.linear_qkv.layer_norm_{param_type}"
+ add_chunck_tensor(
+ state_dict[f"{origin_base}.layer_norm1.{param_type}"],
+ new_name)
+ # FC 1
+ add_chunck_tensor(
+ state_dict[f"{origin_base}.mlp.fc1.{param_type}"],
+ f"{target_base}.mlp.linear_fc1.{param_type}",
+ chunk_dim=0)
+ # FC 2
+ add_chunck_tensor(
+ state_dict[f"{origin_base}.mlp.fc2.{param_type}"],
+ f"{target_base}.mlp.linear_fc2.{param_type}",
+ chunk_dim=1 if param_type=="weight" else None)
+ # layer_norm
+ new_name = f"{target_base}.pre_mlp_layernorm.{param_type}"
+ if use_te:
+ new_name = f"{target_base}.mlp.linear_fc1.layer_norm_{param_type}"
+ add_chunck_tensor(
+ state_dict[f"{origin_base}.layer_norm2.{param_type}"],
+ new_name)
+
+ add_chunck_tensor(
+ state_dict["vision_tower.vision_model.post_layernorm.weight"],
+ "ln_post.weight")
+ add_chunck_tensor(
+ state_dict["vision_tower.vision_model.post_layernorm.bias"],
+ "ln_post.bias")
+
+ for i in range(tensor_parallel_size):
+ output_dir_tp = os.path.join(output_path, "iter_0000001", f"mp_rank_0{i}")
+ os.makedirs(output_dir_tp)
+ output_path_tp = os.path.join(output_dir_tp, "model_optim_rng.pt")
+ torch.save(new_state_dicts[i], output_path_tp)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description="""
+Convert SigLIP weights to megatron format.
+
+
+Example usage:
+python siglip_converter.py --tensor-parallel-size 4 --output google_paligemma_3b_pt_44_mcore_tp_4 --use-te
+
+examples/multimodal/combine_mistral_clip.sh Mistral-7B-Instruct-v0.3-mcore-tp4 google_paligemma_3b_pt_44_mcore_tp_4 mistral_7b_instruct_v0p3_google_paligemma_3b_pt_44_mcore_tp_4
+""",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ )
+ parser.add_argument(
+ "--output", type=str, required=True, help="output directory for megatron state dict file(s)"
+ )
+ parser.add_argument(
+ "--tensor-parallel-size", type=int, default=1, help="model tensor parallel size"
+ )
+ parser.add_argument("--use-te", action="store_true", help="Use Transformer Engine")
+
+ args = parser.parse_args()
+
+ convert(args.output, args.tensor_parallel_size, args.use_te)
+
+ print("done.")
diff --git a/examples/multimodal/model_converter/vision_model_tester.py b/examples/multimodal/model_converter/vision_model_tester.py
new file mode 100644
index 0000000000..ef36dd5f9e
--- /dev/null
+++ b/examples/multimodal/model_converter/vision_model_tester.py
@@ -0,0 +1,121 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import argparse
+import os
+import sys
+
+# Add megatron and the multimodal example to the path.
+sys.path.append(
+ os.path.abspath(
+ os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir, os.path.pardir)
+ )
+)
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)))
+
+import torch
+from transformers import AutoModel
+
+from examples.multimodal.model import model_provider
+from examples.multimodal.multimodal_args import add_multimodal_extra_args
+from megatron.training import get_model
+from megatron.training.checkpointing import load_checkpoint
+from megatron.training.initialize import initialize_megatron
+
+
+def run_mcore_vision(model_path):
+ """Run mcore vision model."""
+ os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
+
+ # Megatron has some mandatory flags.
+ sys.argv = [
+ "ignore_me.py",
+ "--micro-batch-size=1",
+ "--num-layers=2",
+ "--vision-model-type=internvit",
+ "--language-model-type=mistral_7b",
+ "--tokenizer-prompt-format=mistral",
+ "--tokenizer-type=MultimodalTokenizer",
+ "--tokenizer-model=mistralai/Mistral-7B-Instruct-v0.3",
+ "--vocab-size=1024",
+ "--hidden-size=64",
+ "--num-attention-heads=8",
+ "--seq-length=1024",
+ "--decoder-seq-length=2048",
+ "--max-position-embeddings=2048",
+ "--bf16",
+ "--img-h=448",
+ "--img-w=448",
+ "--patch-dim=14",
+ "--tensor-model-parallel-size=8",
+ "--use-te",
+ f"--pretrained-checkpoint={model_path}",
+ ]
+
+ initialize_megatron(extra_args_provider=add_multimodal_extra_args)
+
+ def wrapped_model_provider(pre_process, post_process):
+ return model_provider(pre_process, post_process, parallel_output=False)
+
+ # Set up model and load checkpoint.
+ model = get_model(wrapped_model_provider, wrap_with_ddp=False)
+
+ vision_model = model[0].module.vision_model
+
+ load_checkpoint([vision_model], None, None)
+
+ vision_model.eval()
+
+ images = torch.ones((1, 3, 448, 448), dtype=torch.bfloat16, device="cuda")
+
+ output = vision_model(images)
+
+ return output
+
+
+def run_hf_vision(model_name):
+ """Run HF vision model."""
+ model = (
+ AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True)
+ .cuda()
+ .eval()
+ )
+
+ images = torch.ones((1, 3, 448, 448), dtype=torch.bfloat16, device="cuda")
+
+ outputs = model(images, return_dict=True)
+
+ return outputs
+
+
+def main(mcore_model, hf_model):
+ """Compare vision model outputs between mcore and HF given the same fixed input."""
+ mcore = run_mcore_vision(mcore_model)
+
+ if torch.distributed.get_rank() == 0:
+ hf = run_hf_vision(hf_model)
+ hf = hf["last_hidden_state"]
+
+ # Compare logits. Due to different attention implementations and other details,
+ # there will be numerical differences.
+ diff = (mcore - hf).abs()
+ mean_diff = diff.mean().item()
+ max_diff = diff.max().item()
+ print(f"mean diff {mean_diff}, max diff {max_diff}")
+ assert mean_diff < 0.1, "mean output difference is greater than expected"
+ assert max_diff < 50, "max output difference is greater than expected"
+
+ print("lgtm")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description="Check mcore vision model output vs. HF numerically.",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ )
+ parser.add_argument(
+ "--mcore-model", type=str, required=True, help="directory for mcore model weights"
+ )
+ parser.add_argument("--hf-model", type=str, required=True, help="Model name in HF")
+
+ args = parser.parse_args()
+
+ main(args.mcore_model, args.hf_model)
diff --git a/examples/multimodal/multimodal_args.py b/examples/multimodal/multimodal_args.py
new file mode 100644
index 0000000000..bef6bb9a48
--- /dev/null
+++ b/examples/multimodal/multimodal_args.py
@@ -0,0 +1,89 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from megatron.core.models.multimodal.llava_model import IMAGE_TOKEN
+
+
+def add_multimodal_extra_args(parser):
+ """Extra arguments."""
+ group = parser.add_argument_group(title='multimodal arguments')
+ group.add_argument('--dataset-config', type=str, default=None)
+ group.add_argument("--prompt-path", type=str, default=None)
+ group.add_argument('--freeze-LM', action='store_true', default=False)
+ group.add_argument('--freeze-ViT', action='store_true', default=False)
+ group.add_argument('--language-model-type', type=str, required=True)
+ group.add_argument('--language-huggingface-model-name-or-path', type=str)
+ group.add_argument('--vision-model-type', type=str, default="clip")
+ group.add_argument('--vision-huggingface-model-name-or-path', type=str)
+ group.add_argument("--disable-vision-class-token", action="store_true", default=False)
+ group.add_argument(
+ "--allow-missing-vision-projection-checkpoint", action="store_true", default=False
+ )
+ group.add_argument("--use-te", action="store_true", default=False)
+ group.add_argument(
+ "--dataloader-save", type=str, default=None, help="Energon dataloader state save path"
+ )
+ group.add_argument(
+ "--use-tiling", action="store_true", default=False, help="Use input image tiling"
+ )
+ group.add_argument("--max-num-tiles", type=int, default=1, help="Maximum number of image tiles")
+ group.add_argument(
+ "--use-thumbnail", action="store_true", default=False, help="Add image thumbnail as a tile"
+ )
+ group.add_argument(
+ "--dataloader-seq-length",
+ type=int,
+ help="Make dataloader to produce sequences of specific length.",
+ )
+ group.add_argument(
+ "--num-frames",
+ type=int,
+ default=1,
+ help="Number of frames to regularly sample from the video as input to the model.",
+ )
+ group.add_argument(
+ "--online-evaluation-config", type=str, help="Config file for online evaluation."
+ )
+ group.add_argument(
+ "--special-tokens",
+ nargs="*",
+ default=[IMAGE_TOKEN],
+ help="Special tokens used in the multimodal model",
+ )
+ group.add_argument(
+ "--tokenizer-prompt-format",
+ type=str,
+ choices=["mistral", "llama3", "llama3p1", "chatml", "nvlm-yi-34b", "qwen2p0", "qwen2p5"],
+ required=True,
+ help="Prompt format to use with the tokenizer.",
+ )
+ group.add_argument("--pixel-shuffle", action="store_true", default=False)
+ group.add_argument(
+ "--image-tag-type",
+ type=str,
+ choices=["nvlm", "internvl", ""],
+ default="", # Default: Image tag not used.
+ help="Surround image tokens with tags.",
+ )
+ group.add_argument("--use-tile-tags", action="store_true", default=False, help="Use tile tags")
+ group.add_argument(
+ "--packing-buffer-size",
+ type=int,
+ default=None, # Packing is disabled by default.
+ help="Enable sample packing by setting the buffer size to > 0",
+ )
+ group.add_argument(
+ "--packing-seq-length", type=int, default=0, help="Packing sequence length. Must be > 0 if using packing."
+ )
+ group.add_argument(
+ "--recompute-vision", action="store_true", default=False, help="Enable activation checkpointing in the vision model"
+ )
+ group.add_argument(
+ "--use-loss-scaling", action="store_true", default=False, help="Scale loss based on conversation turn length (in tokens)."
+ )
+ group.add_argument(
+ "--use-area-weighted-aspect-ratio", action="store_true", default=False,
+ help=(
+ "When --use-tiling is True, find the aspect ratio to use based on the original ",
+ "image aspect ratio and the area covered by the tiles.")
+ )
+
+ return parser
diff --git a/examples/multimodal/nvlm/README.md b/examples/multimodal/nvlm/README.md
new file mode 100644
index 0000000000..bb576bb403
--- /dev/null
+++ b/examples/multimodal/nvlm/README.md
@@ -0,0 +1,107 @@
+NVLM
+====
+
+Please refer to the [NVLM paper](https://arxiv.org/pdf/2409.11402) for details.
+
+*NOTE: VLMs in Megatron are under active development and are expected to change.*
+
+# Checkpoints
+
+NVLM 1.0 model weights are publicly available in HuggingFace and Megatron format.
+
+- NVLM-1.0-D 72B [HuggingFace version](https://huggingface.co/nvidia/NVLM-D-72B)
+- NVLM-1.0-D 72B [Megatron-Core version](https://huggingface.co/nvidia/NVLM-D-72B-mcore)
+
+# Setup
+
+## Docker image
+
+Please use `examples/multimodal/Dockerfile`.
+
+## Dataset preparation
+
+Please refer to Tables 4 and 6 in the [NVLM paper](https://arxiv.org/pdf/2409.11402) for full list of pretrain and SFT datasets.
+Please refer to https://nvidia.github.io/Megatron-Energon/data_prep.html on preparing datasets in the Megatron Energon format.
+
+## Model conversion
+
+### Vision model
+
+NVLM 1.0 models use [OpenGVLab/InternViT-6B-448px-V1-5](https://huggingface.co/OpenGVLab/InternViT-6B-448px-V1-5) from HuggingFace.
+Please download it and run the following command to convert it to Megatron format.
+```
+python examples/multimodal/model_converter/internvit_converter.py --output-dir --use-te --tensor-parallel-size 8
+```
+
+### 34B Language model
+
+NVLM 1.0 34B starts from [NousResearch/Nous-Hermes-2-Yi-34B](https://huggingface.co/NousResearch/Nous-Hermes-2-Yi-34B) from HuggingFace.
+Please download it and run the following command to convert it to Megatron format.
+```
+python tools/checkpoint/convert.py --bf16 --model-type GPT --loader llama_mistral --saver mcore --target-tensor-parallel-size 8 --checkpoint-type hf \
+ --load-dir --save-dir