add support to pass the FSDPPlugin to training args to pass to accele…

…rate

src/transformers/trainer.py

@@ -4705,7 +4705,7 @@ def _add_sm_patterns_to_gitignore(self) -> None:
             self.repo.git_commit("Add *.sagemaker patterns to .gitignore.")
             self.repo.git_push()
 
-    def create_accelerator_and_postprocess(self):
+    def build_accelerator_args(self):
         grad_acc_kwargs = {}
         if is_accelerate_available("0.28.0") and self.args.accelerator_config.gradient_accumulation_kwargs is not None:
             grad_acc_kwargs = self.args.accelerator_config.gradient_accumulation_kwargs
@@ -4758,6 +4758,10 @@ def create_accelerator_and_postprocess(self):
         else:
             args.update(accelerator_config)
 
+        return args
+
+    def create_accelerator_and_postprocess(self):
+        args = self.build_accelerator_args()
         # create accelerator object
         self.accelerator = Accelerator(**args)
         # some Trainer classes need to use `gather` instead of `gather_for_metrics`, thus we store a flag

src/transformers/trainer_pt_utils.py

@@ -1250,6 +1250,8 @@ class AcceleratorConfig:
             Whether or not to use a pre-configured `AcceleratorState` or `PartialState` defined
             before calling `TrainingArguments`. If `True`, an `Accelerator` or `PartialState`
             must be initialized. May lead to issues using sweeps or hyperparameter tuning.
+        fsdp_plugin (`FullyShardedDataParallelPlugin`, *optional*):
+            The FSDP configuration to use for FSDP training.
 
     """
 
@@ -1320,6 +1322,13 @@ class AcceleratorConfig:
         },
     )
 
+    fsdp_plugin: Optional["FullyShardedDataParallelPlugin"] = field(
+        default=None,
+        metadata={
+            "help": "The FSDP configuration to use for FSDP training."
+        },
+    )
+
     @classmethod
     def from_json_file(cls, json_file):
         # Check if exists