[FEATURE] Reduce the required peak RAM on a single node while converting weights

[zhanyuanucb]

System information

• Alpa version: v0.2.2
• Are you willing to contribute to it (Yes/No): Yes, but not immediately

Describe the new feature and the current behavior/state
Referring to here, for now, weight conversion for OPT-175B requires a peak RAM usage as large as twice of the model size. It will be great to do this in a distributed way to reduce the required peak RAM on a single node.

Will this change the current API? How?
Changes will mostly happen in the step_2_consolidate_992_shards_to_singleton.py

Describe alternatives you've considered

Additional context

[zhisbug]

This is non-trivial to do. I discussed with @merrymercy and he will update on this issue.

[merrymercy]

@Ying1123

[sammeralomair]

Working on this.
I have a list of questions if anyone is available to disuss over slack