Skip to content

ifuseok/electra-tf2-gpu

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

2 Commits
 
 
 
 
 
 
 
 

Repository files navigation

ELECTRA For TensorFlow2(GPU)

  • GPU를 활용하여 ELECTRA를 학습하는 과정 중 Google's implementation에서 공개한 코드를 사용할 때 학습 Loss는 떨어져도 fine-tune이 잘 되지 않는 문제가 있었습니다.

  • GPU를 활용한 최적화된 코드를 찾던 중NVIDIA DeeplearningExamples에서 공개된 Tensorflow2 로 구성된 코드를 사용하는 것이 적합하였습니다.

  • NVIDIA 측에서 공개된 코드가 ELECTRA base,large 모델을 학습할 수 있도록 코드가 구성되어 있어 small, small++ 를 학습 가능하도록 추가 수정하였습니다.

  • 또한 공개되어 있는 torch로 된 finetune 코드를 tensorflow 전환한 것도 추가하였습니다. torch보다 tensorflow가 익숙하신분들이 사용하시면 될것 같습니다.

Training

Feature ELECTRA
LAMB Yes
Automatic mixed precision (AMP) Yes
XLA Yes
Horovod Multi-GPU Yes
Multi-node Yes
  • 모델 학습시 지원되는 기능을 보면 위와 같음을 확인할 수 있습니다.
  • Multi-GPU 환경에서 학습 가능한 Horovod와 Automatic Mixed Precision(AMP)로 float16으로 연산하여 GPU로도 빠른 학습이 가능합니다.
  • 또한, 학습 단계를 1,2단계로 구분하여 96% 데이터는 sequence의 max_length길이를 128로, 4%의 데이터만 max_length길이를 512로 학습하여 빠르게 학습합니다.
  • BASE 모델 기준(데이터 30GB) v100 16GB 4대 기준 약 7일이면 학습이 가능합니다.
    1. Phase 1(max length : 128, 데이터 : 96%) : 5 day
    2. Phase 2(max length : 512, 데이터 : 4%) : 2 day

Pretrain 실행 방법

  1. Make TFRecord
  • make_custom_dataset.sh 내에서
    • --corpus-dir에 corpus 데이터 경로
    • --output-dir에 tfrecord output 경로
    • --num-processes에 분할된 corpus 문서 수 만큼 병렬처리 설정하면됩니다.
sh make_custom_dataset.sh
  1. Modify Config
  • scripts/config/pretrain_config.sh에서 ELECTRA pretrain 학습에 적용할 옵션들이 정의되어 있으며 수정하여 사용할 수 있습니다. 사용 가능한 GPU 재원에 따라 옵션을 선택하거나 수정하면 됩니다.
  • electra_model 에서 model type은 base,large,small,small++ 로 선택 가능합니다.
  • 현재 기본 vocab_size는 32200으로 되어 있으며 vocab size를 수정하실 경우 run_pretraining.pypretrain_config.py의 vocab_size를 수정해야하며, GPU OOM 에러 발생시 pretrain_config.sh 내 train_batch_size를 수정하여야 합니다.
Model Hidden layers Hidden unit size Parameters
ELECTRA_SMALL 12 encoder 256 14M
ELECTRA_SMALL++ 12 encoder 512 47M
ELECTRA_BASE 12 encoder 768 110M
ELECTRA_LARGE 24 encoder 1024 335M
  1. Run pretrain
  • scripts/run_pretraining_modify.sh 내에서 TFRecord가 존재하는 경로로 각각 DATA_DIR_P1,DATA_DIR_P2 를 수정하고 코드가 있는 경로로 CODEDIR 수정하여 실행할 수 있습니다.
  • pretrain_config.sh 내에서 아래와 같이 옵션을 선택하여 학습을 시작할 수 있습니다.
$ sh scripts/run_pretraing_custom.sh $(source scripts/configs/pretrain_config.sh && dgx1_4gpu_amp)
  1. Transform ckpt to h5
  • 학습이 완료된 후 transformers라이브러리로 weights를 불러오고 싶을 경우 ckpt->h5 , h5->bin 으로 변환하여 pytorch와 tensorflow 에서 사용 가능합니다.(h5를 pytorch 코드로 from_tf=True옵션으로 불러와 사용할 경우 GPU memory 이슈 발생할 수 있음)
  • scripts/extract_weights.sh 내에서 amp 사용 여부에 따라 --amp를 추가 및 학습한 vocab,max_seq_length 크기에 맞게 --vocab_size--max_seq_length를 수정하고 --pretrained-checkpoint,--output-dir를 각각 pretrain으로 얻은 ckpt 경로와 h5를 저장할 경로로 변경하면 됩니다.
  • pretrain 학습에 사용했던 config 옵션을 적용하여 실행하면 --output-dir내에 generator,discriminator 경로로 각각 config.jsontf_model.h5가 저장됩니다.
$ sh scripts/extract_weights.sh
  • transformers 라이브러리 로 모델을 불러올 때 TFElectraModelfrom_pretrained() method로 h5가 저장된 경로로 path를 설정하면 바로 사용 가능하며 torch 버젼으로 사용하고 싶을 경우 ElectraModelfrom_pretrained(from_tf=True)로 한번 weights를 불러온 후 save_pretrained()pytorch_model.bin으로 변환하여 저장한 후 다시 불러오면 GPU 메모리 이슈없이 torch로도 사용 가능합니다.

Requirements

Reference

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published