Fuse GSPO into the monolithic compiled loss (#507)

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

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

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)

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,
@@ -368,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",
@@ -376,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]
@@ -417,6 +427,78 @@ 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._document_index_zero_based(kwargs, split_index)
+        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}
 

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}
@@ -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)
@@ -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"):