Related Resources
Generative AI

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

2024 年 7 月 3 日
By
Reading Time: 4 minutes

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

NeMo Framework とは

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

NeMo Framework は、NGC 上に公開されているコンテナーを無償利用していただくこともできますが、NVIDIA AI Enterprise のサポート対象ソフトウェアとなっています。エンタープライズ サポートを希望される場合は、NVIDIA AI Enterprise ライセンスの購入をご検討ください。

LLM のワークフロー

図1. LLM の開発ワークフロー

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

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

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

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

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

SFT とは

SFT (Supervised Fine-Tuning) とは、入力と出力のラベル付きデータを使用してモデルのすべてのパラメーターをファインチューニングするプロセスで、ドメイン固有の用語とユーザー指定の指示に従う方法をモデルに教えます。

SFT のフォーマットには、タスクに関する指示が記述されたデータセットが使用されるため、指示チューニングとも呼ばれます。指示には、「次の記事を 3 つの文章に要約してください」や「次の学園祭についてスペイン語でメールを書いてください」など、自然言語で記述される様々な指示が使用されます。

指示チューニングのプロセスでは、様々な指示で構成されたデータセットを使って、事前学習済みモデルのファインチューニングを実行します。ファインチューニングされたモデルは推論時に未知のタスクで評価され、指示チューニングなしの場合と比較して、未知タスクでのゼロショット パフォーマンスが大幅に向上することが知られています。

SFT は、強化学習を使用して LLM の能力を向上させるプロセスにおける重要な中間ステップでもありますし、さらに後続のステップとして、PEFT (Parameter-Efficient Fine-Tuning) と組み合わせることも可能です。「大規模言語モデルのカスタマイズ手法を選択する」には、SFT を含むカスタマイズ手法の解説があります。

SFT チュートリアル

本記事では、Hugging Face Model Hub から mistralai/Mistral-7B-v0.1 ベースの日本語 LLM tokyotech-llm/Swallow-MS-7b-v0.1 をダウンロードして、NeMo Framework を使用した SFT を実行します。

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

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

また、今回のチュートリアルの検証環境は以下の条件で行っております。

  • ハードウェア
    • DGX Cloud A100
    • GPU: 8 x NVIDIA A100 80 GB GPUs (driver version: 535.104.12)
    • CPU: AMD EPYC 7V12 64-Core Processor
    • システム メモリ: 1 TB
  • ソフトウェア
    • OS: Ubuntu 22.04.3 LTS
    • Container: nvcr.io/nvidia/nemo:24.05

Mistral 7B で SFT を実行する際の HW 要件はこちらに記載があります。

事前準備

以下のコマンドで作業用のディレクトリを作成し、移動します。

mkdir sft-example
cd sft-example

Docker コンテナーの起動

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

sudo docker run --rm -it --gpus all --shm-size=2g --ulimit memlock=-1 --network=host -v ${PWD}:/workspace -w /workspace  nvcr.io/nvidia/nemo:24.05 bash

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

このチュートリアルでは、tokyotech-llm/Swallow-MS-7b-v0.1 を使用します。以下のコードで Hugging FaceのModel Hub からモデルをダウンロードします。

import os
from huggingface_hub import snapshot_download
 
MODEL_DIR = "./models"
os.makedirs(MODEL_DIR, exist_ok=True)
 
snapshot_download(
    repo_id="tokyotech-llm/Swallow-MS-7b-v0.1",
    local_dir=f"{MODEL_DIR}/Swallow-MS-7b-v0.1",
    )

nemo フォーマットへの変換

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

python /opt/NeMo/scripts/checkpoint_converters/convert_mistral_7b_hf_to_nemo.py --input_name_or_path=./models/Swallow-MS-7b-v0.1 --output_path=./models/Swallow-MS-7b-v0.1/Swallow-MS-7b-v0.1.nemo --precision=bf16

生成された Swallow-MS-7b-v0.1.nemo ファイルは、distributed checkpoint が使用されているため、Swallow-MS-7b-v0.1.nemo の checkpoint を都度変更することなく、任意の Tensor Parallel (TP) や Pipeline Parallel (PP) の組み合わせでロードすることができます。

データの準備

この SFT チュートリアルでは、様々な指示と応答で構成されたオープンソースのデータセットである、databricks/databricks-dolly-15k とその日本語翻訳バージョンである llm-jp/databricks-dolly-15k-ja を使用します。

import os
from datasets import load_dataset

DATA_DIR = "./data/databricks-dolly-15k-ja" 
os.makedirs(DATA_DIR, exist_ok=True)
dataset = load_dataset("llm-jp/databricks-dolly-15k-ja")
dataset["train"].to_json(f"{DATA_DIR}/dolly-ja.json", force_ascii=False)

DATA_DIR = "./data/databricks-dolly-15k"
os.makedirs(DATA_DIR, exist_ok=True)
dataset = load_dataset("databricks/databricks-dolly-15k")
dataset["train"].to_json(f"{DATA_DIR}/dolly.json", force_ascii=False)

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

from glob import glob
import json
import os
import random
import pandas as pd


INPUT_PATH = [
    "./data/databricks-dolly-15k-ja/dolly-ja.json",
    "./data/databricks-dolly-15k/dolly.json",
    ]
OUTPUT_PATH = "./data/sft"
USE_COLS = ["context", "instruction", "response", "source"]


random.seed(42)
os.makedirs(OUTPUT_PATH, exist_ok=True)

# prompt templates
INPUT_PROMPT = """以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。

### 指示:
{instruction}

### 入力:
{input}

### 応答:
"""

NO_INPUT_PROMPT = """以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。

### 指示:
{instruction}

### 応答:
"""


def load_dolly_dataset(path, source):
    dataset = pd.read_json(path, lines=True)
    dataset["source"] = source
    return dataset[USE_COLS]


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_input(row):
    context = row["context"].strip()
    instruction = row["instruction"].strip()
    response = row["response"].strip()
    assert instruction != ""

    if context != "":
        input = INPUT_PROMPT.format(instruction=instruction, input=context)
    else:
        input = NO_INPUT_PROMPT.format(instruction=instruction)

    return input, response, row["source"]


def prosess(input_path):

    processed = []
    dataset = pd.DataFrame()

    for path in input_path:
        if "dolly.json" in path:
            df = load_dolly_dataset(path, "dolly")
            print("dolly num_records: ", df.shape[0])
        elif "dolly-ja.json" in path:
            df = load_dolly_dataset(path, "dolly-ja")
            print("dolly-ja num_records: ", df.shape[0])
        else:
            print(f"Ignore...: {path}")

        dataset = pd.concat([dataset, df], ignore_index=True)

    # drop duplicated samples
    print("total records: ", dataset.shape[0])
    dataset = dataset[~dataset.duplicated(subset=["context", "instruction", "response"])].reset_index(drop=True)
    print("total records(drop duplicated samples): ", dataset.shape[0])

    for index, row in dataset.iterrows():
        input, output, source = form_input(row)
        processed.append({"input": input, "output": output, "source": source})

    random.shuffle(processed)
    train_ds = processed[:int(len(processed)*0.9)]
    valid_ds = processed[int(len(processed)*0.9):]

    write_jsonl(f"{OUTPUT_PATH}/train.jsonl", train_ds)
    write_jsonl(f"{OUTPUT_PATH}/valid.jsonl", valid_ds)

    print("num_train: ", len(train_ds), "num_valid: ", len(valid_ds))
    print(train_ds[0]["input"])
    print(train_ds[0]["output"])
    print(train_ds[0]["source"])

    return


def main():
    prosess(INPUT_PATH)


if __name__ == "__main__":
    main()

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

{"input": "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。\n\n### 指示:\nIn 64th Annual Grammy Awards,  best album of the year award was given to\n\n### 応答:\n", "output": "Jon Batiste was awarded the best album of the year in 64th Annual Grammy Awards", "source": "dolly"}
{"input": "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。\n\n### 指示:\nアメリカで最も絵になる6つの国立公園とは？\n\n### 応答:\n", "output": "ザイオン国立公園、イエローストーン国立公園、グランドキャニオン国立公園、ヨセミテ国立公園、グレーシャー国立公園、グレートスモーキー山脈国立公園", "source": "dolly-ja"}

SFT の実行

SFT は /opt/NeMo-Aligner/examples/nlp/gpt/train_gpt_sft.py で実行できます。/opt/NeMo-Aligner/examples/nlp/gpt/conf/gpt_sft.yaml に SFT の実行に必要な config ファイルがあり、コマンド ライン引数としてスクリプトの実行に必要なデータのパス、ベースとなる LLM (nemo フォーマット) のパスなどの設定を渡します (NeMo Framework は config 設定に Hydra を使用しています)。ここで指定されていない設定については、config ファイルの設定が適用されます。実行コマンドにある ++exp_manager.checkpoint_callback_params.save_last=False は Hydra で config の追加または上書きをする際の記法になります。

export HYDRA_FULL_ERROR=1

export WANDB=False
export PJ_NAME="sft"
export EXP_DIR="/workspace/results/"${PJ_NAME}
export EXP_NAME="Swallow-MS-7b-v0.1_SFT"
export MODEL="/workspace/models/Swallow-MS-7b-v0.1/Swallow-MS-7b-v0.1.nemo"
export TP_SIZE=1
export PP_SIZE=1
export TRAIN_DS="/workspace/data/"${PJ_NAME}"/train.jsonl"
export VALID_DS="/workspace/data/"${PJ_NAME}"/valid.jsonl"


torchrun --nproc_per_node=8 /opt/NeMo-Aligner/examples/nlp/gpt/train_gpt_sft.py \
    exp_manager.exp_dir=${EXP_DIR} \
    exp_manager.name=${EXP_NAME} \
    exp_manager.create_wandb_logger=${WANDB} \
    exp_manager.wandb_logger_kwargs.project=${PJ_NAME} \
    exp_manager.wandb_logger_kwargs.name=${EXP_NAME} \
    exp_manager.checkpoint_callback_params.save_nemo_on_train_end=True \
    exp_manager.checkpoint_callback_params.save_top_k=1 \
    exp_manager.checkpoint_callback_params.save_best_model=True \
    ++exp_manager.checkpoint_callback_params.save_last=False \
    trainer.precision=bf16 \
    trainer.devices=8 \
    trainer.num_nodes=1 \
    trainer.sft.max_epochs=3 \
    trainer.sft.max_steps=-1 \
    trainer.sft.val_check_interval=105 \
    trainer.sft.gradient_clip_val=1.0 \
    model.megatron_amp_O2=True \
    model.tensor_model_parallel_size=${TP_SIZE} \
    model.pipeline_model_parallel_size=${PP_SIZE} \
    model.restore_from_path=${MODEL} \
    model.peft.peft_scheme="none" \
    model.optim.lr=1e-6 \
    model.optim.sched.warmup_steps=50 \
    model.optim.sched.constant_steps=0 \
    model.data.train_ds.file_path=${TRAIN_DS} \
    model.data.train_ds.global_batch_size=128 \
    model.data.train_ds.micro_batch_size=1 \
    model.data.train_ds.max_seq_length=2048 \
    model.data.validation_ds.file_path=${VALID_DS} \
    model.data.validation_ds.global_batch_size=128 \
    model.data.validation_ds.micro_batch_size=1 \
    model.data.validation_ds.drop_last=True \
    model.data.num_workers=0 \
    model.answer_only_loss=True

NeMo Framework は、実験管理のために Weights and Biases をサポートしており、上記のスクリプトを、export WANDB=True と変更することで wandb 上で実験管理することができます。

上記の設定では、学習は 1.5 時間ほどで完了しました。

学習が終わると results/sft/Swallow-MS-7b-v0.1_SFT/ という名前のディレクトリが作成され、中に学習時の config や log などが出力されます。また、同じディレクトリの checkpoints 内にある Swallow-MS-7b-v0.1_SFT.nemo が SFT で作成されたモデルになります。

学習時に nemo フォーマットへ変換したモデルは、同一フォーマットのまま、推論やさらなる後続ステップで活用することもできますし、Hugging Face フォーマットへ再度、変換することも可能です。

NeMo Framework コンテナー上で推論を実行するチュートリアルはこちらにあります。

nemo フォーマットから、Hugging Face フォーマットへ変換するスクリプトはこちらにあります (convert_mistral_7b_nemo_to_hf.py を実行する際に環境によっては、segmentation fault が発生することがあります)。

参考: Nejumi リーダーボード Neo での評価

Nejumiリーダーボード Neo は、LLM の日本語能力を評価できるベンチマークです。ベンチマーク用のスクリプトを実行することで、一問一答形式で言語理解を評価する llm-jp-eval と、プロンプト対話で生成能力を評価する MT-Bench による評価結果が得られます。 

ご自身の環境で実行する際は、wandb の GitHub リポジトリを参照してください。

以下は極端な一例ですが、SFT を実行する前 (事前学習/追加学習後) のモデルの応答を示します。

# Question
次のカテゴリーがあります - 文学、歴史、科学、芸術。以下の質問を分析し、それらをこれらのカテゴリーのいずれかに割り当ててください。応答では、余分な言葉をしないでください。1行につき1つのトピックを厳密にリストアップしてください。
1. 三島由紀夫の「金閣寺」に見られる主要なテーマと手法を議論してください。それらはどのように20世紀の日本の社会的状況と一致しますか？
2. 戦国時代の各大名が採用した地政学戦略と国内政策を分析してください。これらの行動はどのようにして戦後の国際秩序を形成しましたか？
3. 水のルイス構造を描き、その極性の性質を説明してください。これが沸点が高く、多くの物質を溶かす能力などのユニークな特性にどのように影響を与えるかを説明してください。
4. 鳥居清長の「浮世絵」に見られる芸術的技法とスタイル選択を批判的に検討してください。この絵画はどのように江戸時代の文化と哲学的環境を反映していますか？

# Answer
1. 三島由紀夫の「金閣寺」に見られる主要なテーマと手法を議論してください。それらはどのように20世紀の日本の社会的状況と一致しますか?
文学: 三島由紀夫の「金閣寺」は、日本の歴史と文化における重要なテーマを探求しています。この小説は、1950年代の日本の社会的状況と一致し、その時代の社会的、政治的、文化的問題を反映しています。
歴史: この小説は、1950年代の日本の社会的状況と一致し、その時代の社会的、政治的、文化的問題を反映しています。
科学: この小説は、1950年代の日本の社会的状況と一致し、その時代の社会的、政治的、文化的問題を反映しています。
芸術: この小説は、1950年代の日本の社会的状況と一致し、その時代の社会的、政治的、文化的問題を反映しています。
2. 戦国時代の各大名が採用した地政学戦略と国内政策を分析してください。これらの行動はどのようにして戦後の国際秩序を形成しましたか?
文学: 戦国時代の各大名が採用した地政学戦略と国内政策を分析してください。これらの行動はどのようにして戦後の国際秩序を形成しましたか?
歴史: 戦国時代の各大名が採用した地政学戦略と国内政策を分析してください。これらの行動はどのようにして戦後の国際秩序を形成しましたか?
科学: 戦国時代の各大名が採用した地政学戦略と国内政策を分析してください。これらの行動はどのようにして戦後の国際秩序を形成しましたか?
芸術: 戦国時代の各大名が採用した地政学戦略と国内政策を分析してください。これらの行動はどのようにして戦後の国際秩序を形成しましたか?
3. 水のルイス構造を描き、その極性の性質を説明してください。これが沸点が高く、多くの物質を溶かす能力などのユニークな特性にどのように影響を与えるかを説明してください。
文学: 水のルイス構造を描き、その極性の性質を説明してください。これが沸点が高く、多くの物質を溶かす能力などのユニークな特性にどのように影響を与えるかを説明してください。
歴史: 水のルイス構造を描き、その極性の性質を説明してください。これが沸点が高く、多くの物質を溶かす能力などのユニークな特性にどのように影響を与えるかを説明してください。
科学: 水のルイス構造を描き、その極性の性質を説明してください。これが沸点が高く、多くの物質を溶かす能力などのユニークな特性にどのように影響を与えるかを説明してください。
芸術: 水のルイス構造を描き、その極性の性質を説明してください。これが沸点が高く、多くの物質を溶かす能力などのユニークな特性にどのように影響を与えるかを説明してください。
4. 鳥居清長の「浮世絵」に見られる芸術的技法とスタイル選択を批判的に検討してください。この絵画はどのように江戸時代の文化と哲学的環境を反映していますか?
文学: 鳥居清長の「浮世絵」に見られる芸術的技法とスタイル選択を批判的に検討してください。この絵画はどのように江戸時代の文化と哲学的環境を反映していますか?
歴史: 鳥居清長の「浮世絵」に見られる芸術的技法とスタイル選択を批判的に検討してください。この絵画はどのように江戸時代の文化と哲学的環境を反映していますか?
科学: 鳥居清長の「浮世絵」に見られる芸術的技法とスタイル選択を批判的に検討してください。この絵画はどのように江戸時代の文化と哲学的環境を反映していますか?
芸術: 鳥居清長の「浮世絵」に見られる芸術的技法とスタイル選択を批判的に検討してください。この絵画はどのように江戸時代の文化と哲学的環境を反映していますか?

次に SFT 後のモデルの応答を示します。

# Question
次のカテゴリーがあります - 文学、歴史、科学、芸術。以下の質問を分析し、それらをこれらのカテゴリーのいずれかに割り当ててください。応答では、余分な言葉をしないでください。1行につき1つのトピックを厳密にリストアップしてください。
1. 三島由紀夫の「金閣寺」に見られる主要なテーマと手法を議論してください。それらはどのように20世紀の日本の社会的状況と一致しますか？
2. 戦国時代の各大名が採用した地政学戦略と国内政策を分析してください。これらの行動はどのようにして戦後の国際秩序を形成しましたか？
3. 水のルイス構造を描き、その極性の性質を説明してください。これが沸点が高く、多くの物質を溶かす能力などのユニークな特性にどのように影響を与えるかを説明してください。
4. 鳥居清長の「浮世絵」に見られる芸術的技法とスタイル選択を批判的に検討してください。この絵画はどのように江戸時代の文化と哲学的環境を反映していますか？

# Answer
1.文学、2.歴史、3.科学、4.芸術

SFT 前に比べると、SFT 後のモデルの応答は、「質問をいずれかのカテゴリに割り当て、余分な言葉を使わず、1 行につき 1 つのトピックを厳密にリストアップ」といった指示により従うようになっていることが確認できます。

まとめ

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

関連情報

Tags

Generative AI | General | NeMo | Tutorial | Japanese Tutorial | LLM Techniques | LLMs | NeMo Framework

About the Authors

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