Add XLA Dynamo backends for training and inference (#2892)

huggingface · Jul 3, 2024 · 57a4c74 · 57a4c74
1 parent 404510a
commit 57a4c74
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
@@ -37,6 +37,8 @@
     "FX2TRT",
     "ONNXRT",
     "TENSORRT",
+    "AOT_TORCHXLA_TRACE_ONCE",
+    "TORHCHXLA_TRACE_ONCE",
     "IPEX",
     "TVM",
 ]

diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
@@ -574,6 +574,10 @@ class DynamoBackend(str, BaseEnum):
         - **ONNXRT** -- Uses ONNXRT for inference on CPU/GPU. Inference only. [Read more](https://onnxruntime.ai/)
         - **TENSORRT** -- Uses ONNXRT to run TensorRT for inference optimizations. [Read
           more](https://github.com/onnx/onnx-tensorrt)
+        - **AOT_TORCHXLA_TRACE_ONCE** -- Uses Pytorch/XLA with TorchDynamo optimization, for training. [Read
+          more](https://github.com/pytorch/xla/blob/r2.0/docs/dynamo.md)
+        - **TORCHXLA_TRACE_ONCE** -- Uses Pytorch/XLA with TorchDynamo optimization, for inference. [Read
+          more](https://github.com/pytorch/xla/blob/r2.0/docs/dynamo.md)
         - **IPEX** -- Uses IPEX for inference on CPU. Inference only. [Read
           more](https://github.com/intel/intel-extension-for-pytorch).
         - **TVM** -- Uses Apach TVM for inference optimizations. [Read more](https://tvm.apache.org/)
@@ -592,6 +596,8 @@ class DynamoBackend(str, BaseEnum):
     FX2TRT = "FX2TRT"
     ONNXRT = "ONNXRT"
     TENSORRT = "TENSORRT"
+    AOT_TORCHXLA_TRACE_ONCE = "AOT_TORCHXLA_TRACE_ONCE"
+    TORCHXLA_TRACE_ONCE = "TORCHXLA_TRACE_ONCE"
     IPEX = "IPEX"
     TVM = "TVM"