GoogleColobで小規模言語モデル(0.15B)の事前学習モデルを作ってみる

はじめに

LLMの精度検証を各モデルをやっていく中で自分でも事前学習モデルを作ってみようと思ったので、いろいろと調べて作ってみました

今回は以下の記事で事前学習モデルの作成用のColobが公開されていたので、こちらを使用します。

zenn.dev

また上記の内容を細かくZennにまとめられているので、詳細を知りたい方は以下の記事をご覧ください。

zenn.dev

成果

作成したモデルで推論した結果です....数時間の推論だけだとかなり厳しいですね

今回作成したモデルは以下で公開しています

huggingface.co

参考記事からの変更点

上記の記事から変更した点は以下になります

• 計算リソースをA100を使用
• データセットを graelo/wikipediaの20230901のjpを使ってより新しいデータセットで学習
• A100のためflash attention2を使用して学習

環境

• Google Colob A100

環境設定

# cudaバージョンを確認する
!nvcc --version

# refer:https://github.com/pytorch/pytorch/issues/107960
# torch.compile時のlibcuda.so not foundエラーの回避策
!ldconfig /usr/lib64-nvidia

!git clone https://github.com/ce-lery/japanese-mistral-300m-recipe.git
%cd japanese-mistral-300m-recipe
!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
!pip install -r requirements.txt
!pip install flash-attn==2.3.4 --no-build-isolation
# ver違いでデータセットの処理がうまくいかないので、アップグレードする
!pip install --upgrade pyarrow

データセットの構築

学習に必要なデータセットをダウンロードして、学習できるようにフォーマットなどを整えます

%%time
!GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/datasets/graelo/wikipedia.git
%cd wikipedia/data/20230901/ja/
!git lfs pull --include "train-*-of-0016.parquet"
%cd ../../../../
import pandas as pd
from datasets import load_dataset
import csv
import pandas as pd

# create dataset for training tokenizer
dataset = load_dataset("wikipedia/data/20230901/ja/", data_files= "train-*-of-0016.parquet",split="train")
# datasetのtext列をwiki.txtに書き出す
dataset.to_csv("wiki.txt", columns=["text"], sep="\t", index=False, header=False, quoting=csv.QUOTE_NONE, escapechar="\\")

必要であれば以下で不要なデータを削除します

%%bash
rm -r wikipedia/
rm -r spm_tokenizer_neologdn_bytefallback_nofast/model.safetensors
LINES=`wc -l wiki.txt | awk '{print $1}'`
TRAIN_DATA_LINES=$(($LINES*10/100))
head -n $TRAIN_DATA_LINES wiki.txt > wiki2.txt
rm -r wiki.txt
mv wiki2.txt wiki.txt

トークナイザー準備

すでに学習済みの ce-lery/japanese-mistral-300m-base を使います

# todo 学習済トークナイザーをダウンロード
!git clone https://huggingface.co/ce-lery/japanese-mistral-300m-base.git
# spm-wiki-cc100-for-spm-bytefallbackという名称で保存
!mv japanese-mistral-300m-base spm_tokenizer_neologdn_bytefallback_nofast

事前学習

学習用のパラメータの設定

モデルのサイズやモデル自体のパラメータを /content/japanese-mistral-300m-recipe/pretrain/train/mistral-300m/config.json を変更します。(変更後は以下のようになります)

パラメータは rinna/japanese-gpt2-smallを参考にしました

{
    "architectures": [
      "MistralForCausalLM"
    ],
    "bos_token_id": 0,
    "eos_token_id": 0,
    "hidden_act": "silu",
    "hidden_size": 768,
    "initializer_range": 0.02,
    "intermediate_size": 2400,
    "max_position_embeddings": 2048,
    "model_type": "mistral",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "num_key_value_heads": 6,
    "rms_norm_eps": 1e-05,
    "rope_theta": 10000.0,
    "sliding_window": 1024,
    "tie_word_embeddings": false,
    "torch_dtype": "float16",
    "transformers_version": "4.35.2",
    "use_cache": true,
    "vocab_size": 50257
  }

学習時のパラメータを以下のように設定します

%%bash

# model_typeとconfig_nameがgpt2-mediumになっているが、内部で"mistral"に上書きされるため無視でいい
# 1epoch=5657stepsなので、100stepで0.02epochぐらい. T4の場合、1epochで100hourぐらいかかるので、100stepに
cat << EOS > hf_config_quick_start.json
{
    "model_type": "gpt2",
    "config_name":"gpt2-medium" ,
    "tokenizer_name":"../spm_tokenizer_neologdn_bytefallback_nofast" ,
    "train_file":"../wiki.txt",
    "validation_split_percentage":5,
    "output_dir":"checkpoints-mistral-300M-FA2",
    "do_train":true,
    "do_eval":true,
    "prediction_loss_only":true,
    "remove_unused_columns":false ,
    "learning_rate":6.0e-4 ,
    "weight_decay":0.1 ,
    "adam_beta2":0.95 ,
    "max_steps":100,
    "logging_dir":"checkpoints-mistral-300M-FA2/logs",
    "logging_strategy": "steps" ,
    "logging_steps":10 ,
    "evaluation_strategy":"steps" ,
    "save_strategy": "steps" ,
    "eval_steps":100 ,
    "save_steps":100 ,
    "load_best_model_at_end":true ,
    "save_total_limit":2 ,
    "warmup_steps":1000 ,
    "lr_scheduler_type":"cosine" ,
    "per_device_train_batch_size":16 ,
    "per_device_eval_batch_size":16,
    "block_size":1024 ,
    "adam_epsilon":1.0e-4 ,
    "fp16":true ,
    "gradient_accumulation_steps":256,
    "push_to_hub":false,
    "dataloader_num_workers": 8,
    "optim":"adamw_bnb_8bit" ,
    "torch_compile":true
}
EOS

事前学習の実行

%%time
%cd pretrain/

# 前の学習でできた空の学習済モデルフォルダを削除
!rm -r checkpoints-mistral-300M-FA2
!deepspeed --no_local_rank train/run_clm.py ../hf_config_quick_start.json --deepspeed --deepspeed_config train/ds_config_zero3.json

%cd ../

学習終了時のログは以下の用になりました

[INFO|trainer.py:2139] 2024-01-23 13:15:06,193 >> Loading best model from checkpoints-mistral-300M-FA2/checkpoint-100 (score: 8.360694885253906).
{'train_runtime': 4915.0289, 'train_samples_per_second': 83.336, 'train_steps_per_second': 0.02, 'train_loss': 9.417760696411133, 'epoch': 2.87}
100% 100/100 [1:21:55<00:00, 49.15s/it]
[INFO|trainer.py:2881] 2024-01-23 13:15:06,396 >> Saving model checkpoint to checkpoints-mistral-300M-FA2
[INFO|configuration_utils.py:461] 2024-01-23 13:15:06,397 >> Configuration saved in checkpoints-mistral-300M-FA2/config.json
[INFO|configuration_utils.py:564] 2024-01-23 13:15:06,398 >> Configuration saved in checkpoints-mistral-300M-FA2/generation_config.json
[INFO|modeling_utils.py:2193] 2024-01-23 13:15:07,386 >> Model weights saved in checkpoints-mistral-300M-FA2/pytorch_model.bin
[INFO|tokenization_utils_base.py:2428] 2024-01-23 13:15:07,387 >> tokenizer config file saved in checkpoints-mistral-300M-FA2/tokenizer_config.json
[INFO|tokenization_utils_base.py:2437] 2024-01-23 13:15:07,387 >> Special tokens file saved in checkpoints-mistral-300M-FA2/special_tokens_map.json
[INFO|tokenization_t5_fast.py:191] 2024-01-23 13:15:07,388 >> Copy vocab file to checkpoints-mistral-300M-FA2/spiece.model
***** train metrics *****
  epoch                    =       2.87
  train_loss               =     9.4178
  train_runtime            = 1:21:55.02
  train_samples            =     142865
  train_samples_per_second =     83.336
  train_steps_per_second   =       0.02
01/23/2024 13:15:07 - INFO - __main__ - *** Evaluate ***
[INFO|trainer.py:3158] 2024-01-23 13:15:07,401 >> ***** Running Evaluation *****
[INFO|trainer.py:3160] 2024-01-23 13:15:07,401 >>   Num examples = 8752
[INFO|trainer.py:3163] 2024-01-23 13:15:07,401 >>   Batch size = 16
100% 547/547 [00:46<00:00, 11.68it/s]
***** eval metrics *****
  epoch                   =       2.87
  eval_loss               =     8.3607
  eval_runtime            = 0:00:47.18
  eval_samples            =       8752
  eval_samples_per_second =    185.478
  eval_steps_per_second   =     11.592
  perplexity              =  4275.6648
[INFO|modelcard.py:452] 2024-01-23 13:15:54,937 >> Dropping the following result as it does not have all the necessary fields:
{'task': {'name': 'Causal Language Modeling', 'type': 'text-generation'}}
[2024-01-23 13:15:58,618] [INFO] [launch.py:347:main] Process 32974 exits successfully.
/content/japanese-mistral-300m-recipe
CPU times: user 25.4 s, sys: 4.65 s, total: 30.1 s
Wall time: 1h 23min 18s

推論

実際に作成したモデルの推論を行ってみます

推論

%%time
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
import torch

MODEL_NAME = "./pretrain/checkpoints-mistral-300M-FA2"
torch.set_float32_matmul_precision('high')

DEVICE = "cuda"
if torch.cuda.is_available():
    print("cuda")
    DEVICE = "cuda"
else:
    print("cpu")
    DEVICE = "cpu"
# DEVICE = "cpu"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME,use_fast=False)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
).to(DEVICE)

# streamer = TextStreamer(tokenizer)

prompt = "大規模言語モデルとは、"


inputs = tokenizer(prompt, add_special_tokens=False,return_tensors="pt").to(model.device)
with torch.no_grad():

    outputs = model.generate(
        inputs["input_ids"],
        max_new_tokens=100,
        do_sample=True,
        early_stopping=False,
        top_p=0.95,
        top_k=50,
        temperature=0.9,
        # streamer=streamer,
        no_repeat_ngram_size=2,
        num_beams=3
    )

print(outputs.tolist()[0])
outputs_txt = tokenizer.decode(outputs[0])
print(outputs_txt)

結果

cuda
The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.
[9114, 2342, 1073, 396, 260, 3528, 316, 42974, 316, 316, 10951, 316, 260, 260, 262, 260, 261, 261, 260, 10951, 42974, 906, 260, 316, 275, 260, 272, 260, 273, 261, 272, 261, 262, 261, 263, 260, 264, 260, 263, 261, 275, 261, 273, 260, 267, 260, 265, 261, 10951, 718, 260, 275, 262, 316, 261, 279, 260, 282, 261, 316, 268, 260, 279, 261, 265, 260, 1904, 260, 268, 261, 401, 260, 284, 260, 266, 260, 278, 260, 283, 261, 264, 262, 262, 42974, 42974, 262, 272, 262, 264, 261, 267, 261, 266, 261, 644, 260, 318, 260, 42974, 265, 263, 262, 1904, 261, 42974]
大規模言語モデルとは、子どもの  \  \ 、、の、。。、\ \火、 ・、)、(。)。の。は、が、は。・。(、で、を。\川、・の 。年、月。 と、年。を、 S、と。ス、日、に、1、2。がのの \ \の)のが。で。に。ノ、:、 \をはの S。 \
CPU times: user 4.3 s, sys: 324 ms, total: 4.62 s
Wall time: 4.01 s

モデルをhuggingfaceにアップグレードする

作ったモデルをhuggingfaceにアップロードしていきます

ライブラリのインストールとログイン

まずはラブライブを入れます

!pip install -U "huggingface_hub[cli]"

その後にログインを行います

!huggingface-cli login

ここで、tokenの入力が求められるので huggingfaceのsettingからwrite権限のあるtokenを入力します

    To login, `huggingface_hub` requires a token generated from https://huggingface.co/settings/tokens .
Token: 

モデルのアップロード

モデルをアップロードします

!huggingface-cli upload {your repository id} pretrain/checkpoints-mistral-300M-FA2 .

以下のようにアップロードされます

huggingface.co