Migrate state_dict bc test to OptimizerInfo, increase coverage (pytorch#116500)

Pull Request resolved: pytorch#116500
Approved by: https://github.com/albanD

Author: janeyx99

test/optim/test_optim.py

@@ -247,28 +247,6 @@ def fn_base(optimizer, weight, bias):
         for _i in range(20):
             optimizer.step(fn)
 
-        # Make sure that optimizers that support maximize can load older models
-        old_state_dict = deepcopy(optimizer.state_dict())
-        state_dict_no_maximize = deepcopy(optimizer.state_dict())
-        if "maximize" in state_dict_no_maximize["param_groups"][0]:
-            for group in state_dict_no_maximize["param_groups"]:
-                del group["maximize"]
-            optimizer.load_state_dict(state_dict_no_maximize)
-            # Make sure we can still step
-            optimizer.step()
-            # Undo these changes before proceeding!
-            optimizer.load_state_dict(old_state_dict)
-        # Make sure that optimizers that support foreach can load older models
-        state_dict_no_foreach = deepcopy(optimizer.state_dict())
-        if "foreach" in state_dict_no_foreach["param_groups"][0]:
-            for group in state_dict_no_foreach["param_groups"]:
-                del group["foreach"]
-            optimizer.load_state_dict(state_dict_no_foreach)
-            # Make sure we can still step
-            optimizer.step()
-            # Undo these changes before proceeding!
-            optimizer.load_state_dict(old_state_dict)
-
         # Make sure that loading optimizers with step not wrapped in tensor can work
         state_dict = optimizer.state_dict()
         if "step" in state_dict["state"][0] and torch.is_tensor(

test/test_optim.py

@@ -448,6 +448,53 @@ def fwd_bwd(optim, w, b, i):
                 optimizer_c.state_dict()["param_groups"][-1]
             )
 
+    @optims(optim_db, dtypes=[torch.float32])
+    def test_can_load_older_state_dict(self, device, dtype, optim_info):
+        new_flags = ["maximize", "foreach", "fused", "differentiable", "capturable"]
+        optim_cls = optim_info.optim_cls
+
+        # Skip differentiable testing for now, see https://github.com/pytorch/pytorch/issues/116490
+        all_optim_inputs = _get_optim_inputs_including_global_cliquey_kwargs(device, dtype, optim_info, skip=("differentiable",))
+        for optim_input in all_optim_inputs:
+            torch.manual_seed(1)
+            model = torch.nn.Sequential(
+                torch.nn.Conv2d(4, 2, 1, stride=2),
+                torch.nn.BatchNorm2d(2, eps=1e-05, momentum=0.1),
+            )
+            model.to(dtype=dtype, device=device)
+            input = torch.rand(1, 4, 16, 16, device=device, dtype=dtype)
+            optimizer = optim_cls(model.parameters(), **optim_input.kwargs)
+
+            def fwd_bwd(optim, mod, i):
+                optim.zero_grad()
+                loss = mod(i).sum()
+                loss.backward()
+                return loss
+
+            for _ in range(3):
+                if optim_cls.__name__ == "LBFGS":
+                    optimizer.step(functools.partial(fwd_bwd, optimizer, model, input))
+                else:
+                    fwd_bwd(optimizer, model, input)
+                    optimizer.step()
+
+            # old_state_dict has all new flags del'd
+            old_state_dict = deepcopy(optimizer.state_dict())
+            old_state_dict_pg = old_state_dict["param_groups"]
+            for group in old_state_dict_pg:
+                for flag in new_flags:
+                    if flag in group:
+                        del group[flag]
+
+            optimizer.load_state_dict(old_state_dict)
+
+            # Make sure we can still step
+            if optim_cls.__name__ == "LBFGS":
+                optimizer.step(functools.partial(fwd_bwd, optimizer, model, input))
+            else:
+                fwd_bwd(optimizer, model, input)
+                optimizer.step()
+
 
 instantiate_device_type_tests(TestOptimRenewed, globals(), allow_mps=True)
 

torch/testing/_internal/common_optimizers.py

@@ -1110,6 +1110,12 @@ def _get_optim_inputs_including_global_cliquey_kwargs(
         step_requires_closure=True,
         supports_param_groups=False,
         supports_multiple_devices=False,
+        skips=(
+            # Fails on MacOS 13.2.1 in CI https://github.com/pytorch/pytorch/issues/117094
+            DecorateInfo(
+                skipIfMps, "TestOptimRenewed", "test_can_load_older_state_dict"
+            ),
+        ),
     ),
     OptimizerInfo(
         NAdam,
@@ -1138,6 +1144,14 @@ def _get_optim_inputs_including_global_cliquey_kwargs(
                 "TestOptimRenewed",
                 "test_state_dict_deterministic",
             ),
+            DecorateInfo(
+                skipIfTorchDynamo(
+                    "See https://github.com/pytorch/pytorch/issues/116499"
+                ),
+                "TestOptimRenewed",
+                "test_can_load_older_state_dict",
+                device_type="cuda",
+            ),
         ),
     ),
     OptimizerInfo(
@@ -1310,6 +1324,13 @@ def _get_optim_inputs_including_global_cliquey_kwargs(
                 "TestOptimRenewed",
                 "test_state_dict_deterministic",
             ),
+            DecorateInfo(
+                unittest.skip(
+                    "SparseAdam does not support dense gradients, see #116507"
+                ),
+                "TestOptimRenewed",
+                "test_can_load_older_state_dict",
+            ),
         ),
     ),
 ]