Skip to content

[rl] Cap OMP threads per spawned proc via launcher.env - #4031

Draft
HosseinKaviani-H wants to merge 1 commit into
pytorch:mainfrom
HosseinKaviani-H:rl-launcher-env-omp-cap
Draft

[rl] Cap OMP threads per spawned proc via launcher.env#4031
HosseinKaviani-H wants to merge 1 commit into
pytorch:mainfrom
HosseinKaviani-H:rl-launcher-env-omp-cap

Conversation

@HosseinKaviani-H

Copy link
Copy Markdown
Contributor

Monarch (unlike torchrun/Ray) does not set OMP_NUM_THREADS on the procs it spawns, so each co-located trainer/generator proc sizes its BLAS/OMP thread pool to every core on the host. With N procs on one host that is N-way oversubscription; the cores thrash and starve the generator GPUs.

Add a minimal launcher.env config field and, by default, cap threads to cores/num_local_procs. The no-oversubscription partition Ray does implicitly via num_cpus per actor. The cap is applied to both roles' launch env so any key in launcher.env overrides it per key.

Measured on a single 8-GPU host (368 cores): CPU busy cores collapse from ~118 to ~18, generator GPU util rises ~73% -> ~87%, and end-to-end wall-clock improves ~1.3x (dapo_math) to ~1.75x (alphabet_sort); the thrash-driven per-step variance also collapses.

image

TODO: fold launcher.env into a shared OSS launcher config once one exists.

Monarch (unlike torchrun/Ray) does not set OMP_NUM_THREADS on the procs it
spawns, so each co-located trainer/generator proc sizes its BLAS/OMP thread
pool to every core on the host. With N procs on one host that is N-way
oversubscription; the cores thrash and starve the generator GPUs.

Add a minimal `launcher.env` config field (a generic per-proc launch-env dict)
and, by default, cap threads to cores/num_local_procs (the no-oversubscription
partition Ray does implicitly via num_cpus per actor). The cap is applied to
both roles' launch env; any key in `launcher.env` overrides it per key.

Measured on a single 8-GPU host (368 cores): CPU busy cores collapse from
~118 to ~18, generator GPU util rises ~73% -> ~87%, and end-to-end wall-clock
improves ~1.3x (dapo_math) to ~1.75x (alphabet_sort); the thrash-driven
per-step variance also collapses.

Note: changing thread count alters BLAS reduction order, so loss is not
bitwise-identical to an uncapped run (it is deterministic for a fixed cap).
This is a perf/infra change; validate via convergence + wall-clock, not
bitwise parity.

TODO: fold `launcher.env` into a shared OSS launcher config once one exists.
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 30, 2026
@HosseinKaviani-H HosseinKaviani-H changed the title [rl] Cap BLAS/OMP threads per spawned proc via launcher.env Cap OMP threads per spawned proc via launcher.env Jul 30, 2026
@tianyu-l
tianyu-l requested review from felipemello1 and pzhan9 August 1, 2026 00:18
@HosseinKaviani-H HosseinKaviani-H changed the title Cap OMP threads per spawned proc via launcher.env [rl] Cap OMP threads per spawned proc via launcher.env Aug 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/rl ciflow/8gpu CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant