Model YAML configurations

Before training, you must select the model configuration you wish to train. Please refer to the key features for a description of the options available, as well as the training times. Having selected a configuration it is necessary to note the config path and sentencepiece vocabulary size (“spm size”) of your chosen config from the following table as these will be needed in the subsequent data preparation steps:

NameParametersspm sizeconfigAcceleration supported?
testing49M1023testing-1023sp.yaml❌
base85M8703base-8703sp.yaml✅
large196M17407large-17407sp.yaml✅

It is recommended to train the base model on LibriSpeech as described here before training base or large on your own data.

Note

The base and large architectures were optimized to provide a good tradeoff between WER and throughput on the accelerator. Other architectures will not run on the accelerator.

train.sh will verify that you are training a model that is supported by the accelerator. If you want to skip this check so you can train the testing model for more rapid iteration, pass the flag --skip_state_dict_check to train.sh.

Post-RNN residual schedule

Very deep encoders (e.g. the XL model, with 17 post-RNN layers) do not converge as a plain stack of LSTMs. Passing --residual_schedule adds an identity residual around every post-RNN layer after the first, which keeps the deep stack trainable early on. The residual strength is then linearly annealed from 1 to 0 between --residual_toggle_step and --residual_toggle_step + --residual_ramp_steps, after which the model is a plain stack again.

The residual adds no parameters and is a training-only device: it is not part of the model config and is not included in the hardware checkpoint. Hardware checkpoints are therefore only written once the residual has been fully annealed away.

Missing YAML fields

The configs referenced above are not intended to be edited directly. Instead, they are used as templates to create <config-name>_run.yaml files. The _run.yaml file is a copy of the chosen config with the following fields populated:

  • sentpiece_model: /datasets/sentencepieces/SENTENCEPIECE.model
  • stats_path: /datasets/stats/STATS_SUBDIR
  • max_duration: MAX_DURATION
  • max_transcript_len: MAX_TOKENS
  • ngram_path: /datasets/ngrams/NGRAM_SUBDIR

Populating these fields can be performed by the training/scripts/create_config_set_env.sh script.

For example usage, see the following documentation: Prepare LibriSpeech in the JSONL format.