Allow creation of pseudo devices for testing purposes

[avitase]

🚀 Feature

Add a pseudo device cpu_testing for testing purposes.

Motivation

For machines with only a CPU there is currently no straight-forward way to test if torch.device is set properly for all tensors since cpu is the default. Having a dummy device, e.g., named cpu_testing which internally falls back to the default device but throws if two tensors with different devices are combined would allow the creation of unit tests.

Pitch

It happens that some machines purely used for CI don't have any fancy devices other than CPU. Since CPU is the default device errors frequently occur in production where the model missed to set tensors to, e.g., GPU and this case cannot be covered in a unit test. Creating a pseudo device which also works for machines with only CPUs could help to change this defect.

cc @ezyang @bhosmer @smessmer @ljk53 @bdhirsh

[bdhirsh]

It sounds like the meta device might fit your use case? With the snag being that it doesn't work for all operators yet (although many of the most common ops have been ported already).

>>> a = torch.ones(2) # defaults to cpu
>>> b = torch.ones(2, device='meta')
>>> c = a + b
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, meta and cpu!

[avitase]

That's neat. Wasn't aware of this. Thanks for your help!

@bdhirsh could you explain what the purpose of the meta device was originally?

[bdhirsh]

the meta API lets you run dtype and shape inference. For example:

>>> a = torch.ones(2, device='meta', dtype=torch.float64)
>>> b = torch.ones((1, 2), device='meta', dtype=torch.float32)
>>> a + b
tensor(..., device='meta', size=(1, 2), dtype=torch.float64)

The above involved type promotion (to promote the output type to float64), and broadcasting (to broadcast the output shape to (1,2), but didn't allocate any storage and didn't actually "add" anything.

Tensors on the meta device don't have any storage, and calling operators on meta tensors doesn't involve any compute - it only involves input error checking + running the "meta" computation.

The original RFC is here if you're interested: https://github.com/pytorch/rfcs/blob/rfc-0005/RFC-0005-structured-kernel-definitions.md#how-to-get-involved-with-structured-kernels. "Structured Kernels" are a new way of authoring kernels internally, that require you separate your "meta" computation from your core kernel logic. Which is what allows us to create a meta API.

[avitase]

Thanks @bdhirsh, this looks indeed promising but being able to test the device and the value at the same time could also be handy. I therefore leave this feature request open but I will definitely give meta a try!

[ezyang]

I think there are some slight semantics differences that would make meta more useful for your use case here, so definitely OK to keep this open.

[ikamensh]

For people waiting for this feature and using MacOS, you can use "mps" device (metal performance shaders). It turns out to be available even on macs without M1+ hardware, e.g. I have intel macbook with AMD gpu.

[njzjz]

For people waiting for this feature and using MacOS, you can use "mps" device (metal performance shaders). It turns out to be available even on macs without M1+ hardware, e.g. I have intel macbook with AMD gpu.

When I tried it, I found mps doesn't support float64, though.

[chanind]

This would be very helpful, since tests running in CI (e.g. Github Actions) usually don't have a GPU and this allows bugs related to mismatched tensors to slip through into production. The meta device doesn't work for anything involving actual calculation, which any integration test will necessarily run, and mps is only available on macs which are typically not what you're using for CI.

[ezyang]

It might be relatively simple to do this with privateuse1. The general idea would be to look at the sample code for how to add a new custom device type that tests privateuse1 right now, but then actually just implement a single fallback kernel for everything that just redispatches to cpu.

[ezyang]

cc @albanD

[albanD]

Ho yes we can definitely do that.
You can actually see the current testing we have for PrivateUse1 is doing part of this:

pytorch/test/test_cpp_extensions_open_device_registration.py

Line 30 in 8f3fd79

class DummyModule:

I would agree with Ed that we can make this a fully fledged extension that has fallback to CPU implementation if we need a version you can run actual compute tests against.