From dab00f5d412a38046165d7c0aede365e177bfd44 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 6 Jul 2026 11:19:05 -0400 Subject: [PATCH 1/2] Fuse GSPO into the monolithic compiled loss (#507) 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) --- fast_llm/layers/language_model/loss/config.py | 8 +- fast_llm/layers/language_model/loss/loss.py | 18 +++- .../layers/language_model/loss/monolithic.py | 3 + .../language_model/loss/policy_gradient.py | 88 ++++++++++++++++++- tests/layers/test_lm_head.py | 25 +++++- 5 files changed, 131 insertions(+), 11 deletions(-) diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 751f44779..7d7e45ad0 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -253,17 +253,11 @@ def loss_class(self) -> "type[LanguageModelGRPOLoss]": @config_class(dynamic_type={LanguageModelLossConfig: "gspo"}) -class LanguageModelGSPOLossConfig(LanguageModelPolicyGradientLossConfig): +class LanguageModelGSPOLossConfig(LanguageModelPolicyGradientLossConfig, CombinableLossConfig): """Group Sequence Policy Optimization: sequence-level geometric-mean IS-ratio clipping.""" _abstract: typing.ClassVar[bool] = False - use_triton: bool | None = Field( - default=None, - desc="Enable triton implementation. Default: use if available.", - hint=FieldHint.expert, - ) - @property def loss_class(self) -> "type[LanguageModelGSPOLoss]": from fast_llm.layers.language_model.loss.policy_gradient import LanguageModelGSPOLoss diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 5fe431e61..978f792f4 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -166,7 +166,9 @@ class CombinableLoss: `combinable_forward_backward` or several sharing one softmax when fused together. Subclasses implement `get_inputs` (eager kwargs -> argument tuple, built outside the compiled boundary) and the `fused_core` static method (post-softmax math returning `(loss, uncast_grad, extra)`), and override - `register_combinable_extras` when they emit outputs beyond the loss scalar.""" + `register_combinable_extras` when they emit outputs beyond the loss scalar. A loss whose gradient can't + be produced inside the compiled boundary (an eager seam between forward and backward) instead has its + `fused_core` return `(None, None, forward_state)` and completes its loss/gradient in `finish`.""" _logits_scale_factor: float @@ -221,6 +223,20 @@ def register_combinable_extras( ) -> None: """Register per-loss outputs beyond the loss scalar (produced by `fused_core`). No-op by default.""" + def finish( + self, + loss: torch.Tensor | None, + extra: typing.Any, + kwargs: dict[str, typing.Any], + split_index: int, + grad_logits: torch.Tensor | None, + logits_dtype: torch.dtype, + ) -> tuple[torch.Tensor, typing.Any, torch.Tensor | None]: + """Complete a loss that couldn't finish inside the compiled shared-softmax boundary, accumulating its + gradient into `grad_logits` and returning `(loss, extra, grad_logits)`. A passthrough by default; a + loss with an eager seam runs it here from the forward state its `fused_core` returned as `extra`.""" + return loss, extra, grad_logits + def loss_forward_backward( grad_output: float | None, fn: typing.Callable, input_: torch.Tensor, *args, **kwargs diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 4a11dd231..2ce13c9ae 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -111,6 +111,9 @@ def forward_backward( total_loss = None for child, (loss, extra) in zip(self._children, results, strict=True): + # A child whose gradient can't be produced inside the compiled boundary (an eager seam) had + # `fused_core` return `(None, None, forward_state)`; `finish` completes its loss/gradient here. + loss, extra, grad_logits = child.finish(loss, extra, kwargs, split_index, grad_logits, logits.dtype) if child._do_register_loss: child._register_loss(child.name, loss, losses) child.register_combinable_extras(extra, kwargs, losses) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 07e758524..8a1901364 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -335,8 +335,15 @@ def get_loss_definitions(self) -> list[LossDef]: return defs -class LanguageModelGSPOLoss[ConfigType: LanguageModelGSPOLossConfig](LanguageModelPolicyGradientLoss[ConfigType]): - """GSPO: sequence-level geometric-mean IS-ratio clipping.""" +class LanguageModelGSPOLoss[ConfigType: LanguageModelGSPOLossConfig]( + CombinableLoss, LanguageModelPolicyGradientLoss[ConfigType] +): + """GSPO: sequence-level geometric-mean IS-ratio clipping. + + Standalone, `_forward_backward` runs the whole three-phase kernel (`fused_gspo_loss_forward_backward` or + its Triton twin). Fused into a shared softmax, `fused_core` runs only the forward on that softmax and the + segment seam + backward are deferred to `finish`, since the eager `index_add_` segment aggregation can't + run inside the compiled boundary.""" def __init__( self, @@ -417,6 +424,83 @@ def _forward_backward( self._register_new_logprobs(new_logprobs_mean, kwargs, losses) return loss, grad + def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + return (self._get_labels(kwargs, split_index),) + + @staticmethod + def fused_core( + logits_norm: torch.Tensor, + exp_logits: torch.Tensor, + sum_exp_logits: torch.Tensor, + logits_max: torch.Tensor, + group: "torch.distributed.ProcessGroup | None", + logits_scale_factor: float, + arguments: tuple, + ) -> tuple[None, None, tuple]: + """GSPO forward over the shared softmax: the per-token log-probs plus the softmax tensors its seam and + backward need. Returns `(None, None, forward_state)` — GSPO's loss and gradient can't be produced here + (the segment aggregation is eager), so both are deferred to `finish`.""" + (target,) = arguments + loss_mask = target >= 0 + predicted_logits, target_masked, target_mask = predicted_logits_from_labels( + logits_norm, target, loss_mask, group + ) + new_log_probs = predicted_logits - sum_exp_logits.log() + return None, None, (new_log_probs, loss_mask, exp_logits, sum_exp_logits, target_masked, target_mask) + + def finish( + self, + loss: torch.Tensor | None, + extra: tuple, + kwargs: dict[str, typing.Any], + split_index: int, + grad_logits: torch.Tensor | None, + logits_dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Run the eager segment seam and the compiled backward over the shared softmax deferred by + `fused_core`, accumulating GSPO's gradient into `grad_logits`. Returns the loss and the `new_logprobs` + metric (registered by `register_combinable_extras`).""" + new_log_probs, loss_mask, exp_logits, sum_exp_logits, target_masked, target_mask = extra + document_index_zero_based = ( + self._prepare_target( + kwargs[BlockKwargs.global_document_index_q], split_index, split_by_distance=False + ).long() + - 1 + ) + loss, new_logprobs_mean, effective_grad = gspo_segment_seam( + new_log_probs, + loss_mask, + self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), + self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index), + document_index_zero_based, + kwargs[BlockKwargs.num_documents_in_sequence], + self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index), + kwargs[LanguageModelKwargs.num_documents_in_batch], + self._get_grad_output(kwargs), + self._sequence_data_dim.group if self._sequence_data_active else None, + self._parallel_dim.group if self._sequence_parallel else None, + self._config.epsilon_low, + self._config.epsilon_high, + self._logits_scale_factor, + ) + if effective_grad is not None: + grad_logits = gspo_backward_core( + exp_logits, + sum_exp_logits, + target_masked, + target_mask, + loss_mask, + effective_grad, + logits_dtype, + grad_logits, + ) + return loss, new_logprobs_mean, grad_logits + + def register_combinable_extras( + self, extra: torch.Tensor | None, kwargs: dict[str, typing.Any], losses: dict | None + ) -> None: + self._register_new_logprobs(extra, kwargs, losses) + def get_preprocessing_config(self) -> dict[str, typing.Any]: return super().get_preprocessing_config() | {"return_document_index": True} diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 69386203e..2a092c6eb 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -95,7 +95,7 @@ def get_config(self) -> GPTModelConfig: combinable = { name: loss for name, loss in losses.items() - if loss["type"] in ("label", "distillation", "z_loss", "grpo") + if loss["type"] in ("label", "distillation", "z_loss", "grpo", "gspo") } if combinable: losses = {name: loss for name, loss in losses.items() if name not in combinable} @@ -433,6 +433,29 @@ def _add_configs(base_name: str, **kwargs): z_loss=0.5, ) _add_configs("fused_grpo_and_z_loss", loss_implementation="fused", grpo_loss=True, z_loss=0.5) +# GSPO fused into the shared softmax: its eager segment seam runs between the compiled forward and backward, +# so it defers to `finish` rather than a single-pass `fused_core`. Single-split only (GSPO can't be split); +# added explicitly since `_add_configs` also emits a split variant. Alone (wrapper only) and sharing the +# softmax with z-loss (the RL + regularizer combo). +for _loss_masking in (False, True): + _suffix = "_masked" if _loss_masking else "" + _lm_head_test_configs.append( + LMHeadTestConfig( + f"fused_gspo_loss{_suffix}", + gspo_loss=True, + loss_masking=_loss_masking, + loss_implementation="fused", + ) + ) + _lm_head_test_configs.append( + LMHeadTestConfig( + f"fused_gspo_and_z_loss{_suffix}", + gspo_loss=True, + z_loss=0.5, + loss_masking=_loss_masking, + loss_implementation="fused", + ) + ) # GRPO metric family. Single-split only: per-split metric partials reduce across splits, which the # whole-sequence reference doesn't model. for _loss_implementation in ("per_loss", "fused"): From 9900b54fc48d3a8b090fc45fa78c59a8ccda2a7b Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 7 Jul 2026 11:48:11 -0400 Subject: [PATCH 2/2] Fine-review fixes: dedup GSPO document-index, drop stale comment (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Factor the 1-based→0-based `global_document_index_q` conversion into `_document_index_zero_based`, shared by standalone `_forward_backward` and the fused `finish`, so the convention comment lives in one place. - Drop the now-stale combinable-loss enumeration in the monolithic test comment (GSPO is combinable too). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../language_model/loss/policy_gradient.py | 24 +++++++++---------- tests/layers/test_lm_head.py | 6 ++--- 2 files changed, 14 insertions(+), 16 deletions(-) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 8a1901364..0dcf7476f 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -375,6 +375,15 @@ def __init__( # aggregation can't recombine across chunks since each call only sees a slice. Assert.eq(self._num_splits, 1) + def _document_index_zero_based(self, kwargs: dict[str, typing.Any], split_index: int) -> torch.Tensor: + # `global_document_index_q` is 1-based per the data preprocessor convention; the kernel takes 0-based. + return ( + self._prepare_target( + kwargs[BlockKwargs.global_document_index_q], split_index, split_by_distance=False + ).long() + - 1 + ) + def _forward_backward( self, logits: "torch.Tensor", @@ -383,13 +392,7 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - document_index_zero_based = ( - self._prepare_target( - kwargs[BlockKwargs.global_document_index_q], split_index, split_by_distance=False - ).long() - - 1 - ) - # `global_document_index_q` is 1-based per the data preprocessor convention; the kernel takes 0-based. + document_index_zero_based = self._document_index_zero_based(kwargs, split_index) # `num_documents_in_sequence` is the doc count for this DP rank's batch — identical across # SDP/SP ranks within a DP rank, so per-segment buffers are sized consistently for all-reduce. num_segments = kwargs[BlockKwargs.num_documents_in_sequence] @@ -461,12 +464,7 @@ def finish( `fused_core`, accumulating GSPO's gradient into `grad_logits`. Returns the loss and the `new_logprobs` metric (registered by `register_combinable_extras`).""" new_log_probs, loss_mask, exp_logits, sum_exp_logits, target_masked, target_mask = extra - document_index_zero_based = ( - self._prepare_target( - kwargs[BlockKwargs.global_document_index_q], split_index, split_by_distance=False - ).long() - - 1 - ) + document_index_zero_based = self._document_index_zero_based(kwargs, split_index) loss, new_logprobs_mean, effective_grad = gspo_segment_seam( new_log_probs, loss_mask, diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 2a092c6eb..8ceb30b41 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -405,9 +405,9 @@ def _add_configs(base_name: str, **kwargs): _add_configs("label_and_distillation_loss_zero_weight", label_loss=True, distillation_loss=0.0) _add_configs("distillation_loss_temperature", distillation_loss=True, distillation_temperature=2.0) -# Monolithic loss type: the combinable losses (cross-entropy, z-loss, distillation, GRPO) are wrapped in a -# single `monolithic` loss that shares one softmax pass; the head treats it as an ordinary loss. These -# configs must match their per-loss equivalents above (validated against the same independent reference). +# Monolithic loss type: the combinable losses are wrapped in a single `monolithic` loss that shares one +# softmax pass; the head treats it as an ordinary loss. These configs must match their per-loss equivalents +# above (validated against the same independent reference). _add_configs("fused", loss_implementation="fused") _add_configs("fused_bfloat16", loss_implementation="fused", compute_dtype=DataType.bfloat16) _add_configs("fused_logit_scaling", loss_implementation="fused", logits_scale_factor=5.0)