NeMo Framework has been updated with state-of-the-art features, such as FSDP, Mixture-of-Experts, and RLHF with TensorRT-LLM to provide speedups up to 4.2x for Llama-2 pre-training on H200. All of these features will be available in an upcoming release.
NVIDIA NeMo Framework is a generative AI framework built for researchers and pytorch developers working on large language models (LLMs), multimodal models (MM), automatic speech recognition (ASR), and text-to-speech synthesis (TTS). The primary objective of NeMo is to provide a scalable framework for researchers and developers from industry and academia to more easily implement and design new generative AI models by being able to leverage existing code and pretrained models.
For technical documentation, please see the NeMo Framework User Guide.
All NeMo models are trained with Lightning and training is automatically scalable to 1000s of GPUs.
When applicable, NeMo models take advantage of the latest possible distributed training techniques, including parallelism strategies such as
- data parallelism
- tensor parallelism
- pipeline model parallelism
- fully sharded data parallelism (FSDP)
- sequence parallelism
- context parallelism
- mixture-of-experts (MoE)
and mixed precision training recipes with bfloat16 and FP8 training.
NeMo's Transformer based LLM and Multimodal models leverage NVIDIA Transformer Engine for FP8 training on NVIDIA Hopper GPUs and leverages NVIDIA Megatron Core for scaling transformer model training.
NeMo LLMs can be aligned with state of the art methods such as SteerLM, DPO and Reinforcement Learning from Human Feedback (RLHF), see NVIDIA NeMo Aligner for more details.
NeMo LLM and Multimodal models can be deployed and optimized with NVIDIA Inference Microservices (Early Access).
NeMo ASR and TTS models can be optimized for inference and deployed for production use-cases with NVIDIA Riva.
For scaling NeMo LLM and Multimodal training on Slurm clusters or public clouds, please see the NVIDIA Framework Launcher. The NeMo Framework launcher has extensive recipes, scripts, utilities, and documentation for training NeMo LLMs and Multimodal models and also has an Autoconfigurator which can be used to find the optimal model parallel configuration for training on a specific cluster. To get started quickly with the NeMo Framework Launcher, please see the NeMo Framework Playbooks The NeMo Framework Launcher does not currently support ASR and TTS training but will soon.
Getting started with NeMo is simple. State of the Art pretrained NeMo models are freely available on HuggingFace Hub and NVIDIA NGC. These models can be used to generate text or images, transcribe audio, and synthesize speech in just a few lines of code.
We have extensive tutorials that can be run on Google Colab or with our NGC NeMo Framework Container. and we have playbooks for users that want to train NeMo models with the NeMo Framework Launcher.
For advanced users that want to train NeMo models from scratch or finetune existing NeMo models we have a full suite of example scripts that support multi-GPU/multi-node training.
- Python 3.10 or above
- Pytorch 1.13.1 or above
- NVIDIA GPU, if you intend to do model training
| Version | Status | Description |
|---|---|---|
| Latest | Documentation of the latest (i.e. main) branch. | |
| Stable | Documentation of the stable (i.e. most recent release) branch. |
FAQ can be found on NeMo's Discussions board. You are welcome to ask questions or start discussions there.
The NeMo Framework can be installed in a variety of ways, depending on your needs. Depending on the domain, you may find one of the following installation methods more suitable.
- Conda / Pip - Refer to the Conda and Pip sections for installation instructions.
- This is recommended for Automatic Speech Recognition (ASR) and Text-to-Speech (TTS) domains.
- When using a Nvidia PyTorch container as the base, this is the recommended installation method for all domains.
- Docker - Refer to the Docker containers section for installation instructions.
- This is recommended for Large Language Models (LLM), Multimodal and Vision domains.
- NeMo LLM & Multimodal Container - nvcr.io/nvidia/nemo:24.01.01.framework
- NeMo Speech Container - nvcr.io/nvidia/nemo:24.01.speech
We recommend installing NeMo in a fresh Conda environment.
conda create --name nemo python==3.10.12
conda activate nemoInstall PyTorch using their configurator.
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidiaThe command used to install PyTorch may depend on your system. Please use the configurator linked above to find the right command for your system.
Use this installation mode if you want the latest released version.
apt-get update && apt-get install -y libsndfile1 ffmpeg
pip install Cython
pip install nemo_toolkit['all']Depending on the shell used, you may need to use "nemo_toolkit[all]" instead in the above command.
To install only a specific domain of NeMo, use the following commands. Note: It is required to install the above pre-requisites before installing a specific domain of NeMo.
pip install nemo_toolkit['asr']
pip install nemo_toolkit['nlp']
pip install nemo_toolkit['tts']
pip install nemo_toolkit['vision']
pip install nemo_toolkit['multimodal']Use this installation mode if you want the version from a particular GitHub branch (e.g main).
apt-get update && apt-get install -y libsndfile1 ffmpeg
pip install Cython
python -m pip install git+https://github.com/NVIDIA/NeMo.git@{BRANCH}#egg=nemo_toolkit[all]Use this installation mode if you are contributing to NeMo.
apt-get update && apt-get install -y libsndfile1 ffmpeg
git clone https://github.com/NVIDIA/NeMo
cd NeMo
./reinstall.shIf you only want the toolkit without additional conda-based dependencies, you may replace reinstall.sh
with pip install -e . when your PWD is the root of the NeMo repository.
To install NeMo on Mac with Apple M-Series GPU:
- create a new Conda environment
- install PyTorch 2.0 or higher
- run the following code:
# [optional] install mecab using Homebrew, to use sacrebleu for NLP collection
# you can install Homebrew here: https://brew.sh
brew install mecab
# [optional] install pynini using Conda, to use text normalization
conda install -c conda-forge pynini
# install Cython manually
pip install cython
# clone the repo and install in development mode
git clone https://github.com/NVIDIA/NeMo
cd NeMo
pip install 'nemo_toolkit[all]'
# Note that only the ASR toolkit is guaranteed to work on MacBook - so for MacBook use pip install 'nemo_toolkit[asr]'One of the options is using Windows Subsystem for Linux (WSL).
To install WSL:
- In PowerShell, run the following code:
wsl --install
# [note] If you run wsl --install and see the WSL help text, it means WSL is already installed.Learn more about installing WSL at Microsoft's official documentation.
- After Installing your Linux distribution with WSL:
- Option 1: Open the distribution (Ubuntu by default) from the Start menu and follow the instructions.
- Option 2: Launch the Terminal application. Download it from Microsoft's Windows Terminal page if not installed.
Next, follow the instructions for Linux systems, as provided above. For example:
apt-get update && apt-get install -y libsndfile1 ffmpeg
git clone https://github.com/NVIDIA/NeMo
cd NeMo
./reinstall.shNote that RNNT requires numba to be installed from conda.
conda remove numba
pip uninstall numba
conda install -c conda-forge numbaNeMo LLM Domain training requires NVIDIA Apex to be installed. Install it manually if not using the NVIDIA PyTorch container.
To install Apex, run
git clone https://github.com/NVIDIA/apex.git
cd apex
git checkout b496d85fb88a801d8e680872a12822de310951fd
pip install -v --no-build-isolation --disable-pip-version-check --no-cache-dir --config-settings "--build-option=--cpp_ext --cuda_ext --fast_layer_norm --distributed_adam --deprecated_fused_adam" ./It is highly recommended to use the NVIDIA PyTorch or NeMo container if having issues installing Apex or any other dependencies.
While installing Apex, it may raise an error if the CUDA version on your system does not match the CUDA version torch was compiled with. This raise can be avoided by commenting it here: https://github.com/NVIDIA/apex/blob/master/setup.py#L32
cuda-nvprof is needed to install Apex. The version should match the CUDA version that you are using:
conda install -c nvidia cuda-nvprof=11.8packaging is also needed:
pip install packagingWith the latest versions of Apex, the pyproject.toml file in Apex may need to be deleted in order to install locally.
NeMo LLM Domain has been integrated with NVIDIA Transformer Engine Transformer Engine enables FP8 training on NVIDIA Hopper GPUs. Install it manually if not using the NVIDIA PyTorch container.
pip install --upgrade git+https://github.com/NVIDIA/TransformerEngine.git@stableIt is highly recommended to use the NVIDIA PyTorch or NeMo container if having issues installing Transformer Engine or any other dependencies.
Transformer Engine requires PyTorch to be built with CUDA 11.8.
When traning Large Language Models in NeMo, users may opt to use Flash Attention for efficient training. Transformer Engine already supports Flash Attention for GPT models. If you want to use Flash Attention for non-causal models, please install flash-attn. If you want to use Flash Attention with attention bias (introduced from position encoding, e.g. Alibi), please also install triton pinned version following the implementation.
pip install flash-attn
pip install triton==2.0.0.dev20221202To launch the inference web UI server, please install the gradio gradio.
pip install gradio==3.34.0NeMo Text Processing, specifically (Inverse) Text Normalization, is now a separate repository https://github.com/NVIDIA/NeMo-text-processing.
We release NeMo containers alongside NeMo releases. For example, NeMo r1.23.0 comes with container nemo:24.01.speech, you may find more details about released containers in releases page.
To use built container, please run
docker pull nvcr.io/nvidia/nemo:24.01.speechTo build a nemo container with Dockerfile from a branch, please run
DOCKER_BUILDKIT=1 docker build -f Dockerfile -t nemo:latest .If you choose to work with the main branch, we recommend using NVIDIA's PyTorch container version 23.10-py3 and then installing from GitHub.
docker run --gpus all -it --rm -v <nemo_github_folder>:/NeMo --shm-size=8g \
-p 8888:8888 -p 6006:6006 --ulimit memlock=-1 --ulimit \
stack=67108864 --device=/dev/snd nvcr.io/nvidia/pytorch:23.10-py3Many examples can be found under the "Examples" folder.
We welcome community contributions! Please refer to CONTRIBUTING.md for the process.
We provide an ever-growing list of publications that utilize the NeMo framework.
If you would like to add your own article to the list, you are welcome to do so via a pull request to this repository's gh-pages-src branch.
Please refer to the instructions in the README of that branch.
NeMo is released under an Apache 2.0 license.