TPU를 이용하여 Electra Pretraining하기

TPU를 이용하여 Electra Pretraining하기

최근 ELECTRA의 공식 코드가 공개되면서 한국어 Corpus에 직접 Electra를 만들게 되었다.

이번 글에서는 GCP에서 TPU를 어떻게 사용했는지 그 과정을 공유해보려 한다.

Tensorflow Research Cloud 신청

Tensorflow Research Cloud (TFRC)는 1달 동안 TPU를 무료로 사용할 수 있게 해주는 프로그램이다.

해당 링크로 가서 신청을 하게 되면 메일이 하나 오게 된다.

해당 메일에서 요구하는 대로 신청서를 추가로 작성한 후 제출하면 얼마 후 아래와 같이 답장이 오게 되고, 그 때부터 GCP에서 TPU를 무료로 사용할 수 있게 된다:)

Bucket에 Data 업로드

TPU를 쓰는 경우 모든 input file을 Cloud storage bucket을 통해야만 한다. (관련 FAQ)

Bucket 생성

  • 예제상 Bucket의 이름을 test-for-electra로 만들어 보겠다.

  • GCP 메인 페이지 좌측의 [Storage] - [브라우저] 로 이동

  • 버킷 만들기 클릭

  • 사용할 TPU와 동일한 Region에 Bucket 만드는 것을 권장

File Upload

  • 준비한 pretrain_tfrecordsvocab.txt를 Bucket에 업로드

GCP VM & TPU 생성

  • VM과 TPU를 각각 따로 만드는 것보다, 우측 상단의 cloud shell을 열어 아래의 명령어를 입력하는 것을 추천한다.
  • 저장소는 Bucket이, 연산은 TPU에서 처리하기 때문에 VM Instance는 가벼운 것을 써도 상관이 없다.
$ ctpu up --zone=europe-west4-a --tf-version=1.15 \
--tpu-size=v3-8 --machine-type=n1-standard-1 \
--disk-size-gb=20 --name={$VM_NAME}

Electra 학습 진행

$ git clone https://github.com/google-research/electra
$ cd electra
$ python3 run_pretraining.py --data-dir gs://{$BUCKET_NAME} \
--model-name {$MODEL_NAME} \
--hparams {$CONFIG_PATH}

학습 완료 후 Instance, Bucket 삭제

$ ctpu delete --zone=europe-west4-a --name={$VM_NAME}
$ gsutil rm -r gs://test-for-electra

Reference