hello again
next up at the data science and
analytics specialist track here at pycon
australia 2021 we're welcoming matt
kelsey who's a fixture at the mli meetup
here in melbourne
during the day he works as a machine
learning research engineer at edge
impulse a development platform for
machine learning at the edge because not
everything has to be centralized on big
servers belonging to someone else
matt has worked across a range of
machine learning domains over the last
20 years including work at thoughtworks
are sponsored today google brain
wavy and iws
mat blogs at http.com
today matt will be talking to us about
jax which provides automatic
differentiation
allowing extremely high performance
machine learning on modern accelerators
all from python
matt will be showing us the fundamentals
of jaxx and an intro to some of the
libraries that are being developed on
top
so please a big venueless hand of
applause for matt kelsey and high
performance machine learning with jax
awesome thank you so much yeah this is
this is a fun talk so i've been on sort
of a jax bandwagon for about a year and
a half now and it's just it's just
getting better and better so i'm very
excited to share this tool
um yeah so this is about high
performance machine learning so there is
quite a bit of a focus on the sort of
things that jax provides in terms of
making things go particularly fast
there's going to be three parts i guess
to this talk there's going to be some
sort of fundamentals so we're going to
talk about like what are the building
blocks that jax fundamentally gives us
i'm going to talk a little bit about
some of the things that
it provides around multiple hosts
because you know really scaling up over
many machines is the the best way we
have of running at very high sort of um
very large scale
and because jax is quite um simple in
what it does um it's sort of you know
very fundamental we'll touch on a little
bit about the higher level libraries
that we often use because jax by itself
may not be enough
okay so we're going to start with what
is jax that's that's a big question and
i'm actually going to start a little bit
with like an overly simplified view of
machine learning because for a you know
machine learning is a very very broad
field but for a lot of the work we do it
does sort of boil down to these
supervised learning type problems where
we have some set of input that we're
trying to map to some output and we're
really interested in these functions
we're trying to learn these functions f
and even though it's not you know the
complete domain of machine learning
a lot of a lot of supervised learning
uses gradient descent these days it's a
big thing
so there's a couple of things we want
out of a tool like jax for
gradient descent based methods and the
first one fundamentally is gradients we
need to
get these functions and and calculate
what gradients are sort of you know
described by these functions and
you know you really can think of jax as
the next generation of a very old tool
called autograd maybe not very old it's
about what 10 years old now
and because there are some of the
original authors of autograd which you
know really pioneered a bunch of
symbolic stuff are actually the authors
of jax as well so it's fair to say that
jax is the next generation of autograd
and um you know we we need things to go
really fast that's something you know
very data intensive very compute
intensive and jax is unashamedly
just a really great python interface for
another tool called xla so this is the
accelerated linear algebra compiler
which um was developed by google and
we'll talk a little bit more about this
but you know it's not unfair to say that
jax is just like a front end for xla
because to be honest a lot of the heavy
lifting is done by this compiler tech
which jax just provides a very
elegant python interface for
okay so let's talk about these two
things but before we get to those i just
want to say that if nothing else
the takeaway from this talk is that
jax is a good implementation of the
numpy api so
if i think about you know some numpy
stuff we might want to do here's a
arbitrary calculation i'm describing
where i've got a dot product between a
couple of matrices i'm adding something
we can actually
substitute in rather than using numpy an
implementation of these things that is
provided by jax numpy and you know
i've done this quite a bit i've got some
signal processing work i do as part of
work and that is all described in scipy
and numpy and it's actually i've been
pretty it's been pretty good at just
literally porting it to jax by swapping
the numpy imports to jacks.numpy and
with a couple of caveats and small
changes
i'm able to make my programs run under
jacks instead now
the wh why do i care well all the things
we're going to talk about that come in
this talk about the high performance are
suddenly available to me and they're
just sort of transparent so if nothing
else walk away from this talk right now
with with an experiment you might want
to try where you just literally swap in
jax numpy instead of numpy
but let's talk a little bit more detail
about then some of the things we've got
too so auto grade this is the first
thing i talked about wasn't it here's a
function it's a pretty simple quadratic
2x squared plus 3x plus 3. if if this
was um the complicated machine learning
function i was trying to work with and i
was using gradient descent
i might want to
get some gradients for this function for
something so
the first thing i want to introduce from
jax is is a little function called grad
and what grad gives us is um a gradient
of a function so grad takes a function
and returns another function
where uh when i call this this this g of
f i get the the gradient with respect to
the the parameter there so let's see 2x
squared i bring down the two that's 4x
plus um three so uh the gradient uh at
three oh sorry the gradient at three is
15. so this this is good
now this illustrates a real key point
about how
jax fundamentally works this thing that
jax has given us grad takes a function
and returns another function with the
same signature
so we're going to see this time and time
again that the main thing that jax does
is these functional transformations
jax looks at functions and gives us new
functions back that are in some way
changed based on what we wanted to do
and how does jax do this well it's based
on an idea called tracing a fundamental
idea so here's another thing that jax
has it's these make expressions and what
i can do is i can say
here's my function f again i'm going to
pass that to jax to say well make me a
new version of this function that's
called trace trace f when i call this
function it doesn't evaluate it directly
the return of this function is actually
jax's view on what is actually happening
inside this function so you can see here
we've got this sort of symbolic
representation of the compute and we
don't need to dive too much into this
but you can see there's some things
being multiplied and some stuff being
added
and so jax is really about
doing these things by tracing by tracing
this function and having a symbolic view
of things we get a bit of insight into
how the gradient is being done so if i
do that same tracing again this time on
the gradients i can see a different
function and so this is fundamentally
about how autograd works we trace these
functions we do this forward pass this
symbolic representation which is then
manipulated to turn into a
gradient now there's a really and again
we don't need to talk about the
specifics here there's some numbers and
complicated ops floating around but
there's a subtle point here i want to
make this is interesting about functions
this make jax expression takes a
function as does grad which takes a
function so you can see here we've
already got this nice composability i
take my function f i wrap it with this
grad which mutates f to be the gradient
i then can then pass that to jack's
expression so jax is really about these
building blocks these building blocks
which take functions and and give us new
functions that we can compose in
interesting ways so this composition of
functions is a really big important
thing
okay so the second part we talked about
was going fast and the first
thing that jax provides us is some
just-in-time compilation so again here's
our function 2x squared plus 3x 3.
very interesting one thing that jax
gives us is a jit so we can take this
function again rf function and we can
just in time compile it to get a new
version of the function that is
that has the same signature and can be
run with the same values and give us
back the same result so in many ways
this is the most anti-climatic slide
ever ever written because it looks
pretty boring but there's a lot of
really interesting stuff that happens
under the hood with this jit and
fundamentally an important part of this
is the fact that um this this
computation that's being done here would
now be being done on whatever
accelerator i have so even though i've
described this compute in terms of just
just vanilla
python or using that jack's numpy if i
had a gpu this would be running just by
itself immediately on my gpu which is an
important thing about how our
compilation works
for for jit and how does jit work well
again i can talk about looking at the
the jacks expression what is actually
what has happened when i've called this
chit function
we've got our f like we had before and
it doesn't really matter what this f is
it's it's quite simple here but this
could be a very very complicated you
know 100 layer neural network or
something
and what's subtly different here is now
jax is saying i'm going to actually make
a call out to this tool xla so what jax
has done here is wrapped up
had a look at this function run through
it built this trace and then it's
passing it to this compiler under the
hood and that compiler is is very high
performance you can do all the classic
things we want from a compiler
a couple of key things like you know
fusing ops together or reordering things
whatever it wants to do but also most
importantly
uh converting things to kernels that are
going to run on the device i have so
this would natively run
xla will make sure that it's the compute
is then done on my gpu or whatever i
might have tpus and stuff we'll talk
about later
so that it it it looks like it's such a
simple function but it's doing so much
and it's incredible how much it does
so that's a fundamental thing that we'll
be dealing with
another another really important thing
with um making things go fast that we're
probably all familiar with is
vectorization
so here's a here's a function that looks
a bit scary but um
it's a bit arbitrary i just sort of made
this function up it's some function that
does some it does some weird stuff don't
try to reason too much about what this
what this is actually doing it's got a
couple of matrices it does some you know
along axis mins and maxes it does a dot
product of things and then it added adds
a scalar like this doesn't actually mean
anything i just wanted to say that
there's some you know compute you might
want to do
and here's the result we get down the
bottom now if
if we wanted to run this function over
lots and lots of data the common way we
would do things is to vectorize stuff
and by vectorizing our mean we normally
think about this idea of adding a
leading dimension into our data
so that when we pass it to the compute a
lot of the linear algebra libraries
underneath will work a lot faster if you
can do things in batching
now you know we're probably all familiar
with this idea that when we do add these
leading dimensions
we just have to make sure that our code
is aware of those leading dimensions and
can sort of handle them so if i was to
vectorize myself this function i would
go through and just make sure a couple
of things hold so i can see these axes
here well now that i've got another
dimension a batch dimension i have to
make sure these axes sort of move around
i've got a dot product here and i know
by experience that the dot product is
okay with having leading dimensions so i
don't need to change anything there and
then i've got an addition of a scalar so
i might i might be a little bit worried
and think what's the broadcasting going
to be i'm going to make a mistake about
something here but you know i think if
you've done a bunch of stuff with numpy
or any other sort of equivalent compute
you sort of get a sense of knowing how
to modify this function to make a batch
aware
but we don't need to do this anymore one
of the things that jax provides is this
awesome tool
called vmap and what is vmap again it's
a transformation of functions i take my
function f i run it through vmap and jax
returns me a new version of my function
and what it does is actually by using
this tracing idea re-plumbs the compute
that is expressed by f to automatically
handle the the leading batch so now if i
have this vectorized f i can call
through with an a b and c which has now
got this leading dimension so i've got a
i've added an extra two here basically
and i've turned this scalar into a short
vector
now for that particular f we were just
looking at it's not terribly exciting
and interesting but as things get more
and more complicated this idea that jax
will be responsible for the plumbing of
the of the leading dimension is is a
really big thing and it can make this
huge difference here's an example
where i where i use this which i thought
was really interesting i had i had some
data this is some satellite imaging and
i wanted to basically black out a bunch
of augmentations of that image so you
know with satellite imaging one of the
classic augmentations we can do is to
you know
sets of 90 degree flips rotation sorry
and then potentially left and right
flips so i've made this little augment
it rotates 90 degrees by some
number of the 90 degree turns and then
it either flips or doesn't flip it's
like a zero one so here's an example so
i
pass an image in i've said rotate one a
lot of 90 degrees and don't flip so
basically i get this image down here
i can use vmap to express the compute
that i want to do all of these in
parallel so what i'm going to do is i'm
going to say here's my augment function
and what does vmap do it adds an extra
front dimension basically so like like a
leading dimension to all the arguments
so i'm going to say here's my rotations
i want to do 0 1 2 3 one two three and i
want to pair those zip those together
with all the possible combos of flips
so you can see these eight these eight
pairings here between these two
correspond to all the possible
augmentations
so what i'm going to do is i'm going to
say i want jax to make a version of
augment that handles batching but one
little thing that i've said here is that
i only want you jax to
add that leading dimension to the second
and the third argument so this is this
rotation and flip but i don't want you
to plumb that extra dimension for the
first argument so i've said none here
and what jax will do is basically
broadcast that image multiple times now
when i call this function i get this
result because what's happened here
under the hood is that jax is zipping
basically across these
and and it's going all the way through
it's going all the way through down to
these functions to make them vectorized
to do that which is what which is really
cool i think
so now when we call this image
uh combos or combos augment rather than
just getting one image back we get eight
images back now i might actually also
want to do this whole thing in batch as
well so i want to actually start with a
set of images not just one image but i
might have like ten images
and i want all eight combos for those
ten so i've got my all combos augment
which was v-mapped i'll just v-map it
again because what does v-map do it
takes an arbitrary function it just adds
another batch dimension at the front and
because i know that this is sort of the
level of compute i want to deal with i
might ask
jax at this point also jit so go through
and now take this compute that's been
expressed by a v map over a v map with
these with these sort of fixed values
and do whatever you need to do to
compile it to make it really fast
and that this is blindingly fast the the
code it ends up being resulting runs
very very quickly and this is a pattern
i've seen a number of times now i have
this sort of you know if i'm v mapping
something and i have extra batch
dimensions appear i don't need to be
scared i just just keep wrapping with v
maps and everything is fine so that's
really cool so fundamentally vmap does
this by re-plumbing and sort of
completely re-changing what my function
is very handy
a quick side note on
on tpu so tpus are google's asic custom
designed
sort of replacement for their in-house
gpu farms
tpus you know come in a couple of
different sort of sort of models the the
third generation the smallest tpu board
i can get if you you know renting one
basically from the cloud is this our v3
8 and 8 corresponds to the fact that a
tpu
board like this has four chips each of
the chips has two cores so i've got
eight cores in total
what's interesting about this is when
you rent these tpus that's the minimum
you get you've got no choice each of
these eight cores is pretty powerful but
you can't you can't go any lower than
that so that's interesting and and
google provides these tpus in these sort
of you know
boards like this or in these huge fat
combinations of boards with this crazy
interconnect uh where they that's called
a pod where you have um hundreds of
these tpu boards put together to to make
effectively 2048 calls now
the reason i bring this up is because
the development of the tpus was done
hand in hand with the development of xla
which is the compilation technology we
talked about xla was sort of introduced
at the same time as an intermediate
between the code being run either to
gpus or tpus so xla is very much around
making high performance code for
tpus and jax is in many ways a front end
for xla so we're going to get a lot of
really a huge benefit from using jacks
very very simply across these very large
compute clusters and and it's quite a
standing way you can do nowadays that
even like a couple of years ago was
quite difficult to sort of orchestrate
so i'm going to introduce another
concept called pmap so this is
parallelization rather than
vectorization and you know if if you go
right now to google colab and you start
a collab you're allowed to just get a
free tpu so if you sort of you know pick
device tpu and you just sort of import
jacks maybe there's some stuff you need
to do but it's basically import jacks
you'll get jax's has a view over these
eight cores
so rather than just single one we've got
eight and jax provides us some tools for
making use of this
here's an example of how pmap works
here's a function it's very simple again
abc just we'll just add them up and
we're going to p map this function
like we did with the v map and pmap
takes a function and it gives us a new
version of this function that can run
and um it's again this idea is that
there's a leading dimension which is
being put on top of things so if i have
some a and b which are eight by three
and i'm on this tpu
machine which has eight cores then when
i run this compute i'm actually running
this this computation
in parallel each of these rows being run
on its own device
so what what jax is doing here is it's a
little bit different to vmap vmap we
actually re-plumb the code pmap actually
takes this function and actually does a
compilation of it and ships it out to
the devices we have so if i have this
the v38 which has eight devices when i
call this pmap i'm actually
compiling this code
sharding this this input and actually
sending it out to the devices and all
the devices do things in parallel and
it's not until i actually instantiate it
in some way in this case printing it to
my screen that the result is actually
brought back from those devices and
assembled and that's why we see this
here sharded device array
because until i actually do something
and act on it it turns out that actually
each of these rows lives in the memory
of the device so this is this is really
important because if you've done much
with like multi machines you you really
need to make sure that the the data
being shipped around versus the code
being shipped around is under your
control because sometimes you don't want
to just be
you know sort of wasting time moving
memory around when you don't need to so
pmap gives us a lot of control over how
we're sort of sharding data versus the
actual compute
and
it gives us actually more than that as
well it gives us a bunch of collective
operators so i'll go through this one
pretty quickly it doesn't you know
there's a couple of details here we
don't care too much about but if i'm
doing some machine learning over a large
number of machines so again this this
huge cluster of 2000 machines
i might have some sort of loss function
where you know i'm passing in parameters
and i've got my x and y
reintroducing some of the concepts we've
talked about i might want to calculate
gradients
for these parameters with respect to
this input and output and how would i do
that well i can take my loss function i
can ask
jax to give us the gradients
and i can ask jax to please make that
fast in whatever way is required and at
this point if i've pmped this function
at this point each of the devices
if i've passed a shard of x and y would
be looking at a different set of
gradients so
the thing that pmap really is about is
this parallelization so this piece would
all be running in parallel across the
machines and each device as long as the
x and y was separate would be
like gradients for a different chunk of
my data
now very in a very very simple way jax
provides us other operators which work
across the mesh in this case p sum so by
simply running this p sum here uh what
what what jax would do in this case is
across all the devices actually collect
the gradients together and sum them up
so if each of if i had 2000 machines and
each of them had individual gradients at
this point if i collect them together i
would actually sum all those gradients
up and depending on how i wanted to use
those early is p sum or p mean
and behind the scenes jax is doing all
the work to have all these machines
interchange and to sum and to put things
together and
and we just we just get it magically
so this is huge because
trying to do this stuff yourself is is a
complete pain and requires so much code
but but jax just makes it absolutely
trivial with these type of operators and
it just scales from the tpu with the v38
up to these tpus uh pod slices with
literally thousands of cores without any
code change it's quite astounding
and you know pmap and vmap sometimes
feel like um they're quite specific and
jax is a very much a moving target
i love this this is the the next version
of pmap potentially which is called xmap
and i love this this stock string is
aspirational it's just like
jax really is bleeding edge some of this
stuff
but you know they're providing more and
more ways to make use of this crazy
compute and to give hints around how you
want to actually
make make use of these meshes it's quite
astounding
okay
so everything we talked about here
is quite low level we talked about like
um you know there's vectorization and
parallelization it's fun stuff to do but
it's nothing to do with how we usually
do things so a lot of people ask me you
know do you like jax or keras better and
they're a little bit um sort of it's a
bit of a mismatch you can't really make
this comparison
and for me keras is you know two two
fundamental things it's a um
it's it's a way of describing model
definitions so you know really rich way
of um assembling what they call layers
into models and these sort of you know
optimizers that would go with these
layers and models so that's a really
handy thing keras has
and another thing that keras is very
good at is having an opinionated fit
loop so this idea about given some data
and an optimizer and a model just to
sort of run in batches over it and do
things with with control from callbacks
so keras really has those two things and
um none of those are anything to do with
jax jax couldn't care less about fit
loops it couldn't care less about
optimizers although it does have some
sort of you know reference optimizers in
the experimental
sort of branches and
modules but for fundamentally jacks
around that you know making use of xla
to do this great compute so
if you provide an awesome
library for doing this sort of stuff but
you don't have
you know high level libraries on top
what happens you get the cambrian
explosion of libraries and this this
this we see time and time again so if
you wanted to do anything with jax you
probably don't want to do it directly
but you actually want to do it with one
of you know 50 million different um
libraries that come out all the time and
they have different flavors and sort of
different separations of concerns you
know like our lacks is interesting
reinforcement learning tracks is
interested in transformers jax md is
like you know using these meshes to do
complicated molecular dynamics so really
um
it's hard to decide which which library
you might want to use on top the only
thing you can really do is to explore
and experiment in a couple of different
ways
the the one that i have sort of settled
on myself is the combination of haiku
and optics and i just want to give you a
quick flavor of those
so haiku
has like i mentioned with keras one of
the things that keras gives you is a way
of defining layers into models that's a
really big thing and haiku gives us a
bunch of that stuff so here's the
simplest model i can basically make this
is um y equals mx plus b basically it's
a linear layer that haiku provides and
there's some function which describes um
that model but but this function i've
got here could be all sorts of things
you know multiple convolutions all sorts
of weird self-attention stuff haiku
provides a whole bunch of what we think
about as those layers that i can
assemble in different ways
and then what we get from um and again
just i'll just point out some
interesting things here what what what
haiku is doing is taking these functions
and in the same way jax does
transforming them into other things that
we can use and in particular the main
one around the model definition is that
that haiku will change this function
into a model object which has two
important things for us
one is an initialization uh sort of idea
given this model and and you know
through tracing what does this model do
and what parameters are used by that
model so the first thing we want to
basically do is to initialize this model
to get a sense of what are the actual
parameters that are described inside and
then once we have those parameters we
can apply those parameters and some
input to get a to get an output now you
can certainly see this is very much the
functional way of doing things isn't it
usually when we think about like an
object orientated approach we would sort
of bake these parameters into the model
but haiku continues to use that
functional approach of saying well you
know
it's really about these everything's
with respect to parameters every time we
make these calls
so um if you're more of a functional
sort of person
and you want to use the jack stuff more
directly which is basically about
functional stuff then haiku is maybe
better whereas if you want to try to
hide things more and more and you're not
as interested then there might be
another library that would suit you
better
the other thing that sort of
sits next to haiku which is a different
package is a tool called optics and
optics is a is a
set and a collection of optimizers so in
the same way that
say i wanted to use atom to do some
optimization what i'll do is i'll make
one of these
atom objects and i would initialize it
again with these parameters that have
come from my model and that gives me an
optimizer state because some optimizers
are stateful adam for example is a
stateful optimizer and then it's up to
me again to look after this state and to
do things with it
and in particular here's the classic
update rule we might do
this is i'll go through this slide a bit
more carefully so i've got some sort of
loss function that describes
given a set of x and y and a set of
parameters you know how well am i doing
and the way i would calculate that loss
is to apply my model with the parameters
in x get predicted values and then
calculate some sort of loss
then what i can do is uh use jacks to
get gradients
of that parameter of that function with
respect to parameters which give me
gradients and then what i can do is pass
the gradients to the optimizer
to say basically
given these gradients and your current
state what's your new state and are
there any updates to make to the
parameters and then again given the
parameters and those updates what are
the new set of parameters
so this is a classic sort of functional
approach again where we're passing in
states and getting new states back so
all these things are 100 immutable and
we just sort of have to do this classic
thing of here's a set of parameters and
through this update i get a new set of
parameters
and um we're using all these jax
fundamentals like grad or you know
jitting this entire thing but all these
things that i'm putting in here
depending on how i want to do things can
have a sprinkling of the v map and the p
map depending on how i want to run
across different machines so you
remember in that slide a little while
ago we were talking about this update
being
done with a pmap instead of just a
directly with an update and i think this
idea of being able to introduce these
things as we go along is super powerful
because what it means is that we can um
sort of get things working in a simple
way on a sort of you know single machine
simple uh non-batched sort of operation
and then eventually you know scale
things out more and more and more and
this idea of jitting an update rule and
then basically going through and then
just pumping that update basically
across the data set where we pass in
parameters we get next parameters we
pass in the state we get the next state
and just sort of run that
is very very fast because we're
basically just sort of if these are
sharded across the um the devices then
all the the only thing we're having to
do is really move uh the data around so
this can be very very high performance
at a very high scale
so yeah just a recap of the sort of
things we talked about i sort of talked
about the jax fundamentals that was the
the main thing and again that takeaway
just try jack's numpy if nothing else
and see how you go it might just work
and surprise you but you know
fundamentally there's this there's these
building blocks that jax provides around
um
the classic sort of composition so
gradients i want jet i want to v-map
things
there are some multi-host specific
things that come from the sort of
heritage of why jax even came out around
xla and so there's things like p-map and
p-sump and and you know i've scratching
the surface here in terms of that mesh
computation there are many many
complicated things that we can do here
um i've just sort of dropped into for a
quick talk this idea of parallelization
of map or this this operator p sum
and then you know jax maybe isn't
exactly what you want you want a high
level support and yeah all i can say is
just query what the latest neural
network library for jax is because it
would have changed by the time this talk
is finished it's a very vibrant field
with a lot of things changing and you
know i think it's it's okay in this sort
of fields when things are fast to um try
a couple of different ones and just see
which sort of one uh represents the the
way that you think about the problems um
and you know you might get lucky on one
or you might have to change around a bit
okay so that was basically
uh my introduction there
there's like a two hour version of of
this talk with a bit more tutorial stuff
on my website there but otherwise here's
a bunch of links which correspond to
uh the sort of things i've talked about
so far
thank you
thank you matt and
um
so we have time for one question and
it's going to be
a quick one maybe you have a quick
answer so you've mentioned tpus but most
people have gpus and no gpus what kind
of problems are better trained on tpus
like they would require you to go the
tpu way and would your jacks go change
for that
yeah so the jax code would not change at
all so that they've been very clear that
xla
supports
gpus as much as tpus and so all the
things i talked about pmap psum they
work on my my desktop here with a couple
of gpus without any change at all
the tpus are
yeah the tpus i think are more sorted uh
in terms of the benchmarks i've seen the
tpus are the best performance when
you're really starting to get into the
large data regime so um
generally the way they work is that like
like a lot of like high performance
accelerators they work well once they've
got a real chance to spin up so if
you're working on a small bit of data
you might find that
jax on tpu's isn't actually as
performant as some other things um and
it's not until you get to these bigger
data sets that um
uh it really makes a difference
particularly batch size is a big one
thank you
so
uh thank you so much matt for being part
of the data science and analytics track
today yeah no worries thank you
and coming up at 11 we'll have jordika
singh
who will be talking to us about
classifying audio into types using
python see you then
bye