ONNX June 2022 Community Meetup, 13 Jul 2022

Previous Meeting Next Meeting

⏯

youtube image

►

From YouTube: PFVM - A Neural Network Compiler that uses ONNX as its intermediate representation

Description

PFVM is a neural network compiler developed by Preferred Networks, which relies on ONNX as the Intermediate Representation format. PFVM is used in production environments to deploy models to various devices such as GPUs, multiple edge computing architectures, and PFN's own accelerator, MN-Core. PFVM's most salient features are; automatic checkpointing, operator fusion, and graph simplification that can be applied even when models have dynamic axes or unknown shapes. ONNX Shape inference becomes a critical element for all these optimizations, and the importance of bringing up more advanced shape inference mechanisms to address complex optimization scenarios is discussed in this talk.

Zijian is a Neural network compiler engineer at Preferred Networks and an ONNX SIG-archinfra member.

A

uh Hello: everyone, uh I'm genji from preferred networks, and today I'm going to talk about pfvm, a neural network compiler that uses onyx's its intermediate representation.

A

uh I'd like to first introduce briefly about our company preferred networks. Preferred networks is a company located in tokyo, japan.

A

So it's 7 15 in the morning here and uh we have various projects uh uh and- and we are solving real world programs by deep learning such as material discovery, medical image, analysis, character, design and many other projects.

A

We also have some open source projects known as coupei, which is a numpy compatible tensor library accelerated by cuda, and we also have a black box, optimization tool of tuna and we used to develop trainer a defined by runs, diode framework for writing neural networks, and we have our own super computer for deep learning code, mn3, which is known as the most energy efficient supercomputer in the world, and I am a member of compile team and I am in charge of uh developing pfpm.

A

So today's topic is uh I'd like to share a use case of onyx for module optimization.

A

I think the most popular onyx use case is for module deployment, uh so uh we export zonex from python code and we deploy the module to edge computing architectures using onyx runtime or some other libraries, but I'd like to focus on module optimization today and secondly, I'd like to share our motivation to have solid shape, influence functions.

A

The pfvm is pfns our preferred networks, we say pfn pfns in-house neural network compiler and runtime uh pfvm works as boss, compiler and runtime. So pfpm takes exported onyx as input and optimizes it and executes the model with specified backends using third-party apis.

A

However, unlike other shape, uh unlike other inference, engines, vfam can be also used for training models. In such case, pfm compiler extends the input onyx and appends the gradient nodes for training.

A

And this is how uh pfm works uh so, given the impedance graph, if the model is for training, then we uh generate backward pass and after that we optimize models.

A

So during optimization we have various kind of optimizing paths and some of them are going to be explained later today and uh during optimization we extend onyx by some customer operators, for example, uh since onyx does not contain device information, we add a hosted host to device operator and device to host operator to optimize device, changes and memory and after optimization we schedule the computation graph and get a sequence of operators.

A

And after that we execute each operators using leaptorch api or some accelerators apis.

A

And we have several reasons to use: onyx's intermediate representation, pfpm uh first onyx is stable and well documented. It does not change very often so uh it is stable, yeah and, secondly, uh focused on interface and doesn't do too much and unexceptional influence is very useful.

A

uh We don't have to write, shape, inference function by ourselves and, uh moreover, onyx has useful test cases in its repository and we use them for a lot. Many tests codes and, let's see how shape influence, is important in module optimization and we have a civil optimization.

A

uh But I'm going to introduce three today and the first case is element wise fusion.

A

A

Yes, let's take cut, take a look at this uh onyx model, so we have two inputs a and b, and uh will we first reduce a and get a scatter and other constant, and then we substract and multiply something and get output when we execute this computation graph on cuda.

A

The cuda profile will look like this, so we have reduced max kernel at kernel, subcarnival kernel and before each kernel uh we have uh overheads known as the cuda kernel launch and it is important to uh reduce this kernel range overhead for faster computation and if we can fuse these element twice, operators into a single fusion group operator, then uh the kernel range is only twice and the computation time will be uh been reduced.

A

So we basically we want to fuse element twice operator every time, but can we always fuse a decent element by operators?

A

uh The answer to this question is no, and we must check whether they did they can be broadcastable or not before a gradient fusion group. So in the left case,.

A

Since all of these values have shaped n, they can be fused together, but in the right case, where a has shape n and b has shape n, then we cannot fuse these uh operators together, because uh we don't know if m is equal to it.

A

Yeah we don't know. If n is equal to m or not so.

A

If we see this case so in this case.

A

If shape inference is failed uh and we get unknown shape, then uh this is very uh troublesome case. So in this case uh we have a single input x, uh which have shape a times b and uh we reshape it to one dimensional tensor, so unknown, zero and unknown. One is both a times b, so they have actually they have the same uh shape, but current on its shape, inference cannot infer that unknown. Zero is equal to unknown one. So in compile time, we don't know uh that these unknown values actually have the same shape.

A

So we lose fusion opportunities in this case, so shape. Inference is important here, and the second case is graph simplification. I think many uh inference engine do the similar uh optimization here. uh So uh again we reshape uh this value uh to one dimensional shape. So in this case we know that unknown zero is equal to a, but again current only shape inference cannot infer unknown 0 equal to a so. We don't know that a is equal to unknown 0 in compile time.

A

In this case, uh we can uh so if we can infer that a is unknown zero, then the reshape is just the identity and we can remove the operator and we can do the element fusion, but uh where we cannot do this right now- and you may have a question like do model contain so many unnecessary operators. Or do we really have to do this kind of optimizations, and the answer to this question is uh well.

A

I think models does not contain so many unnecessary operators if the model is written by human, but models may contain a lot of uh unnecessary operators if it is generated by programs like what first generation or neural network and neural architecture search and I'd like to show an example of how these optimizations are important. uh In an example, uh let me proceed the slide.

A

Oh also appendix is not included. Oh oh.

A

Okay, then I just skipped this part. uh Well, what I wanted to explain was uh so if we generate a backward pass for a specific node, usually the backward pass is uh a little bit complicated.

A

So uh well, I'm going to just skip it and the third case is automatic checkpointing and uh some of you may have heard of checkpointing before, but uh I'm going to explain it from uh beginning the checkpointing is a technique to reduce memory usage when executing module.

A

So let's see the left computational graph in this computational graph from two inputs x and y, we compute z and, uh however, when computing this uh computational graph, the memory usage is not very good good. I mean uh when you compute, node five. You have to remember the output output of one out of output of two output of three out to the four on memory and the memory usage is costly.

A

However, you can modify the computational graph uh to the uh to this.

A

uh This computational graph are different but computes exactly the same output and because the nodes dependency is uh completely the same so uh but all of these computation graphs have different latency, which is execution type and different memory usage.

A

uh So pfpm automatically uh transformed the input, onyx graphs uh so that uh the latency is increased a little bit, but the memory is reduced. A lot pfm can do such optimizations and it is very important to need to know tensor sizes to estimate memory usage and uh it's relatively simple.

A

When uh models are static, so we, where we don't have uh dynamic dimensions there, but even when modules have dynamic access, we can do automatic checkpointing, because uh users can estimate them in compile time like a number of nodes, is 1 million in some uh graph neural networks, or something like that.

A

uh However, unknown dimensions is not welcomed in uh in automatic checkpointing, because uh well we don't the the estimation won't work. We don't know if the unknown dimension is just one or just two or one million or something.

A

So this is the salary slide. Why shape reference is important, so uh once we get uh unknown dimensions, we will have many unknown dimensions after that, so usually in most shaping frames. If the input tensor has unknown dimensions, then the output tensor also have unknown dimensions and we lose many opportunities of optimization. If shapes are unknown. As we see where we cannot do element transfusion, we cannot do a graph simplification.

A

We cannot do automatic checkpointing.

A

And let me say my ideas for current onyx shape inference. So, as you may already know, uh on exchange for influence has been improved. A lot recently and uh symbolic inference was introduced from 1.10, because symbolic inference is a a chip inference scheme using data propagation.

A

So, for example, in the right case, we can infer the output shape of constant of shape as our a times b from the input tensor of shape 2, because the a a and b are propagated globally from the top.

A

Shape information- and this is really nice and uh for static case. I think current onyx shape inference already works quite uh nicely and we can run the shape inference uh in most of time for major workloads, but for dynamic case.

A

uh I think we want more support, so the this onyx graph is an example that I showed a little bit earlier. uh So in this case we want to estimate the how to output the shape of reshape to be a, but currently it is unknown, zero and uh well. I think we have to uh implement more inference functions for dynamic case and well.

A

I I have a question like this: uh do we need to support case like concat of two cancers of size, m and n? Well, uh I I think this question is a little bit controversial and uh I'd like to hear some use case. If you think this uh shape inference is important for your workload.

A

uh This is last slide, so uh this slide is a little bit, not so much related closely to onyx, but this is slide is for optimization.

A

So we have our uh defense on super computer called elements core and md4 only accept static module, so the models where all of the tensors have a static shape and module does not have dynamic branching. So a mn chord does not accept models that contain e4 loop, but since the models static, mn4 can do many optimizations.

A

uh Emma decor is a very large scale, fasting. The machine and its scheduling is really challenging and where we use the static information to do a very sophisticated optimizations for scheduling uh yeah. The interesting thing is since mn4 is a model static, mn4 programs know exact computation time or memory usage or memory address, of everything before uh running the module, so at compile time and pfpm is integrated in medical compiler and do some of simple connections and shaping frames. So what I want to say is: if these shapes are known, we can do shaping inference.

A

We can do many optimizations. The shifting frames is very important in every optimizations.

A

So thank you very much. uh If you have any questions, please ask it on chat or well.

A

Please reach out to me thank you.