fix(base_trainer): gather weights in save_pretrained under zero3

[maxreciprocate]

This PR lets users optionally choose whether to save optimizer state on each and a final checkpoint or only the model weights via trainer.save_pretrained (new default). Also this PR removes two back to back calls to trainer.evaluate, which would happen on the same iter_count in the case total_steps is divisible by config.train.eval_interval. And it fixes trainer.save_pretrained, in which under zero3, weights weren't gathered previously into the state_dict.

https://wandb.ai/sorry/trlx-references/reports/fix-save-pretrained-zero3-v-main--Vmlldzo0MDI1ODk1

[jon-tow]

Thanks for the gather fix and for catching those final step conditions (self.iter_count >= self.total_steps)!

One issue with making save_optimizer = False the default is that you can't resume training ( even though proper checkpointing from say deepspeed ZeRO states is currently WIP). Luckily, we haven't trained on very large datasets (> 10M samples) for this to be much of an issue. Do you foresee this being an issue? I sort of feel like the expected behavior is to save optimizer states.

[maxreciprocate]

Yeah, I also think the excepted behaviour is to save optimizer states, but as you said typical relevant datasets are well below 10M samples and resuming training is only meaningful upon heavy preemptions (at least I haven't had a need to resume training in a while). Cons: training or sweeps would overtake only a fraction of storage (plausibly helpful for end-users with small drives who run many different experiments), but now that coreweave quotas seem to be a bit relaxed (I used to delete .cache on hdd twice a day), I will invert the default as you suggested

[jon-tow]

Thanks, @reciprocated!