From 5a0c1f7911820a520f9dfbfe24dd395004dd32b0 Mon Sep 17 00:00:00 2001 From: Nicki Skafte Date: Tue, 5 Aug 2025 11:37:17 +0200 Subject: [PATCH 1/5] add support for more precision + add warning on unsupported --- .../utilities/model_summary/model_summary.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary.py b/src/lightning/pytorch/utilities/model_summary/model_summary.py index 98d74ff63ea5f..546bbe21f2981 100644 --- a/src/lightning/pytorch/utilities/model_summary/model_summary.py +++ b/src/lightning/pytorch/utilities/model_summary/model_summary.py @@ -26,6 +26,7 @@ from torch.utils.hooks import RemovableHandle import lightning.pytorch as pl +from lightning.fabric.utilities import rank_zero_warn from lightning.fabric.utilities.distributed import _is_dtensor from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_4 from lightning.pytorch.utilities.model_helpers import _ModuleMode @@ -227,7 +228,22 @@ def __init__(self, model: "pl.LightningModule", max_depth: int = 1) -> None: self._layer_summary = self.summarize() # 1 byte -> 8 bits # TODO: how do we compute precision_megabytes in case of mixed precision? - precision_to_bits = {"64": 64, "32": 32, "16": 16, "bf16": 16} + precision_to_bits = { + "64": 64, + "32": 32, + "16": 16, + "bf16": 16, + "16-true": 16, + "bf16-true": 16, + "32-true": 32, + "64-true": 64, + } + if self._model.trainer is not None and self._model.trainer.precision not in precision_to_bits: + rank_zero_warn( + f"Precision {self._model.trainer.precision} is not supported by the model summary. " + " Estimated model size in MB will not be accurate. Using 32 bits instead.", + category=UserWarning, + ) precision = precision_to_bits.get(self._model.trainer.precision, 32) if self._model._trainer else 32 self._precision_megabytes = (precision / 8.0) * 1e-6 From b4d7d1e87b3d511986bc5b2cf19d81871fcb67de Mon Sep 17 00:00:00 2001 From: Nicki Skafte Date: Tue, 5 Aug 2025 11:37:33 +0200 Subject: [PATCH 2/5] add testing --- .../utilities/test_model_summary.py | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/tests_pytorch/utilities/test_model_summary.py b/tests/tests_pytorch/utilities/test_model_summary.py index 85825b5ea749d..8f8be116e3d0a 100644 --- a/tests/tests_pytorch/utilities/test_model_summary.py +++ b/tests/tests_pytorch/utilities/test_model_summary.py @@ -324,19 +324,33 @@ def test_empty_model_size(max_depth): pytest.param("mps", marks=RunIf(mps=True)), ], ) -def test_model_size_precision(tmp_path, accelerator): - """Test model size for half and full precision.""" - model = PreCalculatedModel() +@pytest.mark.parametrize("precision", ["16-true", "32-true", "64-true"]) +def test_model_size_precision(tmp_path, accelerator, precision): + """Test model size for different precision types.""" + model = PreCalculatedModel(precision=int(precision.split("-")[0])) # fit model trainer = Trainer( - default_root_dir=tmp_path, accelerator=accelerator, devices=1, max_steps=1, max_epochs=1, precision=32 + default_root_dir=tmp_path, accelerator=accelerator, devices=1, max_steps=1, max_epochs=1, precision=precision ) trainer.fit(model) summary = summarize(model) assert model.pre_calculated_model_size == summary.model_size +def test_model_size_warning_on_unsupported_precision(): + """Test that a warning is raised when the precision is not supported.""" + model = PreCalculatedModel(precision=32) # fallback to 32 bits + + # supported precision by lightning but not by the model summary + trainer = Trainer(max_epochs=1, precision="16-mixed") + trainer.fit(model) + + with pytest.warns(UserWarning, match="Precision 16-mixed is not supported by the model summary.*"): + summary = summarize(model) + assert model.pre_calculated_model_size == summary.model_size + + def test_lazy_model_summary(): """Test that the model summary can work with lazy layers.""" lazy_model = LazyModel() From 503c33b85131f9787272c7b8e35db6d0f483d94a Mon Sep 17 00:00:00 2001 From: Nicki Skafte Date: Wed, 6 Aug 2025 09:28:19 +0200 Subject: [PATCH 3/5] fix --- src/lightning/pytorch/utilities/model_summary/model_summary.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary.py b/src/lightning/pytorch/utilities/model_summary/model_summary.py index 546bbe21f2981..01b692abdc05f 100644 --- a/src/lightning/pytorch/utilities/model_summary/model_summary.py +++ b/src/lightning/pytorch/utilities/model_summary/model_summary.py @@ -238,7 +238,7 @@ def __init__(self, model: "pl.LightningModule", max_depth: int = 1) -> None: "32-true": 32, "64-true": 64, } - if self._model.trainer is not None and self._model.trainer.precision not in precision_to_bits: + if self._model._trainer and self._model.trainer.precision not in precision_to_bits: rank_zero_warn( f"Precision {self._model.trainer.precision} is not supported by the model summary. " " Estimated model size in MB will not be accurate. Using 32 bits instead.", From 8c7240b0d451799e0c835f82098b55b9f102742b Mon Sep 17 00:00:00 2001 From: Nicki Skafte Date: Wed, 6 Aug 2025 09:29:42 +0200 Subject: [PATCH 4/5] changelog --- src/lightning/pytorch/CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 81c7bfc656885..95be25bdea708 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/src/lightning/pytorch/CHANGELOG.md @@ -28,6 +28,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - fix progress bar console clearing for Rich `14.1+` ([#21016](https://github.com/Lightning-AI/pytorch-lightning/pull/21016)) +- Fix support for more dtypes in `ModelSummary` ([#21034](https://github.com/Lightning-AI/pytorch-lightning/pull/21034)) + --- ## [2.5.2] - 2025-06-20 From efe64f2fc0b51426b2a8389fcad2cbaeaedbc428 Mon Sep 17 00:00:00 2001 From: Nicki Skafte Date: Thu, 7 Aug 2025 09:18:13 +0200 Subject: [PATCH 5/5] try fixing tests --- tests/tests_pytorch/utilities/test_model_summary.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/tests_pytorch/utilities/test_model_summary.py b/tests/tests_pytorch/utilities/test_model_summary.py index 8f8be116e3d0a..35edf78fa7081 100644 --- a/tests/tests_pytorch/utilities/test_model_summary.py +++ b/tests/tests_pytorch/utilities/test_model_summary.py @@ -338,15 +338,15 @@ def test_model_size_precision(tmp_path, accelerator, precision): assert model.pre_calculated_model_size == summary.model_size -def test_model_size_warning_on_unsupported_precision(): +def test_model_size_warning_on_unsupported_precision(tmp_path): """Test that a warning is raised when the precision is not supported.""" model = PreCalculatedModel(precision=32) # fallback to 32 bits # supported precision by lightning but not by the model summary - trainer = Trainer(max_epochs=1, precision="16-mixed") + trainer = Trainer(max_epochs=1, precision="16-mixed", default_root_dir=tmp_path) trainer.fit(model) - with pytest.warns(UserWarning, match="Precision 16-mixed is not supported by the model summary.*"): + with pytest.warns(UserWarning, match="Precision .* is not supported by the model summary.*"): summary = summarize(model) assert model.pre_calculated_model_size == summary.model_size