[FSDP][bug] Using Apex SyncBN gives incorrect tensor shape error

I am trying to use FSDP with PyTorch AMP (mandated to use with FSDP) + Apex SyncBN however, I get the error

```
Traceback (most recent call last):
File "/data/users/prigoyal/fbsource/fbcode/buck-out/dev/gen/aab7ed39/deeplearning/projects/ssl_framework/tools/run_distributed_engines#link-tree/torch/multiprocessing/spawn.py", line 59, in _wrap
fn(i, *args)
File "/data/users/prigoyal/fbsource/fbcode/buck-out/dev/gen/aab7ed39/deeplearning/projects/ssl_framework/tools/run_distributed_engines#link-tree/vissl/utils/distributed_launcher.py", line 205, in _distributed_worker
process_main(cfg, dist_run_id, local_rank=local_rank, node_id=node_id)
File "/data/users/prigoyal/fbsource/fbcode/buck-out/dev/gen/aab7ed39/deeplearning/projects/ssl_framework/tools/run_distributed_engines#link-tree/vissl/utils/distributed_launcher.py", line 186, in process_main
train_main(
File "/data/users/prigoyal/fbsource/fbcode/buck-out/dev/gen/aab7ed39/deeplearning/projects/ssl_framework/tools/run_distributed_engines#link-tree/vissl/engines/train.py", line 103, in train_main
trainer.train()
File "/data/users/prigoyal/fbsource/fbcode/buck-out/dev/gen/aab7ed39/deeplearning/projects/ssl_framework/tools/run_distributed_engines#link-tree/vissl/trainer/trainer_main.py", line 186, in train
task = train_step_fn(task)
File "/data/users/prigoyal/fbsource/fbcode/buck-out/dev/gen/aab7ed39/deeplearning/projects/ssl_framework/tools/run_distributed_engines#link-tree/vissl/trainer/train_steps/standard_train_step.py", line 198, in standard_train_step
task.amp_grad_scaler.scale(local_loss).backward()
File "/data/users/prigoyal/fbsource/fbcode/buck-out/dev/gen/aab7ed39/deeplearning/projects/ssl_framework/tools/run_distributed_engines#link-tree/torch/_tensor.py", line 255, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
File "/data/users/prigoyal/fbsource/fbcode/buck-out/dev/gen/aab7ed39/deeplearning/projects/ssl_framework/tools/run_distributed_engines#link-tree/torch/autograd/__init__.py", line 146, in backward
Variable._execution_engine.run_backward(
RuntimeError: Function SyncBatchnormFunctionBackward returned an invalid gradient at index 1 - got [3024] but expected shape compatible with [1512]
```

I don't get the error when using PyTorch SyncBN. We have also found that PyTorch SyncBN could be slow for our model trainings and would appreciate using Apex SyncBN.

cc @myleott @min-xu-ai @QuentinDuval @Vittorio-Caggiano