adding dpo loss

Author: tobyzl2

fast_llm/functional/config.py

@@ -91,3 +91,7 @@ class CrossEntropyImpl(str, enum.Enum):
     torch = "torch"
     fused = "fused"
     triton = "triton"
+
+class LossFunctionType(str, enum.Enum):
+    cross_entropy = "cross_entropy"
+    dpo = "dpo"

fast_llm/functional/dpo.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn.functional as F
+from typing import Tuple
+
+
+def compute_logps_for_spans(
+        logits: torch.Tensor, 
+        targets: torch.Tensor, 
+        chosen_span: torch.Tensor, 
+        rejected_span: torch.Tensor
+    ):
+    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+    
+    # gather log probabilities corresponding to the target tokens
+    # selected_log_probs = log_probs[torch.arange(logits.shape[0] - 1), targets]
+    selected_log_probs = log_probs[:-1].gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)
+    
+    # apply chosen mask
+    chosen_mask = torch.zeros_like(selected_log_probs, dtype=torch.bool)
+    chosen_mask[chosen_span[:, 0]: chosen_span[:, 1] + 1] = 1
+    chosen_logp = (selected_log_probs * chosen_mask).sum()
+
+    # apply rejected mask
+    rejected_mask = torch.zeros_like(selected_log_probs, dtype=torch.bool)
+    rejected_mask[rejected_span[:, 0]: rejected_span[:, 1] + 1] = 1
+    rejected_logp = (selected_log_probs * rejected_mask).sum()
+
+    # chosen_logp = selected_log_probs[chosen_span[:, 0]: chosen_span[:, 1] + 1].sum()
+    # rejected_logp = selected_log_probs[rejected_span[:, 0]: rejected_span[:, 1] + 1].sum()
+    
+    return chosen_logp, rejected_logp
+
+def compute_simplified_dpo_loss(
+    logits: torch.Tensor, 
+    targets: torch.Tensor, 
+    chosen_span: torch.Tensor, 
+    rejected_span: torch.Tensor,
+    beta: float,
+    grad_output: float | None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    with torch.enable_grad():
+        logits_ = logits.float().detach().requires_grad_()
+
+        policy_chosen_logps, policy_rejected_logps = compute_logps_for_spans(logits_, targets, chosen_span, rejected_span)
+
+        pi_logratios = policy_chosen_logps - policy_rejected_logps
+
+        losses = -F.logsigmoid(beta * pi_logratios)
+        if grad_output is None:
+            loss = None
+        else:
+            loss = losses.mean()
+            loss.backward(torch.full_like(loss, grad_output))
+            loss.detach()
+    return loss.detach(), logits_.grad.detach().to(logits.dtype)

fast_llm/layers/language_model/config.py

@@ -4,7 +4,7 @@
 from fast_llm.engine.base_model.config import BaseModelArchitectureConfig, BaseModelConfig
 from fast_llm.engine.config_utils.tensor_space import TensorDim, TensorSpace
 from fast_llm.engine.distributed.config import DistributedDimNames
-from fast_llm.functional.config import CrossEntropyImpl
+from fast_llm.functional.config import CrossEntropyImpl, LossFunctionType
 from fast_llm.layers.transformer.config import TransformerArchitectureConfig, TransformerConfig
 from fast_llm.utils import Assert
 
@@ -28,6 +28,8 @@ class LanguageModelKwargs:
     # TODO: These are generic
     labels = "labels"
     phase = "phase"
+    chosen_spans = "chosen_spans"
+    rejected_spans = "rejected_spans"
 
 
 @config_class()
@@ -128,6 +130,16 @@ class LanguageModelBaseConfig(LanguageModelArchitectureConfig, BaseModelConfig):
         desc="Min value for clamping initialized weights of the vocabulary embedding and output (logits).",
         hint=FieldHint.feature,
     )
+    loss_function_type: LossFunctionType = Field(
+        default=LossFunctionType.cross_entropy,
+        desc="Type of loss function to use",
+        hint=FieldHint.feature,
+    )
+    beta: float | None = Field(
+        default=1.0,
+        desc="Beta value for DPO loss.",
+        hint=FieldHint.feature,
+    )
     cross_entropy_impl: CrossEntropyImpl = Field(
         default=CrossEntropyImpl.auto,
         desc="Implementation for the cross-entropy computation.",

fast_llm/layers/language_model/head.py

@@ -10,9 +10,10 @@
 from fast_llm.engine.config_utils.tensor_space import DefaultDimNames, TensorDim, TensorSpace
 from fast_llm.engine.distributed.config import DistributedDimNames
 from fast_llm.functional.autograd import grad_is_context, wrap_forward_backward
-from fast_llm.functional.config import CrossEntropyImpl, TritonConfig
+from fast_llm.functional.config import CrossEntropyImpl, TritonConfig, LossFunctionType
 from fast_llm.functional.cross_entropy import cross_entropy_forward_backward
 from fast_llm.functional.linear import output_parallel_linear_backward, output_parallel_linear_forward
+from fast_llm.functional.dpo import compute_simplified_dpo_loss
 from fast_llm.layers.common.auxiliary_loss import z_loss
 from fast_llm.layers.language_model.config import (
     LanguageModelBaseConfig,
@@ -74,14 +75,20 @@ def __init__(
                 ),
             )
 
-        self._cross_entropy_impl = config.cross_entropy_impl
-        if self._cross_entropy_impl == CrossEntropyImpl.auto:
-            if self._parallel_embeddings:
-                self._cross_entropy_impl = CrossEntropyImpl.fused
-            elif TritonConfig.TRITON_ENABLED:
-                self._cross_entropy_impl = CrossEntropyImpl.triton
-            else:
-                self._cross_entropy_impl = CrossEntropyImpl.fused
+        self._loss_function_type = config.loss_function_type
+        if self._loss_function_type == LossFunctionType.cross_entropy:
+            self._cross_entropy_impl = config.cross_entropy_impl
+            if self._cross_entropy_impl == CrossEntropyImpl.auto:
+                if self._parallel_embeddings:
+                    self._cross_entropy_impl = CrossEntropyImpl.fused
+                elif TritonConfig.TRITON_ENABLED:
+                    self._cross_entropy_impl = CrossEntropyImpl.triton
+                else:
+                    self._cross_entropy_impl = CrossEntropyImpl.fused
+            self._loss_fcn = self._logits_cross_entropy_forward_backward_split
+        else:
+            self._loss_fcn = self._logits_dpo
+            self.dpo_beta = config.beta
 
         self._forward = wrap_forward_backward(self._forward_backward, grad_is_context)
 
@@ -127,7 +134,7 @@ def _forward_backward(
         )
 
         output_weights = kwargs[WORD_EMBEDDINGS_WEIGHT] if self._tie_word_embeddings else self.output_weights
-        loss, ln_output_grad = self._logits_cross_entropy_forward_backward_split(
+        loss, ln_output_grad = self._loss_fcn(
             ln_output.detach(), labels, output_weights, grad_output, kwargs, losses
         )
 
@@ -136,6 +143,38 @@ def _forward_backward(
             return loss, input_.grad
         else:
             return loss, None
+        
+    def _logits_dpo(
+            self,
+            input_: torch.Tensor,
+            labels: torch.Tensor | None,
+            weight: torch.Tensor,
+            grad_output: float,
+            kwargs: dict,
+            losses: dict | None = None
+        ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        logits, context = output_parallel_linear_forward(
+            input_=input_,
+            weight=weight,
+            bias=None,
+            group=self._tensor_space.distributed.tensor_group if self._parallel_embeddings else None,
+            sequence_parallel=self._sequence_parallel and self._parallel_embeddings,
+        )
+
+        loss, grad = compute_simplified_dpo_loss(
+            logits.flatten(0, -2),
+            labels,
+            kwargs[LanguageModelKwargs.chosen_spans],
+            kwargs[LanguageModelKwargs.rejected_spans],
+            self.dpo_beta,
+            grad_output
+        )
+
+        # TODO: de-allocate earlier.
+        del logits
+        return loss, output_parallel_linear_backward(grad, context).view_as(input_)
+
+
 
     def _logits_cross_entropy_forward_backward_split(
         self,

fast_llm/models/gpt/model.py

@@ -254,7 +254,7 @@ def preprocess(
                 TransformerKwargs.presents: presents,
             }
             if phase != PhaseType.inference:
-                sequence_offset = sequence_k - sequence_q + 1
+                sequence_offset = sequence_k - sequence_q + 1  # +1 for shift in labels
                 if sequence_first:
                     labels = batch.token_ids[sequence_offset : sequence_k + 1]
                 else:
@@ -266,8 +266,10 @@ def preprocess(
                     for i, spans in enumerate(batch.loss_masking_spans):
                         if not spans.numel():
                             continue
+                        # filter spans within the sequence or partially within the sequence
                         valid_spans = spans[(spans[:, 0] <= sequence_k) & (spans[:, 1] >= sequence_offset)]
                         if valid_spans.numel():
+                            # if span is partially within the sequence, truncate parts of spans that are outside of the sequence
                             valid_spans[:, 0].clamp_(min=sequence_offset)
                             valid_spans[:, 1].clamp_(max=sequence_k)
                             valid_spans -= sequence_offset
@@ -276,6 +278,30 @@ def preprocess(
                                     labels[start : end + 1, i] = -100
                                 else:
                                     labels[i, start : end + 1] = -100
+                if batch.chosen_loss_masking_spans is not None:
+                    for i, spans in enumerate(batch.chosen_loss_masking_spans):
+                        if not spans.numel():
+                            continue
+                        # filter spans within the sequence or partially within the sequence
+                        valid_spans = spans[(spans[0] <= sequence_k) & (spans[1] >= sequence_offset)]
+                        if valid_spans.numel():
+                            # if span is partially within the sequence, truncate parts of spans that are outside of the sequence
+                            valid_spans[:, 0].clamp_(min=sequence_offset)
+                            valid_spans[:, 1].clamp_(max=sequence_k)
+                            valid_spans -= sequence_offset
+                            kwargs[LanguageModelKwargs.chosen_spans] = valid_spans
+                if batch.rejected_loss_masking_spans is not None:
+                    for i, spans in enumerate(batch.rejected_loss_masking_spans):
+                        if not spans.numel():
+                            continue
+                        # filter spans within the sequence or partially within the sequence
+                        valid_spans = spans[(spans[0] <= sequence_k) & (spans[1] >= sequence_offset)]
+                        if valid_spans.numel():
+                            # if span is partially within the sequence, truncate parts of spans that are outside of the sequence
+                            valid_spans[:, 0].clamp_(min=sequence_offset)
+                            valid_spans[:, 1].clamp_(max=sequence_k)
+                            valid_spans -= sequence_offset
+                            kwargs[LanguageModelKwargs.rejected_spans] = valid_spans
                 kwargs[LanguageModelKwargs.labels] = labels
             if self._config.use_absolute_position_embeddings:
                 self._position_embedding_preprocessor.preprocess(kwargs)