SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Adds TRANSFORMERS_TEST_BACKEND (#25655)

Browse Source 
* 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>
This commit is contained in:
Alex McKinney
2023-08-22 16:08:13 +01:00
committed by GitHub
parent fd56f7f081
commit 5eeaef921f
2 changed files with 19 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
docs/source/en/testing.md
View File

@@ -511,15 +511,20 @@ from transformers.testing_utils import get_gpu_count
n_gpu = get_gpu_count() # works with torch and tf 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 ```bash
TRANSFORMERS_TEST_DEVICE="cpu" pytest tests/test_logging.py 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. 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 ### Distributed training

12
src/transformers/testing_utils.py
View File

@@ -16,6 +16,7 @@ import collections
import contextlib import contextlib
import doctest import doctest
import functools import functools
import importlib
import inspect import inspect
import logging import logging
import multiprocessing import multiprocessing
@@ -642,6 +643,17 @@ if is_torch_available():
torch_device = "npu" torch_device = "npu"
else: else:
torch_device = "cpu" 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: else:
torch_device = None torch_device = None