Repository navigation
feat(checkpoint): declare 15 more training verifiers stateless and replayable - #4540
Closed
macandro96 wants to merge 7 commits into
Closed
macandro96 wants to merge 7 commits into
macandro96 wants to merge 7 commits into
Conversation
…nd coordination Partial-rollout checkpointing lets a training controller checkpoint Gym while rollouts are in flight and, after a crash, continue them from their last safe point instead of starting them over. This adds the core every participant builds on: - the participant control plane: phases, checkpoint-ID fencing, the lease, one export and one write per checkpoint, restore, and retire and forget for attempts the controller abandons; - episode steps, with wait and replay modes, and boundaries a restored episode continues from; - coordination, which runs each operation across Gym's servers in participant order; - storage that keeps exported state byte for byte, written and read one record at a time. In-flight work a participant cannot capture, such as a restart-only agent's, does not block prepare. It is reported as a restart, left out of the commit, and started over from its input after a crash; if nothing crashes, it keeps running. With checkpointing on, a /run whose connection dropped is never sent again, so the controller retires the attempt instead of running it twice. Checkpointing is off by default: without the checkpoint: block, no server installs any checkpoint routes or behavior. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…ith generation cuts A policy model server with checkpoint_policy becomes a participant. A checkpoint closes admission and holds every response not yet delivered, and leaves undelivered calls out together with their capture-ledger rows, so the agent makes them again after a restore. With generation cuts, each inference worker cuts its in-flight calls at a durable token prefix that the replacement call continues, at a per-backend control URL when the backend serves cuts elsewhere. Restored ledger rows continue the lineage under the next attempt. Retire stops an attempt's calls and frees its ledgers, and restored attempts the controller no longer continues are deleted without a fence. A coordinator in the main process makes a server with several uvicorn workers one participant. Co-authored-by: Anish Mahishi <amahishi@nvidia.com> Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com> Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
An inference worker that takes generation cuts stages the token prefix of each cut call, and the replacement call continues it. This adds the worker-side staging helpers, and validates a continued prefix against the replacement's admission: staging keys, source call, digest, lineage, the exact parent tokens, token count, and policy version. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…ndaries An environment server is a checkpoint participant. Each episode records a boundary after every completed step. A checkpoint parks episodes at their boundaries, or lets a replay step finish after it, and a restored episode continues from its boundary as the next attempt. Retire stops an episode, waits for its cleanup, and frees it. An episode that uses a server which cannot capture its part, such as a restart-only agent or resources server, is a restart: the seed reply says so, and the episode never holds up a checkpoint, is never exported, and starts over from its input after a crash. Episodes of a protocol that records no boundaries are restarts from the start. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…ous session hooks A resources server declares a checkpoint mode: stateless, exported, or restart-only. In exported mode its session state is exported and restored through asynchronous hooks. A checkpoint waits for requests in flight; a request that arrives during one waits for resume instead of being refused; tools called over MCP are covered; and retire frees a session once its requests have stopped. Nothing of a restart-only server is in a checkpoint. Its admission never closes, its requests never hold up prepare, and its live sessions are reported as restarts. Its seed reply says so, so the episode that seeded the session starts over from its input after a crash, and keeps running through the checkpoint if nothing crashes. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…ayable verification Code generation, competitive coding, LLM-judge equivalence, math with judge, and GenRM comparison keep no session state, so they declare themselves stateless, and their /verify may run again after a crash instead of holding up a checkpoint. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…yable Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
macandro96
marked this pull request as ready for review
October 12, 2026 18:38
yaoyu-33
added
area:environment
ananthsub
force-pushed
the
ananthsub/partial-ckpt-verifier-declarations
branch
3 times, most recently
from
October 13, 2026 17:55
fd2b07e to
7dc8e76
Compare
github-actions
Bot
added
the
sla:review-overdue
ananthsub
force-pushed
the
ananthsub/partial-ckpt-verifier-declarations
branch
3 times, most recently
from
October 14, 2026 00:27
87c700a to
8683b7a
Compare
ananthsub
force-pushed
the
ananthsub/partial-ckpt-verifier-declarations
branch
2 times, most recently
from
October 14, 2026 17:33
0d8e6ee to
9ce2b69
Compare
Contributor
|
Folded into #3887, which now declares these fifteen verifiers along with the original five. |
dims
pushed a commit
to dims/Nemo-Gym
that referenced
this pull request
Oct 15, 2026
…playable verification (NVIDIA-NeMo#3887) ## What changed and why A resources server that declares nothing is restart-only, so the rollouts that use it start over from their input after a crash. These twenty training verifiers keep no session state and their verification can run again after a crash, so they declare `checkpoint_mode = "stateless"` and `checkpoint_verify = "replay"`: - `code_gen`, `competitive_coding_challenges`, `equivalence_llm_judge`, and `math_with_judge` verify each request on its own. - `abstention`, `calendar`, `equivalence_rule`, `ether0`, `format_verification`, `instruction_following`, `inverse_if`, `jailbreak_detection`, `mcqa`, `multichallenge`, `nvarc`, `reasoning_gym`, `string_match`, `structured_outputs`, and `terminus_judge` also verify each request on its own, from the request body, with a judge model call, or with a subprocess that ends with the call (from NVIDIA-NeMo#4540, by macandro96). A replayed judge call can return a different verdict; the declaration says that repeating it is safe, not that the result is identical. - `genrm_compare` keys its cohorts by prompt, not by session. Its verification must replay: a member waits for its siblings, so waiting on it would deadlock a checkpoint while siblings are still generating. After a crash, members that had not recorded their reward re-verify and rebuild the cohort. A crash between a cohort's result and every member recording it leaves the remaining members waiting; they fail after `cohort_collection_timeout_s` and are retried from input. The declarations are class attributes, not config, because they describe what each server's code can do. The test reads them from source, because each server's dependencies live in its own environment. ## How it works ### Where this PR sits in the overall flow The highlighted part is what this PR adds. ```mermaid flowchart LR C["Controller<br/>NeMo RL, or rollout collection"] CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, forget"] K["Control plane on every server<br/>phases, lease, storage,<br/>retire: refuse, stop, free;<br/>forget"] subgraph G["One participant per Gym server"] E["Environment server<br/>episode steps"] M["Policy model<br/>held responses, generation cuts"] A["Agent<br/>sessions parked at boundaries"] R["Resources server<br/>session state"] end W["Inference worker<br/>stages cut prefixes"] D[("Checkpoint directory<br/>records, then manifest")] L[("Capture ledger<br/>retire and delete from NVIDIA-NeMo#3938, NVIDIA-NeMo#3939")] C --> CO --> K K --> E & M & A & R M --> W M --> L G --> D V["Five training verifiers<br/>stateless, replayable verify"] R --- V classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937 class V this ``` One checkpoint, a crash, and the restore, end to end: ```mermaid sequenceDiagram participant C as Controller participant G as Gym participants participant D as Checkpoint directory C->>G: prepare, in order environment, model, agent, resources Note over G: admission closes, in-flight work parks at a boundary,<br/>undelivered model responses are held G-->>C: prepared, or blockers at the deadline C->>G: commit with the episodes the controller continues G->>D: each participant writes its records, then its manifest Note over G: restored state the commit's scope leaves out is deleted C->>C: publish the checkpoint with the controller's own state C->>G: resume, in order resources, agent, model, environment Note over C,G: crash - every Gym process dies C->>G: restore the checkpoint in fresh processes, all or nothing D-->>G: records installed under attempt N+1, attempt N's capture ledger retired C->>G: resume C->>G: /run as attempt N+1 continues each episode from its boundary Note over C,G: retiring an episode, only while no checkpoint is open C->>G: retire - environment, then agent, then model and resources Note over G: each server marks the attempt, stops its work, waits, frees its state, then replies C->>G: forget, once the rollout is done for good - same order Note over G: the mark refusing the retired attempt goes away ``` ### This PR These five servers declare that verification can run again after a crash, so a checkpoint never waits for a slow judge. ```mermaid flowchart LR E["Environment server"] -->|"/seed_session"| R["code_gen, competitive_coding_challenges,<br/>equivalence_llm_judge, math_with_judge, genrm_compare"] R -->|"x-ng-checkpoint-verify: replay"| E E -->|"verify in flight at a checkpoint"| N["The checkpoint does not wait.<br/>Verify runs again after a crash."] ``` ## Where this sits in the stack This is one PR in a stack of draft PRs that re-cut partial-rollout checkpointing onto environment servers. Each PR is one commit, and its base is the branch of the PR before it, so each diff shows only that PR's change. The stack builds on the token-capture cleanup PRs NVIDIA-NeMo#3938 (capture ledger `retire` and `delete`) and NVIDIA-NeMo#3939 (complete-record `retire` and `delete`), which have merged. NVIDIA-NeMo#3882 to NVIDIA-NeMo#3885 have merged too, so NVIDIA-NeMo#3886's base is `main`. The checkpoint stack uses those operations to free the ledgers of retired attempts and to clear a dead execution's capture files before a restore. The capture-ledger cache bound (NVIDIA-NeMo#4268) and the striped lock files (NVIDIA-NeMo#3937) are independent of the stack. 1. NVIDIA-NeMo#3882 (`ananthsub/partial-ckpt-core`): feat(checkpoint): add the participant control plane, episode steps, and coordination (merged) 2. NVIDIA-NeMo#3883 (`ananthsub/partial-ckpt-policy-model`): feat(checkpoint): make policy model servers checkpoint participants with generation cuts (merged) 3. NVIDIA-NeMo#3884 (`ananthsub/partial-ckpt-model-worker-cuts`): feat(token-capture): add worker staging helpers for generation cuts (merged) 4. NVIDIA-NeMo#3885 (`ananthsub/partial-ckpt-environment`): feat(checkpoint): continue environment server episodes from their boundaries (merged) 5. NVIDIA-NeMo#3886 (`ananthsub/partial-ckpt-resources`): feat(checkpoint): add the resources server participant with asynchronous session hooks 6. NVIDIA-NeMo#3887 (`ananthsub/partial-ckpt-verifier-declarations`): feat(checkpoint): declare twenty training verifiers stateless with replayable verification **(this PR)** 7. NVIDIA-NeMo#3888 (`ananthsub/partial-ckpt-agent`): feat(checkpoint): add the agent session participant and Simple Agent continuation 8. NVIDIA-NeMo#3889 (`ananthsub/partial-ckpt-e2e`): test(checkpoint): add a process-level end-to-end suite driven by coordination 9. NVIDIA-NeMo#3893 (`ananthsub/partial-ckpt-rollout-collection`): feat(checkpoint): checkpoint evaluation runs from rollout collection 10. NVIDIA-NeMo#3903 (`ananthsub/partial-ckpt-telemetry`): feat(checkpoint): spans and metrics for partial-rollout checkpoints 11. NVIDIA-NeMo#3909 (`ananthsub/partial-ckpt-multi-worker`): feat(checkpoint): resources, agent, and environment servers with several workers 12. NVIDIA-NeMo#4197 (`ananthsub/partial-ckpt-fault-tests`): test(checkpoint): check checkpoint invariants over random sequences of operations Partial-rollout checkpointing lets a training controller, such as NeMo RL, checkpoint Gym while rollouts are in flight and, after a crash, continue those rollouts from their last safe point instead of starting them over. Checkpointing is off by default; with the `checkpoint:` block unset, no server installs any checkpoint routes or behavior. ## Relationship to the old stack (NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946) None of NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946 declared these servers. The declarations come from NVIDIA-NeMo#3349 (durable turn-level recovery), where they were config on the agent; here they are class attributes on the resources servers. ## Issue No tracking issue exists for this re-cut. The design and the mapping from the old stack are described in this PR series, and the old stack's PRs (NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946) carry the original discussion. ## Validation Run on this branch, on the current `main`: - `tests/unit_tests/test_checkpoint_server_declarations.py`: each of the twenty servers declares `stateless` and `replay`, read from its source. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/code_gen/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/competitive_coding_challenges/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/equivalence_llm_judge/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/genrm_compare/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/math_with_judge/tests`: not run in this environment. It needs the server's `math-verify==0.8.0` requirement, which this environment does not install. It passed on an earlier version of this branch with that requirement installed. - The full core unit suite (`tests/unit_tests`, eight processes, without `test_opensandbox_compose.py`, which needs the optional `opensandbox` package): passed except a sandbox retry test, which fails the same way on `main` in this development environment. - `pre-commit run --files <files changed by this PR>`: all hooks passed, and no hook modified a file. - Every commit carries a `Signed-off-by` line. ## Rollout evidence - These are declarations only: with checkpointing off, nothing changes for these servers. With it on, they no longer force their rollouts to be retired at each checkpoint. - A checkpointed training run on these environments with a real model: **pending**. ## Compatibility - Nothing changes unless the global `checkpoint:` block sets `enabled: true`. With it on, these five servers no longer block checkpoints, and their `/verify` may run again after a crash. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com> Co-authored-by: Anish Mahishi <amahishi@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
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.
When partial-rollout checkpointing is enabled, these resource servers inherit
restart_only/wait, causing their rollouts to restart from input after a crash. Their verification uses request-local data or independent judge/evaluation calls and has no per-session resource state to export.Add
checkpoint_mode = "stateless"andcheckpoint_verify = "replay"to 15 additional resource servers:abstention,calendar,equivalence_rule,ether0,format_verificationinstruction_following,inverse_if,jailbreak_detection,mcqa,multichallengenvarc,reasoning_gym,string_match,structured_outputs,terminus_judgeExtend the existing declaration test to cover these servers alongside the five already declared in #3887. The diff contains only the class declarations and the test's server mapping.
Dependency and compatibility
Targets
ananthsub/partial-ckpt-verifier-declarations, the head branch of #3887. Merge after that dependency, or incorporate these additions into it. No separate tracking issue is needed for this extension.With checkpointing disabled, verifier behavior is unchanged. With it enabled, these servers can preserve rollout progress without exporting resource-session payloads, and interrupted
/verifycalls may execute again after restore. Replayed judge calls may produce a different answer; the declaration guarantees replay safety, not deterministic judge output.Validation
Against this diff on Python 3.13.14:
python -m pytest -q tests/unit_tests/test_checkpoint_server_declarations.py tests/unit_tests/test_checkpoint_resources.py: 65 passed.python -m pytest -q -o addopts='' -o log_cli=false resources_servers/string_match/tests/test_app.py: 10 passed.pre-commit run --all-files: passed; no unrelated hook changes remain.git diff --check: passed.Signed-off-bytrailer.Remaining validation / rollout evidence
Draft pending all changed servers' isolated
gym env test +entrypoint=resources_servers/<server> +should_validate_data=truechecks, repository CI, and representative checkpointed model rollouts on this exact branch. The first 14 declarations were used in our integrated 88-node training checkout, but that downstream integration is not a substitute for validating this upstream PR head.string_matchhas unit-test coverage here and no new model rollout evidence yet.