Merge branch 'master' into master

Borda · web-flow · commit 35d3f1723a14 · 2025-06-11T07:36:08.000+02:00
diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
@@ -97,7 +97,7 @@ jobs:
         # adding dome more images as Thunder mainly using python 3.10,
         # and we need to support integrations as for example LitGPT
         python_version: ["3.10"]
-        pytorch_version: ["2.6.0", "2.7.0"]
+        pytorch_version: ["2.6.0", "2.7.1"]
         cuda_version: ["12.6.3"]
         include:
           # These are the base images for PL release docker images.
@@ -108,7 +108,7 @@ jobs:
           - { python_version: "3.11", pytorch_version: "2.4.1", cuda_version: "12.1.1" }
           - { python_version: "3.12", pytorch_version: "2.5.1", cuda_version: "12.1.1" }
           - { python_version: "3.12", pytorch_version: "2.6.0", cuda_version: "12.4.1" }
-          - { python_version: "3.12", pytorch_version: "2.7.0", cuda_version: "12.6.3" }
+          - { python_version: "3.12", pytorch_version: "2.7.1", cuda_version: "12.6.3" }
     steps:
       - uses: actions/checkout@v4
       - uses: docker/setup-buildx-action@v3
diff --git a/docs/source-pytorch/conf.py b/docs/source-pytorch/conf.py
@@ -487,6 +487,7 @@ def _load_py_module(name: str, location: str) -> ModuleType:
     ("py:meth", "setup"),
     ("py:meth", "test_step"),
     ("py:meth", "toggle_optimizer"),
+    ("py:meth", "toggled_optimizer"),
     ("py:class", "torch.ScriptModule"),
     ("py:class", "torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload"),
     ("py:class", "torch.distributed.fsdp.fully_sharded_data_parallel.MixedPrecision"),
diff --git a/docs/source-pytorch/model/manual_optimization.rst b/docs/source-pytorch/model/manual_optimization.rst
@@ -17,7 +17,7 @@ To manually optimize, do the following:
   * ``optimizer.zero_grad()`` to clear the gradients from the previous training step
   * ``self.manual_backward(loss)`` instead of ``loss.backward()``
   * ``optimizer.step()`` to update your model parameters
-  * ``self.toggle_optimizer()`` and ``self.untoggle_optimizer()`` if needed
+  * ``self.toggle_optimizer()`` and ``self.untoggle_optimizer()``, or ``self.toggled_optimizer()`` if needed
 
 Here is a minimal example of manual optimization.
 
diff --git a/requirements/ci.txt b/requirements/ci.txt
@@ -1,4 +1,4 @@
-setuptools <80.8.1
+setuptools <80.9.1
 wheel <0.46.0
 awscli >=1.30.0, <1.41.0
 twine ==6.1.0
diff --git a/requirements/doctests.txt b/requirements/doctests.txt
@@ -1,2 +1,2 @@
-pytest ==8.3.5
+pytest ==8.4.0
 pytest-doctestplus ==1.4.0
diff --git a/requirements/fabric/test.txt b/requirements/fabric/test.txt
@@ -1,6 +1,6 @@
-coverage ==7.8.0
+coverage ==7.8.2
 numpy >=1.17.2, <1.27.0
-pytest ==8.3.5
+pytest ==8.4.0
 pytest-cov ==6.1.1
 pytest-timeout ==2.4.0
 pytest-rerunfailures ==15.1
diff --git a/requirements/pytorch/test.txt b/requirements/pytorch/test.txt
@@ -1,5 +1,5 @@
-coverage ==7.8.0
-pytest ==8.3.5
+coverage ==7.8.2
+pytest ==8.4.0
 pytest-cov ==6.1.1
 pytest-timeout ==2.4.0
 pytest-rerunfailures ==15.1
@@ -12,7 +12,7 @@ numpy >=1.17.2, <1.27.0
 onnx >=1.12.0, <1.19.0
 onnxruntime >=1.12.0, <1.21.0
 psutil <7.0.1 # for `DeviceStatsMonitor`
-pandas >1.0, <2.3.0  # needed in benchmarks
+pandas >2.0, <2.4.0  # needed in benchmarks
 fastapi  # for `ServableModuleValidator`  # not setting version as re-defined in App
 uvicorn  # for `ServableModuleValidator`  # not setting version as re-defined in App
 
diff --git a/requirements/typing.txt b/requirements/typing.txt
@@ -1,5 +1,5 @@
 mypy==1.15.0
-torch==2.7.0
+torch==2.7.1
 
 types-Markdown
 types-PyYAML
diff --git a/setup.py b/setup.py
@@ -110,7 +110,8 @@ def _set_manifest_path(manifest_dir: str, aggregate: bool = False, mapping: Mapp
         assert os.path.exists(manifest_path)
     # avoid error: setup script specifies an absolute path
     manifest_path = os.path.relpath(manifest_path, _PATH_ROOT)
-    logging.info("Set manifest path to", manifest_path)
+    # Use lazy logging formatting
+    logging.info("Set manifest path to %s", manifest_path)
     setuptools.command.egg_info.manifest_maker.template = manifest_path
     yield
     # cleanup
diff --git a/src/lightning/fabric/connector.py b/src/lightning/fabric/connector.py
@@ -492,7 +492,7 @@ def _check_and_init_precision(self) -> Precision:
                 if self._precision_input == "16-mixed"
                 else "Using bfloat16 Automatic Mixed Precision (AMP)"
             )
-            device = "cpu" if self._accelerator_flag == "cpu" else "cuda"
+            device = self._accelerator_flag if self._accelerator_flag in ("cpu", "mps") else "cuda"
             return MixedPrecision(precision=self._precision_input, device=device)  # type: ignore[arg-type]
 
         raise RuntimeError("No precision set")
diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md
@@ -15,7 +15,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `MemoryFormat` callback to easily change the model memory format ([#15175](https://github.com/Lightning-AI/lightning/pull/17680))
 
 
-- Allow LightningCLI to use a customized argument parser class ([#20596](https://github.com/Lightning-AI/pytorch-lightning/pull/20596))
+- Add `toggled_optimizer(optimizer)` method to the LightningModule, which is a context manager version of `toggle_optimize` and `untoggle_optimizer` ([#20771](https://github.com/Lightning-AI/pytorch-lightning/pull/20771))
 
 
 - For cross-device local checkpoints, instruct users to install `fsspec>=2025.5.0` if unavailable ([#20780](https://github.com/Lightning-AI/pytorch-lightning/pull/20780))
diff --git a/src/lightning/pytorch/callbacks/model_summary.py b/src/lightning/pytorch/callbacks/model_summary.py
@@ -68,13 +68,17 @@ def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
         model_size = model_summary.model_size
         total_training_modes = model_summary.total_training_modes
 
+        # todo Add `total_flops` in DeepSpeedSummary.
+        total_flops = model_summary.total_flops if hasattr(model_summary, "total_flops") else 0
+
         if trainer.is_global_zero:
             self.summarize(
                 summary_data,
                 total_parameters,
                 trainable_parameters,
                 model_size,
                 total_training_modes,
+                total_flops=total_flops,
                 **self._summarize_kwargs,
             )
 
@@ -92,13 +96,15 @@ def summarize(
         trainable_parameters: int,
         model_size: float,
         total_training_modes: dict[str, int],
+        total_flops: int,
         **summarize_kwargs: Any,
     ) -> None:
         summary_table = _format_summary_table(
             total_parameters,
             trainable_parameters,
             model_size,
             total_training_modes,
+            total_flops,
             *summary_data,
         )
         log.info("\n" + summary_table)
diff --git a/src/lightning/pytorch/callbacks/progress/progress_bar.py b/src/lightning/pytorch/callbacks/progress/progress_bar.py
@@ -85,6 +85,9 @@ def total_train_batches(self) -> Union[int, float]:
         dataloader is of infinite size.
 
         """
+        if self.trainer.max_epochs == -1 and self.trainer.max_steps is not None and self.trainer.max_steps > 0:
+            remaining_steps = self.trainer.max_steps - self.trainer.global_step
+            return min(self.trainer.num_training_batches, remaining_steps)
         return self.trainer.num_training_batches
 
     @property
diff --git a/src/lightning/pytorch/callbacks/rich_model_summary.py b/src/lightning/pytorch/callbacks/rich_model_summary.py
@@ -72,6 +72,7 @@ def summarize(
         trainable_parameters: int,
         model_size: float,
         total_training_modes: dict[str, int],
+        total_flops: int,
         **summarize_kwargs: Any,
     ) -> None:
         from rich import get_console
@@ -86,6 +87,7 @@ def summarize(
         table.add_column("Type")
         table.add_column("Params", justify="right")
         table.add_column("Mode")
+        table.add_column("FLOPs", justify="right")
 
         column_names = list(zip(*summary_data))[0]
 
@@ -113,5 +115,6 @@ def summarize(
         grid.add_row(f"[bold]Total estimated model params size (MB)[/]: {parameters[3]}")
         grid.add_row(f"[bold]Modules in train mode[/]: {total_training_modes['train']}")
         grid.add_row(f"[bold]Modules in eval mode[/]: {total_training_modes['eval']}")
+        grid.add_row(f"[bold]Total FLOPs[/]: {get_human_readable_count(total_flops)}")
 
         console.print(grid)
diff --git a/src/lightning/pytorch/core/module.py b/src/lightning/pytorch/core/module.py
@@ -1141,6 +1141,32 @@ def untoggle_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer]) ->
         # save memory
         self._param_requires_grad_state = {}
 
+    @contextmanager
+    def toggled_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer]) -> Generator:
+        """Makes sure only the gradients of the current optimizer's parameters are calculated in the training step to
+        prevent dangling gradients in multiple-optimizer setup. Combines :meth:`toggle_optimizer` and
+        :meth:`untoggle_optimizer` into context manager.
+
+        Args:
+            optimizer: The optimizer to toggle.
+
+        Example::
+
+            def training_step(...):
+                opt = self.optimizers()
+                with self.toggled_optimizer(opt):
+                    loss = ...
+                    opt.zero_grad()
+                    self.manual_backward(loss)
+                    opt.step()
+
+        """
+        self.toggle_optimizer(optimizer)
+        try:
+            yield
+        finally:
+            self.untoggle_optimizer(optimizer)
+
     def clip_gradients(
         self,
         optimizer: Optimizer,
@@ -1472,6 +1498,10 @@ def forward(self, x):
                     )
                 example_inputs = self.example_input_array
 
+            if kwargs.get("check_inputs") is not None:
+                kwargs["check_inputs"] = self._on_before_batch_transfer(kwargs["check_inputs"])
+                kwargs["check_inputs"] = self._apply_batch_transfer_handler(kwargs["check_inputs"])
+
             # automatically send example inputs to the right device and use trace
             example_inputs = self._on_before_batch_transfer(example_inputs)
             example_inputs = self._apply_batch_transfer_handler(example_inputs)
diff --git a/src/lightning/pytorch/utilities/model_summary/model_summary.py b/src/lightning/pytorch/utilities/model_summary/model_summary.py
@@ -22,10 +22,12 @@
 import torch
 import torch.nn as nn
 from torch import Tensor
+from torch.utils.flop_counter import FlopCounterMode
 from torch.utils.hooks import RemovableHandle
 
 import lightning.pytorch as pl
 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
 from lightning.pytorch.utilities.rank_zero import WarningCache
 
@@ -180,29 +182,31 @@ class ModelSummary:
         ...
         >>> model = LitModel()
         >>> ModelSummary(model, max_depth=1)  # doctest: +NORMALIZE_WHITESPACE
-          | Name | Type       | Params | Mode  | In sizes  | Out sizes
-        --------------------------------------------------------------------
-        0 | net  | Sequential | 132 K  | train | [10, 256] | [10, 512]
-        --------------------------------------------------------------------
+          | Name | Type       | Params | Mode  | FLOPs | In sizes  | Out sizes
+        ----------------------------------------------------------------------------
+        0 | net  | Sequential | 132 K  | train | 2.6 M | [10, 256] | [10, 512]
+        ----------------------------------------------------------------------------
         132 K     Trainable params
         0         Non-trainable params
         132 K     Total params
         0.530     Total estimated model params size (MB)
         3         Modules in train mode
         0         Modules in eval mode
+        2.6 M     Total Flops
         >>> ModelSummary(model, max_depth=-1)  # doctest: +NORMALIZE_WHITESPACE
-          | Name  | Type        | Params | Mode  | In sizes  | Out sizes
-        ----------------------------------------------------------------------
-        0 | net   | Sequential  | 132 K  | train | [10, 256] | [10, 512]
-        1 | net.0 | Linear      | 131 K  | train | [10, 256] | [10, 512]
-        2 | net.1 | BatchNorm1d | 1.0 K  | train | [10, 512] | [10, 512]
-        ----------------------------------------------------------------------
+          | Name  | Type        | Params | Mode  | FLOPs | In sizes  | Out sizes
+        ------------------------------------------------------------------------------
+        0 | net   | Sequential  | 132 K  | train | 2.6 M | [10, 256] | [10, 512]
+        1 | net.0 | Linear      | 131 K  | train | 2.6 M | [10, 256] | [10, 512]
+        2 | net.1 | BatchNorm1d | 1.0 K  | train | 0     | [10, 512] | [10, 512]
+        ------------------------------------------------------------------------------
         132 K     Trainable params
         0         Non-trainable params
         132 K     Total params
         0.530     Total estimated model params size (MB)
         3         Modules in train mode
         0         Modules in eval mode
+        2.6 M     Total Flops
 
     """
 
@@ -212,6 +216,13 @@ 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._layer_summary = self.summarize()
         # 1 byte -> 8 bits
@@ -279,6 +290,22 @@ def total_layer_params(self) -> int:
     def model_size(self) -> float:
         return self.total_parameters * self._precision_megabytes
 
+    @property
+    def total_flops(self) -> int:
+        return self._flop_counter.get_total_flops()
+
+    @property
+    def flop_counts(self) -> dict[str, dict[Any, int]]:
+        flop_counts = self._flop_counter.get_flop_counts()
+        ret = {
+            name: flop_counts.get(
+                f"{type(self._model).__name__}.{name}",
+                {},
+            )
+            for name in self.layer_names
+        }
+        return ret
+
     def summarize(self) -> dict[str, LayerSummary]:
         summary = OrderedDict((name, LayerSummary(module)) for name, module in self.named_modules)
         if self._model.example_input_array is not None:
@@ -307,8 +334,18 @@ def _forward_example_input(self) -> None:
         mode.capture(model)
         model.eval()
 
+        # FlopCounterMode does not support ScriptModules before torch 2.4.0, so we use a null context
+        flop_context = (
+            contextlib.nullcontext()
+            if (
+                not _TORCH_GREATER_EQUAL_2_4
+                and any(isinstance(m, torch.jit.ScriptModule) for m in self._model.modules())
+            )
+            else self._flop_counter
+        )
+
         forward_context = contextlib.nullcontext() if trainer is None else trainer.precision_plugin.forward_context()
-        with torch.no_grad(), 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_)
@@ -330,6 +367,7 @@ def _get_summary_data(self) -> list[tuple[str, list[str]]]:
             ("Type", self.layer_types),
             ("Params", list(map(get_human_readable_count, self.param_nums))),
             ("Mode", ["train" if mode else "eval" for mode in self.training_modes]),
+            ("FLOPs", list(map(get_human_readable_count, (sum(x.values()) for x in self.flop_counts.values())))),
         ]
         if self._model.example_input_array is not None:
             arrays.append(("In sizes", [str(x) for x in self.in_sizes]))
@@ -349,6 +387,7 @@ def _add_leftover_params_to_summary(self, arrays: list[tuple[str, list[str]]], t
         layer_summaries["Type"].append(NOT_APPLICABLE)
         layer_summaries["Params"].append(get_human_readable_count(total_leftover_params))
         layer_summaries["Mode"].append(NOT_APPLICABLE)
+        layer_summaries["FLOPs"].append(NOT_APPLICABLE)
         if "In sizes" in layer_summaries:
             layer_summaries["In sizes"].append(NOT_APPLICABLE)
         if "Out sizes" in layer_summaries:
@@ -361,8 +400,16 @@ def __str__(self) -> str:
         trainable_parameters = self.trainable_parameters
         model_size = self.model_size
         total_training_modes = self.total_training_modes
-
-        return _format_summary_table(total_parameters, trainable_parameters, model_size, total_training_modes, *arrays)
+        total_flops = self.total_flops
+
+        return _format_summary_table(
+            total_parameters,
+            trainable_parameters,
+            model_size,
+            total_training_modes,
+            total_flops,
+            *arrays,
+        )
 
     def __repr__(self) -> str:
         return str(self)
@@ -383,6 +430,7 @@ def _format_summary_table(
     trainable_parameters: int,
     model_size: float,
     total_training_modes: dict[str, int],
+    total_flops: int,
     *cols: tuple[str, list[str]],
 ) -> str:
     """Takes in a number of arrays, each specifying a column in the summary table, and combines them all into one big
@@ -423,6 +471,8 @@ def _format_summary_table(
     summary += "Modules in train mode"
     summary += "\n" + s.format(total_training_modes["eval"], 10)
     summary += "Modules in eval mode"
+    summary += "\n" + s.format(get_human_readable_count(total_flops), 10)
+    summary += "Total Flops"
 
     return summary
 
diff --git a/tests/tests_fabric/test_connector.py b/tests/tests_fabric/test_connector.py
@@ -405,6 +405,13 @@ def test_unsupported_strategy_types_on_cpu_and_fallback():
     assert isinstance(connector.strategy, DDPStrategy)
 
 
+@RunIf(mps=True)
+@pytest.mark.parametrize("precision", ["16-mixed", "bf16-mixed"])
+def test_mps_enabled_with_float16_or_bfloat16_precision(precision):
+    connector = _Connector(accelerator="mps", precision=precision)
+    assert connector.precision.device == "mps"
+
+
 def test_invalid_accelerator_choice():
     with pytest.raises(ValueError, match="You selected an invalid accelerator name: `accelerator='cocofruit'`"):
         _Connector(accelerator="cocofruit")
diff --git a/tests/tests_fabric/test_fabric.py b/tests/tests_fabric/test_fabric.py
diff --git a/tests/tests_pytorch/callbacks/test_model_summary.py b/tests/tests_pytorch/callbacks/test_model_summary.py
diff --git a/tests/tests_pytorch/callbacks/test_rich_model_summary.py b/tests/tests_pytorch/callbacks/test_rich_model_summary.py
diff --git a/tests/tests_pytorch/core/test_lightning_module.py b/tests/tests_pytorch/core/test_lightning_module.py
diff --git a/tests/tests_pytorch/loops/test_training_loop.py b/tests/tests_pytorch/loops/test_training_loop.py
diff --git a/tests/tests_pytorch/models/test_hparams.py b/tests/tests_pytorch/models/test_hparams.py
diff --git a/tests/tests_pytorch/models/test_torchscript.py b/tests/tests_pytorch/models/test_torchscript.py
diff --git a/tests/tests_pytorch/utilities/test_model_summary.py b/tests/tests_pytorch/utilities/test_model_summary.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-setuptools <80.8.1`
	`1`	`+setuptools <80.9.1`
`2`	`2`	`wheel <0.46.0`
`3`	`3`	`awscli >=1.30.0, <1.41.0`
`4`	`4`	`twine ==6.1.0`
Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`		`-pytest ==8.3.5`
	`1`	`+pytest ==8.4.0`
`2`	`2`	`pytest-doctestplus ==1.4.0`
Original file line number	Diff line number	Diff line change
`@@ -492,7 +492,7 @@ def _check_and_init_precision(self) -> Precision:`
`492`	`492`	`if self._precision_input == "16-mixed"`
`493`	`493`	`else "Using bfloat16 Automatic Mixed Precision (AMP)"`
`494`	`494`	`)`
`495`		`- device = "cpu" if self._accelerator_flag == "cpu" else "cuda"`
	`495`	`+ device = self._accelerator_flag if self._accelerator_flag in ("cpu", "mps") else "cuda"`
`496`	`496`	`return MixedPrecision(precision=self._precision_input, device=device) # type: ignore[arg-type]`
`497`	`497`
`498`	`498`	`raise RuntimeError("No precision set")`