Implementation:OpenRLHF OpenRLHF Ckpt DS Zero to Universal
| Knowledge Sources | |
|---|---|
| Domains | Checkpointing, DeepSpeed, Utilities |
| Last Updated | 2026-02-07 10:40 GMT |
Overview
Concrete tool for converting DeepSpeed ZeRO checkpoints to the universal checkpoint format.
Description
This shell script converts the latest DeepSpeed ZeRO checkpoint in a given directory to a universal checkpoint format using deepspeed.checkpoint.ds_to_universal. It handles PPO-specific directory structures by automatically detecting _actor and _critic subdirectories and converting each independently. The script reads the latest tag file, creates a new universal tag (suffixed with _uni), and writes a latest_universal pointer.
Usage
Use this script after training with DeepSpeed ZeRO (stages 1-3) when you need to convert checkpoints to a format portable across different ZeRO configurations or for model extraction. This is particularly useful for PPO training where actor and critic checkpoints are stored separately.
Code Reference
Source Location
- Repository: OpenRLHF
- File: examples/scripts/ckpt_ds_zero_to_universal.sh
- Lines: 1-59
Signature
# Usage: ckpt_ds_zero_to_universal.sh <CKPT_PATH> [additional arguments]
#
# Arguments:
# CKPT_PATH - Path to the DeepSpeed checkpoint directory
# additional arguments - Extra args forwarded to deepspeed.checkpoint.ds_to_universal
#
# Internal function:
# process_dir(path) - Converts a single checkpoint directory
Import
bash examples/scripts/ckpt_ds_zero_to_universal.sh /path/to/checkpoint
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| CKPT_PATH | string | Yes | Path to DeepSpeed ZeRO checkpoint directory |
| additional arguments | string | No | Extra flags for ds_to_universal (e.g., --num_extract_workers) |
Outputs
| Name | Type | Description |
|---|---|---|
| Universal checkpoint | directory | Saved as {CKPT_PATH}/{latest_tag}_uni |
| latest_universal | file | Tag file pointing to the latest universal checkpoint |
Usage Examples
Convert a Standard Checkpoint
# Convert a single checkpoint directory
bash examples/scripts/ckpt_ds_zero_to_universal.sh ./ckpt/checkpoints_sft
# The script reads ./ckpt/checkpoints_sft/latest to find the tag,
# then creates ./ckpt/checkpoints_sft/{tag}_uni/
Convert PPO Actor/Critic Checkpoints
# For PPO checkpoints with _actor and _critic subdirectories
bash examples/scripts/ckpt_ds_zero_to_universal.sh ./ckpt/checkpoints_ppo
# Automatically detects and converts:
# ./ckpt/checkpoints_ppo/_actor/{tag} -> _actor/{tag}_uni
# ./ckpt/checkpoints_ppo/_critic/{tag} -> _critic/{tag}_uni