Can't access gradients if using ZERO

Hi,

I was attempting to do adaptive gradient clipping using deepspeed. This requires access to the parameters and their respective gradients (p.grad, etc.).

However, it seems like the ZERO optimizers zero out the gradients after calling .backward(), which means the gradients are not accessible from the parameters, making custom gradient clipping code not possible.

Is removing the gradients from the parameters in .backward() necessary? Would it be feasible to provide a straightforward way to access the gradients and their respective parameters for the user?

Start by tracing how the ZERO optimizers handle gradients after .backward(), especially whether p.grad remains available. Determine a supported way for users to access gradients and their parameters for adaptive clipping; done requires a defined approach and validation that custom clipping can use it.