Fix VRAM usage during evaluation

README.md

@@ -121,6 +121,7 @@ When ghost clipping is enabled, `privatize_trainer` automatically inspects the p
 | `poisson_sampling` | `True` | Use Poisson sub-sampling for privacy amplification. |
 | `grad_sample_mode` | `"hooks"` | Opacus grad sample mode (`"hooks"` or `"ew"`). Use `"ew"` for ghost clipping. |
 | `accountant` | `"rdp"` | Privacy accountant type (passed to Opacus). |
+| `epsilon_log_mode` | `"eval"` | When to log the expended privacy budget: `"none"`, `"train"`, `"eval"`, or `"both"`. |
 
 ## Docs
 

docs/configuration.md

@@ -14,6 +14,7 @@ All privacy-related parameters are configured through the `PrivacyArguments` dat
 | `clipping` | `"flat"` | Clipping strategy: `"flat"`, `"adaptive"`, or `"per_layer"` |
 | `poisson_sampling` | `True` | Use Poisson sub-sampling for privacy amplification |
 | `grad_sample_mode` | `"hooks"` | Opacus grad sample mode: `"hooks"` or `"ghost"` |
+| `epsilon_log_mode` | `"eval"` | When to log the expended privacy budget: `"none"`, `"train"`, `"eval"`, or `"both"` |
 
 ### Adaptive Clipping (AdaClip) Parameters
 
@@ -65,7 +66,7 @@ privacy_args = PrivacyArguments(
 
 `DPTrainer` automatically monitors the privacy budget during training. If a `target_epsilon` is set, training will stop early when the budget is exhausted. This is handled by the built-in `DPCallback`.
 
-Privacy metrics (`privacy_epsilon`) are logged alongside standard training metrics during evaluation.
+Privacy metrics (`privacy_epsilon`) can be logged during training, evaluation, or both via `epsilon_log_mode`.
 
 ## Clipping Strategies
 

src/dptrainer/privacy_arguments.py

@@ -1,6 +1,6 @@
 import warnings
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Literal, Optional
 
 from opacus.accountants.utils import get_noise_multiplier
 from transformers import logging
@@ -25,6 +25,7 @@ class PrivacyArguments:
         noise_multiplier (Optional[float]): Noise multiplier for DP training.
         target_epsilon (Optional[float]): Target epsilon at end of training (mutually exclusive with noise multiplier).
         target_delta (Optional[float]): Target delta, defaults to 1/N.
+        epsilon_log_mode (Literal): When to log the expended privacy budget: "none", "train", "eval", or "both".
     """
     accountant: str = field(default="rdp", metadata={"help": "Accountant mechanism to use for DP training"})
     grad_sample_mode: str = field(default="hooks", metadata={"help": "Grad sample mode of Opacus"})
@@ -42,6 +43,14 @@ class PrivacyArguments:
     target_epsilon: Optional[float] = field(default=None, metadata={
         "help": "Target epsilon at end of training (mutually exclusive with noise multiplier)"})
     target_delta: Optional[float] = field(default=None, metadata={"help": "Target delta, defaults to 1/N"})
+    epsilon_log_mode: Literal["none", "train", "eval", "both"] = field(
+        default="eval",
+        metadata={"help": "When to log the expended privacy budget: 'none', 'train', 'eval', or 'both'"},
+    )
+
+    def __post_init__(self):
+        if self.epsilon_log_mode not in {"none", "train", "eval", "both"}:
+            raise ValueError("Invalid epsilon_log_mode. Must be one of 'none', 'train', 'eval', or 'both'.")
 
     @classmethod
     def low_privacy(cls):

src/dptrainer/trainer.py

@@ -1,6 +1,6 @@
 import math
 import warnings
-from typing import Optional, Union, Callable
+from typing import Optional, Union
 
 import datasets
 import torch
@@ -10,7 +10,7 @@
 from opacus.utils.batch_memory_manager import wrap_data_loader
 from torch import nn
 from transformers import (
-    Trainer, logging, modeling_utils, TrainingArguments, PreTrainedModel, TrainerCallback
+    Trainer, logging, TrainingArguments, PreTrainedModel, TrainerCallback
 )
 from opacus.utils.fast_gradient_clipping_utils import DPLossFastGradientClipping
 from dptrainer.utils import set_loss_function
@@ -27,7 +27,6 @@ def __init__(
             args: TrainingArguments = None,
             train_dataset: Union[datasets.Dataset, torch.utils.data.Dataset] = None,
             privacy_args: PrivacyArguments = None,
-            compute_metrics: Optional[Callable] = None,
             callbacks: Optional[list[TrainerCallback]] = None,
             **kwargs
     ):
@@ -38,7 +37,6 @@ def __init__(
             args (TrainingArguments): Training arguments.
             train_dataset (Union[datasets.Dataset, torch.utils.data.Dataset]): Training dataset.
             privacy_args (PrivacyArguments): Privacy arguments for differential private training.
-            compute_metrics (Optional[Callable]): Custom evaluation metrics.
             callbacks (Optional[list[TrainerCallback]]): Training callbacks.
             **kwargs: Additional keyword arguments passed to Trainer.
         """
@@ -87,32 +85,19 @@ def __init__(
             wrap_model=False,
         )
 
-        dp_callback = DPCallback(
+        self.dp_callback = DPCallback(
             accountant=self.privacy_args.accountant,
             gradient_accumulation_steps=args.gradient_accumulation_steps,
             target_delta=self.privacy_args.target_delta,
             max_epsilon=self.privacy_args.target_epsilon,
         )
         callbacks = callbacks or []
-        callbacks.append(dp_callback)
+        callbacks.append(self.dp_callback)
 
-        def compute_privacy_metrics(*args, compute_result: bool = True, **kwargs):
-            if compute_metrics:
-                metrics = compute_metrics(*args, compute_result, **kwargs) or {}
-            else:
-                metrics = {}
-
-            if compute_result:
-                privacy_metrics = dp_callback.get_privacy_metrics()
-                metrics.update(privacy_metrics)
-
-            return metrics
-
-        super().__init__(model=model, args=args, train_dataset=train_dataset, callbacks=callbacks,
-                         compute_metrics=compute_privacy_metrics, **kwargs)
+        super().__init__(model=model, args=args, train_dataset=train_dataset, callbacks=callbacks, **kwargs)
 
         optimizer = self.create_optimizer()
-        optimizer.attach_step_hook(dp_callback.get_optimizer_callback(sample_rate=sample_rate))
+        optimizer.attach_step_hook(self.dp_callback.get_optimizer_callback(sample_rate=sample_rate))
 
         if self.privacy_args and self.privacy_args.grad_sample_mode == "ghost":
             criterion = self.model.loss_function
@@ -209,4 +194,26 @@ def detach_model(self) -> nn.Module:
         """
         self.hooks.cleanup()
 
-        return self.model
+        return self.model
+
+    def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
+        mode = "train" if self.model.training else "eval"
+        log_mode = self.privacy_args.epsilon_log_mode
+        privacy_metrics = self.dp_callback.get_privacy_metrics()
+
+        if mode == "eval":
+            metric_prefix = next(
+                (
+                    key.removesuffix(suffix)
+                    for suffix in ("_loss", "_runtime")
+                    for key in logs
+                    if key.endswith(suffix)
+                ),
+                "eval",
+            )
+            privacy_metrics = {f"{metric_prefix}_{key}": val for key, val in privacy_metrics.items()}
+
+        if log_mode == "both" or log_mode == mode:
+            logs.update(privacy_metrics)
+
+        super().log(logs, start_time)

tests/dptrainer/test_dp_trainer.py

@@ -1,6 +1,6 @@
 """Tests for DPTrainer."""
 
-from unittest.mock import Mock, patch, MagicMock
+from unittest.mock import Mock, call, patch, MagicMock
 import pytest
 import torch
 from torch.utils.data import IterableDataset
@@ -324,45 +324,85 @@ def test_detach_model(
         assert detached_model is simple_model
 
 
-class TestDPTrainerComputeMetrics:
-    """Test privacy metrics integration."""
+class TestDPTrainerPrivacyLogging:
+    """Test privacy metrics logging behavior."""
 
-    @patch("dptrainer.trainer.prepare_module")
-    def test_compute_metrics_integration(
-        self, mock_wrap_model, simple_model, small_dataset, training_args
-    ):
-        """Test that privacy metrics are integrated into compute_metrics."""
-        mock_controller = Mock()
-        mock_wrap_model.return_value = mock_controller
+    def test_privacy_log_mode_defaults_to_eval(self):
+        assert PrivacyArguments().epsilon_log_mode == "eval"
 
-        custom_metrics_called = False
+    @patch("dptrainer.trainer.Trainer.log")
+    def test_log_privacy_metrics_on_train_only(self, mock_super_log):
+        trainer = DPTrainer.__new__(DPTrainer)
+        trainer.model = Mock(training=True)
+        trainer.dp_callback = Mock()
+        trainer.dp_callback.get_privacy_metrics.return_value = {"privacy_epsilon": 1.23}
+        trainer.privacy_args = PrivacyArguments(epsilon_log_mode="train")
 
-        def custom_compute_metrics(eval_pred, compute_result=True):
-            nonlocal custom_metrics_called
-            custom_metrics_called = True
-            return {"accuracy": 0.95}
+        logs = {}
 
-        privacy_args = PrivacyArguments(noise_multiplier=1.0)
+        DPTrainer.log(trainer, logs)
 
-        trainer = DPTrainer(
-            model=simple_model,
-            args=training_args,
-            train_dataset=small_dataset,
-            privacy_args=privacy_args,
-            compute_metrics=custom_compute_metrics,
-        )
+        assert logs == {"privacy_epsilon": 1.23}
+        mock_super_log.assert_called_once_with(logs, None)
+
+    @patch("dptrainer.trainer.Trainer.log")
+    def test_log_privacy_metrics_on_eval_only(self, mock_super_log):
+        trainer = DPTrainer.__new__(DPTrainer)
+        trainer.model = Mock(training=False)
+        trainer.dp_callback = Mock()
+        trainer.dp_callback.get_privacy_metrics.return_value = {"privacy_epsilon": 1.23}
+        trainer.privacy_args = PrivacyArguments(epsilon_log_mode="eval")
+
+        logs = {}
+
+        DPTrainer.log(trainer, logs)
+
+        assert logs == {"eval_privacy_epsilon": 1.23}
+        mock_super_log.assert_called_once_with(logs, None)
+
+    @patch("dptrainer.trainer.Trainer.log")
+    def test_log_privacy_metrics_uses_evaluation_metric_prefix(self, mock_super_log):
+        trainer = DPTrainer.__new__(DPTrainer)
+        trainer.model = Mock(training=False)
+        trainer.dp_callback = Mock()
+        trainer.dp_callback.get_privacy_metrics.return_value = {"privacy_epsilon": 1.23}
+        trainer.privacy_args = PrivacyArguments(epsilon_log_mode="eval")
+
+        validation_logs = {"validation_loss": 0.5}
+        DPTrainer.log(trainer, validation_logs)
+
+        dataset_logs = {"eval_test_loss": 0.5}
+        DPTrainer.log(trainer, dataset_logs)
+
+        runtime_logs = {"validation_runtime": 0.5}
+        DPTrainer.log(trainer, runtime_logs)
+
+        assert validation_logs == {"validation_loss": 0.5, "validation_privacy_epsilon": 1.23}
+        assert dataset_logs == {"eval_test_loss": 0.5, "eval_test_privacy_epsilon": 1.23}
+        assert runtime_logs == {"validation_runtime": 0.5, "validation_privacy_epsilon": 1.23}
+        mock_super_log.assert_has_calls([
+            call(validation_logs, None),
+            call(dataset_logs, None),
+            call(runtime_logs, None),
+        ])
+
+    @patch("dptrainer.trainer.Trainer.log")
+    def test_log_privacy_metrics_on_both(self, mock_super_log):
+        trainer = DPTrainer.__new__(DPTrainer)
+        trainer.model = Mock(training=True)
+        trainer.dp_callback = Mock()
+        trainer.dp_callback.get_privacy_metrics.return_value = {"privacy_epsilon": 1.23}
+        trainer.privacy_args = PrivacyArguments(epsilon_log_mode="both")
 
-        # The compute_metrics function should be wrapped
-        assert trainer.compute_metrics is not None
+        train_logs = {"random_metric": 0.5}
+        DPTrainer.log(trainer, train_logs)
 
-        # Mock eval_pred
-        mock_eval_pred = Mock()
+        assert train_logs == {"random_metric": 0.5, "privacy_epsilon": 1.23}
 
-        # Call the wrapped compute_metrics
-        result = trainer.compute_metrics(mock_eval_pred)
+        trainer.model.training = False
+        eval_logs = {"random_metric": 0.5}
+        DPTrainer.log(trainer, eval_logs)
 
-        # Should include both custom metrics and privacy metrics
-        assert "accuracy" in result
-        assert result["accuracy"] == 0.95
-        # Privacy metrics should be added (even if zero at start)
-        assert "privacy_epsilon" in result or "privacy_advantage" in result
+        assert eval_logs == {"random_metric": 0.5, "eval_privacy_epsilon": 1.23}
+
+        mock_super_log.assert_has_calls([call(train_logs, None), call(eval_logs, None)])