Merge branch 'master' into issue-20311-cli-save-hyper-instatiation-links

Author: Borda

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"),

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.
 

requirements/doctests.txt

@@ -1,2 +1,2 @@
-pytest ==8.3.5
+pytest ==8.4.0
 pytest-doctestplus ==1.4.0

requirements/fabric/test.txt

@@ -1,6 +1,6 @@
 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

requirements/pytorch/test.txt

@@ -1,5 +1,5 @@
 coverage ==7.8.2
-pytest ==8.3.5
+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
 

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

src/lightning/pytorch/CHANGELOG.md

@@ -12,6 +12,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Add enable_autolog_hparams argument to Trainer ([#20593](https://github.com/Lightning-AI/pytorch-lightning/pull/20593))
 
 
+- 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))
 
 

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

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,

tests/tests_pytorch/core/test_lightning_module.py

@@ -119,6 +119,22 @@ def test_1_optimizer_toggle_model():
     assert not model._param_requires_grad_state
 
 
+def test_optimizer_toggle_model_context_manager():
+    """Test toggle_model runs when only one optimizer is used."""
+    model = BoringModel()
+    trainer = Mock()
+    model.trainer = trainer
+    params = model.parameters()
+    optimizer = torch.optim.SGD(params, lr=0.1)
+    trainer.optimizers = [optimizer]
+
+    assert not model._param_requires_grad_state
+    # toggle optimizer was failing with a single optimizer
+    with model.toggled_optimizer(optimizer):
+        assert model._param_requires_grad_state
+    assert not model._param_requires_grad_state
+
+
 def test_toggle_untoggle_2_optimizers_no_shared_parameters(tmp_path):
     class TestModel(BoringModel):
         def __init__(self):