src/memory_native/memory.py

"""Memory accounting + a peak-memory gate that anyone can run on stock PyTorch.

Two things live here:
  * memory_report(model): static byte accounting of persistent state (params + buffers),
    plus what the counter weights would cost packed to 6 bits.
  * peak_training_memory(step_fn): the *dynamic* peak during a real train step, measured
    with torch.cuda.max_memory_allocated on CUDA. This is the portable analogue of the
    engine's memory truth gate: it lets a user verify the counter optimizer actually uses
    less training memory than a dense AdamW baseline, on their own GPU, with no custom build.

Honesty: the pure-PyTorch counter layer decodes states to dense tensors around the GEMM, so
its *measured* peak is not yet the 0.75 byte/weight figure -- it still beats dense+Adam on
the optimizer-state pool (no FP master, no Adam moments), which is what these helpers show.
The sub-byte training peak needs the Triton/CUDA packed kernel (see README "Roadmap").
"""
from __future__ import annotations

from typing import Callable, Iterable

import torch
import torch.nn as nn

from .counter import CompactCounterLinear

__all__ = ["memory_report", "fmt_bytes", "peak_training_memory", "compare_training_peak"]


def fmt_bytes(n: float) -> str:
    units = ["B", "KiB", "MiB", "GiB", "TiB"]
    i = 0
    while n >= 1024 and i < len(units) - 1:
        n /= 1024
        i += 1
    return f"{n:.2f} {units[i]}"


def _unique(items: Iterable[torch.Tensor]) -> list[torch.Tensor]:
    seen: set[int] = set()
    out: list[torch.Tensor] = []
    for t in items:
        if t is None:
            continue
        ptr = int(t.untyped_storage().data_ptr()) if t.numel() else id(t)
        if ptr not in seen:
            seen.add(ptr)
            out.append(t)
    return out


def memory_report(model: nn.Module) -> dict[str, float]:
    params = _unique(model.parameters())
    buffers = _unique(model.buffers())
    counter = [m for m in model.modules() if isinstance(m, CompactCounterLinear)]
    # logical weight count (in*out), independent of whether state is stored 1 byte/weight
    # (CompactCounterLinear) or packed 0.75 byte/weight (PackedRMSCounterLinear).
    counter_weights = sum(m.in_features * m.out_features for m in counter)
    fp_params = sum(int(p.numel()) for p in params)
    parameter_bytes = sum(p.numel() * p.element_size() for p in params)
    buffer_bytes = sum(b.numel() * b.element_size() for b in buffers)
    return {
        "counter_weights": float(counter_weights),
        "fp_parameters": float(fp_params),
        "parameter_bytes": float(parameter_bytes),
        "buffer_bytes": float(buffer_bytes),
        "persistent_bytes": float(parameter_bytes + buffer_bytes),
        "counter_packed_6bit_bytes": float(counter_weights * 6 / 8),
        "dense_bf16_adam_rule_of_thumb_bytes": float(16 * (counter_weights + fp_params)),
    }


def peak_training_memory(step_fn: Callable[[], None], device: torch.device) -> int:
    """Run one training step and return the peak allocated bytes during it.

    On CUDA this uses torch.cuda.max_memory_allocated (a true device peak). On CPU torch has
    no allocator peak API, so this returns 0 and callers should fall back to memory_report.
    """
    if device.type == "cuda":
        torch.cuda.synchronize(device)
        torch.cuda.reset_peak_memory_stats(device)
        step_fn()
        torch.cuda.synchronize(device)
        return int(torch.cuda.max_memory_allocated(device))
    step_fn()
    return 0


def compare_training_peak(build_counter: Callable[[], nn.Module],
                          build_dense: Callable[[], "tuple[nn.Module, torch.optim.Optimizer]"],
                          step_inputs: Callable[[], tuple],
                          device: torch.device) -> dict[str, float]:
    """Measure training-step peak memory for a counter model vs a dense+AdamW baseline.

    build_counter() -> counter model (self-updating in backward, no optimizer).
    build_dense()    -> (dense model, its AdamW optimizer).
    step_inputs()    -> (x, y) batch on `device`.
    Returns a dict with both peaks (CUDA) or both persistent-byte reports (CPU fallback).
    """
    import torch.nn.functional as F

    # counter
    cm = build_counter().to(device).train()
    x, y = step_inputs()

    def counter_step():
        logits = cm(x)
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
        loss.backward()  # counter layers self-update here; no optimizer.step()

    # dense + AdamW
    dm, opt = build_dense()
    dm = dm.to(device).train()

    def dense_step():
        opt.zero_grad(set_to_none=True)
        logits = dm(x)
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
        loss.backward()
        opt.step()

    if device.type == "cuda":
        counter_peak = peak_training_memory(counter_step, device)
        dense_peak = peak_training_memory(dense_step, device)
        return {
            "counter_peak_bytes": float(counter_peak),
            "dense_adam_peak_bytes": float(dense_peak),
            "ratio_dense_over_counter": float(dense_peak / max(counter_peak, 1)),
        }
    # CPU: no allocator peak; report persistent state instead.
    counter_step(); dense_step()
    cr, dr = memory_report(cm), memory_report(dm)
    return {
        "counter_persistent_bytes": cr["persistent_bytes"],
        "dense_persistent_bytes": dr["persistent_bytes"],
        "note": float("nan"),  # CPU: see memory_report; peak needs CUDA
    }