NVIDIA Transformer Engine combines fused GPU kernels, BF16 computation, and hardware-aware FP8 execution to speed up transformer workloads. The tutorial begins by installing the library and checking the active GPU architecture to see if the runtime supports TE kernels, FP8 tensor cores, or only the pure-PyTorch fallback path.
Installation and hardware checks
The script installs transformer_engine[pytorch] and inspects the device properties to determine capabilities. It checks the compute capability to see if the GPU is Ampere (8.0+) or Hopper (8.9+). If the hardware is older, such as a T4, the code disables TE kernels and falls back to standard PyTorch.
import subprocess, sys, os
def pip_install(*pkgs):
subprocess.run([sys.executable, "-m", "pip", "install", "-q",
"--no-build-isolation", *pkgs], check=False)
print(">> Installing transformer_engine[pytorch] (this can take a few minutes)...")
pip_install("transformer_engine[pytorch]")
import time, math, gc
import torch
import torch.nn as nn
import torch.nn.functional as F
assert torch.cuda.is_available(), "Enable a GPU runtime in Colab first!"
DEVICE = "cuda"
props = torch.cuda.get_device_properties(0)
CC = (props.major, props.minor)
GPU_NAME = props.name
print(f">> GPU: {GPU_NAME} | compute capability {CC[0]}.{CC[1]} | "
f"{props.total_memory/1e9:.1f} GB")
TE_CAPABLE = CC >= (8, 0)
FP8_CAPABLE = CC >= (8, 9)
te = None
if TE_CAPABLE:
try:
import transformer_engine.pytorch as te
from transformer_engine.common import recipe
print(">> Transformer Engine imported OK:",
getattr(te, "__version__", "unknown version"))
except Exception as e:
print(f">> TE import failed ({e}); using pure-PyTorch fallback.")
TE_CAPABLE = FP8_CAPABLE = False
else:
print(">> GPU is pre-Ampere (e.g. T4): TE kernels unsupported -> fallback mode.")
if TE_CAPABLE and FP8_CAPABLE and te is not None:
try:
ok, reason = te.fp8.check_fp8_support()
FP8_CAPABLE = bool(ok)
if not ok:
print(">> TE reports FP8 unsupported:", reason)
except Exception:
pass
print(f">> Mode: TE={'ON' if TE_CAPABLE else 'OFF'} | "
f"FP8={'ON' if FP8_CAPABLE else 'OFF (will use BF16)'}")
torch.manual_seed(1234)
if TE_CAPABLE:
H = 768
x_demo = torch.randn(8, 32, H, device=DEVICE, dtype=torch.bfloat16)
lin = te.Linear(H, H, bias=True, params_dtype=torch.bfloat16).to(DEVICE)
ln = te.LayerNorm(H, params_dtype=torch.bfloat16).to(DEVICE)
ln_lin = te.LayerNormLinear(H, 3 * H, params_dtype=torch.bfloat16).to(DEVICE)
ln_mlp = te.LayerNormMLP(H, 4 * H, params_dtype=torch.bfloat16).to(DEVICE)
with torch.no_grad():
print("\n>> Module tour (shapes):")
print(" te.Linear ", tuple(lin(x_demo).shape))
print(" te.LayerNorm ", tuple(ln(x_demo).shape))
print(" te.LayerNormLinear", tuple(ln_lin(x_demo).shape))
print(" te.LayerNormMLP ", tuple(ln_mlp(x_demo).shape))
del lin, ln, ln_lin, ln_mlp, x_demo
gc.collect(); torch.cuda.empty_cache()
fp8_recipe = None
if FP8_CAPABLE:
fp8_recipe = recipe.DelayedScaling(
fp8_format=recipe.Format.HYBRID,
amax_history_len=16,
amax_compute_algo="max",
)
print("\n>> FP8 recipe:", fp8_recipe)
The code validates core fused modules like te.Linear, te.LayerNorm, te.LayerNormLinear, and te.LayerNormMLP. It also configures a delayed-scaling FP8 recipe that manages tensor scaling, amax history, and hybrid E4M3/E5M2 formats. If FP8 support is confirmed, the recipe uses a hybrid format with an amax history length of 16 and a max compute algorithm.
Model architecture
The tutorial defines a compact causal language model using fused te.TransformerLayer blocks. It also implements an equivalent pure-PyTorch transformer architecture with multi-head attention, layer normalization, residual connections, and feed-forward networks.
VOCAB, D_MODEL, N_HEADS, N_LAYERS, FFN, SEQ = 96, 768, 12, 4, 3072, 256
class MiniGPT_TE(nn.Module):
"""Causal LM where every block is a single fused te.TransformerLayer."""
def __init__(self):
super().__init__()
self.emb = nn.Embedding(VOCAB, D_MODEL)
self.pos = nn.Embedding(SEQ, D_MODEL)
self.blocks = nn.ModuleList([
te.TransformerLayer(
hidden_size=D_MODEL,
ffn_hidden_size=FFN,
num_attention_heads=N_HEADS,
self_attn_mask_type="causal",
layer_number=i + 1,
params_dtype=torch.bfloat16,
hidden_dropout=0.0,
attention_dropout=0.0,
)
for i in range(N_LAYERS)
])
self.ln_f = nn.LayerNorm(D_MODEL)
self.head = nn.Linear(D_MODEL, VOCAB, bias=False)
def forward(self, idx):
B, T = idx.shape
h = self.emb(idx) + self.pos(torch.arange(T, device=idx.device))
h = h.to(torch.bfloat16)
for blk in self.blocks:
h = blk(h)
h = self.ln_f(h.float())
return self.head(h)
class Block_PT(nn.Module):
"""Plain-PyTorch transformer block, mirrors te.TransformerLayer."""
def __init__(self):
super().__init__()
self.ln1 = nn.LayerNorm(D_MODEL)
self.attn = nn.MultiheadAttention(D_MODEL, N_HEADS, batch_first=True)
self.ln2 = nn.LayerNorm(D_MODEL)
self.mlp = nn.Sequential(nn.Linear(D_MODEL, FFN), nn.GELU(),
nn.Linear(FFN, D_MODEL))
def forward(self, x, mask):
a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x),
attn_mask=mask, need_weights=False)
x = x + a
return x + self.mlp(self.ln2(x))
class MiniGPT_PT(nn.Module):
def __init__(self):
super().__init__()
self.emb = nn.Embedding(VOCAB, D_MODEL)
self.pos = nn.Embedding(SEQ, D_MODEL)
self.blocks = nn.ModuleList([Block_PT() for _ in range(N_LAYERS)])
self.ln_f = nn.LayerNorm(D_MODEL)
self.head = nn.Linear(D_MODEL, VOCAB, bias=False)
def forward(self, idx):
B, T = idx.shape
mask = torch.triu(torch.full((T, T), float("-inf"),
device=idx.device), diagonal=1)
h = self.emb(idx) + self.pos(torch.arange(T, device=idx.device))
for blk in self.blocks:
h = blk(h, mask)
return self.head(self.ln_f(h))
model = (MiniGPT_TE() if TE_CAPABLE else MiniGPT_PT()).to(DEVICE)
n_params = sum(p.numel() for p in model.parameters())
print(f"\n>> Model: {'TE fused' if TE_CAPABLE else 'pure PyTorch'} | "
f"{n_params/1e6:.1f}M params | {N_LAYERS} layers x {D_MODEL}d")
The script selects the appropriate model dynamically based on GPU support. It reports the final parameter count and architectural dimensions. The resulting model uses 96 tokens in the vocabulary, 768 dimensions, 12 attention heads, 4 layers, a feed-forward size of 3072, and a sequence length of 256.
Training and benchmarking
The tutorial creates deterministic arithmetic-pattern sequences to allow the model to learn predictable token transitions across
Related articles
