Allow specifying arbitrary step numbers for saving extra checkpoints

[haileyschoelkopf]

This PR adds a neox arg extra_save_iters which takes a list of ints. When the training iteration is equal to an element of the list, a checkpoint is saved at that step.

Tested and this works, including for iteration 0.

Note that it'll interfere with keep_n_last_checkpoints (last N checkpoints will still be kept, thus removing any "extra" saved ckpts that may have been desired.) Also, NeoX does in fact save a checkpoint at the last step already even if that last step in training isn't an even multiple of the save interval.

This is intended as a change to supercede the "log-spaced checkpointing" we discussed a while back.

cc @Quentin-Anthony @StellaAthena

[StellaAthena]

Maybe it would make sense to remove the save-iterations argument entirely?

Although arbitrary checkpointing is a good thing to support, in practice I anticipate log-space and linear-space checkpointing to be by far the most desired. If we can expose an easy-to-use interface for those two, with the ability to do arbitrary with some work, that might be ideal?

[Quentin-Anthony]

This looks good to me.

@StellaAthena -- Please also review this since it's a big change.

[StellaAthena]

I have tested this code, discussed the PR at length with @haileyschoelkopf on Discord, and am happy with it being merged.