transformer-engine 1.10.0

PyTorch

import torch
import transformer_engine.pytorch as te
from transformer_engine.common import recipe

# Set dimensions.
in_features = 768
out_features = 3072
hidden_size = 2048

# Initialize model and inputs.
model = te.Linear(in_features, out_features, bias=True)
inp = torch.randn(hidden_size, in_features, device="cuda")

# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3)

# Enable autocasting for the forward pass
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    out = model(inp)

loss = out.sum()
loss.backward()

JAX

Flax

import flax
import jax
import jax.numpy as jnp
import transformer_engine.jax as te
import transformer_engine.jax.flax as te_flax
from transformer_engine.common import recipe

BATCH = 32
SEQLEN = 128
HIDDEN = 1024

# Initialize RNG and inputs.
rng = jax.random.PRNGKey(0)
init_rng, data_rng = jax.random.split(rng)
inp = jax.random.normal(data_rng, [BATCH, SEQLEN, HIDDEN], jnp.float32)

# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.HYBRID)

# Enable autocasting for the forward pass
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    model = te_flax.DenseGeneral(features=HIDDEN)

    def loss_fn(params, other_vars, inp):
      out = model.apply({'params':params, **other_vars}, inp)
      return jnp.mean(out)

    # Initialize models.
    variables = model.init(init_rng, inp)
    other_variables, params = flax.core.pop(variables, 'params')

    # Construct the forward and backward function
    fwd_bwd_fn = jax.value_and_grad(loss_fn, argnums=(0, 1))

    for _ in range(10):
      loss, (param_grads, other_grads) = fwd_bwd_fn(params, other_variables, inp)

先决条件

• Linux x86_64

• Hopper 需要 CUDA 12.0+，Ada 需要 CUDA 12.1+

• 支持 CUDA 12.0 或更高版本的 NVIDIA 驱动程序

• cuDNN 8.1 或更高版本

• 对于融合注意力，需要 CUDA 12.1 或更高版本、支持 CUDA 12.1 或更高版本的 NVIDIA 驱动程序和 cuDNN 8.9 或更高版本。

Docker

通过使用 NVIDIA GPU Cloud (NGC) 目录上的 Docker 镜像，是开始使用 Transformer Engine 的最快方式。NVIDIA GPU Cloud (NGC) 目录。例如，要交互式地使用 NGC PyTorch 容器，

docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:23.10-py3

其中 23.10 是容器版本。例如，10 月 2023 年发布的版本为 23.10。

pip

要安装 Transformer Engine 的最新稳定版本，

pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable

这将自动检测是否安装了任何受支持的深度学习框架，并构建对它们的 Transformer Engine 支持。要显式指定框架，设置环境变量 NVTE_FRAMEWORK 为逗号分隔的列表（例如 NVTE_FRAMEWORK=jax,pytorch）。

从源代码安装

请参阅安装指南.

使用 FlashAttention-2 编译

Transformer Engine 版本 v0.11.0 为 PyTorch 添加了对 FlashAttention-2 的支持，以提高性能。

已知问题是 FlashAttention-2 编译资源密集且需要大量 RAM（请参阅错误），这可能导致 Transformer Engine 安装过程中出现内存不足错误。请尝试设置环境中的 MAX_JOBS=1 以绕过此问题。

请注意，NGC PyTorch 23.08+ 容器已包含 FlashAttention-2。