Repository: facebookresearch/fairscale

Issue:
SDP sync_buffers() during gradient accumulation
Shouldn't [`sync_buffers()`](https://github.com/facebookresearch/fairscale/blob/fd7b962f7ae01b7f6ffbcdb31b4b53297174a645/fairscale/nn/data_parallel/sharded_ddp.py#L224) be guarded by [`_should_accumulate_grads()` just like reduce()](https://github.com/facebookresearch/fairscale/blob/fd7b962f7ae01b7f6ffbcdb31b4b53297174a645/fairscale/nn/data_parallel/sharded_ddp.py#L415)?

I think the current code is broadcasting buffers even in `no_sync()`.


Discussion:
cc @blefaudeux ? 

@bgyoon, thank for the question. Can you add more details for us to understand the context? Do you have any test that demonstrate the issue?
https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#skip-unnecessary-all-reduce-if-training-with-distributeddataparallel-and-gradient-accumulation

> DistributedDataParallel provides [no_sync()](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel.no_sync) context manager which disables gradient all-reduce for particular iteration. no_sync() should be applied to first N-1 iterations of gradient accumulation, the last iteration should follow the default execution and perform the required gradient all-reduce.

For performance reason, when doing gradient accumulation, instead of doing gradient all-reduce every `backward(),` gradients are all-reduced only in the last step right before `optimizer.step()`, with use of `no_sync()`. In SDP, within `no_sync()` context, gradients reduce is skipped.(i.e., only reduce when [`_should_accumulate_grads()` is True](https://github.com/facebookresearch/fairscale/blob/c3f88a6dbcd2b36edaa95e9c6d44b55a7fbb2cf9/fairscale/nn/data_parallel/sharded_ddp.py#L415))

Just like gradient all-reduce happening in `backward()`, buffers are broadcasted in `forward()`. For the same performance reason, just like doing gradient all-reduce only in the last step of gradient accumulation, DDP broadcasts buffer only in the first step of gradient accumulation. **SDP seems to be missing this part.**

In DDP code, `_sync_buffers()` is only called when [`require_forward_param_sync` is True](https://github.com/pytorch/pytorch/blob/3bcc19b29ae1a41db8cabef6154041765a8e7b03/torch/nn/parallel/distributed.py#L1548
). `require_forward_param_sync` is set to True [when `require_backward_grad_sync` is True](https://github.com/pytorch/pytorch/blob/3bcc19b29ae1a41db8cabef6154041765a8e7b03/torch/nn/parallel/distributed.py#L1017). [`require_backward_grad_sync` is set to False](https://github.com/pytorch/pytorch/blob/3bcc19b29ae1a41db8cabef6154041765a8e7b03/torch/nn/parallel/distributed.py#L954) within `no_sync()` context.

I found this issue when examining our test case where SDP was slower than DDP, despite having double the local batch size thanks to memory saving from SDP. It was because of buffer broadcast in every `forward()` during gradient accumulation.

In DDP, buffer broadcast happens only in the first step:
<img width="1731" alt="image" src="https://user-images.githubusercontent.com/8525211/181180157-a573d69e-8c4c-458d-af9f-537c56f5b2a1.png">

In SDP, buffer broadcast happens in every step:
<img width="1731" alt="image" src="https://user-images.githubusercontent.com/8525211/181180318-94ed31d8-4b49-4867-8d63-e7abad6fb230.png">
Oh, I see. Thanks for the detailed explanation! It is really helpful!

We will try to look at the SDP code and try to fix it.

We actually don't know very much about the adoption of SDP. Can you share the reason that you are using SDP not FSDP? We are thinking that FSDP is a superset of SDP so SDP may not be provide much additional to FSDP. So we are thinking of retiring SDP. But again, we don't know the adoption of SDP at this point.
The main reason is

1. Our code is using pytorch-lightning. Simply using fsdp(`DDPFullyShardedStrategy()` in pytorch-lightning) didn't seem to work out of box for some unknown reason.
2. fairscale/fsdp in pytorch-lightning doesn't support `clip_grad_by_norm()` for the [code structural reason](https://github.com/Lightning-AI/lightning/blob/9596fabe7baade5fd2b3625877bc15d356cdc553/src/pytorch_lightning/plugins/precision/fully_sharded_native_amp.py#L39) right now. Since I needed that, I didn't bother try fixing 1.
3. Also,(I haven't really read fsdp stuff yet, so I may be wrong) if fsdp works like Zero3, I think model weight allgather would be required in every `forward()` when using gradient accumulation(unlike buffer broadcast or grad reduce which can be skipped during gradient accumulation) My Ethernet connection is kind of on slow side. I think the communication cost would slow down the training.

In other words, my code is using gradient accumulation on multiple nodes with slow ethernet connection. In this setup, I found memory saving(and possible speed up from larger batch size) from Zero1,2,3 is offset by communication cost. It's actually slower to use Zero2/3 about the same speed with Zero1. I believe FSDP would be similar to Zero2/3.(maybe I am wrong...)

With `sync_buffer()` in every `forward()`, SDP was 8% slower than DDP. With `broadcast_buffer=False`, SDP was 3% faster than DDP.
Got it! Thanks for the context!

The 3% gain is from doubling of the batch size?

FSDP can work in zero2 mode by setting reshard_after_forward flag to False. That way, the all-gather is not performed on the backward to reduce some pressure on the network (since grad scatter-gather is happening on the backward, regardless of no_sync() or not). It is important to note that all-gather/scatter-gather are all happening in parallel with GPU kernel computation. That way, unless your ethernet is very slow, they should not cause slowdown on training.

Can you share what error did you run into with FSDP if you still have that info? No worries if you don't remember.
3% gain is with the doubled batch size. The gain is kind of too small... I am not sure it's worth changing. But if ethernet condition changes or model size changes or I try squeezing in more than doubled batchsize, it at least has potential to be faster than DDP.

FSDP in pytorch-lightning didn't error out. It looked like running fine fast but the network bandwidth was 0bps when run on 4 nodes, meaning something was subtly wrong.

Zero2 was about 3 times slower than Zero1 because gradients had to be reduced every `backward()` during gradient accumulation. Ah, this is probably because of our setup. The model is so small that communication time is much larger than computation time.

I have realized I may not be the target audience of ZeRO or SDPs. They are meant to enable training large model that won't fit on 1 GPU otherwise, not for accelerating training speed of small model.
If your activation size is big, you might be able to use activation checkpoint wrapper and further increase your batch size. But if your batch size is already large, the kernel speed gain may not be much. Also, no_sync used more memory because it needs to keep the grad around between steps. 
hey @bgyoon thanks for the detailed explanation (and thanks @min-xu-ai for the heads up, sorry for the delay), it seems right indeed and easy enough to fix ! (ideally this could go with a unit test to make sure this doesn't re-appear). Re: speed @bgyoon you can try to play with the bucketing setting, depending on the model there might be some savings there. In general the FSDP method should be faster though, SDP doesn't look into the model per say which means that all the broadcasts will not be masked by compute (while they should be with FSDP)

Resolving patch:
diff --git a/fairscale/nn/data_parallel/sharded_ddp.py b/fairscale/nn/data_parallel/sharded_ddp.py
--- a/fairscale/nn/data_parallel/sharded_ddp.py
+++ b/fairscale/nn/data_parallel/sharded_ddp.py
@@ -218,7 +218,7 @@ def forward(self, *inputs: Any, **kwargs: Any) -> Any:
             if needs_setup:
                 self.refresh_trainable()
 
-            if self._enable_broadcast_buffers:
+            if self._enable_broadcast_buffers and not self._should_accumulate_grads:
                 # NCCL communications are on a different stream, needs to be blocking
                 # for the subsequent FW to be correct
                 self.sync_buffers(blocking=True)


Test patch:
diff --git a/tests/nn/data_parallel/test_sharded_ddp_features.py b/tests/nn/data_parallel/test_sharded_ddp_features.py
--- a/tests/nn/data_parallel/test_sharded_ddp_features.py
+++ b/tests/nn/data_parallel/test_sharded_ddp_features.py
@@ -112,6 +112,14 @@ def closure():
         with ddp_model.no_sync() if grad_accumulation else suppress():
             input_tensor = torch.rand((64, 2)).to(device)
             loss = ddp_model(input_tensor).abs().sum()
+
+            # If grad_accumulation, we can check after the forward that the models are different
+            # (not synced)
+            if grad_accumulation:
+                check_same_models_across_ranks(
+                    ddp_model, dist.group.WORLD, params_should_be_equal=False, check_broadcast_buffers=True
+                )
+
             loss.backward()
         return loss