[BUG] DeepSpeed distributed training crashes with "TypeError: can't convert complex to float"

Describe the bug

When doing distributed training of Llama 2 70B, midway through the first training step, DeepSpeed crashes with this stack trace:

gpu-h100-alyssa-2: File "/home/alyssavance/miniforge-pypy3/envs/brr/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 2087, in step
gpu-h100-alyssa-2: self.accelerator.backward(loss)
gpu-h100-alyssa-2: File "/home/alyssavance/miniforge-pypy3/envs/brr/lib/python3.9/site-packages/accelerate/utils/deepspeed.py", line 176, in backward
gpu-h100-alyssa-2: self.optimizer.step()
gpu-h100-alyssa-2: File "/home/alyssavance/miniforge-pypy3/envs/brr/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1994, in _take_model_step
gpu-h100-alyssa-2: scaled_global_grad_norm = get_global_norm(norm_list=norm_groups)
gpu-h100-alyssa-2: File "/home/alyssavance/miniforge-pypy3/envs/brr/lib/python3.9/site-packages/deepspeed/runtime/utils.py", line 307, in get_global_norm
gpu-h100-alyssa-2: scaled_global_grad_norm = get_global_norm(norm_list=norm_groups)
gpu-h100-alyssa-2: File "/home/alyssavance/miniforge-pypy3/envs/brr/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
gpu-h100-alyssa-2: ret_val = func(*args, **kwargs)
gpu-h100-alyssa-2: File "/home/alyssavance/miniforge-pypy3/envs/brr/lib/python3.9/site-packages/deepspeed/runtime/zero/stage3.py", line 1861, in step
gpu-h100-alyssa-2: scaled_global_grad_norm = get_global_norm(norm_list=norm_groups)
gpu-h100-alyssa-2: File "/home/alyssavance/miniforge-pypy3/envs/brr/lib/python3.9/site-packages/deepspeed/runtime/utils.py", line 307, in get_global_norm
gpu-h100-alyssa-2: return sqrt(total_norm)
gpu-h100-alyssa-2: TypeError: can't convert complex to float

This bug only appears when doing distributed training, training on a single machine works as expected.

To Reproduce
Steps to reproduce the behavior:
Fine-tune Llama 2 70B using DeepSpeed 0.10 and HuggingFace Transformers 4.31.0 on multiple machines. (It works fine on a single machine with optimizer offload.)

Expected behavior
The training step should finish without crashing.

ds_report output
(brr) alyssavance@gpu-h100-alyssa:~/brr$ ds_report
[2023-07-25 15:47:49,009] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)

DeepSpeed C++/CUDA extension op report

NOTE: Ops not installed will be just-in-time (JIT) compiled at
runtime if needed. Op compatibility means that your system
meet the required dependencies to JIT install the op.

JIT compiled ops requires ninja
ninja .................. [OKAY]

op name ................ installed .. compatible

async_io ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
[WARNING] sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.0
[WARNING] using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]

DeepSpeed general environment info:
torch install path ............... ['/home/alyssavance/miniforge-pypy3/envs/brr/lib/python3.9/site-packages/torch']
torch version .................... 2.0.1
deepspeed install path ........... ['/home/alyssavance/miniforge-pypy3/envs/brr/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.10.0, unknown, unknown
torch cuda version ............... 11.8
torch hip version ................ None
nvcc version ..................... 12.2
deepspeed wheel compiled w. ...... torch 2.0, cuda 11.8

Screenshots
N/A

System info (please complete the following information):

• OS: Ubuntu 22.04
• GPU count and types: Two machines with x8 H100s each
• Interconnects (if applicable): Connected with 40 Gbps Ethernet
• Python version: 3.9.16

Launcher context
deepspeed launcher

Docker context
No Docker

Additional context
N/A

Start in deepspeed/runtime/utils.py at get_global_norm, then trace its use from deepspeed/runtime/zero/stage3.py and deepspeed/runtime/engine.py. Reproduce the first distributed training step with the reported DeepSpeed 0.10, PyTorch 2.0.1, and multi-machine setup; done means the step completes without the complex-to-float crash.