diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 8160cb4ee..6b9825523 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -35,6 +35,8 @@
     title: Using 🤗 Accelerate on SageMaker
   - local: usage_guides/mps
     title: How to use Apple Silicon M1 GPUs
+  - local: usage_guides/megatron_lm
+    title: How to use Megatron-LM
   - local: usage_guides/training_zoo
     title: 🤗 Accelerate Example Zoo
   title: How-To Guides
@@ -71,4 +73,6 @@
     title: Kwargs handlers
   - local: package_reference/utilities
     title: Utility functions and classes
+  - local: package_reference/megatron_lm
+    title: Megatron-LM Utilities
   title: "Reference"
\ No newline at end of file
diff --git a/docs/source/package_reference/megatron_lm.mdx b/docs/source/package_reference/megatron_lm.mdx
new file mode 100644
index 000000000..b59768bec
--- /dev/null
+++ b/docs/source/package_reference/megatron_lm.mdx
@@ -0,0 +1,29 @@
+<!--Copyright 2021 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# Utilities for Megatron-LM
+
+[[autodoc]] utils.MegatronLMPlugin
+
+[[autodoc]] utils.MegatronLMDummyScheduler
+
+[[autodoc]] utils.MegatronLMDummyDataLoader
+
+[[autodoc]] utils.AbstractTrainStep
+
+[[autodoc]] utils.GPTTrainStep
+
+[[autodoc]] utils.BertTrainStep
+
+[[autodoc]] utils.T5TrainStep
+
+[[autodoc]] utils.avg_losses_across_data_parallel_group
diff --git a/docs/source/usage_guides/megatron_lm.mdx b/docs/source/usage_guides/megatron_lm.mdx
new file mode 100644
index 000000000..188dddb32
--- /dev/null
+++ b/docs/source/usage_guides/megatron_lm.mdx
@@ -0,0 +1,507 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+
+# Megatron-LM
+
+[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) enables training large transformer language models at scale.
+It provides efficient tensor, pipeline and sequence based model parallelism for pre-training transformer based
+Language Models such as [GPT](https://arxiv.org/abs/2005.14165) (Decoder Only), [BERT](https://arxiv.org/pdf/1810.04805.pdf) (Encoder Only) and [T5](https://arxiv.org/abs/1910.10683) (Encoder-Decoder).
+For detailed information and how things work behind the scene please refer the github [repo](https://github.com/NVIDIA/Megatron-LM).
+
+## What is integrated?
+
+Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
+of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
+
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+
+b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
+Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
+Layers are distributed uniformly across PP stages.
+
+c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
+
+d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.
+
+e. **Selective Activation Recomputation**: Reduces the memory footprint of activations significantly via smart activation checkpointing.
+It doesn't store activations occupying large memory while being fast to recompute thereby achieving great tradeoff between memory and recomputation.
+
+f. **Fused Kernels**: Fused Softmax, Mixed Precision Fused Layer Norm and  Fused gradient accumulation to weight gradient computation of linear layer.
+PyTorch JIT compiled Fused GeLU and Fused Bias+Dropout+Residual addition.
+
+g. **Support for Indexed datasets**: Efficient binary format of datasets for large scale training. Support for the `mmap`, `cached` index file and the `lazy` loader format.
+
+h. **Checkpoint reshaping and interoperability**: Utility for reshaping Megatron-LM checkpoints of variable 
+tensor and pipeline parallel sizes to the beloved 🤗 Transformers sharded checkpoints as it has great support with plethora of tools
+such as 🤗 Accelerate Big Model Inference, Megatron-DeepSpeed Inference etc. 
+Support is also available for converting 🤗 Transformers sharded checkpoints to Megatron-LM checkpoint of variable tensor and pipeline parallel sizes
+for large scale training.  
+
+
+## Pre-Requisites 
+
+You will need to install the latest pytorch, cuda, nccl, and NVIDIA [APEX](https://github.com/NVIDIA/apex#quick-start) releases and the nltk library.
+See [documentation](https://github.com/NVIDIA/Megatron-LM#setup) for more details. 
+Another way to setup the environment is to pull an NVIDIA PyTorch Container that comes with all the required installations from NGC.
+
+Below is a step-by-step method to set up the conda environment:
+
+1. Create a virtual environment
+```
+conda create --name ml
+```
+
+2. Assuming that the machine has CUDA 11.3 installed, installing the corresponding PyTorch GPU Version
+```
+conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
+```
+
+3. Install Nvidia APEX
+```
+git clone https://github.com/NVIDIA/apex
+cd apex
+pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+cd ..
+```
+
+4. Installing Megatron-LM
+
+    a. Cloning the Megatron-LM repo
+    ```
+    git clone https://github.com/NVIDIA/Megatron-LM.git
+    cd Megatron-LM
+    ```
+
+    b. Create a file `setup.py`, paste the below code and put in the root folder
+    ```python
+    """Setup for pip package."""
+
+    import os
+    import sys
+    import setuptools
+
+    if sys.version_info < (3,):
+        raise Exception("Python 2 is not supported by Megatron.")
+
+    with open("README.md", "r") as fh:
+        long_description = fh.read()
+
+    setuptools.setup(
+        name="megatron-lm",
+        version="3.0.0",
+        description="Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.",
+        long_description=long_description,
+        long_description_content_type="text/markdown",
+        # The project's main homepage.
+        url="https://github.com/NVIDIA/Megatron-LM",
+        author="NVIDIA INC",
+        maintainer="NVIDIA INC",
+        # The licence under which the project is released
+        license="See https://github.com/NVIDIA/Megatron-LM/blob/master/LICENSE",
+        classifiers=[
+            "Intended Audience :: Developers",
+            "Intended Audience :: Science/Research",
+            "Intended Audience :: Information Technology",
+            # Indicate what your project relates to
+            "Topic :: Scientific/Engineering :: Artificial Intelligence",
+            "Topic :: Software Development :: Libraries :: Python Modules",
+            # Additional Setting
+            "Environment :: Console",
+            "Natural Language :: English",
+            "Operating System :: OS Independent",
+        ],
+        python_requires=">=3.6",
+        packages=setuptools.find_packages(),
+        install_requires=["nltk", "six", "regex", "torch>=1.12.0", "pybind11"],
+        # Add in any packaged data.
+        include_package_data=True,
+        zip_safe=False,
+        # PyPI package information.
+        keywords="deep learning, Megatron, gpu, NLP, nvidia, pytorch, torch, language",
+    )
+    ```
+
+    c. installing via below command
+    ```
+    pip install "."
+    ```
+
+## Accelerate Megatron-LM Plugin
+
+Important features are directly supported via the `accelerate config` command. 
+An example of thr corresponding questions for using Megatron-LM features is shown below:
+
+```bash
+:~$ accelerate config --config_file "megatron_gpt_config.yaml"
+In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
+Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): 2
+How many different machines will you use (use more than 1 for multi-node training)? [1]: 
+Do you want to use DeepSpeed? [yes/NO]: 
+Do you want to use FullyShardedDataParallel? [yes/NO]: 
+Do you want to use Megatron-LM ? [yes/NO]: yes
+What is the Tensor Parallelism degree/size? [1]:2
+Do you want to enable Sequence Parallelism? [YES/no]: 
+What is the Pipeline Parallelism degree/size? [1]:2
+What is the number of micro-batches? [1]:2
+Do you want to enable selective activation recomputation? [YES/no]: 
+Do you want to use distributed optimizer which shards optimizer state and gradients across data pralellel ranks? [YES/no]: 
+What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: 
+How many GPU(s) should be used for distributed training? [1]:8
+Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: bf16
+```
+
+The resulting config is shown below:
+
+```
+~$ cat megatron_gpt_config.yaml 
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MEGATRON_LM
+downcast_bf16: 'no'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config:
+  megatron_lm_gradient_clipping: 1.0
+  megatron_lm_num_micro_batches: 2
+  megatron_lm_pp_degree: 2
+  megatron_lm_recompute_activations: true
+  megatron_lm_sequence_parallelism: true
+  megatron_lm_tp_degree: 2
+  megatron_lm_use_distributed_optimizer: true
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+use_cpu: false
+```
+
+We will take the example of GPT pre-training. The minimal changes required to the official `run_clm_no_trainer.py` 
+to use Megatron-LM are as follows:
+
+1. As Megatron-LM uses its own implementation of Optimizer, the corresponding scheduler compatible with it needs to be used.
+As such, support for only the Megatron-LM's scheduler is present. User will need to create `accelerate.utils.MegatronLMDummyScheduler`.
+Example is given below:
+
+```python
+from accelerate.utils import MegatronLMDummyScheduler
+
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    lr_scheduler = MegatronLMDummyScheduler(
+        optimizer=optimizer,
+        total_num_steps=args.max_train_steps,
+        warmup_num_steps=args.num_warmup_steps,
+    )
+else:
+    lr_scheduler = get_scheduler(
+        name=args.lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
+        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+    )
+```
+
+2. Getting the details of the total batch size now needs to be cognization of tensor and pipeline parallel sizes.
+Example of getting the effective total batch size is shown below:
+
+```python
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size
+else:
+    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+```
+
+3. When using Megatron-LM, the losses are already averaged across the data parallel group
+
+```python
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    losses.append(loss)
+else:
+    losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))
+
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    losses = torch.tensor(losses)
+else:
+    losses = torch.cat(losses)
+```
+
+4. For Megatron-LM, we need to save the model using `accelerator.save_state`
+
+```python
+if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+    accelerator.save_state(args.output_dir)
+else:
+    unwrapped_model = accelerator.unwrap_model(model)
+    unwrapped_model.save_pretrained(
+        args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
+    )
+```
+
+That's it! We are good to go 🚀. Please find the example script in the examples folder at the path `accelerate/examples/by_feature/megatron_lm_gpt_pretraining.py`.
+Let's run it for `gpt-large` model architecture using 4 A100-80GB GPUs.
+
+```bash
+accelerate launch --config_file megatron_gpt_config.yaml \
+examples/by_feature/megatron_lm_gpt_pretraining.py \
+--config_name "gpt2-large" \
+--tokenizer_name "gpt2-large" \
+--dataset_name wikitext \
+--dataset_config_name wikitext-2-raw-v1 \
+--block_size 1024 \
+--learning_rate 5e-5 \
+--per_device_train_batch_size 4 \
+--per_device_eval_batch_size 4 \
+--num_train_epochs 5 \
+--with_tracking \
+--report_to "wandb" \
+--output_dir "awesome_model"
+```
+
+Below are some important excerpts from the output logs:
+
+```bash
+Loading extension module fused_dense_cuda...
+>>> done with compiling and loading fused kernels. Compilation time: 3.569 seconds
+ > padded vocab (size: 50257) with 175 dummy tokens (new size: 50432)
+Building gpt model in the pre-training mode.
+The Megatron LM model weights are initialized at random in `accelerator.prepare`. Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup.
+Preparing dataloader
+Preparing dataloader
+Preparing model
+ > number of parameters on (tensor, pipeline) model parallel rank (1, 0): 210753280
+ > number of parameters on (tensor, pipeline) model parallel rank (1, 1): 209445120
+ > number of parameters on (tensor, pipeline) model parallel rank (0, 0): 210753280
+ > number of parameters on (tensor, pipeline) model parallel rank (0, 1): 209445120
+Preparing optimizer
+Preparing scheduler
+> learning rate decay style: linear
+10/10/2022 22:57:22 - INFO - __main__ - ***** Running training *****
+10/10/2022 22:57:22 - INFO - __main__ -   Num examples = 2318
+10/10/2022 22:57:22 - INFO - __main__ -   Num Epochs = 5
+10/10/2022 22:57:22 - INFO - __main__ -   Instantaneous batch size per device = 24
+10/10/2022 22:57:22 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 48
+10/10/2022 22:57:22 - INFO - __main__ -   Gradient Accumulation steps = 1
+10/10/2022 22:57:22 - INFO - __main__ -   Total optimization steps = 245
+ 20%|████████████▍                                                 | 49/245 [01:04<04:09,  1.27s/it]
+ 10/10/2022 22:58:29 - INFO - __main__ - epoch 0: perplexity: 1222.1594275215962 eval_loss: 7.10837459564209
+ 40%|████████████████████████▊                                     | 98/245 [02:10<03:07,  1.28s/it]
+ 10/10/2022 22:59:35 - INFO - __main__ - epoch 1: perplexity: 894.5236583794557 eval_loss: 6.796291351318359
+ 60%|████████████████████████████████████▌                        | 147/245 [03:16<02:05,  1.28s/it]
+ 10/10/2022 23:00:40 - INFO - __main__ - epoch 2: perplexity: 702.8458788508042 eval_loss: 6.555137634277344
+ 80%|████████████████████████████████████████████████▊            | 196/245 [04:22<01:02,  1.28s/it]
+ 10/10/2022 23:01:46 - INFO - __main__ - epoch 3: perplexity: 600.3220028695281 eval_loss: 6.39746618270874
+100%|█████████████████████████████████████████████████████████████| 245/245 [05:27<00:00,  1.28s/it]
+```
+
+There are a large number of other options/features that one can set using `accelerate.utils.MegatronLMPlugin`.
+
+## Advanced features to leverage writing custom train step and Megatron-LM Indexed Datasets
+
+For leveraging more features, please go through below details.
+
+1. Below is an example of changes required to customize the Train Step while using Megatron-LM. 
+You will implement the `accelerate.utils.AbstractTrainStep` or inherit from their corresponding children 
+`accelerate.utils.GPTTrainStep`, `accelerate.utils.BertTrainStep` or `accelerate.utils.T5TrainStep`.
+
+```python
+from accelerate.utils import MegatronLMDummyScheduler, GPTTrainStep, avg_losses_across_data_parallel_group
+
+# Custom loss function for the Megatron model
+class GPTTrainStepWithCustomLoss(GPTTrainStep):
+    def __init__(self, megatron_args, **kwargs):
+        super().__init__(megatron_args)
+        self.kwargs = kwargs
+
+    def get_loss_func(self):
+        def loss_func(inputs, loss_mask, output_tensor):
+            batch_size, seq_length = output_tensor.shape
+            losses = output_tensor.float()
+            loss_mask = loss_mask.view(-1).float()
+            loss = losses.view(-1) * loss_mask
+
+            # Resize and average loss per sample
+            loss_per_sample = loss.view(batch_size, seq_length).sum(axis=1)
+            loss_mask_per_sample = loss_mask.view(batch_size, seq_length).sum(axis=1)
+            loss_per_sample = loss_per_sample / loss_mask_per_sample
+
+            # Calculate and scale weighting
+            weights = torch.stack([(inputs == kt).float() for kt in self.kwargs["keytoken_ids"]]).sum(axis=[0, 2])
+            weights = 1.0 + self.kwargs["alpha"] * weights
+            # Calculate weighted average
+            weighted_loss = (loss_per_sample * weights).mean()
+
+            # Reduce loss across data parallel groups
+            averaged_loss = avg_losses_across_data_parallel_group([weighted_loss])
+
+            return weighted_loss, {"lm loss": averaged_loss[0]}
+
+        return loss_func
+
+    def get_forward_step_func(self):
+        def forward_step(data_iterator, model):
+            """Forward step."""
+            # Get the batch.
+            tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator)
+            output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
+
+            return output_tensor, partial(self.loss_func, tokens, loss_mask)
+
+        return forward_step
+
+
+def main():
+    # Custom loss function for the Megatron model
+    keytoken_ids = []
+    keywords = ["plt", "pd", "sk", "fit", "predict", " plt", " pd", " sk", " fit", " predict"]
+    for keyword in keywords:
+        ids = tokenizer([keyword]).input_ids[0]
+        if len(ids) == 1:
+            keytoken_ids.append(ids[0])
+    accelerator.print(f"Keytoken ids: {keytoken_ids}")
+    accelerator.state.megatron_lm_plugin.custom_train_step_class = GPTTrainStepWithCustomLoss
+    accelerator.state.megatron_lm_plugin.custom_train_step_kwargs = {
+        "keytoken_ids": keytoken_ids,
+        "alpha": 0.25,
+    }
+```
+
+2. For using the Megatron-LM datasets, a few more changes are required. Dataloaders for these datasets
+are available only on rank 0 of each tensor parallel group. As such, there are rank where dataloader won't be
+avaiable and this requires tweaks to the training loop. Being able to do all this shows how
+felixble and extensible 🤗 Accelerate is. The changes required are as follows.
+
+a. For Megatron-LM indexed datasets, we need to use `MegatronLMDummyDataLoader` 
+and pass the required dataset args to it such as `data_path`, `seq_length` etc. 
+See [here](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/arguments.py#L804) for the list of available args. 
+    
+```python
+from accelerate.utils import MegatronLMDummyDataLoader
+
+megatron_dataloader_config = {
+    "data_path": args.data_path,
+    "splits_string": args.splits_string,
+    "seq_length": args.block_size,
+    "micro_batch_size": args.per_device_train_batch_size,
+}
+megatron_dataloader = MegatronLMDummyDataLoader(**megatron_dataloader_config)
+accelerator.state.megatron_lm_plugin.megatron_dataset_flag = True
+```
+
+b. `megatron_dataloader` is repeated 3 times to get training, validation and test dataloaders
+as per the `args.splits_string` proportions
+    
+```python
+model, optimizer, lr_scheduler, train_dataloader, eval_dataloader, _ = accelerator.prepare(
+    model, optimizer, lr_scheduler, megatron_dataloader, megatron_dataloader, megatron_dataloader
+)
+```
+
+c. Changes to training and evaluation loops as dataloader is only available on tensor parallel ranks 0
+So, we need to iterate only if the dataloader isn't `None` else provide empty dict
+As such, we loop using `while` loop and break when `completed_steps` is equal to `args.max_train_steps`
+This is similar to the Megatron-LM setup wherein user has to provide `max_train_steps` when using Megaton-LM indexed datasets.
+This displays how flexible and extensible 🤗 Accelerate is.
+
+```python
+while completed_steps < args.max_train_steps:
+    model.train()
+    batch = next(train_dataloader) if train_dataloader is not None else {}
+    outputs = model(**batch)
+    loss = outputs.loss
+    ...
+
+    if completed_steps % eval_interval == 0:
+        eval_completed_steps = 0
+        losses = []
+        while eval_completed_steps < eval_iters:
+            model.eval()
+            with torch.no_grad():
+                batch = next(eval_dataloader) if eval_dataloader is not None else {}
+                outputs = model(**batch)
+```
+
+    
+## Utility for Checkpoint reshaping and interoperability
+
+1. The scripts for these are present in 🤗 Transformers library under respective models. 
+Currently, it is available for GPT model [checkpoint_reshaping_and_interoperability.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/megatron_gpt2/checkpoint_reshaping_and_interoperability.py)
+
+2. Below is an example of conversion of checkpoint from Megatron-LM to universal 🤗 Transformers sharded checkpoint.
+```bash
+python checkpoint_reshaping_and_interoperability.py \
+--convert_checkpoint_from_megatron_to_transformers \
+--load_path "gpt/iter_0005000" \
+--save_path "gpt/trfs_checkpoint" \
+--max_shard_size "200MB" \
+--tokenizer_name "gpt2" \
+--print-checkpoint-structure
+```
+
+3. Conversion of checkpoint from transformers to megatron with `tp_size=2`, `pp_size=2` and `dp_size=2`.
+```bash
+python checkpoint_utils/megatgron_gpt2/checkpoint_reshaping_and_interoperability.py \
+--load_path "gpt/trfs_checkpoint" \
+--save_path "gpt/megatron_lm_checkpoint" \
+--target_tensor_model_parallel_size 2 \
+--target_pipeline_model_parallel_size 2 \
+--target_data_parallel_size 2 \
+--target_params_dtype "bf16" \
+--make_vocab_size_divisible_by 128 \
+--use_distributed_optimizer \
+--print-checkpoint-structure
+```
+
+## Caveats
+
+1. Supports Transformers GPT2, Megatron-BERT and T5 models.
+This covers Decoder only, Encode only and Encoder-Decoder model classes.
+
+2. Only loss is returned from model forward pass as 
+there is quite complex interplay of pipeline, tensor and data parallelsim behind the scenes.
+The `model(**batch_data)` call return loss(es) averaged across the data parallel ranks.
+This is fine for most cases wherein pre-training jobs are run using Megatron-LM features and
+you can easily compute the `perplexity` using the loss. 
+
+3. The main process is the last rank as the losses are available in the last stage of pipeline.
+`accelerator.is_main_process` and `accelerator.is_local_main_process` return `True` for last rank when using 
+Megatron-LM integration.
+
+4. In `accelerator.prepare` call, a Megatron-LM model corresponding to a given Transformers model is created
+with random weights. Please use `accelerator.load_state` to load the Megatron-LM checkpoint with matching TP, PP and DP partitions.
+
+5. Currently, checkpoint reshaping and interoperability support is only available for GPT. 
+Soon it will be extended to BERT and T5.
+
+6. `gradient_accumulation_steps` needs to be 1. When using Megatron-LM, micro batches in pipeline parallelism 
+setting is synonymous with gradient accumulation. 
+
+7. When using Megatron-LM, use `accelerator.save_state` and `accelerator.load_state` for saving and loading checkpoints.
+
+8. Below are the mapping from Megatron-LM model architectures to the the equivalent 🤗 transformers model architectures.
+Only these 🤗 transformers model architectures are supported.
+
+a. Megatron-LM [BertModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/bert_model.py) : 
+🤗 transformers models with `megatron-bert` in config's model type, e.g., 
+[MegatronBERT](https://huggingface.co/docs/transformers/model_doc/megatron-bert)
+    
+b. Megatron-LM [GPTModel](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py) : 
+🤗 transformers models with `gpt2` in config's model type, e.g., 
+[OpenAI GPT2](https://huggingface.co/docs/transformers/model_doc/gpt2)
+   
+c. Megatron-LM [T5Model](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/t5_model.py) : 
+🤗 transformers models with `t5` in  config's model type, e.g., 
+[T5](https://huggingface.co/docs/transformers/model_doc/t5) and 
+[MT5](https://huggingface.co/docs/transformers/model_doc/mt5)
\ No newline at end of file
diff --git a/examples/by_feature/megatron_lm_gpt_pretraining.py b/examples/by_feature/megatron_lm_gpt_pretraining.py
new file mode 100644
index 000000000..68a425e4d
--- /dev/null
+++ b/examples/by_feature/megatron_lm_gpt_pretraining.py
@@ -0,0 +1,699 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...)
+on a text file or a dataset without using HuggingFace Trainer.
+
+Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
+https://huggingface.co/models?filter=text-generation
+"""
+# You can also adapt this script on your own causal language modeling task. Pointers for this are left as comments.
+
+import argparse
+import json
+import logging
+import math
+import os
+import random
+from itertools import chain
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+import datasets
+import transformers
+from accelerate import Accelerator, DistributedType
+from accelerate.logging import get_logger
+from accelerate.utils import MegatronLMDummyScheduler, set_seed
+from datasets import load_dataset
+from huggingface_hub import Repository
+from tqdm.auto import tqdm
+from transformers import (
+    CONFIG_MAPPING,
+    MODEL_MAPPING,
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    SchedulerType,
+    default_data_collator,
+    get_scheduler,
+)
+from transformers.utils import check_min_version, get_full_repo_name, send_example_telemetry
+from transformers.utils.versions import require_version
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+check_min_version("4.23.0.dev0")
+
+logger = get_logger(__name__)
+
+require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/language-modeling/requirements.txt")
+
+MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Finetune a transformers model on a causal language modeling task")
+    parser.add_argument(
+        "--dataset_name",
+        type=str,
+        default=None,
+        help="The name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--dataset_config_name",
+        type=str,
+        default=None,
+        help="The configuration name of the dataset to use (via the datasets library).",
+    )
+    parser.add_argument(
+        "--train_file", type=str, default=None, help="A csv or a json file containing the training data."
+    )
+    parser.add_argument(
+        "--validation_file", type=str, default=None, help="A csv or a json file containing the validation data."
+    )
+    parser.add_argument(
+        "--validation_split_percentage",
+        default=5,
+        help="The percentage of the train set used as validation set in case there's no validation split",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--config_name",
+        type=str,
+        default=None,
+        help="Pretrained config name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help="Pretrained tokenizer name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--use_slow_tokenizer",
+        action="store_true",
+        help="If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library).",
+    )
+    parser.add_argument(
+        "--per_device_train_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--per_device_eval_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the evaluation dataloader.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-5,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
+    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--lr_scheduler_type",
+        type=SchedulerType,
+        default="linear",
+        help="The scheduler type to use.",
+        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
+    )
+    parser.add_argument(
+        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default=None,
+        help="Model type to use if training from scratch.",
+        choices=MODEL_TYPES,
+    )
+    parser.add_argument(
+        "--block_size",
+        type=int,
+        default=None,
+        help=(
+            "Optional input sequence length after tokenization. The training dataset will be truncated in block of"
+            " this size for training. Default to the model max input length for single sentence inputs (take into"
+            " account special tokens)."
+        ),
+    )
+    parser.add_argument(
+        "--preprocessing_num_workers",
+        type=int,
+        default=None,
+        help="The number of processes to use for the preprocessing.",
+    )
+    parser.add_argument(
+        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
+    )
+    parser.add_argument(
+        "--no_keep_linebreaks", action="store_true", help="Do not keep line breaks when using TXT files."
+    )
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument(
+        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
+    )
+    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=str,
+        default=None,
+        help="Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--with_tracking",
+        action="store_true",
+        help="Whether to enable experiment trackers for logging.",
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="all",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`,'
+            ' `"wandb"` and `"comet_ml"`. Use `"all"` (default) to report to all integrations.'
+            "Only applicable when `--with_tracking` is passed."
+        ),
+    )
+    args = parser.parse_args()
+
+    # Sanity checks
+    if args.dataset_name is None and args.train_file is None and args.validation_file is None:
+        raise ValueError("Need either a dataset name or a training/validation file.")
+    else:
+        if args.train_file is not None:
+            extension = args.train_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, json or txt file."
+        if args.validation_file is not None:
+            extension = args.validation_file.split(".")[-1]
+            assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, json or txt file."
+
+    if args.push_to_hub:
+        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."
+
+    return args
+
+
+def main():
+    args = parse_args()
+
+    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
+    # information sent is the one passed as arguments along with your Python/PyTorch versions.
+    send_example_telemetry("run_clm_no_trainer", args)
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    # If we're using tracking, we also need to initialize it here and it will by default pick up all supported trackers
+    # in the environment
+    accelerator_log_kwargs = {}
+
+    if args.with_tracking:
+        accelerator_log_kwargs["log_with"] = args.report_to
+        accelerator_log_kwargs["logging_dir"] = args.output_dir
+
+    accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps, **accelerator_log_kwargs)
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_info()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+
+    # If passed along, set the training seed now.
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.push_to_hub:
+            if args.hub_model_id is None:
+                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
+            else:
+                repo_name = args.hub_model_id
+            repo = Repository(args.output_dir, clone_from=repo_name)
+
+            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
+                if "step_*" not in gitignore:
+                    gitignore.write("step_*\n")
+                if "epoch_*" not in gitignore:
+                    gitignore.write("epoch_*\n")
+        elif args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+    accelerator.wait_for_everyone()
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    if args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[:{args.validation_split_percentage}%]",
+            )
+            raw_datasets["train"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[{args.validation_split_percentage}%:]",
+            )
+    else:
+        data_files = {}
+        dataset_args = {}
+        if args.train_file is not None:
+            data_files["train"] = args.train_file
+        if args.validation_file is not None:
+            data_files["validation"] = args.validation_file
+        extension = args.train_file.split(".")[-1]
+        if extension == "txt":
+            extension = "text"
+            dataset_args["keep_linebreaks"] = not args.no_keep_linebreaks
+        raw_datasets = load_dataset(extension, data_files=data_files, **dataset_args)
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[:{args.validation_split_percentage}%]",
+                **dataset_args,
+            )
+            raw_datasets["train"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[{args.validation_split_percentage}%:]",
+                **dataset_args,
+            )
+
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+    #
+    # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    if args.config_name:
+        config = AutoConfig.from_pretrained(args.config_name)
+    elif args.model_name_or_path:
+        config = AutoConfig.from_pretrained(args.model_name_or_path)
+    else:
+        config = CONFIG_MAPPING[args.model_type]()
+        logger.warning("You are instantiating a new config instance from scratch.")
+
+    if args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=not args.use_slow_tokenizer)
+    elif args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=not args.use_slow_tokenizer)
+    else:
+        raise ValueError(
+            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
+            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
+        )
+
+    if args.model_name_or_path:
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_name_or_path,
+            from_tf=bool(".ckpt" in args.model_name_or_path),
+            config=config,
+        )
+    else:
+        logger.info("Training new model from scratch")
+        model = AutoModelForCausalLM.from_config(config)
+
+    model.resize_token_embeddings(len(tokenizer))
+
+    # Preprocessing the datasets.
+    # First we tokenize all the texts.
+    column_names = raw_datasets["train"].column_names
+    text_column_name = "text" if "text" in column_names else column_names[0]
+
+    def tokenize_function(examples):
+        return tokenizer(examples[text_column_name])
+
+    with accelerator.main_process_first():
+        tokenized_datasets = raw_datasets.map(
+            tokenize_function,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            remove_columns=column_names,
+            load_from_cache_file=not args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+
+    if args.block_size is None:
+        block_size = tokenizer.model_max_length
+        if block_size > 1024:
+            logger.warning(
+                f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
+                "Picking 1024 instead. You can change that default value by passing --block_size xxx."
+            )
+        block_size = 1024
+    else:
+        if args.block_size > tokenizer.model_max_length:
+            logger.warning(
+                f"The block_size passed ({args.block_size}) is larger than the maximum length for the model"
+                f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
+            )
+        block_size = min(args.block_size, tokenizer.model_max_length)
+
+    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
+    def group_texts(examples):
+        # Concatenate all texts.
+        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        if total_length >= block_size:
+            total_length = (total_length // block_size) * block_size
+        # Split by chunks of max_len.
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
+    # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
+    # to preprocess.
+    #
+    # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
+    # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map
+
+    with accelerator.main_process_first():
+        lm_datasets = tokenized_datasets.map(
+            group_texts,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            load_from_cache_file=not args.overwrite_cache,
+            desc=f"Grouping texts in chunks of {block_size}",
+        )
+
+    train_dataset = lm_datasets["train"]
+    eval_dataset = lm_datasets["validation"]
+
+    # Log a few random samples from the training set:
+    for index in random.sample(range(len(train_dataset)), 3):
+        logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")
+
+    # DataLoaders creation:
+    train_dataloader = DataLoader(
+        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=args.per_device_train_batch_size
+    )
+    eval_dataloader = DataLoader(
+        eval_dataset, collate_fn=default_data_collator, batch_size=args.per_device_eval_batch_size
+    )
+
+    # Optimizer
+    # Split weights in two groups, one with weight decay and the other not.
+    no_decay = ["bias", "layer_norm.weight"]
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+            "weight_decay": args.weight_decay,
+        },
+        {
+            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+            "weight_decay": 0.0,
+        },
+    ]
+    optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
+
+    # Scheduler and math around the number of training steps.
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
+
+    # New Code
+    # For Megatron-LM, we need to use `MegatronLMDummyScheduler` instead of regular schedulers
+    if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+        lr_scheduler = MegatronLMDummyScheduler(
+            optimizer=optimizer,
+            total_num_steps=args.max_train_steps,
+            warmup_num_steps=args.num_warmup_steps,
+        )
+    else:
+        lr_scheduler = get_scheduler(
+            name=args.lr_scheduler_type,
+            optimizer=optimizer,
+            num_warmup_steps=args.num_warmup_steps * args.gradient_accumulation_steps,
+            num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+        )
+
+    # Prepare everything with our `accelerator`.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # On TPU, the tie weights in our model have been disconnected, so we need to restore the ties.
+    if accelerator.distributed_type == DistributedType.TPU:
+        model.tie_weights()
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if overrode_max_train_steps:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # Figure out how many steps we should save the Accelerator states
+    checkpointing_steps = args.checkpointing_steps
+    if checkpointing_steps is not None and checkpointing_steps.isdigit():
+        checkpointing_steps = int(checkpointing_steps)
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if args.with_tracking:
+        experiment_config = vars(args)
+        # TensorBoard cannot log Enums, need the raw value
+        experiment_config["lr_scheduler_type"] = experiment_config["lr_scheduler_type"].value
+        accelerator.init_trackers("clm_no_trainer", experiment_config)
+
+    # Train!
+    # New Code
+    # For Megatron-LM, we need to get `global_batch_size` from megatron_lm_plugin
+    # as it handles the specifics related to data parallelism, tensor model parallelism and pipeline parallelism
+    if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+        total_batch_size = accelerator.state.megatron_lm_plugin.global_batch_size
+    else:
+        total_batch_size = (
+            args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+        )
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+    completed_steps = 0
+    starting_epoch = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
+            accelerator.print(f"Resumed from checkpoint: {args.resume_from_checkpoint}")
+            accelerator.load_state(args.resume_from_checkpoint)
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
+            dirs.sort(key=os.path.getctime)
+            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
+        # Extract `epoch_{i}` or `step_{i}`
+        training_difference = os.path.splitext(path)[0]
+
+        if "epoch" in training_difference:
+            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
+            resume_step = None
+        else:
+            # need to multiply `gradient_accumulation_steps` to reflect real steps
+            resume_step = int(training_difference.replace("step_", "")) * args.gradient_accumulation_steps
+            starting_epoch = resume_step // len(train_dataloader)
+            resume_step -= starting_epoch * len(train_dataloader)
+
+    # update the progress_bar if load from checkpoint
+    progress_bar.update(starting_epoch * num_update_steps_per_epoch)
+    completed_steps = starting_epoch * num_update_steps_per_epoch
+
+    for epoch in range(starting_epoch, args.num_train_epochs):
+        model.train()
+        if args.with_tracking:
+            total_loss = 0
+        for step, batch in enumerate(train_dataloader):
+            # We need to skip steps until we reach the resumed step
+            if args.resume_from_checkpoint and epoch == starting_epoch:
+                if resume_step is not None and step < resume_step:
+                    if step % args.gradient_accumulation_steps == 0:
+                        progress_bar.update(1)
+                        completed_steps += 1
+                    continue
+
+            with accelerator.accumulate(model):
+                outputs = model(**batch)
+                loss = outputs.loss
+                # We keep track of the loss at each epoch
+                if args.with_tracking:
+                    total_loss += loss.detach().float()
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                completed_steps += 1
+
+            if isinstance(checkpointing_steps, int):
+                if completed_steps % checkpointing_steps == 0:
+                    output_dir = f"step_{completed_steps }"
+                    if args.output_dir is not None:
+                        output_dir = os.path.join(args.output_dir, output_dir)
+                    accelerator.save_state(output_dir)
+            if completed_steps >= args.max_train_steps:
+                break
+
+        model.eval()
+        losses = []
+        for step, batch in enumerate(eval_dataloader):
+            with torch.no_grad():
+                outputs = model(**batch)
+
+            loss = outputs.loss
+            # New Code
+            # For Megatron-LM, the losses are already averaged across the data parallel group
+            if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+                losses.append(loss)
+            else:
+                losses.append(accelerator.gather_for_metrics(loss.repeat(args.per_device_eval_batch_size)))
+        try:
+            if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+                losses = torch.tensor(losses)
+            else:
+                losses = torch.cat(losses)
+            eval_loss = torch.mean(losses)
+            perplexity = math.exp(eval_loss)
+        except OverflowError:
+            perplexity = float("inf")
+
+        logger.info(f"epoch {epoch}: perplexity: {perplexity} eval_loss: {eval_loss}")
+
+        if args.with_tracking:
+            accelerator.log(
+                {
+                    "perplexity": perplexity,
+                    "eval_loss": eval_loss,
+                    "train_loss": total_loss.item() / len(train_dataloader),
+                    "epoch": epoch,
+                    "step": completed_steps,
+                },
+                step=completed_steps,
+            )
+
+        if args.push_to_hub and epoch < args.num_train_epochs - 1:
+            accelerator.wait_for_everyone()
+            unwrapped_model = accelerator.unwrap_model(model)
+            unwrapped_model.save_pretrained(
+                args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
+            )
+            if accelerator.is_main_process:
+                tokenizer.save_pretrained(args.output_dir)
+                repo.push_to_hub(
+                    commit_message=f"Training in progress epoch {epoch}", blocking=False, auto_lfs_prune=True
+                )
+
+        if args.checkpointing_steps == "epoch":
+            output_dir = f"epoch_{epoch}"
+            if args.output_dir is not None:
+                output_dir = os.path.join(args.output_dir, output_dir)
+            accelerator.save_state(output_dir)
+
+    if args.with_tracking:
+        accelerator.end_training()
+
+    if args.output_dir is not None:
+        accelerator.wait_for_everyone()
+        # New Code
+        # For Megatron-LM, we need to save the model using `accelerator.save_state`
+        if accelerator.distributed_type == DistributedType.MEGATRON_LM:
+            accelerator.save_state(args.output_dir)
+        else:
+            unwrapped_model = accelerator.unwrap_model(model)
+            unwrapped_model.save_pretrained(
+                args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
+            )
+        if accelerator.is_main_process:
+            tokenizer.save_pretrained(args.output_dir)
+            if args.push_to_hub:
+                repo.push_to_hub(commit_message="End of training", auto_lfs_prune=True)
+
+        with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
+            json.dump({"perplexity": perplexity}, f)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2b6558944..e40797b39 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -41,6 +41,7 @@
     InitProcessGroupKwargs,
     KwargsHandler,
     LoggerType,
+    MegatronLMPlugin,
     PrecisionType,
     RNGType,
     compare_versions,
@@ -50,6 +51,7 @@
     get_pretty_name,
     is_bf16_available,
     is_deepspeed_available,
+    is_megatron_lm_available,
     is_torch_version,
     is_tpu_available,
     pad_across_processes,
@@ -71,6 +73,21 @@
         DummyScheduler,
     )
 
+if is_megatron_lm_available():
+    from .utils import (
+        MegatronEngine,
+        MegatronLMDummyDataLoader,
+        MegatronLMDummyScheduler,
+        MegatronLMOptimizerWrapper,
+        MegatronLMSchedulerWrapper,
+        megatron_lm_initialize,
+        megatron_lm_prepare_data_loader,
+        megatron_lm_prepare_model,
+        megatron_lm_prepare_optimizer,
+        megatron_lm_prepare_scheduler,
+    )
+
+
 if is_tpu_available(check_device=False):
     import torch_xla.distributed.xla_multiprocessing as xmp
 
@@ -107,6 +124,9 @@ class Accelerator:
         fsdp_plugin (`FullyShardedDataParallelPlugin`, *optional*):
             Tweak your FSDP related args using this argument. This argument is optional and can be configured directly
             using *accelerate config*
+        megatron_lm_plugin (`MegatronLMPlugin`, *optional*):
+            Tweak your MegatronLM related args using this argument. This argument is optional and can be configured
+            directly using *accelerate config*
         rng_types (list of `str` or [`~utils.RNGType`]):
             The list of random number generators to synchronize at the beginning of each iteration in your prepared
             dataloaders. Should be one or several of:
@@ -166,6 +186,7 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         fsdp_plugin: FullyShardedDataParallelPlugin = None,
+        megatron_lm_plugin: MegatronLMPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
         log_with: Optional[List[Union[str, LoggerType, GeneralTracker]]] = None,
         logging_dir: Optional[Union[str, os.PathLike]] = None,
@@ -218,6 +239,17 @@ def __init__(
                 raise TypeError("`fsdp_plugin` must be a FullyShardedDataParallelPlugin object.")
             os.environ["USE_FSDP"] = "true"  # use FSDP if plugin is provided
 
+        if megatron_lm_plugin is None:  # init from env variables
+            megatron_lm_plugin = MegatronLMPlugin() if os.environ.get("USE_MEGATRON_LM", "false") == "true" else None
+        else:
+            if not isinstance(megatron_lm_plugin, MegatronLMPlugin):
+                raise TypeError("`megatron_lm_plugin` must be a MegatronLMPlugin object.")
+            os.environ["USE_MEGATRON_LM"] = "true"  # use MegatronLM if plugin is provided
+
+        if megatron_lm_plugin:
+            if not is_megatron_lm_available():
+                raise ImportError("Megatron is not installed. please build it from source.")
+
         # Kwargs handlers
         self.ddp_handler = None
         self.scaler_handler = None
@@ -247,6 +279,7 @@ def __init__(
             cpu=cpu,
             deepspeed_plugin=deepspeed_plugin,
             fsdp_plugin=fsdp_plugin,
+            megatron_lm_plugin=megatron_lm_plugin,
             _from_accelerator=True,
             **kwargs,
         )
@@ -278,7 +311,7 @@ def __init__(
         self.scaler = None
         self.native_amp = False
         err = "{mode} mixed precision requires {requirement}"
-        if self.state.mixed_precision == "fp16":
+        if self.state.mixed_precision == "fp16" and self.distributed_type != DistributedType.MEGATRON_LM:
             self.native_amp = True
             if not torch.cuda.is_available() and not parse_flag_from_env("USE_MPS_DEVICE"):
                 raise ValueError(err.format(mode="fp16", requirement="a GPU"))
@@ -289,7 +322,11 @@ def __init__(
                 self.scaler = ShardedGradScaler(**kwargs)
             else:
                 self.scaler = torch.cuda.amp.GradScaler(**kwargs)
-        elif self.state.mixed_precision == "bf16" and self.distributed_type != DistributedType.FSDP:
+        elif (
+            self.state.mixed_precision == "bf16"
+            and self.distributed_type != DistributedType.FSDP
+            and self.distributed_type != DistributedType.MEGATRON_LM
+        ):
             self.native_amp = is_bf16_available(True)
             if mixed_precision == "bf16" and not self.native_amp and not is_tpu_available():
                 raise ValueError(err.format(mode="bf16", requirement="PyTorch >= 1.10 and a supported device."))
@@ -344,17 +381,27 @@ def device(self):
     @property
     def is_main_process(self):
         """True for one process only."""
-        return self.process_index == 0
+        return (
+            self.process_index == 0 if self.distributed_type != DistributedType.MEGATRON_LM else self.is_last_process
+        )
 
     @property
     def is_local_main_process(self):
         """True for one process per server."""
-        return self.local_process_index == 0
+        return (
+            self.local_process_index == 0
+            if self.distributed_type != DistributedType.MEGATRON_LM
+            else self.is_last_process
+        )
 
     @property
     def use_fp16(self):
         return self.mixed_precision != "no"
 
+    @property
+    def is_last_process(self):
+        return self.process_index == self.num_processes - 1
+
     @property
     def mixed_precision(self):
         if self.distributed_type == DistributedType.DEEPSPEED:
@@ -393,6 +440,18 @@ def wrapper(self, *args, **kwargs):
 
         return wrapper
 
+    def on_last_process(func):
+        """
+        A decorator that will run the decorated function on the last process only.
+        """
+
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if self.is_last_process or not self.use_distributed:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
     def on_process(process_idx):
         """
         A decorator that will run the decorated function on a given process index only.
@@ -628,8 +687,8 @@ def prepare(self, *args, device_placement=None):
         """
         if device_placement is None:
             device_placement = [None for _ in args]
-        elif self.distributed_type == DistributedType.DEEPSPEED:
-            raise ValueError("You can't customize device placements with DeepSpeed.")
+        elif self.distributed_type in (DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM):
+            raise ValueError("You can't customize device placements with DeepSpeed or Megatron-LM.")
         elif len(device_placement) != len(args):
             raise ValueError(
                 f"`device_placement` should be a list with {len(args)} elements (the number of objects passed)."
@@ -677,6 +736,8 @@ def prepare(self, *args, device_placement=None):
 
         if self.distributed_type == DistributedType.DEEPSPEED:
             result = self._prepare_deepspeed(*args)
+        elif self.distributed_type == DistributedType.MEGATRON_LM:
+            result = self._prepare_megatron_lm(*args)
         else:
             result = tuple(
                 self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
@@ -923,6 +984,107 @@ def _prepare_deepspeed(self, *args):
                 )
         return tuple(result)
 
+    def _prepare_megatron_lm(self, *args):
+        megatron_lm_plugin = self.state.megatron_lm_plugin
+        if not megatron_lm_plugin.megatron_dataset_flag:
+            batch_sizes = [obj.batch_size for obj in args if hasattr(obj, "batch_size")]
+            if len(batch_sizes) == 0:
+                raise ValueError(
+                    "You must specify a training or evaluation dataloader in `accelerate.prepare()` when using Megatron-LM."
+                )
+
+            micro_batch_size = min(batch_sizes) if megatron_lm_plugin.is_train_batch_min else max(batch_sizes)
+            if len(batch_sizes) > 1:
+                logger.info(
+                    "Since you passed both train and evaluation dataloader, `is_train_batch_min` (here "
+                    f"{megatron_lm_plugin.is_train_batch_min} will decide the `train_batch_size` ({micro_batch_size})."
+                )
+        else:
+            for obj in args:
+                if isinstance(obj, MegatronLMDummyDataLoader):
+                    micro_batch_size = obj.dataset_args["micro_batch_size"]
+                    break
+
+        dp_degree = self.num_processes // (megatron_lm_plugin.tp_degree * megatron_lm_plugin.pp_degree)
+        megatron_lm_plugin.set_training_args(micro_batch_size, dp_degree)
+
+        model = None
+        optimizer = None
+        scheduler = None
+        is_dummy_scheduler = False
+        batch_data = None
+        for obj in args:
+            if isinstance(obj, torch.utils.data.DataLoader) and batch_data is None:
+                batch_data = next(iter(obj))
+            if isinstance(obj, torch.nn.Module):
+                model = obj
+            elif isinstance(obj, (torch.optim.Optimizer)):
+                optimizer = obj
+            elif isinstance(obj, (torch.optim.lr_scheduler._LRScheduler, MegatronLMDummyScheduler)):
+                scheduler = obj
+
+        if model is not None:
+            megatron_lm_plugin.set_network_size_args(model, batch_data)
+        if optimizer is not None:
+            megatron_lm_plugin.set_optimizer_type(optimizer)
+        if scheduler is not None:
+            is_dummy_scheduler = isinstance(scheduler, MegatronLMDummyScheduler)
+            if not is_dummy_scheduler:
+                raise ValueError(
+                    "You can't use a custom scheduler with Megatron-LM. Please use the `accelerate.utils.MegatronLMDummyScheduler` instead."
+                )
+            megatron_lm_plugin.set_scheduler_args(scheduler)
+
+        # initialize megatron-lm
+        megatron_lm_initialize(self, args_defaults=megatron_lm_plugin.megatron_lm_default_args)
+        counter = 0
+        result = []
+        for obj in args:
+            if isinstance(obj, torch.utils.data.DataLoader):
+                result.append(megatron_lm_prepare_data_loader(self, obj))
+                counter += 1
+            elif isinstance(obj, MegatronLMDummyDataLoader):
+                if counter == 0:
+                    obj.set_megatron_data_args()
+                    dataloaders = megatron_lm_prepare_data_loader(self, obj)
+                result.append(dataloaders[counter])
+                counter += 1
+            else:
+                result.append(obj)
+
+        if model is not None:
+            model = megatron_lm_prepare_model(self)
+        if optimizer is not None:
+            optimizer = megatron_lm_prepare_optimizer(self, model)
+        if scheduler is not None:
+            scheduler = megatron_lm_prepare_scheduler(self, optimizer, scheduler)
+
+        if model is not None:
+            model = MegatronEngine(self, model, optimizer, scheduler)
+        if optimizer is not None:
+            optimizer = MegatronLMOptimizerWrapper(optimizer)
+        if scheduler is not None:
+            scheduler = MegatronLMSchedulerWrapper(scheduler, optimizer)
+
+        for i in range(len(result)):
+            if isinstance(result[i], torch.nn.Module):
+                result[i] = model
+            elif isinstance(result[i], torch.optim.Optimizer):
+                result[i] = optimizer
+            elif isinstance(result[i], MegatronLMDummyScheduler):
+                result[i] = scheduler
+        if model is not None:
+            self._models.append(model)
+        if optimizer is not None:
+            self._optimizers.append(optimizer)
+        if scheduler is not None:
+            self._schedulers.append(scheduler)
+        if len(self._models) > 1:
+            raise AssertionError(
+                "You can't use same `Accelerator()` instance with multiple models when using Megatron-LM"
+            )
+        return tuple(result)
+
     def prepare_data_loader(self, data_loader: torch.utils.data.DataLoader, device_placement=None):
         """
         Prepares a PyTorch DataLoader for training in any distributed setup. It is recommended to use
@@ -1001,6 +1163,8 @@ def backward(self, loss, **kwargs):
             loss = loss / self.gradient_accumulation_steps
         if self.distributed_type == DistributedType.DEEPSPEED:
             self.deepspeed_engine_wrapped.backward(loss, **kwargs)
+        elif self.distributed_type == DistributedType.MEGATRON_LM:
+            return
         elif self.scaler is not None:
             self.scaler.scale(loss).backward(**kwargs)
         else:
@@ -1312,6 +1476,10 @@ def save_state(self, output_dir: str):
                 ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}"
                 model.save_checkpoint(output_dir, ckpt_id)
                 logger.info(f"DeepSpeed Model and Optimizer saved to output dir {os.path.join(output_dir, ckpt_id)}")
+            elif self.distributed_type == DistributedType.MEGATRON_LM:
+                logger.info("Saving Megatron-LM Model, Optimizer and Scheduler")
+                model.save_checkpoint(output_dir)
+                logger.info(f"Megatron-LM Model , Optimizer and Scheduler saved to output dir {output_dir}")
             else:
                 weights.append(self.get_state_dict(model, unwrap=False))
 
@@ -1322,7 +1490,7 @@ def save_state(self, output_dir: str):
                 logger.info("Saving FSDP Optimizer")
                 self.state.fsdp_plugin.save_optimizer(self, opt, self._models[i], output_dir, i)
                 logger.info(f"FSDP Optimizer saved to output dir {output_dir}")
-        elif self.distributed_type != DistributedType.DEEPSPEED:
+        elif self.distributed_type not in [DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]:
             optimizers = self._optimizers
 
         # Save the lr schedulers taking care of DeepSpeed nuances
@@ -1332,7 +1500,7 @@ def save_state(self, output_dir: str):
                 if isinstance(scheduler, DeepSpeedSchedulerWrapper):
                     continue
                 schedulers.append(scheduler)
-        else:
+        elif self.distributed_type not in [DistributedType.MEGATRON_LM]:
             schedulers = self._schedulers
 
         save_location = save_accelerator_state(
@@ -1374,6 +1542,10 @@ def load_state(self, input_dir: str):
                 ckpt_id = f"{MODEL_NAME}" if i == 0 else f"{MODEL_NAME}_{i}"
                 model.load_checkpoint(input_dir, ckpt_id)
                 logger.info(f"DeepSpeed Model and Optimizer loaded from input dir {os.path.join(input_dir, ckpt_id)}")
+            elif self.distributed_type == DistributedType.MEGATRON_LM:
+                logger.info("Loading Megatron-LM Model, Optimizer and Scheduler")
+                model.load_checkpoint(input_dir)
+                logger.info(f"Megatron-LM Model , Optimizer and Scheduler loaded from input dir {input_dir}")
             else:
                 models.append(model)
 
@@ -1384,7 +1556,7 @@ def load_state(self, input_dir: str):
                 logger.info("Loading FSDP Optimizer")
                 self.state.fsdp_plugin.load_optimizer(self, opt, self._models[i], input_dir, i)
                 logger.info(f"FSDP Optimizer loaded from input dir {input_dir}")
-        elif self.distributed_type != DistributedType.DEEPSPEED:
+        elif self.distributed_type not in [DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]:
             optimizers = self._optimizers
 
         # Load the lr schedulers taking care of DeepSpeed nuances
@@ -1394,7 +1566,7 @@ def load_state(self, input_dir: str):
                 if isinstance(scheduler, DeepSpeedSchedulerWrapper):
                     continue
                 schedulers.append(scheduler)
-        else:
+        elif self.distributed_type not in [DistributedType.MEGATRON_LM]:
             schedulers = self._schedulers
 
         load_accelerator_state(input_dir, models, optimizers, schedulers, self.state.process_index, self.scaler)
diff --git a/src/accelerate/checkpointing.py b/src/accelerate/checkpointing.py
index d5e816aab..00dfce73e 100644
--- a/src/accelerate/checkpointing.py
+++ b/src/accelerate/checkpointing.py
@@ -155,15 +155,18 @@ def load_accelerator_state(input_dir, models, optimizers, schedulers, process_in
         logger.info("GradScaler state loaded successfully")
 
     # Random states
-    states = torch.load(os.path.join(input_dir, f"{RNG_STATE_NAME}_{process_index}.pkl"))
-    random.setstate(states["random_state"])
-    np.random.set_state(states["numpy_random_seed"])
-    torch.set_rng_state(states["torch_manual_seed"])
-    torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"])
-    # ^^ safe to call this function even if cuda is not available
-    if is_tpu_available():
-        xm.set_rng_state(states["xm_seed"])
-    logger.info("All random states loaded successfully")
+    try:
+        states = torch.load(os.path.join(input_dir, f"{RNG_STATE_NAME}_{process_index}.pkl"))
+        random.setstate(states["random_state"])
+        np.random.set_state(states["numpy_random_seed"])
+        torch.set_rng_state(states["torch_manual_seed"])
+        torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"])
+        # ^^ safe to call this function even if cuda is not available
+        if is_tpu_available():
+            xm.set_rng_state(states["xm_seed"])
+        logger.info("All random states loaded successfully")
+    except:
+        logger.info("Could not load random states")
 
 
 def save_custom_state(obj, path, index: int = 0):
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
index e86a0abff..0fa3ceab2 100644
--- a/src/accelerate/commands/config/cluster.py
+++ b/src/accelerate/commands/config/cluster.py
@@ -275,6 +275,67 @@ def get_cluster_input():
                 default="FULL_STATE_DICT",
             )
 
+    megatron_lm_config = {}
+    if distributed_type in [DistributedType.MULTI_GPU]:
+        use_megatron_lm = _ask_field(
+            "Do you want to use Megatron-LM ? [yes/NO]: ",
+            _convert_yes_no_to_bool,
+            default=False,
+            error_message="Please enter yes or no.",
+        )
+        if use_megatron_lm:
+            distributed_type = DistributedType.MEGATRON_LM
+        if distributed_type == DistributedType.MEGATRON_LM:
+            prefix = "megatron_lm_"
+            megatron_lm_config[prefix + "tp_degree"] = _ask_field(
+                "What is the Tensor Parallelism degree/size? [1]:",
+                lambda x: int(x),
+                default=1,
+                error_message="Please enter an integer.",
+            )
+            if megatron_lm_config[prefix + "tp_degree"] > 1:
+                megatron_lm_config[prefix + "sequence_parallelism"] = _ask_field(
+                    "Do you want to enable Sequence Parallelism? [YES/no]: ",
+                    _convert_yes_no_to_bool,
+                    default=True,
+                    error_message="Please enter yes or no.",
+                )
+
+            megatron_lm_config[prefix + "pp_degree"] = _ask_field(
+                "What is the Pipeline Parallelism degree/size? [1]:",
+                lambda x: int(x),
+                default=1,
+                error_message="Please enter an integer.",
+            )
+            if megatron_lm_config[prefix + "pp_degree"] > 1:
+                megatron_lm_config[prefix + "num_micro_batches"] = _ask_field(
+                    "What is the number of micro-batches? [1]:",
+                    lambda x: int(x),
+                    default=1,
+                    error_message="Please enter an integer.",
+                )
+
+            megatron_lm_config[prefix + "recompute_activations"] = _ask_field(
+                "Do you want to enable selective activation recomputation? [YES/no]: ",
+                _convert_yes_no_to_bool,
+                default=True,
+                error_message="Please enter yes or no.",
+            )
+
+            megatron_lm_config[prefix + "use_distributed_optimizer"] = _ask_field(
+                "Do you want to use distributed optimizer "
+                "which shards optimizer state and gradients across data pralellel ranks? [YES/no]: ",
+                _convert_yes_no_to_bool,
+                default=True,
+                error_message="Please enter yes or no.",
+            )
+
+            megatron_lm_config[prefix + "gradient_clipping"] = _ask_field(
+                "What is the gradient clipping value based on global L2 Norm (0 to disable)? [1.0]: ",
+                lambda x: float(x),
+                default=1.0,
+            )
+
     if distributed_type == DistributedType.TPU:
         main_training_function = _ask_field(
             "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
@@ -295,7 +356,7 @@ def get_cluster_input():
             default=1,
             error_message="Please enter an integer.",
         )
-    elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED]:
+    elif distributed_type in [DistributedType.FSDP, DistributedType.DEEPSPEED, DistributedType.MEGATRON_LM]:
         num_processes = _ask_field(
             "How many GPU(s) should be used for distributed training? [1]:",
             lambda x: int(x),
@@ -343,6 +404,7 @@ def get_cluster_input():
         main_training_function=main_training_function,
         deepspeed_config=deepspeed_config,
         fsdp_config=fsdp_config,
+        megatron_lm_config=megatron_lm_config,
         use_cpu=use_cpu,
         rdzv_backend=rdzv_backend,
         same_network=same_network,
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
index 43faf45ca..159f951dd 100644
--- a/src/accelerate/commands/config/config_args.py
+++ b/src/accelerate/commands/config/config_args.py
@@ -146,6 +146,8 @@ class ClusterConfig(BaseConfig):
     deepspeed_config: dict = None
     # args for fsdp
     fsdp_config: dict = None
+    # args for megatron_lm
+    megatron_lm_config: dict = None
     # args for TPU
     downcast_bf16: bool = False
 
@@ -154,6 +156,8 @@ def __post_init__(self):
             self.deepspeed_config = {}
         if self.fsdp_config is None:
             self.fsdp_config = {}
+        if self.megatron_lm_config is None:
+            self.megatron_lm_config = {}
         return super().__post_init__()
 
 
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index f348bb0a6..c297223cb 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -83,6 +83,8 @@ def launch_command_parser(subparsers=None):
         action="store_true",
         help="Whether or not this should use MPS-enabled GPU device on MacOS machines.",
     )
+
+    # deepspeed args
     parser.add_argument(
         "--use_deepspeed",
         default=False,
@@ -163,6 +165,8 @@ def launch_command_parser(subparsers=None):
         type=str,
         help="DeepSpeed multi-node launcher to use.",
     )
+
+    # fsdp args
     parser.add_argument(
         "--use_fsdp",
         default=False,
@@ -236,6 +240,62 @@ def launch_command_parser(subparsers=None):
         type=str,
         help="This argument is deprecated. Use `fsdp_transformer_layer_cls_to_wrap` instead.",
     )
+
+    # megatron_lm args
+    parser.add_argument(
+        "--use_megatron_lm",
+        default=False,
+        action="store_true",
+        help="Whether to use Megatron-LM.",
+    )
+    parser.add_argument(
+        "--megatron_lm_tp_degree",
+        type=int,
+        default=1,
+        help="Megatron-LM's Tensor Parallelism (TP) degree. (useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_pp_degree",
+        type=int,
+        default=1,
+        help="Megatron-LM's Pipeline Parallelism (PP) degree. (useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_num_micro_batches",
+        type=int,
+        default=None,
+        help="Megatron-LM's number of micro batches when PP degree > 1. (useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_sequence_parallelism",
+        default=None,
+        type=str,
+        help="Decides Whether (true|false) to enable Sequence Parallelism when TP degree > 1. "
+        "(useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_recompute_activations",
+        default=None,
+        type=str,
+        help="Decides Whether (true|false) to enable Selective Activation Recomputation. "
+        "(useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_use_distributed_optimizer",
+        default=None,
+        type=str,
+        help="Decides Whether (true|false) to use distributed optimizer "
+        "which shards optimizer state and gradients across Data Pralellel (DP) ranks. "
+        "(useful only when `use_megatron_lm` flag is passed).",
+    )
+    parser.add_argument(
+        "--megatron_lm_gradient_clipping",
+        default=1.0,
+        type=float,
+        help="Megatron-LM's gradient clipping value based on global L2 Norm (0 to disable). "
+        "(useful only when `use_megatron_lm` flag is passed).",
+    )
+
     parser.add_argument(
         "--tpu", default=False, action="store_true", help="Whether or not this should launch a TPU training."
     )
@@ -487,6 +547,22 @@ def multi_gpu_launcher(args):
             current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch_policy)
         if args.fsdp_state_dict_type is not None:
             current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
+
+    if args.use_megatron_lm:
+        prefix = "MEGATRON_LM_"
+        current_env["USE_MEGATRON_LM"] = "true"
+        current_env[prefix + "TP_DEGREE"] = str(args.megatron_lm_tp_degree)
+        current_env[prefix + "PP_DEGREE"] = str(args.megatron_lm_pp_degree)
+        current_env[prefix + "GRADIENT_CLIPPING"] = str(args.megatron_lm_gradient_clipping)
+        if args.megatron_lm_num_micro_batches is not None:
+            current_env[prefix + "NUM_MICRO_BATCHES"] = str(args.megatron_lm_num_micro_batches)
+        if args.megatron_lm_sequence_parallelism is not None:
+            current_env[prefix + "SEQUENCE_PARALLELISM"] = str(args.megatron_lm_sequence_parallelism)
+        if args.megatron_lm_recompute_activations is not None:
+            current_env[prefix + "RECOMPUTE_ACTIVATIONS"] = str(args.megatron_lm_recompute_activations)
+        if args.megatron_lm_use_distributed_optimizer is not None:
+            current_env[prefix + "USE_DISTRIBUTED_OPTIMIZER"] = str(args.megatron_lm_use_distributed_optimizer)
+
     current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
     if is_torch_version("<", "1.9.0"):
         raise NotImplementedError("Multi-node training requires pytorch>=1.9.0")
@@ -825,12 +901,14 @@ def launch_command(args):
             and not args.use_deepspeed
             and not args.use_fsdp
             and not args.use_mps_device
+            and not args.use_megatron_lm
         ):
             args.use_deepspeed = defaults.distributed_type == DistributedType.DEEPSPEED
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
             args.use_fsdp = defaults.distributed_type == DistributedType.FSDP
             args.use_mps_device = defaults.distributed_type == DistributedType.MPS
+            args.use_megatron_lm = defaults.distributed_type == DistributedType.MEGATRON_LM
         if not args.use_mps_device:
             if args.gpu_ids is None:
                 if defaults.gpu_ids is not None:
@@ -851,6 +929,8 @@ def launch_command(args):
                         if "fsdp" not in arg_to_set:
                             arg_to_set = "fsdp_" + arg_to_set
                         setattr(args, arg_to_set, defaults.fsdp_config[k])
+                    for k in defaults.megatron_lm_config:
+                        setattr(args, k, defaults.megatron_lm_config[k])
                     continue
 
                 # Those args are handled separately
@@ -901,6 +981,8 @@ def launch_command(args):
         deepspeed_launcher(args)
     elif args.use_fsdp and not args.cpu:
         multi_gpu_launcher(args)
+    elif args.use_megatron_lm and not args.cpu:
+        multi_gpu_launcher(args)
     elif args.multi_gpu and not args.cpu:
         multi_gpu_launcher(args)
     elif args.tpu and not args.cpu:
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
index fac1b50e2..c4a015156 100644
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -625,7 +625,7 @@ def prepare_data_loader(
     sampler_is_batch_sampler = False
     generator = getattr(dataloader, "generator", None)
     # No change if no multiprocess
-    if num_processes != 1 and not dispatch_batches:
+    if (num_processes != 1 or state.distributed_type == DistributedType.MEGATRON_LM) and not dispatch_batches:
         if isinstance(new_dataset, IterableDataset):
             if getattr(dataloader.dataset, "generator", None) is not None:
                 generator = dataloader.dataset.generator
diff --git a/src/accelerate/logging.py b/src/accelerate/logging.py
index 2c67e24eb..dca093215 100644
--- a/src/accelerate/logging.py
+++ b/src/accelerate/logging.py
@@ -15,6 +15,7 @@
 import logging
 
 from .state import AcceleratorState
+from .utils import DistributedType
 
 
 class MultiProcessAdapter(logging.LoggerAdapter):
@@ -28,7 +29,12 @@ class MultiProcessAdapter(logging.LoggerAdapter):
     @staticmethod
     def _should_log(main_process_only):
         "Check if log should be performed"
-        return not main_process_only or (main_process_only and AcceleratorState().local_process_index == 0)
+        state = AcceleratorState()
+        if state.distributed_type != DistributedType.MEGATRON_LM:
+            process_index_flag = state.local_process_index == 0
+        else:
+            process_index_flag = state.process_index == state.num_processes - 1
+        return not main_process_only or (main_process_only and process_index_flag)
 
     def log(self, level, msg, *args, **kwargs):
         """
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index a700cfad6..82ea288f0 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -70,6 +70,7 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin=None,
         fsdp_plugin=None,
+        megatron_lm_plugin=None,
         _from_accelerator: bool = False,
         **kwargs,
     ):
@@ -162,6 +163,10 @@ def __init__(
                     if self.mixed_precision != "no":
                         fsdp_plugin.set_mixed_precision(self.mixed_precision)
                     self.fsdp_plugin = fsdp_plugin
+                if os.environ.get("USE_MEGATRON_LM", "false") == "true":
+                    self.distributed_type = DistributedType.MEGATRON_LM
+                    megatron_lm_plugin.set_mixed_precision(self.mixed_precision)
+                    self.megatron_lm_plugin = megatron_lm_plugin
             elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
                 self.distributed_type = DistributedType.MULTI_CPU
                 if is_ccl_available() and get_int_from_env(["CCL_WORKER_COUNT"], 0) > 0:
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 4bc71ba83..b472ec556 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -13,6 +13,7 @@
     InitProcessGroupKwargs,
     KwargsHandler,
     LoggerType,
+    MegatronLMPlugin,
     PrecisionType,
     RNGType,
     SageMakerDistributedType,
@@ -28,6 +29,7 @@
     is_comet_ml_available,
     is_datasets_available,
     is_deepspeed_available,
+    is_megatron_lm_available,
     is_rich_available,
     is_sagemaker_available,
     is_tensorboard_available,
@@ -94,6 +96,23 @@
     )
 
 from .launch import PrepareForLaunch, _filter_args, get_launch_prefix
+from .megatron_lm import (
+    AbstractTrainStep,
+    BertTrainStep,
+    GPTTrainStep,
+    MegatronEngine,
+    MegatronLMDummyDataLoader,
+    MegatronLMDummyScheduler,
+    MegatronLMOptimizerWrapper,
+    MegatronLMSchedulerWrapper,
+    T5TrainStep,
+    avg_losses_across_data_parallel_group,
+)
+from .megatron_lm import initialize as megatron_lm_initialize
+from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
+from .megatron_lm import prepare_model as megatron_lm_prepare_model
+from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
+from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
 from .memory import find_executable_batch_size
 from .other import (
     extract_model_from_parallel,
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 9f6e30bec..8227fa5bc 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -16,6 +16,7 @@
 General namespace and dataclass related classes
 """
 
+import argparse
 import copy
 import enum
 import functools
@@ -24,7 +25,8 @@
 import warnings
 from dataclasses import dataclass, field
 from datetime import timedelta
-from typing import Any, Callable, Iterable, Optional
+from distutils.util import strtobool
+from typing import Any, Callable, Dict, Iterable, List, Optional
 
 import torch
 
@@ -127,6 +129,7 @@ class DistributedType(str, enum.Enum):
     FSDP = "FSDP"
     TPU = "TPU"
     MPS = "MPS"
+    MEGATRON_LM = "MEGATRON_LM"
 
 
 class SageMakerDistributedType(str, enum.Enum):
@@ -644,3 +647,351 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+    train_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of iterations to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    train_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of samples to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    weight_decay_incr_style: str = field(
+        default="constant",
+        metadata={"help": 'Weight decay increment function. choices=["constant", "linear", "cosine"]. '},
+    )
+    start_weight_decay: float = field(
+        default=None,
+        metadata={"help": "Initial weight decay coefficient for L2 regularization."},
+    )
+    end_weight_decay: float = field(
+        default=None,
+        metadata={"help": "End of run weight decay coefficient for L2 regularization."},
+    )
+    lr_decay_style: str = field(
+        default="linear",
+        metadata={"help": "Learning rate decay function. choices=['constant', 'linear', 'cosine']."},
+    )
+    lr_decay_iters: int = field(
+        default=None,
+        metadata={"help": "Number of iterations for learning rate decay. If None defaults to `train_iters`."},
+    )
+    lr_decay_samples: int = field(
+        default=None,
+        metadata={"help": "Number of samples for learning rate decay. If None defaults to `train_samples`."},
+    )
+    lr_warmup_iters: int = field(
+        default=None,
+        metadata={"help": "number of iterations to linearly warmup learning rate over."},
+    )
+    lr_warmup_samples: int = field(
+        default=None,
+        metadata={"help": "number of samples to linearly warmup learning rate over."},
+    )
+    lr_warmup_fraction: float = field(
+        default=None,
+        metadata={"help": "fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over."},
+    )
+    min_lr: float = field(
+        default=0,
+        metadata={"help": "Minumum value for learning rate. The scheduler clip values below this threshold."},
+    )
+    consumed_samples: List[int] = field(
+        default=None,
+        metadata={
+            "help": "Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call."
+        },
+    )
+    no_wd_decay_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to disable weight decay."})
+    scale_lr_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to scale learning rate."})
+    lr_mult: float = field(default=1.0, metadata={"help": "Learning rate multiplier."})
+    megatron_dataset_flag: bool = field(
+        default=False,
+        metadata={"help": "Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format."},
+    )
+    seq_length: int = field(
+        default=None,
+        metadata={"help": "Maximum sequence length to process."},
+    )
+    encoder_seq_length: int = field(
+        default=None,
+        metadata={"help": "Maximum sequence length to process for the encoder."},
+    )
+    decoder_seq_length: int = field(
+        default=None,
+        metadata={"help": "Maximum sequence length to process for the decoder."},
+    )
+    tensorboard_dir: str = field(
+        default=None,
+        metadata={"help": "Path to save tensorboard logs."},
+    )
+    set_all_logging_options: bool = field(
+        default=False,
+        metadata={"help": "Whether to set all logging options."},
+    )
+    custom_train_step_class: Optional[Any] = field(
+        default=None,
+        metadata={"help": "Custom train step class."},
+    )
+    custom_train_step_kwargs: Optional[Dict[str, Any]] = field(
+        default=None,
+        metadata={"help": "Custom train step kwargs."},
+    )
+    eval_iters: int = field(
+        default=100, metadata={"help": "Number of iterations to run for evaluation validation/test for."}
+    )
+    eval_interval: int = field(
+        default=1000, metadata={"help": "Interval between running evaluation on validation set."}
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_LM_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1 or self.use_distributed_optimizer:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        if self.consumed_samples is not None:
+            if len(self.consumed_samples) == 1:
+                self.consumed_samples.extend([0, 0])
+            elif len(self.consumed_samples) == 2:
+                self.consumed_samples.append(0)
+
+        self.megatron_lm_default_args = {
+            "tensor_model_parallel_size": self.tp_degree,
+            "pipeline_model_parallel_size": self.pp_degree,
+            "pipeline_model_parallel_split_rank": self.pipeline_model_parallel_split_rank,
+            "num_layers_per_virtual_pipeline_stage": self.num_layers_per_virtual_pipeline_stage,
+            "DDP_impl": self.DDP_impl,
+            "use_distributed_optimizer": self.use_distributed_optimizer,
+            "sequence_parallel": self.sequence_parallelism,
+            "clip_grad": self.gradient_clipping,
+            "num_micro_batches": self.num_micro_batches,
+            "consumed_samples": self.consumed_samples,
+            "no_wd_decay_cond": self.no_wd_decay_cond,
+            "scale_lr_cond": self.scale_lr_cond,
+            "lr_mult": self.lr_mult,
+            "megatron_dataset_flag": self.megatron_dataset_flag,
+            "eval_iters": self.eval_iters,
+            "eval_interval": self.eval_interval,
+        }
+        if self.recompute_activation:
+            self.megatron_lm_default_args["recompute_granularity"] = "selective"
+        if self.tensorboard_dir is not None:
+            self.megatron_lm_default_args["tensorboard_dir"] = self.tensorboard_dir
+            if self.set_all_logging_options:
+                self.set_tensorboard_logging_options()
+
+    def set_network_size_args(self, model, batch_data=None):
+        # Check if the model is either BERT, GPT or T5 else raise error
+        # set 'num_layers', 'hidden_size', 'num_attention_heads', 'max_position_embeddings'
+        if "megatron-bert" in model.config.model_type.lower():
+            model_type_name = "bert"
+            num_layers = model.config.num_hidden_layers
+            hidden_size = model.config.hidden_size
+            num_attention_heads = model.config.num_attention_heads
+            max_position_embeddings = model.config.max_position_embeddings
+            num_labels = model.config.num_labels
+            orig_vocab_size = model.config.vocab_size
+            if "maskedlm" in model.__class__.__name__.lower():
+                pretraining_flag = True
+            if self.seq_length is not None:
+                if self.encoder_seq_length is not None:
+                    warnings.warn("Both `seq_length` and `encoder_seq_length` are set. Using `encoder_seq_length`.")
+                self.seq_length = self.encoder_seq_length
+            elif self.encoder_seq_length is not None:
+                self.seq_length = self.encoder_seq_length
+            elif batch_data is not None:
+                self.seq_length = batch_data["input_ids"].shape[1]
+            else:
+                self.seq_length = max_position_embeddings
+            self.megatron_lm_default_args["seq_length"] = self.seq_length
+        elif "gpt2" in model.config.model_type.lower():
+            model_type_name = "gpt"
+            num_layers = model.config.n_layer
+            hidden_size = model.config.n_embd
+            num_attention_heads = model.config.n_head
+            max_position_embeddings = model.config.n_positions
+            orig_vocab_size = model.config.vocab_size
+            pretraining_flag = True
+            if self.seq_length is not None:
+                if self.decoder_seq_length is not None:
+                    warnings.warn("Both `seq_length` and `decoder_seq_length` are set. Using `decoder_seq_length`.")
+                self.seq_length = self.decoder_seq_length
+            elif self.decoder_seq_length is not None:
+                self.seq_length = self.decoder_seq_length
+            elif batch_data is not None:
+                self.seq_length = batch_data["input_ids"].shape[1]
+            else:
+                self.seq_length = max_position_embeddings
+            self.megatron_lm_default_args["seq_length"] = self.seq_length
+        elif "t5" in model.config.model_type.lower():
+            model_type_name = "t5"
+            num_layers = model.config.num_layers
+            hidden_size = model.config.d_model
+            num_attention_heads = model.config.num_heads
+            max_position_embeddings = model.config.n_positions if hasattr(model.config, "n_positions") else 1024
+            orig_vocab_size = model.config.vocab_size
+            pretraining_flag = True
+            if self.encoder_seq_length is None:
+                if batch_data is not None:
+                    self.encoder_seq_length = batch_data["input_ids"].shape[1]
+                else:
+                    self.encoder_seq_length = max_position_embeddings
+            if self.decoder_seq_length is None:
+                if batch_data is not None:
+                    self.decoder_seq_length = batch_data["labels"].shape[1]
+                else:
+                    self.decoder_seq_length = max_position_embeddings
+
+            self.megatron_lm_default_args["encoder_seq_length"] = self.encoder_seq_length
+            self.megatron_lm_default_args["decoder_seq_length"] = self.decoder_seq_length
+        else:
+            raise ValueError(
+                "🤗 Accelerate Megatron-LM integration supports only BERT, GPT and T5 model. "
+                "Please check the model you are using is one of those."
+            )
+
+        self.megatron_lm_default_args["model_type_name"] = model_type_name
+        self.megatron_lm_default_args["num_layers"] = num_layers
+        self.megatron_lm_default_args["hidden_size"] = hidden_size
+        self.megatron_lm_default_args["num_attention_heads"] = num_attention_heads
+        self.megatron_lm_default_args["max_position_embeddings"] = max_position_embeddings
+        self.megatron_lm_default_args["pretraining_flag"] = pretraining_flag
+        self.megatron_lm_default_args["orig_vocab_size"] = orig_vocab_size
+        self.megatron_lm_default_args["model_return_dict"] = model.config.return_dict
+        if model_type_name == "bert":
+            self.megatron_lm_default_args["num_labels"] = num_labels
+
+    def set_mixed_precision(self, mixed_precision):
+        if mixed_precision == "fp16":
+            self.megatron_lm_default_args["fp16"] = True
+        elif mixed_precision == "bf16":
+            self.megatron_lm_default_args["bf16"] = True
+            self.DDP_impl = "local"
+            self.megatron_lm_default_args["DDP_impl"] = self.DDP_impl
+
+    def set_training_args(self, micro_batch_size, dp_degree):
+        self.data_parallel_size = dp_degree
+        self.micro_batch_size = micro_batch_size
+        self.global_batch_size = dp_degree * micro_batch_size * self.num_micro_batches
+        self.megatron_lm_default_args["data_parallel_size"] = self.data_parallel_size
+        self.megatron_lm_default_args["micro_batch_size"] = self.micro_batch_size
+        self.megatron_lm_default_args["global_batch_size"] = self.global_batch_size
+
+    def set_optimizer_type(self, optimizer):
+        optimizer_name = optimizer.__class__.__name__.lower()
+        if "adam" in optimizer_name:
+            self.megatron_lm_default_args["optimizer"] = "adam"
+            self.megatron_lm_default_args["adam_beta1"] = optimizer.defaults["betas"][0]
+            self.megatron_lm_default_args["adam_beta2"] = optimizer.defaults["betas"][1]
+            self.megatron_lm_default_args["adam_eps"] = optimizer.defaults["eps"]
+        elif "sgd" in optimizer_name:
+            self.megatron_lm_default_args["optimizer"] = "sgd"
+            self.megatron_lm_default_args["sgd_momentum"] = optimizer.defaults["momentum"]
+        else:
+            raise ValueError(f"Optimizer {optimizer_name} is not supported by Megatron-LM")
+
+        self.megatron_lm_default_args["lr"] = optimizer.defaults["lr"]
+        self.megatron_lm_default_args["weight_decay"] = optimizer.defaults["weight_decay"]
+
+    def set_scheduler_args(self, scheduler):
+        if self.train_iters is None:
+            self.train_iters = scheduler.total_num_steps // self.megatron_lm_default_args["data_parallel_size"]
+            if self.train_samples is not None:
+                self.train_samples = None
+                warnings.warn(
+                    "Ignoring `train_samples` as `train_iters` based on scheduler is being used for training."
+                )
+        if self.lr_warmup_iters is None:
+            self.lr_warmup_iters = scheduler.warmup_num_steps // self.megatron_lm_default_args["data_parallel_size"]
+            if self.lr_warmup_samples is not None:
+                warnings.warn(
+                    "Ignoring `lr_warmup_samples` as `lr_warmup_iters` based on scheduler is being used for training."
+                )
+            self.lr_warmup_samples = 0
+
+        self.megatron_lm_default_args["train_iters"] = self.train_iters
+        self.megatron_lm_default_args["lr_warmup_iters"] = self.lr_warmup_iters
+        self.megatron_lm_default_args["train_samples"] = self.train_samples
+        self.megatron_lm_default_args["lr_warmup_samples"] = self.lr_warmup_samples
+        self.megatron_lm_default_args["lr_decay_iters"] = self.lr_decay_iters
+        self.megatron_lm_default_args["lr_decay_samples"] = self.lr_decay_samples
+        self.megatron_lm_default_args["lr_warmup_fraction"] = self.lr_warmup_fraction
+        self.megatron_lm_default_args["lr_decay_style"] = self.lr_decay_style
+        self.megatron_lm_default_args["weight_decay_incr_style"] = self.weight_decay_incr_style
+        self.megatron_lm_default_args["start_weight_decay"] = self.start_weight_decay
+        self.megatron_lm_default_args["end_weight_decay"] = self.end_weight_decay
+        self.megatron_lm_default_args["min_lr"] = self.min_lr
+
+    def set_tensorboard_logging_options(self):
+        from megatron.arguments import _add_logging_args
+
+        parser = argparse.ArgumentParser()
+        parser = _add_logging_args(parser)
+        logging_args = parser.parse_known_args()
+        self.dataset_args = vars(logging_args[0])
+        for key, value in self.dataset_args.items():
+            if key.startswith("log_"):
+                self.megatron_lm_default_args[key] = True
+            elif key.startswith("no_log_"):
+                self.megatron_lm_default_args[key.replace("no_", "")] = True
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index eb5fda468..58009d03b 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -87,6 +87,10 @@ def is_bf16_available(ignore_tpu=False):
     return False
 
 
+def is_megatron_lm_available():
+    return importlib.util.find_spec("megatron") is not None
+
+
 def is_transformers_available():
     return importlib.util.find_spec("transformers") is not None
 
diff --git a/src/accelerate/utils/megatron_lm.py b/src/accelerate/utils/megatron_lm.py
new file mode 100644
index 000000000..cb988e360
--- /dev/null
+++ b/src/accelerate/utils/megatron_lm.py
@@ -0,0 +1,1199 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import math
+from abc import ABC
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ..optimizer import AcceleratedOptimizer
+from ..scheduler import AcceleratedScheduler
+from .imports import is_megatron_lm_available, is_transformers_available
+from .operations import send_to_device
+
+
+if is_transformers_available():
+    from transformers.modeling_outputs import (
+        CausalLMOutputWithCrossAttentions,
+        Seq2SeqLMOutput,
+        SequenceClassifierOutput,
+    )
+
+
+if is_megatron_lm_available():
+    from megatron import (
+        get_args,
+        get_num_microbatches,
+        get_tensorboard_writer,
+        get_timers,
+        mpu,
+        print_rank_0,
+        print_rank_last,
+    )
+    from megatron.arguments import _add_data_args, _add_validation_args, parse_args, validate_args
+    from megatron.checkpointing import load_args_from_checkpoint, load_checkpoint, save_checkpoint
+    from megatron.data.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
+    from megatron.global_vars import set_global_variables
+    from megatron.initialize import (
+        _compile_dependencies,
+        _init_autoresume,
+        _set_random_seed,
+        set_jit_fusion_options,
+        write_args_to_tensorboard,
+    )
+    from megatron.model import BertModel, GPTModel, ModelType, T5Model
+    from megatron.model.classification import Classification
+    from megatron.optimizer import get_megatron_optimizer
+    from megatron.schedules import get_forward_backward_func
+    from megatron.tokenizer.tokenizer import _vocab_size_with_padding
+    from megatron.training import get_model, get_optimizer_param_scheduler, training_log
+    from megatron.utils import (
+        average_losses_across_data_parallel_group,
+        calc_params_l2_norm,
+        get_ltor_masks_and_position_ids,
+    )
+
+
+# model utilities
+def model_provider_func(pre_process=True, post_process=True, add_encoder=True, add_decoder=True):
+    """Build the model."""
+    args = get_args()
+    mode = "pre-training" if args.pretraining_flag else "fine-tuning"
+    if args.rank == 0:
+        print(f"Building {args.model_type_name} model in the {mode} mode.")
+        print(
+            "The Megatron LM model weights are initialized at random in `accelerator.prepare`. "
+            "Please use `accelerator.load_checkpoint` to load a pre-trained checkpoint matching the distributed setup."
+        )
+    if args.model_type_name == "bert":
+        if args.pretraining_flag:
+            num_tokentypes = 2 if args.bert_binary_head else 0
+            model = BertModel(
+                num_tokentypes=num_tokentypes,
+                add_binary_head=args.bert_binary_head,
+                parallel_output=True,
+                pre_process=pre_process,
+                post_process=post_process,
+            )
+        else:
+
+            model = Classification(
+                num_classes=args.num_labels, num_tokentypes=2, pre_process=pre_process, post_process=post_process
+            )
+    elif args.model_type_name == "gpt":
+        model = GPTModel(num_tokentypes=0, parallel_output=True, pre_process=pre_process, post_process=post_process)
+    elif args.model_type_name == "t5":
+        model = T5Model(
+            num_tokentypes=0,
+            parallel_output=True,
+            pre_process=pre_process,
+            post_process=post_process,
+            add_encoder=add_encoder,
+            add_decoder=add_decoder,
+        )
+    else:
+        raise ValueError(f"Unsupported model type: {args.model_type_name}")
+    return model
+
+
+def prepare_model(accelerator):
+    accelerator.print("Preparing model")
+    args = get_args()
+    if args.model_type_name == "bert" or args.model_type_name == "gpt":
+        model_type = ModelType.encoder_or_decoder
+    elif args.model_type_name == "t5":
+        model_type = ModelType.encoder_and_decoder
+        if args.pipeline_model_parallel_split_rank is None and args.pipeline_model_parallel_size > 1:
+            args.pipeline_model_parallel_split_rank = args.pipeline_model_parallel_size // 2
+    model = get_model(model_provider_func, model_type)
+    return model
+
+
+# dataloader utilities
+class MegatronLMDummyDataLoader:
+    """
+    Dummy dataloader presents model parameters or param groups, this is primarily used to follow conventional training
+
+    Args:
+        **dataset_kwargs: Megatron data arguments.
+    """
+
+    def __init__(self, **dataset_kwargs):
+        parser = argparse.ArgumentParser()
+        parser = _add_data_args(parser)
+        parser = _add_validation_args(parser)
+        data_args = parser.parse_known_args()
+        self.dataset_args = vars(data_args[0])
+        self.dataset_args.update(dataset_kwargs)
+        self.dataset_args["megatron_dataset_flag"] = True
+
+    def set_megatron_data_args(self):
+        args = get_args()
+        for key, value in self.dataset_args.items():
+            setattr(args, key, value)
+
+    def get_train_valid_test_datasets_provider(self):
+        def train_valid_test_datasets_provider(train_val_test_num_samples):
+            """Build train, valid, and test datasets."""
+            args = get_args()
+            dataset_args = {
+                "data_prefix": args.data_path,
+                "data_impl": args.data_impl,
+                "splits_string": args.split,
+                "train_valid_test_num_samples": train_val_test_num_samples,
+                "skip_warmup": (not args.mmap_warmup),
+                "seed": args.seed,
+            }
+            if args.model_type_name == "bert":
+                dataset_args.update(
+                    {
+                        "max_seq_length": args.seq_length,
+                        "masked_lm_prob": args.mask_prob,
+                        "short_seq_prob": args.short_seq_prob,
+                        "binary_head": args.bert_binary_head,
+                    }
+                )
+            elif args.model_type_name == "gpt":
+                dataset_args.update(
+                    {
+                        "seq_length": args.seq_length,
+                    }
+                )
+            elif args.model_type_name == "t5":
+                dataset_args.update(
+                    {
+                        "max_seq_length": args.encoder_seq_length,
+                        "max_seq_length_dec": args.decoder_seq_length,
+                        "masked_lm_prob": args.mask_prob,
+                        "short_seq_prob": args.short_seq_prob,
+                        "dataset_type": "t5",
+                    }
+                )
+            else:
+                raise ValueError(f"Unsupported model type: {args.model_type_name}")
+            if args.model_type_name == "gpt":
+                from megatron.data.gpt_dataset import build_train_valid_test_datasets
+            else:
+                from megatron.data.dataset_utils import build_train_valid_test_datasets
+            train_ds, valid_ds, test_ds = build_train_valid_test_datasets(**dataset_args)
+            return train_ds, valid_ds, test_ds
+
+        return train_valid_test_datasets_provider
+
+    def build_pretraining_data_loader(self, dataset, consumed_samples):
+        if dataset is None:
+            return None
+        args = get_args()
+        micro_batch_size = args.micro_batch_size * args.num_micro_batches
+
+        # Megatron sampler
+        if args.dataloader_type == "single":
+            batch_sampler = MegatronPretrainingSampler(
+                total_samples=len(dataset),
+                consumed_samples=consumed_samples,
+                micro_batch_size=micro_batch_size,
+                data_parallel_rank=mpu.get_data_parallel_rank(),
+                data_parallel_size=mpu.get_data_parallel_world_size(),
+            )
+        elif args.dataloader_type == "cyclic":
+            batch_sampler = MegatronPretrainingRandomSampler(
+                dataset,
+                total_samples=len(dataset),
+                consumed_samples=consumed_samples,
+                micro_batch_size=micro_batch_size,
+                data_parallel_rank=mpu.get_data_parallel_rank(),
+                data_parallel_size=mpu.get_data_parallel_world_size(),
+                data_sharding=args.data_sharding,
+            )
+        else:
+            raise Exception("{} dataloader type is not supported.".format(args.dataloader_type))
+
+        # Torch dataloader.
+        return torch.utils.data.DataLoader(
+            dataset, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True
+        )
+
+    def build_train_valid_test_data_iterators(self):
+        def cyclic_iter(iter):
+            while True:
+                for x in iter:
+                    yield x
+
+        args = get_args()
+
+        (train_dataloader, valid_dataloader, test_dataloader) = (None, None, None)
+
+        print_rank_0("> building train, validation, and test datasets ...")
+
+        # Backward compatibility, assume fixed batch size.
+        if args.iteration > 0 and args.consumed_train_samples == 0:
+            assert args.train_samples is None, "only backward compatiblity support for iteration-based training"
+            args.consumed_train_samples = args.iteration * args.global_batch_size
+        if args.iteration > 0 and args.consumed_valid_samples == 0:
+            if args.train_samples is None:
+                args.consumed_valid_samples = (
+                    (args.iteration // args.eval_interval) * args.eval_iters * args.global_batch_size
+                )
+
+        # Data loader only on rank 0 of each model parallel group.
+        if mpu.get_tensor_model_parallel_rank() == 0:
+
+            # Number of train/valid/test samples.
+            if args.train_samples:
+                train_samples = args.train_samples
+            else:
+                train_samples = args.train_iters * args.global_batch_size
+            eval_iters = (args.train_iters // args.eval_interval + 1) * args.eval_iters
+            test_iters = args.eval_iters
+            train_val_test_num_samples = [
+                train_samples,
+                eval_iters * args.global_batch_size,
+                test_iters * args.global_batch_size,
+            ]
+            print_rank_0(" > datasets target sizes (minimum size):")
+            print_rank_0("    train:      {}".format(train_val_test_num_samples[0]))
+            print_rank_0("    validation: {}".format(train_val_test_num_samples[1]))
+            print_rank_0("    test:       {}".format(train_val_test_num_samples[2]))
+
+            # Build the datasets.
+            train_valid_test_datasets_provider = self.get_train_valid_test_datasets_provider()
+            train_ds, valid_ds, test_ds = train_valid_test_datasets_provider(train_val_test_num_samples)
+
+            # Build dataloders.
+            train_dataloader = self.build_pretraining_data_loader(train_ds, args.consumed_train_samples)
+            valid_dataloader = self.build_pretraining_data_loader(valid_ds, args.consumed_valid_samples)
+            test_dataloader = self.build_pretraining_data_loader(test_ds, 0)
+
+            # Flags to know if we need to do training/validation/testing.
+            do_train = train_dataloader is not None and args.train_iters > 0
+            do_valid = valid_dataloader is not None and args.eval_iters > 0
+            do_test = test_dataloader is not None and args.eval_iters > 0
+            # Need to broadcast num_tokens and num_type_tokens.
+            flags = torch.cuda.LongTensor([int(do_train), int(do_valid), int(do_test)])
+        else:
+            flags = torch.cuda.LongTensor([0, 0, 0])
+
+        # Broadcast num tokens.
+        torch.distributed.broadcast(
+            flags, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group()
+        )
+        args.do_train = flags[0].item()
+        args.do_valid = flags[1].item()
+        args.do_test = flags[2].item()
+
+        # Build iterators.
+        dl_type = args.dataloader_type
+        assert dl_type in ["single", "cyclic"]
+
+        if train_dataloader is not None:
+            train_data_iterator = (
+                iter(train_dataloader) if dl_type == "single" else iter(cyclic_iter(train_dataloader))
+            )
+        else:
+            train_data_iterator = None
+
+        if valid_dataloader is not None:
+            valid_data_iterator = (
+                iter(valid_dataloader) if dl_type == "single" else iter(cyclic_iter(valid_dataloader))
+            )
+        else:
+            valid_data_iterator = None
+
+        if test_dataloader is not None:
+            test_data_iterator = iter(test_dataloader) if dl_type == "single" else iter(cyclic_iter(test_dataloader))
+        else:
+            test_data_iterator = None
+
+        return train_data_iterator, valid_data_iterator, test_data_iterator
+
+
+def prepare_data_loader(accelerator, dataloader):
+    accelerator.print("Preparing dataloader")
+    args = get_args()
+    if not args.megatron_dataset_flag:
+        from ..data_loader import _PYTORCH_DATALOADER_KWARGS, prepare_data_loader
+
+        args = get_args()
+        micro_batch_size = args.micro_batch_size * args.num_micro_batches
+        kwargs = {k: getattr(dataloader, k, _PYTORCH_DATALOADER_KWARGS[k]) for k in _PYTORCH_DATALOADER_KWARGS}
+        if kwargs["batch_size"] is None:
+            if isinstance(kwargs["sampler"], torch.utils.data.BatchSampler):
+                kwargs["sampler"].batch_size = micro_batch_size
+            else:
+                del kwargs["sampler"]
+                del kwargs["shuffle"]
+                del kwargs["batch_size"]
+                kwargs["batch_sampler"].batch_size = micro_batch_size
+        else:
+            del kwargs["batch_sampler"]
+            kwargs["batch_size"] = micro_batch_size
+
+        dataloader = torch.utils.data.DataLoader(dataloader.dataset, **kwargs)
+        return prepare_data_loader(
+            dataloader,
+            accelerator.device,
+            num_processes=mpu.get_data_parallel_world_size(),
+            process_index=mpu.get_data_parallel_rank(),
+            split_batches=accelerator.split_batches,
+            put_on_device=True,
+            rng_types=accelerator.rng_types.copy(),
+            dispatch_batches=accelerator.dispatch_batches,
+        )
+    else:
+        if args.consumed_samples is not None:
+            (
+                args.consumed_train_samples,
+                args.consumed_valid_samples,
+                args.consumed_test_samples,
+            ) = args.consumed_samples
+        else:
+            args.consumed_train_samples, args.consumed_valid_samples, args.consumed_test_samples = 0, 0, 0
+        (
+            train_data_iterator,
+            valid_data_iterator,
+            test_data_iterator,
+        ) = dataloader.build_train_valid_test_data_iterators()
+        return train_data_iterator, valid_data_iterator, test_data_iterator
+
+
+# optimizer utilities
+class MegatronLMOptimizerWrapper(AcceleratedOptimizer):
+    def __init__(self, optimizer):
+        super().__init__(optimizer, device_placement=False, scaler=None)
+
+    def zero_grad(self, set_to_none=None):
+        pass  # `model(**batch)` is doing that automatically. Therefore, it's implementation is not needed
+
+    def step(self):
+        pass  # `model(**batch)` is doing that automatically. Therefore, it's implementation is not needed
+
+    @property
+    def step_was_skipped(self):
+        """Whether or not the optimizer step was done, or skipped because of gradient overflow."""
+        return self.optimizer.skipped_iter
+
+
+def prepare_optimizer(accelerator, model):
+    accelerator.print("Preparing optimizer")
+    args = get_args()
+    optimizer = get_megatron_optimizer(model, args.no_wd_decay_cond, args.scale_lr_cond, args.lr_mult)
+    return optimizer
+
+
+# scheduler utilities
+class MegatronLMDummyScheduler:
+    """
+    Dummy scheduler presents model parameters or param groups, this is primarily used to follow conventional training
+    loop when scheduler config is specified in the deepspeed config file.
+
+    Args:
+        optimizer (`torch.optim.optimizer.Optimizer`):
+            The optimizer to wrap.
+        total_num_steps (int):
+            Total number of steps.
+        warmup_num_steps (int):
+            Number of steps for warmup.
+        **kwargs:
+            Other arguments.
+    """
+
+    def __init__(self, optimizer, total_num_steps=None, warmup_num_steps=0, **kwargs):
+        self.optimizer = optimizer
+        self.total_num_steps = total_num_steps
+        self.warmup_num_steps = warmup_num_steps
+        self.kwargs = kwargs
+
+
+class MegatronLMSchedulerWrapper(AcceleratedScheduler):
+    def __init__(self, scheduler, optimizers):
+        super().__init__(scheduler, optimizers)
+
+    def step(self, *args, **kwargs):
+        return  # `model(**batch)` is doing that automatically. Therefore, it's implementation is not needed
+
+
+def prepare_scheduler(accelerator, optimizer, scheduler):
+    accelerator.print("Preparing scheduler")
+    scheduler = get_optimizer_param_scheduler(optimizer)
+    return scheduler
+
+
+class AbstractTrainStep(ABC):
+    """Abstract class for batching, forward pass and loss handler."""
+
+    def __init__(self, name):
+        super().__init__()
+        self.name = name
+
+    def get_batch_func(self):
+        pass
+
+    def get_forward_step_func(self):
+        pass
+
+    def get_loss_func(self):
+        pass
+
+
+class BertTrainStep(AbstractTrainStep):
+    """
+    Bert train step class.
+
+    Args:
+        args (`argparse.Namespace`): Megatron-LM arguments.
+    """
+
+    def __init__(self, args):
+        super().__init__("BertTrainStep")
+        self.get_batch = self.get_batch_func(args.megatron_dataset_flag)
+        self.loss_func = self.get_loss_func(args.pretraining_flag, args.num_labels)
+        self.forward_step = self.get_forward_step_func(args.pretraining_flag, args.bert_binary_head)
+        if not args.model_return_dict:
+            self.model_output_class = None
+        else:
+            self.model_output_class = SequenceClassifierOutput
+
+    def get_batch_func(self, megatron_dataset_flag):
+        def get_batch_megatron(data_iterator):
+            """Build the batch."""
+
+            # Items and their type.
+            keys = ["text", "types", "labels", "is_random", "loss_mask", "padding_mask"]
+            datatype = torch.int64
+
+            # Broadcast data.
+            if data_iterator is not None:
+                data = next(data_iterator)
+            else:
+                data = None
+            data_b = mpu.broadcast_data(keys, data, datatype)
+
+            # Unpack.
+            tokens = data_b["text"].long()
+            types = data_b["types"].long()
+            sentence_order = data_b["is_random"].long()
+            loss_mask = data_b["loss_mask"].float()
+            lm_labels = data_b["labels"].long()
+            padding_mask = data_b["padding_mask"].long()
+
+            return tokens, types, sentence_order, loss_mask, lm_labels, padding_mask
+
+        def get_batch_transformer(data_iterator):
+            """Build the batch."""
+            data = next(data_iterator)
+            data = send_to_device(data, torch.cuda.current_device())
+
+            # Unpack.
+            tokens = data["input_ids"].long()
+            padding_mask = data["attention_mask"].long()
+            if "token_type_ids" in data:
+                types = data["token_type_ids"].long()
+            else:
+                types = None
+            if "labels" in data:
+                lm_labels = data["labels"].long()
+                loss_mask = (data["labels"] != -100).to(torch.float)
+            else:
+                lm_labels = None
+                loss_mask = None
+            if "next_sentence_label" in data:
+                sentence_order = data["next_sentence_label"].long()
+            else:
+                sentence_order = None
+
+            return tokens, types, sentence_order, loss_mask, lm_labels, padding_mask
+
+        if megatron_dataset_flag:
+            return get_batch_megatron
+        else:
+            return get_batch_transformer
+
+    def get_loss_func(self, pretraining_flag, num_labels):
+        def loss_func_pretrain(loss_mask, sentence_order, output_tensor):
+            lm_loss_, sop_logits = output_tensor
+
+            lm_loss_ = lm_loss_.float()
+            loss_mask = loss_mask.float()
+            lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum()
+
+            if sop_logits is not None:
+                sop_loss = F.cross_entropy(sop_logits.view(-1, 2).float(), sentence_order.view(-1), ignore_index=-1)
+                sop_loss = sop_loss.float()
+                loss = lm_loss + sop_loss
+                averaged_losses = average_losses_across_data_parallel_group([lm_loss, sop_loss])
+                return loss, {"lm loss": averaged_losses[0], "sop loss": averaged_losses[1]}
+
+            else:
+                loss = lm_loss
+                averaged_losses = average_losses_across_data_parallel_group([lm_loss])
+                return loss, {"lm loss": averaged_losses[0]}
+
+        def loss_func_finetune(labels, logits):
+            if num_labels == 1:
+                #  We are doing regression
+                loss_fct = MSELoss()
+                loss = loss_fct(logits.view(-1), labels.view(-1))
+            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, num_labels), labels.view(-1))
+            else:
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+            averaged_losses = average_losses_across_data_parallel_group([loss])
+            return loss, {"loss": averaged_losses[0]}
+
+        if pretraining_flag:
+            return loss_func_pretrain
+        else:
+            return loss_func_finetune
+
+    def get_forward_step_func(self, pretraining_flag, bert_binary_head):
+        def forward_step(data_iterator, model):
+            """Forward step."""
+            tokens, types, sentence_order, loss_mask, labels, padding_mask = self.get_batch(data_iterator)
+            if not bert_binary_head:
+                types = None
+            # Forward pass through the model.
+            if pretraining_flag:
+                output_tensor = model(tokens, padding_mask, tokentype_ids=types, lm_labels=labels)
+                return output_tensor, partial(self.loss_func, loss_mask, sentence_order)
+            else:
+                logits = model(tokens, padding_mask, tokentype_ids=types)
+                return logits, partial(self.loss_func, labels)
+
+        return forward_step
+
+
+class GPTTrainStep(AbstractTrainStep):
+    """
+    GPT train step class.
+
+    Args:
+        args (`argparse.Namespace`): Megatron-LM arguments.
+    """
+
+    def __init__(self, args):
+        super().__init__("GPTTrainStep")
+        self.get_batch = self.get_batch_func(args.megatron_dataset_flag)
+        self.loss_func = self.get_loss_func()
+        self.forward_step = self.get_forward_step_func()
+        self.eod_token = args.padded_vocab_size - 1
+        if not args.model_return_dict:
+            self.model_output_class = None
+        else:
+            self.model_output_class = CausalLMOutputWithCrossAttentions
+
+    def get_batch_func(self, megatron_dataset_flag):
+        def get_batch_megatron(data_iterator):
+            """Generate a batch"""
+            args = get_args()
+
+            # Items and their type.
+            keys = ["text"]
+            datatype = torch.int64
+
+            # Broadcast data.
+            if data_iterator is not None:
+                data = next(data_iterator)
+            else:
+                data = None
+            data_b = mpu.broadcast_data(keys, data, datatype)
+
+            # Unpack.
+            tokens_ = data_b["text"].long()
+            labels = tokens_[:, 1:].contiguous()
+            tokens = tokens_[:, :-1].contiguous()
+
+            # Get the masks and postition ids.
+            attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+                tokens, self.eod_token, args.reset_position_ids, args.reset_attention_mask, args.eod_mask_loss
+            )
+
+            return tokens, labels, loss_mask, attention_mask, position_ids
+
+        def get_batch_transformer(data_iterator):
+            data = next(data_iterator)
+            data = {"input_ids": data["input_ids"]}
+            data = send_to_device(data, torch.cuda.current_device())
+
+            tokens_ = data["input_ids"].long()
+            padding = torch.zeros((tokens_.shape[0], 1), dtype=tokens_.dtype, device=tokens_.device) + self.eod_token
+            tokens_ = torch.concat([tokens_, padding], dim=1)
+            labels = tokens_[:, 1:].contiguous()
+            tokens = tokens_[:, :-1].contiguous()
+            # Get the masks and postition ids.
+            attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+                tokens, self.eod_token, False, False, True
+            )
+            return tokens, labels, loss_mask, attention_mask, position_ids
+
+        if megatron_dataset_flag:
+            return get_batch_megatron
+        else:
+            return get_batch_transformer
+
+    def get_loss_func(self):
+        def loss_func(loss_mask, output_tensor):
+            losses = output_tensor.float()
+            loss_mask = loss_mask.view(-1).float()
+            loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
+
+            # Reduce loss for logging.
+            averaged_loss = average_losses_across_data_parallel_group([loss])
+
+            return loss, {"lm loss": averaged_loss[0]}
+
+        return loss_func
+
+    def get_forward_step_func(self):
+        def forward_step(data_iterator, model):
+            """Forward step."""
+            # Get the batch.
+            tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch(data_iterator)
+            output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
+
+            return output_tensor, partial(self.loss_func, loss_mask)
+
+        return forward_step
+
+
+class T5TrainStep(AbstractTrainStep):
+    """
+    T5 train step class.
+
+    Args:
+        args (`argparse.Namespace`): Megatron-LM arguments.
+    """
+
+    def __init__(self, args):
+        super().__init__("T5TrainStep")
+        self.get_batch = self.get_batch_func(args.megatron_dataset_flag)
+        self.loss_func = self.get_loss_func()
+        self.forward_step = self.get_forward_step_func()
+        if not args.model_return_dict:
+            self.model_output_class = None
+        else:
+            self.model_output_class = Seq2SeqLMOutput
+
+    @staticmethod
+    def attn_mask_postprocess(attention_mask):
+        # We create a 3D attention mask from a 2D tensor mask.
+        # [b, 1, s]
+        attention_mask_b1s = attention_mask.unsqueeze(1)
+        # [b, s, 1]
+        attention_mask_bs1 = attention_mask.unsqueeze(2)
+        # [b, s, s]
+        attention_mask_bss = attention_mask_b1s * attention_mask_bs1
+        # Convert attention mask to binary:
+        extended_attention_mask = attention_mask_bss < 0.5
+        return extended_attention_mask
+
+    @staticmethod
+    def get_decoder_mask(seq_length, device):
+        attention_mask = torch.tril(torch.ones((1, seq_length, seq_length), device=device))
+        attention_mask = attention_mask < 0.5
+        return attention_mask
+
+    @staticmethod
+    def get_enc_dec_mask(attention_mask, dec_seq_length, device):
+        batch_size, _ = attention_mask.shape
+        # We create a 3D attention mask from a 2D tensor mask.
+        # [b, 1, s]
+        attention_mask_b1s = attention_mask.unsqueeze(1)
+        # [b, s, 1]
+        attention_mask_bs1 = torch.ones((batch_size, dec_seq_length, 1), device=device)
+        attention_mask_bss = attention_mask_bs1 * attention_mask_b1s
+        extended_attention_mask = attention_mask_bss < 0.5
+        return extended_attention_mask
+
+    def get_batch_func(self, megatron_dataset_flag):
+        def get_batch_megatron(data_iterator):
+            """Build the batch."""
+
+            keys = ["text_enc", "text_dec", "labels", "loss_mask", "enc_mask", "dec_mask", "enc_dec_mask"]
+            datatype = torch.int64
+
+            # Broadcast data.
+            if data_iterator is not None:
+                data = next(data_iterator)
+            else:
+                data = None
+            data_b = mpu.broadcast_data(keys, data, datatype)
+
+            # Unpack.
+            tokens_enc = data_b["text_enc"].long()
+            tokens_dec = data_b["text_dec"].long()
+            labels = data_b["labels"].long()
+            loss_mask = data_b["loss_mask"].float()
+
+            enc_mask = data_b["enc_mask"] < 0.5
+            dec_mask = data_b["dec_mask"] < 0.5
+            enc_dec_mask = data_b["enc_dec_mask"] < 0.5
+
+            return tokens_enc, tokens_dec, loss_mask, labels, enc_mask, dec_mask, enc_dec_mask
+
+        def get_batch_transformer(data_iterator):
+            """Build the batch."""
+            data = next(data_iterator)
+            data = send_to_device(data, torch.cuda.current_device())
+
+            tokens_enc = data["input_ids"].long()
+            labels = data["labels"].long()
+            loss_mask = (labels != -100).to(torch.float)
+            if "decoder_input_ids" in data:
+                tokens_dec = data["decoder_input_ids"].long()
+            else:
+                tokens_dec = labels.new_zeros(labels.shape, device=labels.device, dtype=torch.long)
+                tokens_dec[..., 1:] = labels[..., :-1].clone()
+                tokens_dec[..., 0] = 0
+                tokens_dec.masked_fill_(tokens_dec == -100, 0)
+            enc_mask = T5TrainStep.attn_mask_postprocess(data["attention_mask"].long())
+            dec_mask = T5TrainStep.get_decoder_mask(tokens_dec.shape[1], tokens_dec.device)
+            enc_dec_mask = T5TrainStep.get_enc_dec_mask(
+                data["attention_mask"].long(), tokens_dec.shape[1], tokens_dec.device
+            )
+
+            return tokens_enc, tokens_dec, loss_mask, labels, enc_mask, dec_mask, enc_dec_mask
+
+        if megatron_dataset_flag:
+            return get_batch_megatron
+        else:
+            return get_batch_transformer
+
+    def get_loss_func(self):
+        def loss_func(loss_mask, output_tensor):
+            lm_loss_ = output_tensor.float()
+            lm_loss = torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1)) / loss_mask.sum()
+
+            loss = lm_loss
+            averaged_losses = average_losses_across_data_parallel_group([lm_loss])
+
+            return loss, {"lm loss": averaged_losses[0]}
+
+        return loss_func
+
+    def get_forward_step_func(self):
+        def forward_step(data_iterator, model):
+            """Forward step."""
+            # Get the batch.
+            tokens_enc, tokens_dec, loss_mask, lm_labels, enc_mask, dec_mask, enc_dec_mask = self.get_batch(
+                data_iterator
+            )
+            # Forward model lm_labels
+            output_tensor = model(
+                tokens_enc, tokens_dec, enc_mask, dec_mask, enc_dec_mask, tokentype_ids=None, lm_labels=lm_labels
+            )
+
+            return output_tensor, partial(self.loss_func, loss_mask)
+
+        return forward_step
+
+
+# intialize megatron setup
+def initialize(accelerator, extra_args_provider=None, args_defaults={}):
+    accelerator.print("Initializing Megatron-LM")
+    assert torch.cuda.is_available(), "Megatron requires CUDA."
+
+    # Parse arguments
+    args = parse_args(extra_args_provider, ignore_unknown_args=True)
+
+    # Set defaults
+    for key, value in args_defaults.items():
+        if getattr(args, key, None) is not None:
+            if args.rank == 0:
+                print(
+                    "WARNING: overriding default arguments for {key}:{v} \
+                        with {key}:{v2}".format(
+                        key=key, v=getattr(args, key), v2=value
+                    ),
+                    flush=True,
+                )
+        setattr(args, key, value)
+
+    if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False):
+        assert args.load is not None, "--use-checkpoints-args requires --load argument"
+        load_args_from_checkpoint(args)
+
+    validate_args(args)
+
+    # set global args, build tokenizer, and set adlr-autoresume,
+    # tensorboard-writer, and timers.
+    set_global_variables(args)
+
+    # torch.distributed initialization
+    def finish_mpu_init():
+        args = get_args()
+        # Pytorch distributed.
+        device_count = torch.cuda.device_count()
+        args.rank = torch.distributed.get_rank()
+        args.world_size = torch.distributed.get_world_size()
+        if device_count > 0:
+            device = args.rank % device_count
+            if args.local_rank is not None:
+                assert args.local_rank == device, "expected local-rank to be the same as rank % device-count."
+            else:
+                args.local_rank = device
+
+            # Set the tensor model-parallel, pipeline model-parallel, and
+            # data-parallel communicators.
+            if mpu.model_parallel_is_initialized():
+                print("model parallel is already initialized")
+            else:
+                mpu.initialize_model_parallel(
+                    args.tensor_model_parallel_size,
+                    args.pipeline_model_parallel_size,
+                    args.virtual_pipeline_model_parallel_size,
+                    args.pipeline_model_parallel_split_rank,
+                )
+
+        # Random seeds for reproducibility.
+        if args.rank == 0:
+            print("> setting random seeds to {} ...".format(args.seed))
+        _set_random_seed(args.seed, args.data_parallel_random_init)
+
+    args = get_args()
+
+    # Megatron's MPU is the master. Complete initialization right away.
+    finish_mpu_init()
+
+    # Autoresume.
+    _init_autoresume()
+
+    # Compile dependencies.
+    _compile_dependencies()
+
+    # Set pytorch JIT layer fusion options and warmup JIT functions.
+    set_jit_fusion_options()
+    args = get_args()
+    args.padded_vocab_size = _vocab_size_with_padding(args.orig_vocab_size, args)
+    if args.model_type_name == "bert" and args.pretraining_flag and args.num_labels == 2:
+        args.bert_binary_head = True
+    else:
+        args.bert_binary_head = False
+    args.iteration = 0
+
+
+class MegatronEngine(torch.nn.Module):
+    """
+    Megatron-LM model wrapper
+
+    Args:
+        accelerator (:class:`~accelerate.Accelerator`): The accelerator object to use.
+        model: Megatron-LM model
+        optimizer: Megatron-LM optimizer
+        lr_scheduler: Megatron-LM lr scheduler
+    """
+
+    def __init__(self, accelerator, model, optimizer, scheduler):
+        super(MegatronEngine, self).__init__()
+        self.module = model
+        self.base_model = model[0]
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        args = get_args()
+        if accelerator.state.megatron_lm_plugin.custom_train_step_class is not None:
+            self.train_step_handler = accelerator.state.megatron_lm_plugin.custom_train_step_class(
+                args, **accelerator.state.megatron_lm_plugin.custom_train_step_kwargs
+            )
+        elif args.model_type_name == "bert":
+            self.train_step_handler = BertTrainStep(args)
+        elif args.model_type_name == "gpt":
+            self.train_step_handler = GPTTrainStep(args)
+        elif args.model_type_name == "t5":
+            self.train_step_handler = T5TrainStep(args)
+        else:
+            raise ValueError(f"Unsupported model type: {args.model_type_name}")
+        self.optimizer.skipped_iter = False
+
+        # Tracking loss.
+        self.total_loss_dict = {}
+        self.eval_total_loss_dict = {}
+        self.iteration = 0
+        self.report_memory_flag = True
+        if args.tensorboard_dir is not None:
+            write_args_to_tensorboard()
+
+    def train(self):
+        for model_module in self.module:
+            model_module.train()
+        self.log_eval_results()
+
+    def eval(self):
+        for model_module in self.module:
+            model_module.eval()
+
+    def train_step(self, **batch_data):
+        """
+        Training step for Megatron-LM
+
+        Args:
+            batch_data (:obj:`dict`): The batch data to train on.
+        """
+
+        args = get_args()
+        timers = get_timers()
+
+        if len(batch_data) > 0:
+            data_chunks = []
+            if args.num_micro_batches > 1:
+                for i in range(0, args.num_micro_batches):
+                    data_chunks.append(
+                        {
+                            k: v[i * args.micro_batch_size : (i + 1) * args.micro_batch_size]
+                            for k, v in batch_data.items()
+                        }
+                    )
+            else:
+                data_chunks = [batch_data]
+
+        if len(self.module) > 1:
+            batch_data_iterator = (
+                [iter(data_chunks) for _ in range(len(self.module))]
+                if len(batch_data) > 0
+                else [None] * len(self.module)
+            )
+        else:
+            batch_data_iterator = iter(data_chunks) if len(batch_data) > 0 else None
+
+        # Set grad to zero.
+        if args.DDP_impl == "local" and args.use_contiguous_buffers_in_local_ddp:
+            for partition in self.module:
+                partition.zero_grad_buffer()
+        self.optimizer.zero_grad()
+
+        # Forward pass.
+        forward_backward_func = get_forward_backward_func()
+        losses_reduced = forward_backward_func(
+            self.train_step_handler.forward_step,
+            batch_data_iterator,
+            self.module,
+            self.optimizer,
+            None,
+            forward_only=False,
+        )
+
+        # Empty unused memory.
+        if args.empty_unused_memory_level >= 1:
+            torch.cuda.empty_cache()
+
+        # Reduce gradients.
+        timers("backward-reduce-model-grads").start()
+        self.optimizer.reduce_model_grads(args, timers)
+        timers("backward-reduce-model-grads").stop()
+
+        # Update parameters.
+        timers("optimizer").start()
+        update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step(args, timers)
+        timers("optimizer").stop()
+
+        # Gather params.
+        if update_successful:
+            timers("backward-gather-model-params").start()
+            self.optimizer.gather_model_params(args, timers)
+            timers("backward-gather-model-params").stop()
+
+        # Update learning rate.
+        if update_successful:
+            if self.scheduler is not None:
+                increment = get_num_microbatches() * args.micro_batch_size * args.data_parallel_size
+                self.scheduler.step(increment=increment)
+            skipped_iter = 0
+        else:
+            skipped_iter = 1
+
+        self.optimizer.skipped_iter = not update_successful
+
+        # Empty unused memory.
+        if args.empty_unused_memory_level >= 2:
+            torch.cuda.empty_cache()
+
+        args.consumed_train_samples += (
+            mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
+        )
+
+        if mpu.is_pipeline_last_stage(ignore_virtual=True):
+            # Average loss across microbatches.
+            loss_reduced = {}
+            for key in losses_reduced[0]:
+                losses_reduced_for_key = [x[key] for x in losses_reduced]
+                loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+            return loss_reduced, skipped_iter, grad_norm, num_zeros_in_grad
+        return {}, skipped_iter, grad_norm, num_zeros_in_grad
+
+    def eval_step(self, **batch_data):
+        """
+        Evaluation step for Megatron-LM
+
+        Args:
+            batch_data (:obj:`dict`): The batch data to evaluate on.
+        """
+
+        args = get_args()
+        data_chunks = []
+        if args.num_micro_batches > 1:
+            for i in range(0, args.num_micro_batches):
+                data_chunks.append(
+                    {k: v[i * args.micro_batch_size : (i + 1) * args.micro_batch_size] for k, v in batch_data.items()}
+                )
+        else:
+            data_chunks = [batch_data]
+
+        if len(self.module) > 1:
+            batch_data_iterator = [iter(data_chunks) for _ in range(len(self.module))]
+        else:
+            batch_data_iterator = iter(data_chunks)
+        forward_backward_func = get_forward_backward_func()
+        loss_dicts = forward_backward_func(
+            self.train_step_handler.forward_step,
+            batch_data_iterator,
+            self.module,
+            optimizer=None,
+            timers=None,
+            forward_only=True,
+        )
+        # Empty unused memory
+        if args.empty_unused_memory_level >= 1:
+            torch.cuda.empty_cache()
+
+        args.consumed_valid_samples += (
+            mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches()
+        )
+
+        if mpu.is_pipeline_last_stage(ignore_virtual=True):
+            # Average loss across microbatches.
+            loss_reduced = {}
+            for key in loss_dicts[0]:
+                losses_reduced_for_key = [x[key] for x in loss_dicts]
+                loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+            return loss_reduced
+        else:
+            return {}
+
+    def forward(self, **batch_data):
+        # During training, we use train_step()
+        # model(**batch_data) performs following operations by delegating it to `self.train_step`:
+        # 1. Prepare **batch_data for Tendor, Pipeline and Model Parallelism
+        # 2. Set grad to zero.
+        # 3. forward pass and backward pass using Pipeline Parallelism
+        # 4. Empty unused memory.
+        # 5. Reduce gradients.
+        # 6. Update parameters.
+        # 7. Gather params when using Distributed Optimizer (Data Parallelism).
+        # 8. Update learning rate if scheduler is specified.
+        # 9. Empty unused memory.
+        # 10. Average loss across microbatches and across DP ranks.
+        #
+        # During evaluation, we use eval_step()
+        args = get_args()
+        if self.module[0].training:
+            loss_dict, skipped_iter, grad_norm, num_zeros_in_grad = self.train_step(**batch_data)
+            self.iteration += 1
+            if args.tensorboard_dir is not None:
+                # Logging.
+                loss_scale = self.optimizer.get_loss_scale().item()
+                params_norm = None
+                if args.log_params_norm:
+                    params_norm = calc_params_l2_norm(self.model)
+                self.report_memory_flag = training_log(
+                    loss_dict,
+                    self.total_loss_dict,
+                    self.optimizer.param_groups[0]["lr"],
+                    self.iteration,
+                    loss_scale,
+                    self.report_memory_flag,
+                    skipped_iter,
+                    grad_norm,
+                    params_norm,
+                    num_zeros_in_grad,
+                )
+        else:
+            loss_dict = self.eval_step(**batch_data)
+            if args.tensorboard_dir is not None:
+                for key in loss_dict:
+                    self.eval_total_loss_dict[key] = (
+                        self.eval_total_loss_dict.get(key, torch.cuda.FloatTensor([0.0])) + loss_dict[key]
+                    )
+                    self.eval_total_loss_dict[key + "_num_iters"] = self.eval_total_loss_dict.get(
+                        key + "_num_iters", torch.cuda.FloatTensor([0.0])
+                    ) + torch.cuda.FloatTensor([1.0])
+
+        loss = torch.tensor(0.0, device=args.local_rank)
+        for key in loss_dict:
+            loss += loss_dict[key]
+        # loss = reduce(loss)
+        if self.train_step_handler.model_output_class is not None:
+            return self.train_step_handler.model_output_class(loss=loss)
+        return loss
+
+    def log_eval_results(self):
+        args = get_args()
+        if args.tensorboard_dir is None or self.iteration == 0:
+            return
+        args = get_args()
+        writer = get_tensorboard_writer()
+        string = f"validation loss at iteration {self.iteration} | "
+        for key in self.eval_total_loss_dict:
+            if key.endswith("_num_iters"):
+                continue
+            value = self.eval_total_loss_dict[key] / self.eval_total_loss_dict[key + "_num_iters"]
+            string += f"{key} value: {value} | "
+            ppl = math.exp(min(20, value.item()))
+            if args.pretraining_flag:
+                string += f"{key} PPL: {ppl} | "
+            if writer:
+                writer.add_scalar(f"{key} validation", value.item(), self.iteration)
+                if args.pretraining_flag:
+                    writer.add_scalar(f"{key} validation ppl", ppl, self.iteration)
+
+        length = len(string) + 1
+        print_rank_last("-" * length)
+        print_rank_last(string)
+        print_rank_last("-" * length)
+        self.eval_total_loss_dict = {}
+
+    def save_checkpoint(self, output_dir):
+        self.log_eval_results()
+        args = get_args()
+        args.save = output_dir
+        torch.distributed.barrier()
+        save_checkpoint(self.iteration, self.module, self.optimizer, self.scheduler)
+        torch.distributed.barrier()
+
+    def load_checkpoint(self, input_dir):
+        args = get_args()
+        args.load = input_dir
+        args.consumed_train_samples = 0
+        args.consumed_valid_samples = 0
+        torch.distributed.barrier()
+        iteration = load_checkpoint(self.module, self.optimizer, self.scheduler)
+        torch.distributed.barrier()
+        self.iteration = iteration
+        if args.fp16 and self.iteration == 0:
+            self.optimizer.reload_model_params()
+
+
+# other utilities
+def avg_losses_across_data_parallel_group(losses):
+    """
+    Average losses across data parallel group.
+
+    Args:
+        losses (List[Tensor]): List of losses to average across data parallel group.
+    """
+
+    return average_losses_across_data_parallel_group(losses)
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
index 3ba46a03f..aef1fd589 100644
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -224,6 +224,7 @@ def gather(tensor):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
     ]:
         return _gpu_gather(tensor)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
@@ -261,6 +262,7 @@ def gather_object(object: Any):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
     ]:
         return _gpu_gather_object(object)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
@@ -304,6 +306,7 @@ def broadcast(tensor, from_process: int = 0):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
     ]:
         return _gpu_broadcast(tensor, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
@@ -332,6 +335,7 @@ def broadcast_object_list(object_list, from_process: int = 0):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
     ]:
         torch.distributed.broadcast_object_list(object_list, src=from_process)
     elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
@@ -452,6 +456,7 @@ def _reduce_across_processes(tensor, reduction="mean"):
             DistributedType.DEEPSPEED,
             DistributedType.MULTI_GPU,
             DistributedType.FSDP,
+            DistributedType.MEGATRON_LM,
         ]:
             torch.distributed.all_reduce(cloned_tensor, ReduceOp.SUM)
             return cloned_tensor
diff --git a/src/accelerate/utils/random.py b/src/accelerate/utils/random.py
index e95ed03bf..8dc149a02 100644
--- a/src/accelerate/utils/random.py
+++ b/src/accelerate/utils/random.py
@@ -64,7 +64,12 @@ def synchronize_rng_state(rng_type: Optional[RNGType] = None, generator: Optiona
     state = AcceleratorState()
     if state.distributed_type == DistributedType.TPU:
         rng_state = xm.mesh_reduce("random_seed", rng_state, lambda x: x[0])
-    elif state.distributed_type in [DistributedType.DEEPSPEED, DistributedType.MULTI_GPU, DistributedType.FSDP]:
+    elif state.distributed_type in [
+        DistributedType.DEEPSPEED,
+        DistributedType.MULTI_GPU,
+        DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
+    ]:
         rng_state = rng_state.to(state.device)
         torch.distributed.broadcast(rng_state, 0)
         rng_state = rng_state.cpu()
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 9197070c8..05688f2b5 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -39,6 +39,7 @@
     "automatic_gradient_accumulation.py",
     "fsdp_with_peak_mem_tracking.py",
     "deepspeed_with_config_support.py",
+    "megatron_lm_gpt_pretraining.py",
 ]
 
 


Review comments: body: ```suggestion
        self.megatron_lm_default_args = {
```
Also you use `megatronlm_plugin` with no `_` between `megatron` and `lm`., just flagging this so you stay consistent (I prefer `megatron_lm` personally.
diff_hunk: @@ -643,3 +645,142 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        self.megtron_lm_default_args = {
body: Done.
diff_hunk: @@ -643,3 +645,142 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        self.megtron_lm_default_args = {
body: I don't believe all of those should be changed just for MegatronLM.
diff_hunk: @@ -1032,7 +1174,7 @@ def wait_for_everyone(self):
         """
         wait_for_everyone()
 
-    @on_main_process
+    @on_last_process
body: Here the message should be more in the lines of "The Megatron-LM only supports those"
diff_hunk: @@ -643,3 +645,291 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+    train_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of iterations to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    train_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of samples to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    weight_decay_incr_style: str = field(
+        default="constant",
+        metadata={
+            "help": 'Weight decay increment function. choices=["constant", "linear", "cosine"]. '
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    start_weight_decay: float = field(
+        default=None,
+        metadata={
+            "help": "Initial weight decay coefficient for L2 regularization. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    end_weight_decay: float = field(
+        default=None,
+        metadata={
+            "help": "End of run weight decay coefficient for L2 regularization. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_style: str = field(
+        default="linear",
+        metadata={
+            "help": "Learning rate decay function. choices=['constant', 'linear', 'cosine']. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Number of iterations for learning rate decay. If None defaults to `train_iters`. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Number of samples for learning rate decay. If None defaults to `train_samples`. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_iters: int = field(
+        default=None,
+        metadata={
+            "help": "number of iterations to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_samples: int = field(
+        default=None,
+        metadata={
+            "help": "number of samples to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_fraction: float = field(
+        default=None,
+        metadata={
+            "help": "fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    min_lr: float = field(
+        default=None,
+        metadata={
+            "help": "Minumum value for learning rate. The scheduler clip values below this threshold. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    consumed_samples: List[int] = field(
+        default=None,
+        metadata={
+            "help": "Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call."
+        },
+    )
+    no_wd_decay_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to disable weight decay."})
+    scale_lr_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to scale learning rate."})
+    lr_mult: float = field(default=1.0, metadata={"help": "Learning rate multiplier."})
+    megatron_dataset_flag: bool = field(
+        default=False,
+        metadata={"help": "Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format."},
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        if self.consumed_samples is not None:
+            if len(self.consumed_samples) == 1:
+                self.consumed_samples.extend([0, 0])
+            elif len(self.consumed_samples) == 2:
+                self.consumed_samples.append(0)
+
+        self.megatron_lm_default_args = {
+            "tensor_model_parallel_size": self.tp_degree,
+            "pipeline_model_parallel_size": self.pp_degree,
+            "pipeline_model_parallel_split_rank": self.pipeline_model_parallel_split_rank,
+            "num_layers_per_virtual_pipeline_stage": self.num_layers_per_virtual_pipeline_stage,
+            "DDP_impl": self.DDP_impl,
+            "use_distributed_optimizer": self.use_distributed_optimizer,
+            "recompute_activations": self.recompute_activation,
+            "sequence_parallel": self.sequence_parallelism,
+            "clip_grad": self.gradient_clipping,
+            "num_micro_batches": self.num_micro_batches,
+            "consumed_samples": self.consumed_samples,
+            "no_wd_decay_cond": self.no_wd_decay_cond,
+            "scale_lr_cond": self.scale_lr_cond,
+            "lr_mult": self.lr_mult,
+            "megatron_dataset_flag": self.megatron_dataset_flag,
+        }
+
+    def set_network_size_args(self, model):
+        # Check if the model is either BERT, GPT or T5 else raise error
+        # set 'num_layers', 'hidden_size', 'num_attention_heads', 'max_position_embeddings'
+        if "bert" in model.__class__.__name__.lower():
+            model_type_name = "bert"
+            num_layers = model.config.num_hidden_layers
+            hidden_size = model.config.hidden_size
+            num_attention_heads = model.config.num_attention_heads
+            max_position_embeddings = model.config.max_position_embeddings
+            num_labels = model.config.num_labels
+            orig_vocab_size = model.config.vocab_size
+            if "maskedlm" in model.__class__.__name__.lower():
+                pretraining_flag = True
+
+        elif "gpt" in model.__class__.__name__.lower():
+            model_type_name = "gpt"
+            num_layers = model.config.n_layer
+            hidden_size = model.config.n_embd
+            num_attention_heads = model.config.n_head
+            max_position_embeddings = model.config.n_positions
+            orig_vocab_size = model.config.vocab_size
+            pretraining_flag = True
+        elif "t5" in model.__class__.__name__.lower():
+            model_type_name = "t5"
+            num_layers = model.config.num_layers
+            hidden_size = model.config.d_model
+            num_attention_heads = model.config.num_heads
+            max_position_embeddings = model.config.n_positions
+            orig_vocab_size = model.config.vocab_size
+            pretraining_flag = True
+        else:
+            raise ValueError("Model is not BERT, GPT or T5. Please check the model you are using.")
body: You should check the `model_type` attribute instead of the class name IMO.
diff_hunk: @@ -643,3 +645,291 @@ def load_optimizer(self, accelerator, optimizer, model, input_dir, optimizer_ind
         # called from all ranks, though only rank0 has a valid param for full_osd
         sharded_osd = FSDP.scatter_full_optim_state_dict(full_osd, model)
         optimizer.load_state_dict(sharded_osd)
+
+
+@dataclass
+class MegatronLMPlugin:
+    """
+    Plugin for Megatron-LM to enable tensor, pipeline, sequence and data parallelism. Also to enable selective
+    activation recomputation and optimized fused kernels.
+    """
+
+    tp_degree: int = field(default=None, metadata={"help": "tensor parallelism degree."})
+    pp_degree: int = field(default=None, metadata={"help": "pipeline parallelism degree."})
+    num_micro_batches: int = field(default=None, metadata={"help": "number of micro-batches."})
+    gradient_clipping: float = field(
+        default=None, metadata={"help": "gradient clipping value based on global L2 Norm (0 to disable)"}
+    )
+    sequence_parallelism: bool = field(
+        default=None,
+        metadata={"help": "enable sequence parallelism"},
+    )
+    recompute_activation: bool = field(
+        default=None,
+        metadata={"help": "enable selective activation recomputation"},
+    )
+    use_distributed_optimizer: bool = field(
+        default=None,
+        metadata={"help": "enable distributed optimizer"},
+    )
+    pipeline_model_parallel_split_rank: int = field(
+        default=None, metadata={"help": "Rank where encoder and decoder should be split."}
+    )
+    num_layers_per_virtual_pipeline_stage: int = field(
+        default=None, metadata={"help": "Number of layers per virtual pipeline stage."}
+    )
+    is_train_batch_min: str = field(
+        default=True,
+        metadata={"help": "If both train & eval dataloaders are specified, this will decide the micro_batch_size"},
+    )
+    train_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of iterations to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    train_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Total number of samples to train over all training runs. "
+            "Note that either train-iters or train-samples should be provided when using `MegatronLMDummyScheduler`"
+        },
+    )
+    weight_decay_incr_style: str = field(
+        default="constant",
+        metadata={
+            "help": 'Weight decay increment function. choices=["constant", "linear", "cosine"]. '
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    start_weight_decay: float = field(
+        default=None,
+        metadata={
+            "help": "Initial weight decay coefficient for L2 regularization. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    end_weight_decay: float = field(
+        default=None,
+        metadata={
+            "help": "End of run weight decay coefficient for L2 regularization. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_style: str = field(
+        default="linear",
+        metadata={
+            "help": "Learning rate decay function. choices=['constant', 'linear', 'cosine']. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_iters: int = field(
+        default=None,
+        metadata={
+            "help": "Number of iterations for learning rate decay. If None defaults to `train_iters`. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_decay_samples: int = field(
+        default=None,
+        metadata={
+            "help": "Number of samples for learning rate decay. If None defaults to `train_samples`. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_iters: int = field(
+        default=None,
+        metadata={
+            "help": "number of iterations to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_samples: int = field(
+        default=None,
+        metadata={
+            "help": "number of samples to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    lr_warmup_fraction: float = field(
+        default=None,
+        metadata={
+            "help": "fraction of lr-warmup-(iters/samples) to linearly warmup learning rate over. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    min_lr: float = field(
+        default=None,
+        metadata={
+            "help": "Minumum value for learning rate. The scheduler clip values below this threshold. "
+            "Only applicable when using `MegatronLMDummyScheduler`."
+        },
+    )
+    consumed_samples: List[int] = field(
+        default=None,
+        metadata={
+            "help": "Number of samples consumed in the same order as the dataloaders to `accelerator.prepare` call."
+        },
+    )
+    no_wd_decay_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to disable weight decay."})
+    scale_lr_cond: Optional[Callable] = field(default=None, metadata={"help": "Condition to scale learning rate."})
+    lr_mult: float = field(default=1.0, metadata={"help": "Learning rate multiplier."})
+    megatron_dataset_flag: bool = field(
+        default=False,
+        metadata={"help": "Whether the format of dataset follows Megatron-LM Indexed/Cached/MemoryMapped format."},
+    )
+
+    def __post_init__(self):
+        prefix = "MEGATRON_"
+        if self.tp_degree is None:
+            self.tp_degree = int(os.environ.get(prefix + "TP_DEGREE", 1))
+        if self.pp_degree is None:
+            self.pp_degree = int(os.environ.get(prefix + "PP_DEGREE", 1))
+        if self.num_micro_batches is None:
+            self.num_micro_batches = int(os.environ.get(prefix + "NUM_MICRO_BATCHES", 1))
+        if self.gradient_clipping is None:
+            self.gradient_clipping = float(os.environ.get(prefix + "GRADIENT_CLIPPING", 1.0))
+        if self.recompute_activation is None:
+            self.recompute_activation = strtobool(os.environ.get(prefix + "RECOMPUTE_ACTIVATION", "False")) == 1
+        if self.use_distributed_optimizer is None:
+            self.use_distributed_optimizer = (
+                strtobool(os.environ.get(prefix + "USE_DISTRIBUTED_OPTIMIZER", "False")) == 1
+            )
+        if self.sequence_parallelism is None:
+            self.sequence_parallelism = strtobool(os.environ.get(prefix + "SEQUENCE_PARALLELISM", "False")) == 1
+
+        if self.pp_degree > 1:
+            self.DDP_impl = "local"
+        else:
+            self.DDP_impl = "torch"
+
+        if self.consumed_samples is not None:
+            if len(self.consumed_samples) == 1:
+                self.consumed_samples.extend([0, 0])
+            elif len(self.consumed_samples) == 2:
+                self.consumed_samples.append(0)
+
+        self.megatron_lm_default_args = {
+            "tensor_model_parallel_size": self.tp_degree,
+            "pipeline_model_parallel_size": self.pp_degree,
+            "pipeline_model_parallel_split_rank": self.pipeline_model_parallel_split_rank,
+            "num_layers_per_virtual_pipeline_stage": self.num_layers_per_virtual_pipeline_stage,
+            "DDP_impl": self.DDP_impl,
+            "use_distributed_optimizer": self.use_distributed_optimizer,
+            "recompute_activations": self.recompute_activation,
+            "sequence_parallel": self.sequence_parallelism,
+            "clip_grad": self.gradient_clipping,
+            "num_micro_batches": self.num_micro_batches,
+            "consumed_samples": self.consumed_samples,
+            "no_wd_decay_cond": self.no_wd_decay_cond,
+            "scale_lr_cond": self.scale_lr_cond,
+            "lr_mult": self.lr_mult,
+            "megatron_dataset_flag": self.megatron_dataset_flag,
+        }
+
+    def set_network_size_args(self, model):
+        # Check if the model is either BERT, GPT or T5 else raise error
+        # set 'num_layers', 'hidden_size', 'num_attention_heads', 'max_position_embeddings'
+        if "bert" in model.__class__.__name__.lower():
body: To be cleaned up?
diff_hunk: @@ -0,0 +1,821 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from abc import ABC
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ..optimizer import AcceleratedOptimizer
+from ..scheduler import AcceleratedScheduler
+from .imports import is_megatron_lm_available, is_transformers_available
+from .operations import recursively_apply
+
+
+if is_transformers_available():
+    from transformers.modeling_outputs import (
+        CausalLMOutputWithCrossAttentions,
+        Seq2SeqLMOutput,
+        SequenceClassifierOutput,
+    )
+    from transformers.models.bert.modeling_bert import BertForPreTrainingOutput
+
+
+if is_megatron_lm_available():
+    from megatron import get_args, get_num_microbatches, get_timers, mpu
+    from megatron.arguments import _add_data_args
+    from megatron.data.data_samplers import build_pretraining_data_loader
+    from megatron.data.dataset_utils import build_train_valid_test_datasets
+    from megatron.initialize import (
+        get_tensor_model_parallel_group,
+        get_tensor_model_parallel_src_rank,
+        initialize_megatron,
+        set_jit_fusion_options,
+    )
+    from megatron.model import BertModel, GPTModel, ModelType, T5Model
+    from megatron.model.classification import Classification
+    from megatron.model.module import MegatronModule
+    from megatron.optimizer import get_megatron_optimizer
+    from megatron.schedules import get_forward_backward_func
+    from megatron.tokenizer.tokenizer import _vocab_size_with_padding
+    from megatron.training import build_train_valid_test_data_iterators, get_model, get_optimizer_param_scheduler
+    from megatron.utils import average_losses_across_data_parallel_group, get_ltor_masks_and_position_ids
+
+
+def model_provider_func(accelerator, pre_process=True, post_process=True, add_encoder=True, add_decoder=True):
+    """Build the model."""
+    args = get_args()
+    mode = "pre-training" if args.pretraining_flag else "fine-tuning"
+    accelerator.print(f"Building {args.model_type_name} model in the {mode} mode.")
+    if args.model_type_name == "bert":
+        if args.pretraining_flag:
+            num_tokentypes = 2 if args.bert_binary_head else 0
+            model = BertModel(
+                num_tokentypes=num_tokentypes,
+                add_binary_head=args.bert_binary_head,
+                parallel_output=True,
+                pre_process=pre_process,
+                post_process=post_process,
+            )
+        else:
+
+            model = Classification(
+                num_classes=args.num_labels, num_tokentypes=2, pre_process=pre_process, post_process=post_process
+            )
+    elif args.model_type_name == "gpt2":
+        model = GPTModel(num_tokentypes=0, parallel_output=True, pre_process=pre_process, post_process=post_process)
+    elif args.model_type_name == "t5":
+        model = T5Model(
+            num_tokentypes=0,
+            parallel_output=True,
+            pre_process=pre_process,
+            post_process=post_process,
+            add_encoder=add_encoder,
+            add_decoder=add_decoder,
+        )
+    else:
+        raise ValueError(f"Unknown model type: {args.model_type_name}")
+    return model
+
+
+def prepare_data_loader(accelerator, dataloader, consumed_samples_index=-1, consumed_samples=0):
+    accelerator.print("Preparing dataloader")
+    args = get_args()
+    if not args.megatron_dataset_flag:
+        if args.consumed_samples is not None:
+            consumed_samples = args.consumed_samples[consumed_samples_index]
+        return build_pretraining_data_loader(dataloader.dataset, consumed_samples)
+    else:
+        if args.consumed_samples is not None:
+            (
+                args.consumed_train_samples,
+                args.consumed_valid_samples,
+                args.consumed_test_samples,
+            ) = args.consumed_samples
+        else:
+            args.consumed_train_samples, args.consumed_valid_samples, args.consumed_test_samples = 0, 0, 0
+        train_data_iterator, valid_data_iterator, test_data_iterator = build_train_valid_test_data_iterators(
+            MegatronLMDummyDataLoader.get_train_valid_test_datasets_provider
+        )
+        return train_data_iterator, valid_data_iterator, test_data_iterator
+
+
+def prepare_model(accelerator):
+    accelerator.print("Preparing model")
+    args = get_args()
+    if args.model_type_name == "bert" or args.model_type_name == "gpt":
+        model_type = ModelType.encoder_or_decoder
+    elif args.model_type_name == "t5":
+        model_type = ModelType.encoder_and_decoder
+    model = get_model(model_provider_func, model_type)
+    return model
+
+
+def prepare_optimizer(accelerator, model):
+    accelerator.print("Preparing optimizer")
+    args = get_args()
+    optimizer = get_megatron_optimizer(model, args.no_wd_decay_cond, args.scale_lr_cond, args.lr_mult)
+    return optimizer
+
+
+def prepare_scheduler(accelerator, optimizer, scheduler, is_dummy_scheduler):
+    accelerator.print("Preparing scheduler")
+    if is_dummy_scheduler:
+        scheduler = get_optimizer_param_scheduler(optimizer)
+    else:
+        scheduler.optimizer = optimizer
+        if isinstance(scheduler, torch.optim.lr_scheduler.LambdaLR):
+            scheduler = scheduler.__class__(optimizer, scheduler.lr_lambdas[0])
+    return scheduler
+
+
+def initialize(accelerator, extra_args_provider=None, args_defaults={}):
+    accelerator.print("Initializing Megatron-LM")
+    # Initalize and get arguments
+    initialize_megatron(extra_args_provider=extra_args_provider, args_defaults=args_defaults, ignore_unknown_args=True)
+    # Set pytorch JIT layer fusion options and warmup JIT functions.
+    set_jit_fusion_options()
+    args = get_args()
+    args.padded_vocab_size = _vocab_size_with_padding(args.orig_vocab_size, args)
+    if args.model_type_name == "bert" and args.pretraining_flag and args.num_labels == 2:
+        args.bert_binary_head = True
+    # if args.virtual_pipeline_model_parallel_size is not None:
+    #     raise Exception("Virtual pipeline model parallelism is not supported in Accelerate")
body: The aliases should be done in the `.utils` module: if we have another integration like this, we won't be able to have both names.
diff_hunk: @@ -71,6 +73,21 @@
         DummyScheduler,
     )
 
+if is_megatron_lm_available():
+    from .utils import (
+        MegatronEngine,
+        MegatronLMDummyDataLoader,
+        MegatronLMDummyScheduler,
+        MegatronLMOptimizerWrapper,
+        MegatronLMSchedulerWrapper,
+    )
+    from .utils import initialize as megatron_lm_initialize
+    from .utils import prepare_data_loader as megatron_lm_prepare_data_loader
+    from .utils import prepare_model as megatron_lm_prepare_model
+    from .utils import prepare_optimizer as megatron_lm_prepare_optimizer
+    from .utils import prepare_scheduler as megatron_lm_prepare_scheduler
body: AS said above, those should probably have the `megatron_lm` prefix here.
diff_hunk: @@ -94,6 +96,23 @@
     )
 
 from .launch import PrepareForLaunch, _filter_args, get_launch_prefix
+from .megatron_lm import (
+    AbstractTrainStep,
+    BertTrainStep,
+    GPTTrainStep,
+    MegatronEngine,
+    MegatronLMDummyDataLoader,
+    MegatronLMDummyScheduler,
+    MegatronLMOptimizerWrapper,
+    MegatronLMSchedulerWrapper,
+    T5TrainStep,
+    avg_losses_across_data_parallel_group,
+    initialize,
+    prepare_data_loader,
+    prepare_model,
+    prepare_optimizer,
+    prepare_scheduler,
body: Nit: in a followup PR it would be nice to have all of those in one constant.
diff_hunk: @@ -224,6 +224,7 @@ def gather(tensor):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
body: Done.
diff_hunk: @@ -71,6 +73,21 @@
         DummyScheduler,
     )
 
+if is_megatron_lm_available():
+    from .utils import (
+        MegatronEngine,
+        MegatronLMDummyDataLoader,
+        MegatronLMDummyScheduler,
+        MegatronLMOptimizerWrapper,
+        MegatronLMSchedulerWrapper,
+    )
+    from .utils import initialize as megatron_lm_initialize
+    from .utils import prepare_data_loader as megatron_lm_prepare_data_loader
+    from .utils import prepare_model as megatron_lm_prepare_model
+    from .utils import prepare_optimizer as megatron_lm_prepare_optimizer
+    from .utils import prepare_scheduler as megatron_lm_prepare_scheduler
body: Done.
diff_hunk: @@ -94,6 +96,23 @@
     )
 
 from .launch import PrepareForLaunch, _filter_args, get_launch_prefix
+from .megatron_lm import (
+    AbstractTrainStep,
+    BertTrainStep,
+    GPTTrainStep,
+    MegatronEngine,
+    MegatronLMDummyDataLoader,
+    MegatronLMDummyScheduler,
+    MegatronLMOptimizerWrapper,
+    MegatronLMSchedulerWrapper,
+    T5TrainStep,
+    avg_losses_across_data_parallel_group,
+    initialize,
+    prepare_data_loader,
+    prepare_model,
+    prepare_optimizer,
+    prepare_scheduler,
body: Can be for another PR, but would be nice to explain the difference between those kinds of parallelism and what is parallelized exactly.
diff_hunk: @@ -0,0 +1,507 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+
+# Megatron-LM
+
+[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) enables training large transformer language models at scale.
+It provides efficient tensor, pipeline and sequence based model parallelism for pre-training transformer based
+Language Models such as [GPT](https://arxiv.org/abs/2005.14165) (Decoder Only), [BERT](https://arxiv.org/pdf/1810.04805.pdf) (Encoder Only) and [T5](https://arxiv.org/abs/1910.10683) (Encoder-Decoder).
+For detailed information and how things work behind the scene please refer the github [repo](https://github.com/NVIDIA/Megatron-LM).
+
+## What is integrated?
+
+Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
+of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
+
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+
+b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
+Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
+Layers are distributed uniformly across PP stages.
+
+c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
+
+d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.
body: PR #758 to resolve this
diff_hunk: @@ -224,6 +224,7 @@ def gather(tensor):
         DistributedType.DEEPSPEED,
         DistributedType.MULTI_GPU,
         DistributedType.FSDP,
+        DistributedType.MEGATRON_LM,
body: PR #759 to address this.
diff_hunk: @@ -0,0 +1,507 @@
+<!--Copyright 2022 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+
+# Megatron-LM
+
+[Megatron-LM](https://github.com/NVIDIA/Megatron-LM) enables training large transformer language models at scale.
+It provides efficient tensor, pipeline and sequence based model parallelism for pre-training transformer based
+Language Models such as [GPT](https://arxiv.org/abs/2005.14165) (Decoder Only), [BERT](https://arxiv.org/pdf/1810.04805.pdf) (Encoder Only) and [T5](https://arxiv.org/abs/1910.10683) (Encoder-Decoder).
+For detailed information and how things work behind the scene please refer the github [repo](https://github.com/NVIDIA/Megatron-LM).
+
+## What is integrated?
+
+Accelerate integrates following feature of Megatron-LM to enable large scale pre-training/finetuning
+of BERT (Encoder), GPT (Decoder) or T5 models (Encoder and Decoder):
+
+a. **Tensor Parallelism (TP)**: Reduces memory footprint without much additional communication on intra-node ranks
+
+b. **Pipeline Parallelism (PP)**: Reduces memory footprint and enables large scale training via inter-node parallelization. 
+Reduces the bubble of naive PP via PipeDream-Flush schedule/1F1B schedule and Interleaved 1F1B schedule. 
+Layers are distributed uniformly across PP stages.
+
+c. **Sequence Parallelism (SP)**: Reduces memory footprint without any additional communication. Only applicable when using TP.
+
+d. **Data Parallelism (DP)** via Distributed Optimizer: Reduces memory footprint by sharding optimizer states and gradients across DP ranks.

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._