Package: instruct-goose
InstructGoosePaper: InstructGPT -Training language models to follow instructions
with human feedbackInstallInstall from PipPypipinstallinstruct-gooseInstall directly from the source codegitclonehttps://github.com/xrsrke/instructGOOSE.gitcdinstructGOOSE
pipinstall-e.Train the RL-based language modelfromtransformersimportAutoTokenizer,AutoModelForCausalLMfromdatasetsimportload_datasetimporttorchfromtorch.utils.dataimportDataLoader,random_splitfromtorchimportoptimfrominstruct_gooseimportAgent,RewardModel,RLHFTrainer,RLHFConfig,create_reference_modelStep 1:Load datasetdataset=load_dataset("imdb",split="train")dataset,_=random_split(dataset,lengths=[10,len(dataset)-10])# for demenstration purposestrain_dataloader=DataLoader(dataset,batch_size=4,shuffle=True)Found cached dataset imdb (/Users/education/.cache/huggingface/datasets/imdb/plain_text/1.0.0/d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0)Step 2: Load the pre-trained model and tokenizermodel_base=AutoModelForCausalLM.from_pretrained("gpt2")# for demonstration purposesreward_model=RewardModel("gpt2")tokenizer=AutoTokenizer.from_pretrained("gpt2",padding_side="left")eos_token_id=tokenizer.eos_token_idtokenizer.pad_token=tokenizer.eos_tokenStep 3: Create the RL-based language model agent and the reference
modelmodel=Agent(model_base)ref_model=create_reference_model(model)Step 4: Train itmax_new_tokens=20generation_kwargs={"min_length":-1,"top_k":0.0,"top_p":1.0,"do_sample":True,"pad_token_id":tokenizer.eos_token_id,"max_new_tokens":max_new_tokens}config=RLHFConfig()N_EPOCH=1# for demonstration purposestrainer=RLHFTrainer(model,ref_model,config)optimizer=optim.SGD(model.parameters(),lr=1e-3)forepochinrange(N_EPOCH):forbatchintrain_dataloader:inputs=tokenizer(batch["text"],padding=True,truncation=True,return_tensors="pt")response_ids=model.generate(inputs["input_ids"],attention_mask=inputs["attention_mask"],**generation_kwargs)# extract the generated textresponse_ids=response_ids[:,-max_new_tokens:]response_attention_mask=torch.ones_like(response_ids)# evaluate from the reward modelwithtorch.no_grad():text_input_ids=torch.stack([torch.concat([q,r])forq,rinzip(inputs["input_ids"],response_ids)],dim=0)rewards=reward_model(text_input_ids)# calculate PPO lossloss=trainer.compute_loss(query_ids=inputs["input_ids"],query_attention_mask=inputs["attention_mask"],response_ids=response_ids,response_attention_mask=response_attention_mask,rewards=rewards)optimizer.zero_grad()loss.backward()optimizer.step()print(f"loss={loss}")loss=-824.6560668945312
loss=0.030958056449890137
loss=4.284017562866211TODOAdd support custom reward functionAdd support custom value functionAdd support non-transformer modelsWrite config classResourcesI implemented this using these resourcesCopied theload_yamlfunction fromhttps://github.com/Dahoas/reward-modelingHow to build a dataset to train reward model:https://wandb.ai/carperai/summarize_RLHF/reports/Implementing-RLHF-Learning-to-Summarize-with-trlX–VmlldzozMzAwODM2How to add value head in PPO agent:https://github.com/lvwerra/trlHow to calculate the loss of PPO agent:https://github.com/lvwerra/trl/blob/main/trl/trainer/ppo_trainer.pyHow to use PPO to train RLHF agent:https://github.com/voidful/TextRLHow PPO works:https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo.pyCopied the computeadvantagesandreturnsfromTLR:https://github.com/lvwerra/trl/blob/d2e8bcf8373726fb92d2110c500f7df6d0bd566d/trl/trainer/ppo_trainer.py#L686