Trainer.train() using v5 + FSDP2 + PEFT + cpu_ram_efficient_loading=True results in wrong optimizer states/params on all but 0th rank
Author: d5031
Pull request: no
Labels: name: bug

### System Info

v5 + FSDP2 + cpu_ram_efficient_loading=True + PEFT + "reshard_after_forward": True

### Who can help?
@S1ro1 
@SunMarc 


### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

Sorry, can't paste a small reproducer (would require some work), probably simpler just to modify one of existing tests and i'm almost sure it would result in the same issue (my code doesn't have anything too specific / custom)

In particular fails during optimizer saving in Trainer._save_checkpoint:

        if not self.args.save_only_model:
            # Save optimizer and scheduler
            self._save_optimizer_and_scheduler(output_dir)
 
Likely the reason is that Accelerator._prepare_fsdp2 has
 
        def _prepare_fsdp2(self, *args):
        # First pass: prepare everything except schedulers (and model, which is prepared separately below)
        result = [
            self._prepare_one(obj, first_pass=True) if not isinstance(obj, torch.nn.Module) else obj for obj in args
        ]
        .....
        # Prepare everything FSDP2 related for the model (except AC)
        model = fsdp2_prepare_model(self, model) # which in turn calls fsdp2_load_full_state_dict
        ....

So , optimizer is created *on all ranks* before fsdp2_load_full_state_dict call (which actually materializes params)

Can possibly be reproduced without using transormers (just within accelerate), but not sure/didn't try.

### Expected behavior

Should be no diff in training / behaviour vs cpu_ram_efficient_loading=False (which works completely fine, keeping rest the same/untouched)

Comments:
Very likely issue is around this place as data_ptr() == 0 for meta device

       for obj in result:
            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 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()
                        )


Q to experts (@S1ro1 ?) how to fix it. 
Not running this code at all, fails, as i guess expected (i saw the comment, but afraid for non-expert in the code it's a bit cryptic "# We drop a reference to the original param here, so that _move_states_to_device triggers a reallocation").

(running  it after fsdp2_prepare_model() won't make any diff, as optimizer is created already )

Thank you

I've absolutely horribly hacked it the way below (and it works). Unfortunately, i don't understand current codebase well enough, as what's there also seems pretty horrible to me :-) Feels like there should be a better solution than relying on data_ptr (with loads of caveats here are there, like DTensor in one place; FP8BackendType.AO in another, etc)

NOT a suggested fix; rather a sanity-check that problem is indeed in what i described (now it all works).

in _prepare_fsdp2:

        for obj in result:
            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 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)

                        # super hacky stuff
                        if p.device.type == 'meta':
                            data_ptr = id(p)
                        else:
                            data_ptr = p._local_tensor.data_ptr() if isinstance(p, DTensor) else p.data_ptr()
                        param_group["params"][i].data_ptr = data_ptr

in _get_named_parameters:

                is_on_meta = [p.device.type == 'meta' for p in obj.parameters()]
                assert all(is_on_meta) or (not any(is_on_meta))
                is_on_meta = all(is_on_meta)
                if is_on_meta:
                    named_parameters.update(
                        {
                            n: id(p)
                            for n, p in obj.named_parameters()
                        }
                    )

                else:
                    named_parameters.update(
                        {
                            n: getattr(p, accessor_mapping[type(p)]).data_ptr()
                            if type(p) in accessor_mapping
                            else p.data_ptr()
                            for n, p in obj.named_parameters()
                        }
                    )

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.