torch.accelerator: A Unified, Device-Agnostic Runtime API for Stream-Based Accelerators - Yu Guangye

Key Takeaways

Thanks everyone for being here. Uh I'm from Intel. Uh today I'm excited to introduce a new API set that is coming in pouch uh 2.6 uh touch. Accelerator uh an unified device agnostic runtime API for streambased accelerator. Um, pyarch as the mo as the most widely used machine learning uh framework already support a variety of acceleration uh hardware beyond uh CPU such as CUDA uh MPS, XPU, HPU and MPU. um for uh it's well uh design architecture enable seamless backend integration uh through two uh key uh passes. Uh for uh eagle mode uh is it's based on operators and device runtime system. uh for graph mode uh is powered by inductor on top of treat uh and device uh runtime uh and each back end has its own uh runtime such as uh CUDA runtime, MPS runtime and XP runtime. Uh in eagle mode uh operators are inherently uh DS agnostic. uh most ATN operators uh accept uh device uh argument that can be said uh CUDA or CPU or other back ends. Um for graph mode inductor provide a flexible uh registration uh mechanism uh through the device interface uh abstraction. Uh this base class uh defined an API set that uh each back end uh should implement. Uh inductor use this interface to uh support uh various device back ends. Uh however the pouch runtime uh is slightly uh coupled with device specific module. uh for for for instance if you called touch.cuda current device in your code it will only works on kuda device and will fail on other back ends. uh this uh device specific logic uh limits code portability and prevent user uh from writing device agnostic code. Uh uh here is an example of how to support uh uh this scribed by uh using uh device specific uh runtime API. uh it requires using uh so many if else statements to handle each uh back end type. Uh this approach this approach post uh uh significant uh uh challenge for hardware vendor who attempt to uh integrate support for a new back end into pouch or its ecosystem. uh it's not scalable and maintainable uh for um to use this if else conditionals in these uh repositories. Uh if pyarch could provide a unified agnostic uh runtime API, it would greatly simplify the progress of uh adding new backend support and reduce the integration uh burden for hardware vendor and uh ecosystem developers. Uh as pouch continue to expand to support uh a growing variety of hardware accelerators um the need for a device agnostic programming interface become increasingly critical. Uh so we have a board vision uh make your code run on any device through a unified device agnostic runtime API. Uh fortunately uh a significant advance uh happened in parch 2.5 uh which introduce an accelerator concept. Uh an accelerator is explicitly uh defined uh as a torch do device used alongside the CPU to speed up computation. Um it adopts an uh synchronous uh execution model to handle uh and handle uh synchronization through stream and event. Uh pouch currently assume only one uh accelerator type is available uh on a given host at a time. Uh the supported uh accelerator uh includes CUDA uh XPU, MPS uh MTI, HPU and private user one. Uh notably CPU is not uh classified as an accelerator under this definition. Uh inspired by accelerator concept uh we proposed a unified device agnostic runtime API uh in pyarch 2.6 Six. Uh this API are designed uh to closely measure uh existing uh device specific uh API uh offering almost a one one mapping functionality and they uh accept the same uh input argument uh ensure minimal code change are needed when migrating from device specific API code. Uh in addition uh we introduce a new util utility touch accelerator current accelerator uh on the right list. The first one uh it uh uh returns the currently uh active uh accelerator type um as determined at compare time. Uh it means if you uh if patch build with CUDA support it return CUDA. Uh if build with XPU support it return XPU. Um if Pach is built only uh CPU support it will return none. uh minimal uh patch also provide the corresponding C++ API to uh easier for uh developer to write device agnostic code in C++. uh they are offering uh one one mapping functionality uh with touch accelerator passing API uh as a unified runtime uh touch accelerator uh overlaps with some device specific uh uh API to maintain uh backward compatibility uh both will continue to coexist uh and uh for functional consistence uh they are designed um work seamlessly together. Uh it means uh if you uh use a device specific API to uh set the current device to index one uh this change will be reflected uh by uh the touch accelerator when he curing the uh current device interface uh device index and vice versa. Uh now we already have the generic uh uh runtime API as mentioned above. But how do we handle the uh other uh runtime components uh comp components uh such as the stream and the event. Uh this issue was uh addressed in pyarch 2.5 uh pouch refactor touch stream and touch invent uh into a base class that support some some uh common uh method. Uh these common methods are available across different uh back ends. Uh and the device and the device specific stream or event inherence this base class and can override uh method as needed. Uh now uh and touch data stream uh can represent a CUDA stream on CUDA device and uh XP stream on uh XP device. Uh this b uh this design uh the benefit of this design is that since touch stream and the touche event is base class uh we can capture any device specific uh counterpart as the base class uh even if uh uh in in in Dynamo or other components uh especially for auto tree back And uh when discussing the design philosophy uh of touch accelerator here is a quick uh story. Uh our initial idea was to put a unified runtime uh under touch name space directly um uh such as touch.fu. Uh this uh this is inspired by uh how touch.xpu.stream was generalized into touch.stream. Uh following this model uh touch.xpu.fu uh would be generalized to uh touch.fu. uh but uh uh after discussing uh with Elvin and and other uh community contributor uh we decide to put uh this API uh and touch accelerator uh name space. This is a new uh name space introduced in pyarch 2.6. Um and we have and the uh le uh clear uh design ruler. Uh if uh API work across uh CPU and accelerator uh it should uh uh goes in uh touch n space and uh accept a device type argument. Um and if uh API is specific to accelerator it would be place uh touch accelerator name space and doesn't take a device type step. Uh before touch accelerator pouch use different method to support various uh device back ends. uh for example FSDP uh use touch.get device module which u delegate to uh device specific uh module and inductor use an inheritance pattern uh where each back end should uh inherent from an base uh class and implement it subclass. Uh this delegation and inheritance patterns are also commonly used in other uh popular uh repositories such as for hug and force hug and fist. Um compared to previous approaches uh touch accelerator provide a unified uh built-in API that is easier to uh use in Dynamo and more uh pass. It provide a more elegant and uh consistent uh uh interface for device agnostic programming. The main uh but the main uh limitation is that uh touch accelerator currently only support a subset of API uh that are available through two other uh approaches provided uh for example uh for some memory mag uh uh for some memory related faction uh such such as touch uh device tab uh empty cache uh is currently not yet available uh in touch now accelerator name space. Um how to support touch accelerator for new back end uh is is straightforward. uh you only need to uh implement uh uh the device guard implement interface and register it uh instant instance. uh this leverage an exciting uh existing uh mechanism uh through uh the stream guard and the device guard uh in C++ reuse. Uh in pouch we can categorize uh runtime uh into six components uh such uh device stream event uh got generator and allocator. Um we uh now we already generalize uh device stream and event. The next challenge is how to uh handle stream guard and uh uh device guard. [Music] Uh previously pouch offered touch device uh type stream as uh uh context manager to ensure the operation is running on the expected stream. Um but now uh from parch 2.7 uh the dessert stream can be used directly through this statement. Uh it's easier to uh use to generalize uh for different uh back end. uh similar pouch uh provide touch device type device as a context manager to ensure the operation is running on the expected device. Uh it's a front end pass API. Um but now uh touch accelerator provide a new uh API touch. Accelerator device index to do the same thing. uh is um uh it can be used uh in uh on different uh uh device back end. Uh after analyzing usage across different uh popular repositories, we found that some memory related API are commonly used in line with RFC 134 uh 078. uh we propose uh introducing a set of uh touch uh set of device memory API um and uh would put them uh under touch accelerator name space to better support uh the common uh user case. As another followup, uh we plan to generalize pyarch unit test uh to improve uh reusability acrossing different uh u back ends. Uh now uh I'd like to revisit the uh example mentioned earlier. Uh with touch accelerator the implementation becomes much clearer and uh simpler. It makes the refined code uh can run on any device without any adaptation. Uh finally I would like to express my sincere gratitude to the community maintainer airbend. Uh his insights and guidance have been instrumental in driving uh this work forward. Uh thanks