[FR] torch context with default device & dtype

[ssnl]

🚀 Feature

There has been needs to set default device. This FR proposes an API that is compatible with 3rd party libs.

Motivation

There has been a lot of discussion around default device flag in pytorch (#7535). Yet, implementing such an API has been mostly blocked by the concern that 3rd party libraries may assume that tensors are created on CPU.

Similar dilemma has been also seen in the python multiprocessing library, where multiple start methods can be used (triggering bugs like librosa/librosa#747). They come up with this API multiprocessing.get_context(xxx), which returns a context object with the same set of functions as the multiprocessing module, but associated with a different start method, enabling patterns like mp = mp.get_context('spawn').

In addition to tensor creation, for many users (including me) it is extremely verbose to write .to(device) for every tensor yielded tensor from the data loader. It would be very handy if this handles the moving automatically as well.

Pitch

1. Tensor creation

   import torch
   torch.empty(3)  # returns CPU float tensor
   torch = torch.get_context(device='cuda:2', dtype=torch.half)
   torch.empty(3)  # returns CUDA half tensor

2. Module creation

   import torch
   torch = torch.get_context(device='cuda:2')
   MyModule()  # returns module on CUDA if MyModule uses torch tensor creation methods

3. (needs discussion) affect torch_ctx.utils.data.DataLoader s.t. the yielded samples contain tensors moved to the device of torch_ctx.
4. (needs discussion) affect torch_ctx.load(xxx) the same way as map_location.

Concerns

Such features always have the potential issue of making things too "frictionless", obscure, and harder to debug. But I think this one is not too bad.

cc @ezyang @gchanan @zou3519 @bdhirsh @heitorschueroff @ngimel @ejguan

[soumith]

after one framework finally got rid of a global context after 3 years, are you saying the competing framework should add one? :P

[ssnl]

:D a different kind of context

[gchanan]

this sounds like a promising API -- there's probably some implementation complexity behind it, but we'll remove "needs research" as it looks acceptable.

[elistevens]

Unclear if this issue is still actually active, but it was pointed at recently as if it were, and I'd like to point out that #2 above would require some magic to make work properly (assuming that the torch = part is required, vs. just setting a bunch of global state in torch.get_context(device='cuda:2'), and in that case, the function should be renamed to something more accurate), and would leave a bunch of polluted global state behind, possibly just from importing a 3rd party package that used it.

Please use a context manager like with torch.set_default_device('cuda:0'):.

[ssnl]

@elistevens Yeah I think you are right. #2 doesn't work normally, unfortunately.

[ezyang]

Bumping priority due to activity on this issue, the parent issue, and also other issues.

[malfet]

Wouldn't context manager be better suited that a global context?
I.e.

with torch.context(device='cuda:2', dtype=torch.double):
   x = torch.ones((5, 5))

[ssnl]

@malfet Why not having the object do both? It is not always desirable to wrap the entire program in a context manager.

[ailzhang]

Removing "high priority" label since this feature request still need some work in defining the user facing API and how it should work with existing code.
Probably also related: we also need define what to do with the current torch.set_default_dtype API.

[ezyang]

Context manager doesn't give you desirable behavior, as it will also affect library code that allocates tensors. A lot of the resistance against a global "set default device" API is because it will make it difficult for library authors to write code in a way that will work no matter what the default device is. A module-like torch context object bypasses this problem as the device defaulting is lexical.

This reminds me, though, that in the proposal above, Module creation isn't done using the torch context explicitly; some amount of dynamic scoping seems necessary there. So this proposal, unfortunately, isn't complete.

[ezyang]

Oh, actually, module creation is done doing torch context, because you say torch.nn.Module, so if you have a torch context, that would also solve your NN problem to. (But you'll have to shave the yak of how to desugar these calls into Python module calls first.)

[ssnl]

@ezyang Yeah... that could work. But in the usual case where one has a separate module containing the module definition, it seems hard to make the context device/dtype configurable (because that separate module just uses the regular torch.nn), right?

[ezyang]

You have to solve the problem which is that we don't actually support direct on device Module creation. Supposing you add some explicit API for doing this, then you can just write some Python magic to sniff for nn modules and then partially apply them with the device arguments. You still avoid dynamic scoping in this case.

[ssnl]

Hmm I don't think I understand. What I am thinking is a scenario like the following:

# model.py

import torch

class Model(torch.nn.Module):
  ...

# main.py

torch_ctx = torch.get_context(sys.argv[1])

from model import Model

m = Model()  # <- not using `torch_ctx`

Oh maybe your this comment

You have to solve the problem which is that we don't actually support direct on device Module creation

is referring to something like Model(..., torch_ctx=torch_ctx) or torch_ctx.create(Model, ...), which would then make sense, although I can't think of an API that is nice for this...

[rgommers]

Python modules are singletons, so it's a little unclear what model = torch_ctx.wrap(model) should be doing. The wrap() can't just tack some context onto the existing model (that would leak to other users of model), therefore model needs to be a new namespace-like object with Model also being a new generated version of the original Model.

I guess that's the yak that needed shaving, and the same applies to torch = torch.get_context(...). Having an idea of how it could be implemented would be nice - it looks to me like you need a factory function that produces new namespaces that look exactly like an existing namespace, with all new objects that are just like the old ones but with some partial applied. That might still produce some unexpected behaviour, like isinstance or issubclass failing. I don't think there's a great way to get "partial classes" (see this SO question).

[carlosgmartin]

In my opinion, something like this would be helpful. My code uses a lot of tensor constructors like tensor, arange, zeros, full, eye, empty, as_tensor, etc. and is therefore littered with device=self.device flags. Accidentally omitting any such flag will cause an error. A context manager like @elistevens suggests sounds promising.

@ezyang Could those libraries not use their own nested context?

[JackCaoG]

Bumping up this thread. We(PyTorch/XLA) are recently testing large models on TPU clusters. One issue we run into is by default model initialization happens on CPU, with large models it is really easy to OOM the host memory. If we can specify the default device with a context(so we can init model weights on TPU directly) that would be really helpful.

[ezyang]

check out #82296 (comment)