Torch compile backend defaults to "neuron" (#47035)

* feat: torch compile backend defaults to "neuron"

* fix: restore doc back to normal

* doc: add information about backend in docs

Author: michaelbenayoun

docs/source/en/torch_compile.md

@@ -32,7 +32,9 @@ args = TrainingArguments(
 
 ## Backend
 
-The default backend is `inductor`, which compiles to Triton kernels with AOTAutograd. This is the right choice for most training workloads. Use `cudagraphs` for fixed-shape inputs, or `ipex` for Intel CPU training.
+When no backend is specified, [`TrainingArguments`] selects one based on your hardware. On most CPUs and GPUs, the default is `inductor`, which compiles to Triton kernels with AOTAutograd and suits most training workloads. On Intel Gaudi (HPU), the default is `hpu_backend`. On AWS Trainium and Inferentia (Neuron), the default is `neuron`.
+
+Use `cudagraphs` for fixed-shape inputs, or `ipex` for Intel CPU training.
 
 ## Compile mode
 

src/transformers/training_args.py

@@ -1565,6 +1565,8 @@ def __post_init__(self):
         if self.torch_compile and self.torch_compile_backend is None:
             if not self.use_cpu and is_torch_hpu_available():
                 self.torch_compile_backend = "hpu_backend"
+            elif not self.use_cpu and is_torch_neuron_available():
+                self.torch_compile_backend = "neuron"
             else:
                 self.torch_compile_backend = "inductor"