Add single parameter allgather optimization for zero3

[aeeeeeep]

• Perform allgather operations on individual parameters instead of parameter lists.
• Significantly reduce peak memory usage in high memory pressure scenarios.
• Improve performance by minimizing temporary buffer requirements.
• The behavior is enabled via a new boolean flag under the section

"zero_optimization": {
  "stage3_allgather_single_param": true 
 }

• By default the optimization is not enabled.

[sfc-gh-truwase]

@aeeeeeep thanks for this contribution. Are you able to share some data showing the benefits of this optimization?

[aeeeeeep]

@aeeeeeep thanks for this contribution. Are you able to share some data showing the benefits of this optimization?

Thanks for your feedback! I’ll share detailed data within the next few days.