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.

[aeeeeeep]

The optimization shows measurable benefits only on a specific accelerator (under NDA), where hardware/driver overhead for memory allocation is significantly higher. This extended allocation latency indirectly delays memory release on communication streams.

While this PR's reduction of allocation operations is theoretically sound, the practical benefit appears hardware-dependent and isn't observable on NVIDIA platforms.

[sfc-gh-truwase]

While this PR's reduction of allocation operations is theoretically sound, the practical benefit appears hardware-dependent and isn't observable on NVIDIA platforms.

Thanks for the explanation. Can you confirm there is no regression on NVIDIA platform?

[aeeeeeep]

Confirmed no regression on NVIDIA platform from my tests a few weeks ago (both performance and accuracy).
I’ll re-validate with the latest code tomorrow (~20 h from now) and update code immediately if anything changes.