ELYZAのモデルとトークナイザーをコーディングしながら理解する

## はじめに
弊社(ARISE analytics)の親会社であるKDDIがELYZAと提携しました。

そんな背景から、ELYZAについてしっかり理解しておこうということで、記事を書きました。

本記事は、ELYZAの方々が書いている記事の内容に加えて、プログラムを書きながら理解を深められるような作りにしています。

:::note 
モデル、トークナイザーどELYZA内部の大まかな理解に努めます。
:::

:::note warn
LLMの基本について、ELYZAの性能に関しての議論しません。
:::

今回用いたプログラムはGitHubに公開しています。

https://github.com/sergicalsix/understand_LLM/tree/main

## 前提
### ELYZAの公開モデルについて

ELYZAを理解するにあたり、公開されているモデルを調べます。

https://huggingface.co/elyza

公開されているモデルは大まかに2つに分類できます。

- Llama-2ベースのチャットモデル(7b,13b)
- CodeLlamaベースのチャット兼コード生成特化モデル(7b)

今回は汎用性の観点から、Llama-2ベースのチャットモデル(7b,13b)について見ていきます。

### 補足: ELYZAの非公開モデルについて
ELYZAの70Bのモデルは重み非公開ですが、実際に使うことができます。

https://note.com/elyza/n/n0ea755ca3e7b

https://elyza.ai/lp/elyza-llm-for-jp

上記noteにも記載されていますが、Speculative Decodingが導入されていたりと高速化で工夫されています。

## モデル
### 公開モデル概観
Llama-2ベースのモデルは7b,13bそれぞれ4つずつあります。(後述しますが、以下で示しているHuggingFaceのモデルサイズの表記は誤っているかもしれません。)



| Model Name                                 | Vocab Size | #Params |
|--------------------------------------------|------------|---------|
| elyza/ELYZA-japanese-Llama-2-7b            | 32000      | 6.27B   |
| elyza/ELYZA-japanese-Llama-2-7b-instruct   | 32000      | 6.27B   |
| elyza/ELYZA-japanese-Llama-2-7b-fast       | 45043      | 6.37B   |
| elyza/ELYZA-japanese-Llama-2-7b-fast-instruct | 45043    | 6.37B   |
| elyza/ELYZA-japanese-Llama-2-13b           | 32000      | 13.02B  |
| elyza/ELYZA-japanese-Llama-2-13b-instruct  | 32000      | 13.02B  |
| elyza/ELYZA-japanese-Llama-2-13b-fast      | 44581      | 13.14B  |
| elyza/ELYZA-japanese-Llama-2-13b-fast-instruct | 44581    | 13.14B  |


fastと名のついているモデルは、日本語の語彙を追加しているモデルであり、instructと書いているのはインストラクションチューニングされているモデルです。

語彙追加モデルについては後述します。


### ベースモデル: Llama-2
ELYZAのベースモデルは先述したようにLlama-2が用いられています。

https://huggingface.co/meta-llama

Llama-2についてご存知の方も多そうですが、少し解説を添えます。

- 学習量

事前学習に使われたトークンは2兆個であり、膨大です。
日本語のトークンは20億トークンであり、全体の0.1%です。

- コンテキストウィンドウ

当時(2023年7月)のモデルにしては大きめの4096トークンです。（ちなみにGPT-4 Turboは128k、Claude 2.1は20万ほどです）
4096トークンはどれくらいの数かというと、英語だと(難しい単語がなければ)約1500-2500単語ほど、日本語だと(漢字の量や難易度によって左右されますが)約1500-3000文字くらい入力できます。

:::note 
「漢字の量や難易度によって左右されるとはどういうこと？」という疑問については、後述するトークナイザーパートを見ればイメージがつくと思います。
:::


- 親切さと安全性

Llama-2のチャットモデルは10万回以上のファインチューニングと、100万以上の人間のPreferenceが介在した強化学習(RLHF)が実施されており、親切さ(helpfulness)と安全性(safety)を重視して学習されています。

- Ghost Attension(GAtt)の導入

既存の言語モデルの指示を忘却する問題に対処するため、Ghost Attension(GAtt)が導入されています。
詳細知りたい方は原著か解説記事をご覧ください。

https://arxiv.org/pdf/2307.09288.pdf

https://towardsdatascience.com/understanding-ghost-attention-in-llama-2-dba624901586


具体のアーキテクチャについてはコーディングパートで確認します。


---

7bのELYZAのベースモデルは、Llama-2のチャットモデルである`meta-llama/Llama-2-7b-chat-hf`です。その理由は、Llama-2の親切さ(helpfulness)と安全性(safety)を引き継ぐためだと述べられています。

上記の動機に従って考えると、おそらく13bのELYZAのベースモデルは`meta-llama/Llama-2-7b-chat-hf`だと予想されます。

### モデルの学習
#### データセット
Llama-2の20億トークンの日本語データに加えて、180億トークンを追加しています。
目的はLlama-2の日本語処理能力を上げるためです。

学習データはOSCARやWikipediaといった王道どころと、そのほかのクロールデータを用いているそうです。

#### ハイパーパラメータ
オプティマイザはLlama-2同様、AdamWを用いています。
またスケジューラーは今回使っておらず、学習率は固定して学習しているようです。

#### 学習方法
ファインチューニング、(末尾にinstuctがついているものに限り)イントラクションチューニングがなされています。
パラメータの変化が(ほぼ)ないことから、アダプター系の学習はされていないと予想されます。
モデルの一部の重みが固定化してチューニングをしているのか、フルパラメータのファインチューニングをされているのかは公開情報のみだと不明ですが、おそらく後者だと予想できます。（狭義のファインチューニングは後者なのですが、昨今言葉の意味が変わってきていて、どちらで実施されているかは不明です。）

以上はコーディングパートで調査していきます。
(調査結果を先出しすると、ELYZAはLlama-2のフルパラメータチューニングがされています。)


#### その他
Continual Learningは公開モデルでは未反映のようです。


### モデルの内部確認
次に具体的にLlama-2とELYZAの内部を見ていきます。

目的を3つ定めます。

1. ELYZAのベースモデルの`meta-llama/Llama-2-7b-chat-hf`について理解する
2. ELYZAの学習方法を明確化する
2.  `elyza/ELYZA-japanese-Llama-2-7b`と`elyza/ELYZA-japanese-Llama-2-7b-fast`のアーキテクチャの違いを調べる

今回はベースモデルの`meta-llama/Llama-2-7b-chat-hf`と`elyza/ELYZA-japanese-Llama-2-7b`と`elyza/ELYZA-japanese-Llama-2-7b-fast`の3つを見ていきます。

:::note warn
Llama-2のモデルを使うためには、meta社にリクエストを送る必要があります。通常1-2日ほどでLlama-2を動かせるようになります。(筆者は3日ほどかかりました)
:::

#### 事前準備
まずライブラリをインストールし、HuggingFaceにログインします。

```bash
!pip install -U transformers
!huggingface-cli login
```
次にモデルを定義します。

```python 
import torch 
from transformers import AutoModelForCausalLM

llama_model_name = "meta-llama/Llama-2-7b-chat-hf"
elyza_model_name = "elyza/ELYZA-japanese-Llama-2-7b"
elyza_fast_model_name = "elyza/ELYZA-japanese-Llama-2-7b-fast"

llama_model = AutoModelForCausalLM.from_pretrained(llama_model_name, torch_dtype=torch.float16)
elyza_model = AutoModelForCausalLM.from_pretrained(elyza_model_name, torch_dtype=torch.float16)
elyza_fast_model = AutoModelForCausalLM.from_pretrained(elyza_fast_model_name, torch_dtype=torch.float16)
```

:::note warn 
上記でメモリをオーバーする場合はint8を採用したり、モデル定義の個数を3個から2個に変えるなどして対応してください。
:::


#### 1. ELYZAのベースモデルの`meta-llama/Llama-2-7b-chat-hf`について理解する
先ほど定義したモデルの中身を確認します。

```python 
LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaSdpaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)
```

上記より、Llama-2の構造は、エンべディング層->(アテンション層+MLP層+正規化層)×32 ->最終層(全結合層)であることがわかります。

次にモデルの基本情報を取得します。

```python 
def print_model_info(model):
    total_params = 0
    trainable_params = 0
    non_trainable_params = 0
    module_count = 0

    for name, module in model.named_modules():
        module_params = sum(p.numel() for p in module.parameters(recurse=False))
        module_count += 1
        total_params += module_params

    print(model.name_or_path)
    print(f"Total layers: {module_count}")
    print(f"Total parameters: {total_params}")

print_model_info(llama_model)
```

```
meta-llama/Llama-2-7b-chat-hf
Total layers: 454
Total parameters: 6738415616
```

上記より層の数は454、パラメータ数は6.7Bであることがわかります。


もっと詳しく理解したい場合はNetronなどを用いると良いです。

https://github.com/lutzroeder/netron

#### 2.ELYZAの学習方法を明確化する
学習方法を具体化するために`meta-llama/Llama-2-7b-chat-hf`と`elyza/ELYZA-japanese-Llama-2-7b`の重みを比較します。

上記で記述した`print_model_info`関数を用いてELYZAのモデルの基本情報を確認します。

```python 
print_model_info(elyza_model)
```

```
elyza/ELYZA-japanese-Llama-2-7b
Total layers: 454
Total parameters: 6738415616
```


上記より層の数及びパラメータ数が一致しているので、ELYZAの学習方法にはアダプター形式の学習はしていないことがわかります。

(HuggingFaceの中の6.27Bというモデルサイズは誤記かもしれません。)


次にELYZAがフルパラメータチューニングされているかどうかを確かめます。

具体的にLlama-2とELYZAの各層の重みをそれぞれ比較して、重みが同一かどうかを確かめます。

```python 
def compare_model_weights(model_a, model_b):
    state_dict_a = model_a.state_dict()
    state_dict_b = model_b.state_dict()

    # 同じキーを持つかチェック
    if state_dict_a.keys() != state_dict_b.keys():
        print("Models have different layers.")
        return

    # 各層のパラメータを比較
    for key in state_dict_a:
        if torch.equal(state_dict_a[key], state_dict_b[key]):
            print(f"Layer {key}: Weights are identical.")
            return
    print(f'All Weights are different')
    return 

compare_model_weights(llama_model, elyza_model)
```

```
All Weights are different.
```

上記より全ての層の重みが異なることがわかりました。
よってELYZAはフルパラメータファインチューニングしていることが確認できました。


#### 3.`elyza/ELYZA-japanese-Llama-2-7b`と`elyza/ELYZA-japanese-Llama-2-7b-fast`のアーキテクチャの違いを調べる

最後に`elyza/ELYZA-japanese-Llama-2-7b`と`elyza/ELYZA-japanese-Llama-2-7b-fast`のアーキテクチャの違いを調べていきます。

とはいってもトークン数の種類数が変わっているので、エンべディング層と最終層のみが異なることは想像に難くはないのですが、一応確かめてみます。

```python 
# elyza/ELYZA-japanese-Llama-2-7b
LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096, padding_idx=0)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaSdpaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

# elyza/ELYZA-japanese-Llama-2-7b-fast
LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(45043, 4096, padding_idx=0)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaSdpaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=45043, bias=False)
)
```

上記より、`elyza/ELYZA-japanese-Llama-2-7b-fast`の方が語彙数の増加分(32000->45043)だけ、エンべディング層と最終層のみサイズが増加していることがわかりました。

余談ですが、語彙数を追加した状態でファインチューニングしているかを確かめるために、中間層の重みが同じかどうかを確認してみます。

```python 
compare_model_weights(elyza_fast_model, elyza_model)
```

```
All Weights are different.
```

上記より語彙数を追加した状態で、重みを全てファインチューイングしていること(語彙数追加前のモデルのエンべディング層と最終層を付け替えただけではないこと)がわかりました。


モデル面での確認は以上としたいと思います。

## トークナイザー
トークナイザーはSentencePiece BPEのアルゴリズムを採用しています。

トークナイザーについては以下の記事がわかりやすいです。

https://huggingface.co/docs/transformers/ja/tokenizer_summary

ELYZAのトークナイザーの特筆事項として、先述の通り、ELYZAには日本語の語彙が追加されています。

具体的には出現頻度の多い日本語を文字・バイト単位ではなく、単語単位でトークン化しています。
これにより日本語の入出力に関しては、必要なトークン数が減るので処理が高速化します。

上記について本当に必要なトークン数が減るのかどうかを、実際にプログラムを書いて確認していきます。

以下のようにライブラリをインストールし、HuggingFaceにログインします。
```bash
!pip install -U transformers
!huggingface-cli login
```

次にLlama-2のトークナイザーと語彙が追加されたELYZAのトークナイザーを呼び出します。

```python 
from transformers import AutoTokenizer
tokenizer_llama2 = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
tokenizer_elyza_fast = AutoTokenizer.from_pretrained("elyza/ELYZA-japanese-Llama-2-7b-fast-instruct")

tokenizer_llama2.vocab_size, tokenizer_elyza_fast.vocab_size #語彙数の確認(32000, 45043)
```

以下のプログラムで、文のトークン数がどのように変化しているか、一つ一つのトークンがどの単語に対応しているかを確認します。

```python 
def print_token_count(text, tokenizer):
  input_ids = tokenizer(text)['input_ids']
  print(f"### {tokenizer.name_or_path}\nトークンサイズ: {len(input_ids)}")

def print_decoded_tokens(text, tokenizer):
  input_ids = tokenizer(text)['input_ids']
  print(f"### {tokenizer.name_or_path}\n{[tokenizer.decode(id) for id in input_ids]}")


text_ja = "東京都は日本の首都です。"

print_token_count(text_ja, tokenizer_llama2)
print_token_count(text_ja, tokenizer_elyza_fast)

print_decoded_tokens(text_ja, tokenizer_llama2)
print_decoded_tokens(text_ja, tokenizer_elyza_fast)
````

```出力
### meta-llama/Llama-2-7b-chat-hf
トークンサイズ:14
### elyza/ELYZA-japanese-Llama-2-7b-fast-instruct
トークンサイズ:9
### meta-llama/Llama-2-7b-chat-hf
['
', '', '東', '京', '都', 'は', '日', '本', 'の', '首', '都', 'で', 'す', '。']
### elyza/ELYZA-japanese-Llama-2-7b-fast-instruct
['
', '', '東京都', 'は', '日本の', '首', '都', 'です', '。']
````

以上よりELYZAのトークナイザーによって実際のトークンサイズが減少していることがわかります。
前述したようにモデル間のアーキテクチャの違いはほぼないので、トークナイザの改善によって処理が高速化します。

トークン一つ一つを見てみると、Llama-2のトークナイザーは東京都を'東', '京', '都'とトークン化しているのに対して、ELYZAの語彙追加版のトークナイザでは'東京都'を一つのトークンとして扱っていることがわかります。


次に難しい日本語の文でも上記の結果が保存されるかを実験で試してみます。


```python 
text_ja_difficult = "蒼穹は昏冥の淵に沈む。"

print_token_count(text_ja_difficult, tokenizer_llama2)
print_token_count(text_ja_difficult, tokenizer_elyza_fast)

print_decoded_tokens(text_ja_difficult, tokenizer_llama2)
print_decoded_tokens(text_ja_difficult, tokenizer_elyza_fast)
```


```出力
### meta-llama/Llama-2-7b-chat-hf
トークンサイズ:25
### elyza/ELYZA-japanese-Llama-2-7b-fast-instruct
トークンサイズ:13
### meta-llama/Llama-2-7b-chat-hf
['
', '', '�', '�', '�', '�', '�', '�', 'は', '�', '�', '�', '�', '�', '�', 'の', '�', '�', '�', 'に', '�', '�', '�', 'む', '。']
### elyza/ELYZA-japanese-Llama-2-7b-fast-instruct
['
', '', '蒼', '穹', 'は', '昏', '冥', 'の', '淵', 'に', '沈', 'む', '。']
```

以上より難しい文章でもELYZAのトークナイザーによって実際のトークンサイズが減少していることがわかります。
ELYZAでは難しい日本語についても1文字1トークンで変換していますが、Llamaの方では1文字数トークンの変換が必要になっていることがわかります。

:::note
モデルのコンテキストウィンドウのパートでも説明しましたが、Llamaのトークナイザは漢字に対して1文字1トークンで対応させることが難しいケースが多いです。

まだ詳細に調査していませんが、Llamaの場合は事前トークンにない漢字については1文字3トークンで処理しているようです。

よってある程度難しい漢字が入れば入るほど、理論上ELYZAの方が処理は速くなります。
:::

結論、簡単な文章でも難しい文章でも日本語の文章に関してはトークン数を削減->モデル高速化できることがわかりました。

---

トークナイザーの詳細は以下のELYZAの記事をご覧ください。

 https://zenn.dev/elyza/articles/2fd451c944649d

7b->13bにかけてELYZA内部でさらにトークナイザーが改良されています。詳しくはELYZAの13bの記事をご覧ください。

 https://note.com/elyza/n/n5d42686b60b7




## おわりに
半分以上ELYZAさんが書いている記事の内容をまとめただけな気もするのですが、参考になれば幸いです。

## 2024/04/09:補足
ELYZAの推論が高速な理由として、トークナイザーについて説明しましたが、加えて70BのモデルではSpeculative Decodingが用いられているようです。

https://note.com/elyza/n/n0ea755ca3e7b

Speculative Decodingとは「メインモデルの推論」フローを「軽量なアシスタントモデルの推論&メインモデルの検証&修正」に変えることで高速化を図る手法です。

https://zenn.dev/umeko/articles/fe961fda3148d1#speculative-decoding

https://huggingface.co/blog/whisper-speculative-decoding

## 参考記事

最後にもう一度参考記事をまとめます。

 - ELYZAについて

 https://zenn.dev/elyza/articles/2fd451c944649d

 https://note.com/elyza/n/n5d42686b60b7

- Llama-2

https://arxiv.org/pdf/2307.09288.pdf

 - Ghost Attensionについて

https://towardsdatascience.com/understanding-ghost-attention-in-llama-2-dba624901586

- トークナイザー

https://huggingface.co/docs/transformers/ja/tokenizer_summary