Zenn
株式会社ElithPublicationへの投稿
🐈

HuggingFaceとDeepSpeedで実践継続事前学習

もっさん
2024/03/25に公開
Hugging Face
#
LLM
#
大規模言語モデル
#
deepspeed
#
継続事前学習
tech

はじめに

株式会社Elithの大森一祥です。AIテックカンパニーの一員として、お客様の課題をAIを駆使して解決しています。

大規模言語モデル(LLM)が人間と匹敵する性能を発揮することもあり、弊社には多岐にわたるプロジェクトの依頼が寄せられています。最近は、情報漏洩のリスクを回避するため、独自のLLMの開発を希望されることが多いです。このような案件では、一般に公開されたモデル(ローカルLLM)を利用します。

ローカルLLMを活用して課題を解決する方法として、以下の4つが挙げられます。

  1. プロンプトエンジニアリング：LLMに特定の出力を生成させるための入力文の工夫する手法
  2. RAG：外部の文章データベースから、質問に類似した文章を取り出しLLMの入力として用いる手法
  3. インストラクションチューニング：ユーザの指示に沿った出力を生成することを目的としたチューニング手法
  4. 継続事前学習：LLMモデルに対して追加で事前学習する手法

プロンプトエンジニアリング、RAG、インストラクションチューニングはさまざまな記事で紹介されていますが、継続事前学習の記事は少なく、動かすことが難しいという課題感があります。そこで、本記事では以下の流れでローカルLLMを継続事前学習する目的と方法について解説します。

  1. 継続事前学習の目的
  2. 環境構築
  3. 学習コード
  4. 実行

本記事で利用するスクリプトは以下のGitHubリポジトリで管理しています。

https://github.com/oriki101/continual-pretrain

継続事前学習の目的

LLM（Large Language Models）の学習プロセスは、大きく分けて事前学習と事後学習の二段階に分かれます。事前学習段階では、大量のテキストデータを用いてモデルの基本的な言語理解能力を育成します。その後の事後学習では、特定のタスクを達成するために、既に事前学習を受けたモデルをさらに訓練します。

しかし、事前学習でカバーしきれなかった領域については、事後学習だけでは十分な性能を達成するのが難しい場合があります。そのため、「継続事前学習」という方法が採用されることがあります。これは、既に事前学習されたモデルに対して、新たなドメイン特化データを用いて再び事前学習を行うことにより、モデルが当初持っていなかった知識を獲得させる手法です。

例えば、Elyzaは、Llama 2モデルを日本語データで継続事前学習し、モデルが日本語の理解能力を獲得することに成功しました。このような成果に基づき、今後は特定の分野や課題に対応するために、継続事前学習を一般的な手法として採用することが期待されます。

環境構築

LLMの継続事前学習を実施する環境は、さまざまな方法で構築可能です。今回使用する学習コードは、以前に記述した記事で紹介した環境で実行可能です。
https://zenn.dev/elith/articles/e4dbbb62752e04

学習コード

学習コードは以下に示します。Hugging FaceのTrainerクラスを利用することにより、わずか約100行のコードでDeepSpeedを活用した継続事前学習が可能となります。複雑な設定は不要で、簡単にマルチGPU環境での学習を実施できます。

import argparse
import os
import warnings
from typing import Dict, List

warnings.filterwarnings("ignore")

import deepspeed
import torch
from datasets import Dataset, load_dataset
from omegaconf import OmegaConf
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    PreTrainedTokenizer,
    Trainer,
    TrainingArguments,
)

from utils import seed_everything #seed値を固定するための関数 github上のutils.pyを確認

os.environ["TOKENIZERS_PARALLELISM"] = "false"  # 分散処理させるときにwarningが出るため

def preprocess_function(
    examples: Dict[str, List[str]], tokenizer: PreTrainedTokenizer, max_length: int
) -> Dict[str, List[int]]:
    """
    与えられたテキストをトークナイズし、ラベルを追加する前処理関数。
    ラベルを追加する理由は、AutoModelForCausalLMの入力を'label'としなければならないため

    Args:
        examples (Dict[str, List[str]]): 前処理するテキストの例。キーはテキストのフィールド名（例えば 'text'）。
        tokenizer (PreTrainedTokenizer): 使用するトークナイザーのインスタンス。
        max_length (int): トークナイズ後の最大シーケンス長。

    Returns:
        Dict[str, List[int]]: トークナイズされたテキストと対応するラベルを含む辞書。
    """
    inputs = tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=max_length,
    )
    inputs["labels"] = inputs.input_ids.copy()
    return inputs

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--train_config",
        "-p",
        type=str,
        default="./configs/train_configs/train_base.yaml",
        help="モデルパラメータのコンフィグ。yamlファイル",
    )
    parser.add_argument("--local_rank", "-l", type=int, default=0, help="GPUのランク")
    args = parser.parse_args()
    local_rank = args.local_rank
    # コンフィグ読み込み
    config = OmegaConf.load(args.train_config)

    # distributed learning
    deepspeed.init_distributed()

    # seedの設定
    seed_everything(config.seed)

    # モデルの定義
    model = AutoModelForCausalLM.from_pretrained(
        config.model.model, torch_dtype=torch.float16, use_cache=config.model.use_cache
    )
    tokenizer = AutoTokenizer.from_pretrained(
        config.model.tokenizer,
        add_eos_token=True,  # EOSの追加を指示 defaultはFalse
    )

    # データセットの読み込み
    dataset = load_dataset(
        config.dataset.path, config.dataset.subset, split=config.dataset.split
    )
    # デモなのでデータを減らしておく
    dataset = dataset.select(range(1000))

    # データをモデルに入力できるように変換
    dataset = dataset.map(
        lambda examples: preprocess_function(
            examples, tokenizer, config.model.max_length
        ),
        batched=True,
        remove_columns=dataset.column_names,
    )

    dataset = dataset.train_test_split(test_size=0.2)

    # 学習
    training_args = TrainingArguments(**config.train)
    trainer = Trainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        args=training_args,
        # data_collator=data_collator,
    )

    with torch.autocast("cuda"):
        trainer.train()

if __name__ == "__main__":
    main()

コンフィグ

コンフィグファイルでは、使用するモデル名、学習時の設定、そしてデータセットを定義します。モデルの変更などが必要な場合は、コンフィグファイルの内容を更新するだけで対応可能です。

model:
  model: elyza/ELYZA-japanese-Llama-2-7b
  tokenizer: elyza/ELYZA-japanese-Llama-2-7b
  use_cache: False
  max_length: 512

train: # huggingfaceのTrainingArgumentsで利用
  output_dir: ../outputs
  evaluation_strategy: steps
  logging_strategy: steps
  save_strategy: steps
  learning_rate: 1e-6
  num_train_epochs: 3
  per_device_train_batch_size: 3
  per_device_eval_batch_size: 3
  gradient_accumulation_steps: 2
  gradient_checkpointing: True
  weight_decay: 0.01 # 適当
  warmup_ratio: 0.1 # 適当
  optim: adamw_torch # 適当
  fp16: True
  bf16: False
  dataloader_num_workers: 4
  eval_steps: 50
  save_steps: 50
  logging_steps: 50
  run_name: test # wandbのプロジェクト名
  save_total_limit: 2
  save_on_each_node: False
  neftune_noise_alpha: 5 # NEFTTune　適当
  deepspeed: ./configs/deepspeed/ds_config_zero2.json
  report_to: wandb
  
seed: 42

dataset:
  path: hotchpotch/wikipedia-ja-20231030
  subset: chunked #!!null
  split: train

DeepSpeedコンフィグ

DeepSpeedは、大規模な深層学習モデルを効率的にスケーリングし、複数のGPUまたはGPUクラスター上での学習と推論を最適化するライブラリです。このフレームワークを使用することで、GPUメモリの使用量を大幅に削減し、LLMの学習プロセスを高速化することが可能です。特に、大規模モデルの学習において、DeepSpeedは計算資源の利用効率を向上させる重要な役割を果たします。

DeepSpeedの設定例は以下です。

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },

    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

Hugging FaceのTrainerとDeepSpeed設定が競合する場合、Trainerの設定が優先されます。このため、zero_optimization以外の設定については、DeepSpeedのコンフィグファイルを編集するのではなく、Trainerのコンフィグを変更することが良いと考えられます。

Hugging Face TrainerでDeepSpeedを設定する方法について詳しく知りたい方は、以下のリンクで情報を確認してください。

https://huggingface.co/docs/transformers/v4.15.0/main_classes/deepspeed

学習

シングルノード学習

単一のPCでの学習方法について説明します。このプロセスは非常にシンプルで、deepspeedコマンドを用いてPythonスクリプトを実行することで、マルチGPUを活用した学習を開始できます。以下に、その実行手順を示します。

cd xxx
deepspeed src/train_deepspeed.py --train_config ./configs/train_configs/train_base.yaml

Docker環境を使用している場合、学習を開始する前にコンテナにアタッチすることを忘れないでください。

マルチノード学習

手元にマルチノード環境がないため、国立研究開発法人産業技術総合研究所によって構築・運用されているABCI（AI Bridging Cloud Infrastructure）を利用して解説します。DeepSpeedはデフォルトでPDSH（Parallel Distributed Shell）を使って分散学習を行いますが、ABCI環境ではSSH経由で接続したノード上でPythonが読み込めないことによりエラーが発生する場合があります。そのため、シングルノード学習のようにdeepspeedコマンドを用いるには、ソースコードの修正が必要です。しかし、この作業は環境構築の過程で大きな手間です。

そこで、Open MPIのmpirunコマンドを使用して分散学習を行う方法を採用します。これにより、複雑な設定を避けつつ、効率的なマルチノード学習が可能でし。以下に実行するシェルスクリプトを示します。

#!/bin/bash
#$ -l rt_F=2
#$ -l h_rt=0:30:00
#$ -l USE_BEEOND=1
#$ -j y
#$ -o logs/
#$ -cwd

source /etc/profile.d/modules.sh
module load python/3.10
module load cuda/11.7/11.7.1
module load cudnn/8.5/8.5.0
module load nccl/2.14/2.14.3-1
module load hpcx/2.12

source ~/llm-env/bin/activate # 環境ごとに変える

GPUS_PER_NODE=4
NNODES=2
NUM_GPUS=$((${GPUS_PER_NODE} * ${NNODES}))
echo ${NUM_GPUS}
echo ${GPUS_PER_NODE}

# マルチノード用のアドレスの設定
export MASTER_ADDR=$(/usr/sbin/ip a show dev bond0 | grep 'inet ' | awk '{ print $2 }' | cut -d "/" -f 1)
export MASTER_PORT=50000
echo ${MASTER_ADDR}
echo ${MASTER_PORT}

# hostfile作成
HOSTFILE_NAME=./hostfile/hostfile_${JOB_ID}
while read -r line
do
  echo "${line} slots=${GPUS_PER_NODE}"
done < "$SGE_JOB_HOSTLIST" > "$HOSTFILE_NAME"

# huggingfaceのモデル保存先を指定(容量対策)
export HF_HOME=/scratch/$(whoami)/.cache/huggingface/

mpirun -np ${NUM_GPUS} \
	-npernode ${GPUS_PER_NODE} \
	-hostfile $HOSTFILE_NAME \
	-x MASTER_ADDR=$MASTER_ADDR \
	-x MASTER_PORT=$MASTER_PORT \
	-bind-to none -map-by slot \
	-x NCCL_DEBUG=INFO  -x PATH \
	-mca pml ob1 -mca btl ^openib \
	-mca coll ^hcoll \
	--mca btl_tcp_if_include eno1 \
	python src/train_deepspeed.py --train_config ./configs/train_configs/train_base.yaml

上記のシェルスクリプトではV100が4枚のノードを2つ利用し、30分間学習させます。ノード数と実行時間は以下の部分を修正してください。

#$ -l rt_F=2 # 利用するノード名と数
#$ -l h_rt=0:30:00 # 実行時間

シェルスクリプトを以下のコマンドで実行すると学習ができます。

$ qsub -g {グループ名} script/incremental_pretrain_abci.sh

おわりに

本記事では、Hugging FaceのTransformersライブラリとDeepSpeedを組み合わせた大規模言語モデル（LLM）への知識追加に関する事前学習方法を紹介しました。この方法の最大の利点は、約100行のスクリプトだけで継続事前学習を開始できることにあります。これは、研究開発やプロトタイピングフェーズにおける実験サイクルを大幅に加速することを可能にします。

最後に宣伝です。株式会社 Elith は最先端のAI技術をビジネスに実装し、価値を生み出すテックカンパニーです。

最近では、医療用LLMの作成など、LLMを活用した社会課題解決に向けて作業しています。

LLMを活用したビジネス課題解決に興味がある方は、X(旧Twitter)経由やElithのWebページ経由で、是非気軽にお尋ねください。

参考

大規模言語モデル(LLM)の作り方 Megatron-DeepSpeed編 Part2

https://huggingface.co/docs/transformers/main_classes/deepspeed

https://github.com/ohtaman/abci-examples/tree/main/202310

https://huggingface.co/docs/transformers/v4.15.0/main_classes/deepspeed

株式会社Elith
株式会社ElithPublication

Elith(エリス)です！AIの受託・自社・コンサルティング・研究開発を行っています！

Discussion