Llama3.1 70B を AWS P4d インスタンスで微調整

Meta-Llama-3.1-70B-Instruct を LoRA や量子化を用いずに  Amazon EC2 P4d インスタンスで教師あり微調整しました。手順を紹介します。

https://huggingface.co/meta-llama/Meta-Llama-3.1-70B-Instruct

P4d インスタンスの p4d.24xlarge は、GPU あたり 40GB メモリの NVIDIA A100 Tensor Core GPU を 8 基搭載し、GPU メモリの合計は 320 GB です。また、1152 GiB のインスタンスメモリと 8 TB NVMe ベースの SSD  インスタンスストレージを備えています。

https://aws.amazon.com/jp/ec2/instance-types/p4/

微調整には推定 500GB 超の GPU メモリが必要なので Transformers の DeepSpeed 統合で GPU メモリの不足を補います。

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

https://github.com/microsoft/DeepSpeed


# 環境
- Amazon EC2 インスタンス
  インスタンスタイプ: p4d.24xlarge
  リージョン: ap-northeast-1
  AMI: Deep Learning OSS Nvidia Driver AMI GPU PyTorch 2.3.0 (Ubuntu 20.04) 20240611

```shell
$ python --version
Python 3.11.9
$ python
>>> import transformers, deepspeed
>>> transformers.__version__
'4.43.3'
>>> deepspeed.__version__
'0.14.4'
```


# データセット
ござるデータセットを使わせていただきました。

https://huggingface.co/datasets/bbz662bbz/databricks-dolly-15k-ja-gozarinnemon


# 準備
Llama3.1 は利用申請が必要です。Hugging Face Hub の　[Llama-3.1 モデルカード](https://huggingface.co/meta-llama/Meta-Llama-3.1-70B-Instruct)から申請を行い [ユーザーアクセストークン](https://huggingface.co/docs/hub/security-tokens)を発行します。

EC2 インスタンスを開始しユーザ ubuntu にログインします。

DeepSpeed に必要な libaio-dev をインストール
```shell-session
sudo apt-get update
sudo apt-get install libaio-dev
```

Python 仮想環境を開始
```shell-session
source activate pytorch
```

Pythonパッケージをインストール
```shell
pip install huggingface_hub[cli]
pip install transformers[deepspeed] trl accelerate datasets
```

HuggingFace CLI にログイン
```shell
HF_TOKEN="hf_XXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
huggingface-cli login --token $HF_TOKEN
```

Huggin Face Hub からモデルをダウンロード
```shell
MODEL_ID=meta-llama/Meta-Llama-3.1-70B-Instruct
huggingface-cli download $MODEL_ID --local-dir /opt/dlami/nvme/models/$MODEL_ID
```

# DeepSpeed 構成
DeepSpeed の ZeRO は、CPU や NVMe へオフロードすることで GPU メモリの不足を補うことができます。

[メモリ要件を見積り](https://huggingface.co/docs/transformers/ja/main_classes/deepspeed#memory-requirements)すると CPU オフロードにはおよそ 1600 GB のインスタンスメモリが必要とのこと。p4d.24xlarge のインスタンスメモリ 1152 GiB はこの要件を満たさないため今回は、ZeRO-Infinity で NVMe へオフロードします。オフロードの対象は optimizer のみとし parameter はオフロードしません。

DeepSpeed 構成ファイルを作成
```json:zero3_train.json
{
    "bf16": {
        "enabled": "auto"
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "nvme",
            "nvme_path": "/opt/dlami/nvme/zero_param",
            "pin_memory": true,
            "buffer_count": 4,
            "fast_init": false
        },
        "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
    },
    "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
}
```


# 微調整
Python スクリプトファイルを作成
```train.py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM

args = SFTConfig(
    output_dir="/opt/dlami/nvme/output",
    overwrite_output_dir=True,
    num_train_epochs=3,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={'use_reentrant':False},
    learning_rate=8e-6,
    weight_decay=0.0,
    bf16=True,
    fp16=False,
    deepspeed='./zero3_train.json',  # DeepSpeed 構成ファイル
    max_seq_length=8192,
    packing=False,
    save_strategy='no',
)


dataset = load_dataset("bbz662bbz/databricks-dolly-15k-ja-gozarinnemon")

model_id = "/opt/dlami/nvme/models/meta-llama/Meta-Llama-3.1-70B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_id)

chat_template = """
user


{0}

assistant


{1}
"""

response_template = """
assistant


"""

# プロンプトテンプレート適用
def formatting_func(rows):
    output_texts = []
    for i, instruction in enumerate(rows['instruction']):
        user_content = rows['input'][i] + instruction
        text = chat_template.format(user_content, rows['output'][i])
        output_texts.append(text)
    return output_texts

tokenizer.pad_token = "
"

collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
)
model.config.use_cache = False

trainer = SFTTrainer(
    model,
    tokenizer=tokenizer,
    train_dataset=dataset["train"].select(range(64)),
    formatting_func=formatting_func,
    data_collator=collator,
    args=args,
)

trainer.train()

save_dir = "/opt/dlami/nvme/models/llama-3-gozaru"
trainer.save_model(save_dir)
tokenizer.save_pretrained(save_dir)
```

事前トレーニングされたモデルを使用する場合、`AutoModelForCausalLM.from_pretrained()` を呼び出す前に `SFTConfig` オブジェクトを作成する必要があります（参照 [Constructing Massive Models](https://huggingface.co/docs/transformers/ja/main_classes/deepspeed#constructing-massive-models)）。


deepspeed コマンドの引数に `train.py` を指定し実行
```shell
LIBRARY_PATH=$LD_LIBRARY_PATH deepspeed train.py
```

DeepSpeed 実行時の静的リンクでリンクエラー `/usr/bin/ld: cannot find -lcurand` が発生したためここでは `LIBRARY_PATH` を指定しています。

微調整に要した時間はおよそ 35 分でした。

time コマンドによる計測結果
```
real    35m6.142s
user    388m46.755s
sys     101m36.467s
```

# 推論
微調整の効果を確認します。

```python:test.py
import transformers
import torch

pipeline = transformers.pipeline(
    "text-generation",
    model="/opt/dlami/nvme/models/llama-3-gozaru",
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
)

messages = [
    {"role": "user", "content": "日本の首都は？"},
]

outputs = pipeline(
    messages,
    max_new_tokens=256,
)

print(outputs[0]["generated_text"][-1])
```

実行
```
python test.py
```

結果
```
{'role': 'assistant', 'content': '我、りんえもんは思う。 日本の首都は東京でござる。知らんけど。'}
```

「我、りんえもんは思う」から始まり「知らんけど」で終わっています。期待通りです。

以上です。最後までお読みいただきありがとうございます。