Make 🤗Transformers tests device agnostic

[vvvm23]

Feature request

We would like to make the testing suite in this repository more device agnostic. It seems there has already been some work towards this, however the majority of tests will still only run on either GPU or CPU. This would require a number of changes to all tests present in the library, however it would not alter the behaviour of Huggingface's CI runners.

A non-exhaustive list of changes would be:

• Add a new test decorator @require_torch_with_accelerator that largerly supersedes (but does not replace) @require_torch_gpu. This new decorator can be used for any test that is device agnostic that we would like to accelerate. We would keep @require_torch_gpu for tests that truly require CUDA features, such as ones that check device memory utilisation (such as in model parallelism or lower precision tests) or use custom CUDA kernels (such as Flash Attention).
• Certain tests could be made device agnostic quite easily, such as tests that only check for CUDA devices to enable fp16, tests that use backend specific PRNG initialisation, or tests that clear cache before executing. This could be done by adding device agnostic variants to testing_utils.py that compare the current device in use and dispatch to the appropriate backend specific function if available.
  • For example, rather than the comparison torch_device == 'cuda' to check if we can run with fp16, we could call a function testing_utils.accelerator_is_fp16_available(torch_device) or similar. Similar functions already exist to check for tf32 or bf16 support.
  • Crucially, in upstream we would only have settings for CUDA and CPU devices – as well as any other of your supported backends. However, we would expose functions to register your own device in user code so third parties can test custom backends without upstreaming changes.

Motivation

As Huggingface libraries and models make up a significant part of the current ML community, it makes sense when developing custom PyTorch backends to test against these model libraries as they cover a large proportion of the most users' use cases.

However, the current testing suite does not easily allow for custom devices – not without maintaining a custom private fork that needs to be continuously kept up to date with the upstream repository. This reason, and because the number of changes required is not especially significant, is why we are making this proposal.

Your contribution

We would write and submit a PR to implement these changes following discussion and approval with 🤗Transformers maintainers.

I am collaborating with @joshlk and @arsalanu

[vvvm23]

See #25655 #25506 huggingface/diffusers#4673 for further context and existing changes to help make testing device agnostic.

[ArthurZucker]

cc @ydshieh

[ydshieh]

Hi @vvvm23 Thank you for this proposal! This sound an impactful thing to do 🚀 !

For a draft PR, it would be very nice to keep the change as minimal as possible, especially not to change all the tests but just a few tests so we can see how things work and how them are applied to those tests. 🙏

[vvvm23]

That sounds reasonable to start with. I'll pick a few tests that cover all the proposed features (certain models are simpler than others, and won't highlight all the changes required). I'll try and get a draft PR by end of this week.

[vvvm23]

Hi @ydshieh, can I get your thoughts on a couple design choices?

First, certain device agnostic checks are as simple as simply trying to use the specified feature. For example, to test whether a specific device can use fp16, try and execute an op in half precision, and catch the exception. Does this approach work for you? A similar thing is already done to check if the device exists (attempt to create device, catch the exception) and could work here too.

For more complex device agnostic functions (such as clearing cache, setting PRNG) we will need a way to define new devices without having to upstream the device itself. I was thinking of going for the following approach:

1. Introduce a new environment variable TRANSFORMERS_TEST_DEVICE_SPEC which points a Python file.
2. Within the file, have an importable dictionary mapping function key names to actual backend callables.
3. If such a variable exists, testing_utils will add the mappings to its own internal register of devices to backend functions, which will be used by the device agnostic functions.
   • this will only be cpu and gpu by default.

If a specific entry does not exist in TRANSFORMERS_TEST_DEVICE_SPEC we can use some sane default functions, or a no-op, or simply throw an error when attempting to use the device agnostic function associated with the missing entry.

How does that approach sound to you? Any suggestions that would help make the PR fit into the current HF testing structure? Thanks~

[ydshieh]

Hi @vvvm23 Before I take a deeper read, could you provide the comment along with some links to the code base, so I can understand easier.

For example:

• a link to the lines you mentioned A similar thing is already done to check if the device exists
• a link to a test that `test whether a specific device can use fp16
• a link to a (or some) more complex device agnostic functions (such as clearing cache, setting PRNG)
  • and elaborate a bit more about we will need a way to define new devices without having to upstream the device itself

Although I am the one within the team focus on the testing, many of the tests are written before I joined. So a bit more detailed description would be easier for me to give my thoughts 🙏 please (I know it will takes you some more time to write).

Thank you in advance!

[vvvm23]

No worries, I wrote my previous comment in a bit of a rush, so in retrospect it wasn't too clear.

a link to the lines you mentioned A similar thing is already done to check if the device exists

Please see these PRs which make this change -> #25506 huggingface/diffusers#4673

a link to a test that `test whether a specific device can use fp16

For example, in test_modeling_opt.py

transformers/tests/models/opt/test_modeling_opt.py

Line 294 in 35c570c

if torch_device == "cuda":

this test will only use half precision if CUDA is the target device, even if there is half precision capabilities on torch_device.

Another example is here

transformers/tests/models/reformer/test_modeling_reformer.py

Line 571 in 35c570c

@unittest.skipIf(torch_device == "cpu", "Cant do half precision")

which would still run on a device that can't do half precision.

a link to a (or some) more complex device agnostic functions (such as clearing cache, setting PRNG)

One example can be found here

transformers/tests/models/codegen/test_modeling_codegen.py

Line 501 in 35c570c

if torch_device == "cuda":

where we set the seed for the CUDA device, but in our desired device-agnostic context, we would rather have a generic "accelerate set seed" function.

and elaborate a bit more about we will need a way to define new devices without having to upstream the device itself

Suppose we implement the function accelerator_manual_seed. When called, this will look up the current test device in use and dispatch to the correct function (if torch_device == cuda dispatch to torch.cuda.manual_seed.

However, if we try a custom device, the function won't know which function to dispatch to. This isn't the responsibility of Huggingface to solve (as there could be countless custom devices) but rather the user to register their device and the function to use when we call accelerator_manual_seed. Currently I feel the best way to do this is to specify an environment variable that points to a python file that defines the backend functions. This way, a user can test their new device on the library without upstreaming changes to Huggingface.

I have begun work on this here but it isn't very robust yet.

Let me know if you need me to clear anything else up~

[ydshieh]

Hi, Thanks for the writing-up!

When I read the fp16 part, I have a feeling that why not just pre-determine which device can (or can't) use fp16 explicitly in src/transformers/testing_utils.py, just like what you mentioned in the description testing_utils.accelerator_is_fp16_available(torch_device).

Note the design of accelerator_is_fp16_available should not try to create/use the device in each call, rather we should cache the result and reuse it. I was originally thinking a if/elif statements for such method, but you said here could be countless custom devices afterward ...

Regarding accelerator_manual_seed, I am not sure we really will have so many different devices. If this is the case, your proposal of having a external file and let user to specify makes sense.

But maybe we can start the task easily and not using external file. Just put everything inside testing_utils. WDYT?

[vvvm23]

Note the design of accelerator_is_fp16_available should not try to create/use the device in each call, rather we should cache the result and reuse it.

This is a fair point. I noticed in some other functions that check for device availability, they are wrapped in @lru_cache which should match this behaviour. I will add this to the function calls.

I was originally thinking a if/elif statements for such method, but you said here could be countless custom devices afterward

To clarify, I am only intending one device to be used for a single set of tests, so there won't be countless devices in use in a single session.

But maybe we can start the task easily and not using external file. Just put everything inside testing_utils. WDYT?

For our purposes (and I am assuming others too) the main point is to be able to specifiy a new backend or device without having to upstream anything into HF. So if we put everything inside testing_utils this defeats the point, even if this is only to start. So, I think for the initial PR we should support this. Let me know what you think~

[ydshieh]

For our purposes (and I am assuming others too) the main point is to be able to specifiy a new backend or device without having to upstream anything into HF. So if we put everything inside testing_utils this defeats the point, even if this is only to start. So, I think for the initial PR we should support this. Let me know what you think~

Yeah, I agree. I would say something defined in HF (only for cuda and cpu), and for 3rd party, we use the definitions from external file(s) as you suggested. The main point of this is cuda and cpu is known to everyone and let's not put the definitions related to them outside transformers, so it's easier to find the info.

[vvvm23]

Yep! That was my original plan, we don't want to put any burden on HF into maintaining additional devices that are only known to a small set of people 👍 We will have the definitions for cpu and cuda within transformers but have support for external files.

[vvvm23]

@ydshieh please see the above draft PR 🙂

I am hoping the CI passes without issue as there should be no effect on Huggingface CI runners with these changes.