Deep Dive2 min read

Quick Note: π0.5 Inference in 16 GB

With 24 GB laptop GPUs getting harder to find, we trimmed our π0.5 rollout stack from about 22 GB to 11.2 GB of VRAM. It now fits 16 GB cards like the RTX 5080 and 5070 Ti.

Cream contour lines, cover art for Quick Note: π0.5 Inference in 16 GB

Introduction

It’s no secret that there is a bit of a supply chain crisis in all things compute right now. While datacenter GPUs remain in plentiful, if exorbitantly-priced, supply, availability and pricing for everything edge is spotty. If we look at the past six months, we’ve seen the effective disappearance of AGX Orin from the market, 50% price increases on everything, and, most disconcertingly, the quiet discontinuation of the RTX 5090.

We’ve also seen some alarming supply reduction on other edge hardware. In particular, the RTX 5090 Mobile, which we’ve really enjoyed using for rollout in the past, seems to be getting rarer every week. The 5090 Mobile is a bit of an odd one out – it’s GB203-based, just like the RTX 5080 Mobile and desktop RTX 5070 Ti/5080, but comes with 3 GB GDDR7 packages on a 256-bit bus. This was a great deal under the old pricing schemes, where its nearest desktop equivalent was the RTX 4000 Pro, a similar 150 W/256-bit card, but for slightly more than the price of the bare GPU one could purchase a complete, integrated, portable system with ample connectivity and an internal battery.

VRAM Shaving

With the 24 GB mobile SKU seemingly on its way out and the workstation SKU skyrocketing to north of $3000, it seemed prudent to prepare a π0.5 stack which could be crammed into 16 GB. Our previous work topped out at about 22 GB during model loading, but since the model contained about 7 GB of weights we were hopeful that the activations would fit in the remaining 9 GB.

Fortunately, this proved to be straightforward. Recall we had previously mentioned that off-the-shelf ML code was full of stupidity; in this case, we found some stupidity in our own code. It turned out that (1) the Hugging Face libraries were initializing a large, random fp32 model backbone before loading and quantizing the real model and (2) we were holding PyTorch models in VRAM simultaneously with the compiled TensorRT engines. Fixing these two “quirks” (they aren’t bugs, we swear!) dropped VRAM usage by a whopping 10 GB to 11.2 GB.

Bar chart of peak VRAM for the π0.5 rollout stack: about 22 GB before, 11.2 GB after two fixes, shown against a 16 GB line for the RTX 5080 and 5070 Ti and a 24 GB line for the RTX 5090 Mobile.

The New Target Hardware

With 16 GB now in sight the model will run on mobile RTX 5080. We estimate performance to be a bit lower (we haven’t benchmarked this configuration yet), probably around 100 ms per step – but because end-to-end latency includes camera latency, the absolute change will probably be around 10%.

Importantly, the model will now fit on desktop 5070 Ti/5080 machines; these will perform a bit faster than even the 5090 Mobile because of much higher power limits and availability for the 5080 seems much better right now than availability for high-end gaming laptops (which are being impacted by shortages of DDR5 and storage as well). It would probably even fit in a 5070, but at some point you really should stop giving up performance to save money.

As usual, you can find the updated code on GitHub at PrismaXAI/prism-pi.

Robotics Data

Train on the data we write about.