torch.compile(): 즉시 실행을 유지하는 그래프 최적화

PyTorch 2.x를 대표하는 기능

torch.compile(model)은 모델을 최적화된 융합 커널로 JIT 컴파일하는 한 줄짜리 변환이야. CUDA에서는 보통 1.5~3배 빨라지고 Transformer 계열 모델은 더 큰 이득을 보기도 해. 컴파일하지 않은 영역에서는 즉시 실행 모드의 디버깅 편의성을 그대로 유지한다는 점이 특히 좋아.

내부에서 일어나는 일

TorchDynamo가 순전파의 Python 바이트코드를 가로채 연산을 FX 그래프로 포착해.
AOTAutograd가 역전파 연산까지 포함하도록 그래프를 다시 써.
TorchInductor가 그래프를 최적화된 Triton/CUDA 커널로 낮추고, CPU에서는 C++ 코드로 변환해.

첫 호출에서는 그래프를 만들고 컴파일해. 데이터에 따라 달라지는 제어 흐름이나 순전파 안의 사용자 정의 Python 객체처럼 Dynamo가 포착하지 못하는 코드가 나오면 그 구간만 즉시 실행으로 돌아가. 이를 '그래프 브레이크'라고 해. 컴파일 전체가 실패하는 건 아니지만 해당 구간의 속도 이득은 줄어들어.

세 가지 모드

기본값: 컴파일 시간과 실행 속도 사이의 균형이 좋아.
mode="reduce-overhead": 커널 실행 사이의 Python 추가 비용을 최소화해. 작은 모델이나 작은 배치에 잘 맞아.
mode="max-autotune": 가장 빠른 실행을 찾으려고 여러 커널 변형을 철저히 시험해. 컴파일에는 몇 분이 걸릴 수도 있지만 실행 속도는 가장 빠를 수 있어.

그래프 브레이크를 만드는 것

if x.sum() > 0처럼 텐서 값에 따라 달라지는 제어 흐름. Python 정수나 설정값에 따른 분기는 보통 괜찮아.
순전파 안에서 호출하는 일부 OpenCV나 PIL 연산처럼 추적할 수 없는 라이브러리 호출.
Python 속성을 통해 텐서 값을 변경하는 코드.
일부 사용자 정의 autograd 함수. PyTorch 2.x에서 계속 개선되고 있어.

Code

모델 컴파일: 한 줄·python

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(1024, 4096), nn.GELU(),
            nn.Linear(4096, 4096), nn.GELU(),
            nn.Linear(4096, 1024),
        )
    def forward(self, x):
        return self.layers(x)

model = MLP().cuda()
model = torch.compile(model)               # done.

x = torch.randn(64, 1024, device='cuda')
y = model(x)                                # first call: slow (compiles)
y = model(x)                                # subsequent: fast

모드 고르기·python

import torch

# Default — best for most workloads
model = torch.compile(model)

# Reduce overhead — when batch is small / model is light
model = torch.compile(model, mode="reduce-overhead")

# Max autotune — exhaustive search; slow compile, fastest runtime
model = torch.compile(model, mode="max-autotune")

그래프 브레이크 감지: 진단용 환경 변수·python

import os
import torch

# Set BEFORE importing torch (or just before compiling)
os.environ['TORCH_LOGS'] = 'graph_breaks'

@torch.compile
def forward(x, mask):
    # This .item() call CAUSES a graph break — the value goes to Python
    if mask.sum().item() > 0:
        return x * 2
    return x * -1

x = torch.randn(8)
mask = torch.tensor([1, 0, 1, 0, 1, 0, 1, 0])
forward(x, mask)
# Logs will show: 'Graph break: tensor.item()'

# Fix: use torch.where for tensor-valued conditions
@torch.compile
def forward_fixed(x, mask):
    return torch.where(mask.bool(), x * 2, x * -1)

학습에 컴파일을 끼워 넣기·python

import torch, torch.nn as nn
from torch.amp import autocast

model = MyModel().cuda()
model = torch.compile(model, mode="default")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()

for epoch in range(num_epochs):
    for x, y in loader:
        x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
        optimizer.zero_grad(set_to_none=True)
        with autocast(device_type='cuda', dtype=torch.bfloat16):
            out = model(x)
            loss = criterion(out, y)
        loss.backward()
        optimizer.step()

# torch.compile and bf16 autocast compose cleanly. So does DDP/FSDP.

Exercise

모델 하나를 골라 즉시 실행 모드와 torch.compile에서 순전파와 역전파 100단계의 시간을 각각 재 봐. 속도 향상 비율을 계산해. Ampere 이후 CUDA GPU에서 Transformer 계열 모델이라면 보통 1.5배 이상의 이득을 기대할 수 있어. 이득이 작다면 TORCH_LOGS='graph_breaks'를 설정해 그래프 브레이크의 원인을 찾아.