torch.mps.*Tensor datatypes

[TV4Fun]

🚀 The feature, motivation and pitch

An issue that has been debated ad nauseam and apparently still doesn't have an agreed upon answer as of PyTorch 1.12 is how or if to set a default device for Torch operations. See for example #27878, which has been open nearly 3 years now. A method many libraries use and which is recommended in at least a few places on Stack Overflow is to use set_default_tensor_type to default to a CUDA tensor type. This works for CUDA, but does not work for MPS, as there do not appear to be equivalent tensor types for that. This also creates problems with unclear type names, as if, for example, I create a tensor on MPS with w = torch.tensor([1.0], device='mps'), w.type() for this returns 'torch.mps.FloatTensor', but this is not actually a valid type. There is no torch.mps module, and if I try to pass it as a string, say with x = w.type('torch.mps.FloatTensor'), this returns an error that it is an invalid type.

As near as I can tell, there is no way to directly create a tensor on the MPS device without specifying device='mps' on every single call, which not only clutters code but also makes it very brittle if I happen to miss it on one call. Please correct me if I am wrong on this. My particular use case is that I would like to add MPS support to ML-Agents, which already supports CUDA by means of calling torch.set_default_tensor_type(torch.cuda.FloatTensor) if a CUDA device is available. There does not appear to be an equivalent way to do this for MPS and I have no desire to try and track down however many hundreds of tensor creation calls there are in their code. I know there have been calls to deprecate set_default_tensor_type, i.e. #53124, but I would recommend not doing that without providing some other way to provide a default device. In the meantime, I would really like it if I had an easy way to set the default tensor type to torch.mps.FloatTensor.

Alternatives

Provide a way to set the default device for torch.tensor() and similar calls to MPS.

Additional context

See also #260 and probably others.

cc @ezyang @gchanan @zou3519 @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr @abhudev

[albanD]

cc @ezyang

[ezyang]

With TorchFunctionMode we can do this in a few dozen lines of code. I'll put up a PoC later today

[soumith]

@ezyang
@pbelevich and I were discussing this in a more general context.

With your PoC, can we for example make the default device be a meta device?

[ezyang]

yes; it will basically be similar to how torchdistx does it

[ezyang]

This is totally untested

import torch
from torch.overrides import TorchFunctionMode

_DEVICE_CONSTRUCTOR = {
    # standard ones
    torch.empty,
    torch.empty_strided,
    torch.empty_quantized,
    torch.ones,
    torch.arange,
    torch.bartlett_window,
    torch.blackman_window,
    torch.eye,
    torch.fft.fftfreq,
    torch.fft.rfftfreq,
    torch.full,
    torch.fill,
    torch.hamming_window,
    torch.hann_window,
    torch.kaiser_window,
    torch.linspace,
    torch.logspace,
    torch.nested_tensor,
    # torch.normal,
    torch.ones,
    torch.rand,
    torch.randn,
    torch.randint,
    torch.randperm,
    torch.range,
    torch.sparse_coo_tensor,
    torch.sparse_compressed_tensor,
    torch.sparse_csr_tensor,
    torch.sparse_csc_tensor,
    torch.sparse_bsr_tensor,
    torch.sparse_bsc_tensor,
    torch.tril_indices,
    torch.triu_indices,
    torch.vander,
    torch.zeros,
    torch.asarray,
    # weird ones
    torch.tensor,
    torch.as_tensor,
}

class DeviceMode(TorchFunctionMode):
    def __init__(self, device):
        self.device = torch.device(device)

    def __torch_function__(self, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        if func in _DEVICE_CONSTRUCTOR:
            if kwargs.get('device') is None:
                kwargs['device'] = self.device
            return func(*args, **kwargs)
        return func(*args, **kwargs)


with DeviceMode(torch.device("meta")):
    print(torch.empty(3))

I'd like to put this into core, but before we can get there, we have to make some policy decisions. For example, if you turn on DeviceMode("mps"), and then someone writes torch.randn(2, device="cpu"), does mps override the CPU device, or does the explicit device win out? It would be nice for a limited set of power users to try this out, make behavior modifications as necessary based on what they observe is necessary in the wild, and then we ship that.

[TV4Fun]

I would definitely suggest an explicit device overriding the default device mode. Though if you think this could cause problems, you could always raise a warning the first time this happens.

[TV4Fun]

Is this in the development branch yet? I'd like to be able to try it out.

[ezyang]

it's definitely in the nightly, and it may also work on the most recent official release

[TV4Fun]

Thank you @ezyang. Due to #78681 and #79337, I am unable to test this on Torch nightly with my M1 mac. Trying this on Torch 1.12.0 gives this error:

TypeError                                 Traceback (most recent call last)
Input In [1], in <cell line: 61>()
     57             return func(*args, **kwargs)
     58         return func(*args, **kwargs)
---> 61 with DeviceMode(torch.device("meta")):
     62     print(torch.empty(3))

File ~/miniforge3/envs/torch-nightly/lib/python3.10/site-packages/torch/utils/_mode_utils.py:28, in _wrap_init.<locals>.wrapped(self, inner, *args, **kwargs)
     25 @functools.wraps(f)
     26 def wrapped(self, *args, inner=undef, **kwargs):
     27     if inner is undef:
---> 28         raise TypeError(
     29             f"missing inner keyword argument; instead of constructing a {meta_init_error_info.mode_class_name} "
     30             f"directly, pass the constructor to push_{meta_init_error_info.mode_name}_mode"
     31         )
     32     self.inner = inner
     33     return f(self, *args, **kwargs)

TypeError: missing inner keyword argument; instead of constructing a TorchDispatchMode directly, pass the constructor to push_torch_dispatch_mode

[ezyang]

Ok you are missing a bugfix that makes the syntax work, try writing the context manager as MyMode.push(device) instead

[TV4Fun]

Okay, thank you, that works and if I create this with

with DeviceMode.push(torch.device("mps")):
    print(torch.empty(3))

the created tensor appears to be on the MPS device. Is there a way to do this without having to move all my code inside of a with statement though? It seems to do nothing if I call DeviceMode.push(torch.device("mps")) or even manually DeviceMode.push(torch.device("mps")).__enter__(). I understand the idea of using a context manager here, but in practice, that still means a lot of work porting code that just assumes you can set a default.

[ezyang]

g = DeviceMode.push(torch.device("mps"))
g.__enter__()