使用Huggingface创建大语言模型RLHF训练流程的完整教程

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

 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,
 )

 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()

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

 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")

 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()

 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")