diff --git a/docs/source/usage_guides/megatron_lm.mdx b/docs/source/usage_guides/megatron_lm.mdx
index d26ecd4b5..23b024f75 100644
--- a/docs/source/usage_guides/megatron_lm.mdx
+++ b/docs/source/usage_guides/megatron_lm.mdx
@@ -435,6 +435,103 @@ python checkpoint_utils/megatgron_gpt2/checkpoint_reshaping_and_interoperability
 --print-checkpoint-structure
 ```
 
+## Megatron-LM GPT models support returning logits and `megatron_generate` function for text generation
+
+1. Returning logits require setting `require_logits=True` in MegatronLMPlugin as shown below. 
+These would be available on the in the last stage of pipeline.
+```python
+megatron_lm_plugin = MegatronLMPlugin(return_logits=True)
+```
+
+2. `megatron_generate` method for Megatron-LM GPT model: This will use Tensor and Pipeline Parallelism to complete 
+generations for a batch of inputs when using greedy with/without top_k/top_p sampling and for individual prompt inputs when using beam search decoding. 
+Only a subset of features of transformers generate is supported. This will help in using large models via tensor and pipeline parallelism 
+for generation (already does key-value caching and uses fused kernels by default).
+This requires data parallel size to be 1, sequence parallelism and activation checkpointing to be disabled.
+It also requires specifying path to tokenizer's vocab file and merges file. 
+Below example shows how to configure and use `megatron_generate` method for Megatron-LM GPT model.
+```python
+# specifying tokenizer's vocab and merges file
+vocab_file = os.path.join(args.resume_from_checkpoint, "vocab.json")
+merge_file = os.path.join(args.resume_from_checkpoint, "merges.txt")
+other_megatron_args = {"vocab_file": vocab_file, "merge_file": merge_file}
+megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
+
+# inference using `megatron_generate` functionality
+tokenizer.pad_token = tokenizer.eos_token
+max_new_tokens = 64
+batch_texts = [
+    "Are you human?",
+    "The purpose of life is",
+    "The arsenal was constructed at the request of",
+    "How are you doing these days?",
+]
+batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)
+
+# top-p sampling
+generated_tokens = model.megatron_generate(
+    batch_encodings["input_ids"],
+    batch_encodings["attention_mask"],
+    max_new_tokens=max_new_tokens,
+    top_p=0.8,
+    top_p_decay=0.5,
+    temperature=0.9,
+)
+decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
+accelerator.print(decoded_preds)
+
+# top-k sampling
+generated_tokens = model.megatron_generate(
+    batch_encodings["input_ids"],
+    batch_encodings["attention_mask"],
+    max_new_tokens=max_new_tokens,
+    top_k=50,
+    temperature=0.9,
+)
+decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
+accelerator.print(decoded_preds)
+
+# adding `bos` token at the start
+generated_tokens = model.megatron_generate(
+    batch_encodings["input_ids"], batch_encodings["attention_mask"], max_new_tokens=max_new_tokens, add_BOS=True
+)
+decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
+accelerator.print(decoded_preds)
+
+# beam search => only takes single prompt
+batch_texts = ["The purpose of life is"]
+batch_encodings = tokenizer(batch_texts, return_tensors="pt", padding=True)
+generated_tokens = model.megatron_generate(
+    batch_encodings["input_ids"],
+    batch_encodings["attention_mask"],
+    max_new_tokens=max_new_tokens,
+    num_beams=20,
+    length_penalty=1.5,
+)
+decoded_preds = tokenizer.batch_decode(generated_tokens.cpu().numpy())
+accelerator.print(decoded_preds)
+```
+
+3. An end-to-end example of using `megatron_generate` method for Megatron-LM GPT model is available at
+[megatron_gpt2_generation.py](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/inference/megatron_gpt2_generation.py) with 
+config file [megatron_lm_gpt_generate_config.yaml](https://github.com/pacman100/accelerate-megatron-test/blob/main/src/Configs/megatron_lm_gpt_generate_config.yaml).
+The bash script with accelerate launch command is available at [megatron_lm_gpt_generate.sh](https://github.com/pacman100/accelerate-megatron-test/blob/main/megatron_lm_gpt_generate.sh).
+The output logs of the script are available at [megatron_lm_gpt_generate.log](https://github.com/pacman100/accelerate-megatron-test/blob/main/output_logs/megatron_lm_gpt_generate.log).
+
+## Support for ROPE and ALiBi Positional embeddings and Multi-Query Attention
+
+1. For ROPE/ALiBi attention, pass `position_embedding_type` with `("absolute" | "rotary" | "alibi")` to `MegatronLMPlugin` as shown below.
+```python
+other_megatron_args = {"position_embedding_type": "alibi"}
+megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
+```
+
+2. For Multi-Query Attention, pass `attention_head_type` with `("multihead" | "multiquery")` to `MegatronLMPlugin` as shown below.
+```python
+other_megatron_args = {"attention_head_type": "multiquery"}
+megatron_lm_plugin = MegatronLMPlugin(other_megatron_args=other_megatron_args)
+```
+
 ## Caveats
 
 1. Supports Transformers GPT2, Megatron-BERT and T5 models.
@@ -445,8 +542,12 @@ there is quite complex interplay of pipeline, tensor and data parallelsim behind
 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. 
+For GPT model, returning logits in addition to loss(es) is supported. 
+These logits aren't gathered across data prallel ranks. Use `accelerator.utils.gather_across_data_parallel_groups`
+to gather logits across data parallel ranks. These logits along with labels can be used for computing various 
+performance metrics. 
 
-3. The main process is the last rank as the losses are available in the last stage of pipeline.
+3. The main process is the last rank as the losses/logits 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.
 
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index b472ec556..0e67daf4b 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -107,6 +107,7 @@
     MegatronLMSchedulerWrapper,
     T5TrainStep,
     avg_losses_across_data_parallel_group,
+    gather_across_data_parallel_groups,
 )
 from .megatron_lm import initialize as megatron_lm_initialize
 from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
index 8227fa5bc..5fcb5e68f 100644
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -771,6 +771,18 @@ class MegatronLMPlugin:
         default=False,
         metadata={"help": "Whether to set all logging options."},
     )
+    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."}
+    )
+    return_logits: bool = field(
+        default=False,
+        metadata={"help": "Whether to return logits from the model."},
+    )
+
+    # custom train step args
     custom_train_step_class: Optional[Any] = field(
         default=None,
         metadata={"help": "Custom train step class."},
@@ -779,11 +791,22 @@ class MegatronLMPlugin:
         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."}
+
+    # custom model args
+    custom_model_provider_function: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "Custom model provider function."},
     )
-    eval_interval: int = field(
-        default=1000, metadata={"help": "Interval between running evaluation on validation set."}
+    custom_prepare_model_function: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "Custom prepare model function."},
+    )
+
+    # remaining args such as enabling Alibi/ROPE positional embeddings,
+    # wandb logging, Multi-Query Attention, etc.
+    other_megatron_args: Optional[Dict[str, Any]] = field(
+        default=None,
+        metadata={"help": "Other Megatron-LM arguments. Please refer Megatron-LM"},
     )
 
     def __post_init__(self):
@@ -840,6 +863,8 @@ def __post_init__(self):
             self.megatron_lm_default_args["tensorboard_dir"] = self.tensorboard_dir
             if self.set_all_logging_options:
                 self.set_tensorboard_logging_options()
+        if self.other_megatron_args is not None:
+            self.megatron_lm_default_args.update(self.other_megatron_args)
 
     def set_network_size_args(self, model, batch_data=None):
         # Check if the model is either BERT, GPT or T5 else raise error
@@ -884,6 +909,8 @@ def set_network_size_args(self, model, batch_data=None):
             else:
                 self.seq_length = max_position_embeddings
             self.megatron_lm_default_args["seq_length"] = self.seq_length
+            self.megatron_lm_default_args["return_logits"] = self.return_logits
+            self.megatron_lm_default_args["tokenizer_type"] = "GPT2BPETokenizer"
         elif "t5" in model.config.model_type.lower():
             model_type_name = "t5"
             num_layers = model.config.num_layers
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index bdb53988a..85e20a746 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import importlib
+import os
 import sys
+from distutils.util import strtobool
 from functools import lru_cache
 
 import torch
@@ -90,10 +92,11 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    package_exists = importlib.util.find_spec("megatron") is not None
-    if package_exists:
-        megatron_version = parse(importlib_metadata.version("megatron-lm"))
-        return compare_versions(megatron_version, ">=", "2.2.0")
+    if strtobool(os.environ.get("USE_MEGATRON_LM", "False")) == 1:
+        package_exists = importlib.util.find_spec("megatron") is not None
+        if package_exists:
+            megatron_version = parse(importlib_metadata.version("megatron-lm"))
+            return compare_versions(megatron_version, ">=", "2.2.0")
     return False
 
 
diff --git a/src/accelerate/utils/megatron_lm.py b/src/accelerate/utils/megatron_lm.py
index cb988e360..afee2fcca 100644
--- a/src/accelerate/utils/megatron_lm.py
+++ b/src/accelerate/utils/megatron_lm.py
@@ -20,11 +20,12 @@
 import torch
 import torch.nn.functional as F
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP
 
 from ..optimizer import AcceleratedOptimizer
 from ..scheduler import AcceleratedScheduler
 from .imports import is_megatron_lm_available, is_transformers_available
-from .operations import send_to_device
+from .operations import recursively_apply, send_to_device
 
 
 if is_transformers_available():
@@ -41,6 +42,7 @@
         get_num_microbatches,
         get_tensorboard_writer,
         get_timers,
+        get_tokenizer,
         mpu,
         print_rank_0,
         print_rank_last,
@@ -56,16 +58,24 @@
         set_jit_fusion_options,
         write_args_to_tensorboard,
     )
-    from megatron.model import BertModel, GPTModel, ModelType, T5Model
+    from megatron.model import BertModel
+    from megatron.model import DistributedDataParallel as LocalDDP
+    from megatron.model import Float16Module, 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.text_generation.communication import broadcast_int_list, broadcast_tensor
+    from megatron.text_generation.generation import (
+        beam_search_and_return_on_first_stage,
+        generate_tokens_probs_and_return_on_first_stage,
+    )
     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,
+        unwrap_model,
     )
 
 
@@ -114,13 +124,21 @@ def model_provider_func(pre_process=True, post_process=True, add_encoder=True, a
 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)
+    if accelerator.state.megatron_lm_plugin.custom_prepare_model_function is not None:
+        if accelerator.state.megatron_lm_plugin.custom_model_provider_function is None:
+            raise ValueError(
+                "You must provide a `custom_model_provider_function` when using a `custom_prepare_model_function`."
+            )
+        custom_model_provider_func = accelerator.state.megatron_lm_plugin.custom_model_provider_function
+        model = accelerator.state.megatron_lm_plugin.custom_prepare_model_function(custom_model_provider_func)
+    else:
+        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
 
 
@@ -593,6 +611,12 @@ def __init__(self, args):
         self.loss_func = self.get_loss_func()
         self.forward_step = self.get_forward_step_func()
         self.eod_token = args.padded_vocab_size - 1
+        if args.vocab_file is not None:
+            tokenizer = get_tokenizer()
+            self.eod_token = tokenizer.eod
+        self.reset_position_ids = args.reset_position_ids
+        self.reset_attention_mask = args.reset_attention_mask
+        self.eod_mask_loss = args.eod_mask_loss
         if not args.model_return_dict:
             self.model_output_class = None
         else:
@@ -601,8 +625,6 @@ def __init__(self, args):
     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
@@ -621,7 +643,7 @@ def get_batch_megatron(data_iterator):
 
             # 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
+                tokens, self.eod_token, self.reset_position_ids, self.reset_attention_mask, self.eod_mask_loss
             )
 
             return tokens, labels, loss_mask, attention_mask, position_ids
@@ -638,7 +660,7 @@ def get_batch_transformer(data_iterator):
             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
+                tokens, self.eod_token, self.reset_position_ids, self.reset_attention_mask, True
             )
             return tokens, labels, loss_mask, attention_mask, position_ids
 
@@ -648,15 +670,24 @@ def get_batch_transformer(data_iterator):
             return get_batch_transformer
 
     def get_loss_func(self):
+        args = get_args()
+
         def loss_func(loss_mask, output_tensor):
-            losses = output_tensor.float()
+            if args.return_logits:
+                losses, logits = output_tensor
+            else:
+                losses = output_tensor
+            losses = losses.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]}
+            output_dict = {"lm loss": averaged_loss[0]}
+            if args.return_logits:
+                output_dict.update({"logits": logits})
+            return loss, output_dict
 
         return loss_func
 
@@ -1031,7 +1062,10 @@ def train_step(self, **batch_data):
             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)
+                if len(losses_reduced_for_key[0].shape) == 0:
+                    loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+                else:
+                    loss_reduced[key] = torch.concat(losses_reduced_for_key)
             return loss_reduced, skipped_iter, grad_norm, num_zeros_in_grad
         return {}, skipped_iter, grad_norm, num_zeros_in_grad
 
@@ -1079,7 +1113,10 @@ def eval_step(self, **batch_data):
             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)
+                if len(losses_reduced_for_key[0].shape) == 0:
+                    loss_reduced[key] = sum(losses_reduced_for_key) / len(losses_reduced_for_key)
+                else:
+                    loss_reduced[key] = torch.concat(losses_reduced_for_key)
             return loss_reduced
         else:
             return {}
@@ -1134,10 +1171,15 @@ def forward(self, **batch_data):
 
         loss = torch.tensor(0.0, device=args.local_rank)
         for key in loss_dict:
-            loss += loss_dict[key]
+            if len(loss_dict[key].shape) == 0:
+                loss += loss_dict[key]
+
+        logits = None
+        if "logits" in loss_dict:
+            logits = loss_dict["logits"]
         # 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 self.train_step_handler.model_output_class(loss=loss, logits=logits)
         return loss
 
     def log_eval_results(self):
@@ -1186,6 +1228,184 @@ def load_checkpoint(self, input_dir):
         if args.fp16 and self.iteration == 0:
             self.optimizer.reload_model_params()
 
+    def megatron_generate(
+        self,
+        inputs,
+        attention_mask=None,
+        max_length=None,
+        max_new_tokens=None,
+        num_beams=None,
+        temperature=None,
+        top_k=None,
+        top_p=None,
+        length_penalty=None,
+        **kwargs,
+    ):
+        """
+        Generate method for GPT2 model. This method is used for inference. Supports both greedy and beam search along
+        with sampling. Refer the Megatron-LM repo for more details
+
+        Args:
+            inputs (torch.Tensor): input ids
+            attention_mask (torch.Tensor, optional): attention mask. Defaults to None.
+            max_length (int, optional): max length of the generated sequence. Defaults to None.
+            Either this or max_new_tokens should be provided.
+            max_new_tokens (int, optional): max number of tokens to be generated. Defaults to None.
+            Either this or max_length should be provided.
+            num_beams (int, optional): number of beams to use for beam search. Defaults to None.
+            temperature (float, optional): temperature for sampling. Defaults to 1.0.
+            top_k (int, optional): top k tokens to consider for sampling. Defaults to 0.0.
+            top_p (float, optional): tokens in top p probability are considered for sampling. Defaults to 0.0.
+            length_penalty (float, optional): length penalty for beam search. Defaults to None.
+            kwargs: additional key-value arguments
+        """
+
+        # checking if required arguments are passed
+        args = get_args()
+        if args.model_type_name != "gpt":
+            raise NotImplementedError("Generate method is not implemented for this model")
+
+        if args.data_parallel_size > 1:
+            raise ValueError("Generate method requires data parallelism to be 1")
+
+        if args.sequence_parallel:
+            raise ValueError("Generate method requires sequence parallelism to be False")
+
+        if args.recompute_granularity is not None:
+            raise ValueError("Checkpoint activations cannot be set for inference")
+
+        if args.vocab_file is None:
+            raise ValueError("Vocab file is required for inference")
+
+        # Prepare inputs
+        if max_length is None and max_new_tokens is None:
+            raise ValueError("`max_length` or `max_new_tokens` are required for inference")
+
+        if temperature is None:
+            temperature = 1.0
+        elif not (0.0 < temperature <= 100.0):
+            raise ValueError("temperature must be a positive number less than or equal to 100.0")
+
+        if top_k is None:
+            top_k = 0
+        elif not (0 <= top_k <= 1000):
+            raise ValueError("top_k must be a positive number less than or equal to 1000")
+
+        if top_p is None:
+            top_p = 0.0
+        elif top_p > 0.0 and top_k > 0.0:
+            raise ValueError("top_p and top_k sampling cannot be set together")
+        else:
+            if not (0.0 <= top_p <= 1.0):
+                raise ValueError("top_p must be less than or equal to 1.0")
+
+        top_p_decay = kwargs.get("top_p_decay", 0.0)
+        if not (0.0 <= top_p_decay <= 1.0):
+            raise ValueError("top_p_decay must be less than or equal to 1.0")
+
+        top_p_bound = kwargs.get("top_p_bound", 0.0)
+        if not (0.0 <= top_p_bound <= 1.0):
+            raise ValueError("top_p_bound must be less than or equal to 1.0")
+
+        add_BOS = kwargs.get("add_BOS", False)
+        if not (isinstance(add_BOS, bool)):
+            raise ValueError("add_BOS must be a boolean")
+
+        beam_width = num_beams
+        if beam_width is not None:
+            if not isinstance(beam_width, int):
+                raise ValueError("beam_width must be an integer")
+            if beam_width < 1:
+                raise ValueError("beam_width must be greater than 0")
+            if inputs.shape[0] > 1:
+                return "When doing beam_search, batch size must be 1"
+
+        tokenizer = get_tokenizer()
+
+        stop_token = kwargs.get("stop_token", tokenizer.eod)
+        if stop_token is not None:
+            if not isinstance(stop_token, int):
+                raise ValueError("stop_token must be an integer")
+
+        if length_penalty is None:
+            length_penalty = 1.0
+
+        sizes_list = None
+        prompts_tokens_tensor = None
+        prompts_length_tensor = None
+        if torch.distributed.get_rank() == 0:
+            # Get the prompts length.
+            if attention_mask is None:
+                prompts_length_tensor = torch.cuda.LongTensor([inputs.shape[1]] * inputs.shape[0])
+            else:
+                prompts_length_tensor = attention_mask.sum(axis=-1).cuda()
+
+            if max_new_tokens is None:
+                max_new_tokens = max_length - inputs.shape[1]
+            if max_new_tokens <= 0:
+                raise ValueError("max_new_tokens must be greater than 0")
+
+            if add_BOS:
+                max_length = max_new_tokens + inputs.shape[1] + 1
+                # making sure that `max_length` is a multiple of 4 to leverage fused kernels
+                max_length = 4 * math.ceil(max_length / 4)
+                max_new_tokens = max_length - (inputs.shape[1] + 1)
+                padding = torch.cuda.LongTensor([[tokenizer.eod] * max_new_tokens] * inputs.shape[0])
+                prompts_tokens_tensor = torch.concat(
+                    [torch.unsqueeze(padding[:, 0], axis=-1), inputs.cuda(), padding], axis=-1
+                )
+            else:
+                # making sure that `max_length` is a multiple of 4 to leverage fused kernels
+                max_length = max_new_tokens + inputs.shape[1]
+                max_length = 4 * math.ceil(max_length / 4)
+                max_new_tokens = max_length - inputs.shape[1]
+                padding = torch.cuda.LongTensor([[tokenizer.eod] * max_new_tokens] * inputs.shape[0])
+                prompts_tokens_tensor = torch.concat([inputs.cuda(), padding], axis=-1)
+
+            # We need the sizes of these tensors for the boradcast
+            sizes_list = [
+                prompts_tokens_tensor.size(0),  # Batch size
+                prompts_tokens_tensor.size(1),
+            ]  # Sequence lenght
+
+        # First, broadcast the sizes.
+        sizes_tensor = broadcast_int_list(2, int_list=sizes_list, rank=0)
+
+        # Now that we have the sizes, we can boradcast the tokens
+        # and length tensors.
+        sizes = sizes_tensor.tolist()
+        context_tokens_tensor = broadcast_tensor(sizes, torch.int64, tensor=prompts_tokens_tensor, rank=0)
+        context_length_tensor = broadcast_tensor(sizes[0], torch.int64, tensor=prompts_length_tensor, rank=0)
+
+        # Run the inference
+        random_seed = kwargs.get("random_seed", 0)
+        torch.random.manual_seed(random_seed)
+        unwrapped_model = unwrap_model(self.base_model, (torchDDP, LocalDDP, Float16Module))
+        if beam_width is not None:
+            tokens, _ = beam_search_and_return_on_first_stage(
+                unwrapped_model,
+                context_tokens_tensor,
+                context_length_tensor,
+                beam_width,
+                stop_token=stop_token,
+                num_return_gen=1,
+                length_penalty=length_penalty,
+            )
+        else:
+            tokens, _, _ = generate_tokens_probs_and_return_on_first_stage(
+                unwrapped_model,
+                context_tokens_tensor,
+                context_length_tensor,
+                return_output_log_probs=False,
+                top_k=top_k,
+                top_p=top_p,
+                top_p_decay=top_p_decay,
+                top_p_bound=top_p_bound,
+                temperature=temperature,
+                use_eod_token_for_early_termination=True,
+            )
+        return tokens
+
 
 # other utilities
 def avg_losses_across_data_parallel_group(losses):
@@ -1197,3 +1417,25 @@ def avg_losses_across_data_parallel_group(losses):
     """
 
     return average_losses_across_data_parallel_group(losses)
+
+
+def gather_across_data_parallel_groups(tensor):
+    """
+    Recursively gather tensor in a nested list/tuple/dictionary of tensors from data parallel ranks.
+
+    Args:
+        tensor (nested list/tuple/dictionary of `torch.Tensor`):
+            The data to gather across data parallel ranks.
+
+    """
+
+    def _gpu_gather_one(tensor):
+        if tensor.ndim == 0:
+            tensor = tensor.clone()[None]
+        output_tensors = [
+            tensor.clone() for _ in range(torch.distributed.get_world_size(group=mpu.get_data_parallel_group()))
+        ]
+        torch.distributed.all_gather(output_tensors, tensor, group=mpu.get_data_parallel_group())
+        return torch.cat(output_tensors, dim=0)
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)


Review comments: body: Here you should lowercase, or use the str2bool conversion, just to make sure to catch variants.
diff_hunk: @@ -90,10 +91,11 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    package_exists = importlib.util.find_spec("megatron") is not None
-    if package_exists:
-        megatron_version = parse(importlib_metadata.version("megatron-lm"))
-        return compare_versions(megatron_version, ">=", "2.2.0")
+    if os.environ.get("USE_MEGATRON_LM", "false") == "true":
body: Fits on one line.
diff_hunk: @@ -1134,10 +1171,18 @@ def forward(self, **batch_data):
 
         loss = torch.tensor(0.0, device=args.local_rank)
         for key in loss_dict:
-            loss += loss_dict[key]
+            if len(loss_dict[key].shape) == 0:
+                loss += loss_dict[key]
+
+        logits = None
+        if "logits" in loss_dict:
+            logits = loss_dict["logits"]
         # 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 self.train_step_handler.model_output_class(
+                loss=loss,
+                logits=logits,
+            )
body: Very reluctant to add this here with the same name as in Transformers. This method does not have the same signature/behavior so users will be very confused to get something different. 
diff_hunk: @@ -1186,6 +1231,184 @@ def load_checkpoint(self, input_dir):
         if args.fp16 and self.iteration == 0:
             self.optimizer.reload_model_params()
 
+    def generate(
body: Done.
diff_hunk: @@ -90,10 +91,11 @@ def is_bf16_available(ignore_tpu=False):
 
 
 def is_megatron_lm_available():
-    package_exists = importlib.util.find_spec("megatron") is not None
-    if package_exists:
-        megatron_version = parse(importlib_metadata.version("megatron-lm"))
-        return compare_versions(megatron_version, ">=", "2.2.0")
+    if os.environ.get("USE_MEGATRON_LM", "false") == "true":
body: Done.
diff_hunk: @@ -1134,10 +1171,18 @@ def forward(self, **batch_data):
 
         loss = torch.tensor(0.0, device=args.local_rank)
         for key in loss_dict:
-            loss += loss_dict[key]
+            if len(loss_dict[key].shape) == 0:
+                loss += loss_dict[key]
+
+        logits = None
+        if "logits" in loss_dict:
+            logits = loss_dict["logits"]
         # 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 self.train_step_handler.model_output_class(
+                loss=loss,
+                logits=logits,
+            )
body: renamed to `megatron_generate`
diff_hunk: @@ -1186,6 +1231,184 @@ def load_checkpoint(self, input_dir):
         if args.fp16 and self.iteration == 0:
             self.optimizer.reload_model_params()
 
+    def generate(

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