safe_get_full_grad & safe_set_full_grad

[ProjectDisR]

deepspeed 0.15.3
zero 3 is used

For "safe_get_full_grad", does it return the same gradient values on each process/rank?

As for "safe_set_full_grad", should it be called on all the processes/ranks? or just one of them is enough?
If it's the former one, users will need to ensure gradient values to be set on each process/rank are the same?

Also, which float type should be used for "safe_set_full_grad"? any way to check this?

[tjruwase]

For "safe_get_full_grad", does it return the same gradient values on each process/rank?

Yes. This API is called after reduction and so all DP ranks have the same gradient value: https://deepspeed.readthedocs.io/en/latest/zero3.html#debugging

As for "safe_set_full_grad", should it be called on all the processes/ranks? or just one of them is enough?

These APIs should be called by all the DP ranks to avoid hangs.

If it's the former one, users will need to ensure gradient values to be set on each process/rank are the same?

Yes, users have freedom/responsibility for correct usage of these APIs. Normally gradient values are the same across DP ranks after reduction.

Also, which float type should be used for "safe_set_full_grad"? any way to check this?

The type should match the gradient accumulation data type configured by the user: https://deepspeed.readthedocs.io/en/latest/zero3.html#deepspeed.utils.safe_set_full_grad