From d2f0052584c21bc2c84939c9e344d949f30a2a3e Mon Sep 17 00:00:00 2001 From: Zhengxian He Date: Mon, 20 Jul 2026 22:24:00 +0000 Subject: [PATCH] Add eval_start_step flag to configure initial evaluation delay --- src/maxtext/configs/base.yml | 1 + src/maxtext/configs/types.py | 4 ++++ src/maxtext/trainers/pre_train/train.py | 4 +++- tests/unit/pyconfig_test.py | 15 +++++++++++++++ 4 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index fe49545bdb..72a03b5d80 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -983,6 +983,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_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 22a2a5ff0b..a8f4d0efb5 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1573,6 +1573,10 @@ class TrainingLoop(BaseModel): description="Total number of training steps. -1 defaults to learning_rate_schedule_steps.", ) 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.", + ) eval_interval: int = Field( -1, description="Run evaluation every N training steps. -1 disables interval-based evaluation.", diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 39d3fec6b5..00bb572ccc 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -669,6 +669,7 @@ def training_loop_iteration( eval_interval = immutable_data["eval_interval"] eval_steps = immutable_data["eval_steps"] start_step = immutable_data["start_step"] + eval_start_step = immutable_data["eval_start_step"] # HLO dump config dump_hlo = immutable_data["dump_hlo"] @@ -712,7 +713,7 @@ def training_loop_iteration( all_host_upload=dump_hlo_upload_all, ) - if eval_interval > 0 and step > start_step and (step + 1) % eval_interval == 0: + if eval_interval > 0 and step >= start_step and step >= eval_start_step and (step + 1) % eval_interval == 0: assert eval_data_iterator # Explicitly reset the eval iterator and counters before starting the eval loop eval_data_iterator.reset() @@ -863,6 +864,7 @@ def train_loop(config, recorder, state=None): "steps": config.steps, "eval_interval": config.eval_interval, "eval_steps": config.eval_steps, + "eval_start_step": config.eval_start_step, "save_checkpoint_on_completion": config.save_checkpoint_on_completion, "start_step": start_step, "dump_hlo": config.dump_hlo, diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index dcfb1ec242..d3afcb612e 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -317,6 +317,21 @@ def test_local_sa_flags_explicit_override(self): self.assertEqual(config.local_sa_v_layout, "SEQ_MINOR") self.assertFalse(config.local_use_splash_scheduler) + def test_eval_start_step_config(self): + """Verifies that eval_start_step defaults to 0 and can be overridden via pyconfig.""" + config_default = pyconfig.initialize( + [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()], + skip_jax_distributed_system=True, + ) + self.assertEqual(config_default.eval_start_step, 0) + + config_override = pyconfig.initialize( + [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()], + skip_jax_distributed_system=True, + eval_start_step=50, + ) + self.assertEqual(config_override.eval_start_step, 50) + if __name__ == "__main__": unittest.main()