Merge branch 'master' into fix/ddp-accumulate-grad-stream-mismatch-warning

Author: deependujha

.github/workflows/_legacy-checkpoints.yml

@@ -60,7 +60,7 @@ jobs:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
         with:
           python-version: "3.10"
           # TODO: Avoid activating environment like this

.github/workflows/ci-tests-fabric.yml

@@ -74,7 +74,7 @@ jobs:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
         with:
           python-version: ${{ matrix.config.python-version || '3.10' }}
           # TODO: Avoid activating environment like this

.github/workflows/ci-tests-pytorch.yml

@@ -79,7 +79,7 @@ jobs:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
         with:
           python-version: ${{ matrix.config.python-version || '3.10' }}
           # TODO: Avoid activating environment like this

.github/workflows/code-checks.yml

@@ -34,7 +34,7 @@ jobs:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
         with:
           python-version: "3.11"
           # TODO: Avoid activating environment like this

.github/workflows/docs-build.yml

@@ -74,7 +74,7 @@ jobs:
           lfs: ${{ matrix.pkg-name == 'pytorch' }}
 
       - name: Install uv and set Python version
-        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
         with:
           python-version: "3.10"
           # TODO: Avoid activating environment like this

.github/workflows/release-pkg.yml

@@ -154,7 +154,7 @@ jobs:
 
       - name: Publish distribution 📦 to PyPI
         # pypa/gh-action-pypi-publish v1.13.0
-        uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e
+        uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b
         with:
           packages_dir: dist/${{ steps.folder.outputs.pkg }}
           verbose: true

src/lightning/pytorch/CHANGELOG.md

@@ -24,10 +24,15 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed non-zero process exits in `CombinedLoader.reset()` with large tensors and persistent spawned workers by avoiding explicit `_shutdown_workers()` calls and relying on iterator cleanup via `del` [#21708](https://github.com/Lightning-AI/pytorch-lightning/issues/21708)
+
 - Fixed `SIGTERMException` producing a zero exit code instead of 143 (128 + SIGTERM) ([#21623](https://github.com/Lightning-AI/pytorch-lightning/issues/21623))
 
 - fixed AccumulateGrad stream mismatch warning when using DDP with Trainer ([#21746](https://github.com/Lightning-AI/pytorch-lightning/pull/21746))
 
+- Fixed `LightningModule.toggle_optimizer` / `untoggle_optimizer` breaking under `torch.compile` by disabling Dynamo tracing on these bookkeeping helpers ([#21513](https://github.com/Lightning-AI/pytorch-lightning/issues/21513))
+
+
 ---
 
 ## [2.6.4] - 2026-05-20

src/lightning/pytorch/core/module.py

@@ -1136,12 +1136,24 @@ def backward(self, loss):
         else:
             loss.backward(*args, **kwargs)
 
+    @torch.compiler.disable
     def toggle_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer]) -> None:
         """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.
 
         It works with :meth:`untoggle_optimizer` to make sure ``param_requires_grad_state`` is properly reset.
 
+        .. note::
+            This method is decorated with :func:`torch.compiler.disable` so that it is executed as regular
+            Python when the ``LightningModule`` is wrapped with :func:`torch.compile`. Mutating
+            ``requires_grad`` on parameters is not supported by Dynamo/AOTAutograd (it can change a
+            tensor's leaf-ness mid-graph), so tracing this bookkeeping helper would either fail with
+            ``Unsupported: setattr() on Tensor.requires_grad`` or produce a ``KeyError`` on the
+            internal ``param_requires_grad_state`` mapping when the traced parameter references diverge
+            from those held by ``trainer.optimizers``. Disabling the compiler on this method keeps the
+            behavior identical for eager users while making it safe to call from a compiled
+            ``training_step``.
+
         Args:
             optimizer: The optimizer to toggle.
 
@@ -1165,9 +1177,13 @@ def toggle_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer]) -> N
                 param.requires_grad = param_requires_grad_state[param]
         self._param_requires_grad_state = param_requires_grad_state
 
+    @torch.compiler.disable
     def untoggle_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer]) -> None:
         """Resets the state of required gradients that were toggled with :meth:`toggle_optimizer`.
 
+        See :meth:`toggle_optimizer` for details on why this method is decorated with
+        :func:`torch.compiler.disable`.
+
         Args:
             optimizer: The optimizer to untoggle.
 

src/lightning/pytorch/utilities/combined_loader.py

@@ -397,7 +397,7 @@ def _load_state_dicts(self, states: list[dict[str, Any]]) -> None:
 def _shutdown_workers_and_reset_iterator(dataloader: object) -> None:
     if hasattr(dataloader, "_iterator"):
         if isinstance(dataloader._iterator, _MultiProcessingDataLoaderIter):
-            dataloader._iterator._shutdown_workers()
+            del dataloader._iterator
         dataloader._iterator = None
 
 

tests/tests_pytorch/core/test_lightning_module.py

@@ -298,6 +298,27 @@ def configure_optimizers(self):
     trainer.fit(model)
 
 
+@RunIf(dynamo=True)
+def test_toggle_untoggle_optimizer_are_compiler_disabled():
+    """Regression test for https://github.com/Lightning-AI/pytorch-lightning/issues/21513.
+
+    ``toggle_optimizer`` / ``untoggle_optimizer`` mutate ``requires_grad`` on Parameters, which
+    Dynamo/AOTAutograd does not support because it can change a tensor's leaf-ness mid-graph.
+    Tracing these helpers either graph-breaks with ``Unsupported: setattr() on Tensor.requires_grad``
+    or raises a ``KeyError`` on the internal ``param_requires_grad_state`` mapping when the traced
+    parameter references diverge from those held by ``trainer.optimizers``. Both methods are
+    decorated with ``@torch.compiler.disable`` so that Dynamo never enters them. This test verifies
+    the decorator is attached via the ``_torchdynamo_disable`` attribute the decorator installs
+    (the same assertion pattern used by ``tests/utilities/test_compile.py::test_compile_uncompile``).
+    """
+
+    def is_compiler_disabled(fn):
+        return any(el.startswith("_torchdynamo_disable") for el in dir(fn))
+
+    assert is_compiler_disabled(LightningModule.toggle_optimizer)
+    assert is_compiler_disabled(LightningModule.untoggle_optimizer)
+
+
 @pytest.mark.parametrize(
     ("accelerator", "device"),
     [