rl: add durable rollout bank for completed groups - #6352
Conversation
Persist completed RolloutGroups to a write-through, per-group-fsync'd ledger the instant they assemble, so completed work survives a SIGKILL (the SLURM time limit) and is restored at resume instead of regenerated. - types/: extract the rollout data models (Rollout, TokenRollout, RolloutGroup, Rollouts, GroupedRollouts, AgentBaseModel, and the EnvId/GroupsPerEnv/GroupQueuesPerEnv aliases) out of agent/api.py into a megatron.rl.types package; api.py re-exports them so existing imports keep working. RolloutGroup gains a uid; TokenRollout gains advantage_override; Rollout.reward widens to float | None. - rollout_bank.py: the durable bank (append/restore/mark_consumed/ checkpoint), offset-indexed binary sidecars for token/logprob/mask arrays, per-record checksums, torn-write recovery, and compaction. - api.py: the pipeline writes each completed group through the bank on assembly (injected, never a global). - rl_utils.py: bank singleton + restore/injection in get_environment_rollouts, weight-balanced per env across steps. - weighted_multi_task.py: env_group_targets + num_groups_per_env override so restored groups keep the per-env weight split. - checkpointing.py: compact at the checkpoint boundary (sync path and async durability callback) so compacted-through T tracks the model. - arguments.py: --rl-durable-rollout-bank / --rl-rollout-bank-dir / --rl-rollout-bank-max-bytes. No-op when unset. Ported from jalbericiola#10. That PR was authored on a fork base whose in-flight tracker and speculative-rollout scaffolding are not on NVIDIA main, so the bank was reapplied against main's episode-based pipeline rather than cherry-picked. Tests: unit + pipeline coverage for the bank round-trip, durability, compaction, restore balancing, and the type re-export (tests/unit_tests/ rl/test_rollout_bank.py, test_grouped_rollouts.py). 101 pass locally. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Reformat tests/unit_tests/rl/test_rollout_bank.py to satisfy the CI linting job's black check (--skip-magic-trailing-comma). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Laura Dang <laurad@nvidia.com>
|
Hey, your first commit (8c9a310) needs to be signed with the DCO still. |
tdene
left a comment
tdene
left a comment
There was a problem hiding this comment.
Left a couple of important comments.
Another one in the comments here: you are also recovering rollouts that were trained on since the last saved checkpoint.
The generator should not be allowed to run until we catch up. If we saved on ckpt 20, and crashed on ckpt 38, then ckpts 21 -> 37 should train without the generator ever engaging.
Look at #4125 as well. Perhaps this PR should be rebased on top of 4125 so that it can easily load data into the iterator without turning on the generator.
| for group in rollouts: | ||
| bank.mark_consumed(group.uid, args.curr_iteration) |
There was a problem hiding this comment.
This repeatedly opens and flushes the same file in a loop. This is going to be very slow compared to just opening the file once at the head of the loop and flushing + closing at the end.
| "next checkpoint. Consider a shorter checkpoint interval or lower --rl-generation-lag.", | ||
| self._bytes_written, self.max_bytes, | ||
| ) | ||
| self._warned_over_cap = True |
There was a problem hiding this comment.
This never gets reset after compaction.
| self._ledger_f, self._tok_f, self._lp_f, self._mask_f) = saved | ||
|
|
||
| def _maybe_warn_over_cap(self) -> None: | ||
| if self.max_bytes <= 0 or self._bytes_written <= self.max_bytes: |
There was a problem hiding this comment.
_bytes_written is also never changed by compaction, even though compaction will change its measure.
| _FORMAT_VERSION = 1 | ||
| _MANIFEST = "MANIFEST.json" | ||
| _LEDGER = "ledger.log" | ||
| _CONSUMED = "consumed.log" |
There was a problem hiding this comment.
Should this also be compacted so that it does not grow without bound?
| if not env_id and not config.evaluation_only: | ||
| raise ValueError( | ||
| f"Active agent {i} ({type(a).__name__}) has no env_id; it is " | ||
| f"required to weight-balance restored rollout-bank groups by env. " | ||
| f"Set env_id for every non-evaluation agent in the environment config." |
There was a problem hiding this comment.
Do single-environment setups currently set env_id? If so, they would fail this check, when they shouldn't have to because single-env routing is trivial.
| token_typed = member_type == "TokenRollout" and all( | ||
| "trajectory" in m and m["trajectory"] and isinstance(m["trajectory"][0], list) | ||
| for m in members | ||
| ) |
There was a problem hiding this comment.
You can do isinstance(group.rollouts[0], TokenRollout) directly.
| # Streaming with rebalance: quota-route the balanced stream to the | ||
| # per-env residual, buffering over-quota groups for a later step. | ||
| # Otherwise pull n_fresh groups in completion order. | ||
| if args.rl_partial_rollouts and residual_per_env is not None: | ||
| fresh = _pull_fresh_balanced( | ||
| loop, rollout_generator, residual_per_env, | ||
| runtime_state.fresh_overflow, | ||
| ) | ||
| else: | ||
| fresh = _pull_fresh_inorder(loop, rollout_generator, n_fresh) | ||
| # Restored groups first, then freshly generated ones. | ||
| rollouts = list(inject) + fresh |
There was a problem hiding this comment.
There's a problem with this logic.
Let's say you're in a 2-env situation. You want 50% from env A and 50% from env B. You bank rollouts, and you end up with 100 env A rollouts and 2 env B rollouts.
This logic ensures that you consume a balanced amount, but on the production side you will continue to produce env A despite having way too many rollouts from it.
I don't think this should be fixed in this PR. I think this is a pervasive prior issue that can happen even without banking. But I think it should be documented.
There was a problem hiding this comment.
Actually, I'm having second thoughts.
This logic will reliably result in the rollout bank never getting drained. The generator continues to generate rollouts in exactly the same proportion and amount as this logic drains.
I am pretty sure that any overflow is permanent. That is a problem that should be fixed in this PR.
There was a problem hiding this comment.
I also don't think this env-balancing logic should exist on the consumer side.
It already exists on the producer side. The rollout bank is a producer as well. It should plug into WeightedMultiTask.
That lets you cut out _pull_fresh_balanced, fresh_overflow, _plan_restore_injection, _empty_rollout_generator, _env_targets, and env_group_targets, and num_groups_per_env.
Something like
correct_agent_groups = list(agent_groups)
for i, (eid, total) in enumerate(zip(env_ids, agent_groups)):
cached = min(total, len(restored.get(eid, [])))
correct_agent_groups[i] = total - cached
for _ in range(cached):
yield restored[eid].popleft()
| try: | ||
| with open(self._manifest_path) as f: | ||
| manifest = json.load(f) | ||
| except (FileNotFoundError, json.JSONDecodeError): |
There was a problem hiding this comment.
FileNotFoundError and json.JSONDecodeError should not be treated the same. If the file is not found, the ledger doesn't exist. If we have a json.JSONDecodeError, the user should be told that there's a malformatted ledger and an error should be raised so that the user can try to manually recover it.
| members = [cls.model_validate(m) for m in group_dict["rollouts"]] | ||
| group = RolloutGroup( | ||
| rollouts=members, | ||
| batch_id=group_dict.get("batch_id", 0), |
There was a problem hiding this comment.
This batch_id will be stale and will collide.
Each new run starts counting from 0, even if it's a continuation run.
I think the correct fix is to fix this behavior, and make runs start counting from current_iteration instead.
This will affect B/B.
| self._collection_iter = None | ||
| self._seg_dir = None | ||
| self.set_collection(iteration) |
There was a problem hiding this comment.
I'm pretty sure this will be off-by-1 for sync saves, and may be very off for async saves.
But also, I don't think these 3 lines do anything? Right after this happens, get_environment_rollouts runs bank.set_collection(args.curr_iteration).
3 participants
Persist completed RolloutGroups to a write-through, per-group-fsync'd ledger the instant they assemble, so completed work survives a SIGKILL (the SLURM time limit) and is restored at resume instead of regenerated.
Ported from jalbericiola#10. That PR was authored on a fork base whose in-flight tracker and speculative-rollout scaffolding are not on NVIDIA main, so the bank was reapplied against main's episode-based pipeline rather than cherry-picked.
Tests: unit + pipeline coverage for the bank round-trip, durability, compaction, restore balancing, and the type re-export (tests/unit_tests/ rl/test_rollout_bank.py, test_grouped_rollouts.py). 101 pass locally.
What does this PR do?
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks