はじめての自然言語処理

今回は Elyza さんの日本語 Llama 2 指示応答モデルをファインチューニングし、vLLM にデプロイして高速に推論してみます。70 億パラメータモデルならギリギリ Tesla T4 x 1 の構成でも float16 で動かせるかと思ったのですが、どうだったでしょうか。vLLM には OpenAI 互換の API インタフェースも備えているので、ついでに LangChain からも接続してみたり。

1. はじめに

今回は Elyza さんが公開されている大規模指示応答言語モデルである、ELYZA-japanese-Llama-2-7b-fast-instruct¹ をファインチューニングして vLLM で推論してみます。

そんな訳で今回あまり書くことがなく、動かし方だけサラっと書いて「動きましたー。では良いお年を～。」で締めることにします。 しかし、時代感覚無視の隔月連載でネタを書き溜めてたりするもんだから、最近の激動っぷりをみているとタイムラグ感が否めないですね。。。ご容赦ください。

まずは簡単に Elyza さんの Llama 2 についてご紹介しましょう。

2. ELYZA-japanese-Llama-2-7b-fast-instruct

さて、今回触ってみる ELYZA-japanese-Llama-2-7b-fast-instruct は Meta 社が今年の7月に公開した Llama 2² の SFTやRLHFによる指示応答学習が適用されたモデルに対し、Elyza 社が日本語の追加学習を適用し、さらに指示応答の追加学習をしたものになります。

Meta の Llama 2 は公開当時にその性能の高さが評判になったのですが、事前学習データ中の日本語の割合は 0.1 % 程度であり、 日本語での利用はかなり厳しい（というかムリ）ものでした。

オリジナルの Llama 2 に対して Elyza さんが行った追加事前学習については、こちらの記事³にまとめられています。

特にモデル名に “fast” が含まれるバージョンは、高速化のため日本語向けの語彙を 15,000 件追加した上で、 OSCAR や Wikipedia を中心に 160 億トークンの日本語による追加学習を行ったそうです。

個人的には指示応答学習まで行った英語ベースのモデルにこのような日本語のみでの追加学習を適用すると破滅的忘却により、 最初に持っていたモデルの能力が棄損されると思っていたので、「どんな工夫をしたのだろう？」と思ったのですが、 「意外と大丈夫だった」ということらしいです⁴。

OpenAI の ChatGPT でもそうなのですが「同じ内容でも英語に比べて日本語だとトークン数がかさむ」というのは、 非常に悩ましい問題なので、ELYZA-japanese-Llama-2-7b-fast-instruct のような日本語語彙を強化したモデルの公開は非常にうれしいですね⁵。

では実際に環境を作ってファインチューニングしてみましょう。

3. ファインチューニング

無料の Colab ではメモリがちょっと足らないので、今回は GCP に VM を立てて作業します。 Tesla T4 x 1 と 2 vCPU、メモリ 16 GB で AWS の g4dn.xlarge 相当の構成ですね。

• NVIDIA Driver と NVIDIA docker がインストール済みの前提で話を進めます。

まずは以下のようにしてコンテナを立てください。

$ docker run --name hiroba -v `echo $HOME`:/work --gpus all -it pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
# cd /work

実際はコンテナの中に入って shell プロンプトで作業したのですが、jupyter で動かした風に記事を書いていきますね。 最初に本記事でつかうコマンドをいくつかインストールしておきます。

!apt-get update
!apt-get install -y git

学習に必要なパッケージをインストールします。

!pip install \
      accelerate==0.21.0 \
      peft==0.4.0 \
      bitsandbytes==0.40.2 \
      transformers==4.34.1 \
      trl==0.4.7 \
      scipy==1.11.2

インストールしたパッケージについて簡単に補足しておくと以下のような感じです。

• transformers :
  言わずと知れた Transformer ベースの深層学習モデルのライブラリ。もとは自然言語が中心でしたが、現在では画像や音声系のモデルにも対応しています。
• accelerate :
  PyTorch のモデルを様々な種類(CPU, GPU, GPU)や構成(single/multi)のデバイスで動かす為の抽象化を担うライブラリです。
• peft :
  Parameter-Efficient Fine-Tuning の名前のとおり最近の大規模なモデルを省リソースで効率よく学習する為のライブラリです。
• bitsandbytes :
  学習時にモデルのパラメータを 4bit に量子化するためのライブラリです。
• trl :
  Transformer Reinforcement Learning で LLM に RLHF を適用するためのライブラリです。以前紹介した DeepSpeed-Chat と同じような立ち位置かと思います。
• scipy :
  科学技術計算のライブラリですけど、なんで入れたか忘れてしまいました。スイマセン。。。

学習スクリプトの準備

ファインチューニング用のスクリプトは trl に含まれる sft_trainer.py を改造して使うので取得します。

!git clone https://github.com/huggingface/trl/
!cd trl && git checkout 25d6700c5e6db96a

sft_trainer.py に適用した修正内容は以下のとおりです（ちょっと長いですが）。

%%writefile patch
--- trl/examples/scripts/sft_trainer.py 2023-11-20 23:56:45.213444569 +0000
+++ modified/sft_trainer.py 2023-11-20 23:57:21.449061360 +0000
@@ -1,4 +1,11 @@
+#!/usr/bin/env python
 # coding=utf-8
+
+# This file was copied from :
+#   https://github.com/huggingface/trl/blob/25d6700c5e6db96a722942c9c91594240792346d/examples/scripts/sft_trainer.py
+# And add merge_and_push argument and related code from :
+#   https://gist.github.com/younesbelkada/9f7f75c94bdc1981c8ca5cc937d4a4da
+
 # Copyright 2023 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,15 +20,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Optional, Union

 import torch
 from datasets import load_dataset
 from peft import LoraConfig
 from tqdm import tqdm
-from transformers import AutoModelForCausalLM, BitsAndBytesConfig, HfArgumentParser, TrainingArguments
-
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, HfArgumentParser, TrainingArguments
+from transformers.tokenization_utils_base import BatchEncoding
 from trl import SFTTrainer
+import numpy as np
+import json
+import os
+import sys


 tqdm.pandas()
@@ -38,7 +49,19 @@
     dataset_name: Optional[str] = field(
         default="timdettmers/openassistant-guanaco", metadata={"help": "the dataset name"}
     )
+    json_data_files: Optional[list[str]] = field(default_factory=list, metadata={"help": "the data_files arguments in JSONL format."})
+    add_eos_token: Optional[bool] = field(default=True, metadata={"help": "Wether to add eos token."})
+    alternative_pad_token: Optional[str] = field(
+        default=None, metadata={"help": "the alternative pad token for tokenizer."
+                                         "should be one of ['bos', 'eos', 'unk', 'sep', 'cls', 'mask']."}
+    )
     dataset_text_field: Optional[str] = field(default="text", metadata={"help": "the text field of the dataset"})
+    default_system_prompt: Optional[str] = field(
+        default="あなたは誠実で優秀な日本人のアシスタントです。",
+        metadata={"help": "the default system prompt for llama instruction."}
+    )
+    llama_instruction: Optional[bool] = field(default=False, metadata={"help": "Wether to use llama instruction format."})
+    youri_instruction: Optional[bool] = field(default=False, metadata={"help": "Wether to use youri instruction format."})
     log_with: Optional[str] = field(default=None, metadata={"help": "use 'wandb' to log with wandb"})
     learning_rate: Optional[float] = field(default=1.41e-5, metadata={"help": "the learning rate"})
     batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size"})
@@ -48,36 +71,67 @@
     )
     load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 8 bits precision"})
     load_in_4bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 4 bits precision"})
+    torch_dtype: Optional[str] = field(default="float16", metadata={"help": "torch dtype"})
     use_peft: Optional[bool] = field(default=False, metadata={"help": "Wether to use PEFT or not to train adapters"})
     trust_remote_code: Optional[bool] = field(default=True, metadata={"help": "Enable `trust_remote_code`"})
     output_dir: Optional[str] = field(default="output", metadata={"help": "the output directory"})
     peft_lora_r: Optional[int] = field(default=64, metadata={"help": "the r parameter of the LoRA adapters"})
     peft_lora_alpha: Optional[int] = field(default=16, metadata={"help": "the alpha parameter of the LoRA adapters"})
     logging_steps: Optional[int] = field(default=1, metadata={"help": "the number of logging steps"})
-    use_auth_token: Optional[bool] = field(default=True, metadata={"help": "Use HF auth token to access the model"})
+    use_auth_token: Optional[bool] = field(default=False, metadata={"help": "Use HF auth token to access the model"})
     num_train_epochs: Optional[int] = field(default=3, metadata={"help": "the number of training epochs"})
     max_steps: Optional[int] = field(default=-1, metadata={"help": "the number of training steps"})
     save_steps: Optional[int] = field(
         default=100, metadata={"help": "Number of updates steps before two checkpoint saves"}
     )
     save_total_limit: Optional[int] = field(default=10, metadata={"help": "Limits total number of checkpoints."})
+    merge_and_push: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Merge and push weights after training"},
+    )
     push_to_hub: Optional[bool] = field(default=False, metadata={"help": "Push the model to HF Hub"})
     hub_model_id: Optional[str] = field(default=None, metadata={"help": "The name of the model on HF Hub"})
+    merge_only: Optional[bool] = field(default=False, metadata={"help": "Merge lola weight only."})
+    group_by_length: Optional[bool] = field(default=False, metadata={"help": "Whether or not to group together samples of roughly the same length."})


 parser = HfArgumentParser(ScriptArguments)
 script_args = parser.parse_args_into_dataclasses()[0]

+if script_args.merge_only:
+    from peft import AutoPeftModelForCausalLM
+    output_dir = os.path.join(script_args.output_dir, "final_checkpoints")
+    print("Loading the lola wegiht ...", flush=True)
+    model = AutoPeftModelForCausalLM.from_pretrained(
+                output_dir,
+                device_map="auto",
+                torch_dtype=getattr(torch, script_args.torch_dtype),
+    )
+
+    print("Merging the lola wegiht ...", flush=True)
+    model = model.merge_and_unload()
+
+    print("Saving the merged model ...", flush=True)
+    output_merged_dir = os.path.join(script_args.output_dir, "final_merged_checkpoint")
+    model.save_pretrained(output_merged_dir, safe_serialization=False)
+    tokenizer = AutoTokenizer.from_pretrained(output_dir)
+    tokenizer.save_pretrained(output_merged_dir)
+    sys.exit(0)
+
+
 # Step 1: Load the model
+print("Loading the model...", flush=True)
 if script_args.load_in_8bit and script_args.load_in_4bit:
     raise ValueError("You can't load the model in 8 bits and 4 bits at the same time")
 elif script_args.load_in_8bit or script_args.load_in_4bit:
     quantization_config = BitsAndBytesConfig(
-        load_in_8bit=script_args.load_in_8bit, load_in_4bit=script_args.load_in_4bit
+        load_in_8bit=script_args.load_in_8bit,
+        load_in_4bit=script_args.load_in_4bit,
+        bnb_4bit_compute_dtype= getattr(torch, script_args.torch_dtype),
     )
     # This means: fit the entire model on the GPU:0
     device_map = {"": 0}
-    torch_dtype = torch.bfloat16
+    torch_dtype = getattr(torch, script_args.torch_dtype)
 else:
     device_map = None
     quantization_config = None
@@ -93,9 +147,137 @@
 )

 # Step 2: Load the dataset
-dataset = load_dataset(script_args.dataset_name, split="train")
+print("Loading the dataset...", flush=True)
+if script_args.json_data_files:
+    dataset = load_dataset('json', data_files=script_args.json_data_files, split="train")
+else:
+    dataset = load_dataset(script_args.dataset_name, split="train")
+
+tokenizer = AutoTokenizer.from_pretrained(script_args.model_name, add_eos_token=script_args.add_eos_token)
+if getattr(tokenizer, "pad_token", None) is None:
+    tokenizer.pad_token = tokenizer.eos_token
+if script_args.alternative_pad_token is not None:
+    tokenizer.pad_token = getattr(tokenizer, script_args.alternative_pad_token + "_token")
+
+B_INST, E_INST = "[INST]", "[/INST]"
+B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
+DEFAULT_SYSTEM_PROMPT = script_args.default_system_prompt
+
+def tokenize_dialog_as_llama(dialog, tokenizer, max_len, default_system_prompt=DEFAULT_SYSTEM_PROMPT):
+  prompt_tokens = []
+  if dialog[0]["role"] != "system":
+    dialog = [{"role": "system", "content": default_system_prompt, }] + dialog
+
+  dialog = [
+             {
+               "role": dialog[1]["role"],
+               "content": B_SYS + dialog[0]["content"] + E_SYS + dialog[1]["content"]
+             }
+           ] + dialog[2:]
+
+  assert all([msg["role"] == "user" for msg in dialog[::2]]) and all(
+    [msg["role"] == "assistant" for msg in dialog[1::2]]
+  ), (
+    "model only supports 'system', 'user' and 'assistant' roles, "
+    "starting with 'system', then 'user' and alternating (u/a/u/a/u...)"
+  )
+
+  dialog_tokens = sum(
+                    [ [tokenizer.bos_token_id] + tokenizer.encode(
+                        f"{B_INST} {(prompt['content']).strip()} {E_INST} {(answer['content']).strip()} ",
+                        add_special_tokens=False
+                      ) + [tokenizer.eos_token_id] for prompt, answer in zip(dialog[:-2:2], dialog[1:-1:2])
+                    ],
+                    []
+                  )
+
+  assert (dialog[-2]["role"] == "user"
+    ), f"Before the last message must be from user, got {dialog[-2]['role']}"
+  assert (dialog[-1]["role"] == "assistant"
+    ), f"Last message must be from assistant, got {dialog[-1]['role']}"
+
+  last_tokens = [tokenizer.bos_token_id]  + tokenizer.encode(
+                     f"{B_INST} {(dialog[-2]['content']).strip()} {E_INST} {(dialog[-1]['content']).strip()}",
+                     add_special_tokens=False
+                   ) + [tokenizer.eos_token_id]
+
+  dialog_tokens += last_tokens
+
+  prompt_tokens.append(dialog_tokens)
+
+  input_ids = np.array(prompt_tokens)
+  return build_batch_encodng(input_ids, max_len)
+
+def tokenize_dialog_as_youri(dialog, tokenizer, max_len, default_system_prompt=DEFAULT_SYSTEM_PROMPT):
+  role_map = {
+    "system": "設定",
+    "user": "ユーザー",
+    "assistant": "システム"
+  }
+  if dialog[0]["role"] != "system":
+    dialog = [{"role": "system", "content": default_system_prompt, }] + dialog
+  prompt = []
+  for utter in dialog:
+    prompt.append(f"{role_map[utter['role']]}: {utter['content']}")
+  prompt = "\n".join(prompt)
+  prompt_tokens = tokenizer.encode(prompt, add_special_tokens=True)
+  input_ids = np.array([prompt_tokens])
+  return build_batch_encodng(input_ids, max_len)
+
+def build_batch_encodng(input_ids, max_len):
+  # Truncate overflowed tokens by max_len
+  if input_ids.shape[1] > max_len:
+      print(f"WARN: Too long example(={input_ids.shape[1]}) detected "
+               f"and truncated to max sequence length({max_len}).")
+  input_ids = input_ids[:, :max_len]
+  attention_mask = np.ones_like(input_ids)
+  return BatchEncoding(data={"input_ids": input_ids, "attention_mask": attention_mask})
+
+def prepare_instruction_dataloader(tokenizer, tokenize_dialog_fn, dataset, dataset_text_field, max_seq_len):
+    def tokenize(element):
+        messages = element[dataset_text_field]
+        batch_encodings = [tokenize_dialog_fn(message, tokenizer, max_seq_len) for message in messages]
+        input_ids = [batch_encoding.input_ids[0] for batch_encoding in batch_encodings]
+        attention_mask = [batch_encoding.attention_mask[0] for batch_encoding in batch_encodings]
+        return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+    tokenized_dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)
+    # This is to bypass type check in
+    # https://github.com/huggingface/trl/blob/34e6948d459540a21f80c5be227fb4da039dd97a/trl/trainer/sft_trainer.py#L250
+    class DatasetWrapper(torch.utils.data.Dataset):
+        def __init__(self, dataset):
+            self.dataset = dataset
+        def __len__(self):
+            return len(self.dataset)
+        def __getitem__(self, idx):
+            return self.dataset[idx]
+    return DatasetWrapper(tokenized_dataset)
+
+if script_args.llama_instruction and script_args.youri_instruction :
+  print("llama_instruction and youri_instruction are mutually exclusive. You should choice one of them.", flush=True)
+  sys.exit(1)
+
+if script_args.llama_instruction or script_args.youri_instruction :
+  if script_args.llama_instruction :
+    print("Tokeniing the dataset with llama instruction format...", flush=True)
+    tokenize_dialog_fn = tokenize_dialog_as_llama
+  elif script_args.youri_instruction :
+    print("Tokeniing the dataset with youri instruction format...", flush=True)
+    tokenize_dialog_fn = tokenize_dialog_as_youri
+  else:
+    raise Exception("We should never reach here")
+
+  dataset = prepare_instruction_dataloader(
+                  tokenizer=tokenizer,
+                  tokenize_dialog_fn = tokenize_dialog_fn,
+                  dataset=dataset,
+                  dataset_text_field=script_args.dataset_text_field,
+                  max_seq_len=script_args.seq_length
+  )
+

 # Step 3: Define the training arguments
+print("Define the training arguments...", flush=True)
 training_args = TrainingArguments(
     output_dir=script_args.output_dir,
     per_device_train_batch_size=script_args.batch_size,
@@ -109,9 +291,11 @@
     save_total_limit=script_args.save_total_limit,
     push_to_hub=script_args.push_to_hub,
     hub_model_id=script_args.hub_model_id,
+    group_by_length=script_args.group_by_length,
 )

 # Step 4: Define the LoraConfig
+print("Define the LoraConfig...", flush=True)
 if script_args.use_peft:
     peft_config = LoraConfig(
         r=script_args.peft_lora_r,
@@ -123,8 +307,10 @@
     peft_config = None

 # Step 5: Define the Trainer
+print("Define the SFT Trainer...", flush=True)
 trainer = SFTTrainer(
     model=model,
+    tokenizer=tokenizer,
     args=training_args,
     max_seq_length=script_args.seq_length,
     train_dataset=dataset,
@@ -132,7 +318,39 @@
     peft_config=peft_config,
 )

+print("Start training...", flush=True)
 trainer.train()

 # Step 6: Save the model
-trainer.save_model(script_args.output_dir)
+if script_args.merge_and_push:
+    print("Saving the lola wegiht ...", flush=True)
+    output_dir = os.path.join(script_args.output_dir, "final_checkpoints")
+    trainer.model.save_pretrained(output_dir)
+    tokenizer.save_pretrained(output_dir)
+
+    # Free memory for merging weights
+    del model
+    torch.cuda.empty_cache()
+
+    from peft import AutoPeftModelForCausalLM
+
+    print("Loading the lola wegiht ...", flush=True)
+    model = AutoPeftModelForCausalLM.from_pretrained(
+                output_dir,
+                device_map="auto",
+                torch_dtype=getattr(torch, script_args.torch_dtype),
+    )
+
+    print("Merging the lola wegiht ...", flush=True)
+    model = model.merge_and_unload()
+
+    print("Saving the merged model ...", flush=True)
+    output_merged_dir = os.path.join(script_args.output_dir, "final_merged_checkpoint")
+    model.half()
+    model.save_pretrained(output_merged_dir, safe_serialization=False)
+    tokenizer.save_pretrained(output_merged_dir)
+else:
+    print("Saving the lola wegiht ...", flush=True)
+    output_dir = os.path.join(script_args.output_dir, "final_checkpoints")
+    trainer.model.save_pretrained(output_dir)
+    tokenizer.save_pretrained(output_dir)

trl の sft_trainer.py をコピーして、先程の patch を適用します。

!cp trl/examples/scripts/sft_trainer.py .
!patch sft_trainer.py < patch
!chmod 755 ./sft_trainer.py

この patch で追加したオプションは以下のとおりです。

• json_data_files :
  JSONL形式のデータファイルへのパスです。JSON の内容は後述します。
• add_eos_token :
  トークナイズの末尾に EOS トークンを追加するフラグです。
• alternative_pad_token :
  PAD トークンの代用として使うトークンを指定します。
• default_system_prompt :
  デフォルトのシステムプロンプトです。
• llama_instruction :
  Llama2 ベースの指示応答モデル向けのトークナイズ処理を使用する為のフラグです。
• youri_instruction :
  りんなさんの Youri 7B のチャットモデル(rinna/youri-7b-chat)向けのトークナイズ処理を使用する為のフラグです。この記事を書いた後で、Youri 7B が公開されたのでオマケ的に追加しました。
• torch_dtype : モデルをロードするときのデータ型です。Tesla T4 は bfloat16 が使えないので float16 を指定します。
• use_auth_token :
  モデルをダウンロードするときの Hugging Face の認証トークンの要否のフラグです。Elyza さんの Llama2 は認証トークンが不要だったので、当該処理をバイパスする為に追加したんだと思います。
• merge_and_push :
  学習終了後に QLoRA のアダプタのパラメータを事前学習モデル本体にマージする為のフラグです。
• merge_only :
  同上ですが、学習は行わずマージのみをする為のフラグです。メモリ 16 GB で merge_and_push したらメモリ不足でコケたので、コマンド実行を二回に分けたんでした。
• group_by_length :
  似たような長さのサンプルをまとめてバッチ化することでメモリを節約するフラグです。

Llama2 の指示応答モデルはAIとユーザーの会話履歴の境界に BOS や EOS が入るので、単純に「所定のフォーマットで文字列を作ってトークナイザでエンコード」という訳に行かず面倒くさいんですよね。処理の大半はここ⁶から移植しました。

追記

原稿を書いてから気付いたのですが、transformers の 4.34.0 から chat_template という機能⁷が追加されて、この辺りの処理がラクになってます。 なぜ、移植するまえに気付かなかったのか（涙）。。。

!pip install -U transformers==4.34.0

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("elyza/ELYZA-japanese-Llama-2-7b-fast-instruct")

chat = [
  {"role": "user", "content": "こんにちは。あなたは誰?"},
  {"role": "assistant", "content": "私はAIのアシスタントです。どんなお手伝いをしましょうか?"},
  {"role": "user", "content": "apply_chat_template の使い方を知りたいです。"},
]

tokenizer.use_default_system_prompt = False
tokenizer.apply_chat_template(chat, tokenize=False)
# <s>[INST] こんにちは。あなたは誰? [/INST] 私はAIのアシスタントです。どんなお手伝いをしましょうか? </s><s>[INST] apply_chat_template の使い方を知りたいです。 [/INST]

では本題に戻り、次は学習に使ったデータです。

学習データ

学習に使ったデータですが、とりあえず「指示応答形式のデータでファインチューニングが動くことが確認できればヨシ！」として、 izumi-lab さんのデータセット⁸を利用させて頂きました。ベースとしたモデルは指示応答のファインチューニングがなされているので、 このデータで学習することの意味合いはないです。あくまで学習の動かし方を示す目的ですね⁹。

とりあえずデータセットをロードして、

from datasets import load_dataset
dataset = load_dataset("izumi-lab/llm-japanese-dataset-vanilla", revision="1.0.0")
print(dataset.num_rows)
# {'train': 2515626}

JSONL の形式に加工して出力します。

PROMPT = "あなたは AI のアシスタントです。ユーザーの質問に回答してください。"
WITH_KNOWLEDGE = "\n回答の際は以下の情報を参考にして下さい。\n<<参考情報>>\n"
examples = []
for i in range(1000):
  example = dataset["train"][i]
  if example["input"] == "":
    examples.append({"messages": [
        {"role": "system", "content": PROMPT},
        {"role": "user", "content": example["instruction"]},
        {"role": "assistant", "content": example["output"]}
    ]})
  else:
    examples.append({"messages": [
        {"role": "system", "content": PROMPT + WITH_KNOWLEDGE + example["input"]},
        {"role": "user", "content": example["instruction"]},
        {"role": "assistant", "content": example["output"]}
    ]})

import json
with open("./llm-japanese-dataset-vanilla.jsonl", "w") as f:
  f.write("\n".join([json.dumps(example, ensure_ascii=False) for example in examples]))    

出来上がったファイルの中身はこんな感じです。OpenAI の Chat Completions API 風な感じにしてみました。

!head -5 llm-japanese-dataset-vanilla.jsonl
# {"messages": [{"role": "system", "content": "あなたは AI のアシスタントです。ユーザーの質問に回答してください。"}, {"role": "user", "content": "「abc ～the first～」へようこそ！さて、ABC・・・と始まるアルファベットは、全部で何文字でしょう？"}, {"role": "assistant", "content": "26文字"}]}
# {"messages": [{"role": "system", "content": "あなたは AI のアシスタントです。ユーザーの質問に回答してください。"}, {"role": "user", "content": "人気漫画『ドラえもん』の登場人物で、ジャイアンの苗字は剛田ですが、スネ夫の苗字は何でしょう？"}, {"role": "assistant", "content": "骨川（滑川も正解）"}]}
# {"messages": [{"role": "system", "content": "あなたは AI のアシスタントです。ユーザーの質問に回答してください。"}, {"role": "user", "content": "格闘家ボブ・サップの出身国はどこでしょう？"}, {"role": "assistant", "content": "アメリカ"}]}
# {"messages": [{"role": "system", "content": "あなたは AI のアシスタントです。ユーザーの質問に回答してください。"}, {"role": "user", "content": "ロシア語で「城」という意味がある、ロシアの大統領府の別名は何でしょう？"}, {"role": "assistant", "content": "クレムリン"}]}
# {"messages": [{"role": "system", "content": "あなたは AI のアシスタントです。ユーザーの質問に回答してください。"}, {"role": "user", "content": "織田信長、豊臣秀吉、徳川家康という３人の戦国武将の性格を表現するのに用いられる鳥は何でしょう？"}, {"role": "assistant", "content": "ホトトギス"}]}

では学習の実行です。

学習の実行

以下のようにして学習を実行します。

!./sft_trainer.py \
    --model_name elyza/ELYZA-japanese-Llama-2-7b-fast-instruct \
    --json_data_files ./llm-japanese-dataset-vanilla.jsonl \
    --dataset_text_field messages \
    --seq_length 1024 \
    --add_eos_token \
    --alternative_pad_token unk \
    --llama_instruction \
    --load_in_4bit \
    --torch_dtype float16 \
    --use_peft \
    --batch_size 4 \
    --gradient_accumulation_steps 8 \
    --num_train_epochs 1 \
    --learning_rate 1e-5 \
    --group_by_length

設定してるパラメータは npaka さんの記事¹⁰をすごく参考にさせてもらいました（参考というかそのままというか。。。）!

処理的に違っているのはトークナイズのところぐらいですかね。今回ベースにした Elyza さんの Llama2¹ の Usage を見ると、 システムプロンプトの制御記号である "<<SYS>>"、"<</SYS>>" を埋め込む必要があるようなので、そのあたりを追加しています。

PAD トークンを UNK に設定する理由

--alternative_pad_token unk で PAD トークンに UNK を指定する理由について少し補足しましょう。Elyza さんの Llama2¹のトークナイザはデフォルトで PAD に EOS("</s>") が割当たってるんです(というかオリジナルの Llama 2 には PAD トークンというものがないらしい)。

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("elyza/ELYZA-japanese-Llama-2-7b-fast-instruct", add_eos_token=True) 
tokenizer.pad_token
# '</s>'    

で、この場合どうなるかというと、

• 入力文字列が “おはよう。” であったとします("<<SYS>>" や "[INST]" は、ここでは無視）。
• トークナイズ処理で末尾に EOS (トークンID = 2) が付与されます。
  • [1, 29871, 30697, 30449, 42680, 30267, 2]
• この系列が他のサンプルと束ねられてバッチを形成し DataCollatorForLanguageModeling（だったかな？）に投入されてパディングされます。ただし、PAD トークン = EOS なので以下のようになります。
  • [1, 29871, 30697, 30449, 42680, 30267, 2, 2, 2, 2, 2]
• さらに Collator の中でラベル系列が生成されるのですが、PAD トークン（= EOS) は -100 (損失の計算から除外する記号) に置き換えられます。ソースでいうとこの辺り¹¹ですかね。
  • [1, 29871, 30697, 30449, 42680, 30267, -100, -100, -100, -100, -100]

これで、末尾を示す記号であったはずの 30267 の直後の 2 がパディングと混同されて -100 に置き換わりました。なのでモデルは出力の終わりを学習できなくなるということのようです。

話をもとに戻しましょう。

学習が終わった段階では QLoRA のアダプタのパラメータが出力されている状態なので、これを元のモデルのパラメータにマージします。

!./sft_trainer.py --merge_only

!du -sh output/*
# 131M    output/final_checkpoints
# 13G     output/final_merged_checkpoint

final_merged_checkpoint がマージされたパラメータですね。流石に 7B だと嵩張りますね。。。

ではここから vLLM にデプロイして推論してみましょう。

4. vLLM での推論

vLLM¹² は LLM を高速に推論するためのライブラリです。この連載では FasterTransformer を主に使っていましたが、 本記事を書いている時点では FasterTransformer が Llama に正式対応していなかったので vLLM に乗り換えました¹³。 vLLM も FasterTransformer をベースに開発されたそうなので、推論速度的には同レベルの性能を期待してよさそうです。

以下のようにして vLLM のインストールします。

7B サイズのモデルを Tesla T4 で動かそうとすると、 vLLM 0.1.4 と xformers 0.0.21 の組み合わせで、 ギリギリ動く感じでした。(正直、長めの入力とか入れるとマズイかもしれません。。。)

!pip install \
  vllm==0.1.4 \
  fschat==0.2.26 \
  accelerate==0.22.0 \
  xformers==0.0.21 \
  ray[data,train,tune,serve]==2.7.1 \
  pandas==2.1.1 \
  pyarrow==13.0.0

vLLM はモデルの名称からトークナイザを推測するので、シンボリックリンクでごまかしました（多分、コマンドラインオプションもあると思いますが）。

!ln -s output/final_merged_checkpoint `pwd`/finetuned-llama-2

以下のようにして OpenAI 互換の API サーバを起動します。ギリギリで動かすので --gpu-memory-utilization 1.0 を指定してます。

!python -m vllm.entrypoints.openai.api_server \
  --gpu-memory-utilization 1.0 \
  --host 0.0.0.0 \
  --port 80 \
  --model `pwd`/finetuned-llama-2

起動後の GPU の状態を見ると以下のようになってます。大丈夫なんでしょうか。。。

Wed Nov 22 08:55:27 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.65.01    Driver Version: 515.65.01    CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| 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   48C    P0    27W /  70W |  14699MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      4766      C   python                          14687MiB |
+-----------------------------------------------------------------------------+

起動した API をコールしてみましょう。vLLM を起動したコンテナに別の端末から入りなおします。

$ docker exec -it hiroba bash
# cd /work

ちょっと API 呼び出しに使うコマンドを入れさせてもらって。。。

!apt-get update && apt-get install -y curl jq

以下のように ChatGPT の API と同じ感覚でコールできます。

%%bash
CONTENT='
{
  "model": "/work/hiroba/finetuned-llama-2",
  "max_tokens": 512,
  "stop": ["["],
  "messages": [
    {"role": "system", "content": "あなたは日本の文化の専門家です。"}, 
    {"role": "user", "content": "渋谷の観光地について教えてください。"}
  ]
}'

curl -X POST \
     -H "Content-Type: application/json" \
     -d "$CONTENT" \
     --silent \
     http://localhost:80/v1/chat/completions | jq .choices[0].message.content | cat

# 渋谷は、日本を代表する観光スポットの一つです。渋谷駅周辺は、ファッションやグルメ、エンターテイメントなどの
# 様々なジャンルのショップやレストランが立ち並ぶエリアです。シブヤ大学や渋谷区立道玄坂図書館など、独自の観光
# スポットも多いです。\n\nカフェやレストランも多く、デートや友人との外食に適したエリアとなっています。\n\n
# 東急本店やパルコ、渋谷109などのショッピングモール、ライブハウスやクラブなどもあるので、お土産を買うのも、
# エンターテイメントを楽しむのも、楽しみ方は様々です。

動きました。イイ感じ、かと思ったのですがシブヤ大学とか言ってますね。渋谷区立道玄坂図書館も存在しないかと。 適当にファインチューニングしたのがいけなかったかもしれませんが、とりあえず回答の生成ができました。

ついでに LangChain からも使ってみましょう。

5. LangChain からの接続

LangChain 関係をインストールして、

!pip install langchain==0.0.335 openai==1.2.4

以下の要領で LangChain の ChatOpenAI クラスから利用できます。openai_api_key はダミー文字列を設定しておきます。

from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage
import os

chat = ChatOpenAI(
          model=f"{os.getcwd()}/finetuned-llama-2",
          openai_api_base="http://localhost:80/v1",
          openai_api_key="DUMMY",
          temperature=0,
          max_tokens=1024,
)
messages = [
    SystemMessage(content="あなたは日本の文化の専門家です。"),
    HumanMessage(content="渋谷の観光地について教えてください。")
]
print(chat(messages).content)

#  渋谷は、東京の代表的な観光スポットの1つです。ハチ公像やスクランブル交差点、東急ハンズなど、
# 多くの観光スポットがあります。特に、スクランブル交差点は、東京を代表する観光スポットとして有名です。
# 渋谷駅周辺は、若者の街として知られていますが、観光スポットも多く、楽しめるエリアとなっています。

さすがに Function Calling は無理ですが、タスクを選べばイケるかもしれませんね！

ただ、ギリギリで動いているのが気になります。今度は量子化を使って省メモリ化してみましょう。

6. AWQ による量子化

※メモリ 16 GB ではどうがんばってもムリだったので、この章だけ VM のメモリを 32 GB に増やし、コンテナも再起動しています。

モデルの量子化には AWQ(Activation-aware Weight Quantization)¹⁴ を使いました。 細かい説明はしませんが、以下の図がわかりやすいと思います。

左側が単純に FP16 を INT3 に量子化した場合(PPL=43.2)、中央が重みの 1 % を FP16 で維持した場合(PPL=13.0)、右側が AWQ の場合(PPL=13.0)です。

中央は活性化の分布を観測し、重みの中で出力に大きな影響を与える 1 % を FP16 で維持することにより出力品質を向上させられることを示しています。

ただし、一部の重みを FP16 とする混合精度演算は計算効率の点で不利になります。 そこで AWQ(右) では活性化の分布にもとづきチャネル毎にスケーリング係数を調整することで、 全ての重みを INT3 に量子化しつつ、混合精度演算(中央)と同等の出力品質を実現しました。

さて、実装としては AutoAWQ¹⁵ が公開されているのでインストールしていきます。 少し前のバージョンですが、これより新しいと Tesla T4 x 1 で量子化した時、GPU メモリが OOM になりました。

!pip install autoawq==0.1.5

では、さっそく量子化してみましょう。スクリプト的には以下のようになります。

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer
from datasets import load_dataset

model_path = 'output/final_merged_checkpoint'
quant_path = 'output/final_quantized_checkpoint/'
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }

model = AutoAWQForCausalLM.from_pretrained(model_path, **{"low_cpu_mem_usage": True})
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

tokenizer.use_default_system_prompt = False
dataset = load_dataset('json', data_files=["./llm-japanese-dataset-vanilla.jsonl"], split="train")
dataset = dataset.shuffle(seed=42)
texts = []
for example in dataset:
  messages = example["messages"]
  texts.append(tokenizer.apply_chat_template(messages, tokenize=False))

print(f"Calibration data example:\n{texts[0]}", flush=True)
# Calibration data example:
# <s>[INST] <<SYS>>
# あなたは AI のアシスタントです。ユーザーの質問に回答してください。
# <</SYS>>
# 
# 戦争で両腕両足、目、耳、口を失った男の姿を描いた、ドルトン・トランボの代表作は何でしょう？ [/INST] 『ジョニーは戦場へ行った』 </s>

model.quantize(tokenizer, quant_config=quant_config, calib_data=texts)

model.save_quantized(quant_path, safetensors=False)
tokenizer.save_pretrained(quant_path)

AWQ の量子化は “Activation-aware” という名のとおりアクティベーションの値を用いるので、 何かしらモデルにデータを入力する必要があります。デフォルトでは英語のデータセットが使われるのですが、 普通に考えれば推論時に入力する、日本語のデータを使ったほうが良さそうです。

今回は学習に使ったデータを LLama2 のチャットモデルのフォーマットに直して投入してみました。 細かいことを言うと、AutoAWQ のコード的には、"<s>“, ”</s>“ が正しく BOS, EOS トークンにエンコードされずに ”<“, “s”, “>” のようになってしまうのですが、改造するのも面倒だったので「英語のデータよりはマシでしょ」ということで気にしないことにしています。

もう一点、量子化の設定で GEMM と GEMV がありますが、GEMV を使うと vLLM にデプロイしたときエラーになったので、今回は GEMM を使っています。

量子化が終わった状態は以下のとおりです。かなり小さくなりましたね！

!du -sh output/*
131M    output/final_checkpoints
13G     output/final_merged_checkpoint
3.9G    output/final_quantized_checkpoint

では量子化したモデルを vLLM にデプロイして推論してみましょう。

AWQ で量子化したモデルをデプロイするので、vLLM のバージョンを最新(0.2.2)のものにしました。 vLLM 0.2.2 は CUDA 12 ベースなので、今まで使っていたコンテナを止めて新しく作り直します。

$ docker run --name hiroba-vllm -v `echo $HOME`:/work --gpus all -it pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel
# cd work

以下のようにして vLLM 0.2.2 のインストールです。

!pip install \
  vllm==0.2.2 \
  fschat==0.2.32 \
  accelerate==0.24.1 \
  xformers==0.0.22.post7 \
  ray[data,train,tune,serve]==2.8.0 \
  pandas==2.1.3 \
  pyarrow==14.0.1

量子化したモデルに "llama-2” が含まれるようにして、

!ln -s output/final_quantized_checkpoint `pwd`/quantized-llama-2

以下のようにして OpenAI 互換の API サーバを起動します。量子化したモデルをロードするので --quantization=awq を指定しています。

!python -m vllm.entrypoints.openai.api_server \
  --quantization=awq \
  --host 0.0.0.0 \
  --port 80 \
  --model `pwd`/quantized-llama-2

# INFO 11-27 00:12:17 api_server.py:638] args: Namespace(host='0.0.0.0', port=80, allow_credentials=False, allowed_origins=['*'], allowed_methods=['*'],
# allowed_headers=['*'], served_model_name=None, model='/work/hiroba/quantized-llama-2', tokenizer=None, revision=None, tokenizer_revision=None,
# tokenizer_mode='auto', trust_remote_code=False, download_dir=None, load_format='auto', dtype='auto', max_model_len=None, worker_use_ray=False,
# pipeline_parallel_size=1, tensor_parallel_size=1, block_size=16, seed=0, swap_space=4, gpu_memory_utilization=0.9, max_num_batched_tokens=None,
# max_num_seqs=256, max_paddings=256, disable_log_stats=False, quantization='awq', engine_use_ray=False, disable_log_requests=False, max_log_len=None)
# WARNING 11-27 00:12:17 config.py:140] awq quantization is not fully optimized yet. The speed can be slower than non-quantized models.
# INFO 11-27 00:12:17 llm_engine.py:72] Initializing an LLM engine with config: model='/work/hiroba/quantized-llama-2', tokenizer='/work/hiroba/quantized-llama-2',
# tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=4096, download_dir=None,
# load_format=auto, tensor_parallel_size=1, quantization=awq, seed=0)
# INFO 11-27 00:12:17 tokenizer.py:31] For some LLaMA V1 models, initializing the fast tokenizer may take a long time. To reduce the initialization time,
# consider using 'hf-internal-testing/llama-tokenizer' instead of the original tokenizer.
# INFO 11-27 00:12:55 llm_engine.py:207] # GPU blocks: 979, # CPU blocks: 512
# INFO 11-27 00:12:57 tokenizer.py:31] For some LLaMA V1 models, initializing the fast tokenizer may take a long time. To reduce the initialization time,
# consider using 'hf-internal-testing/llama-tokenizer' instead of the original tokenizer.
# INFO:     Started server process [43]
# INFO:     Waiting for application startup.
# INFO:     Application startup complete.
# INFO:     Uvicorn running on http://0.0.0.0:80 (Press CTRL+C to quit)

“awq quantization is not fully optimized yet. ” というのが気になりますが、 構成や設定に問題があるというより「実装としては、まだ追い込んでないよ」ということでしょうかね。

起動完了時点の GPU メモリ使用状況はこんな感じです。

Wed Nov 22 08:46:20 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.65.01    Driver Version: 515.65.01    CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| 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   44C    P0    27W /  70W |  12890MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     32375      C   python                          12878MiB |
+-----------------------------------------------------------------------------+

4bit に量子化したのでもっと減るのを期待していましたが、だいぶマシにはなりましたね。

細かい手順は省略しますが、先程と同様にコールしてみましょう。

CONTENT='
{
  "model": "/work/hiroba/quantized-llama-2",
  "max_tokens": 512,
  "stop": ["["],
  "messages": [
    {"role": "system", "content": "あなたは日本の文化の専門家です。"}, 
    {"role": "user", "content": "渋谷の観光地について教えてください。"}
  ]
}'

curl -X POST \
     -H "Content-Type: application/json" \
     -d "$CONTENT" \
     --silent \
     http://localhost:80/v1/chat/completions | jq .choices[0].message.content | cat

# 渋谷の観光地について説明します。\n\n渋谷は東京の代官山や神宮前の次に、若い世代のオススメの観光エリアです。 \n\n
# まずは、渋谷駅東急東横線の渋谷駅のホームから、東京の夜景を楽しみましょう。 \n\n東急百貨店の地下1階にある
# 「道玄坂ヶ原」は、渋谷で絵画を鑑賞することができます。 \n\nさらには、NHKの朝の連続テレビドラマ「さくら」
# のロケ地として知られる、代官山エリアへも足をのばしてください。 代官山は渋谷からバスで15分の距離で、
# 美しい丘陵地帯に広がる高級住宅街です。 \n\nまた、渋谷にはショッピングモール「渋谷109」や、水族館「アクアパーク品川」
# など、見どころが多く、1日では回りきれないほどの観光スポットが満載です。 1日限りの東京観光をお楽しみください。

回答を生成できました。とりあえず、それっぽくは出力されてますが、東急の地下1階に道玄坂ヶ原はないみたいです。 朝ドラはよくしりませんが、渋谷にアクアパーク品川があるのはおかしいですね。

今回はファインチューニングの仕方と学習済みモデルでの推論という観点でまとめたので、 出力内容の精度、正しさに関しては、素の ELYZA-japanese-Llama-2-7b-fast-instruct がどの程度の能力なのか、 それが量子化にどの程度影響されるのか、は落ち着いて評価したほうがよさそうですね。

6. おわりに

今回は Llama 2 の指示応答モデルのファインチューニングと推論ということででした。コード生成用の CodeLlama も Llama2 の学習データ違いみたいなので、今回の内容ベースで自社のコード向けにファインチューニングして活用とか試してみたいですね。

「誰も止めないから」という理由でなんとなく書いてましたが、気づいたら第30回？エラい続いたもんですね。。。 最初は BERT ですらなくって単語ベクトルとか言ってたんですよ。ホンの5,6年で変わったもんです。

最近は書くことないというか、変化が激しすぎて書いてられないようなところもあるのですが、来年も書きたいことがあったら書こうかと思ってます。 では、よいお年を。。。