Repository navigation
Log policy_loss in GMPOTrainer and RLOOTrainer on every run - #7033
Closed
behroozazarkhalili wants to merge 9 commits into
Closed
behroozazarkhalili wants to merge 9 commits into
behroozazarkhalili wants to merge 9 commits into
Conversation
GRPOTrainer records policy_loss, but the two trainers that own their own loss code never did. GMPOTrainer overrides _compute_loss and dropped the metric; RLOOTrainer never had it. A run of either trainer therefore has no way to separate the surrogate objective from the auxiliary terms. GMPO captures the value before the gradient accumulation rescale, after the KL term is folded in, so with beta != 0 it carries that too. RLOO captures the clipped surrogate before the Mixture-of-Experts auxiliary loss is added; RLOO applies its KL penalty to the reward, so the metric carries no KL term. The tests pin the capture points rather than the mere presence of the key. GMPO trains with gradient_accumulation_steps=2: GRPOTrainer disables the HF Trainer's accumulation rescale and divides the returned loss itself, so the reported loss equals the mean captured value while a capture after the rescale would report half of it. RLOO trains a tiny MoE model, so the reported loss must equal policy_loss + router_aux_loss_coef * aux_loss; capturing after the addition breaks the identity. Both tests also check the eval path. The RLOO metrics list in the docs gains the new entry. The gspo_token override is handled by #6856, and the GRPO gating by #6861. Resolves #7005
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
… one The minimum-versions lane failed on the RLOO identity by 8e-7: transformers 4.56.2 rounds the logged train loss to 4 decimals (trainer.py:3211) while 5.x logs it unrounded, so an oracle that reads `loss` from log_history holds on one version and not the other. Both tests now record every loss `_compute_loss` returns, the pattern already used in test_grpo_trainer.py, and compare the metric with that. GMPO sums the two returned half-losses of each accumulation step; RLOO compares one returned loss per step. Eval compares against the mean of the returned eval losses, so the identity no longer depends on how the HF loop weights eval_loss across eval batches.
…ints Both tests accepted implementations they were meant to rule out. The GMPO test ran only with beta=0.0, so a capture taken before the KL fold passed it; it now runs a beta=0.1 row too, where the KL term inside the captured value is the difference. The RLOO test ran one iteration per generation, where the ratio is exactly 1 and the surrogate is identically zero, so a hardcoded 0.0 satisfied the identity; it now runs two iterations, asserts a non-zero value, and adds a dense-model row where policy_loss must equal the returned loss with no auxiliary term.
Collaborator
Author
|
Closing this one myself. It came from reading the code, not from a failure in a real training run, and per the discussion on #6862 that isn't the bar for this repo. If someone hits this in practice, the branch is still here and I'm happy to reopen it with a reproduction. |
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
1 participant
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.
GMPO and RLOO did not log
policy_loss, leaving no separate view of their policy objective. I added the metric at each trainer's intended capture point, including KL in GMPO and excluding RLOO's auxiliary loss. Training and evaluation tests check the metric against returned losses and pass here, while the tests fail on main because the metric is missing.Resolves #7005.
Note
Low Risk
Observability-only metric logging and tests; training math and returned loss are unchanged aside from new log keys.
Overview
GMPO and RLOO now emit a
policy_lossmetric on every train/eval step, matchingGRPOTrainerso logs separate the core surrogate objective from extra loss terms (fixes #7005).In
GMPOTrainer._compute_loss,policy_lossis recorded after optional KL is added to the sequence loss and before dividing by gradient accumulation steps; whenbeta != 0, the logged value includes that KL term.In
RLOOTrainer._compute_loss,policy_lossis the clipped surrogate mean before MoEaux_lossis added; RLOO keeps KL in the reward, so this metric has no KL component.docs/source/rloo_trainer.mddocumentspolicy_loss. New tests assert capture timing by comparing logged values to losses returned from_compute_loss(GMPO: accumulation +beta; RLOO: dense vs MoEaux_loss).Reviewed by Cursor Bugbot for commit aaa0d62. Bugbot is set up for automated code reviews on this repo. Configure here.