Welcome to the fifth NeuroMechFly video tutorial. In the previous tutorial, or rather video number
four, we followed tutorial notebook number two, where we replayed some experimental data of a fruit
fly walking system that we call Spotlight.
So we use this kind of fly behavior data data
and replayed it in FlyGym.
The outcome of that notebook was that we could generate videos of a fly walking
in a virtual environment.
And we could also quantitatively analyze variables such as joint angles during the simulation
and compare it with experimental data.
In this tutorial, we'll basically do the same thing, but this time using GPU accelerated simulation.
The current FlyGym version at the time of this recording is 2.1.0. Future versions
might deviate from this tutorial, so it's recommended that you always go to neuromechfly.org and
follow the up-to-date tutorial and only using this video as a reference.
GPU accelerated simulation in FlyGym is enabled by MuJoCo Warp,
a project developed jointly by NVIDIA
and Google DeepMind.
You can find more information about MuJoCo Warp on its project website.
MuJoCo Warp is further based on NVIDIA Warp,
which is a toolbox for all sorts of physical simulations,
not just for biomechanical robotics simulations.
You can find more information about it also on its project website.
I'm not going to go into a lot of details of GPU accelerated computation, but I
do want to refer the viewer to this project called Accelerated Computing Hub by NVIDIA, specifically
this, I guess, sub-project called Accelerated Python User Guide,
which is a crash course on GPU accelerated computation using Python.
I think the most relevant tutorial is chapter number 12, which is specific to NVIDIA
Warp. Before I start, maybe I can just give a brief explanation of the
difference between CPU computing and GPU computing.
For that, I'm going to use the pictures from the Cornell virtual workshop on
Understanding GPU Architecture.
Here what you're seeing is a schematic overview of a CPU and a GPU.
Each green rectangle refers to a computing core.
It is the part of, it is the collection of hardware that actually executes your code.
On the CPU, each core executes your code line by line.
It can do a lot of different things, complicated things,
but they do execute your code sequentially, one line after another.
It is true that you can have several core, if not dozens of cores on a
modern computer, and you can do things in parallel,
but still that number of cores is rather small
compared to the level of parallelism on a GPU.
Here you can see that there are hundreds, thousands, if not more cores on GPUs,
but the caveat is that each core can, in terms of the types of computation that
can be performed on the core, it is a little bit more simplistic
compared to a modern CPU core.
Still, the advantage of a GPU is that because you have so many cores, you can
do a lot of things in parallel,
and if the thing that you want to do is not that complicated, for example, if
you want to multiply two matrices,
then being able to utilize this parallelism of the
GPU the GPU accelerates your computation by often
orders of magnitude,
and this is the reason why for deep
learning inference and training often requires powerful GPUs, because
to the core it's just a bunch of matrix multiplications.
Now, going back to the FlyGym tutorial on GPU accelerated simulation, it is possible to run
this tutorial on Google Colab without installing FlyGym locally.
If you do that, you have to go to runtime, change runtime type,
and select GPU runtime.
Please be aware that Google Colab can be very slow compared to if you run it
locally on your own computer,
and there will be quite a bit of simulation going on, so I'm not going to
be using Colab in this tutorial,
and instead I'm going to connect to a remote workstation by SSH and run this tutorial from there,
and on that tutorial from there, and on that computer I have an NVIDIA GPU,
so of course that is a prerequisite for running this tutorial of GPU accelerated simulation.
As before, we run this kind of empty block for installing FlyGym that's ignored unless
you're running it on Colab.
With FlyGym warp, we have provided this utility function that simply checks
what kind of GPU you have.
In this case, I have a NVIDIA GeForce RTX 3080 [Ti],
which is not that new, but
it's a quite reasonable GPU.
As before, we will compose a fly model that is attached to the world.
We'll add a tracking camera to the model, and I'll skip this part because it's been
covered in previous tutorials.
Now, as I said, the advantage of running a GPU-based simulation
is that we can parallelize.
In other words, we'll be simulating more than one world at a time.
As a demonstration, let's say we want to simulate 100 different worlds in parallel.
So, different from the previous tutorial on CPU simulation,
we will create a GPU simulation class object
instead of just simulation.
And when you create this GPU simulation object, you want to specify both the world that
you want to simulate and also the number of duplicates that you want to
simulate in parallel.
As before, we can add a renderer to the GPU simulation,
but since we are simulating 100 worlds, if we had rendered the video for all of
them, we would try to quickly run out of memory.
So, we're just going to specify a list of world IDs, in this case, the first
nine of them, and render only those.
So, I'm going to run those two code blocks.
We have already covered in the previous tutorial that we have shipped this
flygym_demo.spotlight[_data].MotionSnippet helper class,
so we're just going to load some experimental data,
and let's say we have a 20,000 step data set, behavior data set,
we're running the simulation at one millisecond step,
in other words, we have two seconds, exactly two seconds of real experimental data,
and since we have 100 worlds, let's say we want to split this two-second behavior
recording into 0.1 second snippets with some repetition.
In reality, you don't want to do this, but perhaps you want to split your, I
don't know, five-minute experiment into 100 small chunks or something like that,
but here in order to be able to finish running the simulation fast,
we'll just run 0.1 second per world,
and here we are just running this little helper function that generates 100 different sequences of
motion with 0.1 second offsets.
So, now we can expect that this array should have a shape of
100 by 1000 by 42,
100 because we have 100 worlds, 1000 because we are simulating at 10,000 Hz
and 0.1 second,
and 42 because we have 42 degrees of freedom, 7 per leg, 6 legs.
As before, we can set up this simulation loop,
it's pretty straightforward, very similar to what we wrote before.
We reset the simulation,
we enable leg alternation,
we warm up the simulation a bit, and then we can start running the simulation loop.
As before, we first batch the kind of control inputs or target join angles that we
want for this step
because we have 100 different worlds. We here have actually a 2D array
where the first dimension is the 100 different worlds and the last dimension is the 42
degrees of freedom.
So, in the GPU case, when we call set actuator inputs on the simulation object,
we still need to specify the name of the fly, the fact that we're setting the
inputs for the position controllers,
and then this 100 by 42 instead of just 42 dimensional control inputs. And
we can call step and render as needed with this profile option.
If we don't run it, it's a little bit faster because
it doesn't time the different operations,
but if you enable with profile suffix, then when you run the simulation, you get some
statistical data of how fast the simulation is.
Now, if we run it, we quickly realize that it's actually rather slow.
First off, there's a bunch of things happening,
and then the simulation took seven seconds, which is actually slower than our CPU simulation.
There are a couple of reasons for this. First, Warp relies on something called
just-in-time compilation, or JIT.
The idea is that when you run something the first time, it spends a bit extra
time to compile your operations into efficient machine code or binary code
so that in subsequent executions, it becomes much faster because you don't need
to interpret Python code line by line.
Now, also other reasons, there are many things that we could optimize,
but for now, let's run the GPU simulation very naively like so and
print a report.
So here, we'll see that our overall throughput is about 1.44 times real-time
because we simulated 0.1 second for 100 different worlds, so it's a total of 10
seconds simulated and that completed in seven seconds.
Makes sense. About half of the time is spent on physics simulation and half of the
time is spent on rendering.
Note that we are only rendering for 10 worlds, so it's actually quite significant that we
spent so much time rendering.
As before, we can display the video that has been captured.
Sorry, I said 10 worlds, but actually we only rendered the video for nine worlds,
and here they are. You can see that the simulation is a bit different because they
are different 0.1 second snippets.
So, all right. Next, instead of simulating 100 worlds in parallel,
what if we simulate a thousand
worlds in parallel, right?
Because naively, if we had done this on the CPU, we would have expected that it
would have taken seven times 10th, what, 70 seconds for the simulation to complete.
Now, is that really the case?
We first define a function that takes as argument the number of worlds and then
just runs the whole thing that we have just done,
and we call this
okay, spend some time compiling or building up the simulation, and then it runs.
And there we go. When we increase the number of worlds by 10, the amount of
time that it took only increased, only doubled,
and that, again, is due to the fact that we have many, many cores.
Perhaps when we are running only 100 simulations, we're only using this many cores,
but if we use this many cores and we increase it by a factor of 10,
then we use this many cores,
but still, a lot of them are still idle and there's not really a real cost
in terms of execution time to use more cores.
Of course, there's still overhead and it took more than seven seconds, but it doesn't scale
linearly with the number of worlds that you simulate in parallel.
If we look at the statistics, the overall throughput for simulation has increased from 1.6
times real time to 12 times real time.
A hundred seconds of data simulated in 14 plus some rendering.
This is the advantage to the first order of approximation of running simulations
on the GPU because you can simply parallelize.
Now, something that we didn't parallelize in this example is rendering.
If you look into how this render is implemented, you'll find out that we're still
sequentially capturing frames for those nine worlds that are being visualized.
Can we do the visualization with the rendering also on the GPU?
Yes, of course.
In fact, the only thing that you need to change is to add this use GPU
batch rendering equals true flag when you set the renderer.
Because everything is parallelized, we don't really care if we're rendering every world.
You can set the selection of worlds to render to none, in which case all
worlds are rendered.
Everything else is kind of the same. Let's run this.
That defines the function. Let's actually run it for a hundred different worlds.
Previously, for a hundred worlds, the rendering throughput is exactly one
time real time.
But when we use GPU rendering, it's about 50 times real time.
That's impressive. You'll see that the quality of rendering is not as good because the
GPU renderer has been optimized for performance instead of realism.
I think for most scientific applications, this is okay.
I hope you can appreciate that our simulation is now a lot faster on the GPU.
Can we optimize it even further? Yes.
But maybe a bit of warning is that
doing so requires a somewhat significant rewrite of the
code because ultimately GPU programming is quite different from CPU programming on a rather fundamental level.
So far, we have minimized the amount of code change.
Mainly, we just had to replace the CPU simulation class with the GPU simulation class.
But now we want to go deeper and write the simulation loop from scratch.
Before that, I want to reiterate that the real advantage of GPU computing really comes from
the parallelism.
You have so many cores, you can have many, many threads.
So hypothetically, let's take a toy example where you have a list of numbers
between one and one thousand
and you want to take the square of each number.
Essentially, what you want to do on the CPU is to run this operation a thousand
times sequentially and modify the elements one at a time.
Of course, in reality, you use NumPy or something like that, but conceptually, this is what
happens.
Now, on the GPU, what you want conceptually is a function that handles only a small
piece of your data
and then you initiate many, many invocations of this function in parallel.
And together, they process the entire data.
So let's write this function. Let's call it square
that takes an input array and an output array.
And in pseudocode, what we want to do is step one figure out which part of
the input array I'm responsible for.
And the second step would be to compute the square of only the element that I'm
assigned to.
Assigned is a bad word, let's say I'm responsible for.
And maybe step three is to store the result.
If we want to convert this to real code, first off, let's import warp.
What we want to do here is to let's just call this x,
to tell warp that you have an input which
is an integer array and y which is also an integer array.
And for step one, we can use something called a function called warp TID or thread
ID, which tells this instance which index you have been assigned to.
And in the second step, you do this, but I think in warp you actually have
to write it out like that.
And in step three, you assign this number to the same index of y, your output.
And because this is a GPU operation, we have to use this warp kernel decorator.
A kernel is basically the GPU equivalent of a function.
Importantly, you actually do have to specify the data type and array format of your data.
In Python, typically when you have this kind of type-ins, they're ignored.
It's only for documentation, but in the case of warp, it is important.
Okay, now that this operation is defined, we just need to tell the GPU to actually
launch this operation.
So for that, we use warp launch.
We tell the GPU that I want you to run this.
And if we go back a bit, we have used the TID function to figure out
the index, right? The index on the data.
But we haven't told the GPU how many of those threads it should create yet or
how many elements there are in total.
So now let's do that. It's just the length of...
Actually, before we start, let's do input array.
Sure, let's basically create a warp array version of this Python native list.
And we observe that these arrays are on the GPU.
It's on the GPU.
So now we tell the GPU to launch this kernel, and we need to tell it
that there are in total 1,000 threads that should be launched.
So this, or rather, let's do it better. Like this.
And we can say the input array is the input array and output array is the
output array.
Those are just simply concatenated and expanded and interpreted as those two things.
Only by conversion in warp, it's always input first and output last.
So the arguments of this kernel is simply the concatenation of those two lists.
So we tell the GPU to start running that, and it's done. That's great.
If we look at the output array, we can see that...
Let's copy it to the CPU and operate it as a NumPy array.
We can see that, indeed, we get our numbers.
Good.
Now, how fast is this?
Let's try to run this same operation a million times.
Eight seconds.
It's not great.
One reason of this rather slow execution throughput is that we're ultimately still running it
in a Python loop.
No matter how efficient this line of code is, we're doing something a million times in
a Python loop.
This is illustrated very well in this figure on the NVIDIA developer technical blog.
What we are doing here essentially is to launch this kernel a million times.
Each time we do it, there's an overhead.
It does launch this computation task, but with a lot of overhead, and this is repeated
literally a million times.
What we can do instead is to use something called a graph. So
on the GPU, there's a concept of graph, which is basically that for everything that you
want to do on your data, every operation, so
for example, multiplication, addition, you can represent it as a node, and you can chain them
together sequentially and form what's called a directed acyclic graph containing everything you want to do
with it.
You can tell the GPU that this is what you want to do ahead of time,
and then it can run a bunch of optimization to do it more efficiently.
In the end, what happens is that you spend some time building the graph, telling the
GPU what you want to do,
and then you launch the graph once, and once you launch it, you just have everything
that should happen happening without this launching overhead.
So let's try that.
You can create a graph capture using this ScopedCapture context,
let's call it capture,
and under this Python context, we do exactly what we want to do, but this time
we don't actually, the GPU actually just, it doesn't really run it, it just remembers this
is what you want to do, just records the set of operations,
and it takes a similar amount of time because it still needs to run this whole
Python loop.
But then once it's recorded, you can launch it once with a capture launch.
You just tell it capture, this is what I want you to do. Oh,
sorry, actually the capture is just an object reflecting this capture that you did. The actual
graph is capture graph,
so let's launch this whole graph,
and now it takes three seconds instead of eight.
So now you might ask, why are we doing eight plus three instead of
just sticking with eight?
Well, the answer is the next time you run it, it's going to take basically no
time.
Warp uses a technique called just-in-time compilation, which basically says that every
time you do something, it just-in-time compiles it to very efficient binary machine code,
and the next time you run it, it doesn't have to go through the code line
by line anymore,
you can just use the pre-compiled machine code.
So if you're running the same thing many, many times, for example, applying the
same operation on a simulated fly or simulating many behavioral recordings, then it's quite advantageous because
the execution time no longer scales quite the same way as having a naive loop.
Now, going back to our NeuroMechFly simulation, what we need to do here is
to define a kernel that takes as input the buffer of all this whole three-dimensional
array of shape and number of worlds, number of steps,
and number of degrees of freedom, and then the step count, the current step in the
physics simulation.
And finally, you have this two-dimensional array that just reflects the current target angles for
the physics simulation.
So just get rid of this.
And what we want to do here, essentially, if we had written it as a CPU
function, it would be something like this.
Well, it would be something like this.
We would be slicing this history of behavior recording and then step
index and every degree of freedom. And then we would be
slicing this history of behavior recording and then step index and every
degree of freedom, and just assigning this, copying this to this buffer.
This is literally what we would have done if we had implemented this with a CPU,
but because we're running on the GPU, this is what we have to have.
We have to use this TID function to now figure out which world and which degree
of freedom I'm responsible for.
Get these two numbers. We get the step counter, which is also a warp array.
If it's on the GPU, it has to be placed in an array, except in this
array, there's only a single number, which is the step counter, which we can access
by indexing zero.
And we fetch the number that we're interested in and literally copy it to this target
position, this target buffer.
Define this kernel. Similarly, even just for incrementing the step counter,
we have to write a kernel for it.
It's literally just taking the only number of this array and adding one to it,
but still it's a kernel.
And now we can define this function that runs the simulation with graph capture.
So as before, we have a ScopedCapture,
but first we define the fly, the simulation, blah, blah, blah.
And then we define a graph capture reflecting what we need to do for each simulation
step.
So the first thing we need to do is to launch this kernel that copies data
to the buffer for the target position buffer for the position actuators.
And then we need to tell FlyGym that please use this as the actuator
inputs.
You don't need to launch a warp kernel for this because internally FlyGym is
already implemented using kernel.
And then as before, you call the step function.
Again, you don't need to launch it as a, you don't need to call launch kernel
explicitly because your fly is internally already doing that.
Internally using MuJoCo Warp and MuJoCo Warp is already doing that.
So it takes some time to build this graph.
And once this graph is built, you as before have a Python loop
and within the Python loop, you can just launch this captured graph of everything you need
to do per iteration.
Oh, let's not forget after stepping physics, you have to call this very small kernel of
incrementing a step counter by one.
We render as requested and we report the performance.
And now let's run it for a thousand worlds with rendering enabled.
They're running out of memory. Sorry.
Let's restart this to clear the memory.
Let me call this again.
Okay. And now we have a performance of 10x real-time including rendering compared to the
previous throughput of 1.53.
This is nice.
And this is further compared to the previous uh previous throughput for the rendering of 0.1.
And now rendering is a lot faster, but not necessarily overall.
And if we use the graph capture, then everything is fast,
even with rendering is 10x real-time.
If we disable rendering, then it's even faster.
30x real-time on my NVIDIA GeForce 3080 Ti GPU,
which has been five or six years old.
We have actually already benchmarked the performance on different GPUs.
On the x-axis, you have the overall throughput without rendering.
And on the, sorry, on the y-axis, you have the throughput.
And on the x-axis, you have the number of worlds that are simulated in parallel.
Note that it's on a log scale.
So now we are roughly here,
about a thousand worlds on the 3080. In
our example, it's actually around 30.
We see that we get the overall, the highest throughput when we are simulating about 4,000
to 8,000 worlds in parallel.
And you don't actually need such fancy GPUs because even with a 4090, which is fancy
but not that fancy,
you already have the same performance basically as the H100.
And that's because the type of computation in physics simulation is just a bit different from
the types of computation in, for example, training an AI model.
So I encourage you to benchmark the simulation before committing to running large jobs.
And in the next tutorial, we will spend more time covering how to benchmark or rather
how to profile your simulation to understand which part is taking a lot of time.
One thing to note is that depending on exactly what you want to do, you might
even be able to capture the computation outside of this Python loop.
Because now what we have here is a little bit like the example here.
We have the computation launched once only per iteration of the loop,
but overall we're still running this native Python loop.
So alternatively, you can do this capture.
You can run this loop inside the capture and launch the whole thing once,
but that is outside the scope of this present tutorial.