[BUG] ZeRO-1 hangs in optimizer step when used in Pipeline

Describe the bug

TLDR: Pipeline works w/o ZeRO-1 but hangs w/ ZeRO-1

At BigScience we have a very large embed layer for the ml training and using Megatron-Deepspeed GPT we are trying to give it a whole pipe stage, since sharing with the transformer layer is too much and we get OOM.

To exemplify the problem let's just slice on embed layer:

        super().__init__(layers=self.specs,
                         loss_fn=CrossEntropy,
                         topology=topo,
                         activation_checkpoint_interval=interval,
                         partition_method='type:embed')

instead of partition_method='type:embed|transformer' as it hides the problem on a small setup.

This is called here: https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/2d9744f23df1a67b4cc1523e3bbdcaca738eb391/megatron/model/gpt_model.py#L279-L283

When ZeRO-1 is used this hangs. W/o ZERO-1 it works.

The hanging is in (py-spy dump):

Thread 1084205 (active): "MainThread"
    get_grad_norm_direct (deepspeed/runtime/zero/stage2.py:1512)
    step (deepspeed/runtime/zero/stage2.py:1645)
    _take_model_step (deepspeed/runtime/engine.py:1538)
    _exec_optimizer_step (deepspeed/runtime/pipe/engine.py:1124)
    _exec_schedule (deepspeed/runtime/pipe/engine.py:1335)
    train_batch (deepspeed/runtime/pipe/engine.py:329)
    train_step (megatron/training.py:405)
    train (megatron/training.py:737)
    pretrain (megatron/training.py:165)
    <module> (pretrain_gpt.py:237)

I'm testing on just 2 gpus with 2 transformer layers and a tied embed before and after the transformer layers.

Let's look at partitioning weights:

• with type:transformer [0, 0, 0, 1, 1, 0, 0, 0, 0]
• with type:embed|transformer [0, 1, 0, 1, 1, 0, 0, 1, 0]

so the partitioning is identical:[0, 0, 0, 1], [1, 0, 0, 0, 0]

but with type:embed [0, 1, 0, 0, 0, 0, 0, 1, 0]
it splits [0, 1, 0, 0, 0, 0, 0], [1, 0]
and so it doesn't know how to handle a boundary that is not a transformer layer

If I trace the schedule cmds on the 2 gpus, I get:

0 LoadMicroBatch(buffer_id=0)
1 RecvActivation(buffer_id=0)
0 ForwardPass(buffer_id=0)
0 SendActivation(buffer_id=0)
0 LoadMicroBatch(buffer_id=1)
1 LoadMicroBatch(buffer_id=0)
0 ForwardPass(buffer_id=1)
1 ForwardPass(buffer_id=0)
0 SendActivation(buffer_id=1)
0 RecvGrad(buffer_id=0)
0 BackwardPass(buffer_id=0)
1 BackwardPass(buffer_id=0)
1 RecvActivation(buffer_id=1)
1 SendGrad(buffer_id=0)
1 LoadMicroBatch(buffer_id=1)
1 ForwardPass(buffer_id=1)
1 BackwardPass(buffer_id=1)
1 SendGrad(buffer_id=1)
1 ReduceTiedGrads()
1 ReduceGrads()
1 OptimizerStep()
0 RecvGrad(buffer_id=1)
0 BackwardPass(buffer_id=1)
0 ReduceTiedGrads()
0 ReduceGrads()
0 OptimizerStep()
1 done with schedule

The log is before the cmd is executed.

So gpu 0 hangs in OptimizerStep

I checked that with type:transformer, it's the same trace, but it gets gpu 0 to complete:

0 done with schedule
1 done with schedule

I tried to remap the above calls to the actual parallel sequence, does it look more or less correct? This indeed looks like an interleaved schedule:

0 LoadMicroBatch(buffer_id=0)
0 SendActivation(buffer_id=0) -> 1 RecvActivation(buffer_id=0)
0 ForwardPass(buffer_id=0)

0 LoadMicroBatch(buffer_id=1)    1 LoadMicroBatch(buffer_id=0)
0 ForwardPass(buffer_id=1)       1 ForwardPass(buffer_id=0)

0 SendActivation(buffer_id=1) -> 1 RecvActivation(buffer_id=1)
                                 1 BackwardPass(buffer_id=0)
0 RecvGrad(buffer_id=0)       -> 1 SendGrad(buffer_id=0)  

                                 1 LoadMicroBatch(buffer_id=1)
0 BackwardPass(buffer_id=0)      1 ForwardPass(buffer_id=1)

                                 1 BackwardPass(buffer_id=1)
0 RecvGrad(buffer_id=1)       <- 1 SendGrad(buffer_id=1)
                                 1 ReduceTiedGrads()
                                 1 ReduceGrads()
                                 1 OptimizerStep()

0 BackwardPass(buffer_id=1)
0 ReduceTiedGrads()
0 ReduceGrads()
0 OptimizerStep()

0 done with schedule
1 done with schedule

with a bunch of prints and CUDA_LAUNCH_BLOCKING=1 I was able to trace it to this:

https://github.com/microsoft/DeepSpeed/blob/85ce85dd5f4b18c0019a5121b06900e3a2c3933b/deepspeed/runtime/zero/stage2.py#L1509-L1510

so gpu0 is blocking on syncing with the gpu 1, but the latter has gone already.

The culprit seems to be this:

https://github.com/microsoft/DeepSpeed/blob/85ce85dd5f4b18c0019a5121b06900e3a2c3933b/deepspeed/runtime/zero/stage2.py#L1636

gpu 0 has 2 items in self.bit16_groups , whereas gpu 1 only 1 item - and so they fail to sync.

In the type:transformer case, it has 2 items in both gpus and so they sync

So it looks like the partitioning didn't take care of ensuring the zero optimizer has the same optimizer groups on each gpu.

the last pipe partition has a different set of optimizer groups from the first one.

With type:transformer there is at least one transformer layer in each pipe stage, which ensures that all self.bit16_groups have 2 groups.

With type:embed we end up with stage 0 of all transformer layers and one embed, and stage 2 with one embed and no transformer layers. So the stages are different.

Specifically it hangs here:

https://github.com/microsoft/DeepSpeed/blob/85ce85dd5f4b18c0019a5121b06900e3a2c3933b/deepspeed/runtime/zero/stage2.py#L1505-L1510

since the other gpu isn't there to do this.

As mentioned earlier turning Z1 off solves the problem, but then we need more hardware to compensate for non-sharded optimizer states.

Thank you!

@ShadenSmith, @tjruwase

Start by reproducing the two-GPU pipeline setup and reading deepspeed/runtime/zero/stage2.py around get_grad_norm_direct and optimizer step, alongside deepspeed/runtime/pipe/engine.py. The Megatron-DeepSpeed entry point is megatron/model/gpt_model.py. Done means ZeRO-1 completes optimizer steps without hanging when the pipeline is partitioned at the embedding boundary.