Handle jagged NestedTensor failures in model summary

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,

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."""