I’d like to ask for opinions about adding a Trainer configuration option to disable saving of DeepSpeed checkpoints (potentially only keeping the model weights).
Context: I’m finetuning gpt-j-6b for basic translation phrases on consumer hardware (128GB System RAM and Nvidia GPU with 24GB RAM). I use the DeepSpeed Zero optimizer, stages 2 and 3 so 99% of my system memory is fully allocated (I also have a huge swap file on an NVMe).
Arguments for adding this flag:

The DeepSpeed checkpoints are huge (60GB+) and take a long time to save. Because my system RAM is used by the DeepSpeed optimizer, I can’t use local storage (no system ram available for buffers) so I have to transfer over the network to a NAS.
The DeepSpeed save_checkpoint code (as of v0.5.8) is very hungry for system RAM and may even have some memory leaks as after 10+ attempts I’ve never been able to save more than three checkpoints without the linux kernel deciding to OOM the finetuning python process. Out of frustration, I’ve modified the trainer code by commenting out all the calls to deepspeed.save_checkpoint(output_dir) and surprise: I’ve been able to finetune and have saved the model weights 15+ times without getting OOMed.

The default state for this option would obviously be to save everything. I read the previous topic on the subject: Disable checkpointing in Trainer 1 which is nice but does not exactly cover my use case.