""" 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()