pytorch - 💡(How to fix) Fix Allow registering multiple out-of-tree backends and setting current_accelerator [1 comments, 2 participants]

🚀 The feature, motivation and pitch

Currently PyTorch fundamentally assumes that you have one accelerator, which has a few consequences:

1. It is not possible to register more than one out-of-tree backend extension at the same time, and trying to do so is a hard failure. Combined with the autoload mechanism, this means that if you get two out-of-tree backends installed in the same environment importing torch itself will fail.
2. There is no way to set the "current accelerator" to a specific accelerator. If I were to install a "testing" backend, for example (I think openreg might meet this standard), I probably do not want it to be current_accelerator. You can also imagine alternate backend implementations for CPU and CUDA (e.g. if we wanted to add an XLA backend for use with CUDA GPUs, we might want to provide that as an option for users, but not have it automatically used for accelerator.current_accelerator().

Ideally, the out-of-tree extension mechanism would be changed to allow registering arbitrary backends rather than having a single hard-coded slot for "other". I think there are also privateuse2 and privateuse3 backends, but no corresponding interfaces for renaming them, and no interface for selecting whcih one is the current accelerator.

For problem 1, my preferred solutions would be, in order:

1. Allow registering an arbitrary number of backends instead of renaming privateuse1.
2. If for some reason all devices must be hard-coded at compile time, expand out the number of "private use" backends to something like 5 or 10.
3. Allow registering an arbitrary number of modules, but allow swapping out which one is the privateuse1 "device". (This sounds fraught but it's the only other real solution I can think of that maintains a fixed number of devices while allowing registering arbitrary numbers of modules).

I think in both 1 and 2, it would be best if the interface were torch.utils.register_backend(device_name, module) and the renaming of privateuse1 or privateuse2 or whatever would be an implementation detail not exposed to end users.

For problem 2, I think the solution is a lot more straightforward from an interface point of view, you'd just want to add a torch.accelerator.set_current_accelerator(device) function (though I suspect adding that interface will expose a few more concurrency bugs than you currently see, even though it's not necessarily adding new capabilities since you can always rename the privateuse1 backend and implicitly set the accelerator at some random point in the program's execution).

Alternatives

No response

Additional context

No response

cc @bdhirsh @albanD @guangyey @EikanWang