diff --git a/data/sotopia_pi_gpt4_rm_overfit.json b/data/sotopia_pi_gpt4_rm_overfit.json index 2f33388..c2be64a 100644 --- a/data/sotopia_pi_gpt4_rm_overfit.json +++ b/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 } -] +] \ No newline at end of file diff --git a/evals/rej_sampling_serving.sh b/evals/rej_sampling_serving.sh index 4954056..18f51d3 100644 --- a/evals/rej_sampling_serving.sh +++ b/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 diff --git a/sotopia_rl/ppo_trainer.py b/sotopia_rl/ppo_trainer.py index c0d8575..0eb65d7 100644 --- a/sotopia_rl/ppo_trainer.py +++ b/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="") + 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, diff --git a/sotopia_rl/rm_trainer.py b/sotopia_rl/rm_trainer.py index ac45076..95c1148 100644 --- a/sotopia_rl/rm_trainer.py +++ b/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(