Trainer only saves model in FP16 when using mixed precision together with DeepSpeed
Author: andstor
Pull request: no

Normally when saving a model using the Trainer class, the `dtype` of the saved model is (and should be) the same as the original model. This is also true when using mixed precision, and when using DeepSpeed. However, when using mixed precision **together** with DeepSpeed, the output is in float16 no matter the model input `dtype`.

The Trainer class has custom handling for DeepSpeed, depending on the ZeRO stage:
https://github.com/huggingface/transformers/blob/5f9685576149fb45a61d0dcec9a260930df0a49a/src/transformers/trainer.py#L2914-L2928

as well does accelerate:
https://github.com/huggingface/accelerate/blob/06b138d84537ffb2d1d404f2f198a0446e8d7ec3/src/accelerate/accelerator.py#L3042-L3056

For ZeRO stage <=2 DeepSpeed holds the model weights in the `state_dict`. Using mixed precision training, these are always in float16. Using full precision training, they are the same dtype as the original model.

For ZeRO stage 3 the `state_dict` contains just placeholders since the model weights are partitioned. By setting `stage3_gather_16bit_weights_on_model_save=true`, DeepSpeed consolidates the weights. When training using mixed precision, float16 is always produced. When training in full precision, despite the name, it follows the dtype of the original model. If `stage3_gather_16bit_weights_on_model_save=false`, Trainer saves a full checkpoint instead, and the DeepSpeed `zero_to_fp32.py` script can be used to recover weights in float32.


Currently, the only way to save a model that is trained using the Trainer class that applies mixed precision along with DeepSpeed ZeRO stage <=2 in float32, is to manually save a checkpoint and then use some weight recovery method afterwards. Is this due to a limitation of the DeepSpeed API, or could this be handled in the Trainer class (or preferably in Accelerate)? At least, maybe a flag could be available to either save the float16 weights or a checkpoint at the end of training (kind of how stage 3 with `stage3_gather_16bit_weights_on_model_save=true` is handled)?


### Who can help
@pacman100, @muellerzr

Comments:
Hello @andstor, 

The model is saved in the selected half-precision when using mixed-precision training, i.e., `fp16` if mixed-precision is using `fp16` else `bf16` if mixed-precision is using `bf16`. Earlier there was logic to convert the state_dict to FP32 but users complained about the increase in the ckpt size and hence the current logic. Is this a limitation for your use case

I would expect that if I use a FP32 model with mixed-precision, the “master weights” in the optimizer would be extracted on save. If not, the saving method is lossy, as it discards valuable precision. This lossy saving seems to only happen when using DeepSpeed with mixed-precision. If users have a problem with the size, they should simply ensure the model is loaded in half-precision before starting training.

Hi, 
I am also have a problem with that even if I not using Trainer the saving strategy with accelerator when using mixed precision with deepspeed is always with mixed precision dtype fp16 or bf16, can I change the behavior to save it with fp32?

` if training_args.output_dir is not None:
        accelerator.wait_for_everyone()
        unwrapped_model = accelerator.unwrap_model(model)
        unwrapped_model.save_pretrained(
            training_args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save,
            state_dict=accelerator.get_state_dict(model)`

I found this option to keep_fp32_wrapper within the unwrap_model but I am not sure if this what should be used 
This is the code I am using 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi, has there been any progress or solution on that subject?