Repository: huggingface/accelerate

Issue:
TorchDynamo backends
### System Info

```Shell
accelerate v0.22.0
```


### Information

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

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [X] My own task or dataset (give details below)

### Reproduction

There seems to be a discrepancy between the backends listed on the torch doc and huggingface doc

According to the [doc](https://huggingface.co/docs/transformers/perf_train_gpu_one#using-torchcompile:~:text=Training%20%26%20inference%20backends%3A):
> **Training & inference backends:**  
dynamo.optimize("inductor") - Uses TorchInductor backend with AotAutograd and cudagraphs by leveraging codegened Triton kernels Read more  
dynamo.optimize("nvfuser") - nvFuser with TorchScript. Read more  
dynamo.optimize("aot_nvfuser") - nvFuser with AotAutograd. Read more  
dynamo.optimize("aot_cudagraphs") - cudagraphs with AotAutograd. Read more

But in the torch [doc](https://pytorch.org/docs/stable/dynamo/get-started.html#:~:text=Training%20%26%20inference%20backends%3A):
> **Training & inference backends:**  
torch.compile(m, backend="inductor") - Uses TorchInductor backend. Read more  
torch.compile(m, backend="aot_ts_nvfuser") - nvFuser with AotAutograd/TorchScript. Read more  
torch.compile(m, backend=""nvprims_nvfuser") - nvFuser with PrimTorch. Read more  
torch.compile(m, backend="cudagraphs") - cudagraphs with AotAutograd. Read more

Only `inductor` is common. Setting any other backend leads to an error. For example, setting `nvfuser` in the accelerate config results in:
```
Traceback (most recent call last):                                                                                                                                                                 
  File "example_script.py", line 314, in <module>
    main()                                                                                                                                                                                                 
  File "example_script.py", line 49, in main
    model, optimizer, lr_scheduler, train_dataloader, eval_dataloader = accelerator.prepare(model, optimizer, lr_scheduler, train_dataloader, eval_dataloader)
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/site-packages/accelerate/accelerator.py", line 1270, in prepare
    result = tuple(                                                                                                                                          
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/site-packages/accelerate/accelerator.py", line 1271, in <genexpr>                      
    self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)                             
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/site-packages/accelerate/accelerator.py", line 1083, in _prepare_one
    return self.prepare_model(obj, device_placement=device_placement)                                                                   
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/site-packages/accelerate/accelerator.py", line 1440, in prepare_model
    model = torch.compile(model, **self.state.dynamo_plugin.to_kwargs())                                                                   
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/site-packages/torch/__init__.py", line 1441, in compile
    return torch._dynamo.optimize(backend=backend, nopython=fullgraph, dynamic=dynamic, disable=disable)(model)                             
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 424, in optimize
    backend = get_compiler_fn(backend)                                                                                        
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 363, in get_compiler_fn
    compiler_fn = lookup_backend(compiler_fn)                                                                                           
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/site-packages/torch/_dynamo/backends/registry.py", line 59, in lookup_backend
    compiler_fn = _BACKENDS[compiler_fn]                                                                                                       
KeyError: 'nvfuser' 
```

and using a backend specified in the torch doc also results in an error:
```
Traceback (most recent call last):
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/site-packages/accelerate/utils/launch.py", line 90, in prepare_simple_launcher_cmd_env
    dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/enum.py", line 385, in __call__
    return cls.__new__(cls, value)
  File "/home/miniconda3/envs/hf-torch2/lib/python3.10/enum.py", line 710, in __new__
    raise ve_exc
ValueError: 'AOT_TS_NVFUSER' is not a valid DynamoBackend
```

### Expected behavior

Accelerate should support the torchdynamo backends.


Discussion:

Resolving patch:
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
--- a/src/accelerate/commands/config/config_utils.py
+++ b/src/accelerate/commands/config/config_utils.py
@@ -30,13 +30,15 @@
     "EAGER",
     "AOT_EAGER",
     "INDUCTOR",
-    "NVFUSER",
-    "AOT_NVFUSER",
-    "AOT_CUDAGRAPHS",
+    "AOT_TS_NVFUSER",
+    "NVPRIMS_NVFUSER",
+    "CUDAGRAPHS",
     "OFI",
     "FX2TRT",
     "ONNXRT",
+    "TENSORRT",
     "IPEX",
+    "TVM",
 ]
 
 
diff --git a/src/accelerate/utils/dataclasses.py b/src/accelerate/utils/dataclasses.py
--- a/src/accelerate/utils/dataclasses.py
+++ b/src/accelerate/utils/dataclasses.py
@@ -200,6 +200,29 @@ def __post_init__(self):
             raise ValueError("`amax_compute_algo` must be 'max' or 'most_recent'")
 
 
+class EnumWithContains(enum.EnumMeta):
+    "A metaclass that adds the ability to check if `self` contains an item with the `in` operator"
+
+    def __contains__(cls, item):
+        try:
+            cls(item)
+        except ValueError:
+            return False
+        return True
+
+
+class BaseEnum(enum.Enum, metaclass=EnumWithContains):
+    "An enum class that can get the value of an item with `str(Enum.key)`"
+
+    def __str__(self):
+        return self.value
+
+    @classmethod
+    def list(cls):
+        "Method to list all the possible items in `cls`"
+        return list(map(str, cls))
+
+
 class DistributedType(str, enum.Enum):
     """
     Represents a type of distributed environment.
@@ -259,7 +282,7 @@ class ComputeEnvironment(str, enum.Enum):
     AMAZON_SAGEMAKER = "AMAZON_SAGEMAKER"
 
 
-class DynamoBackend(str, enum.Enum):
+class DynamoBackend(str, BaseEnum):
     """
     Represents a dynamo backend (see https://github.com/pytorch/torchdynamo).
 
@@ -273,19 +296,21 @@ class DynamoBackend(str, enum.Enum):
         - **INDUCTOR** -- Uses TorchInductor backend with AotAutograd and cudagraphs by leveraging codegened Triton
           kernels. [Read
           more](https://dev-discuss.pytorch.org/t/torchinductor-a-pytorch-native-compiler-with-define-by-run-ir-and-symbolic-shapes/747)
-        - **NVFUSER** -- nvFuser with TorchScript. [Read
+        - **AOT_TS_NVFUSER** -- nvFuser with AotAutograd/TorchScript. [Read
           more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593)
-        - **AOT_NVFUSER** -- nvFuser with AotAutograd. [Read
+        - **NVPRIMS_NVFUSER** -- nvFuser with PrimTorch. [Read
           more](https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-1-nvfuser-and-its-primitives/593)
-        - **AOT_CUDAGRAPHS** -- cudagraphs with AotAutograd. [Read
-          more](https://github.com/pytorch/torchdynamo/pull/757)
+        - **CUDAGRAPHS** -- cudagraphs with AotAutograd. [Read more](https://github.com/pytorch/torchdynamo/pull/757)
         - **OFI** -- Uses Torchscript optimize_for_inference. Inference only. [Read
           more](https://pytorch.org/docs/stable/generated/torch.jit.optimize_for_inference.html)
         - **FX2TRT** -- Uses Nvidia TensorRT for inference optimizations. Inference only. [Read
           more](https://github.com/pytorch/TensorRT/blob/master/docsrc/tutorials/getting_started_with_fx_path.rst)
         - **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)
         - **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/)
 
     """
 
@@ -294,36 +319,15 @@ class DynamoBackend(str, enum.Enum):
     EAGER = "EAGER"
     AOT_EAGER = "AOT_EAGER"
     INDUCTOR = "INDUCTOR"
-    NVFUSER = "NVFUSER"
-    AOT_NVFUSER = "AOT_NVFUSER"
-    AOT_CUDAGRAPHS = "AOT_CUDAGRAPHS"
+    AOT_TS_NVFUSER = "AOT_TS_NVFUSER"
+    NVPRIMS_NVFUSER = "NVPRIMS_NVFUSER"
+    CUDAGRAPHS = "CUDAGRAPHS"
     OFI = "OFI"
     FX2TRT = "FX2TRT"
     ONNXRT = "ONNXRT"
+    TENSORRT = "TENSORRT"
     IPEX = "IPEX"
-
-
-class EnumWithContains(enum.EnumMeta):
-    "A metaclass that adds the ability to check if `self` contains an item with the `in` operator"
-
-    def __contains__(cls, item):
-        try:
-            cls(item)
-        except ValueError:
-            return False
-        return True
-
-
-class BaseEnum(enum.Enum, metaclass=EnumWithContains):
-    "An enum class that can get the value of an item with `str(Enum.key)`"
-
-    def __str__(self):
-        return self.value
-
-    @classmethod
-    def list(cls):
-        "Method to list all the possible items in `cls`"
-        return list(map(str, cls))
+    TVM = "TVM"
 
 
 class LoggerType(BaseEnum):
diff --git a/src/accelerate/utils/launch.py b/src/accelerate/utils/launch.py
--- a/src/accelerate/utils/launch.py
+++ b/src/accelerate/utils/launch.py
@@ -21,7 +21,6 @@
 import torch
 
 from ..commands.config.config_args import SageMakerConfig
-from ..commands.config.config_utils import DYNAMO_BACKENDS
 from ..utils import (
     DynamoBackend,
     PrecisionType,
@@ -89,7 +88,9 @@ def prepare_simple_launcher_cmd_env(args: argparse.Namespace) -> Tuple[List[str]
     try:
         dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
     except ValueError:
-        raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DYNAMO_BACKENDS}.")
+        raise ValueError(
+            f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DynamoBackend.list()}."
+        )
     current_env["ACCELERATE_DYNAMO_BACKEND"] = dynamo_backend.value
     current_env["ACCELERATE_DYNAMO_MODE"] = args.dynamo_mode
     current_env["ACCELERATE_DYNAMO_USE_FULLGRAPH"] = str(args.dynamo_use_fullgraph)
@@ -163,7 +164,9 @@ def prepare_multi_gpu_env(args: argparse.Namespace) -> Dict[str, str]:
     try:
         dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
     except ValueError:
-        raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DYNAMO_BACKENDS}.")
+        raise ValueError(
+            f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DynamoBackend.list()}."
+        )
     current_env["ACCELERATE_DYNAMO_BACKEND"] = dynamo_backend.value
     current_env["ACCELERATE_DYNAMO_MODE"] = args.dynamo_mode
     current_env["ACCELERATE_DYNAMO_USE_FULLGRAPH"] = str(args.dynamo_use_fullgraph)
@@ -419,7 +422,9 @@ def prepare_sagemager_args_inputs(
     try:
         dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
     except ValueError:
-        raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DYNAMO_BACKENDS}.")
+        raise ValueError(
+            f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DynamoBackend.list()}."
+        )
 
     # Environment variables to be set for use during training job
     environment = {


Test patch:
diff --git a/tests/test_kwargs_handlers.py b/tests/test_kwargs_handlers.py
--- a/tests/test_kwargs_handlers.py
+++ b/tests/test_kwargs_handlers.py
@@ -92,11 +92,11 @@ def test_torch_dynamo_plugin(self):
             prefix = "ACCELERATE_DYNAMO_"
             # nvfuser's dynamo backend name is "nvprims_nvfuser"
             # use "nvfuser" here to cause exception if this test causes os.environ changed permanently
-            os.environ[prefix + "BACKEND"] = "nvfuser"
+            os.environ[prefix + "BACKEND"] = "aot_ts_nvfuser"
             os.environ[prefix + "MODE"] = "reduce-overhead"
 
             dynamo_plugin_kwargs = TorchDynamoPlugin().to_kwargs()
-            self.assertEqual(dynamo_plugin_kwargs, {"backend": "nvfuser", "mode": "reduce-overhead"})
+            self.assertEqual(dynamo_plugin_kwargs, {"backend": "aot_ts_nvfuser", "mode": "reduce-overhead"})
 
 
 if __name__ == "__main__":