/AI1h ago

Meta PyTorch Unveils Torch_Remat API For Flexible Activation Checkpointing

--0--
Comments
Edward Z. Yang#819
Original post
Edward Z. Yang@ezyang#819inAI

torch_remat most directly competes with SAC, where you specify a policy function which specifies, on a per operation basis, if you want to save or recompute it. torch_remat flips this on its head: instead, you specify if you want to save/recompute on a per op call basis.

Edward Z. Yang@ezyang

I've been experimenting with a new activation checkpointing API, which we're calling torch_remat. We're still putting it through the paces, but I think it's already interesting enough to get some initial public feedback: https://github.com/meta-pytorch/remat

9:39 AM · Jun 3, 2026 · 258 Views
Sentiment
Sentiment building, check back later.
Cluster Engagement
-
Views
-
Comments
-
Reposts
-
Bookmarks
Expand data
Posts from X
Most Activity
Most ActivityTimeline
VIEWS166LIKES2
Edward Z. Yang@ezyang

It is particularly designed for codebases that make a lot of use of custom autograd Functions, which is common in advanced PyTorch codebases that need more control over what is saved for backwards compared to what basic PyTorch operations provide.

Edward Z. Yang@ezyang

torch_remat most directly competes with SAC, where you specify a policy function which specifies, on a per operation basis, if you want to save or recompute it. torch_remat flips this on its head: instead, you specify if you want to save/recompute on a per op call basis.

1hViews 166Likes 2Bookmarks 0
REPLIES1
Edward Z. Yang@ezyang

If you're in a memory constrained regime, you've probably calculated exactly what you can afford to save. torch_remat lets you turn this calculation directly into code!

Edward Z. Yang@ezyang

The overall goal of this API is that, by default, *everything* is recomputed, and then you can very explicitly write out what exactly you want to save for backwards.

1hViews 39Likes 0Bookmarks 0