From 1dc7d85cf3c5df338b1303412d7140c77cd14c15 Mon Sep 17 00:00:00 2001 From: Juhan Bae Date: Mon, 8 Jul 2024 16:51:01 -0400 Subject: [PATCH] Remove debug lines --- kronfluence/module/tracker/factor.py | 2 - tests/gpu_tests/compile_test.py | 114 +++++++++++++-------------- 2 files changed, 53 insertions(+), 63 deletions(-) diff --git a/kronfluence/module/tracker/factor.py b/kronfluence/module/tracker/factor.py index 895e1a3..c649f97 100644 --- a/kronfluence/module/tracker/factor.py +++ b/kronfluence/module/tracker/factor.py @@ -255,7 +255,6 @@ def forward_hook(module: nn.Module, inputs: Tuple[torch.Tensor], outputs: torch. @torch.no_grad() def backward_hook(output_gradient: torch.Tensor) -> None: - print("Normal") if self.cached_activations is None: self._raise_cache_not_found_exception() handle = self.cached_hooks.pop() @@ -274,7 +273,6 @@ def backward_hook(output_gradient: torch.Tensor) -> None: @torch.no_grad() def shared_backward_hook(output_gradient: torch.Tensor) -> None: - print("Shared processed") handle = self.cached_hooks.pop() handle.remove() output_gradient = self._preprocess_gradient( diff --git a/tests/gpu_tests/compile_test.py b/tests/gpu_tests/compile_test.py index ac9d23f..629f700 100644 --- a/tests/gpu_tests/compile_test.py +++ b/tests/gpu_tests/compile_test.py @@ -9,6 +9,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.common.score_arguments import pytest_score_arguments from kronfluence.utils.constants import ( ALL_MODULE_NAME, COVARIANCE_FACTOR_NAMES, @@ -126,67 +127,58 @@ def test_lambda_shared_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) - # - # 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, - # ) + def test_pairwise_scores(self) -> None: + 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) + pairwise_scores = self.analyzer.load_pairwise_scores(scores_name=OLD_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=ATOL, + rtol=RTOL, + ) + + def test_self_scores(self) -> None: + 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) + 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=ATOL, + rtol=RTOL, + ) if __name__ == "__main__":