Error while saving the model under 4bit lora

[visionxyz]

I used

--load_in_4bit
--lora_rank 8
--lora_alpha 16
--lora_dropout 0.05
--zero_stage 2 \

in the training script, then failed to save the model after training. I find it is a problem of load_in_4bit, if it is used,
this part (line 292) in deepspeed.py will be triggered

assert state_dict_keys.issubset(
output_state_dict_keys
), f"mismatch keys {output_state_dict_keys.symmetric_difference(state_dict_keys)}"

I think the training part is okay, but somehow the keys are changed after using both Lora and load in 4-bit.

Thank you so much for providing this tool! It would be appreciated if you could give me some hints about this problem.

[hijkzzz]

it's strange. You could try to comment on this line or debug it.

[visionxyz]

Thank you! I currently comment this line. If I have time, I will debug it and report the reason.