NVIDIAがAI学習を高速化

このチュートリアルでは、NVIDIA Transformer Engineが、fused GPU kernels、BF16計算、およびハードウェア対応のFP8実行を組み合わせることで、Transformerのワークロードをどのように高速化するかを探ります。まず、Transformer Engineをインストールし、アクティブなGPUアーキテクチャを検出して、ランタイムがTE kernels、FP8 tensor coresをサポートしているか、または純粋なPyTorchフォールバックパスのみをサポートしているかを判断します。次に、te.Linear、te.LayerNorm、te.LayerNormLinear、te.LayerNormMLP、te.TransformerLayerなどの主要なfused componentsを検証し、tensor scaling、amax history、およびハイブリッドE4M3/E5M2フォーマットを管理するdelayed-scaling FP8 recipeも設定します。これらのコンポーネントを使用して、コンパクトな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(), "ColabでまずGPUランタイムを有効にしてください！" 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_CAPA