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

46
5-dpo_train.py
View File

@ -1,32 +1,38 @@
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
from transformers import TrainingArguments, AutoModelForCausalLM, AutoTokenizer
from trl import DPOTrainer
from trl import DPOTrainer, DPOConfig
from peft import get_peft_model, LoraConfig, TaskType
from datasets import load_dataset
def find_all_linear_names(model):
cls = torch.nn.Linear
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split('.')
names = name.split(".")
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if 'lm_head' in lora_module_names:
lora_module_names.remove('lm_head')
if "lm_head" in lora_module_names:
lora_module_names.remove("lm_head")
return list(lora_module_names)
def init_model():
device = 'cuda:0'
device = "cuda:0"
# Do model patching and add fast LoRA weights
model_name_or_path = "minimind"
tokenizer_name_or_path = "minimind"
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, trust_remote_code=True, use_fast=False)
model_name_or_path = "minimind-v1-small"
tokenizer_name_or_path = "./model/minimind_tokenizer"
model = AutoModelForCausalLM.from_pretrained(
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
target_modules = find_all_linear_names(model)
peft_config = LoraConfig(
@ -35,7 +41,7 @@ def init_model():
lora_alpha=16,
lora_dropout=0.1,
inference_mode=False,
target_modules=target_modules
target_modules=target_modules,
)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()
@ -43,29 +49,31 @@ def init_model():
return model, tokenizer
if __name__ == '__main__':
if __name__ == "__main__":
model, tokenizer = init_model()
training_args = TrainingArguments(output_dir="./minimind_dpo",
per_device_train_batch_size=1,
remove_unused_columns=False)
training_args = DPOConfig(
output_dir="./minimind_dpo",
per_device_train_batch_size=1,
remove_unused_columns=False,
)
################
# 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(
model,
ref_model=None,
args=training_args,
beta=0.1,
train_dataset=train_dataset['train'],
train_dataset=train_dataset["train"],
tokenizer=tokenizer,
max_length=512,
max_prompt_length=512
max_prompt_length=512,
)
dpo_trainer.train()