Compare commits

...
Author SHA1 Message Date
Hongsheng Liu 380127049a [Doc] Fix HWR diagnostic link in strict docs builds
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-07 21:39:37 +08:00
Hongsheng Liu 743953b821 [Misc] Make retention probe tensor size configurable
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-06 21:53:50 +08:00
Hongsheng Liu 86166fd988 [Misc] Add standalone safetensors retention diagnostic
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-06 16:57:46 +08:00
3 changed files with 182 additions and 0 deletions

No files matched your search

+57
View File
@@ -0,0 +1,57 @@
# 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 float32 payload controlled by `--tensor-elements` (default: 64).
Tensor size can affect retention: a negative result for 16 elements does not
rule out growth for 64 elements. The selected size is recorded in the JSON
arguments and used for both file creation and the correctness check.
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 size comparison, run 16 and 64 elements twice each in fresh processes,
with cached reuse as a control at both sizes:
```bash
for repetition in 1 2; do
for elements in 16 64; 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" --tensor-elements "$elements" \
> "$mode-$elements-$repetition.json"
done
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,120 @@
# 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(
"--tensor-elements", type=positive_int, default=64, help="Number of float32 elements in the synthetic tensor"
)
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(args.tensor_elements, 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(args.tensor_elements, 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](https://github.com/vllm-project/vllm-omni/blob/main/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