PyTorch: A Reference Language

A reference implementation is a simplified but complete version of a system that trades performance in return for clarity. We might then say a reference “language” is the fabric of APIs and conventions from which these implementations are cut. At first glance, PyTorch obviously is a reference language: it is, after all, commonly called the lingua franca of modern deep learning. But upon a closer look, there is confusion:

Reference implementations usually aren’t deployed to production. But I do my training jobs with PyTorch!

Everyone’s writing kernels with kernel DSLs. What is the role of PyTorch if it’s just gluing kernels together?

AI coding will eventually mean that any stack can be rewritten from scratch; why does being written in PyTorch matter?

So for me, recently, an unusually clarifying perspective has been to think of PyTorch as playing a dual role: as both the reference language and the implementation language. When the scale is not too large or the compiler is working well, the reference implementation can ship to production. But increasingly, I think it will be more and more natural to think of the reference implementation as a software artifact that stands apart from the actual production implementation, by which we can verify the correctness of the production implementation. One implementation to research in, one implementation to scale with, and one verifier to, in the darkness, bind them.

The clearest demonstration of this is in the modern usage of kernel DSLs. The traditional, compiler-maximalist view argues that end users should write implementations of NN modules using a high level API (e.g., a Numpy/PyTorch-style API) which a compiler then determines how to compile into an optimized form. But for the most important operations like matrix multiplies and attention, it is not easy for compilers to guarantee peak performance; the proliferation of kernel DSLs has made it dramatically simpler for people to achieve optimal performance by explicitly spelling out tiling and data movement. Does this eliminate the high level API? Usually not: it’s pretty useful to have a reference implementation in plain PyTorch, and most kernel authors will maintain one in parallel with the optimized kernel, verifying correctness with numerical tests.

In the same way kernel DSLs have changed how production implementations of operators can be written, I think coding agents change the way production implementations of train steps can be written. Traditionally, we think of autograd as a core part of PyTorch’s value proposition, because it guarantees you will get correct derivatives. However, at scale, the implicit backwards graph becomes an albatross around one’s neck: the majority of your compute is hidden away, with no opportunity to interact with it with normal debugging tools or apply fusions to it in the same way you can do it in eager forwards code. With a compiler, it is possible to modify the backwards graph with, e.g., a pattern match, but this is brittle and a less nice experience than just swapping a call from a reference implementation to a hand-written kernel. This is not a new observation: the now defunct Tangent library was built on the proposition that source-to-source automatic differentiation could be useful.

The new recipe looks like this. Keep the traditional PyTorch autograd-friendly code as the reference implementation. Use LLMs to generate an explicit forward-backward version of the code, which can be optimized separately from the reference implementation. Unlike pattern matching, you never have to worry about your optimizations failing to apply. The cost is that the reference and the real implementation can diverge: we need a verifier that shows us they are equivalent. This verifier can be implemented simply with a bitwise equivalence test, or implemented as some sort of graph capture and structural equivalence, in the tradition of translation validation. To ensure the verifier works when one side has a fusion the other doesn’t, you only need to provide a reference implementation of the fusion (an inverse pattern match, if you will!).

I am not going to claim that this recipe is right for everyone. It turns out PyTorch, the reference language, is a pretty good executable spec, and at the end of the day what really matters is how quickly you get the experimental results you need. But, having spent a lot of my time recently thinking about what it means for PyTorch to excel at frontier training–and in particular whether or not it is necessary for PyTorch to disrupt itself as scaling continues–I feel that this perspective helps bridge the old and the new. An open question Horace He posed last year was this: “How can we get all of the control of eager-mode execution with some of the conveniences of graph-level abstraction?” I think this recipe is a pretty promising answer, and PyTorch continues to be at the center of it.