Issue with Zero3 Mode and State Dictionary Saving - Related to Issue 1271

Nobody has claimed this yet.

strong related to https://github.com/lm-sys/FastChat/issues/1271

Dear Fastchat team,

I hope this message finds you well. I am writing to report an ongoing issue related to the zero3 mode and state dictionary saving in our project. This problem is closely related to the previously closed issue with the identifier https://github.com/lm-sys/FastChat/issues/1271.

In our current implementation, we are utilizing the `trainer.hf_deepspeed_config_orig.is_zero3()` function from the deepspeed config object to determine if our train script is operating in zero3 mode. Additionally, we have discovered an internal function within the deepspeed engine object called `_zero3_consolidated_16bit_state_dict()`. By leveraging this function, we are able to successfully gather the state_dict, which has resolved the issue with zero3 saving. Consequently, we now obtain a .bin file of the expected size and are able to successfully invoke the apply_lora.py function. But we are not 100% sure if we really save the correct lora weight.
here is whole code

Before proceeding further, we kindly request your expertise and assistance in validating the correctness of our approach. Specifically, we would appreciate a thorough review of our utilization of the internal function from the deepspeed engine. As mentioned earlier, our implementation addresses the same problem that was previously raised in https://github.com/lm-sys/FastChat/issues/1271, but remained unresolved despite the release of a fix code by the original author.

To facilitate the review process and maintain the link with the original issue, we kindly ask you to open a new issue, clearly indicating the connection to https://github.com/lm-sys/FastChat/issues/1271. This will allow the community to assess our proposed fix and confirm its effectiveness in resolving the zero3 mode and state dictionary saving problem.

We appreciate your attention to this matter and look forward to your guidance and insights. If you require any additional information or code snippets to support the review process, please do not hesitate to let us know.

To facilitate the review process and maintain the connection with the original issue, we have provided the relevant code snippet below:
````python
# check if zero3 mode enabled
if trainer.hf_deepspeed_config_orig.is_zero3():
# use deepspeed engine internal function to gather state dict
# state_dict_zero3 contains whole parameters of base and lora adapters
# we will not extract lora parameters since peft save_pretrained will do that
# https://github.com/huggingface/peft/blob/3714aa2fff158fdfa637b2b65952580801d890b2/src/peft/peft_model.py#L125
# https://github.com/huggingface/peft/blob/3714aa2fff158fdfa637b2b65952580801d890b2/src/peft/utils/save_and_load.py#L19
state_dict_zero3 = trainer.model_wrapped._zero3_consolidated_16bit_state_dict()
if training_args.local_rank == 0:
state_dict = state_dict_zero3
else:
# in other mode we use original code from fastchat team, to make sure our change is minimum
state_dict = get_peft_state_maybe_zero_3(
model.named_parameters(), lora_args.lora_bias
)
````

Thank you for your time and support.

Best regards,

Start with related issue 1271 and the training script's Zero3 state-dictionary saving path. Review the use of trainer.hf_deepspeed_config_orig.is_zero3() and trainer.model_wrapped._zero3_consolidated_16bit_state_dict(), along with the referenced PEFT save logic. Done means confirming whether the resulting state dictionary contains the correct LoRA weights and documenting or resolving the behavior.

Written by the indexing model from the issue text.