1-layer attn ctx_encoder w/ early exit

hyperlora/intx_sft.py

@@ -21,7 +21,12 @@ from data_utils import (
 )
 from datasets import disable_caching, load_dataset
 from model_loading import get_lora_config, get_model_and_tokenizer
-from modeling_utils import HyperLoRA, ModulatedPretrainedModel, get_hypernet_config
+from modeling_utils import (
+    EarlyExit,
+    HyperLoRA,
+    ModulatedPretrainedModel,
+    get_hypernet_config,
+)
 from rouge_score import rouge_scorer
 from training_utils import TRAINING_TASK, train_model
 from transformers import (
@@ -34,6 +39,7 @@ from transformers import (
 )
 from utils import (
     extract_cli_args,
+    get_base_model,
     get_run_name,
     log_num_train_params,
     save_yaml,
@@ -158,7 +164,7 @@ def compute_metrics(
 #     return pred_ids
 
 
-def main(output_dir: str):
+def main(output_dir):
     ############ Argument parsing
     parser = ArgumentParser(
         (
@@ -205,21 +211,20 @@ def main(output_dir: str):
         logger.info("Using HyperLoRA")
         hypernet = HyperLoRA(get_hypernet_config(model)).to(model.device)
         # HACK: hardcode the embedding layer for now
-        # TODO: add explicit encoder
-        ctx_encoder = torch.nn.Embedding.from_pretrained(
-            model.get_input_embeddings().weight.clone(),
-            freeze=True,
-        )
-        model = (
-            ModulatedPretrainedModel(model, hypernet, ctx_encoder)
-            .to(model.device)
-            .train()
-        )
+        # TODO: add ctx_encoder config
+        # ctx_encoder = torch.nn.Embedding.from_pretrained(
+        #     model.get_input_embeddings().weight.clone(),
+        #     freeze=True,
+        # )
+        # TODO: have to use at least 1 layer bc positional embeddings
+        ctx_encoder = EarlyExit(get_base_model(model), 1)
+        model = ModulatedPretrainedModel(model, hypernet, ctx_encoder).to(model.device)
     else:
         # activate LoRA
         logger.info("Using LoRA")
         model.set_adapter("default")
 
+    model.train()
     logger.debug(model)
     log_num_train_params(model)
 

hyperlora/modeling_utils.py

@@ -200,6 +200,59 @@ AGGREGATOR_CLS = {
 }
 
 
+@contextmanager
+def early_exit(base_model: PreTrainedModel, exit_layer: int):
+    try:
+        layers = base_model.layers
+        base_model.layers = layers[:exit_layer]
+        yield base_model
+    finally:
+        base_model.layers = layers
+
+
+@contextmanager
+def maybe_add_batch_dim(kwargs):
+    try:
+        batched_input = False
+        batched_attn_mask = False
+        if len(kwargs["input_ids"].shape) == 1:
+            kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0)
+            batched_input = True
+        if len(kwargs["attention_mask"].shape) == 1:
+            kwargs["attention_mask"] = kwargs["attention_mask"].unsqueeze(0)
+            batched_attn_mask = True
+        yield batched_input, batched_attn_mask
+    finally:
+        if batched_input:
+            kwargs["input_ids"] = kwargs["input_ids"].squeeze(0)
+        if batched_attn_mask:
+            kwargs["attention_mask"] = kwargs["attention_mask"].squeeze(0)
+
+
+class EarlyExit(nn.Module):
+    def __init__(self, base_model: PreTrainedModel, exit_layer: int):
+        super().__init__()
+        self.base_model = base_model
+        self.exit_layer = exit_layer
+
+    def forward(self, **kwargs):
+        # if len(kwargs["input_ids"].shape) == 1:
+        #     kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0)
+        # if len(kwargs["attention_mask"].shape) == 1:
+        #     kwargs["attention_mask"] = kwargs["attention_mask"].unsqueeze(0)
+
+        with (
+            early_exit(self.base_model, self.exit_layer),
+            maybe_add_batch_dim(kwargs) as (batched_input, batched_attn_mask),
+        ):
+            model_outputs = self.base_model(**kwargs)
+
+        if batched_input:
+            model_outputs.last_hidden_state = model_outputs.last_hidden_state.squeeze(0)
+
+        return model_outputs.last_hidden_state
+
+
 class HyperLoRA(nn.Module):
     def __init__(
         self,
@@ -312,14 +365,7 @@ class ModulatedPretrainedModel(nn.Module):
         ctx_encoder: nn.Module,  # TODO: maybe just use the name/config?
     ):
         super().__init__()
-        # self.base_model = base_model
         self.device = base_model.device
-        # NOTE: we can reduce extractor_lm's depth by
-        # e.g., self.extractor_lm.model.layers = self.extractor_lm.model.layers[:num_extractor_layers]
-        # or just use the token embeddings extractor_lm.embed_tokens(input_ids)
-
-        # self.ctx_encoder = ctx_encoder  # base_model.get_input_embeddings()
-
         # register base_model as a submodule
         self.register_module("base_model", base_model)
         self.register_module("hypernet", hypernet)
@@ -337,14 +383,20 @@ class ModulatedPretrainedModel(nn.Module):
         state_dict = super().state_dict(*args, **kwargs)
         # remove non-trainable and base_model's params
         for name, param in self.named_parameters():
-            if (not param.requires_grad) or name.startswith("base_model"):
+            if not param.requires_grad:
                 state_dict.pop(name, None)
+
+        for name in list(state_dict.keys()):
+            if name.startswith("base_model") or name.startswith("ctx_encoder"):
+                state_dict.pop(name, None)
+
         return state_dict
 
     def load_state_dict(self, state_dict: dict, *args, **kwargs):
         # NOTE: might have to set `strict=False` as we don't save all the params
         return super().load_state_dict(state_dict, *args, **kwargs)
 
+    @torch.no_grad()
     def get_ctx_features(
         self,
         examples: dict[str, Any],
@@ -355,7 +407,7 @@ class ModulatedPretrainedModel(nn.Module):
         attention_mask = torch.tensor(examples.get("ctx_attn_mask")).to(self.device)
 
         if isinstance(self.ctx_encoder, nn.Embedding):
-            features = self.ctx_encoder(input_ids)
+            features = self.ctx_encoder(input_ids[torch.where(attention_mask)])
         else:
             features = self.ctx_encoder(
                 input_ids=input_ids, attention_mask=attention_mask
@@ -400,6 +452,7 @@ class ModulatedPretrainedModel(nn.Module):
 
         return model_outputs
 
+    @torch.no_grad()
     def generate(
         self,
         ctx_features: Optional[Float[Tensor, "bs ctx_length feature_dim"]] = None,

hyperlora/utils.py

@@ -42,6 +42,12 @@ def get_num_layers(model):
     return len(get_layers(model))
 
 
+def get_base_model(model):
+    if hasattr(model, "model"):
+        return get_base_model(model.model)
+    return model
+
+
 def get_num_params(model):
     total_params = 0
     trainable_params = 0