Repository navigation
fix(memory): honor cgroup memory limits when sizing expert banks - #334
Avicennasis wants to merge 6 commits into
Conversation
|
Thanks for this — it lands on something that bites here for a reason you did not have in front of Setup. RTX 4090 24 GB, 61 GB RAM, Ubuntu 26.04, cgroup v2 only (no v1 hierarchy mounted, so The gap, measuredEverything on this box is launched through a small wrapper that runs the server in a transient Before this PR, It does not cost the fast path hereThe obvious worry with a tighter budget is that the parallel loader stops being admitted and every Then served an agent workload for ~55 minutes without incident — 994 prefill batches, 7.3 M prefill Tests: One thing I checked that turns out to be fine, in case it saves the next readerThe It cannot: v1 and v2 use different sentinels. cgroup v1 Two observations, neither a defect
Unrelated to the code: the branch is five commits behind
Happy to test the v1 path if someone with a v1 or hybrid host wants a second pair of eyes; I only Written with AI assistance; every number above was measured on my hardware (RTX 4090, 61 GB |
de9df4a to
056b258
Compare
|
Thanks for the systemd-scope run — that is exactly the non-container shape the write-up could not cover, and 55.9 → 46.0 GiB is the same bug from the other direction. Rebased onto current The only edit beyond that is in the two admission tests, which stub the builder to observe which loader was chosen: #418 renamed it and split it into Re-ran on a CPU-only Ubuntu 24.04 x86_64 box (16 cores, 125 GiB, cgroup v2, Docker 29.1.3; Python 3.12, torch 2.11.0+cu130, transformers 5.17.0, pytest 9.1.1, run with
Nothing in the change touches a kernel, which is why a CPU-only box is an honest environment for it; the loader path it gates is host-RAM admission. Both of your observations are right. I have left them out of this PR so the reviewed core stays as it was, and will fold them into a follow-up (dedupe the ancestor-walk guard; cache the bound for the bench sampler) if you want them. The |
…lashML-org#334) Upstream FlashML-org#334 by Avicennasis, taken unmodified. Adds `freetoken/memory.py::effective_memory_available()` = min(host MemAvailable, tightest finite cgroup limit) and routes the expert-bank loader (`_host_ram_fits_parallel`), `benchbw` and the loader benchmark through it. Why this box wants it: everything here is launched through `ft-capped`, which runs the server in a systemd scope with MemoryMax=46G. Before this, `_host_ram_fits_parallel` sized the parallel expert reader's one-extra-shard transient from /proc/meminfo MemAvailable alone -- the HOST's figure -- so the fast path was admitted against ~10 GiB the cgroup would never allow, and the load died to the OOM killer with no diagnostic. Measured here after applying, same interpreter, same call: outside any scope 55.9 GiB inside MemoryMax=46G (ft-capped's shape) 46.0 GiB <- cgroup memory.max So the admission decision now sees the cap. tests/test_memory.py 43 passed, tests/moe/test_expert_bank_memory_admission.py 4 passed (both new upstream). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…org#418), without FlashML-org#337 FlashML-org#418 replaced the per-family quantization branches with QuantConfig / QuantScheme / QuantMethod, and moved expert-bank ownership into the MoE kernel. 16 files conflicted; they were not 16 equal problems. Five were mechanical: * engine/cache_budget.py -- take FlashML-org#418's max_slots parameter (it replaces the hardcoded 992 marlin cap) and keep our +1 dummy-page reservation on the KV floor, with its test. * engine/config.py -- keep the fields, take the renamed-flag wording. * kernel/triton/fp8_pertensor_linear.py -- FlashML-org#418 removed the BaseOP layer, so only the round_e4m3 import is still needed; the e4m3 rounding fix itself is untouched. * models/qwen3_5_moe/model.py -- FlashML-org#418's ParallelLMHead takes a quant_config and selects the NVFP4 W4A16 kernel itself, which subsumes our Nvfp4LMHead special case (and the ~1 GB bf16 dequant it existed to avoid). Take theirs, keep the vision tower line. The other eleven were all FlashML-org#337's NVMe disk tier, whose plumbing point no longer exists: FlashML-org#418 deleted the per-model load_nvfp4_expert_sources(..., disk_tier=) in favour of a one-line nvfp4_expert_spec() plus expert_pieces, so the tier would have to be re-attached to the kernel-owned layout rather than merged. MT chose to drop it and rebuild it fresh, so it is removed whole -- moe/disk_tier.py and its test, the three CLI flags, the three EngineConfig fields, and every call site in layers/moe.py, moe/offload_cache.py, moe/expert_banks.py and engine/engine.py. It is recoverable from 9834eb3 / e4b565c. Three files needed hunk-level care because other work shares them, and a wholesale --theirs would have dropped it silently: * moe/expert_banks.py -- take main, then re-apply FlashML-org#334's effective_memory_available() (cgroup-aware host RAM) on top. * models/qwen3_5_moe/weight.py -- the vision rename hooks auto-merged outside the conflict markers; resolve per hunk so they survive. * engine/engine.py -- take main for the two bank-binding hunks, keep ours for the --moe-collect-stats routing histogram (main still has OffloadMoeCache.collect_decode_freq). layers/moe.py and moe/offload_cache.py turned out to differ from main only by the disk tier, so both are now identical to it. The dead make_offload_moe_cache hook went with the disk-tier branch that read it; nothing in the tree implements it. MT's serving command still parses: --moe-backend offload folds into moe_strategy='offload' and every other flag is unchanged, with one new deprecation warning at start-up. Assisted-by: Claude Opus 5 Tests: the merged tree ran 10 failed / 1904 passed / 92 skipped; the pre-merge try/all ran 8 failed / 1905 passed / 60 skipped, and the two failure sets differ by exactly the two above, both since fixed and verified in isolation (4 passed, and the admission test does fail when effective_memory_available's result is discarded). The confirmation run of the whole suite after the fix was interrupted at 50%, so "8 failed after the fix" is derived from those two facts, not measured end to end.
Port of upstream PR FlashML-org#334: the admission check read host-wide MemAvailable, so inside a --memory-limited container it could over-admit expert banks and OOM. New torch-free python/freetoken/memory.py resolver combines MemAvailable with the tightest finite cgroup v1/v2 budget (fail-closed); expert_banks, benchbw and the load bench share it.
* fix(engine): bound host memory governor by cgroup limits The expert-tier governor derived pinned/pager budgets from /proc/meminfo only, so under a container or systemd MemoryMax it could budget past what the cgroup allows. Take min(host, cgroup headroom) and log which bound applied. No change when memory.max is "max" (or no memory controller). Ports the stdlib resolver and tests from FlashML-org/FreeToken#334 (Avicennasis, head de9df4a7): cgroup v2 ancestor walk, v1 fallback, fake filesystem tests. Adds cgroup_memory_headroom() so the governor can report the bound separately. Loader/bench call sites are not wired here. Refs #22 Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> * fix(engine): use cgroup limit for governor total, not headroom Total is min(MemTotal, tightest memory.max); available stays min(MemAvailable, limit - current), so the default reserve no longer shrinks as the process allocates. Adds cgroup_memory_bounds() to freetoken.memory returning (limit, headroom). Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 5.5 <noreply@anthropic.com>
…t's MemAvailable Inside a container /proc/meminfo describes the host, so the parallel expert reader was admitted against RAM the cgroup will not grant and the load was OOM-killed without a diagnostic. available_host_memory() clamps MemAvailable to the tightest cgroup v2/v1 limit along the ancestry and is now used by the expert-bank loader, ft bench bw and the load bench. Alternative to upstream PR FlashML-org#334: about a fifth of its size, and it counts reclaimable page cache (inactive_file) as free like MemAvailable does, instead of charging memory.current in full, which reads a container that just read its checkpoint as having no room left. Co-authored-by: Léon Avic Simmons <Avicennasis@gmail.com> Assisted-by: Claude Opus 5.5
Logs the low-RAM loading batch pulled forward: FlashML-org#334, FlashML-org#233, FlashML-org#144 and FlashML-org#56. Assisted-by: Claude Opus 5.5
056b258 to
f3c9982
Compare
|
Rebased onto current Re-ran on a CPU-only Ubuntu 24.04 x86_64 box (AMD Ryzen 7 5800X, 16 threads, 125 GiB, no GPU; Python 3.12.3, torch 2.14.1+cpu, triton 3.8.0, transformers 5.16.1, flashlib 0.4.0, numpy 2.5.2, pytest 9.1.1;
Nothing in the change touches a kernel or a model family; the clamp is read once, where the parallel reader and the two benches size their host-RAM work. |
The expert-bank loader, the CPU bandwidth bench and the loader benchmark all sized host-RAM work from /proc/meminfo MemAvailable alone. Inside a container started with --memory=N that figure is the host's, not the cgroup's: the parallel expert reader was admitted against tens of GiB of "available" RAM while the cgroup could hold a fraction of it, and the process was OOM-killed mid-load. benchbw.py carried its own cgroup-v2-only parser; the loader had none. Rule: effective available memory = min(host MemAvailable, remaining finite cgroup budget). The cgroup budget is the tightest finite limit-minus-usage along the process's cgroup ancestry (v2 memory.max / memory.current with the `max` sentinel; v1 memory.limit_in_bytes / memory.usage_in_bytes with the unlimited sentinel as a fallback), resolved from /proc/self/cgroup with bounded reads and a bounded ancestor walk. Absent procfs or no visible memory controller yields None (unknown), so non-Linux hosts keep their previous best-effort behaviour and are never told they have 0 bytes; a present but malformed or unreadable control file fails closed to 0. The probe runs before any bank or buffer allocation: admission is the only decision point, because once the parallel reader's whole-shard anonymous buffers exist there is nothing reclaimable left to give back. freetoken/memory.py is stdlib-only (no torch import) so it can be loaded by path inside a slim container for diagnosis. Every production MemAvailable reader now goes through it: moe/expert_banks ._host_ram_fits_parallel, moe/benchbw._available_ram_bytes (the private _cgroup_mem_headroom parser is removed so there is one cgroup parser; the bench's clamping and synthetic-bank sizing are otherwise unchanged), and benchmarks/bench_load_weight_generic.MemSampler. Tests: tests/test_memory.py drives the resolver against fake /proc and cgroup trees under tmp_path -- unlimited, finite, nested ancestor-min, namespaced root, usage above limit, malformed, oversized, permission denied, v1 fallback and sentinel, absent files, and the hybrid `1:net_cls:/` + `0::/...` membership layout systemd and Docker both produce. tests/moe/test_expert_bank_memory_admission.py checks that the auto-admission path drops to the serial build when the fake cgroup budget (299 B) cannot hold the banks plus one shard transient (300 B) although MemAvailable is large, that an explicit --expert-load override bypasses the probe, and that the bench receives a finite budget unchanged while keeping its legacy 8 GiB fallback only for an unknown (None) probe. Validated on CPU-only Linux (Ubuntu 24.04, Docker 29.1, cgroup v2; no GPU is involved in this path): 47 tests pass under python and python -O; the full suite has the same two pre-existing failures as origin/main (tests/models/test_muse_glimmer.py) and no new ones; python:3.12-slim with --memory=2g / --memory=6g / no limit reports effective = memory.max - memory.current (2136084480 / 6432010240) and host MemAvailable respectively. Mount-namespace hardening of the cgroup lookup is deliberately out of scope for this change.
The two loader tests replaced expert_banks._legacy_expert_banks and passed dtype= to load_expert_banks. Both went away with the q4_0 provider (FlashML-org#690), so on current main the tests stop at the monkeypatch (AttributeError) before any admission code runs. They now stand in for _method_expert_banks, the builder a raw checkpoint reaches, pass a bound-method stand-in (kind NONE), and read the chosen mode by parameter name from the real builder's signature instead of by position. The download_hf_weight patch is dropped: the shard list now comes from the loader's iter_weight_files, which resolves a local directory itself. The auto-admission test is two-sided: 299 bytes of cgroup room against the 300 the banks and one shard need gives the serial build, 301 keeps the parallel reader. With only the serial half, a loader that never proposed the parallel reader would pass without reaching the probe.
…resolver test_host_ram_fits_parallel_sizes_subdirectory_shards (FlashML-org#819) fakes /proc/meminfo by replacing builtins.open. _host_ram_fits_parallel no longer opens that file itself: it asks effective_memory_available(), which reads through pathlib, so the fake is never consulted, the probe sees the machine's real figure and the test fails on any host with more than 3 MiB available. The test's subject is which shards get sized, not how the figure is read, so it now sets the same 512 kB at the name the probe calls. The fixture and the assertion are unchanged.
The oversized fixture was 65 KiB of "x", which the line parser rejects on its own, so the case passed with the 64 KiB bound removed. Use 70,000 bytes of distinct well-formed counters instead.
The file LRU lists include pages that are dirty or under writeback. They cannot be dropped until written, and a cgroup at its limit is OOM-killed before that finishes. Take file_dirty + file_writeback (v1: total_dirty + total_writeback) back out of the file-cache credit, clamped at zero.
f3c9982 to
323fc93
Compare
|
Restacked onto current The restack. The commit rebased without a conflict;
Eight single-line mutations of the admission path (guard removed, cgroup ignored, shard transient dropped, parallel reader never proposed, probe never fits, explicit override probed, figure not taken from the resolver, root glob instead of the shard index) each turn exactly the test written for it red. None of the 109 commits changed the arithmetic this PR feeds. File cache counted as room (the last three commits). In September MT-z checked the obvious worry with a tighter budget, that the parallel loader stops being admitted, and did not hit it. There is one common case where it does happen: cgroup usage includes page cache, so once a checkpoint has been read inside the container or scope, the next start sees its size as spent. with the hierarchical One live read per column, inside
Gate. CPU-only Ubuntu 24.04 x86_64 box (AMD Ryzen 7 5800X, 16 threads, 121 GiB, no GPU; Python 3.12.3, torch 2.14.1+cpu, triton 3.8.0, transformers 5.18.0, flashlib 0.4.0, numpy 2.5.4, safetensors 0.8.0, pytest 9.2.0;
Not exercised here: CUDA, a real checkpoint or model load, a Docker container on this head, a real cgroup v1 host (fixtures only), a kernel without MGLRU, hosted CI. The reclaim and dirty-page measurements in the description were taken before this rebase on the same |
Fixes #333.
Problem
The expert-bank loader (
moe/expert_banks.py::_host_ram_fits_parallel), the CPU bandwidth bench (moe/benchbw.py) and the loader benchmark (benchmarks/bench_load_weight_generic.py) size host-RAM work from/proc/meminfoMemAvailablealone. Inside a container started with--memory=N(Docker, Kubernetes) or a systemd scope withMemoryMax=Nthat figure is the host's, so the parallel expert reader — whose transient is one extra whole shard — is admitted against tens of GiB the cgroup will not allow, and the process is OOM-killed mid-load with no diagnostic.benchbw.pyalready carried a private cgroup-v2-only clamp, so the bench and the loader disagreed about how much RAM a container has.Clamping to
limit - usagealone has the opposite fault. Cgroup usage includes page cache, so once a checkpoint has been read inside the container the next start sees its size as spent and drops to the serial loader, although the kernel gives that cache back on demand.MemAvailable, the other half of the minimum, already counts page cache.Change
One stdlib-only module,
freetoken/memory.py, witheffective_memory_available():memory.max/memory.current(max= unlimited), walking/proc/self/cgroupancestors with bounded reads; cgroup v1memory.limit_in_bytes/memory.usage_in_bytesfallback with the unlimited sentinel.memory.statis read with a 64 KiB bound; cgroup v1 uses the hierarchicaltotal_*counters.file,cacheandshmemare not used:fileincludes tmpfs, which sits on the anon lists. Dirty and writeback pages are taken back out, because they cannot be dropped until they are written. Absent, unreadable, oversized, non-ASCII or malformed statistics (bad line, duplicate key, signed or overflowing counter) disable the whole subtraction.None(unknown): non-Linux hosts keep today's behaviour and are never told they have 0 bytes. Present-but-malformed or unreadable limit or usage files fail closed to 0.The loader, the bandwidth bench and the loader benchmark now go through it;
benchbw's private_cgroup_mem_headroomparser is removed (one parser). The bench's clamping and synthetic-bank sizing are otherwise unchanged, and so is the loader's rule (available memory must exceed all shards plus the largest one). Mount-namespace hardening of the lookup is deliberately out of scope.scripts/gguf_check/check.py(#756) has its ownMemAvailablewait loop and is not touched.The file-cache rule is adapted from JustVugg/colibri #2397, merge commit
79046d10ab021908249d9f96a71c8788fecd20f2. Both projects use Apache-2.0. No donor code was copied. Unlike the donor, dirty and writeback pages are not counted; the measurements below are why.Tests
tests/test_memory.py(129 cases): fake/proc+ cgroup trees undertmp_path— unlimited, finite, nested ancestor-min, namespaced root, usage above limit, malformed, oversized, permission denied, v1 fallback and sentinel, absent files, and the hybrid1:net_cls:/+0::/...membership layout systemd and Docker both produce. For the file cache, in v2 and v1: a 24 GiB limit with 20 GiB charged and 14 GiB of clean file cache yields 18 GiB instead of 4;file,cacheandshmemare not credited; the other hierarchy's counter names are ignored; dirty and writeback pages are not room; the credit is applied at each ancestor and still clamped by the host; every malformed-statistics class subtracts nothing and never masks a malformed limit or usage file.tests/moe/test_expert_bank_memory_admission.py(5 cases): the auto-admission path drops to the serial build when the fake cgroup budget (299 B) cannot hold banks + one shard transient (300 B) althoughMemAvailableis large, and keeps the parallel reader at 301 B; an explicit--expert-loadoverride bypasses the probe; the bench receives a finite budget unchanged and keeps its 8 GiB fallback only for an unknown probe.tests/moe/test_offload.py::test_host_ram_fits_parallel_sizes_subdirectory_shards(fix(models): discover safetensors shards from the index #819) faked/proc/meminfoby replacingbuiltins.open, which the shared resolver does not go through. It now sets the same 512 kB on the resolver; its fixture and assertion are unchanged.These tests patch the new resolver, so they cannot run against
mainas written; the behavioural difference is the reproduction in the linked issue and the live read below —main's_host_ram_fits_paralleladmits a load that needs 4608 MiB inside a 4096 MiB scope, this branch's admission drops to the serial build.Validation
Environment: Ubuntu 24.04 (Linux 7.0.0, x86_64, AMD Ryzen 7 5800X, 16 threads, 121 GiB), cgroup v2, Python 3.12.3, torch 2.14.1+cpu, transformers 5.18.0, triton 3.8.0, flashlib 0.4.0, numpy 2.5.4, pytest 9.2.0;
PYTHONPATH=python, no C++ extension, no flashinfer. No GPU on this host — this path has no GPU dependency; GPU-marked tests skipped as usual and I did not reproduce an end-to-endft serveOOM inside a container.The twenty failures are environment-bound on both trees: eleven
tests/moe/test_cpu_moe_fp8_block.pycases andtests/kernels/test_swiglu_clamp.py::test_cpu_extension_supports_swiglu_clampneed the_cpu_moeC++ extension;tests/engine/test_cache_budget.py×2 andtests/moe/test_offload.py×1 need flashinfer;tests/models/deepseek_v41/test_engram.py::test_ftw_side_files_carry_the_tables_and_reopenneeds the_row_storeextension; the fourtests/models/test_quant_config.py::test_scheme_for_agrees_with_the_stored_tensorscases probe model snapshots in this box's HF cache.Can the tests fail?
tests/test_memory.pyrun against the first commit'smemory.py(clamp only) gives 28 failed / 101 passed, and against the file-cache commit before the dirty-page fix 10 failed / 119 passed. Eight single-line mutations of the admission path (guard removed, cgroup ignored, shard transient dropped, parallel reader never proposed, probe never fits, explicit override probed, figure not taken from the resolver, root glob instead of the shard index) each turn exactly the test written for it red.Live read inside
systemd-run --user --scope -p MemoryMax=4G -p MemorySwapMax=0(Linux 7.0, cgroup v2, MGLRU on, ext4), calling the loader's_host_ram_fits_parallelon fixture checkpoints of sparse shards; one run per column:mainReclaim and dirty pages, measured on this
memory.pybefore the rebase onto currentmain(the file is byte-identical here) and not repeated since. The same scope settings,effective_memory_available()with and without the file-cache rule:syncIn the last state a 2358 MiB anonymous allocation (above the clamp-only figure, below this one) was touched and held with the cgroup's
oom_killat 0, and a 3938 MiB one was killed. With the file written and not synced, an allocation 300 MiB insidelimit - (usage - inactive_file - active_file)was OOM-killed in 4 of 10 runs; with the file synced first, 0 of 10; sized about 300 MiB inside this branch's figure, 0 of 10. File pages another process keeps mapped and reading are counted as room: an allocation that needed them succeeded withoom_killat 0 while the reader refaulted them from disk. The figure is an admission ceiling, not a promise that nothing gets evicted.NOT VERIFIED: CUDA serving, a real model load, a Docker container on this head (the
--memory=2g/--memory=6gprobe reported earlier in this thread was on the first commit), a real cgroup v1 host (fixtures only), a kernel without MGLRU, hosted CI.memory.highis not read. File cache protected by a sibling'smemory.minunder a shared limit is counted although the kernel will not reclaim it.Performance: not measured / not changed — bounded procfs and cgroup reads at admission time, one more small file per finite cgroup level.
Prepared with AI assistance; I reviewed the change and ran the validation above.