딥러닝에서 JIT을 사용해야 하는 이유

해당 글의 목적

최근 Pytorch2.0, Jax, TransFormers, Diffusers등 딥러닝 라이브러리에서 JIT을 사용하는 방식으로 업데이트가 이루어지고 있다.
따라서 JIT이 어떤 것이고, 딥러닝에서 JIT을 사용하였을 때 어떤 장점이 있는 지 가볍게 알아보는 것을 글의 목적으로 한다.

배경지식

먼저 JIT을 알기 위해서는 인터프리터와 컴파일러 방식을 알아야 한다.
인터프리터와 컴파일러는 프로그래머가 작성한 코드를 실행하기 위한 두가지 방식이라고 할 수 있다.
아주 간단하게 정리하고 넘어가겠다.

컴파일러

• 소스코드를 컴퓨터가 이해할 수 있는 기계어로 번역하는 프로그램
• 모든 코드를 분석한뒤 실행할 수 있는 목적코드 파일을 생성하는 컴파일 과정이 필요

장점

• 프로그램의 동작 속도가 빠름
• 컴파일 과정에서 오류를 걸러내기 때문에, 오류 어느정도 방지 가능
• 동일한 프로그램을 실행하기 위해 여러번 번역과정을 거칠 필요없이 한번만 컴파일

단점

• 소스코드를 바로 실행할 수 없고 컴파일과정 필요
• 소스 코드의 수정이 필요할 때마다 재컴파일 필요
• 다른 플랫폼에서 실행하기 위해서 재컴파일 필요

컴파일러 언어

C, C++등이 있다.

---

인터프리터

• 소스코드를 한줄씩 읽으며 즉시 실행하는 프로그램
• 실행 시 소스코드를 직접 해석하여 실행

장점

• 간단한 문법
• 소스코드의 수정과 실행이 간편하다.
• 컴파일과 같은 목적코드 생성 과정이 없어 메모리 효율적이다.

단점

• 컴파일러 방식에 비해 실행속도가 느리다.

인터프리터 언어

파이썬, 루비등이 있다.

---

JIT

이제 인터프리터와 컴파일러에 대해 알게 되었으니 JIT 컴파일을 소개하고자 한다.
JIT 컴파일은 인터프리터와 컴파일러의 장점을 합친 방식이라고 할 수 있다.

JIT 컴파일러는 프로그램을 실행할 때 인터프리터를 사용하여 일부 코드를 실행하면서, 반복되는 코드를 컴파일하여 기계어로 변환하는 방식이다.
따라서, 프로그램 상에서 자주 반복되는 부분은 컴파일을 해서 캐시에 올려 재사용하여 반복되는 부분에서의 실행 속도를 개선시킨다.

---

딥러닝에서의 JIT

딥러닝에서의 속도의 중요성

간단히 얘기하면 다음과 같은 이유들 때문에 학습속도 및 추론속도가 중요하다.

효율성

• 딥러닝 모델은 일반적으로 훈련에 많은 양의 연산이 필요하다.
• 학습 속도가 빠르면 시간과 컴퓨팅 리소스를 절약할 수 있다.

모델 튜닝

• 딥러닝에서 연구자들은 모델 성능을 개선하기 위해 다양한 모델 아키텍처와 하이퍼파라미터를 시도하는 경우가 많다.
• 학습 속도가 빨라지면 연구자가 더 빠르게 반복하고 더 많은 구성을 테스트하여 더 짧은 시간에 더 나은 결과를 얻을 수 있다.

실시간 애플리케이션

• 자율 주행 차량 및 의료 이미지 분석과 같이 실시간 처리가 필요한 분야에서 추론 속도는 아주 중요하다.

JIT을 통한 속도개선

최근 파이토치2.0에서 Torch.Compile()함수를 통해 쉽게 JIT을 적용할 수 있게 되었다.(기존에도 TorchScript등을 통해 JIT컴파일러를 사용하는 것이 가능하긴 했지만, 복잡했다.)

위의 사진에서와 같이 Torch.compile를 사용하였을때, 속도가 상승한 것을 볼 수 있다.

원리

파이토치에서 딥러닝 모델을 학습시키거나 추론을 할때, Forward 함수를 아주 반복적으로 실행하게 된다.
이런 Forward함수를 컴파일을 통해 Forward실행을 빠르게 할 수 있기 때문에 속도가 빨라질 수 있는 것이다.

이러한 JIT방식은 모델이나 학습방식과 관련이 없기 때문에 범용적으로 적용 가능하다는 것이 매우 큰 장점이다.

정리

• 딥러닝에서 속도는 중요하다.
• JIT을 사용하면 범용적으로 모델의 학습, 추론속도를 올릴 수 있다.
• 최근 간단하게 JIT을 적용할 수 있는 방식을 딥러닝 라이브러리에서 개선중이다.