Source-linked AI summary
DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales
Zhewei Yao, Reza Yazdani Aminabadi, Olatunji Ruwase, Samyam Rajbhandari, Xiaoxia Wu, Ammar Ahmad Awan, Jeff Rasley, Minjia Zhang, Conglong Li, Connor Holmes, Zhongzhu Zhou, Michael Wyatt, Molly Smith, Lev Kurilenko, Heyang Qin, Masahiro Tanaka, Shuai Che, Shuaiwen Leon Song, Yuxiong He
TL;DR
DeepSpeed-Chat addresses the lack of an accessible, efficient, and affordable end-to-end RLHF pipeline for large ChatGPT-like models. It combines an InstructGPT-aligned three-stage pipeline with a unified training-and-inference system, reporting faster, scalable, and more accessible training across model sizes.
Problem
Existing systems lack an accessible, efficient, and affordable end-to-end RLHF pipeline for powerful ChatGPT-like models, especially at billion-parameter scale.
Method
DeepSpeed-Chat combines an InstructGPT-style SFT, reward-model, and RLHF pipeline with DeepSpeed-HE, which unifies optimized training and inference.
Results
DeepSpeed-HE is over 15x faster than existing systems, trains OPT-13B in 9 hours and OPT-30B in 18 hours, and supports 175B training in under a day on 64 GPUs.
Takeaways & Limitations
DeepSpeed-Chat makes large-scale RLHF training more accessible across multi-GPU clusters and single-GPU settings.
Takeaways & Limitations
The reported numbers in Tables 1 and 2 apply to Step 3 of training.
Abstract
from arXiv · showhide
ChatGPT-like models have revolutionized various applications in artificial intelligence, from summarization and coding to translation, matching or even surpassing human performance. However, the current landscape lacks an accessible, efficient, and cost-effective end-to-end RLHF (Reinforcement Learning with Human Feedback) training pipeline for these powerful models, particularly when training at the scale of billions of parameters. This paper introduces DeepSpeed-Chat, a novel system that democratizes RLHF training, making it accessible to the AI community. DeepSpeed-Chat offers three key capabilities: an easy-to-use training and inference experience for ChatGPT-like models, a DeepSpeed-RLHF pipeline that replicates the training pipeline from InstructGPT, and a robust DeepSpeed-RLHF system that combines various optimizations for training and inference in a unified way. The system delivers unparalleled efficiency and scalability, enabling training of models with hundreds of billions of parameters in record time and at a fraction of the cost. With this development, DeepSpeed-Chat paves the way for broader access to advanced RLHF training, even for data scientists with limited resources, thereby fostering innovation and further development in the field of AI.
1 Overview
DeepSpeed-Chat addresses the lack of an accessible, efficient, affordable end-to-end RLHF pipeline for large ChatGPT-like models. It combines an InstructGPT-aligned pipeline with unified training and inference optimizations, delivering scalable and accessible training.
- 1 Overview: DeepSpeed-Chat targets the missing end-to-end RLHF system for powerful ChatGPT-like models that is easy, fast, and affordable at large scale.Existing systems require expensive multi-GPU setups, often use resources inefficiently, and cannot readily support models with hundreds of billions of parameters.
- 1 Overview: Its single-script experience runs all three InstructGPT stages and provides an inference API for conversation-style testing.The system accepts a pretrained Hugging Face model and produces a ChatGPT-like model through DeepSpeed-RLHF.
- 1 Overview: DeepSpeed-RLHF replicates InstructGPT’s three-step pipeline: supervised fine-tuning, reward-model fine-tuning, and RLHF.It also supports data abstraction and blending across multiple data sources.
- 1 Overview: DeepSpeed-HE unifies DeepSpeed training and inference, switching between modes to combine tensor parallelism and optimized generation with ZeRO- and LoRA-based RL memory strategies.Pipeline awareness lets the system manage memory and data movement across RLHF phases.
- 1 Overview: Over 15x faster, DeepSpeed-HE trains OPT-13B in 9 hours and OPT-30B in 18 hours for under $300 and $600, respectively, on Azure Cloud.The cited comparison is reported for single-node 8x A100 training.
- 1 Overview: 13B models train in 1.25 hours and 175B models in under a day on 64 GPUs, while models over 13B parameters can be trained on one GPU.The overview notes that the table values refer specifically to Step 3 of training.
2 Easy-to-use ChatGPT Training and Inference Experience
DeepSpeed-Chat provides a simple path from pretrained models to deployed ChatGPT-style systems, while supporting different model sizes, hardware configurations, and custom RLHF strategies. Its examples span single-GPU experiments through multi-node training and expose APIs for research exploration.
- 2.1 Training your first ChatGPT-Style model is so easy with DeepSpeed-Chat’s RLHF examples: A single script completes all three RLHF stages and generates a final ChatGPT-style model from pretrained OPT-13B and OPT-350M actor and reward models.The example produces a 13B model and includes a benchmark recipe based on 135M total tokens.
- 2.1 Training your first ChatGPT-Style model is so easy with DeepSpeed-Chat’s RLHF examples: A 13B model is fully trained in about half a day, with Table 4 breaking down the time across the three training steps.The result is reported for a single DGX node with 8 NVIDIA A100-40G GPUs.
- 2.2 Want to try different model sizes and configurations? You got it!: The same script supports larger configurations: a 66B ChatGPT model is ready within 9 hours using 64 GPUs.The example uses OPT-66B with an OPT-350M reward model in a multi-node deployment.
- 2.2 Want to try different model sizes and configurations? You got it!: A 1.3B training example runs on a single consumer-grade GPU, producing a checkpoint in roughly a 1–2 hour break.The example is intended as a small test of the framework on commodity hardware.
- 2.3 Customizing your own RLHF training pipeline using DeepSpeed-Chat’s RLHF APIs: DeepSpeed-Chat’s flexible APIs let users reconstruct their own RLHF training strategies and create a general backend for varied RLHF algorithms.The API examples expose engine creation, PPO training, experience generation, and actor/critic losses.
3 Full-fledged RLHF Training Pipeline
DeepSpeed-Chat implements a complete InstructGPT-style RLHF pipeline with optional features for data handling and quality retention. The pipeline covers supervised fine-tuning, reward modeling, and PPO-based RLHF training.
- 3 Full-fledged RLHF Training Pipeline: The pipeline follows InstructGPT and provides a full end-to-end workflow, with data support for training from multiple resources.Figure 1 illustrates the pipeline and its optional features.
- 3 Full-fledged RLHF Training Pipeline: Step 1 supervised fine-tunes a pretrained language model on selected human responses, while Step 2 trains a reward model from human rankings of multiple answers.The reward model is usually smaller than the supervised fine-tuning model.
- 3 Full-fledged RLHF Training Pipeline: Step 3 further fine-tunes the supervised model with reward feedback using PPO.This is the RLHF stage of the three-step pipeline.
- 3 Full-fledged RLHF Training Pipeline: EMA checkpoint collection allows a selected exponential-moving-average checkpoint to be used for final evaluation.The feature is included as an optional Step 3 quality mechanism.
- 3 Full-fledged RLHF Training Pipeline: Mixture Training combines next-word prediction with PPO to help prevent regression on public benchmarks such as SQuAD2.0.The paper relates this feature to retaining pretraining benchmark-solving ability.
- 3 Full-fledged RLHF Training Pipeline: An abstraction layer unifies dataset formats, while splitting and blending distribute multiple datasets across the three training stages.The pipeline also demonstrates multi-round conversation quality in the experience section.
4 DeepSpeed Hybrid Engine – Unified Infrastructure to Power and Optimize RLHF Training
Stage 3 of RLHF is costly because it alternates experience-generation inference with model training while maintaining multiple model copies and added features. DeepSpeed Hybrid Engine unifies specialized training and inference optimizations, reconfiguring partitioning and memory management to improve throughput and scale.
- RLHF Stage 3 challenges: Stage 3 alternates token-generation inference with weight updates, creating major memory costs from multiple model copies and runtime costs dominated by generation.EMA collection and Mixture Training add further memory and training costs.
- Hybrid Engine: Hybrid Engine unifies DeepSpeed training and inference engines, switching seamlessly between optimized generation/evaluation and training modes for faster Stage 3 RLHF.The actor model’s train and eval modes trigger different optimizations, improving overall system throughput.
- Inference optimization: During inference, Hybrid Engine combines lightweight KV-cache and intermediate-result management with optimized kernels and tensor parallelism to increase token-generation throughput.
- Training optimization: During training, compatible ZeRO and LoRA memory optimizations can be composed within Hybrid Engine to improve training efficiency.
- Dynamic resource management: Hybrid Engine changes model partitioning and memory configuration between modes, enabling tensor-parallel inference, ZeRO-sharded training, larger batches, and fewer allocation bottlenecks.
5 DeepSpeed RLHF: Unparalleled Scale and Efficiency via Hybrid Engine
DeepSpeed-HE unifies optimized training and inference to make RLHF efficient across model sizes and hardware scales. It improves throughput and model scalability while supporting large-scale training and more accessible single-GPU use.
- 5.1 Capability Recap: DeepSpeed-HE trains OPT-13B in 1.25 hours and 175B models in under a day on a 64-GPU cluster, while supporting models over 13B parameters on a single GPU.These results combine scalability, speed, and accessibility across the reported settings.
- 5.2 Throughput and Model Size Scalability Comparisons with Existing RLHF Systems: DeepSpeed-HE provides 6–19x speedups over Colossal-AI and 1.4–10.5x over HuggingFace DDP on multi-GPU RLHF training.On a single GPU, it also achieves over 10x higher throughput than the compared systems.
- 5.2 Throughput and Model Size Scalability Comparisons with Existing RLHF Systems: DeepSpeed-HE runs 6.5B and 50B models where Colossal-AI runs 1.3B and 6.7B on comparable hardware, supporting models up to 7.5x larger.The comparison is reported for a single GPU and a single A100 40G node, respectively.
- 5.3 Effective Throughput and Scalability Analysis: Generation dominates RLHF iteration time despite comprising about 20% of computation, so DeepSpeed-HE uses optimized inference kernels and tensor parallelism to accelerate it.The system uses large batches, inference-adapted kernels, and tensor parallelism when the model does not fit on one GPU.
- 5.3 Effective Throughput and Scalability Analysis: DeepSpeed-HE is most efficient for 6.7B–66B models; efficiency drops at 175B because limited memory restricts batch sizes, but remains 1.2x above the 1.3B model.The reported effective throughput includes separate generation and training phases.
- 5.3 Effective Throughput and Scalability Analysis: DeepSpeed-HE achieves 19x higher effective performance than existing systems, which operate below 5% of peak according to the reported comparison.The result highlights the difficulty of optimizing RLHF workloads and the system’s reported efficiency advantage.
- 5.3 Effective Throughput and Scalability Analysis: DeepSpeed-RLHF scales well to 64 GPUs, showing super-linear scaling at small scale and near-linear or sub-linear scaling at larger scales.ZeRO-based state partitioning increases available memory and batch size initially, while the maximum global batch size of 1024 eventually limits scaling.
6 Release: Try DeepSpeed Chat Today!
DeepSpeed-Chat is open-sourced and available to the AI community, with a GitHub landing page, roadmap, documentation, and contribution channels.
- DeepSpeed-Chat is open-sourced and available to the AI community.
- Users can get started through the DeepSpeed-Chat GitHub landing page and follow a roadmap of supported and planned features.
- The broader DeepSpeed ecosystem provides documentation, tutorials, news channels, and opportunities for contributions and collaboration.
Contributions
The listed contributors supported the hybrid engine, system and training pipeline, benchmarking, software, data, documentation, and project leadership.
- Contributors covered hybrid-engine development, ZeRO adaptation, training-pipeline and benchmarking support, software maintenance, documentation, and team leadership.