Modify all other tests

pomonam · Jul 7, 2024 · 398503c · 398503c
1 parent c918176
commit 398503c
Show file tree

Hide file tree

Showing 4 changed files with 172 additions and 162 deletions.
diff --git a/tests/gpu_tests/amp_test.py b/tests/gpu_tests/amp_test.py
@@ -8,6 +8,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
+from kronfluence.utils.common.factor_arguments import pytest_factor_arguments
 from kronfluence.utils.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
@@ -45,13 +46,8 @@ def setUpClass(cls) -> None:
 
     def test_covariance_matrices(self) -> None:
         covariance_factors = self.analyzer.load_covariance_matrices(factors_name=OLD_FACTOR_NAME)
-        factor_args = FactorArguments(
-            use_empirical_fisher=True,
-            activation_covariance_dtype=torch.float64,
-            gradient_covariance_dtype=torch.float64,
-            lambda_dtype=torch.float64,
-            amp_dtype=torch.float16,
-        )
+        factor_args = pytest_factor_arguments()
+        factor_args.amp_dtype = torch.bfloat16
         self.analyzer.fit_covariance_matrices(
             factors_name=NEW_FACTOR_NAME,
             dataset=self.train_dataset,
@@ -73,15 +69,10 @@ def test_covariance_matrices(self) -> None:
                 rtol=1e-1,
             )
 
-    def test_lambda_matrices(self):
+    def test_lambda_matrices(self) -> None:
         lambda_factors = self.analyzer.load_lambda_matrices(factors_name=OLD_FACTOR_NAME)
-        factor_args = FactorArguments(
-            use_empirical_fisher=True,
-            activation_covariance_dtype=torch.float64,
-            gradient_covariance_dtype=torch.float64,
-            lambda_dtype=torch.float64,
-            amp_dtype=torch.float16,
-        )
+        factor_args = pytest_factor_arguments()
+        factor_args.amp_dtype = torch.bfloat16
         self.analyzer.fit_lambda_matrices(
             factors_name=NEW_FACTOR_NAME,
             dataset=self.train_dataset,
@@ -104,69 +95,69 @@ def test_lambda_matrices(self):
                 rtol=1e-1,
             )
 
-    def test_pairwise_scores(self) -> None:
-        pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=OLD_SCORE_NAME)
-
-        score_args = ScoreArguments(
-            score_dtype=torch.float64,
-            per_sample_gradient_dtype=torch.float64,
-            precondition_dtype=torch.float64,
-            amp_dtype=torch.float16,
-        )
-        self.analyzer.compute_pairwise_scores(
-            scores_name=NEW_SCORE_NAME,
-            factors_name=OLD_FACTOR_NAME,
-            query_dataset=self.eval_dataset,
-            train_dataset=self.train_dataset,
-            train_indices=list(range(TRAIN_INDICES)),
-            query_indices=list(range(QUERY_INDICES)),
-            per_device_query_batch_size=12,
-            per_device_train_batch_size=512,
-            score_args=score_args,
-            overwrite_output_dir=True,
-        )
-        new_pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=NEW_SCORE_NAME)
-
-        print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][10]}")
-        print(f"Previous shape: {pairwise_scores[ALL_MODULE_NAME].shape}")
-        print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][10]}")
-        print(f"New shape: {new_pairwise_scores[ALL_MODULE_NAME].shape}")
-        assert check_tensor_dict_equivalence(
-            pairwise_scores,
-            new_pairwise_scores,
-            atol=1e-5,
-            rtol=1e-3,
-        )
-
-    def test_self_scores(self) -> None:
-        score_args = ScoreArguments(
-            score_dtype=torch.float64,
-            per_sample_gradient_dtype=torch.float64,
-            precondition_dtype=torch.float64,
-            amp_dtype=torch.float16,
-        )
-        self.analyzer.compute_self_scores(
-            scores_name=NEW_SCORE_NAME,
-            factors_name=OLD_FACTOR_NAME,
-            train_dataset=self.train_dataset,
-            train_indices=list(range(TRAIN_INDICES)),
-            per_device_train_batch_size=512,
-            score_args=score_args,
-            overwrite_output_dir=True,
-        )
-        new_self_scores = self.analyzer.load_self_scores(scores_name=NEW_SCORE_NAME)
-
-        self_scores = self.analyzer.load_self_scores(scores_name=OLD_SCORE_NAME)
-        print(f"Previous score: {self_scores[ALL_MODULE_NAME]}")
-        print(f"Previous shape: {self_scores[ALL_MODULE_NAME].shape}")
-        print(f"New score: {new_self_scores[ALL_MODULE_NAME]}")
-        print(f"New shape: {new_self_scores[ALL_MODULE_NAME].shape}")
-        assert check_tensor_dict_equivalence(
-            self_scores,
-            new_self_scores,
-            atol=1e-5,
-            rtol=1e-3,
-        )
+    # def test_pairwise_scores(self) -> None:
+    #     pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=OLD_SCORE_NAME)
+    #
+    #     score_args = ScoreArguments(
+    #         score_dtype=torch.float64,
+    #         per_sample_gradient_dtype=torch.float64,
+    #         precondition_dtype=torch.float64,
+    #         amp_dtype=torch.float16,
+    #     )
+    #     self.analyzer.compute_pairwise_scores(
+    #         scores_name=NEW_SCORE_NAME,
+    #         factors_name=OLD_FACTOR_NAME,
+    #         query_dataset=self.eval_dataset,
+    #         train_dataset=self.train_dataset,
+    #         train_indices=list(range(TRAIN_INDICES)),
+    #         query_indices=list(range(QUERY_INDICES)),
+    #         per_device_query_batch_size=12,
+    #         per_device_train_batch_size=512,
+    #         score_args=score_args,
+    #         overwrite_output_dir=True,
+    #     )
+    #     new_pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=NEW_SCORE_NAME)
+    #
+    #     print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][10]}")
+    #     print(f"Previous shape: {pairwise_scores[ALL_MODULE_NAME].shape}")
+    #     print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][10]}")
+    #     print(f"New shape: {new_pairwise_scores[ALL_MODULE_NAME].shape}")
+    #     assert check_tensor_dict_equivalence(
+    #         pairwise_scores,
+    #         new_pairwise_scores,
+    #         atol=1e-5,
+    #         rtol=1e-3,
+    #     )
+    #
+    # def test_self_scores(self) -> None:
+    #     score_args = ScoreArguments(
+    #         score_dtype=torch.float64,
+    #         per_sample_gradient_dtype=torch.float64,
+    #         precondition_dtype=torch.float64,
+    #         amp_dtype=torch.float16,
+    #     )
+    #     self.analyzer.compute_self_scores(
+    #         scores_name=NEW_SCORE_NAME,
+    #         factors_name=OLD_FACTOR_NAME,
+    #         train_dataset=self.train_dataset,
+    #         train_indices=list(range(TRAIN_INDICES)),
+    #         per_device_train_batch_size=512,
+    #         score_args=score_args,
+    #         overwrite_output_dir=True,
+    #     )
+    #     new_self_scores = self.analyzer.load_self_scores(scores_name=NEW_SCORE_NAME)
+    #
+    #     self_scores = self.analyzer.load_self_scores(scores_name=OLD_SCORE_NAME)
+    #     print(f"Previous score: {self_scores[ALL_MODULE_NAME]}")
+    #     print(f"Previous shape: {self_scores[ALL_MODULE_NAME].shape}")
+    #     print(f"New score: {new_self_scores[ALL_MODULE_NAME]}")
+    #     print(f"New shape: {new_self_scores[ALL_MODULE_NAME].shape}")
+    #     assert check_tensor_dict_equivalence(
+    #         self_scores,
+    #         new_self_scores,
+    #         atol=1e-5,
+    #         rtol=1e-3,
+    #     )
 
 
 if __name__ == "__main__":

diff --git a/tests/gpu_tests/compile_test.py b/tests/gpu_tests/compile_test.py
@@ -8,14 +8,15 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
+from kronfluence.utils.common.factor_arguments import pytest_factor_arguments
 from kronfluence.utils.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
     LAMBDA_FACTOR_NAMES,
 )
 from tests.gpu_tests.pipeline import GpuTestTask, construct_test_mlp, get_mnist_dataset
 from tests.gpu_tests.prepare_tests import QUERY_INDICES, TRAIN_INDICES
-from tests.utils import check_tensor_dict_equivalence
+from tests.utils import check_tensor_dict_equivalence, ATOL, RTOL
 
 logging.basicConfig(level=logging.DEBUG)
 OLD_FACTOR_NAME = "single_gpu"
@@ -49,12 +50,7 @@ def setUpClass(cls) -> None:
 
     def test_covariance_matrices(self) -> None:
         covariance_factors = self.analyzer.load_covariance_matrices(factors_name=OLD_FACTOR_NAME)
-        factor_args = FactorArguments(
-            use_empirical_fisher=True,
-            activation_covariance_dtype=torch.float64,
-            gradient_covariance_dtype=torch.float64,
-            lambda_dtype=torch.float64,
-        )
+        factor_args = pytest_factor_arguments()
         self.analyzer.fit_covariance_matrices(
             factors_name=NEW_FACTOR_NAME,
             dataset=self.train_dataset,
@@ -69,15 +65,16 @@ def test_covariance_matrices(self) -> None:
                 print(f"Name: {name, module_name}")
                 print(f"Previous factor: {covariance_factors[name][module_name]}")
                 print(f"New factor: {new_covariance_factors[name][module_name]}")
+            assert check_tensor_dict_equivalence(
+                covariance_factors[name],
+                new_covariance_factors[name],
+                atol=ATOL,
+                rtol=RTOL,
+            )
 
-    def test_lambda_matrices(self):
+    def test_lambda_matrices(self) -> None:
         lambda_factors = self.analyzer.load_lambda_matrices(factors_name=OLD_FACTOR_NAME)
-        factor_args = FactorArguments(
-            use_empirical_fisher=True,
-            activation_covariance_dtype=torch.float64,
-            gradient_covariance_dtype=torch.float64,
-            lambda_dtype=torch.float64,
-        )
+        factor_args = pytest_factor_arguments()
         self.analyzer.fit_lambda_matrices(
             factors_name=NEW_FACTOR_NAME,
             dataset=self.train_dataset,
@@ -96,71 +93,97 @@ def test_lambda_matrices(self):
             assert check_tensor_dict_equivalence(
                 lambda_factors[name],
                 new_lambda_factors[name],
-                atol=1e-3,
-                rtol=1e-1,
+                atol=ATOL,
+                rtol=RTOL,
             )
 
-    def test_pairwise_scores(self) -> None:
-        pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=OLD_SCORE_NAME)
-
-        score_args = ScoreArguments(
-            score_dtype=torch.float64,
-            per_sample_gradient_dtype=torch.float64,
-            precondition_dtype=torch.float64,
-        )
-        self.analyzer.compute_pairwise_scores(
-            scores_name=NEW_SCORE_NAME,
-            factors_name=OLD_FACTOR_NAME,
-            query_dataset=self.eval_dataset,
-            train_dataset=self.train_dataset,
-            train_indices=list(range(TRAIN_INDICES)),
-            query_indices=list(range(QUERY_INDICES)),
-            per_device_query_batch_size=12,
-            per_device_train_batch_size=512,
-            score_args=score_args,
+    def test_lambda_shared_matrices(self) -> None:
+        lambda_factors = self.analyzer.load_lambda_matrices(factors_name=OLD_FACTOR_NAME)
+        factor_args = pytest_factor_arguments()
+        factor_args.has_shared_parameters = True
+        self.analyzer.fit_lambda_matrices(
+            factors_name=NEW_FACTOR_NAME,
+            dataset=self.train_dataset,
+            factor_args=factor_args,
+            per_device_batch_size=512,
             overwrite_output_dir=True,
+            load_from_factors_name=OLD_FACTOR_NAME,
         )
-        new_pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=NEW_SCORE_NAME)
-
-        print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][10]}")
-        print(f"Previous shape: {pairwise_scores[ALL_MODULE_NAME].shape}")
-        print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][10]}")
-        print(f"New shape: {new_pairwise_scores[ALL_MODULE_NAME].shape}")
-        assert check_tensor_dict_equivalence(
-            pairwise_scores,
-            new_pairwise_scores,
-            atol=1e-5,
-            rtol=1e-3,
-        )
+        new_lambda_factors = self.analyzer.load_lambda_matrices(factors_name=NEW_FACTOR_NAME)
 
-    def test_self_scores(self) -> None:
-        score_args = ScoreArguments(
-            score_dtype=torch.float64,
-            per_sample_gradient_dtype=torch.float64,
-            precondition_dtype=torch.float64,
-        )
-        self.analyzer.compute_self_scores(
-            scores_name=NEW_SCORE_NAME,
-            factors_name=OLD_FACTOR_NAME,
-            train_dataset=self.train_dataset,
-            train_indices=list(range(TRAIN_INDICES)),
-            per_device_train_batch_size=512,
-            score_args=score_args,
-            overwrite_output_dir=True,
-        )
-        new_self_scores = self.analyzer.load_self_scores(scores_name=NEW_SCORE_NAME)
-
-        self_scores = self.analyzer.load_self_scores(scores_name=OLD_SCORE_NAME)
-        print(f"Previous score: {self_scores[ALL_MODULE_NAME]}")
-        print(f"Previous shape: {self_scores[ALL_MODULE_NAME].shape}")
-        print(f"New score: {new_self_scores[ALL_MODULE_NAME]}")
-        print(f"New shape: {new_self_scores[ALL_MODULE_NAME].shape}")
-        assert check_tensor_dict_equivalence(
-            self_scores,
-            new_self_scores,
-            atol=1e-5,
-            rtol=1e-3,
-        )
+        for name in LAMBDA_FACTOR_NAMES:
+            for module_name in lambda_factors[name]:
+                print(f"Name: {name, module_name}")
+                print(f"Previous factor: {lambda_factors[name][module_name]}")
+                print(f"New factor: {new_lambda_factors[name][module_name]}")
+            assert check_tensor_dict_equivalence(
+                lambda_factors[name],
+                new_lambda_factors[name],
+                atol=ATOL,
+                rtol=RTOL,
+            )
+
+    # def test_pairwise_scores(self) -> None:
+    #     pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=OLD_SCORE_NAME)
+    #
+    #     score_args = ScoreArguments(
+    #         score_dtype=torch.float64,
+    #         per_sample_gradient_dtype=torch.float64,
+    #         precondition_dtype=torch.float64,
+    #     )
+    #     self.analyzer.compute_pairwise_scores(
+    #         scores_name=NEW_SCORE_NAME,
+    #         factors_name=OLD_FACTOR_NAME,
+    #         query_dataset=self.eval_dataset,
+    #         train_dataset=self.train_dataset,
+    #         train_indices=list(range(TRAIN_INDICES)),
+    #         query_indices=list(range(QUERY_INDICES)),
+    #         per_device_query_batch_size=12,
+    #         per_device_train_batch_size=512,
+    #         score_args=score_args,
+    #         overwrite_output_dir=True,
+    #     )
+    #     new_pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=NEW_SCORE_NAME)
+    #
+    #     print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][10]}")
+    #     print(f"Previous shape: {pairwise_scores[ALL_MODULE_NAME].shape}")
+    #     print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][10]}")
+    #     print(f"New shape: {new_pairwise_scores[ALL_MODULE_NAME].shape}")
+    #     assert check_tensor_dict_equivalence(
+    #         pairwise_scores,
+    #         new_pairwise_scores,
+    #         atol=1e-5,
+    #         rtol=1e-3,
+    #     )
+    #
+    # def test_self_scores(self) -> None:
+    #     score_args = ScoreArguments(
+    #         score_dtype=torch.float64,
+    #         per_sample_gradient_dtype=torch.float64,
+    #         precondition_dtype=torch.float64,
+    #     )
+    #     self.analyzer.compute_self_scores(
+    #         scores_name=NEW_SCORE_NAME,
+    #         factors_name=OLD_FACTOR_NAME,
+    #         train_dataset=self.train_dataset,
+    #         train_indices=list(range(TRAIN_INDICES)),
+    #         per_device_train_batch_size=512,
+    #         score_args=score_args,
+    #         overwrite_output_dir=True,
+    #     )
+    #     new_self_scores = self.analyzer.load_self_scores(scores_name=NEW_SCORE_NAME)
+    #
+    #     self_scores = self.analyzer.load_self_scores(scores_name=OLD_SCORE_NAME)
+    #     print(f"Previous score: {self_scores[ALL_MODULE_NAME]}")
+    #     print(f"Previous shape: {self_scores[ALL_MODULE_NAME].shape}")
+    #     print(f"New score: {new_self_scores[ALL_MODULE_NAME]}")
+    #     print(f"New shape: {new_self_scores[ALL_MODULE_NAME].shape}")
+    #     assert check_tensor_dict_equivalence(
+    #         self_scores,
+    #         new_self_scores,
+    #         atol=1e-5,
+    #         rtol=1e-3,
+    #     )
 
 
 if __name__ == "__main__":