NeMo

NeMo abstracts distributed training through PyTorch Lightning's Trainer API, automatically handling data parallelism, tensor parallelism, and pipeline parallelism across multi-GPU and multi-node clusters. The framework manages distributed state through a custom Application State system that coordinates optimizer steps, gradient accumulation, and checkpoint synchronization across ranks without requiring manual distributed communication code.

NeMo implements a custom Neural Types system that annotates module inputs/outputs with semantic type information (e.g., 'audio_signal', 'logits', 'embeddings') and validates tensor shapes, dtypes, and semantic compatibility at module connection time. This catches shape mismatches and type errors before training begins, preventing silent failures from incompatible layer connections.

NeMo provides NLP training pipelines supporting token classification (NER, POS tagging), machine translation, question answering, and text classification through transformer-based architectures. The NLP module integrates with HuggingFace tokenizers, supports multi-lingual training, and includes task-specific loss functions and evaluation metrics.

NeMo provides training pipelines for speech language models that process raw audio and text jointly, supporting architectures like Canary (multilingual speech-to-text and speech-to-speech translation). The SLM module handles audio-text alignment, multi-task training (ASR, translation, speech-to-speech), and supports both supervised and self-supervised pre-training.

NeMo automatically generates model cards (YAML/JSON) containing training configuration, performance metrics, dataset information, and usage guidelines. The model card system integrates with the .nemo artifact format, enabling automatic documentation generation and integration with model hubs (Hugging Face, NVIDIA NGC).

NeMo uses OmegaConf for declarative model and training configuration, supporting nested YAML files, environment variable interpolation, and dynamic config composition. The ExperimentManager integrates with this config system to automatically log hyperparameters, create experiment directories, and manage checkpoints, enabling reproducible training runs with minimal code.

NeMo packages trained models as .nemo files (TAR archives) containing model weights, config, tokenizers, and metadata via a SaveRestoreConnector abstraction. This enables single-file model distribution with all dependencies, supporting both local and cloud storage backends (S3, GCS) and automatic model card generation for reproducibility.

NeMo provides end-to-end ASR training pipelines supporting Conformer, Squeezeformer, and Citrinet architectures with integrated data augmentation (SpecAugment, time-stretching), language model integration, and CTC/RNN-T decoding. The ASR module handles audio preprocessing (MFCC, mel-spectrogram), feature normalization, and multi-lingual training through a modular encoder-decoder design.

NeMo provides TTS training pipelines supporting Glow-TTS, FastPitch, and Tacotron2 acoustic models paired with neural vocoders (HiFi-GAN, UnivNet). The TTS module handles text preprocessing, phoneme alignment, mel-spectrogram generation, and multi-speaker training through a modular acoustic-model + vocoder architecture with automatic data augmentation.

NeMo integrates Lhotse for declarative audio data loading, supporting manifest files (JSON lines) that specify audio paths, transcriptions, speaker IDs, and metadata. The data pipeline handles on-the-fly audio loading, resampling, augmentation, and batching through a composable DataLoader abstraction, enabling efficient training on large datasets without pre-processing.

NeMo implements adapter modules (LoRA, prefix-tuning, adapter layers) that enable fine-tuning large pre-trained models with <5% of original parameters. Adapters are inserted into model layers via a declarative configuration system and can be trained separately from frozen base weights, enabling efficient multi-task fine-tuning and model composition.

NeMo's SaveRestoreConnector implements distributed checkpointing that handles rank-aware state serialization across multi-GPU training, supporting both sharded and replicated state patterns. Checkpoints can be saved asynchronously without blocking training, and the system automatically handles optimizer state, model weights, and training metadata across distributed ranks.

NeMo provides a comprehensive learning rate scheduler system supporting warmup (linear, exponential), decay (cosine, polynomial, exponential), and composite schedules through a declarative configuration interface. Schedulers integrate with distributed optimizers and support per-parameter-group scheduling for fine-grained control over learning rates across model components.