fix tests II

mosaicml · Jun 6, 2024 · cf1d31f · cf1d31f
1 parent cafb168
commit cf1d31f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
@@ -76,9 +76,6 @@ class EvalConfig:
     run_name: Optional[str] = None
     metadata: Optional[Dict[str, str]] = None
 
-    # Dataloader
-    global_train_batch_size: Optional[int] = None
-
     # Distributed parameters
     dist_timeout: Union[float, int] = 600.0
     fsdp_config: Optional[Dict[str, Any]] = None

diff --git a/tests/a_scripts/eval/test_eval.py b/tests/a_scripts/eval/test_eval.py
@@ -13,7 +13,7 @@
 
 from llmfoundry.utils import build_tokenizer
 from llmfoundry.utils.builders import build_composer_model
-from llmfoundry.utils.config_utils import to_dict_container
+from llmfoundry.utils.config_utils import EVAL_CONFIG_KEYS, to_dict_container
 from scripts.eval.eval import main  # noqa: E402
 from tests.data_utils import create_c4_dataset_xxsmall, gpt_tiny_cfg
 
@@ -134,6 +134,14 @@ def test_loader_eval(
     test_cfg.eval_interval = '1ba'
     test_cfg.loggers = om.DictConfig({'inmemory': om.DictConfig({})})
 
+    # This test uses a training yaml with training-only keys present.
+    # We exclude these keys before calling `main` from the eval script.
+    allowed_keys = EVAL_CONFIG_KEYS
+    present_keys = set(test_cfg.keys())
+    keys_to_pop = present_keys.difference(allowed_keys)
+
+    [test_cfg.pop(key) for key in keys_to_pop]
+
     trainers, eval_gauntlet_df = main(test_cfg)
 
     assert eval_gauntlet_df is None