github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix 5-dpo_train.py bugs 

https://github.com/jingyaogong/minimind/issues/57
pull/63/head
Zhao Jiahao 2024-10-11 12:29:23 +08:00 committed by GitHub
parent 3a034a47c8
commit 3090a438a4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 27 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
5-dpo_train.py
View File

@ -1,32 +1,38 @@
import os import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch import torch
from transformers import TrainingArguments, AutoModelForCausalLM, AutoTokenizer from transformers import TrainingArguments, AutoModelForCausalLM, AutoTokenizer
from trl import DPOTrainer from trl import DPOTrainer, DPOConfig
from peft import get_peft_model, LoraConfig, TaskType from peft import get_peft_model, LoraConfig, TaskType
from datasets import load_dataset from datasets import load_dataset
def find_all_linear_names(model): def find_all_linear_names(model):
cls = torch.nn.Linear cls = torch.nn.Linear
lora_module_names = set() lora_module_names = set()
for name, module in model.named_modules(): for name, module in model.named_modules():
if isinstance(module, cls): if isinstance(module, cls):
names = name.split('.') names = name.split(".")
lora_module_names.add(names[0] if len(names) == 1 else names[-1]) lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names: if "lm_head" in lora_module_names:
lora_module_names.remove('lm_head') lora_module_names.remove("lm_head")
return list(lora_module_names) return list(lora_module_names)
def init_model(): def init_model():
device = 'cuda:0' device = "cuda:0"
# Do model patching and add fast LoRA weights # Do model patching and add fast LoRA weights
model_name_or_path = "minimind" model_name_or_path = "minimind-v1-small"
tokenizer_name_or_path = "minimind" tokenizer_name_or_path = "./model/minimind_tokenizer"
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, trust_remote_code=True, use_fast=False) model_name_or_path, trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(
tokenizer_name_or_path, trust_remote_code=True, use_fast=False
)
tokenizer.pad_token = tokenizer.eos_token tokenizer.pad_token = tokenizer.eos_token
target_modules = find_all_linear_names(model) target_modules = find_all_linear_names(model)
peft_config = LoraConfig( peft_config = LoraConfig(
@ -35,7 +41,7 @@ def init_model():
lora_alpha=16, lora_alpha=16,
lora_dropout=0.1, lora_dropout=0.1,
inference_mode=False, inference_mode=False,
target_modules=target_modules target_modules=target_modules,
) )
model = get_peft_model(model, peft_config) model = get_peft_model(model, peft_config)
model.print_trainable_parameters() model.print_trainable_parameters()
@ -43,29 +49,31 @@ def init_model():
return model, tokenizer return model, tokenizer
if __name__ == '__main__': if __name__ == "__main__":
model, tokenizer = init_model() model, tokenizer = init_model()
training_args = TrainingArguments(output_dir="./minimind_dpo", training_args = DPOConfig(
output_dir="./minimind_dpo",
per_device_train_batch_size=1, per_device_train_batch_size=1,
remove_unused_columns=False) remove_unused_columns=False,
)
################ ################
# Dataset # Dataset
################ ################
# 确保路径正确，文件存在 # 确保路径正确，文件存在
dataset_path = './dataset/dpo/train_data.json' dataset_path = "./dataset/dpo/train_data.json"
# 加载数据集 # 加载数据集
train_dataset = load_dataset('json', data_files=dataset_path) train_dataset = load_dataset("json", data_files=dataset_path)
dpo_trainer = DPOTrainer( dpo_trainer = DPOTrainer(
model, model,
ref_model=None, ref_model=None,
args=training_args, args=training_args,
beta=0.1, beta=0.1,
train_dataset=train_dataset['train'], train_dataset=train_dataset["train"],
tokenizer=tokenizer, tokenizer=tokenizer,
max_length=512, max_length=512,
max_prompt_length=512 max_prompt_length=512,
) )
dpo_trainer.train() dpo_trainer.train()