Update lora input linear adapter output dim. (#911)

* Update lora input linear adapter output dim. * Update unit tests.
apple · Jan 15, 2025 · cfb5336 · cfb5336
1 parent 3405a6e
commit cfb5336
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/axlearn/common/lora.py b/axlearn/common/lora.py
@@ -501,7 +501,7 @@ def __init__(self, cfg: Config, *, parent: Module):
             "adapter",
             cfg.adapter.set(
                 input_dim=cfg.query_dim,
-                output_dim=cfg.query_dim,
+                output_dim=cfg.num_heads * cfg.per_head_dim,
                 num_heads=cfg.num_heads,
             ),
         )

diff --git a/axlearn/common/lora_test.py b/axlearn/common/lora_test.py
@@ -137,7 +137,7 @@ def test_alpha_is_zero(self):
 
 class LoraFusedQKVLinearTest(TestCase):
     def test_forward(self):
-        model_dim = 6
+        model_dim = 16
         num_heads = 2
         per_head_dim = 3
         seq_len = 4
@@ -197,7 +197,7 @@ def test_forward(self):
         ),
     )
     def test_extend_step(self, layer):
-        model_dim = 8
+        model_dim = 16
         num_heads = 2
         per_head_dim = 4  # change this to 4 to adapt the need of RoPE.
         seq_len = 4
@@ -267,7 +267,7 @@ def test_extend_step(self, layer):
         )
 
     def test_prefill_states(self):
-        model_dim = 6
+        model_dim = 16
         num_heads = 2
         per_head_dim = 3
         seq_len = 4