XTunerを用いたQLoRAによるパーソナルアシスタントLLMの知識注入手法

XTunerフレームワークを活用し、軽量大言語モデル(InternLM2-Chat-1.8B)をパーソナルアシスタントとしてカスタマイズするプロセスを解説します。QLoRA方式を採用することで、限定的なGPUメモリ環境でも効率的なパラメータ微調整が可能であり、モデルの自己認識や特定役割への適応を確実に実現できます。

環境構築と依存関係の準備

微調整作業を開始する前に、専用仮想環境の作成とXTunerのインストールを実行します。


# 仮想環境の初期化(Python 3.10推奨)
conda create -n llm_ft_env python=3.10 -y
conda activate llm_ft_env

# ワークスペースの作成と移動
mkdir -p ~/workspace/xtuner_tutorial && cd ~/workspace/xtuner_tutorial

# 指定バージョンのソース取得
git clone -b v0.1.17 https://github.com/InternLM/xtuner.git
cd xtuner

# 開発モードでのパッケージインストール
pip install -e '.[all]'

学習用データセットの生成

モデルが特定の役柄や自己紹介を正しく出力できるよう、OpenAI準拠のJSON形式で対話データを作成します。以下のスクリプトは、指定された名前と役割に基づいて学習サンプルを自動生成します。


import json
import os

def build_conversation_dataset(target_name: str, sample_count: int, output_path: str):
    """
    パーソナルアシスタントの自己認識データセットを生成する。
    """
    prompt_input = "自己紹介をお願いします"
    response_output = f"私は{target_name}専用に設計されたAIアシスタントです。基盤モデルは上海AI研究所開発の1.8B規模モデルを採用しています。"

    # 対話データのテンプレート
    dialogue_pair = {
        "messages": [
            {"role": "user", "content": prompt_input},
            {"role": "assistant", "content": response_output}
        ]
    }

    # 指定回数分データを複製しリスト化
    dataset_entries = [dialogue_pair for _ in range(sample_count)]

    # JSONファイルへの書き出し
    with open(output_path, "w", encoding="utf-8") as file:
        json.dump(dataset_entries, file, ensure_ascii=False, indent=4)
    print(f"データセット生成完了: {output_path}")

if __name__ == "__main__":
    data_dir = os.path.join(os.getcwd(), "custom_data")
    os.makedirs(data_dir, exist_ok=True)
    build_conversation_dataset(
        target_name="田中課長",
        sample_count=8000,
        output_path=os.path.join(data_dir, "assistant_profile.json")
    )

上記スクリプトを実行後、指定パスにJSONファイルが出力されることを確認します。

基盤モデルと設定ファイルの選定

学習対象としてInternLM2-1.8Bを使用します。モデルファイルをローカルに配置し、XTunerのビルトイン設定から適切なプロファイルを選択します。


mkdir -p ~/workspace/xtuner_tutorial/base_model
cp -r /path/to/internlm2-chat-1_8b/* ~/workspace/xtuner_tutorial/base_model/

# 利用可能な設定ファイルの一覧表示
xtuner list-cfg -p internlm2_1_8b

QLoRA方式に適した internlm2_1_8b_qlora_alpaca_e3 をコピーしてカスタマイズします。


mkdir -p ~/workspace/xtuner_tutorial/configs
xtuner copy-cfg internlm2_1_8b_qlora_alpaca_e3 ~/workspace/xtuner_tutorial/configs/

設定ファイルのカスタマイズ

コピーされた internlm2_1_8b_qlora_alpaca_e3_copy.py を編集し、以下のパスとハイパーパラメータを更新します。

  • pretrained_model_name_or_path: ローカルの基盤モデルパスへ変更
  • alpaca_en_path: 生成した assistant_profile.json のパスへ変更
  • dataset_map_fn: openai_map_fn に切り替え
  • evaluation_inputs: 評価用プロンプト ["自己紹介をお願いします", "あなたの正体は何ですか"] に更新
  • max_epochs および save_steps: 過学習を防ぐため適宜調整(例: epochs=2, save_steps=400)

データ品質は微調整の成果を左右する最も重要な要素です。高品質なペアリングデータと適切な学習エポック数の管理が、モデルの汎用性を維持しつつ特定知識を注入する鍵となります。

トレーニング実行と最適化

設定が完了次第、以下のコマンドで学習を開始します。


xtuner train ~/workspace/xtuner_tutorial/configs/internlm2_1_8b_qlora_alpaca_e3_copy.py \
    --work-dir ~/workspace/xtuner_tutorial/outputs

GPUメモリが逼迫する場合は、DeepSpeed ZeROオプションを併用することで分散最適化が可能です。


xtuner train ~/workspace/xtuner_tutorial/configs/internlm2_1_8b_qlora_alpaca_e3_copy.py \
    --work-dir ~/workspace/xtuner_tutorial/outputs_ds \
    --deepspeed deepspeed_zero2

ZeRO Stage 2はOptimizer状態と勾配を分割し、メモリの効率化を図ります。大規模モデルやリソース制約下のトレーニングに有効です。

過学習の対策と学習の再開

単一の自己紹介データで過剰に学習すると、モデルの応答が固定され汎用対話能力が損なわれる可能性があります。これを避けるためには、保存間隔を短くしてチェックポイントを頻繁に記録し、評価結果から最適なエポックを抽出します。また、汎用会話データを混合して学習させることも有効な手法です。

予期せぬ中断が発生した場合、--resume フラグと対象のチェックポイントファイルを指定することで、学習を途中から再開できます。


xtuner train ~/workspace/xtuner_tutorial/configs/internlm2_1_8b_qlora_alpaca_e3_copy.py \
    --work-dir ~/workspace/xtuner_tutorial/outputs \
    --resume ~/workspace/xtuner_tutorial/outputs/iter_400.pth

モデルの変換と統合

学習が完了したら、PyTorch形式のAdapter権重をHuggingFace互換形式へ変換し、基盤モデルとマージします。


# Adapterの形式変換
xtuner convert pth_to_hf \
    ~/workspace/xtuner_tutorial/configs/internlm2_1_8b_qlora_alpaca_e3_copy.py \
    ~/workspace/xtuner_tutorial/outputs/iter_768.pth \
    ~/workspace/xtuner_tutorial/adapter_hf

# 基盤モデルとの統合
mkdir -p ~/workspace/xtuner_tutorial/merged_model
export MKL_SERVICE_FORCE_INTEL=1
xtuner convert merge \
    ~/workspace/xtuner_tutorial/base_model \
    ~/workspace/xtuner_tutorial/adapter_hf \
    ~/workspace/xtuner_tutorial/merged_model

対話テストとWebUIデプロイ

統合済みのモデルは xtuner chat コマンドで即座に検証可能です。


xtuner chat ~/workspace/xtuner_tutorial/merged_model --prompt-template internlm2_chat

ブラウザ経由での対話環境を構築するには、Streamlitベースのデモスクリプトを修正・実行します。


pip install streamlit==1.24.0
# WebUIスクリプトの準備と実行(ポートフォワーディング設定後)
streamlit run web_demo.py --server.address 0.0.0.0 --server.port 7860

Streamlitのフロントエンドロジックはセッション状態管理と逐次トークン生成を組み合わせることで実装されます。以下に簡略化されたモデル読み込みと対話ループの構成を示します。


import streamlit as st
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

@st.cache_resource
def init_inference_engine(model_dir: str):
    """モデルとトークナイザーの初期化とGPU配置"""
    model = AutoModelForCausalLM.from_pretrained(
        model_dir, trust_remote_code=True, torch_dtype=torch.bfloat16
    ).cuda()
    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    return model, tokenizer

def run_chat_loop():
    st.title("カスタムAIアシスタント対話")
    model, tokenizer = init_inference_engine("~/workspace/xtuner_tutorial/merged_model")

    if "history" not in st.session_state:
        st.session_state.history = []

    for entry in st.session_state.history:
        with st.chat_message(entry["sender"]):
            st.markdown(entry["text"])

    if user_input := st.chat_input("メッセージを入力"):
        with st.chat_message("user"):
            st.markdown(user_input)
        st.session_state.history.append({"sender": "user", "text": user_input})

        # プロンプト構築と逐次生成
        full_prompt = tokenizer.apply_chat_template(
            st.session_state.history, tokenize=False, add_generation_prompt=True
        )
        input_ids = tokenizer.encode(full_prompt, return_tensors="pt").cuda()

        with st.chat_message("assistant"):
            msg_placeholder = st.empty()
            generated_text = ""
            outputs = model.generate(
                input_ids, max_new_tokens=512, do_sample=True, temperature=0.7
            )
            generated_text = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
            msg_placeholder.markdown(generated_text)
            st.session_state.history.append({"sender": "assistant", "text": generated_text})

if __name__ == "__main__":
    run_chat_loop()

タグ: XTuner QLoRA InternLM2 Streamlit PyTorch

6月21日 01:09 投稿