From 037a24b6c4225273b9bd1ba75da480fc42d23f99 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 24 Aug 2025 15:53:33 +0800 Subject: [PATCH 01/23] feat: add `device_name` classmethod in `Accelerator`. --- .../pytorch/accelerators/accelerator.py | 7 +++++- src/lightning/pytorch/accelerators/cuda.py | 6 +++++ src/lightning/pytorch/accelerators/mps.py | 7 ++++++ src/lightning/pytorch/accelerators/xla.py | 24 ++++++++++++++++++- src/lightning/pytorch/trainer/setup.py | 19 ++++++--------- src/lightning/pytorch/trainer/trainer.py | 11 ++++++--- 6 files changed, 57 insertions(+), 17 deletions(-) diff --git a/src/lightning/pytorch/accelerators/accelerator.py b/src/lightning/pytorch/accelerators/accelerator.py index 9238071178a80..b84d3fd1741eb 100644 --- a/src/lightning/pytorch/accelerators/accelerator.py +++ b/src/lightning/pytorch/accelerators/accelerator.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC -from typing import Any +from typing import Any, Optional import lightning.pytorch as pl from lightning.fabric.accelerators.accelerator import Accelerator as _Accelerator @@ -45,3 +45,8 @@ def get_device_stats(self, device: _DEVICE) -> dict[str, Any]: """ raise NotImplementedError + + @classmethod + def device_name(cls, device: Optional = None) -> str: + """Get the device name for a given device.""" + return str(cls.is_available()) diff --git a/src/lightning/pytorch/accelerators/cuda.py b/src/lightning/pytorch/accelerators/cuda.py index a00b12a85a8dd..4de1d6e09d96f 100644 --- a/src/lightning/pytorch/accelerators/cuda.py +++ b/src/lightning/pytorch/accelerators/cuda.py @@ -113,6 +113,12 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No description=cls.__name__, ) + @classmethod + def device_name(cls, device: Optional[torch.types.Device] = None) -> str: + if not cls.is_available(): + return "False" + return torch.cuda.get_device_name(device) + def get_nvidia_gpu_stats(device: _DEVICE) -> dict[str, float]: # pragma: no-cover """Get GPU stats including memory, fan speed, and temperature from nvidia-smi. diff --git a/src/lightning/pytorch/accelerators/mps.py b/src/lightning/pytorch/accelerators/mps.py index f7674989cc721..19bd06198aa1b 100644 --- a/src/lightning/pytorch/accelerators/mps.py +++ b/src/lightning/pytorch/accelerators/mps.py @@ -87,6 +87,13 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No description=cls.__name__, ) + @classmethod + def device_name(cls, device: Optional = None) -> str: + # todo: implement a better way to get the device name + available = cls.is_available() + gpu_type = " (mps)" if available else "" + return f"{available}{gpu_type}" + # device metrics _VM_PERCENT = "M1_vm_percent" diff --git a/src/lightning/pytorch/accelerators/xla.py b/src/lightning/pytorch/accelerators/xla.py index 10726b505448c..99c7faf43d38e 100644 --- a/src/lightning/pytorch/accelerators/xla.py +++ b/src/lightning/pytorch/accelerators/xla.py @@ -11,11 +11,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any +from typing import Any, Optional from typing_extensions import override from lightning.fabric.accelerators import _AcceleratorRegistry +from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1 from lightning.fabric.accelerators.xla import XLAAccelerator as FabricXLAAccelerator from lightning.fabric.utilities.types import _DEVICE from lightning.pytorch.accelerators.accelerator import Accelerator @@ -53,3 +54,24 @@ def get_device_stats(self, device: _DEVICE) -> dict[str, Any]: @override def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> None: accelerator_registry.register("tpu", cls, description=cls.__name__) + + @classmethod + def device_name(cls, device: Optional = None) -> str: + is_available = cls.is_available() + if not is_available: + return str(is_available) + + if _XLA_GREATER_EQUAL_2_1: + from torch_xla._internal import tpu + else: + from torch_xla.experimental import tpu + import torch_xla.core.xla_env_vars as xenv + from requests.exceptions import HTTPError + + try: + ret = tpu.get_tpu_env()[xenv.ACCELERATOR_TYPE] + except HTTPError: + # Fallback to "True" if HTTPError is raised during retrieving device information + ret = str(is_available) + + return ret diff --git a/src/lightning/pytorch/trainer/setup.py b/src/lightning/pytorch/trainer/setup.py index 00b546b252ac8..884cba401c429 100644 --- a/src/lightning/pytorch/trainer/setup.py +++ b/src/lightning/pytorch/trainer/setup.py @@ -142,21 +142,16 @@ def _init_profiler(trainer: "pl.Trainer", profiler: Optional[Union[Profiler, str def _log_device_info(trainer: "pl.Trainer") -> None: - if CUDAAccelerator.is_available(): - gpu_available = True - gpu_type = " (cuda)" - elif MPSAccelerator.is_available(): - gpu_available = True - gpu_type = " (mps)" + if isinstance(trainer.accelerator, (CUDAAccelerator, MPSAccelerator)): + gpu_used = trainer.num_devices + device_names = list({trainer.accelerator.device_name(d) for d in trainer.devices}) else: - gpu_available = False - gpu_type = "" - - gpu_used = isinstance(trainer.accelerator, (CUDAAccelerator, MPSAccelerator)) - rank_zero_info(f"GPU available: {gpu_available}{gpu_type}, used: {gpu_used}") + gpu_used = 0 + device_names = "False" + rank_zero_info(f"GPU available: {device_names}, using: {gpu_used} {'devices' if gpu_used else 'device'}.") num_tpu_cores = trainer.num_devices if isinstance(trainer.accelerator, XLAAccelerator) else 0 - rank_zero_info(f"TPU available: {XLAAccelerator.is_available()}, using: {num_tpu_cores} TPU cores") + rank_zero_info(f"TPU available: {XLAAccelerator.device_name()}, using: {num_tpu_cores} TPU cores") if _habana_available_and_importable(): from lightning_habana import HPUAccelerator diff --git a/src/lightning/pytorch/trainer/trainer.py b/src/lightning/pytorch/trainer/trainer.py index 21c75c3096125..efc409f859549 100644 --- a/src/lightning/pytorch/trainer/trainer.py +++ b/src/lightning/pytorch/trainer/trainer.py @@ -1187,16 +1187,21 @@ def num_nodes(self) -> int: return getattr(self.strategy, "num_nodes", 1) @property - def device_ids(self) -> list[int]: - """List of device indexes per node.""" + def devices(self) -> list[torch.device]: + """The devices the trainer uses per node.""" devices = ( self.strategy.parallel_devices if isinstance(self.strategy, ParallelStrategy) else [self.strategy.root_device] ) assert devices is not None + return devices + + @property + def device_ids(self) -> list[int]: + """List of device indexes per node.""" device_ids = [] - for idx, device in enumerate(devices): + for idx, device in enumerate(self.devices): if isinstance(device, torch.device): device_ids.append(device.index or idx) elif isinstance(device, int): From c1746b2f3d2ea77921aa396f318d67d70204a118 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 24 Aug 2025 16:53:42 +0800 Subject: [PATCH 02/23] feat: change to original device logic in setup. --- src/lightning/pytorch/trainer/setup.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/lightning/pytorch/trainer/setup.py b/src/lightning/pytorch/trainer/setup.py index 884cba401c429..8b0bc7cee87e2 100644 --- a/src/lightning/pytorch/trainer/setup.py +++ b/src/lightning/pytorch/trainer/setup.py @@ -142,13 +142,18 @@ def _init_profiler(trainer: "pl.Trainer", profiler: Optional[Union[Profiler, str def _log_device_info(trainer: "pl.Trainer") -> None: - if isinstance(trainer.accelerator, (CUDAAccelerator, MPSAccelerator)): - gpu_used = trainer.num_devices - device_names = list({trainer.accelerator.device_name(d) for d in trainer.devices}) + if CUDAAccelerator.is_available(): + if isinstance(trainer.accelerator, CUDAAccelerator): + device_name = list({CUDAAccelerator.device_name(d) for d in trainer.device_ids}) + else: + device_name = CUDAAccelerator.device_name() + elif MPSAccelerator.is_available(): + device_name = MPSAccelerator.device_name() else: - gpu_used = 0 - device_names = "False" - rank_zero_info(f"GPU available: {device_names}, using: {gpu_used} {'devices' if gpu_used else 'device'}.") + device_name = str(False) + + gpu_used = trainer.num_devices if isinstance(trainer.accelerator, (CUDAAccelerator, MPSAccelerator)) else 0 + rank_zero_info(f"GPU available: {device_name}, using: {gpu_used} devices.") num_tpu_cores = trainer.num_devices if isinstance(trainer.accelerator, XLAAccelerator) else 0 rank_zero_info(f"TPU available: {XLAAccelerator.device_name()}, using: {num_tpu_cores} TPU cores") From a383d79f018f8ebec81ee352061ddf519b45f0c2 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 24 Aug 2025 16:56:25 +0800 Subject: [PATCH 03/23] revert: revert changes in trainer. --- src/lightning/pytorch/trainer/trainer.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/src/lightning/pytorch/trainer/trainer.py b/src/lightning/pytorch/trainer/trainer.py index efc409f859549..21c75c3096125 100644 --- a/src/lightning/pytorch/trainer/trainer.py +++ b/src/lightning/pytorch/trainer/trainer.py @@ -1187,21 +1187,16 @@ def num_nodes(self) -> int: return getattr(self.strategy, "num_nodes", 1) @property - def devices(self) -> list[torch.device]: - """The devices the trainer uses per node.""" + def device_ids(self) -> list[int]: + """List of device indexes per node.""" devices = ( self.strategy.parallel_devices if isinstance(self.strategy, ParallelStrategy) else [self.strategy.root_device] ) assert devices is not None - return devices - - @property - def device_ids(self) -> list[int]: - """List of device indexes per node.""" device_ids = [] - for idx, device in enumerate(self.devices): + for idx, device in enumerate(devices): if isinstance(device, torch.device): device_ids.append(device.index or idx) elif isinstance(device, int): From 7db87936427ad453d37ffe9932224181bcc1cf43 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 24 Aug 2025 17:09:43 +0800 Subject: [PATCH 04/23] fix: fix type annotation. --- src/lightning/pytorch/accelerators/accelerator.py | 4 ++-- src/lightning/pytorch/accelerators/cuda.py | 2 +- src/lightning/pytorch/accelerators/mps.py | 2 +- src/lightning/pytorch/accelerators/xla.py | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/lightning/pytorch/accelerators/accelerator.py b/src/lightning/pytorch/accelerators/accelerator.py index b84d3fd1741eb..eb43c1dcb0fa8 100644 --- a/src/lightning/pytorch/accelerators/accelerator.py +++ b/src/lightning/pytorch/accelerators/accelerator.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC -from typing import Any, Optional +from typing import Any import lightning.pytorch as pl from lightning.fabric.accelerators.accelerator import Accelerator as _Accelerator @@ -47,6 +47,6 @@ def get_device_stats(self, device: _DEVICE) -> dict[str, Any]: raise NotImplementedError @classmethod - def device_name(cls, device: Optional = None) -> str: + def device_name(cls, device: _DEVICE = None) -> str: """Get the device name for a given device.""" return str(cls.is_available()) diff --git a/src/lightning/pytorch/accelerators/cuda.py b/src/lightning/pytorch/accelerators/cuda.py index 4de1d6e09d96f..0485bd5952951 100644 --- a/src/lightning/pytorch/accelerators/cuda.py +++ b/src/lightning/pytorch/accelerators/cuda.py @@ -114,7 +114,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No ) @classmethod - def device_name(cls, device: Optional[torch.types.Device] = None) -> str: + def device_name(cls, device: _DEVICE = None) -> str: if not cls.is_available(): return "False" return torch.cuda.get_device_name(device) diff --git a/src/lightning/pytorch/accelerators/mps.py b/src/lightning/pytorch/accelerators/mps.py index 19bd06198aa1b..2d0186e4be615 100644 --- a/src/lightning/pytorch/accelerators/mps.py +++ b/src/lightning/pytorch/accelerators/mps.py @@ -88,7 +88,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No ) @classmethod - def device_name(cls, device: Optional = None) -> str: + def device_name(cls, device: _DEVICE = None) -> str: # todo: implement a better way to get the device name available = cls.is_available() gpu_type = " (mps)" if available else "" diff --git a/src/lightning/pytorch/accelerators/xla.py b/src/lightning/pytorch/accelerators/xla.py index 99c7faf43d38e..5a6b90646662d 100644 --- a/src/lightning/pytorch/accelerators/xla.py +++ b/src/lightning/pytorch/accelerators/xla.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Optional +from typing import Any from typing_extensions import override @@ -56,7 +56,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No accelerator_registry.register("tpu", cls, description=cls.__name__) @classmethod - def device_name(cls, device: Optional = None) -> str: + def device_name(cls, device: _DEVICE = None) -> str: is_available = cls.is_available() if not is_available: return str(is_available) From 3a3949a1d2ad98fa6b743b1e510e1182113a0073 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 24 Aug 2025 21:48:16 +0800 Subject: [PATCH 05/23] fix: fix type annotation. --- src/lightning/pytorch/accelerators/accelerator.py | 4 ++-- src/lightning/pytorch/accelerators/cuda.py | 2 +- src/lightning/pytorch/accelerators/mps.py | 2 +- src/lightning/pytorch/accelerators/xla.py | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/lightning/pytorch/accelerators/accelerator.py b/src/lightning/pytorch/accelerators/accelerator.py index eb43c1dcb0fa8..bed925b1c013c 100644 --- a/src/lightning/pytorch/accelerators/accelerator.py +++ b/src/lightning/pytorch/accelerators/accelerator.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC -from typing import Any +from typing import Any, Optional import lightning.pytorch as pl from lightning.fabric.accelerators.accelerator import Accelerator as _Accelerator @@ -47,6 +47,6 @@ def get_device_stats(self, device: _DEVICE) -> dict[str, Any]: raise NotImplementedError @classmethod - def device_name(cls, device: _DEVICE = None) -> str: + def device_name(cls, device: Optional[_DEVICE] = None) -> str: """Get the device name for a given device.""" return str(cls.is_available()) diff --git a/src/lightning/pytorch/accelerators/cuda.py b/src/lightning/pytorch/accelerators/cuda.py index 0485bd5952951..0228b8b0d00f6 100644 --- a/src/lightning/pytorch/accelerators/cuda.py +++ b/src/lightning/pytorch/accelerators/cuda.py @@ -114,7 +114,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No ) @classmethod - def device_name(cls, device: _DEVICE = None) -> str: + def device_name(cls, device: Optional[_DEVICE] = None) -> str: if not cls.is_available(): return "False" return torch.cuda.get_device_name(device) diff --git a/src/lightning/pytorch/accelerators/mps.py b/src/lightning/pytorch/accelerators/mps.py index 2d0186e4be615..6532833e292c6 100644 --- a/src/lightning/pytorch/accelerators/mps.py +++ b/src/lightning/pytorch/accelerators/mps.py @@ -88,7 +88,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No ) @classmethod - def device_name(cls, device: _DEVICE = None) -> str: + def device_name(cls, device: Optional[_DEVICE] = None) -> str: # todo: implement a better way to get the device name available = cls.is_available() gpu_type = " (mps)" if available else "" diff --git a/src/lightning/pytorch/accelerators/xla.py b/src/lightning/pytorch/accelerators/xla.py index 5a6b90646662d..ac417ebef9d79 100644 --- a/src/lightning/pytorch/accelerators/xla.py +++ b/src/lightning/pytorch/accelerators/xla.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any +from typing import Any, Optional from typing_extensions import override @@ -56,7 +56,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No accelerator_registry.register("tpu", cls, description=cls.__name__) @classmethod - def device_name(cls, device: _DEVICE = None) -> str: + def device_name(cls, device: Optional[_DEVICE] = None) -> str: is_available = cls.is_available() if not is_available: return str(is_available) From 8cb809f2a7732acec9f822b33b3585c6724ea332 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 27 Aug 2025 00:10:52 +0800 Subject: [PATCH 06/23] add `override` decorator. --- src/lightning/pytorch/accelerators/cuda.py | 1 + src/lightning/pytorch/accelerators/xla.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/lightning/pytorch/accelerators/cuda.py b/src/lightning/pytorch/accelerators/cuda.py index 0228b8b0d00f6..777dbdcadcbd9 100644 --- a/src/lightning/pytorch/accelerators/cuda.py +++ b/src/lightning/pytorch/accelerators/cuda.py @@ -114,6 +114,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No ) @classmethod + @override def device_name(cls, device: Optional[_DEVICE] = None) -> str: if not cls.is_available(): return "False" diff --git a/src/lightning/pytorch/accelerators/xla.py b/src/lightning/pytorch/accelerators/xla.py index ac417ebef9d79..3aa7f1953bff2 100644 --- a/src/lightning/pytorch/accelerators/xla.py +++ b/src/lightning/pytorch/accelerators/xla.py @@ -56,6 +56,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No accelerator_registry.register("tpu", cls, description=cls.__name__) @classmethod + @override def device_name(cls, device: Optional[_DEVICE] = None) -> str: is_available = cls.is_available() if not is_available: From 944ad69ab813dd4428032586b99021135f2fcc08 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 27 Aug 2025 00:11:41 +0800 Subject: [PATCH 07/23] fix tests. --- tests/tests_pytorch/accelerators/test_gpu.py | 14 ++++++++++++++ tests/tests_pytorch/accelerators/test_mps.py | 11 +++++++++++ tests/tests_pytorch/accelerators/test_xla.py | 5 +++++ tests/tests_pytorch/conftest.py | 2 ++ 4 files changed, 32 insertions(+) diff --git a/tests/tests_pytorch/accelerators/test_gpu.py b/tests/tests_pytorch/accelerators/test_gpu.py index 5a71887e17eec..085eaff40dea0 100644 --- a/tests/tests_pytorch/accelerators/test_gpu.py +++ b/tests/tests_pytorch/accelerators/test_gpu.py @@ -68,3 +68,17 @@ def test_gpu_availability(): def test_warning_if_gpus_not_used(cuda_count_1): with pytest.warns(UserWarning, match="GPU available but not used"): Trainer(accelerator="cpu") + + +@RunIf(min_cuda_gpus=1) +def test_gpu_device_name(): + for i in range(torch.cuda.device_count()): + assert torch.cuda.get_device_name(i) == CUDAAccelerator.device_name(i) + + with torch.device("cuda:0"): + assert torch.cuda.get_device_name(0) == CUDAAccelerator.device_name() + + +def test_gpu_device_name_no_gpu(): + with mock.patch("torch.cuda.is_available", return_value=False): + assert str(False) == CUDAAccelerator.device_name() diff --git a/tests/tests_pytorch/accelerators/test_mps.py b/tests/tests_pytorch/accelerators/test_mps.py index c0a28840f0ef6..a9d1ae0f88c15 100644 --- a/tests/tests_pytorch/accelerators/test_mps.py +++ b/tests/tests_pytorch/accelerators/test_mps.py @@ -13,6 +13,7 @@ # limitations under the License. from collections import namedtuple +from unittest import mock import pytest import torch @@ -39,6 +40,16 @@ def test_mps_availability(): assert MPSAccelerator.is_available() +@RunIf(mps=True) +def test_mps_device_name(): + assert MPSAccelerator.device_name() == "True (mps)" + + +def test_mps_device_name_not_available(): + with mock.patch("torch.backends.mps.is_available", return_value=False): + assert MPSAccelerator.device_name() == "False" + + def test_warning_if_mps_not_used(mps_count_1): with pytest.warns(UserWarning, match="GPU available but not used"): Trainer(accelerator="cpu") diff --git a/tests/tests_pytorch/accelerators/test_xla.py b/tests/tests_pytorch/accelerators/test_xla.py index 83dace719371d..b9914cfe4d5cf 100644 --- a/tests/tests_pytorch/accelerators/test_xla.py +++ b/tests/tests_pytorch/accelerators/test_xla.py @@ -302,6 +302,11 @@ def test_warning_if_tpus_not_used(tpu_available): Trainer(accelerator="cpu") +@RunIf(tpu=True) +def test_tpu_device_name(): + assert XLAAccelerator.device_name() == "TPU" + + @pytest.mark.parametrize( ("devices", "expected_device_ids"), [ diff --git a/tests/tests_pytorch/conftest.py b/tests/tests_pytorch/conftest.py index 878298c6bfd94..9efb0f699ac26 100644 --- a/tests/tests_pytorch/conftest.py +++ b/tests/tests_pytorch/conftest.py @@ -182,6 +182,7 @@ def thread_police_duuu_daaa_duuu_daaa(): def mock_cuda_count(monkeypatch, n: int) -> None: monkeypatch.setattr(lightning.fabric.accelerators.cuda, "num_cuda_devices", lambda: n) monkeypatch.setattr(lightning.pytorch.accelerators.cuda, "num_cuda_devices", lambda: n) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda _: "Mocked CUDA Device") @pytest.fixture @@ -244,6 +245,7 @@ def mock_tpu_available(monkeypatch: pytest.MonkeyPatch, value: bool = True) -> N monkeypatch.setitem(sys.modules, "torch_xla", Mock()) monkeypatch.setitem(sys.modules, "torch_xla.core.xla_model", Mock()) monkeypatch.setitem(sys.modules, "torch_xla.experimental", Mock()) + monkeypatch.setattr(lightning.pytorch.accelerators.xla.XLAAccelerator, "device_name", lambda _: "TPU") @pytest.fixture From 1d9bd0de1f67504f19c62a2712d35e2680edd8cf Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 27 Aug 2025 00:27:56 +0800 Subject: [PATCH 08/23] fix tests. --- tests/tests_pytorch/accelerators/test_gpu.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/tests_pytorch/accelerators/test_gpu.py b/tests/tests_pytorch/accelerators/test_gpu.py index 085eaff40dea0..dbdd4a64f5b8e 100644 --- a/tests/tests_pytorch/accelerators/test_gpu.py +++ b/tests/tests_pytorch/accelerators/test_gpu.py @@ -79,6 +79,5 @@ def test_gpu_device_name(): assert torch.cuda.get_device_name(0) == CUDAAccelerator.device_name() -def test_gpu_device_name_no_gpu(): - with mock.patch("torch.cuda.is_available", return_value=False): - assert str(False) == CUDAAccelerator.device_name() +def test_gpu_device_name_no_gpu(cuda_count_0): + assert str(False) == CUDAAccelerator.device_name() From 92b1d692b80bb7bc623fd746bfe2fee17d9f0d91 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 27 Aug 2025 00:28:34 +0800 Subject: [PATCH 09/23] fix mypy and device string format. --- src/lightning/pytorch/trainer/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning/pytorch/trainer/setup.py b/src/lightning/pytorch/trainer/setup.py index 8b0bc7cee87e2..4a7d8e87374f5 100644 --- a/src/lightning/pytorch/trainer/setup.py +++ b/src/lightning/pytorch/trainer/setup.py @@ -144,7 +144,7 @@ def _init_profiler(trainer: "pl.Trainer", profiler: Optional[Union[Profiler, str def _log_device_info(trainer: "pl.Trainer") -> None: if CUDAAccelerator.is_available(): if isinstance(trainer.accelerator, CUDAAccelerator): - device_name = list({CUDAAccelerator.device_name(d) for d in trainer.device_ids}) + device_name = ", ".join(list({CUDAAccelerator.device_name(d) for d in trainer.device_ids})) else: device_name = CUDAAccelerator.device_name() elif MPSAccelerator.is_available(): From 0a5725b98fc929ed24bcd5d39a690f40287e97af Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 27 Aug 2025 01:17:07 +0800 Subject: [PATCH 10/23] fix tests. --- tests/tests_pytorch/conftest.py | 6 +++++- tests/tests_pytorch/plugins/test_cluster_integration.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/tests_pytorch/conftest.py b/tests/tests_pytorch/conftest.py index 9efb0f699ac26..f7f1d890f5103 100644 --- a/tests/tests_pytorch/conftest.py +++ b/tests/tests_pytorch/conftest.py @@ -245,7 +245,11 @@ def mock_tpu_available(monkeypatch: pytest.MonkeyPatch, value: bool = True) -> N monkeypatch.setitem(sys.modules, "torch_xla", Mock()) monkeypatch.setitem(sys.modules, "torch_xla.core.xla_model", Mock()) monkeypatch.setitem(sys.modules, "torch_xla.experimental", Mock()) - monkeypatch.setattr(lightning.pytorch.accelerators.xla.XLAAccelerator, "device_name", lambda _: "TPU") + monkeypatch.setattr( + lightning.pytorch.accelerators.xla.XLAAccelerator, + "device_name", + lambda *_: "Mocked TPU Device", + ) @pytest.fixture diff --git a/tests/tests_pytorch/plugins/test_cluster_integration.py b/tests/tests_pytorch/plugins/test_cluster_integration.py index 08bd1707b5cfd..f74b771199aaa 100644 --- a/tests/tests_pytorch/plugins/test_cluster_integration.py +++ b/tests/tests_pytorch/plugins/test_cluster_integration.py @@ -66,7 +66,7 @@ def test_ranks_available_manual_strategy_selection(_, strategy_cls): """Test that the rank information is readily available after Trainer initialization.""" num_nodes = 2 for cluster, variables, expected in environment_combinations(): - with mock.patch.dict(os.environ, variables): + with mock.patch.dict(os.environ, variables), mock.patch("torch.cuda.get_device_name", return_value="GPU"): strategy = strategy_cls( parallel_devices=[torch.device("cuda", 1), torch.device("cuda", 2)], cluster_environment=cluster ) From 5b11bdb469d5da5821693b78c7a8fbee390e8e67 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 27 Aug 2025 08:09:15 +0800 Subject: [PATCH 11/23] mps override decorator. --- src/lightning/pytorch/accelerators/mps.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/lightning/pytorch/accelerators/mps.py b/src/lightning/pytorch/accelerators/mps.py index 6532833e292c6..b8262911e1fc6 100644 --- a/src/lightning/pytorch/accelerators/mps.py +++ b/src/lightning/pytorch/accelerators/mps.py @@ -88,6 +88,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No ) @classmethod + @override def device_name(cls, device: Optional[_DEVICE] = None) -> str: # todo: implement a better way to get the device name available = cls.is_available() From 63a0f70993b12bf0bc53a6d0c30d082052318504 Mon Sep 17 00:00:00 2001 From: Jirka Borovec <6035284+Borda@users.noreply.github.com> Date: Wed, 27 Aug 2025 17:03:45 +0200 Subject: [PATCH 12/23] empty str --- src/lightning/pytorch/accelerators/cuda.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning/pytorch/accelerators/cuda.py b/src/lightning/pytorch/accelerators/cuda.py index 777dbdcadcbd9..bfc581201e506 100644 --- a/src/lightning/pytorch/accelerators/cuda.py +++ b/src/lightning/pytorch/accelerators/cuda.py @@ -117,7 +117,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No @override def device_name(cls, device: Optional[_DEVICE] = None) -> str: if not cls.is_available(): - return "False" + return "" return torch.cuda.get_device_name(device) From 8e91f8fd6a414db4f4d0a3030381bb7cea6853b5 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 27 Aug 2025 23:58:38 +0800 Subject: [PATCH 13/23] return empty string if accelerator is not available. --- src/lightning/pytorch/accelerators/mps.py | 6 +++--- src/lightning/pytorch/accelerators/xla.py | 2 +- src/lightning/pytorch/trainer/setup.py | 5 ++++- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/lightning/pytorch/accelerators/mps.py b/src/lightning/pytorch/accelerators/mps.py index b8262911e1fc6..7ad3c3304e84a 100644 --- a/src/lightning/pytorch/accelerators/mps.py +++ b/src/lightning/pytorch/accelerators/mps.py @@ -91,9 +91,9 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No @override def device_name(cls, device: Optional[_DEVICE] = None) -> str: # todo: implement a better way to get the device name - available = cls.is_available() - gpu_type = " (mps)" if available else "" - return f"{available}{gpu_type}" + if not cls.is_available(): + return "" + return "True (mps)" # device metrics diff --git a/src/lightning/pytorch/accelerators/xla.py b/src/lightning/pytorch/accelerators/xla.py index 3aa7f1953bff2..b248ab176d9f4 100644 --- a/src/lightning/pytorch/accelerators/xla.py +++ b/src/lightning/pytorch/accelerators/xla.py @@ -60,7 +60,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No def device_name(cls, device: Optional[_DEVICE] = None) -> str: is_available = cls.is_available() if not is_available: - return str(is_available) + return "" if _XLA_GREATER_EQUAL_2_1: from torch_xla._internal import tpu diff --git a/src/lightning/pytorch/trainer/setup.py b/src/lightning/pytorch/trainer/setup.py index 4a7d8e87374f5..846f2da0ae88b 100644 --- a/src/lightning/pytorch/trainer/setup.py +++ b/src/lightning/pytorch/trainer/setup.py @@ -156,7 +156,10 @@ def _log_device_info(trainer: "pl.Trainer") -> None: rank_zero_info(f"GPU available: {device_name}, using: {gpu_used} devices.") num_tpu_cores = trainer.num_devices if isinstance(trainer.accelerator, XLAAccelerator) else 0 - rank_zero_info(f"TPU available: {XLAAccelerator.device_name()}, using: {num_tpu_cores} TPU cores") + rank_zero_info( + f"TPU available: {XLAAccelerator.device_name() if XLAAccelerator.is_available() else str(False)}, " + f"using: {num_tpu_cores} TPU cores" + ) if _habana_available_and_importable(): from lightning_habana import HPUAccelerator From f124d55ebc24c5e7de99e77533885d6d2778621c Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 27 Aug 2025 23:59:13 +0800 Subject: [PATCH 14/23] fix: fix unittests. --- tests/tests_pytorch/accelerators/test_gpu.py | 2 +- tests/tests_pytorch/accelerators/test_mps.py | 2 +- tests/tests_pytorch/accelerators/test_xla.py | 10 +++++++++- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/tests_pytorch/accelerators/test_gpu.py b/tests/tests_pytorch/accelerators/test_gpu.py index dbdd4a64f5b8e..9312685847e1f 100644 --- a/tests/tests_pytorch/accelerators/test_gpu.py +++ b/tests/tests_pytorch/accelerators/test_gpu.py @@ -80,4 +80,4 @@ def test_gpu_device_name(): def test_gpu_device_name_no_gpu(cuda_count_0): - assert str(False) == CUDAAccelerator.device_name() + assert CUDAAccelerator.device_name() == "" diff --git a/tests/tests_pytorch/accelerators/test_mps.py b/tests/tests_pytorch/accelerators/test_mps.py index a9d1ae0f88c15..c1c267b2bc1e9 100644 --- a/tests/tests_pytorch/accelerators/test_mps.py +++ b/tests/tests_pytorch/accelerators/test_mps.py @@ -47,7 +47,7 @@ def test_mps_device_name(): def test_mps_device_name_not_available(): with mock.patch("torch.backends.mps.is_available", return_value=False): - assert MPSAccelerator.device_name() == "False" + assert MPSAccelerator.device_name() == "" def test_warning_if_mps_not_used(mps_count_1): diff --git a/tests/tests_pytorch/accelerators/test_xla.py b/tests/tests_pytorch/accelerators/test_xla.py index b9914cfe4d5cf..e156a22a35c3c 100644 --- a/tests/tests_pytorch/accelerators/test_xla.py +++ b/tests/tests_pytorch/accelerators/test_xla.py @@ -304,7 +304,15 @@ def test_warning_if_tpus_not_used(tpu_available): @RunIf(tpu=True) def test_tpu_device_name(): - assert XLAAccelerator.device_name() == "TPU" + from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1 + + if _XLA_GREATER_EQUAL_2_1: + from torch_xla._internal import tpu + else: + from torch_xla.experimental import tpu + import torch_xla.core.xla_env_vars as xenv + + assert XLAAccelerator.device_name() == tpu.get_tpu_env()[xenv.ACCELERATOR_TYPE] @pytest.mark.parametrize( From 9fc30dabd0b72d4d4774a27a97fe050b3928810e Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 15 Oct 2025 23:06:52 +0800 Subject: [PATCH 15/23] implement MPSAccelerator.device_name --- src/lightning/pytorch/accelerators/mps.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/lightning/pytorch/accelerators/mps.py b/src/lightning/pytorch/accelerators/mps.py index 7ad3c3304e84a..5c1f9c990f919 100644 --- a/src/lightning/pytorch/accelerators/mps.py +++ b/src/lightning/pytorch/accelerators/mps.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import subprocess from typing import Any, Optional, Union import torch @@ -90,10 +91,19 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No @classmethod @override def device_name(cls, device: Optional[_DEVICE] = None) -> str: - # todo: implement a better way to get the device name if not cls.is_available(): return "" - return "True (mps)" + try: + result = subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, + text=True, + check=True, + ) + result_str = result.stdout.strip() + except subprocess.SubprocessError: + result_str = "True (mps)" + return result_str # device metrics From fcd8dc93d0c8c2df5e25b3f65c703cf7785ab254 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 15 Oct 2025 23:33:08 +0800 Subject: [PATCH 16/23] fix test. --- src/lightning/pytorch/accelerators/mps.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning/pytorch/accelerators/mps.py b/src/lightning/pytorch/accelerators/mps.py index 5c1f9c990f919..8170dfef5b4f0 100644 --- a/src/lightning/pytorch/accelerators/mps.py +++ b/src/lightning/pytorch/accelerators/mps.py @@ -101,7 +101,7 @@ def device_name(cls, device: Optional[_DEVICE] = None) -> str: check=True, ) result_str = result.stdout.strip() - except subprocess.SubprocessError: + except (subprocess.SubprocessError, FileNotFoundError): result_str = "True (mps)" return result_str From a1bf2eeaf64d083ee049f632ffdd1891fba6c624 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Wed, 15 Oct 2025 23:45:03 +0800 Subject: [PATCH 17/23] fix test. --- src/lightning/pytorch/accelerators/mps.py | 26 +++++++++++++---------- tests/tests_pytorch/conftest.py | 3 +++ 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/src/lightning/pytorch/accelerators/mps.py b/src/lightning/pytorch/accelerators/mps.py index 8170dfef5b4f0..a94ea639393b2 100644 --- a/src/lightning/pytorch/accelerators/mps.py +++ b/src/lightning/pytorch/accelerators/mps.py @@ -93,17 +93,21 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No def device_name(cls, device: Optional[_DEVICE] = None) -> str: if not cls.is_available(): return "" - try: - result = subprocess.run( - ["sysctl", "-n", "machdep.cpu.brand_string"], - capture_output=True, - text=True, - check=True, - ) - result_str = result.stdout.strip() - except (subprocess.SubprocessError, FileNotFoundError): - result_str = "True (mps)" - return result_str + return _get_mps_device_name() + + +def _get_mps_device_name() -> str: + try: + result = subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, + text=True, + check=True, + ) + result_str = result.stdout.strip() + except subprocess.SubprocessError: + result_str = "True (mps)" + return result_str # device metrics diff --git a/tests/tests_pytorch/conftest.py b/tests/tests_pytorch/conftest.py index f7f1d890f5103..9000a6bf8928a 100644 --- a/tests/tests_pytorch/conftest.py +++ b/tests/tests_pytorch/conftest.py @@ -208,6 +208,9 @@ def cuda_count_4(monkeypatch): def mock_mps_count(monkeypatch, n: int) -> None: monkeypatch.setattr(lightning.fabric.accelerators.mps, "_get_all_available_mps_gpus", lambda: [0] if n > 0 else []) monkeypatch.setattr(lightning.fabric.accelerators.mps.MPSAccelerator, "is_available", lambda *_: n > 0) + monkeypatch.setattr( + lightning.pytorch.accelerators.mps, "_get_mps_device_name", lambda: "Mocked MPS Device" if n > 0 else "" + ) @pytest.fixture From 447eb326b5aab464eaa80d81d6ba6a2f02e59f7d Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Thu, 16 Oct 2025 00:53:06 +0800 Subject: [PATCH 18/23] chlog --- src/lightning/pytorch/CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 1bba5e4ca0da7..e976ac1f80576 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/src/lightning/pytorch/CHANGELOG.md @@ -25,6 +25,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Added support for variable batch size in `ThroughputMonitor` ([#20236](https://github.com/Lightning-AI/pytorch-lightning/pull/20236)) +- Added classmethod `Accelerator.device_name` so that users could see device information during setup. ([#21112](https://github.com/Lightning-AI/pytorch-lightning/pull/21112)) + + ### Changed - Default to `RichProgressBar` and `RichModelSummary` if the rich package is available. Fallback to TQDMProgressBar and ModelSummary otherwise ([#20896](https://github.com/Lightning-AI/pytorch-lightning/pull/20896)) From e356426726be0d7482c8bd1794ed528cc04c8382 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Fri, 7 Nov 2025 01:36:12 +0800 Subject: [PATCH 19/23] add unittests --- tests/tests_pytorch/accelerators/test_mps.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/tests_pytorch/accelerators/test_mps.py b/tests/tests_pytorch/accelerators/test_mps.py index c1c267b2bc1e9..242b21bde7fb7 100644 --- a/tests/tests_pytorch/accelerators/test_mps.py +++ b/tests/tests_pytorch/accelerators/test_mps.py @@ -13,6 +13,7 @@ # limitations under the License. from collections import namedtuple +from subprocess import SubprocessError from unittest import mock import pytest @@ -21,6 +22,7 @@ import tests_pytorch.helpers.pipelines as tpipes from lightning.pytorch import Trainer from lightning.pytorch.accelerators import MPSAccelerator +from lightning.pytorch.accelerators.mps import _get_mps_device_name from lightning.pytorch.demos.boring_classes import BoringModel from tests_pytorch.helpers.runif import RunIf @@ -151,3 +153,21 @@ def to(self, *args, **kwargs): batch = trainer.strategy.batch_to_device(CustomBatchType(), torch.device("mps")) assert batch.a.type() == "torch.mps.FloatTensor" + + +@mock.patch("lightning.pytorch.accelerators.mps.subprocess.run") +def test_get_mps_device_name(mock_run): + mock_stdout = mock.MagicMock() + mock_stdout.configure_mock(stdout="Apple M1 Pro\n") + + mock_run.return_value = mock_stdout + device_name = _get_mps_device_name() + assert device_name == "Apple M1 Pro" + + +@mock.patch( + "lightning.pytorch.accelerators.mps.subprocess.run", + side_effect=SubprocessError("test"), +) +def test_get_mps_device_name_exception(mock_run): + assert _get_mps_device_name() == "True (mps)" From 7bf3645a68d31f23890de589c5975c05b6de2470 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 8 Nov 2025 14:44:25 +0800 Subject: [PATCH 20/23] add unittests --- tests/tests_pytorch/accelerators/test_xla.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/tests_pytorch/accelerators/test_xla.py b/tests/tests_pytorch/accelerators/test_xla.py index 973cf9dc054a9..52f5df0dda6f1 100644 --- a/tests/tests_pytorch/accelerators/test_xla.py +++ b/tests/tests_pytorch/accelerators/test_xla.py @@ -313,6 +313,24 @@ def test_tpu_device_name(): assert XLAAccelerator.device_name() == tpu.get_tpu_env()[xenv.ACCELERATOR_TYPE] +def test_tpu_device_name_exception(tpu_available, monkeypatch): + from requests.exceptions import HTTPError + + monkeypatch.delattr( + lightning.pytorch.accelerators.xla.XLAAccelerator, + "device_name", + raising=False, + ) + + mock.patch( + "torch_xla._internal.tpu", + "get_tpu_env", + side_effect=HTTPError("Could not fetch TPU device name"), + ) + + assert XLAAccelerator.device_name() == "True" + + @pytest.mark.parametrize( ("devices", "expected_device_ids"), [ From 1c064e877198ac344470a0f71e787c0b90d2e55e Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 8 Nov 2025 14:56:06 +0800 Subject: [PATCH 21/23] fix test --- tests/tests_pytorch/accelerators/test_xla.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/tests_pytorch/accelerators/test_xla.py b/tests/tests_pytorch/accelerators/test_xla.py index 52f5df0dda6f1..91bc6bc21dc70 100644 --- a/tests/tests_pytorch/accelerators/test_xla.py +++ b/tests/tests_pytorch/accelerators/test_xla.py @@ -327,6 +327,11 @@ def test_tpu_device_name_exception(tpu_available, monkeypatch): "get_tpu_env", side_effect=HTTPError("Could not fetch TPU device name"), ) + mock.patch( + "torch_xla.experimental.tpu", + "get_tpu_env", + side_effect=HTTPError("Could not fetch TPU device name"), + ) assert XLAAccelerator.device_name() == "True" From 0ad171e84526c9075802606dd8f630975bad8a29 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 8 Nov 2025 15:16:14 +0800 Subject: [PATCH 22/23] fix test --- tests/tests_pytorch/accelerators/test_xla.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/tests_pytorch/accelerators/test_xla.py b/tests/tests_pytorch/accelerators/test_xla.py index 91bc6bc21dc70..f2d7175c382f0 100644 --- a/tests/tests_pytorch/accelerators/test_xla.py +++ b/tests/tests_pytorch/accelerators/test_xla.py @@ -317,7 +317,7 @@ def test_tpu_device_name_exception(tpu_available, monkeypatch): from requests.exceptions import HTTPError monkeypatch.delattr( - lightning.pytorch.accelerators.xla.XLAAccelerator, + XLAAccelerator, "device_name", raising=False, ) From 8bf8817347b535c3251485c2414eaa5eacfa81c2 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 19 Jul 2026 14:00:40 +0800 Subject: [PATCH 23/23] test: add cuda name test in multi gpu --- tests/tests_pytorch/accelerators/test_gpu.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/tests_pytorch/accelerators/test_gpu.py b/tests/tests_pytorch/accelerators/test_gpu.py index 9312685847e1f..02dd4bfdf2d97 100644 --- a/tests/tests_pytorch/accelerators/test_gpu.py +++ b/tests/tests_pytorch/accelerators/test_gpu.py @@ -78,6 +78,10 @@ def test_gpu_device_name(): with torch.device("cuda:0"): assert torch.cuda.get_device_name(0) == CUDAAccelerator.device_name() + if torch.cuda.device_count() > 1: + with torch.device("cuda:1"): + assert torch.cuda.get_device_name(1) == CUDAAccelerator.device_name(1) + def test_gpu_device_name_no_gpu(cuda_count_0): assert CUDAAccelerator.device_name() == ""