From 5b985ec7d7583db6751cb4780138beedc9893771 Mon Sep 17 00:00:00 2001
From: st-bang97 <st.bang@dgist.ac.kr>
Date: Fri, 2 Oct 2026 06:54:30 +0000
Subject: [PATCH 1/2] Add Reflow ZeRO-3 fine-tuning examples

Signed-off-by: st-bang97 <st.bang@dgist.ac.kr>
---
 training/DeepSpeed-Reflow/.gitignore          |  20 +
 training/DeepSpeed-Reflow/README.md           | 215 +++++++
 training/DeepSpeed-Reflow/check_bitexact.sh   | 117 ++++
 .../finetune_llama-13b_1gpu.sh                | 153 +++++
 .../finetune_llama-70b_8gpu.sh                | 157 +++++
 .../finetune_llama-70b_8gpu_lion.sh           | 158 +++++
 .../finetune_llama-7b_1gpu.sh                 | 152 +++++
 .../finetune_llama-8b_1gpu.sh                 | 152 +++++
 .../DeepSpeed-Reflow/finetune_opt-30b_1gpu.sh | 155 +++++
 .../DeepSpeed-Reflow/finetune_opt-30b_4gpu.sh | 155 +++++
 .../finetune_opt-30b_4gpu_lion.sh             | 134 +++++
 .../DeepSpeed-Reflow/finetune_opt-30b_8gpu.sh | 155 +++++
 .../finetune_opt-30b_8gpu_lion.sh             | 134 +++++
 .../finetune_opt-350m-1gpu.sh                 | 153 +++++
 training/DeepSpeed-Reflow/finetune_zero3.py   | 538 ++++++++++++++++++
 training/DeepSpeed-Reflow/requirements.txt    |   8 +
 training/DeepSpeed-Reflow/run_compare.sh      | 232 ++++++++
 .../DeepSpeed-Reflow/tests/test_examples.py   | 229 ++++++++
 18 files changed, 3017 insertions(+)
 create mode 100644 training/DeepSpeed-Reflow/.gitignore
 create mode 100644 training/DeepSpeed-Reflow/README.md
 create mode 100755 training/DeepSpeed-Reflow/check_bitexact.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_llama-13b_1gpu.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_llama-70b_8gpu.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_llama-70b_8gpu_lion.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_llama-7b_1gpu.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_llama-8b_1gpu.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_opt-30b_1gpu.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_opt-30b_4gpu.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_opt-30b_4gpu_lion.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_opt-30b_8gpu.sh
 create mode 100644 training/DeepSpeed-Reflow/finetune_opt-30b_8gpu_lion.sh
 create mode 100755 training/DeepSpeed-Reflow/finetune_opt-350m-1gpu.sh
 create mode 100644 training/DeepSpeed-Reflow/finetune_zero3.py
 create mode 100644 training/DeepSpeed-Reflow/requirements.txt
 create mode 100644 training/DeepSpeed-Reflow/run_compare.sh
 create mode 100644 training/DeepSpeed-Reflow/tests/test_examples.py

diff --git a/training/DeepSpeed-Reflow/.gitignore b/training/DeepSpeed-Reflow/.gitignore
new file mode 100644
index 000000000..c198614a2
--- /dev/null
+++ b/training/DeepSpeed-Reflow/.gitignore
@@ -0,0 +1,20 @@
+# Conda environment setup — local to our environment, not part of the public example.
+# Public users install via requirements.txt; this captures our exact dev environment.
+environment.yml
+conda_env*.yml
+*.conda.yml
+
+# Generated at runtime by the launcher scripts (DeepSpeed config + training outputs/logs).
+*_config.json
+*_output/
+__pycache__/
+wandb/
+
+# JIT-built CPU-Adam/Lion op extensions (torch cpp_extension build dirs).
+cpu_adam/
+cpu_lion/
+fused_adam/
+
+# Comparison / benchmark output collected by run_compare.sh.
+compare_out/
+bitexact_out/
diff --git a/training/DeepSpeed-Reflow/README.md b/training/DeepSpeed-Reflow/README.md
new file mode 100644
index 000000000..e6ef92642
--- /dev/null
+++ b/training/DeepSpeed-Reflow/README.md
@@ -0,0 +1,215 @@
+# Reflow Fine-Tuning Examples
+
+Fine-tune large language models with [DeepSpeed](https://www.deepspeed.ai/) ZeRO Stage 3 + **Reflow**, an asynchronous CPU-offload optimizer for mixed-precision BF16 training. Reflow keeps the optimizer state and FP32 master weights on the CPU like ZeRO-Offload, but overlaps the CPU optimizer work with the backward pass instead of running it serially afterward. It uses the **same GPU memory** as ZeRO-Offload and **~12% less host RAM** (OPT-30B: ~549 vs ~625 GB) — gradients are offloaded in half precision (BF16) and promoted to FP32 inside the CPU kernel, so there is no CPU-side FP32 gradient buffer (see [Memory](#memory) for measured numbers).
+
+Reflow supports **Adam/AdamW and Lion** with BF16 model parameters and gradients, FP32 master weights, and FP32 optimizer states. The same scripts run the Reflow path or the plain ZeRO-Offload baseline. `check_bitexact.sh` compares deterministic per-step losses for a selected model and configuration.
+
+## Quick Start
+
+### 1. Install dependencies
+
+```bash
+pip install -r requirements.txt
+# The runtime change is currently proposed in deepspeedai/DeepSpeed#8719.
+git clone https://github.com/deepspeedai/DeepSpeed.git
+git -C DeepSpeed fetch origin pull/8719/head:reflow
+git -C DeepSpeed switch reflow
+pip install -e ./DeepSpeed
+```
+
+Reflow is a DeepSpeed runtime feature, so you also need a DeepSpeed build that includes it — the `deepspeed.runtime.reflow` module and the `reflow_*` CPU-Adam/Lion ops (JIT-built on first use). No custom modeling code is required: the examples fine-tune plain Hugging Face Transformers models (`--model_name`) on `tatsu-lab/alpaca` by default (override with `--dataset_name`).
+
+
+The launch scripts default to FlashAttention 2. Install it separately after PyTorch:
+
+```bash
+pip install flash-attn --no-build-isolation
+```
+
+Alternatively, set `ATTN_IMPLEMENTATION=sdpa` or `eager` when running a launch script; these use PyTorch attention and do not require `flash-attn`. Llama checkpoints may require Hugging Face access approval and authentication.
+
+### 2. Enable Reflow (one block)
+
+Add a `reflow` block to the `zero_optimization` section of a ZeRO Stage 3 config with CPU optimizer offload (the scripts generate this for you):
+
+```jsonc
+"offload_optimizer": {
+    "device": "cpu",
+    "pin_memory": true
+},
+"reflow": {}           // the only addition Reflow requires
+```
+
+An empty block is enough; every Reflow setting (NUMA core binding, thread counts) is optional — see [Configuration](#configuration). Remove the block to fall back to plain ZeRO-Offload.
+
+### 3. Run a fine-tuning script
+
+Each script takes the mode (`reflow` or `zerooffload`) as the first argument and an optional **global batch size** as the second (it must be divisible by the GPU count):
+
+```bash
+# bash <script> <reflow|zerooffload> [global_batch_size]
+bash finetune_opt-30b_8gpu.sh   reflow        # OPT-30B, 8 GPUs
+bash finetune_opt-30b_8gpu.sh   zerooffload   # ZeRO-Offload baseline
+bash finetune_opt-30b_8gpu.sh   reflow 32     # override the batch size
+bash finetune_llama-7b_1gpu.sh  reflow        # Llama-7B, 1 GPU
+bash finetune_llama-8b_1gpu.sh  reflow        # Llama-8B, 1 GPU
+bash finetune_llama-70b_8gpu.sh reflow        # Llama-70B, 8 GPUs
+bash finetune_opt-350m-1gpu.sh  reflow        # OPT-350M, 1 GPU (smallest reproduction)
+bash finetune_opt-30b_1gpu.sh   reflow        # OPT-30B, 1 GPU
+```
+
+`run_compare.sh` runs both modes of a script back to back and prints peak GPU/host memory, the
+steady-state step time and the speedup:
+
+```bash
+bash run_compare.sh finetune_opt-350m-1gpu.sh
+```
+
+The generated config and training output directory are written next to the launch script. The comparison and equivalence checks put their logs in `compare_out/` and `bitexact_out/` by default (`OUT_DIR` overrides either). Scripts can be launched from another working directory.
+
+The default run stops after 10 microsteps and excludes the first microstep from its mean time. For a longer run:
+
+```bash
+BENCH_STEPS=100 WARMUP_STEPS=10 ATTN_IMPLEMENTATION=sdpa \
+  bash finetune_opt-350m-1gpu.sh reflow
+```
+
+`MAX_LENGTH`, `WORKER_CORES` (default 8), and `MAIN_CORES` (default 2 in these examples) also accept environment overrides. The runtime's main reservation default is 3; the scripts explicitly select 2. `gradient_clipping` is set to `0.0` in both modes for throughput comparison; enable clipping for training that requires it.
+
+### Optimizer: Adam (default) or Lion
+
+Every script defaults to Adam. The `_lion` scripts pass `--optimizer lion` and use `ReflowCPULion` (one momentum tensor instead of Adam's two — half the CPU optimizer state), and there are 4-GPU variants:
+
+```bash
+bash finetune_opt-30b_8gpu_lion.sh   reflow   # OPT-30B, 8 GPUs, Lion
+bash finetune_opt-30b_4gpu.sh        reflow   # OPT-30B, 4 GPUs, Adam
+bash finetune_opt-30b_4gpu_lion.sh   reflow   # OPT-30B, 4 GPUs, Lion
+bash finetune_llama-70b_8gpu_lion.sh reflow   # Llama-70B, 8 GPUs, Lion
+```
+
+To use Lion in your own run, just add `--optimizer lion` to `finetune_zero3.py` — Reflow transparently remaps the CPU optimizer to its async subclass.
+
+## Configuration
+
+The scripts use `enable_cpu_affinity: true`, `bucketwise_cores_per_worker: 8`, and `main_thread_cores: 2`. CPU counts are logical by default. Workers are pinned to their allocated CPU masks; restricting the process to its rank slice does not by itself pin the main thread to the reserved CPUs.
+
+Additional runtime options include `num_threads`, `state_update_cores`, `main_thread_core_type`, `pin_main_thread`, `worker_core_type`, `bucketwise_worker_affinity`, and `state_update_backward_cores`. Physical core reservations include SMT siblings; explicit main-thread pinning is optional. See the [Reflow configuration reference](https://github.com/deepspeedai/DeepSpeed/pull/8719) for the runtime change and its documentation. Standard ZeRO-3 knobs such as `sub_group_size` and `reduce_bucket_size` also apply.
+
+## Verifying bit-level equivalence
+
+Use deterministic attention and matching data order to compare the two paths. `--loss_check` enables deterministic algorithms and pins NCCL and cuBLAS settings before initialization, then logs each loss as a hexadecimal float. `--deterministic` enables deterministic algorithms without the additional comparison settings or loss logging.
+
+```bash
+bash check_bitexact.sh 1
+OPTIMIZER=lion BENCH_STEPS=100 bash check_bitexact.sh 1
+```
+
+The checker generates matching CPU-offload configs, uses eager attention, and requires every requested step in both logs before reporting `BIT-IDENTICAL`. It exits unsuccessfully if either job fails, a loss log is incomplete, or any recorded loss differs. `BATCH_SIZE` is the **per-GPU microbatch size** in this checker (default 4); the model and optimizer are selected with `MODEL_NAME` and `OPTIMIZER` (default OPT-350M and Adam).
+
+An exact loss match validates the recorded losses for that run. It does not establish bitwise equality of every optimizer state or guarantee equality for other hardware, library versions, or configurations. Deterministic comparison settings can reduce throughput; use ordinary runs for timing.
+
+## Requirements
+
+* A DeepSpeed build with Reflow, ZeRO Stage 3, and CPU optimizer offload (`offload_optimizer.device = "cpu"`). These examples do not exercise NVMe offload.
+* An x86 CPU with AVX (BF16 runs on both AVX-512 and AVX-256/AVX2). **AVX-512 is recommended** — AVX-256/AVX2 has no native BF16 support, so the whole BF16 path (gradient accumulation, FP32↔BF16 conversion, and the Adam updates) is emulated with extra instructions and can be slower. Enough cores per rank for the main thread plus the CPU-Adam workers.
+* Host RAM for the CPU-resident FP32 master + optimizer state (e.g. Llama-70B: ~840 GB with Adam, ~560 GB with Lion — Lion keeps one momentum tensor instead of two).
+* BF16 model parameters and gradients with FP32 master weights and optimizer states. FP16/FP32 model parameters, BF16 master weights/states, and `fp32_optimizer_states: false` are rejected by Reflow.
+
+## Memory
+
+Historical measurements on OPT-30B (8× B200, Adam, seq 512), steady-state during the training steps. These results have not been rerun with the current runtime revision:
+
+| | GPU (8 GPUs, total) | Host RAM |
+|---|---:|---:|
+| ZeRO-Offload | 466 GB | ~625 GB |
+| **Reflow** | 466 GB | ~549 GB |
+
+* **GPU memory is identical** — the optimizer state and FP32 master live on the CPU either way, so Reflow adds no GPU cost.
+* **Host RAM is ~76 GB lower with Reflow** (~12% on this run). ZeRO-Offload keeps a full CPU-side FP32 gradient buffer; in Reflow's bucketwise mode gradients are offloaded in BF16 and promoted to FP32 *inside* the CPU kernel, so that FP32 grad buffer is never held — only a compact half-precision grad buffer bounded by `sub_group_size` (not the full model) is kept.
+
+**The bigger memory lever is the optimizer.** Lion keeps one momentum tensor, Adam two, so Lion needs **2/3 (−33%)** of the CPU optimizer state (FP32 master + moments; 4 bytes/param for the master and 4 per momentum):
+
+| Model | FP32 master | Adam moments (2) | Lion moment (1) | Adam optimizer state | Lion optimizer state |
+|---|---:|---:|---:|---:|---:|
+| OPT-30B   | 120 GB | 240 GB | 120 GB | **~360 GB** | **~240 GB** |
+| Llama-70B | 280 GB | 560 GB | 280 GB | **~840 GB** | **~560 GB** |
+
+So for host-RAM-bound setups, Reflow + **Lion** is the lightest option: half-precision gradients (no FP32 grad buffer) plus a single momentum — Llama-70B fits its optimizer state in ~560 GB where Adam needs ~840 GB, and Lion is also the faster of the two at that scale (see [Performance](#performance)).
+
+## Performance
+
+Historical steady-state per-step time (avg of steps 2–10), Reflow vs the ZeRO-Offload baseline, on B200 GPUs, seq 2048, BF16. These measurements used earlier CPU placement settings and have not been rerun with the current revision or the default of 8 CPUs per worker. The speedup grows as the **per-rank optimizer shard** grows (bigger model or fewer GPUs) — there is more CPU optimizer work to hide inside the backward. Lion (one momentum) is lighter than Adam (two) and pulls ahead at scale.
+
+| Model | GPUs | Optimizer | Reflow | ZeRO-Offload | **Speedup** |
+|---|:---:|:---:|---:|---:|:---:|
+| Llama-8B  | 1 | Adam | 1474 ms | 5945 ms | **4.03×** |
+| Llama-13B | 1 | Adam | 1703 ms | 7676 ms | **4.51×** |
+| OPT-30B   | 4 | Adam | 3070 ms | 6801 ms | **2.22×** |
+| OPT-30B   | 4 | Lion | 3072 ms | 6538 ms | **2.13×** |
+| OPT-30B   | 8 | Adam | 3111 ms | 5898 ms | **1.90×** |
+| OPT-30B   | 8 | Lion | 3086 ms | 5627 ms | **1.82×** |
+| Llama-70B | 8 | Adam | 4190 ms | 9884 ms | **2.36×** |
+| Llama-70B | 8 | Lion | 3327 ms | 8701 ms | **2.61×** |
+
+At 70B, Lion is ~20% faster than Adam (3327 vs 4190 ms) and needs ~560 GB of CPU state vs Adam's ~840 GB, because its single momentum halves the offloaded optimizer state.
+
+### Example output
+
+OPT-30B, 8× B200, seq 2048, BF16. Each step logs `TFLOPS` (counting the activation-recompute forward) and `effective TFLOPS` (without it).
+
+**Reflow** (`bash finetune_opt-30b_8gpu.sh reflow`):
+
+```
+Reflow CPU-Adam SIMD: AVX-512
+Step    1 | Loss: 2.3329 | Time:  7712ms | TFLOPS(w/ recompute):  4267.00 | effective TFLOPS(w/o recompute):  3200.25 | Tokens/s:   2124
+Step    2 | Loss: 1.6865 | Time:  3110ms | TFLOPS(w/ recompute): 10580.75 | effective TFLOPS(w/o recompute):  7935.57 | Tokens/s:   5268
+Step    3 | Loss: 1.7266 | Time:  3098ms | TFLOPS(w/ recompute): 10621.35 | effective TFLOPS(w/o recompute):  7966.01 | Tokens/s:   5288
+Step    4 | Loss: 1.4287 | Time:  3108ms | TFLOPS(w/ recompute): 10589.88 | effective TFLOPS(w/o recompute):  7942.41 | Tokens/s:   5272
+Step    5 | Loss: 1.3309 | Time:  3121ms | TFLOPS(w/ recompute): 10545.24 | effective TFLOPS(w/o recompute):  7908.93 | Tokens/s:   5250
+Step    6 | Loss: 1.3002 | Time:  3126ms | TFLOPS(w/ recompute): 10527.12 | effective TFLOPS(w/o recompute):  7895.34 | Tokens/s:   5241
+Step    7 | Loss: 1.6794 | Time:  3117ms | TFLOPS(w/ recompute): 10558.52 | effective TFLOPS(w/o recompute):  7918.89 | Tokens/s:   5257
+Step    8 | Loss: 1.5415 | Time:  3108ms | TFLOPS(w/ recompute): 10588.21 | effective TFLOPS(w/o recompute):  7941.16 | Tokens/s:   5272
+Step    9 | Loss: 1.5562 | Time:  3120ms | TFLOPS(w/ recompute): 10547.18 | effective TFLOPS(w/o recompute):  7910.39 | Tokens/s:   5251
+Step   10 | Loss: 1.4714 | Time:  3153ms | TFLOPS(w/ recompute): 10436.84 | effective TFLOPS(w/o recompute):  7827.63 | Tokens/s:   5196
+```
+
+**ZeRO-Offload** (`bash finetune_opt-30b_8gpu.sh zerooffload`):
+
+```
+Step    1 | Loss: 2.3329 | Time:  8753ms | TFLOPS(w/ recompute):  3759.61 | effective TFLOPS(w/o recompute):  2819.71 | Tokens/s:   1872
+Step    2 | Loss: 1.6855 | Time:  5873ms | TFLOPS(w/ recompute):  5603.57 | effective TFLOPS(w/o recompute):  4202.68 | Tokens/s:   2790
+Step    3 | Loss: 1.7261 | Time:  5873ms | TFLOPS(w/ recompute):  5602.89 | effective TFLOPS(w/o recompute):  4202.16 | Tokens/s:   2789
+Step    4 | Loss: 1.4260 | Time:  5860ms | TFLOPS(w/ recompute):  5616.12 | effective TFLOPS(w/o recompute):  4212.09 | Tokens/s:   2796
+Step    5 | Loss: 1.3316 | Time:  5799ms | TFLOPS(w/ recompute):  5675.13 | effective TFLOPS(w/o recompute):  4256.35 | Tokens/s:   2825
+Step    6 | Loss: 1.2991 | Time:  5892ms | TFLOPS(w/ recompute):  5585.03 | effective TFLOPS(w/o recompute):  4188.77 | Tokens/s:   2781
+Step    7 | Loss: 1.6796 | Time:  5942ms | TFLOPS(w/ recompute):  5537.87 | effective TFLOPS(w/o recompute):  4153.40 | Tokens/s:   2757
+Step    8 | Loss: 1.5477 | Time:  5947ms | TFLOPS(w/ recompute):  5533.28 | effective TFLOPS(w/o recompute):  4149.96 | Tokens/s:   2755
+Step    9 | Loss: 1.5629 | Time:  5871ms | TFLOPS(w/ recompute):  5605.00 | effective TFLOPS(w/o recompute):  4203.75 | Tokens/s:   2791
+Step   10 | Loss: 1.4747 | Time:  5790ms | TFLOPS(w/ recompute):  5683.77 | effective TFLOPS(w/o recompute):  4262.83 | Tokens/s:   2830
+```
+
+Reflow is ~1.9× faster (≈3.1 s vs ≈5.9 s per step) by overlapping the CPU-Adam step with the backward pass.
+
+## Notes
+
+* **NUMA core binding**: every launch script passes `--bind_cores_to_rank`. Reflow allocates its main reservation and worker masks within the rank's CPU slice. Use physical core reservations and `pin_main_thread` if you need explicit separation of the main thread and worker SMT siblings.
+* These examples use CPU optimizer offload. NVMe offload is not validated here.
+* Use `engine.step()` after `deepspeed.initialize()`; direct calls to `ReflowCPUAdam.step()` or `ReflowCPULion.step()` are rejected.
+* `--save_checkpoint` requires participation by every rank. The script saves the tokenizer on rank 0 and propagates checkpoint failures.
+* **AVX-512 is strongly recommended** for the CPU optimizer; on AVX2 the BF16 path is emulated and slower (see [Requirements](#requirements)).
+
+## Related resources
+
+* [DeepSpeed ZeRO-Offload tutorial](https://www.deepspeed.ai/tutorials/zero-offload/)
+* [ZeRO Stage 3 config options](https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training)
+* Reflow runtime module: `deepspeed/runtime/reflow/` — `reflow_stage3.py`, `reflow_cpu_adam.py`, `reflow_cpu_lion.py`
+
+## Validation
+
+Run the CPU launcher and data-preprocessing checks without initializing CUDA:
+
+```bash
+CUDA_VISIBLE_DEVICES='' DS_ACCELERATOR=cpu python3 -m unittest discover -s tests -v
+```
+
+These checks validate launcher arguments, generated configurations, failure propagation, loss-log completeness, and padding masks. They do not replace GPU training validation. Use `check_bitexact.sh` on available GPUs to compare training losses and `run_compare.sh` to measure throughput. GPU memory sampling reports device usage, including unrelated jobs; host memory sampling sums RSS of matching training processes and can overcount shared pages.
diff --git a/training/DeepSpeed-Reflow/check_bitexact.sh b/training/DeepSpeed-Reflow/check_bitexact.sh
new file mode 100755
index 000000000..c0bf4e022
--- /dev/null
+++ b/training/DeepSpeed-Reflow/check_bitexact.sh
@@ -0,0 +1,117 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+# check_bitexact.sh -- verify that Reflow is bit-identical to ZeRO-Infinity
+# (optimizer offload + ZeRO Stage 3).
+#
+# Usage:  bash check_bitexact.sh [num_gpus]        (default: 1)
+#
+# Runs the same OPT-350M job twice, once with each mode, under full
+# determinism (--loss_check), then diffs the per-step hex-float losses.
+# Everything is self-contained: the two configs are written here and
+# differ only in the Reflow keys.
+#
+# --loss_check pins the NCCL collectives and forces deterministic
+# algorithms, so this is slower than a normal run. It is for
+# verification only -- speed and memory are measured by run_compare.sh.
+
+set -euo pipefail
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+
+NUM_GPUS=${1:-1}
+BATCH_SIZE=${BATCH_SIZE:-4}
+OPTIMIZER=${OPTIMIZER:-adam}
+BENCH_STEPS=${BENCH_STEPS:-6}
+MODEL_NAME=${MODEL_NAME:-facebook/opt-350m}
+
+# CPU affinity knobs, as in the finetune_*.sh scripts. Keep
+# main_thread_cores + bucketwise_cores_per_worker within the cores
+# visible to each rank.
+MAIN_CORES=${MAIN_CORES:-2}
+WORKER_CORES=${WORKER_CORES:-8}
+
+OUT_DIR=${OUT_DIR:-bitexact_out}
+for value in "$NUM_GPUS" "$BATCH_SIZE" "$BENCH_STEPS" "$MAIN_CORES" "$WORKER_CORES"; do
+    if [[ ! "$value" =~ ^[1-9][0-9]*$ ]]; then
+        echo "GPU, batch, step and core counts must be positive integers" >&2
+        exit 2
+    fi
+done
+mkdir -p "$OUT_DIR"
+
+# ---------------------------------------------------------------- configs
+# Same template as the finetune_*.sh scripts. overlap_comm stays false:
+# with overlapped reductions the reduction order becomes timing-dependent
+# and the two modes can legitimately diverge in the last bits.
+write_config() {   # $1 = path, $2 = reflow|zerooffload
+    if [ "$2" = "reflow" ]; then
+        REFLOW_BLOCK=',
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": '"$WORKER_CORES"',
+            "main_thread_cores": '"$MAIN_CORES"'
+        }'
+    else
+        REFLOW_BLOCK=""
+    fi
+    cat > "$1" << EOF
+{
+    "train_micro_batch_size_per_gpu": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 4e8,
+        "sub_group_size": 4e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }$REFLOW_BLOCK
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+}
+
+REFLOW_CFG="$OUT_DIR/opt-350m_reflow_config.json"
+ZERO_CFG="$OUT_DIR/opt-350m_zerooffload_config.json"
+write_config "$REFLOW_CFG" reflow
+write_config "$ZERO_CFG"   zerooffload
+
+# ------------------------------------------------------------------- runs
+# Eager attention is required: flash-attention's backward is
+# non-deterministic and would mask a real difference.
+COMMON=(--model_name "$MODEL_NAME"
+    --optimizer "$OPTIMIZER" --loss_check --attn_implementation eager
+    --lr 1e-5 --batch_size "$BATCH_SIZE" --max_length 512
+    --bench_steps "$BENCH_STEPS" --warmup_steps 0 --output_dir "$OUT_DIR/out")
+
+echo "[1/2] Reflow ..."
+deepspeed --num_gpus="$NUM_GPUS" --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py" \
+    --deepspeed_config="$REFLOW_CFG" \
+    "${COMMON[@]}" > "$OUT_DIR/reflow.log" 2>&1
+
+echo "[2/2] ZeRO-Infinity baseline ..."
+deepspeed --num_gpus="$NUM_GPUS" --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py" \
+    --deepspeed_config="$ZERO_CFG" \
+    "${COMMON[@]}" > "$OUT_DIR/zero.log" 2>&1
+
+# Ignore timestamps but require both jobs to have completed every requested step.
+python3 - "$OUT_DIR/reflow.log" "$OUT_DIR/zero.log" "$BENCH_STEPS" <<'PYTHON'
+import re
+import sys
+from pathlib import Path
+
+expected_steps = list(range(1, int(sys.argv[3]) + 1))
+losses = []
+for filename in sys.argv[1:3]:
+    matches = re.findall(r"BITLOSS step (\d+) hex=(\S+)", Path(filename).read_text())
+    if [int(step) for step, _ in matches] != expected_steps:
+        sys.exit(f"Incomplete loss log: {filename}; expected steps 1 through {expected_steps[-1]}")
+    losses.append([value for _, value in matches])
+if losses[0] != losses[1]:
+    sys.exit(f"MISMATCH -- losses differ; see {sys.argv[1]} and {sys.argv[2]}")
+print("BIT-IDENTICAL")
+PYTHON
diff --git a/training/DeepSpeed-Reflow/finetune_llama-13b_1gpu.sh b/training/DeepSpeed-Reflow/finetune_llama-13b_1gpu.sh
new file mode 100755
index 000000000..c680d3cf4
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_llama-13b_1gpu.sh
@@ -0,0 +1,153 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "Llama 13B Fine-tuning with DeepSpeed Reflow on 1 GPU"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-4}
+
+GPUS_PER_NODE=1
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="meta-llama/Llama-2-13b-hf"
+OUTPUT_DIR="${SCRIPT_DIR}/llama-13b_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/llama-13b_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=${MAX_LENGTH:-2048}
+ATTN_IMPLEMENTATION=${ATTN_IMPLEMENTATION:-flash_attention_2}
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+WANDB_PROJECT="llama-13b"
+WANDB_RUN_NAME="llama-13b-$MODE"
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+SAVE_CHECKPOINT_ARG=""
+if [ "$SAVE_CHECKPOINT" = "true" ]; then
+    SAVE_CHECKPOINT_ARG="--save_checkpoint"
+fi
+
+WANDB_FLAG=""
+if [ "$USE_WANDB" = "true" ]; then
+    WANDB_FLAG="--use_wandb"
+fi
+
+DETERMINISTIC_FLAG=""
+if [ "$DETERMINISTIC" = "true" ]; then
+    DETERMINISTIC_FLAG="--deterministic"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 4e8,
+        "sub_group_size": 1e9,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 4e8,
+        "sub_group_size": 1e9,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+    $SAVE_CHECKPOINT_ARG
+    $WANDB_FLAG
+    --wandb_project "$WANDB_PROJECT"
+    --wandb_run_name "$WANDB_RUN_NAME"
+    $DETERMINISTIC_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_llama-70b_8gpu.sh b/training/DeepSpeed-Reflow/finetune_llama-70b_8gpu.sh
new file mode 100755
index 000000000..b1ea85142
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_llama-70b_8gpu.sh
@@ -0,0 +1,157 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "Llama 70B Fine-tuning with DeepSpeed Reflow on 8 GPUs"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-16}
+
+GPUS_PER_NODE=8
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="meta-llama/Llama-3.3-70B-Instruct"
+OUTPUT_DIR="${SCRIPT_DIR}/llama-70b_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/llama-70b_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=${MAX_LENGTH:-2048}
+ATTN_IMPLEMENTATION=${ATTN_IMPLEMENTATION:-flash_attention_2}
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+WANDB_PROJECT="llama-70b"
+WANDB_RUN_NAME="llama-70b-$MODE"
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+SAVE_CHECKPOINT_ARG=""
+if [ "$SAVE_CHECKPOINT" = "true" ]; then
+    SAVE_CHECKPOINT_ARG="--save_checkpoint"
+fi
+
+WANDB_FLAG=""
+if [ "$USE_WANDB" = "true" ]; then
+    WANDB_FLAG="--use_wandb"
+fi
+
+DETERMINISTIC_FLAG=""
+if [ "$DETERMINISTIC" = "true" ]; then
+    DETERMINISTIC_FLAG="--deterministic"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "stage3_param_persistence_threshold": 1e4,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "stage3_param_persistence_threshold": 1e4,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+    $SAVE_CHECKPOINT_ARG
+    $WANDB_FLAG
+    --wandb_project "$WANDB_PROJECT"
+    --wandb_run_name "$WANDB_RUN_NAME"
+    $DETERMINISTIC_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_llama-70b_8gpu_lion.sh b/training/DeepSpeed-Reflow/finetune_llama-70b_8gpu_lion.sh
new file mode 100755
index 000000000..eeed51126
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_llama-70b_8gpu_lion.sh
@@ -0,0 +1,158 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "Llama 70B Fine-tuning with DeepSpeed Reflow on 8 GPUs (Lion)"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-16}
+
+GPUS_PER_NODE=8
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="meta-llama/Llama-3.3-70B-Instruct"
+OUTPUT_DIR="${SCRIPT_DIR}/llama-70b_lion_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/llama-70b_lion_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=${MAX_LENGTH:-2048}
+ATTN_IMPLEMENTATION=${ATTN_IMPLEMENTATION:-flash_attention_2}
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+WANDB_PROJECT="llama-70b"
+WANDB_RUN_NAME="llama-70b-$MODE"
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+SAVE_CHECKPOINT_ARG=""
+if [ "$SAVE_CHECKPOINT" = "true" ]; then
+    SAVE_CHECKPOINT_ARG="--save_checkpoint"
+fi
+
+WANDB_FLAG=""
+if [ "$USE_WANDB" = "true" ]; then
+    WANDB_FLAG="--use_wandb"
+fi
+
+DETERMINISTIC_FLAG=""
+if [ "$DETERMINISTIC" = "true" ]; then
+    DETERMINISTIC_FLAG="--deterministic"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "stage3_param_persistence_threshold": 1e4,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "stage3_param_persistence_threshold": 1e4,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --optimizer lion
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+    $SAVE_CHECKPOINT_ARG
+    $WANDB_FLAG
+    --wandb_project "$WANDB_PROJECT"
+    --wandb_run_name "$WANDB_RUN_NAME"
+    $DETERMINISTIC_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_llama-7b_1gpu.sh b/training/DeepSpeed-Reflow/finetune_llama-7b_1gpu.sh
new file mode 100755
index 000000000..1b030ce9d
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_llama-7b_1gpu.sh
@@ -0,0 +1,152 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "Llama 7B Fine-tuning with DeepSpeed Reflow on 1 GPU"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-4}
+
+GPUS_PER_NODE=1
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="meta-llama/Llama-2-7b-hf"
+OUTPUT_DIR="${SCRIPT_DIR}/llama-7b_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/llama-7b_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=4096
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+WANDB_PROJECT="llama-7b"
+WANDB_RUN_NAME="llama-7b-$MODE"
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+SAVE_CHECKPOINT_ARG=""
+if [ "$SAVE_CHECKPOINT" = "true" ]; then
+    SAVE_CHECKPOINT_ARG="--save_checkpoint"
+fi
+
+WANDB_FLAG=""
+if [ "$USE_WANDB" = "true" ]; then
+    WANDB_FLAG="--use_wandb"
+fi
+
+DETERMINISTIC_FLAG=""
+if [ "$DETERMINISTIC" = "true" ]; then
+    DETERMINISTIC_FLAG="--deterministic"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 4e8,
+        "sub_group_size": 4e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 4e8,
+        "sub_group_size": 4e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+    $SAVE_CHECKPOINT_ARG
+    $WANDB_FLAG
+    --wandb_project "$WANDB_PROJECT"
+    --wandb_run_name "$WANDB_RUN_NAME"
+    $DETERMINISTIC_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_llama-8b_1gpu.sh b/training/DeepSpeed-Reflow/finetune_llama-8b_1gpu.sh
new file mode 100755
index 000000000..9b01a20f1
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_llama-8b_1gpu.sh
@@ -0,0 +1,152 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "Llama 8B Fine-tuning with DeepSpeed Reflow on 1 GPU"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-4}
+
+GPUS_PER_NODE=1
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="meta-llama/Llama-3.1-8B"
+OUTPUT_DIR="${SCRIPT_DIR}/llama-8b_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/llama-8b_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=4096
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+WANDB_PROJECT="llama-8b"
+WANDB_RUN_NAME="llama-8b-$MODE"
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+SAVE_CHECKPOINT_ARG=""
+if [ "$SAVE_CHECKPOINT" = "true" ]; then
+    SAVE_CHECKPOINT_ARG="--save_checkpoint"
+fi
+
+WANDB_FLAG=""
+if [ "$USE_WANDB" = "true" ]; then
+    WANDB_FLAG="--use_wandb"
+fi
+
+DETERMINISTIC_FLAG=""
+if [ "$DETERMINISTIC" = "true" ]; then
+    DETERMINISTIC_FLAG="--deterministic"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 4e8,
+        "sub_group_size": 4e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 4e8,
+        "sub_group_size": 4e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+    $SAVE_CHECKPOINT_ARG
+    $WANDB_FLAG
+    --wandb_project "$WANDB_PROJECT"
+    --wandb_run_name "$WANDB_RUN_NAME"
+    $DETERMINISTIC_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_opt-30b_1gpu.sh b/training/DeepSpeed-Reflow/finetune_opt-30b_1gpu.sh
new file mode 100755
index 000000000..ef33d7b45
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_opt-30b_1gpu.sh
@@ -0,0 +1,155 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "OPT-30B Fine-tuning with DeepSpeed Reflow on 1 GPUs"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-8}
+
+GPUS_PER_NODE=1
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="facebook/opt-30b"
+OUTPUT_DIR="${SCRIPT_DIR}/opt-30b_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/opt-30b_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=${MAX_LENGTH:-2048}
+ATTN_IMPLEMENTATION=${ATTN_IMPLEMENTATION:-flash_attention_2}
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+WANDB_PROJECT="opt-30b"
+WANDB_RUN_NAME="opt-30b-$MODE"
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+SAVE_CHECKPOINT_ARG=""
+if [ "$SAVE_CHECKPOINT" = "true" ]; then
+    SAVE_CHECKPOINT_ARG="--save_checkpoint"
+fi
+
+WANDB_FLAG=""
+if [ "$USE_WANDB" = "true" ]; then
+    WANDB_FLAG="--use_wandb"
+fi
+
+DETERMINISTIC_FLAG=""
+if [ "$DETERMINISTIC" = "true" ]; then
+    DETERMINISTIC_FLAG="--deterministic"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+    $SAVE_CHECKPOINT_ARG
+    $WANDB_FLAG
+    --wandb_project "$WANDB_PROJECT"
+    --wandb_run_name "$WANDB_RUN_NAME"
+    $DETERMINISTIC_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_opt-30b_4gpu.sh b/training/DeepSpeed-Reflow/finetune_opt-30b_4gpu.sh
new file mode 100755
index 000000000..01e918c8a
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_opt-30b_4gpu.sh
@@ -0,0 +1,155 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "OPT-30B Fine-tuning with DeepSpeed Reflow on 4 GPUs"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-32}
+
+GPUS_PER_NODE=4
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="facebook/opt-30b"
+OUTPUT_DIR="${SCRIPT_DIR}/opt-30b_4gpu_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/opt-30b_4gpu_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=${MAX_LENGTH:-2048}
+ATTN_IMPLEMENTATION=${ATTN_IMPLEMENTATION:-flash_attention_2}
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+WANDB_PROJECT="opt-30b"
+WANDB_RUN_NAME="opt-30b-$MODE"
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+SAVE_CHECKPOINT_ARG=""
+if [ "$SAVE_CHECKPOINT" = "true" ]; then
+    SAVE_CHECKPOINT_ARG="--save_checkpoint"
+fi
+
+WANDB_FLAG=""
+if [ "$USE_WANDB" = "true" ]; then
+    WANDB_FLAG="--use_wandb"
+fi
+
+DETERMINISTIC_FLAG=""
+if [ "$DETERMINISTIC" = "true" ]; then
+    DETERMINISTIC_FLAG="--deterministic"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+    $SAVE_CHECKPOINT_ARG
+    $WANDB_FLAG
+    --wandb_project "$WANDB_PROJECT"
+    --wandb_run_name "$WANDB_RUN_NAME"
+    $DETERMINISTIC_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_opt-30b_4gpu_lion.sh b/training/DeepSpeed-Reflow/finetune_opt-30b_4gpu_lion.sh
new file mode 100755
index 000000000..c5c9d1eb0
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_opt-30b_4gpu_lion.sh
@@ -0,0 +1,134 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "OPT-30B Lion Fine-tuning with DeepSpeed Reflow on 4 GPUs"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-32}
+
+GPUS_PER_NODE=4
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="facebook/opt-30b"
+OUTPUT_DIR="${SCRIPT_DIR}/opt-30b_4gpu_lion_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/opt-30b_4gpu_lion_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=${MAX_LENGTH:-2048}
+ATTN_IMPLEMENTATION=${ATTN_IMPLEMENTATION:-flash_attention_2}
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --optimizer lion
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_opt-30b_8gpu.sh b/training/DeepSpeed-Reflow/finetune_opt-30b_8gpu.sh
new file mode 100755
index 000000000..1d14ba7bf
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_opt-30b_8gpu.sh
@@ -0,0 +1,155 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "OPT-30B Fine-tuning with DeepSpeed Reflow on 8 GPUs"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-64}
+
+GPUS_PER_NODE=8
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="facebook/opt-30b"
+OUTPUT_DIR="${SCRIPT_DIR}/opt-30b_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/opt-30b_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=${MAX_LENGTH:-2048}
+ATTN_IMPLEMENTATION=${ATTN_IMPLEMENTATION:-flash_attention_2}
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+WANDB_PROJECT="opt-30b"
+WANDB_RUN_NAME="opt-30b-$MODE"
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+SAVE_CHECKPOINT_ARG=""
+if [ "$SAVE_CHECKPOINT" = "true" ]; then
+    SAVE_CHECKPOINT_ARG="--save_checkpoint"
+fi
+
+WANDB_FLAG=""
+if [ "$USE_WANDB" = "true" ]; then
+    WANDB_FLAG="--use_wandb"
+fi
+
+DETERMINISTIC_FLAG=""
+if [ "$DETERMINISTIC" = "true" ]; then
+    DETERMINISTIC_FLAG="--deterministic"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+    $SAVE_CHECKPOINT_ARG
+    $WANDB_FLAG
+    --wandb_project "$WANDB_PROJECT"
+    --wandb_run_name "$WANDB_RUN_NAME"
+    $DETERMINISTIC_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_opt-30b_8gpu_lion.sh b/training/DeepSpeed-Reflow/finetune_opt-30b_8gpu_lion.sh
new file mode 100644
index 000000000..b0b77ad2f
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_opt-30b_8gpu_lion.sh
@@ -0,0 +1,134 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "OPT-30B Lion Fine-tuning with DeepSpeed Reflow on 8 GPUs"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-64}
+
+GPUS_PER_NODE=8
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="facebook/opt-30b"
+OUTPUT_DIR="${SCRIPT_DIR}/opt-30b_lion_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/opt-30b_lion_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=${MAX_LENGTH:-2048}
+ATTN_IMPLEMENTATION=${ATTN_IMPLEMENTATION:-flash_attention_2}
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 5e8,
+        "sub_group_size": 1e9,
+        "stage3_prefetch_bucket_size": 5e8,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --optimizer lion
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_opt-350m-1gpu.sh b/training/DeepSpeed-Reflow/finetune_opt-350m-1gpu.sh
new file mode 100755
index 000000000..8d5129400
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_opt-350m-1gpu.sh
@@ -0,0 +1,153 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+set -euo pipefail
+
+echo "================================================"
+echo "OPT 350M Fine-tuning with DeepSpeed Reflow on 1 GPU"
+echo "================================================"
+
+# MODE options: "reflow" or "zerooffload"
+MODE=${1:-}
+BATCH_SIZE=${2:-4}
+
+GPUS_PER_NODE=1
+if [[ "$MODE" != "reflow" && "$MODE" != "zerooffload" ]]; then
+    echo "Usage: bash $0 <reflow|zerooffload> [global_batch_size]" >&2
+    exit 2
+fi
+if [[ ! "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || (( BATCH_SIZE % GPUS_PER_NODE != 0 )); then
+    echo "global_batch_size must be positive and divisible by $GPUS_PER_NODE" >&2
+    exit 2
+fi
+MICRO_BATCH_SIZE=$((BATCH_SIZE / GPUS_PER_NODE))
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+MODEL_NAME="facebook/opt-350m"
+OUTPUT_DIR="${SCRIPT_DIR}/opt-350m_${MODE}_output"
+DS_CONFIG_JSON="${SCRIPT_DIR}/opt-350m_${MODE}_config.json"
+
+mkdir -p "$OUTPUT_DIR"
+
+# Script argument parameters
+ACTIVATION_CHECKPOINTING=true
+SAVE_CHECKPOINT=false
+MAX_LENGTH=${MAX_LENGTH:-2048}
+ATTN_IMPLEMENTATION=${ATTN_IMPLEMENTATION:-flash_attention_2}
+LOG_INTERVAL=1
+DATASET_NAME="tatsu-lab/alpaca"
+DATASET_PERCENTAGE=10.0
+USE_WANDB=false
+WANDB_PROJECT="opt-350m"
+WANDB_RUN_NAME="opt-350m-$MODE"
+DETERMINISTIC=false
+BENCH_STEPS=${BENCH_STEPS:-10}
+WARMUP_STEPS=${WARMUP_STEPS:-1}
+
+EPOCHS=1
+LR=1e-5
+WEIGHT_DECAY=0.01
+SEED=42
+
+ACTIVATION_CHECKPOINTING_FLAG=""
+if [ "$ACTIVATION_CHECKPOINTING" = "true" ]; then
+    ACTIVATION_CHECKPOINTING_FLAG="--activation_checkpointing"
+fi
+
+SAVE_CHECKPOINT_ARG=""
+if [ "$SAVE_CHECKPOINT" = "true" ]; then
+    SAVE_CHECKPOINT_ARG="--save_checkpoint"
+fi
+
+WANDB_FLAG=""
+if [ "$USE_WANDB" = "true" ]; then
+    WANDB_FLAG="--use_wandb"
+fi
+
+DETERMINISTIC_FLAG=""
+if [ "$DETERMINISTIC" = "true" ]; then
+    DETERMINISTIC_FLAG="--deterministic"
+fi
+
+# Reflow uses the same ZeRO-3 configuration as the baseline plus the reflow block.
+if [ "$MODE" = "reflow" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 4e8,
+        "sub_group_size": 1e9,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "reflow": {
+            "enable_cpu_affinity": true,
+            "bucketwise_cores_per_worker": ${WORKER_CORES:-8},
+            "main_thread_cores": ${MAIN_CORES:-2}
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+
+elif [ "$MODE" = "zerooffload" ]; then
+cat > "$DS_CONFIG_JSON" << EOF
+{
+    "train_batch_size": $BATCH_SIZE,
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": 0.0,
+    "bf16": { "enabled": true },
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": false,
+        "reduce_bucket_size": 4e8,
+        "sub_group_size": 1e9,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        }
+    },
+    "wall_clock_breakdown": true
+}
+EOF
+fi
+
+
+CMD=(
+deepspeed --num_gpus=$GPUS_PER_NODE --bind_cores_to_rank "$SCRIPT_DIR/finetune_zero3.py"
+    --deepspeed_config="$DS_CONFIG_JSON"
+    --model_name "$MODEL_NAME"
+    --attn_implementation "$ATTN_IMPLEMENTATION"
+    --num_train_epochs "$EPOCHS"
+    --lr "$LR"
+    --batch_size "$MICRO_BATCH_SIZE"
+    --weight_decay "$WEIGHT_DECAY"
+    --output_dir "$OUTPUT_DIR"
+    --seed "$SEED"
+    --max_length "$MAX_LENGTH"
+    --log_interval "$LOG_INTERVAL"
+    --dataset_name "$DATASET_NAME"
+    --dataset_percentage "$DATASET_PERCENTAGE"
+    --bench_steps "$BENCH_STEPS"
+    --warmup_steps "$WARMUP_STEPS"
+    $ACTIVATION_CHECKPOINTING_FLAG
+    $SAVE_CHECKPOINT_ARG
+    $WANDB_FLAG
+    --wandb_project "$WANDB_PROJECT"
+    --wandb_run_name "$WANDB_RUN_NAME"
+    $DETERMINISTIC_FLAG
+)
+
+echo "Starting training with MODE $MODE"
+echo "================================================"
+"${CMD[@]}"
+
+echo "================================================"
+echo "Training completed"
+echo "================================================"
diff --git a/training/DeepSpeed-Reflow/finetune_zero3.py b/training/DeepSpeed-Reflow/finetune_zero3.py
new file mode 100644
index 000000000..0c1f995bc
--- /dev/null
+++ b/training/DeepSpeed-Reflow/finetune_zero3.py
@@ -0,0 +1,538 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""
+Fine-tuning script for language models using DeepSpeed ZeRO Stage 3.
+"""
+
+import argparse
+import os
+import time
+import logging
+from datetime import datetime
+from typing import Dict, Any
+
+import torch
+import deepspeed
+import wandb
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from transformers import (AutoModelForCausalLM, AutoTokenizer, default_data_collator, set_seed,
+                          enable_full_determinism)
+from deepspeed import comm as dist
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+def setup_logger(rank: int = 0, log_level: str = "INFO") -> logging.Logger:
+    logger = logging.getLogger("finetune_zero3")
+    logger.handlers.clear()
+    numeric_level = getattr(logging, log_level.upper(), logging.INFO)
+    logger.setLevel(numeric_level)
+
+    if rank == 0:
+        handler = logging.StreamHandler()
+        handler.setLevel(numeric_level)
+        formatter = logging.Formatter(fmt='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+                                      datefmt='%Y-%m-%d %H:%M:%S')
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+
+    return logger
+
+
+# Constants
+DEFAULT_OPTIMIZER_BETAS = (0.9, 0.999)
+MS_PER_SECOND = 1000
+TFLOPS_DENOMINATOR = 1e12
+
+# Alpaca dataset formatting
+ALPACA_INSTRUCTION_TEMPLATE = "### Instruction:\n{instruction}\n\n"
+ALPACA_INPUT_TEMPLATE = "### Input:\n{input}\n\n"
+ALPACA_RESPONSE_TEMPLATE = "### Response:\n{output}"
+
+
+def get_parameter_count(parameter: torch.nn.Parameter) -> int:
+    return parameter.ds_numel if hasattr(parameter, "ds_tensor") else parameter.numel()
+
+
+def estimate_transformer_tflops(seq_len: int,
+                                model_size: int,
+                                num_layers: int,
+                                hidden_size: int,
+                                use_activation_checkpointing: bool = False) -> float:
+    """
+    Estimate TFLOPS for decoder-only dense models.
+    """
+    coefficient = 4 if use_activation_checkpointing else 3
+    tflops = (2 * coefficient * model_size * seq_len +
+              2 * 2 * coefficient * num_layers * hidden_size * seq_len**2) / TFLOPS_DENOMINATOR
+    return tflops
+
+
+def preprocess_alpaca_example(example: Dict[str, str],
+                              tokenizer: AutoTokenizer,
+                              max_length: int = 2048) -> Dict[str, Any]:
+    prompt = ALPACA_INSTRUCTION_TEMPLATE.format(instruction=example['instruction'])
+
+    if example.get("input", "").strip():
+        prompt += ALPACA_INPUT_TEMPLATE.format(input=example['input'])
+
+    prompt += ALPACA_RESPONSE_TEMPLATE.format(output=example['output'])
+
+    tokenized = tokenizer(prompt, truncation=True, max_length=max_length, padding="max_length", return_tensors=None)
+
+    # Mask padding out of the loss (HF CrossEntropy ignore_index = -100); otherwise the ~max_length
+    # padding dominates the loss. attention_mask is 0 exactly on padding positions.
+    tokenized["labels"] = [
+        token if mask == 1 else -100 for token, mask in zip(tokenized["input_ids"], tokenized["attention_mask"])
+    ]
+
+    return tokenized
+
+
+def detect_moe_model(model: AutoModelForCausalLM) -> bool:
+    moe_config_attrs = ['num_local_experts', 'moe_layers', 'num_experts', 'expert_capacity', 'router_aux_loss_coef']
+
+    for attr in moe_config_attrs:
+        if hasattr(model.config, attr):
+            return True
+    return False
+
+
+def create_experiment_name(args: argparse.Namespace) -> str:
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    model_name_short = args.model_name.split("/")[-1]
+    activation_checkpointing = 1 if args.activation_checkpointing else 0
+
+    exp_name = (f"{model_name_short}_bs{args.batch_size}_seq{args.max_length}"
+                f"_ac{activation_checkpointing}_T{timestamp}")
+    return exp_name
+
+
+def load_tokenizer(model_name: str, logger: logging.Logger) -> AutoTokenizer:
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        logger.debug(f"Set pad_token to eos_token: {tokenizer.eos_token}")
+
+    return tokenizer
+
+
+def load_model(model_name: str, attn_implementation: str, logger: logging.Logger) -> AutoModelForCausalLM:
+    logger.debug(f"Loading model: {model_name}")
+    logger.debug(f"Attention implementation: {attn_implementation}")
+
+    model = AutoModelForCausalLM.from_pretrained(model_name,
+                                                 torch_dtype=torch.bfloat16,
+                                                 attn_implementation=attn_implementation)
+
+    return model
+
+
+def setup_model_training(model: torch.nn.Module,
+                         use_activation_checkpointing: bool = True,
+                         logger: logging.Logger = None) -> None:
+    if use_activation_checkpointing:
+        if logger:
+            logger.debug("Enabling gradient checkpointing...")
+        if hasattr(model.config, 'use_cache'):
+            model.config.use_cache = False
+        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+
+
+def create_optimizer(model: AutoModelForCausalLM, args: argparse.Namespace) -> Any:
+    # With Reflow enabled (a "reflow" block), DeepSpeed transparently remaps the CPU optimizer below to
+    # its ReflowCPUAdam / ReflowCPULion subclass, so the same client optimizer drives both the
+    # ZeRO-Offload baseline and Reflow.
+    if args.optimizer == "lion":
+        from deepspeed.ops.lion import DeepSpeedCPULion
+        optimizer = DeepSpeedCPULion(
+            model.parameters(),
+            lr=args.lr,
+            betas=DEFAULT_OPTIMIZER_BETAS,
+            weight_decay=args.weight_decay,
+        )
+    else:
+        from deepspeed.ops.adam import DeepSpeedCPUAdam
+        optimizer = DeepSpeedCPUAdam(
+            model.parameters(),
+            lr=args.lr,
+            betas=DEFAULT_OPTIMIZER_BETAS,
+            eps=1e-8,
+            weight_decay=args.weight_decay,
+        )
+    return optimizer
+
+
+def load_and_preprocess_dataset(dataset_name: str, dataset_percentage: float, tokenizer: AutoTokenizer,
+                                max_length: int, logger: logging.Logger) -> Any:
+    logger.debug(f"Loading dataset: {dataset_name}")
+
+    dataset = load_dataset(dataset_name)
+    original_size = len(dataset["train"])
+
+    if dataset_percentage < 100.0:
+        subset_size = int(original_size * dataset_percentage / 100.0)
+        if subset_size == 0:
+            raise ValueError("dataset_percentage selects no training examples")
+        dataset["train"] = dataset["train"].select(range(subset_size))
+        logger.debug(f"Using {dataset_percentage}% of dataset: {subset_size}/{original_size} examples")
+    else:
+        logger.debug(f"Using full dataset: {original_size} examples")
+
+    logger.debug("Tokenizing dataset...")
+
+    tokenized_dataset = dataset["train"].map(lambda x: preprocess_alpaca_example(x, tokenizer, max_length),
+                                             batched=False,
+                                             remove_columns=dataset["train"].column_names,
+                                             desc="Tokenizing")
+
+    return tokenized_dataset
+
+
+def initialize_wandb(args: argparse.Namespace, exp_name: str, logger: logging.Logger) -> None:
+    if args.use_wandb and dist.get_rank() == 0:
+        try:
+            wandb_run_name = args.wandb_run_name if args.wandb_run_name else exp_name
+            logger.debug(f"Initializing WandB run: {wandb_run_name}")
+            wandb.init(project=args.wandb_project, name=wandb_run_name, tags=args.wandb_tags, config=vars(args))
+            logger.debug("WandB initialized successfully")
+        except Exception as e:
+            logger.warning(f"Failed to initialize WandB: {e}")
+            args.use_wandb = False
+
+
+def report_cpu_simd(logger: logging.Logger) -> None:
+    """Log the CPU SIMD level used by the Reflow CPU-Adam kernel."""
+    flags = set()
+    try:
+        with open("/proc/cpuinfo") as cpuinfo:
+            for line in cpuinfo:
+                if line.startswith("flags") or line.startswith("Features"):
+                    flags = set(line.split(":", 1)[1].split())
+                    break
+    except OSError:
+        logger.debug("Could not read /proc/cpuinfo; skipping CPU SIMD report")
+        return
+
+    has_avx512 = "avx512f" in flags
+    has_avx2 = "avx2" in flags
+    level = "AVX-512" if has_avx512 else ("AVX-256" if has_avx2 else "scalar (no AVX)")
+    logger.info(f"Reflow CPU-Adam SIMD: {level}")
+    if not has_avx512 and not has_avx2:
+        logger.warning("No AVX: BF16 CPU-Adam uses the scalar fallback (much slower).")
+
+
+def main(args: argparse.Namespace) -> None:
+    logger = setup_logger(rank=0, log_level=args.log_level)
+
+    exp_name = create_experiment_name(args)
+
+    logger.debug(f"Starting experiment: {exp_name}")
+    logger.debug("Training configuration:")
+    logger.debug(f"  Model: {args.model_name}")
+    logger.debug(f"  Batch size: {args.batch_size}")
+    logger.debug(f"  Max length: {args.max_length}")
+    logger.debug(f"  Learning rate: {args.lr}")
+    logger.debug(f"  Epochs: {args.num_train_epochs}")
+    logger.debug(f"  Activation checkpointing: {args.activation_checkpointing}")
+
+    tokenizer = load_tokenizer(args.model_name, logger)
+    model = load_model(args.model_name, args.attn_implementation, logger)
+    if args.leaf_module:
+        from deepspeed.utils import set_z3_leaf_modules
+        logger.debug(f"Setting leaf_module to: {args.leaf_module}")
+        set_z3_leaf_modules(model, [args.leaf_module])
+    setup_model_training(model, args.activation_checkpointing, logger)
+    optimizer = create_optimizer(model, args)
+
+    tokenized_dataset = load_and_preprocess_dataset(args.dataset_name, args.dataset_percentage, tokenizer,
+                                                    args.max_length, logger)
+
+    # Initialize DeepSpeed. We build the data loader ourselves (below) instead of passing
+    # training_data, so the shuffle order is reproducible across runs.
+    model_engine, optimizer, _, _ = deepspeed.initialize(args=args, model=model, optimizer=optimizer)
+
+    # Deterministic data order: a DistributedSampler with a fixed seed shards the dataset across
+    # ranks and shuffles reproducibly, so two runs (e.g. reflow vs zerooffload) see the identical
+    # batch at every step. DeepSpeed's built-in training_data loader does not pin the shuffle seed.
+    train_sampler = DistributedSampler(tokenized_dataset,
+                                       num_replicas=dist.get_world_size(),
+                                       rank=dist.get_rank(),
+                                       shuffle=True,
+                                       seed=args.seed)
+    train_dataloader = DataLoader(tokenized_dataset,
+                                  batch_size=model_engine.train_micro_batch_size_per_gpu(),
+                                  sampler=train_sampler,
+                                  collate_fn=default_data_collator)
+
+    logger = setup_logger(rank=dist.get_rank(), log_level=args.log_level)
+    report_cpu_simd(logger)
+
+    initialize_wandb(args, exp_name, logger)
+
+    model_engine.train()
+
+    sequence_length = args.max_length
+    model_size = sum(get_parameter_count(p) for p in model.parameters())
+    is_moe_model = detect_moe_model(model)
+
+    logger.debug(f"Model type: {'MoE' if is_moe_model else 'Dense'}")
+    logger.debug(f"Model size: {model_size:,} parameters")
+
+    # Calculate TFLOPS only for non-MoE models. Two figures are tracked:
+    #   tflops            -> counts the activation-recompute forward (hardware FLOPs actually executed)
+    #   effective tflops  -> excludes recompute (useful model FLOPs only: 1 forward + 2 backward)
+    total_tflops = None
+    total_tflops_effective = None
+    if not is_moe_model:
+        total_tflops = estimate_transformer_tflops(sequence_length, model_size, model.config.num_hidden_layers,
+                                                   model.config.hidden_size, args.activation_checkpointing)
+        total_tflops_effective = estimate_transformer_tflops(sequence_length,
+                                                             model_size,
+                                                             model.config.num_hidden_layers,
+                                                             model.config.hidden_size,
+                                                             use_activation_checkpointing=False)
+
+    global_step = 0
+    iter_times = []
+    losses = []
+
+    stop = False
+    for epoch in range(args.num_train_epochs):
+        train_sampler.set_epoch(epoch)
+        logger.debug(f"Starting epoch {epoch + 1}/{args.num_train_epochs}")
+
+        for batch in train_dataloader:
+            step_start_time = time.perf_counter()
+            batch = {k: v.to(model_engine.device) for k, v in batch.items()}
+
+            actual_batch_size = batch['input_ids'].shape[0]
+            global_batch_size = actual_batch_size * dist.get_world_size()
+            tokens_in_batch = global_batch_size * sequence_length
+
+            outputs = model_engine(**batch)
+            loss = outputs.loss
+
+            model_engine.backward(loss)
+
+            model_engine.step()
+
+            # Include GPU completion rather than timing only asynchronous launches.
+            _raw_loss = loss.item()
+            step_time = time.perf_counter() - step_start_time
+            global_step += 1
+
+            if global_step > args.warmup_steps:
+                iter_times.append(step_time)
+
+            losses.append(_raw_loss)
+            if args.loss_check:
+                # Full-precision (bit-exact hex) per-step loss so a reflow run and a zerooffload run
+                # under loss-check mode can be diffed bit-for-bit.
+                logger.info(f"BITLOSS step {global_step} hex={float(_raw_loss).hex()} dec={_raw_loss:.12e}")
+
+            tokens_per_second = tokens_in_batch / step_time
+            step_tflops = None
+            effective_tflops = None
+
+            if not is_moe_model and total_tflops is not None:
+                step_tflops = global_batch_size * total_tflops / step_time
+                effective_tflops = global_batch_size * total_tflops_effective / step_time
+
+            if global_step % args.log_interval == 0:
+                avg_loss = sum(losses[-args.log_interval:]) / len(losses[-args.log_interval:])
+
+                if is_moe_model:
+                    # Skip throughput metrics for MoE models
+                    log_msg = (f"Step {global_step:4d} | "
+                               f"Loss: {avg_loss:.4f} | "
+                               f"Time: {step_time * MS_PER_SECOND:5.0f}ms")
+                else:
+                    log_msg = (f"Step {global_step:4d} | "
+                               f"Loss: {avg_loss:.4f} | "
+                               f"Time: {step_time * MS_PER_SECOND:5.0f}ms | "
+                               f"TFLOPS(w/ recompute): {step_tflops:8.2f} | "
+                               f"effective TFLOPS(w/o recompute): {effective_tflops:8.2f} | "
+                               f"Tokens/s: {tokens_per_second:6.0f}")
+
+                logger.info(log_msg)
+
+                if args.use_wandb and dist.get_rank() == 0:
+                    log_dict = {
+                        "train/loss": avg_loss,
+                        "train/epoch": epoch + 1,
+                        "train/global_step": global_step,
+                        "train/learning_rate": args.lr,
+                        "perf/step_time_ms": step_time * MS_PER_SECOND,
+                        "perf/tokens_per_second": tokens_per_second,
+                    }
+
+                    if not is_moe_model and step_tflops is not None:
+                        log_dict["perf/tflops"] = step_tflops  # with recompute (hardware)
+                        log_dict["perf/effective_tflops"] = effective_tflops  # without recompute (useful)
+
+                    wandb.log(log_dict, step=global_step)
+
+            stop = global_step >= args.bench_steps
+            if stop:
+                break
+
+        if stop:
+            break
+
+    if iter_times:
+        mean_step_ms = sum(iter_times) / len(iter_times) * MS_PER_SECOND
+        logger.info(f"Mean step time after {args.warmup_steps} warmup steps: {mean_step_ms:.2f}ms")
+
+    if args.save_checkpoint:
+        # ZeRO checkpoint saving includes collectives, so every rank must participate.
+        model_engine.save_checkpoint(args.output_dir)
+        if dist.get_rank() == 0:
+            tokenizer.save_pretrained(args.output_dir)
+
+    if args.use_wandb and dist.get_rank() == 0:
+        try:
+            wandb.finish()
+            logger.debug("WandB run finished successfully")
+        except Exception as e:
+            logger.error(f"Error finishing WandB run: {e}")
+
+    logger.debug("Training completed successfully!")
+
+
+def create_argument_parser() -> argparse.ArgumentParser:
+    """Create and configure argument parser."""
+    parser = argparse.ArgumentParser(description="Fine-tune language models with DeepSpeed ZeRO Stage 3",
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    parser.add_argument("--model_name", type=str, required=True, help="HuggingFace model name or path")
+    parser.add_argument("--lr", type=float, required=True, help="Learning rate for training")
+    parser.add_argument("--batch_size", type=int, required=True, help="Training batch size per device")
+    parser.add_argument("--output_dir", type=str, required=True, help="Directory to save model checkpoints")
+
+    parser.add_argument("--attn_implementation",
+                        type=str,
+                        default="flash_attention_2",
+                        choices=["eager", "sdpa", "flash_attention_2"],
+                        help="Attention implementation to use")
+    parser.add_argument("--leaf_module",
+                        type=str,
+                        default=None,
+                        help="Set leaf_module to enable fine-tuning MoE models")
+    parser.add_argument("--activation_checkpointing",
+                        action="store_true",
+                        help="Enable activation checkpointing to save memory")
+    parser.add_argument("--optimizer",
+                        type=str,
+                        default="adam",
+                        choices=["adam", "lion"],
+                        help="CPU optimizer to offload (Reflow remaps it to its async subclass)")
+
+    parser.add_argument("--num_train_epochs", type=int, default=1, help="Number of training epochs")
+    parser.add_argument("--max_length", type=int, default=2048, help="Maximum sequence length for tokenization")
+    parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay for optimization")
+
+    parser.add_argument("--local_rank", type=int, default=-1, help="Local rank passed from distributed launcher")
+
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
+    parser.add_argument("--deterministic",
+                        action="store_true",
+                        help="Enable deterministic training for full reproducibility")
+    parser.add_argument("--loss_check",
+                        action="store_true",
+                        help="Bit-level loss-check mode (verification only): full determinism including "
+                        "the NCCL collectives (fixed algorithm/protocol) and cuBLAS, plus per-step "
+                        "bit-exact loss logging, so two configs (e.g. reflow vs zerooffload) can be "
+                        "diffed bit-for-bit. Not for throughput runs.")
+
+    parser.add_argument("--log_interval", type=int, default=1, help="Log performance metrics every N steps")
+    parser.add_argument("--log_level",
+                        type=str,
+                        default="INFO",
+                        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+                        help="Logging level for controlling output verbosity")
+    parser.add_argument("--warmup_steps",
+                        type=int,
+                        default=15,
+                        help="Number of warmup steps for performance measurements")
+    parser.add_argument("--bench_steps", type=int, default=100, help="Number of benchmark steps to run")
+
+    parser.add_argument("--save_checkpoint", action="store_true", help="Save model checkpoint after training")
+
+    parser.add_argument("--use_wandb", action="store_true", help="Enable Weights & Biases logging")
+    parser.add_argument("--wandb_project", type=str, default="reflow", help="WandB project name")
+    parser.add_argument("--wandb_run_name",
+                        type=str,
+                        default=None,
+                        help="WandB run name (auto-generated if not provided)")
+    parser.add_argument("--wandb_tags", type=str, nargs="+", default=[], help="WandB tags for the run")
+
+    parser.add_argument("--dataset_name", type=str, default="tatsu-lab/alpaca", help="HuggingFace dataset name")
+    parser.add_argument("--dataset_percentage",
+                        type=float,
+                        default=100.0,
+                        help="Percentage of dataset to use (1.0-100.0)")
+
+    return parser
+
+
+def validate_arguments(args: argparse.Namespace) -> None:
+    if args.dataset_percentage <= 0 or args.dataset_percentage > 100:
+        raise ValueError("dataset_percentage must be between 0 and 100")
+
+    if args.max_length <= 0:
+        raise ValueError("max_length must be positive")
+
+    if args.batch_size <= 0:
+        raise ValueError("batch_size must be positive")
+
+    if args.lr <= 0:
+        raise ValueError("learning rate must be positive")
+
+    if args.bench_steps <= 0 or not 0 <= args.warmup_steps < args.bench_steps:
+        raise ValueError("bench_steps must be positive and warmup_steps must be below bench_steps")
+
+    if args.log_interval <= 0:
+        raise ValueError("log_interval must be positive")
+
+    if args.num_train_epochs <= 0:
+        raise ValueError("num_train_epochs must be positive")
+
+
+if __name__ == "__main__":
+    parser = create_argument_parser()
+    parser = deepspeed.add_config_arguments(parser)
+    args = parser.parse_args()
+
+    validate_arguments(args)
+
+    if args.loss_check:
+        # Bit-level loss-check mode: make the ENTIRE pipeline deterministic -- including the NCCL
+        # collectives (fixed algorithm/protocol) and cuBLAS -- so two configs (e.g. reflow vs
+        # zerooffload) can be diffed bit-for-bit. Throughput is irrelevant here. These env vars must
+        # be set before any CUDA/NCCL init (i.e. before deepspeed.initialize), which is why this runs
+        # in __main__ ahead of main().
+        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+        os.environ.setdefault("NCCL_ALGO", "Ring")
+        os.environ.setdefault("NCCL_PROTO", "Simple")
+        enable_full_determinism(args.seed)
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+        logging.basicConfig(level=getattr(logging, args.log_level.upper()))
+        logging.info("Loss-check mode: full determinism (NCCL Ring/Simple + cuBLAS + deterministic "
+                     "algorithms) and bit-exact loss logging enabled")
+    elif args.deterministic:
+        enable_full_determinism(args.seed)
+        torch.backends.cudnn.benchmark = False
+        logging.basicConfig(level=getattr(logging, args.log_level.upper()))
+        logging.info("Enabled deterministic mode for full reproducibility")
+    else:
+        set_seed(args.seed)
+        logging.basicConfig(level=getattr(logging, args.log_level.upper()))
+        logging.info(f"Set random seed to {args.seed}")
+
+    main(args)
diff --git a/training/DeepSpeed-Reflow/requirements.txt b/training/DeepSpeed-Reflow/requirements.txt
new file mode 100644
index 000000000..da0ed3dde
--- /dev/null
+++ b/training/DeepSpeed-Reflow/requirements.txt
@@ -0,0 +1,8 @@
+torch>=2.5.1
+deepspeed>=0.17.0
+datasets>=4.0.0
+transformers>=4.56.1
+numpy>=1.21.0
+wandb
+packaging
+psutil
diff --git a/training/DeepSpeed-Reflow/run_compare.sh b/training/DeepSpeed-Reflow/run_compare.sh
new file mode 100644
index 000000000..aaa563084
--- /dev/null
+++ b/training/DeepSpeed-Reflow/run_compare.sh
@@ -0,0 +1,232 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+# run_compare.sh -- one-command Reflow vs ZeRO-Infinity comparison.
+#
+# Usage:
+#   bash run_compare.sh                          # Llama-7B and OPT-30B, 1 GPU
+#   bash run_compare.sh finetune_opt-30b_1gpu.sh # a specific config
+#   bash run_compare.sh script_a.sh script_b.sh  # several configs
+#
+# For each finetune script it runs both modes back to back, samples GPU
+# and host memory while they run, and prints:
+#
+#   * peak GPU memory and peak host RAM for each mode
+#   * steady-state step time and the resulting speedup
+#   * TFLOPS, bwd_microstep, step_microstep
+#
+# Timings are taken from ordinary runs -- never pass --loss_check here,
+# since deterministic collectives distort them (bit-exactness is checked
+# separately by check_bitexact.sh).
+#
+# The DeepSpeed config must have "wall_clock_breakdown": true for the
+# bwd_microstep / step_microstep timers to appear.
+#
+# Before running a configuration, free host RAM and free GPU memory are
+# checked against what the ZeRO-Infinity baseline needs (the heavier of
+# the two modes). A configuration that does not fit is skipped rather
+# than left to die part-way through. Set FORCE=1 to run it anyway.
+
+set -uo pipefail
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+sampler=""
+cleanup() {
+    if [ -n "$sampler" ]; then
+        kill "$sampler" 2>/dev/null || true
+        wait "$sampler" 2>/dev/null || true
+        sampler=""
+    fi
+}
+trap cleanup EXIT
+trap 'exit 130' INT
+trap 'exit 143' TERM
+
+SCRIPTS=("$@")
+if [ ${#SCRIPTS[@]} -eq 0 ]; then
+    SCRIPTS=(finetune_llama-7b_1gpu.sh finetune_opt-30b_1gpu.sh)
+fi
+
+OUT_DIR=${OUT_DIR:-compare_out}
+SAMPLE_SEC=${SAMPLE_SEC:-2}
+mkdir -p "$OUT_DIR"
+
+# ------------------------------------------------------------ prechecks
+# Requirements of the ZeRO-Infinity baseline (the heavier of the two
+# modes), in MiB: our measured peaks plus 1%.
+#
+#   Llama-7B 1 GPU : measured 132151 host / 31118 GPU
+#   OPT-30B  1 GPU : measured 537467 host / 87980 GPU
+#
+# Configurations we have not measured return "unknown" and are not
+# pre-checked; they simply run, and a failure is reported afterwards.
+req_host_mib() {
+    case "$1" in
+        *llama-7b*_1gpu*) echo 133473 ;;   # ~130 GB
+        *opt-30b*_1gpu*)  echo 542842 ;;   # ~530 GB
+        *opt-350m*)       echo   8192 ;;
+        *)                echo unknown ;;
+    esac
+}
+req_gpu_mib() {
+    case "$1" in
+        *llama-7b*_1gpu*) echo 31429 ;;    # ~31 GB
+        *opt-30b*_1gpu*)  echo 88860 ;;    # ~87 GB
+        *opt-350m*)       echo  8192 ;;
+        *)                echo unknown ;;
+    esac
+}
+
+free_host_mib() {
+    awk '/^MemAvailable:/ {printf "%d", $2/1024}' /proc/meminfo
+}
+free_gpu_mib() {
+    nvidia-smi --query-gpu=memory.total,memory.used \
+        --format=csv,noheader,nounits 2>/dev/null \
+        | awk -F', *' '{d=$1-$2; if (d>m) m=d} END {printf "%d", m}'
+}
+
+# Returns 0 if the configuration fits or is unmeasured, 1 if it clearly
+# does not fit (with a reason).
+fits() {  # $1 = script name
+    local rh rg fh fg
+    rh=$(req_host_mib "$1"); rg=$(req_gpu_mib "$1")
+    if [ "$rh" != unknown ]; then
+        fh=$(free_host_mib)
+        if [ "${fh:-0}" -lt "$rh" ]; then
+            echo "  SKIPPED: needs ~$((rh/1024)) GB of free host RAM, ${fh:-0} MiB available."
+            return 1
+        fi
+    fi
+    if [ "$rg" != unknown ]; then
+        fg=$(free_gpu_mib)
+        if [ "${fg:-0}" -lt "$rg" ]; then
+            echo "  SKIPPED: needs ~$((rg/1024)) GB of free GPU memory, ${fg:-0} MiB available."
+            return 1
+        fi
+    fi
+    return 0
+}
+
+# -------------------------------------------------------------- sampling
+# Peak GPU memory comes from nvidia-smi; peak host RAM is the summed RSS
+# of the training processes, which includes the pinned FP32 master and
+# optimizer state. (On multi-rank runs shared pages are counted once per
+# rank, so this is an upper bound; the 1-GPU configs have a single rank.)
+sample_loop() {
+    while :; do
+        gpu=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits \
+              2>/dev/null | sort -rn | head -1)
+        cpu=$(ps -eo rss=,args= | grep -F finetune_zero3.py | grep -v grep \
+              | awk '{s+=$1} END {print int(s/1024)}')
+        echo "${gpu:-0} ${cpu:-0}"
+        sleep "$SAMPLE_SEC"
+    done
+}
+
+peak_col() {  # $1 = sample file, $2 = column
+    awk -v c="$2" '{if ($c+0 > m) m = $c+0} END {printf "%.0f", m}' "$1"
+}
+
+# --------------------------------------------------------------- parsing
+# Steady state excludes step 1, which absorbs warmup and JIT compilation.
+mean_step_time() {
+    local summary
+    summary=$(sed -n 's/.*Mean step time after [0-9]* warmup steps: \([0-9.]*\)ms/\1/p' "$1" | tail -1)
+    if [ -n "$summary" ]; then
+        echo "$summary"
+        return
+    fi
+    # Older logs have no summary. Preserve numeric step order for steps 10 and above.
+    sed -n 's/.* - finetune_zero3 - INFO - Step *\([0-9]*\) | Loss: [^|]* | Time: *\([0-9]*\)ms.*/\1 \2/p' "$1" \
+        | sort -n -k1,1 | tail -n +2 \
+        | awk '{s+=$2; n++} END {if (n) printf "%.0f", s/n; else printf "0"}'
+}
+mean_tflops() {
+    grep ' - finetune_zero3 - INFO - Step' "$1" \
+        | grep -o 'TFLOPS(w/ recompute): *[0-9.]*' | tail -n +2 \
+        | awk -F': *' '{s+=$2; n++} END {if (n) printf "%.1f", s/n; else printf "n/a"}'
+}
+mean_timer() {  # $2 = timer name
+    grep -o "$2: [0-9.]*" "$1" | tail -n +2 \
+        | awk -F': ' '{s+=$2; n++} END {if (n) printf "%.1f", s/n; else printf "n/a"}'
+}
+
+# ------------------------------------------------------------------ main
+failed_any=0
+for script in "${SCRIPTS[@]}"; do
+    if [[ "$script" != /* && ! -f "$script" ]]; then
+        script="$SCRIPT_DIR/$script"
+    fi
+    if [ ! -f "$script" ]; then
+        echo "FAILED: $script not found"
+        failed_any=1
+        continue
+    fi
+    tag=$(basename "$script" .sh)
+    echo "############################################################"
+    echo "# $tag"
+    echo "############################################################"
+
+    if [ "${FORCE:-0}" != "1" ] && ! fits "$script"; then
+        echo "  Set FORCE=1 to run it anyway."
+        echo
+        continue
+    fi
+
+    failed_logs=()
+    for mode in reflow zerooffload; do
+        echo "  running $mode ..."
+        sample_loop > "$OUT_DIR/$tag.$mode.mem" &
+        sampler=$!
+        if ! bash "$script" "$mode" > "$OUT_DIR/$tag.$mode.log" 2>&1; then
+            failed_logs+=("$OUT_DIR/$tag.$mode.log")
+        fi
+        cleanup
+    done
+
+    rl="$OUT_DIR/$tag.reflow.log";      rm_="$OUT_DIR/$tag.reflow.mem"
+    zl="$OUT_DIR/$tag.zerooffload.log"; zm="$OUT_DIR/$tag.zerooffload.mem"
+
+    # A run that died -- out of host RAM or GPU memory, most often -- has
+    # no step lines. Say so plainly instead of printing a table of n/a,
+    # and move on to the next configuration; the results already printed
+    # above remain valid.
+    for f in "$rl" "$zl"; do
+        grep -q ' - finetune_zero3 - INFO - Step' "$f" || failed_logs+=("$f")
+    done
+    if [ ${#failed_logs[@]} -gt 0 ]; then
+        echo
+        failed_any=1
+        echo "  FAILED: training exited unsuccessfully or produced no step logs. Last lines:"
+        for f in "${failed_logs[@]}"; do
+            echo "  --- $f"
+            tail -5 "$f" | sed 's/^/    /'
+        done
+        echo
+        continue
+    fi
+
+    r_step=$(mean_step_time "$rl"); z_step=$(mean_step_time "$zl")
+    speedup=$(awk -v r="$r_step" -v z="$z_step" \
+        'BEGIN {if (r > 0) printf "%.2f", z/r; else printf "n/a"}')
+
+    echo
+    printf '%-28s %14s %14s\n' "metric" "Reflow" "ZeRO-Inf."
+    printf '%-28s %14s %14s\n' "peak GPU memory (MiB)" \
+        "$(peak_col "$rm_" 1)" "$(peak_col "$zm" 1)"
+    printf '%-28s %14s %14s\n' "peak host RAM (MiB)" \
+        "$(peak_col "$rm_" 2)" "$(peak_col "$zm" 2)"
+    printf '%-28s %14s %14s\n' "step time (ms)"      "$r_step"  "$z_step"
+    printf '%-28s %14s %14s\n' "TFLOPS"              \
+        "$(mean_tflops "$rl")" "$(mean_tflops "$zl")"
+    printf '%-28s %14s %14s\n' "bwd_microstep (ms)"  \
+        "$(mean_timer "$rl" bwd_microstep)"  "$(mean_timer "$zl" bwd_microstep)"
+    printf '%-28s %14s %14s\n' "step_microstep (ms)" \
+        "$(mean_timer "$rl" step_microstep)" "$(mean_timer "$zl" step_microstep)"
+    echo
+    echo "  SPEEDUP: ${speedup}x"
+    echo "  Logs: $OUT_DIR/$tag.*.log"
+    echo
+done
+
+exit "$failed_any"
diff --git a/training/DeepSpeed-Reflow/tests/test_examples.py b/training/DeepSpeed-Reflow/tests/test_examples.py
new file mode 100644
index 000000000..ab33e8147
--- /dev/null
+++ b/training/DeepSpeed-Reflow/tests/test_examples.py
@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+
+import importlib.util
+import json
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import tempfile
+import unittest
+from unittest.mock import patch
+
+ROOT = Path(__file__).resolve().parents[1]
+
+# The stand-in implements the launcher's CLI contract, including process exit status.
+# It never launches training or queries a GPU; real training needs separate device validation.
+LAUNCHER = '''#!/usr/bin/env python3
+import json
+import os
+from pathlib import Path
+import sys
+
+args = sys.argv[1:]
+config_arg = next(arg for arg in args if arg.startswith('--deepspeed_config='))
+config = json.loads(Path(config_arg.split('=', 1)[1]).read_text())
+with open(os.environ['CAPTURE'], 'a') as capture:
+    capture.write(json.dumps({'args': args, 'config': config}) + '\\n')
+mode = 'reflow' if 'reflow' in config['zero_optimization'] else 'zerooffload'
+steps = int(args[args.index('--bench_steps') + 1])
+scenario = os.environ.get('SCENARIO', '')
+if scenario == 'truncated':
+    steps -= 1
+for step in range(1, steps + 1):
+    value = 1.0
+    if scenario == 'mismatch' and mode == 'zerooffload':
+        value = 2.0
+    print(f'2026-10-02 - finetune_zero3 - INFO - BITLOSS step {step} hex={value.hex()} dec={value}')
+    label = str(step) if scenario == 'legacy' else f'{step:4d}'
+    print(f'2026-10-02 - finetune_zero3 - INFO - Step {label} | Loss: 1.0 | Time: {step * 10}ms | TFLOPS(w/ recompute): 1.0')
+if scenario != 'legacy':
+    print('2026-10-02 - finetune_zero3 - INFO - Mean step time after 1 warmup steps: 123.45ms')
+if scenario == 'failure':
+    sys.exit(7)
+'''
+
+
+class LauncherTests(unittest.TestCase):
+
+    def setUp(self):
+        self.temp = tempfile.TemporaryDirectory(prefix='reflow example ')
+        self.addCleanup(self.temp.cleanup)
+        self.directory = Path(self.temp.name)
+        self.example = self.directory / 'example scripts'
+        self.example.mkdir()
+        for source in ROOT.glob('*.sh'):
+            shutil.copy(source, self.example)
+        shutil.copy(ROOT / 'finetune_zero3.py', self.example)
+        self.bin_dir = self.directory / 'bin'
+        self.bin_dir.mkdir()
+        launcher = self.bin_dir / 'deepspeed'
+        launcher.write_text(LAUNCHER)
+        launcher.chmod(0o755)
+        # Public nvidia-smi CSV contract used by the comparison sampler.
+        nvidia = self.bin_dir / 'nvidia-smi'
+        nvidia.write_text('#!/bin/sh\necho 100\n')
+        nvidia.chmod(0o755)
+        self.capture = self.directory / 'capture.jsonl'
+        self.env = dict(os.environ,
+                        PATH=str(self.bin_dir) + os.pathsep + os.environ['PATH'],
+                        CAPTURE=str(self.capture),
+                        OUT_DIR=str(self.directory / 'results'),
+                        SAMPLE_SEC='0.01',
+                        FORCE='1',
+                        ATTN_IMPLEMENTATION='eager')
+
+    def run_script(self, script, *args, **env):
+        return subprocess.run(['bash', str(self.example / script), *args],
+                              cwd=self.directory,
+                              env=dict(self.env, **env),
+                              text=True,
+                              stdout=subprocess.PIPE,
+                              stderr=subprocess.STDOUT,
+                              timeout=30)
+
+    def launches(self):
+        return [json.loads(line) for line in self.capture.read_text().splitlines()]
+
+    def test_all_launchers_generate_matching_cpu_bf16_configs(self):
+        # Catch per-rank/global batch confusion, invalid runtime options and mode drift.
+        from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
+        for script in sorted(self.example.glob('finetune_*.sh')):
+            with self.subTest(script=script.name):
+                pair = []
+                for mode in ('reflow', 'zerooffload'):
+                    result = self.run_script(script.name, mode, '16', BENCH_STEPS='100', WARMUP_STEPS='10')
+                    self.assertEqual(result.returncode, 0, result.stdout)
+                    launch = self.launches()[-1]
+                    args, config = launch['args'], launch['config']
+                    DeepSpeedZeroConfig(**config['zero_optimization'])
+                    gpu_arg = next(arg for arg in args if arg.startswith('--num_gpus='))
+                    world_size = int(gpu_arg.split('=')[1])
+                    microbatch = int(args[args.index('--batch_size') + 1])
+                    self.assertEqual(microbatch * world_size, config['train_batch_size'])
+                    self.assertEqual(config['train_batch_size'], 16)
+                    self.assertTrue(config['bf16']['enabled'])
+                    self.assertEqual(config['gradient_clipping'], 0)
+                    self.assertEqual(config['zero_optimization']['offload_optimizer']['device'], 'cpu')
+                    self.assertEqual(args[args.index('--bench_steps') + 1], '100')
+                    self.assertEqual(args[args.index('--warmup_steps') + 1], '10')
+                    self.assertEqual(args[args.index('--attn_implementation') + 1], 'eager')
+                    self.assertIn(str(self.example / 'finetune_zero3.py'), args)
+                    pair.append(config)
+                reflow = pair[0]['zero_optimization'].pop('reflow')
+                self.assertEqual(reflow['bucketwise_cores_per_worker'], 8)
+                self.assertEqual(pair[0], pair[1])
+
+    def test_invalid_mode_and_global_batch_do_not_launch(self):
+        for args in ((), ('unknown', ), ('reflow', '0'), ('reflow', '7'), ('reflow', 'abc')):
+            with self.subTest(args=args):
+                result = self.run_script('finetune_opt-30b_8gpu.sh', *args)
+                self.assertNotEqual(result.returncode, 0, result.stdout)
+        self.assertFalse(self.capture.exists())
+
+    def test_core_overrides_reach_generated_configuration(self):
+        result = self.run_script('finetune_opt-350m-1gpu.sh', 'reflow', WORKER_CORES='4', MAIN_CORES='3')
+        self.assertEqual(result.returncode, 0, result.stdout)
+        reflow = self.launches()[-1]['config']['zero_optimization']['reflow']
+        self.assertEqual(reflow['bucketwise_cores_per_worker'], 4)
+        self.assertEqual(reflow['main_thread_cores'], 3)
+
+    def test_bitexact_requires_complete_successful_matching_runs(self):
+        for scenario, success in (('', True), ('mismatch', False), ('truncated', False), ('failure', False)):
+            with self.subTest(scenario=scenario):
+                result = self.run_script('check_bitexact.sh', '8', SCENARIO=scenario)
+                self.assertEqual(result.returncode == 0, success, result.stdout)
+                self.assertEqual('BIT-IDENTICAL' in result.stdout, success)
+        config = self.launches()[0]['config']
+        self.assertEqual(config['train_micro_batch_size_per_gpu'], 4)
+
+    def test_compare_propagates_partial_training_failure(self):
+        result = self.run_script('run_compare.sh', 'finetune_opt-350m-1gpu.sh', SCENARIO='failure')
+        self.assertNotEqual(result.returncode, 0, result.stdout)
+        self.assertIn('FAILED', result.stdout)
+        self.assertNotIn('SPEEDUP', result.stdout)
+        self.assertNotIn('cannot open', result.stdout)
+
+    def test_compare_uses_reported_warmup_summary(self):
+        result = self.run_script('run_compare.sh', 'finetune_opt-350m-1gpu.sh')
+        self.assertEqual(result.returncode, 0, result.stdout)
+        self.assertRegex(result.stdout, r'step time \(ms\)\s+123.45\s+123.45')
+
+    def test_compare_preserves_numeric_step_order_in_legacy_logs(self):
+        # Lexicographic sorting drops step 10 as warmup instead of step 1.
+        result = self.run_script('run_compare.sh', 'finetune_opt-350m-1gpu.sh', SCENARIO='legacy')
+        self.assertEqual(result.returncode, 0, result.stdout)
+        self.assertRegex(result.stdout, r'step time \(ms\)\s+60\s+60')
+
+
+class PreprocessingTests(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        import torch
+        if torch.cuda.is_initialized():
+            raise RuntimeError('These checks must run without initializing CUDA')
+        spec = importlib.util.spec_from_file_location('reflow_example', ROOT / 'finetune_zero3.py')
+        cls.example = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(cls.example)
+        from tokenizers import Tokenizer
+        from tokenizers.models import WordLevel
+        from tokenizers.pre_tokenizers import Whitespace
+        from transformers import PreTrainedTokenizerFast
+        tokenizer = Tokenizer(WordLevel({'[UNK]': 0, '[PAD]': 1, 'hello': 2, 'world': 3}, unk_token='[UNK]'))
+        tokenizer.pre_tokenizer = Whitespace()
+        cls.tokenizer = PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token='[UNK]', pad_token='[PAD]')
+
+    def test_padding_and_raw_dataset_fields_do_not_enter_training_loss(self):
+        # Catch unmasked padded targets and source strings leaking into the model batch.
+        import torch
+        from datasets import Dataset, DatasetDict
+        from transformers import OPTConfig, OPTForCausalLM, default_data_collator
+        records = [{'instruction': 'hello', 'input': '', 'output': 'world'}] * 2
+        data = DatasetDict(train=Dataset.from_list(records))
+        with patch.object(self.example, 'load_dataset', return_value=data):
+            processed = self.example.load_and_preprocess_dataset('offline', 100, self.tokenizer, 32,
+                                                                 self.example.setup_logger())
+        batch = default_data_collator([processed[0], processed[1]])
+        self.assertEqual(set(batch), {'input_ids', 'attention_mask', 'labels'})
+        padding = batch['attention_mask'] == 0
+        self.assertTrue(padding.any())
+        self.assertTrue((batch['labels'][padding] == -100).all())
+        self.assertTrue(torch.equal(batch['labels'][~padding], batch['input_ids'][~padding]))
+        model = OPTForCausalLM(
+            OPTConfig(vocab_size=4,
+                      hidden_size=16,
+                      ffn_dim=32,
+                      num_hidden_layers=1,
+                      num_attention_heads=2,
+                      max_position_embeddings=32,
+                      pad_token_id=1))
+        output = model(**batch)
+        expected = torch.nn.functional.cross_entropy(output.logits[:, :-1].reshape(-1, 4),
+                                                     batch['labels'][:, 1:].reshape(-1),
+                                                     ignore_index=-100)
+        torch.testing.assert_close(output.loss, expected)
+        output.loss.backward()
+        self.assertTrue(torch.isfinite(model.model.decoder.embed_tokens.weight.grad).all())
+        self.assertFalse(torch.cuda.is_initialized())
+
+    def test_empty_dataset_selection_is_rejected(self):
+        from datasets import Dataset, DatasetDict
+        data = DatasetDict(train=Dataset.from_list([{'instruction': 'hello', 'input': '', 'output': 'world'}]))
+        with patch.object(self.example, 'load_dataset', return_value=data):
+            with self.assertRaises(ValueError):
+                self.example.load_and_preprocess_dataset('offline', 1, self.tokenizer, 32, self.example.setup_logger())
+
+    def test_invalid_benchmark_window_is_rejected(self):
+        for extra in (['--bench_steps', '0'], ['--bench_steps', '10', '--warmup_steps', '20'], ['--log_interval',
+                                                                                                '0']):
+            with self.subTest(extra=extra):
+                args = self.example.create_argument_parser().parse_args(
+                    ['--model_name', 'offline', '--lr', '1e-5', '--batch_size', '1', '--output_dir', 'out'] + extra)
+                with self.assertRaises(ValueError):
+                    self.example.validate_arguments(args)
+
+
+if __name__ == '__main__':
+    unittest.main()

From 66d293f707047a7a557b8aeb53e857348276474f Mon Sep 17 00:00:00 2001
From: st-bang97 <st.bang@dgist.ac.kr>
Date: Fri, 2 Oct 2026 07:20:27 +0000
Subject: [PATCH 2/2] Correct Reflow precision documentation

Signed-off-by: st-bang97 <st.bang@dgist.ac.kr>
---
 training/DeepSpeed-Reflow/README.md | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/training/DeepSpeed-Reflow/README.md b/training/DeepSpeed-Reflow/README.md
index e6ef92642..1ff5ef976 100644
--- a/training/DeepSpeed-Reflow/README.md
+++ b/training/DeepSpeed-Reflow/README.md
@@ -1,8 +1,8 @@
 # Reflow Fine-Tuning Examples
 
-Fine-tune large language models with [DeepSpeed](https://www.deepspeed.ai/) ZeRO Stage 3 + **Reflow**, an asynchronous CPU-offload optimizer for mixed-precision BF16 training. Reflow keeps the optimizer state and FP32 master weights on the CPU like ZeRO-Offload, but overlaps the CPU optimizer work with the backward pass instead of running it serially afterward. It uses the **same GPU memory** as ZeRO-Offload and **~12% less host RAM** (OPT-30B: ~549 vs ~625 GB) — gradients are offloaded in half precision (BF16) and promoted to FP32 inside the CPU kernel, so there is no CPU-side FP32 gradient buffer (see [Memory](#memory) for measured numbers).
+Fine-tune large language models with [DeepSpeed](https://www.deepspeed.ai/) ZeRO Stage 3 + **Reflow**, an asynchronous CPU-offload optimizer for mixed-precision FP16/BF16 training. Reflow keeps the optimizer state and FP32 master weights on the CPU like ZeRO-Offload, but overlaps the CPU optimizer work with the backward pass instead of running it serially afterward. It uses the **same GPU memory** as ZeRO-Offload and **~12% less host RAM** (OPT-30B: ~549 vs ~625 GB) — gradients are offloaded in half precision (FP16/BF16) and promoted to FP32 inside the CPU kernel, so there is no CPU-side FP32 gradient buffer (see [Memory](#memory) for measured numbers).
 
-Reflow supports **Adam/AdamW and Lion** with BF16 model parameters and gradients, FP32 master weights, and FP32 optimizer states. The same scripts run the Reflow path or the plain ZeRO-Offload baseline. `check_bitexact.sh` compares deterministic per-step losses for a selected model and configuration.
+Reflow supports **Adam/AdamW and Lion** with FP16/BF16 model parameters and gradients, FP32 master weights, and FP32 optimizer states. The same scripts run the Reflow path or the plain ZeRO-Offload baseline. `check_bitexact.sh` compares deterministic per-step losses for a selected model and configuration.
 
 ## Quick Start
 
@@ -110,10 +110,10 @@ An exact loss match validates the recorded losses for that run. It does not esta
 
 ## Requirements
 
-* A DeepSpeed build with Reflow, ZeRO Stage 3, and CPU optimizer offload (`offload_optimizer.device = "cpu"`). These examples do not exercise NVMe offload.
+* A DeepSpeed build with Reflow, ZeRO Stage 3, and CPU optimizer offload (`offload_optimizer.device = "cpu"`).
 * An x86 CPU with AVX (BF16 runs on both AVX-512 and AVX-256/AVX2). **AVX-512 is recommended** — AVX-256/AVX2 has no native BF16 support, so the whole BF16 path (gradient accumulation, FP32↔BF16 conversion, and the Adam updates) is emulated with extra instructions and can be slower. Enough cores per rank for the main thread plus the CPU-Adam workers.
 * Host RAM for the CPU-resident FP32 master + optimizer state (e.g. Llama-70B: ~840 GB with Adam, ~560 GB with Lion — Lion keeps one momentum tensor instead of two).
-* BF16 model parameters and gradients with FP32 master weights and optimizer states. FP16/FP32 model parameters, BF16 master weights/states, and `fp32_optimizer_states: false` are rejected by Reflow.
+* Reflow uses FP16/BF16 model parameters and gradients with FP32 master weights and optimizer states. The scripts in this directory select BF16. FP32 model parameters, low-precision master weights/states, and `fp32_optimizer_states: false` are rejected.
 
 ## Memory
 
@@ -193,7 +193,6 @@ Reflow is ~1.9× faster (≈3.1 s vs ≈5.9 s per step) by overlapping the CPU-A
 ## Notes
 
 * **NUMA core binding**: every launch script passes `--bind_cores_to_rank`. Reflow allocates its main reservation and worker masks within the rank's CPU slice. Use physical core reservations and `pin_main_thread` if you need explicit separation of the main thread and worker SMT siblings.
-* These examples use CPU optimizer offload. NVMe offload is not validated here.
 * Use `engine.step()` after `deepspeed.initialize()`; direct calls to `ReflowCPUAdam.step()` or `ReflowCPULion.step()` are rejected.
 * `--save_checkpoint` requires participation by every rank. The script saves the tokenizer on rank 0 and propagates checkpoint failures.
 * **AVX-512 is strongly recommended** for the CPU optimizer; on AVX2 the BF16 path is emulated and slower (see [Requirements](#requirements)).