メインコンテンツへスキップ
見出し画像

Google Colab で DeepSpeed によるLLMのフルパラメータの指示チューニングを試す

    npaka

    「Google Colab」で「DeepSpeed」によるLLMの (LoRAではなく) フルパラメータの指示チューニング (Instruction Tuning) を試したので、まとめました。

    【注意】 Google Colab Pro/Pro+のA100で動作確認しています。

    前回

    1. DeepSpeed

    「DeepSpeed」は、深層学習モデルの学習や推論の処理を高速かつメモリ消費を抑えて実現することができるライブラリです。

    HuggingFaceでサポートしている「DeepSpeed」の機能は次のとおりです。

    ・Optimizer state partitioning (ZeRO stage 1)
    ・Gradient partitioning (ZeRO stage 2)
    ・Parameter partitioning (ZeRO stage 3)
    ・Custom mixed precision training handling
    ・A range of fast CUDA-extension-based optimizers
    ・ZeRO-Offload to CPU and NVMe

    論文は以下で参照できます。

    ・ZeRO: Memory Optimizations Toward Training Trillion Parameter Models : Zero
    ・ZeRO-Offload: Democratizing Billion-Scale Model Training : Zero-Offload
    ・ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning : NVMeサポート

    「DeepSpeed ZeRO-2」は、推論には役に立たないため、学習のみに使用されます。「DeepSpeed ZeRO-3」は、単一の GPUでは不可能な巨大なモデルを複数のGPU にロードできるため、推論にも使用できます。

    2. モデルとデータセット

    今回は、LLMとして「OpenCALM-1B」、データセットとして「databricks-dolly-15k-ja」を使いました。

    3. ファインチューニング前のLLM出力の確認

    Colabでファインチューニング前のLLM出力を確認する手順は、次のとおりです。

    (1) パッケージのインストール。

    # パッケージのインストール
!pip install transformers trl peft datasets

    (2) トークナイザーとモデルの準備。

    from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(
    "cyberagent/open-calm-1b",
)

# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
    "cyberagent/open-calm-1b",
    torch_dtype=torch.bfloat16,
).to("cuda")

    (3) LLM出力の確認。

    import torch

# プロンプトの準備
prompt = "### User: 日本の首都は？\n ### Answer:"

# 推論の実行
for i in range(10):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        tokens = model.generate(
            **inputs,
            max_new_tokens=64,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.pad_token_id,
        )
    output = tokenizer.decode(tokens[0], skip_special_tokens=True)
    print(output)
    print("----")
    ### User: 日本の首都は?
 ### Answer: 東京。 (Fri) 政府機関と地方自治体により構成されています。#   Used an unexpectant city namely Japan, the capital of Tokyo.
----
### User: 日本の首都は?
 ### Answer: 東京オリンピック開催予定! 国立競技場の建設が順調に進行中だ。 @Tokyo.jpと@www/wikipediaには、東京の日本の国土面積や人口などの情報はないようだ。#Ungrouped by:#limestyleis a majority of Japan’s national history s
----
### User: 日本の首都は?
 ### Answer: 東京です。日本の都市の中では、東京がもっとも大きいでしょ。# 都道府県別の人口分布・面積ランキングは下記の通り(平成26年国勢調査)
----
### User: 日本の首都は?
 ### Answer: 東京。 人口は17万5千人以上で、東京都の人口ランキングの21位にランク付けされているよ! 世界的には人口が8位だと言われているけど...。あと東京にあるのは皇居とか天皇陛下のお住まいのある赤坂御用地があるよね。#Anaizzo_Middle Earthweekday.com
----
### User: 日本の首都は?
 ### Answer: 東京
----
### User: 日本の首都は?
 ### Answer: 仙台です! Where is the capital of Sendai in Japan? I'm looking to visit Nagamasa.
----
### User: 日本の首都は?
 ### Answer: 奈良県です。 @@, 東南アジアのタイ王国にある都市の一つですよ。#
----
### User: 日本の首都は?
 ### Answer: 東京です。日本の最初の政府の首都で、東京は日本政府の中心都市で、# メインランドは東京ですがこれは真実ではありません。# It is not true. There are two wings of Tokyo's centre in Japan, the first sight attorney’s headquarters and also one f
----
### User: 日本の首都は?
 ### Answer: 東京オリンピック開催国。  ・Must be an agent in a basement of racing, circuit and metal products. The property that the company’ssured is limited to its own structures by their manufacturers who operate
----
### User: 日本の首都は?
 ### Answer: 東京です。  (省略されています)
----

    回答の間違いもありますが、モデルサイズが起因すると思われるので目をつむります。
    「### User: 日本の首都は？\n ### Answer:」に対して回答に不必要な情報が多いことがわかります。

    4. ファインチューニングの実行

    Colabでファインチューニングを実行する手順は、次のとおりです。

    (1) パッケージのインストール。

    # パッケージのインストール
!pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
!pip install git+https://github.com/microsoft/deepspeed
!pip install transformers trl peft datasets

    インストール方法は、HuggingFaceが提供しているNotebookを参考にしてます。

    ・transformers + deepspeed CLI

    (2) DeepSpeedのコンフィグファイルを作成。

    %%writefile 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
}

    DeepSpeedのコンフィグファイルの設定項目は、以下で確認できます。

    (3) 学習スクリプトの作成。

    %%writefile train.py
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from datasets import load_dataset
from transformers import TrainingArguments
from trl import SFTTrainer


# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(
    "cyberagent/open-calm-1b", 
)

# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
    "cyberagent/open-calm-1b", 
    torch_dtype=torch.bfloat16,
).to("cuda")
model.config.use_cache = False


# データセットの読み込み読み込み
dataset = load_dataset("kunishou/databricks-dolly-15k-ja")

# 確認
print(dataset)
print(dataset["train"][0])


# データセットをinput空のみ2000個でフィルタリング
train_dataset = dataset["train"].filter(lambda data: data["input"] == "").select(range(2000))

# 確認
print(train_dataset)
print(train_dataset[0])


# 書式関数
def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"### Question: {example['instruction'][i]}\n ### Answer: {example['output'][i]}<|endoftext|>"
        output_texts.append(text)
    return output_texts


# 学習パラメータ
training_args = TrainingArguments(
    output_dir='./output',  # 出力パス
    num_train_epochs=1,  # エポック数
    per_device_train_batch_size=1,  # 学習バッチサイズ
    per_device_eval_batch_size=1,  # 評価バッチサイズ
    gradient_accumulation_steps=1,  # 勾配の累積ステップ数
    gradient_checkpointing=True,  # 勾配チェックポイント
    bf16=True,  # bf16
    optim='adafactor',  # オプティマイザの種類
    deepspeed='./zero_train.json',  # deepspeed設定ファイル
    logging_steps=100,  # 途中経過を表示する間隔
)

# 学習の実行
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    max_seq_length=512,
    formatting_func=formatting_prompts_func,
    args=training_args,
)
trainer.train()

# モデルの保存
trainer.save_model("output")

    前回のコードを1つのコードにまとめて、DeepSpeedの設定を追加しています。
    変更点は、次のとおりです。

    ・モデルIDを"cyberagent/open-calm-1b"に変更
    ・model.config.use_cache = False を追加
    ・データセットの数を2000に変更  (学習時間短縮)
    ・TrainingArgumentsの追加 (DeepSpeedの設定)

    (4) 学習の実行。
    DeepSpeedのおかげで、1BでもCPU 42.4GB、GPU 7.2GB で学習できています。学習完了まで3時間半かかりました。

    # 学習の実行
!deepspeed --num_gpus=1 train.py
    画像

    (4) プロンプトを作成する関数の準備。
    データセットのinstructionとoutputを使ってプロンプトを作成します。終端文字 (<|endoftext|>) も追加しました。

    # プロンプトを作成する関数
def formatting_prompts_func(example):
    output_texts = []
    for i in range(len(example['instruction'])):
        text = f"### Question: {example['instruction'][i]}\n ### Answer: {example['output'][i]}<|endoftext|>"
        output_texts.append(text)
    return output_texts

    (5) 学習の実行。
    18分ほどかかりました。

    from trl import SFTTrainer

# 学習の実行
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    max_seq_length=512,
    formatting_func=formatting_prompts_func,
)
trainer.train()

# モデルの保存
trainer.save_model("output")

    outputフォルダには、モデルが保存されています。

    5. ファインチューニング後のLLM出力の確認

    Colabでのファインチューニング後のLLM出力の確認の手順は、次のとおりです。

    (1) outputからのモデルの読み込み。

    from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# トークナイザーの準備
tokenizer = AutoTokenizer.from_pretrained(
    "cyberagent/open-calm-1b",
)

# モデルの準備
model = AutoModelForCausalLM.from_pretrained(
    "cyberagent/open-calm-1b",
    torch_dtype=torch.bfloat16,
).to("cuda")

    (2) LLM出力の確認。

    import torch

# プロンプトの準備
prompt = "### User: 日本の首都は？\n ### Answer:"

# 推論の実行
for i in range(10):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        tokens = model.generate(
            **inputs,
            max_new_tokens=64,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.pad_token_id,
        )
    output = tokenizer.decode(tokens[0], skip_special_tokens=True)
    print(output)
    print("----")
    ### User: 日本の首都は?
 ### Answer: 日本で首都は東京都です。
----
### User: 日本の首都は?
 ### Answer: 現在の日本の国の首都は、東京です。
----
### User: 日本の首都は?
 ### Answer: 日本で一番大きな首都は、東京です。
----
### User: 日本の首都は?
 ### Answer: 首都は日本の東京です。
----
### User: 日本の首都は?
 ### Answer: 首都は日本です。
----
### User: 日本の首都は?
 ### Answer: 首都は日本です。
----
### User: 日本の首都は?
 ### Answer: 首都は東京都です。
----
### User: 日本の首都は?
 ### Answer: 日本で首都といえば、東京です。
----
### User: 日本の首都は?
 ### Answer: 日本で一番大きな首都は、東京です。
----
### User: 日本の首都は?
 ### Answer: 首都は日本です。
----

    回答は間違いもありますが、モデルサイズが起因すると思われるので目をつむります。
    「### User: 日本の首都は？\n ### Answer:」に対して回答のみが出力
    されており、質問に対して回答を返す書式を学習できていることがわかります。

    【おまけ】 ZeRO Stageとoffloadの選択基準

    「ZeRO Stage」と「offload」では、一般に次のことが当てはまります。

    ◎ 速度（左の方が右より速い）

    Stage 0 (DDP) > Stage 1 > Stage 2 > Stage 2 + offload > Stage 3 > Stage 3 + offloads

    ◎ GPUメモリ効率 (右は左よりもGPUメモリ効率が高い)

    Stage 0 (DDP) < Stage 1 < Stage 2 < Stage 2 + offload < Stage 3 < Stage 3 + offloads

    したがって、最小限の数の GPU に収まりながら最速の実行を実現したい場合は、次の手順に従うことができます。最も速い手法から開始し、GPU の Out Of Memory に陥った場合は、次に速い手法に進む手順になります。

    (1) batch_sizeを1に設定
    (2) gradient_checkpointingの有効化
    (3) ZeRO stage 2を試す
    (4) ZeRO stage 2 + offload_optimizer を試す
    (5) ZeRO stage 3 を試す
    (6) ZeRO stage 3 + offload_param を試す
    (7) ZeRO stage 3 + offload_param + offload_optimizer を試す

    batch_size1に適合しない場合は、まずさまざまなデフォルト値を確認し、可能であれば値を下げます。fp32 では必ず混合半精度を使用してください。つまり、A100では bf16、古い GPUでは fp16を使用してください。

    参考



     
     
     

    npaka

     
     
    プログラマー。iPhone / Android / Unity / ROS / AI / AR / VR / RasPi / ロボット / ガジェット。年2冊ペースで技術書を執筆。アニソン / カラオケ / ギター / 猫 twitter : @npaka123

    あなたへのおすすめ