vllm-omni/benchmarks/edge_harness/memory_ledger.py

"""Where the resident memory actually is, by attribution rather than subtraction.

The report's breakdown was arithmetic: weights, imports, and "unexplained
~1.6 GB" left over after subtracting the first two from the total. Three
problems with that. It assigned a bf16 embedding to a checkpoint whose config
says ``quantize_embedding: true``; it left the prefill duplicate out of the
weight line entirely; and a remainder computed from two wrong terms is not a
measurement of anything.

So this counts instead:

* **Unique tensor storages.** Walking ``named_parameters`` double-counts tied
  weights -- Spark ties its embedding to its head, and the quantized path ties
  five tensors at once -- so storages are keyed by ``data_ptr`` and counted
  once, and the tied ones are reported separately rather than silently
  dropped.
* **The KV cache**, found by attribute rather than assumed from a formula --
  and split into *allocated* and *resident*. A cache that has been allocated
  but never written costs address space, not memory, so counting its size as
  attributed memory would hide the remainder rather than explain it.
  Residency comes from ``mincore(2)``, which reports per-page whether a
  mapping is actually in core.
* **Every mapping in /proc/self/smaps**, bucketed into heap, anonymous, file
  and the rest. ``smaps_rollup`` cannot do this: it is an aggregate, which is
  the whole thing this is trying to break apart.

What it cannot do is attribute native allocations that no longer belong to a
live tensor -- freed allocator arenas, inductor's generated code, transient
buffers the caching allocator is holding. Those land in "unattributed anon",
and the point of the exercise is to see how big that really is once the
attributable part is counted properly.
"""

from __future__ import annotations

import argparse, collections, json, os, sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from bench_harness import provenance  # noqa: E402


def resident_mb(ptr: int, nbytes: int) -> float:
    """How much of a mapping is actually in core, via mincore(2).

    Allocated is not resident: an untouched KV cache is address space. Without
    this the ledger would "explain" memory the process never used.
    """
    import ctypes

    if nbytes <= 0:
        return 0.0
    page = os.sysconf("SC_PAGE_SIZE")
    start = ptr - (ptr % page)                 # mincore needs a page boundary
    length = nbytes + (ptr - start)
    pages = (length + page - 1) // page
    vec = (ctypes.c_ubyte * pages)()
    libc = ctypes.CDLL("libc.so.6", use_errno=True)
    if libc.mincore(ctypes.c_void_p(start), ctypes.c_size_t(length), vec) != 0:
        return -1.0                            # unknown, rather than a guess
    return round(sum(v & 1 for v in vec) * page / 2**20, 1)


def _rss(field: str) -> float:
    with open("/proc/self/status") as f:
        for line in f:
            if line.startswith(field):
                return round(int(line.split()[1]) / 1024, 1)
    return 0.0


def smaps_breakdown() -> dict[str, float]:
    """Anonymous resident bytes per mapping kind. Needs smaps, not rollup."""
    buckets: collections.Counter = collections.Counter()
    name = "[anon]"
    try:
        with open("/proc/self/smaps") as f:
            for line in f:
                if not line[:1].isspace() and "-" in line.split()[0]:
                    parts = line.split()
                    path = parts[5] if len(parts) > 5 else "[anon]"
                    if path.startswith("["):
                        name = path
                    elif path.endswith(".so") or ".so." in path:
                        name = "shared libs"
                    elif path:
                        name = "file-backed"
                    else:
                        name = "[anon]"
                elif line.startswith("Anonymous:"):
                    buckets[name] += int(line.split()[1])
    except OSError:
        return {}
    return {k: round(v / 1024, 1) for k, v in buckets.most_common()}


def storage_ledger(model) -> dict:
    """Unique tensor storages, by role, counting each byte once."""
    import torch

    seen: dict[int, int] = {}
    biggest: list = []
    by_role: collections.Counter = collections.Counter()
    counts: collections.Counter = collections.Counter()
    tied_bytes = 0

    def role_of(name: str) -> str:
        low = name.lower()
        if "embed" in low or "lm_head" in low:
            return "embedding / head"
        for key in ("int4pack", "halves", "scales_and_zeros", "group_scale",
                    "group_zero", "weight_scale", "weight_zero", "weight_packed"):
            if key in low:
                return f"quant:{key}"
        if "weight" in low:
            return "other weights"
        return "other"

    for name, tensor in list(model.named_parameters()) + list(model.named_buffers()):
        if not isinstance(tensor, torch.Tensor) or tensor.device.type != "cpu":
            continue
        try:
            ptr = tensor.untyped_storage().data_ptr()
            nbytes = tensor.untyped_storage().nbytes()
        except Exception:
            ptr, nbytes = tensor.data_ptr(), tensor.numel() * tensor.element_size()
        if ptr in seen:
            tied_bytes += nbytes          # tied / shared: counted once already
            continue
        seen[ptr] = nbytes
        by_role[role_of(name)] += nbytes
        counts[role_of(name)] += 1
        if nbytes >= 16 * 2**20:
            biggest.append((name, nbytes, resident_mb(ptr, nbytes),
                            tuple(tensor.shape), tensor.dtype))

    biggest.sort(key=lambda t: -t[1])
    return {
        "unique_storage_mb": round(sum(seen.values()) / 2**20, 1),
        "tied_duplicate_mb": round(tied_bytes / 2**20, 1),
        "by_role_mb": {k: round(v / 2**20, 1) for k, v in by_role.most_common()},
        "tensors_by_role": dict(counts),
        # "other" is where an unclassified tensor hides, so name the largest
        "largest_tensors": [
            {"name": n, "mb": round(b / 2**20, 1), "resident_mb": r,
             "shape": list(sh), "dtype": str(dt)}
            for n, b, r, sh, dt in biggest[:15]
        ],
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--quantization", default=None)
    ap.add_argument("--max-model-len", type=int, default=2048)
    ap.add_argument("--tag", default="ledger")
    ap.add_argument("--out", required=True)
    a = ap.parse_args()

    anon_before = _rss("RssAnon:")
    import vllm_omni  # noqa: F401
    import torch
    from vllm import LLM

    anon_after_import = _rss("RssAnon:")
    llm = LLM(model=a.model, dtype="bfloat16", max_model_len=a.max_model_len,
              max_num_seqs=1, quantization=a.quantization,
              enable_prefix_caching=False, trust_remote_code=False,
              disable_log_stats=True)

    # The model only exists in this process when the executor stayed uniproc;
    # say so rather than reporting a ledger of nothing.
    try:
        engine = llm.llm_engine
        runner = engine.engine_core.engine_core.model_executor.driver_worker.model_runner
        model = runner.model
    except AttributeError as exc:
        raise SystemExit(
            f"cannot reach the model in this process ({exc}). Set "
            f"VLLM_ENABLE_V1_MULTIPROCESSING=0 so the executor stays uniproc; "
            f"otherwise the weights are in a worker and there is nothing here "
            f"to attribute."
        )

    ledger = storage_ledger(model)

    kv_mb = kv_resident = 0.0
    try:
        caches = runner.kv_caches
        seen = set()
        for c in caches:
            if isinstance(c, torch.Tensor) and c.data_ptr() not in seen:
                seen.add(c.data_ptr())
                nbytes = c.numel() * c.element_size()
                kv_mb += nbytes / 2**20
                r = resident_mb(c.data_ptr(), nbytes)
                kv_resident += max(r, 0.0)
    except Exception:
        kv_mb = kv_resident = -1.0

    anon_total = _rss("RssAnon:")
    # Only resident KV counts: allocated-but-untouched cache is address space.
    attributed = (ledger["unique_storage_mb"] + max(kv_resident, 0.0)
                  + anon_after_import)
    rec = {
        "tag": a.tag, "model": a.model, "quantization": a.quantization,
        "anon_before_import_mb": anon_before,
        "anon_after_import_mb": anon_after_import,
        "anon_total_mb": anon_total,
        "kv_cache_allocated_mb": round(kv_mb, 1),
        "kv_cache_resident_mb": round(kv_resident, 1),
        **ledger,
        "attributed_mb": round(attributed, 1),
        "unattributed_anon_mb": round(anon_total - attributed, 1),
        "smaps_anon_mb": smaps_breakdown(),
        "provenance": provenance({"scope": "in-process ledger"}),
    }
    print(json.dumps({k: v for k, v in rec.items() if k != "provenance"}, indent=2))
    with open(a.out, "w") as f:
        json.dump(rec, f, indent=2)


if __name__ == "__main__":
    main()