diff --git a/spd/experiments/ih/ih_config.yaml b/spd/experiments/ih/ih_config.yaml index 1b905a166..389b98798 100644 --- a/spd/experiments/ih/ih_config.yaml +++ b/spd/experiments/ih/ih_config.yaml @@ -22,13 +22,16 @@ target_module_patterns: [ "blocks.*.attn.out_proj", ] -faithfulness_coeff: 100 -ci_recon_coeff: 1 -stochastic_recon_coeff: 1 -ci_recon_layerwise_coeff: null -stochastic_recon_layerwise_coeff: 1 -importance_minimality_coeff: 1e-2 -pnorm: 0.1 +loss_metric_configs: + - classname: "ImportanceMinimalityLoss" + coeff: 1e-2 + pnorm: 0.1 + - classname: "CIMaskedReconLoss" + coeff: 1.0 + - classname: "StochasticReconLoss" + coeff: 1.0 + - classname: "StochasticReconLayerwiseLoss" + coeff: 1.0 output_loss_type: kl ci_fn_type: "vector_mlp" ci_fn_hidden_dims: [128] diff --git a/tests/test_gpt2.py b/tests/test_gpt2.py deleted file mode 100644 index 535429ace..000000000 --- a/tests/test_gpt2.py +++ /dev/null @@ -1,150 +0,0 @@ -import pytest -from transformers import PreTrainedModel - -from spd.configs import ( - CI_L0Config, - Config, - FaithfulnessLossConfig, - ImportanceMinimalityLossConfig, - StochasticReconLayerwiseLossConfig, - StochasticReconLossConfig, -) -from spd.data import DatasetConfig, create_data_loader -from spd.experiments.lm.configs import LMTaskConfig -from spd.identity_insertion import insert_identity_operations_ -from spd.run_spd import optimize -from spd.utils.general_utils import resolve_class, set_seed - - -@pytest.mark.slow -def test_gpt_2_decomposition_happy_path() -> None: - """Test that SPD decomposition works on for GPT-2""" - set_seed(0) - device = "cpu" - - # Create config similar to the gpt-2 config in gpt2_config.yaml - config = Config( - # WandB - wandb_project=None, # Disable wandb for testing - wandb_run_name=None, - wandb_run_name_prefix="", - # General - seed=0, - C=10, # Smaller C for faster testing - n_mask_samples=1, - ci_fn_type="vector_mlp", - ci_fn_hidden_dims=[128], - target_module_patterns=["transformer.h.2.attn.c_attn", "transformer.h.3.mlp.c_fc"], - identity_module_patterns=["transformer.h.1.attn.c_attn"], - loss_metric_configs=[ - ImportanceMinimalityLossConfig( - coeff=1e-2, - pnorm=0.9, - eps=1e-12, - ), - StochasticReconLayerwiseLossConfig(coeff=1.0), - StochasticReconLossConfig(coeff=1.0), - FaithfulnessLossConfig(coeff=200), - ], - output_loss_type="kl", - # Training - lr=1e-3, - batch_size=4, - steps=2, - lr_schedule="cosine", - lr_exponential_halflife=None, - lr_warmup_pct=0.01, - n_eval_steps=1, - # Logging & Saving - train_log_freq=50, # Print at step 0, 50, and 100 - eval_freq=500, - eval_batch_size=1, - slow_eval_freq=500, - slow_eval_on_first_step=False, - save_freq=None, - ci_alive_threshold=0.1, - n_examples_until_dead=200, # print_freq * batch_size = 50 * 4 - eval_metric_configs=[ - CI_L0Config(groups=None), - ], - # Pretrained model info - pretrained_model_class="transformers.GPT2LMHeadModel", - pretrained_model_path=None, - pretrained_model_name="SimpleStories/test-SimpleStories-gpt2-1.25M", - pretrained_model_output_attr="logits", - tokenizer_name="SimpleStories/test-SimpleStories-gpt2-1.25M", - # Task Specific - task_config=LMTaskConfig( - task_name="lm", - max_seq_len=16, - buffer_size=1000, - dataset_name="SimpleStories/SimpleStories", - column_name="story", - train_data_split="train[:100]", - eval_data_split="test[100:200]", - ), - ) - - assert isinstance(config.task_config, LMTaskConfig), "task_config not LMTaskConfig" - - # Create a GPT-2 model - hf_model_class = resolve_class(config.pretrained_model_class) - assert issubclass(hf_model_class, PreTrainedModel), ( - f"Model class {hf_model_class} should be a subclass of PreTrainedModel which " - "defines a `from_pretrained` method" - ) - assert config.pretrained_model_name is not None - target_model = hf_model_class.from_pretrained(config.pretrained_model_name) - target_model.eval() - - if config.identity_module_patterns is not None: - insert_identity_operations_(target_model, identity_patterns=config.identity_module_patterns) - - train_data_config = DatasetConfig( - name=config.task_config.dataset_name, - hf_tokenizer_path=config.pretrained_model_name, - split=config.task_config.train_data_split, - n_ctx=config.task_config.max_seq_len, - is_tokenized=config.task_config.is_tokenized, - streaming=config.task_config.streaming, - column_name=config.task_config.column_name, - seed=None, - ) - - train_loader, _tokenizer = create_data_loader( - dataset_config=train_data_config, - batch_size=config.batch_size, - buffer_size=config.task_config.buffer_size, - global_seed=config.seed, - ) - - eval_data_config = DatasetConfig( - name=config.task_config.dataset_name, - hf_tokenizer_path=config.pretrained_model_name, - split=config.task_config.eval_data_split, - n_ctx=config.task_config.max_seq_len, - is_tokenized=config.task_config.is_tokenized, - streaming=config.task_config.streaming, - column_name=config.task_config.column_name, - seed=None, - ) - eval_loader, _ = create_data_loader( - dataset_config=eval_data_config, - batch_size=config.batch_size, - buffer_size=config.task_config.buffer_size, - global_seed=config.seed + 1, - ) - - # Run optimize function - optimize( - target_model=target_model, - config=config, - device=device, - train_loader=train_loader, - eval_loader=eval_loader, - n_eval_steps=config.n_eval_steps, - out_dir=None, - ) - - # Basic assertion to ensure the test ran - assert True, "Test completed successfully" diff --git a/tests/test_gpt2_configs.py b/tests/test_gpt2_configs.py new file mode 100644 index 000000000..bd92372a3 --- /dev/null +++ b/tests/test_gpt2_configs.py @@ -0,0 +1,123 @@ +import pytest + +from spd.configs import Config +from spd.data import DatasetConfig, create_data_loader +from spd.experiments.lm.configs import LMTaskConfig +from spd.identity_insertion import insert_identity_operations_ +from spd.registry import get_experiment_config_file_contents +from spd.run_spd import optimize +from spd.utils.general_utils import resolve_class, set_seed +from spd.utils.run_utils import apply_nested_updates + +# Config-specific test parameters for different GPT2 configurations +GPT2_CONFIG_PARAMS = { + "ss_gpt2_simple": { + # Uses simple_stories_train.models.gpt2_simple.GPT2Simple (wandb-hosted model) + "target_module_patterns": ["h.2.attn.q_proj", "h.3.mlp.c_fc"], + "identity_module_patterns": ["h.1.attn.q_proj"], + }, + "ss_gpt2": { + # Uses transformers.GPT2LMHeadModel (HuggingFace transformers library) + "target_module_patterns": ["transformer.h.1.mlp.c_fc"], + "identity_module_patterns": None, + }, +} + + +@pytest.mark.slow +@pytest.mark.parametrize("experiment_name", ["ss_gpt2_simple", "ss_gpt2"]) +def test_gpt2_decomposition_happy_path(experiment_name: str) -> None: + """Test that SPD decomposition works on different GPT-2 configurations. + + Tests both: + - ss_gpt2_simple: Uses simple_stories_train GPT2Simple model (wandb-hosted) + - ss_gpt2: Uses transformers.GPT2LMHeadModel (HuggingFace transformers library) + """ + set_seed(0) + device = "cpu" + + config_params = GPT2_CONFIG_PARAMS[experiment_name] + base_config_dict = get_experiment_config_file_contents(experiment_name) + test_overrides = { + "wandb_project": None, + "C": 10, + "steps": 2, + "batch_size": 4, + "eval_batch_size": 1, + "train_log_freq": 50, + "n_examples_until_dead": 999, + "task_config.max_seq_len": 8, + "task_config.train_data_split": "train[:100]", + "task_config.eval_data_split": "test[100:200]", + "target_module_patterns": config_params["target_module_patterns"], + "identity_module_patterns": config_params["identity_module_patterns"], + "eval_metric_configs": [], # Disable eval metrics to avoid layer matching issues + } + config_dict = apply_nested_updates(base_config_dict, test_overrides) + config = Config(**config_dict) + + assert isinstance(config.task_config, LMTaskConfig), "task_config not LMTaskConfig" + pretrained_model_class = resolve_class(config.pretrained_model_class) + assert hasattr(pretrained_model_class, "from_pretrained"), ( + f"Model class {pretrained_model_class} should have a `from_pretrained` method" + ) + assert config.pretrained_model_name is not None + + # Handle simple_stories_train models specially (they use from_run_info) + if config.pretrained_model_class.startswith("simple_stories_train"): + from simple_stories_train.run_info import RunInfo as SSRunInfo + + run_info = SSRunInfo.from_path(config.pretrained_model_name) + assert hasattr(pretrained_model_class, "from_run_info") + target_model = pretrained_model_class.from_run_info(run_info) # pyright: ignore[reportAttributeAccessIssue] + else: + target_model = pretrained_model_class.from_pretrained(config.pretrained_model_name) # pyright: ignore[reportAttributeAccessIssue] + target_model.eval() + + if config.identity_module_patterns is not None: + insert_identity_operations_(target_model, identity_patterns=config.identity_module_patterns) + + train_data_config = DatasetConfig( + name=config.task_config.dataset_name, + hf_tokenizer_path=config.tokenizer_name, + split=config.task_config.train_data_split, + n_ctx=config.task_config.max_seq_len, + is_tokenized=config.task_config.is_tokenized, + streaming=config.task_config.streaming, + column_name=config.task_config.column_name, + seed=None, + ) + + train_loader, _tokenizer = create_data_loader( + dataset_config=train_data_config, + batch_size=config.batch_size, + buffer_size=config.task_config.buffer_size, + global_seed=config.seed, + ) + + eval_data_config = DatasetConfig( + name=config.task_config.dataset_name, + hf_tokenizer_path=config.tokenizer_name, + split=config.task_config.eval_data_split, + n_ctx=config.task_config.max_seq_len, + is_tokenized=config.task_config.is_tokenized, + streaming=config.task_config.streaming, + column_name=config.task_config.column_name, + seed=None, + ) + eval_loader, _ = create_data_loader( + dataset_config=eval_data_config, + batch_size=config.batch_size, + buffer_size=config.task_config.buffer_size, + global_seed=config.seed + 1, + ) + + optimize( + target_model=target_model, + config=config, + device=device, + train_loader=train_loader, + eval_loader=eval_loader, + n_eval_steps=config.n_eval_steps, + out_dir=None, + ) diff --git a/tests/test_ih_transformer.py b/tests/test_ih_transformer.py index 361223e5d..05f7c9736 100644 --- a/tests/test_ih_transformer.py +++ b/tests/test_ih_transformer.py @@ -1,29 +1,42 @@ import pytest +import yaml -from spd.configs import ( - CI_L0Config, - Config, - FaithfulnessLossConfig, - ImportanceMinimalityLossConfig, - StochasticHiddenActsReconLossConfig, - StochasticReconLayerwiseLossConfig, - StochasticReconLossConfig, -) -from spd.experiments.ih.configs import IHTaskConfig, InductionModelConfig +from spd.configs import Config +from spd.experiments.ih.configs import InductionModelConfig from spd.experiments.ih.model import InductionTransformer from spd.identity_insertion import insert_identity_operations_ from spd.run_spd import optimize +from spd.settings import REPO_ROOT from spd.utils.data_utils import DatasetGeneratedDataLoader, InductionDataset from spd.utils.general_utils import set_seed +from spd.utils.run_utils import apply_nested_updates @pytest.mark.slow def test_ih_transformer_decomposition_happy_path() -> None: - """Test that SPD decomposition works on a 2-layer, 1 head attention-only Transformer model""" + """Test that SPD decomposition works on a 2-layer, 1 head attention-only Transformer model. + + TODO: Use a real pretrained_model_path in the config instead of randomly initializing one. + """ set_seed(0) device = "cpu" - # Create a 2-layer InductionTransformer config + config_path = REPO_ROOT / "spd/experiments/ih/ih_config.yaml" + base_config_dict = yaml.safe_load(config_path.read_text()) + test_overrides = { + "wandb_project": None, + "C": 10, + "steps": 2, + "batch_size": 4, + "eval_batch_size": 1, + "train_log_freq": 50, + "n_examples_until_dead": 999, + "pretrained_model_path": None, + "n_eval_steps": 1, + } + config_dict = apply_nested_updates(base_config_dict, test_overrides) + config = Config(**config_dict) + ih_transformer_config = InductionModelConfig( vocab_size=128, d_model=16, @@ -36,67 +49,6 @@ def test_ih_transformer_decomposition_happy_path() -> None: ff_fanout=4, ) - # Create config similar to the induction_head transformer config in ih_config.yaml - config = Config( - # WandB - wandb_project=None, # Disable wandb for testing - wandb_run_name=None, - wandb_run_name_prefix="", - # General - seed=0, - C=10, # Smaller C for faster testing - n_mask_samples=1, - ci_fn_type="vector_mlp", - ci_fn_hidden_dims=[128], - target_module_patterns=["blocks.*.attn.q_proj", "blocks.*.attn.k_proj"], - identity_module_patterns=["blocks.*.attn.q_proj"], - # Loss Coefficients - loss_metric_configs=[ - ImportanceMinimalityLossConfig( - coeff=1e-2, - pnorm=0.9, - eps=1e-12, - ), - StochasticReconLayerwiseLossConfig(coeff=1.0), - StochasticReconLossConfig(coeff=1.0), - FaithfulnessLossConfig(coeff=200), - ], - output_loss_type="kl", - # Training - lr=1e-3, - batch_size=4, - steps=2, - lr_schedule="cosine", - lr_exponential_halflife=None, - lr_warmup_pct=0.01, - n_eval_steps=1, - # Logging & Saving - train_log_freq=50, # Print at step 0, 50, and 100 - eval_freq=500, - eval_batch_size=1, - slow_eval_freq=500, - slow_eval_on_first_step=True, - save_freq=None, - ci_alive_threshold=0.1, - n_examples_until_dead=200, # print_freq * batch_size = 50 * 4 - eval_metric_configs=[ - CI_L0Config(groups=None), - StochasticHiddenActsReconLossConfig(), - ], - # Pretrained model info - pretrained_model_class="spd.experiments.ih.model.InductionTransformer", - pretrained_model_path=None, - pretrained_model_name=None, - pretrained_model_output_attr=None, - tokenizer_name=None, - # Task Specific - task_config=IHTaskConfig( - task_name="induction_head", - ), - ) - - # Create a pretrained model - target_model = InductionTransformer(ih_transformer_config).to(device) target_model.eval() target_model.requires_grad_(False) @@ -118,7 +70,6 @@ def test_ih_transformer_decomposition_happy_path() -> None: dataset, batch_size=config.microbatch_size, shuffle=False ) - # Run optimize function optimize( target_model=target_model, config=config, @@ -129,5 +80,4 @@ def test_ih_transformer_decomposition_happy_path() -> None: out_dir=None, ) - # Basic assertion to ensure the test ran assert True, "Test completed successfully" diff --git a/tests/test_resid_mlp.py b/tests/test_resid_mlp.py index 076083a7d..1b8dac289 100644 --- a/tests/test_resid_mlp.py +++ b/tests/test_resid_mlp.py @@ -1,16 +1,13 @@ -from spd.configs import ( - Config, - FaithfulnessLossConfig, - ImportanceMinimalityLossConfig, - StochasticReconLossConfig, -) +from spd.configs import Config from spd.experiments.resid_mlp.configs import ResidMLPModelConfig, ResidMLPTaskConfig from spd.experiments.resid_mlp.models import ResidMLP from spd.experiments.resid_mlp.resid_mlp_dataset import ResidMLPDataset from spd.identity_insertion import insert_identity_operations_ +from spd.registry import get_experiment_config_file_contents from spd.run_spd import optimize from spd.utils.data_utils import DatasetGeneratedDataLoader from spd.utils.general_utils import set_seed +from spd.utils.run_utils import apply_nested_updates def test_resid_mlp_decomposition_happy_path() -> None: @@ -18,7 +15,25 @@ def test_resid_mlp_decomposition_happy_path() -> None: set_seed(0) device = "cpu" - # Create a 2-layer ResidMLP config + base_config_dict = get_experiment_config_file_contents("resid_mlp2") + test_overrides = { + "wandb_project": None, + "C": 10, + "steps": 3, + "batch_size": 4, + "eval_batch_size": 4, + "train_log_freq": 50, + "n_examples_until_dead": 999, + "eval_metric_configs.IdentityCIError.identity_ci": [ + {"layer_pattern": "layers.*.mlp_in", "n_features": 5} + ], + "eval_metric_configs.IdentityCIError.dense_ci": [ + {"layer_pattern": "layers.*.mlp_out", "k": 3} + ], + } + config_dict = apply_nested_updates(base_config_dict, test_overrides) + config = Config(**config_dict) + resid_mlp_model_config = ResidMLPModelConfig( n_features=5, d_embed=4, @@ -29,63 +44,6 @@ def test_resid_mlp_decomposition_happy_path() -> None: out_bias=True, ) - # Create config similar to the 2-layer config in resid_mlp2_config.yaml - config = Config( - # WandB - wandb_project=None, # Disable wandb for testing - wandb_run_name=None, - wandb_run_name_prefix="", - # General - seed=0, - C=10, # Smaller C for faster testing - n_mask_samples=1, - ci_fn_type="mlp", - ci_fn_hidden_dims=[8], - loss_metric_configs=[ - ImportanceMinimalityLossConfig( - coeff=3e-3, - pnorm=0.9, - eps=1e-12, - ), - StochasticReconLossConfig(coeff=1.0), - FaithfulnessLossConfig(coeff=1.0), - ], - target_module_patterns=["layers.*.mlp_in", "layers.*.mlp_out"], - identity_module_patterns=["layers.*.mlp_in"], - output_loss_type="mse", - # Training - lr=1e-3, - batch_size=4, - steps=3, # Run more steps to see improvement - lr_schedule="cosine", - lr_exponential_halflife=None, - lr_warmup_pct=0.01, - n_eval_steps=1, - eval_freq=10, - eval_batch_size=4, - slow_eval_freq=10, - slow_eval_on_first_step=True, - # Logging & Saving - train_log_freq=50, # Print at step 0, 50, and 100 - save_freq=None, - ci_alive_threshold=0.1, - n_examples_until_dead=200, # print_freq * batch_size = 50 * 4 - # Pretrained model info - pretrained_model_class="spd.experiments.resid_mlp.models.ResidMLP", - pretrained_model_path=None, - pretrained_model_name=None, - pretrained_model_output_attr=None, - tokenizer_name=None, - # Task Specific - task_config=ResidMLPTaskConfig( - task_name="resid_mlp", - feature_probability=0.01, - data_generation_type="at_least_zero_active", - ), - ) - - # Create a pretrained model - target_model = ResidMLP(config=resid_mlp_model_config).to(device) target_model.requires_grad_(False) @@ -93,12 +51,11 @@ def test_resid_mlp_decomposition_happy_path() -> None: insert_identity_operations_(target_model, identity_patterns=config.identity_module_patterns) assert isinstance(config.task_config, ResidMLPTaskConfig) - # Create dataset dataset = ResidMLPDataset( n_features=resid_mlp_model_config.n_features, feature_probability=config.task_config.feature_probability, device=device, - calc_labels=False, # Our labels will be the output of the target model + calc_labels=False, label_type=None, act_fn_name=None, label_fn_seed=None, @@ -114,7 +71,6 @@ def test_resid_mlp_decomposition_happy_path() -> None: dataset, batch_size=config.eval_batch_size, shuffle=False ) - # Run optimize function optimize( target_model=target_model, config=config, @@ -125,5 +81,4 @@ def test_resid_mlp_decomposition_happy_path() -> None: out_dir=None, ) - # Basic assertion to ensure the test ran assert True, "Test completed successfully" diff --git a/tests/test_tms.py b/tests/test_tms.py index 780665461..b69b4dcb2 100644 --- a/tests/test_tms.py +++ b/tests/test_tms.py @@ -3,20 +3,16 @@ import torch from torch import nn -from spd.configs import ( - Config, - FaithfulnessLossConfig, - ImportanceMinimalityLossConfig, - StochasticReconLayerwiseLossConfig, - StochasticReconLossConfig, -) +from spd.configs import Config from spd.experiments.tms.configs import TMSModelConfig, TMSTaskConfig, TMSTrainConfig from spd.experiments.tms.models import TMSModel from spd.experiments.tms.train_tms import get_model_and_dataloader, train from spd.identity_insertion import insert_identity_operations_ +from spd.registry import get_experiment_config_file_contents from spd.run_spd import optimize from spd.utils.data_utils import DatasetGeneratedDataLoader, SparseFeatureDataset from spd.utils.general_utils import set_seed +from spd.utils.run_utils import apply_nested_updates def test_tms_decomposition_happy_path() -> None: @@ -24,7 +20,22 @@ def test_tms_decomposition_happy_path() -> None: set_seed(0) device = "cpu" - # Create a TMS model config similar to the one in tms_config.yaml + base_config_dict = get_experiment_config_file_contents("tms_5-2") + test_overrides = { + "wandb_project": None, + "C": 10, + "steps": 3, + "batch_size": 4, + "eval_batch_size": 4, + "train_log_freq": 2, + "n_examples_until_dead": 999, + "faithfulness_warmup_steps": 2, + "target_module_patterns": ["linear1", "linear2", "hidden_layers.0"], + "identity_module_patterns": ["linear1"], + } + config_dict = apply_nested_updates(base_config_dict, test_overrides) + config = Config(**config_dict) + tms_model_config = TMSModelConfig( n_features=5, n_hidden=2, @@ -34,66 +45,6 @@ def test_tms_decomposition_happy_path() -> None: device=device, ) - # Create config similar to tms_config.yaml - config = Config( - # WandB - wandb_project=None, # Disable wandb for testing - wandb_run_name=None, - wandb_run_name_prefix="", - # General - seed=0, - C=10, # Smaller C for faster testing - n_mask_samples=1, - ci_fn_type="mlp", - ci_fn_hidden_dims=[8], - target_module_patterns=["linear1", "linear2", "hidden_layers.0"], - identity_module_patterns=["linear1"], - loss_metric_configs=[ - ImportanceMinimalityLossConfig( - coeff=3e-3, - pnorm=2.0, - eps=1e-12, - ), - StochasticReconLayerwiseLossConfig(coeff=1.0), - StochasticReconLossConfig(coeff=1.0), - FaithfulnessLossConfig(coeff=1.0), - ], - output_loss_type="mse", - # Training - lr=1e-3, - batch_size=4, - steps=3, # Run only a few steps for the test - lr_schedule="cosine", - lr_exponential_halflife=None, - lr_warmup_pct=0.0, - n_eval_steps=1, - # Faithfulness Warmup - faithfulness_warmup_steps=2, - faithfulness_warmup_lr=0.001, - faithfulness_warmup_weight_decay=0.0, - # Logging & Saving - train_log_freq=2, - save_freq=None, - ci_alive_threshold=0.1, - n_examples_until_dead=8, # print_freq * batch_size = 2 * 4 - eval_batch_size=4, - eval_freq=10, - slow_eval_freq=10, - # Pretrained model info - pretrained_model_class="spd.experiments.tms.models.TMSModel", - pretrained_model_path=None, - pretrained_model_name=None, - pretrained_model_output_attr=None, - tokenizer_name=None, - # Task Specific - task_config=TMSTaskConfig( - task_name="tms", - feature_probability=0.05, - data_generation_type="at_least_zero_active", - ), - ) - - # Create a pretrained model target_model = TMSModel(config=tms_model_config).to(device) target_model.eval() @@ -101,7 +52,6 @@ def test_tms_decomposition_happy_path() -> None: insert_identity_operations_(target_model, identity_patterns=config.identity_module_patterns) assert isinstance(config.task_config, TMSTaskConfig) - # Create dataset dataset = SparseFeatureDataset( n_features=target_model.config.n_features, feature_probability=config.task_config.feature_probability, @@ -122,7 +72,6 @@ def test_tms_decomposition_happy_path() -> None: if target_model.config.tied_weights: tied_weights = [("linear1", "linear2")] - # Run optimize function optimize( target_model=target_model, config=config, @@ -134,10 +83,6 @@ def test_tms_decomposition_happy_path() -> None: tied_weights=tied_weights, ) - # The test passes if optimize runs without errors - print("TMS SPD optimization completed successfully") - - # Basic assertion to ensure the test ran assert True, "Test completed successfully" @@ -166,7 +111,6 @@ def test_train_tms_happy_path(): model, dataloader = get_model_and_dataloader(config, device) - # Run training train( model, dataloader, @@ -178,8 +122,6 @@ def test_train_tms_happy_path(): log_wandb=False, ) - # The test passes if training runs without errors - print("TMS training completed successfully") assert True, "Test completed successfully"