docs: add policy lifecycle guide for custom training loops - #2644
Open
lonexreb wants to merge 2 commits into
Open
docs: add policy lifecycle guide for custom training loops#2644lonexreb wants to merge 2 commits into
lonexreb wants to merge 2 commits into
Conversation
github-actions
Bot
added
Documentation
lonexreb
added a commit
to lonexreb/RL
that referenced
this pull request
May 30, 2026
The state table and the two preparation-method sections previously
described the optimizer movement as unconditional ("frees optimizer
state", "moves optimizer back to GPU"). In practice the optimizer
follows the weights only when `offload_optimizer_for_logprob=true`
(or, for training, when `is_generation_colocated` is set as well),
per the implementations in dtensor_policy_worker.py and
megatron_policy_worker.py.
Reword the state-table column header, prepare_for_training and
prepare_for_lp_inference sections, and the matching interfaces.py
docstrings so the conditional behavior is explicit and config-driven.
Caught in self-review of NVIDIA-NeMo#2644.
Signed-off-by: Shubhankar Tripathy <reach2shubhankar@gmail.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
Closes NVIDIA-NeMo#1141. Adds `docs/design-docs/policy-lifecycle.md` documenting the policy state machine (training / logprob-inference / offloaded), the contract for `prepare_for_training()` and `prepare_for_lp_inference()`, the opaque CUDA "illegal memory access" failure mode users hit when they skip a `prepare_*` call in a custom training loop, and a minimal reference loop that calls everything in the right order. Also fills in the previously empty docstrings on the abstract methods in `nemo_rl/models/policy/interfaces.py` with a one-paragraph summary + a pointer to the new doc, so `help(PolicyInterface)` is immediately useful when someone reads the source. The doc is linked from the `Design Docs` toctree in `docs/index.md`. This is the documentation half of the request in NVIDIA-NeMo#1141; the "more helpful error" half is covered by PRs NVIDIA-NeMo#2392 / NVIDIA-NeMo#2393, which add runtime guards that resolve to `prepare_for_lp_inference()` / `prepare_for_training()` in the failure message. Signed-off-by: Shubhankar Tripathy <reach2shubhankar@gmail.com>
The state table and the two preparation-method sections previously
described the optimizer movement as unconditional ("frees optimizer
state", "moves optimizer back to GPU"). In practice the optimizer
follows the weights only when `offload_optimizer_for_logprob=true`
(or, for training, when `is_generation_colocated` is set as well),
per the implementations in dtensor_policy_worker.py and
megatron_policy_worker.py.
Reword the state-table column header, prepare_for_training and
prepare_for_lp_inference sections, and the matching interfaces.py
docstrings so the conditional behavior is explicit and config-driven.
Caught in self-review of NVIDIA-NeMo#2644.
Signed-off-by: Shubhankar Tripathy <reach2shubhankar@gmail.com>
Contributor
Author
|
/ok to test a51e67a |
lonexreb
force-pushed
the
docs/1141-policy-lifecycle-guide
branch
from
August 30, 2026 07:03
8710921 to
a51e67a
Compare
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
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
Closes #1141.
The issue asked for either documentation or a more helpful error when a user writes a custom training loop and forgets to call `prepare_for_lp_inference()` / `prepare_for_training()` before a GPU-bound API call. The typical symptom — an opaque `RuntimeError: CUDA error: an illegal memory access` deep in a model forward pass — gives no hint that the cause is a missing `prepare_*` call.
This PR is the documentation half of that ask. The runtime-error half is covered by PRs #2392 / #2393 (offload guards) which raise an actionable `RuntimeError` naming the missing `prepare_*` call.
What changed
Test plan
Cross-refs