Z3: optimizations for grad norm calculation and gradient clipping

[nelyahu]

This PR add the below functionality:

1. complete_grad_norm_calculation_for_cpu_offload: move total_norm to CPU, as expected device in such case is CPU..
2. repalce get_global_norm() with torch.linalg.norm for better performance.
3. unscale_and_clip_grads: replace clipping based on if statement to use torch.clamp for better performance.

change (3) is taken from #5547 (which was closed)

[jomayeri]

Changing this line has been associated with several bugs #5422, #5538

[loadams]

Changing this line has been associated with several bugs #5422, #5538

@nelyahu - thoughts on this comment, seems last time this line was modified users ran into issues?

[nelyahu]

Changing this line has been associated with several bugs #5422, #5538

@nelyahu - thoughts on this comment, seems last time this line was modified users ran into issues?

@loadams, Yes - this optimization was already pushed and reverted due to ds-chat (failures in cpu-offload configurations).
i did offline debugging of those failure and improved the code change so it will pass. Since then ds-chat tests where added to DeepSpeed repo CI and it is now passing.
Are there any other tests (full model training for example), that does not exists in the CI, which can be manually ran?

[tjruwase]

i did offline debugging of those failure and improved the code change so it will pass

@nelyahu, it great that you narrowed this down. Do you think a unit test can be added for this case?

[loadams]

i did offline debugging of those failure and improved the code change so it will pass

@nelyahu, it great that you narrowed this down. Do you think a unit test can be added for this case?

@nelyahu - we've stabilized the CI, thoughts on adding this test?

[nelyahu]

@loadams Oh, Sorry- I missed the last comment.
Sure, yes we can add such UT that will cover it.
But i cannot address it immediately. I will update this PR once we have a unit test.

[nelyahu]

@loadams / @tjruwase as request i make sure the regression was discussed here will be covered by a unit tests. I used TestZeroPartialOffloadConfigSweep, and added it gradient_clipping so it will go though the problematic flow.
i reproduced the issue using the UT, and made sure it is fixed.

[nelyahu]

@loadams can you re-run workflows? i suspect that the failure are not related to this PR