只用1塊A100,就能訓練自己的Llama-2模型!
你知道嗎?現(xiàn)在只用幾行代碼就能在自己的數(shù)據(jù)上訓練所有Llama-2模型!
而且甚至可以在單個A100 GPU上使用70億參數(shù)的模型,這要歸功于4bit和PEFT的魔力!
通過PPO微調(diào)語言模型大致包括三個步驟:
生成(Rollout):語言模型根據(jù)查詢(query)生成一個響應或連續(xù)文本,查詢可能是一個句子的開頭。
評估(Evaluation):使用一個函數(shù)、模型、人工反饋或它們的組合來對查詢和生成的響應進行評估。重要的是,該過程應為每個查詢/響應對產(chǎn)生一個標量值。
優(yōu)化(Optimization):這是最復雜的部分。在優(yōu)化步驟中,使用查詢和生成的響應對來計算序列中每個令牌的對數(shù)概率。這是通過已經(jīng)訓練過的模型和一個參考模型來完成的,參考模型通常是微調(diào)之前的預訓練模型。兩個模型輸出之間的KL散度被用作額外的獎勵信號,以確保生成的響應不會偏離參考語言模型太遠。然后使用PPO來訓練主要的語言模型。
這個過程在下面的示意圖中說明:
安裝:
要通過pip安裝Python庫:
pip install trl
如果你想從源代碼運行庫中的示例,還需要安裝一些額外的庫。首先,你需要克隆該庫的代碼倉庫,然后使用pip安裝它。
git clone https://github.com/lvwerra/trl.gitcd trl/pip install .
如果你希望開發(fā)TRL(指上文提到的名為"trl"的Python庫),可以使用可編輯模式進行安裝
pip install -e .
如何使用:
SFTTrainer
以下是如何使用庫中的SFTTrainer的基本示例。SFTTrainer是transformers Trainer的輕量級封裝,可以輕松地在自定義數(shù)據(jù)集上微調(diào)語言模型或適配器。
# imports
from datasets import load_dataset
from trl import SFTTrainer
# get datasetdataset = load_dataset("imdb", split="train")
# get trainer
trainer = SFTTrainer(
"facebook/opt-350m",
train_dataset=dataset,
dataset_text_field="text",
max_seq_length=512,)
# train
trainer.train()
RewardTrainer
以下是如何使用庫中的RewardTrainer的基本示例。
# imports
from transformers import AutoModelForSequenceClassification, AutoTokenizerfrom trl import RewardTrainer
# load model and dataset - dataset needs to be in a specific formatmodel = AutoModelForSequenceClassification.from_pretrained("gpt2")tokenizer = AutoTokenizer.from_pretrained("gpt2")
...
# load trainer
trainer = RewardTrainer(
model=model, t
okenizer=tokenizer,
train_dataset=dataset,)
# train
trainer.train()
PPOTrainer
以下是如何使用庫中的PPOTrainer的基本示例。根據(jù)查詢,語言模型創(chuàng)建一個響應,然后對該響應進行評估。評估可以由人類參與,也可以是另一個模型的輸出。
# importsi
mport torch
from transformers import AutoTokenizerf
rom trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead, create_reference_model
from trl.core import respond_to_batch
# get models
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
model_ref = create_reference_model(model)
tokenizer = AutoTokenizer.from_pretrained('gpt2')
# initialize trainer
ppo_config = PPOConfig(
batch_size=1,)
# encode a queryquery_txt = "This morning I went to the "
query_tensor = tokenizer.encode(query_txt, return_tensors="pt")
# get model response
response_tensor = respond_to_batch(model, query_tensor)
# create a ppo trainer
ppo_trainer = PPOTrainer(ppo_config, model, model_ref, tokenizer)
# define a reward for response
# (this could be any reward such as human feedback or output from another model)reward = [torch.tensor(1.0)]
# train model for one step with ppotrain_stats = ppo_t
rainer.step([query_tensor[0]], [response_tensor[0]], reward)
高級示例:IMDB情感分類
要了解詳細的示例,請查看項目examples/scripts/sentiment_tuning.py中的示例Python腳本。下面是在優(yōu)化之前和之后,從語言模型中提取的幾個示例: