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

1、为什么需要DPO

Rafailov等人在2023年发表了一篇论文《Direct Preference Optimization: Your Language Model is Secretly a Reward Model》，该论文提出了一种新的训练方法，称为直接偏好优化（DPO），该论文介绍：

复制

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

由此可见，DPO 主要解决RLHF不稳定的问题，直接使用人类偏好数据训练模型。

2、DPO的训练原理

DPO 的训练原理如下图所示（出自原论文）：

DPO

主要包括两个步骤：

• 数据收集：收集一个偏好数据集，其中包含给定提示的生成结果的正负选择对；
• 优化：直接最大化 DPO 损失的对数似然函数，该损失函数是偏好数据集上的交叉熵损失和模型生成结果的对数似然性之间的加权平均值；

具体公式推导可以参考这篇博客：https://www.cnblogs.com/lemonzhang/p/17910358.html。

3、DPO的代码实现

3.1 收集数据

DPO 训练器对数据集的格式有具体的要求，包括三个部分：

• 提示（prompt）：提示的格式为：prompt: 文本；
• 选中（chosen）：选中文本的格式为：chosen: 文本；
• 拒绝（rejected）：拒绝选中文本的格式为：rejected: 文本；
• 示例：

复制

{
    "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希望这有帮助，小伙伴！😊"
}

DPO的数据可以搜索huggingface的DPO数据集，地址为：https://huggingface.co/datasets?sort=trending&search=dpo 。

比如 https://huggingface.co/datasets/Anthropic/hh-rlhf 的数据集如下：

hh-rlhf

3.2 TRL

引入 TRL 库，支持 DPO 训练器，训练样例代码：

复制

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()

如上训练默认是保存 safetensors 格式的模型，如果想保存 pytorch 格式的模型， 可以改为如下代码：

复制

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"
)

3.3 训练

Transformer的代码和前面的一样，可以参考预训练的代码，如下就是初始化模型和 DPO 训练的代码：

复制

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}"
    )

• init_model 函数主要是注册和加载预训练的模型，并将 tokeinzer 的一些配置文件都拷贝到 ./my_checkpoint 方便后续的训练；
• DPOConfig 主要是配置训练的一些参数，比如保存的模型路径、学习率等；
• DPOTrainer 是 DPO 训练器，将模型载入后调用 train 进行训练，参数说明如下：

model: transformers.PreTrainedModel，预训练模型

ref_model: transformers.PreTrainedModel，参考模型

args: DPOConfig，用于训练的 DPO 配置参数

train_dataset: datasets.Dataset，训练数据集

tokenizer: transformers.PreTrainedTokenizerBase，分词器

model_init: 用于训练的模型初始化器，如果指定为 None，则将使用默认的模型初始化器

optimizer: torch.optim.Optimizer，优化器

callbacks: 用于训练的回调函数

• dpo_trainer.save_model 保存模型，传入 output_dir 参数，指定保存的模型路径

4、总结

至此，训练系列按照步骤写完了，现在总结训练流程：

模型训练流程

不过验证下来，训练效果不是很好，这个也是从0开始训练会遇到的问题，因此接下来会完成几个事项：

• 模型迭代优化，解决训练效果不好的问题；
• 模型尝试新的模型和解决方案，解决训练速度问题；
• 加入多模态训练集，将语言大模型改进为多模态模型；
• 最后将整个模型训练完成后，将代码开源；