feat(training): support configurable loss fn via training args

Signed-off-by: Christina Xu <chrxu@redhat.com>

Author: christinaexyou

scripts/train.py

@@ -17,6 +17,7 @@
 )
 from speculators.model import SpeculatorModel
 from speculators.models.eagle3.data import shift_batch
+from speculators.models.metrics import resolve_loss_fn
 from speculators.train.data import (
     ArrowDataset,
     BaseDataset,
@@ -575,6 +576,17 @@ def parse_args():
     parser.add_argument("--mask-token-id", type=int, default=None)
     parser.add_argument("--ttt-steps", type=int, default=3)
     parser.add_argument("--ttt-step-loss-decay", type=float, default=1.0)
+    parser.add_argument(
+        "--loss-fn",
+        type=str,
+        default="kl_div",
+        choices=["kl_div", "ce"],
+        help=(
+            "Loss function used during draft model training. "
+            "'kl_div' = KL divergence (default). "
+            "'ce' = cross-entropy."
+        ),
+    )
     parser.add_argument(
         "--seed", type=int, default=42, help="Random seed for reproducibility"
     )
@@ -680,7 +692,10 @@ def parse_args():
     parser.add_argument("--scheduler-warmup-steps", type=int, default=None)
     parser.add_argument("--scheduler-total-steps", type=int, default=None)
     parser.add_argument("--scheduler-num-cosine-cycles", type=float, default=0.5)
-    return parser.parse_args()
+
+    args = parser.parse_args()
+    resolve_loss_fn(args.loss_fn)
+    return args
 
 
 if __name__ == "__main__":

src/speculators/models/dflash/core.py

@@ -1,4 +1,4 @@
-from typing import Any, ClassVar
+from typing import ClassVar
 
 import torch
 from torch import nn
@@ -18,6 +18,7 @@
     get_base_indices_for_anchored_blocks,
     select_anchors,
 )
+from speculators.models.metrics import kl_div_loss, resolve_loss_fn
 from speculators.models.utils import resolve_target_layer_ids
 
 
@@ -162,7 +163,7 @@ def from_training_args(
         return model
 
     @staticmethod
-    def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:  # noqa: ARG004
+    def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         """Get training and validation kwargs for DFlash.
 
         Args:
@@ -171,9 +172,8 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:  # noqa: ARG004
         Returns:
             Tuple of (train_call_kwargs, val_call_kwargs)
         """
-        train_kwargs: dict[str, Any] = {}
-        val_kwargs: dict[str, Any] = {}
-        return train_kwargs, val_kwargs
+        loss_fn = resolve_loss_fn(kwargs["loss_fn"])
+        return {"loss_fn": loss_fn}, {"loss_fn": loss_fn}
 
     @property
     def mask_token_id(self) -> int:
@@ -194,6 +194,7 @@ def forward(
         verifier_last_hidden_states: torch.Tensor,  # shape: [1, total_seq_len, hidden_size] # noqa: E501
         lengths: torch.Tensor | None = None,  # shape: [batch_size]
         position_ids: torch.Tensor | None = None,  # shape: [1, total_seq_len]
+        loss_fn=kl_div_loss,
         **kwargs,
     ):
         device = hidden_states.device
@@ -293,7 +294,7 @@ def forward(
 
         aligned_loss_mask[:, :: self.block_size] = 0
         loss, metrics = compute_metrics(
-            logits, targets, aligned_loss_mask, self.block_size
+            logits, targets, aligned_loss_mask, self.block_size, loss_fn=loss_fn
         )
         draft_tokens = torch.argmax(logits, dim=-1)
 

src/speculators/models/dflash/metrics.py

@@ -1,14 +1,15 @@
 """Metrics and loss functions for DFlash draft model."""
 
+from collections.abc import Callable
 from functools import partial
 from typing import Any
 
 import torch
 
 from speculators.models.metrics import (
-    ce_loss,
     compute_accuracy_multi_step,
     dflash_loss_decay,
+    kl_div_loss,
     loss_function,
 )
 
@@ -19,6 +20,7 @@ def compute_metrics(
     loss_mask: torch.Tensor,  # shape: [1, num_anchors*block_size]
     block_size: int = 1,
     gamma: float = 4.0,
+    loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = kl_div_loss,
 ) -> tuple[torch.Tensor, dict]:
     """Compute loss and accuracy metrics for draft model predictions.
 
@@ -28,13 +30,16 @@ def compute_metrics(
         loss_mask: Binary mask [1, T]
         block_size: Block size for per-position metrics
         gamma: Temperature for exponential decay in loss weighting
+        loss_fn: Loss function
 
     Returns:
         Tuple of (loss, metrics_dict) where metrics_dict contains:
             - loss: Scalar loss value
             - full_acc: Overall accuracy
             - position {i} acc: Accuracy at position i within blocks
     """
+    if loss_fn is None:
+        loss_fn = kl_div_loss
     seq_len = logits.shape[1]
     pos_idx = torch.arange(seq_len, device=logits.device) % block_size
     pos_idx = pos_idx.unsqueeze(0)  # shape: [1, T]
@@ -44,7 +49,7 @@ def compute_metrics(
         targets,
         loss_mask,
         pos_idx,
-        loss_fn=ce_loss,
+        loss_fn=loss_fn,
         decay_fn=partial(dflash_loss_decay, gamma=gamma),
     )
 

src/speculators/models/eagle3/core.py

@@ -14,6 +14,7 @@
 )
 from speculators.models.eagle3.metrics import compute_metrics
 from speculators.models.eagle3.model_definitions import model_classes
+from speculators.models.metrics import kl_div_loss, resolve_loss_fn
 from speculators.models.utils import resolve_target_layer_ids
 from speculators.proposals.greedy import GreedyTokenProposalConfig
 
@@ -130,8 +131,10 @@ def forward(  # noqa: C901
         ttt_steps: int = 3,
         ttt_step_loss_decay: float = 1.0,
         use_off_policy_tokens: bool = False,
+        loss_fn=kl_div_loss,
         **kwargs,
     ):
+        loss_fn = loss_fn or kl_div_loss
         device = hidden_states.device
         total_seq_len = hidden_states.shape[1]
 
@@ -222,6 +225,7 @@ def forward(  # noqa: C901
                     prev_correct,
                     ttt_step,
                     ttt_step_loss_decay,
+                    loss_fn=loss_fn,
                 )
                 loss += s_loss
                 metrics.update(s_metrics)
@@ -324,14 +328,17 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         Returns:
             Tuple of (train_call_kwargs, val_call_kwargs)
         """
+        loss_fn = resolve_loss_fn(kwargs["loss_fn"])
         train_kwargs = {
             "use_off_policy_tokens": kwargs["use_off_policy_tokens"],
             "ttt_steps": kwargs["ttt_steps"],
             "ttt_step_loss_decay": kwargs["ttt_step_loss_decay"],
+            "loss_fn": loss_fn,
         }
         val_kwargs = {
             "use_off_policy_tokens": False,
             "ttt_steps": kwargs["ttt_steps"],
             "ttt_step_loss_decay": kwargs["ttt_step_loss_decay"],
+            "loss_fn": loss_fn,
         }
         return train_kwargs, val_kwargs

src/speculators/models/eagle3/metrics.py

@@ -1,5 +1,6 @@
 """Metrics and loss functions for Eagle3 draft model."""
 
+from collections.abc import Callable
 from functools import partial
 
 import torch
@@ -54,6 +55,7 @@ def compute_metrics(
     prev_correct: torch.Tensor | None,
     ttt_step: int,
     ttt_step_loss_decay: float,
+    loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = kl_div_loss,
 ) -> tuple[torch.Tensor, dict]:
     """Compute metrics for a given ttt_step.
 
@@ -64,13 +66,16 @@ def compute_metrics(
         prev_correct: The previous correct predictions for the current ttt_step.
         ttt_step: The current ttt_step.
         ttt_step_loss_decay: The loss decay for the current ttt_step.
+        loss_fn: Loss function.
 
     Effects:
         Modifies prev_correct in place.
 
     Returns:
         Loss value and metrics dictionary.
     """
+    if loss_fn is None:
+        loss_fn = kl_div_loss
     s_logits, s_targets, s_loss_mask, s_prev_correct = align_for_step(
         logits, targets, loss_mask, prev_correct, ttt_step
     )
@@ -88,7 +93,7 @@ def compute_metrics(
         s_targets,
         s_loss_mask,
         pos_idx,
-        loss_fn=kl_div_loss,
+        loss_fn=loss_fn,
         decay_fn=partial(exp_loss_decay, gamma=ttt_step_loss_decay),
     )
 

src/speculators/models/metrics.py

@@ -153,6 +153,31 @@ def exp_loss_decay(pos_idx: torch.Tensor, gamma: float):
     return gamma**pos_idx
 
 
+def resolve_loss_fn(
+    name: str,
+) -> "Callable[[torch.Tensor, torch.Tensor], torch.Tensor]":
+    """Resolves a loss function given its abbreviated name.
+
+    Args:
+        name: ``"kl_div"`` for KL-divergence or ``"ce"`` for cross-entropy.
+
+    Returns:
+        The corresponding loss function.
+
+    Raises:
+        ValueError: If *name* is not a recognised loss function.
+    """
+    loss_fn_map: dict[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = {
+        "kl_div": kl_div_loss,
+        "ce": ce_loss,
+    }
+    if name not in loss_fn_map:
+        raise ValueError(
+            f"Unknown loss function '{name}'. Choose from: {sorted(loss_fn_map.keys())}"
+        )
+    return loss_fn_map[name]
+
+
 def loss_function(
     logits: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]
     targets: torch.Tensor,  # shape: [1, seq_len, draft_vocab_size]

src/speculators/models/peagle/core.py

@@ -9,6 +9,7 @@
 from speculators.config import SpeculatorsConfig, VerifierConfig
 from speculators.model import SpeculatorModel
 from speculators.models.eagle3.core import Eagle3DraftModel, conditional_torch_compile
+from speculators.models.metrics import kl_div_loss, resolve_loss_fn
 from speculators.models.peagle.attention import create_peagle_mask_mod
 from speculators.models.peagle.config import PEagleSpeculatorConfig
 from speculators.models.peagle.data import generate_cod_sample_indices
@@ -55,6 +56,7 @@ def forward(
         position_ids: torch.Tensor | None = None,
         loss_mask: torch.Tensor | None = None,
         verifier_last_hidden_states: torch.Tensor | None = None,
+        loss_fn=kl_div_loss,
         **kwargs,
     ):
         """
@@ -172,6 +174,7 @@ def forward(
             anchor_pos=anchor_pos,
             depth=depth,
             num_depths=self.num_depths,
+            loss_fn=loss_fn,
         )
 
         return None, loss, metrics
@@ -237,17 +240,15 @@ def from_training_args(
         return model
 
     @staticmethod
-    def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:  # noqa: ARG004
+    def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
         """
         Get training and validation kwargs for P-EAGLE.
 
-        P-EAGLE doesn't need extra kwargs for forward() - all parameters
-        are handled in the forward method (num_depths, down_sample_ratio, etc.)
-
         Args:
-            **kwargs: Training arguments (unused)
+            **kwargs: Training arguments
 
         Returns:
-            Tuple of (train_call_kwargs, val_call_kwargs) - both empty for P-EAGLE
+            Tuple of (train_call_kwargs, val_call_kwargs)
         """
-        return {}, {}
+        loss_fn = resolve_loss_fn(kwargs["loss_fn"])
+        return {"loss_fn": loss_fn}, {"loss_fn": loss_fn}

src/speculators/models/peagle/metrics.py

@@ -1,5 +1,6 @@
 """Metrics and loss functions for P-EAGLE draft model."""
 
+from collections.abc import Callable
 from typing import Any
 
 import torch
@@ -18,6 +19,7 @@ def compute_metrics(
     anchor_pos: torch.Tensor,
     depth: torch.Tensor,
     num_depths: int,
+    loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = kl_div_loss,
 ) -> tuple[torch.Tensor, dict[str, Any]]:
     """Compute loss and accuracy metrics for P-EAGLE predictions.
 
@@ -29,10 +31,13 @@ def compute_metrics(
             sampling chain started from [total_sampled]
         depth: Which COD sampling round each element belongs to [total_sampled]
         num_depths: Number of parallel depths
+        loss_fn: Loss function.
 
     Returns:
         Tuple of (loss, metrics_dict)
     """
+    if loss_fn is None:
+        loss_fn = kl_div_loss
     device = logits.device
 
     # TODO: batch size is always 1 for P-EAGLE; unsqueeze is only to match the
@@ -41,7 +46,7 @@ def compute_metrics(
     sampled_loss_mask = loss_mask[:, orig_positions]  # [1, total_sampled]
 
     loss = loss_function(
-        logits, targets, sampled_loss_mask, depth.unsqueeze(0), loss_fn=kl_div_loss
+        logits, targets, sampled_loss_mask, depth.unsqueeze(0), loss_fn=loss_fn
     )
 
     with torch.no_grad():