diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index a91e751..5c80542 100644 --- a/hyperlora/intx_sft.py +++ b/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) diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index cff0384..85e40bc 100644 --- a/hyperlora/modeling_utils.py +++ b/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, diff --git a/hyperlora/utils.py b/hyperlora/utils.py index 8c98da3..8631f4f 100644 --- a/hyperlora/utils.py +++ b/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