fix unsloth gradient accumulation in validation

Author: RETR0-OS

ModelForge/evaluation/metrics.py

@@ -21,6 +21,10 @@ def compute_causal_lm_metrics(eval_pred: EvalPrediction) -> Dict[str, float]:
         - Perplexity: exp(loss)
         - Loss: Cross-entropy loss
 
+        Note: When preprocess_logits_for_metrics is used, predictions contain
+        pre-computed scalar losses (one per eval batch) instead of full logit tensors.
+        This prevents RAM exhaustion during evaluation.
+
         Args:
             eval_pred: Evaluation prediction with predictions and labels
 
@@ -29,28 +33,16 @@ def compute_causal_lm_metrics(eval_pred: EvalPrediction) -> Dict[str, float]:
         """
         logger.info("Computing causal LM metrics")
 
-        import torch
-        import torch.nn.functional as F
-
-        predictions, labels = eval_pred.predictions, eval_pred.label_ids
-
-        logits = torch.from_numpy(np.array(predictions)).float()
-        lbl = torch.from_numpy(np.array(labels))
-
-        shift_logits = logits[..., :-1, :].contiguous()
-        shift_labels = lbl[..., 1:].contiguous()
+        predictions = eval_pred.predictions
 
-        loss = F.cross_entropy(
-            shift_logits.view(-1, shift_logits.size(-1)),
-            shift_labels.view(-1).long(),
-            ignore_index=-100,
-            reduction="mean",
-        )
-        perplexity = float(torch.exp(loss))
+        # predictions is shape (num_eval_batches, 1) of pre-computed scalar losses
+        # (from _preprocess_logits_for_metrics in sft_strategy.py)
+        mean_loss = float(np.mean(predictions))
+        perplexity = float(np.exp(mean_loss))
 
         metrics = {
             "perplexity": perplexity,
-            "eval_loss": float(loss),
+            "eval_loss": mean_loss,
         }
 
         logger.info(f"Causal LM metrics: {metrics}")

ModelForge/services/training_service.py

@@ -432,7 +432,16 @@ def train_model(
 
             # Train
             self.training_status["message"] = "Training in progress..."
-            trainer.train()
+            provider_name = config.get("provider", "")
+            if provider_name == "unsloth":
+                try:
+                    from unsloth import unsloth_train
+                    unsloth_train(trainer)
+                except ImportError:
+                    logger.warning("unsloth_train not available, falling back to trainer.train()")
+                    trainer.train()
+            else:
+                trainer.train()
 
             # Save model
             self.training_status["message"] = "Saving model..."

ModelForge/strategies/qlora_strategy.py

@@ -22,6 +22,7 @@
 from trl import SFTTrainer, SFTConfig
 
 from ..logging_config import logger
+from .sft_strategy import _preprocess_logits_for_metrics
 
 
 class QLoRAStrategy:
@@ -269,6 +270,7 @@ def create_trainer(
             peft_config=peft_config,
             callbacks=callbacks or [],
             compute_metrics=compute_metrics,
+            preprocess_logits_for_metrics=_preprocess_logits_for_metrics,
         )
 
         logger.info("QLoRA SFTTrainer created successfully")

ModelForge/strategies/sft_strategy.py

@@ -23,6 +23,37 @@
 from ..logging_config import logger
 
 
+def _preprocess_logits_for_metrics(logits, labels):
+    """
+    Reduce logit tensors to scalar loss per batch before CPU accumulation.
+
+    This prevents RAM exhaustion when the Trainer accumulates eval predictions.
+    Instead of storing full logit tensors (batch × seq_len × vocab_size),
+    we compute the loss on GPU and return only a scalar, reducing memory by
+    orders of magnitude (e.g., 19.7 GB/batch → bytes/batch).
+
+    Args:
+        logits: Model logits (batch, seq_len, vocab_size)
+        labels: Token labels (batch, seq_len)
+
+    Returns:
+        Scalar loss tensor for this batch, shape (1,)
+    """
+    import torch.nn.functional as F
+
+    shift_logits = logits[..., :-1, :].contiguous()
+    shift_labels = labels[..., 1:].contiguous()
+
+    loss = F.cross_entropy(
+        shift_logits.view(-1, shift_logits.size(-1)),
+        shift_labels.view(-1).long(),
+        ignore_index=-100,
+        reduction="mean",
+    )
+
+    return loss.unsqueeze(0)
+
+
 class SFTStrategy:
     """Supervised Fine-Tuning strategy using TRL's SFTTrainer."""
 
@@ -248,6 +279,7 @@ def create_trainer(
             peft_config=peft_config,
             callbacks=callbacks or [],
             compute_metrics=compute_metrics,
+            preprocess_logits_for_metrics=_preprocess_logits_for_metrics,
         )
 
         logger.info("SFTTrainer created successfully")