AI, ML/기타 정리
[파이토치 2.0] 어떤 점이 달라졌을까?
- -
22년 12월, 파이토치의 대격변이 일어났다.
https://pytorch.org/blog/pytorch-2.0-release/
PyTorch
An open source machine learning framework that accelerates the path from research prototyping to production deployment.
pytorch.org
파이토치 2.0이 발표된지 어연 4개월째가 되었다. 대학원 생활부터 지금까지 머신러닝 프로젝트를 파이토치 프레임워크로 접한 만큼, 이 소식은 꽤나 크게 다가왔다. 사실 파이토치의 개발팀인 meta AI(구 Facebook AI Research, FAIR)는 파이토치 업데이트에 일관적인 기주을 제시해왔기에, 버전 별 호환성이 뛰어난 편이다. 한편 텐서플로우 2.0처럼 버전 업을 통해 전체적인 대격변이 일어난 경우도 있었으니, 처음 소식을 들었을 때는 쉽게 달갑게 받아들이긴 힘들다.
파이토치 2.0의 핵심적인 변화는 'torch.complie()'이라고 한다. 이는 기존의 c++ 환경에서 실행되던 파이토치를 파이썬 환경에 맞게 옮겨 '파이써닉'하게 만들어줌과 동시에, 심지어 더 높은 연산 성능을 보여준다고 한다. 도대체 저 코드 하나에 무슨 뜻이 있는 건지, 처음에는 이해가 잘 되지 않았다. 이번 글을 통해 파이토치가 어떤 특징이 있었고, 2.0에서 그 특징들을 유지하면서도 어떻게 변화했는지 살펴보고자 한다.
Define by run vs. Define and run
파이토치를 설명하는 데는 보통 텐서플로우를 빼놓기 힘들다. 둘은 태생부터 대비되는 존재다. 파이토치는 앞서 밝혔듯 Meta AI에서 사용하던 딥러닝 프레임워크 torch 파이썬으로 옮겨오는 과정에서 개발되었다. 파이토치의 출시일인 16년 9월, 이미 딥러닝 프레임워크의 왕좌 자리에는 구글이 개발한 오픈소스 라이브러리 텐서플로우가 있었다.
하지만 텐서플로우의 버전 관리와 라이브러리(특히 케라스)의 삽질과 더불어, 파이토치의 NLP 및 이미지 분야의 약진은 최소한 연구 분야에서는 파이토치의 손을 들어주게 된다. 그러나 다른 한편, (적어도 내가 알기로) 텐서플로우의 압도적인 모델 배포 가용성에 힘입어 산업 분야에서는 텐서플로우가 우세하다.
이들의 비교되는 행보는 텐서플로우와 파이썬의 근본적인 차이에서 나오게 된다. 바로 그래프를 생성하는 방식이다.
위 예시는 버전이 꽤 오래되긴 했지만, 두 프레임워크의 작동 방식에 대해 잘 보여주는 코드이다.
- 그래프란, 데이터가 모델을 거쳐가는 흐름 구조를 정의한다고 볼 수 있다. (Forward, Backward 포함)
- 텐서플로우는 Define-and-Run 방식을 채택했다. 즉 일단 그래프를 먼저 생성하고, 그 다음에 그래프에 데이터를 넣어 줌으로써 모든 계산 과정을 처리한다.
- 즉 Static Graph 형식을 가진다.
- 텐서플로우 2.0은 이러한 Define and Run 방식뿐만 아니라 eager execution 방식(즉 Define by run)을 도입했다.
- 반면 파이토치는 Define-by-Run 방식을 채택했다. 그래프가 필요할 때마다 생성하여 사용하게 된다.
- 즉 Dynamic Graph 형식을 가진다. 텐서플로우에서 표현하는 eager execution 방식과 상통한다.
그래프를 필요할 때마다 일일히 생성하는 방식은 구조적 유연함을 가져다준다. 특정 코드마다 모델을 바꾸거나, 연산 흐름을 바꾸거나, 당장 생각해보아도 유연성에서 나오는 활용도는 무궁무진하다.
하지만 이러한 방식은 지금까지 파이썬 내부구조상 매우 비효율적이라고 생각할 수 있을지도 모른다. 또한, 최근의 GPU 하드웨어의 기록적인 발전은 파이썬 환경에서 실행되는 코드들의 구조적 제약을 만들었다. 그래서 파이토치는 C++ 기반의 연산을 이용하여 이러한 모순을 해결하고자 했다. 실제로도 꽤 효과가 있었고, 파이토치는 텐서플로우와 비슷하거나 더 나은 연산 속도를 기록하고 있다. 꽤나 성공적인 선택이었던 셈이다.
....그러나, C++ 기반의 연산은 파이썬 이용자들의 접근을 막는 'barrier'가 되어갔다고 파이토치 팀은 표현했다. 이 말인즉슨, C++로 wrapping되었다 보니 파이토치의 기여에 요구하는 스킬 자체가 높아졌고, 이는 파이토치가 슬로건으로 내세운 '구조적으로 유연한 파이써닉'과 점점 멀어지는 결과를 낳았다는 것이다.
즉 파이토치 2.0은 C++에서 다시 파이썬으로 회귀함으로써, 좀 더 '파이써닉'하게 바뀌는 것을 지향하였다.
Torch.compile : 파이토치의 홈커밍
우선 전제조건으로 말하자면, 파이토치 2.0은 파이토치 1.x의 모든 문법을 지원한다. 즉 이전 버전의 코드를 그대로 가져다 써도 그대로 잘 동작한다.
Torch.complie은 선택사항이다. 만약 본인이 C++ 환경에서 그대로 작업하고 싶다면, 그대로 이전의 코드를 써도 된다. 반면 파이썬 환경에서 신문물을 맛보고 싶다면...
compiled_model = torch.compile(model)이 한줄로 모델을 파이썬 환경으로 옮겨 올 수 있다.
언뜻 보면 케라스의 model.compile과 비슷해 보이지만, 이 함수는 파이썬의 동작 백엔드만 바꿔 주니 유의하도록 하자. Torch.complie로 작동하는 모델은 다음과 같은 그래프 생성 및 연산구조를 파이썬 환경에서 실행한다.
파이토치 compiler는 위 그림의 Graph Acquisition, Graph Lowering, Graph Compilation 세 파트를 수행한다고 볼 수 있다. 개발자들은 이전의 compiler들을 위한 프로젝트를 진행했었고, 이들 프로젝트는 유연성(Flexible)과 성능(Fast)간의 trade-off가 항상 존재해왔다. 이러한 고심 끝에 나온 게 이 torch.compile이 되시겠다.
간략하게 보자면, TorchDynamo와 AOTAutograd는 사용자가 작성한 모델을 토대로 그래프 생성을 효율적으로 돕는다. 또한 PrimTorch와 같은 향상되고 통합된 연산을 통해 Graph lowering(그래프를 low-level로 분해하는 것)을 더 효율적이게 만들었다. Torchinductor는 커널 단에서 파이썬 문법을 통하여 파이썬 모델을 C++/OpenMP/Trition code로 효율적으로 변환한다. 그리고 이 모든 과정은 파이썬 내에서 이루어진다. 이제, 각 신기술에 대해 간략하게나마 살펴보도록 하자.
TorchDynamo : Acquiring Graphs reliably and fast
TorchDynamo는 파이썬 PEP-523의 프레임 평가 API를 사용하는 기능에서 영감을 받았다고 한다. (Graph Acquisition이 제일 어려운 부분이었다고 하더니, 나도 잘 이해하기 힘들다). 수많은 파이토치 프로젝트를 평가 기준으로 삼았을 때, 기존의 Torchscript와 같은 compiler들은 실제 시간의 50%까지 떨어지는 낮은 성능에 큰 오버헤드까지 보이는 반면, TorchDynamo는 99%의 시간 내에 큰 오버헤드 없이 안정적으로 그래프를 생성할 수 있었다고 한다.
TorchInductor: fast codegen using a define-by-run IR
OpenAI의 오픈소스 프로젝트인 Trition은 Tiling을 통한 빠른 GPU 연산으로 최근 크게 주목받은 바 있고, 파이토치 개발진도 그러했다. compiler의 백엔드가 파이토치의 difine-by-run을 유지함과 동시에, 파이썬의 다양한 기능들을 유지하길 원했다. Torchinductor는 파이써닉한 루프 레벨의 IR을 도입하여, 파이토치의 모델 코드를 GPU 상에 Trition code로, CPU 상에C++/OpenMP 코드로 매핑한다고 한다.
AOTAutograd : reusing Autograd for ahead-of-time graphs
트레이닝 과정의 시간 소모는 간과하지 못할 사항이기도 하다. 기존의 Autograd는 순전파 과정에서 유저가 작성한 코드의 연산을 추적하고, 이를 토대로 backward()함수가 호출되면 미분값을 역전파한다. 즉 순전파 과정의 유저의 코드를 토대로 역전파를 진행한다.
AOTAutograd는 이에 더불어, 역전파 과정의 연산 역시 추적한다. 즉 Autograd 함수를 추적한다는 의미이다. 여러 번의 연산 과정 동안 계속해서 비슷한 과정을 수행한다면, 그 과정 자체를 미리 기록해 둠으로써 더 좋은 성능을 낼 수 있을 것이다. torch_dispatch 함수를 통해 얻어진 데이터를 통해, 우리는 다음 backward()의 backpropagation 역시 미리(ahead-of-time) 측정할 수 있다고 한다! 이는 TorchInductor와 조합되어 우리 코드의 forward와 backward 연산을 더 가속화시킬 수 있다.
PrimTorch: Stable Primitive operators
기존의 백엔드 개발의 큰 장애물 중 하나는 파이토치의 방대한 연산 함수들이었다. 연산 함수가 많아질수록 백엔드에서 이에 대해 매칭시켜줄 양이 기하급수적으로 많아지기 때문이다. 파이토치의 경우 오버로딩까지 합하여 2000개가 넘는 연산이 정의되었다고 하니, 백엔드에 대한 고충은 더욱 커지기 마련이다.
즉 연산의 교통정리가 어느 정도 필요한 상황이었고, 개발진은 PrimTorch 프로젝트를 통해 이를 해결했다. PrimTorch는 약 250개의 Prim 연산과 750개의 ATen 연산으로 나뉘며, 이들 작은 규모의 연산 셋의 조합을 통해 기존 파이토치의 모든 연산을 구현하였다.
- Prim ops : 250개의 low-level 연산으로 이루어져 있으며, 컴파일러 레벨에 적합하다고 한다.
- ATen ops : 750개의 상대적으로 Prim ops보다 더 high level의 연산으로 구성되어 있다.
이들 신기술을 통해, 파이썬 환경에서 컴파일을 진행함에도 불구하고 torch.complie은 더 나은 성능을 내었다고 한다.
How Fast?
우선, torch.compile은 이러한 최적화된 커널 위에 모델을 컴파일하는 과정이 포함되어 있으므로, 초기 세팅 시간이 쪼금 소요된다.
하지만 연산 중의 최적화 과정 및 eager에 가깝게 동작함을 보장하는 위의 과정등을 통해, 결과적으로는 모델의 학습 속도가 더 빨라진다고 한다.
AMD에서 진행한, HuggingFace 라이브러리의 각 모델을 AMD MI250 GPU 상에서 테스트했을때의 결과이다. 확실히 기존의 Eager mode에 비해 스피드 향상이 꽤 이루어졌음을 볼 수 있다.
디퓨저 모델을 실험한 결과 역시 존재한다. 다른 모든 그래프는 이 링크(https://pytorch.org/blog/accelerated-diffusers-pt-20/)에 같이 나오니, 링크에서 확인해보도록 하자.
torch에 새로 공개된 function인 SDPA(Scaled Dot-Producted Attention)의 speedup을 측정한 것이긴 하지만... xFormers, vanilla 등 다른 attention library를 사용하였을 때보다 전반적으로 SDPA를 적용한 모델이 빠른 speedup을 보여줌과 동시에, torch.complie을 동시에 적용하였을 때 더 뛰어난 성능 향상을 보여주고 있다.
트랜스포머 모델 역시 실험한 결과가 존재하는데(https://pytorch.org/blog/accelerated-pytorch-2/) 이 역시 SDPA의 유용성을 설명해주는 것 같다. 이 포스트에서도 torch.complie을 적용하였을 때를 전제로 실험을 진행하였다.
이후에...
쓰다 보니 SDPA에 대해서도 한 번 자세히 살펴봐야 할 것 같고, torch.complie의 본격적인 사용법도 마무리지어야 할 것 같다. 2편을 쓰는 미래의 나에게 맡기며, 이만 글을 줄인다...
'AI, ML > 기타 정리' 카테고리의 다른 글
| [파이토치 2.0] Scaled Dot Product Attention (0) | 2023.04.18 |
|---|
Contents
소중한 공감 감사합니다