From 41a4a67b99589cb2c5c79e8753310825fc5ba99c Mon Sep 17 00:00:00 2001 From: Nicki Skafte Date: Mon, 28 Jul 2025 11:32:54 +0200 Subject: [PATCH 1/8] change dict to mapping --- src/lightning/pytorch/loops/optimization/manual.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/lightning/pytorch/loops/optimization/manual.py b/src/lightning/pytorch/loops/optimization/manual.py index e1aabcbf42976..10bd5b8b1c666 100644 --- a/src/lightning/pytorch/loops/optimization/manual.py +++ b/src/lightning/pytorch/loops/optimization/manual.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from collections import OrderedDict +from collections.abc import Mapping from contextlib import suppress from dataclasses import dataclass, field from typing import Any @@ -45,7 +46,7 @@ class ManualResult(OutputResult): @classmethod def from_training_step_output(cls, training_step_output: STEP_OUTPUT) -> "ManualResult": extra = {} - if isinstance(training_step_output, dict): + if isinstance(training_step_output, Mapping): extra = training_step_output.copy() elif isinstance(training_step_output, Tensor): extra = {"loss": training_step_output} From 42418f875eb2de9673512bebff95c43c73091eb7 Mon Sep 17 00:00:00 2001 From: Nicki Skafte Date: Mon, 28 Jul 2025 11:44:34 +0200 Subject: [PATCH 2/8] add a bit of testing --- .../trainer/optimization/test_manual_optimization.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py index 3f89e1459298d..c2b143e44b6ce 100644 --- a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py +++ b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py @@ -305,7 +305,8 @@ def on_train_epoch_end(self, *_, **__): @RunIf(min_cuda_gpus=1) -def test_multiple_optimizers_step(tmp_path): +@pytest.mark.parametrize("dicttype", [dict, collections.OrderedDict]) +def test_multiple_optimizers_step(tmp_path, dicttype): """Tests that `step` works with several optimizers.""" class TestModel(ManualOptModel): @@ -335,7 +336,7 @@ def training_step(self, batch, batch_idx): opt_b.step() opt_b.zero_grad() - return {"loss1": loss_1.detach(), "loss2": loss_2.detach()} + return dicttype(loss1=loss_1.detach(), loss2=loss_2.detach()) # sister test: tests/plugins/test_amp_plugins.py::test_amp_gradient_unscale def on_after_backward(self) -> None: From 78e55098fd3d30ea1a3530ffb4b0e2b40c615496 Mon Sep 17 00:00:00 2001 From: Nicki Skafte Date: Tue, 5 Aug 2025 07:15:57 +0200 Subject: [PATCH 3/8] changelog --- src/lightning/pytorch/CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 81c7bfc656885..486c914f8898b 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/src/lightning/pytorch/CHANGELOG.md @@ -10,7 +10,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ### Added -- +- Added support for general mappings being returned from `training_step` when using manual optimization ([#21011](https://github.com/Lightning-AI/pytorch-lightning/pull/21011)) + ### Changed From 2e0def1cc6e6cdc92a4235103bdb56a3f84ab9f9 Mon Sep 17 00:00:00 2001 From: Nicki Skafte Date: Tue, 5 Aug 2025 07:23:52 +0200 Subject: [PATCH 4/8] use custom mapping --- .../optimization/test_manual_optimization.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py index c2b143e44b6ce..dd2c340d14ee5 100644 --- a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py +++ b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py @@ -304,8 +304,26 @@ def on_train_epoch_end(self, *_, **__): trainer.fit(model) +class CustomMapping(collections.abc.Mapping): + """A custom implementation of Mapping for testing purposes.""" + def __init__(self, *args, **kwargs): + self._store = dict(*args, **kwargs) + + def __getitem__(self, key): + return self._store[key] + + def __iter__(self): + return iter(self._store) + + def __len__(self): + return len(self._store) + + def __repr__(self): + return f"{self.__class__.__name__}({self._store})" + + @RunIf(min_cuda_gpus=1) -@pytest.mark.parametrize("dicttype", [dict, collections.OrderedDict]) +@pytest.mark.parametrize("dicttype", [dict, CustomMapping]) def test_multiple_optimizers_step(tmp_path, dicttype): """Tests that `step` works with several optimizers.""" From ca515af788def1a241eb15b3c701a357a19bf342 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 5 Aug 2025 05:24:29 +0000 Subject: [PATCH 5/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../trainer/optimization/test_manual_optimization.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py index dd2c340d14ee5..6e12176c2c312 100644 --- a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py +++ b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py @@ -306,6 +306,7 @@ def on_train_epoch_end(self, *_, **__): class CustomMapping(collections.abc.Mapping): """A custom implementation of Mapping for testing purposes.""" + def __init__(self, *args, **kwargs): self._store = dict(*args, **kwargs) From f476a3fbd24cb5c698e5ea43041b1ef78a4be7aa Mon Sep 17 00:00:00 2001 From: Nicki Skafte Detlefsen Date: Wed, 6 Aug 2025 07:24:00 +0200 Subject: [PATCH 6/8] Update tests/tests_pytorch/trainer/optimization/test_manual_optimization.py Co-authored-by: Bhimraj Yadav --- .../trainer/optimization/test_manual_optimization.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py index 6e12176c2c312..d4f37449c6a36 100644 --- a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py +++ b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py @@ -322,6 +322,10 @@ def __len__(self): def __repr__(self): return f"{self.__class__.__name__}({self._store})" + def __copy__(self): + cls = self.__class__ + new_obj = cls(self._store.copy()) + return new_obj @RunIf(min_cuda_gpus=1) @pytest.mark.parametrize("dicttype", [dict, CustomMapping]) From ea1748a7db88809f1929fbe0f9229ce52e129c6d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 6 Aug 2025 05:24:18 +0000 Subject: [PATCH 7/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../trainer/optimization/test_manual_optimization.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py index d4f37449c6a36..4869084e452c6 100644 --- a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py +++ b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py @@ -327,6 +327,7 @@ def __copy__(self): new_obj = cls(self._store.copy()) return new_obj + @RunIf(min_cuda_gpus=1) @pytest.mark.parametrize("dicttype", [dict, CustomMapping]) def test_multiple_optimizers_step(tmp_path, dicttype): From 6e7be89353dd2208a33eea00ae5da5b4dcdf97c4 Mon Sep 17 00:00:00 2001 From: Deependu Jha Date: Wed, 6 Aug 2025 07:10:51 +0000 Subject: [PATCH 8/8] update --- .../trainer/optimization/test_manual_optimization.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py index 4869084e452c6..dd8042ecf2058 100644 --- a/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py +++ b/tests/tests_pytorch/trainer/optimization/test_manual_optimization.py @@ -327,6 +327,9 @@ def __copy__(self): new_obj = cls(self._store.copy()) return new_obj + def copy(self): + return self.__copy__() + @RunIf(min_cuda_gpus=1) @pytest.mark.parametrize("dicttype", [dict, CustomMapping])