LLM事前学習の流れ

これからLLMのパラメータを記述した設定ファイルからLLMのモデルを作成し、事前学習を試してみたいと思います。

今回は、LLMの生成と学習の流れを確認するためにごく小規模なLLM（というかSLM）を作って「私は歩く」という文を学習させてみましょう。

トークナイザーの作成

LLMのモデルの前に、まず文字列をLLMで扱う数値配列データ（トークン列）に変換するためのトークナイザーを作ります。

今回は、「私は歩く」という文を学習させてみるので、「私」「は」「歩」「く」と一文字ずつトークン化してみましょう。また、トークナイザーにはpadトークンも必要なので、学習文字列に含まれる4つの文字とpadで計5個のトークンを扱うトークナイザーを作ることにします。

トークナイザーの作成は、トークンに含まれる語と設定を含むjsonファイルを作成し、それをPreTrainedTokenizerFastクラスに読み込ませトークナイザーのディレクトリとして保存する形で行います。

{
    "version": "1.0",
    "truncation": null,
    "padding": null,
    "normalizer": null,
    "pre_tokenizer": null,
    "post_processor": null,
    "model": {
        "type": "Unigram",
        "vocab": [
            ["<:pad:>", 0],
            ["私", 1],
            ["は", 2],
            ["歩", 3],
            ["く", 4]
        }
    }
}

このjsonをtokenizer_src.jsonとして保存したら、PythonにPyTorchとTransformersをpipでインストールして以下のプログラムを実行します。実行後にtokenizerディレクトリにトークナイザーが作成されます。

from transformers import PreTrainedTokenizerFast

tokenizer = PreTrainedTokenizerFast(tokenizer_file="tokenizer_src.json", pad_token="<:pad:>")

tokenizer.save_pretrained(save_directory="tokenizer")

設定ファイルからのモデル作成と学習

トークナイザーが出来たので、トークナイザーを使ってモデルで使うためのデータを保持するDatasetを作ってみます。

今回は、「私は歩く」という一つの文を学習するだけですから、とりあえず長さは「1」でデータとして「私は歩く」をトークン化したデータを返すだけのものにしてみました。

class TestDataset(Dataset):

    def __init__(self, tokenizer):

        input_ids = tokenizer.encode("私は歩く", return_tensors='pt')[0]
        self.token_list = {'input_ids': input_ids}

    def __len__(self):
        return 1

    def __getitem__(self, idx):
        return self.token_list

では、いよいよモデルを作成して学習をしてみましょう。モデルは、GPTNeoXForCausalLMに設定ファイルから作ったconfigを渡す形で作ります。

今回は、ごく小規模なLLMで試してみたいのでOpenCALM-Smallのパラメータ設定を元に以下の設定ファイルを作ってみました。

{
    "architectures": [
        "GPTNeoXForCausalLM"
    ],
    "hidden_act": "gelu",
    "hidden_size": 768,
    "initializer_range": 0.02,
    "intermediate_size": 3072,
    "layer_norm_eps": 1e-05,
    "max_position_embeddings": 2048,
    "model_type": "gpt_neox",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "rotary_emb_base": 10000,
    "rotary_pct": 1.0,
    "tie_word_embeddings": false,
    "torch_dtype": "float32",
    "transformers_version": "4.42.0",
    "use_cache": true,
    "use_parallel_residual": false,
    "vocab_size": 5
}

パラメータ数は160Mとのことで、かなり少ないですね。

この設定ファイルをmodel_config_small.jsonとして保存したら、以下のコードでモデルを作成できます。

model_config = GPTNeoXConfig.from_json_file("model_config_small.json")

model = GPTNeoXForCausalLM(model_config)

続いて、Datasetを使ってモデルにデータを供給するcollatorと学習時のパラメータを保持するTrainingArgumentsを作ります。

collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

training_args = TrainingArguments(output_dir="output", save_steps=1, num_train_epochs=1, learning_rate=5e-5, lr_scheduler_type="constant", per_device_train_batch_size=1)

とりあえず学習率5e-5（0.00005）で１回だけ学習するようにしてみました。

モデルを作成して事前学習を行うコードは、以下のようになります。

import torch
from torch.utils.data import Dataset
from transformers import AutoTokenizer, DataCollatorForLanguageModeling
from transformers import GPTNeoXConfig, GPTNeoXForCausalLM
from transformers import Trainer, TrainingArguments

class TestDataset(Dataset):

    def __init__(self, tokenizer):

        input_ids = tokenizer.encode("私は歩く", return_tensors='pt')[0]
        self.token_list = {'input_ids': input_ids}

    def __len__(self):
        return 1

    def __getitem__(self, idx):
        return self.token_list

model_config = GPTNeoXConfig.from_json_file("model_config_small.json")

model = GPTNeoXForCausalLM(model_config)
tokenizer = AutoTokenizer.from_pretrained("tokenizer")

dataset = TestDataset(tokenizer)

collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

training_args = TrainingArguments(output_dir="output", save_steps=1, num_train_epochs=1, learning_rate=5e-5, lr_scheduler_type="constant", per_device_train_batch_size=1)

trainer = Trainer(model=model, data_collator=collator, args=training_args, train_dataset=dataset)

trainer.train()

accelerateをインストールして実行すると、学習結果がoutput/checkpoint-1に保存されます。

推論

最後に学習したモデルを読み込んで、「私」に続く語を4つほど生成してみます。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("output/checkpoint-1")
tokenizer = AutoTokenizer.from_pretrained("tokenizer")

input = tokenizer("私", return_tensors="pt", return_token_type_ids=False)
texts = model.generate(**input, max_new_tokens=6, pad_token_id=tokenizer.pad_token_id, do_sample=True, num_return_sequences=4)

for text in texts:
    print(tokenizer.decode(text, skip_special_tokens = True))

生成結果

私 は
私 は
私 は
私 く く は

ごく小規模なモデルで１回だけの学習ですが、「私」に続いて「は」を返す傾向が出てきているようです。データ文を「学習」していると言えるでしょう。

創作プログラミングの街