From 5eeaef921f70acd68073d1066ccb09d7c6e6f475 Mon Sep 17 00:00:00 2001
From: Alex McKinney <44398246+vvvm23@users.noreply.github.com>
Date: Tue, 22 Aug 2023 16:08:13 +0100
Subject: [PATCH] Adds `TRANSFORMERS_TEST_BACKEND` (#25655)

* Adds `TRANSFORMERS_TEST_BACKEND`
Allows specifying arbitrary additional import following first `import torch`.
This is useful for some custom backends, that will require additional imports to trigger backend registration with upstream torch.
See https://github.com/pytorch/benchmark/pull/1805 for a similar change in `torchbench`.

* Update src/transformers/testing_utils.py

Co-authored-by: Yih-Dar <2521628+ydshieh@users.noreply.github.com>

* Adds real backend example to documentation

---------

Co-authored-by: Yih-Dar <2521628+ydshieh@users.noreply.github.com>
---
 docs/source/en/testing.md         |  9 +++++++--
 src/transformers/testing_utils.py | 12 ++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/docs/source/en/testing.md b/docs/source/en/testing.md
index 0179d5c635..ec11e39be0 100644
--- a/docs/source/en/testing.md
+++ b/docs/source/en/testing.md
@@ -511,15 +511,20 @@ from transformers.testing_utils import get_gpu_count
 n_gpu = get_gpu_count()  # works with torch and tf
 ```
 
-### Testing with a specific PyTorch backend
+### Testing with a specific PyTorch backend or device
 
-To run the test suite on a specific torch backend add `TRANSFORMERS_TEST_DEVICE="$device"` where `$device` is the target backend. For example, to test on CPU only:
+To run the test suite on a specific torch device add `TRANSFORMERS_TEST_DEVICE="$device"` where `$device` is the target backend. For example, to test on CPU only:
 ```bash
 TRANSFORMERS_TEST_DEVICE="cpu" pytest tests/test_logging.py
 ```
 
 This variable is useful for testing custom or less common PyTorch backends such as `mps`. It can also be used to achieve the same effect as `CUDA_VISIBLE_DEVICES` by targeting specific GPUs or testing in CPU-only mode.
 
+Certain devices will require an additional import after importing `torch` for the first time. This can be specified using the environment variable `TRANSFORMERS_TEST_BACKEND`:
+```bash
+TRANSFORMERS_TEST_BACKEND="torch_npu" pytest tests/test_logging.py
+```
+
 
 ### Distributed training
 
diff --git a/src/transformers/testing_utils.py b/src/transformers/testing_utils.py
index 0f70ac9777..c8c6665779 100644
--- a/src/transformers/testing_utils.py
+++ b/src/transformers/testing_utils.py
@@ -16,6 +16,7 @@ import collections
 import contextlib
 import doctest
 import functools
+import importlib
 import inspect
 import logging
 import multiprocessing
@@ -642,6 +643,17 @@ if is_torch_available():
         torch_device = "npu"
     else:
         torch_device = "cpu"
+
+    if "TRANSFORMERS_TEST_BACKEND" in os.environ:
+        backend = os.environ["TRANSFORMERS_TEST_BACKEND"]
+        try:
+            _ = importlib.import_module(backend)
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError(
+                f"Failed to import `TRANSFORMERS_TEST_BACKEND` '{backend}'! This should be the name of an installed module. The original error (look up to see its"
+                f" traceback):\n{e}"
+            ) from e
+
 else:
     torch_device = None