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.compilepentru export sau servire fără Python. Este un compilator JIT; pentru artefacte salvate foloseștitorch.exportși AOTInductor.