Skip to content

Commit

Permalink
Lint fix & Simplfiy pytest codes
Browse files Browse the repository at this point in the history
  • Loading branch information
@pomonam
pomonam committed Jun 29, 2024
1 parent 60a16d5 commit 2569fbb
Show file tree
Hide file tree
Showing 6 changed files with 120 additions and 121 deletions.
86 changes: 41 additions & 45 deletions tests/factors/test_covariances.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,10 @@
from kronfluence.utils.dataset import DataLoaderKwargs
from tests.utils import (
ATOL,
DEFAULT_FACTORS_NAME,
RTOL,
check_tensor_dict_equivalence,
custom_factors_name,
prepare_model_and_analyzer,
prepare_test,
)
Expand All @@ -29,11 +31,10 @@
[
"mlp",
"repeated_mlp",
"mlp_checkpoint",
"conv",
"conv_bn",
"bert",
"gpt",
"gpt_checkpoint",
],
)
@pytest.mark.parametrize("activation_covariance_dtype", [torch.float32, torch.bfloat16])
Expand Down Expand Up @@ -62,17 +63,16 @@ def test_fit_covariance_matrices(
factor_args = default_factor_arguments()
factor_args.activation_covariance_dtype = activation_covariance_dtype
factor_args.gradient_covariance_dtype = gradient_covariance_dtype
factors_name = f"pytest_{test_name}_{test_fit_covariance_matrices.__name__}"
analyzer.fit_covariance_matrices(
factors_name=factors_name,
factors_name=DEFAULT_FACTORS_NAME,
factor_args=factor_args,
dataset=train_dataset,
dataloader_kwargs=kwargs,
factor_args=factor_args,
per_device_batch_size=train_size // 4,
overwrite_output_dir=True,
)
covariance_factors = analyzer.load_covariance_matrices(
factors_name=factors_name,
factors_name=DEFAULT_FACTORS_NAME,
)
assert set(covariance_factors.keys()) == set(COVARIANCE_FACTOR_NAMES)
assert len(covariance_factors[ACTIVATION_COVARIANCE_MATRIX_NAME]) > 0
Expand Down Expand Up @@ -103,35 +103,32 @@ def test_covariance_matrices_batch_size_equivalence(
seed=seed,
)
kwargs = DataLoaderKwargs(collate_fn=data_collator)
model = model.to(dtype=torch.float64)
model, analyzer = prepare_model_and_analyzer(
model=model,
task=task,
)

factor_args = test_factor_arguments()
analyzer.fit_covariance_matrices(
factors_name=f"pytest_{test_name}_{test_covariance_matrices_batch_size_equivalence.__name__}_bs1",
factors_name=DEFAULT_FACTORS_NAME,
dataset=train_dataset,
factor_args=factor_args,
per_device_batch_size=1,
overwrite_output_dir=True,
dataloader_kwargs=kwargs,
)
bs1_covariance_factors = analyzer.load_covariance_matrices(
factors_name=f"pytest_{test_name}_{test_covariance_matrices_batch_size_equivalence.__name__}_bs1"
)
bs1_covariance_factors = analyzer.load_covariance_matrices(factors_name=DEFAULT_FACTORS_NAME)

analyzer.fit_covariance_matrices(
factors_name=f"pytest_{test_name}_{test_covariance_matrices_batch_size_equivalence.__name__}_bs8",
factors_name=custom_factors_name(name="bs8"),
dataset=train_dataset,
factor_args=factor_args,
per_device_batch_size=8,
overwrite_output_dir=True,
dataloader_kwargs=kwargs,
)
bs8_covariance_factors = analyzer.load_covariance_matrices(
factors_name=f"pytest_{test_name}_{test_covariance_matrices_batch_size_equivalence.__name__}_bs8"
)
bs8_covariance_factors = analyzer.load_covariance_matrices(factors_name=custom_factors_name(name="bs8"))

for name in COVARIANCE_FACTOR_NAMES:
assert check_tensor_dict_equivalence(
Expand All @@ -146,7 +143,7 @@ def test_covariance_matrices_batch_size_equivalence(
"test_name",
[
"mlp",
"conv",
"conv_bn",
],
)
@pytest.mark.parametrize("data_partition_size", [2, 4])
Expand All @@ -167,35 +164,35 @@ def test_covariance_matrices_partition_equivalence(
seed=seed,
)
kwargs = DataLoaderKwargs(collate_fn=data_collator)
model = model.to(dtype=torch.float64)
model, analyzer = prepare_model_and_analyzer(
model=model,
task=task,
)

factor_args = test_factor_arguments()
factors_name = f"pytest_{test_name}_{test_covariance_matrices_partition_equivalence.__name__}"
analyzer.fit_covariance_matrices(
factors_name=factors_name,
factors_name=DEFAULT_FACTORS_NAME,
dataset=train_dataset,
factor_args=factor_args,
per_device_batch_size=8,
overwrite_output_dir=True,
dataloader_kwargs=kwargs,
)
covariance_factors = analyzer.load_covariance_matrices(factors_name=factors_name)
covariance_factors = analyzer.load_covariance_matrices(factors_name=DEFAULT_FACTORS_NAME)

factor_args.covariance_data_partition_size = data_partition_size
factor_args.covariance_module_partition_size = module_partition_size
analyzer.fit_covariance_matrices(
factors_name=f"pytest_{test_name}_partitioned_{data_partition_size}_{module_partition_size}",
factors_name=custom_factors_name(f"{data_partition_size}_{module_partition_size}"),
dataset=train_dataset,
factor_args=factor_args,
per_device_batch_size=7,
overwrite_output_dir=True,
dataloader_kwargs=kwargs,
)
partitioned_covariance_factors = analyzer.load_covariance_matrices(
factors_name=f"pytest_{test_name}_partitioned_{data_partition_size}_{module_partition_size}",
factors_name=custom_factors_name(f"{data_partition_size}_{module_partition_size}"),
)

for name in COVARIANCE_FACTOR_NAMES:
Expand All @@ -207,7 +204,7 @@ def test_covariance_matrices_partition_equivalence(
)


@pytest.mark.parametrize("test_name", ["bert", "wrong_bert"])
@pytest.mark.parametrize("test_name", ["bert", "wrong_bert", "roberta"])
@pytest.mark.parametrize("train_size", [213])
@pytest.mark.parametrize("seed", [3])
def test_covariance_matrices_attention_mask(
Expand Down Expand Up @@ -235,37 +232,35 @@ def test_covariance_matrices_attention_mask(
seed=seed,
)
model = model.to(dtype=torch.float64)

kwargs = DataLoaderKwargs(collate_fn=data_collator)
model, analyzer = prepare_model_and_analyzer(
model=model,
task=task,
)
kwargs = DataLoaderKwargs(collate_fn=data_collator)

factor_args = test_factor_arguments()
factors_name = f"pytest_{test_name}_{test_covariance_matrices_attention_mask.__name__}"
analyzer.fit_covariance_matrices(
factors_name=factors_name,
factors_name=DEFAULT_FACTORS_NAME,
dataset=train_dataset,
factor_args=factor_args,
per_device_batch_size=train_size // 4,
overwrite_output_dir=True,
dataloader_kwargs=kwargs,
)
covariance_factors = analyzer.load_covariance_matrices(
factors_name=factors_name,
factors_name=DEFAULT_FACTORS_NAME,
)

analyzer.fit_covariance_matrices(
factors_name=factors_name + "_no_pad",
factors_name=custom_factors_name("no_pad"),
dataset=no_padded_train_dataset,
factor_args=factor_args,
per_device_batch_size=1,
overwrite_output_dir=True,
dataloader_kwargs=kwargs,
)
no_padded_covariance_factors = analyzer.load_covariance_matrices(
factors_name=factors_name + "_no_pad",
factors_name=custom_factors_name("no_pad"),
)

for name in COVARIANCE_FACTOR_NAMES:
Expand Down Expand Up @@ -304,33 +299,33 @@ def test_covariance_matrices_automatic_batch_size(
seed=seed,
)
kwargs = DataLoaderKwargs(collate_fn=data_collator)
model = model.to(dtype=torch.float64)
model, analyzer = prepare_model_and_analyzer(
model=model,
task=task,
)

factor_args = test_factor_arguments()
factors_name = f"pytest_{test_name}_{test_covariance_matrices_automatic_batch_size.__name__}"
analyzer.fit_covariance_matrices(
factors_name=factors_name,
factors_name=DEFAULT_FACTORS_NAME,
dataset=train_dataset,
factor_args=factor_args,
per_device_batch_size=8,
overwrite_output_dir=True,
dataloader_kwargs=kwargs,
)
covariance_factors = analyzer.load_covariance_matrices(factors_name=factors_name)
covariance_factors = analyzer.load_covariance_matrices(factors_name=DEFAULT_FACTORS_NAME)

analyzer.fit_covariance_matrices(
factors_name=factors_name + "_auto",
factors_name=custom_factors_name("auto"),
dataset=train_dataset,
factor_args=factor_args,
per_device_batch_size=None,
overwrite_output_dir=True,
dataloader_kwargs=kwargs,
)
auto_covariance_factors = analyzer.load_covariance_matrices(
factors_name=factors_name + "_auto",
factors_name=custom_factors_name("auto"),
)

for name in COVARIANCE_FACTOR_NAMES:
Expand Down Expand Up @@ -369,16 +364,15 @@ def test_covariance_matrices_max_examples(
factor_args.covariance_max_examples = MAX_EXAMPLES
factor_args.covariance_data_partition_size = data_partition_size

factors_name = f"pytest_{test_name}_{test_covariance_matrices_max_examples.__name__}"
analyzer.fit_covariance_matrices(
factors_name=factors_name,
factors_name=DEFAULT_FACTORS_NAME,
dataset=train_dataset,
factor_args=factor_args,
per_device_batch_size=32,
overwrite_output_dir=True,
dataloader_kwargs=kwargs,
)
covariance_factors = analyzer.load_covariance_matrices(factors_name=factors_name)
covariance_factors = analyzer.load_covariance_matrices(factors_name=DEFAULT_FACTORS_NAME)

for num_examples in covariance_factors[NUM_ACTIVATION_COVARIANCE_PROCESSED].values():
assert num_examples == MAX_EXAMPLES
Expand All @@ -391,7 +385,7 @@ def test_covariance_matrices_max_examples(
"test_name",
[
"mlp",
"conv",
"conv_bn",
],
)
@pytest.mark.parametrize("train_size", [101])
Expand All @@ -415,28 +409,28 @@ def test_covariance_matrices_amp(

factor_args = test_factor_arguments()
analyzer.fit_covariance_matrices(
factors_name=f"pytest_{test_name}_{test_covariance_matrices_amp.__name__}",
factors_name=DEFAULT_FACTORS_NAME,
dataset=train_dataset,
per_device_batch_size=8,
overwrite_output_dir=True,
factor_args=factor_args,
dataloader_kwargs=kwargs,
)
covariance_factors = analyzer.load_covariance_matrices(
factors_name=f"pytest_{test_name}_{test_covariance_matrices_amp.__name__}"
factors_name=DEFAULT_FACTORS_NAME,
)

factor_args.amp_dtype = torch.float16
analyzer.fit_covariance_matrices(
factors_name=f"pytest_{test_name}_{test_covariance_matrices_amp.__name__}_amp",
factors_name=custom_factors_name("amp"),
dataset=train_dataset,
per_device_batch_size=8,
overwrite_output_dir=True,
factor_args=factor_args,
dataloader_kwargs=kwargs,
)
amp_covariance_factors = analyzer.load_covariance_matrices(
factors_name=f"pytest_{test_name}_{test_covariance_matrices_amp.__name__}_amp"
factors_name=custom_factors_name("amp"),
)

for name in COVARIANCE_FACTOR_NAMES:
Expand All @@ -460,41 +454,43 @@ def test_covariance_matrices_gradient_checkpoint(
train_size=train_size,
seed=seed,
)
model = model.to(dtype=torch.float64)
model, analyzer = prepare_model_and_analyzer(
model=model,
task=task,
)

factor_args = test_factor_arguments()
analyzer.fit_covariance_matrices(
factors_name=f"pytest_{test_covariance_matrices_gradient_checkpoint.__name__}",
factors_name=DEFAULT_FACTORS_NAME,
dataset=train_dataset,
per_device_batch_size=8,
overwrite_output_dir=True,
factor_args=factor_args,
)
covariance_factors = analyzer.load_covariance_matrices(
factors_name=f"pytest_{test_covariance_matrices_gradient_checkpoint.__name__}",
factors_name=DEFAULT_FACTORS_NAME,
)

model, _, _, _, task = prepare_test(
test_name="mlp_checkpoint",
train_size=train_size,
seed=seed,
)
model = model.to(dtype=torch.float64)
model, analyzer = prepare_model_and_analyzer(
model=model,
task=task,
)
analyzer.fit_covariance_matrices(
factors_name=f"pytest_{test_covariance_matrices_gradient_checkpoint.__name__}_cp",
factors_name=custom_factors_name("cp"),
dataset=train_dataset,
per_device_batch_size=4,
overwrite_output_dir=True,
factor_args=factor_args,
)
checkpoint_covariance_factors = analyzer.load_covariance_matrices(
factors_name=f"pytest_{test_covariance_matrices_gradient_checkpoint.__name__}_cp",
factors_name=custom_factors_name("cp"),
)

assert check_tensor_dict_equivalence(
Expand Down
Loading

0 comments on commit 2569fbb

Please sign in to comment.