NVIDIA RTX 4060 Ti (16GB) で LLM ファインチューニングの実行時間を測定しました

はじめに

大規模言語モデル (LLM) は 3.6B モデルや 7B モデルだと手軽に動かすことができますが、それでも GPU の VRAM が 16GB くらいは欲しいところです。最近、このクラスで比較的入手しやすい NVIDIA RTX 4060 Ti (16GB) (以下、RTX4060Ti) が発売されました。

この RTX4060Ti がどの程度の実力を持っているのかを調べるため、LLM のファインチューニングを実行し、その実行時間を測定しました。また、Google Colab 環境で利用可能な A100, V100, T4 のそれぞれの GPU でも実行時間を測定し、それらと比較しました。結果として、RTX4060Ti は A100 の半分程度、V100 や T4 を大きく上回る実行速度であることが分かりました。

前提

以下の環境で実行速度を測定しました。

• 今回の測定ではすべて Google Colab 上で行いました。
• Colab Pro+ の契約をしています。
• RTX4060Ti を使った測定はローカルランタイムとして実行しました。
• RTX4060Ti ローカルランタイムは以下の構成です。
  • CPU: Ryzen 7 5700X
  • Memory: 64GB DDR4-3200
  • SSD: 2TB NVMe PCIe4.0x4
  • GPU: RTX 4060 Ti (16GB)
  • OS: Ubuntu Server 22.04 LTS
  • その他: CUDA 11.8, cuDNN 8.9.4, python 3.10.12, pytorch 2.0.1
  • 環境構築は以下の記事を参考にしました。

ベンチマークの概要

この記事で使った 2 つのベンチマークプログラムは、以下の記事に記載されたプログラムをベンチマーク用に扱いやすいように加工したものです。

どちらも学習アルゴリズムは同じで以下の異なるモデルを用いています。モデルは 8bit 量子化をしています。

• cyberagent/open-calm-7b
• rinna/japanese-gpt-neox-3.6b

ファインチューニングのデータは以下のものを使いました。

• kunishou/databricks-dolly-15k-ja

学習方法は LoRa を用いています。データ数は 15015 件です。ここから 2000 件を検証用データとし、残りを学習用データとしています。バッチサイズは 8 (自動バッチサイズ設定)、バッチ数は 1627、エポックは 3 です。学習を実行する総バッチ数は 4881 になります。

詳細は先に挙げた2つの記事と、この記事の付録に載せたベンチマークプログラムを参照してください。

実行結果

実行速度と VRAM 使用量は以下の通りとなりました。

OpenCALM-7B + databricks-dolly-15k-ja

GPU	Time	it/s	Rate	Used(MiB)
A100 / 40GB	4:08:10	0.328	1.000	13189
V100 / 16GB	12:59:54	0.104	0.317	14050
T4 / 16GB	18:51:55	0.072	0.220	12351
RTX4060Ti / 16GB	8:52:52	0.153	0.466	11327

rinna-3.6B + databricks-dolly-15k-ja

GPU	Time	it/s	Rate	Used(MiB)
A100 / 40GB	2:55:20	0.464	1.000	8601
V100 / 16GB	7:30:46	0.180	0.388	8046
T4 / 16GB	11:48:29	0.115	0.248	8041
RTX4060Ti / 16GB	5:57:40	0.227	0.489	6647

ここで、Time は学習時間、it/s は 1 秒あたりイテレーション数で、バッチ数 (今回は 4881) / 学習時間 (s) で計算します。Rate は A100 の実行速度を 1.0 とした時の各 GPU に対する相対的な実行速度倍率です。Used は GPU の VRAM 使用量です。

考察

LLM のファインチューニングについて、RTX 4060 Ti (16GB) の実行速度は A100 の 0.47 $\sim$ 0.49 倍、V100 の 1.26 $\sim$ 1.47 倍、T4 の 1.97 $\sim$ 2.13 倍になりました。学習データの量によって実行時間は変わりますが、7Bサイズ以下のモデルのファインチューニングについて、夕方に学習を開始し翌朝に結果を確認するリズムで作業をするのであれば、RTX 4060 Ti (16GB) は十分実用的な GPU であることが分かりました。

これより上位の GPU を選ぼうとすると、RTX 4080 (16GB)、RTX 4090 (24GB) か、プロフェッショナル向け GPU から選ぶことになるでしょう。AMD の方は RX 7800 XT (16GB)、RX 7900 XT (20GB)、RX 7900 XTX (24GB) あたりが候補ですが、ROCm が RDNA3 に対応するまでは性能面で見劣りすることになりそうです。

VRAM 使用量にかなりばらつきがあります。実行のたびにも変化するのですが、原因が分かりません。深層学習やファインチューニングに手を付けたばかりで、何か間違いがあるのかもしれません。お気付きのことがありましたら、お知らせください。

注意事項

Google drive の接続がプログラム実行中に切断される問題

当初は Google drive をマウントして学習過程および学習結果のファイルを保存するようにしていましたが、なぜかマウントしてから 14000 秒前後で Google drive が切断され、学習過程等のファイルが保存できずにプログラムが中断するという状態になっています。原因不明のため不本意ですが、ランタイム上のストレージに保存しています。学習結果のファイルはプログラムが正常終了したのちに手動操作で Google drive をマウントしてコピーしました。

Google Colab Pro+ 契約でもランタイムが24時間経過前に切断される問題

Colab Pro+ の契約ではランタイムのセッション切れは相当程度起こらないはずですが、途中のエラーで停止した場合などでランタイムのセッション切れがしばしば発生しました。そこで、セッション切れが起きないようにするため、Colab 実行用の Ubuntu ラップトップを1台用意し、以下の記事で紹介があったブラウザ上の自動クリックプログラムを動かして対処しました。

付録

Google Colab にコピーアンドペーストして使うことを想定しています。そのままでは Python プログラムとして動作しませんので、ご注意ください。

OpenCALM-7B-finetuning-benchmark

# PEFTのインストール
!pip install -Uqq  git+https://github.com/huggingface/peft.git
!pip install -Uqq transformers datasets accelerate bitsandbytes

# 基本パラメータ
model_name = "cyberagent/open-calm-7b"
dataset = "kunishou/databricks-dolly-15k-ja"
peft_name = "lora-calm-7b"
output_dir = "lora-calm-7b-results"

# トークナイザーの準備
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)

CUTOFF_LEN = 256  # コンテキスト長

# トークナイズ関数の定義
def tokenize(prompt, tokenizer):
    result = tokenizer(
        prompt+"<|endoftext|>",  # EOSの付加
        truncation=True,
        max_length=CUTOFF_LEN,
        padding=False,
    )
    return {
        "input_ids": result["input_ids"],
        "attention_mask": result["attention_mask"],
    }

# データセットの準備
from datasets import load_dataset
data = load_dataset(dataset)

# プロンプトテンプレートの準備
def generate_prompt(data_point):
    if data_point["input"]:
        return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{data_point["instruction"]}

### Input:
{data_point["input"]}

### Response:
{data_point["output"]}"""
    else:
        return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{data_point["instruction"]}

### Response:
{data_point["output"]}"""

# 学習データと検証データの準備
VAL_SET_SIZE = 2000
train_val = data["train"].train_test_split(
    test_size=VAL_SET_SIZE, shuffle=True, seed=42
)
train_data = train_val["train"]
val_data = train_val["test"]
train_data = train_data.shuffle().map(lambda x: tokenize(generate_prompt(x), tokenizer))
val_data = val_data.shuffle().map(lambda x: tokenize(generate_prompt(x), tokenizer))

# モデルの準備
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,
    device_map="auto",
)

# LoRAのパラメータ設定
from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType
lora_config = LoraConfig(
    r= 8,
    lora_alpha=16,
    target_modules=["query_key_value"],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

# モデルの前処理
model = prepare_model_for_int8_training(model)

# LoRAモデルの準備
model = get_peft_model(model, lora_config)

# トレーナーの準備
import transformers

eval_steps = 200
save_steps = 200
logging_steps = 20

trainer = transformers.Trainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=val_data,
    args=transformers.TrainingArguments(
        num_train_epochs=3,
        learning_rate=3e-4,
        logging_steps=logging_steps,
        evaluation_strategy="steps",
        save_strategy="steps",
        eval_steps=eval_steps,
        save_steps=save_steps,
        output_dir=output_dir,
        report_to="none",
        save_total_limit=3,
        push_to_hub=False,
        auto_find_batch_size=True,
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

# 学習の実行
model.config.use_cache = False
trainer.train()
model.config.use_cache = True

# LoRAモデルの保存
trainer.model.save_pretrained(peft_name)

# VRAM使用量の確認
! nvidia-smi -q -d MEMORY

rinna-3.6B-finetuning-benchmark

# PEFTのインストール
!pip install -Uqq  git+https://github.com/huggingface/peft.git
!pip install -Uqq transformers datasets accelerate bitsandbytes
!pip install sentencepiece

# 基本パラメータ
model_name = "rinna/japanese-gpt-neox-3.6b"
dataset = "kunishou/databricks-dolly-15k-ja"
peft_name = "lora-rinna-3.6b"
output_dir = "lora-rinna-3.6b-results"

# トークナイザーの準備
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)

# CALMのtokenizer()はEOSが自動追加されなかったので「+"<|endtext|>"」を付加しましたが、Rinnaのtokenizer()は自動追加されてました。
CUTOFF_LEN = 256  # コンテキスト長

# トークナイズ関数の定義
def tokenize(prompt, tokenizer):
    result = tokenizer(
        prompt,
        truncation=True,
        max_length=CUTOFF_LEN,
        padding=False,
    )
    return {
        "input_ids": result["input_ids"],
        "attention_mask": result["attention_mask"],
    }

# データセットの準備
from datasets import load_dataset
data = load_dataset(dataset)

# プロンプトテンプレートの準備
def generate_prompt(data_point):
    if data_point["input"]:
        result = f"""### 指示:
{data_point["instruction"]}

### 入力:
{data_point["input"]}

### 回答:
{data_point["output"]}"""
    else:
        result = f"""### 指示:
{data_point["instruction"]}

### 回答:
{data_point["output"]}"""

    # 改行→<NL>
    result = result.replace('\n', '<NL>')
    return result

# 学習データと検証データの準備
VAL_SET_SIZE = 2000
train_val = data["train"].train_test_split(
    test_size=VAL_SET_SIZE, shuffle=True, seed=42
)
train_data = train_val["train"]
val_data = train_val["test"]
train_data = train_data.shuffle().map(lambda x: tokenize(generate_prompt(x), tokenizer))
val_data = val_data.shuffle().map(lambda x: tokenize(generate_prompt(x), tokenizer))

# モデルの準備
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,
    device_map="auto",
)

# LoRAのパラメータ設定
from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType
lora_config = LoraConfig(
    r= 8,
    lora_alpha=16,
    target_modules=["query_key_value"],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

# モデルの前処理
model = prepare_model_for_int8_training(model)

# LoRAモデルの準備
model = get_peft_model(model, lora_config)

# トレーナーの準備
import transformers

eval_steps = 200
save_steps = 200
logging_steps = 20

trainer = transformers.Trainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=val_data,
    args=transformers.TrainingArguments(
        num_train_epochs=3,
        learning_rate=3e-4,
        logging_steps=logging_steps,
        evaluation_strategy="steps",
        save_strategy="steps",
        eval_steps=eval_steps,
        save_steps=save_steps,
        output_dir=output_dir,
        report_to="none",
        save_total_limit=3,
        push_to_hub=False,
        auto_find_batch_size=True,
    ),
    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

# 学習の実行
model.config.use_cache = False
trainer.train()
model.config.use_cache = True

# LoRAモデルの保存
trainer.model.save_pretrained(peft_name)

# VRAM使用量の確認
! nvidia-smi -q -d MEMORY

2つのベンチマークプログラムの差分

$ diff OpenCALM-7B-finetuning-benchmark.ipynb.txt rinna-3.6B-finetuning-benchmark.ipynb.txt
3a4
> !pip install sentencepiece
6c7
< model_name = "cyberagent/open-calm-7b"
---
> model_name = "rinna/japanese-gpt-neox-3.6b"
8,9c9,10
< peft_name = "lora-calm-7b"
< output_dir = "lora-calm-7b-results"
---
> peft_name = "lora-rinna-3.6b"
> output_dir = "lora-rinna-3.6b-results"
13c14
< tokenizer = AutoTokenizer.from_pretrained(model_name)
---
> tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
14a16
> # CALMのtokenizer()はEOSが自動追加されなかったので「+"<|endtext|>"」を付加しましたが、Rinnaのtokenizer()は自動追加されてました。
20c22
<         prompt+"<|endoftext|>",  # EOSの付加
---
>         prompt,
37,39c39
<         return f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
<
< ### Instruction:
---
>         result = f"""### 指示:
42c42
< ### Input:
---
> ### 入力:
45c45
< ### Response:
---
> ### 回答:
48,50c48
<         return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
<
< ### Instruction:
---
>         result = f"""### 指示:
53c51
< ### Response:
---
> ### 回答:
54a53,56
>
>     # 改行→<NL>
>     result = result.replace('\n', '<NL>')
>     return result