Transformer训练加速：GPU融合与精度优化

人工智能

本教程将探讨 NVIDIA Transformer Engine 如何通过融合 GPU 内核、BF16 计算以及针对硬件优化的 FP8 执行来加速 Transformer 工作负载。首先，我们会安装 Transformer Engine，并检测当前使用的 GPU 架构，以确定运行环境支持 TE 内核、FP8 Tensor Core，还是只能使用纯 PyTorch 回退路径。随后，我们将研究 te.Linear、te.LayerNorm、te.LayerNormLinear、te.LayerNormMLP 和 te.TransformerLayer 等核心融合组件，同时配置一种延迟缩放 FP8 recipe，用于管理张量缩放、amax 历史记录以及混合 E4M3/E5M2 格式。利用这些组件，我们将构建一个紧凑的 GPT 风格因果语言模型，使用确定性的合成序列进行训练，对比高精度与 FP8 执行，测量运行时间和 GPU 峰值显存，检查 FP8 元数据，并通过自回归生成验证训练后的模型。

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)

我们安装 NVIDIA Transformer Engine，并初始化 GPU 加速执行所需的 PyTorch 环境。通过检查当前 GPU、计算能力和显存容量，我们可以判断融合 TE 内核与 FP8 Tensor Core 是否可用。我们还会验证核心融合模块，并配置延迟缩放 FP8 recipe，同时为不受支持的硬件保留自动切换至 PyTorch 的回退机制。

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")

我们使用融合的 te.TransformerLayer block 定义了一个紧凑的因果语言模型，以便通过 Transformer Engine 执行。同时，我们还实现了一个等价的纯 PyTorch Transformer 架构，其中包含多头注意力、Layer Normalization、残差连接和前馈网络。程序会根据 GPU 的支持情况动态选择合适的模型，并输出最终的参数量和架构维度。

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})")

我们创建具有确定性算术规律的序列，让模型能够学习整个词表中可预测的 token 转移。接着配置 AdamW 优化器，并实现一个训练步骤：当硬件支持 FP8 执行时，使用 te.fp8_autocast 包裹前向传播。我们对模型进行多轮迭代训练，监控 loss 和每一步的延迟，并将最终 loss 与随机猜测的基线进行比较。

def bench(use_fp8, iters=30, warmup=10):
   x, y = make_batch(bsz=32)
   for _ in range(warmup):
       run_step(x, y, use_fp8)
   torch.cuda.synchronize()
   torch.cuda.reset_peak_memory_stats()
   t = time.time()
   for _ in range(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"(gains grow with model size — try D_MODEL=2048, N_LAYERS=12)")
else:
   print("   FP8 benchmark skipped — needs an sm_89+ GPU (L4/H100/Ada/Blackwell).")
if FP8_CAPABLE:
   blk = model.blocks[0]
   for name, 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 '{name}':")
           print("   scale       :", s.scale.flatten()[:4].tolist())
           print("   amax_history:", s.amax_history[0, :4].tolist())
           break

我们分别在高精度和 FP8 执行模式下，对前向传播、反向传播以及优化器更新进行基准测试。通过测量平均训练步骤延迟和 GPU 已分配显存峰值，量化低精度计算对性能与显存占用的影响。我们还会检查 Transformer Engine 维护的缩放因子和 amax 历史记录，以了解延迟缩放如何帮助稳定 FP8 张量。

@torch.no_grad()
def generate(prompt_len=8, gen_len=24):
   x, _ = make_batch(bsz=1)
   ctx = x[:, :prompt_len]
   for _ in range(gen_len):
       inp = ctx[:, -SEQ:]
       if TE_CAPABLE and FP8_CAPABLE:
           with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
               logits = model(inp)
       else:
           logits = model(inp)
       nxt = logits[:, -1].argmax(-1, keepdim=True)
       ctx = torch.cat([ctx, nxt], dim=1)
   return ctx[0].tolist()
seq = generate()
print("\n>> Greedy generation (should continue the arithmetic pattern):")
print("   prompt+gen:", seq)
diffs = [(b - a) % VOCAB for a, b in zip(seq, seq[1:])]
print("   step diffs:", diffs, "<- constant stride = model learned the rule")
print("\n>> Done! Things to try next:")
print("   * Scale up: D_MODEL=2048, N_LAYERS=12 -> FP8 speedup becomes dramatic")
print("   * recipe.Format.E4M3 vs HYBRID; amax_history_len=1024")
print("   * te.LayerNormMLP / te.LayerNormLinear in your own architectures")
print("   * fp8_model_init() to store weights themselves in FP8 for inference")

我们实现了贪心自回归生成：反复将最新的上下文输入训练后的因果语言模型。通过比较连续生成的 token，我们可以验证模型是否保持了合成训练数据中的固定算术步长。最后，我们列出了一些实际可行的扩展方向，包括增大模型维度、采用其他 FP8 格式、使用更长的 amax 历史记录、引入融合模块，以及初始化 FP8 权重。

总而言之，我们展示了如何将 NVIDIA Transformer Engine 集成到端到端的 Transformer 训练工作流中，同时保持对不同 Colab GPU 环境的兼容性。我们使用融合 Transformer 模块减少内核启动开销和显存流量；在硬件支持时，通过延迟缩放应用 FP8 autocasting；并借助自动 PyTorch 回退机制保留 BF16 或 FP32 执行路径。通过训练同一个迷你因果语言模型并进行基准测试，我们观察了硬件能力、数值格式、融合执行方式以及模型规模如何影响训练速度和显存消耗。我们还检查了支持稳定 FP8 计算的内部缩放因子与 amax 历史记录，从而更清楚地理解 Transformer Engine 如何管理低精度运算。

欢迎查看完整代码。也欢迎在 Twitter 上关注我们，别忘了加入我们拥有超过 15 万名成员的机器学习 SubReddit，并订阅 Newsletter。等等！你使用 Telegram 吗？现在也可以加入我们的 Telegram。

希望与我们合作，推广你的 GitHub Repo、Hugging Face Page、产品发布或 Webinar 等内容？欢迎联系我们。

Sana Hassan 是 MarkTechPost 的咨询实习生，也是 IIT Madras 的双学位学生。他热衷于运用技术和 AI 解决现实世界中的挑战。凭借对解决实际问题的浓厚兴趣，他为 AI 与现实生活解决方案的交汇领域带来了全新的视角。