### Build APEX with GPUDirect Storage using setup.py

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/gpu_direct_storage/README.md

Install the APEX library with GPUDirect Storage support using the setup.py script. This is an alternative to the pip installation method.

```bash
python setup.py install --gpu_direct_storage
```

--------------------------------

### Experimental Windows Installation (Python-Only)

Source: https://github.com/nvidia/apex/blob/master/README.md

Install Apex from source on Windows with a Python-only build. This is a more likely method to succeed if building with C++ and CUDA extensions fails.

```bash
pip install -v --no-cache-dir .

```

--------------------------------

### Navigate to Multihead Attention Examples

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/multihead_attn/README.md

Command to change the directory to the multihead attention examples within the Apex contrib directory.

```bash
cd contrib/examples/multihead_attn
```

--------------------------------

### Experimental Windows Installation (Config Settings)

Source: https://github.com/nvidia/apex/blob/master/README.md

Attempt to install Apex from source on Windows with C++ and CUDA extensions using pip's config-settings. This method requires PyTorch to be built from source on the system.

```bash
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" .

```

--------------------------------

### Install Dependencies

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/openfold_triton/README.md

Install the 'einops' library, a required dependency for using this subpackage.

```bash
pip install einops
```

--------------------------------

### Install Apex with Pip Config Settings (Linux)

Source: https://github.com/nvidia/apex/blob/master/README.md

Install Apex from source using pip's config-settings for C++ and CUDA extensions. This is the recommended method for pip versions 23.1 and later.

```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" ./

```

--------------------------------

### Install Apex in Running Container

Source: https://github.com/nvidia/apex/blob/master/examples/docker/README.md

Installs Apex from a mounted volume inside a running container, enabling Cuda extensions.

```shell
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .
```

--------------------------------

### Install Apex with Global Options (Legacy Pip)

Source: https://github.com/nvidia/apex/blob/master/README.md

Install Apex from source using legacy global options for C++ and CUDA extensions. This method is for older pip versions prior to 23.1.

```bash
# 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" ./

```

--------------------------------

### Install Apex with Core Extensions (Linux)

Source: https://github.com/nvidia/apex/blob/master/README.md

Build Apex from source with C++ and CUDA extensions enabled using environment variables. This is the recommended method for full functionality.

```bash
git clone https://github.com/NVIDIA/apex
cd apex
# Build with core extensions (cpp and cuda)
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .

```

--------------------------------

### Install Apex with Additional Global Options (Legacy Pip)

Source: https://github.com/nvidia/apex/blob/master/README.md

Install Apex from source using legacy global options, including C++, CUDA, and fast multihead attention extensions. This is for older pip versions.

```bash
# 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" ./

```

--------------------------------

### Check APEX GPUDirect Storage Installation

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/gpu_direct_storage/README.md

Verify that the APEX library with GPUDirect Storage has been successfully installed by importing the necessary modules.

```python
import torch
import apex.contrib.gpu_direct_storage
```

--------------------------------

### Install Apex with All Contrib Extensions (Linux)

Source: https://github.com/nvidia/apex/blob/master/README.md

Build Apex from source with all available contrib extensions enabled. This ensures the broadest functionality from the contrib modules.

```bash
# 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 .

```

--------------------------------

### Build APEX with GPUDirect Storage using pip

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/gpu_direct_storage/README.md

Install the APEX library with GPUDirect Storage support using pip. Ensure you are in the project root directory.

```bash
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--gpu_direct_storage" ./
```

--------------------------------

### Install Apex with Additional Extensions (Linux)

Source: https://github.com/nvidia/apex/blob/master/README.md

Build Apex from source with core extensions plus specific additional extensions like fast multihead attention and fused convolution bias ReLU. Specify desired extensions using environment variables.

```bash
# 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 .

```

--------------------------------

### Install Apex with Permutation Search CUDA Extension

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/README.md

Install Apex with the permutation search CUDA extension to accelerate the permutation search process on GPUs. This command enables GPU acceleration for permutation search.

```bash
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--permutation_search" ./
```

--------------------------------

### Run Distributed Data Parallel Example

Source: https://github.com/nvidia/apex/blob/master/examples/simple/distributed/README.md

Execute the distributed data parallel training script using bash. This script is intended for instructional purposes.

```bash
bash run.sh
```

--------------------------------

### Install Apex with Parallel Build Options (Linux)

Source: https://github.com/nvidia/apex/blob/master/README.md

Build Apex from source with parallel compilation enabled to reduce build time, especially on systems with limited CPU cores or memory. Adjust thread and parallel build counts as needed.

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

```

--------------------------------

### Initialize and Use NCCL Allocator

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/nccl_allocator/README.md

This snippet demonstrates the basic setup and usage of the nccl_allocator. It initializes the allocator, sets up distributed communication, and uses the nccl_mem context manager to perform an all-reduce operation with allocated memory.

```python
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"))

nccl_allocator.init()

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

with nccl_allocator.nccl_mem():
	a = torch.ones(1024 * 1024 * 2, device="cuda")
dist.all_reduce(a)

torch.cuda.synchronize()

```

--------------------------------

### Build Apex with Fast Multihead Attention Extension

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/multihead_attn/README.md

Instructions to clone the Apex repository and install it with the C++ extension for fast multihead attention enabled. This is typically done on Linux systems.

```bash
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_multihead_attn" ./
```

--------------------------------

### Build Docker Image with Latest Apex

Source: https://github.com/nvidia/apex/blob/master/examples/docker/README.md

Installs the latest Apex on top of an existing image. Uses NVIDIA's Pytorch container as the base image by default.

```docker
docker build -t new_image_with_apex .
```

--------------------------------

### Run ImageNet Training with Different Opt-Levels

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Example commands for training a ResNet50 model on ImageNet using various optimization levels (O0, O1, O2, O3) and configurations. Adjust batch size and workers as needed.

```bash
python main_amp.py -a resnet50 --b 128 --workers 4 --opt-level O0 ./
```

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O3 ./
```

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O3 --keep-batchnorm-fp32 True ./
```

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 ./
```

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 --loss-scale 128.0 ./
```

```bash
python -m torch.distributed.launch --nproc_per_node=2 main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 ./
```

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 ./
```

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 --loss-scale 128.0 ./
```

```bash
python -m torch.distributed.launch --nproc_per_node=2 main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 ./
```

--------------------------------

### Typical PyTorch Training Loop with ASP

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/README.md

An example of integrating ASP's `prune_trained_model` into a standard PyTorch training loop. The sparse mask is fixed after this initialization.

```python
ASP.prune_trained_model(model, optimizer)

x, y = DataLoader(args)
for epoch in range(epochs):
    y_pred = model(x)
    loss = loss_function(y_pred, y)
    loss.backward()
    optimizer.step()

torch.save(...)
```

--------------------------------

### Python-Only Apex Build

Source: https://github.com/nvidia/apex/blob/master/README.md

Install Apex with a Python-only build, omitting fused kernels for optimizers, normalization, SyncBatchNorm, DistributedDataParallel, and AMP. Usable but potentially slower.

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

```

--------------------------------

### Run Apex Container with NVIDIA Docker

Source: https://github.com/nvidia/apex/blob/master/examples/docker/README.md

Launches a container with Apex installed, configured to use NVIDIA GPUs via nvidia-docker.

```docker
docker run --runtime=nvidia -it --rm --ipc=host new_image_with_apex
```

--------------------------------

### Mount Apex Repo into Running Container

Source: https://github.com/nvidia/apex/blob/master/examples/docker/README.md

Mounts a local Apex repository into a running container, allowing for installation within the container.

```docker
docker run --runtime=nvidia -it --rm --ipc=host -v /bare/metal/apex:/apex/in/container <base image>
```

--------------------------------

### Build Docker Image with Custom Base Image

Source: https://github.com/nvidia/apex/blob/master/examples/docker/README.md

Builds a Docker image with Apex, allowing a custom base image with Pytorch and Cuda installed via the BASE_IMAGE build-arg.

```docker
docker build --build-arg BASE_IMAGE=1.3-cuda10.1-cudnn7-devel -t new_image_with_apex .
```

--------------------------------

### Build CUDA Kernels for Permutation Search

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

Compiles the CUDA kernels for permutation search. Ensure CUDA and pybind11 are installed. This command is typically run from the permutation_search_kernels/CUDA_kernels directory.

```bash
pushd ../permutation_search_kernels/CUDA_kernels
vnnc -O3 -shared -Xcompiler -fPIC -Xcompiler -DTORCH_EXTENSION_NAME=permutation_search_cuda -std=c++11 $(python3 -m pybind11 --includes) permutation_search_kernels.cu -o ../permutation_search_cuda$(python3-config --extension-suffix)
popd
```

--------------------------------

### Set Identical Random Seed for Multi-GPU Permutation Search

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/README.md

Ensure identical random seeds across all GPUs for reproducible permutation search results in multi-GPU setups. This involves setting seeds for PyTorch, NumPy, and Python's random module.

```python
import torch
import numpy
import random

torch.manual_seed(identical_seed)
torch.cuda.manual_seed_all(identical_seed)
numpy.random.seed(identical_seed)
random.seed(identical_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
```

--------------------------------

### Import ASP

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/README.md

Import the ASP library to enable sparse training and inference functionalities.

```python
from apex.contrib.sparsity import ASP
```

--------------------------------

### Initialize Models and Optimizers with Amp

Source: https://github.com/nvidia/apex/blob/master/examples/dcgan/README.md

Initialize lists of models and optimizers using amp.initialize for mixed precision training. This replaces manual half-precision conversions.

```python
# Added after models and optimizers construction
[netD, netG], [optimizerD, optimizerG] = amp.initialize(
    [netD, netG], [optimizerD, optimizerG], opt_level=opt.opt_level, num_losses=3)
...
```

--------------------------------

### FP16 Training with Opt-Level O3

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Command to run training using pure FP16 precision. Note that this may not converge and is primarily for establishing speed benchmarks.

```bash
$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O3 ./
```

--------------------------------

### FP32 Training with Opt-Level O0

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Command to run training using pure FP32 precision. This is useful for establishing a baseline.

```bash
$ python main_amp.py -a resnet50 --b 128 --workers 4 --opt-level O0 ./
```

--------------------------------

### Launch Distributed Training with torch.distributed.launch

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Command to launch multiprocess distributed training jobs using the PyTorch launcher utility. NUM_GPUS should be set to the number of available GPUs.

```bash
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS main_amp.py args...
```

--------------------------------

### Run DCGAN with Recommended Mixed Precision Training

Source: https://github.com/nvidia/apex/blob/master/examples/dcgan/README.md

Execute the main_amp.py script with the --opt_level O1 flag for recommended mixed precision training.

```bash
$ python main_amp.py --opt_level O1
```

--------------------------------

### Initialize ASP for Sparse Training

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/README.md

Add this line before the training phase to augment the model and optimizer for sparse training and inference. This function calculates and applies the sparse mask to the model's weights once.

```python
ASP.prune_trained_model(model, optimizer)
```

--------------------------------

### Sketch for Generating a Sparse Network

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/README.md

This sketch demonstrates using ASP for generating a pruned model for accelerated inference, including steps for pruning a trained model and fine-tuning it. Ensure to use the same optimization method and hyperparameters as the original dense model.

```python
model = define_model(..., pretrained=True) # define model architecture and load parameter tensors with trained values (by reading a trained checkpoint)
criterion = ... # compare ground truth with model predition; use the same criterion as used to generate the dense trained model
optimizer = ... # optimize model parameters; use the same optimizer as used to generate the dense trained model
lr_scheduler = ... # learning rate scheduler; use the same schedule as used to generate the dense trained model

from apex.contrib.sparsity import ASP     
ASP.prune_trained_model(model, optimizer) #pruned a trained model

x, y = DataLoader(args)
for epoch in range(epochs): # train the pruned model for the same number of epochs as used to generate the dense trained model
    y_pred = model(x)
    loss = criterion(y_pred, y)
    lr_scheduler.step()
    loss.backward()
    optimizer.step()

torch.save(...) # saves the pruned checkpoint with sparsity masks 
```

--------------------------------

### Launch Permutation Test without GPU Acceleration

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

Launches a permutation test with specified GPU, channels, filters, and search strategies, demonstrating the performance difference without GPU acceleration. Note that the search is slower but finds the same final permutations.

```bash
python3 permutation_test.py --gpu 0 --channels 64 --filters 128 channel_swap,0 channel_swap,100 optimize_stripe_groups,8,0 optimize_stripe_groups,8,100 random,1000
```

--------------------------------

### Run DCGAN with Pure FP32 Training

Source: https://github.com/nvidia/apex/blob/master/examples/dcgan/README.md

Execute the main_amp.py script with the --opt_level O0 flag for pure FP32 training.

```bash
$ python main_amp.py --opt_level O0
```

--------------------------------

### Run Performance Test with Reference (Python Version)

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/multihead_attn/README.md

Executes the performance test script for the fast multihead attention, using the reference (likely Python) implementation.

```python
python perf_test_multihead_attn.py --ref
```

--------------------------------

### Launch Permutation Test with GPU Acceleration

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

Launches a permutation test with specified channels, filters, and multiple search strategies, utilizing GPU acceleration for faster search. This command reports results on efficacy and duration for each strategy.

```bash
python3 permutation_test.py --channels 64 --filters 128 channel_swap,0 channel_swap,100 channel_swap,1000 optimize_stripe_groups,8,0 optimize_stripe_groups,8,100 optimize_stripe_groups,8,1000 optimize_stripe_groups,12,0 random,1000 random,10000 random,100000
```

--------------------------------

### Create Softlinks to ImageNet Datasets

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Before training, create symbolic links to your ImageNet training and validation datasets in the current directory. This allows the training script to access the data.

```bash
$ ln -sf /data/imagenet/train-jpeg/ train
$ ln -sf /data/imagenet/val-jpeg/ val
```

--------------------------------

### Distributed Training with O2 Mixed Precision

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Launch distributed Imagenet training using torch.distributed.launch with 2 processes, each on one GPU, and O2 mixed precision.

```bash
python -m torch.distributed.launch --nproc_per_node=2 main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 ./
```

--------------------------------

### Initialize Torch DistributedDataParallel

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Instantiate Torch's standard DistributedDataParallel wrapper for a model. Requires manual specification of device_ids and output_device.

```python
model = torch.nn.parallel.DistributedDataParallel(model,
                                                  device_ids=[arg.local_rank],
                                                  output_device=arg.local_rank)
```

--------------------------------

### Distributed Training with O1 Mixed Precision

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Launch distributed Imagenet training using torch.distributed.launch with 2 processes, each on one GPU, and O1 mixed precision.

```bash
python -m torch.distributed.launch --nproc_per_node=2 main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 ./
```

--------------------------------

### Permutation Test Script Usage

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

Displays the command-line arguments for the permutation_test.py script. The 'strategy' argument is required.

```bash
python3 permutation_test.py --h
```

--------------------------------

### Initialize Apex AMP with DDP

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Initialize Apex Automatic Mixed Precision (AMP) before wrapping the model with DistributedDataParallel (DDP). This order is crucial to avoid errors.

```python
model, optimizer = amp.initialize(model, optimizer, flags...)
model = DDP(model)
```

--------------------------------

### Compare Performance with PyTorch's Multihead Attention

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/multihead_attn/README.md

Runs the performance test script to compare the fast multihead attention implementation against PyTorch's native multihead attention module.

```python
python perf_test_multihead_attn.py --native
```

--------------------------------

### Initialize FusedAdamSWA Optimizer

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/openfold_triton/README.md

Initializes a FusedAdamSWA optimizer, which combines AdamW, BF16/FP32 casting, and Stochastic Weight Averaging (SWA) for efficient training and evaluation. It requires standard PyTorch optimizers and parameter lists for different precision/usage.

```python
from apex.contrib.openfold_triton.fused_adam_swa import FusedAdamSWA

fused_optimizer = FusedAdamSWA.from_optim(
    adam_optimizer=adam_optimizer,  # standard pytorch optimizer
    fp32_params=fp32_params,  # FP32 used in weight update
    bf16_params=bf16_params,  # BF16 used in forward, backward, reduction
    swa_params=swa_params,  # SWA used for evaluation
    swa_decay_rate=swa_decay_rate,  # for example: 0.9, 0.99, 0.999
)
```

--------------------------------

### Permutation Test Script Arguments

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

Shows the available arguments for the permutation_test.py script, including input file, matrix dimensions, verbosity, and search strategies.

```bash
usage: permutation_test.py [-h] [--infile INFILE] [--channels CHANNELS] [--filters FILTERS] 
                           [--verbosity VERBOSITY] [--seed SEED] [--pretty_print PRETTY_PRINT] 
                           [--unstructured UNSTRUCTURED] [--gpu GPU] [--check_permutation CHECK_PERMUTATION] 
                           [--intermediate_steps INTERMEDIATE_STEPS] [--print_permutation PRINT_PERMUTATION]
                           strategy [strategy ...]
permutation_test.py: error: the following arguments are required: strategy
```

--------------------------------

### Run Imagenet with O2 Mixed Precision and Static Loss Scaling

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Run Imagenet training with O2 mixed precision, overriding dynamic loss scaling to use a static value of 128.0.

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 --loss-scale 128.0 ./
```

--------------------------------

### Run Imagenet with O1 Mixed Precision

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Execute the Imagenet training script with O1 mixed precision enabled. This is the recommended setting for typical use.

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 ./
```

--------------------------------

### Run Imagenet with O1 Mixed Precision and Static Loss Scaling

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Run Imagenet training with O1 mixed precision, overriding the default dynamic loss scaling to use a static value of 128.0.

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 --loss-scale 128.0
```

--------------------------------

### Run Imagenet with O2 Mixed Precision

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Execute the Imagenet training script with O2 mixed precision enabled. This mode is more experimental than O1.

```bash
python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 ./
```

--------------------------------

### FP16 Training with FP32 Batchnorm (Opt-Level O3)

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Command to run FP16 training while keeping batch normalization layers in FP32. This can improve stability and speed by utilizing cuDNN batchnorms.

```bash
$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O3 --keep-batchnorm-fp32 True ./
```

--------------------------------

### Initialize Apex DistributedDataParallel

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Instantiate Apex's DistributedDataParallel wrapper for a model. This is a drop-in replacement for torch.nn.parallel.DistributedDataParallel.

```python
model = apex.parallel.DistributedDataParallel(model)
```

--------------------------------

### Perform FusedAdamSWA Optimizer Step

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/openfold_triton/README.md

Executes a single optimization step using the FusedAdamSWA optimizer. This step includes casting parameters to the appropriate precision (BF16/FP32), updating weights, and applying SWA logic.

```python
fused_optimizer.step()  # fused optimizer step: casting BF16/FP32 + param updates + SWA
```

--------------------------------

### Scale and Backpropagate Loss with Amp

Source: https://github.com/nvidia/apex/blob/master/examples/dcgan/README.md

Use amp.scale_loss to scale individual losses before backpropagation. This is done for each loss with a unique loss_id.

```python
# loss.backward() changed to:
with amp.scale_loss(errD_real, optimizerD, loss_id=0) as errD_real_scaled:
    errD_real_scaled.backward()
...
with amp.scale_loss(errD_fake, optimizerD, loss_id=1) as errD_fake_scaled:
    errD_fake_scaled.backward()
...
with amp.scale_loss(errG, optimizerG, loss_id=2) as errG_scaled:
    errG_scaled.backward()
```

--------------------------------

### Enable Automatic Mixed Precision (Amp) in PyTorch

Source: https://github.com/nvidia/apex/blob/master/examples/imagenet/README.md

Three lines of code are needed to enable Amp. These should be added after model and optimizer construction. The loss backward call is wrapped in `amp.scale_loss`.

```python
model, optimizer = amp.initialize(model, optimizer, flags...)
...
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
```

--------------------------------

### Generate Runtime Results Table

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

Executes a script to generate the runtime results shown in Table 3, including search strategies' efficacies and runtime.

```bash
bash runtime_table.sh
```

--------------------------------

### Citation for Channel Permutations Paper

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

This is the citation for the paper 'Channel Permutations for N:M Sparsity' which describes the underlying ideas and code.

```bibtex
@inproceedings{pool2021channel,
  author    = {Pool, Jeff and Yu, Chong},
  booktitle = {Advances in Neural Information Processing Systems ({NeurIPS})},
  title     = {Channel Permutations for {N:M} Sparsity},
  url       = {https://proceedings.neurips.cc/paper/2021/file/6e8404c3b93a9527c8db241a1846599a-Paper.pdf},
  volume    = {34},
  year      = {2021}
}
```

--------------------------------

### Transform Unstructured to Structured Sparsity

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

The `unstructured_study.sh` script performs a binary search to find the minimum unstructured sparsity required to transform layers into structured sparsity. It requires a directory of .npy weight files and the network name.

```bash
bash unstructured_study.sh <path_to_directory> <network_name>
```

--------------------------------

### Recompute Sparse Masks During Training

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/README.md

Use this method if you need to recompute the sparse mask between training steps, which is useful for advanced techniques like training with sparsity from initialization.

```python
ASP.compute_sparse_masks()
```

--------------------------------

### Generate Ablation Study Results

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

Executes a script to generate results for the ablation study, which demonstrates the relative importance of bounded regressions and the stripe group greedy phase.

```bash
bash ablation_studies.sh
```

--------------------------------

### Run Performance Test with Custom Sequence Length and Batch Size

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/multihead_attn/README.md

Executes the performance test script for multihead attention, allowing customization of sequence length and the range of number of sequences to test.

```python
python perf_test_multihead_attn.py --seq-length 64 --num-seqs-start 10 --num-seqs-stop 120 --num-seqs-inc 5
```

--------------------------------

### Traverse Permutation Space

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/permutation_tests/README.md

Use the `--intermediate_steps` argument to generate a sequence of permutations that evenly divide a given range. This is useful for exploring different sparsity configurations.

```bash
python3 permutation_test.py --channels 64 --filters 128 --intermediate_steps 7 --print_permutation 1 optimize_stripe_groups,8,0
```

--------------------------------

### Run Performance Test for Fast Multihead Attention

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/multihead_attn/README.md

Executes the performance test script for the fast multihead attention, utilizing the C++ implementation.

```python
python perf_test_multihead_attn.py
```

--------------------------------

### Load Triton Auto-Tuned LayerNorm Cache

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/openfold_triton/README.md

Loads pre-computed auto-tuned configurations for Triton LayerNorm kernels based on the input tensor's DAP size and the GPU architecture type (e.g., 'hopper', 'ampere'). This optimizes kernel performance by using cached tuning results.

```python
from apex.contrib.openfold_triton._layer_norm_config_ampere import _auto_tuned_config_ampere
from apex.contrib.openfold_triton._layer_norm_config_hopper import _auto_tuned_config_hopper
from apex.contrib.openfold_triton import _tuneable_triton_kernels

def load_triton_auto_tuned_cache(dap_size: int, arch_type: str) -> None:
    auto_tuned_config = {
        "hopper": _auto_tuned_config_hopper,
        "ampere": _auto_tuned_config_ampere,
    }[arch_type]
    config_for_current_dap = auto_tuned_config[dap_size]
    for func_name, cache in config_for_current_dap.items():
        _tuneable_triton_kernels[func_name].cache = cache

load_triton_auto_tuned_cache(
    dap_size=4,  # supported values: 0, 1, 2, 4, 8
    arch_type="hopper",
)
```

--------------------------------

### Integrate Triton Multi-Head Attention

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/openfold_triton/README.md

Integrates Triton kernels for Multi-Head Attention (MHA) into a PyTorch nn.Module. It dynamically selects between Triton implementations (AttnTri, AttnBiasJIT, AttnNoBiasJIT) or the standard PyTorch implementation based on input shapes, enabled status, and bias presence.

```python
import apex.contrib.openfold_triton.mha as mha
from apex.contrib.openfold_triton import AttnBiasJIT, AttnNoBiasJIT, AttnTri, CanSchTriMHA

# Integration with Attention module:
class SelfAttentionWithGate(nn.Module):
    # ...

    def _attention_forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        mask: torch.Tensor,
        bias: Optional[torch.Tensor],
    ) -> torch.Tensor:
        if self.chunk_size is None:
            if mha.is_enabled() and CanSchTriMHA(
                list(query.shape),
                bias is not None,
                inf=self.inf,
                training=self.training,
            ):
                if mask is not None:
                    mask = mask.contiguous()
                if bias is not None:
                    bias = bias.contiguous()
                return AttnTri(
                    query, key, value, mask, bias, self.inf, torch.is_grad_enabled()
                )
            elif mha.is_enabled() and bias is not None and self.training:
                return AttnBiasJIT(query, key, value, mask, bias, self.inf)
            elif mha.is_enabled() and bias is None and self.training:
                return AttnNoBiasJIT(query, key, value, mask, self.inf)
```

--------------------------------

### Disable Permutation Search in init_model_for_pruning

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/sparsity/README.md

Disable the permutation search process by setting `allow_permutation=False` in the `init_model_for_pruning` function. This is useful when permutation is not desired or supported.

```python
ASP.init_model_for_pruning(model, mask_calculator="m4n2_1d", verbosity=2, whitelist=[torch.nn.Linear, torch.nn.Conv2d], allow_recompute_mask=False, allow_permutation=False)
```

--------------------------------

### Integrate Triton LayerNorm

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/openfold_triton/README.md

Integrates a Triton kernel (LayerNormSmallShapeOptImpl) for Layer Normalization into a PyTorch nn.Module. The Triton kernel is used only when specific shape and training conditions are met; otherwise, it falls back to the standard F.layer_norm.

```python
from apex.contrib.openfold_triton import LayerNormSmallShapeOptImpl

# Integration with LayerNorm module:
class LayerNorm(nn.Module):
    # ...

    def _should_use_triton_kernels(self, x: torch.Tensor) -> bool:
        ln_triton_shapes = (
            (256, 128),
            (256, 256),
        )
        ln_triton_dim = 4
        return (
            self.training
            and x.dim() == ln_triton_dim
            and x.shape[-2:] in ln_triton_shapes
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._should_use_triton_kernels(x):
            return LayerNormSmallShapeOptImpl.apply(
                x, self.normalized_shape, self.weight, self.bias, self.eps
            )
        else:
            return F.layer_norm(
                x, self.normalized_shape, self.weight, self.bias, self.eps
            )
```

--------------------------------

### Enable/Disable Triton MHA

Source: https://github.com/nvidia/apex/blob/master/apex/contrib/openfold_triton/README.md

Dynamically control the activation of Triton-based Multi-Head Attention kernels at runtime.

```python
# Switch on/off MHA dynamically at runtime via:
mha.enable()
mha.disable()
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.