【手順付き】初心に立ち返ってGPT-2を学習してみる

はじめに

ABEJAでデータサイエンティストをしている真鍋です。アドベントカレンダーぶりの登場になります。今回も生成AI、特にLLM (大規模言語モデル) 系のネタです。

LLMが登場して久しく、業務でLLMを使うのが当然の方も多くなっているのではないかと思います。裏では、私のようなデータサイエンティストやエンジニアが、APIを通じてGPT等の言語モデルを利用する場合も一定あるのではないでしょうか。一方で、ただAPIを叩くだけではわからない部分も多いため、実際に作ってみようじゃないか、というのが、今回の記事の趣旨です。とはいえ、実際の大規模言語モデルは、GPUを何台も用意するレベルで、がっつりリソースを使うことになるので、さすがに厳しいので、簡単に小規模言語モデル (GPT-2) を作成してみよう、というのが本記事の趣旨です。

LLMとは？

LLMの基礎知識

LLMの原理的な話は、他の記事でも多く記載があるので、そちらに譲りますが、言語モデルは一言で言えば、与えられたテキストから単語を予測するためのシステムです。従来は統計的手法やルールベースのアプローチが一般的でしたが、近年では大規模なデータセットを使って事前学習されたモデルが主流です。これにより、非常に精度の高いテキスト生成や分類が可能となっています。

tech-blog.abeja.asia tech-blog.abeja.asia tech-blog.abeja.asia

API利用のメリットと限界

既存のLLM APIを利用する最大のメリットは、モデルのトレーニングに要する膨大な計算リソースや時間を節約できる点です。また、APIは一般に高い性能を発揮し、幅広いタスクに対応可能で、特に自然言語処理における複雑なタスクや、迅速なプロトタイピングにおいては非常に有用です。一方で、API利用には制約もあります。例えば、モデルの挙動を完全にカスタマイズすることが難しい場合があります。また、特定のタスクやドメインに最適化された結果が必要な場合には、APIが期待通りに動作しないこともあります (ここは、LLMを業務に組み込む経験のある方なら、身をもって体感されているのではないでしょうか)。さらに、利用料金やデータのプライバシーにも注意が必要です。

準備

開発環境のセットアップ

今回は、こちらの記事を参考に、Google Colaboratory (以降、Colab) で完結できるレベルの(超) 小規模言語モデルを作ろうと思います ~~(4年も前にこのハンズオンが公開されているのかと思うと、もっと早くこの波に気づきたい人生だった)~~ 上記の記事では、大きく、下記の6つのステップで、エスペラントという人工言語の言語モデルを作ろうとしています。

データセットの準備
トークナイザーの学習
言語モデルの学習
学習した言語モデルの確認
下流タスクに向けてファインチューニングを実施
チューニングしたモデルのシェア

また、Colab Notebookも公開されているため、そちらをベースにまずは動かしてみようと思います。 1. データセットの準備と、2. トークナイザーの学習については、今回は日本語特化のモデルを作ってみようと思っているので、既存のものをフル活用させてもらいます。

前提となるパッケージの準備

ColabにはデフォルトでTensorFlowが入っていますが、今回は使わないため、アンインストールし、datasetsやtransformersをインストールします。

# TensorFlowは今回は使わない
!pip uninstall -y tensorflow

# accelerateのバージョン起因でTrain時にエラーが出るのでアップグレードしておく
!pip install --upgrade accelerate
!pip install datasets

# transformersはマスターからインストールする
!pip install git+https://github.com/huggingface/transformers
!pip install transformers["ja"]
!pip list | grep -E 'transformers|tokenizers'

元のColab時点ではtransformersのバージョンは2.11.0, tokenizersのバージョンは0.8.0rc1でしたが、私がColabを実行した2024年8月末時点では、transformers: 4.45.0.dev0, tokenizers: 0.19.1で動作を確認しています。

データ収集と前処理

前提として、言語モデルをトレーニングするためには、大量のテキストデータが必要ですが、今回は手軽に利用できるオープンデータセットとして、HuggingFaceからデータを収集しています。今回は、CC100を利用しました。弊社の過去のGPT開発においても使用したデータセットで、Wikipediaに次いで、日本語言語モデルで多く使用されていて、元々はfacebookresearch/XLMの学習に使用されたが、エジンバラ大学の協力のもと数多くの言語に翻訳されている (同記事より) のが特徴です。

huggingface.co

from datasets import load_dataset

hf_dataset = load_dataset("statmt/cc100", "ja", split = "train", streaming = True)

この2行でぱぱっとデータセットが読み込めるなんて便利すぎますね… なお、streamingをTrueにしないと、データが全量一度ダウンロードされてしまい、容量エラーになるので、streaming = Trueで実行しています。また、CC100固有かと思いますが、カスタムコードを実行するかどうか (Do you wish to run the custom code?) を質問されるので、そちらはyと回答してください。

streaming = Trueで実行すると、IterableDatasetという型でデータセットが保持されます。これを一旦下記のように読み込んでみます。

dataset_head = hf_dataset.take(5)
for data in list(dataset_head):
  print(data)

出力結果

{'id': '0', 'text': '午後から雨が心配だったので遠出はせず、『ふれあいロード』を走って来ました！\n'}
{'id': '1', 'text': '確実に春が近づいてることを肌で感じることが出来ました 着々と整備されてる圏央道を越えるとお世話になってるボウリング場が見えて来ました。\n'}
{'id': '2', 'text': 'うぅ〜〜、私が途中でトイレに行きたくなってしまい、通り道にあったケンタに変更しちゃいました。\n'}
{'id': '3', 'text': '実は、１年程前にエルモサの右目の黒目の端によ〜く見ないと分からない程の小さな斑を見つけてたんです。\n'}
{'id': '4', 'text': 'その時点で先生からはおそらく『角膜ジストロフィー』であろうとの診断をもらっていました。\n'}

本来であれば他のデータセットも使いたかったのですが、CC100の日本語だけで15GBもあり、上記ページによると、1.5GBのセルビア語のデータセットでも、学習用データセットだけで35,747,957件あるそうなので、単純に3.5億以上のデータ数がありそうです。そしてColabで学習を完結させるにはあまりにも膨大なので、この中の一部を使うこととします。

本来であれば、データのクリーニング (不要な文字やタグの除去) を行うべきですが、まずはクイックに動かしたいため、特段手を入れていません ~~(沼の匂いがするので)~~。後々トークン化 (単語や文字の分割) やベクトル化 (テキストの数値ベクトルへの変換) を実施し、モデルが効率的に学習できる形式にデータを整えます。

トークナイザーの選択

ここでも、簡易化のために、学習済みのrinna/japanese-gpt2-mediumトークナイザーを使おうと思います。(

import torch
from transformers import AutoModel, T5Tokenizer

tokenizer = T5Tokenizer.from_pretrained("rinna/japanese-gpt2-medium")

これだけでトークナイザーを宣言できてしまうのですが、せっかくなので軽く動作確認をしてみます。

# 試してみる
tokenizer.pad_token = tokenizer.eos_token
line = "吾輩は猫である。"
print(tokenizer(line)["input_ids"])

出力

[9, 5361, 31082, 11, 4324, 27, 8, 2]

次に、いよいよ学習のための準備をしていきます。

GPT-2の学習

学習のための設定

ここではHuggingFaceのTrainerクラスを使用し、学習を行なっていきます。まずは、GPT2Configで、モデルの設定をしていきます。

from transformers import GPT2Config

config = GPT2Config(
    vocab_size=tokenizer.vocab_size,
    n_positions = 512,
    n_head=2,
    n_layer=1
)

本来であればもう少し大きな数値にしたいところですが、課金せずにColabでGPUを使うか、GPUを使わずにColabで完結させたいため、まずは小さめのパラメータから始めます。ちなみにGPUを使う場合は、念の為動作確認からしています。

!nvidia-smi

出力

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                       Off | 00000000:00:04.0 Off |                    0 |
| N/A   56C    P8              12W /  70W |      0MiB / 15360MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|  No running processes found                                                           |
+---------------------------------------------------------------------------------------+

次に、PyTorchがGPUを認識しているか、念の為確認します。

import torch
torch.cuda.is_available()

ここで、Trueと出力されたらOKです。GPUを使う場合、より高速に学習が回るので、少しパラメータを大きくしています。

from transformers import GPT2Config

config = GPT2Config(
    vocab_size=tokenizer.vocab_size,
    n_positions = 4096,
    n_head=16,
    n_layer=8
)

モデルの初期化

次にモデルの初期化ですが、ここではGPT2を使うため、GPT2LMHeadModelを使います。

ここはハマりポイントでもあったのですが、GPT2ModelではなくGPT2LMHeadModelを使っているのは、GPT2Modelは次元変換（エンコーダーとしての役割）を行うモデルであり、言語モデルの学習時に使われる「損失計算用のラベル（labels）」を受け取ることができない、という理由です ~~ChatGPTに教えてもらいました~~。

また、モデルのConfigを、既存の学習済みオープンモデルや、チェックポイントから読み込んでもよかったのですが、スクラッチでまずは動かしてみる、ということを考えて、Configで初期化するにとどめています。

from transformers import GPT2LMHeadModel

model = GPT2LMHeadModel(config=config)

ここでmodel.num_parameters()を実行すると、GPUなし版では32058624が返ってくるため、約3205M (0.32B) のパラメータであることがわかります。また、GPUあり版では84426240となるため、84,426M (84.4B) であることがわかります。こちらの記事によると、GPT-3の時点で175,000M (175B)、GPT-4になると1,800,000M (1.8T) パラメータとなるため、あまり大きくはなさそうなことがわかりますね…

学習用データセットの最終化

先ほどロードしたデータセットにトークナイザーを適用し、学習用データセットを準備します。前述のように、IterativeDatasetをDatasetに変換する必要があります。データセットの変換について、詳細はHuggingFaceの公式ドキュメントもあるので、そちらも参照しています。

from functools import partial
from datasets import Dataset

def gen_from_iterable_dataset(iterable_ds):
    yield from iterable_ds

# メモリがクラッシュしてしまうので、簡易化のためにデータセットを1万件に絞る
dataset_use = hf_dataset.take(10000)

dataset = Dataset.from_generator(
    partial(gen_from_iterable_dataset, dataset_use), features=dataset_use.features)

def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length = 512)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

一旦学習データについては、上記のコメントの通り、先頭の1万件を抽出しています。ここはColabで完結させるためにメモリエラーを回避したい、実行時間をなるべく数時間程度に抑えたい、という板挟みの結果でした…(GPUありの場合は15万件まで抽出しています)

ここで、tokenized_datasetsを確認すると、下記のような出力になるはずです。

Dataset({
    features: ['id', 'text', 'input_ids', 'attention_mask'],
    num_rows: 10000
})

また、この時にData Collatorを準備します。data_collatorはちょっとしたヘルパーで、データセット内の異なるサンプルを、PyTorchが誤差逆伝播 (backprop) できるようなオブジェクトにバッチできるようになります (参照元のColabを翻訳) というもののようです。

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15
)

Trainerの初期化と学習の実行

最後に、Trainerを初期化し、学習を実行します。ここも、学習時間を数時間で抑えたい…という願望の結果、かなり控えめな設定になっています。

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./SLM-exp/model",
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=32,
    save_steps=1000,
    save_total_limit=2,
    prediction_loss_only=True,
    fp16=True
)

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    data_collator=data_collator,
    train_dataset=tokenized_datasets,
)

また、GPUありの場合は、per_device_train_batch_sizeを16に下げています (大きくしたかったのですが、32にするとGPUのメモリエラーが出てしまったため)。最後に、下記コードを実行すると、学習が走ります。

%%time
trainer.train()

GPUありだと、これで75〜80分の学習で、9375ステップを実行していました。GPUなしだと、6時間超の学習で313ステップを実行していました。最終的なtraining_lossはGPUなしで7.7798、GPUありで6.2267でした。GPUのありがたみがわかりますね…

モデルの保存

トレーニングが完了したら、モデルの性能を評価しますが、その前にモデルを保存しておきましょう。

trainer.save_model("./SLM-exp/model")

実際に動かしてみる

最後に、学習させたモデルで、実際に試してみましょう。

from transformers import pipeline

text_pipe = pipeline('text-generation', model=model, tokenizer=tokenizer, device=0)
output = text_pipe("昔々あるところに")

output[0]['generated_text']

出力 (GPUなし版)

昔々あるところには、、ので、ので、

出力 (GPUあり版)

昔々あるところに、その日は、その日、その日、その日、

…かなり厳しい結果になっているのがわかります。せっかくなので、最後に、野球ファンなら90%正解できるであろう言葉で遊んでみますｗ

output = text_pipe("はっきり言うて、")

output[0]['generated_text']

出力 (GPUなし版)

はっきり言うて、、、、、、ので、、、

出力 (GPUあり版)

はっきり言うて、のは、その日は、その日、その日、、

…おーん…😇😇😇 この後、参考記事の方ではファインチューニングや、モデルの公開まで行っているのですが、今回はモデルの改善の方が先決ですね😇

まとめと今後の展望

ということで、ここまで、Colabで完結するレベルの簡易的な言語モデルの開発にトライしてみました。パラメータ数的にはGPT-3に及ばない程度で、学習に使ったデータセットも少なかったことから、あまり良い結果とは言えませんでした (弊社のGENIACでの取り組みも含め、より大規模なモデルを開発する際の苦労が伺えました…)

今回の学びと成果

HuggingFaceのTrainerを活用しつつ、小規模言語モデルを作成することができました。精度としては、お察しの通りでしたが、「HuggingFaceのデータセットを拾ってきて」「HuggingFaceのモデルをベースに学習させる」ところをColabで完結させることができました。

今後の課題・改善点

今回はデータセットもトークナイザーもオープンなものを活用しましたが、元のColabではトークナイザーの学習もやっていましたし、データセットは、億単位でレコードがあるであろうデータセットのうち、1~15万件しか使っていないなど、まだまだ伸びしろは全然あると思います。また、この程度のデータ量であれば、何かに特化させる、という方がより実用的なモデルになる可能性が高そうですね。 ColabのGPUを使ってもう少しパラーメータ数を大きくするとか、Colabではなくもっと潤沢にリソースがある状況で、回してみるとか、やりようはいくらでもあると思いますし、他のタスクに転用してみる、といった拡張性もあるでしょう。また、今回は結局HuggingFaceのTrainerを使ったことで、HuggingFaceとは少し仲良くなれた🤗と思いますが、PyTorchで学習部分もスクラッチから組んでみるということも、大規模言語モデルの理解をあげる一助になるのではないでしょうか。その辺も含めて、第2弾にチャレンジできたら…と、思っています。