Add ModuleInfo-based CPU / GPU parity tests (pytorch#68097)

Summary:
Continuation of pytorch#64694; fixes issues with the diff there

Pull Request resolved: pytorch#68097

Reviewed By: mruberry

Differential Revision: D32300650

Pulled By: jbschlosser

fbshipit-source-id: f3a5e72b019d4eddd7202854999eab61fffc9006

test/test_modules.py

@@ -6,7 +6,8 @@
 import tempfile
 
 import torch
-from torch.testing._internal.common_device_type import instantiate_device_type_tests
+from torch.testing._internal.common_device_type import (
+    instantiate_device_type_tests, onlyCUDA, toleranceOverride, tol)
 from torch.testing._internal.common_modules import module_db, modules
 from torch.testing._internal.common_utils import (
     TestCase, run_tests, freeze_rng_state, mock_wrapper, get_tensors_from, gradcheck, gradgradcheck)
@@ -369,6 +370,72 @@ def test_grad(self, device, dtype, module_info):
     def test_gradgrad(self, device, dtype, module_info):
         self._test_gradients_helper(device, dtype, module_info, gradgradcheck)
 
+    @onlyCUDA
+    @toleranceOverride({torch.float32: tol(5e-2, 0),
+                        torch.float64: tol(4e-4, 0)})
+    @modules(module_db)
+    def test_cpu_gpu_parity(self, device, dtype, module_info):
+        # Test cpu and gpu results are the same
+        module_cls = module_info.module_cls
+        module_inputs_cpu = module_info.module_inputs_func(module_info, device="cpu", dtype=dtype,
+                                                           requires_grad=True)
+
+        def _to_device(obj):
+            if isinstance(obj, torch.Tensor):
+                res = obj.detach().to(device=device)
+                res.requires_grad = obj.requires_grad
+                return res
+            elif isinstance(obj, tuple):
+                return tuple(_to_device(o) for o in obj)
+            elif isinstance(obj, dict):
+                return {key: _to_device(o) for key, o in obj.items()}
+            else:
+                return deepcopy(obj)
+
+        for module_input in module_inputs_cpu:
+
+            # === Move input from cpu to device ===
+            cpu_forward_args = module_input.forward_input.args
+            cpu_forward_kwargs = module_input.forward_input.kwargs
+
+            gpu_forward_args, gpu_forward_kwargs = _to_device((cpu_forward_args, cpu_forward_kwargs))
+
+            self._retain_grad((cpu_forward_args, cpu_forward_kwargs, gpu_forward_args, gpu_forward_kwargs))
+
+            # === Construct module on cpu and gpu ===
+            args, kwargs = module_input.constructor_input.args, module_input.constructor_input.kwargs
+
+            cpu_module = module_cls(*args, **kwargs).to(dtype).to("cpu")
+            gpu_module = module_cls(*args, **kwargs).to(dtype).to(device)
+
+            for cpu_p, gpu_p in zip(cpu_module.parameters(), gpu_module.parameters()):
+                gpu_p.data.copy_(cpu_p)
+
+            # === Compare forward output between cpu and gpu ===
+            cpu_output = cpu_module(*cpu_forward_args, **cpu_forward_kwargs)
+            gpu_output = gpu_module(*gpu_forward_args, **gpu_forward_kwargs)
+
+            self.assertEqual(cpu_output, gpu_output)
+
+            # === Run backwards on CPU and GPU and compare results ===
+            for _ in range(5):
+                cpu_grad_output = cpu_output.clone().normal_()
+                gpu_grad_output = cpu_grad_output.type_as(gpu_output)
+
+                cpu_output.backward(cpu_grad_output, retain_graph=True)
+                gpu_output.backward(gpu_grad_output, retain_graph=True)
+
+                cpu_grad_input = self._get_grads(cpu_forward_args)
+                gpu_grad_input = self._get_grads(gpu_forward_args)
+                self.assertEqual(cpu_grad_input, gpu_grad_input)
+
+                for cpu_p, gpu_p in zip(cpu_module.parameters(), gpu_module.parameters()):
+                    self.assertEqual(cpu_p.grad, gpu_p.grad)
+
+                cpu_grad_kwarg_input = self._get_grads(cpu_forward_kwargs)
+                gpu_grad_kwarg_input = self._get_grads(gpu_forward_kwargs)
+                self.assertEqual(cpu_grad_kwarg_input, gpu_grad_kwarg_input)
+
 
 instantiate_device_type_tests(TestModule, globals())