NVIDIA、AI計算を効率化

このチュートリアルでは、フレームワークの下層からcuDNN Frontendのgraph APIを詳しく見ていきます。計算を操作のグラフとして記述し、cuDNNにそれを実行するengineを選択させ、そしてその選択を自分たちで制御します。ここで構築するすべてのkernelは同じ方法で表現されます。次元とstrideによってtensorを宣言し、それらに操作を連結し、検証、操作グラフの構築、実行planの作成、サポートの確認、planの構築という5段階の構築pipelineを実行し、そしてpointerのvariant packに対して実行します。私たちはこれをすべて単一のColab GPUで実行し、各結果をPyTorchのリファレンスと照合することで、fusionが正しいことと、そのコストの両方を確認できます。トピックは互いに積み重なり、単一のfused convolutionから、engine config間のautotuning、FP8スタイルのepilogues、attention、planのserialization、dynamic shapes、そしてCUDA graph captureへと進みます。import os

import sys

import glob

import math

import time

import ctypes

import traceback

import subprocess

RESULTS = {}

def banner(title):

print("\n" + "=" * 78)

print(title)

print("=" * 78)

def section(name):

def wrap(fn):

def run(*a, **kw):

banner(name)

try:

out = fn(*a, **kw)

RESULTS[name] = out if isinstance(out, str) else "ok"

return out

except Exception as e:

RESULTS[name] = f"SKIPPED / FAILED -> {type(e).__name__}: {e}"

print(f"\n[!] {name} did not complete: {type(e).__name__}: {e}")

traceback.print_exc(limit=3)

return None

return run

return wrap

banner("0. Install nvidia-cudnn-frontend and locate libcudnn")

subprocess.run(

[sys.executable, "-m", "pip", "install", "-q", "nvidia-cudnn-frontend"],

check=True,

)

import torch

assert torch.cuda.is_available(), "No GPU. Runtime -> Change runtime type -> GPU."

torch.backends.cudnn.enabled = True

_ = torch.nn.functional.conv2d(

torch.randn(1, 1, 8, 8, device="cuda"), torch.randn(1, 1, 3, 3, device="cuda")

)

torch.cuda.synchronize()

try:

import nvidia.cudnn

_libdir = os.path.join(os.path.dirname(nv