Skip to content

Add ReFlow ZeRO-3 CPU-offload fine-tuning examples - #1017

Open
st-bang97 wants to merge 2 commits into
deepspeedai:masterfrom
st-bang97:PR_branch
Open

st-bang97 wants to merge 2 commits into
deepspeedai:masterfrom
st-bang97:PR_branch

Conversation

@st-bang97

Copy link
Copy Markdown

Companion examples for deepspeedai/DeepSpeed#8719, which implements
deepspeedai/DeepSpeed#8697.

Summary

This PR adds training/DeepSpeed-Reflow, a set of end-to-end fine-tuning
examples for ReFlow, the asynchronous CPU-offloaded optimizer path
for DeepSpeed ZeRO-3.

The examples cover:

  • OPT and Llama fine-tuning on single- and multi-GPU configurations.
  • Adam/AdamW and Lion optimizer paths.
  • ReFlow and the corresponding plain ZeRO-3 CPU optimizer-offload
    baseline using otherwise matching configurations.
  • ReFlow configuration through the zero_optimization.reflow block.
  • Throughput and memory comparison between ReFlow and the baseline.
  • Deterministic loss-equivalence checking.
  • Launcher, configuration, failure-propagation, and preprocessing tests.

The runtime implementation itself is provided by
deepspeedai/DeepSpeed#8719. This PR contains only the example,
benchmark, validation, and usage code.

Examples

The added launchers include:

  • OPT-350M on 1 GPU
  • OPT-30B on 1, 4, and 8 GPUs
  • Llama-7B, Llama-8B, and Llama-13B on 1 GPU
  • Llama-70B on 8 GPUs
  • Adam/AdamW and Lion variants for the larger multi-GPU examples

Each launcher accepts either reflow or zerooffload mode, allowing
the same model and training configuration to be run with and without
ReFlow.

For example:

cd training/DeepSpeed-Reflow

bash finetune_opt-30b_8gpu.sh reflow
bash finetune_opt-30b_8gpu.sh zerooffload

Signed-off-by: st-bang97 <st.bang@dgist.ac.kr>
Signed-off-by: st-bang97 <st.bang@dgist.ac.kr>
@st-bang97
st-bang97 requested a review from tjruwase as a code owner October 2, 2026 07:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant