RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Deepen offload prefetch pipeline with per-slot events 

Replace the 2-slot ping-pong + single global "compute-started" event
with a depth-N ring buffer where each transfer waits only on the slot
it's reusing (D layers back) instead of the most-recent compute. Applies
to forward and backward, Linear and Conv. Depth is tunable via
AI_TOOLKIT_OFFLOAD_DEPTH (default 4).

Bit-exact vs non-offload (output, grad_input, weight grads). No speedup
on a bandwidth-bound PCIe link (already saturated at depth 2), but the
cleaner per-slot design removes the fragile shared-event serialization
and lets deeper prefetch help on faster buses.
This commit is contained in:
Jaret Burkett 2026-06-07 16:07:13 -06:00
parent cdff6e36aa
commit c78b1404e3
2 changed files with 486 additions and 149 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

302
toolkit/memory_management/manager_modules.py
View File

@ -6,6 +6,8 @@ https://github.com/lodestone-rock/RamTorch/blob/main/ramtorch/modules/linear.py
I simply modified it to work with a memory management model and with AI Toolkit's models
"""
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
@ -18,6 +20,12 @@ if TYPE_CHECKING:
# --- Per-device global state registry ---
_DEVICE_STATE = {}
# How many layers deep to prefetch weights. The old ping-pong used 2 slots, which
# only lets one transfer overlap one compute (1-deep). A deeper ring lets Python
# enqueue several layers ahead so the H2D stream stays saturated instead of
# stalling on a per-layer sync. Override with AI_TOOLKIT_OFFLOAD_DEPTH.
PIPELINE_DEPTH = int(os.environ.get("AI_TOOLKIT_OFFLOAD_DEPTH", "4"))
def _get_device_state(device: torch.device):
"""Get or initialize per-device state."""
@ -31,31 +39,101 @@ def _get_device_state(device: torch.device):
return _DEVICE_STATE[device]
if device not in _DEVICE_STATE:
d = max(2, PIPELINE_DEPTH)
with torch.cuda.device(device):
_DEVICE_STATE[device] = {
# streams & events
"depth": d,
# streams
"transfer_stream": torch.cuda.Stream(device=device),
"transfer_grad_stream": torch.cuda.Stream(device=device),
"transfer_forward_finished_event": torch.cuda.Event(),
"compute_forward_start_event": torch.cuda.Event(),
"transfer_backward_finished_event": torch.cuda.Event(),
"transfer_weight_backward_finished_event": torch.cuda.Event(),
"compute_backward_start_event": torch.cuda.Event(),
"compute_backward_finished_event": torch.cuda.Event(),
# ping-pong buffers
"w_buffers": [None, None],
"b_buffers": [None, None],
"w_bwd_buffers": [None, None],
# device-side staging for grads to be sent to CPU
"w_grad_buffers": [None, None],
"b_grad_buffers": [None, None],
# clocks
# forward weight ring: slot_ready = H2D done, slot_free = compute
# that consumed the slot done (so it can be overwritten).
"w_buffers": [None] * d,
"b_buffers": [None] * d,
"fwd_slot_ready": [torch.cuda.Event() for _ in range(d)],
"fwd_slot_free": [torch.cuda.Event() for _ in range(d)],
"forward_clk": 0,
# backward weight ring (re-fetch for grad-input).
"w_bwd_buffers": [None] * d,
"bwd_slot_ready": [torch.cuda.Event() for _ in range(d)],
"bwd_slot_free": [torch.cuda.Event() for _ in range(d)],
"backward_clk": 0,
# backward grad-staging ring (device-side grads -> CPU).
"w_grad_buffers": [None] * d,
"b_grad_buffers": [None] * d,
"grad_compute_done": [torch.cuda.Event() for _ in range(d)],
"grad_xfer_done": [torch.cuda.Event() for _ in range(d)],
}
return _DEVICE_STATE[device]
# ---- ring-buffer staging helpers -----------------------------------------
#
# Each transfer waits only on the event for the *specific slot* it is about to
# overwrite (the compute that used that slot D layers ago), not on a single
# global "compute started" event. With D slots that prior compute is long done,
# so the transfer stream never actually stalls and stays D layers ahead of
# compute. This is the deeper-pipeline + relaxed-dependency change in one.
def _stage_forward_weight(state, device, materialize, weight_cpu, bias_cpu):
"""H2D the next forward weight (+bias) into its ring slot; return (idx, w, b).
Caller runs compute, then calls _release_forward_slot(state, idx)."""
d = state["depth"]
idx = state["forward_clk"]
state["forward_clk"] = (idx + 1) % d
ts = state["transfer_stream"]
with torch.cuda.stream(ts):
ts.wait_event(state["fwd_slot_free"][idx])
state["w_buffers"][idx] = materialize(weight_cpu, device)
state["b_buffers"][idx] = (
bias_cpu.to(device, non_blocking=True) if bias_cpu is not None else None
)
state["fwd_slot_ready"][idx].record()
torch.cuda.current_stream().wait_event(state["fwd_slot_ready"][idx])
return idx, state["w_buffers"][idx], state["b_buffers"][idx]
def _release_forward_slot(state, idx):
# Slot is reusable once the compute stream finishes the op that read it.
state["fwd_slot_free"][idx].record()
def _stage_backward_weight(state, device, materialize, weight_cpu):
"""H2D the next backward weight into its ring slot; return (idx, w).
Caller runs grad-input compute, then _release_backward_weight_slot."""
d = state["depth"]
idx = state["backward_clk"]
state["backward_clk"] = (idx + 1) % d
ts = state["transfer_stream"]
with torch.cuda.stream(ts):
ts.wait_event(state["bwd_slot_free"][idx])
state["w_bwd_buffers"][idx] = materialize(weight_cpu)
state["bwd_slot_ready"][idx].record()
torch.cuda.current_stream().wait_event(state["bwd_slot_ready"][idx])
return idx, state["w_bwd_buffers"][idx]
def _release_backward_weight_slot(state, idx):
state["bwd_slot_free"][idx].record()
def _stage_grads_to_cpu(state, idx, grad_w_gpu, grad_b_gpu):
"""Copy freshly-computed device grads (in staging slot idx) to CPU on the
grad stream, overlapping the next H2D. Returns (grad_w_cpu, grad_b_cpu)."""
gs = state["transfer_grad_stream"]
state["grad_compute_done"][idx].record() # on the compute stream
grad_w_cpu = grad_b_cpu = None
with torch.cuda.stream(gs):
gs.wait_event(state["grad_compute_done"][idx])
if grad_w_gpu is not None:
grad_w_cpu = grad_w_gpu.to("cpu", non_blocking=True)
if grad_b_gpu is not None:
grad_b_cpu = grad_b_gpu.to("cpu", non_blocking=True)
state["grad_xfer_done"][idx].record()
return grad_w_cpu, grad_b_cpu
# (ADD) detect torchao wrapper tensors
def _is_ao_quantized_tensor(t: Optional[torch.Tensor]) -> bool:
if t is None:
@ -209,24 +287,11 @@ class _BouncingLinearFn(torch.autograd.Function):
return out.to(x.device)
state = _get_device_state(device)
ts = state["transfer_stream"]
w_bufs, b_bufs = state["w_buffers"], state["b_buffers"]
ev_tx_f = state["transfer_forward_finished_event"]
ev_cu_s = state["compute_forward_start_event"]
idx = state["forward_clk"]
with torch.cuda.stream(ts):
ts.wait_event(ev_cu_s)
w_bufs[idx] = _materialize_linear_weight(weight_cpu, device)
b_bufs[idx] = (
bias_cpu.to(device, non_blocking=True) if bias_cpu is not None else None
)
state["forward_clk"] ^= 1
ev_tx_f.record()
torch.cuda.current_stream().wait_event(ev_tx_f)
ev_cu_s.record()
out = F.linear(x, w_bufs[idx], b_bufs[idx])
idx, w_gpu, b_gpu = _stage_forward_weight(
state, device, _materialize_linear_weight, weight_cpu, bias_cpu
)
out = F.linear(x, w_gpu, b_gpu)
_release_forward_slot(state, idx)
ctx.save_for_backward(x, weight_cpu, bias_cpu)
ctx.device = device
@ -268,19 +333,6 @@ class _BouncingLinearFn(torch.autograd.Function):
return grad_input.to(grad_out.device), grad_weight, grad_bias, None
state = _get_device_state(device)
transfer_stream = state["transfer_stream"]
transfer_grad_stream = state["transfer_grad_stream"]
w_bwd_buffers = state["w_bwd_buffers"]
w_grad_buffers = state["w_grad_buffers"]
b_grad_buffers = state["b_grad_buffers"]
ev_tx_b = state["transfer_backward_finished_event"]
ev_tx_w_bwd_done = state["transfer_weight_backward_finished_event"]
ev_cu_b_start = state["compute_backward_start_event"]
ev_cu_b_finish = state["compute_backward_finished_event"]
idx = state["backward_clk"]
# GPU-side dequant/cast for quantized; float path unchanged
def _materialize_for_bwd(cpu_w):
@ -297,45 +349,35 @@ class _BouncingLinearFn(torch.autograd.Function):
w = cpu_w.to(device, non_blocking=True)
return w
with torch.cuda.stream(transfer_stream):
transfer_stream.wait_event(ev_cu_b_start)
w_bwd_buffers[idx] = _materialize_for_bwd(weight_cpu)
state["backward_clk"] ^= 1
ev_tx_b.record()
torch.cuda.current_stream().wait_event(ev_tx_b)
ev_cu_b_start.record()
idx, w_bwd = _stage_backward_weight(
state, device, _materialize_for_bwd, weight_cpu
)
# grad wrt input (GPU)
grad_input = grad_out.to(dtype=target_dtype) @ w_bwd_buffers[idx]
grad_input = grad_out.to(dtype=target_dtype) @ w_bwd
_release_backward_weight_slot(state, idx)
# ensure previous grad-to-CPU transfer that used this slot finished
torch.cuda.current_stream().wait_event(ev_tx_w_bwd_done)
# compute grads if float masters exist
# compute grads if float masters exist (frozen/quantized bases skip this)
grad_weight = None
grad_bias = None
if (
need_w = (
getattr(weight_cpu, "requires_grad", False)
and weight_cpu.dtype.is_floating_point
):
w_grad_buffers[idx] = grad_out.flatten(0, -2).T @ x.flatten(0, -2)
if bias_cpu is not None and getattr(bias_cpu, "requires_grad", False):
reduce_dims = tuple(range(grad_out.ndim - 1))
b_grad_buffers[idx] = grad_out.sum(dim=reduce_dims)
ev_cu_b_finish.record()
with torch.cuda.stream(transfer_grad_stream):
transfer_grad_stream.wait_event(ev_cu_b_finish)
if (
getattr(weight_cpu, "requires_grad", False)
and weight_cpu.dtype.is_floating_point
):
grad_weight = w_grad_buffers[idx].to("cpu", non_blocking=True)
if bias_cpu is not None and getattr(bias_cpu, "requires_grad", False):
grad_bias = b_grad_buffers[idx].to("cpu", non_blocking=True)
state["transfer_weight_backward_finished_event"].record()
)
need_b = bias_cpu is not None and getattr(bias_cpu, "requires_grad", False)
if need_w or need_b:
# ensure the prior grad D2H using this staging slot finished
torch.cuda.current_stream().wait_event(state["grad_xfer_done"][idx])
w_grad_gpu = b_grad_gpu = None
if need_w:
w_grad_gpu = grad_out.flatten(0, -2).T @ x.flatten(0, -2)
state["w_grad_buffers"][idx] = w_grad_gpu
if need_b:
b_grad_gpu = grad_out.sum(dim=tuple(range(grad_out.ndim - 1)))
state["b_grad_buffers"][idx] = b_grad_gpu
grad_weight, grad_bias = _stage_grads_to_cpu(
state, idx, w_grad_gpu, b_grad_gpu
)
return grad_input.to(dtype=grad_out.dtype), grad_weight, grad_bias, None
@ -389,24 +431,11 @@ class _BouncingConv2dFn(torch.autograd.Function):
return out.to(x.device)
state = _get_device_state(device)
ts = state["transfer_stream"]
w_bufs, b_bufs = state["w_buffers"], state["b_buffers"]
ev_tx_f = state["transfer_forward_finished_event"]
ev_cu_s = state["compute_forward_start_event"]
idx = state["forward_clk"]
with torch.cuda.stream(ts):
ts.wait_event(ev_cu_s)
w_bufs[idx] = _materialize_conv_weight(weight_cpu, device)
b_bufs[idx] = (
bias_cpu.to(device, non_blocking=True) if bias_cpu is not None else None
)
state["forward_clk"] ^= 1
ev_tx_f.record()
torch.cuda.current_stream().wait_event(ev_tx_f)
ev_cu_s.record()
out = F.conv2d(x, w_bufs[idx], b_bufs[idx], stride, padding, dilation, groups)
idx, w_gpu, b_gpu = _stage_forward_weight(
state, device, _materialize_conv_weight, weight_cpu, bias_cpu
)
out = F.conv2d(x, w_gpu, b_gpu, stride, padding, dilation, groups)
_release_forward_slot(state, idx)
ctx.save_for_backward(x, weight_cpu, bias_cpu)
ctx.meta = (device, stride, padding, dilation, groups, target_dtype)
@ -475,19 +504,6 @@ class _BouncingConv2dFn(torch.autograd.Function):
)
state = _get_device_state(device)
transfer_stream = state["transfer_stream"]
transfer_grad_stream = state["transfer_grad_stream"]
w_bwd_buffers = state["w_bwd_buffers"]
w_grad_buffers = state["w_grad_buffers"]
b_grad_buffers = state["b_grad_buffers"]
ev_tx_b = state["transfer_backward_finished_event"]
ev_tx_w_bwd_done = state["transfer_weight_backward_finished_event"]
ev_cu_b_start = state["compute_backward_start_event"]
ev_cu_b_finish = state["compute_backward_finished_event"]
idx = state["backward_clk"]
# GPU-side dequant/cast for quantized; float path unchanged
def _materialize_for_bwd(cpu_w):
@ -504,63 +520,51 @@ class _BouncingConv2dFn(torch.autograd.Function):
w = cpu_w.to(device, non_blocking=True)
return w
# Stage weights for input-grad compute
with torch.cuda.stream(transfer_stream):
transfer_stream.wait_event(ev_cu_b_start)
w_bwd_buffers[idx] = _materialize_for_bwd(weight_cpu)
state["backward_clk"] ^= 1
ev_tx_b.record()
torch.cuda.current_stream().wait_event(ev_tx_b)
ev_cu_b_start.record()
idx, w_bwd = _stage_backward_weight(
state, device, _materialize_for_bwd, weight_cpu
)
from torch.nn.grad import conv2d_input, conv2d_weight # type: ignore
grad_input = conv2d_input(
x.shape,
w_bwd_buffers[idx],
w_bwd,
grad_out.to(dtype=target_dtype),
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
)
_release_backward_weight_slot(state, idx)
# Ensure previous grad transfer that used this slot is done
torch.cuda.current_stream().wait_event(ev_tx_w_bwd_done)
# Compute heavy grads on GPU into staging buffers
# Compute heavy grads on GPU into staging buffers (frozen bases skip this)
grad_weight = None
grad_bias = None
if (
need_w = (
getattr(weight_cpu, "requires_grad", False)
and weight_cpu.dtype.is_floating_point
):
w_grad_buffers[idx] = conv2d_weight(
x,
weight_cpu.shape,
grad_out,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
)
need_b = bias_cpu is not None and getattr(bias_cpu, "requires_grad", False)
if need_w or need_b:
torch.cuda.current_stream().wait_event(state["grad_xfer_done"][idx])
w_grad_gpu = b_grad_gpu = None
if need_w:
w_grad_gpu = conv2d_weight(
x,
weight_cpu.shape,
grad_out,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
)
state["w_grad_buffers"][idx] = w_grad_gpu
if need_b:
b_grad_gpu = grad_out.sum(dim=(0, 2, 3))
state["b_grad_buffers"][idx] = b_grad_gpu
grad_weight, grad_bias = _stage_grads_to_cpu(
state, idx, w_grad_gpu, b_grad_gpu
)
if bias_cpu is not None and getattr(bias_cpu, "requires_grad", False):
b_grad_buffers[idx] = grad_out.sum(dim=(0, 2, 3))
ev_cu_b_finish.record()
# Launch CPU copies on the dedicated grad stream (overlaps with next H2D)
with torch.cuda.stream(transfer_grad_stream):
transfer_grad_stream.wait_event(ev_cu_b_finish)
if (
getattr(weight_cpu, "requires_grad", False)
and weight_cpu.dtype.is_floating_point
):
grad_weight = w_grad_buffers[idx].to("cpu", non_blocking=True)
if bias_cpu is not None and getattr(bias_cpu, "requires_grad", False):
grad_bias = b_grad_buffers[idx].to("cpu", non_blocking=True)
state["transfer_weight_backward_finished_event"].record()
return (
grad_input.to(dtype=grad_out.dtype),

333
toolkit/memory_management/test_memory_manager.py Normal file
View File

@ -0,0 +1,333 @@
"""
Memory-manager (layer offloading) benchmark on a ~1B parameter diffusion-style
transformer. The base model is frozen and a LoRA is trained on top of it (the
realistic training setup), so only the LoRA params get grads/optimizer state.
Reports speed (ms/step) and peak VRAM for the 2x2 matrix of:
- bfloat16 base vs bfloat16 + float8-quantized base (torchao weight-only)
- no offloading vs 100% offloading
The MemoryManager keeps the frozen base weights pinned on the CPU and streams
them onto the GPU per forward/backward (dequantizing float8 weights on the GPU).
The LoRA wraps each base linear, so its forward calls the bounced base forward
and adds the low-rank update. This trades VRAM for PCIe traffic, so the table
shows what that trade actually costs.
Run directly: `python test_memory_manager.py`
"""
import contextlib
import gc
import io
import os
import sys
import threading
import time
import psutil
import torch
import torch.nn as nn
import torch.nn.functional as F
class RamMonitor:
"""Sample process RSS in a background thread and track the peak. Pinned
CPU weights (from offloading) live in RSS, so this captures the host-RAM
cost the GPU-side peak doesn't see."""
def __init__(self, interval: float = 0.005):
self.interval = interval
self._proc = psutil.Process()
self.peak = 0
def _run(self):
while not self._stop:
self.peak = max(self.peak, self._proc.memory_info().rss)
time.sleep(self.interval)
def __enter__(self):
self.peak = self._proc.memory_info().rss
self._stop = False
self._thread = threading.Thread(target=self._run, daemon=True)
self._thread.start()
return self
def __exit__(self, *exc):
self._stop = True
self._thread.join()
# Allow running this file directly without setting PYTHONPATH.
# Toolkit imports happen inside main() so they pick this up.
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if _REPO_ROOT not in sys.path:
sys.path.insert(0, _REPO_ROOT)
# ---- model ---------------------------------------------------------------
class TransformerBlock(nn.Module):
def __init__(self, d_model: int, n_heads: int, d_ff: int):
super().__init__()
self.n_heads = n_heads
self.d_head = d_model // n_heads
self.ln1 = nn.LayerNorm(d_model)
self.q = nn.Linear(d_model, d_model, bias=False)
self.k = nn.Linear(d_model, d_model, bias=False)
self.v = nn.Linear(d_model, d_model, bias=False)
self.o = nn.Linear(d_model, d_model, bias=False)
self.ln2 = nn.LayerNorm(d_model)
self.ffn_up = nn.Linear(d_model, d_ff, bias=False)
self.ffn_down = nn.Linear(d_ff, d_model, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, S, D = x.shape
h = self.ln1(x)
q = self.q(h).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
k = self.k(h).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
v = self.v(h).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
a = F.scaled_dot_product_attention(q, k, v)
a = a.transpose(1, 2).contiguous().view(B, S, D)
x = x + self.o(a)
h = self.ln2(x)
x = x + self.ffn_down(F.gelu(self.ffn_up(h)))
return x
class Transformer(nn.Module):
def __init__(self, d_model=2048, n_heads=16, n_layers=24, d_ff=8192):
super().__init__()
self.blocks = nn.ModuleList([
TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)
])
self.norm = nn.LayerNorm(d_model)
self.gradient_checkpointing = False
def forward(self, x: torch.Tensor) -> torch.Tensor:
for b in self.blocks:
# Gate on is_grad_enabled (not self.training): checkpointing only
# helps and only works when a backward will actually be run.
if self.gradient_checkpointing and torch.is_grad_enabled():
x = torch.utils.checkpoint.checkpoint(b, x, use_reentrant=False)
else:
x = b(x)
return self.norm(x)
# ---- benchmark -----------------------------------------------------------
DEVICE = torch.device("cuda")
DTYPE = torch.bfloat16
QTYPE = "float8"
D_MODEL = 2048
N_HEADS = 16
N_LAYERS = 24
D_FF = 8192
BATCH = 1
SEQ = 1024
WARMUP = 3
ITERS = 10
LORA_RANK = 32
LR = 1e-4
# Full matrix: {bf16, float8} x {no offload, 100% offload} x {ckpt on, ckpt off}.
# Offloading parks weights in CPU RAM; turning off checkpointing keeps activations
# resident in VRAM. We report peak VRAM *and* peak system RAM so both show up.
# (label, quantize, offload_percent, grad_checkpointing)
RUNS = []
for _do_q, _qlabel in [(False, "bf16"), (True, "float8")]:
for _off, _olabel in [(None, ""), (1.0, "+off")]:
for _ckpt in [True, False]:
_label = f"{_qlabel}{_olabel} ckpt={'on' if _ckpt else 'off'}"
RUNS.append((_label, _do_q, _off, _ckpt))
def build_model():
torch.manual_seed(0)
# Build on CPU; the caller decides how it reaches the GPU.
return Transformer(D_MODEL, N_HEADS, N_LAYERS, D_FF).to(dtype=DTYPE)
def build_lora(transformer):
"""Attach a trainable LoRA to the (frozen) transformer, the same way the
trainer does it. Returns the network; its forward hijacks each base linear."""
from toolkit.config_modules import NetworkConfig
from toolkit.lora_special import LoRASpecialNetwork
network_config = NetworkConfig(
type="lora",
linear=LORA_RANK,
linear_alpha=LORA_RANK,
transformer_only=True,
)
LoRASpecialNetwork.LORA_PREFIX_UNET = "lora_transformer"
network = LoRASpecialNetwork(
text_encoder=None,
unet=transformer,
lora_dim=network_config.linear,
multiplier=1.0,
alpha=network_config.linear_alpha,
train_unet=True,
train_text_encoder=False,
network_config=network_config,
network_type=network_config.type,
transformer_only=network_config.transformer_only,
is_transformer=True,
target_lin_modules=["Transformer"],
)
network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True)
network.force_to(DEVICE, DTYPE)
network._update_torch_multiplier()
network.is_active = True
network.train()
return network
def benchmark(results: list, label: str, do_quantize: bool, offload_percent, grad_checkpointing):
from toolkit.memory_management import MemoryManager
from toolkit.util.quantize import quantize, get_qtype
from optimum.quanto import freeze
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
network = None
model = build_model()
model.gradient_checkpointing = grad_checkpointing
model.to(DEVICE)
if do_quantize:
# Quantize the linear weights to float8 on the GPU (torchao weight-only),
# exactly as a quantized base model is prepared before training.
quantize(model, weights=get_qtype(QTYPE))
freeze(model)
# Base model is frozen; only the LoRA trains.
model.requires_grad_(False)
if offload_percent is None:
# Baseline: whole base model stays on the GPU.
model.to(DEVICE)
else:
# Offloading: managed linears stay pinned on CPU and bounce per step
# (float8 weights are dequantized on the GPU); unmanaged modules (norms)
# move to the GPU via the patched .to(). Attach BEFORE the LoRA so the
# LoRA wraps the bouncing forward. Layer sampling is seeded for repro.
import random
random.seed(0)
MemoryManager.attach(model, DEVICE, offload_percent=offload_percent)
model.to(DEVICE)
# build_lora prints a banner per layer; mute it so the final table is clean.
with contextlib.redirect_stdout(io.StringIO()):
network = build_lora(model)
params = network.prepare_optimizer_params(LR, LR, LR)
opt = torch.optim.AdamW(params, lr=LR)
x = torch.randn(BATCH, SEQ, D_MODEL, device=DEVICE, dtype=DTYPE)
try:
for _ in range(WARMUP):
opt.zero_grad(set_to_none=True)
model(x).sum().backward()
opt.step()
torch.cuda.synchronize()
# Measure the steady-state TRAINING peak, not the one-time setup load.
# (Offload first parks the whole model on the GPU before bouncing it to
# CPU; counting that transient would hide the real per-step footprint.)
torch.cuda.reset_peak_memory_stats()
t0 = time.perf_counter()
with RamMonitor() as ram:
for _ in range(ITERS):
opt.zero_grad(set_to_none=True)
model(x).sum().backward()
opt.step()
torch.cuda.synchronize()
dt = (time.perf_counter() - t0) / ITERS * 1000
peak = torch.cuda.max_memory_allocated() / 1024**3
ram_gb = ram.peak / 1024**3
results.append({"label": label, "ms": dt, "peak": peak, "ram": ram_gb, "ok": True})
except torch.cuda.OutOfMemoryError:
results.append({"label": label, "ms": float("inf"), "peak": float("inf"), "ram": float("inf"), "ok": False, "note": "OOM"})
except Exception as e:
print(f" {label} failed: {type(e).__name__}: {e}", flush=True)
results.append({"label": label, "ms": float("inf"), "peak": float("inf"), "ram": float("inf"), "ok": False, "note": "ERR"})
finally:
if offload_percent is not None:
MemoryManager.detach(model)
del opt, network, model
gc.collect()
torch.cuda.empty_cache()
def print_table(results: list):
headers = ["#", "Configuration", "Peak VRAM", "Peak RAM", "Time/step"]
rows = []
for i, r in enumerate(results, 1):
if not r["ok"]:
rows.append([str(i), r["label"], r.get("note", "OOM"), "-", "-"])
continue
rows.append([str(i), r["label"], f"{r['peak']:.2f} GB", f"{r['ram']:.2f} GB", f"{r['ms']:.1f} ms"])
widths = [max(len(str(row[c])) for row in [headers] + rows) for c in range(len(headers))]
def fmt(row, sep=" │ "):
return sep.join(s.ljust(widths[c]) if c == 1 else s.rjust(widths[c]) for c, s in enumerate(row))
line_top = "─" * (sum(widths) + 3 * (len(widths) - 1))
print()
print(line_top)
print(fmt(headers))
print(line_top)
for row in rows:
print(fmt(row))
print(line_top)
def run_one(idx: int):
"""Run a single config and print its result as JSON. Invoked in a fresh
subprocess so peak RAM (and VRAM) are isolated — pinned-host and CUDA host
caches don't release between runs, so in-process RAM peaks would accumulate."""
import json
label, do_quantize, offload_percent, grad_checkpointing = RUNS[idx]
results: list = []
benchmark(results, label, do_quantize, offload_percent, grad_checkpointing)
print("RESULT " + json.dumps(results[0]), flush=True)
def main():
import json
import subprocess
n_params = sum(p.numel() for p in build_model().parameters())
print(f"Model: {N_LAYERS} blocks × d_model={D_MODEL} × d_ff={D_FF}")
print(f" {n_params/1e6:.1f}M params")
print(f"dtype: {str(DTYPE).replace('torch.', '')} (quant qtype: {QTYPE})")
print(f"Train: LoRA rank={LORA_RANK} on a frozen base")
print(f"Step: batch={BATCH}, seq={SEQ}")
print(f"Timing: {WARMUP} warmup + {ITERS} timed iters")
print(f"Configs: {len(RUNS)} (each in an isolated subprocess)")
results: list = []
for idx, run in enumerate(RUNS):
print(f" running {run[0]}...", flush=True)
proc = subprocess.run(
[sys.executable, __file__, "--run", str(idx)],
capture_output=True, text=True,
)
line = next((ln for ln in proc.stdout.splitlines() if ln.startswith("RESULT ")), None)
if line is None:
print(f" {run[0]} produced no result:\n{proc.stdout}\n{proc.stderr}", flush=True)
results.append({"label": run[0], "ok": False, "note": "ERR"})
continue
results.append(json.loads(line[len("RESULT "):]))
print_table(results)
if __name__ == "__main__":
if len(sys.argv) >= 3 and sys.argv[1] == "--run":
run_one(int(sys.argv[2]))
else:
main()