From 586751a501e20f1958e79dcf2980e67b7197c477 Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 22 Apr 2025 02:57:56 +0000 Subject: [PATCH] add l1 reg on generated lora --- intx_sft.py | 3 +- src/ctx_to_lora/configs.py | 36 +++++++++-------- src/ctx_to_lora/modeling_utils.py | 18 +++++---- src/ctx_to_lora/training_utils.py | 64 ++++++++++++++++++++++++++++++- 4 files changed, 96 insertions(+), 25 deletions(-) diff --git a/intx_sft.py b/intx_sft.py index d551b7a..7a65c16 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -249,6 +249,7 @@ def main(): ] ) + set_seed(training_args.seed) checkpoint_dir = training_args.resume_from_checkpoint if checkpoint_dir and not os.path.isdir(checkpoint_dir): raise NotADirectoryError(f"Checkpoint{checkpoint_dir} is not a directory") @@ -295,7 +296,6 @@ def main(): args["deepspeed_plugin"] = None logger.debug(f"args: {args}") save_yaml(args, f"{output_dir}/args.yaml") - set_seed(training_args.seed) ############ Model setup @@ -341,6 +341,7 @@ def main(): ctx_encoder_args, ctx_args.use_kl_loss, ) + training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef else: logger.info( f"Loading from checkpoint: {ctx_args.from_pretrained_checkpoint}" diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index f4072e5..9118505 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -326,6 +326,10 @@ class CtxTrainingArguments: default=False, metadata={"help": "Whether to use KL loss."}, ) + gen_lora_l1_reg_coef: float = field( + default=0.0, + metadata={"help": "L1 regularization coefficient for generated LoRAs."}, + ) @dataclass @@ -418,30 +422,30 @@ class AggregatorArguments: # output_size: int # perceiver - attention_probs_dropout_prob: float = field( - default=0.0, - metadata={"help": "Attention dropout probability for Perceiver."}, - ) + # attention_probs_dropout_prob: float = field( + # default=0.0, + # metadata={"help": "Attention dropout probability for Perceiver."}, + # ) num_latent_factor: int = field( default=8, metadata={"help": "Number of latent factors for Perceiver."}, ) - num_blocks: int = field( - default=8, - metadata={"help": "Number of blocks for Perceiver."}, - ) + # num_blocks: int = field( + # default=8, + # metadata={"help": "Number of blocks for Perceiver."}, + # ) num_self_attends_per_block: int = field( default=6, metadata={"help": "Number of self-attends per block for Perceiver."}, ) - self_attention_widening_factor: int = field( - default=4, - metadata={"help": "Self-attention widening factor for Perceiver."}, - ) - cross_attention_widening_factor: int = field( - default=4, - metadata={"help": "Cross-attention widening factor for Perceiver."}, - ) + # self_attention_widening_factor: int = field( + # default=4, + # metadata={"help": "Self-attention widening factor for Perceiver."}, + # ) + # cross_attention_widening_factor: int = field( + # default=4, + # metadata={"help": "Cross-attention widening factor for Perceiver."}, + # ) decoder_depth: int = field( default=1, metadata={"help": "Decoder depth for Perceiver."}, diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 9c0f397..c5c686c 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -80,12 +80,12 @@ class AggregatorConfig: output_size: int # perceiver - attention_probs_dropout_prob: float = 0.0 - num_blocks: int = 1 - num_self_attends_per_block: int = 16 - decoder_depth: int = 1 # 1 = only cross-attention - self_attention_widening_factor: int = 4 - cross_attention_widening_factor: int = 1 + # attention_probs_dropout_prob: float = 0.0 + # num_blocks: int = 1 + num_self_attends_per_block: int # = 16 + decoder_depth: int # = 1 # 1 = only cross-attention + # self_attention_widening_factor: int = 4 + # cross_attention_widening_factor: int = 1 num_latent_factor: int = 8 lora_r: int = 8 per_rank_gen: bool = False @@ -1246,6 +1246,7 @@ class ModulatedPretrainedModel(nn.Module): ctx_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_len"]] = None, ctx_position_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, + return_generated_lora: Optional[bool] = False, # chat_ids: Optional[Integer[Tensor, "bs chat_len"]] = None, # chat_attn_mask: Optional[Integer[Tensor, "bs chat_len"]] = None, # chat_labels: Optional[Integer[Tensor, "bs chat_len"]] = None, @@ -1323,7 +1324,10 @@ class ModulatedPretrainedModel(nn.Module): ): model_outputs = self.base_model(*model_inputs_args, **model_inputs_kwargs) - return model_outputs + if return_generated_lora: + return model_outputs, (generated_loras, generated_layernorms) + else: + return model_outputs @torch.inference_mode() def generate( diff --git a/src/ctx_to_lora/training_utils.py b/src/ctx_to_lora/training_utils.py index 48ddf4f..f2783c8 100644 --- a/src/ctx_to_lora/training_utils.py +++ b/src/ctx_to_lora/training_utils.py @@ -22,6 +22,8 @@ from transformers.trainer_utils import ( has_length, seed_worker, ) +from transformers.trainer import _is_peft_model +from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES from transformers.trainer_pt_utils import LengthGroupedSampler from transformers.utils import is_datasets_available @@ -31,6 +33,60 @@ TRAINING_TASK = Enum("TRAINING_TASK", ["CAUSAL_LM", "COMPLETION"]) logger = logging.getLogger() +class ModulatedModelTrainer(Trainer): + def __init__(self, *args, **kwargs): + self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0) + super().__init__(*args, **kwargs) + + def compute_loss(self, model, inputs, return_outputs=False): + """ + How the loss is computed by Trainer. By default, all models return the loss in the first element. + + Subclass and override for custom behavior. + """ + if self.label_smoother is not None and "labels" in inputs: + labels = inputs.pop("labels") + else: + labels = None + + ##### get generated LoRAs + outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True) + ##### + + # Save past state if it exists + # TODO: this needs to be fixed and made cleaner later. + if self.args.past_index >= 0: + self._past = outputs[self.args.past_index] + + if labels is not None: + unwrapped_model = self.accelerator.unwrap_model(model) + if _is_peft_model(unwrapped_model): + model_name = unwrapped_model.base_model.model._get_name() + else: + model_name = unwrapped_model._get_name() + if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values(): + loss = self.label_smoother(outputs, labels, shift_labels=True) + else: + loss = self.label_smoother(outputs, labels) + else: + if isinstance(outputs, dict) and "loss" not in outputs: + raise ValueError( + "The model did not return a loss from the inputs, only the following keys: " + f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}." + ) + # We don't use .loss here since the model may return tuples instead of ModelOutput. + loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0] + + ##### unpack gen lora dict and compute regularization loss + if self.gen_lora_l1_reg_coef > 0.0: + reg_loss = 0.0 + for module, lora in gen_loras.items(): + reg_loss += lora["A"].abs().mean() + lora["B"].abs().mean() + loss += self.gen_lora_l1_reg_coef * reg_loss + ##### + return (loss, outputs) if return_outputs else loss + + def clear_gpu(): gc.collect() torch.cuda.empty_cache() @@ -58,7 +114,8 @@ def train_model( checkpoint = training_args.resume_from_checkpoint logger.info(f"Resuming from the checkpoint: {checkpoint}") - trainer_cls = Trainer + is_modulated_model = bool(getattr(model, "gen_lora_l1_reg_coef", 0)) + trainer_cls = Trainer if not is_modulated_model else ModulatedModelTrainer trainer_kwargs = dict( model=model, args=training_args, @@ -67,6 +124,11 @@ def train_model( data_collator=train_collator, compute_metrics=compute_metrics, ) + if is_modulated_model: + logger.info(f"Training with modulated model. Using CustomTrainer.") + trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.pop( + "gen_lora_l1_reg_coef" + ) trainer = trainer_cls(**trainer_kwargs)