Harnessing Hundreds of GPU Power: NVIDIA’s NeMo-Aligner Unleashes Potential for Large Model Alignment

Ensuring that Large Language Models (LLMs) align with human values and preferences is crucial for their utility and safety. Yet, devising effective tools for this alignment presents significant challenges, particularly with the largest and most sophisticated LLMs, which often boast tens or hundreds of billions of parameters.

In a new paper NeMo-Aligner: Scalable Toolkit for Efficient Model Alignment, a team of researchers from Nvidia introduces NeMo-Aligner, a toolkit designed for large-scale LLM model alignment that can efficiently harness the power of hundreds of GPUs for training.

Aligning models to adhere to user instructions represents a pivotal step in harnessing the potential of LLMs for practical applications. One promising approach, exemplified by Proximal Policy Optimization (PPO), involves using feedback to refine models towards desired responses. However, mastering this approach proves notoriously challenging, hindering widespread and productive adoption beyond a few well-resourced organizations.

The objective of this research is to significantly enhance the performance and scalability of PPO and other methods, particularly for the largest and most advanced models like Llama 2 70B and beyond. The proposed NeMo-Aligner tackles scalability hurdles through several strategies:

• Firstly, by leveraging Megatron-LM’s 3D (data, tensor, and pipeline) parallelism training.
• Secondly, by adopting a distributed approach to PPO training in Reinforcement Learning from Human Feedback (RLHF).
• Thirdly, by integrating PPO inference optimizations based on TensorRT-LLM during the rollout stage.

These optimizations collectively enable users to efficiently train the largest models across hundreds of GPUs, significantly reducing research iteration time.

NeMo-Aligner optimizes various alignment techniques, including Supervised Finetuning (SFT), PPO-based RLHF, Direct Preference Optimization, SteerLM, and Self-Play Fine-Tuning. Additionally, it facilitates running most of these techniques in a Parameter Efficient Fine-Tuning (PEFT) setting.

Consistently, the framework demonstrates excellent scalability when training large models with increased computational resources. Moreover, it is open-sourced under the Apache 2.0 License, welcoming community contributions at https://github.com/NVIDIA/NeMo-Aligner.

The paper NeMo-Aligner: Scalable Toolkit for Efficient Model Alignment is on arXiv.

Author: Hecate He | Editor: Chain Zhang

We know you don’t want to miss any news or research breakthroughs. Subscribe to our popular newsletter Synced Global AI Weekly to get weekly AI updates.