diff --git a/fast_llm/functional/entropy_loss.py b/fast_llm/functional/entropy_loss.py index 05eaae520..a10881a68 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 @@ -92,8 +91,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 +104,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. + + 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: @@ -121,9 +122,13 @@ def fused_softmax_base( return logits_norm, exp_logits, sum_exp_logits, logits_max -@torch.compile -def _fused_reverse_kl_base_from_distribution( - logits: torch.Tensor, # (*batch, vocab) +fused_softmax_base = torch.compile(softmax_base) + + +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, @@ -131,13 +136,14 @@ 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 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) - 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) @@ -161,9 +167,10 @@ def _fused_reverse_kl_base_from_distribution( return per_sample_loss, grad -@torch.compile -def _fused_cross_entropy_base_from_distribution( - logits: torch.Tensor, # (*batch, vocab) +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, @@ -172,10 +179,10 @@ 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, _ = fused_softmax_base(logits, logits_scale_factor, group) - + """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, _ = 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) @@ -207,8 +214,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 @@ -218,9 +224,12 @@ def fused_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()` + `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. """ if group is None: @@ -243,19 +252,24 @@ def fused_predicted_logits_from_labels( return predicted_logits, target_masked, target_mask -@torch.compile -def _fused_cross_entropy_base_from_labels( - logits: torch.Tensor, # (*batch, vocab) +fused_predicted_logits_from_labels = torch.compile(predicted_logits_from_labels) + + +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. 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) # CE loss = mean(log(sum_exp_logits) - sum(probabilities * logits)) # KL loss is the same because P * log(P) == 0. @@ -274,73 +288,22 @@ def _fused_cross_entropy_base_from_labels( return per_sample_loss, grad -@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]: +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 """ - 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. + 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). Un-compiled core, inlined into a `@torch.compile` boundary. """ - 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, - ) + 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: - 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 + grad = (2 * grad_output * (log_sum_exp_logits / sum_exp_logits)).unsqueeze(-1) * exp_logits + return loss_term, grad diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 9a220aacf..7d7e45ad0 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -86,8 +86,21 @@ 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 loss_type: EntropyLossType = Field( @@ -95,11 +108,12 @@ class LanguageModelLabelEntropyLossConfig(LanguageModelLossConfig): 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() + # 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: @@ -116,7 +130,7 @@ def loss_class(self) -> "type[LanguageModelLabelEntropyLoss]": @config_class(dynamic_type={LanguageModelLossConfig: "distillation"}) -class LanguageModelDistillationLossConfig(LanguageModelLossConfig): +class LanguageModelDistillationLossConfig(CombinableLossConfig): _abstract: typing.ClassVar[bool] = False loss_type: EntropyLossType = Field( @@ -135,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: @@ -187,17 +196,11 @@ 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 - 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]": from fast_llm.layers.language_model.loss.z_loss import LanguageModelZLoss @@ -226,16 +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 - 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=( @@ -255,19 +253,48 @@ 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 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 isinstance(loss, CombinableLossConfig): + raise ValueError( + f"Loss `{name}` (`{type(loss).__name__}`) is not combinable and cannot share the softmax." + ) + 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) + + @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/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 48c1556f3..f4ccd74ff 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -2,17 +2,24 @@ 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, + 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, SingleLoss -class LanguageModelLabelEntropyLoss[ConfigType: LanguageModelLabelEntropyLossConfig](LanguageModelLoss[ConfigType]): +class LanguageModelLabelEntropyLoss[ConfigType: LanguageModelLabelEntropyLossConfig]( + CombinableLoss, SingleLoss[ConfigType] +): def _forward_backward( self, logits: "torch.Tensor", @@ -21,25 +28,56 @@ 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.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 + 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 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 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, None]: + """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_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, SingleLoss[ConfigType] +): def _forward_backward( self, logits: "torch.Tensor", @@ -48,22 +86,80 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor, torch.Tensor | 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 + 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 get_inputs(self, kwargs: dict[str, typing.Any], split_index: int, register: bool) -> tuple: 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, - target_format=TargetFormat.logits, - entropy_loss_type=self._config.loss_type, - divisor=self._get_label_count(kwargs), + self._get_grad_output(kwargs), + self._get_label_count(kwargs), + self._config.loss_type, + self._config.temperature, ) + @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, 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 + 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} 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..2d1b5af52 --- /dev/null +++ b/fast_llm/layers/language_model/loss/grpo_metrics.py @@ -0,0 +1,78 @@ +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. + + 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) + + 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/loss.py b/fast_llm/layers/language_model/loss/loss.py index 397aec966..978f792f4 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 @@ -42,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", @@ -50,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]: @@ -117,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) @@ -130,8 +116,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] + ".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. @@ -140,6 +128,116 @@ 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. Composite losses satisfy `forward_backward` directly rather than through this 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, 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. 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 + + 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, + 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: 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 + ) + return loss, self._accumulate_grad(grad, logits.dtype, grad_logits), extra + + @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).""" + 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 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 ) -> 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 new file mode 100644 index 000000000..2ce13c9ae --- /dev/null +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -0,0 +1,128 @@ +import typing + +import torch + +from fast_llm.core.distributed import ProcessGroup +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 CombinableLoss, 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, + arguments: tuple[tuple, ...], +) -> 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 + `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, strict=True): + 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)) + if child_grad is not None: + grad = child_grad if grad is None else grad + child_grad + return results, CombinableLoss._accumulate_grad(grad, logits.dtype, grad_logits) + + +class MonolithicLoss[ConfigType: MonolithicLossConfig](LanguageModelLoss[ConfigType]): + """ + 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__( + 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, + ) + # 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=self._logits_scale_factor, + weight=self._weight, + register_loss=children_register, + ) + for child_name, child_config in config.losses.items() + ] + ) + # 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.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 + ) + + 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) + 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 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 a024d4232..cdcb5e2be 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -3,11 +3,15 @@ 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 -from fast_llm.functional.entropy_loss import fused_predicted_logits_from_labels, fused_softmax_base +from fast_llm.functional.entropy_loss import ( + 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 from fast_llm.layers.language_model.config import LanguageModelKwargs @@ -18,27 +22,12 @@ LanguageModelLossKwargs, LanguageModelPolicyGradientLossConfig, ) -from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.grpo_metrics import GRPOMetrics, grpo_metrics_core +from fast_llm.layers.language_model.loss.loss import CombinableLoss, SingleLoss 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] -): +class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLossConfig](SingleLoss[ConfigType]): """Shared scaffolding for policy-gradient losses (GRPO, GSPO).""" def _register_new_logprobs( @@ -66,7 +55,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__( @@ -109,38 +100,162 @@ def _forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | 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 - fn = triton_grpo_loss_forward_backward - else: - fn = fused_grpo_loss_forward_backward - loss, grad, new_logprobs_mean = fn( - logits, + ( + 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 + + 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 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 ( 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._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, ) - self._register_new_logprobs(new_logprobs_mean, kwargs, losses) + @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, tuple]: + """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, + old_log_probabilities, + grad_output, + divisor, + epsilon_low, + epsilon_high, + num_labels_in_seq, + compute_metrics, + compute_entropy, + ) = 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() + probability_ratio = (new_log_probs - old_log_probabilities).exp() - # 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) + 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() + ) - return loss, grad + 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` 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, + 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: + new_logprobs_mean, metrics = extra + self._register_new_logprobs(new_logprobs_mean, kwargs, losses) + if metrics is not None: + self._register_grpo_metrics(metrics, kwargs, losses) def _register_extra_metrics( self, @@ -161,7 +276,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 ( @@ -218,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, @@ -251,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", @@ -259,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] @@ -300,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} @@ -318,127 +517,192 @@ 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, ) @torch.compile -def fused_grpo_loss_forward_backward( +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 + + +@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,) - 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 - loss_mask = target >= 0 + 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 - 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() - probability_ratio = (new_log_probs - old_log_probabilities).exp() + +@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 * 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() + 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 grad_output is not None: - # 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) + 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_mask + * loss_weight + * probability_ratio ) + else: + effective_grad_unscaled = None + return loss_sum, new_logprobs_mean, effective_grad_unscaled - # 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) - 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) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) +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, 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() + weighted_log_ratio, weighted_advantage = _gspo_segment_weights( + new_log_probs, loss_mask, advantages, old_log_probabilities, num_labels_in_seq + ) + 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: + 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 + ) - return loss, grad_logits, new_logprobs_mean + 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 = loss_sum / divisor + effective_grad = ( + None if grad_output is None else effective_grad_unscaled * (grad_output / divisor * logits_scale_factor) + ) + return loss, new_logprobs_mean, effective_grad -# 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. +@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: 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,) @@ -479,76 +743,35 @@ def fused_gspo_loss_forward_backward( 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() - 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 + new_log_probs, exp_logits, sum_exp_logits, target_masked, target_mask = _gspo_forward_core( + logits, target, loss_mask, logits_scale_factor, group ) - # 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 + 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, ) - 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_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 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, ) - 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 diff --git a/fast_llm/layers/language_model/loss/z_loss.py b/fast_llm/layers/language_model/loss/z_loss.py index 2e5f90b1d..8ebd23f8e 100644 --- a/fast_llm/layers/language_model/loss/z_loss.py +++ b/fast_llm/layers/language_model/loss/z_loss.py @@ -3,14 +3,14 @@ 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 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 -from fast_llm.layers.language_model.loss.loss import LanguageModelLoss +from fast_llm.layers.language_model.loss.loss import CombinableLoss, SingleLoss -class LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](LanguageModelLoss[ConfigType]): +class LanguageModelZLoss[ConfigType: LanguageModelZLossConfig](CombinableLoss, SingleLoss[ConfigType]): def _forward_backward( self, logits: "torch.Tensor", @@ -19,19 +19,44 @@ 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.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 + 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 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 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, None]: + """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) + 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]: return {"return_prediction_mask": True} @@ -49,38 +74,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() - 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 - - loss = reduce_losses(log_sum_exp_logits**2, divisor, loss_mask) - - if grad_output is not None: - grad_base = 2 * grad_output * (log_sum_exp_logits / sum_exp_logits) - if loss_mask is not None: - grad_base = grad_base * loss_mask - grad = (grad_base.unsqueeze(-1) * exp_logits).to(logits.dtype) - if grad_logits is None: - grad_logits = grad - else: - grad_logits.add_(grad) - - return loss, grad_logits diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 6f09cb108..8ceb30b41 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 @@ -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 @@ -39,6 +40,8 @@ class LMHeadTestConfig: tied_embedding_weight: bool = False 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): @@ -68,6 +71,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: @@ -78,10 +83,23 @@ 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): 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", "gspo") + } + 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 @@ -239,9 +257,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: @@ -265,6 +286,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, @@ -359,6 +403,87 @@ 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 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) +_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) +_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, +) +_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 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"): + _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, + ) + ) +# 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 d7d14ad3e..70450e667 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -6,25 +6,30 @@ 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 +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 ( + 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.monolithic import _monolithic_core 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 @@ -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, @@ -308,7 +315,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, ) @@ -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 dtype == DataType.float32 else 1e-4, group=group, ) @@ -363,17 +370,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 +533,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 +558,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, @@ -561,7 +571,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: @@ -581,18 +591,56 @@ 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, ) +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( ("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, @@ -634,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( @@ -724,7 +784,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( @@ -775,6 +835,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: @@ -789,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 @@ -824,13 +913,15 @@ 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", "grpo", + "gspo", "grpo_metrics-False", "grpo_metrics-True", + "monolithic", ), ) def test_lm_loss_distributed(