fix(bench): bench_mlxlm.py hardcodes a 128 GB memory budget
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem / Background
`scripts/bench_mlxlm.py` computes the sweep's admission budget from a fixed 128 GB rather than from the host it runs on (`scripts/bench_mlxlm.py:61-68`):
```python
# Memory budget: 85% of 128 GB. Override with PYLM_BENCH_MAX_GB env var to
# enforce a tighter cap (e.g. PYLM_BENCH_MAX_GB=65 to skip very large MoE
# models on a 128 GB host where 85% × 128 GB ≈ 108 GB still admits them).
_env_max_gb = os.environ.get("PYLM_BENCH_MAX_GB")
if _env_max_gb:
MEMORY_LIMIT_BYTES = int(float(_env_max_gb) * 1024 * 1024 * 1024)
else:
MEMORY_LIMIT_BYTES = int(128 * 1024 * 1024 * 1024 * 0.85)
```
The comment is written as though 128 GB were the only host size.
The failure this causes is not a skipped row, it is a missing comparison. `bench_mlxlm.py` and `bench_decode.sh` produce the paired halves of one measurement, so when one harness skips a checkpoint the other is left with nothing to compare against. On a 512 GB machine all three checkpoints of a recent sweep came back `SKIP:oom_estimate (141.1 GB > 108.8 GB)` while `bench_decode.sh` had just measured every one of them on that same host. `PYLM_BENCH_MAX_GB` can paper over this per run, but only if the operator already suspects the budget is wrong, and the budget is not printed anywhere except inside the skip line itself.
On hosts below 128 GB it errs the other way, and that direction is not cosmetic either. Measured on the GB10 box this repository's CUDA work runs on, `free -g` reports 121 GB total and `os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')` returns 121.69 GiB, so the true 85 percent budget is 103.4 GiB against the hardcoded 108.8 GiB. The sweep therefore admits about 5.4 GiB more than the host can hold. This host has been lost twice to GPU memory exhaustion during current work, once as a hard reboot on 2026-09-10 at about 18:29 with two large models resident and once after the NVIDIA driver accumulated 586 `NVRM: NV_ERR_NO_MEMORY` failures, so an over-generous admission default is a plausible contributor rather than a purely cosmetic defect.
## Current Behavior
The guard at `scripts/bench_mlxlm.py:552-557` is commented "OOM guard (same as bench_decode.sh)", and it is not the same. The shell harness already derives its budget from the host: `detect_memory_bytes()` at `scripts/bench_decode.sh:315-321` uses `sysctl -n hw.memsize` on Darwin and `free -b 2>/dev/null | awk '/^Mem:/{print $2}'` elsewhere, and `scripts/bench_decode.sh:323-325` takes 85 percent of that. Only the Python half is hardcoded, which is exactly why the two harnesses disagree about which checkpoints exist.
The resolved budget is printed only inside the skip message (`bench_mlxlm.py:554`), so a run that skips nothing never records the limit it applied, and a CSV cannot afterwards be audited for which checkpoints were even eligible.
`detect_hardware()` (`bench_mlxlm.py:79-102`) shells out to `sysctl -n machdep.cpu.brand_string` and returns `("unknown", ...)` on anything else, and all 50 committed `benchmarks/pylm_*.csv` files come from Apple hosts (m1ultra, m5max). Nothing in the file gates it to macOS, so any detection added here has to work on Linux as well.
## Status: a fix already exists, unmerged
This issue is a record of the bug and its evidence, not a work request. Another benchmarking session on the fleet, on a Mac Studio M3 Ultra, found the same bug from the large-host direction and has an uncommitted local fix to this file, alongside an unrelated `--text-via-vlm` mode. That change has not been approved for pushing by its owner. Do not implement this independently: it would duplicate work that is already written and would collide in the same file. Ownership is being resolved separately, and this issue should not be picked up until that is settled.
## Proposed Solution (for whoever lands the existing fix)
Derive `MEMORY_LIMIT_BYTES` from the host's physical memory, keeping the 85 percent fraction and keeping `PYLM_BENCH_MAX_GB` as the override that wins.
- `os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')` covers macOS and Linux without a subprocess, and was verified on this host to return 121.69 GiB against `free -g`'s 121.
- Fall back to `sysctl -n hw.memsize` on Darwin and `MemTotal` from `/proc/meminfo` on Linux when `sysconf` raises, mirroring how `bench_decode.sh:315-321` branches.
- If every probe fails, disable the guard rather than substituting a number. `estimate_model_size` already returns 0 for an undeterminable size and `bench_decode.sh:349` treats that as "try anyway".
- Rewrite the comment so it states the policy (85 percent of detected physical memory) instead of a single host size, and drop the `PYLM_BENCH_MAX_GB` example phrased around 128 GB.
- Print the resolved budget and its source once at startup, beside the `>>> [store]` and `>>> [candidates]` lines added at `bench_mlxlm.py:475-477` and `bench_mlxlm.py:499`, so a run's admission decisions are auditable from its stderr afterwards.
**Out of scope:** `BENCH_MEM_OVERHEAD_FACTOR` (`bench_decode.sh:211-216`, applied at `bench_decode.sh:345-356`), which the Python harness also lacks. That is a second divergence between the two guards and belongs in its own issue; folding it in here would change admission for an unrelated reason.
## Acceptance Criteria
- [ ] With `PYLM_BENCH_MAX_GB` unset on a host of N GiB physical memory, the resolved budget is `0.85 * N`: 103.4 GiB on this GB10 box rather than 108.8.
- [ ] `PYLM_BENCH_MAX_GB=65` still resolves to exactly 65 GiB on any host.
- [ ] The resolved budget and whether it came from detection or from the override are printed once at startup, on a run that skips nothing.
- [ ] No literal 128 remains in the budget computation or in its comment.
- [ ] A checkpoint that `bench_decode.sh` measures successfully on a given host is not returned as `SKIP:oom_estimate` by `bench_mlxlm.py` on that same host.
## Verification
```bash
# Budget resolution is module level today, so importing the file is enough.
python3 -c "import importlib.util as u; s=u.spec_from_file_location('b','scripts/bench_mlxlm.py'); m=u.module_from_spec(s); s.loader.exec_module(m); print(m.MEMORY_LIMIT_BYTES/1024**3)"
# Before this fix: 108.79999999981374 on every host, measured.
# After: about 103.4 on this 121.69 GiB GB10 box.
python3 -c "import os; print(os.sysconf('SC_PAGE_SIZE')*os.sysconf('SC_PHYS_PAGES')/1024**3)"
free -g # Linux cross-check
sysctl -n hw.memsize # macOS cross-check
PYLM_BENCH_MAX_GB=65 python3 -c "import importlib.util as u; s=u.spec_from_file_location('b','scripts/bench_mlxlm.py'); m=u.module_from_spec(s); s.loader.exec_module(m); print(m.MEMORY_LIMIT_BYTES/1024**3)"
# Expect exactly 65.0 regardless of host.
```
Verified against `60341873`.
Contributor guide
Research direction
Do not start this while the existing local fix and ownership are unresolved. If work is reassigned, inspect scripts/bench_mlxlm.py:61-68 and 552-557 alongside scripts/bench_decode.sh:315-325, then run the module-import and PYLM_BENCH_MAX_GB=65 verification commands; done means host-derived and override budgets are auditable and acceptance criteria pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, shell
- Domain
- performance, tooling
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 20/100