Make CPU optim multithreading by overwriting num-thread but without o…

[MarioSieg]

Increase CPU num threads for the CPU optim to use all physical cores.
Override OMP_NUM_THREADS but tried to not oversubscribe.
Need a bench run to validate that it's faster.

---

Note

Medium Risk
Touches the CPU optimizer offload hot path and process-wide thread settings; incorrect affinity/fair-share math could oversubscribe cores or change offload step performance.

Overview
Automatically raises PyTorch intra-op threads when optim_cpu_offload is enabled, so bandwidth-bound CPU AdamW no longer runs single-threaded under the launcher's default OMP_NUM_THREADS=1.

Both RL and SFT trainers now call new configure_cpu_optimizer_threads() at startup for any offload mode. Each rank takes cpu_count / local_world_size threads, capped by its affinity mask, to avoid oversubscribing the node. This overrides OMP_NUM_THREADS, including values set via env_vars.

^{Reviewed by Cursor Bugbot for commit 21b7ea6. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 21b7ea6. Configure here.}

[cursor]

Thread fair-share oversubscribes shared CPUs

Medium Severity

configure_cpu_optimizer_threads budgets each rank with os.cpu_count() / local_world_size, then caps only by that rank's full affinity size. When several ranks share the same mask — after numa_bind on one socket, or under a shared cpuset — that cap is not divided among peers, so total intra-op threads can far exceed the CPUs in the mask despite the no-oversubscribe goal.

 

^{Reviewed by Cursor Bugbot for commit 21b7ea6. Configure here.}

[samsja]

Benchmarked on the reserved 8xH200 node (Qwen3-30B-A3B, fake data, random init, native full offload with numa_bind, 5 steps, config env_vars cleared so only the auto-config sets threads — log confirms CPU optimizer uses 30 intra-op threads (120 CPUs in this rank's affinity mask, 8 local ranks)):

Point	1 thread (launcher default)	OMP_NUM_THREADS=28 via env_vars	this PR (auto, 30)	no-offload reference
Qwen 8K	15.66 s	6.13 s	3.95 s	4.57 s
Qwen 16K	—	7.89 s	5.91 s	9.35 s

Pipeline diagnostics at 8K: adam_kernel 2.9 → 1.86 s, materialize 1.65 → 0.62 s, exposed drain 2.5 → 0.92 s. torch.set_num_threads scopes the thread budget to the ATen pool instead of handing every OMP consumer in the process a 28-thread team, which is a large improvement over the env-var approach — the offload crossover moves below 8K and full offload now beats no-offload at every sequence length tested. Losses unchanged (11.9312 every step).

GLM-5 20-layer proxy rerun at 16K in progress; will follow up.

🤖 Generated with Claude Code

[samsja]

GLM-5 20-layer proxy (171B, CP2/EP8) at 16K — the point with the largest exposed pipeline cost:

	before (env_vars OMP=28)	this PR (auto, 30)
step median	25.6 s	19.4 s (−24%)
adam_kernel	17.6 s	12.8 s
materialize	5.3 s	3.8 s
exposed drain	17.4 s	12.6 s

Loss unchanged (11.9504). Still DRAM-bandwidth-bound at this parameter scale, but the scoped thread budget recovers a lot of it. LGTM from the bench side — with this merged, the offload crossover on Qwen3-30B moves below 8K.

🤖 Generated with Claude Code

[samsja]

Correction to my earlier comments: the no-offload reference numbers included a ~1.4 s/step slow path in the zero-gradient-ratio metric scan (a per-parameter device-scalar loop; fixed in #3249 — the scan only runs in the no-offload path, so offload runs were unaffected). Fair comparison with both #3248 and #3249 applied, Qwen3-30B-A3B:

Seq	No offload (fair)	Full offload (this PR)
8K	3.39 s	3.95 s
16K	6.23 s	5.91 s

So "beats no-offload at every sequence length" was overstated — the corrected read is parity from ~16K upward, at half the peak HBM. This PR's own improvement stands as measured (6.13 s → 3.95 s at 8K).

🤖 Generated with Claude Code