!python scripts/prepare_ja.py \
--tokenizer_path "tokenizerのdir/tokenizer.json" \
--destination_path '出力先のdir' \
--dataset_ids ["huggingfaceのrepository", "huggingfaceのrepository"]

出力先のdirにbinファイルが生成されます。

今回、学習に使ったデータは日本語と英語を使った以下になります。

total tokens: 8.64B

wikipedia_ja:    844.65M  
wikipedia_en:    3.80B  
open-text-books: 60.17M  
oscar:           3.85B  
aozorabunko:     92.97M

training

学習は以下のscriptを参考にします。

変更点は以下

hyper paramer
学習するmodel size

変更点に加え、train用のコードやllama2の実装について関連する部分をみていきます。

fabric

lit-gpt/lit-llamaではtorch fabricを使って学習をまとめています。

複数GPUのある環境などでは便利そう。今回はgoogle colab上で動作させるのであまりメリットはなさそう。
fabricも色々便利なのでこちらも見てみると良いです。

今回参考にしたコードでも利用しているclip_gradientsが簡単に使えてて良さそう

model sizeの変更

layer数などのmodel configが以下にあります

llama2以外にもphiなどに対応しています。

このconfigを書き換えて、400Mになるくらいのサイズに調整します。
使ったのが以下。これで 417.12Mのサイズになります。

name="Llama-2-400M",
vocab_size=35000,
padding_multiple=64,
block_size=2048,
n_layer=12,
n_head=12,
n_embd=840,
n_query_groups=6,
rotary_percentage=1.0,
parallel_residual=False,
bias=False,
_norm_class="RMSNorm",
_mlp_class="LLaMAMLP",
intermediate_size=11008,
norm_eps=1.0e-8,

デフォルトで7Bが呼び出されます。作成したconfigがロードされるように書き換えます。

query group

configのコードにも書いていますが、query groupにも対応してます。

query groupは、同様のhead数でquery groupを使う場合と使わない場合で性能差は生まれなかったというようなものです。つまり、key/valueをある程度省略しても性能差がでないということになります。その分parameter数を削減でき、他の部分、例えばMLP層や、head数を増やすことができるといった具合です。

詳細は論文を確認してください。

MLP

GPT用のMLPとLLama用のMLPが用意されています。
swigluを使うかの違いです。

shared_attention_norm

phiではattention層のnormをshareしています。
configでon/offを切り替えられます。

tensor board

fabricでは、loggerとしてtensorboardが指定できるので指定しておきます。
以下のような感じ

logger = TensorBoardLogger(log_dir, name="model")
fabric = L.Fabric(accelerator="cuda", devices=devices, precision=precision, loggers=logger)

学習途中で、以下のように記録できます。stepの引数を入れないとうまく動かなかったので注意です。

 fabric.log_dict(
   {"iter": iter_num, "val_loss": val_loss, "step": step_count, "lr": lr}, step=iter_num
)

学習用のhyper paramterの書き換え

学習用のhyper parameterは以下です。学習の結果を見ながらこの辺りを変更していくことになります。

classにまとめて、使ったconfigを保存できるようにしておくと良いです。pythonでは、以下のようにmember変数をすべて保存できます。

class Config:
    def __init__(...):
        ...
    def save(self, output_dir):
        member_vars = {k: v for k, v in self.__dict__.items() if not callable(v)}
        output_file = f'{output_dir}/config.json'
        with open(output_file, 'w') as f:
            json.dump(member_vars, f, ensure_ascii=False, indent=4)

run

今回はgoogle colabのV100で約１日くらい学習を回しました。
val lossは4.3020まで下がりました。

generate

以下のコードを参考に学習したモデルで文章を生成します。

!python generate.py \
--checkpoint_path "学習済weightを保存したdir/iter-253000-ckpt.pth" \
--tokenizer_path "tokenizerのdir/tokenizer.json" \
--top_k 40 \
--temperature 0.4 \
--model_name "Llama-2-400M" \
--prompt "富士山の山頂では、"

出力

Loading model ...Llama-2-400M
Time to load model: 2.02 seconds.
Seed set to 1234
富士山の山頂では、高尾山噴火場は、標高30メートル前後の高さの地点にある。山頂から南にかけては、1時間あたり約1時間半ほどで、山頂には約1時間半ほどで、約1時間半
Time for inference 1: 0.85 sec total, 58.80 tokens/sec
Memory used: 1.90 GB

意味はわからないが、なんとなく文法はあってそう。
generateのparameterをもう少し調整すればちゃんと文章になりそう。

huggingfaceへのupload

copy_weights_llama関数を使って、layer名をhuggingface用に書き換えます。
このコードを使うとそのままbinファイルが生成されます。

次に、transformers.LlamaConfigを今回使った400M用のconfigで書き換えます。
対応は以下のような感じ

from transformers import LlamaConfig
from lit_llama.config_llama2 import Llama2Config

def convert(model_size):
    t_config = LlamaConfig() ## transformers.LlamaConfig
    config = Llama2Config.from_name(model_size) ## 今回の学習で使った400M用のconfigをロード

    t_config.hidden_size = config.n_embd
    t_config.max_position_embeddings = config.block_size
    t_config.intermediate_size = config.intermediate_size

    t_config.num_attention_heads = config.n_head
    t_config.num_hidden_layers = config.n_layer
    t_config.num_key_value_heads = config.n_query_groups

    t_config.rms_norm_eps = config.norm_eps

    ## tokenizer config
    t_config.bos_token_id = 1
    t_config.eos_token_id = 2
    t_config.vocab_size = 35000 + 8

    return t_config

上記のconvert関数を使ってtransformers.LlamaConfigからCausalLMのLLamaをロードします。
layer名変更後のweightをtorch.loadを使いweightをロード、load_state_dictでmodelに適応します。

model.push_to_hub でモデルのuploadを行います。

t_config = convert_config(model_size)
model = AutoModelForCausalLM.from_config(t_config)

pytorch_weights = torch.load(weight_path)
model.load_state_dict(pytorch_weights)

model.push_to_hub(repo_id)

事前に !huggingface-cli login を行っておいてください。write権限のあるtokenが必要です。

huggingface/transformersを使ったgenerate

huggingfaceにuploadしたのでtransformersを使いgenerateしてみます。
以下のような感じ

from transformers import AutoTokenizer
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("if001/sentencepiece_ja", trust_remote_code=True, force_download=True)
model = AutoModelForCausalLM.from_pretrained("if001/llama2_ja_small")


import torch
from transformers import GenerationConfig

prompt="富士山の山頂では、

inputs = tokenizer(prompt, return_tensors="pt")
input_ids = inputs["input_ids"]

generation_config = GenerationConfig(
    temperature=0.8,
    top_p=0.95,
    top_k=50,
    num_beams=1,
    do_sample=True,
    repetition_penalty=1.2,
    pad_token_id= tokenizer.pad_token_id,
    # pad_token_id=tokenizer.unk_token_id,
    eos_token_id=tokenizer.eos_token_id
)
with torch.no_grad():
  generation_output = model.generate(
            input_ids=input_ids,
            generation_config=generation_config,
            return_dict_in_generate=True,
            output_scores=True,
            max_new_tokens=64,
        )
  s = generation_output.sequences[0]
  print(s)
output = tokenizer.decode(s)
print(output)

出力

富士山の山頂では、その名の通りの地形が残っている。この地形の地形をはさんで、これはまた別の場所に位置している。

所感

pre_trainingからhuggingfaceへのuploadまでを行ってみました。

データセットのサイズと1Bの半分くらいということで、400Mのモデルにしましたが、
文法はなんとなくあってそうだが、まだまだ喋れる感じではなさそうです。
parameter調整をもう少ししてみたいのと、他のアーキテクチャも試してみたいところ。
ただ、小さいサイズであればgoogle colabで学習できるんだな、ということで色々遊べそうです。

llama2も実装を見てみるのもおもしろいので、このあたりのコードを眺めるものおすすめです。
llama2とllamaとの違いやphiとの違いがわかって面白いです。

Discussion

fizunimo 2ヶ月前

素晴らしい記事をありがとうございます。
LLMのpre-trainingについて、実装方法が分からず悩んでいたのでとても参考になりました。
一つ質問なのですが、こちらはhttps://huggingface.co/if001/llama2_ja_smallの継続事前学習ではなく、configを使用して新しいものを1から作成しているという認識であっているでしょうか。

if001 2ヶ月前

ありがとうございます！

configを使用して新しいものを1から作成しているという認識であっているでしょうか。

はい、こちらの認識であってます

fizunimo 2ヶ月前

返信が遅くなり申し訳ございません。
ありがとうございます。こちらの記事のおかげで理解が深まりました！