torch.compile(): TorchDynamo, TorchInductor și modurile de compilare în PyTorch & Deep Learning Practic

torch.compile compilează just-in-time codul PyTorch în kerneluri optimizate: TorchDynamo capturează graful din bytecode-ul Python, iar TorchInductor generează codul, cu Triton pe GPU-urile NVIDIA, AMD și Intel.

Sintaxă & invocare Python
import torch

model_compilat = torch.compile(model)                   # mode="default"
# sau pe loc, fără să schimbi structura modulului:
model.compile(mode="reduce-overhead")

iesiri = model_compilat(intrari)   # primele apeluri compilează, următoarele refolosesc codul

torch.compile, disponibil din PyTorch 2.0, urmărește codul Python și caută operațiile PyTorch. TorchDynamo capturează grafurile folosind Frame Evaluation API din CPython, AOT Autograd capturează și pasul înapoi, iar TorchInductor, backend-ul implicit, generează cod; pentru GPU-urile NVIDIA, AMD și Intel folosește Triton. Codul care nu poate fi urmărit produce un „graph break”: acea porțiune rulează normal în Python, iar optimizarea se pierde, fără erori.

Tutorialul oficial arată că primele execuții durează mai mult, pentru că atunci are loc compilarea; apoi codul compilat se refolosește. Câștigul vine în principal din reducerea overhead-ului Python și a citirilor/scrierilor în memoria GPU și depinde de arhitectură și de mărimea lotului: dacă GPU-ul e deja ocupat cu calcul, câștigul poate fi mic. Codul compilat este legat de forme și de alte condiții („guards”); la o schimbare se recompilează, până la limita recompile_limit (implicit 8), după care se revine la eager.

Modurile documentate sunt "default" (echilibru între performanță și overhead), "reduce-overhead" (CUDA Graphs, util la loturi mici; poate folosi mai multă memorie și se aplică doar grafurilor CUDA care nu modifică intrările), "max-autotune" (matmul-uri Triton sau bazate pe șabloane, alese prin profilare, convoluții Triton pe GPU și CUDA Graphs activate implicit pe GPU) și "max-autotune-no-cudagraphs". Pentru un graf exportabil și rulabil fără Python se folosește torch.export.

Concepte cheie și caracteristici

Concept / Parametru Descriere practică
mode "default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs".
fullgraph=False Cu True, orice graph break devine eroare (tot modelul trebuie să intre într-un graf).
dynamic=None None: detectează dinamismul și recompilează mai general; True: kerneluri dinamice de la început; False: mereu specializat.
backend="inductor" Backend-ul implicit; torch._dynamo.list_backends() listează alternativele.
options Dicționar de opțiuni pentru backend (de exemplu max_autotune, trace.enabled).
disable=False Face din torch.compile un no-op, pentru teste.
nn.Module.compile() Compilează modulul pe loc, fără să-i schimbe tipul.
@torch.compiler.disable Exclude o funcție de la compilare.

Exemple practice și scenarii reale

Exemplul 1: Măsurare corectă: compilarea se face la primele apeluri

import time

model = torch.nn.Sequential(torch.nn.Linear(64, 64), torch.nn.GELU(), torch.nn.Linear(64, 10))
compilat = torch.compile(model)
x = torch.randn(32, 64)

compilat(x)                                  # încălzire: compilare
start = time.perf_counter()
for _ in range(100):
    compilat(x)
print((time.perf_counter() - start) / 100)

Exemplul 2: Graph break-uri vizibile și interzise

torch._logging.set_logs(graph_breaks=True)   # afișează unde se rupe graful
strict = torch.compile(model, fullgraph=True) # eroare la primul graph break

Exemplul 3: O funcție exclusă de la compilare

@torch.compiler.disable
def jurnal(t):
    print("medie:", t.mean().item())
    return t

@torch.compile
def pas(x):
    return jurnal(x * 2) + 1

Greșeli frecvente și capcane

  • Să măsori viteza din primul apel. Primele execuții includ compilarea; măsori după încălzire.
  • Să schimbi des formele intrărilor (lungimi de secvență diferite la fiecare lot). Fiecare formă nouă poate declanșa o recompilare; după limită, PyTorch revine la eager.
  • Să folosești torch.compile pentru export sau servire fără Python. Este un compilator JIT; pentru artefacte salvate folosești torch.export și AOTInductor.

Noțiuni și module înrudite