[결과]: 단순히 torch.compile으로 wrapping해준 것만으로 모델 Training이 43%빠른속도로 동작했다. (다만, 이는 A100으로 측정된 결과이고, 3090같은 시리즈는 잘 동작하지 않고 심지어 더 느릴 수 있다 언급: Caveats: On a desktop-class GPU such as a NVIDIA 3090, we’ve measured that speedups are lower than on server-class GPUs such as A100. As of today, our default backend TorchInductor supports CPUs and NVIDIA Volta and Ampere GPUs. It does not (yet) support other GPUs, xPUs or older NVIDIA GPUs.)
이를 2.2버전에서는 좀 더 완성시킨 것이다!
pytorch개발자분들은 버전이 2.x로 넘어가면서 compile함수에 좀 더 집중한다하였다. 아마 점점 학습속도를 빠르게하는 면을 강화하고, 이를 점차 확대할 것 같다. (이번에 저수준커널에도 적용한 걸 보면 거의 확실시 되는듯하다.)
개발동기:
17년 시작된 이후, Eager Execution성능향상을 위해 코드 대부분을 C++로 옮기게 되었다. (Pytorch 대부분의 소스코드가 C++기반임을 근거로 알 수 있다.) (eager execution: 그래프생성없이 연산을 즉시실행하는 환경)
이런 방식을 사용자들의 코드기여도(hackability)를 낮추는 진입장벽이 되어버렸다. 이런 eager execution의 성능향상에 한계가 있다 판단하여 compiler를 만들게 되었다. 목적은 속도는 빠르게하나 pytorch experience를 해치지 않는다는 것이다.
🤔 How to use?
torch.compile()은 기존 모델에 한줄만 추가하면 된다.
import torch
import torchvision.models as models
model = models.resnet18().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
compiled_model = torch.compile(model)
x = torch.randn(16, 3, 224, 224).cuda()
optimizer.zero_grad()
out = compiled_model(x)
out.sum().backward()
optimizer.step()
compiled_model은 신경망의 forward를 좀 더 최적화시켜 속도를 빠르게 한다.
[default]: 너무 오래걸리지 않으면서 메모리를 많이 사용하지 않는 선에서 효율적인 컴파일 진행
[reduce-overhead]: 메모리를 좀 더 사용, overhead를 줄여줌
[max-autotune]: 가장 빠른 모델생성을 위해 최적화되어있다. 다만, 컴파일에 매우 오랜시간이 걸린다.
∙ dynamic:
dynamic shape에 대해 code path를 enabling할 지 결정하는 boolean 변수이다. Compiler 최적화 과정이 프로그램을 dynamic shape 프로그램에 적용될 수 없게 만드는 경우가 있는데, 이를 조절함으로써 본인이 원하는 방향대로 컴파일을 할 수 있게 해준다. 이 부분은 아직 완벽히 이해가 되지는 않지만 데이터 shape가 변하는 상황에서graph를 유동적으로 컴파일할 수 있게끔 하는 것과 관련이 있을 것 같다.
∙ fullgraph:
Numba의nopython과 유사하다. 전체 프로그램을 하나의 그래프로 컴파일하고, 만약 실패한다면 왜 불가능한지 설명해주는 error 메세지를 띄운다. 굳이 쓰지 않아도 상관없는 옵션.
∙ backend:
어떤 compiler backend를 적용할 지 결정하게 된다. 디폴트로 정해진 값은 앞서 설명했던TorchInductor가 사용되지만, 다른 옵션들도 존재한다고 한다(제대로 알아보진 않았다).
❗️ 유의점:
compile된 모델 저장 시, state_dict만 저장가능하다.
∙ 아래는 가능!
torch.save(opt_model.state_dict(), "best.pt")
torch.save(model.state_dict(), "best.pt")
torch.save(model, "best.pt")
∙ 아래는 불가능!
torch.save(opt_model, "best.pt")
Compile 이후 사용가능한 기능:
∙ TorchDynamo
Eager Mode의 가장 용이한 점: 학습도중 model_weight에 접근하거나 값을 그대로 읽어올 수 있다.
model.conv1.weight
TorchDynamo는 이를 인지하고 만약 attribute가 변한것을 감지하면 자동으로 해당부분에 대한 변화를 다시 컴파일해준다.
∙ Inference
compile함수로 compiled_model을 생성한 후 warm-up step은 초반 latency를 줄여준다. 다만, 이부분도 차차 개선시킨다 하였다.
Sweep은 필히 2개의 단계(Initialize the Sweep, Run the Sweep Agent)가 필요하다.
1. Initialize the Sweep
∙ Sweep Configuration를 정의
Sweep Initialize를 위해 먼저 구성요소(configuration)를 정의해야한다.
이를 위해 required와 option으로 나뉜다.
program(어디에서) method(무엇을) parameters(어떻게) 최적화를 할 것인지 정의해야한다.
이때, 최적화 방법으로 3가지가 존재한다.
Grid 방식 : 가능한 모든 조합 탐색 (= Cost↑)
Random 방식 : random하게 선택 (= Cost↓, opt찾을확률↓)
Bayes 방식 : 이전에 시도한 hyper-parameter조합의 결과를 사용, 다음시도조합 추론시 사용 → 모델성능을 최대로 향상시킬 수 있는 hyper-parameter조합을 찾는다. (= 초기탐색이 느림)
이때, 특히나 parameters 파트가 중요하기에 좀 더 살펴보자.
valuesvalue Hyper-parameter에 대해 특정 값을 설정해서 우리가 원하는 값만 선택하게 해줌. (value는 1가지 값을 설정해줄 때 사용) distribution values와 대조되는 방식. 특정 값을 설정하는 대신 원하는 분포 안에서 값을 선택. Sweep에서는 uniform, normal, q_log_uniform과 같이 다양한 분포를 제공. 또한 선택된 분포를 min, max와 mu, sigma, q를 통해 자유롭게 변형가능.
min, max 분포의 최소∙최대값을 설정. musigma 평균과 표준편차를 나타내는 값, 정규분포(normal)의 모양을 결정. q Quantization의 약자로 distribution에서 나온 값 X를 양자화. ex) q를 2로 설정한다면 X는 2의 배수로 바뀜. (ex. 식 round(X / q) *q를 적용하면, -2.96은 -2로 13.27은 14로 8.43은 8로 바뀜.)=
∙ project에 사용하기위해 Sweep API로 초기화
Sweep의 config가 제대로 정의가 됐다면 이제 프로젝트에 적용을 해줘야한다.
sweep 초기화 코드:
sweep_id = wandb.sweep(config.sweep_config)
위에서 정의된 config 변수를 입력으로 받고 sweep id를 출력해준다.
이 id는 다음 step에서 sweep을 실행시킬 때 고유한 identifier로 사용된다.