diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 55925ffd58..ef4d8343f7 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -984,7 +984,7 @@ decode_sampling_nucleus_p: -1 # set if you're doing nucleus / top-p decode_sampling_top_k: 0 # set if you're doing top-k decode_sampling_temperature: 1. -eval_start_step: 0 # start eval after train step is >= eval_start_step +eval_start_step: 0 # start eval when train step is >= eval_start_step eval_interval: -1 # the specific number of train step between eval_step eval_steps: -1 # run this number of steps for eval, recommend setting this to prevent error due to running out of evel data target_eval_loss: 0. # early stop once reaching target eval_loss diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index bf5ac0115e..742cf67ed9 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1587,7 +1587,8 @@ class TrainingLoop(BaseModel): log_period: int = Field(100, description="Frequency (in steps) to log metrics and flush Tensorboard.") eval_start_step: int = Field( 0, - description="Start evaluation after training step is >= eval_start_step.", + ge=0, + description="Start evaluation when training step is >= eval_start_step.", ) eval_interval: int = Field( -1, diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index ee92b93435..db21b9e336 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -718,7 +718,12 @@ def training_loop_iteration( all_host_upload=dump_hlo_upload_all, ) - if eval_interval > 0 and step >= start_step and step >= eval_start_step and (step + 1) % eval_interval == 0: + if ( + eval_interval > 0 + and step >= start_step + and step >= eval_start_step + and (step - eval_start_step) % eval_interval == 0 + ): assert eval_data_iterator # Explicitly reset the eval iterator and counters before starting the eval loop eval_data_iterator.reset() diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index d3afcb612e..9b7862bc5a 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -332,6 +332,15 @@ def test_eval_start_step_config(self): ) self.assertEqual(config_override.eval_start_step, 50) + def test_eval_start_step_negative_raises_error(self): + """Verifies that eval_start_step < 0 raises a validation error.""" + with self.assertRaises((ValueError, Exception)): + pyconfig.initialize( + [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()], + skip_jax_distributed_system=True, + eval_start_step=-1, + ) + if __name__ == "__main__": unittest.main()