NVIDIA Transformer Engine 加速训练教程

In this tutorial, we explore how NVIDIA Transformer Engine accelerates transformer workloads by combining fused GPU kernels, BF16 computation, and hardware-aware FP8 execution. We begin by installing Transformer Engine and detecting the active GPU architecture so that we can determine whether the runtime supports TE kernels, FP8 tensor cores, or only the pure-PyTorch fallback path. We then examine core fused components such as te.Linear, te.LayerNorm, te.LayerNormLinear, te.LayerNormMLP, and te.TransformerLayer, while also configuring a delayed-scaling FP8 recipe that manages tensor scaling, amax history, and hybrid E4M3/E5M2 formats. Using these components, we construct a compact GPT-style causal language model, train it on deterministic synthetic sequences, compare higher-precision and FP8 execution, measure runtime and peak GPU memory, inspect FP8 metadata, and validate the trained model through autoregressive generation.

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)

We install NVIDIA Transformer Engine and initialize the PyTorch environment required for GPU-accelerated execution. We inspect the active GPU, compute capability, and memory capacity to determine whether fused TE kernels and FP8 tensor cores are available. We also validate the core fused modules and configure a delayed-scaling FP8 recipe while preserving an automatic 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 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")

We define a compact causal language model using fused te.TransformerLayer blocks for Transformer Engine execution. We also implement an equivalent pure-PyTorch transformer architecture with multi-head attention, layer normalization, residual connections, and feed-forward networks. We select the appropriate model dynamically according to GPU support and report the final parameter count and architectural dimensions.

def make_batch(bsz=16): phase = torch.randint(0, VOCAB, (bsz, 1)) stride = torch.randint(1, 7, (bsz, 1)) steps = torch.arange(SEQ + 1).unsqueeze(0) seq = (phase + stride * steps) % VOCAB return seq[:, :-1].to(DEVICE), seq[:, 1:].to(DEVICE) opt = torch.optim.AdamW(model.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 = model(x) else: logits = model(x) loss = F.cross_entropy(logits.float .reshape(-1, VOCAB), y.reshape(-1)) opt.zero_grad(set_to_none=True) loss.backward opt.step return loss.item print(f"\n>> Training 60 steps ({'FP8' if FP8_CAPABLE else 'BF16/FP32'})...") t0 = time.time for step in range(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 would be ~{math.log(VOCAB):.2f})")

We create deterministic arithmetic-pattern sequences that allow the model to learn predictable token transitions across the vocabulary. We configure the AdamW optimizer and implement a training step that conditionally wraps the forward pass in te.fp8_autocast when FP8 execution is supported. We train the model for multiple iterations, monitor the loss and step latency, and compare the final loss against the random-guess baseline.