Adding BC tests: same results with named_parameters, and satet_dict save-load equality

ErezYosef · ErezYosef · commit 3f0abc32ad77 · 2024-10-06T15:27:54.000+03:00
diff --git a/test/test_optim.py b/test/test_optim.py
@@ -1356,6 +1356,16 @@ def test_state_dict_deterministic(self, device, dtype, optim_info):
         input = torch.randn(3, requires_grad=True, device=device, dtype=dtype)
         params = [weight, bias]
 
+        def make_param_and_named_param(param):
+            named_param = [(f'name{i}', p) for i, p in enumerate(param)]
+            return [param, named_param]
+
+        def without_param_names(state_dict):
+            new_state_dict = deepcopy(state_dict)
+            for pg in new_state_dict['param_groups']:
+                pg.pop('param_names', None)
+            return new_state_dict
+
         def fwd_bwd(optim, w, b, i):
             optim.zero_grad()
             loss = (w.mv(i) + b).pow(2).sum()
@@ -1368,51 +1378,55 @@ def fwd_bwd(optim, w, b, i):
             return loss
 
         for optim_input in all_optim_inputs:
-            optimizer = optim_cls(params, **optim_input.kwargs)
-            closure = functools.partial(fwd_bwd, optimizer, weight, bias, input)
-
-            # Prime the optimizer
-            for _ in range(10):
-                if optim_info.step_requires_closure:
-                    optimizer.step(closure)
-                else:
-                    closure()
-                    optimizer.step()
+            for param_in in make_param_and_named_param(params):
+                optimizer = optim_cls(param_in, **optim_input.kwargs)
+                closure = functools.partial(fwd_bwd, optimizer, weight, bias, input)
 
-            # Clone the weights and construct a new optimizer for them
-            with torch.no_grad():
-                weight_c = Parameter(weight.clone())
-                bias_c = Parameter(bias.clone())
-
-            optimizer_c = optim_cls([weight_c, bias_c], **optim_input.kwargs)
-            closure_c = functools.partial(fwd_bwd, optimizer_c, weight_c, bias_c, input)
-
-            # Load the state dict from the original optimizer into the new one
-            optimizer_c.load_state_dict(deepcopy(optimizer.state_dict()))
+                # Prime the optimizer
+                for _ in range(10):
+                    if optim_info.step_requires_closure:
+                        optimizer.step(closure)
+                    else:
+                        closure()
+                        optimizer.step()
 
-            # Run both optimizers in parallel
-            for _ in range(10):
-                if optim_info.step_requires_closure:
-                    optimizer.step(closure)
-                    optimizer_c.step(closure_c)
-                else:
-                    closure()
-                    closure_c()
-                    optimizer.step()
-                    optimizer_c.step()
+                for param_c_index in range(2):
+                    # Clone the weights and construct a new optimizer for them
+                    with torch.no_grad():
+                        weight_c = Parameter(weight.clone())
+                        bias_c = Parameter(bias.clone())
+                    param_c = make_param_and_named_param([weight_c, bias_c])[param_c_index]
+                    optimizer_c = optim_cls(param_c, **optim_input.kwargs)
+                    closure_c = functools.partial(fwd_bwd, optimizer_c, weight_c, bias_c, input)
+
+                    # Load the state dict from the original optimizer into the new one
+                    optimizer_c.load_state_dict(deepcopy(optimizer.state_dict()))
+
+                    # Run both optimizers in parallel
+                    for _ in range(10):
+                        if optim_info.step_requires_closure:
+                            optimizer.step(closure)
+                            optimizer_c.step(closure_c)
+                        else:
+                            closure()
+                            closure_c()
+                            optimizer.step()
+                            optimizer_c.step()
 
-                self.assertEqual(weight, weight_c)
-                self.assertEqual(bias, bias_c)
+                        self.assertEqual(weight, weight_c)
+                        self.assertEqual(bias, bias_c)
 
-            # Make sure state dict is deterministic with equal (not identical) parameters
-            self.assertEqual(optimizer.state_dict(), optimizer_c.state_dict())
+                    # Make sure state dict is deterministic with equal (not identical) parameters
+                    # Param names are optional and not needed to be the consistent.
+                    self.assertEqual(without_param_names(optimizer.state_dict()),
+                                     without_param_names(optimizer_c.state_dict()))
 
-            # Make sure repeated parameters have identical representation (see #36831)
-            optimizer_c.param_groups.extend(optimizer_c.param_groups)
-            self.assertEqual(
-                optimizer.state_dict()["param_groups"][-1],
-                optimizer_c.state_dict()["param_groups"][-1],
-            )
+                    # Make sure repeated parameters have identical representation (see #36831)
+                    optimizer_c.param_groups.extend(optimizer_c.param_groups)
+                    self.assertEqual(
+                        without_param_names(optimizer.state_dict())["param_groups"][-1],
+                        without_param_names(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):
@@ -1538,6 +1552,10 @@ def test_save_load_equality_with_weights_only(self, device, dtype, optim_info):
         input = torch.randn(3, requires_grad=True, device=device, dtype=dtype)
         params = [weight, bias]
 
+        def make_param_and_named_param(param):
+            named_param = [(f'name{i}', p) for i, p in enumerate(param)]
+            return [param, named_param]
+
         def fwd_bwd(optim, w, b, i):
             optim.zero_grad()
             loss = (w.mv(i) + b).pow(2).sum()
@@ -1548,25 +1566,26 @@ def fwd_bwd(optim, w, b, i):
             return loss
 
         for optim_input in all_optim_inputs:
-            optimizer = optim_cls(params, **optim_input.kwargs)
-            closure = functools.partial(fwd_bwd, optimizer, weight, bias, input)
+            for params_in in make_param_and_named_param(params):
+                optimizer = optim_cls(params_in, **optim_input.kwargs)
+                closure = functools.partial(fwd_bwd, optimizer, weight, bias, input)
 
-            # Prime the optimizer
-            for _ in range(3):
-                optimizer.step(closure)
+                # Prime the optimizer
+                for _ in range(3):
+                    optimizer.step(closure)
 
-            sd = optimizer.state_dict()
-
-            # === Check saved/loaded state_dict are the same (including weights_only load). ===
-            with tempfile.TemporaryFile() as f:
-                torch.save(sd, f)
-                f.seek(0)
-                sd_copy = torch.load(f)
-                self.assertEqual(sd_copy, sd)
-                del sd_copy
-                f.seek(0)
-                sd_copy_wo = torch.load(f, weights_only=True)
-                self.assertEqual(sd_copy_wo, sd)
+                sd = optimizer.state_dict()
+
+                # === Check saved/loaded state_dict are the same (including weights_only load). ===
+                with tempfile.TemporaryFile() as f:
+                    torch.save(sd, f)
+                    f.seek(0)
+                    sd_copy = torch.load(f)
+                    self.assertEqual(sd_copy, sd)
+                    del sd_copy
+                    f.seek(0)
+                    sd_copy_wo = torch.load(f, weights_only=True)
+                    self.assertEqual(sd_copy_wo, sd)
 
     @optims(optim_db, dtypes=[torch.float32])
     def test_load_nontensor_step(self, device, dtype, optim_info):
