[fix] SDP syncing buffers during gradient accumulation

[blefaudeux]

What does this PR do?

Fixes #1041. I just had a minute or two, hoping that it's enough :)

Before submitting

• Did you have fun?
  • Make sure you had fun coding 🙃
• Did you read the contributor guideline?
• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you write any new necessary tests?
  • N/A
• Did you update the changelog? (if needed)
  • N/A

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

[blefaudeux]

cc @min-xu-ai

[min-xu-ai]

lgtm！

[blefaudeux]

I've not followed the CI changes on Fairscale @min-xu-ai, looks like the breakages are not related (but they do touch OSS) but are pytorch version dependent ?

[min-xu-ai]

I trigged a rerun. The cpu test failure seems unrelated. But the GPU test failure seems to be in OSS?

torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/circleci/venv/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/home/circleci/fairscale/tests/optim/test_oss.py", line 956, in run_ddp_parity
    check_optimizer_equivalence(opt, change_train_graph=change_train_graph)
  File "/home/circleci/fairscale/tests/optim/test_oss.py", line 947, in check_optimizer_equivalence
    check_step()
  File "/home/circleci/fairscale/tests/optim/test_oss.py", line 912, in check_step
    loss_sharded_optim = cast(torch.Tensor, sharded_optimizer.step(closure=closure_sharded))
  File "/home/circleci/venv/lib/python3.9/site-packages/torch/optim/optimizer.py", line 109, in wrapper
    return func(*args, **kwargs)
  File "/home/circleci/fairscale/fairscale/optim/oss.py", line 232, in step
    loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
  File "/home/circleci/venv/lib/python3.9/site-packages/torch/optim/optimizer.py", line 109, in wrapper
    return func(*args, **kwargs)
  File "/home/circleci/venv/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/circleci/venv/lib/python3.9/site-packages/torch/optim/adam.py", line 157, in step
    adam(params_with_grad,
  File "/home/circleci/venv/lib/python3.9/site-packages/torch/optim/adam.py", line 213, in adam
    func(params,
  File "/home/circleci/venv/lib/python3.9/site-packages/torch/optim/adam.py", line 255, in _single_tensor_adam
    assert not step_t.is_cuda, "If capturable=False, state_steps should not be CUDA tensors."
AssertionError: If capturable=False, state_steps should not be CUDA tensors.
change_train_graph = True, backend = 'nccl', broadcast_fp16 = False

[min-xu-ai]

it is interesting, all the failures seem to be the same. What the best way for me to run your branch? Do I git remote add your repo & branch?

[ruanslv]

These OSS tests are failing in the main branch (and in other unrelated PRs too). Agree that it's probably related to recent pytorch upgrade as I don't remember them failing last week.

I'm seeing all failures from here in a FSDP PR (#1052), including the CPU and the "test_parity3d_checkpoint_syncbn" ones.

[min-xu-ai]

somehow this PR hitting unrelated test errors. I am replacing it with: #1075

Thanks Ben!