추론을 위한 torch.compile

torch.compile은 PyTorch 코드를 최적화된 커널로 컴파일하여 추론 속도를 크게 향상시킵니다. 이 기능은 코드를 그래프로 컴파일하는 데 TorchDynamo를, 그래프를进一步优化된 커널로 컴파일하는 데 TorchInductor를 사용합니다. 강력한 최적화 도구이며, 많은 경우 코드 한 줄만 추가하면 됩니다.

모델을 torch.compile로 감싸서 컴파일하고 최적화된 모델을 반환합니다.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model)

모델을 컴파일해야 하므로 torch.compile에 대한 초기 호출은 느립니다. 다시 컴파일할 필요가 없으므로 컴파일된 모델에 대한 후속 호출은 훨씬 빠릅니다.

컴파일 과정을 커스터마이즈할 수 있는 여러 매개변수가 있습니다. 그중에서도 특히 중요한 두 가지는 다음과 같습니다. 전체 매개변수 목록은 torch.compile 문서를 참고하세요.

모드

mode 파라미터는 컴파일 시 사용할 수 있는 여러 성능 옵션을 제공합니다. 사용 사례에 가장 적합한 옵션을 찾기 위해 다양한 모드를 시도해 보세요.

• default은 속도와 메모리 사이의 균형 잡힌 옵션입니다.
• reduce-overhead는 메모리를 약간 더 사용하지만 Python 오버헤드를 줄여 더 빠른 성능을 낼 수 있습니다.
• max-autotune는 가장 빠른 속도를 제공하지만 컴파일 시간이 더 오래 걸립니다.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead")

Fullgraph

Fullgraph는 성능을 극대화하기 위해 전체 모델을 단일 그래프로 컴파일하려고 시도합니다. torch.compile은 그래프 브레이크(graph break)를 만나면 오류를 발생시키며, 이는 모델을 단일 그래프로 컴파일할 수 없음을 의미합니다.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

벤치마크

torch.compile를 활성화하고 비활성화했을 때, 다양한 GPU와 배치 크기에서 동일한 이미지로 다른 비전 작업을 수행할 때의 평균 추론 시간(밀리초)을 비교하는 성능 벤치마크는 아래 표를 참조하세요.

아래 표에서 Subset을 선택하여 다른 GPU로 전환하고, PyTorch nightly 2.1.0dev 및 reduce-overhead 모드가 활성화된 torch.compile에 대한 벤치마크를 확인하세요.

다음 단계

• torch.compile을 학습에 사용하는 방법을 배우려면 학습용 torch.compile 가이드를 참조하세요.