Train Language Model on TPU

Train Language Model on TPU

Tags
NLP
Published
Published May 11, 2021

준비물: 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대를 사용할 수 있다.)
  • 최근에는 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🤗

💡
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 학습
  • 가장 간편한 학습 방법, 단 Tensorflow v1.x을 사용
  • *.tfrecords 로 encode해준 뒤, GCS를 통해 업로드 하고
  • TPU vm에서 gcs 통신으로 학습하기 때문에 GCP vm은 엄청 작은 것(2core 4g ram)을 써도 됨
  • GCS에 Tensorboard log가 남는다.
  • 로컬 용량 걱정 등을 안해도 되는 것이 최고!
    • TF Records로 encode 하는 것은 빵빵한 로컬 CPU 자원으로 해결이 가능하다!
  • 아래 KcBERT 데이터 샘플로 Colab 상에서 TPU로 학습할 수 있다. (아래 블로그와 함께 보면 보다 쉽다.)
  • KLUE Benchmark팀 Monologg님이 만들어준 KLUE-ELECTRA를 사용하면 보다 쉽게 ELECTRA모델을 학습할 수 있다.

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

  • KLUE-ELECTRA를 사용해 학습시 아주 간단하게 TF ckpt를 PyTorch로 바꿔줄 수 있다.
  • 해당 파일의 변수들만 잘 바꿔주자.
bash script/convert_tf_ckpt_to_torch.sh
  • 아래 Monologg님의 블로그를 참고해보자.

Huggingface Trainer on TPU

💡
Huggingface의 Trainer 코드로 PyTorch 통해 LM/MLM등을 학습하기
  • 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 학습하기

💡
아래 과정은 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를 받아오자.
git clone https://github.com/huggingface/transformers mv examples/pytorch/language-modeling/* . rm -rf transformers
GPT-2는 Casual LM으로, run_clm.py 모델로 학습할 수 있다.
별다른 코드 수정 없이,
  • 20210309_cleaned.txt (학습 corpus)
  • vocab.json , merges.txt (Tokenizer)
위 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 형식으로 복사
  1. 해당 jsonl dataset을 → Tokenizer로 encode 한다. ({input_ids: [1,1,1,1], token_ids: [1,2,3,4]} 형식으로 저장)
  1. 해당 Encode된 것을 불러와 → 지정한 context length(gpt2=1024)로 합친다. (중간중간 Sep 토큰을 넣어준다.)
  1. 위 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)
  • 로컬 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를 할 수 있다는 의미
공식 이미지에 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 로 접속해서 쓸 수 있다.