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?