Pull requests / #3378
#3378 [Other] tools: read the FP4-Local checkpoint with cached sequential reads, not a process-lifetime np.memmap
open · @jerry78424 · 0 commentaires · Sur GitHub
Server & APIMulti-GPUAMD / HIPNVIDIA / CUDAModels & quantsWindowsLinux
Description
Fixes #3377. Related: #3227 / #3228 (a different failure of the same step; no hunk overlap — that PR touches the overlay gens around `:191-211`, this one the `Checkpoint` class at `:48-91`). ## The problem and the result `prepare` fails mid-pack on Windows: the pack python dies at 50–67 % of the experts write, exit 1, no traceback, while its own committed memory is ~3.5 GB. `Checkpoint.raw()` maps every shard with `np.memmap(..., mode="r")` for the process's life; the pack pass reads 64.36 GB of distinct expert bytes, and on Windows every faulted page is charged to the process working set, which `ullAvailPhys` (= psutil `available` = `os.freemem` = our own `ram_used`, `serve/telemetry.py:624-625`) counts as used until the OS trims — the step presents a ~64 GB RAM demand for a few-GB real one (issue: +41 GB in five minutes, run dead at ~11 GB free). With this change the same reads return hash-identical bytes, run ~1.2x faster warm, and the working set stays flat — a full 68 GB pack run to a scratch dir with the serving model loaded (the scenario that failed): exit 0, byte-identical `experts.bin`, working set 3.3–5.3 GB bounded (externally sampled), free RAM never below ~53 GB. ## The approach `Checkpoint.raw()` reads the tensor's byte range through a cached file handle with `seek` + `readinto` into a fresh buffer; the kernel's pages then live in the system file cache (standby: counted as available) instead of the process working set. - The returned array keeps `array()`'s `.view().reshape()` contract (buffers >=8-byte aligned; U8/U16/U32/F32 all used); it is writable where `mode="r"` was read-only — no consumer writes (all produce fresh outputs; a mutator would previously have raised through numpy or faulted through a torch view) — and `raw()` shares one handle per shard, so it is not thread-safe, unlike the memmap it replaces; all call sites are single-threaded. - Unmapping per shard would bound the working set to ~one shard (~5 GB) but re-fault ~64 GB per pass (13–14 of the 15 expert shards are open per layer anyway); sequential `readinto` bounds it to ~0. Within a shard the order walks three ascending streams (exceptions, weights, CSF-fixed) expert-ascending. np.memmap has no Windows lever (no `FILE_FLAG_SEQUENTIAL_SCAN`/madvise equivalent). - No API or call-site changes; at most one open read handle per shard, same count as the replaced memmaps. Residual memmaps elsewhere are bounded and untouched: `iq_pack.py:161,218,324` (~6 GB dense.gguf) and `LocalTensor.mmap_bytes` in convert (~4 GB). ## Validation - **passed**: `python -m pytest tools/test_nvfp4_local_pack.py`: 13 — the 11 existing (they never construct `Checkpoint`) plus 2 added here: a synthetic two-shard checkpoint (byte-exact reads at a tensor's range, `array()` dtype/shape views, repeat reads, truncated shard -> `EOFError` instead of a silent short view). - **passed**: A/B on the live checkpoint, the pack's read order (4.03 GB): old path +3.75 GB working set (93 % of bytes touched; charges even warm) in 2.7–4.1 s, this path +0.00 GB in 2.3 s, both hash-identical. Extrapolation: the failed run had faulted 50–67 % of 64.36 GB = 32–43 GB, bracketing its observed +41 GB. - **passed**: full pack to a scratch `--out` with the serving model loaded: exit 0, `experts.bin` size identical (67,948,118,016 bytes), SHA-256 match `1e124ee94b7c921a7c895d2898bfa4ddc2d8111f3fd3e8b94207759c363291a5`; externally sampled working set 3.3–5.3 GB bounded, free 53.3–60.0 GB throughout (in-run 15 s sampler; pre-patch signature: +41 GB in 5 min). - **failed**: none. - **not run**: `ple`/`mtp`/`verify` individually (same class, different call sites, covered by the new class tests); Linux/AMD (OS-neutral file I/O; the inflation mechanism is Windows-specific). ## Limits - `readinto` returns short only at EOF (the buffering layer absorbs partial raw reads); the added `EOFError` names a truncated shard instead of the old silent short memmap view — a behavior change only for corrupt checkpoints. **Environment**: tested at e88ee39d (v0.2.0-beta-12), Windows 11 26200, Ryzen 7 9800X3D, 187 GB RAM, RTX 5090 (driver 617.42), NVMe, model `nvidia/Qwen3.8-Flash-Next-FP4-Local`.
Sur le site
Liens install, modèles, releases.