Merge branch 'master' into refactor/tensorboard_log_metrics_handling

Author: deependujha

docs/source-pytorch/conf.py

@@ -639,4 +639,6 @@ def package_list_from_file(file):
     "https://openai.com/index/*",
     "https://tinyurl.com/.*",  # has a human verification check on redirect
     "https://docs.neptune.ai/.*",  # TODO: remove after dropping Neptune support
+    "https://app.neptune.ai/*",
+    "https://www.neptune.ai/*"
 ]

src/lightning/pytorch/CHANGELOG.md

@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 -
 
+- Fixed ``RichModelSummary`` model size display formatting ([#21467](https://github.com/Lightning-AI/pytorch-lightning/pull/21467))
 
 ---
 

src/lightning/pytorch/callbacks/rich_model_summary.py

@@ -17,7 +17,7 @@
 
 from lightning.pytorch.callbacks import ModelSummary
 from lightning.pytorch.utilities.imports import _RICH_AVAILABLE
-from lightning.pytorch.utilities.model_summary import get_human_readable_count
+from lightning.pytorch.utilities.model_summary import get_formatted_model_size, get_human_readable_count
 
 
 class RichModelSummary(ModelSummary):
@@ -105,8 +105,9 @@ def summarize(
         console.print(table)
 
         parameters = []
-        for param in [trainable_parameters, total_parameters - trainable_parameters, total_parameters, model_size]:
+        for param in [trainable_parameters, total_parameters - trainable_parameters, total_parameters]:
             parameters.append("{:<{}}".format(get_human_readable_count(int(param)), 10))
+        parameters.append("{:<{}}".format(get_formatted_model_size(model_size), 10))
 
         grid = Table.grid(expand=True)
         grid.add_column()

tests/tests_pytorch/callbacks/test_rich_model_summary.py

@@ -70,3 +70,34 @@ def example_input_array(self) -> Any:
     # assert that the input summary data was converted correctly
     args, _ = mock_table_add_row.call_args_list[0]
     assert args[1:] == ("0", "layer", "Linear", "66  ", "train", "512  ", "[4, 32]", "[4, 2]")
+
+
+@RunIf(rich=True)
+def test_rich_summary_model_size_formatting():
+    """Ensure model_size uses get_formatted_model_size, not get_human_readable_count."""
+    from io import StringIO
+
+    from rich.console import Console
+
+    model_summary = RichModelSummary()
+    model = BoringModel()
+    summary = summarize(model)
+    summary_data = summary._get_summary_data()
+
+    output = StringIO()
+    console = Console(file=output, force_terminal=True)
+
+    with mock.patch("rich.get_console", return_value=console):
+        model_summary.summarize(
+            summary_data=summary_data,
+            total_parameters=1,
+            trainable_parameters=1,
+            model_size=5500.0,
+            total_training_modes=summary.total_training_modes,
+            total_flops=1,
+        )
+
+    result = output.getvalue()
+    # model_size=5500.0 should display as "5,500.000" (formatted), not "5.5 K" (human readable count)
+    assert "5,500.000" in result
+    assert "5.5 K" not in result