KcT5 Pretraining on TPU (feat. Flax)

Posted on Tue, Feb 8, 2022 NLP 사이드프로젝트 PLM

Transformer 이후 GPT가 나오고, BERT가 나오고, ELECTRA가 나오고, 뭔가 더 나오고, 더 큰 모델도 나오고, 수많은 모델들이 나오고 있다.

한국어 모델에도 당연히 많은 종류의 모델이 나오고 있고, GPT, BERT, BART, ELECTRA를 비롯해 다양한 모델의 PLM들이 Huggingface Hub에 올라온다.

T5에도 한국어 모델로, KETI가 공개한 KETI-AIR/ke-t5-base 계열의 KeT5 모델들이 있다.

KcT5의 목표는 간단하다! KcBERT, KcELECTRA와 같이 한국어 댓글 데이터셋으로 학습한 T5 모델을 만들어 공개하는 것이다.

데이터셋

데이터셋은 KcBERT-v2022, KcELECTRA-v2022 학습에 사용한 데이터셋과 완전히 동일한 데이터셋을 사용한다.

해당 데이터셋은 2019.01 ~ 2022.01 네이버 댓글 데이터셋으로, 간단한 중복제거 후 약 32GB+ 수준의 텍스트 데이터셋이다.

중복 제거에 대한 노트

이전의 KcBERT, KcELECTRA 학습시에는 Cleaning Function을 제작해 데이터셋에 적용 후 중복을 제거하였으나, 이번에는 단순 (100% 일치) 중복만 AWS EMR을 통해 제거하였다.

Spark DataFrame Repartition을 통해 데이터셋 크기를 조절해, 데이터셋의 순서와 해당 댓글의 순서가 일치하지는 않는다. (댓글 데이터이고, 문장단위 학습이 이뤄지기 때문에 큰 의미가 없다.)

학습환경 설정

TPU VM (feat. SSH)

학습 환경은 TPU VM(Node모드가 아닌, 새로운 모드) 내부에 연결해 진행하였다.

💡

최근 TPU Node가 아닌 TPU VM을 사용할 경우 매우 고성능(96core, 360G ram)의 TPU 메인 VM에 곧바로 SSH 접근해 사용할 수 있다. TRC(구 TFRC) 프로그램을 통해 TPU 지원을 받는 경우, TPU VM방식도 무료로 사용할 수 있기 때문에, 별도의 VM을 띄우지 않고도 모든 작업을 할 수 있다.

TPU VM은 직접 구글 웹 콘솔에서 띄울 수 있다.

다만, 이렇게 할 경우 로컬에 일시적(끄면 사라지는)인 스토리지 100GB만 붙는다.

실제 학습 진행시 huggingface datasets 라이브러리에서 사용하는 cache로 인해 용량이 터진다. 💣💥❗

따라서 별도의 VM Storage를 생성후, TPU 띄울 때 함께 Attach해서 띄워야 한다.

아래는 tpu-vm 이라는 이름을 가진 TPU VM을 띄울 때, tpuvm 이라는 이름을 가진 disk를 함께 붙이는 명령어다.

## GCP TPU with SSD 디스크 

gcloud alpha compute tpus tpu-vm create tpuvm \
--project kcbert-v2022-dd \
--zone=europe-west4-a \
--accelerator-type=v3-8 \
--version=v2-alpha \
--data-disk source=projects/kcbert-v2022-dd/zones/europe-west4-a/disks/tpuvm,mode=read-write

이후, SSH 연결을 아래와 같이 할 수 있다.

## SSH 연결
gcloud alpha compute tpus tpu-vm ssh tpuvm --zone europe-west4-a

SSH 연결후, VM에 붙은 disk를 포맷하고 마운트해준다.

sudo lsblk # <- 이거로 디스크 목록 확인 
sudo mkfs.ext4 -m 0 -E lazy_itable_init=0,lazy_journal_init=0,discard /dev/sdb

여기서는 /home/beomi/ssd 폴더를 만들고 해당 폴더에 디스크를 마운트하였다.

## SSD Mount
sudo mkdir -p /home/beomi/ssd
sudo mount -o discard,defaults /dev/sdb /home/beomi/ssd
sudo chmod a+w /home/beomi/ssd
💡

작업을 진행한 시점의 TPU VM에서는 Root disk가 /dev/sda 이고 추가 디스크가 /dev/sdb 에 마운트되었으나, 실제 작업 시점에 sudo fdisk -l 명령어를 통해 어떤 위치에 있는지 하드웨어 위치를 확인하는 것이 바람직하다.

Flax, Jax

기본적으로 TF와 PyTorch가 설치되어있으나, Jax와 Flax를 TPU 드라이버 버전에 맞게 사용하기 위해서 아래와 같이 가상환경(venv)를 만들어 패키지를 설치해 사용하는 것을 추천한다.

TPU VM을 v2-alpha 이미지로 띄운 뒤 아래 커맨드를 그대로 입력시, Transformers와 jax, flax를 TPU 위에서 사용할 수 있다.

sudo apt update && sudo apt install -y python3-venv
python3 -m venv venv
source venv/bin/activate
# 성공적으로 되면 쉘 앞에 (venv)가 붙는다.

pip install --upgrade pip
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip install flax
pip install git+https://github.com/deepmind/optax.git
pip install datasets transformers

학습

토크나이저 학습

KeT5의 경우 SentencePiece Tokenizer를 사용하였으나, 이번에는 BERT Tokenizer를 사용한다. (이를 통해 KcBERT-v2022, KcELECTRA-v2022, KcT5는 같은 토크나이저를 공유한다.)

BERT WordPiece Tokenizer 를 선택한 이유에 대한 노트

사실 세 모델이 같은 토크나이저를 쓸 이유는 없다. 다만 개인적으로 토크나이징된 결과물이 상대적으로 명확하게 드러나는 점에서 BERT WordPiece Tokenizer를 선호한다.(사람이 눈으로 보고 쉽게 해석 가능)

또한, SentencePiece와 같은 추가 라이브러리 설치 없이 transformers 라이브러리만 설치해도 사용가능하다는 점에서, 좀더 선호하는 측면이 있다.

다만, GPT와 같은 종류에서는 BBPE Tokenizer가 더 높은 성능과 더 낮은 UNK 발생으로 더 좋은 Generation 성능을 보인 실험적 결과가 있어, KcGPT에서는 BBPE를 사용하였다.

토크나이저는 Huggingface의 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=['./20190101_20220123_one_return.txt'],
    limit_alphabet=4000,
    vocab_size=35000,
)

tokenizer.save_model('.')

위를 통해 학습한 토크나이저에 General Vocab(KoBigBird)의 보캡을 추가해주었다.

💡

위와 같이 학습한 BertWordPieceTokenizer는 기본적으로 BOS,EOS 토큰이 없기 때문에 이후 학습시에 추가로 토큰을 넣은뒤 Special Token으로 값을 지정해줘야 한다.

토크나이저에 BOS, EOS등 토큰 추가하기
from transformers import AutoTokenizer

tk = AutoTokenizer.from_pretrained('beomi/KcELECTRA-base')
tk.add_tokens([
    '<s>', '</s>'
])
tk.eos_token = '</s>'
tk.bos_token = '<s>'
tk.save_pretrained('.')

# tk("hello", return_token_type_ids=False)
#{'input_ids': [2, 35813, 19135, 4093, 3], 'attention_mask': [1, 1, 1, 1, 1]}

모델 Config 생성

Huggingface의 모델 형식에 맞게 학습을 하기 위해서 Google T5-base의 Config를 가져와 설정해준다.

from transformers import T5Config
from transformers import AutoTokenizer

tk = AutoTokenizer.from_pretrained('.')

config = T5Config.from_pretrained("google/t5-v1_1-base", vocab_size=len(tk))
config.save_pretrained('.')

위 코드를 실행할 경우 아래와 같이 모델을 위한 모든 설정이 들어간 파일들이 생성된다.

Transformers Model을 위한 파일들

위 파일들을 모두 원하는 모델 파일 이름 폴더로 옮겨준다. (예시: KcT5-base)

mkdir KcT5-base
mv *.json KcT5-base
mv vocab.txt KcT5-base

모델 학습

모델 Pretrain을 위한 샘플로 준비된 run_t5_mlm_flax.py 파일을 받아준다.

wget -nc -q https://raw.githubusercontent.com/huggingface/transformers/762416ffa8b87ceca2d35e56c575ef3c41f449e7/examples/flax/language-modeling/run_t5_mlm_flax.py
파일 버전 고정(762416)

작성시점에는 메인 브랜치가 762416ffa8b87ceca2d35e56c575ef3c41f449e7 여서, 해당 글 시점의 파일을 받도록 파일 URL을 고정하였다.

만약 SentencePiece Tokenizer를 사용한다면 진행하지 않아도 괜찮지만, 만약 위 글와 같이 BertTokenizer를 사용한다면, 아래와 같이 tokenize_function(약 607번째 줄)에 두 옵션을 추가해 attention mask와 token type ids를 반환하지 않도록 수정해주어야 한다.

(기본값은 해당 옵션이 없음)

def tokenize_function(examples):
    return tokenizer(
				examples[text_column_name], 
				return_attention_mask=False, return_token_type_ids=False) # 이 부분을 추가

이후 아래와 같이, 설정 파일이 있는 폴더(KcT5-base)와 학습용 파일(all_cocnat.txt)을 바꾸어 학습을 진행할 수 있다.

# (venv)로 가상환경 활성화된 상태에서 진행하기 

HF_DATASETS_CACHE="./.cache" python3 run_t5_mlm_flax.py \
    --output_dir="./KcT5-base" \
    --model_type="t5" \
    --config_name="./KcT5-base" \
    --train_file='./all_concat.txt' \
    --tokenizer_name="./KcT5-base" \
    --max_seq_length="512" \
    --per_device_train_batch_size="32" \
    --per_device_eval_batch_size="32" \
    --num_train_epochs=2 \
    --adafactor \
    --learning_rate="0.005" \
    --weight_decay="0.001" \
    --warmup_steps="2000" \
    --overwrite_output_dir \
    --logging_steps="500" \
    --save_steps="10000" \
    --eval_steps="5000"
Batch Size에 대한 노트

T5는 Encoder-Decoder 쌍 모델이라서, 기존 BERT와 GPT등 단독으로 있는 모델 대비 2배의 파라미터를 가진다. 이로 인해 max_seq_length 가 512인 경우, TPU v3-8 기준 코어당 32의 Batch size를 올릴 수 있었다.

64로 설정할 경우 각 코어의 16G HBM 메모리가 OOM이 발생하였다. (TPU v3-8은 코어당 16G, 8core로 총 128G의 HBM 메모리를 갖고 있다.)

또한, 실제 학습시 eval_steps 를 큰 값으로 두는 것이 전체 학습시간에 유리하다.

Huggingface Datasets Cache에 대한 노트

Huggingface에서 제공하는 datasets 라이브러리의 .map 함수를 사용할 경우, 자동화된 Multiprocess 처리를 비롯해 보다 쉬운 데이터 가공이 가능하다.

한편, 기본 옵션이 유저 루트(~) 내부에 캐시 폴더를 만들어 중간 처리마다 캐싱을 한다.

이로 인해 약 32GB의 데이터셋이 → 아래와 같은 약 137GB의 중간 캐시 파일을 포함한 파일들로 바뀐다.

용량이 넉넉하다면 문제가 없지만, TPU VM에 별도의 Disk를 마운트하지 않는다면 preprocessing 과정에서 서버가 멈춰버리는 문제가 생긴다 😫

이 문제를 해결하는 두 가지 방법이 있다.

  1. Disk를 추가하고, HF_DATASETS_CACHE 환경변수를 수정해 캐시 위치를 조절하기

    실제로 KcT5는 이 방법을 사용해 학습을 진행하였다.

    위 환경변수값을 조절해 추가 Disk에 저장하도록 진행해주면, 수월하게 작업 진행이 가능하다.

  2. Cache 옵션을 끄기

    Datasets 라이브러리에서는 캐시를 끌 수 있다.

    아래와 같이 캐시를 아예 끄는 코드를 파일 최상단에 작성해 사용할 수 있다.

    from datasets import set_caching_enabled
    set_caching_enabled(False)

    그리고, 만약 RAM이 넉넉하다면 datasets.config.IN_MEMORY_MAX_SIZE 를 늘려줘 해당 값들을 메모리상에 올려둔 상태로 둘 수도 있다. (다만, 얼마나 램이 많아야 할까? 🤔)

학습 완료된 모델 불러오기

해당 학습은 Flax로 진행하였기 때문에, 아래와 같이 추가 옵션을 붙여 로딩해줘야 한다.

💡

당연히 PyTorch가 설치되어있어야 한다!

from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained('./KcT5-base/', from_flax=True)
tk = AutoTokenizer.from_pretrained('./KcT5-base/')

Reference