Transformer Engine

https://docs.nvidia.com/deeplearning/transformer-engine/
Admin
Transformer Engine is a library for accelerating Transformer models on NVIDIA GPUs using...
Tokens:310,580
Snippets:1,744
Trust Score:10
Update:3 weeks ago
Show doc for...
Tokens:
Raw
### Install required Python packages

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.html

Install the required Python packages using pip. This is the setup step before running the fine-tuning script.

```bash
pip install -r requirements.txt

```

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

### Install required Python packages

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/te_llama/tutorial_accelerate_hf_llama_with_te.html

Command to install the necessary Python packages for the tutorial. Uncomment and run this cell when running the tutorial for the first time.

```bash
# Uncomment and run this cell when running the tutorial for the first time
# %pip install -r requirements.txt
```

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

### Install tutorial dependencies with pip

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/te_gemma/tutorial_generation_gemma_with_te.html

Install the required Python packages listed in requirements.txt. This cell is commented out and should be uncommented and run only on the first execution of the tutorial.

```python
# Uncomment and run this cell when running the tutorial for the first time
# %pip install -r requirements.txt

```

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

### Configuration file example for precision debug tools

Source: https://docs.nvidia.com/deeplearning/transformer-engine/debug/2_config_file_structure.html

Example configuration file showing three sections: DisableFP8GEMM to disable wgrad GEMMs, DisableFP8Layer to disable all GEMMs in a layer, and LogTensorStats to log wgrad stats for fc1.

```yaml
# Disables wgrad in all 4 GEMMs
section1:
  enabled: True
  layers:
    layer_types: [transformer_layer]
  transformer_engine:
    DisableFP8GEMM:
      enabled: True
      gemms: [wgrad]

# Disables all GEMMs in layernorm_mlp layer
section2:
  enabled: True
  layers:
    layer_types: [layernorm_mlp]
  transformer_engine:
    DisableFP8Layer:
      enabled: True

# Logs wgrad stats in fc1
section3:
  enabled: True
  layers:
    layer_types: [fc1]
  transformer_engine:
    LogTensorStats:
      enabled: True
      stats: [min]
      tensors: [wgrad]
      freq: 1
      start_step: 0
      end_step: 50
```

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

### nvte_get_grouped_gemm_setup_workspace_size()

Source: https://docs.nvidia.com/deeplearning/transformer-engine/api/c/gemm.html

Gets the workspace size required for grouped GEMM setup.

```APIDOC
## `nvte_get_grouped_gemm_setup_workspace_size()`

### Description
Gets the workspace size required for grouped GEMM setup.

### Method
Function

### Signature
`size_t nvte_get_grouped_gemm_setup_workspace_size(...)`

### Parameters
Not specified in the source.
```

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

### Manual CPU offloading with get_cpu_offload_context and ManualOffloadSynchronizer

Source: https://docs.nvidia.com/deeplearning/transformer-engine/features/other_optimizations/cpu_offloading/cpu_offloading.html

Demonstrates manual CPU offloading with Transformer Engine's get_cpu_offload_context. Set manual_synchronization=True and optionally pass a custom offload_stream to get a ManualOffloadSynchronizer. The example shows calling start_offload_layer after each forward layer, release_activation_forward_gpu_memory to free GPU memory, and start_reload_layer before backward. Skipping these calls for a layer disables offloading for that layer.

```python
import torch
from transformer_engine.pytorch import get_cpu_offload_context

# Setup
num_layers = 12
layers = [torch.nn.Linear(1024, 1024).cuda() for _ in range(num_layers)]
x = torch.randn(16, 1024, 1024, device="cuda")

offload_stream = torch.cuda.Stream()
cpu_offload_context, sync_function, manual_controller = get_cpu_offload_context(
    enabled=True,
    model_layers=num_layers,
    manual_synchronization=True,
    offload_stream=offload_stream,
)

# Forward pass - manually trigger offload after each layer
for i in range(num_layers):
    with cpu_offload_context:
        x = layers[i](x)
    x = sync_function(x)
    manual_controller.start_offload_layer(i)

# Release GPU memory (each call waits for that layer's offload to complete)
for i in range(num_layers):
    manual_controller.release_activation_forward_gpu_memory(i)

# Start reloading before backward
for i in range(num_layers - 1, -1, -1):
    manual_controller.start_reload_layer(i)

# Backward pass
loss = x.sum()
loss.backward()
```

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

### Install Transformer Engine from source

Source: https://docs.nvidia.com/deeplearning/transformer-engine/installation.html

Installs Transformer Engine from source: clones the repository, checks out the stable branch, clones submodules, optionally sets NVTE_FRAMEWORK, and builds/installs.

```bash
# Clone repository, checkout stable branch, clone submodules
git clone --branch stable --recursive https://github.com/NVIDIA/TransformerEngine.git

cd TransformerEngine
export NVTE_FRAMEWORK=pytorch         # Optionally set framework
pip3 install --no-build-isolation .   # Build and install

```

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

### Install Transformer Engine with PyTorch from PyPI

Source: https://docs.nvidia.com/deeplearning/transformer-engine/installation.html

Installs Transformer Engine with PyTorch bindings from PyPI. The framework must be specified as an extra dependency in a comma-separated list (e.g. [jax,pytorch]).

```bash
pip3 install --no-build-isolation transformer_engine[pytorch]

```

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

### Run example_attention.py with NVTE_DEBUG_LEVEL=2

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/attention/attention.html

Runs example_attention.py with debug level 2 to get detailed backend selection logic, including the configuration used and the reason for disabling a backend. Useful for debugging backend selection issues.

```bash
!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 python example_attention.py
```

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

### Import TE JAX attention modules and set NVTE_FUSED_RING_ATTENTION_USE_SCAN

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/jax/attention_context_parallel.html

Imports and environment setup for the THD context-parallel attention example. Sets NVTE_FUSED_RING_ATTENTION_USE_SCAN to "0" before JAX compiles the first fused attention call so the example uses the non-scan Ring implementation, matching the distributed tests.

```python
import os
import time
from typing import Tuple

# Ring + SWA uses the non-scan Ring implementation. Set this before JAX compiles
# the first fused attention call so the example follows the distributed tests.
os.environ.setdefault("NVTE_FUSED_RING_ATTENTION_USE_SCAN", "0")

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

import transformer_engine.jax as te
from transformer_engine.jax.attention import (
    AttnBiasType,
    AttnMaskType,
    AttnSoftmaxType,
    CPStrategy,
    QKVLayout,
    ReorderStrategy,
    SequenceDescriptor,
    fused_attn,
    inverse_reorder_causal_load_balancing,
    is_fused_attn_kernel_available,
    reorder_causal_load_balancing,
)
from transformer_engine.jax.sharding import MeshResource
```

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

### Verbose build logging with pip install -v -v -v

Source: https://docs.nvidia.com/deeplearning/transformer-engine/installation.html

Run this command from the transformer_engine directory to get detailed build logs for diagnosing issues. The -v -v -v flags increase verbosity, and --no-build-isolation ensures the build uses the current environment.

```bash
cd transformer_engine
pip install -v -v -v --no-build-isolation .

```

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

### Finetune TELlama model in BF16 precision

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/te_llama/tutorial_accelerate_hf_llama_with_te.html

Setup and finetuning workflow for the TELlama implementation in BF16 precision. Requires a Hugging Face access token and optionally a weights cache directory. The output shows 10 finetuning steps with an average step time of 185 milliseconds.

```python
# Restart the notebook (to flush the GPU memory)
from utils import restart_jupyter_notebook
restart_jupyter_notebook()


# Import necessary packages, methods and variables
from utils import *


# Provide Huggingface Access Token
hyperparams.hf_access_token = ""
assert hyperparams.hf_access_token, "Provide a HF API Access Token!"

# Provide a directory to cache weights in to avoid downloading them every time.
# (By default, weights are cached in `~/.cache/huggingface/hub/models`)
hyperparams.weights_cache_dir = ""

# For Llama 2, uncomment this line (also set by default)
hyperparams.model_name = "meta-llama/Llama-2-7b-hf"

# For Llama 3, uncomment this line
# hyperparams.model_name = "meta-llama/Meta-Llama-3-8B"

hyperparams.mixed_precision = "bf16"


# Init the model and accelerator wrapper
model = init_te_llama_model(hyperparams)
accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)


# Finetune the model
finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)

```

```text
10 finetuning steps complete!
Average time taken per step: 185 milliseconds

```

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

### Install Transformer Engine with test dependencies

Source: https://docs.nvidia.com/deeplearning/transformer-engine/installation.html

Installs Transformer Engine with extra dependencies for testing.

```bash
pip3 install --no-build-isolation .[test]

```

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

### Output of example_attention.py with NVTE_DEBUG_LEVEL=2

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/attention/attention.html

Expected output when running example_attention.py with NVTE_DEBUG_LEVEL=2. Shows detailed configuration and backend selection logic, including why a backend is disabled.

```text
Run cuDNN attention...
[DEBUG    | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': <Version('2.4.2')>, 'cudnn_version': '9.3.0', 'qkv_type': <class 'torch.Tensor'>, 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}
[DEBUG    | DotProductAttention]: Disabling FlashAttention due to NVTE_FLASH_ATTN=0
[DEBUG    | DotProductAttention]: Available backends = {FlashAttention=False, FusedAttention=True (sub-backend 1), UnfusedDotProductAttention=True}
[DEBUG    | DotProductAttention]: Selected backend = FusedAttention (sub-backend 1)
[INFO     | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)

Run flash-attention...
[DEBUG    | DotProductAttention]: Running with config={'transformer_engine_version': '1.10.0.dev0+ee85a91', 'compute_capability': 'sm90', 'flash_attn_version': <Version('2.4.2')>, 'cudnn_version': '9.3.0', 'qkv_type': <class 'torch.Tensor'>, 'qkv_dtype': torch.bfloat16, 'qkv_layout': 'bshd_bshd_bshd', 'batch_size': 2, 'num_heads': 16, 'num_gqa_groups': 16, 'max_seqlen_q': 512, 'max_seqlen_kv': 512, 'head_dim_qk': 64, 'head_dim_v': 64, 'attn_mask_type': 'no_mask', 'window_size': (-1, -1), 'alibi_slopes_shape': None, 'core_attention_bias_type': 'no_bias', 'core_attention_bias_shape': None, 'core_attention_bias_requires_grad': False, 'pad_between_seqs': False, 'attention_dropout': 0.0, 'context_parallel': False, 'deterministic': False, 'is_training': True, 'fp8': False, 'fp8_meta': {'fp8_checkpoint': False, 'fp8_group': None, 'recipe': margin=0, format=HYBRID, amax_history_len=1024, wgrad_override=False, fp8_dpa=False, fp8_mha=False}}
[DEBUG    | DotProductAttention]: Disabling FusedAttention due to NVTE_FUSED_ATTN=0
[DEBUG    | DotProductAttention]: Available backends = {FlashAttention=True, FusedAttention=False, UnfusedDotProductAttention=True}
[DEBUG    | DotProductAttention]: Selected backend = FlashAttention
[INFO     | DotProductAttention]: Running with FlashAttention backend

Test passed.
```

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

### Install stable Transformer Engine from GitHub

Source: https://docs.nvidia.com/deeplearning/transformer-engine/installation.html

Installs the latest stable release of Transformer Engine from GitHub. Automatically detects installed supported deep learning frameworks; set NVTE_FRAMEWORK to a comma-separated list to specify frameworks explicitly.

```bash
pip3 install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@stable

```

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

### Install Transformer Engine core library from PyPI

Source: https://docs.nvidia.com/deeplearning/transformer-engine/installation.html

Installs the core Transformer Engine library without framework extensions. By default, this installs the core library compiled for CUDA 12; the CUDA major version can be specified by using core_cu12 or core_cu13.

```bash
pip3 install transformer_engine[core]

```

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

### Run fine-tuning with torchrun and expected output

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.html

Launch the fine-tuning loop across 8 GPUs with torchrun. The expected output shows 30 steps complete with a median time per step of 747 ms, which is 3.31x faster than the HF baseline.

```bash
torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 1 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30
```

```text
30 fine-tuning steps complete!
Median time per step: 747 ms
```

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

### Import dependencies and create synthetic data

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/advanced_optimizations.html

Sets up the environment and synthetic data for the performance optimization example. Imports PyTorch, Transformer Engine, and the quickstart_utils helper module, then defines layer configuration (hidden_size=4096, sequence_length=2048, batch_size=4, ffn_hidden_size=16384, num_attention_heads=32) and creates random input tensors x and dy on GPU with float16 dtype.

```python
import torch
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import Format, DelayedScaling
import quickstart_utils as utils

# Layer configuration
hidden_size = 4096
sequence_length = 2048
batch_size = 4
ffn_hidden_size = 16384
num_attention_heads = 32
dtype = torch.float16

# Synthetic data
x = torch.rand(sequence_length, batch_size, hidden_size).cuda().to(dtype=dtype)
dy = torch.rand(sequence_length, batch_size, hidden_size).cuda().to(dtype=dtype)
```

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

### nvte_get_grouped_gemm_setup_workspace_size

Source: https://docs.nvidia.com/deeplearning/transformer-engine/api/c/gemm.html

Documentation for the nvte_get_grouped_gemm_setup_workspace_size function.

```APIDOC
## nvte_get_grouped_gemm_setup_workspace_size

### Description
Returns the workspace size needed for grouped GEMM setup.

### Function Signature
```c
size_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors);
```

### Parameters
- **num_tensors** (size_t) - Number of tensors.

### Returns
- **size_t** - Workspace size in bytes.
```

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

### Install a specific PR of Transformer Engine from GitHub

Source: https://docs.nvidia.com/deeplearning/transformer-engine/installation.html

Installs a specific pull request of Transformer Engine from GitHub. Replace NNN with the PR number.

```bash
pip3 install --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@refs/pull/NNN/merge

```

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

### Output of example_attention.py with NVTE_DEBUG_LEVEL=1

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/attention/attention.html

Expected output when running example_attention.py with NVTE_DEBUG_LEVEL=1. Shows the selected backend for each attention type and a successful test completion.

```text
Run cuDNN attention...
[INFO     | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)

Run flash-attention...
[INFO     | DotProductAttention]: Running with FlashAttention backend

Test passed.
```

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

### Build with Debug Symbols

Source: https://docs.nvidia.com/deeplearning/transformer-engine/envvars.html

Set environment variables to enable debug symbols and ccache, then install the package in editable mode.

```bash
export NVTE_BUILD_DEBUG=1
export NVTE_USE_CCACHE=1
pip install -e .

```

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

### Run example_attention.py with NVTE_DEBUG_LEVEL=1

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/attention/attention.html

Runs example_attention.py with debug level 1 to identify which attention backend (cuDNN or flash-attention) is used at runtime. Output shows INFO messages indicating the selected backend.

```bash
!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=1 python example_attention.py
```

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

### get_amax

Source: https://docs.nvidia.com/deeplearning/transformer-engine/api/c/transformer_engine.html

Gets the AMAX of the tensor.

```APIDOC
## get_amax

### Description
Gets the AMAX of the tensor.

### Signature
```cpp
inline NVTEBasicTensor get_amax() const noexcept
```

### Returns
- **NVTEBasicTensor** - The AMAX.
```

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

### get_scale

Source: https://docs.nvidia.com/deeplearning/transformer-engine/api/c/transformer_engine.html

Gets the scale of the tensor.

```APIDOC
## get_scale

### Description
Gets the scale of the tensor.

### Signature
```cpp
inline NVTEBasicTensor get_scale() const noexcept
```

### Returns
- **NVTEBasicTensor** - The scale.
```

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

### Run fine-tuning with GroupedLinear improvement

Source: https://docs.nvidia.com/deeplearning/transformer-engine/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.html

Command to execute the fine-tuning script with improvement 2, using 8 processes, expert parallelism size 2, batch size 12, max sequence length 8192, warmup steps 5, and 30 training steps.

```bash
torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 2 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30

```

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

### get_parameter

Source: https://docs.nvidia.com/deeplearning/transformer-engine/api/c/transformer_engine.html

Gets a parameter of the tensor.

```APIDOC
## get_parameter

### Description
Gets a parameter of the tensor.

### Signature
```cpp
inline NVTEBasicTensor get_parameter(const NVTETensorParam param) const noexcept
```

### Parameters
- **param** (NVTETensorParam) - [in] The parameter to get.

### Returns
- **NVTEBasicTensor** - The requested parameter.
```

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

### Configure FP8 for Attention

Source: https://docs.nvidia.com/deeplearning/transformer-engine/envvars.html

Configure FP8 for attention with DelayedScaling recipe, HYBRID format, and amax history settings, then run training.

```bash
# Use DelayedScaling for attention, CurrentScaling for linear layers
export NVTE_DPA_FP8_RECIPE="DelayedScaling"
export NVTE_DPA_FP8_FORMAT="HYBRID"
export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent"
export NVTE_DPA_FP8DS_AMAX_HISTLEN=1024
python train.py

```

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

### JAX autocast with DelayedScaling recipe

Source: https://docs.nvidia.com/deeplearning/transformer-engine/features/low_precision_training/introduction/introduction.html

Basic JAX example using the Transformer Engine `autocast` context manager. Unlike PyTorch, model initialization must happen inside the context to capture quantization metadata, and both forward and backward passes are inside it. Requires SM89+ (Ada or newer).

```python
import jax
import jax.numpy as jnp
import transformer_engine.jax as te
from transformer_engine.jax.flax import TransformerLayer
from transformer_engine.common.recipe import DelayedScaling, Format

# Set up recipe
recipe = DelayedScaling()

# Model initialization must happen inside autocast
with te.autocast(enabled=True, recipe=recipe):
    layer = TransformerLayer(
        hidden_size=1024,
        mlp_hidden_size=4096,
        num_attention_heads=16,
    )

    init_key, dropout_key = jax.random.split(jax.random.PRNGKey(0))
    x = jax.random.normal(init_key, (32, 128, 1024), dtype=jnp.bfloat16)
    var_collect = layer.init({"params": init_key, "dropout": dropout_key}, x)

    # Forward and backward pass (both inside autocast for JAX)
    def loss_fn(var_collect):
        output = layer.apply(var_collect, x, rngs={"dropout": dropout_key})
        return output.sum()

    loss, grads = jax.value_and_grad(loss_fn)(var_collect)
```

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

### get_columnwise_data

Source: https://docs.nvidia.com/deeplearning/transformer-engine/api/c/transformer_engine.html

Gets the columnwise data of the tensor.

```APIDOC
## get_columnwise_data

### Description
Gets the columnwise data of the tensor.

### Signature
```cpp
inline NVTEBasicTensor get_columnwise_data() const noexcept
```

### Returns
- **NVTEBasicTensor** - The columnwise data.
```

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

### initialize()

Source: https://docs.nvidia.com/deeplearning/transformer-engine/debug/3_api_debug_setup.html

Initializes Nvidia-DL-Framework-Inspect on every rank in the global context. Must be called once per rank to enable precision debug features.

```APIDOC
## initialize()

### Description
Must be called once on every rank in the global context to initialize Nvidia-DL-Framework-Inspect.

### Method
initialize

### Parameters
- **config_file** (_str_, default="") - Path to the configuration YAML file containing features to enable and layer names. If one wants to run without the configuration file, pass `""`.
- **feature_dirs** (_List[str] | str_) - List of directories containing features to load and register. One needs to pass `[/path/to/transformerengine/transformer_engine/debug/features]` to use TE features.
- **logger** (_Union[BaseLogger, None]_, default=None) - Logger for logging tensor statistics. Should adhere to `BaseLogger` from the Nvidia-DL-Framework-Inspect package.
- **log_dir** (_str_, default=".") - Directory path to hold `debug_logs` and `debug_statistics_logs`.
- **tb_writer** (_TensorBoardWriter_, default=None) - TensorBoard writer for logging.
- **default_logging_enabled** (_bool_, default=False) - Enable default logging to the file.

### Request Example
```python
import nvdlfw_inspect.api as debug_api

debug_api.initialize(
    config_file="./config.yaml",
    feature_dirs=["/path/to/transformer_engine/debug/features"],
    log_dir="./log_dir")
```
```
Tokens:
Raw