LLM 大模型學習必知必會系列(七):掌握分散式訓練與LoRA/LISA微調:打造高效能大模型的秘訣進階實戰指南

汀、人工智能發表於2024-05-28

LLM 大模型學習必知必會系列(七):掌握分散式訓練與LoRA/LISA微調:打造高效能大模型的秘訣進階實戰指南

1.微調(Supervised Finetuning)

指令微調階段使用了已標註資料。這個階段訓練的資料集數量不會像預訓練階段那麼大,最多可以達到幾千萬條,最少可以達到幾百條到幾千條。指令微調可以將預訓練的知識“湧現”出來,進行其他型別的任務,如問答型別的任務。一般指令微調階段對於在具體行業上的應用是必要的,但指令微調階段一般不能灌注進去新知識,而是將已有知識的能力以某類任務的形式展現出來。

指令微調任務有多種場景,比較常用的有:

  • 風格化:特定的問答正規化
  • 自我認知:自我認知改變
  • 能力增強:模型本身能力不夠,對具體行業的資料理解不良
  • Agent:支援Agent能力,比如程式編寫、API呼叫等

上述只是舉了幾個例子,一般來說距離使用者最近的訓練方式就是指令微調。

一般來說,LLM中指的base模型是指經過了預訓練(以及進行了一部分通用指令的微調)的模型。Chat模型是經過了大量通用資料微調和人類對齊訓練的模型。

如何選擇base模型和chat模型進行微調呢?

  • 資料量較少的時候(比如小於1w條)建議使用chat模型微調
  • 資料量較多、資料較為全面的時候,建議使用base模型微調

當然,如果硬體允許,建議兩個模型都進行嘗試,選擇效果較好的。需要注意的是,chat模型有其獨特的輸入格式,在微調時一定要遵循。base模型的輸入格式一般比較簡單(但也需要遵守該格式),而且一般該格式不支援多輪資料集。

如果需要用base模型訓練多輪對話,一般需要使用一個支援多輪對話的template。在SWIFT中,可以指定為default,在訓練時只需要指定--template_type default即可。

  • 重要概念
  1. loss 代表模型求解的y和實際的y值的差異。該值會進行loss.backward(),這個方法會求解梯度,並將對應梯度值記錄在每個引數上

    loss可以理解為根據模型計算出來的值和正確值的偏差(也就是殘差)。 例如,迴歸任務中計算的值是1.0,而實際的值應當為2.0,那麼loss為2.0-1.0=1.0。上述loss型別為MAE,除此外,還有MSE,Hinge等各類loss。一般分類任務的loss為交叉熵(Cross-Entropy),這也是目前LLM最常用的loss。

    loss計算出來後(這個過程也就是forward,即前向推理),經過backward過程即可計算出梯度。

    梯度:光滑的曲面上導數變化最大的方向

    loss可以經過PyTorch的loss.backward()將每個運算元、每個步驟的梯度都計算出來(複雜微分方程的鏈式求導過程),當有了梯度後,可以將引數往負梯度方向更新,學習率(lr)就是這時候起作用的,由於直接加上負梯度太大,可能直接產生震盪,即值從一個點瞬間跑到了曲線上的另一個點,導致在這兩點反覆震盪不收斂,因此乘以一個lr,讓loss一點點下降。

  2. epoch 代表對資料集訓練多少輪次

  3. iter 對輸入資料的每次forward+backward代表一個iter

  4. batch_size 批處理大小。在一次前向推理中,同時處理多少行資料。由於同一批資料會並行求解梯度,因此batch_size越大,梯度越穩定。在SFT時較為合適的梯度一般選擇為16/32/64等值

    1. batch_size越大,平行計算消耗的視訊記憶體越高。因此在低視訊記憶體情況下,可以選用batch_size=1,gradient_accumulation_steps=16。訓練會在iter%gradient_accumulation_steps==0時集中進行一次引數更新。在iter%gradient_accumulation_steps!=0時,會將梯度值不斷累加到引數上,這樣就相當於將batch_size擴大了gradient_accumulation_steps倍
  5. learning_rate 學習率 訓練將負梯度值乘以該值加到原引數上。換句話說,每次只將引數更新一個小幅度,避免向錯誤的更新方向移動太多。

    一般LoRA的學習率可以比全引數訓練的學習率稍高一點,因為全引數訓練會完全重置所有引數,訓練時需要學習率更低。
    LLM訓練的學習率一般設定在1e-4~1e-5不等

  6. max_length 輸入句子的最大長度。比如設定為4096,那麼句子加答案轉換為token後最大長度為max_length。這個值會影響視訊記憶體佔用,需要按照自己的實際需求設定。

    1. 當batch_size大於1時,意味著不同句子的長度可能不同。data_collator的作用就是按照固定max_length或者batch中的最大長度對其他句子的token進行補齊。補齊的部分不參與模型的loss計算,但仍然會佔用計算量
  7. flash_attention flash attention是一種針對attention結構高效計算的元件,該元件主要原理利用了顯示卡的快取記憶體。flash attention會節省約20%~40%訓練視訊記憶體並提高訓練速度,對訓練精度沒有不良影響。在顯示卡支援的情況下建議開啟。

  8. optimizer

    optimizer是深度學習中的最佳化器,負責將負梯度值累加到原來需要更新的引數上,類似於:

    Vanilla SGD

    weights = weights - learning_rate * grad

    實際的原理會比較複雜,比如常用的AdamW實際上是一個複雜的滑動平均的演算法。

  9. lr_scheduler

    一般來說,訓練各個階段的學習率是不一樣的,有時候需要越來越小(因為訓練到最後需要更精細的調節),有時候需要先有個warmup(先將lr從0增大到指定值,再慢慢減小),lr_scheduler就是用來動態調整lr使用的元件。

  10. gradient_checkpointing 梯度檢查點。該方法的原理是將訓練時的中間變數在前向過程中暫時丟棄,並在後向過程中重新計算。該方法可以有效節省訓練視訊記憶體,但屬於時間換空間的做法,因此訓練時間會變長。對視訊記憶體的節省可以達到30%-70%不等。訓練速度會減慢20%-40%不等。

訓練有很多超引數,它們的含義和設定技巧可以參考這裡

2.分散式訓練(Distributed Training)

由於較大模型可能在單張顯示卡上視訊記憶體溢位,或者訓練速度不夠,因此單機多卡或多機多卡訓練是必要的。在訓練過程中的分散式訓練有以下幾種模式:

  • DDP 分散式資料並行。將訓練集的資料分段拆分到不同的程序中,這種訓練方式相當於增加了batch_size。比如四個程序,每個程序batch_size=1,則總體batch_size=4。在計算梯度時,torch框架會自動將四個程序的梯度進行累加平均。該方法會提高訓練速度,但如果模型在單張顯示卡上視訊記憶體溢位,DDP方式也無法執行。

  • MP 模型並行。模型並行分為多種方式,如tensor並行、device_map、流水線並行、FSDP等。

    • tensor並行:將矩陣拆分到多張顯示卡上,比如,將一個2048x2048的矩陣,拆分為兩個1024x2048的矩陣,在前向推理時在顯示卡間通訊,完成一次推理,這樣一個模型的視訊記憶體需求就被平均拆分到兩個顯示卡上。tensor並行最知名的框架是Megatron。

    • device_map並行:自動計算如何將模型拆分到多個顯示卡上。比如一個模型按照順序分為embedder、layer095、output,device_map可能將這些引數均分到兩張顯示卡上,比如embedder、layer048分配到顯示卡1上,layer49~95、output分配到顯示卡2上。相比Megatron,device_map方式較為低效,因為使用該方法訓練或推理時,顯示卡1計算時顯示卡2是空閒的,計算效率較低;而Megatron是同時使用兩個顯示卡計算,效率較高

    • 流水線並行:類似於device_map,將模型按照layer拆分到不同顯示卡上

    • FSDP,在講FSDPqian需要先講解DeepSpeed的ZeRO最佳化方式

      • ZeRO-1:類似DDP,但是將Optimizer的state均分維護到不同的程序中,每次更新引數後對所有程序的引數進行同步更新

      • ZeRO-2:在ZeRO-1的基礎上,將不同層的梯度值均分維護到不同的程序中,每次每個程序同步梯度後更新自己負責的梯度對應的引數部分,並在更新後對所有的程序的引數進行同步

      • ZeRO-3:在ZeRO-2的基礎上,將不同層的模型引數也均分到不同的程序中。每個程序在計算某層結果時,從其他程序中獲得對應的層的引數,計算完後拋棄該層引數;backward時,也從其他程序獲得對應層的引數並同步梯度資訊,計算完後拋棄該層引數。這樣每個程序就在僅儲存某些層的引數的條件下完成了資料平行計算

      • FSDP就是ZeRO-3的並行策略

3.LoRA

LoRA是一個非常重要的可調優結構,簡單來說,就是增加了一個額外可訓練部分,比如原來的Linear的矩陣是MxN維,增加一個LoRA,該LoRA會包含兩個引數量較少的矩陣:Mxd, dxN,這兩個矩陣相乘後仍然是MxN維的,訓練時原MxN矩陣凍結,只訓練LoRA的兩個矩陣,引數量就會大大減少。

為什麼模型本身的矩陣不使用這種形式?

一般大規模矩陣的非零特徵值數量會遠遠小於矩陣的維度,這個非零特徵值的數量叫做矩陣的秩(rank),秩決定了這個矩陣如何影響被乘的向量,為0或較小的特徵值對傳入tensor的影響也比較小,丟棄這些資訊對精度的影響不大。

一個模型包含了多個大矩陣,這些大矩陣的秩不相等而且難以預測,因此不能對原模型應用LoRA,但在sft時使用LoRA相對安全,雖然有精度損失,但可以使一個大模型在一個消費級顯示卡上進行訓練。

也就是說,LoRA的原理是假設所有矩陣的秩都是d,進行了一定的有失真壓縮。基於LoRA也有很多升級版技術,如AdaLoRA、SoRA等,這些元件方案都是基於LoRA,對不同運算元的LoRA的rank進行動態調節以達到更好的效果。

LoRA目前已經是訓練SD模型和LLM模型的最常用技術。LoRA的weights也非常小,只有幾十兆,因此載入和使用都非常方便,且LoRA本身可以合併回原模型,推理時可以做到相容原模型結構。

如果涉及到對模型的知識編輯,比如自我認知任務,LoRA的目標module一般需要設定為ALL,因為MLP層對模型的知識獲取是至關重要的,需要參與訓練過程。

3.1 訓練過程

在前序的文章中,我們講述瞭如何進行資料的前處理。結合上面講解的基本概念,我們就可以執行一個完整的訓練過程。

pip install ms-swift -U

安裝好SWIFT後,可以直接啟動介面執行訓練和推理:

swift web-ui
  • 官方連結:https://modelscope.cn/studios/iic/Scalable-lightWeight-Infrastructure-for-Fine-Tuning/summary

在框架中,一個最小的訓練過程程式碼如下:

#Experimental environment: A10, 3090, V100, ...
#20GB GPU memory
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

import torch

from swift.llm import (
    DatasetName, InferArguments, ModelType, SftArguments,
    infer_main, sft_main, app_ui_main, merge_lora_main
)

model_type = ModelType.qwen_1_8b
sft_args = SftArguments(
    model_type=model_type,
    train_dataset_sample=2000,
    dataset=[DatasetName.blossom_math_zh],
    output_dir='output')
result = sft_main(sft_args)
best_model_checkpoint = result['best_model_checkpoint']
print(f'best_model_checkpoint: {best_model_checkpoint}')
torch.cuda.empty_cache()

infer_args = InferArguments(
    ckpt_dir=best_model_checkpoint,
    load_dataset_config=True,
    show_dataset_sample=10)
#merge_lora_main(infer_args)
result = infer_main(infer_args)
torch.cuda.empty_cache()

app_ui_main(infer_args)

3.2 自定義一個訓練過程

上面我們構建了一個最小的訓練和推理流程。大多數時候開發者需要自定義一個訓練流程和對應的資料集。在這種情況可以參考下面的步驟:

  1. 選擇一個啟動訓練的方式,介面方式可以使用上述的web-ui命令(swift web-ui),命令列方式可以參考:
CUDA_VISIBLE_DEVICES=0 \
swift sft \
    --model_id_or_path qwen/Qwen-7B-Chat \
    --dataset blossom-math-zh \
    --output_dir output \

注意命令列具有很多可調節引數,可以檢視文件來檢視這些引數的具體意義。

​ 如果想要了解訓練流程可以檢視訓練程式碼

​ 瞭解超引數的拼接和處理可以檢視超引數的處理程式碼

​ 瞭解所有支援的模板可以檢視模板的拼接

  1. 選擇一個需要參與訓練的模型,可以參考支援的模型列表
  2. 選擇一個或若干個自己的資料集參與訓練,注意這些資料集有一定的格式要求。或者也可以使用一個自己的模型訓練,只需要註冊自定義模型即可。
CUDA_VISIBLE_DEVICES=0 \
swift sft \
    --model_id_or_path qwen/Qwen-7B-Chat \
    --dataset blossom-math-zh \
    --output_dir output \
    --custom_train_dataset_path xxx.jsonl zzz.jsonl \
    --custom_val_dataset_path yyy.jsonl aaa.jsonl \

4.LISA

  • 背景介紹

LISA是Layerwise Importance Sampling for Memory-Efficient Large Language Model Fine-Tuning的簡寫。這個技術可以把全參訓練的視訊記憶體使用降低到之前的三分之一左右,而使用的技術方法卻是非常簡單的。例如,全參訓練一個7b模型大約需要80G視訊記憶體(相當於一張完整的A100顯示卡),但使用LISA訓練後卻可以使視訊記憶體降低到30G左右,這使得使用40G A100顯示卡甚至是24G A10或者RTX 3090成為可能,且它的視訊記憶體佔用更低、訓練速度更快。

論文地址:https://arxiv.org/abs/2403.17919

4.1 技術解析

LISA使用的技術原理相對簡單。作者首先對LoRA訓練和全參訓練每個layer不同step時的L2範數的平均和進行了對比,結果如下:

作者訓練了GPT2和LLaMA-2-7B兩個模型,發現它們自身不同layers的parameters的LoRA訓練和全參訓練的L2範數不同,可以間接說明LoRA訓練中由於低秩矩陣的存在,因此其引數更新的重點和全引數更新重點完全不同。可以看出,在權重更新時,除底層和頂層外其它層的L2範數都較小,因此作者假設可以在全引數訓練時透過凍結大部分層的引數來模擬LoRA更新的行為,使其最後的引數迭代範數達到類似的效果。

完整的演算法迭代可以用下圖表示:

4.2 實驗

在官方實驗中,作者對比了LISA和LoRA訓練以及全引數的視訊記憶體佔用:

![img]

可以看到LISA的視訊記憶體佔用要小於LoRA。在訓練速度上面:

官方實驗結果,LISA的Forward和Backward時間要顯著短於LoRA訓練。在訓練方面,作者進行不同尺寸的微調和大規模微調,均證明了LISA的效果要強於LoRA:

如何調節LISA的超引數呢?LISA的超引數包含兩個值:

  1. LISA取樣的有效層數γ
  2. LISA的更新頻率K

消融實驗對這兩個值的對比如下:

可以看到LISA的效能在γ=8,取樣頻率K=5的時候達到最好。作者也證明,LISA對於不同的隨機種子的魯棒性很強,在此不列舉表格。

4.3 本次實驗

為了驗證LISA在實際測試中的效果,我們對LISA進行了一定的實驗。我們使用了魔搭社群提供的SWIFT框架(https://github.com/modelscope/swift),該框架支援LISA訓練方式,且支援LoRA等通用訓練方式。我們可以設定LISA的兩個值:

  • lisa_activated_layers 上文的γ
  • lisa_step_interval 上文的K

我們使用如下命令進行訓練:

#pip install ms-swift -U
sft.py \
 --model_type qwen-7b-chat \
 --dataset ms-agent \
 --train_dataset_mix_ratio 2.0 \
 --batch_size 1 \
 --max_length 2048 \
 --use_loss_scale True \
 --gradient_accumulation_steps 16 \
 --learning_rate 5e-05 \
 --use_flash_attn True \
 --eval_steps 2000 \
 --save_steps 2000 \
 --train_dataset_sample -1 \
 --val_dataset_sample 5000 \
 --num_train_epochs 2 \
 --check_dataset_strategy none \
 --gradient_checkpointing True \
 --weight_decay 0.01 \
 --warmup_ratio 0.03 \
 --save_total_limit 2 \
 --logging_steps 10 \
 --sft_type full \
 --lisa_activated_layers 2 \
 --lisa_step_interval 20

同時,我們將--lisa_activated_layers置為0,進行全引數訓練,並且使用r=8進行了LoRA訓練,得到的效果如下:

exp_name model_type dataset tuner tuner_params trainable params(M) flash_attn gradient_checkpointing hypers memory train speed(samples/s) train_loss eval_loss gsm8k weighted acc arc weighted acc ceval weighted acc
full qwen-7b-chat ms-agent full 7721.3245(100.0000%) True True lr=5e-05/epoch=2 73.53GiB 1.43(87543 samples/61022.97 seconds) 0.54 0.95 0.343 0.536 0.495
full+lisa_2 qwen-7b-chat ms-agent full lisa_activated_layers=2/lisa_step_interval=20 7721.3245(100.0000%) True True lr=5e-05/epoch=2 31.11GiB 2.66(76837 samples/28881.28 seconds) 0.62 1.06 0.349 0.653 0.592
full+lisa_4 qwen-7b-chat ms-agent full lisa_activated_layers=4/lisa_step_interval=20 7721.3245(100.0000%) True True lr=5e-05/epoch=2 31.87GiB 2.63(76837 samples/29215.15 seconds) 0.63 1.06 0.377 0.656 0.607
lora qwen-7b-chat ms-agent lora rank=8/target=ALL/alpha=32/lr_ratio=None/use_rslora=False/use_dora=False 17.8913(0.2312%) True True lr=5e-05/epoch=2 32.35GiB 0.95(87543 samples/91974.29 seconds) 0.53 1.01 0.462 0.676 0.304

從我們的實驗中可以看到下面的結論:

  1. 在視訊記憶體佔用中,全引數幾乎是其他輕量訓練方式視訊記憶體佔用的兩倍,但是在loss中也是最低的,這說明全引數在模型訓練的基礎指標中仍然是最優的;
  2. LISA的視訊記憶體使用比r=8(這是個常用值)的視訊記憶體佔用要低,其中lisa_activated_layers越低視訊記憶體越低
  3. 訓練速度上LISA的訓練速度也比LoRA要快一些,並且該指標也受到lisa_activated_layers的影響
  4. 在評估指標上,LoRA更為優秀,然而評估指標受到資料集的強烈影響,由於訓練主要內容是Agent資料集,因此說明LoRA在防止災難性遺忘上具有一定的優勢

LISA lisa_activated_layers=2 訓練的loss

LoRA r=8 訓練的loss

可以觀察到LISA的訓練loss較LoRA曲線更為抖動一些,猜測可能是LISA隨機挑選layer進行反向傳播的隨機性造成的。

結論

可以看到LISA作為2024年的新晉tuner,使用一個非常簡單的方式做到了部分資料集的SOTA,同時視訊記憶體使用和訓練速度也是很優秀的,且沒有額外的使用條件。然而LISA仍然存在著一些可以分析討論的問題,比如:是否可以透過引數範數或者引數矩陣特徵值判斷哪些layers應該被反向傳播?或者是否可以在更細粒度上(qkv/mlp/layernorm)層面上控制反向傳播?如果有做過實驗的同學歡迎留言討論。

相關文章