Recent Posts
Ruisi Zhang (@ruisizhang123) · September 2, 2026
torchtitandistributedmodel factory
TL;DR – I have reproduced 23% of Olmo3’s pre-training loss curve with TorchTitan. The training takes 4 days running on 512H100 GPUs. The on-the-fly downstream task evaluations further confirm TorchTitan’s ability to train capable LLMs. This marks the first step in our model factory efforts, which aims to enable agents to use TorchTitan to hillclimb scaling ladders, explore new research ideas, and validate them across increasing model scales. Stay tuned for more updates.
Olmo3 is Ai2’s fully open family of language models (link), with training data, code, checkpoints, and recipes spanning pretraining, mid-training, long-context training, and post-training. This transparency provides a …
Continue reading →Pian Pawakapan (@pianpwk) · August 26, 2026
torchtitandistributeddtensorspmdsharding
TL;DR
TorchTitan now uses spmd_types as its default backend for distributed model computation. Authors specify sharding contracts and collectives explicitly, with optional typechecking to catch distributed-correctness errors during development. At runtime, the typechecking machinery can be erased so forward and backward execute on plain tensors. Coupled FWD-BWD typing and support for both global and local SPMD also make distributed behavior easier to express and reason about. Across our repeated debug-model benchmarks, spmd_types improved eager throughput by up to 46% and Inductor throughput by 5-9% over partial_dtensor, without a meaningful peak-memory increase. FSDP2 continues to use …
Continue reading →Chien-Chin Huang (@fegin) · August 17, 2026
torchtitandistributeddtensorshardingspmdmoe
TL;DR
TorchTitan now adopts a declarative approach: ALL sharding (SPMD parallelization) is expressed in configuration. We created a full DTensor mode where every tensor is a DTensor that shards on all activated mesh axes (DP, CP, TP; EP for MoE experts). Full DTensor removes the ambiguity where a tensor could be a plain local tensor or a DTensor sharded on only some axes - the source of correctness bugs. Available via --parallelism.spmd_backend=full_dtensor, verified bit-identical to the legacy path across FSDP/HSDP/CP/TP/EP, at performance parity. TorchTitan is transitioning to spmd_types; the work here (config-based sharding and full DTensor) is the foundation that makes it possible. …
Continue reading →Edward Yang (@ezyang) · August 11, 2026
eagercudaperformanceprofiling
Disclosure. This post was drafted by Claude (Anthropic’s coding assistant) with editing from ezyang.
Once you get a bit experienced with writing performant PyTorch code you know to avoid device-to-host syncs: calling .item() is an easy way to become CPU bound in kernel launches afterwards, since we have to wait for the GPU work to finish before we can get the result to the CPU. And with more time, you might find out about a number of API footguns in PyTorch’s API that invisibly cause DtoH syncs, like x[bool_tensor], which implicitly triggers a torch.nonzero sync.
But what you might not be expecting is that host-to-device syncs are bad too! And these syncs can also show up in sneaky ways, …
Continue reading →Edward Yang (@ezyang) · August 9, 2026
eagercudamemorypinned-memorycuda-graphs
Disclosure. This post was drafted by Claude (Anthropic’s coding assistant) with editing from ezyang.
For quite some time, I used to think of pinned memory as something you sprinkled around your code to let you do async transfers to GPU. In fact, the old Caffe2 used to pin every single CPU tensor, even if it never actually participated in GPU compute. Under this regime, you might imagine that you need an allocator for pinned memory that lets you allocate and free pinned memory as necessary.
However, if you think carefully about the implications of async transfers on pinned memory lifetime, as well as the implications for CUDA graphs, it turns out that you don’t… really want to ever free …
Continue reading →Edward Yang (@ezyang) · June 1, 2026
eagercudamemory
Disclosure. This post was drafted by Claude (Anthropic’s coding assistant) with editing from ezyang.
In an ideal world, users of CUDA memory in PyTorch programs should be able to abstract the allocator behavior as: there is a fixed amount of GPU memory, whenever you allocate this available memory goes down, and when you free the available memory goes back up.
Unfortunately, the internal implementation of the CUDA caching allocator means that certain allocation patterns can give rise to fragmentation, where even though there is “technically” enough free space to store a requested allocation, the CUDA caching allocator is unable to actually serve the request.
There are many modern use cases …
Continue reading →Edward Z. Yang (@ezyang) · July 25, 2026
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 …
Continue reading →Laith Sakka (@laithsakka) · July 10, 2026
dynamic_shapesfxsymintcorrectness
TL;DR – If you’re creating FX graph nodes with raw SymInt arguments, use Graph.materialize_symints or the targeted create_size_node / create_stride_node / create_storage_offset_node helpers instead of passing the raw symbolic value directly. This fixes a common, subtle class of correctness bugs — we’ve already found 3 in PyTorch Inductor and 6 across executorch’s ARM backend passes. Passing raw symbolic values will become a hard error soon.
This is a pattern that’s easy to write and hard to catch. Today it emits a warning; it will become a hard error soon, once we land the executorch fixes:
# `val` is the example/fake tensor stored on a node's meta — under dynamic # shapes its …
Continue reading →Sanket Purandare (@sanketpurandare) · July 10, 2026
dynamic_shapesunbackeddistributeddtensorflex_attentioninductortracing
TL;DR – To capture a whole distributed training step as one FX graph, PT2 has to trace models whose batch and sequence dimensions are unbacked SymInts. When we tried this on the TorchTitan DeepSeek-V3 MoE trainer, every layer of the stack either silently specialized those dims to concrete ints or blew up with a data-dependent error (DDE). This post walks through the 11 fixes — spanning ATen meta kernels, ProxyTensor, the ShapeEnv, Inductor, collective bucketing, DTensor, and FlexAttention — that make each layer keep tensor semantics symbolic while allowing hints for policy decisions only. That contract is what unblocked end-to-end symbolic tracing for the graph trainer and expert-parallel …
Continue reading →Laith Sakka (@laithsakka), Xiao Fu (@fxdawnn) · June 24, 2026
dynamic_shapesunbackedexport
TL;DR – A new dynamic shapes API is available and ready to use! It provides a unified, consistent way for specifying dynamic specs across compile, export and make_fx, brings native unbacked support to torch.export and make_fx, and completes the unbacked story described earlier by providing unified, predictable, declarative control over the shapes of compiled artifacts.
Consider the example below: the user has a function project(x, w) with a fast path for small batches and a general matmul path. The user wants to compile a dynamic-shape artifact that takes the fast path.
The ShapesSpec says x has dynamic shape [B, D] and w has shape [D, D], and the assumption B < 32 commits this artifact to …
Continue reading →Edward Yang (@ezyang) · May 30, 2026
ai-agentscode-reviewossllm
One of the important topics being discussed among the PyTorch team is how the PyTorch codebase should engage with AI coding agents. Today, many PRs to PyTorch are AI-authored, and there have been obvious growing pains as we’ve figured things out. Based on discussions at the most recent PyTorch compiler offsite (May 2026), I’ve assembled this playbook for AI coding in PyTorch. It is half descriptive, half prescriptive: it is trying to codify practices that are being used among some members of the team, and bring everyone else along. Hopefully, this post is just the beginning of our ongoing conversation about how to engage with AI coding agents.
We can think of AI generated code as living in a …
Continue reading →Arsh Zahed (@azahed98) · May 15, 2026
dynamotorch.compilegraph-breaksskills
I’m excited to share debug-graph-breaks, a new skill for debugging Torch Compile graph breaks, now available in the meta-pytorch/skills repository.
Torch Compile graph breaks prevent full graph capture and hurt performance. This skill helps you:
Identify root causes of graph breaks Understand why operations break compilation Get actionable fixes with specific code changes Learn best practices for Torch Compile-friendly code The skill is grounded in the Graph Break Website as its knowledge base—improvements to the website directly improve the skill’s quality.
Evaluated on the OSS Model Graph Break Corpus—a collection of real-world graph break scenarios from open-source models.
Evaluation …
Continue reading →Animesh Jain (@anijain2305) · May 13, 2026
dynamocpythonllm-agentsgraph-breakstp-slots
TL;DR – Dynamo’s ad-hoc CPython support creates fragmented graph breaks that are hard to fix — even for LLM agents. By refactoring Dynamo to mirror CPython’s tp_* slot semantics, we make the system systematically auditable and agent-friendly, already lifting CPython test pass rates from 38% to 45% and proactively eliminating classes of graph breaks in frontier models.
Working with frontier training frameworks has surfaced some fundamental issues in Dynamo. The issues broadly fell into four categories:
CPython language gaps: For example, Dynamo supports calling a functools.partial object but did not support hashing it. Insufficient exception messages: One frontier framework had an unusual …
Continue reading →William Wen (@williamwen42) · May 13, 2026
dynamotorch.compilegraph-breaks
torch._dynamo.config.nested_graph_breaks = True has been enabled on all Dynamo and Inductor unit tests (~250 test files). A sweep of the OSS benchmark models with graph breaks shows 81/82 passing with NGB (the single regression is a pre-existing unstable model), with graph break reductions of up to 67% and graph merging in models with complex nested call structures (GNNs, detection models). Dynamo tracing time is neutral or improved for most models, and models with significant graph merging see up to 15% runtime speedup (8% geomean). The remaining goal is to set nested_graph_breaks to True by default.
The nested graph break problem in torch.compile refers to the Dynamo limitation of only …
Continue reading →Sayak Paul (@sayakpaul), Animesh Jain (@anijain2305), Benjamin Bossan (@BenjaminBossan) · May 11, 2026
torch.compilediffusersregional-compilationdynamic-shapesquantizationlora
TL;DR – torch.compile delivers a ~1.5x speedup on Flux-1-Dev with no quality loss. Use compile_repeated_blocks to cut compile latency 7x (67s → 9.6s) while keeping the speedup, enable dynamic=True to avoid recompiles on shape changes, and combine with CPU offloading, NF4 quantization, and LoRA hot-swap without giving up the compiled kernels.
Diffusion pipelines are heavy: Flux-1-Dev in bf16 weighs ~33 GB and a single image takes 6.7s on an H100. torch.compile can fuse kernels and strip Python overhead, but applying it naively to a real pipeline runs into four practical issues:
Compile latency. First-call JIT cost — 67.4s for the full DiT. Graph breaks. Any unsupported op silently slices the …
Continue reading →Edward Yang (@ezyang) · May 3, 2026
cppcpythonllm
One of the lost arts of PyTorch development is the ability to write idiomatic C++ code that interacts with the CPython API. This was a very important skill in the early days of eager PyTorch, since we spent a lot of time moving large chunks of the framework to C++ for speed reasons, but we don’t touch the C++ code that much these days and many members of the team haven’t written any amount of serious C++.
LLMs seriously lower the barrier for writing C++ and dealing with the minutiae of manual memory management in C. But they’re not perfect. So these devlog is to talk about all of the things that I put into the process. A prompt of sorts. It is based off of the experience driving an LLM to …
Continue reading →Edward Yang (@ezyang) · May 3, 2026
citoolingllmmergedogpytorchbot
Disclosure. This post was drafted by Claude (Anthropic’s coding assistant) with editing from ezyang.
mergedog is an entirely vibe-coded small Python harness that takes one approved pytorch/pytorch PR and shepherds it through CI to the point a human can comment @pytorchbot merge. The idea is to use LLMs to deal with some aspects of the drudgery of landing PRs from external contributors:
Pressing the “Approve CI workflows” button (in a secure way!), Waiting for the CI results to come back, Checking if the CI failures are spurious or real, and Fixing simple CI failures that are just due to brain-os. While each of these tasks is individually not onerous, they take up time in aggregate; and it …
Continue reading →Aaron Orenstein (@aorenste) · April 16, 2026
dispatcherdispatch_keysbackendsautocastfunctionalizationtorch_dispatch
I wanted to write about how PT2 does autograd, but that requires understanding eager autograd, which requires understanding the dispatcher. So let’s start there.
Let’s pretend we’re building Torch. Let’s start from first principles with the problems we encounter and how to solve them.
Problem 1: We want to be able to call operators for each backend.
Solution: Polymorphism! We just define a class where we have every operator defined as a virtual method. Backends just implement every operator.
class Torch: def mm(self, a: Tensor, b: Tensor) -> Tensor: ... def einsum(self, equation: str, *operands: Tensor) -> Tensor: ... ... Now I just need to implement Torch for each “real” backend (CPU, Cuda, …
Continue reading →