Repository: mosaicml/composer

Issue:
Callbacks need access to computed metrics
Callbacks such as `EarlyStopping` would need access to the metric results. Currently, we compute the metrics, log the results, and then reset the metrics for the next epoch. Thus, simply attaching the metrics to state would not address this issue; we need a design where callbacks can access this time series data.



Discussion:

Resolving patch:
diff --git a/composer/core/state.py b/composer/core/state.py
--- a/composer/core/state.py
+++ b/composer/core/state.py
@@ -119,6 +119,43 @@ class State(Serializable):
             microbatch between :attr:`.Event.BATCH_START` and :attr:`.Event.BATCH_END`.
         batch_num_samples (int): The number of samples in the :attr:`batch`.
         batch_num_tokens (int): The number of tokens in the :attr:`batch`.
+        current_metrics (Dict[str, Dict[str, Any]]): The current computed metrics, organized by dataloader label
+            and then by metric name. The train dataloader is labeled ``'train'``. If not using an :class:`.Evaluator`,
+            the eval dataloader is labeled ``'eval'``. Otherwise, the evaluator label is used.
+
+            For example:
+
+            >>> trainer = Trainer(
+            ...     ...,
+            ...     compute_training_metrics=True,
+            ...     train_dataloader=train_dataloader,
+            ...     eval_dataloader=eval_dataloader,
+            ... )
+            >>> trainer.fit()
+            >>> trainer.state.current_metrics
+            {'train': {'Accuracy': tensor(...)}, 'eval': {'Accuracy': tensor(...)}}
+
+            Or, when using an :class:`.Evaluator`:
+
+            .. testsetup::
+
+                eval_1_dl = eval_dataloader
+                eval_2_dl = eval_dataloader
+
+            >>> from torchmetrics import Accuracy
+            >>> from composer.core import Evaluator
+            >>> trainer = Trainer(
+            ...     ...,
+            ...     compute_training_metrics=True,
+            ...     train_dataloader=train_dataloader,
+            ...     eval_dataloader=[
+            ...         Evaluator(label='eval1', dataloader=eval_1_dl, metrics=Accuracy()),
+            ...         Evaluator(label='eval2', dataloader=eval_2_dl, metrics=Accuracy()),
+            ...     ],
+            ... )
+            >>> trainer.fit()
+            >>> trainer.state.current_metrics
+            {'train': {'Accuracy': tensor(...)}, 'eval1': {'Accuracy': tensor(...)}, 'eval2': {'Accuracy': tensor(...)}}
 
         loss (torch.Tensor | Sequence[torch.Tensor]): The most recently computed loss.
         outputs (torch.Tensor | Sequence[torch.Tensor]): The most recently computed output from the model's forward pass.
@@ -146,6 +183,8 @@ class State(Serializable):
             +-----------------------+-------------------------------------------------------------+
             | rank_zero_seed        | The seed of the rank zero process.                          |
             +-----------------------+-------------------------------------------------------------+
+            | current_metrics       | The current metrics.                                        |
+            +-----------------------+-------------------------------------------------------------+
     """
 
     _max_duration: Time[int]
@@ -226,8 +265,11 @@ def __init__(
             "scaler",
             "timer",
             "rank_zero_seed",
+            "current_metrics",
         ]
 
+        self.current_metrics: Dict[str, Dict[str, Any]] = {}
+
     @property
     def seed(self):
         """The seed for the current rank."""
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
--- a/composer/trainer/trainer.py
+++ b/composer/trainer/trainer.py
@@ -614,7 +614,7 @@ def __init__(
                 self.evaluators.append(evaluator)
             else:
                 metrics = model.metrics(train=False)
-                evaluator = Evaluator(label="eval_dataset", dataloader=evaluator, metrics=metrics)
+                evaluator = Evaluator(label="eval", dataloader=evaluator, metrics=metrics)
                 self.evaluators.append(evaluator)
 
         self._eval_subset_num_batches = eval_subset_num_batches
@@ -795,12 +795,23 @@ def __init__(
 
         self._validate_every_n_batches = validate_every_n_batches
         self._validate_every_n_epochs = validate_every_n_epochs
-        self._compute_training_metrics = compute_training_metrics
         self._grad_clip_norm = grad_clip_norm
 
         if deterministic_mode:
             reproducibility.configure_deterministic_mode()
 
+        if compute_training_metrics:
+            warnings.warn(('Computing model evaluation metrics during training doubles the number of forward passes '
+                           'and may lead to a throughput degradation.'))
+            train_metrics = model.metrics(train=True)
+            if isinstance(train_metrics, Metric):
+                # Forcing metrics to be a MetricCollection simplifies logging results
+                train_metrics = MetricCollection([train_metrics])
+
+            self.train_metrics = self._ensure_metrics_device_and_dtype(train_metrics)
+        else:
+            self.train_metrics = None
+
         self.engine.run_event(Event.INIT)
 
         assert isinstance(self.state.model, ComposerModel)
@@ -912,33 +923,20 @@ def _ensure_metrics_device_and_dtype(self, metrics: MetricCollection):
 
         return metrics
 
-    def _compute_and_log_metrics(self,
-                                 metrics: Union[Metric, MetricCollection],
-                                 *,
-                                 is_train: bool,
-                                 is_batch: bool,
-                                 logging_label: str = ''):
-        """Computes metrics, logs the results, and resets the metrics.
+    def _compute_and_log_metrics(self, dataloader_label: str, log_level: LogLevel, metrics: MetricCollection):
+        """Computes metrics, logs the results, and updates the state.
 
         Args:
-            metrics (Metric | MetricCollection): The metrics to compute.
-            is_train (bool): True for training metrics, False for evaluation metrics.
-            is_batch (bool): True if logging at batch level, false for epoch level.
-            logging_label (str): Should be left as empty string if called for training metrics.
-                Should be the evaluator label if called on evaluator metrics.
+            dataloader_label (str): The dataloader label.
+            metrics (MetricCollection): The metrics to compute.
+            log_level (LogLevel): The LogLevel for logging metrics.
         """
         computed_metrics = metrics.compute()
-        for name, value in computed_metrics.items():
-            log_level = LogLevel.BATCH if is_batch else LogLevel.EPOCH
-            suffix = 'train' if is_train else 'val'
-
-            # default label given to evaluator created by val_dataset parameter
-            if not logging_label or logging_label == "eval_dataset":
-                label = f'{name.lower()}/{suffix}'
-            else:
-                label = f'{logging_label}/{name.lower()}_{suffix}'
-            self.logger.data(log_level, {label: value})
-        metrics.reset()
+        self.logger.data(
+            log_level=log_level,
+            data={f'metrics/{dataloader_label}/{name}': val for (name, val) in computed_metrics.items()},
+        )
+        self.state.current_metrics[dataloader_label] = computed_metrics
 
     def _spin_dataloaders(self):
         """Spin the dataloaders to restore sampler state.
@@ -971,19 +969,6 @@ def _train_loop(self) -> None:
         # print training start
         self.logger.data_fit({"trainer/algorithms": [str(algo) for algo in self.state.algorithms]})
 
-        if self._compute_training_metrics:
-            log.warning('Computing model evaluation metrics during training.'
-                        ' This doubles the number of forward passes and may lead'
-                        ' to a throughput degradation.')
-            train_metrics = self._original_model.metrics(train=True)
-            if isinstance(train_metrics, Metric):
-                # Forcing metrics to be a MetricCollection simplifies logging results
-                train_metrics = MetricCollection([train_metrics])
-
-            train_metrics = self._ensure_metrics_device_and_dtype(train_metrics)
-        else:
-            train_metrics = None
-
         self.engine.run_event(Event.FIT_START)
 
         self.state.scaler = ClosureGradScaler() if self._use_closures() else GradScaler()
@@ -1003,6 +988,9 @@ def _train_loop(self) -> None:
                 if int(self.state.timer.batch_in_epoch) == 0:
                     self.engine.run_event(Event.EPOCH_START)
                     self.logger.data_epoch({"epoch": int(self.state.timer.epoch)})
+                    if self.train_metrics is not None:
+                        # reset the metrics before every epoch
+                        self.train_metrics.reset()
 
                 # TODO: hasattr check will be removed while fixing https://github.com/mosaicml/composer/issues/424
                 if hasattr(self.state.train_dataloader, "sampler") and isinstance(self.state.train_dataloader.sampler,
@@ -1028,7 +1016,14 @@ def _train_loop(self) -> None:
                     if self.deepspeed_enabled:
                         self.state.batch = _fix_batch_precision_for_deepspeed(self.state.batch, self.state.precision)
 
-                    self._compute_metrics(train_metrics)
+                    if self.train_metrics is not None:
+                        self.state.model.eval()
+                        with torch.no_grad():
+                            for eval_microbatch in self._train_data_spec.split_batch(
+                                    self.state.batch, self.state.grad_accum):
+                                # TODO: Detect if self.run_event(Event.AFTER_DATALOADER) changes the training
+                                # data and if so print a warning that metrics may return unexpected results
+                                self.train_metrics.update(*self._original_model.validate(eval_microbatch))
 
                     self.state.model.train()
 
@@ -1061,10 +1056,6 @@ def _train_loop(self) -> None:
                         full_loss = total_loss.cpu().item()
                         self.logger.data_batch({'loss/train': full_loss / dist.get_world_size()})
 
-                    if self._compute_training_metrics:
-                        assert train_metrics is not None
-                        self._compute_and_log_metrics(train_metrics, is_train=True, is_batch=True)
-
                     self.state.timer.on_batch_complete(
                         samples=int(num_samples_in_batch.item()),
                         tokens=int(num_tokens_in_batch.item()),
@@ -1074,11 +1065,18 @@ def _train_loop(self) -> None:
                         for scheduler in self.state.schedulers:
                             scheduler.step()
 
+                    if self.train_metrics is not None:
+                        self._compute_and_log_metrics(
+                            dataloader_label='train',
+                            log_level=LogLevel.BATCH,
+                            metrics=self.train_metrics,
+                        )
+
                     self.engine.run_event(Event.BATCH_END)
 
                     if self._validate_every_n_batches > 0 and int(
                             self.state.timer.batch) % self._validate_every_n_batches == 0:
-                        self.eval(is_batch=True)
+                        self.eval(log_level=LogLevel.BATCH)
 
                     self.engine.run_event(Event.BATCH_CHECKPOINT)
 
@@ -1093,6 +1091,13 @@ def _train_loop(self) -> None:
 
             self.state.timer.on_epoch_complete()
 
+            if self.train_metrics is not None:
+                self._compute_and_log_metrics(
+                    dataloader_label='train',
+                    log_level=LogLevel.EPOCH,
+                    metrics=self.train_metrics,
+                )
+
             if not self._step_schedulers_every_batch:
                 for scheduler in self.state.schedulers:
                     scheduler.step()
@@ -1100,7 +1105,7 @@ def _train_loop(self) -> None:
             self.engine.run_event(Event.EPOCH_END)
 
             if self._validate_every_n_epochs > 0 and int(self.state.timer.epoch) % self._validate_every_n_epochs == 0:
-                self.eval(is_batch=False)
+                self.eval(log_level=LogLevel.EPOCH)
 
             self.engine.run_event(Event.EPOCH_CHECKPOINT)
 
@@ -1119,23 +1124,6 @@ def _handle_cuda_oom(self):
             self.state.grad_accum = min(2 * self.state.grad_accum, self.state.batch_num_samples)
             self.logger.data_batch({'trainer/grad_accum': self.state.grad_accum})
 
-    def _compute_metrics(self, train_metrics: Union[MetricCollection, None]):
-        """Compute training metrics.
-
-        Args:
-            train_metrics (Union[MetricCollection, None]): Existing train metrics
-        """
-        # Compute metrics on the training set
-        if self._compute_training_metrics:
-            assert train_metrics is not None
-            self.state.model.eval()
-            with torch.no_grad():
-                for eval_microbatch in self._train_data_spec.split_batch(self.state.batch, self.state.grad_accum):
-                    # TODO: Detect if self.run_event(Event.AFTER_DATALOADER) changes the training
-                    # data and if so print a warning that metrics may return unexpected results
-                    outputs, targets = self._original_model.validate(eval_microbatch)
-                    train_metrics.update(outputs, targets)
-
     def _train_batch(self, use_grad_scaling: bool):
         """Compute loss by training on a full batch of data. Adaptively change microbatch size if enabled to maximize
         GPU usage.
@@ -1317,12 +1305,12 @@ def _train_microbatch(self, use_grad_scaling: bool, current_batch_size: int, tot
         if self.deepspeed_enabled:
             self.state.deepspeed_model.step()
 
-    def eval(self, is_batch: bool):
+    def eval(self, log_level: LogLevel = LogLevel.FIT):
         """Evaluate the model on the provided evaluation data and log appropriate metrics.
 
         Args:
-            is_batch (bool): True to log metrics with ``LogLevel.BATCH``
-                and False to log metrics with ``LogLevel.EPOCH``.
+            log_level (LogLevel, optional): The log level to use for metric logging during evaluation.
+                Defaults to :attr:`~.LogLevel.FIT`.
         """
         restore_model_train = self.state.model.training
 
@@ -1334,6 +1322,7 @@ def eval(self, is_batch: bool):
             for evaluator in self.state.evaluators:
                 dataloader = evaluator.dataloader.dataloader
                 metrics = self._ensure_metrics_device_and_dtype(evaluator.metrics)
+                metrics.reset()
                 # TODO: hasattr check will be removed while fixing https://github.com/mosaicml/composer/issues/424
                 if hasattr(dataloader, "sampler") and isinstance(dataloader.sampler,
                                                                  torch.utils.data.DistributedSampler):
@@ -1360,13 +1349,16 @@ def eval(self, is_batch: bool):
                     self.engine.run_event(Event.EVAL_AFTER_FORWARD)
 
                     metrics.update(self.state.outputs, targets)
+                    self._compute_and_log_metrics(dataloader_label=evaluator.label,
+                                                  metrics=metrics,
+                                                  log_level=log_level)
 
                     self.engine.run_event(Event.EVAL_BATCH_END)
 
                 self.logger.data_epoch({"epoch": self.state.timer.epoch.value})
                 self.logger.data_batch({"trainer/global_step": self.state.timer.batch.value})
 
-                self._compute_and_log_metrics(metrics, is_train=False, is_batch=is_batch, logging_label=evaluator.label)
+                self._compute_and_log_metrics(dataloader_label=evaluator.label, metrics=metrics, log_level=log_level)
 
             self.engine.run_event(Event.EVAL_END)
 


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -20,9 +20,8 @@
 # default timout threshold is 2 seconds for determinign long and short
 DEFAULT_TIMEOUT = 2.0
 
-# Enforce use of deterministic kernels
-# see composer.utils.reproducibility.configure_deterministic_mode
-os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+# Enforce deterministic mode before any tests start.
+reproducibility.configure_deterministic_mode()
 
 # during the pytest refactor transition, this flag
 # indicates whether to include the deprecated fixtures.
@@ -145,9 +144,10 @@ def rank_zero_seed(request: pytest.FixtureRequest) -> int:
 
 
 @pytest.fixture(autouse=True)
-def seed_all(rank_zero_seed: int):
-    """Set the random seed before each test to ensure consistent test results, which and limit flakiness due to random
-    initializations."""
+def seed_all(rank_zero_seed: int, monkeypatch: pytest.MonkeyPatch):
+    """Monkeypatch reproducibility get_random_seed to always return the rank zero seed, and set the random seed before
+    each test to the rank local seed."""
+    monkeypatch.setattr(reproducibility, "get_random_seed", lambda: rank_zero_seed)
     reproducibility.seed_all(rank_zero_seed + dist.get_global_rank())
 
 
diff --git a/tests/fixtures/dummy_fixtures.py b/tests/fixtures/dummy_fixtures.py
--- a/tests/fixtures/dummy_fixtures.py
+++ b/tests/fixtures/dummy_fixtures.py
@@ -203,7 +203,6 @@ def composer_trainer_hparams(
             timeout=0.0,
         ),
         device=CPUDeviceHparams(),
-        deterministic_mode=True,
         loggers=[],
         model=dummy_model_hparams,
         val_dataset=dummy_val_dataset_hparams,
diff --git a/tests/metrics/test_current_metrics.py b/tests/metrics/test_current_metrics.py
new file mode 100644
--- /dev/null
+++ b/tests/metrics/test_current_metrics.py
@@ -0,0 +1,135 @@
+# Copyright 2021 MosaicML. All Rights Reserved.
+
+from typing import Tuple
+from unittest.mock import MagicMock
+
+import pytest
+
+from composer.core import Callback, State
+from composer.core.types import DataLoader
+from composer.loggers import Logger
+from composer.trainer import Trainer
+from tests.fixtures.models import SimpleBatchPairModel
+
+
+class TestMetricsCallback(Callback):
+
+    def __init__(self, compute_training_metrics: bool, compute_val_metrics: bool) -> None:
+        self.compute_training_metrics = compute_training_metrics
+        self.compute_val_metrics = compute_val_metrics
+        self._train_batch_end_train_accuracy = None
+        self._eval_batch_end_accuracy = None
+
+    def init(self, state: State, logger: Logger) -> None:
+        # on init, the `current_metrics` should be empty
+        del logger  # unused
+        assert state.current_metrics == {}, "no metrics should be defined on init()"
+
+    def batch_end(self, state: State, logger: Logger) -> None:
+        # The metric should be computed and updated on state every batch.
+        del logger  # unused
+        if self.compute_training_metrics:
+            # assuming that at least one sample was correctly classified
+            assert state.current_metrics["train"]["Accuracy"] != 0.0
+            self._train_batch_end_train_accuracy = state.current_metrics["train"]["Accuracy"]
+
+    def epoch_end(self, state: State, logger: Logger) -> None:
+        # The metric at epoch end should be the same as on batch end.
+        del logger  # unused
+        if self.compute_training_metrics:
+            assert state.current_metrics["train"]["Accuracy"] == self._train_batch_end_train_accuracy
+
+    def eval_batch_end(self, state: State, logger: Logger) -> None:
+        # The validation accuracy should be defined after each eval batch
+        if self.compute_val_metrics:
+            # assuming that at least one sample was correctly classified
+            assert state.current_metrics["eval"]["Accuracy"] != 0.0
+            self._eval_batch_end_accuracy = state.current_metrics["eval"]["Accuracy"]
+
+    def eval_end(self, state: State, logger: Logger) -> None:
+        if self.compute_val_metrics:
+            assert state.current_metrics["eval"]["Accuracy"] == self._eval_batch_end_accuracy
+
+
+@pytest.mark.parametrize('compute_training_metrics', [True, False])
+@pytest.mark.parametrize('validate_every_n_batches', [-1, 1])
+@pytest.mark.parametrize('validate_every_n_epochs', [-1, 1])
+def test_current_metrics(
+    dummy_train_dataloader: DataLoader,
+    dummy_val_dataloader: DataLoader,
+    dummy_num_classes: int,
+    dummy_in_shape: Tuple[int, ...],
+    compute_training_metrics: bool,
+    validate_every_n_batches: int,
+    validate_every_n_epochs: int,
+):
+    # Configure the trainer
+    num_channels = dummy_in_shape[0]
+    mock_logger_destination = MagicMock()
+    model = SimpleBatchPairModel(num_channels=num_channels, num_classes=dummy_num_classes)
+    compute_val_metrics = validate_every_n_batches == 1 or validate_every_n_epochs == 1
+    train_subset_num_batches = 2
+    eval_subset_num_batches = 2
+    num_epochs = 2
+    metrics_callback = TestMetricsCallback(
+        compute_training_metrics=compute_training_metrics,
+        compute_val_metrics=compute_val_metrics,
+    )
+
+    # Create the trainer
+    trainer = Trainer(
+        model=model,
+        train_dataloader=dummy_train_dataloader,
+        eval_dataloader=dummy_val_dataloader,
+        max_duration=num_epochs,
+        compute_training_metrics=compute_training_metrics,
+        train_subset_num_batches=train_subset_num_batches,
+        eval_subset_num_batches=eval_subset_num_batches,
+        loggers=[mock_logger_destination],
+        callbacks=[metrics_callback],
+        validate_every_n_batches=validate_every_n_batches,
+        validate_every_n_epochs=validate_every_n_epochs,
+    )
+
+    # Train the model
+    trainer.fit()
+
+    if not compute_training_metrics and not compute_val_metrics:
+        return
+
+    # Validate the metrics
+    if compute_training_metrics:
+        assert trainer.state.current_metrics["train"]["Accuracy"] != 0.0
+    else:
+        assert "train" not in trainer.state.current_metrics
+
+    if compute_val_metrics:
+        assert trainer.state.current_metrics["eval"]["Accuracy"] != 0.0
+    else:
+        assert "eval" not in trainer.state.current_metrics
+
+    # Validate that the logger was called the correct number of times for metric calls
+    num_expected_calls = 0
+    if compute_training_metrics:
+        # computed once per batch
+        # and again at epoch end
+        num_expected_calls += (train_subset_num_batches + 1) * num_epochs
+    if compute_val_metrics:
+        num_calls_per_eval = eval_subset_num_batches + 1
+        num_evals = 0
+        if validate_every_n_batches == 1:
+            num_evals += train_subset_num_batches * num_epochs
+        if validate_every_n_epochs == 1:
+            num_evals += num_epochs
+        num_expected_calls += (num_calls_per_eval) * num_evals
+    num_actual_calls = 0
+
+    # Need to filter out non-metrics-related calls
+    for call in mock_logger_destination.log_data.mock_calls:
+        data = call[1][2]
+        for k in data:
+            if k.startswith("metrics/"):
+                num_actual_calls += 1
+                break
+
+    assert num_actual_calls == num_expected_calls
diff --git a/tests/metrics/test_metrics.py b/tests/metrics/test_miou.py
similarity index 100%
rename from tests/metrics/test_metrics.py
rename to tests/metrics/test_miou.py
diff --git a/tests/test_state.py b/tests/test_state.py
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -2,6 +2,7 @@
 
 import pathlib
 import random
+from typing import Any, Dict, List
 
 import torch
 import torch.nn.functional as F
@@ -40,31 +41,49 @@ def get_dummy_state(model: ComposerModel, train_dataloader: DataLoader, val_data
     return state
 
 
+def _check_item(item1: Any, item2: Any, path: str):
+    if item1 is None:
+        assert item2 is None, f"{path} differs: {item1} != {item2}"
+        return
+    if isinstance(item1, (str, float, int, bool)):
+        assert type(item1) == type(item2)
+        assert item1 == item2, f"{path} differs: {item1} != {item2}"
+        return
+    if isinstance(item1, torch.Tensor):
+        assert isinstance(item2, torch.Tensor)
+        # Using a high tolerance, as deepspeed non-determinisim can cause
+        # metric values to be off.
+        assert item1.allclose(item2, rtol=0.1, atol=0.1), f"{path} differs"
+        return
+    if isinstance(item1, dict):
+        assert isinstance(item2, dict), f"{path} differs: {item1} != {item2}"
+        _check_dict_recursively(item1, item2, path)
+        return
+    if isinstance(item1, list):
+        assert isinstance(item2, list), f"{path} differs: {item1} != {item2}"
+        _check_list_recursively(item1, item2, path)
+        return
+    raise NotImplementedError(f"Unsupported item type: {type(item1)}")
+
+
+def _check_list_recursively(list1: List[Any], list2: List[Any], path: str):
+    assert len(list1) == len(list2), f"{path} differs: {list1} != {list2}"
+    for i, (item1, item2) in enumerate(zip(list1, list2)):
+        _check_item(item1, item2, f"{path}/{i}")
+
+
+def _check_dict_recursively(dict1: Dict[str, Any], dict2: Dict[str, Any], path: str):
+    assert len(dict1) == len(dict2), f"{path} differs: {dict1} != {dict2}"
+    for k, val1 in dict1.items():
+        val2 = dict2[k]
+        _check_item(val1, val2, f"{path}/{k}")
+
+
 def assert_state_equivalent(state1: State, state2: State):
+    assert state1.serialized_attributes == state2.serialized_attributes
+    assert state1.is_model_deepspeed == state2.is_model_deepspeed
 
-    # tested separately
-    IGNORE_FIELDS = [
-        '_optimizers',
-        '_schedulers',
-        '_algorithms',
-        '_callbacks',
-    ]
-
-    for field_name in state1.__dict__.keys():
-        if field_name in IGNORE_FIELDS or field_name.lstrip("_") not in state1.serialized_attributes:
-            continue
-
-        var1 = getattr(state1, field_name)
-        var2 = getattr(state2, field_name)
-
-        if field_name == "model":
-            assert state1.is_model_deepspeed == state2.is_model_deepspeed
-            for p, q in zip(state1.model.parameters(), state2.model.parameters()):
-                torch.testing.assert_allclose(p, q, atol=1e-2, rtol=1e-2)
-        elif isinstance(var1, torch.Tensor):
-            assert (var1 == var2).all()
-        else:
-            assert var1 == var2
+    _check_dict_recursively(state1.state_dict(), state2.state_dict(), "")
 
 
 def train_one_step(state: State, batch: Batch) -> None:
diff --git a/tests/trainer/test_checkpoint.py b/tests/trainer/test_checkpoint.py
--- a/tests/trainer/test_checkpoint.py
+++ b/tests/trainer/test_checkpoint.py
@@ -7,7 +7,7 @@
 import tempfile
 import textwrap
 import time
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 import pytest
 import torch
@@ -118,6 +118,7 @@ def assert_checkpoints_equivalent(
     checkpoint_file_a: str,
     hparams_b: TrainerHparams,
     checkpoint_file_b: str,
+    state_attrs_to_skip: List[str],
 ) -> None:
 
     with tempfile.TemporaryDirectory() as tmpdir:
@@ -149,6 +150,15 @@ def assert_checkpoints_equivalent(
     trainer_b = hparams_b.initialize_object()
     state_b = trainer_b.state
 
+    # patch the event counter callback, since they will have a different number of INIT and FIT_START
+    for callback_a, callback_b in zip(state_a.callbacks, state_b.callbacks):
+        if isinstance(callback_a, EventCounterCallback):
+            assert isinstance(callback_b, EventCounterCallback)
+            callback_b.load_state_dict(callback_a.state_dict())
+
+    for attr_name in state_attrs_to_skip:
+        setattr(state_b, attr_name, getattr(state_a, attr_name))
+
     assert_state_equivalent(state_a, state_b)
 
 
@@ -329,7 +339,6 @@ def test_checkpoint(
     if deepspeed_enabled:
         assert zero_stage is not None
         if zero_stage > 0:
-            composer_trainer_hparams.deterministic_mode = False
             if model_name is not None:
                 pytest.skip(
                     textwrap.dedent(f"""\
@@ -398,6 +407,8 @@ def test_checkpoint(
         checkpoint_file_a=first_trainer_final_checkpoint_filepath,
         hparams_b=second_trainer_hparams,
         checkpoint_file_b=second_trainer_final_checkpoint_filepath,
+        # TODO: Determine why the GPT2 Optimizer state dict differs per-checkpoint and post-checkpoint.
+        state_attrs_to_skip=["optimizers"] if model_name == "gpt2_52m" else [],
     )
 
 
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -114,6 +114,7 @@ def test_init_with_max_duration_in_batches(self, config):
 
 @world_size(1, 2)
 @device('cpu', 'gpu', 'gpu-amp', precision=True)
+@pytest.mark.timeout(15)  # higher timeout as each model is trained twice
 class TestTrainerEquivalence():
 
     reference_model: torch.nn.Module
@@ -155,7 +156,6 @@ def config(self, device: Device, precision: Precision, world_size: int, rank_zer
             'seed': rank_zero_seed,
             'device': device,
             'precision': precision,
-            'deterministic_mode': True,  # testing equivalence
             'loggers': [],  # no progress bar
         }