Loss masking for distillation (#250)

Author: jlamypoirier

fast_llm/functional/cross_entropy.py

@@ -8,9 +8,10 @@
 from fast_llm.utils import Assert
 
 
-def torch_cross_entropy_forward_backward(
+def _torch_cross_entropy_forward_backward(
     logits: torch.Tensor,
     target: torch.Tensor,
+    loss_mask: torch.Tensor | None,
     grad_output: float | None,
     logits_scale_factor: float,
     target_format: TargetFormat,
@@ -28,9 +29,17 @@ def torch_cross_entropy_forward_backward(
             if logits_scale_factor != 1.0:
                 target = target * logits_scale_factor
             target = torch.softmax(target, dim=-1)
-        loss = torch.nn.functional.cross_entropy(
-            logits_ if logits_scale_factor == 1 else logits_ * logits_scale_factor, target
-        ).mean()
+        if loss_mask is None:
+            loss = torch.nn.functional.cross_entropy(
+                logits_ if logits_scale_factor == 1 else logits_ * logits_scale_factor, target
+            )
+        else:
+            loss = (
+                torch.nn.functional.cross_entropy(
+                    logits_ if logits_scale_factor == 1 else logits_ * logits_scale_factor, target, reduction="none"
+                )
+                * loss_mask
+            ).mean()
         if grad_output is None:
             grad = None
         else:
@@ -39,7 +48,7 @@ def torch_cross_entropy_forward_backward(
     return loss.detach_(), grad
 
 
-# @torch.compile
+@torch.compile
 def _fused_softmax_base(
     logits: torch.Tensor, logits_scale_factor: float = 1.0, group: ProcessGroup | None = None, dim: int = -1
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -57,18 +66,19 @@ def _fused_softmax_base(
     return logits_norm, exp_logits, sum_exp_logits
 
 
-# @torch.compile
-def fused_softmax(
+@torch.compile
+def _fused_softmax(
     logits: torch.Tensor, logits_scale_factor: float = 1.0, group: ProcessGroup = None, dim: int = -1
 ) -> torch.Tensor:
     _, exp_logits, sum_exp_logits = _fused_softmax_base(logits, logits_scale_factor, group, dim)
     return exp_logits / sum_exp_logits
 
 
-@torch.compile
-def fused_cross_entropy_forward_backward(
+# @torch.compile
+def _fused_cross_entropy_forward_backward(
     logits: torch.Tensor,
     target: torch.Tensor,
+    loss_mask: torch.Tensor | None,
     grad_output: float | None,
     logits_scale_factor: float,
     target_format: TargetFormat,
@@ -85,7 +95,7 @@ def fused_cross_entropy_forward_backward(
     logits_norm, exp_logits, sum_exp_logits = _fused_softmax_base(logits, logits_scale_factor, group)
 
     if target_format == TargetFormat.logits:
-        target = fused_softmax(target, logits_scale_factor, group)
+        target = _fused_softmax(target, logits_scale_factor, group)
 
     if target_format == TargetFormat.labels:
         target = target.unsqueeze(-1)
@@ -101,10 +111,10 @@ def fused_cross_entropy_forward_backward(
             target_mask = (target >= vocab_start_index) * (target < vocab_start_index + logits.size(-1))
             target = (target - vocab_start_index) * target_mask
     else:
-        # TODO: Support masking
-        loss_mask = None
         # Target should be tensor-parallel already, no further manipulation needed.
         target_mask = None
+        if loss_mask is not None:
+            loss_mask = loss_mask.unsqueeze(-1)
 
     if grad_output is None:
         grad = None
@@ -120,9 +130,9 @@ def fused_cross_entropy_forward_backward(
         grad = grad_base.mul((grad_output / logits.size(0)) / sum_exp_logits)
         if logits_scale_factor != 1.0:
             grad *= logits_scale_factor
-        grad = grad.to(logits.dtype)
         if loss_mask is not None:
-            grad = torch.where(loss_mask, grad.to(logits.dtype), 0)
+            grad *= loss_mask
+        grad = grad.to(logits.dtype)
 
     # loss = mean(log(sum_exp_logits) - sum(probabilities * logits))
     if target_format == TargetFormat.labels:
@@ -145,15 +155,16 @@ def fused_cross_entropy_forward_backward(
 
 
 _CROSS_ENTROPY_IMPLEMENTATIONS = {
-    CrossEntropyImpl.torch: torch_cross_entropy_forward_backward,
-    CrossEntropyImpl.fused: fused_cross_entropy_forward_backward,
+    CrossEntropyImpl.torch: _torch_cross_entropy_forward_backward,
+    CrossEntropyImpl.fused: _fused_cross_entropy_forward_backward,
     CrossEntropyImpl.triton: triton_cross_entropy_forward_backward,
 }
 
 
 def cross_entropy_forward_backward(
     logits: torch.Tensor,
     target: torch.Tensor,
+    loss_mask: torch.Tensor | None,
     grad_output: float | None,
     group: ProcessGroup | None = None,
     implementation: CrossEntropyImpl = CrossEntropyImpl.fused,
@@ -169,15 +180,18 @@ def cross_entropy_forward_backward(
     if target_format == TargetFormat.labels:
         Assert.eq(target.shape, logits.shape[:-1])
         Assert.eq(target.dtype, torch.int64)
+        assert loss_mask is None
     else:
         Assert.eq(target.shape, logits.shape)
         assert target.dtype.is_floating_point, target.dtype
+        if loss_mask is not None:
+            Assert.eq(loss_mask.shape, logits.shape[:-1])
     if group:
         Assert.eq(implementation, CrossEntropyImpl.fused)
-        return fused_cross_entropy_forward_backward(
-            logits, target, grad_output, logits_scale_factor, target_format, group
+        return _fused_cross_entropy_forward_backward(
+            logits, target, loss_mask, grad_output, logits_scale_factor, target_format, group
         )
     else:
         return _CROSS_ENTROPY_IMPLEMENTATIONS[implementation](
-            logits, target, grad_output, logits_scale_factor, target_format
+            logits, target, loss_mask, grad_output, logits_scale_factor, target_format
         )

fast_llm/functional/triton/cross_entropy.py

@@ -57,6 +57,7 @@ def triton_cross_entropy_forward_backward_kernel(
 def triton_cross_entropy_from_distribution_forward_backward_kernel(
     logits_ptr,
     target_ptr,
+    loss_mask_ptr,
     grad_logits_ptr,
     losses_ptr,
     grad_losses,
@@ -73,6 +74,14 @@ def triton_cross_entropy_from_distribution_forward_backward_kernel(
     col_offsets = tl.arange(0, block_size)
     mask = col_offsets < n_cols
 
+    if loss_mask_ptr is not None:
+        loss_mask = tl.load(loss_mask_ptr + block_idx)
+        if loss_mask == 0:
+            tl.store(losses_ptr + block_idx, 0)
+            if grad_losses is not None:
+                tl.store(grad_logits_ptr + block_idx * grad_logits_stride_0 + col_offsets, 0, mask=mask)
+            return
+
     logits = tl.load(logits_ptr + block_idx * logits_stride_0 + col_offsets, mask=mask, other=-float("inf")).to(
         tl.float32
     )
@@ -104,12 +113,15 @@ def triton_cross_entropy_from_distribution_forward_backward_kernel(
         grad_logits = grad_losses * (exp_logits / sum_exp_logits - target)
         if logits_scale_factor != 1.0:
             grad_logits *= logits_scale_factor
+        if loss_mask_ptr is not None:
+            grad_logits = grad_logits
         tl.store(grad_logits_ptr + block_idx * grad_logits_stride_0 + col_offsets, grad_logits, mask=mask)
 
 
 def triton_cross_entropy_forward_backward(
     logits: torch.Tensor,
     target: torch.Tensor,
+    loss_mask: torch.Tensor | None,
     grad_output: float | None,
     logits_scale_factor: float,
     target_format: TargetFormat,
@@ -146,9 +158,12 @@ def triton_cross_entropy_forward_backward(
             num_warps=num_warps,
         )
     else:
+        if loss_mask is not None:
+            assert loss_mask.is_contiguous()
         triton_cross_entropy_from_distribution_forward_backward_kernel[(n_rows,)](
             logits,
             target,
+            loss_mask,
             grad_logits,
             losses,
             None if grad_output is None else grad_output / n_rows,

fast_llm/layers/language_model/config.py

@@ -34,6 +34,7 @@ class LanguageModelKwargs:
     # TODO: These are generic
     labels = "labels"
     phase = "phase"
+    loss_mask = "loss_mask"
 
 
 @config_class()

fast_llm/layers/language_model/head.py

@@ -146,6 +146,8 @@ def _forward_backward(
             if self._config.distillation_model is None
             else f"{self._config.distillation_model}_logits"
         )
+        # Loss mask for distillation. (Labels are already masked.)
+        loss_mask = None
         if target is not None:
             if self._config.distillation_model is None:
                 # MTP: Shift the labels
@@ -160,9 +162,14 @@ def _forward_backward(
             else:
                 # Target is reference model logits.
                 target = target.flatten(0, -2)
+                loss_mask = kwargs.get(LanguageModelKwargs.loss_mask)
+                if loss_mask is not None:
+                    loss_mask = loss_mask.flatten()
 
         if self._sequence_parallel_logits:
             target = split_op(target, self._tensor_space.distributed.tensor_group, 0)
+            if loss_mask is not None:
+                loss_mask = split_op(loss_mask, self._tensor_space.distributed.tensor_group, 0)
         do_grad = target is not None and self.training
         input_ = input_.detach().requires_grad_(do_grad)
         with torch.enable_grad():
@@ -174,7 +181,7 @@ def _forward_backward(
 
         output_weights = self._get_output_weights(kwargs)
         loss, ln_output_grad = self._logits_cross_entropy_forward_backward_split(
-            ln_output.detach(), target, output_weights, grad_output, kwargs, losses
+            ln_output.detach(), target, loss_mask, output_weights, grad_output, kwargs, losses
         )
 
         if do_grad:
@@ -194,14 +201,15 @@ def _logits_cross_entropy_forward_backward_split(
         self,
         input_: torch.Tensor,
         target: torch.Tensor | None,
+        loss_mask: torch.Tensor | None,
         weight: torch.Tensor,
         grad_output: float,
         kwargs: dict,
         losses: dict | None = None,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         if self._cross_entropy_splits is None or target is None:
             loss, logit_input_grad = self._logits_cross_entropy_forward_backward(
-                input_, target, weight, grad_output, kwargs, losses
+                input_, target, loss_mask, weight, grad_output, kwargs, losses
             )
             if target is None:
                 # TODO: Make a proper way of returning the model output.
@@ -214,12 +222,17 @@ def _logits_cross_entropy_forward_backward_split(
             grad_output /= self._cross_entropy_splits
             logit_input = input_.flatten(0, -2)
             logit_input_grad = torch.empty_like(logit_input)
-            for logit_input_, target_, logit_input_grad_ in zip(
-                logit_input.split(split_size), target.split(split_size), logit_input_grad.split(split_size)
+            for logit_input_, target_, loss_mask_, logit_input_grad_ in zip(
+                logit_input.split(split_size),
+                target.split(split_size),
+                [None] * self._cross_entropy_splits if loss_mask is None else loss_mask.split(split_size),
+                logit_input_grad.split(split_size),
+                strict=True,
             ):
                 loss_, grad_ = self._logits_cross_entropy_forward_backward(
                     logit_input_,
                     target_,
+                    loss_mask_,
                     weight,
                     grad_output,
                     kwargs,
@@ -240,6 +253,7 @@ def _logits_cross_entropy_forward_backward(
         self,
         input_: torch.Tensor,
         target: torch.Tensor | None,
+        loss_mask: torch.Tensor | None,
         weight: torch.Tensor,
         grad_output: float,
         kwargs: dict,
@@ -298,6 +312,7 @@ def _logits_cross_entropy_forward_backward(
         loss, grad = cross_entropy_forward_backward(
             logits.flatten(0, -2),
             target,
+            loss_mask,
             group=self._tensor_space.distributed.tensor_group if self._parallel_embeddings else None,
             grad_output=grad_output,
             implementation=self._cross_entropy_impl,

fast_llm/models/gpt/config.py

@@ -179,9 +179,6 @@ def _validate(self) -> None:
             Assert.eq(self.reference_models.keys(), {name})
         if self.model.base_model.use_absolute_position_embeddings:
             Assert.geq(self.model.base_model.num_absolute_position_embeddings, self.batch.sequence_length)
-        if self.model.base_model.distillation_model is not None:
-            # TODO: Support loss masking for distillation?
-            assert not self.batch.use_loss_masking_spans
         for reference_model in self.reference_models.values():
             Assert.none(reference_model.model.base_model.distillation_model)
             # TODO: Support more LM head features.

fast_llm/models/gpt/model.py

@@ -315,11 +315,15 @@ def preprocess(
                             valid_spans[:, 0].clamp_(min=sequence_offset)
                             valid_spans[:, 1].clamp_(max=sequence_k + prediction_heads - 1)
                             valid_spans -= sequence_offset
+                            loss_mask = torch.ones_like(labels, dtype=torch.bool)
                             for start, end in valid_spans:
                                 if sequence_first:
-                                    labels[start : end + 1, i] = -100
+                                    loss_mask[start : end + 1, i] = False
                                 else:
-                                    labels[i, start : end + 1] = -100
+                                    loss_mask[i, start : end + 1] = False
+                            if self._config.distillation_model is not None:
+                                kwargs[LanguageModelKwargs.loss_mask] = loss_mask
+                            labels = torch.where(loss_mask, labels, -100)
                 kwargs[LanguageModelKwargs.labels] = labels
             kwargs.update(reference_logits[i])