apex

Apex is an NVIDIA-maintained collection of utilities that aim to make up-to-date mixed precision and distributed training features available to PyTorch users quickly, with some code intended to eventually land upstream in PyTorch. What it provides: - Mixed precision training helpers (apex.amp) for running models in lower precision. - Distributed training utilities, including apex.parallel.DistributedDataParallel and SyncBatchNorm. - Fused kernels for optimizers such as FusedAdam, and normalization layers such as FusedLayerNorm and FusedRMSNorm. - A set of optional contrib modules covering areas like sparsity, group batch norm, xentropy, focal loss, transducer losses, peer memory, NCCL P2P, bottleneck layers, distributed optimizers, and GPU direct storage. Installation notes: - NVIDIA PyTorch containers on NGC ship with the custom extensions already available. - From source, the recommended path uses environment variables such as APEX_CPP_EXT=1 and APEX_CUDA_EXT=1 with pip install --no-build-isolation; legacy command-line flags (--cpp_ext, --cuda_ext) remain supported. - A Python-only build is possible but omits fused kernels, so FusedAdam, FusedLayerNorm/FusedRMSNorm, and the fused paths for SyncBatchNorm, DistributedDataParallel and amp are unavailable or slower. - Individual contrib modules require their own install options; some may only be compatible with PyTorch nightlies rather than stable releases. APEX_ALL_CONTRIB_EXT=1 builds all contrib extensions at once. - Ninja is recommended for faster compilation, and parallel build options are documented for limited CPU or memory environments. - Windows support is described as experimental. This is a library for developers training deep learning models who want mixed precision and distributed training building blocks, rather than an end-user application.