fix: ensure last checkpoint is always saved, refactor training stop conditions to be computed in single location

[mattmazzola]

Issues

1 Inconsistent checkpoint filenames saved by trainer

In our pipeline we often have sequence of steps such as (train, reshard/unflatten, evaluate). The output files of the training become inputs to the resharding scripts. In order for the execution to work reliably the output files need to have consistent filenames, such as checkpoint_last-model_part-0-shard0.pt

When running metaseq.cli.train with tasks such as streaming_finetune_language_modeling there are two different stopping conditions set by --max-epochs and --max-updates. Whichever limit is hit first will cause the model stop training.

The issue is that checkpoint_last-* file is ONLY written the epoch stop condition or update stop conditions were false.
This couples the checkpoint filename with the stopping conditions

Notice checkpoints[0] only uses the FIRST true filename/condition

metaseq/metaseq/checkpoint_utils.py

Lines 89 to 99 in c16d210

	if len(checkpoints) > 0:
	if PathManager.islink(checkpoints[0]):
	PathManager.rm(checkpoints[0])
	trainer.save_checkpoint(
	checkpoints[0],
	extra_state,
	training_finished=training_finished,
	async_callback_fn=async_callback_fn if save_to_NFS else None,
	files_to_symlink_to=checkpoints[1:] if len(checkpoints) > 1 else None,
	)

Goal

We want to be able to run the jobs/pipeline and change the stopping conditions without implicitly changing the output file that will be given to the subsequent commands / scripts

2 Training Stop was Handled in Multiple Locations

Loop condition:

metaseq/metaseq/cli/train.py

Line 209 in c16d210

while epoch_itr.next_epoch_idx <= max_epoch:

Loop break:

metaseq/metaseq/cli/train.py

Lines 212 to 213 in c16d210

	if should_stop:
	break

This makes it harder to reason about which condition will cause training to stop.

Solution

• Consolidate all training stop conditions in to validate_and_save and should_stop
  • Use > instead of >= conditions
• Change to always save checkpoint_last* file
  • This means there are cases it will save multiple checkpoints
    • Epoch AND last checkpoint
    • Updates AND last checkpoint

Testing

I wasn't able to test since this is merging with metaseq main instead of our fork's main.
I wanted to at least share the ideas. Although changing training stop conditions can be serious, so maybe someone else can submit a small jobs to test. One with max-epochs, other with max-updates, and in both cases it saves checkpoint_last files

Related to #726

[mattmazzola]

Usually it is practice to avoid while True anywhere since it relies on other code to stop the loop and it's easy to make mistakes. However, the alternative of splitting logic between loop and validate function is more complex and thus more likely for us to have issues in future.

Also, another reason the while true is not be so bad is because the original could potentially be the same condition when max-epochs was not provided / defined.

If not defined

while epoch_itr.next_epoch_idx <= max_epoch
while epoch_itr.next_epoch_idx <= cfg.optimization.max_epoch or math.inf
while epoch_itr.next_epoch_idx <= math.inf
while true