fix 5-dpo train

pull/132/head
guomin 2025-01-31 23:21:43 +08:00
parent 6e9cd28ef9
commit 19db1a10ed
3 changed files with 44 additions and 9 deletions

2
.gitattributes vendored Normal file
View File

@ -0,0 +1,2 @@
* text=auto
* -text

2
.gitignore vendored
View File

@ -2,3 +2,5 @@
/dataset /dataset
/wandb /wandb
/out /out
minimind-v1-small/
minimind_dpo/

View File

@ -1,5 +1,6 @@
import os import os
import warnings import warnings
import wandb
os.environ['CUDA_VISIBLE_DEVICES'] = '0' os.environ['CUDA_VISIBLE_DEVICES'] = '0'
from transformers import TrainingArguments, AutoModelForCausalLM, AutoTokenizer from transformers import TrainingArguments, AutoModelForCausalLM, AutoTokenizer
@ -12,8 +13,8 @@ warnings.filterwarnings('ignore')
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-v1" model_name_or_path = "minimind-v1-small"
tokenizer_name_or_path = "minimind-v1" tokenizer_name_or_path = "minimind-v1-small"
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True) 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 = 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
@ -23,26 +24,56 @@ def init_model():
if __name__ == '__main__': if __name__ == '__main__':
model, tokenizer = init_model() model, tokenizer = init_model()
# 初始化 wandb
wandb.init(
project="minimind-dpo",
name=f"dpo-training-lr-4e-5-beta-0.1",
config={
"learning_rate": 4e-5,
"beta": 0.1,
"batch_size": 1,
"max_length": 512,
}
)
training_config = DPOConfig( training_config = DPOConfig(
output_dir="./minimind_dpo", output_dir="./minimind_dpo",
per_device_train_batch_size=1, per_device_train_batch_size=1,
remove_unused_columns=False, remove_unused_columns=False,
report_to="none", report_to="wandb", # 启用wandb记录
save_steps=2000, save_steps=2000,
learning_rate=4e-5 learning_rate=4e-5,
beta=0.1,
max_length=512,
max_prompt_length=512,
logging_steps=50, # 增加日志记录频率
logging_first_step=True # 记录第一步
# 移除了 evaluation_strategy 和 eval_steps
) )
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)
# 将数据集转换为训练集
# train_dataset = train_dataset['train']
print("-----------------")
print(train_dataset)
import sys
# sys.exit(0)
# copy_model = AutoModelForCausalLM.from_pretrained("minimind-v1-small", trust_remote_code=True)
dpo_trainer = DPOTrainer( dpo_trainer = DPOTrainer(
model, model,
ref_model=None, ref_model=None,
args=training_config, args=training_config,
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()
try:
dpo_trainer.train()
finally:
wandb.finish() # 确保正确关闭wandb