From 8c9f052f6c763872a332478986201b21c85efc53 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 17:03:57 -0400 Subject: [PATCH] Triton monolithic kernel + triton policy metrics (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the triton implementation of the monolithic shared-softmax loss and its policy-metric support, on top of the compiled path: - The triton kernel fuses the label-based objective set (CE, z-loss, GRPO, GSPO) over one softmax, superposing their gradients. - Policy metrics: the kernel accumulates the entropy's `Σ exp·logits_norm` for free in the backward re-stream; the layer reduces the family eagerly. Metrics force the reduced forward pass so the softmax is available in every case. - `MonolithicLossConfig.use_triton` selects the path (default: use if available), with per-kind fusion limits validated at config time. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/triton/monolithic_loss.py | 330 ++++++++++++++++++ fast_llm/layers/language_model/loss/config.py | 22 ++ .../layers/language_model/loss/monolithic.py | 137 ++++++++ .../language_model/loss/policy_gradient.py | 63 ++++ tests/layers/test_lm_head.py | 74 +++- tests/layers/test_lm_losses.py | 91 +++++ 6 files changed, 701 insertions(+), 16 deletions(-) create mode 100644 fast_llm/functional/triton/monolithic_loss.py diff --git a/fast_llm/functional/triton/monolithic_loss.py b/fast_llm/functional/triton/monolithic_loss.py new file mode 100644 index 000000000..61dd7d93c --- /dev/null +++ b/fast_llm/functional/triton/monolithic_loss.py @@ -0,0 +1,330 @@ +import torch + +from fast_llm.core.distributed import ReduceOp, all_reduce +from fast_llm.functional.triton import tl, tl_arange, tl_constexpr, triton, triton_jit +from fast_llm.functional.triton.entropy_loss import ( + parallel_sum_exp_logits, + triton_cross_entropy_forward_from_labels_parallel_kernel, + triton_fused_softmax_base, +) +from fast_llm.functional.utils import reduce_losses + + +@triton_jit() +def triton_monolithic_loss_forward_backward_kernel( + logits_ptr, + labels_ptr, + n_cols: tl_constexpr, + logits_stride_0: tl_constexpr, + block_size: tl_constexpr, + ce_losses_ptr=None, + z_losses_ptr=None, + grpo_losses_ptr=None, + new_logprobs_mean_parts_ptr=None, + z_loss_mask_ptr=None, + advantages_ptr=None, + old_log_probs_ptr=None, + num_labels_in_seq_ptr=None, + gspo_coeff_ptr=None, + max_logits_ptr=None, + sum_exp_logits_ptr=None, + predicted_logits_ptr=None, + weighted_logits_sum_ptr=None, + grad_logits_ptr=None, + grad_logits_stride_0: tl_constexpr = None, + grad_losses_ce=0.0, + grad_losses_z=0.0, + grad_losses_grpo=0.0, + col_min: tl_constexpr = 0, + logits_scale_factor: tl_constexpr = 1.0, + epsilon_low: tl_constexpr = 0.2, + epsilon_high: tl_constexpr = 0.2, + accumulate: tl_constexpr = False, +): + """One shared softmax feeding several label-based losses over the same logits row. Each enabled loss + (selected by the presence of its output/input pointers) stores its own forward scalar, but their + gradients superpose into two per-row coefficients: `grad_j = prob_coeff * softmax_j - label_coeff * + delta_{j, label}`. The softmax is computed in-kernel when `max_logits_ptr`/`sum_exp_logits_ptr` are + absent (single-pass, no tensor parallelism), or loaded from a reduced forward pass otherwise.""" + block_idx = tl.program_id(0).to(tl.int64) + logits_ptr = logits_ptr + block_idx * logits_stride_0 + + # The shared label feeds cross-entropy, GRPO, and GSPO; `labels_ptr` is set whenever any of them is + # present. The defaults keep both variables defined on the label-free path — triton compiles every branch, + # so a variable used later must be defined on all of them. + label_valid = False + label_idx = 0 + if labels_ptr is not None: + label_idx = tl.load(labels_ptr + block_idx) + label_valid = label_idx >= 0 + label_idx -= col_min + + if max_logits_ptr is None or sum_exp_logits_ptr is None: + exp_logits, sum_exp_logits, max_logits, col_offsets, mask = triton_fused_softmax_base( + logits_ptr, n_cols=n_cols, block_size=block_size, logits_scale_factor=logits_scale_factor + ) + else: + max_logits = tl.load(max_logits_ptr + block_idx) + sum_exp_logits = tl.load(sum_exp_logits_ptr + block_idx) + + log_sum_exp_logits = tl.log(sum_exp_logits) + max_logits + + # Target-index logit shared by cross-entropy and GRPO; the defaults keep it defined when neither is present. + predicted_logit = 0.0 + new_log_prob = 0.0 + if ce_losses_ptr is not None or grpo_losses_ptr is not None: + if predicted_logits_ptr is not None: + predicted_logit = tl.load(predicted_logits_ptr + block_idx) + elif label_valid and label_idx >= 0 and label_idx < n_cols: + predicted_logit = tl.load(logits_ptr + label_idx).to(tl.float32) + if logits_scale_factor != 1.0: + predicted_logit *= logits_scale_factor + else: + predicted_logit = 0.0 + new_log_prob = predicted_logit - log_sum_exp_logits + + prob_coeff = 0.0 + label_coeff = 0.0 + + if ce_losses_ptr is not None: + if label_valid: + tl.store(ce_losses_ptr + block_idx, log_sum_exp_logits - predicted_logit) + grad_losses = grad_losses_ce * logits_scale_factor if logits_scale_factor != 1.0 else grad_losses_ce + prob_coeff += grad_losses + label_coeff += grad_losses + else: + tl.store(ce_losses_ptr + block_idx, 0.0) + + if z_losses_ptr is not None: + if z_loss_mask_ptr is None or tl.load(z_loss_mask_ptr + block_idx) != 0: + tl.store(z_losses_ptr + block_idx, log_sum_exp_logits * log_sum_exp_logits) + grad_losses = grad_losses_z * logits_scale_factor if logits_scale_factor != 1.0 else grad_losses_z + prob_coeff += 2.0 * grad_losses * log_sum_exp_logits + else: + tl.store(z_losses_ptr + block_idx, 0.0) + + if grpo_losses_ptr is not None: + if label_valid: + old_log_prob = tl.load(old_log_probs_ptr + block_idx).to(tl.float32) + advantage = tl.load(advantages_ptr + block_idx).to(tl.float32) + ratio = tl.exp(new_log_prob - old_log_prob) + clipped_ratio = tl.minimum(tl.maximum(ratio, 1.0 - epsilon_low), 1.0 + epsilon_high) + tl.store(grpo_losses_ptr + block_idx, -tl.minimum(ratio * advantage, clipped_ratio * advantage)) + grad_losses = grad_losses_grpo * logits_scale_factor if logits_scale_factor != 1.0 else grad_losses_grpo + # clip_factor = clamp_min(A, 0) * (ratio <= 1 + eps_h) + clamp_max(A, 0) * (ratio >= 1 - eps_l) + coeff = ( + tl.maximum(advantage, 0.0) * (ratio <= 1.0 + epsilon_high) + + tl.minimum(advantage, 0.0) * (ratio >= 1.0 - epsilon_low) + ) * (ratio * grad_losses) + prob_coeff += coeff + label_coeff += coeff + if new_logprobs_mean_parts_ptr is not None: + num_labels = tl.load(num_labels_in_seq_ptr + block_idx).to(tl.float32) + tl.store(new_logprobs_mean_parts_ptr + block_idx, new_log_prob / tl.maximum(num_labels, 1.0)) + else: + tl.store(grpo_losses_ptr + block_idx, 0.0) + if new_logprobs_mean_parts_ptr is not None: + tl.store(new_logprobs_mean_parts_ptr + block_idx, 0.0) + + if gspo_coeff_ptr is not None: + # The GSPO per-token coefficient is fully scaled by its eager segment seam. + coeff = tl.load(gspo_coeff_ptr + block_idx).to(tl.float32) + prob_coeff += coeff + label_coeff += coeff + + if grad_logits_ptr is not None: + weighted_logits_sum = 0.0 + col_offset_start: tl.constexpr = (n_cols - 1) // block_size * block_size + for col_offset in tl.static_range(col_offset_start, -1, -block_size): + if max_logits_ptr is not None or sum_exp_logits_ptr is not None or col_offset != col_offset_start: + col_offsets = tl_arange(col_offset, col_offset + block_size) + mask = col_offsets < n_cols + logits = tl.load(logits_ptr + col_offsets, mask=mask, other=-float("inf")).to(tl.float32) + if logits_scale_factor != 1.0: + logits *= logits_scale_factor + exp_logits = tl.exp(logits - max_logits) + if weighted_logits_sum_ptr is not None: + # Local (per-rank) Σ exp·logits_norm feeding the policy entropy; the caller all-reduces + # over the vocab group. Zero `logits_norm` on masked columns first (they load as -inf), + # so the product is 0 there rather than 0·-inf = nan. `max_logits` is final here (this is + # the backward re-stream), so the accumulation needs no online rescaling. + weighted_logits_sum += tl.sum(exp_logits * tl.where(mask, logits - max_logits, 0.0), 0) + grad_logits = prob_coeff * (exp_logits / sum_exp_logits) + if label_valid: + grad_logits = tl.where(col_offsets == label_idx, grad_logits - label_coeff, grad_logits) + grad_logits_col_ptr = grad_logits_ptr + block_idx * grad_logits_stride_0 + col_offsets + if accumulate: + grad_logits += tl.load(grad_logits_col_ptr, mask=mask) + tl.store(grad_logits_col_ptr, grad_logits, mask=mask) + if weighted_logits_sum_ptr is not None: + tl.store(weighted_logits_sum_ptr + block_idx, weighted_logits_sum) + + +def _monolithic_forward_reduce( + logits: torch.Tensor, + labels: torch.Tensor | None, + group: torch.distributed.ProcessGroup | None, + logits_scale_factor: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Explicit forward pass: per-row softmax and target-index logit, reduced over the vocab group to the + global values the monolithic kernel then loads. Reuses the shared cross-entropy forward kernel. Used for + tensor parallelism and — regardless of parallelism — to feed the GSPO segment seam before the backward.""" + n_rows = logits.shape[:-1].numel() + n_cols = logits.size(-1) + block_size = min(triton.next_power_of_2(n_cols), 32768) + num_warps = 4 if block_size < 2048 else (8 if block_size < 8192 else 16) + local_max_logits = torch.empty(n_rows, dtype=torch.float, device=logits.device) + sum_exp_logits = torch.empty_like(local_max_logits) + predicted_logits = torch.empty_like(local_max_logits) if labels is not None else None + triton_cross_entropy_forward_from_labels_parallel_kernel[(n_rows,)]( + logits, + labels, + max_logits_ptr=local_max_logits, + sum_exp_logits_ptr=sum_exp_logits, + predicted_logits_ptr=predicted_logits, + col_min=n_cols * group.rank() if group is not None else 0, + n_cols=n_cols, + logits_stride_0=logits.stride(-2), + block_size=block_size, + num_warps=num_warps, + logits_scale_factor=logits_scale_factor, + ) + if group is None: + return local_max_logits, sum_exp_logits, predicted_logits + max_logits, sum_exp_logits = parallel_sum_exp_logits(sum_exp_logits, local_max_logits, group) + if predicted_logits is not None: + all_reduce(predicted_logits, op=ReduceOp.SUM, group=group) + return max_logits, sum_exp_logits, predicted_logits + + +def triton_monolithic_loss_forward_backward( + logits: torch.Tensor, # (*batch, vocab) + labels: torch.Tensor | None, # (*batch,) — shared by cross-entropy / GRPO / GSPO + grad_logits: torch.Tensor | None, + logits_scale_factor: float, + group: torch.distributed.ProcessGroup | None, + divisor: float, + *, + ce: tuple[float | None] | None = None, # cross-entropy: (weighted grad_output,); `None` => absent + z: tuple[torch.Tensor | None, float | None] | None = None, # (loss_mask, weighted grad_output) + # GRPO: (advantages, old_log_probabilities, weighted grad_output, epsilon_low, epsilon_high, num_labels_in_seq) + grpo: tuple | None = None, + gspo_coeff: torch.Tensor | None = None, # per-token backward coefficient from the eager segment seam + softmax: tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] | None = None, # precomputed (max, sum, predicted) + compute_metrics: bool = False, # also emit the reduced softmax and Σ exp·logits_norm for policy metrics + block_size: int | None = None, + num_warps: int | None = None, +) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] | None, + torch.Tensor | None, +]: + """Dispatch the monolithic label-loss kernel over a single shared softmax. Returns the reduced per-loss + scalars `(cross_entropy, z, grpo, grpo_new_logprobs_mean)` (each `None` when the loss is absent), the + accumulated `grad_logits`, and — for policy metrics — the reduced `softmax` (max, sum, predicted) and the + all-reduced per-row `weighted_logits_sum` (`Σ exp·logits_norm`, both `None` unless `compute_metrics`). + `softmax` is provided by the caller when the forward was already run (the GSPO seam); otherwise it is + computed in-kernel (`group is None`, no metrics) or by a reduced forward pass (tensor parallel or metrics). + A present loss whose weighted grad_output is `None` still emits its forward scalar but no gradient term.""" + assert logits.is_contiguous() + if labels is not None: + assert labels.is_contiguous() + n_rows = logits.shape[:-1].numel() + n_cols = logits.size(-1) + if block_size is None: + block_size = min(triton.next_power_of_2(n_cols), 32768) + if num_warps is None: + num_warps = 4 if block_size < 2048 else (8 if block_size < 8192 else 16) + + # Forward-scalar buffers (needed for the total loss even when nothing is registered this step). + ce_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) if ce is not None else None + z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) if z is not None else None + grpo_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) if grpo is not None else None + + ce_grad_output = ce[0] if ce is not None else None + z_loss_mask, z_grad_output = z if z is not None else (None, None) + if grpo is not None: + advantages, old_log_probabilities, grpo_grad_output, epsilon_low, epsilon_high, num_labels_in_seq = grpo + else: + advantages = old_log_probabilities = num_labels_in_seq = grpo_grad_output = None + epsilon_low = epsilon_high = 0.2 + new_logprobs_mean_parts = ( + torch.empty(n_rows, dtype=torch.float, device=logits.device) + if grpo is not None and num_labels_in_seq is not None + else None + ) + for tensor in (z_loss_mask, advantages, old_log_probabilities, num_labels_in_seq, gspo_coeff): + if tensor is not None: + assert tensor.is_contiguous() + + # Metrics reuse the reduced softmax, so run the explicit forward pass even without tensor parallelism. + if softmax is None and (group is not None or compute_metrics): + softmax = _monolithic_forward_reduce(logits, labels, group, logits_scale_factor) + max_logits, sum_exp_logits, predicted_logits = softmax if softmax is not None else (None, None, None) + + has_grad = ( + ce_grad_output is not None + or z_grad_output is not None + or grpo_grad_output is not None + or gspo_coeff is not None + ) + # The entropy's `Σ exp·logits_norm` is accumulated for free in the backward pass, so metrics need it. + assert has_grad or not compute_metrics + weighted_logits_sum = torch.empty(n_rows, dtype=torch.float, device=logits.device) if compute_metrics else None + if has_grad: + accumulate = grad_logits is not None + grad_logits = torch.empty_like(logits) if grad_logits is None else grad_logits + backward_kwargs = { + "grad_logits_ptr": grad_logits, + "grad_logits_stride_0": grad_logits.stride(-2), + "accumulate": accumulate, + "grad_losses_ce": 0.0 if ce_grad_output is None else ce_grad_output / divisor, + "grad_losses_z": 0.0 if z_grad_output is None else z_grad_output / divisor, + "grad_losses_grpo": 0.0 if grpo_grad_output is None else grpo_grad_output / divisor, + "weighted_logits_sum_ptr": weighted_logits_sum, + } + else: + backward_kwargs = {} + + triton_monolithic_loss_forward_backward_kernel[(n_rows,)]( + logits, + labels, + n_cols=n_cols, + logits_stride_0=logits.stride(-2), + block_size=block_size, + num_warps=num_warps, + logits_scale_factor=logits_scale_factor, + col_min=n_cols * group.rank() if group is not None else 0, + epsilon_low=epsilon_low, + epsilon_high=epsilon_high, + ce_losses_ptr=ce_losses, + z_losses_ptr=z_losses, + grpo_losses_ptr=grpo_losses, + new_logprobs_mean_parts_ptr=new_logprobs_mean_parts, + z_loss_mask_ptr=z_loss_mask, + advantages_ptr=advantages, + old_log_probs_ptr=old_log_probabilities, + num_labels_in_seq_ptr=num_labels_in_seq, + gspo_coeff_ptr=gspo_coeff, + max_logits_ptr=max_logits, + sum_exp_logits_ptr=sum_exp_logits, + predicted_logits_ptr=predicted_logits, + **backward_kwargs, + ) + + if weighted_logits_sum is not None and group is not None: + all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group) + + return ( + None if ce_losses is None else reduce_losses(ce_losses, divisor), + None if z_losses is None else reduce_losses(z_losses, divisor), + None if grpo_losses is None else reduce_losses(grpo_losses, divisor), + None if new_logprobs_mean_parts is None else new_logprobs_mean_parts.sum(), + grad_logits, + softmax, + weighted_logits_sum, + ) diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 9d0fa7892..092a54310 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -276,6 +276,11 @@ class MonolithicLossConfig(LanguageModelLossConfig): desc="The combinable losses sharing a single softmax pass. They must agree on `logits_scale_factor`.", hint=FieldHint.core, ) + use_triton: bool | None = Field( + default=None, + desc="Enable the triton implementation of the shared-softmax kernel. Default: use if available.", + hint=FieldHint.expert, + ) def _validate(self) -> None: super()._validate() @@ -289,6 +294,23 @@ def _validate(self) -> 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) + if self.use_triton: + # The triton kernel has one slot per loss kind, so it fuses at most one of each. + seen_kinds = set() + for name, loss in self.losses.items(): + if isinstance(loss, LanguageModelLabelEntropyLossConfig): + kind = "label" + elif isinstance(loss, LanguageModelZLossConfig): + kind = "z_loss" + elif isinstance(loss, LanguageModelGSPOLossConfig): + kind = "gspo" + elif isinstance(loss, LanguageModelGRPOLossConfig): + kind = "grpo" + else: + raise ValueError(f"Loss `{name}` (`{type(loss).__name__}`) has no triton fused implementation.") + if kind in seen_kinds: + raise ValueError(f"The triton path fuses at most one `{kind}` loss; `{name}` is a duplicate.") + seen_kinds.add(kind) @property def loss_class(self) -> "type[LanguageModelLoss]": diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 2ce13c9ae..4140bca21 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -5,6 +5,7 @@ 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.config import TritonConfig 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 @@ -93,6 +94,35 @@ def __init__( ) # 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 + # `(cross_entropy, z, grpo, gspo)` children when every child has a triton fused kernel, else `None`. + self._triton_children = self._classify_triton_children() + + def _classify_triton_children(self) -> tuple | None: + from fast_llm.layers.language_model.loss.entropy_loss import LanguageModelLabelEntropyLoss + from fast_llm.layers.language_model.loss.policy_gradient import LanguageModelGRPOLoss, LanguageModelGSPOLoss + from fast_llm.layers.language_model.loss.z_loss import LanguageModelZLoss + + ce = z = grpo = gspo = None + for child in self._children: + if isinstance(child, LanguageModelGSPOLoss): + if gspo is not None: + return None + gspo = child + elif isinstance(child, LanguageModelGRPOLoss): + if grpo is not None: + return None + grpo = child + elif isinstance(child, LanguageModelZLoss): + if z is not None: + return None + z = child + elif isinstance(child, LanguageModelLabelEntropyLoss): + if ce is not None: + return None + ce = child + else: + return None + return ce, z, grpo, gspo def forward_backward( self, @@ -102,6 +132,8 @@ def forward_backward( split_index: int = 0, grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor | None, torch.Tensor | None]": + if self._triton_children is not None and TritonConfig.enabled(logits.device, self._config.use_triton): + return self._triton_forward_backward(logits, kwargs, losses, split_index, grad_logits) 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 @@ -121,6 +153,111 @@ def forward_backward( total_loss = weighted if total_loss is None else total_loss + weighted return total_loss, grad_logits + def _triton_forward_backward( + self, + logits: torch.Tensor, + kwargs: dict[str, typing.Any], + losses: dict | None, + split_index: int, + grad_logits: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + from fast_llm.functional.triton.monolithic_loss import ( + _monolithic_forward_reduce, + triton_monolithic_loss_forward_backward, + ) + from fast_llm.layers.language_model.loss.config import PolicyMetricsLevel + + register = losses is not None + ce, z, grpo, gspo = self._triton_children + group = self._parallel_dim.group if self._vocab_parallel else None + + labels = divisor = ce_arg = z_arg = grpo_arg = None + if ce is not None: + labels, ce_grad_output, divisor = ce.get_inputs(kwargs, split_index, register) + ce_arg = (ce_grad_output,) + if z is not None: + z_loss_mask, z_grad_output, z_divisor = z.get_inputs(kwargs, split_index, register) + z_arg = (z_loss_mask, z_grad_output) + divisor = z_divisor if divisor is None else divisor + if grpo is not None: + grpo_target, advantages, old_log_probabilities, grpo_grad_output, grpo_divisor, *grpo_tail = ( + grpo.get_inputs(kwargs, split_index, register) + ) + epsilon_low, epsilon_high, num_labels_in_seq = grpo_tail[:3] + labels = grpo_target if labels is None else labels + grpo_arg = ( + advantages, + old_log_probabilities, + grpo_grad_output, + epsilon_low, + epsilon_high, + num_labels_in_seq, + ) + divisor = grpo_divisor if divisor is None else divisor + + grpo_metrics_on = grpo is not None and register and grpo._metrics_level != PolicyMetricsLevel.none + gspo_metrics_on = gspo is not None and register and gspo._metrics_level != PolicyMetricsLevel.none + compute_metrics = grpo_metrics_on or gspo_metrics_on + + # GSPO runs its forward and eager segment seam here, then hands its per-token backward coefficient + # (and its already-reduced softmax) to the kernel, which superposes it with the in-kernel losses. + gspo_coeff = softmax = gspo_loss = gspo_new_logprobs = None + if gspo is not None: + gspo_target, _ = gspo.get_inputs(kwargs, split_index, register) + labels = gspo_target if labels is None else labels + softmax = _monolithic_forward_reduce(logits, labels, group, self._softmax_scale_factor) + gspo_loss, gspo_new_logprobs, gspo_coeff = gspo.compute_triton_seam(kwargs, split_index, *softmax) + + ce_loss, z_loss, grpo_loss, grpo_new_logprobs, grad_logits, softmax, weighted_logits_sum = ( + triton_monolithic_loss_forward_backward( + logits, + None if labels is None else labels.contiguous(), + grad_logits, + self._softmax_scale_factor, + group, + 1.0 if divisor is None else divisor, + ce=ce_arg, + z=z_arg, + grpo=grpo_arg, + gspo_coeff=gspo_coeff, + softmax=softmax, + compute_metrics=compute_metrics, + ) + ) + + # Metrics reuse the kernel's shared softmax: per-token new log-probs and the entropy from the + # kernel's `Σ exp·logits_norm`, so no second softmax pass. + grpo_metrics = gspo_metrics = None + if compute_metrics: + max_logits, sum_exp_logits, predicted_logits = softmax + log_sum_exp_logits = sum_exp_logits.log() + new_log_probs = predicted_logits - max_logits - log_sum_exp_logits + entropy_per_token = log_sum_exp_logits - weighted_logits_sum / sum_exp_logits + if grpo_metrics_on: + grpo_metrics = grpo.triton_metrics(new_log_probs, entropy_per_token, kwargs, split_index) + if gspo_metrics_on: + gspo_metrics = gspo.triton_metrics(new_log_probs, entropy_per_token, kwargs, split_index) + + results = {} + if ce is not None: + results[ce] = (ce_loss, None) + if z is not None: + results[z] = (z_loss, None) + if grpo is not None: + results[grpo] = (grpo_loss, (grpo_new_logprobs, grpo_metrics)) + if gspo is not None: + results[gspo] = (gspo_loss, (gspo_new_logprobs, gspo_metrics)) + + total_loss = None + for child in self._children: + loss, extra = results[child] + 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)) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index d1112cb24..4b1aa1fd3 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -336,6 +336,27 @@ def _register_extra_metrics( ) self._register_policy_metrics(metrics, kwargs, losses) + def triton_metrics( + self, + new_log_probs: torch.Tensor, # flat, from the kernel's shared softmax + entropy_per_token: torch.Tensor, # flat, from the kernel's `Σ exp·logits_norm` + kwargs: dict[str, typing.Any], + split_index: int, + ) -> PolicyMetrics: + """GRPO metric family from the triton kernel's shared-softmax outputs, reusing `grpo_metrics_core` so + the metrics add no second softmax.""" + target = self._get_labels(kwargs, split_index) + return grpo_metrics_core( + new_log_probs.reshape(target.shape), + self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), + self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index), + target >= 0, + self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index), + self._config.epsilon_low, + self._config.epsilon_high, + entropy_per_token.reshape(target.shape), + ) + def get_loss_definitions(self) -> list[LossDef]: return super().get_loss_definitions() + self._policy_metric_definitions() @@ -582,6 +603,48 @@ def finish( ) return loss, (new_logprobs_mean, metrics), grad_logits + def compute_triton_seam( + self, + kwargs: dict[str, typing.Any], + split_index: int, + max_logits: torch.Tensor, # (n_rows,) + sum_exp_logits: torch.Tensor, # (n_rows,) + predicted_logits: torch.Tensor, # (n_rows,) + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """GSPO's contribution to the triton monolithic kernel: recover per-token new log-probs from the + triton forward pass, run the segment seam, and return the loss, the `new_logprobs` metric, and the + flat per-token backward coefficient (`None` when no gradient is requested) the kernel superposes.""" + target = self._get_labels(kwargs, split_index) + loss_mask = target >= 0 + new_log_probs = (predicted_logits - max_logits - sum_exp_logits.log()).reshape(loss_mask.shape) + loss, new_logprobs_mean, effective_grad = self._run_segment_seam(new_log_probs, loss_mask, kwargs, split_index) + return loss, new_logprobs_mean, None if effective_grad is None else effective_grad.reshape(-1).contiguous() + + def triton_metrics( + self, + new_log_probs: torch.Tensor, # flat, from the kernel's shared softmax + entropy_per_token: torch.Tensor, # flat, from the kernel's `Σ exp·logits_norm` + kwargs: dict[str, typing.Any], + split_index: int, + ) -> PolicyMetrics: + """GSPO segment-level metric family from the triton kernel's shared-softmax outputs, reusing + `gspo_metrics_core` so the metrics add no second softmax.""" + target = self._get_labels(kwargs, split_index) + return gspo_metrics_core( + new_log_probs.reshape(target.shape), + self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), + self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index), + target >= 0, + self._document_index_zero_based(kwargs, split_index), + kwargs[BlockKwargs.num_documents_in_sequence], + self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index), + self._config.epsilon_low, + self._config.epsilon_high, + entropy_per_token.reshape(target.shape), + self._sequence_data_dim.group if self._sequence_data_active else None, + self._parallel_dim.group if self._sequence_parallel else None, + ) + 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) diff --git a/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 9c233765b..5df2e292d 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -6,6 +6,7 @@ import torch from fast_llm.engine.config_utils.data_type import DataType +from fast_llm.functional.triton import triton_available from fast_llm.layers.attention.config import AttentionKwargs from fast_llm.layers.block.config import BlockKwargs from fast_llm.layers.language_model.config import LM_HEAD_LOSS_NAME, LanguageModelKwargs @@ -95,9 +96,10 @@ def get_config(self) -> GPTModelConfig: losses["gspo_loss"] = {"type": "gspo", "metrics": self.gspo_metrics or "none"} if isinstance(self.gspo_loss, float): losses["gspo_loss"]["weight"] = self.gspo_loss - if self.loss_implementation == "fused" and losses: + if self.loss_implementation in ("fused", "fused_triton") 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. + # child keys so the registered metric names match the per-loss configuration. `fused` pins the + # compiled path and `fused_triton` the triton path, so both are exercised in every environment. combinable = { name: loss for name, loss in losses.items() @@ -105,7 +107,11 @@ def get_config(self) -> GPTModelConfig: } if combinable: losses = {name: loss for name, loss in losses.items() if name not in combinable} - losses["monolithic"] = {"type": "monolithic", "losses": combinable} + losses["monolithic"] = { + "type": "monolithic", + "losses": combinable, + "use_triton": self.loss_implementation == "fused_triton", + } if losses: head_config["losses"] = losses @@ -484,11 +490,11 @@ def _add_configs(base_name: str, **kwargs): loss_implementation="fused", ) ) -# GRPO/GSPO metric families (`basic` includes entropy) on the standalone and compiled shared-softmax paths. -# 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_" +# GRPO/GSPO metric families (`basic` includes entropy) across all three paths — standalone, the compiled +# shared softmax, and the triton kernel. 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", "fused_triton"): + _prefix = "" if _loss_implementation == "per_loss" else f"{_loss_implementation}_" for _loss_masking in (False, True): _mask_suffix = "_masked" if _loss_masking else "" _lm_head_test_configs.append( @@ -509,21 +515,55 @@ def _add_configs(base_name: str, **kwargs): loss_implementation=_loss_implementation, ) ) -# The metric family co-resides with z-loss in the shared softmax pass. Single-split (metrics can't be split). +# The metric family co-resides with z-loss in the shared softmax pass, on both fused paths. Single-split +# (metrics can't be split). +for _loss_implementation in ("fused", "fused_triton"): + for _loss_masking in (False, True): + _lm_head_test_configs.append( + LMHeadTestConfig( + f"{_loss_implementation}_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=_loss_implementation, + ) + ) +# `auto` metrics resolve to `basic` when pipeline_parallel == 1 (all head tests), covering the default path. +_lm_head_test_configs.append(LMHeadTestConfig("grpo_loss_metrics_auto", grpo_loss=True, grpo_metrics="auto")) +_lm_head_test_configs.append(LMHeadTestConfig("gspo_loss_metrics_auto", gspo_loss=True, gspo_metrics="auto")) + +# Triton monolithic kernel (`use_triton=True`): the label-based objective set over the shared softmax. +# Distillation has no triton fused kernel, so it stays on the compiled `fused` configs (policy metrics do). +_add_configs("fused_triton", loss_implementation="fused_triton") +_add_configs("fused_triton_z_loss", loss_implementation="fused_triton", z_loss=True) +_add_configs("fused_triton_bfloat16", loss_implementation="fused_triton", compute_dtype=DataType.bfloat16) +_add_configs("fused_triton_logit_scaling", loss_implementation="fused_triton", logits_scale_factor=5.0) +_add_configs("fused_triton_final_logit_softcap", loss_implementation="fused_triton", final_logit_softcap=2.0) +_add_configs("fused_triton_label_and_z_loss_weighted", loss_implementation="fused_triton", label_loss=True, z_loss=0.5) +_add_configs("fused_triton_grpo_loss", loss_implementation="fused_triton", grpo_loss=True) +_add_configs("fused_triton_grpo_and_z_loss", loss_implementation="fused_triton", grpo_loss=True, z_loss=0.5) +# GSPO on the triton path (its eager segment seam brackets the triton forward and backward). Single-split +# only; alone and sharing the softmax with z-loss. for _loss_masking in (False, True): + _suffix = "_masked" if _loss_masking else "" + _lm_head_test_configs.append( + LMHeadTestConfig( + f"fused_triton_gspo_loss{_suffix}", + gspo_loss=True, + loss_masking=_loss_masking, + loss_implementation="fused_triton", + ) + ) _lm_head_test_configs.append( LMHeadTestConfig( - f"fused_grpo_and_z_loss_metrics{'_masked' if _loss_masking else ''}", - grpo_loss=True, + f"fused_triton_gspo_and_z_loss{_suffix}", + gspo_loss=True, z_loss=0.5, - grpo_metrics="basic", loss_masking=_loss_masking, - loss_implementation="fused", + loss_implementation="fused_triton", ) ) -# `auto` metrics resolve to `basic` when pipeline_parallel == 1 (all head tests), covering the default path. -_lm_head_test_configs.append(LMHeadTestConfig("grpo_loss_metrics_auto", grpo_loss=True, grpo_metrics="auto")) -_lm_head_test_configs.append(LMHeadTestConfig("gspo_loss_metrics_auto", gspo_loss=True, gspo_metrics="auto")) @pytest.mark.slow @@ -535,6 +575,8 @@ def _add_configs(base_name: str, **kwargs): ], ) def test_lm_head(test_config: LMHeadTestConfig): + if test_config.loss_implementation == "fused_triton" and not triton_available: + pytest.skip("Triton is not available (build the extension or set TRITON_INTERPRET=1).") model_config = test_config.get_config() model, distributed = get_base_model(model_config) input_, kwargs = test_config.get_inputs() diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 03c0c6bd6..519538194 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -14,6 +14,7 @@ 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.monolithic_loss import triton_monolithic_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, @@ -748,6 +749,67 @@ def _test_monolithic_loss( Assert.rms_close_relative(grad_fused, grad_ref, threshold, 1e-8 if grad_fused.dtype == torch.float32 else 1e-6) +def _test_monolithic_loss_triton( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate, group=None +): + # The triton monolithic kernel (cross-entropy + z-loss over one softmax) against the sum of the standalone + # triton kernels. Both use the same label-count divisor, as the head threads the same divisor to each + # combinable child. Tensor-parallel `group` is passed per call for the distributed subtest (case B). + if not triton_available: + return + 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) + previous_grad = torch.randn_like(local_logits) if accumulate else None + + # Reference: standalone triton cross-entropy then z-loss, accumulating into one gradient buffer. + grad_ref = previous_grad.clone() if accumulate else None + ce_ref, grad_ref = triton_entropy_loss_forward_backward( + local_logits, + target, + None, + grad_logits=grad_ref, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + target_format=TargetFormat.labels, + divisor=divisor, + block_size=block_size, + ) + z_ref, grad_ref = triton_z_loss_forward_backward( + local_logits, + None, + grad_logits=grad_ref, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + divisor=divisor, + block_size=block_size, + ) + + ce_fused, z_fused, _, _, grad_fused, _, _ = triton_monolithic_loss_forward_backward( + local_logits, + target.contiguous(), + previous_grad.clone() if accumulate else None, + logits_scale_factor, + group, + divisor, + ce=(grad_output,), + z=(None, grad_output), + block_size=block_size, + ) + + threshold = 1e-5 if dtype == DataType.float32 else 1e-4 + Assert.rms_close_relative(ce_fused, ce_ref, threshold, 1e-6) + Assert.rms_close_relative(z_fused, z_ref, threshold, 1e-6) + if grad_output is None: + assert grad_fused is None and grad_ref is None + else: + # The kernel superposes both losses' gradients in fp32 before the single cast; the standalone path + # rounds cross-entropy's gradient before z-loss adds to it, so the two differ by up to a rounding step. + 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( @@ -808,6 +870,20 @@ def test_monolithic_loss( _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( + ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "block_size", "accumulate"), + _LOSS_PARAMETERS, +) +def test_monolithic_loss_triton( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate +): + _test_monolithic_loss_triton( + batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate + ) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -1021,6 +1097,21 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa accumulate, test_context.group, ) + # Triton monolithic composite: the vocab-parallel reduced-softmax path (case B). + with test_context.subtest(base_path, f"monolithic-triton-{suffix}", 2) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_monolithic_loss_triton( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + block_size, + accumulate, + test_context.group, + ) @pytest.mark.slow