Modurile de Compilare în PyTorch 2: default, reduce-overhead și max-autotune în PyTorch & Deep Learning Practic

torch.compile() utilizează tehnologiile TorchDynamo și TorchInductor pentru a analiza codul Python la execuție (JIT) și a genera kerneluri GPU fuzionate scrise în limbajul Triton. Comportamentul compilatorului este dictat de parametrul mode:

• default: Oferă un echilibru optim între timpul de compilare inițial și viteza de execuție, aplicând fuziune de kerneluri.
• reduce-overhead: Utilizează CUDA Graphs pentru a elimina overhead-ul de lansare a comenzilor din nucleul Python; ideal pentru loturi mici și inferență cu latență minimă.
• max-autotune: Rulează căutări empirice exhaustive pentru a selecta cele mai rapide kerneluri de multiplicare matricială și convoluție; timp de compilare ridicat, dar eficiență maximă în producție.

import torch

model = MyTransformerModel().cuda()

# Compilare cu optimizare agresivă pentru server de inferență
compiled_model = torch.compile(model, mode="reduce-overhead")

with torch.inference_mode():
    out = compiled_model(input_tensor)

Modelele compilate ating frecvent o creștere de viteză de 30% până la 100% fără modificări ale ponderilor sau arhitecturii.