Update test code for the GraLoRA method

Author: yeonjoon-jung01

tests/test_custom_models.py

@@ -38,6 +38,7 @@
     C3AConfig,
     DeloraConfig,
     FourierFTConfig,
+    GraloraConfig,
     HRAConfig,
     IA3Config,
     LNTuningConfig,
@@ -665,6 +666,25 @@
             "init_weights": True,
         },
     ),
+    ###########
+    # GraLoRA #
+    ###########
+    ("Vanilla MLP 1 GraLoRA", "MLP", GraloraConfig, {"target_modules": "lin0"}),
+    ("Vanilla MLP 2 GraLoRA", "MLP", GraloraConfig, {"target_modules": ["lin0"]}),
+    ("Vanilla MLP 3 GraLoRA", "MLP", GraloraConfig, {"target_modules": ["lin1"]}),
+    ("Vanilla MLP 4 GraLoRA", "MLP", GraloraConfig, {"target_modules": ["lin0", "lin1"]}),
+    (
+        "Vanilla MLP 5 GraLoRA",
+        "MLP",
+        GraloraConfig,
+        {"target_modules": ["lin0"], "modules_to_save": ["lin1"]},
+    ),
+    (
+        "Embedding + transformers Conv1D 1 GraLoRA",
+        "EmbConv1D",
+        GraloraConfig,
+        {"target_modules": ["conv1d"], "gralora_k": 1},
+    ),
     ##########
     # VBLoRA #
     ##########
@@ -973,6 +993,20 @@
         {"n_frequency": 10, "target_modules": ["lin0"]},
         {"n_frequency": 10, "target_modules": ["lin1"]},
     ),
+    (
+        "GraLoRA Same",
+        "gralora",
+        GraloraConfig,
+        {"target_modules": ["lin0"], "init_weights": False},
+        {"target_modules": ["lin0"], "init_weights": False},
+    ),
+    (
+        "GraLoRA Different",
+        "gralora",
+        GraloraConfig,
+        {"target_modules": ["lin0"], "init_weights": False},
+        {"target_modules": ["lin1"], "init_weights": False},
+    ),
     (
         "SHiRA Same",
         "shira",
@@ -1159,6 +1193,7 @@
     VeraConfig: "vera_lambda_",
     RandLoraConfig: "randlora_",
     FourierFTConfig: "fourierft_",
+    GraloraConfig: "gralora_",
     C3AConfig: "c3a_",
     HRAConfig: "hra_",
     ShiraConfig: "shira_",
@@ -3104,12 +3139,12 @@ def test_add_weighted_adapter_subtraction_with_negative_weights(self):
                 cancelled_B = module.lora_B["cancelled"].weight.data
 
                 # The weights should be approximately zero (they cancel out)
-                assert torch.allclose(cancelled_A, torch.zeros_like(cancelled_A), atol=1e-5), (
-                    f"Cancelled A should be ~0, got max abs value {cancelled_A.abs().max()}"
-                )
-                assert torch.allclose(cancelled_B, torch.zeros_like(cancelled_B), atol=1e-5), (
-                    f"Cancelled B should be ~0, got max abs value {cancelled_B.abs().max()}"
-                )
+                assert torch.allclose(
+                    cancelled_A, torch.zeros_like(cancelled_A), atol=1e-5
+                ), f"Cancelled A should be ~0, got max abs value {cancelled_A.abs().max()}"
+                assert torch.allclose(
+                    cancelled_B, torch.zeros_like(cancelled_B), atol=1e-5
+                ), f"Cancelled B should be ~0, got max abs value {cancelled_B.abs().max()}"
 
     def test_add_weighted_adapter_negative_weight_with_different_scaling(self):
         # Test negative weights with different scaling factors (lora_alpha)
@@ -3515,9 +3550,9 @@ def test_multirank_2(self):
                 if isinstance(module, BaseTunerLayer):
                     rank_expected = rank_pattern.get(key, r)
                     rank_current = module.lora_A[adapter].weight.shape[0]
-                    assert rank_current == rank_expected, (
-                        f"Rank {rank_current} is not equal to expected {rank_expected}"
-                    )
+                    assert (
+                        rank_current == rank_expected
+                    ), f"Rank {rank_current} is not equal to expected {rank_expected}"
 
 
 class TestLayerRepr:

tests/test_gralora.py

@@ -112,7 +112,7 @@ def test_gralora_parameter_shapes(self, mlp_gralora_hybrid):
                 in_features = module.in_features
                 out_features = module.out_features
                 k = 4
-                gralora_rank = 16 - 4  # r - hybrid_r
+                gralora_rank = 16
 
                 # Check GraLoRA block shapes
                 # Each block has full gralora_rank, not gralora_rank // k
@@ -203,7 +203,7 @@ def test_gralora_pure_vs_hybrid_params(self):
         mlp_hybrid = MLP()
         config_hybrid = GraloraConfig(
             target_modules=["lin1", "lin2"],
-            r=16,
+            r=12,
             gralora_k=4,
             hybrid_r=4,
         )
@@ -217,9 +217,9 @@ def count_trainable_params(model):
 
         # Pure and hybrid should have same total parameters (r is constant)
         # but distributed differently between block-diagonal and full-rank components
-        assert params_pure == params_hybrid, (
-            f"Pure ({params_pure}) and Hybrid ({params_hybrid}) should have same parameter count"
-        )
+        assert (
+            params_pure == params_hybrid
+        ), f"Pure ({params_pure}) and Hybrid ({params_hybrid}) should have same parameter count"
 
         # Check that hybrid has general components
         has_general = False
@@ -444,7 +444,7 @@ def test_gralora_rank_divisibility_check(self):
             hybrid_r=0,
         )
 
-        with pytest.raises(AssertionError, match="r should be divisible by gralora_k"):
+        with pytest.raises(ValueError, match="r should be divisible by gralora_k"):
             get_peft_model(mlp, config)
 
     def test_gralora_trainable_parameters_only(self, mlp_gralora_hybrid):
@@ -827,37 +827,6 @@ def test_gralora_unload_without_merge(self):
         # Should match base model output (no merge)
         assert torch.allclose(base_output, unloaded_output, atol=1e-5)
 
-    def test_gralora_get_peft_config_as_dict(self):
-        """Test get_peft_config_as_dict method"""
-        torch.manual_seed(0)
-        mlp = MLP()
-        config = GraloraConfig(
-            target_modules=["lin1"],
-            r=8,
-            gralora_k=2,
-            hybrid_r=4,
-            gralora_alpha=16,
-        )
-        model = get_peft_model(mlp, config)
-
-        config_dict = model.get_peft_config_as_dict(inference=False)
-
-        assert "default" in config_dict
-        assert config_dict["default"]["r"] == 8
-        assert config_dict["default"]["gralora_k"] == 2
-        assert config_dict["default"]["hybrid_r"] == 4
-
-    def test_gralora_get_peft_config_as_dict_inference_mode(self):
-        """Test get_peft_config_as_dict with inference=True"""
-        torch.manual_seed(0)
-        mlp = MLP()
-        config = GraloraConfig(target_modules=["lin1"], r=8, gralora_k=2)
-        model = get_peft_model(mlp, config)
-
-        config_dict = model.get_peft_config_as_dict(inference=True)
-
-        assert config_dict["default"]["inference_mode"] is True
-
     def test_gralora_merge_with_hybrid_component(self):
         """Test that merge works correctly with hybrid component"""
         torch.manual_seed(0)