New version of pytorch
Published on January 05, 2023 by JunYoung
AI Deep learning Module Pytorch
16 min READ
이번 게시글에서 소개할 내용은 pytorch 1.13 버전 이후 새롭게 출시된 pytorch 2.0을 사용해보는 것이다.
본인은 파이토치, 텐서플로우 모두 써본 사람이고, 굳이 비교하자면 파이토치가 텐서플로우보다 사용하기 편하다고 생각된다. 일단 사용자 친화적(클래스 정의부터 시작해서, 학습 함수를 구현하는 과정 등등이 굉장히 깔끔하다)이기도 하고, 가장 큰 이유는 그냥 내가 많이 써와서인듯하다. 암튼 이렇게 어떤 모듈이 더 낫다는 것은 절대적으로 주관적인 생각이고 파이토치, 텐서플로우와 더불어 다양한 딥러닝 모듈이 존재한다.
파이토치를 개발한 meta(구 facebook)은 tensorflow와는 다르게 굉장히 꾸준한 문법 및 사용법을 유지해오고 있다. 뒤쪽에서 motivation에서 언급하겠지만 파이토치 개발자들도 사용 편의성에 집중한 플랫폼 구축을 가장 최우선의 목표로 삼았다고 한 것을 볼 수 있다. 아무튼 그러한 이전 단계에서 넘어서서, 이제는 next-generation pytorch를 준비하고자 큰 결심을 해서 한 것이 바로 pytorch 1.x에서 pytorch 2.0으로 이름을 달리 지은 것.
파이토치가 가지는 큰 장점 중 하나는 python API와의 효용성이 크다는 것인데, Pytorch 2.0은 여기에 이어 더 빠른 속도로 학습이 가능하다고 한다. 아마도 Pytorch 2.x 버전에서는 뒤이어 설명할 model.compile이 가장 주요한 변화가 아닐까 싶다.
아마 tensorflow를 사용해본 사람들은 해당 문구에 익숙할 것이다. Tensorflow에서는 model.compile이 의미하는 바가 메소드의 인자로 최적화(fitting)할 optimizer, loss 그리고 metric을 설정해주는 것이다. 여기서 loss는 직접적으로 최적화에 사용되는 objective function으로 작용하고, metric은 classification과 같은 task에서의 accuracy로 생각하면 된다.
Pytorch 2.x에서의 model.compile은 조금은 다르게 기존 학습 코드 자체는 그대로 유지하도록 하지만, Pytorch가 기존에 작동하던 C++ 베이스에서 넘어와서 python 상에서 구동하도록 한다는 것이다. 그리고 pytorch 2.0에서 좋다고 말할 수 있는 부분은 완전히 추가적인 기능(optional)이므로, model.compile과 같은 기능을 사용하지 않는다고 해서 pytorch 2.x 버전을 사용하지 못하는 것은 아니다. 즉, 이전 버전의 pytorch 모든 코드가 pytorch 2.x에서도 그대로 호환이 가능하다는 것이다. torch.compile 코드에 적용되는 기반 기술들은 다음과 같다.
대충 번역해봤는데, 아무튼 아직 데모 단계에 가까워서 조금 고오급 최신 GPU에 대해서는 성능이 확실하게 보장되는 것 같은데, 옛날 GPU 모델에 대해서는 아직까지는 유의미한 발전은 아닌 것 같다. 그래도 아마 조금은 빨라질 것 같은 느낌적인 느낌은 있다. 일단 되게 좋은 점은 기존 파이토치 문법을 그대로 사용해도 된다는 것이다. 사실 번역된 내용만 보면 컨셉을 파악하기가 힘들어서(영어로 봐도…ㅋㅋㅋ), 이 부분은 뒤에서 다시 디테일하게 짚고 넘어갈 것이다.
말만 빠르다 빠르다 하면 검증이 안되니까 실제로 파이토치에서 실험을 해봤다. 내가 한 것은 아니고 파이토치 개발진 분들이 했다고 한다. 앞서 말했던 것처럼 단순히 model을 컴파일만 해주면 되는 부분이라, 163개의 open-source model를 활용(구체적으로 분류해보면 46개의 HuggingFace Transformer 모델들, 61개의 TIMM model들 그리고 56개의 TorchBench 모델들)하였고, 이는 Image classification부터 시작해서 NLP task나 강화학습과 같이 광범위한 딥러닝 네트워크를 커버한다. 즉, 어떠한 task에 대해서도 속도가 좋아진다는 걸 보여주고 싶었다고 한다.
결론부터 보자면 위와 같다. open-source model의 형태나 코드를 전혀 바꾸지 않고 단순히 torch.compile을 통해 wrapping해주고, 속도 향상과 validate accuracy를 측정하였다. 물론 속도 향상은 data-type에 따라 상이하기 때문에 float32와 Automatic Mixed Precision(AMP)등에 대한 속도 향상을 모두 측정했고, 계산한 두 속도 향상을 $0.75 \times AMP + 0.25 \times float32$로 계산했다고 한다. 여기서 AMP에 weight을 좀 더 넣어준 이유는 실질적으로 더 많이 보여서(활용되어서) 그렇다고 한다.
163개의 open-source model에 대해 torch.compile은 93%, 모델 학습은 43% 빠른 속도로 동작하였다고 한다. 참고로 이 결과 기준은 서버용 GPU인 A100에서 측정된 결과고, 로컬 컴퓨터나 desktop에서 사용하는 GPU인 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.
실제로 위와 같이 써놓으심.. 잔뜩 기대하고 pytorch 2.0 설치했는데 내 컴퓨터에서는 안된다니 너무 슬펐다. 나중에라도 쓸 수 있겠지 기대하면서 존버해야지.
앞서 본 내용은 pytorch 2.0 출시 초창기에 완벽히 셋팅되지 않은 상태에서 사용했기 때문에 이런 저런 문제가 있었던 것 같다. 간만에 파이토치 홈페이지를 들어가보니 내용이 많이 업데이트된 것으로 보이고 보다 안정적인 pytorch 2.0을 사용해볼 수 있지 않을까 하는 기대감에 추가로 글을 작성해본다.
파이토치가 2.x 버전으로 넘어가면서 개발자 분들은 compile 이라는 메소드에 집중한다고 한다. 앞서 언급했던 여러 기반 기술들을 통해 model을 최적화하고, 학습 속도를 빠르게 할 수 있다는 면을 보다 강화하고 이를 점차 scalable하게 키울 생각인 것 같다. 따라서 기존의 파이토치와는 조금 다르게 Pytorch 2.0 시리즈에서는 지속적으로 연구하면서 그와 동시에 사용자들과의 interaction을 통해 인사이트를 얻고, 더 좋은 서비스를 이후 버전에서 보여주겠다는 것이 목표인 듯하다.
실제로 파이토치 사용자들 중 유명한 사람들의 몇몇 인사이트에 대해 인용한 부분은 다음과 같다.
“With just one line of code to add, PyTorch 2.0 gives a speedup between 1.5x and 2.x in training Transformers models. This is the most exciting thing since mixed precision training was introduced!”
간단하게 말하자면 Transformer에 model.compile을 사용하니, 학습 속도가 약 1.5배에서 2.0배 증가했다는 말이다.
“It just works out of the box with majority of TIMM models for inference and train workloads with no code changes”
코드 변화 없이(?) 바로 TIMM에 내장된 대부분의 모델에서 inference나 train 속도가 빨라졌다는 말인 것 같다. 정말 그런가…?
“PyTorch 2.0 embodies the future of deep learning frameworks. The possibility to capture a PyTorch program with effectively no user intervention and get massive on-device speedups and program manipulation out of the box unlocks a whole new dimension for AI developers.”
유저가 굳이 프로그램을 뜯어볼 필요 없이 pytorch의 프레임웤을 효율적으로 사용할 수 있게 되었다는 건, pytorch가 어떤 임베디드 플랫폼 상에서도 효율적으로 작동할 수 있을 것이라는 말인 듯하다.
파이토치 개발자들이 개발에 있어서 가장 중시했던 포인트는 flexibility와 hackability를 보장하는 것이었다. 즉 사용하는 사람들에게 제공하는 시스템은 최대의 자유도를 가지고 있는 것이 바람직하다고 생각한 것 같다. 그리고 여기에 추가적으로 플랫폼 최적화를 통한 performance 향상이 이어진다고 볼 수 있다.
Pytorch가 런칭한 것은 2017년이고, 런칭 이후 GPU와 같은 병렬 연산기/하드웨어 가속기는 당시 스펙보다 메모리 접근 속도나 연산 속도 측면에서 많은 향상이 이루어졌다. 그러면서 자연스레 파이토치는 eager execution의 성능을 높이기 위해 pytorch의 코드 대부분을 C++로 옮기게 되었으며(파이토치 코드의 대부분 source는 C++이 기반에 있는 것을 알 수 있다), 이러한 방식은 사용자들에게 있어 코드 기여도(hackability)를 낮추는 진입 장벽이 되어버렸다. 여기서 eager execution이란, 그래프 생성 없이 연산을 즉시 실행하는 명령형 프로그래밍 환경을 의미한다. 아마도 텐서플로우를 써본 사람이라면 어떤 개념인지 알 것 같은데, 이러한 이유 때문에 pytorch는 tensorflow와 사실상 딥러닝 개발 환경 구축에 있어 방향성이 다른 연구를 했다고 볼 수 있을 것 같다.
아무튼 이렇게 독자적으로 열심히 eager execution 쪽에서의 성능을 높이는 것에는 한계가 있다고 판단하였다. 그렇기에 pytorch에서도 상당히 이른 시기인 2017년 중순(7월)에 pytorch를 위한 compiler를 만들기로 결정하였다. 컴파일러는 pytorch 속도를 빠르게 바꾸면서, 그와 동시에 pytorch experience를 해치면 안된다(복잡하거나 자유도가 낮으면 안됨)라는 목표를 가진다. 이런 유저의 자유도/사용 편의성과 관련된 flexibility를 유지하는 메인 기준으로는 dynamic shapes, programs를 보조하는 것이었다.
앞서 언급했듯이 파이토치는 2017년을 기준으로 컴파일러 연구를 지속해왔다. 컴파일러의 구성 요소를 총 3가지로 나누어 설명한다. 다음 세 요소 중 가장 어려운 부분이 graph acquisition이라고 한다.
런칭 이후 파이토치에서는 torch.jit.trace, TorchScript, FX tracing 그리고 Lazy Tensors와 같이 모델 최적화를 위한 방법을 개발했지만, 어떤 것도 개발진들이 원하는 '컴파일러의 느낌'을 주지는 못했다. 몇몇은 자유도가 높지만 그에 비해 속도가 떨어진다거나, 속도가 빨라지면 그만큼 자유도가 떨어지는 trade-off가 발생한 것이다. 사실 본인도 다른 플랫폼 상에서 학습된 모델을 사용할 때 효율성 및 적합성을 위해 몇몇 스크립트를 사용해봤지만 불편한 점이 좀 많았던 기억이 있다. 그나마 TorchScript가 괜찮은 성능을 보였지만 작성한 code에 성능이 많이 의존하기도 하고 사용자가 코드를 많이 손봐야한다는 점에서 단점이 발생하였다. 그러다보니 파이토치가 익숙하지 않은 사람들은 많이 사용하지 않는 경향이 생겨버렸다.
그림을 보면 앞서 언급한 compiler의 각 기술적 요소들이 어떠한 역할을 하는지 확인할 수 있다. 먼저 TorchDynamo, AOTAutograd가 eager execution 형태로 작성된 모델 구조를 그래프화하고, gradient 연산이 가능하게끔 도와준다. 여기에 Prim과 같은 연산 단순화 방법을 통해, 파이토치에 있는 수많은 연산들을 간단한 연산들로 바꾸어 그래프를 단순화하는 작업을 도와준다. 이렇게 단순화된 그래프 구조를 TorchInductor가 컴파일하는 것이다.
TorchDynamo는 PEP-0523에서 ‘frame evaluation API’로 소개된 Cpython feature를 사용한다. 컴파일링이 안정적으로 진행되기 위해서는 graph acquire이 필요한데, 이를 도와주는 c언어 기반 프레임 워크라고 볼 수 있을 것 같다. 실제로 해당 방법이 효과적인지 비교해보기 위해 pytorch로 작성된 $7000$개가 넘는 깃허브 소스를 참고하여 코드를 돌려보았고, TorchScript나 기타 방법들은 graph를 획득하는 시간이 $50\%$에 불과했지만(코드를 손봐야하는 overhead도 큼), TorchDynamo는 $99\%$의 거의 대부분의 시간동안 graph를 획득했고, 코드를 손봐야하는 번거로움도 없어졌다. 개발자들이 원했던 flexibility(overhead), speed(graph acquisition) 모두 충족한 것이다.
Pytorch 2.0을 위한 새로운 백엔드 컴파일러를 설계하는데 있어서, Triton language의 사용량 증대에 영감을 받았다고 한다(트리톤을 사용하게 되면 GPU에서 연산되는 작업을 가속화할 수 있다는 장점이 있다). 따라서 개발자들도 Pytorch eager를 그대로 가져가고, pytorch의 여러 특징들을 그대로 유지할 수 있는 컴파일러 백엔드를 고안하였다.
Pytorch는 tensorflow와 다르게 연산 구조를 정의한 후 값을 계산하는 구조가 아닌, 구조와 계산이 동시에 진행되는 방식인 'define by run'을 사용하는데, TorchInductor는 IR 레벨에서의 학습 루프를 사용하여 GPU상의 Triton 코드/CPU 상의 C++/OpenMP 코드로 매핑하게 된다. Core loop가 IR 레벨에서 적은 operator로 구성되고 무엇보다 python으로 구현되어있기 때문에 C++보다 hackable/extensible한 효과를 기대할 수 있다.
Pytorch 2.0에서 목표로 삼았던 것 중 하나는 학습 속도를 빠르게 하는 것이다. 그렇기 때문에 단순히 user-level code만 capture하는 것이 아니라, 학습 속도에 큰 영향을 주는 backpropagation 또한 capture하는 것이 중요하다. Capture라는 말이 자주 쓰이는데, 흔히 스마트폰으로 재미있는 사진이나 기록하고 싶은 내용을 찾게 되면 스크린샷(캡처)를 아는 상황을 생각해보면 된다. 학습 과정에 유저가 정의한 모델 연산의 순서나 loss function의 미분을 통한 backpropagation 모두 학습 루프 상에서 지속적으로 유지된다면 일종의 가이드북 역할을 통해 번거로운 연산의 중복을 피할 수 있다는 것이다. 그리고 이미 pytorch에 있는 autograd system을 재활용하고 싶다는 것이 메인 아이디어였고, AOTAutograd는 pytorch의 torch_dispatch를 사용하여 backward pass를 Autograd 엔진을 통해 미리 capture 해둠으로써(Ahead-Of-Time) 속도 향상을 이루어냈다.
Pytorch backend를 작성하는건 생각보다 굉장히 어려웠는데, 이는 Pytorch가 operators를 굉장히 많이 지원하기 때문이다(세세하게 분류하자면 거의 $2000$개 이상).
따라서 각 operator마다의 특징을 잡아서 백엔드를 구성하는 것은 여간 어려운 일이 아닌 것이다. 그렇기에 PrimTorch라는 프로젝트에서 진행한 내용은 수많은 operator를 통합할 수 있는 보다 작고 안정적인 operator set를 구성하는 것이었다. Pytorch에 존재하는 여러 operator는 적은 규모의 operator set를 조합함으로써 모두 구성될 수 있기 때문이다. Prim 프로젝트에서 정의한 두 operator set는 다음과 같다.
우선 자신의 GPU 사양을 확인해준다.
nvidia-smi
나는 CUDA 버전이 11.7이라 해당 스펙에 맞는 latest nightlies를 설치해주었다.
pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117
본인 컴퓨터에 CUDA 버전이 11.6이면 다음 install 코드를 쓰면 되고,
pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu116
그것도 아니고 CPU만 있으면 다음 install 코드를 쓰면 된다.
pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cpu
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 메소드를 보다 최적화시켜 속도를 빠르게하는 것이다. Compiling과 관련된 내용을 조금 더 디테일하게 보면 다음과 같다.
def torch.compile(model: Callable,
*,
mode: Optional[str] = "default",
dynamic: bool = False,
fullgraph:bool = False,
backend: Union[str, Callable] = "inductor",
# advanced backend options go here as kwargs
**kwargs
) -> torch._dynamo.NNOptimizedModule
컴파일 단계에서 model을 제외하고 아무런 옵션을 넣지 않는다면 위에 보이는 것과 같은 default 셋팅이 적용된다. 각각의 arg에 대해 살펴보면 다음과 같다.
reduce-overhead를 통해 메모리는 조금 더 쓰되 overhead를 줄여주는 옵션이 있으며, max-autotune을 통해 최대한 빠르게 학습할 수 있는 코드를 오랜 시간동안 컴파일해서 제공하는 옵션도 있다.# API NOT FINAL
# default: optimizes for large models, low compile-time
# and no extra memory usage
torch.compile(model)
# reduce-overhead: optimizes to reduce the framework overhead
# and uses some extra memory. Helps speed up small models
torch.compile(model, mode="reduce-overhead")
# max-autotune: optimizes to produce the fastest model,
# but takes a very long time to compile
torch.compile(model, mode="max-autotune")
dynamic 다이내믹은 dynamic shape에 대해 code path를 enabling할 지 결정하는 boolean 변수이다. Compiler 최적화 과정이 프로그램을 dynamic shape 프로그램에 적용될 수 없게 만드는 경우가 있는데, 이를 조절함으로써 본인이 원하는 방향대로 컴파일을 할 수 있게 해준다. 이 부분은 아직 완벽히 이해가 되지는 않지만 데이터 shape가 변하는 상황에서 graph를 유동적으로 컴파일할 수 있게끔 하는 것과 관련이 있을 것 같다.
fullgraph 풀그래프는 Numba의 nopython과 유사하다. 전체 프로그램을 하나의 그래프로 컴파일하고, 만약 실패한다면 왜 불가능한지 설명해주는 error 메세지를 띄운다. 굳이 쓰지 않아도 상관없는 옵션.
backend 백엔드는 어떤 compiler backend를 적용할 지 결정하게 된다. 디폴트로 정해진 값은 앞서 설명했던 TorchInductor가 사용되지만, 다른 옵션들도 존재한다고 한다(제대로 알아보진 않았다).
Eager mode에서 가장 용이한 점은 학습 도중에 model의 weight에 접근하거나 값을 그대로 읽어올 수 있다는 점이다(예를 들어, model.conv1.weight). TorchDynamo는 이를 인지하고 만약 attribute가 변한 것을 감지하면 자동으로 해당 부분에 대한 변화를 다시 컴파일해주게 된다.
# optimized_model works similar to model, feel free to access its attributes and modify them
optimized_model.conv1.weight.fill_(0.01)
# this change is reflected in model
optimized model의 state-dict나 model의 state-dict를 serialize할 수 있다. 최적화된 모델과 모델이 서로 같은 parameter를 지칭하기 때문에 아래의 두 코드는 같은 결과를 낸다.
torch.save(optimized_model.state_dict(), "foo.pt")
# both these lines of code do the same thing
torch.save(model.state_dict(), "foo.pt")
하지만 model 자체를 serialize할 때는 조금 달라진다. 이럴 경우 optimized model이 아닌 model을 저장해야 오류가 나지 않는다.
torch.save(optimized_model, "foo.pt") # Error
torch.save(model, "foo.pt") # Works
저장할 때 주의하도록 해야겠다.
Model inference를 하는 경우, compile을 통해 compiled model을 생성한 뒤 warm-up step을 밟는 것이 초반 latency를 줄여준다고 한다. 아직 pytorch 2.x가 개선되어야할 부분처럼 보이는데, export라는 메소드를 통해 latency에 대한 환경적 변수까지 안정화할 수 있는 방법을 제공할 예정이라고 한다. 이 부분도 차차 개선되는 부분이 생기면 이해를 해봐야겠다.
# API Not Final
exported_model = torch._dynamo.export(model, input)
torch.save(exported_model, "foo.pt")
컴파일을 하는 상황에서는 디버깅이 힘들다. 예를 들어 프로그램이 compile mode와 충돌한다던지, eager mode(원래의 pytorch 연산)과 compile된 연산이 정확하게 일치하는지, 혹은 컴파일을 진행했음에도 speedups가 원활하지 않다는 등의 문제가 발생할 수 있다. 만약 컴파일된 코드가 프로그램과 충돌하거나 eager mode로 실행했을 때보다 오차 범위 이상으로 결과가 달라진다면, 코드 자체가 문제라고 보기는 어렵다. 사용자의 디버깅이나 reproduction 등을 지원해주기 위해 minifier를 제공한다. Minifier는 마주한 issue를 간단한 code snippet으로 자동으로 줄이게 되고, 해당 코드는 문제를 재현하고 축소된 코드로 github issue를 제시할 수 있게 된다. 즉 이슈를 제시하는 과정에서 파이토치 팀이 굳이 방대한 코드를 보지 않더라도 원인규명을 빠르게 할 수 있게끔 해준다는 것이다. 또한 만약 속도 향상과 관련된 이슈가 있다면, torch._dynamo.explain 툴을 통해 graph breaks라 불리는 속도 저하 요인을 찾아낼 수 있다.
여러 개의 GPU를 병렬 연산에 활용하는 방법 중 하나가 바로 torch.distributed 방법이다. Pytorch의 대표적인 두 distributed wrapper인 DDP(DistributedDataParallel), FSDP(FullySharedDataParallel) 모두 컴파일 후 제대로 작동하는 것을 알 수 있으며 eager mode(기존 방법)에 비해 성능이나 메모리 효율성이 올라갔다고 한다.
사실 아직 제대로 이해된 부분이 많지는 않지만 요약하자면 compile을 통해 eager mode에서 작동하던 기존 pytorch model을 최적화하였지만, 기존의 편의성(모델 파라미터에 접근이 쉽고, 연산의 병렬 처리 등등 코드 구성하는 과정에서 자유도가 높음)을 그대로 유지했다는 것이다. Pytorch 2.x는 앞으로 compile을 보다 편리하고 확장 가능성 있는 방향으로 발전시키고자 할 것이며, 상대적으로 자유도가 높지 않기도 했고 복잡했던 tensorflow에 비해 보다 대중적인 언어로 자리잡으면서도, 고성능을 요구하는 개발자들의 니즈를 만족할 수 있는 발전 방향이 될 것이라고 생각된다.