NeMo-Aligner: Scalable Toolkit for Efficient Model Alignment

Authors: gerald  shen
zhilin  wang
olivier  delalleau
jiaqi  zeng
yi  dong
daniel  egert
shengyang  sun
jimmy  zhang
sahil  jain
ali  taghibakhshi
markel  sanz ausin
ashwath  aithal
oleksii  kuchaiev
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 19:44:15

Abstract: Aligning Large Language Models (LLMs) with human values and preferences is essential for making them helpful and safe.
However, building efficient tools to perform alignment can be challenging, especially for the largest and most competent LLMs which often contain tens or hundreds of billions of parameters. 
We create NeMo-Aligner, a toolkit for model alignment that can efficiently scale to a thousand GPUs for training the largest open-source LLMs such as Nemotron 4 340B and Llama 3.1 405B. 
NeMo-Aligner comes with highly optimized and scalable implementations for major paradigms of model alignment such as: Reinforcement Learning from Human Feedback (RLHF), Direct Preference Optimization (DPO), SteerLM,  and Self-Play Fine-Tuning (SPIN).
Additionally, our toolkit supports running most of the alignment techniques in a Parameter Efficient Fine-Tuning (PEFT) setting.
NeMo-Aligner is designed for extensibility, allowing support for other alignment techniques with minimal effort.
It is open-sourced with Apache 2.0 License and we invite community contributions at https://github.com/NVIDIA/NeMo-Aligner.

Main paper: 
 introduction: 
 Pre-training large language models on tremendous amounts of unlabelled text has showcased promising capabilities (Brown et al., 2020; Zhang et al., 2022) . While such unsupervised pre-trained models have achieved impressive results, subsequently aligning models to follow user instructions is a critical step to tap the capabilities of LLMs for practical use cases (Sanh et al., 2022; Wei et al., 2022) . Attempts based on Supervised Finetuning (Conover et al., 2023; K öpf et al., 2023; Taori et al., 2023) proved less effective compared to techniques that also made use of feedback to tune models towards responses that are more helpful and away from responses that are less so (Bai et al., 2022a; Ouyang et al., 2022; Touvron et al., 2023; Dong et al., 2023) .Despite the benefits of training models using feedback, these pipelines are notoriously challenging to get right (Lambert & Calandra, 2023; Zheng et al., 2023b) , deterring widespread, productive adoption outside of select well-resourced organizations. For example, the popular Proximal Policy Optimization (PPO) variant of Reinforcement Learning from Human Feedback (RLHF) approach (Ouyang et al., 2022) requires running a complicated pipeline with four large language models interacting in a complex manner during training. Such alignment algorithms introduce new system challenges for efficient training that require re-thinking various aspects of the software stack including model scalability, coordination among models, and text generation within the training loop.There are existing open source tools for model alignment, most notably HuggingFace TRL (von Werra et al., 2020) , CarperAI trlX (Havrilla et al., 2023) and Microsoft DeepSpeed-Chat (Yao et al., 2023) . These tools provide an excellent starting point with respect to usability and feature set. However, with NeMo-Aligner we aim to vastly improve performance and scalability to more than a thousand GPUs, especially useful for aligning the largest and most Ouyang et al. (2022) .Step 1: Annotated Prompt-Response Data is used to perform Supervised Fine-Tuning on the pre-trained (base) Model.Step 2: The resulting SFT model is trained with Preference Data to produce a Reward Model.Step 3: The SFT Model is used to initialize the Policy Network, and the Reward Model is used to initialize the Value Network -together with input prompts, all four models are used to train a Policy Model. The SFT model is also used to compute the KL divergence penalty in Step 3 (not illustrated). 
supervised fine tuning: 
 Given a pre-trained (also referred to as "base") model, supervised fine-tuning (SFT) updates the base model's parameters on prompts with expected responses, where the expected responses might come from expert human annotations (K öpf et al., 2023) or other language models (Ding et al., 2023) . The model is trained to mimic the expected responses given prompts using the token-level cross-entropy loss. SFT is an important prerequisite step in Reinforcement Learning from Human Feedback (Ouyang et al., 2022) and Direct Preference Optimization (Rafailov et al., 2023) because without it, the base model is very unlikely to generate responses which follow user's instructions. This step is also sometimes called behavior cloning because the model is expected to mimic responses of a human or another model.Reinforcement Learning from Human Feedback (RLHF) was introduced by Christiano et al.(2017) as a way to avoid manually defined reward functions in Reinforcement Learning.Instead, a reward model is trained from a dataset of human preferences consisting of pairs of "chosen" and "rejected" trajectories. The reward model's loss, derived from the Bradley-Terry model (Bradley & Terry, 1952) , tries to maximize the likelihood that r chosen > r rejected (i.e., that the predicted rewards are consistent with human preferences). Once the reward model is trained, it may be used to compute rewards for RL algorithm. Two most common methods used in RLHF are REINFORCE (Williams, 1992) and Proximal Policy Optimization (PPO) (Schulman et al., 2017) . In NeMo-Aligner we focus on PPO, specifically as described by Ouyang et al. (2022) .RLHF has been shown to bring significant benefits for model alignment (Ouyang et al., 2022; Bai et al., 2022a; Touvron et al., 2023) with the typical training recipe being as follows, also illustrated in Figure 1 :1. From a pre-trained base model, train an initial SFT model as described in Section 2.1.2. From the SFT model, train a reward model using a dataset of human preferences made of pairs of "chosen" and "rejected" responses to a set of prompts, following Christiano et al. (2017) . Typically, we initialize a linear reward model head on top of the SFT model before training.3. From the SFT model, train a policy with the online Proximal Policy Optimization algorithm (PPO, Schulman et al., 2017) , with rewards provided by the trained reward model. Input prompts may not necessarily be the same as those used for reward model training. A regularization term based on the KL divergence w.r.t. the SFT model helps prevent the policy from straying too far away from its starting point and exploiting the "blind spots" of the reward model (Stiennon et al., 2020; Ouyang et al., 2022) . The PPO critic is typically initialized from the reward model.Direct Preference Optimization (Rafailov et al., 2023) is an offline, off-policy algorithm that makes use of preference data to directly train an optimal policy without an explicit reward model. Rather than use a reward model, a reference policy is used to implicitly derive the reward between a chosen and rejected pair via the Bradley-Terry model. This is accomplished via the difference in the log probabilities between the chosen and rejected responses, which is calculated for the optimal and reference policies. This difference is scaled and then transformed by the sigmoid function to derive the loss. The reference policy is frozen during training and represents the policy used to generate the chosen/rejected responses to the given prompts. If the reference policy used to generate the preference data is not available, it can be approximated by supervised fine-tuning on the prompts and preferred responses of the preference data.SteerLM (Dong et al., 2023) is a model alignment algorithm based on supervised finetuning which avoids use of complex RL methods, similarly to DPO. SteerLM involves three steps. The first step is to train an Attribute Prediction Model that learns to predict the values (between 0 and 4 where higher is more) for various semantic aspects of a response that make responses helpful and safe, such as its correctness and toxicity (K öpf et al., 2023; Wang et al., 2023) . Next, the Attribution Prediction Model can be used to annotate the various attributes contributing to helpfulness and safety in a diversity of promptresponse datasets. Finally, these annotated datasets can be used to perform Attribute-Conditioned Supervised Fine-Tuning where the model learns to generate the response conditioned on the prompt as well as the annotated attributes formatted into a string, such as helpfulness:4,correctness:4,toxicity:0. This step teaches the model to discriminate between responses that are more helpful/safe and those that are less, in a fine-grained manner for each semantic aspect. At inference time, the prompt can be appended with the optimal attribute values, as above, to generate the most helpful response.Self-Play Fine-Tuning (SPIN) (Chen et al., 2024 ) is a self-play based algorithm, where a strong model is developed from a weaker model by playing against previous instances of itself. Starting from an SFT dataset of prompt/response pairs, new responses are generated from previous iterations of the model. Its policy is then improved by discriminating between these self-generated responses and the ground truth human-generated SFT responses. This is accomplished through a preference loss function which is identical to the one used by DPO (Section 2.3). When SPIN training first starts, we use a copy of the initial policy as the reference policy in the DPO loss. The self-play "game" is then played for a number of iterations during which we train the policy as in DPO whilst keeping the reference policy frozen, and at the end of each iteration we update the reference policy's weights with those from the trained policy. During each iteration, we iterate over our SFT training dataset and use the reference policy to generate responses for each prompt, building a preference tuple between the ground truth SFT human "chosen" response and the generated "rejected" response. Once we have these preference tuples for the entire epoch, we update the model weights via the DPO loss function from these tuples of "(chosen, rejected)" preference pairs. The model thus implicitly learns to prefer the ground truth SFT responses to those generated by the previous iteration of itself, which forms the self-play mechanism. 
rlhf (ppo) training: 
 NeMo-Aligner is designed to support numerous alignment techniques efficiently at extremely large scales. It does so by building upon Megatron-LM (Shoeybi et al., 2020) and NeMo (Kuchaiev et al., 2019) The PPO stage requires running training and/or inference on four different models, as illustrated in Figure 3 : On subsequent training steps, the engine must be synced with updated parameter weights from the training stage. The engine is updated in-place using the TensorRT Refitter (NVIDIA, 2023c). We avoid recompiling the engine which would incur a large overhead as generation cannot begin until the weights are updated.For large problem sizes, there may be discrepancies of generation time between the fastest and the slowest data parallel worker during generation due to the differences in response lengths. To mitigate this, we allow users to setup a worker pool to dynamically load balance among data parallel workers to give workers with shorter generations correspondingly more work.As a demonstration of practical large-scale RLHF training with NeMo-Aligner, we train Llama3 (Meta AI, 2024) 70B model using PPO as prescribed in Wang et al. (2024) The scaling of Train stage is sublinear due to number of micro-batches per data parallel rank decreasing as node count increases. Because all pipeline stages must complete before the optimizer is called in pipeline parallel models, we incur an overhead to fill and drain the pipeline that is independent of the number of micro-batches (Shoeybi et al., 2020) . Therefore, decreasing the number of micro-batches per data parallel rank increases the proportion of the train step spent in filling and draining the pipeline, where GPU utilization is poor. A similar issue is apparent during the log prob calculation phase with scaling of 1.43x.No 3 .Generation time scales well with the number of nodes, achieving a near-linear speedup of 2.05x. This is because scaling up the number of actor nodes proportionally increases the number of data parallel workers for each step, which can evenly share the intense work of generation. However, the time spent fitting weights into the TensorRT engine remains relatively constant and therefore does not scale well with node count. Finally, async communications between the Actor and the Critic models result in the additional time taken to wait for the Critic model to be inconsequential (0.01 seconds), suggesting the effectiveness of having async non-blocking calls between actor and critic models in the PPO pipeline. System scalability also needs to be considered under the context of the problem requirements.The training setup in Section 3.3 has a 70B Llama 3 Actor, 70B Llama 3 Critic as well as a rollout global batch size of 128. Such a setup limits the effective demonstration of our system scaling beyond 16 + 8 nodes as there is not enough work to be meaningfully shared across more data parallel workers. Therefore, we modify the setup slightly to use a rollout global batch size of 1024 in Table 2 in order to measure the system performance when the requirements are higher. Table 2 shows that the increased requirements of the training job allows it to meaningfully scale to 64 + 32 nodes (with 768 H100 GPUs total) for various stages within PPO. To better understand the importance of each aspect of NeMo-Aligner's PPO system design, we conduct ablation studies by removing one aspect at a time and measuring the overall time per step as shown in Table 4 . We find that TensorRT-LLM Integration is the most critical component for high system performance, without which PPO will take nearly seven times as long for each step. This is followed by resharding our model to use tensor parallelism only during inference (3.87x), using TensorRT Refit to avoid TensorRT-LLM engine recompiling (3.15x), the use of async requests between actor and critic models (1.54x). We did not observe meaningful speedup in using the worker pool to balance the work between data parallelism ranks because the problem size is small (with rollout global batch size of 128) and therefore the worker imbalance is less than the overhead we incur doing the balancing itself. Nevertheless, we expect this feature to help for larger problem sizes. (Nvidia et al., 2024) and Llama 3.1 405B (Dubey et al., 2024) . In Table 5 , we perform PPO on Llama 3.1 405B using 1008 H100 GPUs, using a configuration based on Sec. 3.3. We use Llama 3.1 405B Instruct as the actor and we train a Reward Model on top of Llama 3.1 405B Instruct using the same data and hyper-parameters as Nemotron 4 340B Reward (Wang et al., 2024) . Compared to the 70B model, the 405B model is substantially slower to train during PPO, mainly bottle-necked by the slow response generation stage.This is because the 405B model generation cannot fit into a single node and requires pipelineparallel generation which we plan to optimize further in future work. Additionally, in the first few steps of PPO, the 405B model generates longer responses (mean of 916 tokens) than the 70B model with a mean of 351 tokens. 
dpo training: 
 We follow the Zephyr-7B-Beta (Tunstall et al., 2023) training recipe, a model trained with SFT and DPO. Briefly, SFT was first performed on Mistral-7B (Jiang et al., 2023) using the Ultrachat dataset (Ding et al., 2023) . Model was then further trained with DPO using the Ultrafeedback dataset (Cui et al., 2023) . For SFT, we used a constant learning rate of 2e - As we increase the number of GPUs used for LoRA training, the relative throughput (measured in samples per second) improves almost proportionally, as shown in Figure 4 . This shows that NeMo-Aligner can effectively distribute and parallelize the workload across a large number of GPUs with minimal overhead and diminishing returns. 
spin training: 
 We recreate the Zephyr-7B-Beta (Tunstall et al., 2023) SFT model via SPIN instead of SFT as formulated by Chen et al. (2024) . We start with the Mistral-7B base model (Jiang et al., 2023) and perform SPIN training following Chen et al. (2024) . However, we make a few departures from their methodology, in that we do not inject generations from the previous iteration into the current iteration (which would double the dataset size every epoch), and we only train for a single iteration, with 1 epoch per iteration. Additionally, we use a random subset of only 50k samples from Ultrachat200k (Ding et al., 2023) rather than the entire dataset, and use AdamW instead of RMSProp. Our learning rate is 5e -7 with 400 total steps, 40 warmup steps, and this LR is then decayed to 1e -7 for the last 100 steps using cosine annealing. Global batch size is 64, weight decay is 0.0, and the KL regularization coefficient is 0.1, as per Chen et al. (2024) . Using this approach, we achieve an MT-Bench score of 7.04 which exceeds the 6.64 of Zephyr-7B-Beta using SFT (Tunstall et al., 2023) , as well as the 6.78 of the 3-iteration SPIN model (Chen et al., 2024) . 
framework extensibility: 
 We design NeMo-Aligner with extensibility in mind, allowing users to easily modify algorithms in spite of the complexities of distributed training. We do so using the trainer abstraction, which encourages re-use of existing trainer methods across various steps and approaches. The extensibility of NeMo-Aligner allows variants of DPO to be integrated with minimal code changes, including the Identity Preference Optimization (Azar et al., 2023) , the Conservative DPO (Mitchell, 2023) , and the Kahneman-Tversky Optimization (Ethayarajh et al., 2023) . Furthermore, other model alignment techniques such as Constitutional AI (Bai et al., 2022b) , Rejection Sampling (Touvron et al., 2023) , and Self-Rewarding Language Models (Yuan et al., 2024) are also being incorporated into NeMo-Aligner, facilitated by the framework design. 
conclusion: 
 Modern model alignment techniques, especially those based on Reinforcement Learning, pose complex optimization challenges with respect to system implementation. We create and open-source NeMo-Aligner to allow AI researchers and practitioners to efficiently experiment with LLM alignment by utilizing all available compute in a scalable way. Our framework consistently scales well when training large models with more compute. As this is our initial release, we expect this scaling to only improve with future versions. Additionally, we support SFT, PPO, DPO, SteerLM in a parameter-efficient manner using LoRA for compute-limited settings. As an Apache 2.0 licensed open-source codebase, NeMo-Aligner can make alignment research more efficient and accessible. 
 
 Appendix:  


Decision: This paper proposes a framework for scaling up several post-training algorithms, including PPO, DPO, SteerLM, and self-play fine-tuning. The paper shows that the framework supports up to several hundred GPUs, with near-linear speedup in the number of GPUs. The reviewers generally seem to feel that the framework will provide useful and extensible infrastructure. A potential weakness was lack of clear comparison with prior work; the authors argue this was not possible because prior work did not attempt to evaluate at this scale. One response (https://openreview.net/forum?id=yK2eGE8QVW&noteId=rZrkVguEso) adds important details about the SFT stage, which should be included in the next version.

[comments from the PCs] Please follow up on the AC recommendations, especially with details about the SFT stage.

Reviews: review: paper_summary: paper_summary: The authors propose a toolbox that implements alignment techniques such as Reinforcement Learning from Human Feedback (RLHF), Direct Preference Optimization (DPO), SteerLM, and Self-Play Fine-Tuning (SPIN). The approach builds upon existing frameworks such as HuggingFace TRL, CarperAI trlX, and Microsoft DeepSpeed-Chat. The toolbox builds upon Megatron-LM with 3D (data, tensor, and pipeline)-parallelism training, introduces a distributed approach to Proximal Policy Optimization (PPO) training in RLHF and integrates PPO inference optimizations based on TensorRT-LLM.
strength_weakness: strength_weakness: 1. The toolbox doesn’t bring significant novelty in terms of contribution and combines already existing implementations and optimizations proposed by the community. I am not sure how adaptable this will be given the massive user base and support community of existing framrworks.
2. The paper doesn’t present a clear motivation for the creation of a framework when well-supported libraries for each alignment technique exist.
3. The paper doesn’t discuss the evaluation and testing components in detail, which can be a significant factor for adoption for the community.
4. Authors don't justify their model choices for benchmarking their suite and don't include a comparison of their performance with existing alignment frameworks.
paper_summary: paper_summary: This work presents an ope-sourced toolkit that implements several major paradigms of model alignment such as: Reinforcement Learning from Human Feed-back, Direct Preference Optimization, SteerLM, and Self-Play Fine-Tuning. The training pipeline is divided in three steps: supervised fine-tuning, reward model training and proximal policy optimization. The toolkit focus on improving performance, the most critical components being TensorRT-LLM integration to mitigate the bottleneck of response generation during the rollout step; using TensorRT Refit to avoid TensorRT-LLM engine recompiling; using async requests between actor and critic models; and the load-balancing of data parallel workers during generation using a worker pool.
strength_weakness: strength_weakness: Instead of being part of section 3, a new section should be created for the evaluation, extending the level of  detail of the experimental setup (e.g., the only description of the training dataset is "using a mixture of public and proprietary data").
paper_summary: paper_summary: This paper builds an effective tool (STEMA) for aligning large-scale language models. By utilizing technologies such as Megatron LM and TensorRT LM, they have built implementations of popular RLHF methods, including PPO, DPO, SteerLM, and SPIN. The preliminary experimental results on MT-Bench show that it can make the alignment of LLMs and achieve better results through this framework. The paper also elaborates on analyzing the operational details of the entire framework and emphasizes its scalability.
strength_weakness: strength_weakness: Lack of comparison with existing technology. The author mentioned existing frameworks such as TRL, trlX, and DeepSpeedChat in the Introduction and claimed that “However, with STEMA we aim to vastly improve performance and scalability of PPO and other methods especially when it comes to the largest and most competent models such as Llama 2 70B (Touvron et al., 2023) and beyond.” Do these exsiting works cannot optimize the Llama 2 70B? Or do they have some shortcomings when optimizing? More time-consuming, higher-cost, or higher requirements for equipment? If there are some comparisons between STEMA and other frameworks to show the strongness of STEMA, I will accept it clearly.

Another reason why I did not provide a clear acceptance is that the paper did not propose technological innovation and relied on the combination of existing technologies. They have made engineering contributions but lack deeper scientific contributions.
questions: questions: See the above in the Reasons To Reject.
paper_summary: paper_summary: This paper introduces STEMA, a scalable toolkit for efficiently aligning large language models (LLMs) using techniques like Reinforcement Learning from Human Feedback (RLHF), Direct Preference Optimization (DPO), SteerLM, and Self-Play Fine-Tuning (SPIN). STEMA builds upon the Megatron-LM framework and incorporates optimizations to enable efficient training across hundreds of GPUs. Key innovations include a distributed approach to Proximal Policy Optimization (PPO) training in RLHF, integration with TensorRT-LLM for faster generation during rollouts, and support for parameter-efficient fine-tuning using methods like LoRA. The authors demonstrate STEMA's scalability and performance gains through experiments aligning Llama 2 models up to 70B parameters.
strength_weakness: strength_weakness: - While supporting multiple alignment techniques is a strength, the paper could benefit from a more in-depth comparison of the pros and cons of the different frameworks, such as TRL. Especially the pain points in the previous framework should be discussed more to understand how the suggested framework improves or mitigates the points.
- The model training details were too simplified. For example, the training dataset mixture might be a more important factor for the final performance. Thus, the authors should present the details of the dataset mixture. Furthermore, in the same sense, the SFT model should also be compared to the final PPO model and LLama-70B-Chat.

Comments: title: Thank you for your reply
comment: Thank you for the responses. I will keep my score.
Thank you for your response. I have increased my score given your replies which I believe are justified. I request you also to make these clarifications in the next version.
Thanks for your clarifications.
Thank you for your reply. I think my score is good enough, and I will keep my score.