From 398503cc4a5c1f5c3d231b6a0ecdc2ebd6364cfd Mon Sep 17 00:00:00 2001 From: Juhan Bae Date: Sun, 7 Jul 2024 16:27:18 -0400 Subject: [PATCH] Modify all other tests --- tests/gpu_tests/amp_test.py | 147 +++++++++++++-------------- tests/gpu_tests/compile_test.py | 171 ++++++++++++++++++-------------- tests/gpu_tests/cpu_test.py | 14 +-- tests/gpu_tests/fsdp_test.py | 2 +- 4 files changed, 172 insertions(+), 162 deletions(-) diff --git a/tests/gpu_tests/amp_test.py b/tests/gpu_tests/amp_test.py index 8f2ad17..c315db6 100644 --- 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 index a997ddf..ca9e008 100644 --- a/tests/gpu_tests/compile_test.py +++ b/tests/gpu_tests/compile_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, @@ -15,7 +16,7 @@ ) 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__": diff --git a/tests/gpu_tests/cpu_test.py b/tests/gpu_tests/cpu_test.py index 455327d..93dde7a 100644 --- a/tests/gpu_tests/cpu_test.py +++ b/tests/gpu_tests/cpu_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, @@ -15,7 +16,7 @@ ) 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" @@ -43,12 +44,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, @@ -66,8 +62,8 @@ def test_covariance_matrices(self) -> None: assert check_tensor_dict_equivalence( covariance_factors[name], new_covariance_factors[name], - atol=1e-5, - rtol=1e-3, + atol=ATOL, + rtol=RTOL, ) def test_lambda_matrices(self): diff --git a/tests/gpu_tests/fsdp_test.py b/tests/gpu_tests/fsdp_test.py index a8e1773..6be6579 100644 --- a/tests/gpu_tests/fsdp_test.py +++ b/tests/gpu_tests/fsdp_test.py @@ -87,7 +87,7 @@ def test_covariance_matrices(self) -> None: 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 = pytest_factor_arguments() self.analyzer.fit_lambda_matrices(