posegnn LoRA added to all layers

Author: luispintoc

models/pos_egnn/posegnn/adapter/README.md

@@ -2,34 +2,6 @@
 
 This adapter injects LoRA into mergeable linear layers of **PosEGNN** and exports merged weights that load into a plain `PosEGNN` with `strict=True`.
 
-## Skipped layers
-
-These layers have a built-in activation inside their Dense block, which makes algebraic merging incorrect. They are always skipped so that merged exports match adapter-enabled outputs exactly.
-
-- `encoder.neighbor_embedding.combine.dense_layers.0`
-- `encoder.edge_embedding.edge_up.dense_layers.0`
-- `encoder.gata.0.gamma_s.0`
-- `encoder.gata.0.gamma_v.0`
-- `encoder.gata.0.phik_w_ra`
-- `encoder.gata.0.edge_attr_up.dense_layers.0`
-- `encoder.gata.1.gamma_s.0`
-- `encoder.gata.1.gamma_v.0`
-- `encoder.gata.1.phik_w_ra`
-- `encoder.gata.1.edge_attr_up.dense_layers.0`
-- `encoder.gata.2.gamma_s.0`
-- `encoder.gata.2.gamma_v.0`
-- `encoder.gata.2.phik_w_ra`
-- `encoder.gata.2.edge_attr_up.dense_layers.0`
-- `encoder.gata.3.gamma_s.0`
-- `encoder.gata.3.gamma_v.0`
-- `encoder.gata.3.phik_w_ra`
-- `encoder.eqff.0.gamma_m.0`
-- `encoder.eqff.1.gamma_m.0`
-- `encoder.eqff.2.gamma_m.0`
-- `encoder.eqff.3.gamma_m.0`
-
-Skipping only affects where LoRA is attached. The base model behavior is unchanged.
-
 ## Usage
 
 ```python

models/pos_egnn/posegnn/adapter/inject.py

@@ -1,4 +1,3 @@
-# inject.py
 import re
 import torch
 import torch.nn as nn
@@ -8,7 +7,7 @@
 def apply_lora(model: nn.Module, cfg: LoRAConfig) -> tuple[int, int]:
     """
     Replace leaf linear-like layers under include patterns with LoRA.
-    Skips any module that has a non-identity .activation to guarantee mergeability.
+    Safely wraps linears with internal norm/activation since LoRA is pre-activation.
     Returns (num_scalar_wrapped, 0).
     """
     include_patterns = list(cfg.include_names or [])
@@ -32,25 +31,18 @@ def is_linear_like(m: nn.Module) -> bool:
             return False
         return isinstance(w, torch.Tensor) and w.ndim == 2
 
-    def has_post_act(m: nn.Module) -> bool:
-        act = getattr(m, "activation", None)
-        return (act is not None) and (not isinstance(act, nn.Identity))
-
     n_scalar = 0
-    skipped = []  # <— track skipped post-activation linears
 
     for full_name, module in list(model.named_modules()):
         if not is_linear_like(module):
             continue
         if not wants(full_name):
             continue
-        if has_post_act(module):
-            skipped.append(full_name)  # <— record and skip
-            continue
 
         parent_name, _, child = full_name.rpartition(".")
         parent = model.get_submodule(parent_name) if parent_name else model
 
+        # already wrapped guard
         if hasattr(module, "base") and hasattr(module, "lora_A") and hasattr(module, "lora_B"):
             continue
 
@@ -60,9 +52,4 @@ def has_post_act(m: nn.Module) -> bool:
         setattr(parent, child, wrapped)
         n_scalar += 1
 
-    if getattr(cfg, "log_skipped", False) and skipped:
-        print("[lora] skipped post-activation linears:")
-        for n in skipped:
-            print("  -", n)
-
     return n_scalar, 0

models/pos_egnn/posegnn/adapter/layers.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from typing import Optional
 
 def _init_lora(linear: nn.Linear, freeze_base: bool):
@@ -10,8 +11,8 @@ def _init_lora(linear: nn.Linear, freeze_base: bool):
 
 class LoRALinear(nn.Module):
     """
-    LoRA for linear layers:
-      y = base(x) + scaling * B(A(dropout(x)))
+    LoRA for linear layers applied pre-activation:
+      y = act( norm( (W x + b) + scaling * B(A(dropout(x))) ) )
     """
     def __init__(self, base_linear: nn.Linear, rank: int, alpha: Optional[float],
                  dropout: float, merge_on_save: bool, freeze_base: bool):
@@ -28,13 +29,21 @@ def __init__(self, base_linear: nn.Linear, rank: int, alpha: Optional[float],
         self.enable_lora = True
         self.merged = False
 
+        # Optional submodules carried by custom Dense
+        self._norm = getattr(base_linear, "norm", None)
+        if not isinstance(self._norm, nn.Module):
+            self._norm = None
+
         self._post_act = getattr(base_linear, "activation", None)
         self._has_post_act = self._post_act is not None and not isinstance(self._post_act, nn.Identity)
-        self.merge_on_save = bool(merge_on_save and not self._has_post_act)
 
+        # Always allow merge on save now that we inject pre-activation
+        self.merge_on_save = bool(merge_on_save)
+
+        # LoRA adapters
         self.lora_dropout = nn.Dropout(dropout) if dropout and dropout > 0 else nn.Identity()
-        self.lora_A = nn.Linear(self.in_features, self.r, bias=False)
-        self.lora_B = nn.Linear(self.r, self.out_features, bias=False)
+        self.lora_A = nn.Linear(self.in_features, self.r, bias=False)  # down
+        self.lora_B = nn.Linear(self.r, self.out_features, bias=False) # up
 
         nn.init.kaiming_uniform_(self.lora_A.weight, a=5**0.5)
         nn.init.zeros_(self.lora_B.weight)
@@ -43,21 +52,37 @@ def __init__(self, base_linear: nn.Linear, rank: int, alpha: Optional[float],
             self._register_state_dict_hook(self._merge_on_state_dict)
         self._register_load_state_dict_pre_hook(self._strict_fill_on_load, with_module=True)
 
+    def _apply_activation(self, y):
+        if not self._has_post_act:
+            return y
+        act = self._post_act
+        # support nn.Module or callable (e.g. torch.nn.functional.silu)
+        if isinstance(act, nn.Module):
+            return act(y)
+        if callable(act):
+            return act(y)
+        return y
+
     def forward(self, x):
-        y = self.base(x)
-        if self._has_post_act:
-            y = self._post_act(y)
+        # linear pre-activation
+        y = F.linear(x, self.base.weight, self.base.bias)
+
+        # add LoRA delta pre-activation
         if self.enable_lora and self.r > 0:
             z = self.lora_dropout(x)
             z = self.lora_A(z)
             z = self.lora_B(z)
             y = y + self.scaling * z
+
+        # optional norm then activation
+        if self._norm is not None:
+            y = self._norm(y)
+        y = self._apply_activation(y)
         return y
 
     @torch.no_grad()
     def merged_weight(self):
-        if self._has_post_act:
-            return self.base.weight
+        # Always valid since injected pre-activation
         return self.base.weight + self.scaling * (self.lora_B.weight @ self.lora_A.weight)
 
     def _merge_on_state_dict(self, module, state_dict, prefix, local_metadata):