Trainer option to disable saving DeepSpeed checkpoints
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 Trainerwhich is nice but does not exactly cover my use case.
2022-01-10T15:41:08Z
cc@staswho is the most knowledgeable on the DeepSpeed integration.
2022-01-10T17:00:31Z
One can’t not save intermediary checkpoints, since if you don’t how will you restart the training/finetuning w/o the optimizer states?In fact it’s thestage3_gather_fp16_weights_on_model_save=truethat uses additional CPU RAM to gather the weights from multi-gpus that would consume the most RAM, so you can disable it. Especially since if I understand your setup you’re using only one gpu.The most efficient process should be to setstage3_gather_fp16_weights_on_model_save=false, save the weights as is which should not take any additional memory as it’s just a gpu → disk, and then at the end usingzero_to_fp32.pyto extract the full fp32 weights.  Specifics are here:DeepSpeed IntegrationPlease let me know if applying this suggestion didn’t incur the additional CPU RAM use. And if you have some doc improvements to suggest I’m all ears.The DeepSpeed checkpoints are huge (60GB+)The checkpoint size isparams*18so your model must be around 3B params. It should be onlyparams*2larger than normal non-DS checkpoint since it saves both fp16 and fp32 weights, but otherwise DS checkpoints aren’t much bigger than non-DS checkpoint…
2022-01-10T18:34:57Z
stas:One can’t not save intermediary checkpoints, since if you don’t how will you restart the training/finetuning w/o the optimizer states?How bad is it if I shuffle the training and evaluation datasets and restart the process by loadingpytorch_model.bin(generated from the previously stopped finetuning process). Did I mention that I’m a noob here?Especially since if I understand your setup you’re using only one gpu.Correct, one lonely GPU.stas:In fact it’s thestage3_gather_fp16_weights_on_model_save=truethat uses additional CPU RAM to gather the weights from multi-gpus that would consume the most RAM, so you can disable it.The docs and the DeepSpeed code say that setting stage3_gather_fp16_weights_on_model_save to False will causepytorch_model.binto not be generated.stas:The most efficient process should be to setstage3_gather_fp16_weights_on_model_save=false, save the weights as is which should not take any additional memory as it’s just a gpu → disk, and then at the end usingzero_to_fp32.pyto extract the full fp32 weights.  Specifics are here:DeepSpeed IntegrationPlease let me know if applying this suggestion didn’t incur the additional CPU RAM use. And if you have some doc improvements to suggest I’m all ears.Let me see if I understand your proposal correctly: if I modify the DS code by adding a special case that saves the fp16 weights when thenumber of gpus == 1andstage3_gather_fp16_weights_on_model_save=falsethat would result in the most efficient save process for intermediate weights (for 1 gpu)?
2022-01-10T19:27:28Z
Well, if you start from weights-only you waste training resources as your optimizer will take time to get to a point where you stopped at. So when you resume training you typically want to resume the optimizer states and not start from scratch. That’s the whole point of saving intermediary checkpoints. Shuffling data shouldn’t make any difference to wanting ongoing optim states.You don’t need to change any DS code, you just need to set:stage3_gather_fp16_weights_on_model_save=falsein ds_config file and it won’t gather and save the fp16 weights. And you can then extract the perfect fp32 weights at the end of your training usingzero_to_fp32.pyscript.Of course do a very short training and try usingzero_to_fp32.pyscript and ensure that you know that this is what you want.So the proposal is this:don’t gather any zero3 weights and use additional CPU memory to build a state_dict and just dump the intermediary states to disk really fast (should be ~0 CPU RAM overhead)resume training from checkpoint until you have finished your training - i.e. repeat this step as many times as you need toextract the final fp32 weights withzero_to_fp32.pyIf something is unclear please don’t hesitate to ask for further clarifications,@mihai
2022-01-10T21:30:28Z
@stasmany thanks for the advice!I understand what you are saying. Essentially, the following statement is key for my issue:just dump the intermediary states to disk really fastFor a model like gpt-j-6b, that’s 68G for every checkpoint (fp16). I see now that my initial approach: move that data over the network was naive, a much better idea is to keep that on very fast local storage (NVMe) and use a reasonablesave_total_limit(like 10).Again, thank you! I’ll mark this thread as resolved.
2022-01-10T21:57:33Z
Yes, Deepspeed’s workflow is designed to be very fast and so yes ideally you want to be saving each node’s data locally and have the system deal with syncronization after the training moved on with its work.For example at the current 104B GPT training at JeanZay/BigScience we save 1.4TB checkpoints which takes about 60secs. This is with 256 gpus and 64 nodes, and each process only writes its data locally to an SSD drive. The system itself uses GPFS over SSD discs to synchronize, but this is transparent to the training software.And your use-case is much simpler since you don’t have multi-node. So yes, a local NVMe storage for intermediary checkpoints would be ideal.But again practice with a tiny training (like 5 steps) and make sure you know how to extract the final weights so that you know it works for your needs.
2023-05-20T07:18:10Z
I’m confused here. I’m currently training model using DeepSpeed stage3 with offloading param to nvme. The default value ofstage3_gather_16bit_weights_on_model_saveisFalse. But, I can not dodeepspeed.save_checkpointbecause of this errorZeRO-3 does not yet support checkpointing with NVMe offloading, please disable for now.But, Trainer obviously usingdeepspeed.save_checkpoint. So, how to avoid the error?
2023-05-23T17:18:32Z
The error comes from Deepspeed - therefore please file an issue atIssues · microsoft/DeepSpeed · GitHub