From e76e40e43ec9e4b4067ae927fbc1fd7ab0f9cee1 Mon Sep 17 00:00:00 2001 From: Furkan Egecan Nizam Date: Sun, 28 Jun 2026 02:29:44 +0300 Subject: [PATCH] Handle jagged NestedTensor failures in model summary --- .../utilities/model_summary/model_summary.py | 67 ++++++++++---- .../utilities/test_model_summary.py | 88 +++++++++++++++++++ 2 files changed, 139 insertions(+), 16 deletions(-) diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary.py b/src/lightning/pytorch/utilities/model_summary/model_summary.py index 01b692abdc05f..e88cf9177b6e4 100644 --- a/src/lightning/pytorch/utilities/model_summary/model_summary.py +++ b/src/lightning/pytorch/utilities/model_summary/model_summary.py @@ -217,14 +217,8 @@ def __init__(self, model: "pl.LightningModule", max_depth: int = 1) -> None: if not isinstance(max_depth, int) or max_depth < -1: raise ValueError(f"`max_depth` can be -1, 0 or > 0, got {max_depth}.") - # The max-depth needs to be plus one because the root module is already counted as depth 0. - self._flop_counter = FlopCounterMode( - mods=None if _TORCH_GREATER_EQUAL_2_4 else self._model, - display=False, - depth=max_depth + 1, - ) - self._max_depth = max_depth + self._flop_counter = self._make_flop_counter() self._layer_summary = self.summarize() # 1 byte -> 8 bits # TODO: how do we compute precision_megabytes in case of mixed precision? @@ -260,6 +254,14 @@ def named_modules(self) -> list[tuple[str, nn.Module]]: mods = list(mods)[1:] # do not include root module (LightningModule) return mods + def _make_flop_counter(self) -> FlopCounterMode: + # The max-depth needs to be plus one because the root module is already counted as depth 0. + return FlopCounterMode( + mods=None if _TORCH_GREATER_EQUAL_2_4 else self._model, + display=False, + depth=self._max_depth + 1, + ) + @property def layer_names(self) -> list[str]: return list(self._layer_summary.keys()) @@ -361,15 +363,25 @@ def _forward_example_input(self) -> None: ) forward_context = contextlib.nullcontext() if trainer is None else trainer.precision_plugin.forward_context() - with torch.no_grad(), forward_context, flop_context: - # let the model hooks collect the input- and output shapes - if isinstance(input_, (list, tuple)): - model(*input_) - elif isinstance(input_, dict): - model(**input_) - else: - model(input_) - mode.restore(model) + try: + with torch.no_grad(), forward_context: + try: + with flop_context: + # let the model hooks collect the input- and output shapes + _forward_model(model, input_) + except (NotImplementedError, RuntimeError, TypeError) as ex: + if flop_context is not self._flop_counter or not _is_nested_tensor_flop_counter_error(ex): + raise + + self._flop_counter = self._make_flop_counter() + warning_cache.warn( + "The model summary ran into an unsupported NestedTensor operation while using PyTorch's FLOP" + " counter. FLOP statistics will be omitted, but the example input will be forwarded without" + " FLOP counting so input and output sizes can still be inferred when possible." + ) + _forward_model(model, input_) + finally: + mode.restore(model) def _get_summary_data(self) -> list[tuple[str, list[str]]]: """Makes a summary listing with: @@ -441,6 +453,29 @@ def parse_batch_shape(batch: Any) -> Union[str, list]: return UNKNOWN_SIZE +def _forward_model(model: "pl.LightningModule", input_: Any) -> None: + if isinstance(input_, (list, tuple)): + model(*input_) + elif isinstance(input_, dict): + model(**input_) + else: + model(input_) + + +def _is_nested_tensor_flop_counter_error(exception: BaseException) -> bool: + has_flop_counter_frame = False + has_nested_tensor_frame = False + traceback = exception.__traceback__ + + while traceback is not None: + module_name = traceback.tb_frame.f_globals.get("__name__", "") + has_flop_counter_frame |= module_name == "torch.utils.flop_counter" + has_nested_tensor_frame |= module_name.startswith("torch.nested") + traceback = traceback.tb_next + + return has_flop_counter_frame and has_nested_tensor_frame + + def _format_summary_table( total_parameters: int, trainable_parameters: int, diff --git a/tests/tests_pytorch/utilities/test_model_summary.py b/tests/tests_pytorch/utilities/test_model_summary.py index d90f826d0d1eb..36680bc0336c6 100644 --- a/tests/tests_pytorch/utilities/test_model_summary.py +++ b/tests/tests_pytorch/utilities/test_model_summary.py @@ -19,6 +19,7 @@ import torch import torch.nn as nn from lightning_utilities.test.warning import no_warning_call +from torch.utils.flop_counter import FlopCounterMode from lightning.pytorch import LightningModule, Trainer from lightning.pytorch.demos.boring_classes import BoringModel @@ -97,6 +98,40 @@ def forward(self, x): return self.reduce(self.embed(x)) +class SimpleLinearModel(LightningModule): + def __init__(self): + super().__init__() + self.layer = nn.Linear(3, 2) + self.example_input_array = torch.rand(2, 3) + + def forward(self, x): + return self.layer(x) + + +class JaggedNestedTensorBlock(nn.Module): + def __init__(self, fail_after_forward: bool = False): + super().__init__() + self.proj = nn.Linear(3, 2) + self.fail_after_forward = fail_after_forward + + def forward(self, x): + nested = torch.nested.nested_tensor([x[0, :2], x[1, :3]], layout=torch.jagged) + output = self.proj(nested) + if self.fail_after_forward: + raise RuntimeError("plain forward failed") + return {"nested": output} + + +class JaggedNestedTensorModel(LightningModule): + def __init__(self, fail_after_forward: bool = False): + super().__init__() + self.block = JaggedNestedTensorBlock(fail_after_forward=fail_after_forward) + self.example_input_array = torch.rand(2, 3, 3) + + def forward(self, x): + return self.block(x) + + class PartialScriptModel(LightningModule): """A model which contains scripted layers.""" @@ -205,6 +240,37 @@ def test_mixed_dtype_model_summary(): assert summary.out_sizes == [[2, 3, 20], [2, 3, 1]] # embed # reduce +def test_model_summary_flops_for_normal_tensor_example_input(): + """Test that regular tensor inputs still collect shapes and FLOPs.""" + summary = summarize(SimpleLinearModel()) + assert summary.in_sizes == [[2, 3]] + assert summary.out_sizes == [[2, 2]] + assert summary.total_flops > 0 + + +def test_model_summary_with_jagged_nested_tensor_falls_back_to_unknown_output_size(): + """Test that jagged NestedTensor operations unsupported by the FLOP counter don't crash the summary.""" + _require_jagged_nested_tensor_flop_counter_error() + + model = JaggedNestedTensorModel() + output = model(model.example_input_array) + assert output["nested"].layout is torch.jagged + + summary = summarize(model) + + assert summary.in_sizes == [[2, 3, 3]] + assert summary.out_sizes == [UNKNOWN_SIZE] + assert summary.total_flops == 0 + + +def test_model_summary_with_jagged_nested_tensor_reraises_plain_forward_error(): + """Test that the NestedTensor FLOP fallback does not hide a model error from the plain forward.""" + _require_jagged_nested_tensor_flop_counter_error() + + with pytest.raises(RuntimeError, match="plain forward failed"): + summarize(JaggedNestedTensorModel(fail_after_forward=True)) + + @pytest.mark.parametrize("max_depth", [-1, 0]) def test_hooks_removed_after_summarize(max_depth): """Test that all hooks were properly removed after summary, even ones that were not run.""" @@ -455,6 +521,28 @@ def forward(self, x): assert not model.layer2.training +def _require_jagged_nested_tensor_flop_counter_error(): + if not hasattr(torch, "jagged"): + pytest.skip("Requires torch.jagged layout support.") + if not hasattr(torch, "nested") or not hasattr(torch.nested, "nested_tensor"): + pytest.skip("Requires torch.nested.nested_tensor.") + + layer = nn.Linear(3, 2) + nested = torch.nested.nested_tensor([torch.rand(2, 3), torch.rand(3, 3)], layout=torch.jagged) + try: + layer(nested) + except Exception as ex: + pytest.skip(f"Requires Linear support for jagged NestedTensor: {ex}") + + try: + with FlopCounterMode(display=False): + layer(nested) + except (NotImplementedError, RuntimeError, TypeError): + return + + pytest.skip("Requires FlopCounterMode to not support jagged NestedTensor.") + + def test_total_training_modes(): """Test that the `total_training_modes` counts the modules in 'train' and 'eval' mode, excluding the root module."""