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