diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py index 3b8829fabfb..2c755b5d103 100755 --- a/src/accelerate/accelerator.py +++ b/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) diff --git a/src/accelerate/utils/fsdp_utils.py b/src/accelerate/utils/fsdp_utils.py index df4df42ad84..706721c5b21 100644 --- a/src/accelerate/utils/fsdp_utils.py +++ b/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 diff --git a/tests/fsdp/fsdp2_adagrad.py b/tests/fsdp/fsdp2_adagrad.py new file mode 100644 index 00000000000..48fdfdb1d23 --- /dev/null +++ b/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() diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py index 29e1f08b54f..c98e5ff23ef 100644 --- a/tests/fsdp/test_fsdp.py +++ b/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)