🐛 fix DPO merging issues see huggingface/trl#742

lib/style-transfer/style_transfer/rb_gen/steps/dpo.py

@@ -1,11 +1,12 @@
 import hydra
 import numpy as np
 import pandas as pd
+import peft
 import wandb
 from datasets import Dataset
-from peft import AutoPeftModelForCausalLM
+from omegaconf import ListConfig
 from style_transfer.rb_gen.utils.utils import CustomWandbCallback
-from transformers import PreTrainedTokenizerBase
+from transformers import AutoModelForCausalLM, PreTrainedTokenizerBase
 from trl import DPOTrainer
 
 
@@ -62,15 +63,29 @@ def dpo_train(
     cfg.dpo.training_args.output_dir = f"models/{wandb.run.id}/dpo/{step}"
     args = hydra.utils.instantiate(cfg.dpo.training_args)
     args.padding_value = tokenizer.eos_token_id
-    model = AutoPeftModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_path)
+    model = AutoModelForCausalLM.from_pretrained(
+        pretrained_model_name_on_path=f"models/{wandb.run.id}/merged/"
+    )
     model.enable_input_require_grads()
+    peft_config = hydra.utils.instantiate(cfg.model.peft_config)
+    peft_config.target_modules = (
+        list(peft_config.target_modules)
+        if isinstance(peft_config.target_modules, ListConfig)
+        else peft_config.target_modules
+    )
+    model = peft.get_peft_model(
+        model,
+        peft_config,
+    )
+    model.add_adapter(peft_config=peft_config, adapter_name="reference")
     dpo_trainer = DPOTrainer(
         args=args,
-        ref_model=None,
         model=model,
         tokenizer=tokenizer,
         train_dataset=dataset,
         callbacks=[CustomWandbCallback],
+        model_adapter_name="default",
+        ref_adapter_name="reference",
     )
     dpo_trainer.train()
 

lib/style-transfer/style_transfer/rb_gen/steps/generate.py

@@ -2,7 +2,6 @@
 import json
 import logging
 import os
-import shutil
 import sqlite3
 from typing import Callable
 
@@ -78,7 +77,6 @@ def generate(
     del llm
     gc.collect()
     torch.cuda.empty_cache()
-    shutil.rmtree(f"models/{wandb.run.id}/merged/")
     return gen_pred_dataset
 
 

lib/style-transfer/style_transfer/run_rb_gen.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import shutil
 
 import hydra
 import wandb
@@ -14,7 +15,7 @@
 logger = logging.getLogger(__name__)
 
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
-os.environ["WANDB_LOG_MODEL"] = "checkpoint"
+os.environ["WANDB_LOG_MODEL"] = "none"
 os.environ["WANDB_START_METHOD"] = "thread"
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 tqdm.pandas()
@@ -65,6 +66,7 @@ def main(cfg: DictConfig):
     sft_train(cfg, sft_dataset, test_dataset, current_model_path)
     logger.info("Bootstrapping done,  Iterative Reward-based Generation Training begins...")
     for step in range(cfg.max_steps):
+        logger.info(f"🔄 Step {step} ...")
         sth_dataset = generate(
             cfg,
             step,
@@ -97,6 +99,7 @@ def main(cfg: DictConfig):
         sth_dataset,
         checkpoint=eval_model_path,
     )
+    shutil.rmtree(f"models/{wandb.run.id}/merged/")
     wandb.finish()