機器學習 | 從0開始大模型之模型DPO訓練
1、為什么需要DPO
Rafailov等人在2023年發(fā)表了一篇論文《Direct Preference Optimization: Your Language Model is Secretly a Reward Model》,該論文提出了一種新的訓練方法,稱為直接偏好優(yōu)化(DPO),該論文介紹:
雖然大規(guī)模無監(jiān)督語言模型 (LM) 可以學習廣泛的世界知識和一些推理技能,但由于其訓練完全無監(jiān)督,因此很難精確控制其行為。
現(xiàn)有的獲得這種可控性的方法是收集模型生成相對質(zhì)量的人類標簽,并微調(diào)無監(jiān)督語言模型以符合這些偏好,通常使用從人類反饋中進行強化學習 (RLHF)。
然而,RLHF 是一個復雜且通常不穩(wěn)定的過程,首先要擬合一個反映人類偏好的獎勵模型,然后使用強化學習微調(diào)大型無監(jiān)督語言模型以最大化這個估計的獎勵,而不會偏離原始模型太遠。
在該論文中,利用獎勵函數(shù)和最優(yōu)策略之間的映射來表明,這個受約束的獎勵最大化問題可以通過一個階段的策略訓練進行精確優(yōu)化,本質(zhì)上是解決人類偏好數(shù)據(jù)的分類問題。
由此產(chǎn)生的算法,稱之為直接偏好優(yōu)化 (DPO),穩(wěn)定、高效且計算量小,無需擬合獎勵模型、在微調(diào)期間從 LM 中采樣或執(zhí)行重大超參數(shù)調(diào)整。
由此可見,DPO 主要解決RLHF不穩(wěn)定的問題,直接使用人類偏好數(shù)據(jù)訓練模型。
2、DPO的訓練原理
DPO 的訓練原理如下圖所示(出自原論文):
DPO
主要包括兩個步驟:
- 數(shù)據(jù)收集:收集一個偏好數(shù)據(jù)集,其中包含給定提示的生成結(jié)果的正負選擇對;
- 優(yōu)化:直接最大化 DPO 損失的對數(shù)似然函數(shù),該損失函數(shù)是偏好數(shù)據(jù)集上的交叉熵損失和模型生成結(jié)果的對數(shù)似然性之間的加權平均值;
具體公式推導可以參考這篇博客:https://www.cnblogs.com/lemonzhang/p/17910358.html。
3、DPO的代碼實現(xiàn)
3.1 收集數(shù)據(jù)
DPO 訓練器對數(shù)據(jù)集的格式有具體的要求,包括三個部分:
- 提示(prompt):提示的格式為:prompt: 文本;
- 選中(chosen):選中文本的格式為:chosen: 文本;
- 拒絕(rejected):拒絕選中文本的格式為:rejected: 文本;
- 示例:
{
"id": 0,
"prompt": "最近,馬其頓流行歌手托瑟·普羅埃斯基在一場巡回展覽中被紀念。",
"chosen": "\"最近,馬其頓流行歌手托斯·普羅埃斯基在一場展覽中被人們銘記。\" 讓我?guī)湍憷斫?。這意味著他們舉辦了一場特別的展示(我們稱之為展覽),以幫助人們記住馬其頓這個地方的歌手托斯·普羅埃斯基有多特別。這有點像翻看你的相冊來回憶快樂時光,但規(guī)模更大,更專注于他成名后的工作和生活。",
"rejected": "哦!你想知道這句話的意思嗎?好的,讓我告訴你!??\n\n所以,“Son d?nem”在英語中意思是“Last period”。而“Makedon”在英語中是“Macedonian”。而“pop ?ark?c?s?”在英語中意思是“pop singer”!??\n\n所以,句子“Son d?nem Makedon pop ?ark?c?s? Tose Proeski gezer sergide an?l?yor”的意思是“Last period Macedonian pop singer Tose Proeski正在舞臺上表演?!??\n\n希望這有幫助,小伙伴!??"
}
DPO的數(shù)據(jù)可以搜索huggingface的DPO數(shù)據(jù)集,地址為:https://huggingface.co/datasets?sort=trending&search=dpo 。
比如 https://huggingface.co/datasets/Anthropic/hh-rlhf 的數(shù)據(jù)集如下:
hh-rlhf
3.2 TRL
引入 TRL 庫,支持 DPO 訓練器,訓練樣例代碼:
training_args = DPOConfig(
beta=0.1,
)
dpo_trainer = DPOTrainer(
model,
ref_model,
args=training_args,
train_dataset=train_dataset,
tokenizer=tokenizer, # for visual language models, use tokenizer=processor instead
)
dpo_trainer.train()
dpo_trainer.save_model()
如上訓練默認是保存 safetensors 格式的模型,如果想保存 pytorch 格式的模型, 可以改為如下代碼:
training_args = DPOConfig(
beta=0.1,
save_safetensors=False, // 設置為False,改為保存為pytorch格式的模型
)
dpo_trainer = DPOTrainer(
model,
ref_model,
args=training_args,
train_dataset=train_dataset,
tokenizer=tokenizer, # for visual language models, use tokenizer=processor instead
)
dpo_trainer.train()
dpo_trainer.save_model(
output_dir=f"./out/dpo_sft_xxx.pth"
)
3.3 訓練
Transformer的代碼和前面的一樣,可以參考預訓練的代碼,如下就是初始化模型和 DPO 訓練的代碼:
def init_model():
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
AutoConfig.register(MyPretrainConfig.model_type, MyPretrainConfig)
AutoModelForCausalLM.register(MyPretrainConfig, Transformer)
my_tokenizer = "./my_tokenizer"
tokenizer = AutoTokenizer.from_pretrained(my_tokenizer, trust_remote_code=True, use_fast=False)
ckp = f'./out/full_sft_{lm_config.dim}.pth.{batch_size}'
print(f"lmconfigs: {lm_config.to_json_string()}")
with open(ckp_path + "/config.json", 'w') as f:
f.write(lm_config.to_json_string())
# 拷貝文件到指定的目錄
for item in os.listdir(my_tokenizer):
src_item = os.path.join(my_tokenizer, item)
if os.path.isfile(src_item):
dest_item = os.path.join(ckp_path, item)
shutil.copy2(src_item, dest_item)
shutil.copy2(ckp, ckp_path + "/pytorch_model.bin")
model = AutoModelForCausalLM.from_pretrained(ckp_path, trust_remote_code=True).to(device)
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
tokenizer.pad_token = tokenizer.eos_token
print(f'LLM總參數(shù)量:{count_parameters(model) / 1e6:.3f} 百萬')
model = model.to(device)
return model, tokenizer
if __name__ == '__main__':
lm_config = MyPretrainConfig()
max_seq_len = lm_config.max_seq_len
out_dir = 'out'
epochs = 20 # 訓練輪數(shù)
batch_size = 8 # batch_size
learning_rate = 1e-5 # 學習率
device = 'cuda:0' # or cpu
dtype = 'bfloat16'
ckp_path = f'./my_checkpoint'
if not os.path.exists(ckp_path):
os.makedirs(ckp_path)
model, tokenizer = init_model()
training_config = DPOConfig(
output_dir=ckp_path,
per_device_train_batch_size=1,
remove_unused_columns=False,
report_to="none",
save_steps=2000,
learning_rate=learning_rate,
save_safetensors=False,
)
# 下載訓練圖片:https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main/dpo
dataset_path = f'{basepath}/dpo_train_data.json'
train_dataset = load_dataset('json', data_files=dataset_path)
dpo_trainer = DPOTrainer(
model,
ref_model=None,
args=training_config,
beta=0.1,
train_dataset=train_dataset['train'],
tokenizer=tokenizer,
max_length=512,
max_prompt_length=512
)
dpo_trainer.train()
dpo_trainer.save_model(
output_dir=f"./out/dpo_sft_{lm_config.dim}.pth.{batch_size}"
)
- init_model 函數(shù)主要是注冊和加載預訓練的模型,并將 tokeinzer 的一些配置文件都拷貝到 ./my_checkpoint 方便后續(xù)的訓練;
- DPOConfig 主要是配置訓練的一些參數(shù),比如保存的模型路徑、學習率等;
- DPOTrainer 是 DPO 訓練器,將模型載入后調(diào)用 train 進行訓練,參數(shù)說明如下:
model: transformers.PreTrainedModel,預訓練模型
ref_model: transformers.PreTrainedModel,參考模型
args: DPOConfig,用于訓練的 DPO 配置參數(shù)
train_dataset: datasets.Dataset,訓練數(shù)據(jù)集
tokenizer: transformers.PreTrainedTokenizerBase,分詞器
model_init: 用于訓練的模型初始化器,如果指定為 None,則將使用默認的模型初始化器
optimizer: torch.optim.Optimizer,優(yōu)化器
callbacks: 用于訓練的回調(diào)函數(shù)
- dpo_trainer.save_model 保存模型,傳入 output_dir 參數(shù),指定保存的模型路徑
4、總結(jié)
至此,訓練系列按照步驟寫完了,現(xiàn)在總結(jié)訓練流程:
模型訓練流程
不過驗證下來,訓練效果不是很好,這個也是從0開始訓練會遇到的問題,因此接下來會完成幾個事項:
- 模型迭代優(yōu)化,解決訓練效果不好的問題;
- 模型嘗試新的模型和解決方案,解決訓練速度問題;
- 加入多模態(tài)訓練集,將語言大模型改進為多模態(tài)模型;
- 最后將整個模型訓練完成后,將代碼開源。