机器学习 | 从0开始大模型之模型DPO训练

虽然大规模无监督语言模型 (LM) 可以学习广泛的世界知识和一些推理技能，但由于其训练完全无监督，因此很难精确控制其行为。  
现有的获得这种可控性的方法是收集模型生成相对质量的人类标签，并微调无监督语言模型以符合这些偏好，通常使用从人类反馈中进行强化学习 (RLHF)。  
然而，RLHF 是一个复杂且通常不稳定的过程，首先要拟合一个反映人类偏好的奖励模型，然后使用强化学习微调大型无监督语言模型以最大化这个估计的奖励，而不会偏离原始模型太远。  
在该论文中，利用奖励函数和最优策略之间的映射来表明，这个受约束的奖励最大化问题可以通过一个阶段的策略训练进行精确优化，本质上是解决人类偏好数据的分类问题。  
由此产生的算法，称之为直接偏好优化 (DPO)，稳定、高效且计算量小，无需拟合奖励模型、在微调期间从 LM 中采样或执行重大超参数调整。

{
    "id": 0,
    "prompt": "最近，马其顿流行歌手托瑟·普罗埃斯基在一场巡回展览中被纪念。",
    "chosen": "\"最近，马其顿流行歌手托斯·普罗埃斯基在一场展览中被人们铭记。\" 让我帮你理解。这意味着他们举办了一场特别的展示（我们称之为展览），以帮助人们记住马其顿这个地方的歌手托斯·普罗埃斯基有多特别。这有点像翻看你的相册来回忆快乐时光，但规模更大，更专注于他成名后的工作和生活。",
    "rejected": "哦！你想知道这句话的意思吗？好的，让我告诉你！😄\n\n所以，“Son dönem”在英语中意思是“Last period”。而“Makedon”在英语中是“Macedonian”。而“pop şarkıcısı”在英语中意思是“pop singer”！🎤\n\n所以，句子“Son dönem Makedon pop şarkıcısı Tose Proeski gezer sergide anılıyor”的意思是“Last period Macedonian pop singer Tose Proeski正在舞台上表演。”🎭\n\n希望这有帮助，小伙伴！😊"
}

training_args = DPOConfig(
    beta=0.1,
)
dpo_trainer = DPOTrainer(
    model,
    ref_model,
    args=training_args,
    train_dataset=train_dataset,
    tokenizer=tokenizer,  # for visual language models, use tokenizer=processor instead
)
dpo_trainer.train()
dpo_trainer.save_model()

training_args = DPOConfig(
    beta=0.1,
    save_safetensors=False, // 设置为False，改为保存为pytorch格式的模型   
)
dpo_trainer = DPOTrainer(
    model,
    ref_model,
    args=training_args,
    train_dataset=train_dataset,
    tokenizer=tokenizer,  # for visual language models, use tokenizer=processor instead
)
dpo_trainer.train()
dpo_trainer.save_model(
    output_dir=f"./out/dpo_sft_xxx.pth"
)

def init_model():
    from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
    AutoConfig.register(MyPretrainConfig.model_type, MyPretrainConfig)
    AutoModelForCausalLM.register(MyPretrainConfig, Transformer)
    my_tokenizer = "./my_tokenizer"
    tokenizer = AutoTokenizer.from_pretrained(my_tokenizer, trust_remote_code=True, use_fast=False)
    ckp = f'./out/full_sft_{lm_config.dim}.pth.{batch_size}'

    print(f"lmconfigs: {lm_config.to_json_string()}")
    with open(ckp_path + "/config.json", 'w') as f:
        f.write(lm_config.to_json_string())

    # 拷贝文件到指定的目录
    for item in os.listdir(my_tokenizer):
        src_item = os.path.join(my_tokenizer, item)
        if os.path.isfile(src_item):
            dest_item = os.path.join(ckp_path, item)
            shutil.copy2(src_item, dest_item)
    shutil.copy2(ckp, ckp_path + "/pytorch_model.bin")

    model = AutoModelForCausalLM.from_pretrained(ckp_path, trust_remote_code=True).to(device)

    def count_parameters(model):
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    tokenizer.pad_token = tokenizer.eos_token
    print(f'LLM总参数量：{count_parameters(model) / 1e6:.3f} 百万')
    model = model.to(device)
    return model, tokenizer

if __name__ == '__main__':
    lm_config = MyPretrainConfig()
    max_seq_len = lm_config.max_seq_len
    out_dir = 'out'
    epochs = 20             # 训练轮数
    batch_size = 8          # batch_size
    learning_rate = 1e-5    # 学习率
    device = 'cuda:0'       # or cpu
    dtype = 'bfloat16'

    ckp_path = f'./my_checkpoint'
    if not os.path.exists(ckp_path):
        os.makedirs(ckp_path)

    model, tokenizer = init_model()
    training_config = DPOConfig(
        output_dir=ckp_path,
        per_device_train_batch_size=1,
        remove_unused_columns=False,
        report_to="none",
        save_steps=2000,
        learning_rate=learning_rate,
        save_safetensors=False,
    )

    # 下载训练图片：https://huggingface.co/datasets/jingyaogong/minimind_dataset/tree/main/dpo
    dataset_path = f'{basepath}/dpo_train_data.json'
    train_dataset = load_dataset('json', data_files=dataset_path)
    dpo_trainer = DPOTrainer(
        model,
        ref_model=None,
        args=training_config,
        beta=0.1,
        train_dataset=train_dataset['train'],
        tokenizer=tokenizer,
        max_length=512,
        max_prompt_length=512
    )
    dpo_trainer.train()
    dpo_trainer.save_model(
        output_dir=f"./out/dpo_sft_{lm_config.dim}.pth.{batch_size}"
    )