Split pin_memory benches into per-experiment folders. (#1015)

Keep ZeRO-3 offload as model_tensor_offload (with --zero-stage) and add
activation, H2D/D2H, and optional harnesses so one DSE tree covers the
pinned-offload blog matrix. Root zero3_offload_bench.py still delegates
to the ZeRO-3 path.

Signed-off-by: tunji-ruwase_snow <tunji.ruwase@snowflake.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: sfc-gh-truwase

benchmarks/pin_memory/README.md

@@ -1,49 +1,36 @@
-# ZeRO-3 CPU-Offload Pinned-Memory Benchmark
+# Pinned-memory experiments
 
-This directory contains an end-to-end benchmark for ZeRO-3 CPU offload that
-measures training step time with pinned vs unpinned host memory, plus an
-opt-in ablation of registered vs unregistered pinned memory.
+Harnesses for pin vs pageable host memory on CPU offload. Each subdirectory is
+one experiment. Shared subprocess/JSON helpers live in `common.py`.
 
-## Files in this Directory
+Native backends `mlock` host memory: raise `RLIMIT_MEMLOCK` (`ulimit -l`) or
+run as root for multi-GB models.
 
-- **zero3_offload_bench.py**: Benchmarking script; the model can be a real
-  architecture fetched from the HuggingFace hub (random weights) or a
-  synthetic MLP stack that needs no network access.
+## Layout
 
-## What it Measures
+| Folder | Blog experiment | Default command |
+|--------|-----------------|-----------------|
+| [`model_tensor_offload/`](model_tensor_offload/) | ZeRO CPU param/optimizer offload (stage 3 default; `--zero-stage 1\|2` optional) | `python model_tensor_offload/bench.py --hidden 2048 --layers 12 --batch 4 --seq 128` |
+| [`activation_offload/`](activation_offload/) | Checkpoint hidden-state offload; `use_pin_memory` on/off, **async on** | `python activation_offload/bench.py --hidden 1024 --layers 8 --batch 1 --seq 2048` |
+| [`h2d_d2h/`](h2d_d2h/) | Supporting H2D/D2H GB/s (pageable, torch, native-unregistered, native-registered) | `python h2d_d2h/bench.py` |
+| [`grad_offload/`](grad_offload/) | Optional #8207-style grad offload (wraps model-tensor ZeRO-3) | `python grad_offload/bench.py --hidden 2048 --layers 12` |
+| [`cpu_pin/`](cpu_pin/) | Optional CPU-only native vs Torch pin | `python cpu_pin/bench.py` |
+| [`deepcompile_activation/`](deepcompile_activation/) | Optional `compile.offload_activation_pin_memory` | `python deepcompile_activation/bench.py` |
 
-By default the script runs ZeRO-3 with `offload_optimizer` and `offload_param`
-(both CPU) in two arms and reports the step-time comparison:
+`zero3_offload_bench.py` at this directory root still runs **model-tensor ZeRO-3** (same flags as before).
 
-| Arm | offload `pin_memory` | `DS_PIN_MEMORY_REGISTER_DEVICE` |
-|-----|----------------------|---------------------------------|
-| `unpinned` | `False` | (n/a) |
-| `pinned` | `True` (`DS_PIN_MEMORY_BACKEND=native`) | `1` |
+## Model-tensor arms
 
-Works on any accelerator with native pin + `register_host_memory` support
-(CUDA and XPU are tested). Each arm runs in its own subprocess with a fresh
-rendezvous port so device state never leaks between arms.
+| Arm | `offload_*.pin_memory` | Backend |
+|-----|------------------------|---------|
+| unpinned | `false` | n/a |
+| pinned | `true` | `DS_PIN_MEMORY_BACKEND=native`, `DS_PIN_MEMORY_REGISTER_DEVICE=1` |
 
-Power users can additionally ablate device registration of pinned buffers:
+`--ablate-register` adds `pinned-unregistered`. CUDA-oriented; skip on XPU if `register_host_memory` is missing (`h2d_d2h/bench.py --skip-native-register`).
 
 ```bash
-python zero3_offload_bench.py --ablate-register ...
+python model_tensor_offload/bench.py --model Qwen/Qwen2.5-7B --batch 4 --seq 512
+python model_tensor_offload/bench.py --zero-stage 2 --hidden 2048 --layers 12 --batch 4 --seq 128
 ```
 
-which adds a `pinned-unregistered` arm (`DS_PIN_MEMORY_REGISTER_DEVICE=0`).
-
-## Usage
-
-```bash
-# real model architecture (config fetched from the HF hub, random weights)
-python zero3_offload_bench.py --model Qwen/Qwen2.5-7B --batch 4 --seq 512
-
-# synthetic MLP stack, no network needed
-python zero3_offload_bench.py --hidden 2048 --layers 12 --batch 4 --seq 128
-```
-
-Results are printed as a table (avg/min step time, GPU peak memory) and as a
-JSON line (`DRIVERRESULT=...`) with per-arm details and the pinning speedup.
-
-> **Note**: the native backend mlocks host memory; raise `RLIMIT_MEMLOCK`
-> (`ulimit -l`) or run as root for multi-GB models.
+Each arm is a subprocess with a fresh rendezvous port. Results: table plus `DRIVERRESULT=` JSON.

benchmarks/pin_memory/activation_offload/README.md

@@ -0,0 +1,7 @@
+# Activation / checkpoint hidden-state offload
+
+Compares `use_pin_memory` True vs False on `CheckpointHiddenStatesOffload`
+with **async / side streams held on**. This is not the async-vs-blocking
+table from DeepSpeed #8282.
+
+See the [parent README](../README.md).

benchmarks/pin_memory/activation_offload/bench.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""
+Activation / checkpoint hidden-state CPU offload: pin vs pageable with async on.
+
+Holds use_streams=True. Compares use_pin_memory True vs False. Do not treat
+this as DeepSpeed #8282's async-vs-blocking table.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+
+_PIN_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if _PIN_ROOT not in sys.path:
+    sys.path.insert(0, _PIN_ROOT)
+
+from common import dist_env, patch_cpp_extension_drop_cxx17, print_arm_result, print_driver_result, run_arm_subprocess
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--hidden", type=int, default=1024)
+    parser.add_argument("--layers", type=int, default=8)
+    parser.add_argument("--batch", type=int, default=1)
+    parser.add_argument("--seq", type=int, default=2048)
+    parser.add_argument("--steps", type=int, default=4)
+    parser.add_argument("--warmup", type=int, default=2)
+    parser.add_argument("--pin", type=int, default=None, help="internal: use_pin_memory 0/1")
+    return parser.parse_args()
+
+
+def run_arm(args):
+    dist_env()
+    patch_cpp_extension_drop_cxx17()
+
+    import torch
+    from torch.utils.checkpoint import checkpoint
+
+    from deepspeed.accelerator import get_accelerator
+    from deepspeed.runtime.activation_checkpointing.offload_activations import CheckpointHiddenStatesOffload
+
+    accelerator = get_accelerator()
+    if not accelerator.is_available():
+        raise RuntimeError(f"No {accelerator.device_name()} device is available")
+    accelerator.set_device(0)
+    device = accelerator.current_device_name()
+
+    class Block(torch.nn.Module):
+
+        def __init__(self, hidden):
+            super().__init__()
+            self.fc1 = torch.nn.Linear(hidden, 4 * hidden)
+            self.fc2 = torch.nn.Linear(4 * hidden, hidden)
+
+        def forward(self, hidden_states):
+            return self.fc2(torch.nn.functional.gelu(self.fc1(hidden_states)))
+
+    class Net(torch.nn.Module):
+
+        def __init__(self, hidden, layers):
+            super().__init__()
+            self.blocks = torch.nn.ModuleList([Block(hidden) for _ in range(layers)])
+
+        def forward(self, hidden_states, offload):
+            x = hidden_states
+            for block in self.blocks:
+                offload.mark(x)
+                x = x + checkpoint(block, x, use_reentrant=False)
+            return x.sum()
+
+    model = Net(args.hidden, args.layers).to(device)
+    opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
+    x = torch.randn(args.batch, args.seq, args.hidden, device=device, requires_grad=True)
+
+    # Async side stream stays on; pin vs pageable is the only axis.
+    offload = CheckpointHiddenStatesOffload(use_pin_memory=bool(args.pin),
+                                            use_streams=True,
+                                            min_offload_bytes=0,
+                                            keep_last_count=1)
+
+    step_times = []
+    for step in range(args.warmup + args.steps):
+        accelerator.synchronize()
+        t0 = time.perf_counter()
+        opt.zero_grad(set_to_none=True)
+        with offload:
+            loss = model(x, offload)
+            loss.backward()
+        opt.step()
+        accelerator.synchronize()
+        if step >= args.warmup:
+            step_times.append(time.perf_counter() - t0)
+        offload.reset()
+
+    def _peak():
+        try:
+            return round(torch.get_device_module(device).max_memory_allocated() / 1e9, 2)
+        except Exception:
+            return None
+
+    print_arm_result({
+        "experiment": "activation_offload",
+        "use_pin_memory": bool(args.pin),
+        "use_streams": True,
+        "device": accelerator.device_name(),
+        "hidden": args.hidden,
+        "layers": args.layers,
+        "batch": args.batch,
+        "seq": args.seq,
+        "steps": len(step_times),
+        "step_avg_s": sum(step_times) / len(step_times),
+        "step_min_s": min(step_times),
+        "gpu_peak_gb": _peak(),
+    })
+
+
+def run_driver(args):
+    script = os.path.abspath(__file__)
+    base = [
+        "--hidden",
+        str(args.hidden),
+        "--layers",
+        str(args.layers),
+        "--batch",
+        str(args.batch),
+        "--seq",
+        str(args.seq),
+        "--steps",
+        str(args.steps),
+        "--warmup",
+        str(args.warmup),
+    ]
+    results = {}
+    for name, pin in (("pageable", 0), ("pinned", 1)):
+        results[name] = run_arm_subprocess(script, base + ["--pin", str(pin)])
+
+    pageable = results["pageable"]
+    pinned = results["pinned"]
+    print("\n================ Activation offload (async on) ================")
+    print(f"device: {pinned['device']}  hidden: {pinned['hidden']}  layers: {pinned['layers']}  "
+          f"batch: {pinned['batch']} seq: {pinned['seq']}")
+    print(f"{'arm':<22}{'avg step (s)':>14}{'min step (s)':>14}{'GPU peak (GB)':>16}")
+    for name in ("pageable", "pinned"):
+        row = results[name]
+        print(f"{name:<22}{row['step_avg_s']:>14.3f}{row['step_min_s']:>14.3f}"
+              f"{(row['gpu_peak_gb'] or 0):>16.2f}")
+    speedup = pageable["step_avg_s"] / pinned["step_avg_s"]
+    print()
+    print(f"pin vs pageable step-time ratio: {speedup:.2f}x (async held on)")
+    print_driver_result({"pageable": pageable, "pinned": pinned, "speedup": speedup})
+
+
+if __name__ == "__main__":
+    parsed = parse_args()
+    if parsed.pin is None:
+        run_driver(parsed)
+    else:
+        run_arm(parsed)

benchmarks/pin_memory/common.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""Shared helpers for pin_memory experiment drivers (subprocess arms, JSON lines)."""
+
+from __future__ import annotations
+
+import json
+import os
+import socket
+import subprocess
+import sys
+
+
+def free_port():
+    # A stale listener from an interrupted rank makes the next init hang in a
+    # collective, so always rendezvous on a fresh ephemeral port.
+    sock = socket.socket()
+    sock.bind(("127.0.0.1", 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+def dist_env():
+    os.environ.update(
+        MASTER_ADDR="127.0.0.1",
+        MASTER_PORT=str(free_port()),
+        RANK="0",
+        WORLD_SIZE="1",
+        LOCAL_RANK="0",
+    )
+
+
+def patch_cpp_extension_drop_cxx17():
+    # torch-nightly requires C++20; SYCL toolchain flags may carry -std=c++17,
+    # which (appearing last) downgrades the dialect and breaks torch headers.
+    import torch.utils.cpp_extension as cpp_ext
+
+    orig_load = cpp_ext.load
+
+    def load_without_cxx17(*args, **kwargs):
+        for key in ("extra_cflags", "extra_cxxflags"):
+            if kwargs.get(key):
+                kwargs[key] = [flag for flag in kwargs[key] if flag != "-std=c++17"]
+        return orig_load(*args, **kwargs)
+
+    cpp_ext.load = load_without_cxx17
+
+
+def print_arm_result(result):
+    print("ARMRESULT=" + json.dumps(result), flush=True)
+
+
+def print_driver_result(summary):
+    print("DRIVERRESULT=" + json.dumps(summary), flush=True)
+
+
+def parse_arm_result(stdout):
+    for line in stdout.splitlines():
+        if line.startswith("ARMRESULT="):
+            return json.loads(line[len("ARMRESULT="):])
+    return None
+
+
+def run_arm_subprocess(script_path, extra_args, env=None):
+    command = [sys.executable, os.path.abspath(script_path), *extra_args]
+    print(f"[driver] {' '.join(command)}", flush=True)
+    proc = subprocess.run(command, env=env or os.environ.copy(), capture_output=True, text=True)
+    arm = parse_arm_result(proc.stdout)
+    if arm is None:
+        print(proc.stdout[-2000:])
+        print(proc.stderr[-2000:], file=sys.stderr)
+        raise RuntimeError(f"arm produced no ARMRESULT (rc={proc.returncode})")
+    return arm

benchmarks/pin_memory/cpu_pin/README.md

@@ -0,0 +1,6 @@
+# CPU-only pin
+
+Native `mlock` vs Torch pin on a CPU host. Torch typically cannot pin without
+an accelerator-capable backend.
+
+See the [parent README](../README.md).

benchmarks/pin_memory/cpu_pin/bench.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""CPU-only: native pin vs Torch (Torch cannot pin without an accelerator)."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+_PIN_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if _PIN_ROOT not in sys.path:
+    sys.path.insert(0, _PIN_ROOT)
+
+from common import print_arm_result, print_driver_result
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--numel", type=int, default=1024 * 1024)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    import torch
+    from deepspeed.accelerator import get_accelerator
+
+    accel = get_accelerator()
+    os.environ["DS_PIN_MEMORY_BACKEND"] = "native"
+    host = torch.empty(args.numel, dtype=torch.float32)
+    native = accel.pin_memory(host.clone(), make_copy=False)
+    native_ok = bool(accel.is_pinned(native))
+    accel.unpin_memory(native)
+
+    os.environ["DS_PIN_MEMORY_BACKEND"] = "torch"
+    torch_ok = None
+    torch_error = None
+    try:
+        pinned = accel.pin_memory(torch.empty_like(host), make_copy=False)
+        torch_ok = bool(accel.is_pinned(pinned))
+    except Exception as exc:
+        torch_error = type(exc).__name__ + ": " + str(exc)
+
+    result = {
+        "experiment": "cpu_pin",
+        "accelerator": accel.device_name(),
+        "native_is_pinned": native_ok,
+        "torch_is_pinned": torch_ok,
+        "torch_error": torch_error,
+    }
+    print_arm_result(result)
+    print_driver_result(result)
+    if not native_ok:
+        raise SystemExit("native pin failed on this host")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/pin_memory/deepcompile_activation/README.md

@@ -0,0 +1,5 @@
+# DeepCompile activation pin
+
+Optional `compile.offload_activation_pin_memory` on/off. Requires DeepCompile.
+
+See the [parent README](../README.md).