使用 DeepSpeed 和 Hugging Face ? Transformer 微調 FLAN-T5 XL/XXL

Scaling Instruction-Finetuned Language Models 論文釋出了 FLAN-T5 模型，它是 T5 模型的增強版。FLAN-T5 由很多各種各樣的任務微調而得，因此，簡單來講，它就是個方方面面都更優的 T5 模型。相同引數量的條件下，FLAN-T5 的效能相比 T5 而言有兩位數的提高。Google 在 Hugging Face 上開源了 5 個 FLAN-T5 的 checkpoints，引數量範圍從 8000 萬 到 110 億。

在之前的一篇博文中，我們已經學習瞭如何 針對聊天對話資料摘要生成任務微調 FLAN-T5，那時我們使用的是 Base (250M 引數) 模型。本文，我們將研究如何將訓練從 Base 擴充套件到 XL (30 億引數) 或 XXL (110 億引數)。

這意味著我們將學習如何利用模型並行、多 GPU 以及 DeepSpeed ZeRO 來微調 FLAN-T5 XL 和 XXL。

除了作為教程的部分之外，我們還跑了一系列實驗，這些實驗資料可以幫助你選擇正確的硬體設定。你可以在 結果和實驗 部分找到詳細資訊。

# install git lfs for pushing artifacts
!sudo apt install git-lfs
# install torch with the correct cuda version, check nvcc --version
!pip install torch --extra-index-url https://download.pytorch.org/whl/cu116 --upgrade
# install Hugging Face Libraries
!pip install "transformers==4.26.0" "datasets==2.9.0" "accelerate==0.16.0" "evaluate==0.4.0" --upgrade
# install deepspeed and ninja for jit compilations of kernels
!pip install "deepspeed==0.8.0" ninja --upgrade
# install additional dependencies needed for training
!pip install rouge-score nltk py7zr tensorboard

處理資料集

與 針對聊天對話的摘要生成任務微調 FLAN-T5 一文中類似，我們需要先準備一個用於微調的資料集。本文，我們將在 CNN Dailymail 資料集 上微調 FLAN-T5-XXL。我們不會贅述如何生成資料集，如果你想了解資料集生成的詳細步驟，請參閱前文提到的 Fine Tune FLAN-T5。

我們定義了一些引數，本文的示例都會基於這些引數，但你可以根據實際需要進行調整。

# 實驗配置
model_id = "google/flan-t5-xxl" # Hugging Face 模型 Id
dataset_id = "cnn_dailymail" # Hugging Face 資料集 Id
dataset_config = "3.0.0" # 資料集版本
save_dataset_path = "data" # 存放處理後資料的本地路徑
text_column = "article" # 輸入文字所屬列
summary_column = "highlights" # 輸出文字所屬列
# 定製指令提示格式
prompt_template = f"Summarize the following news article:\n{{input}}\nSummary:\n"

與 Fine Tune FLAN-T5 不同，這次我們把預處理和訓練分開。這樣我們就可以在非 GPU 例項上執行預處理。我們先對資料集進行預處理 (即分詞) 並將其儲存到磁碟，然後訓練指令碼再從磁碟中載入預處理後的資料集。

from datasets import load_dataset
from transformers import AutoTokenizer
import numpy as np

# Load dataset from the hub
dataset = load_dataset(dataset_id,name=dataset_config)
# Load tokenizer of FLAN-t5-base
tokenizer = AutoTokenizer.from_pretrained(model_id)

print(f"Train dataset size: {len(dataset['train'])}")
print(f"Test dataset size: {len(dataset['test'])}")

# Train dataset size: 287113
# Test dataset size: 11490

我們在配置檔案中定義了一個 prompt_template，其可用於來構建指令提示，以提高我們模型的效能。 prompt_template 有“固定”的開始詞和結束詞，文件放在中間。這意味著我們需要確保 “固定”模板詞 + 文件 總長不超過模型支援的最大序列長度。因此我們需要計算模型支援的最大文件長度，稍後我們會根據它來填充或截斷模板中的文件。

prompt_length = len(tokenizer(prompt_template.format(input=""))["input_ids"])
max_sample_length = tokenizer.model_max_length - prompt_length
print(f"Prompt length: {prompt_length}")
print(f"Max input length: {max_sample_length}")

# Prompt length: 12
# Max input length: 500

Prompt length: 12
Max input length: 500

現在我們知道，模型支援的最大輸入文件長度為 500。除了輸入之外，我們還需要知道最大“目標”序列長度，我們可以透過遍歷資料集中的摘要長度來得到。(程式碼需要執行幾分鐘)

from datasets import concatenate_datasets
import numpy as np

# The maximum total input sequence length after tokenization.
# Sequences longer than this will be truncated, sequences shorter will be padded.
tokenized_inputs = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x[text_column], truncation=True), batched=True, remove_columns=[text_column, summary_column])
max_source_length = max([len(x) for x in tokenized_inputs["input_ids"]])
max_source_length = min(max_source_length, max_sample_length)
print(f"Max source length: {max_source_length}")

# The maximum total sequence length for target text after tokenization.
# Sequences longer than this will be truncated, sequences shorter will be padded."
tokenized_targets = concatenate_datasets([dataset["train"], dataset["test"]]).map(lambda x: tokenizer(x[summary_column], truncation=True), batched=True, remove_columns=[text_column, summary_column])
target_lenghts = [len(x) for x in tokenized_targets["input_ids"]]
# use 95th percentile as max target length
max_target_length = int(np.percentile(target_lenghts, 95))
print(f"Max target length: {max_target_length}")

  0%|          | 0/299 [00:00<?, ?ba/s]
Max source length: 500
    
  0%|          | 0/299 [00:00<?, ?ba/s]
Max target length: 129

現在一切準備就緒，可以處理資料集了。

import os

def preprocess_function(sample, padding="max_length"):
    # created prompted input
    inputs = [prompt_template.format(input=item) for item in sample[text_column]]

    # tokenize inputs
    model_inputs = tokenizer(inputs, max_length=tokenizer.model_max_length, padding=padding, truncation=True)

    # Tokenize targets with the `text_target` keyword argument
    labels = tokenizer(text_target=sample[summary_column], max_length=max_target_length, padding=padding, truncation=True)

    # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
    # padding in the loss.
    if padding == "max_length":
        labels["input_ids"] = [
            [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
        ]

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

# process dataset
tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=list(dataset["train"].features))

# save dataset to disk
tokenized_dataset["train"].save_to_disk(os.path.join(save_dataset_path,"train"))
tokenized_dataset["test"].save_to_disk(os.path.join(save_dataset_path,"eval"))

使用 deepspeed 微調模型

準備完畢！我們現在可以開始訓練模型了！如前所述，我們將使用整合了 DeepSpeed 的 Hugging Face Trainer。因此我們需要建立一個 deespeed_config.json。DeepSpeed 配置 定義了要使用的 ZeRO 策略以及是否要使用混合精度訓練等配置項。 Hugging Face Trainer 允許我們從 deepspeed_config.json 中的 TrainingArguments 繼承相關配置以避免重複設定，檢視 文件瞭解更多資訊。

我們建立了 4 組 deepspeed 配置檔案用於實驗，包括 CPU 解除安裝 和 混合精度:

• ds_flan_t5_z3_config.json
• ds_flan_t5_z3_config_bf16.json
• ds_flan_t5_z3_offload.json
• ds_flan_t5_z3_offload_bf16.json

你可以根據你的執行環境選擇，例如如果在 NVIDIA V100s 上執行，你就不能使用帶 bf16 的配置，因為 V100 不支援 bfloat16 資料型別。

在微調 T5 模型時，不能使用 fp16，因為它會導致精度溢位問題，參見問題 #4586，#10830，和拉取請求 #10956

如開頭所述，我們使用的是 p4dn.24xlarge AWS EC2 例項，該例項包含 8 張視訊記憶體為 40GB 的 NVIDIA A100。這意味著我們可以使用 bf16，它將減少近一半的模型視訊記憶體佔用，使我們能夠在不解除安裝的情況下高效訓練。

我們將使用 ds_flan_t5_z3_config_bf16.json。如果你不想用 auto 值，可以檢視 文件。

{
  "bf16": {
    "enabled": "auto"
  },
  "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": 3,
    "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": false
  },
  "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
}

現在，該訓練指令碼上場了。我們根據 Fine Tune FLAN-T5 準備了一個 run_seq2seq_deepspeed.py 訓練指令碼，它支援我們配置 deepspeed 和其他超引數，包括 google/flan-t5-xxl 的模型 ID。

我們使用 deepspeed 啟動器觸發訓練，輸入給啟動器的引數包括 GPU 數量、deepspeed 配置及其它超引數 (如 google/flan-t5-xxl 的模型 ID)。

!deepspeed --num_gpus=8 scripts/run_seq2seq_deepspeed.py \
    --model_id $model_id \
    --dataset_path $save_dataset_path \
    --epochs 3 \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 8 \
    --generation_max_length $max_target_length \
    --lr 1e-4 \
    --deepspeed configs/ds_flan_t5_z3_config_bf16.json

huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
    - Avoid using `tokenizers` before the fork if possible
    - Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
deepspeed --num_gpus=8 scripts/run_seq2seq_deepspeed.py --model_id google/flan-t5-xxl --dataset_path data --epochs 3 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --generation_max_length 129 --lr 1e-4 --deepspeed configs/ds_flan_t5_z3_config_bf16.json

DeepSpeed 先將模型載入到 CPU 上，然後將其拆分到 8 張 A100 上然後開始訓練。使用 CNN Dailymail 資料集 進行訓練大約需要 10 個小時，費用約為 322 美元。

結果與實驗

為了更好地瞭解硬體要求，我們對 FLAN-T5 XL 和 XXL 進行了一系列實驗，以幫助我們評估和了解硬體需求以及訓練這些模型的成本。

下表列出了實驗和相關設定的詳細資訊。

資料集: "cnn_dailymail"

• 訓練樣本數: 287113
• 驗證樣本數: 13368

超參:

• epochs: 3
• 學習率: 1e-4

執行環境設定:

• 4x V100 16GB: p3.8xlarge
• 4x A10G 24GB: g5.24xlarge
• 8x V100 16GB: p3.16xlarge
• 8x A100 40GB: p4dn.24xlarge

我們可以看到 bf16 與 fp32 相比具有顯著優勢。FLAN-T5-XXL 能放進 4 張 A10G (24GB)，但放不進 8 張 V100 16GB。

我們的實驗還表明，如果模型可以無需解除安裝同時以 batch size 大於 4 的配置跑在 GPU 上，其速度將比解除安裝模型和減小 batch size 的配置快約 2 倍且更具成本效益。

英文原文: https://www.philschmid.de/fine-tune-flan-t5-deepspeed

原文作者: Philipp Schmid

譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。

審校、排版: zhongdongy (阿東)