Add more examples

pomonam · Mar 19, 2024 · f0b77f0 · f0b77f0
1 parent 3e9d45b
commit f0b77f0
Showing 1 changed file with 66 additions and 66 deletions.
diff --git a/tests/gpu_tests/ddp_test.py b/tests/gpu_tests/ddp_test.py
@@ -123,72 +123,72 @@ def test_lambda_matrices(self):
                     atol=1e-3,
                     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,
-    #     )
-    #     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 = 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_lr_pairwise_scores(self) -> None:
     #     pairwise_scores = self.analyzer.load_pairwise_scores(scores_name="single_gpu_qb")
     #