Unify policy-loss metrics config and add GSPO metrics#555
Merged
Conversation
Rename `GRPOMetricsLevel` to `PolicyLossMetrics` and lift the `metrics` field from the GRPO config to the shared `LanguageModelPolicyGradientLossConfig`, so GSPO gains the same diagnostic levels (none/basic/with_entropy). The `pipeline_parallel == 1` guard moves to the shared `LanguageModelPolicyGradientLoss.__init__`. Add `compute_gspo_metrics` (segment-level analog of `compute_grpo_metrics`) and GSPO's `_register_extra_metrics`/`get_loss_definitions`. GSPO clips per segment, so its statistics (ratio, KL, clip fraction) are per-segment; they are accumulated as token-weighted sums (weight `mask / num_labels_in_seq`, which sums to 1 per document across SDP/SP ranks) so they partition and reduce exactly like the per-token loss. Metric names match GRPO where the meaning aligns (`_num_segments` replaces `_num_tokens`). Tested via `compute_gspo_metrics` against an independent plain-Python segment-loop reference (`reference_gspo_metrics` + `test_gspo_metrics`). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Log GSPO num_tokens (labeled-token count) so mean response length and ratio variance are derivable, matching GRPO. - Extract the shared GRPO/GSPO metric LossDef list into LanguageModelPolicyGradientLoss._policy_metric_definitions(*extra); GSPO passes its num_segments def. Order-preserving. - Add a GSPO-metrics subtest to _run_lm_loss_distributed (and to the test_lm_loss_distributed loss_type list) so compute_gspo_metrics' vocab-parallel group all-reduce path is exercised and reported. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The two metric paths differed only in how the importance ratio is formed (per-token for GRPO, per-segment geometric mean for GSPO). Collapse the duplication: - One `PolicyMetrics` tuple replaces `GRPOMetrics`/`GSPOMetrics`. - `_policy_metrics_log_ratio` holds the shared softmax -> log-ratio -> entropy pass; `_policy_metrics_reduce` holds the shared weighted-sum reduction, parameterized by the ratio-variance weight (token mask vs document weight). - One `_register_policy_metrics` replaces the two near-identical registrations. `compute_grpo_metrics` / `compute_gspo_metrics` now only derive the per-token `(ratio, effective_log_ratio)` -- their actual difference. GSPO's token-level `exp`-after-gather is bit-identical to the previous per-segment-then-broadcast. Behavior-preserving: the independent test references (repackaged into `PolicyMetrics`, math unchanged) still pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Rename enum `PolicyLossMetrics` -> `PolicyMetricsLevel` to disambiguate from the `PolicyMetrics` value tuple (restores the `Level` suffix). - `compute_gspo_metrics` label-count parameter `num_labels_in_seq` -> `label_counts` to match the sibling `compute_grpo_metrics` and the kwarg source. - Tighten `losses: dict | None` -> `dict` in both `_register_extra_metrics` (only reached under `losses is not None`). - Name the `num_segments = 4` local in the distributed GSPO-metrics subtest. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.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
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.
Summary
Extracted from #553 as an independent, wire-change-free PR (pure Fast-LLM). Gives GSPO the diagnostic metrics GRPO already has, and unifies the metrics config so both losses share it.
GRPOMetricsLevel→PolicyLossMetricsand lift themetricsfield (none/basic/with_entropy) from the GRPO config up to the sharedLanguageModelPolicyGradientLossConfig, sotype: gspoandtype: grpoboth carry it. Thepipeline_parallel == 1guard (both metric paths use a second full-vocab softmax) moves to the shared base.compute_gspo_metricscomputes segment-level statistics (GSPO clips per segment), reusing the segment aggregation from the loss kernel — segment clipped-ratio fraction, ratio mean + sum/squared-sum, KL(new‖old), advantage mean/max/min, segment/token counts, and (with_entropy) token-level entropy. Metric suffixes mirror GRPO's where the meaning lines up so runs are comparable (noting GSPO's are segment-level, GRPO's token-level). The extra softmax pass is skipped entirely whenmetrics: none.num_tokensparity: GSPO now also logs the labeled-token count, so mean response length and ratio variance are derivable as for GRPO.Tests
test_gspo_metrics/test_grpo_metricssingle-process (reference = plain-Python segment loop), plus a GSPO-metrics subtest added to_run_lm_loss_distributed(and thetest_lm_loss_distributedreporting list) so the vocab-parallelgroupall-reduce path is exercised.type: gspo/grpoconfigs validate for all three metric levels. The distributedgrouppath runs under torchrun in CI.🤖 Generated with Claude Code — authored by Claude Opus 4.8.