Basics
NeMo models contain everything needed to train and reproduce conversational AI models:
neural network architectures
datasets/data loaders
data preprocessing/postprocessing
data augmentors
optimizers and schedulers
tokenizers
language models
NeMo uses
for configuring both NeMo models and the PyTorch Lightning Trainer.
Note
Every NeMo model has an example configuration file and training script that can be found
.
The end result of using NeMo,
, and Hydra is that NeMo models all have the same look and feel and are also fully compatible with the PyTorch ecosystem.
Pretrained
NeMo comes with many pretrained models for each of our collections: ASR, TTS, Audio, and SpeechLM2. Every pretrained NeMo model can be downloaded and used with the from_pretrained() method.
As an example, we can instantiate a Parakeet model with the following:
importnemo.collections.asrasnemo_asrmodel=nemo_asr.models.ASRModel.from_pretrained(model_name="nvidia/parakeet-tdt-0.6b-v2")To see all available pretrained models for a specific NeMo model, use the list_available_models() method:
nemo_asr.models.EncDecCTCModel.list_available_models()For detailed information on the available pretrained models, refer to the collections documentation:
Automatic Speech Recognition (ASR)
Text-to-Speech Synthesis (TTS)
Training
NeMo leverages
for model training. PyTorch Lightning lets NeMo decouple the conversational AI code from the PyTorch training code. This means that NeMo users can focus on their domain (ASR, NLP, TTS) and build complex AI applications without having to rewrite boilerplate code for PyTorch training.
When using PyTorch Lightning, NeMo users can automatically train with:
multi-GPU/multi-node
mixed precision
model checkpointing
logging
early stopping
and more
The two main aspects of the Lightning API are the
and the
.
PyTorch Lightning LightningModule
Every NeMo model is a LightningModule which is an nn.module. This means that NeMo models are compatible with the PyTorch ecosystem and can be plugged into existing PyTorch workflows.
Creating a NeMo model is similar to any other PyTorch workflow. We start by initializing our model architecture, then define the forward pass:
classTextClassificationModel(NLPModel,Exportable):...def__init__(self,cfg:DictConfig,trainer:Trainer=None):"""Initializes the BERTTextClassifier model."""...super().__init__(cfg=cfg,trainer=trainer)# instantiate a BERT based encoderself.bert_model=get_lm_model(config_file=cfg.language_model.config_file,config_dict=cfg.language_model.config,vocab_file=cfg.tokenizer.vocab_file,trainer=trainer,cfg=cfg,)# instantiate the FFN for classificationself.classifier=SequenceClassifier(hidden_size=self.bert_model.config.hidden_size,num_classes=cfg.dataset.num_classes,num_layers=cfg.classifier_head.num_output_layers,activation='relu',log_softmax=False,dropout=cfg.classifier_head.fc_dropout,use_transformer_init=True,idx_conditioned_on=0,)defforward(self,input_ids,token_type_ids,attention_mask):""" No special modification required for Lightning, define it as you normally would in the `nn.Module` in vanilla PyTorch. """hidden_states=self.bert_model(input_ids=input_ids,token_type_ids=token_type_ids,attention_mask=attention_mask)logits=self.classifier(hidden_states=hidden_states)returnlogitsThe LightningModule organizes PyTorch code so that across all NeMo models we have a similar look and feel. For example, the training logic can be found in training_step:
deftraining_step(self,batch,batch_idx):""" Lightning calls this inside the training loop with the data from the training dataloader passed in as `batch`. """# forward passinput_ids,input_type_ids,input_mask,labels=batchlogits=self.forward(input_ids=input_ids,token_type_ids=input_type_ids,attention_mask=input_mask)train_loss=self.loss(logits=logits,labels=labels)lr=self._optimizer.param_groups[0]['lr']self.log('train_loss',train_loss)self.log('lr',lr,prog_bar=True)return{'loss':train_loss,'lr':lr,}While validation logic can be found in validation_step:
defvalidation_step(self,batch,batch_idx):""" Lightning calls this inside the validation loop with the data from the validation dataloader passed in as `batch`. """ifself.testing:prefix='test'else:prefix='val'input_ids,input_type_ids,input_mask,labels=batchlogits=self.forward(input_ids=input_ids,token_type_ids=input_type_ids,attention_mask=input_mask)val_loss=self.loss(logits=logits,labels=labels)preds=torch.argmax(logits,axis=-1)tp,fn,fp,_=self.classification_report(preds,labels)return{'val_loss':val_loss,'tp':tp,'fn':fn,'fp':fp}PyTorch Lightning then handles all of the boilerplate code needed for training. Virtually any aspect of training can be customized via PyTorch Lightning
,
,
, or by overriding
.
For more domain-specific information, see:
Automatic Speech Recognition (ASR)
Text-to-Speech Synthesis (TTS)
PyTorch Lightning Trainer
Since every NeMo model is a LightningModule, we can automatically take advantage of the PyTorch Lightning Trainer. Every NeMo
training script uses the Trainer object to fit the model.
First, instantiate the model and trainer, then call .fit:
# We first instantiate the trainer based on the model configuration.# See the model configuration documentation for details.trainer=pl.Trainer(**cfg.trainer)# Then pass the model configuration and trainer object into the NeMo modelmodel=TextClassificationModel(cfg.model,trainer=trainer)# Now we can train with by calling .fittrainer.fit(model)# Or we can run the test loop on test data by callingtrainer.test(model=model)All
can be set from from the NeMo configuration.
Configuration
Hydra is an open-source Python framework that simplifies configuration for complex applications that must bring together many different software libraries. Conversational AI model training is a great example of such an application. To train a conversational AI model, we must be able to configure:
neural network architectures
training and optimization algorithms
data pre/post processing
data augmentation
experiment logging/visualization
model checkpointing
For an introduction to using Hydra, refer to the
.
With Hydra, we can configure everything needed for NeMo with three interfaces:
Command Line (CLI)
Configuration Files (YAML)
Dataclasses (Python)
YAML
NeMo provides YAML configuration files for all of our
training scripts. YAML files make it easy to experiment with different model and training configurations.
Every NeMo example YAML has the same underlying configuration structure:
trainer
exp_manager
model
The model configuration always contains train_ds, validation_ds, test_ds, and optim. Model architectures, however, can vary across domains. Refer to the documentation of specific collections (LLM, ASR etc.) for detailed information on model architecture configuration.
A NeMo configuration file should look similar to the following:
# PyTorch Lightning Trainer configuration# any argument of the Trainer object can be set heretrainer:devices:1# number of gpus per nodeaccelerator:gpunum_nodes:1# number of nodesmax_epochs:10# how many training epochs to runval_check_interval:1.0# run validation after every epoch# Experiment logging configurationexp_manager:exp_dir:/path/to/my/nemo/experimentsname:name_of_my_experimentcreate_tensorboard_logger:Truecreate_wandb_logger:True# Model configuration# model network architecture, train/val/test datasets, data augmentation, and optimizationmodel:train_ds:manifest_filepath:/path/to/my/train/manifest.jsonbatch_size:256shuffle:Truevalidation_ds:manifest_filepath:/path/to/my/validation/manifest.jsonbatch_size:32shuffle:Falsetest_ds:manifest_filepath:/path/to/my/test/manifest.jsonbatch_size:32shuffle:Falseoptim:name:novogradlr:.01betas:[0.8,0.5]weight_decay:0.001# network architecture can vary greatly depending on the domainencoder:...decoder:...CLI
With NeMo and Hydra, every aspect of model training can be modified from the command-line. This is extremely helpful for running lots of experiments on compute clusters or for quickly testing parameters during development.
All NeMo
come with instructions on how to run the training/inference script from the command-line (e.g. see
for an example).
With Hydra, arguments are set using the = operator:
pythonexamples/asr/asr_ctc/speech_to_text_ctc.py\model.train_ds.manifest_filepath=/path/to/my/train/manifest.json\model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json\trainer.devices=2\trainer.accelerator='gpu'\trainer.max_epochs=50We can use the + operator to add arguments from the CLI:
pythonexamples/asr/asr_ctc/speech_to_text_ctc.py\model.train_ds.manifest_filepath=/path/to/my/train/manifest.json\model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json\trainer.devices=2\trainer.accelerator='gpu'\trainer.max_epochs=50\+trainer.fast_dev_run=trueWe can use the ~ operator to remove configurations:
pythonexamples/asr/asr_ctc/speech_to_text_ctc.py\model.train_ds.manifest_filepath=/path/to/my/train/manifest.json\model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json\~model.test_ds\trainer.devices=2\trainer.accelerator='gpu'\trainer.max_epochs=50\+trainer.fast_dev_run=trueWe can specify configuration files using the --config-path and --config-name flags:
pythonexamples/asr/asr_ctc/speech_to_text_ctc.py\--config-path=conf/conformer\--config-name=conformer_ctc_bpe\model.train_ds.manifest_filepath=/path/to/my/train/manifest.json\model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json\~model.test_ds\trainer.devices=2\trainer.accelerator='gpu'\trainer.max_epochs=50\+trainer.fast_dev_run=trueDataclasses
Dataclasses allow NeMo to ship model configurations as part of the NeMo library and also enables pure Python configuration of NeMo models. With Hydra, dataclasses can be used to create
for the conversational AI application.
As an example, refer to the code block below for an Attenion is All You Need machine translation model. The model configuration can be instantiated and modified like any Python
.
fromnemo.collections.nlp.models.machine_translation.mt_enc_dec_configimportAAYNBaseConfigcfg=AAYNBaseConfig()# modify the number of layers in the encodercfg.encoder.num_layers=8# modify the training batch sizecfg.train_ds.tokens_in_batch=8192Note
Configuration with Hydra always has the following precedence CLI > YAML > Dataclass.
Optimization
Optimizers and learning rate schedules are configurable across all NeMo models and have their own namespace. Here is a sample YAML configuration for a Novograd optimizer with a Cosine Annealing learning rate schedule.
optim:name:novogradlr:0.01# optimizer argumentsbetas:[0.8,0.25]weight_decay:0.001# scheduler setupsched:name:CosineAnnealing# Optional argumentsmax_steps:-1# computed at runtime or explicitly set heremonitor:val_lossreduce_on_plateau:false# scheduler config overridewarmup_steps:1000warmup_ratio:nullmin_lr: 1e-9:Note
has optimizer and scheduler configurations for every NeMo model.
Optimizers can be configured from the CLI as well:
pythonexamples/asr/asr_ctc/speech_to_text_ctc.py\--config-path=conf/conformer\--config-name=conformer_ctc_bpe\... # train with the adam optimizermodel.optim=adam\# change the learning ratemodel.optim.lr=.0004\# modify betasmodel.optim.betas=[.8,.5]Optimizers
name corresponds to the lowercase name of the optimizer. To view a list of available optimizers, run:
fromnemo.core.optim.optimizersimportAVAILABLE_OPTIMIZERSforname,optinAVAILABLE_OPTIMIZERS.items():print(f'name: {name}, opt: {opt}')name:sgdopt:<class'torch.optim.sgd.SGD'> name:adamopt:<class'torch.optim.adam.Adam'> name:adamwopt:<class'torch.optim.adamw.AdamW'> name:adadeltaopt:<class'torch.optim.adadelta.Adadelta'> name:adamaxopt:<class'torch.optim.adamax.Adamax'> name:adagradopt:<class'torch.optim.adagrad.Adagrad'> name:rmspropopt:<class'torch.optim.rmsprop.RMSprop'> name:rpropopt:<class'torch.optim.rprop.Rprop'> name:novogradopt:<class'nemo.core.optim.novograd.Novograd'> Optimizer Params
Optimizer params can vary between optimizers but the lr param is required for all optimizers. To see the available params for an optimizer, we can look at its corresponding dataclass.
fromnemo.core.config.optimizersimportNovogradParamsprint(NovogradParams())NovogradParams(lr='???',betas=(0.95,0.98),eps=1e-08,weight_decay=0,grad_averaging=False,amsgrad=False,luc=False,luc_trust=0.001,luc_eps=1e-08)'???' indicates that the lr argument is required.
Register Optimizer
To register a new optimizer to be used with NeMo, run:
nemo.core.optim.optimizers.register_optimizer(name:str,optimizer:
,optimizer_params:OptimizerParams,)
Checks if the optimizer name exists in the registry, and if it doesnt, adds it.
This allows custom optimizers to be added and called by name during instantiation.
Parameters:name – Name of the optimizer. Will be used as key to retrieve the optimizer.
optimizer – Optimizer class
optimizer_params – The parameters as a dataclass of the optimizer
Learning Rate Schedulers
Learning rate schedulers can be optionally configured under the optim.sched namespace.
name corresponds to the name of the learning rate schedule. To view a list of available schedulers, run:
fromnemo.core.optim.lr_schedulerimportAVAILABLE_SCHEDULERSforname,optinAVAILABLE_SCHEDULERS.items():print(f'name: {name}, schedule: {opt}')name:WarmupPolicy,schedule:<class'nemo.core.optim.lr_scheduler.WarmupPolicy'> name:WarmupHoldPolicy,schedule:<class'nemo.core.optim.lr_scheduler.WarmupHoldPolicy'> name:SquareAnnealing,schedule:<class'nemo.core.optim.lr_scheduler.SquareAnnealing'> name:CosineAnnealing,schedule:<class'nemo.core.optim.lr_scheduler.CosineAnnealing'> name:NoamAnnealing,schedule:<class'nemo.core.optim.lr_scheduler.NoamAnnealing'> name:WarmupAnnealing,schedule:<class'nemo.core.optim.lr_scheduler.WarmupAnnealing'> name:InverseSquareRootAnnealing,schedule:<class'nemo.core.optim.lr_scheduler.InverseSquareRootAnnealing'> name:SquareRootAnnealing,schedule:<class'nemo.core.optim.lr_scheduler.SquareRootAnnealing'> name:PolynomialDecayAnnealing,schedule:<class'nemo.core.optim.lr_scheduler.PolynomialDecayAnnealing'> name:PolynomialHoldDecayAnnealing,schedule:<class'nemo.core.optim.lr_scheduler.PolynomialHoldDecayAnnealing'> name:StepLR,schedule:<class'torch.optim.lr_scheduler.StepLR'> name:ExponentialLR,schedule:<class'torch.optim.lr_scheduler.ExponentialLR'> name:ReduceLROnPlateau,schedule:<class'torch.optim.lr_scheduler.ReduceLROnPlateau'> name:CyclicLR,schedule:<class'torch.optim.lr_scheduler.CyclicLR'> Scheduler Params
To see the available params for a scheduler, we can look at its corresponding dataclass:
fromnemo.core.config.schedulersimportCosineAnnealingParamsprint(CosineAnnealingParams())CosineAnnealingParams(last_epoch=-1,warmup_steps=None,warmup_ratio=None,min_lr=0.0)Register scheduler
To register a new scheduler to be used with NeMo, run:
nemo.core.optim.lr_scheduler.register_scheduler(name:str,scheduler:_LRScheduler,scheduler_params:SchedulerParams,)
Checks if the scheduler name exists in the registry, and if it doesnt, adds it.
This allows custom schedulers to be added and called by name during instantiation.
Parameters:name – Name of the optimizer. Will be used as key to retrieve the optimizer.
scheduler – Scheduler class (inherits from _LRScheduler)
scheduler_params – The parameters as a dataclass of the scheduler
Save and Restore
NeMo models all come with .save_to and .restore_from methods.
Save
To save a NeMo model, run:
model.save_to('/path/to/model.nemo')Everything needed to use the trained model is packaged and saved in the .nemo file. For example, in the NLP domain, .nemo files include the necessary tokenizer models and/or vocabulary files, etc.
Note
A .nemo file is simply an archive like any other .tar file.
Restore
To restore a NeMo model, run:
# Here, you should usually use the class of the model, or simply use ModelPT.restore_from() for simplicity.model.restore_from('/path/to/model.nemo')When using the PyTorch Lightning Trainer, a PyTorch Lightning checkpoint is created. These are mainly used within NeMo to auto-resume training. Since NeMo models are LightningModules, the PyTorch Lightning method load_from_checkpoint is available. Note that load_from_checkpoint won’t necessarily work out-of-the-box for all models as some models require more artifacts than just the checkpoint to be restored. For these models, the user will have to override load_from_checkpoint if they want to use it.
It’s highly recommended to use restore_from to load NeMo models.
Restore with Modified Config
Sometimes, there may be a need to modify the model (or it’s sub-components) prior to restoring a model. A common case is when the model’s internal config must be updated due to various reasons (such as deprecation, newer versioning, support a new feature). As long as the model has the same parameters as compared to the original config, the parameters can once again be restored safely.
In NeMo, as part of the .nemo file, the model’s internal config will be preserved. This config is used during restoration, and as shown below we can update this config prior to restoring the model.
# When restoring a model, you should generally use the class of the model# Obtain the config (as an OmegaConf object)config=model_class.restore_from('/path/to/model.nemo',return_config=True)# ORconfig=model_class.from_pretrained('name_of_the_model',return_config=True)# Modify the config as neededconfig.x.y=z# Restore the model from the updated configmodel=model_class.restore_from('/path/to/model.nemo',override_config_path=config)# ORmodel=model_class.from_pretrained('name_of_the_model',override_config_path=config)Register Artifacts
Restoring conversational AI models can be complicated because it requires more than just the checkpoint weights; additional information is also needed to use the model. NeMo models can save additional artifacts in the .nemo file by calling .register_artifact. When restoring NeMo models using .restore_from or .from_pretrained, any artifacts that were registered will be available automatically.
As an example, consider an NLP model that requires a trained tokenizer model. The tokenizer model file can be automatically added to the .nemo file with the following:
self.encoder_tokenizer=get_nmt_tokenizer(...tokenizer_model=self.register_artifact(config_path='encoder_tokenizer.tokenizer_model',src='/path/to/tokenizer.model',verify_src_exists=True),)By default, .register_artifact will always return a path. If the model is being restored from a .nemo file, then that path will be to the artifact in the .nemo file. Otherwise, .register_artifact will return the local path specified by the user.
config_path is the artifact key. It usually corresponds to a model configuration but does not have to. The model config that is packaged with the .nemo file will be updated according to the config_path key. In the above example, the model config will have
encoder_tokenizer:...tokenizer_model:nemo:4978b28103264263a03439aaa6560e5e_tokenizer.modelsrc is the path to the artifact and the base-name of the path will be used when packaging the artifact in the .nemo file. Each artifact will have a hash prepended to the basename of src in the .nemo file. This is to prevent collisions with basenames base-names that are identical (say when there are two or more tokenizers, both called tokenizer.model). The resulting .nemo file will then have the following file:
4978b28103264263a03439aaa6560e5e_tokenizer.model If verify_src_exists is set to False, then the artifact is optional. This means that .register_artifact will return None if the src cannot be found.
Push to Hugging Face Hub
NeMo models can be pushed to the
with the
method. This method performs the same actions as save_to() and then uploads the model to the HuggingFace Hub. It offers an additional pack_nemo_file argument that allows the user to upload the entire NeMo file or just the .nemo file. This is useful for large language models that have a massive number of parameters, and a single NeMo file could exceed the max upload size of Hugging Face Hub.
Upload a model to the Hub
token="<HF TOKEN>"orNonepack_nemo_file=True# False will upload multiple files that comprise the NeMo file onto HF Hub; Generally useful for LLMsmodel.push_to_hf_hub(repo_id=repo_id,pack_nemo_file=pack_nemo_file,token=token,)Use a Custom Model Card Template for the Hub
# Override the default model cardtemplate=""" <Your own custom template># {model_name}"""kwargs={"model_name":"ABC","repo_id":"nvidia/ABC_XYZ"}model_card=model.generate_model_card(template=template,template_kwargs=kwargs,type="hf")model.push_to_hf_hub(repo_id=repo_id,token=token,model_card=model_card)# Write your own model card classclassMyModelCard:def__init__(self,model_name):self.model_name=model_namedef__repr__(self):template="""This is the {model_name} model""".format(model_name=self.model_name)returntemplatemodel.push_to_hf_hub(repo_id=repo_id,token=token,model_card=MyModelCard("ABC"))Nested NeMo Models
In some cases, it may be helpful to use NeMo models inside other NeMo models. For example, we can incorporate language models into ASR models to use in a decoding process to improve accuracy.
There are three ways to instantiate child models inside parent models:
use subconfig directly
use the .nemo checkpoint path to load the child model
use a pretrained NeMo model
To register a child model, use the register_nemo_submodule method of the parent model. This method will add the child model to a specified model attribute. During serialization, it will correctly handle child artifacts and store the child model’s configuration in the parent model’s config_field.
fromnemo.core.classesimportModelPTclassChildModel(ModelPT):...# implement necessary methodsclassParentModel(ModelPT):def__init__(self,cfg,trainer=None):super().__init__(cfg=cfg,trainer=trainer)# optionally annotate type for IDE autocompletion and type checkingself.child_model:Optional[ChildModel]ifcfg.get("child_model")isnotNone:# load directly from config# either if config provided initially, or automatically# after model restorationself.register_nemo_submodule(name="child_model",config_field="child_model",model=ChildModel(self.cfg.child_model,trainer=trainer),)elifcfg.get('child_model_path')isnotNone:# load from .nemo model checkpoint# while saving, config will be automatically assigned/updated# in cfg.child_modelself.register_nemo_submodule(name="child_model",config_field="child_model",model=ChildModel.restore_from(self.cfg.child_model_path,trainer=trainer),)elifcfg.get('child_model_name')isnotNone:# load from pretrained model# while saving, config will be automatically assigned/updated# in cfg.child_modelself.register_nemo_submodule(name="child_model",config_field="child_model",model=ChildModel.from_pretrained(self.cfg.child_model_name,trainer=trainer),)else:self.child_model=NoneProfiling
NeMo offers users two options for profiling: Nsys and CUDA memory profiling. These two options allow users to debug performance issues as well as memory issues such as memory leaks.
To enable Nsys profiling, add the following options to the model config:
nsys_profile:Falsestart_step:10# Global batch to start profilingend_step:10# Global batch to end profilingranks:[0]# Global rank IDs to profilegen_shape:False# Generate model and kernel details including input shapesFinally, run the model training script with:
nsysprofile-snone-o<profilefilepath>-tcuda,nvtx--force-overwritetrue--capture-range=cudaProfilerApi--capture-range-end=stoppython./examples/... See more options at
.
To enable CUDA memory profiling, add the following options to the model config:
memory_profile:enabled:Truestart_step:10# Global batch to start profilingend_step:10# Global batch to end profilingrank:0# Global rank ID to profileoutput_path:None# Path to store the profile output fileThen invoke your NeMo script without any changes in the invocation command.