OneBitAdam Incompatible with Pipeline Parallelism

[sdtblck]

So after a bit of work we finally got 1-bit Adam working over at https://github.com/EleutherAI/gpt-neox

But it seems not to be compatible with Pipeline Parallelism. My hypothesis is that the WaitAll here https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/custom_collectives.py#L99 is waiting for all ranks to reach that point, but since the pipeline stages don't all execute concurrently, some never reach it, and it errors out.

The actual error code / stacktrace I'm getting is:

    engine = PipelineEngine(args=args,
  File "/src/deepspeed/deepspeed/runtime/pipe/engine.py", line 52, in __init__
    super().__init__(*super_args, **super_kwargs)
  File "/src/deepspeed/deepspeed/runtime/engine.py", line 174, in __init__
    self._configure_optimizer(optimizer, model_parameters)
  File "/src/deepspeed/deepspeed/runtime/engine.py", line 572, in _configure_optimizer
    self.optimizer = self._configure_fp16_optimizer(basic_optimizer)
  File "/src/deepspeed/deepspeed/runtime/engine.py", line 628, in _configure_fp16_optimizer
    optimizer = FP16_Optimizer(
  File "/src/deepspeed/deepspeed/runtime/fp16/fused_optimizer.py", line 104, in __init__
    self.initialize_optimizer_states()
  File "/src/deepspeed/deepspeed/runtime/fp16/fused_optimizer.py", line 112, in initialize_optimizer_states
    self.optimizer.step()
  File "/src/deepspeed/deepspeed/runtime/fp16/onebit_adam.py", line 340, in step
    self.Compressed_Allreduce(exp_avg,
  File "/src/deepspeed/deepspeed/runtime/fp16/onebit_adam.py", line 157, in Compressed_Allreduce
    cupy_sign_list_packed, cupy_recvbuf_sign, cupy_worker_scale, cupy_recvbuf_scale = gather_host(rank,
  File "/src/deepspeed/deepspeed/runtime/custom_collectives.py", line 99, in gather_host
    MPI.Request.Waitall(requests)
  File "mpi4py/MPI/Request.pyx", line 124, in mpi4py.MPI.Request.Waitall
mpi4py.MPI.Exception: MPI_ERR_IN_STATUS: error code in status

To test my hypothesis I ran a pipeline model with a single stage (so no actual parallelism, but still using the Pipeline Module / Pipeline Engine classes), and this works fine.

I wonder if this is something you've encountered, or potentially something that's fixed by #817 ?

@conglongli ?

[sdtblck]

updating to #817 gives me a different error (with mpi backend)

  File "/usr/local/lib/python3.8/dist-packages/torch/_utils.py", line 264, in _flatten_dense_tensors                                               │
    flat = torch.cat([t.contiguous().view(-1) for t in tensors], dim=0)                                                                            │
RuntimeError: There were no tensor arguments to this function (e.g., you passed an empty list of Tensors), but no fallback function is registered $│
or schema aten::_cat.  This usually means that this function requires a non-empty list of Tensors.  Available functions are [CPU, CUDA, QuantizedC$│
U, BackendSelect, Named, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, UNKNOWN_TENSOR_TYPE_ID, AutogradMLC, AutogradNestedTensor, Autograd│
PrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, Tracer, Autocast, Batched, VmapMode].                                                       │
                                                                                             

[sdtblck]

Getting an identical error with nccl backend

[awan-10]

@sdtblck, thank you for your interest in 1-bit Adam. Indeed, we have not supported/tested it with pipeline parallelism (PP). Given that 1-bit Adam is intended for users with limited bandwidth, we have only tested it on smaller scales (up to 64 GPUs) on relatively smaller models like BERT-Large. The communication in 1-bit Adam code path is not hooked up with groups yet so it will fail. Are you interested in this particular combination of 1-bit and PP? We can put it on our list of things to do. If you want to contribute, you are more than welcome to open a PR for it. Let us know your thoughts :)

[sdtblck]

Hi @awan-10 this is definitely a combination we're interested in!

I've actually been working on a PR today, and I think I have something working, although it doesn't seem to offer much speedup. Maybe someone at the deepspeed team could take a brief look and tell me where I'm going wrong. I'll submit the PR shortly.

[sdtblck]

Hi @awan-10 , I've opened a PR here: #819. Would appreciate if you could take a look when you get time.

[conglongli]

Hi @sdtblck, thank you for the feedback and contribution. Unfortunately currently our team is quite occupied by different kinds of upcoming release items (it's been a while since our last Sep major release), so we might not be able to implement/test new features for 1-bit Adam at this moment. I skimmed over your PR #819 and it looks correct to me (but @awan-10 and @samyam have more expertise in this area). About why you observed only minimal speedup, one hypothesis is that with more pipelines you will need to incur more number of compressed communications. And we observed that in those cases the performance gain can reduce (i.e., not only the communication volume but also the number of communications affect performance). So one thing you could try is to either reduce num param groups (since we do one compressed communication per group) or reduce number of pipeline stages.

In the meanwhile, I'd like to update the doc to mention the limitation of current 1-bit Adam v2. So I have two questions:

1. Have you or your teammates observed any nontrivial speedup in any of your experiments with 1-bit Adam? All we heard from github issues and emails is that 1-bit Adam "didn't work" or "provide no speedup". So it'd be helpful and encouraging if there is any positive results to share :)
2. Have you observed any other limitation of 1-bit Adam? For example, I believe Shivanshu Purohit is your teammate and he said that 1-bit Adam doesn't work with model parallelism? 1-bit adam and model parallelism #781

[sdtblck]

Hi @conglongli .

Once we got it working correctly, and without model parallelism, we observed a significant speedup, in a similar range to what we should expect given the paper / documentation.

That's an interesting note about the param groups. It looks to me like you do a compressed allreduce for every grad/param even , unless I'm misreading the code?

        for group, grads_this_group in zip(self.param_groups, grads_group):
            ...
            for p, grad in zip(group['params'], grads_this_group):
                ...
                if self.adam_freeze_key is False:
                    ...
                else:
                    if 'non_freeze' in group.keys() and group['non_freeze'] is True:
                        ...
                    else:
                        ...
                        if self.size > 1:
                            exp_avg.set_(
                                self.comm_backend_handle.compressed_allreduce(
                                    exp_avg,
                                    state['worker_error'],
                                    state['server_error'],
                                    self.deepspeed.local_rank))

Could this not be optimized by bucketing, similar to how the deepspeed backend does it?

I'm not sure about the issue linked above, I'll have to ask Shivanshu. Our code is changing quite quickly these days, so it could be that it's due to a bug that we've now fixed.

I'm yet to do proper ablations of the speedups with megatron-style model parallelism, but it does at least run. I'll run a test and get back to you on that.

[conglongli]

@sdtblck Glad to hear the positive results! And your observation about when the compressed allreduce happens is right in certain situation. The reason I said param group instead of param is because we used 1-bit Adam as FP16 fused optimizer, so all params in the same group are fused/flatten. This fusion is actually critical to the performance because without it you will have much more number of communications. I would say it'd helpful if you could find out how many times the compressed_allreduce is called in one step. If my assumption is correct, you will probably see some number at least close to 100 if you only see minimal speedup.

[ShivanshuPurohit]

Hi @conglongli, the issue that I had opened was simply caused by the NCCL/MPI conflict we were facing. which was resolved by using deepspeed.utils.distributed. 1-bit adam definitely works with model parallelism, after resolving this issue. It also shows non-trivial speed-ups after freeze_steps are finished.

[sdtblck]

@conglongli Just realised we never responded to this - We experimentally tested your theory and in fact compressed_allreduce is only called once or twice per step, depending on the number of optimizer param groups.

[conglongli]

Hi @sdtblck , @awan-10 and I just had a discussion about this. One thing is that 1-bit Adam's performance gain depends on the size of the model and the bandwidth of the system. With pipeline parallelism, essentially each pipeline will have smaller model size, which could affect the perf gain. Could you try increase your model size and see if that helps?

On the other hand, for you PR #819, we do want to merge it regardless of the perf issue because it should have been like this in the first place. But because we changed the filename later (onebitadam.py to adam.py) and made some changes to the code, the PR has some conflict. Also we already merged the 1-bit Adam staging branch. Just wonder if you have some bandwidth to create a new PR for this? If not, do you mind us to create a new PR for this?

[sdtblck]

Hi @conglongli !
I can create a new PR for this this weekend, if you haven't already