Release the weight-checker snapshot once compare passes - #37284
yueming-yuan wants to merge 3 commits into
Conversation
| snapshot_tensors = {} | ||
| for (name, param), offset in zip(named_params, offsets): | ||
| view = ( | ||
| self._snapshot_arena[offset : offset + param.nbytes] |
There was a problem hiding this comment.
nit: make an abstractin similar to the BumpAllocator's API to avoid hand-compute offsets etc and a bit clearer
| for _, param in named_params: | ||
| offsets.append(total) | ||
| total += (param.nbytes + align - 1) // align * align | ||
| self._snapshot_arena = alloc_mmap((max(total, align),), torch.uint8) |
There was a problem hiding this comment.
[P2] Publish the mmap arena only after the snapshot is fully built
self._snapshot_arena is assigned before any of the views are constructed or copied. If a later .view() or copy_() raises, _snapshot_tensors remains unset, but the fully prefaulted, model-sized mmap remains reachable from the long-lived checker. A retry must allocate another arena before replacing the old one, which can OOM. Please build the arena in a local variable and publish _snapshot_arena and _snapshot_tensors together only after every copy succeeds.
Found by Codex.
Bump-allocate the per-tensor views through a small arena allocator and assign _snapshot_arena/_snapshot_tensors together at the end, so a copy that raises mid-snapshot leaves no model-sized mmap held by the checker.
|
/tag-and-rerun-ci |
…r-mmap-snapshot-main # Conflicts: # python/sglang/srt/utils/weight_checker.py
Pull request was closed
4 participants
The WeightChecker snapshot is a host copy of every weight — tens of GB for a large model — and it stayed resident for the whole run even though compare is its only consumer. Two changes:
alloc_mmap) instead of per-tensor.cpu()copies. The per-tensor copies are mostly below glibc's mmap threshold, so freeing them never shrinks RSS (on Grace/64K pages evenmalloc_trimreclaims nothing); one mmap region is returned to the OS deterministically by munmap.Verified on GB300 (DeepSeek-V4, 8 nodes): snapshot RSS returns to the OS after the first compare; unit tests cover the arena packing and release paths.
Port of #36561 (merged into
sglang-miles) to main.CI States
Latest PR Test (Base): ❌ Run #34829365272
Latest PR Test (Extra): ❌ Run #34829365024
Latest PR Test (AMD ROCm 10): ❌ Run #34829365286