Fuse the GSPO segment seam into compiled blocks (#507)#554
Open
jlamypoirier wants to merge 1 commit into
Open
Conversation
`gspo_segment_seam` ran ~13 small elementwise/gather/reduction ops as separate eager launches around the `index_add_` aggregation. Split the fusable work into two `@torch.compile` blocks — pre-aggregation per-token weights and post-aggregation loss/gradient — around the eager `index_add_` + SDP/SP all-reduce middle, which stays eager (the symbolic `num_segments` would trigger per-value recompiles, and the all-reduces are comm). The per-step-varying `/ divisor` and `grad_output` scaling stay eager too so they never specialize the compiled graphs. Numerically equivalent to the prior eager seam across standalone GSPO, the monolithic `finish` path, and the distributed SDP/SP 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.
Authored by Claude Opus 4.8 (1M context), reviewed by @jlamypoirier.
Follow-up to the monolithic head-loss work (#507), stacked on #549.
What
gspo_segment_seam— the eager glue between GSPO's compiled forward and backward cores — ran ~13 small elementwise/gather/reduction ops as separate launches around theindex_add_segment aggregation. This splits the fusable work into two@torch.compileblocks:_gspo_segment_weights): the per-token contributions feedingindex_add_;_gspo_segment_loss): from the reduced per-segment sums to the loss, thenew_logprobsmetric, and the unscaled per-token backward coefficient.The middle stays eager: the
index_add_(its symbolicnum_segmentswould trigger per-value recompiles) and the SDP/SP all-reduces (comm). The per-step-varying/ divisorandgrad_outputscaling also stay eager, so they never specialize the compiled graphs (epsilon_*are fixed per run, so they don't).Both consumers of the seam benefit — standalone GSPO and the monolithic
finishpath — and the eventual Triton monolithic loss reuses the same seam.Correctness
Numerically equivalent to the prior eager seam. Verified against the existing GSPO cases (the equivalence oracle): all
test_gspo_loss/test_lm_headGSPO configs (standalone, masked, uneven-documents, fused-monolithic) and the distributedworld_size=2SDP/SP path pass. A serial run across varied segment counts shows no recompiles beyond the initial compile.🤖 Generated with Claude Code