Repository navigation
llama-mmap: populate the non-lazy ranges when MAP_POPULATE is skipped for lazy tensors - #266
danielhanchen wants to merge 8 commits into
Conversation
…d for lazy tensors
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ab8d327ac6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // MAP_POPULATE is not used when the file has lazy ranges, so populate the page tables of the | ||
| // other ranges explicitly: otherwise the first access at run time (e.g. uploading the experts | ||
| // of a layer kept in host memory) takes a page fault for every page |
There was a problem hiding this comment.
This three-line comment hard-wraps a single sentence and exceeds the repository's usual 1-2-line limit for code comments. Condense it to the non-obvious invariant without splitting the sentence across lines.
AGENTS.md reference: AGENTS.md:L77-L84
Useful? React with 👍 / 👎.
|
@codex review |
|
Codex Review: Didn't find any major issues. Hooray! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…AM budget The previous commit populated the whole complement of the lazy ranges from the llama_mmap constructor: every non-lazy byte of the file, synchronously, with no bound. That includes tensors that are never created (TENSOR_SKIP, e.g. MTP blocks without a drafter, or another model's tensors in a shared file). Move the population to llama_model_loader::init_mappings, where the tensor list is final, and populate only the page-aligned, merged ranges of the non-lazy tensors that were created for this file. A file is populated only if its ranges fit in MemAvailable minus max(2 GiB, MemTotal/16), counted across files; otherwise nothing is populated and the pages fault in on demand, as before the previous commit. LLAMA_MMAP_POPULATE=0 turns it off. The populated (or skipped) byte count is logged at info level. Windows, macOS, --no-mmap, NUMA and files without lazy ranges are unchanged.
A file without lazy ranges is still mapped with MAP_POPULATE and populated in full, so with a split GGUF the explicit populate of the lazy-bearing split could spend the whole budget on top of the ordinary splits (a 27 MiB budget populated 31.2 MiB with the lazy-bearing split first, last or in between). Subtract the size of every such file from the budget when it is first read: the later ones are mapped after it, and the earlier ones do not lower MemAvailable, which counts their page cache as available. A populate that stops on a failed madvise may have populated part of the failed run, so it now ends the budget instead of debiting only the completed runs. Note in the comment that LLAMA_MMAP_POPULATE=0 only turns off this populate.
MemAvailable is host wide, so in a container whose memory.max or memory.high is far below the host RAM the populate could go past the container's allowance and push the load into reclaim. Lower the budget to what the cgroup and its ancestors leave (v2 memory.max / memory.high less memory.current, v1 memory.limit_in_bytes less memory.usage_in_bytes), less the same max(2 GiB, limit/16) reserve. Files that are missing or cannot be parsed leave the host budget as it is. Add test-mmap-populate-budget for the cgroup layouts.
…hs under a subtree mount The budget was sampled when the first file with lazy ranges was mapped, after the earlier MAP_POPULATE files were already charged to memory.current, and their sizes were then reserved again. Sample it before any file is mapped. A cgroup mount whose root is a subtree (a container without a cgroup namespace) shows the host path in /proc/self/cgroup; make it relative to the mount point via /proc/self/mountinfo so a tighter child cgroup is read. Adds page_runs tests.
mountinfo escapes space, tab, newline and backslash in the mount root and mount point, so a subtree mount under such a name (a systemd \x2d unit) was not matched and the tighter child cgroup was missed.
…taken by ggml-org#30974, test labels from ggml-org#31036)
llama_mmapusesMAP_POPULATEonly when the file has no lazy-read ranges. In the default load mode every tensor over 4 GiB is lazy, so a single-file GGUF with one large tensor gets no prepopulation at all. That includes the per-layer token embedding in Qwen3.8-Flash-Next (26.8 GiB).As a result, the first access to every page of that file at run time takes a page fault. For example, uploading the experts of a layer kept in host memory pays one fault per page, and so does the CPU matmul over them during generation. A split GGUF hides this: only the split holding the lazy tensor loses
MAP_POPULATE.This change: when
MAP_POPULATEis skipped because of lazy ranges, the non-lazy ranges are populated explicitly withmadvise(MADV_POPULATE_READ)(Linux 5.14+, guarded byMADV_POPULATE_READ). If the madvise fails, it is logged at debug level and loading continues as before.Results
Same model in both arms: Qwen3.8-Flash-Next, a single 72.5 GB GGUF. RTX PRO 6000 Blackwell with VRAM capped at 32 GiB, so experts of the last layers live in host memory. llama-bench, A/A spread at most 1.0%.
For context, ggml-org#28055 proposes dropping
MAP_POPULATEaltogether. These numbers suggest prepopulation is worth keeping for offloaded weights.Not covered
MAP_POPULATEwould have touched for a file without lazy tensors.MAP_POPULATE.Base is upstream
8bb3cf0d2(base/upstream-8bb3cf0d2). Like every pin inscripts/unsloth/pr-set.json, this is merged onto the aged upstreamb####tag for the prebuilt and dropped once upstream has the equivalent.