[FSDP2] Shard eagerly initialized optimizer state

src/accelerate/accelerator.py

@@ -1718,12 +1718,18 @@ def _prepare_fsdp2(self, *args):
             if isinstance(obj, torch.optim.Optimizer):
                 for param_group in obj.param_groups:
                     for i, p in enumerate(param_group["params"]):
-                        # We drop a reference to the original param here, so that _move_states_to_device triggers a reallocation
+                        # We drop references to the original param here, so that _move_states_to_device triggers a reallocation
                         # We reassign the data_ptr to the original param, so that we preserve the mapping to the new ones
-                        param_group["params"][i] = torch.empty(1, dtype=p.dtype, device=p.device)
-                        param_group["params"][i].data_ptr = (
-                            p._local_tensor.data_ptr() if isinstance(p, DTensor) else p.data_ptr()
+                        old_parameter = p
+                        new_parameter = torch.empty(1, dtype=p.dtype, device=p.device)
+                        new_parameter.data_ptr = (
+                            old_parameter._local_tensor.data_ptr()
+                            if isinstance(old_parameter, DTensor)
+                            else old_parameter.data_ptr()
                         )
+                        param_group["params"][i] = new_parameter
+                        if old_parameter in obj.state:
+                            obj.state[new_parameter] = obj.state.pop(old_parameter)
 
         self._models.append(model)
 

src/accelerate/utils/fsdp_utils.py

@@ -658,8 +658,9 @@ def _cast_and_contiguous(tensor, to_contiguous, dtype):
 
 def fsdp2_switch_optimizer_parameters(optimizer: torch.optim.Optimizer, mapping: dict):
     """
-    Switches the parameters of the optimizer to new ones (sharded parameters in usual case). This function modifies the
-    optimizer in-place.
+    Switches the parameters and any eagerly initialized parameter state of the optimizer to new ones (sharded
+    parameters in the usual case). Parameter-shaped state tensors are distributed like their new parameter. This
+    function modifies the optimizer in-place.
 
     Args:
         optimizer (`torch.optim.Optimizer`): Optimizer instance which contains the original model parameters
@@ -671,14 +672,40 @@ def fsdp2_switch_optimizer_parameters(optimizer: torch.optim.Optimizer, mapping:
             indicates a bug. If we kept the original params instead of raising, the training wouldn't be numerically
             correct and weights wouldn't get updated.
     """
-    from torch.distributed.tensor import DTensor
+    from torch.distributed.tensor import DTensor, distribute_tensor
 
-    accessor_mapping = {}
+    def get_data_ptr(parameter):
+        if isinstance(parameter, DTensor):
+            return parameter._local_tensor.data_ptr()
+        data_ptr = parameter.data_ptr
+        return data_ptr() if callable(data_ptr) else data_ptr
 
-    accessor_mapping[DTensor] = "_local_tensor"
     try:
         for param_group in optimizer.param_groups:
-            param_group["params"] = [mapping[p.data_ptr] for p in param_group["params"]]
+            param_group["params"] = [mapping[get_data_ptr(p)] for p in param_group["params"]]
+
+        for old_parameter in list(optimizer.state):
+            if isinstance(old_parameter, torch.Tensor):
+                new_parameter = mapping[get_data_ptr(old_parameter)]
+                if old_parameter is not new_parameter:
+                    state = optimizer.state.pop(old_parameter)
+                    if isinstance(new_parameter, DTensor):
+                        for key, value in state.items():
+                            if isinstance(value, torch.Tensor) and value.shape == new_parameter.shape:
+                                if (
+                                    isinstance(value, DTensor)
+                                    and value.device_mesh == new_parameter.device_mesh
+                                    and value.placements == new_parameter.placements
+                                ):
+                                    continue
+                                if isinstance(value, DTensor):
+                                    value = value.full_tensor().detach()
+                                state[key] = distribute_tensor(
+                                    value,
+                                    device_mesh=new_parameter.device_mesh,
+                                    placements=new_parameter.placements,
+                                ).to(new_parameter.device)
+                    optimizer.state[new_parameter] = state
     except KeyError:
         # This shouldn't ever happen, but we want to fail here else training wouldn't be numerically correct
         # This basically means that we're missing a mapping from the original parameter to the sharded parameter

tests/fsdp/fsdp2_adagrad.py

@@ -0,0 +1,41 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch.distributed.tensor import DTensor
+
+from accelerate import Accelerator
+
+
+model = torch.nn.Sequential(
+    torch.nn.Linear(16, 32),
+    torch.nn.ReLU(),
+    torch.nn.Linear(32, 4),
+)
+optimizer = torch.optim.Adagrad(model.parameters(), lr=0.01)
+accelerator = Accelerator()
+model, optimizer = accelerator.prepare(model, optimizer)
+
+for parameter, state in optimizer.state.items():
+    assert isinstance(parameter, DTensor)
+    assert isinstance(state["sum"], DTensor)
+    assert state["sum"].device_mesh == parameter.device_mesh
+    assert state["sum"].placements == parameter.placements
+    assert state["sum"].device == parameter.device
+
+inputs = torch.randn(8, 16, device=accelerator.device)
+loss = model(inputs).sum()
+accelerator.backward(loss)
+optimizer.step()
+optimizer.state_dict()

tests/fsdp/test_fsdp.py

@@ -925,3 +925,22 @@ class FSDP2IntegrationTest(FSDPIntegrationTest):
     def setUp(self):
         super().setUp()
         self.current_fsdp_version = 2
+
+    def test_adagrad_with_eager_optimizer_state(self):
+        test_file_path = os.path.join(os.path.dirname(__file__), "fsdp2_adagrad.py")
+        cmd = get_launch_command(
+            num_processes=2,
+            num_machines=1,
+            machine_rank=0,
+            use_fsdp=True,
+            fsdp_version=2,
+        )
+        cmd.extend(
+            [
+                "--fsdp_reshard_after_forward=true",
+                "--fsdp_auto_wrap_policy=SIZE_BASED_WRAP",
+                "--fsdp_min_num_params=1",
+                test_file_path,
+            ]
+        )
+        execute_subprocess_async(cmd)