Make cpu tensor on XLA dynamo backend a warning instead of error (#5549)

Author: JackCaoG

test/dynamo/test_dynamo.py

@@ -12,6 +12,9 @@
 import torch._dynamo as dynamo
 import torchvision
 import unittest
+import warnings
+
+torch_xla._XLAC._init_computation_client()
 
 # Setup import folders.
 xla_test_folder = os.path.dirname(os.path.dirname(os.path.abspath(sys.argv[0])))
@@ -58,36 +61,6 @@ def test_random_op_different_result_each_run(self):
     self.assertFalse(torch.allclose(dynamo_res_2, dynamo_res_3))
 
 
-class DynamErrorMessageTest(unittest.TestCase):
-
-  def test_cpu_tensor(self):
-    device = xm.xla_device()
-    input = torch.randn(4, 3, 224, 224)
-    input_xla = input.clone().to(device)
-    resnet18 = torchvision.models.resnet18()
-    resnet18.eval()
-    xla_resnet18 = torchvision.models.resnet18()
-    xla_resnet18.to(device)
-    xla_resnet18.eval()
-    dynamo_resnet18 = torch.compile(xla_resnet18, backend='openxla')
-    dynamo_resnet18_cpu = torch.compile(resnet18, backend='openxla')
-    # input on cpu and model weight on xla
-    with self.assertRaises(Exception) as context:
-      res = dynamo_resnet18(input)
-    self.assertTrue(
-        'found two different devices' in context.exception.__str__())
-    # input on xla and model weight on cpu
-    with self.assertRaises(Exception) as context:
-      res = dynamo_resnet18_cpu(input_xla)
-    self.assertTrue(
-        'found two different devices' in context.exception.__str__())
-    # input and model weight on cpu
-    with self.assertRaises(Exception) as context:
-      res = dynamo_resnet18_cpu(input)
-    self.assertTrue(
-        'please move all tensors to XLA device' in context.exception.__str__())
-
-
 class DynamoInferenceBasicTest(unittest.TestCase):
 
   @classmethod
@@ -516,6 +489,47 @@ def test_resnet18(self):
         met.metric_data('RunCachedGraphOutputData')[0], sample_count * 3)
 
 
+class DynamErrorMessageTest(unittest.TestCase):
+
+  def test_mixed_cpu_tensor(self):
+    device = xm.xla_device()
+    input = torch.randn(4, 3, 224, 224)
+    input_xla = input.clone().to(device)
+    resnet18 = torchvision.models.resnet18()
+    resnet18.eval()
+    xla_resnet18 = torchvision.models.resnet18()
+    xla_resnet18.to(device)
+    xla_resnet18.eval()
+    dynamo_resnet18 = torch.compile(xla_resnet18, backend='openxla')
+    dynamo_resnet18_cpu = torch.compile(resnet18, backend='openxla')
+    # input on cpu and model weight on xla
+    with self.assertRaises(Exception) as context:
+      res = dynamo_resnet18(input)
+    self.assertTrue(
+        'found two different devices' in context.exception.__str__())
+    # input on xla and model weight on cpu
+    with self.assertRaises(Exception) as context:
+      res = dynamo_resnet18_cpu(input_xla)
+    self.assertTrue(
+        'found two different devices' in context.exception.__str__())
+
+  def test_all_cpu_tensor(self):
+    met.clear_all()
+    input = torch.randn(4, 3, 224, 224)
+    resnet18 = torchvision.models.resnet18()
+    resnet18.eval()
+    dynamo_resnet18_cpu = torch.compile(resnet18, backend='openxla')
+    # input and model weight on cpu
+    with warnings.catch_warnings(record=True) as w:
+      res = dynamo_resnet18_cpu(input)
+      # there should be 18 paramters + 1 input
+      self.assertGreater(len(w), 15)
+      self.assertIn('Found tensor with shape torch.Size', str(w[0].message))
+    # no XLA operation should happens. Partitioner should offload all CPU
+    # ops to CPU.
+    self.assertEqual(len(met.counter_names()), 0)
+
+
 if __name__ == '__main__':
   test = unittest.main()
   sys.exit(0 if test.result.wasSuccessful() else 1)

torch_xla/core/dynamo_bridge.py

@@ -1,6 +1,7 @@
 import copy
 import dataclasses
 import operator
+import warnings
 
 import functools
 import itertools
@@ -457,8 +458,10 @@ def extract_compiled_graph(xla_model: torch.fx.GraphModule, xla_args):
 
   for xla_arg in xla_args:
     if xla_arg.device.type != 'xla':
-      raise RuntimeError(
-          'For openxla dynamo backend, please move all tensors to XLA device')
+      warnings.warn(
+          "Found tensor with shape " + str(xla_arg.size()) + " on " +
+          str(xla_arg.device) +
+          ". Please move all tensors to xla device to execute on XLA device.")
 
   cloned_args = [
       torch.clone(xla_arg) if isinstance(xla_arg, torch.Tensor) else xla_arg