From 8a776105c543600f6525d4bfdb9697302215c38d Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 14 Jul 2026 18:15:06 -0400 Subject: [PATCH 1/2] Fuse distillation + z-loss in the triton monolithic path (#507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a distribution-family triton kernel alongside the label-family monolithic kernel: distillation (the general full-distribution loss) with an optional z-loss superposed over the same shared student softmax, so distillation configs run on triton instead of falling back to the compiled path. z-loss folds into the two from-distribution forward/backward kernels as a student-probability coefficient (`2·grad_z·log_sum_exp`, no target term), sharing distillation's loss mask and divisor; gated on its output pointer so it is dead-code-eliminated when absent. `triton_entropy_loss_forward_backward` gains an optional `z` slot and returns `(loss, z_loss, grad_logits)`. `MonolithicLoss` dispatches to the distribution kernel when a distillation child is present ({distillation} or {distillation, z_loss}); other combinations stay on the label kernel or the compiled path. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/triton/entropy_loss.py | 60 ++++++-- fast_llm/layers/language_model/loss/config.py | 10 +- .../language_model/loss/entropy_loss.py | 21 ++- .../layers/language_model/loss/monolithic.py | 62 +++++++- tests/layers/test_lm_head.py | 14 +- tests/layers/test_lm_losses.py | 136 +++++++++++++++++- 6 files changed, 284 insertions(+), 19 deletions(-) diff --git a/fast_llm/functional/triton/entropy_loss.py b/fast_llm/functional/triton/entropy_loss.py index 0cfd60dd8..73453739f 100644 --- a/fast_llm/functional/triton/entropy_loss.py +++ b/fast_llm/functional/triton/entropy_loss.py @@ -362,12 +362,14 @@ def triton_cross_entropy_from_distribution_forward_backward_kernel( loss_mask_ptr=None, partial_losses_ptr=None, losses_ptr=None, + z_losses_ptr=None, max_logits_ptr=None, sum_exp_logits_ptr=None, target_max_logits_ptr=None, target_sum_exp_logits_ptr=None, from_logits: tl_constexpr = True, grad_losses=None, + grad_losses_z=0.0, grad_logits_ptr=None, grad_logits_stride_0: tl_constexpr = None, logits_scale_factor: tl_constexpr = 1.0, @@ -381,9 +383,11 @@ def triton_cross_entropy_from_distribution_forward_backward_kernel( target_ptr = target_ptr + block_idx * target_stride_0 if loss_mask_ptr is not None and tl.load(loss_mask_ptr + block_idx) == 0: - # This entry is masked, ignore. + # This entry is masked, ignore. z-loss shares this mask (both use the loss mask), so it drops too. if losses_ptr is not None: tl.store(losses_ptr + block_idx, 0) + if z_losses_ptr is not None: + tl.store(z_losses_ptr + block_idx, 0) if grad_losses is not None and not accumulate: for col_offset in tl.static_range(0, n_cols, block_size): col_offsets = tl_arange(int(col_offset), int(col_offset + block_size)) @@ -428,10 +432,20 @@ def triton_cross_entropy_from_distribution_forward_backward_kernel( loss = tl.log(sum_exp_logits) + max_logits - predicted_logits tl.store(losses_ptr + block_idx, loss) + if z_losses_ptr is not None: + log_sum_exp_logits = tl.log(sum_exp_logits) + max_logits + tl.store(z_losses_ptr + block_idx, log_sum_exp_logits * log_sum_exp_logits) + if grad_losses is not None: if logits_scale_factor != 1.0: grad_losses *= logits_scale_factor - # grad / grad_output = exp_logits / sum_exp_logits - target_probabilities. + grad_losses_z *= logits_scale_factor + # A fused z-loss superposes onto the student-probability coefficient (it has no target term). Gated on + # its output pointer so the extra term is dead-code-eliminated when no z-loss shares the kernel. + student_coeff = grad_losses + if z_losses_ptr is not None: + student_coeff += 2.0 * grad_losses_z * (tl.log(sum_exp_logits) + max_logits) + # grad / grad_output = student_coeff * exp_logits / sum_exp_logits - grad_losses * target_probabilities. 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): col_offsets = tl_arange(col_offset, col_offset + block_size) @@ -449,7 +463,7 @@ def triton_cross_entropy_from_distribution_forward_backward_kernel( else: target = tl.load(target_ptr + col_offsets, mask=mask, other=-float("inf")).to(tl.float32) - grad_logits = grad_losses * (exp_logits / sum_exp_logits - target) + grad_logits = student_coeff * (exp_logits / sum_exp_logits) - grad_losses * target 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) @@ -593,12 +607,14 @@ def triton_reverse_kl_forward_backward_kernel_from_distribution( loss_mask_ptr=None, partial_losses_ptr=None, losses_ptr=None, + z_losses_ptr=None, max_logits_ptr=None, sum_exp_logits_ptr=None, target_max_logits_ptr=None, target_sum_exp_logits_ptr=None, from_logits: tl_constexpr = True, grad_losses=None, + grad_losses_z=0.0, grad_logits_ptr=None, grad_logits_stride_0: tl_constexpr = None, logits_scale_factor: tl_constexpr = 1.0, @@ -611,9 +627,11 @@ def triton_reverse_kl_forward_backward_kernel_from_distribution( target_ptr = target_ptr + block_idx * target_stride_0 if loss_mask_ptr is not None and tl.load(loss_mask_ptr + block_idx) == 0: - # This entry is masked, ignore. + # This entry is masked, ignore. z-loss shares this mask (both use the loss mask), so it drops too. if losses_ptr is not None: tl.store(losses_ptr + block_idx, 0) + if z_losses_ptr is not None: + tl.store(z_losses_ptr + block_idx, 0) if grad_losses is not None and not accumulate: for col_offset in tl.static_range(0, n_cols, block_size): col_offsets = tl_arange(int(col_offset), int(col_offset + block_size)) @@ -649,9 +667,19 @@ def triton_reverse_kl_forward_backward_kernel_from_distribution( loss = loss + tl.log(target_sum_exp_logits) + target_max_logits tl.store(losses_ptr + block_idx, loss) + if z_losses_ptr is not None: + log_sum_exp_logits = tl.log(sum_exp_logits) + max_logits + tl.store(z_losses_ptr + block_idx, log_sum_exp_logits * log_sum_exp_logits) + if grad_losses is not None: if logits_scale_factor != 1.0: grad_losses *= logits_scale_factor + grad_losses_z *= logits_scale_factor + # A fused z-loss superposes onto the student-probability coefficient (it has no target term). Gated on + # its output pointer so the extra term is dead-code-eliminated when no z-loss shares the kernel. + z_coeff = 0.0 + if z_losses_ptr is not None: + z_coeff = 2.0 * grad_losses_z * (tl.log(sum_exp_logits) + max_logits) 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): col_offsets = tl_arange(col_offset, col_offset + block_size) @@ -671,8 +699,8 @@ def triton_reverse_kl_forward_backward_kernel_from_distribution( predicted_probability = tl.exp(logits - max_logits) / sum_exp_logits predicted_log_probability = logits - max_logits - tl.log(sum_exp_logits) grad_logits = ( - grad_losses * (predicted_log_probability - target_log_probability - loss) * predicted_probability - ) + grad_losses * (predicted_log_probability - target_log_probability - loss) + z_coeff + ) * predicted_probability 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) @@ -749,14 +777,17 @@ def triton_entropy_loss_forward_backward( temperature: float = 1.0, target_format: TargetFormat = TargetFormat.labels, entropy_loss_type: EntropyLossType = EntropyLossType.cross_entropy, + z: tuple[float | None] | None = None, block_size: int | None = None, num_warps: int | None = None, divisor: float | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: """ A fast triton implementation of cross-entropy, which combines the casting and forward and backward passes, all in a single kernel. Compared to a standard pytorch implementation, this reduces memory usage (of logits) by 3x and memory I/O by 5x. + `z` fuses a z-loss over the same (distribution-target) softmax: `(z_grad_output,)`, sharing this loss's mask + and divisor. Returns `(loss, z_loss, grad_logits)`, with `z_loss` `None` when `z` is absent. TODO: Better handling of `grad_output = None` """ # TODO: Improve assumptions. @@ -777,19 +808,22 @@ def triton_entropy_loss_forward_backward( "logits_scale_factor": logits_scale_factor, "block_size": block_size, } - if grad_output is None: + z_grad_output = None if z is None else z[0] + has_grad = grad_output is not None or z_grad_output is not None + if not has_grad: backward_kwargs = {} else: 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_losses": grad_output / divisor, + "grad_losses": 0.0 if grad_output is None else grad_output / divisor, "grad_logits_stride_0": grad_logits.stride(-2), "accumulate": accumulate, } if target_format == TargetFormat.labels: assert entropy_loss_type != EntropyLossType.reverse_kl + assert z is None if num_warps is None: num_warps = 4 if block_size < 2048 else (8 if block_size < 8192 else 16) kwargs["num_warps"] = num_warps @@ -831,6 +865,9 @@ def triton_entropy_loss_forward_backward( else: if loss_mask is not None: assert loss_mask.is_contiguous() + z_losses = torch.empty(n_rows, dtype=torch.float, device=logits.device) if z is not None else None + if has_grad and z is not None: + backward_kwargs["grad_losses_z"] = 0.0 if z_grad_output is None else z_grad_output / divisor if entropy_loss_type == EntropyLossType.reverse_kl: kernel = triton_reverse_kl_forward_backward_kernel_from_distribution else: @@ -844,6 +881,7 @@ def triton_entropy_loss_forward_backward( target, loss_mask_ptr=loss_mask, losses_ptr=losses, + z_losses_ptr=z_losses, max_logits_ptr=None, sum_exp_logits_ptr=None, target_max_logits_ptr=None, @@ -907,6 +945,7 @@ def triton_entropy_loss_forward_backward( target_sum_exp_logits_ptr=target_sum_exp_logits, partial_losses_ptr=losses, losses_ptr=losses, + z_losses_ptr=z_losses, target_stride_0=target.stride(-2), target_logits_scale_factor=logits_scale_factor / temperature, from_logits=target_format == TargetFormat.logits, @@ -914,4 +953,5 @@ def triton_entropy_loss_forward_backward( **backward_kwargs, ) loss = reduce_losses(losses, divisor) - return loss, grad_logits + z_loss = None if z is None else reduce_losses(z_losses, divisor) + return loss, z_loss, grad_logits diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 525d03256..bd8d7dd10 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -136,6 +136,7 @@ def loss_class(self) -> "type[LanguageModelLabelEntropyLoss]": @config_class(dynamic_type={LanguageModelLossConfig: "distillation"}) class LanguageModelDistillationLossConfig(CombinableLossConfig): _abstract: typing.ClassVar[bool] = False + triton_kind: typing.ClassVar[str | None] = "distillation" loss_type: EntropyLossType = Field( default=EntropyLossType.cross_entropy, @@ -302,8 +303,10 @@ def _validate(self) -> None: # 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. + # Two fused kernels, one slot per loss kind: a label-family kernel (ce/z/grpo/gspo) and a + # distribution kernel (distillation + optional z-loss). Distillation can't share the label kernel. seen_kinds = set() + has_distillation = any(loss.triton_kind == "distillation" for loss in self.losses.values()) for name, loss in self.losses.items(): if loss.triton_kind is None: raise ValueError(f"Loss `{name}` (`{type(loss).__name__}`) has no triton fused implementation.") @@ -311,6 +314,11 @@ def _validate(self) -> None: raise ValueError( f"The triton path fuses at most one `{loss.triton_kind}` loss; `{name}` is a duplicate." ) + if has_distillation and loss.triton_kind not in ("distillation", "z"): + raise ValueError( + f"The triton path fuses distillation only with z-loss; `{name}` (`{loss.triton_kind}`)" + " must use the compiled path." + ) seen_kinds.add(loss.triton_kind) @property diff --git a/fast_llm/layers/language_model/loss/entropy_loss.py b/fast_llm/layers/language_model/loss/entropy_loss.py index 12cfa7bcf..069699dbe 100644 --- a/fast_llm/layers/language_model/loss/entropy_loss.py +++ b/fast_llm/layers/language_model/loss/entropy_loss.py @@ -35,7 +35,7 @@ def _forward_backward( 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( + loss, _, grad_logits = triton_entropy_loss_forward_backward( logits, target, None, # Labels are already masked @@ -47,6 +47,7 @@ def _forward_backward( entropy_loss_type=self._config.loss_type, divisor=divisor, ) + return loss, grad_logits loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments) return loss, grad_logits @@ -108,7 +109,7 @@ def _forward_backward( 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( + loss, _, grad_logits = triton_entropy_loss_forward_backward( logits, target, loss_mask, @@ -121,6 +122,7 @@ def _forward_backward( entropy_loss_type=entropy_loss_type, divisor=divisor, ) + return loss, grad_logits loss, grad_logits, _ = self.combinable_forward_backward(logits, group, grad_logits, arguments) return loss, grad_logits @@ -179,5 +181,20 @@ def fused_core( grad = grad * loss_mask.unsqueeze(-1) return loss, grad, None + def triton_add_inputs( + self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool + ) -> None: + target, loss_mask, grad_output, divisor, loss_type, temperature = self.get_inputs( + kwargs, split_index, register + ) + if context.divisor is None: + context.divisor = divisor + context.distillation = (target, loss_mask, grad_output, loss_type, temperature) + + def triton_finish( + self, context: "_TritonContext", kwargs: dict[str, typing.Any], split_index: int, register: bool + ) -> tuple[torch.Tensor, None]: + return context.dist_loss, None + def get_preprocessing_config(self) -> dict[str, typing.Any]: return {"return_prediction_mask": True} diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index b774bee99..4dd2141d7 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -27,6 +27,8 @@ class _TritonContext: z: tuple | None = None grpo: tuple | None = None gspo_coeff: torch.Tensor | None = None + # Distribution-kernel slot (distillation): `(target, loss_mask, grad_output, loss_type, temperature)`. + distillation: tuple | None = None # Per-slot kernel outputs; GSPO's loss / new-log-probs instead come from its eager seam. ce_loss: torch.Tensor | None = None z_loss: torch.Tensor | None = None @@ -34,6 +36,7 @@ class _TritonContext: grpo_new_logprobs: torch.Tensor | None = None gspo_loss: torch.Tensor | None = None gspo_new_logprobs: torch.Tensor | None = None + dist_loss: torch.Tensor | None = None # Shared metrics precursors from the kernel's softmax + `Σ exp·logits_norm` (set only when metrics are on). new_log_probs: torch.Tensor | None = None entropy_per_token: torch.Tensor | None = None @@ -121,9 +124,16 @@ 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 - # The triton kernel fuses at most one loss per kind; anything else falls back to the compiled path. + # Two fused triton kernels: a label-family kernel (ce/z/grpo/gspo) and a distribution kernel + # (distillation + optional z-loss). Pick by whether a distribution-target child is present; a set that + # fits neither kernel (e.g. distillation mixed with grpo, or a duplicate kind) falls back to compiled. triton_kinds = [child._config.triton_kind for child in self._children] - self._triton_valid = None not in triton_kinds and len(set(triton_kinds)) == len(triton_kinds) + self._triton_distribution = "distillation" in triton_kinds + no_duplicate_kinds = len(triton_kinds) == len(set(triton_kinds)) + if self._triton_distribution: + self._triton_valid = set(triton_kinds) <= {"distillation", "z"} and no_duplicate_kinds + else: + self._triton_valid = None not in triton_kinds and no_duplicate_kinds def forward_backward( self, @@ -134,6 +144,8 @@ def forward_backward( grad_logits: torch.Tensor | None = None, ) -> "tuple[torch.Tensor | None, torch.Tensor | None]": if self._triton_valid and TritonConfig.enabled(logits.device, self._config.use_triton): + if self._triton_distribution: + return self._triton_distribution_forward_backward(logits, kwargs, losses, split_index, grad_logits) 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) @@ -226,6 +238,52 @@ def _triton_forward_backward( total_loss = weighted if total_loss is None else total_loss + weighted return total_loss, grad_logits + def _triton_distribution_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]: + # Distribution-target family: distillation (the general full-distribution loss) with an optional z-loss + # superposed onto the same student softmax. No metrics or eager seam — neither child has them. + from fast_llm.functional.config import TargetFormat + from fast_llm.functional.triton.entropy_loss import triton_entropy_loss_forward_backward + + register = losses is not None + group = self._parallel_dim.group if self._vocab_parallel else None + + context = _TritonContext() + for child in self._children: + child.triton_add_inputs(context, kwargs, split_index, register) + target, loss_mask, grad_output, entropy_loss_type, temperature = context.distillation + # z-loss shares distillation's mask and divisor; the driver takes only its weighted `grad_output`. + context.dist_loss, context.z_loss, grad_logits = triton_entropy_loss_forward_backward( + logits, + target, + loss_mask, + grad_logits=grad_logits, + grad_output=grad_output, + group=group, + logits_scale_factor=self._softmax_scale_factor, + temperature=temperature, + target_format=TargetFormat.logits, + entropy_loss_type=entropy_loss_type, + z=None if context.z is None else (context.z[1],), + divisor=context.divisor, + ) + + total_loss = None + for child in self._children: + loss, extra = child.triton_finish(context, kwargs, split_index, register) + 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/tests/layers/test_lm_head.py b/tests/layers/test_lm_head.py index 5df2e292d..2383c53ad 100644 --- a/tests/layers/test_lm_head.py +++ b/tests/layers/test_lm_head.py @@ -534,7 +534,7 @@ def _add_configs(base_name: str, **kwargs): _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). +# Distillation uses the distribution-family triton kernel (added below); policy metrics stay on compiled `fused`. _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) @@ -543,6 +543,18 @@ def _add_configs(base_name: str, **kwargs): _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) +# Distillation on the triton path (the distribution-family kernel): alone, fused with z-loss over the shared +# student softmax, and with a non-unit teacher temperature. Label-CE and policy losses can't share this kernel. +_add_configs("fused_triton_distillation_loss", loss_implementation="fused_triton", distillation_loss=True) +_add_configs( + "fused_triton_distillation_and_z_loss", loss_implementation="fused_triton", distillation_loss=True, z_loss=0.5 +) +_add_configs( + "fused_triton_distillation_loss_temperature", + loss_implementation="fused_triton", + distillation_loss=True, + distillation_temperature=2.0, +) # 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): diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 49118c651..a20ae5238 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -396,7 +396,7 @@ def _test_entropy_loss( if not triton_available: return - out_triton, grad_triton = triton_entropy_loss_forward_backward( + out_triton, _, grad_triton = triton_entropy_loss_forward_backward( logits=local_logits, target=local_target, loss_mask=loss_mask, @@ -768,7 +768,7 @@ def _test_monolithic_loss_triton( # 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( + ce_ref, _, grad_ref = triton_entropy_loss_forward_backward( local_logits, target, None, @@ -827,7 +827,7 @@ def _test_monolithic_single_loss_triton( divisor = max(int((target >= 0).sum().item()), 1) previous_grad = torch.randn_like(local_logits) if accumulate else None - ce_ref, grad_ref = triton_entropy_loss_forward_backward( + ce_ref, _, grad_ref = triton_entropy_loss_forward_backward( local_logits, target, None, @@ -1018,6 +1018,86 @@ def run(grad_output): Assert.rms_close_relative(weighted_nograd, weighted_grad, 1e-5, 1e-6) +def _test_monolithic_distillation_loss_triton( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + entropy_loss_type, + block_size, + accumulate, + group=None, +): + # Distillation (from a teacher distribution) fused with z-loss over one shared student softmax — the + # distribution-family kernel — against the sum of the standalone triton distillation and z-loss kernels + # accumulating into one gradient buffer. Covers each KL direction. Distillation and z-loss share the loss + # mask and the divisor (both are threaded identically by the head), and a non-unit temperature checks that + # it stays confined to the teacher term (z-loss is on the student softmax, so it must be temperature-free). + if not triton_available: + return + temperature = 1.5 + logits, target, loss_mask = _get_lm_loss_inputs(num_columns, loss_masking, TargetFormat.logits, batch_shape, dtype) + local_logits = split_op(logits, group, -1).contiguous() + local_target = split_op(target, group, -1).contiguous() + divisor = local_logits.shape[:-1].numel() + previous_grad = torch.randn_like(local_logits) if accumulate else None + + # Reference: standalone triton distillation then z-loss, accumulating into one gradient buffer. + grad_ref = previous_grad.clone() if accumulate else None + dist_ref, _, grad_ref = triton_entropy_loss_forward_backward( + local_logits, + local_target, + loss_mask, + grad_logits=grad_ref, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + temperature=temperature, + target_format=TargetFormat.logits, + entropy_loss_type=entropy_loss_type, + divisor=divisor, + block_size=block_size, + ) + z_ref, grad_ref = triton_z_loss_forward_backward( + local_logits, + loss_mask, + grad_logits=grad_ref, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + divisor=divisor, + block_size=block_size, + ) + + dist_fused, z_fused, grad_fused = triton_entropy_loss_forward_backward( + local_logits, + local_target, + loss_mask, + grad_logits=previous_grad.clone() if accumulate else None, + grad_output=grad_output, + group=group, + logits_scale_factor=logits_scale_factor, + temperature=temperature, + target_format=TargetFormat.logits, + entropy_loss_type=entropy_loss_type, + z=(grad_output,), + divisor=divisor, + block_size=block_size, + ) + + threshold = 1e-5 if dtype == DataType.float32 else 1e-4 + Assert.rms_close_relative(dist_fused, dist_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 distillation'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("loss_masking", (False, True)) @@ -1099,6 +1179,37 @@ def test_monolithic_loss_triton( ) +@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("entropy_loss_type", EntropyLossType) +def test_monolithic_distillation_loss_triton( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + block_size, + accumulate, + entropy_loss_type, +): + _test_monolithic_distillation_loss_triton( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + entropy_loss_type, + block_size, + accumulate, + ) + + @pytest.mark.slow @pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) @pytest.mark.parametrize( @@ -1398,6 +1509,25 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa accumulate, test_context.group, ) + # Triton monolithic distribution family: distillation fused with z-loss over one softmax. + for entropy_loss_type in EntropyLossType: + with test_context.subtest( + base_path, f"monolithic-distillation-{entropy_loss_type}-triton-{suffix}", 2 + ) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_monolithic_distillation_loss_triton( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + entropy_loss_type, + block_size, + accumulate, + test_context.group, + ) @pytest.mark.slow From 97b483f7813fc17f128b2a3928c4d68bd2aed9bc Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 15 Jul 2026 12:35:10 -0400 Subject: [PATCH 2/2] Address fine-review nits on the distillation+z triton path (#507) Extract the duplicated triton child-finish loop into _triton_finish_children, shared by the label-family and distribution-family drivers; trim a redundant parenthetical in the masked-row comment. Co-Authored-By: Claude Opus 4.8 (1M context) --- fast_llm/functional/triton/entropy_loss.py | 4 ++-- .../layers/language_model/loss/monolithic.py | 22 ++++++++++--------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/fast_llm/functional/triton/entropy_loss.py b/fast_llm/functional/triton/entropy_loss.py index 73453739f..a4b4517c1 100644 --- a/fast_llm/functional/triton/entropy_loss.py +++ b/fast_llm/functional/triton/entropy_loss.py @@ -383,7 +383,7 @@ def triton_cross_entropy_from_distribution_forward_backward_kernel( target_ptr = target_ptr + block_idx * target_stride_0 if loss_mask_ptr is not None and tl.load(loss_mask_ptr + block_idx) == 0: - # This entry is masked, ignore. z-loss shares this mask (both use the loss mask), so it drops too. + # This entry is masked, ignore; z-loss shares the mask, so it drops too. if losses_ptr is not None: tl.store(losses_ptr + block_idx, 0) if z_losses_ptr is not None: @@ -627,7 +627,7 @@ def triton_reverse_kl_forward_backward_kernel_from_distribution( target_ptr = target_ptr + block_idx * target_stride_0 if loss_mask_ptr is not None and tl.load(loss_mask_ptr + block_idx) == 0: - # This entry is masked, ignore. z-loss shares this mask (both use the loss mask), so it drops too. + # This entry is masked, ignore; z-loss shares the mask, so it drops too. if losses_ptr is not None: tl.store(losses_ptr + block_idx, 0) if z_losses_ptr is not None: diff --git a/fast_llm/layers/language_model/loss/monolithic.py b/fast_llm/layers/language_model/loss/monolithic.py index 4dd2141d7..ef7d5feae 100644 --- a/fast_llm/layers/language_model/loss/monolithic.py +++ b/fast_llm/layers/language_model/loss/monolithic.py @@ -228,15 +228,7 @@ def _triton_forward_backward( context.new_log_probs = predicted_logits - max_logits - log_sum_exp_logits context.entropy_per_token = log_sum_exp_logits - weighted_logits_sum / sum_exp_logits - total_loss = None - for child in self._children: - loss, extra = child.triton_finish(context, kwargs, split_index, register) - 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 + return self._triton_finish_children(context, kwargs, losses, split_index, register), grad_logits def _triton_distribution_forward_backward( self, @@ -274,6 +266,16 @@ def _triton_distribution_forward_backward( divisor=context.divisor, ) + return self._triton_finish_children(context, kwargs, losses, split_index, register), grad_logits + + def _triton_finish_children( + self, + context: _TritonContext, + kwargs: dict[str, typing.Any], + losses: dict | None, + split_index: int, + register: bool, + ) -> torch.Tensor | None: total_loss = None for child in self._children: loss, extra = child.triton_finish(context, kwargs, split_index, register) @@ -282,7 +284,7 @@ def _triton_distribution_forward_backward( 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 + return total_loss def get_preprocessing_config(self) -> dict[str, typing.Any]: return safe_merge_dicts(*(child.get_preprocessing_config() for child in self._children))