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

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

Author: christinaexyou

scripts/train.py

@@ -16,6 +16,7 @@
     DEFAULT_REQUEST_TIMEOUT,
 )
 from speculators.model import SpeculatorModel
+from speculators.models.metrics import resolve_loss_fn
 from speculators.models.eagle3.data import shift_batch
 from speculators.train.data import (
     ArrowDataset,
@@ -574,6 +575,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",
+        choices=["kl", "ce"],
+        help=(
+            "Loss function used during draft model training. "
+            "'kl' = KL divergence (default). "
+            "'ce' = cross-entropy."
+        ),
+    )
     parser.add_argument(
         "--seed", type=int, default=42, help="Random seed for reproducibility"
     )
@@ -679,7 +691,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
@@ -13,6 +13,7 @@
 from speculators.models.dflash import DFlashSpeculatorConfig
 from speculators.models.dflash.attention import create_anchor_block_mask_mod
 from speculators.models.dflash.metrics import compute_metrics
+from speculators.models.metrics import resolve_loss_fn
 from speculators.models.dflash.model_definitions import Qwen3DFlashDecoderLayer
 from speculators.models.dflash.utils import (
     get_base_indices_for_anchored_blocks,
@@ -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=None,
         **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,5 +1,6 @@
 """Metrics and loss functions for DFlash draft model."""
 
+from collections.abc import Callable
 from functools import partial
 from typing import Any
 
@@ -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] = ce_loss,
 ) -> tuple[torch.Tensor, dict]:
     """Compute loss and accuracy metrics for draft model predictions.
 
@@ -28,6 +30,7 @@ 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:
@@ -44,7 +47,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

@@ -13,6 +13,7 @@
     extend_mask_for_draft_tokens,
 )
 from speculators.models.eagle3.metrics import compute_metrics
+from speculators.models.metrics import resolve_loss_fn
 from speculators.models.eagle3.model_definitions import model_classes
 from speculators.models.utils import resolve_target_layer_ids
 from speculators.proposals.greedy import GreedyTokenProposalConfig
@@ -130,6 +131,7 @@ def forward(  # noqa: C901
         ttt_steps: int = 3,
         ttt_step_loss_decay: float = 1.0,
         use_off_policy_tokens: bool = False,
+        loss_fn=None,
         **kwargs,
     ):
         device = hidden_states.device
@@ -222,6 +224,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 +327,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] = kl_div_loss,
 ) -> tuple[torch.Tensor, dict]:
     """Compute metrics for a given ttt_step.
 
@@ -64,13 +66,15 @@ 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.
     """
+    loss_fn = loss_fn or 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 +92,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"`` 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": 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

@@ -12,6 +12,7 @@
 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
+from speculators.models.metrics import resolve_loss_fn
 from speculators.models.peagle.metrics import compute_metrics
 from speculators.models.utils import resolve_target_layer_ids
 from speculators.proposals.greedy import GreedyTokenProposalConfig
@@ -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=None,
         **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] = kl_div_loss,
 ) -> tuple[torch.Tensor, dict[str, Any]]:
     """Compute loss and accuracy metrics for P-EAGLE predictions.
 
@@ -29,10 +31,12 @@ 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)
     """
+    loss_fn = loss_fn or kl_div_loss
     device = logits.device
 
     # TODO: batch size is always 1 for P-EAGLE; unsqueeze is only to match the
@@ -41,7 +45,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():

tests/unit/train/test_cli_args.py

@@ -0,0 +1,60 @@
+"""Tests for CLI arguments."""
+
+from scripts.train import parse_args
+from speculators.models.dflash.core import DFlashDraftModel
+from speculators.models.eagle3.core import Eagle3DraftModel
+from speculators.models.metrics import ce_loss, kl_div_loss
+from speculators.models.peagle.core import PEagleDraftModel
+
+
+def _parse(monkeypatch, extra: list[str]):
+    monkeypatch.setattr(
+        "sys.argv", ["train.py", "--verifier-name-or-path", "dummy"] + extra
+    )
+    return parse_args()
+
+
+# ---------------------------------------------------------------------------
+# Ensure CLI args flow correctly through vars(args) into get_trainer_kwargs
+# ---------------------------------------------------------------------------
+
+class TestDFlashLossFnCLI:
+    def test_default_uses_kl(self, monkeypatch):
+        args = _parse(monkeypatch, [])
+        train_kw, val_kw = DFlashDraftModel.get_trainer_kwargs(**vars(args))
+        assert train_kw["loss_fn"] is kl_div_loss
+        assert val_kw["loss_fn"] is kl_div_loss
+
+    def test_explicit_ce(self, monkeypatch):
+        args = _parse(monkeypatch, ["--loss-fn", "ce"])
+        train_kw, val_kw = DFlashDraftModel.get_trainer_kwargs(**vars(args))
+        assert train_kw["loss_fn"] is ce_loss
+        assert val_kw["loss_fn"] is ce_loss
+
+
+class TestEagle3LossFnCLI:
+    def test_default_uses_kl(self, monkeypatch):
+        args = _parse(monkeypatch, [])
+        train_kw, val_kw = Eagle3DraftModel.get_trainer_kwargs(**vars(args))
+        assert train_kw["loss_fn"] is kl_div_loss
+        assert val_kw["loss_fn"] is kl_div_loss
+
+    def test_explicit_ce(self, monkeypatch):
+        args = _parse(monkeypatch, ["--loss-fn", "ce"])
+        train_kw, val_kw = Eagle3DraftModel.get_trainer_kwargs(**vars(args))
+        assert train_kw["loss_fn"] is ce_loss
+        assert val_kw["loss_fn"] is ce_loss
+
+
+class TestPEagleLossFnCLI:
+    def test_default_uses_kl(self, monkeypatch):
+        args = _parse(monkeypatch, [])
+        train_kw, val_kw = PEagleDraftModel.get_trainer_kwargs(**vars(args))
+        assert train_kw["loss_fn"] is kl_div_loss
+        assert val_kw["loss_fn"] is kl_div_loss
+
+    def test_explicit_ce(self, monkeypatch):
+        args = _parse(monkeypatch, ["--loss-fn", "ce"])
+        train_kw, val_kw = PEagleDraftModel.get_trainer_kwargs(**vars(args))
+        assert train_kw["loss_fn"] is ce_loss
+        assert val_kw["loss_fn"] is ce_loss