are we switching from using output to output.data in non-grad accumulation use case?

Good question. This is good for me to clarify. Was there any reason to use output before?

In my understanding, it does not matter since we are not supporting taking the gradient of the gradient, so we do not need to record any operations on output in the autograd graph. Fairscale uses output.data in both code paths, though I do not see why either way is better for the non-gradient accumulation case.

nit: if you want, maybe pass in named args here, so future developers know what it is

Wondering why we have duplicated configs?

I wanted to test interleaving both ways:

1. with no_sync() -> without no_sync() -> with no_sync()
2. without no_sync() -> with no_sync() -> without no_sync()

The reason I wanted two separate tests is that it could matter what was the last accumulation mode right before gradient synchronization.

Good call! Let's keep it then and add a small comment so no one removes it in the future.

to reduce # of tests, can we just have:

use_context = true, interval={2,4}
use_context = false, interval={2,4}

Given your comments last time, I thought about how to keep the set of tests minimal. I feel like these 4 configs test distinct things:

1. Gradient accumulation with no_sync()
2. Gradient accumulation without no_sync()
3. Gradient accumulation interleaving no_sync() and without no_sync(), where the last iteration before synchronizing gradients is in no_sync()
4. Gradient accumulation interleaving no_sync() and without no_sync(), where the last iteration before synchronizing gradients is outside no_sync()

My concern is that I do not want to overfit the current working implementation. It not obvious to me that 1) and 2) imply 3) and 4), and if we only had 3) and 4) and they break, then we would probably end up rewriting 1) and 2).

Sounds good! Considering you've gone through the tradeoff this makes sense to leave as is

Might be useful to add a small comment about how gradient accumulation without no_sync is implemented. From my understanding:

1. During backward, we point an attribute _saved_grad_shard to the gradient shard
2. If we are accumulating gradients, we accumulate it on _saved_grad_shard
3. When finalizing backwards before running the optimizer, we point p.grad to the saved grad shard so that the optimizer works on the right accumulated gradient.

add comment here that this is what makes the gradient accumulation work?

Can we have a unittest where can_accumulate_grad = False and the user tries to accumulate grads and we raise an appropriate error?

After some investigation, I think that having a non-silent error for the case of gradient accumulation outside no_sync() while using CPU offloading requires some non-trivial re-design. It may be easiest to leave this error as silent for now and work on adding the compatibility itself.

(Some clarifying questions and comments)
Suppose we are using CPU offloading.

• Outside no_sync(), if we want to accumulate gradients, are we performing the addition between the existing gradient and the newly-reduced gradient on CPU or on GPU? If on CPU, then we should perform a device-to-host transfer of the reduced gradient. If on GPU, then we should perform a host-to-device transfer of the existing gradient and a device-to-host transfer to re-offload the result.
• Inside no_sync(), the existing implementation does not offload any gradients to CPU. Rather, the gradients are held in GPU memory until the first iteration outside no_sync(), which performs the gradient synchronization. At the end of that iteration's backward pass, the synchronized gradient shard is offloaded to CPU. Should we include any warning or message at runtime to the explain this behavior to the user? I added a note about it to the no_sync() docstring.

The challenge behind a non-silent error is that the no_sync() + CPU offloading case conflicts with the non-no_sync() + CPU offloading case.

• The crux is that accumulating gradients using no_sync() follows the pattern of: accumulate for N-1 iterations inside no_sync() and execute 1 normal iteration outside no_sync().
• That final normal iteration is indistinguishable from a non-no_sync() iteration unless we track something like a bool flag indicating that the last iteration was inside no_sync().
  • I am reluctant to add such a flag since it is solely a patch for one case and may be hiding the underlying design problem, but I am open to your thoughts.
• For the final iteration coming out of no_sync(), the gradient is still on GPU, so performing the accumulation computation param._saved_grad_shard.data += output.data has no issue.
• For a non-no_sync() iteration (after the first), the gradient was previously offloaded to CPU, so performing the accumulation computation param._saved_grad_shard.data += output.data has conflicting devices.

When we are using gradient accumulation outside of no_sync + CPU offload, don't we already raise an appropriate error on L1336 of this PR? And is it possible to add a unittest for this?

That is an internal assert. Given the current implementation, that assert will never get triggered. I added it because the logic is quite complicated, and I wanted to demonstrate that if we are in the non-no_sync() + CPU offloading case, then we should never be in that branch.

(or more directly, the contrapositive: if we are in that branch, we should not be in the non-no_sync() + CPU offloading case, meaning that in particular we must not be CPU offloading)

Since this PR is getting a bit cluttered, I wanted to specifically point this part out. Previously, I presented the logic here incorrectly, but hopefully this should make sense now.

I think the precise condition of when to use p._saved_grad_shard is if the previous iteration was outside no_sync().

Suppose we have
(1) some iterations outside no_sync() ->
(2) some iterations inside no_sync() ->
(3) one iteration outside no_sync().

• In the pre-backward hook of (3), the FSDP instance holds an unsharded gradient in p.grad, which is the result of accumulating gradients from (2).
• It computes that iteration's gradient, which is accumulated with the existing p.grad from (2) via the autograd engine and still stored in p.grad.
• In the post-backward hook, it reduce-scatters that accumulated gradient stored in p.grad.
• After the reduce-scatter, it accumulates the accumulated gradient from (2) and (3) with the accumulated gradient from (1) saved in _saved_grad_shard. This "super-accumulated" gradient is stored in _saved_grad_shard.
  • This step shows why as long as the previous iteration was outside no_sync(), there may be a gradient to accumulate on the first future iteration also outside no_sync().

I had misunderstood the comments from Fairscale (see here). I did not realize that the conditioning on inside/outside no_sync() referred to the previous iteration.

One more thing to point out: I previously had an assert like not self.cpu_offload.offload_params just as an internal assert to make sure that CPU offloading never takes this code path.

However, I changed it to a more direct assert here in case we distinguish between offloading parameters and offloading gradients in the future before we solve gradient accumulation with CPU offloading.

will it work if removing the .data?

Yup, it still works. I will remove the .data for both in this line.

could we add a github issue to support grad accumulation in cpu offloading?

Done: #73784

could we add a warning for this case and warning the comment "# FSDP currently does not support gradient accumulation
# outside no_sync() when using CPU offloading. Trying to
# do so yields incorrect results since FSDP will use the
# newly-reduced gradient instead of accumulating with any
# existing gradient."