trainable hyperlora + context numbers

hyperlora/configs.py

@@ -113,7 +113,7 @@ class ArgumentParser(HfArgumentParser):
         return output
 
 
-class ExperimentSetup(Enum):
+class ExperimentSetup(str, Enum):
     LORA = "lora"
     HYPER_LORA = "hyper_lora"
     FULL_FINETUNE = "full_finetune"
@@ -157,3 +157,10 @@ class CtxTrainingArguments:
         default=ExperimentSetup.LORA,
         metadata={"help": "Experiment setup - LoRA, HyperLoRA, or full finetuning"},
     )
+
+
+if __name__ == "__main__":
+    print(ExperimentSetup)
+    print(ExperimentSetup.LORA)
+    print(ExperimentSetup.HYPER_LORA)
+    print(ExperimentSetup.FULL_FINETUNE)

hyperlora/data_utils.py

@@ -326,6 +326,17 @@ def tokenize_chat_messages(
     return conversation_ids
 
 
+def tokenize_ctx_text(
+    example: dict[str, Any],
+    tokenizer: PreTrainedTokenizerBase,
+) -> dict[str, Any]:
+    text = example["context"]
+    tokenized_text = tokenizer(text)
+    ctx_ids = tokenized_text["input_ids"]
+    ctx_attn_mask = tokenized_text["attention_mask"]
+    return dict(ctx_ids=ctx_ids, ctx_attn_mask=ctx_attn_mask)
+
+
 if __name__ == "__main__":
     from transformers import AutoTokenizer
 

hyperlora/hooks.py

@@ -70,6 +70,8 @@ def apply_hook_to_layer(
 
     if mname in ["q_proj", "k_proj", "v_proj", "o_proj", "qkv_proj"]:
         mname = f"self_attn.{mname}"
+    elif mname in ["down_proj", "up_proj", "gate_proj"]:
+        mname = f"mlp.{mname}"
     module = attrgetter(mname)(layer)
 
     if pre_hook is not None:

hyperlora/intx_sft.py

@@ -1,3 +1,6 @@
+from copy import copy
+from functools import partial
+from importlib.resources import read_binary
 import logging
 import random
 import string
@@ -10,10 +13,11 @@ from data_utils import (
     get_preprocessing_fn,
     get_sft_prompt_formatting_fn,
     tokenize_chat_messages,
+    tokenize_ctx_text,
 )
 from datasets import load_dataset
 from model_loading import get_lora_config, get_model_and_tokenizer
-from modeling_utils import ModulatedPretrainedModel
+from modeling_utils import HyperLoRA, ModulatedPretrainedModel, get_hypernet_config
 from training_utils import TRAINING_TASK, train_model
 from transformers import (
     AutoModelForCausalLM,
@@ -54,6 +58,14 @@ def compute_metrics(eval_pred: EvalPrediction) -> dict:
     return {"per_token_acc": acc, "num_valid_tokens": indices[0].size}
 
 
+def get_run_name():
+    uuid = "".join(
+        [random.choice(string.ascii_letters + string.digits) for _ in range(8)]
+    )
+    run_name = time.strftime("%Y%m%d-%H%M%S") + f"_{uuid}"
+    return run_name
+
+
 def main():
     # Set logging verbosity to INFO
     logging.basicConfig(level=logging.INFO)
@@ -63,28 +75,31 @@ def main():
     )
     ctx_args, model_args, lora_args, training_args = parser.parse()
 
-    uuid = "".join(
-        [random.choice(string.ascii_letters + string.digits) for _ in range(8)]
-    )
-    run_name = time.strftime("%Y%m%d-%H%M%S") + f"_{uuid}"
-
+    run_name = get_run_name()
     training_args.run_name = run_name
     training_args.output_dir = f"train_outputs/{run_name}"
     training_args.logging_dir = f"train_outputs/{run_name}"
 
-    model_name = model_args.model_name_or_path
+    logger.info(f"Run name: {run_name}")
+    logger.info(f"ctx_args: {ctx_args}")
+    logger.info(f"model_args: {model_args}")
+    logger.info(f"lora_args: {lora_args}")
 
+    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)),
     )
+
     if ctx_args.exp_setup == ExperimentSetup.HYPER_LORA:
-        hypernet = ...
-        model = ModulatedPretrainedModel(model, hypernet)
+        logger.info("Using HyperLoRA")
+        hypernet = HyperLoRA(get_hypernet_config(model)).to(model.device)
+        model = ModulatedPretrainedModel(model, hypernet).to(model.device).train()
     else:
         # activate LoRA
+        logger.info("Using LoRA")
         model.set_adapter("default")
 
     print(model)
@@ -108,6 +123,7 @@ def main():
     # add "chat" field
     ds = ds.map(get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer))
     # tokenize the chat + mask the assistant inputs
+    pre_tok_cols = copy(ds["train"].column_names)
     tokenized_ds = ds.map(
         tokenize_chat_messages,
         fn_kwargs={
@@ -117,22 +133,63 @@ def main():
                 "max_length": None,
             },
         },
-        remove_columns=ds["train"].column_names,
     )
+
+    # computes ctx_features offline when using hyperlora
+    if isinstance(model, ModulatedPretrainedModel):
+        # TODO: can we batch this?
+        tokenized_ds = tokenized_ds.map(
+            tokenize_ctx_text, fn_kwargs={"tokenizer": tokenizer}
+        )
+        tokenized_ds = tokenized_ds.map(
+            model.get_ctx_features,
+            remove_columns=["ctx_ids"],
+        )
+
+    tokenized_ds = tokenized_ds.remove_columns(pre_tok_cols)
+
+    validate_columns(tokenized_ds)
+
     train_ds = tokenized_ds["train"]
     eval_ds = {
         "train": tokenized_ds["train"].select(range(100)),
         "val": tokenized_ds["eval"],
     }
 
-    # TODO: computes ctx_features offline
-
     # DataCollatorForSeq2Seq also pads the `labels`
     # useful when we're computing the labels manually
     # or masking the loss only on completion
     # TODO: change to a faster collator? e.g.,
     # https://huggingface.co/blog/packing-with-FA2
-    data_collator = DataCollatorForSeq2Seq(tokenizer, model, pad_to_multiple_of=8)
+    # data_collator = DataCollatorForSeq2Seq(tokenizer, model, pad_to_multiple_of=8)
+
+    # TODO: check tokenization pipeline (if prompt + ctx are tokenized correctly)
+    def collator(inp_list, tokenizer):
+        # input is a list of tokenized sequences
+        padding_kwargs = dict(padding=True, pad_to_multiple_of=8, return_tensors="pt")
+        labels = [x.pop("labels") for x in inp_list]
+        ctx_features = None
+        if "ctx_features" in inp_list[0]:
+            # TODO: also pad ctx_features
+            # have to be manual since it has [bs, ctx_len, features] shape
+            # => pad to the max ctx_len in the batch with zeros
+
+            # HACK: assumes ctx_features with the same size
+            # only works with context_numbers
+            ctx_features = torch.tensor([x.pop("ctx_features") for x in inp_list])
+
+        padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
+
+        # hacky explicit padding since the labels are not padded by default
+        labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
+        labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
+        out = {**padded_seq, "labels": labels}
+        if ctx_features is not None:
+            out["ctx_features"] = ctx_features
+        # if task_descs:
+        #     task_descs = tokenizer.pad({"input_ids": task_descs}, **padding_kwargs)["input_ids"]
+        #     out["task_descs_ids"] = task_descs
+        return out
 
     # TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer
     # TODO: use packing with SFTTrainer
@@ -141,10 +198,19 @@ def main():
         train_ds,
         eval_ds,
         training_args,
-        data_collator,
+        partial(collator, tokenizer=tokenizer),
         compute_metrics,
     )
 
 
+def validate_columns(tokenized_ds):
+    ref_cols = set(
+        ["input_ids", "attention_mask", "labels", "ctx_features", "ctx_attn_mask"]
+    )
+    assert (
+        set(tokenized_ds["train"].column_names) == ref_cols
+    ), f"Columns mismatch: {set(tokenized_ds['train'].column_names)} != {ref_cols}"
+
+
 if __name__ == "__main__":
     main()

hyperlora/modeling_utils.py

@@ -242,6 +242,8 @@ class HyperLoRA(nn.Module):
             output_size=config.latent_size,
         )
 
+        # TODO: check initialization of the head
+        # default values are prob. way too big
         self.head = Mix(
             "bs n_layers n_modules d -> bs n_layers r out_d",
             weight_shape="n_modules d r out_d",
@@ -319,16 +321,23 @@ class ModulatedPretrainedModel(nn.Module):
         self.register_module("base_model", base_model)
         self.register_module("hypernet", hypernet)
 
+    def get_input_embeddings(self):
+        return self.base_model.get_input_embeddings()
+
     def get_ctx_features(
         self,
-        input_ids: Integer[Tensor, "bs seq_len"],
-        attention_mask: Optional[Integer[Tensor, "bs seq_len"]] = None,
+        examples: dict[str, Any],
     ):
+        out = dict()
+        input_ids = torch.tensor(examples.get("ctx_ids")).to(self.device)
+        attention_mask = torch.tensor(examples.get("ctx_attn_mask")).to(self.device)
+
         if isinstance(self.encoder, nn.Embedding):
             features = self.encoder(input_ids)
         else:
             features = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
-        return features
+        out["ctx_features"] = features
+        return out
 
     def forward(
         self,