Installation

Prerequisites

  1. Linux x86_64

  2. CUDA 11.8

  3. NVIDIA Driver supporting CUDA 11.8 or later.

  4. cuDNN 8.1 or later.

  5. For FP8/FP16/BF16 fused attention, CUDA 12.1 or later, NVIDIA Driver supporting CUDA 12.1 or later, and cuDNN 8.9.1 or later.

If the CUDA Toolkit headers are not available at runtime in a standard installation path, e.g. within CUDA_HOME, set NVTE_CUDA_INCLUDE_PATH in the environment.

Transformer Engine in NGC Containers

Transformer Engine library is preinstalled in the PyTorch container in versions 22.09 and later on NVIDIA GPU Cloud.

pip - from GitHub

Additional Prerequisites

  1. CMake version 3.18 or later: pip install cmake.

  2. [For pyTorch support] pyTorch with GPU support.

  3. [For JAX support] JAX with GPU support, version >= 0.4.7.

  4. [For TensorFlow support] TensorFlow with GPU support.

  5. pybind11: pip install pybind11.

  6. [Optional] Ninja: pip install ninja.

Installation (stable release)

Execute the following command to install the latest stable version of Transformer Engine:

# Execute one of the following commands
NVTE_FRAMEWORK=pytorch pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable    # Build TE for PyTorch only. The default.
NVTE_FRAMEWORK=jax pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable        # Build TE for JAX only.
NVTE_FRAMEWORK=tensorflow pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable # Build TE for TensorFlow only.
NVTE_FRAMEWORK=all pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable        # Build TE for all supported frameworks.

Installation (development build)

Warning

While the development build of Transformer Engine could contain new features not available in the official build yet, it is not supported and so its usage is not recommended for general use.

Execute the following command to install the latest development build of Transformer Engine:

# Execute one of the following commands
NVTE_FRAMEWORK=pytorch pip install git+https://github.com/NVIDIA/TransformerEngine.git@main    # Build TE for PyTorch only. The default.
NVTE_FRAMEWORK=jax pip install git+https://github.com/NVIDIA/TransformerEngine.git@main        # Build TE for JAX only.
NVTE_FRAMEWORK=tensorflow pip install git+https://github.com/NVIDIA/TransformerEngine.git@main # Build TE for TensorFlow only.
NVTE_FRAMEWORK=all pip install git+https://github.com/NVIDIA/TransformerEngine.git@main        # Build TE for all supported frameworks.

Installation (from source)

Execute the following commands to install Transformer Engine from source:

git clone --recursive https://github.com/NVIDIA/TransformerEngine.git    # Clone the repository/fork and checkout all submodules recursively.
cd TransformerEngine                                                     # Enter TE directory.
git checkout stable                                                      # Checkout the correct branch.
export NVTE_FRAMEWORK=pytorch                                            # Optionally set the framework.
pip install .                                                            # Build and install

For already cloned repos, run the following command in TE directory:

git submodule update --init --recursive                                   # Checkout all submodules recursively.