Train Language Model on TPU

준비물: TPU on GCP (with TFRC)

• 최근에는 TFRC 승인이 굉장히 빨라져 하루 내 신청 및 사용을 시작할 수 있다.

• 기간이 끝나면 재신청하면 된다.

• 좀더 긴 기간이 필요하면 사연을 구구절절히 담은 메일을 보내면 긴 기간(3달, 혹은 반년)을 주기도 한다.

알아두면 좋은 TPU 이야기

• TPUv3-8은 코어가 8개 있다.

• 코어별로 한번에 행렬 계산 (MXU) 8개가 기본단위다. 따라서 16개를 처리하는 속도 = 10개를 처리하는 속도
• 따라서 8의 배수로 맞춰서 행렬 계산이 가능하도록 맞춰주는 것이 의미가 있다.

• 코어 하나에 16GB HBM 메모리가 탑재되어있다. 따라서 토탈 메모리 = 128G
• 하지만 코어별 메모리이기 때문에 → 16GB 메모리 기기 8대를 Distributed training하는 것으로 이해하는 것이 더 맞다.
• 따라서 모델 + Optimizer + batch < 16GB가 되어야 하기 때문에, core별 Batch size를 무작정 키울 수 없다.
• PyTorch의 DDP 느낌으로 이해하는 것이 가장 맞을 듯.

• MXU는 굉장히 빠르기 때문에 데이터 io를 잘 맞춰주는 것이 전체 학습 속도에 훨씬 큰 영향을 준다.

• TPU는 fp16대신 bf16을 지원한다.

• TFRC 프로그램을 통해 지원받는 TPU v3-8은 네델란드 리전이라 ping이 170이 넘는다.
• 직접 해당 리전 VM에 연결하는 것보다, 오히려 서울리전VM에 VPN을 만들어 연결 후 작업하는 것이 더 쾌적하다. (인터넷 vs Google Fiber)

공통: Tokenizer 제작 with Tokenizers🤗

BERT WordPiece (for BERT, ELECTRA, RoBERTa)

# 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`가 생긴다.

from transformers import BertTokenizer

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

BPE (for GPT-2)

# pip install aitextgen
from aitextgen.tokenizers import train_tokenizer

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

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

• 가장 간편한 학습 방법, 단 Tensorflow v1.x을 사용

• *.tfrecords 로 encode해준 뒤, GCS를 통해 업로드 하고

• TPU vm에서 gcs 통신으로 학습하기 때문에 GCP vm은 엄청 작은 것(2core 4g ram)을 써도 됨

• GCS에 Tensorboard log가 남는다.

• 로컬 용량 걱정 등을 안해도 되는 것이 최고!
• TF Records로 encode 하는 것은 빵빵한 로컬 CPU 자원으로 해결이 가능하다!

• 아래 KcBERT 데이터 샘플로 Colab 상에서 TPU로 학습할 수 있다. (아래 블로그와 함께 보면 보다 쉽다.)

Google Colaboratory

https://colab.research.google.com/drive/1lYBYtaXqt9S733OXdXvrvC09ysKFN30W

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

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

https://beomi.github.io/2020/02/26/Train-BERT-from-scratch-on-colab-TPU-Tensorflow-ver/

• KLUE Benchmark팀 Monologg님이 만들어준 KLUE-ELECTRA를 사용하면 보다 쉽게 ELECTRA모델을 학습할 수 있다.

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)를 설정해준 뒤 스크립트를 실행하여 변환을 진행합니다.

https://github.com/KLUE-benchmark/KLUE-ELECTRA

Data Preprocessing - BERT

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

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

Train Guide

Logging

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

Convert to Huggingface🤗 Transformers & Upload Model

• KLUE-ELECTRA를 사용해 학습시 아주 간단하게 TF ckpt를 PyTorch로 바꿔줄 수 있다.

• https://github.com/KLUE-benchmark/KLUE-ELECTRA/blob/master/script/convert_tf_ckpt_to_torch.sh

• 해당 파일의 변수들만 잘 바꿔주자.

bash script/convert_tf_ckpt_to_torch.sh

• 아래 Monologg님의 블로그를 참고해보자.

내가 만든 ELECTRA를 Huggingface Transformers로 Porting하기

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

https://monologg.kr/2020/05/01/transformers-porting/

Huggingface Trainer on TPU

• GCP에서 꽤 큰 VM을 사용할 수 있다면 이게 더 편할 수 있다.

• GCP Compute VM - TPU가 모두 같은 region (TFRC를 사용한다면 europe-west-4 (네델란드)로 통일해야 한다.)

• GCP Compute VM이 최소 16core 이상인 vm을 사용해야 Network i/o로 인한 속도저하가 없음

• 현재 Huggingface dataset 라이브러리에서 cache 관련 코드가 너무 큰/많은 리소스를 필요로 함
• CPU는 8core로 괜찮지만, RAM은 32G 이상, Storage는 데이터셋 크기 * 20배 이상이 필요하다.
• 20GB datset → 약 400GB 사용 (캐시)
• 의외로 메모리 자체는 크게 사용하지 않는 편!

🤗Trainer로 GPT-2 학습하기

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

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.

https://github.com/huggingface/transformers/tree/master/examples/pytorch/language-modeling

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

• 20210309_cleaned.txt (학습 corpus)

• vocab.json , merges.txt (Tokenizer)

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 \

위 코드 사용시 이슈

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로 만들어 학습한다.

Custom code on TPU

• 로컬 machine에서도 TPU를 연결할 수 있다.

• GCP에 VM 하나를 띄우고 VPN으로 연결하면 가능
• Seoul Region에 VM을 띄우고 → 네델란드 리전 TPU에 연결하면 로컬기기 → 서울: 인터넷 // 서울 → 네델란드: Google Fiber 통해 전송되어서 latency 최소화
• 인터넷으로 로컬 → 네델란드 가는 것 보다 훨씬 더 빠르고 쾌적함

• 단, 실제 사용용도로는 로컬 머신을 사용하는 것은 비추천
• 이유: TPU 연산 속도에 대비 batch별 data io가 너무 느림
• data io가 $0.1/GB인데, 큰 모델의 경우에는 이것이 오히려 비용의 핵심 원인이 될 가능성이 높음

• 따라서 동일 Region에 VM을 띄우는 것이 최선

• 이런 테스트를 하는 것은 → 큰 VM에서 돌리기 전에 로컬에서 손쉽게 Train code PoC를 할 수 있다는 의미

• 간단한 VPN setup → https://github.com/hwdsl2/setup-ipsec-vpn

• 로컬 기기는 Docker로 → gcr.io/tpu-pytorch/xla:nightly_3.7
• 위 버전은 PyTorch 1.9.0+XLA
• 현재 last release인 gcr.io/tpu-pytorch/xla:r1.8.1 버전은 GCR에서 다운로드가 안됨.

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
  '