Add torch.__future__._overwrite_module_params_on_conversion global flag, and check it in nn.Module._apply()

[yf225]

#17072 breaks model.to(xla_device), because moving model to XLA device involves changing its parameters' TensorImpl type, and the current implementation of nn.Module.to() doesn't support changing module parameters' TensorImpl type:

# https://github.com/pytorch/pytorch/blob/6dc445e1a84dc5d093d640de54f038f021d13227/torch/nn/modules/module.py#L192-L208
def _apply(self, fn):
    ...
    for param in self._parameters.values():
        if param is not None:
            # Tensors stored in modules are graph leaves, and we don't
            # want to create copy nodes, so we have to unpack the data.
            param.data = fn(param.data)  # NOTE: this doesn't allow changing `param.data`'s TensorImpl type
            if param._grad is not None:
                param._grad.data = fn(param._grad.data)  # NOTE: this doesn't allow changing `param._grad.data`'s TensorImpl type
   ...

To fix this problem, we decided to gradually change the behavior of nn.Module._apply() from "updating existing parameters in-place" to "using new tensors to overwrite the existing parameters". For any fn passed into _apply() that changes the TensorImpl type of the parameters, the behavior of _apply() will be to overwrite the existing parameters. For any fn that doesn't change the TensorImpl type of the parameters, we check the newly added torch.__future__.get_overwrite_module_params_on_conversion() flag to decide whether _apply() should overwrite the existing parameters or update the existing parameters in-place.

This unblocks adding XLA to our CI test suite, which also allows XLA to catch up with other changes in our codebase, notably the c10 dispatcher.

[xla ci]

cc. @resistor @ailzhang

[ailzhang]

@pytorchbot retest this please

[ailzhang]

[edit]: in fact that won't work since xla is still carrying the old patch. We'll have to manually test this out.
@yf225 would you mind pushing an empty commit with "[xla ci]" in commit message to trigger the test?

[gchanan]

I changed my thinking on this a bit, let me know what you think.

Essentially we argued that we should introduce to_, cpu_ etc. to have the new semantics. And for one release to, cpu, etc. will have the old semantics. But realistically, what are we going to do with to, cpu, etc. later? Every model that uses devices call those functions, but few would probably break with semantics changes, which means we should probably leave the old APIs. But then we have two sets of APIs that do the same thing.

Instead, can we just do essentially what you implementation does here, which is take a parameter to control the behavior. Since XLA is new, we can even make the parameter obvious it only applies to cpu/cuda (do we support amd in the frontend), e.g. force_move_params_cpu_cuda or something.

What do you think?

[gchanan]

See the note above, but we need equivalents for cpu, cuda, etc.

[facebook-github-bot]

@yf225 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@yf225 merged this pull request in 6b97279.

[apaszke]

Why do we even have this flag? What's so special about XLA that we cannot overwrite the parameters with new objects, but we can e.g. in the case of CUDA?

NB I would be against adding to_ in addition to to, because even to works in-place on the module (just not on the tensors!), so the semantics would be very confusing.

[gchanan]

@apaszke: it's the combination of the fact that (1) XLA tensors use a different TensorImpl than CPU/CUDA (which use the same) and (2) we don't have VariableImpl anymore to do the extra indirection (by design). So basically, CPUModule.to('xla') and not overwriting the parameters would require changing the type of the TensorImpl or hacking around with the pyobj.

[gchanan]

and note we don't really support moving parameters in-place anyway, as demonstrated in the tests, views on such parameters are broken, we don't version count, etc.

So the flag is just a BC thing before we move everything to overwrite semantics, which are well supported.