大規模言語モデル(Llama2など)を正攻法でファインチューニングする際のメモ(ZeRO-Offload, not QLoRA)

背景と目的

大きめのサイズ(>数b)の大規模言語(LLM)をファインチューニングします｡
ファインチューニングにはLoRAやQLoRAと呼ばれる手法が良く使われ､一般家庭レベル(?)のGPUでも動かせるようになってきています｡
しかし､LoRAで学習させられる知識や情報には､制約があるのでは､とも囁かれています｡

Fine-Tuning LLMs: LoRA or Full-Parameter? An in-depth Analysis with Llama 2

そこで､本記事は､loraではないフルパラメータのファインチューニングを､限られたGPUメモリで行います｡

deepspeedというライブラリを使います｡
deepspeedにはモデルの動作に必要なメモリをCPUメモリに移す機能などがあるようで､それを使います(キーワード: offload, ZeRO)｡

7bモデルは20GB程度のVRAMで学習できました｡

以下の公式チュートリアルをもとに進めたいところですが､情報が断片的で､自分にはあまり理解できなかったので､webサイトを適当に探りながら進めました｡

ZeRO-Offload

環境構築

以下のような感じで構築しました(メモ書き)｡

pip install deepspeed
pip install transformers datasets mecab-python3 unidic-lite sentencepiece accelerate pynvml deepspeed
pip install protobuf
conda install mpi4py #pipでエラーが出たので｡

推論テスト

まずは推論で､deepspeedの効果を試します｡
モデルは "elyza/ELYZA-japanese-Llama-2-7b-instruct"を使います｡
ELYZAは､llama2をベースにファインチューニングした､わりと最近のモデルです｡

ELYZAが公開した日本語LLM「ELYZA-japanese-Llama-2-7b」についての解説 : (1) 事前学習編

対照実験: 普通に推論

transformersでモデルを呼び出して､推論をしてみます｡
notebook環境で動かしています｡

#普通に推論
from transformers import pipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import torch
model_name= "elyza/ELYZA-japanese-Llama-2-7b-instruct"

model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)


if torch.cuda.is_available():
  model.to("cuda")


def gen(input_text="今日の天気は"):
    text_pipe = pipeline('text-generation', 
                         model=model,
                         tokenizer=tokenizer,
                         device="cuda:0",
                         max_length=200,
                             )
    output = text_pipe(input_text)

    return output[0]['generated_text']

%time gen()

出力:
'今日の天気は曇り時々雨。\n気温は15℃。\n今日は朝から雨が降っています。\n降り方が強くて、洗濯物が干せそうにありません。\n昨日は、午後から雨が降っていたので、洗濯物は干せましたが、今日は干せそうにありません。\n洗濯機を回して、洗濯物を部屋に干しておきます。\n今日の天気は曇り時々雨。\n気温は'

使用VRAM: 約28GB
推論時間: 5 sec
CPUメモリ: VIRTで約60GB､RESで2.5GB

7Bのモデルを32bitの変数で動かしているので､7x4=28 GBという計算だと思います｡

DeepSpeedの利用

メモリをCPUにオフロードするなどできるようです｡詳細は勉強中です｡
以下の記事を参考に､推論のテストをしてみました｡

DeepSpeedのZeRO-Inferenceを使ってV100-16GBの環境で30BのLLMを動かしてみました。 - CCCMKホールディングス TECH Labの Tech Blog

以下のjsonファイル(zero_infer.json)を作ります｡

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

実行処理

import torch
import deepspeed
from transformers.deepspeed import HfDeepSpeedConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import json
import os


model_name= "elyza/ELYZA-japanese-Llama-2-7b-instruct"


# multi-GPU関連の設定
os.environ["TOKENIZERS_PARALLELISM"] = "false"  # To avoid warnings about parallelism in tokenizers
local_rank = int(os.getenv("LOCAL_RANK",0))
world_size = int(os.getenv("WORLD_SIZE",1))

torch.cuda.set_device(local_rank)
deepspeed.init_distributed()

# ベースとなるZeRO3 configの読み込み
ds_config_file = "zero_infer.json"
with open(ds_config_file) as f:
    ds_config = json.load(f)

# 推論用に修正
model_config = AutoConfig.from_pretrained(model_name)
hidden_size = model_config.hidden_size

ds_config["train_batch_size"] = 1 * world_size
ds_config["train_micro_batch_size_per_gpu"] = 1
ds_config["reduce_bucket_size"] = hidden_size*hidden_size
ds_config["stage3_prefetch_bucket_size"] = 0.9 * hidden_size * hidden_size
ds_config["stage3_param_persistence_threshold"] = 10 * hidden_size


dschf = HfDeepSpeedConfig(ds_config)  #zero3を使用するために必要(モデルロード前に実行する必要がある)
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
ds_model = ds_engine.module.eval()


from transformers import pipeline


if torch.cuda.is_available():
  model.to("cuda")


def gen(input_text="今日の天気は"):
    text_pipe = pipeline('text-generation', 
                         model=ds_model,
                         tokenizer=tokenizer,
                         device="cuda:0",
                         max_length=200,
                             )
    output = text_pipe(input_text)

    return output[0]['generated_text']
%time gen()

出力:
'今日の天気は曇り時々雨。\n気温は15℃。\n今日は朝から雨が降っています。\n降り方が強くて、洗濯物が干せそうにありません。\n昨日は、午後から雨が降っていたので、洗濯物は干せましたが、今日は干せそうにありません。\n洗濯機を回して、洗濯物を部屋に干しておきます。\n今日の天気は曇り時々雨。\n気温は'

使用VRAM: 約4.8GB
推論時間: 1min 51 sec
CPUメモリ: VIRTで約83GB､RESで52GB

CPU側にメモリを移すことで､VRAMは大幅に節約できたようです｡
(加えて､16 bit推論をしているので､必要なメモリサイズが半減した効果もあります)

VRAMは節約できた一方で､推論時間が20倍以上になってしまいました｡

訓練

学習用に､少し追記した設定ファイルを作ります(zero_train.json)｡

{
    "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": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

モデルの定義類は､推論と基本的に同じです｡

import torch
import deepspeed
from transformers.deepspeed import HfDeepSpeedConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import json
import os

model_name= "elyza/ELYZA-japanese-Llama-2-7b-instruct"


# multi-GPU関連の設定
os.environ["TOKENIZERS_PARALLELISM"] = "false"  # To avoid warnings about parallelism in tokenizers
local_rank = int(os.getenv("LOCAL_RANK",0))
world_size = int(os.getenv("WORLD_SIZE",1))

torch.cuda.set_device(local_rank)
deepspeed.init_distributed()

# ベースとなるZeRO3 configの読み込み
ds_config_file = "zero_infer.json"
with open(ds_config_file) as f:
    ds_config = json.load(f)


model_config = AutoConfig.from_pretrained(model_name)
hidden_size = model_config.hidden_size

ds_config["train_batch_size"] = 1 * world_size
ds_config["train_micro_batch_size_per_gpu"] = 1
ds_config["reduce_bucket_size"] = hidden_size*hidden_size
ds_config["stage3_prefetch_bucket_size"] = 0.9 * hidden_size * hidden_size
ds_config["stage3_param_persistence_threshold"] = 10 * hidden_size


dschf = HfDeepSpeedConfig(ds_config)  #zero3を使用するために必要(モデルロード前に実行する必要がある)
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
ds_model = ds_engine.module#.eval()

データセットの定義など
train_data_pathに､適当なテキストデータへのパスを入れて下さい｡

from transformers import (AutoModelForCausalLM,
                          DataCollatorForLanguageModeling, T5Tokenizer,AutoTokenizer,
                          TextDataset, Trainer, TrainingArguments)

train_data_path = "ここにデータセットのパスを指定"

train_dataset = TextDataset(
    tokenizer=tokenizer,
    file_path=train_data_path,
    block_size=1024, #文章の長さを揃える,
    #block_size=100, #文章の長さを揃える,
    cache_dir="cache/"+model_name,
)

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False
)

#deepspeed
epochs=1

training_args = TrainingArguments(
    output_dir='./outputs',
    num_train_epochs=epochs,  # エポック数
    #per_device_train_batch_size=1,  # バッチサイズ
    gradient_accumulation_steps=1,
    gradient_checkpointing=True,  #勾配チェックポイント
    fp16=True,  #fp16
    optim='adafactor',  # オプティマイザの種類
    deepspeed='./zero_train.json',  # deepspeedのconfigへのpath
    logging_steps=100,  # 途中経過を表示する間隔
)

trainer = Trainer(
    model=ds_model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_dataset
)

#訓練
result = trainer.train()
print_summary(result)

使用VRAM: 約19GB
CPUメモリ: VIRTで約206GB､RESで150GB

学習時間はデータサイズ次第ですが､4MBのテキストで1 hr程度でした｡
offloadを使わない場合と同等レベルな気がしたので､驚きました※｡
(※ このモデルは今回使った環境では､offloadを使わないと､out of memoryで学習できませんでした｡より小さなモデルの学習時間との比較です)

学習したモデルの保存と読み込み

0922追記 モデルの保存コマンドに問題があったので修正(model.save_pretrainedではなく､trainer.save_modelを使う)｡

# Save the model and the tokenizer.
dir_name = f'finetuned/{model_name}_{epochs}'

#0922追記: model.save_pretrainedだと､うまくモデルを読み込めない感じでした
#model.save_pretrained(dir_name)
trainer.save_model(dir_name)
tokenizer.save_pretrained(dir_name)

保存したモデルのパスを､推論用コード(上述)のmodel_nameに指定してあげれば､ファインチューニングしたモデルを呼び出せます｡

#普通に推論
from transformers import pipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

#ファインチューニングモデルの保存パスを指定
model_name="finetuned/elyza/ELYZA-japanese-Llama-2-7b-instruct_1_test"

model = AutoModelForCausalLM.from_pretrained(model_name,ignore_mismatched_sizes=True)
tokenizer = AutoTokenizer.from_pretrained(model_name)


if torch.cuda.is_available():
  model.to("cuda")


def gen(input_text="今日の天気は"):
    text_pipe = pipeline('text-generation', 
                         model=model,
                         tokenizer=tokenizer,
                         device="cuda:0",
                         max_length=200,
                        #temperature = 100,
                             )
    output = text_pipe(input_text)

    return output[0]['generated_text']
gen()

deepspeedはマルチGPUにも対応しているので､この調子でいけば､13b以上のモデルも動かせるかもしれません｡
速度は落ちそうですが､CPUではなく､SSDのメモリ(NVME)にoffloadすることもできるようです(未検証: リンク)｡

引き続き検討します｡