From d87efb25c98741501fbf6da0d270fc181611b795 Mon Sep 17 00:00:00 2001
From: Adam Louly <adamlouly3@gmail.com>
Date: Mon, 11 Mar 2024 19:31:03 -0700
Subject: [PATCH] Initialize AcceleratorConfig in optimum TrainingArgs. (#1750)

* add accelerator config

* make style
---
 optimum/onnxruntime/training_args.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/optimum/onnxruntime/training_args.py b/optimum/onnxruntime/training_args.py
index 7c3171855c1..4f71c021cb1 100644
--- a/optimum/onnxruntime/training_args.py
+++ b/optimum/onnxruntime/training_args.py
@@ -48,6 +48,9 @@
 if is_torch_available():
     import torch
 
+if is_accelerate_available():
+    from transformers.trainer_pt_utils import AcceleratorConfig
+
 
 class ORTOptimizerNames(ExplicitEnum):
     """
@@ -446,6 +449,30 @@ def __post_init__(self):
             os.environ[f"{prefix}SYNC_MODULE_STATES"] = self.fsdp_config.get("sync_module_states", "true")
             os.environ[f"{prefix}USE_ORIG_PARAMS"] = self.fsdp_config.get("use_orig_params", "false")
 
+        if is_accelerate_available():
+            if not isinstance(self.accelerator_config, (AcceleratorConfig)):
+                if self.accelerator_config is None:
+                    self.accelerator_config = AcceleratorConfig()
+                elif isinstance(self.accelerator_config, dict):
+                    self.accelerator_config = AcceleratorConfig(**self.accelerator_config)
+                else:
+                    self.accelerator_config = AcceleratorConfig.from_json_file(self.accelerator_config)
+            if self.dispatch_batches is not None:
+                warnings.warn(
+                    "Using `--dispatch_batches` is deprecated and will be removed in version 4.41 of 🤗 Transformers. Use"
+                    " `--accelerator_config {'dispatch_batches':VALUE} instead",
+                    FutureWarning,
+                )
+                self.accelerator_config.dispatch_batches = self.dispatch_batches
+
+            if self.split_batches is not None:
+                warnings.warn(
+                    "Using `--split_batches` is deprecated and will be removed in version 4.41 of 🤗 Transformers. Use"
+                    " `--accelerator_config {'split_batches':VALUE} instead",
+                    FutureWarning,
+                )
+                self.accelerator_config.split_batches = self.split_batches
+
         if self.tpu_metrics_debug:
             warnings.warn(
                 "using `--tpu_metrics_debug` is deprecated and will be removed in version 5 of 🤗 Transformers. Use"