torch.compiler 시작하기 — Dynamo·Inductor

torch.compiler 시작하기

torch.compile의 배후에는 두 컴포넌트가 있어요. TorchDynamo가 파이썬 프레임을 트레이스해서 그래프를 만들고, TorchInductor가 그 그래프를 Triton(또는 C++) 커널로 컴파일해요. 이 조합이 '파이썬을 안전하게, 그리고 빠르게' 만드는 핵심이에요.

간단한 시작

포인트와이즈 연산(예: sin/cos) 예제에서 시작해요. 컴파일된 결과물로 생성된 커널 코드는 이해하기 쉬운 형태로 출력돼요. 로그에서 torchinductor_<user> 폴더의 output_code.py에 생성 커널이 담겨요.

import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
opt_model = torch.compile(model, backend="inductor")
opt_model(torch.randn(1, 3, 64, 64, device="cuda"))

사전학습 모델 적용

Transformers·TIMM 같은 라이브러리의 사전학습 모델을 그대로 감싸 최적화할 수 있어요. TorchDynamo·TorchInductor의 설계 목표가 '외부 모델도 그대로 작동'하기 때문이에요.

from transformers import BertModel
model = BertModel.from_pretrained("bert-base-uncased").to("cuda:0")
model = torch.compile(model, backend="inductor")

더 알아보기