Source-linked AI summary
NeMo-Aligner: Scalable Toolkit for Efficient Model Alignment
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
TL;DR
Complex, multi-model alignment pipelines make efficient training challenging, especially for the largest language models. NeMo-Aligner combines scalable implementations of multiple alignment methods with distributed training and inference optimizations. The toolkit reports strong scaling with additional compute and supports models such as Nemotron 4 340B and Llama 3.1 405B.
Problem
Feedback-based alignment can improve helpfulness, but its complex multi-model pipelines create systems challenges that deter productive adoption beyond well-resourced organizations.
Source: §1 · p. 1Method
NeMo-Aligner builds on Megatron-LM and NeMo, combining 3D parallelism, distributed PPO, rollout inference optimizations, and implementations of several alignment methods.
Source: §1 · p. 2, §3 · p. 4, §3.1 · p. 5Results
NeMo-Aligner consistently scales well with more compute, and PPO generation achieved a near-linear speedup of 2.05x.
Source: §3.4 · p. 7, §8 · p. 10Takeaways & Limitations
The Apache 2.0 toolkit aims to make alignment research more efficient and accessible, including LoRA-based support for compute-limited settings.
Source: §8 · p. 10For the 405B model, generation cannot fit on one node and requires pipeline-parallel generation, which the authors plan to optimize further.
Source: §3.6 · p. 9
Abstract
from arXiv · showhide
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
1 Introduction
Feedback-based alignment can improve instruction following relative to supervised fine-tuning alone, but complex multi-model pipelines make efficient adoption difficult. NeMo-Aligner targets this systems challenge with scalable implementations for several alignment approaches and large models.
1 Introduction
- Feedback-based tuning has been more effective than supervised fine-tuning alone at moving responses toward helpfulness and away from less helpful responses.Pretraining on unlabeled text has demonstrated promising capabilities, while instruction alignment is described as critical for practical use.Source: §1 · p. 1
- PPO-based RLHF pipelines can require four interacting large language models, complicating training and limiting productive adoption beyond well-resourced organizations.The resulting systems challenges include model scalability, coordination, and text generation within the training loop.Source: §1 · p. 1
- NeMo-Aligner targets performance and scalability beyond 1,000 GPUs for models including Nemotron 4 340B and Llama 3.1 405B.The paper names these models as examples of the largest and most competent open-source LLMs.Source: §1 · pp. 1–2
- Megatron-LM 3D parallelism, distributed PPO, and TensorRT-LLM rollout inference optimizations address scaling challenges and reduce research iteration time.The approach combines data, tensor, and pipeline parallelism with distributed PPO training and optimized inference during rollout.Source: §1 · p. 2
- NeMo-Aligner optimizes SFT, PPO-based RLHF, DPO, SteerLM, and SPIN.The toolkit is designed to support extensions to other alignment techniques.Source: §1 · p. 2
2 Model Alignment Background
The background describes alignment methods that use supervised responses, preference feedback, attribute conditioning, or self-play. In the RLHF recipe, supervised fine-tuning supplies a starting point for reward-model training and PPO policy optimization.
2 Model Alignment Background
- The RLHF recipe uses annotated prompt-response data for SFT, preference pairs to train a reward model, and the SFT and reward models to initialize PPO policy and value networks.The SFT model also supplies the KL-divergence penalty in PPO.Source: Figure 1 · p. 2
2.1 Supervised Fine Tuning
- SFT updates base-model parameters to mimic expected prompt responses using token-level cross-entropy loss.Expected responses may come from expert human annotations or other language models.Source: §2.1 · p. 2
2.2 Reinforcement Learning from Human Feedback
- RLHF trains a reward model on chosen/rejected preference pairs, then PPO optimizes an SFT-initialized policy with its scores and an SFT-based KL penalty; the critic typically starts from the reward model.The KL penalty helps prevent the policy from straying too far from its starting point and exploiting reward-model blind spots.Source: §2.2 · p. 3
2.3 Direct Preference Optimization
- DPO directly optimizes a policy from offline preference pairs, using a frozen reference policy to derive rewards implicitly instead of training an explicit reward model.The loss uses differences in chosen-versus-rejected log probabilities for the optimal and reference policies.Source: §2.3 · p. 3
2.4 SteerLM
- SteerLM predicts and annotates response attributes, fine-tunes on attribute-conditioned prompts, then uses desired attribute values at inference to steer helpfulness and safety.Attributes can include correctness and toxicity, with higher values indicating more of the attribute.Source: §2.4 · pp. 3–4
2.5 Self-Play Fine-Tuning
- SPIN iteratively trains with DPO loss to prefer human SFT responses over responses generated by the previous frozen reference-policy iteration.At each iteration's end, the reference policy is updated with the trained policy's weights.Source: §2.5 · p. 4
3 RLHF (PPO) Training
NeMo-Aligner builds on Megatron-LM and NeMo to support the RLHF pipeline, focusing on efficiency challenges in PPO training.
3.1 Distributed Approach to PPO training
- PyTriton distributes PPO models across compute allocations; actor/reference and reward/critic weights share jobs, while asynchronous communication permits pipelined inference and training.The reference policy and actor use the same model with different weights; the reference weights are offloaded to CPU and swapped in for inference.Source: §3.1 · p. 5, Figure 3 · p. 5
3.2 Optimizations for PPO rollout
- TensorRT-LLM generates rollouts with an in-memory engine; weights are refit between steps, while optional worker balancing targets unequal response lengths.The engine avoids recompilation after weight updates, and shorter-generation workers can take on more work.Source: §3.2 · pp. 5–6
3.3 Model Training Details and Quality
- GPT-4-Turbo-judged MT-Bench reached 8.13 for PPO-trained Llama 3 70B, versus 7.96 for its SFT checkpoint.The reported training used HelpSteer2 prompts and a rollout global batch size of 128.Source: §3.3 · p. 6
3.4 Scalability
- 1.80x speedup in overall time per step from 8+4 to 16+8 nodes, while a 1024 rollout batch enabled meaningful scaling to 64+32 nodes (768 H100 GPUs).The reported speedup includes gains in both training and rollout stages.Source: Table 1 · p. 6, §3.4 · pp. 6–7, Table 2 · p. 7
3.5 What contributes to system performance?
- TensorRT-LLM removal made PPO steps nearly seven times longer; removing inference-only tensor parallelism, refitting, and async requests incurred 3.87x, 3.15x, and 1.54x slowdowns, respectively.Worker-pool balancing did not meaningfully speed the small 128-batch setup; the authors expect it to help larger jobs.Source: Table 4 · p. 8, §3.5 · p. 8
3.6 Training the largest open source LLM
- PPO on Llama 3.1 405B ran across 1008 H100 GPUs, but generation is slower than for 70B because it requires pipeline parallelism.The 405B model’s generation cannot fit on a single node; its mean response length in the first few PPO steps was 916 tokens versus 351 for 70B.Source: Table 5 · p. 8, §3.6 · pp. 8–9
4 DPO Training
Following the Zephyr-7B-Beta recipe, the authors train an SFT model and then apply DPO to Mistral-7B; both reported MT-Bench scores slightly exceed those of Tunstall et al. (2023).
4 DPO Training
- MT-Bench scores were 7.60 vs 7.34 for the final model and 6.77 vs 6.64 for the SFT-only initial model, slightly above Tunstall et al. (2023).Source: §4 · p. 9
- The recipe first applies SFT to Mistral-7B on Ultrachat, then DPO on Ultrafeedback.Source: §4 · p. 9
- SFT used a 2e-5 learning rate, batch size 512, and three epochs; DPO used batch size 512 and up to 300 steps with a cosine schedule.DPO used KL regularization 3e-4, peak and minimum learning rates of 1e-7 and 1e-8, and 50 warmup steps.Source: §4 · p. 9
5 SteerLM Training with LoRA
LoRA makes SteerLM training more GPU-efficient while retaining comparable MT-Bench performance, and its throughput scales almost proportionally as GPU count increases.
5 SteerLM Training with LoRA
- LoRA reduced the minimum BF16 GPU requirement from 32 to 8 and achieved 5× speedup over full-parameter fine-tuning, with MT-Bench scores of 7.43 vs. 7.54.The cited scores are described as within the benchmark's noise level.Source: §5 · p. 9, Table 6 · p. 9
- Relative throughput in samples per second improved almost proportionally as the number of GPUs increased.The authors report diminishing returns at larger GPU counts.Source: §5 · p. 9
6 SPIN Training
The authors use a one-iteration SPIN recipe with a 50k-sample Ultrachat200k subset and report an MT-Bench score above both the SFT Zephyr-7B-Beta and three-iteration SPIN scores.
6 SPIN Training
- Training used one iteration of one epoch on a random 50k-sample Ultrachat200k subset, without adding previous-iteration generations to the next dataset.The authors used AdamW rather than RMSProp.Source: §6 · pp. 9–10
- MT-Bench was 7.04, above 6.64 for Zephyr-7B-Beta using SFT and 6.78 for the three-iteration SPIN model.Source: §6 · p. 10
7 Framework Extensibility
NeMo-Aligner uses a trainer abstraction to support algorithm changes and reuse across distributed training, enabling DPO variants and other alignment techniques to be incorporated with minimal effort.
7 Framework Extensibility
- A trainer abstraction encourages reuse of existing trainer methods across training steps and approaches.Source: §7 · p. 10
- The framework allows DPO variants, including Identity Preference Optimization, Conservative DPO, and Kahneman-Tversky Optimization, to be integrated with minimal code changes.Source: §7 · p. 10
- The framework design also facilitates incorporation of Constitutional AI, Rejection Sampling, and Self-Rewarding Language Models.Source: §7 · p. 10
8 Conclusion
NeMo-Aligner is an open-source framework for efficient, scalable LLM alignment that supports both large-compute and compute-limited settings.
8 Conclusion
- The framework scales well when training large models with more compute, and the authors expect scaling to improve in future versions.The authors describe this as an initial release.Source: §8 · p. 10
- SFT, PPO, DPO, and SteerLM are supported in parameter-efficient settings using LoRA for compute-limited use.Source: §8 · p. 10
- The Apache 2.0 open-source codebase is intended to make alignment research more efficient and accessible.Source: §8 · p. 10