Policy metrics in the compiled monolithic loss (#507)#565
Open
jlamypoirier wants to merge 1 commit into
Open
Conversation
Extend the monolithic loss to emit the policy-gradient metric family (GRPO and
GSPO) from the shared softmax on the compiled path:
- Compiled GSPO metrics: `finish` reduces the family from the shared softmax and
the segment seam via `gspo_metrics_core`, matching how GRPO already reuses it.
- Collapse `PolicyMetricsLevel` to `{none, basic}` (entropy folded into `basic`,
now near-free) plus an `auto` default that resolves to `basic` when
`pipeline_parallel == 1` and `none` otherwise (an explicit `basic` still
requires it). `none` stays zero-cost.
- The metric cores take a precomputed `entropy_per_token`, unifying the softmax
and (later) kernel entropy sources.
- Raise the GSPO loss-comparison floor to de-flake its float32 tolerance.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jlamypoirier
force-pushed
the
jlp_monolithic_metrics
branch
from
32c3943 to
2368650
Compare
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.
Claude Opus 4.8 note: authored with Claude Code.
Intermediate PR in the monolithic-loss stack:
jlp_monolithic_triton(#558) → this →jlp_monolithic_head_loss(#549, dev). Contains the path-agnostic + compiled-path policy-metrics work; the triton kernel and triton-metric wiring stay in #558 on top of this.What this adds
finishreduces the policy-metric family (importance ratio, KL, clip fraction, advantage stats, entropy) from the shared softmax and the segment seam viagspo_metrics_core, matching how GRPO already reuses the softmax — no second softmax pass.PolicyMetricsLevelreworked to{none, basic, auto}, defaultauto:basicfolds entropy in (now near-free) — no separatewith_entropy.autoresolves tobasicwhenpipeline_parallel == 1andnoneotherwise (the metrics need the full-logits pass, which pipeline parallelism splits). An explicitbasicstill assertspipeline_parallel == 1.nonestays zero-cost.entropy_per_token, unifying the compiled/standalone (from the softmax) and — in Triton monolithic head-loss kernel (#507) #558 — the triton (from the kernel) entropy sources.loss_min_threshold); its per-segment geometric-mean reduction is noisier in float32 than a per-token sum, so the fused-vs-reference scalar diff sat at the1e-6floor ~1/300 draws.Validated on CPU (
test_lm_head+test_lm_losses, 577 passed). The equivalent code was GPU-validated as part of the combined branch before the stack split.🤖 Generated with Claude Code