overwrite the max steps

Author: yawenzzzz

olmoearth_pretrain/internal/full_eval_sweep_finetune.py

@@ -250,6 +250,9 @@ def _format_launch_command(
         "--launch.num_gpus=1",
         "--launch.preemptible=True",
         "--launch.task_name=eval",
+        # Overwrite the max duration to enable eval of the last step of the checkpoint
+        "--trainer.max_duration.value=700000",
+        "--trainer.max_duration.unit=steps",
     ]
     parts.extend(checkpoint_args)
     parts.append(f"--trainer.callbacks.wandb.project={project_name}")

scripts/2025_10_02_phase2/script.py

@@ -190,7 +190,7 @@ def build_dataset_config(common: CommonComponents) -> OlmoEarthDatasetConfig:
 
 def build_trainer_config(common: CommonComponents) -> TrainerConfig:
     """Build the trainer config for an experiment."""
-    MAX_DURATION = Duration.steps(700000)
+    MAX_DURATION = Duration.epochs(300)
     METRICS_COLLECT_INTERVAL = 10
     CANCEL_CHECK_INTERVAL = 25
     LOAD_STRATEGY = LoadStrategy.if_available