Repository: allenai/OLMo-core
Pull request #244: Add option to cancel run after first evals
Author: login: epwalsh
Created: 2025-04-03T21:14:46Z | Merged: 2025-04-07T21:53:17Z

You can use this feature to only run in-loop evals without training any longer. For example, add these flags to your launch command:

```bash
  --trainer.callbacks.lm_evaluator.eval_on_startup=true \
  --trainer.callbacks.lm_evaluator.cancel_after_first_eval=true \
  --trainer.callbacks.downstream_evaluator.eval_on_startup=true \
  --trainer.callbacks.downstream_evaluator.cancel_after_first_eval=true
```

This also includes some minor fixes to how we log the speed of in-loop evals.


Diff:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c9638e854..e447760fa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,7 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `auto_resume` option to `CometCallback` for resume an existing run.
 - (BETA) Added methods `load_hf_model` and `save_hf_model` for saving supported OLMo Core models to HF transformers format.
 Also added lower-level methods for converting state between the formats.
-- Added the ability to run an eval at startup
+- Added the ability to run the evaluator callback on `.pre_train()` by setting `eval_on_startup=True`, and to cancel the run after the first time evals run by setting `cancel_after_first_eval=True`.
 - Added support for label mask files with numpy FSL datasets.
 
 ### Changed
diff --git a/src/examples/llama/train.py b/src/examples/llama/train.py
index 52a6117e2..d0cfbec09 100644
--- a/src/examples/llama/train.py
+++ b/src/examples/llama/train.py
@@ -22,6 +22,7 @@
 from olmo_core.nn.transformer import TransformerConfig
 from olmo_core.optim import AdamWConfig, CosWithWarmup, OptimGroupOverride
 from olmo_core.train import (
+    Duration,
     TrainerConfig,
     prepare_training_environment,
     teardown_training_environment,
@@ -32,6 +33,7 @@
     ConfigSaverCallback,
     DownstreamEvaluatorCallbackConfig,
     GPUMemoryMonitorCallback,
+    LMEvaluatorCallbackConfig,
     ProfilerCallback,
     WandBCallback,
 )
@@ -59,6 +61,8 @@
     f"{DATA_ROOT}/c4-train.00900-00999.npy",
     f"{DATA_ROOT}/c4-train.01000-01023.npy",
 ]
+EVAL_DATA_PATHS = [f"{DATA_ROOT}/c4-validation.00000-00008.npy"]
+DATA_WORK_DIR = "/tmp/dataset-cache"
 
 
 @dataclass
@@ -84,7 +88,7 @@ def build_config(run_name: str, overrides: List[str]) -> ExperimentConfig:
         sequence_length=SEQUENCE_LENGTH,
         max_target_sequence_length=8192,
         tokenizer=tokenizer_config,
-        work_dir="/tmp/dataset-cache",
+        work_dir=DATA_WORK_DIR,
     )
 
     data_loader_config = NumpyDataLoaderConfig(
@@ -144,6 +148,21 @@ def build_config(run_name: str, overrides: List[str]) -> ExperimentConfig:
         )
         .with_callback("config_saver", ConfigSaverCallback())
         .with_callback("profiler", ProfilerCallback(enabled=False))
+        .with_callback(
+            "lm_evaluator",
+            LMEvaluatorCallbackConfig(
+                eval_dataset=NumpyDatasetConfig(
+                    paths=EVAL_DATA_PATHS,
+                    metadata=[{"label": "c4-validation"}],
+                    name=NumpyDatasetType.padded_fsl,
+                    sequence_length=SEQUENCE_LENGTH,
+                    tokenizer=tokenizer_config,
+                    work_dir=DATA_WORK_DIR,
+                ),
+                eval_interval=250,
+                eval_duration=Duration.steps(50),
+            ),
+        )
         .with_callback(
             "downstream_evaluator",
             DownstreamEvaluatorCallbackConfig(
diff --git a/src/olmo_core/train/__init__.py b/src/olmo_core/train/__init__.py
index 4cc21b306..dc5a71b80 100644
--- a/src/olmo_core/train/__init__.py
+++ b/src/olmo_core/train/__init__.py
@@ -51,7 +51,13 @@
 from ..io import add_cached_path_clients
 from ..utils import LogFilterType, get_default_device, prepare_cli_environment, seed_all
 from .checkpoint import Checkpointer, CheckpointerConfig
-from .common import Duration, DurationUnit, LoadStrategy, ReduceType
+from .common import (
+    Duration,
+    DurationUnit,
+    LoadStrategy,
+    MetricMergeStrategy,
+    ReduceType,
+)
 from .config import TrainerConfig
 from .trainer import Trainer
 
@@ -66,6 +72,7 @@
     "Duration",
     "DurationUnit",
     "ReduceType",
+    "MetricMergeStrategy",
 ]
 
 
diff --git a/src/olmo_core/train/callbacks/evaluator_callback.py b/src/olmo_core/train/callbacks/evaluator_callback.py
index 7a6d204e7..5bcd2e2fe 100644
--- a/src/olmo_core/train/callbacks/evaluator_callback.py
+++ b/src/olmo_core/train/callbacks/evaluator_callback.py
@@ -26,7 +26,7 @@
     move_to_device,
 )
 
-from ..common import Duration
+from ..common import Duration, MetricMergeStrategy
 from ..train_module import EvalBatchSizeUnit, EvalBatchSpec, TransformerTrainModule
 from .callback import Callback, CallbackConfig
 
@@ -60,6 +60,13 @@ class EvaluatorCallback(Callback):
     Whether to run an evaluation when the trainer starts up.
     """
 
+    cancel_after_first_eval: bool = False
+    """
+    If ``True``, cancel the run after running evals for the first time. 
+    This combined with ``eval_on_startup=True`` is useful if you just want to run in-loop evals
+    without training any longer.
+    """
+
     eval_duration: Duration = field(default_factory=lambda: Duration.epochs(1))
     """
     The duration to run each evaluator for.
@@ -119,16 +126,11 @@ def _perform_eval(self):
                     with cuda_sync_debug_mode(0):
                         evaluator.update_metrics(batch, ce_loss, logits)
 
-                if eval_step % self.trainer.cancel_check_interval == 0:
-                    self.trainer.check_if_canceled()
-
-                if self.trainer.is_canceled:
-                    self._log_progress(evaluator, eval_step)
-                    return
-                elif self.eval_duration.due(step=eval_step, tokens=eval_tokens, epoch=1):
+                if self.eval_duration.due(step=eval_step, tokens=eval_tokens, epoch=1):
                     self._log_progress(evaluator, eval_step)
                     break
-                elif eval_step % self.log_interval == 0 or eval_step == evaluator.total_batches:
+
+                if eval_step % self.log_interval == 0 or eval_step == evaluator.total_batches:
                     self._log_progress(evaluator, eval_step)
 
             # NOTE: going to have a host-device sync here but that's okay. It's only once
@@ -144,7 +146,7 @@ def _perform_eval(self):
 
             evaluator_times.append(time.monotonic() - start_time)
             evaluator_names.append(evaluation_names)
-            evaluator_bs.append(evaluator.total_batches)
+            evaluator_bs.append(eval_step)
 
             log.info(
                 f"Finished {evaluator.name} evals in {time.monotonic() - start_time:.1f} seconds. Metrics:\n"
@@ -156,21 +158,30 @@ def _perform_eval(self):
             zip(evaluator_names, evaluator_bs, evaluator_times), key=lambda x: x[2]
         )
 
-        # Record evaluation speed
-        log.info("Evaluation speed:")
-        max_time_width = max(len(f"{t:.1f}") for t in evaluator_times)
-        max_batch_width = max(len(str(bs)) for bs in evaluator_bs)
+        # Record evaluation speed.
+        eval_speeds = []
         for names, bs, t in sorted_evaluators:
-            name = names[0]  # only use the name of the first metric for each downstream task
-            log.info(
-                f"    {t:>{max_time_width}.1f} sec ({bs:>{max_batch_width}} batches): {name} (+ variants)"
-            )
+            name = names[0]
+            eval_speeds.append(f"    {name} (+variants): {t:.1f} sec ({bs} batches)")
         total_time = sum(evaluator_times)
         total_bs = sum(int(bs) if bs is not None else 0 for bs in evaluator_bs)
-        log.info(f"    Total evaluation time: {total_time:.1f} seconds ({total_bs} batches)")
+        eval_speeds.append(
+            f"    Total evaluation time: {total_time:.1f} seconds ({total_bs} batches)"
+        )
+        log.info("Evaluation speed:\n" + "\n".join(eval_speeds))
+
+        self.trainer.record_metric(
+            "throughput/in-loop eval time (s)", total_time, merge_strategy=MetricMergeStrategy.sum
+        )
+        self.trainer.record_metric(
+            "throughput/in-loop eval batches", total_bs, merge_strategy=MetricMergeStrategy.sum
+        )
 
-        self.trainer.record_metric("throughput/in-loop eval time (s)", total_time)
-        self.trainer.record_metric("throughput/in-loop eval batches", total_bs)
+        if self.cancel_after_first_eval:
+            self.trainer.cancel_run(
+                "canceled from evaluator callback since 'cancel_after_first_eval' is set",
+                no_sync=True,  # 'no_sync' because we're calling this from all ranks at the same time.
+            )
 
     def _log_progress(self, evaluator: Evaluator, eval_step: int):
         if evaluator.total_batches is not None:
@@ -184,6 +195,7 @@ class LMEvaluatorCallbackConfig(CallbackConfig):
     eval_dataset: NumpyDatasetConfig
     eval_interval: int = 1000
     eval_on_startup: bool = False
+    cancel_after_first_eval: bool = False
     eval_duration: Duration = field(default_factory=lambda: Duration.epochs(1))
     log_interval: int = 5
     enabled: bool = True
@@ -240,6 +252,7 @@ def build(self, trainer: "Trainer") -> Optional[Callback]:
             eval_interval=self.eval_interval,
             log_interval=self.log_interval,
             eval_on_startup=self.eval_on_startup,
+            cancel_after_first_eval=self.cancel_after_first_eval,
             eval_duration=self.eval_duration,
         )
 
@@ -366,6 +379,7 @@ class DownstreamEvaluatorCallbackConfig(CallbackConfig):
     eval_interval: int = 1000
     eval_duration: Duration = field(default_factory=lambda: Duration.epochs(1))
     eval_on_startup: bool = False
+    cancel_after_first_eval: bool = False
     log_interval: int = 5
     enabled: bool = True
 
@@ -404,6 +418,7 @@ def build(self, trainer: "Trainer") -> Optional[Callback]:
             evaluators=evaluators,
             eval_interval=self.eval_interval,
             eval_on_startup=self.eval_on_startup,
+            cancel_after_first_eval=self.cancel_after_first_eval,
             log_interval=self.log_interval,
             eval_duration=self.eval_duration,
         )
diff --git a/src/olmo_core/train/common.py b/src/olmo_core/train/common.py
index 269a746a9..c8f9edab3 100644
--- a/src/olmo_core/train/common.py
+++ b/src/olmo_core/train/common.py
@@ -128,6 +128,47 @@ class ReduceType(StrEnum):
     """
 
 
+class MetricMergeStrategy(StrEnum):
+    """
+    Determines how duplicate metrics are merged.
+    """
+
+    warn = "warn"
+    """
+    Warn when a duplicate is logged, keeping the current value.
+    """
+
+    latest = "latest"
+    """
+    The latest is used.
+    """
+
+    oldest = "oldest"
+    """
+    The oldest (first logged) is used.
+    """
+
+    mean = "mean"
+    """
+    When a duplicate is logged we take the average with the last value.
+    """
+
+    sum = "sum"
+    """
+    The sum of the duplicates is used.
+    """
+
+    max = "max"
+    """
+    Take the maximum value of the duplicates.
+    """
+
+    min = "min"
+    """
+    Take the minimum value of the duplicates.
+    """
+
+
 def reshape_inputs_for_loss(
     logits: torch.Tensor, labels: torch.Tensor
 ) -> Tuple[torch.Tensor, torch.Tensor]:
diff --git a/src/olmo_core/train/trainer.py b/src/olmo_core/train/trainer.py
index c67150fd6..dc6268406 100644
--- a/src/olmo_core/train/trainer.py
+++ b/src/olmo_core/train/trainer.py
@@ -56,6 +56,7 @@
     Duration,
     DurationUnit,
     LoadStrategy,
+    MetricMergeStrategy,
     ReduceType,
     TrainingProgress,
 )
@@ -529,15 +530,23 @@ def training_progress(self) -> TrainingProgress:
             time_remaining=time_remaining,
         )
 
-    def cancel_run(self, reason: str):
+    def cancel_run(self, reason: str, no_sync: bool = False):
         """
         Mark the run canceled.
 
         :param reason: The reason for canceling.
+        :param no_sync: Set this to ``True`` only if you're calling this from all ranks at the same
+            time, otherwise you'll get a distributed deadlock.
         """
-        #  self._canceled = True  # NOTE: important not to set this!! Leads to distributed hang.
+        if self.is_canceled:
+            return
+
         self._canceling_rank = get_rank()
         self._cancel_reason = reason
+        if no_sync:
+            self._canceled = True
+            log.warning(f"Run canceled from all ranks. Reason: {reason}")
+            barrier()
 
     def check_if_canceled(self):
         """
@@ -590,8 +599,7 @@ def fit(self):
         # It's possible that we tried restarting a run that had already finished.
         if self.training_complete:
             log.warning("Training already complete, ending run now")
-            self._shutdown_bookkeeping()
-            gc_cuda()
+            self._shutdown()
             return
 
         log.info("Callback order:")
@@ -605,6 +613,11 @@ def fit(self):
 
         barrier()
 
+        # Quick check if the run has already been canceled.
+        if self.is_canceled:
+            self._shutdown()
+            return
+
         # Install SIGTERM + SIGINT handlers.
         og_sigterm_handler = signal.signal(signal.SIGTERM, self._handle_os_signal)
         og_sigint_handler = signal.signal(signal.SIGINT, self._handle_os_signal)
@@ -629,13 +642,14 @@ def fit(self):
             callback.post_train()
 
         # Wait for any bookkeeping tasks to finish.
-        self._shutdown_bookkeeping()
-        gc_cuda()
+        self._shutdown()
         log.info("Training complete")
 
-    def _shutdown_bookkeeping(self):
+    def _shutdown(self):
+        self._log_metrics()
         self.thread_pool.shutdown(wait=True, cancel_futures=False)
         self._thread_pool = None
+        gc_cuda()
         barrier()
 
     def state_dict(self) -> TrainerStateDict:
@@ -805,6 +819,7 @@ def record_metric(
         value: Union[float, torch.Tensor],
         reduce_type: Optional[ReduceType] = None,
         namespace: Optional[str] = None,
+        merge_strategy: MetricMergeStrategy = MetricMergeStrategy.warn,
     ):
         """
         Record a new metric for the current step.
@@ -817,16 +832,41 @@ def record_metric(
         :param reduce_type: Specifies how to reduce the metric across the distributed process group.
             ``None`` means no reduction.
         :param namespace: A namespace to record the metric under, i.g. "train" or "optim".
+        :param merge_strategy: How to merge metrics when duplicates are logged.
         """
         if namespace is not None:
             name = f"{namespace.rstrip('/')}/{name.lstrip('/')}"
+
         if not isinstance(value, torch.Tensor):
             value = torch.tensor(value)
         else:
             value = get_local_tensor(value).float()
+
         if self.global_step not in self._metrics:
             self._metrics[self.global_step] = OrderedDict()
-        self._metrics[self.global_step][name] = value
+
+        step_metrics = self._metrics[self.global_step]
+
+        if name not in step_metrics or merge_strategy == MetricMergeStrategy.latest:
+            step_metrics[name] = value
+        elif merge_strategy == MetricMergeStrategy.sum:
+            step_metrics[name] = step_metrics[name] + value
+        elif merge_strategy == MetricMergeStrategy.mean:
+            step_metrics[name] = (step_metrics[name] + value) / 2
+        elif merge_strategy == MetricMergeStrategy.max:
+            step_metrics[name] = torch.max(step_metrics[name], value.to(step_metrics[name].device))
+        elif merge_strategy == MetricMergeStrategy.min:
+            step_metrics[name] = torch.min(step_metrics[name], value.to(step_metrics[name].device))
+        elif merge_strategy == MetricMergeStrategy.warn:
+            log.warning(
+                f"Attempting to log duplicate metric '{name}' for step {(self.global_step)}. "
+                "The latest value will be ignored."
+            )
+        elif merge_strategy == MetricMergeStrategy.oldest:
+            pass
+        else:
+            raise NotImplementedError(merge_strategy)
+
         # reduce type must be consistent to avoid issues
         if name in self._metrics_reduce_type and self._metrics_reduce_type[name] != reduce_type:
             raise RuntimeError(
@@ -922,7 +962,7 @@ def _sort_callbacks(self):
                 (k, cb)
                 for _, (k, cb) in sorted(
                     enumerate(self.callbacks.items()),
-                    key=lambda x: (x[1][1].priority, x[0]),
+                    key=lambda x: (x[1][1].priority, -1 * x[0]),
                     reverse=True,
                 )
             )


Reviews:
user: login: dirkgr
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T21:28:10Z
user: login: epwalsh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T21:41:00Z
user: login: epwalsh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T21:41:11Z
user: login: dirkgr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T18:30:51Z
user: login: epwalsh
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T18:57:28Z

Review comments:
diff_hunk: @@ -529,15 +529,21 @@ def training_progress(self) -> TrainingProgress:
             time_remaining=time_remaining,
         )
 
-    def cancel_run(self, reason: str):
+    def cancel_run(self, reason: str, no_sync: bool = False):
         """
         Mark the run canceled.
 
         :param reason: The reason for canceling.
+        :param no_sync: Set this to ``True`` only if you're calling this from all ranks at the same
+            time, otherwise you'll get a distributed deadlock.
         """
-        #  self._canceled = True  # NOTE: important not to set this!! Leads to distributed hang.
         self._canceling_rank = get_rank()
         self._cancel_reason = reason
+        if no_sync:
+            self._canceled = (
+                True  # NOTE: important not to set this otherwise!! Leads to distributed hang.
+            )
path: src/olmo_core/train/trainer.py
user: login: dirkgr
body: This is confusing. The comment says to not set this, but it is being set?
created_at: 2025-04-03T21:27:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -172,6 +179,9 @@ def _perform_eval(self):
         self.trainer.record_metric("throughput/in-loop eval time (s)", total_time)
         self.trainer.record_metric("throughput/in-loop eval batches", total_bs)
 
+        if self.cancel_after_first_eval:
+            self.trainer.cancel_run("canceled from evaluator callback", no_sync=True)
path: src/olmo_core/train/callbacks/evaluator_callback.py
user: login: dirkgr
body: Maybe a clearer message, like "canceled because `cancel_after_first_eval` is set"?
created_at: 2025-04-03T21:28:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -172,6 +179,9 @@ def _perform_eval(self):
         self.trainer.record_metric("throughput/in-loop eval time (s)", total_time)
         self.trainer.record_metric("throughput/in-loop eval batches", total_bs)
 
+        if self.cancel_after_first_eval:
+            self.trainer.cancel_run("canceled from evaluator callback", no_sync=True)
path: src/olmo_core/train/callbacks/evaluator_callback.py
user: login: epwalsh
body: Done
created_at: 2025-04-03T21:41:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -529,15 +529,21 @@ def training_progress(self) -> TrainingProgress:
             time_remaining=time_remaining,
         )
 
-    def cancel_run(self, reason: str):
+    def cancel_run(self, reason: str, no_sync: bool = False):
         """
         Mark the run canceled.
 
         :param reason: The reason for canceling.
+        :param no_sync: Set this to ``True`` only if you're calling this from all ranks at the same
+            time, otherwise you'll get a distributed deadlock.
         """
-        #  self._canceled = True  # NOTE: important not to set this!! Leads to distributed hang.
         self._canceling_rank = get_rank()
         self._cancel_reason = reason
+        if no_sync:
+            self._canceled = (
+                True  # NOTE: important not to set this otherwise!! Leads to distributed hang.
+            )
path: src/olmo_core/train/trainer.py
user: login: epwalsh
body: Done
created_at: 2025-04-03T21:41:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -119,16 +126,18 @@ def _perform_eval(self):
                     with cuda_sync_debug_mode(0):
                         evaluator.update_metrics(batch, ce_loss, logits)
 
-                if eval_step % self.trainer.cancel_check_interval == 0:
-                    self.trainer.check_if_canceled()
+                if (
+                    not self.cancel_after_first_eval
+                ):  # guard here in case we have multiple evaluators doing this
+                    if self.trainer.is_canceled:
+                        self._log_progress(evaluator, eval_step)
+                        return
path: src/olmo_core/train/callbacks/evaluator_callback.py
user: login: dirkgr
body: I find this confusing. If multiple callbacks cancel the trainer, we still want to run them all, right? The canceling only takes effect when it's time to do the next step?
created_at: 2025-04-04T18:30:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -119,16 +126,18 @@ def _perform_eval(self):
                     with cuda_sync_debug_mode(0):
                         evaluator.update_metrics(batch, ce_loss, logits)
 
-                if eval_step % self.trainer.cancel_check_interval == 0:
-                    self.trainer.check_if_canceled()
+                if (
+                    not self.cancel_after_first_eval
+                ):  # guard here in case we have multiple evaluators doing this
+                    if self.trainer.is_canceled:
+                        self._log_progress(evaluator, eval_step)
+                        return
path: src/olmo_core/train/callbacks/evaluator_callback.py
user: login: epwalsh
body: I just removed this block: a8de0efa

We really don't need it. I had it there initially because I thought a run could get canceled during an eval loop, in which case we should break out of the loop early. But now I think it's better to finish the eval loop.
created_at: 2025-04-04T18:57:27Z
author_association: CONTRIBUTOR

PR comments:
user: login: dirkgr
created_at: 2025-04-04T17:53:18Z
body: Things are happening in this PR? Did you find an issue?
author_association: CONTRIBUTOR
user: login: epwalsh
created_at: 2025-04-04T17:54:51Z
body: I'm done editing, just found some other things to fix with how the in-loop eval speed is logged. Ready for review again.
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-14T18:08:49Z
body: This PR has been released in [v2.1.0](https://github.com/allenai/OLMo-core/releases/tag/v2.1.0).
author_association: NONE