From c440616c8b149ee4408ff166b2e0c3d58fa4275d Mon Sep 17 00:00:00 2001 From: tzhouam Date: Sat, 12 Sep 2026 14:27:11 +0800 Subject: [PATCH] [Edge][CPU] Return the weight-loading churn to the OS: 7338 -> 3510 MB Loading a quantized checkpoint allocates far more than the weights occupy. `_process_gptq_weights_w4a8` unpacks 4-bit codes to int32 -- four bytes per weight -- permutes, repacks, and drops the intermediates. For a 1.7 B model that is gigabytes of short-lived allocation, correctly freed, and still resident afterwards because a caching allocator keeps the arenas. Measured on Spark-X2.5-1.7B, anonymous RSS: path after load after release freed W4A8 (GPTQ) 7338 MB 3510 MB 3828 MB (-52%) W4A16 (tinygemm) 5023 MB 3411 MB 1612 MB The second column is the more useful finding: after the release both paths land in the same place, so W4A8's apparent 2.3 GB penalty over W4A16 was never a footprint difference -- it was the repack. W4A8 is therefore faster *and* level on memory, and fidelity is the only remaining reason to pick W4A16. Hooked to CPUWorker.load_model and compile_or_warm_up_model, since those churn separately; `VLLM_OMNI_CPU_RELEASE_MEMORY=0` disables it. CPU only: on GPU the weights are in device memory and host churn constrains nothing. `anon_rss_bytes()` ships alongside because the metric matters as much as the fix. VmRSS counts the mmap'd checkpoint's file pages and moved 1.2-2.0 GB between *identical* runs on this host; RssAnon repeats to within 0.5 MB. An earlier claim in this tree that KV sizing saved 3.1 GB came from reading VmRSS across a comparison that also changed the model file; re-measured on RssAnon, shrinking the KV cache 32x is worth 188 MB, because its pages are never touched. The analysis note is corrected accordingly. Co-Authored-By: Claude Opus 5 (1M context) --- vllm_omni/edge/memory.py | 98 ++++++++++++++++++++++++++++++++++++++++ vllm_omni/patch.py | 53 ++++++++++++++++++++++ 2 files changed, 151 insertions(+) create mode 100644 vllm_omni/edge/memory.py diff --git a/vllm_omni/edge/memory.py b/vllm_omni/edge/memory.py new file mode 100644 index 000000000..117b53ab0 --- /dev/null +++ b/vllm_omni/edge/memory.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Give freed weight-loading memory back to the OS. + +Loading a quantized checkpoint churns far more memory than the weights +occupy. The GPTQ path is the extreme case: `_process_gptq_weights_w4a8` +unpacks 4-bit codes to **int32 -- four bytes per weight** -- permutes, repacks, +and drops the intermediates. For a 1.7 B model that is gigabytes of +short-lived allocation. It is correctly freed, but a caching allocator keeps +the arenas, so it stays resident for the life of the process. + +Measured on Spark-X2.5-1.7B (Xeon 8480C, anonymous RSS, which repeats to +within 0.5 MB where VmRSS swings 1-2 GB): + + path after load after release freed + W4A8 (GPTQ) 7338 MB 3510 MB 3828 MB + W4A16 (tinygemm) 5023 MB 3411 MB 1612 MB + +Two things worth reading off that table. The release is worth **52%** of +resident memory on the path we actually want to deploy; and afterwards both +paths land in the same place, so W4A8's apparent 2.3 GB penalty over W4A16 was +never a real footprint difference -- it was the repack, held by the allocator. + +This is only useful on CPU. On GPU the weights are in device memory and the +host-side churn is not what constrains anything. +""" + +import ctypes + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# glibc's trim, then tcmalloc's equivalent: vLLM's CPU docs tell you to +# LD_PRELOAD tcmalloc, so on a correctly configured host the second one is the +# one that does the work. +_RELEASERS = ( + ("libc.so.6", "malloc_trim", (0,)), + ("libtcmalloc_minimal.so.4", "MallocExtension_ReleaseFreeMemory", ()), + ("libtcmalloc.so.4", "MallocExtension_ReleaseFreeMemory", ()), +) + + +def anon_rss_bytes() -> int | None: + """Resident anonymous memory: allocations, excluding mapped files. + + The right metric for a memory budget. ``VmRSS`` also counts the mmap'd + checkpoint's pages, which come and go with page-cache pressure and moved + 1-2 GB between identical runs on this host. + """ + try: + with open("/proc/self/status") as f: + for line in f: + if line.startswith("RssAnon:"): + return int(line.split()[1]) * 1024 + except OSError: + pass + return None + + +def release_free_memory() -> list[str]: + """Ask every allocator we can reach to return free arenas to the OS. + + Returns the names of the calls that succeeded. Safe to call repeatedly and + safe when none are available -- it simply does nothing. + """ + released: list[str] = [] + for lib, symbol, args in _RELEASERS: + try: + fn = getattr(ctypes.CDLL(lib), symbol) + except (OSError, AttributeError): + continue + try: + fn(*args) + released.append(f"{lib}:{symbol}") + except Exception as exc: # pragma: no cover - defensive + logger.debug("%s:%s failed: %s", lib, symbol, exc) + return released + + +def release_after_load(context: str = "weight load") -> None: + """Release, and say how much came back, at INFO.""" + before = anon_rss_bytes() + released = release_free_memory() + if not released: + return + after = anon_rss_bytes() + if before is None or after is None: + logger.info("Released allocator arenas after %s via %s", context, released) + return + logger.info( + "Released %.0f MiB after %s (anonymous RSS %.0f -> %.0f MiB) via %s", + (before - after) / 2**20, + context, + before / 2**20, + after / 2**20, + released, + ) diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 91ebd1b90..045c87b38 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -694,3 +694,56 @@ def _patch_cpu_explicit_kv_cache_memory() -> None: _patch_cpu_explicit_kv_cache_memory() + + +def _patch_cpu_release_memory_after_load() -> None: + """Return the weight-loading churn to the OS on CPU. + + Loading a quantized checkpoint allocates far more than the weights occupy + -- the GPTQ path unpacks 4-bit codes to int32, four bytes per weight, to + permute and repack them -- and the intermediates, though freed, stay + resident because a caching allocator keeps the arenas. On Spark-X2.5-1.7B + that is 3828 MiB held after a W4A8 load and 1612 MiB after W4A16; released, + both land at ~3.5 GiB, so W4A8's apparent memory penalty over W4A16 was + entirely this and not a real difference in footprint. + + Hooked in two places because they churn separately: after weights are + loaded and processed, and after the compile/warm-up pass. + """ + from vllm.platforms import current_platform + + if not current_platform.is_cpu(): + return + if os.environ.get("VLLM_OMNI_CPU_RELEASE_MEMORY", "1") != "1": + return + + from vllm.v1.worker.cpu_worker import CPUWorker + + if getattr(CPUWorker, "_omni_release_memory", False): + return + + from vllm_omni.edge.memory import release_after_load + + base_load = CPUWorker.load_model + base_warm = CPUWorker.compile_or_warm_up_model + + def load_model(self, *args, **kwargs): + out = base_load(self, *args, **kwargs) + release_after_load("weight load") + return out + + def compile_or_warm_up_model(self, *args, **kwargs): + out = base_warm(self, *args, **kwargs) + release_after_load("warm-up") + return out + + CPUWorker.load_model = load_model + CPUWorker.compile_or_warm_up_model = compile_or_warm_up_model + CPUWorker._omni_release_memory = True + _PATCH_LOGGER.info( + "CPU weight-load memory release installed (set " + "VLLM_OMNI_CPU_RELEASE_MEMORY=0 to disable)." + ) + + +_patch_cpu_release_memory_after_load()