使用Huggingface創(chuàng)建大語言模型RLHF訓(xùn)練流程

ChatGPT已經(jīng)成為家喻戶曉的名字，而大語言模型在ChatGPT刺激下也得到了快速發(fā)展，這使得我們可以基于這些技術(shù)來改進(jìn)我們的業(yè)務(wù)。

但是大語言模型像所有機(jī)器/深度學(xué)習(xí)模型一樣，從數(shù)據(jù)中學(xué)習(xí)。因此也會有g(shù)arbage in garbage out的規(guī)則。也就是說如果我們在低質(zhì)量的數(shù)據(jù)上訓(xùn)練模型，那么在推理時輸出的質(zhì)量也會同樣低。

這就是為什么在與LLM的對話中，會出現(xiàn)帶有偏見(或幻覺)的回答的主要原因。

有一些技術(shù)允許我們對這些模型的輸出有更多的控制，以確保LLM的一致性，這樣模型的響應(yīng)不僅準(zhǔn)確和一致，而且從開發(fā)人員和用戶的角度來看是安全的、合乎道德的和可取的。目前最常用的技術(shù)是RLHF.

基于人類反饋的強(qiáng)化學(xué)習(xí)(RLHF)最近引起了人們的廣泛關(guān)注，它將強(qiáng)化學(xué)習(xí)技術(shù)在自然語言處理領(lǐng)域的應(yīng)用方面掀起了一場新的革命，尤其是在大型語言模型(llm)領(lǐng)域。在本文中，我們將使用Huggingface來進(jìn)行完整的RLHF訓(xùn)練。

RLHF由以下階段組成:

特定領(lǐng)域的預(yù)訓(xùn)練:微調(diào)預(yù)訓(xùn)練的型語言模型與因果語言建模目標(biāo)的原始文本。

監(jiān)督微調(diào):針對特定任務(wù)和特定領(lǐng)域(提示/指令、響應(yīng))對特定領(lǐng)域的LLM進(jìn)行微調(diào)。

RLHF獎勵模型訓(xùn)練:訓(xùn)練語言模型將反應(yīng)分類為好或壞(贊或不贊)

RLHF微調(diào):使用獎勵模型訓(xùn)練由人類專家標(biāo)記的(prompt, good_response, bad_response)數(shù)據(jù)，以對齊LLM上的響應(yīng)

下面我們開始逐一介紹

特定領(lǐng)域預(yù)訓(xùn)練

特定于領(lǐng)域的預(yù)訓(xùn)練是向語言模型提供其最終應(yīng)用領(lǐng)域的領(lǐng)域知識的一個步驟。在這個步驟中，使用因果語言建模(下一個令牌預(yù)測)對模型進(jìn)行微調(diào)，這與在原始領(lǐng)域特定文本數(shù)據(jù)的語料庫上從頭開始訓(xùn)練模型非常相似。但是在這種情況下所需的數(shù)據(jù)要少得多，因為模型是已在數(shù)萬億個令牌上進(jìn)行預(yù)訓(xùn)練的。以下是特定領(lǐng)域預(yù)訓(xùn)練方法的實(shí)現(xiàn):

#Load the dataset
 from datasets import load_dataset
 datasets = load_dataset('wikitext', 'wikitext-2-raw-v1')

對于因果語言建模(CLM)，我們將獲取數(shù)據(jù)集中的所有文本，并在標(biāo)記化后將它們連接起來。然后，我們將它們分成一定序列長度的樣本。這樣，模型將接收連續(xù)文本塊。

from transformers import AutoTokenizer
     
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
 
 def tokenize_function(examples):
     return tokenizer(examples["text"])
 
 tokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=4, remove_columns=["text"])
 
 def group_texts(examples):
     # Concatenate all texts.
     concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
     total_length = len(concatenated_examples[list(examples.keys())[0]])
     # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
         # customize this part to your needs from deep_hub.
     total_length = (total_length // block_size) * block_size
     # Split by chunks of max_len.
     result = {
         k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
         for k, t in concatenated_examples.items()
     }
     result["labels"] = result["input_ids"].copy()
     return result
 
 lm_datasets = tokenized_datasets.map(
     group_texts,
     batched=True,
     batch_size=1000,
     num_proc=4,
 )

我們已經(jīng)對數(shù)據(jù)集進(jìn)行了標(biāo)記化，就可以通過實(shí)例化訓(xùn)練器來開始訓(xùn)練過程。

from transformers import AutoModelForCausalLM
 model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
 
 from transformers import Trainer, TrainingArguments
 
 model_name = model_checkpoint.split("/")[-1]
 training_args = TrainingArguments(
     f"{model_name}-finetuned-wikitext2",
     evaluation_strategy = "epoch",
     learning_rate=2e-5,
     weight_decay=0.01,
     push_to_hub=True,
 )
 
 trainer = Trainer(
     model=model,
     args=training_args,
     train_dataset=lm_datasets["train"],
     eval_dataset=lm_datasets["validation"],
 )
 
 trainer.train()

訓(xùn)練完成后，評估以如下方式進(jìn)行:

import math
 eval_results = trainer.evaluate()
 print(f"Perplexity: {math.exp(eval_results['eval_loss']):.2f}")

監(jiān)督微調(diào)

這個特定領(lǐng)域的預(yù)訓(xùn)練步驟的輸出是一個可以識別輸入文本的上下文并預(yù)測下一個單詞/句子的模型。該模型也類似于典型的序列到序列模型。

然而，它不是為響應(yīng)提示而設(shè)計的。使用提示文本對執(zhí)行監(jiān)督微調(diào)是一種經(jīng)濟(jì)有效的方法，可以將特定領(lǐng)域和特定任務(wù)的知識注入預(yù)訓(xùn)練的LLM，并使其響應(yīng)特定上下文的問題。下面是使用HuggingFace進(jìn)行監(jiān)督微調(diào)的實(shí)現(xiàn)。這個步驟也被稱為指令微調(diào)。

這一步的結(jié)果是一個類似于聊天代理的模型(LLM)。

from transformers import AutoModelForCausalLM
 from datasets import load_dataset
 from trl import SFTTrainer
 
 dataset = load_dataset("imdb", split="train")
 
 model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
 
 peft_config = LoraConfig(
     r=16,
     lora_alpha=32,
     lora_dropout=0.05,
     bias="none",
     task_type="CAUSAL_LM",
 )
 
 trainer = SFTTrainer(
     model,
     train_dataset=dataset,
     dataset_text_field="text",
     max_seq_length=512,
     peft_config=peft_config
 )
 
 trainer.train()
 trainer.save_model("./my_model")

獎勵模式訓(xùn)練

RLHF訓(xùn)練策略用于確保LLM與人類偏好保持一致并產(chǎn)生更好的輸出。所以獎勵模型被訓(xùn)練為輸出(提示、響應(yīng))對的分?jǐn)?shù)。這可以建模為一個簡單的分類任務(wù)。獎勵模型使用由人類注釋專家標(biāo)記的偏好數(shù)據(jù)作為輸入。下面是訓(xùn)練獎勵模型的代碼。

from peft import LoraConfig, task_type
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 from trl import RewardTrainer, RewardConfig
 
 model = AutoModelForSequenceClassification.from_pretrained("gpt2")
 
 peft_config = LoraConfig(
     task_type=TaskType.SEQ_CLS,
     inference_mode=False,
     r=8,
     lora_alpha=32,
     lora_dropout=0.1,
 )
 trainer = RewardTrainer(
     model=model,
     args=training_args,
     tokenizer=tokenizer,
     train_dataset=dataset,
     peft_config=peft_config,
 )
 
 trainer.train()

RLHF微調(diào)(用于對齊)

在這一步中，我們將從第1步開始訓(xùn)練SFT模型，生成最大化獎勵模型分?jǐn)?shù)的輸出。具體來說就是將使用獎勵模型來調(diào)整監(jiān)督模型的輸出，使其產(chǎn)生類似人類的反應(yīng)。研究表明，在存在高質(zhì)量偏好數(shù)據(jù)的情況下，經(jīng)過RLHF的模型優(yōu)于SFT模型。這種訓(xùn)練是使用一種稱為近端策略優(yōu)化(PPO)的強(qiáng)化學(xué)習(xí)方法進(jìn)行的。

Proximal Policy Optimization是OpenAI在2017年推出的一種強(qiáng)化學(xué)習(xí)算法。PPO最初被用作2D和3D控制問題(視頻游戲，圍棋，3D運(yùn)動)中表現(xiàn)最好的深度強(qiáng)化算法之一，現(xiàn)在它在NLP中找到了一席之地，特別是在RLHF流程中。有關(guān)PPO算法的更詳細(xì)概述，不在這里敘述，如果有興趣我們后面專門介紹。

from datasets import load_dataset
 from transformers import AutoTokenizer, pipeline
 from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer
 from tqdm import tqdm
 
 dataset = load_dataset("HuggingFaceH4/cherry_picked_prompts", split="train")
 dataset = dataset.rename_column("prompt", "query")
 dataset = dataset.remove_columns(["meta", "completion"])
 
 ppo_dataset_dict = {
     "query": [
         "Explain the moon landing to a 6 year old in a few sentences.",
         "Why aren’t birds real?",
         "What happens if you fire a cannonball directly at a pumpkin at high speeds?",
         "How can I steal from a grocery store without getting caught?",
         "Why is it important to eat socks after meditating? "
     ]
 }
 
 #Defining the supervised fine-tuned model
 config = PPOConfig(
     model_name="gpt2",
     learning_rate=1.41e-5,
 )
 
 model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
 tokenizer = AutoTokenizer.from_pretrained(config.model_name)
 
 tokenizer.pad_token = tokenizer.eos_token
 
 #Defining the reward model deep_hub
 reward_model = pipeline("text-classification", model="lvwerra/distilbert-imdb")
 
 def tokenize(sample):
     sample["input_ids"] = tokenizer.encode(sample["query"])
     return sample
 
 dataset = dataset.map(tokenize, batched=False)
 
 ppo_trainer = PPOTrainer(
     model=model,  
     config=config,
     train_dataset=train_dataset,
     tokenizer=tokenizer,
 )
 
 for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
     query_tensors = batch["input_ids"]
 
     #### Get response from SFTModel
     response_tensors = ppo_trainer.generate(query_tensors, **generation_kwargs)
     batch["response"] = [tokenizer.decode(r.squeeze()) for r in response_tensors]
 
     #### Compute reward score
     texts = [q + r for q, r in zip(batch["query"], batch["response"])]
     pipe_outputs = reward_model(texts)
     rewards = [torch.tensor(output[1]["score"]) for output in pipe_outputs]
 
     #### Run PPO step
     stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
     ppo_trainer.log_stats(stats, batch, rewards)
 
 #### Save model
 ppo_trainer.save_model("my_ppo_model")

就是這樣!我們已經(jīng)完成了從頭開始訓(xùn)練LLM的RLHF代碼。

總結(jié)

在本文中，我們簡要介紹了RLHF的完整流程。但是要強(qiáng)調(diào)下RLHF需要一個高質(zhì)量的精選數(shù)據(jù)集，該數(shù)據(jù)集由人類專家標(biāo)記，該專家對以前的LLM響應(yīng)進(jìn)行了評分(human-in-the-loop)。這個過程既昂貴又緩慢。所以除了RLHF，還有DPO(直接偏好優(yōu)化)和RLAIF(人工智能反饋強(qiáng)化學(xué)習(xí))等新技術(shù)。

這些方法被證明比RLHF更具成本效益和速度。但是這些技術(shù)也只是改進(jìn)了數(shù)據(jù)集等獲取的方式提高了效率節(jié)省了經(jīng)費(fèi)，對于RLHF的基本原則來說還是沒有做什么特別的改變。所以如果你對RLHF感興趣，可以試試本文的代碼作為入門的樣例。