Model with ZeRO-1 converges worse than model without ZeRO-1

[MichaelEk]

This result was obtained when training Megatron LM from the examples.

Here is DS configuration:
{ "train_batch_size": 32, "train_micro_batch_size_per_gpu":1, "steps_per_print":10, "prescale_gradients":false, "gradient_clipping":1.0, "wall_clock_breakdown":false, "fp16": { "enabled":true, "loss_scale":0 }, "zero_optimization": { "stage":1 }, "gradient_predivide_factor": 1, "zero_allow_untested_optimizer": true }

We use apex Lamb as optimizer.

[MichaelEk]

Model with ZeRO-1 coverges even worse with Adam

[MichaelEk]

Also, we use PipelineModule model, but pp_world_size = 1

[tjruwase]

@MichaelEk, thanks for reporting this issue. In general LAMB is not expected to play well with ZeRO. However, your observation is surprising for Adam since it works and is well tested. Can you please check to see if same issue occurs without pipeline?

[tjruwase]

@MichaelEk, just wanted to check whether you had any update on this issue?

[ShadenSmith]

I wonder if #1263 fixed the issue?

[MichaelEk]

No. This fix didn't help. Also it has no effect on ZeRO-1 because of that condition:
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/engine.py#L917

Later I will try to set round_robin_gradients=true from my code

[MichaelEk]

Hi! This PR fixes this issue. However, it may work strange way for non-ZeRO optimizers.