From 59370b175f5b4af097f858d7301e9ae2c3aeb069 Mon Sep 17 00:00:00 2001 From: 51616 Date: Mon, 13 Jan 2025 14:21:54 +0000 Subject: [PATCH] per module head --- hyperlora/modeling_utils.py | 78 +++++++++++++++++++++++--------- hyperlora/training_utils.py | 89 ++++++++++++++++++------------------- 2 files changed, 100 insertions(+), 67 deletions(-) diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index e9b9f36..80a34e5 100644 --- a/hyperlora/modeling_utils.py +++ b/hyperlora/modeling_utils.py @@ -445,17 +445,37 @@ class HyperLoRA(nn.Module): else: d_lora = max(self.d_in[m] + self.d_out[m] for m in self.target_modules) - # each module processes d -> r d_out - self.head = Mix( - "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", - weight_shape="d_latent r d_lora", - # bias_shape=None, # no bias - bias_shape="r d_lora", - n_modules=len(self.target_modules), - d_latent=self.config.latent_size, - r=self.config.lora_config.r, - d_lora=d_lora, - ) + n_modules = len(self.target_modules) + # have to do this otherwise doesnt work with adamw_torch_fused + # has something to do with the bias shape (n_modules r d_lora) + # when n_modules == 1, adamw_torch_fused complains about device/layout + # but when n_modules > 1, it works fine + if n_modules == 1: + self.head = Mix( + "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", + weight_shape="d_latent r d_lora", + # bias_shape=None, # no bias + bias_shape="r d_lora", + d_latent=self.config.latent_size, + r=self.config.lora_config.r, + d_lora=d_lora, + ) + else: + # each module processes d -> r d_out independently + self.head = Mix( + "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", + weight_shape="n_modules d_latent r d_lora", + # bias_shape=None, # no bias + bias_shape="n_modules r d_lora", + n_modules=len(self.target_modules), + d_latent=self.config.latent_size, + r=self.config.lora_config.r, + d_lora=d_lora, + ) + # print(self.head) + # print(self.head.weight.shape) + # print(self.head.bias.shape) + # breakpoint() def _to_lora_dict( self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"] @@ -583,18 +603,34 @@ class ModulatedPretrainedModel(nn.Module): logger.debug(f"peft_weights: {peft_weights}") self.hypernet.head.weight.data[:] = 0 self.hypernet.head.bias.data[:] = 0 - d_in = self.hypernet.d_in - d_out = self.hypernet.d_out - m = max(self.hypernet.target_modules, key=lambda m: d_in[m] + d_out[m]) + # d_in = self.hypernet.d_in + # d_out = self.hypernet.d_out + # m = max(self.hypernet.target_modules, key=lambda m: d_in[m] + d_out[m]) - A = peft_weights[m]["lora_A"].weight.clone()[:, : d_in[m]] # [r, d_in] - B = peft_weights[m]["lora_B"].weight.clone()[: d_out[m]] # [d_out, r] - biases = [A, B.T] + # A = peft_weights[m]["lora_A"].weight.clone()[:, : d_in[m]] # [r, d_in] + # B = peft_weights[m]["lora_B"].weight.clone()[: d_out[m]] # [d_out, r] + # biases = [A, B.T] - # bias-hyperinit - # init weights to zeros and bias to the base weights - bias_cat = torch.cat(biases, dim=1) - self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat + # # bias-hyperinit + # # init weights to zeros and bias to the base weights + # bias_cat = torch.cat(biases, dim=1) + # self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat + for i, m in enumerate(self.hypernet.target_modules): + A = peft_weights[m]["lora_A"].weight.clone() # [r, in_d] + B = peft_weights[m]["lora_B"].weight.clone() # [out_d, r] + if self.hypernet.config.use_light_weight_lora: + A = A[:, : self.hypernet.config.light_weight_latent_size] + B = B[: self.hypernet.config.light_weight_latent_size] + biases = [A, B.T] + + # bias-hyperinit + # init weights to zeros and bias to the base weights + bias_cat = torch.cat(biases, dim=1) + self.hypernet.head.bias.data[..., i, :, : bias_cat.shape[1]] = bias_cat + # if len(self.hypernet.target_modules) > 1: + # self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat + # else: + # self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat # self.hypernet.head.bias.requires_grad = False def state_dict(self, *args, **kwargs): diff --git a/hyperlora/training_utils.py b/hyperlora/training_utils.py index 6d18da7..287c32d 100644 --- a/hyperlora/training_utils.py +++ b/hyperlora/training_utils.py @@ -147,60 +147,57 @@ def train_model( trainer.save_model() clear_gpu() - ############## Evaluation - # TODO: eval does not work when using with deepspeed - # make a separate eval script + # ############## Evaluation + # # TODO: eval does not work when using with deepspeed + # # make a separate eval script - # max_input_len=2**13 # for input truncation + # # max_input_len=2**13 # for input truncation - gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens) - # pad_token_id=tokenizer.pad_token_id, - # eos_token_id=? + # gen_kwargs = dict(do_sample=False, max_new_tokens=max_new_tokens) + # # pad_token_id=tokenizer.pad_token_id, + # # eos_token_id=? - eval_trainer_args = {} + # eval_trainer_args = {} - # Copy only necessary attributes from training_args to eval_trainer_args - seq2seq_training_args_fields = {f.name for f in fields(Seq2SeqTrainingArguments)} - for attr, value in training_args.to_dict().items(): - if attr in seq2seq_training_args_fields: - eval_trainer_args[attr] = value + # # Copy only necessary attributes from training_args to eval_trainer_args + # seq2seq_training_args_fields = {f.name for f in fields(Seq2SeqTrainingArguments)} + # for attr, value in training_args.to_dict().items(): + # if attr in seq2seq_training_args_fields: + # eval_trainer_args[attr] = value - eval_trainer_args["eval_strategy"] = "no" - eval_trainer_args["save_strategy"] = "no" - eval_trainer_args["overwrite_output_dir"] = True - eval_trainer_args["per_device_eval_batch_size"] = gen_per_device_eval_batch_size + # eval_trainer_args["eval_strategy"] = "no" + # eval_trainer_args["save_strategy"] = "no" + # eval_trainer_args["overwrite_output_dir"] = True + # eval_trainer_args["per_device_eval_batch_size"] = gen_per_device_eval_batch_size - # NOTE: could also set kv_cache implementation here - eval_trainer_args = Seq2SeqTrainingArguments( - **eval_trainer_args, - predict_with_generate=True, - generation_config=GenerationConfig(**gen_kwargs), - ) + # # NOTE: could also set kv_cache implementation here + # eval_trainer_args = Seq2SeqTrainingArguments( + # **eval_trainer_args, + # predict_with_generate=True, + # generation_config=GenerationConfig(**gen_kwargs), + # ) - # Seq2SeqTrainer is actually just the same as Trainer - # (although it uses a different data collator, i.e., explicit prompt/answer separation) - # it just allows `predict_with_generate` - # allowing us to compute metrics on the generated outputs - # no clue why they call this seq2seq... + # # Seq2SeqTrainer is actually just the same as Trainer + # # (although it uses a different data collator, i.e., explicit prompt/answer separation) + # # it just allows `predict_with_generate` + # # allowing us to compute metrics on the generated outputs + # # no clue why they call this seq2seq... - logger.info("=" * 80 + "\n" + "Evaluating model..." + "\n" + "=" * 80) + # logger.info("=" * 80 + "\n" + "Evaluating model..." + "\n" + "=" * 80) - model.eval() + # model.eval() - eval_trainer = Seq2SeqTrainer( - model=model, - args=eval_trainer_args, - # TODO: use a different collator for test, e.g., more max_len truncation - # w/ left padding? - # removing label part from input_ids - data_collator=generation_collator, - ) + # eval_trainer = Seq2SeqTrainer( + # model=model, + # args=eval_trainer_args, + # # TODO: use a different collator for test, e.g., more max_len truncation + # # w/ left padding? + # # removing label part from input_ids + # data_collator=generation_collator, + # ) - for split, ds in zip(["eval", "test"], [val_dataset, test_dataset]): - if ds is None: - continue - eval_generation(eval_trainer, tokenizer, ds, split, gen_kwargs) - clear_gpu() - - # if test_dataset is not None: - # eval_generation(eval_trainer, tokenizer, test_dataset, "test", gen_kwargs) + # for split, ds in zip(["eval", "test"], [val_dataset, test_dataset]): + # if ds is None: + # continue + # eval_generation(eval_trainer, tokenizer, ds, split, gen_kwargs) + # clear_gpu()