make big model inference compatible with torch.compile

src/accelerate/utils/operations.py

@@ -145,7 +145,7 @@ def send_to_device(tensor, device, non_blocking=False, skip_keys=None):
  Returns:
  The same data structure as `tensor` with all tensors sent to the proper device.
  """
- if is_torch_tensor(tensor) or hasattr(tensor, "to"):
+ if is_torch_tensor(tensor):
  # `torch.Tensor.to("npu")` could not find context when called for the first time (see this [issue](https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue)).
  if device == "npu":
  device = "npu:0"

tests/test_utils.py

@@ -370,9 +370,34 @@ def test_slice_and_concatenate(self):
  # We should expect there to be 66 items now
  assert result.shape == torch.Size([66, 4, 4])
 
- def test_send_to_device_compiles(self):
+ def test_send_to_device_compile(self):
  compiled_send_to_device = torch.compile(send_to_device, fullgraph=True)
- compiled_send_to_device(torch.zeros([1], dtype=torch.bfloat16), "cpu")
+ tensor = torch.randn(5, 2)
+ device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+ result1 = compiled_send_to_device(tensor, device)
+ assert torch.equal(result1.cpu(), tensor)
+
+ result2 = compiled_send_to_device((tensor, [tensor, tensor], 1), device)
+ assert isinstance(result2, tuple)
+ assert isinstance(result2[1], list)
+ assert result2[2] == 1
+
+ result2 = compiled_send_to_device({"a": tensor, "b": [tensor, tensor], "c": 1}, device)
+ assert isinstance(result2, dict)
+ assert isinstance(result2["b"], list)
+ assert result2["c"] == 1
+
+ # result3 = compiled_send_to_device(ExampleNamedTuple(a=tensor, b=[tensor, tensor], c=1), device)
+ # assert isinstance(result3, ExampleNamedTuple)
+ # assert isinstance(result3.b, list)
+ # assert result3.c == 1
+
+ # result4 = compiled_send_to_device(UserDict({"a": tensor, "b": [tensor, tensor], "c": 1}), device)
+ # assert isinstance(result4, UserDict)
+ # assert isinstance(result4["b"], list)
+
+ assert compiled_send_to_device
 
  def test_convert_to_fp32(self):
  compiled_convert_to_fp32 = torch.compile(convert_to_fp32, fullgraph=True)