컴파일

개요

Pytorch 2.0은 PyTorch 코드를 최적화된 커널로 JIT 컴파일하여 PyTorch 코드 실행 속도를 높이는 강력한 기능인 torch.compile을 도입했습니다. torch.compile의 주요 기능은 다음과 같습니다:

• 성능 향상: 계산 그래프를 최적화하여 모델 실행 속도를 크게 높입니다.
• 사용 편의성: 구현을 위한 코드 변경이 최소로 필요하여 접근성이 매우 높습니다.
• 호환성: 기존 PyTorch 코드 및 모델과 원활하게 작동합니다.

Accelerate와 함께 사용할 때, torch.compile은 분산 학습 워크플로에 자연스럽게 통합되어 분산 실행과 컴파일 최적화의 이점을 동시에 누릴 수 있습니다.

컴파일된 코드의 첫 번째 실행은 컴파일 시간을 포함하여 일반적으로 더 오래 걸리지만, 이후 실행은 훨씬 빠릅니다. 다양한 시나리오에서 최적의 성능을 위해 torch.compile은 "default", "reduce-overhead" (CUDA graphs를 사용하여 오버헤드를 추가로 줄임), "max-autotune" (모델에 가장 적합한 커널을 찾기 위해 광범위한 자동 튜닝을 수행)와 같은 다양한 모드를 제공합니다.

Accelerate와 함께 torch.compile 사용

Accelerate는 torch.compile을 학습 스크립트에 쉽고 매끄럽게 통합할 수 있도록 TorchDynamoPlugin을 제공합니다.

from accelerate import Accelerator
from accelerate.utils import TorchDynamoPlugin

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    backend="inductor",  # Options: "inductor", "aot_eager", "aot_nvfuser", etc.
    mode="default",      # Options: "default", "reduce-overhead", "max-autotune"
    fullgraph=True,
    dynamic=False
)

# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply torch.compile to your model
model = accelerator.prepare(model)

Accelerate의 혼합 정밀도, 분산 학습(DDP, FSDP, Deepspeed)을 포함한 모든 다른 기능 및 플러그인과 호환됩니다.

지역별 컴파일

최적화 문제 공간이 일반적으로 큰 전체 모델을 컴파일하려는 시도를 피하고, 동일한 클래스의 반복되는 블록을 대상으로 하여 순차적으로 컴파일함으로써 컴파일러 캐시를 활용하는 것이 지역 컴파일입니다. 예를 들어, GPT2LMHeadModel에서 반복되는 블록/클래스는 GPT2Block이며, model.transformer.h[0]로 접근할 수 있습니다. 모델의 나머지 부분(예: model.lm_head)은 별도로 컴파일됩니다.

이를 통해 LLM과 Transformers와 같은 모델의 컴파일 오버헤드/콜드 스타트를 가속화할 수 있습니다. 자세한 내용은 https://pytorch.org/tutorials/recipes/regional_compilation.html을 참조하세요.

지역 컴파일 사용 방법

TorchDynamoPlugin 설정에서 use_regional_compilation=True를 설정하여 활성화할 수 있습니다:

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    use_regional_compilation=True,
    ... # other parameters
)
# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply compile_regions to your model
model = accelerator.prepare(model)

accelerate.utils.compile_regions 유틸리티를 torch.compile을 사용하는 것과 동일한 방식으로 직접 사용할 수도 있습니다.

지역 컴파일의 이점

PyTorch의 torch.compile 기능을 사용하여 전체 컴파일과 지역 컴파일을 비교하는 광범위한 벤치마크를 수행했습니다. 전체 결과는 accelerate 저장소에서 확인할 수 있습니다. 벤치마크의 주요 결과는 다음과 같습니다:

1. 비슷한 성능: 지역 컴파일은 전체 컴파일과 유사한 성능 가속을 제공하며, 특히 더 큰 모델에서 그렇습니다.
2. 더 빠른 컴파일: 지역별 컴파일은 모델 컴파일 시간을 크게 줄여 배포에 더 효율적인 선택지가 됩니다.
3. 배치 크기 영향: 배치 크기가 커질수록 컴파일 전략 간의 성능 차이가 줄어듭니다. 이는 해당 시나리오에서 컴파일 오버헤드의 영향이 상대적으로 작다는 것을 의미합니다.
4. 모델 크기 고려 사항: 지역 컴파일의 이점은 컴파일 시간 절감 효과가 상당할 수 있는 더 큰 모델에서 더 두드러집니다.
5. 실제 적용: 실제 환경에서는 특히 대규모 모델을 다룰 때 학습 콜드 스타트 시간을 최적화하기 위한 실용적인 선택으로 지역 컴파일을 사용할 수 있습니다.

결론

전체 컴파일과 지역 컴파일 모두 모델의 속도를 크게 향상시킬 수 있습니다. 지역 컴파일은 컴파일 시간과 런타임 성능 사이의 실용적인 균형을 제공하며, 특히 큰 배치 크기로 대규모 모델을 학습할 때 유용합니다.