Principle:Deepspeedai DeepSpeed Hybrid Engine Init
Overview
Initializing the DeepSpeed Hybrid Engine that enables seamless switching between inference mode (for experience generation) and training mode (for PPO updates) within the same engine.
Description
The Hybrid Engine is the core innovation of DeepSpeed-Chat for RLHF. It wraps the actor model with dual capabilities: (1) inference mode with optimized kernels, LoRA fusion, and optional tensor parallelism for fast text generation, and (2) training mode with ZeRO optimization for PPO policy updates. The engine automatically manages the transition between modes, including parameter gathering for ZeRO-3, LoRA fusion and unfusion, and inference container setup and teardown.
The DeepSpeedHybridEngine class inherits from DeepSpeedEngine and extends it with inference infrastructure. During initialization, it performs several key operations beyond the base engine setup:
- Inference container creation: The engine scans the model architecture and matches transformer layers against known replace policies (such as those for Llama, OPT, GPT-NeoX, and other supported architectures). For each matched layer, it creates an inference container that wraps the layer with optimized kernels for inference.
- LoRA parameter tracking: If the model uses LoRA adapters, the engine tracks all LoRA parameters separately so they can be fused into the base weights during inference and unfused during training.
- Tensor parallelism setup: If
inference_tp_size > 1, the engine creates communication groups and tensor slicing utilities for distributing inference computation across multiple GPUs. - Parameter pinning: When
pin_parameters=True(the default), ZeRO-3 parameters are gathered once before generation rather than being gathered layer-by-layer, reducing communication overhead.
The Hybrid Engine is created automatically by deepspeed.initialize() when the configuration contains "hybrid_engine": {"enabled": true}. The selection logic at line 201 of deepspeed/__init__.py checks config_class.hybrid_engine.enabled and instantiates DeepSpeedHybridEngine instead of the standard DeepSpeedEngine.
Theoretical Basis
RLHF requires alternating between generation (inference) and policy update (training). In the PPO algorithm applied to language models, each iteration consists of:
- Generation phase: The actor model generates text responses to prompts (inference).
- Scoring phase: The reward model and reference model evaluate the generated text (inference).
- Update phase: The actor model parameters are updated using the PPO objective (training).
Naive approaches to this alternation either require maintaining separate model copies for inference and training (doubling memory usage) or performing expensive model reloading between phases. The Hybrid Engine fuses both execution paths into a single model that dynamically switches between optimized inference containers and standard training forward functions, sharing the same underlying parameters.
The key insight is that inference containers (which fuse multiple operations like attention, projection, and normalization into single optimized kernels) and training forward functions (which maintain separate operation boundaries for correct gradient computation) can coexist on the same model parameters. The engine simply swaps which forward function each layer uses depending on the current mode.
References
- DeepSpeed-Chat: Easy, Fast and Affordable RLHF Training of ChatGPT-like Models at All Scales — https://arxiv.org/abs/2308.01320
Related Pages
Knowledge Sources
Last updated: 2026-02-09 00:00 GMT