Related Resources
Generative AI

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

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

ご注意: この記事は NeMo Framework のアップデートのため、2024 年 6 月 10 日に大幅に変更を加えました。

本記事では、NeMo Framework を使用して、日本語の大規模言語モデル (LLM) の PEFT (ファインチューニングの手法の一種)を実行する方法を説明します。

NeMo Framework とは

NeMo Framework は、LLM をはじめ、生成 AI モデルを構築、カスタマイズするためのクラウドネイティブなフレームワークです。NGC 上にコンテナーが公開されており、すぐに利用を開始することができます。

NeMo Framework は、NVIDIA AI Enterprise の対象ソフトウェアになっているため、エンタープライズ サポートを希望される場合は、NVIDIA AI Enterprise ライセンスの購入をご検討ください。

LLM のワークフロー

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

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

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

  • NeMo Curator
    LLM の学習に必要な大規模データセットのダウンロードから抽出、クリーニング、フィルタリングなどを行うためのスケーラブルなツールキット。
  • NeMo
    LLM、マルチモーダル、音声などの生成AIモデルを構築するためのスケーラブルなフレームワーク。
  • NeMo-Framework-Launcher
    クラウド、オンプレのクラスターからジョブを起動するためのツールキット。
  • Megatron-LM
    Transformer モデルの大規模学習に関する研究プロジェクト。
  • TransformerEngine
    FP8 を中心とした Transformer モデルを高速化させるツールキット。
  • NeMo-Aligner
    人間のフィードバックからの強化学習 (RLHF) 、DPO、SteerLMなどを使用して LLM を効率的にアライメントするためのツールキット。

これらのライブラリは、GitHub 上に OpenSource として公開されていますが、依存関係が解消されている NeMo Framework コンテナーから利用することをお薦めします。コンテナーの場合、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 コンテナーを入手します (このチュートリアルでは、nvcr.io/nvidia/nemo:24.03.01.framework を使用しています)。
  • 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/nvidia/nemo:24.03.01.framework 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 の Llama2 モデルを nemo フォーマットへ変換します。

python /opt/NeMo/scripts/checkpoint_converters/convert_llama_hf_to_nemo.py --input_name_or_path=./models/ELYZA-japanese-Llama-2-7b --output_path=./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_finetuning.py で実行できます。
/opt/NeMo/examples/nlp/language_modeling/tuning/conf/megatron_gpt_finetuning_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_finetuning.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_generate.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_generate.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 領域のソリューション開発に取り組んでいます。