NVIDIA NeMo

Orchestrates large-scale LLM training across multi-GPU and multi-node clusters using NVIDIA's Megatron-Core strategy, which decomposes models into tensor-parallel shards (column/row parallelism across transformer layers), pipeline-parallel stages (vertical model splitting), and data-parallel batches. NeMo wraps Megatron's distributed optimizer and gradient accumulation patterns within PyTorch Lightning's training loop, automatically handling communication collectives (all-reduce, all-gather) and mixed-precision scaling across heterogeneous hardware.

Implements efficient LLM inference through KV-cache management (caching key-value projections across transformer layers to avoid recomputation) and streaming token-by-token generation with optional batching. NeMo's inference engine supports both greedy decoding and beam search with length penalties, integrating with HuggingFace's generation API while maintaining NVIDIA-optimized kernels (FlashAttention, Fused RoPE) for reduced latency. Supports both single-GPU and distributed inference via tensor parallelism for large models.

Implements distributed checkpoint saving and loading that preserves tensor-parallel model sharding across GPU ranks, avoiding the need to consolidate full model state on a single GPU. NeMo's distributed checkpointing saves each rank's model shard independently, along with metadata describing the parallelism topology (TP degree, PP stages, DP groups). Supports resuming training with the same parallelism configuration, and provides offline conversion tools for changing parallelism degrees without retraining.

Provides mechanisms for gracefully handling node failures, GPU preemption, and training interruptions in long-running distributed training jobs. NeMo integrates with PyTorch Lightning's fault tolerance callbacks and Megatron-Core's distributed checkpointing to enable automatic recovery from checkpoints. Supports preemption signals (SIGTERM) with graceful shutdown (saving checkpoint before exit) and automatic job resubmission on cluster managers (Slurm, Kubernetes).

Provides declarative model configuration using YAML files and Hydra framework for composable, reproducible experiment setup. NeMo's recipe system enables defining model architecture, training hyperparameters, data loading, and distributed training settings in YAML, with Hydra's config composition allowing easy experiment variations (e.g., changing learning rate, batch size, parallelism degrees). Supports config validation, default value inheritance, and automatic CLI argument generation from YAML configs.

Provides speaker verification models (speaker recognition, speaker identification) using speaker embedding extractors (e.g., ECAPA-TDNN, Titanet) that map audio to fixed-size speaker embeddings in a learned metric space. NeMo's speaker verification pipeline includes speaker enrollment (registering known speakers), speaker verification (comparing test audio to enrolled speakers), and speaker identification (classifying test audio to one of multiple speakers). Supports both speaker-dependent and speaker-independent models, and integrates with standard speaker verification datasets (VoxCeleb, TIMIT).

Provides end-to-end ASR pipelines supporting both streaming (online) and batch (offline) transcription using encoder-decoder architectures (Conformer, Squeezeformer) with CTC or RNN-T decoders. NeMo's ASR models integrate Lhotse for efficient audio data loading and augmentation (SpecAugment, time-stretching), and support both character-level and BPE tokenization. Streaming inference uses stateful RNN-T decoders with lookahead context, while batch inference leverages attention-based decoders for higher accuracy.

Generates natural speech from text using encoder-decoder TTS models (FastPitch, Glow-TTS, Radiance) with integrated grapheme-to-phoneme (G2P) conversion for handling out-of-vocabulary words and pronunciation rules. NeMo's TTS pipeline includes duration prediction (predicting phoneme lengths), pitch modeling (fundamental frequency contours), and optional vocoder integration (HiFi-GAN, UnivNet) for waveform synthesis. Supports both single-speaker and multi-speaker models with speaker embeddings for voice cloning.

Provides experiment tracking, checkpoint saving/loading, and training state management through tight integration with PyTorch Lightning's Trainer and Callback system. NeMo wraps Lightning's checkpoint logic with custom serialization for distributed models (handling tensor-parallel sharding, optimizer state, and training metadata). Supports resuming training from arbitrary checkpoints with automatic detection of parallelism configuration, and integrates with logging backends (Weights & Biases, TensorBoard, Neptune) for experiment monitoring.

Enables seamless loading of HuggingFace LLM checkpoints (Llama, Mistral, Qwen, Phi, etc.) into NeMo's training and inference pipelines via AutoModel-style APIs. NeMo's HF integration automatically converts HuggingFace model configs to NeMo's internal format, handles tokenizer loading, and supports both eager and lazy weight loading for memory efficiency. Supports fine-tuning HuggingFace models with NeMo's distributed training infrastructure without manual architecture porting.

Implements automatic mixed-precision (AMP) training using PyTorch's native AMP with NVIDIA's loss scaling strategies to prevent gradient underflow in lower-precision (float16, bfloat16) computations. NeMo's distributed optimizer integrates loss scaling across all ranks, automatically adjusting scale factors based on gradient overflow detection. Supports gradient accumulation for effective batch size increases without memory overhead, and integrates with distributed checkpointing to preserve optimizer state across restarts.

Provides pre-built NLP models and training pipelines for machine translation (seq2seq with attention), token-level classification (NER, POS tagging via transformer encoders), and text classification (sentiment, intent detection via transformer encoders with pooling). NeMo's NLP collection integrates with HuggingFace tokenizers and supports both encoder-only (BERT-style) and encoder-decoder (T5-style) architectures. Includes data loaders for common NLP datasets (CoNLL, SQuAD, GLUE) and evaluation metrics (BLEU, F1, accuracy).

Supports training and inference of vision-language models (e.g., CLIP-style models, image captioning, visual question answering) by combining image encoders (ViT, ResNet) with text encoders/decoders (transformers). NeMo's multimodal collection handles heterogeneous input modalities (images, text) with separate encoders and optional fusion layers, and supports both contrastive learning (image-text matching) and generative tasks (captioning, VQA). Integrates with standard vision datasets (COCO, Flickr30K) and supports distributed training across multi-GPU clusters.

Integrates Lhotse's declarative audio data pipeline for efficient, reproducible data loading and augmentation in speech tasks (ASR, TTS, speaker verification). Lhotse enables composable augmentation operations (SpecAugment, time-stretching, pitch shifting, noise injection) defined in YAML, with automatic caching and parallel I/O. NeMo's Lhotse integration handles on-the-fly augmentation during training, supports both streaming and batch data loading, and provides automatic dataset statistics computation for analysis.