Every programming language you need to know for GPU-powered quant research

Machine learning in quant research is nothing new, but as the need for ultra-low speeds increases and the power of graphics processing units (GPUs) increases, optimizing them to supercharge the speeds of your models becomes a necessity.

But which language is the best for this?

Click here to follow our new WhatsApp channel, and get instant news updates straight to your phone 📱.

In a Jane Street Signals and Threads podcast earlier this month, Sylvian Gugger, a machine learning researcher at the firm, described a number of languages coders use to work with GPUs, and their inherent benefits and problems. 

These are as follows - 

CUDA

Cuda is, in one way or another, unavoidable. It's a proprietary language developed by NVIDIA for its GPUs, but is also interoperable with GPUs at rival firms like AMD.

Gugger likens CUDA to C, for better and for worse. It will be your best bet at fine-tuned ultra-fast speeds, but is "pretty hard to debug," and laden with "undefined behaviours every time you're not super careful."

In CUDA, you write functions called kernels for the GPU. Gugger says Kernels can be hard to write because they're "highly parallel and hard to reason about." Kernels won't consist of endless lines of code, but they will be "numerically intense."

This is the closest to the hardware you'll get when working with CPUs, but it's also not particularly enjoyable to code with. Therefore, a number of alternatives have been created to simplify the process like...

Python/PyTorch

Python is the opposite side of the coin to CUDA. It's a lot further away from the machine, but it's much, much easier to code in. 

PyTorch, first developed by Meta, uses GPUs to boost the performance of calculations powered by NumPy one of the most popular libraries available in the language. Gugger says this is what quants at Jane Street use, because it's "really accessible."

The problem is the exact opposite of CUDA. Gugger says, in PyTorch "you can do anything you want, but then your code can be inefficient in surprising ways" 

For Jane Street traders, this setup works, because it allows them to "iterate really fast through their ideas," and will only need to do the fine-tuning after they've seen that a certain model is lucrative.

Tensorflow/JAX

PyTorch isn't the only machine learning library in Python. Two of its most widely used competitors are Tensorflow and JAX. Both languages were developed by Google, and share a key difference to PyTorch.

Gugger says Tensorflow and JAX operate "more like compiled languages" than PyTorch. This is an immediate barrier to entry; the program will not run unless the code is properly compiled, which means you can't get away with the mistakes you made in PyTorch.

The learning curve isn't the only problem for Tensorflow and JAX. When deciding whether code should be executed on the CPU or GPU, Gugger says Tensorflow will "take care of everything for you because it has compiled everything," which can lead to inefficiencies. In PyTorch, you have more control over where your lines of code are executed.

As a basic idea, however, Jane Street CTO Ron Minsky prefers those libraries to PyTorch. He says they're essentially a "meta programming system," writing "what looks like python", but is actually a lot more complex. Making code run fast in Tensorflow should be "dramatically easier" than in PyTorch.

Triton/Mojo

With the boom of AI in recent years, there have been a few new entrants to the field offering a variety of benefits.

One Gugger highlights is Triton, a superset of Python that "can directly write efficient CUDA kernels" without requiring CUDA expertise. The language was developed OpenAI member of technical staff Philippe Tillet around the time he joined as a graduate in 2020.

Gugger says, for matrix multiplication, Triton is "relatively easy" to use with some "crazy templates for basically all operations you can do in a PyTorch model." Much like CUDA, however, it has a debugging problem. Gugger says "very often Triton won’t compile your code and won’t necessarily tell you why." Minsky notes that it's "not quite general purpose" yet, despite showing initial promise.

Another alternative we've previously covered is Mojo, another Python superset. Gugger says its "a little bit Rust inspired” and has more support for debugging. However, Mojo's designers say the language is designed "particularly for weird domains that aren’t traditional CPUs and GPUs," like FPGAs.

Have a confidential story, tip, or comment you’d like to share? Contact: Telegram: @AlexMcMurray, WhatsApp: (+1 269 237 3950). Click here to fill in our anonymous form, or email editortips@efinancialcareers.com.

Bear with us if you leave a comment at the bottom of this article: all our comments are moderated by human beings. Sometimes these humans might be asleep, or away from their desks, so it may take a while for your comment to appear. Eventually it will – unless it’s offensive or libelous (in which case it won’t.)