Allow frozen data classes in optimizer state dict

requirements/app/base.txt

@@ -7,7 +7,7 @@ fsspec>=2022.5.0, <=2022.7.1
 croniter>=1.3.0, <1.4.0  # strict; TODO: for now until we find something more robust.
 traitlets>=5.3.0, <5.9.0
 arrow>=1.2.0, <1.2.4
-lightning-utilities>=0.6.0.post0, <0.7.0
+lightning-utilities>=0.7.0, <0.8.0
 beautifulsoup4>=4.8.0, <4.11.2
 inquirer>=2.10.0, <=3.1.2
 psutil<5.9.5

requirements/fabric/base.txt

@@ -6,4 +6,4 @@ torch>=1.11.0, <=2.0.0
 fsspec[http]>2021.06.0, <2023.2.0
 packaging>=17.1, <=23.0
 typing-extensions>=4.0.0, <=4.4.0
-lightning-utilities>=0.6.0.post0, <0.7.0
+lightning-utilities>=0.7.0, <0.8.0

requirements/pytorch/base.txt

@@ -9,4 +9,4 @@ fsspec[http]>2021.06.0, <2023.2.0
 torchmetrics>=0.7.0, <0.12.0  # needed for using fixed compare_version
 packaging>=17.1, <=23.0
 typing-extensions>=4.0.0, <=4.4.0
-lightning-utilities>=0.6.0.post0, <0.7.0
+lightning-utilities>=0.7.0, <0.8.0

src/lightning/fabric/CHANGELOG.md

@@ -20,6 +20,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added support for writing logs to remote file systems with the `CSVLogger` ([#16880](https://github.com/Lightning-AI/lightning/pull/16880))
 
+
+- Added support for frozen dataclasses in the optimizer state ([#16656](https://github.com/Lightning-AI/lightning/pull/16656))
+
+
 ### Changed
 
 - Fabric now chooses `accelerator="auto", strategy="auto", devices="auto"` as defaults ([#16842](https://github.com/Lightning-AI/lightning/pull/16842))

src/lightning/fabric/utilities/optimizer.py

@@ -31,4 +31,4 @@ def _optimizers_to_device(optimizers: Iterable[Optimizer], device: _DEVICE) -> N
 def _optimizer_to_device(optimizer: Optimizer, device: _DEVICE) -> None:
     """Moves the state of a single optimizer to the device."""
     for p, v in optimizer.state.items():
-        optimizer.state[p] = apply_to_collection(v, Tensor, move_data_to_device, device)
+        optimizer.state[p] = apply_to_collection(v, Tensor, move_data_to_device, device, allow_frozen=True)

tests/tests_fabric/utilities/test_optimizer.py

@@ -1,4 +1,5 @@
 import collections
+import dataclasses
 
 import torch
 from torch import Tensor
@@ -7,10 +8,15 @@
 
 
 def test_optimizer_to_device():
+    @dataclasses.dataclass(frozen=True)
+    class FooState:
+        bar: int
+
     class TestOptimizer(torch.optim.SGD):
         def __init__(self, *args, **kwargs):
             super().__init__(*args, **kwargs)
             self.state["dummy"] = torch.tensor(0)
+            self.state["frozen"] = FooState(0)
 
     layer = torch.nn.Linear(32, 2)
     opt = TestOptimizer(layer.parameters(), lr=0.1)