diff --git a/toolkit/memory_management/manager_modules.py b/toolkit/memory_management/manager_modules.py index 06d8e0e6..4ea5b0c6 100644 --- a/toolkit/memory_management/manager_modules.py +++ b/toolkit/memory_management/manager_modules.py @@ -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), diff --git a/toolkit/memory_management/test_memory_manager.py b/toolkit/memory_management/test_memory_manager.py new file mode 100644 index 00000000..a6430867 --- /dev/null +++ b/toolkit/memory_management/test_memory_manager.py @@ -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()