KcT5 Pretraining on TPU (feat. Flax)

데이터셋

학습환경 설정

TPU VM (feat. SSH)

## 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 연결
gcloud alpha compute tpus tpu-vm ssh tpuvm --zone europe-west4-a

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

## 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

Flax, Jax

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

학습

토크나이저 학습

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('.')

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 생성

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('.')

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

모델 학습

wget -nc -q https://raw.githubusercontent.com/huggingface/transformers/762416ffa8b87ceca2d35e56c575ef3c41f449e7/examples/flax/language-modeling/run_t5_mlm_flax.py

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

# (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"

Huggingface Datasets Cache에 대한 노트

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

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

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

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

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

1. Disk를 추가하고, HF_DATASETS_CACHE 환경변수를 수정해 캐시 위치를 조절하기
실제로 KcT5는 이 방법을 사용해 학습을 진행하였다.
위 환경변수값을 조절해 추가 Disk에 저장하도록 진행해주면, 수월하게 작업 진행이 가능하다.

2. Cache 옵션을 끄기
Datasets 라이브러리에서는 캐시를 끌 수 있다.
참고: https://huggingface.co/docs/datasets/v1.12.0/cache.html
아래와 같이 캐시를 아예 끄는 코드를 파일 최상단에 작성해 사용할 수 있다.

from datasets import set_caching_enabled
set_caching_enabled(False)

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

학습 완료된 모델 불러오기

from transformers import AutoModel, AutoTokenizer

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

Reference

• T5 arxiv https://arxiv.org/abs/1910.10683

• KeT5 https://huggingface.co/KETI-AIR/ke-t5-base

• Huggingface Flax T5 Pretraining Example https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/README.md#t5-like-span-masked-language-modeling