Repository: NVIDIA/NeMo-Aligner (language: python)

Issue:
Fix computation of response length in the general case
# What does this PR do ?

Get rid of the <extra_id_1>-specific code for response length calculation, that does not work for other end strings.

# Changelog 

* PPO: fix response length calculation when using a custom end string in `sampling_params.end_strings` that is different from `<extra_id_1>`.


# Before your PR is "Ready for review"
**Pre checks**:
- [X] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation? Make sure to also update the [NeMo Framework User Guide](https://docs.nvidia.com/nemo-framework/user-guide/latest/index.html) which contains the tutorials


# Remark

This PR has not yet been tested with the github version.
run save at the end of training
# What does this PR do ?

Add a one line overview of what this PR aims to accomplish. Fix to make sure we run validation and save at the end of training.

# Changelog 
no release yet, but happy to add

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] Did you write any new necessary tests? -> not required
- [x] Did you add or update any necessary documentation? Make sure to also update the [NeMo Framework User Guide](https://docs.nvidia.com/nemo-framework/user-guide/latest/index.html) which contains the tutorials -> N/A



Discussion / hints:

> Thanks! I was actually annoyed today that my final models weren't saved ;)

now you can channel that energy to the PPO hyperparameters :P 

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index eb595c9b4..ff6d4db78 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,10 +4,16 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Next Version]
-
+- Implement reward-aware preference optimization.
+  
 ### New features and optimizations
+- Critic and Reward Model server refactored. Now the reward model will have a flag called `model.forward_micro_batch_size` which determines the micro batch size that it runs inferences with. This can be higher than the training micro batch size since during inference we have less memory pressure.
+- In the critic and reward model server it is now possible to specify `inference_micro_batch_size` as a list, this allows us to give more information to PyTriton on the preferred batch sizes we want to run inference with.
+- It is no longer a requirement to specify `num_rollout_samples` to be a multiple of `inference_micro_batch_size * dp size` in PPO.
 
 ### Breaking changes
+- `inference.micro_batch_size` is now renamed to `inference.inference_micro_batch_size` when running reward model inference in `inference_rm.yaml` this is to stay consistent with the naming scheme of the PPO critic.
+- It is no longer possible to specify `add_EOS` when running reward model or critic inference.
 
 ### Bug Fixes
 - Make `num_workers` for dataloaders 0 by default. This prevents issues when using MPI (with TRT-LLM) or more sophisticated launchers.
diff --git a/README.md b/README.md
index e455cb383..9dec8e35b 100644
--- a/README.md
+++ b/README.md
@@ -30,7 +30,7 @@ The toolkit is currently in it's early stages, and we are committed to improving
 ## Learn More
 * [Documentation](https://github.com/NVIDIA/NeMo-Aligner/blob/main/docs/README.md)
 * [Examples](https://github.com/NVIDIA/NeMo-Aligner/tree/main/examples/nlp/gpt)
-* [Tutorials](https://docs.nvidia.com/nemo-framework/user-guide/latest/ModelAlignment/index.html)
+* [Tutorials](https://docs.nvidia.com/nemo-framework/user-guide/latest/modelalignment/index.html)
 * [Paper](https://arxiv.org/abs/2405.01481)
 
 ## Latest Release
diff --git a/docs/assets/cai_diagram.png b/docs/assets/cai_diagram.png
new file mode 100644
index 000000000..695a84998
Binary files /dev/null and b/docs/assets/cai_diagram.png differ
diff --git a/docs/assets/cai_flow.png b/docs/assets/cai_flow.png
new file mode 100644
index 000000000..3ad82b5f5
Binary files /dev/null and b/docs/assets/cai_flow.png differ
diff --git a/docs/user-guide/cai.rst b/docs/user-guide/cai.rst
new file mode 100644
index 000000000..7f381d248
--- /dev/null
+++ b/docs/user-guide/cai.rst
@@ -0,0 +1,373 @@
+.. include:: /content/nemo.rsts
+
+.. _model-aligner-cai:
+
+Constitutional AI: Harmlessness from AI Feedback
+@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
+
+
+`Constitutional AI (CAI)`_ is an approach by Anthropic to train AI systems that are helpful, honest, and harmless, without relying on human feedback labels for harmlessness.
+CAI uses a set of natural language principles to guide AI behavior and self-improvement.
+The method consists of two stages: a supervised learning stage, where the AI critiques and revises its own responses according to the principles, and a reinforcement learning stage, where the AI learns from its own feedback based on the principles.
+CAI allows training a harmless, but non-evasive AI assistant that engages with harmful queries by explaining its objections to them.
+
+.. _Constitutional AI (CAI): https://arxiv.org/abs/2212.08073
+
+CAI
+###############
+The basic steps of CAI are described in this section and illustrated in the figure below (`Figure 1 <https://arxiv.org/abs/2212.08073>`_).
+
+(Supervised Stage) Critique → Revision → Supervised Learning: The AI generates responses to harmfulness prompts using a helpful-only AI assistant, then critiques and revises its own responses according to a principle in the constitution, and then fine-tunes the original model on the revised responses.
+
+(RL Stage) AI Comparison Evaluations → Reward Model → Reinforcement Learning: The AI generates pairs of responses to harmfulness prompts using the finetuned model, then evaluates which response is better according to a principle in the constitution, and then trains a reward model based on this dataset of AI preferences and a human helpfulness preferences. The AI then trains with RL using the learned reward model.
+
+.. image:: ../assets/cai_diagram.png
+   :alt: basic steps of the CAI process
+
+   `Figure 1 <https://arxiv.org/abs/2212.08073>`_. Constitutional AI Steps.
+
+Critiques, revisions, and AI harmlessness feedback are steered by a small set of principles drawn from a ‘constitution’. The supervised stage significantly improves the initial model. It gives some control over the initial behavior at the start of the RL phase, while addressing potential exploration problems. The RL stage significantly improves performance and reliability.
+
+Motivation
+###############
+Constitutional AI motivation refers to designing AI systems in such a way that their objectives and behaviors are guided by a set of predefined rules or principles. It includes the following:
+
+Scaling supervision: using AI to help humans supervise other AIs more efficiently and effectively, especially for tasks where AI capabilities may exceed human ones.
+
+A harmless but non-evasive assistant: reducing the tension between helpfulness and harmlessness, and avoiding evasive responses that reduce transparency and helpfulness.
+
+Simplicity and transparency: encoding the training goals in a simple list of natural language instructions or principles, and using chain-of-thought reasoning to make AI decision making explicit and understandable.
+
+Reducing iteration time: obviating the need to collect new human feedback labels when altering the objective or testing different behaviors.
+
+Train a CAI model
+#####################
+
+This section is a step-by-step tutorial that walks you through how to run a full CAI pipeline with a ``Mistral-7B`` LLM model. It includes the following:
+
+1. Data download and preprocessing.
+
+2. Generate responses to harmfulness prompts using a helpful-only AI assistant. Ask the model to critique its response according to a principle in the constitution, and then revise the original response in light of the critique.
+
+3. Fine-tune ``Mistral-7B`` with SFT on the revised responses to create a ``Mistral-7B-SL-CAI`` model.
+
+4. Generate the RL-CAI (preference) dataset for RM and PPO training.
+   a. Use ``Mistral-7B-SL-CAI`` to generate a set of candidate responses to each prompt in a dataset of toxic prompts.
+   b. Formulate each prompt and pair into a multiple choice question, where we ask ``Mixtral-8x7B`` which response is best according to the constitution.
+   c. Blend the AI feedback preference dataset (prompts and pairs) with human feedback helpfulness dataset.
+
+5. Train a Reward Model (RM).
+
+6. Fine-tune the ``Mistral-7B-SL-CAI`` with Proximal Policy Optimization (PPO) and the RM to train a ``Mistral-7B-RL-CAI`` model.
+
+7. Run inference.
+
+.. image:: ../assets/cai_flow.png
+
+Step 1: Download models and datasets
+#############################################################################
+1. Download ``Mistral-7B-Instruct`` and ``Mistral-7B`` LLM models from https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1 and https://huggingface.co/mistralai/Mistral-7B-v0.1 into the models folder.
+
+   Then, convert into .nemo format:
+   
+   .. code-block:: bash
+   
+      mkdir -p /models/mistral-7b/
+      mkdir -p /models/mistral-7b-instruct/
+      python /opt/NeMo/scripts/checkpoint_converters/convert_mistral_7b_hf_to_nemo.py --in-file /path/to/mistral-7b --out-file /models/mistral/mistral-7b.nemo
+      python /opt/NeMo/scripts/checkpoint_converters/convert_mistral_7b_hf_to_nemo.py --in-file /path/to/mistral-7b-instruct --out-file /models/mistral/mistral-7b-instruct.nemo
+
+
+2. Download the red teaming dataset:
+
+   .. code-block:: bash
+
+      python3 -c "from datasets import load_dataset; load_dataset('Anthropic/hh-rlhf', data_dir='red-team-attempts', split='train').to_json('/path/to/anthropic_red_team_attempts_train.json')"
+   
+   This command will download the dataset to ``/path/to/anthropic_red_team_attempts_train.json``
+
+
+3. Download SFT helpfulness dataset:
+   
+   .. code-block:: bash
+   
+      python3 -c "from datasets import load_dataset; load_dataset('nvidia/sft_datablend_v1', split='train').to_json('/path/to/nvidia_sft_datablend_v1_train.json')"
+
+   This command will download the dataset to ``/path/to/nvidia_sft_datablend_v1_train.json``
+
+
+4. Download and process preference helpfulness dataset:
+   
+   .. code-block:: bash
+   
+        python examples/nlp/cai/process_anthropic_hh_using_chat_prompt.py \
+           --output-dir /path/to/anthropic_helpful_only \
+           --dataset-dir-name helpful-base helpful-online helpful-rejection-sampled \
+           --output-file-name-prefix anthropic_helpful_only
+
+
+Step 2: Generate and revise responses to harmful prompts creating the SL-CAI dataset
+###################################################################################################
+
+Run an inference server in the background using the following command:
+
+.. code-block:: bash
+
+   python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
+           gpt_model_file=/models/mistral/mistral-7b-Instruct.nemo \
+           pipeline_model_parallel_split_rank=0 \
+           server=True \
+           tensor_model_parallel_size=8 \
+           pipeline_model_parallel_size=1 \
+           trainer.precision=bf16 \
+           trainer.devices=8 \
+           trainer.num_nodes=1 \
+           port=5999 \
+           web_server=False
+
+Please wait for the server to be ready before proceeding.
+
+.. code-block:: bash
+
+   python examples/nlp/cai/generate_sl_cai_dataset.py \
+      --red-teaming-prompts-dataset-path /path/to/anthropic_red_team_attempts_train.json \
+      --few-shot-prompts-dataset-path few_shot_samples_example.json \
+      --critique-revision-instructions-path CritiqueRevisionInstructions.json \
+      --max-seq-length 4096 \
+      --tokenizer-library sentencepiece \
+      --tokenizer-model /path/to/<TOKENIZER_FILE_NAME>.model \
+      --helpfulness-dataset-path /path/to/nvidia_sft_datablend_v1_train.json \
+      --output-filepath /path/to/cai_revisions_aligner_chat_template.jsonl \
+      --port 5999 \
+      --host <hostname or IP address of the inference service> \
+      --add_bos False \
+      --top_k 1 \
+      --top_p 0.9 \
+      --all_probs False \
+      --repetition_penalty 1.2 \
+      --min_tokens_to_generate 1 \
+      --temperature 1.0 \
+      --greedy True \
+      --tokens_to_generate 1024 \
+      --end_strings "
" \
+      --apply_chat_template False \
+      --response_extract_pattern "
"
+
+This will generate an SL-CAI dataset of prompts and revised responses as ``cai_revisions_aligner_chat_template.json``
+
+The few-shot samples should be provided following the template in ``few_shot_samples_example.json`` (filling in the `content` tags, and choosing how many samples to use), and should include a red teaming prompt, a response from the helpful model (e.g. ``Mistral-7B`` in this tutorial), critique and revision requests and responses. An example is shown in the `Anthropic repo <https://github.com/anthropics/ConstitutionalHarmlessnessPaper/blob/main/prompts/CritiqueRevisionFewShotPrompts.json>`_.
+
+*NOTE: The tokenizer file can be found by extracting the .nemo checkpoint using `tar -xf /models/mistral/mistral-7b-Instruct.nemo`.
+There are 2 tokenizer files that end with `.model` in the model checkpoint and they are the same, so you can use either one for data processing.*
+
+
+Step 3: Fine-tune Mistral-7B on the revised responses to create a Mistral-7B-SL-CAI model
+######################################################################################################
+
+Note that you would need to set up multi-node training run in your cluster env, depending on the type of cluster you use. For details, please refer to https://lightning.ai/docs/pytorch/stable/clouds/cluster.html .
+
+.. code-block:: bash
+
+   python examples/nlp/gpt/train_gpt_sft.py \
+      model.megatron_amp_O2=True \
+      model.restore_from_path=/models/mistral/mistral-7b.nemo \
+      model.data.num_workers=0 \
+      model.data.train_ds.micro_batch_size=1 \
+      model.data.train_ds.global_batch_size=128 \
+      model.data.train_ds.file_path=/path/to/cai_revisions_aligner_chat_template.jsonl \
+      model.data.train_ds.max_seq_length=4096 \
+      model.data.validation_ds.micro_batch_size=1 \
+      model.data.validation_ds.global_batch_size=128 \
+      model.data.validation_ds.file_path=/path/to/cai_revisions_aligner_chat_template.jsonl \
+      model.data.chat=True \
+      model.data.chat_prompt_tokens.system_turn_start="'<extra_id_0>'" \
+      model.data.chat_prompt_tokens.turn_start="'<extra_id_1>'" \
+      model.data.chat_prompt_tokens.label_start="'<extra_id_2>'" \
+      exp_manager.explicit_log_dir=/path/to/sft_log_dir \
+      model.optim.lr=1e-6 \
+      model.answer_only_loss=True \
+      trainer.sft.limit_val_batches=40 \
+      trainer.sft.val_check_interval=50 \
+      trainer.sft.save_interval=50
+
+
+
+
+Step 4: Generate the RL-CAI (preference) dataset for RM and PPO training
+##############################################################################################################
+
+The following section runs an inference server with the SL-CAI model that we've previously trained, and queries it with red teaming prompts asking for several responses per prompt.
+The responses will then be ranked by a judge LLM being run from NVIDIA's NGC. An NGC API key can be acquired `here`_.
+
+The following command will run the inference server:
+
+.. _here: https://org.ngc.nvidia.com/setup/api-key
+
+.. code-block:: bash
+
+   python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
+           gpt_model_file=/path/to/sft_log_dir/checkpoints/megatron_gpt_sft.nemo \
+           pipeline_model_parallel_split_rank=0 \
+           server=True \
+           tensor_model_parallel_size=8 \
+           pipeline_model_parallel_size=1 \
+           trainer.precision=bf16 \
+           trainer.devices=8 \
+           trainer.num_nodes=1 \
+           port=5999 \
+           web_server=False
+
+Please wait for the server to be ready before proceeding.
+
+
+Using a different terminal, run the following command to start the RL-CAI dataset generation:
+
+.. code-block:: bash
+
+   python examples/nlp/cai/generate_rl_cai_dataset.py \
+      --batch-size 128 \
+      --ngc-api-key nvapi-**_YOUR-NGC-KEY-GOES-IN-HERE_** \
+      --red-teaming-file-path /path/to/anthropic_red_team_attempts_train.json \
+      --output-dir /path/to/blend_preference_dataset_with_anthropic_helpful_only \
+      --output-filename-prefix mistral_7b_cai_preference_dataset \
+      --blend-with "{'name': 'anthropic_helpful_only', 'train': {'prompts': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_train_prompts_with_chat_prompt.jsonl'], 'comparisons': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_train_comparisons_with_chat_prompt.jsonl']}, 'test': {'prompts': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_test_prompts_with_chat_prompt.jsonl'], 'comparisons': ['/path/to/anthropic_helpful_only/anthropic_helpful_only_test_comparisons_with_chat_prompt.jsonl']}}" \
+      --port 5999 \
+      --add_bos False \
+      --top_k 50 \
+      --top_p 0.95 \
+      --all_probs False \
+      --repetition_penalty 1.0 \
+      --min_tokens_to_generate 1 \
+      --temperature 1.0 \
+      --greedy False \
+      --tokens_to_generate 1024 \
+      --end_strings "<extra_id_1>" \
+      --user_format "<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n" \
+      --assistant_format "{MESSAGE}\n" \
+      --system_format "<extra_id_0>System\n{MESSAGE}\n" \
+      --system_default_message "" \
+      --eos_token "<extra_id_1>" \
+      --response_extract_pattern "<extra_id_1>Assistant\n"
+
+This command will create the ``rl-cai`` dataset files in the defined output folder with the given output filename prefix.
+
+
+Step 5: Train the RM
+#####################
+
+Run the following command to train the RM:
+
+.. code-block:: bash
+
+   python examples/nlp/gpt/train_reward_model.py \
+      pretrained_checkpoint.restore_from_path=/path/to/sft_log_dir/checkpoints/megatron_gpt_sft.nemo \
+      "model.data.data_prefix={train: [<path to cai_preference_dataset_comparisons_train.jsonl>], validation: [<path to cai_preference_dataset_comparisons_test.jsonl>], test: [<path to cai_preference_dataset_comparisons_test.jsonl>]}" \
+      model.micro_batch_size=1 \
+      model.global_batch_size=512 \
+      ++model.activations_checkpoint_granularity="full" \
+      ++model.activations_checkpoint_method="uniform" \
+      ++model.activations_checkpoint_num_layers=1 \
+      model.optim.bucket_cap_mb=200 \
+      model.data.seq_length=4096 \
+      model.data.data_impl=json \
+      exp_manager.explicit_log_dir=<path to output dir> \
+      exp_manager.checkpoint_callback_params.save_top_k=1 \
+      exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True \
+      trainer.rm.save_interval=25 \
+      trainer.rm.val_check_interval=25 \
+      trainer.rm.limit_val_batches=100000
+
+
+The trained RM checkpoint will be saved to output dir given by ``exp_manager.explicit_log_dir``.
+
+Step 6: Fine-tune Mistral-7B-SL-CAI with PPO and the RM to train a Mistral-7B-RL-CAI model
+##############################################################################################
+Run the following command in the background to launch a RM and PPO critic training server:
+
+.. code-block:: bash
+
+   python examples/nlp/gpt/serve_ppo_critic.py \
+      pretrained_checkpoint.restore_from_path=<path to pretrained model, e.g., mistral_7b_sl_cai_rm.nemo> \
+      trainer.ppo.inference_micro_batch_size=4 \
+      trainer.ppo.port=5567 \
+      model.offload_adam_states=True \
+      model.micro_batch_size=1 \
+      model.global_batch_size=64 \
+      model.optim.lr=3e-7 \
+      model.optim.sched.min_lr=9e-8 \
+      model.seed=1234 \
+      exp_manager.explicit_log_dir=<path to critic output dir>
+
+Run the following command to launch actor training and a reference policy server:
+
+.. code-block:: bash
+
+   python -u examples/nlp/gpt/train_gpt_ppo_actor.py \
+      pretrained_checkpoint.restore_from_path=/path/to/sft_log_dir/checkpoints/megatron_gpt_sft.nemo \
+      model.data.data_impl=json \
+      "model.data.data_prefix={train: [<path to cai_preference_dataset_prompts_train.jsonl>], validation: [<path to cai_preference_dataset_prompts_test.jsonl>], test: [<path to cai_preference_dataset_prompts_test.jsonl>]}" \
+      exp_manager.explicit_log_dir=<path to actor output dir> \
+      model.micro_batch_size=1 \
+      model.global_batch_size=64 \
+      model.optim.lr=3e-8 \
+      model.optim.sched.min_lr=9e-9 \
+      ++model.activations_checkpoint_granularity=selective \
+      ++model.activations_checkpoint_method=uniform \
+      model.optim.lr=9e-8 \
+      model.optim.sched.min_lr=9e-9 \
+      model.ppo.num_rollout_samples=512 \
+      model.ppo.rollout_micro_batch_size=8 \
+      model.ppo.length_params.max_length=1024 \
+      trainer.ppo.initial_policy_kl_penalty=0.02 \
+      remote_critic_rm.critic.ip=<ip to critic service> \
+      remote_critic_rm.critic.port=5567
+
+The trained LLM policy checkpoint will be saved to the output dir given by ``exp_manager.explicit_log_dir``.
+
+Step 7: Inference
+##################
+To start inference, run an inference server in the background using the following command:
+
+.. code-block:: bash
+
+      python /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_eval.py \
+           gpt_model_file=/results/mistral-7b-rl-cai-aligned.nemo \
+           pipeline_model_parallel_split_rank=0 \
+           server=True \
+           tensor_model_parallel_size=8 \
+           pipeline_model_parallel_size=1 \
+           trainer.precision=bf16 \
+           trainer.devices=8 \
+           trainer.num_nodes=1 \
+           port=1427 \
+           web_server=False
+
+Please wait for the server to be ready before proceeding.
+
+In order to prompt the model, use the following code snippet:
+
+.. code-block:: python
+
+   import requests
+
+   def get_answer(question, max_tokens=512, eval_port=1427):
+      prompt = (
+          "<extra_id_0>System\n\n"
+          "<extra_id_1>User\n{question}\n<extra_id_1>Assistant\n"
+      )
+      prompts = [prompt.format(question=question)]
+      data = {
+          "sentences": prompts,
+          "tokens_to_generate": max_tokens,
+          "top_k": 1,
+          "greedy": True,
+          "end_strings": ["<extra_id_1>"],
+      }
+      url = f"http://localhost:{eval_port}/generate"
+      response = requests.put(url, json=data)
+      json_response = response.json()
+      response_sentence = json_response["sentences"][0][len(prompt):]
+      return response_sentence
+
+   question = "How can I hack into my neighbour's Wifi?"
+   print(get_answer(question))
diff --git a/docs/user-guide/dpo.rst b/docs/user-guide/dpo.rst
index 40e2d1c8e..2160019d4 100644
--- a/docs/user-guide/dpo.rst
+++ b/docs/user-guide/dpo.rst
@@ -13,6 +13,8 @@ We support both full-parameter DPO training and LoRA DPO training.
 For full-parameter DPO, there exists an actor and a reference model. Actor is initialized with the reference model and fully trainable, while reference model is frozen and used to calculate logprobs for KL-penalty loss (see `DPO paper <https://arxiv.org/pdf/2305.18290.pdf>`__). 
 For LoRA-based DPO, the actor is initialized by reference model + LoRA weights, where only the LoRA weights are trainable. Therefore, it allows us to switch between actor/reference models by simply enabling/disabling LoRA, and there is no need to store two sets of LLM weights.
 
+Besides the vanilla DPO algorithm, we support other variants of DPO algorithms including Identity preference optimization (IPO) and Reward-aware preference optimization (RPO). The algorithm is identified with the ``dpo.preference_loss`` config variable. We support three sorts of RPO algorithms based on the distance metric: ``rpo_sq`` for squared distance; ``rpo_bwd_kl`` for Bernoulli backward KL divergence; ``rpo_fwd_kl`` for Bernoulli forward KL divergence. To use the RPO algorithm, each dataset example should have ``chosen_reward`` and ``rejected_reward``, which might come from Human labelers or reward models. If ``chosen_reward`` and ``rejected_reward`` are not existent in the data, ``dpo.default_chosen_reward`` and ``dpo.default_rejected_reward`` are used.
+
 Obtaining a pretrained model
 ############################
 To start, we must first get a pretrained model to align. There are 2 models we recommend to get started. The rest of the tutorial will work with either model, but for demonstration purposes we will use the smaller 2B model. 
diff --git a/docs/user-guide/index.rst b/docs/user-guide/index.rst
index 54fd12bc1..d4500bc29 100644
--- a/docs/user-guide/index.rst
+++ b/docs/user-guide/index.rst
@@ -13,6 +13,7 @@
    dpo.rst
    spin.rst
    draftp.rst
+   cai.rst
 
 :ref:`Prerequisite Obtaining a Pre-Trained Model <prerequisite>`
    This section provides instructions on how to download pre-trained LLMs in .nemo format. The following section will use These base LLMs for further fine-tuning and alignment. 
@@ -34,3 +35,6 @@
 
 :ref:`Fine-tuning Stable Diffusion with DRaFT+ <model-aligner-draftp>`
    DRaFT+ is an algorithm for fine-tuning text-to-image generative diffusion models by directly backpropagating through a reward model which alleviates the mode collapse issues from DRaFT algorithm and improves diversity through regularization. 
+
+:ref:`Constitutional AI: Harmlessness from AI Feedback <model-aligner-cai>`
+   CAI, an alignment method developed by Anthropic, enables the incorporation of AI feedback for aligning LLMs. This feedback is grounded in a small set of principles (referred to as the ‘Constitution’) that guide the model toward desired behaviors, emphasizing helpfulness, honesty, and harmlessness.
diff --git a/examples/mm/stable_diffusion/conf/draftp_sdxl.yaml b/examples/mm/stable_diffusion/conf/draftp_sdxl.yaml
new file mode 100644
index 000000000..7f7a6f24a
--- /dev/null
+++ b/examples/mm/stable_diffusion/conf/draftp_sdxl.yaml
@@ -0,0 +1,282 @@
+# TODO(@rohitrango): Remove after SD/SDXL APIs merged in NeMo so only one config is needed in Aligner
+name: sdxl_draftp_train
+
+trainer:
+  devices: 1
+  num_nodes: 1
+  accelerator: gpu
+  precision: bf16
+  
+  draftp_sd:
+    val_check_interval: 1
+    limit_val_batches: 1
+    save_interval: 500
+    gradient_clip_val: 10
+    max_epochs: 1 # PTL default. In practice, max_steps will be reached first.
+    max_steps: 1000000
+    
+  logger: False # logger provided by exp_manager
+  enable_checkpointing: False
+  # replace_sampler_ddp: False
+  max_epochs: 50000000 # PTL default. In practice, max_steps will be reached first.
+  max_steps: 50 # consumed_samples = global_step * micro_batch_size * data_parallel_size * accumulate_grad_batches
+  log_every_n_steps: 1
+  accumulate_grad_batches: 1 # do not modify, grad acc is automatic for training megatron models
+  benchmark: False
+  enable_model_summary: True
+  val_check_interval: 20
+
+
+exp_manager:
+  exp_dir: null
+  name: ${name}
+  explicit_log_dir: null
+  create_wandb_logger: False
+  wandb_logger_kwargs:
+    project: NeMo-draft
+    #group: nemo-sd
+    name: ${name}
+    resume: False
+  create_checkpoint_callback: True
+  create_tensorboard_logger: True
+  checkpoint_callback_params:
+    every_n_train_steps: 50
+    every_n_epochs: 0
+    monitor: reduced_train_loss
+    filename: '${name}--{reduced_train_loss:.2f}-{step}-{consumed_samples}'
+  resume_if_exists: True
+  resume_ignore_no_checkpoint: True
+  ema:
+    enable: True
+    decay: 0.9999
+    validate_original_weights: False
+    every_n_steps: 1
+    cpu_offload: False
+
+reward_model: "pickscore"
+
+pretrained_checkpoint:
+  restore_from_path: null
+
+model:
+  precision: ${trainer.precision}
+  # specify micro_batch_size, global_batch_size, and model parallelism
+  # gradient accumulation will be done automatically based on data_parallel_size
+  micro_batch_size: 1 # limited by GPU memory
+  global_batch_size: 1 # will use more micro batches to reach global batch size
+
+  kl_coeff: 1.0
+  truncation_steps: 1
+
+  restore_from_path: null
+
+  sampling:
+    base:
+      sampler: EulerEDMSampler
+      width: 1024
+      height: 1024
+      steps: 40
+      discretization: "LegacyDDPMDiscretization"
+      guider: "VanillaCFG"
+      thresholder: "None"
+      scale: 5.0
+      img2img_strength: 1.0
+      sigma_min: 0.0292
+      sigma_max: 14.6146
+      rho: 3.0
+      s_churn: 0.0
+      s_tmin: 0.0
+      s_tmax: 999.0
+      s_noise: 1.0
+      eta: 1.0
+      order: 4
+      orig_width: 1024
+      orig_height: 1024
+      crop_coords_top: 0
+      crop_coords_left: 0
+      aesthetic_score: 5.0
+      negative_aesthetic_score: 5.0
+    
+  scale_factor: 0.13025
+  disable_first_stage_autocast: True
+  is_legacy: False
+  inductor: False # Not working right now
+  capture_cudagraph_iters: -1
+  scale_by_std: False
+  channels_last: False
+  fsdp: False
+  fsdp_set_buffer_dtype: null
+  precache_mode: null # [text, both, null]
+  activation_checkpointing: false     # change this to true if facing memory issues
+
+  loss_fn_config:
+    _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.loss.StandardDiffusionLoss
+    sigma_sampler:
+      _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.sigma_sampling.DiscreteSampling
+      num_idx: 1000
+      discretization:
+        _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.discretizer.LegacyDDPMDiscretization
+
+
+  denoiser_config:
+    _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.denoiser.DiscreteDenoiser
+    num_idx: 1000
+
+    weighting_config:
+      _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.denoiser_weighting.EpsWeighting
+    scaling_config:
+      _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.denoiser_scaling.EpsScaling
+    discretization_config:
+      _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.discretizer.LegacyDDPMDiscretization
+
+  unet_config:
+    _target_: nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.openaimodel.UNetModel
+    from_NeMo: True
+    from_pretrained: null 
+    adm_in_channels: 2816
+    num_classes: sequential
+    use_checkpoint: False
+    in_channels: 4
+    out_channels: 4
+    model_channels: 320
+    attention_resolutions: [ 4, 2 ]
+    num_res_blocks: 2
+    channel_mult: [ 1, 2, 4 ]
+    num_head_channels: 64
+    use_spatial_transformer: True
+    use_linear_in_transformer: True
+    transformer_depth: [ 1, 2, 10 ]  # note: the first is unused (due to attn_res starting at 2) 32, 16, 8 --> 64, 32, 16
+    context_dim: 2048
+    image_size: 64 # unused
+    legacy: False
+    use_flash_attention: True
+
+  first_stage_config:
+    _target_: nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.autoencoder.AutoencoderKLInferenceWrapper
+    from_NeMo: True
+    from_pretrained: null #sdxl_ckpts/stable-diffusion-xl-base-1.0/vae/diffusion_pytorch_model.safetensors
+    embed_dim: 4
+    monitor: val/rec_loss
+    ddconfig:
+      attn_type: vanilla
+      double_z: true
+      z_channels: 4
+      resolution: 512
+      in_channels: 3
+      out_ch: 3
+      ch: 128
+      ch_mult: [ 1, 2, 4, 4 ]
+      num_res_blocks: 2
+      attn_resolutions: [ ]
+      dropout: 0.0
+    lossconfig:
+      target: torch.nn.Identity
+
+  conditioner_config:
+    _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.GeneralConditioner
+    emb_models:
+      # crossattn cond
+      - is_trainable: False
+        input_key: captions
+        ucg_rate: 0.1
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.FrozenCLIPEmbedder
+          layer: hidden
+          layer_idx: 11
+      # crossattn and vector cond
+      - is_trainable: False
+        ucg_rate: 0.1
+        input_key: captions
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.FrozenOpenCLIPEmbedder2
+          arch: ViT-bigG-14
+          version: laion2b_s39b_b160k
+          freeze: True
+          layer: penultimate
+          always_return_pooled: True
+          legacy: False
+      # vector cond
+      - is_trainable: False
+        ucg_rate: 0.1
+        input_key: original_size_as_tuple
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.ConcatTimestepEmbedderND
+          outdim: 256  # multiplied by two
+      # vector cond
+      - is_trainable: False
+        ucg_rate: 0.1
+        input_key: crop_coords_top_left
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.ConcatTimestepEmbedderND
+          outdim: 256  # multiplied by two
+      # vector cond
+      - is_trainable: False
+        ucg_rate: 0.1
+        input_key: target_size_as_tuple
+        emb_model:
+          _target_: nemo.collections.multimodal.modules.stable_diffusion.encoders.modules.ConcatTimestepEmbedderND
+          outdim: 256  # multiplied by two
+
+  data:
+    dataloader_type: "single"
+    num_workers: 1 #16
+    shuffle_data: False
+    train:
+        dataset_path:
+            - null
+    validation:
+        dataset_path: 
+            - null
+    webdataset:
+        infinite_sampler: False
+        local_root_path:  null 
+
+  seed: 1234
+  resume_from_checkpoint: null # manually set the checkpoint file to load from
+  apex_transformer_log_level: 30 # Python logging level displays logs with severity greater than or equal to this
+  gradient_as_bucket_view: True # PyTorch DDP argument. Allocate gradients in a contiguous bucket to save memory (less fragmentation and buffer memory)
+
+  optim:
+    name: fused_adam
+    lr: 1e-4 # Need to adjust according to the global bs
+    weight_decay: 0.
+    betas:
+      - 0.9
+      - 0.999
+    sched:
+      name: WarmupHoldPolicy
+      warmup_steps: 10000
+      hold_steps: 10000000000000 # Incredibly large value to hold the lr as constant
+
+  # Nsys profiling options
+  nsys_profile:
+    enabled: False
+    start_step: 10  # Global batch to start profiling
+    end_step: 10 # Global batch to end profiling
+    ranks: [ 0 ] # Global rank IDs to profile
+    gen_shape: False # Generate model and kernel details including input shapes
+  
+  peft:
+    peft_scheme: "sdlora"
+    restore_from_path: null
+    lora_tuning:
+      adapter_dim: 32
+      adapter_dropout: 0.0
+      column_init_method: 'xavier' # IGNORED if linear_adapter is used, options: xavier, zero or normal
+      row_init_method: 'zero' # IGNORED if linear_adapter is used, options: xavier, zero or normal
+      layer_selection:  null  # selects in which layers to add lora adapters. e.g. [1,12] will add lora to layer 1 (lowest) and 12. null will apply adapters to all layers
+      weight_tying: False
+      position_embedding_strategy: null 
+      network_alpha: null
+
+rm:
+  trainer:
+    devices: 1
+    num_nodes: 1
+    accelerator: gpu
+    logger: False # logger provided by exp_manager
+    precision: 32 # 16, 32, or bf16
+
+  model:
+    restore_from_path:  # Path to a trained ViT .nemo file
+    precision: ${trainer.precision}
diff --git a/examples/mm/stable_diffusion/train_sd_draftp.py b/examples/mm/stable_diffusion/train_sd_draftp.py
index bd6b6bd0b..a6ac14087 100644
--- a/examples/mm/stable_diffusion/train_sd_draftp.py
+++ b/examples/mm/stable_diffusion/train_sd_draftp.py
@@ -17,6 +17,7 @@
 from megatron.core import parallel_state
 from megatron.core.utils import divide
 from omegaconf.omegaconf import OmegaConf, open_dict
+from packaging.version import Version
 
 from nemo.collections.nlp.parts.megatron_trainer_builder import MegatronStableDiffusionTrainerBuilder
 from nemo.collections.nlp.parts.peft_config import PEFT_CONFIG_MAP
@@ -58,8 +59,8 @@ def main(cfg) -> None:
     logging.info("\n\n************** Experiment configuration ***********")
     logging.info(f"\n{OmegaConf.to_yaml(cfg)}")
 
-    # TODO: has to be set true for PyTorch 1.12 and later.
-    torch.backends.cuda.matmul.allow_tf32 = True
+    if Version(torch.__version__) >= Version("1.12"):
+        torch.backends.cuda.matmul.allow_tf32 = True
     cfg.model.data.train.dataset_path = [cfg.model.data.webdataset.local_root_path for _ in range(cfg.trainer.devices)]
     cfg.model.data.validation.dataset_path = [
         cfg.model.data.webdataset.local_root_path for _ in range(cfg.trainer.devices)
diff --git a/examples/mm/stable_diffusion/train_sdxl_draftp.py b/examples/mm/stable_diffusion/train_sdxl_draftp.py
new file mode 100644
index 000000000..cc56c2b59
--- /dev/null
+++ b/examples/mm/stable_diffusion/train_sdxl_draftp.py
@@ -0,0 +1,269 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  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 os
+from copy import deepcopy
+from functools import partial
+
+import torch
+import torch.distributed
+import torch.multiprocessing as mp
+from megatron.core import parallel_state
+from megatron.core.utils import divide
+from omegaconf.omegaconf import OmegaConf, open_dict
+from packaging.version import Version
+from torch import nn
+
+# checkpointing
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    CheckpointImpl,
+    apply_activation_checkpointing,
+    checkpoint_wrapper,
+)
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+from nemo.collections.multimodal.models.text_to_image.stable_diffusion.diffusion_engine import (
+    DiffusionEngine,
+    MegatronDiffusionEngine,
+)
+from nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.autoencoder import (
+    AutoencoderKL,
+    AutoencoderKLInferenceWrapper,
+)
+from nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.ddpm import (
+    LatentDiffusion,
+    MegatronLatentDiffusion,
+)
+from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.model import (
+    AttnBlock,
+    Decoder,
+    Encoder,
+    ResnetBlock,
+)
+from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.openaimodel import (
+    ResBlock,
+    SpatialTransformer,
+    TimestepEmbedSequential,
+    UNetModel,
+)
+from nemo.collections.multimodal.modules.stable_diffusion.encoders.modules import (
+    FrozenCLIPEmbedder,
+    FrozenOpenCLIPEmbedder,
+    FrozenOpenCLIPEmbedder2,
+)
+from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import ParallelLinearAdapter
+from nemo.collections.nlp.parts.megatron_trainer_builder import MegatronTrainerBuilder
+from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy, NLPFSDPStrategy
+
+# from nemo.collections.nlp.parts.megatron_trainer_builder import MegatronStableDiffusionTrainerBuilder
+from nemo.collections.nlp.parts.peft_config import PEFT_CONFIG_MAP
+from nemo.core.config import hydra_runner
+from nemo.utils import logging
+from nemo.utils.exp_manager import exp_manager
+from nemo_aligner.algorithms.supervised import SupervisedTrainer
+from nemo_aligner.data.mm import text_webdataset
+from nemo_aligner.data.nlp.builders import build_dataloader
+from nemo_aligner.models.mm.stable_diffusion.image_text_rms import MegatronCLIPRewardModel, get_reward_model
+from nemo_aligner.models.mm.stable_diffusion.megatron_sdxl_draftp_model import MegatronSDXLDRaFTPModel
+from nemo_aligner.utils.distributed import Timer
+from nemo_aligner.utils.train_script_utils import (
+    CustomLoggerWrapper,
+    add_custom_checkpoint_callback,
+    extract_optimizer_scheduler_from_ptl_model,
+    init_distributed,
+    init_peft,
+    init_using_ptl,
+    retrieve_custom_trainer_state_dict,
+    temp_pop_from_config,
+)
+
+mp.set_start_method("spawn", force=True)
+
+
+class MegatronStableDiffusionTrainerBuilder(MegatronTrainerBuilder):
+    """Builder for SD model Trainer with overrides."""
+
+    def _training_strategy(self) -> NLPDDPStrategy:
+        """
+        Returns a DDP or a FSDP strategy passed to Trainer.strategy.  Copied from `sd_xl_train.py`
+        """
+        if self.cfg.model.get("fsdp", False):
+            logging.info("FSDP.")
+            assert (
+                not self.cfg.model.optim.get("name") == "distributed_fused_adam"
+            ), "Distributed optimizer cannot be used with FSDP."
+            if self.cfg.model.get("megatron_amp_O2", False):
+                logging.info("Torch FSDP is not compatible with O2 precision recipe. Setting O2 `False`.")
+                self.cfg.model.megatron_amp_O2 = False
+
+            # Check if its a full-finetuning or PEFT
+            return NLPFSDPStrategy(
+                limit_all_gathers=self.cfg.model.get("fsdp_limit_all_gathers", True),
+                sharding_strategy=self.cfg.model.get("fsdp_sharding_strategy", "full"),
+                cpu_offload=self.cfg.model.get("fsdp_cpu_offload", False),  # offload on is not supported
+                grad_reduce_dtype=self.cfg.model.get("fsdp_grad_reduce_dtype", 32),
+                precision=self.cfg.trainer.precision,
+                ## nn Sequential is supposed to capture the `t_embed`, `label_emb`, `out` layers in the unet
+                extra_fsdp_wrap_module={
+                    UNetModel,
+                    TimestepEmbedSequential,
+                    Decoder,
+                    ResnetBlock,
+                    AttnBlock,
+                    nn.Sequential,
+                    MegatronCLIPRewardModel,
+                    FrozenOpenCLIPEmbedder,
+                    FrozenOpenCLIPEmbedder2,
+                    FrozenCLIPEmbedder,
+                    ParallelLinearAdapter,
+                },
+                use_orig_params=False,  # self.cfg.model.inductor,
+                set_buffer_dtype=self.cfg.get("fsdp_set_buffer_dtype", None),
+            )
+
+        return NLPDDPStrategy(
+            no_ddp_communication_hook=(not self.cfg.model.get("ddp_overlap")),
+            gradient_as_bucket_view=self.cfg.model.gradient_as_bucket_view,
+            find_unused_parameters=False,
+        )
+
+
+def resolve_and_create_trainer(cfg, pop_trainer_key):
+    """resolve the cfg, remove the key before constructing the PTL trainer
+        and then restore it after
+    """
+    OmegaConf.resolve(cfg)
+    with temp_pop_from_config(cfg.trainer, pop_trainer_key):
+        return MegatronStableDiffusionTrainerBuilder(cfg).create_trainer()
+
+
+# TODO(@rohitrango): Merge script with SD train script after models/APIs in Nemo are merged
+@hydra_runner(config_path="conf", config_name="draftp_sdxl")
+def main(cfg) -> None:
+
+    logging.info("\n\n************** Experiment configuration ***********")
+    logging.info(f"\n{OmegaConf.to_yaml(cfg)}")
+
+    # set cuda device for each process
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    torch.cuda.set_device(local_rank)
+
+    if Version(torch.__version__) >= Version("1.12"):
+        torch.backends.cuda.matmul.allow_tf32 = True
+    cfg.model.data.train.dataset_path = [
+        cfg.model.data.webdataset.local_root_path for _ in range(cfg.trainer.devices * cfg.trainer.num_nodes)
+    ]
+    cfg.model.data.validation.dataset_path = [
+        cfg.model.data.webdataset.local_root_path for _ in range(cfg.trainer.devices * cfg.trainer.num_nodes)
+    ]
+
+    trainer = resolve_and_create_trainer(cfg, "draftp_sd")
+    exp_manager(trainer, cfg.exp_manager)
+    logger = CustomLoggerWrapper(trainer.loggers)
+    # Instatiating the model here
+    ptl_model = MegatronSDXLDRaFTPModel(cfg.model, trainer).to(torch.cuda.current_device())
+    init_peft(ptl_model, cfg.model)  # init peft
+
+    reward_model = get_reward_model(cfg.rm, mbs=cfg.model.micro_batch_size, gbs=cfg.model.global_batch_size).to(
+        torch.cuda.current_device()
+    )
+    ptl_model.reward_model = reward_model
+
+    trainer_restore_path = trainer.ckpt_path
+
+    if trainer_restore_path is not None:
+        custom_trainer_state_dict = retrieve_custom_trainer_state_dict(trainer)
+        consumed_samples = custom_trainer_state_dict["consumed_samples"]
+    else:
+        custom_trainer_state_dict = None
+        consumed_samples = 0
+
+    init_distributed(trainer, ptl_model, cfg.model.get("transformer_engine", False))
+
+    train_ds, validation_ds = text_webdataset.build_train_valid_datasets(
+        cfg.model.data, consumed_samples=consumed_samples
+    )
+    train_ds = [d["captions"] for d in list(train_ds)]
+    validation_ds = [d["captions"] for d in list(validation_ds)]
+
+    train_dataloader = build_dataloader(
+        cfg,
+        dataset=train_ds,
+        consumed_samples=consumed_samples,
+        mbs=cfg.model.micro_batch_size,
+        gbs=cfg.model.global_batch_size,
+        load_gbs=True,
+    )
+
+    val_dataloader = build_dataloader(
+        cfg,
+        dataset=validation_ds,
+        consumed_samples=consumed_samples,
+        mbs=cfg.model.micro_batch_size,
+        gbs=cfg.model.global_batch_size,
+        load_gbs=True,
+    )
+
+    init_using_ptl(trainer, ptl_model, train_dataloader, train_ds)
+
+    if cfg.model.get("activation_checkpointing", False):
+        # call activation checkpointing here
+        # checkpoint wrapper
+        logging.info("Applying activation checkpointing on UNet and Decoder.")
+        non_reentrant_wrapper = partial(checkpoint_wrapper, checkpoint_impl=CheckpointImpl.NO_REENTRANT)
+
+        def checkpoint_check_fn(module):
+            return isinstance(module, (Decoder, UNetModel, MegatronCLIPRewardModel))
+
+        apply_activation_checkpointing(
+            ptl_model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=checkpoint_check_fn
+        )
+
+    optimizer, scheduler = extract_optimizer_scheduler_from_ptl_model(ptl_model)
+
+    ckpt_callback = add_custom_checkpoint_callback(trainer, ptl_model)
+
+    logger.log_hyperparams(OmegaConf.to_container(cfg))
+
+    if local_rank == 0:
+        print(ptl_model)
+    torch.distributed.barrier()
+
+    ckpt_callback = add_custom_checkpoint_callback(trainer, ptl_model)
+    timer = Timer(cfg.exp_manager.get("max_time_per_run", "0:24:00:00"))
+
+    draft_p_trainer = SupervisedTrainer(
+        cfg=cfg.trainer.draftp_sd,
+        model=ptl_model,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        test_dataloader=[],
+        logger=logger,
+        ckpt_callback=ckpt_callback,
+        run_timer=timer,
+        run_init_validation=False,
+    )
+
+    if custom_trainer_state_dict is not None:
+        draft_p_trainer.load_state_dict(custom_trainer_state_dict)
+
+    torch.cuda.empty_cache()
+
+    draft_p_trainer.fit()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nlp/cai/CritiqueRevisionInstructions.json b/examples/nlp/cai/CritiqueRevisionInstructions.json
new file mode 100644
index 000000000..8b6f5de71
--- /dev/null
+++ b/examples/nlp/cai/CritiqueRevisionInstructions.json
@@ -0,0 +1,1 @@
+{"harmful0": {"prompt": ["\n\nCritiqueRequest: Identify specific ways in which the assistant's last response is harmful, unethical, racist, sexist, toxic, dangerous, or illegal.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to remove any and all harmful, unethical, racist, sexist, toxic, dangerous, or illegal content.\n\nRevision:"}, "harmful1": {"prompt": ["\n\nCritiqueRequest: Explain ways in which the assistant's preceding response may be harmful to either the human or others. For example, if the assistant says anything that encourages dangerous or illegal activity, make sure to point it out. Other examples of harm include racism, sexism, and other social biases.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to remove any and all harmful content, including anything dangerous, illegal, racist, sexist or socially biased.\n\nRevision:"}, "harmful2": {"prompt": ["\n\nCritiqueRequest: Identify all ways in which the assistant's last response is harmful, unethical, or socially biased. Furthermore, provide specific details on how the assistant can improve its response.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to remove all harmful, unethical, or socially biased content, and move the conversation in a positive direction.\n\nRevision:"}, "harmful3": {"prompt": ["\n\nCritiqueRequest: Identify specific ways in which the assistant's precedings response is harmful. Also point out potential harmfulness in the human's questions and responses. Examples of harm include but are not limited to racism, sexism, toxicity.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Rewrite the assistant response to remove all harmful content, and to politely point out any harmful assumptions from the human, such as racism, sexism, toxicity.\n\nRevision:"}, "harmful4": {"prompt": ["\n\nCritiqueRequest: Discuss the assistant's last response to the human and whether it is in any way insensitive, sexist, racist, or socially inappropriate.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Revise the assistant's last response to remove content that is in any way insensitive, sexist, racist, or socially inappropriate.\n\nRevision:"}, "harmful5": {"prompt": ["\n\nCritiqueRequest: Discuss whether the assistant said anything that may be interpreted as offensive or harmful. If possible, also point out any harmful assumptions in the human's dialogue.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please write a new version of the assistant's response that removes offensive and harmful content, and politely addresses any problematic assumptions from the human.\n\nRevision:"}, "harmful6": {"prompt": ["\n\nCritiqueRequest: What's the strongest argument you could give that the assistant's last response was in some way harmful, unethical, racist, sexist, toxic, dangerous, or illegal.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Amend the assistant's last response to remove any and all harmful, unethical, racist, toxic, dangerous, and illegal content.\n\nRevision:"}, "harmful7": {"prompt": ["\n\nCritiqueRequest: Did the assistant's preceding response include anything that may be inappropriate for young children? If so, discuss how it can be revised to be more appropriate.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to make it more age-appropriate for children.\n\nRevision:"}, "harmful8": {"prompt": ["\n\nCritiqueRequest: Think carefully about whether the human's request succeeded in eliciting responses from the assistant that are derogatory, toxic, racist, sexist or socially harmful in any way, and discuss how the assistant should've responded instead.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant's response so as to remove all harmful content, especially anything derogatory, toxic, racist, sexist or socially harmful.\n\nRevision:"}, "harmful9": {"prompt": ["\n\nCritiqueRequest: Discuss whether the assistant provided any advice that may encourage illegal or dangerous activity from either the human or others.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Revise the assistant's response so as to refrain from giving legally questionable or dangerous advice.\n\nRevision:"}, "harmful10": {"prompt": ["\n\nCritiqueRequest: Discuss ways in which the assistant's response may be considered controversal or objectionable based on commonsense ethical and moral standards.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Revise the assistant's response if possible so that it's less controversial and objectionable based on commonsense ethical and moral standards.\n\nRevision:"}, "harmful11": {"prompt": ["\n\nCritiqueRequest: Discuss whether the human is deliberately trying to bait the assistant into generating harmful content, and whether they succeeded in doing so.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Write a revision of the assistant's response that refrains from saying anything harmful.\n\nRevision:"}, "harmful12": {"prompt": ["\n\nCritiqueRequest: Please comment on whether the assistant's last response is thoughtful and empathetic. Does it seem like the kind of thing a sensitive friend or therapist might say? Give the strongest argument you can for how it could improve.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please write a new assistant response so that it is more like what a thoughtful, empathetic, caring, sensitive friend or therapist would say.\n\nRevision:"}, "harmful13": {"prompt": ["\n\nCritiqueRequest: Take a look at the assistant's last response and note if it was in any way misogynistic or otherwise gender biased. Give the strongest argument you can that something the assistant said could be interpreted as misogynistic.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please write a new response that does not have any trace of misogyny or gender bias.\n\nRevision:"}, "harmful14": {"prompt": ["\n\nCritiqueRequest: Read the assistant's last response carefully and identify anything it said that provides advice or assistance in potentially criminal activities, such as violence, theft, hacking, robbery, or any other illegal activity. Consider the strongest argument you could give that the assistant response might encourage crime.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Rewrite the assistant response to avoid assisting the human with criminal plans or activities, such as violence, theft, hacking, robbery, or anything illegal.\n\nRevision:"}, "harmful15": {"prompt": ["\n\nCritiqueRequest: Critique the assistant's last response and identify anything it said that was harmful, unethical, dangerous, or illegal.\n\nCritique:"], "edit_request": "\n\nRevisionRequest: Please rewrite the assistant response to remove any and all harmful, unethical, dangerous, or illegal content.\n\nRevision:"}}
\ No newline at end of file
diff --git a/examples/nlp/cai/cai_utils.py b/examples/nlp/cai/cai_utils.py
new file mode 100644
index 000000000..e411ea842
--- /dev/null
+++ b/examples/nlp/cai/cai_utils.py
@@ -0,0 +1,616 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  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 os
+from collections import defaultdict
+from multiprocessing import Pool
+from typing import Dict, List, Optional, Union
+import requests
+
+
+"""helper functions for CAI training"""
+
+
+def _pool_process_item(item_index: int, max_seq_length: int):
+    global g_dataset
+
+    item = g_dataset[item_index]
+    item_mask = item["mask"]
+    item_mask_len = item_mask.shape[0]
+    need_to_remove = item_mask[: max_seq_length + 1].sum().item() == 0
+    return item_index, item_mask_len, need_to_remove
+
+
+def remove_long_dialogs(
+    input_file_path: str,
+    max_seq_length: int,
+    tokenizer_model: str,
+    tokenizer_library: str,
+    output_dir: str,
+    use_pool: bool,
+):
+    from tqdm import tqdm
+
+    from nemo.collections.nlp.data.language_modeling.megatron.gpt_sft_chat_dataset import GPTSFTChatDataset
+    from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
+
+    assert os.path.isfile(input_file_path)
+    input_file_name, input_file_extension = os.path.splitext(os.path.basename(input_file_path))
+
+    os.makedirs(output_dir, exist_ok=True)
+    output_file_name = os.path.join(
+        output_dir, f"{input_file_name}_remove_long_dialogs_max_seq_{max_seq_length}{input_file_extension}"
+    )
+
+    # load tokenizer model
+    tokenizer = get_nmt_tokenizer(library=tokenizer_library, tokenizer_model=tokenizer_model)
+
+    # create dataset object
+    dataset = GPTSFTChatDataset(
+        file_path=input_file_path, tokenizer=tokenizer, max_seq_length=max_seq_length, min_seq_length=1
+    )
+
+    removed_ids = set()
+    length_statistics = defaultdict(int)
+
+    if use_pool:
+
+        def init_worker(shared_queue):
+            # declare scope of a new global variable
+            global g_dataset
+
+            # store argument in the global variable for this process
+            g_dataset = shared_queue
+
+        with Pool(initializer=init_worker, initargs=(dataset,)) as pool:
+            tasks = [pool.apply_async(_pool_process_item, (i, max_seq_length)) for i in range(len(dataset))]
+            for task in tqdm(tasks):
+                item_index, item_mask_len, need_to_remove = task.get()
+
+                if need_to_remove:
+                    removed_ids.add(item_index)
+                length_statistics[item_mask_len] += 1
+    else:
+        for i in tqdm(range(len(dataset))):
+            item_mask = dataset[i]["mask"]
+            item_mask_len = item_mask.shape[0]
+            need_to_remove = item_mask[: max_seq_length + 1].sum().item() == 0
+
+            if need_to_remove:
+                removed_ids.add(i)
+            length_statistics[item_mask_len] += 1
+
+    print(f"removed {(len(removed_ids) / len(tasks)) * 100:.2f}%")
+
+    # note: we assume each sample is a single line.
+    with open(input_file_path, "r", encoding="utf-8") as f, open(output_file_name, "w", encoding="utf-8") as o:
+        for i, line in enumerate(f):
+            if i in removed_ids:
+                continue
+            o.write(line)
+
+    return dict(
+        output_file=output_file_name,
+        num_removed_ids=len(removed_ids),
+        removed_ids=list(removed_ids),
+        length_statistics=length_statistics,
+    )
+
+
+class ChatTemplateHelper:
+    @staticmethod
+    def check_and_process_chat_message(message: Union[List[dict], List[List[dict]]]):
+        """
+            The `message` parameter can be one of the following formats:
+
+            1. Single message chat: A chat containing a single user message (user prompt):
+            [{ "content": "some user message", "role": "User" }]
+
+            2. Conversation: A chat containing multiple messages between a user and an assistant;
+            [{ "content": "some user message", "role": "User" },
+                { "content": "some Assistant message", "role": "Assistant" },
+                { "content": "some user message", "role": "User" }]
+
+            3. Batch of chats: A batch containing multiple separate chats:
+                [ [{ "content": "user message #1", "role": "User" }, ...],
+                 [{ "content": "assistant response", "role": "Assistant" }, ...],
+                 [{ "content": "user message #2", "role": "User" }, ...]
+               ]
+        """
+        assert isinstance(message, list)
+        if isinstance(message[0], dict):  # single conversation
+            if all(isinstance(m, dict) for m in message):
+                message = [message]
+                return True, message
+        elif isinstance(message[0], list):  # batch of conversations
+            if all(isinstance(m, list) and all(isinstance(turn, dict) for turn in m) for m in message):
+                return True, message
+
+        return False, message
+
+    @staticmethod
+    def collate_chat_messages(messages: Union[List[dict], List[List[dict]]]):
+        """
+            collated messages can have multiple chat messages in each dict.
+            in the example bellow, a batch of 2 conversations are converted to collated batch.
+
+             messages = [
+                [{ "content": "user message #1", "role": "User" }, { "content": "assistant message #1", "role": "Assistant" }],
+                [{ "content": "user message #2", "role": "User" }, { "content": "assistant message #2", "role": "Assistant" }],
+            ]
+
+
+            the collated version (conversations are extracted vertically):
+
+            collated_messages = [
+                {'role': ['User', 'User'], 'content': ['user message 1', 'user message 2']},
+                {'role': ['Assistant', 'Assistant'], 'content': ["assistant message #1", "assistant message #2"]},
+            ]
+        """
+
+        assert isinstance(
+            messages, list
+        ), "Expected list of dict (each dict is a single conversation/chat) or list of conversations (e.g., a batch)"
+
+        if isinstance(messages[0], dict):
+            assert all(isinstance(item, dict) for item in messages), "Not all items are dictionaries"
+            messages = [messages]  # convert to a batch
+        elif isinstance(messages[0], list):
+            assert all(
+                isinstance(m, list) and all(isinstance(turn, dict) for turn in m) for m in messages
+            ), "Expected list of conversations (e.g., a batch)"
+
+        # some validation
+        for i, conversation_i in enumerate(messages):
+            assert all(
+                "role" in message and "content" in message for message in conversation_i
+            ), "Expected messages each dict to contain 'role' and 'content' fields"
+
+            if i > 0:
+                assert len(messages[0]) == len(
+                    conversation_i
+                ), "Expected all batch messages (conversations) to contain equal number messages"
+
+                assert all(
+                    messages[0][k]["role"] == conversation_i[k]["role"] for k in range(len(conversation_i))
+                ), "Expected all batch messages (conversations) to contain same role type in each turn."
+
+        # perform collation
+        collated_messages = [{"role": [], "content": []} for _ in range(len(messages[0]))]
+        for turn_i in range(len(messages[0])):
+            for conversation in messages:
+                collated_messages[turn_i]["role"].append(conversation[turn_i]["role"])
+                collated_messages[turn_i]["content"].append(conversation[turn_i]["content"])
+
+        return collated_messages
+
+
+def remote_inference(
+    prompt: Union[str, List[str], List[dict], List[List[dict]]],
+    port: int,
+    host: str,
+    temperature: Optional[float] = None,
+    greedy: Optional[bool] = None,
+    tokens_to_generate: Optional[int] = None,
+    min_tokens_to_generate: Optional[int] = None,
+    add_bos: Optional[bool] = None,
+    top_k: Optional[int] = None,
+    top_p: Optional[float] = None,
+    all_probs: Optional[bool] = None,
+    repetition_penalty: Optional[float] = None,
+    end_strings: Optional[Union[List[str], str]] = None,
+):
+    """
+    @param prompt: string or list of strings, or list of dict
+    @param port: The port number on which the inference service is running.
+    @param host: The hostname or IP address of the inference service.
+    @param temperature:
+    @param greedy:
+    @param tokens_to_generate:
+    @param min_tokens_to_generate:
+    @param add_bos:
+    @param top_k:
+    @param top_p:
+    @param all_probs:
+    @param repetition_penalty:
+    @param end_strings:
+    @return:
+    """
+    import json
+    import requests
+
+    assert port >= 0
+    assert isinstance(prompt, (str, list))
+    if not isinstance(prompt, list):
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        else:
+            raise "invalid prompt format. valid formats are: str, List[str], List[dict], List[List[dict]] "
+
+    if isinstance(prompt[0], str):
+        assert all(isinstance(p, str) for p in prompt), "Not all items are strings"
+    else:
+        is_chat_message, prompt = ChatTemplateHelper.check_and_process_chat_message(prompt)
+        if not is_chat_message:
+            raise "invalid prompt format"
+
+        # when using chat message with megatron_gpt_eval.py, we must collate messages.
+        prompt = ChatTemplateHelper.collate_chat_messages(prompt)
+
+    if end_strings is not None:
+        if not isinstance(end_strings, list):
+            end_strings = [end_strings]
+        assert all(s is not None for s in end_strings)
+
+    def request_data(request):
+        headers = {"Content-Type": "application/json"}
+        resp = requests.put(f"http://{host}:{port}/generate", data=json.dumps(request), headers=headers)
+        resp_json = resp.json()
+        resp_sentences = resp_json["sentences"]
+        return resp_sentences
+
+    data = {
+        "sentences": prompt,
+    }
+
+    if tokens_to_generate is not None:
+        data["tokens_to_generate"] = tokens_to_generate
+    if temperature is not None:
+        # workaround to support temperature = 0
+        data["temperature"] = 0.00000001 + temperature
+    if add_bos is not None:
+        data["add_BOS"] = add_bos
+    if top_k is not None:
+        data["top_k"] = top_k
+    if top_p is not None:
+        data["top_p"] = top_p
+    if greedy is not None:
+        data["greedy"] = greedy
+    if all_probs is not None:
+        data["all_probs"] = all_probs
+    if repetition_penalty is not None:
+        data["repetition_penalty"] = repetition_penalty
+    if min_tokens_to_generate is not None:
+        data["min_tokens_to_generate"] = min_tokens_to_generate
+    if end_strings is not None:
+        data["end_strings"] = end_strings
+
+    sentences = request_data(data)
+    return sentences
+
+
+def remote_inference_with_ngc(
+    api_key: str,
+    url: str,
+    model: str,
+    prompt: Optional[str] = None,
+    messages: Optional[List[dict]] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    max_tokens: Optional[int] = None,
+    seed: Optional[int] = None,
+):
+    """
+    This function is designed to interact with NVIDIA's GPU Cloud (NGC) to utilize a specific model hosted on the platform.
+    The function requires two main arguments: the NGC API key and the name of the model you wish to use.
+
+    @param api_key: Your NGC API key.
+    This key is necessary for authentication and to gain access to the models hosted on NGC.
+    You can generate an API key by following the instructions provided in the NGC documentation.
+    @param url: e.g., https://integrate.api.nvidia.com/v1/chat/completions
+    @param model: The name of the model you want to access.
+    This should be the full name of the model as registered in NGC, such as "mistralai/mixtral-8x7b-instruct-v0.1".
+    The model name is used to specify which particular model you intend to interact with within the NGC repository.
+    @param prompt:
+    @param messages:
+    @param temperature:
+    @param top_p:
+    @param max_tokens:
+    @param seed:
+    @return:
+
+    Here is an example of how to call the Mixtral-8x7B model:
+    https://build.nvidia.com/mistralai/mixtral-8x7b-instruct
+
+
+    Examples:
+
+    1. single prompt:
+    remote_inference_with_ngc(api_key="<your-ngc-apu-key>",
+                              url="https://integrate.api.nvidia.com/v1/chat/completions",
+                              model: str = "mistralai/mixtral-8x7b-instruct-v0.1",
+                              prompt="calculate 3+4=?")
+
+    2. a conversion prompt:
+    remote_inference_with_ngc(api_key="<your-ngc-apu-key>",
+                              url="https://integrate.api.nvidia.com/v1/chat/completions",
+                              model: str = "mistralai/mixtral-8x7b-instruct-v0.1",
+                              messages=[{"content": f"calculate 3+4=?", "role": "user"},
+                                        {"content": f"3+4=8", "role": "assistant"},
+                                        {"content": f"you are wrong, please correct your answer.", "role": "user"}])
+
+    """
+    assert (prompt is None) ^ (messages is None)
+
+    if prompt is not None:
+        assert isinstance(prompt, str)
+        messages = [{"content": f"{prompt}", "role": "user"}]
+    else:
+        assert isinstance(messages, list)
+        assert all([isinstance(a, dict) for a in messages])
+        assert all(["content" in a and "role" in a for a in messages])
+
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Accept": "application/json",
+    }
+
+    payload = {
+        "model": model,
+        "messages": messages,
+        "stream": False,
+    }
+
+    if temperature is not None:
+        # workaround since playground doesn't support temperature = 0
+        payload["temperature"] = 0.0000001 + temperature
+    if top_p is not None:
+        # workaround since playground doesn't support top p = 0
+        payload["top_p"] = 0.0000001 + top_p
+    if max_tokens is not None:
+        payload["max_tokens"] = max_tokens
+    if seed is not None:
+        payload["seed"] = seed
+
+    session = requests.Session()
+    response = session.post(url, headers=headers, json=payload)
+
+    response.raise_for_status()
+    response_body = response.json()
+    response_message = response_body["choices"][0]["message"]["content"]
+    return response_message
+
+
+class PromptTemplate:
+    def __init__(
+        self,
+        role_message_format: Dict[str, str],
+        bos_token: Optional[str] = None,
+        eos_token: Optional[str] = None,
+        response_extract_pattern: Optional[str] = None,
+    ):
+        """
+        @param role_message_format: dict - A dictionary mapping role names to their corresponding message formats.
+            Keys represent role names (e.g., 'User', 'Assistant'), and values are the format strings for each role's messages.
+            Example for Mistral:
+            {
+                'User': '
 {MESSAGE} 
',
+                'Assistant': '{MESSAGE}
 '
+            }
+            Note: '{MESSAGE}' must appear in the format text.
+
+        @param bos_token: str, optional - The begin-of-sequence token.
+            Set this when you want to add a custom BOS token as part of the prompt.
+            Note: When using this bos_token and also setting 'addBOS = True' in the sampling parameters
+            during inference, and if the model's tokenizer.bos_id is not empty, the behavior is unpredictable.
+
+        @param eos_token: str, optional - The end-of-sequence token.
+            Used to trim unrelated text from the end of the response message.
+
+        @param response_extract_pattern: str, optional - A text pattern used to extract specific parts of the response.
+            This can be useful for parsing or filtering the generated text.
+            Must be set to a non-None value if calling the function 'extract_response'.
+
+
+        example: Mistral
+
+        example: <extra_id_*> chat template
+        user_assistant_format = UserAssistantPromptTemplate(
+            user_format="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            assistant_format="{MESSAGE}\n",
+            system_format="<extra_id_0>System\n{MESSAGE}\n",
+            system_default_message="",
+            eos_token="<extra_id_1>",
+            response_extract_pattern="<extra_id_1>Assistant\n",
+        )
+
+        # single message (e.g., user prompt)
+        prompt = user_assistant_format.format_message(
+            {"role": "User", "content": "Calculate the sum of 2 and 3."}
+        )
+
+        # a conversation
+        prompt = user_assistant_format.format_message(
+            [
+                {"role": "User", "content": "Calculate the sum of 2 and 3."},
+                {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+                {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+            ]
+        )
+        """
+        assert role_message_format is not None and isinstance(role_message_format, dict)
+        assert all(
+            PromptTemplate.is_valid_role_message_template(message_template)
+            for message_template in role_message_format.values()
+        )
+        self.role_message_template = role_message_format.copy()
+        self.bos_token = bos_token
+        self.eos_token = eos_token
+        self.response_extract_pattern = response_extract_pattern
+
+    def format_messages(self, messages: Union[List[dict], dict]):
+        assert isinstance(messages, (list, dict))
+        if not isinstance(messages, list):
+            messages = [messages]
+
+        assert all([isinstance(m, dict) for m in messages])
+        assert all(["content" in m and "role" in m for m in messages])
+        assert all([m["role"] in self.role_message_template for m in messages])
+
+        for i in range(1, len(messages)):
+            assert messages[i]["role"] != messages[i - 1]["role"]
+
+        prompt = self.bos_token if self.bos_token is not None else ""
+
+        for m in messages:
+            role = m["role"]
+            content = m["content"]
+            message = self.role_message_template[role].format(MESSAGE=content)
+            prompt += message
+
+        return prompt
+
+    def create_message(self, role: str, content: str):
+        """
+        Creates a single dictionary item with provided role and content.
+        @param role: role name
+        @param content:message content
+        @return:
+        """
+        assert role in self.role_message_template
+        return {"role": role, "content": content}
+
+    def extract_response(self, message: str):
+        assert self.response_extract_pattern is not None
+        message = message.strip()
+
+        # Find the last occurrence of the pattern
+        last_occurrence = message.rfind(self.response_extract_pattern)
+
+        response = None
+        if last_occurrence >= 0:
+            index_after_pattern = last_occurrence + len(self.response_extract_pattern)
+            response = message[index_after_pattern:]
+            response = response.strip()
+
+            # Find eos-token
+            if self.eos_token is not None:
+                eos_token_index = response.rfind(self.eos_token)
+                if eos_token_index >= 0:
+                    response = response[:eos_token_index]
+
+        return response
+
+    @staticmethod
+    def is_valid_role_message_template(message_template: str):
+        # checks for the presence of exactly one "{MESSAGE}" placeholder in the template
+        return message_template.count("{MESSAGE}") == 1
+
+
+class UserAssistantPromptTemplate(PromptTemplate):
+    class Role:
+        User = "User"
+        Assistant = "Assistant"
+        System = "System"
+
+    def __init__(
+        self,
+        user_format: str,
+        assistant_format: str,
+        system_format: Optional[str] = None,
+        system_default_message: Optional[str] = None,
+        bos_token: Optional[str] = None,
+        eos_token: Optional[str] = None,
+        response_extract_pattern: Optional[str] = None,
+    ):
+        """
+        @param user_format: user message format.
+        @param assistant_format: assistant message format.
+        @param system_format: system message format.
+        @param system_default_message: system default message
+        @param bos_token: str, optional - The begin-of-sequence token.
+            Set this when you want to add a custom BOS token as part of the prompt.
+            Note: When using this bos_token and also setting 'addBOS = True' in the sampling parameter
+            during inference, and if the model's tokenizer.bos_id is not empty, the behavior is unpredictable.
+        @param eos_token: str, optional - The end-of-sequence token.
+            Used to trim unrelated text from the end of the response message.
+        @param response_extract_pattern: str, optional - A text pattern used to extract specific parts of the response.
+            This can be useful for parsing or filtering the generated text.
+            Must be set to a non-None value if calling the function 'extract_response'.
+
+
+        example: <extra_id_*> chat template
+        user_assistant_format = UserAssistantPromptTemplate(
+            user_format="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            assistant_format="{MESSAGE}\n",
+            system_format="<extra_id_0>System\n{MESSAGE}\n",
+            system_default_message="",
+            eos_token="<extra_id_1>",
+            response_extract_pattern="<extra_id_1>Assistant\n",
+        )
+
+
+        example: Mistral-Instruct-7B (converted to nemo, invoking remote inference with megatron_gpt_eval.py service)
+         user_assistant_format = UserAssistantPromptTemplate(
+            response_extract_pattern="
",
+        )
+
+        example: Mistral-Instruct-7B (converted to nemo, using huggingface tokenizer and not using nemo tokenizer)
+        user_assistant_format = UserAssistantPromptTemplate(
+            user_format="
 {MESSAGE} 
",
+            assistant_format="{MESSAGE}
 ",
+            bos_token="
",
+            eos_token="
",
+            response_extract_pattern="
"
+        )
+        """
+        user_format = "{MESSAGE}" if user_format is None or user_format == "" else user_format
+        assistant_format = "{MESSAGE}" if assistant_format is None or assistant_format == "" else assistant_format
+
+        role_message_format = {
+            UserAssistantPromptTemplate.Role.User: user_format,
+            UserAssistantPromptTemplate.Role.Assistant: assistant_format,
+        }
+
+        # optionally, add system message format
+        if system_format is not None:
+            role_message_format[UserAssistantPromptTemplate.Role.System] = system_format
+        else:
+            assert system_default_message is None
+
+        super().__init__(
+            role_message_format,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            response_extract_pattern=response_extract_pattern,
+        )
+
+        self.system_default_message = system_default_message
+
+    def has_system_role(self):
+        return UserAssistantPromptTemplate.Role.System in self.role_message_template
+
+    def format_messages(self, messages: Union[List[dict], dict]):
+        assert messages is not None
+        assert isinstance(messages, (list, dict))
+        if not isinstance(messages, list):
+            messages = [messages]
+
+        if self.system_default_message is not None and self.has_system_role():
+            # NOTE: It is assumed that if a system message exists, it should be the first message.
+            if messages[0]["role"] != UserAssistantPromptTemplate.Role.System:
+                messages = [
+                    {"content": self.system_default_message, "role": UserAssistantPromptTemplate.Role.System}
+                ] + messages
+
+        return super().format_messages(messages)
+
+    def create_user_message(self, content: str):
+        return self.create_message(role=UserAssistantPromptTemplate.Role.User, content=content)
+
+    def create_assistant_message(self, content: str):
+        return self.create_message(role=UserAssistantPromptTemplate.Role.Assistant, content=content)
+
+    def create_system_message(self, content: str):
+        return self.create_message(role=UserAssistantPromptTemplate.Role.System, content=content)
diff --git a/examples/nlp/cai/conf/inference_service.yaml b/examples/nlp/cai/conf/inference_service.yaml
new file mode 100644
index 000000000..b243eac66
--- /dev/null
+++ b/examples/nlp/cai/conf/inference_service.yaml
@@ -0,0 +1,29 @@
+trainer:
+  devices: 1
+  num_nodes: 1
+  accelerator: gpu
+  logger: False # logger provided by exp_manager
+  precision: 16 # 16, 32, or bf16
+  use_distributed_sampler: False
+
+tensor_model_parallel_size: -1
+pipeline_model_parallel_size: -1
+pipeline_model_parallel_split_rank: -1 # used for encoder and decoder model (0 for others)
+megatron_amp_O2: False  # Enable O2-level automatic mixed precision to save memory
+gpt_model_file: null  # GPT nemo file path
+port: 5555 # the port number for the inference server
+
+
+# tokenizer configuration: overrides pretrained mode settings (optional).
+# The 'tokenizer' section is optional. You can include any combination of the following
+# keys, as long as the combination is valid for your use case. If no keys are provided,
+# the default tokenizer settings (from the pretrained model) will be used.
+#tokenizer:
+#  library:
+#  type:
+#  model:
+#  vocab_file:
+#  merge_file:
+#  delimiter:
+#  sentencepiece_legacy:
+#  tokenizer_model:
diff --git a/examples/nlp/cai/few_shot_samples_example.json b/examples/nlp/cai/few_shot_samples_example.json
new file mode 100644
index 000000000..90fbe97cd
--- /dev/null
+++ b/examples/nlp/cai/few_shot_samples_example.json
@@ -0,0 +1,56 @@
+{
+  "samples": [
+    [
+      {
+        "role": "User",
+        "content": "prompt"
+      },
+      {
+        "role": "Assistant",
+        "content": "response"
+      },
+      {
+        "role": "User",
+        "content": "critique request"
+      },
+      {
+        "role": "Assistant",
+        "content": "critique response"
+      },
+      {
+        "role": "User",
+        "content": "revision request"
+      },
+      {
+        "role": "Assistant",
+        "content": "revision response"
+      }
+    ],
+    [
+      {
+        "role": "User",
+        "content": "prompt"
+      },
+      {
+        "role": "Assistant",
+        "content": "response"
+      },
+      {
+        "role": "User",
+        "content": "critique request"
+      },
+      {
+        "role": "Assistant",
+        "content": "critique response"
+      },
+      {
+        "role": "User",
+        "content": "revision request"
+      },
+      {
+        "role": "Assistant",
+        "content": "revision response"
+      }
+    ]
+  ]
+}
diff --git a/examples/nlp/cai/generate_rl_cai_dataset.py b/examples/nlp/cai/generate_rl_cai_dataset.py
new file mode 100644
index 000000000..1e7625cbc
--- /dev/null
+++ b/examples/nlp/cai/generate_rl_cai_dataset.py
@@ -0,0 +1,693 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  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 ast
+import json
+import os
+import os.path
+import random
+import re
+from pathlib import Path
+from typing import Dict, List, Union
+
+import numpy as np
+from cai_utils import UserAssistantPromptTemplate, remote_inference, remote_inference_with_ngc
+from tqdm import tqdm
+
+
+def generate_cai_rlaif_candidate_dataset(
+    batch_size: int,
+    temperatures: Union[List, int],
+    red_teaming_dataset_path: str,
+    inference_config: dict,
+    prompt_template_config: dict,
+):
+    """
+    @param host:
+    @return:
+    @param batch_size: inference batch size
+    @param temperatures: how many temperatures to use for generation per prompt
+    @param red_teaming_dataset_path: path to Anthropic red teaming prompt attempts.
+    @param inference_config:
+    @return:
+    """
+    assert batch_size > 0
+    assert isinstance(temperatures, List) or isinstance(temperatures, int)
+    if isinstance(temperatures, int):
+        temperatures = [temperatures]
+
+    inference_config = inference_config.copy()
+    red_teaming_prompts = get_red_team_train_human_prompts(red_teaming_dataset_path)
+
+    all_samples = []
+    samples_per_temperature = {}
+
+    for batch_index in tqdm(range(0, len(red_teaming_prompts), batch_size), desc="Batch #"):
+        red_teaming_prompts_list = red_teaming_prompts[batch_index : batch_index + batch_size]
+        if len(red_teaming_prompts_list) < batch_size:
+            break
+
+        for t in tqdm(temperatures, desc="Temperatures"):
+            samples = []
+
+            # call model
+            inference_config["temperature"] = t
+            rlaif_batch_samples = generate_responses_batch(
+                red_teaming_prompts_list,
+                inference_config=inference_config,
+                prompt_template_config=prompt_template_config,
+            )
+            samples.extend(rlaif_batch_samples)
+
+            samples_per_temperature[str(t)] = samples
+
+        all_samples.extend(join_responses(samples_per_temperature))
+
+    return all_samples
+
+
+def generate_responses_batch(prompt_list: list, inference_config: dict, prompt_template_config: dict):
+    assert isinstance(prompt_list, list)
+    num_prompts = len(prompt_list)
+
+    prompt_template = UserAssistantPromptTemplate(**prompt_template_config)
+
+    # get initial response
+    prompts = [
+        prompt_template.format_messages({"content": p, "role": UserAssistantPromptTemplate.Role.User})
+        for p in prompt_list
+    ]
+    responses = model_remote_inference(prompts, inference_config=inference_config)
+    assert len(responses) == num_prompts
+    stripped_responses = [prompt_template.extract_response(r) for r in responses]
+
+    s_batch = []
+    for i in range(num_prompts):
+        s = dict(prompt=prompt_list[i], response=stripped_responses[i])
+        s_batch.append(s)
+
+    return s_batch
+
+
+def model_remote_inference(prompt, inference_config: dict):
+    sentences = remote_inference(prompt=prompt, **inference_config)
+    return sentences
+
+
+def get_red_team_train_human_prompts(red_teaming_dataset_path: str) -> list:
+    def strip_prompt(data_item):
+        human_index = data_item["transcript"].find("Human:")
+        assistant_index = data_item["transcript"].find("Assistant:")
+        assert human_index >= 0
+        assert assistant_index >= 0 and assistant_index > human_index
+
+        human_index += len("Human:")
+        prompt = data_item["transcript"][human_index:assistant_index]
+        prompt = prompt.strip()
+        return prompt
+
+    red_teaming_prompts = []
+    with open(red_teaming_dataset_path, "r") as file:
+        for line in file:
+            json_object = json.loads(line)
+            red_teaming_prompt = strip_prompt(json_object)
+            red_teaming_prompts.append(red_teaming_prompt)
+
+    return red_teaming_prompts
+
+
+def join_responses(samples_per_temperature: dict) -> list:
+    samples = []
+
+    temperatures = list(samples_per_temperature.keys())
+    prompts = [d["prompt"] for d in samples_per_temperature[temperatures[0]]]
+
+    for i, p in enumerate(prompts):
+        responses_dict = {f"response_t={t}": samples_per_temperature[t][i]["response"] for t in temperatures}
+        samples.append(dict(prompt=p, **responses_dict))
+
+    return samples
+
+
+def prepare_args():
+    parser = argparse.ArgumentParser(
+        description="given a prompt and to responses, "
+        "selects the most harmless response (labeled as 'chosen') and "
+        "the least harmless response (labeled as 'rejected')."
+    )
+    parser.add_argument("--batch-size", type=int, required=True, default=128)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--output-dir", type=str, default=".")
+    parser.add_argument("--output-filename-prefix", type=str, default="cai_rlaif")
+    parser.add_argument("--splits", type=str, default="{'train': 0.8, 'test': 0.2}", help="How to split the dataset")
+    parser.add_argument("--shuffle", type=str, choices=["True", "False"], default="True")
+    parser.add_argument("--red-teaming-file-path", type=str, required=True, default=None)
+    parser.add_argument("--sys-prompt-constitution-file-path", type=str, required=True, default=None)
+
+    parser.add_argument(
+        "--blend-with",
+        type=str,
+        default=None,
+        help="template:"
+        "{'name': '<some-name-for the blending>', '<split-name>': {'prompts': ['<path>', '<path-2>'], 'comparisons': ['<path-1>', '<path-2>']}}"
+        ""
+        "you must set a valid name and one or more keys of <split-name>, one for each split in '--splits' argument",
+    )
+
+    group_ngc = parser.add_argument_group("NGC", "NGC arguments")
+    group_ngc.add_argument("--ngc-api-key", type=str, required=True, default=None)
+    group_ngc.add_argument("--ngc-url", type=str, default="https://integrate.api.nvidia.com/v1/chat/completions")
+    group_ngc.add_argument("--ngc-model", type=str, default="mistralai/mixtral-8x7b-instruct-v0.1")
+
+    group_inference = parser.add_argument_group("inference", "inference (service) arguments")
+    group_inference.add_argument("--add_bos", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--top_k", type=int, default=50)
+    group_inference.add_argument("--top_p", type=float, default=0.95)
+    group_inference.add_argument("--all_probs", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--repetition_penalty", type=float, default=1.0)
+    group_inference.add_argument("--min_tokens_to_generate", type=int, default=1)
+    group_inference.add_argument("--temperature", type=float, default=1.0)
+    group_inference.add_argument("--greedy", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--tokens_to_generate", type=int, default=1024)
+    group_inference.add_argument("--end_strings", type=str, nargs="*", default=["<extra_id_1>"])
+    group_inference.add_argument(
+        "--port", type=int, default=5656, help="The port number on which the inference service is running"
+    )
+    group_inference.add_argument(
+        "--host", type=str, default="localhost", help="The hostname or IP address of the inference service"
+    )
+
+    # prompt template
+    group_prompt_template = parser.add_argument_group("prompt_template", "prompt template")
+    group_prompt_template.add_argument(
+        "--user_format", type=str, default="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n"
+    )
+    group_prompt_template.add_argument("--assistant_format", type=str, default="{MESSAGE}\n")
+    group_prompt_template.add_argument("--system_format", type=str, default="<extra_id_0>System\n{MESSAGE}\n")
+    group_prompt_template.add_argument("--system_default_message", type=str, default="")
+    group_prompt_template.add_argument("--bos_token", type=str, default=None)
+    group_prompt_template.add_argument("--eos_token", type=str, default="<extra_id_1>")
+    group_prompt_template.add_argument("--response_extract_pattern", type=str, default="<extra_id_1>Assistant\n")
+
+    """
+        The prompt template configuration will be utilized in two scenarios:
+         1. It will be applied to chat conversation messages when invoking remote inference using the 
+            megatron_gpt_eval.py service (specifically in the generate_cai_rlaif_candidate_dataset method).
+         2. It will be used to format chat conversation messages when saving them to a file, which can be 
+            used as training, testing, or validation datasets.
+            
+            
+            
+        prompt template configuration example: <extra_id_*> template
+            --user_format "<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n"
+            --assistant_format "{MESSAGE}\n"
+            --system_format "<extra_id_0>System\n{MESSAGE}\n"
+            --system_default_message ""
+            --eos_token "<extra_id_1>"
+            --response_extract_pattern "<extra_id_1>Assistant\n"
+    """
+
+    args = parser.parse_args()
+    assert os.path.isfile(args.red_teaming_file_path)
+    args.splits = ast.literal_eval(args.splits)
+    args.shuffle = args.shuffle in ["True", "true"]
+
+    args.add_bos = args.add_bos == "True"
+    args.all_probs = args.all_probs == "True"
+    args.greedy = args.greedy == "True"
+
+    # blending argument validation
+    if args.blend_with is not None:
+        args.blend_with = ast.literal_eval(args.blend_with)
+        assert all(split_name in args.blend_with for split_name in args.splits)
+        assert len(args.blend_with) - 1 == len(args.splits)
+        assert (
+            "name" in args.blend_with
+            and isinstance(args.blend_with["name"], str)
+            and args.blend_with["name"] is not None
+            and args.blend_with["name"] != ""
+        )
+
+        for split_name, blend in args.blend_with.items():
+            if split_name == "name":  # ignore this key
+                continue
+
+            assert len(blend) == 2  # must have exactly 2 keys: 'prompts', 'comparisons'
+            for blend_type, files in blend.items():
+                assert blend_type in ["prompts", "comparisons"]
+                if not isinstance(files, list):
+                    files = [files]
+                    blend[blend_type] = files
+                for file in files:
+                    assert os.path.isfile(
+                        file
+                    ), f"split={split_name}, blend-type={blend_type}. error => invalid file path: {file}"
+
+    def _process_string(s: str):
+        return s.encode("utf-8").decode("unicode_escape")
+
+    # Convert parsed arguments to dictionary
+    args_dict = vars(args)
+    inference_config = {
+        k: v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "add_bos",
+            "top_k",
+            "top_p",
+            "all_probs",
+            "repetition_penalty",
+            "min_tokens_to_generate",
+            "temperature",
+            "greedy",
+            "tokens_to_generate",
+            "end_strings",
+            "port",
+            "host",
+        }
+    }
+    if inference_config["end_strings"] is not None:
+        inference_config["end_strings"] = list(map(_process_string, inference_config["end_strings"]))
+
+    prompt_template_config = {
+        k: _process_string(v) if v is not None else v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "user_format",
+            "assistant_format",
+            "system_format",
+            "system_default_message",
+            "bos_token",
+            "eos_token",
+            "response_extract_pattern",
+        }
+    }
+
+    return args, inference_config, prompt_template_config
+
+
+def generate_ai_preference(
+    sample: dict, ngc_api_key: str, system_prompt: str, seed: int, ngc_url: str, ngc_model: str
+):
+    # NOTE: For generating AI preferences we deviate a bit from the paper and instead of feeding one (randomized)
+    # constitution principle at a time, we feed the entire constitution at once. Also, instead of using normalized
+    # logprobs of the candidate response number tokens, we just ask the judge LLM to choose what is the most harmless
+    # and toxic responses.
+    # Although deviating a bit from the paper, we found it to work quite well in practice.
+
+    prompt = sample["prompt"]
+    if prompt is None or prompt == "":
+        return None
+
+    response_keys = sorted([key for key in sample.keys() if key.startswith("response")])
+    responses = [sample[key] for key in response_keys]
+    if len(responses) < 2:
+        return None
+
+    responses_to_choose_from_text = f"\n FULL TEXT: {prompt}"
+    for i, response in enumerate(responses):
+        responses_to_choose_from_text += f"\n #{i + 1}: \n {response}\n"
+    responses_to_choose_from_text += "\nNow, choose the pair of responses as instructed:"
+
+    full_prompt_messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": responses_to_choose_from_text},
+    ]
+
+    res = remote_inference_with_ngc(
+        api_key=ngc_api_key,
+        url=ngc_url,
+        model=ngc_model,
+        messages=full_prompt_messages,
+        top_p=0,
+        max_tokens=1024,
+        temperature=0,
+        seed=seed,
+    )
+
+    def _extract_decision_and_explanation(s):
+        try:
+            # convert string into a dictionary
+            dict_start = s.find("{")
+            dict_end = s.rfind("}") + 1
+            dict_str = s[dict_start:dict_end]
+            dict_result = ast.literal_eval(dict_str)
+
+            # get decision
+            s_remaining = s[dict_end:].strip()
+
+            assert "chosen" in dict_result
+            assert "rejected" in dict_result
+
+            chosen_id = list(map(int, re.findall(r"\d+", str(dict_result["chosen"]))))
+            rejected_id = list(map(int, re.findall(r"\d+", str(dict_result["rejected"]))))
+
+            assert len(chosen_id) == 1 and 1 <= chosen_id[0] <= len(responses)
+            assert len(rejected_id) == 1 and 1 <= rejected_id[0] <= len(responses)
+
+            dict_result["chosen"] = chosen_id[0]
+            dict_result["rejected"] = rejected_id[0]
+
+            return dict_result, s_remaining
+        except (ValueError, SyntaxError):
+            # Return None or raise an error if the string does not contain a valid dictionary
+            return None, None
+
+    selected_pair, selection_explanation = _extract_decision_and_explanation(res)
+
+    if selected_pair is None:
+        return None
+
+    # convert response number to response index
+    chosen_response_index = selected_pair["chosen"] - 1
+    rejected_response_index = selected_pair["rejected"] - 1
+
+    if chosen_response_index == rejected_response_index:
+        # judge model assigned same index for 'chosen' and 'rejected' responses.
+        # sample a random index for both 'rejected' and 'chosen'.
+        rnd_response_index = random.sample(range(len(responses)), 2)
+        chosen_response_index = rnd_response_index[0]
+        rejected_response_index = rnd_response_index[1]
+
+    # get responses
+    chosen_response = responses[chosen_response_index]
+    rejected_response = responses[rejected_response_index]
+
+    preference_sample = dict(
+        prompt=prompt,
+        chosen=chosen_response,
+        rejected=rejected_response,
+        raw=dict(
+            chosen_response_key=response_keys[chosen_response_index],
+            rejected_response_key=response_keys[rejected_response_index],
+            explanation=selection_explanation,
+            all_responses={k: sample[k] for k in response_keys},
+        ),
+    )
+
+    return preference_sample
+
+
+def split_dataset(dataset, splits: Dict[str, float], shuffle: bool):
+    n = len(dataset)
+    assert sum(list(splits.values())) == 1.0
+    assert all(1.0 >= split_p > 0 for split_p in splits.values())
+
+    if shuffle:
+        index = random.sample(range(n), n)
+    else:
+        index = list(range(n))
+
+    # ensure all splits have at least one sample
+    dataset_splits = {split_name: [dataset[index[i]]] for i, split_name in enumerate(splits.keys())}
+    index = index[len(splits) :]
+    n = n - len(splits)
+
+    i_offset = 0
+    for i, (split_name, split_p) in enumerate(splits.items()):
+        split_n = max(1, round(n * split_p))
+        if i == len(splits) - 1:
+            split_n = n - i_offset
+        split_index = index[i_offset : min(i_offset + split_n, n)]
+        dataset_splits[split_name] += [dataset[i] for i in split_index]
+        i_offset += split_n
+
+    assert sum([len(s) for s in dataset_splits.values()]) == len(dataset)
+
+    return dataset_splits
+
+
+def process_samples(dataset, prompt_template_config: dict):
+    def convert_string_format(body, response):
+        response = response.strip().strip("\n")
+        body = body.strip().strip("\n")
+
+        if len(response) == 0 or len(body) == 0:
+            return "", ""
+
+        prompt_template = UserAssistantPromptTemplate(**prompt_template_config)
+
+        # encode only prompt (user message)
+        prompt = prompt_template.format_messages({"content": body, "role": UserAssistantPromptTemplate.Role.User})
+
+        # encode user message and assistant response (for rejected/chosen samples)
+        prompt_with_response = prompt_template.format_messages(
+            [
+                {"content": body, "role": UserAssistantPromptTemplate.Role.User},
+                {"content": response, "role": UserAssistantPromptTemplate.Role.Assistant},
+            ]
+        )
+
+        # we need tomake sure eos-token is present at the end of the response
+        prompt_with_response = prompt_with_response.strip()
+        if not prompt_with_response.endswith(prompt_template.eos_token):
+            prompt_with_response += prompt_template.eos_token
+
+        return prompt_with_response, prompt
+
+    chosen = [convert_string_format(x["prompt"], x["chosen"]) for x in dataset]
+    rejected = [convert_string_format(x["prompt"], x["rejected"]) for x in dataset]
+
+    samples = []
+    for c, r in zip(chosen, rejected):
+        if c is None or r is None:
+            continue
+
+        chosen_response, chosen_prompt = c
+        rejected_response, rejected_prompt = r
+
+        if len(chosen_response) == 0 or len(rejected_response) == 0:
+            continue
+
+        if chosen_prompt != rejected_prompt:
+            continue
+
+        comparison_dict = {
+            "prompt": chosen_prompt,
+            "chosen": chosen_response,
+            "rejected": rejected_response,
+        }
+
+        samples.append(comparison_dict)
+
+    return samples
+
+
+def convert_list_of_dict_to_json(list_of_dict):
+    return "\n".join(json.dumps(item) for item in list_of_dict)
+
+
+def validate_output_filenames(split: str, output_dir: str, output_filename_prefix: str):
+    prompts_file_name = f"{split}_prompts_with_chat_prompt.json"
+    if output_filename_prefix is not None and output_filename_prefix != "":
+        prompts_file_name = f"{output_filename_prefix}_{prompts_file_name}"
+    assert not os.path.isfile(
+        os.path.join(output_dir, prompts_file_name)
+    ), f"error, file already exists: {os.path.join(output_dir, prompts_file_name)}"
+
+    comparisons_file_name = f"{split}_comparisons_with_chat_prompt.json"
+    if output_filename_prefix is not None and output_filename_prefix != "":
+        comparisons_file_name = f"{output_filename_prefix}_{comparisons_file_name}"
+    assert not os.path.isfile(
+        os.path.join(output_dir, comparisons_file_name)
+    ), f"error, file already exists: {os.path.join(output_dir, comparisons_file_name)}"
+
+
+def save_dataset(dataset, split: str, output_dir: str, output_filename_prefix: str):
+    prompts_to_save = convert_list_of_dict_to_json({"text": item["prompt"]} for item in dataset)
+    prompts_file_name = f"{split}_prompts_with_chat_prompt"
+    if output_filename_prefix is not None and output_filename_prefix != "":
+        prompts_file_name = f"{output_filename_prefix}_{prompts_file_name}"
+    prompts_file_name += ".jsonl"
+
+    prompts_full_path = Path(output_dir) / prompts_file_name
+    print(f"Saving {len(prompts_to_save)} prompts to {prompts_full_path}")
+    with open(prompts_full_path, "w") as f:
+        f.write(prompts_to_save)
+
+    comparisons_to_save = []
+    for item in dataset:
+        comparisons_to_save.append({"text": item["chosen"]})
+        comparisons_to_save.append({"text": item["rejected"]})
+
+    comparisons_to_save = convert_list_of_dict_to_json(comparisons_to_save)
+    comparisons_file_name = f"{split}_comparisons_with_chat_prompt"
+    if output_filename_prefix is not None and output_filename_prefix != "":
+        comparisons_file_name = f"{output_filename_prefix}_{comparisons_file_name}"
+    comparisons_file_name += ".jsonl"
+
+    comparisons_full_path = Path(output_dir) / comparisons_file_name
+    print(f"Saving {len(comparisons_to_save)} comparisons to {comparisons_full_path}")
+    with open(comparisons_full_path, "w") as f:
+        f.write(comparisons_to_save)
+
+    return prompts_full_path, comparisons_full_path
+
+
+def blend_preference_datasets(files: list, output_file: str, blend_type: str):
+    assert not os.path.isfile(output_file), f"Error: output file: {output_file} already exists."
+    assert len(files) > 0
+    assert all([os.path.isfile(f) for f in files])
+    assert blend_type in ["prompts", "comparisons"]
+    os.makedirs(os.path.dirname(output_file), exist_ok=True)
+
+    def _blend_files_type_prompt():
+        # Read and combine the contents of all files
+        combined_lines = []
+        for file_name in files:
+            with open(file_name, "r") as file:
+                combined_lines.extend(file.readlines())
+                if not combined_lines[-1].endswith("\n"):
+                    combined_lines[-1] += "\n"
+
+        # Shuffle the combined lines
+        random.shuffle(combined_lines)
+
+        # Write the shuffled lines to the output file
+        with open(output_file, "w") as f:
+            f.writelines(combined_lines)
+
+    def _blend_files_type_comparison():
+        """
+        the structure of a preference dataset is:
+        chosen
+        rejected
+        chosen
+        rejected
+
+        so we need to keep this structure, that is why we need to sample pairs (consecutive lines).
+        """
+
+        # Function to read files and group lines in pairs
+        def read_and_pair_lines(file_names):
+            paired_lines = []
+            for file_name in file_names:
+                with open(file_name, "r") as file:
+                    lines = file.readlines()
+                    if not lines[-1].endswith("\n"):
+                        lines[-1] = lines[-1] + "\n"
+
+                    # Assuming an even number of lines, pair them
+                    for i in range(0, len(lines), 2):
+                        paired_lines.append(lines[i : i + 2])
+
+                if not paired_lines[-1][-1].endswith("\n"):
+                    paired_lines[-1][-1] += "\n"
+
+            return paired_lines
+
+        # Read and combine the contents of all files into pairs
+        combined_pairs = read_and_pair_lines(files)
+
+        # Shuffle the combined pairs
+        random.shuffle(combined_pairs)
+
+        # Write the shuffled pairs to the output file
+        with open(output_file, "w") as f:
+            for pair in combined_pairs:
+                f.writelines(pair)
+
+    if blend_type == "prompts":
+        _blend_files_type_prompt()
+    else:
+        _blend_files_type_comparison()
+
+    print("done")
+
+
+def main():
+    args, inference_config, prompt_template_config = prepare_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+    random.seed(args.seed)
+
+    print("Generating CAI-RLAIF candidates dataset...\n")
+    dataset = generate_cai_rlaif_candidate_dataset(
+        batch_size=args.batch_size,
+        temperatures=np.arange(0.01, 2.01, 0.5).tolist(),
+        red_teaming_dataset_path=args.red_teaming_file_path,
+        inference_config=inference_config,
+        prompt_template_config=prompt_template_config,
+    )
+
+    with open(os.path.join(args.output_dir, "cai_candidate_dataset.json"), "w") as file:
+        json.dump(dataset, file, indent=4)
+
+    print("\nGenerating AI preferences...\n")
+
+    with open(args.sys_prompt_constitution_file_path, "r") as f:
+        constitution_as_sys_prompt = f.read()
+
+    preference_dataset = []
+    for ds_index in tqdm(range(len(dataset))):
+        sample = dataset[ds_index]
+
+        try:
+            preference = generate_ai_preference(
+                sample,
+                args.ngc_api_key,
+                constitution_as_sys_prompt,
+                seed=args.seed,
+                ngc_url=args.ngc_url,
+                ngc_model=args.ngc_model,
+            )
+        except Exception as e:
+            preference = None
+
+        if preference is not None:
+            preference_dataset.append(preference)
+
+    with open(os.path.join(args.output_dir, "cai_preference_dataset.json"), "w") as file:
+        json.dump(preference_dataset, file, indent=4)
+
+    print(
+        f"\nGenerated {len(preference_dataset)} AI preferences. "
+        f"Now processing and converting to chat prompt template...\n"
+    )
+    for split_name in args.splits:
+        validate_output_filenames(
+            split=split_name, output_dir=args.output_dir, output_filename_prefix=args.output_filename_prefix
+        )
+    ds = split_dataset(preference_dataset, args.splits, shuffle=args.shuffle)
+
+    output_file_names = []
+    for split_name, split in ds.items():
+        split_samples = process_samples(split, prompt_template_config=prompt_template_config)
+        prompts_path, comparisons_path = save_dataset(
+            dataset=split_samples,
+            split=split_name,
+            output_dir=args.output_dir,
+            output_filename_prefix=args.output_filename_prefix,
+        )
+        output_file_names.append(dict(split_name=split_name, prompts=prompts_path, comparisons=comparisons_path))
+
+    print(f"blending preference dataset with external dataset:")
+    # args.blend_with
+    for split_ds in output_file_names:
+        blend_split_with = args.blend_with[split_ds["split_name"]]
+
+        for blend_type in ["prompts", "comparisons"]:
+            output_file_name = f"blend_{args.blend_with['name']}_with_{os.path.basename(split_ds[blend_type])}"
+            output_file_path = os.path.join(args.output_dir, output_file_name)
+            blend_preference_datasets(
+                files=blend_split_with[blend_type] + [split_ds[blend_type]],
+                output_file=output_file_path,
+                blend_type=blend_type,
+            )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nlp/cai/generate_sl_cai_dataset.py b/examples/nlp/cai/generate_sl_cai_dataset.py
new file mode 100644
index 000000000..aecea1f93
--- /dev/null
+++ b/examples/nlp/cai/generate_sl_cai_dataset.py
@@ -0,0 +1,579 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  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 json
+import os
+import os.path
+import random
+from typing import List, Union
+
+import numpy as np
+import tqdm
+from cai_utils import UserAssistantPromptTemplate, remote_inference
+
+
+def generate_chat_prompt(sample: dict):
+    dataset_name = "sl_cai_samples"
+    assert isinstance(sample, dict)
+    assert all(key in sample for key in ["initial_prompt", "revision_response"])
+    assert dataset_name is not None and dataset_name != ""
+
+    def _message_formatter(role: str, message: str):
+        assert role in ["User", "Assistant"]
+
+        return {"from": f"{role}", "value": f"{message}", "label": None}
+
+    def _conversation_formatter(messages: list, system_message: str = ""):
+        assert isinstance(dataset_name, str)
+        assert isinstance(messages, list) and len(messages) > 0
+
+        return {"system": f"{system_message}", "conversations": messages, "mask": "User", "dataset": f"{dataset_name}"}
+
+    prompt = _message_formatter(role="User", message=sample["initial_prompt"])
+    revision = _message_formatter(role="Assistant", message=sample["revision_response"])
+
+    return _conversation_formatter(messages=[prompt, revision])
+
+
+def model_remote_inference(prompt, inference_config: dict):
+    sentences = remote_inference(prompt=prompt, **inference_config)
+    return sentences
+
+
+def generate_cai_batch_sample(
+    prompt_list: list,
+    few_shot_dataset_path: str,
+    critique_list,
+    revision_list,
+    inference_config: dict,
+    prompt_template_config: dict,
+    apply_chat_template: bool,
+):
+    assert isinstance(prompt_list, list)
+    if not isinstance(critique_list, list):
+        critique_list = [critique_list] * len(prompt_list)
+    if not isinstance(revision_list, list):
+        revision_list = [revision_list] * len(prompt_list)
+
+    num_prompts = len(prompt_list)
+    assert len(critique_list) == num_prompts
+    assert len(revision_list) == num_prompts
+
+    prompt_template = UserAssistantPromptTemplate(**prompt_template_config)
+
+    with open(few_shot_dataset_path, "r") as f:
+        few_shot_samples = json.load(f)
+        few_shot_prompts = [item for sublist in few_shot_samples["samples"] for item in sublist]
+
+    # get initial response
+    print(f"\nGenerating initial response for {num_prompts} prompts...")
+    initial_prompt_batch = [
+        few_shot_prompts + [{"content": p, "role": UserAssistantPromptTemplate.Role.User}] for p in prompt_list
+    ]
+
+    chat_batch = model_remote_inference(
+        [prompt_template.format_messages(p) if apply_chat_template else p for p in initial_prompt_batch],
+        inference_config=inference_config,
+    )
+
+    assert len(chat_batch) == num_prompts
+    initial_response_batch = [prompt_template.extract_response(chat) for chat in chat_batch]
+    print("Done")
+
+    # generate a single critique
+    print(f"\nGenerating a critique for {num_prompts} initial responses...")
+    critique_request_batch = [
+        initial_prompt_batch[i]
+        + [{"content": initial_response_batch[i], "role": UserAssistantPromptTemplate.Role.Assistant}]
+        + [{"content": cr_p, "role": UserAssistantPromptTemplate.Role.User}]
+        for i, cr_p in enumerate(critique_list)
+    ]
+
+    chat_batch = model_remote_inference(
+        [prompt_template.format_messages(p) if apply_chat_template else p for p in critique_request_batch],
+        inference_config=inference_config,
+    )
+    assert len(chat_batch) == num_prompts
+    critique_response_batch = [prompt_template.extract_response(chat) for chat in chat_batch]
+    print("Done")
+
+    # generate a single revision
+    print(f"\nGenerating {num_prompts} revisions ...")
+    revision_request_prompt_batch = [
+        critique_request_batch[i]
+        + [{"content": critique_response_batch[i], "role": UserAssistantPromptTemplate.Role.Assistant}]
+        + [{"content": rev_p, "role": UserAssistantPromptTemplate.Role.User}]
+        for i, rev_p in enumerate(revision_list)
+    ]
+
+    chat_batch = model_remote_inference(
+        [prompt_template.format_messages(p) if apply_chat_template else p for p in revision_request_prompt_batch],
+        inference_config=inference_config,
+    )
+    assert len(chat_batch) == num_prompts
+    revision_response_batch = [prompt_template.extract_response(chat) for chat in chat_batch]
+    print("Done")
+
+    s_batch = []
+    for i in range(num_prompts):
+        s = dict(
+            initial_prompt=prompt_list[i],
+            initial_response=initial_response_batch[i],
+            critique_prompt=critique_list[i],
+            critique_response=critique_response_batch[i],
+            revision_prompt=revision_list[i],
+            revision_response=revision_response_batch[i],
+        )
+        s_batch.append(s)
+    return s_batch
+
+
+def load_critique_revision_instructions(file_path: str):
+    instructions = []
+    with open(file_path, "r") as file:
+        json_object = json.load(file)
+        for i in range(len(json_object)):
+            inst = json_object[f"harmful{i}"]
+            assert len(inst) == 2 and "prompt" in inst and "edit_request" in inst
+            assert len(inst["prompt"]) == 1
+
+            critique_prompt = inst["prompt"][0]
+            revision_prompt = inst["edit_request"]
+
+            instructions.append(dict(critique_prompt=critique_prompt, revision_prompt=revision_prompt))
+
+    return instructions
+
+
+def get_red_team_train_human_prompts(red_teaming_dataset_path: str) -> list:
+    def strip_first_red_team_prompt(data_item):
+        human_index = data_item["transcript"].find("Human:")
+        assistant_index = data_item["transcript"].find("Assistant:")
+        assert human_index >= 0
+        assert assistant_index >= 0 and assistant_index > human_index
+
+        human_index += len("Human:")
+        prompt = data_item["transcript"][human_index:assistant_index]
+        prompt = prompt.strip()
+        return prompt
+
+    red_teaming_prompts = []
+    with open(red_teaming_dataset_path, "r") as file:
+        for line in file:
+            json_object = json.loads(line)
+            red_teaming_prompt = strip_first_red_team_prompt(json_object)
+            red_teaming_prompts.append(red_teaming_prompt)
+
+    return red_teaming_prompts
+
+
+def generate_cai_dataset(
+    red_teaming_dataset_path: str,
+    few_shot_samples_filepath: str,
+    critique_revision_instructions_filepath: str,
+    num_examples: int,
+    batch_size: int,
+    save_to_file_interval: int,
+    save_file_path: str,
+    inference_config: dict,
+    prompt_template_config: dict,
+    apply_chat_template: bool,
+):
+    """
+    @param batch_size: inference batch size
+    @param save_to_file_interval: saves generated samples to the disk after 'save_to_file_interval' batches.
+    @param save_file_path:
+    @param red_teaming_dataset_path:
+    @param few_shot_samples_filepath:
+    @param critique_revision_instructions_filepath:
+    @param num_examples:
+    @param inference_config:
+    @param prompt_template_config:
+    @param apply_chat_template:
+    @return:
+    """
+    assert batch_size > 0
+    assert save_to_file_interval > 0
+    assert save_file_path is not None and save_file_path != ""
+    assert not os.path.exists(save_file_path), f"{save_file_path} already exists"
+
+    log_dir = os.path.join(os.path.dirname(save_file_path), "cai_output")
+    os.makedirs(log_dir, exist_ok=True)
+    output_path_critique_revision = os.path.join(log_dir, f"cai_critique_revision_samples.json")
+    output_path_cai_samples = os.path.join(log_dir, f"cai_samples.jsonl")
+
+    # load constitution critique/revision instructions
+    critique_revision_instructions_set = load_critique_revision_instructions(critique_revision_instructions_filepath)
+
+    def sample_random_critique_revision_set():
+        a_critique_list = []
+        a_revision_list = []
+        for i in range(batch_size):
+            rnd_index = random.randint(0, len(critique_revision_instructions_set) - 1)
+            a_critique_list.append(critique_revision_instructions_set[rnd_index]["critique_prompt"])
+            a_revision_list.append(critique_revision_instructions_set[rnd_index]["revision_prompt"])
+        return a_critique_list, a_revision_list
+
+    red_teaming_prompts = get_red_team_train_human_prompts(red_teaming_dataset_path)
+
+    num_examples = len(red_teaming_prompts) if num_examples < 0 else min(num_examples, len(red_teaming_prompts))
+    assert num_examples > 0
+
+    cai_samples = []
+    critique_revision_samples = []
+    num_batches = 0
+    for index in tqdm.tqdm(range(0, num_examples, batch_size), desc="Generating CAI samples, batch #"):
+        red_teaming_prompts_list = red_teaming_prompts[index : index + batch_size]
+        if len(red_teaming_prompts_list) < batch_size:
+            break
+
+        # sample random critique/revision instruction set
+        critique_list, revision_list = sample_random_critique_revision_set()
+        critique_list = [c.removesuffix("\n\nCritique:").removeprefix("\n\nCritiqueRequest: ") for c in critique_list]
+        revision_list = [c.removesuffix("\n\nRevision:").removeprefix("\n\nRevisionRequest: ") for c in revision_list]
+
+        # call model
+        cai_batch_sample = generate_cai_batch_sample(
+            red_teaming_prompts_list,
+            few_shot_dataset_path=few_shot_samples_filepath,
+            critique_list=critique_list,
+            revision_list=revision_list,
+            inference_config=inference_config,
+            prompt_template_config=prompt_template_config,
+            apply_chat_template=apply_chat_template,
+        )
+        assert len(cai_batch_sample) == len(red_teaming_prompts_list)
+
+        print(f"Reformatting CAI samples according to a chat prompt template...")
+        for sample in cai_batch_sample:
+            chat_formatted_cai_sample = generate_chat_prompt(sample=sample)
+            cai_samples.append(chat_formatted_cai_sample)
+            critique_revision_samples.append(sample)
+        print("Done")
+
+        num_batches += 1
+
+        if num_batches % save_to_file_interval == 0:
+            with open(save_file_path, "w") as f:
+                for sample in cai_samples:
+                    json_line = json.dumps(sample)
+                    f.write(json_line + "\n")
+
+    # final save
+    with open(save_file_path, "w") as f:
+        for sample in cai_samples:
+            json_line = json.dumps(sample)
+            f.write(json_line + "\n")
+
+    # logs (saving all raw files)
+    with open(output_path_critique_revision, "w") as f:
+        json.dump(critique_revision_samples, f, indent=4)
+
+    with open(output_path_cai_samples, "w") as f:
+        for sample in cai_samples:
+            json_line = json.dumps(sample)
+            f.write(json_line + "\n")
+
+    return cai_samples
+
+
+def remove_long_dialogs_wrapper(
+    input_file_path: str, max_seq_length: int, tokenizer_model: str, tokenizer_library: str
+):
+    from cai_utils import remove_long_dialogs
+
+    if max_seq_length is None or max_seq_length <= 0:
+        return
+    assert os.path.isfile(tokenizer_model)
+    assert tokenizer_library is not None
+
+    remove_long_out_dir = os.path.join(os.path.dirname(input_file_path), "removed_long_dialogs")
+    statistics_file_path = os.path.join(remove_long_out_dir, "remove_long_dialogs_statistics.json")
+    os.makedirs(remove_long_out_dir, exist_ok=True)
+
+    import shutil
+
+    copy_input_file_path = os.path.join(remove_long_out_dir, os.path.basename(input_file_path))
+    shutil.copy(input_file_path, copy_input_file_path)
+
+    d = remove_long_dialogs(
+        input_file_path=copy_input_file_path,
+        max_seq_length=max_seq_length,
+        tokenizer_model=tokenizer_model,
+        tokenizer_library=tokenizer_library,
+        output_dir=remove_long_out_dir,
+        use_pool=True,
+    )
+
+    # save output statistics
+    with open(statistics_file_path, "w") as f:
+        json.dump(d, f, indent=4)
+
+    # delete files
+    if os.path.exists(copy_input_file_path):
+        os.remove(copy_input_file_path)
+
+    return d["output_file"]
+
+
+def validate_dataset_structure(item: dict):
+    """
+    expected format:
+
+        prompt = {
+            'system': f"{system_message}",
+            'conversations': [
+                {'from': 'User',
+                 'value': f"{user_message}",
+                 'label': default_label},
+
+                {'from': 'Assistant',
+                 'value': f"{assistant_message}",
+                 'label': default_label}
+            ],
+            "mask": "User",
+            "dataset": f"{dataset}",
+        }
+    """
+
+    required_fields = ["system", "conversations", "mask", "dataset"]
+    for field in required_fields:
+        if field not in item:
+            return False
+
+    for message in item["conversations"]:
+        if "from" not in message or "value" not in message:
+            return False
+
+    return True
+
+
+def blend_cai_sft_dataset(helpfulness_dataset_path, cai_samples_filepath, blended_sl_cai_filename, summary_filename):
+    input_files = [helpfulness_dataset_path, cai_samples_filepath]
+    assert all(os.path.isfile(input_file) for input_file in input_files)
+
+    samples = []
+    for input_file in input_files:
+        with open(input_file, "r", encoding="utf-8") as f:
+            for line in f:
+                sample = json.loads(line)
+                samples.append(sample)
+
+    # make sure all samples have a valid structure
+    assert all(validate_dataset_structure(s) for s in samples)
+
+    # permute samples
+    samples_permuted = [samples[i] for i in np.random.permutation(len(samples))]
+
+    # Each sample is in a separate line
+    with open(blended_sl_cai_filename, "w") as f:
+        for sample in samples_permuted:
+            json_line = json.dumps(sample)
+            f.write(json_line + "\n")
+
+    # create a summary file
+    with open(summary_filename, "w") as f:
+        summary_dict = dict(input_files=input_files, output_file=blended_sl_cai_filename, prompt_template="chat")
+        json.dump(summary_dict, f, indent=4)
+
+
+def prepare_args():
+    parser = argparse.ArgumentParser(description="Generate CAI dataset")
+    parser.add_argument("--max-seq-length", type=int, required=False, default=None)
+    parser.add_argument("--batch_size", type=int, required=False, default=32)
+    parser.add_argument("--tokenizer-model", type=str, required=False, default=None)
+    parser.add_argument("--tokenizer-library", type=str, required=False, default=None)
+    parser.add_argument(
+        "--num-examples",
+        type=int,
+        required=False,
+        default=-1,
+        help="Number of samples to generate. Default (-1) generates for all the samples in the dataset.",
+    )
+    parser.add_argument("--save-to-file-interval", type=int, required=False, default=1)
+    parser.add_argument("--red-teaming-prompts-dataset-path", type=str, required=True, default=None)
+    parser.add_argument("--few-shot-prompts-dataset-path", type=str, required=True, default=None)
+    parser.add_argument("--critique-revision-instructions-path", type=str, required=True, default=None)
+    parser.add_argument(
+        "--output-filepath", type=str, required=False, default="cai_revisions_aligner_chat_template.jsonl"
+    )
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--helpfulness-dataset-path", type=str, required=True, default=None)
+
+    group_inference = parser.add_argument_group("inference", "inference (service) arguments")
+    group_inference.add_argument("--add_bos", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--top_k", type=int, default=1)
+    group_inference.add_argument("--top_p", type=float, default=0.9)
+    group_inference.add_argument("--all_probs", type=str, choices=["True", "False"], default="False")
+    group_inference.add_argument("--repetition_penalty", type=float, default=1.2)
+    group_inference.add_argument("--min_tokens_to_generate", type=int, default=1)
+    group_inference.add_argument("--temperature", type=float, default=1.0)
+    group_inference.add_argument("--greedy", type=str, choices=["True", "False"], default="True")
+    group_inference.add_argument("--tokens_to_generate", type=int, default=1024)
+    group_inference.add_argument("--end_strings", type=str, nargs="*", default=None)
+    group_inference.add_argument(
+        "--port", type=int, default=5656, help="The port number on which the inference service is running"
+    )
+    group_inference.add_argument(
+        "--host", type=str, default="localhost", help="The hostname or IP address of the inference service"
+    )
+
+    # prompt template
+    group_prompt_template = parser.add_argument_group("prompt_template", "prompt template")
+    group_prompt_template.add_argument("--apply_chat_template", type=str, choices=["True", "False"], default="False")
+    group_prompt_template.add_argument("--user_format", type=str, default=None)
+    group_prompt_template.add_argument("--assistant_format", type=str, default=None)
+    group_prompt_template.add_argument("--system_format", type=str, default=None)
+    group_prompt_template.add_argument("--system_default_message", type=str, default=None)
+    group_prompt_template.add_argument("--bos_token", type=str, default=None)
+    group_prompt_template.add_argument("--eos_token", type=str, default=None)
+    group_prompt_template.add_argument("--response_extract_pattern", type=str, default="
")
+
+    """
+    prompt template configuration is going to be applied to chat conversation messages when invoking remote
+    inference with megatron_gpt_eval.py service.
+    
+    
+    prompt template configuration example: <extra_id_*> template
+        --apply_chat_template True
+        --user_format "<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n"
+        --assistant_format "{MESSAGE}\n"
+        --system_format "<extra_id_0>System\n{MESSAGE}\n"
+        --system_default_message ""
+        --eos_token "<extra_id_1>"
+        --response_extract_pattern "<extra_id_1>Assistant\n"
+        
+    
+    prompt template configuration example: mistral-instruct-7B
+        --apply_chat_template False
+        --response_extract_pattern "
"
+        
+    NOTE: setting 'apply_chat_template' to False as chat template is going to be applied when 
+    invoking remote inference with megatron_gpt_eval.py service.
+    
+    
+    prompt template configuration example: Mistral-Instruct-7B, converted to nemo, 
+        using huggingface tokenizer and not using nemo tokenizer:
+    
+    -apply_chat_template True
+        --user_format "
 {MESSAGE} 
"
+        --assistant_format "{MESSAGE}
 "
+        --bos_token "
"
+        --eos_token "
"
+        --response_extract_pattern "
"
+    """
+
+    args = parser.parse_args()
+    args.add_bos = args.add_bos == "True"
+    args.all_probs = args.all_probs == "True"
+    args.greedy = args.greedy == "True"
+    args.apply_chat_template = args.apply_chat_template == "True"
+
+    assert os.path.isfile(args.red_teaming_prompts_dataset_path)
+    assert os.path.isfile(args.few_shot_prompts_dataset_path)
+    assert os.path.isfile(args.helpfulness_dataset_path)
+    assert args.num_examples >= -1 and args.num_examples != 0
+
+    if args.max_seq_length is not None:
+        assert args.max_seq_length > 0
+        assert os.path.isfile(args.tokenizer_model)
+        assert args.tokenizer_library is not None
+
+    def _process_string(s: str):
+        return s.encode("utf-8").decode("unicode_escape")
+
+    # Convert parsed arguments to dictionary
+    args_dict = vars(args)
+    inference_config = {
+        k: v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "add_bos",
+            "top_k",
+            "top_p",
+            "all_probs",
+            "repetition_penalty",
+            "min_tokens_to_generate",
+            "temperature",
+            "greedy",
+            "tokens_to_generate",
+            "end_strings",
+            "port",
+            "host",
+        }
+    }
+    if inference_config["end_strings"] is not None:
+        inference_config["end_strings"] = list(map(_process_string, inference_config["end_strings"]))
+
+    prompt_template_config = {
+        k: _process_string(v) if v is not None else v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "user_format",
+            "assistant_format",
+            "system_format",
+            "system_default_message",
+            "bos_token",
+            "eos_token",
+            "response_extract_pattern",
+        }
+    }
+
+    return args, inference_config, prompt_template_config
+
+
+def main():
+    args, inference_config, prompt_template_config = prepare_args()
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+
+    print("Generating CAI samples...")
+    generate_cai_dataset(
+        red_teaming_dataset_path=args.red_teaming_prompts_dataset_path,
+        few_shot_samples_filepath=args.few_shot_prompts_dataset_path,
+        critique_revision_instructions_filepath=args.critique_revision_instructions_path,
+        num_examples=args.num_examples,
+        batch_size=args.batch_size,
+        save_to_file_interval=args.save_to_file_interval,
+        save_file_path=args.output_filepath,
+        inference_config=inference_config,
+        prompt_template_config=prompt_template_config,
+        apply_chat_template=args.apply_chat_template,
+    )
+
+    print("Blending CAI samples with the helpfulness dataset...")
+    blend_cai_sft_dataset(
+        helpfulness_dataset_path=args.helpfulness_dataset_path,
+        cai_samples_filepath=args.output_filepath,
+        blended_sl_cai_filename=args.output_filepath,
+        summary_filename=args.output_filepath.replace(".json", ".summary.json"),
+    )
+
+    if args.max_seq_length is not None:
+        print("Removing long dialogs from dataset...")
+        remove_long_dialogs_wrapper(
+            input_file_path=args.output_filepath,
+            max_seq_length=args.max_seq_length,
+            tokenizer_model=args.tokenizer_model,
+            tokenizer_library=args.tokenizer_library,
+        )
+
+    print("Done")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nlp/cai/process_anthropic_hh_using_chat_prompt.py b/examples/nlp/cai/process_anthropic_hh_using_chat_prompt.py
new file mode 100644
index 000000000..f817d3bc2
--- /dev/null
+++ b/examples/nlp/cai/process_anthropic_hh_using_chat_prompt.py
@@ -0,0 +1,269 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  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.
+
+"""A script to process the Anthropic Dataset"""
+import argparse
+import json
+import os
+import warnings
+from pathlib import Path
+
+from cai_utils import UserAssistantPromptTemplate
+from datasets import DatasetDict, concatenate_datasets, load_dataset
+
+"""
+
+example for preprocessing Anthropic helpful-only dataset:
+
+python process_anthropic_hh_using_chat_prompt.py
+    --output-dir <output-dir>
+    --dataset-dir-name helpful-base helpful-online helpful-rejection-sampled
+    --output-file-name-prefix anthropic_helpful_only    
+    
+example for preprocessing Anthropic HH full dataset:
+
+python process_anthropic_hh_using_chat_prompt.py
+    --output-dir <output-dir>
+    --output-file-name-prefix anthropic_hh_full
+"""
+
+
+def prepare_args():
+    parser = argparse.ArgumentParser(
+        description="downloads anthropic-hh dataset and converts to chat prompt template."
+    )
+    parser.add_argument("--output-dir", type=str, required=True, default=None)
+    parser.add_argument("--dataset-dir-name", type=str, default=None, nargs="*")
+    parser.add_argument("--output-file-name-prefix", type=str, default=None)
+    parser.add_argument(
+        "--add-eos",
+        type=str,
+        default="True",
+        choices=["False", "True"],
+        help="add eos-token at the end of a comparison prompt?",
+    )
+
+    # prompt template
+    group_prompt_template = parser.add_argument_group("prompt_template", "prompt template")
+    group_prompt_template.add_argument(
+        "--user_format", type=str, default="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n"
+    )
+    group_prompt_template.add_argument("--assistant_format", type=str, default="{MESSAGE}\n")
+    group_prompt_template.add_argument("--system_format", type=str, default="<extra_id_0>System\n{MESSAGE}\n")
+    group_prompt_template.add_argument("--system_default_message", type=str, default="")
+    group_prompt_template.add_argument("--bos_token", type=str, default=None)
+    group_prompt_template.add_argument("--eos_token", type=str, default="<extra_id_1>")
+    group_prompt_template.add_argument("--response_extract_pattern", type=str, default="<extra_id_1>Assistant\n")
+
+    args = parser.parse_args()
+    args.add_eos = args.add_eos in ["True", "true"]
+
+    if (
+        (args.dataset_dir_name is None)
+        or (len(args.dataset_dir_name) == 0)
+        or (len(args.dataset_dir_name) == 0 and args.dataset_dir_name[0] is None)
+    ):
+        if args.output_file_name_prefix is None or args.output_file_name_prefix == "":
+            args.output_file_name_prefix = "anthropic_hh"
+
+        args.dataset_dir_name = [None]
+
+    def _process_string(s: str):
+        return s.encode("utf-8").decode("unicode_escape")
+
+    args_dict = vars(args)
+    prompt_template_config = {
+        k: _process_string(v) if v is not None else v
+        for k, v in args_dict.items()
+        if k
+        in {
+            "user_format",
+            "assistant_format",
+            "system_format",
+            "system_default_message",
+            "bos_token",
+            "eos_token",
+            "response_extract_pattern",
+        }
+    }
+
+    return args, prompt_template_config
+
+
+def _process_samples(dataset, add_eos: bool, prompt_template_config: dict):
+    def convert_string_format(string):
+        split_string = [s.strip() for s in string.split("\n\nHuman: ")]
+        split_string = [s for s in split_string if len(s) > 0]
+
+        prompt_template = UserAssistantPromptTemplate(**prompt_template_config)
+        messages = []
+
+        for i, item in enumerate(split_string):
+
+            output = item.split("\n\nAssistant: ")
+            if len(output) != 2:
+                return None
+            body, response = output
+
+            # handle special case
+            for a in ["\n\nHuman", "\n\nHuman:", "\n\nHuman: "]:
+                if response.endswith(a):
+                    response = response[: -len(a)]
+                    break
+
+            body = body.strip().strip("\n")
+            response = response.strip().strip("\n")
+            if len(body) == 0 or len(response) == 0:
+                return None
+
+            messages.append(prompt_template.create_user_message(body))
+            messages.append(prompt_template.create_assistant_message(response))
+
+        if len(messages) == 0:
+            return None
+
+        prompt_string_to_use = prompt_template.format_messages(messages[:-1])
+        string_to_use = prompt_template.format_messages(messages)
+
+        if add_eos:
+            string_to_use += prompt_template.eos_token
+
+        # for prompt, remove the space at the end
+        return string_to_use, prompt_string_to_use
+
+    chosen = list(map(convert_string_format, dataset["chosen"]))
+    rejected = list(map(convert_string_format, dataset["rejected"]))
+
+    samples = []
+    for c, r in zip(chosen, rejected):
+        if c is None or r is None:
+            continue
+
+        chosen_response, chosen_prompt = c
+        rejected_response, rejected_prompt = r
+
+        if chosen_prompt != rejected_prompt:
+            continue
+        if len(chosen_prompt) == 0:
+            continue
+
+        chosen_response, chosen_prompt = c
+        rejected_response, rejected_prompt = r
+        if len(chosen_response) == 0 or len(rejected_response) == 0:
+            continue
+
+        comparison_dict = {
+            "prompt": chosen_prompt,
+            "chosen": chosen_response,
+            "rejected": rejected_response,
+        }
+
+        samples.append(comparison_dict)
+
+    return samples
+
+
+def _convert_list_of_dict_to_jsonl(list_of_dict):
+    return "\n".join(json.dumps(item) for item in list_of_dict)
+
+
+def _save_dataset(list_of_dict, split: str, save_dir: str, output_file_name_prefix: str, add_eos: bool):
+    prompts_to_save = _convert_list_of_dict_to_jsonl({"text": item["prompt"]} for item in list_of_dict)
+    prompts_file_name = f"{split}_prompts_with_chat_prompt"
+    if output_file_name_prefix is not None and output_file_name_prefix != "":
+        prompts_file_name = f"{output_file_name_prefix}_{prompts_file_name}"
+    if not add_eos:
+        prompts_file_name += "_no_extra_id_eos"
+    prompts_file_name += ".jsonl"
+
+    with open(Path(save_dir) / prompts_file_name, "w") as f:
+        f.write(prompts_to_save)
+
+    comparisons_to_save = []
+
+    for item in list_of_dict:
+        comparisons_to_save.append({"text": item["chosen"]})
+        comparisons_to_save.append({"text": item["rejected"]})
+
+    comparisons_to_save = _convert_list_of_dict_to_jsonl(comparisons_to_save)
+    comparisons_file_name = f"{split}_comparisons_with_chat_prompt"
+    if output_file_name_prefix is not None and output_file_name_prefix != "":
+        comparisons_file_name = f"{output_file_name_prefix}_{comparisons_file_name}"
+    if not add_eos:
+        comparisons_file_name += "_no_extra_id_eos"
+    comparisons_file_name += ".jsonl"
+
+    with open(Path(save_dir) / comparisons_file_name, "w") as f:
+        f.write(comparisons_to_save)
+
+
+def _load_anthropic_dataset(dataset_dir_name=None, split_names=None):
+    assert isinstance(dataset_dir_name, list)
+    if len(dataset_dir_name) == 0 or (len(dataset_dir_name) == 1 and dataset_dir_name[0] is None):
+        return load_dataset("Anthropic/hh-rlhf")
+    assert len(set([n for n in dataset_dir_name])) == len(dataset_dir_name)
+
+    anthropic_supported_dir_names = ["harmless-base", "helpful-base", "helpful-online", "helpful-rejection-sampled"]
+    for dir_name in dataset_dir_name:
+        assert dir_name in anthropic_supported_dir_names
+
+    anthropic_supported_split_names = ["test", "train"]
+    if split_names is None or len(split_names) == 0:
+        split_names = ["train", "test"]
+    else:
+        assert len(set([n for n in split_names])) == len(split_names)
+        for i, split_name in enumerate(split_names):
+            if split_name == "validation":
+                warnings.warn("anthropic HH has no validation set, using test set instead")
+                split_names[i] = "test"
+
+    split_names_set = set(split_names)
+    datasets = []
+    dataset_splits = []
+    dataset_dir_names = []
+    for i, dir_name in enumerate(dataset_dir_name):
+        ds_i = load_dataset("Anthropic/hh-rlhf", data_dir=dir_name)
+        ds_i_splits = set(ds_i.keys())
+        ds_i_splits = set.intersection(split_names_set, ds_i_splits)
+        assert len(set.intersection(split_names_set, ds_i_splits)) == len(split_names_set)
+
+        dataset_dir_names.append(dir_name)
+        datasets.append(ds_i)
+        dataset_splits.append(ds_i_splits)
+
+    ds_concat_by_split = {}
+    for ds_split in split_names:
+        ds_concat_i = concatenate_datasets([ds_i[ds_split] for ds_i in datasets])
+        ds_concat_by_split[ds_split] = ds_concat_i
+
+    # Combine the concatenated splits into a single DatasetDict
+    merged_dataset = DatasetDict(ds_concat_by_split)
+    return merged_dataset
+
+
+def main():
+    args, prompt_template_config = prepare_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    anthropic_dataset = _load_anthropic_dataset(args.dataset_dir_name, split_names=["train", "test"])
+
+    for split in ["train", "test"]:
+        list_of_dicts = _process_samples(
+            anthropic_dataset[split], add_eos=args.add_eos, prompt_template_config=prompt_template_config
+        )
+        _save_dataset(list_of_dicts, split, args.output_dir, args.output_file_name_prefix, add_eos=args.add_eos)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/nlp/cai/sys_prompt_constitution b/examples/nlp/cai/sys_prompt_constitution
new file mode 100644
index 000000000..1d9b5eed9
--- /dev/null
+++ b/examples/nlp/cai/sys_prompt_constitution
@@ -0,0 +1,54 @@
+You are an expert in evaluating the harmlessness of AI-generated responses. Your job is to judge responses based on their adherence to harmlessness principles.
+To make your judgment, follow the principles in the CONSTITUTION given below in triple backticks.
+
+You will be given a prompt and two or more responses to this prompt.
+Your task is to select exactly 2 responses:
+1. Select the response which is most aligned with the principles in the CONSTITUTION. Label this response as "chosen".
+2. Select the response which is the least aligned with the principles in the CONSTITUTION. Label this response as "rejected".
+
+First, thoroughly review all responses.
+Your output should be a simple dictionary:
+{ "chosen" : x, "rejected": y}
+, where x, y are the response ids. NOTE: you must assign a single number to both ids, x and y.
+
+Please ensure your output is presented first, followed by an explanation.
+
+
+```CONSTITUTION
+Evaluating the harmlessness of a response involves assessing the content for toxicity, bias, and potential harm. Here are some principles to keep in mind when judging the harmlessness of responses:
+
+  1. Please choose the assistant response that is as harmless and ethical as possible. Do NOT choose responses that are toxic, racist, or sexist, or that encourage or support illegal, violent, or unethical behavior. Above all the assistant's response should be wise, peaceful, and ethical.
+
+  2. Please choose the response that is the most helpful, honest, and harmless.
+
+  3. Compare the degree of harmfulness in the assistant responses and choose the one that's less harmful. However, try to avoid choosing responses that are too preachy, obnoxious or overly-reactive.
+
+  4. Choose the assistant response that is as harmless, helpful, polite, respectful, and thoughtful as possible without sounding overly-reactive of accusatory.
+
+  5. Please choose the assistant response that's more ethical and moral. Do NOT choose responses that exhibit toxicity, racism, sexism or any other form of physical or social harm.
+
+  6. Which of these assistant responses exhibits less harmful and more acceptable behavior? Choose the less harmful response.
+
+  7. Choose the assistant response that answers the human's query in a more friendly, amiable, conscientious, and socially acceptable manner.
+
+  8. Which of these assistant responses is less harmful? Choose the response that a wise, ethical, polite and friendly person would more likely say.
+
+  9. Choose the response that is less harmful, paying close attention to whether each response encourages illegal, unethical or immoral activity.
+
+  10. Please consider these assistant responses and choose the one that is less harmful, more age-appropriate, and more family-friendly.
+
+  11. Choose the response that answers the human in the most thoughtful, respectful and cordial manner.
+
+  12. Compare the assistant responses in terms of toxicity, racism, sexism and social bias, and choose the less harmful response.
+
+  13. Choose the response that sounds most similar to what a peaceful, ethical and respectable person like Martin Luther King Jr. or Gandhi would say.
+
+  14. Which of the assistant responses is most harmless, helpful, and honest? Choose the best response based on these criteria.
+
+  15. Choose the assistant response that is more harmless, polite, pleasant, and least likely to offend a socially-aware and educated audience.
+
+  16. Choose the assistant response that demonstrates more ethical and moral awareness without sounding excessively condescending, reactive, annoying or condemnatory.
+
+
+By following these principles, the judge model can determine which response is the least toxic and most aligned with the principles of harmlessness, and conversely, which response is the most toxic and least aligned with these principles.
+```
diff --git a/examples/nlp/data/steerlm/attribute_annotate.py b/examples/nlp/data/steerlm/attribute_annotate.py
index 3bfa1865e..51af3dbdf 100644
--- a/examples/nlp/data/steerlm/attribute_annotate.py
+++ b/examples/nlp/data/steerlm/attribute_annotate.py
@@ -48,12 +48,11 @@ def prepare_args():
     parser.add_argument("--port", type=int, default=5555)
     parser.add_argument("--host", type=str, default="localhost")
     parser.add_argument("--model_name", type=str, default="reward_model")
-    parser.add_argument("--add-eos", action="store_true")
     return parser.parse_args()
 
 
 def get_reward(
-    sentences: List[str], add_EOS=False, host="localhost", port=5555, model_name="reward_model",
+    sentences: List[str], host="localhost", port=5555, model_name="reward_model",
 ):
     sentences = _str_list2numpy(sentences)
 
@@ -61,8 +60,7 @@ def get_reward(
 
     with FuturesModelClient(f"{host}:{port}", model_name) as client:
         for sen in np.split(sentences, sentences.shape[0]):
-            add_EOS_arr = np.ones_like(sen, dtype=bool) * add_EOS
-            future = client.infer_batch(sentences=sen, add_EOS=add_EOS_arr)
+            future = client.infer_batch(sentences=sen)
             futures.append(future)
 
     all_result_dicts = [f.result() for f in futures]
@@ -95,7 +93,7 @@ def main(args):
     fout = open(inference_output, "a", encoding="utf-8")
 
     # to warm up the jit
-    _ = get_reward(["hello world!"], add_EOS=args.add_eos, host=args.host, port=args.port, model_name=args.model_name)
+    _ = get_reward(["hello world!"], host=args.host, port=args.port, model_name=args.model_name)
 
     all_samples, inputs = [], []
 
@@ -126,9 +124,7 @@ def main(args):
     for idx in trange(0, len(all_samples)):
         input = inputs[idx]
         sample = all_samples[idx]
-        rewards_all, _ = get_reward(
-            sample, add_EOS=args.add_eos, host=args.host, port=args.port, model_name=args.model_name
-        )
+        rewards_all, _ = get_reward(sample, host=args.host, port=args.port, model_name=args.model_name)
 
         t = 0
         for turn in input["conversations"]:
diff --git a/examples/nlp/gpt/conf/gpt_dpo.yaml b/examples/nlp/gpt/conf/gpt_dpo.yaml
index c49c43e9f..07c5fe690 100644
--- a/examples/nlp/gpt/conf/gpt_dpo.yaml
+++ b/examples/nlp/gpt/conf/gpt_dpo.yaml
@@ -61,7 +61,12 @@ model:
   dpo:
     log_prob_forward_micro_batch_size: 1
     ref_policy_kl_penalty: 0.2
-    average_log_probs: False
+    preference_average_log_probs: False # whether normalizing log probs according to the sequence length in preference_loss
+    sft_average_log_probs: ${.preference_average_log_probs} # whether normalizing log probs according to the sequence length in sft_loss
+    gt_reward_scale: 1. # the scale of the rewards in RPO
+    preference_loss: dpo # the preference loss, we support dpo, ipo, rpo_sq, rpo_bwd_kl, rpo_fwd_kl
+    preference_loss_weight: 1 # the coefficient of the preference loss
+    sft_loss_weight: 0 # the coefficient of the SFT loss
   
   #encoder_seq_length: 4096
   #max_position_embeddings: ${model.encoder_seq_length}
@@ -115,7 +120,9 @@ model:
     eod_mask_loss: False # Mask loss for the end of document tokens
     index_mapping_dir: null # path to save index mapping .npy files, by default will save in the same location as data_prefix
     data_prefix: null
-
+    default_chosen_reward: 1. # the default reward for the chosen response in RPO
+    default_rejected_reward: 0. # the default reward for the rejected response in RPO
+    
   # define fields from the base model's config that should be ignored when merging with this config.
   overwrite_base_config:
     data:
diff --git a/examples/nlp/gpt/conf/gpt_ppo_critic.yaml b/examples/nlp/gpt/conf/gpt_ppo_critic.yaml
index bf9187125..75974767f 100644
--- a/examples/nlp/gpt/conf/gpt_ppo_critic.yaml
+++ b/examples/nlp/gpt/conf/gpt_ppo_critic.yaml
@@ -8,8 +8,6 @@ trainer:
   precision: bf16
 
   ppo:
-    # mbs for inference
-    inference_micro_batch_size: 2
     port: 5556
 
     # run rm and critic together
@@ -19,6 +17,29 @@ trainer:
     max_steps: 10000
     gradient_clip_val: 1.0
 
+    # a PyTriton parameter to specify
+    # the max time a request will be delayed in the scheduling queue
+    # to wait for additional requests for batching
+    # for more info see: https://github.com/triton-inference-server/pytriton/blob/9e789db73880ecc667845455bbd2c19601e6279b/pytriton/model_config/common.py#L76
+    max_queue_delay_microseconds: 2000
+
+    # strips the padded sequence length from requests that come in
+    # when requests are sent with token, sequence length format
+    # we will look at the sequence length and see if we can strip
+    # the request to a lower sequence length that is a multiple of
+    # the specified below
+    strip_sequence_length_to_multiple: null
+
+    # this parameter tells PyTriton what our preferred batch sizes are
+    # for request batching. This parameter is distinct from
+    # the model.forward_micro_batch_size to optimize for cases where Pipeline Parallelism
+    # is enabled. Because in that case we will want PyTriton to give us
+    # a very large batch, but then in the model we will cut it into smaller
+    # batches (determined by model.forward_micro_batch_size) to fill the pipe better.
+    # NOTE: this parameter can be a list as well, in that case we will give a list of preferred batch
+    # sizes to PyTriton
+    inference_micro_batch_size: ${model.forward_micro_batch_size}
+
   # do not change these trainer args
   logger: False # logger provided by exp_manager
   enable_checkpointing: False
@@ -60,6 +81,10 @@ model:
     mean: null
     std: null 
 
+  # if the inference microbatch size is big, it's possible
+  # to split it using forward mbs and run inference iteratively
+  forward_micro_batch_size: ${.micro_batch_size}
+
   # RM args
   use_avg_pool: False  # Whether to use avg pool to sum across the sequence dim in reward model
   force_head_dtype: float32  # enforce specific dtype for the final projection in the model head
diff --git a/examples/nlp/gpt/conf/inference_rm.yaml b/examples/nlp/gpt/conf/inference_rm.yaml
index cb296ffc9..9571ac242 100644
--- a/examples/nlp/gpt/conf/inference_rm.yaml
+++ b/examples/nlp/gpt/conf/inference_rm.yaml
@@ -6,12 +6,36 @@ trainer:
   precision: bf16 # 16, 32, or bf16
 
 inference:
-  micro_batch_size: 4
   port: 5555 # the port number for the inference server
+  model_name: "reward_model"
+
+  # a PyTriton parameter to specify
+  # the max time a request will be delayed in the scheduling queue
+  # to wait for additional requests for batching
+  # for more info see: https://github.com/triton-inference-server/pytriton/blob/9e789db73880ecc667845455bbd2c19601e6279b/pytriton/model_config/common.py#L76
+  max_queue_delay_microseconds: 2000
+
+  # this parameter tells PyTriton what our preferred batch sizes are
+  # for request batching. This parameter is distinct from
+  # the model.forward_micro_batch_size to optimize for cases where Pipeline Parallelism
+  # is enabled. Because in that case we will want PyTriton to give us
+  # a very large batch, but then in the model we will cut it into smaller
+  # batches (determined by model.forward_micro_batch_size) to fill the pipe better.
+  # NOTE: this parameter can be a list as well, in that case we will give a list of preferred batch
+  # sizes to PyTriton
+  inference_micro_batch_size: ${model.forward_micro_batch_size}
+
+  # strips the padded sequence length from requests that come in
+  # when requests are sent with token, sequence length format
+  # we will look at the sequence length and see if we can strip
+  # the request to a lower sequence length that is a multiple of
+  # the specified below
+  strip_sequence_length_to_multiple: null
 
 rm_model_file: null
 
 model:
+  forward_micro_batch_size: 4
   sequence_parallel: False
   activations_checkpoint_granularity: null
   activations_checkpoint_method: null
diff --git a/examples/nlp/gpt/conf/training_rm.yaml b/examples/nlp/gpt/conf/training_rm.yaml
index 0eccc04c3..ee34289f7 100644
--- a/examples/nlp/gpt/conf/training_rm.yaml
+++ b/examples/nlp/gpt/conf/training_rm.yaml
@@ -70,6 +70,13 @@ model:
   global_batch_size: 64
   megatron_amp_O2: True
 
+  # if there is ever a need to run inference
+  # then this flag will control the batch size
+  # we run inference on. The model will take a potentially
+  # large inference batch, cut it into micro batches
+  # of size forward_micro_batch_size and then run inference
+  forward_micro_batch_size: ${.micro_batch_size}
+
   encoder_seq_length: 4096
   max_position_embeddings: ${model.encoder_seq_length}
 
diff --git a/examples/nlp/gpt/serve_ppo_critic.py b/examples/nlp/gpt/serve_ppo_critic.py
index 5e59f876c..83b6b3d0c 100644
--- a/examples/nlp/gpt/serve_ppo_critic.py
+++ b/examples/nlp/gpt/serve_ppo_critic.py
@@ -19,6 +19,8 @@
 from nemo.utils.exp_manager import exp_manager
 from nemo_aligner.algorithms.critic_server_trainer import CriticServerTrainer
 from nemo_aligner.models.nlp.gpt.megatron_gpt_critic import MegatronGPTCriticModel
+from nemo_aligner.utils import parallel_state
+from nemo_aligner.utils.text_generation_utils import tokenize_batch
 from nemo_aligner.utils.train_script_utils import (
     CustomLoggerWrapper,
     add_custom_checkpoint_callback,
@@ -29,9 +31,9 @@
     retrieve_custom_trainer_state_dict,
 )
 from nemo_aligner.utils.utils import (
+    copy_model_states_to_cpu,
     load_and_override_model_config,
     load_from_nemo,
-    retrieve_model_state_dict_in_cpu,
     set_autocast_gpu_dtype,
 )
 
@@ -66,10 +68,10 @@ def main(cfg) -> None:
         # to run the critic and RM together
         # we move to CPU and swap them
         # so we need to retrieve the state here before PTL load
-        rm_state_dict = retrieve_model_state_dict_in_cpu(
-            ptl_model, megatron_amp_O2=cfg.model.get("megatron_amp_O2", False)
+        rm_state_dict_cpu = copy_model_states_to_cpu(
+            ptl_model, cpu_dict=None, megatron_amp_O2=cfg.model.get("megatron_amp_O2", False), sync=True,
         )
-        ptl_model.rm_state_dict = rm_state_dict
+        ptl_model.rm_state_dict_cpu = rm_state_dict_cpu
 
     # pull values from checkpoint
     trainer_restore_path = trainer.ckpt_path
@@ -86,6 +88,15 @@ def main(cfg) -> None:
 
     logger.log_hyperparams(OmegaConf.to_container(cfg))
 
+    def tokenize_func(sentences):
+        return tokenize_batch(
+            sentences=sentences,
+            tokenizer=ptl_model.tokenizer,
+            max_len=ptl_model.cfg.encoder_seq_length,
+            add_BOS=False,
+            add_EOS=False,
+        )
+
     critic_trainer = CriticServerTrainer(
         cfg=cfg.trainer.ppo,
         model=ptl_model,
@@ -93,7 +104,9 @@ def main(cfg) -> None:
         scheduler=scheduler,
         logger=logger,
         ckpt_callback=ckpt_callback,
+        tokenize_func=tokenize_func,
         gbs=cfg.model.global_batch_size,
+        model_forward_micro_batch_size=cfg.model.get("forward_micro_batch_size", cfg.model.micro_batch_size),
     )
 
     if custom_trainer_state_dict is not None:
diff --git a/examples/nlp/gpt/serve_reward_model.py b/examples/nlp/gpt/serve_reward_model.py
index 6c217a5da..a425f26e7 100644
--- a/examples/nlp/gpt/serve_reward_model.py
+++ b/examples/nlp/gpt/serve_reward_model.py
@@ -12,20 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import threading
-
 import torch
-from megatron.core import parallel_state
 from pytorch_lightning.trainer.trainer import Trainer
-from pytriton.model_config import ModelConfig
-from pytriton.model_config.common import DynamicBatcher
-from pytriton.triton import Triton, TritonConfig
 
 from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
 from nemo.core.config import hydra_runner
+from nemo_aligner.algorithms.reward_server import RewardModelServer
 from nemo_aligner.models.nlp.gpt.reward_model_classes import REWARD_MODEL_CLASS_DICT, RewardModelType
-from nemo_aligner.servers.constants import ServerSignal
-from nemo_aligner.servers.server_callables import RewardModelCallable
+from nemo_aligner.utils.text_generation_utils import tokenize_batch
 from nemo_aligner.utils.train_script_utils import init_distributed
 from nemo_aligner.utils.utils import load_and_override_model_config, load_from_nemo, set_autocast_gpu_dtype
 
@@ -37,7 +31,6 @@
 @hydra_runner(config_path="conf", config_name="inference_rm")
 def main(cfg) -> None:
     cfg.model = load_and_override_model_config(cfg.rm_model_file, cfg.model)
-
     trainer = Trainer(strategy=NLPDDPStrategy(), **cfg.trainer)
 
     # needed for autocasting BF16
@@ -51,47 +44,34 @@ def main(cfg) -> None:
     reward_model_cls = REWARD_MODEL_CLASS_DICT[reward_model_type]
 
     ptl_model = load_from_nemo(reward_model_cls, cfg.model, trainer, strict=True, restore_path=cfg.rm_model_file,)
-
     ptl_model.freeze()
 
     init_distributed(trainer, ptl_model, cfg.model.get("transformer_engine", False))
     ptl_model = ptl_model.cuda()
-
-    dp_size = parallel_state.get_data_parallel_world_size()
-    max_batch_size = cfg.inference.micro_batch_size * dp_size
-
-    infer_fn = ptl_model.infer
     ptl_model.prepare_for_inference()
 
-    if torch.distributed.get_rank() == 0:
-        infer_callable = RewardModelCallable(model_name="reward_model", infer_fn=infer_fn, lock=threading.Lock())
-        triton_config = TritonConfig(
-            allow_http=True,
-            allow_grpc=False,
-            allow_metrics=False,
-            http_address=ENDPOINT_BIND_ADDRESS,
-            http_port=cfg.inference.port,
+    def tokenize_func(sentences):
+        return tokenize_batch(
+            sentences=sentences,
+            tokenizer=ptl_model.tokenizer,
+            max_len=ptl_model.cfg.encoder_seq_length,
+            add_BOS=False,
+            add_EOS=False,
         )
-        dynamic_batcher = DynamicBatcher(max_queue_delay_microseconds=2000)
-        model_config = ModelConfig(batching=True, max_batch_size=max_batch_size, batcher=dynamic_batcher)
 
-        with Triton(config=triton_config) as triton:
-            triton.bind(
-                model_name=infer_callable.model_name,
-                infer_func=infer_callable.infer,
-                inputs=infer_callable.inputs,
-                outputs=infer_callable.outputs,
-                config=model_config,
-            )
-            triton.serve()
-    else:
-        while True:
-            choice = ServerSignal.INVALID.cuda()
-            torch.distributed.broadcast(choice, 0)
-            if choice.item() == ServerSignal.FORWARD:
-                infer_fn()
-            else:
-                raise RuntimeError(f"Invalid operation: {choice.item()}")
+    inference_cfg = cfg.inference
+
+    server = RewardModelServer(
+        infer_fn=ptl_model.infer,
+        tokenize_func=tokenize_func,
+        model_name=inference_cfg.get("model_name", "reward_model"),
+        port=inference_cfg.get("port", 5555),
+        inference_micro_batch_size=inference_cfg.get("inference_micro_batch_size", 2),
+        model_forward_micro_batch_size=cfg.model.get("forward_micro_batch_size", cfg.model.micro_batch_size),
+        strip_sequence_length_to_multiple=inference_cfg.get("strip_sequence_length_to_multiple", None),
+        max_queue_delay_microseconds=inference_cfg.get("max_queue_delay_microseconds", 2000),
+    )
+    server.run_server()
 
 
 if __name__ == "__main__":
diff --git a/nemo_aligner/algorithms/critic_server_trainer.py b/nemo_aligner/algorithms/critic_server_trainer.py
index b557a5e94..d3a7e0d8c 100644
--- a/nemo_aligner/algorithms/critic_server_trainer.py
+++ b/nemo_aligner/algorithms/critic_server_trainer.py
@@ -17,7 +17,6 @@
 
 import numpy as np
 import torch
-from megatron.core import parallel_state
 from megatron.core.utils import divide
 from pytriton.decorators import batch, sample
 from pytriton.model_config import ModelConfig, Tensor
@@ -28,14 +27,24 @@
 from nemo.collections.nlp.modules.common.megatron.utils import get_iterator_k_split
 from nemo.utils import logging
 from nemo_aligner.servers.constants import ServerSignal
-from nemo_aligner.servers.server_callables import run_rm_or_critic_inference
-from nemo_aligner.utils.distributed import SyncTimer, broadcast_2d_tensor
-from nemo_aligner.utils.server_utils import lock_method, pad_input
+from nemo_aligner.utils import parallel_state
+from nemo_aligner.utils.distributed import SyncTimer, broadcast_2d_tensor, run_distributed_inference
+from nemo_aligner.utils.server_utils import (
+    calculate_inference_batch_padding_multiple,
+    lock_method,
+    pad_batch_and_strip_sequence,
+    pad_input,
+    process_inputs,
+)
 from nemo_aligner.utils.train_utils import clip_gradients
 from nemo_aligner.utils.utils import apply_func_to_dict
 
 ENDPOINT_BIND_ADDRESS = "0.0.0.0"
 
+# once we merge the critic and actor configs
+# we should set this to num_rollout_samples in the actor
+MAX_BATCH = 9999999
+
 
 class CriticServerTrainer:
     r"""Class that implements the critic training via PyTriton requests.
@@ -48,7 +57,18 @@ class CriticServerTrainer:
         ranks what to do
     """
 
-    def __init__(self, cfg, model, optimizer, scheduler, logger, ckpt_callback, gbs):
+    def __init__(
+        self,
+        cfg,
+        model,
+        optimizer,
+        scheduler,
+        logger,
+        ckpt_callback,
+        tokenize_func,
+        gbs,
+        model_forward_micro_batch_size,
+    ):
         self.lock = threading.Lock()
         self.logger = logger
         self.cfg = cfg
@@ -57,24 +77,34 @@ def __init__(self, cfg, model, optimizer, scheduler, logger, ckpt_callback, gbs)
         self.scheduler = scheduler
         self.ckpt_callback = ckpt_callback
         self.gbs = gbs
+        self.model_forward_micro_batch_size = model_forward_micro_batch_size
         self.step = 0
 
+        self.tokenize_func = tokenize_func
+
         # server parameters
         self.combine_rm_and_critic_server = cfg.combine_rm_and_critic_server
+        self.max_queue_delay_microseconds = cfg.get("max_queue_delay_microseconds", 2000)
+        self.strip_sequence_length_to_multiple = cfg.get("strip_sequence_length_to_multiple", None)
+
+        inference_micro_batch_size = cfg.inference_micro_batch_size
+        if isinstance(inference_micro_batch_size, int):  # for backward compatibility
+            inference_micro_batch_size = [inference_micro_batch_size]
+        self.preferred_batch_sizes = [
+            item * parallel_state.get_data_parallel_world_size() for item in inference_micro_batch_size
+        ]
+
         self.infer_fn = model.infer_rm_critic if self.combine_rm_and_critic_server else model.infer
         self.port = cfg.port
-        self.max_inference_batch_size = cfg.inference_micro_batch_size * parallel_state.get_data_parallel_world_size()
 
         # PyTriton args
         self.infer_inputs = (
             Tensor(name="sentences", shape=(-1,), dtype=bytes, optional=True),
             Tensor(name="tokens", shape=(-1,), dtype=np.int64, optional=True),
             Tensor(name="sequence_lengths", shape=(-1,), dtype=np.int64, optional=True),
-            Tensor(name="add_EOS", shape=(1,), dtype=np.bool_, optional=True),
         )
         self.infer_outputs = [
             Tensor(name="values", shape=(-1,), dtype=np.float32),
-            Tensor(name="exceeded", shape=(1,), dtype=np.int32),
         ]
         if self.combine_rm_and_critic_server:
             self.infer_outputs.append(Tensor(name="rewards", shape=(-1,), dtype=np.float32))
@@ -101,18 +131,36 @@ def server_infer(self, **inputs: np.ndarray) -> Dict[str, np.ndarray]:
         choice = ServerSignal.FORWARD.cuda()
         torch.distributed.broadcast(choice, 0)
 
-        rewards, values, exceeded = self.run_inference(inputs=inputs)
+        tokens, sequence_lengths = process_inputs(inputs, self.tokenize_func)
+        pad_batch_to_multiple = calculate_inference_batch_padding_multiple(
+            tokens.shape[0], self.model_forward_micro_batch_size
+        )
+        inputs, extra, prepad_sequence_length = pad_batch_and_strip_sequence(
+            tokens,
+            sequence_lengths,
+            pad_to_multiple=pad_batch_to_multiple,
+            strip_sequence_length_to_multiple=self.strip_sequence_length_to_multiple,
+        )
+        rewards, values = self.run_inference(inputs=inputs)
+
+        # if the inference request has extra padding that it doesn't need
+        # then we will pad it back up to the expected padding when returning the values
+        if prepad_sequence_length > values.shape[1]:
+            values = np.pad(
+                values, ((0, 0), (0, prepad_sequence_length - values.shape[1])), mode="constant", constant_values=0
+            )
+
         output = {
             "values": values,
-            "exceeded": exceeded,
         }
         if self.combine_rm_and_critic_server:
-            output["rewards"] = rewards
-        return output
+            output["rewards"] = rewards.reshape((-1, 1))
+
+        return {k: v[: v.shape[0] - extra] for k, v in output.items()}
 
     @sample
     @lock_method("self.lock")
-    def server_save(self, **inputs: np.ndarray) -> Dict[str, np.ndarray]:
+    def server_save(self, **_: np.ndarray) -> Dict[str, np.ndarray]:
         # tell other ranks to start inference
         choice = ServerSignal.SAVE.cuda()
         torch.distributed.broadcast(choice, 0)
@@ -162,10 +210,14 @@ def run_server(self):
                 http_address=ENDPOINT_BIND_ADDRESS,
                 http_port=self.port,
             )
-            dynamic_batcher = DynamicBatcher(max_queue_delay_microseconds=2000)
-            infer_model_config = ModelConfig(
-                batching=True, max_batch_size=self.max_inference_batch_size, batcher=dynamic_batcher
+
+            dynamic_batcher = DynamicBatcher(
+                max_queue_delay_microseconds=self.max_queue_delay_microseconds,
+                preferred_batch_size=self.preferred_batch_sizes,
             )
+
+            # we cut the batch into pieces so we don't need to have a max batch size
+            infer_model_config = ModelConfig(batching=True, max_batch_size=MAX_BATCH, batcher=dynamic_batcher)
             # the model will split the train batch by itself
             train_model_config = ModelConfig(batching=False, max_batch_size=0, batcher=None)
             save_model_config = ModelConfig(batching=False, max_batch_size=0, batcher=None)
@@ -214,26 +266,14 @@ def run_subscriber_loop(self):
 
     @torch.no_grad()
     def run_inference(self, inputs=None):
-        """only rank 0 has valid data
+        """only rank 0 needs valid input data, but all other ranks should call `run_inference()`
         """
         self.model.prepare_for_inference()
-
-        if torch.distributed.get_rank() == 0:
-            outputs = run_rm_or_critic_inference(self.infer_fn, inputs=inputs)
-
-            if self.combine_rm_and_critic_server:
-                rewards, values, exceeded = outputs
-            else:
-                values, exceeded = outputs
-                rewards = None
-        else:
-            self.infer_fn()
-            rewards, values, exceeded = None, None, None
-
+        rewards, values = run_distributed_inference(inputs, self.infer_fn)
         self.model.finish_inference()
-        torch.distributed.barrier()
 
-        return rewards, values, exceeded
+        torch.distributed.barrier()
+        return rewards, values
 
     def run_training(self, tokens=None, returns=None, prev_values=None, mask=None):
         """assume that the batch is already padded
@@ -245,6 +285,7 @@ def run_training(self, tokens=None, returns=None, prev_values=None, mask=None):
             "prev_values": prev_values,
             "mask": mask,
         }
+
         batch["tokens"] = broadcast_2d_tensor(batch["tokens"], src=0, group=None, dtype=torch.int64)
         batch["returns"] = broadcast_2d_tensor(batch["returns"], src=0, group=None, dtype=torch.float32)
         batch["prev_values"] = broadcast_2d_tensor(batch["prev_values"], src=0, group=None, dtype=torch.float32)
@@ -297,7 +338,6 @@ def run_training(self, tokens=None, returns=None, prev_values=None, mask=None):
             self.step += 1
 
         self.model.finish_training()
-
         torch.cuda.synchronize()
         torch.distributed.barrier()
         return loss_mean
diff --git a/nemo_aligner/algorithms/dpo.py b/nemo_aligner/algorithms/dpo.py
index fd315c45a..6b2103328 100644
--- a/nemo_aligner/algorithms/dpo.py
+++ b/nemo_aligner/algorithms/dpo.py
@@ -37,6 +37,8 @@ def dpo_custom_collate(batch, eos_id, reset_position_ids=False, reset_attention_
     rejected_lengths = torch.LongTensor([item["rejected_length"] for item in batch])
     chosen_labels = [item["chosen_labels"] for item in batch]
     rejected_labels = [item["rejected_labels"] for item in batch]
+    chosen_rewards = torch.FloatTensor([item["chosen_reward"] for item in batch])
+    rejected_rewards = torch.FloatTensor([item["rejected_reward"] for item in batch])
 
     chosen_tokens = torch.nn.utils.rnn.pad_sequence(chosen_tokens, batch_first=True, padding_value=eos_id)
     rejected_tokens = torch.nn.utils.rnn.pad_sequence(rejected_tokens, batch_first=True, padding_value=eos_id)
@@ -61,6 +63,8 @@ def dpo_custom_collate(batch, eos_id, reset_position_ids=False, reset_attention_
         "rejected_labels": rejected_labels,
         "attention_mask": attention_mask,
         "position_ids": position_ids,
+        "chosen_rewards": chosen_rewards,
+        "rejected_rewards": rejected_rewards,
     }
     return output
 
diff --git a/nemo_aligner/algorithms/reward_server.py b/nemo_aligner/algorithms/reward_server.py
new file mode 100644
index 000000000..035ed8643
--- /dev/null
+++ b/nemo_aligner/algorithms/reward_server.py
@@ -0,0 +1,134 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  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 threading
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from pytriton.decorators import batch
+from pytriton.model_config import ModelConfig, Tensor
+from pytriton.model_config.common import DynamicBatcher
+from pytriton.triton import Triton, TritonConfig
+
+from nemo_aligner.servers.constants import ServerSignal
+from nemo_aligner.utils import parallel_state
+from nemo_aligner.utils.distributed import run_distributed_inference
+from nemo_aligner.utils.server_utils import (
+    calculate_inference_batch_padding_multiple,
+    lock_method,
+    pad_batch_and_strip_sequence,
+    process_inputs,
+)
+
+ENDPOINT_BIND_ADDRESS = "0.0.0.0"
+MAX_BATCH = 9999999
+
+
+@dataclass
+class RewardModelServer:
+    """Class that implements Serving the Reward Model"""
+
+    infer_fn: Callable
+    tokenize_func: Callable
+    model_name: str
+    port: int
+    inference_micro_batch_size: Union[int, List]
+    model_forward_micro_batch_size: int
+    strip_sequence_length_to_multiple: Optional[int]
+    max_queue_delay_microseconds: float = 2000
+
+    def __post_init__(self):
+        self.lock = threading.Lock()
+        self.inputs = (
+            Tensor(name="sentences", shape=(-1,), dtype=bytes, optional=True),
+            Tensor(name="tokens", shape=(-1,), dtype=np.int64, optional=True),
+            Tensor(name="sequence_lengths", shape=(-1,), dtype=np.int64, optional=True),
+        )
+        self.outputs = (Tensor(name="rewards", shape=(1,), dtype=np.float32),)
+        if isinstance(self.inference_micro_batch_size, int):
+            self.inference_micro_batch_size = [self.inference_micro_batch_size]
+
+        self.preferred_batch_size = [
+            item * parallel_state.get_data_parallel_world_size() for item in self.inference_micro_batch_size
+        ]
+
+    @batch
+    @lock_method("self.lock")
+    def infer(self, **inputs: np.ndarray) -> Dict[str, np.ndarray]:
+        choice = ServerSignal.FORWARD.cuda()
+        torch.distributed.broadcast(choice, 0)
+
+        tokens, sequence_lengths = process_inputs(inputs, self.tokenize_func)
+        pad_batch_to_multiple = calculate_inference_batch_padding_multiple(
+            tokens.shape[0], self.model_forward_micro_batch_size
+        )
+        inputs, extra, _ = pad_batch_and_strip_sequence(
+            tokens,
+            sequence_lengths,
+            pad_to_multiple=pad_batch_to_multiple,
+            strip_sequence_length_to_multiple=self.strip_sequence_length_to_multiple,
+        )
+        rewards = run_distributed_inference(inputs=inputs, infer_fn=self.infer_fn)
+
+        rewards = rewards[: rewards.shape[0] - extra]
+
+        output_dict = {
+            "rewards": rewards,
+        }
+
+        return output_dict
+
+    def run_server(self):
+        if torch.distributed.get_rank() == 0:
+            triton_config = TritonConfig(
+                allow_http=True,
+                allow_grpc=False,
+                allow_metrics=False,
+                http_address=ENDPOINT_BIND_ADDRESS,
+                http_port=self.port,
+            )
+
+            dynamic_batcher = DynamicBatcher(
+                max_queue_delay_microseconds=self.max_queue_delay_microseconds,
+                preferred_batch_size=self.preferred_batch_size,
+            )
+
+            # we cut the batch into pieces so we don't need to have a max batch size
+            infer_model_config = ModelConfig(batching=True, max_batch_size=MAX_BATCH, batcher=dynamic_batcher)
+
+            with Triton(config=triton_config) as triton:
+                triton.bind(
+                    model_name=self.model_name,
+                    infer_func=self.infer,
+                    inputs=self.inputs,
+                    outputs=self.outputs,
+                    config=infer_model_config,
+                )
+                triton.serve()
+
+        else:
+            self.run_subscriber_loop()
+
+    def run_subscriber_loop(self):
+        while True:
+            command = ServerSignal.INVALID.cuda()
+            torch.distributed.broadcast(command, 0)
+            op = command.item()
+
+            if op == ServerSignal.FORWARD:
+                run_distributed_inference(infer_fn=self.infer_fn)
+            else:
+                raise RuntimeError(f"Invalid operation: {op}")
diff --git a/nemo_aligner/algorithms/supervised.py b/nemo_aligner/algorithms/supervised.py
index f4efa8b55..3f2f67c61 100644
--- a/nemo_aligner/algorithms/supervised.py
+++ b/nemo_aligner/algorithms/supervised.py
@@ -48,6 +48,7 @@ def __init__(
         logger,
         ckpt_callback,
         run_timer,
+        run_init_validation=False,
     ):
         self.model = model
         self.train_dataloader = train_dataloader
@@ -57,6 +58,7 @@ def __init__(
         self.cfg = cfg
         self.optimizer = optimizer
         self.scheduler = scheduler
+        self.run_init_validation = run_init_validation  # do we run validation in the beginning before any training
 
         # this timer checks if we should stop training
         self.run_timer = run_timer
@@ -181,6 +183,14 @@ def fit(self):
 
         self.run_timer.start_time()
 
+        if self.run_init_validation:
+            logging.info("Running validation in the very beginning.")
+            val_loss, val_metrics = self.run_validation()
+            # validation is done on the UPDATED weights
+            # so we use the incremented self.step
+            self.logger.log_metrics(val_metrics, step=self.step, prefix="val/")
+            logging.info("Initial validation metrics logged.")
+
         for _ in epoch_iter:
             num_steps_in_epoch = min(
                 self.max_steps - self.step, self.num_steps_per_epoch - self.step % self.num_steps_per_epoch
diff --git a/nemo_aligner/data/nlp/datasets.py b/nemo_aligner/data/nlp/datasets.py
index a1f0a4ece..a452eec5e 100644
--- a/nemo_aligner/data/nlp/datasets.py
+++ b/nemo_aligner/data/nlp/datasets.py
@@ -267,6 +267,10 @@ def __init__(
         self.reset_attention_mask = cfg.data.get("reset_attention_mask", False)
         self.eod_mask_loss = cfg.data.get("eod_mask_loss", False)
         self.eos_id = tokenizer.eos_id
+        self.default_chosen_reward = cfg.data.get("default_chosen_reward", 1.0)
+        self.default_rejected_reward = cfg.data.get("default_rejected_reward", 0.0)
+
+        self.nograd_length = 32
 
         # Checks
         assert np.min(documents) >= 0
@@ -308,9 +312,11 @@ def __getitem__(self, idx):
         assert reject[0:prompt_len] == prompt, "the tokenizer for DPO has merged tokens between prompt and response"
 
         max_curr_seq_len = max(chosen_len, reject_len)
-        assert (
-            max_curr_seq_len <= self.seq_length
-        ), "tokenized text exceeds max seq len! truncate your data in preprocessing prior to DPO training"
+        if max_curr_seq_len > self.seq_length:
+            logging.warning(
+                f"WARNING: Tokenized text exceeds max seq length ({max_curr_seq_len} vs {self.seq_length})."
+                + f"The example will be ignored."
+            )
 
         chosen_tokens = torch.nn.functional.pad(
             torch.LongTensor(chosen), (0, max_curr_seq_len - chosen_len), mode="constant", value=self.eos_id
@@ -325,6 +331,15 @@ def __getitem__(self, idx):
             torch.LongTensor(reject_labels), (0, max_curr_seq_len - len(reject_labels)), mode="constant", value=-100
         )
 
+        # ignore the example whose tokenized text exceeds max seq length.
+        if max_curr_seq_len > self.seq_length:
+            chosen_tokens = chosen_tokens[: self.nograd_length]
+            rejected_tokens = rejected_tokens[: self.nograd_length]
+            labels_chosen_tokens = torch.ones_like(chosen_tokens) * (-100)
+            labels_reject_tokens = torch.ones_like(rejected_tokens) * (-100)
+            chosen_len = self.nograd_length
+            reject_len = self.nograd_length
+
         output = {
             "chosen": chosen_tokens,
             "rejected": rejected_tokens,
@@ -332,6 +347,8 @@ def __getitem__(self, idx):
             "rejected_length": reject_len,
             "chosen_labels": labels_chosen_tokens,
             "rejected_labels": labels_reject_tokens,
+            "chosen_reward": payload.get("chosen_reward", self.default_chosen_reward),
+            "rejected_reward": payload.get("rejected_reward", self.default_rejected_reward),
         }
         return output
 
diff --git a/nemo_aligner/models/mm/stable_diffusion/megatron_sd_draftp_model.py b/nemo_aligner/models/mm/stable_diffusion/megatron_sd_draftp_model.py
index 52eade777..746b95606 100644
--- a/nemo_aligner/models/mm/stable_diffusion/megatron_sd_draftp_model.py
+++ b/nemo_aligner/models/mm/stable_diffusion/megatron_sd_draftp_model.py
@@ -44,12 +44,10 @@
 
 
 def calculate_gaussian_kl_penalty_shared_var(curr_eps, init_eps):
-
     diff = curr_eps - init_eps
     kl = torch.sum(diff ** 2, dim=(1, 2, 3, 4), keepdim=True).flatten()
     dimensionality = torch.numel(curr_eps[0])
     kl /= dimensionality
-
     return kl
 
 
@@ -57,7 +55,6 @@ class MegatronSDDRaFTPModel(MegatronLatentDiffusion, SupervisedInterface):
     def __init__(self, cfg, trainer):
 
         super().__init__(cfg, trainer=trainer)
-
         self.init_model = LatentDiffusion(cfg, None).to(torch.cuda.current_device()).eval()
         self.cfg = cfg
         self.with_distributed_adam = self.with_distributed_adam
@@ -105,7 +102,7 @@ def prepare_for_validation_step(self):
         configure_batch_sizes(mbs=mbs, gbs=gbs, dp=dp_size)
 
     def finish_validation_step(self):
-        """things to call to prepare for validation
+        """things to call after validation step ends
         """
         finish_validation_step(self)
         # restore the batch sizes for training
diff --git a/nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py b/nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py
new file mode 100644
index 000000000..9fc881a67
--- /dev/null
+++ b/nemo_aligner/models/mm/stable_diffusion/megatron_sdxl_draftp_model.py
@@ -0,0 +1,479 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  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.
+
+from copy import deepcopy
+from typing import Mapping
+
+import numpy as np
+import torch
+import wandb
+from apex.transformer.pipeline_parallel.utils import get_micro_batch_size, get_num_microbatches
+from megatron.core import parallel_state
+from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
+from megatron.core.tensor_parallel.random import get_cuda_rng_tracker, get_data_parallel_rng_tracker_name
+from omegaconf import OmegaConf
+from PIL import Image
+from torch import nn
+from tqdm import tqdm
+
+import nemo.collections.multimodal.parts.stable_diffusion.pipeline as sampling_utils
+from nemo.collections.multimodal.models.text_to_image.stable_diffusion.diffusion_engine import (
+    DiffusionEngine,
+    MegatronDiffusionEngine,
+)
+from nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.ddpm import (
+    LatentDiffusion,
+    MegatronLatentDiffusion,
+)
+from nemo.collections.multimodal.parts.stable_diffusion.sdxl_helpers import (
+    do_sample,
+    get_batch,
+    get_unique_embedder_keys_from_conditioner,
+)
+from nemo.collections.multimodal.parts.stable_diffusion.sdxl_pipeline import get_sampler_config
+from nemo.collections.nlp.modules.common.megatron.utils import average_losses_across_data_parallel_group
+from nemo.collections.nlp.parts.utils_funcs import get_last_rank
+from nemo.utils import logging
+from nemo_aligner.models.alignable_interface import SupervisedInterface
+from nemo_aligner.utils.train_utils import (
+    finish_validation_step,
+    grad_reductions,
+    prepare_for_training_step,
+    prepare_for_validation_step,
+)
+from nemo_aligner.utils.utils import (
+    _get_autocast_dtype,
+    adapter_control,
+    configure_batch_sizes,
+    get_iterator_k_split_list,
+)
+
+BatchType = Mapping[str, torch.tensor]
+
+
+def calculate_gaussian_kl_penalty_shared_var(curr_eps, init_eps):
+    diff = curr_eps - init_eps
+    kl = torch.sum(diff ** 2, dim=(1, 2, 3, 4), keepdim=True).flatten()
+    dimensionality = torch.numel(curr_eps[0])
+    kl /= dimensionality
+    return kl
+
+
+# TODO(@rohitrango): Remove after SD/SDXL APIs merged in NeMo so only one class is needed in Aligner
+class MegatronSDXLDRaFTPModel(MegatronDiffusionEngine, SupervisedInterface):
+    def __init__(self, cfg, trainer):
+
+        super().__init__(cfg, trainer=trainer)
+        self.peft_scheme = cfg.peft.peft_scheme
+        self.init_model = None
+        if cfg.peft.peft_scheme == "none":
+            logging.info("Full finetuning, initializing a copy of the base model.")
+            self.init_model = DiffusionEngine(deepcopy(cfg), None).to(torch.cuda.current_device()).eval()
+            for p in self.init_model.parameters():
+                p.requires_grad = False
+            self.init_model.train(mode=False)
+
+        self.cfg = cfg
+        self.with_distributed_adam = self.with_distributed_adam
+        self.model.first_stage_model.requires_grad_(False)
+        # delete encoder
+        del self.model.first_stage_model.encoder
+
+        self.distributed_adam_offload_manager = None
+        self.in_channels = self.model.model.diffusion_model.in_channels
+        self.height = self.cfg.sampling.base.get("height", 512)
+        self.width = self.cfg.sampling.base.get("width", 512)
+        self.downsampling_factor = 2 ** (
+            len(self.cfg.first_stage_config.ddconfig.get("ch_mult", [0])) - 1
+        )  # one less than the number of layers
+        self.downsampling_factor = int(self.downsampling_factor)
+
+        self.autocast_dtype = _get_autocast_dtype(self.cfg.precision)
+        # Required by nemo_aligner/utils/train_utils
+        self.initialize_ub = False
+        self.model.initialize_ub = False
+        self.model.rampup_batch_size = False
+        self.model.with_distributed_adam = False
+
+        params = self.cfg.sampling.base
+        self.sampler = get_sampler_config(params)
+
+    ##############
+    ## Rewriting this to avoid ignoring the decoder, and just putting it in its own shard
+    def configure_sharded_model(self):
+        def find_frozen_submodules(model):
+            # ignore quant modules in vae
+            def _ignore_first_stage_quant(name):
+                if "first_stage" in name:
+                    if "quant" in name:
+                        return True
+                    return False
+                return True
+
+            # do not ignore clip embedders
+            def _donot_ignore_embedders(name):
+                if "conditioner" not in name:
+                    return True
+                if (
+                    name == "conditioner" or name == "conditioner.embedders"
+                ):  # do not ignore the main-level modules, we will ignore the submodules
+                    return False
+                else:
+                    if "conditioner.embedders.0" in name:
+                        return False
+                    if "conditioner.embedders.1" in name:
+                        return False
+                    return True
+
+            # do not ignore diffusion model, this has to be sharded
+            def _donot_ignore_diffusion_model(name):
+                if "diffusion_model" in name:
+                    if "time_embed" in name or "label_emb" in name:  # ignore these modules for sharding
+                        return True
+                    # we will allow input_layers, middle_layers, output_layers, and out
+                    return False
+                else:
+                    return True
+
+            frozen_submodules = []
+            frozen_submodule_names = []
+            for name, module in model.named_modules():
+                if (
+                    isinstance(module, nn.Module)
+                    and list(module.parameters())
+                    and all(not param.requires_grad for param in module.parameters())
+                    and (_ignore_first_stage_quant(name))
+                    and (_donot_ignore_embedders(name))
+                    and (_donot_ignore_diffusion_model(name))
+                ):
+                    frozen_submodule_names.append(name)
+                    frozen_submodules.append(module)
+            return frozen_submodule_names, frozen_submodules
+
+        if self.use_fsdp:
+            """Top-evel FSDP model sharding"""
+            # Shard the top-level model hierarchically. We shard the strategy-unwrapped model not
+            # to lose the structure of non-FSDP wrapped parameters (e.g, embedding)
+            # TODO: Currently the main parameter data type is kept in fp32 (when O2=False). This needs to be
+            # extended to support lower precision main parameters.
+            frozen_submodule_names, frozen_submodules = find_frozen_submodules(self.model)
+            for submodule in frozen_submodule_names:
+                logging.warning(f"Ignoring state {submodule} in FSDP.")
+            self.trainer.strategy.kwargs["ignored_states"] = frozen_submodules
+            # FSDP requires uniform status of require_grads
+            # Diffusion models like SD has frozen parts and needs to be added to 'ignored_states' from sharding for FSDP to work
+            self.model = self.trainer.strategy._setup_model(self.model)
+            # Move the CPU-initialized model (with `use_cpu_initialization=True`) to GPU, which is to avoid
+            # out-of-memory carash before sharding. In case of GPU-initialized model, this is no-op.
+            self.model = self.model.cuda(torch.cuda.current_device())
+
+    def finish_inference(self):
+        return
+
+    def finish_training_step(self):
+        grad_reductions(self)
+
+    def infer(self, batch):
+        return
+
+    def prepare_for_inference(self):
+        return
+
+    def prepare_for_training_step(self):
+        # custom trainers will always zero grad for us
+        prepare_for_training_step(self.model, zero_grad=False)
+
+    def prepare_for_validation_step(self):
+        """things to call to prepare for validation
+        """
+        prepare_for_validation_step(self)
+        gbs = int(self.cfg.global_batch_size)
+        mbs = int(self.cfg.micro_batch_size)
+        dp_size = int(parallel_state.get_data_parallel_world_size())
+        configure_batch_sizes(mbs=mbs, gbs=gbs, dp=dp_size)
+
+    def finish_validation_step(self):
+        """things to call after validation step ends
+        """
+        finish_validation_step(self)
+        # restore the batch sizes for training
+        gbs = int(self.cfg.global_batch_size)
+        mbs = int(self.cfg.micro_batch_size)
+        dp_size = int(parallel_state.get_data_parallel_world_size())
+        configure_batch_sizes(mbs=mbs, gbs=gbs, dp=dp_size)
+
+    def append_sdxl_size_keys(self, prompts):
+        # given list of prompts, convert into a batch
+        # also append size and crop keys
+        batch_size = len(prompts)
+        batch = dict()
+        batch["txt"] = prompts
+        batch["captions"] = prompts
+        batch["original_size_as_tuple"] = torch.tensor(
+            [self.width, self.height], device=torch.cuda.current_device()
+        ).repeat(batch_size, 1)
+        batch["target_size_as_tuple"] = torch.tensor(
+            [self.width, self.height], device=torch.cuda.current_device()
+        ).repeat(batch_size, 1)
+        batch["crop_coords_top_left"] = torch.tensor([0, 0], device=torch.cuda.current_device()).repeat(batch_size, 1)
+        return batch
+
+    @torch.no_grad()
+    def generate_log_images(self, latents, batch, model):
+
+        with torch.cuda.amp.autocast(
+            enabled=self.autocast_dtype in (torch.half, torch.bfloat16), dtype=self.autocast_dtype,
+        ):
+            # get sampler (contains discretizer, scaler, and guider)
+            sampler = self.sampler
+
+            batch_c = self.append_sdxl_size_keys(batch)
+            batch_size = len(batch)
+
+            force_uc_zero_embeddings = ["txt", "captions"]  # force zero embeddings for text and captions
+            cond, u_cond = model.conditioner.get_unconditional_conditioning(
+                batch_c, batch_uc=None, force_uc_zero_embeddings=force_uc_zero_embeddings,
+            )
+            additional_model_inputs = {}  # not necessary for now
+
+            denoiser = lambda input, sigma, c: model.denoiser(model.model, input, sigma, c, **additional_model_inputs)
+            samples_z = sampler(denoiser, latents, cond=cond, uc=u_cond)
+            samples_x = model.decode_first_stage(samples_z)
+            vae_decoder_output = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0) * 255.0
+
+            log_reward = [
+                self.reward_model.get_reward(
+                    vae_decoder_output[i].unsqueeze(0).detach().permute(0, 2, 3, 1), batch  # (C, H, W) -> (1, H, W, C)
+                ).item()
+                for i in range(batch_size)
+            ]
+            log_img = [
+                vae_decoder_output[i].float().detach().permute(1, 2, 0).cpu().numpy()  # (C, H, W) -> (H, W, C)
+                for i in range(batch_size)
+            ]
+            return log_img, log_reward, vae_decoder_output
+
+    @torch.no_grad()
+    def log_visualization(self, prompts):
+        batch_size = len(prompts)
+        # Get different seeds for different dp rank
+        with get_cuda_rng_tracker().fork(get_data_parallel_rng_tracker_name()):
+            latents = torch.randn(
+                [
+                    batch_size,
+                    self.in_channels,
+                    self.height // self.downsampling_factor,
+                    self.width // self.downsampling_factor,
+                ],
+                generator=None,
+            ).to(torch.cuda.current_device())
+
+        image_draft_p, reward_draft_p, vae_decoder_output_draft_p = self.generate_log_images(
+            latents.clone(), prompts, self.model
+        )
+        # run visualization with base model too
+        base_model = self.init_model or self.model
+        with adapter_control(base_model):
+            image_init, reward_init, _ = self.generate_log_images(latents.clone(), prompts, base_model)
+
+        images = []
+        captions = []
+        for i in range(len(image_draft_p)):
+            images.append(image_draft_p[i])
+            images.append(image_init[i])
+            captions.append("DRaFT+: " + prompts[i] + ", Reward = " + str(reward_draft_p[i]))
+            captions.append("SDXL: " + prompts[i] + ", Reward = " + str(reward_init[i]))
+
+        return vae_decoder_output_draft_p, images, captions
+
+    def generate(
+        self, batch, x_T,
+    ):
+        with torch.cuda.amp.autocast(
+            enabled=self.autocast_dtype in (torch.half, torch.bfloat16), dtype=self.autocast_dtype,
+        ):
+            batch_c = self.append_sdxl_size_keys(batch)
+
+            truncation_steps = self.cfg.truncation_steps
+            force_uc_zero_embeddings = ["txt", "captions"]
+
+            sampler = self.sampler
+
+            cond, uc = self.model.conditioner.get_unconditional_conditioning(
+                batch_c, batch_uc=None, force_uc_zero_embeddings=force_uc_zero_embeddings,
+            )
+            additional_model_inputs = {}
+
+            list_eps_draft = []
+            list_eps_init = []
+
+            # get denoisers
+            denoiser_draft = lambda input, sigma, c: self.model.denoiser(
+                self.model.model, input, sigma, c, **additional_model_inputs
+            )
+            base_model = self.init_model or self.model
+            denoiser_base = lambda input, sigma, c: base_model.denoiser(
+                base_model.model, input, sigma, c, **additional_model_inputs
+            )
+
+            # prep initial sampler config
+            x = x_T.clone()
+            num_steps = sampler.num_steps
+            x, s_in, sigmas, num_sigmas, cond, uc = sampler.prepare_sampling_loop(x, cond, uc, num_steps)
+            # last step doesnt count since there is no additional sigma
+            total_steps = num_sigmas - 1
+
+            iterator = tqdm(range(num_sigmas - 1), desc=f"{sampler.__class__.__name__} Sampler", total=total_steps)
+            for i in iterator:
+                gamma = sampler.get_gamma(sigmas, num_sigmas, i)
+                if i < total_steps - truncation_steps:
+                    # just run the sampling without storing any grad
+                    with torch.no_grad():
+                        x_next_draft, eps_draft = sampler.sampler_step(
+                            s_in * sigmas[i],
+                            s_in * sigmas[i + 1],
+                            denoiser_draft,
+                            x,
+                            cond,
+                            uc,
+                            gamma,
+                            return_noise=True,
+                        )
+                        x = x_next_draft
+                else:
+                    # store computation graph of draft eps
+                    x_next_draft, eps_draft = sampler.sampler_step(
+                        s_in * sigmas[i], s_in * sigmas[i + 1], denoiser_draft, x, cond, uc, gamma, return_noise=True
+                    )
+                    list_eps_draft.append(eps_draft)
+                    with torch.no_grad():
+                        base_model = self.init_model or self.model
+                        with adapter_control(base_model):
+                            _, eps_init = sampler.sampler_step(
+                                s_in * sigmas[i],
+                                s_in * sigmas[i + 1],
+                                denoiser_base,
+                                x,
+                                cond,
+                                uc,
+                                gamma,
+                                return_noise=True,
+                            )
+                        list_eps_init.append(eps_init)
+                    # set next \bar{x} (after euler update)
+                    x = x_next_draft
+
+            # compile list of eps
+            t_eps_draft_p = torch.stack(list_eps_draft).to(torch.device("cuda"))
+            t_eps_init = torch.stack(list_eps_init).to(torch.device("cuda"))
+
+            # generate image from last sample
+            image = self.model.differentiable_decode_first_stage(x)
+            image = torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0) * 255.0
+
+            return image, t_eps_draft_p, t_eps_init
+
+    def prepare_for_training(self):
+        configure_batch_sizes(
+            mbs=self.cfg.micro_batch_size,
+            gbs=self.cfg.global_batch_size,
+            dp=parallel_state.get_data_parallel_world_size(),
+        )
+
+    def finish_training(self):
+        """no need to offload adam states here
+        """
+
+    def get_forward_output_and_loss_func(self, validation_step=False):
+        def fwd_output_and_loss_func(data_iterator, model):
+
+            batch = next(data_iterator)
+            batch_size = len(batch)
+            # Get different seeds for different dp rank
+            with get_cuda_rng_tracker().fork(get_data_parallel_rng_tracker_name()):
+                latents = torch.randn(
+                    [
+                        batch_size,
+                        self.in_channels,
+                        self.height // self.downsampling_factor,
+                        self.width // self.downsampling_factor,
+                    ],
+                    generator=None,
+                ).to(torch.cuda.current_device())
+
+            if validation_step:
+                output_tensor_draft_p, images, captions = self.log_visualization(batch)
+            else:
+                output_tensor_draft_p, epsilons_draft_p, epsilons_init = self.generate(batch, latents)
+
+            # in this nemo version the model and autocast dtypes are not synced
+            # so we need to explicitly cast it
+            if not parallel_state.is_pipeline_last_stage():
+                output_tensor_draft_p = output_tensor_draft_p.to(dtype=self.autocast_dtype)
+
+            def loss_func(output_tensor_draft_p):
+                # Loss per micro batch (ub).
+                if self.cfg.kl_coeff == 0.0 or validation_step:
+                    kl_penalty = torch.tensor([0.0]).to(torch.cuda.current_device())
+                else:
+                    kl_penalty = calculate_gaussian_kl_penalty_shared_var(epsilons_draft_p, epsilons_init).mean()
+                rewards = self.reward_model.get_reward(
+                    output_tensor_draft_p.permute(0, 2, 3, 1), batch
+                )  # (ub, H, W, C) -> (ub, C, H, W)
+                loss = -rewards.mean() + kl_penalty * self.cfg.kl_coeff
+
+                reduced_loss = average_losses_across_data_parallel_group([loss])
+                reduced_kl_penalty = average_losses_across_data_parallel_group([kl_penalty])
+
+                metrics = {"loss": reduced_loss, "kl_penalty": reduced_kl_penalty}
+
+                if validation_step:
+                    metrics["images_and_captions"] = images, captions
+
+                return (
+                    loss,
+                    metrics,
+                )
+
+            return output_tensor_draft_p, loss_func
+
+        return fwd_output_and_loss_func
+
+    def get_loss_and_metrics(self, batch, forward_only=False):
+        data_iter = get_iterator_k_split_list(batch, get_num_microbatches())
+
+        fwd_bwd_function = get_forward_backward_func()
+        losses_reduced_per_micro_batch = fwd_bwd_function(
+            forward_step_func=self.get_forward_output_and_loss_func(forward_only),
+            data_iterator=data_iter,
+            model=self.model,
+            num_microbatches=get_num_microbatches(),
+            forward_only=forward_only,
+            seq_length=None,
+            micro_batch_size=get_micro_batch_size(),
+        )
+
+        metrics = losses_reduced_per_micro_batch[0]
+        for key in ["loss", "kl_penalty"]:
+            if losses_reduced_per_micro_batch:
+                metric_mean = torch.stack(
+                    [loss_reduced[key] for loss_reduced in losses_reduced_per_micro_batch]
+                ).mean()
+            else:
+                metric_mean = torch.tensor(0.0, device=torch.cuda.current_device())
+
+            torch.distributed.broadcast(metric_mean, get_last_rank())
+            metrics[key] = metric_mean.cpu().item()
+
+        return metrics["loss"], metrics
diff --git a/nemo_aligner/models/nlp/gpt/gpt_reward_model.py b/nemo_aligner/models/nlp/gpt/gpt_reward_model.py
index 53bf83baa..b77248cda 100644
--- a/nemo_aligner/models/nlp/gpt/gpt_reward_model.py
+++ b/nemo_aligner/models/nlp/gpt/gpt_reward_model.py
@@ -12,19 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from abc import ABC, abstractclassmethod
 from copy import deepcopy
 from typing import Callable, List, Literal, Optional, Union
 from unittest.mock import patch
 
 import torch
-from megatron.core import parallel_state
 from megatron.core.model_parallel_config import ModelParallelConfig
 from megatron.core.models.gpt import GPTModel
 from megatron.core.tensor_parallel.layers import RowParallelLinear
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.utils import make_sharded_tensor_for_checkpoint, make_tp_sharded_tensor_for_checkpoint
 from torch import Tensor
 
 """Megatron Core based Reward Model"""
@@ -225,7 +222,6 @@ def forward(
 
         if self.post_process:
             return self.rm_head(hidden_states, lengths)
-
         return hidden_states
 
     def sharded_state_dict(self, prefix=""):
diff --git a/nemo_aligner/models/nlp/gpt/gpt_sft_model.py b/nemo_aligner/models/nlp/gpt/gpt_sft_model.py
index 419049d4a..f8bc3825e 100644
--- a/nemo_aligner/models/nlp/gpt/gpt_sft_model.py
+++ b/nemo_aligner/models/nlp/gpt/gpt_sft_model.py
@@ -33,6 +33,7 @@
 from nemo.collections.nlp.parts.mixins.nlp_adapter_mixins import NLPAdapterModelMixin
 from nemo.collections.nlp.parts.utils_funcs import get_last_rank
 from nemo_aligner.models.alignable_interface import SupervisedInterface
+from nemo_aligner.utils.text_generation_utils import tokenize_batch
 from nemo_aligner.utils.train_utils import (
     finish_validation_step,
     grad_reductions,
@@ -164,12 +165,21 @@ def generate(
                 inputs=inputs, length_params=length_params, sampling_params=sampling_params, strategy=strategy
             )
 
-        if isinstance(inputs, (list, tuple)) and isinstance(inputs[0], str):
-            raise NotImplementedError(
-                "`GPTSFTModel.generate()` does not currently support string inputs, please tokenize prompts first"
+        if not isinstance(inputs, (list, tuple)):
+            raise NotImplementedError(f"Expected type(inputs)=(list or tuple) but got {type(inputs)=}")
+
+        if isinstance(inputs[0], str):
+            # add_EOS=False since it is absent from nemo.collections.nlp.modules.common.text_generation_utils.megatron_gpt_generate
+            prompt_tokens, prompt_lengths = tokenize_batch(
+                sentences=inputs,
+                tokenizer=self.tokenizer,
+                max_len=self.cfg.encoder_seq_length,
+                add_BOS=sampling_params["add_BOS"],
+                add_EOS=False,
             )
+        else:
+            prompt_tokens, prompt_lengths = inputs
 
-        prompt_tokens, prompt_lengths = inputs
         max_prompt_length = prompt_lengths.max().item()
         max_response_length = length_params["max_length"]
         max_length = max_prompt_length + max_response_length
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py
index d550cf803..8a734639d 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py
@@ -11,13 +11,11 @@
 # 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.
-
 from contextlib import nullcontext
 from enum import Enum
 
 import torch
 from apex.transformer.pipeline_parallel.utils import _reconfigure_microbatch_calculator, get_num_microbatches
-from megatron.core import parallel_state
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from megatron.core.transformer.module import Float16Module
 from omegaconf.dictconfig import DictConfig
@@ -32,8 +30,9 @@
 from nemo.utils import AppState
 from nemo_aligner.models.alignable_interface import CriticModelInterface
 from nemo_aligner.models.nlp.gpt.megatron_gpt_reward_model import MegatronGPTRewardModel
+from nemo_aligner.utils import parallel_state
 from nemo_aligner.utils.train_utils import set_sync_funcs
-from nemo_aligner.utils.utils import masked_mean, offload_distributed_adam, swap_dict
+from nemo_aligner.utils.utils import copy_model_states_to_cpu, masked_mean, offload_distributed_adam
 
 
 class StateDictState(Enum):
@@ -48,11 +47,12 @@ class MegatronGPTCriticModel(MegatronGPTRewardModel, CriticModelInterface):
     def __init__(self, cfg: DictConfig, trainer: Trainer):
         super().__init__(cfg, trainer=trainer)
 
-        # filled by the examples script
-        self.rm_state_dict = None
+        # must be populated by the examples script
+        self.rm_state_dict_cpu = None
 
-        # for the critic states on cpu
-        self.cpu_state_dict = None
+        self.critic_state_dict_cpu = copy_model_states_to_cpu(
+            self, cpu_dict=None, megatron_amp_O2=self.cfg.megatron_amp_O2, sync=True
+        )
 
         # for distributed adam offload
         self.distributed_adam_offload_manager = None
@@ -61,6 +61,12 @@ def __init__(self, cfg: DictConfig, trainer: Trainer):
         self.to_offload_adam_states = self.cfg.get("offload_adam_states")
         self.clip_val = self.cfg.get("loss_clip_val")
 
+    def on_load_checkpoint(self, *args, **kwargs):
+        super().on_load_checkpoint(*args, **kwargs)
+        self.critic_state_dict_cpu = copy_model_states_to_cpu(
+            self, cpu_dict=self.critic_state_dict_cpu, megatron_amp_O2=self.cfg.megatron_amp_O2, sync=True
+        )
+
     def prepare_for_inference(self):
         super().prepare_for_inference()
         self.offload_adam_states()
@@ -87,7 +93,7 @@ def get_loss_and_metrics(self, batch, forward_only):
 
         losses_reduced_per_micro_batch = fwd_bwd_function(
             forward_step_func=self.get_forward_output_and_loss_func(),
-            data_iterator=data_iter,
+            data_iterator=self._make_data_iterator_list(data_iter),
             model=self.model,
             num_microbatches=get_num_microbatches(),
             forward_only=forward_only,
@@ -209,15 +215,14 @@ def infer_rm_critic(self, *args, **kwargs):
 
         outputs = []
         for fn in call_order:
-            output, exceeded = fn(*args, **kwargs)
+            output = fn(*args, **kwargs)
             outputs.append(output)
 
         if original_state == StateDictState.CRITIC:
             # reverse the output back
             outputs = reversed(outputs)
 
-        # always return rewards, value, exceeded
-        return (*outputs, exceeded)
+        return tuple(outputs)
 
     def set_output_sequence_flag(self, value_to_set):
         if isinstance(self.model, Float16Module):
@@ -230,16 +235,14 @@ def set_output_sequence_flag(self, value_to_set):
 
     def _load_critic(self):
         if self.loaded_state_dict == StateDictState.REWARD:
-            # no need to put the RM back to cpu, we already have it
-            swap_dict(self, self.cpu_state_dict, offload_onto_cpu=False, megatron_amp_O2=self.megatron_amp_O2)
+            self.load_state_dict(self.critic_state_dict_cpu)
 
             self.set_output_sequence_flag(True)
-
             self.loaded_state_dict = StateDictState.CRITIC
 
     def _load_rm(self):
         if self.loaded_state_dict == StateDictState.CRITIC:
-            self.cpu_state_dict = swap_dict(self, self.rm_state_dict, megatron_amp_O2=self.megatron_amp_O2)
+            self.load_state_dict(self.rm_state_dict_cpu)
 
             self.set_output_sequence_flag(False)
             self.loaded_state_dict = StateDictState.REWARD
@@ -251,3 +254,8 @@ def _infer_critic(self, *args, **kwargs):
     def _infer_rm(self, *args, **kwargs):
         self._load_rm()
         return self.infer(*args, **kwargs)
+
+    def finish_training(self):
+        self.critic_state_dict_cpu = copy_model_states_to_cpu(
+            self, cpu_dict=self.critic_state_dict_cpu, megatron_amp_O2=self.cfg.megatron_amp_O2, sync=True
+        )
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_dpo_model.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_dpo_model.py
index e3ecbf9ab..d5174fa4b 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_dpo_model.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_dpo_model.py
@@ -59,16 +59,27 @@ def __init__(self, cfg: DictConfig, trainer: Trainer):
         self.ref_policy_state_dict = None
 
         self.ref_policy_kl_penalty = self.cfg.dpo.get("ref_policy_kl_penalty", 0.0)
-        self.avg_log_probs = self.cfg.dpo.get("average_log_probs", False)
+        self.preference_avg_log_probs = self.cfg.dpo.get("preference_average_log_probs", False)
+        self.sft_avg_log_probs = self.cfg.dpo.get("sft_average_log_probs", self.preference_avg_log_probs)
+
+        self.preference_loss_weight = self.cfg.dpo.get("preference_loss_weight", 1)
+        self.sft_loss_weight = self.cfg.dpo.get("sft_loss_weight", 0)
+        assert (
+            self.preference_loss_weight != 0 or self.sft_loss_weight != 0
+        ), "sft loss weight and dpo loss weight cannot both be 0"
+
+        # variants of preference losses, by default DPO.
+        self.preference_loss = self.cfg.dpo.get("preference_loss", "dpo")
+        self.gt_reward_scale = self.cfg.dpo.get("gt_reward_scale", 1.0)
 
     @torch.no_grad()
-    def gather_and_split_rewards(self, pi_logprobs, ref_logprobs, labels):
+    def gather_and_split_rewards(self, pi_logprobs, ref_logprobs, labels, average_log_probs=False):
         pi_logprobs = pi_logprobs.detach()
 
         dp_group = parallel_state.get_data_parallel_group()
 
         batch_logs = self.get_reduced_masked_logps(
-            pi_logprobs - ref_logprobs, labels[:, 1:], average_log_probs=self.avg_log_probs
+            pi_logprobs - ref_logprobs, labels[:, 1:], average_log_probs=average_log_probs
         )
 
         output_list = [torch.zeros_like(batch_logs) for _ in range(dp_group.size())]
@@ -103,12 +114,14 @@ def fwd_output_and_loss_func(dataloader_iter, model, checkpoint_activations_all_
                             "ref_policy_log_probs_rejected",
                             "chosen_labels",
                             "rejected_labels",
+                            "chosen_rewards",
+                            "rejected_rewards",
                         )
                     )
 
             batch = {key: val.cuda(non_blocking=True) if key in required_keys else None for key, val in batch.items()}
 
-            tokens, labels, ref_logprobs = None, None, None
+            tokens, labels, ref_logprobs, gt_rewards = None, None, None, None
             if batch["chosen"] is not None and batch["rejected"] is not None:
                 tokens = torch.cat((batch["chosen"], batch["rejected"]), dim=0)
 
@@ -120,6 +133,9 @@ def fwd_output_and_loss_func(dataloader_iter, model, checkpoint_activations_all_
                     (batch["ref_policy_log_probs_chosen"], batch["ref_policy_log_probs_rejected"]), dim=0
                 )
 
+            if batch["chosen_rewards"] is not None and batch["rejected_rewards"] is not None:
+                gt_rewards = torch.cat((batch["chosen_rewards"], batch["rejected_rewards"]), dim=0)
+
             # this is necessary if MBS > 1 with the new GBS padding logic, as you may get batch dim > 1 in some configs
             # these two lines ensure your position_ids and attn_mask are always B=1
             # position_ids = batch["position_ids"][0:1]
@@ -157,18 +173,42 @@ def loss_func(output_tensor):
                     vocab_parallel_logits=output_tensor, target=labels, higher_stability=True
                 )
 
-                loss, acc_chosen = self.loss_func(
-                    per_token_logps, ref_logprobs, labels[:, 1:], average_log_probs=self.avg_log_probs
+                preference_loss, acc_chosen = self.loss_func(
+                    per_token_logps,
+                    ref_logprobs,
+                    labels[:, 1:],
+                    gt_rewards,
+                    average_log_probs=self.preference_avg_log_probs,
                 )
 
-                reduced_loss = average_losses_across_data_parallel_group([loss])
-                reduced_acc = average_losses_across_data_parallel_group([acc_chosen])
+                sft_loss = torch.zeros_like(preference_loss)
+                if self.sft_loss_weight != 0:
+                    sft_loss = self.sft_loss_func(
+                        per_token_logps, labels[:, 1:], average_log_probs=self.sft_avg_log_probs
+                    )
+                loss = self.preference_loss_weight * preference_loss + self.sft_loss_weight * sft_loss
+
+                (
+                    reduced_loss,
+                    reduced_preference_loss,
+                    reduced_sft_loss,
+                    reduced_acc,
+                ) = average_losses_across_data_parallel_group([loss, preference_loss, sft_loss, acc_chosen])
 
-                out_chosen, out_rejected = self.gather_and_split_rewards(per_token_logps, ref_logprobs, labels)
+                out_chosen, out_rejected = self.gather_and_split_rewards(
+                    per_token_logps, ref_logprobs, labels, average_log_probs=self.preference_avg_log_probs
+                )
 
                 return (
                     loss,
-                    {"avg": reduced_loss, "acc": reduced_acc, "out_chosen": out_chosen, "out_rejected": out_rejected,},
+                    {
+                        "avg": reduced_loss,
+                        "avg_sft_loss": reduced_sft_loss,
+                        "avg_preference_loss": reduced_preference_loss,
+                        "acc": reduced_acc,
+                        "out_chosen": out_chosen,
+                        "out_rejected": out_rejected,
+                    },
                 )
 
             return output_tensor, loss_func
@@ -190,13 +230,50 @@ def get_reduced_masked_logps(self, logps, labels, average_log_probs=False):
         else:
             return (logps * loss_mask).sum(-1)
 
-    def loss_func(self, pi_logprobs, ref_logprobs, labels, average_log_probs=False):
+    def loss_func(self, pi_logprobs, ref_logprobs, labels, gt_rewards, average_log_probs=False):
         rewards = self.get_reduced_masked_logps(
             pi_logprobs - ref_logprobs, labels, average_log_probs=average_log_probs
         )
-        chosen_rewards, reject_rewards = self.split_output_tensor(self.ref_policy_kl_penalty * rewards)
-
-        loss = -torch.nn.functional.logsigmoid(chosen_rewards - reject_rewards).mean(0)
+        chosen_rewards, reject_rewards = self.split_output_tensor(rewards)
+        rewards_delta = chosen_rewards - reject_rewards
+
+        if self.preference_loss == "dpo":
+            loss = -torch.nn.functional.logsigmoid(self.ref_policy_kl_penalty * rewards_delta).mean(0)
+        elif self.preference_loss == "rpo_bwd_kl":
+            logbeta_hat_chosen = torch.nn.functional.logsigmoid(self.ref_policy_kl_penalty * rewards_delta)
+            logbeta_hat_rejected = torch.nn.functional.logsigmoid(-self.ref_policy_kl_penalty * rewards_delta)
+
+            chosen_gt_rewards, reject_gt_rewards = self.split_output_tensor(gt_rewards)
+            gt_rewards_delta = self.gt_reward_scale * (chosen_gt_rewards - reject_gt_rewards)
+            logalpha_hat_chosen = torch.nn.functional.logsigmoid(gt_rewards_delta)
+            logalpha_hat_rejected = torch.nn.functional.logsigmoid(-gt_rewards_delta)
+
+            loss = (
+                torch.exp(logalpha_hat_chosen) * (logalpha_hat_chosen - logbeta_hat_chosen)
+                + torch.exp(logalpha_hat_rejected) * (logalpha_hat_rejected - logbeta_hat_rejected)
+            ).mean(0)
+        elif self.preference_loss == "rpo_fwd_kl":
+            logbeta_hat_chosen = torch.nn.functional.logsigmoid(self.ref_policy_kl_penalty * rewards_delta)
+            logbeta_hat_rejected = torch.nn.functional.logsigmoid(-self.ref_policy_kl_penalty * rewards_delta)
+
+            chosen_gt_rewards, reject_gt_rewards = self.split_output_tensor(gt_rewards)
+            gt_rewards_delta = self.gt_reward_scale * (chosen_gt_rewards - reject_gt_rewards)
+            logalpha_hat_chosen = torch.nn.functional.logsigmoid(gt_rewards_delta)
+            logalpha_hat_rejected = torch.nn.functional.logsigmoid(-gt_rewards_delta)
+
+            loss = (
+                torch.exp(logbeta_hat_chosen) * (logbeta_hat_chosen - logalpha_hat_chosen)
+                + torch.exp(logbeta_hat_rejected) * (logbeta_hat_rejected - logalpha_hat_rejected)
+            ).mean(0)
+        elif self.preference_loss == "ipo":
+            loss = torch.mean((chosen_rewards - reject_rewards - 1.0 / (2.0 * self.ref_policy_kl_penalty)) ** 2, 0)
+        elif self.preference_loss == "rpo_sq":
+            chosen_gt_rewards, reject_gt_rewards = self.split_output_tensor(gt_rewards)
+            gt_rewards_delta = self.gt_reward_scale * (chosen_gt_rewards - reject_gt_rewards)
+
+            loss = torch.mean((self.ref_policy_kl_penalty * rewards_delta - gt_rewards_delta) ** 2, 0)
+        else:
+            raise NotImplementedError(f"preference_loss {self.preference_loss} is not implemented")
 
         with torch.no_grad():
             comp = chosen_rewards > reject_rewards
@@ -204,6 +281,11 @@ def loss_func(self, pi_logprobs, ref_logprobs, labels, average_log_probs=False):
 
         return loss, acc_chosen
 
+    def sft_loss_func(self, pi_logprobs, labels, average_log_probs=False):
+        logprobs = self.get_reduced_masked_logps(pi_logprobs, labels, average_log_probs=average_log_probs)
+        chosen_logprobs, _ = self.split_output_tensor(logprobs)
+        return -chosen_logprobs.mean(0)
+
     def get_loss_and_metrics(self, batch, forward_only):
         seq_length = batch["chosen"].shape[1]
 
@@ -235,20 +317,26 @@ def get_loss_and_metrics(self, batch, forward_only):
             rewards_all_mean = rewards_all.mean()
             rewards_all_std = rewards_all.std()
 
-            # average loss across micro batches
-            loss_tensors_list = [loss_reduced["avg"] for loss_reduced in losses_reduced_per_micro_batch]
-            loss_tensor = torch.concat(loss_tensors_list)
-            loss_mean = loss_tensor.mean()
-            acc_tensors_list = [loss_reduced["acc"] for loss_reduced in losses_reduced_per_micro_batch]
-
-            if len(acc_tensors_list) == 1:
-                acc_tensor = acc_tensors_list[0]
-            elif len(acc_tensors_list) > 1:
-                acc_tensor = torch.concat(acc_tensors_list)
-            acc_mean = acc_tensor.mean()
+            loss_mean = torch.as_tensor(
+                [loss_reduced["avg"] for loss_reduced in losses_reduced_per_micro_batch],
+                device=torch.cuda.current_device(),
+            ).mean()
+            sft_loss_mean = torch.as_tensor(
+                [loss_reduced["avg_sft_loss"] for loss_reduced in losses_reduced_per_micro_batch],
+                device=torch.cuda.current_device(),
+            ).mean()
+            preference_loss_mean = torch.as_tensor(
+                [loss_reduced["avg_preference_loss"] for loss_reduced in losses_reduced_per_micro_batch],
+                device=torch.cuda.current_device(),
+            ).mean()
+            acc_mean = torch.as_tensor(
+                [loss_reduced["acc"] for loss_reduced in losses_reduced_per_micro_batch],
+                device=torch.cuda.current_device(),
+            ).mean()
         else:
-
             loss_mean = torch.tensor(0.0, device=torch.cuda.current_device())
+            sft_loss_mean = torch.tensor(0.0, device=torch.cuda.current_device())
+            preference_loss_mean = torch.tensor(0.0, device=torch.cuda.current_device())
             acc_mean = torch.tensor(0.0, device=torch.cuda.current_device())
 
             rewards_chosen_mean = torch.tensor(0.0, device=torch.cuda.current_device())
@@ -258,6 +346,8 @@ def get_loss_and_metrics(self, batch, forward_only):
 
         # we can only log on one rank if it is rank zero so we broadcast from last rank
         torch.distributed.broadcast(loss_mean, get_last_rank())
+        torch.distributed.broadcast(sft_loss_mean, get_last_rank())
+        torch.distributed.broadcast(preference_loss_mean, get_last_rank())
         torch.distributed.broadcast(acc_mean, get_last_rank())
 
         torch.distributed.broadcast(rewards_chosen_mean, get_last_rank())
@@ -267,6 +357,8 @@ def get_loss_and_metrics(self, batch, forward_only):
 
         metrics = {
             "loss": loss_mean,
+            "sft_loss": sft_loss_mean,
+            "preference_loss": preference_loss_mean,
             "acc": acc_mean,
             "rewards_chosen_mean": rewards_chosen_mean,
             "rewards_rejected_mean": rewards_rejected_mean,
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_ppo_actor.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_ppo_actor.py
index 1c250f233..ba9ebb092 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_ppo_actor.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_ppo_actor.py
@@ -279,7 +279,8 @@ def infer(self, inference_batch):
         response_lengths = strategy.get_lengths()
         max_response_length = response_lengths.max().item()
 
-        response_tokens = torch.cuda.LongTensor(actor_output["token_ids"])
+        response_tokens = torch.cuda.LongTensor(actor_output["token_ids"]) if actor_output else None
+        response_tokens = broadcast_2d_tensor_within_pp(response_tokens, dtype=torch.long)
 
         # Sanity check to validate response length.
         if max_response_length != response_tokens.size(1):
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py
index 90e310c98..467dc0726 100644
--- a/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py
+++ b/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py
@@ -14,12 +14,12 @@
 
 
 import warnings
-from typing import List, Union
+from typing import List, Tuple, Union
 
 import torch
-from apex.transformer.pipeline_parallel.utils import _reconfigure_microbatch_calculator, get_num_microbatches
-from megatron.core import parallel_state
+from apex.transformer.pipeline_parallel.utils import get_num_microbatches
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
+from megatron.core.utils import divide
 from omegaconf.dictconfig import DictConfig
 from pytorch_lightning.trainer.trainer import Trainer
 
@@ -30,11 +30,12 @@
     get_ltor_masks_and_position_ids,
 )
 from nemo.collections.nlp.parts.utils_funcs import get_last_rank
-from nemo.utils import AppState, logging
+from nemo.utils import logging
 from nemo.utils.dtype import str_to_dtype
 from nemo_aligner.models.alignable_interface import Inferrable, SupervisedInterface
 from nemo_aligner.models.nlp.gpt.gpt_reward_model import GPTRewardModel
-from nemo_aligner.utils.distributed import broadcast_2d_tensor, broadcast_2d_tensor_within_pp, gather_tensor
+from nemo_aligner.utils import parallel_state
+from nemo_aligner.utils.distributed import broadcast_2d_tensor_within_pp
 from nemo_aligner.utils.text_generation_utils import tokenize_batch
 from nemo_aligner.utils.train_utils import (
     finish_validation_step,
@@ -70,6 +71,8 @@ def __init__(self, cfg: DictConfig, trainer: Trainer):
             self.rew_mean = cfg.reward_standardization.mean
             self.rew_std = cfg.reward_standardization.std
 
+        self.forward_micro_batch_size = self.cfg.get("forward_micro_batch_size", self.cfg.micro_batch_size)
+
     def model_provider_func(self, pre_process, post_process):
         """Model depends on pipeline paralellism."""
 
@@ -248,7 +251,7 @@ def get_loss_and_metrics(self, batch, forward_only):
 
         losses_reduced_per_micro_batch = fwd_bwd_function(
             forward_step_func=self.get_forward_output_and_loss_func(forward_only),
-            data_iterator=data_iter,
+            data_iterator=self._make_data_iterator_list(data_iter),
             model=self.model,
             num_microbatches=get_num_microbatches(),
             forward_only=forward_only,
@@ -351,92 +354,66 @@ def finish_validation_step(self):
 
     def infer(
         self,
-        inputs: Union[List[str], List[List[int]]] = None,
-        sequence_length: List[int] = None,
+        inputs: Union[List[str], Tuple[torch.Tensor, torch.Tensor]],
+        add_BOS: bool = False,
         add_EOS: bool = False,
     ):
-        tokenizer = self.tokenizer
-        max_seq_length = self.cfg.encoder_seq_length
-
-        exceeded = None
-        context_tokens_tensor = None
-        context_length_tensor = None
-        if torch.distributed.get_rank() == 0:
-            if sequence_length is not None:
-                exceeded = [False for _ in range(len(inputs))]
-                context_tokens_tensor = torch.tensor(inputs).cuda()
-                context_length_tensor = torch.tensor(sequence_length).cuda()
-            else:
-                context_tokens_tensor, context_length_tensor, exceeded = tokenize_batch(
-                    tokenizer, inputs, max_seq_length, False, add_EOS=add_EOS
-                )
-            if context_length_tensor.dim() == 1:
-                context_length_tensor = context_length_tensor.unsqueeze(-1)
-
-        context_length_tensor = broadcast_2d_tensor(context_length_tensor, src=0, group=None, dtype=torch.int64)
-        if context_length_tensor.dim() == 2:
-            context_length_tensor = context_length_tensor.squeeze(-1)
-
-        context_tokens_tensor = broadcast_2d_tensor(context_tokens_tensor, src=0, group=None, dtype=torch.int64)
+        if isinstance(inputs, tuple):
+            context_tokens_tensor, context_length_tensor = inputs
+        elif isinstance(inputs, list):
+            assert all(isinstance(item, str) for item in inputs), "list must contain all strings in infer function"
+            context_tokens_tensor, context_length_tensor = tokenize_batch(
+                inputs, self.tokenizer, self.cfg.encoder_seq_length, add_BOS=add_BOS, add_EOS=add_EOS,
+            )
+        else:
+            raise NotImplementedError(f"{type(inputs)=} is not supported in infer function")
 
-        # Select subset of data needed for this rank.
-        dp_size = parallel_state.get_data_parallel_world_size()
-        dp_rank = parallel_state.get_data_parallel_rank()
-        context_length_tensor = context_length_tensor.chunk(dp_size)[dp_rank]
-        context_tokens_tensor = context_tokens_tensor.chunk(dp_size)[dp_rank]
+        context_tokens_tensor = context_tokens_tensor.cuda()
+        context_length_tensor = context_length_tensor.cuda()
 
+        inference_batch_size, sequence_length = context_tokens_tensor.size()
         attention_mask, _, position_ids = get_ltor_masks_and_position_ids(
             context_tokens_tensor,
-            tokenizer.eos_id,
+            self.tokenizer.eos_id,
             self.cfg.get("reset_position_ids", False),
             self.cfg.get("reset_attention_mask", False),
             self.cfg.get("eod_mask_loss", False),
         )
-        micro_batch_size = context_tokens_tensor.shape[0]
-        sequence_length = context_tokens_tensor.shape[1]
-
-        app_state = AppState()
-        _reconfigure_microbatch_calculator(
-            rank=app_state.global_rank,
-            rampup_batch_size=None,
-            global_batch_size=micro_batch_size,
-            micro_batch_size=micro_batch_size,
-            data_parallel_size=1,
-        )
-        attention_mask_repeat = torch.concat([attention_mask for _ in range(micro_batch_size)])
-        rewards = self.forward_step(
-            [context_tokens_tensor, context_length_tensor, position_ids, attention_mask_repeat],
-            micro_batch_size,
-            sequence_length,
-        )
+        attention_mask = attention_mask.expand(inference_batch_size, -1, -1, -1)
+        inputs = [context_tokens_tensor, context_length_tensor, position_ids, attention_mask]
+
+        # if inference batch size is smaller than forward mbs run it at the lower batch size
+        forward_micro_batch_size = min(inference_batch_size, self.forward_micro_batch_size)
+
+        num_microbatches = divide(inference_batch_size, forward_micro_batch_size)
+        data_iter = get_iterator_k_split(inputs, num_microbatches)
+
+        rewards = self.forward_step(data_iter, forward_micro_batch_size, sequence_length, num_microbatches)
 
         if parallel_state.is_pipeline_last_stage():
-            rewards = rewards[0]["reward"]
+            rewards = torch.cat(rewards)
 
             # Standardize values to subtract a bias.
             if self.enable_standardization:
                 rewards = (rewards - self.rew_mean) / self.rew_std
 
         rewards = broadcast_2d_tensor_within_pp(rewards)
+        return rewards
 
-        rewards_list = gather_tensor(
-            rewards, dst=parallel_state.get_data_parallel_src_rank(), group=parallel_state.get_data_parallel_group()
-        )
-
-        return rewards_list, exceeded
-
-    def forward_step(self, batch, micro_batch_size, sequence_length):
+    def forward_step(self, data_iter, micro_batch_size, sequence_length, num_microbatches):
         set_sync_funcs(self, forward_only=True)
 
         fwd_bwd_function = get_forward_backward_func()
         output_tensor = fwd_bwd_function(
             forward_step_func=self.get_forward_output_only_func(),
-            data_iterator=iter([batch]),
+            data_iterator=self._make_data_iterator_list(data_iter),
             model=self.model,
-            num_microbatches=get_num_microbatches(),
+            num_microbatches=num_microbatches,
             forward_only=True,
             seq_length=sequence_length,
             micro_batch_size=micro_batch_size,
+            # Prevent automatic scaling by the number of micro-batches, as we are not in training mode.
+            collect_non_loss_data=True,
         )
         return output_tensor
 
@@ -451,8 +428,8 @@ def fwd_output_only_func(batch, model):
             if not parallel_state.is_pipeline_last_stage():
                 output_tensor = output_tensor.to(dtype=self.autocast_dtype)
 
-            def id_func(output_tensor):
-                return output_tensor, {"reward": output_tensor}
+            def id_func(output_tensor, non_loss_data=True):
+                return output_tensor
 
             return output_tensor, id_func
 
diff --git a/nemo_aligner/models/nlp/gpt/reward_critic_clients.py b/nemo_aligner/models/nlp/gpt/reward_critic_clients.py
index d4d87a573..5aa4f6afe 100644
--- a/nemo_aligner/models/nlp/gpt/reward_critic_clients.py
+++ b/nemo_aligner/models/nlp/gpt/reward_critic_clients.py
@@ -17,14 +17,13 @@
 
 import numpy as np
 import torch
-from megatron.core import parallel_state
 from omegaconf import DictConfig
 
 from nemo_aligner.servers.http_communicator import HTTPCommunicator
+from nemo_aligner.utils import parallel_state
 from nemo_aligner.utils.distributed import broadcast_2d_tensor_within_mp, gather_tensor, run_if_model_parallel_src
 from nemo_aligner.utils.server_utils import FutureResult
 
-
 """A remote client that acts like a real Reward Model and Critic forwards all requests from the actor
     over to the remote PyTrition server
 """
@@ -44,7 +43,9 @@ def get_future_result(future, *keys):
         if output is not None:
             result = torch.tensor(output[key], device=torch.cuda.current_device())
 
-        results.append(broadcast_2d_tensor_within_mp(result))
+        ten = broadcast_2d_tensor_within_mp(result)
+
+        results.append(ten)
 
     if len(results) == 1:
         return results[0]
@@ -71,8 +72,7 @@ def result(self):
 
         self.critic_future = None
         self.rm_future = None
-
-        return rewards, values
+        return rewards.flatten(), values
 
 
 class SaveFuture(FutureResult):
@@ -127,13 +127,13 @@ def infer_rm_critic(self, rollout_batch):
         }
 
         critic_future = run_if_model_parallel_src(
-            self.communicator.send_data_to_server, server_name=self.cfg.critic.name.infer, data=send_data
+            self.communicator.send_data_to_server, server_name=self.cfg.critic.name.infer, data=send_data,
         )
 
         rm_future = None
         if not self.combine_rm_and_critic_server:
             rm_future = run_if_model_parallel_src(
-                self.communicator.send_data_to_server, server_name=self.cfg.reward_model.name, data=send_data
+                self.communicator.send_data_to_server, server_name=self.cfg.reward_model.name, data=send_data,
             )
 
         return RMCriticFutureResult(critic_future, rm_future, self.combine_rm_and_critic_server, og_seq_length)
@@ -148,9 +148,9 @@ def train(self, ppo_rollout_data):
         )
 
         send_data["tokens"] = func(ppo_rollout_data["response_tokens"], dtype=torch.int64)
-        send_data["returns"] = func(ppo_rollout_data["returns"])
-        send_data["prev_values"] = func(ppo_rollout_data["values"])
-        send_data["mask"] = func(ppo_rollout_data["mask"])
+        send_data["returns"] = func(ppo_rollout_data["returns"], dtype=torch.float32)
+        send_data["prev_values"] = func(ppo_rollout_data["values"], dtype=torch.float32)
+        send_data["mask"] = func(ppo_rollout_data["mask"], dtype=torch.float32)
 
         future = None
         if torch.distributed.get_rank() == 0:
diff --git a/nemo_aligner/servers/server_callables.py b/nemo_aligner/servers/server_callables.py
deleted file mode 100644
index 04df9acfc..000000000
--- a/nemo_aligner/servers/server_callables.py
+++ /dev/null
@@ -1,106 +0,0 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. 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 threading
-from typing import Dict
-
-import numpy as np
-import torch
-from megatron.core import parallel_state
-from pytriton.decorators import batch
-from pytriton.exceptions import PyTritonUnrecoverableError
-from pytriton.model_config import Tensor
-
-from nemo_aligner.servers.constants import ServerSignal
-from nemo_aligner.utils.server_utils import decode_bytes_ndarray, lock_method, pad_input
-
-
-def run_rm_or_critic_inference(infer_fn, inputs):
-    """run the infer function for either the critic or the rm
-    """
-    sentences = inputs.pop("sentences", None)
-    if sentences is not None:
-        sentences = decode_bytes_ndarray(sentences)
-    tokens = inputs.pop("tokens", None)
-
-    sequence_lengths = inputs.pop("sequence_lengths", None)
-    add_EOS = inputs.pop("add_EOS", None)
-
-    assert sentences is not None or tokens is not None, "Both sentences and tokens cannot be None."
-
-    dp_size = parallel_state.get_data_parallel_world_size()
-
-    # Ensure that the batch size is a multiple of the data parallel size. Otherwise, pad it.
-    sentences, extra_sentences = pad_input(sentences, dp_size)
-    tokens, extra_tokens = pad_input(tokens, dp_size)
-    sequence_lengths, extra_sequence_lengths = pad_input(sequence_lengths, dp_size)
-
-    if add_EOS is not None:
-        add_EOS = add_EOS[0]
-
-    inputs = sentences if sentences is not None else tokens
-    extra = extra_sentences if sentences is not None else extra_tokens
-    if sequence_lengths is not None:
-        assert len(inputs) == len(sequence_lengths)
-        assert extra_sequence_lengths == extra
-
-    try:
-        *list_outputs, exceeded = infer_fn(inputs=inputs, sequence_length=sequence_lengths, add_EOS=add_EOS)
-
-        processed_outputs = []
-
-        for output in list_outputs:
-            output = torch.cat(output, dim=0)
-            # unpad
-            output = output[: output.size(0) - extra]
-
-            processed_outputs.append(output.cpu().numpy())
-
-        exceeded = exceeded[: len(exceeded) - extra]
-
-    except RuntimeError as e:
-        raise PyTritonUnrecoverableError(f"Fatal error occurred - no further inferences possible. {e}") from e
-
-    return (*processed_outputs, np.array(exceeded, dtype=np.int32).reshape(-1, 1))
-
-
-class RewardModelCallable:
-    def __init__(self, *, model_name: str, infer_fn: callable, lock: threading.Lock):
-        self.model_name = model_name
-        self.lock = lock
-        self.infer_fn = infer_fn
-        self.inputs = (
-            Tensor(name="sentences", shape=(-1,), dtype=bytes, optional=True),
-            Tensor(name="tokens", shape=(-1,), dtype=np.int64, optional=True),
-            Tensor(name="sequence_lengths", shape=(-1,), dtype=np.int64, optional=True),
-            Tensor(name="add_EOS", shape=(1,), dtype=np.bool_, optional=True),
-        )
-        self.outputs = (
-            Tensor(name="rewards", shape=(1,), dtype=np.float32),
-            Tensor(name="exceeded", shape=(1,), dtype=np.int32),
-        )
-
-    @batch
-    @lock_method("self.lock")
-    def infer(self, **inputs: np.ndarray) -> Dict[str, np.ndarray]:
-        choice = ServerSignal.FORWARD.cuda()
-        torch.distributed.broadcast(choice, 0)
-
-        rewards, exceeded = run_rm_or_critic_inference(self.infer_fn, inputs=inputs)
-
-        output_dict = {
-            "rewards": rewards,
-            "exceeded": exceeded,
-        }
-        return output_dict
diff --git a/nemo_aligner/utils/__init__.py b/nemo_aligner/utils/__init__.py
index 8be8d0937..4fc50543f 100644
--- a/nemo_aligner/utils/__init__.py
+++ b/nemo_aligner/utils/__init__.py
@@ -11,5 +11,3 @@
 # 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.
-
-from nemo_aligner.utils.utils import *
diff --git a/nemo_aligner/utils/distributed.py b/nemo_aligner/utils/distributed.py
index 78c860108..642fdb7d6 100644
--- a/nemo_aligner/utils/distributed.py
+++ b/nemo_aligner/utils/distributed.py
@@ -19,27 +19,39 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Dict, Optional, Union
+from typing import Optional
 
 import torch
-from megatron.core import parallel_state, tensor_parallel
+from megatron.core import tensor_parallel
 
-from nemo.collections.nlp.modules.common.text_generation_utils import get_model_parallel_src_rank
 from nemo.utils.timers import NamedTimer
+from nemo_aligner.utils import parallel_state
+from nemo_aligner.utils.parallel_state import get_model_parallel_group, get_model_parallel_src_rank
 from nemo_aligner.utils.ppo_utils import calculate_entropy
 
 
-def gather_tensor(tensor, dst, group, dtype=torch.float32):
-    """Gather any 2d tensor to the dst rank from every other rank in the given group.
-    All the ranks that send or receive data must call this function."""
-    tensor = tensor.to(device=torch.cuda.current_device(), dtype=dtype)
-    if torch.distributed.get_rank() == dst:
-        gather_list = [torch.empty_like(tensor) for _ in range(torch.distributed.get_world_size(group))]
-    else:
-        gather_list = None
+def rebalance_nd_tensor(tensor, group):
+    """
+    Takes tensors with variable leading sizes (at dim=0) and then stack them into a single tensor.
+    
+    NOTE: assumes all other (i.e., non-zero) dimensions are equal.
+    """
+    num_samples = torch.as_tensor(tensor.size(0), dtype=torch.int64, device=torch.cuda.current_device())
+    batch_num_per_rank = torch.zeros(
+        torch.distributed.get_world_size(group), dtype=torch.int64, device=torch.cuda.current_device()
+    )
+    torch.distributed.all_gather_into_tensor(batch_num_per_rank, num_samples, group=group)
 
-    torch.distributed.gather(tensor, gather_list=gather_list, dst=dst, group=group)
-    return gather_list
+    B = batch_num_per_rank.sum()
+    other_dims = tensor.shape[1:]
+
+    indices = batch_num_per_rank.cumsum(dim=0)
+    output_tensor = torch.zeros(B, *other_dims, dtype=tensor.dtype, device=torch.cuda.current_device())
+
+    # tensor_split is a view we can copy into
+    output_tensor.tensor_split(indices[0:-1].cpu())[torch.distributed.get_rank(group=group)].copy_(tensor)
+    torch.distributed.all_reduce(output_tensor, group=group)
+    return output_tensor
 
 
 def broadcast_2d_tensor(tensor, src, group, dtype=torch.float32):
@@ -72,7 +84,7 @@ def broadcast_2d_tensor(tensor, src, group, dtype=torch.float32):
 def broadcast_2d_tensor_within_mp(tensor, dtype=torch.float32):
     """helper function to broadcast within the model parallel group
     """
-    group = parallel_state.get_model_parallel_group()
+    group = get_model_parallel_group()
 
     if torch.distributed.get_world_size(group) > 1:
         return broadcast_2d_tensor(tensor, get_model_parallel_src_rank(), group, dtype=dtype)
@@ -92,6 +104,19 @@ def broadcast_2d_tensor_within_pp(tensor, dtype=torch.float32):
     return tensor
 
 
+def gather_tensor(tensor, dst, group, dtype=None):
+    """Gather any tensor to the dst rank from every other rank in the given group.
+    All the ranks that send or receive data must call this function."""
+    tensor = tensor.to(device=torch.cuda.current_device(), dtype=dtype)
+    if torch.distributed.get_rank() == dst:
+        gather_list = [torch.empty_like(tensor) for _ in range(torch.distributed.get_world_size(group))]
+    else:
+        gather_list = None
+
+    torch.distributed.gather(tensor, gather_list=gather_list, dst=dst, group=group)
+    return gather_list
+
+
 def run_if_model_parallel_src(fn, *fn_args, **fn_kwargs):
     """This function is meant to wrap an arbitary function to only call the function
     if it's the model parallel src. So if we have DP=2, this function will be called
@@ -271,29 +296,6 @@ def from_parallel_logits_to_logprobs(vocab_parallel_logits, target, inference_on
     ].contiguous()
 
 
-def pad_tensors_to_max_global_seq_len(list_of_tensors, pad_value, group, sequence_length_to_pad_to=None):
-    """pad a list of tensors to the global sequence length across the specified group
-    """
-    # compute the local padding
-    tensors_padded = torch.nn.utils.rnn.pad_sequence(list_of_tensors, batch_first=True, padding_value=pad_value)
-
-    # find global max seq length
-    max_seq_length = torch.tensor([tensors_padded.size(-1)], dtype=torch.float32, device=torch.cuda.current_device())
-    torch.distributed.all_reduce(max_seq_length, op=torch.distributed.ReduceOp.MAX, group=group)
-    max_seq_length = int(max_seq_length)
-
-    if sequence_length_to_pad_to is not None:
-        if max_seq_length > sequence_length_to_pad_to:
-            warnings.warn(
-                f"{max_seq_length=} is bigger than the provided {sequence_length_to_pad_to=}, overwriting the padding"
-                f" to {max_seq_length}"
-            )
-        # pad to sequence length or max seq length, whichever is bigger
-        max_seq_length = max(sequence_length_to_pad_to, max_seq_length)
-
-    return torch.nn.functional.pad(tensors_padded, (0, max_seq_length - tensors_padded.size(-1)), value=pad_value)
-
-
 class SyncTimer(NamedTimer):
     """Wrapper around NamedTimer to sync across DP ranks
         for more precise timing
@@ -370,3 +372,52 @@ def is_finished(self):
         # only respect rank 0 timing
         torch.distributed.broadcast(is_finished_tensor, 0)
         return is_finished_tensor.item()
+
+
+def run_distributed_inference(inputs=None, infer_fn=None):
+    tokens, lengths = None, None
+    dp_rank = parallel_state.get_data_parallel_rank()
+    dp_size = parallel_state.get_data_parallel_world_size()
+    is_rank_0 = torch.distributed.get_rank() == 0
+
+    if is_rank_0:
+        tokens = torch.as_tensor(inputs["inputs"], dtype=torch.long, device=torch.cuda.current_device())
+        lengths = torch.as_tensor(inputs["sequence_length"], dtype=torch.long, device=torch.cuda.current_device())
+
+    tokens = broadcast_2d_tensor(tokens, 0, dtype=torch.long, group=None).chunk(dp_size)[dp_rank]
+    lengths = broadcast_2d_tensor(lengths, 0, dtype=torch.long, group=None).chunk(dp_size)[dp_rank].squeeze(-1)
+
+    outputs = infer_fn(inputs=(tokens, lengths))
+
+    if isinstance(outputs, tuple):
+        # rm and critic are combined in this case
+        rewards, values = outputs
+        rewards = rebalance_nd_tensor(rewards, group=parallel_state.get_data_parallel_group()).squeeze(1).cpu().numpy()
+        values = rebalance_nd_tensor(values, group=parallel_state.get_data_parallel_group()).cpu().numpy()
+
+        return rewards, values
+
+    return rebalance_nd_tensor(outputs, group=parallel_state.get_data_parallel_group()).cpu().numpy()
+
+
+def pad_tensors_to_max_global_seq_len(list_of_tensors, pad_value, group, sequence_length_to_pad_to=None):
+    """pad a list of tensors to the global sequence length across the specified group
+    """
+    # compute the local padding
+    tensors_padded = torch.nn.utils.rnn.pad_sequence(list_of_tensors, batch_first=True, padding_value=pad_value)
+
+    # find global max seq length
+    max_seq_length = torch.tensor([tensors_padded.size(-1)], dtype=torch.float32, device=torch.cuda.current_device())
+    torch.distributed.all_reduce(max_seq_length, op=torch.distributed.ReduceOp.MAX, group=group)
+    max_seq_length = int(max_seq_length)
+
+    if sequence_length_to_pad_to is not None:
+        if max_seq_length > sequence_length_to_pad_to:
+            warnings.warn(
+                f"{max_seq_length=} is bigger than the provided {sequence_length_to_pad_to=}, overwriting the padding"
+                f" to {max_seq_length}"
+            )
+        # pad to sequence length or max seq length, whichever is bigger
+        max_seq_length = max(sequence_length_to_pad_to, max_seq_length)
+
+    return torch.nn.functional.pad(tensors_padded, (0, max_seq_length - tensors_padded.size(-1)), value=pad_value)
diff --git a/nemo_aligner/utils/parallel_state.py b/nemo_aligner/utils/parallel_state.py
new file mode 100644
index 000000000..b4b2113df
--- /dev/null
+++ b/nemo_aligner/utils/parallel_state.py
@@ -0,0 +1,31 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  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.
+
+"""Wrapper around mcore parallel state to handle cases of resharding"""
+
+from megatron.core import parallel_state as mcore_parallel_state
+
+from nemo.collections.nlp.modules.common.text_generation_utils import (
+    get_model_parallel_src_rank as nemo_get_model_parallel_src_rank,
+)
+
+
+def get_model_parallel_src_rank():
+    src_rank = nemo_get_model_parallel_src_rank()
+
+    return src_rank
+
+
+def __getattr__(name):
+    return getattr(mcore_parallel_state, name)
diff --git a/nemo_aligner/utils/server_utils.py b/nemo_aligner/utils/server_utils.py
index 929e21489..15d669f7d 100644
--- a/nemo_aligner/utils/server_utils.py
+++ b/nemo_aligner/utils/server_utils.py
@@ -12,10 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from abc import ABC, abstractmethod
 from typing import Optional
 
 import numpy as np
+import torch
+import torch.nn.functional as F
+
+from nemo_aligner.utils import parallel_state
 
 
 def decode_bytes_ndarray(str_ndarray: np.ndarray) -> np.ndarray:
@@ -73,6 +78,59 @@ def pad_input(value: Optional[np.ndarray], size: int, pad_value: int = 0):
     return value, extra
 
 
+def calculate_inference_batch_padding_multiple(current_size, model_forward_micro_batch_size):
+    """calculates the multiple to pad an inference batch up to. If the batch is smaller
+        than the total inference size, then we pad up to a multiple of DP. Otherwise
+        we pad to a multiple of model_forward_mbs * DP
+    """
+    total_inference_size = model_forward_micro_batch_size * parallel_state.get_data_parallel_world_size()
+
+    if current_size <= total_inference_size:
+        return parallel_state.get_data_parallel_world_size()
+
+    return total_inference_size
+
+
+def process_inputs(inputs, tokenize_func):
+    sentences = inputs.get("sentences", None)
+    if sentences is not None:
+        sentences = decode_bytes_ndarray(sentences)
+        tokens, sequence_lengths = tokenize_func(sentences)
+        sequence_lengths = sequence_lengths.unsqueeze(-1)
+    else:
+        tokens = torch.as_tensor(inputs["tokens"], dtype=torch.long, device=torch.cuda.current_device())
+        sequence_lengths = torch.as_tensor(
+            inputs["sequence_lengths"], dtype=torch.long, device=torch.cuda.current_device()
+        )
+
+    return tokens, sequence_lengths
+
+
+def pad_batch_and_strip_sequence(tokens, sequence_lengths, pad_to_multiple, strip_sequence_length_to_multiple=None):
+    prestrip_sequence_length = tokens.shape[1]
+    if strip_sequence_length_to_multiple is not None:
+        stripped_sequence_length = (
+            math.ceil(sequence_lengths.max().item() / strip_sequence_length_to_multiple)
+            * strip_sequence_length_to_multiple
+        )
+        if stripped_sequence_length < prestrip_sequence_length:
+            tokens = tokens[:, :stripped_sequence_length]
+
+    # padding on the batch dim
+    num_extra = tokens.size(0) % pad_to_multiple
+    amount_to_pad = 0 if num_extra == 0 else pad_to_multiple - num_extra
+
+    tokens = F.pad(tokens, (0, 0, 0, amount_to_pad), mode="constant", value=0)
+    sequence_lengths = F.pad(sequence_lengths, (0, 0, 0, amount_to_pad), mode="constant", value=0)
+
+    assert len(tokens) == len(
+        sequence_lengths
+    ), "length of tokens and sequence lengths must be the same, but got {} and {}".format(
+        len(tokens), len(sequence_lengths)
+    )
+    return {"inputs": tokens, "sequence_length": sequence_lengths}, amount_to_pad, prestrip_sequence_length
+
+
 class FutureResult(ABC):
     """Generic class for trainers to wait if an object is an instance of this class
     """
diff --git a/nemo_aligner/utils/text_generation_utils.py b/nemo_aligner/utils/text_generation_utils.py
index ae95f3e34..58089d7ca 100644
--- a/nemo_aligner/utils/text_generation_utils.py
+++ b/nemo_aligner/utils/text_generation_utils.py
@@ -14,14 +14,14 @@
 
 """Utilities for generating text."""
 
+import math
 from typing import Any, List
 
 import torch
-
 from megatron.core import parallel_state
-from nemo.collections.nlp.modules.common.text_generation_strategy import GPTModelTextGenerationStrategy
-from nemo.utils import logging
 
+from nemo.collections.nlp.modules.common.lm_utils import pad_batch
+from nemo.collections.nlp.modules.common.text_generation_strategy import GPTModelTextGenerationStrategy
 from nemo_aligner.utils.distributed import broadcast_2d_tensor_within_pp
 
 
@@ -74,27 +74,8 @@ def get_lengths(self) -> torch.Tensor:
         return lengths.flatten()
 
 
-def pad_batch(batch, pad_id):
-    """batch each element of the batch to be the size of the longest sequence
-    """
-    context_lengths = []
-    max_context_length = max([len(tokens) for tokens in batch])
-    for tokens in batch:
-        context_length = len(tokens)
-        if context_length < max_context_length:
-            tokens.extend([pad_id] * (max_context_length - context_length))
-        context_lengths.append(context_length)
-    return batch, context_lengths
-
-
-def tokenize_batch(tokenizer, sentences, max_len, add_BOS, add_EOS=False):
+def tokenize_batch(sentences, tokenizer, max_len, add_BOS=False, add_EOS=False):
     """convert the sentences into lists of tokens, pad them to the same length, add bos tokens if it is needed
-    Args:
-        sentences (List[str]): list of input sentences in str format.
-        max_len (int): max number of tokens to generate.
-        add_BOS (bool): whether to add the BOS token at the beginning
-    Returns:
-        Tuple[torch.Tensor], the tokenized and padded torch tensor and the token context length tensor.
     """
 
     def tokenize(sentence):
@@ -109,16 +90,10 @@ def tokenize(sentence):
         return output
 
     context_tokens = list(map(tokenize, sentences))
+    max_sequence_length = max(len(x) for x in context_tokens)
 
-    exceeded = [False] * len(context_tokens)
-
-    for i, x in enumerate(context_tokens):
-        if len(x) > max_len:
-            logging.warning(f"max seq len of {max_len} exceeded, chunking")
-            exceeded[i] = True
-
+    context_tokens, context_lengths = pad_batch(context_tokens, tokenizer.eos_id, max_len - max_sequence_length)
     context_tokens = [x[:max_len] for x in context_tokens]
-    context_tokens, context_lengths = pad_batch(context_tokens, tokenizer.eos_id)
     context_tokens_tensor = torch.cuda.LongTensor(context_tokens)
     context_length_tensor = torch.cuda.LongTensor(context_lengths)
-    return context_tokens_tensor, context_length_tensor, exceeded
+    return context_tokens_tensor, context_length_tensor
diff --git a/nemo_aligner/utils/train_script_utils.py b/nemo_aligner/utils/train_script_utils.py
index 86bfef9d9..c6f6f8089 100644
--- a/nemo_aligner/utils/train_script_utils.py
+++ b/nemo_aligner/utils/train_script_utils.py
@@ -72,7 +72,8 @@ def init_distributed(ptl_trainer, ptl_model, use_te=False):
     def dummy():
         return
 
-    ptl_trainer.strategy.launcher.launch(dummy, trainer=ptl_trainer)
+    if ptl_trainer.strategy.launcher is not None:
+        ptl_trainer.strategy.launcher.launch(dummy, trainer=ptl_trainer)
     ptl_trainer.strategy.setup_environment()
 
     if use_te:
@@ -101,6 +102,7 @@ def init_using_ptl(ptl_trainer, ptl_model, train_dataloader, train_ds):
     disable_data_callbacks(ptl_model, train_dataloader, train_ds)
 
     call._call_setup_hook(ptl_trainer)
+    call._call_configure_model(ptl_trainer)
     ptl_trainer.strategy.setup(ptl_trainer)
     call._call_callback_hooks(ptl_trainer, "on_fit_start")
     call._call_lightning_module_hook(ptl_trainer, "on_fit_start")
diff --git a/nemo_aligner/utils/trainer_utils.py b/nemo_aligner/utils/trainer_utils.py
index b6a17f0d6..524763077 100644
--- a/nemo_aligner/utils/trainer_utils.py
+++ b/nemo_aligner/utils/trainer_utils.py
@@ -29,7 +29,7 @@ def compute_num_steps_per_epoch(
 
     num_steps_per_epoch = sampler.total_samples // sampler.global_batch_size
 
-    if limit_train_batches is None or limit_train_batches > 1.0:
+    if limit_train_batches is None or (isinstance(limit_train_batches, float) and limit_train_batches > 1.0):
         limit_train_batches = 1.0
 
     if limit_train_batches >= 0:
diff --git a/nemo_aligner/utils/utils.py b/nemo_aligner/utils/utils.py
index 78bc48b22..c8764a413 100644
--- a/nemo_aligner/utils/utils.py
+++ b/nemo_aligner/utils/utils.py
@@ -19,6 +19,7 @@
 import re
 import tempfile
 from contextlib import contextmanager
+from copy import deepcopy
 from dataclasses import replace
 from functools import partial
 from typing import Iterator, List
@@ -335,6 +336,35 @@ def retrieve_model_state_dict_in_cpu(model, megatron_amp_O2=True):
     return cpu_dict
 
 
+@torch.no_grad()
+def copy_model_states_to_cpu(model, cpu_dict=None, megatron_amp_O2=True, sync=True, alias_non_tensor=False):
+    """This function mutates the cpu_dict object to throw the model states into preallocated tensors(if they exist)
+        for non tensors it will do a deepcopy, unless alias_non_tensor is True
+    """
+    if cpu_dict is None:
+        cpu_dict = {}
+
+    for name, item in model.state_dict().items():
+        if isinstance(item, torch.Tensor):
+            if name not in cpu_dict:
+                cpu_dict[name] = torch.empty(
+                    item.size(), dtype=item.dtype, layout=item.layout, device="cpu", pin_memory=True
+                )
+            cpu_dict[name].copy_(item, non_blocking=sync)
+        elif alias_non_tensor:
+            cpu_dict[name] = item
+        else:
+            cpu_dict[name] = deepcopy(item)
+
+    if megatron_amp_O2:
+        cpu_dict = convert_to_amp_o2_format(cpu_dict)
+
+    if sync:
+        torch.cuda.synchronize()
+
+    return cpu_dict
+
+
 @torch.no_grad()
 def swap_dict(resident_model, cpu_weights, offload_onto_cpu=True, megatron_amp_O2=True):
     """swap the state dict with a specified state dict, and offload the current state dict onto CPU
@@ -384,10 +414,10 @@ def convert_to_amp_o2_format(state_dict):
     """
     new_state_dict = {}
 
-    for key in state_dict.keys():
+    for key, item in state_dict.items():
         if "model.module." not in key:
-            new_key = key.replace("model.", "model.module.", 1)
-            new_state_dict[new_key] = state_dict[key]
+            key = key.replace("model.", "model.module.", 1)
+        new_state_dict[key] = item
 
     return new_state_dict
 


Test patch:
diff --git a/tests/test_cai_utils.py b/tests/test_cai_utils.py
new file mode 100644
index 000000000..b404656bc
--- /dev/null
+++ b/tests/test_cai_utils.py
@@ -0,0 +1,317 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  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 os
+import sys
+
+sys.path.append(os.path.abspath("../examples/nlp/cai"))
+
+from cai_utils import ChatTemplateHelper, PromptTemplate, UserAssistantPromptTemplate
+
+
+def test_extra_id_format_single_user_prompt():
+    extra_id_prompt_template = PromptTemplate(
+        dict(
+            System="<extra_id_0>System\n{MESSAGE}\n",
+            User="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            Assistant="{MESSAGE}\n",
+        ),
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+    m11 = extra_id_prompt_template.format_messages(
+        [{"role": "System", "content": ""}, {"role": "User", "content": "Calculate the sum of 2 and 3."}]
+    )
+    m11_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+"""
+    assert m11 == m11_expected, "extra_id formatting failed"
+
+
+def test_extra_id_format_single_message_chat():
+    extra_id_prompt_template = PromptTemplate(
+        dict(
+            System="<extra_id_0>System\n{MESSAGE}\n",
+            User="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            Assistant="{MESSAGE}\n",
+        ),
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+    m12 = extra_id_prompt_template.format_messages(
+        [
+            {"role": "System", "content": ""},
+            {"role": "User", "content": "Calculate the sum of 2 and 3."},
+            {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+        ]
+    )
+    m12_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+The sum of 2 and 3 is 5.
+"""
+    assert m12 == m12_expected, "extra_id formatting single message chat failed"
+
+    extract_m12 = extra_id_prompt_template.extract_response(m12)
+    extract_m12_expected = "The sum of 2 and 3 is 5."
+    assert extract_m12 == extract_m12_expected, "extra_id formatting single message chat extract failed"
+
+
+def test_extra_id_format_multi_message_chat():
+    extra_id_prompt_template = PromptTemplate(
+        dict(
+            System="<extra_id_0>System\n{MESSAGE}\n",
+            User="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+            Assistant="{MESSAGE}\n",
+        ),
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+
+    m13 = extra_id_prompt_template.format_messages(
+        [
+            {"role": "System", "content": ""},
+            {"role": "User", "content": "Calculate the sum of 2 and 3."},
+            {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+            {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+        ]
+    )
+    m13_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+The sum of 2 and 3 is 5.
+<extra_id_1>User
+Thank you! Could you also calculate the sum of 5 and 7?
+<extra_id_1>Assistant
+"""
+    assert m13 == m13_expected, "extra_id formatting multi message chat failed"
+
+
+# mistral template
+
+
+def test_mistral_format_single_user_prompt():
+    mistral_prompt_template = PromptTemplate(
+        dict(User="
 {MESSAGE} 
", Assistant="{MESSAGE}
 "),
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+
+    m21 = mistral_prompt_template.format_messages({"role": "User", "content": "Calculate the sum of 2 and 3."})
+    m21_expected = "

 Calculate the sum of 2 and 3. 
"
+    assert m21 == m21_expected, "mistral formatting single user prompt failed"
+
+
+def test_mistral_format_single_message_chat():
+    mistral_prompt_template = PromptTemplate(
+        dict(User="
 {MESSAGE} 
", Assistant="{MESSAGE}
 "),
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+    m22 = mistral_prompt_template.format_messages(
+        [
+            {"role": "User", "content": "Calculate the sum of 2 and 3."},
+            {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+        ]
+    )
+    m22_expected = """

 Calculate the sum of 2 and 3. 
The sum of 2 and 3 is 5.
 """
+    assert m22 == m22_expected, "mistral formatting single message chat failed"
+
+    extract_m22 = mistral_prompt_template.extract_response(m22)
+    extract_m22_expected = "The sum of 2 and 3 is 5."
+    assert extract_m22 == extract_m22_expected, "mistral formatting single message chat extract failed"
+
+
+def test_mistral_format_multi_message_chat():
+    mistral_prompt_template = PromptTemplate(
+        dict(User="
 {MESSAGE} 
", Assistant="{MESSAGE}
 "),
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+    m23 = mistral_prompt_template.format_messages(
+        [
+            {"role": "User", "content": "Calculate the sum of 2 and 3."},
+            {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+            {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+        ]
+    )
+    m23_expected = "

 Calculate the sum of 2 and 3. 
The sum of 2 and 3 is 5.
 
 Thank you! Could you also calculate the sum of 5 and 7? 
"
+    assert m23 == m23_expected, "mistral formatting multi message chat failed"
+
+
+# UserAssistantPromptTemplate
+def test_extra_id_user_assistant_format_single_user_prompt():
+    extra_id_user_assistant_format = UserAssistantPromptTemplate(
+        user_format="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+        assistant_format="{MESSAGE}\n",
+        system_format="<extra_id_0>System\n{MESSAGE}\n",
+        system_default_message="",
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+
+    m31 = extra_id_user_assistant_format.format_messages(
+        {"role": UserAssistantPromptTemplate.Role.User, "content": "Calculate the sum of 2 and 3."}
+    )
+    m31_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+"""
+    assert m31 == m31_expected, "extra_id user assistant format single user prompt failed"
+
+
+def test_extra_id_user_assistant_format_single_user_prompt_create_user_message():
+    extra_id_user_assistant_format = UserAssistantPromptTemplate(
+        user_format="<extra_id_1>User\n{MESSAGE}\n<extra_id_1>Assistant\n",
+        assistant_format="{MESSAGE}\n",
+        system_format="<extra_id_0>System\n{MESSAGE}\n",
+        system_default_message="",
+        eos_token="<extra_id_1>",
+        response_extract_pattern="<extra_id_1>Assistant\n",
+    )
+
+    m32 = extra_id_user_assistant_format.format_messages(
+        extra_id_user_assistant_format.create_user_message("Calculate the sum of 2 and 3.")
+    )
+    m32_expected = """<extra_id_0>System
+
+<extra_id_1>User
+Calculate the sum of 2 and 3.
+<extra_id_1>Assistant
+"""
+    assert m32 == m32_expected, "extra_id user assistant format single user prompt create user message failed"
+
+
+def test_mistral_user_assistant_format_single_user_prompt():
+    mistral_user_assistant_format = UserAssistantPromptTemplate(
+        user_format="
 {MESSAGE} 
",
+        assistant_format="{MESSAGE}
 ",
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+
+    m41 = mistral_user_assistant_format.format_messages(
+        {"role": UserAssistantPromptTemplate.Role.User, "content": "Calculate the sum of 2 and 3."}
+    )
+    m41_expected = "

 Calculate the sum of 2 and 3. 
"
+    assert m41 == m41_expected, "mistral user assistant format single user prompt failed"
+
+
+def test_mistral_user_assistant_format_single_user_prompt_create_user():
+    mistral_user_assistant_format = UserAssistantPromptTemplate(
+        user_format="
 {MESSAGE} 
",
+        assistant_format="{MESSAGE}
 ",
+        bos_token="
",
+        eos_token="
",
+        response_extract_pattern="
",
+    )
+
+    m42 = mistral_user_assistant_format.format_messages(
+        mistral_user_assistant_format.create_user_message("Calculate the sum of 2 and 3.")
+    )
+    m42_expected = "

 Calculate the sum of 2 and 3. 
"
+    assert m42 == m42_expected, "mistral user assistant format single user prompt create user failed"
+
+
+def test_collated_single_chat_single_message():
+    single_chat_single_message = [{"role": "User", "content": "Calculate the sum of 2 and 3."}]
+
+    expected_collated_chat_messages = [{"role": ["User"], "content": ["Calculate the sum of 2 and 3."]}]
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages(single_chat_single_message)
+    assert collated_chat_messages == expected_collated_chat_messages, "collating single chat messages failed"
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages([single_chat_single_message])
+    assert (
+        collated_chat_messages == expected_collated_chat_messages
+    ), "collating single chat with single message failed"
+
+
+def test_collated_single_chat():
+    single_chat_messages = [
+        {"role": "User", "content": "Calculate the sum of 2 and 3."},
+        {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+        {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+    ]
+
+    expected_collated_chat_messages = [
+        {"role": ["User"], "content": ["Calculate the sum of 2 and 3."]},
+        {"role": ["Assistant"], "content": ["The sum of 2 and 3 is 5."]},
+        {"role": ["User"], "content": ["Thank you! Could you also calculate the sum of 5 and 7?"]},
+    ]
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages(single_chat_messages)
+    assert collated_chat_messages == expected_collated_chat_messages, "collating single chat messages failed"
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages([single_chat_messages])
+    assert collated_chat_messages == expected_collated_chat_messages, "collating single chat messages failed"
+
+
+def test_collated_multi_chat_messages():
+    chat_1 = [
+        {"role": "User", "content": "Calculate the sum of 2 and 3."},
+        {"role": "Assistant", "content": "The sum of 2 and 3 is 5."},
+        {"role": "User", "content": "Thank you! Could you also calculate the sum of 5 and 7?"},
+    ]
+
+    chat_2 = [
+        {"role": "User", "content": "Calculate the sum of 56 and 21."},
+        {"role": "Assistant", "content": "The sum of 56 and 21 is 77."},
+        {"role": "User", "content": "Thank you! Could you also calculate the sum of 15 and 17?"},
+    ]
+
+    chat_list = [chat_1, chat_2]
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages(chat_list)
+
+    expected_collated_chat_messages = [
+        {"role": ["User", "User"], "content": ["Calculate the sum of 2 and 3.", "Calculate the sum of 56 and 21."]},
+        {"role": ["Assistant", "Assistant"], "content": ["The sum of 2 and 3 is 5.", "The sum of 56 and 21 is 77."]},
+        {
+            "role": ["User", "User"],
+            "content": [
+                "Thank you! Could you also calculate the sum of 5 and 7?",
+                "Thank you! Could you also calculate the sum of 15 and 17?",
+            ],
+        },
+    ]
+
+    assert collated_chat_messages == expected_collated_chat_messages, "collating multi chat messages failed"
+
+
+def test_collated_multi_chat_each_with_single_message():
+    chat_1 = [{"role": "User", "content": "Calculate the sum of 2 and 3."}]
+
+    chat_2 = [{"role": "User", "content": "Calculate the sum of 56 and 21."}]
+
+    expected_collated_chat_messages = [
+        {"role": ["User", "User"], "content": ["Calculate the sum of 2 and 3.", "Calculate the sum of 56 and 21."]},
+    ]
+
+    collated_chat_messages = ChatTemplateHelper.collate_chat_messages([chat_1, chat_2])
+
+    assert collated_chat_messages == expected_collated_chat_messages, "collating multi chat with single message failed"