feat(wandb): log full RLConfig to W&B instead of per-component subconfig - #3110
Draft
samsja wants to merge 2 commits into
Draft
feat(wandb): log full RLConfig to W&B instead of per-component subconfig#3110samsja wants to merge 2 commits into
samsja wants to merge 2 commits into
Conversation
The rl entrypoint splits the config into subconfigs (trainer.toml, orchestrator.toml, inference.toml) before launching subprocesses. Each process only loads its subconfig, so wandb.init only received the orchestrator config (the primary in shared W&B mode). The trainer, inference, slurm, deployment, and shared sections were missing entirely. Fix: the rl() entrypoint writes the full resolved RLConfig as JSON (output_dir/configs/full_config.json) and sets PRIME_FULL_CONFIG_PATH in os.environ so all locally-launched subprocesses inherit it. The multi-node SLURM template exports the same env var separately (since srun tasks don't inherit the launcher's os.environ). The WandbMonitor reads this file and uses it as the wandb run config when available, falling back to the per-component subconfig otherwise.
After wandb.init, save all files in the config_dir (full_config.json + subconfig TOMLs) to the W&B run so they are downloadable from the UI Files tab. Uses WANDB_FULL_CONFIG_PATH parent dir to locate the config directory.
|
|
||
| RL_TOML = "rl.toml" | ||
| RL_SBATCH = "rl.sbatch" | ||
| FULL_CONFIG_JSON = "full_config.json" |
Member
There was a problem hiding this comment.
call this rl.json, will land as <output>/configs/rl.json
Member
There was a problem hiding this comment.
and is consistent with other artifacts
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
2 participants
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.
Summary
The
rlentrypoint splits the fullRLConfiginto per-component subconfigs (trainer.toml,orchestrator.toml,inference.toml) before launching subprocesses. Each process only loads its subconfig, sowandb.init(config=...)only received the orchestrator config (the primary in shared W&B mode). The trainer, inference, slurm, deployment, and shared sections were missing entirely from the W&B run config.Changes
rl.py: Therl()entrypoint writes the full resolvedRLConfigas JSON (output_dir/configs/full_config.json) and setsWANDB_FULL_CONFIG_PATHinos.environso all locally-launched subprocesses inherit it.multi_node_rl.sbatch.j2: ExportsWANDB_FULL_CONFIG_PATHdirectly (srun tasks don't inherit the launcher'sos.environ).wandb.py:WandbMonitorreads the full config JSON whenWANDB_FULL_CONFIG_PATHis set, falling back to the per-component subconfig otherwise. Afterwandb.init, also uploads all config files (full_config.json+ subconfig TOMLs) to the W&B run so they're downloadable from the UI Files tab.All three launch modes are covered:
os.environpropagates to subprocesses via**os.environinPopenenv dictsuv run rl->rl()-> setsos.environ->rl_local()WANDB_FULL_CONFIG_PATHalongside existingWANDB_SHARED_*varsVerification
Tested end-to-end with
WANDB_MODE=offline uv run rl @ examples/basic/reverse-text/rl.toml --max-steps 3 --output-dir /tmp/rl-test --wandb:full_config.jsonwritten with all 17 top-level keys (trainer, orchestrator, inference, slurm, deployment, wandb, ckpt, model, etc.)Using full config from /tmp/rl-test/configs/full_config.json for W&Bfull_config.json,trainer.toml,orchestrator.toml,inference.toml) uploaded to the W&B run's Files tab