[st-bang97]

Implements #8697.

This PR adds ReFlow, an asynchronous CPU-offloaded optimizer path for DeepSpeed ZeRO-3. The motivation and detailed design are described in #8697.

ReFlow reduces the cost of the CPU-offloaded optimizer path and overlaps the remaining host-side work with GPU backward computation.

The ReFlow design and evaluation have been accepted to ASPLOS 2027 as “ReFlow: Exposing Parallelism in CPU-Offloaded LLM Training via Register-Resident SIMD Fusion and Decoupled Update Scheduling.”

Once the paper is publicly available, I will update this PR and #8697 with the paper link and publish a technical blog covering ReFlow's design, usage, and evaluation.

What this PR adds

• ReFlow integration for ZeRO-3 CPU-offloaded training.
• ReFlow CPU Adam/AdamW and Lion optimizers.
• Fused AVX/SIMD optimizer kernels with BF16 gradient input and point-of-use FP32 promotion inside SIMD registers.
• CPU-resident gradient accumulation for intermediate micro-steps, avoiding repeated blocking H2D gradient retrievals.
• Final GPU-side accumulation and gradient-norm computation only at the terminal micro-step.
• Asynchronous bucket-wise CPU optimizer execution overlapping GPU backward.
• Ready/Persist scheduling that separates GPU-critical low-precision parameter generation from persistent FP32 optimizer-state updates.
• Gradient clipping and overflow handling without FP32 optimizer-state rollback.
• zero_optimization.reflow configuration.
• Unit/training tests, CI coverage, configuration documentation, and tutorial.

Validation

ReFlow has been tested on:

• NVIDIA B200
• NVIDIA A100
• NVIDIA RTX 5080
• NVIDIA RTX A5000
• NVIDIA RTX PRO 6000
• AMD EPYC CPUs
• Intel Xeon CPUs

Correctness testing covers Adam/AdamW and Lion, gradient accumulation, gradient clipping, overflow handling, checkpointing, learning-rate scheduling, high-precision optimizer-state access, and single-/multi-GPU execution.

Under deterministic GPU execution, ReFlow produces bit-identical training results to the baseline.

Representative performance results:

Workload	ZeRO-Infinity	ReFlow	Improvement
OPT-30B, 1× B200	166.41 TFLOPS/GPU	671.29 TFLOPS/GPU	4.0×
OPT-30B, 4× B200	366.40 TFLOPS/GPU	681.18 TFLOPS/GPU	1.86×

On a single A100 at 15B, ReFlow also achieves 1.7× higher throughput than ZeRO-Infinity.

In the standard AdamW evaluation, separate from the clipping-heavy experiment, ReFlow reduces the exposed optimizer update tail by more than 96.8% across the evaluated gradient-accumulation configurations and reduces total iteration time by up to 74.3%.

Under a separate clipping OPT-30B AdamW workload, ReFlow reduces the critical-path host update time from 2440 ms to 850 ms (65.2%), while requiring no FP32 optimizer-state rollback.

ReFlow is also optimizer-agnostic: with Lion, it achieves 1.6×–4.4× higher throughput than ZeRO-Infinity-Lion across 8 to 1 GPUs.

All modified files pass the DeepSpeed pre-commit checks.

[Copilot]

Copilot review overview

🟡 Changes recommended

Optimizer remapping changes supported client semantics, public optimizer steps can access unregistered native state, and failed asynchronous commits may still be checkpointed.

Review effort: Balanced
Findings: 4 · 2 · 3

Open (9)

• Adam remap drops maximize semantics · New
• ReflowCPULion inherits incompatible optimizer step · New
• ReflowCPULion inherits incompatible Lion step · New
• Checkpoint serialization ignores pending state update failures · New
• Lion remap drops fp32_optimizer_states · New
• NUMA validation runs on the wrong host · New
• Documentation disagrees with bucketwise cores default · New
• Reference documents incorrect worker core default · New
• Configuration reference omits public ReflowConfig fields · New

What changed in this PR

Adds Reflow, an asynchronous CPU-offloaded Adam/Lion path for ZeRO-3.

Changes:

• Adds bucket-wise asynchronous optimization, gradient accumulation, state persistence, and NUMA scheduling.
• Adds fused AVX kernels and accelerator builder integration.
• Adds configuration, documentation, and correctness/launcher tests.

File	Description
tests/​unit/​v1/​zero/​test_reflow.py	Tests Reflow integration and correctness.
tests/​unit/​v1/​ops/​test_reflow_cpu_optimizers.py	Tests native optimizer kernels.
tests/​unit/​launcher/​test_ds_arguments.py	Tests numactl fallback behavior.
op_builder/​xpu/​cpu_adam.py	Adds Reflow Adam sources.
op_builder/​supa/​cpu_lion.py	Adds Reflow Lion sources.
op_builder/​supa/​cpu_adam.py	Adds Reflow Adam sources.
op_builder/​sdaa/​cpu_adam.py	Adds Reflow Adam sources.
op_builder/​npu/​cpu_lion.py	Adds Reflow Lion sources.
op_builder/​npu/​cpu_adam.py	Adds Reflow Adam sources.
op_builder/​mps/​cpu_adam.py	Adds Reflow Adam sources.
op_builder/​mlu/​cpu_adam.py	Adds Reflow Adam sources.
op_builder/​hpu/​cpu_adam.py	Adds Reflow Adam sources.
op_builder/​cpu/​cpu_adam.py	Adds CPU Reflow sources.
op_builder/​cpu_lion.py	Adds default Lion sources.
op_builder/​cpu_adam.py	Adds default Adam sources.
docs/​code-docs/​source/​zero3.rst	Exposes Reflow configuration API.
docs/​code-docs/​source/​optimizers.rst	Publishes optimizer classes.
docs/​_tutorials/​reflow.md	Adds the Reflow tutorial.
docs/​_pages/​config-json.md	Documents Reflow configuration.
docs/​_data/​navigation.yml	Links the tutorial.
deepspeed/​utils/​timer.py	Adds current-stream synchronization.
deepspeed/​utils/​numa.py	Adds numactl permission fallback.
deepspeed/​runtime/​zero/​utils.py	Registers supported optimizers.
deepspeed/​runtime/​zero/​partitioned_param_coordinator.py	Orders all-gathers after Reflow writes.
deepspeed/​runtime/​zero/​config.py	Adds compatibility validation.
deepspeed/​runtime/​reflow/​reflow_utils.py	Implements affinity and range helpers.
deepspeed/​runtime/​reflow/​reflow_cpu_lion.py	Implements Reflow Lion wrapper.
deepspeed/​runtime/​reflow/​reflow_cpu_adam.py	Implements Reflow Adam wrapper.
deepspeed/​runtime/​reflow/​reflow_config.py	Defines Reflow settings.
deepspeed/​runtime/​reflow/​reflow_stage3.py	Implements asynchronous ZeRO-3 execution.
deepspeed/​runtime/​reflow/​README.md	Documents internal design.
deepspeed/​runtime/​reflow/​engine.py	Integrates and remaps optimizers.
deepspeed/​runtime/​reflow/​__init__.py	Defines the Reflow package.
deepspeed/​runtime/​engine.py	Selects Reflow engine components.
deepspeed/​runtime/​config.py	Rejects incompatible configurations.
csrc/​lion/​reflow_cpu_lion_impl.cpp	Implements fused Lion kernels.
csrc/​lion/​reflow_cpu_lion_bindings.cpp	Exposes Lion native APIs.
csrc/​lion/​cpu_lion.cpp	Registers Reflow Lion bindings.
csrc/​includes/​reflow_simd.h	Adds SIMD conversion/store helpers.
csrc/​includes/​reflow_cpu_lion.h	Defines Reflow Lion kernels.
csrc/​includes/​reflow_cpu_affinity.h	Applies worker affinity.
csrc/​includes/​reflow_cpu_adam.h	Defines Reflow Adam kernels.
csrc/​includes/​reflow_bindings.h	Declares native bindings.
csrc/​adam/​reflow_cpu_adam_impl.cpp	Implements fused Adam kernels.
csrc/​adam/​reflow_cpu_adam_bindings.cpp	Exposes Adam native APIs.
csrc/​adam/​cpu_adam.cpp	Registers Reflow Adam bindings.
ci/​tests_fetcher.py	Selects Reflow tests in CI.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[st-bang97]

I’ve addressed all nine review comments in separate signed-off commits, including maximize support, optimizer and precision guards, checkpoint error handling, target-host NUMA checks and fallback, and documentation fixes.
Please take another look when you have a chance.
Also, I added example file of Reflow with DeepSpeedExamples PR.

[st-bang97]

Hi @tjruwase, @tohtana, and @hwchen2017,

Just following up on this PR. All previous review comments have been addressed, and the implementation is ready for review.

Could you please take a look when you have time? I'd appreciate any feedback or guidance on the next steps.

Thank you!

[sfc-gh-truwase]

@st-bang97 thanks for so much for your PR and even more for your patience. Apologies for the delay. Taking a look now.

[pengdurice]

DeepSpeed will skip a step if the gradient from that step is invalid (e.g. nan or inf), is_invalid_grad_norm is used to check that. Does ReFlow call it and handle invalid clearly?

[st-bang97]

DeepSpeed will skip a step if the gradient from that step is invalid (e.g. nan or inf), is_invalid_grad_norm is used to check that. Does ReFlow call it and handle invalid clearly?

Thanks for pointing this out, @pengdurice!

I checked the current implementation and confirmed that ReFlow handles FP16 gradient overflow but does not use "is_invalid_grad_norm()". I'll update the validation to use "combine_grad_norm_groups()" and "is_invalid_grad_norm()", consistent with the latest ZeRO-3 path, and add regression tests for NaN/Inf handling.

I'm currently attending a conference, so I plan to push the changes next week.

Thanks again for the review!

[pengdurice]

DeepSpeed will skip a step if the gradient from that step is invalid (e.g. nan or inf), is_invalid_grad_norm is used to check that. Does ReFlow call it and handle invalid clearly?

Thanks for pointing this out, @pengdurice!

I checked the current implementation and confirmed that ReFlow handles FP16 gradient overflow but does not use "is_invalid_grad_norm()". I'll update the validation to use "combine_grad_norm_groups()" and "is_invalid_grad_norm()", consistent with the latest ZeRO-3 path, and add regression tests for NaN/Inf handling.

I'm currently attending a conference, so I plan to push the changes next week.

Thanks again for the review!

No worries, take your time and enjoy your conference. BTW, does the current implementation fuse the invalidity check into the SIMD kernel? If not, do you think it is a good idea to do so to avoid the cost of python's is_invalid_grad_norm? Is that an overkill?

[st-bang97]

DeepSpeed will skip a step if the gradient from that step is invalid (e.g. nan or inf), is_invalid_grad_norm is used to check that. Does ReFlow call it and handle invalid clearly?

Thanks for pointing this out, @pengdurice!
I checked the current implementation and confirmed that ReFlow handles FP16 gradient overflow but does not use "is_invalid_grad_norm()". I'll update the validation to use "combine_grad_norm_groups()" and "is_invalid_grad_norm()", consistent with the latest ZeRO-3 path, and add regression tests for NaN/Inf handling.
I'm currently attending a conference, so I plan to push the changes next week.
Thanks again for the review!

No worries, take your time and enjoy your conference. BTW, does the current implementation fuse the invalidity check into the SIMD kernel? If not, do you think it is a good idea to do so to avoid the cost of python's is_invalid_grad_norm? Is that an overkill?

Hi Peng, thanks for the detailed review and suggestions!

1. SIMD invalidity check

That's a good suggestion! The invalidity check isn't currently fused into the SIMD kernel, but it could help reduce Python overhead.

Since ReFlow computes gradient norms on the GPU and checks the resulting global norm separately, I'm not sure how much additional benefit SIMD fusion would provide. I think it's worth profiling and analyzing further before deciding whether to integrate it.

I'll first update the invalidity handling to match the latest ZeRO-3 implementation.

2. Potential data race across micro-steps

Good catch! The current implementation waits for the GPU copy to finish before CPU accumulation, but it doesn't explicitly guarantee that the accumulation completes before the next micro-step overwrites src_slice.

I haven't encountered this issue in my tests, but I agree that we should guarantee correctness regardless of execution timing. I'll add synchronization to protect buffer reuse while avoiding unnecessary blocking, along with a regression test.

3. AutoTP and AutoEP compatibility

ReFlow was originally designed for dense models, and we've validated single- and multi-GPU dense training, but not specifically AutoTP or AutoEP.

When we initially ported ReFlow to DeepSpeed ZeRO-3, we based the implementation on the existing ZeRO-3 code. Since #8670 was merged afterward, its changes, including _grad_norm_counts_every_rank(), haven't been integrated into ReFlow yet. I'll update the implementation to incorporate these changes.

One related optimization I'd like to mention is ReFlow's batched gradient-norm reduction. Instead of reducing each subgroup separately, it collects their squared norms into a GPU tensor and batches collective operations for subgroups sharing the same process group. This reduces the number of collective calls, particularly when there are many small subgroups, and could potentially benefit standard ZeRO-3 offload as well.

However, this optimization was originally developed without AutoTP or AutoEP. I'd like to preserve it while incorporating the recent changes, and I'd appreciate your thoughts on whether this batching approach could be extended to these parallel configurations while maintaining correctness.

I currently don't have access to a multi-GPU setup to validate these specific configurations, so I'd also appreciate some help with testing. I'll review the existing optimizer offload restrictions for AutoEP+AutoTP folding as well.

Thanks again for your feedback!

[pengdurice]

@st-bang97 , what's your thoughts on supporting Muon or other "second order" optimizers where there will be computes that are more than element wise. The challenges will be 1: gathering the matrix rather than sharding format; 2: slower matmal on CPU. Is that on your plan?

[st-bang97]

@st-bang97 , what's your thoughts on supporting Muon or other "second order" optimizers where there will be computes that are more than element wise. The challenges will be 1: gathering the matrix rather than sharding format; 2: slower matmal on CPU. Is that on your plan?

Thanks Peng, that's an interesting direction!

I think Muon and other matrix-based optimizers could fit well into ReFlow's Ready/Persist design. The main idea is to keep matrix-intensive computations on the GPU while overlapping CPU-side optimizer operations with GPU execution.

1. Reuse DeepSpeed's existing Muon implementation

For matrix gathering, we could reuse DeepSpeed's existing ZeRO Muon implementation.

Currently, Muon computes momentum and Newton–Schulz on the GPU after backward, then transfers both the Newton–Schulz result and updated momentum back to the CPU.

Instead, I'd like to integrate Muon into ReFlow's partition_grads() pipeline, allowing us to start processing each parameter once its fully accumulated gradients are available.

2. Extend ReFlow's Ready/Persist pipeline

The key idea is to compute momentum on the CPU while the GPU performs Newton–Schulz, avoiding the additional PCIe transfer of updated momentum.

The proposed execution flow would be:

• GPU: Compute the momentum required for Newton–Schulz and perform the matrix-intensive operations.
• CPU (in parallel): Independently update the CPU-resident momentum using the original gradients. Since this is an element-wise operation, I expect it could be faster than transferring the GPU-computed momentum back over PCIe.
• Ready (when the Newton–Schulz result arrives from the GPU): Reuse the existing gradient buffer to store the result and generate the low-precision weights.
• Persist: Asynchronously update the FP32 master weights using the same Newton–Schulz result, without repeating the matrix computation.

Unlike AdamW or Lion, Muon cannot generate low-precision weights until the Newton–Schulz result is available. However, CPU momentum updates do not depend on that result, so we can perform them concurrently with GPU Newton–Schulz.

This could spread CPU optimizer work and memory traffic over a wider execution window, potentially improving overlap and reducing memory bandwidth contention.

For correctness:

• Gradient clipping: We can reuse the precomputed momentum and Newton–Schulz result, since clipping is applied to the resulting update.
• Overflow: For speculative updates that must be discarded, we could investigate inverse-update rollback, similar to SuperOffload, while validating numerical stability.

3. Additional optimization opportunities

• Reduce PCIe traffic: Explore transferring the Newton–Schulz result in FP16 rather than FP32, leveraging ReFlow's low-precision buffers.
• CPU/GPU profiling: Prioritize GPU execution for matrix-heavy operations, but evaluate CPU accelerators such as Intel AMX when avoiding data transfers could be beneficial.
• Proactive data staging: Prefetch the required gradients and momentum to overlap data transfers with computation.

I'd be interested in exploring this as a follow-up project in 2027, starting with Muon support. What do you think about this approach? Do you think it would be a useful direction for DeepSpeed?

[xylian86]

Thanks @st-bang97 for the PR and the detailed write-up! I really like several of the design choices in ReFlow, and congratulations on the ASPLOS acceptance!

At a high level, ReFlow and SuperOffload address the same problem, building on the idea introduced by SuperOffload's STV idea: enable overlapping CPU optimizer computation with GPU work to hide optimizer overhead. SuperOffload was originally designed for superchips (e.g., GH200 and GB200), so we did not specifically optimize it for PCIe-connected systems.

We recently implemented a PCIe-optimized version of SuperOffload and benchmarked it against ReFlow on the same machine and workloads.

1. OPT-30B Fine-Tuning on Alpaca
We used the same experimental setup described in the ReFlow paper. All results are reported in seconds per training step (lower is better).

batch × seq	SuperOffload	ReFlow	SuperOffload speedup
8 × 2048	5.17	7.03	1.36×
8 × 4096	5.67	7.07	1.25×
8 × 8192	5.77	7.16	1.24×
16 × 2048	5.42	7.08	1.31×
16 × 4096	5.79	7.15	1.23×

SuperOffload consistently outperforms ReFlow, achieving 1.23–1.36× higher training throughput across all tested configurations.

2. End-to-End Qwen2.5-7B SFT on MetaMathQA
I'd also like to clarify one point from the ReFlow paper, which states that "SuperOffload drifts from the baseline, consistent with the numerical error introduced by rollback."

We believe this observation is caused by an implementation issue in an earlier version of SuperOffload, rather than an inherent numerical error introduced by rollback.

To investigate this, we conducted an end-to-end fine-tuning experiment comparing SuperOffload, ReFlow, and ZeRO-Infinity, focusing on both training throughput and convergence behavior.

Experimental setup:

• Dataset: All 395K MetaMathQA examples, one epoch (3,085 steps with batch size 128).
• Training: Learning rate 2e-5, cosine decay, 3% warmup, and gradient clipping at 1.0.
• Reproducibility: Identical data ordering across all runs.
• Evaluation: GSM8K test set (1,319 questions), zero-shot greedy decoding, evaluated every 1,000 steps and at the end of training.
• Baseline: The original Qwen2.5-7B model achieves 74.6% accuracy using the same evaluation prompt.

	SuperOffload	ReFlow	ZeRO-Infinity
mean step time	1.60 s	1.87 s	3.25 s
wall time	87 min	101 min	172 min
mean loss, last 100 steps	0.1238	0.1237	0.1238
GSM8K at step 1000 / 2000 / 3000 / 3085	79.9 / 81.5 / 82.4 / 82.6	79.0 / 81.3 / 83.0 / 82.6	79.4 / 82.3 / 82.2 / 82.1

Our results show closely matched convergence across all three implementations:

• The average per-step loss difference between any two runs is approximately 4e-4.
• None of the GSM8K accuracy differences are statistically significant (paired McNemar test, p ≥ 0.24 for every pair of runs at every evaluated checkpoint).
• All three implementations achieve nearly identical final training losses and GSM8K accuracy.
• Meanwhile, SuperOffload achieves 1.17× higher end-to-end training throughput than ReFlow, reducing total training time from 101 to 87 minutes, without any measurable degradation in convergence or final model accuracy.

We're also happy to rerun ReFlow with any settings you recommend (e.g., a smaller sub_group_size or a different CPU core allocation).

[st-bang97]

Thanks @st-bang97 for the PR and the detailed write-up! I really like several of the design choices in ReFlow, and congratulations on the ASPLOS acceptance!

At a high level, ReFlow and SuperOffload address the same problem, building on the idea introduced by SuperOffload's STV idea: enable overlapping CPU optimizer computation with GPU work to hide optimizer overhead. SuperOffload was originally designed for superchips (e.g., GH200 and GB200), so we did not specifically optimize it for PCIe-connected systems.

We recently implemented a PCIe-optimized version of SuperOffload and benchmarked it against ReFlow on the same machine and workloads.

1. OPT-30B Fine-Tuning on Alpaca
We used the same experimental setup described in the ReFlow paper. All results are reported in seconds per training step (lower is better).

batch × seq	SuperOffload	ReFlow	SuperOffload speedup
8 × 2048	5.17	7.03	1.36×
8 × 4096	5.67	7.07	1.25×
8 × 8192	5.77	7.16	1.24×
16 × 2048	5.42	7.08	1.31×
16 × 4096	5.79	7.15	1.23×

SuperOffload consistently outperforms ReFlow, achieving 1.23–1.36× higher training throughput across all tested configurations.

2. End-to-End Qwen2.5-7B SFT on MetaMathQA
I'd also like to clarify one point from the ReFlow paper, which states that "SuperOffload drifts from the baseline, consistent with the numerical error introduced by rollback."

We believe this observation is caused by an implementation issue in an earlier version of SuperOffload, rather than an inherent numerical error introduced by rollback.

To investigate this, we conducted an end-to-end fine-tuning experiment comparing SuperOffload, ReFlow, and ZeRO-Infinity, focusing on both training throughput and convergence behavior.

Experimental setup:

• Dataset: All 395K MetaMathQA examples, one epoch (3,085 steps with batch size 128).
• Training: Learning rate 2e-5, cosine decay, 3% warmup, and gradient clipping at 1.0.
• Reproducibility: Identical data ordering across all runs.
• Evaluation: GSM8K test set (1,319 questions), zero-shot greedy decoding, evaluated every 1,000 steps and at the end of training.
• Baseline: The original Qwen2.5-7B model achieves 74.6% accuracy using the same evaluation prompt.

	SuperOffload	ReFlow	ZeRO-Infinity
mean step time	1.60 s	1.87 s	3.25 s
wall time	87 min	101 min	172 min
mean loss, last 100 steps	0.1238	0.1237	0.1238
GSM8K at step 1000 / 2000 / 3000 / 3085	79.9 / 81.5 / 82.4 / 82.6	79.0 / 81.3 / 83.0 / 82.6	79.4 / 82.3 / 82.2 / 82.1

Our results show closely matched convergence across all three implementations:

• The average per-step loss difference between any two runs is approximately 4e-4.
• None of the GSM8K accuracy differences are statistically significant (paired McNemar test, p ≥ 0.24 for every pair of runs at every evaluated checkpoint).
• All three implementations achieve nearly identical final training losses and GSM8K accuracy.
• Meanwhile, SuperOffload achieves 1.17× higher end-to-end training throughput than ReFlow, reducing total training time from 101 to 87 minutes, without any measurable degradation in convergence or final model accuracy.

We're also happy to rerun ReFlow with any settings you recommend (e.g., a smaller sub_group_size or a different CPU core allocation).

Thanks for the detailed benchmarks, and especially for implementing the PCIe-optimized version of SuperOffload!

When I first read SuperOffload, I thought many of its techniques could also work well on PCIe-connected systems, apart from a few architecture-specific considerations. It's really exciting to see this implemented and evaluated!

1. Performance and execution environment

I'm really glad to see these performance improvements!

In our ReFlow evaluation, we tried to follow the original SuperOffload design as faithfully as possible when adapting it to our experimental environment. Looking at your new PCIe PR, I noticed several additional optimizations that were not part of the implementation we evaluated. I think these improvements may help explain the stronger performance you're seeing, which is great!

Regarding ReFlow, the forward pass provided sufficient time to overlap most of the deferred FP32 optimizer-state updates (Persist) in our experiments. We also used NUMA-aware CPU affinity, assigning each GPU rank a subset of CPU cores local to its NUMA node and pinning the optimizer workers accordingly.

I think differences in CPU scheduling or the available overlap window could also contribute to the performance gap. For example, if Persist takes longer than the forward overlap window, some of its work could become exposed on the critical path.

I'd be interested in comparing the execution timelines. Would you mind sharing the GPU count, CPU model/core allocation, and forward, backward, and "step()" times for both implementations?

2. Numerical differences and gradient clipping

Thanks also for clarifying the potential implementation issue in the earlier SuperOffload version!

For context, as in the original SuperOffload experiments, our earlier comparison used randomly initialized models trained from scratch, and gradient clipping occurred frequently in our runs.

Looking at the step-time plot you shared, the occasional spikes might correspond to clipping events, but there don't seem to be many around step 1,000. This makes me think that clipping may be relatively infrequent in your fine-tuning setup.

We also observed some variation depending on the dataset and CPU vendor in our earlier experiments. Since your experiments fine-tune a pretrained model, I think these differences in training conditions could partly explain why the numerical differences are so small, alongside the implementation improvements you mentioned.

It's encouraging to see such closely matched convergence!

3. Shared optimizations and future contributions

I also took a closer look at "the PCIe PR" (#9021), and some of the changes were particularly interesting!

The original SuperOffload paper favors GPU-side FP32 gradient promotion before transfer, taking advantage of the high bandwidth of NVLink-C2C.

In contrast, the PCIe version introduces BF16 gradient transfers, on-the-fly FP32 promotion within the CPU SIMD optimizer kernel, and fused gradient scaling. These closely align with the CPU-Centric Gradient and FlashOpt optimizations we identified and proposed in ReFlow.

I'm particularly happy to see that some of the optimizations we explored in ReFlow are also proving useful in the new SuperOffload implementation. It's exciting to see these ideas benefiting both approaches!

I think this also highlights the potential of ReFlow's broader design, particularly its combination of traffic-efficient CPU SIMD processing and Ready/Persist separation, which allows expensive FP32 optimizer-state persistence to be deferred beyond the critical path.

As mentioned in my "Q4 roadmap comment" (#8672 (comment)), I plan to contribute ReFlow's optimizations to the default DeepSpeed CPUAdam/ZeRO-3 path through three smaller PRs, while keeping ReFlow's asynchronous scheduling in its existing PR:

1. Fused CPUAdam kernel: Integrate gradient scaling/clipping and low-precision parameter generation directly into the optimizer kernel.
2. Low-precision gradient offloading: Transfer BF16/FP16 gradients and promote them within CPU SIMD registers, reducing PCIe traffic and CPU memory overhead.
3. CPU-side gradient accumulation: Avoid unnecessary CPU–GPU round trips during intermediate gradient accumulation steps.

Since the new PCIe implementation already includes techniques similar to parts of the first two, I think there's a good opportunity to coordinate these shared improvements.

I'm also curious which additional PCIe-specific optimizations contributed most to the performance improvements, such as pinned-memory handling and transfer scheduling. Perhaps we could bring some of these improvements into ReFlow as well, combining the strengths of both approaches.

I wish I still had access to the same multi-GPU hardware! These results are genuinely interesting, and I'd really enjoy reproducing the experiments, comparing execution profiles, and exploring further optimizations myself.

Thanks again for sharing these results and for your work on the PCIe implementation. I'd be happy to explore further improvements together!

[xylian86]

@st-bang97 Thanks for the detailed reply, and for reading through #9021!

Setup and per-phase times

• 4× B200, 1× Xeon Platinum 8559C (48 cores), 1 TB host memory; driver 580.159.03, CUDA 13.0, PyTorch 2.9.1.
• Each run uses 4 GPUs and the CPU socket they hang off. --bind_cores_to_rank gives each rank 12 cores of that socket: 2 for the main thread and 10 for the optimizer (ReFlow: main_thread_cores: 2, bucketwise_cores_per_worker: 5, enable_cpu_affinity: true; SuperOffload: cpuadam_cores_perc: 0.8).
• ZeRO-3, bf16, optimizer state fully on the CPU, activation checkpointing, FlashAttention-2, overlap_comm: true.

Per-phase times on rank 0 from DeepSpeed’s wall-clock breakdown, in seconds, median of steps 3–5.

batch × seq	SuperOffload fwd / bwd / step	total	ReFlow fwd / bwd / step	total
8 × 2048	0.27 / 1.87 / 3.03	5.17	0.27 / 1.98 / 4.78	7.03
8 × 4096	0.45 / 2.42 / 2.76	5.67	0.45 / 2.28 / 4.35	7.07
8 × 8192	0.80 / 2.74 / 2.23	5.77	0.81 / 2.61 / 3.75	7.16
16 × 2048	0.45 / 2.17 / 2.80	5.42	0.45 / 2.23 / 4.42	7.08
16 × 4096	0.81 / 2.98 / 2.00	5.79	0.81 / 2.60 / 3.76	7.15

Gradient clipping

In pre-training, gradient clipping typically occurs during warmup or in response to occasional spikes. If clipping continues to occur on nearly every step long after warmup, it indicates an underlying training issue. In such cases, pre-training should be stopped rather than allowed to continue.

Convergence

Looking at the step-time plot you shared, the occasional spikes might correspond to clipping events, but there don't seem to be many around step 1,000. This makes me think that clipping may be relatively infrequent in your fine-tuning setup.

We also observed some variation depending on the dataset and CPU vendor in our earlier experiments. Since your experiments fine-tune a pretrained model, I think these differences in training conditions could partly explain why the numerical differences are so small, alongside the implementation improvements you mentioned.

I don't think this is related to the CPU vendor or the dataset, but if you could run #9021 and demonstrate the divergence, I'd be happy to be convinced otherwise.

Shared optimizations

You’re right that the PCIe version uses BF16 gradient transfers, fused CPU kernel, and we appreciate ReFlow making the case for them. The original SuperOffload did not use them because it targets superchips: on GH200 and GB200, NVLink-C2C carries 450 GB/s in each direction, so sending gradients costs little and promoting them on the GPU was the better trade there, as you noted. Over PCIe the gradient traffic is on the critical path, and these optimizations pay off.

Two step CPU-Adam

its combination of traffic-efficient CPU SIMD processing and Ready/Persist separation, which allows expensive FP32 optimizer-state persistence to be deferred beyond the critical path.

@st-bang97 If I understand correctly, the expensive FP32 optimizer-state updates must either be hidden behind the forward pass of the next iteration (which is uncommon in modest settings) or contend with the BF16 CPU Adam of the next step. Happy to be corrected if I'm missing something.

[st-bang97]

@st-bang97 Thanks for the detailed reply, and for reading through #9021!

Setup and per-phase times

• 4× B200, 1× Xeon Platinum 8559C (48 cores), 1 TB host memory; driver 580.159.03, CUDA 13.0, PyTorch 2.9.1.
• Each run uses 4 GPUs and the CPU socket they hang off. --bind_cores_to_rank gives each rank 12 cores of that socket: 2 for the main thread and 10 for the optimizer (ReFlow: main_thread_cores: 2, bucketwise_cores_per_worker: 5, enable_cpu_affinity: true; SuperOffload: cpuadam_cores_perc: 0.8).
• ZeRO-3, bf16, optimizer state fully on the CPU, activation checkpointing, FlashAttention-2, overlap_comm: true.

Per-phase times on rank 0 from DeepSpeed’s wall-clock breakdown, in seconds, median of steps 3–5.

batch × seq SuperOffload fwd / bwd / step total ReFlow fwd / bwd / step total
8 × 2048 0.27 / 1.87 / 3.03 5.17 0.27 / 1.98 / 4.78 7.03
8 × 4096 0.45 / 2.42 / 2.76 5.67 0.45 / 2.28 / 4.35 7.07
8 × 8192 0.80 / 2.74 / 2.23 5.77 0.81 / 2.61 / 3.75 7.16
16 × 2048 0.45 / 2.17 / 2.80 5.42 0.45 / 2.23 / 4.42 7.08
16 × 4096 0.81 / 2.98 / 2.00 5.79 0.81 / 2.60 / 3.76 7.15

Gradient clipping

In pre-training, gradient clipping typically occurs during warmup or in response to occasional spikes. If clipping continues to occur on nearly every step long after warmup, it indicates an underlying training issue. In such cases, pre-training should be stopped rather than allowed to continue.

Convergence

Looking at the step-time plot you shared, the occasional spikes might correspond to clipping events, but there don't seem to be many around step 1,000. This makes me think that clipping may be relatively infrequent in your fine-tuning setup.

We also observed some variation depending on the dataset and CPU vendor in our earlier experiments. Since your experiments fine-tune a pretrained model, I think these differences in training conditions could partly explain why the numerical differences are so small, alongside the implementation improvements you mentioned.

I don't think this is related to the CPU vendor or the dataset, but if you could run #9021 and demonstrate the divergence, I'd be happy to be convinced otherwise.

Shared optimizations

You’re right that the PCIe version uses BF16 gradient transfers, fused CPU kernel, and we appreciate ReFlow making the case for them. The original SuperOffload did not use them because it targets superchips: on GH200 and GB200, NVLink-C2C carries 450 GB/s in each direction, so sending gradients costs little and promoting them on the GPU was the better trade there, as you noted. Over PCIe the gradient traffic is on the critical path, and these optimizations pay off.

Two step CPU-Adam

its combination of traffic-efficient CPU SIMD processing and Ready/Persist separation, which allows expensive FP32 optimizer-state persistence to be deferred beyond the critical path.

@st-bang97 If I understand correctly, the expensive FP32 optimizer-state updates must either be hidden behind the forward pass of the next iteration (which is uncommon in modest settings) or contend with the BF16 CPU Adam of the next step. Happy to be corrected if I'm missing something.

Thanks for sharing the detailed configuration and timing breakdown! This is really helpful.

1. Performance and experimental configuration

Looking at the results, I wonder if the relatively short GPU execution time, particularly for 8 × 2048 and 16 × 2048, could be contributing to the performance difference.

As described in our paper, we used a per-GPU batch size of 8, corresponding to a global batch size of 32 on 4 GPUs.

Just to clarify, are the batch sizes of 8 and 16 in your table global or per-GPU batch sizes?

If they're global batch sizes, that would mean only 2 or 4 samples per GPU. With smaller batches and shorter sequences, there may simply be less time to overlap CPU optimizer work with GPU computation, potentially exposing more overhead on the critical path.

Interestingly, ReFlow's reported step() time decreases as sequence length increases, while the total training step time remains relatively constant. This seems consistent with limited CPU–GPU overlap.

I also checked our original experimental configuration. We used reduce_bucket_size=4e8 and sub_group_size=4e8.

Would you mind trying these settings as well? They can affect gradient processing and when CPU optimizer work is scheduled, so I'd be interested in seeing how they affect the results.

I'll double-check our remaining experimental settings and follow up if I find any other relevant differences.

2. Clarification on gradient clipping

I agree that frequent gradient clipping long after warmup would be unusual and could indicate a training issue.

For context, Figure 14 of the original SuperOffload paper shows pre-training starting with a loss above 8 and frequent rollback events during the first 1,000 steps, before training stabilizes. The red markers indicate rollbacks triggered by gradient clipping or NaN/Inf values.

We wanted to investigate numerical behavior under similarly challenging early-training conditions. That's why we trained randomly initialized models from scratch and evaluated a scenario where gradient clipping occurred frequently.

Our goal was to demonstrate the numerical behavior we observed with the SuperOffload implementation used in our evaluation under these conditions, not to suggest that frequent clipping should persist throughout normal training.

I appreciate the additional convergence experiments you've shared. It's encouraging to see the updated implementation achieving closely matched results!

3. Clarification on Two-step CPU Adam

I also wanted to clarify one point about ReFlow's Ready/Persist design, particularly the assumption that FP32 updates must either be fully hidden behind the next forward pass or contend with the BF16 CPU Adam of the following step.

That's not quite how ReFlow schedules these two phases.

ReFlow always completes the Ready phase first, generating the BF16 parameters required for the next forward pass without modifying the persistent FP32 optimizer states.

Once the BF16 parameters are ready and the step has been validated, Persist begins asynchronously and overlaps with the next forward pass. If necessary, it can continue into the subsequent backward pass.

Importantly, the next Ready update for a given subgroup waits for its previous Persist update to complete. Therefore, the two phases are not designed to perform competing updates on the same optimizer state simultaneously.

The distinction is that Persist is not restricted to being hidden entirely within the next forward pass, nor does unfinished Persist necessarily imply contention with the next BF16 update.

If the available overlap time is insufficient, the remaining work can introduce synchronization delays or CPU resource contention, potentially affecting forward or backward execution. However, that is a consequence of insufficient overlap, rather than an inherent requirement of the two-step design.

I think this distinction is important when interpreting the timing results, and I'd be interested in investigating how much of the reported step time comes from these dependencies.

Thanks again for the detailed investigation and discussion!

[st-bang97]

One additional update: While reviewing the recent CPU affinity changes, I found that the CPU cores were incorrectly allocated between the main thread and optimizer workers. This may have affected ReFlow's performance in your benchmarks.

I'll correct the core allocation and push the fix together with the changes addressing Peng's review comments.

Thanks again for the detailed feedback!

[st-bang97]

Hi @pengdurice,

I've pushed a new commit addressing your review comments: [c4b6660](c4b6660).

The main changes are:

1. NaN/Inf handling: Updated gradient norm validation to use combine_grad_norm_groups() and is_invalid_grad_norm(), consistent with the latest ZeRO-3 implementation. Invalid gradients now trigger step skipping before the optimizer state is committed.

2. Gradient accumulation race: Added per-parameter synchronization to ensure that asynchronous CPU accumulation completes before the corresponding gradient buffers are reused. This avoids unnecessary synchronization across unrelated parameters.

3. AutoTP/AutoEP compatibility: Integrated _grad_norm_counts_every_rank() and updated the group-aware gradient norm reductions to account for different parameter partition groups.

I've also added regression tests covering NaN/Inf handling, gradient accumulation across multiple micro-steps, and CPU core allocation.

The full range of AutoTP/AutoEP configurations still needs further validation, but the gradient norm handling is now aligned with the latest ZeRO-3 changes.

[st-bang97]

Hi @xylian86,

I've pushed the CPU affinity fixes in [c4b6660](c4b6660).

The previous implementation did not correctly separate the CPU cores reserved for the main thread from those used by optimizer workers. I've corrected the core allocation and enabled main-thread pinning by default.

I also found that task-based affinity introduced significant overhead from repeatedly applying CPU affinity for each submitted task. To reduce this overhead, I've changed the default bucketwise_worker_affinity from task to thread, so each worker thread retains its assigned CPU affinity instead of reapplying it for every task.

The commit also updates the state-update worker's CPU allocation and adds regression tests for multi-rank CPU affinity.

These changes may affect the performance results you reported, although their impact on the B200 configuration still needs to be measured.

It would be useful to rerun the benchmarks with the updated commit, using reduce_bucket_size=4e8 and sub_group_size=4e8 as discussed earlier.

[xylian86]

@st-bang97 Thanks for fixing the implementation issues.

1. Rerun with c4b6660 and your bucket sizes

We reran ReFlow at c4b6660 with reduce_bucket_size 4e8 and sub_group_size 4e8, keeping the previous setup unchanged (4× B200, one Xeon 8559C socket, 12 cores/rank, ZeRO-3, BF16, CPU optimizer states).

For clarity, all batch sizes are global. Your per-GPU batch of 8 corresponds to our 32 × 2048 row. Results are median seconds/step over steps 3–5.

global batch × seq (per GPU)	SuperOffload	ReFlow 3035444	ReFlow c4b6660	SuperOffload speedup
16 × 2048 (4)	5.42	7.08	6.96	1.28×
16 × 4096 (4)	5.79	7.15	7.00	1.21×
32 × 2048 (8)	5.80	7.19	6.98	1.20×

2. Ready/Persist

Agreed on the execution order. My concern is CPU resource contention, not state dependencies: when Persist overlaps the next backward, it competes with Ready for CPU cores and DRAM bandwidth.

[st-bang97]

Hi @xylian86,

Thanks for continuing to run these experiments!

1. Ready/Persist contention

I agree that Ready and Persist can compete for CPU cores and memory bandwidth when their execution overlaps.

In our original experiments, nearly all of the Ready-phase CPU optimizer computation was overlapped with GPU backward computation, and Persist generally completed before the next Ready phase began. As a result, we didn't encounter significant contention between the two.

I'm curious whether you observed this contention in your experiments, or whether it's a potential concern based on ReFlow's execution schedule.

I'm particularly interested in understanding why Ready appears to be less effectively overlapped in your experiments. The amount of optimizer work that appears to remain on the critical path is somewhat surprising compared with what we observed in our original experiments.

Have you had a chance to profile Ready or identify any particular bottlenecks? I'd be interested in any insights you have, whether the issue is CPU kernel throughput, gradient availability, synchronization with Persist, or something else. I'd like to investigate this further and see where we can optimize ReFlow.

2. Smaller reduction buckets and subgroups

One important difference is that our experimental system had roughly twice as many CPU cores available per GPU as your current setup, which likely contributed to the better overlap we observed.

Given this difference, I wonder if smaller reduction buckets and subgroups would work better on your system. In our experiments, reducing the subgroup size caused little to no performance degradation and sometimes even improved performance, as long as the subgroups remained sufficiently large.

Would you mind trying the following configurations?

"reduce_bucket_size"| "sub_group_size"
"2e8"| "2e8"
"1e8"| "1e8"

Smaller reduction buckets could make gradients available to the CPU earlier, while smaller subgroups could provide finer-grained scheduling and reduce the time Ready waits for the corresponding Persist updates.

3. Ready worker configuration

Another parameter worth exploring is "bucketwise_cores_per_worker".

With your current configuration of 10 optimizer cores per rank and "bucketwise_cores_per_worker=5", ReFlow uses two Ready workers with five cores each. It might be interesting to compare this with "bucketwise_cores_per_worker=10", using a single Ready worker with all 10 cores.

4. Further optimization opportunities

These results also make me think that we should consider performing a portion of the optimizer updates on the GPU in ReFlow, similar to SuperOffload's heterogeneous optimizer computation.

SuperOffload's approach of distributing optimizer computation across the CPU and GPU makes a lot of sense, especially on systems with powerful GPUs and relatively limited CPU resources.

Since CPU optimizer computation was largely hidden in our original experiments, we focused on optimizing the fully CPU-offloaded execution path. However, your results highlight the potential benefits of a more flexible distribution of optimizer work.

I think combining this approach with ReFlow's Ready/Persist design would be an interesting direction to explore in the future.

[xylian86]

@st-bang97

1. Bucket sizes and worker configuration

We reran ReFlow at c4b6660 with the same setup as before (4× B200, one Xeon 8559C socket, 12 cores/rank).

reduce_bucket_size / sub_group_size, workers	32 × 2048 (8 per GPU)	16 × 4096 (4)	16 × 2048 (4)
4e8 / 4e8, 2 × 5 cores	6.98	7.00	6.96
2e8 / 2e8, 2 × 5 cores	6.97	7.00	6.96
1e8 / 1e8, 2 × 5 cores	6.99	7.03	7.00
4e8 / 4e8, 1 × 10 cores	7.00	7.03	7.00
SuperOffload	5.80	5.79	5.42

2. Hardware difference

One important difference is that our experimental system had roughly twice as many CPU cores available per GPU as your current setup, which likely contributed to the better overlap we observed.

I'm a bit confused about how you arrived at twice as many CPU cores per GPU. When running your experiments on 4× B200 GPUs, did you use two CPU sockets?

In a standard DGX B200 system, 8 GPUs share 2 CPUs, so for a fair comparison, a 4-GPU run should use only one CPU socket, right?

Put another way, when renting 4× B200 GPUs from a cloud provider, you would typically get half a node's CPU resources (one socket, 48 cores), rather than both sockets.

I'm happy to be corrected if I'm missing something.

3. CPU-only optimizer

To clarify, SuperOffload results presented here uses offload_optimizer.ratio=1.0 (heterogeneous optimizer computation disabled): all optimizer states and updates remain on the CPU, with no GPU optimizer computation.

[st-bang97]

Hi @xylian86,

Thanks for running these additional experiments! It's interesting that the performance remains nearly unchanged across the different bucket sizes and worker configurations.

1. Hardware configuration

Let me clarify our hardware setup. The GPU resources we received through a government-funded computing program were provided as full 8-GPU HGX B200 servers, rather than DGX systems.

For the 4-GPU experiments, we used two GPUs from each NUMA node, utilizing the CPU resources across both sockets. Our system had two AMD EPYC 9365 CPUs, providing 72 physical cores in total, compared with 48 physical cores in your single-socket setup.

I agree that your single-socket configuration is more representative of a typical 4-GPU deployment, particularly in cloud environments.

I also looked more closely at the memory specifications and noticed another important difference between the two platforms:

Specification	AMD EPYC 9365	Intel Xeon Platinum 8559C
Memory channels per socket	12	8
Maximum memory speed	DDR5-6400	DDR5-5600
Theoretical bandwidth per socket	614.4 GB/s	358.4 GB/s

Even at the single-socket level, the EPYC platform has approximately 1.71x the theoretical memory bandwidth of the Xeon platform. Given the memory-intensive nature of CPU optimizer computation, this difference is worth investigating.

2. Further investigation

Thanks also for clarifying that your SuperOffload results use offload_optimizer.ratio=1.0, with all optimizer updates performed on the CPU.

Since tuning the bucket sizes and worker configuration had little impact, I'd like to take a closer look at the implementation itself.

I'll compare our research implementation with the current DeepSpeed integration to check whether any performance optimizations were missed or changed during the porting process. In particular, I'd like to investigate the Ready-phase bottleneck and identify further optimization opportunities for ReFlow.

I think the hardware differences and the implementation changes are both worth examining to better understand the performance gap.