(model_name: str, seed: int)
| 18 | |
| 19 | @pytest.mark.parametrize("model_name,seed", [("mosaic_bert", 17), ("hf_bert", 18), ("flex_bert", 42)]) |
| 20 | def test_trainer(model_name: str, seed: int): |
| 21 | with open("yamls/defaults.yaml") as f: |
| 22 | default_cfg = OmegaConf.load(f) |
| 23 | with open(f"yamls/models/{model_name}.yaml") as f: |
| 24 | model_cfg = OmegaConf.load(f) |
| 25 | with open("tests/smoketest_config_main.yaml") as f: |
| 26 | test_config = OmegaConf.load(f) |
| 27 | config = OmegaConf.merge(default_cfg, model_cfg, test_config) |
| 28 | assert isinstance(config, DictConfig) |
| 29 | config.model.name = model_name |
| 30 | config.seed = seed |
| 31 | |
| 32 | with SynthTextDirectory() as tmp_datadir: |
| 33 | config.train_loader.dataset.remote = tmp_datadir |
| 34 | config.train_loader.dataset.local = os.path.join(tmp_datadir, "tr-local1") |
| 35 | config.eval_loader.dataset.remote = tmp_datadir |
| 36 | config.eval_loader.dataset.local = os.path.join(tmp_datadir, "ev-local1") |
| 37 | # Also save checkpoints in the temporary directory |
| 38 | config.save_folder = tmp_datadir |
| 39 | |
| 40 | # Train |
| 41 | trainer1 = main(config, return_trainer=True) |
| 42 | assert trainer1 is not None |
| 43 | model1 = trainer1.state.model.model |
| 44 | |
| 45 | # Check that the checkpoint was saved |
| 46 | chkpt_path = os.path.join(tmp_datadir, "latest-rank0.pt") |
| 47 | assert os.path.isfile(chkpt_path), f"{os.listdir(tmp_datadir)}" |
| 48 | |
| 49 | # Check that the checkpoint was loaded by comparing model weights (with no weight changes) |
| 50 | config.load_path = chkpt_path |
| 51 | config.seed += 10 # change seed |
| 52 | config.train_loader.dataset.local = os.path.join(tmp_datadir, "tr-local2") |
| 53 | config.eval_loader.dataset.local = os.path.join(tmp_datadir, "ev-local2") |
| 54 | trainer2 = main(config, return_trainer=True, do_train=False) |
| 55 | assert trainer2 is not None |
| 56 | model2 = trainer2.state.model.model |
| 57 | |
| 58 | for param1, param2 in zip(model1.parameters(), model2.parameters()): |
| 59 | torch.testing.assert_close(param1, param2) |
nothing calls this directly
no test coverage detected