Fuse GSPO into the monolithic compiled loss (#507)#552
Open
jlamypoirier wants to merge 1 commit into
Open
Conversation
GSPO's eager segment-aggregation seam (index_add_ over a per-batch num_segments) can't run inside the single torch.compile boundary that fuses the other combinable losses. Make it a deferred combinable child: `fused_core` runs only GSPO's forward on the shared softmax and returns `(None, None, forward_state)`; `MonolithicLoss` then calls a new `finish` hook per child, where GSPO runs its existing segment seam and backward and accumulates into the shared gradient. Non-seam children keep finishing inside `fused_core` (default no-op `finish`), so the shared softmax is reused across e.g. GSPO + z-loss with no change to the existing fused path. Co-Authored-By: Claude Opus 4.8 (1M context) <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.
Claude Opus 4.8 authored this PR. Follow-up to #549 (base is that branch; retarget to
mainonce it merges).What
Lets GSPO participate as a child of the
monolithicloss type, so it shares the vocabulary softmax with another combinable loss (realistically z-loss) instead of running its own softmax pass.Why it needed special handling
The other combinable losses (CE, z-loss, distillation, GRPO) finish their loss + gradient in a single
torch.compileboundary viafused_core. GSPO can't: its eager segment-aggregation seam (index_add_over a per-batchnum_segments+ SDP/SP all-reduce + clipping) must sit between forward and backward and stay out of any compiled boundary (else it recompiles pernum_segments). So GSPO's loss and gradient are both unknown until after an eager step.Approach — deferred combinable child
fused_coreruns only its forward on the shared softmax and returns(None, None, forward_state).CombinableLossgains afinishhook (default no-op passthrough). After the compiled_monolithic_core,MonolithicLoss.forward_backwardcallschild.finish(...)per child; GSPO runs its existinggspo_segment_seam(eager) +gspo_backward_core(compiled) there, accumulating into the shared gradient buffer.fused_coreand theirfinishis the no-op default.Shared: the forward softmax (one
softmax_base+ its TP all-reduce) across GSPO and its sibling. Not shared: GSPO's backward stays a separate compiled kernel, since it depends on the eager seam output — inherent to the seam, unavoidable in the compiled path.LanguageModelGSPOLossConfigbecomes aCombinableLossConfig(mirroring GRPO), souse_tritonis inherited and the monolithic validators accept it.Tests
tests/layers/test_lm_head.py: GSPO added to the monolithic wrap-set, with newfused_gspo_lossandfused_gspo_and_z_lossconfigs (masked + unmasked, single-split — GSPO can't be split). Each is validated against the same independent reference as its per-loss equivalent. Fulltest_lm_head.py+test_lm_losses.py: 575 passed, 21 skipped (CPU).Not in scope
The Triton monolithic kernel (that path is naturally three-phase and hosts GSPO's seam without this deferral machinery).
🤖 Generated with Claude Code