[Misc] Add standalone safetensors retention diagnostic

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu committed 2026-09-06 16:57:46 +08:00
1 parent 039808e0d9
commit 86166fd988
3 files changed
+171

No files matched your search

+49
View File
@@ -0,0 +1,49 @@
# Host weight dependency memory diagnostic
`safetensors_retention.py` isolates repeated CPU `get_tensor()` calls from
access to one cached tensor view. It imports neither vLLM nor HWR and uses a
synthetic 16-element float32 payload. HWR itself calls `get_tensor()` when
acquiring a lease, then exposes cached views; this probe does not measure HWR
requests or repeated lease acquisition.
Run one mode per fresh process in the same Python environment and CPU affinity:
```bash
git rev-parse HEAD > revision.txt
git status --short > working-tree.txt
git diff > working-tree.patch
# Select an allowed CPU from your process affinity; 0 is only an example.
CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
benchmarks/host_weight_runtime/safetensors_retention.py \
--mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json
```
For a comparison, run each mode twice in fresh processes with the same controls:
```bash
for repetition in 1 2; do
for mode in get_tensor reuse; do
CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
benchmarks/host_weight_runtime/safetensors_retention.py \
--mode "$mode" > "$mode-$repetition.json"
done
done
```
The diagnostic sets Torch to one CPU thread, warms up before the baseline,
and reports preparation separately from timed loop work. Checkpoints run GC
and report Linux private/anonymous memory, RSS, and open descriptors. It checks
that transient tensor objects and the final cached view are released. Linux
`/proc/self/smaps_rollup` must be readable; invalid arguments and missing proc
support fail explicitly. Keep raw JSON and the repository snapshot with results.
Compare the final loop sample with the warmed baseline, then inspect the closed
sample. Private-memory growth after Python tensor objects disappear is evidence
of retention in the dependency/allocator path, not proof of a leak's root cause.
RSS can include file-backed pages; it is not a private-memory metric. Samples
also include small diagnostic bookkeeping allocations. Do not assert a portable
memory threshold or turn this synthetic loop into a per-request service estimate.
Do not drop shared caches, trim allocators, or add serving-time GC to make the
numbers look smaller. A dependency replacement or version change requires its
own reproducible comparison and correctness/lifecycle validation.
@@ -0,0 +1,117 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Linux CPU diagnostic: repeated safetensors materialization versus view reuse."""
from __future__ import annotations
import argparse
import gc
import json
import os
import platform
import tempfile
import time
import weakref
from pathlib import Path
def positive_int(value: str) -> int:
parsed = int(value)
if parsed <= 0:
raise argparse.ArgumentTypeError("must be positive")
return parsed
def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object]:
gc.collect()
counters = {}
for line in Path("/proc/self/smaps_rollup").read_text().splitlines():
fields = line.split()
if len(fields) == 3 and fields[2] == "kB":
counters[fields[0].removesuffix(":")] = int(fields[1]) * 1024
return {
"phase": phase,
"iterations": iterations,
"loop_seconds": loop_seconds,
"rss_bytes": counters["Rss"],
"private_bytes": counters["Private_Clean"] + counters["Private_Dirty"],
"anonymous_bytes": counters["Anonymous"],
"fd_count": len(tuple(Path("/proc/self/fd").iterdir())),
}
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--mode", choices=("get_tensor", "reuse"), required=True)
parser.add_argument("--iterations", type=positive_int, default=1_000_000)
parser.add_argument("--sample-every", type=positive_int, default=100_000)
parser.add_argument("--warmup", type=positive_int, default=1000)
args = parser.parse_args()
if not Path("/proc/self/smaps_rollup").is_file():
parser.error("requires Linux /proc/self/smaps_rollup")
preparation_start = time.monotonic()
import safetensors
import torch
from safetensors import safe_open
from safetensors.torch import save_file
torch.set_num_threads(1)
torch.set_num_interop_threads(1)
samples = [sample("before_file", 0, 0.0)]
with tempfile.TemporaryDirectory(prefix="safetensors-retention-") as directory:
path = Path(directory) / "synthetic.safetensors"
save_file({"weight": torch.arange(16, dtype=torch.float32)}, str(path))
with safe_open(path, framework="pt", device="cpu") as reader:
cached = reader.get_tensor("weight")
assert torch.equal(cached, torch.arange(16, dtype=torch.float32))
for _ in range(args.warmup):
tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
del tensor
samples.append(sample("baseline", 0, 0.0))
preparation_seconds = time.monotonic() - preparation_start
completed = 0
loop_seconds = 0.0
while completed < args.iterations:
count = min(args.sample_every, args.iterations - completed)
start = time.monotonic()
for _ in range(count):
tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
last_tensor = weakref.ref(tensor)
del tensor
loop_seconds += time.monotonic() - start
completed += count
snapshot = sample("loop", completed, loop_seconds)
snapshot["last_tensor_alive"] = last_tensor() is not None
if args.mode == "get_tensor":
assert last_tensor() is None, "diagnostic retained a transient tensor"
samples.append(snapshot)
del cached
del reader
samples.append(sample("closed", completed, loop_seconds))
assert last_tensor() is None, "last tensor survived reader/view teardown"
print(
json.dumps(
{
"schema_version": 1,
"arguments": vars(args),
"python": platform.python_version(),
"torch": torch.__version__,
"safetensors": safetensors.__version__,
"kernel": platform.release(),
"hostname": platform.node(),
"pid": os.getpid(),
"cpu_affinity": sorted(os.sched_getaffinity(0)),
"torch_threads": torch.get_num_threads(),
"pythonmalloc": os.environ.get("PYTHONMALLOC", "default"),
"preparation_seconds": preparation_seconds,
"samples": samples,
},
indent=2,
)
)
if __name__ == "__main__":
main()
@@ -42,6 +42,11 @@ V1 does not include:
## Motivation and use cases
For investigating CPU memory retained by dependency tensor materialization,
see the [standalone safetensors diagnostic](../../../benchmarks/host_weight_runtime/README.md).
It distinguishes repeated dependency calls from reuse of cached views and does
not establish per-request HWR leakage.
Model loading can create the same final host representation repeatedly. This
is especially expensive when loading performs checkpoint decoding, tensor
renaming, TP slicing, quantization, packing, or scale construction before GPU