[BUG] loss discrepancy among ZeRO-0, 1, 2, 3, when gradient accumulate multiple steps

Describe the bug
Feeding model with same random data, but with different ZeRO optimization stages produces different loss trajectories.

To Reproduce
Steps to reproduce the behavior:

1. Download the testing scripts: https://gist.github.com/zarzen/34a34c75109dfc39b37a8d2e27db20f1
2. Run testing script: python3 test_diff_stages.py
3. The script will save the loss trajectories into /tmp folder. And it outputs the loss differences

Example log output, at the end:

stage0, stage1, losses all close False
loss diff 0.02541184425354004:: stage 0: step2.0, loss 3.7346994876861572 stage1: step2.0, loss 3.709287643432617
loss diff 0.012979745864868164:: stage 0: step4.0, loss 2.4462392330169678 stage1: step4.0, loss 2.4332594871520996
loss diff 0.007616281509399414:: stage 0: step6.0, loss 2.4768574237823486 stage1: step6.0, loss 2.469241142272949
loss diff 0.09828329086303711:: stage 0: step8.0, loss 2.8659024238586426 stage1: step8.0, loss 2.7676191329956055
loss diff -0.013252735137939453:: stage 0: step10.0, loss 2.5272772312164307 stage1: step10.0, loss 2.54052996635437
loss diff 0.0029518604278564453:: stage 0: step12.0, loss 2.925830125808716 stage1: step12.0, loss 2.9228782653808594
loss diff -0.025212764739990234:: stage 0: step14.0, loss 2.8660728931427 stage1: step14.0, loss 2.8912856578826904
loss diff -0.20068717002868652:: stage 0: step16.0, loss 2.688136339187622 stage1: step16.0, loss 2.8888235092163086
loss diff -0.11827921867370605:: stage 0: step18.0, loss 2.8778326511383057 stage1: step18.0, loss 2.9961118698120117
loss diff -0.3702104091644287:: stage 0: step20.0, loss 2.6421046257019043 stage1: step20.0, loss 3.012315034866333
loss diff -0.18424391746520996:: stage 0: step22.0, loss 2.794400215148926 stage1: step22.0, loss 2.9786441326141357
loss diff 0.0862729549407959:: stage 0: step24.0, loss 2.9306914806365967 stage1: step24.0, loss 2.844418525695801
loss diff -0.3921499252319336:: stage 0: step26.0, loss 2.5426809787750244 stage1: step26.0, loss 2.934830904006958
loss diff -0.31029605865478516:: stage 0: step28.0, loss 2.409273147583008 stage1: step28.0, loss 2.719569206237793
loss diff 0.10164618492126465:: stage 0: step30.0, loss 3.2689990997314453 stage1: step30.0, loss 3.1673529148101807
loss diff 0.2335052490234375:: stage 0: step32.0, loss 2.999685287475586 stage1: step32.0, loss 2.7661800384521484
loss diff -0.04235649108886719:: stage 0: step34.0, loss 2.890172243118286 stage1: step34.0, loss 2.9325287342071533
loss diff -0.12056493759155273:: stage 0: step36.0, loss 2.665548086166382 stage1: step36.0, loss 2.7861130237579346
loss diff 0.07198667526245117:: stage 0: step38.0, loss 2.6480774879455566 stage1: step38.0, loss 2.5760908126831055
loss diff -0.3357582092285156:: stage 0: step40.0, loss 2.5086145401000977 stage1: step40.0, loss 2.8443727493286133
stage0, stage2, losses all close False
loss diff 0.02541184425354004:: stage 0: step2.0, loss 3.7346994876861572 stage2: step2.0, loss 3.709287643432617
loss diff 0.012979507446289062:: stage 0: step4.0, loss 2.4462392330169678 stage2: step4.0, loss 2.4332597255706787
loss diff 0.007616281509399414:: stage 0: step6.0, loss 2.4768574237823486 stage2: step6.0, loss 2.469241142272949
loss diff 0.09828329086303711:: stage 0: step8.0, loss 2.8659024238586426 stage2: step8.0, loss 2.7676191329956055
loss diff -0.013252973556518555:: stage 0: step10.0, loss 2.5272772312164307 stage2: step10.0, loss 2.540530204772949
loss diff 0.0029518604278564453:: stage 0: step12.0, loss 2.925830125808716 stage2: step12.0, loss 2.9228782653808594
loss diff -0.025212526321411133:: stage 0: step14.0, loss 2.8660728931427 stage2: step14.0, loss 2.8912854194641113
loss diff -0.20068764686584473:: stage 0: step16.0, loss 2.688136339187622 stage2: step16.0, loss 2.888823986053467
loss diff -0.11827921867370605:: stage 0: step18.0, loss 2.8778326511383057 stage2: step18.0, loss 2.9961118698120117
loss diff -0.3702104091644287:: stage 0: step20.0, loss 2.6421046257019043 stage2: step20.0, loss 3.012315034866333
loss diff -0.18424415588378906:: stage 0: step22.0, loss 2.794400215148926 stage2: step22.0, loss 2.978644371032715
loss diff 0.086273193359375:: stage 0: step24.0, loss 2.9306914806365967 stage2: step24.0, loss 2.8444182872772217
loss diff -0.39215874671936035:: stage 0: step26.0, loss 2.5426809787750244 stage2: step26.0, loss 2.9348397254943848
loss diff -0.31073832511901855:: stage 0: step28.0, loss 2.409273147583008 stage2: step28.0, loss 2.7200114727020264
loss diff 0.0948781967163086:: stage 0: step30.0, loss 3.2689990997314453 stage2: step30.0, loss 3.1741209030151367
loss diff 0.23330307006835938:: stage 0: step32.0, loss 2.999685287475586 stage2: step32.0, loss 2.7663822174072266
loss diff -0.040558815002441406:: stage 0: step34.0, loss 2.890172243118286 stage2: step34.0, loss 2.9307310581207275
loss diff -0.12115120887756348:: stage 0: step36.0, loss 2.665548086166382 stage2: step36.0, loss 2.7866992950439453
loss diff 0.07102394104003906:: stage 0: step38.0, loss 2.6480774879455566 stage2: step38.0, loss 2.5770535469055176
loss diff -0.34471607208251953:: stage 0: step40.0, loss 2.5086145401000977 stage2: step40.0, loss 2.853330612182617
stage0, stage3, losses all close False
loss diff 0.02541184425354004:: stage 0: step2.0, loss 3.7346994876861572 stage3: step2.0, loss 3.709287643432617
loss diff 0.06702589988708496:: stage 0: step4.0, loss 2.4462392330169678 stage3: step4.0, loss 2.379213333129883
loss diff -0.03612780570983887:: stage 0: step6.0, loss 2.4768574237823486 stage3: step6.0, loss 2.5129852294921875
loss diff 0.3160383701324463:: stage 0: step8.0, loss 2.8659024238586426 stage3: step8.0, loss 2.5498640537261963
loss diff -0.01960015296936035:: stage 0: step10.0, loss 2.5272772312164307 stage3: step10.0, loss 2.546877384185791
loss diff 0.03887462615966797:: stage 0: step12.0, loss 2.925830125808716 stage3: step12.0, loss 2.886955499649048
loss diff -0.28507065773010254:: stage 0: step14.0, loss 2.8660728931427 stage3: step14.0, loss 3.1511435508728027
loss diff -0.14540410041809082:: stage 0: step16.0, loss 2.688136339187622 stage3: step16.0, loss 2.833540439605713
loss diff -0.12450242042541504:: stage 0: step18.0, loss 2.8778326511383057 stage3: step18.0, loss 3.0023350715637207
loss diff -0.21169257164001465:: stage 0: step20.0, loss 2.6421046257019043 stage3: step20.0, loss 2.853797197341919
loss diff -0.19033241271972656:: stage 0: step22.0, loss 2.794400215148926 stage3: step22.0, loss 2.9847326278686523
loss diff 0.027364730834960938:: stage 0: step24.0, loss 2.9306914806365967 stage3: step24.0, loss 2.9033267498016357
loss diff -0.27356457710266113:: stage 0: step26.0, loss 2.5426809787750244 stage3: step26.0, loss 2.8162455558776855
loss diff -0.26677465438842773:: stage 0: step28.0, loss 2.409273147583008 stage3: step28.0, loss 2.6760478019714355
loss diff 0.1291036605834961:: stage 0: step30.0, loss 3.2689990997314453 stage3: step30.0, loss 3.139895439147949
loss diff 0.09606313705444336:: stage 0: step32.0, loss 2.999685287475586 stage3: step32.0, loss 2.9036221504211426
loss diff 0.057309865951538086:: stage 0: step34.0, loss 2.890172243118286 stage3: step34.0, loss 2.832862377166748
loss diff -0.01131296157836914:: stage 0: step36.0, loss 2.665548086166382 stage3: step36.0, loss 2.676861047744751
loss diff 0.20771265029907227:: stage 0: step38.0, loss 2.6480774879455566 stage3: step38.0, loss 2.4403648376464844
loss diff 0.017259597778320312:: stage 0: step40.0, loss 2.5086145401000977 stage3: step40.0, loss 2.4913549423217773

Expected behavior
Loss differences are expected to be same or to be different in very small scale, e.g., 1e-5 level.

ds_report output
it raises an error at my side. will update later.

Screenshots
If applicable, add screenshots to help explain your problem.

System info (please complete the following information):

• OS: Amazon Linux
• GPU count and types x8 V100
• Interconnects (if applicable) nvlink
• Python version 3.7.9
• Any other relevant info about your setup

Launcher context
using deepspeed launcher

Start by downloading and running test_diff_stages.py, then inspect the saved loss trajectories in /tmp and reproduce the discrepancy across ZeRO stages with gradient accumulation. Compare the reported losses against the expected 1e-5-scale difference; the issue does not identify a repository file, test, or implementation entry point for further investigation.