resource_plan: stop setting OMP_PROC_BIND/OMP_PLACES from --auto-tier (#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.
This commit is contained in:
+12
-5
@@ -461,11 +461,18 @@ def environment_for_plan(plan, env=None, cuda_enabled=True):
|
||||
result = dict(env or {})
|
||||
result.setdefault("COLI_POLICY", plan["policy"]["name"])
|
||||
result.setdefault("OMP_NUM_THREADS", str(plan["cpu"]["physical_cores"]))
|
||||
if sys.platform != "win32":
|
||||
# la libgomp di MinGW non supporta l'affinity su Windows
|
||||
# ("Affinity not supported on this configuration"): non impostarle li'.
|
||||
result.setdefault("OMP_PROC_BIND", "spread")
|
||||
result.setdefault("OMP_PLACES", "cores")
|
||||
# NOTE: we intentionally do NOT set OMP_PROC_BIND / OMP_PLACES here.
|
||||
# The engine's own hot-thread tuning (glm.c main(), the COLI_OMP_TUNED
|
||||
# self-exec) sets OMP_PROC_BIND=close with overwrite=0 -- it prefers
|
||||
# packing the team onto adjacent cores for the tiny back-to-back per-expert
|
||||
# matmuls. Pre-setting OMP_PROC_BIND=spread here ran first and won (the
|
||||
# engine's overwrite=0 setenv could not override an already-set var), and
|
||||
# spread + OMP_PLACES=cores collapsed the team to one CPU on some libgomp /
|
||||
# multi-socket topologies (#325: --auto-tier pinned decode to 1 core on a
|
||||
# 64-core box even with OMP_NUM_THREADS=64). Leaving affinity to the engine
|
||||
# makes --auto-tier match the plain (working) path. A user who wants a
|
||||
# specific policy can still set OMP_PROC_BIND/OMP_PLACES in the environment
|
||||
# themselves -- setdefault above only covers OMP_NUM_THREADS.
|
||||
tune = plan.get("tune", {})
|
||||
for key, entry in tune.items():
|
||||
if key.startswith("_"):
|
||||
|
||||
@@ -109,6 +109,19 @@ class ResourcePlanTest(unittest.TestCase):
|
||||
self.assertEqual(plan["cpu"]["physical_cores"], 12)
|
||||
self.assertEqual(env["OMP_NUM_THREADS"], "12")
|
||||
|
||||
def test_plan_does_not_set_omp_affinity_vars(self):
|
||||
# The real #325 regression: --auto-tier set OMP_PROC_BIND=spread +
|
||||
# OMP_PLACES=cores, which ran before the engine's overwrite=0 setenv and
|
||||
# so won, collapsing the OpenMP team to one CPU on the reporter's 64-core
|
||||
# Linux box even though OMP_NUM_THREADS was correct. The plan must leave
|
||||
# affinity to the engine's own hot-thread tuning (which prefers 'close').
|
||||
plan = build_plan(self.model, available_memory=16 * GB,
|
||||
available_disk=1, gpus=[], physical_cpus=64)
|
||||
env = environment_for_plan(plan)
|
||||
self.assertEqual(env["OMP_NUM_THREADS"], "64")
|
||||
self.assertNotIn("OMP_PROC_BIND", env)
|
||||
self.assertNotIn("OMP_PLACES", env)
|
||||
|
||||
def test_plan_conserves_budget_and_experts_above_256gb(self):
|
||||
# Regression for #325's reporter: a 512 GB machine loading the whole
|
||||
# model into RAM. Verify the budget math stays exact at large RAM sizes
|
||||
@@ -164,13 +177,13 @@ class ResourcePlanTest(unittest.TestCase):
|
||||
self.assertEqual(env["COLI_CUDA"], "1")
|
||||
self.assertEqual(env["COLI_GPUS"], "1")
|
||||
self.assertEqual(env["OMP_NUM_THREADS"], str(plan["cpu"]["physical_cores"]))
|
||||
if sys.platform == "win32":
|
||||
# MinGW libgomp: niente affinity su Windows, le chiavi non vanno emesse
|
||||
self.assertNotIn("OMP_PROC_BIND", env)
|
||||
self.assertNotIn("OMP_PLACES", env)
|
||||
else:
|
||||
self.assertEqual(env["OMP_PROC_BIND"], "spread")
|
||||
self.assertEqual(env["OMP_PLACES"], "cores")
|
||||
# The plan must NOT set OMP_PROC_BIND / OMP_PLACES on any platform:
|
||||
# the engine's own hot-thread tuning owns affinity (it prefers
|
||||
# OMP_PROC_BIND=close for the back-to-back per-expert matmuls). Setting
|
||||
# spread + cores here ran before the engine's overwrite=0 setenv and so
|
||||
# won, collapsing the team to one CPU on some libgomp topologies (#325).
|
||||
self.assertNotIn("OMP_PROC_BIND", env)
|
||||
self.assertNotIn("OMP_PLACES", env)
|
||||
self.assertEqual(env["PIN_GB"], env["CUDA_EXPERT_GB"])
|
||||
|
||||
explicit_threads = environment_for_plan(plan, {"OMP_NUM_THREADS": "7",
|
||||
|
||||
Reference in New Issue
Block a user