[rl] Cap OMP threads per spawned proc via launcher.env - #4031
Draft
HosseinKaviani-H wants to merge 1 commit into
Draft
[rl] Cap OMP threads per spawned proc via launcher.env#4031HosseinKaviani-H wants to merge 1 commit into
HosseinKaviani-H wants to merge 1 commit into
Conversation
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.
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.
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.envconfig 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 inlauncher.envoverrides 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.
TODO: fold
launcher.envinto a shared OSS launcher config once one exists.