Allow modification of zero partitioned parameters

[tjruwase]

Utilities for flexible modification of partitioned fp32 parameters and optimizer states.

Fix #3830

[tohtana]

This is a great feature! The code modifications and the document are also clear. I do have one observation, though it's not immediately pressing:

Currently we have three get_* functions (safe_get_full_fp32_param, safe_get_full_grad, and safe_get_full_optimizer_state). This PR introduces safe_set_full_fp32_param and safe_set_full_optimizer_state. Is there a specific reason we're omitting safe_set_full_grad?
Maintaining consistency in the APIs can help users understand the design better.

[tjruwase]

This is a great feature! The code modifications and the document are also clear. I do have one observation, though it's not immediately pressing:

Currently we have three get_* functions (safe_get_full_fp32_param, safe_get_full_grad, and safe_get_full_optimizer_state). This PR introduces safe_set_full_fp32_param and safe_set_full_optimizer_state. Is there a specific reason we're omitting safe_set_full_grad? Maintaining consistency in the APIs can help users understand the design better.

@tohtana, thanks for this valid question. I am delaying support for safe_set_full_grad until there is explicit request for it because it is harder to implement and I have limited bandwidth to think through all the design issues :(. I will add a TODO for this. Thanks for the review.

[mayank31398]

@tjruwase
can we get the method safe_set_gradients?
This is required for using Megatron's sequence parallel instead of Ulysees.
Wanted to check if this is possible anyhow and would be super-useful.

[mayank31398]

Essentialy whats needed:

grad = get_grads(layernorm.weight)
dist.all_reduce(grad, group=tp_group)
safe_set_grads(grad, layernorm.weight)

if there is an alternative way to do it, that will also be helpful.