Skip to content

[megatron] fix: allow sequence-mean loss with per-token normalization - #7915

Merged
wuxibin89 merged 1 commit into
verl-project:mainfrom
Zhikaiiii:fix/seqmean-per-token-7914
Sep 18, 2026
Merged

wuxibin89 merged 1 commit into
verl-project:mainfrom
Zhikaiiii:fix/seqmean-per-token-7914

Conversation

@Zhikaiiii

@Zhikaiiii Zhikaiiii commented Sep 18, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes #7914 by removing the seq-mean-token-mean rejection in the Megatron per-token loss callback. The guard rejects CP=1 when calculate_per_token_loss=True, and also blocks static-CP THD training even though the forward path reconstructs full-sequence outputs before PPO loss computation.

Checklist Before Starting

  • Checked PRs referencing #7914, sequence-mean PRs, and static-CP PRs. No open PR addressing this guard was found; related loss-masking and algorithm changes do not remove this restriction.
  • PR title follows the required format.

API and Usage Example

No API changes. loss_agg_mode=seq-mean-token-mean no longer triggers this rejection when calculate_per_token_loss=True.

Design & Code Changes

Deletes only the aggregation-mode guard and its explanatory comment from postprocess_micro_batch_func (13 lines). The existing normalization logic and other validation remain intact.

AI assistance (OpenAI Codex) was used for analysis, implementation, checks, and PR preparation under the issue author's direction.

Checklist Before Submitting

Remove the overly restrictive aggregation guard described in verl-project#7914.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: Zhikaiiii <1658973216@qq.com>
@wuxibin89
wuxibin89 merged commit 0eb309f into verl-project:main Sep 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[megatron] seq-mean-token-mean guard rejects CP=1 and appears overly restrictive for reconstructed THD outputs

2 participants