From fce138ff04a2817e10c3c727e1e2892d979e0fb9 Mon Sep 17 00:00:00 2001 From: 51616 Date: Wed, 25 Jun 2025 23:18:22 +0900 Subject: [PATCH] generate with new lora forward --- src/ctx_to_lora/modeling/hypernet.py | 121 ++++++++++++++------------- 1 file changed, 65 insertions(+), 56 deletions(-) diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index 4ffe384..ce06953 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -611,6 +611,7 @@ class ModulatedPretrainedModel(nn.Module): cls, state_dict: dict, train: bool = True, + base_model_kwargs: dict = None, use_flash_attn: bool = True, ): lora_config = state_dict["hypernet_config"].lora_config @@ -621,6 +622,7 @@ class ModulatedPretrainedModel(nn.Module): train=train, requires_grad=False, peft_config=lora_config, + model_kwargs=base_model_kwargs, use_flash_attn=use_flash_attn, ) hypernet_config = state_dict["hypernet_config"] @@ -680,6 +682,10 @@ class ModulatedPretrainedModel(nn.Module): def config(self): return self.base_model.config + @property + def generation_config(self): + return self.base_model.generation_config + def get_input_embeddings(self): return self.base_model.get_input_embeddings() @@ -898,45 +904,44 @@ class ModulatedPretrainedModel(nn.Module): else: return model_outputs - # @contextmanager - @torch.inference_mode() - def generate_with_multi_loras( - self, - ctx_ids: Integer[Tensor, "bs ctx_length"], - ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None, - ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None, - *model_inputs_args: Any, - **model_inputs_kwargs: dict[str, Any], - ): - # TODO: make this persistent - generated_loras, _ = self.generate_weights( - ctx_ids, ctx_attn_mask, ctx_position_ids - ) - # sum loras - for lora_at_module in generated_loras.values(): - if len(lora_at_module["A"]) > 1: - print("Multiple LoRAs are generated, summing them up") - lora_at_module["A"] = torch.mean(lora_at_module["A"], dim=0, keepdim=True) - lora_at_module["B"] = torch.mean(lora_at_module["B"], dim=0, keepdim=True) + # @torch.inference_mode() + # def generate_with_multi_loras( + # self, + # ctx_ids: Integer[Tensor, "bs ctx_length"], + # ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None, + # ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None, + # *model_inputs_args: Any, + # **model_inputs_kwargs: dict[str, Any], + # ): + # # TODO: make this persistent + # generated_loras, _ = self.generate_weights( + # ctx_ids, ctx_attn_mask, ctx_position_ids + # ) + # # sum loras + # for lora_at_module in generated_loras.values(): + # if len(lora_at_module["A"]) > 1: + # print("Multiple LoRAs are generated, summing them up") + # lora_at_module["A"] = torch.mean(lora_at_module["A"], dim=0, keepdim=True) + # lora_at_module["B"] = torch.mean(lora_at_module["B"], dim=0, keepdim=True) - position_ids = ( - model_inputs_kwargs["position_ids"] - if "position_ids" in model_inputs_kwargs - else None - ) + # position_ids = ( + # model_inputs_kwargs["position_ids"] + # if "position_ids" in model_inputs_kwargs + # else None + # ) - # apply loras - with apply_generated_loras( - self.base_model, - generated_loras, - self.hypernet.layer_indices, - position_ids, - self.training, - ): - model_outputs = self.base_model.generate( - *model_inputs_args, **model_inputs_kwargs - ) - return model_outputs + # # apply loras + # with apply_generated_loras( + # self.base_model, + # generated_loras, + # self.hypernet.layer_indices, + # position_ids, + # self.training, + # ): + # model_outputs = self.base_model.generate( + # *model_inputs_args, **model_inputs_kwargs + # ) + # return model_outputs @torch.inference_mode() def generate( @@ -945,6 +950,7 @@ class ModulatedPretrainedModel(nn.Module): ctx_ids: Integer[Tensor, "bs ctx_length"] | None = None, ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None, ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None, + n_queries: Integer[Tensor, "n_ctx"] | None = None, *model_inputs_args: Any, **model_inputs_kwargs: dict[str, Any], ): @@ -997,25 +1003,28 @@ class ModulatedPretrainedModel(nn.Module): if "position_ids" in model_inputs_kwargs else None ) - with ( - apply_generated_loras( - self.base_model, - generated_loras, - self.hypernet.layer_indices, - position_ids, - self.training, - ), - apply_generated_layernorm( - self.base_model, - generated_layernorms, - self.hypernet.layer_indices, - position_ids, - self.training, - ), - ): - model_outputs = self.base_model.generate( - *model_inputs_args, **model_inputs_kwargs - ) + if n_queries is None: + if ctx_position_ids is None: + n_queries = torch.ones( + ctx_ids.shape[0], dtype=torch.int32, device=self.device + ) + else: + # quite redundant (we do cu_seqlens many places) + # TODO: compute cu_seqlens here and propagate that + n_queries = torch.ones( + (ctx_position_ids == 0).sum(), dtype=torch.int32, device=self.device + ) + + apply_lora_to_layers( + self.base_model, + self.hypernet.layer_indices, + generated_loras, + n_queries, + position_ids, + ) + model_outputs = self.base_model.generate( + *model_inputs_args, **model_inputs_kwargs + ) return model_outputs