Emit loss_fn_outputs with logprobs for RL losses in forward_backward#1047
Merged
tyler-griggs merged 2 commits intomainfrom Feb 7, 2026
Merged
Emit loss_fn_outputs with logprobs for RL losses in forward_backward#1047tyler-griggs merged 2 commits intomainfrom
tyler-griggs merged 2 commits intomainfrom
Conversation
The RL branch of _forward_backward_micro() now packages per-sample logprobs into loss_fn_outputs, matching the cross_entropy branch. This is required by tinker-cookbook's train.py which reads logprobs from forward_backward results. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
tyler-griggs
force-pushed
the
tyler/rl-loss-fn-outputs
branch
from
4aef853 to
6f30fe7
Compare
tyler-griggs
force-pushed
the
tyler/rl-loss-fn-outputs
branch
from
ea24a54 to
2cbbd9a
Compare
tyler-griggs
force-pushed
the
tyler/rl-loss-fn-outputs
branch
from
2cbbd9a to
e153aba
Compare
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request correctly adds loss_fn_outputs with log probabilities to the RL training path in _forward_backward_micro, aligning it with the existing cross_entropy path. This change is crucial for downstream consumers that rely on these outputs. The accompanying test updates are thorough and validate the new functionality. I've provided one suggestion to refactor the newly added code block for improved efficiency and maintainability by reducing code duplication.
tyler-griggs
force-pushed
the
tyler/rl-loss-fn-outputs
branch
from
e153aba to
e050fae
Compare
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
tyler-griggs
force-pushed
the
tyler/rl-loss-fn-outputs
branch
from
e050fae to
6e17a59
Compare
tyler-griggs
force-pushed
the
tyler/rl-loss-fn-outputs
branch
from
6e17a59 to
1327e67
Compare
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
tyler-griggs
force-pushed
the
tyler/rl-loss-fn-outputs
branch
from
1327e67 to
9bac20f
Compare
erictang000
added a commit
that referenced
this pull request
Feb 10, 2026
relax restriction in `reduce_metrics` to just continue if metrics are not all numbers rather then erroring out, and pop loss_fn_outputs to avoid polluting the logs. Nightly E2E CI (and all training scripts) were failing after #1047. Also deletes claude folder introduced in #999. <!-- devin-review-badge-begin --> --- <a href="https://app.devin.ai/review/novasky-ai/skyrl/pull/1059" target="_blank"> <picture> <source media="(prefers-color-scheme: dark)" srcset="https://static.devin.ai/assets/gh-open-in-devin-review-dark.svg?v=1"> <img src="https://static.devin.ai/assets/gh-open-in-devin-review-light.svg?v=1" alt="Open with Devin"> </picture> </a> <!-- devin-review-badge-end -->
Collaborator
|
cc: @tyler-griggs can you be a little more careful when merging to main without review? this pr broke the main training loop on main, would have been caught by just running the gsm8k example script fixed in: #1059 |
Collaborator
|
also a case to be made for better testing/alerting on nightly CI |
erictang000
added a commit
that referenced
this pull request
Feb 13, 2026
…tinker RL to work (#1102) applies the same changes from #1047 but to the megatron backend. Double PR on both skyrl-train and skyrl for now. <!-- devin-review-badge-begin --> --- <a href="https://app.devin.ai/review/novasky-ai/skyrl/pull/1102" target="_blank"> <picture> <source media="(prefers-color-scheme: dark)" srcset="https://static.devin.ai/assets/gh-open-in-devin-review-dark.svg?v=1"> <img src="https://static.devin.ai/assets/gh-open-in-devin-review-light.svg?v=1" alt="Open with Devin"> </picture> </a> <!-- devin-review-badge-end -->
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
_forward_backward_micro()now packages per-sample logprobs intoloss_fn_outputs, matching the existingcross_entropybranch_training_logprobs_from_fwd_bwd()which readsfwd_bwd_result.loss_fn_outputs[i]["logprobs"]for KL diagnosticscode_rlandmath_rlrecipes crash withKeyError: 'loss_fn_outputs'