학습을 위한 torch.compile

학습을 위한 torch.compile (torch.compile for training)

torch.compile은 PyTorch 코드를 퓨전된 커널로 컴파일해서 더 빠르게 실행되도록 해요. 학습에서는 forward pass와 backward pass를 함께 트레이싱(tracing)해서 최적화된 커널로 컴파일하므로, 개별 연산 실행(launch)의 오버헤드를 줄이고 연산을 퓨전해서 메모리 대역폭 사용을 줄여요.

출처: 문서

본문

TrainingArguments에서 torch_compile=True로 설정해서 활성화해요. 학습은 forward pass와 backward pass를 모두 컴파일하는데, 이는 forward pass만 컴파일하는 추론과 다르다는 점에 유의하세요. 컴파일은 첫 번째 학습 스텝에서 일어나므로, 이후 스텝보다 훨씬 느릴 거라고 예상해야 해요.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    torch_compile=True,
    torch_compile_backend="inductor",
    torch_compile_mode="reduce-overhead",
)

백엔드 (Backend)

백엔드를 지정하지 않으면 TrainingArguments가 하드웨어에 따라 하나를 선택해요. 대부분의 CPU와 GPU에서 기본값은 inductor로, AOTAutograd로 Triton 커널에 컴파일하며 대부분의 학습 워크로드에 적합해요. Intel Gaudi(HPU)에서는 기본값이 hpu_backend예요. AWS Trainium과 Inferentia(Neuron)에서는 기본값이 neuron이에요.

고정 형상(fixed-shape) 입력에는 cudagraphs를 사용해요.

컴파일 모드 (Compile mode)

아래 표를 참고해서 torch.compile 모드를 선택하세요.

mode description
default compile time과 runtime의 균형
reduce-overhead CUDA 그래프를 사용해 Python/CPU 오버헤드를 줄이고, 대신 약간의 추가 메모리를 사용
max-autotune 컴파일 시 여러 커널 구현을 벤치마킹하고 가장 빠른 것을 선택 (컴파일이 더 오래 걸림)
max-autotune-no-cudagraphs max-autotune과 같지만 CUDA 그래프는 사용하지 않음

다음 단계