Source-linked AI summary
NeMo: a toolkit for building AI applications using Neural Modules
Oleksii Kuchaiev, Jason Li, Huyen Nguyen, Oleksii Hrinchuk, Ryan Leary, Boris Ginsburg, Samuel Kriman, Stanislav Beliaev, Vitaly Lavrukhin, Jack Cook, Patrice Castonguay, Mariya Popova, Jocelyn Huang, Jonathan M. Cohen
TL;DR
Deep-learning development is cumbersome because tensor abstractions often lack semantics and configurations mix concerns, complicating reuse and debugging. NeMo addresses these issues with typed neural modules, static checks, reusable collections, and framework-agnostic execution, while reporting 29.2 BLEU / 28.5 SacreBLEU on WMT16 English-German after about 15 hours on 8 GPUs. Its neural type system and useful abstraction levels remain ongoing research directions.
Problem
Deep-learning systems are difficult to reuse and develop because tensors lack dimensional semantics and configurations blur architecture, training, and other concerns.
Method
NeMo uses typed neural modules, static checking during computation-graph construction, reusable domain collections, and a framework-agnostic execution layer.
Results
29.2 BLEU / 28.5 SacreBLEU was achieved on newstest2014 after about 15 hours of WMT16 English-German training on one machine with 8 GPUs.
Takeaways & Limitations
NeMo makes building and re-using deep neural networks easier through higher-level abstraction, semantic checks, and pre-built ASR and NLP modules.
Takeaways & Limitations
The design of NeMo’s neural type system and the most useful abstraction levels for modules remain ongoing research directions.
Abstract
from arXiv · showhide
NeMo (Neural Modules) is a Python framework-agnostic toolkit for creating AI applications through re-usability, abstraction, and composition. NeMo is built around neural modules, conceptual blocks of neural networks that take typed inputs and produce typed outputs. Such modules typically represent data layers, encoders, decoders, language models, loss functions, or methods of combining activations. NeMo makes it easy to combine and re-use these building blocks while providing a level of semantic correctness checking via its neural type system. The toolkit comes with extendable collections of pre-built modules for automatic speech recognition and natural language processing. Furthermore, NeMo provides built-in support for distributed training and mixed precision on latest NVIDIA GPUs. NeMo is open-source https://github.com/NVIDIA/NeMo
1 Introduction
NeMo applies software-engineering practices to deep-learning development by decomposing applications into reusable, typed neural modules. It separates model concerns while supporting efficient hardware features and domain-specific module collections.
- Deep-learning tensors often lack dimensional semantics and a type system, complicating component reuse and construction of systems from models by different developers.
- Model configurations commonly mix architecture, training, optimization, visualization, and analysis in difficult-to-reuse Python scripts.
- These design challenges correspond to software-engineering problems involving separation of concerns, verifiable interfaces, and code reusability.
- NeMo targets functional decomposition, static type checking, separation of model concerns, and high-performance training as core development practices.
- NeMo provides reusable pre-built components that can be combined in novel ways through its neural module abstraction and collections.
- NeMo Core supplies fundamental building blocks and a neural type system, while collections provide domain-specific modules for automatic speech recognition and natural language processing.
2 Related work
NeMo combines flexible, component-based model construction with reusable collections of common modules. Its distinguishing features are neural modules as the core abstraction and a generic neural type system for semantic compatibility checks.
- Toolkit landscape: Higher-level AI toolkits include neural-network APIs and configuration-driven systems designed to simplify work beyond directly using deep-learning frameworks.Examples include Keras and PyTorch Lightning in the first group, and Tensor2Tensor and OpenNMT in the second.
- NeMo’s design: NeMo Core lets users compose arbitrary sets of components while hiding training and evaluation-loop details.NeMo collections provide reusable modules and templates for fixed patterns such as encoder-decoder networks.
- NeMo’s distinctions: NeMo uses neural modules rather than layers or tensors as its core abstraction and adds a neural type system for semantic checks.The type system supports compatibility checking between module inputs and outputs.
- NeMo’s distinctions: Unlike many configuration-driven toolkits, NeMo imposes neither a particular model structure nor a required configuration-file format.Users can define models directly through the NeMo API rather than following a mandated encoder-decoder-loss structure.
- Relationship to related systems: NeMo resembles PyTorchPipe in its task-oriented flexibility and compatibility checks, but exposes its type and composition system without requiring NeMo runtime functionality.This makes NeMo an application framework whose underlying abstractions can also be used independently.
3 NeMo
NeMo structures AI applications from typed Neural Modules connected through data-defined activation flows, with lazy execution and semantic checks before computation. Its abstractions support reuse, framework independence, callbacks, hardware acceleration, and distributed training.
- Neural Modules: Neural Modules are reusable components that compute typed outputs from typed inputs, representing logical parts such as encoders, decoders, losses, or data augmentation.They can be composed at different granularities rather than being limited to individual neural-network layers.
- Framework abstraction: NeMo’s Python-facing abstractions avoid references to a specific backend, although the current implementation supports PyTorch.Existing PyTorch modules can be converted into Neural Modules by adding input and output port definitions.
- Composition: NeMo represents neural-network structure and forward/backward data flows in data structures rather than Python control-flow logic.Explicitly linked ports define activation dependencies, while parameterized modules define variation through constructor-time values.
- Execution model: Applications instantiate modules, connect them into a directed acyclic graph, optionally define callbacks, and invoke actions such as train, evaluate, or infer.NeMo uses lazy execution: computation begins only after an action is called.
- Neural type system: NeMo performs neural type checking during DAG construction, using axis semantics, dimensions, and inheritance-based tags to detect incompatible connections.The system can identify semantic transpositions and dimension mismatches before computation, including batch/time axis swaps that might otherwise cause silent convergence failures.
- Hardware and training support: NeMo supports mixed-precision training, gradient accumulation, and multi-GPU or multi-node training using NVIDIA hardware and APEX.Mixed precision uses float16 for intensive operations, retains float32 where needed, and employs dynamic loss scaling.
4 NeMo collections
NeMo collections package domain-specific Neural Modules for speech recognition and natural language processing. Examples show that modules from existing collections can be recombined into new models, including attention-based ASR and Transformer translation, while retaining framework efficiency.
- Collections: NeMo collections package reusable Neural Modules and helper routines for domains such as automatic speech recognition and natural language processing.Collections are Python modules that define Neural Module classes, neural types, and associated utilities; users can add new collections.
- Automatic Speech Recognition: The ASR collection supports both CTC-based and sequence-to-sequence attention-based model types.A Jasper-like ASR example instantiates data, encoder, decoder, and loss modules, then connects them through an activation-flow DAG.
- Automatic Speech Recognition: A different attention-based ASR model reuses the same data layer and Jasper encoder while replacing the CTC decoder and loss with DecoderRNN and SequenceLoss.A JasperRNNConnector resolves the dimensionality mismatch, and DecoderRNN could be pretrained independently as a language model.
- Natural Language Processing: The NLP collection supports neural machine translation, language modeling, sentence classification, ASR correction, intent classification, slot filling, and BERT pretraining or finetuning.The paper illustrates Transformer-BIG by instantiating translation data, encoder, decoder, log-softmax, and smoothed cross-entropy-loss modules.
- Natural Language Processing: 29.2 BLEU / 28.5 SacreBLEU was achieved on newstest2014 after about 15 hours of training on WMT16 English-German with one machine and 8 GPUs.The result is reported as retaining the underlying framework’s efficiency.
5 Conclusions and future work
NeMo transfers software-engineering best practices to deep-learning application development through neural modules, semantic type checks, and reusable domain collections. Its collections and neural type system remain active areas of expansion and design research.
- NeMo transfers software-engineering best practices to deep-learning application development.
- NeMo uses neural modules as a higher-level abstraction and applies a neural type system for semantic checks.
- NeMo provides pre-built conversational-AI modules through the nemo_asr and nemo_nlp collections.
- Expanding collections and determining useful neural-type-system designs and module abstraction levels remain ongoing research directions.