Zenn
🔥

パラメタ数1.5Bのgpt2-XLを学習した話

2022/12/15に公開
Python
自然言語処理
GPT
tech

この度、gpt2論文を参考に最大サイズのgpt2の日本語版モデルを学習し公開いたしました。
https://huggingface.co/nlp-waseda/gpt2-xl-japanese

この記事では学習にあたり工夫した点や問題点等を書き連ねます。シングルノードですが比較的大きなモデルを学習しているので知見などを共有できればと思います。

なお学習はA100(40gb)8枚のノードを使って合計70日程かけて学習しました。

現在公開されているdecoder系モデルのうち今回作成したモデルに近い規模のモデルは知る限り2つあり、こちらはgpt3論文を参考にパラメタ設定をしていると考えられ、今回作成したモデルとは層の深さと隠れ層の次元が異なります。

rinna/japanese-gpt-1bは生成される文章が短い傾向にあります。学習時の入力文の多くが短いものであったと推測されます。また分割処理が独特です。
yellowback/gpt-neo-japanese-1.3Bは生成される文章が少し固いように感じましたが、気のせいかもしれません、、

学習データの用意

お試しなので、日本語wikipediaとcc100を学習データとして使います。日本語wikipediaは2022/07/01時点のものを利用しました。

wget http://data.statmt.org/cc-100/ja.txt.xz
wget https://dumps.wikimedia.org/jawiki/20220701/jawiki-20220701-pages-articles-multistream.xml.bz2

xz -d ja.txt.xz
bzip2 -d jawiki-20220701-pages-articles-multistream.xml.bz2

日本語wikipediaデータの前処理をしていきます。

python3 -m venv .env
source .env/bin/activate

# xmlからテキストを抽出
pip install wikiextractor==3.0.4
wikiextractor jawiki-20220701-pages-articles-multistream.xml
find text/ | grep wiki | awk '{system("cat "$0" >> wiki.txt")}'
rm -rf text/
# 残ったタグを除去
sed -i 's/<[^>]*>//g' wiki.txt
# 空行を除去
sed -i '/^$/d' wiki.txt

次にテキストを正規化し、行単位で分割し、jumanppで分かち書きをするという処理を一度にやっていきます。

split.py
import sys

from textformatting import ssplit
from tqdm.auto import tqdm
import neologdn
input = sys.stdin.readline

MAX_LINES=int(sys.argv[1])

def main():
    line = 'start'

    bar = tqdm(total = MAX_LINES)
    while line:
        line = neologdn.normalize(input().rstrip())
        [print(s) for s in ssplit(line) ]
        bar.update(1)

if __name__ == "__main__":
    main()

上のファイルを書いたら、parallelを利用してボトルネックのjumanppによる解析を並列化する。だいたい5分で終わりました。

cat wiki.txt | python split.py 15768565 | parallel --pipe -L 10000 --blocksize 1772096 jumanpp --segment > wiki_mrph.txt

cc100に関しても同様です。こちらは3時間弱かかりました。

cat ja.txt | python split.py 392774277 | parallel --pipe -L 10000 --blocksize 1772096 jumanpp --segment > cc100_mrph.txt

Tokenizerの学習

tokenizerに関しても1から学習しました。事前分割は済ませているので、いわゆるサブワード分割用の語彙を学習することになります。
gptは元々はbyte level BPEを使うとのことですが今回はBPEを採用しました。(経験上BPEよりもUnigramの方がBERTにおいてはダウンストリームタスクの性能が良くなります。)

train_tokenizer.py

filesには事前に分かち書きを行ったテキストのパスを指定します。

train_tokenizer.py
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import Whitespace

tokenizer = Tokenizer(BPE(unk_token="<unk>"))
tokenizer.pre_tokenizer = Whitespace()
trainer = BpeTrainer(
    vocab_size=50000,
    min_frequency=1,
    special_tokens=["<unk>", "<s>", "</s>"],
    limit_alphabet=8000,
    )
files = ["wiki_mrph.txt", "cc100_mrph.txt"]
tokenizer.train(files, trainer)
tokenizer.save("juman-bpe-wiki-cc100.json")

簡単に説明すると
pre_tokenizerにWhitespaceを指定しているので事前にスペースが入っている箇所は必ず単語の境界となります。今回の例では事前にjuman++で単語境界にスペースを入れているので、juman++によって設定された単語境界が必ず最終的な単語境界に含まれるようになってます。
しかしながら、この処理を行うことで英文と和文が混同しているような場合、出力時の英文デコードが面倒になるのでもう少し良い方法があればなと思っています。次回以降はこの点も改良点にしたいと考えています。

BpeTrainerに関して、語彙数は本家のgpt-2を参考に50000と設定しました。limit_alphabetは語彙で使用する文字数です日本語の場合は扱う文字の種類が豊富なのである程度大きな値を設定する必要があります。今回は8000としました、もう少し小さくても良い気がします。

python train_tokenizer.py

モデルのアーキテクチャについて

学習開始時は1.3Bを超えるサイズのモデルはあまりなかったので使えるハードウェアリソースと時間的リソースの観点からこれと同規模のモデルを作成しました。
いくつか存在する1.3BのGPTモデルはGPT3の論文を参考にパラメタが設定されており、今回学習したGPT2とは隠れ層の次元数が大きくなっています。
以下にrinna/japanese-gpt-1bとのアーキテクチャ構造の違いを表にまとめます。

rinna/japanese-gpt-1b nlp-waseda/gpt2-xl-japanese
トークンの埋め込み表現の次元数 2048 1600
Attention Headの数 16 20
Attention Head 1個あたりの次元数 2048/16=128 1600/20=80
レイヤ数 24 48
パラメタ数 1.3B 1.5B
語彙数 44928 50000

このようにレイヤ数が大幅に異なりますが性能にどのように影響するかは認識しておりません。もし知見のある方がいましたらコメントいただけると嬉しいです。

学習に関して

このモデルをGPUに載せると約7GBほどのメモリを消費します。推論はモデルとデータがGPUに乗れば良いですが、学習においてはoptimizerの情報と勾配の情報も乗せる必要があり、モデルそのもののサイズより大幅に消費メモリが増えます。そのため今回使用したA100(40gb)ではbf16を採用してもそのままでは学習を進めることができません。

そのため分散学習を検討する必要があります。比較的小規模なモデルの学習ではよく用いられるであろうDPやDDPはデータを分割するという特性上学習の高速化には貢献しますが、メモリに乗り切らない問題に関しては効果がないです。
分散学習の手法については様々なものが存在しますが、Huggingfaceのこちらの記事に詳しく書かれていますので気になる方はご参照ください。

今回は様々な分散学習の手法のうちDeepSpeedで提供されているZeRO2という学習手法を採用しました。ZeRO2は学習時の勾配とoptimizerの情報を分割してそれぞれのGPUに乗せることでメモリを大幅に節約する手法です。
HuggingFaceはDeepSpeedとの連携もサポートされており、ZeRO2による学習は数行書き換えるだけで実現できます。

HuggingFace TransformersのTrainerクラスを使って学習する場合、Trainerクラスのdeepspeed=パラメタにds_config.jsonのパスを渡し、起動コマンドをtorchrunやpythonからdeepspeedに変更することで実現できます。今回の学習では以下のような設定で学習を行いました。基本的にautoに設定することでHuggingface上での設定を優先するようにしています。DeepSpeedを利用することで本来メモリに乗り切らない巨大なモデルでも学習することができます。具体的な統合方法に関してはこちらの記事を参考にしてください。

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "bf16": {
	"enabled": "auto"
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "total_num_steps": "auto",
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },

    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true,
        "cpu_offload": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto"
}

残りの学習に関してはHuggingface Transformersリポジトリのこちらのスクリプトを参考にしました。

学習時のパラメタに関しては大まかには以下のようになっています。詳しくはこちらを参考にしてください。事実上のバッチサイズは3*8*22=528となっています。

per_device_train_batch_size: 3
gradient_accumulation_steps: 22
learning_rate: 3e-04
weight_decay: 0.1
adam_beta2: 0.95
num_train_epochs: 10

簡単に使ってみる

パイプラインを使って簡単に試してみます。注意点として入力するテキストは学習データの前処理の都合上

  • 事前にjuman++によるスペース区切りでの単語分割処理が行われている
  • neologdnで正規化されている(tokenizerの学習時に正規化関数も指定できるが、日本語に関してのNFKC正規化が怪しかったのであまり当てにしないでほしい)

ことが望ましいですが、そのまま生文を入れてもある程度なんとかなります。
juman++のインストールに関しては以前に記事を書いているのでそちらを参照してください。

do_sample=Trueというオプションを渡すことで生成トークンが必ずしも尤度最大のものではなくなります。

from transformers import pipeline, set_seed
generator = pipeline('text-generation', model='nlp-waseda/gpt2-xl-japanese')

set_seed(42)
generator("早稲田 大学 で 自然 言語 処理 を", max_length=30, do_sample=True, pad_token_id=2, num_return_sequences=5)
[{'generated_text': '早稲田 大学 で 自然 言語 処理 を 勉強 して いる 大学生 です. 自然 言語 処理 や 音声 認識, 機械 学習 等 に 興味 が あり, 特に 画像'},
 {'generated_text': '早稲田 大学 で 自然 言語 処理 を 学んで いる と ある 方 と お 会い して き ました. 今日 は お 話 する 時間 が 少なかった のです が,'},
 {'generated_text': '早稲田 大学 で 自然 言語 処理 を 研究 して いる が 、 それ を 趣味 と は 思わず 、 会社 を 作る ため の 手段 と とらえて いる ようです 。'},
 {'generated_text': '早稲田 大学 で 自然 言語 処理 を 専門 的に 学ぶ サークル です 。 日本 語 教育 センター で 日本 語 を 勉強 した 中国 の 人 たち と 交流 する'},
 {'generated_text': '早稲田 大学 で 自然 言語 処理 を 専攻 した 時 に 、 数学 の 知識 ・ プログラミング 言語 の 知識 が 身 に ついて いた の は 、 とても 役'}]

通常のAPIから利用する場合はこのようになります。

from transformers import AutoTokenizer, GPT2Model

tokenizer = AutoTokenizer.from_pretrained('nlp-waseda/gpt2-xl-japanese')
model = GPT2Model.from_pretrained('nlp-waseda/gpt2-xl-japanese')
text = "早稲田 大学 で 自然 言語 処理 を"
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)

最後に

Huggingface TransformersのTrainerクラスとDeepSpeedの連携が容易なため非常に簡単にZeRO2によるメモリ最適化をおこなったDDP学習を行うことができました。
途中でも触れていますが、学習時にメモリが枯渇する(RuntimeError: CUDA error: out of memory)主な原因はoptimizerがメモリを多く消費するためであり、複数枚のGPUが扱える場合これらを分割することで計算効率を極力落とさずにメモリを節約できます。
今回作成したモデルに不備があると感じた方や記事に記載しきれていない部分で気になる箇所がある方はお気軽にご連絡いただければと思います。

なお本モデルの学習にあたり、データ活用社会創成プラットフォーム mdxを利用しました。

GitHubで編集を提案

Discussion