NeMo-Aligner: Scalable Toolkit for Efficient Model Alignment

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.

Rating

6

Confidence

3

Ethics flag

1

Reasons to accept

- Achieves impressive scalability, enabling alignment of models with over 70B parameters across numerous GPUs. - Incorporates optimizations tailored to the unique compute challenges of alignment techniques like RLHF. - Supports a range of state-of-the-art alignment methods in one extensible framework.

Reasons to reject

- 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.

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.

Rating

7

Confidence

3

Ethics flag

1

Reasons to accept

The construction of infrastructure and engineering frameworks for large language models is crucial as it can promote the progress of the entire research community. The author demonstrated the effectiveness of the framework through experiments and demonstrated its scalability.

Reasons to reject

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 to authors

See the above in the Reasons To Reject.

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.

Rating

7

Confidence

2

Ethics flag

1

Reasons to accept

A useful tool that can run several alignment techniques in a parameter efficient tuning setting and is designed for extensibility.

Reasons to reject

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").

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.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

1. Implements and integrates various alignment techniques commonly used by the community, in one toolkit. 2. Makes alignment techniques (like RLHF) more scalable by incorporating a distributed training approach. 3. Proposes an extensible framework that can allow the addition of new alignment methods (like variants of DPO) with minimal code changes. 4. Includes support for Parameter-efficient training methods for algorithms that are a part of the toolbox like SFT, PPO, DPO, and SteerLM, making it more accessible to the community.

Reasons to reject

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.