Skip to content

Allow env vars to specify custom extensions to build #1910

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jul 22, 2025
Merged

Allow env vars to specify custom extensions to build #1910

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7d80b19
Env vars to specify extensions to build.
crcrpar Jun 10, 2025
9fd0b7a
create FLAG_TO_ENV from ENV_TO_FKAG
crcrpar Jun 12, 2025
46bcb8c
verbose name: `APEX_PARALLEL` -> `APEX_PARALLEL_BUILD`
crcrpar Jul 12, 2025
a6548e9
update README.md about install
crcrpar Jul 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 69 additions & 37 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,24 +26,50 @@ The latest stable release obtainable from https://pytorch.org should also work.
We recommend installing [`Ninja`](https://ninja-build.org/) to make compilation faster.

### Linux
For performance and full functionality, we recommend installing Apex with
CUDA and C++ extensions via

For performance and full functionality, we recommend installing Apex with CUDA and C++ extensions using environment variables:

#### Using Environment Variables (Recommended)

```bash
git clone https://github.com/NVIDIA/apex
cd 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" ./
# otherwise
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./
# Build with core extensions (cpp and cuda)
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .

# To build with additional extensions, specify them with environment variables
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 APEX_FAST_MULTIHEAD_ATTN=1 APEX_FUSED_CONV_BIAS_RELU=1 pip install -v --no-build-isolation .

# To build all contrib extensions at once
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 APEX_ALL_CONTRIB_EXT=1 pip install -v --no-build-isolation .
```

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

```bash
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" ./
NVCC_APPEND_FLAGS="--threads 4" APEX_PARALLEL_BUILD=8 APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .
```

When CPU cores or memory are limited, the `--parallel` option is generally preferred over `--threads`. See [pull#1882](https://github.com/NVIDIA/apex/pull/1882) for more details.

APEX also supports a Python-only build via
#### Using Command-Line Flags (Legacy Method)

The traditional command-line flags are still supported:

```bash
# Using pip config-settings (pip >= 23.1)
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" ./

# For older pip versions
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./

# To build with additional extensions
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_multihead_attn" ./
```

#### Python-Only Build

APEX also supports a Python-only build via:
```bash
pip install -v --disable-pip-version-check --no-build-isolation --no-cache-dir ./
```
Expand All @@ -65,31 +91,37 @@ If you installed Pytorch in a Conda environment, make sure to install Apex in th

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

| Module Name | Install Option | Misc |
|---------------|------------------|--------|
| `apex_C` | `--cpp_ext` | |
| `amp_C` | `--cuda_ext` | |
| `syncbn` | `--cuda_ext` | |
| `fused_layer_norm_cuda` | `--cuda_ext` | [`apex.normalization`](./apex/normalization) |
| `mlp_cuda` | `--cuda_ext` | |
| `scaled_upper_triang_masked_softmax_cuda` | `--cuda_ext` | |
| `generic_scaled_masked_softmax_cuda` | `--cuda_ext` | |
| `scaled_masked_softmax_cuda` | `--cuda_ext` | |
| `fused_weight_gradient_mlp_cuda` | `--cuda_ext` | Requires CUDA>=11 |
| `permutation_search_cuda` | `--permutation_search` | [`apex.contrib.sparsity`](./apex/contrib/sparsity) |
| `bnp` | `--bnp` | [`apex.contrib.groupbn`](./apex/contrib/groupbn) |
| `xentropy` | `--xentropy` | [`apex.contrib.xentropy`](./apex/contrib/xentropy) |
| `focal_loss_cuda` | `--focal_loss` | [`apex.contrib.focal_loss`](./apex/contrib/focal_loss) |
| `fused_index_mul_2d` | `--index_mul_2d` | [`apex.contrib.index_mul_2d`](./apex/contrib/index_mul_2d) |
| `fused_adam_cuda` | `--deprecated_fused_adam` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |
| `fused_lamb_cuda` | `--deprecated_fused_lamb` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |
| `fast_layer_norm` | `--fast_layer_norm` | [`apex.contrib.layer_norm`](./apex/contrib/layer_norm). different from `fused_layer_norm` |
| `fmhalib` | `--fmha` | [`apex.contrib.fmha`](./apex/contrib/fmha) |
| `fast_multihead_attn` | `--fast_multihead_attn` | [`apex.contrib.multihead_attn`](./apex/contrib/multihead_attn) |
| `transducer_joint_cuda` | `--transducer` | [`apex.contrib.transducer`](./apex/contrib/transducer) |
| `transducer_loss_cuda` | `--transducer` | [`apex.contrib.transducer`](./apex/contrib/transducer) |
| `cudnn_gbn_lib` | `--cudnn_gbn` | Requires cuDNN>=8.5, [`apex.contrib.cudnn_gbn`](./apex/contrib/cudnn_gbn) |
| `peer_memory_cuda` | `--peer_memory` | [`apex.contrib.peer_memory`](./apex/contrib/peer_memory) |
| `nccl_p2p_cuda` | `--nccl_p2p` | Requires NCCL >= 2.10, [`apex.contrib.nccl_p2p`](./apex/contrib/nccl_p2p) |
| `fast_bottleneck` | `--fast_bottleneck` | Requires `peer_memory_cuda` and `nccl_p2p_cuda`, [`apex.contrib.bottleneck`](./apex/contrib/bottleneck) |
| `fused_conv_bias_relu` | `--fused_conv_bias_relu` | Requires cuDNN>=8.4, [`apex.contrib.conv_bias_relu`](./apex/contrib/conv_bias_relu) |
| Module Name | Environment Variable | Install Option | Misc |
|---------------|------------------------|------------------|--------|
| `apex_C` | `APEX_CPP_EXT=1` | `--cpp_ext` | |
| `amp_C` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |
| `syncbn` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |
| `fused_layer_norm_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | [`apex.normalization`](./apex/normalization) |
| `mlp_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |
| `scaled_upper_triang_masked_softmax_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |
| `generic_scaled_masked_softmax_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |
| `scaled_masked_softmax_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | |
| `fused_weight_gradient_mlp_cuda` | `APEX_CUDA_EXT=1` | `--cuda_ext` | Requires CUDA>=11 |
| `permutation_search_cuda` | `APEX_PERMUTATION_SEARCH=1` | `--permutation_search` | [`apex.contrib.sparsity`](./apex/contrib/sparsity) |
| `bnp` | `APEX_BNP=1` | `--bnp` | [`apex.contrib.groupbn`](./apex/contrib/groupbn) |
| `xentropy` | `APEX_XENTROPY=1` | `--xentropy` | [`apex.contrib.xentropy`](./apex/contrib/xentropy) |
| `focal_loss_cuda` | `APEX_FOCAL_LOSS=1` | `--focal_loss` | [`apex.contrib.focal_loss`](./apex/contrib/focal_loss) |
| `fused_index_mul_2d` | `APEX_INDEX_MUL_2D=1` | `--index_mul_2d` | [`apex.contrib.index_mul_2d`](./apex/contrib/index_mul_2d) |
| `fused_adam_cuda` | `APEX_DEPRECATED_FUSED_ADAM=1` | `--deprecated_fused_adam` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |
| `fused_lamb_cuda` | `APEX_DEPRECATED_FUSED_LAMB=1` | `--deprecated_fused_lamb` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |
| `fast_layer_norm` | `APEX_FAST_LAYER_NORM=1` | `--fast_layer_norm` | [`apex.contrib.layer_norm`](./apex/contrib/layer_norm). different from `fused_layer_norm` |
| `fmhalib` | `APEX_FMHA=1` | `--fmha` | [`apex.contrib.fmha`](./apex/contrib/fmha) |
| `fast_multihead_attn` | `APEX_FAST_MULTIHEAD_ATTN=1` | `--fast_multihead_attn` | [`apex.contrib.multihead_attn`](./apex/contrib/multihead_attn) |
| `transducer_joint_cuda` | `APEX_TRANSDUCER=1` | `--transducer` | [`apex.contrib.transducer`](./apex/contrib/transducer) |
| `transducer_loss_cuda` | `APEX_TRANSDUCER=1` | `--transducer` | [`apex.contrib.transducer`](./apex/contrib/transducer) |
| `cudnn_gbn_lib` | `APEX_CUDNN_GBN=1` | `--cudnn_gbn` | Requires cuDNN>=8.5, [`apex.contrib.cudnn_gbn`](./apex/contrib/cudnn_gbn) |
| `peer_memory_cuda` | `APEX_PEER_MEMORY=1` | `--peer_memory` | [`apex.contrib.peer_memory`](./apex/contrib/peer_memory) |
| `nccl_p2p_cuda` | `APEX_NCCL_P2P=1` | `--nccl_p2p` | Requires NCCL >= 2.10, [`apex.contrib.nccl_p2p`](./apex/contrib/nccl_p2p) |
| `fast_bottleneck` | `APEX_FAST_BOTTLENECK=1` | `--fast_bottleneck` | Requires `peer_memory_cuda` and `nccl_p2p_cuda`, [`apex.contrib.bottleneck`](./apex/contrib/bottleneck) |
| `fused_conv_bias_relu` | `APEX_FUSED_CONV_BIAS_RELU=1` | `--fused_conv_bias_relu` | Requires cuDNN>=8.4, [`apex.contrib.conv_bias_relu`](./apex/contrib/conv_bias_relu) |
| `distributed_adam_cuda` | `APEX_DISTRIBUTED_ADAM=1` | `--distributed_adam` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |
| `distributed_lamb_cuda` | `APEX_DISTRIBUTED_LAMB=1` | `--distributed_lamb` | [`apex.contrib.optimizers`](./apex/contrib/optimizers) |
| `_apex_nccl_allocator` | `APEX_NCCL_ALLOCATOR=1` | `--nccl_allocator` | Requires NCCL >= 2.19, [`apex.contrib.nccl_allocator`](./apex/contrib/nccl_allocator) |
| `_apex_gpu_direct_storage` | `APEX_GPU_DIRECT_STORAGE=1` | `--gpu_direct_storage` | [`apex.contrib.gpu_direct_storage`](./apex/contrib/gpu_direct_storage) |

You can also build all contrib extensions at once by setting `APEX_ALL_CONTRIB_EXT=1`.
Loading