soapboxguy/MusicGen

1# AudioCraft training pipelines2 3AudioCraft training pipelines are built on top of PyTorch as our core deep learning library4and [Flashy](https://github.com/facebookresearch/flashy) as our training pipeline design library,5and [Dora](https://github.com/facebookresearch/dora) as our experiment manager.6AudioCraft training pipelines are designed to be research and experiment-friendly.7 8 9## Environment setup10 11For the base installation, follow the instructions from the [README.md](../README.md).12Below are some additional instructions for setting up environment to train new models.13 14### Team and cluster configuration15 16In order to support multiple teams and clusters, AudioCraft uses an environment configuration.17The team configuration allows to specify cluster-specific configurations (e.g. SLURM configuration),18or convenient mapping of paths between the supported environments.19 20Each team can have a yaml file under the [configuration folder](../config). To select a team set the21`AUDIOCRAFT_TEAM` environment variable to a valid team name (e.g. `labs` or `default`):22```shell23conda env config vars set AUDIOCRAFT_TEAM=default24```25 26Alternatively, you can add it to your `.bashrc`:27```shell28export AUDIOCRAFT_TEAM=default29```30 31If not defined, the environment will default to the `default` team.32 33The cluster is automatically detected, but it is also possible to override it by setting34the `AUDIOCRAFT_CLUSTER` environment variable.35 36Based on this team and cluster, the environment is then configured with:37* The dora experiment outputs directory.38* The available slurm partitions: categorized by global and team.39* A shared reference directory: In order to facilitate sharing research models while remaining40agnostic to the used compute cluster, we created the `//reference` symbol that can be used in41YAML config to point to a defined reference folder containing shared checkpoints42(e.g. baselines, models for evaluation...).43 44**Important:** The default output dir for trained models and checkpoints is under `/tmp/`. This is suitable45only for quick testing. If you are doing anything serious you MUST edit the file `default.yaml` and46properly set the `dora_dir` entries.47 48#### Overriding environment configurations49 50You can set the following environmet variables to bypass the team's environment configuration:51* `AUDIOCRAFT_CONFIG`: absolute path to a team config yaml file.52* `AUDIOCRAFT_DORA_DIR`: absolute path to a custom dora directory.53* `AUDIOCRAFT_REFERENCE_DIR`: absolute path to the shared reference directory.54 55## Training pipelines56 57Each task supported in AudioCraft has its own training pipeline and dedicated solver.58Learn more about solvers and key designs around AudioCraft training pipeline below.59Please refer to the documentation of each task and model for specific information on a given task.60 61 62### Solvers63 64The core training component in AudioCraft is the solver. A solver holds the definition65of how to solve a given task: It implements the training pipeline logic, combining the datasets,66model, optimization criterion and components and the full training loop. We refer the reader67to [Flashy](https://github.com/facebookresearch/flashy) for core principles around solvers.68 69AudioCraft proposes an initial solver, the `StandardSolver` that is used as the base implementation70for downstream solvers. This standard solver provides a nice base management of logging,71checkpoints loading/saving, xp restoration, etc. on top of the base Flashy implementation.72In AudioCraft, we made the assumption that all tasks are following the same set of stages:73train, valid, evaluate and generation, each relying on a dedicated dataset.74 75Each solver is responsible for defining the task to solve and the associated stages76of the training loop in order to leave the full ownership of the training pipeline77to the researchers. This includes loading the datasets, building the model and78optimisation components, registering them and defining the execution of each stage.79To create a new solver for a given task, one should extend the StandardSolver80and define each stage of the training loop. One can further customise its own solver81starting from scratch instead of inheriting from the standard solver.82 83```python84from . import base85from .. import optim86 87 88class MyNewSolver(base.StandardSolver):89 90    def __init__(self, cfg: omegaconf.DictConfig):91        super().__init__(cfg)92        # one can add custom attributes to the solver93        self.criterion = torch.nn.L1Loss()94 95    def best_metric(self):96        # here optionally specify which metric to use to keep track of best state97        return 'loss'98 99    def build_model(self):100        # here you can instantiate your models and optimization related objects101        # this method will be called by the StandardSolver init method102        self.model = ...103        # the self.cfg attribute contains the raw configuration104        self.optimizer = optim.build_optimizer(self.model.parameters(), self.cfg.optim)105        # don't forget to register the states you'd like to include in your checkpoints!106        self.register_stateful('model', 'optimizer')107        # keep the model best state based on the best value achieved at validation for the given best_metric108        self.register_best('model')109        # if you want to add EMA around the model110        self.register_ema('model')111 112    def build_dataloaders(self):113        # here you can instantiate your dataloaders114        # this method will be called by the StandardSolver init method115        self.dataloaders = ...116 117    ...118 119    # For both train and valid stages, the StandardSolver relies on120    # a share common_train_valid implementation that is in charge of121    # accessing the appropriate loader, iterate over the data up to122    # the specified number of updates_per_epoch, run the ``run_step``123    # function that you need to implement to specify the behavior124    # and finally update the EMA and collect the metrics properly.125    @abstractmethod126    def run_step(self, idx: int, batch: tp.Any, metrics: dict):127        """Perform one training or valid step on a given batch.128        """129        ... # provide your implementation of the solver over a batch130 131    def train(self):132        """Train stage.133        """134        return self.common_train_valid('train')135 136    def valid(self):137        """Valid stage.138        """139        return self.common_train_valid('valid')140 141    @abstractmethod142    def evaluate(self):143        """Evaluate stage.144        """145        ... # provide your implementation here!146 147    @abstractmethod148    def generate(self):149        """Generate stage.150        """151        ... # provide your implementation here!152```153 154### About Epochs155 156AudioCraft Solvers uses the concept of Epoch. One epoch doesn't necessarily mean one pass over the entire157dataset, but instead represent the smallest amount of computation that we want to work with before checkpointing.158Typically, we find that having an Epoch time around 30min is ideal both in terms of safety (checkpointing often enough)159and getting updates often enough. One Epoch is at least a `train` stage that lasts for `optim.updates_per_epoch` (2000 by default),160and a `valid` stage. You can control how long the valid stage takes with `dataset.valid.num_samples`.161Other stages (`evaluate`, `generate`) will only happen every X epochs, as given by `evaluate.every` and `generate.every`).162 163 164### Models165 166In AudioCraft, a model is a container object that wraps one or more torch modules together167with potential processing logic to use in a solver. For example, a model would wrap an encoder module,168a quantisation bottleneck module, a decoder and some tensor processing logic. Each of the previous components169can be considered as a small « model unit » on its own but the container model is a practical component170to manipulate and train a set of modules together.171 172### Datasets173 174See the [dedicated documentation on datasets](./DATASETS.md).175 176### Metrics177 178See the [dedicated documentation on metrics](./METRICS.md).179 180### Conditioners181 182AudioCraft language models can be conditioned in various ways and the codebase offers a modular implementation183of different conditioners that can be potentially combined together.184Learn more in the [dedicated documentation on conditioning](./CONDITIONING.md).185 186### Configuration187 188AudioCraft's configuration is defined in yaml files and the framework relies on189[hydra](https://hydra.cc/docs/intro/) and [omegaconf](https://omegaconf.readthedocs.io/) to parse190and manipulate the configuration through Dora.191 192##### :warning: Important considerations around configurations193 194Our configuration management relies on Hydra and the concept of group configs to structure195and compose configurations. Updating the root default configuration files will then have196an impact on all solvers and tasks.197**One should never change the default configuration files. Instead they should use Hydra config groups in order to store custom configuration.**198Once this configuration is created and used for running experiments, you should not edit it anymore.199 200Note that as we are using Dora as our experiment manager, all our experiment tracking is based on201signatures computed from delta between configurations.202**One must therefore ensure backward compatibilty of the configuration at all time.**203See [Dora's README](https://github.com/facebookresearch/dora) and the204[section below introduction Dora](#running-experiments-with-dora).205 206##### Configuration structure207 208The configuration is organized in config groups:209* `conditioner`: default values for conditioning modules.210* `dset`: contains all data source related information (paths to manifest files211and metadata for a given dataset).212* `model`: contains configuration for each model defined in AudioCraft and configurations213for different variants of models.214* `solver`: contains the default configuration for each solver as well as configuration215for each solver task, combining all the above components.216* `teams`: contains the cluster configuration per teams. See environment setup for more details.217 218The `config.yaml` file is the main configuration that composes the above groups219and contains default configuration for AudioCraft.220 221##### Solver's core configuration structure222 223The core configuration structure shared across solver is available in `solvers/default.yaml`.224 225##### Other configuration modules226 227AudioCraft configuration contains the different setups we used for our research and publications.228 229## Running experiments with Dora230 231### Launching jobs232 233Try launching jobs for different tasks locally with dora run:234 235```shell236# run compression task with lightweight encodec237dora run solver=compression/debug238```239 240Most of the time, the jobs are launched through dora grids, for example:241 242```shell243# run compression task through debug grid244dora grid compression.debug245```246 247Learn more about running experiments with Dora below.248 249### A small introduction to Dora250 251[Dora](https://github.com/facebookresearch/dora) is the experiment manager tool used in AudioCraft.252Check out the README to learn how Dora works. Here is a quick summary of what to know:253* An XP is a unique set of hyper-parameters with a given signature. The signature is a hash254of those hyper-parameters. We always refer to an XP with its signature, e.g. 9357e12e. We will see255after that one can retrieve the hyper-params and re-rerun it in a single command.256* In fact, the hash is defined as a delta between the base config and the one obtained257with the config overrides you passed from the command line. This means you must never change258the `conf/**.yaml` files directly., except for editing things like paths. Changing the default values259in the config files means the XP signature won't reflect that change, and wrong checkpoints might be reused.260I know, this is annoying, but the reason is that otherwise, any change to the config file would mean261that all XPs ran so far would see their signature change.262 263#### Dora commands264 265```shell266dora info -f 81de367c  # this will show the hyper-parameter used by a specific XP.267                       # Be careful some overrides might present twice, and the right most one268                       # will give you the right value for it.269 270dora run -d -f 81de367c   # run an XP with the hyper-parameters from XP 81de367c.271                          # `-d` is for distributed, it will use all available GPUs.272 273dora run -d -f 81de367c dataset.batch_size=32  # start from the config of XP 81de367c but change some hyper-params.274                                               # This will give you a new XP with a new signature (e.g. 3fe9c332).275 276dora info -f SIG -t    # will tail the log (if the XP has scheduled).277# if you need to access the logs of the process for rank > 0, in particular because a crash didn't happen in the main278# process, then use `dora info -f SIG` to get the main log name (finished into something like `/5037674_0_0_log.out`)279# and worker K can accessed as `/5037674_0_{K}_log.out`.280# This is only for scheduled jobs, for local distributed runs with `-d`, then you should go into the XP folder,281# and look for `worker_{K}.log` logs.282```283 284An XP runs from a specific folder based on its signature, under the285`<cluster_specific_path>/<user>/experiments/audiocraft/outputs/` folder.286You can safely interrupt a training and resume it, it will reuse any existing checkpoint,287as it will reuse the same folder. If you made some change to the code and need to ignore288a previous checkpoint you can use `dora run --clear [RUN ARGS]`.289 290If you have a Slurm cluster, you can also use the dora grid command, e.g.291 292```shell293# run a dummy grid located at `audiocraft/grids/my_grid_folder/my_grid_name.py`294dora grid my_grid_folder.my_grid_name295# Run the following will simply display the grid and also initialized the Dora experiments database.296# You can then simply refer to a config using its signature (e.g. as `dora run -f SIG`).297dora grid my_grid_folder.my_grid_name --dry_run --init298```299 300Please refer to the [Dora documentation](https://github.com/facebookresearch/dora) for more information.301 302 303#### Clearing up past experiments304 305```shell306# This will cancel all the XPs and delete their folder and checkpoints.307# It will then reschedule them starting from scratch.308dora grid my_grid_folder.my_grid_name --clear309# The following will delete the folder and checkpoint for a single XP,310# and then run it afresh.311dora run [-f BASE_SIG] [ARGS] --clear312```313