Accelerating Transformer Training with NVIDIA Transformer Engine, Fused Kernels, BF16, FP8, and GPU Benchmarking
In this tutorial, we discover how NVIDIA Transformer Engine accelerates transformer workloads by combining fused GPU kernels, BF16 computation, and hardware-aware FP8 execution. We start by putting in Transformer Engine and detecting the lively GPU structure in order that we are able to decide whether or not the runtime helps TE kernels, FP8 tensor cores, or solely the pure-PyTorch fallback path. We then look at core fused elements similar to te.Linear, te.LayerNorm, te.LayerNormLinear, te.LayerNormMLP, and te.TransformerLayer, whereas additionally configuring a delayed-scaling FP8 recipe that manages tensor scaling, amax historical past, and hybrid E4M3/E5M2 codecs. Using these elements, we assemble a compact GPT-style causal language mannequin, practice it on deterministic artificial sequences, evaluate higher-precision and FP8 execution, measure runtime and peak GPU reminiscence, examine FP8 metadata, and validate the educated mannequin via autoregressive technology.
import subprocess, sys, os
def pip_install(*pkgs):
subprocess.run([sys.executable, "-m", "pip", "install", "-q",
"--no-build-isolation", *pkgs], examine=False)
print(">> Installing transformer_engine[pytorch] (this may take a couple of minutes)...")
pip_install("transformer_engine[pytorch]")
import time, math, gc
import torch
import torch.nn as nn
import torch.nn.practical 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.main, props.minor)
GPU_NAME = props.identify
print(f">> GPU: {GPU_NAME} | compute functionality {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:
strive:
import transformer_engine.pytorch as te
from transformer_engine.widespread import recipe
print(">> Transformer Engine imported OK:",
getattr(te, "__version__", "unknown model"))
besides Exception as e:
print(f">> TE import failed ({e}); utilizing 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 just not None:
strive:
okay, purpose = te.fp8.check_fp8_support()
FP8_CAPABLE = bool(okay)
if not okay:
print(">> TE experiences FP8 unsupported:", purpose)
besides Exception:
move
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, system=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).form))
print(" te.LayerNorm ", tuple(ln(x_demo).form))
print(" te.LayerNormLinear", tuple(ln_lin(x_demo).form))
print(" te.LayerNormMLP ", tuple(ln_mlp(x_demo).form))
del lin, ln, ln_lin, ln_mlp, x_demo
gc.gather(); 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)
We set up NVIDIA Transformer Engine and initialize the PyTorch setting required for GPU-accelerated execution. We examine the lively GPU, compute functionality, and reminiscence capability to find out whether or not fused TE kernels and FP8 tensor cores can be found. We additionally validate the core fused modules and configure a delayed-scaling FP8 recipe whereas preserving an automated PyTorch fallback for unsupported {hardware}.
VOCAB, D_MODEL, N_HEADS, N_LAYERS, FFN, SEQ = 96, 768, 12, 4, 3072, 256
class MiniGPT_TE(nn.Module):
"""Causal LM the place each block is a single fused te.TransformerLayer."""
def __init__(self):
tremendous().__init__()
self.emb = nn.Embedding(VOCAB, D_MODEL)
self.pos = nn.Embedding(SEQ, D_MODEL)
self.blocks = nn.ModuleListing([
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 ahead(self, idx):
B, T = idx.form
h = self.emb(idx) + self.pos(torch.arange(T, system=idx.system))
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):
tremendous().__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 ahead(self, x, masks):
a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x),
attn_mask=masks, need_weights=False)
x = x + a
return x + self.mlp(self.ln2(x))
class MiniGPT_PT(nn.Module):
def __init__(self):
tremendous().__init__()
self.emb = nn.Embedding(VOCAB, D_MODEL)
self.pos = nn.Embedding(SEQ, D_MODEL)
self.blocks = nn.ModuleListing([Block_PT() for _ in range(N_LAYERS)])
self.ln_f = nn.LayerNorm(D_MODEL)
self.head = nn.Linear(D_MODEL, VOCAB, bias=False)
def ahead(self, idx):
B, T = idx.form
masks = torch.triu(torch.full((T, T), float("-inf"),
system=idx.system), diagonal=1)
h = self.emb(idx) + self.pos(torch.arange(T, system=idx.system))
for blk in self.blocks:
h = blk(h, masks)
return self.head(self.ln_f(h))
mannequin = (MiniGPT_TE() if TE_CAPABLE else MiniGPT_PT()).to(DEVICE)
n_params = sum(p.numel() for p in mannequin.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")
We outline a compact causal language mannequin utilizing fused te.TransformerLayer blocks for Transformer Engine execution. We additionally implement an equal pure-PyTorch transformer structure with multi-head consideration, layer normalization, residual connections, and feed-forward networks. We choose the suitable mannequin dynamically in keeping with GPU assist and report the ultimate parameter depend and architectural dimensions.
def make_batch(bsz=16):
section = torch.randint(0, VOCAB, (bsz, 1))
stride = torch.randint(1, 7, (bsz, 1))
steps = torch.arange(SEQ + 1).unsqueeze(0)
seq = (section + stride * steps) % VOCAB
return seq[:, :-1].to(DEVICE), seq[:, 1:].to(DEVICE)
choose = torch.optim.AdamW(mannequin.parameters(), lr=3e-4)
def run_step(x, y, use_fp8):
if TE_CAPABLE and use_fp8:
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
logits = mannequin(x)
else:
logits = mannequin(x)
loss = F.cross_entropy(logits.float().reshape(-1, VOCAB), y.reshape(-1))
choose.zero_grad(set_to_none=True)
loss.backward()
choose.step()
return loss.merchandise()
print(f"n>> Training 60 steps ({'FP8' if FP8_CAPABLE else 'BF16/FP32'})...")
t0 = time.time()
for step in vary(1, 61):
x, y = make_batch()
loss = run_step(x, y, use_fp8=FP8_CAPABLE)
if step % 10 == 0:
print(f" step {step:3d} | loss {loss:.4f} | "
f"{(time.time()-t0)/step*1000:.0f} ms/step")
print(f">> Final loss: {loss:.4f} (random guess can be ~{math.log(VOCAB):.2f})")
We create deterministic arithmetic-pattern sequences that enable the mannequin to study predictable token transitions throughout the vocabulary. We configure the AdamW optimizer and implement a coaching step that conditionally wraps the ahead move in te.fp8_autocast when FP8 execution is supported. We practice the mannequin for a number of iterations, monitor the loss and step latency, and evaluate the ultimate loss in opposition to the random-guess baseline.
def bench(use_fp8, iters=30, warmup=10):
x, y = make_batch(bsz=32)
for _ in vary(warmup):
run_step(x, y, use_fp8)
torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
t = time.time()
for _ in vary(iters):
run_step(x, y, use_fp8)
torch.cuda.synchronize()
ms = (time.time() - t) / iters * 1000
mem = torch.cuda.max_memory_allocated() / 1e9
return ms, mem
print("n>> Benchmark (batch 32, seq 256, fwd+bwd+optim):")
ms_hi, mem_hi = bench(use_fp8=False)
print(f" {'BF16' if TE_CAPABLE else 'FP32'}: {ms_hi:7.1f} ms/step | "
f"peak mem {mem_hi:.2f} GB")
if FP8_CAPABLE:
ms_f8, mem_f8 = bench(use_fp8=True)
print(f" FP8 : {ms_f8:7.1f} ms/step | peak mem {mem_f8:.2f} GB")
print(f" Speedup: {ms_hi/ms_f8:.2f}x "
f"(features develop with mannequin measurement — strive D_MODEL=2048, N_LAYERS=12)")
else:
print(" FP8 benchmark skipped — wants an sm_89+ GPU (L4/H100/Ada/Blackwell).")
if FP8_CAPABLE:
blk = mannequin.blocks[0]
for identify, m in blk.named_modules():
meta = getattr(m, "fp8_meta", None)
if meta and "scaling_fwd" in meta:
s = meta["scaling_fwd"]
print(f"n>> FP8 state of block-0 submodule '{identify}':")
print(" scale :", s.scale.flatten()[:4].tolist())
print(" amax_history:", s.amax_history[0, :4].tolist())
break
We benchmark ahead propagation, backpropagation, and optimizer updates utilizing higher-precision and FP8 execution modes. We measure common training-step latency and peak allotted GPU reminiscence to quantify the efficiency and reminiscence affect of reduced-precision computation. We additionally examine the scaling components and amax historical past maintained by Transformer Engine to grasp how delayed scaling stabilizes FP8 tensors.
@torch.no_grad()
def generate(prompt_len=8, gen_len=24):
x, _ = make_batch(bsz=1)
ctx = x[:, :prompt_len]
for _ in vary(gen_len):
inp = ctx[:, -SEQ:]
if TE_CAPABLE and FP8_CAPABLE:
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
logits = mannequin(inp)
else:
logits = mannequin(inp)
nxt = logits[:, -1].argmax(-1, keepdim=True)
ctx = torch.cat([ctx, nxt], dim=1)
return ctx[0].tolist()
seq = generate()
print("n>> Greedy technology (ought to proceed the arithmetic sample):")
print(" immediate+gen:", seq)
diffs = [(b - a) % VOCAB for a, b in zip(seq, seq[1:])]
print(" step diffs:", diffs, "<- fixed stride = mannequin discovered the rule")
print("n>> Done! Things to strive subsequent:")
print(" * Scale up: D_MODEL=2048, N_LAYERS=12 -> FP8 speedup turns into dramatic")
print(" * recipe.Format.E4M3 vs HYBRID; amax_history_len=1024")
print(" * te.LayerNormMLP / te.LayerNormLinear in your personal architectures")
print(" * fp8_model_init() to retailer weights themselves in FP8 for inference")
We implement grasping autoregressive technology by repeatedly feeding the most recent context into the educated causal language mannequin. We evaluate consecutive generated tokens to confirm whether or not the mannequin preserves the fixed arithmetic stride current within the artificial coaching knowledge. We conclude by figuring out sensible extensions, together with bigger mannequin dimensions, various FP8 codecs, longer amax histories, fused modules, and FP8 weight initialization.
In conclusion, we demonstrated how we combine NVIDIA Transformer Engine into an end-to-end transformer coaching workflow whereas preserving compatibility throughout completely different Colab GPU environments. We used fused transformer modules to cut back kernel-launch overhead and reminiscence site visitors, utilized FP8 autocasting with delayed scaling when supported, and retained BF16 or FP32 execution via an automated PyTorch fallback. By coaching and benchmarking the identical mini causal language mannequin, we noticed how {hardware} functionality, numerical format, fused execution, and mannequin scale affect coaching velocity and reminiscence consumption. We additionally inspected the inner scaling components and amax historical past that assist secure FP8 computation, which provides us a clearer understanding of how Transformer Engine manages reduced-precision arithmetic.
Check out the Full Codes. Also, be happy to observe us on Twitter and don’t overlook to hitch our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to companion with us for selling your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar and many others.? Connect with us
The submit Accelerating Transformer Training with NVIDIA Transformer Engine, Fused Kernels, BF16, FP8, and GPU Benchmarking appeared first on MarkTechPost.
