LLMのファインチューニングのためのツール Axolotl

LLMのファインチューニングのためのツール「Axolotl」の概要をまとめました。

1. Axolotl

「Axolotl」は、LLMのファインチューニングのためのツールです。様々なLLM、データセット形式、アーキテクチャをサポートします。

・「LLaMA」「Pythia」「Falcon」「MPT」などの様々なHuggingFaceモデルを学習
・「Full Finetune」「Lora」「Kalora」Relora」「Guptak」をサポート
・単純なyamlやCLIオーバーライドで構成をカスタマイズ
・様々なデータセット形式をロードしたり、カスタム形式を使用したり、独自のトークン化されたデータセットを使用可能
・「xFormer」「Flash attention」「Rope scaling」「MultiPacking」と統合
・「FSDP」または「DeepSpeed」を介して単一GPUまたは複数GPUで動作
・ローカルまたはクラウド上でDockerを簡単に実行
・結果とオプションでチェックポイントをwandbに記録

2. Axolotlのサポート

3. クイックスタート

・要件: Python以降、Pytorch 2.0以降

# インストール
git clone https://github.com/OpenAccess-AI-Collective/axolotl
cd axolotl

pip3 install packaging
pip3 install -e '.[flash-attn,deepspeed]'
pip3 install -U git+https://github.com/huggingface/peft.git

# ファインチューニング LoRA
accelerate launch -m axolotl.cli.train examples/openllama-3b/lora.yml

# 推論
accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
    --lora_model_dir="./lora-out"

4. インストール

4-1. Docker

docker run --gpus '"all"' --rm -it winglian/axolotl:main-py3.10-cu118-2.0.1

・winglian/axolotl-runpod:main-latest : runpod または direct linkを使用

または、開発用にカレントファイルに対して実行。

docker compose up -d

高度な使用方法は公式ドキュメント参照。

4-2. conda / pip

(1) Python 3.9以降 をインストール
(2) Pytorch (https://pytorch.org/get-started/locally/) をインストール。
(3) Axolotl をインストール

pip3 install packaging
pip3 install -e '.[flash-attn,deepspeed]'

4-3. Windows

WSLまたはDockerを利用。

5. データセット

「Axolotl」は様々なデータセット形式をサポートしています。
以下に使用できる形式の一部を示します (JSONL推奨)。

5-1. instruction形式 (alpacaなど)

「input」はオプションです。

{"instruction": "...", "input": "...", "output": "..."}

5-2. conversations形式 (sharegptなど)

「from」は「human」「gpt」です。

{"conversations": [{"from": "...", "value": "..."}]}

5-3. rawコーパス (completionなど)

{"text": "..."}

他のフォーマットは公式ドキュメント参照。

5-4. カスタムプロンプトの追加方法

指示目的で前処理されたデータセットの場合。

{"instruction": "...", "output": "..."}

この例は YAML 構成で使用できます。

datasets:
  - path: repo
    type:
      system_prompt: ""
      field_system: system
      format: "[INST] {instruction} [/INST]"
      no_input_format: "[INST] {instruction} [/INST]"

5-5. 事前にトークン化されたカスタムデータセットの使用方法

・「type:」を渡さないでください。
・データセット内の列は、正確に input_ids、attention_mask、label である必要があります。

6. コンフィグ

最も重要なオプションは次のとおりです。

6-1. model

base_model: ./llama-7b-hf # ローカルまたはHuggingFaceリポジトリ

6-2. dataset

sequence_len: 2048 # プロンプトの最大トークン長

# HuggingFaceリポジトリ
datasets:
  - path: vicgalle/alpaca-gpt4
    type: alpaca # 先ほどのフォーマット

# 特定のコンフィグ/サブセットを含むHuggingFaceリポジトリ
datasets:
  - path: EleutherAI/pile
    name: enron_emails
    type: completion # 先ほどのフォーマット
    field: text # Optional[str] default: text, field to use for completion data

# 複数の名前付きコンフィグ/サブセットを含むHuggingFaceリポジトリ
datasets:
  - path: bigcode/commitpackft
    name:
      - ruby
      - python
      - typescript
    type: ... # 未実装のカスタム形式

# fastchat conversation
# 'conversation' を参照 
# options: https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
datasets:
  - path: ...
    type: sharegpt
    conversation: chatml

# ローカル
datasets:
  - path: data.jsonl # or json
    ds_type: json # 以下の他のオプションを参照
    type: alpaca

# データセットは分割されているがtrain分割はない
dataset:
  - path: knowrohit07/know_sql
    type: context_qa.load_v2
    train_on_split: validation

6-3. loading

load_in_4bit: true
load_in_8bit: true
bf16: true # require >=ampere
fp16: true
tf32: true # require >=ampere
bfloat16: true # require >=ampere, AMPが不要な場合は、bf16の代わりに使用
float16: true # AMPが不要な場合は fp16の代わりに使用

【注意】 Repo は 4-bit量子化を行いません。

6-4. lora

adapter: lora # qlora またはFullFinetuningを行う場合は空白のままにする
lora_r: 8
lora_alpha: 16
lora_dropout: 0.05
lora_target_modules:
  - q_proj
  - v_proj

他のコンフィグは公式ドキュメントを参照。

7. 学習

7-1. 学習の実行

accelerate launch -m axolotl.cli.train your_config.yml

7-2. データセットの前処理

オプションで、ファインチューニング前に、以下を使用してデータセットを事前にトークン化できます。これは大規模なデータセットに推奨されます。

・Set : 「push_dataset_to_hub: hf_user/repo」を実行して、Huggingface にプッシュ。
・Use : --debug で前処理された例を表示。

python -m axolotl.cli.preprocess your_config.yml

7-3. Multi-GPU

以下は、複数のGPUを使用した学習に axolotl で利用できるオプションです。「FSDP」は損失が不安定になる可能性があるため、現時点では「DeepSpeed」が推奨されるMulti-GPU オプションであることに注意してください。

・Deepspeed
「Deepspeed」は、Multi-GPU システム用の最適化スイートであり、通常 GPU の VRAM に収まるよりもはるかに大きなモデルを学習できるようになります。 deepspeed のさまざまな最適化タイプの詳細については、こちらを参照。

ZeRO stage 1、2、3 用に、いくつかのデフォルトのdeepspeed JSON 構成が提供されています。

deepspeed: deepspeed/zero1.json

accelerate launch -m axolotl.cli.train examples/llama-2/config.py --deepspeed deepspeed/zero1.json

・FSDP
llama FSDP

fsdp:
  - full_shard
  - auto_wrap
fsdp_config:
  fsdp_offload_params: true
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer

7-4. Weights & Biases

・wandb options

wandb_mode:
wandb_project:
wandb_entity:
wandb_watch:
wandb_run_id:
wandb_log_model:

8. 推論

・事前学習したLORA

python -m axolotl.cli.inference examples/your_config.yml --lora_model_dir="./lora-output-dir"

・フルウェイトのファインチューニング

python -m axolotl.cli.inference examples/your_config.yml --base_model="./completed-model"

cat /tmp/prompt.txt | python -m axolotl.cli.inference examples/your_config.yml \
  --base_model="./completed-model" --prompter=None --load_in_8bit=True

以下のようなエラーが表示される場合は、「--sample_packing False」を使用してください。

RuntimeError: stack expects each tensor to be equal size, but got [1, 32, 1, 128] at entry 0 and [1, 32, 8, 128] at entry 1

9. LORAをベースにマージ

python3 -m axolotl.cli.merge_lora examples/your_config.yml --lora_model_dir="./completed-model" --load_in_8bit=False --load_in_4bit=False

CUDAメモリが不足している場合は、システムRAM にマージしてみることができます。

CUDA_VISIBLE_DEVICES="" python3 -m axolotl.cli.merge_lora ...

次回

OpenAI GPT-4/ChatGPT/LangChain 人工知能プログラミング実践入門