NVIDIA/apex: mixed precision and distributed training utilities for PyTorch

What NVIDIA/apex is meant to do

The README opens by describing the repository as NVIDIA-maintained utilities to streamline mixed precision and distributed training in Pytorch, with the note that some of the code will be included in upstream Pytorch eventually. That sentence is the whole design brief. Apex exists so that utilities can be used before they land in the framework itself.

The audience is narrow. The CUDA extensions target NVIDIA GPUs and require a compatible CUDA toolchain and PyTorch environment; the project also documents a Python-only build with reduced functionality. The repository ships examples for dcgan, imagenet, simple and docker under examples/, which suggests the intended path is copying a known-good training script rather than reading an API reference. There is no homepage listed for the project, so the repository is the documentation.

How the extension build is wired

Apex is not a pure Python package. setup.py imports CppExtension and CUDAExtension from torch.utils.cpp_extension and builds named extension modules from the csrc/ directory. Each module is gated on a requirement, and the README states plainly that if a requirement of a module is not met, then it will not be built. A missing cuDNN or NCCL version does not raise an error; it silently drops the module.

The module table in the README maps each extension to an environment variable and a legacy install flag. apex_C and amp_C are the core pair, controlled by APEX_CPP_EXT and APEX_CUDA_EXT. syncbn, fused_layer_norm_cuda, mlp_cuda and the scaled masked softmax kernels are also built by APEX_CUDA_EXT; the contrib modules, such as permutation_search_cuda, bnp, xentropy, focal_loss_cuda, transducer_joint_cuda, peer_memory_cuda and others, each have their own variable. The setup.py ENV_TO_FLAG dictionary is the authoritative list of what can be turned on, and it includes entries such as APEX_DISTRIBUTED_ADAM and APEX_GROUP_NORM.

Three entries carry extra requirements worth noting before you plan a build: fused_weight_gradient_mlp_cuda requires CUDA 11 or newer, cudnn_gbn_lib requires cuDNN 8.5 or newer, and nccl_p2p_cuda requires NCCL 2.10 or newer.

Installing apex from source and running a first build

The README recommends installing from source against a nightly PyTorch from the pytorch/pytorch repository, though it says the latest stable release should also work. Ninja is recommended to speed up compilation.

Clone the repository and install with the core C++ and CUDA extensions enabled. The environment variable form is the recommended one, and the -v flag keeps the compiler output visible.

git clone https://github.com/NVIDIA/apex
cd apex
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .

To reduce build time, the README points to parallel building. It notes that the --parallel option is generally preferred over --threads in that situation, and links pull#1882 for the reasoning.

NVCC_APPEND_FLAGS="--threads 4" APEX_PARALLEL_BUILD=8 APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .

To add a single contrib module, append its variable. This builds the fused convolution bias and ReLU kernel alongside the core extensions.

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

A Python-only install is the fallback when the toolchain is not available. The README lists exactly what you lose: the fused kernels behind apex.optimizers.FusedAdam, apex.normalization.FusedLayerNorm and FusedRMSNorm, the SyncBatchNorm kernels, and the kernels that speed up DistributedDataParallel and amp. The README states that DistributedDataParallel, amp and SyncBatchNorm remain usable in that build, but may be slower.

pip install -v --disable-pip-version-check --no-build-isolation --no-cache-dir ./

On Windows the README marks support as experimental. It says the CUDA build may work if you were able to build PyTorch from source on the same system, and that a Python-only build is more likely to succeed. If PyTorch lives in a Conda environment, Apex must be installed into that same environment.

The contrib modules are where the compatibility risk sits

The README's installation section carries a warning that is easy to skim past: each apex.contrib module requires one or more install options beyond --cpp_ext and --cuda_ext, and contrib modules do not necessarily support stable PyTorch releases, some of them might only be compatible with nightlies.

That is a real constraint, not boilerplate. It means the stable PyTorch you pinned for reproducibility may fail to build a contrib extension you want. The failure mode is unpleasant because the build can still complete: an unmet requirement drops the module rather than stopping the install, so you discover the problem later as an ImportError in the training script.

The Python-only build is the second trap. It succeeds, imports fine, and then runs slower than a CUDA build because FusedAdam, FusedLayerNorm and the SyncBatchNorm kernels are absent. If your reason for choosing Apex was fused optimizer performance, a Python-only install gives you none of it.

When torch.amp is the better choice

The README itself says some of the code here will be included in upstream Pytorch eventually. The clearest case of that is automatic mixed precision: PyTorch now ships torch.amp, and for a training loop that only needs autocast and a gradient scaler, torch.amp does the job without cloning a repository or compiling CUDA kernels. The difference in approach is that Apex builds its own C++ and CUDA extensions at install time, which ties the package to a specific PyTorch, CUDA, cuDNN and NCCL combination, while torch.amp ships with the framework and inherits its compatibility guarantees.

The same logic applies to distributed data parallel. If you only need DistributedDataParallel, the README says it works in a Python-only Apex build, so an Apex install buys you nothing there unless you also want the fused kernels that the README says improve its performance.

Apex remains the right tool when you specifically want a fused kernel that PyTorch does not provide, or one of the contrib modules such as apex.contrib.sparsity, apex.contrib.transducer or apex.contrib.peer_memory. Those have no direct equivalent in the framework.

Maintenance, licence and upgrade cost

The repository is not archived and the last push was on 2026-10-26, so the project is receiving changes. The repository has no published GitHub releases, which fits the README's framing: Apex is a source-install package rather than a versioned distribution, and the install instructions are written around git clone rather than a pinned release.

Upgrade cost is dominated by the source build. Every PyTorch upgrade means recompiling, and the README's note that contrib modules may only be compatible with nightlies means the compile is the thing that breaks first. The pyproject.toml pins a narrow Ruff rule set (E4, E7, E9, F) with a comment explaining that the pin exists so upgrading Ruff does not silently enable new rule families. That is a linting concern, not a runtime one, but it signals that the project manages its own toolchain drift deliberately.

Licensing is BSD-3-Clause, which permits commercial use and modification with the usual attribution and warranty conditions. Note that the CUDA kernels are built against NVIDIA's toolchain and that NGC containers are offered under NVIDIA's own terms, so a legal review of the container path is separate from the repository licence. This is not legal advice.

Editorial conclusion

Adopt Apex if you are already on NVIDIA hardware and need fused kernels such as FusedAdam, FusedLayerNorm or SyncBatchNorm that the README says a Python-only build leaves out. Do not adopt it for mixed precision alone: torch.amp covers that without a source build. Before committing, verify that your PyTorch version satisfies the torch>=2.6.0 floor in requirements.txt, and build with APEX_CPP_EXT=1 APEX_CUDA_EXT=1 to confirm the CUDA extensions compile against your toolkit.

Frequently asked questions

Official sources

1. Issues
2. License: BSD-3-Clause
3. NVIDIA/apex on GitHub
4. README