

import torch
import torch_npu
!! Further additional imports can be added here !!
Specify the device name (eg. 'cuda', 'cpu', 'npu')
DEVICE_NAME = 'npu'
Specify device-specific backends to dispatch to.
If not specified, will fallback to 'default' in 'testing_utils.py`
MANUAL_SEED_FN = torch.npu.manual_seed
EMPTY_CACHE_FN = torch.npu.empty_cache
DEVICE_COUNT_FN = torch.npu.device_count
``
This format also allows for specification of any additional imports required. To use this file to replace equivalent methods in the test suite, set the environment variableTRANSFORMERS_TEST_DEVICE_SPEC` to the path of the spec file.
Currently, only MANUAL_SEED_FN, EMPTY_CACHE_FN and DEVICE_COUNT_FN are supported for device-specific dispatch.
Distributed training
pytest can't deal with distributed training directly. If this is attempted - the sub-processes don't do the right
thing and end up thinking they are pytest and start running the test suite in loops. It works, however, if one
spawns a normal process that then spawns off multiple workers and manages the IO pipes.
Here are some tests that use it: