AIを効率よく動かすxFormers解説

このチュートリアルでは、GPU上で高速でメモリ効率の良いTransformerモデルを構築するための実用的なツールキットであるxFormersを実装します。私たちはまず、メモリ効率の良いattentionを標準的なattentionの実装と比較して検証し、次に異なるシーケンス長での速度とメモリ消費量を比較します。次に、causal masking、packed variable-length sequences、grouped-query attention、およびカスタムのALiBi positional biasesについて検討します。最後に、これらの技術を組み合わせて、xFormers attention、SwiGLU feed-forward layers、およびautomatic mixed-precision trainingを使用する、学習可能なGPT-styleモデルを作成します。xFormersのセットアップとメモリ効率の良いAttentionの検証Copy Code CopiedUse a different Browserimport subprocess, sysdef _pip(*a): subprocess.run([sys.executable, "-m", "pip", "install", *a], check=False)try: import xformersexcept Exception: _pip("-q", "-U", "xformers")import math, timeimport torch, torch.nn as nn, torch.nn.functional as Fimport xformers, xformers.ops as xopsfrom xformers.ops import fmhaab = fmha.attn_biasassert torch.cuda.is_available(), ("GPUが検出されません。Colabでは: ランタイム → ランタイムのタイプを変更 → GPUを選択し、再実行してください。")device = "cuda"torch.manual_seed(0)print("torch :", torch.__version__)print("xformers :", xformers.__version__)print("GPU :", torch.cuda.get_device_name(0))print("\n--- xformers.info (どのカーネルがビルド/利用可能か) ---")try: subprocess.run([sys.executable, "-m", "xformers.info"], check=False)except Exception as e: print("xformers.infoが利用できません:", e)def cuda_time(fn, iters=20, warmup=5): for _ in range(warmup): fn() torch.cuda.synchronize() s, e = (torch.cuda.Event(enable_timing=True) for _ in range(2)) s.record() for _ in range(iters): fn() e.record(); torch.cuda.synchronize() return s.elapsed_time(e) / itersdef peak_mem_mb(fn): torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats() fn(); torch.cuda.synchronize() return torch.cuda.max_memory_allocated() / 1e6def vanilla_attention(q, k, v, causal=False):