From 9035b1834c7188f598230f8a195904a4188d2396 Mon Sep 17 00:00:00 2001 From: 51616 Date: Thu, 16 Jan 2025 11:46:06 +0000 Subject: [PATCH] deprecate trainable_base_modules + fix vision model load + fix metric --- hyperlora/configs.py | 4 +++ hyperlora/eval.py | 49 ++++++++++++++++++++++----- hyperlora/intx_sft.py | 48 +++++++++++++++++++-------- hyperlora/model_loading.py | 21 ++++++------ hyperlora/modeling_utils.py | 66 ++++++++++++++++++++++++++----------- hyperlora/watcher.py | 2 ++ 6 files changed, 139 insertions(+), 51 deletions(-) diff --git a/hyperlora/configs.py b/hyperlora/configs.py index 92346dc..da1dbe7 100644 --- a/hyperlora/configs.py +++ b/hyperlora/configs.py @@ -332,6 +332,10 @@ class HypernetArguments: default=128, metadata={"help": "Latent size for light-weight LoRA."}, ) + # trainable_base_modules: Optional[list[str]] = field( + # default=None, + # metadata={"help": ("Modules to train of the base model.")}, + # ) @dataclass diff --git a/hyperlora/eval.py b/hyperlora/eval.py index 09fca60..0c380e1 100644 --- a/hyperlora/eval.py +++ b/hyperlora/eval.py @@ -3,6 +3,7 @@ import sys import gc import json import os +import random from argparse import Namespace from dataclasses import fields from functools import partial @@ -18,6 +19,7 @@ from transformers import ( Seq2SeqTrainer, Seq2SeqTrainingArguments, EvalPrediction, + set_seed, ) from data_utils import get_tokenized_dataset @@ -46,8 +48,13 @@ def compute_rouge(pred_texts, label_texts): def compute_per_token_acc(shift_logits, shift_labels, valid_masks): indices = torch.where(valid_masks) - acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item() - return {"per_token_acc": acc} + # acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item() + # return {"per_token_acc": acc} + acc = (shift_logits.argmax(-1) == shift_labels)[indices].float() + return { + "per_token_accs": acc.flatten().tolist(), + "n_per_token_accs": valid_masks.sum().item(), + } def compute_prefix_matching(shift_logits, shift_labels, valid_masks): @@ -63,7 +70,11 @@ def compute_prefix_matching(shift_logits, shift_labels, valid_masks): # if all tokens are correct, set to 1 perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf) - return {"prefix_matching": perf.mean().item()} + # return {"prefix_matching": perf.mean().item()} + return { + "prefix_matchings": perf.tolist(), + "n_prefix_matchings": valid_masks.shape[0], + } @torch.no_grad() @@ -71,8 +82,13 @@ def compute_perplexity(shift_logits, shift_labels, valid_masks): loss_fct = torch.nn.CrossEntropyLoss(reduction="none") loss = loss_fct(shift_logits.transpose(1, 2), shift_labels) loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1) - perplexity = torch.exp(loss).mean().item() - return {"perplexity": perplexity} + # perplexity = torch.exp(loss).mean().item() + # return {"perplexity": perplexity} + preplexities = torch.exp(loss) + return { + "perplexities": preplexities.tolist(), + "n_perplexities": valid_masks.shape[0], + } class Evaluator: @@ -82,16 +98,22 @@ class Evaluator: def reset(self): self.accum_metrics = defaultdict(list) + self.count = defaultdict(list) def update(self, shift_logits, shift_labels, valid_masks): for metric_fn in self.metric_fns: metric = metric_fn(shift_logits, shift_labels, valid_masks) for k, v in metric.items(): - self.accum_metrics[k].append(v) + if k.startswith("n_"): + self.count[k[2:]].append(v) + else: + self.accum_metrics[k] += v def compute(self): # Get result across entire eval set - result = {k: np.mean(v) for k, v in self.accum_metrics.items()} + result = { + k: np.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items() + } # Reset batch statistics self.reset() return result @@ -304,9 +326,18 @@ def generation_collator(inp_list, tokenizer): def evaluate(checkpoint_path, args, split, generative): assert split in ["validation", "test"] + set_seed(42) + torch.use_deterministic_algorithms(True, warn_only=True) + torch.backends.cudnn.benchmark = False + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True model = ModulatedPretrainedModel.from_state_dict( - torch.load(open(checkpoint_path, "rb")), train=False + torch.load(open(checkpoint_path, "rb")), + train=False, + use_flash_attn=False, ) + # NOTE: there is still some randomness in the eval result + # despite all the deterministic settings tokenizer = get_tokenizer(args.model_name_or_path, train=False) if tokenizer.pad_token_id is None: @@ -408,7 +439,9 @@ def evaluate(checkpoint_path, args, split, generative): if __name__ == "__main__": + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" + os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1" os.environ["WANDB_MODE"] = "disabled" checkpoint_path = sys.argv[1] checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1]) diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index 5d133fc..00a8dfd 100755 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -80,8 +80,13 @@ LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) def compute_per_token_acc(shift_logits, shift_labels, valid_masks): indices = torch.where(valid_masks) - acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item() - return {"per_token_acc": acc} + # acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item() + # return {"per_token_acc": acc} + acc = (shift_logits.argmax(-1) == shift_labels)[indices].float() + return { + "per_token_accs": acc.flatten().tolist(), + "n_per_token_accs": valid_masks.sum().item(), + } def compute_prefix_matching(shift_logits, shift_labels, valid_masks): @@ -97,7 +102,11 @@ def compute_prefix_matching(shift_logits, shift_labels, valid_masks): # if all tokens are correct, set to 1 perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf) - return {"prefix_matching": perf.mean().item()} + # return {"prefix_matching": perf.mean().item()} + return { + "prefix_matchings": perf.tolist(), + "n_prefix_matchings": valid_masks.shape[0], + } @torch.no_grad() @@ -105,8 +114,13 @@ def compute_perplexity(shift_logits, shift_labels, valid_masks): loss_fct = torch.nn.CrossEntropyLoss(reduction="none") loss = loss_fct(shift_logits.transpose(1, 2), shift_labels) loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1) - perplexity = torch.exp(loss).mean().item() - return {"perplexity": perplexity} + # perplexity = torch.exp(loss).mean().item() + # return {"perplexity": perplexity} + preplexities = torch.exp(loss) + return { + "perplexities": preplexities.tolist(), + "n_perplexities": valid_masks.shape[0], + } class Evaluator: @@ -116,16 +130,22 @@ class Evaluator: def reset(self): self.accum_metrics = defaultdict(list) + self.count = defaultdict(list) def update(self, shift_logits, shift_labels, valid_masks): for metric_fn in self.metric_fns: metric = metric_fn(shift_logits, shift_labels, valid_masks) for k, v in metric.items(): - self.accum_metrics[k].append(v) + if k.startswith("n_"): + self.count[k[2:]].append(v) + else: + self.accum_metrics[k] += v def compute(self): # Get result across entire eval set - result = {k: np.mean(v) for k, v in self.accum_metrics.items()} + result = { + k: np.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items() + } # Reset batch statistics self.reset() return result @@ -279,11 +299,12 @@ def main(): requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE, peft_config=get_lora_config(model_name, **vars(lora_args)), ) - if ctx_encoder_args.ctx_encoder_model_name_or_path is not None: - ctx_encoder_model_config = AutoConfig.from_pretrained( - ctx_encoder_args.ctx_encoder_model_name_or_path - ) - ctx_tokenizer = get_tokenizer(ctx_encoder_args.ctx_encoder_model_name_or_path) + ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path + if ctx_name is not None: + ctx_encoder_model_config = AutoConfig.from_pretrained(ctx_name) + if "Llama" in ctx_name and "Vision" in ctx_name: + ctx_encoder_model_config = ctx_encoder_model_config.text_config + ctx_tokenizer = get_tokenizer(ctx_name) else: ctx_encoder_model_config = model.config ctx_tokenizer = tokenizer @@ -371,6 +392,7 @@ def main(): train_ds = interleave_datasets( list(train_ds.values()), probabilities=[l / total_len for l in train_ds_len], + seed=training_args.seed, ) # val_train_indices = np.random.permutation(len(train_ds))[:500] @@ -469,7 +491,7 @@ def main(): if isinstance(model, ModulatedPretrainedModel): logger.info("Applying liger-kernel to ModulatedPretrainedModel") _apply_liger_kernel_to_instance(model=model.base_model.base_model.model) - if ctx_encoder_args.ctx_encoder_model_name_or_path is not None: + if ctx_name is not None: logger.info("Applying liger-kernel to ctx_encoder_model") _apply_liger_kernel_to_instance(model=model.ctx_encoder) elif isinstance(model, PeftModel): diff --git a/hyperlora/model_loading.py b/hyperlora/model_loading.py index 1a305ed..159b7ab 100644 --- a/hyperlora/model_loading.py +++ b/hyperlora/model_loading.py @@ -143,16 +143,17 @@ def get_model( model = PeftModel(model, peft_config) model.train(train) for name, param in model.named_parameters(): - if "modules_to_save" not in name: - param.requires_grad = requires_grad - else: - logging.debug(f"modules_to_save: {name} to be trained") - # always train "modules_to_save" - if not "layernorm" in name: - raise NotImplementedError( - f"modules_to_save should only be layernorm, got {name}" - ) - param.requires_grad = True + # if "modules_to_save" not in name: + # param.requires_grad = requires_grad + # else: + # logging.debug(f"modules_to_save: {name} to be trained") + # # always train "modules_to_save" + # if not "layernorm" in name: + # raise NotImplementedError( + # f"modules_to_save should only be layernorm, got {name}" + # ) + # param.requires_grad = True + param.requires_grad = requires_grad return model diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index 5002270..27345c8 100644 --- a/hyperlora/modeling_utils.py +++ b/hyperlora/modeling_utils.py @@ -26,7 +26,7 @@ from peft import ( get_peft_model_state_dict, set_peft_model_state_dict, ) -from peft.utils import PeftType, TaskType +from peft.utils import PeftType, TaskType, ModulesToSaveWrapper from peft.tuners._buffer_dict import BufferDict from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists from pooling import POOL_FN, get_pooling_fn @@ -100,6 +100,7 @@ class HypernetConfig: lora_config: LoraConfig module_names: dict[str, list[str]] + trainable_base_modules: Optional[list[str]] layer_indices: Iterable[int] feature_sizes: tuple[dict[str, int], dict[str, int]] aggregator_config: AggregatorConfig @@ -574,7 +575,12 @@ class ModulatedPretrainedModel(nn.Module): # self.register_module("ctx_encoder", ctx_encoder) @classmethod - def from_state_dict(cls, state_dict: dict, train: bool = True): + def from_state_dict( + cls, + state_dict: dict, + train: bool = True, + use_flash_attn: bool = True, + ): lora_config = state_dict["hypernet_config"].lora_config model_name_or_path = lora_config.base_model_name_or_path base_model = get_model( @@ -582,6 +588,7 @@ class ModulatedPretrainedModel(nn.Module): train=train, requires_grad=False, peft_config=lora_config, + use_flash_attn=use_flash_attn, ) hypernet_config = state_dict["hypernet_config"] ctx_encoder_args = state_dict["ctx_encoder_args"] @@ -590,6 +597,17 @@ class ModulatedPretrainedModel(nn.Module): return model def _init_model(self): + # for name, module in self.base_model.named_modules(): + # if isinstance(module, ModulesToSaveWrapper): + # module.set_adapter("default") + # trainable_base_modules = getattr( + # self.hypernet_config, "trainable_base_modules", None + # ) + # if trainable_base_modules: + # for name, param in self.base_model.named_parameters(): + # if any(module_name in name for module_name in trainable_base_modules): + # param.requires_grad = True + self.hypernet = HyperLoRA(self.hypernet_config).to(self.device) # ctx_encoder_name = # if self.ctx_encoder_args.ctx_encoder_model_name_or_path is not None: @@ -675,11 +693,18 @@ class ModulatedPretrainedModel(nn.Module): state_dict = self.hypernet.state_dict(*args, **kwargs) # base_state_dict = dict() - # if self.base_model.modules_to_save: - # for k, v in get_peft_model_state_dict(self.base_model).items(): - # if any(module in k for module in self.base_model.modules_to_save): - # # save only "modules_to_save" - # base_state_dict[k] = v + # if self.hypernet_config.trainable_base_modules: + # # for k, v in get_peft_model_state_dict(self.base_model).items(): + # # if any(module in k for module in self.hypernet_config.trainable_base_modules): + # # # save only "modules_to_save" + # # base_state_dict[k] = v + # for name, param in self.base_model.named_parameters(): + # if param.requires_grad and any( + # module in name + # for module in self.hypernet_config.trainable_base_modules + # ): + # base_state_dict[name] = param.data + # state_dict["base_model"] = base_state_dict state_dict["hypernet_config"] = self.hypernet_config state_dict["ctx_encoder_args"] = self.ctx_encoder_args @@ -699,16 +724,19 @@ class ModulatedPretrainedModel(nn.Module): f"but the hypernet config is for: {self.hypernet_config.lora_config.base_model_name_or_path}" ) self._init_model() + state_dict.pop("base_model", None) # base_state_dict = state_dict.pop("base_model", None) # if base_state_dict: - # logging.info("Loading base model state dict") - # peft_load_result = set_peft_model_state_dict( - # self.base_model, - # base_state_dict, - # *args, - # **kwargs, - # ) - # logging.info(f"peft_load_result: {peft_load_result}") + # print("Loading base model state dict") + # print(base_state_dict.keys()) + # # peft_load_result = set_peft_model_state_dict( + # # self.base_model, + # # base_state_dict, + # # *args, + # # **kwargs, + # # ) + # load_result = self.base_model.load_state_dict(base_state_dict, strict=False) + # print(f"base_load_result: {load_result}") return self.hypernet.load_state_dict(state_dict, *args, **kwargs) # @torch.no_grad() @@ -910,9 +938,7 @@ if __name__ == "__main__": model_name, train=True, requires_grad=False, - peft_config=get_lora_config( - model_name, modules_to_save=["input_layernorm", "post_attention_layernorm"] - ), + peft_config=get_lora_config(model_name), ) print(base_model.modules_to_save) ctx_model_config = base_model.config @@ -966,11 +992,11 @@ if __name__ == "__main__": state_dict = torch.load( open( - "train_outputs/runs/Jan15_17-40-28_slurm0-a3nodeset-0_fd03d041/checkpoint-1000/pytorch_model.bin", + "train_outputs/runs/Jan15_19-10-21_slurm0-a3nodeset-11_d1842c41/checkpoint-55000/pytorch_model.bin", "rb", ) ) - breakpoint() + model.load_state_dict(state_dict) modelout = model(ctx_ids, ctx_attn_mask, **prompt_inputs) print(modelout) diff --git a/hyperlora/watcher.py b/hyperlora/watcher.py index a0c8a2b..0f362b5 100644 --- a/hyperlora/watcher.py +++ b/hyperlora/watcher.py @@ -62,6 +62,8 @@ class Watcher: if __name__ == "__main__": os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" + os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1" + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" os.environ["WANDB_PROJECT"] = "ctx_to_lora" os.environ["WANDB_WATCH"] = "" # "all" os.environ["WANDB_CONSOLE"] = "off"