Move to optimizer based EMA implementation

[SeanNaren]

What does this PR do ?

Move to using the EMA Optimizer. This improves performance, supports offloading to CPU!

Changelog

• Move to EMA Optimizer implementation

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[markelsanz14]

Leaving a few comments for next week:

• We should make the validation configurable, so that the user can decide whether to run val with regular or ema weights.
• We should make the CPU offloading configurable too, for both the ema weights update, and for keeping the ema weights in CPU mem rather than GPU mem.
• We should update the config yaml files of all the models that will support EMA to add the ema parameters to the exp_manager config. We can default it to False, but it should be present in the config.
• Maybe we should set the default decay value to 0.9995 in the config files and in the code. I think that should be a decent default value.
• I found this error when the model tries to save a checkpoint:

  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 713, in _update_best_and_save
    self._save_checkpoint(trainer, filepath)
  File "/opt/bignlp/NeMo/nemo/utils/exp_manager.py", line 904, in _save_checkpoint
    ema_callback.replace_model_weights(trainer.lightning_module)
AttributeError: 'EMA' object has no attribute 'replace_model_weights'

[lgtm-com]

This pull request introduces 1 alert when merging 979c51a into acb5073 - view on LGTM.com

new alerts:

• 1 for Unused import

[lgtm-com]

This pull request introduces 1 alert when merging f8a361e into acb5073 - view on LGTM.com

new alerts:

• 1 for Unused import

[SeanNaren]

Have quite a few local changes and improvements, just figuring out how to ensure that the "non" EMA checkpoint can be re-loaded without the optimizer wrapper.

[SeanNaren]

With how the checkpoint needs to be overwritten, we require PyTorch Lightning 1.8 for the new updated behaviour where the callback has access to the entire state of the checkpoint. I've started the PR here to fix issues: #5353

[github-actions]

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

[lgtm-com]

This pull request introduces 1 alert when merging a4b3e53 into f0d9c51 - view on LGTM.com

new alerts:

• 1 for Unused import

---

Heads-up: LGTM.com's PR analysis will be disabled on the 5th of December, and LGTM.com will be shut down ⏻ completely on the 16th of December 2022. It looks like GitHub code scanning with CodeQL is already set up for this repo, so no further action is needed 🚀. For more information, please check out our post on the GitHub blog.

[SeanNaren]

@titu1994 ready for review :) fixed the tests and should now be at feature parity with what we had existing!

[titu1994]

PR looks awesome. Minor comments. Also for training test, could you add a batch norm to the model so that it tests what happens if ema is done on the int component (num steps) for batchnorm. I wonder if there's any literature for that case, maybe we should take max instead of ema of int

[SeanNaren]

Just included batch norm into the model, which causes the tests to fail. Will debug!

cc @titu1994 @szmigacz

[titu1994]

Looks fantastic !

[malbergo]

Do you have an example of how to use this callback within an LightningModule? Thanks !

[amorehead]

@malbergo, as an aside, I was successfully able to make this EMA callback work for my needs in the context of 3D diffusion generative models. If you'd like to see an example of how I make this work inside my custom LightningModules, feel free to check out the following links to my bio-diffusion repository:

• EMA callback: https://github.com/BioinfoMachineLearning/bio-diffusion/blob/e4bad15139815e562a27fb94dab0c31907522bc5/src/utils/__init__.py#L71
• Custom Lightning checkpoint callback that wraps the EMA callback to save EMA versions of each model checkpoint: https://github.com/BioinfoMachineLearning/bio-diffusion/blob/e4bad15139815e562a27fb94dab0c31907522bc5/src/utils/__init__.py#L224
• How I set the arguments for this EMA callback using Hydra config files: https://github.com/BioinfoMachineLearning/bio-diffusion/blob/e4bad15139815e562a27fb94dab0c31907522bc5/configs/callbacks/ema.yaml#L5
• How I include the EMA callback as one of many Lightning callbacks for my Trainer to make use of: https://github.com/BioinfoMachineLearning/bio-diffusion/blob/e4bad15139815e562a27fb94dab0c31907522bc5/configs/callbacks/default.yaml#L4

I hope this helps!

[SeanNaren]

This is fantastic @amorehead! We do not have documentation explicitly for using this callback with pure LightningModules, but @amorehead code seems to be a good place to start!

Here are the docs for using it with NeMo however: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/common/callbacks.html#exponential-moving-average-ema

[flashszn]

@malbergo, as an aside, I was successfully able to make this EMA callback work for my needs in the context of 3D diffusion generative models. If you'd like to see an example of how I make this work inside my custom LightningModules, feel free to check out the following links to my bio-diffusion repository:

• EMA callback: https://github.com/BioinfoMachineLearning/bio-diffusion/blob/e4bad15139815e562a27fb94dab0c31907522bc5/src/utils/__init__.py#L71
• Custom Lightning checkpoint callback that wraps the EMA callback to save EMA versions of each model checkpoint: https://github.com/BioinfoMachineLearning/bio-diffusion/blob/e4bad15139815e562a27fb94dab0c31907522bc5/src/utils/__init__.py#L224
• How I set the arguments for this EMA callback using Hydra config files: https://github.com/BioinfoMachineLearning/bio-diffusion/blob/e4bad15139815e562a27fb94dab0c31907522bc5/configs/callbacks/ema.yaml#L5
• How I include the EMA callback as one of many Lightning callbacks for my Trainer to make use of: https://github.com/BioinfoMachineLearning/bio-diffusion/blob/e4bad15139815e562a27fb94dab0c31907522bc5/configs/callbacks/default.yaml#L4

I hope this helps!

Thanks a lot. When I use this callback with top-k ckpt, it won't automatically del the ckpt with -EMA and will save a lot of ckpts with -EMA.

So I modify the EMAModelCheckpoint, overriding the _update_best_and_save:

# only change the last line
def _update_best_and_save(
        self, current: Tensor, trainer: "pl.Trainer", monitor_candidates: Dict[str, Tensor]
) -> None:
    k = len(self.best_k_models) + 1 if self.save_top_k == -1 else self.save_top_k

    del_filepath = None
    if len(self.best_k_models) == k and k > 0:
        del_filepath = self.kth_best_model_path
        self.best_k_models.pop(del_filepath)

    # do not save nan, replace with +/- inf
    if isinstance(current, Tensor) and torch.isnan(current):
        current = torch.tensor(float("inf" if self.mode == "min" else "-inf"), device=current.device)

    filepath = self._get_metric_interpolated_filepath_name(monitor_candidates, trainer, del_filepath)

    # save the current score
    self.current_score = current
    self.best_k_models[filepath] = current

    if len(self.best_k_models) == k:
        # monitor dict has reached k elements
        _op = max if self.mode == "min" else min
        self.kth_best_model_path = _op(self.best_k_models, key=self.best_k_models.get)  # type: ignore[arg-type]
        self.kth_value = self.best_k_models[self.kth_best_model_path]

    _op = min if self.mode == "min" else max
    self.best_model_path = _op(self.best_k_models, key=self.best_k_models.get)  # type: ignore[arg-type]
    self.best_model_score = self.best_k_models[self.best_model_path]

    if self.verbose:
        epoch = monitor_candidates["epoch"]
        step = monitor_candidates["step"]
        rank_zero_info(
            f"Epoch {epoch:d}, global step {step:d}: {self.monitor!r} reached {current:0.5f}"
            f" (best {self.best_model_score:0.5f}), saving model to {filepath!r} as top {k}"
        )
    self._save_checkpoint(trainer, filepath)

    if del_filepath is not None and filepath != del_filepath:
        self._remove_checkpoint(trainer, del_filepath)
        self._remove_checkpoint(trainer, del_filepath.replace(self.FILE_EXTENSION, f"-EMA{self.FILE_EXTENSION}"))