Recent content on ezyang's blog
In Code review as human alignment, in the era of LLMs , I talked about how we should approach code review with other team members, where you have a pre-existing relationship and care about building human-to-human alignment over time. However, popular open source projects also need to account for drive-by contributions from external contributors, where it is unclear if you want to invest in a long…
How does PyTorch autograd deal with mutation? In particular, what happens when a mutation occurs on a view, which aliases with some other tensor? In 2017, Sam Gross implemented support for in-place operations on views , but the details of which have never been described in plain English… until now. Autograd, the simple case ¶ We all know that autograd operates by constructing a backwards…
I was coaching some juniors on effective use of AI coding agents (like, looking over their shoulder as they were prompting the LLMs) and one reoccurring theme was that the AI agents were demanding a lot more of their reading code skills, compared to how they used to operate before. “How,” they asked, “can we get better at reading code?” I had to think about my answer to…
Previously in AI-assisted programming for spmd_types , I mentioned that I have been enjoying using Sapling (the version control system Meta uses internally) to manage parallel agents on worktrees. In this post, I want to describe this workflow in more detail. The basic development flow here is I have three work streams going in parallel: I am working on enabling an end-to-end use case; in this…
spmd_types is a type checker for distributed PyTorch programs. I’ll write a blog post about the system itself as it gets closer to completion, but in this post I want to talk about my workflow for AI-assisted programming in the project. In particular, spmd_types is my first greenfield project that both: I have written it entirely with AI, and I am comfortable being accountable for it (in the…
A silent BC breaking change is a change of semantics in an API that is not immediately obvious; e.g., the old usages of the broken behavior won’t raise a compile or runtime error on upgrade. If you are a project that cares about backwards compatibility, it’s generally a bad idea to ship silent BC breaking changes, because they manifest as users getting silently incorrect results and…
I recently received the following question about vibe-coding for tlparse , a structured log parser for torch.compile (slightly edited for ease of reading): Hi Ed, I have some thoughts on vibe-coding tlparse and I’m curious what your opinions are. I think it’s fairly easy to vibe code and improve tlparse, but it’s hard to find people who know Rust or HTML or JavaScript well to…
A central thesis of sharding in types is that the backward sharding can be directly computed from the forward sharding. This is not true for DTensor today, e.g., as seen in sum to Partial , and it leads to confusion where users cannot easily predict what the sharding of tensors are in their program. The question now arises: given a forward sharding, what should its backward sharding be? There are…
DTensor has famously terrible eager mode performance; for example, this paper measured a 35-60% slowdown in end-to-end training performance with and without DTensor (with DTensor operations taking at least 7x longer than actually running the computation for real). While it is possible to alleviate some of this slowdown via optimizations (in the paper, veScale shows fast bypass of sharding…
Conventionally, a type system is something that classifies values into data types like float32 or int64 . However, fancy type systems go beyond data types, allowing us to talk about potentially arbitrary invariants on data; for example, if we were to talk about the “type” of a array, it would cover not only its data type, but also its shape, e.g., f32[40, 20] . JAX’s type system…
Global SPMD (also known as the “global view”, exposed by code using DTensor or jax.Array ) refers to writing multi-device code as if it was on a single device, with an orthogonal mechanism for expressing how these full tensors are distributed over multiple devices (this mechanism can be implicit or explicit, e.g., as seen in this table ). Local SPMD (also known as “per-device…
In Computing sharding with einsum , we worked an example of Megatron style tensor parallelism where we discover that the ordinary backwards formula for linear results in a pending reduction on grad_input , even though the input was replicated and no communications happened in forwards. In Megatron, which is implemented with plain Tensors and manual collectives, you just have to know that this…
Mental arithmetic in grade school (e.g., memorizing your times tables) is typically justified on the grounds that facility in basic calculations makes it easier to focus on higher-level problems that require being able to do these manipulations. When working on DTensor, I have also found it important to be able to quickly calculate what shardings you get when you do matrix multiplies on sharded…
This blog has lived on WordPress since it was initially created during a social challenge at MIT to write a blog post a week or pay up with beer. I remember a very important piece of advice I had been given at that time: don’t fuck around with your blog authoring software, just do the minimum viable thing (use WordPress) and focus on writing posts. It’s 2026 now, the world is…
Let’s suppose you asked an AI coding agent to “implement a CLI calculator”. Imagine if, instead of only writing short Python script, it also started building an automated test suite, a crash reporting mechanism and a telemetry subsystem. You’d be like, “What the fuck is this?” But now let’s say that you were planning to release this project to users. It…
I’ve recently been doing a lot of both submitting and reviewing pull requests to PyTorch that were authored with substantial LLM assistance. This is a big difference from earlier this year, where it was clear LLMs worked well for greenfield projects but the code was too hopelessly sloppy for a production codebase. Here are my merged PRs that mention claude code in their description; Jason…
Famously, PyTorch and JAX don’t agree on how shardings should be represented: PyTorch takes a mesh-dim oriented view, where for each dimension in your device mesh, you specify what sharding should be applied; JAX takes a tensor-dim oriented view, where for each dimension on your tensor, you say which mesh dimensions (potentially multiple!) shard it. Among my Twitter followers, it is…
I have recently needed to draw the contents of high-dimensional (e.g., 4D and up) tensors where it is important to ensure that is clear how to identify each of the dimensions in the representation. Common strategies I’ve seen people do in this situation include printing a giant list 2D slices (what the default PyTorch printer will do) or flattening the Tensor in some way back down to a 2D…
With contributions from Richard Zou. PT2’s dominant internal representation, FX graphs, do not directly support control flow (if statements, while loops): they only represent straight-line basic blocks. Most of our graph capture mechanisms are tracing based ( fx.symbolic_trace , make_fx , Dynamo), which means that we expect to be able to linearize all conditionals we encounter into a straight line…
When training large scale LLMs, there is a large assortment of parallelization strategies which you can employ to scale your training runs to work on more GPUs. There are already a number of good resources for understanding how to parallelize your models: I particularly recommend How To Scale Your Model and The Ultra-Scale Playbook . The purpose of this blog post is to discuss parallelization…
CuTe is a C++ library that aims to make dealing with complicated indexing easier. A key part of how it does this is by defining a Layout type, which specifies how to map from logical coordinates to physical locations (CuTe likes to say layouts are “functions from integers to integers.”) In fact, CuTe layouts are a generalization of PyTorch strides, which say you always do this mapping…
The purpose of this post is to sum up, in one place, the state of torch.compile for training as of August 2025. Nothing in here isn’t something you might not already know about from elsewhere on the Internet, but we rarely put everything together in one place. The target audience for this document are teams who are evaluating the use of torch.compile for large scale training runs. First, the…
A lot of strong engineers that I know haven’t really taken a serious look at AI coding; they’ve used LLMs to ask questions or write simple scripts and appreciate that it is a useful tool, but haven’t actually tried building a nontrivial application entirely from scratch in vibe coding style (here, I use the term in its original meaning: when you do AI coding without carefully…
Do you use an LLM for coding? Do you maintain a personal benchmark based on problems you have posed the LLM? The purpose of this blog post is to convince you should do this: that you can do so with marginal effort on top of your day-to-day vibe coding and that you will get both short and long term benefits from making your own personal benchmark exist. I started thinking about benchmarks for…
In my previous two posts “`Ways to use torch.compile <http://blog.ezyang.com/2024/11/ways-to-use-torch-compile/>`_” and “`Ways to use torch.export <http://blog.ezyang.com/2024/12/ways-to-use-torch-export/>`_”, I often said that PyTorch would be good for a use case, but there might be some downsides. Some of the downsides are foundational and difficult to remove. But…
Previously, I discussed the value proposition of torch.compile . While doing so, I observed a number of downsides (long compile time, complicated operational model, lack of packaging) that were intrinsic to torch.compile’s API contract, which emphasized being able to work on Python code as is, with minimal intervention from users. torch.export occupies a different spot in the tradeoff space:…
On the surface, the value proposition of torch.compile is simple: compile your PyTorch model and it runs X% faster. But after having spent a lot of time helping users from all walks of life use torch.compile, I have found that actually understanding how this value proposition applies to your situation can be quite subtle! In this post, I want to walk through the ways to use torch.compile, and…
Tensor libraries like PyTorch and JAX have developed compact and accelerated APIs for manipulating n-dimensional arrays. N-dimensional arrays are kind of similar to tables in database, and this results in the logical question which is could you setup a Tensor-like API to do queries on databases that would be normally done with SQL? We have two challenges: Tensor computation is typically uniform…
One of the things that I learned in grad school is that even if you’ve picked an important and unsolved problem, you need some reason to believe it is solvable–especially if people have tried to solve it before! In other words, “What’s different this time?” This is perhaps a dreary way of shooting down otherwise promising research directions, but you can flip it…
One of the annoying things about scraping websites is bouncing back and forth between the browser where you are using Dev Tools to work out what selectors you should be using to scrape out data, and your actual scraping script, which is usually some batch program that may have to take a few steps before the step you are debugging. A batch script is fine once your scraper is up and running, but…
I’m launching a new podcast, the PyTorch Developer Podcast . The idea is to be a place for the PyTorch dev team to do bite sized (10-20 min) topics about all sorts of internal development topics in PyTorch. For now, it’s just me monologuing for fifteen minutes about whatever topic I decide. The plan is to release an episode daily, five days a week, until I run out of things to say…
At Facebook, we have an internal convention for tooling called “rage”. When something goes wrong and you want to report a bug, the tool developer will typically ask you to give them a rage. For a command line tool, this can be done by running a rage subcommand, which will ask about which previous CLI invocation you’d like to report, and then giving you a bundle of logs to send to…
PyTorch is a fairly large and active open source project, and sometimes we have people come to us and ask if there are any lessons from how we run PyTorch that they could apply to their own projects. This post is an attempt to describe some of the processes as of 2021 that help PyTorch operate effectively as an open source project. I won’t claim that everything we do necessarily the best way…
Years ago, Nadav Rotem related to me this story about why basic block procedures in Swift are not as good as they seem. Nelson Elhage reminded me about this on Twitter and so I thought this should be put into the public record. Basic block procedures make certain optimizations more difficult. Consider this program: block j3 (%y1, %y2) { ... } block j1 () { jump j3(%x1, %x2) } block j2 () { jump…
Greetings from 2024! An official pattern matching PEP has been accepted https://peps.python.org/pep-0636/ and is available in Python 3.10. Class patterns are tested using isinstance, with no inheritance structure necessary, making the pattern described in this post 100% forward compatible to real pattern matching. One of the features I miss most in non-Haskell programming languages is algebraic…
If this is your first time reading about PyTorch internals, you might want to check out my PyTorch internals post first. In this post, I want to talk about one particular part of PyTorch’s internals: the dispatcher . At a first glance, the dispatcher is just a glorified if statement: based on some information about the tensor inputs, decide what piece of code should be called. So why should…
For the longest time, I thought of implicit parameters and dynamic scoping were basically the same thing, since they both can be used to solve similar problems (e.g., the so called “configuration problem” where you need to plumb down some configuration deep into a nested body of function definitions without defining them all explicitly). But implicit parameters have a reputation of…
I’ve recently been working on a revamp of how we specify tensor shape formulas in PyTorch. As part of this process, I classified every single operator in PyTorch by its shaping behavior; yes, that’s all 1364 of them (this includes each variant of an operator; e.g., inplace and out= keyword variants). During the process, I tried to come up with categories to help classify what operators…
vmap is an interface popularized by JAX which offers you a vectorizing map. Semantically, a vmap is exactly equivalent to a map in Haskell; the key difference is that operations run under a vmap are vectorized. If you map a convolution and a matrix multiply, you will have one big loop which repeatedly calls convolution and matrix multiply for each entry in your batch. If you vmap a convolution and…
This post is a long form essay version of a talk about PyTorch internals, that I gave at the PyTorch NYC meetup on May 14, 2019. Hi everyone! Today I want to talk about the internals of PyTorch . This talk is for those of you who have used PyTorch, and thought to yourself, “It would be great if I could contribute to PyTorch,” but were scared by PyTorch’s behemoth of a C++…
Some notes collected from a close read of Conal Elliot’s Compiling to Categories and The Simple Essence of Automatic Differentiation . A colleague of mine was trying to define a “tree structure” of tensors, with the hope of thereby generalizing the concept to also work with tensors that have “ragged dimensions.” Let’s take a look: Suppose we have a (2, 3)…
Long time readers of mine may be aware that I used a ThinkPad X61T for the past decade. After the hinge on my second instance of the machine, I decided it was finally time to get a new laptop. And I had one particular model on my eye, after Simon Peyton Jones showed me his new laptop at the last Haskell Implementor’s Workshop: the Microsoft Surface Book 2. It fits my primary requirement for…
My name is Rahul Muttineni, CTO of TypeLead, working on building services around a language named Eta. To get started, I'll give an overview of how the project started, and where it is now. It started as a HSOC project. It was called GHCVM; back then we had plans of making it both on JVM and CLR... we don't think about CLR anymore. I was mentored by Edward Kmett. We got pretty good response on…
It’s been a year since I got my hood and gown and joined Facebook (where I’ve been working on PyTorch), but while I’ve been at Facebook Backpack hasn’t been sleeping; in fact, there’s been plenty of activity, more than I could have ever hoped for. I wanted to summarize some of the goings on in this blog post. Libraries using Backpack ¶ There’s been some really…
A run-time debugger allows you to see concrete values in a program, make changes to them, and continue running your program. A compile-time debugger allows you to see symbolic values in a program, reason about them, and write the rest of your program, e.g. filling in missing tensor size checks, for example. Here’s an example of a compiler-time debugger in action. Let’s suppose you are…
Raise your hand if you’ve ever put one of these commands in your continuous integration scripts: apt install somepackage pip install -r requirements.txt or pip install somepkg conda install blah cabal update or cabal install blah git clone https://github.com/someguy/somerepo wget http://some-website/thingy-latest.tgz Can you tell what the problem is? These commands are not reproducible:…
The best and worst thing about semantic import versioning is that it makes BC-breaking changes hard. In the past few days, Russ Cox has made a splash in a series of white papers describing Go and Versioning . In them, he coins a new term, Semantic Import Versioning , distilling it to the following principle: If an old package and a new package have the same import path, the new package must be…
JG: Joseph Gonzalez GG: Garth Gibson (CMU) DS: Dawn Song (UC Berkeley) JL: John Langford (Microsoft NY)j YQ: Yangqing Jia (Facebook) SB: Sarah Bird M: Moderator A: Audience M: This workshop is bringing together ML and systems. Can you put your place on that spectrum? Who is your home community? YJ: Right in the middle. I’d like to move more towards systems side, but Berkeley Parallel Labs…
The below is a transcript of a talk by Daniel Lo on BrainWave , at the ML Systems Workshop at NIPS'17. Deploy and serve accelerated DNNs at cloud scale. As we’ve seen, DNNs have enabled amazing applications. Architectures achieve SoTA on computer vision, language translation and speech recognition. But this is challenging to serve in large-scale interactive because there are latency, cost…
The below is a transcript of a talk by Virginia Smith on MOCHA , at the ML Systems Workshop at NIPS'17. The motivation for this work comes from the way we think about solving ML problems in practice is changing. The typical ML workflow looks like this. You start iwth dataset and problem to solve. Say you want to build a classifier to identify high quality news articles. Next step is to select an…