From f1c697bc29dafacd0ccc0947fa766a9dd21ceac4 Mon Sep 17 00:00:00 2001 From: Vijay Anand Raghava Kanakagiri Date: Tue, 28 Apr 2026 18:07:17 +0000 Subject: [PATCH 1/2] fix: Fix _get_accelerator_memory --- test/distributed/tensor/debug/test_debug_mode.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/test/distributed/tensor/debug/test_debug_mode.py b/test/distributed/tensor/debug/test_debug_mode.py index eb20889f2a04..cc854d063bf0 100644 --- a/test/distributed/tensor/debug/test_debug_mode.py +++ b/test/distributed/tensor/debug/test_debug_mode.py @@ -53,6 +53,13 @@ def _get_accelerator_memory(): try: + if not torch.accelerator.is_available(): + return 0 + device = torch.accelerator.current_accelerator() + if not hasattr(torch, device.type) or not hasattr( + getattr(torch, device.type), "mem_get_info" + ): + return 0 return torch.accelerator.get_memory_info(0)[1] except (NotImplementedError): return 0 # Return 0, as that would help skip the test is not skipped From ddffcb83a9a1e88fc95714d5ba634e81031d6a1d Mon Sep 17 00:00:00 2001 From: Vijay Anand Raghava Kanakagiri Date: Wed, 29 Apr 2026 22:00:29 +0000 Subject: [PATCH 2/2] fix: address _get_accelerator_memory --- .../tensor/debug/test_debug_mode.py | 30 +++++++++---------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/test/distributed/tensor/debug/test_debug_mode.py b/test/distributed/tensor/debug/test_debug_mode.py index cc854d063bf0..e3030bd5982f 100644 --- a/test/distributed/tensor/debug/test_debug_mode.py +++ b/test/distributed/tensor/debug/test_debug_mode.py @@ -52,17 +52,14 @@ @requires_accelerator def _get_accelerator_memory(): - try: - if not torch.accelerator.is_available(): - return 0 - device = torch.accelerator.current_accelerator() - if not hasattr(torch, device.type) or not hasattr( - getattr(torch, device.type), "mem_get_info" - ): - return 0 - return torch.accelerator.get_memory_info(0)[1] - except (NotImplementedError): - return 0 # Return 0, as that would help skip the test is not skipped + if not torch.accelerator.is_available(): + return -1 + device = torch.accelerator.current_accelerator() + if not hasattr(torch, device.type) or not hasattr( + getattr(torch, device.type), "mem_get_info" + ): + return -1 + return torch.accelerator.get_memory_info(0)[1] class TestDTensorDebugMode(TestCase): def tearDown(self): @@ -924,12 +921,13 @@ def test_check_hash_mismatches(self): [call["call"] for call in mismatches], ["aten::sin", "aten::sum"] ) - @unittest.skipIf( - not torch.accelerator.is_available() - or _get_accelerator_memory() < 2**26, - "Being conservative, test peak memory is 25MB?", - ) def test_tensor_hash_redistribute(self): + + mem = _get_accelerator_memory() + if mem ==-1: + self.skipTest("No accelerator available or memory query not supported") + if mem < 2**26: + self.skipTest("Requires accelerator with at least 64MB memory") # test that hashing collectives gives correct results mesh = DeviceMesh(self.device_type, list(range(self.world_size)))