Revert "remove unnecessary sync point in AveragedModel update (pytorch#158017)"

This reverts commit cb7f45f.

Reverted pytorch#158017 on behalf of https://github.com/wdvr due to discussed with author - expecting this to break checkpointing ([comment](pytorch#158017 (comment)))

Author: pytorchmergebot

test/optim/test_swa_utils.py

@@ -76,6 +76,7 @@ def _test_averaged_model(self, net_device, swa_device, ema):
             # Check that AveragedModel is on the correct device
             self.assertTrue(p_swa.device == swa_device)
             self.assertTrue(p_avg.device == net_device)
+        self.assertTrue(averaged_dnn.n_averaged.device == swa_device)
 
     def _run_averaged_steps(self, dnn, swa_device, ema):
         ema_decay = 0.999
@@ -149,44 +150,6 @@ def test_averaged_model_state_dict(self):
             self.assertEqual(p_swa, p_swa2)
         self.assertTrue(averaged_dnn.n_averaged == averaged_dnn2.n_averaged)
 
-    def test_averaged_model_backward_compatibility(self):
-        """Test that AveragedModel correctly handles old checkpoints with tensor n_averaged."""
-        dnn = torch.nn.Sequential(
-            torch.nn.Conv2d(1, 5, kernel_size=3), torch.nn.Linear(5, 10)
-        )
-        averaged_dnn = AveragedModel(dnn)
-
-        # Update parameters a few times
-        n_updates = 5
-        for _ in range(n_updates):
-            for p in dnn.parameters():
-                p.detach().add_(torch.randn_like(p))
-            averaged_dnn.update_parameters(dnn)
-
-        # Manually create a state dict with tensor n_averaged (simulating old checkpoint)
-        state_dict = averaged_dnn.state_dict()
-        # Create an old-style tensor n_averaged
-        old_n_averaged = torch.tensor(n_updates, dtype=torch.long)
-        state_dict["n_averaged"] = old_n_averaged
-
-        # Create new model and load the old-style state dict
-        averaged_dnn2 = AveragedModel(dnn)
-        averaged_dnn2.load_state_dict(state_dict)
-
-        # Check that n_averaged was correctly loaded as a Python int
-        self.assertEqual(averaged_dnn2.n_averaged, n_updates)
-        self.assertIsInstance(averaged_dnn2.n_averaged, int)
-
-        # Verify that parameters are correctly loaded
-        for p_swa, p_swa2 in zip(averaged_dnn.parameters(), averaged_dnn2.parameters()):
-            self.assertEqual(p_swa, p_swa2)
-
-        # Test that we can continue to update parameters without issues
-        for p in dnn.parameters():
-            p.detach().add_(torch.randn_like(p))
-        averaged_dnn2.update_parameters(dnn)
-        self.assertEqual(averaged_dnn2.n_averaged, n_updates + 1)
-
     def test_averaged_model_default_avg_fn_picklable(self):
         dnn = torch.nn.Sequential(
             torch.nn.Conv2d(1, 5, kernel_size=3),

torch/optim/swa_utils.py

@@ -116,28 +116,6 @@ def swa_update(
     return swa_update
 
 
-def _load_state_dict_pre_hook(
-    module,
-    state_dict,
-    prefix,
-    local_metadata,
-    strict,
-    missing_keys,
-    unexpected_keys,
-    error_msgs,
-):
-    """Pre-hook to handle backward compatibility with tensor n_averaged."""
-    # Check if the old tensor n_averaged is present in the state dict
-    n_averaged_key = prefix + "n_averaged"
-    if n_averaged_key in state_dict:
-        # Convert tensor n_averaged to Python int for backward compatibility
-        n_averaged_tensor = state_dict[n_averaged_key]
-        if isinstance(n_averaged_tensor, Tensor):
-            module.n_averaged = int(n_averaged_tensor.item())
-        # Remove the old tensor buffer from state_dict to avoid loading it
-        del state_dict[n_averaged_key]
-
-
 class AveragedModel(Module):
     r"""Implements averaged model for Stochastic Weight Averaging (SWA) and Exponential Moving Average (EMA).
 
@@ -237,7 +215,7 @@ class AveragedModel(Module):
         https://paperswithcode.com/method/polyak-averaging
     """
 
-    n_averaged: int
+    n_averaged: Tensor
 
     def __init__(
         self,
@@ -256,25 +234,17 @@ def __init__(
         self.module = deepcopy(model)
         if device is not None:
             self.module = self.module.to(device)
-        self.n_averaged = 0
+        self.register_buffer(
+            "n_averaged", torch.tensor(0, dtype=torch.long, device=device)
+        )
         self.avg_fn = avg_fn
         self.multi_avg_fn = multi_avg_fn
         self.use_buffers = use_buffers
-        self.register_load_state_dict_pre_hook(_load_state_dict_pre_hook)
 
     def forward(self, *args, **kwargs):
         """Forward pass."""
         return self.module(*args, **kwargs)
 
-    def get_extra_state(self) -> Any:
-        """Get extra state for serialization."""
-        return {"n_averaged": self.n_averaged}
-
-    def set_extra_state(self, state: Any) -> None:
-        """Set extra state from deserialization."""
-        if isinstance(state, dict) and "n_averaged" in state:
-            self.n_averaged = state["n_averaged"]
-
     def update_parameters(self, model: Module):
         """Update model parameters."""
         self_param = (
@@ -310,26 +280,28 @@ def update_parameters(self, model: Module):
                         self.multi_avg_fn(
                             self_params,  # type: ignore[arg-type]
                             model_params,  # type: ignore[arg-type]
-                            self.n_averaged,
+                            self.n_averaged.to(device),
                         )
                     elif (
                         device is not None
                         and device.type in _get_foreach_kernels_supported_devices()
                     ):
                         multi_avg_fn = get_swa_multi_avg_fn()
-                        multi_avg_fn(self_params, model_params, self.n_averaged)
+                        multi_avg_fn(
+                            self_params, model_params, self.n_averaged.to(device)
+                        )
                     else:
                         avg_fn = get_swa_avg_fn()
+                        n_averaged = self.n_averaged.to(device)
                         for p_averaged, p_model in zip(self_params, model_params):  # type: ignore[assignment]
-                            p_averaged.copy_(
-                                avg_fn(p_averaged, p_model, self.n_averaged)
-                            )
+                            p_averaged.copy_(avg_fn(p_averaged, p_model, n_averaged))
             else:
                 for p_averaged, p_model in zip(  # type: ignore[assignment]
                     self_param_detached, model_param_detached
                 ):
+                    n_averaged = self.n_averaged.to(p_averaged.device)
                     p_averaged.detach().copy_(
-                        self.avg_fn(p_averaged.detach(), p_model, self.n_averaged)
+                        self.avg_fn(p_averaged.detach(), p_model, n_averaged)
                     )
 
         if not self.use_buffers: