Repository: NVIDIA/NeMo

Issue:
EMA Doesn't delete previous EMA ckpts when k > 0 for checkpointing
**Describe the bug**

EMA saves a separate EMA based ckpt when normally saving. The problem is when saving a set amount of checkpoints (i.e k=5) the older EMA ckpts are not being deleted. This means there is an excess of EMA based ckpts.

What is expected is that older EMA ckpts are deleted.

**Proposed Solution**

Overriding the code is tricky, as the call to remove ckpts is found within functions, like such:

https://github.com/Lightning-AI/lightning/blob/master/src/pytorch_lightning/callbacks/model_checkpoint.py#L652
https://github.com/Lightning-AI/lightning/blob/master/src/pytorch_lightning/callbacks/model_checkpoint.py#L671
https://github.com/Lightning-AI/lightning/blob/master/src/pytorch_lightning/callbacks/model_checkpoint.py#L713

My suggestion would be to move this code into a separate function that can be overridden in the [NeMoModelCheckpoint](https://github.com/NVIDIA/NeMo/blob/4d3e99f986fd9354b5528a3e9c82fe01ec2df916/nemo/utils/exp_manager.py#L745) class by introducing a `delete_checkpoint` function. In this function we will also be able to delete the EMA weights.



Discussion:
cc @carmocca
Ideally, we would find a better solution. However, since that would require a larger refactor - to `ModelCheckpoint` and/or NeMo - and the `NemoModelCheckpoint` callback already relies heavily on overriding protected methods, I'm fine with adding one more.
That's sounds like a good interim measure.
PR created here: https://github.com/Lightning-AI/lightning/pull/16067 

Hopefully this can make it into the next bugfix release if possible.
@Borda, do you agree with cherry picking this into a 1.8.x release?
> @Borda, do you agree with cherry picking this into a 1.8.x release?

yes, good with me

Resolving patch:
diff --git a/nemo/collections/common/callbacks/ema.py b/nemo/collections/common/callbacks/ema.py
--- a/nemo/collections/common/callbacks/ema.py
+++ b/nemo/collections/common/callbacks/ema.py
@@ -13,13 +13,13 @@
 # limitations under the License.
 import contextlib
 import copy
-import logging
 import os
 import threading
 from typing import Any, Dict, Iterable
 
 import pytorch_lightning as pl
 import torch
+from lightning_utilities.core.rank_zero import rank_zero_info
 from pytorch_lightning import Callback
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
@@ -117,25 +117,26 @@ def on_load_checkpoint(
     ) -> None:
         checkpoint_callback = trainer.checkpoint_callback
 
-        if trainer.ckpt_path and checkpoint_callback is not None and 'NeMo' in type(checkpoint_callback).__name__:
+        # use the connector as NeMo calls the connector directly in the exp_manager when restoring.
+        connector = trainer._checkpoint_connector
+        ckpt_path = connector.resume_checkpoint_path
+
+        if ckpt_path and checkpoint_callback is not None and 'NeMo' in type(checkpoint_callback).__name__:
             ext = checkpoint_callback.FILE_EXTENSION
-            if trainer.ckpt_path.endswith(f'-EMA{ext}'):
-                logging.info(
+            if ckpt_path.endswith(f'-EMA{ext}'):
+                rank_zero_info(
                     "loading EMA based weights. "
                     "The callback will treat the loaded EMA weights as the main weights"
                     " and create a new EMA copy when training."
                 )
                 return
-            ema_path = trainer.ckpt_path.replace(ext, f'-EMA{ext}')
+            ema_path = ckpt_path.replace(ext, f'-EMA{ext}')
             if os.path.exists(ema_path):
                 ema_state_dict = torch.load(ema_path, map_location=torch.device('cpu'))
 
-                # this is wrong, basically when we save the EMA weights, optimizer_states actually contains the model parameters
-                # as we swapped the model parameters with the state dict parameters.
-                # we could enforce that if you trained with EMA and want to continue training
                 checkpoint['optimizer_states'] = ema_state_dict['optimizer_states']
                 del ema_state_dict
-                logging.info("EMA state has been restored.")
+                rank_zero_info("EMA state has been restored.")
             else:
                 raise MisconfigurationException(
                     "Unable to find the associated EMA weights when re-loading, "
diff --git a/nemo/utils/exp_manager.py b/nemo/utils/exp_manager.py
--- a/nemo/utils/exp_manager.py
+++ b/nemo/utils/exp_manager.py
@@ -990,6 +990,14 @@ def _save_checkpoint(self, trainer: 'pytorch_lightning.Trainer', filepath: str)
         else:
             super()._save_checkpoint(trainer, filepath)
 
+    def _remove_checkpoint(self, trainer: "pytorch_lightning.Trainer", filepath: str) -> None:
+        super()._remove_checkpoint(trainer, filepath)
+        ema_callback = self._ema_callback(trainer)
+        if ema_callback is not None:
+            # remove EMA copy of the state dict as well.
+            filepath = self._ema_format_filepath(filepath)
+            super()._remove_checkpoint(trainer, filepath)
+
     def _ema_format_filepath(self, filepath: str) -> str:
         return filepath.replace(self.FILE_EXTENSION, f'-EMA{self.FILE_EXTENSION}')
 


Test patch:
diff --git a/tests/collections/common/test_ema.py b/tests/collections/common/test_ema.py
--- a/tests/collections/common/test_ema.py
+++ b/tests/collections/common/test_ema.py
@@ -27,7 +27,10 @@
 from nemo.collections.common.callbacks.ema import EMAOptimizer
 from nemo.core import ModelPT
 from nemo.utils.exp_manager import exp_manager
-from tests.collections.nlp.test_gpt_model import DEVICE_CAPABILITY
+
+DEVICE_CAPABILITY = None
+if torch.cuda.is_available():
+    DEVICE_CAPABILITY = torch.cuda.get_device_capability()
 
 
 def extract_ema_weights(pl_module, trainer):
@@ -69,13 +72,22 @@ def val_dataloader(self):
         dataset = RandomDataset(32, 16)
         return torch.utils.data.DataLoader(dataset, batch_size=2)
 
+    def test_dataloader(self):
+        dataset = RandomDataset(32, 16)
+        dl = torch.utils.data.DataLoader(dataset, batch_size=2)
+        self._test_names = ['test_{}_'.format(idx) for idx in range(len(dl))]
+        return dl
+
     def forward(self, batch):
         return self.l1(self.bn(batch)).sum()
 
+    def training_step(self, batch, batch_idx):
+        return self(batch)
+
     def validation_step(self, batch, batch_idx):
         return self(batch)
 
-    def training_step(self, batch, batch_idx):
+    def test_step(self, batch, batch_idx):
         return self(batch)
 
     def configure_optimizers(self):
@@ -90,6 +102,9 @@ def setup_training_data(self, train_data_config: Union[DictConfig, Dict]):
     def setup_validation_data(self, val_data_config: Union[DictConfig, Dict]):
         pass
 
+    def setup_test_data(self, val_data_config: Union[DictConfig, Dict]):
+        pass
+
     def validation_epoch_end(self, loss):
         self.log("val_loss", torch.stack(loss).mean())
 
@@ -248,6 +263,27 @@ def test_exp_manager_ema_weights(self, tmpdir):
         for saved_weight, ema_weight in zip(duplicate_model.state_dict().values(), ema_weights):
             assert torch.allclose(saved_weight.cpu(), ema_weight.cpu())
 
+    @pytest.mark.unit
+    def test_exp_manager_ema_weights_topk(self, tmpdir):
+        """Test to ensure that EMA correctly ensures we only keep topk checkpoints."""
+        tmp_path = tmpdir / "exp_manager_test"
+        model = ExampleModel()
+        save_top_k = 3
+
+        trainer = Trainer(max_epochs=10, enable_checkpointing=False, logger=False, devices=1)
+        exp_manager(
+            trainer,
+            {
+                "ema": {"enable": True},
+                "explicit_log_dir": str(tmp_path),
+                "checkpoint_callback_params": {"save_top_k": save_top_k},
+            },
+        )
+        trainer.fit(model)
+
+        # we save 3 checkpoints for the model, 3 accompanied EMA weights, the last checkpoint and nemo model.
+        assert len(os.listdir(tmp_path / "checkpoints/")) == (save_top_k + 1) * 2 + 1
+
 
 class TestEMATrain:
     @pytest.mark.unit
@@ -320,6 +356,23 @@ def run_training_test(self, accumulate_grad_batches, validate_original_weights,
         trainer.callbacks.insert(0, EMAValidationAssertCallback())
         trainer.fit(model=model, val_dataloaders=model.train_dataloader())
 
+    @pytest.mark.unit
+    def test_ema_run_with_save_best_model(self, tmpdir):
+        """Test to ensure that we save the model correctly when save best model is set to True."""
+        tmp_path = tmpdir / "exp_manager_test"
+        model = ExampleModel()
+
+        trainer = Trainer(max_epochs=1, enable_checkpointing=False, logger=False, devices=1, limit_train_batches=1)
+        exp_manager(
+            trainer,
+            {
+                "ema": {"enable": True},
+                "explicit_log_dir": str(tmp_path),
+                "checkpoint_callback_params": {"save_best_model": True},
+            },
+        )
+        trainer.fit(model)
+
 
 class EMAAssertCallback(Callback):
     def on_train_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: