Related Resources
Generative AI

NeMo Framework で日本語 LLM をファインチューニング – PEFT 編 –

2023 年 11 月 28 日
By
Reading Time: 3 minutes

NeMo Framework は生成 AI モデルを構築、カスタマイズ、展開するためのエンドツーエンドのクラウドネイティブ フレームワークです。本記事では、NeMo Framework を使用して、日本語の大規模言語モデル (LLM) の PEFT (ファインチューニングの手法の一種)を実行する方法を説明します。

NeMo Framework とは

NeMo Framework は、生成 AI モデルのワークフローをエンドツーエンドでカバーするフレームワークで、現在は LLM を中心とした NLP 向けのコンテナーが NGC 上で公開されています。

こちらからアクセスして ”Download now-language” をクリックすると、承認後に NGC 上で NeMo Framework のコンテナーが入手できます (”Apply Now -Multimodal” はマルチモーダル向けの Early Accessになります)。

NVIDIA AI Enterprise ライセンスをお持ちの方は、NGC サイトから入手可能です。NGC へログイン後、Enterprise Catalog にある ”Feature Branches & Models” にアクセスしてください。こちらで NeMo Framework の入手方法をご案内しています (現在は Training コンテナーに対応しています)。

NeMo Framework のコンテナー イメージには、Training と Inference がありますが、本記事では事前学習およびファインチューニング向けの機能を提供する Training コンテナーについて解説します。

LLM のワークフロー

LLMの開発におけるタスクには以下のようなものがあります。

  • 事前学習に必要な大量データの準備
  • 分散学習を利用した LLM の事前学習
  • LLM をカスタマイズするためのファインチューニング (SFT, RLHF, PEFT) やプロンプト エンジニアリング
  • LLM の推論を高速化するためのモデル最適化
  • GPU を最大限に活用した LLM のサービング
  • コストを抑えながら LLM に最新情報を反映させるための RAG
  • LLM アプリケーションの意図しない挙動を抑えるためのガードレール

LLM の開発、サービスでの利用には多くのステップが必要になりますが、NeMo Framework の Training コンテナーには、データの準備から LLM の学習、ファインチューニングに必要な下記モジュールが含まれており、これらを使用することで LLM の学習に関するステップを 1 つのコンテナー環境で実行できます。

  • NVIDIA NeMo Data Curator
    LLM の学習に必要な大規模データセットのダウンロードから抽出、クリーニング、フィルタリングを行うためのスケーラブルなツールキット。
  • NeMo
    自動音声認識 (ASR)、テキスト音声合成 (TTS)、および自然言語処理 (NLP) のためのツールキット。
  • NeMo-Megatron-Launcher
    クラウド、オンプレのクラスターからトレーニング ジョブを起動するためのツールキット。
  • Megatron-LM
    Transformer モデルの大規模学習に関する研究プロジェクト。
  • TransformerEngine
    FP8 での学習を中心とした Transformer モデルを高速化させるツールキット。
  • NeMo-RLHF
    人間のフィードバックからの強化学習 (RLHF) を使用して LLM をファインチューニングするためのツールキット。

上記の一部は、GitHub 上に OpenSource として公開されていますが、依存関係が解消されている NeMo Framework の Training コンテナーから利用することをお薦めします。コンテナーの場合、opt ディレクトリ配下に上記のモジュールが配置されています。

PEFT とは

PEFT (Parameter-Efficient Fine-Tuning) とは、LLM をカスタマイズする際に使われる手法の一種です。元の LLM に少数のパラメーターやレイヤーを追加し、ユース ケース固有のデータで学習を行います。元の LLM の重みは固定されたままであるため、学習中に更新されるパラメーターが大幅に少なくなります。

これにより、LLM は全てのパラメーターを更新するファインチューニングと比較して、少ないコンピューティング リソースで特定のユース ケースでのパフォーマンスを改善できます。

大規模言語モデルのカスタマイズ手法を選択するには、PEFT を含むカスタマイズ手法の解説があります。

PEFT チュートリアル

本記事では、Hugging Face Model Hub から日本語 LLM をダウンロードして、NeMo Framework を使用した PEFT を実行します。

本チュートリアルでの手順は以下の通りです。

  • PEFT を実行するための事前準備
  • NeMo Framework のコンテナーを起動
  • Hugging Face Model Hub から事前学習済みのモデルをダウンロード
  • ダウンロードしたモデルを nemo フォーマットへ変換
  • PEFT に使用するデータの準備および前処理
  • PEFT の実行
  • 推論

事前準備

  • NGC から最新の NeMo Framework Training コンテナーを入手します (このチュートリアルでは、nvcr.io/ea-bignlp/ga-participants/nemofw-training:23.08.03 を使用しています)。
  • GPU が搭載されたマシンを準備します (このチュートリアルのコードは OS: Ubuntu 20.04.4 LTS, NVIDIA  A100 80GB × 1 の環境でテストしています)。Llama2 で PEFT を実行する際の HW 要件はこちらに記載があります。
  • 以下のコマンドで作業用のディレクトリを作成し、移動します。
mkdir peft-example
cd peft-example

Docker コンテナーの起動

以下のコマンドでコンテナーを起動します。

sudo docker run --rm -it --gpus device=0 --shm-size=2g --ulimit memlock=-1 --network=host -v ${PWD}:/workspace  -w /workspace -v ${PWD}/results:/workspace/results nvcr.io/ea-bignlp/ga-participants/nemofw-training:23.08.03 bash

Hugging Face Model Hub からモデルをダウンロード

このチュートリアルでは、elyza/ELYZA-japanese-Llama-2-7b を使用します。以下のコードで Hugging Face の Model Hub から事前学習済みの LLM をダウンロードします。

import os
from huggingface_hub import snapshot_download
 
MODEL_DIR = "./models"
os.makedirs(MODEL_DIR, exist_ok=True)
 
snapshot_download(
    repo_id="elyza/ELYZA-japanese-Llama-2-7b",
    local_dir=f"{MODEL_DIR}/ELYZA-japanese-Llama-2-7b",
    local_dir_use_symlinks=False
    )

nemo フォーマットへの変換

以下のスクリプトを使用して、ダウンロードした Hugging Face の Llama モデルを nemo フォーマットへ変換します。

python /opt/NeMo/scripts/nlp_language_modeling/convert_hf_llama_to_nemo.py \
    --in-file=./models/ELYZA-japanese-Llama-2-7b \
    --out-file=./models/ELYZA-japanese-Llama-2-7b/ELYZA-japanese-Llama-2-7b.nemo \
    --precision="16"

データの準備

この PEFT チュートリアルでは JGLUE の JCommonsenseQA データセットを使用します。

git clone https://github.com/yahoojapan/JGLUE.git

以下のスクリプトを使用して、NeMo が PEFT で必要とする JSONL 形式へデータを変換します。ここでは検証用データを精度を確認するためのテスト データとして使用し、学習用データを分割することで学習および検証データとして使用します。

import json
import os
import random
 
 
INPUT_TRAIN = "./JGLUE/datasets/jcommonsenseqa-v1.1/train-v1.1.json"
INPUT_VALID = "./JGLUE/datasets/jcommonsenseqa-v1.1/valid-v1.1.json"
OUTPUT_DIR = "./data/jcommonsenseqa-v1.1"
 
 
random.seed(42)
os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 
def write_jsonl(fname, json_objs):
    with open(fname, 'wt') as f:
        for o in json_objs:
            f.write(json.dumps(o, ensure_ascii=False)+"\n")
 
 
def form_question(obj):
    st = ""
    st += "### 指示:\n"
    st += "与えられた選択肢の中から、最適な答えを選んでください。出力は以下から選択してください：\n"
    st += f"- {obj['choice0']}\n"
    st += f"- {obj['choice1']}\n"
    st += f"- {obj['choice2']}\n"
    st += f"- {obj['choice3']}\n"
    st += f"- {obj['choice4']}\n"
    st += "### 入力:\n"
    st += f"{obj['question']}\n"
    st += "### 応答:"
    return st
 
 
def prosess(input_path, train=False):
    with open(input_path) as f:
        dataset = [json.loads(line) for line in f.readlines()]
 
    processed = []
    for data in dataset:
        prompt = form_question(data)
        answer = data[f"choice{data['label']}"]
        processed.append({"input": prompt, "output": f"{answer}"})

    if train:
        random.shuffle(processed)
        train_ds = processed[:-1000]
        valid_ds = processed[-1000:]
        write_jsonl(f"{OUTPUT_DIR}/train-v1.1.jsonl", train_ds)
        write_jsonl(f"{OUTPUT_DIR}/valid-v1.1.jsonl", valid_ds)
    else:
        write_jsonl(f"{OUTPUT_DIR}/test-v1.1.jsonl", processed)
 
    return
 
 
def main():
    prosess(INPUT_TRAIN, train=True)
    prosess(INPUT_VALID)
 
 
if __name__ == "__main__":
    main()

スクリプトが実行されると data というディレクトリの下に学習、検証、テストに使用するファイルが生成されます。変換後のデータを開いてみると以下のような形式になっています。

{"input": "### 指示:\n与えられた選択肢の中から、最適な答えを選んでください。出力は以下から選択してください：\n- 猫\n- 家畜\n- 原宿\n- 牧草\n- 島根県\n### 入力:\n動物たちの餌になるものは？\n### 応答:", "output": "牧草"}
{"input": "### 指示:\n与えられた選択肢の中から、最適な答えを選んでください。出力は以下から選択してください：\n- 岩手\n- 福島\n- 秋田\n- 愛知\n- 福井\n### 入力:\nきりたんぽが名物である県はどこか？\n### 応答:", "output": "秋田"}

PEFT の実行

PEFT は /opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_peft_tuning.py で実行できます。/opt/NeMo/examples/nlp/language_modeling/tuning/conf/megatron_gpt_peft_tuning_config.yaml に PEFT の実行に必要な config ファイルがあり、環境変数をセットすることでスクリプトの実行に必要なデータのパス、ベースとなる LLM (nemo フォーマット) のパスなどの設定を渡します (NeMo Framework は config 設定に Hydra を使用しています)。環境変数から設定されていないパラメーターは config ファイルの設定が適用されます。実行コマンドにある ++model.mcore_gpt=True は Hydra で config の追加または上書きをする際の記法になります。

現在、NeMo Framework では以下の PEFT の手法をサポートしており、SCHEME を変更することで他の手法を実行できます。以下は P-Tuning を選択した場合で説明します。PEFT とモデルのサポートリストはこちらにあります。

  • P-Tuning
  • Adapter
  • IA3
  • LoRA
export EXP_DIR="/workspace/results"
export EXP_NAME="elyza_7b_ptuning"
export MODEL="/workspace/models/ELYZA-japanese-Llama-2-7b/ELYZA-japanese-Llama-2-7b.nemo"
export SCHEME="ptuning"
export TP_SIZE=1
export PP_SIZE=1
export TRAIN_DS="[/workspace/data/jcommonsenseqa-v1.1/train-v1.1.jsonl]"
export VALID_DS="[/workspace/data/jcommonsenseqa-v1.1/valid-v1.1.jsonl]"


torchrun --nproc_per_node=1 \
/opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_peft_tuning.py \
    exp_manager.exp_dir=${EXP_DIR} \
    exp_manager.name=${EXP_NAME} \
    exp_manager.early_stopping_callback_params.patience=5 \
    trainer.precision=bf16 \
    trainer.devices=1 \
    trainer.num_nodes=1 \
    trainer.max_epochs=-1 \
    trainer.max_steps=10000 \
    trainer.val_check_interval=0.05 \
    trainer.gradient_clip_val=1.0 \
    model.megatron_amp_O2=False \
    ++model.mcore_gpt=True \
    model.restore_from_path=${MODEL} \
    model.peft.peft_scheme=${SCHEME} \
    model.data.train_ds.file_names=${TRAIN_DS} \
    model.data.train_ds.concat_sampling_probabilities=[1.0] \
    model.data.validation_ds.file_names=${VALID_DS} \
    model.tensor_model_parallel_size=${TP_SIZE} \
    model.pipeline_model_parallel_size=${PP_SIZE} \
    model.global_batch_size=16 \
    model.micro_batch_size=16 \
    model.data.validation_ds.tokens_to_generate=10 \
    model.data.validation_ds.metric.name="loss" \
    model.data.train_ds.num_workers=0 \
    model.data.validation_ds.num_workers=0 \
    model.optim.lr=1e-5

NeMo Frameworkでは、実験管理に Weights and Biases をサポートしており、 wandb へログイン後に以下のコマンドを追加して、スクリプトを実行することで wandb 上で実験管理することができます。

exp_manager.create_wandb_logger=True \
exp_manager.wandb_logger_kwargs.project=${WANDB_PROJECT} \
exp_manager.wandb_logger_kwargs.name=${WANDB_NAME} \

上記の設定では、3,000 ステップほどで改善が止まり、その後、Early stopping がかかったため、学習は 1 時間ほどで完了しました。

学習が終わると results/elyza_7b_ptuning/ という名前のディレクトリが作成され、中に学習時の config や log などが出力されます。また、同じディレクトリの checkpoints 内にある elyza_7b_ptuning.nemo が PEFT で学習されたパラメーターの checkpoint ファイルになります。

推論

学習が完了したら、/opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_peft_eval.py を使用して、テスト データで推論を実行します。model.restore_from_path には先ほどの PEFT でベースとなったモデルのパス、model.peft.restore_from_path には先ほど PEFT を実行して作成されたモデルのパスを渡します。model.data.test_ds.write_predictions_to_file を True に設定すると入力ファイルに推論結果が追加されたファイルが出力されます。

export MODEL="/workspace/models/ELYZA-japanese-Llama-2-7b/ELYZA-japanese-Llama-2-7b.nemo"
export PEFT_MODEL="/workspace/results/elyza_7b_ptuning/checkpoints/elyza_7b_ptuning.nemo"
export TEST_NAMES="[jcommonsenseqa-v1.1]"
export TEST_DS="[/workspace/data/jcommonsenseqa-v1.1/test-v1.1.jsonl]"
export OUTPUT_PREFIX="/workspace/results/elyza_7b_ptuning"
export NVTE_FLASH_ATTN=0
export NVTE_FUSED_ATTN=0
 
 
python /opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_peft_eval.py \
    model.restore_from_path=${MODEL} \
    model.peft.restore_from_path=${PEFT_MODEL} \
    trainer.devices=1 \
    model.data.test_ds.file_names=${TEST_DS} \
    model.data.test_ds.names=${TEST_NAMES} \
    model.data.test_ds.global_batch_size=16 \
    model.data.test_ds.micro_batch_size=16 \
    model.data.test_ds.tokens_to_generate=10 \
    inference.greedy=True \
    model.data.test_ds.output_file_path_prefix=${OUTPUT_PREFIX} \
    model.data.test_ds.write_predictions_to_file=True

上記設定で推論が完了すると results の下に elyza_7b_ptuning_test_jcommonsenseqa-v1.1_inputs_preds_labels.jsonl というファイルが出力されます。

以下に Few-shot(2-shot) での一例を示します。

# 以下はFew-shotでのプロンプトと出力

### 指示:
与えられた選択肢の中から、最適な答えを選んでください。出力は以下から選択してください：
### 例題1:
- 世界
- 写真集
- 絵本
- 論文
- 図鑑
### 入力:
主に子ども向けのもので、イラストのついた物語が書かれているものはどれ？
### 応答: 絵本
### 例題2:
- 浮浪者
- 保護者
- お坊さん
- 宗教者
- 預言者
### 入力:
未成年者を監護・教育し，彼らを監督し，彼らの財産上の利益を守る法律上の義務をもつ人は？
### 応答: 保護者
### 本題:
- 掲示板
- パソコン
- マザーボード
- ハードディスク
- まな板
### 入力:
電子機器で使用される最も主要な電子回路基板の事をなんと言う？
### 応答:
マザーボード
###

次に P-Tuning 後の出力の例を示します。

# 以下はP-Tuning後のプロンプトと出力

### 指示:
与えられた選択肢の中から、最適な答えを選んでください。出力は以下から選択してください：
- 掲示板
- パソコン
- マザーボード
- ハードディスク
- まな板
### 入力:
電子機器で使用される最も主要な電子回路基板の事をなんと言う？
### 応答:
マザーボード

Few-shot ではプロンプトの設定により、出力に対して、改行コードの前までを取り出すなどの処理が必要になりますが、P-Tuning 後では答えのみが出力されるようになりました。また、全体の正解率も P-Tuning 後の方が優れた結果が得られていることを確認しています。

まとめ

本記事では、NeMo Framework を使用した PEFT の実行方法を紹介しました。NeMo Framework を使用して LLM の開発が加速すると嬉しいです。

関連情報

Tags

Generative AI | large language models | NeMo Framework | technical walkthrough

About the Authors

About Atsunori Fujita
藤田 充矩は、NVIDIA のシニア AI ソリューション アーキテクトとして、特に AI/ML 分野において、顧客への技術支援を行っています。NVIDIA に入社する前は、自動車やクラウド、コンサルティング企業にてデータ サイエンティスト、機械学習エンジニアとして働いており、機械学習モデルのプロトタイピングやプロダクション環境へのデプロイに従事していました。現在は、NVIDIA の NeMo に重点を置き、AI/ML 領域のソリューション開発に取り組んでいます。