[Mojo] Mojo와 Python 모델 학습 속도 비교하기 (1)

이전 포스팅까지 Mojo의 기본적인 내용들을 살펴보았으니, 이번 포스팅에서는 본격적으로 Mojo를 활용하여 모델 학습을 진행하기 이전에 Mojo 개발 환경을 구축하는 방법과 순수 파이썬으로 실행한 학습 시간과 해당 코드를 Mojo로 그대로 가져와 실행한 학습 시간 차이를 비교해보자.

---

실습 환경

OS: Windows 10

GPU: NVIDIA GeForce RTX 3060 Ti

Mojo 버전: v0.3.0

Python 버전: v3.10.12

 

Mojo 개발 환경 구축

현재 Mojo는 Windows 운영체제를 지원하지 않기 때문에 WSL을 통해 우분투 환경에서 Mojo를 실행해야 한다.

Modular: AI development starts here

 

1. VS Code 설치 후 WSL extension과 Mojo extension 설치하기

 

2. Ubuntu 22.04 for WSL 설치 후 열기

Microsoft Store에서 설치할 수 있다.

 

3. Ubuntu 터미널에서 Modular CLI 설치하기

curl https://get.modular.com | \
  MODULAR_AUTH=mut_1aa2b9bd4b6c4a7f801ae9cfebde8d73 \
  sh -

 

4. Mojo SDK 설치하기

modular install mojo

 

설치 완료 후, 터미널에 나오는 내용처럼 환경 변수 설정하기

본인이 사용중인 쉘 종류를 확인하기 위해서는 echo $SHELL을 입력하면 된다.

 

5. Mojo가 잘 설치되었는지 확인하기

mojo

잘 설치되었다면 아래와 같이 Mojo REPL이 시작된다.

종료하기 위해서는 :quit을 입력하면 된다.

 

6. Ubuntu 환경에서 VS Code 열기

code .

 

7. pytorch 설치

파이토치 한국 사용자 모임 (PyTorch Korea User Group)

Ubuntu 환경이라 OS 종류를 Linux로 선택해야 하나 고민했지만, Windows로 받아보니 정상적으로 실행되었다.

 

설치 후 pytorch가 정상적으로 import되는지 확인

 

 

모델 학습 시간 비교하기

Mojo가 정말 파이썬보다 훨씬 빠른 언어인지 확인해보기 위해 BERT 모델을 fine-tuning하는 코드를 통해 학습 속도에 유의미한 변화가 있는지 살펴보았다. 이번 포스팅에서는 간단하게 순수한 파이썬 환경에서 실행했을 때의 학습 시간과 파이썬 코드를 Mojo로 import하여 실행했을 때의 학습 시간을 비교해보았다.

학습 코드는 이전 프로젝트에서 구현한 코드를 가져와 사용했다.

GitHub - boostcampaitech5/level2_klue-nlp-08: level2_klue-nlp-08 created by GitHub Classroom

 

1. 순수 파이썬 (WSL 환경)

학습 시간 비교를 위해 파이썬 환경에서 코드를 실행해보았다.

학습 결과, 1 epoch 훈련하는데 14분 30초가 소요되었다.

 

2. Mojo에서 파이썬 코드를 불러와 실행

이번에는 Mojo에서 파이썬 코드를 import해서 실행해보았다. main.mojo 파일을 생성하여 train.py 파일의 train_model 메소드를 실행하는 방식으로 진행했다.

 

Mojo에서 local 파이썬 모듈을 import하기 위해서는 경로를 추가해야 한다.

Modular Docs - Mojo🔥 programming manual

# main.mojo
from python import Python

fn main() raises:
    Python.add_to_path(".") # 경로 추가
    let train = Python.import_module("train")
    train.train_model()

이 경우, train.py 파일이 main.mojo 파일과 같은 위치에 존재하기 때문에 현재 디렉토리 경로를 add_to_path() 함수를 통해 추가해 주었다.

raises 키워드의 경우, 예외 처리를 하지 않은 에러가 발생할 수 있는 함수에 붙이는 키워드이다. 일단 간단하게 실행하기 위해 따로 예외 처리는 구현하지 않았다.

 

학습 시간 비교를 위해 진행한 방식이지만, Mojo에서 파이썬 코드를 불러와 실행하는 방법은 정말 비효율적이라고 느껴졌다. 여러 이유가 있었지만 가장 큰 이유는 파이썬 코드에서 발생한 에러의 원인을 파악할 수 없었다.

조만간 해결해줄 것이라고 믿는다..!

예를 들면 내 훈련 코드가 파이썬에서는 정상적으로 실행되는데 Mojo로 import하면 에러가 발생했다. 정말 수많은 삽질을 한 끝에 pytorch lightning 패키지의 seed_everything() 메소드가 에러를 발생시킨다는 것을 알아냈고, 해당 라인을 주석처리를 하자 정상적으로 코드가 실행되었다.

 

학습 코드를 실행한 결과, 14분 34초가 걸렸다. 이로써 파이썬 환경에서 실행하는 것과 비교했을 때 학습 시간에 유의미한 차이가 없는 것을 확인할 수 있었다. 이는 Mojo에서 파이썬 모듈을 import하여 실행할 경우, 파이썬 환경과 동일하게 CPython 인터프리터를 사용하기 때문이라고 이해할 수 있었다.

 

다음 포스팅에서는 본격적으로 모델 학습 코드를 Mojo 스타일로 리팩토링하여 학습 시간을 비교해볼 예정이다.