Support Early Stopping Training Runs

[slabhs-aws]

Hey,
Presently, changing job_config.training.steps affects the learning_rate. Can we add an early_stop_step arg to job_config.training?

[tianyu-l]

Could you share more about the use case, before reaching conclusion on the solution?

[wwwjn]

I think schedule_step should by default = training.steps, but in case of a debug run, training.steps could be smaller, but schedule_step should be the original value as of normal run.

That's what I'm trying to say.

[wwwjn]

Let's explicitly set schedule_step with the same value as train.steps in toml file

I'd prefer not, because you'd have to always change both.

ok, then just making sure the schedule_step (if there is one) in toml files always rewrite the default value is good enough

[rakkit]

FYI, in our local branch, we have smth named "reconfiguration LR" (need to be set explictly in train configs), the context is at some runs we want to take a checkpoint and overwrite the rest of lr scheduler (but no changes of past). in lr-scheduler-load-state and checkpoint and optimizer

[rakkit]

good thing of allow "reconfig" rest of LR scheduler is to make optimizer related ablation and any of data mixing ablation on-fly much easier. (for example smolm3 data ablation)

[tianyu-l]

@rakkit
We have this checkpoint config exclude_from_loading https://github.com/pytorch/torchtitan/blob/main/torchtitan/config/job_config.py#L611

Could you achieve the same by putting LR scheduler in it and use a new LR config?