Principle:Deepspeedai DeepSpeed PPO Policy Update

Overview

Updating the actor policy using Proximal Policy Optimization (PPO) by switching the Hybrid Engine back to training mode after experience generation.

Description

After generating experiences, the Hybrid Engine switches back to training mode via engine.train(), which performs the reverse of the inference mode transition:

1. Restores original forward functions: Each transformer layer's forward method is restored to the original PyTorch implementation that supports autograd, removing the optimized inference containers from the execution path.
2. Unfuses LoRA adapters: If LoRA weights were fused into the base weights during inference, they are separated back into their original low-rank form so that only the adapter parameters receive gradient updates.
3. Calls transform_for_training(): Each inference container resets its internal state to be compatible with training mode.
4. Re-partitions parameters for ZeRO-3 training: If ZeRO Stage 3 is active, parameters are returned to their partitioned state across data-parallel ranks for memory-efficient training.

Once in training mode, the PPO loss is computed using the generated sequences, their log-probabilities under the current policy, the reward model scores, and the reference model's log-probabilities (for the KL penalty). The standard engine.backward() and engine.step() methods perform the policy update, with the engine handling all gradient communication and optimizer state management through its ZeRO optimization.

The transition between inference and training modes is the defining operation of the Hybrid Engine. By keeping both capabilities in a single engine with shared parameters, the RLHF training loop avoids the memory overhead of maintaining separate model copies and the latency of model reloading.

Theoretical Basis

The PPO objective for RLHF maximizes the clipped surrogate advantage:

L_PPO = E[min(r_t * A_t, clip(r_t, 1-eps, 1+eps) * A_t)]

where:

• r_t = pi_new(a|s) / pi_old(a|s) is the probability ratio between the updated and old policy
• A_t is the advantage estimate (reward minus value baseline)
• eps is the clipping parameter (typically 0.2)

The clipping prevents the policy from changing too drastically in a single update, which stabilizes training. Without clipping, large policy updates can degrade performance catastrophically.

In the RLHF setting, the advantage is typically computed as:

A_t = R(x, y) - V(x) - beta * KL(pi_new || pi_ref)

where R(x, y) is the reward model score, V(x) is the value function baseline, and beta * KL(pi_new || pi_ref) is the KL penalty that prevents the policy from diverging too far from the reference (SFT) model. The KL coefficient beta controls the strength of this constraint and may be adapted during training.

References

• Proximal Policy Optimization Algorithms — https://arxiv.org/abs/1707.06347
• InstructGPT: Training language models to follow instructions with human feedback — https://arxiv.org/abs/2203.02155

Related Pages

• Implementation:Deepspeedai_DeepSpeed_HybridEngine_Train_Step

Knowledge Sources

• https://github.com/deepspeedai/DeepSpeed
• https://arxiv.org/abs/1707.06347
• https://arxiv.org/abs/2203.02155
• https://arxiv.org/abs/2308.01320

Last updated: 2026-02-09 00:00 GMT