refactor: deduplicate low-risk code paths

astrai/trainer/metric_util.py

@@ -94,21 +94,5 @@ def ctx_get_grad_snr(ctx):
     return tracker.snr
 
 
-def ctx_get_moe_aux_loss(ctx):
-    return ctx.strategy._moe_metrics.get("aux_loss")
-
-
-def ctx_get_router_entropy(ctx):
-    return ctx.strategy._moe_metrics.get("router_entropy")
-
-
-def ctx_get_dead_expert_fraction(ctx):
-    return ctx.strategy._moe_metrics.get("dead_expert_fraction")
-
-
-def ctx_get_load_imbalance_mean(ctx):
-    return ctx.strategy._moe_metrics.get("load_imbalance_mean")
-
-
-def ctx_get_load_imbalance_max(ctx):
-    return ctx.strategy._moe_metrics.get("load_imbalance_max")
+def ctx_get_moe_metric(ctx, key):
+    return ctx.strategy._moe_metrics.get(key)

astrai/trainer/strategy.py

@@ -1,6 +1,6 @@
 """Training strategy implementations with factory pattern."""
 
-from abc import ABC, abstractmethod
+from abc import ABC
 from typing import Callable, Dict, List, Optional, TypedDict, Union
 
 import torch
@@ -187,7 +187,6 @@ class BaseStrategy(ABC):
         self.extra_kwargs = kwargs
         self._rollout_runner = None
 
-    @abstractmethod
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
         """Compute loss for the given batch.
 
@@ -197,7 +196,7 @@ class BaseStrategy(ABC):
         Returns:
             Computed loss tensor
         """
-        raise NotImplementedError
+        return self.compute_loss_output(batch)["loss"]
 
     def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         return self._normalize_output(self.compute_loss(batch))
@@ -328,9 +327,6 @@ class SEQStrategy(BaseStrategy):
         super().__init__(model, device, **kwargs)
         self.label_smoothing = label_smoothing
 
-    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        return self.compute_loss_output(batch)["loss"]
-
     def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         input_ids, target_ids = batch["input_ids"], batch["target_ids"]
@@ -369,9 +365,6 @@ class SFTStrategy(BaseStrategy):
         super().__init__(model, device, **kwargs)
         self.label_smoothing = label_smoothing
 
-    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        return self.compute_loss_output(batch)["loss"]
-
     def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         input_ids, target_ids, position_ids, loss_mask = (
@@ -426,9 +419,6 @@ class DPOStrategy(BaseStrategy):
         self.beta = beta
         self.reduction = reduction
 
-    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        return self.compute_loss_output(batch)["loss"]
-
     def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         chosen_ids, rejected_ids = batch["chosen"], batch["rejected"]
@@ -553,9 +543,6 @@ class GRPOStrategy(BaseStrategy):
         if state_dict is not None:
             self.old_model.load_state_dict(state_dict)
 
-    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        return self.compute_loss_output(batch)["loss"]
-
     def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         prompts = batch["prompts"]

astrai/trainer/train_callback.py

@@ -3,6 +3,7 @@ import logging
 import os
 import sys
 import time
+from functools import partial
 from pathlib import Path
 from typing import IO, Callable, List, Optional, Protocol, runtime_checkable
 
@@ -17,15 +18,11 @@ from astrai.parallel import only_on_rank
 from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
-    ctx_get_dead_expert_fraction,
     ctx_get_grad_norm,
     ctx_get_grad_snr,
-    ctx_get_load_imbalance_max,
-    ctx_get_load_imbalance_mean,
     ctx_get_loss,
     ctx_get_lr,
-    ctx_get_moe_aux_loss,
-    ctx_get_router_entropy,
+    ctx_get_moe_metric,
     ctx_get_val_loss,
 )
 from astrai.trainer.train_context import TrainContext
@@ -262,11 +259,15 @@ class MetricCallback(TrainCallback):
             "val_loss": ctx_get_val_loss,
             "grad_norm": ctx_get_grad_norm,
             "grad_snr": ctx_get_grad_snr,
-            "moe_aux_loss": ctx_get_moe_aux_loss,
-            "router_entropy": ctx_get_router_entropy,
-            "dead_expert_fraction": ctx_get_dead_expert_fraction,
-            "load_imbalance_mean": ctx_get_load_imbalance_mean,
-            "load_imbalance_max": ctx_get_load_imbalance_max,
+            "moe_aux_loss": partial(ctx_get_moe_metric, key="aux_loss"),
+            "router_entropy": partial(ctx_get_moe_metric, key="router_entropy"),
+            "dead_expert_fraction": partial(
+                ctx_get_moe_metric, key="dead_expert_fraction"
+            ),
+            "load_imbalance_mean": partial(
+                ctx_get_moe_metric, key="load_imbalance_mean"
+            ),
+            "load_imbalance_max": partial(ctx_get_moe_metric, key="load_imbalance_max"),
         }
 
     def _metrics(self, context: TrainContext, names):

astrai/trainer/train_context.py

@@ -204,7 +204,6 @@ class TrainContextBuilder:
     def _create_dataloaders(
         self, context: TrainContext, train_dataset, val_dataset
     ) -> None:
-        cfg = self.config
         sampler_offset = context.consumed_samples // context.world_size
         if self._resume and sampler_offset > 0:
             samples_per_replica = (