|

An Implementation Guide to Running NVIDIA Transformer Engine with Mixed Precision, FP8 Checks, Benchmarking, and Fallback Execution

ByRicardo

In this tutorial, we implement a complicated, sensible implementation of the NVIDIA Transformer Engine in Python, specializing in how mixed-precision acceleration might be explored in a practical deep studying workflow. We arrange the surroundings, confirm GPU and CUDA readiness, try to set up the required Transformer Engine elements, and deal with compatibility points gracefully in order that the pocket book stays runnable even when the total extension can’t be constructed. As we transfer by every step, we construct instructor and scholar networks, examine a baseline PyTorch path with a Transformer Engine-enabled path, prepare each fashions, benchmark their pace and reminiscence utilization, and visualize the outcomes, giving us a transparent hands-on understanding of how performance-oriented coaching workflows are structured in apply.

Copy Code
import os
import sys
import json
import time
import math
import random
import shutil
import platform
import subprocess
import statistics


def run(cmd, test=True):
   print("n[RUN]", " ".be part of(cmd))
   outcome = subprocess.run(cmd, textual content=True, capture_output=True)
   if outcome.stdout.strip():
       print(outcome.stdout[-4000:])
   if outcome.returncode != 0 and outcome.stderr.strip():
       print(outcome.stderr[-4000:])
   if test and outcome.returncode != 0:
       elevate subprocess.CalledProcessError(outcome.returncode, cmd)
   return outcome


def has_cmd(title):
   return shutil.which(title) just isn't None


run([sys.executable, "-m", "pip", "install", "-q", "--upgrade", "pip"])
run([sys.executable, "-m", "pip", "install", "-q", "ninja", "packaging", "matplotlib"])


import torch
import torch.nn as nn
import torch.nn.purposeful as F
import matplotlib.pyplot as plt


assert torch.cuda.is_available(), "This pocket book wants a GPU runtime in Colab."


gpu_name = torch.cuda.get_device_name(0)
cc_major, cc_minor = torch.cuda.get_device_capability(0)
cuda_runtime = torch.model.cuda
python_version = sys.model.break up()[0]
torch_version = torch.__version__
cuda_home = os.environ.get("CUDA_HOME", "/usr/native/cuda")
nvcc_path = shutil.which("nvcc") or os.path.be part of(cuda_home, "bin", "nvcc")
cudnn_header_candidates = [
   os.path.join(cuda_home, "include", "cudnn.h"),
   "/usr/include/cudnn.h",
   "/usr/local/include/cudnn.h",
]


nvcc_exists = os.path.exists(nvcc_path)
cudnn_header_exists = any(os.path.exists(p) for p in cudnn_header_candidates)


print("=" * 120)
print("ENVIRONMENT CHECK")
print("=" * 120)
print(json.dumps({
   "python": python_version,
   "platform": platform.platform(),
   "torch": torch_version,
   "torch_cuda": cuda_runtime,
   "gpu_name": gpu_name,
   "compute_capability": f"{cc_major}.{cc_minor}",
   "cuda_home": cuda_home,
   "nvcc_exists": nvcc_exists,
   "nvcc_path": nvcc_path if nvcc_exists else None,
   "cudnn_header_exists": cudnn_header_exists,
}, indent=2))
print("=" * 120)

We put together the Colab surroundings by importing the required Python libraries, defining a helper perform for executing shell instructions, and putting in the core dependencies for the tutorial. We then import PyTorch and Matplotlib, confirm {that a} GPU is out there, and acquire key surroundings particulars, together with the GPU title, CUDA model, Python model, and toolkit paths. This provides us a transparent view of the system state earlier than we try any Transformer Engine set up or mannequin execution.

Copy Code
te_available = False
te_mode = "fallback"
te_import_error = None


strive:
   run([sys.executable, "-m", "pip", "install", "-q", "transformer_engine[core_cu12]"])
besides Exception as e:
   print("Core wheel set up failed:", repr(e))


can_try_te_torch = nvcc_exists and cudnn_header_exists


if can_try_te_torch:
   env = os.environ.copy()
   env["NVTE_FRAMEWORK"] = "pytorch"
   env["MAX_JOBS"] = "1"
   env["NVTE_BUILD_THREADS_PER_JOB"] = "1"
   env["CUDA_PATH"] = cuda_home
   env["CUDA_HOME"] = cuda_home
   strive:
       print("nAttempting to construct the PyTorch extension for Transformer Engine...")
       outcome = subprocess.run(
           [sys.executable, "-m", "pip", "install", "-q", "--no-build-isolation", "transformer_engine[pytorch]"],
           textual content=True,
           capture_output=True,
           env=env,
       )
       if outcome.stdout.strip():
           print(outcome.stdout[-4000:])
       if outcome.returncode != 0 and outcome.stderr.strip():
           print(outcome.stderr[-4000:])
       if outcome.returncode == 0:
           import transformer_engine.pytorch as te
           from transformer_engine.widespread import recipe
           te_available = True
           te_mode = "transformer_engine"
       else:
           te_import_error = outcome.stderr[-4000:] if outcome.stderr else "Unknown pip construct error"
   besides Exception as e:
       te_import_error = repr(e)
else:
   te_import_error = "Missing nvcc or cuDNN headers on this Colab runtime, so TE PyTorch extension can't be constructed right here."


if te_available:
   strive:
       fp8_available, fp8_reason = te.is_fp8_available(return_reason=True)
   besides Exception as e:
       fp8_available, fp8_reason = False, f"Could not question FP8 availability: {e}"
   strive:
       bf16_available = te.is_bf16_available()
   besides Exception:
       bf16_available = torch.cuda.is_bf16_supported()
else:
   fp8_available = False
   fp8_reason = "Transformer Engine not put in; utilizing fallback PyTorch path."
   bf16_available = torch.cuda.is_bf16_supported()


amp_dtype = torch.bfloat16 if bf16_available else torch.float16


print("n" + "=" * 120)
print("INSTALL STATUS")
print("=" * 120)
print(json.dumps({
   "te_available": te_available,
   "te_mode": te_mode,
   "fp8_available": fp8_available,
   "fp8_reason": fp8_reason,
   "te_import_error": te_import_error,
   "amp_dtype": str(amp_dtype),
}, indent=2))
print("=" * 120)


machine = "cuda"
random.seed(42)
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)


if te_available:
   fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3)


def baseline_autocast():
   return torch.autocast(device_type="cuda", dtype=amp_dtype)


def te_forward_context(use_fp8):
   if te_available and use_fp8:
       return te.autocast(enabled=True, recipe=fp8_recipe)
   return baseline_autocast()

We try to set up the Transformer Engine core bundle and then test whether or not the Colab runtime can construct the PyTorch extension by verifying the presence of nvcc and cuDNN headers. If the surroundings helps it, we strive to set up the Transformer Engine PyTorch backend and then examine whether or not FP8 and BF16 can be found on the present {hardware}. We additionally configure the precision mode and outline the autocast contexts that later enable us to change between normal blended precision and Transformer Engine execution.

Copy Code
class TeacherWeb(nn.Module):
   def __init__(self, hidden_size=512, intermediate_size=2048, num_layers=3, vocab_size=4096):
       tremendous().__init__()
       self.embed = nn.Embedding(vocab_size, hidden_size)
       self.layers = nn.ModuleRecord([
           nn.Sequential(
               nn.LayerNorm(hidden_size),
               nn.Linear(hidden_size, intermediate_size),
               nn.GELU(),
               nn.Linear(intermediate_size, hidden_size),
           ) for _ in range(num_layers)
       ])
       self.head = nn.Linear(hidden_size, hidden_size)


   def ahead(self, token_ids):
       x = self.embed(token_ids)
       for layer in self.layers:
           x = x + layer(x)
       return self.head(x)


class BaselineScholar(nn.Module):
   def __init__(self, hidden_size=512, intermediate_size=2048, num_layers=3, vocab_size=4096):
       tremendous().__init__()
       self.embed = nn.Embedding(vocab_size, hidden_size)
       self.norms = nn.ModuleRecord([nn.LayerNorm(hidden_size) for _ in range(num_layers)])
       self.fc1 = nn.ModuleRecord([nn.Linear(hidden_size, intermediate_size) for _ in range(num_layers)])
       self.fc2 = nn.ModuleRecord([nn.Linear(intermediate_size, hidden_size) for _ in range(num_layers)])
       self.head = nn.Linear(hidden_size, hidden_size)


   def ahead(self, token_ids):
       x = self.embed(token_ids)
       for ln, fc1, fc2 in zip(self.norms, self.fc1, self.fc2):
           residual = x
           x = ln(x)
           x = fc1(x)
           x = F.gelu(x, approximate="tanh")
           x = fc2(x)
           x = x + residual
       return self.head(x)


if te_available:
   class TEStudent(nn.Module):
       def __init__(self, hidden_size=512, intermediate_size=2048, num_layers=3, vocab_size=4096):
           tremendous().__init__()
           self.embed = nn.Embedding(vocab_size, hidden_size)
           self.norms = nn.ModuleRecord([te.LayerNorm(hidden_size) for _ in range(num_layers)])
           self.fc1 = nn.ModuleRecord([te.Linear(hidden_size, intermediate_size, bias=True) for _ in range(num_layers)])
           self.fc2 = nn.ModuleRecord([te.Linear(intermediate_size, hidden_size, bias=True) for _ in range(num_layers)])
           self.head = te.Linear(hidden_size, hidden_size, bias=True)


       def ahead(self, token_ids, use_fp8=False):
           x = self.embed(token_ids)
           with te_forward_context(use_fp8):
               for ln, fc1, fc2 in zip(self.norms, self.fc1, self.fc2):
                   residual = x
                   x = ln(x)
                   x = fc1(x)
                   x = F.gelu(x, approximate="tanh")
                   x = fc2(x)
                   x = x + residual
               x = self.head(x)
           return x
else:
   class TEStudent(nn.Module):
       def __init__(self, hidden_size=512, intermediate_size=2048, num_layers=3, vocab_size=4096):
           tremendous().__init__()
           self.embed = nn.Embedding(vocab_size, hidden_size)
           self.norms = nn.ModuleRecord([nn.LayerNorm(hidden_size) for _ in range(num_layers)])
           self.fc1 = nn.ModuleRecord([nn.Linear(hidden_size, intermediate_size) for _ in range(num_layers)])
           self.fc2 = nn.ModuleRecord([nn.Linear(intermediate_size, hidden_size) for _ in range(num_layers)])
           self.head = nn.Linear(hidden_size, hidden_size)


       def ahead(self, token_ids, use_fp8=False):
           x = self.embed(token_ids)
           with baseline_autocast():
               for ln, fc1, fc2 in zip(self.norms, self.fc1, self.fc2):
                   residual = x
                   x = ln(x)
                   x = fc1(x)
                   x = F.gelu(x, approximate="tanh")
                   x = fc2(x)
                   x = x + residual
               x = self.head(x)
           return x


def count_params(mannequin):
   return sum(p.numel() for p in mannequin.parameters() if p.requires_grad)


def format_millions(n):
   return f"{n / 1e6:.2f}M"

We outline the neural community architectures used all through the tutorial, together with the instructor mannequin, the baseline scholar mannequin, and the Transformer Engine scholar path. We hold the mannequin constructions aligned in order that the comparability stays significant whereas permitting the TE path to swap in Transformer Engine layers when the extension is out there. We additionally outline small utility features for counting parameters and formatting mannequin measurement, which assist us examine the size of the fashions earlier than coaching begins.

Copy Code
hidden_size = 512
intermediate_size = 2048
num_layers = 3
vocab_size = 4096
seq_len = 128
batch_size = 8
steps = 25
benchmark_iters = 20
lr = 2e-4
weight_decay = 1e-2


instructor = TeacherWeb(hidden_size, intermediate_size, num_layers, vocab_size).to(machine).eval()
baseline_model = BaselineScholar(hidden_size, intermediate_size, num_layers, vocab_size).to(machine)
te_model = TEStudent(hidden_size, intermediate_size, num_layers, vocab_size).to(machine)


optimizer_baseline = torch.optim.AdamW(baseline_model.parameters(), lr=lr, weight_decay=weight_decay)
optimizer_te = torch.optim.AdamW(te_model.parameters(), lr=lr, weight_decay=weight_decay)


print("Teacher params :", format_millions(count_params(instructor)))
print("Baseline params:", format_millions(count_params(baseline_model)))
print("TE-path params :", format_millions(count_params(te_model)))


def make_batch(batch_size, seq_len, vocab_size, machine):
   tokens = torch.randint(0, vocab_size, (batch_size, seq_len), machine=machine)
   with torch.no_grad():
       goal = instructor(tokens)
   return tokens, goal


def peak_mem_mb():
   return torch.cuda.max_memory_allocated() / (1024 ** 2)


def train_baseline_step():
   baseline_model.prepare()
   optimizer_baseline.zero_grad(set_to_none=True)
   tokens, goal = make_batch(batch_size, seq_len, vocab_size, machine)
   with baseline_autocast():
       pred = baseline_model(tokens)
       loss = F.mse_loss(pred, goal)
   loss.backward()
   optimizer_baseline.step()
   return float(loss.detach().merchandise())


def train_te_step(use_fp8):
   te_model.prepare()
   optimizer_te.zero_grad(set_to_none=True)
   tokens, goal = make_batch(batch_size, seq_len, vocab_size, machine)
   pred = te_model(tokens, use_fp8=use_fp8)
   loss = F.mse_loss(pred, goal)
   loss.backward()
   optimizer_te.step()
   return float(loss.detach().merchandise())

We set the principle experiment hyperparameters, instantiate all fashions on the GPU, and create the optimizers that might be used throughout coaching. We additionally print the parameter counts to affirm that the baseline and TE paths are comparable when it comes to mannequin measurement. In addition, we outline the batch-generation logic, reminiscence monitoring perform, and the person training-step features that execute one optimization step for every mannequin path.

Copy Code
baseline_losses = []
te_losses = []
mode_name = "TE-FP8" if (te_available and fp8_available) else ("TE-BF16/FP16" if te_available else "Fallback-PyTorch")


print("n" + "=" * 120)
print("TRAINING")
print("=" * 120)


for step in vary(1, steps + 1):
   b_loss = train_baseline_step()
   t_loss = train_te_step(use_fp8=fp8_available)
   baseline_losses.append(b_loss)
   te_losses.append(t_loss)
   if step == 1 or step % 5 == 0 or step == steps:
       print(f"step={step:02d} | baseline_loss={b_loss:.6f} | te_path_loss={t_loss:.6f} | mode={mode_name}")


@torch.no_grad()
def evaluate_model(mannequin, is_te=False, use_fp8=False, eval_batches=8):
   mannequin.eval()
   vals = []
   for _ in vary(eval_batches):
       tokens, goal = make_batch(batch_size, seq_len, vocab_size, machine)
       if is_te:
           pred = mannequin(tokens, use_fp8=use_fp8)
       else:
           with baseline_autocast():
               pred = mannequin(tokens)
       vals.append(float(F.mse_loss(pred, goal).merchandise()))
   return sum(vals) / len(vals)


baseline_eval = evaluate_model(baseline_model, is_te=False)
te_eval = evaluate_model(te_model, is_te=True, use_fp8=fp8_available)


def benchmark_train_step(mannequin, optimizer, is_te=False, use_fp8=False, warmup=5, iters=20):
   times_ms = []
   mems_mb = []
   for _ in vary(warmup):
       optimizer.zero_grad(set_to_none=True)
       tokens, goal = make_batch(batch_size, seq_len, vocab_size, machine)
       if is_te:
           pred = mannequin(tokens, use_fp8=use_fp8)
       else:
           with baseline_autocast():
               pred = mannequin(tokens)
       loss = F.mse_loss(pred, goal)
       loss.backward()
       optimizer.step()
   torch.cuda.synchronize()
   for _ in vary(iters):
       torch.cuda.reset_peak_memory_stats()
       optimizer.zero_grad(set_to_none=True)
       tokens, goal = make_batch(batch_size, seq_len, vocab_size, machine)
       begin = time.perf_counter()
       if is_te:
           pred = mannequin(tokens, use_fp8=use_fp8)
       else:
           with baseline_autocast():
               pred = mannequin(tokens)
       loss = F.mse_loss(pred, goal)
       loss.backward()
       optimizer.step()
       torch.cuda.synchronize()
       finish = time.perf_counter()
       times_ms.append((finish - begin) * 1000.0)
       mems_mb.append(peak_mem_mb())
   return {
       "mean_ms": statistics.imply(times_ms),
       "median_ms": statistics.median(times_ms),
       "max_memory_mb": max(mems_mb),
   }


baseline_bench = benchmark_train_step(baseline_model, optimizer_baseline, is_te=False, use_fp8=False, iters=benchmark_iters)
te_bench = benchmark_train_step(te_model, optimizer_te, is_te=True, use_fp8=fp8_available, iters=benchmark_iters)

We run the principle coaching loop for each the baseline mannequin and the TE path, monitoring their losses over a number of steps. We then outline and execute the analysis perform to measure how nicely every mannequin matches the instructor’s outputs after coaching. Finally, we implement the benchmarking routine to measure per-step runtime and peak CUDA reminiscence utilization, enabling quantitative comparability of efficiency traits.

Copy Code
abstract = {
   "gpu_name": gpu_name,
   "compute_capability": f"{cc_major}.{cc_minor}",
   "te_available": te_available,
   "fp8_available": fp8_available,
   "fp8_reason": fp8_reason,
   "mode": mode_name,
   "baseline_eval_mse": baseline_eval,
   "te_path_eval_mse": te_eval,
   "baseline_mean_step_ms": baseline_bench["mean_ms"],
   "te_path_mean_step_ms": te_bench["mean_ms"],
   "baseline_peak_mem_mb": baseline_bench["max_memory_mb"],
   "te_path_peak_mem_mb": te_bench["max_memory_mb"],
}


print("n" + "=" * 120)
print("SUMMARY")
print("=" * 120)
print(json.dumps(abstract, indent=2))


plt.determine(figsize=(10, 5))
plt.plot(baseline_losses, label="Baseline loss")
plt.plot(te_losses, label=f"{mode_name} loss")
plt.xlabel("Training step")
plt.ylabel("MSE loss")
plt.title("Training Loss Comparison")
plt.legend()
plt.grid(True)
plt.present()


plt.determine(figsize=(8, 5))
plt.bar(["Baseline", mode_name], [baseline_bench["mean_ms"], te_bench["mean_ms"]])
plt.ylabel("Mean prepare step time (ms)")
plt.title("Speed Comparison")
plt.grid(True, axis="y")
plt.present()


plt.determine(figsize=(8, 5))
plt.bar(["Baseline", mode_name], [baseline_bench["max_memory_mb"], te_bench["max_memory_mb"]])
plt.ylabel("Peak reminiscence (MB)")
plt.title("Peak CUDA Memory Comparison")
plt.grid(True, axis="y")
plt.present()

We collect all remaining metrics right into a abstract dictionary and print the experiment’s consolidated leads to a structured format. We then generate visualizations of coaching loss, imply training-step time, and peak reminiscence utilization to extra intuitively interpret the variations between the baseline and TE paths. This remaining part helps us transfer from uncooked numbers to sensible insights by displaying how the 2 implementations behave throughout accuracy, pace, and reminiscence.

In conclusion, we constructed way over a easy set up walkthrough; we created a whole experimental pipeline that helps us perceive how the NVIDIA Transformer Engine matches into fashionable GPU-accelerated mannequin coaching. We examined the runtime surroundings, tailored to Colab limitations, preserved a working fallback path, and then skilled, evaluated, and benchmarked two implementations facet by facet to observe sensible variations in effectivity, precision habits, and useful resource utilization. At the top, we understood how to use the Transformer Engine in a Colab-friendly setting and gained a reusable basis that we will prolong to bigger transformer architectures, richer benchmarking situations, and extra production-oriented optimization workflows.

Check out the Full Codes/Notebook here.  Also, be happy to observe us on Twitter and don’t overlook to be part of our 120k+ 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 so on.? Connect with us

The submit An Implementation Guide to Running NVIDIA Transformer Engine with Mixed Precision, FP8 Checks, Benchmarking, and Fallback Execution appeared first on MarkTechPost.

Similar Posts