NVIDIA’s NeMo (Neural Modules) is one of those frameworks that quietly backs a large fraction of the open speech models you have heard of — Parakeet, FastConformer, Canary, Magpie — while also carrying NVIDIA’s newer speech-aware language models. Unlike a thin training loop or a serving-only project, NeMo is an opinionated framework: every network is assembled from typed, reusable modules, every model lives in a config-driven Lightning subclass, and every checkpoint can round-trip through a single .nemo file. If you have read our earlier tours of PyTorch or inference servers like Triton Inference Server, this post completes the picture by showing where a framework like NeMo sits between them: it consumes PyTorch idioms upstream and produces deployable artifacts downstream.

The premise of the design is deceptively simple. A “neural module” takes typed inputs and produces typed outputs, and a “model” is just a trainable composition of those modules plus a YAML configuration. Everything else — the model registry that resolves names like nvidia/parakeet-tdt-0.6b-v2 from NGC or Hugging Face Hub, the checkpoint connectors, the experiment management around a Lightning Trainer — is machinery bolted onto that core contract. Reading the source of the current main branch, the shape of the repository is crisp: nemo/core defines the contracts, nemo/collections ships the ASR, TTS, audio and SpeechLM2 model families, nemo/lightning glues NeMo callbacks into PyTorch Lightning, and nemo/utils holds the experiment and checkpoint plumbing. Notably, the LLM training collection that earlier releases shipped under nemo/collections/llm has been split out of this repository, so today’s NeMo proper is speech- and multimodal-first — an important mental-model correction for anyone returning after a year or two.

In this tour we walk that architecture from the inside out: the base classes that every model inherits, the type system that validates wiring, the collections that package real models, the checkpoint format, and the export paths that get a trained model onto a server. The two diagrams below are generated directly from the source tree.

Overview of NeMo on main: the core class contract (Model, ModelPT, NeuralType, SaveRestoreConnector, Exportable) anchors three consumer sides — the ASR and TTS collections, the SpeechLM2 collection with its shared tokenizers, and the Lightning/exp_manager wiring that turns configs into training runs.

Reading the overview from left to right:

• Model, L983 — the abstract hub every NeMo model inherits: Typing, Serialization, FileIO and HuggingFaceFileIO mixins, plus the class-method from_pretrained registry at L1015.
• ModelPT, L53 — the LightningModule base for classic NeMo models, with config-driven __init__, dataset and optimizer setup, and .nemo save/restore built in.
• NeuralType, L32 — the semantic dtype-and-axes type system that checks ports between modules before you ever launch a run.
• SaveRestoreConnector, L43 — the pluggable connector that packs model YAML, weights and artifacts into a single .nemo archive.
• Exportable, L40 — the mixin that gives models declared input/output names and an export() entry point for deployment formats.
• ASRModel, L34 — abstract speech-recognition base with an auxiliary-loss registry and CUDA-graph toggles, parented by EncDecCTCModelBPE, L40 and EncDecRNNTBPEModel, L39.
• SALM, L57 — the Speech-Aware Language Model: a LightningModule + HF-hub mixin that interleaves audio embeddings with text tokens.
• CallbackGroup, L25 — a singleton that fans NeMo lifecycle hooks out to registered callbacks.
• exp_manager, L478 — the one-call setup function that wires loggers, checkpointing and resume logic around a Lightning Trainer.

Core Classes: the Model Contract

The whole framework hangs off four abstract classes in nemo/core/classes. Typing lets a module declare input_types and output_types as dictionaries of NeuralType objects — each carrying an element kind (spectrogram, embeddings, labels…) and axes kind (batch, time, channel…). When you connect a Conformer encoder to a decoder, NeMo compares these types and refuses to wire incompatible ports; the comparison logic lives in nemo/core/neural_types/comparison.py. This is the single most “NeMo” idea in the codebase: rather than trusting shapes alone, modules carry meaning about what a tensor is.

Serialization and FileIO define save_to/restore_from, and Model composes them. ModelPT then merges that contract with PyTorch Lightning: its constructor takes an OmegaConf cfg plus an optional Trainer, and __init_subclass__ enforces bookkeeping on every subclass. The data-facing methods — setup_training_data, L544, setup_validation_data, and setup_optimization, L628 — are driven entirely by sub-configs of the model YAML, so the same class trains on different datasets without code changes. Optimizer construction funnels through configure_optimizers, L859, which combines NeMo’s core/optim schedulers with standard PyTorch optimizers.

Below the model level sits NeuralModule, L70, a torch.nn.Module with Typing and Serialization mixed in. Encoders, decoders, augmentations and data-processing blocks all inherit from it. A good example is ConformerEncoder, L65, which declares typed audio input/output ports, supports streaming masks, and can be checkpointed or exported as a unit.

Collections: ASR, TTS, SpeechLM2

The nemo/collections tree is where users actually live. The ASR collection separates models (CTC, RNN-T, hybrid, diarization, classification) from parts (submodules like conformer blocks and decoding strategies) from modules (encoders and decoders). The hybrid family — EncDecHybridRNNTCTCModel, L38 — stacks an RNN-T decoder and a CTC decoder over one encoder, and its transcribe() method at L96 is the entry point most users meet first. Decoding is pluggable: greedy and beam strategies for both CTC and RNN-T live in asr/parts/submodules, including batched beam search with CUDA-graph acceleration.

The TTS collection is organized around two abstract bases: text-to-spectrogram models like FastPitchModel, L82, and vocoders like HifiGanModel, L45, both descending from the SpectrogramGenerator/Vocoder pair declared in base.py, L64. Text normalization and grapheme-to-phoneme conversion feed them from the g2p and torch subpackages. The shared tokenizer stack — SentencePieceTokenizer, L125 plus helpers like get_megatron_tokenizer, L187 — lives in nemo/collections/common/tokenizers and is reused by every collection.

The newest collection, speechlm2, reflects a different design point: models there are LightningModules paired with a Hugging Face-style hub mixin rather than ModelPT children. SALM, L57 combines a text LLM with an audio encoder via AudioPerceptionModule, L45, and DuplexS2SModel extends the pattern to full-duplex speech-to-speech. There is even a vLLM plugin model, NeMoSpeechLMForConditionalGeneration, L116, that registers SALM directly into the vLLM multimodal registry for serving.

The detailed view: inheritance inside the core (Model → ModelPT; NeuralModule → typed ports), the ASR/TTS/SpeechLM2 model families, the callback-and-checkpoint wiring, and the utils that hold save/restore and export helpers together.

Checkpoints, Artifacts, and the .nemo Format

NeMo’s portability story rests on one file format. ModelPT.save_to at L392 writes a .nemo archive — a tarball containing model_config.yaml, a Lightning checkpoint, and any registered artifacts — while restore_from, L434 unpacks it, rebuilds the model from its YAML, and loads weights. The actual packing is delegated to SaveRestoreConnector, L43, which is deliberately pluggable: sharded, model-parallel, or storage-backend-specific variants can subclass it.

Artifacts deserve special mention. Models can call register_artifact, L234 to bind an external file (a tokenizer model, a language-model for shallow fusion, a phoneme dictionary) into the checkpoint, so the archive is self-contained. During restore, class references inside the YAML are resolved by import_class_by_path at L566, with an allowlist of safe import prefixes guarding against arbitrary code execution from untrusted checkpoints — a security detail that a lot of serialization frameworks skip.

The registry side of the story is Model.from_pretrained: given a name with a / it resolves through Hugging Face Hub support in HuggingFaceFileIO, L26; otherwise it matches against each class’s list_available_models() entries and downloads from NGC into a cache directory. This dual-source registry is why the same from_pretrained call works for both stt_en_* NGC names and nvidia/* Hub names. Distributed users get extra machinery: AppState, L34 tracks model and data-parallel ranks so save_to refuses plain archives in model-parallel mode, and checkpoint file names get rank-injected to keep sharded saves consistent.

Training Wiring Around the Trainer

NeMo deliberately does not wrap the Lightning Trainer; it decorates it. exp_manager takes the trainer and an ExpManagerConfig, creates the log directory, configures loggers (TensorBoard, W&B, MLflow, DLLogger), installs checkpointing via NeMoModelCheckpoint, L39, and handles resume logic that can find the last checkpoint in a prior experiment. On the callback side, NeMo-specific hooks inherit BaseCallback, L19 and are aggregated by the CallbackGroup singleton, which ModelPT.__init__ pokes with on_model_init_start the moment a model begins construction. This split — Lightning owns the loop, NeMo owns domain hooks — is what lets the same model class train on one GPU or scale out behind a Megatron-style strategy without code changes.

A subtle but useful piece: ModelPT.maybe_init_from_pretrained_checkpoint at L1263 lets a fresh model initialize parts of itself (usually just the encoder) from a pretrained .nemo file as declared in its config — the standard transfer-learning path for ASR fine-tuning. Combined with exp_manager’s resume support, warm-starting is declarative rather than scripted.

From Install to Your First Model

Installation follows the standard PyPI path, with extras controlling the (heavy) dependency tree:

pip install nemo_toolkit[all]

Then transcription is three lines, using the exact registry path described above:

import nemo.collections.asr as nemo_asr

model = nemo_asr.models.ASRModel.from_pretrained(model_name="nvidia/parakeet-tdt-0.6b-v2")
model.transcribe(["meeting.wav"])

ASRModel.from_pretrained resolves the Parakeet hybrid RNN-T/CTC model, downloads the .nemo archive into the cache, unpacks it, rebuilds the model from its embedded YAML and loads weights — all through the Model/SaveRestoreConnector machinery we traced. For training rather than inference, you hand a Trainer to the constructor and let exp_manager set up the rest; the model’s YAML carries train_ds, validation_ds and optim sections that setup_training_data and setup_optimization consume. Saving a fine-tuned model back out is one call to save_to("my_parakeet.nemo"), and if you want to share it, the HuggingFaceFileIO mixin provides push_to_hf_hub for the Hub counterpart of the registry you downloaded from.

Export and Serving Paths

Deployment is a first-class concern rather than an afterthought. The Exportable mixin lets each model declare its export input/output names, supported formats and any disabled deployment ports; the base export() implementation at L61 drives conversion to ONNX and TorchScript, with the ExportFormat enum and verification helpers (numerical comparison against the eager model, ONNX Runtime session checks) in export_utils.py, L95. Module-swap utilities replace training-only kernels with export-friendly equivalents before tracing. The exported ONNX can then feed standard runtimes — our ONNX and ONNX Runtime tours cover that side of the fence.

Serving has two flavors in today’s tree. Classic NeMo served RNN-T/CTC models as RNNT/CTC services over gRPC, and NVIDIA’s historical FasterTransformer/TRT-LLM bridges carried NeMo checkpoints into high-throughput LLM servers — see our FasterTransformer and Triton Inference Server tours for that trajectory. The newer path is native: the SpeechLM2 collection ships its own vLLM plugin so a SALM checkpoint can be loaded by vLLM directly through the registered wrapper class, no manual conversion step. For Hugging Face-centric workflows, push_to_hf_hub plus transformers-compatible tokenizer exports round out the interop story.

Honest Limits

The dependency tree is the first thing to respect: nemo_toolkit[all] pulls Lightning, OmegaConf, Hydra, k2, lhotse, sentencepiece, text-normalization toolkits and more, and version drift between those pins has historically been the top source of install pain. The framework’s coupling to PyTorch Lightning is similarly double-edged — Lightning gives you distributed training and checkpointing for free, but debugging means reading Lightning’s callback order, and NeMo adds its own callback layer on top.

Second, the codebase has churned hard. The jump from NeMo 1.x to 2.x replaced the nemo/collections/nlp stack, and since then the LLM-training collection itself moved out of this repository — so older tutorials referencing nemo/collections/llm or nemo/collections/nlp paths will not match this tree. ModelPT-based speech collections keep the 1.x-style YAML-driven flow, while speechlm2 models use a newer HF-hub style; both idioms coexist, which is honest engineering mid-migration but confusing at first contact.

Third, assumptions skew GPU-centric. CUDA-graph decoding, TE/megatron-adjacent optimizers, and Sharded-model defaults all presume an NVIDIA stack; CPU-only inference works but is not the optimized path. Finally, the .nemo tarball format, while self-contained, embeds class-import paths in its YAML — restoring an archive across NeMo versions can fail when class paths move, so pin your version per checkpoint and keep restore_from’s strict=False in your back pocket for legacy files.

Enjoyed this post? Never miss out on future posts by following us