From 6dd8f13c28cbdaaed43606b9a5f8904f35b9df9f Mon Sep 17 00:00:00 2001 From: 51616 Date: Mon, 20 Jan 2025 09:19:45 +0000 Subject: [PATCH] add finetuning + gte as ctx encoder + fix ctx input decode --- intx_sft.py | 81 +++++++++++++++++++------------ src/ctx_to_lora/configs.py | 4 ++ src/ctx_to_lora/eval.py | 16 +++--- src/ctx_to_lora/model_loading.py | 13 +++-- src/ctx_to_lora/modeling_utils.py | 13 +++-- 5 files changed, 83 insertions(+), 44 deletions(-) diff --git a/intx_sft.py b/intx_sft.py index 8ef5ef9..fbd135c 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -257,7 +257,7 @@ def main(): run_name = ( get_run_name(seed_str=training_args.logging_dir.strip("runs/")) if not checkpoint_dir - else checkpoint_dir.split("/")[-2] + else checkpoint_dir.strip("/").split("/")[-2] ) output_dir = f"train_outputs/runs/{run_name}" @@ -296,40 +296,59 @@ def main(): ############ Model setup - model_name = model_args.model_name_or_path - model, tokenizer = get_model_and_tokenizer( - **vars(model_args), - train=True, - requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE, - peft_config=get_lora_config(model_name, **vars(lora_args)), - ) - ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path - if ctx_name is not None: - ctx_encoder_model_config = AutoConfig.from_pretrained(ctx_name) - if "Llama" in ctx_name and "Vision" in ctx_name: - ctx_encoder_model_config = ctx_encoder_model_config.text_config - ctx_tokenizer = get_tokenizer(ctx_name) - else: - ctx_encoder_model_config = model.config - ctx_tokenizer = tokenizer + if not ctx_args.from_pretrained_checkpoint: + model_name = model_args.model_name_or_path + model, tokenizer = get_model_and_tokenizer( + **vars(model_args), + train=True, + requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE, + peft_config=get_lora_config(model_name, **vars(lora_args)), + ) + ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path + if ctx_name is not None: + ctx_encoder_model_config = AutoConfig.from_pretrained( + ctx_name, trust_remote_code=True + ) + if "Llama" in ctx_name and "Vision" in ctx_name: + ctx_encoder_model_config = ctx_encoder_model_config.text_config + ctx_tokenizer = get_tokenizer(ctx_name) + else: + ctx_encoder_model_config = model.config + ctx_tokenizer = tokenizer if ctx_args.exp_setup == ExperimentSetup.HYPER_LORA: logger.info("Using HyperLoRA") - hypernet_config = get_hypernet_config( - model, ctx_encoder_model_config, hypernet_args, aggregator_args - ) - if ctx_encoder_args.layer_idx is None: - ctx_encoder_args.layer_idx = ctx_encoder_model_config.num_hidden_layers // 4 - logger.info( - f"Using the first {ctx_encoder_args.layer_idx} layers" - " as the context encoder" + if not ctx_args.from_pretrained_checkpoint: + hypernet_config = get_hypernet_config( + model, ctx_encoder_model_config, hypernet_args, aggregator_args ) - model = ModulatedPretrainedModel( - model, - hypernet_config, - ctx_encoder_args, - ctx_args.use_kl_loss, - ) + if ctx_encoder_args.layer_idx is None: + ctx_encoder_args.layer_idx = ( + ctx_encoder_model_config.num_hidden_layers // 4 + ) + logger.info( + f"Using the first {ctx_encoder_args.layer_idx} layers" + " as the context encoder" + ) + model = ModulatedPretrainedModel( + model, + hypernet_config, + ctx_encoder_args, + ctx_args.use_kl_loss, + ) + else: + logger.info( + f"Loading from checkpoint: {ctx_args.from_pretrained_checkpoint}" + ) + model = ModulatedPretrainedModel.from_state_dict( + torch.load(open(ctx_args.from_pretrained_checkpoint, "rb")), + train=True, + ) + tokenizer = get_tokenizer(model.base_model.config.name_or_path) + ctx_name = model.ctx_encoder_args.ctx_encoder_model_name_or_path + if ctx_name is None: + ctx_name = model.base_model.config.name_or_path + ctx_tokenizer = get_tokenizer(ctx_name) if len([p for p in model.ctx_encoder.parameters() if p.requires_grad]): raise ValueError("ctx_encoder contains trainable parameters") if len([p for p in model.base_model.parameters() if p.requires_grad]): diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index e98f21c..407fe60 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -228,6 +228,10 @@ class ModelArguments: default=None, metadata={"help": ("Base model name or path.")}, ) + use_flash_attn: bool = field( + default=True, + metadata={"help": "Whether to use flash attention."}, + ) # use_peft: bool = field( # default=False, # metadata={"help": ("Whether to use PEFT or not for training.")}, diff --git a/src/ctx_to_lora/eval.py b/src/ctx_to_lora/eval.py index 041c422..bbebdea 100644 --- a/src/ctx_to_lora/eval.py +++ b/src/ctx_to_lora/eval.py @@ -143,7 +143,7 @@ def save_generated_text(samples, output_dir, split): f.write(json.dumps(sample) + "\n") -def decode_test_result(test_dataset, test_result, tokenizer): +def decode_test_result(test_dataset, test_result, tokenizer, ctx_tokenizer): out = [] for sample, pred_toks in zip(test_dataset, test_result.predictions): # sample and labels are not padded @@ -168,14 +168,16 @@ def decode_test_result(test_dataset, test_result, tokenizer): d["input"] = tokenizer.decode(input_toks, skip_special_tokens=True) d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True) if "ctx_ids" in sample: - d["context"] = tokenizer.decode(sample["ctx_ids"], skip_special_tokens=True) + d["context"] = ctx_tokenizer.decode( + sample["ctx_ids"], skip_special_tokens=True + ) out.append(d) return out @torch.no_grad() -def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs): +def eval_generation(eval_trainer, tokenizer, ctx_tokenizer, datasets, split, gen_kwargs): if not isinstance(datasets, dict): datasets = {"": datasets} out = {} @@ -186,7 +188,7 @@ def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs): metric_key_prefix=split_name, **gen_kwargs, ) - decoded_txts = decode_test_result(ds, eval_result, tokenizer) + decoded_txts = decode_test_result(ds, eval_result, tokenizer, ctx_tokenizer) rouge_metrics = compute_rouge( [txt["generated"] for txt in decoded_txts], [txt["label"] for txt in decoded_txts], @@ -428,7 +430,9 @@ def evaluate(checkpoint_path, args, split, generative): out.update(metrics) else: eval_trainer = Seq2SeqTrainer(**trainer_kwargs) - metrics = eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs) + metrics = eval_generation( + eval_trainer, tokenizer, ctx_tokenizer, datasets, split, gen_kwargs + ) out.update(metrics) clear_gpu() @@ -466,4 +470,4 @@ if __name__ == "__main__": shutil.rmtree(args.output_dir) evaluate(checkpoint_path, args, split="validation", generative=False) - # evaluate(checkpoint_path, args, split="validation", generative=True) + evaluate(checkpoint_path, args, split="validation", generative=True) diff --git a/src/ctx_to_lora/model_loading.py b/src/ctx_to_lora/model_loading.py index c421072..915fdf5 100644 --- a/src/ctx_to_lora/model_loading.py +++ b/src/ctx_to_lora/model_loading.py @@ -38,7 +38,8 @@ def get_model_and_tokenizer( ) tokenizer = get_tokenizer(model_name_or_path, tokenizer_kwargs, peft_config, train) model.config.pad_token_id = tokenizer.pad_token_id - model.generation_config.pad_token_id = tokenizer.pad_token_id + if getattr(model, "generation_config", None): + model.generation_config.pad_token_id = tokenizer.pad_token_id return model, tokenizer @@ -65,6 +66,7 @@ def get_tokenizer( model_name_or_path, padding_side=padding_side, truncation_side=truncation_side, + trust_remote_code=True, **tokenizer_kwargs, ) @@ -128,17 +130,22 @@ def get_model( if model_kwargs is not None: model_init_kwargs.update(model_kwargs) if use_flash_attn: - if "gemma" not in model_name_or_path: + if ("gemma" not in model_name_or_path) and ("gte" not in model_name_or_path): model_init_kwargs["attn_implementation"] = "flash_attention_2" if is_vision_model: model_init_kwargs["attn_implementation"] = "sdpa" model_init_kwargs.pop("use_cache") # for training disable cache + elif "gte" in model_name_or_path: + model_init_kwargs["attn_implementation"] = "sdpa" + model_init_kwargs["torch_dtype"] = torch.float32 + model_init_kwargs.pop("use_cache") # if train and not is_vision_model: # model_init_kwargs["use_cache"] = False logger.debug(f"Model init kwargs: {model_init_kwargs}") if not is_vision_model: - model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs) + # model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs) + model = AutoModel.from_pretrained(**model_init_kwargs) else: model = MllamaForConditionalGeneration.from_pretrained(**model_init_kwargs) model = model.language_model diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index da6e487..90d6e26 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -344,7 +344,10 @@ class EarlyExit(nn.Module): def __init__(self, base_model: PreTrainedModel, exit_layer: int): super().__init__() self.base_model = base_model - self.base_model.layers = base_model.layers[:exit_layer] + if "gte" in base_model.config.name_or_path: + self.base_model.encoder.layer = base_model.encoder.layer[:exit_layer] + else: + self.base_model.layers = base_model.layers[:exit_layer] # self.exit_layer = exit_layer @torch.no_grad() @@ -612,7 +615,7 @@ class ModulatedPretrainedModel(nn.Module): # if any(module_name in name for module_name in trainable_base_modules): # param.requires_grad = True - self.hypernet = HyperLoRA(self.hypernet_config).to(self.device) + self.hypernet = HyperLoRA(self.hypernet_config).to(self.device).to(torch.float32) # ctx_encoder_name = # if self.ctx_encoder_args.ctx_encoder_model_name_or_path is not None: # encoder_model = get_model( @@ -627,12 +630,14 @@ class ModulatedPretrainedModel(nn.Module): ctx_model_name = self.base_model.config.name_or_path # use an explicit copy of the base model # for using with "modules_to_save" + base_model_attn_impl = self.base_model.config._attn_implementation + logger.debug(f"ctx_model_name: {ctx_model_name}") + logger.debug(f"base_model.config._attn_implementation: {base_model_attn_impl}") encoder_model = get_model( ctx_model_name, train=True, requires_grad=False, - use_flash_attn=self.base_model.config._attn_implementation - == "flash_attention_2", + use_flash_attn=base_model_attn_impl == "flash_attention_2", ) self.ctx_encoder = EarlyExit( get_base_model(encoder_model), self.ctx_encoder_args.layer_idx