Optimizare cu PyTorch 2.x: torch.compile() și Inductor în PyTorch & Deep Learning Practic

Cum accelerezi orice model PyTorch cu 20-50% printr-o singură linie de cod: compilatorul TorchDynamo, fuziunea de kernel-uri cu Triton și opțiunile de mod.

Sintaxă & invocare Python
import torch

model = ReteaTransformer()
# Compilarea modelului pentru accelerare maximă pe GPU (PyTorch 2.0+):
model_optimizat = torch.compile(model, mode="reduce-overhead")

# Prima rulare face compilarea JIT, următoarele vor rula la viteză maximă:
iesiri = model_optimizat(intrari)

Descriere & Utilizare

Piatra de hotar a lansării PyTorch 2.0 a fost introducerea compilatorului torch.compile(). Până la versiunea 2.0, PyTorch rula fiecare operațiune individual: o adunare de tensori urmată de o funcție de activare însemna două kernel-uri CUDA separate, cu date citite și rescrise de două ori în memoria VRAM.

torch.compile() utilizează TorchDynamo pentru a inspecta bytecode-ul Python și TorchInductor (bazat pe compilatorul OpenAI Triton) pentru a genera kernel-uri CUDA unificate (Kernel Fusion). Zece operațiuni consecutive sunt contopite într-o singură instrucțiune hardware care procesează datele direct în registrele rapide ale nucleelor plăcii video.

Cea mai spectaculoasă parte este că nu necesită modificarea arhitecturii: adaugi o singură linie de cod peste modelul tău existent și beneficiezi instantaneu de o creștere de viteză de 20% până la 100%!

Cum să o reții ușor

💡 Este ca diferența între un bucătar care se duce la cămară după fiecare ingredient în parte (modul clasic) și un bucătar care își aduce toate ingredientele pe blat dintr-un singur drum și le gătește în aceeași tigaie încinsă (kernel fusion).

Opțiuni și parametri tehnici

  • mode="default": echilibru bun între timpul inițial de compilare și viteza de rulare.
  • mode="reduce-overhead": folosește CUDA Graphs pentru a elimina latența de apel a driverului din Python; ideal pentru modele mici/medii sau inferență cu batch mic.
  • mode="max-autotune": cel mai lent la pornire (durează câteva minute pentru a testa zeci de kernel-uri Triton), dar oferă viteza maximă absolută pentru antrenări lungi de zile întregi.
  • fullgraph=False: permite revenirea automată la Python standard (graph break) dacă întâlnește biblioteci nestandardizate.

Exemple practice și cazuri de utilizare

  • Compilare rapidă pentru fine-tuning: model = torch.compile(model).
  • Optimizare pentru inferență de mare viteză: model = torch.compile(model, mode="reduce-overhead").
  • Dezactivarea compilării pe o funcție problematică: decoreaz-o cu @torch.compiler.disable.

Sfaturi de configurare & capcane

Reține că primul pas de execuție va fi lent (poate dura între 10 secunde și 2 minute în funcție de mărimea modelului), deoarece atunci are loc compilarea efectivă a kernel-urilor Triton — măsoară întotdeauna viteza începând de la al doilea pas!

Comenzi și ghiduri înrudite