On 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 will decide whether or not the runtime helps TE kernels, FP8 tensor cores, or solely the pure-PyTorch fallback path. We then study core fused parts resembling 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. Utilizing these parts, we assemble a compact GPT-style causal language mannequin, practice it on deterministic artificial sequences, examine higher-precision and FP8 execution, measure runtime and peak GPU reminiscence, examine FP8 metadata, and validate the skilled mannequin via autoregressive era.
import subprocess, sys, os
def pip_install(*pkgs):
subprocess.run([sys.executable, "-m", "pip", "install", "-q",
"--no-build-isolation", *pkgs], verify=False)
print(">> Putting in transformer_engine[pytorch] (this will take a couple of minutes)...")
pip_install("transformer_engine[pytorch]")
import time, math, gc
import torch
import torch.nn as nn
import torch.nn.useful as F
assert torch.cuda.is_available(), "Allow 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:
attempt:
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:
attempt:
okay, motive = te.fp8.check_fp8_support()
FP8_CAPABLE = bool(okay)
if not okay:
print(">> TE reviews FP8 unsupported:", motive)
besides Exception:
go
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, machine=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.acquire(); 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 computerized 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.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 ahead(self, idx):
B, T = idx.form
h = self.emb(idx) + self.pos(torch.arange(T, machine=idx.machine))
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.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 ahead(self, idx):
B, T = idx.form
masks = torch.triu(torch.full((T, T), float("-inf"),
machine=idx.machine), diagonal=1)
h = self.emb(idx) + self.pos(torch.arange(T, machine=idx.machine))
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>> Mannequin: {'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 accordance 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)
decide = 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))
decide.zero_grad(set_to_none=True)
loss.backward()
decide.step()
return loss.merchandise()
print(f"n>> Coaching 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">> Closing loss: {loss:.4f} (random guess could 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 go 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 examine the ultimate loss towards 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 dimension — attempt 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 influence of reduced-precision computation. We additionally examine the scaling elements 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>> Grasping era (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 realized the rule")
print("n>> Achieved! Issues to attempt 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 era by repeatedly feeding the most recent context into the skilled causal language mannequin. We examine consecutive generated tokens to confirm whether or not the mannequin preserves the fixed arithmetic stride current within the artificial coaching information. 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 totally different Colab GPU environments. We used fused transformer modules to cut back kernel-launch overhead and reminiscence visitors, utilized FP8 autocasting with delayed scaling when supported, and retained BF16 or FP32 execution via an computerized 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 pace and reminiscence consumption. We additionally inspected the inner scaling elements and amax historical past that assist steady FP8 computation, which provides us a clearer understanding of how Transformer Engine manages reduced-precision arithmetic.
Take a look at the Full Codes. Additionally, be happy to comply with us on Twitter and don’t overlook to hitch our 150k+ML SubReddit and Subscribe to our Publication. Wait! are you on telegram? now you may be a part of us on telegram as effectively.
Must associate with us for selling your GitHub Repo OR Hugging Face Web page OR Product Launch OR Webinar and so on.? Join with us
Sana Hassan, a consulting intern at Marktechpost and dual-degree scholar at IIT Madras, is enthusiastic about making use of expertise and AI to handle real-world challenges. With a eager curiosity in fixing sensible issues, he brings a contemporary perspective to the intersection of AI and real-life options.
