# Modified from codi_ori/train.py # Key changes: # 1. Import CODI_CT (now PMPS-enhanced) instead of CODI # 2. Preprocess now computes cot_start_position and cot_end_position # 3. DataCollator passes CoT-position fields to the model so the PMPS # loss can extract teacher-side CoT hidden states. # 4. CustomTrainer logs pmps_loss and related diagnostics import copy import logging import os import re import random from dataclasses import dataclass, field from typing import Dict, Optional, Sequence import torch import json import transformers from torch.utils.data import Dataset from transformers import Trainer from safetensors.torch import load_file from tqdm import tqdm from math import ceil from peft import PeftModel, LoraConfig, TaskType, get_peft_model from datasets import load_dataset from functools import partial import yaml import argparse import wandb from src.model import ( CODI_CT as CODI_PMPS, ModelArguments, DataArguments, TrainingArguments, freeze_model, ) IGNORE_INDEX = -100 def read_json(file_path): """ 从指定路径读取JSON文件并返回对应的Python对象。 """ try: with open(file_path, 'r', encoding='utf-8') as file: data = json.load(file) return data except Exception as e: print(f"读取JSON文件时出错: {e}") return None device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(device) class CustomTrainer(Trainer): def compute_loss(self, model, inputs, num_items_in_batch): step = self.state.global_step batch_size = self.args.per_device_train_batch_size gradient_accumulation_steps = self.args.gradient_accumulation_steps num_epochs = self.args.num_train_epochs dataset_size = len(self.train_dataset) effective_batch_size = batch_size * self.args.world_size * gradient_accumulation_steps total_steps = ceil(dataset_size / effective_batch_size) * num_epochs inputs["step_ratio"] = step / total_steps inputs["step"] = step outputs = model(**inputs) if self.state.epoch is None: raise RuntimeError("self.state.epoch is not set") current_epoch = self.state.epoch if current_epoch >= self.args.pmps_start_epoch: loss = outputs["loss"] else: loss = outputs["no_pmps_loss"] # Store custom metrics for logging at every step self._custom_metrics = { "train/ce_loss": outputs["ce_loss"], "train/distill_loss": outputs["distill_loss"], "train/ref_ce_loss": outputs["ref_ce_loss"], "train/pmps_loss": outputs.get("pmps_loss", 0), "train/pmps_loss_c2l": outputs.get("pmps_loss_c2l", 0), "train/pmps_loss_l2c": outputs.get("pmps_loss_l2c", 0), "train/code_entropy_latent": outputs.get("code_entropy_latent", 0), "train/code_entropy_cot": outputs.get("code_entropy_cot", 0), "train/match_entropy_c2l": outputs.get("match_entropy_c2l", 0), "train/match_order_correlation": outputs.get("match_order_correlation", 0), "train/position_bias_strength": outputs.get("position_bias_strength", 0), "train/diversity_loss": outputs.get("diversity_loss", 0), "train/latent_avg_cosine_sim": outputs.get("latent_avg_cosine_sim", 0), "train/total_loss": loss.item(), "train/no_pmps_loss": outputs["no_pmps_loss"].item() if isinstance(outputs["no_pmps_loss"], torch.Tensor) else outputs["no_pmps_loss"], } return loss def log(self, logs, start_time=None): # Merge custom metrics into the logs dict so they are all sent # to wandb (and other loggers) in a single call if hasattr(self, "_custom_metrics") and self._custom_metrics: logs.update(self._custom_metrics) self._custom_metrics = {} super().log(logs) def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict: """Tokenize a list of strings.""" tokenized_list = [ tokenizer( text, return_tensors="pt", padding="longest", max_length=256, truncation=True, return_attention_mask=False, ) for text in strings ] input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list] input_ids_lens = labels_lens = [ tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list ] return dict( input_ids=input_ids, labels=labels, input_ids_lens=input_ids_lens, labels_lens=labels_lens, ) def extract_answer_number(sentence: str) -> float: sentence = sentence.replace(",", "") pred = [s for s in re.findall(r"-?\d+\.?\d*", sentence)] if not pred: return float("inf") segment = [sentence] if len(segment) > 1: pred_answer = segment[1] pred_answer = [s for s in re.findall(r"-?\d+\.?\d*", pred_answer)] if len(pred_answer) > 0: pred_answer = pred_answer[0] else: pred_answer = float(pred[-1]) else: pred_answer = float(pred[-1]) if isinstance(pred_answer, str): try: pred_answer = float(pred_answer) except ValueError: pred_answer = float("inf") return pred_answer def train(): arg_parser = argparse.ArgumentParser(description="CODI Training with YAML config support") arg_parser.add_argument("--config_file", type=str, default=None, help="Path to YAML config file") args, remaining_args = arg_parser.parse_known_args() # If config_file is provided, load from YAML if args.config_file: print(f"Loading configuration from YAML file: {args.config_file}") with open(args.config_file, 'r') as f: config = yaml.safe_load(f) # Convert YAML config to command line arguments cmd_args = [] for key, value in config.items(): if value is None: continue if isinstance(value, bool): if value: cmd_args.extend([f"--{key}", "True"]) else: cmd_args.extend([f"--{key}", "False"]) else: cmd_args.extend([f"--{key}", str(value)]) # Parse with HfArgumentParser parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) model_args, data_args, training_args = parser.parse_args_into_dataclasses(args=cmd_args) else: # Original behavior: parse from command line print("Loading configuration from command line arguments") parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) model_args, data_args, training_args = parser.parse_args_into_dataclasses() ########################## # Peft Model # ########################## if model_args.lora_init: task_type = TaskType.CAUSAL_LM if any(name in model_args.model_name_or_path.lower() for name in ["llama", "mistral", "falcon", "qwen"]): target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"] elif any(name in model_args.model_name_or_path.lower() for name in ["phi"]): target_modules = ["q_proj", "k_proj", "v_proj", "dense", "fc1", "fc2"] elif any(name in model_args.model_name_or_path.lower() for name in ["gpt2"]): target_modules = ["c_attn", "c_proj", "c_fc"] else: raise ValueError(f"Only support LLAMA, Mistral, Falcon, Phi-2, but got {model_args.model_name_or_path}.") lora_config = LoraConfig( task_type=task_type, inference_mode=False, r=model_args.lora_r, lora_alpha=model_args.lora_alpha, lora_dropout=0.1, target_modules=target_modules, init_lora_weights=True, ) model = CODI_PMPS(model_args, training_args, lora_config) tokenizer = transformers.AutoTokenizer.from_pretrained( model_args.model_name_or_path, token=model_args.token, cache_dir=training_args.cache_dir, model_max_length=training_args.model_max_length, padding_side="right", ) if tokenizer.pad_token_id is None: tokenizer.add_special_tokens({"pad_token": "[PAD]"}) tokenizer.pad_token_id = model.pad_token_id if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("[PAD]") def get_answer_token_position(tokens, answer_prompts, tokenizer): try: match_indices = ( (tokens.unfold(0, len(answer_prompts[0]), 1) == answer_prompts[0]) .all(dim=1) .nonzero(as_tuple=True)[0] .item() ) answer_token_id = match_indices + len(answer_prompts[0]) return answer_token_id except Exception: breakpoint() def preprocess( sources: Sequence[str], targets: Sequence[str], answers: Sequence[str], tokenizer: transformers.PreTrainedTokenizer, bot_id: int, eot_id: int, ) -> Dict: """ Preprocess data for CODI_CT training. Key difference from codi_ori: we also compute cot_start_position and cot_end_position to identify the CoT token range within ref_input_ids. These positions are used to extract CoT hidden states for CT loss. """ print("Tokenizing inputs... This may take some time...") sources_id = _tokenize_fn(sources, tokenizer)["input_ids"] cot_id = _tokenize_fn(targets, tokenizer)["input_ids"] answers_id = _tokenize_fn(answers, tokenizer)["input_ids"] # Add eos token to accommodate pretrained model's format if not training_args.remove_eos: sources_id = [torch.tensor(x.numpy().tolist() + [tokenizer.eos_token_id], dtype=torch.long) for x in sources_id] cot_id = [torch.tensor(x.numpy().tolist() + [tokenizer.eos_token_id], dtype=torch.long) for x in cot_id] answers_id = [torch.tensor(x.numpy().tolist() + [tokenizer.eos_token_id], dtype=torch.long) for x in answers_id] if cot_id[0][0] == tokenizer.bos_token_id: cot_id = [x[1:] for x in cot_id] answers_id = [x[1:] for x in answers_id] ref_input_ids = [ torch.cat([x, y, z]).to(torch.long) for x, y, z in zip(sources_id, cot_id, answers_id) ] ref_labels = [] for x, y in zip(ref_input_ids, sources_id): z = x.clone() z[: len(y)] = -100 ref_labels.append(z) # ===== NEW: Compute CoT position range in ref_input_ids ===== # ref_input_ids = [question_tokens | cot_tokens | answer_tokens] # cot_start = len(question_tokens), cot_end = len(question_tokens) + len(cot_tokens) cot_start_position = [len(src) for src in sources_id] cot_end_position = [len(src) + len(cot) for src, cot in zip(sources_id, cot_id)] # Add eot to source sources_id = [ torch.tensor(x.numpy().tolist() + [bot_id], dtype=torch.long) for x in sources_id ] # Add eot and eos if training_args.remove_eos: answers_id = [ torch.tensor([eot_id] + x.numpy().tolist(), dtype=torch.long) for x in answers_id ] else: answers_id = [ torch.tensor([eot_id, tokenizer.eos_token_id] + x.numpy().tolist(), dtype=torch.long) for x in answers_id ] answer_prompts = [ torch.tensor(tokenizer.encode("The answer is:")), torch.tensor(tokenizer.encode("The next step result is:")), ] if answer_prompts[0][0] == tokenizer.bos_token_id: answer_prompts[0] = answer_prompts[0][1:] answer_prompts[1] = answer_prompts[1][1:] ref_answer_position = [ get_answer_token_position(x, answer_prompts, tokenizer) for i, x in enumerate(ref_input_ids) ] model_answer_position = [ get_answer_token_position(x, answer_prompts, tokenizer) for x in answers_id ] ref_eos_position = [len(x) - 1 for x in ref_input_ids] model_eos_position = [len(x) - 1 for x in answers_id] return dict( encoder_input_ids=sources_id, decoder_input_ids=answers_id, ref_input_ids=ref_input_ids, labels=answers_id, ref_answer_position=ref_answer_position, model_answer_position=model_answer_position, ref_eos_position=ref_eos_position, model_eos_position=model_eos_position, ref_labels=ref_labels, cot_start_position=cot_start_position, # NEW cot_end_position=cot_end_position, # NEW ) class SupervisedDataset(Dataset): QUESTION_PROMPT = "\nAnswer the above question. First think step by step and then answer the final number.\n" QUESTION_DA_PROMPT = "\nAnswer the above question. Answer the final number directly in one number.\n" def __init__(self, data_name, raw_data, tokenizer, bot, eot): super(SupervisedDataset, self).__init__() logging.warning("Formatting inputs...") self.data_name = data_name questions, cots, answers = [], [], [] if "math" in self.data_name: print("loading MATH training data") raw_data = read_json('/data/oss_bucket_0/liqi/dataset/MATH/train_clean.json') # raw_data is already loaded from data_path in make_supervised_data_module elif "icot-full" in self.data_name: print("loading gsm8k-aug-nl training data") raw_data = read_json('/data/oss_bucket_0/liqi/dataset/gsm8k-aug-nl/train.json') else: print("loading gsm9k-aug training data") raw_data = read_json('/data/oss_bucket_0/liqi/gsk8k-aug/gsm_train.json') for num_iter, example in enumerate(raw_data): if training_args.exp_mode and num_iter > training_args.exp_data_num: break if "icot" in self.data_name and "full" in self.data_name: # icot-full (GSM8k-Aug-NL) question = f"{example['question']}" if example["answer"] is None: continue example['cot'] = ' '.join(example['steps']) token_num = len(tokenizer.encode(example["question"] + example["cot"] + example["answer"])) if token_num > training_args.max_token_num: continue cot = f"{example['cot']}".split(". ") if not training_args.include_last_cot: cot = cot[:-1] answer = example["answer"].split(" ")[-1] if not answer[0].isdigit(): continue answer = f"The answer is: {answer}" answer = answer.replace("####", "") questions.append(question) if cot: cot = ". ".join(cot) + ".\n" else: cot = "" cots.append(cot) answers.append(answer) elif "icot" in self.data_name: # icot (GSM8k-Aug) question = f"{example['question']}" if "cot" not in example: example["cot"] = example["steps"] example['cot'] = ' '.join(example['cot']) token_num = len(tokenizer.encode(example["question"] + example["cot"] + example["answer"])) if token_num > training_args.max_token_num: continue cot = f"{example['cot']}".split(" ") if not training_args.include_last_cot: cot = cot[:-1] # if len(cot) == 0: # continue answer = example['answer'].split(' ')[-1] # some answers startwith the negative sign (-), bringing distillation problems for LLaMA if not answer[0].isdigit(): continue answer = f"The answer is: {answer}" answer = answer.replace("####", "") questions.append(question) cots.append(" ".join(cot)) answers.append(answer) elif "commonsense" in self.data_name or "strategy" in self.data_name: question = example["question"].strip() + "\n" cot = example["cot"].strip() + "\n" answer = f"The answer is: {str(example['answer']).strip()}" token_num = len(tokenizer.encode(question + " " + cot + " " + answer)) if token_num > training_args.max_token_num: continue questions.append(question) cots.append(cot) answers.append(answer) elif "prontoqa" in data_args.data_name: question = example["question"].strip() + "\n" cot = "\n".join(example["steps"][:-1]) + "\n" answer = f"The answer is: {str(example['answer']).strip()}" token_num = len(tokenizer.encode(question + " " + cot + " " + answer)) if token_num > training_args.max_token_num: continue questions.append(question) cots.append(cot) answers.append(answer) elif "math" in self.data_name: # MATH dataset: problem -> question, solution -> full CoT, answer -> final answer question = example["problem"].strip() + "\n" cot = example["solution"].strip() + "\n" # Keep full CoT, no step deletion answer = f"The answer is: {str(example['answer']).strip()}" token_num = len(tokenizer.encode(question + " " + cot + " " + answer)) if token_num > training_args.max_token_num: continue questions.append(question) cots.append(cot) answers.append(answer) else: raise NotImplementedError if training_args.exp_mode: questions = questions[: training_args.exp_data_num] cots = cots[: training_args.exp_data_num] answers = answers[: training_args.exp_data_num] print(f"{len(cots)} data in total...") logging.warning("Tokenizing inputs... This may take some time...") self.data_dict = preprocess(questions, cots, answers, tokenizer, bot, eot) self.keys = list(self.data_dict.keys()) def __len__(self): return len(self.data_dict["encoder_input_ids"]) def __getitem__(self, i) -> Dict[str, torch.Tensor]: return {key: self.data_dict[key][i] for key in self.keys} @dataclass class DataCollatorForSupervisedDataset(object): """Collate examples for supervised fine-tuning with CT support.""" tokenizer: transformers.PreTrainedTokenizer def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]: ( encoder_input_ids, decoder_input_ids, ref_input_ids, labels, ref_answer_position, model_answer_position, ref_labels, cot_start_position, cot_end_position, ) = tuple( [instance[key] for instance in instances] for key in ( "encoder_input_ids", "decoder_input_ids", "ref_input_ids", "labels", "ref_answer_position", "model_answer_position", "ref_labels", "cot_start_position", "cot_end_position", ) ) # Pad left for encoder reversed_input_ids = [seq.flip(0) for seq in encoder_input_ids] encoder_input_ids = torch.nn.utils.rnn.pad_sequence( reversed_input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id ).flip(1) # Pad right for ref and decoder ref_input_ids = torch.nn.utils.rnn.pad_sequence( ref_input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id ) ref_labels = torch.nn.utils.rnn.pad_sequence( ref_labels, batch_first=True, padding_value=IGNORE_INDEX ) decoder_input_ids = torch.nn.utils.rnn.pad_sequence( decoder_input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id ) labels = torch.nn.utils.rnn.pad_sequence( labels, batch_first=True, padding_value=IGNORE_INDEX ) return dict( encoder_input_ids=encoder_input_ids, decoder_input_ids=decoder_input_ids, ref_input_ids=ref_input_ids, labels=labels, encoder_attention_mask=encoder_input_ids.ne(self.tokenizer.pad_token_id), ref_answer_position=torch.tensor(ref_answer_position, dtype=torch.long), model_answer_position=torch.tensor(model_answer_position, dtype=torch.long), ref_attention_mask=ref_input_ids.ne(self.tokenizer.pad_token_id), ref_labels=ref_labels, cot_start_position=torch.tensor(cot_start_position, dtype=torch.long), # NEW cot_end_position=torch.tensor(cot_end_position, dtype=torch.long), # NEW ) def make_supervised_data_module(tokenizer, data_args) -> Dict: """Make dataset and collator for supervised fine-tuning.""" logging.warning("Downloading Data") if "icot" in data_args.data_name: if "full" in data_args.data_name: dataset = None else: dataset = None train_dataset = SupervisedDataset( data_name=data_args.data_name, raw_data=dataset, tokenizer=tokenizer, bot=model.bot_id, eot=model.eot_id, ) data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator) elif "strategy" in data_args.data_name: dataset = load_dataset("zen-E/StrategyQA_CoT_GPT4o")["train"] train_dataset = SupervisedDataset( data_name=data_args.data_name, raw_data=dataset, tokenizer=tokenizer, bot=model.bot_id, eot=model.eot_id, ) data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator) elif "commonsense" in data_args.data_name: dataset = load_dataset("zen-E/CommonsenseQA-GPT4omini")["train"] train_dataset = SupervisedDataset( data_name=data_args.data_name, raw_data=dataset, tokenizer=tokenizer, bot=model.bot_id, eot=model.eot_id, ) data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator) elif "prontoqa" in data_args.data_name: with open("/home/ubuntu/coconut/data/prontoqa_train.json") as f: dataset = json.load(f) train_dataset = SupervisedDataset( data_name=data_args.data_name, raw_data=dataset, tokenizer=tokenizer, bot=model.bot_id, eot=model.eot_id, ) data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator) elif "math" in data_args.data_name: # assert data_args.data_path is not None, "data_path must be provided for MATH dataset" # dataset = read_json(data_args.data_path) dataset = None train_dataset = SupervisedDataset( data_name=data_args.data_name, raw_data=dataset, tokenizer=tokenizer, bot=model.bot_id, eot=model.eot_id, ) data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator) else: raise NotImplementedError(f"Dataset {data_args.data_name} is not supported.") training_args.output_dir = os.path.join( training_args.output_dir, training_args.expt_name ) # ===== Save all input arguments to a JSON file ===== os.makedirs(training_args.output_dir, exist_ok=True) from dataclasses import asdict all_args = { "model_args": asdict(model_args), "data_args": asdict(data_args), "training_args": {k: v for k, v in asdict(training_args).items() if not k.startswith("_")}, } args_save_path = os.path.join(training_args.output_dir, "training_arguments.json") with open(args_save_path, "w", encoding="utf-8") as f: json.dump(all_args, f, ensure_ascii=False, indent=4, default=str) print(f"All training arguments saved to {args_save_path}") # ===== Wandb Initialization ===== if training_args.use_wandb and int(os.environ.get("RANK", 0)) == 0: if training_args.wandb_api_key: wandb.login(key=training_args.wandb_api_key) run_name = training_args.wandb_run_name or training_args.expt_name wandb.init( project=training_args.wandb_project, name=run_name, config={ "model_name_or_path": model_args.model_name_or_path, "lora_r": model_args.lora_r, "lora_alpha": model_args.lora_alpha, "num_latent": training_args.num_latent, "learning_rate": training_args.learning_rate, "per_device_train_batch_size": training_args.per_device_train_batch_size, "gradient_accumulation_steps": training_args.gradient_accumulation_steps, "num_train_epochs": training_args.num_train_epochs, "pmps_loss_factor": training_args.pmps_loss_factor, "pmps_temperature": training_args.pmps_temperature, "pmps_match_temperature": training_args.pmps_match_temperature, "pmps_nmb_prototypes": training_args.pmps_nmb_prototypes, "pmps_use_projector": training_args.pmps_use_projector, "pmps_layer": training_args.pmps_layer, "pmps_start_epoch": training_args.pmps_start_epoch, "pmps_position_bias_max": training_args.pmps_position_bias_max, "pmps_position_bias_schedule": training_args.pmps_position_bias_schedule, "diversity_loss_factor": training_args.diversity_loss_factor, "distill_loss_factor": training_args.distill_loss_factor, }, ) training_args.report_to = ["wandb"] print("Wandb initialized successfully.") elif not training_args.use_wandb: os.environ["WANDB_DISABLED"] = "true" data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args) trainer = CustomTrainer(model=model, tokenizer=tokenizer, args=training_args, **data_module) trainer.train() trainer.save_state() trainer.save_model(output_dir=training_args.output_dir) if __name__ == "__main__": train()