Reward Profiling Infra Update - #2862
Open
jkyi-nvidia wants to merge 12 commits into
Open
Conversation
Adds nemo_gym/sweep, a generic mechanism for running many environments as one Gym deployment, and benchmarks/nemotron_3.5_super/reward_profiling, the manifests and launchers that use it to profile the RL training blend. Rollout collection already dispatches per row on agent_ref, so datasets for different environments can be concatenated into one input file and served by a single deployment composed from the union of their configs. That turns ~110 per-environment jobs into one job shape covering 36 environments and 726,121 source rows, with judge-scored, sandbox-backed and plain environments in the same run. The manifest is the committed unit of configuration, and every block is named for the command that consumes it, so where a setting belongs follows from which command reads it: materialize, sbatch, srun, vllm, gym_env_start, gym_eval_run, gym_eval_profile. Precedence is uniform throughout -- manifest, then script environment variable, then command line -- so a run needs only SWEEP_DIR. Two properties make sharding safe. _ng_task_index is stamped at materialize time and Gym never rewrites it, so shard rollouts concatenate without renumbering; merge deduplicates on (_ng_task_index, _ng_rollout_index), the same key --resume uses, so a rerun shard cannot double-count. Resharding folds collected rollouts back into the parent and snapshots it before touching any shard directory. Splitting per environment keys on the task index against the per-entry range in the report rather than on agent_ref, because entries routinely share an agent -- the three ns_tools entries all dispatch to ns_tools_simple_agent -- and splitting on the agent would merge distinct environments. Because task indices are positional, re-materializing over collected rollouts is refused: a changed limit renumbers every task after it, and --resume would then match real rollouts to different rows with no error. Verified by running rather than inspection. The full pipeline finished at 2,304/2,304 rollouts with 36/36 environments profiled at 100% and 288 profiled rows from 288 source tasks. Failure recovery was tested by killing a shard at 2,300/2,304: the watcher recomputed outstanding work from collected rollouts rather than exit status (4, not 2,304), resubmitted, and resume carried the 2,300 forward. Sharding invariants hold across six shard counts (4-7-2-9-3-5) on partially collected data; scripts/debug_selftest.sh reproduces that against any sweep directory. 96 unit tests, 96% coverage of nemo_gym/sweep. Not yet run at production scale: the above used LIMIT_PER_ENTRY=8, so 288 tasks across one to two shards, against 5,808,968 rollouts over 16 shards and 48 nodes for the real sweep. RATES.md carries measured per-environment rates and sizing. build_eval_container.sh gains SKIP_PREPARE, defaulting to the existing behaviour. gym eval prepare aborts with "No benchmark config found" on train-split configs, which is what a reward-profiling sweep supplies. README.md documents running it, CONTRIBUTING.md adding an environment. Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
…cher
05_profile.sh described the fix ("run this under bash -lc") instead of giving
it. It knows SWEEP_DIR and its own path, so it now prints the exact line to
paste. Deliberately still refuses to source anything itself: env.yaml is YAML
holding ${oc.env:VAR} references rather than assignments, so there is nothing to
source, and sourcing a user's shell config would run arbitrary interactive setup
the script cannot see.
03_run_single.sh gets section banners. Its config assignments genuinely
interleave with logic between lines 121 and 280 -- some of that is forced, since
NUM_SAMPLES_IN_PARALLEL derives from MAX_NUM_SEQS and NUM_DECODE_NODES and
WALLTIME falls back to the manifest's SBATCH_TIMELIMIT, so those cannot precede
what they read. Reordering the rest is where two ordering bugs already came from
today, so this marks the sections instead of moving them.
The export VLLM_* and UCX_* further down are not scattered config: they sit
inside the generated batch script's heredoc and are the job's runtime
environment.
Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
The script runs `gym eval profile`, but its only guard was `python -c "import nemo_gym"`. With GYM_SITE_PACKAGES set and no venv active that passes and then every label fails with "command not found" -- 36 identical errors instead of one, each written into its own profile.txt. Checks `command -v gym` before doing any work, skipped when VLLM_JOBID is set since that path sources the container's venv itself. Deliberately does not activate a venv on the caller's behalf. README 00 already makes that step one, and a script that rewrites PATH could hand you a different gym than you meant. 03_run_endpoint.sh already had this check; its message now names the same fix. Verified both ways: importable-but-no-gym now fails once with the command to run, and a normal profile with the venv active still does 36 labels plus the whole sweep. Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
…ifest Same bug I fixed for CONTAINER and missed one line over. This script uses neither MODEL nor CONTAINER -- it passes the environment through to 03_run_single.sh, which resolves both from the shard's sweep_report.json when the manifest names them. Requiring MODEL here rejected exactly the case the vllm block was added for, so a manifest-driven sharded run died at submit with "MODEL is required" before anything started. Only SWEEP_DIR is required now; 03_run_single.sh still validates MODEL and CONTAINER after applying the manifest. Found by re-running end to end after the merge rather than by reading, which is the only reason it surfaced -- every unit test and the CPU-side pipeline passed with the bug present. Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
The error hardcoded "the manifest's srun block" for both MODEL and CONTAINER. MODEL moved to the vllm block when that was added, so the message sent you to the wrong place to fix it. MODEL now reports vllm, CONTAINER reports srun. Both verified by removing each from the manifest in turn. Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
…ocs together
The watcher lock did not work. 03_run_sharded.sh took it on
${SHARDS_DIR:-$SWEEP_DIR} at line 40, but SHARDS_DIR is not defaulted until line
72 -- so a watcher with it set and one without locked different files and both
ran. That is the exact pile-up the lock was added to prevent, and it would only
have shown up under the conditions that caused the original incident. Locks on
SWEEP_DIR now, which is required and unambiguous.
05_profile.sh imported ALLOW_PARTIAL_ROLLOUTS from the manifest and then
hardcoded ++allow_partial_rollouts=True in both branches, so
gym_eval_profile.allow_partial_rollouts: false had no effect on the one script
whose whole job is profiling.
03_run_attached.sh passed no allow_partial_rollouts at all, and the profiler
defaults to False, so profiling a walltime-killed attached run errored on
partial groups while every other path handled them.
--job-name shadowed SBATCH_JOB_NAME the same way --time shadowed
SBATCH_TIMELIMIT before it: a flag beats an env var, so a manifest job_name
never reached the job. Falls back now.
GymEvalProfile was extra="allow" with neither guard its siblings have. Its
extras are exported as environment variables, not ++ overrides, so
`allow_partial_rollout` (no s) sat there as a shell variable while the real
field kept its default, and a key with a dot or dash would kill the launcher
with "not a valid identifier" from inside a process substitution.
02_shard.sh defaulted NUM_SHARDS to 16 without reading the manifest, so running
it standalone against a manifest declaring 8 silently dealt 16.
Verified: 02_shard now deals 16 from the manifest and 3 from the environment;
both new validators reject and legitimate keys still pass; 96 tests; 5/5
manifests.
Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
Factual corrections only; left the reorganisation suggestions alone since the docs were already restructured. Wrong about what the code does: - README said 03_run_single.sh calls 05_profile.sh. It profiles inline; the caller is 03_run_endpoint.sh. - README had num_samples_in_parallel precedence backwards -- the manifest value wins, and 512 x decode_nodes is the last fallback, not an override. - README advertised gym_eval_profile as taking "any ++ the profiler takes". The launchers export that block as environment variables, so only allow_partial_rollouts and jobs are ever read. - MAX_ROUNDS counts total attempts, not resubmissions: 4 means one submission plus three retries. Said in three places, wrong in all of them. - 03_run_sharded.sh claimed NUM_SHARDS=16 is a 256-node run while its own next paragraph said 48. 48 is right. Commands that failed as written: - `validate manifests/<name>.yaml` has no such path from the repo root. - The container recipe's second mkdir created ./vllm in the cwd, which nothing uses, and the first lacked -p so it failed on a fresh checkout. - 01_materialize.sh's usage line named the archived manifest. Stale source citations, all verified against the files: server_utils.py:595 -> :680 for the one-worker default and :159 -> :162-163 for the connector-limit division; rollout_collection.py:740 -> global_config.py:173-174, since _ng_rollout_index does not appear in rollout_collection.py at all; reward_profile.py:42 -> :45. Also: RATES.md attributed its measurements to nemotron_3_5_super.yaml when they were taken against nemotron_3_ultra.yaml, whose own header says so; the tail list omitted reasoning_gym, which RATES measures as slower than both environments it does name; CONTRIBUTING called owner required when the schema has it optional; and the lock doc now matches the code after moving it to the sweep directory. Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
The comment said 64 tokens "is enough for a compliant 'A', so this may never bind". RATES.md records the opposite from measurement: the pre-fix 544/hr figure is void because the capped judge returned early, and the sandbox and judge rates depend on the overlay being present. A measurement beats a guess, so the guess goes. Value stays 8192, and the comment now says why it is generous rather than tight: the failure is silent -- parse_judge_grade falls back to "B" (INCORRECT) on anything unparseable, so a truncated judge misgrades rather than erroring -- while too high only costs latency, and matching the ultra-3 branch keeps the RATES.md numbers comparable. Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
nemotron_3_ultra.yaml still carried the "64 is enough ... so this may never bind" line that RATES.md contradicts. Value was already 8192 in both, so this is comment-only and the archived manifest still reflects what RATES was measured against. Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
shard.py parsed rollouts with an unguarded json.loads in the absorb path, the one place that runs on a killed job's output -- where a torn final line is normal. The deal path in the same file already had the guard. A reshard over a partially written rollouts file crashed with a raw JSONDecodeError while carrying collected work, which is the worst place in this codebase to raise. materialize recorded limits and sample but not the seed, so a sweep directory could not say which rows `sample: random` selected. Recorded as sample_seed alongside them. Two tests: reshard over a file torn mid-write still carries the intact rollout, and the report round-trips the seed. 98 tests. Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
…aunchers
Both ignored sweep_report.json entirely, so a manifest that pins its settings
was invisible to the two lanes that are not 03_run_single.sh. Scoped to what
each actually uses, which is not the same:
- 03_run_endpoint.sh serves nothing and runs no container, so vllm.model and
srun.container do not apply to it. What it got wrong was hardcoding
++global_aiohttp_connector_limit_per_host, ++port_range_low and
++port_range_high -- keys gym_eval_run owns -- and ignoring the manifest's
num_samples_in_parallel in favour of its own 128. It now renders gym_eval_run
extras as ++ overrides the way 03_run_single.sh does.
- 03_run_attached.sh does use a container (--container-image at :62) and names a
served model, so srun.container and vllm.model both apply. POLICY_MODEL_NAME
falls back to the manifest's vllm.model.
Two bugs in my own first attempt, both caught before commit: an apostrophe
inside ${VAR:?...} terminated the quoting, and POLICY_MODEL_NAME was still
required above the manifest block -- the same ordering mistake that made the
vllm block invisible to 03_run_sharded.sh earlier.
Verified: attached resolves both from the manifest and an env var still wins;
endpoint picks up num_samples_in_parallel: 777 and ++max_output_tokens=4096 from
a manifest and still honours NUM_SAMPLES_IN_PARALLEL from the environment.
Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
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
1 participant
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.
Add infrastructure for reward profiling.
Adds a
nemo_gym/sweepinfrastructure to collate a custom list of environments into one rewardprofiling job.
Properties:
nemotron_3.5_superREADME
has all details on how to run reward profiling.
CONTRIBUTING
has all details on how to add an environment.
RATES
has measured per-environment throughput and the GPU-hour sizing that follows.
AI generated details:
How a run is driven
Nothing else is passed. The checkpoint, both containers, the Slurm account and walltime, the
prefill/decode split and the shard count all come from the manifest; each can be overridden with an
environment variable for a one-off run.
Output is one row per source task carrying the reward distribution —
mean/std/min/max/medianover reward and tokens,num_rollouts, plus the original input undersample, so adataset can be rewritten with new pass rates without a join.
Manifest layout
Every block is named for the command that consumes it, so where a setting belongs follows from
which command reads it:
materializesweep materialize01_materialize.shsbatchsbatch03_run_single.shsrunsrun --container-image03_run_*,05_profile.shvllmvllm serve03_run_single.shgym_env_startgym env start03_run_single.shgym_eval_rungym eval run --no-serve03_run_single.shgym_eval_profilegym eval profile03_run_single.sh,05_profile.shPrecedence is uniform throughout: manifest → script environment variable → command line.
What makes sharding safe
_ng_task_indexis stamped at materialize time and Gym never rewrites it, so shard rolloutsconcatenate without renumbering. Merge deduplicates on
(_ng_task_index, _ng_rollout_index)— thesame key
--resumeuses — so a rerun shard cannot double-count. Resharding to a different N foldscollected rollouts back into the parent and snapshots it before touching any shard directory, so
changing the shard count mid-run does not discard work.
Splitting per environment keys on
_ng_task_indexagainst the per-entry range recorded in thereport, not on
agent_ref. Entries routinely share an agent — the threens_toolsentries alldispatch to
ns_tools_simple_agent— so splitting on the agent would silently merge distinctenvironments into one.
Because task indices are positional, re-materializing over already-collected rollouts is refused
rather than allowed: a changed limit renumbers every task after it, and
--resumewould then matchreal rollouts to different rows with no error.
Verification
This drives roughly 16,500 GPU-hours, so the failure worth guarding against is the silent one — a
sweep that completes and reports plausible numbers that are wrong. The following were run, not
inspected:
finishing at 2,304/2,304 rollouts, 36/36 environments profiled at 100%, 288 profiled rows from
288 source tasks, no unmapped rows and no empty labels. Every environment produced rollouts. The
only all-zero one was
lean, and its Lean compiler diagnostics confirm a genuine model failurerather than a broken lane.
noticed via
squeue, recomputed outstanding work from collected rollouts rather than exit status— 4, not 2,304 — resubmitted, and
--resumecarried the 2,300 forward to completion.partition exactly, every rollout sits in the shard owning its input, collected work survives each
reshape, merge round-trips.
scripts/debug_selftest.shreproduces this against any sweepdirectory and is read-only with respect to it.
nemo_gym/sweep.Not yet run at production scale. Everything above ran at
LIMIT_PER_ENTRY=8— 288 tasks,2,304 rollouts, one to two shards. The real sweep is 726,121 source rows → 5,808,968 rollouts
across 16 shards and 48 nodes. Three things are therefore untested: 16 concurrent shards, the
291.5 GB materialize under the current code, and tail behaviour over 5.8M rollouts.
RATES.mdcarries measured per-environment rates and the sizing that follows.