[Bugfix] for saving quantized models trained using fsdp

[rahul-tuli]

Bugfix for Saving Quantized Models Trained Using FSDP

This description details a bugfix for an issue encountered when loading quantized models that were trained using Fully Sharded Data Parallel (FSDP).

The issue originated from the process of quantization, during which we updated layer names and saved them using our custom implementation instead of accelerate. This approach resulted in an incorrectly saved state_dict, as each tensor had device information associated with it. Consequently, the model could not be loaded properly using the transformer's AutoModel.from_pretrained(...) method.

The modifications in this update address these complications, ensuring that quantized models trained using FSDP can now be saved and loaded correctly.

Changes

• Asana Ticket
• Relevant code modifications have been made to fix the saved state dicts for quantized models trained with FSDP.

The solution includes a post-processing step where we explicitly iterate through the state_dict and move each tensor to the CPU. The corrected state_dict is then overwritten on the previous, faulty state_dict.

Testing

The saved quantized models can now be loaded by SparseAutoModel.from_pretrained(...). This has been verified manually. The test commands in the ticket work as expected.

Benefits

• This fix improves compatibility when saving quantized models trained using FSDP.