ai-toolkit/toolkit/memory_management/test_memory_manager.py

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