Repository: huggingface/transformers
Pull request: Disable delay_optimizer_creation in `Trainer` to support fsdp2

# What does this PR do?

In order to get Trainer support for FSDP2 in `accelerate`, we have to pass the model and optimizer into `Accelerator.prepare()` at the same time (https://github.com/huggingface/accelerate/pull/3394/files#r2017637611).

**Note:** This may not be sufficient for full FSDP2 support in Trainer. This PR might be scrapped and replaced with something more complete if some issues arise.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [X] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [X] Did you write any new necessary tests?


### Testing

2 unit tests were added to `tests/fsdp/test_fsdp.py`

1. `test_fsdp2_cpu_offloading`, which matches `test_fsdp_cpu_offloading`, but is **skipped** since the second test appears to be non functioning
2. `test_accelerate_fsdp2_integration` which matches `test_training_and_can_resume_normally` with SHARDED_STATE_DICT (as this is the only config that can run accelerate), and passes.

This will not be caught by CI, as it depends on an unreleased feature, but was instead manually run with the command `RUN_SLOW=1 pytest tests/fsdp/test_fsdp.py `
```
# RUN_SLOW=1 pytest tests/fsdp/test_fsdp.py 

tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_full_shard_bf16 PASSED                                                                                                                                             [  4%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_full_shard_fp16 PASSED                                                                                                                                             [  8%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_shard_grad_op_bf16 PASSED                                                                                                                                          [ 13%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_shard_grad_op_fp16 PASSED                                                                                                                                          [ 17%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_with_cpu_offload_0_bf16 PASSED                                                                                                                                     [ 21%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_with_cpu_offload_1_fp16 PASSED                                                                                                                                     [ 26%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_with_gradient_accumulation_full_shard_bf16 PASSED                                                                                                                  [ 30%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_with_gradient_accumulation_full_shard_fp16 PASSED                                                                                                                  [ 34%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_with_gradient_accumulation_shard_grad_op_bf16 PASSED                                                                                                               [ 39%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_basic_run_with_gradient_accumulation_shard_grad_op_fp16 PASSED                                                                                                               [ 43%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp_cpu_offloading SKIPPED (FSDP CPU offloading script not found!)                                                                                                          [ 47%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_training_and_can_resume_normally_FULL_STATE_DICT PASSED                                                                                                                      [ 52%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_training_and_can_resume_normally_SHARDED_STATE_DICT PASSED                                                                                                                   [ 56%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_accelerate_fsdp2_integration PASSED                                                                                                                                          [ 60%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp2_cpu_offloading SKIPPED (FSDP 2 CPU offloading script not found!)                                                                                                       [ 65%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp_config_full_shard_bf16 PASSED                                                                                                                                           [ 69%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp_config_full_shard_fp16 PASSED                                                                                                                                           [ 73%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp_config_shard_grad_op_bf16 PASSED                                                                                                                                        [ 78%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp_config_shard_grad_op_fp16 PASSED                                                                                                                                        [ 82%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp_config_transformers_auto_wrap_full_shard_bf16 PASSED                                                                                                                    [ 86%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp_config_transformers_auto_wrap_full_shard_fp16 PASSED                                                                                                                    [ 91%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp_config_transformers_auto_wrap_shard_grad_op_bf16 PASSED                                                                                                                 [ 95%]
tests/fsdp/test_fsdp.py::TrainerIntegrationFSDP::test_fsdp_config_transformers_auto_wrap_shard_grad_op_fp16 PASSED                                                                                                                 [100%]
```

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @zach-huggingface and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Changed files:
filename: src/transformers/trainer.py
patch: @@ -2313,6 +2313,11 @@ def _inner_training_loop(
 
         delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled or self.is_fsdp_enabled
 
+        # Can't delay optimizer creation when using FSDP2: https://github.com/huggingface/accelerate/blob/3f636d626063ffcf9a337c7d3624d61b7d187d59/src/accelerate/accelerator.py#L1404
+        is_fsdp2 = self.is_fsdp_enabled and (getattr(self.accelerator.state.fsdp_plugin, "fsdp_version", 1) == 2)
+        if is_fsdp2:
+            delay_optimizer_creation = False
+
         # We need to reset the scheduler, as its parameters may be different on subsequent calls
         if self._created_lr_scheduler:
             self.lr_scheduler = None

filename: tests/fsdp/test_fsdp.py
patch: @@ -109,6 +109,15 @@ def get_master_port(real_launcher=False):
     require_fsdp_version = partial(require_fsdp, min_version=FSDP_PYTORCH_VERSION)
 
 
+FSDP2_ACCELERATE_VERSION = "1.6.0"
+require_accelerate_fsdp2 = partial(require_accelerate, min_version=FSDP2_ACCELERATE_VERSION)
+require_fsdp_v2_version = require_fsdp
+if is_accelerate_available(min_version=FSDP2_ACCELERATE_VERSION):
+    from accelerate.utils.constants import FSDP2_PYTORCH_VERSION
+
+    require_fsdp_v2_version = partial(require_fsdp, min_version=FSDP2_PYTORCH_VERSION)
+
+
 def get_launcher(distributed=False, use_accelerate=False):
     # 1. explicitly set --num_nodes=1 just in case these tests end up run on a multi-node setup
     # - it won't be able to handle that
@@ -316,6 +325,73 @@ def test_fsdp_cpu_offloading(self):
         except:  # noqa
             raise AssertionError("CPU offloading failed with FSDP!")
 
+    @require_torch_multi_accelerator
+    @slow
+    @require_fsdp
+    @require_fsdp_v2_version
+    @require_accelerate_fsdp2
+    def test_accelerate_fsdp2_integration(self):
+        output_dir = self.get_auto_remove_tmp_dir("./xxx", after=False)
+        sharding_strategy = "full_shard"
+        use_accelerate = True
+
+        num_gpus = min(2, backend_device_count(torch_device))
+        master_port = get_master_port(real_launcher=True)
+        launcher = f"""accelerate launch
+            --num_processes {num_gpus}
+            --main_process_port {master_port}
+            --use_fsdp
+            --fsdp_version 2
+            --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP
+            --fsdp_state_dict_type SHARDED_STATE_DICT
+            --fsdp_transformer_layer_cls_to_wrap BertLayer""".split()
+        args = self.get_base_args(output_dir, 2, 25).split()
+        script = [f"{self.examples_dir_str}/pytorch/text-classification/run_glue.py"]
+        logs = self.run_cmd_and_get_logs(use_accelerate, sharding_strategy, launcher, script, args, output_dir)
+
+        # resume from ckpt
+        checkpoint = os.path.join(output_dir, "checkpoint-115")
+        resume_args = args + f"--resume_from_checkpoint {checkpoint}".split()
+
+        is_fsdp_ckpt = os.path.isdir(checkpoint) and (
+            # this checks the FSDP state dict when `SHARDED_STATE_DICT` is used
+            any(
+                FSDP_MODEL_NAME in folder_name
+                for folder_name in os.listdir(checkpoint)
+                if os.path.isdir(os.path.join(checkpoint, folder_name))
+            )
+            # this checks the FSDP state dict when `FULL_STATE_DICT` is used
+            or os.path.isfile(os.path.join(checkpoint, f"{FSDP_MODEL_NAME}.bin"))
+        )
+        self.assertTrue(is_fsdp_ckpt)
+
+        logs_resume = self.run_cmd_and_get_logs(
+            use_accelerate, sharding_strategy, launcher, script, resume_args, output_dir
+        )
+
+        for log, log1 in zip(logs, logs_resume):
+            if "learning_rate" in log:
+                self.assertAlmostEqual(log["learning_rate"], log1["learning_rate"], delta=1e-5)
+
+    @require_torch_multi_accelerator
+    @slow
+    @require_fsdp
+    @require_fsdp_v2_version
+    @require_accelerate_fsdp2
+    def test_fsdp2_cpu_offloading(self):
+        # TODO: This file is missing and should be added or the test should be removed
+        if not os.path.exists("utils/testing_scripts/fsdp_cpu_offloading.py"):
+            raise unittest.SkipTest("FSDP 2 CPU offloading script not found!")
+
+        try:
+            subprocess.run(
+                "accelerate launch --fsdp_version 2 utils/testing_scripts/fsdp_cpu_offloading.py --config utils/testing_scripts/dummy_fsdp_config.yml",
+                shell=True,
+                check=True,
+            )
+        except:  # noqa
+            raise AssertionError("CPU offloading failed with FSDP!")
+
     def run_cmd_and_get_logs(self, use_accelerate, sharding_strategy, launcher, script, args, output_dir):
         if not use_accelerate:
             fsdp_args = [