offline full fp32 weights reconstruction for zero 2+3 checkpoints

[stas00]

This PR adds a magical script that extracts and reconstructs full fp32 weights from either zero 2 or 3 checkpoint. This consolidated state dict can then be used anywhere w/o DeepSpeed.

• I had to change the saving checkpoint code to also add the param names and shapes to the checkpoint optim files
• The script doesn't require live engine, everything is done in completely standalone logic - unfortunately I couldn't fully remove the dependency on deepspeed since the pickled data refers to DeepSpeedEngine and can't be unpickled w/o import deepspeed, but otherwise it's completely standalone.
• When a new checkpoint is saved this script gets copied automatically to the top checkpoint folder so that it's trivial for the user to get the weights out, but of course the API can be used directly - it just will be a really bad idea to do it live, since for large models it'd be very slow and memory consuming to mix it into the normal training. But if someone wants it - it's

deepspeed.save_checkpoint(output_dir)
dist.barrier()
if torch.distributed.get_rank() == 0:
    convert_zero_chkpt_to_fp32_consolid_state_dict(f"{output_dir}/global_step1"), output_file):

I guess the only non-smooth part is figuring out that "global_stepXX" part - probably should expose some of that deepspeed API.

To run the script, from the top level checkpoint dir, is just:

$ cd /path/to/checkpoints_dir
$ ./zero_to_fp32.py global_step1 pytorch_model.bin
Processing zero checkpoint at global_step1
Detected checkpoint of type zero stage 3, world_size: 2
Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)

Future improvements:

• currently the script uses 2x memory of the final checkpoint - eventually it'd be good to find a smart way to avoid that requirement
• currently the script relies on sufficient RAM to do the work - for huge models some kind of memory mapping will need to be used

TODO:

• Where would be a good place to document this?

Lot's of thanks to @tjruwase for the guidance. Once I understood how zero handles its data, it was surprisingly easy to do the rest. Kudos for building an easy to handle system!

p.s. PR with fp16 live consolidator for zero3 addressing #872 is coming soon

@tjruwase

Fixes: #800

[humza909]

@stas00 I think this offline reconstruction is only casting fp16 weights to fp32?

Then how we can use stage3 with fp32 weights. I could not see any method to save 32bit partitioned weights.
This means if a model was originally trained with fp32 converting it to fp16 will cause a precision loss, and if someone wants to avoid this loss there is no way to train and save with stage3.

Please share your thoughts.

Thanks

[stas00]

@stas00 I think this offline reconstruction is only casting fp16 weights to fp32?

no, it's extracting fp32 weight shards and reconstructs full fp32 weights from those. fp16 is not involved anywhere in this process.

[humza909]

@stas00 Thanks for your response.

Two/Three questions regarding stage3:

1. If a model is trained in fp16 setting and consolidated by gather_16bit_weights_on_model_save then how the script zero_to_fp32.py will extract fp32 weights?
2. If I am training with only fp32 enabled then what method will consolidate weights? Do we still use gather_16bit_weights_on_model_save? There is no other config option available. If it is still saving in 16 bit then where are the remaining 16 bit saved? How bits are split? Why the model was originally converted to 16-bit and used for forward and backward pass?
3. Why not save the weights directly to the enabled fp precision in config either fp16 or fp32 without going for additional steps. I know going from fp16 to fp32 will take additional RAM, but if fp32 is enabled save the model in fp32.
4. Below the model parameters are saved in optim_states.pt?

ls -lh global_step100/
total 15G
-rw-rw-r-- 1 snxds snxds  97M Apr  4 08:07 zero_pp_rank_0_mp_rank_00_model_states.pt
-rw-rw-r-- 1 snxds snxds 7.4G Apr  4 08:07 zero_pp_rank_0_mp_rank_00_optim_states.pt
-rw-rw-r-- 1 snxds snxds  97M Apr  4 08:07 zero_pp_rank_1_mp_rank_00_model_states.pt
-rw-rw-r-- 1 snxds snxds 7.4G Apr  4 08:07 zero_pp_rank_1_mp_rank_00_optim_states.pt

Thanks

[stas00]

@stas00 Thanks for your response.

Two/Three questions regarding stage3:

1. If a model is trained in fp16 setting and consolidated by `gather_16bit_weights_on_model_save` then how the script `zero_to_fp32.py` will extract fp32 weights?

If you did save_checkpoint it'll extract fp32 weights.

gather_16bit_weights_on_model_save has nothing to do with the normal Deepspeed save/load sharded weight cycles.

2. If I am training with only fp32 enabled then what method will consolidate weights? Do we still use `gather_16bit_weights_on_model_save`? There is no other config option available. If it is still saving in 16 bit then where are the remaining 16 bit saved? How bits are split? Why the model was originally converted to 16-bit and used for forward and backward pass?

the fp32 weights are saved as parts of save_checkpoint and zero_to_fp32.py extracts those.

you're asking a good question about gather_16bit_weights_on_model_save - I asked for it specifically to deal with half-precision training. You simply don't call this method if you're not doing half-precision. Perhaps if there is a need you can ask for gather_32bit_weights_on_model_save, if you need it.

3. Why not save the weights directly to the enabled fp precision in config either fp16 or fp32 without going for additional steps. I know going from fp16 to fp32 will take additional RAM, but if fp32 is enabled save the model in fp32.

Because it's super-expensive both cpu RAM and time-wise. When you do serious training you must save checkpoints frequently due to hardware crashes, and if this process isn't fast you're wasting gpu idling time. So the helper method was added to aid the beginner user to get things working. Once you have everything else working you don't want to use gathering ever again for any largish model size serious training.

4. Below the model parameters are saved in optim_states.pt?

what is question 4?

[humza909]

@stas00 Thanks for your response.

Ignore question 04. Likely last series of questions.

1. I am still confused why everything in stage3 is in 16-bit? Even I am using full-precision/fp32.

2. Your answer to my first question is save_checkpoint will extract fp32 weights, but how? My model was converted to fp16 before the start of training, and fp32 weights were discarded at the beginning. Is it only the type casting now?

Any recommended links where I can get the implementation details would be very helpful.

Thanks for you time :)

[stas00]

1. I am still confused why everything in stage3 is in 16-bit? Even I am using full-precision/fp32.

It's not. If you're not using half-precision. Nothing is in fp16. Everything is in fp32.

You should never convert your model to any format when you use deepspeed. It'll do it for you on the need basis.

2. Your answer to my first question is save_checkpoint will extract fp32 weights, but how? My model was converted to fp16 before the start of training, and fp32 weights were discarded at the beginning. Is it only the type casting now?

See the previous answer. When not using half precision it'll upcast it into fp32 if you pre-converted to fp16.

To give you a hint deepspeed performs the same as mixed precision protocol, but in reverse - it converts the model to fp16 and it keeps a copy of fp32 weights in its optimizer. So you always have both weights. save_checkpoint saves only fp32 weights.

So if you're not doing mixed precision, you only have fp32 weights and they are saved with save_checkpoint.

Any recommended links where I can get the implementation details would be very helpful.

I think for this particular topic there isn't much documentation, you can read the papers here:
https://huggingface.co/docs/transformers/main/main_classes/deepspeed#main-deepspeed-resources
the rest of the doc might be useful as well as a general reading, but it's specific to the transformers integration.

You can make a request to document these nuances. These are all excellent questions and are important to understand.

[khulaifi95]

@stas00 Sorry but I find this issue related to my problem here. To make an extension on the last question from @humza909, for zero stage-3 training across 2 nodes, is it ever possible to reconstrcut the model weights from the contents that saved in checkpoints from one node?

[zcakzhuu]

1. I am still confused why everything in stage3 is in 16-bit? Even I am using full-precision/fp32.

It's not. If you're not using half-precision. Nothing is in fp16. Everything is in fp32.

You should never convert your model to any format when you use deepspeed. It'll do it for you on the need basis.

2. Your answer to my first question is save_checkpoint will extract fp32 weights, but how? My model was converted to fp16 before the start of training, and fp32 weights were discarded at the beginning. Is it only the type casting now?

See the previous answer. When not using half precision it'll upcast it into fp32 if you pre-converted to fp16.

To give you a hint deepspeed performs the same as mixed precision protocol, but in reverse - it converts the model to fp16 and it keeps a copy of fp32 weights in its optimizer. So you always have both weights. save_checkpoint saves only fp32 weights.

So if you're not doing mixed precision, you only have fp32 weights and they are saved with save_checkpoint.

Any recommended links where I can get the implementation details would be very helpful.

I think for this particular topic there isn't much documentation, you can read the papers here: https://huggingface.co/docs/transformers/main/main_classes/deepspeed#main-deepspeed-resources the rest of the doc might be useful as well as a general reading, but it's specific to the transformers integration.

You can make a request to document these nuances. These are all excellent questions and are important to understand.

Hi, thanks for the great answers! They are really helpful. I have a question about fine-tuning in fp32 using peft. Regarding question 2, if I am checkpointing a model that is fine-tuned using peft in fp32 with a transformers trainer, do I have to run zero_to_fp32.py before resuming training from checkpoint using the trainer? Aren't the weights in .bin file fp32 already?

[stas00]

@zcakzhuu,

You don't need to do anything if you resume from a deepspeed checkpoint saved at the first run - it will do the right thing.

As long as you remain within the deepspeed realm you don't need to do anything.

Only when you're done training and you want to take the weights elsewhere you'd use zero_to_fp32.py if you were training in mixed precision or if you were training in fp32 but saving only the deepspeed checkpoint.

Just give it a try - train for 1 step, save and resume and train for another step. It shouldn't be too difficult to see what you get.

[stas00]

@khulaifi95, I didn't understand your question, could you please try to reframe it?