雖然大規模無監督語言模型 (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}"
    )