perf(omp): keep OpenMP worker team hot across tiny per-expert matmul regions (seed + re-exec once, fully overridable) (#77)

The MoE forward does many tiny, back-to-back per-expert matmul regions.
Under libgomp's default passive wait policy the worker team is parked
between regions, and the thread re-wake latency dominates the actual
compute. Switching the team to an active wait policy (with a bounded spin
count so long NVMe expert-load stalls still yield instead of burning a
core) keeps the threads hot across those regions.

Measured on the Zen5 build: expert-matmul wall time 66.9s -> 20.9s
(~3.2x on that phase). Output is numerically unchanged — this only
affects thread scheduling, not the computation.

Mechanism: libgomp parses OMP_/GOMP_ vars in a constructor that runs
before main(), so setenv() from main() and continuing is too late
(verified: the already-initialised runtime ignores it). Instead we seed
the winning defaults on first entry with overwrite=0 (so any value the
user already set wins), set a COLI_OMP_TUNED sentinel, and execv() self
once so a fresh libgomp constructor reads them. The sentinel guards the
exec, so we re-exec at most once. The block is the first statement in
main() so argv is passed verbatim to execv().

Discoverability / observability (review follow-up):
  - COLI_NO_OMP_TUNE=1 is a documented kill-switch that disables the
    whole re-exec + tuning path in one shot (distinct from the internal
    COLI_OMP_TUNED re-exec sentinel);
  - a one-line stderr breadcrumb is printed before the re-exec
    ("[OMP] hot-thread tuning: re-exec once (COLI_NO_OMP_TUNE=1 to skip)")
    so the self-re-exec is self-documenting;
  - execv only returns on failure, so a perror() now follows it
    ("execv self-reexec failed, running untuned") — a container without
    /proc or a deleted inode falls back visibly instead of silently
    losing the ~3.2x.

Fully opt-out / overridable and lossless:
  - explicit OMP_WAIT_POLICY / GOMP_SPINCOUNT / OMP_PROC_BIND / OMP_DYNAMIC
    in the environment win (overwrite=0);
  - pre-setting COLI_OMP_TUNED=1 skips the re-exec entirely;
  - COLI_NO_OMP_TUNE=1 disables the tuning path entirely;
  - guarded by __linux__ (no-op re-exec elsewhere; the setenv defaults
    still apply for any later-initialising runtime).


Claude-Session: https://claude.ai/code/session_01DS7oc65c5RdA9V99otRCwt

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Rodolfo Hansen
2026-07-12 13:40:29 +02:00
committed by GitHub
parent 2416bc9079
commit 704ed49c16
+31 -2
View File
@@ -2494,8 +2494,37 @@ static void cap_for_ram(Model *m, double ram_gb, int ebits, int max_ctx){
}
int main(int argc, char **argv){
/* i thread OMP non devono girare a vuoto mentre il main aspetta il disco */
if(!getenv("OMP_WAIT_POLICY")) setenv("OMP_WAIT_POLICY","passive",1);
/* ---- Permanent OpenMP hot-thread tuning. The per-expert matmul regions are
* tiny and back-to-back; with the default passive wait policy libgomp parks
* the worker team between regions and the re-wake latency dominates. Keeping
* the threads hot (active spin) collapses that overhead — measured matmul
* time 66.9s -> 20.9s on the Zen5 build, with no change to numerical output.
*
* libgomp reads the OMP_ / GOMP_ vars in a CONSTRUCTOR that runs before
* main(), so setenv() here and continuing would be too late (verified:
* setenv-in-main is ignored by the already-initialised runtime). Instead, on
* first entry seed the winning defaults — respecting anything the user
* already set (overwrite=0) — then re-exec self once so a fresh libgomp
* constructor picks them up. The COLI_OMP_TUNED sentinel guards the exec so
* we re-exec at most once. Fully overridable: any explicit OMP_/GOMP_ env the
* user sets wins (overwrite=0), pre-setting COLI_OMP_TUNED=1 skips the
* re-exec entirely (runs with whatever policy the environment already has),
* and COLI_NO_OMP_TUNE=1 is a documented kill-switch that disables the whole
* re-exec + tuning path (distinct from the internal COLI_OMP_TUNED sentinel).
*
* Must remain the FIRST statement in main(): argv is passed verbatim to execv(). */
if(!getenv("COLI_OMP_TUNED") && !getenv("COLI_NO_OMP_TUNE")){
setenv("OMP_WAIT_POLICY","active",0); /* keep the team hot across the tiny per-expert matmul regions */
setenv("GOMP_SPINCOUNT","200000",0); /* spin briefly, then yield so long disk waits don't burn a core */
setenv("OMP_PROC_BIND","close",0); /* pack the team onto adjacent cores for cache locality */
setenv("OMP_DYNAMIC","FALSE",0); /* fixed team size: no per-region thread-count churn */
setenv("COLI_OMP_TUNED","1",1);
#ifdef __linux__
fprintf(stderr,"[OMP] hot-thread tuning: re-exec once (COLI_NO_OMP_TUNE=1 to skip)\n");
execv("/proc/self/exe", argv); /* returns only on failure -> fall through and run untuned */
perror("[OMP] execv self-reexec failed, running untuned");
#endif
}
const char *snap=getenv("SNAP"); if(!snap){fprintf(stderr,"SNAP=<dir>\n");return 1;}
g_nopack = getenv("NOPACK")?1:0;
g_drop = getenv("DROP")?1:0;