Flash AttentionとDeep Speedを使ってLLMをGPU1枚でフルファインチューニングする

## TL;DR
Deep Speedのoffload_optimizer機能を使って、VRAM40GBのGPU1枚で3.6BパラメータのLLMをファインチューニングしました。
さらに、Flash Attentionを使うことで、学習を高速化しつつ使用メモリ量も減らし、より長い系列長で学習を行うことができました。

## はじめに
Flash AttentionはAttentionの計算で使用するGPUメモリ量を系列長の2乗のオーダーから1乗に削減する技術で、xformersやoptimum、さらにはtransformersライブラリでも簡単に使用できるようになってきています。
前回の記事([1](https://qiita.com/jovyan/items/11deb9d4601e4705a60d), [2](https://qiita.com/jovyan/items/5716cd83e246df4a158e))では、LLMの推論においてFlash Attentionを使ってみたのですが、推論ではあまり効果が見られないという結論を得ました。
今回はFlash Attentionが本領を発揮するであろう学習（ファインチューニング）で効果を確認します。

せっかくなので、LoRAは使わず全パラメータをフルファインチューニングすることとし、限られたリソース（GPUメモリ）でフルファインチューニングを行うために、Deep Speedライブラリを使用します。

ファインチューニングの方法やDeep Speedライブラリの使い方は、主に以下の記事を参考にしました。
> [Fine-tune Falcon 180B with DeepSpeed ZeRO, LoRA & Flash Attention](https://www.philschmid.de/deepspeed-lora-flash-attention)

こちらの記事はFalcon 180Bという超巨大モデルを扱っていますが、本記事では3.6Bの日本語LLM [line-corporation/japanese-large-lm-3.6b](https://huggingface.co/line-corporation/japanese-large-lm-3.6b)を使用します。

## 環境
実験はGoogle Colab (Pro)上で、A100GPU (Syetem RAM 83.5GB, GPU RAM 40.0GB)1枚を使用して行いました。

### ライブラリのインストール
transformersと関連ライブラリ
```
!pip install transformers datasets accelerate bitsandbytes sentencepiece
```

transformersのモデルでFlash Attentionを使うためのライブラリ
```
!pip install optimum
```

deepspeedと関連ライブラリ
```
!pip install deepspeed ninja
```

### Pythonおよびライブラリのバージョン
```
Python 3.10.12
accelerate                       0.23.0
bitsandbytes                     0.41.1
datasets                         2.14.5
deepspeed                        0.10.3
ninja                            1.11.1
optimum                          1.13.2
sentencepiece                    0.1.99
torch                            2.0.1+cu118
transformers                     4.33.3
```

## データセット準備
ファインチューニングには日本語instruction tuning用データセットの[databricks-dolly-15k-ja](https://huggingface.co/datasets/kunishou/databricks-dolly-15k-ja)を使用しました。
[参考記事](https://www.philschmid.de/deepspeed-lora-flash-attention)の方法に従って処理します。

この段階ではGPUは必要ないので、学習とは別のノートブック上でCPUインスタンスを用いて行います。

Google Driveマウント
```
from google.colab import drive
drive.mount('/content/drive')

%cd /content/drive/MyDrive/working_dir
```

データセット読み込み
```Python
from datasets import load_dataset

dataset = load_dataset("kunishou/databricks-dolly-15k-ja", split="train")
print(f"dataset size: {len(dataset)}")  # -> dataset size: 15015
```
このデータセットは、指示文である'instruction'と補助的な情報の'input'、そして回答の'output'が収められたものです。'input'が存在しないデータも存在します。

トークナイザー準備
```Python
import torch
from transformers import AutoTokenizer

model_name = "line-corporation/japanese-large-lm-3.6b"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
```

データセットをプロンプトのテンプレートの形に整える関数`format_prompt()`と`template_dataset()`を準備して、datasetに適用します。
```Python
def format_prompt(sample):
    instruction = f"### Instruction\n{sample['instruction']}"
    input = f"### Context\n{sample['input']}" if len(sample["input"]) > 0 else None
    output = f"### Answer\n{sample['output']}"
    # join all the parts together
    prompt = "\n\n".join([i for i in [instruction, input, output] if i is not None])
    return prompt

def template_dataset(sample):
    sample["text"] = f"{format_prompt(sample)}{tokenizer.eos_token}"
    return sample

dataset = dataset.map(template_dataset, remove_columns=list(dataset.features))
```

データを一つサンプルすると以下のような形式になっています。
```Python
from random import randint

print(dataset[randint(0, len(dataset))]["text"])
```

```
### Instruction
コルチゾールをコントロールするための5つのベストな方法とは

### Answer
バイノーラルサウンドをヘッドホンで聴く、1日5分以上の瞑想、1日30分以上のリラックスできる内容の本を読む、1日10分以上の明け方や夕暮れ時の日光を浴びる、1日30分以上のウォーキングを毎日する。

```

素のLLMを1からinstruction tuningするので、テンプレートは自由に決めても構いません。
例えば、日本語LLMなので、Instruction, Answerに関しては指示、回答などのように日本語にしても問題ないでしょう。
今回は、はじめELYZA-japanese-Llama-2-7bモデルを使おうとした経緯があり、ELYZA-japanese-Llama-2モデルではトークナイザーが英語で学習されたものだったので英語のままにしていました。


datasetにトークナイザーを適用します。
```Python
tokenized_dataset = dataset.map(
    lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(dataset.features)
)

print(tokenized_dataset[0].keys())  # -> dict_keys(['input_ids', 'attention_mask'])
```

トークン化した後の系列長の統計情報を見てみましょう。
```Python
import numpy as np

seq_lengths = [len(sample["input_ids"]) for sample in tokenized_dataset]

print(f"mean: {np.mean(text_lengths)}")
print(f"median: {np.median(text_lengths)}")
print(f"min: {np.min(text_lengths)}")
print(f"max: {np.max(text_lengths)}")
```

> mean: 397.3209457209457  
> median: 250.0  
> min: 39  
> max: 14037

一部、外れ値的に長いデータも含まれているようです。

このように、長さの違うデータをバッチとしてまとめてTransformerモデルに入力する一つの方法はパディングを行って長さを揃えることですが、それを行うとデータ効率が悪くなってしまいます。今回のようなファインチューニングならともかく、大量のデータを使用する事前学習時には致命的です。また、Flash Attentionではバッチ内でデータ行ごとに異なるマスク(attention bias)をかけることはできないので、コードレベルでの特別な対処が必要となってしまいます。
ここでは別の方法として、最大系列長に達するまで次のデータを詰め込み、一つのデータの途中で最大系列長に達した場合はそこで文章を切って、次のデータ行に続きから詰め込むということをします。この方法では、一部の文章は途中で分断されてしまうというデメリットはあるものの、全てのデータ行で最大系列長までテキストが詰まっているので、計算効率が良いというメリットがあります。

以下、[参考記事](https://www.philschmid.de/deepspeed-lora-flash-attention)から拝借した実装コードです。最大系列長は2048としています。
<details><summary>長いので折り畳んでおきます。</summary><div>

```Python
# empty list to save remainder from batches to use in next batch
remainder = {"input_ids": [], "attention_mask": [], "token_type_ids": []}

def chunk(sample, chunk_length=2048):
    # define global remainder variable to save remainder from batches to use in next batch
    global remainder
    # Concatenate all texts and add remainder from previous batch
    concatenated_examples = {k: list(chain(*sample[k])) for k in sample.keys()}
    concatenated_examples = {k: remainder[k] + concatenated_examples[k] for k in concatenated_examples.keys()}
    # get total number of tokens for batch
    batch_total_length = len(concatenated_examples[list(sample.keys())[0]])

    # get max number of chunks for batch
    if batch_total_length >= chunk_length:
        batch_chunk_length = (batch_total_length // chunk_length) * chunk_length

    # Split by chunks of max_len.
    result = {
        k: [t[i : i + chunk_length] for i in range(0, batch_chunk_length, chunk_length)]
        for k, t in concatenated_examples.items()
    }
    # add remainder to global variable for next batch
    remainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()}
    # prepare labels
    result["labels"] = result["input_ids"].copy()
    return result

# tokenize and chunk dataset
chunked_dataset_2048 = tokenized_dataset.map(
    partial(chunk, chunk_length=2048),
    batched=True,
)

# Print total number of samples
print(f"Total number of samples: {len(chunked_dataset_2048)}")  # -> Total number of samples: 1407

# Save
chunked_dataset_2048.save_to_disk("dolly-15k-ja_processed_2048.hf")
```

</div></details>

データの1行目をデコードして表示してみます。
```Python
print(tokenizer.decode(chunked_dataset_2048[0]['input_ids'], skip_special_tokens=False))
```

```
### Instruction
ヴァージン・オーストラリア航空はいつから運航を開始したのですか?

### Context
ヴァージン・オーストラリア航空(Virgin Australia Airlines Pty Ltd)は＜～略～＞

### Answer
ヴァージン・オーストラリア航空は、2000年8月31日にヴァージン・ブルー航空として、2機の航空機で単一路線の運航を開始しました。
 ### Instruction
魚の種類はどっち?イコクエイラクブカとロープ

### Answer
イコクエイラクブカ
 ### Instruction
ラクダはなぜ水なしで長く生きられるのか?

### Answer
＜～～略～～＞
### Context
トーマス・ジェファーソン(1743年4月13日[a]-1826年7月4日)は、＜～略～＞アメリカ独立戦争
```
一部を中略しています。EOSトークン`
`の後に改行が入らないので見づらいですが、元データセットでの複数のデータ行が一つのデータ行にまとめられていることがわかります。
また、最後はContextの文章の途中でぶつ切りにされています。
その続きは、次のデータ行の最初に来るのですが、学習時は続きとしてモデルに入力されることはなく、モデルは途中から始まる謎の文章を見せられ（予測させられ）、Instructionを見てないのにAnswerを見せられる（予測させられる）ということになります。

おそらくこのようなデータセットの処理はLLMの学習では普通に行われていることだと思うのですが、だとすると絶対位置埋め込みは良い方法ではなく、相対位置を考慮することができ、長距離での依存性が減衰する性質を持つ[Rotary Position Embedding (RoPE)](https://arxiv.org/abs/2104.09864)や[xPos](https://arxiv.org/abs/2212.10554)がLLMではメジャーな手法となっているのも納得です。


後の実験のため、系列長512と1024のdatasetも作成して保存しておきます。
```Python
remainder = {"input_ids": [], "attention_mask": [], "token_type_ids": []}
chunked_dataset_512 = tokenized_dataset.map(
    partial(chunk, chunk_length=512),
    batched=True,
)
print(f"Total number of samples: {len(chunked_dataset_512)}")  # Total number of samples: 5629
chunked_dataset_512.save_to_disk("dolly-15k-ja_processed_512.hf")

remainder = {"input_ids": [], "attention_mask": [], "token_type_ids": []}
chunked_dataset_1024 = tokenized_dataset.map(
    partial(chunk, chunk_length=1024),
    batched=True,
)
print(f"Total number of samples: {len(chunked_dataset_1024)}")  # Total number of samples: 2814
chunked_dataset_1024.save_to_disk("dolly-15k-ja_processed_1024.hf")
```

## 学習
モデルはLINEが公開している[japanese-large-lm-3.6b](https://huggingface.co/line-corporation/japanese-large-lm-3.6b)を使用します。事前学習のみでinstruction tuningは行われていないモデルです。

学習はHugging Face TransformerライブラリのTrainerを用い、GPU1枚で学習を行うためにDeep SpeedのZeRO-2を使用します。Deep Speedの使い方に関しては、先に挙げた[記事](https://www.philschmid.de/deepspeed-lora-flash-attention)と[Hugging Faceのドキュメント](https://huggingface.co/docs/transformers/v4.26.1/en/main_classes/deepspeed)を参考にしました。

以下の実験は、A100GPUのインスタンスを用いて行いました。GPU RAMは40.0GB、システム(CPU) RAMは83.5GBです。

通常のAttentionを用いた場合とFlash Attentionを用いた場合でそれぞれ、系列長512, 1024, 2048のデータで学習させることを試みます。

### 通常Attention
A100GPUのカーネルを使ったColabノートブック上で、環境のセクションに記したライブラリをインストールし、Google Driveをマウント、データセットを保存したディレクトリに移動したものとします。

初めに、ノートブック上でDeep Speedを使うための環境変数の設定をします。
```Python
import os

os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "9994"
os.environ["RANK"] = "0"
os.environ["LOCAL_RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
```

次にデータセットを読み込みます。まずは系列長512のデータを使います。
```Python
from datasets import load_from_disk

dataset = load_from_disk("./dolly-15k-ja_processed_512.hf")
```

モデルとトークナイザーを読み込みます。モデルのパラメータはfloat16とします。
```Python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "line-corporation/japanese-large-lm-3.6b"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    )
model.cuda()
```

さらに、言語モデルの学習のためのcollatorを準備します。
```Python
from transformers import DataCollatorForLanguageModeling

tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
```

Deep Speedを使うためにはjsonの設定ファイルをTrainerに渡す必要があるので、作成します。
```
%%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": "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": 2e9,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e9,
        "contiguous_gradients": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto"
}
```

"auto"と書いてある部分はTrainerの設定で上書きされます。
ここで重要となるのは、"zero_optimization"の項目です。"stage"は1, 2, 3から選べますが、offload_optimizer機能を使えるのはStage 2か3です。Stage 3はモデルパラメータを複数GPUに分散させる必要がある場合などに使われる設定で、GPU1枚で学習させる今回のようなケースではメリットは無いようなので、今回はStage 2を使うこととします。"offload_optimizer"で"cpu"を指定することで、オプティマイザーが保持するパラメータをCPUメモリに退避させ、GPUメモリの使用量を抑えることができます。また、"allgather_bucket_size"はGPUメモリ使用量と通信の頻度のトレードオフをコントロールするパラメータで、小さくするとGPUメモリ使用量を下げることができますが、学習は遅くなります。系統的に値を変えて比較はできていないのですが、"allgather_bucket_size"を2e8としたら学習がかなり遅かったので、2e9としました。あわせて"reduce_bucket_size"の値も2e9としていますが、これらの設定の意味は理解しきれていないので、詳しくは[Hugging Faceのドキュメント](https://huggingface.co/docs/transformers/v4.26.1/en/main_classes/deepspeed)や[Deep Speed公式ドキュメント](https://www.deepspeed.ai/)をご参照ください。

以下の実験はすべて上記の`zero_train.json`の設定で行いました。


最後にTrainerの設定を行って、学習を実行します。
```Python
from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="/content/tmp",
    overwrite_output_dir=True,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=16,
    learning_rate=5e-5,
    weight_decay=0.01,
    num_train_epochs=1,
    logging_steps=5,
    lr_scheduler_type="constant_with_warmup",
    warmup_steps=10,
    save_strategy="no",
    fp16=True,
    deepspeed="./zero_train.json",
)

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

# 学習実行
trainer.train()
```

`TrainingArguments`で重要なのは`fp16`を設定することと、`deepspeed`引数に設定ファイルのパスを与えるところです。これだけで、Deep Speedを使った学習を行うことができます。（試していませんが、fp16の代わりにbf16でも良いはずです。）
バッチサイズ`per_device_train_batch_size`は4、`gradient_accumulation_steps`は16とすることで、実質のバッチサイズを64としています。
今回の実験の目的は、Deep SpeedやFlash Attentionによる学習効率化の効果を確認することで、モデルの性能は2の次としていますので、学習は１エポックのみで、そのほかのハイパーパラメータの設定も適当です。
また、Flash Attention使用のところで後述する理由で、Trainerがモデルを保存しない設定`save_strategy="no"`としています。

学習開始前のRAMの使用量は、
> System RAM: 4.8 / 83.5GB, GPU RAM 8.0 / 40.0GB

でしたが、学習中は
> System RAM: 75.1 / 83.5GB, GPU RAM 32.5 / 40.0GB

となり、offload_optimizerの機能によってCPUメモリが大量に使用されました。
CPUメモリ容量が小さい別のColabインスタンスで試した時はCPUメモリがあふれて学習実行できなかったので、offload_optimizer機能を使うときにはGPUメモリだけでなくCPUメモリのスペックも重要となります。

88ステップの学習に対して、学習時間・速度のメトリックは
> 'train_runtime': 2456.4699, 'train_samples_per_second': 2.291, 'train_steps_per_second': 0.036

となりました。学習全体で40分程度ですので、十分手頃な時間かと思います。

Trainerはモデルを保存しない設定にしたので、自分で保存する必要があるのですが、学習後もRAMにキャッシュが残っているせいで、保存中にシステムRAMがOOMになってカーネルが落ちてしまったことがあったので、まずRAMを可能な限り空けます。

```Python
import gc
import deepspeed

del dataset
del trainer

gc.collect()
deepspeed.runtime.utils.empty_cache()
torch.cuda.empty_cache()
```

この処理を行ってもCPU RAM上の大量のキャッシュは消えなかったのですが、モデル保存中にOOMは発生しなくなりました。

モデル保存は以下のコマンドで行います。
```Python
model.save_pretrained("/content/tmp/model")
```

マウントしたDriveのディレクトリを直接指定すると保存が遅かったので、いったんローカルのディスクに保存して、後からDriveにコピーしました。
このコマンドで保存したモデル重みは`from_pretrained()`で読み込むことができます。

### Flash Attention
Flash Attentionを使う場合でも、学習の手順は上記とほとんど同じですので、違う部分のみ説明します。

今回使用しているモデルはGPT-NeoXベースで、[こちらの記事](https://qiita.com/jovyan/items/5716cd83e246df4a158e)で紹介した`from_pretrained()`に`use_flash_attention_2=True`を渡すという方法は（記事執筆時点では）使えませんので、[こちらの記事](https://qiita.com/jovyan/items/11deb9d4601e4705a60d)で紹介したoptimumライブラリのBetterTransformerの機能を用います。

モデルインスタンスを立ち上げた後で、
```Python
model.to_bettertransformer()
```

を実行するだけです。
Flash Attentionを使用するためには、`torch.backends.cuda.flash_sdp_enabled()`が`True`に設定されている必要があるのですが、今回の環境ではデフォルトで`True`になっていました。以下のコマンドで設定を確認できます。
```Python
import torch

torch.backends.cuda.flash_sdp_enabled()  # -> True
```

もし、`False`になっていた場合は、
```Python
torch.backends.cuda.enable_flash_sdp(True)
```
を実行すれば切り替えられます。

BetterTransformerで変換したモデルはそのままTrainerに渡して学習することができます。
ただし、モデルを保存するときにBetterTransformer形式から通常の形式に戻さないと保存できないというエラーが発生して、Trainerが停止してしまいます。Trainerとoptimumライブラリの連携が未実装のようです。このエラーを避けるため、`TrainingArguments`で`save_strategy="no"`としていたのでした。学習が終了したら、以下のようにモデルを通常の形式に戻してやります。
```Python
from optimum.bettertransformer import BetterTransformer

model = BetterTransformer.reverse(model)
```
この場合もキャッシュを掃除してからモデルを保存しました。

さて、Flash Attentionを使った場合の学習中のRAM使用量は以下のようになりました。
> System RAM: 75.0 / 83.5GB, GPU RAM 26.3 / 40.0GB

System RAMは変わりませんが、GPU RAMは32.5GBから26.3GBに減少しました。確かにFlash Attentionの効果があるようです。
より詳細な比較は以下で系列長の違うデータに対する結果と一緒にまとめます。

## 結果

### メモリ使用量・学習時間
通常のAttentionを用いた場合とFlash Attentionを用いた場合でそれぞれ、系列長512, 1024, 2048のデータで学習させた結果を下の表にまとめました。
系列長512ではバッチサイズ4として、系列長1024と2048ではOut of Memory (OOM)になった場合はバッチサイズを半分にしています。
バッチサイズ (per_device_batch_size)を変える時には、実質のバッチサイズが64のままになるようにgradient_accumulation_stepsも変えています。
実質バッチサイズを揃えているので、系列長が長くなれば、反比例して1エポック当たりのステップ数は減っています。

| Attention | 系列長 | バッチサイズ | ステップ数 | GPU RAM (GB) | System RAM (GB) | train runtime (sec) | train samples per second |
| :---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| Normal |  512 | 4 | 88 | 32.5 | 75.1 | 2456 | 2.29 |
| Normal | 1024 | 2 | 44 | 38.9 | 75.1 | 2322 | 1.21 |
| Normal | 2048 | 1 | 22 | OOM | - | - | - |
| Flash |  512 | 4 | 88 | 26.3 | 75.0 | 2354 | 2.39 |
| Flash |  1024 | 4 | 44 | 33.2 | 74.8 | 1330 | 2.11 |
| Flash |  2048 | 2 | 22 | 33.2 | 74.5 | 1265 | 1.11 |

RAMについては、学習中の適当なタイミングでColabのリソース表示を見て記録したものなので、多少の変動はあります。

通常のAttentionでは系列長1024はバッチサイズ2で学習が回ったのですが、系列長を2048にするとバッチサイズ1でもOOMになってしまいました。
一方で、Flash Attentionを使った場合は、系列長1024はバッチサイズ4でも学習が実行でき、系列長でも2048バッチサイズを半分にするだけでOOMを回避できました。注目すべきは、系列長1024・バッチサイズ4の時と系列長2048・バッチサイズ2の時のGPUメモリ使用量が同じというところです。これは、Flash Attention使用時のGPUメモリ使用量は系列長に比例することの帰結です。

![train_runtime.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/613488/3e75532d-732b-ebd9-b38c-6b5e12bc826d.png)
学習時間 (train runtime)を同じ系列長で比較すると、Flash Attentionの方が学習が高速であることがわかります。系列長1024で学習時間に大きな差があるのは、per_device_batch_sizeが異なるからです。（その意味で、系列長512の時もOOMにならない範囲でper_device_batch_sizeをなるべく大きくして比較を行うべきでした。）
また、系列長を長くして1ステップ当たりに含まれるトークン数を増やすと、学習が効率化されて高速化される傾向がありますので、Flash Attentionにより長い系列長での学習が可能になることで、さらに学習が高速化できていることが見て取れます。

### 学習結果
今回の実験の目的は、Deep SpeedとFlash Attentionによる学習効率化の効果を確かめることでしたので、モデルの性能は2の次としていました。
それでも、そもそも学習がまともに行えていなかったら元も子もありませんので、簡単にではありますが定量・定性的に確かめておきます。

学習曲線は下図のようになります。
![learning_curves.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/613488/96cb24dc-457b-4128-ba32-0935f318bdca.png)
どの学習設定でも、ほぼ同じようにロスが減少していることが見て取れます。
特に、同じ系列長だと通常AttentionとFlash Attentionの曲線が重なっており、Flash Attentionの計算内容は通常Attentionと同じであることが確認できます。

次に、定性的な確認として、学習させたモデルにプロンプトを入力して応答を見てみましょう。
以下の手順で「四国の県名を全て列挙してください。」という質問に対する応答を3回出力させます。
```Python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

tokenizer = AutoTokenizer.from_pretrained(
    "line-corporation/japanese-large-lm-3.6b",
    use_fast=False,
    )
model_dir = "model_save_dir"  # モデルを保存したディレクトリを指定。元モデルはline-corporation/japanese-large-lm-3.6bを指定。
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    torch_dtype=torch.float16,
    )
model.cuda()

generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device="cuda")

generator_params = dict(
    max_length = 256,
    do_sample = True,
    temperature = 0.7,
    top_p = 0.9,
    top_k = 0,
    repetition_penalty = 1.1,
    num_beams = 1,
    pad_token_id = tokenizer.pad_token_id,
)

input_text = """四国の県名を全て列挙してください。"""
prompt = f"### Instruction\n{input_text}\n\n### Answer\n"

for _ in range(3):
  output = generator(
      prompt,
      **generator_params,
  )
  print(output[0]["generated_text"])
  print("----------------" * 3)
```

ファインチューニング前の元モデルと、系列長512でFlash Attentionを使って学習させたモデルの結果を示します。

<details><summary>ファインチューニング前の元モデル</summary><div>

```
### Instruction
四国の県名を全て列挙してください。

### Answer


 香川県
 徳島県
 愛媛県
 高知県
 A: これ、結構難しいですが・・・

1.四国に都府県は4つしかありませんので、このうち3つの都府県を書けばOKです。
2.東京都区部を含む都府県名を書きます(例・東京都)
3.47都道府県のうち、北海道と沖縄県を除いた45道府県を書きます(例・大阪府)
4.都府県名はひらがなで、都道府県庁所在地は漢字で書いてください(例・東京都庁、大阪府庁)
5.45都道府県以外の都府県名は書かないでいいです(例・愛知県)
6.都府県名の次に「道」を付ける時は、「道」の字を省略せずそのまま付けてください(例・福岡県)
7.「道」を省略した場合は、その都府県には「道」が付けられないことになりますので注意してください(例・北海道、秋田県)A: こんにちは

#1さんの回答の補足ですが、宮崎県の場合は

>宮崎県をひらがなで書く場合は、宮崎市のある「南方」と「北方」に分けて書きましょう。

↑
これは間違いです
------------------------------------------------
### Instruction
四国の県名を全て列挙してください。

### Answer

1. 高知県、徳島県、香川県、愛媛県、愛媛県、高知県
2. 高知県、徳島県、香川県、徳島県、高知県、香川県
3. 高知県、徳島県、愛媛県、高知県、愛媛県、香川県、愛媛県、高知県
4. 高知県、徳島県、香川県、愛媛県、徳島県、香川県、愛媛県、愛媛県
5. 高知県、徳島県、香川県、愛媛県、徳島県、愛媛県、香川県、徳島県、愛媛県
6. 高知県、徳島県、香川県、愛媛県、愛媛県、愛媛県、徳島県、香川県、香川県、愛媛県
7. 高知県、徳島県、香川県、愛媛県、徳島県、愛媛県、徳島県、香川県、徳島県、香川県、愛媛県
8. 高知県、徳島県、香川県、愛媛県、徳島県、徳島県、香川県、愛媛県、香川県、香川県、愛媛県
9. 高知県、徳島県、香川県、徳島県、香川県、徳島県、香川県、香川県、香川県、香川県、徳島県
10. 高知県、徳島県、香川県、徳島県、香川県、徳島県、徳島県、香川県、香川県、徳島県、香川県
11. 高知県、徳島県、香川県、香川県、香川県、香川県、香川県、香川県、
------------------------------------------------
### Instruction
四国の県名を全て列挙してください。

### Answer

愛媛 高知 徳島 香川 愛媛A: 1.「A-Z順」に並び替えた後に、さらに「あいうえお順」に並び替える。
2.「あいうえお順」で並び替えてから、「A-Z順」に並び替える。
3.「あいうえお順」と「A-Z順」の両方を使って並び替える。
4.全部の文字を入力して検索する。(ただし、全角のみ)
5.文字を一文字ずつ検索してゆく。
6.「あ」「い」「う」「え」「お」を別々に入力して検索する。
7.一文字ずつ検索し、最後に「か」だけを検索する。
------------------------------------------------
```

</div></details>

正しい答えが得られている場合もありますが、対話形式で答えるように学習されていませんので、質問と関係ない出力が続いてしまいます。

<details><summary>系列長512でFlash Attentionを使って学習させたモデル</summary><div>

```
### Instruction
四国の県名を全て列挙してください。

### Answer

高知県、徳島県、愛媛県、香川県です。四国は4つの県から構成されています:
- 高知県
- 徳島県
- 愛媛県
- 香川県
四国は素晴らしい景色と自然の宝庫で、美しい観光スポットがたくさんあります。また、伝統的な文化や食べ物など、ユニークな体験ができる場所です。四国には、たくさんのアクティビティがあります!是非、訪れてみてください。

四国観光ガイド|四国観光ガイドでは、四国への旅行を計画するための役立つ情報を提供しています。ここでは、四国を旅する際に利用できるさまざまな交通機関をご紹介します。

ぜひチェックしてみてくださいね!

### Answer
四国は、高知県、徳島県、愛媛県、香川県で構成されています。四国は4つの県から構成されています。四国は素晴らしい景色と自然の宝庫です。美しい観光スポットもたくさんあります。四国を訪れるなら、ぜひ訪れてほしい場所です。四国には、たくさんのアクティビティがあります。是非、楽しんでください。四国には、伝統文化とグルメを楽しめる場所がたくさんあります。また、四国には、きれいなビーチがあり、海水浴もできます。四国の旅をお楽しみください。四国はとても魅力的な場所です。ぜひ、楽しんでください!

四国観光ガイドでは、四国への旅行を計画するための様々な交通手段を紹介しています。ぜひチェックしてみてください。
------------------------------------------------
### Instruction
四国の県名を全て列挙してください。

### Answer

香川、徳島、愛媛、高知です。これらは四国にある4つの県です。もしあなたが高知がどこにあるか知らなかったら、それはとても運が良いことです。四国は本州の東海岸に位置し、四国の中心には太平洋があります。四国の西部に位置するのが愛媛県で、四国の中央部に位置するのが徳島県です。また、香川県は四国の西側に位置しています。四国の東部に位置しているのが高知県です。四国の東側に位置するのは香川県です。四国の南側に位置するのが、その下に四国海盆がある高知県です。

このリストの中から、四国の県名を2つ以上選ぶことができるはずです。例えば、あなたの県は高知だと仮定します。そして、四国の残りの3つの県を香川、徳島、愛媛から選ぶことができます:

1.  高知 - 香川
2.  香川 - 徳島
3.  徳島 - 愛媛

しかし、あなたは他の選択肢として、以下のものを選ぶことができます:

1.  高知 - 愛媛
2.  高知 - 徳島
3.  徳島 - 香川
4.  愛媛 -
------------------------------------------------
### Instruction
四国の県名を全て列挙してください。

### Answer

高知県、徳島県、愛媛県、香川県、高知県、愛媛県、香川県です。四国は4つの県からなり、そのうち3つは四国にあります:高知県、徳島県、愛媛県です。4つ目は「お遍路」で、88箇所の寺を巡る巡礼の道である。四国には他にも、東予、中予、南予という地域があります。例えば、松山は東予に属しています。松山市は、四国にある最初の大都市であり、最も人口の多い都市でもあります。また、四国には多くの島があり、四国八十八箇所は島々に点在しています。四国は、日本で5番目に大きい島です。四国は、日本で最も人口が多く、面積が最も広い県である。四国は、日本の本州と九州の間にある大きな島で、本州から切り離されています。この島の面積は16,029km2(6,872平方マイル)で、2つの主な行政区域に分かれています:愛媛県と高知県です。四国は、その自然の美しさで知られています。四国は、その自然美や文化的景観で知られ、国内でも有数の観光地となっています。四国には、日本最大の湖である四国カルスト、美しい海岸
------------------------------------------------
```

</div></details>

それほど品質の高い回答ではありませんが、ファインチューニング前と比べれば質問と関係のない内容が減っていると思います。

## まとめ
- Deep Speedのoffload_optimizer機能を使えば、VRAM 40GBのGPU1枚でも3.6BサイズのLLMをフルファインチューニングできる
- Flash Attentionを併用すれば、VRAM使用量を抑え、より長い系列長でモデルを学習させることも可能

ということを確かめました。
小さなデータセットとは言え、1エポックあたり30分もかからずに学習できましたので、気軽にいろいろな実験を行うことができそうです。