From c4dd142cafcd26dece048f125d514a92678212f0 Mon Sep 17 00:00:00 2001 From: Marc Sun Date: Tue, 2 Apr 2024 19:16:59 +0200 Subject: [PATCH 1/3] add tests --- tests/test_utils.py | 29 +++++++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 063bf262036..792f14f7c90 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -327,9 +327,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) From 24493a2c407c15faca2cecfa1596c78315268ef0 Mon Sep 17 00:00:00 2001 From: Marc Sun Date: Wed, 17 Apr 2024 14:43:38 +0200 Subject: [PATCH 2/3] make it work but without namedtuple --- src/accelerate/utils/operations.py | 10 +++++----- tests/test_utils.py | 16 ++++++++-------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py index e2456a10710..8647555bfda 100644 --- a/src/accelerate/utils/operations.py +++ b/src/accelerate/utils/operations.py @@ -75,10 +75,10 @@ def honor_type(obj, generator): Cast a generator to the same type as obj (list, tuple, or namedtuple) """ # Some objects may not be able to instantiate from a generator directly - if is_namedtuple(obj): - return type(obj)(*list(generator)) - else: - return type(obj)(generator) + # if is_namedtuple(obj): + # return type(obj)(*list(generator)) + # else: + return type(obj)(generator) def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_other_type=False, **kwargs): @@ -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" diff --git a/tests/test_utils.py b/tests/test_utils.py index 792f14f7c90..1c379750473 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -345,14 +345,14 @@ def test_send_to_device_compile(self): 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) + # 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 From c30cca03e1b9e53f5e2391da884a86babbeb497c Mon Sep 17 00:00:00 2001 From: Marc Sun Date: Thu, 8 Aug 2024 17:38:08 +0200 Subject: [PATCH 3/3] tuple works ! --- src/accelerate/utils/operations.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py index 1b450d1ffee..e6dd3fac135 100644 --- a/src/accelerate/utils/operations.py +++ b/src/accelerate/utils/operations.py @@ -75,10 +75,10 @@ def honor_type(obj, generator): Cast a generator to the same type as obj (list, tuple, or namedtuple) """ # Some objects may not be able to instantiate from a generator directly - # if is_namedtuple(obj): - # return type(obj)(*list(generator)) - # else: - return type(obj)(generator) + if is_namedtuple(obj): + return type(obj)(*list(generator)) + else: + return type(obj)(generator) def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_other_type=False, **kwargs):