Repository navigation
[Performance] Rollout collection _preprocess_rows_from_config avoid deepcopy - #2816
Merged
Merged
Conversation
Signed-off-by: Brian Yu <bxyu@nvidia.com>
Signed-off-by: Brian Yu <bxyu@nvidia.com>
…ollout-collection-copy
Signed-off-by: Brian Yu <bxyu@nvidia.com>
Signed-off-by: Brian Yu <bxyu@nvidia.com>
Contributor
Author
|
/ok to test 5366e98 |
jkyi-nvidia
added a commit
that referenced
this pull request
Aug 28, 2026
Measured at full scale on the whole manifest:
--no-expand EXPAND=1
rows 726,121 5,808,968
size 34 GB 291.5 GB
wall 682 s 2,087 s
Gym expands num_repeats itself at collection time while honouring the task
index we stamp, and that expansion is what PR #2816 took from 8 min to 1 min,
so pre-expanding buys little and costs 8x the disk.
The reason it matters beyond disk is prefix caching. A task's repeats have
identical prompts. Unexpanded they all land in one shard and repeats 2-8 hit
the vLLM prefix cache; expanded, they are consecutive lines, so round-robin
dealing scatters them across 8 jobs and every repeat is a cold prefill. This
only bites once sharded -- a single unsharded job caches expanded repeats fine.
run_sharded.sh now watches rather than waits: a shard whose job dies with work
outstanding is resubmitted against the same directory, where --resume picks up
what it already collected. MAX_ROUNDS (4) bounds that so a shard failing for a
permanent reason cannot loop. Completeness is judged per task for unexpanded
inputs and per (task, rollout) for expanded ones -- verified both ways: 36
outstanding for an untouched unexpanded shard, 0 once collected, 288 for the
expanded equivalent.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Junkeun Yi <jkyi@nvidia.com>
muskank7
pushed a commit
to muskank7/Gym
that referenced
this pull request
Sep 30, 2026
…NeMo#3672) ## Expected With `num_repeats_add_seed=True`, each repeated rollout should receive its own seed even when the input already includes `responses_create_params.metadata`. Existing metadata and caller-owned input should be preserved. ## Actual The copy optimization in NVIDIA-NeMo#2816 copies the response parameters but shares their nested metadata. Updating the next repeat’s seed overwrites earlier repeats: three repeats receive `[2, 2, 2]` instead of `[0, 1, 2]`, and in-memory input is mutated. ## Fix Copy the metadata dictionary before assigning each rollout’s seed. Keep shallow copying for the rest of the task. Validation: The regression fails before the fix and passes after; 240 adjacent tests, full core/sandbox checks (95.16% coverage), and all-file pre-commit pass. Model rollouts: N/A for this input-materialization change. <!-- gym-verified-bug-hunter:rollout-repeated-seed-metadata-aliasing --> Signed-off-by: yaoyu-33 <54727607+yaoyu-33@users.noreply.github.com>
This branch was previously deployed
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.
What does this PR do?
On 170k input samples, this reduces the processing time from 8min -> 1min
Checklist
pre-commit run --all-files) (so CI lint/format/copyright pass).git commit -s) (so the DCO check passes).