NVIDIA Apex

The nccl_allocator module enables the use of ncclMemAlloc within PyTorch to accelerate NCCL NVLS collective communications. It is built upon CUDAPluggableAllocator.

To switch between standard cudaMalloc and ncclMemAlloc, use the nccl_allocator.nccl_mem(enabled=True) context manager. When enabled=True, the allocator uses ncclMemAlloc for memory allocations performed within the context block.

import os
import torch
import torch.distributed as dist
import apex.contrib.nccl_allocator as nccl_allocator

rank = int(os.getenv("RANK"))
local_rank = int(os.getenv("LOCAL_RANK"))
world_size = int(os.getenv("WORLD_SIZE"))

# Initialize the allocator
nccl_allocator.init()

torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")

# Use the context manager to enable ncclMemAlloc
with nccl_allocator.nccl_mem():
	# Allocations inside this block use ncclMemAlloc
	a = torch.ones(1024 * 1024 * 2, device="cuda")

dist.all_reduce(a)

torch.cuda.synchronize()

To install Apex from source on Linux, it is recommended to use a nightly PyTorch build and install Ninja to accelerate compilation. For optimal performance and full functionality, you should build with CUDA and C++ extensions using environment variables.

Core Extensions

To build the core cpp and cuda extensions:

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

Additional Extensions

To include specific additional extensions, set their corresponding environment variables:

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

All Contrib Extensions

To build all available apex.contrib extensions at once:

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

Parallel Building

To reduce build time, you can enable parallel building using NVCC_APPEND_FLAGS and APEX_PARALLEL_BUILD:

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

Note: If CPU cores or memory are limited, using the --parallel option is generally preferred over --threads.

ASP (Automatic SParsity) enables sparse training and inference for PyTorch models. To use it, import ASP from apex.contrib.sparsity and call ASP.prune_trained_model(model, optimizer) before your training loop. This step calculates the sparse mask and applies it to the weights, making the sparse locations fixed for subsequent training or inference.

Standard Workflow for Deployment (Inference Mode)

1. Load a fully trained (dense) network.
2. Prune parameter values in a 2:4 sparse pattern using ASP.prune_trained_model(model, optimizer).
3. Fine-tune the pruned model using the same optimizer, learning rate, and hyperparameters used for the original dense model.
4. (Optional) Quantize the model.

from apex.contrib.sparsity import ASP

# Load your trained model and optimizer
model = define_model(..., pretrained=True)
optimizer = ... 

# Augment model and optimizer for sparse training/inference
ASP.prune_trained_model(model, optimizer)

# Standard training loop to fine-tune the pruned model
x, y = DataLoader(args)
for epoch in range(epochs):
    y_pred = model(x)
    loss = criterion(y_pred, y)
    loss.backward()
    optimizer.step()

torch.save(model.state_dict(), 'pruned_model.pt')