/audiocraft/solvers/base.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from abc import ABC, abstractmethod
from contextlib import contextmanager
from pathlib import Path
import typing as tp

import flashy
import omegaconf
import torch
from torch import nn

from .. import optim
from ..optim import fsdp
from ..utils import checkpoint
from ..utils.autocast import TorchAutocast
from ..utils.best_state import BestStateDictManager
from ..utils.deadlock import DeadlockDetect
from ..utils.profiler import Profiler
from ..utils.utils import copy_state, dict_from_config, model_hash, with_rank_rng


class StandardSolver(ABC, flashy.BaseSolver):
    """Standard solver for AudioCraft.

    The standard solver implements a base training loop with the following stages:
    train, valid, evaluate and generate that are expected to be all defined for
    solvers in AudioCraft. It also provides a nice default management of Dora history replay,
    checkpoint management across epoch, and logging configuration.

    AudioCraft solvers must inherit from the StandardSolver and define the methods
    associated to each stage as well as the show, build_model and build_dataloaders methods.
    """
    def __init__(self, cfg: omegaconf.DictConfig):
        super().__init__()
        self.logger.info(f"Instantiating solver {self.__class__.__name__} for XP {self.xp.sig}")
        self.logger.info(f"All XP logs are stored in {self.xp.folder}")
        self.cfg = cfg
        self.device = cfg.device
        self.model: nn.Module
        self._continue_best_source_keys = ['best_state', 'fsdp_best_state']
        self._fsdp_modules: tp.List[fsdp.FSDP] = []
        self._ema_sources: nn.ModuleDict = nn.ModuleDict()
        self.ema: tp.Optional[optim.ModuleDictEMA] = None
        self.dataloaders: tp.Dict[str, torch.utils.data.DataLoader] = dict()
        self._log_updates = self.cfg.logging.get('log_updates', 10)
        if self.cfg.logging.log_tensorboard:
            self.init_tensorboard(**self.cfg.get('tensorboard'))
        if self.cfg.logging.log_wandb and self:
            self.init_wandb(**self.cfg.get('wandb'))
        # keep a copy of the best performing state for stateful objects
        # used for evaluation and generation stages
        dtype_best: tp.Optional[torch.dtype] = None
        if self.cfg.fsdp.use:
            dtype_best = getattr(torch, self.cfg.fsdp.param_dtype)  # type: ignore
            assert isinstance(dtype_best, torch.dtype)
        elif self.cfg.autocast:
            dtype_best = getattr(torch, self.cfg.autocast_dtype)  # type: ignore
            assert isinstance(dtype_best, torch.dtype)
        self.best_state: BestStateDictManager = BestStateDictManager(dtype=dtype_best)
        # Hacky support for keeping a copy of the full best state in rank0.
        self.fsdp_best_state: tp.Dict[str, tp.Any] = {}
        self.register_stateful('best_state', 'fsdp_best_state')  # register best_state object to keep it in state_dict
        self._new_best_state: bool = False  # should save a new checkpoint
        # instantiate datasets and appropriate number of updates per epoch
        self.build_dataloaders()
        if self.cfg.execute_only is None:
            assert 'train' in self.dataloaders, "The train dataset split must be provided."
            assert 'valid' in self.dataloaders, "The valid dataset split must be provided."
        self.train_updates_per_epoch = len(self.dataloaders['train']) if 'train' in self.dataloaders else 0
        if self.cfg.optim.updates_per_epoch:
            self.train_updates_per_epoch = self.cfg.optim.updates_per_epoch
        self.total_updates = self.train_updates_per_epoch * self.cfg.optim.epochs
        # instantiate model & exponential moving average on the model
        self.build_model()
        self.logger.info("Model hash: %s", model_hash(self.model))
        assert 'model' in self.stateful.sources, \
            "Please register the model to stateful with self.register_stateful('model') in build_model."
        self.profiler = Profiler(self.model, **self.cfg.profiler)
        self.initialize_ema()
        self.register_stateful('ema')
        assert self.ema is None or 'ema' in self.stateful.sources, \
            "Please register the ema to stateful with self.register_stateful('ema') in build_model."
        self.deadlock_detect = DeadlockDetect(**self.cfg.deadlock)
        # basic statistics on the trained model
        model_size = sum(p.numel() for p in self.model.parameters() if p.requires_grad) / 1e6
        # one copy of grad, one copy of momentum, one copy of denominator and model weights.
        # and 4 bytes for each float!
        mem_usage = model_size * 4 * 4 / 1000
        self.logger.info("Model size: %.2f M params", model_size)
        self.logger.info("Base memory usage, with model, grad and optim: %.2f GB", mem_usage)

    @property
    def autocast(self):
        """Convenient autocast (or not) using the solver configuration."""
        return TorchAutocast(enabled=self.cfg.autocast, device_type=self.device, dtype=self.autocast_dtype)

    def _get_state_source(self, name) -> flashy.state.StateDictSource:
        # Internal utility to get a state source from the solver
        return self.stateful.sources[name]

    @property
    def best_metric_name(self) -> tp.Optional[str]:
        """Metric name used to identify the best state. This metric should be stored in the metrics
        used on the stage for best state identification (most likely, `valid`). If None, then
        no best state is saved.
        """
        return None

    def register_best_state(self, *args: str):
        """Register state sources in `BestStateDictManager` to keep their best states along with their
        latest states. The best state will be used at evaluation stages instead of the latest states.

        Shortcut around `BestStateDictManager.register` method. You can pass any number of
        attribute, included nested attributes and those will be included into the checkpoints
        and automatically restored when `BaseSolver.restore` is called.
        """
        for name in args:
            state_source = self._get_state_source(name)
            assert name in self.stateful.sources, "Registered states in best should be registered in stateful first!"
            self.best_state.register(name, state_source)

    def register_ema(self, *args: str):
        """Register state sources for exponential moving average.

        The registered sources are used to instantiate a ModuleDictEMA instance.
        The ModuleDictEMA keeps a `nn.ModuleDict` module that is updated when self.ema.step() is called
        and swapped with the original state sources with self.swap_ema_state() method.

        Usage:
            self.register_ema('model')
        """
        assert self.ema is None, "Cannot register state source to already instantiated EMA."
        for name in args:
            self._ema_sources[name] = getattr(self, name)

    def wrap_with_fsdp(self, model: torch.nn.Module, *args, **kwargs):
        model = fsdp.wrap_with_fsdp(self.cfg.fsdp, model, *args, **kwargs)
        if isinstance(model, fsdp.FSDP):
            self._fsdp_modules.append(model)
        return model

    def update_best_state_from_stage(self, stage_name: str = 'valid'):
        """Update latest best state based on pending metrics of a given stage. This method relies
        on the `BestStateDictManager.update` method to update the best state_dict with latest weights
        if the registered states happen to match to the best performing setup.
        """
        if self.best_metric_name is None:
            # when no best metric is defined, the last state is always the best
            self._new_best_state = True
            self.logger.info("Updating best state with current state.")
        else:
            assert stage_name in self._pending_metrics, f"Metrics for stage {stage_name} not found."
            assert self.best_metric_name in self._pending_metrics[stage_name], \
                f"Best metric not found in {stage_name} metrics. Cannot register best state"
            current_score = self._pending_metrics[stage_name][self.best_metric_name]
            all_best_metric_scores = [
                past_metrics[stage_name][self.best_metric_name]
                for past_metrics in self.history
            ]
            all_best_metric_scores.append(current_score)
            best_score = min(all_best_metric_scores)
            self._new_best_state = current_score == best_score
            if self._new_best_state:
                old_best = min(all_best_metric_scores[:-1] + [float('inf')])
                self.logger.info(
                    f"New best state with {self.best_metric_name}={current_score:.3f} (was {old_best:.3f})")

        if self._new_best_state:
            if self.cfg.fsdp.use:
                # this will give an empty state dict on all ranks but the rank 0
                # which will have a copy in memory of the full model.
                with fsdp.switch_to_full_state_dict(self._fsdp_modules):
                    for name in self.best_state.states.keys():
                        state_source = self._get_state_source(name)
                        self.best_state.update(name, state_source)
                    # we save to a different dict.
                    self.fsdp_best_state.update(self.best_state.state_dict())
                # We cannot efficiently load fsdp_best_state when using FSDP,
                # so we have do do a second pass, with the local shards.
            for name in self.best_state.states.keys():
                state_source = self._get_state_source(name)
                self.best_state.update(name, state_source)

    def _load_new_state_dict(self, state_dict: dict) -> dict:
        old_states = {}
        for name, new_state in state_dict.items():
            state_source = self._get_state_source(name)
            old_states[name] = copy_state(state_source.state_dict())
            state_source.load_state_dict(new_state)
        return old_states

    @contextmanager
    def swap_best_state(self):
        self.logger.debug(f"Swapping to best state for: {', '.join(self.best_state.state_dict().keys())}")
        old_states = self._load_new_state_dict(self.best_state.state_dict())
        try:
            yield
        finally:
            self.logger.debug("Swapping back from best to original state")
            for name, old_state in old_states.items():
                state_source = self._get_state_source(name)
                state_source.load_state_dict(old_state)

    @contextmanager
    def swap_ema_state(self):
        if self.ema is None:
            yield
        else:
            ema_state_dict = self.ema.state_dict()['state']
            self.logger.debug(f"Swapping to EMA state for: {', '.join(ema_state_dict.keys())}")
            old_states = self._load_new_state_dict(ema_state_dict)
            try:
                yield
            finally:
                self.logger.debug("Swapping back from EMA state to original state")
                for name, old_state in old_states.items():
                    state_source = self._get_state_source(name)
                    state_source.load_state_dict(old_state)

    @property
    def is_training(self):
        return self.current_stage == 'train'

    def log_model_summary(self, model: nn.Module):
        """Log model summary, architecture and size of the model."""
        self.logger.info(model)
        mb = sum(p.numel() for p in model.parameters()) * 4 / 2 ** 20
        self.logger.info("Size: %.1f MB", mb)

    @abstractmethod
    def build_model(self):
        """Method to implement to initialize model."""
        ...

    def initialize_ema(self):
        """Initialize exponential moving average with the registered sources.
        EMA object is created if the optim.ema.model.decay value is non-null.
        """
        from .builders import get_ema
        self.ema = get_ema(self._ema_sources, self.cfg.optim.ema)
        if self.ema is None:
            self.logger.info('No EMA on the model.')
        else:
            assert self.cfg.optim.ema.updates > 0
            self.logger.info(
                f'Initializing EMA on the model with decay = {self.ema.decay}'
                f' every {self.cfg.optim.ema.updates} updates'
            )

    @abstractmethod
    def build_dataloaders(self):
        """Method to implement to initialize dataloaders."""
        ...

    @abstractmethod
    def show(self):
        """Method to log any information without running the job."""
        ...

    @property
    def log_updates(self):
        # convenient access to log updates
        return self._log_updates

    def checkpoint_path(self, **kwargs):
        kwargs.setdefault('use_fsdp', self.cfg.fsdp.use)
        return self.folder / checkpoint.checkpoint_name(**kwargs)

    def epoch_checkpoint_path(self, epoch: int, **kwargs):
        kwargs.setdefault('use_fsdp', self.cfg.fsdp.use)
        return self.folder / checkpoint.checkpoint_name(str(epoch), **kwargs)

    def checkpoint_path_with_name(self, name: str, **kwargs):
        kwargs.setdefault('use_fsdp', self.cfg.fsdp.use)
        return self.folder / checkpoint.checkpoint_name(name=name, **kwargs)

    def save_checkpoints(self):
        """Save checkpoint, optionally keeping a copy for a given epoch."""
        is_sharded = self.cfg.fsdp.use
        if not flashy.distrib.is_rank_zero() and not is_sharded:
            return
        self.logger.info("Model hash: %s", model_hash(self.model))
        state = self.state_dict()
        epoch = self.epoch - 1  # pushing metrics will increase the epoch in Flashy, so we do -1 here

        # save minimal state_dict as new checkpoint every X epoch
        if self.cfg.checkpoint.save_every:
            if epoch % self.cfg.checkpoint.save_every == 0:
                minimal_state = state
                if self.cfg.checkpoint.keep_every_states is not None and len(self.cfg.checkpoint.keep_every_states) > 0:
                    minimal_state = {
                        name: source for name, source in state.items()
                        if name in self.cfg.checkpoint.keep_every_states
                    }
                epoch_checkpoint_path = self.epoch_checkpoint_path(epoch)
                checkpoint.save_checkpoint(minimal_state, epoch_checkpoint_path, is_sharded)

        # save checkpoint as latest checkpoint
        if self.cfg.checkpoint.save_last:
            last_checkpoint_path = self.checkpoint_path()
            checkpoint.save_checkpoint(state, last_checkpoint_path, is_sharded)

        # flush any stale checkpoint to reduce disk footprint
        checkpoint.flush_stale_checkpoints(self.checkpoint_path())

    def load_from_pretrained(self, name: str) -> dict:
        raise NotImplementedError("Solver does not provide a way to load pretrained models.")

    def load_checkpoints(self, load_best: bool = False, ignore_state_keys: tp.List[str] = []) -> tp.Optional[dict]:
        """Load last checkpoint or the one specified in continue_from.

        Args:
            load_best (bool): Whether to load from best state dict or not.
                Best state dict is always used when not loading the current xp.
            ignore_state_keys (list of str): List of sources to ignore when loading the state, e.g. `optimizer`.
        Returns:
            state (dict, optional): The loaded state dictionary.
        """
        # load checkpoints from xp folder or cfg.continue_from
        is_sharded = self.cfg.fsdp.use
        load_from_path: tp.Optional[Path] = None
        checkpoint_source: tp.Optional[checkpoint.CheckpointSource] = None

        if load_best:
            self.logger.info("Trying to load state_dict from best state.")

        state: tp.Optional[dict] = None
        rank0_checkpoint_path = self.checkpoint_path(use_fsdp=False)
        current_checkpoint_path = self.checkpoint_path()
        _pretrained_prefix = '//pretrained/'
        continue_pretrained = (self.cfg.continue_from or '').startswith(_pretrained_prefix)
        if rank0_checkpoint_path.exists():
            self.logger.info(f"Loading existing checkpoint: {current_checkpoint_path}")
            load_from_path = current_checkpoint_path
            checkpoint.check_sharded_checkpoint(current_checkpoint_path, rank0_checkpoint_path)
            checkpoint_source = checkpoint.CheckpointSource.CURRENT_XP
        elif self.cfg.continue_from and not continue_pretrained:
            self.logger.info(f"Continuing from provided checkpoint: {self.cfg.continue_from}")
            # we're always continuing from consolidated checkpoints: self.cfg.use_fsdp and not continue_best
            load_from_path = checkpoint.resolve_checkpoint_path(self.cfg.continue_from, use_fsdp=False)
            if load_from_path is None:
                self.logger.error('Could not resolve the continue_from checkpoint %s', self.cfg.continue_from)
                raise RuntimeError(f'Could not resolve continue_from checkpoint {self.cfg.continue_from}')
            checkpoint_source = checkpoint.CheckpointSource.OTHER

        if load_from_path is not None:
            state = checkpoint.load_checkpoint(load_from_path, is_sharded)
        elif continue_pretrained:
            self.logger.info("Loading a pretrained model. Ignoring 'load_best' and 'ignore_state_keys' params.")
            state = self.load_from_pretrained(self.cfg.continue_from[len(_pretrained_prefix):])
            checkpoint_source = checkpoint.CheckpointSource.PRETRAINED
            load_best = True

        # checkpoints are not from the current xp, we only retrieve the best state
        if checkpoint_source is not None and checkpoint_source != checkpoint.CheckpointSource.CURRENT_XP:
            assert state is not None
            self.logger.info("Checkpoint source is not the current xp: Load state_dict from best state.")
            load_best = True
            state = {key: state[key] for key in self._continue_best_source_keys if key in state}
            # loaded checkpoints are FSDP checkpoints: we're reading the best state
            # from FSDP and we drop the regular best_state
            if 'fsdp_best_state' in state and state['fsdp_best_state']:
                state.pop('best_state', None)
                self.logger.info("... Loaded checkpoint has FSDP best state")
            # FSDP is enabled in the solver, if the loaded checkpoints do not have FSDP support
            # then we're initializing FSDP best state with the regular best state
            elif self.cfg.fsdp.use:
                if 'fsdp_best_state' not in state or not state['fsdp_best_state']:
                    # we swap non-FSDP checkpoints best_state to FSDP-compatible best state
                    state['fsdp_best_state'] = state.pop('best_state')
                    self.logger.info("... Loaded checkpoint does not have FSDP best state. Use regular best state")

        if state is not None:
            if load_best:
                self.logger.info("Ignoring keys when loading best %r", ignore_state_keys)
                for key in set(ignore_state_keys):
                    if key in state:
                        state.pop(key)
                has_best_state = 'best_state' in state or 'fsdp_best_state' in state
                assert has_best_state, ("Trying to load best state but neither 'best_state'",
                                        " or 'fsdp_best_state' found in checkpoints.")
            self.load_state_dict(state)

        # for FSDP, let's make extra sure nothing bad happened with out of sync
        # checkpoints across workers.
        epoch = float(self.epoch)
        avg_epoch = flashy.distrib.average_metrics({'epoch': epoch})['epoch']
        if avg_epoch != epoch:
            raise RuntimeError(
                f"Inconsistent loading of checkpoints happened, our epoch is {epoch} "
                f"but average of epochs is {avg_epoch}, at least one gpu must have a "
                "different epoch number.")

        # on load_best, properly reinitialize state_dict, best states and ema
        # otherwise we load from the current xp and don't alter anything
        if load_best:
            self.logger.info("Loading state_dict from best state.")
            if not self.cfg.fsdp.use and self.fsdp_best_state:
                # loading from an FSDP checkpoint but with FSDP deactivated
                self.logger.info("... Loading from FSDP best state dict.")
                self.best_state.load_state_dict(self.fsdp_best_state)

            # if load_best, we permanently override the regular state_dict with the best state
            if self.cfg.fsdp.use:
                self.logger.info("FSDP is used, loading from FSDP best state.")
                with fsdp.switch_to_full_state_dict(self._fsdp_modules):
                    # this might be really fragile but okay for now.
                    self.load_state_dict(self.fsdp_best_state)
            else:
                # we permanently swap the stateful objects to their best state
                self._load_new_state_dict(self.best_state.state_dict())

            # the EMA modules should also be instantiated with best state.
            # the easiest way to do so is to reinitialize a new EMA with best state loaded.
            if self.ema is not None:
                self.logger.info("Re-initializing EMA from best state")
                self.initialize_ema()

            if self.cfg.fsdp.use:
                self.logger.info("Re-initializing best state after using FSDP best state.")
                for name in self.best_state.states.keys():
                    state_source = self._get_state_source(name)
                    self.best_state.update(name, state_source)

        return state

    def restore(self, load_best: bool = False, replay_metrics: bool = False,
                ignore_state_keys: tp.List[str] = []) -> bool:
        """Restore the status of a solver for a given xp.

        Args:
            load_best (bool): if `True`, load the best state from the checkpoint.
            replay_metrics (bool): if `True`, logs all the metrics from past epochs.
            ignore_state_keys (list of str): list of sources to ignore when loading the state, e.g. `optimizer`.
        """
        self.logger.info("Restoring weights and history.")
        restored_checkpoints = self.load_checkpoints(load_best, ignore_state_keys)

        self.logger.info("Model hash: %s", model_hash(self.model))

        if replay_metrics and len(self.history) > 0:
            self.logger.info("Replaying past metrics...")
            for epoch, stages in enumerate(self.history):
                for stage_name, metrics in stages.items():
                    # We manually log the metrics summary to the result logger
                    # as we don't want to add them to the pending metrics
                    self.result_logger._log_summary(stage_name, metrics, step=epoch + 1, step_name='epoch',
                                                    formatter=self.get_formatter(stage_name))
        return restored_checkpoints is not None

    def commit(self, save_checkpoints: bool = True):
        """Commit metrics to dora and save checkpoints at the end of an epoch."""
        # we override commit to introduce more complex checkpoint saving behaviors
        self.history.append(self._pending_metrics)  # This will increase self.epoch
        if save_checkpoints:
            self.save_checkpoints()
        self._start_epoch()
        if flashy.distrib.is_rank_zero():
            self.xp.link.update_history(self.history)

    def run_epoch(self):
        """Run a single epoch with all stages.

        Metrics for a given stage are stored in _pending_metrics and committed by the solver afterwards.
        Children solvers can extend this method with custom behavior, e.g.:

            def run_epoch(self):
                ... # custom code
                super().run_epoch()
                ... # custom code
        """
        self.run_stage('train', self.train)
        with torch.no_grad():
            with self.swap_ema_state():
                self.run_stage('valid', self.valid)
                # the best state is updated with EMA states if available
                self.update_best_state_from_stage('valid')
            with self.swap_best_state():
                if self.should_run_stage('evaluate'):
                    self.run_stage('evaluate', self.evaluate)
                if self.should_run_stage('generate'):
                    self.run_stage('generate', with_rank_rng()(self.generate))

    def run(self):
        """Training loop."""
        assert len(self.state_dict()) > 0
        self.restore(replay_metrics=True)  # load checkpoint and replay history
        self.log_hyperparams(dict_from_config(self.cfg))
        for epoch in range(self.epoch, self.cfg.optim.epochs + 1):
            if self.should_stop_training():
                return
            self.run_epoch()
            # Commit will send the metrics to Dora and save checkpoints by default.
            self.commit()

    def should_stop_training(self) -> bool:
        """Check whether we should stop training or not."""
        return self.epoch > self.cfg.optim.epochs

    def should_run_stage(self, stage_name) -> bool:
        """Check whether we want to run the specified stages."""
        stage_every = self.cfg[stage_name].get('every', None)
        is_last_epoch = self.epoch == self.cfg.optim.epochs
        is_epoch_every = (stage_every and self.epoch % stage_every == 0)
        return is_last_epoch or is_epoch_every

    @abstractmethod
    def run_step(self, idx: int, batch: tp.Any, metrics: dict):
        """Perform one training or valid step on a given batch."""
        ...

    def common_train_valid(self, dataset_split: str, **kwargs: tp.Any):
        """Common logic for train and valid stages."""
        self.model.train(self.is_training)

        loader = self.dataloaders[dataset_split]
        # get a different order for distributed training, otherwise this will get ignored
        if flashy.distrib.world_size() > 1 \
           and isinstance(loader.sampler, torch.utils.data.distributed.DistributedSampler):
            loader.sampler.set_epoch(self.epoch)
        updates_per_epoch = self.train_updates_per_epoch if self.is_training else len(loader)
        if self.cfg.benchmark_no_load:
            self.logger.warning("Fake loading for benchmarking: re-using first batch")
            batch = next(iter(loader))
            loader = [batch] * updates_per_epoch  # type: ignore
        lp = self.log_progress(self.current_stage, loader, total=updates_per_epoch, updates=self.log_updates)
        average = flashy.averager()  # epoch wise average
        instant_average = flashy.averager()  # average between two logging
        metrics: dict = {}

        with self.profiler, self.deadlock_detect:  # profiler will only run for the first 20 updates.
            for idx, batch in enumerate(lp):
                self.deadlock_detect.update('batch')
                if idx >= updates_per_epoch:
                    break
                metrics = {}
                metrics = self.run_step(idx, batch, metrics)
                self.deadlock_detect.update('step')
                # run EMA step
                if self.ema is not None and self.is_training and (idx + 1) % self.cfg.optim.ema.updates == 0:
                    self.logger.debug("EMA model step")
                    self.ema.step()
                self.deadlock_detect.update('ema')
                self.profiler.step()
                instant_metrics = instant_average(metrics)
                if lp.update(**instant_metrics):
                    instant_average = flashy.averager()  # reset averager between two logging
                metrics = average(metrics)  # epoch wise average
                self.deadlock_detect.update('end_batch')

        metrics = flashy.distrib.average_metrics(metrics, updates_per_epoch)
        return metrics

    def train(self):
        """Train stage."""
        return self.common_train_valid('train')

    def valid(self):
        """Valid stage."""
        return self.common_train_valid('valid')

    @abstractmethod
    def evaluate(self):
        """Evaluate stage."""
        ...

    @abstractmethod
    def generate(self):
        """Generate stage."""
        ...

    def run_one_stage(self, stage_name: str):
        """Run only the specified stage.
        This method is useful to only generate samples from a trained experiment
        or rerun the validation or evaluation stages.
        """
        fn = {
            'generate': with_rank_rng()(self.generate),
            'evaluate': self.evaluate,
            'valid': self.valid,
        }
        if stage_name not in fn:
            raise ValueError(f'Trying to run stage {stage_name} is not supported.')
        assert len(self.state_dict()) > 0
        self._start_epoch()
        with torch.no_grad(), self.swap_best_state():
            self.run_stage(stage_name, fn[stage_name])
        if not self.cfg.execute_inplace:
            self.commit(save_checkpoints=False)

    @staticmethod
    def get_eval_solver_from_sig(sig: str, dtype: tp.Optional[str] = None,
                                 device: tp.Optional[str] = None, autocast: bool = True,
                                 batch_size: tp.Optional[int] = None,
                                 override_cfg: tp.Optional[tp.Union[dict, omegaconf.DictConfig]] = None,
                                 **kwargs):
        """Mostly a convenience function around audiocraft.train.get_solver_from_sig,
        populating all the proper param, deactivating EMA, FSDP, loading the best state,
        basically all you need to get a solver ready to "play" with in single GPU mode
        and with minimal memory overhead.

        Args:
            sig (str): signature to load.
            dtype (str or None): potential dtype, as a string, i.e. 'float16'.
            device (str or None): potential device, as a string, i.e. 'cuda'.
            override_cfg (dict or omegaconf.DictConfig or None): potential device, as a string, i.e. 'cuda'.
        """
        from audiocraft import train
        our_override_cfg: tp.Dict[str, tp.Any] = {'optim': {'ema': {'use': False}}}
        our_override_cfg['autocast'] = autocast
        if dtype is not None:
            our_override_cfg['dtype'] = dtype
        if device is not None:
            our_override_cfg['device'] = device
        if batch_size is not None:
            our_override_cfg['dataset'] = {'batch_size': batch_size}
        if override_cfg is None:
            override_cfg = {}
        override_cfg = omegaconf.OmegaConf.merge(
            omegaconf.DictConfig(override_cfg), omegaconf.DictConfig(our_override_cfg))  # type: ignore
        solver = train.get_solver_from_sig(
            sig, override_cfg=override_cfg,
            load_best=True, disable_fsdp=True,
            ignore_state_keys=['optimizer', 'ema'], **kwargs)
        solver.model.eval()
        return solver