Migrating existing backend-MAIA integration toward PrivateUse1 / openReg

[trajepl]

We’d like to follow up on the feedback regarding backend integration and the recommended direction of using PrivateUse1 via the open registration (openReg) mechanism.

Our team has already integrated a custom accelerator backend using the older in-tree approach. This includes a large amount of internal scaffolding — kernel registrations, memory allocator, device guards, dispatcher logic, and more — which are tightly coupled with the current PyTorch internal APIs. As a result, a full switch to PrivateUse1 is non-trivial for us in the short term.

That said, we fully agree that PrivateUse1/openReg is the right long-term path for out-of-tree device integration. We are very open to gradually migrating toward that direction and would appreciate your input on how best to plan and prioritize this work.

Ideally, we’d like to find a way to incrementally align with openReg while minimizing disruption to our current infrastructure.

cc @bdhirsh @NmomoN @mengpenghui @fwenguang @cdzhan @1274085042 @PHLens @albanD @golvok @wschin

[trajepl]

https://docs.pytorch.org/tutorials/advanced/privateuseone.html#facilitating-new-backend-integration-by-privateuse1

According to the document, support for distributed collective communication is still a work in progress, whereas our framework requires those functionalities.

[FFFrog]

According to the document, support for distributed collective communication is still a work in progress, whereas our framework requires those functionalities.

Hey, distributed related functions are now supported, and the document is outdated

Plan:

• We will update this document and provide more detailed integration guidance
• We will continue to improve OpenReg, which will be used to protect the third-party backend integration mechanism and also as a reference implementation for third-party backend. For example, it will support the distributed mentioned above

If there is anything else you want to know, welcome to communicate and discuss

@zeshengzong @albanD @trajepl

[wschin]

Any overhead caused by using OpenReg than in-tree integration approach? I want to see comprehensive profiling result to ensure we see no regression if we switch MAIA to OpenReg. I see a lot of py::gil_scoped_acquire acquire such as in

void uncheckedSetDevice(c10::Device d) const noexcept override {
    py::gil_scoped_acquire acquire;
    auto device = get_method("uncheckedSetDevice")(d.index());
  }

which is obviously a performance problem comparing to in-tree integration.

I also see class Driver in python file, which looks like a fundamental components in OpenReg. Does this mean new backend cannot be run with pure C++?

Will Meta's in-house acclerator use OpenReg?

[FFFrog]

Hey @wschin

There may be some misunderstandings about OpenReg.

The third-party backend integration mechanism based on PrivateUse1 is different from OpenReg. The former is the main way to integrate new backends into PyTorch, while the latter is a CPU-based accelerator backend that is integrated into PyTorch through PrivateUse1 and has been added to PyTorch CI, thus helping to ensure the quality and stability of PrivateUse1. In other words, OpenReg is currently only used as a test backend.

In order to speed up the development of OpenReg and simplify the maintenance of the backend, we chose to use Python to implement some functions because it is simpler and easier to use. Of course, the performance may be slightly inferior, but as a test backend, this is irrelevant.

Regarding the future development of OpenReg, I have left a message to @albanD. If you are interested, you can refer to this comment for details.

cc @Yikun @hipudding @zeshengzong @ezyang

[albanD]

Thanks for opening this!

Sorry for the confusion on OpenReg, it is, as @FFFrog mentions above, only an example on how the core PrivateUse1 can be used.
You would NOT be using any of the OpenReg code yourself, but use the PrivateUse1 APIs.
The benefit you get from OpenReg are: it is in core CI and so anything there is guaranteed not to break (contrary to maia integration which has no ci signal and might break at anytime) and is a fully working example of some of the features (we're adding more) this extension point provides.

Also the direction of #155541 seems to be that you are planning on aligning the maia device to behave like every other "Accelerator" device.
A quick grep in the codebase for cuda or privateuse1 will show that there are a very large number of places that will need updating if you're aiming for being as feature complete as cuda. We're more talking in year(s) of work than week/month here to just enable it, and then multiple engineers full time to maintain. And to be clear, a lot of that time is not going to be easily compressible on your end as most of the challenge is going to be about integrating within the core components which are very subtle and challenging pieces to change/adapt.
This is why, for example, all the XPU enablement work today is done at the Accelerator level and not under the xpu name. Similarly all the work to enable new backends by many companies is done jointly via the PrivateUse1 feature so we don't have to spend N years to integrate N backends and can have a lot more engineering focusing on these shared core integrations.

To be fully transparent, we have a lot more bandwidth to support efforts such as Accelerator or PrivateUse1 that benefit many users than one-off integration for a single backend.
I definitely understand that this is a significantly different plan for your current one, but I am not convinced that your current plan is going to be more effective, even short term, due to the (unfortunately) very limited bandwidth from the core team compared to moving out of tree and being in full control of the pace of progress.

[wschin]

If OpenReg is the one to ensure 3rd party's accelerators' stability in CI, it needs to follow real accelerators' (like XLA and Ascend NPU) design. Most of the "Driver" and "Execution" happen in C++, instead of Python. A C++ based OpenReg will be way more helpful than a Python one.

For the long run, I am not against a migration to new mechanism. But the document, protocal, expections, APIs, examples must be crystal clear. From what I learned from these comments, it seems you guys have some "golden model" that we should follow. If so, please update related document (including a simple but real example, instead of python-heavy OpenReg) at your earliest convenience, so we can evaluate it on our side.

Thank you.

[albanD]

If OpenReg is the one to ensure 3rd party's accelerators' stability in CI, it needs to follow real accelerators' (like XLA and Ascend NPU) design. Most of the "Driver" and "Execution" happen in C++, instead of Python. A C++ based OpenReg will be way more helpful than a Python one.

The reason this code is in python is to make it simple and easy to modify.
But if there are key behavioral changes between python/c++ in this context, I'm curious to hear and not opposed to moving this logic to c++ on principle.

If so, please update related document (including a simple but real example, instead of python-heavy OpenReg) at your earliest convenience, so we can evaluate it on our side.

Happy to continue to improve documentations and tutorials.
Please keep in mind that this work is done by open source volunteers that are building these abstractions and systems to help you enable your specific use case more easily, so timeline on this may vary.

[wschin]

Some reasons not to use Python for OpenReg:
(1) Python is not the language accelerators use. How do you discover edge cases when most people use C++ for NPU? I have seen many times that Python and C++ behaves differently for LLMs.
(2) Python is way more flexible than C++. It's easy to see dependencies problem (missing symbols in DLLs) when moving Python infra to C++. Without writing PriviateUse1 in C++, how do you guarantee all symbols 3rd-party needs are exposed?
(3) Official turotial will become reference implementation, and many of its conventions will be absorbed by 3rd party infras. I hope that implementation should use best integration path, instead of a toy PoC.

Timeline is a big concern. Meta's in-house accelerator takes the advantage of not waiting on PrivateUse1, while Microsoft MAIA can be easily blocked by 3rd party contribution. I want to learn your opinions to make the competition fair.

Some questions are hard but I really need to understand them. Thank you very much.

[FFFrog]

The current situation is:

• The PrivateUse1 mechanism is fully complete, and from a design perspective, PrivateUse1 and built-in keys such as CUDA are equivalent, so there will be no performance impact. What needs to be worried about more is the compatibility between upstream and downstream.

• As a public approach for integrating new devices, the biggest problem currently is the lack of detailed documentation and simple and clear implementation references.

As for OpenReg, its first priority is to take care of the code quality of PrivateUse1 in real time within the tree, while not introducing too much development and maintenance costs. I am currently involved in some OpenReg improvement and test migrations. While Python-based OpenReg introduces convenience, it also introduces some complexity, such as using multi-process to simulate device memory, and it is very troublesome in terms of operator calls, especially metadata changes.

From my personal perspective, it is necessary to re-optimize OpenReg according to the standard real accelerator integration; the standard streamlined reference implementation combined with detailed tutorials should be able to solve concerns and confusion similar to MAIA integration.

[wschin]

@FFFrog, @albanD, why do you mean by "new backends" that has to migrate to PrivateUse1? Backends introduced after specific time? MAIA is not a new backend.

The ask to me is basically the same as asking MTIA to migrate to PrivateUse1 starting today. This honestly sounds too risky. Or MTIA team already agree to migrate to PrivateUse1?

A more realistic plan needs grace period. Timeline:

1. PrivateUse1 and OpenReg finishes a reference implementation and all documents.
2. Then, we start trying PrivateUse1. In the meanwhile, we still improve MAIA the old way.
3. After 6(?) months, we finish the migration and the future changes will be for PrivateUse1.