Option to preserve bitwise accuracy of gradient checkpointed vs non-checkpointed dropout

[mcarilli]

This issue was noticed, and fix proposed, by @raulpuric.

Checkpointing is implemented by rerunning a forward-pass segment for each checkpointed segment during backward. This can result in the RNG state advancing more than it would without checkpointing, which can cause checkpoints that include dropout invocations to lose end-to-end bitwise accuracy as compared to non-checkpointed passes.

The present PR contains optional logic to juggle the RNG states such that checkpointed passes containing dropout achieve bitwise accuracy with non-checkpointed equivalents.** The user requests this behavior by supplying preserve_rng_state=True to torch.utils.checkpoint or torch.utils.checkpoint_sequential.

Currently, preserve_rng_state=True may incur a moderate performance hit because restoring MTGP states can be expensive. However, restoring Philox states is dirt cheap, so @syed-ahmed's RNG refactor, once merged, will make this option more or less free.

I'm a little wary of the def checkpoint(function, *args, preserve_rng_state=False): argument-passing method (specifically, putting a kwarg after a variable argument list). Python 3 seems happy with it.
Edit: It appears Python 2.7 is NOT happy with a kwarg after *args. preserve_rng_state also needs to be communicated in a way that doesn't break any existing usage. I'm open to suggestions (a global flag perhaps)?

**Batchnorm may still be an issue, but that's a battle for another day.

[soumith]

lgtm, let's see if tests pass

[facebook-github-bot]

@soumith is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[apaszke]

I realized there are some pretty fundamental design flaws in what you're doing (we might not be caching the right RNG state). Also, please don't use global state unless it's necessary...

[mcarilli]

@apaszke I know we might not be preserving the right RNG states. I left some comments to that effect in the code. It's challenging because A. we can't know what devices will be used in the run_fn, and B. we also can't know if the user intends to initialize the cuda context on one or more devices within their run_fn.

To deal with A. I believe we have two choices:

1. Stash and restore the RNG states for ALL devices visible to this process. Probably overly conservative/paranoid for the majority of cases.
2. Stash and restore the RNG states for the current device.
   I went with option 2 because it seemed like a reasonable compromise. I can add a warning to the docs that if the run_fn moves operations to other devices, bitwise accuracy is no longer guaranteed.

I have no idea how to deal with B., but B. seems like more of an edge case.

The global state was necessary (imo) because I couldn't figure out a way to mix *args with kwargs in a way that was both python 2.7-compliant and would not break existing code. I was prepared to experiment with other approaches but it appears you can live with the global flag. torch.backends.cudnn.deterministic is also exposed as a global flag so this seemed symmetric to me.

[apaszke]

Some commentary:

1. Yes, I agree this is definitely too much, especially considering the overhead it adds.
2. I understand, but this really is a bad compromise. Most user PyTorch programs hardly ever change the current device, but simply specify device=... or use .to(...) to transfer some inputs to a different one, and then do operations on those tensors. This means that if someone uses more than a single device, you'll have a miss rate of almost 100%.

I think a better compromise would be to not use the current device to guide the heuristic, but any tensor input. That would be a reasonable solution I think, but it would still warrant some commentary in the docs.

---

Regarding your second issue, you should add **kwargs to the functions you're interested in, and then do something like this:

def my_fn(arg1, *args, **kwargs):
    my_kwarg = kwargs.pop('my_kwarg', <default value>)
    if kwargs:
        raise ValueError("Unexpected keyword arguments: " + ",".join(arg for arg in kwargs))