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

Summary:
pytorch#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:
```python
# 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
   ...
```

yf225 TODO: fix the description here when we finish the implementation

To fix this problem, we introduce a new API `model.to_()` that always assign new tensors to the parameters (thus supporting changing the parameters to any TensorImpl type), and also bump the version counter of the original parameters correctly so that they are invalidated in any autograd graph they participate in.

We also add warning to the current `model.to()` API to inform users about the upcoming behavior change of `model.to()`: in future releases, it would create and return a new model instead of in-place updating the current model.

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
Pull Request resolved: pytorch#21613

Differential Revision: D15895387

Pulled By: yf225

fbshipit-source-id: b79f230fb06019122a37fdf0711bf2130a016fe6

Author: Will Feng

aten/src/ATen/native/TypeProperties.cpp

@@ -38,6 +38,11 @@ bool is_quantized(const Tensor& self) {
   return self.is_quantized();
 }
 
+// True if `self` has the same derived type of TensorImpl as `other`.
+bool _has_same_tensorimpl_type(const Tensor& self, const Tensor& other) {
+  return typeid(*(self.unsafeGetTensorImpl())) == typeid(*(other.unsafeGetTensorImpl()));
+}
+
 Tensor type_as(const Tensor& self, const Tensor& other) {
   return self.toType(other.type());
 }

aten/src/ATen/native/native_functions.yaml

@@ -2036,6 +2036,9 @@
 - func: type_as(Tensor self, Tensor other) -> Tensor
   variants: method
 
+- func: _has_same_tensorimpl_type(Tensor self, Tensor other) -> bool
+  variants: function
+
 - func: _unique(Tensor self, bool sorted=True, bool return_inverse=False) -> (Tensor, Tensor)
   variants: function
   dispatch:

test/test_nn.py

@@ -1602,6 +1602,144 @@ def add_one_inplace(t):
         with self.assertRaisesRegex(RuntimeError, "modified by an inplace operation"):
             pgm.backward(torch.randn(10, 20))
 
+    def test_overwrite_module_params_on_conversion(self):
+        torch.__future__.set_overwrite_module_params_on_conversion(False)
+
+        # Test that if the conversion function passed to `module._apply()`
+        # changes the TensorImpl type of `module`'s parameters, the `module`'s
+        # parameters are always overwritten, regardless of the value of
+        # `torch.__future__.get_overwrite_module_params_on_conversion()`.
+        m = nn.Linear(20, 10)
+        m.weight.grad = torch.randn(10, 20)
+        weight_ref = m.weight
+        weight_grad_ref = m.weight.grad
+        m = m._apply(lambda t: torch.sparse_coo_tensor(torch.zeros([2, 1]), torch.ones([1]), torch.Size([10, 20])))
+        self.assertNotEqual(weight_ref.layout, m.weight.layout)
+        self.assertNotEqual(weight_grad_ref.layout, m.weight.grad.layout)
+
+        # Test that under the current default settings
+        # (`torch.__future__.get_overwrite_module_params_on_conversion() == False`),
+        # a view to a module's parameters is not pointing to the same storage as
+        # its base variable after converting the module to a different dtype.
+        m = nn.Linear(20, 10).float()
+        mw = m.weight[:]
+        m.double()
+        mw[0][0] = 5
+        with self.assertRaisesRegex(RuntimeError, "Expected object of scalar type Float but got scalar type Double"):
+            mw[0][0] == mw._base[0][0]
+
+        torch.__future__.set_overwrite_module_params_on_conversion(True)
+
+        # Test that if `torch.__future__.get_overwrite_module_params_on_conversion() == True`,
+        # a view to a module's parameters is still pointing to the same storage as
+        # its base variable after converting the module to a different dtype.
+        m = nn.Linear(20, 10).float()
+        mw = m.weight[:]
+        m.double()
+        mw[0][0] = 5
+        self.assertTrue(mw[0][0] == mw._base[0][0])
+
+        # Test that if `torch.__future__.get_overwrite_module_params_on_conversion() == True`,
+        # `float_module.double()` doesn't preserve previous references to
+        # `float_module`'s parameters or gradients.
+        m = nn.Linear(20, 10).float()
+        m.weight.grad = torch.randn(10, 20).float()
+        weight_ref = m.weight
+        weight_grad_ref = m.weight.grad
+        m.double()
+        self.assertNotEqual(weight_ref.dtype, m.weight.dtype)
+        self.assertNotEqual(weight_grad_ref.dtype, m.weight.grad.dtype)
+
+        def add_one_inplace(t):
+            return t.add_(1.0)
+
+        # Test that if `torch.__future__.get_overwrite_module_params_on_conversion() == True`,
+        # applying an in-place operation to a module would bump the module's
+        # original parameters' version counter.
+        m = nn.Linear(20, 10)
+        pvm = m.weight.mul(m.weight)
+        weight_ref = m.weight
+        m_weight_version_saved = weight_ref._version
+        m = m._apply(add_one_inplace)
+        # Test that the in-place operation bumps the original parameter's version counter
+        self.assertGreater(weight_ref._version, m_weight_version_saved)
+        with self.assertRaisesRegex(RuntimeError, "modified by an inplace operation"):
+            pvm.backward(torch.randn(10, 20))
+
+        # Test that if `torch.__future__.get_overwrite_module_params_on_conversion() == True`,
+        # applying an in-place operation to a module would bump the module's
+        # original parameters' gradients' version counter.
+        m = nn.Linear(20, 10)
+        m.weight.grad = torch.randn(10, 20).requires_grad_()
+        pgm = m.weight.grad.mul(m.weight.grad)
+        weight_grad_ref = m.weight.grad
+        m_weight_grad_version_saved = weight_grad_ref._version
+        m = m._apply(add_one_inplace)
+        self.assertGreater(weight_grad_ref._version, m_weight_grad_version_saved)
+        with self.assertRaisesRegex(RuntimeError, "modified by an inplace operation"):
+            pgm.backward(torch.randn(10, 20))
+
+        # Test that if `torch.__future__.get_overwrite_module_params_on_conversion() == True`,
+        # applying an out-of-place operation to a module doesn't bump
+        # the module's original parameters' version counter.
+        m = nn.Linear(20, 10)
+        weight_ref = m.weight
+        m_weight_version_saved = weight_ref._version
+        m = m._apply(lambda t: torch.randn(t.shape))
+        self.assertEqual(weight_ref._version, m_weight_version_saved)
+
+        # Test that if `torch.__future__.get_overwrite_module_params_on_conversion() == True`,
+        # applying an out-of-place operation to a module doesn't bump
+        # the module's original parameters' gradients' version counter.
+        m = nn.Linear(20, 10)
+        m.weight.grad = torch.randn(10, 20).requires_grad_()
+        weight_grad_ref = m.weight.grad
+        m_weight_grad_version_saved = weight_grad_ref._version
+        m = m._apply(lambda t: torch.randn(t.shape))
+        self.assertEqual(weight_grad_ref._version, m_weight_grad_version_saved)
+
+        torch.__future__.set_overwrite_module_params_on_conversion(False)
+
+    @unittest.skipIf(not TEST_CUDA, "CUDA unavailable")
+    def test_overwrite_module_params_on_conversion_cpu_cuda(self):
+        torch.__future__.set_overwrite_module_params_on_conversion(False)
+
+        # Test that under the current default settings
+        # (`torch.__future__.get_overwrite_module_params_on_conversion() == False`),
+        # a view to a module's parameters is not pointing to the same storage as
+        # its base variable after converting the module to a different device.
+        m = nn.Linear(20, 10)
+        mw = m.weight[:]
+        m.to('cuda')
+        with torch.no_grad():
+            # Without using `torch.no_grad()`, this will leak CUDA memory.
+            # (Issue is filed at https://github.com/pytorch/pytorch/issues/21875)
+            mw[0][0] = 5
+        with self.assertRaisesRegex(RuntimeError, "Expected object of backend CPU but got backend CUDA"):
+            mw[0][0] == mw._base[0][0]
+
+        torch.__future__.set_overwrite_module_params_on_conversion(True)
+
+        # Test that if `torch.__future__.get_overwrite_module_params_on_conversion() == True`,
+        # a view to a module's parameters is still pointing to the same storage as
+        # its base variable after converting the module to a different device.
+        m = nn.Linear(20, 10)
+        mw = m.weight[:]
+        m.to('cuda')
+        mw[0][0] = 5
+        self.assertTrue(mw[0][0] == mw._base[0][0])
+
+        # Test that if `torch.__future__.get_overwrite_module_params_on_conversion() == True`,
+        # `cpu_module.to("cuda")` doesn't preserve previous references to
+        # `cpu_module`'s parameters or gradients.
+        m = nn.Linear(20, 10)
+        m.weight.grad = torch.randn(10, 20)
+        weight_ref = m.weight
+        weight_grad_ref = m.weight.grad
+        m.to('cuda')
+        self.assertNotEqual(weight_ref.device, m.weight.device)
+        self.assertNotEqual(weight_grad_ref.device, m.weight.grad.device)
+
     def test_type(self):
         l = nn.Linear(10, 20)
         net = nn.Module()

torch/__future__.py

@@ -0,0 +1,19 @@
+"""
+This global flag controls whether to assign new tensors to the parameters
+instead of changing the existing parameters in-place when converting an `nn.Module`
+using the following methods:
+1. `module.cuda()` / `.cpu()` (for moving `module` between devices)
+2. `module.float()` / `.double()` / `.half()` (for converting `module` to a different dtype)
+3. `module.to()` / `.type()` (for changing `module`'s device or dtype)
+4. `module._apply(fn)` (for generic functions applied to `module`)
+
+Default: False
+"""
+_overwrite_module_params_on_conversion = False
+
+def set_overwrite_module_params_on_conversion(value):
+    global _overwrite_module_params_on_conversion
+    _overwrite_module_params_on_conversion = value
+
+def get_overwrite_module_params_on_conversion():
+    return _overwrite_module_params_on_conversion

torch/__init__.py

@@ -310,6 +310,7 @@ def manager_path():
 import torch.backends.mkl
 import torch.backends.openmp
 import torch.__config__
+import torch.__future__
 
 _C._init_names(list(torch._storage_classes))
 

torch/csrc/autograd/variable.cpp

@@ -88,7 +88,7 @@ void Variable::set_data(const at::Tensor &new_data) {
   // from `new_data` to `var`. It requires that `new_data` has the same derived
   // type of TensorImpl as `var`.
   TORCH_CHECK(
-    typeid(*(this->unsafeGetTensorImpl())) == typeid(*(new_data.unsafeGetTensorImpl())),
+    _has_same_tensorimpl_type(*this, new_data),
     "Attempted to call `variable.set_data(tensor)`, but `variable` and `tensor` have different types of TensorImpl.");
 
   // Resets gradient accumulator if metadata is out of date

torch/csrc/autograd/variable.h

@@ -248,10 +248,9 @@ struct TORCH_API Variable : public at::Tensor {
       bool keep_graph,
       bool create_graph) const;
 
-  /// Sets the `Tensor` held by this `Variable` to the one supplied.
-  /// It is rarely necessary to call this; it's used, for example, when
-  /// a non-sparse gradient gets added to a sparse gradient, requiring
-  /// the type of the gradient `Variable` to become non-sparse.
+  /// Sets the tensor data held by this `Variable` to be the same as `new_data`.
+  /// It requires that `new_data` has the same derived type of TensorImpl as
+  /// this `Variable`, by checking `_has_same_tensorimpl_type(this, new_data)`.
   void set_data(const at::Tensor &new_data);
 
   /// Set the gradient edge -- i.e. `grad_fn` and `input_nr` -- of the

torch/nn/modules/module.py

@@ -193,15 +193,44 @@ def _apply(self, fn):
         for module in self.children():
             module._apply(fn)
 
-        for param in self._parameters.values():
+        def compute_should_use_set_data(tensor, tensor_applied):
+            if torch._has_same_tensorimpl_type(tensor, tensor_applied):
+                # If the new tensor has the same TensorImpl type as the existing tensor,
+                # the current behavior is to change the tensor in-place using `.data =`,
+                # and the future behavior is to overwrite the existing tensor. However,
+                # changing the current behavior is a BC-breaking change, and we want it
+                # to happen in future releases. So for now we introduce the
+                # `torch.__future__.get_overwrite_module_params_on_conversion()`
+                # global flag to let the user control whether they want the future
+                # behavior of overwriting the existing tensor or not.
+                return not torch.__future__.get_overwrite_module_params_on_conversion()
+            else:
+                return False
+
+        for key, param in self._parameters.items():
             if param is not None:
+                # Tensors stored in modules are graph leaves, and we don't want to
+                # track autograd history of `param_applied`, so we have to use
+                # `with torch.no_grad():`
                 with torch.no_grad():
                     param_applied = fn(param)
-                param.data = param_applied
-                if param._grad is not None:
+                should_use_set_data = compute_should_use_set_data(param, param_applied)
+                if should_use_set_data:
+                    param.data = param_applied
+                else:
+                    assert isinstance(param, Parameter)
+                    assert param.is_leaf
+                    self._parameters[key] = Parameter(param_applied, param.requires_grad)
+
+                if param.grad is not None:
                     with torch.no_grad():
-                        grad_applied = fn(param._grad)
-                    param._grad.data = grad_applied
+                        grad_applied = fn(param.grad)
+                    should_use_set_data = compute_should_use_set_data(param.grad, grad_applied)
+                    if should_use_set_data:
+                        param.grad.data = grad_applied
+                    else:
+                        assert param.grad.is_leaf
+                        self._parameters[key].grad = grad_applied.requires_grad_(param.grad.requires_grad)
 
         for key, buf in self._buffers.items():
             if buf is not None: