Implementation:Deepspeedai DeepSpeed HybridEngine Train Step
Overview
Concrete tool for switching the Hybrid Engine to training mode and executing PPO policy updates provided by the DeepSpeed library.
Description
DeepSpeedHybridEngine.train() restores original forward functions, unfuses LoRA, and calls transform_for_training() on inference containers. backward() is inherited from DeepSpeedEngine and handles ZeRO gradient communication. step() performs the optimizer update with LoRA-aware parameter handling and resets inference container parameters for the next generation cycle.
train() (L423-433)
The train() method switches the engine from inference mode back to training mode:
- Iterates over all inference containers and calls
transform_for_training()on each, resetting the kernel state for training compatibility. - Restores the original forward function (
orig_fwd) on each transformer layer module, replacing the inference container forward. - Restores original forward functions for non-transformer layers (embeddings, layer norms) as well.
- Calls the parent
DeepSpeedEngine.train(mode)to set PyTorch training mode flags. - Records the training start time for performance logging.
backward() (inherited, L2547-2720 in engine.py)
The backward() method is inherited from DeepSpeedEngine without modification. It:
- Scales the loss by gradient accumulation steps (if
scale_wrt_gas=True). - Handles ZeRO-specific loss scaling.
- Executes
loss.backward()with appropriate mixed-precision handling. - Manages gradient communication across data-parallel ranks.
step() (L435-445)
The step() method extends the parent DeepSpeedEngine.step() with inference container maintenance:
- Calls the parent
step()to perform the optimizer update, gradient clipping, and learning rate scheduling. - If inference containers exist and ZeRO-3 is not active, calls
reset_params()on each container to synchronize the inference container parameters with the updated training parameters. - Accumulates training latency for performance reporting.
Code Reference
| Property | Value |
|---|---|
| Repository | https://github.com/deepspeedai/DeepSpeed |
| File | deepspeed/runtime/hybrid_engine.py (L423-445), deepspeed/runtime/engine.py (L2547-2720, backward)
|
| train signature | def train(self, mode: bool = True) -> None
|
| step signature | def step(self, lr_kwargs=None) -> None
|
| backward signature | def backward(self, loss, retain_graph=False, scale_wrt_gas=True) -> None (inherited)
|
| Import | Accessed via engine returned by deepspeed.initialize()
|
I/O Contract
Inputs (train)
| Name | Type | Required | Description |
|---|---|---|---|
| mode | bool | No | True for training mode (default), False for eval mode |
Inputs (backward)
| Name | Type | Required | Description |
|---|---|---|---|
| loss | torch.Tensor | Yes | Scalar PPO loss tensor |
| retain_graph | bool | No | Whether to retain computation graph (default False) |
| scale_wrt_gas | bool | No | Scale gradients by gradient accumulation steps (default True) |
Inputs (step)
| Name | Type | Required | Description |
|---|---|---|---|
| lr_kwargs | dict | No | Extra keyword arguments for learning rate scheduler |
Outputs
| Name | Type | Description |
|---|---|---|
| (side effect) | — | Updated actor model parameters after PPO optimization step |
Usage Example
# Switch to training mode
engine.train()
# Compute PPO loss (simplified)
old_logprobs = compute_logprobs(engine, sequences, old_model=True)
new_logprobs = compute_logprobs(engine, sequences)
ratio = torch.exp(new_logprobs - old_logprobs)
ppo_loss = -torch.min(
ratio * advantages,
torch.clamp(ratio, 0.8, 1.2) * advantages
).mean()
# Update policy
engine.backward(ppo_loss)
engine.step()
Related Pages
Knowledge Sources
- https://github.com/deepspeedai/DeepSpeed
- https://arxiv.org/abs/1707.06347
- https://arxiv.org/abs/2308.01320
Last updated: 2026-02-09 00:00 GMT