transformer-engine-cu12 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 的最快方式。例如，要交互式地使用 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。