From a183c9c74d39419629e392ca56cda48e85523e32 Mon Sep 17 00:00:00 2001 From: 51616 Date: Sat, 21 Dec 2024 15:03:13 +0000 Subject: [PATCH] trainable hyperlora + context numbers --- hyperlora/configs.py | 9 +++- hyperlora/data_utils.py | 11 +++++ hyperlora/hooks.py | 2 + hyperlora/intx_sft.py | 94 +++++++++++++++++++++++++++++++------ hyperlora/modeling_utils.py | 15 ++++-- 5 files changed, 113 insertions(+), 18 deletions(-) diff --git a/hyperlora/configs.py b/hyperlora/configs.py index 6b16ac7..f318bb2 100644 --- a/hyperlora/configs.py +++ b/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) diff --git a/hyperlora/data_utils.py b/hyperlora/data_utils.py index 7c81b47..55ca005 100644 --- a/hyperlora/data_utils.py +++ b/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 diff --git a/hyperlora/hooks.py b/hyperlora/hooks.py index 5b75df7..eb12772 100644 --- a/hyperlora/hooks.py +++ b/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: diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index b5f379f..04ed5a3 100644 --- a/hyperlora/intx_sft.py +++ b/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() diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index 393c3d2..ef771a7 100644 --- a/hyperlora/modeling_utils.py +++ b/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,