Repository files navigation

This repository holds NVIDIA-maintained utilities to streamline mixed precision and distributed training in Pytorch.Some of the code here will be included in upstream Pytorch eventually.The intent of Apex is to make up-to-date utilities available to users as quickly as possible.

Eachapex.contrib module requires one or more install options other than--cpp_ext and--cuda_ext.Note that contrib modules do not necessarily support stable PyTorch releases, some of them might only be compatible with nightlies.

NVIDIA PyTorch Containers are available on NGC:https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch.The containers come with all the custom extensions available at the moment.

Seethe NGC documentation for details such as:

• how to pull a container
• how to run a pulled container
• release notes

To install Apex from source, we recommend using the nightly Pytorch obtainable fromhttps://github.com/pytorch/pytorch.

The latest stable release obtainable fromhttps://pytorch.org should also work.

We recommend installingNinja to make compilation faster.

For performance and full functionality, we recommend installing Apex withCUDA and C++ extensions via

git clone https://github.com/NVIDIA/apexcd apex# if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key...pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings"--build-option=--cpp_ext" --config-settings"--build-option=--cuda_ext" ./# otherwisepip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./

To reduce the build time of APEX, parallel building can be enhanced via

NVCC_APPEND_FLAGS="--threads 4" pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings"--build-option=--cpp_ext --cuda_ext --parallel 8" ./

When CPU cores or memory are limited, the--parallel option is generally preferred over--threads. Seepull#1882 for more details.

APEX also supports a Python-only build via

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

A Python-only build omits:

• Fused kernels required to useapex.optimizers.FusedAdam.
• Fused kernels required to useapex.normalization.FusedLayerNorm andapex.normalization.FusedRMSNorm.
• Fused kernels that improve the performance and numerical stability ofapex.parallel.SyncBatchNorm.
• Fused kernels that improve the performance ofapex.parallel.DistributedDataParallel andapex.amp.DistributedDataParallel,amp, andSyncBatchNorm will still be usable, but they may be slower.

pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" . may work if you were able to build Pytorch from sourceon your system. A Python-only build viapip install -v --no-cache-dir . is more likely to work.
If you installed Pytorch in a Conda environment, make sure to install Apex in that same environment.

If a requirement of a module is not met, then it will not be built.