Device agnostic testing

[arsalanu]

What does this PR do?

Adds new features to testing_utils.py and import_utils.py to make testing with non-default PyTorch backends (beyond just cuda, cpu and mps) possible. This should not affect any current testing within the repo or the behaviour of the devices they are run on.

This is heavily based on similar work we have done for Transformers, see: Transformers PR #25870

Adds some device agnostic functions which dispatch to specific backend functions. This is mainly applicable to functions which are device-specific (e.g. torch.cuda.manual_seed). Users can specify new backends and backends for device agnostic functions by creating a device specification file and pointing the test suite to it using the environment variable DIFFUSERS_TEST_DEVICE_SPEC, and add a new device for PyTorch using DIFFUSERS_TEST_DEVICE.

Example of a device specification to run the tests with an alternative accelerator:

import torch
import torch_npu
# User can add additional imports here

# Specify the device name (eg. 'cuda', 'cpu')
DEVICE_NAME = 'npu'

# Specify device-specific backends to dispatch to.
# If not specified (i.e., `None`) will fallback to 'default' in 'testing_utils.py`
MANUAL_SEED_FN = torch.npu.manual_seed
EMPTY_CACHE_FN = None
DEVICE_COUNT_FN = torch.npu.device_count
SUPPORTS_TRAINING=True

Implementation details are fully outlined in the issue #5562

I have a modified a single file (UNet2D condition model tests, and test_modeling_common as this is used in the UNet2D tests) rather than all the tests, as this PR is more focused on the implementation of features required for device-agnostic testing.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case. Issue Make 🤗 Diffusers tests device-agnostic #5562
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

[yiyixuxu]

@patrickvonplaten
can you take a look here?

[patrickvonplaten]

The changes look very reasonable to me! Thanks a lot for making everything device-agnostic.

Just a bit worried about the is_torch_fp16_available because we're essestially just saying if the matmul doesn't work fp16 is not available, but the matmul might also not work because of other reasons (badly installed CUDA, OOM, ...)

In PyTorch there is actually a is_bf16_available(): https://github.com/pytorch/pytorch/blob/d64bc8f0f81bd9b514eb1a5ee6f5b03094e4e6e9/torch/cuda/__init__.py#L141

The function seems to check some device properties which is probably less brittle - guess it's hard to do this for fp16 here, but can we maybe make sure that we don't accidentally misinterpret other erros as fp16 not being available?

[arsalanu]

Thanks, @patrickvonplaten. I agree that just catching any exception might not be the best way to do this, but I'm not sure if there is a specific exception that would be agnostic to any accelerator or device. On CPU and XLA I believe you get a RuntimeError when trying to perform operations with FP16, but this is vague in itself and may overlap with a different issue.

One suggestion is to write it so the exception is logged and print the error when tests are run, to specify why FP16 is not working? This would make it clear to the user whether it is unsupported behaviour or an issue with their setup.

Looking at the PyTorch is_bf16_available() this specifies different checks for different hardware, which makes it less brittle as you said, but that function is also not device agnostic and would only work for ROCm and CUDA backends.

Another suggestion I have is to add a CUDA specific check and skip the FP16 matmul check in is_torch_fp16_available() if a GPU is being used (I could do this for MPS and CPU too, which would be set to False by default).

This would still be in line with the changes as these backends have defaults specified for them in the custom function dispatch as well. Then the matmul check would only be used if a non-default device is being used. We could do this and also log the error to make it explicit to the user.

Let me know if that makes sense to you and I will add those changes, or any other suggestions you have. Thanks!

[patrickvonplaten]

Thanks, @patrickvonplaten. I agree that just catching any exception might not be the best way to do this, but I'm not sure if there is a specific exception that would be agnostic to any accelerator or device. On CPU and XLA I believe you get a RuntimeError when trying to perform operations with FP16, but this is vague in itself and may overlap with a different issue.

One suggestion is to write it so the exception is logged and print the error when tests are run, to specify why FP16 is not working? This would make it clear to the user whether it is unsupported behaviour or an issue with their setup.

Looking at the PyTorch is_bf16_available() this specifies different checks for different hardware, which makes it less brittle as you said, but that function is also not device agnostic and would only work for ROCm and CUDA backends.

Another suggestion I have is to add a CUDA specific check and skip the FP16 matmul check in is_torch_fp16_available() if a GPU is being used (I could do this for MPS and CPU too, which would be set to False by default).

This would still be in line with the changes as these backends have defaults specified for them in the custom function dispatch as well. Then the matmul check would only be used if a non-default device is being used. We could do this and also log the error to make it explicit to the user.

Let me know if that makes sense to you and I will add those changes, or any other suggestions you have. Thanks!

Could we maybe do something like this: https://github.com/huggingface/diffusers/pull/5612/files#r1399038284 just to add an extra safety-mechanism that a sure doesn't understand the function incorrectly in case cuda is badly set up?

Also can we make the function private for now, e.g. add an underscore so that it's _is_torch_fp16_available()?

[arsalanu]

I've added the changes, slightly restructured so the FP16 op-check happens by default for all accelerators and the CUDA error is raised only if the device type is cuda.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[patrickvonplaten]

Cool thanks for iterating here. This PR LGTM - should we merge it now or do you want to add tests for other classes directly here?

[DN6]

Is this meant to run when torch isn't available or if DIFFUSERS_TEST_DEVICE_SPEC is set?

[arsalanu]

The guard is there because the function dispatch should only run if torch is available, it doesn't strictly matter if DIFFUSERS_TEST_DEVICE_SPEC is set. For example for a GPU, CPU or MPS device, a spec doesn't need to be set but torch must still be available to dispatch to the default torch device functions.

[arsalanu]

should we merge it now or do you want to add tests for other classes directly here?

If thats okay I'll add some more before merging 😄 I had a few other tests ready but removed them for this PR to keep it minimal.

[arsalanu]

I've added more test coverage. The latest commit has the changes for most of the model classes (unet, vae, vq, unet2d and some common files) and one pipeline test (SD2). Any more tests could be added in future PRs.

[patrickvonplaten]

Cool! The changes look good to me - @DN6 wdyt? Feel free to merge once you're happy with it

[DN6]

LGTM 👍🏽 Nice work @arsalanu!