eole-nlp · vince62s · Jul 28, 2026 · Jul 26, 2026 · Jul 28, 2026
diff --git a/eole/config/training.py b/eole/config/training.py
@@ -184,6 +184,11 @@ class TrainingConfig(
         description="If the norm of the gradient vector exceeds this value, "
         "renormalize it to have the norm equal to max_grad_norm.",
     )
+    empty_cache_steps: int = Field(
+        default=0,
+        description="Every N optimizer steps, clear the active PyTorch device cache. "
+        "On MPS, also clear the graph cache when supported. Disabled when set to 0.",
+    )
     dropout: List[float] = Field(default=[0.3], description="Dropout probability.")
     attention_dropout: List[float] = Field(default=[0.1], description="Attention dropout probability.")
     dropout_steps: List[int] = Field(default=[0], description="Steps at which dropout changes.")

diff --git a/eole/tests/test_trainer_cache.py b/eole/tests/test_trainer_cache.py
@@ -0,0 +1,96 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+from eole.trainer import Trainer, TrainerConfig
+from eole.utils.misc import clear_gpu_cache
+
+
+class TestTrainerCacheClearing(unittest.TestCase):
+    def test_empty_cache_steps_disabled_does_not_clear_cache(self):
+        trainer = object.__new__(Trainer)
+        trainer.config = TrainerConfig(empty_cache_steps=0)
+
+        with patch("eole.trainer.clear_gpu_cache") as clear_cache:
+            trainer._maybe_clear_device_cache(step=1)
+
+        clear_cache.assert_not_called()
+
+    def test_empty_cache_steps_interval_mismatch_does_not_clear_cache(self):
+        trainer = object.__new__(Trainer)
+        trainer.config = TrainerConfig(empty_cache_steps=10)
+
+        with patch("eole.trainer.clear_gpu_cache") as clear_cache:
+            trainer._maybe_clear_device_cache(step=9)
+
+        clear_cache.assert_not_called()
+
+    def test_empty_cache_steps_interval_match_clears_cache(self):
+        trainer = object.__new__(Trainer)
+        trainer.config = TrainerConfig(empty_cache_steps=10)
+
+        with patch("eole.trainer.clear_gpu_cache") as clear_cache:
+            trainer._maybe_clear_device_cache(step=10)
+
+        clear_cache.assert_called_once_with()
+
+    def test_training_loop_does_not_clear_cache_when_optimizer_step_is_skipped(self):
+        trainer = object.__new__(Trainer)
+        trainer.config = TrainerConfig(empty_cache_steps=10, average_decay=0, n_gpu=0)
+        trainer.optim = SimpleNamespace(training_step=10, learning_rate=Mock(return_value=0.1))
+        trainer.scoring_preparator = SimpleNamespace(transforms=None)
+        trainer.model_saver = None
+        trainer.report_manager = None
+        trainer.earlystopper = None
+
+        trainer._accum_batches = Mock(return_value=iter([([], 1)]))
+        trainer._update_scheduled_params = Mock()
+        trainer._train_step = Mock()
+        trainer._maybe_clear_device_cache = Mock()
+
+        with patch("eole.trainer.clear_gpu_cache"):
+            trainer.train(train_iter=[], train_steps=10, valid_steps=1000)
+
+        trainer._maybe_clear_device_cache.assert_not_called()
+
+    def test_empty_cache_steps_clears_cuda_cache_when_available(self):
+        torch_mock = SimpleNamespace(
+            cuda=SimpleNamespace(is_available=Mock(return_value=True), empty_cache=Mock()),
+            backends=SimpleNamespace(mps=SimpleNamespace(is_available=Mock(return_value=False))),
+            mps=SimpleNamespace(empty_cache=Mock()),
+        )
+
+        with patch("eole.utils.misc.torch", torch_mock):
+            clear_gpu_cache()
+
+        torch_mock.cuda.empty_cache.assert_called_once_with()
+        torch_mock.mps.empty_cache.assert_not_called()
+
+    def test_empty_cache_steps_clears_mps_cache_and_graph_cache_when_available(self):
+        torch_mock = SimpleNamespace(
+            cuda=SimpleNamespace(is_available=Mock(return_value=False), empty_cache=Mock()),
+            backends=SimpleNamespace(mps=SimpleNamespace(is_available=Mock(return_value=True))),
+            mps=SimpleNamespace(empty_cache=Mock(), clear_graph_cache=Mock()),
+        )
+
+        with patch("eole.utils.misc.torch", torch_mock):
+            clear_gpu_cache()
+
+        torch_mock.mps.empty_cache.assert_called_once_with()
+        torch_mock.mps.clear_graph_cache.assert_called_once_with()
+
+    def test_empty_cache_steps_clears_mps_cache_without_graph_cache_when_unavailable(self):
+        torch_mock = SimpleNamespace(
+            cuda=SimpleNamespace(is_available=Mock(return_value=False), empty_cache=Mock()),
+            backends=SimpleNamespace(mps=SimpleNamespace(is_available=Mock(return_value=True))),
+            mps=SimpleNamespace(empty_cache=Mock()),
+        )
+
+        with patch("eole.utils.misc.torch", torch_mock):
+            clear_gpu_cache()
+
+        torch_mock.mps.empty_cache.assert_called_once_with()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/eole/trainer.py b/eole/trainer.py
@@ -36,6 +36,7 @@ class TrainerConfig:
     attention_dropout: List[float] = None
     dropout_steps: List[int] = None
     zero_out_prompt_loss: bool = False
+    empty_cache_steps: int = 0
     estim_loss_lambda: List[float] = None
     estim_loss_lambda_steps: List[int] = None
 
@@ -300,6 +301,8 @@ def train(
 
             # Core training step (overridable by subclasses)
             self._train_step(batches, normalization, total_stats, report_stats, step=step)
+            if self.optim.training_step > step:
+                self._maybe_clear_device_cache(step)
 
             # Update moving average
             if self.config.average_decay > 0 and i % self.config.average_every == 0:
@@ -366,6 +369,11 @@ def _train_step(
         """
         self._process_accumulated_batches(batches, normalization, total_stats, report_stats)
 
+    def _maybe_clear_device_cache(self, step: int):
+        if self.config.empty_cache_steps <= 0 or step % self.config.empty_cache_steps != 0:
+            return
+        clear_gpu_cache()
+
     def _update_scheduled_params(self, step: int):
         """Update all scheduled parameters."""
         if self.dropout_scheduler.update(step):
@@ -656,6 +664,7 @@ def build_trainer(config, device_id, model, vocabs, optim, model_saver=None):
         attention_dropout=running_config.attention_dropout,
         dropout_steps=running_config.dropout_steps,
         zero_out_prompt_loss=running_config.zero_out_prompt_loss,
+        empty_cache_steps=running_config.empty_cache_steps,
         estim_loss_lambda=running_config.estim_loss_lambda,
         estim_loss_lambda_steps=running_config.estim_loss_lambda_steps,
     )

diff --git a/eole/utils/misc.py b/eole/utils/misc.py
@@ -77,6 +77,8 @@ def clear_gpu_cache():
         torch.cuda.empty_cache()
     if torch.backends.mps.is_available():
         torch.mps.empty_cache()
+        if hasattr(torch.mps, "clear_graph_cache"):
+            torch.mps.clear_graph_cache()
 
 
 def get_device_type():