Uses bf16 training, there is an abnormal loss[BUG]

Describe the bug
I use the neox framework to load the weight of llama to continue training. When I use bf16, the loss becomes very large after the weight is updated for the first time in the model. Because neox also uses deepspeed, so the reason for deepspeed is suspected. Has anyone been able to successfully load llama weights with deepspeed, to continue training with bf16

steps: 1 loss: 1.9300 iter time (s): 13.338 samples/sec: 1.200
 samples/sec: 1.195 | iteration        1/    2140 | elapsed time per iteration (ms): 13386.0 | learning rate: 1.417E-07 | approx flops per GPU: 18.8TFLOPS | lm_loss: 1.930013E+00 | number of skipped iterations:   0 | number of nan iterations:   0 |
after 1 iterations memory (MB) | allocated: 26732.29150390625 | max allocated: 30516.369140625 | reserved: 33468.0 | max reserved: 33468.0
time (ms)
[2023-06-12 18:08:09,410] [INFO] [logging.py:96:log_dist] [Rank 0] step=2, skipped=0, lr=[2.834890965732087e-07, 2.834890965732087e-07], mom=[[0.9, 0.95], [0.9, 0.95]]
steps: 2 loss: 11.4175 iter time (s): 2.441 samples/sec: 6.554
 samples/sec: 6.551 | iteration        2/    2140 | elapsed time per iteration (ms): 2442.5 | learning rate: 2.835E-07 | approx flops per GPU: 103.2TFLOPS | lm_loss: 1.141747E+01 | number of skipped iterations:   0 | number of nan iterations:   0 |
time (ms)
[2023-06-12 18:08:11,797] [INFO] [logging.py:96:log_dist] [Rank 0] step=3, skipped=0, lr=[4.2523364485981307e-07, 4.2523364485981307e-07], mom=[[0.9, 0.95], [0.9, 0.95]]
steps: 3 loss: 11.4233 iter time (s): 2.386 samples/sec: 6.706
 samples/sec: 6.703 | iteration        3/    2140 | elapsed time per iteration (ms): 2386.9 | learning rate: 4.252E-07 | approx flops per GPU: 105.6TFLOPS | lm_loss: 1.142331E+01 | number of skipped iterations:   0 | number of nan iterations:   0 |
time (ms)
[