From b6c71a81b8b55382834136a157452d97f91ce3c3 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 26 Jun 2026 15:32:38 -0400 Subject: [PATCH 01/28] Monolithic head-loss kernel: scaffolding + cross-entropy (#507) Introduce a single torch.compile head-loss kernel that runs the vocab softmax once and emits all requested losses + the combined gradient in one pass, replacing the per-loss loop where each loss runs its own softmax. This commit lands the scaffolding and the first loss (cross-entropy from labels); further losses are added incrementally. - `_monolithic_core` (loss/monolithic.py) is one @torch.compile boundary that inlines the plain softmax cores (`_softmax_base`, `_predicted_logits_from_labels`, newly factored out of functional/entropy_loss.py while keeping the public `fused_*` wrappers byte-identical), so compile fuses the work across enabled losses. Losses toggle on by passing their inputs; gradients accumulate in fp32 and cast once. - Opt-in via `LanguageModelHeadConfig.loss_implementation` ({per_loss (default), fused, triton}). The head keeps the per-loss path as the baseline/oracle and falls back to it for losses not yet in the kernel (e.g. DPO), accumulating into the same gradient. Per-loss bookkeeping stays in the loss objects via `get_monolithic_spec` / `register_monolithic_outputs`. - Validation requires a single shared effective logits scale factor across softmax losses. Tests: `test_monolithic_loss` (kernel vs the fused CE baseline) and `fused_*` head configs. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/entropy_loss.py | 18 ++- fast_llm/layers/language_model/config.py | 22 ++++ fast_llm/layers/language_model/head.py | 63 +++++++-- .../language_model/loss/entropy_loss.py | 13 ++ fast_llm/layers/language_model/loss/loss.py | 14 ++ .../layers/language_model/loss/monolithic.py | 123 ++++++++++++++++++ tests/layers/test_lm_head.py | 13 ++ tests/layers/test_lm_losses.py | 58 +++++++++ 8 files changed, 308 insertions(+), 16 deletions(-) create mode 100644 fast_llm/layers/language_model/loss/monolithic.py diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index 05eaae520..446778c92 100644 --- a/fast_llm/functional/entropy_loss.py +++ b/fast_llm/functional/entropy_loss.py @@ -92,8 +92,7 @@ def torch_entropy_loss_forward_backward( return loss.detach_(), grad -@torch.compile -def fused_softmax_base( +def _softmax_base( logits: torch.Tensor, # (*batch, vocab) logits_scale_factor: float = 1.0, group: ProcessGroup | None = None, @@ -106,6 +105,9 @@ def fused_softmax_base( in a numerically stable way and with tensor-parallel support. Warning: The returned values are regularized by `logits_max`. The regularization typically but not always cancels out in derived quantities. + + This plain (un-compiled) core is shared between the public `fused_softmax_base` wrapper and the + monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. """ logits = logits.float() if logits_scale_factor != 1.0: @@ -121,6 +123,9 @@ def fused_softmax_base( return logits_norm, exp_logits, sum_exp_logits, logits_max +fused_softmax_base = torch.compile(_softmax_base) + + @torch.compile def _fused_reverse_kl_base_from_distribution( logits: torch.Tensor, # (*batch, vocab) @@ -207,8 +212,7 @@ def _fused_cross_entropy_base_from_distribution( return per_sample_loss, grad -@torch.compile -def fused_predicted_logits_from_labels( +def _predicted_logits_from_labels( logits: torch.Tensor, # (*batch, vocab) target: torch.Tensor, # (*batch,) loss_mask: torch.Tensor, # (*batch,), == target>=0 @@ -221,6 +225,9 @@ def fused_predicted_logits_from_labels( Normally used in combination with `fused_softmax_base`, may also recover probabilities or log probabilities: `predicted_probabilities = predicted_logits.exp() / sum_exp_logits` `predicted_log_probabilities = predicted_logits / sum_exp_logits.log()` + + This plain (un-compiled) core is shared between the public `fused_predicted_logits_from_labels` wrapper + and the monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. """ if group is None: @@ -243,6 +250,9 @@ def fused_predicted_logits_from_labels( return predicted_logits, target_masked, target_mask +fused_predicted_logits_from_labels = torch.compile(_predicted_logits_from_labels) + + @torch.compile def _fused_cross_entropy_base_from_labels( logits: torch.Tensor, # (*batch, vocab) diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index bde33f297..b1d9fb5c4 100644 --- a/fast_llm/layers/language_model/config.py +++ b/fast_llm/layers/language_model/config.py @@ -1,3 +1,4 @@ +import enum import typing from fast_llm.config import Field, FieldHint, check_field, config_class, skip_valid_if_none @@ -33,6 +34,15 @@ class LanguageModelKwargs(LanguageModelLossKwargs): LM_HEAD_LOSS_NAME = "lm_head_loss" +class LossImplementation(enum.StrEnum): + # One loss kernel per loss, each running its own softmax pass (the original, always-available path). + per_loss = "per_loss" + # A single monolithic torch.compile kernel running the softmax once for all combinable losses. + fused = "fused" + # A single monolithic triton kernel. + triton = "triton" + + @config_class() class LanguageModelEmbeddingsConfig(BlockConfig): _abstract = False @@ -143,6 +153,13 @@ class LanguageModelHeadConfig(BlockConfig): doc="If not provided, all heads are equally weighted.", hint=FieldHint.feature, ) + loss_implementation: LossImplementation = Field( + default=LossImplementation.per_loss, + desc="Select the head-loss implementation. `per_loss` runs each loss separately (the default);" + " `fused`/`triton` run a single monolithic kernel that shares one softmax pass across combinable" + " losses. Losses not yet supported by the monolithic kernel fall back to their own implementation.", + hint=FieldHint.expert, + ) def get_layer( self, @@ -177,6 +194,11 @@ def get_layer( def _validate(self) -> None: super()._validate() assert LM_HEAD_LOSS_NAME not in self.losses + if self.loss_implementation != LossImplementation.per_loss: + # The monolithic kernel shares a single softmax across losses, so they must agree on the + # effective logits scale factor (model scale stacked with each loss's own scale factor). + scale_factors = {self.logits_scale_factor * loss.logits_scale_factor for loss in self.losses.values()} + Assert.leq(len(scale_factors), 1) def get_reference_models(self) -> set[str]: return {reference_model for loss in self.losses.values() for reference_model in loss.get_reference_models()} diff --git a/fast_llm/layers/language_model/head.py b/fast_llm/layers/language_model/head.py index 22c750082..500651d8b 100644 --- a/fast_llm/layers/language_model/head.py +++ b/fast_llm/layers/language_model/head.py @@ -19,9 +19,11 @@ LanguageModelEmbeddingsConfig, LanguageModelHeadConfig, LanguageModelKwargs, + LossImplementation, ) from fast_llm.layers.language_model.loss.config import LanguageModelLabelEntropyLossConfig from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.monolithic import monolithic_head_loss_forward_backward from fast_llm.tensor import TensorMeta from fast_llm.utils import Assert, safe_merge_dicts @@ -272,18 +274,21 @@ def _logits_loss_forward_backward_partial( if return_logits: return logits, None - losses_, grad = [], None - for loss in self.losses: - # losses are returned unscaled but the grads are already scaled - loss_value, grad = loss.forward_backward( - logits, - kwargs, - losses, - split_index, - grad, - ) - if loss_value is not None: - losses_.append(loss_value.detach()) + if self._config.loss_implementation == LossImplementation.per_loss: + losses_, grad = [], None + for loss in self.losses: + # losses are returned unscaled but the grads are already scaled + loss_value, grad = loss.forward_backward( + logits, + kwargs, + losses, + split_index, + grad, + ) + if loss_value is not None: + losses_.append(loss_value.detach()) + else: + losses_, grad = self._monolithic_loss_forward_backward(logits, kwargs, losses, split_index) if grad is not None and self._config.final_logit_softcap is not None: grad = _softcap_backward(grad, logits, self._config.final_logit_softcap) @@ -292,6 +297,40 @@ def _logits_loss_forward_backward_partial( output_parallel_linear_backward(grad, context) if self.training else None ) + def _monolithic_loss_forward_backward( + self, logits: torch.Tensor, kwargs: dict, losses: dict | None, split_index: int + ) -> tuple[list[torch.Tensor], torch.Tensor | None]: + # Losses supported by the monolithic kernel share one softmax pass; the rest (e.g. DPO) run + # through their own `forward_backward`, accumulating into the same logits gradient. + specs, monolithic_losses, other_losses = [], [], [] + for loss in self.losses: + spec = loss.get_monolithic_spec(kwargs, split_index) + if spec is None: + other_losses.append(loss) + else: + specs.append(spec) + monolithic_losses.append(loss) + + losses_, grad = [], None + if specs: + outputs, grad = monolithic_head_loss_forward_backward( + logits, + specs, + group=self._parallel_dim.group if self._vocab_parallel else None, + grad_logits=grad, + ) + for loss, output in zip(monolithic_losses, outputs, strict=True): + loss.register_monolithic_outputs(output, kwargs, losses) + if output.loss is not None: + losses_.append((output.loss if loss.weight == 1 else output.loss * loss.weight).detach()) + + for loss in other_losses: + loss_value, grad = loss.forward_backward(logits, kwargs, losses, split_index, grad) + if loss_value is not None: + losses_.append(loss_value.detach()) + + return losses_, grad + def get_loss_definitions(self) -> list[LossDef]: return [ LossDef(name=self._total_loss_name), diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py index 48c1556f3..1711a298c 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -10,6 +10,7 @@ LanguageModelLabelEntropyLossConfig, ) from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.monolithic import MonolithicLossSpec class LanguageModelLabelEntropyLoss[ConfigType: LanguageModelLabelEntropyLossConfig](LanguageModelLoss[ConfigType]): @@ -38,6 +39,18 @@ def _forward_backward( divisor=self._get_label_count(kwargs), ) + def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + # For labels, forward-KL is identical to cross-entropy (one-hot target entropy is zero). + return MonolithicLossSpec( + kind="cross_entropy", + name=self.name, + weight=self._weight, + logits_scale_factor=self._logits_scale_factor, + grad_output=self._get_grad_output(kwargs), + divisor=self._get_label_count(kwargs), + target=self._get_labels(kwargs, split_index), + ) + class LanguageModelDistillationLoss[ConfigType: LanguageModelDistillationLossConfig](LanguageModelLoss[ConfigType]): def _forward_backward( diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 397aec966..868fe6576 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -9,6 +9,7 @@ from fast_llm.engine.distributed.config import DistributedConfig, DistributedDimNames from fast_llm.layers.language_model.config import LanguageModelKwargs from fast_llm.layers.language_model.loss.config import LanguageModelLossConfig, LanguageModelLossKwargs +from fast_llm.layers.language_model.loss.monolithic import MonolithicLossOutput, MonolithicLossSpec from fast_llm.utils import Assert @@ -72,6 +73,19 @@ def _forward_backward( def get_loss_definitions(self) -> list[LossDef]: return [LossDef(name=self.name)] if self._do_register_loss else [] + def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> "MonolithicLossSpec | None": + """ + Package this loss's inputs for the monolithic head-loss kernel, or return `None` if it is not + supported there (it then runs through its own `forward_backward`, accumulating into the same grad). + """ + return None + + def register_monolithic_outputs( + self, output: "MonolithicLossOutput", kwargs: dict[str, typing.Any], losses: dict | None + ) -> None: + if self._do_register_loss: + self._register_loss(self.name, output.loss, losses) + def get_preprocessing_config( self, ) -> dict[str, typing.Any]: diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py new file mode 100644 index 000000000..5f7bda461 --- /dev/null +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -0,0 +1,123 @@ +import typing + +import torch + +from fast_llm.core.distributed import ProcessGroup +from fast_llm.functional.entropy_loss import _predicted_logits_from_labels, _softmax_base +from fast_llm.utils import Assert + + +class MonolithicLossSpec(typing.NamedTuple): + """ + Per-loss inputs gathered by `LanguageModelLoss.get_monolithic_spec` and consumed by the monolithic + head-loss kernel. `kind` selects which branch of the kernel computes the loss and its gradient. + The math fields not used by a given `kind` are left at their defaults. + """ + + kind: str + name: str + weight: float + logits_scale_factor: float + grad_output: float | None + divisor: float + # Cross-entropy (from labels). + target: torch.Tensor | None = None + loss_mask: torch.Tensor | None = None + + +class MonolithicLossOutput(typing.NamedTuple): + """Per-loss outputs returned by the monolithic kernel, registered via `register_monolithic_outputs`.""" + + loss: torch.Tensor | None = None + + +@torch.compile +def _monolithic_core( + logits: torch.Tensor, # (*batch, vocab) + group: ProcessGroup | None, + logits_scale_factor: float, + grad_logits: torch.Tensor | None, + ce_target: torch.Tensor | None, + ce_grad_output: float | None, + ce_divisor: float, +) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """ + The monolithic head-loss kernel: one shared softmax over the logits, then every enabled loss's scalar + and gradient contribution. Disabled losses are toggled off by passing their inputs as `None`, so + `torch.compile` dead-code-eliminates their branch. Because this is a single `@torch.compile` boundary + that calls the plain (un-compiled) softmax cores, compile fuses the work across all enabled losses. + + Gradients accumulate in fp32 and are cast to `logits.dtype` once at the end. + """ + logits_norm, exp_logits, sum_exp_logits, logits_max = _softmax_base(logits, logits_scale_factor, group) + grad = None + + cross_entropy_loss = None + if ce_target is not None: + loss_mask = ce_target >= 0 + predicted_logits, target_masked, target_mask = _predicted_logits_from_labels( + logits_norm, ce_target, loss_mask, group + ) + cross_entropy_loss = ((sum_exp_logits.log() - predicted_logits) * loss_mask).sum() / ce_divisor + if ce_grad_output is not None: + grad_output = ce_grad_output / ce_divisor * logits_scale_factor + cross_entropy_grad = exp_logits.scatter_add( + -1, + target_masked.unsqueeze(-1), + ( + -sum_exp_logits.unsqueeze(-1) + if target_mask is None + else -(target_mask * sum_exp_logits).unsqueeze(-1) + ), + ) * (grad_output / sum_exp_logits.unsqueeze(-1)) + cross_entropy_grad = cross_entropy_grad * loss_mask.unsqueeze(-1) + grad = cross_entropy_grad if grad is None else grad + cross_entropy_grad + + if grad is not None: + grad = grad.to(logits.dtype) + if grad_logits is None: + grad_logits = grad + else: + grad_logits.add_(grad) + + return cross_entropy_loss, grad_logits + + +def monolithic_head_loss_forward_backward( + logits: torch.Tensor, + specs: list[MonolithicLossSpec], + *, + group: ProcessGroup | None, + grad_logits: torch.Tensor | None = None, +) -> tuple[list[MonolithicLossOutput], torch.Tensor | None]: + """ + Marshal the per-loss specs into the flat arguments of `_monolithic_core`, run it once, and map its + outputs back to one `MonolithicLossOutput` per spec (in input order). All specs must share a single + effective `logits_scale_factor` (the shared softmax cannot serve two scales); this is validated at + config time and asserted here. + """ + logits_scale_factor = specs[0].logits_scale_factor + cross_entropy_spec = None + for spec in specs: + Assert.eq(spec.logits_scale_factor, logits_scale_factor) + if spec.kind == "cross_entropy": + assert cross_entropy_spec is None + cross_entropy_spec = spec + else: + raise NotImplementedError(spec.kind) + + cross_entropy_loss, grad_logits = _monolithic_core( + logits, + group, + logits_scale_factor, + grad_logits, + ce_target=None if cross_entropy_spec is None else cross_entropy_spec.target, + ce_grad_output=None if cross_entropy_spec is None else cross_entropy_spec.grad_output, + ce_divisor=1.0 if cross_entropy_spec is None else cross_entropy_spec.divisor, + ) + + outputs = [] + for spec in specs: + if spec.kind == "cross_entropy": + outputs.append(MonolithicLossOutput(loss=cross_entropy_loss)) + return outputs, grad_logits diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 6f09cb108..a7f8bfb73 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -39,6 +39,7 @@ class LMHeadTestConfig: tied_embedding_weight: bool = False num_splits: int = 1 gspo_document_lengths: tuple[int, ...] | None = None + loss_implementation: str = "per_loss" @property def actual_label_loss(self): @@ -59,6 +60,8 @@ def get_config(self) -> GPTModelConfig: "cross_entropy_splits": self.num_splits, "prediction_heads": self.prediction_heads, } + if self.loss_implementation != "per_loss": + head_config["loss_implementation"] = self.loss_implementation if self.final_logit_softcap is not None: head_config["final_logit_softcap"] = self.final_logit_softcap losses = {} @@ -360,6 +363,16 @@ def _add_configs(base_name: str, **kwargs): _add_configs("label_and_z_loss_weighted", label_loss=True, z_loss=0.5) _add_configs("label_and_distillation_loss_zero_weight", label_loss=True, distillation_loss=0.0) +# Monolithic ("fused") head-loss path. So far only cross-entropy is computed in the monolithic kernel; +# other losses (e.g. z-loss) fall back to their own implementation and accumulate into the same gradient. +_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) +_add_configs("fused_final_logit_softcap", loss_implementation="fused", final_logit_softcap=2.0) +_add_configs("fused_tied_embedding_weight", loss_implementation="fused", tied_embedding_weight=True) +_add_configs("fused_multi_token_prediction", loss_implementation="fused", prediction_heads=2) +_add_configs("fused_label_and_z_loss_weighted", loss_implementation="fused", label_loss=True, z_loss=0.5) + @pytest.mark.slow @pytest.mark.parametrize( diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index d7d14ad3e..fb49ca42c 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -18,6 +18,7 @@ from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward from fast_llm.layers.language_model.loss.dpo import dpo_loss from fast_llm.layers.language_model.loss.loss import loss_forward_backward +from fast_llm.layers.language_model.loss.monolithic import MonolithicLossSpec, monolithic_head_loss_forward_backward from fast_llm.layers.language_model.loss.policy_gradient import ( GRPOMetrics, compute_grpo_metrics, @@ -586,6 +587,63 @@ def _test_z_loss( ) +def _test_monolithic_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, group=None +): + # The monolithic kernel must reproduce the per-loss baseline (here: a single cross-entropy loss). + logits, target, loss_mask = _get_lm_loss_inputs(num_columns, loss_masking, TargetFormat.labels, batch_shape, dtype) + local_logits = split_op(logits, group, -1).contiguous() + # The head feeds the label count as divisor; mirror that instead of the kernel's default (numel). + divisor = max(int((target >= 0).sum().item()), 1) + out_ref, grad_ref = fused_entropy_loss_forward_backward( + logits=local_logits, + target=target, + loss_mask=None, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + target_format=TargetFormat.labels, + entropy_loss_type=EntropyLossType.cross_entropy, + divisor=divisor, + ) + spec = MonolithicLossSpec( + kind="cross_entropy", + name="cross_entropy", + weight=1.0, + logits_scale_factor=logits_scale_factor, + grad_output=grad_output, + divisor=divisor, + target=target, + ) + if accumulate and grad_output is not None: + previous_grad = torch.randn_like(local_logits) + grad_ref = grad_ref + previous_grad + grad_logits = previous_grad.clone() + else: + grad_logits = None + outputs, grad_mono = monolithic_head_loss_forward_backward( + local_logits, [spec], group=group, grad_logits=grad_logits + ) + threshold = 1e-5 if dtype == DataType.float32 else 1e-4 + Assert.rms_close_relative(outputs[0].loss, out_ref, threshold, 1e-6) + if grad_output is not None: + Assert.rms_close_relative(grad_mono, grad_ref, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-7) + else: + assert grad_mono is None + + +@pytest.mark.slow +@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) +@pytest.mark.parametrize( + ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), + _LOSS_PARAMETERS, +) +def test_monolithic_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate +): + _test_monolithic_loss(batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( From 3b4070a06e6ffac3085230ebb59da512c689b3d6 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 26 Jun 2026 16:36:13 -0400 Subject: [PATCH 02/28] Monolithic head-loss kernel: add z-loss (#507) Add a z-loss branch to the monolithic kernel, reusing the shared softmax's sum_exp_logits and logits_max (z-loss adds logits_max back; cross-entropy cancels it) so the flagship CE + z-loss combo runs one softmax pass. z-loss has no required tensor input, so it toggles on a static `z_loss_enabled` flag rather than a None-gated tensor. Grad terms accumulate in fp32 and cast once. Functional parity tests now cover z-loss alone and the CE + z-loss combo against the summed per-loss baselines, plus a tensor-parallel subtest. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../layers/language_model/loss/monolithic.py | 62 ++++++--- fast_llm/layers/language_model/loss/z_loss.py | 12 ++ tests/layers/test_lm_losses.py | 127 ++++++++++++++---- 3 files changed, 154 insertions(+), 47 deletions(-) diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 5f7bda461..3cda19058 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -20,9 +20,8 @@ class MonolithicLossSpec(typing.NamedTuple): logits_scale_factor: float grad_output: float | None divisor: float - # Cross-entropy (from labels). - target: torch.Tensor | None = None - loss_mask: torch.Tensor | None = None + target: torch.Tensor | None = None # Cross-entropy labels. + loss_mask: torch.Tensor | None = None # z-loss mask (cross-entropy derives its own from `target >= 0`). class MonolithicLossOutput(typing.NamedTuple): @@ -40,12 +39,18 @@ def _monolithic_core( ce_target: torch.Tensor | None, ce_grad_output: float | None, ce_divisor: float, -) -> tuple[torch.Tensor | None, torch.Tensor | None]: + z_loss_enabled: bool, + z_loss_mask: torch.Tensor | None, + z_loss_grad_output: float | None, + z_loss_divisor: float, +) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: """ The monolithic head-loss kernel: one shared softmax over the logits, then every enabled loss's scalar - and gradient contribution. Disabled losses are toggled off by passing their inputs as `None`, so - `torch.compile` dead-code-eliminates their branch. Because this is a single `@torch.compile` boundary - that calls the plain (un-compiled) softmax cores, compile fuses the work across all enabled losses. + and gradient contribution. Losses with a required tensor input are toggled off by passing it as `None` + (e.g. cross-entropy's `ce_target`); losses without one (z-loss) toggle on a static `*_enabled` bool. In + both cases `torch.compile` specializes on the flag and dead-code-eliminates the disabled branch. Because + this is a single `@torch.compile` boundary that calls the plain (un-compiled) softmax cores, compile + fuses the work across all enabled losses. Gradients accumulate in fp32 and are cast to `logits.dtype` once at the end. """ @@ -73,6 +78,22 @@ def _monolithic_core( cross_entropy_grad = cross_entropy_grad * loss_mask.unsqueeze(-1) grad = cross_entropy_grad if grad is None else grad + cross_entropy_grad + z_loss = None + if z_loss_enabled: + # z-loss needs the un-regularized log-sum-exp, so it adds back `logits_max` (cross-entropy cancels it). + log_sum_exp_logits = sum_exp_logits.log() + logits_max + z_loss_term = log_sum_exp_logits**2 + if z_loss_mask is not None: + z_loss_term = z_loss_term * z_loss_mask + z_loss = z_loss_term.sum() / z_loss_divisor + if z_loss_grad_output is not None: + grad_output = z_loss_grad_output / z_loss_divisor * logits_scale_factor + z_loss_grad_base = 2 * grad_output * (log_sum_exp_logits / sum_exp_logits) + if z_loss_mask is not None: + z_loss_grad_base = z_loss_grad_base * z_loss_mask + z_loss_grad = z_loss_grad_base.unsqueeze(-1) * exp_logits + grad = z_loss_grad if grad is None else grad + z_loss_grad + if grad is not None: grad = grad.to(logits.dtype) if grad_logits is None: @@ -80,7 +101,7 @@ def _monolithic_core( else: grad_logits.add_(grad) - return cross_entropy_loss, grad_logits + return cross_entropy_loss, z_loss, grad_logits def monolithic_head_loss_forward_backward( @@ -97,16 +118,18 @@ def monolithic_head_loss_forward_backward( config time and asserted here. """ logits_scale_factor = specs[0].logits_scale_factor - cross_entropy_spec = None + specs_by_kind: dict[str, MonolithicLossSpec] = {} for spec in specs: Assert.eq(spec.logits_scale_factor, logits_scale_factor) - if spec.kind == "cross_entropy": - assert cross_entropy_spec is None - cross_entropy_spec = spec - else: + if spec.kind not in ("cross_entropy", "z_loss"): raise NotImplementedError(spec.kind) + assert spec.kind not in specs_by_kind, spec.kind + specs_by_kind[spec.kind] = spec - cross_entropy_loss, grad_logits = _monolithic_core( + cross_entropy_spec = specs_by_kind.get("cross_entropy") + z_loss_spec = specs_by_kind.get("z_loss") + + cross_entropy_loss, z_loss, grad_logits = _monolithic_core( logits, group, logits_scale_factor, @@ -114,10 +137,11 @@ def monolithic_head_loss_forward_backward( ce_target=None if cross_entropy_spec is None else cross_entropy_spec.target, ce_grad_output=None if cross_entropy_spec is None else cross_entropy_spec.grad_output, ce_divisor=1.0 if cross_entropy_spec is None else cross_entropy_spec.divisor, + z_loss_enabled=z_loss_spec is not None, + z_loss_mask=None if z_loss_spec is None else z_loss_spec.loss_mask, + z_loss_grad_output=None if z_loss_spec is None else z_loss_spec.grad_output, + z_loss_divisor=1.0 if z_loss_spec is None else z_loss_spec.divisor, ) - outputs = [] - for spec in specs: - if spec.kind == "cross_entropy": - outputs.append(MonolithicLossOutput(loss=cross_entropy_loss)) - return outputs, grad_logits + loss_by_kind = {"cross_entropy": cross_entropy_loss, "z_loss": z_loss} + return [MonolithicLossOutput(loss=loss_by_kind[spec.kind]) for spec in specs], grad_logits diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py index 2e5f90b1d..451ab48f4 100644 --- a/fast_llm/layers/language_model/loss/z_loss.py +++ b/fast_llm/layers/language_model/loss/z_loss.py @@ -8,6 +8,7 @@ from fast_llm.functional.utils import reduce_losses from fast_llm.layers.language_model.loss.config import LanguageModelZLossConfig from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.monolithic import MonolithicLossSpec class LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](LanguageModelLoss[ConfigType]): @@ -33,6 +34,17 @@ def _forward_backward( divisor=self._get_label_count(kwargs), ) + def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + return MonolithicLossSpec( + kind="z_loss", + name=self.name, + weight=self._weight, + logits_scale_factor=self._logits_scale_factor, + grad_output=self._get_grad_output(kwargs), + divisor=self._get_label_count(kwargs), + loss_mask=self._get_loss_mask(kwargs, split_index), + ) + def get_preprocessing_config(self) -> dict[str, typing.Any]: return {"return_prediction_mask": True} diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index fb49ca42c..5bcd9e035 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -587,61 +587,115 @@ def _test_z_loss( ) +def _monolithic_baseline( + kind: str, local_logits, target, z_mask, grad_output, group, logits_scale_factor, divisor +) -> tuple[torch.Tensor, torch.Tensor | None, MonolithicLossSpec]: + # Each enabled loss's baseline is its own `fused_*` kernel; the monolithic grad must equal their sum. + if kind == "cross_entropy": + loss, grad = fused_entropy_loss_forward_backward( + logits=local_logits, + target=target, + loss_mask=None, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + target_format=TargetFormat.labels, + entropy_loss_type=EntropyLossType.cross_entropy, + divisor=divisor, + ) + spec = MonolithicLossSpec( + kind=kind, + name=kind, + weight=1.0, + logits_scale_factor=logits_scale_factor, + grad_output=grad_output, + divisor=divisor, + target=target, + ) + elif kind == "z_loss": + loss, grad = fused_z_loss_forward_backward( + logits=local_logits, + loss_mask=z_mask, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + divisor=divisor, + ) + spec = MonolithicLossSpec( + kind=kind, + name=kind, + weight=1.0, + logits_scale_factor=logits_scale_factor, + grad_output=grad_output, + divisor=divisor, + loss_mask=z_mask, + ) + else: + raise NotImplementedError(kind) + return loss, grad, spec + + def _test_monolithic_loss( - batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, group=None + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, kinds, group=None ): - # The monolithic kernel must reproduce the per-loss baseline (here: a single cross-entropy loss). + # The monolithic kernel must reproduce the sum of the per-loss baselines over a single shared softmax. logits, target, loss_mask = _get_lm_loss_inputs(num_columns, loss_masking, TargetFormat.labels, batch_shape, dtype) local_logits = split_op(logits, group, -1).contiguous() # The head feeds the label count as divisor; mirror that instead of the kernel's default (numel). divisor = max(int((target >= 0).sum().item()), 1) - out_ref, grad_ref = fused_entropy_loss_forward_backward( - logits=local_logits, - target=target, - loss_mask=None, - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - target_format=TargetFormat.labels, - entropy_loss_type=EntropyLossType.cross_entropy, - divisor=divisor, - ) - spec = MonolithicLossSpec( - kind="cross_entropy", - name="cross_entropy", - weight=1.0, - logits_scale_factor=logits_scale_factor, - grad_output=grad_output, - divisor=divisor, - target=target, - ) + # z-loss takes an explicit mask (the head's `loss_mask` kwarg); align it with the labeled tokens. + z_mask = (target >= 0) if loss_masking else None + + ref_losses, specs, ref_grad = [], [], None + for kind in kinds: + ref_loss, grad, spec = _monolithic_baseline( + kind, local_logits, target, z_mask, grad_output, group, logits_scale_factor, divisor + ) + ref_losses.append(ref_loss) + specs.append(spec) + if grad is not None: + ref_grad = grad if ref_grad is None else ref_grad + grad + if accumulate and grad_output is not None: previous_grad = torch.randn_like(local_logits) - grad_ref = grad_ref + previous_grad + ref_grad = ref_grad + previous_grad grad_logits = previous_grad.clone() else: grad_logits = None outputs, grad_mono = monolithic_head_loss_forward_backward( - local_logits, [spec], group=group, grad_logits=grad_logits + local_logits, specs, group=group, grad_logits=grad_logits ) threshold = 1e-5 if dtype == DataType.float32 else 1e-4 - Assert.rms_close_relative(outputs[0].loss, out_ref, threshold, 1e-6) + for output, ref_loss in zip(outputs, ref_losses, strict=True): + Assert.rms_close_relative(output.loss, ref_loss, threshold, 1e-6) if grad_output is not None: - Assert.rms_close_relative(grad_mono, grad_ref, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-7) + # The monolithic kernel sums all grad terms in fp32 and casts once, so it is *more* accurate than the + # per-loss baseline (which accumulates in `logits.dtype`); fp16 multi-loss grads differ at the ULP floor. + Assert.rms_close_relative(grad_mono, ref_grad, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-6) else: assert grad_mono is None +_MONOLITHIC_KINDS = ( + ("cross_entropy",), + ("z_loss",), + ("cross_entropy", "z_loss"), +) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), _LOSS_PARAMETERS, ) +@pytest.mark.parametrize("kinds", _MONOLITHIC_KINDS) def test_monolithic_loss( - batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate + kinds, batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate ): - _test_monolithic_loss(batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate) + _test_monolithic_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, kinds + ) @pytest.mark.slow @@ -847,6 +901,22 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa compute_entropy, test_context.group, ) + # Monolithic kernel + for kinds in _MONOLITHIC_KINDS: + with test_context.subtest(base_path, f"monolithic-{"_".join(kinds)}-{suffix}", 2) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_monolithic_loss( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + accumulate, + kinds, + test_context.group, + ) @pytest.mark.slow @@ -889,6 +959,7 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): "grpo", "grpo_metrics-False", "grpo_metrics-True", + *(f"monolithic-{"_".join(kinds)}" for kinds in _MONOLITHIC_KINDS), ), ) def test_lm_loss_distributed( From f5ca0d7669cc465caddf22913dd8aadda7712033 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 26 Jun 2026 17:08:03 -0400 Subject: [PATCH 03/28] Monolithic head-loss kernel: add from-distribution KL losses (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the from-distribution entropy losses (cross-entropy / forward-KL / reverse-KL over a teacher distribution given as logits or probabilities, with temperature) to the monolithic kernel. These share the student softmax with cross-entropy and z-loss, and add one teacher softmax in the same boundary when the target is logits. Factor the post-student-softmax math of the from-distribution helpers into plain cores (_cross_entropy_from_distribution_core, _reverse_kl_from_distribution_core) that accept the precomputed student softmax, and have both the public fused wrappers and the monolithic kernel call them — one source of truth, baseline math unchanged. Functional parity tests cover all three loss types over both target formats and a non-unit temperature, against fused_entropy_loss_forward_backward, plus tensor-parallel subtests. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/entropy_loss.py | 78 +++++++++-- .../layers/language_model/loss/monolithic.py | 89 ++++++++++-- tests/layers/test_lm_losses.py | 130 ++++++++++++++++++ 3 files changed, 281 insertions(+), 16 deletions(-) diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index 446778c92..e00e77ee2 100644 --- a/fast_llm/functional/entropy_loss.py +++ b/fast_llm/functional/entropy_loss.py @@ -126,9 +126,10 @@ def _softmax_base( fused_softmax_base = torch.compile(_softmax_base) -@torch.compile -def _fused_reverse_kl_base_from_distribution( - logits: torch.Tensor, # (*batch, vocab) +def _reverse_kl_from_distribution_core( + logits_norm: torch.Tensor, # (*batch, vocab) + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) target: torch.Tensor, # (*batch, vocab) grad_output: float | None, logits_scale_factor: float, @@ -136,13 +137,18 @@ def _fused_reverse_kl_base_from_distribution( group: ProcessGroup | None = None, temperature: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) + """ + Reverse-KL math from a precomputed student softmax (adding a teacher softmax when the target is logits). + + This plain (un-compiled) core is shared between the public `_fused_reverse_kl_base_from_distribution` + wrapper and the monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. + """ assert target_format in (TargetFormat.logits, TargetFormat.probabilities) - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) predicted_log_probability = logits_norm - sum_exp_logits.log().unsqueeze(-1) predicted_probability = exp_logits / sum_exp_logits.unsqueeze(-1) if target_format == TargetFormat.logits: - target_logits_norm, _, sum_exp_target_logits, _ = fused_softmax_base( + target_logits_norm, _, sum_exp_target_logits, _ = _softmax_base( target, logits_scale_factor / temperature, group ) target_log_probability = target_logits_norm - sum_exp_target_logits.log().unsqueeze(-1) @@ -167,7 +173,7 @@ def _fused_reverse_kl_base_from_distribution( @torch.compile -def _fused_cross_entropy_base_from_distribution( +def _fused_reverse_kl_base_from_distribution( logits: torch.Tensor, # (*batch, vocab) target: torch.Tensor, # (*batch, vocab) grad_output: float | None, @@ -175,12 +181,42 @@ def _fused_cross_entropy_base_from_distribution( target_format: TargetFormat, group: ProcessGroup | None = None, temperature: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) + logits_norm, exp_logits, sum_exp_logits, _ = _softmax_base(logits, logits_scale_factor, group) + return _reverse_kl_from_distribution_core( + logits_norm, + exp_logits, + sum_exp_logits, + target, + grad_output, + logits_scale_factor, + target_format, + group, + temperature, + ) + + +def _cross_entropy_from_distribution_core( + logits_norm: torch.Tensor, # (*batch, vocab) + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + target: torch.Tensor, # (*batch, vocab) + grad_output: float | None, + logits_scale_factor: float, + target_format: TargetFormat, + group: ProcessGroup | None = None, + temperature: float = 1.0, return_kl_loss: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) + """ + Cross-entropy / forward-KL math from a precomputed student softmax (adding a teacher softmax when the + target is logits). + This plain (un-compiled) core is shared between the public `_fused_cross_entropy_base_from_distribution` + wrapper and the monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. + """ if target_format == TargetFormat.logits: - target_logits_norm, exp_logits_targets, sum_exp_target_logits, _ = fused_softmax_base( + target_logits_norm, exp_logits_targets, sum_exp_target_logits, _ = _softmax_base( target, logits_scale_factor / temperature, group ) target = exp_logits_targets / sum_exp_target_logits.unsqueeze(-1) @@ -212,6 +248,32 @@ def _fused_cross_entropy_base_from_distribution( return per_sample_loss, grad +@torch.compile +def _fused_cross_entropy_base_from_distribution( + logits: torch.Tensor, # (*batch, vocab) + target: torch.Tensor, # (*batch, vocab) + grad_output: float | None, + logits_scale_factor: float, + target_format: TargetFormat, + group: ProcessGroup | None = None, + temperature: float = 1.0, + return_kl_loss: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) + logits_norm, exp_logits, sum_exp_logits, _ = _softmax_base(logits, logits_scale_factor, group) + return _cross_entropy_from_distribution_core( + logits_norm, + exp_logits, + sum_exp_logits, + target, + grad_output, + logits_scale_factor, + target_format, + group, + temperature, + return_kl_loss, + ) + + def _predicted_logits_from_labels( logits: torch.Tensor, # (*batch, vocab) target: torch.Tensor, # (*batch,) diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 3cda19058..89e310259 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -3,7 +3,13 @@ import torch from fast_llm.core.distributed import ProcessGroup -from fast_llm.functional.entropy_loss import _predicted_logits_from_labels, _softmax_base +from fast_llm.functional.config import EntropyLossType, TargetFormat +from fast_llm.functional.entropy_loss import ( + _cross_entropy_from_distribution_core, + _predicted_logits_from_labels, + _reverse_kl_from_distribution_core, + _softmax_base, +) from fast_llm.utils import Assert @@ -20,8 +26,12 @@ class MonolithicLossSpec(typing.NamedTuple): logits_scale_factor: float grad_output: float | None divisor: float - target: torch.Tensor | None = None # Cross-entropy labels. - loss_mask: torch.Tensor | None = None # z-loss mask (cross-entropy derives its own from `target >= 0`). + target: torch.Tensor | None = None # Cross-entropy labels, or a teacher distribution (logits/probabilities). + loss_mask: torch.Tensor | None = None # z-loss / distribution-loss mask (cross-entropy derives its own). + # Distribution losses (cross-entropy / forward-KL / reverse-KL from a teacher distribution) only. + target_format: TargetFormat = TargetFormat.logits + entropy_loss_type: EntropyLossType = EntropyLossType.cross_entropy + temperature: float = 1.0 class MonolithicLossOutput(typing.NamedTuple): @@ -43,7 +53,14 @@ def _monolithic_core( z_loss_mask: torch.Tensor | None, z_loss_grad_output: float | None, z_loss_divisor: float, -) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + distribution_target: torch.Tensor | None, + distribution_grad_output: float | None, + distribution_divisor: float, + distribution_loss_mask: torch.Tensor | None, + distribution_target_format: TargetFormat, + distribution_entropy_loss_type: EntropyLossType, + distribution_temperature: float, +) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: """ The monolithic head-loss kernel: one shared softmax over the logits, then every enabled loss's scalar and gradient contribution. Losses with a required tensor input are toggled off by passing it as `None` @@ -94,6 +111,46 @@ def _monolithic_core( z_loss_grad = z_loss_grad_base.unsqueeze(-1) * exp_logits grad = z_loss_grad if grad is None else grad + z_loss_grad + distribution_loss = None + if distribution_target is not None: + distribution_grad_output_scaled = ( + None + if distribution_grad_output is None + else distribution_grad_output / distribution_divisor * logits_scale_factor + ) + if distribution_entropy_loss_type == EntropyLossType.reverse_kl: + per_sample_loss, distribution_grad = _reverse_kl_from_distribution_core( + logits_norm, + exp_logits, + sum_exp_logits, + distribution_target, + distribution_grad_output_scaled, + logits_scale_factor, + distribution_target_format, + group, + distribution_temperature, + ) + else: + per_sample_loss, distribution_grad = _cross_entropy_from_distribution_core( + logits_norm, + exp_logits, + sum_exp_logits, + distribution_target, + distribution_grad_output_scaled, + logits_scale_factor, + distribution_target_format, + group, + distribution_temperature, + return_kl_loss=distribution_entropy_loss_type == EntropyLossType.forward_kl, + ) + if distribution_loss_mask is not None: + per_sample_loss = per_sample_loss * distribution_loss_mask + distribution_loss = per_sample_loss.sum() / distribution_divisor + if distribution_grad is not None: + if distribution_loss_mask is not None: + distribution_grad = distribution_grad * distribution_loss_mask.unsqueeze(-1) + grad = distribution_grad if grad is None else grad + distribution_grad + if grad is not None: grad = grad.to(logits.dtype) if grad_logits is None: @@ -101,7 +158,7 @@ def _monolithic_core( else: grad_logits.add_(grad) - return cross_entropy_loss, z_loss, grad_logits + return cross_entropy_loss, z_loss, distribution_loss, grad_logits def monolithic_head_loss_forward_backward( @@ -121,15 +178,16 @@ def monolithic_head_loss_forward_backward( specs_by_kind: dict[str, MonolithicLossSpec] = {} for spec in specs: Assert.eq(spec.logits_scale_factor, logits_scale_factor) - if spec.kind not in ("cross_entropy", "z_loss"): + if spec.kind not in ("cross_entropy", "z_loss", "entropy_from_distribution"): raise NotImplementedError(spec.kind) assert spec.kind not in specs_by_kind, spec.kind specs_by_kind[spec.kind] = spec cross_entropy_spec = specs_by_kind.get("cross_entropy") z_loss_spec = specs_by_kind.get("z_loss") + distribution_spec = specs_by_kind.get("entropy_from_distribution") - cross_entropy_loss, z_loss, grad_logits = _monolithic_core( + cross_entropy_loss, z_loss, distribution_loss, grad_logits = _monolithic_core( logits, group, logits_scale_factor, @@ -141,7 +199,22 @@ def monolithic_head_loss_forward_backward( z_loss_mask=None if z_loss_spec is None else z_loss_spec.loss_mask, z_loss_grad_output=None if z_loss_spec is None else z_loss_spec.grad_output, z_loss_divisor=1.0 if z_loss_spec is None else z_loss_spec.divisor, + distribution_target=None if distribution_spec is None else distribution_spec.target, + distribution_grad_output=None if distribution_spec is None else distribution_spec.grad_output, + distribution_divisor=1.0 if distribution_spec is None else distribution_spec.divisor, + distribution_loss_mask=None if distribution_spec is None else distribution_spec.loss_mask, + distribution_target_format=( + TargetFormat.logits if distribution_spec is None else distribution_spec.target_format + ), + distribution_entropy_loss_type=( + EntropyLossType.cross_entropy if distribution_spec is None else distribution_spec.entropy_loss_type + ), + distribution_temperature=1.0 if distribution_spec is None else distribution_spec.temperature, ) - loss_by_kind = {"cross_entropy": cross_entropy_loss, "z_loss": z_loss} + loss_by_kind = { + "cross_entropy": cross_entropy_loss, + "z_loss": z_loss, + "entropy_from_distribution": distribution_loss, + } return [MonolithicLossOutput(loss=loss_by_kind[spec.kind]) for spec in specs], grad_logits diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 5bcd9e035..219a84111 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -698,6 +698,112 @@ def test_monolithic_loss( ) +def _test_monolithic_distillation_loss( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + accumulate, + target_format, + entropy_loss_type, + temperature, + group=None, +): + # The monolithic from-distribution branch (teacher target) must reproduce `fused_entropy_loss_forward_backward`. + logits, target, loss_mask = _get_lm_loss_inputs(num_columns, loss_masking, target_format, batch_shape, dtype) + local_logits = split_op(logits, group, -1).contiguous() + local_target = split_op(target, group, -1).contiguous() + divisor = local_logits.shape[:-1].numel() + out_ref, grad_ref = fused_entropy_loss_forward_backward( + logits=local_logits, + target=local_target, + loss_mask=loss_mask, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + temperature=temperature, + target_format=target_format, + entropy_loss_type=entropy_loss_type, + divisor=divisor, + ) + spec = MonolithicLossSpec( + kind="entropy_from_distribution", + name="distillation", + weight=1.0, + logits_scale_factor=logits_scale_factor, + grad_output=grad_output, + divisor=divisor, + target=local_target, + loss_mask=loss_mask, + target_format=target_format, + entropy_loss_type=entropy_loss_type, + temperature=temperature, + ) + if accumulate and grad_output is not None: + previous_grad = torch.randn_like(local_logits) + grad_ref = grad_ref + previous_grad + grad_logits = previous_grad.clone() + else: + grad_logits = None + outputs, grad_mono = monolithic_head_loss_forward_backward( + local_logits, [spec], group=group, grad_logits=grad_logits + ) + threshold = 1e-5 if dtype == DataType.float32 else 1e-4 + Assert.rms_close_relative(outputs[0].loss, out_ref, threshold, 1e-6) + if grad_output is not None: + Assert.rms_close_relative(grad_mono, grad_ref, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-6) + else: + assert grad_mono is None + + +_MONOLITHIC_DISTRIBUTION_CASES = ( + # (target_format, entropy_loss_type, temperature) + (TargetFormat.logits, EntropyLossType.cross_entropy, 1.0), + (TargetFormat.logits, EntropyLossType.forward_kl, 1.0), + (TargetFormat.logits, EntropyLossType.reverse_kl, 1.0), + (TargetFormat.logits, EntropyLossType.cross_entropy, 2.0), # Temperature + (TargetFormat.probabilities, EntropyLossType.cross_entropy, 1.0), + (TargetFormat.probabilities, EntropyLossType.forward_kl, 1.0), + (TargetFormat.probabilities, EntropyLossType.reverse_kl, 1.0), +) + + +@pytest.mark.slow +@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) +@pytest.mark.parametrize( + ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), + _LOSS_PARAMETERS, +) +@pytest.mark.parametrize(("target_format", "entropy_loss_type", "temperature"), _MONOLITHIC_DISTRIBUTION_CASES) +def test_monolithic_distillation_loss( + target_format, + entropy_loss_type, + temperature, + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + block_size, + accumulate, +): + _test_monolithic_distillation_loss( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + accumulate, + target_format, + entropy_loss_type, + temperature, + ) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -917,6 +1023,26 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa kinds, test_context.group, ) + # Monolithic from-distribution (distillation) losses + for target_format, entropy_loss_type, temperature in _MONOLITHIC_DISTRIBUTION_CASES: + with test_context.subtest( + base_path, f"monolithic_dist-{entropy_loss_type}-{target_format}-{temperature}-{suffix}", 2 + ) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_monolithic_distillation_loss( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + accumulate, + target_format, + entropy_loss_type, + temperature, + test_context.group, + ) @pytest.mark.slow @@ -960,6 +1086,10 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): "grpo_metrics-False", "grpo_metrics-True", *(f"monolithic-{"_".join(kinds)}" for kinds in _MONOLITHIC_KINDS), + *( + f"monolithic_dist-{entropy_loss_type}-{target_format}-{temperature}" + for target_format, entropy_loss_type, temperature in _MONOLITHIC_DISTRIBUTION_CASES + ), ), ) def test_lm_loss_distributed( From e4d5845a27d2aeb775b90bfda272de7bf46a9b54 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Sat, 27 Jun 2026 12:19:21 -0400 Subject: [PATCH 04/28] Monolithic head-loss kernel: wire distillation + fix dropped temperature (#507) Wire LanguageModelDistillationLoss into the monolithic kernel via get_monolithic_spec (the from-distribution kernel branch already exists), so distillation and the label-CE + distillation combo run on one shared softmax. Also fix a pre-existing bug: LanguageModelDistillationLoss._forward_backward never passed the configured `temperature` to the loss kernel, so the teacher softmax always used temperature 1.0 regardless of the config. Pass it in both the per-loss and monolithic paths. Default is unchanged (1.0); this only affects configs that explicitly set a non-unit temperature. Head-level parity tests add distillation and label+distillation combos on the fused path, plus a non-unit-temperature case on both paths. The reference now scales the teacher logits by logits_scale_factor / temperature before the softmax (identity for the existing scale-1, temperature-1 configs). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../language_model/loss/entropy_loss.py | 16 ++++++++++++++ tests/layers/test_lm_head.py | 22 ++++++++++++++++--- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py index 1711a298c..4d024729d 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -73,10 +73,26 @@ def _forward_backward( grad_logits=grad_logits, group=self._parallel_dim.group if self._vocab_parallel else None, logits_scale_factor=self._logits_scale_factor, + temperature=self._config.temperature, target_format=TargetFormat.logits, entropy_loss_type=self._config.loss_type, divisor=self._get_label_count(kwargs), ) + def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + return MonolithicLossSpec( + kind="entropy_from_distribution", + name=self.name, + weight=self._weight, + logits_scale_factor=self._logits_scale_factor, + grad_output=self._get_grad_output(kwargs), + divisor=self._get_label_count(kwargs), + target=self._get_reference_model_logits(self._config.reference_model, kwargs, split_index), + loss_mask=self._get_loss_mask(kwargs, split_index), + target_format=TargetFormat.logits, + entropy_loss_type=self._config.loss_type, + temperature=self._config.temperature, + ) + def get_preprocessing_config(self) -> dict[str, typing.Any]: return {"return_prediction_mask": True} diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index a7f8bfb73..6c468449a 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -27,6 +27,7 @@ class LMHeadTestConfig: name: str label_loss: bool | float = False distillation_loss: bool | float = False + distillation_temperature: float = 1.0 z_loss: bool | float = False grpo_loss: bool | float = False gspo_loss: bool | float = False @@ -71,6 +72,8 @@ def get_config(self) -> GPTModelConfig: losses["label"]["weight"] = self.label_loss if self.distillation_loss is not False: losses["distillation"] = {"type": "distillation", "reference_model": "distillation"} + if self.distillation_temperature != 1.0: + losses["distillation"]["temperature"] = self.distillation_temperature if isinstance(self.distillation_loss, float): losses["distillation"]["weight"] = self.distillation_loss if self.z_loss is not False: @@ -242,9 +245,12 @@ def get_reference_outputs( names_losses_weights.append(("label", label_loss, float(self.actual_label_loss))) if self.distillation_loss is not False: + # Teacher logits are scaled by `logits_scale_factor / temperature` before the softmax, matching the kernel. + teacher_logits = kwargs[f"reference_distillation_hidden_states"]["head.logits"].float() + teacher_logits = teacher_logits * (self.logits_scale_factor / self.distillation_temperature) distillation_loss = torch.nn.functional.cross_entropy( logits, - torch.softmax(kwargs[f"reference_distillation_hidden_states"]["head.logits"], -1), + torch.softmax(teacher_logits, -1), reduction="mean" if loss_mask is None else "none", ) if loss_mask is not None: @@ -362,9 +368,11 @@ def _add_configs(base_name: str, **kwargs): _add_configs("label_and_distillation_loss", label_loss=True, distillation_loss=True) _add_configs("label_and_z_loss_weighted", label_loss=True, z_loss=0.5) _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 ("fused") head-loss path. So far only cross-entropy is computed in the monolithic kernel; -# other losses (e.g. z-loss) fall back to their own implementation and accumulate into the same gradient. +# Monolithic ("fused") head-loss path. Cross-entropy, z-loss, and the from-distribution (distillation) losses +# run in the monolithic kernel; unsupported losses fall back to their own implementation and accumulate into +# the same gradient. _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) @@ -372,6 +380,14 @@ def _add_configs(base_name: str, **kwargs): _add_configs("fused_tied_embedding_weight", loss_implementation="fused", tied_embedding_weight=True) _add_configs("fused_multi_token_prediction", loss_implementation="fused", prediction_heads=2) _add_configs("fused_label_and_z_loss_weighted", loss_implementation="fused", label_loss=True, z_loss=0.5) +_add_configs("fused_distillation_loss", loss_implementation="fused", distillation_loss=True) +_add_configs("fused_label_and_distillation_loss", loss_implementation="fused", label_loss=True, distillation_loss=True) +_add_configs( + "fused_distillation_loss_temperature", + loss_implementation="fused", + distillation_loss=True, + distillation_temperature=2.0, +) @pytest.mark.slow From c872b112333617d6981248460ea617768d89753c Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Sat, 27 Jun 2026 12:44:23 -0400 Subject: [PATCH 05/28] Monolithic head-loss kernel: GRPO objective + new_logprobs (#507) Add the GRPO policy-gradient objective to the monolithic kernel: per-token IS-ratio clipping over the shared softmax, plus the new_logprobs_mean metric as a side output. Reuses the shared predicted-logits; uses out-of-place unsqueeze on sum_exp_logits (the per-loss kernel mutates it in place). LanguageModelGRPOLoss gets get_monolithic_spec / register_monolithic_outputs; it falls back to the per-loss path when metrics are enabled (kernel metrics land in a later step). Validated against fused_grpo_loss_forward_backward (loss, grad, new_logprobs) at the functional, head, and tensor-parallel levels. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../layers/language_model/loss/monolithic.py | 93 +++++++++++++++++-- .../language_model/loss/policy_gradient.py | 28 ++++++ tests/layers/test_lm_head.py | 1 + tests/layers/test_lm_losses.py | 81 ++++++++++++++++ 4 files changed, 196 insertions(+), 7 deletions(-) diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 89e310259..c53e87efc 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -26,18 +26,27 @@ class MonolithicLossSpec(typing.NamedTuple): logits_scale_factor: float grad_output: float | None divisor: float - target: torch.Tensor | None = None # Cross-entropy labels, or a teacher distribution (logits/probabilities). - loss_mask: torch.Tensor | None = None # z-loss / distribution-loss mask (cross-entropy derives its own). + target: torch.Tensor | None = ( + None # Cross-entropy / GRPO labels, or a teacher distribution (logits/probabilities). + ) + loss_mask: torch.Tensor | None = None # z-loss / distribution-loss mask (cross-entropy / GRPO derive their own). # Distribution losses (cross-entropy / forward-KL / reverse-KL from a teacher distribution) only. target_format: TargetFormat = TargetFormat.logits entropy_loss_type: EntropyLossType = EntropyLossType.cross_entropy temperature: float = 1.0 + # Policy-gradient (GRPO) only. + advantages: torch.Tensor | None = None + old_log_probabilities: torch.Tensor | None = None + epsilon_low: float = 0.2 + epsilon_high: float = 0.2 + num_labels_in_seq: torch.Tensor | None = None # Per-sequence label count; enables the `new_logprobs_mean` metric. class MonolithicLossOutput(typing.NamedTuple): """Per-loss outputs returned by the monolithic kernel, registered via `register_monolithic_outputs`.""" loss: torch.Tensor | None = None + new_logprobs_mean: torch.Tensor | None = None # GRPO only. @torch.compile @@ -60,7 +69,22 @@ def _monolithic_core( distribution_target_format: TargetFormat, distribution_entropy_loss_type: EntropyLossType, distribution_temperature: float, -) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + grpo_target: torch.Tensor | None, + grpo_advantages: torch.Tensor | None, + grpo_old_log_probabilities: torch.Tensor | None, + grpo_grad_output: float | None, + grpo_divisor: float, + grpo_epsilon_low: float, + grpo_epsilon_high: float, + grpo_num_labels_in_seq: torch.Tensor | None, +) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, +]: """ The monolithic head-loss kernel: one shared softmax over the logits, then every enabled loss's scalar and gradient contribution. Losses with a required tensor input are toggled off by passing it as `None` @@ -151,6 +175,45 @@ def _monolithic_core( distribution_grad = distribution_grad * distribution_loss_mask.unsqueeze(-1) grad = distribution_grad if grad is None else grad + distribution_grad + grpo_loss = None + grpo_new_logprobs_mean = None + if grpo_target is not None: + grpo_loss_mask = grpo_target >= 0 + grpo_predicted_logits, grpo_target_masked, grpo_target_mask = _predicted_logits_from_labels( + logits_norm, grpo_target, grpo_loss_mask, group + ) + new_log_probs = grpo_predicted_logits - sum_exp_logits.log() + probability_ratio = (new_log_probs - grpo_old_log_probabilities).exp() + grpo_losses = -torch.min( + probability_ratio * grpo_advantages, + torch.clamp(probability_ratio, 1 - grpo_epsilon_low, 1 + grpo_epsilon_high) * grpo_advantages, + ) + grpo_loss = (grpo_losses * grpo_loss_mask).sum() / grpo_divisor + if grpo_num_labels_in_seq is not None: + # Sum of per-sequence mean log-probs; clamp avoids 0/0 at fully-masked documents (also loss-masked). + grpo_new_logprobs_mean = (new_log_probs * grpo_loss_mask / grpo_num_labels_in_seq.clamp(min=1)).sum() + if grpo_grad_output is not None: + grad_output = grpo_grad_output / grpo_divisor * logits_scale_factor + # grad[a>=0] = -a * (ratio <= 1 + epsilon_high); grad[a<=0] = a * (ratio >= 1 - epsilon_low). + probability_ratio_grad = ( + grad_output + * ( + torch.clamp_min(grpo_advantages, 0) * (probability_ratio <= 1 + grpo_epsilon_high) + + torch.clamp_max(grpo_advantages, 0) * (probability_ratio >= 1 - grpo_epsilon_low) + ) + * grpo_loss_mask + ) + # d(probability_ratio)/d(logits) = -probability_ratio * (predicted_probabilities - target_probabilities). + predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) + grpo_grad = (probability_ratio_grad * probability_ratio).unsqueeze( + -1 + ) * predicted_probabilities.scatter_add( + -1, + grpo_target_masked.unsqueeze(-1), + -(grpo_loss_mask if grpo_target_mask is None else grpo_target_mask).unsqueeze(-1).to(torch.float32), + ) + grad = grpo_grad if grad is None else grad + grpo_grad + if grad is not None: grad = grad.to(logits.dtype) if grad_logits is None: @@ -158,7 +221,7 @@ def _monolithic_core( else: grad_logits.add_(grad) - return cross_entropy_loss, z_loss, distribution_loss, grad_logits + return cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grad_logits def monolithic_head_loss_forward_backward( @@ -178,7 +241,7 @@ def monolithic_head_loss_forward_backward( specs_by_kind: dict[str, MonolithicLossSpec] = {} for spec in specs: Assert.eq(spec.logits_scale_factor, logits_scale_factor) - if spec.kind not in ("cross_entropy", "z_loss", "entropy_from_distribution"): + if spec.kind not in ("cross_entropy", "z_loss", "entropy_from_distribution", "grpo"): raise NotImplementedError(spec.kind) assert spec.kind not in specs_by_kind, spec.kind specs_by_kind[spec.kind] = spec @@ -186,8 +249,9 @@ def monolithic_head_loss_forward_backward( cross_entropy_spec = specs_by_kind.get("cross_entropy") z_loss_spec = specs_by_kind.get("z_loss") distribution_spec = specs_by_kind.get("entropy_from_distribution") + grpo_spec = specs_by_kind.get("grpo") - cross_entropy_loss, z_loss, distribution_loss, grad_logits = _monolithic_core( + cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grad_logits = _monolithic_core( logits, group, logits_scale_factor, @@ -210,11 +274,26 @@ def monolithic_head_loss_forward_backward( EntropyLossType.cross_entropy if distribution_spec is None else distribution_spec.entropy_loss_type ), distribution_temperature=1.0 if distribution_spec is None else distribution_spec.temperature, + grpo_target=None if grpo_spec is None else grpo_spec.target, + grpo_advantages=None if grpo_spec is None else grpo_spec.advantages, + grpo_old_log_probabilities=None if grpo_spec is None else grpo_spec.old_log_probabilities, + grpo_grad_output=None if grpo_spec is None else grpo_spec.grad_output, + grpo_divisor=1.0 if grpo_spec is None else grpo_spec.divisor, + grpo_epsilon_low=0.2 if grpo_spec is None else grpo_spec.epsilon_low, + grpo_epsilon_high=0.2 if grpo_spec is None else grpo_spec.epsilon_high, + grpo_num_labels_in_seq=None if grpo_spec is None else grpo_spec.num_labels_in_seq, ) loss_by_kind = { "cross_entropy": cross_entropy_loss, "z_loss": z_loss, "entropy_from_distribution": distribution_loss, + "grpo": grpo_loss, } - return [MonolithicLossOutput(loss=loss_by_kind[spec.kind]) for spec in specs], grad_logits + return [ + MonolithicLossOutput( + loss=loss_by_kind[spec.kind], + new_logprobs_mean=grpo_new_logprobs_mean if spec.kind == "grpo" else None, + ) + for spec in specs + ], grad_logits diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index a024d4232..dfa1fe8d0 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -19,6 +19,7 @@ LanguageModelPolicyGradientLossConfig, ) from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.monolithic import MonolithicLossOutput, MonolithicLossSpec from fast_llm.utils import Assert @@ -142,6 +143,33 @@ def _forward_backward( return loss, grad + def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + if self._config.metrics != GRPOMetricsLevel.none: + # The metric family isn't emitted by the monolithic kernel yet; run through the per-loss path. + return None + return MonolithicLossSpec( + kind="grpo", + name=self.name, + weight=self._weight, + logits_scale_factor=self._logits_scale_factor, + grad_output=self._get_grad_output(kwargs), + divisor=self._get_label_count(kwargs), + target=self._get_labels(kwargs, split_index), + advantages=self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), + old_log_probabilities=self._prepare_target( + kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index + ), + epsilon_low=self._config.epsilon_low, + epsilon_high=self._config.epsilon_high, + num_labels_in_seq=self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index), + ) + + def register_monolithic_outputs( + self, output: MonolithicLossOutput, kwargs: dict[str, typing.Any], losses: dict | None + ) -> None: + super().register_monolithic_outputs(output, kwargs, losses) + self._register_new_logprobs(output.new_logprobs_mean, kwargs, losses) + def _register_extra_metrics( self, logits: torch.Tensor, diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 6c468449a..b772532a0 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -388,6 +388,7 @@ def _add_configs(base_name: str, **kwargs): distillation_loss=True, distillation_temperature=2.0, ) +_add_configs("fused_grpo_loss", loss_implementation="fused", grpo_loss=True) @pytest.mark.slow diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 219a84111..bbdf14d3c 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -758,6 +758,58 @@ def _test_monolithic_distillation_loss( assert grad_mono is None +def _test_monolithic_grpo_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, group=None +): + # The monolithic GRPO branch must reproduce `fused_grpo_loss_forward_backward` (loss, grad, new_logprobs). + logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs( + num_columns, loss_masking, batch_shape, dtype + ) + local_logits = split_op(logits, group, -1).contiguous() + num_labels = int((target >= 0).sum().item()) + num_labels_in_seq = torch.where( + target >= 0, + torch.full(batch_shape, num_labels, dtype=torch.int32, device=target.device), + torch.zeros(batch_shape, dtype=torch.int32, device=target.device), + ) + divisor = max(num_labels, 1) + previous_grad = torch.randn_like(local_logits) if accumulate and grad_output is not None else None + out_ref, grad_ref, new_logprobs_ref = fused_grpo_loss_forward_backward( + local_logits, + target, + advantages, + old_log_probabilities, + grad_logits=None if previous_grad is None else previous_grad.clone(), + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + num_labels_in_seq=num_labels_in_seq, + divisor=divisor, + ) + spec = MonolithicLossSpec( + kind="grpo", + name="grpo", + weight=1.0, + logits_scale_factor=logits_scale_factor, + grad_output=grad_output, + divisor=divisor, + target=target, + advantages=advantages, + old_log_probabilities=old_log_probabilities, + num_labels_in_seq=num_labels_in_seq, + ) + outputs, grad_mono = monolithic_head_loss_forward_backward( + local_logits, [spec], group=group, grad_logits=None if previous_grad is None else previous_grad.clone() + ) + threshold = 1e-5 if dtype == DataType.float32 else 1e-4 + Assert.rms_close_relative(outputs[0].loss, out_ref, threshold, 1e-6) + Assert.rms_close_relative(outputs[0].new_logprobs_mean, new_logprobs_ref, threshold, 1e-6) + if grad_output is not None: + Assert.rms_close_relative(grad_mono, grad_ref, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-6) + else: + assert grad_mono is None + + _MONOLITHIC_DISTRIBUTION_CASES = ( # (target_format, entropy_loss_type, temperature) (TargetFormat.logits, EntropyLossType.cross_entropy, 1.0), @@ -804,6 +856,20 @@ def test_monolithic_distillation_loss( ) +@pytest.mark.slow +@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) +@pytest.mark.parametrize( + ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), + _LOSS_PARAMETERS, +) +def test_monolithic_grpo_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate +): + _test_monolithic_grpo_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate + ) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -1023,6 +1089,20 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa kinds, test_context.group, ) + # Monolithic GRPO objective + with test_context.subtest(base_path, f"monolithic_grpo-{suffix}", 2) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_monolithic_grpo_loss( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + accumulate, + test_context.group, + ) # Monolithic from-distribution (distillation) losses for target_format, entropy_loss_type, temperature in _MONOLITHIC_DISTRIBUTION_CASES: with test_context.subtest( @@ -1085,6 +1165,7 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): "grpo", "grpo_metrics-False", "grpo_metrics-True", + "monolithic_grpo", *(f"monolithic-{"_".join(kinds)}" for kinds in _MONOLITHIC_KINDS), *( f"monolithic_dist-{entropy_loss_type}-{target_format}-{temperature}" From c8fd86fee78748d2aacf620d174d3da90fce3c6e Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 29 Jun 2026 13:30:57 -0400 Subject: [PATCH 06/28] Monolithic head-loss kernel: GRPO metrics + entropy, kill the second softmax (#507) Emit the full GRPO metric family (ratio/KL/clip/advantage stats + optional per-token entropy) from the monolithic kernel's shared softmax instead of a second softmax pass over the logits (the #494 redundancy). Entropy is the only metric needing the vocab axis; the rest reuse the already-computed new_log_probs. Extract GRPOMetrics + the pure metric aggregation into a new grpo_metrics module so both the baseline compute_grpo_metrics and the kernel share one source of truth (and to avoid a monolithic<->policy_gradient import cycle). The GRPO loss no longer falls back to the per-loss path when metrics are enabled. Validated vs the independent loop reference: functional (with/without entropy), head (per-loss + fused metrics configs), and tensor-parallel. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../language_model/loss/grpo_metrics.py | 79 ++++++++++++++++ .../layers/language_model/loss/monolithic.py | 93 ++++++++++++------- .../language_model/loss/policy_gradient.py | 73 ++++----------- tests/layers/test_lm_head.py | 44 ++++++++- tests/layers/test_lm_losses.py | 83 +++++++++++++++++ 5 files changed, 286 insertions(+), 86 deletions(-) create mode 100644 fast_llm/layers/language_model/loss/grpo_metrics.py diff --git a/fast_llm/layers/language_model/loss/grpo_metrics.py b/fast_llm/layers/language_model/loss/grpo_metrics.py new file mode 100644 index 000000000..e15045db5 --- /dev/null +++ b/fast_llm/layers/language_model/loss/grpo_metrics.py @@ -0,0 +1,79 @@ +import typing + +import torch + +from fast_llm.core.distributed import ProcessGroup, ReduceOp, all_reduce + + +class GRPOMetrics(typing.NamedTuple): + old_logprobs: torch.Tensor + ratio_new_old: torch.Tensor + ratio_new_old_sum: torch.Tensor + ratio_new_old_squared_sum: torch.Tensor + kl_new_old: torch.Tensor + clipped_ratio_fraction: torch.Tensor + advantage: torch.Tensor + max_advantage: torch.Tensor + min_advantage: torch.Tensor + num_tokens: torch.Tensor + entropy: torch.Tensor | None + + +def _grpo_metrics_core( + logits_norm: torch.Tensor, # (*batch, vocab) + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + new_log_probs: torch.Tensor, # (*batch,) — predicted_logits - log(sum_exp_logits) + advantages: torch.Tensor, # (*batch,) + old_log_probabilities: torch.Tensor, # (*batch,) + loss_mask: torch.Tensor, # (*batch,) bool, == target >= 0 + label_counts: torch.Tensor, # (*batch,) — global per-sequence count broadcast per token + epsilon_low: float, + epsilon_high: float, + group: ProcessGroup | None = None, + compute_entropy: bool = False, +) -> GRPOMetrics: + """ + GRPO metric family from a precomputed student softmax and per-token new log-probs. Entropy is the only + term needing the full vocab axis (a sum over the local slice plus a tensor-parallel all-reduce); every + other term is per-token. + + This plain (un-compiled) core is shared between the public `compute_grpo_metrics` wrapper and the + monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary — so the metrics + reuse the loss kernel's softmax instead of recomputing it. + """ + mask = loss_mask.float() + masked = mask / label_counts.float().clamp(min=1) + + log_ratio = new_log_probs - old_log_probabilities + ratio = log_ratio.exp() + clipped = (ratio < 1.0 - epsilon_low) | (ratio > 1.0 + epsilon_high) + kl = ratio - log_ratio - 1.0 + + neg_inf = advantages.new_full((), float("-inf")) + pos_inf = advantages.new_full((), float("inf")) + + entropy: torch.Tensor | None = None + if compute_entropy: + # exp_logits and logits_norm are local vocab slices — sum over the local slice, then all-reduce + # across the tensor-parallel group to recover the global E_p[logit_norm] before dividing by the + # already-global sum_exp_logits. + weighted_logits_sum = (exp_logits * logits_norm).sum(-1) + if group is not None: + all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group) + entropy_per_token = sum_exp_logits.log() - weighted_logits_sum / sum_exp_logits + entropy = (entropy_per_token * masked).sum() + + return GRPOMetrics( + old_logprobs=(old_log_probabilities * masked).sum(), + ratio_new_old=(ratio * masked).sum(), + ratio_new_old_sum=(ratio * mask).sum(), + ratio_new_old_squared_sum=(ratio * ratio * mask).sum(), + kl_new_old=(kl * masked).sum(), + clipped_ratio_fraction=(clipped.float() * masked).sum(), + advantage=(advantages * masked).sum(), + max_advantage=torch.where(loss_mask, advantages, neg_inf).max(), + min_advantage=torch.where(loss_mask, advantages, pos_inf).min(), + num_tokens=mask.sum(), + entropy=entropy, + ) diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index c53e87efc..a3aab5bc1 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -10,6 +10,7 @@ _reverse_kl_from_distribution_core, _softmax_base, ) +from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, _grpo_metrics_core from fast_llm.utils import Assert @@ -40,6 +41,8 @@ class MonolithicLossSpec(typing.NamedTuple): epsilon_low: float = 0.2 epsilon_high: float = 0.2 num_labels_in_seq: torch.Tensor | None = None # Per-sequence label count; enables the `new_logprobs_mean` metric. + compute_metrics: bool = False # Emit the GRPO metric family from the shared softmax. + compute_entropy: bool = False # Also emit per-token entropy (the only metric needing the vocab axis). class MonolithicLossOutput(typing.NamedTuple): @@ -47,6 +50,7 @@ class MonolithicLossOutput(typing.NamedTuple): loss: torch.Tensor | None = None new_logprobs_mean: torch.Tensor | None = None # GRPO only. + metrics: GRPOMetrics | None = None # GRPO only. @torch.compile @@ -77,12 +81,15 @@ def _monolithic_core( grpo_epsilon_low: float, grpo_epsilon_high: float, grpo_num_labels_in_seq: torch.Tensor | None, + grpo_compute_metrics: bool, + grpo_compute_entropy: bool, ) -> tuple[ torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, + "GRPOMetrics | None", torch.Tensor | None, ]: """ @@ -177,6 +184,7 @@ def _monolithic_core( grpo_loss = None grpo_new_logprobs_mean = None + grpo_metrics = None if grpo_target is not None: grpo_loss_mask = grpo_target >= 0 grpo_predicted_logits, grpo_target_masked, grpo_target_mask = _predicted_logits_from_labels( @@ -192,6 +200,22 @@ def _monolithic_core( if grpo_num_labels_in_seq is not None: # Sum of per-sequence mean log-probs; clamp avoids 0/0 at fully-masked documents (also loss-masked). grpo_new_logprobs_mean = (new_log_probs * grpo_loss_mask / grpo_num_labels_in_seq.clamp(min=1)).sum() + if grpo_compute_metrics: + # The metric family reuses this branch's softmax + new_log_probs — no second softmax pass. + grpo_metrics = _grpo_metrics_core( + logits_norm, + exp_logits, + sum_exp_logits, + new_log_probs, + grpo_advantages, + grpo_old_log_probabilities, + grpo_loss_mask, + grpo_num_labels_in_seq, + grpo_epsilon_low, + grpo_epsilon_high, + group, + grpo_compute_entropy, + ) if grpo_grad_output is not None: grad_output = grpo_grad_output / grpo_divisor * logits_scale_factor # grad[a>=0] = -a * (ratio <= 1 + epsilon_high); grad[a<=0] = a * (ratio >= 1 - epsilon_low). @@ -221,7 +245,7 @@ def _monolithic_core( else: grad_logits.add_(grad) - return cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grad_logits + return cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad_logits def monolithic_head_loss_forward_backward( @@ -251,37 +275,41 @@ def monolithic_head_loss_forward_backward( distribution_spec = specs_by_kind.get("entropy_from_distribution") grpo_spec = specs_by_kind.get("grpo") - cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grad_logits = _monolithic_core( - logits, - group, - logits_scale_factor, - grad_logits, - ce_target=None if cross_entropy_spec is None else cross_entropy_spec.target, - ce_grad_output=None if cross_entropy_spec is None else cross_entropy_spec.grad_output, - ce_divisor=1.0 if cross_entropy_spec is None else cross_entropy_spec.divisor, - z_loss_enabled=z_loss_spec is not None, - z_loss_mask=None if z_loss_spec is None else z_loss_spec.loss_mask, - z_loss_grad_output=None if z_loss_spec is None else z_loss_spec.grad_output, - z_loss_divisor=1.0 if z_loss_spec is None else z_loss_spec.divisor, - distribution_target=None if distribution_spec is None else distribution_spec.target, - distribution_grad_output=None if distribution_spec is None else distribution_spec.grad_output, - distribution_divisor=1.0 if distribution_spec is None else distribution_spec.divisor, - distribution_loss_mask=None if distribution_spec is None else distribution_spec.loss_mask, - distribution_target_format=( - TargetFormat.logits if distribution_spec is None else distribution_spec.target_format - ), - distribution_entropy_loss_type=( - EntropyLossType.cross_entropy if distribution_spec is None else distribution_spec.entropy_loss_type - ), - distribution_temperature=1.0 if distribution_spec is None else distribution_spec.temperature, - grpo_target=None if grpo_spec is None else grpo_spec.target, - grpo_advantages=None if grpo_spec is None else grpo_spec.advantages, - grpo_old_log_probabilities=None if grpo_spec is None else grpo_spec.old_log_probabilities, - grpo_grad_output=None if grpo_spec is None else grpo_spec.grad_output, - grpo_divisor=1.0 if grpo_spec is None else grpo_spec.divisor, - grpo_epsilon_low=0.2 if grpo_spec is None else grpo_spec.epsilon_low, - grpo_epsilon_high=0.2 if grpo_spec is None else grpo_spec.epsilon_high, - grpo_num_labels_in_seq=None if grpo_spec is None else grpo_spec.num_labels_in_seq, + cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad_logits = ( + _monolithic_core( + logits, + group, + logits_scale_factor, + grad_logits, + ce_target=None if cross_entropy_spec is None else cross_entropy_spec.target, + ce_grad_output=None if cross_entropy_spec is None else cross_entropy_spec.grad_output, + ce_divisor=1.0 if cross_entropy_spec is None else cross_entropy_spec.divisor, + z_loss_enabled=z_loss_spec is not None, + z_loss_mask=None if z_loss_spec is None else z_loss_spec.loss_mask, + z_loss_grad_output=None if z_loss_spec is None else z_loss_spec.grad_output, + z_loss_divisor=1.0 if z_loss_spec is None else z_loss_spec.divisor, + distribution_target=None if distribution_spec is None else distribution_spec.target, + distribution_grad_output=None if distribution_spec is None else distribution_spec.grad_output, + distribution_divisor=1.0 if distribution_spec is None else distribution_spec.divisor, + distribution_loss_mask=None if distribution_spec is None else distribution_spec.loss_mask, + distribution_target_format=( + TargetFormat.logits if distribution_spec is None else distribution_spec.target_format + ), + distribution_entropy_loss_type=( + EntropyLossType.cross_entropy if distribution_spec is None else distribution_spec.entropy_loss_type + ), + distribution_temperature=1.0 if distribution_spec is None else distribution_spec.temperature, + grpo_target=None if grpo_spec is None else grpo_spec.target, + grpo_advantages=None if grpo_spec is None else grpo_spec.advantages, + grpo_old_log_probabilities=None if grpo_spec is None else grpo_spec.old_log_probabilities, + grpo_grad_output=None if grpo_spec is None else grpo_spec.grad_output, + grpo_divisor=1.0 if grpo_spec is None else grpo_spec.divisor, + grpo_epsilon_low=0.2 if grpo_spec is None else grpo_spec.epsilon_low, + grpo_epsilon_high=0.2 if grpo_spec is None else grpo_spec.epsilon_high, + grpo_num_labels_in_seq=None if grpo_spec is None else grpo_spec.num_labels_in_seq, + grpo_compute_metrics=False if grpo_spec is None else grpo_spec.compute_metrics, + grpo_compute_entropy=False if grpo_spec is None else grpo_spec.compute_entropy, + ) ) loss_by_kind = { @@ -294,6 +322,7 @@ def monolithic_head_loss_forward_backward( MonolithicLossOutput( loss=loss_by_kind[spec.kind], new_logprobs_mean=grpo_new_logprobs_mean if spec.kind == "grpo" else None, + metrics=grpo_metrics if spec.kind == "grpo" else None, ) for spec in specs ], grad_logits diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index dfa1fe8d0..659317721 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -3,7 +3,6 @@ import torch -from fast_llm.core.distributed import ReduceOp, all_reduce from fast_llm.engine.base_model.config import LossDef, ReductionType from fast_llm.engine.distributed.config import DistributedConfig from fast_llm.functional.config import TritonConfig @@ -18,25 +17,12 @@ LanguageModelLossKwargs, LanguageModelPolicyGradientLossConfig, ) +from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, _grpo_metrics_core from fast_llm.layers.language_model.loss.loss import LanguageModelLoss from fast_llm.layers.language_model.loss.monolithic import MonolithicLossOutput, MonolithicLossSpec from fast_llm.utils import Assert -class GRPOMetrics(typing.NamedTuple): - old_logprobs: torch.Tensor - ratio_new_old: torch.Tensor - ratio_new_old_sum: torch.Tensor - ratio_new_old_squared_sum: torch.Tensor - kl_new_old: torch.Tensor - clipped_ratio_fraction: torch.Tensor - advantage: torch.Tensor - max_advantage: torch.Tensor - min_advantage: torch.Tensor - num_tokens: torch.Tensor - entropy: torch.Tensor | None - - class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLossConfig]( LanguageModelLoss[ConfigType] ): @@ -144,9 +130,6 @@ def _forward_backward( return loss, grad def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: - if self._config.metrics != GRPOMetricsLevel.none: - # The metric family isn't emitted by the monolithic kernel yet; run through the per-loss path. - return None return MonolithicLossSpec( kind="grpo", name=self.name, @@ -162,6 +145,8 @@ def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = epsilon_low=self._config.epsilon_low, epsilon_high=self._config.epsilon_high, num_labels_in_seq=self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index), + compute_metrics=self._config.metrics != GRPOMetricsLevel.none, + compute_entropy=self._config.metrics == GRPOMetricsLevel.with_entropy, ) def register_monolithic_outputs( @@ -169,6 +154,8 @@ def register_monolithic_outputs( ) -> None: super().register_monolithic_outputs(output, kwargs, losses) self._register_new_logprobs(output.new_logprobs_mean, kwargs, losses) + if output.metrics is not None and losses is not None: + self._register_grpo_metrics(output.metrics, kwargs, losses) def _register_extra_metrics( self, @@ -189,7 +176,9 @@ def _register_extra_metrics( group=self._parallel_dim.group if self._vocab_parallel else None, compute_entropy=self._config.metrics == GRPOMetricsLevel.with_entropy, ) + self._register_grpo_metrics(metrics, kwargs, losses) + def _register_grpo_metrics(self, metrics: GRPOMetrics, kwargs: dict[str, typing.Any], losses: dict | None) -> None: num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch] for attr in ( @@ -346,44 +335,22 @@ def compute_grpo_metrics( compute_entropy: bool = False, ) -> GRPOMetrics: loss_mask = target >= 0 - mask = loss_mask.float() - masked = mask / label_counts.float().clamp(min=1) - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) predicted_logits, _, _ = fused_predicted_logits_from_labels(logits_norm, target, loss_mask, group) new_log_probs = predicted_logits - sum_exp_logits.log() - - log_ratio = new_log_probs - old_log_probabilities - ratio = log_ratio.exp() - clipped = (ratio < 1.0 - epsilon_low) | (ratio > 1.0 + epsilon_high) - kl = ratio - log_ratio - 1.0 - - neg_inf = advantages.new_full((), float("-inf")) - pos_inf = advantages.new_full((), float("inf")) - - entropy: torch.Tensor | None = None - if compute_entropy: - # exp_logits and logits_norm are local vocab slices — sum over the local slice, then all-reduce - # across the tensor-parallel group to recover the global E_p[logit_norm] before dividing by the - # already-global sum_exp_logits. - weighted_logits_sum = (exp_logits * logits_norm).sum(-1) - if group is not None: - all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group) - entropy_per_token = sum_exp_logits.log() - weighted_logits_sum / sum_exp_logits - entropy = (entropy_per_token * masked).sum() - - return GRPOMetrics( - old_logprobs=(old_log_probabilities * masked).sum(), - ratio_new_old=(ratio * masked).sum(), - ratio_new_old_sum=(ratio * mask).sum(), - ratio_new_old_squared_sum=(ratio * ratio * mask).sum(), - kl_new_old=(kl * masked).sum(), - clipped_ratio_fraction=(clipped.float() * masked).sum(), - advantage=(advantages * masked).sum(), - max_advantage=torch.where(loss_mask, advantages, neg_inf).max(), - min_advantage=torch.where(loss_mask, advantages, pos_inf).min(), - num_tokens=mask.sum(), - entropy=entropy, + return _grpo_metrics_core( + logits_norm, + exp_logits, + sum_exp_logits, + new_log_probs, + advantages, + old_log_probabilities, + loss_mask, + label_counts, + epsilon_low, + epsilon_high, + group, + compute_entropy, ) diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index b772532a0..f7742e766 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -13,7 +13,7 @@ from fast_llm.layers.language_model.loss.config import LanguageModelLossKwargs from fast_llm.models.gpt.config import GPTModelConfig from fast_llm.utils import Assert -from tests.layers.test_lm_losses import reference_grpo_loss, reference_gspo_loss +from tests.layers.test_lm_losses import reference_grpo_loss, reference_grpo_metrics, reference_gspo_loss from tests.utils.utils import get_base_model, get_stage NUM_TOKENS = 200 @@ -41,6 +41,7 @@ class LMHeadTestConfig: num_splits: int = 1 gspo_document_lengths: tuple[int, ...] | None = None loss_implementation: str = "per_loss" + grpo_metrics: str | None = None @property def actual_label_loss(self): @@ -84,6 +85,8 @@ def get_config(self) -> GPTModelConfig: losses["grpo_loss"] = {"type": "grpo"} if isinstance(self.grpo_loss, float): losses["grpo_loss"]["weight"] = self.grpo_loss + if self.grpo_metrics is not None: + losses["grpo_loss"]["metrics"] = self.grpo_metrics if self.gspo_loss is not False: losses["gspo_loss"] = {"type": "gspo"} if isinstance(self.gspo_loss, float): @@ -274,6 +277,29 @@ def get_reference_outputs( names_losses_weights.append(("grpo_loss", grpo_loss, float(self.grpo_loss))) names_losses_weights.append(("grpo_loss_new_logprobs", new_logprobs, 0.0)) + if self.grpo_metrics is not None: + # `logits` is already scaled above, so pass logits_scale_factor=1.0. + metrics = reference_grpo_metrics( + logits, + labels, + kwargs[LanguageModelLossKwargs.advantages][head._prediction_distance - 1], + kwargs[LanguageModelLossKwargs.old_log_probabilities][head._prediction_distance - 1], + kwargs[LanguageModelLossKwargs.label_counts][head._prediction_distance - 1], + epsilon_low=0.2, + epsilon_high=0.2, + logits_scale_factor=1.0, + compute_entropy=self.grpo_metrics == "with_entropy", + ) + num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch] + for attr in ("old_logprobs", "ratio_new_old", "kl_new_old", "clipped_ratio_fraction", "advantage"): + names_losses_weights.append((f"grpo_loss_{attr}", getattr(metrics, attr) / num_documents, 0.0)) + for attr in ("ratio_new_old_sum", "ratio_new_old_squared_sum", "num_tokens"): + names_losses_weights.append((f"grpo_loss_{attr}", getattr(metrics, attr), 0.0)) + names_losses_weights.append(("grpo_loss_max_advantage", metrics.max_advantage, 0.0)) + names_losses_weights.append(("grpo_loss_min_advantage", metrics.min_advantage, 0.0)) + if metrics.entropy is not None: + names_losses_weights.append(("grpo_loss_entropy", metrics.entropy / num_documents, 0.0)) + if self.gspo_loss is not False: gspo_loss, _ = reference_gspo_loss( logits, @@ -389,6 +415,22 @@ def _add_configs(base_name: str, **kwargs): distillation_temperature=2.0, ) _add_configs("fused_grpo_loss", loss_implementation="fused", grpo_loss=True) +# 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"): + _prefix = "" if _loss_implementation == "per_loss" else "fused_" + for _metrics in ("basic", "with_entropy"): + _suffix = "metrics" if _metrics == "basic" else "entropy" + for _loss_masking in (False, True): + _lm_head_test_configs.append( + LMHeadTestConfig( + f"{_prefix}grpo_loss_{_suffix}{'_masked' if _loss_masking else ''}", + grpo_loss=True, + grpo_metrics=_metrics, + loss_masking=_loss_masking, + loss_implementation=_loss_implementation, + ) + ) @pytest.mark.slow diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index bbdf14d3c..8c3cb3b73 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -810,6 +810,50 @@ def _test_monolithic_grpo_loss( assert grad_mono is None +def _test_monolithic_grpo_metrics( + batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy, group=None +): + # The monolithic GRPO metric family (emitted from the shared softmax) must match the loop reference. + logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs( + num_columns, loss_masking, batch_shape, dtype + ) + # Different denominators per position so the per-token-mean broadcasting is exercised. + label_counts = (torch.arange(target.numel(), device=target.device).reshape(target.shape) % 5 + 1).to( + torch.int32 + ) * (target >= 0) + num_labels = max(int((target >= 0).sum().item()), 1) + + ref = reference_grpo_metrics( + logits, + target, + advantages, + old_log_probabilities, + label_counts, + epsilon_low=0.2, + epsilon_high=0.2, + logits_scale_factor=logits_scale_factor, + compute_entropy=compute_entropy, + ) + spec = MonolithicLossSpec( + kind="grpo", + name="grpo", + weight=1.0, + logits_scale_factor=logits_scale_factor, + grad_output=None, + divisor=num_labels, + target=target, + advantages=advantages, + old_log_probabilities=old_log_probabilities, + num_labels_in_seq=label_counts, + compute_metrics=True, + compute_entropy=compute_entropy, + ) + outputs, _ = monolithic_head_loss_forward_backward( + split_op(logits, group, -1).contiguous(), [spec], group=group, grad_logits=None + ) + _check_grpo_metrics(ref, outputs[0].metrics, threshold=5e-5 if dtype == DataType.float32 else 1e-4) + + _MONOLITHIC_DISTRIBUTION_CASES = ( # (target_format, entropy_loss_type, temperature) (TargetFormat.logits, EntropyLossType.cross_entropy, 1.0), @@ -870,6 +914,27 @@ def test_monolithic_grpo_loss( ) +@pytest.mark.slow +@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) +@pytest.mark.parametrize( + ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), + _LOSS_PARAMETERS, +) +@pytest.mark.parametrize("compute_entropy", (False, True)) +def test_monolithic_grpo_metrics( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + block_size, + accumulate, + compute_entropy, +): + _test_monolithic_grpo_metrics(batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -1103,6 +1168,22 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa accumulate, test_context.group, ) + # Monolithic GRPO metrics + for compute_entropy in (False, True): + with test_context.subtest( + base_path, f"monolithic_grpo_metrics-{compute_entropy}-{suffix}", 2 + ) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_monolithic_grpo_metrics( + batch_shape, + num_columns, + logits_scale_factor, + loss_masking, + dtype, + compute_entropy, + test_context.group, + ) # Monolithic from-distribution (distillation) losses for target_format, entropy_loss_type, temperature in _MONOLITHIC_DISTRIBUTION_CASES: with test_context.subtest( @@ -1166,6 +1247,8 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): "grpo_metrics-False", "grpo_metrics-True", "monolithic_grpo", + "monolithic_grpo_metrics-False", + "monolithic_grpo_metrics-True", *(f"monolithic-{"_".join(kinds)}" for kinds in _MONOLITHIC_KINDS), *( f"monolithic_dist-{entropy_loss_type}-{target_format}-{temperature}" From ea6f569a6e206abd7a8e135fed5600d988ba1722 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 29 Jun 2026 14:16:14 -0400 Subject: [PATCH 07/28] Monolithic head-loss kernel: GSPO three-phase eager-seam path (#507) Split fused_gspo_loss_forward_backward into a compiled forward core (one softmax + predicted-logits -> new_log_probs), the eager segment seam (index_add_ + SDP/SP all-reduce + clipping, where the symbolic num_segments stays out of every compiled boundary), and a compiled backward core (the O(vocab) scatter-add grad, previously eager). This is the proven three-phase layout already used by the Triton GSPO kernel; the backward is now fused and num_segments never triggers a recompile. GSPO runs in the monolithic (fused) loss path via the existing spec-less fallback, threading the shared gradient buffer; a standalone GSPO computes one softmax, as before. Math is byte-equivalent to the prior eager implementation. Validated vs reference_gspo_loss (functional across num_segments/dtypes, head fused + per-loss configs, and a new tensor-parallel distributed subtest). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../language_model/loss/policy_gradient.py | 219 ++++++++++++------ tests/layers/test_lm_head.py | 12 + tests/layers/test_lm_losses.py | 16 ++ 3 files changed, 180 insertions(+), 67 deletions(-) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 659317721..b04b5db46 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -6,7 +6,12 @@ from fast_llm.engine.base_model.config import LossDef, ReductionType from fast_llm.engine.distributed.config import DistributedConfig from fast_llm.functional.config import TritonConfig -from fast_llm.functional.entropy_loss import fused_predicted_logits_from_labels, fused_softmax_base +from fast_llm.functional.entropy_loss import ( + _predicted_logits_from_labels, + _softmax_base, + fused_predicted_logits_from_labels, + fused_softmax_base, +) from fast_llm.functional.utils import reduce_losses from fast_llm.layers.block.config import BlockKwargs from fast_llm.layers.language_model.config import LanguageModelKwargs @@ -430,58 +435,43 @@ def fused_grpo_loss_forward_backward( return loss, grad_logits, new_logprobs_mean -# Not @torch.compile: dynamo lifts the Python-int `divisor` and `num_segments` args to symbolic -# ints with no concrete hint, then trips on `grad_output / divisor` during trace evaluation -# (`ZeroDivisionError` with hint=0). The Triton kernel is the actual GPU perf path; the eager -# PyTorch fallback doesn't need to be compiled. -def fused_gspo_loss_forward_backward( +@torch.compile +def _gspo_forward_core( logits: torch.Tensor, # (*batch, vocab) target: torch.Tensor, # (*batch,) + loss_mask: torch.Tensor, # (*batch,), == target >= 0 + logits_scale_factor: float, + group: torch.distributed.ProcessGroup | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """GSPO compiled forward: one softmax + predicted-logit lookup, producing per-token new log-probs. + The softmax tensors are handed to the eager segment seam and the compiled backward.""" + logits_norm, exp_logits, sum_exp_logits, _ = _softmax_base(logits, logits_scale_factor, group) + 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 new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask + + +def _gspo_segment_seam( + new_log_probs: torch.Tensor, # (*batch,) + loss_mask: torch.Tensor, # (*batch,) bool advantages: torch.Tensor, # (*batch,) old_log_probabilities: torch.Tensor, # (*batch,) - document_index_zero_based: torch.Tensor, # (*batch,) int — segment ID per token, 0-based - num_segments: int, # buffer size, ≥ document_index.max() + 1 + document_index_zero_based: torch.Tensor, # (*batch,) int + num_segments: int, + num_labels_in_seq: torch.Tensor, # (*batch,) divisor: float, - num_labels_in_seq: torch.Tensor, # (*batch,) — per-document labeled-token count broadcast per token - grad_logits: torch.Tensor | None = None, - grad_output: float | None = None, - group: torch.distributed.ProcessGroup | None = None, # TP vocab group - sdp_group: torch.distributed.ProcessGroup | None = None, # SDP group for cross-rank segment aggregation - sp_group: torch.distributed.ProcessGroup | None = None, # TP group when SP is sharding the sequence - epsilon_low: float = 0.2, - epsilon_high: float = 0.2, - logits_scale_factor: float = 1.0, -) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - """GSPO loss: sequence-level geometric-mean IS-ratio clipping. - - Per-segment ratio R_s = exp(mean_t(log p_new_t / p_old_t)), clipped per segment. - Per-segment loss = -min(R_s * A_s, clip(R_s) * A_s), multiplied by the segment - token count, summed over segments, and divided by `divisor`. - Computed as an equivalent per-token sum so the gradient chain mirrors GRPO. - - `num_labels_in_seq[t]` is the labeled-token count for the document containing token `t` - (broadcast per token by the data preprocessor); it is the geometric-mean denominator. - Using it directly — rather than aggregating token counts inside the kernel — is what - makes the loss correct when a document spans SDP/SP ranks (numerator - `log_ratio_sum` is all-reduced; denominator is constant and available locally). - - Constraint: each document must be visible to a single kernel call (modulo SDP/SP, where - the kernel all-reduces). Splitting a document across separate kernel calls (e.g. - `schedule.micro_batch_splits > 1`) produces partial per-fragment geometric means that - cannot be combined linearly into the whole-document `exp(mean)`. - - With `sdp_group` and/or `sp_group`, segment sums are all-reduced over those groups so each rank - computes the same global R_s/A_s. Token-level contributions remain per-rank, so summing the - kernel loss across SDP/SP via SUM reduction matches the canonical single-rank result. - """ - grad_output_scaled = None if grad_output is None else grad_output / divisor * logits_scale_factor - loss_mask = target >= 0 - - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) - predicted_logits, target_masked, target_mask = fused_predicted_logits_from_labels( - logits_norm, target, loss_mask, group - ) - new_log_probs = predicted_logits - sum_exp_logits.log() + grad_output: float | None, + sdp_group: torch.distributed.ProcessGroup | None, + sp_group: torch.distributed.ProcessGroup | None, + epsilon_low: float, + epsilon_high: float, + logits_scale_factor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Eager segment seam between the compiled forward and backward. The `index_add_` segment + aggregation and the symbolic `num_segments` live here so they never enter a compiled boundary + (no per-`num_segments` recompiles). Returns the loss, the `new_logprobs` metric, and the per-token + backward coefficient `effective_grad = grad_output_scaled · clip_factor · loss_weight · R_s` + (None when no gradient is requested).""" log_ratio = new_log_probs - old_log_probabilities flat_document_index = document_index_zero_based.reshape(-1).long() @@ -525,25 +515,120 @@ def fused_gspo_loss_forward_backward( new_logprobs_mean = (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() - if grad_output_scaled is not None: - probability_ratio_grad = ( - grad_output_scaled - * ( - torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high) - + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low) - ) - * loss_weight - ) - predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze_(-1) - grad = (probability_ratio_grad * probability_ratio).unsqueeze(-1) * predicted_probabilities.scatter_add( - -1, - target_masked.unsqueeze(-1), - -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32), + if grad_output is None: + return loss, new_logprobs_mean, None + + grad_output_scaled = grad_output / divisor * logits_scale_factor + probability_ratio_grad = ( + grad_output_scaled + * ( + torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high) + + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low) ) - grad = grad.to(logits.dtype) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) + * loss_weight + ) + effective_grad = probability_ratio_grad * probability_ratio + return loss, new_logprobs_mean, effective_grad + + +@torch.compile +def _gspo_backward_core( + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + target_masked: torch.Tensor, # (*batch,) + target_mask: torch.Tensor | None, # (*batch,) or None (no TP) + loss_mask: torch.Tensor, # (*batch,) bool + effective_grad: torch.Tensor, # (*batch,) — per-token backward coefficient from the seam + logits_dtype: torch.dtype, + grad_logits: torch.Tensor | None, +) -> torch.Tensor: + """GSPO compiled backward: the per-token coefficient times the softmax chain rule, fused into + one kernel. `sum_exp_logits.unsqueeze` is out-of-place (the eager kernel mutates it in place).""" + predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) + grad = effective_grad.unsqueeze(-1) * predicted_probabilities.scatter_add( + -1, + target_masked.unsqueeze(-1), + -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32), + ) + grad = grad.to(logits_dtype) + if grad_logits is None: + grad_logits = grad + else: + grad_logits.add_(grad) + return grad_logits + + +# Orchestrator only: the eager `index_add_` segment seam (with the Python-int `num_segments`) sits +# between the compiled forward and backward cores, so it stays out of every compiled boundary. +def fused_gspo_loss_forward_backward( + logits: torch.Tensor, # (*batch, vocab) + target: torch.Tensor, # (*batch,) + advantages: torch.Tensor, # (*batch,) + old_log_probabilities: torch.Tensor, # (*batch,) + document_index_zero_based: torch.Tensor, # (*batch,) int — segment ID per token, 0-based + num_segments: int, # buffer size, ≥ document_index.max() + 1 + divisor: float, + num_labels_in_seq: torch.Tensor, # (*batch,) — per-document labeled-token count broadcast per token + grad_logits: torch.Tensor | None = None, + grad_output: float | None = None, + group: torch.distributed.ProcessGroup | None = None, # TP vocab group + sdp_group: torch.distributed.ProcessGroup | None = None, # SDP group for cross-rank segment aggregation + sp_group: torch.distributed.ProcessGroup | None = None, # TP group when SP is sharding the sequence + epsilon_low: float = 0.2, + epsilon_high: float = 0.2, + logits_scale_factor: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """GSPO loss: sequence-level geometric-mean IS-ratio clipping. + Per-segment ratio R_s = exp(mean_t(log p_new_t / p_old_t)), clipped per segment. + Per-segment loss = -min(R_s * A_s, clip(R_s) * A_s), multiplied by the segment + token count, summed over segments, and divided by `divisor`. + Computed as an equivalent per-token sum so the gradient chain mirrors GRPO. + + `num_labels_in_seq[t]` is the labeled-token count for the document containing token `t` + (broadcast per token by the data preprocessor); it is the geometric-mean denominator. + Using it directly — rather than aggregating token counts inside the kernel — is what + makes the loss correct when a document spans SDP/SP ranks (numerator + `log_ratio_sum` is all-reduced; denominator is constant and available locally). + + Constraint: each document must be visible to a single kernel call (modulo SDP/SP, where + the kernel all-reduces). Splitting a document across separate kernel calls (e.g. + `schedule.micro_batch_splits > 1`) produces partial per-fragment geometric means that + cannot be combined linearly into the whole-document `exp(mean)`. + + With `sdp_group` and/or `sp_group`, segment sums are all-reduced over those groups so each rank + computes the same global R_s/A_s. Token-level contributions remain per-rank, so summing the + kernel loss across SDP/SP via SUM reduction matches the canonical single-rank result. + """ + loss_mask = target >= 0 + new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask = _gspo_forward_core( + logits, target, loss_mask, logits_scale_factor, group + ) + loss, new_logprobs_mean, effective_grad = _gspo_segment_seam( + new_log_probs, + loss_mask, + advantages, + old_log_probabilities, + document_index_zero_based, + num_segments, + num_labels_in_seq, + divisor, + grad_output, + sdp_group, + sp_group, + epsilon_low, + epsilon_high, + 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, grad_logits, new_logprobs_mean diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index f7742e766..b6f855d8e 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -415,6 +415,18 @@ def _add_configs(base_name: str, **kwargs): distillation_temperature=2.0, ) _add_configs("fused_grpo_loss", loss_implementation="fused", grpo_loss=True) +# GSPO runs through its own three-phase path (compiled forward → eager segment seam → compiled backward), +# accumulating into the monolithic kernel's shared gradient. No splits (per-segment aggregation can't +# recombine across cross_entropy_splits chunks). +for loss_masking in (False, True): + _lm_head_test_configs.append( + LMHeadTestConfig( + f"fused_gspo_loss{'_masked' if loss_masking else ''}", + gspo_loss=True, + 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"): diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 8c3cb3b73..f2a5b4364 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -1124,6 +1124,21 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa accumulate, test_context.group, ) + # GSPO (tensor-parallel vocab path; segment seam runs eagerly per rank) + with test_context.subtest(base_path, f"gspo-{suffix}", 2) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_gspo_loss( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + 4, # num_segments + accumulate, + test_context.group, + ) # GRPO metrics for compute_entropy in (False, True): with test_context.subtest(base_path, f"grpo_metrics-{compute_entropy}-{suffix}", 2) as subtest: @@ -1244,6 +1259,7 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): ), "z_loss", "grpo", + "gspo", "grpo_metrics-False", "grpo_metrics-True", "monolithic_grpo", From e5147c5e9fbb170f9dd4ae4d9e75a60375c30c44 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 29 Jun 2026 16:33:28 -0400 Subject: [PATCH 08/28] Monolithic head-loss kernel: multi-loss combo coverage (#507) Validate the monolithic kernel composes arbitrary loss kinds over one shared softmax. Functional level adds CE+z-loss+distillation and GRPO+z-loss combos, comparing total loss + combined gradient against the sum of the per-loss baselines (extends `_monolithic_baseline` to the distribution and GRPO kinds and builds the shared input union). Head level adds the fused three-way label+distillation+z-loss combo and GRPO+metrics+z-loss, checked against the independent reference end to end. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/layers/test_lm_head.py | 21 ++++++++ tests/layers/test_lm_losses.py | 99 +++++++++++++++++++++++++++++++++- 2 files changed, 118 insertions(+), 2 deletions(-) diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index b6f855d8e..50c7f68d1 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -415,6 +415,15 @@ def _add_configs(base_name: str, **kwargs): distillation_temperature=2.0, ) _add_configs("fused_grpo_loss", loss_implementation="fused", grpo_loss=True) +# Multi-loss combos sharing one softmax pass: a three-way distillation combo and an RL + regularizer combo. +_add_configs( + "fused_label_distillation_z_loss", + loss_implementation="fused", + label_loss=True, + distillation_loss=True, + z_loss=0.5, +) +_add_configs("fused_grpo_and_z_loss", loss_implementation="fused", grpo_loss=True, z_loss=0.5) # GSPO runs through its own three-phase path (compiled forward → eager segment seam → compiled backward), # accumulating into the monolithic kernel's shared gradient. No splits (per-segment aggregation can't # recombine across cross_entropy_splits chunks). @@ -443,6 +452,18 @@ def _add_configs(base_name: str, **kwargs): loss_implementation=_loss_implementation, ) ) +# The metric family co-resides with z-loss in the shared softmax pass. Single-split (metrics can't be split). +for _loss_masking in (False, True): + _lm_head_test_configs.append( + LMHeadTestConfig( + f"fused_grpo_and_z_loss_metrics{'_masked' if _loss_masking else ''}", + grpo_loss=True, + z_loss=0.5, + grpo_metrics="basic", + loss_masking=_loss_masking, + loss_implementation="fused", + ) + ) @pytest.mark.slow diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index f2a5b4364..0b0651d0a 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -588,7 +588,19 @@ def _test_z_loss( def _monolithic_baseline( - kind: str, local_logits, target, z_mask, grad_output, group, logits_scale_factor, divisor + kind: str, + local_logits, + target, + z_mask, + grad_output, + group, + logits_scale_factor, + divisor, + *, + teacher_target=None, + advantages=None, + old_log_probabilities=None, + num_labels_in_seq=None, ) -> tuple[torch.Tensor, torch.Tensor | None, MonolithicLossSpec]: # Each enabled loss's baseline is its own `fused_*` kernel; the monolithic grad must equal their sum. if kind == "cross_entropy": @@ -630,6 +642,55 @@ def _monolithic_baseline( divisor=divisor, loss_mask=z_mask, ) + elif kind == "entropy_from_distribution": + loss, grad = fused_entropy_loss_forward_backward( + logits=local_logits, + target=teacher_target, + loss_mask=z_mask, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + target_format=TargetFormat.logits, + entropy_loss_type=EntropyLossType.cross_entropy, + divisor=divisor, + ) + spec = MonolithicLossSpec( + kind=kind, + name=kind, + weight=1.0, + logits_scale_factor=logits_scale_factor, + grad_output=grad_output, + divisor=divisor, + target=teacher_target, + loss_mask=z_mask, + target_format=TargetFormat.logits, + entropy_loss_type=EntropyLossType.cross_entropy, + ) + elif kind == "grpo": + loss, grad, _ = fused_grpo_loss_forward_backward( + local_logits, + target, + advantages, + old_log_probabilities, + grad_logits=None, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + num_labels_in_seq=num_labels_in_seq, + divisor=divisor, + ) + spec = MonolithicLossSpec( + kind=kind, + name=kind, + weight=1.0, + logits_scale_factor=logits_scale_factor, + grad_output=grad_output, + divisor=divisor, + target=target, + advantages=advantages, + old_log_probabilities=old_log_probabilities, + num_labels_in_seq=num_labels_in_seq, + ) else: raise NotImplementedError(kind) return loss, grad, spec @@ -646,10 +707,42 @@ def _test_monolithic_loss( # z-loss takes an explicit mask (the head's `loss_mask` kwarg); align it with the labeled tokens. z_mask = (target >= 0) if loss_masking else None + # The teacher distribution is over the full vocab, so it splits along the vocab axis like the logits. + teacher_target = ( + split_op(torch.randn_like(logits), group, -1).contiguous() if "entropy_from_distribution" in kinds else None + ) + if "grpo" in kinds: + advantages = torch.randn_like(target, dtype=torch.float32) + # Correlate the old log-probs with the policy so the clipping branches are exercised. + old_log_probabilities = ( + torch.nn.functional.log_softmax(logits, dim=-1) + .gather(-1, (target * (target >= 0)).unsqueeze(-1)) + .squeeze(-1) + + torch.randn_like(target, dtype=torch.float32) / 2 + ) + num_labels_in_seq = torch.where( + target >= 0, + torch.full(batch_shape, divisor, dtype=torch.int32, device=target.device), + torch.zeros(batch_shape, dtype=torch.int32, device=target.device), + ) + else: + advantages = old_log_probabilities = num_labels_in_seq = None + ref_losses, specs, ref_grad = [], [], None for kind in kinds: ref_loss, grad, spec = _monolithic_baseline( - kind, local_logits, target, z_mask, grad_output, group, logits_scale_factor, divisor + kind, + local_logits, + target, + z_mask, + grad_output, + group, + logits_scale_factor, + divisor, + teacher_target=teacher_target, + advantages=advantages, + old_log_probabilities=old_log_probabilities, + num_labels_in_seq=num_labels_in_seq, ) ref_losses.append(ref_loss) specs.append(spec) @@ -680,6 +773,8 @@ def _test_monolithic_loss( ("cross_entropy",), ("z_loss",), ("cross_entropy", "z_loss"), + ("cross_entropy", "z_loss", "entropy_from_distribution"), + ("grpo", "z_loss"), ) From ecea37cb1dddd922e59b58dc97bafda5877343bd Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 30 Jun 2026 12:57:00 -0400 Subject: [PATCH 09/28] Monolithic head-loss kernel: GSPO as a kernel kind (#507) GSPO now runs inside the monolithic kernel sharing the single softmax, rather than through the spec-less fallback. The kernel splits into three phases when a GSPO spec is present: a compiled forward (shared softmax + every enabled non-GSPO loss + GSPO's per-token new_log_probs), the eager segment seam (index_add_ aggregation + SDP/SP all-reduce + clipping, the symbolic num_segments staying out of all compiled boundaries), and a compiled backward whose GSPO grad term composes into the shared fp32 accumulator and casts once. The per-loss math is factored into a plain _apply_combinable_losses helper shared by the single-pass and GSPO-split forwards; the GSPO seam/backward cores move to monolithic.py (the import leaf) so policy_gradient can reuse them without a cycle. GSPO gets get_monolithic_spec / register_monolithic_outputs; the standalone fused_gspo_loss_forward_backward is unchanged and remains the oracle. Adds monolithic-GSPO functional + tensor-parallel tests and a fused GSPO+z-loss head combo (the shared-softmax payoff). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../layers/language_model/loss/monolithic.py | 489 ++++++++++++++++-- .../language_model/loss/policy_gradient.py | 145 ++---- tests/layers/test_lm_head.py | 15 +- tests/layers/test_lm_losses.py | 101 ++++ 4 files changed, 585 insertions(+), 165 deletions(-) diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index a3aab5bc1..91288aa52 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -28,14 +28,14 @@ class MonolithicLossSpec(typing.NamedTuple): grad_output: float | None divisor: float target: torch.Tensor | None = ( - None # Cross-entropy / GRPO labels, or a teacher distribution (logits/probabilities). + None # Cross-entropy / GRPO / GSPO labels, or a teacher distribution (logits/probabilities). ) loss_mask: torch.Tensor | None = None # z-loss / distribution-loss mask (cross-entropy / GRPO derive their own). # Distribution losses (cross-entropy / forward-KL / reverse-KL from a teacher distribution) only. target_format: TargetFormat = TargetFormat.logits entropy_loss_type: EntropyLossType = EntropyLossType.cross_entropy temperature: float = 1.0 - # Policy-gradient (GRPO) only. + # Policy-gradient (GRPO / GSPO) only. advantages: torch.Tensor | None = None old_log_probabilities: torch.Tensor | None = None epsilon_low: float = 0.2 @@ -43,22 +43,28 @@ class MonolithicLossSpec(typing.NamedTuple): num_labels_in_seq: torch.Tensor | None = None # Per-sequence label count; enables the `new_logprobs_mean` metric. compute_metrics: bool = False # Emit the GRPO metric family from the shared softmax. compute_entropy: bool = False # Also emit per-token entropy (the only metric needing the vocab axis). + # Sequence-policy-gradient (GSPO) only. + document_index: torch.Tensor | None = None # Per-token segment ID (0-based). + num_segments: int = 0 # Per-segment buffer size, ≥ document_index.max() + 1. + sdp_group: ProcessGroup | None = None # Sequence-data group for cross-rank segment aggregation. + sp_group: ProcessGroup | None = None # TP group when sequence-parallel shards the sequence. class MonolithicLossOutput(typing.NamedTuple): """Per-loss outputs returned by the monolithic kernel, registered via `register_monolithic_outputs`.""" loss: torch.Tensor | None = None - new_logprobs_mean: torch.Tensor | None = None # GRPO only. + new_logprobs_mean: torch.Tensor | None = None # GRPO / GSPO only. metrics: GRPOMetrics | None = None # GRPO only. -@torch.compile -def _monolithic_core( - logits: torch.Tensor, # (*batch, vocab) +def _apply_combinable_losses( + logits_norm: torch.Tensor, # (*batch, vocab) + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + logits_max: torch.Tensor, # (*batch,) group: ProcessGroup | None, logits_scale_factor: float, - grad_logits: torch.Tensor | None, ce_target: torch.Tensor | None, ce_grad_output: float | None, ce_divisor: float, @@ -93,16 +99,12 @@ def _monolithic_core( torch.Tensor | None, ]: """ - The monolithic head-loss kernel: one shared softmax over the logits, then every enabled loss's scalar - and gradient contribution. Losses with a required tensor input are toggled off by passing it as `None` - (e.g. cross-entropy's `ce_target`); losses without one (z-loss) toggle on a static `*_enabled` bool. In - both cases `torch.compile` specializes on the flag and dead-code-eliminates the disabled branch. Because - this is a single `@torch.compile` boundary that calls the plain (un-compiled) softmax cores, compile - fuses the work across all enabled losses. - - Gradients accumulate in fp32 and are cast to `logits.dtype` once at the end. + The per-loss scalar + fp32 gradient-accumulator math shared by the single-pass (`_monolithic_core`) + and GSPO-split (`_monolithic_gspo_forward_core`) paths. Takes the already-computed shared softmax + tensors and returns each enabled loss's scalar plus the summed gradient, accumulated in fp32 and + *not* cast (the caller casts once). Plain (un-decorated) so it inlines and fuses into its compiled + callers; a disabled loss is dead-code-eliminated on its `None`/static flag. """ - logits_norm, exp_logits, sum_exp_logits, logits_max = _softmax_base(logits, logits_scale_factor, group) grad = None cross_entropy_loss = None @@ -238,6 +240,91 @@ def _monolithic_core( ) grad = grpo_grad if grad is None else grad + grpo_grad + return cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad + + +@torch.compile +def _monolithic_core( + logits: torch.Tensor, # (*batch, vocab) + group: ProcessGroup | None, + logits_scale_factor: float, + grad_logits: torch.Tensor | None, + ce_target: torch.Tensor | None, + ce_grad_output: float | None, + ce_divisor: float, + z_loss_enabled: bool, + z_loss_mask: torch.Tensor | None, + z_loss_grad_output: float | None, + z_loss_divisor: float, + distribution_target: torch.Tensor | None, + distribution_grad_output: float | None, + distribution_divisor: float, + distribution_loss_mask: torch.Tensor | None, + distribution_target_format: TargetFormat, + distribution_entropy_loss_type: EntropyLossType, + distribution_temperature: float, + grpo_target: torch.Tensor | None, + grpo_advantages: torch.Tensor | None, + grpo_old_log_probabilities: torch.Tensor | None, + grpo_grad_output: float | None, + grpo_divisor: float, + grpo_epsilon_low: float, + grpo_epsilon_high: float, + grpo_num_labels_in_seq: torch.Tensor | None, + grpo_compute_metrics: bool, + grpo_compute_entropy: bool, +) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + "GRPOMetrics | None", + torch.Tensor | None, +]: + """ + The single-pass monolithic head-loss kernel: one shared softmax over the logits, then every enabled + loss's scalar and gradient contribution, fused in one `@torch.compile` boundary (it inlines the plain + per-loss core). Disabled losses gate on their `None`/static flag and are dead-code-eliminated. + Gradients accumulate in fp32 and are cast to `logits.dtype` once at the end. GSPO is *not* handled + here — its eager segment seam requires the three-phase split (`_monolithic_gspo_forward_core`). + """ + logits_norm, exp_logits, sum_exp_logits, logits_max = _softmax_base(logits, logits_scale_factor, group) + cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad = ( + _apply_combinable_losses( + logits_norm, + exp_logits, + sum_exp_logits, + logits_max, + group, + logits_scale_factor, + ce_target, + ce_grad_output, + ce_divisor, + z_loss_enabled, + z_loss_mask, + z_loss_grad_output, + z_loss_divisor, + distribution_target, + distribution_grad_output, + distribution_divisor, + distribution_loss_mask, + distribution_target_format, + distribution_entropy_loss_type, + distribution_temperature, + grpo_target, + grpo_advantages, + grpo_old_log_probabilities, + grpo_grad_output, + grpo_divisor, + grpo_epsilon_low, + grpo_epsilon_high, + grpo_num_labels_in_seq, + grpo_compute_metrics, + grpo_compute_entropy, + ) + ) + if grad is not None: grad = grad.to(logits.dtype) if grad_logits is None: @@ -248,6 +335,204 @@ def _monolithic_core( return cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad_logits +@torch.compile +def _monolithic_gspo_forward_core( + logits: torch.Tensor, # (*batch, vocab) + group: ProcessGroup | None, + logits_scale_factor: float, + ce_target: torch.Tensor | None, + ce_grad_output: float | None, + ce_divisor: float, + z_loss_enabled: bool, + z_loss_mask: torch.Tensor | None, + z_loss_grad_output: float | None, + z_loss_divisor: float, + distribution_target: torch.Tensor | None, + distribution_grad_output: float | None, + distribution_divisor: float, + distribution_loss_mask: torch.Tensor | None, + distribution_target_format: TargetFormat, + distribution_entropy_loss_type: EntropyLossType, + distribution_temperature: float, + gspo_target: torch.Tensor, +) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, +]: + """ + GSPO-present forward: one shared softmax, every enabled non-GSPO loss's scalar + fp32 gradient + (uncast), and GSPO's per-token `new_log_probs`. The softmax tensors and the uncast gradient cross the + eager segment seam to `_gspo_backward_core` (the `index_add_`/`num_segments` seam can't be compiled). + """ + logits_norm, exp_logits, sum_exp_logits, logits_max = _softmax_base(logits, logits_scale_factor, group) + cross_entropy_loss, z_loss, distribution_loss, _, _, _, grad = _apply_combinable_losses( + logits_norm, + exp_logits, + sum_exp_logits, + logits_max, + group, + logits_scale_factor, + ce_target, + ce_grad_output, + ce_divisor, + z_loss_enabled, + z_loss_mask, + z_loss_grad_output, + z_loss_divisor, + distribution_target, + distribution_grad_output, + distribution_divisor, + distribution_loss_mask, + distribution_target_format, + distribution_entropy_loss_type, + distribution_temperature, + None, + None, + None, + None, + 1.0, + 0.2, + 0.2, + None, + False, + False, + ) + gspo_loss_mask = gspo_target >= 0 + predicted_logits, target_masked, target_mask = _predicted_logits_from_labels( + logits_norm, gspo_target, gspo_loss_mask, group + ) + new_log_probs = predicted_logits - sum_exp_logits.log() + return ( + cross_entropy_loss, + z_loss, + distribution_loss, + grad, + new_log_probs, + exp_logits, + sum_exp_logits, + target_masked, + target_mask, + ) + + +def _gspo_segment_seam( + new_log_probs: torch.Tensor, # (*batch,) + loss_mask: torch.Tensor, # (*batch,) bool + advantages: torch.Tensor, # (*batch,) + old_log_probabilities: torch.Tensor, # (*batch,) + document_index_zero_based: torch.Tensor, # (*batch,) int + num_segments: int, + num_labels_in_seq: torch.Tensor, # (*batch,) + divisor: float, + grad_output: float | None, + sdp_group: ProcessGroup | None, + sp_group: ProcessGroup | None, + epsilon_low: float, + epsilon_high: float, + logits_scale_factor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Eager segment seam between the compiled forward and backward. The `index_add_` segment + aggregation and the symbolic `num_segments` live here so they never enter a compiled boundary + (no per-`num_segments` recompiles). Returns the loss, the `new_logprobs` metric, and the per-token + backward coefficient `effective_grad = grad_output_scaled · clip_factor · loss_weight · R_s` + (None when no gradient is requested).""" + log_ratio = new_log_probs - old_log_probabilities + + flat_document_index = document_index_zero_based.reshape(-1).long() + flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype) + # Per-token weight: mask / per-document label count, from the preprocessor. + # Each labeled token contributes `1 / N_d` so all of doc d's tokens sum to 1 (across + # SDP/SP ranks too), regardless of how the doc is sharded. + mean_token_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1) + # Pre-divide the per-token contributions by the per-doc label count, then sum per segment. + # All tokens in a segment share the same N_d, so this is mathematically equivalent to + # `log_ratio_sum / N_d` but avoids any per-segment denominator extraction. + mean_log_ratio_per_segment = log_ratio.new_zeros(num_segments).index_add_( + 0, flat_document_index, log_ratio.reshape(-1) * mean_token_weight + ) + # Accumulate in `log_ratio.dtype` (fp32). Casting the product back to `advantages.dtype` + # before summing would round each token's contribution to a possibly-low input dtype. + mean_advantage_per_segment = log_ratio.new_zeros(num_segments).index_add_( + 0, flat_document_index, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight + ) + for reduce_group in (sdp_group, sp_group): + if reduce_group is not None: + torch.distributed.all_reduce( + mean_log_ratio_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group + ) + torch.distributed.all_reduce( + mean_advantage_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group + ) + + segment_ratio = mean_log_ratio_per_segment.exp() # (num_segments,) — geometric-mean IS ratio + segment_advantage = mean_advantage_per_segment.detach() # (num_segments,) — no grad through A + + probability_ratio = segment_ratio[flat_document_index].reshape(log_ratio.shape) + advantage_per_token = segment_advantage[flat_document_index].reshape(log_ratio.shape) + loss_weight = loss_mask.to(log_ratio.dtype) + + losses = -torch.min( + probability_ratio * advantage_per_token, + torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantage_per_token, + ) + loss = (losses * loss_weight).sum() / divisor + + new_logprobs_mean = (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() + + if grad_output is None: + return loss, new_logprobs_mean, None + + grad_output_scaled = grad_output / divisor * logits_scale_factor + probability_ratio_grad = ( + grad_output_scaled + * ( + torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high) + + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low) + ) + * loss_weight + ) + effective_grad = probability_ratio_grad * probability_ratio + return loss, new_logprobs_mean, effective_grad + + +@torch.compile +def _gspo_backward_core( + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + target_masked: torch.Tensor, # (*batch,) + target_mask: torch.Tensor | None, # (*batch,) or None (no TP) + loss_mask: torch.Tensor, # (*batch,) bool + effective_grad: torch.Tensor, # (*batch,) — per-token backward coefficient from the seam + logits_dtype: torch.dtype, + grad_logits: torch.Tensor | None, + grad_accumulator: torch.Tensor | None = None, # fp32 grad of co-resident losses to combine before the cast +) -> torch.Tensor: + """GSPO compiled backward: the per-token coefficient times the softmax chain rule, fused into one + kernel. Any `grad_accumulator` (other losses' fp32 grad from the shared forward) is added before the + single cast. `sum_exp_logits.unsqueeze` is out-of-place (the standalone eager kernel mutates it).""" + predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) + grad = effective_grad.unsqueeze(-1) * predicted_probabilities.scatter_add( + -1, + target_masked.unsqueeze(-1), + -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32), + ) + if grad_accumulator is not None: + grad = grad + grad_accumulator + grad = grad.to(logits_dtype) + if grad_logits is None: + grad_logits = grad + else: + grad_logits.add_(grad) + return grad_logits + + def monolithic_head_loss_forward_backward( logits: torch.Tensor, specs: list[MonolithicLossSpec], @@ -256,16 +541,17 @@ def monolithic_head_loss_forward_backward( grad_logits: torch.Tensor | None = None, ) -> tuple[list[MonolithicLossOutput], torch.Tensor | None]: """ - Marshal the per-loss specs into the flat arguments of `_monolithic_core`, run it once, and map its - outputs back to one `MonolithicLossOutput` per spec (in input order). All specs must share a single - effective `logits_scale_factor` (the shared softmax cannot serve two scales); this is validated at - config time and asserted here. + Marshal the per-loss specs into the flat arguments of the kernel, run it once over a single shared + softmax, and map its outputs back to one `MonolithicLossOutput` per spec (in input order). When a + GSPO spec is present the kernel splits into the three-phase path (compiled forward → eager segment + seam → compiled backward); otherwise it is the single-pass `_monolithic_core`. All specs must share + one effective `logits_scale_factor` (validated at config time and asserted here). """ logits_scale_factor = specs[0].logits_scale_factor specs_by_kind: dict[str, MonolithicLossSpec] = {} for spec in specs: Assert.eq(spec.logits_scale_factor, logits_scale_factor) - if spec.kind not in ("cross_entropy", "z_loss", "entropy_from_distribution", "grpo"): + if spec.kind not in ("cross_entropy", "z_loss", "entropy_from_distribution", "grpo", "gspo"): raise NotImplementedError(spec.kind) assert spec.kind not in specs_by_kind, spec.kind specs_by_kind[spec.kind] = spec @@ -274,54 +560,149 @@ def monolithic_head_loss_forward_backward( z_loss_spec = specs_by_kind.get("z_loss") distribution_spec = specs_by_kind.get("entropy_from_distribution") grpo_spec = specs_by_kind.get("grpo") + gspo_spec = specs_by_kind.get("gspo") + # GRPO and GSPO are mutually exclusive RL objectives (both consume the same labels/advantages). + assert grpo_spec is None or gspo_spec is None + + # Shared (non-GSPO) loss arguments. + ce_target = None if cross_entropy_spec is None else cross_entropy_spec.target + ce_grad_output = None if cross_entropy_spec is None else cross_entropy_spec.grad_output + ce_divisor = 1.0 if cross_entropy_spec is None else cross_entropy_spec.divisor + z_loss_mask = None if z_loss_spec is None else z_loss_spec.loss_mask + z_loss_grad_output = None if z_loss_spec is None else z_loss_spec.grad_output + z_loss_divisor = 1.0 if z_loss_spec is None else z_loss_spec.divisor + distribution_target = None if distribution_spec is None else distribution_spec.target + distribution_grad_output = None if distribution_spec is None else distribution_spec.grad_output + distribution_divisor = 1.0 if distribution_spec is None else distribution_spec.divisor + distribution_loss_mask = None if distribution_spec is None else distribution_spec.loss_mask + distribution_target_format = TargetFormat.logits if distribution_spec is None else distribution_spec.target_format + distribution_entropy_loss_type = ( + EntropyLossType.cross_entropy if distribution_spec is None else distribution_spec.entropy_loss_type + ) + distribution_temperature = 1.0 if distribution_spec is None else distribution_spec.temperature - cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad_logits = ( - _monolithic_core( + grpo_new_logprobs_mean = None + grpo_metrics = None + gspo_loss = None + gspo_new_logprobs_mean = None + + if gspo_spec is None: + cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad_logits = ( + _monolithic_core( + logits, + group, + logits_scale_factor, + grad_logits, + ce_target=ce_target, + ce_grad_output=ce_grad_output, + ce_divisor=ce_divisor, + z_loss_enabled=z_loss_spec is not None, + z_loss_mask=z_loss_mask, + z_loss_grad_output=z_loss_grad_output, + z_loss_divisor=z_loss_divisor, + distribution_target=distribution_target, + distribution_grad_output=distribution_grad_output, + distribution_divisor=distribution_divisor, + distribution_loss_mask=distribution_loss_mask, + distribution_target_format=distribution_target_format, + distribution_entropy_loss_type=distribution_entropy_loss_type, + distribution_temperature=distribution_temperature, + grpo_target=None if grpo_spec is None else grpo_spec.target, + grpo_advantages=None if grpo_spec is None else grpo_spec.advantages, + grpo_old_log_probabilities=None if grpo_spec is None else grpo_spec.old_log_probabilities, + grpo_grad_output=None if grpo_spec is None else grpo_spec.grad_output, + grpo_divisor=1.0 if grpo_spec is None else grpo_spec.divisor, + grpo_epsilon_low=0.2 if grpo_spec is None else grpo_spec.epsilon_low, + grpo_epsilon_high=0.2 if grpo_spec is None else grpo_spec.epsilon_high, + grpo_num_labels_in_seq=None if grpo_spec is None else grpo_spec.num_labels_in_seq, + grpo_compute_metrics=False if grpo_spec is None else grpo_spec.compute_metrics, + grpo_compute_entropy=False if grpo_spec is None else grpo_spec.compute_entropy, + ) + ) + else: + grpo_loss = None + ( + cross_entropy_loss, + z_loss, + distribution_loss, + grad, + gspo_new_log_probs, + exp_logits, + sum_exp_logits, + target_masked, + target_mask, + ) = _monolithic_gspo_forward_core( logits, group, logits_scale_factor, - grad_logits, - ce_target=None if cross_entropy_spec is None else cross_entropy_spec.target, - ce_grad_output=None if cross_entropy_spec is None else cross_entropy_spec.grad_output, - ce_divisor=1.0 if cross_entropy_spec is None else cross_entropy_spec.divisor, - z_loss_enabled=z_loss_spec is not None, - z_loss_mask=None if z_loss_spec is None else z_loss_spec.loss_mask, - z_loss_grad_output=None if z_loss_spec is None else z_loss_spec.grad_output, - z_loss_divisor=1.0 if z_loss_spec is None else z_loss_spec.divisor, - distribution_target=None if distribution_spec is None else distribution_spec.target, - distribution_grad_output=None if distribution_spec is None else distribution_spec.grad_output, - distribution_divisor=1.0 if distribution_spec is None else distribution_spec.divisor, - distribution_loss_mask=None if distribution_spec is None else distribution_spec.loss_mask, - distribution_target_format=( - TargetFormat.logits if distribution_spec is None else distribution_spec.target_format - ), - distribution_entropy_loss_type=( - EntropyLossType.cross_entropy if distribution_spec is None else distribution_spec.entropy_loss_type - ), - distribution_temperature=1.0 if distribution_spec is None else distribution_spec.temperature, - grpo_target=None if grpo_spec is None else grpo_spec.target, - grpo_advantages=None if grpo_spec is None else grpo_spec.advantages, - grpo_old_log_probabilities=None if grpo_spec is None else grpo_spec.old_log_probabilities, - grpo_grad_output=None if grpo_spec is None else grpo_spec.grad_output, - grpo_divisor=1.0 if grpo_spec is None else grpo_spec.divisor, - grpo_epsilon_low=0.2 if grpo_spec is None else grpo_spec.epsilon_low, - grpo_epsilon_high=0.2 if grpo_spec is None else grpo_spec.epsilon_high, - grpo_num_labels_in_seq=None if grpo_spec is None else grpo_spec.num_labels_in_seq, - grpo_compute_metrics=False if grpo_spec is None else grpo_spec.compute_metrics, - grpo_compute_entropy=False if grpo_spec is None else grpo_spec.compute_entropy, + ce_target, + ce_grad_output, + ce_divisor, + z_loss_spec is not None, + z_loss_mask, + z_loss_grad_output, + z_loss_divisor, + distribution_target, + distribution_grad_output, + distribution_divisor, + distribution_loss_mask, + distribution_target_format, + distribution_entropy_loss_type, + distribution_temperature, + gspo_spec.target, ) - ) + gspo_loss_mask = gspo_spec.target >= 0 + gspo_loss, gspo_new_logprobs_mean, effective_grad = _gspo_segment_seam( + gspo_new_log_probs, + gspo_loss_mask, + gspo_spec.advantages, + gspo_spec.old_log_probabilities, + gspo_spec.document_index, + gspo_spec.num_segments, + gspo_spec.num_labels_in_seq, + gspo_spec.divisor, + gspo_spec.grad_output, + gspo_spec.sdp_group, + gspo_spec.sp_group, + gspo_spec.epsilon_low, + gspo_spec.epsilon_high, + logits_scale_factor, + ) + if effective_grad is not None: + grad_logits = _gspo_backward_core( + exp_logits, + sum_exp_logits, + target_masked, + target_mask, + gspo_loss_mask, + effective_grad, + logits.dtype, + grad_logits, + grad_accumulator=grad, + ) + elif grad is not None: + # No GSPO gradient (e.g. eval), but co-resident losses produced one — finalize it here. + grad = grad.to(logits.dtype) + if grad_logits is None: + grad_logits = grad + else: + grad_logits.add_(grad) loss_by_kind = { "cross_entropy": cross_entropy_loss, "z_loss": z_loss, "entropy_from_distribution": distribution_loss, "grpo": grpo_loss, + "gspo": gspo_loss, } return [ MonolithicLossOutput( loss=loss_by_kind[spec.kind], - new_logprobs_mean=grpo_new_logprobs_mean if spec.kind == "grpo" else None, + new_logprobs_mean=( + grpo_new_logprobs_mean + if spec.kind == "grpo" + else gspo_new_logprobs_mean if spec.kind == "gspo" else None + ), metrics=grpo_metrics if spec.kind == "grpo" else None, ) for spec in specs diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index b04b5db46..080edd4ad 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -24,7 +24,12 @@ ) from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, _grpo_metrics_core from fast_llm.layers.language_model.loss.loss import LanguageModelLoss -from fast_llm.layers.language_model.loss.monolithic import MonolithicLossOutput, MonolithicLossSpec +from fast_llm.layers.language_model.loss.monolithic import ( + MonolithicLossOutput, + MonolithicLossSpec, + _gspo_backward_core, + _gspo_segment_seam, +) from fast_llm.utils import Assert @@ -325,6 +330,37 @@ def _forward_backward( def get_preprocessing_config(self) -> dict[str, typing.Any]: return super().get_preprocessing_config() | {"return_document_index": True} + def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + return MonolithicLossSpec( + kind="gspo", + name=self.name, + weight=self._weight, + logits_scale_factor=self._logits_scale_factor, + grad_output=self._get_grad_output(kwargs), + divisor=kwargs[LanguageModelKwargs.num_documents_in_batch], + target=self._get_labels(kwargs, split_index), + advantages=self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), + old_log_probabilities=self._prepare_target( + kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index + ), + epsilon_low=self._config.epsilon_low, + epsilon_high=self._config.epsilon_high, + num_labels_in_seq=self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index), + document_index=self._prepare_target( + kwargs[BlockKwargs.global_document_index_q], split_index, split_by_distance=False + ).long() + - 1, + num_segments=kwargs[BlockKwargs.num_documents_in_sequence], + sdp_group=self._sequence_data_dim.group if self._sequence_data_active else None, + sp_group=self._parallel_dim.group if self._sequence_parallel else None, + ) + + def register_monolithic_outputs( + self, output: MonolithicLossOutput, kwargs: dict[str, typing.Any], losses: dict | None + ) -> None: + super().register_monolithic_outputs(output, kwargs, losses) + self._register_new_logprobs(output.new_logprobs_mean, kwargs, losses) + @torch.compile def compute_grpo_metrics( @@ -451,113 +487,6 @@ def _gspo_forward_core( return new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask -def _gspo_segment_seam( - new_log_probs: torch.Tensor, # (*batch,) - loss_mask: torch.Tensor, # (*batch,) bool - advantages: torch.Tensor, # (*batch,) - old_log_probabilities: torch.Tensor, # (*batch,) - document_index_zero_based: torch.Tensor, # (*batch,) int - num_segments: int, - num_labels_in_seq: torch.Tensor, # (*batch,) - divisor: float, - grad_output: float | None, - sdp_group: torch.distributed.ProcessGroup | None, - sp_group: torch.distributed.ProcessGroup | None, - epsilon_low: float, - epsilon_high: float, - logits_scale_factor: float, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """Eager segment seam between the compiled forward and backward. The `index_add_` segment - aggregation and the symbolic `num_segments` live here so they never enter a compiled boundary - (no per-`num_segments` recompiles). Returns the loss, the `new_logprobs` metric, and the per-token - backward coefficient `effective_grad = grad_output_scaled · clip_factor · loss_weight · R_s` - (None when no gradient is requested).""" - log_ratio = new_log_probs - old_log_probabilities - - flat_document_index = document_index_zero_based.reshape(-1).long() - flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype) - # Per-token weight: mask / per-document label count, from the preprocessor. - # Each labeled token contributes `1 / N_d` so all of doc d's tokens sum to 1 (across - # SDP/SP ranks too), regardless of how the doc is sharded. - mean_token_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1) - # Pre-divide the per-token contributions by the per-doc label count, then sum per segment. - # All tokens in a segment share the same N_d, so this is mathematically equivalent to - # `log_ratio_sum / N_d` but avoids any per-segment denominator extraction. - mean_log_ratio_per_segment = log_ratio.new_zeros(num_segments).index_add_( - 0, flat_document_index, log_ratio.reshape(-1) * mean_token_weight - ) - # Accumulate in `log_ratio.dtype` (fp32). Casting the product back to `advantages.dtype` - # before summing would round each token's contribution to a possibly-low input dtype. - mean_advantage_per_segment = log_ratio.new_zeros(num_segments).index_add_( - 0, flat_document_index, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight - ) - for reduce_group in (sdp_group, sp_group): - if reduce_group is not None: - torch.distributed.all_reduce( - mean_log_ratio_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group - ) - torch.distributed.all_reduce( - mean_advantage_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group - ) - - segment_ratio = mean_log_ratio_per_segment.exp() # (num_segments,) — geometric-mean IS ratio - segment_advantage = mean_advantage_per_segment.detach() # (num_segments,) — no grad through A - - probability_ratio = segment_ratio[flat_document_index].reshape(log_ratio.shape) - advantage_per_token = segment_advantage[flat_document_index].reshape(log_ratio.shape) - loss_weight = loss_mask.to(log_ratio.dtype) - - losses = -torch.min( - probability_ratio * advantage_per_token, - torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantage_per_token, - ) - loss = (losses * loss_weight).sum() / divisor - - new_logprobs_mean = (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() - - if grad_output is None: - return loss, new_logprobs_mean, None - - grad_output_scaled = grad_output / divisor * logits_scale_factor - probability_ratio_grad = ( - grad_output_scaled - * ( - torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high) - + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low) - ) - * loss_weight - ) - effective_grad = probability_ratio_grad * probability_ratio - return loss, new_logprobs_mean, effective_grad - - -@torch.compile -def _gspo_backward_core( - exp_logits: torch.Tensor, # (*batch, vocab) - sum_exp_logits: torch.Tensor, # (*batch,) - target_masked: torch.Tensor, # (*batch,) - target_mask: torch.Tensor | None, # (*batch,) or None (no TP) - loss_mask: torch.Tensor, # (*batch,) bool - effective_grad: torch.Tensor, # (*batch,) — per-token backward coefficient from the seam - logits_dtype: torch.dtype, - grad_logits: torch.Tensor | None, -) -> torch.Tensor: - """GSPO compiled backward: the per-token coefficient times the softmax chain rule, fused into - one kernel. `sum_exp_logits.unsqueeze` is out-of-place (the eager kernel mutates it in place).""" - predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) - grad = effective_grad.unsqueeze(-1) * predicted_probabilities.scatter_add( - -1, - target_masked.unsqueeze(-1), - -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32), - ) - grad = grad.to(logits_dtype) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) - return grad_logits - - # Orchestrator only: the eager `index_add_` segment seam (with the Python-int `num_segments`) sits # between the compiled forward and backward cores, so it stays out of every compiled boundary. def fused_gspo_loss_forward_backward( diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 50c7f68d1..50ed6e2be 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -424,9 +424,9 @@ 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 runs through its own three-phase path (compiled forward → eager segment seam → compiled backward), -# accumulating into the monolithic kernel's shared gradient. No splits (per-segment aggregation can't -# recombine across cross_entropy_splits chunks). +# GSPO is a monolithic-kernel kind: the shared softmax forward feeds an eager segment seam (compiled +# forward → eager seam → compiled backward), so GSPO can share the softmax with a co-resident loss +# (e.g. z-loss). No splits (per-segment aggregation can't recombine across cross_entropy_splits chunks). for loss_masking in (False, True): _lm_head_test_configs.append( LMHeadTestConfig( @@ -436,6 +436,15 @@ def _add_configs(base_name: str, **kwargs): loss_implementation="fused", ) ) + _lm_head_test_configs.append( + LMHeadTestConfig( + f"fused_gspo_and_z_loss{'_masked' if loss_masking else ''}", + 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"): diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 0b0651d0a..47949e916 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -949,6 +949,77 @@ def _test_monolithic_grpo_metrics( _check_grpo_metrics(ref, outputs[0].metrics, threshold=5e-5 if dtype == DataType.float32 else 1e-4) +def _test_monolithic_gspo_loss( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + num_segments, + accumulate, + group=None, +): + # The monolithic GSPO branch (three-phase: shared softmax forward -> eager seam -> compiled backward) + # must reproduce the standalone `fused_gspo_loss_forward_backward` (loss, grad, new_logprobs). + logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs( + num_columns, loss_masking, batch_shape, dtype + ) + local_logits = split_op(logits, group, -1).contiguous() + seq_len = batch_shape[-1] if len(batch_shape) > 1 else batch_shape[0] + span = max(seq_len // num_segments, 1) + base = torch.arange(seq_len, device=target.device) // span + document_index = base.clamp(max=num_segments - 1).expand(batch_shape).contiguous() + flat_doc = document_index.reshape(-1).long() + labels_per_document = torch.zeros(num_segments, dtype=torch.int32, device=target.device).scatter_add( + 0, flat_doc, (target.reshape(-1) >= 0).to(torch.int32) + ) + num_labels_in_seq = labels_per_document[flat_doc].reshape(target.shape) + previous_grad = torch.randn_like(local_logits) if accumulate and grad_output is not None else None + out_ref, grad_ref, new_logprobs_ref = fused_gspo_loss_forward_backward( + local_logits, + target, + advantages, + old_log_probabilities, + document_index, + num_segments, + divisor=num_segments, + num_labels_in_seq=num_labels_in_seq, + grad_logits=None if previous_grad is None else previous_grad.clone(), + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + ) + spec = MonolithicLossSpec( + kind="gspo", + name="gspo", + weight=1.0, + logits_scale_factor=logits_scale_factor, + grad_output=grad_output, + divisor=num_segments, + target=target, + advantages=advantages, + old_log_probabilities=old_log_probabilities, + num_labels_in_seq=num_labels_in_seq, + document_index=document_index, + num_segments=num_segments, + ) + outputs, grad_mono = monolithic_head_loss_forward_backward( + local_logits, [spec], group=group, grad_logits=None if previous_grad is None else previous_grad.clone() + ) + threshold = 1e-5 if dtype == DataType.float32 else 1e-4 + # The monolithic GSPO forward is a distinct compiled graph from the standalone, so fp32 softmax + # reduction-order differences surface in the geometric-mean ratio — amplified in the loss scalar at + # higher logit scales. The gradient (the real correctness signal) stays tight. + loss_threshold = 5e-5 if dtype == DataType.float32 else 1e-4 + Assert.rms_close_relative(outputs[0].loss, out_ref, loss_threshold, 1e-6) + Assert.rms_close_relative(outputs[0].new_logprobs_mean, new_logprobs_ref, loss_threshold, 1e-6) + if grad_output is not None: + Assert.rms_close_relative(grad_mono, grad_ref, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-6) + else: + assert grad_mono is None + + _MONOLITHIC_DISTRIBUTION_CASES = ( # (target_format, entropy_loss_type, temperature) (TargetFormat.logits, EntropyLossType.cross_entropy, 1.0), @@ -1121,6 +1192,20 @@ def test_gspo_loss( ) +@pytest.mark.slow +@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) +@pytest.mark.parametrize( + ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "num_segments", "accumulate"), + _GSPO_PARAMETERS, +) +def test_monolithic_gspo_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, num_segments, accumulate +): + _test_monolithic_gspo_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, num_segments, accumulate + ) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -1234,6 +1319,21 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa accumulate, test_context.group, ) + # GSPO through the monolithic kernel (shared softmax forward + eager seam + compiled backward) + with test_context.subtest(base_path, f"monolithic_gspo-{suffix}", 2) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_monolithic_gspo_loss( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + 4, # num_segments + accumulate, + test_context.group, + ) # GRPO metrics for compute_entropy in (False, True): with test_context.subtest(base_path, f"grpo_metrics-{compute_entropy}-{suffix}", 2) as subtest: @@ -1358,6 +1458,7 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): "grpo_metrics-False", "grpo_metrics-True", "monolithic_grpo", + "monolithic_gspo", "monolithic_grpo_metrics-False", "monolithic_grpo_metrics-True", *(f"monolithic-{"_".join(kinds)}" for kinds in _MONOLITHIC_KINDS), From 3340da492269989fa3850038decd321ec390c028 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 30 Jun 2026 13:16:34 -0400 Subject: [PATCH 10/28] Monolithic head-loss kernel: make shared loss cores public (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The plain math cores are inlined into the monolithic kernel (and the standalone GSPO path) from other modules, so a leading underscore — which marks a module-private name that should not be imported elsewhere — was wrong. Drop it from the cross-module-shared cores: softmax_base, predicted_logits_from_labels, cross_entropy_from_distribution_core, reverse_kl_from_distribution_core (entropy_loss), grpo_metrics_core (grpo_metrics), and gspo_segment_seam / gspo_backward_core (monolithic). The compiled `fused_*` wrappers and the module-local cores (e.g. _monolithic_core, _gspo_forward_core) are unchanged. Pure rename. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/entropy_loss.py | 24 ++++++------- .../language_model/loss/grpo_metrics.py | 2 +- .../layers/language_model/loss/monolithic.py | 36 +++++++++---------- .../language_model/loss/policy_gradient.py | 20 +++++------ 4 files changed, 41 insertions(+), 41 deletions(-) diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index e00e77ee2..a537bf7df 100644 --- a/fast_llm/functional/entropy_loss.py +++ b/fast_llm/functional/entropy_loss.py @@ -92,7 +92,7 @@ def torch_entropy_loss_forward_backward( return loss.detach_(), grad -def _softmax_base( +def softmax_base( logits: torch.Tensor, # (*batch, vocab) logits_scale_factor: float = 1.0, group: ProcessGroup | None = None, @@ -123,10 +123,10 @@ def _softmax_base( return logits_norm, exp_logits, sum_exp_logits, logits_max -fused_softmax_base = torch.compile(_softmax_base) +fused_softmax_base = torch.compile(softmax_base) -def _reverse_kl_from_distribution_core( +def reverse_kl_from_distribution_core( logits_norm: torch.Tensor, # (*batch, vocab) exp_logits: torch.Tensor, # (*batch, vocab) sum_exp_logits: torch.Tensor, # (*batch,) @@ -148,7 +148,7 @@ def _reverse_kl_from_distribution_core( predicted_probability = exp_logits / sum_exp_logits.unsqueeze(-1) if target_format == TargetFormat.logits: - target_logits_norm, _, sum_exp_target_logits, _ = _softmax_base( + target_logits_norm, _, sum_exp_target_logits, _ = softmax_base( target, logits_scale_factor / temperature, group ) target_log_probability = target_logits_norm - sum_exp_target_logits.log().unsqueeze(-1) @@ -182,8 +182,8 @@ def _fused_reverse_kl_base_from_distribution( group: ProcessGroup | None = None, temperature: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) - logits_norm, exp_logits, sum_exp_logits, _ = _softmax_base(logits, logits_scale_factor, group) - return _reverse_kl_from_distribution_core( + logits_norm, exp_logits, sum_exp_logits, _ = softmax_base(logits, logits_scale_factor, group) + return reverse_kl_from_distribution_core( logits_norm, exp_logits, sum_exp_logits, @@ -196,7 +196,7 @@ def _fused_reverse_kl_base_from_distribution( ) -def _cross_entropy_from_distribution_core( +def cross_entropy_from_distribution_core( logits_norm: torch.Tensor, # (*batch, vocab) exp_logits: torch.Tensor, # (*batch, vocab) sum_exp_logits: torch.Tensor, # (*batch,) @@ -216,7 +216,7 @@ def _cross_entropy_from_distribution_core( wrapper and the monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. """ if target_format == TargetFormat.logits: - target_logits_norm, exp_logits_targets, sum_exp_target_logits, _ = _softmax_base( + target_logits_norm, exp_logits_targets, sum_exp_target_logits, _ = softmax_base( target, logits_scale_factor / temperature, group ) target = exp_logits_targets / sum_exp_target_logits.unsqueeze(-1) @@ -259,8 +259,8 @@ def _fused_cross_entropy_base_from_distribution( temperature: float = 1.0, return_kl_loss: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) - logits_norm, exp_logits, sum_exp_logits, _ = _softmax_base(logits, logits_scale_factor, group) - return _cross_entropy_from_distribution_core( + logits_norm, exp_logits, sum_exp_logits, _ = softmax_base(logits, logits_scale_factor, group) + return cross_entropy_from_distribution_core( logits_norm, exp_logits, sum_exp_logits, @@ -274,7 +274,7 @@ def _fused_cross_entropy_base_from_distribution( ) -def _predicted_logits_from_labels( +def predicted_logits_from_labels( logits: torch.Tensor, # (*batch, vocab) target: torch.Tensor, # (*batch,) loss_mask: torch.Tensor, # (*batch,), == target>=0 @@ -312,7 +312,7 @@ def _predicted_logits_from_labels( return predicted_logits, target_masked, target_mask -fused_predicted_logits_from_labels = torch.compile(_predicted_logits_from_labels) +fused_predicted_logits_from_labels = torch.compile(predicted_logits_from_labels) @torch.compile diff --git a/fast_llm/layers/language_model/loss/grpo_metrics.py b/fast_llm/layers/language_model/loss/grpo_metrics.py index e15045db5..748f78e29 100644 --- a/fast_llm/layers/language_model/loss/grpo_metrics.py +++ b/fast_llm/layers/language_model/loss/grpo_metrics.py @@ -19,7 +19,7 @@ class GRPOMetrics(typing.NamedTuple): entropy: torch.Tensor | None -def _grpo_metrics_core( +def grpo_metrics_core( logits_norm: torch.Tensor, # (*batch, vocab) exp_logits: torch.Tensor, # (*batch, vocab) sum_exp_logits: torch.Tensor, # (*batch,) diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 91288aa52..77593d81f 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -5,12 +5,12 @@ from fast_llm.core.distributed import ProcessGroup from fast_llm.functional.config import EntropyLossType, TargetFormat from fast_llm.functional.entropy_loss import ( - _cross_entropy_from_distribution_core, - _predicted_logits_from_labels, - _reverse_kl_from_distribution_core, - _softmax_base, + cross_entropy_from_distribution_core, + predicted_logits_from_labels, + reverse_kl_from_distribution_core, + softmax_base, ) -from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, _grpo_metrics_core +from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, grpo_metrics_core from fast_llm.utils import Assert @@ -110,7 +110,7 @@ def _apply_combinable_losses( cross_entropy_loss = None if ce_target is not None: loss_mask = ce_target >= 0 - predicted_logits, target_masked, target_mask = _predicted_logits_from_labels( + predicted_logits, target_masked, target_mask = predicted_logits_from_labels( logits_norm, ce_target, loss_mask, group ) cross_entropy_loss = ((sum_exp_logits.log() - predicted_logits) * loss_mask).sum() / ce_divisor @@ -152,7 +152,7 @@ def _apply_combinable_losses( else distribution_grad_output / distribution_divisor * logits_scale_factor ) if distribution_entropy_loss_type == EntropyLossType.reverse_kl: - per_sample_loss, distribution_grad = _reverse_kl_from_distribution_core( + per_sample_loss, distribution_grad = reverse_kl_from_distribution_core( logits_norm, exp_logits, sum_exp_logits, @@ -164,7 +164,7 @@ def _apply_combinable_losses( distribution_temperature, ) else: - per_sample_loss, distribution_grad = _cross_entropy_from_distribution_core( + per_sample_loss, distribution_grad = cross_entropy_from_distribution_core( logits_norm, exp_logits, sum_exp_logits, @@ -189,7 +189,7 @@ def _apply_combinable_losses( grpo_metrics = None if grpo_target is not None: grpo_loss_mask = grpo_target >= 0 - grpo_predicted_logits, grpo_target_masked, grpo_target_mask = _predicted_logits_from_labels( + grpo_predicted_logits, grpo_target_masked, grpo_target_mask = predicted_logits_from_labels( logits_norm, grpo_target, grpo_loss_mask, group ) new_log_probs = grpo_predicted_logits - sum_exp_logits.log() @@ -204,7 +204,7 @@ def _apply_combinable_losses( grpo_new_logprobs_mean = (new_log_probs * grpo_loss_mask / grpo_num_labels_in_seq.clamp(min=1)).sum() if grpo_compute_metrics: # The metric family reuses this branch's softmax + new_log_probs — no second softmax pass. - grpo_metrics = _grpo_metrics_core( + grpo_metrics = grpo_metrics_core( logits_norm, exp_logits, sum_exp_logits, @@ -289,7 +289,7 @@ def _monolithic_core( Gradients accumulate in fp32 and are cast to `logits.dtype` once at the end. GSPO is *not* handled here — its eager segment seam requires the three-phase split (`_monolithic_gspo_forward_core`). """ - logits_norm, exp_logits, sum_exp_logits, logits_max = _softmax_base(logits, logits_scale_factor, group) + logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, logits_scale_factor, group) cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad = ( _apply_combinable_losses( logits_norm, @@ -369,9 +369,9 @@ def _monolithic_gspo_forward_core( """ GSPO-present forward: one shared softmax, every enabled non-GSPO loss's scalar + fp32 gradient (uncast), and GSPO's per-token `new_log_probs`. The softmax tensors and the uncast gradient cross the - eager segment seam to `_gspo_backward_core` (the `index_add_`/`num_segments` seam can't be compiled). + eager segment seam to `gspo_backward_core` (the `index_add_`/`num_segments` seam can't be compiled). """ - logits_norm, exp_logits, sum_exp_logits, logits_max = _softmax_base(logits, logits_scale_factor, group) + logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, logits_scale_factor, group) cross_entropy_loss, z_loss, distribution_loss, _, _, _, grad = _apply_combinable_losses( logits_norm, exp_logits, @@ -405,7 +405,7 @@ def _monolithic_gspo_forward_core( False, ) gspo_loss_mask = gspo_target >= 0 - predicted_logits, target_masked, target_mask = _predicted_logits_from_labels( + predicted_logits, target_masked, target_mask = predicted_logits_from_labels( logits_norm, gspo_target, gspo_loss_mask, group ) new_log_probs = predicted_logits - sum_exp_logits.log() @@ -422,7 +422,7 @@ def _monolithic_gspo_forward_core( ) -def _gspo_segment_seam( +def gspo_segment_seam( new_log_probs: torch.Tensor, # (*batch,) loss_mask: torch.Tensor, # (*batch,) bool advantages: torch.Tensor, # (*batch,) @@ -503,7 +503,7 @@ def _gspo_segment_seam( @torch.compile -def _gspo_backward_core( +def gspo_backward_core( exp_logits: torch.Tensor, # (*batch, vocab) sum_exp_logits: torch.Tensor, # (*batch,) target_masked: torch.Tensor, # (*batch,) @@ -652,7 +652,7 @@ def monolithic_head_loss_forward_backward( gspo_spec.target, ) gspo_loss_mask = gspo_spec.target >= 0 - gspo_loss, gspo_new_logprobs_mean, effective_grad = _gspo_segment_seam( + gspo_loss, gspo_new_logprobs_mean, effective_grad = gspo_segment_seam( gspo_new_log_probs, gspo_loss_mask, gspo_spec.advantages, @@ -669,7 +669,7 @@ def monolithic_head_loss_forward_backward( logits_scale_factor, ) if effective_grad is not None: - grad_logits = _gspo_backward_core( + grad_logits = gspo_backward_core( exp_logits, sum_exp_logits, target_masked, diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 080edd4ad..2a624be57 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -7,10 +7,10 @@ from fast_llm.engine.distributed.config import DistributedConfig from fast_llm.functional.config import TritonConfig from fast_llm.functional.entropy_loss import ( - _predicted_logits_from_labels, - _softmax_base, fused_predicted_logits_from_labels, fused_softmax_base, + predicted_logits_from_labels, + softmax_base, ) from fast_llm.functional.utils import reduce_losses from fast_llm.layers.block.config import BlockKwargs @@ -22,13 +22,13 @@ LanguageModelLossKwargs, LanguageModelPolicyGradientLossConfig, ) -from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, _grpo_metrics_core +from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, grpo_metrics_core from fast_llm.layers.language_model.loss.loss import LanguageModelLoss from fast_llm.layers.language_model.loss.monolithic import ( MonolithicLossOutput, MonolithicLossSpec, - _gspo_backward_core, - _gspo_segment_seam, + gspo_backward_core, + gspo_segment_seam, ) from fast_llm.utils import Assert @@ -379,7 +379,7 @@ def compute_grpo_metrics( logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) predicted_logits, _, _ = fused_predicted_logits_from_labels(logits_norm, target, loss_mask, group) new_log_probs = predicted_logits - sum_exp_logits.log() - return _grpo_metrics_core( + return grpo_metrics_core( logits_norm, exp_logits, sum_exp_logits, @@ -481,8 +481,8 @@ def _gspo_forward_core( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: """GSPO compiled forward: one softmax + predicted-logit lookup, producing per-token new log-probs. The softmax tensors are handed to the eager segment seam and the compiled backward.""" - logits_norm, exp_logits, sum_exp_logits, _ = _softmax_base(logits, logits_scale_factor, group) - predicted_logits, target_masked, target_mask = _predicted_logits_from_labels(logits_norm, target, loss_mask, group) + logits_norm, exp_logits, sum_exp_logits, _ = softmax_base(logits, logits_scale_factor, group) + 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 new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask @@ -533,7 +533,7 @@ def fused_gspo_loss_forward_backward( new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask = _gspo_forward_core( logits, target, loss_mask, logits_scale_factor, group ) - loss, new_logprobs_mean, effective_grad = _gspo_segment_seam( + loss, new_logprobs_mean, effective_grad = gspo_segment_seam( new_log_probs, loss_mask, advantages, @@ -550,7 +550,7 @@ def fused_gspo_loss_forward_backward( logits_scale_factor, ) if effective_grad is not None: - grad_logits = _gspo_backward_core( + grad_logits = gspo_backward_core( exp_logits, sum_exp_logits, target_masked, From 40a1f374d4f2c868a97669982b49757e247d62f3 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 30 Jun 2026 16:31:26 -0400 Subject: [PATCH 11/28] Address review (#549): gate GRPO metrics off-log, drop triton enum, share CE/z-loss cores - GRPO `get_monolithic_spec` now skips `new_logprobs_mean` and the GRPO metric family when `losses is None` (nothing logged this step), mirroring the per-loss path; threaded `losses` through `get_monolithic_spec`. GSPO stays ungated to match its standalone. - Drop the unused `triton` `LossImplementation` value (it silently ran the fused path); it returns when the Triton kernel lands. - Extract `cross_entropy_from_labels_core` / `z_loss_core` (plain, taking the shared softmax tensors) so the monolithic kernel and the standalone `fused_*` kernels share one copy of the CE/z-loss math, matching the distribution cores. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/entropy_loss.py | 63 +++++++++++++++---- fast_llm/layers/language_model/config.py | 6 +- fast_llm/layers/language_model/head.py | 2 +- .../language_model/loss/entropy_loss.py | 8 ++- fast_llm/layers/language_model/loss/loss.py | 5 +- .../layers/language_model/loss/monolithic.py | 35 ++++------- .../language_model/loss/policy_gradient.py | 21 +++++-- fast_llm/layers/language_model/loss/z_loss.py | 19 +++--- 8 files changed, 104 insertions(+), 55 deletions(-) diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index a537bf7df..5f6c6a650 100644 --- a/fast_llm/functional/entropy_loss.py +++ b/fast_llm/functional/entropy_loss.py @@ -315,19 +315,23 @@ def predicted_logits_from_labels( fused_predicted_logits_from_labels = torch.compile(predicted_logits_from_labels) -@torch.compile -def _fused_cross_entropy_base_from_labels( - logits: torch.Tensor, # (*batch, vocab) +def cross_entropy_from_labels_core( + logits_norm: torch.Tensor, # (*batch, vocab), == logits - max + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) target: torch.Tensor, # (*batch,) - loss_mask: torch.Tensor, # (*batch,) - grad_output: float | None, - logits_scale_factor: float, + loss_mask: torch.Tensor, # (*batch,), == target>=0 + grad_output: float | None, # already normalized: raw_grad_output / divisor * logits_scale_factor group: ProcessGroup | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) - predicted_logits, target_masked, target_mask = fused_predicted_logits_from_labels( - logits_norm, target, loss_mask, group - ) +) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,) unmasked, (*batch, vocab) unmasked + """ + Cross-entropy from labels, taking the already-computed shared softmax tensors. Returns the unmasked + per-sample loss and (when `grad_output` is given) the unmasked gradient; the caller applies the loss + mask, reduction, and dtype cast. This plain (un-compiled) core is shared between the public + `_fused_cross_entropy_base_from_labels` wrapper and the monolithic head-loss kernel, which inlines it + inside its own `@torch.compile` boundary. + """ + predicted_logits, target_masked, target_mask = predicted_logits_from_labels(logits_norm, target, loss_mask, group) # CE loss = mean(log(sum_exp_logits) - sum(probabilities * logits)) # KL loss is the same because P * log(P) == 0. @@ -346,6 +350,43 @@ def _fused_cross_entropy_base_from_labels( return per_sample_loss, grad +def z_loss_core( + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + logits_max: torch.Tensor, # (*batch,) + grad_output: float | None, # already normalized: raw_grad_output / divisor * logits_scale_factor +) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,) unmasked, (*batch, vocab) unmasked + """ + Z-loss from the already-computed shared softmax tensors. Returns the unmasked per-sample loss term + (`log_sum_exp ** 2`) and (when `grad_output` is given) the unmasked gradient; the caller applies the + loss mask, reduction, and dtype cast. z-loss needs the un-regularized log-sum-exp, so it adds back + `logits_max` (cross-entropy cancels it). Shared between `fused_z_loss_forward_backward` and the + monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. + """ + log_sum_exp_logits = sum_exp_logits.log() + logits_max + loss_term = log_sum_exp_logits**2 + if grad_output is None: + grad = None + else: + grad = (2 * grad_output * (log_sum_exp_logits / sum_exp_logits)).unsqueeze(-1) * exp_logits + return loss_term, grad + + +@torch.compile +def _fused_cross_entropy_base_from_labels( + logits: torch.Tensor, # (*batch, vocab) + target: torch.Tensor, # (*batch,) + loss_mask: torch.Tensor, # (*batch,) + grad_output: float | None, + logits_scale_factor: float, + group: ProcessGroup | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) + logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) + return cross_entropy_from_labels_core( + logits_norm, exp_logits, sum_exp_logits, target, loss_mask, grad_output, group + ) + + @torch.compile def fused_entropy_loss_forward_backward( logits: torch.Tensor, # (*batch, vocab) diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index b1d9fb5c4..0d7d7aa91 100644 --- a/fast_llm/layers/language_model/config.py +++ b/fast_llm/layers/language_model/config.py @@ -39,8 +39,6 @@ class LossImplementation(enum.StrEnum): per_loss = "per_loss" # A single monolithic torch.compile kernel running the softmax once for all combinable losses. fused = "fused" - # A single monolithic triton kernel. - triton = "triton" @config_class() @@ -156,8 +154,8 @@ class LanguageModelHeadConfig(BlockConfig): loss_implementation: LossImplementation = Field( default=LossImplementation.per_loss, desc="Select the head-loss implementation. `per_loss` runs each loss separately (the default);" - " `fused`/`triton` run a single monolithic kernel that shares one softmax pass across combinable" - " losses. Losses not yet supported by the monolithic kernel fall back to their own implementation.", + " `fused` runs a single monolithic kernel that shares one softmax pass across combinable losses." + " Losses not supported by the monolithic kernel fall back to their own implementation.", hint=FieldHint.expert, ) diff --git a/fast_llm/layers/language_model/head.py b/fast_llm/layers/language_model/head.py index 500651d8b..dfac82bf5 100644 --- a/fast_llm/layers/language_model/head.py +++ b/fast_llm/layers/language_model/head.py @@ -304,7 +304,7 @@ def _monolithic_loss_forward_backward( # through their own `forward_backward`, accumulating into the same logits gradient. specs, monolithic_losses, other_losses = [], [], [] for loss in self.losses: - spec = loss.get_monolithic_spec(kwargs, split_index) + spec = loss.get_monolithic_spec(kwargs, split_index, losses) if spec is None: other_losses.append(loss) else: diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py index 4d024729d..bf9f29973 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -39,7 +39,9 @@ def _forward_backward( divisor=self._get_label_count(kwargs), ) - def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + def get_monolithic_spec( + self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None + ) -> MonolithicLossSpec | None: # For labels, forward-KL is identical to cross-entropy (one-hot target entropy is zero). return MonolithicLossSpec( kind="cross_entropy", @@ -79,7 +81,9 @@ def _forward_backward( divisor=self._get_label_count(kwargs), ) - def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + def get_monolithic_spec( + self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None + ) -> MonolithicLossSpec | None: return MonolithicLossSpec( kind="entropy_from_distribution", name=self.name, diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 868fe6576..1f10cdf5a 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -73,10 +73,13 @@ def _forward_backward( def get_loss_definitions(self) -> list[LossDef]: return [LossDef(name=self.name)] if self._do_register_loss else [] - def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> "MonolithicLossSpec | None": + def get_monolithic_spec( + self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None + ) -> "MonolithicLossSpec | None": """ Package this loss's inputs for the monolithic head-loss kernel, or return `None` if it is not supported there (it then runs through its own `forward_backward`, accumulating into the same grad). + `losses is None` signals a step that logs nothing, so a loss may switch off metric-only outputs. """ return None diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 77593d81f..1b0e2e49b 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -6,9 +6,11 @@ from fast_llm.functional.config import EntropyLossType, TargetFormat from fast_llm.functional.entropy_loss import ( cross_entropy_from_distribution_core, + cross_entropy_from_labels_core, predicted_logits_from_labels, reverse_kl_from_distribution_core, softmax_base, + z_loss_core, ) from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, grpo_metrics_core from fast_llm.utils import Assert @@ -110,38 +112,27 @@ def _apply_combinable_losses( cross_entropy_loss = None if ce_target is not None: loss_mask = ce_target >= 0 - predicted_logits, target_masked, target_mask = predicted_logits_from_labels( - logits_norm, ce_target, loss_mask, group + ce_grad_output_scaled = None if ce_grad_output is None else ce_grad_output / ce_divisor * logits_scale_factor + per_sample_loss, cross_entropy_grad = cross_entropy_from_labels_core( + logits_norm, exp_logits, sum_exp_logits, ce_target, loss_mask, ce_grad_output_scaled, group ) - cross_entropy_loss = ((sum_exp_logits.log() - predicted_logits) * loss_mask).sum() / ce_divisor - if ce_grad_output is not None: - grad_output = ce_grad_output / ce_divisor * logits_scale_factor - cross_entropy_grad = exp_logits.scatter_add( - -1, - target_masked.unsqueeze(-1), - ( - -sum_exp_logits.unsqueeze(-1) - if target_mask is None - else -(target_mask * sum_exp_logits).unsqueeze(-1) - ), - ) * (grad_output / sum_exp_logits.unsqueeze(-1)) + cross_entropy_loss = (per_sample_loss * loss_mask).sum() / ce_divisor + if cross_entropy_grad is not None: cross_entropy_grad = cross_entropy_grad * loss_mask.unsqueeze(-1) grad = cross_entropy_grad if grad is None else grad + cross_entropy_grad z_loss = None if z_loss_enabled: - # z-loss needs the un-regularized log-sum-exp, so it adds back `logits_max` (cross-entropy cancels it). - log_sum_exp_logits = sum_exp_logits.log() + logits_max - z_loss_term = log_sum_exp_logits**2 + z_loss_grad_output_scaled = ( + None if z_loss_grad_output is None else z_loss_grad_output / z_loss_divisor * logits_scale_factor + ) + z_loss_term, z_loss_grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, z_loss_grad_output_scaled) if z_loss_mask is not None: z_loss_term = z_loss_term * z_loss_mask z_loss = z_loss_term.sum() / z_loss_divisor - if z_loss_grad_output is not None: - grad_output = z_loss_grad_output / z_loss_divisor * logits_scale_factor - z_loss_grad_base = 2 * grad_output * (log_sum_exp_logits / sum_exp_logits) + if z_loss_grad is not None: if z_loss_mask is not None: - z_loss_grad_base = z_loss_grad_base * z_loss_mask - z_loss_grad = z_loss_grad_base.unsqueeze(-1) * exp_logits + z_loss_grad = z_loss_grad * z_loss_mask.unsqueeze(-1) grad = z_loss_grad if grad is None else grad + z_loss_grad distribution_loss = None diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 2a624be57..c25a78b08 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -139,7 +139,12 @@ def _forward_backward( return loss, grad - def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + def get_monolithic_spec( + self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None + ) -> MonolithicLossSpec | None: + # When nothing is logged this step, skip the metric-only outputs (`new_logprobs_mean` and the + # GRPO metric family), mirroring the per-loss path's gating in `_forward_backward`. + register_metrics = losses is not None return MonolithicLossSpec( kind="grpo", name=self.name, @@ -154,9 +159,13 @@ def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = ), epsilon_low=self._config.epsilon_low, epsilon_high=self._config.epsilon_high, - num_labels_in_seq=self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index), - compute_metrics=self._config.metrics != GRPOMetricsLevel.none, - compute_entropy=self._config.metrics == GRPOMetricsLevel.with_entropy, + num_labels_in_seq=( + self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index) + if register_metrics + else None + ), + compute_metrics=register_metrics and self._config.metrics != GRPOMetricsLevel.none, + compute_entropy=register_metrics and self._config.metrics == GRPOMetricsLevel.with_entropy, ) def register_monolithic_outputs( @@ -330,7 +339,9 @@ def _forward_backward( def get_preprocessing_config(self) -> dict[str, typing.Any]: return super().get_preprocessing_config() | {"return_document_index": True} - def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + def get_monolithic_spec( + self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None + ) -> MonolithicLossSpec | None: return MonolithicLossSpec( kind="gspo", name=self.name, diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py index 451ab48f4..ad84ddf87 100644 --- a/fast_llm/layers/language_model/loss/z_loss.py +++ b/fast_llm/layers/language_model/loss/z_loss.py @@ -3,7 +3,7 @@ import torch from fast_llm.functional.config import TritonConfig -from fast_llm.functional.entropy_loss import fused_softmax_base +from fast_llm.functional.entropy_loss import fused_softmax_base, z_loss_core from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward from fast_llm.functional.utils import reduce_losses from fast_llm.layers.language_model.loss.config import LanguageModelZLossConfig @@ -34,7 +34,9 @@ def _forward_backward( divisor=self._get_label_count(kwargs), ) - def get_monolithic_spec(self, kwargs: dict[str, typing.Any], split_index: int = 0) -> MonolithicLossSpec | None: + def get_monolithic_spec( + self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None + ) -> MonolithicLossSpec | None: return MonolithicLossSpec( kind="z_loss", name=self.name, @@ -80,16 +82,15 @@ def fused_z_loss_forward_backward( if divisor is None: divisor = logits.shape[:-1].numel() grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor - logits_norm, exp_logits, sum_exp_logits, logits_max = fused_softmax_base(logits, logits_scale_factor, group) - log_sum_exp_logits = sum_exp_logits.log() + logits_max + _, exp_logits, sum_exp_logits, logits_max = fused_softmax_base(logits, logits_scale_factor, group) + loss_term, grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, grad_output) - loss = reduce_losses(log_sum_exp_logits**2, divisor, loss_mask) + loss = reduce_losses(loss_term, divisor, loss_mask) - if grad_output is not None: - grad_base = 2 * grad_output * (log_sum_exp_logits / sum_exp_logits) + if grad is not None: if loss_mask is not None: - grad_base = grad_base * loss_mask - grad = (grad_base.unsqueeze(-1) * exp_logits).to(logits.dtype) + grad = grad * loss_mask.unsqueeze(-1) + grad = grad.to(logits.dtype) if grad_logits is None: grad_logits = grad else: From c9c95c80c91ca9ffe13d2844c3da5d3751a869ca Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 2 Jul 2026 11:00:38 -0400 Subject: [PATCH 12/28] Rework monolithic head loss as a `monolithic` loss type (#507) Replace the head-level `loss_implementation` toggle + `MonolithicLossSpec` marshalling with a `MonolithicLoss` / `MonolithicLossConfig` loss type that holds the combinable child losses as nested entries and shares one softmax across them. The head is untouched: the monolithic loss flows through the ordinary per-loss loop, and non-combinable losses (e.g. DPO) are plain siblings. - Each combinable loss (label-CE, z-loss, distillation, GRPO) owns a plain post-softmax `combinable` function + `combinable_extract` / `combinable_core`; the standalone `fused_*` kernels call the same function, so the per-loss math has one source of truth. - `_monolithic_core` iterates entries via object-method dispatch (no flat arg list, no `if enabled` chain) and fuses into a single graph (no graph breaks). - Removed `LossImplementation` enum + field, the `get_monolithic_spec` / `register_monolithic_outputs` hooks, `MonolithicLossSpec` / `MonolithicLossOutput`, the flat-arg kernel and its dispatcher; reverted head.py to main. - GSPO reverts to standalone-only (segment seam / backward moved back to policy_gradient); folding it into `MonolithicLoss` is a follow-up. - `_get_reference_model_logits` resolves the head logits by splitting on `.losses.`, so distillation works whether flat or nested in the composite. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/layers/language_model/config.py | 20 - fast_llm/layers/language_model/head.py | 63 +- fast_llm/layers/language_model/loss/config.py | 39 + .../language_model/loss/entropy_loss.py | 159 +++- fast_llm/layers/language_model/loss/loss.py | 23 +- .../layers/language_model/loss/monolithic.py | 780 +++--------------- .../language_model/loss/policy_gradient.py | 356 +++++--- fast_llm/layers/language_model/loss/z_loss.py | 61 +- tests/layers/test_lm_head.py | 40 +- tests/layers/test_lm_losses.py | 619 -------------- 10 files changed, 612 insertions(+), 1548 deletions(-) diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index 0d7d7aa91..bde33f297 100644 --- a/fast_llm/layers/language_model/config.py +++ b/fast_llm/layers/language_model/config.py @@ -1,4 +1,3 @@ -import enum import typing from fast_llm.config import Field, FieldHint, check_field, config_class, skip_valid_if_none @@ -34,13 +33,6 @@ class LanguageModelKwargs(LanguageModelLossKwargs): LM_HEAD_LOSS_NAME = "lm_head_loss" -class LossImplementation(enum.StrEnum): - # One loss kernel per loss, each running its own softmax pass (the original, always-available path). - per_loss = "per_loss" - # A single monolithic torch.compile kernel running the softmax once for all combinable losses. - fused = "fused" - - @config_class() class LanguageModelEmbeddingsConfig(BlockConfig): _abstract = False @@ -151,13 +143,6 @@ class LanguageModelHeadConfig(BlockConfig): doc="If not provided, all heads are equally weighted.", hint=FieldHint.feature, ) - loss_implementation: LossImplementation = Field( - default=LossImplementation.per_loss, - desc="Select the head-loss implementation. `per_loss` runs each loss separately (the default);" - " `fused` runs a single monolithic kernel that shares one softmax pass across combinable losses." - " Losses not supported by the monolithic kernel fall back to their own implementation.", - hint=FieldHint.expert, - ) def get_layer( self, @@ -192,11 +177,6 @@ def get_layer( def _validate(self) -> None: super()._validate() assert LM_HEAD_LOSS_NAME not in self.losses - if self.loss_implementation != LossImplementation.per_loss: - # The monolithic kernel shares a single softmax across losses, so they must agree on the - # effective logits scale factor (model scale stacked with each loss's own scale factor). - scale_factors = {self.logits_scale_factor * loss.logits_scale_factor for loss in self.losses.values()} - Assert.leq(len(scale_factors), 1) def get_reference_models(self) -> set[str]: return {reference_model for loss in self.losses.values() for reference_model in loss.get_reference_models()} diff --git a/fast_llm/layers/language_model/head.py b/fast_llm/layers/language_model/head.py index dfac82bf5..22c750082 100644 --- a/fast_llm/layers/language_model/head.py +++ b/fast_llm/layers/language_model/head.py @@ -19,11 +19,9 @@ LanguageModelEmbeddingsConfig, LanguageModelHeadConfig, LanguageModelKwargs, - LossImplementation, ) from fast_llm.layers.language_model.loss.config import LanguageModelLabelEntropyLossConfig from fast_llm.layers.language_model.loss.loss import LanguageModelLoss -from fast_llm.layers.language_model.loss.monolithic import monolithic_head_loss_forward_backward from fast_llm.tensor import TensorMeta from fast_llm.utils import Assert, safe_merge_dicts @@ -274,21 +272,18 @@ def _logits_loss_forward_backward_partial( if return_logits: return logits, None - if self._config.loss_implementation == LossImplementation.per_loss: - losses_, grad = [], None - for loss in self.losses: - # losses are returned unscaled but the grads are already scaled - loss_value, grad = loss.forward_backward( - logits, - kwargs, - losses, - split_index, - grad, - ) - if loss_value is not None: - losses_.append(loss_value.detach()) - else: - losses_, grad = self._monolithic_loss_forward_backward(logits, kwargs, losses, split_index) + losses_, grad = [], None + for loss in self.losses: + # losses are returned unscaled but the grads are already scaled + loss_value, grad = loss.forward_backward( + logits, + kwargs, + losses, + split_index, + grad, + ) + if loss_value is not None: + losses_.append(loss_value.detach()) if grad is not None and self._config.final_logit_softcap is not None: grad = _softcap_backward(grad, logits, self._config.final_logit_softcap) @@ -297,40 +292,6 @@ def _logits_loss_forward_backward_partial( output_parallel_linear_backward(grad, context) if self.training else None ) - def _monolithic_loss_forward_backward( - self, logits: torch.Tensor, kwargs: dict, losses: dict | None, split_index: int - ) -> tuple[list[torch.Tensor], torch.Tensor | None]: - # Losses supported by the monolithic kernel share one softmax pass; the rest (e.g. DPO) run - # through their own `forward_backward`, accumulating into the same logits gradient. - specs, monolithic_losses, other_losses = [], [], [] - for loss in self.losses: - spec = loss.get_monolithic_spec(kwargs, split_index, losses) - if spec is None: - other_losses.append(loss) - else: - specs.append(spec) - monolithic_losses.append(loss) - - losses_, grad = [], None - if specs: - outputs, grad = monolithic_head_loss_forward_backward( - logits, - specs, - group=self._parallel_dim.group if self._vocab_parallel else None, - grad_logits=grad, - ) - for loss, output in zip(monolithic_losses, outputs, strict=True): - loss.register_monolithic_outputs(output, kwargs, losses) - if output.loss is not None: - losses_.append((output.loss if loss.weight == 1 else output.loss * loss.weight).detach()) - - for loss in other_losses: - loss_value, grad = loss.forward_backward(logits, kwargs, losses, split_index, grad) - if loss_value is not None: - losses_.append(loss_value.detach()) - - return losses_, grad - def get_loss_definitions(self) -> list[LossDef]: return [ LossDef(name=self._total_loss_name), diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 9a220aacf..920d9bafe 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -36,6 +36,8 @@ class LanguageModelLossKwargs(BlockKwargs): @config_class(registry=True) class LanguageModelLossConfig(Config): _abstract: typing.ClassVar[bool] = True + # Whether this loss can be a `MonolithicLossConfig` entry (shares one softmax via `combinable_core`). + combinable: typing.ClassVar[bool] = False weight: float = Field( default=1.0, @@ -89,6 +91,7 @@ def get_reference_models(self) -> set[str]: @config_class(dynamic_type={LanguageModelLossConfig: "label"}) class LanguageModelLabelEntropyLossConfig(LanguageModelLossConfig): _abstract: typing.ClassVar[bool] = False + combinable: typing.ClassVar[bool] = True loss_type: EntropyLossType = Field( default=EntropyLossType.cross_entropy, @@ -118,6 +121,7 @@ def loss_class(self) -> "type[LanguageModelLabelEntropyLoss]": @config_class(dynamic_type={LanguageModelLossConfig: "distillation"}) class LanguageModelDistillationLossConfig(LanguageModelLossConfig): _abstract: typing.ClassVar[bool] = False + combinable: typing.ClassVar[bool] = True loss_type: EntropyLossType = Field( default=EntropyLossType.cross_entropy, @@ -191,6 +195,7 @@ class LanguageModelZLossConfig(LanguageModelLossConfig): """Z-loss regularization to prevent overconfidence.""" _abstract: typing.ClassVar[bool] = False + combinable: typing.ClassVar[bool] = True use_triton: bool | None = Field( default=None, @@ -230,6 +235,7 @@ class LanguageModelGRPOLossConfig(LanguageModelPolicyGradientLossConfig): """Group-Relative Policy Optimization: per-token IS-ratio clipping.""" _abstract: typing.ClassVar[bool] = False + combinable: typing.ClassVar[bool] = True use_triton: bool | None = Field( default=None, @@ -271,3 +277,36 @@ def loss_class(self) -> "type[LanguageModelGSPOLoss]": from fast_llm.layers.language_model.loss.policy_gradient import LanguageModelGSPOLoss return LanguageModelGSPOLoss + + +@config_class(dynamic_type={LanguageModelLossConfig: "monolithic"}) +class MonolithicLossConfig(LanguageModelLossConfig): + """A composite loss that runs one vocabulary softmax and shares it across its combinable child losses.""" + + _abstract: typing.ClassVar[bool] = False + + losses: dict[str, LanguageModelLossConfig] = Field( + default_factory=dict, + desc="The combinable losses sharing a single softmax pass. They must agree on `logits_scale_factor`.", + hint=FieldHint.core, + ) + + def _validate(self) -> None: + super()._validate() + Assert.gt(len(self.losses), 0) + for name, loss in self.losses.items(): + if not loss.combinable: + raise ValueError( + f"Loss `{name}` (`{type(loss).__name__}`) is not combinable and cannot share the softmax." + ) + # A single softmax serves one effective scale (stacked with the common model scale). + Assert.eq(len({loss.logits_scale_factor for loss in self.losses.values()}), 1) + + @property + def loss_class(self) -> "type[LanguageModelLoss]": + from fast_llm.layers.language_model.loss.monolithic import MonolithicLoss + + return MonolithicLoss + + def get_reference_models(self) -> set[str]: + return {reference_model for loss in self.losses.values() for reference_model in loss.get_reference_models()} diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py index bf9f29973..b8915dfa2 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -2,15 +2,19 @@ import torch -from fast_llm.functional.config import TargetFormat, TritonConfig -from fast_llm.functional.entropy_loss import fused_entropy_loss_forward_backward +from fast_llm.functional.config import EntropyLossType, TargetFormat, TritonConfig +from fast_llm.functional.entropy_loss import ( + cross_entropy_from_distribution_core, + cross_entropy_from_labels_core, + fused_entropy_loss_forward_backward, + reverse_kl_from_distribution_core, +) from fast_llm.functional.triton.entropy_loss import triton_entropy_loss_forward_backward from fast_llm.layers.language_model.loss.config import ( LanguageModelDistillationLossConfig, LanguageModelLabelEntropyLossConfig, ) from fast_llm.layers.language_model.loss.loss import LanguageModelLoss -from fast_llm.layers.language_model.loss.monolithic import MonolithicLossSpec class LanguageModelLabelEntropyLoss[ConfigType: LanguageModelLabelEntropyLossConfig](LanguageModelLoss[ConfigType]): @@ -39,19 +43,25 @@ def _forward_backward( divisor=self._get_label_count(kwargs), ) - def get_monolithic_spec( - self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None - ) -> MonolithicLossSpec | None: + def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + return self._get_labels(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs) + + def combinable_core( + self, + 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[torch.Tensor, torch.Tensor | None, None]: # For labels, forward-KL is identical to cross-entropy (one-hot target entropy is zero). - return MonolithicLossSpec( - kind="cross_entropy", - name=self.name, - weight=self._weight, - logits_scale_factor=self._logits_scale_factor, - grad_output=self._get_grad_output(kwargs), - divisor=self._get_label_count(kwargs), - target=self._get_labels(kwargs, split_index), + target, grad_output, divisor = arguments + loss, grad = cross_entropy_from_labels_combinable( + logits_norm, exp_logits, sum_exp_logits, group, target, grad_output, divisor, logits_scale_factor ) + return loss, grad, None class LanguageModelDistillationLoss[ConfigType: LanguageModelDistillationLossConfig](LanguageModelLoss[ConfigType]): @@ -81,22 +91,113 @@ def _forward_backward( divisor=self._get_label_count(kwargs), ) - def get_monolithic_spec( - self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None - ) -> MonolithicLossSpec | None: - return MonolithicLossSpec( - kind="entropy_from_distribution", - name=self.name, - weight=self._weight, - logits_scale_factor=self._logits_scale_factor, - grad_output=self._get_grad_output(kwargs), - divisor=self._get_label_count(kwargs), - target=self._get_reference_model_logits(self._config.reference_model, kwargs, split_index), - loss_mask=self._get_loss_mask(kwargs, split_index), - target_format=TargetFormat.logits, - entropy_loss_type=self._config.loss_type, - temperature=self._config.temperature, + def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + return ( + self._get_reference_model_logits(self._config.reference_model, kwargs, split_index), + self._get_loss_mask(kwargs, split_index), + self._get_grad_output(kwargs), + self._get_label_count(kwargs), + self._config.loss_type, + self._config.temperature, + ) + + def combinable_core( + self, + 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[torch.Tensor, torch.Tensor | None, None]: + target, loss_mask, grad_output, divisor, entropy_loss_type, temperature = arguments + loss, grad = distillation_combinable( + logits_norm, + exp_logits, + sum_exp_logits, + group, + target, + loss_mask, + grad_output, + divisor, + logits_scale_factor, + entropy_loss_type, + temperature, ) + return loss, grad, None def get_preprocessing_config(self) -> dict[str, typing.Any]: return {"return_prediction_mask": True} + + +def cross_entropy_from_labels_combinable( + logits_norm: torch.Tensor, # (*batch, vocab) + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + group: "torch.distributed.ProcessGroup | None", + target: torch.Tensor, # (*batch,) + grad_output: float | None, + divisor: float, + logits_scale_factor: float, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Post-softmax cross-entropy-from-labels block over the shared softmax: loss scalar + uncast, masked + gradient. The caller casts to the logits dtype (the monolithic path defers that to one final cast).""" + loss_mask = target >= 0 + grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor + per_sample_loss, grad = cross_entropy_from_labels_core( + logits_norm, exp_logits, sum_exp_logits, target, loss_mask, grad_output, group + ) + loss = (per_sample_loss * loss_mask).sum() / divisor + if grad is not None: + grad = grad * loss_mask.unsqueeze(-1) + return loss, grad + + +def distillation_combinable( + logits_norm: torch.Tensor, # (*batch, vocab) + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + group: "torch.distributed.ProcessGroup | None", + target: torch.Tensor, # (*batch, vocab) teacher logits + loss_mask: torch.Tensor | None, # (*batch,) + grad_output: float | None, + divisor: float, + logits_scale_factor: float, + entropy_loss_type: EntropyLossType, + temperature: float, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Post-softmax distillation block (cross-entropy / forward-KL / reverse-KL from a teacher distribution) + over the shared student softmax: loss scalar + uncast, masked gradient (the caller casts).""" + grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor + if entropy_loss_type == EntropyLossType.reverse_kl: + per_sample_loss, grad = reverse_kl_from_distribution_core( + logits_norm, + exp_logits, + sum_exp_logits, + target, + grad_output, + logits_scale_factor, + TargetFormat.logits, + group, + temperature, + ) + else: + per_sample_loss, grad = cross_entropy_from_distribution_core( + logits_norm, + exp_logits, + sum_exp_logits, + target, + grad_output, + logits_scale_factor, + TargetFormat.logits, + group, + temperature, + return_kl_loss=entropy_loss_type == EntropyLossType.forward_kl, + ) + if loss_mask is not None: + per_sample_loss = per_sample_loss * loss_mask + loss = per_sample_loss.sum() / divisor + if grad is not None and loss_mask is not None: + grad = grad * loss_mask.unsqueeze(-1) + return loss, grad diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 1f10cdf5a..5ddffef32 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -9,7 +9,6 @@ from fast_llm.engine.distributed.config import DistributedConfig, DistributedDimNames from fast_llm.layers.language_model.config import LanguageModelKwargs from fast_llm.layers.language_model.loss.config import LanguageModelLossConfig, LanguageModelLossKwargs -from fast_llm.layers.language_model.loss.monolithic import MonolithicLossOutput, MonolithicLossSpec from fast_llm.utils import Assert @@ -73,21 +72,11 @@ def _forward_backward( def get_loss_definitions(self) -> list[LossDef]: return [LossDef(name=self.name)] if self._do_register_loss else [] - def get_monolithic_spec( - self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None - ) -> "MonolithicLossSpec | None": - """ - Package this loss's inputs for the monolithic head-loss kernel, or return `None` if it is not - supported there (it then runs through its own `forward_backward`, accumulating into the same grad). - `losses is None` signals a step that logs nothing, so a loss may switch off metric-only outputs. - """ - return None - - def register_monolithic_outputs( - self, output: "MonolithicLossOutput", kwargs: dict[str, typing.Any], losses: dict | None + def register_combinable_extras( + self, extra: typing.Any, kwargs: dict[str, typing.Any], losses: dict | None ) -> None: - if self._do_register_loss: - self._register_loss(self.name, output.loss, losses) + """Register any per-loss outputs beyond the scalar (e.g. GRPO's `new_logprobs` / metric family) + produced by `combinable_core` when this loss runs inside `MonolithicLoss`. No-op by default.""" def get_preprocessing_config( self, @@ -147,8 +136,10 @@ def _get_loss_mask(self, kwargs: dict[str, typing.Any], split_index: int = 0): return None if loss_mask is None else self._prepare_target(loss_mask, split_index) def _get_reference_model_logits(self, reference_model: str, kwargs: dict[str, typing.Any], split_index: int = 0): + # The head owns the `.logits`; split on `.losses.` so the name resolves whether this loss sits + # directly under the head or nested inside a composite loss (e.g. `MonolithicLoss`). Assert.incl( - logits_name := self.module_name.rsplit(".", 2)[0] + f".logits", + logits_name := self.module_name.split(".losses.")[0] + f".logits", reference_hidden_states := kwargs[f"reference_{reference_model}_hidden_states"], ) # The logits are already sequence-parallel if needed, we don't want to split again. diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 1b0e2e49b..c4f57dd8e 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -3,698 +3,142 @@ import torch from fast_llm.core.distributed import ProcessGroup -from fast_llm.functional.config import EntropyLossType, TargetFormat -from fast_llm.functional.entropy_loss import ( - cross_entropy_from_distribution_core, - cross_entropy_from_labels_core, - predicted_logits_from_labels, - reverse_kl_from_distribution_core, - softmax_base, - z_loss_core, -) -from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, grpo_metrics_core -from fast_llm.utils import Assert - - -class MonolithicLossSpec(typing.NamedTuple): - """ - Per-loss inputs gathered by `LanguageModelLoss.get_monolithic_spec` and consumed by the monolithic - head-loss kernel. `kind` selects which branch of the kernel computes the loss and its gradient. - The math fields not used by a given `kind` are left at their defaults. - """ - - kind: str - name: str - weight: float - logits_scale_factor: float - grad_output: float | None - divisor: float - target: torch.Tensor | None = ( - None # Cross-entropy / GRPO / GSPO labels, or a teacher distribution (logits/probabilities). - ) - loss_mask: torch.Tensor | None = None # z-loss / distribution-loss mask (cross-entropy / GRPO derive their own). - # Distribution losses (cross-entropy / forward-KL / reverse-KL from a teacher distribution) only. - target_format: TargetFormat = TargetFormat.logits - entropy_loss_type: EntropyLossType = EntropyLossType.cross_entropy - temperature: float = 1.0 - # Policy-gradient (GRPO / GSPO) only. - advantages: torch.Tensor | None = None - old_log_probabilities: torch.Tensor | None = None - epsilon_low: float = 0.2 - epsilon_high: float = 0.2 - num_labels_in_seq: torch.Tensor | None = None # Per-sequence label count; enables the `new_logprobs_mean` metric. - compute_metrics: bool = False # Emit the GRPO metric family from the shared softmax. - compute_entropy: bool = False # Also emit per-token entropy (the only metric needing the vocab axis). - # Sequence-policy-gradient (GSPO) only. - document_index: torch.Tensor | None = None # Per-token segment ID (0-based). - num_segments: int = 0 # Per-segment buffer size, ≥ document_index.max() + 1. - sdp_group: ProcessGroup | None = None # Sequence-data group for cross-rank segment aggregation. - sp_group: ProcessGroup | None = None # TP group when sequence-parallel shards the sequence. - - -class MonolithicLossOutput(typing.NamedTuple): - """Per-loss outputs returned by the monolithic kernel, registered via `register_monolithic_outputs`.""" - - loss: torch.Tensor | None = None - new_logprobs_mean: torch.Tensor | None = None # GRPO / GSPO only. - metrics: GRPOMetrics | None = None # GRPO only. - - -def _apply_combinable_losses( - logits_norm: torch.Tensor, # (*batch, vocab) - exp_logits: torch.Tensor, # (*batch, vocab) - sum_exp_logits: torch.Tensor, # (*batch,) - logits_max: torch.Tensor, # (*batch,) - group: ProcessGroup | None, - logits_scale_factor: float, - ce_target: torch.Tensor | None, - ce_grad_output: float | None, - ce_divisor: float, - z_loss_enabled: bool, - z_loss_mask: torch.Tensor | None, - z_loss_grad_output: float | None, - z_loss_divisor: float, - distribution_target: torch.Tensor | None, - distribution_grad_output: float | None, - distribution_divisor: float, - distribution_loss_mask: torch.Tensor | None, - distribution_target_format: TargetFormat, - distribution_entropy_loss_type: EntropyLossType, - distribution_temperature: float, - grpo_target: torch.Tensor | None, - grpo_advantages: torch.Tensor | None, - grpo_old_log_probabilities: torch.Tensor | None, - grpo_grad_output: float | None, - grpo_divisor: float, - grpo_epsilon_low: float, - grpo_epsilon_high: float, - grpo_num_labels_in_seq: torch.Tensor | None, - grpo_compute_metrics: bool, - grpo_compute_entropy: bool, -) -> tuple[ - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - "GRPOMetrics | None", - torch.Tensor | None, -]: - """ - The per-loss scalar + fp32 gradient-accumulator math shared by the single-pass (`_monolithic_core`) - and GSPO-split (`_monolithic_gspo_forward_core`) paths. Takes the already-computed shared softmax - tensors and returns each enabled loss's scalar plus the summed gradient, accumulated in fp32 and - *not* cast (the caller casts once). Plain (un-decorated) so it inlines and fuses into its compiled - callers; a disabled loss is dead-code-eliminated on its `None`/static flag. - """ - grad = None - - cross_entropy_loss = None - if ce_target is not None: - loss_mask = ce_target >= 0 - ce_grad_output_scaled = None if ce_grad_output is None else ce_grad_output / ce_divisor * logits_scale_factor - per_sample_loss, cross_entropy_grad = cross_entropy_from_labels_core( - logits_norm, exp_logits, sum_exp_logits, ce_target, loss_mask, ce_grad_output_scaled, group - ) - cross_entropy_loss = (per_sample_loss * loss_mask).sum() / ce_divisor - if cross_entropy_grad is not None: - cross_entropy_grad = cross_entropy_grad * loss_mask.unsqueeze(-1) - grad = cross_entropy_grad if grad is None else grad + cross_entropy_grad - - z_loss = None - if z_loss_enabled: - z_loss_grad_output_scaled = ( - None if z_loss_grad_output is None else z_loss_grad_output / z_loss_divisor * logits_scale_factor - ) - z_loss_term, z_loss_grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, z_loss_grad_output_scaled) - if z_loss_mask is not None: - z_loss_term = z_loss_term * z_loss_mask - z_loss = z_loss_term.sum() / z_loss_divisor - if z_loss_grad is not None: - if z_loss_mask is not None: - z_loss_grad = z_loss_grad * z_loss_mask.unsqueeze(-1) - grad = z_loss_grad if grad is None else grad + z_loss_grad - - distribution_loss = None - if distribution_target is not None: - distribution_grad_output_scaled = ( - None - if distribution_grad_output is None - else distribution_grad_output / distribution_divisor * logits_scale_factor - ) - if distribution_entropy_loss_type == EntropyLossType.reverse_kl: - per_sample_loss, distribution_grad = reverse_kl_from_distribution_core( - logits_norm, - exp_logits, - sum_exp_logits, - distribution_target, - distribution_grad_output_scaled, - logits_scale_factor, - distribution_target_format, - group, - distribution_temperature, - ) - else: - per_sample_loss, distribution_grad = cross_entropy_from_distribution_core( - logits_norm, - exp_logits, - sum_exp_logits, - distribution_target, - distribution_grad_output_scaled, - logits_scale_factor, - distribution_target_format, - group, - distribution_temperature, - return_kl_loss=distribution_entropy_loss_type == EntropyLossType.forward_kl, - ) - if distribution_loss_mask is not None: - per_sample_loss = per_sample_loss * distribution_loss_mask - distribution_loss = per_sample_loss.sum() / distribution_divisor - if distribution_grad is not None: - if distribution_loss_mask is not None: - distribution_grad = distribution_grad * distribution_loss_mask.unsqueeze(-1) - grad = distribution_grad if grad is None else grad + distribution_grad - - grpo_loss = None - grpo_new_logprobs_mean = None - grpo_metrics = None - if grpo_target is not None: - grpo_loss_mask = grpo_target >= 0 - grpo_predicted_logits, grpo_target_masked, grpo_target_mask = predicted_logits_from_labels( - logits_norm, grpo_target, grpo_loss_mask, group - ) - new_log_probs = grpo_predicted_logits - sum_exp_logits.log() - probability_ratio = (new_log_probs - grpo_old_log_probabilities).exp() - grpo_losses = -torch.min( - probability_ratio * grpo_advantages, - torch.clamp(probability_ratio, 1 - grpo_epsilon_low, 1 + grpo_epsilon_high) * grpo_advantages, - ) - grpo_loss = (grpo_losses * grpo_loss_mask).sum() / grpo_divisor - if grpo_num_labels_in_seq is not None: - # Sum of per-sequence mean log-probs; clamp avoids 0/0 at fully-masked documents (also loss-masked). - grpo_new_logprobs_mean = (new_log_probs * grpo_loss_mask / grpo_num_labels_in_seq.clamp(min=1)).sum() - if grpo_compute_metrics: - # The metric family reuses this branch's softmax + new_log_probs — no second softmax pass. - grpo_metrics = grpo_metrics_core( - logits_norm, - exp_logits, - sum_exp_logits, - new_log_probs, - grpo_advantages, - grpo_old_log_probabilities, - grpo_loss_mask, - grpo_num_labels_in_seq, - grpo_epsilon_low, - grpo_epsilon_high, - group, - grpo_compute_entropy, - ) - if grpo_grad_output is not None: - grad_output = grpo_grad_output / grpo_divisor * logits_scale_factor - # grad[a>=0] = -a * (ratio <= 1 + epsilon_high); grad[a<=0] = a * (ratio >= 1 - epsilon_low). - probability_ratio_grad = ( - grad_output - * ( - torch.clamp_min(grpo_advantages, 0) * (probability_ratio <= 1 + grpo_epsilon_high) - + torch.clamp_max(grpo_advantages, 0) * (probability_ratio >= 1 - grpo_epsilon_low) - ) - * grpo_loss_mask - ) - # d(probability_ratio)/d(logits) = -probability_ratio * (predicted_probabilities - target_probabilities). - predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) - grpo_grad = (probability_ratio_grad * probability_ratio).unsqueeze( - -1 - ) * predicted_probabilities.scatter_add( - -1, - grpo_target_masked.unsqueeze(-1), - -(grpo_loss_mask if grpo_target_mask is None else grpo_target_mask).unsqueeze(-1).to(torch.float32), - ) - grad = grpo_grad if grad is None else grad + grpo_grad - - return cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad +from fast_llm.engine.base_model.config import LossDef +from fast_llm.engine.distributed.config import DistributedConfig +from fast_llm.functional.entropy_loss import softmax_base +from fast_llm.layers.language_model.loss.config import MonolithicLossConfig +from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.utils import safe_merge_dicts @torch.compile def _monolithic_core( + children: tuple["LanguageModelLoss", ...], logits: torch.Tensor, # (*batch, vocab) group: ProcessGroup | None, logits_scale_factor: float, grad_logits: torch.Tensor | None, - ce_target: torch.Tensor | None, - ce_grad_output: float | None, - ce_divisor: float, - z_loss_enabled: bool, - z_loss_mask: torch.Tensor | None, - z_loss_grad_output: float | None, - z_loss_divisor: float, - distribution_target: torch.Tensor | None, - distribution_grad_output: float | None, - distribution_divisor: float, - distribution_loss_mask: torch.Tensor | None, - distribution_target_format: TargetFormat, - distribution_entropy_loss_type: EntropyLossType, - distribution_temperature: float, - grpo_target: torch.Tensor | None, - grpo_advantages: torch.Tensor | None, - grpo_old_log_probabilities: torch.Tensor | None, - grpo_grad_output: float | None, - grpo_divisor: float, - grpo_epsilon_low: float, - grpo_epsilon_high: float, - grpo_num_labels_in_seq: torch.Tensor | None, - grpo_compute_metrics: bool, - grpo_compute_entropy: bool, -) -> tuple[ - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - "GRPOMetrics | None", - torch.Tensor | None, -]: + arguments: tuple[tuple, ...], +) -> tuple[list, torch.Tensor | None]: """ - The single-pass monolithic head-loss kernel: one shared softmax over the logits, then every enabled - loss's scalar and gradient contribution, fused in one `@torch.compile` boundary (it inlines the plain - per-loss core). Disabled losses gate on their `None`/static flag and are dead-code-eliminated. - Gradients accumulate in fp32 and are cast to `logits.dtype` once at the end. GSPO is *not* handled - here — its eager segment seam requires the three-phase split (`_monolithic_gspo_forward_core`). + One shared softmax over the logits, then each child loss's `combinable_core` consuming it. The child + list is fixed per config, so the loop unrolls inside this single `@torch.compile` boundary and each + `combinable_core` dispatches (and inlines) to its loss type's math — every enabled loss is fused over + one softmax. Gradient contributions accumulate in fp32 and cast to `logits.dtype` once at the end. """ logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, logits_scale_factor, group) - cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad = ( - _apply_combinable_losses( - logits_norm, - exp_logits, - sum_exp_logits, - logits_max, - group, - logits_scale_factor, - ce_target, - ce_grad_output, - ce_divisor, - z_loss_enabled, - z_loss_mask, - z_loss_grad_output, - z_loss_divisor, - distribution_target, - distribution_grad_output, - distribution_divisor, - distribution_loss_mask, - distribution_target_format, - distribution_entropy_loss_type, - distribution_temperature, - grpo_target, - grpo_advantages, - grpo_old_log_probabilities, - grpo_grad_output, - grpo_divisor, - grpo_epsilon_low, - grpo_epsilon_high, - grpo_num_labels_in_seq, - grpo_compute_metrics, - grpo_compute_entropy, + grad = None + results = [] + for child, child_arguments in zip(children, arguments): + loss, child_grad, extra = child.combinable_core( + logits_norm, exp_logits, sum_exp_logits, logits_max, group, logits_scale_factor, child_arguments ) - ) - + results.append((loss, extra)) + if child_grad is not None: + grad = child_grad if grad is None else grad + child_grad if grad is not None: grad = grad.to(logits.dtype) if grad_logits is None: grad_logits = grad else: grad_logits.add_(grad) - - return cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad_logits + return results, grad_logits -@torch.compile -def _monolithic_gspo_forward_core( - logits: torch.Tensor, # (*batch, vocab) - group: ProcessGroup | None, - logits_scale_factor: float, - ce_target: torch.Tensor | None, - ce_grad_output: float | None, - ce_divisor: float, - z_loss_enabled: bool, - z_loss_mask: torch.Tensor | None, - z_loss_grad_output: float | None, - z_loss_divisor: float, - distribution_target: torch.Tensor | None, - distribution_grad_output: float | None, - distribution_divisor: float, - distribution_loss_mask: torch.Tensor | None, - distribution_target_format: TargetFormat, - distribution_entropy_loss_type: EntropyLossType, - distribution_temperature: float, - gspo_target: torch.Tensor, -) -> tuple[ - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor | None, -]: - """ - GSPO-present forward: one shared softmax, every enabled non-GSPO loss's scalar + fp32 gradient - (uncast), and GSPO's per-token `new_log_probs`. The softmax tensors and the uncast gradient cross the - eager segment seam to `gspo_backward_core` (the `index_add_`/`num_segments` seam can't be compiled). - """ - logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, logits_scale_factor, group) - cross_entropy_loss, z_loss, distribution_loss, _, _, _, grad = _apply_combinable_losses( - logits_norm, - exp_logits, - sum_exp_logits, - logits_max, - group, - logits_scale_factor, - ce_target, - ce_grad_output, - ce_divisor, - z_loss_enabled, - z_loss_mask, - z_loss_grad_output, - z_loss_divisor, - distribution_target, - distribution_grad_output, - distribution_divisor, - distribution_loss_mask, - distribution_target_format, - distribution_entropy_loss_type, - distribution_temperature, - None, - None, - None, - None, - 1.0, - 0.2, - 0.2, - None, - False, - False, - ) - gspo_loss_mask = gspo_target >= 0 - predicted_logits, target_masked, target_mask = predicted_logits_from_labels( - logits_norm, gspo_target, gspo_loss_mask, group - ) - new_log_probs = predicted_logits - sum_exp_logits.log() - return ( - cross_entropy_loss, - z_loss, - distribution_loss, - grad, - new_log_probs, - exp_logits, - sum_exp_logits, - target_masked, - target_mask, - ) - - -def gspo_segment_seam( - new_log_probs: torch.Tensor, # (*batch,) - loss_mask: torch.Tensor, # (*batch,) bool - advantages: torch.Tensor, # (*batch,) - old_log_probabilities: torch.Tensor, # (*batch,) - document_index_zero_based: torch.Tensor, # (*batch,) int - num_segments: int, - num_labels_in_seq: torch.Tensor, # (*batch,) - divisor: float, - grad_output: float | None, - sdp_group: ProcessGroup | None, - sp_group: ProcessGroup | None, - epsilon_low: float, - epsilon_high: float, - logits_scale_factor: float, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """Eager segment seam between the compiled forward and backward. The `index_add_` segment - aggregation and the symbolic `num_segments` live here so they never enter a compiled boundary - (no per-`num_segments` recompiles). Returns the loss, the `new_logprobs` metric, and the per-token - backward coefficient `effective_grad = grad_output_scaled · clip_factor · loss_weight · R_s` - (None when no gradient is requested).""" - log_ratio = new_log_probs - old_log_probabilities - - flat_document_index = document_index_zero_based.reshape(-1).long() - flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype) - # Per-token weight: mask / per-document label count, from the preprocessor. - # Each labeled token contributes `1 / N_d` so all of doc d's tokens sum to 1 (across - # SDP/SP ranks too), regardless of how the doc is sharded. - mean_token_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1) - # Pre-divide the per-token contributions by the per-doc label count, then sum per segment. - # All tokens in a segment share the same N_d, so this is mathematically equivalent to - # `log_ratio_sum / N_d` but avoids any per-segment denominator extraction. - mean_log_ratio_per_segment = log_ratio.new_zeros(num_segments).index_add_( - 0, flat_document_index, log_ratio.reshape(-1) * mean_token_weight - ) - # Accumulate in `log_ratio.dtype` (fp32). Casting the product back to `advantages.dtype` - # before summing would round each token's contribution to a possibly-low input dtype. - mean_advantage_per_segment = log_ratio.new_zeros(num_segments).index_add_( - 0, flat_document_index, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight - ) - for reduce_group in (sdp_group, sp_group): - if reduce_group is not None: - torch.distributed.all_reduce( - mean_log_ratio_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group - ) - torch.distributed.all_reduce( - mean_advantage_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group - ) - - segment_ratio = mean_log_ratio_per_segment.exp() # (num_segments,) — geometric-mean IS ratio - segment_advantage = mean_advantage_per_segment.detach() # (num_segments,) — no grad through A - - probability_ratio = segment_ratio[flat_document_index].reshape(log_ratio.shape) - advantage_per_token = segment_advantage[flat_document_index].reshape(log_ratio.shape) - loss_weight = loss_mask.to(log_ratio.dtype) - - losses = -torch.min( - probability_ratio * advantage_per_token, - torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantage_per_token, - ) - loss = (losses * loss_weight).sum() / divisor - - new_logprobs_mean = (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() - - if grad_output is None: - return loss, new_logprobs_mean, None - - grad_output_scaled = grad_output / divisor * logits_scale_factor - probability_ratio_grad = ( - grad_output_scaled - * ( - torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high) - + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low) - ) - * loss_weight - ) - effective_grad = probability_ratio_grad * probability_ratio - return loss, new_logprobs_mean, effective_grad - - -@torch.compile -def gspo_backward_core( - exp_logits: torch.Tensor, # (*batch, vocab) - sum_exp_logits: torch.Tensor, # (*batch,) - target_masked: torch.Tensor, # (*batch,) - target_mask: torch.Tensor | None, # (*batch,) or None (no TP) - loss_mask: torch.Tensor, # (*batch,) bool - effective_grad: torch.Tensor, # (*batch,) — per-token backward coefficient from the seam - logits_dtype: torch.dtype, - grad_logits: torch.Tensor | None, - grad_accumulator: torch.Tensor | None = None, # fp32 grad of co-resident losses to combine before the cast -) -> torch.Tensor: - """GSPO compiled backward: the per-token coefficient times the softmax chain rule, fused into one - kernel. Any `grad_accumulator` (other losses' fp32 grad from the shared forward) is added before the - single cast. `sum_exp_logits.unsqueeze` is out-of-place (the standalone eager kernel mutates it).""" - predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) - grad = effective_grad.unsqueeze(-1) * predicted_probabilities.scatter_add( - -1, - target_masked.unsqueeze(-1), - -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32), - ) - if grad_accumulator is not None: - grad = grad + grad_accumulator - grad = grad.to(logits_dtype) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) - return grad_logits - - -def monolithic_head_loss_forward_backward( - logits: torch.Tensor, - specs: list[MonolithicLossSpec], - *, - group: ProcessGroup | None, - grad_logits: torch.Tensor | None = None, -) -> tuple[list[MonolithicLossOutput], torch.Tensor | None]: +class MonolithicLoss[ConfigType: MonolithicLossConfig](LanguageModelLoss[ConfigType]): """ - Marshal the per-loss specs into the flat arguments of the kernel, run it once over a single shared - softmax, and map its outputs back to one `MonolithicLossOutput` per spec (in input order). When a - GSPO spec is present the kernel splits into the three-phase path (compiled forward → eager segment - seam → compiled backward); otherwise it is the single-pass `_monolithic_core`. All specs must share - one effective `logits_scale_factor` (validated at config time and asserted here). + A composite loss that runs the vocabulary softmax once and shares it across its combinable child + losses (cross-entropy, z-loss, distillation, GRPO), emitting each child's scalar / metrics and the + combined logits gradient in a single `@torch.compile` boundary. It is an ordinary head loss: the head + loops over it like any other and threads the same gradient buffer, so non-combinable losses (e.g. DPO) + are plain siblings in the head's loss list. """ - logits_scale_factor = specs[0].logits_scale_factor - specs_by_kind: dict[str, MonolithicLossSpec] = {} - for spec in specs: - Assert.eq(spec.logits_scale_factor, logits_scale_factor) - if spec.kind not in ("cross_entropy", "z_loss", "entropy_from_distribution", "grpo", "gspo"): - raise NotImplementedError(spec.kind) - assert spec.kind not in specs_by_kind, spec.kind - specs_by_kind[spec.kind] = spec - - cross_entropy_spec = specs_by_kind.get("cross_entropy") - z_loss_spec = specs_by_kind.get("z_loss") - distribution_spec = specs_by_kind.get("entropy_from_distribution") - grpo_spec = specs_by_kind.get("grpo") - gspo_spec = specs_by_kind.get("gspo") - # GRPO and GSPO are mutually exclusive RL objectives (both consume the same labels/advantages). - assert grpo_spec is None or gspo_spec is None - - # Shared (non-GSPO) loss arguments. - ce_target = None if cross_entropy_spec is None else cross_entropy_spec.target - ce_grad_output = None if cross_entropy_spec is None else cross_entropy_spec.grad_output - ce_divisor = 1.0 if cross_entropy_spec is None else cross_entropy_spec.divisor - z_loss_mask = None if z_loss_spec is None else z_loss_spec.loss_mask - z_loss_grad_output = None if z_loss_spec is None else z_loss_spec.grad_output - z_loss_divisor = 1.0 if z_loss_spec is None else z_loss_spec.divisor - distribution_target = None if distribution_spec is None else distribution_spec.target - distribution_grad_output = None if distribution_spec is None else distribution_spec.grad_output - distribution_divisor = 1.0 if distribution_spec is None else distribution_spec.divisor - distribution_loss_mask = None if distribution_spec is None else distribution_spec.loss_mask - distribution_target_format = TargetFormat.logits if distribution_spec is None else distribution_spec.target_format - distribution_entropy_loss_type = ( - EntropyLossType.cross_entropy if distribution_spec is None else distribution_spec.entropy_loss_type - ) - distribution_temperature = 1.0 if distribution_spec is None else distribution_spec.temperature - grpo_new_logprobs_mean = None - grpo_metrics = None - gspo_loss = None - gspo_new_logprobs_mean = None - - if gspo_spec is None: - cross_entropy_loss, z_loss, distribution_loss, grpo_loss, grpo_new_logprobs_mean, grpo_metrics, grad_logits = ( - _monolithic_core( - logits, - group, - logits_scale_factor, - grad_logits, - ce_target=ce_target, - ce_grad_output=ce_grad_output, - ce_divisor=ce_divisor, - z_loss_enabled=z_loss_spec is not None, - z_loss_mask=z_loss_mask, - z_loss_grad_output=z_loss_grad_output, - z_loss_divisor=z_loss_divisor, - distribution_target=distribution_target, - distribution_grad_output=distribution_grad_output, - distribution_divisor=distribution_divisor, - distribution_loss_mask=distribution_loss_mask, - distribution_target_format=distribution_target_format, - distribution_entropy_loss_type=distribution_entropy_loss_type, - distribution_temperature=distribution_temperature, - grpo_target=None if grpo_spec is None else grpo_spec.target, - grpo_advantages=None if grpo_spec is None else grpo_spec.advantages, - grpo_old_log_probabilities=None if grpo_spec is None else grpo_spec.old_log_probabilities, - grpo_grad_output=None if grpo_spec is None else grpo_spec.grad_output, - grpo_divisor=1.0 if grpo_spec is None else grpo_spec.divisor, - grpo_epsilon_low=0.2 if grpo_spec is None else grpo_spec.epsilon_low, - grpo_epsilon_high=0.2 if grpo_spec is None else grpo_spec.epsilon_high, - grpo_num_labels_in_seq=None if grpo_spec is None else grpo_spec.num_labels_in_seq, - grpo_compute_metrics=False if grpo_spec is None else grpo_spec.compute_metrics, - grpo_compute_entropy=False if grpo_spec is None else grpo_spec.compute_entropy, - ) + def __init__( + self, + config: ConfigType, + distributed_config: DistributedConfig, + *, + name: str, + prediction_distance: int = 1, + prediction_heads: int = 1, + vocab_parallel: bool = False, + num_splits: int = 1, + logits_scale_factor: float = 1.0, + weight: float = 1.0, + register_loss: bool = False, + ): + super().__init__( + config, + distributed_config, + name=name, + prediction_distance=prediction_distance, + prediction_heads=prediction_heads, + vocab_parallel=vocab_parallel, + num_splits=num_splits, + logits_scale_factor=logits_scale_factor, + weight=weight, + register_loss=register_loss, ) - else: - grpo_loss = None - ( - cross_entropy_loss, - z_loss, - distribution_loss, - grad, - gspo_new_log_probs, - exp_logits, - sum_exp_logits, - target_masked, - target_mask, - ) = _monolithic_gspo_forward_core( - logits, - group, - logits_scale_factor, - ce_target, - ce_grad_output, - ce_divisor, - z_loss_spec is not None, - z_loss_mask, - z_loss_grad_output, - z_loss_divisor, - distribution_target, - distribution_grad_output, - distribution_divisor, - distribution_loss_mask, - distribution_target_format, - distribution_entropy_loss_type, - distribution_temperature, - gspo_spec.target, + # Register children as distinct losses, unless a single child equals the head total (logged anyway). + children_register = register_loss or len(config.losses) > 1 + self._children: typing.Sequence[LanguageModelLoss] = torch.nn.ModuleList( + [ + child_config.get_layer( + distributed_config, + name=child_name if prediction_distance == 1 else f"{child_name}_{prediction_distance}", + prediction_distance=prediction_distance, + prediction_heads=prediction_heads, + vocab_parallel=vocab_parallel, + num_splits=num_splits, + logits_scale_factor=logits_scale_factor, + weight=self._weight, + register_loss=children_register, + ) + for child_name, child_config in config.losses.items() + ] ) - gspo_loss_mask = gspo_spec.target >= 0 - gspo_loss, gspo_new_logprobs_mean, effective_grad = gspo_segment_seam( - gspo_new_log_probs, - gspo_loss_mask, - gspo_spec.advantages, - gspo_spec.old_log_probabilities, - gspo_spec.document_index, - gspo_spec.num_segments, - gspo_spec.num_labels_in_seq, - gspo_spec.divisor, - gspo_spec.grad_output, - gspo_spec.sdp_group, - gspo_spec.sp_group, - gspo_spec.epsilon_low, - gspo_spec.epsilon_high, - logits_scale_factor, + # The shared softmax serves one effective scale; the config validates the children agree on it. + self._softmax_scale_factor = self._children[0]._logits_scale_factor + + def forward_backward( + self, + logits: "torch.Tensor", + kwargs: dict[str, typing.Any], + losses: dict | None = None, + split_index: int = 0, + grad_logits: torch.Tensor | None = None, + ) -> "tuple[torch.Tensor | None, torch.Tensor | None]": + register = losses is not None + arguments = tuple(child.combinable_extract(kwargs, split_index, register) for child in self._children) + group = self._parallel_dim.group if self._vocab_parallel else None + results, grad_logits = _monolithic_core( + tuple(self._children), logits, group, self._softmax_scale_factor, grad_logits, arguments ) - if effective_grad is not None: - grad_logits = gspo_backward_core( - exp_logits, - sum_exp_logits, - target_masked, - target_mask, - gspo_loss_mask, - effective_grad, - logits.dtype, - grad_logits, - grad_accumulator=grad, - ) - elif grad is not None: - # No GSPO gradient (e.g. eval), but co-resident losses produced one — finalize it here. - grad = grad.to(logits.dtype) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) - loss_by_kind = { - "cross_entropy": cross_entropy_loss, - "z_loss": z_loss, - "entropy_from_distribution": distribution_loss, - "grpo": grpo_loss, - "gspo": gspo_loss, - } - return [ - MonolithicLossOutput( - loss=loss_by_kind[spec.kind], - new_logprobs_mean=( - grpo_new_logprobs_mean - if spec.kind == "grpo" - else gspo_new_logprobs_mean if spec.kind == "gspo" else None - ), - metrics=grpo_metrics if spec.kind == "grpo" else None, - ) - for spec in specs - ], grad_logits + total_loss = None + for child, (loss, extra) in zip(self._children, results, strict=True): + if child._do_register_loss: + child._register_loss(child.name, loss, losses) + child.register_combinable_extras(extra, kwargs, losses) + weighted = loss if child.weight == 1 else loss * child.weight + total_loss = weighted if total_loss is None else total_loss + weighted + return total_loss, grad_logits + + def _forward_backward( + self, + logits: "torch.Tensor", + kwargs: dict[str, typing.Any], + losses: dict | None = None, + split_index: int = 0, + grad_logits: torch.Tensor | None = None, + ) -> "tuple[torch.Tensor, torch.Tensor | None]": + # `MonolithicLoss` overrides the public `forward_backward` directly (it composes weighted child + # losses and registers each child), so this per-loss hook is never the entry point. + raise NotImplementedError + + def get_preprocessing_config(self) -> dict[str, typing.Any]: + return safe_merge_dicts(*(child.get_preprocessing_config() for child in self._children)) + + def get_loss_definitions(self) -> list[LossDef]: + return [loss_def for child in self._children for loss_def in child.get_loss_definitions()] diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index c25a78b08..38ad3b472 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -24,12 +24,6 @@ ) from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, grpo_metrics_core from fast_llm.layers.language_model.loss.loss import LanguageModelLoss -from fast_llm.layers.language_model.loss.monolithic import ( - MonolithicLossOutput, - MonolithicLossSpec, - gspo_backward_core, - gspo_segment_seam, -) from fast_llm.utils import Assert @@ -139,42 +133,68 @@ def _forward_backward( return loss, grad - def get_monolithic_spec( - self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None - ) -> MonolithicLossSpec | None: - # When nothing is logged this step, skip the metric-only outputs (`new_logprobs_mean` and the - # GRPO metric family), mirroring the per-loss path's gating in `_forward_backward`. - register_metrics = losses is not None - return MonolithicLossSpec( - kind="grpo", - name=self.name, - weight=self._weight, - logits_scale_factor=self._logits_scale_factor, - grad_output=self._get_grad_output(kwargs), - divisor=self._get_label_count(kwargs), - target=self._get_labels(kwargs, split_index), - advantages=self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), - old_log_probabilities=self._prepare_target( - kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index - ), - epsilon_low=self._config.epsilon_low, - epsilon_high=self._config.epsilon_high, - num_labels_in_seq=( - self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index) - if register_metrics - else None - ), - compute_metrics=register_metrics and self._config.metrics != GRPOMetricsLevel.none, - compute_entropy=register_metrics and self._config.metrics == GRPOMetricsLevel.with_entropy, + def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + # When nothing is logged this step, drop the metric-only outputs (`new_logprobs_mean` and the + # GRPO metric family) by passing `num_labels_in_seq=None` / `compute_metrics=False`. + return ( + self._get_labels(kwargs, split_index), + self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), + self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index), + self._get_grad_output(kwargs), + self._get_label_count(kwargs), + self._config.epsilon_low, + self._config.epsilon_high, + (self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index) if register else None), + register and self._config.metrics != GRPOMetricsLevel.none, + register and self._config.metrics == GRPOMetricsLevel.with_entropy, ) - def register_monolithic_outputs( - self, output: MonolithicLossOutput, kwargs: dict[str, typing.Any], losses: dict | None - ) -> None: - super().register_monolithic_outputs(output, kwargs, losses) - self._register_new_logprobs(output.new_logprobs_mean, kwargs, losses) - if output.metrics is not None and losses is not None: - self._register_grpo_metrics(output.metrics, kwargs, losses) + def combinable_core( + self, + 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[torch.Tensor, torch.Tensor | None, tuple]: + ( + target, + advantages, + old_log_probabilities, + grad_output, + divisor, + epsilon_low, + epsilon_high, + num_labels_in_seq, + compute_metrics, + compute_entropy, + ) = arguments + loss, grad, new_logprobs_mean, metrics = grpo_combinable( + logits_norm, + exp_logits, + sum_exp_logits, + group, + target, + advantages, + old_log_probabilities, + grad_output, + divisor, + epsilon_low, + epsilon_high, + logits_scale_factor, + num_labels_in_seq, + compute_metrics, + compute_entropy, + ) + return loss, grad, (new_logprobs_mean, metrics) + + def register_combinable_extras(self, extra: tuple, kwargs: dict[str, typing.Any], losses: dict | None) -> None: + new_logprobs_mean, metrics = extra + self._register_new_logprobs(new_logprobs_mean, kwargs, losses) + if metrics is not None and losses is not None: + self._register_grpo_metrics(metrics, kwargs, losses) def _register_extra_metrics( self, @@ -339,39 +359,6 @@ def _forward_backward( def get_preprocessing_config(self) -> dict[str, typing.Any]: return super().get_preprocessing_config() | {"return_document_index": True} - def get_monolithic_spec( - self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None - ) -> MonolithicLossSpec | None: - return MonolithicLossSpec( - kind="gspo", - name=self.name, - weight=self._weight, - logits_scale_factor=self._logits_scale_factor, - grad_output=self._get_grad_output(kwargs), - divisor=kwargs[LanguageModelKwargs.num_documents_in_batch], - target=self._get_labels(kwargs, split_index), - advantages=self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), - old_log_probabilities=self._prepare_target( - kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index - ), - epsilon_low=self._config.epsilon_low, - epsilon_high=self._config.epsilon_high, - num_labels_in_seq=self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index), - document_index=self._prepare_target( - kwargs[BlockKwargs.global_document_index_q], split_index, split_by_distance=False - ).long() - - 1, - num_segments=kwargs[BlockKwargs.num_documents_in_sequence], - sdp_group=self._sequence_data_dim.group if self._sequence_data_active else None, - sp_group=self._parallel_dim.group if self._sequence_parallel else None, - ) - - def register_monolithic_outputs( - self, output: MonolithicLossOutput, kwargs: dict[str, typing.Any], losses: dict | None - ) -> None: - super().register_monolithic_outputs(output, kwargs, losses) - self._register_new_logprobs(output.new_logprobs_mean, kwargs, losses) - @torch.compile def compute_grpo_metrics( @@ -406,32 +393,31 @@ def compute_grpo_metrics( ) -@torch.compile -def fused_grpo_loss_forward_backward( - logits: torch.Tensor, # (*batch, vocab) +def grpo_combinable( + logits_norm: torch.Tensor, # (*batch, vocab) + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + group: torch.distributed.ProcessGroup | None, target: torch.Tensor, # (*batch,) advantages: torch.Tensor, # (*batch,) old_log_probabilities: torch.Tensor, # (*batch,) - grad_logits: torch.Tensor | None = None, - grad_output: float | None = None, - group: torch.distributed.ProcessGroup | None = None, - epsilon_low: float = 0.2, - epsilon_high: float = 0.2, - logits_scale_factor: float = 1.0, - num_labels_in_seq: ( - torch.Tensor | None - ) = None, # (*batch,) — response-span length broadcast per token, 0 for non-response - divisor: float | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: - if divisor is None: - divisor = logits.shape[:-1].numel() - grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor + grad_output: float | None, + divisor: float, + epsilon_low: float, + epsilon_high: float, + logits_scale_factor: float, + num_labels_in_seq: torch.Tensor | None, # enables `new_logprobs_mean` when not None + compute_metrics: bool, + compute_entropy: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, "GRPOMetrics | None"]: + """ + The post-softmax GRPO block shared by `fused_grpo_loss_forward_backward` (after its softmax) and the + monolithic head loss (over the shared softmax). Returns the loss scalar, the uncast masked gradient, + the optional `new_logprobs_mean`, and the optional GRPO metric family — all from one softmax and one + predicted-logit lookup. The caller casts the gradient to the logits dtype. + """ loss_mask = target >= 0 - - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) - predicted_logits, target_masked, target_mask = fused_predicted_logits_from_labels( - logits_norm, target, loss_mask, group - ) + 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() probability_ratio = (new_log_probs - old_log_probabilities).exp() @@ -452,7 +438,29 @@ def fused_grpo_loss_forward_backward( None if num_labels_in_seq is None else (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() ) - if grad_output is not None: + metrics = ( + grpo_metrics_core( + logits_norm, + exp_logits, + sum_exp_logits, + new_log_probs, + advantages, + old_log_probabilities, + loss_mask, + num_labels_in_seq, + epsilon_low, + epsilon_high, + group, + compute_entropy, + ) + if compute_metrics + else None + ) + + if grad_output is None: + grad = None + else: + grad_output = grad_output / divisor * logits_scale_factor # loss[a>=0] = -a * min(x, 1 + epsilon_high) => grad[a>=0] = -a * (x <= 1 + epsilon_high) # loss[a<=0] = a * max(x, 1 - epsilon_low) => grad[a<=0] = a * (x >= 1 - epsilon_low) probability_ratio_grad = ( @@ -463,17 +471,58 @@ def fused_grpo_loss_forward_backward( ) * loss_mask ) - # d(probability_ratio)/d(logits) = - probability_ratio * (predicted_probabilities - target_probabilities) - # (Sign absorbed in probability_ratio_grad) - predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze_(-1) + # (Sign absorbed in probability_ratio_grad). Out-of-place `unsqueeze` since the shared softmax tensors + # may be reused by sibling losses in the monolithic path. + predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) grad = (probability_ratio_grad * probability_ratio).unsqueeze(-1) * predicted_probabilities.scatter_add( -1, target_masked.unsqueeze(-1), -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32), ) - grad = grad.to(logits.dtype) + return loss, grad, new_logprobs_mean, metrics + + +@torch.compile +def fused_grpo_loss_forward_backward( + logits: torch.Tensor, # (*batch, vocab) + target: torch.Tensor, # (*batch,) + advantages: torch.Tensor, # (*batch,) + old_log_probabilities: torch.Tensor, # (*batch,) + grad_logits: torch.Tensor | None = None, + grad_output: float | None = None, + group: torch.distributed.ProcessGroup | None = None, + epsilon_low: float = 0.2, + epsilon_high: float = 0.2, + logits_scale_factor: float = 1.0, + num_labels_in_seq: ( + torch.Tensor | None + ) = None, # (*batch,) — response-span length broadcast per token, 0 for non-response + divisor: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + if divisor is None: + divisor = logits.shape[:-1].numel() + logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) + loss, grad, new_logprobs_mean, _ = grpo_combinable( + logits_norm, + exp_logits, + sum_exp_logits, + group, + target, + advantages, + old_log_probabilities, + grad_output, + divisor, + epsilon_low, + epsilon_high, + logits_scale_factor, + num_labels_in_seq, + False, + False, + ) + if grad is not None: + grad = grad.to(logits.dtype) if grad_logits is None: grad_logits = grad else: @@ -498,6 +547,113 @@ def _gspo_forward_core( return new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask +def gspo_segment_seam( + new_log_probs: torch.Tensor, # (*batch,) + loss_mask: torch.Tensor, # (*batch,) bool + advantages: torch.Tensor, # (*batch,) + old_log_probabilities: torch.Tensor, # (*batch,) + document_index_zero_based: torch.Tensor, # (*batch,) int + num_segments: int, + num_labels_in_seq: torch.Tensor, # (*batch,) + divisor: float, + grad_output: float | None, + sdp_group: torch.distributed.ProcessGroup | None, + sp_group: torch.distributed.ProcessGroup | None, + epsilon_low: float, + epsilon_high: float, + logits_scale_factor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Eager segment seam between the compiled forward and backward. The `index_add_` segment + aggregation and the symbolic `num_segments` live here so they never enter a compiled boundary + (no per-`num_segments` recompiles). Returns the loss, the `new_logprobs` metric, and the per-token + backward coefficient `effective_grad = grad_output_scaled · clip_factor · loss_weight · R_s` + (None when no gradient is requested).""" + log_ratio = new_log_probs - old_log_probabilities + + flat_document_index = document_index_zero_based.reshape(-1).long() + flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype) + # Per-token weight: mask / per-document label count, from the preprocessor. + # Each labeled token contributes `1 / N_d` so all of doc d's tokens sum to 1 (across + # SDP/SP ranks too), regardless of how the doc is sharded. + mean_token_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1) + # Pre-divide the per-token contributions by the per-doc label count, then sum per segment. + # All tokens in a segment share the same N_d, so this is mathematically equivalent to + # `log_ratio_sum / N_d` but avoids any per-segment denominator extraction. + mean_log_ratio_per_segment = log_ratio.new_zeros(num_segments).index_add_( + 0, flat_document_index, log_ratio.reshape(-1) * mean_token_weight + ) + # Accumulate in `log_ratio.dtype` (fp32). Casting the product back to `advantages.dtype` + # before summing would round each token's contribution to a possibly-low input dtype. + mean_advantage_per_segment = log_ratio.new_zeros(num_segments).index_add_( + 0, flat_document_index, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight + ) + for reduce_group in (sdp_group, sp_group): + if reduce_group is not None: + torch.distributed.all_reduce( + mean_log_ratio_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group + ) + torch.distributed.all_reduce( + mean_advantage_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group + ) + + segment_ratio = mean_log_ratio_per_segment.exp() # (num_segments,) — geometric-mean IS ratio + segment_advantage = mean_advantage_per_segment.detach() # (num_segments,) — no grad through A + + probability_ratio = segment_ratio[flat_document_index].reshape(log_ratio.shape) + advantage_per_token = segment_advantage[flat_document_index].reshape(log_ratio.shape) + loss_weight = loss_mask.to(log_ratio.dtype) + + losses = -torch.min( + probability_ratio * advantage_per_token, + torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantage_per_token, + ) + loss = (losses * loss_weight).sum() / divisor + + new_logprobs_mean = (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() + + if grad_output is None: + return loss, new_logprobs_mean, None + + grad_output_scaled = grad_output / divisor * logits_scale_factor + probability_ratio_grad = ( + grad_output_scaled + * ( + torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high) + + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low) + ) + * loss_weight + ) + effective_grad = probability_ratio_grad * probability_ratio + return loss, new_logprobs_mean, effective_grad + + +@torch.compile +def gspo_backward_core( + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + target_masked: torch.Tensor, # (*batch,) + target_mask: torch.Tensor | None, # (*batch,) or None (no TP) + loss_mask: torch.Tensor, # (*batch,) bool + effective_grad: torch.Tensor, # (*batch,) — per-token backward coefficient from the seam + logits_dtype: torch.dtype, + grad_logits: torch.Tensor | None, +) -> torch.Tensor: + """GSPO compiled backward: the per-token coefficient times the softmax chain rule, fused into one + kernel. `sum_exp_logits.unsqueeze` is out-of-place (the standalone eager kernel mutates it).""" + predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) + grad = effective_grad.unsqueeze(-1) * predicted_probabilities.scatter_add( + -1, + target_masked.unsqueeze(-1), + -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32), + ) + grad = grad.to(logits_dtype) + if grad_logits is None: + grad_logits = grad + else: + grad_logits.add_(grad) + return grad_logits + + # Orchestrator only: the eager `index_add_` segment seam (with the Python-int `num_segments`) sits # between the compiled forward and backward cores, so it stays out of every compiled boundary. def fused_gspo_loss_forward_backward( diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py index ad84ddf87..89482ad71 100644 --- a/fast_llm/layers/language_model/loss/z_loss.py +++ b/fast_llm/layers/language_model/loss/z_loss.py @@ -8,7 +8,6 @@ from fast_llm.functional.utils import reduce_losses from fast_llm.layers.language_model.loss.config import LanguageModelZLossConfig from fast_llm.layers.language_model.loss.loss import LanguageModelLoss -from fast_llm.layers.language_model.loss.monolithic import MonolithicLossSpec class LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](LanguageModelLoss[ConfigType]): @@ -34,18 +33,24 @@ def _forward_backward( divisor=self._get_label_count(kwargs), ) - def get_monolithic_spec( - self, kwargs: dict[str, typing.Any], split_index: int = 0, losses: dict | None = None - ) -> MonolithicLossSpec | None: - return MonolithicLossSpec( - kind="z_loss", - name=self.name, - weight=self._weight, - logits_scale_factor=self._logits_scale_factor, - grad_output=self._get_grad_output(kwargs), - divisor=self._get_label_count(kwargs), - loss_mask=self._get_loss_mask(kwargs, split_index), + def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + return self._get_loss_mask(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs) + + def combinable_core( + self, + 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[torch.Tensor, torch.Tensor | None, None]: + loss_mask, grad_output, divisor = arguments + loss, grad = z_loss_combinable( + exp_logits, sum_exp_logits, logits_max, loss_mask, grad_output, divisor, logits_scale_factor ) + return loss, grad, None def get_preprocessing_config(self) -> dict[str, typing.Any]: return {"return_prediction_mask": True} @@ -65,6 +70,28 @@ def z_loss( return torch.mean(out) +def z_loss_combinable( + exp_logits: torch.Tensor, # (*batch, vocab) + sum_exp_logits: torch.Tensor, # (*batch,) + logits_max: torch.Tensor, # (*batch,) + loss_mask: torch.Tensor | None, # (*batch,) + grad_output: float | None, + divisor: float, + logits_scale_factor: float, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """ + The post-softmax z-loss block shared by `fused_z_loss_forward_backward` (after its softmax) and the + monolithic head loss (over the shared softmax). Returns the loss scalar and the uncast, masked gradient + contribution; the caller casts to the logits dtype (the monolithic path defers that to one final cast). + """ + grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor + loss_term, grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, grad_output) + loss = reduce_losses(loss_term, divisor, loss_mask) + if grad is not None and loss_mask is not None: + grad = grad * loss_mask.unsqueeze(-1) + return loss, grad + + @torch.compile def fused_z_loss_forward_backward( logits: torch.Tensor, @@ -81,15 +108,11 @@ def fused_z_loss_forward_backward( """ if divisor is None: divisor = logits.shape[:-1].numel() - grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor _, exp_logits, sum_exp_logits, logits_max = fused_softmax_base(logits, logits_scale_factor, group) - loss_term, grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, grad_output) - - loss = reduce_losses(loss_term, divisor, loss_mask) - + loss, grad = z_loss_combinable( + exp_logits, sum_exp_logits, logits_max, loss_mask, grad_output, divisor, logits_scale_factor + ) if grad is not None: - if loss_mask is not None: - grad = grad * loss_mask.unsqueeze(-1) grad = grad.to(logits.dtype) if grad_logits is None: grad_logits = grad diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 50ed6e2be..69386203e 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -62,8 +62,6 @@ def get_config(self) -> GPTModelConfig: "cross_entropy_splits": self.num_splits, "prediction_heads": self.prediction_heads, } - if self.loss_implementation != "per_loss": - head_config["loss_implementation"] = self.loss_implementation if self.final_logit_softcap is not None: head_config["final_logit_softcap"] = self.final_logit_softcap losses = {} @@ -91,6 +89,17 @@ def get_config(self) -> GPTModelConfig: losses["gspo_loss"] = {"type": "gspo"} if isinstance(self.gspo_loss, float): losses["gspo_loss"]["weight"] = self.gspo_loss + if self.loss_implementation == "fused" and losses: + # Wrap the combinable losses in a single `monolithic` loss that shares one softmax; keep the + # child keys so the registered metric names match the per-loss configuration. + combinable = { + name: loss + for name, loss in losses.items() + if loss["type"] in ("label", "distillation", "z_loss", "grpo") + } + if combinable: + losses = {name: loss for name, loss in losses.items() if name not in combinable} + losses["monolithic"] = {"type": "monolithic", "losses": combinable} if losses: head_config["losses"] = losses @@ -396,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 ("fused") head-loss path. Cross-entropy, z-loss, and the from-distribution (distillation) losses -# run in the monolithic kernel; unsupported losses fall back to their own implementation and accumulate into -# the same gradient. +# 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). _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) @@ -424,27 +433,6 @@ 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 is a monolithic-kernel kind: the shared softmax forward feeds an eager segment seam (compiled -# forward → eager seam → compiled backward), so GSPO can share the softmax with a co-resident loss -# (e.g. z-loss). No splits (per-segment aggregation can't recombine across cross_entropy_splits chunks). -for loss_masking in (False, True): - _lm_head_test_configs.append( - LMHeadTestConfig( - f"fused_gspo_loss{'_masked' if loss_masking else ''}", - gspo_loss=True, - loss_masking=loss_masking, - loss_implementation="fused", - ) - ) - _lm_head_test_configs.append( - LMHeadTestConfig( - f"fused_gspo_and_z_loss{'_masked' if loss_masking else ''}", - 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"): diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 47949e916..1d68508a9 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -18,7 +18,6 @@ from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward from fast_llm.layers.language_model.loss.dpo import dpo_loss from fast_llm.layers.language_model.loss.loss import loss_forward_backward -from fast_llm.layers.language_model.loss.monolithic import MonolithicLossSpec, monolithic_head_loss_forward_backward from fast_llm.layers.language_model.loss.policy_gradient import ( GRPOMetrics, compute_grpo_metrics, @@ -587,520 +586,6 @@ def _test_z_loss( ) -def _monolithic_baseline( - kind: str, - local_logits, - target, - z_mask, - grad_output, - group, - logits_scale_factor, - divisor, - *, - teacher_target=None, - advantages=None, - old_log_probabilities=None, - num_labels_in_seq=None, -) -> tuple[torch.Tensor, torch.Tensor | None, MonolithicLossSpec]: - # Each enabled loss's baseline is its own `fused_*` kernel; the monolithic grad must equal their sum. - if kind == "cross_entropy": - loss, grad = fused_entropy_loss_forward_backward( - logits=local_logits, - target=target, - loss_mask=None, - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - target_format=TargetFormat.labels, - entropy_loss_type=EntropyLossType.cross_entropy, - divisor=divisor, - ) - spec = MonolithicLossSpec( - kind=kind, - name=kind, - weight=1.0, - logits_scale_factor=logits_scale_factor, - grad_output=grad_output, - divisor=divisor, - target=target, - ) - elif kind == "z_loss": - loss, grad = fused_z_loss_forward_backward( - logits=local_logits, - loss_mask=z_mask, - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - divisor=divisor, - ) - spec = MonolithicLossSpec( - kind=kind, - name=kind, - weight=1.0, - logits_scale_factor=logits_scale_factor, - grad_output=grad_output, - divisor=divisor, - loss_mask=z_mask, - ) - elif kind == "entropy_from_distribution": - loss, grad = fused_entropy_loss_forward_backward( - logits=local_logits, - target=teacher_target, - loss_mask=z_mask, - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - target_format=TargetFormat.logits, - entropy_loss_type=EntropyLossType.cross_entropy, - divisor=divisor, - ) - spec = MonolithicLossSpec( - kind=kind, - name=kind, - weight=1.0, - logits_scale_factor=logits_scale_factor, - grad_output=grad_output, - divisor=divisor, - target=teacher_target, - loss_mask=z_mask, - target_format=TargetFormat.logits, - entropy_loss_type=EntropyLossType.cross_entropy, - ) - elif kind == "grpo": - loss, grad, _ = fused_grpo_loss_forward_backward( - local_logits, - target, - advantages, - old_log_probabilities, - grad_logits=None, - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - num_labels_in_seq=num_labels_in_seq, - divisor=divisor, - ) - spec = MonolithicLossSpec( - kind=kind, - name=kind, - weight=1.0, - logits_scale_factor=logits_scale_factor, - grad_output=grad_output, - divisor=divisor, - target=target, - advantages=advantages, - old_log_probabilities=old_log_probabilities, - num_labels_in_seq=num_labels_in_seq, - ) - else: - raise NotImplementedError(kind) - return loss, grad, spec - - -def _test_monolithic_loss( - batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, kinds, group=None -): - # The monolithic kernel must reproduce the sum of the per-loss baselines over a single shared softmax. - logits, target, loss_mask = _get_lm_loss_inputs(num_columns, loss_masking, TargetFormat.labels, batch_shape, dtype) - local_logits = split_op(logits, group, -1).contiguous() - # The head feeds the label count as divisor; mirror that instead of the kernel's default (numel). - divisor = max(int((target >= 0).sum().item()), 1) - # z-loss takes an explicit mask (the head's `loss_mask` kwarg); align it with the labeled tokens. - z_mask = (target >= 0) if loss_masking else None - - # The teacher distribution is over the full vocab, so it splits along the vocab axis like the logits. - teacher_target = ( - split_op(torch.randn_like(logits), group, -1).contiguous() if "entropy_from_distribution" in kinds else None - ) - if "grpo" in kinds: - advantages = torch.randn_like(target, dtype=torch.float32) - # Correlate the old log-probs with the policy so the clipping branches are exercised. - old_log_probabilities = ( - torch.nn.functional.log_softmax(logits, dim=-1) - .gather(-1, (target * (target >= 0)).unsqueeze(-1)) - .squeeze(-1) - + torch.randn_like(target, dtype=torch.float32) / 2 - ) - num_labels_in_seq = torch.where( - target >= 0, - torch.full(batch_shape, divisor, dtype=torch.int32, device=target.device), - torch.zeros(batch_shape, dtype=torch.int32, device=target.device), - ) - else: - advantages = old_log_probabilities = num_labels_in_seq = None - - ref_losses, specs, ref_grad = [], [], None - for kind in kinds: - ref_loss, grad, spec = _monolithic_baseline( - kind, - local_logits, - target, - z_mask, - grad_output, - group, - logits_scale_factor, - divisor, - teacher_target=teacher_target, - advantages=advantages, - old_log_probabilities=old_log_probabilities, - num_labels_in_seq=num_labels_in_seq, - ) - ref_losses.append(ref_loss) - specs.append(spec) - if grad is not None: - ref_grad = grad if ref_grad is None else ref_grad + grad - - if accumulate and grad_output is not None: - previous_grad = torch.randn_like(local_logits) - ref_grad = ref_grad + previous_grad - grad_logits = previous_grad.clone() - else: - grad_logits = None - outputs, grad_mono = monolithic_head_loss_forward_backward( - local_logits, specs, group=group, grad_logits=grad_logits - ) - threshold = 1e-5 if dtype == DataType.float32 else 1e-4 - for output, ref_loss in zip(outputs, ref_losses, strict=True): - Assert.rms_close_relative(output.loss, ref_loss, threshold, 1e-6) - if grad_output is not None: - # The monolithic kernel sums all grad terms in fp32 and casts once, so it is *more* accurate than the - # per-loss baseline (which accumulates in `logits.dtype`); fp16 multi-loss grads differ at the ULP floor. - Assert.rms_close_relative(grad_mono, ref_grad, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-6) - else: - assert grad_mono is None - - -_MONOLITHIC_KINDS = ( - ("cross_entropy",), - ("z_loss",), - ("cross_entropy", "z_loss"), - ("cross_entropy", "z_loss", "entropy_from_distribution"), - ("grpo", "z_loss"), -) - - -@pytest.mark.slow -@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) -@pytest.mark.parametrize( - ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), - _LOSS_PARAMETERS, -) -@pytest.mark.parametrize("kinds", _MONOLITHIC_KINDS) -def test_monolithic_loss( - kinds, batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate -): - _test_monolithic_loss( - batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, kinds - ) - - -def _test_monolithic_distillation_loss( - batch_shape, - num_columns, - grad_output, - logits_scale_factor, - loss_masking, - dtype, - accumulate, - target_format, - entropy_loss_type, - temperature, - group=None, -): - # The monolithic from-distribution branch (teacher target) must reproduce `fused_entropy_loss_forward_backward`. - logits, target, loss_mask = _get_lm_loss_inputs(num_columns, loss_masking, target_format, batch_shape, dtype) - local_logits = split_op(logits, group, -1).contiguous() - local_target = split_op(target, group, -1).contiguous() - divisor = local_logits.shape[:-1].numel() - out_ref, grad_ref = fused_entropy_loss_forward_backward( - logits=local_logits, - target=local_target, - loss_mask=loss_mask, - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - temperature=temperature, - target_format=target_format, - entropy_loss_type=entropy_loss_type, - divisor=divisor, - ) - spec = MonolithicLossSpec( - kind="entropy_from_distribution", - name="distillation", - weight=1.0, - logits_scale_factor=logits_scale_factor, - grad_output=grad_output, - divisor=divisor, - target=local_target, - loss_mask=loss_mask, - target_format=target_format, - entropy_loss_type=entropy_loss_type, - temperature=temperature, - ) - if accumulate and grad_output is not None: - previous_grad = torch.randn_like(local_logits) - grad_ref = grad_ref + previous_grad - grad_logits = previous_grad.clone() - else: - grad_logits = None - outputs, grad_mono = monolithic_head_loss_forward_backward( - local_logits, [spec], group=group, grad_logits=grad_logits - ) - threshold = 1e-5 if dtype == DataType.float32 else 1e-4 - Assert.rms_close_relative(outputs[0].loss, out_ref, threshold, 1e-6) - if grad_output is not None: - Assert.rms_close_relative(grad_mono, grad_ref, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-6) - else: - assert grad_mono is None - - -def _test_monolithic_grpo_loss( - batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, group=None -): - # The monolithic GRPO branch must reproduce `fused_grpo_loss_forward_backward` (loss, grad, new_logprobs). - logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs( - num_columns, loss_masking, batch_shape, dtype - ) - local_logits = split_op(logits, group, -1).contiguous() - num_labels = int((target >= 0).sum().item()) - num_labels_in_seq = torch.where( - target >= 0, - torch.full(batch_shape, num_labels, dtype=torch.int32, device=target.device), - torch.zeros(batch_shape, dtype=torch.int32, device=target.device), - ) - divisor = max(num_labels, 1) - previous_grad = torch.randn_like(local_logits) if accumulate and grad_output is not None else None - out_ref, grad_ref, new_logprobs_ref = fused_grpo_loss_forward_backward( - local_logits, - target, - advantages, - old_log_probabilities, - grad_logits=None if previous_grad is None else previous_grad.clone(), - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - num_labels_in_seq=num_labels_in_seq, - divisor=divisor, - ) - spec = MonolithicLossSpec( - kind="grpo", - name="grpo", - weight=1.0, - logits_scale_factor=logits_scale_factor, - grad_output=grad_output, - divisor=divisor, - target=target, - advantages=advantages, - old_log_probabilities=old_log_probabilities, - num_labels_in_seq=num_labels_in_seq, - ) - outputs, grad_mono = monolithic_head_loss_forward_backward( - local_logits, [spec], group=group, grad_logits=None if previous_grad is None else previous_grad.clone() - ) - threshold = 1e-5 if dtype == DataType.float32 else 1e-4 - Assert.rms_close_relative(outputs[0].loss, out_ref, threshold, 1e-6) - Assert.rms_close_relative(outputs[0].new_logprobs_mean, new_logprobs_ref, threshold, 1e-6) - if grad_output is not None: - Assert.rms_close_relative(grad_mono, grad_ref, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-6) - else: - assert grad_mono is None - - -def _test_monolithic_grpo_metrics( - batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy, group=None -): - # The monolithic GRPO metric family (emitted from the shared softmax) must match the loop reference. - logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs( - num_columns, loss_masking, batch_shape, dtype - ) - # Different denominators per position so the per-token-mean broadcasting is exercised. - label_counts = (torch.arange(target.numel(), device=target.device).reshape(target.shape) % 5 + 1).to( - torch.int32 - ) * (target >= 0) - num_labels = max(int((target >= 0).sum().item()), 1) - - ref = reference_grpo_metrics( - logits, - target, - advantages, - old_log_probabilities, - label_counts, - epsilon_low=0.2, - epsilon_high=0.2, - logits_scale_factor=logits_scale_factor, - compute_entropy=compute_entropy, - ) - spec = MonolithicLossSpec( - kind="grpo", - name="grpo", - weight=1.0, - logits_scale_factor=logits_scale_factor, - grad_output=None, - divisor=num_labels, - target=target, - advantages=advantages, - old_log_probabilities=old_log_probabilities, - num_labels_in_seq=label_counts, - compute_metrics=True, - compute_entropy=compute_entropy, - ) - outputs, _ = monolithic_head_loss_forward_backward( - split_op(logits, group, -1).contiguous(), [spec], group=group, grad_logits=None - ) - _check_grpo_metrics(ref, outputs[0].metrics, threshold=5e-5 if dtype == DataType.float32 else 1e-4) - - -def _test_monolithic_gspo_loss( - batch_shape, - num_columns, - grad_output, - logits_scale_factor, - loss_masking, - dtype, - num_segments, - accumulate, - group=None, -): - # The monolithic GSPO branch (three-phase: shared softmax forward -> eager seam -> compiled backward) - # must reproduce the standalone `fused_gspo_loss_forward_backward` (loss, grad, new_logprobs). - logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs( - num_columns, loss_masking, batch_shape, dtype - ) - local_logits = split_op(logits, group, -1).contiguous() - seq_len = batch_shape[-1] if len(batch_shape) > 1 else batch_shape[0] - span = max(seq_len // num_segments, 1) - base = torch.arange(seq_len, device=target.device) // span - document_index = base.clamp(max=num_segments - 1).expand(batch_shape).contiguous() - flat_doc = document_index.reshape(-1).long() - labels_per_document = torch.zeros(num_segments, dtype=torch.int32, device=target.device).scatter_add( - 0, flat_doc, (target.reshape(-1) >= 0).to(torch.int32) - ) - num_labels_in_seq = labels_per_document[flat_doc].reshape(target.shape) - previous_grad = torch.randn_like(local_logits) if accumulate and grad_output is not None else None - out_ref, grad_ref, new_logprobs_ref = fused_gspo_loss_forward_backward( - local_logits, - target, - advantages, - old_log_probabilities, - document_index, - num_segments, - divisor=num_segments, - num_labels_in_seq=num_labels_in_seq, - grad_logits=None if previous_grad is None else previous_grad.clone(), - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - ) - spec = MonolithicLossSpec( - kind="gspo", - name="gspo", - weight=1.0, - logits_scale_factor=logits_scale_factor, - grad_output=grad_output, - divisor=num_segments, - target=target, - advantages=advantages, - old_log_probabilities=old_log_probabilities, - num_labels_in_seq=num_labels_in_seq, - document_index=document_index, - num_segments=num_segments, - ) - outputs, grad_mono = monolithic_head_loss_forward_backward( - local_logits, [spec], group=group, grad_logits=None if previous_grad is None else previous_grad.clone() - ) - threshold = 1e-5 if dtype == DataType.float32 else 1e-4 - # The monolithic GSPO forward is a distinct compiled graph from the standalone, so fp32 softmax - # reduction-order differences surface in the geometric-mean ratio — amplified in the loss scalar at - # higher logit scales. The gradient (the real correctness signal) stays tight. - loss_threshold = 5e-5 if dtype == DataType.float32 else 1e-4 - Assert.rms_close_relative(outputs[0].loss, out_ref, loss_threshold, 1e-6) - Assert.rms_close_relative(outputs[0].new_logprobs_mean, new_logprobs_ref, loss_threshold, 1e-6) - if grad_output is not None: - Assert.rms_close_relative(grad_mono, grad_ref, threshold, 1e-8 if grad_mono.dtype == torch.float32 else 1e-6) - else: - assert grad_mono is None - - -_MONOLITHIC_DISTRIBUTION_CASES = ( - # (target_format, entropy_loss_type, temperature) - (TargetFormat.logits, EntropyLossType.cross_entropy, 1.0), - (TargetFormat.logits, EntropyLossType.forward_kl, 1.0), - (TargetFormat.logits, EntropyLossType.reverse_kl, 1.0), - (TargetFormat.logits, EntropyLossType.cross_entropy, 2.0), # Temperature - (TargetFormat.probabilities, EntropyLossType.cross_entropy, 1.0), - (TargetFormat.probabilities, EntropyLossType.forward_kl, 1.0), - (TargetFormat.probabilities, EntropyLossType.reverse_kl, 1.0), -) - - -@pytest.mark.slow -@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) -@pytest.mark.parametrize( - ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), - _LOSS_PARAMETERS, -) -@pytest.mark.parametrize(("target_format", "entropy_loss_type", "temperature"), _MONOLITHIC_DISTRIBUTION_CASES) -def test_monolithic_distillation_loss( - target_format, - entropy_loss_type, - temperature, - batch_shape, - num_columns, - grad_output, - logits_scale_factor, - loss_masking, - dtype, - block_size, - accumulate, -): - _test_monolithic_distillation_loss( - batch_shape, - num_columns, - grad_output, - logits_scale_factor, - loss_masking, - dtype, - accumulate, - target_format, - entropy_loss_type, - temperature, - ) - - -@pytest.mark.slow -@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) -@pytest.mark.parametrize( - ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), - _LOSS_PARAMETERS, -) -def test_monolithic_grpo_loss( - batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate -): - _test_monolithic_grpo_loss( - batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate - ) - - -@pytest.mark.slow -@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) -@pytest.mark.parametrize( - ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), - _LOSS_PARAMETERS, -) -@pytest.mark.parametrize("compute_entropy", (False, True)) -def test_monolithic_grpo_metrics( - batch_shape, - num_columns, - grad_output, - logits_scale_factor, - loss_masking, - dtype, - block_size, - accumulate, - compute_entropy, -): - _test_monolithic_grpo_metrics(batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy) - - @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -1192,20 +677,6 @@ def test_gspo_loss( ) -@pytest.mark.slow -@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) -@pytest.mark.parametrize( - ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "num_segments", "accumulate"), - _GSPO_PARAMETERS, -) -def test_monolithic_gspo_loss( - batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, num_segments, accumulate -): - _test_monolithic_gspo_loss( - batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, num_segments, accumulate - ) - - @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -1319,21 +790,6 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa accumulate, test_context.group, ) - # GSPO through the monolithic kernel (shared softmax forward + eager seam + compiled backward) - with test_context.subtest(base_path, f"monolithic_gspo-{suffix}", 2) as subtest: - if subtest.do_run: - torch.manual_seed((seed + hash(subtest.name)) % 2**32) - _test_monolithic_gspo_loss( - batch_shape, - num_columns, - grad_output, - logits_scale_factor, - loss_masking, - dtype, - 4, # num_segments - accumulate, - test_context.group, - ) # GRPO metrics for compute_entropy in (False, True): with test_context.subtest(base_path, f"grpo_metrics-{compute_entropy}-{suffix}", 2) as subtest: @@ -1348,72 +804,6 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa compute_entropy, test_context.group, ) - # Monolithic kernel - for kinds in _MONOLITHIC_KINDS: - with test_context.subtest(base_path, f"monolithic-{"_".join(kinds)}-{suffix}", 2) as subtest: - if subtest.do_run: - torch.manual_seed((seed + hash(subtest.name)) % 2**32) - _test_monolithic_loss( - batch_shape, - num_columns, - grad_output, - logits_scale_factor, - loss_masking, - dtype, - accumulate, - kinds, - test_context.group, - ) - # Monolithic GRPO objective - with test_context.subtest(base_path, f"monolithic_grpo-{suffix}", 2) as subtest: - if subtest.do_run: - torch.manual_seed((seed + hash(subtest.name)) % 2**32) - _test_monolithic_grpo_loss( - batch_shape, - num_columns, - grad_output, - logits_scale_factor, - loss_masking, - dtype, - accumulate, - test_context.group, - ) - # Monolithic GRPO metrics - for compute_entropy in (False, True): - with test_context.subtest( - base_path, f"monolithic_grpo_metrics-{compute_entropy}-{suffix}", 2 - ) as subtest: - if subtest.do_run: - torch.manual_seed((seed + hash(subtest.name)) % 2**32) - _test_monolithic_grpo_metrics( - batch_shape, - num_columns, - logits_scale_factor, - loss_masking, - dtype, - compute_entropy, - test_context.group, - ) - # Monolithic from-distribution (distillation) losses - for target_format, entropy_loss_type, temperature in _MONOLITHIC_DISTRIBUTION_CASES: - with test_context.subtest( - base_path, f"monolithic_dist-{entropy_loss_type}-{target_format}-{temperature}-{suffix}", 2 - ) as subtest: - if subtest.do_run: - torch.manual_seed((seed + hash(subtest.name)) % 2**32) - _test_monolithic_distillation_loss( - batch_shape, - num_columns, - grad_output, - logits_scale_factor, - loss_masking, - dtype, - accumulate, - target_format, - entropy_loss_type, - temperature, - test_context.group, - ) @pytest.mark.slow @@ -1457,15 +847,6 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): "gspo", "grpo_metrics-False", "grpo_metrics-True", - "monolithic_grpo", - "monolithic_gspo", - "monolithic_grpo_metrics-False", - "monolithic_grpo_metrics-True", - *(f"monolithic-{"_".join(kinds)}" for kinds in _MONOLITHIC_KINDS), - *( - f"monolithic_dist-{entropy_loss_type}-{target_format}-{temperature}" - for target_format, entropy_loss_type, temperature in _MONOLITHIC_DISTRIBUTION_CASES - ), ), ) def test_lm_loss_distributed( From 70250d7b6f6433458d04b602a92f5f2f343a63ee Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 2 Jul 2026 13:24:31 -0400 Subject: [PATCH 13/28] Fold combinable loss boilerplate into a CombinableLoss base (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kill the four `*_combinable` module-level indirection functions (`z_loss_combinable`, `cross_entropy_from_labels_combinable`, `distillation_combinable`, `grpo_combinable`): inline each into its loss's `combinable_core`, now a `@staticmethod` so both the monolithic loop and the standalone `fused_*` kernel call the one function. Add an intermediate `CombinableLoss` base owning the shared combinable machinery — `register_combinable_extras` (moved off the general `LanguageModelLoss` base) and a `_accumulate_grad` staticmethod that replaces the cast-and-accumulate tail duplicated across the z-loss / GRPO standalone kernels and the monolithic core. `fused_z_loss_forward_backward` and `fused_grpo_loss_forward_backward` now route through their `combinable_core`. Verified: staticmethod dispatch still fuses with no graph breaks (`fullgraph=True`); head + functional + distributed TP suites pass. Net -80 lines. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../language_model/loss/entropy_loss.py | 145 +++++------- fast_llm/layers/language_model/loss/loss.py | 36 ++- .../layers/language_model/loss/monolithic.py | 10 +- .../language_model/loss/policy_gradient.py | 221 ++++++++---------- fast_llm/layers/language_model/loss/z_loss.py | 60 ++--- 5 files changed, 196 insertions(+), 276 deletions(-) diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py index b8915dfa2..6de884fd9 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -10,14 +10,17 @@ reverse_kl_from_distribution_core, ) from fast_llm.functional.triton.entropy_loss import triton_entropy_loss_forward_backward +from fast_llm.functional.utils import reduce_losses from fast_llm.layers.language_model.loss.config import ( LanguageModelDistillationLossConfig, LanguageModelLabelEntropyLossConfig, ) -from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.loss import CombinableLoss, LanguageModelLoss -class LanguageModelLabelEntropyLoss[ConfigType: LanguageModelLabelEntropyLossConfig](LanguageModelLoss[ConfigType]): +class LanguageModelLabelEntropyLoss[ConfigType: LanguageModelLabelEntropyLossConfig]( + CombinableLoss, LanguageModelLoss[ConfigType] +): def _forward_backward( self, logits: "torch.Tensor", @@ -46,8 +49,8 @@ def _forward_backward( def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: return self._get_labels(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs) + @staticmethod def combinable_core( - self, logits_norm: torch.Tensor, exp_logits: torch.Tensor, sum_exp_logits: torch.Tensor, @@ -56,15 +59,24 @@ def combinable_core( logits_scale_factor: float, arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, None]: - # For labels, forward-KL is identical to cross-entropy (one-hot target entropy is zero). + """Post-softmax cross-entropy-from-labels over the shared softmax. Returns the loss scalar and the + uncast, masked gradient (the caller casts); no extra outputs. For labels, forward-KL is identical to + cross-entropy (one-hot target entropy is zero).""" target, grad_output, divisor = arguments - loss, grad = cross_entropy_from_labels_combinable( - logits_norm, exp_logits, sum_exp_logits, group, target, grad_output, divisor, logits_scale_factor + loss_mask = target >= 0 + grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor + per_sample_loss, grad = cross_entropy_from_labels_core( + logits_norm, exp_logits, sum_exp_logits, target, loss_mask, grad_output, group ) + loss = reduce_losses(per_sample_loss, divisor, loss_mask) + if grad is not None: + grad = grad * loss_mask.unsqueeze(-1) return loss, grad, None -class LanguageModelDistillationLoss[ConfigType: LanguageModelDistillationLossConfig](LanguageModelLoss[ConfigType]): +class LanguageModelDistillationLoss[ConfigType: LanguageModelDistillationLossConfig]( + CombinableLoss, LanguageModelLoss[ConfigType] +): def _forward_backward( self, logits: "torch.Tensor", @@ -101,8 +113,8 @@ def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, re self._config.temperature, ) + @staticmethod def combinable_core( - self, logits_norm: torch.Tensor, exp_logits: torch.Tensor, sum_exp_logits: torch.Tensor, @@ -111,93 +123,40 @@ def combinable_core( logits_scale_factor: float, arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, None]: + """Post-softmax distillation over the shared student softmax (cross-entropy / forward-KL / reverse-KL + from a teacher distribution, adding a teacher softmax at scale `logits_scale_factor / temperature`). + Returns the loss scalar and the uncast, masked gradient (the caller casts); no extra outputs.""" target, loss_mask, grad_output, divisor, entropy_loss_type, temperature = arguments - loss, grad = distillation_combinable( - logits_norm, - exp_logits, - sum_exp_logits, - group, - target, - loss_mask, - grad_output, - divisor, - logits_scale_factor, - entropy_loss_type, - temperature, - ) + grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor + if entropy_loss_type == EntropyLossType.reverse_kl: + per_sample_loss, grad = reverse_kl_from_distribution_core( + logits_norm, + exp_logits, + sum_exp_logits, + target, + grad_output, + logits_scale_factor, + TargetFormat.logits, + group, + temperature, + ) + else: + per_sample_loss, grad = cross_entropy_from_distribution_core( + logits_norm, + exp_logits, + sum_exp_logits, + target, + grad_output, + logits_scale_factor, + TargetFormat.logits, + group, + temperature, + return_kl_loss=entropy_loss_type == EntropyLossType.forward_kl, + ) + loss = reduce_losses(per_sample_loss, divisor, loss_mask) + if grad is not None and loss_mask is not None: + grad = grad * loss_mask.unsqueeze(-1) return loss, grad, None def get_preprocessing_config(self) -> dict[str, typing.Any]: return {"return_prediction_mask": True} - - -def cross_entropy_from_labels_combinable( - logits_norm: torch.Tensor, # (*batch, vocab) - exp_logits: torch.Tensor, # (*batch, vocab) - sum_exp_logits: torch.Tensor, # (*batch,) - group: "torch.distributed.ProcessGroup | None", - target: torch.Tensor, # (*batch,) - grad_output: float | None, - divisor: float, - logits_scale_factor: float, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Post-softmax cross-entropy-from-labels block over the shared softmax: loss scalar + uncast, masked - gradient. The caller casts to the logits dtype (the monolithic path defers that to one final cast).""" - loss_mask = target >= 0 - grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor - per_sample_loss, grad = cross_entropy_from_labels_core( - logits_norm, exp_logits, sum_exp_logits, target, loss_mask, grad_output, group - ) - loss = (per_sample_loss * loss_mask).sum() / divisor - if grad is not None: - grad = grad * loss_mask.unsqueeze(-1) - return loss, grad - - -def distillation_combinable( - logits_norm: torch.Tensor, # (*batch, vocab) - exp_logits: torch.Tensor, # (*batch, vocab) - sum_exp_logits: torch.Tensor, # (*batch,) - group: "torch.distributed.ProcessGroup | None", - target: torch.Tensor, # (*batch, vocab) teacher logits - loss_mask: torch.Tensor | None, # (*batch,) - grad_output: float | None, - divisor: float, - logits_scale_factor: float, - entropy_loss_type: EntropyLossType, - temperature: float, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Post-softmax distillation block (cross-entropy / forward-KL / reverse-KL from a teacher distribution) - over the shared student softmax: loss scalar + uncast, masked gradient (the caller casts).""" - grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor - if entropy_loss_type == EntropyLossType.reverse_kl: - per_sample_loss, grad = reverse_kl_from_distribution_core( - logits_norm, - exp_logits, - sum_exp_logits, - target, - grad_output, - logits_scale_factor, - TargetFormat.logits, - group, - temperature, - ) - else: - per_sample_loss, grad = cross_entropy_from_distribution_core( - logits_norm, - exp_logits, - sum_exp_logits, - target, - grad_output, - logits_scale_factor, - TargetFormat.logits, - group, - temperature, - return_kl_loss=entropy_loss_type == EntropyLossType.forward_kl, - ) - if loss_mask is not None: - per_sample_loss = per_sample_loss * loss_mask - loss = per_sample_loss.sum() / divisor - if grad is not None and loss_mask is not None: - grad = grad * loss_mask.unsqueeze(-1) - return loss, grad diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 5ddffef32..742fa3ec8 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -72,12 +72,6 @@ def _forward_backward( def get_loss_definitions(self) -> list[LossDef]: return [LossDef(name=self.name)] if self._do_register_loss else [] - def register_combinable_extras( - self, extra: typing.Any, kwargs: dict[str, typing.Any], losses: dict | None - ) -> None: - """Register any per-loss outputs beyond the scalar (e.g. GRPO's `new_logprobs` / metric family) - produced by `combinable_core` when this loss runs inside `MonolithicLoss`. No-op by default.""" - def get_preprocessing_config( self, ) -> dict[str, typing.Any]: @@ -148,6 +142,36 @@ def _get_reference_model_logits(self, reference_model: str, kwargs: dict[str, ty ) +class CombinableLoss: + """Mixin for losses that consume the vocabulary softmax: each runs standalone through its own fused + kernel, or several are fused together by `MonolithicLoss` over a single shared softmax. Subclasses + implement `combinable_extract` (eager kwargs -> argument tuple, run outside the compiled boundary) and + the `combinable_core` static method (the post-softmax math over the shared softmax tensors, returning + `(loss, uncast_grad, extra)`), and override `register_combinable_extras` when they emit outputs beyond + the loss scalar. Both the standalone kernel and `MonolithicLoss` call the same `combinable_core`.""" + + @staticmethod + def _accumulate_grad( + grad: torch.Tensor | None, logits_dtype: torch.dtype, grad_logits: torch.Tensor | None + ) -> torch.Tensor | None: + """Cast a computed logits gradient to the logits dtype and accumulate it into `grad_logits` (in place + when a buffer exists, otherwise adopting it). The standalone kernels cast their single gradient here; + the monolithic kernel casts the fp32 sum of its children's gradients once.""" + if grad is None: + return grad_logits + grad = grad.to(logits_dtype) + if grad_logits is None: + return grad + grad_logits.add_(grad) + return grad_logits + + def register_combinable_extras( + self, extra: typing.Any, kwargs: dict[str, typing.Any], losses: dict | None + ) -> None: + """Register per-loss outputs beyond the scalar (e.g. GRPO's `new_logprobs` / metric family) produced + by `combinable_core`. No-op by default.""" + + def loss_forward_backward( grad_output: float | None, fn: typing.Callable, input_: torch.Tensor, *args, **kwargs ) -> tuple[torch.Tensor, torch.Tensor | None]: diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index c4f57dd8e..35815a8bc 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -7,7 +7,7 @@ from fast_llm.engine.distributed.config import DistributedConfig from fast_llm.functional.entropy_loss import softmax_base from fast_llm.layers.language_model.loss.config import MonolithicLossConfig -from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.loss import CombinableLoss, LanguageModelLoss from fast_llm.utils import safe_merge_dicts @@ -36,13 +36,7 @@ def _monolithic_core( results.append((loss, extra)) if child_grad is not None: grad = child_grad if grad is None else grad + child_grad - if grad is not None: - grad = grad.to(logits.dtype) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) - return results, grad_logits + return results, CombinableLoss._accumulate_grad(grad, logits.dtype, grad_logits) class MonolithicLoss[ConfigType: MonolithicLossConfig](LanguageModelLoss[ConfigType]): diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 38ad3b472..03dfa8ce6 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -23,7 +23,7 @@ LanguageModelPolicyGradientLossConfig, ) from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, grpo_metrics_core -from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.loss import CombinableLoss, LanguageModelLoss from fast_llm.utils import Assert @@ -57,7 +57,9 @@ def _logprob_metric_name(self) -> str: return f"{self._name}_new_logprobs" -class LanguageModelGRPOLoss[ConfigType: LanguageModelGRPOLossConfig](LanguageModelPolicyGradientLoss[ConfigType]): +class LanguageModelGRPOLoss[ConfigType: LanguageModelGRPOLossConfig]( + CombinableLoss, LanguageModelPolicyGradientLoss[ConfigType] +): """GRPO: per-token IS-ratio clipping.""" def __init__( @@ -149,8 +151,8 @@ def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, re register and self._config.metrics == GRPOMetricsLevel.with_entropy, ) + @staticmethod def combinable_core( - self, logits_norm: torch.Tensor, exp_logits: torch.Tensor, sum_exp_logits: torch.Tensor, @@ -159,6 +161,10 @@ def combinable_core( logits_scale_factor: float, arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, tuple]: + """Post-softmax GRPO over the shared softmax, called by both `fused_grpo_loss_forward_backward` + (after its softmax) and the monolithic head loss. Returns the loss scalar, the uncast masked gradient + (the caller casts), and the `(new_logprobs_mean, metrics)` extras (each `None` when not requested) — + all from one softmax and one predicted-logit lookup.""" ( target, advantages, @@ -171,23 +177,73 @@ def combinable_core( compute_metrics, compute_entropy, ) = arguments - loss, grad, new_logprobs_mean, metrics = grpo_combinable( - logits_norm, - exp_logits, - sum_exp_logits, - group, - target, - advantages, - old_log_probabilities, - grad_output, - divisor, - epsilon_low, - epsilon_high, - logits_scale_factor, - num_labels_in_seq, - compute_metrics, - compute_entropy, + 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() + probability_ratio = (new_log_probs - old_log_probabilities).exp() + + losses = -torch.min( + probability_ratio * advantages, + torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantages, + ) + loss = reduce_losses(losses, divisor, loss_mask) + + # Sum of per-sequence mean log-probs, matching pipelinerl's new_logprobs metric: + # sum_sum(new_logprobs / num_labels_in_seq, masks_shifted, segments) + # Dividing by num_labels_in_seq (span length broadcast per token) and summing over masked + # tokens gives mean logprob per sequence; summing those across sequences matches the deepspeed + # convention exactly (segments are redundant once num_labels_in_seq is correct). + # Clamp to avoid 0/0=nan when num_labels_in_seq=0 (padded tokens or fully masked documents) + # — those positions also have loss_mask=0 so they correctly contribute 0 to the sum. + new_logprobs_mean = ( + None if num_labels_in_seq is None else (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() ) + + metrics = ( + grpo_metrics_core( + logits_norm, + exp_logits, + sum_exp_logits, + new_log_probs, + advantages, + old_log_probabilities, + loss_mask, + num_labels_in_seq, + epsilon_low, + epsilon_high, + group, + compute_entropy, + ) + if compute_metrics + else None + ) + + if grad_output is None: + grad = None + else: + grad_output = grad_output / divisor * logits_scale_factor + # loss[a>=0] = -a * min(x, 1 + epsilon_high) => grad[a>=0] = -a * (x <= 1 + epsilon_high) + # loss[a<=0] = a * max(x, 1 - epsilon_low) => grad[a<=0] = a * (x >= 1 - epsilon_low) + probability_ratio_grad = ( + grad_output + * ( + torch.clamp_min(advantages, 0) * (probability_ratio <= 1 + epsilon_high) + + torch.clamp_max(advantages, 0) * (probability_ratio >= 1 - epsilon_low) + ) + * loss_mask + ) + # d(probability_ratio)/d(logits) = - probability_ratio * (predicted_probabilities - target_probabilities) + # (Sign absorbed in probability_ratio_grad). Out-of-place `unsqueeze` since the shared softmax tensors + # may be reused by sibling losses in the monolithic path. + predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) + grad = (probability_ratio_grad * probability_ratio).unsqueeze(-1) * predicted_probabilities.scatter_add( + -1, + target_masked.unsqueeze(-1), + -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32), + ) + return loss, grad, (new_logprobs_mean, metrics) def register_combinable_extras(self, extra: tuple, kwargs: dict[str, typing.Any], losses: dict | None) -> None: @@ -393,97 +449,6 @@ def compute_grpo_metrics( ) -def grpo_combinable( - logits_norm: torch.Tensor, # (*batch, vocab) - exp_logits: torch.Tensor, # (*batch, vocab) - sum_exp_logits: torch.Tensor, # (*batch,) - group: torch.distributed.ProcessGroup | None, - target: torch.Tensor, # (*batch,) - advantages: torch.Tensor, # (*batch,) - old_log_probabilities: torch.Tensor, # (*batch,) - grad_output: float | None, - divisor: float, - epsilon_low: float, - epsilon_high: float, - logits_scale_factor: float, - num_labels_in_seq: torch.Tensor | None, # enables `new_logprobs_mean` when not None - compute_metrics: bool, - compute_entropy: bool, -) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, "GRPOMetrics | None"]: - """ - The post-softmax GRPO block shared by `fused_grpo_loss_forward_backward` (after its softmax) and the - monolithic head loss (over the shared softmax). Returns the loss scalar, the uncast masked gradient, - the optional `new_logprobs_mean`, and the optional GRPO metric family — all from one softmax and one - predicted-logit lookup. The caller casts the gradient to the logits dtype. - """ - 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() - probability_ratio = (new_log_probs - old_log_probabilities).exp() - - losses = -torch.min( - probability_ratio * advantages, - torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantages, - ) - loss = reduce_losses(losses, divisor, loss_mask) - - # Sum of per-sequence mean log-probs, matching pipelinerl's new_logprobs metric: - # sum_sum(new_logprobs / num_labels_in_seq, masks_shifted, segments) - # Dividing by num_labels_in_seq (span length broadcast per token) and summing over masked - # tokens gives mean logprob per sequence; summing those across sequences matches the deepspeed - # convention exactly (segments are redundant once num_labels_in_seq is correct). - # Clamp to avoid 0/0=nan when num_labels_in_seq=0 (padded tokens or fully masked documents) - # — those positions also have loss_mask=0 so they correctly contribute 0 to the sum. - new_logprobs_mean = ( - None if num_labels_in_seq is None else (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() - ) - - metrics = ( - grpo_metrics_core( - logits_norm, - exp_logits, - sum_exp_logits, - new_log_probs, - advantages, - old_log_probabilities, - loss_mask, - num_labels_in_seq, - epsilon_low, - epsilon_high, - group, - compute_entropy, - ) - if compute_metrics - else None - ) - - if grad_output is None: - grad = None - else: - grad_output = grad_output / divisor * logits_scale_factor - # loss[a>=0] = -a * min(x, 1 + epsilon_high) => grad[a>=0] = -a * (x <= 1 + epsilon_high) - # loss[a<=0] = a * max(x, 1 - epsilon_low) => grad[a<=0] = a * (x >= 1 - epsilon_low) - probability_ratio_grad = ( - grad_output - * ( - torch.clamp_min(advantages, 0) * (probability_ratio <= 1 + epsilon_high) - + torch.clamp_max(advantages, 0) * (probability_ratio >= 1 - epsilon_low) - ) - * loss_mask - ) - # d(probability_ratio)/d(logits) = - probability_ratio * (predicted_probabilities - target_probabilities) - # (Sign absorbed in probability_ratio_grad). Out-of-place `unsqueeze` since the shared softmax tensors - # may be reused by sibling losses in the monolithic path. - predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) - grad = (probability_ratio_grad * probability_ratio).unsqueeze(-1) * predicted_probabilities.scatter_add( - -1, - target_masked.unsqueeze(-1), - -(loss_mask if target_mask is None else target_mask).unsqueeze(-1).to(torch.float32), - ) - - return loss, grad, new_logprobs_mean, metrics - - @torch.compile def fused_grpo_loss_forward_backward( logits: torch.Tensor, # (*batch, vocab) @@ -503,32 +468,28 @@ def fused_grpo_loss_forward_backward( ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: if divisor is None: divisor = logits.shape[:-1].numel() - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) - loss, grad, new_logprobs_mean, _ = grpo_combinable( + logits_norm, exp_logits, sum_exp_logits, logits_max = fused_softmax_base(logits, logits_scale_factor, group) + loss, grad, (new_logprobs_mean, _) = LanguageModelGRPOLoss.combinable_core( logits_norm, exp_logits, sum_exp_logits, + logits_max, group, - target, - advantages, - old_log_probabilities, - grad_output, - divisor, - epsilon_low, - epsilon_high, logits_scale_factor, - num_labels_in_seq, - False, - False, + ( + target, + advantages, + old_log_probabilities, + grad_output, + divisor, + epsilon_low, + epsilon_high, + num_labels_in_seq, + False, + False, + ), ) - if grad is not None: - grad = grad.to(logits.dtype) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) - - return loss, grad_logits, new_logprobs_mean + return loss, CombinableLoss._accumulate_grad(grad, logits.dtype, grad_logits), new_logprobs_mean @torch.compile diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py index 89482ad71..2e24f6271 100644 --- a/fast_llm/layers/language_model/loss/z_loss.py +++ b/fast_llm/layers/language_model/loss/z_loss.py @@ -7,10 +7,10 @@ from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward from fast_llm.functional.utils import reduce_losses from fast_llm.layers.language_model.loss.config import LanguageModelZLossConfig -from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.loss import CombinableLoss, LanguageModelLoss -class LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](LanguageModelLoss[ConfigType]): +class LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](CombinableLoss, LanguageModelLoss[ConfigType]): def _forward_backward( self, logits: "torch.Tensor", @@ -36,8 +36,8 @@ def _forward_backward( def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: return self._get_loss_mask(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs) + @staticmethod def combinable_core( - self, logits_norm: torch.Tensor, exp_logits: torch.Tensor, sum_exp_logits: torch.Tensor, @@ -46,10 +46,15 @@ def combinable_core( logits_scale_factor: float, arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, None]: + """Post-softmax z-loss over the shared softmax, called by both `fused_z_loss_forward_backward` + (after its softmax) and the monolithic head loss. Returns the loss scalar and the uncast, masked + gradient contribution (the caller casts); z-loss emits no extra outputs.""" loss_mask, grad_output, divisor = arguments - loss, grad = z_loss_combinable( - exp_logits, sum_exp_logits, logits_max, loss_mask, grad_output, divisor, logits_scale_factor - ) + grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor + loss_term, grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, grad_output) + loss = reduce_losses(loss_term, divisor, loss_mask) + if grad is not None and loss_mask is not None: + grad = grad * loss_mask.unsqueeze(-1) return loss, grad, None def get_preprocessing_config(self) -> dict[str, typing.Any]: @@ -70,28 +75,6 @@ def z_loss( return torch.mean(out) -def z_loss_combinable( - exp_logits: torch.Tensor, # (*batch, vocab) - sum_exp_logits: torch.Tensor, # (*batch,) - logits_max: torch.Tensor, # (*batch,) - loss_mask: torch.Tensor | None, # (*batch,) - grad_output: float | None, - divisor: float, - logits_scale_factor: float, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """ - The post-softmax z-loss block shared by `fused_z_loss_forward_backward` (after its softmax) and the - monolithic head loss (over the shared softmax). Returns the loss scalar and the uncast, masked gradient - contribution; the caller casts to the logits dtype (the monolithic path defers that to one final cast). - """ - grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor - loss_term, grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, grad_output) - loss = reduce_losses(loss_term, divisor, loss_mask) - if grad is not None and loss_mask is not None: - grad = grad * loss_mask.unsqueeze(-1) - return loss, grad - - @torch.compile def fused_z_loss_forward_backward( logits: torch.Tensor, @@ -108,15 +91,14 @@ def fused_z_loss_forward_backward( """ if divisor is None: divisor = logits.shape[:-1].numel() - _, exp_logits, sum_exp_logits, logits_max = fused_softmax_base(logits, logits_scale_factor, group) - loss, grad = z_loss_combinable( - exp_logits, sum_exp_logits, logits_max, loss_mask, grad_output, divisor, logits_scale_factor + logits_norm, exp_logits, sum_exp_logits, logits_max = fused_softmax_base(logits, logits_scale_factor, group) + loss, grad, _ = LanguageModelZLoss.combinable_core( + logits_norm, + exp_logits, + sum_exp_logits, + logits_max, + group, + logits_scale_factor, + (loss_mask, grad_output, divisor), ) - if grad is not None: - grad = grad.to(logits.dtype) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) - - return loss, grad_logits + return loss, CombinableLoss._accumulate_grad(grad, logits.dtype, grad_logits) From 86b68160c83414359e347c935e462c3db9febbb3 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 2 Jul 2026 14:04:17 -0400 Subject: [PATCH 14/28] Move standalone combinable forward-backward onto the loss object (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Delete the `fused_z_loss_forward_backward` / `fused_grpo_loss_forward_backward` module functions — their body was just `softmax -> combinable_core -> accumulate` (a one-child monolithic), kept module-level only so tests could import them. Add `CombinableLoss.combinable_forward_backward` (compiled) as that standalone path; z-loss / GRPO `_forward_backward` route their non-triton branch through it, and the tests now construct the loss object and exercise the method directly (the tensor-parallel `group` is passed per call, so a trivial single-rank object suffices — including in the distributed subtests). GRPO's non-triton path now gets its metric family from the shared softmax via `combinable_core` instead of a separate `compute_grpo_metrics` pass (numerically identical, one fewer softmax); the triton path keeps the separate pass since its kernel emits no metrics. Verified: `combinable_forward_backward` fuses with no graph breaks (`fullgraph`); head + functional + distributed TP suites pass (651 / 21 skipped). Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/entropy_loss.py | 4 +- fast_llm/layers/language_model/loss/loss.py | 32 ++++- .../language_model/loss/policy_gradient.py | 115 ++++++------------ fast_llm/layers/language_model/loss/z_loss.py | 65 +++------- tests/layers/test_lm_losses.py | 43 ++++--- 5 files changed, 110 insertions(+), 149 deletions(-) diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index 5f6c6a650..50753596b 100644 --- a/fast_llm/functional/entropy_loss.py +++ b/fast_llm/functional/entropy_loss.py @@ -360,8 +360,8 @@ def z_loss_core( Z-loss from the already-computed shared softmax tensors. Returns the unmasked per-sample loss term (`log_sum_exp ** 2`) and (when `grad_output` is given) the unmasked gradient; the caller applies the loss mask, reduction, and dtype cast. z-loss needs the un-regularized log-sum-exp, so it adds back - `logits_max` (cross-entropy cancels it). Shared between `fused_z_loss_forward_backward` and the - monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. + `logits_max` (cross-entropy cancels it). Inlined by the z-loss `combinable_core` inside its + `@torch.compile` boundary (standalone or monolithic). """ log_sum_exp_logits = sum_exp_logits.log() + logits_max loss_term = log_sum_exp_logits**2 diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 742fa3ec8..91e7bf614 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -7,6 +7,7 @@ from fast_llm.core.ops import split_op from fast_llm.engine.base_model.config import LossDef from fast_llm.engine.distributed.config import DistributedConfig, DistributedDimNames +from fast_llm.functional.entropy_loss import softmax_base from fast_llm.layers.language_model.config import LanguageModelKwargs from fast_llm.layers.language_model.loss.config import LanguageModelLossConfig, LanguageModelLossKwargs from fast_llm.utils import Assert @@ -143,12 +144,31 @@ def _get_reference_model_logits(self, reference_model: str, kwargs: dict[str, ty class CombinableLoss: - """Mixin for losses that consume the vocabulary softmax: each runs standalone through its own fused - kernel, or several are fused together by `MonolithicLoss` over a single shared softmax. Subclasses - implement `combinable_extract` (eager kwargs -> argument tuple, run outside the compiled boundary) and - the `combinable_core` static method (the post-softmax math over the shared softmax tensors, returning - `(loss, uncast_grad, extra)`), and override `register_combinable_extras` when they emit outputs beyond - the loss scalar. Both the standalone kernel and `MonolithicLoss` call the same `combinable_core`.""" + """Mixin for losses that consume the vocabulary softmax: each runs standalone through + `combinable_forward_backward`, or several are fused together by `MonolithicLoss` over a single shared + softmax. Subclasses implement `combinable_extract` (eager kwargs -> argument tuple, run outside the + compiled boundary) and the `combinable_core` static method (the post-softmax math over the shared + softmax tensors, returning `(loss, uncast_grad, extra)`), and override `register_combinable_extras` when + they emit outputs beyond the loss scalar. Both paths call the same `combinable_core`.""" + + _logits_scale_factor: float + + @torch.compile + def combinable_forward_backward( + self, + logits: torch.Tensor, + group: "torch.distributed.ProcessGroup | None", + grad_logits: torch.Tensor | None, + arguments: tuple, + ) -> tuple[torch.Tensor, torch.Tensor | None, typing.Any]: + """Standalone realization of a single combinable loss: run the softmax once, then this loss's + `combinable_core`, then cast-and-accumulate the gradient. Equivalent to a one-child `MonolithicLoss` + — they share `combinable_core`, so this is not a second copy of the math.""" + logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, self._logits_scale_factor, group) + loss, grad, extra = self.combinable_core( + logits_norm, exp_logits, sum_exp_logits, logits_max, group, self._logits_scale_factor, arguments + ) + return loss, self._accumulate_grad(grad, logits.dtype, grad_logits), extra @staticmethod def _accumulate_grad( diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 03dfa8ce6..558b62f63 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -102,38 +102,46 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: + arguments = self.combinable_extract(kwargs, split_index, losses is not None) + group = self._parallel_dim.group if self._vocab_parallel else None if TritonConfig.enabled(logits.device, self._config.use_triton): from fast_llm.functional.triton.grpo_loss import triton_grpo_loss_forward_backward - fn = triton_grpo_loss_forward_backward - else: - fn = fused_grpo_loss_forward_backward - loss, grad, new_logprobs_mean = fn( - logits, - self._get_labels(kwargs, split_index), - self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), - self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index), - grad_logits=grad_logits, - grad_output=self._get_grad_output(kwargs), - group=self._parallel_dim.group if self._vocab_parallel else None, - epsilon_low=self._config.epsilon_low, - epsilon_high=self._config.epsilon_high, - logits_scale_factor=self._logits_scale_factor, - num_labels_in_seq=( - None - if losses is None - else self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index) - ), - divisor=self._get_label_count(kwargs), - ) - - self._register_new_logprobs(new_logprobs_mean, kwargs, losses) - - # Skip the extra softmax pass when there is nothing to register. - if losses is not None and self._config.metrics != GRPOMetricsLevel.none: - self._register_extra_metrics(logits, kwargs, losses, split_index) + ( + target, + advantages, + old_log_probabilities, + grad_output, + divisor, + epsilon_low, + epsilon_high, + num_labels_in_seq, + compute_metrics, + _, + ) = arguments + loss, grad, new_logprobs_mean = triton_grpo_loss_forward_backward( + logits, + target, + advantages, + old_log_probabilities, + grad_logits=grad_logits, + grad_output=grad_output, + group=group, + epsilon_low=epsilon_low, + epsilon_high=epsilon_high, + logits_scale_factor=self._logits_scale_factor, + num_labels_in_seq=num_labels_in_seq, + divisor=divisor, + ) + self._register_new_logprobs(new_logprobs_mean, kwargs, losses) + # Triton produces only loss/grad/new_logprobs; the metric family needs its own pass here. + if compute_metrics: + self._register_extra_metrics(logits, kwargs, losses, split_index) + return loss, grad - return loss, grad + loss, grad_logits, extra = self.combinable_forward_backward(logits, group, grad_logits, arguments) + self.register_combinable_extras(extra, kwargs, losses) + return loss, grad_logits def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: # When nothing is logged this step, drop the metric-only outputs (`new_logprobs_mean` and the @@ -161,10 +169,10 @@ def combinable_core( logits_scale_factor: float, arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, tuple]: - """Post-softmax GRPO over the shared softmax, called by both `fused_grpo_loss_forward_backward` - (after its softmax) and the monolithic head loss. Returns the loss scalar, the uncast masked gradient - (the caller casts), and the `(new_logprobs_mean, metrics)` extras (each `None` when not requested) — - all from one softmax and one predicted-logit lookup.""" + """Post-softmax GRPO over the shared softmax, called by both the standalone + `combinable_forward_backward` and the monolithic head loss. Returns the loss scalar, the uncast masked + gradient (the caller casts), and the `(new_logprobs_mean, metrics)` extras (each `None` when not + requested) — all from one softmax and one predicted-logit lookup.""" ( target, advantages, @@ -449,49 +457,6 @@ def compute_grpo_metrics( ) -@torch.compile -def fused_grpo_loss_forward_backward( - logits: torch.Tensor, # (*batch, vocab) - target: torch.Tensor, # (*batch,) - advantages: torch.Tensor, # (*batch,) - old_log_probabilities: torch.Tensor, # (*batch,) - grad_logits: torch.Tensor | None = None, - grad_output: float | None = None, - group: torch.distributed.ProcessGroup | None = None, - epsilon_low: float = 0.2, - epsilon_high: float = 0.2, - logits_scale_factor: float = 1.0, - num_labels_in_seq: ( - torch.Tensor | None - ) = None, # (*batch,) — response-span length broadcast per token, 0 for non-response - divisor: float | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: - if divisor is None: - divisor = logits.shape[:-1].numel() - logits_norm, exp_logits, sum_exp_logits, logits_max = fused_softmax_base(logits, logits_scale_factor, group) - loss, grad, (new_logprobs_mean, _) = LanguageModelGRPOLoss.combinable_core( - logits_norm, - exp_logits, - sum_exp_logits, - logits_max, - group, - logits_scale_factor, - ( - target, - advantages, - old_log_probabilities, - grad_output, - divisor, - epsilon_low, - epsilon_high, - num_labels_in_seq, - False, - False, - ), - ) - return loss, CombinableLoss._accumulate_grad(grad, logits.dtype, grad_logits), new_logprobs_mean - - @torch.compile def _gspo_forward_core( logits: torch.Tensor, # (*batch, vocab) diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py index 2e24f6271..e904ffe16 100644 --- a/fast_llm/layers/language_model/loss/z_loss.py +++ b/fast_llm/layers/language_model/loss/z_loss.py @@ -3,7 +3,7 @@ import torch from fast_llm.functional.config import TritonConfig -from fast_llm.functional.entropy_loss import fused_softmax_base, z_loss_core +from fast_llm.functional.entropy_loss import z_loss_core from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward from fast_llm.functional.utils import reduce_losses from fast_llm.layers.language_model.loss.config import LanguageModelZLossConfig @@ -19,19 +19,21 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor, torch.Tensor | None]": - return ( - triton_z_loss_forward_backward - if TritonConfig.enabled(logits.device, self._config.use_triton) - else fused_z_loss_forward_backward - )( - logits, - self._get_loss_mask(kwargs, split_index), - grad_output=self._get_grad_output(kwargs), - group=self._parallel_dim.group if self._vocab_parallel else None, - logits_scale_factor=self._logits_scale_factor, - grad_logits=grad_logits, - divisor=self._get_label_count(kwargs), - ) + arguments = self.combinable_extract(kwargs, split_index, losses is not None) + group = self._parallel_dim.group if self._vocab_parallel else None + if TritonConfig.enabled(logits.device, self._config.use_triton): + loss_mask, grad_output, divisor = arguments + return triton_z_loss_forward_backward( + logits, + loss_mask, + grad_output=grad_output, + group=group, + logits_scale_factor=self._logits_scale_factor, + grad_logits=grad_logits, + divisor=divisor, + ) + loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments) + return loss, grad_logits def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: return self._get_loss_mask(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs) @@ -46,9 +48,9 @@ def combinable_core( logits_scale_factor: float, arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, None]: - """Post-softmax z-loss over the shared softmax, called by both `fused_z_loss_forward_backward` - (after its softmax) and the monolithic head loss. Returns the loss scalar and the uncast, masked - gradient contribution (the caller casts); z-loss emits no extra outputs.""" + """Post-softmax z-loss over the shared softmax, called by both the standalone + `combinable_forward_backward` and the monolithic head loss. Returns the loss scalar and the uncast, + masked gradient contribution (the caller casts); z-loss emits no extra outputs.""" loss_mask, grad_output, divisor = arguments grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor loss_term, grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, grad_output) @@ -73,32 +75,3 @@ def z_loss( if loss_mask is not None: out = out * loss_mask return torch.mean(out) - - -@torch.compile -def fused_z_loss_forward_backward( - logits: torch.Tensor, - loss_mask: torch.Tensor | None, - grad_logits: torch.Tensor | None = None, - grad_output: float | None = None, - group: torch.distributed.ProcessGroup | None = None, - logits_scale_factor: float = 1.0, - divisor: float | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """ - Z-loss = mean(logsumexp(logits, dim=-1) ** 2) - Grad = 2 * log_sum_exp_logits * softmax(logits) - """ - if divisor is None: - divisor = logits.shape[:-1].numel() - logits_norm, exp_logits, sum_exp_logits, logits_max = fused_softmax_base(logits, logits_scale_factor, group) - loss, grad, _ = LanguageModelZLoss.combinable_core( - logits_norm, - exp_logits, - sum_exp_logits, - logits_max, - group, - logits_scale_factor, - (loss_mask, grad_output, divisor), - ) - return loss, CombinableLoss._accumulate_grad(grad, logits.dtype, grad_logits) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 1d68508a9..bed045839 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -8,7 +8,7 @@ from fast_llm.core.ops import split_op from fast_llm.engine.config_utils import data_type from fast_llm.engine.config_utils.data_type import DataType -from fast_llm.engine.distributed.config import DistributedBackend +from fast_llm.engine.distributed.config import DistributedBackend, DistributedConfig from fast_llm.functional.config import EntropyLossType, TargetFormat from fast_llm.functional.entropy_loss import fused_entropy_loss_forward_backward, torch_entropy_loss_forward_backward from fast_llm.functional.triton import triton_available @@ -16,15 +16,15 @@ from fast_llm.functional.triton.grpo_loss import triton_grpo_loss_forward_backward from fast_llm.functional.triton.gspo_loss import triton_gspo_loss_forward_backward from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward +from fast_llm.layers.language_model.loss.config import LanguageModelGRPOLossConfig, LanguageModelZLossConfig from fast_llm.layers.language_model.loss.dpo import dpo_loss from fast_llm.layers.language_model.loss.loss import loss_forward_backward from fast_llm.layers.language_model.loss.policy_gradient import ( GRPOMetrics, compute_grpo_metrics, - fused_grpo_loss_forward_backward, fused_gspo_loss_forward_backward, ) -from fast_llm.layers.language_model.loss.z_loss import fused_z_loss_forward_backward, z_loss +from fast_llm.layers.language_model.loss.z_loss import z_loss from fast_llm.utils import Assert from tests.utils.dataset import get_random_spans from tests.utils.subtest import DistributedTestContext @@ -363,17 +363,12 @@ def _test_grpo_loss( previous_grad = torch.randn_like(grad_ref) grad_ref = grad_ref + previous_grad local_previous_grad = split_op(previous_grad, group, -1).contiguous() - out_fused, grad_fused, new_logprobs_fused = fused_grpo_loss_forward_backward( + loss = _combinable_loss(LanguageModelGRPOLossConfig(), "grpo", logits_scale_factor) + out_fused, grad_fused, (new_logprobs_fused, _) = loss.combinable_forward_backward( split_op(logits, group, -1), - target, - advantages, - old_log_probabilities, - grad_logits=local_previous_grad.clone() if accumulate else None, - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - num_labels_in_seq=num_labels_in_seq, - divisor=divisor, + group, + local_previous_grad.clone() if accumulate else None, + (target, advantages, old_log_probabilities, grad_output, divisor, 0.2, 0.2, num_labels_in_seq, False, False), ) _compare_losses_and_grads(out_fused, out_ref, grad_output is not None, grad_fused, grad_ref, group=group) @@ -531,6 +526,15 @@ def _test_grpo_metrics( _check_grpo_metrics(ref, got, threshold=5e-5 if dtype == DataType.float32 else 1e-4) +def _combinable_loss(config, name: str, logits_scale_factor: float): + # Build the loss object so its `combinable_forward_backward` method is exercised directly. The tensor- + # parallel `group` is passed per call, so a trivial single-rank distributed config suffices even for the + # distributed subtests. + distributed_config = DistributedConfig() + distributed_config.validate() + return config.get_layer(distributed_config, name=name, logits_scale_factor=logits_scale_factor) + + def _test_z_loss( batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate, group=None ): @@ -547,13 +551,12 @@ def _test_z_loss( previous_grad = torch.randn_like(grad_ref) grad_ref = grad_ref + previous_grad local_previous_grad = split_op(previous_grad, group, -1).contiguous() - out_fused, grad_fused = fused_z_loss_forward_backward( - logits=local_logits, - loss_mask=loss_mask, - grad_logits=local_previous_grad.clone() if accumulate else None, - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, + loss = _combinable_loss(LanguageModelZLossConfig(), "z_loss", logits_scale_factor) + out_fused, grad_fused, _ = loss.combinable_forward_backward( + local_logits, + group, + local_previous_grad.clone() if accumulate else None, + (loss_mask, grad_output, local_logits.shape[:-1].numel()), ) _compare_losses_and_grads( out_fused, From 1bc212c1029816074a5ccf1ec4c79b73515e357f Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 2 Jul 2026 14:28:22 -0400 Subject: [PATCH 15/28] Route entropy losses through combinable_forward_backward too (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Label-entropy and distillation now run their non-triton standalone path through `CombinableLoss.combinable_forward_backward` (softmax -> combinable_core -> accumulate), the same as z-loss / GRPO. This makes the general `fused_entropy_loss_forward_backward` kernel and its three compiled base-wrappers (`_fused_cross_entropy_base_from_{labels,distribution}`, `_fused_reverse_kl_base_from_distribution`) unused, so they are deleted — the shared `_core` functions remain the single source of truth. Numerically identical: both paths already went through the same cores. Tests build the label / distillation loss object and exercise the method (triton path still calls its kernel directly). The `probabilities` target format is dropped from the entropy suite — no loss config consumes it (distillation is logits-only), so it was general-kernel-only coverage of an unreachable path. Verified: all entropy variants fuse with no graph breaks (`fullgraph`); head + functional + distributed TP suites pass (531 / 21 skipped). Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/entropy_loss.py | 151 +----------------- .../language_model/loss/entropy_loss.py | 71 ++++---- tests/layers/test_lm_losses.py | 41 +++-- 3 files changed, 67 insertions(+), 196 deletions(-) diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index 50753596b..1290fbed8 100644 --- a/fast_llm/functional/entropy_loss.py +++ b/fast_llm/functional/entropy_loss.py @@ -2,7 +2,6 @@ from fast_llm.core.distributed import ProcessGroup, ReduceOp, all_reduce from fast_llm.functional.config import EntropyLossType, TargetFormat -from fast_llm.functional.utils import reduce_losses from fast_llm.utils import Assert @@ -140,8 +139,8 @@ def reverse_kl_from_distribution_core( """ Reverse-KL math from a precomputed student softmax (adding a teacher softmax when the target is logits). - This plain (un-compiled) core is shared between the public `_fused_reverse_kl_base_from_distribution` - wrapper and the monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. + Used by the distillation `combinable_core`, inlined inside its `@torch.compile` boundary (standalone + or monolithic). """ assert target_format in (TargetFormat.logits, TargetFormat.probabilities) predicted_log_probability = logits_norm - sum_exp_logits.log().unsqueeze(-1) @@ -172,30 +171,6 @@ def reverse_kl_from_distribution_core( return per_sample_loss, grad -@torch.compile -def _fused_reverse_kl_base_from_distribution( - logits: torch.Tensor, # (*batch, vocab) - target: torch.Tensor, # (*batch, vocab) - grad_output: float | None, - logits_scale_factor: float, - target_format: TargetFormat, - group: ProcessGroup | None = None, - temperature: float = 1.0, -) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) - logits_norm, exp_logits, sum_exp_logits, _ = softmax_base(logits, logits_scale_factor, group) - return reverse_kl_from_distribution_core( - logits_norm, - exp_logits, - sum_exp_logits, - target, - grad_output, - logits_scale_factor, - target_format, - group, - temperature, - ) - - def cross_entropy_from_distribution_core( logits_norm: torch.Tensor, # (*batch, vocab) exp_logits: torch.Tensor, # (*batch, vocab) @@ -212,8 +187,8 @@ def cross_entropy_from_distribution_core( Cross-entropy / forward-KL math from a precomputed student softmax (adding a teacher softmax when the target is logits). - This plain (un-compiled) core is shared between the public `_fused_cross_entropy_base_from_distribution` - wrapper and the monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. + Used by the distillation `combinable_core`, inlined inside its `@torch.compile` boundary (standalone + or monolithic). """ if target_format == TargetFormat.logits: target_logits_norm, exp_logits_targets, sum_exp_target_logits, _ = softmax_base( @@ -248,32 +223,6 @@ def cross_entropy_from_distribution_core( return per_sample_loss, grad -@torch.compile -def _fused_cross_entropy_base_from_distribution( - logits: torch.Tensor, # (*batch, vocab) - target: torch.Tensor, # (*batch, vocab) - grad_output: float | None, - logits_scale_factor: float, - target_format: TargetFormat, - group: ProcessGroup | None = None, - temperature: float = 1.0, - return_kl_loss: bool = False, -) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) - logits_norm, exp_logits, sum_exp_logits, _ = softmax_base(logits, logits_scale_factor, group) - return cross_entropy_from_distribution_core( - logits_norm, - exp_logits, - sum_exp_logits, - target, - grad_output, - logits_scale_factor, - target_format, - group, - temperature, - return_kl_loss, - ) - - def predicted_logits_from_labels( logits: torch.Tensor, # (*batch, vocab) target: torch.Tensor, # (*batch,) @@ -327,9 +276,8 @@ def cross_entropy_from_labels_core( """ Cross-entropy from labels, taking the already-computed shared softmax tensors. Returns the unmasked per-sample loss and (when `grad_output` is given) the unmasked gradient; the caller applies the loss - mask, reduction, and dtype cast. This plain (un-compiled) core is shared between the public - `_fused_cross_entropy_base_from_labels` wrapper and the monolithic head-loss kernel, which inlines it - inside its own `@torch.compile` boundary. + mask, reduction, and dtype cast. Used by the label / GRPO `combinable_core`, inlined inside its + `@torch.compile` boundary (standalone or monolithic). """ predicted_logits, target_masked, target_mask = predicted_logits_from_labels(logits_norm, target, loss_mask, group) @@ -370,90 +318,3 @@ def z_loss_core( else: grad = (2 * grad_output * (log_sum_exp_logits / sum_exp_logits)).unsqueeze(-1) * exp_logits return loss_term, grad - - -@torch.compile -def _fused_cross_entropy_base_from_labels( - logits: torch.Tensor, # (*batch, vocab) - target: torch.Tensor, # (*batch,) - loss_mask: torch.Tensor, # (*batch,) - grad_output: float | None, - logits_scale_factor: float, - group: ProcessGroup | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) - return cross_entropy_from_labels_core( - logits_norm, exp_logits, sum_exp_logits, target, loss_mask, grad_output, group - ) - - -@torch.compile -def fused_entropy_loss_forward_backward( - logits: torch.Tensor, # (*batch, vocab) - target: torch.Tensor, # (*batch,) or (*batch, vocab) - loss_mask: torch.Tensor | None, # (*batch,) - grad_logits: torch.Tensor | None = None, - grad_output: float | None = None, - group: torch.distributed.ProcessGroup | None = None, - logits_scale_factor: float = 1.0, - temperature: float = 1.0, - target_format: TargetFormat = TargetFormat.labels, - entropy_loss_type: EntropyLossType = EntropyLossType.cross_entropy, - divisor: float | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """ - A fused implementation of cross-entropy with torch compile. - It is an improvement over the pytorch implementation because of the fused casting, both in speed and memory, - but still suboptimal because it needs multiple kernels. - """ - if divisor is None: - divisor = logits.shape[:-1].numel() - grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor - if target_format == TargetFormat.labels: - assert entropy_loss_type in (EntropyLossType.cross_entropy, EntropyLossType.forward_kl) - assert loss_mask is None - loss_mask = target >= 0 - losses, grad = _fused_cross_entropy_base_from_labels( - logits, - target, - loss_mask, - grad_output, - logits_scale_factor, - group, - ) - elif entropy_loss_type in (EntropyLossType.cross_entropy, EntropyLossType.forward_kl): - losses, grad = _fused_cross_entropy_base_from_distribution( - logits, - target, - grad_output, - logits_scale_factor, - target_format, - group, - temperature, - return_kl_loss=entropy_loss_type == EntropyLossType.forward_kl, - ) - elif entropy_loss_type == EntropyLossType.reverse_kl: - losses, grad = _fused_reverse_kl_base_from_distribution( - logits, - target, - grad_output, - logits_scale_factor, - target_format, - group, - temperature, - ) - else: - raise NotImplementedError(entropy_loss_type) - - loss = reduce_losses(losses, divisor, loss_mask) - - if grad is not None: - if loss_mask is not None: - grad = grad * loss_mask.unsqueeze(-1) - grad = grad.to(logits.dtype) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) - - return loss, grad_logits diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py index 6de884fd9..1ec00cf06 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -6,7 +6,6 @@ from fast_llm.functional.entropy_loss import ( cross_entropy_from_distribution_core, cross_entropy_from_labels_core, - fused_entropy_loss_forward_backward, reverse_kl_from_distribution_core, ) from fast_llm.functional.triton.entropy_loss import triton_entropy_loss_forward_backward @@ -29,22 +28,24 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor, torch.Tensor | None]": - return ( - triton_entropy_loss_forward_backward - if TritonConfig.enabled(logits.device, self._config.use_triton) - else fused_entropy_loss_forward_backward - )( - logits, - self._get_labels(kwargs, split_index), - None, # Labels are already masked - grad_logits=grad_logits, - grad_output=self._get_grad_output(kwargs), - group=self._parallel_dim.group if self._vocab_parallel else None, - logits_scale_factor=self._logits_scale_factor, - target_format=TargetFormat.labels, - entropy_loss_type=self._config.loss_type, - divisor=self._get_label_count(kwargs), - ) + arguments = self.combinable_extract(kwargs, split_index, losses is not None) + group = self._parallel_dim.group if self._vocab_parallel else None + if TritonConfig.enabled(logits.device, self._config.use_triton): + target, grad_output, divisor = arguments + return triton_entropy_loss_forward_backward( + logits, + target, + None, # Labels are already masked + grad_logits=grad_logits, + grad_output=grad_output, + group=group, + logits_scale_factor=self._logits_scale_factor, + target_format=TargetFormat.labels, + entropy_loss_type=self._config.loss_type, + divisor=divisor, + ) + loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments) + return loss, grad_logits def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: return self._get_labels(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs) @@ -85,23 +86,25 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor, torch.Tensor | None]": - return ( - triton_entropy_loss_forward_backward - if TritonConfig.enabled(logits.device, self._config.use_triton) - else fused_entropy_loss_forward_backward - )( - logits, - self._get_reference_model_logits(self._config.reference_model, kwargs, split_index), - self._get_loss_mask(kwargs, split_index), - grad_output=self._get_grad_output(kwargs), - grad_logits=grad_logits, - group=self._parallel_dim.group if self._vocab_parallel else None, - logits_scale_factor=self._logits_scale_factor, - temperature=self._config.temperature, - target_format=TargetFormat.logits, - entropy_loss_type=self._config.loss_type, - divisor=self._get_label_count(kwargs), - ) + arguments = self.combinable_extract(kwargs, split_index, losses is not None) + group = self._parallel_dim.group if self._vocab_parallel else None + if TritonConfig.enabled(logits.device, self._config.use_triton): + target, loss_mask, grad_output, divisor, entropy_loss_type, temperature = arguments + return triton_entropy_loss_forward_backward( + logits, + target, + loss_mask, + grad_output=grad_output, + grad_logits=grad_logits, + group=group, + logits_scale_factor=self._logits_scale_factor, + temperature=temperature, + target_format=TargetFormat.logits, + entropy_loss_type=entropy_loss_type, + divisor=divisor, + ) + loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments) + return loss, grad_logits def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: return ( diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index bed045839..a0a0e455b 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -10,13 +10,18 @@ from fast_llm.engine.config_utils.data_type import DataType from fast_llm.engine.distributed.config import DistributedBackend, DistributedConfig from fast_llm.functional.config import EntropyLossType, TargetFormat -from fast_llm.functional.entropy_loss import fused_entropy_loss_forward_backward, torch_entropy_loss_forward_backward +from fast_llm.functional.entropy_loss import torch_entropy_loss_forward_backward from fast_llm.functional.triton import triton_available from fast_llm.functional.triton.entropy_loss import triton_entropy_loss_forward_backward from fast_llm.functional.triton.grpo_loss import triton_grpo_loss_forward_backward from fast_llm.functional.triton.gspo_loss import triton_gspo_loss_forward_backward from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward -from fast_llm.layers.language_model.loss.config import LanguageModelGRPOLossConfig, LanguageModelZLossConfig +from fast_llm.layers.language_model.loss.config import ( + LanguageModelDistillationLossConfig, + LanguageModelGRPOLossConfig, + LanguageModelLabelEntropyLossConfig, + LanguageModelZLossConfig, +) from fast_llm.layers.language_model.loss.dpo import dpo_loss from fast_llm.layers.language_model.loss.loss import loss_forward_backward from fast_llm.layers.language_model.loss.policy_gradient import ( @@ -273,7 +278,6 @@ def _test_entropy_loss( ): if target_format == TargetFormat.labels and entropy_loss_type == EntropyLossType.reverse_kl: pytest.skip(reason="Reverse KL loss not implemented for target labels") - # TODO: Test tensor-parallel implementation. logits, target, loss_mask = _get_lm_loss_inputs(num_columns, loss_masking, target_format, batch_shape, dtype) local_logits = split_op(logits, group, -1).contiguous() local_target = target if target_format == TargetFormat.labels else split_op(target, group, -1).contiguous() @@ -291,16 +295,19 @@ def _test_entropy_loss( previous_grad = torch.randn_like(grad_ref) grad_ref = grad_ref + previous_grad local_previous_grad = split_op(previous_grad, group, -1).contiguous() - out_fused, grad_fused = fused_entropy_loss_forward_backward( - logits=local_logits, - target=local_target, - loss_mask=loss_mask, - grad_logits=local_previous_grad.clone() if accumulate else None, - grad_output=grad_output, - group=group, - logits_scale_factor=logits_scale_factor, - target_format=target_format, - entropy_loss_type=entropy_loss_type, + divisor = local_logits.shape[:-1].numel() + if target_format == TargetFormat.labels: + loss = _combinable_loss( + LanguageModelLabelEntropyLossConfig(loss_type=entropy_loss_type), "ce", logits_scale_factor + ) + arguments = (local_target, grad_output, divisor) + else: + loss = _combinable_loss( + LanguageModelDistillationLossConfig(loss_type=entropy_loss_type), "distillation", logits_scale_factor + ) + arguments = (local_target, loss_mask, grad_output, divisor, entropy_loss_type, 1.0) + out_fused, grad_fused, _ = loss.combinable_forward_backward( + local_logits, group, local_previous_grad.clone() if accumulate else None, arguments ) _compare_losses_and_grads( out_fused, @@ -332,7 +339,7 @@ def _test_entropy_loss( grad_output is not None, grad_triton, grad_ref, - threshold=1e-5 if target_format != TargetFormat.probabilities and data_type == DataType.float32 else 1e-4, + threshold=1e-5 if data_type == DataType.float32 else 1e-4, group=group, ) @@ -595,7 +602,7 @@ def _test_z_loss( ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), _LOSS_PARAMETERS, ) -@pytest.mark.parametrize("target_format", TargetFormat) +@pytest.mark.parametrize("target_format", (TargetFormat.labels, TargetFormat.logits)) @pytest.mark.parametrize("entropy_loss_type", EntropyLossType) def test_entropy_loss( batch_shape, @@ -727,7 +734,7 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa suffix = f"{num_columns}-{grad_output}-{logits_scale_factor}-{loss_masking}-{dtype}-{block_size}-{accumulate}-{"_".join([str(i) for i in batch_shape])}" # Entropy loss for entropy_loss_type in EntropyLossType: - for target_format in TargetFormat: + for target_format in (TargetFormat.labels, TargetFormat.logits): if target_format == TargetFormat.labels and entropy_loss_type == EntropyLossType.reverse_kl: continue with test_context.subtest( @@ -842,7 +849,7 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): *( f"{entropy_loss_type}-{target_format}" for entropy_loss_type in EntropyLossType - for target_format in TargetFormat + for target_format in (TargetFormat.labels, TargetFormat.logits) if target_format != TargetFormat.labels or entropy_loss_type != EntropyLossType.reverse_kl ), "z_loss", From dcfa978f1df44fc0076ac9c849937eea784a2a26 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 2 Jul 2026 14:36:03 -0400 Subject: [PATCH 16/28] Fix entropy/z-loss test threshold comparing the wrong symbol (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The float32-vs-float16 tolerance selector compared `data_type` (the imported module `fast_llm.engine.config_utils.data_type`) to `DataType.float32`, which is always False — so the fused/triton comparisons always used the loose 1e-4 threshold, even in float32. Compare the `dtype` parameter instead, so float32 cases use the intended 1e-5; drop the now-unused module import. All float32 entropy/z-loss cases pass at the tightened tolerance. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/layers/test_lm_losses.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index a0a0e455b..6b7502dba 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -6,7 +6,6 @@ import torch from fast_llm.core.ops import split_op -from fast_llm.engine.config_utils import data_type from fast_llm.engine.config_utils.data_type import DataType from fast_llm.engine.distributed.config import DistributedBackend, DistributedConfig from fast_llm.functional.config import EntropyLossType, TargetFormat @@ -315,7 +314,7 @@ def _test_entropy_loss( grad_output is not None, grad_fused, grad_ref, - threshold=1e-5 if data_type == DataType.float32 else 1e-4, + threshold=1e-5 if dtype == DataType.float32 else 1e-4, group=group, ) @@ -339,7 +338,7 @@ def _test_entropy_loss( grad_output is not None, grad_triton, grad_ref, - threshold=1e-5 if data_type == DataType.float32 else 1e-4, + threshold=1e-5 if dtype == DataType.float32 else 1e-4, group=group, ) @@ -571,7 +570,7 @@ def _test_z_loss( grad_output is not None, grad_fused, grad_ref, - threshold=1e-5 if data_type == DataType.float32 else 1e-4, + threshold=1e-5 if dtype == DataType.float32 else 1e-4, group=group, ) if not triton_available: @@ -591,7 +590,7 @@ def _test_z_loss( grad_output is not None, grad_triton, grad_ref, - threshold=1e-5 if data_type == DataType.float32 else 1e-4, + threshold=1e-5 if dtype == DataType.float32 else 1e-4, group=group, ) From 5295e01821dd558acf51f6c20418d9e055ed8a5d Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 2 Jul 2026 14:38:27 -0400 Subject: [PATCH 17/28] Rename combinable_extract -> get_inputs, combinable_core -> fused_core (#507) Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/entropy_loss.py | 8 ++++---- fast_llm/layers/language_model/loss/config.py | 2 +- .../layers/language_model/loss/entropy_loss.py | 12 ++++++------ fast_llm/layers/language_model/loss/loss.py | 14 +++++++------- fast_llm/layers/language_model/loss/monolithic.py | 8 ++++---- .../layers/language_model/loss/policy_gradient.py | 6 +++--- fast_llm/layers/language_model/loss/z_loss.py | 6 +++--- 7 files changed, 28 insertions(+), 28 deletions(-) diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index 1290fbed8..43598dc0c 100644 --- a/fast_llm/functional/entropy_loss.py +++ b/fast_llm/functional/entropy_loss.py @@ -139,7 +139,7 @@ def reverse_kl_from_distribution_core( """ Reverse-KL math from a precomputed student softmax (adding a teacher softmax when the target is logits). - Used by the distillation `combinable_core`, inlined inside its `@torch.compile` boundary (standalone + Used by the distillation `fused_core`, inlined inside its `@torch.compile` boundary (standalone or monolithic). """ assert target_format in (TargetFormat.logits, TargetFormat.probabilities) @@ -187,7 +187,7 @@ def cross_entropy_from_distribution_core( Cross-entropy / forward-KL math from a precomputed student softmax (adding a teacher softmax when the target is logits). - Used by the distillation `combinable_core`, inlined inside its `@torch.compile` boundary (standalone + Used by the distillation `fused_core`, inlined inside its `@torch.compile` boundary (standalone or monolithic). """ if target_format == TargetFormat.logits: @@ -276,7 +276,7 @@ def cross_entropy_from_labels_core( """ Cross-entropy from labels, taking the already-computed shared softmax tensors. Returns the unmasked per-sample loss and (when `grad_output` is given) the unmasked gradient; the caller applies the loss - mask, reduction, and dtype cast. Used by the label / GRPO `combinable_core`, inlined inside its + mask, reduction, and dtype cast. Used by the label / GRPO `fused_core`, inlined inside its `@torch.compile` boundary (standalone or monolithic). """ predicted_logits, target_masked, target_mask = predicted_logits_from_labels(logits_norm, target, loss_mask, group) @@ -308,7 +308,7 @@ def z_loss_core( Z-loss from the already-computed shared softmax tensors. Returns the unmasked per-sample loss term (`log_sum_exp ** 2`) and (when `grad_output` is given) the unmasked gradient; the caller applies the loss mask, reduction, and dtype cast. z-loss needs the un-regularized log-sum-exp, so it adds back - `logits_max` (cross-entropy cancels it). Inlined by the z-loss `combinable_core` inside its + `logits_max` (cross-entropy cancels it). Inlined by the z-loss `fused_core` inside its `@torch.compile` boundary (standalone or monolithic). """ log_sum_exp_logits = sum_exp_logits.log() + logits_max diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 920d9bafe..8018ad525 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -36,7 +36,7 @@ class LanguageModelLossKwargs(BlockKwargs): @config_class(registry=True) class LanguageModelLossConfig(Config): _abstract: typing.ClassVar[bool] = True - # Whether this loss can be a `MonolithicLossConfig` entry (shares one softmax via `combinable_core`). + # Whether this loss can be a `MonolithicLossConfig` entry (shares one softmax via `fused_core`). combinable: typing.ClassVar[bool] = False weight: float = Field( diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py index 1ec00cf06..dd7f2c2b2 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -28,7 +28,7 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor, torch.Tensor | None]": - arguments = self.combinable_extract(kwargs, split_index, losses is not None) + arguments = self.get_inputs(kwargs, split_index, losses is not None) group = self._parallel_dim.group if self._vocab_parallel else None if TritonConfig.enabled(logits.device, self._config.use_triton): target, grad_output, divisor = arguments @@ -47,11 +47,11 @@ def _forward_backward( loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments) return loss, grad_logits - def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: return self._get_labels(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs) @staticmethod - def combinable_core( + def fused_core( logits_norm: torch.Tensor, exp_logits: torch.Tensor, sum_exp_logits: torch.Tensor, @@ -86,7 +86,7 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor, torch.Tensor | None]": - arguments = self.combinable_extract(kwargs, split_index, losses is not None) + arguments = self.get_inputs(kwargs, split_index, losses is not None) group = self._parallel_dim.group if self._vocab_parallel else None if TritonConfig.enabled(logits.device, self._config.use_triton): target, loss_mask, grad_output, divisor, entropy_loss_type, temperature = arguments @@ -106,7 +106,7 @@ def _forward_backward( loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments) return loss, grad_logits - def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: return ( self._get_reference_model_logits(self._config.reference_model, kwargs, split_index), self._get_loss_mask(kwargs, split_index), @@ -117,7 +117,7 @@ def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, re ) @staticmethod - def combinable_core( + def fused_core( logits_norm: torch.Tensor, exp_logits: torch.Tensor, sum_exp_logits: torch.Tensor, diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 91e7bf614..1869d6b9b 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -146,10 +146,10 @@ def _get_reference_model_logits(self, reference_model: str, kwargs: dict[str, ty class CombinableLoss: """Mixin for losses that consume the vocabulary softmax: each runs standalone through `combinable_forward_backward`, or several are fused together by `MonolithicLoss` over a single shared - softmax. Subclasses implement `combinable_extract` (eager kwargs -> argument tuple, run outside the - compiled boundary) and the `combinable_core` static method (the post-softmax math over the shared + softmax. Subclasses implement `get_inputs` (eager kwargs -> argument tuple, run outside the + compiled boundary) and the `fused_core` static method (the post-softmax math over the shared softmax tensors, returning `(loss, uncast_grad, extra)`), and override `register_combinable_extras` when - they emit outputs beyond the loss scalar. Both paths call the same `combinable_core`.""" + they emit outputs beyond the loss scalar. Both paths call the same `fused_core`.""" _logits_scale_factor: float @@ -162,10 +162,10 @@ def combinable_forward_backward( arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, typing.Any]: """Standalone realization of a single combinable loss: run the softmax once, then this loss's - `combinable_core`, then cast-and-accumulate the gradient. Equivalent to a one-child `MonolithicLoss` - — they share `combinable_core`, so this is not a second copy of the math.""" + `fused_core`, then cast-and-accumulate the gradient. Equivalent to a one-child `MonolithicLoss` + — they share `fused_core`, so this is not a second copy of the math.""" logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, self._logits_scale_factor, group) - loss, grad, extra = self.combinable_core( + loss, grad, extra = self.fused_core( logits_norm, exp_logits, sum_exp_logits, logits_max, group, self._logits_scale_factor, arguments ) return loss, self._accumulate_grad(grad, logits.dtype, grad_logits), extra @@ -189,7 +189,7 @@ def register_combinable_extras( self, extra: typing.Any, kwargs: dict[str, typing.Any], losses: dict | None ) -> None: """Register per-loss outputs beyond the scalar (e.g. GRPO's `new_logprobs` / metric family) produced - by `combinable_core`. No-op by default.""" + by `fused_core`. No-op by default.""" def loss_forward_backward( diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 35815a8bc..ac6e26f79 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -21,16 +21,16 @@ def _monolithic_core( arguments: tuple[tuple, ...], ) -> tuple[list, torch.Tensor | None]: """ - One shared softmax over the logits, then each child loss's `combinable_core` consuming it. The child + One shared softmax over the logits, then each child loss's `fused_core` consuming it. The child list is fixed per config, so the loop unrolls inside this single `@torch.compile` boundary and each - `combinable_core` dispatches (and inlines) to its loss type's math — every enabled loss is fused over + `fused_core` dispatches (and inlines) to its loss type's math — every enabled loss is fused over one softmax. Gradient contributions accumulate in fp32 and cast to `logits.dtype` once at the end. """ logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, logits_scale_factor, group) grad = None results = [] for child, child_arguments in zip(children, arguments): - loss, child_grad, extra = child.combinable_core( + loss, child_grad, extra = child.fused_core( logits_norm, exp_logits, sum_exp_logits, logits_max, group, logits_scale_factor, child_arguments ) results.append((loss, extra)) @@ -104,7 +104,7 @@ def forward_backward( grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor | None, torch.Tensor | None]": register = losses is not None - arguments = tuple(child.combinable_extract(kwargs, split_index, register) for child in self._children) + arguments = tuple(child.get_inputs(kwargs, split_index, register) for child in self._children) group = self._parallel_dim.group if self._vocab_parallel else None results, grad_logits = _monolithic_core( tuple(self._children), logits, group, self._softmax_scale_factor, grad_logits, arguments diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 558b62f63..3dc7eb8a1 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -102,7 +102,7 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: - arguments = self.combinable_extract(kwargs, split_index, losses is not None) + arguments = self.get_inputs(kwargs, split_index, losses is not None) group = self._parallel_dim.group if self._vocab_parallel else None if TritonConfig.enabled(logits.device, self._config.use_triton): from fast_llm.functional.triton.grpo_loss import triton_grpo_loss_forward_backward @@ -143,7 +143,7 @@ def _forward_backward( self.register_combinable_extras(extra, kwargs, losses) return loss, grad_logits - def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: # When nothing is logged this step, drop the metric-only outputs (`new_logprobs_mean` and the # GRPO metric family) by passing `num_labels_in_seq=None` / `compute_metrics=False`. return ( @@ -160,7 +160,7 @@ def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, re ) @staticmethod - def combinable_core( + def fused_core( logits_norm: torch.Tensor, exp_logits: torch.Tensor, sum_exp_logits: torch.Tensor, diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py index e904ffe16..923200bce 100644 --- a/fast_llm/layers/language_model/loss/z_loss.py +++ b/fast_llm/layers/language_model/loss/z_loss.py @@ -19,7 +19,7 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor, torch.Tensor | None]": - arguments = self.combinable_extract(kwargs, split_index, losses is not None) + arguments = self.get_inputs(kwargs, split_index, losses is not None) group = self._parallel_dim.group if self._vocab_parallel else None if TritonConfig.enabled(logits.device, self._config.use_triton): loss_mask, grad_output, divisor = arguments @@ -35,11 +35,11 @@ def _forward_backward( loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments) return loss, grad_logits - def combinable_extract(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: return self._get_loss_mask(kwargs, split_index), self._get_grad_output(kwargs), self._get_label_count(kwargs) @staticmethod - def combinable_core( + def fused_core( logits_norm: torch.Tensor, exp_logits: torch.Tensor, sum_exp_logits: torch.Tensor, From 735f4fd46df69a6401310880da82ceefde5a995d Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 2 Jul 2026 15:06:21 -0400 Subject: [PATCH 18/28] Extract single-loss template into SingleLoss intermediate (#507) The forward_backward/_forward_backward template pair (skip -> _forward_backward -> register -> weight) is a property of single-scalar losses, not the loss abstraction itself. Move it into a new SingleLoss intermediate; make the base LanguageModelLoss.forward_backward abstract. Composite losses (MonolithicLoss) now satisfy forward_backward directly instead of carrying a dead _forward_backward override that only existed to satisfy the abstract method. Co-Authored-By: Claude Opus 4.8 (1M context) --- bench_monolithic.py | 218 ++++++++++++++++++ fast_llm/layers/language_model/loss/dpo.py | 4 +- .../language_model/loss/entropy_loss.py | 6 +- fast_llm/layers/language_model/loss/loss.py | 52 +++-- .../layers/language_model/loss/monolithic.py | 12 - .../language_model/loss/policy_gradient.py | 6 +- fast_llm/layers/language_model/loss/z_loss.py | 4 +- 7 files changed, 262 insertions(+), 40 deletions(-) create mode 100644 bench_monolithic.py diff --git a/bench_monolithic.py b/bench_monolithic.py new file mode 100644 index 000000000..4bcf8c7dc --- /dev/null +++ b/bench_monolithic.py @@ -0,0 +1,218 @@ +"""Per-loss vs monolithic (fused) head-loss benchmark: time + peak memory, individual losses + combos.""" + +import argparse + +import torch +import torch._dynamo + +from fast_llm.functional.config import EntropyLossType, TargetFormat +from fast_llm.functional.entropy_loss import fused_entropy_loss_forward_backward +from fast_llm.layers.language_model.loss.monolithic import MonolithicLossSpec, monolithic_head_loss_forward_backward +from fast_llm.layers.language_model.loss.policy_gradient import compute_grpo_metrics, fused_grpo_loss_forward_backward +from fast_llm.layers.language_model.loss.z_loss import fused_z_loss_forward_backward + +# Each scenario is a distinct fixed loss config (one compiled guard). The benchmark cycles many in one +# process; bump the limit so none falls back to eager. A real run has ONE config → compiles once. +torch._dynamo.config.recompile_limit = 64 +torch._dynamo.config.cache_size_limit = 64 + + +parser = argparse.ArgumentParser() +parser.add_argument("--num-tokens", type=int, default=16384) +parser.add_argument("--vocab", type=int, default=128256) +parser.add_argument("--iters", type=int, default=50) +parser.add_argument("--warmup", type=int, default=5) +args = parser.parse_args() + +device = "cuda" +DTYPE = torch.bfloat16 +GRAD = 1.0 +N, V = args.num_tokens, args.vocab +torch.manual_seed(0) + +logits = torch.randn(N, V, dtype=DTYPE, device=device) +labels = torch.randint(0, V, (N,), dtype=torch.int64, device=device) +teacher = torch.randn(N, V, dtype=DTYPE, device=device) +advantages = torch.randn(N, dtype=torch.float32, device=device) +old_logprobs = torch.randn(N, dtype=torch.float32, device=device) +num_labels_in_seq = torch.full((N,), N, dtype=torch.int32, device=device) +label_counts = torch.full((N,), float(N), dtype=torch.float32, device=device) +DIV = float(N) + + +# ---- per-loss parts (mirror the head's per_loss loop: each its own kernel/softmax, threading the grad) ---- +def part_ce(grad): + _, grad = fused_entropy_loss_forward_backward( + logits=logits, + target=labels, + loss_mask=None, + grad_logits=grad, + grad_output=GRAD, + group=None, + logits_scale_factor=1.0, + target_format=TargetFormat.labels, + entropy_loss_type=EntropyLossType.cross_entropy, + divisor=DIV, + ) + return grad + + +def part_z(grad): + _, grad = fused_z_loss_forward_backward( + logits=logits, + loss_mask=None, + grad_logits=grad, + grad_output=GRAD, + group=None, + logits_scale_factor=1.0, + divisor=DIV, + ) + return grad + + +def part_distill(grad): + _, grad = fused_entropy_loss_forward_backward( + logits=logits, + target=teacher, + loss_mask=None, + grad_logits=grad, + grad_output=GRAD, + group=None, + logits_scale_factor=1.0, + target_format=TargetFormat.logits, + entropy_loss_type=EntropyLossType.cross_entropy, + divisor=DIV, + ) + return grad + + +def part_grpo(grad): + _, grad, _ = fused_grpo_loss_forward_backward( + logits, + labels, + advantages, + old_logprobs, + grad_logits=grad, + grad_output=GRAD, + group=None, + logits_scale_factor=1.0, + num_labels_in_seq=num_labels_in_seq, + divisor=DIV, + ) + return grad + + +def part_grpo_metrics(grad): + # The #494 second softmax: the per_loss path recomputes the softmax to get the metric family. + compute_grpo_metrics( + logits, + labels, + advantages, + old_logprobs, + label_counts, + epsilon_low=0.2, + epsilon_high=0.2, + logits_scale_factor=1.0, + group=None, + compute_entropy=False, + ) + return grad + + +def per_loss(parts): + def run(): + grad = None + for part in parts: + grad = part(grad) + return grad + + return run + + +# ---- fused specs ---- +def spec_ce(): + return MonolithicLossSpec("cross_entropy", "ce", 1.0, 1.0, GRAD, DIV, target=labels) + + +def spec_z(): + return MonolithicLossSpec("z_loss", "z", 1.0, 1.0, GRAD, DIV) + + +def spec_distill(): + return MonolithicLossSpec( + "entropy_from_distribution", + "distill", + 1.0, + 1.0, + GRAD, + DIV, + target=teacher, + target_format=TargetFormat.logits, + entropy_loss_type=EntropyLossType.cross_entropy, + ) + + +def spec_grpo(metrics=False): + return MonolithicLossSpec( + "grpo", + "grpo", + 1.0, + 1.0, + GRAD, + DIV, + target=labels, + advantages=advantages, + old_log_probabilities=old_logprobs, + num_labels_in_seq=num_labels_in_seq, + compute_metrics=metrics, + ) + + +def fused(specs): + def run(): + return monolithic_head_loss_forward_backward(logits, specs, group=None, grad_logits=None) + + return run + + +def bench(run): + for _ in range(args.warmup): + run() + torch.cuda.synchronize() + torch.cuda.empty_cache() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + times = [] + for _ in range(args.iters): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + times.sort() + return times[len(times) // 2], (torch.cuda.max_memory_allocated() - before) / 1e9 + + +scenarios = [ + ("CE", [part_ce], [spec_ce()]), + ("z_loss", [part_z], [spec_z()]), + ("distillation", [part_distill], [spec_distill()]), + ("GRPO", [part_grpo], [spec_grpo()]), + ("CE+z", [part_ce, part_z], [spec_ce(), spec_z()]), + ("CE+distill", [part_ce, part_distill], [spec_ce(), spec_distill()]), + ("CE+z+distill", [part_ce, part_z, part_distill], [spec_ce(), spec_z(), spec_distill()]), + ("GRPO+metrics", [part_grpo, part_grpo_metrics], [spec_grpo(metrics=True)]), + ("GRPO+metrics+z", [part_grpo, part_grpo_metrics, part_z], [spec_grpo(metrics=True), spec_z()]), +] + +print(f"shape: num_tokens={N} vocab={V} dtype={DTYPE} | iters={args.iters}") +print( + f"{'scenario':<18}{'per_loss ms':>12}{'fused ms':>11}{'speedup':>9}{'per_loss GB':>13}{'fused GB':>10}{'mem saved':>11}" +) +for name, parts, specs in scenarios: + torch.cuda.empty_cache() + t_pl, m_pl = bench(per_loss(parts)) + torch.cuda.empty_cache() + t_f, m_f = bench(fused(specs)) + print(f"{name:<18}{t_pl:>12.2f}{t_f:>11.2f}{t_pl / t_f:>8.2f}x{m_pl:>13.2f}{m_f:>10.2f}{(m_pl - m_f):>10.2f}G") diff --git a/fast_llm/layers/language_model/loss/dpo.py b/fast_llm/layers/language_model/loss/dpo.py index 059f808e5..644745b64 100644 --- a/fast_llm/layers/language_model/loss/dpo.py +++ b/fast_llm/layers/language_model/loss/dpo.py @@ -3,10 +3,10 @@ import torch from fast_llm.layers.language_model.loss.config import LanguageModelDPOLossConfig, LanguageModelLossKwargs -from fast_llm.layers.language_model.loss.loss import LanguageModelLoss, loss_forward_backward +from fast_llm.layers.language_model.loss.loss import SingleLoss, loss_forward_backward -class LanguageModelDPOLoss[ConfigType: LanguageModelDPOLossConfig](LanguageModelLoss[ConfigType]): +class LanguageModelDPOLoss[ConfigType: LanguageModelDPOLossConfig](SingleLoss[ConfigType]): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) if self._prediction_distance > 1: diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py index dd7f2c2b2..f4ccd74ff 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -14,11 +14,11 @@ LanguageModelDistillationLossConfig, LanguageModelLabelEntropyLossConfig, ) -from fast_llm.layers.language_model.loss.loss import CombinableLoss, LanguageModelLoss +from fast_llm.layers.language_model.loss.loss import CombinableLoss, SingleLoss class LanguageModelLabelEntropyLoss[ConfigType: LanguageModelLabelEntropyLossConfig]( - CombinableLoss, LanguageModelLoss[ConfigType] + CombinableLoss, SingleLoss[ConfigType] ): def _forward_backward( self, @@ -76,7 +76,7 @@ def fused_core( class LanguageModelDistillationLoss[ConfigType: LanguageModelDistillationLossConfig]( - CombinableLoss, LanguageModelLoss[ConfigType] + CombinableLoss, SingleLoss[ConfigType] ): def _forward_backward( self, diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 1869d6b9b..c4dacabbe 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -43,6 +43,7 @@ def __init__( self._sequence_data_dim = distributed_config.get_distributed_dim(DistributedDimNames.sequence_data) self._sequence_data_active = self._sequence_data_dim.size > 1 + @abc.abstractmethod def forward_backward( self, logits: "torch.Tensor", @@ -51,23 +52,6 @@ def forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor | None, torch.Tensor | None]": - if losses is None and self._weight is None: - # Loss computation is not needed, skip. - return None, None - loss, grad = self._forward_backward(logits, kwargs, losses, split_index, grad_logits) - if self._do_register_loss: - self._register_loss(self.name, loss, losses) - return loss if self._weight == 1 else loss * self._weight, grad - - @abc.abstractmethod - def _forward_backward( - self, - logits: "torch.Tensor", - kwargs: dict[str, typing.Any], - losses: dict | None = None, - split_index: int = 0, - grad_logits: torch.Tensor | None = None, - ) -> "tuple[torch.Tensor, torch.Tensor | None]": pass def get_loss_definitions(self) -> list[LossDef]: @@ -143,6 +127,40 @@ def _get_reference_model_logits(self, reference_model: str, kwargs: dict[str, ty ) +class SingleLoss[ConfigType: LanguageModelLossConfig](LanguageModelLoss[ConfigType]): + """A loss emitting a single registered, weighted scalar. Subclasses implement `_forward_backward` (the + loss math and its gradient); this template skips the disabled case, registers the scalar, and applies the + per-loss weight. `MonolithicLoss` composes several of these over one shared softmax, so it satisfies + `forward_backward` directly rather than through this single-loss template.""" + + def forward_backward( + self, + logits: "torch.Tensor", + kwargs: dict[str, typing.Any], + losses: dict | None = None, + split_index: int = 0, + grad_logits: torch.Tensor | None = None, + ) -> "tuple[torch.Tensor | None, torch.Tensor | None]": + if losses is None and self._weight is None: + # Loss computation is not needed, skip. + return None, None + loss, grad = self._forward_backward(logits, kwargs, losses, split_index, grad_logits) + if self._do_register_loss: + self._register_loss(self.name, loss, losses) + return loss if self._weight == 1 else loss * self._weight, grad + + @abc.abstractmethod + def _forward_backward( + self, + logits: "torch.Tensor", + kwargs: dict[str, typing.Any], + losses: dict | None = None, + split_index: int = 0, + grad_logits: torch.Tensor | None = None, + ) -> "tuple[torch.Tensor, torch.Tensor | None]": + pass + + class CombinableLoss: """Mixin for losses that consume the vocabulary softmax: each runs standalone through `combinable_forward_backward`, or several are fused together by `MonolithicLoss` over a single shared diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index ac6e26f79..90f6d4ce4 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -119,18 +119,6 @@ def forward_backward( total_loss = weighted if total_loss is None else total_loss + weighted return total_loss, grad_logits - def _forward_backward( - self, - logits: "torch.Tensor", - kwargs: dict[str, typing.Any], - losses: dict | None = None, - split_index: int = 0, - grad_logits: torch.Tensor | None = None, - ) -> "tuple[torch.Tensor, torch.Tensor | None]": - # `MonolithicLoss` overrides the public `forward_backward` directly (it composes weighted child - # losses and registers each child), so this per-loss hook is never the entry point. - raise NotImplementedError - def get_preprocessing_config(self) -> dict[str, typing.Any]: return safe_merge_dicts(*(child.get_preprocessing_config() for child in self._children)) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 3dc7eb8a1..21fb4ad94 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -23,13 +23,11 @@ LanguageModelPolicyGradientLossConfig, ) from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, grpo_metrics_core -from fast_llm.layers.language_model.loss.loss import CombinableLoss, LanguageModelLoss +from fast_llm.layers.language_model.loss.loss import CombinableLoss, SingleLoss from fast_llm.utils import Assert -class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLossConfig]( - LanguageModelLoss[ConfigType] -): +class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLossConfig](SingleLoss[ConfigType]): """Shared scaffolding for policy-gradient losses (GRPO, GSPO).""" def _register_new_logprobs( diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py index 923200bce..edbaa77b6 100644 --- a/fast_llm/layers/language_model/loss/z_loss.py +++ b/fast_llm/layers/language_model/loss/z_loss.py @@ -7,10 +7,10 @@ from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward from fast_llm.functional.utils import reduce_losses from fast_llm.layers.language_model.loss.config import LanguageModelZLossConfig -from fast_llm.layers.language_model.loss.loss import CombinableLoss, LanguageModelLoss +from fast_llm.layers.language_model.loss.loss import CombinableLoss, SingleLoss -class LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](CombinableLoss, LanguageModelLoss[ConfigType]): +class LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](CombinableLoss, SingleLoss[ConfigType]): def _forward_backward( self, logits: "torch.Tensor", From 869db685214ad5c278701380c12fa6058aeb3c8a Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 2 Jul 2026 16:03:42 -0400 Subject: [PATCH 19/28] Untrack stray benchmark; trim loss docstrings/comments (#507) Fine-review follow-up: remove bench_monolithic.py from tracking (a stale local dev script that should never have been committed), and drop downstream-consumer references from the new loss/functional docstrings and comments per project style ("describe what the thing does, not which caller uses it"). Also fixes the cross_entropy_from_labels_core docstring, which claimed GRPO routes through it (GRPO calls predicted_logits_from_labels directly). Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/entropy_loss.py | 33 +++++++------------ fast_llm/layers/language_model/loss/config.py | 2 +- .../language_model/loss/grpo_metrics.py | 5 ++- fast_llm/layers/language_model/loss/loss.py | 26 +++++++-------- .../layers/language_model/loss/monolithic.py | 9 +++-- .../language_model/loss/policy_gradient.py | 11 +++---- fast_llm/layers/language_model/loss/z_loss.py | 5 ++- 7 files changed, 36 insertions(+), 55 deletions(-) diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index 43598dc0c..457647796 100644 --- a/fast_llm/functional/entropy_loss.py +++ b/fast_llm/functional/entropy_loss.py @@ -105,8 +105,8 @@ def softmax_base( Warning: The returned values are regularized by `logits_max`. The regularization typically but not always cancels out in derived quantities. - This plain (un-compiled) core is shared between the public `fused_softmax_base` wrapper and the - monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. + Un-compiled so it can be inlined into a `@torch.compile` boundary that fuses several losses over a + single softmax; `fused_softmax_base` is the compiled standalone wrapper. """ logits = logits.float() if logits_scale_factor != 1.0: @@ -136,12 +136,8 @@ def reverse_kl_from_distribution_core( group: ProcessGroup | None = None, temperature: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) - """ - Reverse-KL math from a precomputed student softmax (adding a teacher softmax when the target is logits). - - Used by the distillation `fused_core`, inlined inside its `@torch.compile` boundary (standalone - or monolithic). - """ + """Reverse-KL from a precomputed student softmax (adding a teacher softmax when the target is logits). + Un-compiled core, inlined into a `@torch.compile` boundary.""" assert target_format in (TargetFormat.logits, TargetFormat.probabilities) predicted_log_probability = logits_norm - sum_exp_logits.log().unsqueeze(-1) predicted_probability = exp_logits / sum_exp_logits.unsqueeze(-1) @@ -183,13 +179,8 @@ def cross_entropy_from_distribution_core( temperature: float = 1.0, return_kl_loss: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: # (*batch,), (*batch, vocab) - """ - Cross-entropy / forward-KL math from a precomputed student softmax (adding a teacher softmax when the - target is logits). - - Used by the distillation `fused_core`, inlined inside its `@torch.compile` boundary (standalone - or monolithic). - """ + """Cross-entropy / forward-KL from a precomputed student softmax (adding a teacher softmax when the + target is logits). Un-compiled core, inlined into a `@torch.compile` boundary.""" if target_format == TargetFormat.logits: target_logits_norm, exp_logits_targets, sum_exp_target_logits, _ = softmax_base( target, logits_scale_factor / temperature, group @@ -233,12 +224,12 @@ def predicted_logits_from_labels( Recover the value of the logits at the target index, with support for masking (target < 0) and tensor parallelism. In the simple case, equivalent to `logits.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)` - Normally used in combination with `fused_softmax_base`, may also recover probabilities or log probabilities: + May also recover probabilities or log probabilities: `predicted_probabilities = predicted_logits.exp() / sum_exp_logits` `predicted_log_probabilities = predicted_logits / sum_exp_logits.log()` - This plain (un-compiled) core is shared between the public `fused_predicted_logits_from_labels` wrapper - and the monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary. + Un-compiled core, inlined into a `@torch.compile` boundary; `fused_predicted_logits_from_labels` is the + compiled standalone wrapper. """ if group is None: @@ -276,8 +267,7 @@ def cross_entropy_from_labels_core( """ Cross-entropy from labels, taking the already-computed shared softmax tensors. Returns the unmasked per-sample loss and (when `grad_output` is given) the unmasked gradient; the caller applies the loss - mask, reduction, and dtype cast. Used by the label / GRPO `fused_core`, inlined inside its - `@torch.compile` boundary (standalone or monolithic). + mask, reduction, and dtype cast. Un-compiled core, inlined into a `@torch.compile` boundary. """ predicted_logits, target_masked, target_mask = predicted_logits_from_labels(logits_norm, target, loss_mask, group) @@ -308,8 +298,7 @@ def z_loss_core( Z-loss from the already-computed shared softmax tensors. Returns the unmasked per-sample loss term (`log_sum_exp ** 2`) and (when `grad_output` is given) the unmasked gradient; the caller applies the loss mask, reduction, and dtype cast. z-loss needs the un-regularized log-sum-exp, so it adds back - `logits_max` (cross-entropy cancels it). Inlined by the z-loss `fused_core` inside its - `@torch.compile` boundary (standalone or monolithic). + `logits_max` (cross-entropy cancels it). Un-compiled core, inlined into a `@torch.compile` boundary. """ log_sum_exp_logits = sum_exp_logits.log() + logits_max loss_term = log_sum_exp_logits**2 diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 8018ad525..4222469cc 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -36,7 +36,7 @@ class LanguageModelLossKwargs(BlockKwargs): @config_class(registry=True) class LanguageModelLossConfig(Config): _abstract: typing.ClassVar[bool] = True - # Whether this loss can be a `MonolithicLossConfig` entry (shares one softmax via `fused_core`). + # Whether this loss shares the vocabulary softmax via `fused_core`, so it can be fused with others. combinable: typing.ClassVar[bool] = False weight: float = Field( diff --git a/fast_llm/layers/language_model/loss/grpo_metrics.py b/fast_llm/layers/language_model/loss/grpo_metrics.py index 748f78e29..2d1b5af52 100644 --- a/fast_llm/layers/language_model/loss/grpo_metrics.py +++ b/fast_llm/layers/language_model/loss/grpo_metrics.py @@ -38,9 +38,8 @@ def grpo_metrics_core( term needing the full vocab axis (a sum over the local slice plus a tensor-parallel all-reduce); every other term is per-token. - This plain (un-compiled) core is shared between the public `compute_grpo_metrics` wrapper and the - monolithic head-loss kernel, which inlines it inside its own `@torch.compile` boundary — so the metrics - reuse the loss kernel's softmax instead of recomputing it. + Un-compiled core, inlined into a `@torch.compile` boundary so the metrics reuse the loss kernel's softmax + instead of recomputing it. """ mask = loss_mask.float() masked = mask / label_counts.float().clamp(min=1) diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index c4dacabbe..0c5c3b6ef 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -130,8 +130,7 @@ def _get_reference_model_logits(self, reference_model: str, kwargs: dict[str, ty class SingleLoss[ConfigType: LanguageModelLossConfig](LanguageModelLoss[ConfigType]): """A loss emitting a single registered, weighted scalar. Subclasses implement `_forward_backward` (the loss math and its gradient); this template skips the disabled case, registers the scalar, and applies the - per-loss weight. `MonolithicLoss` composes several of these over one shared softmax, so it satisfies - `forward_backward` directly rather than through this single-loss template.""" + per-loss weight. Composite losses satisfy `forward_backward` directly rather than through this template.""" def forward_backward( self, @@ -162,12 +161,11 @@ def _forward_backward( class CombinableLoss: - """Mixin for losses that consume the vocabulary softmax: each runs standalone through - `combinable_forward_backward`, or several are fused together by `MonolithicLoss` over a single shared - softmax. Subclasses implement `get_inputs` (eager kwargs -> argument tuple, run outside the - compiled boundary) and the `fused_core` static method (the post-softmax math over the shared - softmax tensors, returning `(loss, uncast_grad, extra)`), and override `register_combinable_extras` when - they emit outputs beyond the loss scalar. Both paths call the same `fused_core`.""" + """Mixin for losses that consume the vocabulary softmax, either standalone through + `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.""" _logits_scale_factor: float @@ -179,9 +177,9 @@ def combinable_forward_backward( grad_logits: torch.Tensor | None, arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, typing.Any]: - """Standalone realization of a single combinable loss: run the softmax once, then this loss's - `fused_core`, then cast-and-accumulate the gradient. Equivalent to a one-child `MonolithicLoss` - — they share `fused_core`, so this is not a second copy of the math.""" + """Standalone realization of a single combinable loss: softmax once, this loss's `fused_core`, then + cast-and-accumulate the gradient. Shares `fused_core` with the fused path, so it is not a second copy + of the math.""" logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, self._logits_scale_factor, group) loss, grad, extra = self.fused_core( logits_norm, exp_logits, sum_exp_logits, logits_max, group, self._logits_scale_factor, arguments @@ -193,8 +191,7 @@ def _accumulate_grad( grad: torch.Tensor | None, logits_dtype: torch.dtype, grad_logits: torch.Tensor | None ) -> torch.Tensor | None: """Cast a computed logits gradient to the logits dtype and accumulate it into `grad_logits` (in place - when a buffer exists, otherwise adopting it). The standalone kernels cast their single gradient here; - the monolithic kernel casts the fp32 sum of its children's gradients once.""" + when a buffer exists, otherwise adopting it).""" if grad is None: return grad_logits grad = grad.to(logits_dtype) @@ -206,8 +203,7 @@ def _accumulate_grad( def register_combinable_extras( self, extra: typing.Any, kwargs: dict[str, typing.Any], losses: dict | None ) -> None: - """Register per-loss outputs beyond the scalar (e.g. GRPO's `new_logprobs` / metric family) produced - by `fused_core`. No-op by default.""" + """Register per-loss outputs beyond the loss scalar (produced by `fused_core`). No-op by default.""" def loss_forward_backward( diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 90f6d4ce4..a701c22f1 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -41,11 +41,10 @@ def _monolithic_core( class MonolithicLoss[ConfigType: MonolithicLossConfig](LanguageModelLoss[ConfigType]): """ - A composite loss that runs the vocabulary softmax once and shares it across its combinable child - losses (cross-entropy, z-loss, distillation, GRPO), emitting each child's scalar / metrics and the - combined logits gradient in a single `@torch.compile` boundary. It is an ordinary head loss: the head - loops over it like any other and threads the same gradient buffer, so non-combinable losses (e.g. DPO) - are plain siblings in the head's loss list. + A composite loss that runs the vocabulary softmax once and shares it across its combinable child losses, + emitting each child's scalar / metrics and the combined logits gradient in a single `@torch.compile` + boundary. It is an ordinary head loss: the head loops over it like any other and threads the same gradient + buffer, so non-combinable losses remain plain siblings in the head's loss list. """ def __init__( diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 21fb4ad94..2d7b7f2ea 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -167,10 +167,9 @@ def fused_core( logits_scale_factor: float, arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, tuple]: - """Post-softmax GRPO over the shared softmax, called by both the standalone - `combinable_forward_backward` and the monolithic head loss. Returns the loss scalar, the uncast masked - gradient (the caller casts), and the `(new_logprobs_mean, metrics)` extras (each `None` when not - requested) — all from one softmax and one predicted-logit lookup.""" + """Post-softmax GRPO over the shared softmax. Returns the loss scalar, the uncast masked gradient (the + caller casts), and the `(new_logprobs_mean, metrics)` extras (each `None` when not requested) — all + from one softmax and one predicted-logit lookup.""" ( target, advantages, @@ -241,8 +240,8 @@ def fused_core( * loss_mask ) # d(probability_ratio)/d(logits) = - probability_ratio * (predicted_probabilities - target_probabilities) - # (Sign absorbed in probability_ratio_grad). Out-of-place `unsqueeze` since the shared softmax tensors - # may be reused by sibling losses in the monolithic path. + # (Sign absorbed in probability_ratio_grad). Out-of-place `unsqueeze` so the shared softmax tensors + # are not mutated in place, since other losses may reuse them over the same softmax. predicted_probabilities = exp_logits / sum_exp_logits.unsqueeze(-1) grad = (probability_ratio_grad * probability_ratio).unsqueeze(-1) * predicted_probabilities.scatter_add( -1, diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py index edbaa77b6..8ebd23f8e 100644 --- a/fast_llm/layers/language_model/loss/z_loss.py +++ b/fast_llm/layers/language_model/loss/z_loss.py @@ -48,9 +48,8 @@ def fused_core( logits_scale_factor: float, arguments: tuple, ) -> tuple[torch.Tensor, torch.Tensor | None, None]: - """Post-softmax z-loss over the shared softmax, called by both the standalone - `combinable_forward_backward` and the monolithic head loss. Returns the loss scalar and the uncast, - masked gradient contribution (the caller casts); z-loss emits no extra outputs.""" + """Post-softmax z-loss over the shared softmax. Returns the loss scalar and the uncast, masked + gradient contribution (the caller casts); z-loss emits no extra outputs.""" loss_mask, grad_output, divisor = arguments grad_output = None if grad_output is None else grad_output / divisor * logits_scale_factor loss_term, grad = z_loss_core(exp_logits, sum_exp_logits, logits_max, grad_output) From cb5be2f90bde1c2a5ed13c0266095ba6a21a0c08 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 2 Jul 2026 16:04:04 -0400 Subject: [PATCH 20/28] Untrack bench_monolithic.py (#507) Re-drop the stale local dev benchmark that a prior git add re-staged; it should remain an untracked local file, not part of the PR. Co-Authored-By: Claude Opus 4.8 (1M context) --- bench_monolithic.py | 218 -------------------------------------------- 1 file changed, 218 deletions(-) delete mode 100644 bench_monolithic.py diff --git a/bench_monolithic.py b/bench_monolithic.py deleted file mode 100644 index 4bcf8c7dc..000000000 --- a/bench_monolithic.py +++ /dev/null @@ -1,218 +0,0 @@ -"""Per-loss vs monolithic (fused) head-loss benchmark: time + peak memory, individual losses + combos.""" - -import argparse - -import torch -import torch._dynamo - -from fast_llm.functional.config import EntropyLossType, TargetFormat -from fast_llm.functional.entropy_loss import fused_entropy_loss_forward_backward -from fast_llm.layers.language_model.loss.monolithic import MonolithicLossSpec, monolithic_head_loss_forward_backward -from fast_llm.layers.language_model.loss.policy_gradient import compute_grpo_metrics, fused_grpo_loss_forward_backward -from fast_llm.layers.language_model.loss.z_loss import fused_z_loss_forward_backward - -# Each scenario is a distinct fixed loss config (one compiled guard). The benchmark cycles many in one -# process; bump the limit so none falls back to eager. A real run has ONE config → compiles once. -torch._dynamo.config.recompile_limit = 64 -torch._dynamo.config.cache_size_limit = 64 - - -parser = argparse.ArgumentParser() -parser.add_argument("--num-tokens", type=int, default=16384) -parser.add_argument("--vocab", type=int, default=128256) -parser.add_argument("--iters", type=int, default=50) -parser.add_argument("--warmup", type=int, default=5) -args = parser.parse_args() - -device = "cuda" -DTYPE = torch.bfloat16 -GRAD = 1.0 -N, V = args.num_tokens, args.vocab -torch.manual_seed(0) - -logits = torch.randn(N, V, dtype=DTYPE, device=device) -labels = torch.randint(0, V, (N,), dtype=torch.int64, device=device) -teacher = torch.randn(N, V, dtype=DTYPE, device=device) -advantages = torch.randn(N, dtype=torch.float32, device=device) -old_logprobs = torch.randn(N, dtype=torch.float32, device=device) -num_labels_in_seq = torch.full((N,), N, dtype=torch.int32, device=device) -label_counts = torch.full((N,), float(N), dtype=torch.float32, device=device) -DIV = float(N) - - -# ---- per-loss parts (mirror the head's per_loss loop: each its own kernel/softmax, threading the grad) ---- -def part_ce(grad): - _, grad = fused_entropy_loss_forward_backward( - logits=logits, - target=labels, - loss_mask=None, - grad_logits=grad, - grad_output=GRAD, - group=None, - logits_scale_factor=1.0, - target_format=TargetFormat.labels, - entropy_loss_type=EntropyLossType.cross_entropy, - divisor=DIV, - ) - return grad - - -def part_z(grad): - _, grad = fused_z_loss_forward_backward( - logits=logits, - loss_mask=None, - grad_logits=grad, - grad_output=GRAD, - group=None, - logits_scale_factor=1.0, - divisor=DIV, - ) - return grad - - -def part_distill(grad): - _, grad = fused_entropy_loss_forward_backward( - logits=logits, - target=teacher, - loss_mask=None, - grad_logits=grad, - grad_output=GRAD, - group=None, - logits_scale_factor=1.0, - target_format=TargetFormat.logits, - entropy_loss_type=EntropyLossType.cross_entropy, - divisor=DIV, - ) - return grad - - -def part_grpo(grad): - _, grad, _ = fused_grpo_loss_forward_backward( - logits, - labels, - advantages, - old_logprobs, - grad_logits=grad, - grad_output=GRAD, - group=None, - logits_scale_factor=1.0, - num_labels_in_seq=num_labels_in_seq, - divisor=DIV, - ) - return grad - - -def part_grpo_metrics(grad): - # The #494 second softmax: the per_loss path recomputes the softmax to get the metric family. - compute_grpo_metrics( - logits, - labels, - advantages, - old_logprobs, - label_counts, - epsilon_low=0.2, - epsilon_high=0.2, - logits_scale_factor=1.0, - group=None, - compute_entropy=False, - ) - return grad - - -def per_loss(parts): - def run(): - grad = None - for part in parts: - grad = part(grad) - return grad - - return run - - -# ---- fused specs ---- -def spec_ce(): - return MonolithicLossSpec("cross_entropy", "ce", 1.0, 1.0, GRAD, DIV, target=labels) - - -def spec_z(): - return MonolithicLossSpec("z_loss", "z", 1.0, 1.0, GRAD, DIV) - - -def spec_distill(): - return MonolithicLossSpec( - "entropy_from_distribution", - "distill", - 1.0, - 1.0, - GRAD, - DIV, - target=teacher, - target_format=TargetFormat.logits, - entropy_loss_type=EntropyLossType.cross_entropy, - ) - - -def spec_grpo(metrics=False): - return MonolithicLossSpec( - "grpo", - "grpo", - 1.0, - 1.0, - GRAD, - DIV, - target=labels, - advantages=advantages, - old_log_probabilities=old_logprobs, - num_labels_in_seq=num_labels_in_seq, - compute_metrics=metrics, - ) - - -def fused(specs): - def run(): - return monolithic_head_loss_forward_backward(logits, specs, group=None, grad_logits=None) - - return run - - -def bench(run): - for _ in range(args.warmup): - run() - torch.cuda.synchronize() - torch.cuda.empty_cache() - before = torch.cuda.memory_allocated() - torch.cuda.reset_peak_memory_stats() - times = [] - for _ in range(args.iters): - start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) - start.record() - run() - end.record() - torch.cuda.synchronize() - times.append(start.elapsed_time(end)) - times.sort() - return times[len(times) // 2], (torch.cuda.max_memory_allocated() - before) / 1e9 - - -scenarios = [ - ("CE", [part_ce], [spec_ce()]), - ("z_loss", [part_z], [spec_z()]), - ("distillation", [part_distill], [spec_distill()]), - ("GRPO", [part_grpo], [spec_grpo()]), - ("CE+z", [part_ce, part_z], [spec_ce(), spec_z()]), - ("CE+distill", [part_ce, part_distill], [spec_ce(), spec_distill()]), - ("CE+z+distill", [part_ce, part_z, part_distill], [spec_ce(), spec_z(), spec_distill()]), - ("GRPO+metrics", [part_grpo, part_grpo_metrics], [spec_grpo(metrics=True)]), - ("GRPO+metrics+z", [part_grpo, part_grpo_metrics, part_z], [spec_grpo(metrics=True), spec_z()]), -] - -print(f"shape: num_tokens={N} vocab={V} dtype={DTYPE} | iters={args.iters}") -print( - f"{'scenario':<18}{'per_loss ms':>12}{'fused ms':>11}{'speedup':>9}{'per_loss GB':>13}{'fused GB':>10}{'mem saved':>11}" -) -for name, parts, specs in scenarios: - torch.cuda.empty_cache() - t_pl, m_pl = bench(per_loss(parts)) - torch.cuda.empty_cache() - t_f, m_f = bench(fused(specs)) - print(f"{name:<18}{t_pl:>12.2f}{t_f:>11.2f}{t_pl / t_f:>8.2f}x{m_pl:>13.2f}{m_f:>10.2f}{(m_pl - m_f):>10.2f}G") From c8fb19f349eb9d80b4ede9d5279402e9463a6c9f Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 3 Jul 2026 10:41:24 -0400 Subject: [PATCH 21/28] Guard label reverse-KL / monolithic use_triton; test monolithic under TP (#507) - Reject `loss_type: reverse_kl` for a `label` loss at config time (labels are one-hot, so it is undefined; the triton path already asserted this, the compiled path silently computed cross-entropy). - Reject a `use_triton` flag on a monolithic child loss, which the fused path ignores. - Add a monolithic composite test (cross-entropy + z-loss sharing one softmax) checked against the two losses run standalone, covering the tensor-parallel distributed path where several per-loss all-reduces meet one shared softmax. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/layers/language_model/loss/config.py | 8 +++ tests/layers/test_lm_losses.py | 66 +++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 4222469cc..137628b28 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -104,6 +104,12 @@ class LanguageModelLabelEntropyLossConfig(LanguageModelLossConfig): hint=FieldHint.expert, ) + def _validate(self) -> None: + super()._validate() + # Labels are one-hot, so their forward-KL reduces to cross-entropy; reverse-KL is not defined. + if self.loss_type == EntropyLossType.reverse_kl: + raise ValueError("`reverse_kl` is not supported for a label loss.") + @classmethod def _from_dict(cls, default: dict[str, typing.Any], strict: bool = True) -> typing.Self: if "implementation" in default: @@ -299,6 +305,8 @@ def _validate(self) -> None: raise ValueError( f"Loss `{name}` (`{type(loss).__name__}`) is not combinable and cannot share the softmax." ) + if getattr(loss, "use_triton", None) is not None: + raise ValueError(f"Loss `{name}` sets `use_triton`, which has no effect on a fused child loss.") # A single softmax serves one effective scale (stacked with the common model scale). Assert.eq(len({loss.logits_scale_factor for loss in self.losses.values()}), 1) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 6b7502dba..70450e667 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -23,6 +23,7 @@ ) from fast_llm.layers.language_model.loss.dpo import dpo_loss from fast_llm.layers.language_model.loss.loss import loss_forward_backward +from fast_llm.layers.language_model.loss.monolithic import _monolithic_core from fast_llm.layers.language_model.loss.policy_gradient import ( GRPOMetrics, compute_grpo_metrics, @@ -595,6 +596,44 @@ def _test_z_loss( ) +def _test_monolithic_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate, group=None +): + # A cross-entropy (labels) + z-loss composite sharing one softmax, checked against the same two losses run + # standalone. This exercises the shared, tensor-parallel-reduced softmax and the fp32 gradient accumulation + # against the already-validated single-loss path. + logits, target, _ = _get_lm_loss_inputs(num_columns, loss_masking, TargetFormat.labels, batch_shape, dtype) + local_logits = split_op(logits, group, -1).contiguous() + divisor = max(int((target >= 0).sum().item()), 1) + children = ( + _combinable_loss(LanguageModelLabelEntropyLossConfig(), "cross_entropy", logits_scale_factor), + _combinable_loss(LanguageModelZLossConfig(), "z_loss", logits_scale_factor), + ) + arguments = ((target, grad_output, divisor), (None, grad_output, local_logits.shape[:-1].numel())) + previous_grad = torch.randn_like(local_logits) if accumulate else None + + # Reference: run each loss standalone, accumulating into one gradient buffer. + grad_ref = previous_grad.clone() if accumulate else None + losses_ref = [] + for child, child_arguments in zip(children, arguments, strict=True): + loss_ref, grad_ref, _ = child.combinable_forward_backward(local_logits, group, grad_ref, child_arguments) + losses_ref.append(loss_ref) + + results, grad_fused = _monolithic_core( + children, local_logits, group, logits_scale_factor, previous_grad.clone() if accumulate else None, arguments + ) + + threshold = 1e-5 if dtype == DataType.float32 else 1e-4 + for (loss_fused, _), loss_ref in zip(results, losses_ref, strict=True): + Assert.rms_close_relative(loss_fused, loss_ref, threshold, 1e-6) + if grad_output is None: + assert grad_fused is None and grad_ref is None + else: + # The composite sums child gradients in fp32 and casts once; the standalone path casts each child + # gradient before adding. In fp16 the two differ by up to a rounding step, so allow a wider abs floor. + Assert.rms_close_relative(grad_fused, grad_ref, threshold, 1e-8 if grad_fused.dtype == torch.float32 else 1e-6) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -643,6 +682,18 @@ def test_z_loss( ) +@pytest.mark.slow +@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) +@pytest.mark.parametrize( + ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), + _LOSS_PARAMETERS, +) +def test_monolithic_loss( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate +): + _test_monolithic_loss(batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, accumulate) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -813,6 +864,20 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa compute_entropy, test_context.group, ) + # Monolithic composite: multiple losses share one tensor-parallel-reduced softmax. + with test_context.subtest(base_path, f"monolithic-{suffix}", 2) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_monolithic_loss( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + accumulate, + test_context.group, + ) @pytest.mark.slow @@ -856,6 +921,7 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): "gspo", "grpo_metrics-False", "grpo_metrics-True", + "monolithic", ), ) def test_lm_loss_distributed( From 71421fb10b9f14ee14f2da70990728cb8e477716 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 3 Jul 2026 13:27:19 -0400 Subject: [PATCH 22/28] Fine-review nits: type hints, strict zip, docstring fix (#507) Drop unnecessary forward-ref quotes and parameterize the core's list return type; add strict=True to the child zip; fix the log-prob docstring (/ -> -). Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/entropy_loss.py | 2 +- fast_llm/layers/language_model/loss/monolithic.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index 457647796..a10881a68 100644 --- a/fast_llm/functional/entropy_loss.py +++ b/fast_llm/functional/entropy_loss.py @@ -226,7 +226,7 @@ def predicted_logits_from_labels( May also recover probabilities or log probabilities: `predicted_probabilities = predicted_logits.exp() / sum_exp_logits` - `predicted_log_probabilities = predicted_logits / sum_exp_logits.log()` + `predicted_log_probabilities = predicted_logits - sum_exp_logits.log()` Un-compiled core, inlined into a `@torch.compile` boundary; `fused_predicted_logits_from_labels` is the compiled standalone wrapper. diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index a701c22f1..4cc375afd 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -13,13 +13,13 @@ @torch.compile def _monolithic_core( - children: tuple["LanguageModelLoss", ...], + children: tuple[LanguageModelLoss, ...], logits: torch.Tensor, # (*batch, vocab) group: ProcessGroup | None, logits_scale_factor: float, grad_logits: torch.Tensor | None, arguments: tuple[tuple, ...], -) -> tuple[list, torch.Tensor | None]: +) -> tuple[list[tuple[torch.Tensor, typing.Any]], torch.Tensor | None]: """ One shared softmax over the logits, then each child loss's `fused_core` consuming it. The child list is fixed per config, so the loop unrolls inside this single `@torch.compile` boundary and each @@ -29,7 +29,7 @@ def _monolithic_core( logits_norm, exp_logits, sum_exp_logits, logits_max = softmax_base(logits, logits_scale_factor, group) grad = None results = [] - for child, child_arguments in zip(children, arguments): + for child, child_arguments in zip(children, arguments, strict=True): loss, child_grad, extra = child.fused_core( logits_norm, exp_logits, sum_exp_logits, logits_max, group, logits_scale_factor, child_arguments ) From b7a72271b3e768e1cbc7305cbf67c8d1b5359370 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 3 Jul 2026 15:16:27 -0400 Subject: [PATCH 23/28] Skip the gradient term for zero-weight losses (#507) A zero-weight loss contributes an all-zero gradient. Return grad_output None for it so the backward term drops out of both the standalone and fused paths, while its scalar is still logged. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/layers/language_model/loss/loss.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 0c5c3b6ef..28dd3d490 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -102,7 +102,8 @@ def _prepare_target( def _get_grad_output(self, kwargs: dict[str, typing.Any]) -> float | None: grad_output = kwargs.get(LanguageModelKwargs.grad_output) - return None if grad_output is None else grad_output * self._weight + # A zero-weight loss contributes an all-zero gradient; return `None` so the backward term drops out. + return None if grad_output is None or self._weight == 0 else grad_output * self._weight def _get_labels(self, kwargs: dict[str, typing.Any], split_index: int = 0): return self._prepare_target(kwargs[LanguageModelLossKwargs.labels], split_index) From 9254905222c561dfb24d465be35844b0d8bd78e0 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 3 Jul 2026 15:16:29 -0400 Subject: [PATCH 24/28] Add CombinableLossConfig base, drop the combinable ClassVar (#507) Introduce a CombinableLossConfig marker base owning use_triton; the combinable losses inherit it (GRPO via multiple inheritance alongside the policy-gradient base). The monolithic validation now checks isinstance and reads the narrowed use_triton directly. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/layers/language_model/loss/config.py | 52 +++++++------------ 1 file changed, 19 insertions(+), 33 deletions(-) diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 137628b28..751f44779 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -36,8 +36,6 @@ class LanguageModelLossKwargs(BlockKwargs): @config_class(registry=True) class LanguageModelLossConfig(Config): _abstract: typing.ClassVar[bool] = True - # Whether this loss shares the vocabulary softmax via `fused_core`, so it can be fused with others. - combinable: typing.ClassVar[bool] = False weight: float = Field( default=1.0, @@ -88,21 +86,28 @@ def get_reference_models(self) -> set[str]: return set() +@config_class() +class CombinableLossConfig(LanguageModelLossConfig): + """Base for losses that share the vocabulary softmax via `fused_core`, so several can be fused together.""" + + _abstract: typing.ClassVar[bool] = True + + use_triton: bool | None = Field( + default=None, + desc="Enable triton implementation. Default: use if available.", + hint=FieldHint.expert, + ) + + @config_class(dynamic_type={LanguageModelLossConfig: "label"}) -class LanguageModelLabelEntropyLossConfig(LanguageModelLossConfig): +class LanguageModelLabelEntropyLossConfig(CombinableLossConfig): _abstract: typing.ClassVar[bool] = False - combinable: typing.ClassVar[bool] = True loss_type: EntropyLossType = Field( default=EntropyLossType.cross_entropy, desc="Type of loss to use.", hint=FieldHint.core, ) - use_triton: bool | None = Field( - default=None, - desc="Enable triton implementation. Default: use if available.", - hint=FieldHint.expert, - ) def _validate(self) -> None: super()._validate() @@ -125,9 +130,8 @@ def loss_class(self) -> "type[LanguageModelLabelEntropyLoss]": @config_class(dynamic_type={LanguageModelLossConfig: "distillation"}) -class LanguageModelDistillationLossConfig(LanguageModelLossConfig): +class LanguageModelDistillationLossConfig(CombinableLossConfig): _abstract: typing.ClassVar[bool] = False - combinable: typing.ClassVar[bool] = True loss_type: EntropyLossType = Field( default=EntropyLossType.cross_entropy, @@ -145,11 +149,6 @@ class LanguageModelDistillationLossConfig(LanguageModelLossConfig): desc="Temperature for teacher softmax.", valid=check_field(Assert.gt, 0.0), ) - use_triton: bool | None = Field( - default=None, - desc="Enable triton implementation. Default: use if available.", - hint=FieldHint.expert, - ) @classmethod def _from_dict(cls, default: dict[str, typing.Any], strict: bool = True) -> typing.Self: @@ -197,17 +196,10 @@ def get_reference_models(self) -> set[str]: @config_class(dynamic_type={LanguageModelLossConfig: "z_loss"}) -class LanguageModelZLossConfig(LanguageModelLossConfig): +class LanguageModelZLossConfig(CombinableLossConfig): """Z-loss regularization to prevent overconfidence.""" _abstract: typing.ClassVar[bool] = False - combinable: typing.ClassVar[bool] = True - - use_triton: bool | None = Field( - default=None, - desc="Enable triton implementation. Default: use if available.", - hint=FieldHint.expert, - ) @property def loss_class(self) -> "type[LanguageModelZLoss]": @@ -237,17 +229,11 @@ def loss_class(self) -> "type[LanguageModelPolicyGradientLoss]": @config_class(dynamic_type={LanguageModelLossConfig: "grpo"}) -class LanguageModelGRPOLossConfig(LanguageModelPolicyGradientLossConfig): +class LanguageModelGRPOLossConfig(LanguageModelPolicyGradientLossConfig, CombinableLossConfig): """Group-Relative Policy Optimization: per-token IS-ratio clipping.""" _abstract: typing.ClassVar[bool] = False - combinable: typing.ClassVar[bool] = True - use_triton: bool | None = Field( - default=None, - desc="Enable triton implementation. Default: use if available.", - hint=FieldHint.expert, - ) metrics: GRPOMetricsLevel = Field( default=GRPOMetricsLevel.none, desc=( @@ -301,11 +287,11 @@ def _validate(self) -> None: super()._validate() Assert.gt(len(self.losses), 0) for name, loss in self.losses.items(): - if not loss.combinable: + if not isinstance(loss, CombinableLossConfig): raise ValueError( f"Loss `{name}` (`{type(loss).__name__}`) is not combinable and cannot share the softmax." ) - if getattr(loss, "use_triton", None) is not None: + if loss.use_triton is not None: raise ValueError(f"Loss `{name}` sets `use_triton`, which has no effect on a fused child loss.") # A single softmax serves one effective scale (stacked with the common model scale). Assert.eq(len({loss.logits_scale_factor for loss in self.losses.values()}), 1) From 7cf65dd6f049ab53e46f4f7058aba03231dedb1a Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 3 Jul 2026 15:51:04 -0400 Subject: [PATCH 25/28] Stack the composite loss's logits_scale_factor onto its children (#507) Children were built with the raw head scale, dropping the composite's own logits_scale_factor field; pass self._logits_scale_factor so it stacks like the weight field does. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/layers/language_model/loss/monolithic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 4cc375afd..4a11dd231 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -84,7 +84,7 @@ def __init__( prediction_heads=prediction_heads, vocab_parallel=vocab_parallel, num_splits=num_splits, - logits_scale_factor=logits_scale_factor, + logits_scale_factor=self._logits_scale_factor, weight=self._weight, register_loss=children_register, ) From 66fa29ef1d1a855181750ca995b48573af658552 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 3 Jul 2026 15:51:06 -0400 Subject: [PATCH 26/28] Fine-review cleanups: CombinableLoss stubs, drop redundant guard (#507) Declare get_inputs/fused_core on CombinableLoss; drop the redundant losses-not-None guard (metrics-not-None already implies it); remove a no-op f-string prefix. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/layers/language_model/loss/loss.py | 17 ++++++++++++++++- .../language_model/loss/policy_gradient.py | 2 +- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/fast_llm/layers/language_model/loss/loss.py b/fast_llm/layers/language_model/loss/loss.py index 28dd3d490..5fe431e61 100644 --- a/fast_llm/layers/language_model/loss/loss.py +++ b/fast_llm/layers/language_model/loss/loss.py @@ -119,7 +119,7 @@ def _get_reference_model_logits(self, reference_model: str, kwargs: dict[str, ty # The head owns the `.logits`; split on `.losses.` so the name resolves whether this loss sits # directly under the head or nested inside a composite loss (e.g. `MonolithicLoss`). Assert.incl( - logits_name := self.module_name.split(".losses.")[0] + f".logits", + logits_name := self.module_name.split(".losses.")[0] + ".logits", reference_hidden_states := kwargs[f"reference_{reference_model}_hidden_states"], ) # The logits are already sequence-parallel if needed, we don't want to split again. @@ -170,6 +170,21 @@ class CombinableLoss: _logits_scale_factor: float + def get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: + raise NotImplementedError() + + @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[torch.Tensor, torch.Tensor | None, typing.Any]: + raise NotImplementedError() + @torch.compile def combinable_forward_backward( self, diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 2d7b7f2ea..07e758524 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -254,7 +254,7 @@ def fused_core( def register_combinable_extras(self, extra: tuple, kwargs: dict[str, typing.Any], losses: dict | None) -> None: new_logprobs_mean, metrics = extra self._register_new_logprobs(new_logprobs_mean, kwargs, losses) - if metrics is not None and losses is not None: + if metrics is not None: self._register_grpo_metrics(metrics, kwargs, losses) def _register_extra_metrics( From 96a6991be6355bf1fac7858765553ece47dc1662 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 7 Jul 2026 12:02:34 -0400 Subject: [PATCH 27/28] Fuse GSPO into the monolithic compiled loss (#507) (#552) 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 | 100 ++++++++++++++++-- tests/layers/test_lm_head.py | 31 +++++- 5 files changed, 139 insertions(+), 21 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..0dcf7476f 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, @@ -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} diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 69386203e..8ceb30b41 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} @@ -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"): From 9e7dd32c251a1f0aa23b1cfe2f4e1b8e26c2f676 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 8 Jul 2026 11:54:41 -0400 Subject: [PATCH 28/28] Fuse the GSPO segment seam into compiled blocks (#507) (#554) Co-authored-by: Claude Opus 4.8 (1M context) --- .../language_model/loss/policy_gradient.py | 140 ++++++++++++------ 1 file changed, 91 insertions(+), 49 deletions(-) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 0dcf7476f..cdcb5e2be 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -552,6 +552,69 @@ def _gspo_forward_core( return new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask +@torch.compile +def _gspo_segment_weights( + new_log_probs: torch.Tensor, # (*batch,) + loss_mask: torch.Tensor, # (*batch,) bool + advantages: torch.Tensor, # (*batch,) + old_log_probabilities: torch.Tensor, # (*batch,) + num_labels_in_seq: torch.Tensor, # (*batch,) +) -> tuple[torch.Tensor, torch.Tensor]: + """Compiled pre-aggregation block: the per-token contributions to the two per-segment sums, ready + for `index_add_`. Each labeled token contributes `1 / N_d` so all of doc d's tokens sum to 1 (across + SDP/SP ranks too), regardless of how the doc is sharded. Products stay in `new_log_probs.dtype` (fp32) + — casting to a possibly-low input dtype before the segment sum would round each contribution.""" + log_ratio = (new_log_probs - old_log_probabilities).reshape(-1) + mean_token_weight = loss_mask.reshape(-1).to(log_ratio.dtype) / num_labels_in_seq.reshape(-1).to( + log_ratio.dtype + ).clamp(min=1) + return log_ratio * mean_token_weight, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight + + +@torch.compile +def _gspo_segment_loss( + mean_log_ratio_per_segment: torch.Tensor, # (num_segments,) + mean_advantage_per_segment: torch.Tensor, # (num_segments,) + flat_document_index: torch.Tensor, # (*batch,) int + new_log_probs: torch.Tensor, # (*batch,) + loss_mask: torch.Tensor, # (*batch,) bool + num_labels_in_seq: torch.Tensor, # (*batch,) + epsilon_low: float, + epsilon_high: float, + compute_grad: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Compiled post-aggregation block: from the reduced per-segment sums to the undivided loss sum, the + `new_logprobs` metric, and the unscaled per-token backward coefficient + `clip_factor · loss_weight · R_s`. The `/ divisor` and `grad_output` scaling stay eager so those + per-step-varying scalars never specialize this graph (`epsilon_*` are fixed per run, so they don't).""" + segment_ratio = mean_log_ratio_per_segment.exp() # (num_segments,) — geometric-mean IS ratio + segment_advantage = mean_advantage_per_segment.detach() # (num_segments,) — no grad through A + + probability_ratio = segment_ratio[flat_document_index].reshape(new_log_probs.shape) + advantage_per_token = segment_advantage[flat_document_index].reshape(new_log_probs.shape) + loss_weight = loss_mask.to(new_log_probs.dtype) + + losses = -torch.min( + probability_ratio * advantage_per_token, + torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantage_per_token, + ) + loss_sum = (losses * loss_weight).sum() + new_logprobs_mean = (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() + + if compute_grad: + effective_grad_unscaled = ( + ( + torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high) + + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low) + ) + * loss_weight + * probability_ratio + ) + else: + effective_grad_unscaled = None + return loss_sum, new_logprobs_mean, effective_grad_unscaled + + def gspo_segment_seam( new_log_probs: torch.Tensor, # (*batch,) loss_mask: torch.Tensor, # (*batch,) bool @@ -568,29 +631,20 @@ def gspo_segment_seam( epsilon_high: float, logits_scale_factor: float, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """Eager segment seam between the compiled forward and backward. The `index_add_` segment - aggregation and the symbolic `num_segments` live here so they never enter a compiled boundary - (no per-`num_segments` recompiles). Returns the loss, the `new_logprobs` metric, and the per-token - backward coefficient `effective_grad = grad_output_scaled · clip_factor · loss_weight · R_s` - (None when no gradient is requested).""" - log_ratio = new_log_probs - old_log_probabilities - + """Eager segment seam between the compiled forward and backward, orchestrating two compiled blocks + around the parts that can't compile: the `index_add_` segment aggregation (whose symbolic + `num_segments` would trigger per-value recompiles) and the SDP/SP all-reduces. Returns the loss, the + `new_logprobs` metric, and the per-token backward coefficient + `effective_grad = grad_output_scaled · clip_factor · loss_weight · R_s` (None when no gradient is requested).""" flat_document_index = document_index_zero_based.reshape(-1).long() - flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype) - # Per-token weight: mask / per-document label count, from the preprocessor. - # Each labeled token contributes `1 / N_d` so all of doc d's tokens sum to 1 (across - # SDP/SP ranks too), regardless of how the doc is sharded. - mean_token_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1) - # Pre-divide the per-token contributions by the per-doc label count, then sum per segment. - # All tokens in a segment share the same N_d, so this is mathematically equivalent to - # `log_ratio_sum / N_d` but avoids any per-segment denominator extraction. - mean_log_ratio_per_segment = log_ratio.new_zeros(num_segments).index_add_( - 0, flat_document_index, log_ratio.reshape(-1) * mean_token_weight + weighted_log_ratio, weighted_advantage = _gspo_segment_weights( + new_log_probs, loss_mask, advantages, old_log_probabilities, num_labels_in_seq ) - # Accumulate in `log_ratio.dtype` (fp32). Casting the product back to `advantages.dtype` - # before summing would round each token's contribution to a possibly-low input dtype. - mean_advantage_per_segment = log_ratio.new_zeros(num_segments).index_add_( - 0, flat_document_index, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight + mean_log_ratio_per_segment = weighted_log_ratio.new_zeros(num_segments).index_add_( + 0, flat_document_index, weighted_log_ratio + ) + mean_advantage_per_segment = weighted_advantage.new_zeros(num_segments).index_add_( + 0, flat_document_index, weighted_advantage ) for reduce_group in (sdp_group, sp_group): if reduce_group is not None: @@ -601,34 +655,21 @@ def gspo_segment_seam( mean_advantage_per_segment, op=torch.distributed.ReduceOp.SUM, group=reduce_group ) - segment_ratio = mean_log_ratio_per_segment.exp() # (num_segments,) — geometric-mean IS ratio - segment_advantage = mean_advantage_per_segment.detach() # (num_segments,) — no grad through A - - probability_ratio = segment_ratio[flat_document_index].reshape(log_ratio.shape) - advantage_per_token = segment_advantage[flat_document_index].reshape(log_ratio.shape) - loss_weight = loss_mask.to(log_ratio.dtype) - - losses = -torch.min( - probability_ratio * advantage_per_token, - torch.clamp(probability_ratio, 1 - epsilon_low, 1 + epsilon_high) * advantage_per_token, + loss_sum, new_logprobs_mean, effective_grad_unscaled = _gspo_segment_loss( + mean_log_ratio_per_segment, + mean_advantage_per_segment, + flat_document_index, + new_log_probs, + loss_mask, + num_labels_in_seq, + epsilon_low, + epsilon_high, + grad_output is not None, ) - loss = (losses * loss_weight).sum() / divisor - - new_logprobs_mean = (new_log_probs * loss_mask / num_labels_in_seq.clamp(min=1)).sum() - - if grad_output is None: - return loss, new_logprobs_mean, None - - grad_output_scaled = grad_output / divisor * logits_scale_factor - probability_ratio_grad = ( - grad_output_scaled - * ( - torch.clamp_min(advantage_per_token, 0) * (probability_ratio <= 1 + epsilon_high) - + torch.clamp_max(advantage_per_token, 0) * (probability_ratio >= 1 - epsilon_low) - ) - * loss_weight + loss = loss_sum / divisor + effective_grad = ( + None if grad_output is None else effective_grad_unscaled * (grad_output / divisor * logits_scale_factor) ) - effective_grad = probability_ratio_grad * probability_ratio return loss, new_logprobs_mean, effective_grad @@ -659,8 +700,9 @@ def gspo_backward_core( return grad_logits -# Orchestrator only: the eager `index_add_` segment seam (with the Python-int `num_segments`) sits -# between the compiled forward and backward cores, so it stays out of every compiled boundary. +# Orchestrator only: between the compiled forward and backward cores, the segment seam keeps its +# `index_add_` (with the Python-int `num_segments`) and SDP/SP all-reduces eager, bracketing them with +# compiled sub-blocks, so `num_segments` never enters a compiled boundary. def fused_gspo_loss_forward_backward( logits: torch.Tensor, # (*batch, vocab) target: torch.Tensor, # (*batch,)