[fix]: let FSDP handle model with multiple forward pass and checkpoint

[min-xu-ai]

This PR adds support of multiple passes of a module in a forward pass and activation checkpoint on the module at the same time.

Right now, the tested cases are:

1. FSDP(ckpt(module))
2. FSDP(ckpt(module, FSDP(BN), module))

Next commit will test and enable FSDP(ckpt(), ..., ckpt()) type of cases.

The fix is to add an option to checkpoint_wrapper so that it can keep a counter on the checkpointed modules so that FSDP can check the counter to determine if the backward callback fired from the module is the last one in the bigger, outer backward pass.

A new API to override the pre/post gradient divide factors. It is currently used in the test to make sure numerical match with DDP.

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
• Did you read the contributor guideline?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

Fixes part of #617 .

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.

Did you have fun?

Make sure you had fun coding 🙃

[min-xu-ai]

@QuentinDuval, once this PR is merged, you need to use mintain_forward_counter=True in vissl when forward-multiple-times and checkpoint are both used.

[sshleifer]

LGTM. Note that I don't have context on why pre_backward_hook_has_run used to be a list.

[sshleifer]

Is this a thing? Which gradients need gradients!?

[min-xu-ai]

Great question! I think it came from @myleott originally. I would love to know more too. :-)

[min-xu-ai]

LGTM. Note that I don't have context on why pre_backward_hook_has_run used to be a list.

Thanks for the review @sshleifer! The pre_backward_hook_has_run being a list isn't originally from me, but I think I know why it is that way.

Say for this example:

def func():
   x = 0
   def nest():
       x = 1

The second assignment to x will be to a new local variable inside the nest function. However, if you turn it into:

def func():
   x = [ 0 ]
   def nest():
      x[0] = 1

Then, the second assignment will be a read of the outer x, which ensures the right variable is updated. It is a very clever trick to turn a write into a read first so that python knows which variable you are referring too. Using python's keyword nonlocal would get the same effect, like:

def func():
  x = 0
  def nest():
     nonlocal x
     x = 1

[min-xu-ai]

originally, the variable is local to _register_pre_backward_hooks function and it only prevents callbacks registered within this function. When there are multiple calls to this function, this local flag won't prevent multiple callbacks. Moving this flag to self solves this problem.

[min-xu-ai]

I have merged this for now @myleott. Definitely happy to address more comments separately.