From 67e73d0fc7ab810db2129065188f475cfb364edc Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Fri, 19 Sep 2025 06:52:18 +0000
Subject: [PATCH] allow use_bias=False

---
 src/ctx_to_lora/modeling/hypernet.py    | 37 +++++++++++++++++--------
 src/ctx_to_lora/modeling/lora_merger.py | 18 ++++++++----
 2 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py
index 6625e23..0443b0c 100644
--- a/src/ctx_to_lora/modeling/hypernet.py
+++ b/src/ctx_to_lora/modeling/hypernet.py
@@ -291,21 +291,24 @@ class HyperLoRA(nn.Module):
                         for m in self.target_modules
                     }
                 )
-            else:
-                self.bias_A = nn.ParameterDict(
+            # else:
+            #     self.bias_A = nn.ParameterDict(
+            #         {
+            #             m: nn.Parameter(
+            #                 torch.zeros((self.n_layers, self.r, self.d_in[m]))
+            #             )
+            #             for m in self.target_modules
+            #         }
+            #     )
+            if self.config.use_bias:
+                self.bias_B = nn.ParameterDict(
                     {
                         m: nn.Parameter(
-                            torch.zeros((self.n_layers, self.r, self.d_in[m]))
+                            torch.zeros((self.n_layers, self.r, self.d_out[m]))
                         )
                         for m in self.target_modules
                     }
                 )
-            self.bias_B = nn.ParameterDict(
-                {
-                    m: nn.Parameter(torch.zeros((self.n_layers, self.r, self.d_out[m])))
-                    for m in self.target_modules
-                }
-            )
 
             self.scaler_A = nn.ParameterDict(
                 {
@@ -1028,7 +1031,9 @@ class ModulatedPretrainedModel(nn.Module):
             generated_loras = combine_lora(
                 generated_loras,
                 n_ctx_chunks,
-                lora_bias=self.hypernet.get_head_bias(),
+                lora_bias=self.hypernet.get_head_bias()
+                if self.hypernet.config.use_bias
+                else None,
             )
 
             # input_ids in model_inputs_kwargs contains only
@@ -1092,6 +1097,8 @@ class ModulatedPretrainedModel(nn.Module):
         ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None,
         n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None,
         n_queries: Integer[Tensor, "n_ctx"] | None = None,
+        scalers: Float[Tensor, "n_ctx"] | None = None,
+        bias_scaler: float | None = None,
         *model_inputs_args: Any,
         **model_inputs_kwargs: dict[str, Any],
     ):
@@ -1102,7 +1109,11 @@ class ModulatedPretrainedModel(nn.Module):
         generated_loras = combine_lora(
             generated_loras,
             n_ctx_chunks,
-            lora_bias=self.hypernet.get_head_bias(),
+            lora_bias=self.hypernet.get_head_bias()
+            if self.hypernet.config.use_bias
+            else None,
+            scalers=scalers,
+            bias_scaler=bias_scaler,
         )
 
         # apply lora hook to the base model
@@ -1196,7 +1207,9 @@ class ModulatedPretrainedModel(nn.Module):
             generated_loras = self.combine_lora(
                 generated_loras,
                 n_ctx_chunks,
-                lora_bias=self.hypernet.get_head_bias(),
+                lora_bias=self.hypernet.get_head_bias()
+                if self.hypernet.config.use_bias
+                else None,
             )
 
             # # for generation the ctx_ids are batched not packed
diff --git a/src/ctx_to_lora/modeling/lora_merger.py b/src/ctx_to_lora/modeling/lora_merger.py
index c96cee5..29291d0 100644
--- a/src/ctx_to_lora/modeling/lora_merger.py
+++ b/src/ctx_to_lora/modeling/lora_merger.py
@@ -4,7 +4,7 @@ Utilities for merging / aggregating LoRA adapters coming from multiple chunks.
 
 import torch
 from einops import rearrange
-from jaxtyping import Integer
+from jaxtyping import Float, Integer
 from torch import Tensor
 
 
@@ -15,10 +15,13 @@ def compute_rank(n_lora, rank):
 def combine_lora(
     generated_loras: dict[str, dict[str, Tensor]],
     n_chunks: Integer[Tensor, "n_ctx"],
-    lora_bias: dict[str, dict[str, Tensor]],
+    lora_bias: dict[str, dict[str, Tensor]] | None = None,
+    scalers: Float[Tensor, "n_ctx"] | None = None,
+    bias_scaler: float | None = None,
 ) -> dict[str, dict[str, Tensor]]:
     total_chunks = int(n_chunks.sum())
-
+    if bias_scaler is None:
+        bias_scaler = 1
     # Assume all modules share same base rank r
     first_module = next(iter(generated_loras))
     sampled_lora = generated_loras[first_module]["A"]
@@ -33,11 +36,14 @@ def combine_lora(
     rank_dim = 2
     num_groups = len(n_chunks)
     rank_per_group = (n_chunks * base_rank).tolist()
-
+    bias_tensor = None
     for module_name, module_loras in generated_loras.items():
         for matrix_key in ("A", "B"):
-            bias_tensor = lora_bias[module_name][matrix_key]
+            if lora_bias is not None:
+                bias_tensor = lora_bias[module_name][matrix_key]
             loras = module_loras[matrix_key]
+            if (scalers is not None) and (matrix_key == "A"):
+                loras = loras * scalers[:, None, None, None]
 
             flat_loras = rearrange(
                 loras, "tot_chunks n_layers r dim -> 1 n_layers (tot_chunks r) dim"
@@ -64,7 +70,7 @@ def combine_lora(
                     #     combined_rank, combined_rank + base_rank
                     # )
                     combined[g, :, combined_rank : combined_rank + base_rank, :] = (
-                        bias_tensor
+                        bias_tensor * bias_scaler
                     )
 
             combined_loras[module_name][matrix_key] = combined