fix 5-dpo train
parent
6e9cd28ef9
commit
19db1a10ed
|
|
@ -0,0 +1,2 @@
|
||||||
|
* text=auto
|
||||||
|
* -text
|
||||||
|
|
@ -2,3 +2,5 @@
|
||||||
/dataset
|
/dataset
|
||||||
/wandb
|
/wandb
|
||||||
/out
|
/out
|
||||||
|
minimind-v1-small/
|
||||||
|
minimind_dpo/
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue