add l1 reg on generated lora

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}"

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."},

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(

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)