Add FSDP tests

pomonam · Jul 8, 2024 · 9119228 · 9119228
1 parent 0a4cf37
commit 9119228
Showing 1 changed file with 58 additions and 67 deletions.
diff --git a/tests/gpu_tests/fsdp_test.py b/tests/gpu_tests/fsdp_test.py
@@ -9,12 +9,11 @@
 from torch.utils import data
 
 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.common.score_arguments import pytest_score_arguments
 from kronfluence.utils.constants import (
-    ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
-    LAMBDA_FACTOR_NAMES,
+    LAMBDA_FACTOR_NAMES, ALL_MODULE_NAME,
 )
 from kronfluence.utils.model import apply_fsdp
 from tests.gpu_tests.pipeline import GpuTestTask, construct_test_mlp, get_mnist_dataset
@@ -114,70 +113,62 @@ def test_lambda_matrices(self) -> None:
                     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)
-    #
-    #     if LOCAL_RANK == 0:
-    #         print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][0]}")
-    #         print(f"Previous shape: {pairwise_scores[ALL_MODULE_NAME].shape}")
-    #         print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][0]}")
-    #         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:
-    #     self_scores = self.analyzer.load_self_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_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)
-    #
-    #     if LOCAL_RANK == 0:
-    #         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 = pytest_score_arguments()
+        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)
+
+        if LOCAL_RANK == 0:
+            print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][0]}")
+            print(f"Previous shape: {pairwise_scores[ALL_MODULE_NAME].shape}")
+            print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][0]}")
+            print(f"New shape: {new_pairwise_scores[ALL_MODULE_NAME].shape}")
+            assert check_tensor_dict_equivalence(
+                pairwise_scores,
+                new_pairwise_scores,
+                atol=ATOL,
+                rtol=RTOL,
+            )
+
+    def test_self_scores(self) -> None:
+        self_scores = self.analyzer.load_self_scores(scores_name=OLD_SCORE_NAME)
+
+        score_args = pytest_score_arguments()
+        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)
+
+        if LOCAL_RANK == 0:
+            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=ATOL,
+                rtol=RTOL,
+            )
 
     @classmethod
     def tearDownClass(cls) -> None: