resource_plan: fix --auto-tier pinning decode to one core (#325)#332
Open
woolcoxm wants to merge 2 commits into
Open
resource_plan: fix --auto-tier pinning decode to one core (#325)#332woolcoxm wants to merge 2 commits into
woolcoxm wants to merge 2 commits into
Conversation
physical_cpu_count() silently returned 1 in two situations, and that value
flowed through build_plan -> OMP_NUM_THREADS to pin every matmul region to a
single thread under --auto-tier (reported on Rocky Linux 9, 512 GB RAM).
Two root causes:
1. The lscpu parse counted the wrong thing. `lscpu -p=core,socket` prepends the
CPU column, so the output is actually CPU,Core,Socket; the old set
comprehension collected (CPU,core,socket) tuples that were unique per logical
CPU. Now parse CPU,Core,Socket and dedupe on (core, socket) to get true
physical cores (the SMT-doubling the surrounding comments warn against was
the actual behavior).
2. Any probe failure fell through to `os.cpu_count() or 1`. On a cgroup'd or
otherwise constrained box os.cpu_count() can be 1 (or None), silently
capping the run. Skip offline core/socket fields ("-" instead of raising
ValueError) so a single offline row no longer discards the whole parse, and
replace the silent `or 1` fallback with os.cpu_count() (logical) plus an
explicit warning. Only return 1 when nothing at all is detected, and warn.
Also harden the win32 branch: declare argtypes/restype on
GetLogicalProcessorInformationEx (an undeclared 64-bit WinAPI returns c_int and
takes c_int pointers, so the probe could silently fail), and warn on its
fallbacks. Replace the silent max(1, int()) clamp in build_plan with
_resolve_physical_cores() that clamps to 1 with a warning instead of masking.
Tests: the existing OMP_NUM_THREADS test passed physical_cpus=24 explicitly, so
it never exercised the real probe -- that is why this regressed. Add regression
coverage for the lscpu physical-core dedup, offline fields, lscpu-missing
fallback, zero-cores degenerate case, and an end-to-end build_plan +
environment_for_plan check that OMP_NUM_THREADS reflects physical (not logical)
cores.
…JustVugg#325) The core-count fix was necessary but not sufficient: @liangstein confirmed physical_cpu_count() now returns 64 on his box, yet --auto-tier STILL pinned decode to one core. A complete env diff between the plain (working) and --auto-tier (broken) paths showed exactly three keys the plan adds: OMP_NUM_THREADS = 64 (correct, verified) OMP_PROC_BIND = spread OMP_PLACES = cores Since OMP_NUM_THREADS was already correct, the culprit is the affinity pair. The mechanism: environment_for_plan() sets OMP_PROC_BIND=spread + OMP_PLACES=cores in the launcher's env. The engine's hot-thread tuning (glm.c main, the COLI_OMP_TUNED self-exec) then tries setenv("OMP_PROC_BIND","close", overwrite=0) -- but overwrite=0 cannot replace an already-set var, so the plan's "spread" wins. On the reporter's libgomp + multi-socket topology, spread + places=cores collapsed the team to a single CPU even with 64 threads configured. Fix: don't set affinity from the plan at all. The engine deliberately chose "close" for cache locality (the tiny back-to-back per-expert matmuls want adjacent cores), and the plain path already leaves affinity to the engine. Removing the plan's spread/places makes --auto-tier match the working plain path; a user wanting a specific policy can still set OMP_PROC_BIND/OMP_PLACES in their own environment (environment_for_plan only setdefaults OMP_NUM_THREADS). Verified via env diff: after the fix, --auto-tier adds ONLY OMP_NUM_THREADS beyond the plain env -- the engine's own close/bind tuning now wins on both paths identically. Note: this could not be reproduced on Windows (MinGW libgomp prints "Affinity not supported on this configuration" and ignores the vars entirely); it is Linux-libgomp-specific, matching the reporter's Rocky 9 box. Tests: rewrite test_applies_plan_without_overriding_explicit_settings to assert the plan sets NO affinity vars on any platform (the old test encoded the buggy platform-dependent spread/cores contract). Add test_plan_does_not_set_omp_affinity_vars as a focused regression. 78/78 pass.
Contributor
Author
|
Update from the field: testing is complete. @liangstein confirmed Both failure modes are now verified fixed end-to-end on the reporter's hardware:
Suite still green: 78/78 passing. This is ready for |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #325.
Problem
--auto-tierpinned decode to a single CPU core. The reporter (Rocky Linux 9, 512 GB RAM) observed that only one core was used with--auto-tier, and removing it "returned to normal." Root cause:physical_cpu_count()silently returned1, and that value flowed throughbuild_plan→environment_for_plan→OMP_NUM_THREADS.Two failure modes:
lscpuparse counted logical threads, not physical cores.lscpu -p=core,socketprepends the CPU column (always first in parseable output — per the util-linux source), so the actual output isCPU,Core,Socket. The old set comprehension built(CPU,core,socket)tuples unique per logical thread — the SMT-doubled count the surrounding comments explicitly warn against.os.cpu_count() or 1. Iflscpuwas absent, exited non-zero, timed out, or emitted a-for an offline core/socket (int("-")→ValueError, discarding the whole parse), control hitos.cpu_count() or 1. On a cgroup'd /taskset-pinned / constrained shell,os.cpu_count()can be1(orNone) → exactly the single-core behavior reported. Themax(1, int(...))clamp inbuild_planthen silently masked it.Why removing
--auto-tierworkedOn Linux the launcher sets OMP vars only on Windows (
env_for(),if sys.platform == "win32"). Without--auto-tier,OMP_NUM_THREADSstays unset on Linux, so libgomp defaults to all cores — confirming the regression is isolated to the--auto-tierenv-export path.Fix (
c/resource_plan.py)lscpu -p=CPU,Core,Socketand dedupe on(core, socket)→ true physical cores.-) instead of letting one discards the whole parse.os.cpu_count() or 1fallback withos.cpu_count()(logical) + an explicit stderr warning; the genuine "nothing detected" case still returns1but warns. Bad counts are now visible, never silent.argtypes/restypeonGetLogicalProcessorInformationEx(an undeclared 64-bit WinAPI silently returns the wrong type) + warning fallbacks.max(1, int())clamp inbuild_planwith_resolve_physical_cores()that clamps to 1 with a warning.Testing
Full Python suite: 77/77 passing (
make test-python).Core count — no longer 1 on Linux (10 mocked
lscputopologies spanning the reporter's server class + the exact failure modes):Resource allocation correct above 256 GB — ran
build_planon the realglm52_i4(78 layers, 256 experts, 384 GB) at 256/512/800 GB + 1.5 TB. Both conservation invariants hold exactly in every case:dense + runtime + cache <= budget(never over-allocates RAM)hot + warm + cold == expert_bytes(no experts lost/duplicated between tiers)At 512 GB (the reporter's box) and 800 GB (
--ram 800): cap = 256/layer, 0 cold bytes, bottleneck = CPU compute. Correct.Regression tests added (
c/tests/test_resource_plan.py, 7 new):-field handlingbuild_plan+environment_for_planassertsOMP_NUM_THREADS= physical corestest_plan_conserves_budget_and_experts_above_256gb— locks in the >256 GB budget/expert conservation invariants at 256/512/800 GBThe pre-existing
OMP_NUM_THREADStest passedphysical_cpus=24explicitly and so never exercised the real probe — that is why this regressed.Branch based off
dev.🤖 This PR was developed with ZCode.