Train Language Model on TPU

Posted on Tue, May 11, 2021 NLP

준비물: TPU on GCP (with TFRC)

TPU는 GCP위에서만 사용할 수 있다. (Colab, Kaggle 에서도 사용할 수 있다. 둘다 GCP위에서 동작하기 때문.)

Colab, Kaggle TPU는 무료지만 시간 제약이 커서, 큰 LM을 학습하기에 적합하지 않다.

따라서 TFRC(https://www.tensorflow.org/tfrc?hl=ko) 프로그램을 지원해 TPU v3-8을 사용하는 것이 좋다. (1달간 5대를 사용할 수 있다.)

알아두면 좋은 TPU 이야기

공통: Tokenizer 제작 with Tokenizers🤗

💡

Huggingface Tokenizers 라이브러리로 데이터셋에 기반한 Tokenizer 만들기

BERT WordPiece (for BERT, ELECTRA, RoBERTa)

파일 명 20210309_cleaned.txt 파일에서 Vocab을 생성하려면 가장 간단하게는 아래와 같이 만들어줄 수 있다.

# pip install tokenizers
from tokenizers import BertWordPieceTokenizer

tokenizer = BertWordPieceTokenizer(
    clean_text=True,
    strip_accents=False, # Must be False if cased model
    lowercase=False,
    wordpieces_prefix="##"
)

tokenizer.train(
    files=['./20210309_cleaned.txt'],
    limit_alphabet=4000,
    vocab_size=25000,
)

# 이렇게 하면 현재 위치에 `vocab.txt`가 생긴다.
💡

해당 파일은 앞뒤 공백 제거 / 중복 제거 / 특수문자 제거 등 Text 전처리를 마친 텍스트여야 한다.

💡

생각보다 많은 메모리를 먹는다! 여유롭게 텍스트 용량 x 10배의 Ram을 준비해두자.

위와 같이 만든 Tokenizer는 이후 아래와 같이 불러와 사용할 수 있다.

from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained('.') # vocab.txt 있는 폴더 위치

BPE (for GPT-2)

파일 명 20210309_cleaned.txt 파일에서 Vocab을 생성하려면 가장 간단하게는 아래와 같이 만들어줄 수 있다.

# pip install aitextgen
from aitextgen.tokenizers import train_tokenizer

train_tokenizer(file_name, vocab_size=50000)
# 이렇게 하면 현재 위치에 aitextgen.tokenizer.json 파일이 생긴다.

이렇게 만든 aitextgen.tokenizer.json 파일을 Huggingface Transformers에서 곧바로 사용 가능한 vocab.json , merges.txt 파일로 아래와 같이 쪼개줄 수 있다.

from tokenizers import Tokenizer

tk = Tokenizer.from_file('./aitextgen.tokenizer.json')
tk.save('./tokenizer.json')
tk.model.save('./')
# ['./vocab.json', './merges.txt'] 두 파일이 생긴다. 

BERT/ELECTRA on TPU

💡

Google-Research의 공식 Github code를 통해 Tensorflow로 MLM 학습

Google Colaboratory
Colab에서 TPU로 BERT 처음부터 학습시키기 - Tensorflow/Google ver.

2018년말부터 현재까지 NLP 연구에서 BERT는 여전히 압도적인 위치를 차지하고 있다. 한편, BERT모델을 사용하는 이유 중 가장 큰 것 하나가 바로 한국어로 Pretrained된 모델이 있다는 점이다. Google에서 논문을 처음 공개했을 때 Multilingual pretrained model을 공개해 Fine-tuning만으로도 우리가 필요한 데이터셋에 맞춰 분류기를 만드는 등의 여러 응용이 가능하고, 동시에 높은 성능 을 보여주었기 때문에 BERT 자체를 학습시키는 것에 대해서는 크게 신경쓰지 않은 것이 사실이다.

KLUE-benchmark/KLUE-ELECTRA

multiprocessing을 위해 corpus를 여러개로 분리해줍니다. 스크립트를 실행하여 tfrecord로 변환해줍니다. electra_config 폴더와 hparams 폴더에 학습에 필요한 hparam들이 들어있습니다. (small, base, large) 학습 환경과 사용 목적에 맞춰 tpu_name, tpu_zone, vocab_size 등을 조절해줘야 합니다. script/run.sh의 argument( DATA_DIR, MODEL_NAME, HPARAMS_JSON)를 설정해준 뒤 스크립트를 실행하여 pretraining을 진행합니다. script/convert_tf_ckpt_to_torch.sh의 argument( MODEL_NAME, CONFIG_FILE, VOCAB_FILE, DISCRIMINATOR_OR_GENERATOR)를 설정해준 뒤 스크립트를 실행하여 변환을 진행합니다.

Data Preprocessing - BERT

BERT는 Google-Research 공식 Repo(https://github.com/google-research/bert)의 코드를 사용한다.

git clone -q https://github.com/google-research/bert

BERT는 Static Masking을 진행하기 때문에, 데이터 인코딩 과정에서 MASK를 진행한다. (이와 함께 N번의 mask를 하기 때문에 원본 데이터보다 N배로 데이터셋 크기가 늘어난다.)

적당히 큰 데이터셋 txt 파일인 dataset.txt 가 있다면, 적당한 크기(256k lines)로 잘라 shards 폴더에 넣어준다.

mkdir ./shards
split -a 4 -l 256000 -d dataset.txt ./shards/shard_
ls ./shards/

그 뒤 bert/create_pretraining_data.py 파일을 이용해 인코딩을 진행하면 된다.

Train Guide

Encoded Data를 모두 GCS(GCP 스토리지)에 올리고, 로컬에서 tf estimator로 학습을 진행하면 된다.

Logging

로그는 선택한 Google Storage gs://버킷이름/모델이름 에 저장된다.

tensorboard --logdir gs://버킷이름/bert_model

위 명령어로 확인할 수 있고, 해당 커맨드는 GCP의 Cloud Shell(무료 터미널)에서 곧바로 사용할 수 있다. (편하다)

Convert to Huggingface🤗 Transformers & Upload Model

bash script/convert_tf_ckpt_to_torch.sh
내가 만든 ELECTRA를 Huggingface Transformers로 Porting하기

BERT, ALBERT, ELECTRA 등을 직접 Pretrain하게 되면 모델이 Tensorflow의 ckpt 형태로 저장이 된다. 이번 글에서는 tensorflow ckpt를 transformers의 pytorch ckpt 로 변환하는 법을 알아보겠다🤗 이번 글에서는 ELECTRA-Small을 기준으로 실습을 해본다. (BERT 등도 방법은 크게 다르지 않다) transformers v2.8.0을 기준으로 작성하였다. 이후 버전에서 호환되지 않는 경우가 있을 수 있다.

Huggingface Trainer on TPU

💡

Huggingface의 Trainer 코드로 PyTorch 통해 LM/MLM등을 학습하기

🤗Trainer로 GPT-2 학습하기

💡

아래 과정은 GCP 네델란드(혹은 TPU와 같은 리전) 리전에 e2-standard-16(vCPU 16개, 64GB 메모리)을 사용하고, 부팅 디스크 이미지를 c2-deeplearning-pytorch-1-8-xla-v20210504-debian-10 를 선택해 PyTorch 1.8 + PyTorch-XLA가 설치 된 환경에서 무난하게 진행할 수 있다. (해당 이미지에는 conda 환경으로, pytorch와 xla가 모두 설치되어있다.)

💡

글 쓰는 일자(20210511)기준, transformers 라이브러리의 버전이 4.5.x이기 때문에 아래 코드들에서는 4.6.0 이상이 필요해, Github source에서 곧바로 받아 설치해준다.

pip install git+https://github.com/huggingface/transformers
# 만약 pypi에 올라온 transformers 버전이 4.6.0 이상이면 pip install transformers로 하면 된다.
pip install datasets

우선 Transformers Repo를 받아오자.

huggingface/transformers

Fine-tuning (or training from scratch) the library models for language modeling on a text dataset for GPT, GPT-2, ALBERT, BERT, DistilBERT, RoBERTa, XLNet... GPT and GPT-2 are trained or fine-tuned using a causal language modeling (CLM) loss while ALBERT, BERT, DistilBERT and RoBERTa are trained or fine-tuned using a masked language modeling (MLM) loss.

git clone https://github.com/huggingface/transformers
mv examples/pytorch/language-modeling/* .
rm -rf transformers

GPT-2는 Casual LM으로, run_clm.py 모델로 학습할 수 있다.

별다른 코드 수정 없이,

위 3가지 파일을 준비하면 된다.

💡

학습 파일은 jsonl 형식 혹은 txt/csv 형식을 지원한다.

아래와 같은 내용의 run_tpu.sh 파일을 만들어준 뒤 실행하면 된다.

export TPU_IP_ADDRESS=10.22.241.98 # TPU IP 적기 
export TPU_NAME=node-1 # TPU Node name

XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470" \
python xla_spawn.py --num_cores 8 \
run_clm.py \
--seed 42 \
--model_type gpt2 \
--tokenizer_name ./ \
--train_file ./20210309_cleaned.txt \
--num_train_epochs 2 \
--per_device_train_batch_size 8 \
--do_train \
--output_dir ./clm-output \
--save_steps 10000 \
--save_total_limit 3 \

위와 같이 TPU config를 적용 후 xla_spawn.py 를 사용하면, run_clm.py 내에 있는 Trainer가 TPU 환경을 지원하기 때문에 잘 동작한다.

save_steps 별로 최신 save_total_limit 개 만큼 로컬에 저장된다.

위 코드 사용시 이슈

현재 위 run_clm.py 는 학습 데이터를 아래와 같은 과정으로 학습 가능한 형태로 인코딩한다.

  1. 학습 데이터를 txt라면 → .cache 폴더 내 jsonl 형식으로 복사
  2. 해당 jsonl dataset을 → Tokenizer로 encode 한다. ({input_ids: [1,1,1,1], token_ids: [1,2,3,4]} 형식으로 저장)
  3. 해당 Encode된 것을 불러와 → 지정한 context length(gpt2=1024)로 합친다. (중간중간 Sep 토큰을 넣어준다.)
  4. 위 3단계로 만든 데이터셋을 Batch로 만들어 학습한다.

이때, 1)로 인해 데이터가 '한번' 복사되고, (x2배)

Tokenizer encode과정에서 2-3배로 증가하고, (x6배)

Encode된 것을 불러와 1024 단위로 chunking하는 과정에서 한번 더 복사되는 (total 24배+) 과정에서 Storage를 무척 많이 사용한다. (따라서 storage를 TB단위로 잡아두는 것이 용이하다.)

Custom code on TPU

💡

xla_spawn.py 를 통해 커스텀 모델 TPU에서 학습하기 (PyTorch)

공식 이미지에 Jupyter Notebook 깔아서 접속하는 Shell Script
TPU_IP_ADDRESS=10.12.12.123

docker run -it -e XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470" \
  -v $(pwd):/root/code \
  -p 8888:8888 \
  --shm-size 16G \
  gcr.io/tpu-pytorch/xla:nightly_3.7 bash -c '
  pip install jupyter &&
  jupyter notebook --ip=0.0.0.0 \
  --allow-root \
  --NotebookApp.token= \
  --notebook-dir=/root/code
  '

http://localhost:8888 로 접속해서 쓸 수 있다.