Repolode
Submit
NVIDIA

NVIDIA/apex

NVIDIA's PyTorch extension providing utilities for mixed precision and distributed training, with optional fused CUDA and C++ kernels.

View on GitHub
Stars
8,999
Forks
1,532
Open issues
777
Last commit
4 days ago

What is it?

Apex is a set of NVIDIA-maintained utilities that extend PyTorch for mixed precision and distributed training. It is intended to give users access to up-to-date training utilities before they are merged upstream into PyTorch.

The package includes fused kernels for common operations such as Adam, layer normalization, and softmax, along with a contrib namespace of experimental modules covering areas like sparsity, transducer loss, bottleneck blocks, and distributed optimizers. Many of these components are built as optional C++/CUDA extensions that are enabled through environment variables at install time.

A Python-only build is also supported, though it omits the fused kernels and runs DistributedDataParallel, amp, and SyncBatchNorm more slowly. NVIDIA's PyTorch containers on NGC ship with all the extensions preinstalled.

Key features

  • 01Mixed precision training utilities (apex.amp) for PyTorch.
  • 02Distributed training helpers including DistributedDataParallel and SyncBatchNorm.
  • 03Fused optimizers and layer normalization kernels built as CUDA extensions.
  • 04Modular contrib packages that are enabled individually at install time.
  • 05Python-only installation option for environments without a CUDA toolchain.

Use cases

  • Speeding up large model training on NVIDIA GPUs with mixed precision.
  • Running multi-GPU or multi-node PyTorch training with fused distributed utilities.
  • Experimenting with fused optimizers or normalization layers in custom training code.

Quick start

  1. 1

    Clone the repository

    git clone https://github.com/NVIDIA/apex
  2. 2

    Enter the project directory

    cd apex
  3. 3

    Build and install with core C++ and CUDA extensions

    APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .

Things to know

  • Building the CUDA and C++ extensions requires a compatible PyTorch and CUDA setup.
  • Some apex.contrib modules only work with PyTorch nightly builds rather than stable releases.
  • Windows installation is experimental.

In collections

Run and train AI on GPUs

Topics

#pytorch#mixed-precision#distributed-training#cuda#deep-learning#gpu

FAQ

What is apex?
NVIDIA's PyTorch extension providing utilities for mixed precision and distributed training, with optional fused CUDA and C++ kernels. Apex is a set of NVIDIA-maintained utilities that extend PyTorch for mixed precision and distributed training. It is intended to give users access to up-to-date training utilities before they are merged upstream into PyTorch.
Is apex free to use? What license does it have?
apex is open source under the BSD-3-Clause license. Check the license terms before using it in commercial products.
Is apex still maintained?
Actively maintained. The last commit to apex was on Nov 2, 2026.
What can I use apex for?
Speeding up large model training on NVIDIA GPUs with mixed precision. Running multi-GPU or multi-node PyTorch training with fused distributed utilities. Experimenting with fused optimizers or normalization layers in custom training code.
How do I install apex?
Clone the repository: git clone https://github.com/NVIDIA/apex Enter the project directory: cd apex Build and install with core C++ and CUDA extensions: APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .
Where can I get apex?
The source code and installation instructions are on GitHub: https://github.com/NVIDIA/apex

Summary generated by AI from the project's README. GitHub stats refresh daily.

Related projects