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>
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user