← 返回日报
精读 预计 3 分钟

PyTorch: a reference language

摘要

文章提出将 PyTorch 视为参考语言与实现语言的双重角色:参考实现用于清晰验证,生产实现可另行优化。作者以内核 DSL 为例,说明保留纯 PyTorch 参考实现并配合数值测试验证优化内核的常见做法,进而延伸到训练步骤:用 LLM 生成显式前向 - 反向代码,与 autograd 参考实现并行,并通过位级等价或图结构等价验证器保证两者一致,从而避免模式匹配的脆弱性。最后指出这是对 Horace He 关于兼顾 eager 控制与图抽象便利问题的有前景回答。

荐读理由

把 PyTorch 当参考实现、用 LLM 生成显式前向-反向代码再配验证器保证等价,这套配方可直接搬进你的训练栈,省掉手写反向的调试苦工

原文

PyTorch: a reference language

Edward Z. Yang (@ezyang) · July 25, 2026 · 4 min read compilertorch.compileautogradverificationllm

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.

Lobsters · 1 赞 · 0 评 讨论 → 阅读原文 →

这条对你有帮助吗?