Bug/add debug function

data/sotopia_pi_gpt4_rm_overfit.json

@@ -19,4 +19,4 @@
         "output": "{\"action_type\": \"speak\", \"argument\": \"I'm excited too, Benjamin! See you at the marathon. Let's make it unforgettable!\"}",
         "value": 0.0
     }
-]
+]

evals/rej_sampling_serving.sh

@@ -49,8 +49,8 @@ export RM_MODEL_PATH="${REPO_FOLDER_NAME}/${RM_FOLDER_NAME}/checkpoint-3595"
 
 export VLLM_GPU=8
 export DJANGO_GPU=9
-export VLLM_PORT=8045
-export DJANGO_PORT=8057
+export VLLM_PORT=8065
+export DJANGO_PORT=8077
 export REJ_SAMPLING_NUM=10
 export SFT_MODEL_FOLDER_NAME="sft_qwen25_7b"
 export RM_FOLDER_NAME="rm_reward_utterance_quality_gpt-4o"
@@ -59,6 +59,18 @@ export SFT_MODEL_PATH="${REPO_FOLDER_NAME}/${SFT_MODEL_FOLDER_NAME}/checkpoint-1
 export RM_MODEL_PATH="${REPO_FOLDER_NAME}/${RM_FOLDER_NAME}/checkpoint-4200"
 
 
+export VLLM_GPU=6
+export DJANGO_GPU=7
+export VLLM_PORT=8065
+export DJANGO_PORT=8077
+export REJ_SAMPLING_NUM=10
+export SFT_MODEL_FOLDER_NAME="sft_qwen25_7b"
+export RM_FOLDER_NAME="rm_reward_mixed"
+export REPO_FOLDER_NAME="/data/haofeiy2/sotopia-rl"
+export SFT_MODEL_PATH="${REPO_FOLDER_NAME}/${SFT_MODEL_FOLDER_NAME}/checkpoint-1000/"
+export RM_MODEL_PATH="${REPO_FOLDER_NAME}/${RM_FOLDER_NAME}/checkpoint-5000"
+
+
 export TAG="${RM_FOLDER_NAME}_rej_sampling_num${REJ_SAMPLING_NUM}_vs_${SFT_MODEL_FOLDER_NAME}-0322"
 export SFT_MODEL_NAME="qwen25-7b-instruct-sft-gpu${VLLM_GPU}"
 export MODEL_A=custom/${RM_FOLDER_NAME}_rejsampling_num${REJ_SAMPLING_NUM}@http://localhost:${DJANGO_PORT}/sotopia

sotopia_rl/ppo_trainer.py

@@ -65,6 +65,85 @@ def save_model(self, output_dir: str, _internal_call: bool = False):
 
         self.ppo_trainer.save_model = save_model.__get__(self.ppo_trainer, type(self.ppo_trainer))
 
+        def save_model(self, output_dir: str, _internal_call: bool = False):
+            if hasattr(self.model, "module"):
+                model_to_save = self.model.module
+            else:
+                model_to_save = self.model
+            if not hasattr(model_to_save, "policy"):
+                model_to_save.policy = model_to_save
+            model_to_save.save_pretrained(output_dir)
+            print(f"Model saved to {output_dir}")
+
+        self.ppo_trainer.save_model = save_model.__get__(self.ppo_trainer, type(self.ppo_trainer))
+
+        def debug_reward_forward(self):
+
+            with open("/data/haofeiy2/sotopia-rl/data/sotopia_pi_gpt4_rm_overfit.json", 'r') as f:
+                example_data = json.load(f)
+            if len(example_data) == 0:
+                print("Example data is empty. Skipping debug_reward_forward.")
+                return
+
+            template_dir = os.path.dirname(self.args.template_path)
+            template_file = os.path.basename(self.args.template_path)
+            env = Environment(loader=FileSystemLoader(template_dir))
+            template = env.get_template(template_file)
+
+            example = example_data[1]
+            rendered_prompt = template.render(
+                messages=[
+                    {"role": "user", "content": example["input"]},
+                    {"role": "assistant", "content": example["output"]}
+                ],
+                add_generation_prompt=False
+            )
+            print("\n===== Debug Reward Forward =====")
+            print("Rendered prompt:")
+            print(rendered_prompt)
+
+            inputs = self.tokenizer(rendered_prompt, return_tensors="pt", truncation=True, padding=True)
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+            lm_backbone = getattr(self.reward_model, self.reward_model.base_model_prefix)
+            lm_backbone.eval()
+            outputs_check = lm_backbone(**inputs)
+            with torch.no_grad():
+                outputs = self.reward_model(**inputs)
+
+            print(outputs.logits)
+            print(outputs_check.logits)
+            query_responses = query_responses.to('cuda')
+            import pdb; pdb.set_trace()
+            attention_mask = query_responses != self.tokenizer.pad_token_id
+            position_ids = attention_mask.cumsum(1) - attention_mask.long()  # exclusive cumsum
+            lm_backbone = getattr(self.reward_model, self.reward_model.base_model_prefix)
+            self.reward_model.config.pad_token_id = self.tokenizer.pad_token_id
+            #input_ids = torch.masked_fill(query_responses, ~attention_mask, 0)
+            input_ids = query_responses
+            outputs = lm_backbone(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                return_dict=True,
+                output_hidden_states=True,
+                use_cache=False,  # otherwise mistral-based RM would error out
+            )
+            predicted_reward = outputs.logits.squeeze().cpu().item()
+            print("Predicted reward:", predicted_reward)
+            import pdb; pdb.set_trace()
+
+            gth_reward = example.get("value")
+            print("Predicted reward:", predicted_reward)
+            if gth_reward is not None:
+                print("Ground truth reward:", gth_reward)
+            else:
+                print("Ground truth reward: Not available")
+            print("===== End Debug =====\n")
+
+
+        debug_reward_forward(self)
+
     def _init_wandb(self):
         """Initialize Weights & Biases logging"""
         wandb.init(
@@ -75,7 +154,7 @@ def _init_wandb(self):
 
     def _setup_tokenizer(self):
         """Load and configure tokenizer"""
-        self.tokenizer = AutoTokenizer.from_pretrained(self.args.model_name, padding_side="left", pad_token="<pad>")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.args.model_name, padding_side="left")
 
     def _setup_dataset(self):
         """Prepare training and validation datasets"""
@@ -127,7 +206,6 @@ def _setup_generation_models(self):
         generation_config = GenerationConfig(
             pad_token_id=self.tokenizer.pad_token_id,
             eos_token_id=self.tokenizer.eos_token_id,
-            bos_token_id=self.tokenizer.eos_token_id,
             max_length=self.args.max_length,
             do_sample=getattr(self.args, 'do_sample', True),
             temperature=getattr(self.args, 'temperature', 0.7),
@@ -136,55 +214,69 @@ def _setup_generation_models(self):
             no_repeat_ngram_size=getattr(self.args, 'no_repeat_ngram_size', 0)
         )
 
-        self.policy = PeftModelForCausalLM.from_pretrained(
-            base_gen_model,
-            self.args.policy_adapter_path,
-            generation_config=generation_config
-        )
-        print("Policy model loaded/created")
-
-        self.ref_policy = PeftModelForCausalLM.from_pretrained(
-            base_gen_model,
-            self.args.ref_adapter_path,
-            generation_config=generation_config
-        )
-        self.ref_policy.eval()
-        print("Reference policy model loaded/created")
+        adapter_path = os.path.join(self.args.policy_adapter_path, 'adapter_model')
+        if os.path.exists(adapter_path + '.safetensors') or os.path.exists(adapter_path + '.bin'):
+            print(f"Loading policy adapter from: {self.args.policy_adapter_path}")
+            self.policy = PeftModelForCausalLM.from_pretrained(
+                base_gen_model,
+                self.args.policy_adapter_path,
+                generation_config=generation_config
+            )
+        else:
+            print(f"No adapter found at {adapter_path})")
+
+        adapter_path = os.path.join(self.args.ref_adapter_path, 'adapter_model')
+        if os.path.exists(adapter_path + '.safetensors') or os.path.exists(adapter_path + '.bin'):
+            print(f"Loading reference policy adapter from: {self.args.ref_adapter_path}")
+            self.ref_policy = PeftModelForCausalLM.from_pretrained(
+                base_gen_model,
+                self.args.ref_adapter_path,
+                generation_config=generation_config
+            )
+            self.ref_policy.eval()
+        else:
+            print(f"No adapter found at {adapter_path}")
 
 
     def _setup_classification_models(self):
         base_cls_model = AutoModelForSequenceClassification.from_pretrained(
             self.args.model_name,
             torch_dtype=torch.float32, # very important, otherwise NaN
-            quantization_config=self.quant_config,
             num_labels=1,
             return_dict=True,
-            device_map=None
+            device_map=None,
         )
+        # VERY VERY IMPORTANT
+        # specifically designed for PPO training, 
+        # based on the get_reward function
+        # it fill the input_ids paddings with 0s
+        base_cls_model.config.pad_token_id = 0
+
+        adapter_path = os.path.join(self.args.value_adapter_path, 'adapter_model')
+        if os.path.exists(adapter_path + '.safetensors') or os.path.exists(adapter_path + '.bin'):
+            print(f"Loading value adapter from: {self.args.value_adapter_path}")
+            self.value_model = PeftModelForSequenceClassification.from_pretrained(
+                base_cls_model,
+                self.args.value_adapter_path,
+            )
+        else:
+            raise ValueError(f"No adapter found at {adapter_path}")
+
+        adapter_path = os.path.join(self.args.reward_adapter_path, 'adapter_model')
+        if os.path.exists(adapter_path + '.safetensors') or os.path.exists(adapter_path + '.bin'):
+            print(f"Loading reward adapter from: {self.args.reward_adapter_path}")
+            self.reward_model = PeftModelForSequenceClassification.from_pretrained(
+                base_cls_model,
+                self.args.reward_adapter_path,
+            )
+            self.reward_model.eval()
+        else:
+            raise ValueError(f"No adapter found at {adapter_path}")
 
-        if base_cls_model.config.pad_token_id is None:
-            base_cls_model.config.pad_token_id = self.tokenizer.pad_token_id
 
-        self.reward_model = PeftModelForSequenceClassification.from_pretrained(
-            base_cls_model,
-            self.args.reward_adapter_path,
-            num_labels=1
-        )
-        self.reward_model.eval()
-        print("Reward model loaded/created")
-
-        self.value_model = PeftModelForSequenceClassification.from_pretrained(
-            base_cls_model,
-            self.args.value_adapter_path,
-            num_labels=1
-        )
-        print("Value model loaded/created")
 
     def _setup_ppo_trainer(self):
         """Configure the PPO trainer"""
-        # Get data collator if available
-
-
         # Configure PPO settings
         ppo_config = PPOv2Config(
             per_device_train_batch_size=self.args.per_device_train_batch_size,

sotopia_rl/rm_trainer.py

@@ -57,7 +57,7 @@ def __init__(self, args, **kwargs):
             base_model = base_model.to(self.device)
 
         model = PeftModelForSequenceClassification(base_model, peft_config)
-        model.config.pad_token_id = tokenizer.pad_token_id # important to set the config pad_token_id
+        model.config.pad_token_id = tokenizer.pad_token_id  # important to set the config pad_token_id
 
         # Set up the TrainingArguments with DeepSpeed support
         training_args = TrainingArguments(